diff --git a/AGENTS.md b/AGENTS.md index 670cf7653..2573ca737 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -4,11 +4,10 @@ TinyAgents is a Rust 2024 virtual workspace rooted at `Cargo.toml`. It has no compatibility facade crate: consumers depend directly on the focused package -that owns an API. The public packages are `crates/tinyagents-graph/` (durable -typed state graphs), `crates/tinyagents-harness/` (provider-neutral model +that owns an API. The public packages are `crates/tinyagents-harness/` (provider-neutral model calls, tools, middleware, and streaming), `crates/tinyagents-registry/` (the named capability catalog), `crates/tinyagents-session/` (durable session data), -`crates/tinyagents-tasks/` (graph-independent detached-task registry, +`crates/tinyagents-tasks/` (detached-task registry, task stores and orchestration controls), and `crates/tinyagents-orchestration/` (direct and durable subagent invocation, tools, sessions, lifecycle logic, and focused tests). @@ -29,24 +28,22 @@ smallest useful API. Cargo features are package-local. `tinyagents-harness` exposes `sqlite`, `multimodal`, `png-optimize`, `media`, `claude-code`, `langfuse`, `storage-drivers`, and `tracing`; `tinyagents-harness-tools` exposes default -`todos`, `goals`, and `skills`, plus opt-in `subagents`; `tinyagents-graph` -exposes `sqlite`, `storage-drivers`, and `tracing`; session exposes +`todos`, `goals`, and `skills`, plus opt-in `subagents`; session exposes `storage-drivers` and `tracing`; registry exposes `tracing`. The `storage-drivers` features put each store on `tinystoragedrivers` ports. The harness's own `tracing` instrumentation is always compiled in; the `tracing` features forward to `tinytools-agent/tracing`, which enables that crate's `debug!`/`warn!` diagnostics (native-call parsing, malformed arguments). -Integration tests are in `crates/tinyagents-integration-tests/tests/`, covering serialization, graph routing, +Integration tests are in `crates/tinyagents-integration-tests/tests/`, covering serialization, registry binding, streaming, and provider contracts (including live, network-gated tests such as `tests/live_*.rs`). Subagent-specific E2E and live tests live with their owner in `crates/tinyagents-orchestration/tests/`. Runnable usage examples are in -`crates/tinyagents-integration-tests/examples/`, especially `basic_graph.rs`. +`crates/tinyagents-integration-tests/examples/`, especially `openai_chat.rs`. Design notes and module-level specifications live in `docs/`, with `docs/spec/README.md` as the top-level architecture -reference and `docs/modules/` holding per-surface design docs (`graph/`, -`harness/`, `registry/`). A `wiki/` +reference and `docs/modules/` holding per-surface design docs (`harness/` `registry/`). A `wiki/` git submodule holds the published GitHub wiki pages; do not edit it as part of unrelated work, and commit its pointer update separately when it does change. @@ -59,7 +56,7 @@ change. tests, and examples, treating warnings as failures. - `cargo build --workspace --all-targets`: compile all crate targets. - `cargo test --workspace`: run the full test suite. -- `cargo run -p tinyagents-integration-tests --example basic_graph`: run the bundled graph execution example. +- `cargo run -p tinyagents-integration-tests --example openai_chat`: run the bundled graph execution example. Run commands from the repository root unless a future workspace layout changes the crate location. @@ -77,7 +74,7 @@ have a predictable surface. Place integration tests in `crates/tinyagents-integration-tests/tests/` and use descriptive test names such as `serializes_chat_messages`. Add focused tests when changing serialization, -graph routing, tool invocation, or public model request/response shapes. For +tool invocation, or public model request/response shapes. For async behavior, use the existing `tokio` dev dependency rather than introducing another runtime. diff --git a/CHANGELOG.md b/CHANGELOG.md index 969e94347..c601df038 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,18 @@ # Changelog +## 3.0.0 — 2026-10-22 + +**Breaking.** This release completes the TinyAgents slim-down: generic tool +types and tool-call parsing live in TinyTools, while host-specific execution +policy and adapters remain with their hosts. Public crates and modules removed +during the 2.x migration are no longer part of the TinyAgents workspace. + +- Bump the surviving TinyAgents workspace crates to 3.0.0. +- Update the nested TinyTools pin to the 0.6-compatible epic revision. + +The unreleased entries below describe the changes accumulated before this +major release. + ## Unreleased: `AgentProgress::ContextCompacting` **Breaking.** Added `AgentProgress::ContextCompacting` to the public diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 2d26d3586..fb463eb4c 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -17,12 +17,12 @@ cargo test --workspace The bundled example should also run: ```sh -cargo run -p tinyagents-integration-tests --example basic_graph +cargo run -p tinyagents-integration-tests --example openai_chat ``` Cargo features are package-local. `tinyagents-harness` exposes `sqlite`, -`tools`, `multimodal`, and `tracing`; `tinyagents-graph` exposes `sqlite` and -`tracing`; `tinyagents-registry` and `tinyagents-session` expose `tracing`. +`tools`, `multimodal`, and `tracing`; `tinyagents-registry` and +`tinyagents-session` expose `tracing`. To build with a feature enabled, pass it on the relevant package: ```sh @@ -89,7 +89,7 @@ Use concise imperative commit subjects, for example: ```text Add graph route validation tests -Document graph runtime safety boundary +Document harness safety boundary ``` Avoid mixing formatting, refactors, and behavior changes unless they are diff --git a/Cargo.lock b/Cargo.lock index b1aad213a..3313c5266 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -156,7 +156,7 @@ dependencies = [ "maybe-owned", "rustix", "rustix-linux-procfs", - "windows-sys 0.61.2", + "windows-sys 0.60.2", "winx", ] @@ -351,7 +351,7 @@ dependencies = [ "libc", "option-ext", "redox_users", - "windows-sys 0.61.2", + "windows-sys 0.59.0", ] [[package]] @@ -384,7 +384,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" dependencies = [ "libc", - "windows-sys 0.61.2", + "windows-sys 0.59.0", ] [[package]] @@ -1327,7 +1327,7 @@ dependencies = [ "once_cell", "socket2", "tracing", - "windows-sys 0.61.2", + "windows-sys 0.59.0", ] [[package]] @@ -1567,7 +1567,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys", - "windows-sys 0.61.2", + "windows-sys 0.59.0", ] [[package]] @@ -1796,7 +1796,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c3d1e2c7f27f8d4cb10542a02c49005dbd6e93095799d6f3be745fae9f8fedd4" dependencies = [ "libc", - "windows-sys 0.61.2", + "windows-sys 0.60.2", ] [[package]] @@ -1892,7 +1892,7 @@ dependencies = [ "getrandom 0.4.3", "once_cell", "rustix", - "windows-sys 0.61.2", + "windows-sys 0.59.0", ] [[package]] @@ -1917,7 +1917,7 @@ dependencies = [ [[package]] name = "tinyagents-definition" -version = "2.11.0" +version = "3.0.0" dependencies = [ "anyhow", "async-trait", @@ -1930,27 +1930,9 @@ dependencies = [ "tracing", ] -[[package]] -name = "tinyagents-graph" -version = "2.11.0" -dependencies = [ - "async-trait", - "futures", - "serde", - "serde_json", - "tinyagents-harness", - "tinyagents-storage-kit", - "tinyagents-tasks", - "tinyinference-llm", - "tinystoragedrivers-testkit", - "tinytools", - "tokio", - "tracing", -] - [[package]] name = "tinyagents-harness" -version = "2.11.0" +version = "3.0.0" dependencies = [ "anyhow", "async-trait", @@ -1985,8 +1967,8 @@ dependencies = [ "tinystoragedrivers-types", "tinytools", "tinytools-agent", - "tinytools-std", "tokio", + "toml 0.8.23", "tracing", "uuid", "wait-timeout", @@ -1995,14 +1977,25 @@ dependencies = [ [[package]] name = "tinyagents-harness-tools" -version = "2.11.0" +version = "3.0.0" dependencies = [ "anyhow", "async-trait", + "base64 0.23.1", "chrono", "chrono-tz", + "fs2", + "futures-util", + "glob", + "libc", + "log", + "parking_lot", + "regex", + "reqwest", + "rustix", "serde", "serde_json", + "sha2 0.11.1", "tempfile", "tinyagents-harness", "tinyagents-orchestration", @@ -2011,11 +2004,13 @@ dependencies = [ "tinytools", "tokio", "tracing", + "uuid", + "walkdir", ] [[package]] name = "tinyagents-hooks" -version = "2.11.0" +version = "3.0.0" dependencies = [ "anyhow", "async-trait", @@ -2034,7 +2029,7 @@ dependencies = [ [[package]] name = "tinyagents-integration-tests" -version = "2.11.0" +version = "3.0.0" dependencies = [ "anyhow", "async-trait", @@ -2045,7 +2040,6 @@ dependencies = [ "serde_json", "tempfile", "tinyagents-definition", - "tinyagents-graph", "tinyagents-harness", "tinyagents-harness-tools", "tinyagents-orchestration", @@ -2066,7 +2060,7 @@ dependencies = [ [[package]] name = "tinyagents-live" -version = "2.11.0" +version = "3.0.0" dependencies = [ "anyhow", "async-trait", @@ -2079,9 +2073,8 @@ dependencies = [ [[package]] name = "tinyagents-orchestration" -version = "2.11.0" +version = "3.0.0" dependencies = [ - "anyhow", "async-trait", "chrono", "dotenvy", @@ -2093,13 +2086,15 @@ dependencies = [ "tempfile", "thiserror", "tinyagents-definition", - "tinyagents-graph", "tinyagents-harness", "tinyagents-registry", + "tinyagents-runtime", "tinyagents-session", "tinyagents-storage-kit", "tinyagents-tasks", "tinyinference-llm", + "tinystoragedrivers-connect", + "tinystoragedrivers-document", "tinystoragedrivers-testkit", "tinytools", "tokio", @@ -2109,7 +2104,7 @@ dependencies = [ [[package]] name = "tinyagents-registry" -version = "2.11.0" +version = "3.0.0" dependencies = [ "anyhow", "async-trait", @@ -2126,12 +2121,13 @@ dependencies = [ [[package]] name = "tinyagents-runtime" -version = "2.11.0" +version = "3.0.0" dependencies = [ "anyhow", "async-trait", "chrono", "serde_json", + "sha2 0.11.1", "tempfile", "thiserror", "tinyagents-harness", @@ -2140,7 +2136,10 @@ dependencies = [ "tinyagents-session", "tinyagents-storage-kit", "tinyinference-llm", + "tinystoragedrivers", "tinystoragedrivers-api", + "tinystoragedrivers-connect", + "tinystoragedrivers-document", "tinystoragedrivers-testkit", "tinystoragedrivers-types", "tinytools", @@ -2150,7 +2149,7 @@ dependencies = [ [[package]] name = "tinyagents-session" -version = "2.11.0" +version = "3.0.0" dependencies = [ "anyhow", "async-trait", @@ -2166,6 +2165,8 @@ dependencies = [ "tinyagents-storage-kit", "tinyagents-tasks", "tinystoragedrivers-api", + "tinystoragedrivers-connect", + "tinystoragedrivers-document", "tinystoragedrivers-testkit", "tinystoragedrivers-types", "tinytools-agent", @@ -2176,7 +2177,7 @@ dependencies = [ [[package]] name = "tinyagents-storage-kit" -version = "2.11.0" +version = "3.0.0" dependencies = [ "data-encoding", "serde", @@ -2192,7 +2193,7 @@ dependencies = [ [[package]] name = "tinyagents-tasks" -version = "2.11.0" +version = "3.0.0" dependencies = [ "serde", "serde_json", @@ -2205,6 +2206,7 @@ dependencies = [ "tinystoragedrivers-testkit", "tokio", "tracing", + "uuid", ] [[package]] @@ -2342,6 +2344,21 @@ dependencies = [ "url", ] +[[package]] +name = "tinystoragedrivers" +version = "2.1.0" +source = "git+https://github.com/tinyhumansai/tinystoragedrivers?rev=05eaada962decbab33bafa4763d2bd1a67f6b456#05eaada962decbab33bafa4763d2bd1a67f6b456" +dependencies = [ + "tinystoragedrivers-api", + "tinystoragedrivers-blob", + "tinystoragedrivers-connect", + "tinystoragedrivers-document", + "tinystoragedrivers-kv", + "tinystoragedrivers-queue", + "tinystoragedrivers-secrets", + "tinystoragedrivers-types", +] + [[package]] name = "tinystoragedrivers-api" version = "2.1.0" @@ -2366,7 +2383,9 @@ dependencies = [ "futures-core", "serde", "serde_json", + "sha2 0.10.9", "tinystoragedrivers-api", + "tinystoragedrivers-connect", "tinystoragedrivers-types", ] @@ -2439,6 +2458,7 @@ dependencies = [ "serde", "serde_json", "tinystoragedrivers-api", + "tinystoragedrivers-connect", "tinystoragedrivers-types", ] @@ -2531,30 +2551,6 @@ dependencies = [ "tracing", ] -[[package]] -name = "tinytools-std" -version = "0.6.0" -dependencies = [ - "anyhow", - "async-trait", - "base64 0.23.1", - "fs2", - "futures-util", - "glob", - "libc", - "log", - "parking_lot", - "regex", - "reqwest", - "rustix", - "serde_json", - "sha2 0.11.1", - "tinytools", - "tokio", - "tracing", - "walkdir", -] - [[package]] name = "tinyvec" version = "1.13.3" @@ -3096,7 +3092,7 @@ version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" dependencies = [ - "windows-sys 0.61.2", + "windows-sys 0.59.0", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index cbd165d83..7e0045104 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -3,7 +3,6 @@ members = ["crates/*"] default-members = [ "crates/tinyagents-harness", "crates/tinyagents-harness-tools", - "crates/tinyagents-graph", "crates/tinyagents-tasks", "crates/tinyagents-registry", "crates/tinyagents-session", @@ -16,7 +15,7 @@ exclude = ["vendor", "worktrees", "integration-proofs/run-ledger-mongodb"] resolver = "3" [workspace.package] -version = "2.11.0" +version = "3.0.0" edition = "2024" license = "GPL-3.0-only" repository = "https://github.com/tinyhumansai/tinyagents" @@ -41,6 +40,7 @@ tracing = "0.1" tinystoragedrivers-types = { version = "2.1.0", git = "https://github.com/tinyhumansai/tinystoragedrivers", rev = "05eaada962decbab33bafa4763d2bd1a67f6b456" } tinystoragedrivers-api = { version = "2.1.0", git = "https://github.com/tinyhumansai/tinystoragedrivers", rev = "05eaada962decbab33bafa4763d2bd1a67f6b456" } tinystoragedrivers-testkit = { version = "2.1.0", git = "https://github.com/tinyhumansai/tinystoragedrivers", rev = "05eaada962decbab33bafa4763d2bd1a67f6b456" } +tinystoragedrivers = { version = "2.1.0", git = "https://github.com/tinyhumansai/tinystoragedrivers", rev = "05eaada962decbab33bafa4763d2bd1a67f6b456", default-features = false, features = ["sqlite"] } tinystoragedrivers-document = { version = "2.1.0", git = "https://github.com/tinyhumansai/tinystoragedrivers", rev = "05eaada962decbab33bafa4763d2bd1a67f6b456", features = ["sqlite", "file"] } tinystoragedrivers-connect = { version = "2.1.0", git = "https://github.com/tinyhumansai/tinystoragedrivers", rev = "05eaada962decbab33bafa4763d2bd1a67f6b456", features = ["sqlite"] } tinyagents-storage-kit = { path = "crates/tinyagents-storage-kit" } diff --git a/README.md b/README.md index 751c12b45..331e9d244 100644 --- a/README.md +++ b/README.md @@ -9,13 +9,10 @@ License: GPL v3

-TinyAgents is a small, provider-neutral agent harness for Rust, plus a durable -typed state-graph runtime. It takes its shape from -[LangChain](https://www.langchain.com/) (models, tools, middleware, structured -output, streaming, usage/cost) and -[LangGraph](https://www.langchain.com/langgraph) (`START`/`END`, nodes, -conditional edges, channels/reducers, checkpoints, interrupts, subgraphs, time -travel) — rebuilt as ordinary, typed Rust with no hidden magic. +TinyAgents is a small, provider-neutral agent harness for Rust. It takes its +shape from [LangChain](https://www.langchain.com/) (models, tools, middleware, +structured output, streaming, usage/cost) — rebuilt as ordinary, typed Rust +with no hidden magic. It is for Rust services that need to call models and tools in a loop, want that loop to be resumable and inspectable, and would rather not carry a @@ -30,24 +27,19 @@ TinyAgents is a Cargo workspace, not one crate. Depend on the pieces you need: caching, and a Claude Code CLI model adapter with stream-json, session, authentication, and MCP endpoint support. Features: `sqlite`, `storage-drivers`, `multimodal`, `png-optimize` (opt-in PNG re-compression), `media`, - `claude-code` and `langfuse` (both default), `tracing` (forwards to `tinytools-agent/tracing` for its parsing diagnostics; the harness's own instrumentation is always compiled in). + `claude-code` and `langfuse` (both default), `json-schema` (opt-in `JsonSchema` for context config types), and `tracing` (forwards to `tinytools-agent/tracing` for its parsing diagnostics; the harness's own instrumentation is always compiled in). - **`tinyagents-harness-tools`** — reusable implementations for the harness: goals and budgets, per-thread todos, skill tool packs, subagent control tools, time and clarification tools, plan exit, per-call binding, and per-thread locks. Features: `todos`, `goals`, `skills` (all default), and `subagents` (opt-in). See the [crate guide](crates/tinyagents-harness-tools/README.md). -- **`tinyagents-graph`** — a LangGraph-style durable, typed state graph: - `START`/`END`, nodes, conditional edges, `Send` fanout, reducers/channels, - checkpoints, interrupts, subgraphs, and time travel. Features: `sqlite`, - `tracing`. -- **`tinyagents-tasks`** — graph-independent detached-task machinery: task +- **`tinyagents-tasks`** — detached-task machinery: task model, `TaskStore` (in-memory and JSONL), `DetachedTaskRegistry`, restart reconciliation, the durable completion router (push a finished child to its parent, redelivered after a restart), and the `spawn`/`await`/`cancel`/`steer` orchestration tools. - Re-exported from `tinyagents_graph::orchestration` for compatibility. - **`tinyagents-registry`** — a named capability catalog (models, tools, - agents, graphs, and routers), plus an offline model price/capability catalog. + agents, 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, plus JSONL transcripts and the `threads` chat thread/message store. Transcript spend separates charged and @@ -79,7 +71,6 @@ None of the crates are published to crates.io (`publish = false` in every ```toml [dependencies] tinyagents-harness = { git = "https://github.com/tinyhumansai/tinyagents", package = "tinyagents-harness" } -tinyagents-graph = { git = "https://github.com/tinyhumansai/tinyagents", package = "tinyagents-graph" } tinyagents-registry = { git = "https://github.com/tinyhumansai/tinyagents", package = "tinyagents-registry" } # The code samples below build `Message` and provider types directly from # TinyInference, the message/model crate TinyAgents is built on. It is a @@ -87,49 +78,6 @@ tinyagents-registry = { git = "https://github.com/tinyhumansai/tinyagents", pack tinyinference-llm = { git = "https://github.com/tinyhumansai/tinyinference", package = "tinyinference-llm" } ``` -A minimal typed graph — a whole-state agent/tool loop (trimmed from -[`examples/basic_graph.rs`](crates/tinyagents-integration-tests/examples/basic_graph.rs)): - -```rust -use tinyagents_graph::*; -use tinyinference_llm::message::Message; - -#[derive(Clone, Debug)] -struct AgentState { - messages: Vec, - needs_tool: bool, -} - -let graph = GraphBuilder::::overwrite() - .add_node("agent", |mut state: AgentState, _ctx: NodeContext| async move { - state.messages.push(Message::assistant("checking the local tool")); - Ok(NodeResult::Update(state)) - }) - .add_node("tool", |mut state: AgentState, _ctx: NodeContext| async move { - state.messages.push(Message::tool("echo", "tool result")); - state.needs_tool = false; - Ok(NodeResult::Update(state)) - }) - .set_entry("agent") - .add_conditional_edges( - "agent", - |state: &AgentState| if state.needs_tool { "tool".to_string() } else { "done".to_string() }, - [("tool", "tool"), ("done", END)], - ) - .add_edge("tool", "agent") - .compile()?; - -let run = graph.run(AgentState { messages: vec![], needs_tool: true }).await?; -``` - -Run it for real: - -```sh -git clone git@github.com:tinyhumansai/tinyagents.git -cd tinyagents -cargo run -p tinyagents-integration-tests --example basic_graph -``` - A one-shot model call through the harness (`export OPENAI_API_KEY=...` then `cargo run -p tinyagents-integration-tests --example openai_chat`): @@ -141,7 +89,10 @@ use tinyinference_llm::providers::openai::OpenAiModel; let model = OpenAiModel::from_env()?; let mut harness: AgentHarness<()> = AgentHarness::new(); -harness.register_model("openai", Arc::new(model)).set_default_model("openai"); +harness.with_model_resolver(Arc::new(tinyagents_harness::host::FixedModelResolver::named( + "openai", + Arc::new(model), +))); let run = harness .invoke_default(&(), vec![Message::user("What is a Rust trait?")]) @@ -149,16 +100,6 @@ let run = harness println!("{}", run.text().unwrap_or_default()); ``` -## Graph runtime - -`tinyagents-graph` is a durable, typed state graph modeled on LangGraph: -`START`/`END` markers, nodes, static and conditional edges, `Command`-based -routing, `Send` fanout, reducers over named channels, checkpointing (with an -optional `sqlite` backend), interrupts, streaming events, topology export, and -replay/time travel across superstep boundaries. A node can embed another -compiled graph as a subgraph, so a whole workflow can appear as a single step -inside a larger one. - ## Harness `tinyagents-harness` runs the model/tool agent loop: provider-neutral model @@ -180,14 +121,12 @@ changing the outer transcript or treating a preview as a write. `tinyagents_orchestration::subagent::SubAgentJobRegistry`. The `subagents` feature in `tinyagents-harness-tools` exposes `tinyagents_harness_tools::subagents`, which provides -`DriverSubAgentTool`, the host-planned typed dispatcher: its `background` mode -starts a child and returns a stable job id, while `inline` mode waits for the -final result. That module also provides the `SubAgentJobsTool` and -`SubAgentMessageTool` declarations over the shared registry. `SubAgentSession` -covers retained post-completion -conversations, while `SubagentDriver` coordinates durable lifecycle -preparation, execution, pause, resume, and persistence. Teams and workflow DAGs -are intentionally outside this focused crate. +`DriverSubAgentTool`, the host-planned typed dispatcher. Every call starts an +asynchronous job and returns its stable job id; `SubAgentJobsTool` provides +status, wait, cancellation, and steering over the shared registry, while +`SubAgentMessageTool` delivers messages. `SubagentDriver` coordinates durable +lifecycle preparation, execution, pause, resume, and persistence. Teams and +workflow DAGs are intentionally outside this focused crate. ## Session runtime @@ -235,8 +174,8 @@ All live in **`openai_graph_agent`** — provider-backed chat, tool calling, structured output, and a graph-driven agent (all need `OPENAI_API_KEY`). - **`subconscious_loop`** — an offline, testable autonomous closed-loop - harness (see its own - [README](crates/tinyagents-integration-tests/examples/subconscious_loop/README.md)). + harness (see the source and test in + `crates/tinyagents-integration-tests/examples/subconscious_loop/`). ## Documentation diff --git a/ROADMAP.md b/ROADMAP.md index 669675a9a..b14912539 100644 --- a/ROADMAP.md +++ b/ROADMAP.md @@ -8,15 +8,11 @@ build toward a production-grade Rust agent runtime. - typed harness model calls, tools, middleware, structured output, streaming, usage/cost tracking, retry/limits, cache, memory/embeddings, sub-agents, and steering (`harness/`) -- durable typed state graph runtime: `START`/`END`, nodes, conditional - routing, `Command`s, fan-out, reducers/channels, checkpoints, interrupts, - subgraphs, streaming, and topology export (`graph/`) - per-thread `ThreadGoal` and `TodoList` productivity primitives, exposed as harness tools -- named capability registry (models, tools, agents, graphs, stores, +- named capability registry (models, tools, agents, stores, middleware, policy) bound by name (`registry/`) -- an optional SQLite-backed checkpointer (`sqlite` feature) -- an embedded Langfuse client and graph exporter for observability +- an embedded Langfuse client for observability ## Near-Term Work @@ -31,14 +27,11 @@ build toward a production-grade Rust agent runtime. ranks the correctness fixes and feature gaps found by comparing TinyAgents with LangGraph, Pydantic AI and pi -## Parallel Agents And Sub-Agents +## Sub-Agents -Shipped: forked child contexts, shared caches with explicit isolation policy, -child event namespaces, parent/child run ids, deterministic reducer-based -merges, optional/blocking/race/quorum/fallback/compare policies, and -resumable checkpoints across parallel branches. Ongoing work focuses on -hardening edge cases surfaced by the `e2e_parallel_*` and `live_subagent_*` -test suites. +Shipped: child contexts, shared caches with explicit isolation policy, +child event namespaces, parent/child run ids, policy-controlled detached +work, and durable child task lifecycle through `tinyagents-orchestration`. ## Stability diff --git a/crates/tinyagents-definition/src/execution_spec.rs b/crates/tinyagents-definition/src/execution_spec.rs index 4ddd2ba65..0fb0a0327 100644 --- a/crates/tinyagents-definition/src/execution_spec.rs +++ b/crates/tinyagents-definition/src/execution_spec.rs @@ -74,12 +74,7 @@ pub enum SandboxMode { /// The tools through which an agent spawns sub-agents of its own choosing. /// A named scope that lists none of them delegates only to what its host /// declared, never to an inherited roster. -pub const SPAWN_TOOL_NAMES: &[&str] = &[ - "spawn_subagent", - "spawn_async_subagent", - "spawn_parallel_agents", - "spawn_worker_thread", -]; +pub const SPAWN_TOOL_NAMES: &[&str] = &["spawn_subagent"]; impl ToolScope { /// Whether this is a named scope that lists one of [`SPAWN_TOOL_NAMES`]. diff --git a/crates/tinyagents-definition/src/execution_spec_tests.rs b/crates/tinyagents-definition/src/execution_spec_tests.rs index 54a4c8ad8..0dae6e2ee 100644 --- a/crates/tinyagents-definition/src/execution_spec_tests.rs +++ b/crates/tinyagents-definition/src/execution_spec_tests.rs @@ -4,6 +4,15 @@ fn named(names: &[&str]) -> ToolScope { ToolScope::Named(names.iter().map(|n| n.to_string()).collect()) } +#[test] +fn spawn_tool_roster_contains_only_the_async_subagent_tool() { + assert_eq!(SPAWN_TOOL_NAMES, &["spawn_subagent"]); + assert!(named(&["spawn_subagent"]).names_a_spawn_tool()); + assert!(!named(&["spawn_async_subagent"]).names_a_spawn_tool()); + assert!(!named(&["spawn_parallel_agents"]).names_a_spawn_tool()); + assert!(!named(&["spawn_worker_thread"]).names_a_spawn_tool()); +} + #[test] fn a_named_scope_listing_a_spawn_tool_is_detected() { for spawn in SPAWN_TOOL_NAMES { diff --git a/crates/tinyagents-definition/tests/fixtures/openhuman/orchestrator.toml b/crates/tinyagents-definition/tests/fixtures/openhuman/orchestrator.toml index ca2f8dc0f..403fb345c 100644 --- a/crates/tinyagents-definition/tests/fixtures/openhuman/orchestrator.toml +++ b/crates/tinyagents-definition/tests/fixtures/openhuman/orchestrator.toml @@ -243,7 +243,7 @@ named = [ # listing it here (and teaching it in prompt.md) only trained the model # to call a tool that does not exist. # Sub-agent surface, deliberately three tools (#5701). A sub-agent is - # always async: `spawn_async_subagent` returns a task id immediately and + # always async: `spawn_subagent` returns a task id immediately and # the runner delivers the result back on an idle system turn via # `orchestration::background_delivery`. Nothing has to block for it, so # the collection surface that used to sit here — `wait`, `wait_loop`, @@ -254,7 +254,7 @@ named = [ # were fleet-management affordances for a fleet the model no longer # enumerates; a worker that needs different instructions is cheaper to # respawn, and idle sessions are reaped against the registry cap anyway. - "spawn_async_subagent", + "spawn_subagent", # Fleet awareness. Re-enumerates the workers this parent has in flight # (running / awaiting_user / completed-uncollected / failed) plus reusable # durable sessions — the recovery move when the `[async_subagent_ref]` diff --git a/crates/tinyagents-graph/Cargo.toml b/crates/tinyagents-graph/Cargo.toml deleted file mode 100644 index 9a43cd133..000000000 --- a/crates/tinyagents-graph/Cargo.toml +++ /dev/null @@ -1,45 +0,0 @@ -[package] -publish = false -name = "tinyagents-graph" -version.workspace = true -edition.workspace = true -license.workspace = true -repository.workspace = true -rust-version.workspace = true -description = "Durable typed state graphs for TinyAgents." - -[dependencies] -async-trait = { workspace = true } -futures = { workspace = true } -serde = { workspace = true } -serde_json = { workspace = true } -tinyagents-tasks = { path = "../tinyagents-tasks", version = "2.11.0", default-features = false } -tinyagents-harness = { path = "../tinyagents-harness", version = "2.11.0", default-features = false, features = [ - "langfuse", -] } -tinyinference-llm = { path = "../../vendor/tinyinference/crates/tinyinference-llm", version = "0.5.0" } -tinytools = { path = "../../vendor/tinytools/crates/tinytools", version = "0.6.0" } -tokio = { workspace = true, features = ["sync", "time", "macros", "rt", "rt-multi-thread", "fs"] } -tracing = { workspace = true } -# Storage ports (no database client) for `DriverCheckpointer`. -tinyagents-storage-kit = { workspace = true } - -[features] -default = [] -# Retained as a no-op: SQLite persistence moved to the tinystoragedrivers -# SQLite driver, which a host passes to `DriverCheckpointer` / `DriverTaskCache`. -sqlite = [] -# `DriverCheckpointer`: checkpoints, pending writes and leases on any -# tinystoragedrivers backend the host opened. -storage-drivers = [] -# 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] -tinystoragedrivers-testkit = { 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/README.md b/crates/tinyagents-graph/src/README.md deleted file mode 100644 index 82106d49a..000000000 --- a/crates/tinyagents-graph/src/README.md +++ /dev/null @@ -1,81 +0,0 @@ -# graph - -TinyAgents' durable workflow runtime (LangGraph-style), and one of the -load-bearing surfaces of the crate's recursive architecture. - -Because a node can embed another compiled graph (`subgraph`) or invoke a -sub-agent, **graphs run graphs** and orchestration recurses while every step -stays typed, checkpointed, and observable. Workflows driven from a host -orchestrator use these same runtime types. - -Each submodule keeps type definitions in `types.rs`, behavior in `mod.rs`, and -unit tests in a sibling `_tests.rs` (`mod_tests.rs` beside a `mod.rs`); complex submodules additionally -carry their own `README.md` — see the module map below. - -## Module map - -| Module | Concern | -| --- | --- | -| `builder` | Authoring/compile contract: `GraphBuilder` accumulates nodes, edges, conditional routing, and a reducer; `.compile()` validates topology and freezes it into an immutable `CompiledGraph`. See [`builder/README.md`](builder/README.md). | -| `channel` | Channel-per-field state model (additive) — per-field merge rules and concurrent-write conflict detection, running on the unmodified executor. See [`channel/README.md`](channel/README.md). | -| `checkpoint` | The `Checkpointer` trait and backends (file, sqlite, in-memory) — durability that makes runs resumable and time-travelable. See [`checkpoint/README.md`](checkpoint/README.md). | -| `command` | `Command`, `Interrupt`, `NodeResult`, `RouteTarget`, `Send` — the vocabulary a node handler returns to update state, route, interrupt, or fan out. See [`command/README.md`](command/README.md). | -| `compiled` | The superstep executor, `CompiledGraph` — sequential/parallel steps, node retry, resumable failure, run/resume/state APIs. See [`compiled/README.md`](compiled/README.md). | -| `dag` | Structural validation (duplicates, dangling edges, cycles) of a borrowed dependency-node view, independent of the graph runtime. See [`dag/README.md`](dag/README.md). | -| `export` | Graph introspection/visualization: topology extraction, Mermaid/JSON export, validation reports. | -| `goals` | A durable per-thread goal (single "completion contract"), continuation loop, and harness tools. See [`goals/README.md`](goals/README.md). | -| `observability` | Durable graph observability: journals, status stores, the journaling sink, latency/health rollups, Langfuse export. See [`observability/README.md`](observability/README.md). | -| `orchestration` | Detached-task runtime types backed by a `TaskStore`. The runtime lives in `tinyagents-tasks` ([README](../../tinyagents-tasks/src/README.md)); this module re-exports its current public surface for compatibility. Model-facing orchestration tool adapters were removed. | -| `parallel` | `map_reduce` — ordered, bounded-concurrency parallel map/reduce with a configurable failure policy, independent of the graph executor. | -| `recursion` | Recursion policy and depth tracking: `RecursionFrame`/`RecursionPolicy`/`RecursionStack`/`RunTree` bound and observe nested graph/subgraph/sub-agent recursion. See [`recursion/README.md`](recursion/README.md). | -| `reducer` | `StateReducer`/`Reducer` implementations (overwrite, append, min/max, set-union, closures) that fold branch updates into committed state at a superstep boundary. See [`reducer/README.md`](reducer/README.md). | -| `status` | `GraphRunStatus` — a compact run-status snapshot. See [`status/README.md`](status/README.md). | -| `stream` | `GraphEvent`, `GraphEventSink`, and streaming modes — the live, in-process event surface `observability` makes durable. See [`stream/README.md`](stream/README.md). | -| `subagent_node` | Embeds a harness agent as a graph node (the graph-level analogue of `subgraph`, but for agents instead of graphs). | -| `subgraph` | Embeds a `CompiledGraph` as a node (shared-state or adapter mode) — graph-level recursion. See [`subgraph/README.md`](subgraph/README.md). | -| `testkit` | Deterministic node doubles, event recorder, fluent run assertions, storage conformance suites. See [`testkit/README.md`](testkit/README.md). | -| `todos` | A per-thread todo checklist and its harness tool. See [`todos/README.md`](todos/README.md). | - -## How the pieces fit together - -1. **Author**: `GraphBuilder` (in `builder`) accumulates nodes (`NodeHandler`), - edges, conditional `Route`s, and a `StateReducer`, then `.compile()` - validates the topology and produces an immutable `CompiledGraph` (in - `compiled`). -2. **Run**: `CompiledGraph::run`/`run_with_thread` drives supersteps. Each - step's node handlers return `NodeResult`s built from `command` vocabulary - (`Command::Update`, `Command::Goto`, `Interrupt`, `Send` for fanout); the - step boundary folds results through the configured reducer (`reducer`, or - `channel` for per-field semantics). -3. **Persist**: when a `Checkpointer` (`checkpoint`) is configured, each - boundary is checkpointed, making the run resumable, forkable, and - time-travelable; `status::GraphRunStatus` gives a compact live snapshot. -4. **Observe**: `stream::GraphEvent`s fan out live through a `GraphEventSink`; - `observability` makes that history durable (journals, status stores, - Langfuse export) and derives latency/health rollups. -5. **Recurse**: a node can embed another compiled graph (`subgraph`) or a - harness agent (`subagent_node`), or hand off named child work through - `orchestration`'s tool surface; `recursion` tracks depth and the run tree - across all three so nesting stays bounded and observable. -6. **Test**: `testkit` provides node doubles, a recorder, fluent assertions, - and storage conformance suites so both hand-written graphs and new storage - backends can be verified against the same contracts. - -## Errors - -Graph-specific failures surface through the crate's shared `Result`/ -`TinyAgentsError` (see `crate::error`); the most graph-specific variants are -`TinyAgentsError::RecursionLimit` (nesting/loop depth exceeded) and -`TinyAgentsError::InvalidConcurrentUpdate` (a same-step concurrent write to a -channel that does not allow it — see `channel`). - -## Operational constraints - -- Checkpoints are written **only at superstep boundaries**, never mid-node — - a node handler must be safe to re-run from scratch after a crash. -- Subgraph and sub-agent embedding always extends the child's checkpoint - namespace with the embedding node id; inspecting a child's checkpoints - directly requires accounting for that namespace suffix. -- Parallel-step determinism depends on branches never mutating shared state - outside their own snapshot clone; the reducer, not lock ordering, is what - resolves conflicting concurrent writes. diff --git a/crates/tinyagents-graph/src/agent_loop/compile.rs b/crates/tinyagents-graph/src/agent_loop/compile.rs deleted file mode 100644 index 4c0669880..000000000 --- a/crates/tinyagents-graph/src/agent_loop/compile.rs +++ /dev/null @@ -1,172 +0,0 @@ -//! [`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 ctx_guard = rt.ctx.lock().await; - let mut run_guard = rt.run.lock().await; - runtime::settle_node(&harness, &mut ctx_guard, &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 deleted file mode 100644 index ab31865ff..000000000 --- a/crates/tinyagents-graph/src/agent_loop/driver.rs +++ /dev/null @@ -1,343 +0,0 @@ -//! [`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::{self, 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 tinyagents_harness::terminal::{TerminalOutcome, TerminalReason}; -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() - }; - ctx.reset_turn_tracker(loop_state.messages.len()); - let mut current: &str = node::PLAN; - let mut limit_stop = false; - let mut limit_kind = None; - - 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, ctx, 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(), - )); - } - }; - limit_stop |= loop_state.limit_stop; - // Latch the first kind: a later command must not erase it. - limit_kind = limit_kind.or(loop_state.limit_kind); - 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), - } - }; - - let terminal = match &outcome { - Ok(None) => { - let reason = if limit_stop { - TerminalOutcome::limit_reached( - limit_kind, - format!( - "stopped with the partial run: {} limit reached", - limit_kind.map_or("run", |kind| kind.as_str()) - ), - ) - } else { - TerminalOutcome::completed() - }; - Some(reason.with_provider_started(ctx.provider_started())) - } - 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 outcome = if let Some(summary) = ctx.take_halted_by_guard() { - TerminalOutcome::halted(summary) - } else { - TerminalOutcome::new( - TerminalReason::Paused, - reason - .clone() - .unwrap_or_else(|| format!("paused at node `{}`", interrupt.node)), - ) - }; - Some(outcome.with_provider_started(ctx.provider_started())) - } - Err(error) => { - use tinyagents_harness::terminal::TimeoutPhase; - let in_model_call = ctx.model_call_failed() || ctx.active_model_call.is_some(); - let site = if in_model_call && ctx.call_provider_started() { - TimeoutPhase::Provider - } else if in_model_call { - TimeoutPhase::BeforeProvider - } else if ctx.provider_started() { - TimeoutPhase::AfterTurn - } else { - TimeoutPhase::BeforeProvider - }; - // A `LimitExceeded` carries only text; the cap that tripped was - // announced by a `LimitReached` event just before. - let kind = matches!(error, TinyAgentsError::LimitExceeded(_)) - .then(|| ctx.peek_last_limit()) - .flatten(); - let mut outcome = TerminalOutcome::from_error(error, site).with_limit_kind(kind); - // A failed summarizer already received a provider response; keep the - // match for summarizers that report usage but never marked dispatch - // (host-defined `Summarizer` impls cannot call `mark_dispatched`). - outcome.provider_started = ctx.provider_started() - || matches!(error, TinyAgentsError::SummarizationUsage { .. }); - Some(outcome) - } - }; - // Announce whatever the last turn appended and close it, on every exit - // path, as the direct loop does before the transcript moves onto the - // run (`run.messages` is the transcript as of the last node boundary). - phases::lifecycle_close_turn(harness, ctx, &run.messages); - run.terminal = terminal.clone(); - status.mark_running(HarnessPhase::Middleware); - let after_agent = harness.middleware().run_after_agent(ctx, state, run).await; - // `after_agent` may post-process `run.messages`; announce anything it - // appended (even if it then failed), as the direct loop does. - phases::lifecycle_flush(harness, ctx, &run.messages); - if let Err(hook_error) = after_agent { - if outcome.is_err() { - // The originating node failure stays authoritative, as in the - // direct loop; the hook still ran for its cleanup. - tracing::warn!( - target: "tinyagents::agent_loop", - error = %hook_error, - "[agent_loop] after_agent failed after a node error; keeping the node error" - ); - } else { - return Err(hook_error); - } - } - - // `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) => { - // `after_agent` may have replaced the outcome; report the final one. - let outcome = run - .terminal - .clone() - .or(terminal) - .expect("terminal set before after_agent"); - let record = ctx.emit(AgentEvent::RunCompleted { - run_id: ctx.run_id().clone(), - outcome: Some(outcome), - }); - 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); - if !matches!( - run.terminal.as_ref().map(|outcome| outcome.reason), - Some(TerminalReason::Halted) - ) { - 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 deleted file mode 100644 index 431aa1b04..000000000 --- a/crates/tinyagents-graph/src/agent_loop/iter.rs +++ /dev/null @@ -1,287 +0,0 @@ -//! [`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 ctx_guard, &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 deleted file mode 100644 index 4ff4ec5ac..000000000 --- a/crates/tinyagents-graph/src/agent_loop/mod.rs +++ /dev/null @@ -1,133 +0,0 @@ -//! 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). [`tinyagents_harness::runtime::RunPolicy::execution`] -//! defaults to [`tinyagents_harness::runtime::LoopExecution::Direct`], so -//! every existing caller is unaffected unless it opts in. -//! -//! # Lifecycle events -//! -//! The node bodies announce `TurnStarted`, `TurnCompleted` and -//! `MessageAppended` at the same points as the direct loop, through -//! [`tinyagents_harness::agent_loop::phases`]'s `lifecycle_*` functions: pending -//! appends and the turn open right before `ModelStarted`, the assistant reply -//! after the model call, the tool results (and the turn close) after a tool -//! batch, and the final close in `settle`. The transcript present at the first -//! `plan` activation is the seed and is never announced. Nested tool calls never -//! reach the transcript, so they never produce `MessageAppended`. -//! -//! # 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`]'s `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 deleted file mode 100644 index eae609349..000000000 --- a/crates/tinyagents-graph/src/agent_loop/runtime.rs +++ /dev/null @@ -1,951 +0,0 @@ -//! 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 { - // Reached only when the wrap onion elected to call the provider. - ctx.mark_provider_started(); - 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); - } - // Everything on the transcript when the loop is first entered is input (or - // a resumed run's already-announced history); only later appends are - // announced. A no-op after the first activation. - phases::lifecycle_seed(ctx, loop_state.messages.len()); - 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()); - } - // Same run-policy reasoning default the direct loop attaches. - if let Some(reasoning) = harness.policy().default_reasoning.as_ref() - && !reasoning.is_empty() - { - request.reasoning = Some(reasoning.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`. -async fn model_node_inner( - 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; - loop_state.limit_stop = true; - loop_state.limit_kind = Some(tinyagents_harness::events::LimitKind::ModelCalls); - 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())); - } - - phases::lifecycle_seed(ctx, loop_state.messages.len()); - phases::lifecycle_resume(ctx, loop_state.turn, None); - let entry_len = loop_state.messages.len(); - 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 call_id = CallId::new(format!("{}-model-{}", ctx.run_id(), run.model_calls + 1)); - - let mut request = request; - let resolution_cache = Arc::new(std::sync::Mutex::new(Some(( - request.model.clone(), - request.model_hints.clone(), - request.required_capabilities.clone(), - binding, - )))); - // Mirror the direct loop: `before_model` middleware reads the target - // model's profile (e.g. to avoid new system messages on a model that - // hoists them, #6962). - ctx.model_profile = resolution_cache - .lock() - .unwrap() - .as_ref() - .and_then(|(_, _, _, binding)| binding.model.profile().cloned()); - let profile_cache = resolution_cache.clone(); - harness - .middleware() - .run_before_model_with_profile( - ctx, - app_state, - &mut request, - harness, - move |harness, _ctx, request| { - let profile_cache = profile_cache.clone(); - Box::pin(async move { - let key = ( - request.model.clone(), - request.model_hints.clone(), - request.required_capabilities.clone(), - ); - if let Some((cached_model, cached_hints, cached_capabilities, binding)) = - profile_cache.lock().unwrap().as_ref() - && *cached_model == key.0 - && *cached_hints == key.1 - && *cached_capabilities == key.2 - { - return Ok(binding.model.profile().cloned()); - } - let binding = harness.models().resolve_request(request, None, None); - let profile = binding - .as_ref() - .and_then(|binding| binding.model.profile().cloned()); - *profile_cache.lock().unwrap() = - binding.map(|binding| (key.0, key.1, key.2, binding)); - Ok(profile) - }) - }, - ) - .await?; - - // Middleware may select a different registered model or add required - // capabilities. Resolve again so graph execution dispatches the request - // that middleware actually prepared, just like the direct loop. - let key = ( - request.model.clone(), - request.model_hints.clone(), - request.required_capabilities.clone(), - ); - let binding = resolution_cache - .lock() - .unwrap() - .take() - .and_then( - |(cached_model, cached_hints, cached_capabilities, binding)| { - (cached_model == key.0 && cached_hints == key.1 && cached_capabilities == key.2) - .then_some(binding) - }, - ) - .or_else(|| harness.models().resolve_request(&request, None, None)) - .ok_or_else(|| { - TinyAgentsError::ModelNotFound( - request.model.clone().unwrap_or_else(|| "".into()), - ) - })?; - ctx.model_profile = binding.model.profile().cloned(); - tinyagents_harness::middleware::library::rehome_ephemeral_system_instructions( - &mut request, - ctx.model_profile.as_ref(), - ); - let model_name = binding.resolved.name.clone(); - - // Mirror the direct loop: a named effort picks up the resolved model's - // tuned `thinking_level_map` entry, unless an explicit budget is pinned. - 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()); - } - - // Same point as the direct loop: pending appends (steering) are announced, - // the previous turn closed, and this one opened, just before `ModelStarted`. - let turn = phases::lifecycle_start_turn(harness, ctx, &loop_state.messages); - tracing::debug!( - target: "tinyagents::agent_loop", - run_id = %ctx.run_id(), - turn, - "[graph_loop] turn started" - ); - - let started_record = ctx.emit(AgentEvent::ModelStarted { - call_id: call_id.clone(), - model: model_name.clone(), - }); - status.set_last_event(started_record.id); - status.active_model_call = Some(call_id.clone()); - ctx.active_model_call = Some(call_id.clone()); - ctx.begin_model_call(); - let base = DirectModelBase { - model: binding.model.as_ref(), - }; - let wrapped = match harness - .middleware() - .run_wrapped_model(ctx, app_state, request, &base) - .await - { - Ok(wrapped) => wrapped, - Err(error) => { - status.active_model_call = None; - ctx.active_model_call = None; - ctx.mark_model_call_failed(); - return Err(error); - } - }; - let (mut response, wrap_control) = wrapped.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; - status.active_model_call = None; - ctx.active_model_call = None; - 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; - phases::lifecycle_flush(harness, ctx, &loop_state.messages); - - 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() { - let result = apply_control(ctx, &mut loop_state, control, node::MODEL, route); - retract_on_interrupt(harness, ctx, &result, &loop_state.messages, entry_len, true); - return result; - } - // 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 `model` node: [`model_node_inner`], closing the turn it opened if it -/// fails. `GraphLoopDriver` closes on every exit as well, but `LoopIter` and -/// the compiled graph propagate a node error with no epilogue. -pub(crate) async fn model_node( - harness: &AgentHarness, - app_state: &State, - ctx: &mut RunContext, - run: &mut AgentRun, - status: &mut HarnessRunStatus, - loop_state: LoopState, -) -> Result> -where - State: Send + Sync, - Ctx: Send + Sync, -{ - let entry = loop_state.messages.clone(); - let result = model_node_inner(harness, app_state, ctx, run, status, loop_state).await; - if result.is_err() { - phases::lifecycle_close_turn(harness, ctx, &entry); - } - result -} - -/// The `tools` node: [`tools_node_inner`], announcing the results of calls that -/// ran before a failure and closing the turn if it fails. -pub(crate) async fn tools_node( - harness: &AgentHarness, - app_state: &State, - ctx: &mut RunContext, - run: &mut AgentRun, - status: &mut HarnessRunStatus, - loop_state: LoopState, -) -> Result> -where - State: Send + Sync, - Ctx: Send + Sync, -{ - let entry = loop_state.messages.clone(); - let result = tools_node_inner(harness, app_state, ctx, run, status, loop_state).await; - if result.is_err() { - // A batch error leaves the partial transcript on the run. - let messages = if run.messages.len() > entry.len() { - &run.messages - } else { - &entry - }; - phases::lifecycle_close_turn(harness, ctx, messages); - } - result -} - -/// 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. -async fn tools_node_inner( - 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, -{ - phases::lifecycle_seed(ctx, loop_state.messages.len()); - let entry_len = loop_state.messages.len(); - // The turn that issued these calls began at (or before) the assistant - // message; a fresh runtime resuming an interrupted batch has no open turn. - phases::lifecycle_resume(ctx, loop_state.turn, Some(entry_len.saturating_sub(1))); - 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; - let outcome = match outcome { - Ok(outcome) => outcome, - Err(error) - if matches!( - harness.policy().limits.behavior, - tinyagents_harness::limits::LimitBehavior::StopWithPartial - ) && matches!(error, TinyAgentsError::LimitExceeded(_)) - && ctx.peek_last_limit() - == Some(tinyagents_harness::events::LimitKind::ToolCalls) => - { - // Only the tool-cap admission announces `LimitReached(ToolCalls)` - // (and `record_tool_call` cleared any earlier kind first), so a - // `LimitExceeded` raised by middleware is not a partial stop. - loop_state.limit_stop = true; - loop_state.limit_kind = Some(tinyagents_harness::events::LimitKind::ToolCalls); - 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)); - } - Err(error) => { - // Results of calls that ran before the failure are on the node's - // transcript, which the error path would otherwise drop; keep them - // on the run, as the direct loop does, so the driver's final - // lifecycle close announces and counts them. - run.messages = loop_state.messages.clone(); - return Err(error); - } - }; - 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() { - let result = apply_control(ctx, &mut loop_state, control, node::TOOLS, node::PLAN); - // The tool turn stays open on an interrupt: the re-run closes it with the - // results it produces (a fresh runtime re-opens it via `lifecycle_resume`). - if !retract_on_interrupt( - harness, - ctx, - &result, - &loop_state.messages, - entry_len, - false, - ) { - // Every tool result of this batch is on the transcript: announce - // them and close the turn, as the direct loop does after its batch. - phases::lifecycle_close_turn(harness, ctx, &loop_state.messages); - } - return result; - } - phases::lifecycle_close_turn(harness, ctx, &loop_state.messages); - - 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. -/// An interrupted node's state is discarded and the node re-runs from its entry -/// state on resume, so the appends it already announced are retracted: the -/// re-run announces them again, and events never name a message the kept -/// transcript lacks. -fn retract_on_interrupt( - harness: &AgentHarness, - ctx: &mut RunContext, - result: &Result>, - messages: &[tinyinference_llm::message::Message], - entry_len: usize, - close_turn: bool, -) -> bool { - if !matches!(result, Ok(NodeResult::Interrupt(_))) { - return false; - } - tracing::debug!( - target: "tinyagents::agent_loop", - run_id = %ctx.run_id(), - entry_len, - "[graph_loop] node interrupted; retracting its announced appends" - ); - phases::lifecycle_retract(ctx, entry_len); - // A model node's turn has no results to wait for, and a fresh runtime - // cannot carry this tracker's open turn over, so close it. A tools node - // leaves its turn open for the re-run to close with the real results. - if close_turn { - phases::lifecycle_close_turn(harness, ctx, &messages[..entry_len.min(messages.len())]); - } - true -} - -pub(crate) async fn settle_node( - harness: &AgentHarness, - ctx: &mut RunContext, - run: &mut AgentRun, - mut loop_state: LoopState, -) -> Result> -where - State: Send + Sync, - Ctx: Send + Sync, -{ - phases::lifecycle_seed(ctx, loop_state.messages.len()); - // The final assistant message is the last append of its turn; close the - // turn before any output-retry prompt is pushed (that prompt is announced - // with the next turn's start). - phases::lifecycle_close_turn(harness, ctx, &loop_state.messages); - 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 deleted file mode 100644 index 6d6c67f5d..000000000 --- a/crates/tinyagents-graph/src/agent_loop/types.rs +++ /dev/null @@ -1,126 +0,0 @@ -//! 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. A run that ended in an error reports it - /// through the driver's `Result`, not through this struct. - 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, - /// Set when the run finished because a call cap tripped under - /// `LimitBehavior::StopWithPartial`, so the driver can report - /// `TerminalReason::LimitReached` instead of a plain completion. - #[serde(default)] - pub(crate) limit_stop: bool, - /// Which cap tripped when `limit_stop` is set (`None` when unknown). - #[serde(default)] - pub(crate) limit_kind: Option, -} - -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 deleted file mode 100644 index 8ec7f5032..000000000 --- a/crates/tinyagents-graph/src/builder/README.md +++ /dev/null @@ -1,95 +0,0 @@ -# graph::builder - -The authoring/compile contract: `GraphBuilder` accumulates nodes, edges, -conditional routing, and a reducer; `.compile()` validates the topology and -freezes it into an immutable `CompiledGraph`. - -This is the entry point every workflow uses before it can run. Because a node handler can itself -drive another compiled graph or a sub-agent, the builder is also what a -recursively-generated sub-workflow compiles through, one level down. - -## Public surface - -### `GraphBuilder` - -- `new()` / `overwrite()` (the latter is `State == Update` with the built-in - overwrite reducer) — construction. -- `set_reducer`, `set_entry`, `set_finish` — required wiring: a reducer and an - entry node are validated at `compile()`. -- `add_node`, `add_edge`, `add_sequence`, `add_waiting_edge`, - `add_conditional_edges`, `add_barrier_relief` — topology construction. -- `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` — 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 - `GraphDefaults`. -- `compile() -> Result>` — validates topology - (missing nodes, dangling edges, START/END misuse, command-routing conflicts, - missing reducer) and freezes the builder into an immutable, cheaply-clonable - `CompiledGraph`. - -### Supporting types (`types.rs`) - -- `START` / `END` — reserved virtual entry/terminal node names. -- `NodeHandler` / `NodeFuture` — the boxed async - handler shape every node closure is stored as. -- `NodeContext` — per-task runtime context passed to a handler: run identity, - step, resume value, fork identity (`ForkId`) in a parallel step, `send_arg` - for `Send`/`GraphInput`-scheduled activations, recursion frames, the - `ChildRunSink`, and an optional `AgentInvocationBinding`. -- `ForkId` — identifies one branch of a concurrent superstep (`branch_index` - + `node`). -- `Route` — an optional newtype for a conditional-route label, for routers - that prefer a typed value over a bare string. -- `GraphDefaults` — an optional bundle of per-graph defaults applied in one - `set_defaults` call; every field is additive (`Some` overrides, `None` - leaves the current value). -- `RouterFn` — the conditional-router closure shape. -- `BarrierRelief` (in `mod.rs`) — a relief registration for a mixed fan-in - barrier, so a barrier fed by one unconditional and one conditional-only - predecessor doesn't deadlock when the conditional branch isn't taken. - -## Files - -| File | Role | -| --- | --- | -| `types.rs` | `GraphBuilder` struct, `NodeContext`, `ForkId`, `Route`, `GraphDefaults`, `NodeHandler`/`NodeFuture`/`RouterFn` aliases, `START`/`END`, and the crate-private `NodeMeta`/`BuilderNode`/`Branch`. | -| `mod.rs` | All `GraphBuilder` methods (construction, topology, configuration, `compile`) and `BarrierRelief`. | -| `mod_tests.rs` | Unit tests for the compile contract: reducer requirement, START/END validation, missing-node/route detection, command-routing conflicts. | - -## How it fits together - -- `compile()` is the sole producer of `compiled::CompiledGraph` — see - `compiled/README.md` for what runs after a builder is compiled. -- `command::Command`/`RouteTarget` are what a node handler returns at - runtime; `mark_command_routing`/`with_command_destinations` only declare - the *intent* to route that way so `compile()` can reject a mixed topology - and `export` can draw the advisory destinations — the runtime always - resolves the real successor from the emitted `Command`. -- `reducer::StateReducer` is the trait `set_reducer` accepts; `overwrite()` - is shorthand for the built-in `OverwriteStateReducer`. -- `recursion::RecursionFrame` values seeded into a `NodeContext` come from - `subgraph`/`subagent_node`, not from this module — the builder only carries - the `recursion_limit` cap through to the compiled graph. -## Operational constraints - -- `compile()` fails closed: no reducer, a missing/self-referential entry, - dangling edges, START/END misuse, or a node that mixes command routing with - static/conditional edges all fail validation rather than compiling a graph - that would misbehave at runtime. -- A `BarrierRelief`'s `relief_node` must actually be one of `barrier_node`'s - registered waiting predecessors (via `add_waiting_edge`) — a relief - registered against a barrier with no matching waiting registration is a - silent no-op at execution time, not a compile error. -- `add_sequence` only wires edges between nodes that must already exist via - `add_node`; it does not itself add nodes. diff --git a/crates/tinyagents-graph/src/builder/mod.rs b/crates/tinyagents-graph/src/builder/mod.rs deleted file mode 100644 index e07088467..000000000 --- a/crates/tinyagents-graph/src/builder/mod.rs +++ /dev/null @@ -1,771 +0,0 @@ -//! Durable graph builder and compile contract. -//! -//! This is the authoring entry point for the recursive graph runtime: a -//! [`GraphBuilder`] accumulates nodes, edges, conditional routing, and a reducer, -//! and [`GraphBuilder::compile`] validates that topology and freezes it into an -//! immutable [`crate::CompiledGraph`]. Because a node handler can itself -//! drive another compiled graph or a sub-agent, the same builder API is what -//! hand-written Rust and runtime assembly use when constructing a workflow -//! that may recurse into sub-workflows. -//! -//! See `types` for the builder data types. `compile` validates the topology -//! and freezes it into an immutable [`crate::CompiledGraph`]. - -mod policy; -mod types; - -pub use policy::{CacheKeyFn, NodeCachePolicy, NodePolicy, OnErrorFn}; -pub(crate) use types::{Branch, BuilderNode, NodeMeta, UpdateCodec}; -pub use types::{ - END, ForkId, GraphBuilder, GraphDefaults, IdleClock, NodeContext, NodeFuture, NodeHandler, - Route, RouterFn, START, -}; - -use std::collections::{HashMap, HashSet}; -use std::future::Future; -use std::sync::Arc; - -use crate::command::NodeResult; -use crate::compiled::CompiledGraph; -use crate::reducer::{OverwriteStateReducer, StateReducer}; -use crate::{Result, TinyAgentsError}; -use tinyagents_harness::ids::{GraphId, NodeId}; - -/// A relief registration for a mixed fan-in barrier. -/// -/// A waiting/barrier node (see [`GraphBuilder::add_waiting_edge`]) normally -/// activates only once *every* registered predecessor has arrived. When one -/// of those predecessors (`relief_node`) is only reachable via a conditional -/// branch out of `source`, and `source` routes elsewhere instead, that -/// predecessor never runs and the barrier would wait forever. -/// -/// A [`BarrierRelief`] fixes that without weakening the barrier into a plain -/// edge (which would let a *taken* branch's downstream data race the merge -/// and get silently dropped): when `source` completes a superstep without -/// routing to `relief_node`, the executor registers a phantom arrival of -/// `relief_node` at `barrier_node`, so the barrier can still clear on the -/// predecessors that actually ran. -#[derive(Clone, Debug, PartialEq, Eq)] -pub struct BarrierRelief { - /// The brancher node whose conditional routing determines whether - /// `relief_node` runs. - pub source: NodeId, - /// The conditional-only predecessor of `barrier_node` that `source` may - /// or may not route to. - pub relief_node: NodeId, - /// The mixed fan-in (all-waiting) node gated on `relief_node`'s arrival. - pub barrier_node: NodeId, -} - -impl Default for GraphBuilder -where - State: Clone + Send + Sync + 'static, - Update: Send + 'static, -{ - fn default() -> Self { - Self::new() - } -} - -impl GraphBuilder -where - State: Clone + Send + Sync + 'static, - Update: Send + 'static, -{ - /// Creates an empty builder with a generated graph id and a default - /// recursion limit of 50. A reducer must be set before [`Self::compile`]. - pub fn new() -> Self { - Self { - graph_id: GraphId::new(format!("graph-{}", tinyagents_harness::ids::next_seq())), - name: None, - 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(), - reducer: None, - recursion_limit: 50, - parallel: false, - 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. - pub fn set_defaults(mut self, defaults: GraphDefaults) -> Self { - if let Some(limit) = defaults.recursion_limit { - self.recursion_limit = limit; - } - if let Some(parallel) = defaults.parallel { - self.parallel = parallel; - } - if let Some(max) = defaults.max_concurrency { - self.max_concurrency = Some(max); - } - if let Some(timeout) = defaults.node_timeout { - self.node_timeout = Some(timeout); - } - self - } - - /// Bounds the number of branches run concurrently within a single superstep - /// (only meaningful with [`Self::with_parallel`] enabled). The executor runs - /// the active set in chunks of at most `n` futures, so at most `n` node - /// handlers are in flight at once. `0` is treated as unbounded. - pub fn with_max_concurrency(mut self, n: usize) -> Self { - self.max_concurrency = (n > 0).then_some(n); - self - } - - /// Sets a default wall-clock timeout applied to every node handler. A node - /// whose future does not resolve within `timeout` fails the run with - /// [`crate::TinyAgentsError::Timeout`]. - pub fn with_node_timeout(mut self, timeout: std::time::Duration) -> Self { - self.node_timeout = Some(timeout); - self - } - - /// Enables or disables concurrent execution of the active node set within a - /// superstep. Defaults to `false` (sequential), which preserves the exact - /// milestone-1 ordering and semantics. - /// - /// When enabled, a superstep with more than one active node runs every - /// branch concurrently via `futures::future::join_all`. Each branch gets its - /// own cloned `State` snapshot (`State: Clone`) and a distinct - /// [`ForkId`] on its [`NodeContext`]. Branch results are still folded into - /// the reducer in deterministic active-set order at the step boundary, so a - /// downstream node always observes the same merged state regardless of which - /// branch finished first. See [`crate::CompiledGraph`] for the full - /// concurrency and interrupt semantics. - pub fn with_parallel(mut self, parallel: bool) -> Self { - self.parallel = parallel; - self - } - - /// Overrides the graph id. - pub fn with_graph_id(mut self, id: impl Into) -> Self { - self.graph_id = id.into(); - self - } - - /// Sets an optional human-readable graph name surfaced by the topology - /// export. The `graph_id` remains the stable identifier; the name is purely - /// descriptive (e.g. for diagrams a model authored). - pub fn with_name(mut self, name: impl Into) -> Self { - self.name = Some(name.into()); - self - } - - /// Overrides the recursion limit (max number of supersteps). - pub fn with_recursion_limit(mut self, limit: usize) -> Self { - self.recursion_limit = limit; - self - } - - /// Sets the state reducer used to merge partial updates at step boundaries. - pub fn set_reducer(mut self, reducer: R) -> Self - where - R: StateReducer + 'static, - { - self.reducer = Some(Arc::new(reducer)); - 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, - { - self.nodes.insert( - 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 { - handler: Arc::new(move |state, ctx| Box::pin(handler(state, ctx))), - }, - ); - self - } - - /// 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::push_edge(&mut self.edges, from.into(), to.into()); - self - } - - /// Adds a chain of direct edges over the given nodes: `add_sequence([a, b, - /// c])` is equivalent to `add_edge(a, b).add_edge(b, c)`. The nodes must - /// already have been added with [`Self::add_node`]; this only wires edges. A - /// sequence of fewer than two nodes adds no edges. - pub fn add_sequence(mut self, nodes: I) -> Self - where - I: IntoIterator, - N: Into, - { - let nodes: Vec = nodes.into_iter().map(Into::into).collect(); - for pair in nodes.windows(2) { - 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 - /// different supersteps. This is the join/fan-in primitive for diamond - /// topologies where several branches must finish before a synthesis node - /// runs. Calling it repeatedly with the same `to` accumulates the required - /// predecessor set. - pub fn add_waiting_edge(mut self, from: impl Into, to: impl Into) -> Self { - let from = from.into(); - let to = to.into(); - Self::push_edge(&mut self.edges, from.clone(), to.clone()); - self.waiting.entry(to).or_default().insert(from); - self - } - - /// Registers a barrier relief for a mixed fan-in barrier. - /// - /// When `source` completes a superstep without routing to `relief_node` - /// — because `relief_node` sits behind a conditional branch `source` did - /// not take this run — this records a phantom arrival of `relief_node` - /// at `barrier_node`'s waiting-edge barrier (registered separately via - /// [`Self::add_waiting_edge`]), so an all-waiting merge downstream of a - /// mixed fan-in (one unconditionally reachable predecessor plus one - /// reachable only via a conditional branch) can still clear on its - /// remaining predecessors instead of deadlocking on a branch that will - /// never run. - /// - /// `relief_node` must be one of `barrier_node`'s registered waiting - /// predecessors; a relief for a barrier with no matching waiting - /// registration is a no-op at execution time. - pub fn add_barrier_relief( - mut self, - source: impl Into, - relief_node: impl Into, - barrier_node: impl Into, - ) -> Self { - self.barrier_reliefs.push(BarrierRelief { - source: source.into(), - relief_node: relief_node.into(), - barrier_node: barrier_node.into(), - }); - self - } - - /// Sets the entry node, i.e. `add_edge(START, node)`. - pub fn set_entry(self, node: impl Into) -> Self { - self.add_edge(START, node) - } - - /// Marks `node` as terminal, i.e. `add_edge(node, END)`. - pub fn set_finish(self, node: impl Into) -> Self { - self.add_edge(node, END) - } - - /// Adds conditional edges: a router closure mapped against a label table. - /// - /// Both the router's return value and the route-table labels are - /// `impl ToString`, so a user-defined route enum that implements `Display` - /// (or the [`Route`] newtype) can be used directly without manual - /// `.to_string()` calls. Plain `&str`/`String` labels still work unchanged — - /// the label is resolved against the table by its string form at the step - /// boundary. - pub fn add_conditional_edges( - mut self, - from: impl Into, - router: F, - routes: I, - ) -> Self - where - F: Fn(&State) -> R + Send + Sync + 'static, - R: ToString, - I: IntoIterator, - K: ToString, - V: Into, - { - let routes = routes - .into_iter() - .map(|(k, v)| (k.to_string(), v.into())) - .collect(); - self.branches.insert( - from.into(), - Branch { - 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. - pub fn mark_command_routing(mut self, node: impl Into) -> Self { - self.command_nodes.insert(node.into()); - self - } - - /// Records the declared `goto` destination hints for a command-routing node. - /// - /// These are advisory only — the runtime always resolves the real successor - /// from the [`crate::Command`] a node emits — but they let the - /// [topology export](crate::export) draw and validate the set of - /// nodes a command node may jump to. Implies [`Self::mark_command_routing`]. - pub fn with_command_destinations( - mut self, - node: impl Into, - destinations: I, - ) -> Self - where - I: IntoIterator, - N: Into, - { - let node = node.into(); - self.command_nodes.insert(node.clone()); - let dests = destinations.into_iter().map(Into::into).collect(); - self.node_meta.entry(node).or_default().command_destinations = dests; - self - } - - /// Sets a human-readable kind for `node` (e.g. `model`, `tool`, `subgraph`) - /// surfaced as [`crate::NodeInfo::kind`] in the export. - pub fn with_node_kind(mut self, node: impl Into, kind: impl Into) -> Self { - self.node_meta.entry(node.into()).or_default().kind = Some(kind.into()); - self - } - - /// Records a sorted, free-form `key=value` annotation on `node`, carried - /// verbatim into [`crate::NodeInfo::metadata`]. - pub fn with_node_metadata( - mut self, - node: impl Into, - key: impl Into, - value: impl Into, - ) -> Self { - self.node_meta - .entry(node.into()) - .or_default() - .metadata - .insert(key.into(), value.into()); - self - } - - /// Marks `node` as a subgraph-embedding node for the export (an introspection - /// marker only; it does not change how the node executes). - pub fn mark_subgraph(mut self, node: impl Into) -> Self { - self.node_meta.entry(node.into()).or_default().subgraph = true; - self - } - - /// 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 - } - - /// 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 { - let node = node.into(); - self.node_meta.entry(node.clone()).or_default().deferred = true; - self.node_policies.entry(node).or_default().defer = true; - self - } - - /// Validates topology and freezes the graph into a [`CompiledGraph`]. - pub fn compile(self) -> Result> { - if self.reducer.is_none() { - return Err(TinyAgentsError::Validation( - "no state reducer set; call set_reducer (or GraphBuilder::overwrite)".to_string(), - )); - } - - // 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( - "START cannot route directly to END".to_string(), - )); - } - self.require_node(&entry)?; - - // static edges - for (from, targets) in &self.edges { - if from.as_str() != START { - self.require_node(from)?; - } - 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)?; - if self.edges.contains_key(from) { - return Err(TinyAgentsError::Validation(format!( - "node `{from}` has both a static edge and conditional edges" - ))); - } - for target in branch.routes.values() { - if target.as_str() != END { - self.require_node(target)?; - } - } - } - - // barrier/waiting edges: every source and target must exist - for (to, froms) in &self.waiting { - self.require_node(to)?; - for from in froms { - self.require_node(from)?; - } - } - - // 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)?; - if self.edges.contains_key(node) || self.branches.contains_key(node) { - return Err(TinyAgentsError::Validation(format!( - "node `{node}` declares command routing but also has static/conditional edges" - ))); - } - } - - let Self { - graph_id, - name, - nodes, - edges, - branches, - route_label_checks: _, - command_nodes, - waiting, - reducer, - recursion_limit, - parallel, - max_concurrency, - node_timeout, - node_meta, - barrier_reliefs, - node_policies, - node_defaults, - interrupt_before, - interrupt_after, - update_codec, - } = self; - - Ok(CompiledGraph::from_parts( - graph_id, - name, - nodes, - edges, - branches, - command_nodes, - waiting, - entry, - reducer.expect("reducer presence checked above"), - recursion_limit, - parallel, - max_concurrency, - 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<()> { - if self.nodes.contains_key(id) { - Ok(()) - } else { - Err(TinyAgentsError::MissingNode(id.to_string())) - } - } -} - -impl GraphBuilder -where - State: Clone + Send + Sync + 'static, -{ - /// Creates a builder that uses whole-state overwrite updates — the - /// milestone-1 default where each node returns the full next state. - pub fn overwrite() -> Self { - Self::new().set_reducer(OverwriteStateReducer) - } -} - -#[cfg(test)] -#[path = "mod_tests.rs"] -mod test; diff --git a/crates/tinyagents-graph/src/builder/mod_tests.rs b/crates/tinyagents-graph/src/builder/mod_tests.rs deleted file mode 100644 index c0821444f..000000000 --- a/crates/tinyagents-graph/src/builder/mod_tests.rs +++ /dev/null @@ -1,173 +0,0 @@ -//! Unit tests for the graph builder and its compile contract: reducer -//! requirement, START/END validation, missing-node/route detection, and the -//! rejection of nodes that mix command routing with static/conditional edges. - -use super::*; -use crate::TinyAgentsError; -use crate::command::NodeResult; - -type S = i32; - -#[test] -fn compile_requires_reducer() { - let err = GraphBuilder::::new() - .add_node("a", |s: S, _c: NodeContext| async move { - Ok(NodeResult::Update(s)) - }) - .set_entry("a") - .compile() - .unwrap_err(); - assert!(matches!(err, TinyAgentsError::Validation(_))); -} - -#[test] -fn compile_requires_start() { - let err = GraphBuilder::::overwrite() - .add_node("a", |s: S, _c: NodeContext| async move { - Ok(NodeResult::Update(s)) - }) - .compile() - .unwrap_err(); - assert!(matches!(err, TinyAgentsError::MissingStart)); -} - -#[test] -fn compile_rejects_missing_edge_target() { - let err = GraphBuilder::::overwrite() - .add_node("a", |s: S, _c: NodeContext| async move { - Ok(NodeResult::Update(s)) - }) - .set_entry("a") - .add_edge("a", "missing") - .compile() - .unwrap_err(); - assert!(matches!(err, TinyAgentsError::MissingNode(n) if n == "missing")); -} - -#[test] -fn compile_rejects_command_routing_with_edges() { - 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_edge("a", "b") - .set_finish("b") - .mark_command_routing("a") - .compile() - .unwrap_err(); - assert!(matches!(err, TinyAgentsError::Validation(_))); -} - -#[test] -fn compile_rejects_static_and_conditional_on_same_node() { - 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_edge("a", "b") - .add_conditional_edges("a", |_s: &S| "x".to_string(), [("x", "b")]) - .set_finish("b") - .compile() - .unwrap_err(); - 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() - .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_edge("a", "b") - .set_finish("b") - .compile(); - assert!(compiled.is_ok()); -} diff --git a/crates/tinyagents-graph/src/builder/policy.rs b/crates/tinyagents-graph/src/builder/policy.rs deleted file mode 100644 index 724cb5194..000000000 --- a/crates/tinyagents-graph/src/builder/policy.rs +++ /dev/null @@ -1,219 +0,0 @@ -//! 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/types.rs b/crates/tinyagents-graph/src/builder/types.rs deleted file mode 100644 index 71eb604f9..000000000 --- a/crates/tinyagents-graph/src/builder/types.rs +++ /dev/null @@ -1,561 +0,0 @@ -//! Builder types for the durable graph. -//! -//! Everything here is data accumulated by a [`GraphBuilder`]; the behavior -//! that acts on it (adding nodes/edges, validation, `compile`) lives in -//! `builder/mod.rs`. [`GraphBuilder::compile`](super::GraphBuilder::compile) -//! consumes a [`GraphBuilder`] and produces a [`crate::CompiledGraph`], whose -//! own fields largely mirror the ones declared here (see -//! `compiled::types::CompiledGraph`). - -use std::collections::{BTreeMap, HashMap, HashSet}; -use std::future::Future; -use std::pin::Pin; -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, TaskId, ThreadId}; - -/// The reserved virtual entry node. -pub const START: &str = "__start__"; -/// The reserved virtual terminal node. -pub const END: &str = "__end__"; - -/// Boxed future produced by a durable node handler. -pub type NodeFuture = Pin>> + Send>>; - -/// A durable node handler: receives a state snapshot and per-task context, -/// returns a [`NodeResult`]. -/// -/// Internally every handler receives the step's committed state as an -/// `Arc` (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(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. -/// -/// 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. -/// -/// When a graph compiled with [`crate::GraphBuilder::with_parallel`] -/// runs more than one active node in a single superstep, every branch executes -/// against its own cloned `State` snapshot and receives a distinct `ForkId` on -/// its [`NodeContext`]. The `branch_index` is the branch's position in the -/// deterministically-ordered active set, so a handler can tell which fork it is -/// (e.g. to seed per-fork randomness or pick a strategy) and the executor can -/// keep reducer application reproducible regardless of completion order. -/// -/// In sequential mode (the default), and in a parallel step that happens to -/// have a single active node, `NodeContext::fork` is `None`. -#[derive(Clone, Debug, PartialEq, Eq)] -pub struct ForkId { - /// The branch's index in the superstep's active set (0-based, stable). - pub branch_index: usize, - /// The node executing on this branch. - pub node: NodeId, -} - -impl ForkId { - /// Creates a fork id for `node` at `branch_index` within the active set. - pub fn new(branch_index: usize, node: NodeId) -> Self { - Self { branch_index, node } - } -} - -/// 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 -/// optional `resume` value. On a normal run `resume` is `None`; when a run is -/// resumed after an interrupt, the interrupted node is re-run with `resume` set -/// to the value carried by the resume command. -#[derive(Clone)] -pub struct NodeContext { - /// The graph currently executing this node. - pub graph_id: GraphId, - /// The node being executed. - pub node_id: NodeId, - /// The current run id. - pub run_id: RunId, - /// The thread id when checkpointing is enabled. - pub thread_id: Option, - /// The 1-based superstep number. - pub step: usize, - /// Resume value supplied by `CompiledGraph::resume`, if any. - pub resume: Option, - /// The branch identity when this node runs as one fork of a concurrent - /// (fan-out) superstep; `None` in sequential mode or single-node steps. - pub fork: Option, - /// The per-invocation argument when this activation was scheduled by a - /// [`crate::Send`] packet or seeded through - /// [`crate::GraphInput`]; `None` for normal edge/`goto` activations. - /// This is how map-reduce / search-fanout branches and external graph - /// inputs receive custom data that differs from the graph's shared - /// committed state. - /// - /// `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. - pub root_run_id: Option, - /// The enclosing run's live recursion stack (root-first). A subgraph node - /// seeds an embedded child graph with these frames so the child extends the - /// parent's recursion tree instead of starting a fresh one. - pub recursion_frames: Vec, - /// A per-run collector the executor provides so a subgraph node can report - /// the [`ChildRun`](crate::ChildRun) it spawned back to the enclosing - /// run; `None` when no executor sink is attached (e.g. a hand-built context). - pub child_runs: Option, - /// 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 { - fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - formatter - .debug_struct("NodeContext") - .field("graph_id", &self.graph_id) - .field("node_id", &self.node_id) - .field("run_id", &self.run_id) - .field("thread_id", &self.thread_id) - .field("step", &self.step) - .field("resume", &self.resume) - .field("fork", &self.fork) - .field("send_arg", &self.send_arg) - .field("root_run_id", &self.root_run_id) - .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() - } -} - -/// Behavior-free, introspectable metadata attached to a node by the builder. -/// -/// Markers and free-form metadata recorded here never affect execution; they -/// exist so the [topology export](crate::export) can describe what a node -/// *is* (a subgraph embedding, an interrupt point, a deferred join, …) without -/// inspecting the node's opaque handler closure. All fields are optional and -/// additive: an unset [`NodeMeta`] (the [`Default`]) contributes nothing. -#[derive(Clone, Debug, Default)] -pub(crate) struct NodeMeta { - /// A human-readable node kind (e.g. `model`, `tool`, `subgraph`). - pub(crate) kind: Option, - /// The node pauses the run before/at execution (an interrupt point). - pub(crate) interrupt: bool, - /// The node is a deferred join — it activates only after the rest of the - /// active frontier has drained (a barrier-style synthesis node). - pub(crate) deferred: bool, - /// The node embeds and runs a child graph (a subgraph node). - pub(crate) subgraph: bool, - /// Declared `goto` destination hints for a command-routing node, in the - /// order they were registered. Purely advisory: the runtime resolves the - /// real target from the emitted [`crate::Command`] at runtime. - pub(crate) command_destinations: Vec, - /// Arbitrary, sorted key/value annotations carried into the export. - pub(crate) metadata: BTreeMap, -} - -/// 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 { - pub(crate) handler: Arc>, -} - -impl Clone for BuilderNode { - fn clone(&self) -> Self { - Self { - handler: self.handler.clone(), - } - } -} - -/// A small newtype wrapper for a conditional-route label. -/// -/// Routers may return any `impl ToString` label (a plain `&str`/`String`, or a -/// user-defined route enum that implements `Display`). `Route` is an optional -/// ergonomic helper for building route tables and for routers that prefer to -/// return a typed value instead of a bare string; it stringifies via -/// [`ToString`] at the route-table boundary. -#[derive(Clone, Debug, PartialEq, Eq, Hash)] -pub struct Route(pub String); - -impl Route { - /// Wraps any `impl ToString` (e.g. a route enum with `Display`) as a label. - pub fn new(label: impl ToString) -> Self { - Self(label.to_string()) - } - - /// The label as a string slice. - pub fn as_str(&self) -> &str { - &self.0 - } -} - -impl std::fmt::Display for Route { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.write_str(&self.0) - } -} - -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`]. -/// -/// Every field is optional; only the `Some` fields override the builder's -/// current configuration, so partial defaults compose with explicit -/// `with_*` calls. All fields are opt-in and additive — an unset -/// [`GraphDefaults`] (the [`Default`]) changes nothing. -#[derive(Clone, Debug, Default)] -pub struct GraphDefaults { - /// Maximum number of supersteps before [`crate::TinyAgentsError::RecursionLimit`]. - pub recursion_limit: Option, - /// Whether the active node set of a superstep runs concurrently. - pub parallel: Option, - /// Upper bound on the number of branches run concurrently within one step - /// (only meaningful when `parallel` is enabled). `None` means unbounded. - pub max_concurrency: Option, - /// Default wall-clock timeout applied to every node handler; on elapse the - /// run fails with [`crate::TinyAgentsError::Timeout`]. `None` means no - /// per-node timeout. - pub node_timeout: Option, -} - -/// Conditional routing for a node: a router function plus its route table. -pub(crate) struct Branch { - pub(crate) router: Arc>, - pub(crate) routes: HashMap, -} - -impl Clone for Branch { - fn clone(&self) -> Self { - Self { - router: self.router.clone(), - routes: self.routes.clone(), - } - } -} - -/// A mutable, ergonomic builder for a durable state graph. -/// -/// `State` is the committed graph state; `Update` is the partial-update type -/// merged through the configured [`StateReducer`]. For whole-state graphs use -/// `Update == State` together with the overwrite reducer (see -/// [`super::GraphBuilder::overwrite`]). -pub struct GraphBuilder { - pub(crate) graph_id: GraphId, - /// Optional human-readable graph name carried into the topology export - /// (the `graph_id` remains the stable identifier). - pub(crate) name: Option, - pub(crate) nodes: 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. - pub(crate) waiting: HashMap>, - /// Mixed fan-in barrier relief registrations; see - /// [`super::BarrierRelief`]. - pub(crate) barrier_reliefs: Vec, - pub(crate) reducer: Option>>, - pub(crate) recursion_limit: usize, - /// When true, active nodes in a superstep run concurrently (fan-out). - pub(crate) parallel: bool, - /// Upper bound on concurrently-running branches per step (`None` = unbounded). - pub(crate) max_concurrency: Option, - /// Default per-node handler timeout (`None` = no timeout). - 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/drivers.rs b/crates/tinyagents-graph/src/cache/drivers.rs deleted file mode 100644 index c6835a574..000000000 --- a/crates/tinyagents-graph/src/cache/drivers.rs +++ /dev/null @@ -1,169 +0,0 @@ -//! A [`TaskCache`] over `tinystoragedrivers` document ports. -//! -//! [`DriverTaskCache`] keeps cached node results on whatever backend the host -//! opened. The document handle is already bound to a tenant scope, so two -//! tenants never see each other's entries even when graph, node and input hash -//! coincide. -//! -//! # Layout -//! -//! Entries are documents of the shared `tsd.kv` collection, owned by `graph`: -//! `{ owner, ns: "task_cache:", key, value }` with `key` = -//! `Key::of([node, hash])`, under `Key::of([owner, ns, key])`, so any hash -//! length is accepted and no two entries meet. A graph's entries share a -//! namespace, which is what [`clear`](TaskCache::clear) drops. A cache built -//! with [`with_namespace`](DriverTaskCache::with_namespace) uses -//! `task_cache::` so independent caches can share a backend. -//! -//! # Expiry -//! -//! The deadline is the record's `expires_ms` metadata, an absolute Unix epoch -//! millisecond instant that outlives the process. The driver enforces it: an -//! expired entry reads as absent, which this cache reports as a miss. - -use std::sync::Arc; -use std::time::Duration; - -use async_trait::async_trait; -use serde::{Deserialize, Serialize}; -use serde_json::Value; -use tinyagents_storage_kit::{ - CollectionSpec, DocumentStore, Filter, IndexSpec, Key, Repo, StorageError, WriteOptions, - primitives, -}; -use tokio::sync::OnceCell; - -use super::{TaskCache, TaskCacheKey}; -use crate::storage::{DEFAULT_CACHE_NAMESPACE, KV_OWNER, task_cache_ns}; -use crate::{Result, TinyAgentsError}; -use tinyagents_harness::ids::{GraphId, now_ms}; - -/// Map a storage driver failure onto the graph error type. -fn map_error(error: StorageError) -> TinyAgentsError { - TinyAgentsError::Graph(format!("task cache storage: {error}")) -} - -/// A `tsd.kv` document body. -#[derive(Debug, Serialize, Deserialize)] -struct Entry { - owner: String, - ns: String, - key: String, - value: Value, -} - -/// A [`TaskCache`] that stores entries in a driver [`DocumentStore`]. -#[derive(Debug, Clone)] -pub struct DriverTaskCache { - docs: Arc, - namespace: String, - declared: Arc>, -} - -impl DriverTaskCache { - /// Cache in the default namespace. - pub fn new(docs: Arc) -> Self { - Self::with_namespace(docs, DEFAULT_CACHE_NAMESPACE) - } - - /// Cache in `namespace`, so several independent caches can share a backend. - pub fn with_namespace(docs: Arc, namespace: impl Into) -> Self { - Self { - docs, - namespace: namespace.into(), - declared: Arc::new(OnceCell::new()), - } - } - - /// The declaration every cache needs: the shared `tsd.kv` collection. - pub fn collection_spec() -> CollectionSpec { - CollectionSpec::new(primitives::KV) - .index(IndexSpec::new("kv_by_key", ["owner", "ns", "key"])) - } - - /// The `tsd.kv` namespace of `graph`'s entries. - fn kv_ns(&self, graph: &str) -> String { - task_cache_ns(&self.namespace, graph) - } - - /// `(entry id, kv key)` of `key`. - fn ids(&self, key: &TaskCacheKey) -> Result<(String, String)> { - let kv_key = Key::of([key.node_id.as_str(), key.hash.as_str()]) - .map_err(map_error)? - .as_str() - .to_owned(); - let id = Key::of([KV_OWNER, &self.kv_ns(key.graph_id.as_str()), &kv_key]) - .map_err(map_error)? - .as_str() - .to_owned(); - Ok((id, kv_key)) - } - - fn repo(&self, options: WriteOptions) -> Result> { - Ok(Repo::new(&self.docs, primitives::KV) - .map_err(map_error)? - .with_options(options)) - } - - async fn declared(&self) -> Result<()> { - self.declared - .get_or_try_init(|| async { - self.docs - .ensure_collection(&Self::collection_spec()) - .await - .map_err(map_error) - }) - .await - .map(|_| ()) - } -} - -#[async_trait] -impl TaskCache for DriverTaskCache { - async fn get(&self, key: &TaskCacheKey) -> Result> { - self.declared().await?; - let (id, _) = self.ids(key)?; - Ok(self - .repo(WriteOptions::new())? - .get(&id) - .await - .map_err(map_error)? - .map(|found| found.doc.value)) - } - - async fn put(&self, key: &TaskCacheKey, value: Value, ttl: Option) -> Result<()> { - self.declared().await?; - let (id, kv_key) = self.ids(key)?; - let mut options = WriteOptions::new(); - if let Some(ttl) = ttl { - let at = now_ms().saturating_add(ttl.as_millis() as u64); - options = options.expires_ms(i64::try_from(at).unwrap_or(i64::MAX)); - } - let entry = Entry { - owner: KV_OWNER.to_owned(), - ns: self.kv_ns(key.graph_id.as_str()), - key: kv_key, - value, - }; - self.repo(options)? - .put(&id, &entry) - .await - .map(|_| ()) - .map_err(map_error) - } - - async fn clear(&self, graph_id: &GraphId) -> Result<()> { - self.declared().await?; - let filter = Filter::eq("owner", KV_OWNER) - .and(Filter::eq("ns", self.kv_ns(graph_id.as_str()).as_str())); - self.docs - .delete_where(primitives::KV, &filter) - .await - .map(|_| ()) - .map_err(map_error) - } -} - -#[cfg(test)] -#[path = "drivers_tests.rs"] -mod tests; diff --git a/crates/tinyagents-graph/src/cache/drivers_tests.rs b/crates/tinyagents-graph/src/cache/drivers_tests.rs deleted file mode 100644 index d889ba3a9..000000000 --- a/crates/tinyagents-graph/src/cache/drivers_tests.rs +++ /dev/null @@ -1,95 +0,0 @@ -use super::*; -use crate::testkit::conformance::task_cache_contract; -use serde_json::json; -use tinyagents_harness::ids::NodeId; -use tinyagents_storage_kit::{Scope, StorageBackend}; -use tinystoragedrivers_testkit::MemoryBackend; - -fn docs(storage: &dyn StorageBackend, scope: &str) -> Arc { - Arc::clone( - storage - .for_scope(&Scope::new(scope).unwrap()) - .unwrap() - .documents(), - ) -} - -fn key(graph: &str, node: &str, hash: &str) -> TaskCacheKey { - TaskCacheKey::new(GraphId::new(graph), NodeId::from(node), hash) -} - -#[tokio::test] -async fn memory_backend_passes_the_task_cache_contract() { - let storage = MemoryBackend::new(); - task_cache_contract(DriverTaskCache::new(docs(&storage, "local"))).await; -} - -#[tokio::test] -async fn testkit_backend_passes_the_task_cache_contract() { - let storage = MemoryBackend::new(); - task_cache_contract(DriverTaskCache::new(docs(&storage, "local"))).await; -} - -#[tokio::test] -async fn entries_persist_across_driver_handles_on_one_backend() { - let storage = MemoryBackend::new(); - { - let cache = DriverTaskCache::new(docs(&storage, "local")); - cache - .put(&key("g", "n", "h"), json!(7), None) - .await - .unwrap(); - } - let reopened = DriverTaskCache::new(docs(&storage, "local")); - assert_eq!( - reopened.get(&key("g", "n", "h")).await.unwrap(), - Some(json!(7)) - ); -} - -#[tokio::test] -async fn scopes_and_namespaces_keep_entries_apart() { - let storage = MemoryBackend::new(); - let alice = DriverTaskCache::new(docs(&storage, "alice")); - let bob = DriverTaskCache::new(docs(&storage, "bob")); - let other = DriverTaskCache::with_namespace(docs(&storage, "alice"), "other_cache"); - alice - .put(&key("g", "n", "h"), json!(1), None) - .await - .unwrap(); - assert_eq!(bob.get(&key("g", "n", "h")).await.unwrap(), None); - assert_eq!(other.get(&key("g", "n", "h")).await.unwrap(), None); - // Clearing one tenant's graph never reaches another tenant's entry. - bob.put(&key("g", "n", "h"), json!(2), None).await.unwrap(); - alice.clear(&GraphId::new("g")).await.unwrap(); - assert_eq!(bob.get(&key("g", "n", "h")).await.unwrap(), Some(json!(2))); -} - -#[tokio::test] -async fn long_input_hashes_fit_the_document_id_limit() { - let storage = MemoryBackend::new(); - let cache = DriverTaskCache::new(docs(&storage, "local")); - let long = key("g", "n", &"x".repeat(10_000)); - cache.put(&long, json!("big"), None).await.unwrap(); - assert_eq!(cache.get(&long).await.unwrap(), Some(json!("big"))); -} - -#[tokio::test] -async fn a_backend_that_cannot_keep_expiry_refuses_a_ttl_instead_of_caching_forever() { - let storage = MemoryBackend::new(); - let wrapped: Arc = Arc::new(tinyagents_harness::testkit::WithoutNativeTtl( - docs(&storage, "local"), - )); - let cache = DriverTaskCache::new(wrapped); - let error = cache - .put(&key("g", "n", "h"), json!(1), Some(Duration::from_secs(60))) - .await - .unwrap_err(); - assert!(matches!(error, TinyAgentsError::Graph(_)), "{error:?}"); - assert_eq!(cache.get(&key("g", "n", "h")).await.unwrap(), None); - // Without a TTL the entry needs no metadata and is stored. - cache - .put(&key("g", "n", "h"), json!(1), None) - .await - .unwrap(); -} diff --git a/crates/tinyagents-graph/src/cache/memory.rs b/crates/tinyagents-graph/src/cache/memory.rs deleted file mode 100644 index 066a6ff92..000000000 --- a/crates/tinyagents-graph/src/cache/memory.rs +++ /dev/null @@ -1,66 +0,0 @@ -//! 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 deleted file mode 100644 index bf331b0e5..000000000 --- a/crates/tinyagents-graph/src/cache/mod.rs +++ /dev/null @@ -1,26 +0,0 @@ -//! 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 [`DriverTaskCache`], which keeps entries on any `tinystoragedrivers` -//! backend (SQLite, MongoDB, ...) the host opened. - -mod drivers; -mod memory; -mod types; - -pub use drivers::DriverTaskCache; -pub use memory::InMemoryTaskCache; -pub(crate) use types::CachedNode; -pub use types::{TaskCache, TaskCacheKey}; - -#[cfg(test)] -#[path = "mod_tests.rs"] -mod test; diff --git a/crates/tinyagents-graph/src/cache/mod_tests.rs b/crates/tinyagents-graph/src/cache/mod_tests.rs deleted file mode 100644 index 3e7ecff21..000000000 --- a/crates/tinyagents-graph/src/cache/mod_tests.rs +++ /dev/null @@ -1,59 +0,0 @@ -//! Unit tests for the task cache backends: in-memory round trip, TTL -//! expiry, and per-graph `clear`. The driver-backed cache has its own tests. - -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)) - ); -} - -#[tokio::test] -async fn in_memory_passes_the_task_cache_contract() { - crate::testkit::conformance::task_cache_contract(InMemoryTaskCache::new()).await; -} diff --git a/crates/tinyagents-graph/src/cache/types.rs b/crates/tinyagents-graph/src/cache/types.rs deleted file mode 100644 index 8c9199879..000000000 --- a/crates/tinyagents-graph/src/cache/types.rs +++ /dev/null @@ -1,91 +0,0 @@ -//! 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/README.md b/crates/tinyagents-graph/src/channel/README.md deleted file mode 100644 index 78a41e910..000000000 --- a/crates/tinyagents-graph/src/channel/README.md +++ /dev/null @@ -1,81 +0,0 @@ -# graph::channel - -Channel-per-field state model (additive) for the graph runtime — an opt-in -alternative to plain whole-state overwrite that gives individual state fields -their own merge semantics and concurrent-write conflict detection. - -`ChannelState` implements `graph::reducer::StateReducer`, so a channel graph -runs on the **unchanged executor**: the executor folds a superstep's branch -results one at a time (`state = reducer.apply(state, update)` per branch's -`ChannelUpdate`), and because `ChannelState` is itself a reducer, each `apply` -dispatches every write in the update to the owning channel's `Channel::merge`. -No executor changes were needed to add this model. - -## Channel kinds and merge rules - -- **Aggregate channels** — `Topic`, `BinaryAggregate`, `Delta`, `Messages`, - `Barrier`, `NamedBarrier`. `Channel::allows_concurrent` is `true`: when two - fan-out branches write the same channel in one superstep, both writes fold - in deterministic active-set index order. -- **Overwrite channels** — `LastValue`, `Ephemeral`, `Untracked`. - `allows_concurrent` is `false`: a second same-step write to one of these - raises `TinyAgentsError::InvalidConcurrentUpdate`, because there is no - deterministic winner to pick. - -## Concurrent-write conflict detection - -Because the executor applies a step's updates as a contiguous batch, "same -step" is tracked by stamping each `ChannelUpdate` with the node's `ctx.step` -via `ChannelUpdate::at_step`. When updates are stamped, the reducer resets its -per-step bookkeeping (and clears `Ephemeral` channels) whenever the step -number advances. - -Unstamped updates are each treated as their own step — last-value writes -always win, with no conflict detection and no ephemeral clearing — so -existing whole-state habits (a node that writes without stamping) keep -working and conflict detection is strictly **opt-in**. - -## Public surface - -- `Channel` (trait) — defines `merge` and `allows_concurrent` for one field's - storage/merge policy. -- `LastValue` — overwrite semantics; last write in a step wins, concurrent - writes conflict. -- `Ephemeral` — like `LastValue` but cleared at the start of every step - (useful for per-step scratch signals). -- `Untracked` — overwrite semantics, opts a field entirely out of conflict - detection even when stamped. -- `Topic` / `Messages` — append-only aggregate channels; concurrent - writes in one step all append. -- `BinaryAggregate` — combines concurrent writes with a user-supplied - binary operator. -- `Delta` — accumulates numeric/structural deltas across concurrent - writes. -- `Barrier` / `NamedBarrier` — join-coordination channels used to detect when - every expected fan-out branch has arrived. -- `ChannelSet` — the map of named channels making up a graph's channel-typed - state. -- `ChannelState` — wraps a `ChannelSet` and implements `StateReducer`; the - type a `CompiledGraph` is parameterized over when using the channel model. -- `ChannelUpdate` — the per-branch write payload; `ChannelUpdate::at_step(n)` - opts it into step-stamped conflict detection. - -## Files - -| File | Role | -| --- | --- | -| `types.rs` | Type definitions and the channel/model overview. | -| `mod.rs` | `Channel::merge` rules, `ChannelSet` map operations, the `ChannelState` ⇒ `StateReducer` bridge. | -| `mod_tests.rs` | Unit tests (each channel kind's merge rule, conflict detection, step-stamping behavior). | - -## Operational constraints - -- Conflict detection only fires for stamped updates (`ChannelUpdate::at_step`) - — a node handler that forgets to stamp silently loses the safety net and - falls back to last-write-wins. -- `Ephemeral` channels are cleared on step advance, **not** on read — reading - one after the step in which it was written but before the next step still - sees the value. -- Mixing channel-typed state with plain whole-state overwrite in the same - graph is possible (channels are just another reducer) but conflates two - conflict-detection models; prefer one model per graph for predictability. diff --git a/crates/tinyagents-graph/src/channel/mod.rs b/crates/tinyagents-graph/src/channel/mod.rs deleted file mode 100644 index fcd525fcd..000000000 --- a/crates/tinyagents-graph/src/channel/mod.rs +++ /dev/null @@ -1,914 +0,0 @@ -//! Channel-per-field state model (additive). -//! -//! See `types` for the type definitions and the high-level model. This file -//! supplies the concrete [`Channel`] merge rules, the [`ChannelSet`] map -//! operations, and the [`ChannelState`] ⇒ [`StateReducer`] bridge that lets a -//! channel graph run on the existing executor. -//! -//! ## How a channel graph runs on the unchanged executor -//! -//! The executor folds a superstep's branch results one at a time: -//! `state = reducer.apply(state, update)` for each branch's -//! [`ChannelUpdate`]. [`ChannelState`] is its own reducer, so each `apply` -//! dispatches every write in the update to the owning channel's -//! [`Channel::merge`]. -//! -//! ## Concurrent-write conflict detection -//! -//! When two fan-out branches write the *same* channel in *one* superstep, the -//! merge must decide whether that is legal: -//! -//! - **Aggregate channels** ([`Topic`], [`BinaryAggregate`], [`Delta`], -//! [`Messages`], [`Barrier`], [`NamedBarrier`]) set -//! [`Channel::allows_concurrent`] to `true`; both writes fold in -//! deterministic active-set index order. -//! - **Overwrite channels** ([`LastValue`], [`Ephemeral`], [`Untracked`]) -//! return `false`; a second same-step write to such a channel raises -//! [`TinyAgentsError::InvalidConcurrentUpdate`] because there is no -//! deterministic winner. -//! -//! Because the executor applies a step's updates as a contiguous batch, "same -//! step" is tracked by stamping each [`ChannelUpdate`] with the node's -//! `ctx.step` via [`ChannelUpdate::at_step`]. When updates are stamped, the -//! reducer resets its per-step bookkeeping (and clears [`Ephemeral`] channels) -//! whenever the step number advances. Unstamped updates are each treated as -//! their own step (last-value writes always win, no conflict detection and no -//! 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, 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; -use crate::{Result, TinyAgentsError}; - -// --- Channel merge rules --- - -impl Channel for LastValue { - fn kind(&self) -> &'static str { - "last_value" - } - - fn merge(&self, _current: Option, incoming: Value) -> Result { - Ok(incoming) - } - - fn clone_box(&self) -> Box { - Box::new(*self) - } -} - -impl Channel for Topic { - fn kind(&self) -> &'static str { - "topic" - } - - fn merge(&self, current: Option, incoming: Value) -> Result { - // Reuse the existing array in place instead of cloning it per merge. - let mut list = match current { - Some(Value::Array(items)) => items, - Some(other) => vec![other], - None => Vec::new(), - }; - match incoming { - Value::Array(items) => list.extend(items), - other => list.push(other), - } - Ok(Value::Array(list)) - } - - fn allows_concurrent(&self) -> bool { - true - } - - fn clone_box(&self) -> Box { - Box::new(*self) - } -} - -impl Channel for Delta { - fn kind(&self) -> &'static str { - "delta" - } - - fn merge(&self, current: Option, incoming: Value) -> Result { - let add_err = - || TinyAgentsError::Graph("Delta channel only accepts numeric writes".to_string()); - let incoming_num = incoming.as_f64().ok_or_else(add_err)?; - let Some(current) = current else { - return Ok(incoming); - }; - let current_num = current.as_f64().ok_or_else(add_err)?; - - // Stay in integer space when both operands are integers. - if current.is_i64() && incoming.is_i64() { - let sum = current.as_i64().unwrap() + incoming.as_i64().unwrap(); - return Ok(Value::from(sum)); - } - Ok(Value::from(current_num + incoming_num)) - } - - fn allows_concurrent(&self) -> bool { - true - } - - fn clone_box(&self) -> Box { - Box::new(*self) - } -} - -impl Channel for Messages { - fn kind(&self) -> &'static str { - "messages" - } - - fn merge(&self, current: Option, incoming: Value) -> Result { - // Reuse the existing array in place instead of cloning it per merge. - let mut list = match current { - Some(Value::Array(items)) => items, - Some(_) => { - return Err(TinyAgentsError::Graph( - "Messages channel value must be a JSON array".to_string(), - )); - } - None => Vec::new(), - }; - let incoming = match incoming { - Value::Array(items) => items, - other => vec![other], - }; - // Build an id -> index map over the existing list once (O(existing)) so - // each incoming message is an O(1) lookup instead of a linear scan. - // Previously this dedup was O(existing x incoming), which bit at a few - // thousand messages. - let mut index: HashMap = list - .iter() - .enumerate() - .filter_map(|(i, existing)| { - existing - .get("id") - .and_then(Value::as_str) - .map(|id| (id.to_string(), i)) - }) - .collect(); - for msg in incoming { - match msg.get("id").and_then(Value::as_str).map(str::to_string) { - // Keyed message: replace the same id in place, or append and - // remember its position for later incoming writes. - Some(id) => match index.get(&id) { - Some(&i) => list[i] = msg, - None => { - index.insert(id, list.len()); - list.push(msg); - } - }, - // Unkeyed message: always appended (unchanged behavior). - None => list.push(msg), - } - } - Ok(Value::Array(list)) - } - - fn allows_concurrent(&self) -> bool { - true - } - - fn clone_box(&self) -> Box { - Box::new(*self) - } -} - -impl Channel for Ephemeral { - fn kind(&self) -> &'static str { - "ephemeral" - } - - fn merge(&self, _current: Option, incoming: Value) -> Result { - Ok(incoming) - } - - fn is_ephemeral(&self) -> bool { - true - } - - fn clone_box(&self) -> Box { - Box::new(*self) - } -} - -impl Channel for Untracked { - fn kind(&self) -> &'static str { - "untracked" - } - - fn merge(&self, _current: Option, incoming: Value) -> Result { - Ok(incoming) - } - - fn is_tracked(&self) -> bool { - false - } - - fn clone_box(&self) -> Box { - Box::new(*self) - } -} - -impl Barrier { - /// Creates a count-based barrier that is ready after `expected` arrivals. - pub fn new(expected: usize) -> Self { - Self { expected } - } -} - -impl Channel for Barrier { - fn kind(&self) -> &'static str { - "barrier" - } - - fn merge(&self, current: Option, incoming: Value) -> Result { - // Reuse the accumulated array in place instead of cloning it per merge. - let mut list = match current { - Some(Value::Array(items)) => items, - Some(other) => vec![other], - None => Vec::new(), - }; - match incoming { - Value::Array(items) => list.extend(items), - other => list.push(other), - } - Ok(Value::Array(list)) - } - - fn config(&self) -> Value { - serde_json::json!({ "expected": self.expected }) - } - - fn allows_concurrent(&self) -> bool { - true - } - - fn is_ready(&self, current: Option<&Value>) -> bool { - current - .and_then(Value::as_array) - .map(|items| items.len() >= self.expected) - .unwrap_or(self.expected == 0) - } - - fn clone_box(&self) -> Box { - Box::new(*self) - } -} - -impl NamedBarrier { - /// Creates a name-based barrier that is ready once every name has arrived. - pub fn new(expected: impl IntoIterator>) -> Self { - Self { - expected: expected.into_iter().map(Into::into).collect(), - } - } -} - -impl Channel for NamedBarrier { - fn kind(&self) -> &'static str { - "named_barrier" - } - - fn merge(&self, current: Option, incoming: Value) -> Result { - // Reuse the accumulated object in place instead of cloning it per merge. - let mut map = match current { - Some(Value::Object(map)) => map, - Some(_) => { - return Err(TinyAgentsError::Graph( - "NamedBarrier channel value must be a JSON object".to_string(), - )); - } - None => serde_json::Map::new(), - }; - let Value::Object(incoming) = incoming else { - return Err(TinyAgentsError::Graph( - "NamedBarrier writes must be JSON objects of named arrivals".to_string(), - )); - }; - for (key, value) in incoming { - map.insert(key, value); - } - Ok(Value::Object(map)) - } - - fn config(&self) -> Value { - serde_json::json!({ "expected": self.expected }) - } - - fn allows_concurrent(&self) -> bool { - true - } - - fn is_ready(&self, current: Option<&Value>) -> bool { - let Some(Value::Object(map)) = current else { - return self.expected.is_empty(); - }; - self.expected.iter().all(|name| map.contains_key(name)) - } - - fn clone_box(&self) -> Box { - Box::new(self.clone()) - } -} - -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`]. 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 { - fn kind(&self) -> &'static str { - "binary_aggregate" - } - - fn merge(&self, current: Option, incoming: Value) -> Result { - match current { - Some(current) => (self.fold)(current, incoming), - None => Ok(incoming), - } - } - - fn config(&self) -> Value { - ReducerRegistry::config_for(self.reducer_name.as_deref()) - } - - fn allows_concurrent(&self) -> bool { - true - } - - fn clone_box(&self) -> Box { - Box::new(self.clone()) - } -} - -/// 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 { - /// Creates an empty channel set. - pub fn new() -> Self { - Self::default() - } - - /// Registers `channel` under `name`, returning the set for chaining. - pub fn with_channel( - mut self, - name: impl Into, - channel: impl Channel + 'static, - ) -> Self { - self.add_channel(name, channel); - self - } - - /// Registers `channel` under `name`. - pub fn add_channel(&mut self, name: impl Into, channel: impl Channel + 'static) { - 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) - } - - /// Whether `name` is a registered channel. - pub fn contains(&self, name: &str) -> bool { - self.channels.contains_key(name) - } - - /// Whether the channel `name` permits concurrent same-step writes. Errors - /// if `name` is not a registered channel. - pub fn allows_concurrent(&self, name: &str) -> Result { - self.channel(name).map(|c| c.allows_concurrent()) - } - - /// Whether the barrier (or other) channel `name` has received everything it - /// is waiting for. Non-barrier channels are always ready. Errors if `name` - /// is not registered. - pub fn is_ready(&self, name: &str) -> Result { - let channel = self.channel(name)?; - Ok(channel.is_ready(self.values.get(name))) - } - - /// Folds `value` into the channel `name` via its merge rule. Errors with - /// [`TinyAgentsError::Graph`] if `name` is not a registered channel. - /// - /// The unknown-channel check runs *before* any state is touched. If a - /// registered channel's [`Channel::merge`] rejects the write (e.g. a - /// [`Delta`] receiving a non-numeric value), the channel's prior value is - /// dropped — a rejected write leaves the channel unset. This matches the - /// executor's reducer contract, where a merge error discards the whole - /// [`ChannelState`] for that step regardless. - pub fn apply_update(&mut self, name: &str, value: Value) -> Result<()> { - // Field-level borrows (channels immutable, values mutable) so the - // current value can be *moved* into `merge` — accumulating channels then - // fold in place rather than cloning the whole accumulated value. - let channel = self - .channels - .get(name) - .map(AsRef::as_ref) - .ok_or_else(|| TinyAgentsError::Graph(format!("unknown channel `{name}`")))?; - let current = self.values.remove(name); - let merged = channel.merge(current, value)?; - self.values.insert(name.to_string(), merged); - 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 { - self.values - .iter() - .filter(|(name, _)| { - self.channels - .get(*name) - .map(|c| c.is_tracked()) - .unwrap_or(true) - }) - .map(|(name, value)| (name.clone(), value.clone())) - .collect() - } - - /// Clears the value of every [`Ephemeral`] channel. Called at the start of a - /// new step by [`ChannelState`]. - pub(crate) fn clear_ephemeral(&mut self) { - let ephemeral: Vec = self - .channels - .iter() - .filter(|(_, c)| c.is_ephemeral()) - .map(|(name, _)| name.clone()) - .collect(); - for name in ephemeral { - self.values.remove(&name); - } - } - - fn channel(&self, name: &str) -> Result<&dyn Channel> { - self.channels - .get(name) - .map(AsRef::as_ref) - .ok_or_else(|| TinyAgentsError::Graph(format!("unknown channel `{name}`"))) - } -} - -/// 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 { - /// Creates an empty update. - pub fn new() -> Self { - Self::default() - } - - /// 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(), 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 - } - - /// Stamps the update with the producing node's superstep (`ctx.step`), - /// enabling same-step concurrent-write conflict detection and ephemeral - /// clearing. Without a stamp each update is treated as its own step. - pub fn at_step(mut self, step: usize) -> Self { - self.step = Some(step); - self - } - - /// Whether the update carries no writes. - pub fn is_empty(&self) -> bool { - self.writes.is_empty() - } -} - -// --- ChannelState --- - -impl ChannelState { - /// Creates a state with no channels. - pub fn new() -> Self { - Self::default() - } - - /// Registers `channel` under `name`, returning the state for chaining. Use - /// this to declare a graph's channel schema before running. - pub fn with_channel( - mut self, - name: impl Into, - channel: impl Channel + 'static, - ) -> Self { - self.set.add_channel(name, channel); - 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 - } - - /// Returns the current value of channel `name`, if written. - pub fn get(&self, name: &str) -> Option<&Value> { - self.set.get(name) - } - - /// Returns the tracked channel values (see [`ChannelSet::snapshot`]). - pub fn snapshot(&self) -> BTreeMap { - self.set.snapshot() - } - - /// Whether channel `name` is a satisfied barrier (see - /// [`ChannelSet::is_ready`]). - pub fn is_ready(&self, name: &str) -> Result { - self.set.is_ready(name) - } - - /// Folds a [`ChannelUpdate`] into this state, dispatching each write to its - /// channel's merge rule. This is the core reducer step. - /// - /// When the update is stamped (via [`ChannelUpdate::at_step`]) with a step - /// number that differs from the last one seen, the per-step write tracking - /// is reset and [`Ephemeral`] channels are cleared before the writes apply. - /// A second write to a non-aggregate channel within the same stamped step - /// raises [`TinyAgentsError::InvalidConcurrentUpdate`]. - pub fn merge(mut self, update: ChannelUpdate) -> Result { - match update.step { - 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(); - } - } - - // Distinct channels touched by this single update (a node writing the - // same channel twice in one update is last-wins, not a conflict). - let mut distinct: Vec<&str> = Vec::new(); - for (name, _) in &update.writes { - if !distinct.contains(&name.as_str()) { - distinct.push(name.as_str()); - } - } - - // Validate before mutating so a conflicting step never commits partial - // writes. - for name in &distinct { - let allows = self.set.allows_concurrent(name)?; - let count = self.step_writes.get(*name).copied().unwrap_or(0) + 1; - if count > 1 && !allows { - return Err(TinyAgentsError::InvalidConcurrentUpdate(format!( - "channel `{name}` received {count} concurrent writes in one step but is not an aggregate channel" - ))); - } - } - - let touched: HashSet = distinct.iter().map(|n| n.to_string()).collect(); - 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; - } - // `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 -/// (merge rules live in the running `state`'s [`ChannelSet`]), so any -/// `ChannelState` may be passed to `set_reducer`. -impl StateReducer for ChannelState { - fn apply(&self, state: ChannelState, update: ChannelUpdate) -> Result { - state.merge(update) - } -} - -#[cfg(test)] -#[path = "mod_tests.rs"] -mod test; diff --git a/crates/tinyagents-graph/src/channel/mod_tests.rs b/crates/tinyagents-graph/src/channel/mod_tests.rs deleted file mode 100644 index ba65fcce4..000000000 --- a/crates/tinyagents-graph/src/channel/mod_tests.rs +++ /dev/null @@ -1,936 +0,0 @@ -//! Unit tests for the channel-per-field state model: each channel's merge rule, -//! `ChannelSet` dispatch/snapshot, `ChannelState` round-trips through a real -//! graph (fan-out into aggregate channels), and concurrent-write conflict -//! detection on non-aggregate channels. - -use super::*; -use crate::TinyAgentsError; -use crate::builder::{GraphBuilder, NodeContext}; -use crate::command::{Command, NodeResult}; -use serde_json::{Value, json}; -use tinystoragedrivers_testkit::MemoryBackend; - -// --- per-channel merge rules --- - -#[test] -fn last_value_overwrites() { - let c = LastValue; - assert_eq!(c.merge(None, json!(1)).unwrap(), json!(1)); - assert_eq!(c.merge(Some(json!(1)), json!(2)).unwrap(), json!(2)); - assert!(!c.allows_concurrent()); -} - -#[test] -fn topic_appends_scalars_and_arrays() { - let c = Topic; - let v = c.merge(None, json!("a")).unwrap(); - let v = c.merge(Some(v), json!("b")).unwrap(); - let v = c.merge(Some(v), json!(["c", "d"])).unwrap(); - assert_eq!(v, json!(["a", "b", "c", "d"])); - assert!(c.allows_concurrent()); -} - -#[test] -fn delta_accumulates_numbers() { - let c = Delta; - let v = c.merge(None, json!(2)).unwrap(); - let v = c.merge(Some(v), json!(3)).unwrap(); - assert_eq!(v, json!(5)); - // integer + float promotes to float - let v = c.merge(Some(v), json!(0.5)).unwrap(); - assert_eq!(v, json!(5.5)); - assert!(c.merge(Some(json!(1)), json!("x")).is_err()); -} - -#[test] -fn messages_merge_by_id() { - let c = Messages; - let v = c.merge(None, json!([{"id": "1", "text": "hi"}])).unwrap(); - let v = c - .merge(Some(v), json!([{"id": "2", "text": "yo"}])) - .unwrap(); - // upsert existing id 1 - let v = c - .merge(Some(v), json!({"id": "1", "text": "hello"})) - .unwrap(); - assert_eq!( - v, - json!([{"id": "1", "text": "hello"}, {"id": "2", "text": "yo"}]) - ); - assert!(c.allows_concurrent()); -} - -#[test] -fn messages_merge_dedup_map_preserves_order_and_appends_unkeyed() { - // Exercise the id->index map path: a single batch that upserts an existing - // id, appends a new id, and appends an unkeyed message, then a follow-up - // upsert of the id introduced by that batch. Existing order is preserved and - // dedup is by id only. - let c = Messages; - let v = c - .merge( - None, - json!([{"id": "a", "text": "1"}, {"id": "b", "text": "2"}]), - ) - .unwrap(); - let v = c - .merge( - Some(v), - json!([ - {"id": "a", "text": "1-updated"}, - {"id": "c", "text": "3"}, - {"text": "no-id"}, - ]), - ) - .unwrap(); - // Upserting an id first seen in the previous batch must land on that message. - let v = c - .merge(Some(v), json!({"id": "c", "text": "3-updated"})) - .unwrap(); - assert_eq!( - v, - json!([ - {"id": "a", "text": "1-updated"}, - {"id": "b", "text": "2"}, - {"id": "c", "text": "3-updated"}, - {"text": "no-id"}, - ]) - ); -} - -#[test] -fn ephemeral_overwrites_and_is_marked() { - let c = Ephemeral; - assert_eq!(c.merge(Some(json!(1)), json!(2)).unwrap(), json!(2)); - assert!(c.is_ephemeral()); - assert!(!c.allows_concurrent()); -} - -#[test] -fn untracked_is_not_tracked() { - let c = Untracked; - assert!(!c.is_tracked()); - assert_eq!(c.merge(None, json!("x")).unwrap(), json!("x")); -} - -#[test] -fn binary_aggregate_folds_with_closure() { - let c = BinaryAggregate::new(|a: Value, b: Value| { - Ok(json!(a.as_i64().unwrap() * b.as_i64().unwrap())) - }); - let v = c.merge(None, json!(2)).unwrap(); - let v = c.merge(Some(v), json!(3)).unwrap(); - let v = c.merge(Some(v), json!(4)).unwrap(); - assert_eq!(v, json!(24)); - assert!(c.allows_concurrent()); -} - -#[test] -fn binary_aggregate_from_reducer() { - // Any `Reducer` can back an aggregate channel; here a closure reducer - // that keeps the numeric maximum. - let c = BinaryAggregate::from_reducer(crate::ClosureReducer::new(|a: Value, b: Value| { - Ok(if b.as_f64() > a.as_f64() { b } else { a }) - })); - let v = c.merge(None, json!(3)).unwrap(); - let v = c.merge(Some(v), json!(7)).unwrap(); - let v = c.merge(Some(v), json!(1)).unwrap(); - assert_eq!(v, json!(7)); -} - -#[test] -fn count_barrier_readiness() { - let c = Barrier::new(2); - let v = c.merge(None, json!("a")).unwrap(); - assert!(!c.is_ready(Some(&v))); - let v = c.merge(Some(v), json!("b")).unwrap(); - assert!(c.is_ready(Some(&v))); -} - -#[test] -fn named_barrier_readiness() { - let c = NamedBarrier::new(["left", "right"]); - let v = c.merge(None, json!({"left": 1})).unwrap(); - assert!(!c.is_ready(Some(&v))); - let v = c.merge(Some(v), json!({"right": 2})).unwrap(); - assert!(c.is_ready(Some(&v))); - assert!(c.allows_concurrent()); -} - -// --- ChannelSet --- - -#[test] -fn channel_set_apply_and_snapshot_excludes_untracked() { - let mut set = ChannelSet::new() - .with_channel("count", Delta) - .with_channel("scratch", Untracked); - set.apply_update("count", json!(2)).unwrap(); - set.apply_update("count", json!(5)).unwrap(); - set.apply_update("scratch", json!("temp")).unwrap(); - - assert_eq!(set.get("count"), Some(&json!(7))); - let snap = set.snapshot(); - assert_eq!(snap.get("count"), Some(&json!(7))); - assert!( - !snap.contains_key("scratch"), - "untracked excluded from snapshot" - ); -} - -#[test] -fn channel_set_unknown_channel_errors() { - let mut set = ChannelSet::new(); - assert!(matches!( - set.apply_update("nope", json!(1)), - Err(TinyAgentsError::Graph(_)) - )); -} - -#[test] -fn apply_update_accumulates_barrier_by_moving_current() { - // Exercises the owned-`current` merge path through `apply_update`: repeated - // writes to a barrier fold into the same accumulated array without cloning, - // and readiness reflects the accumulated arrivals. - let mut set = ChannelSet::new().with_channel("join", Barrier::new(3)); - set.apply_update("join", json!("a")).unwrap(); - set.apply_update("join", json!("b")).unwrap(); - assert!(!set.is_ready("join").unwrap()); - set.apply_update("join", json!("c")).unwrap(); - assert!(set.is_ready("join").unwrap()); - assert_eq!(set.get("join"), Some(&json!(["a", "b", "c"]))); -} - -#[test] -fn apply_update_rejected_write_clears_channel_value() { - // A registered channel whose merge rejects the write leaves the channel - // unset (the current value was moved into the failed merge). The unknown- - // channel guard, by contrast, runs before any state is touched. - let mut set = ChannelSet::new().with_channel("n", Delta); - set.apply_update("n", json!(5)).unwrap(); - assert_eq!(set.get("n"), Some(&json!(5))); - - // Non-numeric write is rejected by Delta's merge; prior value is dropped. - assert!(set.apply_update("n", json!("not-a-number")).is_err()); - assert_eq!( - set.get("n"), - None, - "rejected write leaves the channel unset" - ); -} - -// --- ChannelState graph round-trips --- - -fn aggregate_state() -> ChannelState { - ChannelState::new() - .with_channel("results", Topic) - .with_channel("total", Delta) -} - -/// `start` fans out to `a` and `b` in one parallel superstep; both write the -/// shared aggregate channels and their writes merge deterministically. -#[tokio::test] -async fn fan_out_merges_both_branches_into_aggregate_channels() { - let graph = GraphBuilder::::new() - .set_reducer(ChannelState::new()) - .with_parallel(true) - .add_node("start", |_s: ChannelState, _c: NodeContext| async move { - Ok(NodeResult::Command(Command::goto(["a", "b"]))) - }) - .add_node("a", |_s: ChannelState, c: NodeContext| async move { - Ok(NodeResult::Update( - ChannelUpdate::new() - .set("results", "a") - .set("total", 10) - .at_step(c.step), - )) - }) - .add_node("b", |_s: ChannelState, c: NodeContext| async move { - Ok(NodeResult::Update( - ChannelUpdate::new() - .set("results", "b") - .set("total", 5) - .at_step(c.step), - )) - }) - .mark_command_routing("start") - .set_entry("start") - .set_finish("a") - .set_finish("b") - .compile() - .unwrap(); - - let exec = graph.run(aggregate_state()).await.unwrap(); - // Topic merges both branches in active-set index order; Delta sums them. - assert_eq!(exec.state.get("results"), Some(&json!(["a", "b"]))); - assert_eq!(exec.state.get("total"), Some(&json!(15))); -} - -/// Two concurrent branches writing the same `LastValue` channel in one step is -/// a conflict. -#[tokio::test] -async fn concurrent_last_value_writes_conflict() { - let graph = GraphBuilder::::new() - .set_reducer(ChannelState::new()) - .with_parallel(true) - .add_node("start", |_s: ChannelState, _c: NodeContext| async move { - Ok(NodeResult::Command(Command::goto(["a", "b"]))) - }) - .add_node("a", |_s: ChannelState, c: NodeContext| async move { - Ok(NodeResult::Update( - ChannelUpdate::new().set("winner", "a").at_step(c.step), - )) - }) - .add_node("b", |_s: ChannelState, c: NodeContext| async move { - Ok(NodeResult::Update( - ChannelUpdate::new().set("winner", "b").at_step(c.step), - )) - }) - .mark_command_routing("start") - .set_entry("start") - .set_finish("a") - .set_finish("b") - .compile() - .unwrap(); - - let initial = ChannelState::new().with_channel("winner", LastValue); - let err = graph.run(initial).await.unwrap_err(); - assert!(matches!(err, TinyAgentsError::InvalidConcurrentUpdate(_))); -} - -/// A `LastValue` channel overwritten across *different* steps is fine — only -/// same-step concurrent writes conflict. -#[tokio::test] -async fn sequential_last_value_overwrites_across_steps() { - let graph = GraphBuilder::::new() - .set_reducer(ChannelState::new()) - .add_node("a", |_s: ChannelState, c: NodeContext| async move { - Ok(NodeResult::Update( - ChannelUpdate::new().set("v", "first").at_step(c.step), - )) - }) - .add_node("b", |_s: ChannelState, c: NodeContext| async move { - Ok(NodeResult::Update( - ChannelUpdate::new().set("v", "second").at_step(c.step), - )) - }) - .add_edge("a", "b") - .set_entry("a") - .set_finish("b") - .compile() - .unwrap(); - - let initial = ChannelState::new().with_channel("v", LastValue); - let exec = graph.run(initial).await.unwrap(); - assert_eq!(exec.state.get("v"), Some(&json!("second"))); -} - -/// An `Ephemeral` channel written in one step is visible to the next step's -/// node, then cleared when the step advances. -#[tokio::test] -async fn ephemeral_channel_is_cleared_on_next_step() { - let graph = GraphBuilder::::new() - .set_reducer(ChannelState::new()) - .add_node("a", |_s: ChannelState, c: NodeContext| async move { - Ok(NodeResult::Update( - ChannelUpdate::new().set("tmp", 1).at_step(c.step), - )) - }) - .add_node("b", |s: ChannelState, c: NodeContext| async move { - // tmp is still visible to this (next-step) node. - let seen = s.get("tmp").cloned().unwrap_or(json!(null)); - Ok(NodeResult::Update( - ChannelUpdate::new().set("saw", seen).at_step(c.step), - )) - }) - .add_edge("a", "b") - .set_entry("a") - .set_finish("b") - .compile() - .unwrap(); - - let initial = ChannelState::new() - .with_channel("tmp", Ephemeral) - .with_channel("saw", LastValue); - let exec = graph.run(initial).await.unwrap(); - assert_eq!(exec.state.get("saw"), Some(&json!(1)), "tmp was visible"); - assert_eq!( - exec.state.get("tmp"), - None, - "tmp cleared after step advanced" - ); - // Snapshot reflects only the tracked, surviving channels. - let snap = exec.state.snapshot(); - assert_eq!(snap.get("saw"), Some(&json!(1))); - assert!(!snap.contains_key("tmp")); -} - -/// A single node writing the same non-aggregate channel twice in one update is -/// last-wins, not a conflict. -#[test] -fn single_update_repeat_write_is_last_wins() { - let state = ChannelState::new().with_channel("v", LastValue); - let merged = state - .merge(ChannelUpdate::new().set("v", 1).set("v", 2).at_step(1)) - .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_a_durable_checkpointer() { - let checkpoint_backend = MemoryBackend::new(); - use crate::checkpoint::Checkpointer; - 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 cp: Arc> = Arc::new( - crate::test_support::dir_checkpointer::(&checkpoint_backend), - ); - 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 - // reopen: 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( - crate::test_support::dir_checkpointer::(&checkpoint_backend), - ); - 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"]))); - } -} - -// --- 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/registry.rs b/crates/tinyagents-graph/src/channel/registry.rs deleted file mode 100644 index 49127c0f3..000000000 --- a/crates/tinyagents-graph/src/channel/registry.rs +++ /dev/null @@ -1,157 +0,0 @@ -//! 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 - /// `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/types.rs b/crates/tinyagents-graph/src/channel/types.rs deleted file mode 100644 index 3ef887fcb..000000000 --- a/crates/tinyagents-graph/src/channel/types.rs +++ /dev/null @@ -1,325 +0,0 @@ -//! Channel-per-field state types. -//! -//! This module defines the *channel* state model described in -//! `docs/modules/graph/state-channels.md`. It is an **additive** alternative to -//! the monolithic `State` + [`crate::StateReducer`] path: instead of one -//! whole-state value with one reducer, state is split into independently-named -//! *channels*, each owning its own current value and its own binary merge rule. -//! -//! The model is built from four pieces: -//! -//! - [`Channel`]: the per-key merge policy (overwrite, aggregate, append, …). -//! - [`ChannelSet`]: a named map of channels plus their current values. -//! - [`ChannelUpdate`]: a batch of `(channel_name, value)` writes a node returns. -//! - [`ChannelState`]: a concrete graph `State` wrapping a [`ChannelSet`] that -//! implements [`crate::StateReducer`], so a -//! channel graph is just `GraphBuilder` running on -//! the unchanged executor. -//! -//! Values are [`serde_json::Value`]-backed for generality so channels compose -//! with checkpointing, export, and language-defined nodes without per-graph type -//! parameters. - -use std::collections::{BTreeMap, HashMap}; - -use serde::{Deserialize, Serialize}; -use serde_json::Value; - -use crate::Result; - -/// A single named state channel: it owns the merge rule that folds an incoming -/// update value into the channel's current value at a superstep boundary. -/// -/// Channels are object-safe and value-typed ([`serde_json::Value`]) so a -/// [`ChannelSet`] can hold a heterogeneous map of `Box`. Each -/// channel decides: -/// -/// - [`Channel::merge`] — how an incoming write combines with the current value. -/// - [`Channel::allows_concurrent`] — whether two branches in the *same* step may -/// both write this channel (aggregates: yes; last-value: no, see -/// [`crate::TinyAgentsError::InvalidConcurrentUpdate`]). -/// - [`Channel::is_ephemeral`] — whether the value is cleared at the start of the -/// next step (one-shot channels). -/// - [`Channel::is_tracked`] — whether the value appears in [`ChannelSet::snapshot`] -/// and is considered part of durable state. -/// - [`Channel::is_ready`] — barrier readiness (defaults to always ready). -pub trait Channel: Send + Sync { - /// A short, stable kind tag used for debugging and topology export. - fn kind(&self) -> &'static str; - - /// Folds `incoming` into the channel's `current` value, returning the new - /// value. `current` is `None` the first time the channel is written. - /// - /// Takes `current` **by value** so accumulating channels (topics, barriers, - /// message logs) can reuse the existing backing `Vec`/`Map` in place instead - /// of cloning the entire accumulated value on every merge, which was - /// 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 - /// [`crate::TinyAgentsError::InvalidConcurrentUpdate`] on a same-step clash. - fn allows_concurrent(&self) -> bool { - false - } - - /// Whether the channel's value is cleared at the start of the next step. - fn is_ephemeral(&self) -> bool { - false - } - - /// Whether the channel participates in [`ChannelSet::snapshot`] / durable - /// state. [`Untracked`] returns `false`. - fn is_tracked(&self) -> bool { - true - } - - /// Barrier readiness: whether the channel has received all the inputs it is - /// waiting for. Non-barrier channels are always ready. - fn is_ready(&self, _current: Option<&Value>) -> bool { - true - } - - /// Clones the channel into a fresh box (enables `Clone` for [`ChannelSet`]). - fn clone_box(&self) -> Box; -} - -impl Clone for Box { - fn clone(&self) -> Self { - self.clone_box() - } -} - -impl std::fmt::Debug for Box { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("Channel") - .field("kind", &self.kind()) - .finish() - } -} - -/// Overwrite channel: each write replaces the value (last-value semantics). -/// -/// Rejects concurrent same-step writes from multiple branches with -/// [`crate::TinyAgentsError::InvalidConcurrentUpdate`]. -#[derive(Clone, Copy, Debug, Default)] -pub struct LastValue; - -/// Append channel: accumulates writes into a JSON array across steps. -/// -/// A scalar write is pushed as one element; an array write extends the list. -/// Allows concurrent same-step writes (the order follows deterministic -/// active-set index order). -#[derive(Clone, Copy, Debug, Default)] -pub struct Topic; - -/// Numeric accumulator: each write is added to the running total. -/// -/// Integer writes stay integers; any float write promotes the total to a float. -/// Allows concurrent same-step writes. -#[derive(Clone, Copy, Debug, Default)] -pub struct Delta; - -/// Message-merge channel: maintains a JSON array of message objects deduplicated -/// by their `id` field. An incoming message whose `id` matches an existing entry -/// replaces it in place; otherwise it is appended. Allows concurrent same-step -/// writes. -#[derive(Clone, Copy, Debug, Default)] -pub struct Messages; - -/// One-shot overwrite channel whose value is cleared at the start of the next -/// step (see [`ChannelUpdate::at_step`] for how step boundaries are detected). -#[derive(Clone, Copy, Debug, Default)] -pub struct Ephemeral; - -/// Overwrite channel excluded from [`ChannelSet::snapshot`] and durable-state -/// views. Useful for scratch values that should not be checkpointed. -#[derive(Clone, Copy, Debug, Default)] -pub struct Untracked; - -/// Count-based barrier: accumulates writes into a JSON array and is *ready* only -/// once it has collected at least `expected` arrivals. Allows concurrent -/// same-step writes (fan-in is the whole point). -#[derive(Clone, Copy, Debug)] -pub struct Barrier { - /// Number of arrivals required before [`Channel::is_ready`] returns `true`. - pub expected: usize, -} - -/// Name-based barrier: accumulates writes into a JSON object keyed by arrival -/// name and is *ready* only once every name in `expected` has arrived. Each -/// incoming write is a JSON object whose keys are merged into the accumulator. -/// Allows concurrent same-step writes. -#[derive(Clone, Debug)] -pub struct NamedBarrier { - /// The set of names that must all arrive before the barrier is ready. - pub expected: Vec, -} - -/// Binary-aggregate channel: folds writes through a user-supplied binary -/// closure (append, add, min, max, custom). The first write becomes the value -/// directly; subsequent writes are `fold(current, incoming)`. Allows concurrent -/// same-step writes. -#[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 { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("BinaryAggregate").finish_non_exhaustive() - } -} - -/// A named map of channels plus their current [`serde_json::Value`]s. -/// -/// The channel definitions (merge rules) live alongside the values, so the set -/// is self-describing: merging an update only needs the set itself. Construct -/// one with [`ChannelSet::new`] and register channels with -/// [`ChannelSet::with_channel`]; feed writes through -/// [`ChannelSet::apply_update`]; read the durable view with -/// [`ChannelSet::snapshot`]. -#[derive(Default)] -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 { - fn clone(&self) -> Self { - Self { - channels: self - .channels - .iter() - .map(|(k, v)| (k.clone(), v.clone_box())) - .collect(), - values: self.values.clone(), - delta_channels: self.delta_channels.clone(), - } - } -} - -impl std::fmt::Debug for ChannelSet { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - let kinds: BTreeMap<&str, &'static str> = self - .channels - .iter() - .map(|(k, v)| (k.as_str(), v.kind())) - .collect(); - f.debug_struct("ChannelSet") - .field("channels", &kinds) - .field("values", &self.values) - .finish() - } -} - -/// 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. -/// -/// `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, ChannelWrite)>, - pub(crate) step: Option, -} - -/// A concrete graph `State` wrapping a [`ChannelSet`]. -/// -/// `ChannelState` implements [`crate::StateReducer`] for itself, so a channel-based graph is built directly with -/// `GraphBuilder` and runs on the unchanged -/// executor: each superstep the executor folds every branch's -/// [`ChannelUpdate`] into the committed `ChannelState` through this reducer, -/// dispatching each write to its channel's merge rule. -/// -/// 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, 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/README.md b/crates/tinyagents-graph/src/checkpoint/README.md deleted file mode 100644 index 667c3569c..000000000 --- a/crates/tinyagents-graph/src/checkpoint/README.md +++ /dev/null @@ -1,112 +0,0 @@ -# graph::checkpoint - -The `Checkpointer` trait and its backends — the durability layer that makes -the recursive graph runtime resumable and time-travelable. - -In a recursive-language-model harness, runs nest: a graph node can run another -compiled graph, which can run another, each producing its own state. -Checkpointing snapshots every level of that tree at superstep boundaries and -keys them by `thread_id` / `namespace` so a parent and its embedded subgraphs -never collide (see `graph::subgraph`). Persisting committed state at each -boundary is what lets a run be paused on an interrupt, resumed later, forked, -or replayed for time-travel debugging. - -Checkpoints are written **at superstep boundaries only — never mid-node** — -so resuming always reruns a node from its start. There is no partial-node -durability; a node handler must be safe to re-run from scratch if the process -crashes mid-execution. - -## Public surface - -### The trait - -`Checkpointer` (async, `Send + Sync`): - -- `put(checkpoint) -> Result` — persists a checkpoint, returns - its id. -- `get(thread_id, checkpoint_id: Option<&str>) -> Result>>` - — loads a checkpoint; `None` id loads the latest for the thread. -- A namespace-scoped variant of `get` restricts the lookup to checkpoints - whose stored namespace matches, which is what keeps a parent run and the - subgraphs it embeds — sharing a thread id but differing in namespace — from - loading each other's checkpoints on resume or inspection. -- Additional methods for listing checkpoint history and pruning lineage (see - `mod.rs` for the full trait). - -### Backends - -- `DriverCheckpointer` (`drivers.rs`) — durable implementation over a - `tinystoragedrivers` `DocumentStore`, so the host's backend (SQLite, MongoDB, - memory) holds the checkpoints. The file and SQLite checkpointers were removed - in favour of it. -- An in-memory backend is also available for tests (see `mod.rs`/`mod_tests.rs`); - suitable only for a single process's lifetime. - -### Types (`types.rs`) - -- `Checkpoint` — the persisted record: committed `state`, `next_nodes` - (the pending activation set), `interrupts`, source, and lineage pointers. - `Checkpoint::to_metadata()` projects the listing-relevant fields into - `CheckpointMetadata` so a `StateSnapshot` (from `graph::compiled`) and - `Checkpointer::list` always agree on what a checkpoint "is." -- `CheckpointTuple` — a checkpoint plus its `config` and - `parent_config`, the shape returned by `get`/history lookups. -- `CheckpointConfig` — addresses a checkpoint by `thread_id` (+ optional - `checkpoint_id`); `CheckpointConfig::latest(thread_id)` is the common case. -- `CheckpointMetadata` — the compact, listing-facing view of a checkpoint. -- `CheckpointSource` — why a checkpoint was written (superstep boundary, - interrupt, resumable-failure boundary, ...); round-trips through - `as_str()` / `parse()`. -- `DurabilityMode` — how aggressively the executor checkpoints (e.g. every - boundary vs. only on interrupt/failure); set via - `CompiledGraph::with_durability`. -- `PendingActivation` — a scheduled-but-not-yet-run node activation persisted - across a checkpoint boundary. -- `BarrierArrivals` — tracks which parallel branches have arrived at a join - barrier, persisted so a resumed run doesn't re-run already-arrived branches. -- `PendingWrite` — a reducer write buffered but not yet folded into committed - state at the point the checkpoint was taken. - -## Files - -| File | Role | -| --- | --- | -| `types.rs` | `Checkpoint`, `CheckpointTuple`, `CheckpointConfig`, `CheckpointMetadata`, `CheckpointSource`, `DurabilityMode`, `PendingActivation`, `BarrierArrivals`, `PendingWrite`. | -| `mod.rs` | The `Checkpointer` trait and the in-memory backend. | -| `drivers.rs` | `DriverCheckpointer` over a driver `DocumentStore`. | -| `mod_tests.rs` | Unit tests (put/get round-trips, namespace scoping, history, pruning). | - -## Operational constraints - -- Namespace scoping is load-bearing for subgraph isolation: a checkpointer - implementation that ignores namespace when a caller asks for it will let a - parent and a subgraph run collide on resume. Any new backend must honor the - namespace-scoped `get`/`list` contract. -- `DriverCheckpointer` on a durable backend is safe across process - restarts; the in-memory backend is not — never use it where resumability - after a crash is required. -- Checkpoint ids must be collision-free across process restarts (see - `graph::compiled`'s `next_checkpoint_id`, which delegates to - `harness::ids::new_checkpoint_id` rather than a process-local counter) or - lineage pruning and time-travel resume can corrupt. -- `DurabilityMode` trades persistence frequency for write volume — a mode that - only checkpoints on interrupt/failure means a mid-run crash loses all - progress since the last such boundary, not just the current node. - -## `DriverCheckpointer` (`storage-drivers` feature) - -`DriverCheckpointer` stores checkpoints, pending writes and execution leases in -a [tinystoragedrivers](https://github.com/tinyhumansai/tinystoragedrivers) -`DocumentStore`. Hosts use it to put graph durability on whichever backend they -opened, such as SQLite or MongoDB. The handle is bound to a tenant scope, so -tenants that share a thread id still never see each other's checkpoints. - -It is the one checkpointer: every user shares `graph.checkpoints`, -`graph.threads` and `graph.writes` and tells itself apart by a namespace -(`DriverCheckpointer::with_namespace`, default `graph`; `delegation` and `flows` -for the delegation graph and workflows). Leases live in `tsd.leases`, claimed -through `Leases`, whose epoch is the write fence. See -`docs/modules/graph/storage-schema.md`. - -It passes the same `testkit::conformance` contracts as the bundled backends: -checkpointer, writes, lineage and concurrency. diff --git a/crates/tinyagents-graph/src/checkpoint/drivers.rs b/crates/tinyagents-graph/src/checkpoint/drivers.rs deleted file mode 100644 index d662cf351..000000000 --- a/crates/tinyagents-graph/src/checkpoint/drivers.rs +++ /dev/null @@ -1,531 +0,0 @@ -//! A [`Checkpointer`] over `tinystoragedrivers` document ports. -//! -//! [`DriverCheckpointer`] puts graph checkpoints, their pending writes and the -//! per-thread execution lease on whatever backend the host opened (SQLite on a -//! desktop, MongoDB in the cloud, memory in tests). The document handle is -//! already bound to a scope, so two tenants' threads never meet even when they -//! share a thread id. -//! -//! # Layout -//! -//! There is one checkpointer: every user of it shares the same three -//! collections and tells itself apart by a namespace (`ns`, default `graph`; -//! the delegation graph uses `delegation` and workflows `flows`). The namespace -//! is part of every id and index, so two namespaces never meet. -//! -//! - `graph.checkpoints`: one document per stored checkpoint, -//! `{ns, thread, subgraph, seq, checkpoint_id, checkpoint}` under -//! `Key::of([ns, thread, seq])`. `subgraph` is the encoded subgraph -//! namespace, so scoped reads filter on it; `checkpoint` is the checkpoint -//! itself. `seq` is a per-thread insertion counter, so listing is a query -//! sorted on it and duplicate checkpoint ids resolve to the latest write, -//! exactly like the append-only backends. -//! - `graph.threads`: one counter document per thread, `{ns, thread, next}`, -//! advanced with [`Repo::update`] so concurrent writers never share a `seq`. -//! - `graph.writes`: one document per `(thread, subgraph, checkpoint)` holding -//! the merged pending writes (see [`merge_writes`]). -//! - Leases live in `tsd.leases` under owner `graph` and key -//! `Key::of([ns, thread])`, claimed through [`Leases`]; the lease's epoch is -//! the write fence (see [`DriverCheckpointer::lease_key`]). - -use std::marker::PhantomData; -use std::sync::Arc; -use std::time::Duration; - -use async_trait::async_trait; -use serde::Serialize; -use serde::de::DeserializeOwned; -use serde_json::Value; -use tinyagents_harness::error::{Result, TinyAgentsError}; -use tinyagents_storage_kit::{ - DocumentStore, DocumentStoreExt, ErrorKind, Filter, Leases, Query, Repo, ScopedStorage, Sort, - StorageError, Versioned, WriteOptions, -}; -use tokio::sync::OnceCell; - -use super::{ - Checkpoint, CheckpointConfig, CheckpointId, CheckpointMetadata, Checkpointer, PendingWrite, - decode_json_err, merge_writes, require_checkpoint_id, -}; -use crate::storage::{ - CHECKPOINTS, CheckpointBody, DEFAULT_NAMESPACE, LEASE_OWNER, SCHEMA_VERSION, THREADS, - ThreadBody, WRITES, WritesBody, checkpoint_id_key, checkpoints_spec, lease_key, leases_spec, - subgraph_key, thread_id_key, threads_spec, writes_id_key, writes_spec, -}; - -/// Map a storage driver failure onto the graph's checkpoint error. -fn map_error(error: StorageError) -> TinyAgentsError { - TinyAgentsError::Checkpoint(format!("storage driver: {error}")) -} - -/// A lease ttl in whole milliseconds; a lease must outlive its own claim, so -/// the shortest is one. -fn ttl_ms(ttl: Duration) -> i64 { - i64::try_from(ttl.as_millis()).unwrap_or(i64::MAX).max(1) -} - -/// A [`Checkpointer`] that stores everything in a driver [`DocumentStore`]. -pub struct DriverCheckpointer { - docs: Arc, - ns: String, - checkpoints: Repo, - threads: Repo, - writes: Repo, - leases: Leases, - declared: Arc>, - _state: PhantomData State>, -} - -impl Clone for DriverCheckpointer { - fn clone(&self) -> Self { - Self { - docs: Arc::clone(&self.docs), - ns: self.ns.clone(), - checkpoints: self.checkpoints.clone(), - threads: self.threads.clone(), - writes: self.writes.clone(), - leases: self.leases.clone(), - declared: Arc::clone(&self.declared), - _state: PhantomData, - } - } -} - -impl std::fmt::Debug for DriverCheckpointer { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("DriverCheckpointer") - .field("ns", &self.ns) - .finish_non_exhaustive() - } -} - -impl DriverCheckpointer { - /// Open the scope's graph storage (running the graph's migrations once per - /// scope, see [`crate::storage::open_storage`]) and store checkpoints in it - /// under the namespace `ns`. Prefer this over [`Self::with_namespace`]. - /// - /// # Errors - /// - /// [`TinyAgentsError::Checkpoint`] when the migration fails. - pub async fn open(storage: ScopedStorage, ns: impl Into) -> Result { - let storage = crate::storage::open_storage(storage) - .await - .map_err(map_error)?; - Ok(Self::with_namespace(Arc::clone(storage.documents()), ns)) - } - - /// [`Self::open`] under the namespace [`DEFAULT_NAMESPACE`] (`graph`). - /// - /// # Errors - /// - /// [`TinyAgentsError::Checkpoint`] when the migration fails. - pub async fn open_default(storage: ScopedStorage) -> Result { - Self::open(storage, DEFAULT_NAMESPACE).await - } - - /// Store checkpoints in the `graph.*` collections of `docs`, under the - /// namespace [`DEFAULT_NAMESPACE`] (`graph`). The caller must have opened - /// the scope's graph storage first; prefer [`Self::open_default`]. - pub fn new(docs: Arc) -> Self { - Self::with_namespace(docs, DEFAULT_NAMESPACE) - } - - /// Store checkpoints in the `graph.*` collections of `docs` under the - /// namespace `ns`, so several independent users of the checkpointer (the - /// delegation graph as `delegation`, workflows as `flows`) can share one - /// backend without meeting. The caller must have opened the scope's graph - /// storage first; prefer [`Self::open`]. - /// - /// # Panics - /// - /// Never: the collection names are valid constants. - #[allow(clippy::expect_used, reason = "the collection names are constants")] - pub fn with_namespace(docs: Arc, ns: impl Into) -> Self { - let options = WriteOptions::new().schema(SCHEMA_VERSION); - fn repo(docs: &Arc, name: &str, options: WriteOptions) -> Repo { - Repo::new(docs, name) - .expect("graph collections are valid names") - .with_options(options) - } - Self { - checkpoints: repo(&docs, CHECKPOINTS, options), - threads: repo(&docs, THREADS, options), - writes: repo(&docs, WRITES, options), - leases: Leases::new(&docs), - docs, - ns: ns.into(), - declared: Arc::new(OnceCell::new()), - _state: PhantomData, - } - } - - /// This checkpointer's namespace. - pub fn namespace(&self) -> &str { - &self.ns - } - - /// The `tsd.leases` owner and key of `thread`'s execution lease. A host may - /// use these identifiers with its own conditional-write policy. - pub fn lease_key(&self, thread: &str) -> (&'static str, String) { - (LEASE_OWNER, lease_key(&self.ns, thread)) - } - - async fn declared(&self) -> Result<()> { - self.declared - .get_or_try_init(|| async { - for spec in [ - checkpoints_spec(), - threads_spec(), - writes_spec(), - leases_spec(), - ] { - self.docs - .ensure_collection(&spec) - .await - .map_err(map_error)?; - } - Ok::<(), TinyAgentsError>(()) - }) - .await - .map(|_| ()) - } - - /// Reserve the next insertion sequence number for `thread`. - async fn next_seq(&self, thread: &str) -> Result { - let (stored, _) = self - .threads - .update(&thread_id_key(&self.ns, thread), |current| { - Ok(Some(ThreadBody { - ns: self.ns.clone(), - thread: thread.to_owned(), - next: current.map_or(1, |counter| counter.next + 1), - })) - }) - .await - .map_err(map_error)? - .ok_or_else(|| { - TinyAgentsError::Checkpoint(format!( - "could not reserve a checkpoint sequence for thread `{thread}`" - )) - })?; - Ok(stored.next - 1) - } - - /// The filter selecting `thread`'s documents in this namespace. - fn thread_filter(&self, thread: &str) -> Filter { - Filter::eq("ns", self.ns.as_str()).and(Filter::eq("thread", thread)) - } - - /// Every stored checkpoint document of `thread`, in insertion order. - async fn thread_docs(&self, thread: &str) -> Result>> { - self.declared().await?; - let query = Query::filter(self.thread_filter(thread)).sort(Sort::asc("seq")); - self.docs - .query_all(CHECKPOINTS, &query) - .await - .map_err(map_error) - } - - async fn drop_writes(&self, filter: Filter) -> Result<()> { - self.docs - .delete_where(WRITES, &filter) - .await - .map(|_| ()) - .map_err(map_error) - } - - /// The newest checkpoint document matching `filter`. - async fn newest(&self, filter: Filter) -> Result>> { - let query = Query::filter(filter).sort(Sort::desc("seq")).limit(1); - let page = self - .docs - .query(CHECKPOINTS, &query) - .await - .map_err(map_error)?; - Ok(page.items.into_iter().next()) - } -} - -impl DriverCheckpointer -where - State: DeserializeOwned, -{ - fn decode(stored: Versioned) -> Result> { - let record = stored.doc.get("checkpoint").cloned().unwrap_or(Value::Null); - // Through the shared classifier, so a `State` that no longer decodes - // is tagged `[schema]` exactly as the file and SQLite backends tag it - // (durable delegations prune such a checkpoint and start fresh). - // (`"record"` names the whole-checkpoint decode to the classifier.) - let mut checkpoint: Checkpoint = serde_json::from_value(record) - .map_err(|error| decode_json_err("storage driver checkpointer", "record", error))?; - checkpoint.normalize(); - Ok(checkpoint) - } -} - -#[async_trait] -impl Checkpointer for DriverCheckpointer -where - State: Serialize + DeserializeOwned + Send + Sync + 'static, -{ - async fn put(&self, checkpoint: Checkpoint) -> Result { - self.declared().await?; - let id = CheckpointId::new(checkpoint.checkpoint_id.clone()); - let seq = self.next_seq(&checkpoint.thread_id).await?; - let body = CheckpointBody { - ns: self.ns.clone(), - thread: checkpoint.thread_id.clone(), - subgraph: subgraph_key(&checkpoint.namespace), - seq, - checkpoint_id: checkpoint.checkpoint_id.clone(), - checkpoint: serde_json::to_value(&checkpoint)?, - }; - self.checkpoints - .insert( - &checkpoint_id_key(&self.ns, &checkpoint.thread_id, seq), - &body, - ) - .await - .map_err(map_error)?; - Ok(id) - } - - async fn get( - &self, - thread_id: &str, - checkpoint_id: Option<&str>, - ) -> Result>> { - self.declared().await?; - let mut filter = self.thread_filter(thread_id); - if let Some(id) = checkpoint_id { - filter = filter.and(Filter::eq("checkpoint_id", id)); - } - self.newest(filter).await?.map(Self::decode).transpose() - } - - /// One indexed query on the stored subgraph namespace, so a parent run and - /// a subgraph sharing a thread never load each other's checkpoints, even - /// when they reuse a checkpoint id. - async fn get_scoped( - &self, - thread_id: &str, - checkpoint_id: Option<&str>, - namespace: &[String], - ) -> Result>> { - self.declared().await?; - let mut filter = self - .thread_filter(thread_id) - .and(Filter::eq("subgraph", subgraph_key(namespace))); - if let Some(id) = checkpoint_id { - filter = filter.and(Filter::eq("checkpoint_id", id)); - } - self.newest(filter).await?.map(Self::decode).transpose() - } - - /// Metadata only: the stored `State` is skipped rather than decoded, so - /// listing a long thread costs its headers, not its states. - async fn list(&self, thread_id: &str) -> Result> { - self.thread_docs(thread_id) - .await? - .into_iter() - .map(|stored| { - let record = stored.doc.get("checkpoint").cloned().unwrap_or(Value::Null); - let mut header: Checkpoint = serde_json::from_value(record) - .map_err(|error| { - decode_json_err("storage driver checkpointer", "record", error) - })?; - header.normalize(); - Ok(header.to_metadata()) - }) - .collect() - } - - async fn get_thread(&self, thread_id: &str) -> Result>> { - self.thread_docs(thread_id) - .await? - .into_iter() - .map(Self::decode) - .collect() - } - - async fn list_threads(&self) -> Result> { - self.declared().await?; - let counters = self - .docs - .query_all(THREADS, &Query::filter(Filter::eq("ns", self.ns.as_str()))) - .await - .map_err(map_error)?; - let mut threads = Vec::new(); - for counter in counters { - let Some(thread) = counter.doc.get("thread").and_then(Value::as_str) else { - continue; - }; - let live = self - .docs - .count(CHECKPOINTS, &self.thread_filter(thread)) - .await - .map_err(map_error)?; - if live > 0 { - threads.push(thread.to_owned()); - } - } - Ok(threads) - } - - async fn delete_thread(&self, thread_id: &str) -> Result<()> { - self.declared().await?; - self.docs - .delete_where(CHECKPOINTS, &self.thread_filter(thread_id)) - .await - .map_err(map_error)?; - // Keep the sequence counter: a later thread of the same name continues - // after it, so its records never interleave with stale cursors. - self.drop_writes(self.thread_filter(thread_id)).await - } - - async fn delete_checkpoints(&self, thread_id: &str, ids: &[String]) -> Result { - if ids.is_empty() { - return Ok(0); - } - self.declared().await?; - let filter = self.thread_filter(thread_id).and(Filter::one_of( - "checkpoint_id", - ids.iter().map(String::as_str), - )); - let removed = self - .docs - .delete_where(CHECKPOINTS, &filter) - .await - .map_err(map_error)?; - self.drop_writes(filter).await?; - Ok(usize::try_from(removed).unwrap_or(usize::MAX)) - } - - async fn put_writes(&self, config: &CheckpointConfig, writes: &[PendingWrite]) -> Result<()> { - let checkpoint_id = require_checkpoint_id(config)?; - if writes.is_empty() { - return Ok(()); - } - self.declared().await?; - let subgraph = subgraph_key(&config.namespace); - let id = writes_id_key(&self.ns, &config.thread_id, &subgraph, &checkpoint_id); - self.writes - .update(&id, |current| { - let mut stored = current.map(|body| body.writes).unwrap_or_default(); - merge_writes(&mut stored, writes); - Ok(Some(WritesBody { - ns: self.ns.clone(), - thread: config.thread_id.clone(), - subgraph: subgraph.clone(), - checkpoint_id: checkpoint_id.clone(), - writes: stored, - })) - }) - .await - .map(|_| ()) - .map_err(|error| { - if error.kind() == ErrorKind::Conflict { - TinyAgentsError::Checkpoint(format!( - "could not record pending writes for thread `{}`", - config.thread_id - )) - } else { - map_error(error) - } - }) - } - - async fn get_writes(&self, config: &CheckpointConfig) -> Result> { - let Some(checkpoint_id) = self.resolve_write_target(config).await? else { - return Ok(Vec::new()); - }; - self.declared().await?; - let id = writes_id_key( - &self.ns, - &config.thread_id, - &subgraph_key(&config.namespace), - &checkpoint_id, - ); - Ok(self - .writes - .get(&id) - .await - .map_err(map_error)? - .map(|found| found.doc.writes) - .unwrap_or_default()) - } - - async fn try_claim(&self, thread: &str, owner: &str, ttl: Duration) -> Result { - self.declared().await?; - match self - .leases - .acquire( - LEASE_OWNER, - &lease_key(&self.ns, thread), - owner, - ttl_ms(ttl), - ) - .await - { - Ok(lease) => Ok(lease.is_some()), - // A claim that kept losing its swap did not get the lease. - Err(error) if error.kind() == ErrorKind::Conflict => Ok(false), - Err(error) => Err(map_error(error)), - } - } - - async fn renew(&self, thread: &str, owner: &str, ttl: Duration) -> Result { - self.declared().await?; - let key = lease_key(&self.ns, thread); - let Some(held) = self - .leases - .current(LEASE_OWNER, &key) - .await - .map_err(map_error)? - else { - return Ok(false); - }; - if held.holder != owner { - return Ok(false); - } - match self - .leases - .renew(LEASE_OWNER, &key, owner, held.epoch, ttl_ms(ttl)) - .await - { - Ok(lease) => Ok(lease.is_some()), - Err(error) if error.kind() == ErrorKind::Conflict => Ok(false), - Err(error) => Err(map_error(error)), - } - } - - async fn release(&self, thread: &str, owner: &str) -> Result<()> { - self.declared().await?; - let key = lease_key(&self.ns, thread); - let Some(held) = self - .leases - .current(LEASE_OWNER, &key) - .await - .map_err(map_error)? - else { - return Ok(()); - }; - if held.holder != owner { - return Ok(()); - } - match self - .leases - .release(LEASE_OWNER, &key, owner, held.epoch) - .await - { - Ok(_) => Ok(()), - // Someone else reclaimed it in between; it is no longer ours. - Err(error) if error.kind() == ErrorKind::Conflict => Ok(()), - Err(error) => Err(map_error(error)), - } - } -} - -#[cfg(test)] -#[path = "drivers_tests.rs"] -mod tests; diff --git a/crates/tinyagents-graph/src/checkpoint/drivers_tests.rs b/crates/tinyagents-graph/src/checkpoint/drivers_tests.rs deleted file mode 100644 index 256689955..000000000 --- a/crates/tinyagents-graph/src/checkpoint/drivers_tests.rs +++ /dev/null @@ -1,371 +0,0 @@ -use std::time::Duration; - -use super::*; -use crate::storage::subgraph_key; -use crate::testkit::conformance::{ - checkpointer_concurrent_contract, checkpointer_contract, checkpointer_lineage_contract, - checkpointer_writes_contract, -}; -use serde::Serialize; -use serde_json::json; -use tinyagents_storage_kit::{Precondition, Scope, StorageBackend}; -use tinystoragedrivers_testkit::MemoryBackend; - -fn docs(storage: &MemoryBackend, scope: &str) -> Arc { - Arc::clone( - storage - .for_scope(&Scope::new(scope).unwrap()) - .unwrap() - .documents(), - ) -} - -fn sample(thread: &str, id: &str, _parent: Option<&str>, step: i32) -> Checkpoint { - Checkpoint::new(step, Vec::new()) - .with_thread_id(thread.to_string()) - .with_checkpoint_id(id.to_string()) -} - -fn checkpointer() -> DriverCheckpointer { - DriverCheckpointer::new(docs(&MemoryBackend::new(), "local")) -} - -#[tokio::test] -async fn passes_the_checkpointer_contract() { - checkpointer_contract(checkpointer()).await; -} - -#[tokio::test] -async fn passes_the_writes_contract() { - checkpointer_writes_contract(checkpointer()).await; -} - -#[tokio::test] -async fn passes_the_lineage_contract() { - checkpointer_lineage_contract(checkpointer()).await; -} - -#[tokio::test] -async fn passes_the_concurrent_contract() { - checkpointer_concurrent_contract(Arc::new(checkpointer())).await; -} - -#[tokio::test] -async fn scopes_and_namespaces_keep_threads_apart() { - let storage = MemoryBackend::new(); - let alice = DriverCheckpointer::::new(docs(&storage, "alice")); - let bob = DriverCheckpointer::::new(docs(&storage, "bob")); - let other = DriverCheckpointer::::with_namespace(docs(&storage, "alice"), "other"); - let checkpoint = sample("t", "c1", None, 1); - alice.put(checkpoint).await.unwrap(); - assert!(bob.get("t", None).await.unwrap().is_none()); - assert!(other.get("t", None).await.unwrap().is_none()); - assert_eq!(alice.list_threads().await.unwrap(), vec!["t".to_owned()]); - assert!(bob.list_threads().await.unwrap().is_empty()); - assert!(format!("{alice:?}").contains("graph")); - let _ = alice.clone(); -} - -#[tokio::test] -async fn leases_follow_the_claim_protocol() { - let cp = checkpointer(); - let minute = Duration::from_secs(60); - assert!(cp.try_claim("t", "a", minute).await.unwrap()); - assert!( - !cp.try_claim("t", "b", minute).await.unwrap(), - "live lease is refused" - ); - assert!( - cp.try_claim("t", "a", minute).await.unwrap(), - "same owner re-claims" - ); - assert!(cp.renew("t", "a", minute).await.unwrap()); - assert!( - !cp.renew("t", "b", minute).await.unwrap(), - "only the owner renews" - ); - assert!(!cp.renew("missing", "a", minute).await.unwrap()); - - cp.release("t", "b").await.unwrap(); - assert!( - !cp.try_claim("t", "b", minute).await.unwrap(), - "foreign release is a no-op" - ); - cp.release("t", "a").await.unwrap(); - cp.release("t", "a").await.unwrap(); - assert!( - cp.try_claim("t", "b", minute).await.unwrap(), - "released lease is free" - ); - - assert!(cp.try_claim("z", "dead", Duration::ZERO).await.unwrap()); - tokio::time::sleep(Duration::from_millis(5)).await; - assert!( - !cp.renew("z", "dead", minute).await.unwrap(), - "expired lease cannot renew" - ); - assert!( - cp.try_claim("z", "new", minute).await.unwrap(), - "expired lease is reclaimable" - ); -} - -#[test] -fn namespaces_encode_injectively() { - let ns = |parts: &[&str]| parts.iter().map(|p| (*p).to_string()).collect::>(); - assert_ne!( - subgraph_key(&ns(&["a", "b"])), - subgraph_key(&ns(&["a\u{1f}b"])) - ); - assert_ne!(subgraph_key(&ns(&["a", "b"])), subgraph_key(&ns(&["a/b"]))); - assert_ne!(subgraph_key(&ns(&["ab"])), subgraph_key(&ns(&["a", "b"]))); - assert_eq!(subgraph_key(&[]), ""); - assert_ne!( - subgraph_key(&ns(&[""])), - subgraph_key(&[]), - "the empty segment is not the root" - ); -} - -#[tokio::test] -async fn scoped_reads_stay_in_their_namespace_when_ids_repeat() { - let saver = checkpointer(); - let child = vec!["sub".to_string()]; - saver.put(sample("t", "same", None, 1)).await.unwrap(); - saver - .put(sample("t", "same", None, 2).with_namespace(child.clone())) - .await - .unwrap(); - let root = saver - .get_scoped("t", Some("same"), &[]) - .await - .unwrap() - .unwrap(); - assert_eq!( - root.state, 1, - "the root never loads the subgraph's checkpoint" - ); - let nested = saver - .get_scoped("t", Some("same"), &child) - .await - .unwrap() - .unwrap(); - assert_eq!(nested.state, 2); - assert_eq!( - saver - .get_scoped("t", None, &[]) - .await - .unwrap() - .unwrap() - .state, - 1 - ); - assert!( - saver - .get_scoped("t", None, &["other".to_string()]) - .await - .unwrap() - .is_none() - ); -} - -#[tokio::test] -async fn pending_writes_never_merge_across_lookalike_namespaces() { - let saver = checkpointer(); - let config = |namespace: Vec| CheckpointConfig { - thread_id: "t".to_string(), - checkpoint_id: Some("c".to_string()), - namespace, - }; - let split = config(vec!["a".to_string(), "b".to_string()]); - let joined = config(vec!["a\u{1f}b".to_string()]); - saver - .put_writes( - &split, - &[PendingWrite::data("n", "task", 0, "out", json!("split"))], - ) - .await - .unwrap(); - assert!(saver.get_writes(&joined).await.unwrap().is_empty()); - assert_eq!(saver.get_writes(&split).await.unwrap().len(), 1); -} - -#[test] -fn driver_failures_are_checkpoint_errors() { - let error = map_error(StorageError::unavailable("busy")); - assert!(matches!(error, TinyAgentsError::Checkpoint(ref m) if m.contains("busy"))); -} - -#[tokio::test] -async fn a_corrupt_record_is_a_checkpoint_error() { - let storage = MemoryBackend::new(); - let docs = docs(&storage, "local"); - let cp = DriverCheckpointer::::new(Arc::clone(&docs)); - cp.put(sample("t", "c1", None, 1)).await.unwrap(); - let stored = cp.thread_docs("t").await.unwrap().remove(0); - docs.put( - "graph.checkpoints", - &stored.id, - json!({"ns": "graph", "thread": "t", "subgraph": "", "seq": 0, "checkpoint_id": "c1", "checkpoint": "nope"}), - Precondition::None, - ) - .await - .unwrap(); - let error = cp.get("t", None).await.unwrap_err(); - assert!(matches!(error, TinyAgentsError::Checkpoint(_)), "{error:?}"); -} - -/// A `State` whose `Deserialize` counts every call, so a test can assert how -/// many full states a read decoded rather than only what it returned. -#[derive(Debug, Clone, Serialize)] -struct CountedState(i32); - -static STATE_DECODES: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0); - -impl<'de> serde::Deserialize<'de> for CountedState { - fn deserialize>( - deserializer: D, - ) -> std::result::Result { - STATE_DECODES.fetch_add(1, std::sync::atomic::Ordering::SeqCst); - i32::deserialize(deserializer).map(CountedState) - } -} - -#[tokio::test] -async fn list_and_limited_history_decode_only_what_they_return() { - use std::sync::atomic::Ordering; - let cp = DriverCheckpointer::::new(docs(&MemoryBackend::new(), "local")); - let mut parent: Option = None; - for step in 0..40 { - let id = format!("c{step}"); - let checkpoint = Checkpoint::new(CountedState(step), Vec::new()) - .with_thread_id("t".to_string()) - .with_checkpoint_id(id.clone()) - .with_parent_checkpoint_id(parent.clone()) - .with_metadata(json!({ "source": "loop", "step": step })); - cp.put(checkpoint).await.unwrap(); - parent = Some(id); - } - - STATE_DECODES.store(0, Ordering::SeqCst); - let listed = cp.list("t").await.unwrap(); - assert_eq!(listed.len(), 40); - assert_eq!(listed[0].checkpoint_id, "c0"); - assert_eq!(listed[39].checkpoint_id, "c39"); - assert_eq!( - STATE_DECODES.load(Ordering::SeqCst), - 0, - "listing must not deserialize any record's full State" - ); - - let recent = cp.state_history("t", &[], Some(1)).await.unwrap(); - assert_eq!(recent.len(), 1); - assert_eq!(recent[0].checkpoint.checkpoint_id, "c39"); - assert_eq!( - STATE_DECODES.load(Ordering::SeqCst), - 1, - "a one-record history decodes one record, not the whole chain" - ); -} - -#[tokio::test] -async fn put_with_writes_persists_the_checkpoint_and_its_writes() { - let cp = checkpointer(); - let writes = vec![PendingWrite { - node: tinyagents_harness::ids::NodeId::from("n"), - task_id: tinyagents_harness::ids::TaskId::from("task-1"), - idx: 0, - channel: "out".to_string(), - payload: json!("hi"), - }]; - let id = cp - .put_with_writes(sample("t", "c1", None, 1), &writes) - .await - .unwrap(); - assert_eq!(id.as_str(), "c1"); - assert!(cp.get("t", Some("c1")).await.unwrap().is_some()); - let config = CheckpointConfig { - thread_id: "t".to_string(), - checkpoint_id: Some("c1".to_string()), - namespace: vec![], - }; - let stored = cp.get_writes(&config).await.unwrap(); - assert_eq!(stored.len(), 1); - assert_eq!(stored[0].channel, "out"); -} - -#[tokio::test] -async fn documents_are_v2_bodies_with_the_namespace_in_ids_and_fields() { - let storage = MemoryBackend::new(); - let docs = docs(&storage, "local"); - let delegation = DriverCheckpointer::::with_namespace(Arc::clone(&docs), "delegation"); - let flows = DriverCheckpointer::::with_namespace(Arc::clone(&docs), "flows"); - delegation.put(sample("a/b", "c1", None, 1)).await.unwrap(); - flows.put(sample("a/b", "c1", None, 2)).await.unwrap(); - - let row = docs - .get("graph.checkpoints", &format!("delegation/a%2Fb/{:020}", 0)) - .await - .unwrap() - .expect("id is Key::of([ns, thread, seq])"); - assert_eq!(row.doc["ns"], "delegation"); - assert_eq!(row.doc["thread"], "a/b"); - assert_eq!(row.doc["subgraph"], ""); - assert_eq!(row.doc["seq"], 0); - assert_eq!(row.doc["checkpoint_id"], "c1"); - assert!( - row.doc["checkpoint"].is_object(), - "typed, not a JSON string" - ); - assert_eq!(row.meta.schema, 4); - let counter = docs - .get("graph.threads", "flows/a%2Fb") - .await - .unwrap() - .expect("id is Key::of([ns, thread])"); - assert_eq!( - (counter.doc["ns"].as_str(), counter.doc["next"].as_u64()), - (Some("flows"), Some(1)) - ); - assert_eq!(delegation.get("a/b", None).await.unwrap().unwrap().state, 1); - assert_eq!(flows.get("a/b", None).await.unwrap().unwrap().state, 2); - assert_eq!(delegation.list_threads().await.unwrap(), ["a/b"]); -} - -#[tokio::test] -async fn the_lease_lives_in_tsd_leases_and_its_epoch_moves_on_every_new_tenure() { - let storage = MemoryBackend::new(); - let docs = docs(&storage, "local"); - let cp = DriverCheckpointer::::with_namespace(Arc::clone(&docs), "flows"); - let minute = Duration::from_secs(60); - let (owner, key) = cp.lease_key("t"); - let leases = tinyagents_storage_kit::Leases::new(&docs); - - assert!(cp.try_claim("t", "a", minute).await.unwrap()); - let first = leases.current(owner, &key).await.unwrap().unwrap(); - assert_eq!((first.holder.as_str(), first.epoch), ("a", 1)); - assert!(cp.try_claim("t", "a", minute).await.unwrap()); - assert_eq!( - leases.current(owner, &key).await.unwrap().unwrap().epoch, - 1, - "re-claim keeps the epoch" - ); - - cp.release("t", "a").await.unwrap(); - assert!(leases.current(owner, &key).await.unwrap().is_none()); - assert!(cp.try_claim("t", "a", minute).await.unwrap()); - let second = leases.current(owner, &key).await.unwrap().unwrap(); - assert_eq!( - second.epoch, 2, - "a new tenure of the same holder is a new epoch" - ); - assert!( - !cp.try_claim("t", "b", minute).await.unwrap(), - "namespaces share no lease, but a thread's claim is exclusive" - ); - let other = DriverCheckpointer::::with_namespace(Arc::clone(&docs), "delegation"); - assert!( - other.try_claim("t", "b", minute).await.unwrap(), - "another namespace's thread is free" - ); -} diff --git a/crates/tinyagents-graph/src/checkpoint/mod.rs b/crates/tinyagents-graph/src/checkpoint/mod.rs deleted file mode 100644 index 909f080fc..000000000 --- a/crates/tinyagents-graph/src/checkpoint/mod.rs +++ /dev/null @@ -1,852 +0,0 @@ -//! Checkpointer trait and in-memory backend — the durability layer that makes -//! the recursive graph runtime resumable and time-travelable. -//! -//! In a recursive-language-model harness, runs nest: a graph node can run -//! another compiled graph, which can run another, each producing its own state. -//! Checkpointing snapshots every level of that tree at superstep boundaries and -//! keys them by `thread_id`/`namespace` so a parent and its embedded subgraphs -//! never collide (see [`crate::subgraph`]). Persisting committed state at -//! each boundary is what lets a run be paused on an interrupt, resumed later, -//! forked, or replayed for time-travel debugging. -//! -//! See `types` for the checkpoint record definitions. Checkpoints are written -//! at superstep boundaries only — never mid-node — so resuming always reruns a -//! node from its start. - -mod drivers; -mod types; - -pub use drivers::DriverCheckpointer; -pub use types::{ - 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}; -use std::sync::{Arc, Mutex}; - -use async_trait::async_trait; - -use crate::{Result, TinyAgentsError}; -use tinyagents_harness::ids::CheckpointId; - -/// Builds a `TinyAgentsError::Checkpoint` for a failed `serde_json` decode, -/// tagging the message with whether the failure looks like a genuine -/// shape/schema mismatch or ambiguous/likely corruption. -/// -/// `serde_json::Error::classify()` distinguishes `Category::Data` (the JSON -/// parsed fine but didn't match the target type — a missing field, a renamed -/// variant, a type change: exactly what a legacy or newer schema produces) -/// from `Category::Syntax` / `Category::Eof` / `Category::Io` (the bytes -/// themselves are malformed or truncated — a write torn by a crash, disk -/// corruption, or a partial copy). Both backends previously tagged every -/// failure here identically as `"decode …"`, and a caller (see -/// `graph::delegation::run::is_incompatible_checkpoint_error`) that expired -/// any checkpoint whose error mentioned "decode" would delete a corrupted -/// record indistinguishably from an outdated one — silently discarding -/// evidence of real storage corruption instead of surfacing it. This tag is -/// the stable discriminator that lets the caller tell the two apart. -/// -/// `backend` is the existing `"sqlite checkpointer" / "file checkpointer"` -/// prefix; `what` names the field being decoded (e.g. `"record"`, -/// `"next_nodes"`) with no leading `"decode "` — this function supplies that -/// wording once, alongside the tag. -/// -/// **Only `what == "record"` — the whole-`Checkpoint` decode, where -/// `State` is the one thing this module knows can legitimately evolve across -/// a caller's schema versions — is eligible for `[schema]` at all.** -/// `"namespace"`, `"next_nodes"`, `"header"` and `"write payload"` are -/// checkpoint-internal engine metadata with no versioned shape of their own; -/// a `Category::Data` failure decoding any of those is corruption by -/// construction; classifying it as `[schema]` and pruning would have been -/// simply wrong. -/// -/// Even for `"record"`, `Category::Data` is a HEURISTIC, not a proof: this -/// module is generic over `State` (see the module docs) and has no way to -/// positively identify a specific caller's known-legacy shape from a byte -/// stream alone — a decode failure produced by data corruption that happens -/// to still be syntactically valid JSON (e.g. a `state.schema_version` field -/// whose value was flipped from a number to a string) also classifies as -/// `Category::Data`, indistinguishable at this layer from a genuine older -/// schema. This is an accepted, bounded residual risk, not a claim of -/// perfect precision: the worst case is an unnecessary fresh restart from -/// `plan` for what was actually corruption, which is a strict improvement -/// over the state before this classifier existed (100% of decode failures, -/// corrupt or not, were pruned identically). Closing the remaining gap -/// precisely would need the generic `Checkpointer` trait to expose the raw -/// bytes on a failed decode so a caller could apply its OWN positive-shape -/// check — a real API addition, tracked as follow-up rather than done here. -pub(super) fn decode_json_err( - backend: &str, - what: &str, - err: serde_json::Error, -) -> TinyAgentsError { - let tag = if what == "record" { - match err.classify() { - serde_json::error::Category::Data => "schema", - serde_json::error::Category::Syntax - | serde_json::error::Category::Eof - | serde_json::error::Category::Io => "corrupt", - } - } else { - "corrupt" - }; - TinyAgentsError::Checkpoint(format!("{backend}: decode [{tag}] {what}: {err}")) -} - -/// Persists and retrieves graph checkpoints keyed by thread. -#[async_trait] -pub trait Checkpointer: Send + Sync -where - State: Send + Sync + 'static, -{ - /// Persists a checkpoint and returns its id. - async fn put(&self, checkpoint: Checkpoint) -> Result; - - /// Loads a checkpoint for a thread. When `checkpoint_id` is `None`, returns - /// the latest checkpoint for the thread. - async fn get( - &self, - thread_id: &str, - checkpoint_id: Option<&str>, - ) -> Result>>; - - /// Loads a checkpoint for a thread scoped to `namespace`. - /// - /// Like [`Checkpointer::get`], but only considers checkpoints whose stored - /// namespace equals `namespace`. This is what keeps a parent run and the - /// subgraphs it embeds — which share a thread id but differ in namespace — - /// from loading each other's checkpoints on resume/inspection. With - /// `checkpoint_id == None` the latest checkpoint *in that namespace* is - /// 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 - /// `FileCheckpointer` and `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, - thread_id: &str, - checkpoint_id: Option<&str>, - namespace: &[String], - ) -> Result>> { - let metas = self.list(thread_id).await?; - let target: Option = match checkpoint_id { - Some(id) => metas - .iter() - .rev() - .find(|m| m.checkpoint_id == id && m.namespace.as_slice() == namespace) - .map(|m| m.checkpoint_id.clone()), - None => metas - .iter() - .rev() - .find(|m| m.namespace.as_slice() == namespace) - .map(|m| m.checkpoint_id.clone()), - }; - match target { - Some(id) => self.get(thread_id, Some(&id)).await, - None => Ok(None), - } - } - - /// Lists checkpoint metadata for a thread in insertion order. - async fn list(&self, thread_id: &str) -> Result>; - - // ---- Pending writes ---------------------------------------------------- - // - // The partial-failure protocol. A superstep can fail after some of its - // tasks have already run; without a per-task record of what they wrote, a - // resume cannot tell "already ran" from "not yet run" and re-executes their - // side effects. `put_writes` records that work against the checkpoint it - // belongs to; `get_writes` reads it back, and `get_tuple` surfaces it as - // `CheckpointTuple::pending_writes` so resume can skip completed tasks. - // - // Both carry default no-op bodies so an out-of-tree `Checkpointer` keeps - // compiling: such a backend simply never persists writes, which is exactly - // the behaviour every backend had before the protocol existed. - - /// Records `writes` against the checkpoint addressed by `config`. - /// - /// `config.checkpoint_id` must name a specific checkpoint — writes belong - /// to the boundary they were produced at, so a `None` id has no meaning and - /// backends reject it. - /// - /// Idempotency follows [`PendingWrite`]'s replace-vs-ignore rule: a data - /// write (`idx >= 0`) whose `(task_id, idx)` is already stored is ignored, - /// while a control-plane write (`idx < 0`) replaces the stored value. Both - /// are implemented through [`merge_writes`], so every backend agrees. - /// - /// The default body is a no-op returning `Ok(())`. - async fn put_writes(&self, _config: &CheckpointConfig, _writes: &[PendingWrite]) -> Result<()> { - 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` did, 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 = checkpoint.config(); - 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. - /// - /// Returns an empty vec for an unknown checkpoint or one that has no - /// writes. When `config.checkpoint_id` is `None` the latest checkpoint in - /// `config.namespace` is resolved first. - /// - /// The default body returns an empty vec. - async fn get_writes(&self, _config: &CheckpointConfig) -> Result> { - 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 - /// bug), a read may legitimately mean "the latest checkpoint in this - /// namespace" — the same relaxation [`Checkpointer::get`] makes. Returns - /// `None` when the thread/namespace has no checkpoint at all. - async fn resolve_write_target(&self, config: &CheckpointConfig) -> Result> { - match &config.checkpoint_id { - Some(id) => Ok(Some(id.clone())), - None => Ok(self - .get_scoped(&config.thread_id, None, &config.namespace) - .await? - .map(|c| c.checkpoint_id)), - } - } - - /// Loads every checkpoint stored under `thread_id`, in listing order. - /// - /// This is the bulk-read companion to [`Checkpointer::list`]: it returns - /// full [`Checkpoint`] records (including state) rather than metadata, so - /// whole-thread operations such as [`Checkpointer::copy_thread`] can read - /// a thread once instead of issuing one [`Checkpointer::get`] per - /// checkpoint. - /// - /// The default is composed from [`Checkpointer::list`] + - /// [`Checkpointer::get`], which re-resolves each id individually (and, on - /// a backend whose `get` scans the whole thread, is O(H²)). Every bundled - /// backend overrides it with a single-pass read; other backends should do - /// the same when they can. - async fn get_thread(&self, thread_id: &str) -> Result>> { - let metas = self.list(thread_id).await?; - let mut out = Vec::with_capacity(metas.len()); - for meta in metas { - if let Some(checkpoint) = self.get(thread_id, Some(&meta.checkpoint_id)).await? { - out.push(checkpoint); - } - } - Ok(out) - } - - /// Loads a [`CheckpointTuple`] — the checkpoint plus its addressing config, - /// its parent's config, and the pending writes carried with it. - /// - /// Composed from [`Checkpointer::get`] so every backend gets it for free; - /// override it only when a backend can build the tuple more cheaply. When - /// `config.checkpoint_id` is `None` the latest checkpoint is returned. - async fn get_tuple(&self, config: CheckpointConfig) -> Result>> { - let Some(checkpoint) = self - .get_scoped( - &config.thread_id, - config.checkpoint_id.as_deref(), - &config.namespace, - ) - .await? - else { - return Ok(None); - }; - let resolved = checkpoint.config(); - let parent_config = checkpoint.parent_config(); - let pending_writes = self.resolved_writes(&resolved, &checkpoint).await?; - Ok(Some(CheckpointTuple { - config: resolved, - checkpoint, - parent_config, - pending_writes, - })) - } - - /// The writes to surface on a tuple for `checkpoint`. - /// - /// Prefers the separately persisted [`Checkpointer::get_writes`] records — - /// the authoritative partial-failure ledger — and falls back to the - /// checkpoint record's own inline `pending_writes` for backends that do not - /// implement the write protocol (whose `get_writes` default returns empty) - /// and for records written before it existed. - async fn resolved_writes( - &self, - config: &CheckpointConfig, - checkpoint: &Checkpoint, - ) -> Result> { - let stored = self.get_writes(config).await?; - if stored.is_empty() { - Ok(checkpoint.pending_writes.clone()) - } else { - Ok(stored) - } - } - - /// Returns a thread's checkpoint lineage newest-first, following each - /// checkpoint's `parent_checkpoint_id` from the latest checkpoint in - /// `namespace`. `limit` caps the number of tuples returned (the most recent - /// ones). - /// - /// The default walks [`Checkpointer::get_tuple`] once per hop, so a backend - /// whose scoped lookup re-reads the whole thread is O(H²) over the lineage. - /// Every bundled backend overrides it to read the thread (or the - /// namespace's rows) once and walk the lineage in memory, so none of the - /// three is in that class: the JSONL backend parses its file once, and the - /// SQLite backend issues one indexed range query — the - /// `(thread_id, namespace, seq)` index is what makes the namespace scope - /// expressible in SQL at all, and without it that backend silently fell - /// back to this default. The observable result is identical to iterating - /// `get_tuple` by parent pointer. - /// - /// The walk carries a **visited set**. `parent_checkpoint_id` is caller-set - /// data, not a structurally enforced acyclic pointer: a hand-written - /// checkpoint, a bad fork, or a `copy_thread` that reused ids can point a - /// record at itself or at one of its descendants. With `limit == None` an - /// unguarded walk then never terminates — it does not merely return a wrong - /// answer, it hangs the caller. Revisiting an id ends the walk. - /// `FileCheckpointer` had always had this guard (its in-memory `remove` - /// doubles as one); this makes it uniform. - async fn state_history( - &self, - thread_id: &str, - namespace: &[String], - limit: Option, - ) -> Result>> { - let mut out = Vec::new(); - let mut cursor: Option = None; - let mut visited: HashSet = HashSet::new(); - loop { - if let Some(limit) = limit - && out.len() >= limit - { - break; - } - let config = CheckpointConfig { - thread_id: thread_id.to_string(), - checkpoint_id: cursor.clone(), - namespace: namespace.to_vec(), - }; - let Some(tuple) = self.get_tuple(config).await? else { - break; - }; - if !visited.insert(tuple.checkpoint.checkpoint_id.clone()) { - tracing::warn!( - "[checkpoint] state_history: lineage cycle at checkpoint `{}` \ - (thread `{thread_id}`); truncating the walk", - tuple.checkpoint.checkpoint_id - ); - break; - } - let parent = tuple.checkpoint.parent_checkpoint_id.clone(); - out.push(tuple); - match parent { - Some(parent) => cursor = Some(parent), - None => break, - } - } - 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`, - // `delete_checkpoints`) have no default body. The higher-level operations - // (`delete_by_run`, `copy_thread`, `prune`) are composed from those plus - // the existing `list`/`get_thread`/`put` surface, so every backend - // inherits them for free and only implements the three storage primitives - // (overriding `get_thread` when a single-pass bulk read is available). - - /// Lists the ids of every thread that currently has at least one checkpoint. - /// - /// Order is backend-defined. Storage-specific: there is no default body. - async fn list_threads(&self) -> Result>; - - /// Deletes every checkpoint stored under `thread_id`. - /// - /// A no-op (still `Ok`) when the thread is unknown. Storage-specific. - async fn delete_thread(&self, thread_id: &str) -> Result<()>; - - /// Low-level primitive: removes the named checkpoints from `thread_id`, - /// returning how many were actually removed. - /// - /// Ids not present are ignored. The default thread operations - /// ([`Checkpointer::delete_by_run`], [`Checkpointer::prune`]) are built on - /// top of this. Storage-specific: there is no default body. - async fn delete_checkpoints(&self, thread_id: &str, ids: &[String]) -> Result; - - /// Deletes every checkpoint in `thread_id` stamped with `run_id`, returning - /// the number removed. - /// - /// Run ids are recorded on checkpoints by the executor; records that predate - /// run-id stamping (or were written manually) carry `None` and are never - /// matched. Composed from [`Checkpointer::list`] + - /// [`Checkpointer::delete_checkpoints`]. - async fn delete_by_run(&self, thread_id: &str, run_id: &str) -> Result { - let ids: Vec = self - .list(thread_id) - .await? - .into_iter() - .filter(|m| m.run_id.as_deref() == Some(run_id)) - .map(|m| m.checkpoint_id) - .collect(); - self.delete_checkpoints(thread_id, &ids).await - } - - /// Deep-copies every checkpoint from `source_thread` into `target_thread`, - /// rewriting only the `thread_id` while preserving each record's - /// `checkpoint_id` and `parent_checkpoint_id`. - /// - /// Because checkpoint ids are unique only within a thread, reusing them in - /// the target keeps the parent lineage spine intact, so time-travel and - /// resume walk the copied thread exactly as they would the source. Records - /// are copied in listing order so parents always precede their children. - /// Composed from [`Checkpointer::get_thread`] + [`Checkpointer::put`], so - /// the source thread is read once (a bulk read, not one - /// [`Checkpointer::get`] per checkpoint). - /// - /// # The target must be empty - /// - /// Copying preserves each record's `checkpoint_id` — that is what keeps the - /// lineage spine intact — so appending a lineage onto a thread that already - /// has one produces a file/table containing two disjoint lineages *with - /// reused ids*. Every subsequent `get(Some(id))` then resolves to whichever - /// copy was written last and the parent walk crosses between lineages: the - /// thread is silently corrupt, with no error at the point of damage. - /// - /// So a non-empty target is **rejected** rather than merged into. Callers - /// that genuinely want to overwrite call [`Checkpointer::delete_thread`] - /// first, which makes the destructive intent explicit. Copying an empty or - /// unknown source thread is a no-op (still `Ok`). - async fn copy_thread(&self, source_thread: &str, target_thread: &str) -> Result<()> { - let existing = self.list(target_thread).await?; - if !existing.is_empty() { - return Err(TinyAgentsError::Checkpoint(format!( - "copy_thread: target thread `{target_thread}` already has {} checkpoint(s); \ - copying would interleave two lineages with reused checkpoint ids. \ - Delete the target first if replacing it is intended.", - existing.len() - ))); - } - for mut checkpoint in self.get_thread(source_thread).await? { - let source_config = CheckpointConfig { - thread_id: source_thread.to_string(), - checkpoint_id: Some(checkpoint.checkpoint_id.clone()), - namespace: checkpoint.namespace.clone(), - }; - let writes = self.get_writes(&source_config).await?; - checkpoint.thread_id = target_thread.to_string(); - let target_config = CheckpointConfig { - thread_id: target_thread.to_string(), - checkpoint_id: Some(checkpoint.checkpoint_id.clone()), - namespace: checkpoint.namespace.clone(), - }; - self.put(checkpoint).await?; - if !writes.is_empty() { - self.put_writes(&target_config, &writes).await?; - } - } - Ok(()) - } - - /// Prunes old checkpoints from `thread_id`, retaining the most recent - /// `keep_last` plus everything they depend on, and returns the number - /// deleted. - /// - /// Strategy (lineage- and delta-safe): - /// - /// 1. Protect the most recent `keep_last` checkpoints (listing order) *of - /// every namespace present in the thread*. An embedded subgraph writes - /// its checkpoints under the parent's thread id but its own namespace, - /// and its lineage is disjoint from the parent's (no parent-namespace - /// record ever references a child-namespace id), so a thread-wide - /// recency window would delete the child lineage outright and leave the - /// thread unresumable. - /// 2. Walk the `parent_checkpoint_id` chain of every protected checkpoint - /// and protect every ancestor reached. This is what honors the - /// delta-channel warning: a kept checkpoint that only stores a delta (or - /// depends on an ancestor's pending writes / snapshot) keeps its entire - /// ancestor chain, so it can never be left dangling without the state it - /// needs to be reconstructed or resumed. - /// 3. Delete every checkpoint not in the protected set. - /// - /// `keep_last == 0` is treated as `keep_last == 1`: the latest checkpoint - /// (and its ancestors) is always retained so the thread stays resumable. - /// Composed from [`Checkpointer::list`] + [`Checkpointer::delete_checkpoints`]. - async fn prune(&self, thread_id: &str, keep_last: usize) -> Result { - let metas = self.list(thread_id).await?; - if metas.is_empty() { - return Ok(0); - } - let keep_last = keep_last.max(1); - - // Index by id so ancestor walks are O(depth). - let mut parent_of: HashMap<&str, Option<&str>> = HashMap::new(); - for m in &metas { - parent_of.insert(m.checkpoint_id.as_str(), m.parent_checkpoint_id.as_deref()); - } - // Group by namespace so each lineage (the root run and every embedded - // subgraph run sharing this thread) gets its own recency window. - let mut by_namespace: HashMap<&Vec, Vec<&CheckpointMetadata>> = HashMap::new(); - for m in &metas { - by_namespace.entry(&m.namespace).or_default().push(m); - } - - let mut protected: HashSet = HashSet::new(); - // Step 1: the recency window, per namespace. - for group in by_namespace.values() { - for m in group.iter().rev().take(keep_last) { - protected.insert(m.checkpoint_id.clone()); - } - } - // Step 2: expand to every ancestor of a protected checkpoint. - let window: Vec = protected.iter().cloned().collect(); - for start in window { - let mut cursor = parent_of.get(start.as_str()).copied().flatten(); - while let Some(parent) = cursor { - if !protected.insert(parent.to_string()) { - break; // already protected — its chain is too. - } - cursor = parent_of.get(parent).copied().flatten(); - } - } - - // Step 3: delete the rest. - let to_delete: Vec = metas - .iter() - .map(|m| m.checkpoint_id.clone()) - .filter(|id| !protected.contains(id)) - .collect(); - self.delete_checkpoints(thread_id, &to_delete).await - } -} - -/// An in-memory [`Checkpointer`] backed by an `Arc>`. -/// -/// Cheap to clone; clones share the same underlying store. -pub struct InMemoryCheckpointer { - inner: Arc>>>>, - /// Pending writes keyed by `(thread_id, namespace, checkpoint_id)` — the - /// same identity the SQL backends use as a primary key prefix. - writes: Arc>>>, -} - -/// The address a batch of pending writes is filed under. -type WriteKey = (String, Vec, String); - -impl InMemoryCheckpointer { - /// Creates an empty checkpointer. - pub fn new() -> Self { - Self { - inner: Arc::new(Mutex::new(HashMap::new())), - writes: Arc::new(Mutex::new(HashMap::new())), - } - } - - /// Returns the number of checkpoints stored for a thread. - pub fn count(&self, thread_id: &str) -> usize { - self.inner - .lock() - .map(|m| m.get(thread_id).map(|v| v.len()).unwrap_or(0)) - .unwrap_or(0) - } -} - -impl Default for InMemoryCheckpointer { - fn default() -> Self { - Self::new() - } -} - -impl Clone for InMemoryCheckpointer { - fn clone(&self) -> Self { - Self { - inner: self.inner.clone(), - writes: self.writes.clone(), - } - } -} - -fn lock_err() -> TinyAgentsError { - TinyAgentsError::Checkpoint("in-memory checkpointer lock poisoned".to_string()) -} - -fn metadata_of(c: &Checkpoint) -> CheckpointMetadata { - c.to_metadata() -} - -#[async_trait] -impl Checkpointer for InMemoryCheckpointer -where - State: Clone + Send + Sync + 'static, -{ - async fn put(&self, checkpoint: Checkpoint) -> Result { - let id = CheckpointId::new(checkpoint.checkpoint_id.clone()); - let mut map = self.inner.lock().map_err(|_| lock_err())?; - map.entry(checkpoint.thread_id.clone()) - .or_default() - .push(checkpoint); - Ok(id) - } - - async fn get( - &self, - thread_id: &str, - checkpoint_id: Option<&str>, - ) -> Result>> { - let map = self.inner.lock().map_err(|_| lock_err())?; - let Some(list) = map.get(thread_id) else { - return Ok(None); - }; - // Duplicate-id lookup resolves to the *last* written record, matching - // the append-only file/sqlite backends (and `get(None)`, which returns - // the latest). Pinning one semantic keeps the three backends - // interchangeable — see the checkpointer conformance suite. - let found = match checkpoint_id { - Some(id) => list.iter().rfind(|c| c.checkpoint_id == id), - None => list.last(), - }; - Ok(found.cloned().map(|mut c| { - c.normalize(); - c - })) - } - - async fn list(&self, thread_id: &str) -> Result> { - let map = self.inner.lock().map_err(|_| lock_err())?; - Ok(map - .get(thread_id) - .map(|list| list.iter().map(metadata_of).collect()) - .unwrap_or_default()) - } - - async fn get_thread(&self, thread_id: &str) -> Result>> { - // 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())?; - 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> { - let map = self.inner.lock().map_err(|_| lock_err())?; - Ok(map - .iter() - .filter(|(_, list)| !list.is_empty()) - .map(|(thread, _)| thread.clone()) - .collect()) - } - - async fn delete_thread(&self, thread_id: &str) -> Result<()> { - let mut map = self.inner.lock().map_err(|_| lock_err())?; - map.remove(thread_id); - // Writes are keyed by thread too, and deleting a thread must not leave - // its write ledger behind for a later thread of the same name to - // inherit. The conformance suite asserts exactly this. - let mut writes = self.writes.lock().map_err(|_| lock_err())?; - writes.retain(|(thread, _, _), _| thread != thread_id); - Ok(()) - } - - async fn delete_checkpoints(&self, thread_id: &str, ids: &[String]) -> Result { - if ids.is_empty() { - return Ok(0); - } - let drop: HashSet<&str> = ids.iter().map(String::as_str).collect(); - let mut map = self.inner.lock().map_err(|_| lock_err())?; - let Some(list) = map.get_mut(thread_id) else { - return Ok(0); - }; - let before = list.len(); - list.retain(|c| !drop.contains(c.checkpoint_id.as_str())); - let removed = before - list.len(); - drop_writes_for(&self.writes, thread_id, &drop)?; - Ok(removed) - } - - async fn put_writes(&self, config: &CheckpointConfig, writes: &[PendingWrite]) -> Result<()> { - let checkpoint_id = require_checkpoint_id(config)?; - if writes.is_empty() { - return Ok(()); - } - let key: WriteKey = ( - config.thread_id.clone(), - config.namespace.clone(), - checkpoint_id, - ); - let mut map = self.writes.lock().map_err(|_| lock_err())?; - let slot = map.entry(key).or_default(); - let changed = merge_writes(slot, writes); - tracing::debug!( - "[checkpoint:memory] put_writes thread={} checkpoint={:?} offered={} stored={}", - config.thread_id, - config.checkpoint_id, - writes.len(), - changed - ); - Ok(()) - } - - async fn get_writes(&self, config: &CheckpointConfig) -> Result> { - let Some(checkpoint_id) = self.resolve_write_target(config).await? else { - return Ok(Vec::new()); - }; - let key: WriteKey = ( - config.thread_id.clone(), - config.namespace.clone(), - checkpoint_id, - ); - let map = self.writes.lock().map_err(|_| lock_err())?; - Ok(map.get(&key).cloned().unwrap_or_default()) - } -} - -/// Removes the write ledgers of `ids` within `thread_id`. -fn drop_writes_for( - writes: &Mutex>>, - thread_id: &str, - ids: &HashSet<&str>, -) -> Result<()> { - let mut map = writes.lock().map_err(|_| lock_err())?; - map.retain(|(thread, _, checkpoint), _| { - thread != thread_id || !ids.contains(checkpoint.as_str()) - }); - Ok(()) -} - -/// Extracts the checkpoint id a `put_writes` call addresses. -/// -/// A write belongs to the boundary that produced it, so an unaddressed -/// (`None`) id is a caller bug rather than "the latest": silently filing the -/// writes against whatever checkpoint happens to be newest is precisely the -/// corruption the protocol exists to prevent. -pub(crate) fn require_checkpoint_id(config: &CheckpointConfig) -> Result { - config.checkpoint_id.clone().ok_or_else(|| { - TinyAgentsError::Checkpoint(format!( - "put_writes requires an explicit checkpoint_id (thread `{}`)", - config.thread_id - )) - }) -} - -#[cfg(test)] -#[path = "mod_tests.rs"] -mod test; diff --git a/crates/tinyagents-graph/src/checkpoint/mod_tests.rs b/crates/tinyagents-graph/src/checkpoint/mod_tests.rs deleted file mode 100644 index 1b0e61051..000000000 --- a/crates/tinyagents-graph/src/checkpoint/mod_tests.rs +++ /dev/null @@ -1,475 +0,0 @@ -//! Unit tests for the in-memory checkpointer: `put`/`get`/`list` roundtrips -//! (including latest-vs-specific lookup and missing threads) and the shared -//! storage guarantee across cheap clones. - -use super::*; -use serde_json::json; -use tinyagents_harness::ids::NodeId; - -fn checkpoint(thread: &str, id: &str, parent: Option<&str>, step: usize) -> Checkpoint { - 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] -async fn put_get_list_roundtrip() { - let cp = InMemoryCheckpointer::::new(); - - cp.put(checkpoint("t1", "c1", None, 1)).await.unwrap(); - cp.put(checkpoint("t1", "c2", Some("c1"), 2)).await.unwrap(); - - // latest - let latest = cp.get("t1", None).await.unwrap().unwrap(); - assert_eq!(latest.checkpoint_id, "c2"); - assert_eq!(latest.state, 2); - - // specific - let first = cp.get("t1", Some("c1")).await.unwrap().unwrap(); - assert_eq!(first.checkpoint_id, "c1"); - - // missing thread - assert!(cp.get("other", None).await.unwrap().is_none()); - - // list - let list = cp.list("t1").await.unwrap(); - assert_eq!(list.len(), 2); - assert_eq!(list[0].checkpoint_id, "c1"); - assert_eq!(list[1].parent_checkpoint_id.as_deref(), Some("c1")); - assert_eq!(list[1].step, 2); -} - -#[test] -fn legacy_checkpoint_json_without_new_fields_still_loads() { - // Back-compat: a checkpoint serialized before `pending_activations` / - // `barrier_arrivals` existed must deserialize, defaulting those fields so - // resume falls back to `next_nodes` and an empty barrier set. - let legacy = json!({ - "thread_id": "t", - "checkpoint_id": "c1", - "run_id": null, - "parent_checkpoint_id": null, - "namespace": [], - "state": 7, - "next_nodes": ["a", "b"], - "completed_tasks": [], - "pending_writes": [], - "interrupts": [], - "metadata": { "source": "loop", "step": 1 } - }); - 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::new( - 1i32, - vec![super::PendingActivation { - node: NodeId::from("w"), - send_arg: Some(std::sync::Arc::new(json!({ "item": 42 }))), - task_id: tinyagents_harness::ids::TaskId::from("1:0:w"), - }], - ) - .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(); - 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")]); -} - -#[tokio::test] -async fn clones_share_storage() { - let cp = InMemoryCheckpointer::::new(); - let cp2 = cp.clone(); - cp.put(checkpoint("t", "c1", None, 1)).await.unwrap(); - assert_eq!(cp2.count("t"), 1); -} - -#[test] -fn checkpoint_source_roundtrips_string_and_display() { - for src in [ - CheckpointSource::Input, - CheckpointSource::Loop, - CheckpointSource::Update, - CheckpointSource::Fork, - ] { - let s = src.to_string(); - assert_eq!(s, src.as_str()); - assert_eq!(CheckpointSource::parse(&s), Some(src)); - // serde wire form matches the Display/string form. - let json = serde_json::to_string(&src).unwrap(); - assert_eq!(json, format!("\"{s}\"")); - } - assert_eq!(CheckpointSource::parse("nope"), None); -} - -#[test] -fn durability_mode_defaults_to_sync() { - assert_eq!(DurabilityMode::default(), DurabilityMode::Sync); -} - -#[tokio::test] -async fn list_metadata_parses_source_enum() { - let cp = InMemoryCheckpointer::::new(); - let mut c = checkpoint("t1", "c1", None, 0); - c.metadata = json!({ "source": "input", "step": 0 }); - cp.put(c).await.unwrap(); - // Unknown/missing source falls back to `loop`. - let mut c2 = checkpoint("t1", "c2", Some("c1"), 1); - c2.metadata = json!({ "step": 1 }); - cp.put(c2).await.unwrap(); - - let list = cp.list("t1").await.unwrap(); - assert_eq!(list[0].source, CheckpointSource::Input); - assert_eq!(list[1].source, CheckpointSource::Loop); -} - -#[tokio::test] -async fn get_tuple_composes_config_and_parent() { - let cp = InMemoryCheckpointer::::new(); - cp.put(checkpoint("t1", "c1", None, 1)).await.unwrap(); - cp.put(checkpoint("t1", "c2", Some("c1"), 2)).await.unwrap(); - - // Latest tuple resolves the concrete id and its parent config. - let tuple = cp - .get_tuple(CheckpointConfig::latest("t1")) - .await - .unwrap() - .unwrap(); - assert_eq!(tuple.config.checkpoint_id.as_deref(), Some("c2")); - assert_eq!(tuple.checkpoint.checkpoint_id, "c2"); - let parent = tuple.parent_config.unwrap(); - assert_eq!(parent.checkpoint_id.as_deref(), Some("c1")); - assert_eq!(parent.thread_id, "t1"); - - // The root checkpoint has no parent config. - let root = cp - .get_tuple(CheckpointConfig { - thread_id: "t1".to_string(), - checkpoint_id: Some("c1".to_string()), - namespace: vec![], - }) - .await - .unwrap() - .unwrap(); - assert!(root.parent_config.is_none()); - - // Missing thread yields no tuple. - assert!( - cp.get_tuple(CheckpointConfig::latest("missing")) - .await - .unwrap() - .is_none() - ); -} - -#[tokio::test] -async fn list_threads_and_delete_thread() { - let cp = InMemoryCheckpointer::::new(); - cp.put(checkpoint("a", "a1", None, 1)).await.unwrap(); - cp.put(checkpoint("b", "b1", None, 1)).await.unwrap(); - - let mut threads = cp.list_threads().await.unwrap(); - threads.sort(); - assert_eq!(threads, vec!["a".to_string(), "b".to_string()]); - - cp.delete_thread("a").await.unwrap(); - assert_eq!(cp.list_threads().await.unwrap(), vec!["b".to_string()]); - assert!(cp.get("a", None).await.unwrap().is_none()); - // Deleting a missing thread is a no-op. - cp.delete_thread("missing").await.unwrap(); -} - -#[tokio::test] -async fn delete_by_run_removes_only_matching_run() { - let cp = InMemoryCheckpointer::::new(); - let mut c1 = checkpoint("t", "c1", None, 1); - c1.run_id = Some("run-1".to_string()); - let mut c2 = checkpoint("t", "c2", Some("c1"), 2); - c2.run_id = Some("run-2".to_string()); - let mut c3 = checkpoint("t", "c3", Some("c2"), 3); - c3.run_id = Some("run-2".to_string()); - cp.put(c1).await.unwrap(); - cp.put(c2).await.unwrap(); - cp.put(c3).await.unwrap(); - - let removed = cp.delete_by_run("t", "run-2").await.unwrap(); - assert_eq!(removed, 2); - let remaining: Vec = cp - .list("t") - .await - .unwrap() - .into_iter() - .map(|m| m.checkpoint_id) - .collect(); - assert_eq!(remaining, vec!["c1".to_string()]); - // Records with no run id are never matched. - assert_eq!(cp.delete_by_run("t", "run-1").await.unwrap(), 1); -} - -#[tokio::test] -async fn copy_thread_preserves_lineage() { - let cp = InMemoryCheckpointer::::new(); - cp.put(checkpoint("src", "c1", None, 1)).await.unwrap(); - cp.put(checkpoint("src", "c2", Some("c1"), 2)) - .await - .unwrap(); - cp.put(checkpoint("src", "c3", Some("c2"), 3)) - .await - .unwrap(); - - cp.copy_thread("src", "dst").await.unwrap(); - - // The source is untouched. - assert_eq!(cp.count("src"), 3); - - // The target carries the same records (ids + parent chain) under a new - // thread id, so time-travel walks the copied thread identically. - let copied = cp.list("dst").await.unwrap(); - assert_eq!(copied.len(), 3); - assert!(copied.iter().all(|m| m.thread_id == "dst")); - assert_eq!(copied[0].checkpoint_id, "c1"); - assert_eq!(copied[0].parent_checkpoint_id, None); - assert_eq!(copied[2].checkpoint_id, "c3"); - assert_eq!(copied[2].parent_checkpoint_id.as_deref(), Some("c2")); - - // The copied checkpoint's state and addressing are intact. - let tip = cp.get("dst", None).await.unwrap().unwrap(); - assert_eq!(tip.thread_id, "dst"); - assert_eq!(tip.state, 3); -} - -#[tokio::test] -async fn get_thread_returns_full_records_in_listing_order() { - let cp = InMemoryCheckpointer::::new(); - cp.put(checkpoint("t", "c1", None, 1)).await.unwrap(); - cp.put(checkpoint("t", "c2", Some("c1"), 2)).await.unwrap(); - cp.put(checkpoint("t", "c3", Some("c2"), 3)).await.unwrap(); - - let records = cp.get_thread("t").await.unwrap(); - assert_eq!(records.len(), 3); - // Full checkpoints (state included), in insertion order. - let ids: Vec<&str> = records.iter().map(|c| c.checkpoint_id.as_str()).collect(); - assert_eq!(ids, vec!["c1", "c2", "c3"]); - assert_eq!(records[2].state, 3); - assert_eq!(records[2].parent_checkpoint_id.as_deref(), Some("c2")); - - // Unknown threads read as empty. - assert!(cp.get_thread("missing").await.unwrap().is_empty()); -} - -#[tokio::test] -async fn prune_keeps_window_and_full_ancestor_chain() { - let cp = InMemoryCheckpointer::::new(); - // Linear lineage c1 <- c2 <- c3 <- c4 <- c5. - cp.put(checkpoint("t", "c1", None, 1)).await.unwrap(); - cp.put(checkpoint("t", "c2", Some("c1"), 2)).await.unwrap(); - cp.put(checkpoint("t", "c3", Some("c2"), 3)).await.unwrap(); - cp.put(checkpoint("t", "c4", Some("c3"), 4)).await.unwrap(); - cp.put(checkpoint("t", "c5", Some("c4"), 5)).await.unwrap(); - - // Keep the last 2 (c4, c5). Their ancestor chain (c3, c2, c1) must be - // retained too — a linear lineage protects everything, deleting nothing. - let removed = cp.prune("t", 2).await.unwrap(); - assert_eq!(removed, 0); - assert_eq!(cp.count("t"), 5); -} - -#[tokio::test] -async fn prune_drops_off_lineage_branches() { - let cp = InMemoryCheckpointer::::new(); - // c1 is the shared root. A dead fork b2 branches off c1 and is never an - // ancestor of the kept tip; the live spine is c1 <- m2 <- m3. - cp.put(checkpoint("t", "c1", None, 1)).await.unwrap(); - cp.put(checkpoint("t", "b2", Some("c1"), 2)).await.unwrap(); - cp.put(checkpoint("t", "m2", Some("c1"), 3)).await.unwrap(); - cp.put(checkpoint("t", "m3", Some("m2"), 4)).await.unwrap(); - - // Keep the last 1 (m3). Protected = {m3} ∪ ancestors {m2, c1}. The dead - // fork b2 is not an ancestor of anything kept, so it is pruned, but the - // ancestor chain a kept delta depends on (m2, c1) survives. - let removed = cp.prune("t", 1).await.unwrap(); - assert_eq!(removed, 1); - let remaining: std::collections::HashSet = cp - .list("t") - .await - .unwrap() - .into_iter() - .map(|m| m.checkpoint_id) - .collect(); - assert_eq!( - remaining, - ["c1", "m2", "m3"].iter().map(|s| s.to_string()).collect() - ); -} - -#[tokio::test] -async fn prune_zero_keeps_latest_and_its_chain() { - let cp = InMemoryCheckpointer::::new(); - cp.put(checkpoint("t", "c1", None, 1)).await.unwrap(); - cp.put(checkpoint("t", "c2", Some("c1"), 2)).await.unwrap(); - - // keep_last == 0 is clamped to 1: the latest checkpoint (and its ancestor - // chain) is always retained so the thread stays resumable. - let removed = cp.prune("t", 0).await.unwrap(); - assert_eq!(removed, 0); - assert_eq!(cp.count("t"), 2); -} - -#[tokio::test] -async fn prune_keeps_a_window_per_namespace() { - let cp = InMemoryCheckpointer::::new(); - // An embedded subgraph writes under the parent's thread but its own - // namespace, interleaved with the parent's records. The two lineages are - // disjoint — no parent record ever references a child id — so a - // thread-wide recency window would delete the child lineage outright and - // leave the thread unresumable. - let sub = |id: &str, parent: Option<&str>, step: usize| { - let mut c = checkpoint("t", id, parent, step); - c.namespace = vec!["sub".to_string()]; - c - }; - cp.put(checkpoint("t", "p1", None, 1)).await.unwrap(); - cp.put(sub("s1", None, 1)).await.unwrap(); - cp.put(sub("s2", Some("s1"), 2)).await.unwrap(); - cp.put(checkpoint("t", "p2", Some("p1"), 2)).await.unwrap(); - - // Keep the last 1 of each namespace plus its ancestors: {p2, p1} and - // {s2, s1} — nothing is deleted, and the subgraph stays resolvable. - let removed = cp.prune("t", 1).await.unwrap(); - assert_eq!(removed, 0); - let child = cp - .get_scoped("t", None, &["sub".to_string()]) - .await - .unwrap() - .expect("the subgraph namespace must stay resumable after prune"); - assert_eq!(child.checkpoint_id, "s2"); -} - -#[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); -} - -/// v1 upgrade alignment: `completed_routes` pairs with `completed_tasks` by -/// position (a shorter list pads with empty routing), and a non-empty -/// `pending_activations` wins over the coarser `next_nodes`. -#[test] -fn v1_records_align_completed_routes_and_prefer_pending_activations() { - use crate::command::RouteTarget; - use tinyagents_harness::ids::TaskId; - - let mut cp = Checkpoint::new(0i32, Vec::new()); - cp.version = 1; - cp.next_nodes = vec![NodeId::from("stale")]; - cp.pending_activations = Some(vec![super::PendingActivation { - node: NodeId::from("w"), - send_arg: None, - task_id: TaskId::from("1:0:w"), - }]); - cp.completed_tasks = vec![NodeId::from("a"), NodeId::from("b")]; - cp.completed_routes = vec![vec![RouteTarget::Node(NodeId::from("x"))]]; - cp.normalize(); - - assert_eq!(cp.version, CHECKPOINT_FORMAT_VERSION); - assert_eq!(cp.tasks.len(), 1); - assert_eq!(cp.tasks[0].node, NodeId::from("w")); - assert_eq!(cp.tasks[0].task_id, TaskId::from("1:0:w")); - let completed: Vec<_> = cp - .completed - .iter() - .map(|task| (task.node.as_str(), task.routes.clone())) - .collect(); - assert_eq!( - completed, - vec![ - ("a", vec![RouteTarget::Node(NodeId::from("x"))]), - ("b", Vec::new()), - ] - ); -} diff --git a/crates/tinyagents-graph/src/checkpoint/types.rs b/crates/tinyagents-graph/src/checkpoint/types.rs deleted file mode 100644 index 6c530ff9f..000000000 --- a/crates/tinyagents-graph/src/checkpoint/types.rs +++ /dev/null @@ -1,960 +0,0 @@ -//! Checkpoint records and metadata — the persisted snapshots that make every -//! level of a recursive graph run resumable and forkable. -//! -//! Checkpoints are graph-runtime persistence, separate from harness memory and -//! long-term stores. They are written at superstep boundaries only — never -//! mid-node — because rerunning a node from its start is far easier to reason -//! about than suspending an async Rust stack, and it matches interrupt/resume -//! semantics exactly. -//! -//! Each record carries a `thread_id` lineage key, a `parent_checkpoint_id` -//! chain (the spine that time-travel and forking walk), and a `namespace` that -//! 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, 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. -/// -/// Mirrors the documented metadata `source` taxonomy: a checkpoint is produced -/// by the initial graph `input`, a normal superstep `loop` boundary, a manual -/// `update` (a state write attributed through the reducers), or a `fork` that -/// branches a thread for time-travel. -#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, serde::Serialize, serde::Deserialize)] -#[serde(rename_all = "lowercase")] -pub enum CheckpointSource { - /// The initial state supplied when a run starts. - Input, - /// A normal superstep boundary in the execution loop. - Loop, - /// A manual state update written through the channel reducers. - Update, - /// A fork that branches a thread for time-travel/replay. - Fork, -} - -impl CheckpointSource { - /// The lowercase wire/string form used in checkpoint metadata. - pub fn as_str(&self) -> &'static str { - match self { - CheckpointSource::Input => "input", - CheckpointSource::Loop => "loop", - CheckpointSource::Update => "update", - CheckpointSource::Fork => "fork", - } - } - - /// Parses a source string, returning `None` for unknown values. - pub fn parse(s: &str) -> Option { - match s { - "input" => Some(CheckpointSource::Input), - "loop" => Some(CheckpointSource::Loop), - "update" => Some(CheckpointSource::Update), - "fork" => Some(CheckpointSource::Fork), - _ => None, - } - } -} - -impl fmt::Display for CheckpointSource { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.write_str(self.as_str()) - } -} - -/// When committed checkpoints are persisted relative to graph execution. -/// -/// The default is [`DurabilityMode::Sync`], which preserves today's behavior. -#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, serde::Serialize, serde::Deserialize)] -#[serde(rename_all = "lowercase")] -pub enum DurabilityMode { - /// Persist a checkpoint before the next step starts. The boundary state is - /// durable before any successor node runs — the strongest guarantee. - #[default] - Sync, - /// Persist off the critical path: the boundary checkpoint write is handed - /// to a spawned background task while the next step executes, so superstep - /// latency does not pay for checkpoint I/O. - /// - /// Failure semantics: a background write error is **not** silently lost. - /// The executor records it and fails the run at the next durability - /// boundary that observes it; at the terminal boundary (and at any - /// interrupt boundary) every in-flight write is awaited, so a completed - /// run's result always reflects its checkpoints' persistence. The final - /// (terminal) and interrupt checkpoints themselves are always written - /// synchronously. The `CheckpointSaved` event for a background write is - /// emitted when the write completes, so its ordering relative to later - /// step events is not deterministic. Outside a tokio runtime this mode - /// degrades to [`DurabilityMode::Sync`]. - Async, - /// Persist only the final checkpoint when the graph exits (or pauses on an - /// interrupt). Intermediate boundaries are not written, trading - /// resumability granularity for fewer writes. - Exit, -} - -/// Coordinates that address a checkpoint within a thread. -/// -/// `checkpoint_id` of `None` selects the latest checkpoint for the thread; -/// `namespace` scopes nested subgraph checkpoints so a parent run and its -/// embedded child graphs never collide. -#[derive(Clone, Debug, Default, PartialEq, Eq, serde::Serialize, serde::Deserialize)] -pub struct CheckpointConfig { - /// Thread lineage key. - pub thread_id: String, - /// Specific checkpoint to address, or `None` for the latest. - pub checkpoint_id: Option, - /// Namespace scoping for nested subgraph checkpoints. - pub namespace: Vec, -} - -impl CheckpointConfig { - /// Builds a config addressing the latest checkpoint of `thread_id` at the - /// root namespace. - pub fn latest(thread_id: impl Into) -> Self { - Self { - thread_id: thread_id.into(), - checkpoint_id: None, - namespace: Vec::new(), - } - } -} - -/// The documented core persistence unit: a checkpoint together with its config, -/// the config of its parent, and the per-task pending writes preserved with it. -/// -/// Backends compose this from `get` + `list` via -/// [`Checkpointer::get_tuple`](crate::Checkpointer::get_tuple). -#[derive(Clone, Debug)] -pub struct CheckpointTuple { - /// Config that addresses this checkpoint. - pub config: CheckpointConfig, - /// The checkpoint record itself. - pub checkpoint: Checkpoint, - /// Config addressing the parent checkpoint, when one exists. - pub parent_config: Option, - /// Pending writes carried by the checkpoint. - pub pending_writes: Vec, -} - -/// A persisted snapshot of a graph run at a superstep boundary. -/// -/// Derives `Serialize`/`Deserialize` with serde's conditional bounds: a -/// `Checkpoint` is (de)serializable exactly when `State` is, which is -/// what lets file-backed backends such as -/// `FileCheckpointer` round-trip whole records -/// 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. - pub checkpoint_id: String, - /// The run that produced this checkpoint, when known. - /// - /// Optional and back-compatible: pre-existing records and manual snapshots - /// may leave it `None`. The executor stamps it so checkpoints can be deleted - /// by run id via [`Checkpointer::delete_by_run`](crate::Checkpointer::delete_by_run). - pub run_id: Option, - /// The previous checkpoint id in the thread lineage. - pub parent_checkpoint_id: Option, - /// Namespace scoping for nested subgraph checkpoints. - pub namespace: Vec, - /// Committed graph state at this boundary. - pub state: State, - /// 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, - /// Barrier (waiting-edge) arrivals accumulated across supersteps, persisted - /// so a join node's precondition survives an interrupt/failure + resume. - /// - /// `#[serde(default)]` for back-compat: older checkpoints load with an - /// 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 -/// resume plus the optional per-invocation [`Send`](crate::Send) -/// argument that scheduled it. -/// -/// The durable counterpart of the executor's in-flight activation. Persisting -/// the `send_arg` is what lets a map-reduce fanout survive an interrupt/failure -/// boundary — without it every pending worker re-runs with no argument. -#[derive(Clone, Debug, PartialEq, serde::Serialize, serde::Deserialize)] -pub struct PendingActivation { - /// The node scheduled to run on resume. - 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>, - /// 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. 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: -/// the predecessors that have already routed to it but whose join has not yet -/// fired. -#[derive(Clone, Debug, PartialEq, serde::Serialize, serde::Deserialize)] -pub struct BarrierArrivals { - /// The waiting/join node. - pub node: NodeId, - /// The predecessor nodes that have arrived so far. - pub arrived: Vec, -} - -impl Checkpoint { - /// The address of this checkpoint: its thread, its id, and its namespace. - pub fn config(&self) -> CheckpointConfig { - CheckpointConfig { - thread_id: self.thread_id.clone(), - checkpoint_id: Some(self.checkpoint_id.clone()), - namespace: self.namespace.clone(), - } - } - - /// The address of the checkpoint this one was written after, if any. - pub fn parent_config(&self) -> Option { - self.parent_checkpoint_id - .as_ref() - .map(|parent| CheckpointConfig { - thread_id: self.thread_id.clone(), - checkpoint_id: Some(parent.clone()), - namespace: self.namespace.clone(), - }) - } - - /// 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 [`crate::checkpoint::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`), 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. 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 - .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 = 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, - has_interrupts: !self.interrupts.is_empty(), - source, - step, - } - } -} - -/// The `idx` reserved for a task's **resume** control-plane write. -/// -/// LangGraph reserves *negative* indices for control-plane channels -/// (`WRITES_IDX_MAP`), which is what distinguishes an upsert from an append: -/// see [`PendingWrite::is_control_plane`]. -pub const WRITES_IDX_RESUME: i64 = -1; - -/// The `idx` reserved for a task's **error** control-plane write. -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 -/// -/// A superstep can fail *after* some of its tasks have already run. The -/// boundary checkpoint records that those tasks completed, but without a record -/// of what they wrote, a resume has no way to tell "this task already ran" from -/// "this task has not run yet" — so it re-runs them, and any side effect they -/// performed happens twice. Writes are therefore persisted per task through -/// [`Checkpointer::put_writes`](crate::Checkpointer::put_writes) and read -/// back into [`CheckpointTuple::pending_writes`], which is what resume consults -/// to skip already-completed work. -/// -/// # Identity -/// -/// A write is addressed by `(thread_id, namespace, checkpoint_id, task_id, -/// idx)`, mirroring the primary key LangGraph's SQL checkpointers use. Within -/// one checkpoint the `(task_id, idx)` pair is unique: re-putting the same pair -/// never produces a second row. -/// -/// # Control-plane writes upsert; data writes are append-once -/// -/// `idx >= 0` is an ordinary data write, emitted once per task in emission -/// order. Re-putting it is **ignored** (insert-or-ignore), so a retried -/// `put_writes` is idempotent. -/// -/// `idx < 0` marks a control-plane write — resume values, errors, interrupts — -/// which by construction there is at most one of per task and whose value -/// legitimately changes on a retry. Re-putting it **replaces** the stored value -/// (insert-or-replace). Use the `WRITES_IDX_*` constants rather than raw -/// negative numbers. -/// -/// # Back-compatibility -/// -/// `task_id`, `idx` and `channel` carry `#[serde(default)]`, so checkpoint -/// records written before the write protocol existed still deserialize (as an -/// anonymous data write at index `0`). -#[derive(Clone, Debug, PartialEq, serde::Serialize, serde::Deserialize)] -pub struct PendingWrite { - /// The node that produced the write. - pub node: NodeId, - /// The task that produced the write, unique within a superstep. - /// - /// 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 = "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)] - pub idx: i64, - /// The channel (state field / node output slot) the write targets. - /// - /// Free-form and backend-opaque; it exists so writes stay distinguishable - /// per channel, and because write isolation is asserted per channel *and* - /// namespace in the conformance suite. - #[serde(default)] - pub channel: String, - /// The serialized write payload. - /// - /// May be [`serde_json::Value::Null`] when the producing runtime cannot - /// serialize its update type. The graph executor is in exactly that - /// position — a graph's `Update` carries no `Serialize` bound — so it - /// records writes as *completion markers*: the applied value is already - /// durable in the checkpoint's `state`, and the write record's job is to - /// answer "did this task already run?". - pub payload: serde_json::Value, -} - -impl PendingWrite { - /// Builds an ordinary data write for `task_id` at position `idx`. - pub fn data( - node: impl Into, - task_id: impl Into, - idx: i64, - channel: impl Into, - payload: serde_json::Value, - ) -> Self { - Self { - node: node.into(), - task_id: task_id.into(), - idx, - channel: channel.into(), - payload, - } - } - - /// 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 { - let node = node.into(); - let channel = node.as_str().to_string(); - Self { - node, - task_id: task_id.into(), - idx: 0, - channel, - payload: serde_json::Value::Null, - } - } - - /// 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) { - (self.task_id.as_str(), self.idx) - } -} - -/// Merges `incoming` into `existing`, applying the replace-vs-ignore rule. -/// -/// Shared by every backend so the three of them cannot drift on the one part of -/// the write protocol that is easy to get subtly wrong: -/// -/// - an incoming **control-plane** write (`idx < 0`) replaces any stored write -/// with the same `(task_id, idx)`; -/// - an incoming **data** write (`idx >= 0`) is ignored when that pair is -/// already stored. -/// -/// Returns the number of entries that were actually inserted or replaced, which -/// backends use for logging. -pub fn merge_writes(existing: &mut Vec, incoming: &[PendingWrite]) -> usize { - let mut changed = 0; - for write in incoming { - match existing - .iter_mut() - .find(|w| w.identity() == write.identity()) - { - Some(slot) => { - if write.is_control_plane() { - *slot = write.clone(); - changed += 1; - } - // Data writes are append-once: a duplicate is a no-op. - } - None => { - existing.push(write.clone()); - changed += 1; - } - } - } - changed -} - -/// Lightweight checkpoint summary returned by `Checkpointer::list`. -/// -/// Listing must not require deserializing full graph state, so metadata is kept -/// separate from the [`Checkpoint`] state payload. -#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)] -pub struct CheckpointMetadata { - /// Thread lineage key. - pub thread_id: String, - /// Checkpoint id. - pub checkpoint_id: String, - /// The run that produced this checkpoint, when known. - pub run_id: Option, - /// Parent checkpoint id. - pub parent_checkpoint_id: Option, - /// Namespace scoping. - pub namespace: Vec, - /// Nodes to run on resume. - pub next_nodes: Vec, - /// Whether the checkpoint carries pending interrupts. - pub has_interrupts: bool, - /// Checkpoint source: `input`, `loop`, `update`, or `fork`. - pub source: CheckpointSource, - /// The superstep number that produced the checkpoint. - pub step: usize, -} diff --git a/crates/tinyagents-graph/src/command/README.md b/crates/tinyagents-graph/src/command/README.md deleted file mode 100644 index 3047d5077..000000000 --- a/crates/tinyagents-graph/src/command/README.md +++ /dev/null @@ -1,64 +0,0 @@ -# graph::command - -The vocabulary a node handler returns to steer the durable executor: partial -state updates, explicit routing, dynamic fanout, and human-in-the-loop -interrupts. - -A node in the durable runtime no longer returns whole state directly. It -returns a [`NodeResult`], and this module is where every shape that result can -take is defined and constructed. `compiled::routing` consumes `Command::goto` -targets when deciding the next superstep's active set; `compiled::executor` -matches on `NodeResult::Interrupt` to pause a run and persist a resumable -checkpoint. - -## Public surface - -- `NodeResult` — what a node handler returns: `Update` (merge through - the reducer), `Command` (update + explicit routing + optional resume value), - or `Interrupt` (pause for human input). -- `Command` — combines an optional partial update with explicit `goto` - routing (one or more targets, overriding static/conditional edges) and an - optional resume value. Constructors: `Command::new`, `::goto`, `::send`, - `::update`, `::resume`; builder methods `with_update`, `with_goto`, - `with_sends`, `with_resume`. -- `RouteTarget` — one routing target: `Node` (plain activation against shared - state) or `Send` (a fanout packet carrying a per-invocation argument). -- `Send` — schedules a node for the next superstep with a custom `arg` - delivered via `NodeContext::send_arg`, independent of the shared committed - state; the map-reduce / search-fanout / per-item-scoring primitive. A single - `Command::send` call may target the same node many times, each with its own - `arg`. -- `Interrupt` — a human-in-the-loop pause point (`id`, `node`, `payload`). - Requires a checkpointer: the executor persists a checkpoint at the - boundary and returns control to the caller; `CompiledGraph::resume` re-runs - the interrupted node with the resume value attached. - -## Files - -| File | Role | -| --- | --- | -| `types.rs` | `NodeResult`, `Command`, `RouteTarget`, `Send`, `Interrupt` definitions. | -| `mod.rs` | Constructors and builder methods on those types; interrupt id generation. | -| `mod_tests.rs` | Unit tests for command/interrupt construction. | - -## How it fits together - -- `builder` nodes are `NodeHandler` closures whose `Result` - future resolves to a `NodeResult` — this module is the return type - of every node in the graph. -- `compiled::routing::route` resolves `Command::goto` targets (when a node - returned a command) ahead of static/conditional edges; `RouteTarget::Send` - entries become distinct activations of the same node, each carrying its own - argument, rather than being deduplicated like plain `Node` targets. -- `compiled::executor` checkpoints and pauses a run on `NodeResult::Interrupt`, - and resumes it by feeding a `Command::resume` value back into the - interrupted node's `NodeContext::resume`. - -## Operational constraints - -- An interrupt with no configured checkpointer (or no thread id) is rejected - before it can pause a run — see `compiled::routing::require_interrupt_durability` - — because an interrupt that cannot be persisted cannot be resumed. -- `Send` targets are never deduplicated by node id the way plain `goto` - targets are; a node that wants to fan out N times to the same target must - emit N distinct `Send` packets, each with its own `arg`. diff --git a/crates/tinyagents-graph/src/command/mod.rs b/crates/tinyagents-graph/src/command/mod.rs deleted file mode 100644 index 5da961db7..000000000 --- a/crates/tinyagents-graph/src/command/mod.rs +++ /dev/null @@ -1,184 +0,0 @@ -//! Command, interrupt, and node-result constructors. -//! -//! Commands are how a node steers the recursive runtime from the inside: a -//! [`Command`] couples a partial state update with explicit `goto` routing -//! (one or many targets — the fanout primitive), so a node — including one that -//! has just run a sub-agent or a subgraph — can dynamically decide which nodes -//! execute next instead of relying only on static edges. [`Interrupt`] pauses a -//! run for human-in-the-loop input, which the durable executor checkpoints and -//! later resumes. -//! -//! See `types` for the definitions. - -mod types; - -pub use types::{Command, Interrupt, NodeResult, RouteTarget, Send}; - -use tinyagents_harness::ids::{NodeId, TaskId}; - -impl Command { - /// Creates an empty command (no update, no routing, no resume). - pub fn new() -> Self { - Self { - update: None, - goto: Vec::new(), - resume: None, - resume_by_task: std::collections::HashMap::new(), - } - } - - /// Creates a command that routes to one or more explicit node targets. - pub fn goto(targets: impl IntoIterator>) -> Self { - Self { - update: None, - goto: targets - .into_iter() - .map(|t| RouteTarget::Node(t.into())) - .collect(), - resume: None, - resume_by_task: std::collections::HashMap::new(), - } - } - - /// Creates a command that fans out to one or more [`Send`] packets, each - /// delivering a custom per-invocation argument to its target node. This is - /// the map-reduce / per-branch-custom-input primitive; targets may repeat. - pub fn send(sends: impl IntoIterator) -> Self { - Self { - update: None, - goto: sends.into_iter().map(RouteTarget::Send).collect(), - resume: None, - resume_by_task: std::collections::HashMap::new(), - } - } - - /// Creates a command carrying a partial state update. - pub fn update(update: Update) -> Self { - Self { - update: Some(update), - goto: Vec::new(), - resume: None, - resume_by_task: std::collections::HashMap::new(), - } - } - - /// Creates a resume command carrying a value for an interrupted node. - pub fn resume(value: serde_json::Value) -> Self { - Self { - 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(), - } - } - - /// Attaches a partial update to this command. - pub fn with_update(mut self, update: Update) -> Self { - self.update = Some(update); - self - } - - /// Appends explicit node routing targets to this command. - pub fn with_goto(mut self, targets: impl IntoIterator>) -> Self { - self.goto - .extend(targets.into_iter().map(|t| RouteTarget::Node(t.into()))); - self - } - - /// Appends [`Send`] fanout packets to this command's routing targets. - pub fn with_sends(mut self, sends: impl IntoIterator) -> Self { - self.goto.extend(sends.into_iter().map(RouteTarget::Send)); - self - } - - /// Attaches a resume value to this command. - pub fn with_resume(mut self, value: serde_json::Value) -> Self { - self.resume = Some(value); - self - } -} - -impl Default for Command { - fn default() -> Self { - Self::new() - } -} - -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 id = format!( - "interrupt-{node}-{}-{}", - tinyagents_harness::ids::process_nonce(), - tinyagents_harness::ids::next_seq() - ); - Self { - id, - node, - payload, - task_id: None, - response_schema: None, - } - } - - /// Creates an interrupt with a caller-supplied id. - pub fn with_id( - id: impl Into, - node: impl Into, - payload: serde_json::Value, - ) -> Self { - Self { - 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)] -#[path = "mod_tests.rs"] -mod test; diff --git a/crates/tinyagents-graph/src/command/mod_tests.rs b/crates/tinyagents-graph/src/command/mod_tests.rs deleted file mode 100644 index 1347a9d09..000000000 --- a/crates/tinyagents-graph/src/command/mod_tests.rs +++ /dev/null @@ -1,94 +0,0 @@ -//! Unit tests for command/interrupt constructors: building commands with -//! updates, `goto` routing, and resume values, plus interrupt id generation. - -use super::*; -use serde_json::json; - -#[test] -fn command_constructors() { - let c: Command = Command::goto(["a", "b"]); - assert_eq!(c.goto.len(), 2); - assert!(c.update.is_none()); - - let c = Command::update(7).with_goto(["next"]); - assert_eq!(c.update, Some(7)); - assert_eq!(c.goto[0].node().as_str(), "next"); - assert!(c.goto[0].send_arg().is_none()); - - let c: Command = Command::resume(json!({ "approved": true })); - assert_eq!(c.resume, Some(json!({ "approved": true }))); -} - -#[test] -fn send_packets_carry_args() { - let c: Command = Command::send([ - Send::new("worker", json!({ "item": 1 })), - Send::new("worker", json!({ "item": 2 })), - ]); - assert_eq!(c.goto.len(), 2); - assert_eq!(c.goto[0].node().as_str(), "worker"); - assert_eq!(c.goto[0].send_arg(), Some(&json!({ "item": 1 }))); - assert_eq!(c.goto[1].send_arg(), Some(&json!({ "item": 2 }))); - - // `with_sends` appends onto an update-bearing command. - let c = Command::update(9).with_sends([Send::new("score", json!("x"))]); - assert_eq!(c.update, Some(9)); - assert_eq!(c.goto.len(), 1); - assert_eq!(c.goto[0].send_arg(), Some(&json!("x"))); -} - -#[test] -fn interrupt_ids_are_unique() { - let a = Interrupt::new("review", json!("a")); - let b = Interrupt::new("review", json!("b")); - assert_ne!(a.id, b.id); - assert_eq!(a.node.as_str(), "review"); - - 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 deleted file mode 100644 index ad2fb28e2..000000000 --- a/crates/tinyagents-graph/src/command/types.rs +++ /dev/null @@ -1,155 +0,0 @@ -//! Partial-update node results, commands, and interrupts. -//! -//! In the durable execution model a node no longer returns whole state. It -//! returns a [`NodeResult`], which is one of: -//! -//! - [`NodeResult::Update`]: a partial update merged through the graph reducer -//! at the superstep boundary. -//! - [`NodeResult::Command`]: a [`Command`] combining an optional update with -//! explicit routing (`goto`) and/or an interrupt resume value. -//! - [`NodeResult::Interrupt`]: an [`Interrupt`] that pauses the run for -//! human-in-the-loop input. - -use tinyagents_harness::ids::{NodeId, TaskId}; - -/// The outcome of running a durable graph node. -#[derive(Clone, Debug)] -pub enum NodeResult { - /// A partial state update to merge through the reducer. - Update(Update), - /// A command that may update state and/or route explicitly. - Command(Command), - /// An interrupt that pauses execution until a resume command arrives. - Interrupt(Interrupt), -} - -/// A dynamic-fanout packet: schedule `node` for the next superstep with a -/// custom `arg` delivered to it through [`crate::NodeContext::send_arg`], -/// independent of the graph's main committed state. -/// -/// `Send` is the primitive for map-reduce, search fanout, parallel tool calls, -/// and per-item scoring: a node can emit one `Send` per work item — even many -/// 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, PartialEq, serde::Serialize, serde::Deserialize)] -pub struct Send { - /// The node to schedule. - pub node: NodeId, - /// The per-invocation input delivered via `NodeContext::send_arg`. - pub arg: serde_json::Value, -} - -impl Send { - /// Creates a `Send` scheduling `node` with `arg`. - pub fn new(node: impl Into, arg: serde_json::Value) -> Self { - Self { - node: node.into(), - arg, - } - } -} - -/// 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`]). -/// -/// 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), - /// Schedule the node with a custom per-invocation argument. - Send(Send), -} - -impl RouteTarget { - /// The destination node id regardless of target kind. - pub fn node(&self) -> &NodeId { - match self { - RouteTarget::Node(node) => node, - RouteTarget::Send(send) => &send.node, - } - } - - /// The per-invocation argument when this is a [`Send`], else `None`. - pub fn send_arg(&self) -> Option<&serde_json::Value> { - match self { - RouteTarget::Node(_) => None, - RouteTarget::Send(send) => Some(&send.arg), - } - } -} - -/// A routing/update/resume directive returned from a node. -/// -/// Commands combine three orthogonal effects: -/// -/// - `update`: a partial state update applied through the reducer. -/// - `goto`: explicit next-step targets — plain node activations and/or [`Send`] -/// fanout packets — overriding static/conditional edges. -/// - `resume`: a value paired with an interrupt resume (set on the caller side). -#[derive(Clone, Debug)] -pub struct Command { - /// Optional partial update applied before routing. - pub update: Option, - /// Explicit routing targets for the next superstep. Each entry is either a - /// 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. -/// -/// Interrupts require a checkpointer. When a node returns an interrupt the -/// executor persists a checkpoint at the boundary and returns control to the -/// caller; `CompiledGraph::resume` re-runs the interrupted node. -#[derive(Clone, Debug, PartialEq, serde::Serialize, serde::Deserialize)] -pub struct Interrupt { - /// Stable id for matching a resume value to this interrupt. - pub id: String, - /// The node that emitted the 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/README.md b/crates/tinyagents-graph/src/compiled/README.md deleted file mode 100644 index d63e0c236..000000000 --- a/crates/tinyagents-graph/src/compiled/README.md +++ /dev/null @@ -1,154 +0,0 @@ -# graph::compiled - -The superstep executor for the durable graph — `CompiledGraph`. - -This is the engine that makes the recursive runtime durable: it drives a -compiled graph in checkpointed supersteps, and because a node handler may -recurse into another compiled graph (a subgraph) or a sub-agent, every level -of that recursion is observed through the same step/boundary/checkpoint -discipline — child runs roll their state, events, and interrupts up through -the parent's reducer and checkpointer. - -## Superstep loop - -Each step: - -1. Take the active node set. -2. Run each active node against the committed state snapshot. -3. Collect updates / commands / interrupts. -4. Apply the reducer at the step boundary. -5. Persist a checkpoint at the boundary (when a checkpointer is configured). - Under `DurabilityMode::Async` non-terminal boundary writes run on spawned - background tasks; a failed write fails the run at the next durability - boundary, and all in-flight writes are drained at the terminal / interrupt / - failure boundary so the run result reflects persistence failures (see - `docs/modules/graph/checkpointing.md`). -6. Select the next active set. - -The loop stops when the active set empties, every branch reaches `END`, an -interrupt pauses the run, or the recursion limit is hit (a deterministic -`TinyAgentsError::RecursionLimit`). - -## Sequential vs. parallel steps - -By default execution is sequential within a step. When the graph is compiled -with `GraphBuilder::with_parallel`, a step with more than one active node runs -every branch concurrently via `futures::future::join_all`, but the data flow -is identical to the sequential case: each branch reads the same committed -snapshot (its own clone), and results fold into the reducer in deterministic -active-set order at the step boundary — the merged state is reproducible -regardless of which branch finishes first. - -Concurrency and interrupt semantics: - -- 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 folded in active-set index order — the reducer is the - fan-in/join, with lower-index branches' updates applied first. -- 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 (below) - and aborts; an interrupt persists a checkpoint whose pending nodes are that - branch and every later active node. -- 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. - -## Network resilience and resumable failures - -Two opt-in mechanisms make a run durable under transient failure and -restartable after a hard one: - -- **Node retry** (`CompiledGraph::with_node_retry`) — a node whose handler - fails with a retryable error (`harness::retry::is_retryable` — the model/tool - transient class) is re-run from its start up to the policy's attempt cap, - emitting `GraphEvent::NodeRetryScheduled` and sleeping the opt-in backoff - between attempts. A single network blip is absorbed without touching the - run. -- **Resumable failure** — when a handler fails beyond the retry budget (or the - error is non-retryable), the executor does not discard the step. On a - checkpointed thread it folds the branches that already completed into - committed state and persists a failure-boundary checkpoint whose - `next_nodes` schedule the failed node (and the not-yet-run tail) for a later - `CompiledGraph::resume` / `CompiledGraph::retry`, with the error and failed - node stamped into checkpoint metadata. The run reports `Failed` (carrying - that checkpoint id) and returns the error. A caller can restart it as-is, or - edit state with `CompiledGraph::update_state` before resuming to continue on - operator feedback. Without a checkpointer the run aborts immediately, as - before. - -## Public surface - -### Construction / configuration (builders, taken by value) - -`with_checkpointer`, `with_event_sink`, `with_durability`, `with_node_retry`, -`with_namespace`, `with_recursion_policy`, `with_recursion_frames`, -`with_recursion_node`, `with_event_journal`, `with_status_store`. - -Accessors: `graph_id()`, `name()`, `namespace()`. - -### Running - -- `run(state)` — fresh, un-checkpointed run (or checkpointed with an - auto-generated thread if a checkpointer is set). -- `run_with_inputs(..)` / `run_with_thread(thread_id, state)` / - `run_with_thread_inputs(..)` — thread-scoped variants for checkpointed runs - and multi-entry-point graphs. -- `resume(..)` / `resume_from(..)` — continue a checkpointed thread from its - last (or a specified) checkpoint, including past an interrupt or a resumable - failure boundary. -- `retry(thread_id)` — re-run the failed node(s) recorded in the last - failure-boundary checkpoint. -- `run_with_agent_binding(..)`, `run_with_thread_agent_binding(..)`, - `resume_with_agent_binding(..)`, `resume_from_with_agent_binding(..)`, and - `retry_with_agent_binding(..)` — execution-scoped variants for graphs that - reach a `SubAgentNode`. The supplied `AgentInvocationBinding` is forwarded - to resumed nodes and nested subgraphs, but is never stored on the reusable - graph or in a checkpoint; an unbound sub-agent continuation fails closed. - -### State inspection / time travel - -- `get_state(..)` / `get_state_history(..)` — read the current or historical - `StateSnapshot` for a thread. -- `update_state(..)` / `bulk_update_state(..)` — operator-driven state edits - between runs (e.g. before a `resume`/`retry`). -- `fork_state(..)` — branch a new thread from an existing checkpoint. - -### Types - -- `CompiledGraph` — the executor itself. -- `GraphExecution` — the result of a run: final/paused state, status, - interrupts, run tree. -- `GraphInput` — seeds a run at a specific node with a payload - (`GraphInput::start` / `::new` / `::node`). -- `StateSnapshot` — a point-in-time view (`values`, `tasks`, - `next_nodes`, `config`, `metadata`, `parent_config`, `pending_interrupts`) - returned by state-inspection calls, projected from a `CheckpointTuple` via - `Checkpoint::to_metadata` so it always matches what - `Checkpointer::list` reports. -- `ResumeTarget` — selects which checkpoint/branch a resume targets. - -## Files - -| File | Role | -| --- | --- | -| `types.rs` | `CompiledGraph`, `GraphExecution`, `GraphInput`, `StateSnapshot`, `ResumeTarget`. | -| `mod.rs` | Module wiring: shared helpers (`Activation`, checkpoint-id/snapshot projection, barrier persistence) and the builder/configuration `impl` (`with_*`, `emit`). | -| `executor.rs` | The superstep loop and run/resume entry points (`run`, `resume`, `retry`, `execute`, `execute_run`, node retry, resumable-failure persistence). | -| `state_api.rs` | State inspection / time travel (`get_state`, `get_state_history`, `update_state`, `bulk_update_state`, `fork_state`). | -| `routing.rs` | Resolving a completed step's active set into the next superstep's activations (`route`, `route_completed`, interrupt-durability preconditions). | -| `mod_tests.rs` | Unit tests (sequential/parallel steps, retry, resumable failure, resume/fork, state history). | - -## Operational constraints - -- Node retry and resumable failure both require a checkpointer to have - observable durability; without one, retries still happen in-process but a - hard failure aborts the run instead of leaving a resumable boundary. -- Parallel-step determinism depends on branches never mutating shared state - outside their own snapshot clone — a node handler that reaches around the - snapshot (e.g. into shared interior-mutable state) breaks the "index order - resolves conflicts" guarantee. -- `id` generation (`next_checkpoint_id`) is collision-free across process - restarts by delegating to `harness::ids::new_checkpoint_id`, not a - process-local counter — this matters for resumed threads restarted in a new - process. diff --git a/crates/tinyagents-graph/src/compiled/boundary.rs b/crates/tinyagents-graph/src/compiled/boundary.rs deleted file mode 100644 index b61231ff6..000000000 --- a/crates/tinyagents-graph/src/compiled/boundary.rs +++ /dev/null @@ -1,1067 +0,0 @@ -//! 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, -} - -/// Which requested stop a boundary is handling. -#[derive(Clone, Copy)] -enum StopKind { - Cancelled, - Drained, -} - -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> { - self.finish_stop_boundary(ctx, active, state, StopKind::Cancelled) - .await - } - - /// 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> { - self.finish_stop_boundary(ctx, active, state, StopKind::Drained) - .await - } - - /// Shared body of the cancellation and drain boundaries: both are - /// requested stops, so a lost Async background write must not turn either - /// into a hard error (the drain and the checkpoint write are best effort). - /// They persist `active` as a resumable checkpoint, record the matching - /// terminal status, emit the matching event, and return `Ok` with no - /// interrupts; only `kind` differs. - async fn finish_stop_boundary( - &self, - ctx: &mut RunCtx<'_, State, Update>, - active: &[Activation], - state: &State, - kind: StopKind, - ) -> Result> { - ctx.disarm_drop_guard(); - let _ = ctx.async_writes.drain().await; - let (marker, terminal) = match kind { - StopKind::Cancelled => ("cancelled", ExecutionStatus::Cancelled), - StopKind::Drained => ("drained", ExecutionStatus::Drained), - }; - let checkpoint_id = self - .persist_stop_checkpoint(ctx, state, active, marker) - .await - .unwrap_or(None); - - let mut status = ctx.base_status(); - status.status = terminal; - 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(match kind { - StopKind::Cancelled => GraphEvent::RunCancelled { - run_id: ctx.run_id.clone(), - }, - StopKind::Drained => 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: matches!(kind, StopKind::Drained), - }) - } - - /// 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 = checkpoint.config(); - 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 = checkpoint.config(); - 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 = checkpoint.config(); - 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 = checkpoint.config(); - - 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_tests.rs b/crates/tinyagents-graph/src/compiled/drain_tests.rs deleted file mode 100644 index 322c45275..000000000 --- a/crates/tinyagents-graph/src/compiled/drain_tests.rs +++ /dev/null @@ -1,267 +0,0 @@ -//! 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_tests.rs b/crates/tinyagents-graph/src/compiled/durable_task_tests.rs deleted file mode 100644 index d7ca81358..000000000 --- a/crates/tinyagents-graph/src/compiled/durable_task_tests.rs +++ /dev/null @@ -1,298 +0,0 @@ -//! `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, 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; -use tinystoragedrivers_testkit::MemoryBackend; - -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_reopening_the_checkpointer() { - let checkpoint_backend = MemoryBackend::new(); - let effects = Arc::new(AtomicUsize::new(0)); - { - let cp: Arc> = Arc::new( - crate::test_support::dir_checkpointer::(&checkpoint_backend), - ); - 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(crate::test_support::dir_checkpointer::( - &checkpoint_backend, - )); - 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); -} diff --git a/crates/tinyagents-graph/src/compiled/durable_tests.rs b/crates/tinyagents-graph/src/compiled/durable_tests.rs deleted file mode 100644 index 9589d8414..000000000 --- a/crates/tinyagents-graph/src/compiled/durable_tests.rs +++ /dev/null @@ -1,123 +0,0 @@ -//! Durable executor tests: interrupt→resume and failure→retry across dropped -//! and reopened driver handles on an explicitly owned testkit backend. -//! Physical persistence is covered by the storage engine's own tests. - -use super::*; -use crate::builder::{GraphBuilder, NodeContext}; -use crate::checkpoint::Checkpointer; -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 tinystoragedrivers_testkit::MemoryBackend; - -/// 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() -} - -// ── DriverCheckpointer over testkit ports ──────────────────────────────── - -#[tokio::test] -async fn driver_backend_interrupt_then_reopen_then_resume() { - let checkpoint_backend = MemoryBackend::new(); - - // Write phase runs the graph to its interrupt point. - { - let cp: Arc> = Arc::new( - crate::test_support::dir_checkpointer::(&checkpoint_backend), - ); - 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); - // Drop the graph and its checkpointer. - } - - // Read/resume phase uses a new driver over the same backend. - let cp2: Arc> = Arc::new(crate::test_support::dir_checkpointer::( - &checkpoint_backend, - )); - 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 driver_backend_failure_then_reopen_then_retry() { - let checkpoint_backend = MemoryBackend::new(); - let attempts = Arc::new(AtomicUsize::new(0)); - - // Write phase: the node fails once and the run aborts, leaving a - // resumable failure checkpoint. - { - let cp: Arc> = Arc::new( - crate::test_support::dir_checkpointer::(&checkpoint_backend), - ); - 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); - // Drop the graph and its checkpointer. - } - - // Read/retry phase uses a new driver over the same backend. The transient - // condition has cleared by the time this attempt runs. - let cp2: Arc> = Arc::new(crate::test_support::dir_checkpointer::( - &checkpoint_backend, - )); - 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 deleted file mode 100644 index 951d59de3..000000000 --- a/crates/tinyagents-graph/src/compiled/executor.rs +++ /dev/null @@ -1,765 +0,0 @@ -//! Public run/resume entry points and the superstep execution engine. -//! -//! Split out of `compiled/mod.rs`; see that module's doc comment for the -//! full executor design (superstep loop, concurrency, and resumable-failure -//! 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::execution_locks::RunLockMap; -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 RunLockMap { - static LOCKS: OnceLock = OnceLock::new(); - LOCKS.get_or_init(RunLockMap::default) -} - -/// 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, - Update: Send + 'static, -{ - /// Runs the graph to completion (or to an interrupt) without a thread. - /// - /// 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(RunSeed::fresh( - state, - vec![Activation::node(self.entry.clone())], - 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 - } - - /// Runs one execution with an explicit host-bound recursive-agent binding. - /// - /// The binding is scoped to this run and descendants spawned from it; it is - /// never retained by this reusable graph value. - pub async fn run_with_agent_binding( - &self, - state: State, - binding: crate::subagent_node::AgentInvocationBinding, - ) -> Result> { - self.execute( - RunSeed::fresh(state, vec![Activation::node(self.entry.clone())], None) - .with_binding(binding), - ) - .await - } - - /// Runs the graph with one or more external inputs in the first superstep. - /// - /// [`GraphInput::start`] targets the graph's compiled entry node, preserving - /// the usual `START -> entry` contract for user input. Additional inputs may - /// target any real node directly, so separate LLM/tool loops can be seeded - /// together. Inputs are not deduplicated: two inputs aimed at the same node - /// produce two separate activations, each with its own - /// [`NodeContext::send_arg`](crate::NodeContext::send_arg). - pub async fn run_with_inputs( - &self, - state: State, - inputs: impl IntoIterator, - ) -> Result> { - let active = self.initial_inputs(inputs)?; - self.execute(RunSeed::fresh(state, active, None)).await - } - - /// Runs the graph under a thread id, persisting checkpoints at every - /// superstep boundary when a checkpointer is configured. - pub async fn run_with_thread( - &self, - thread_id: impl Into, - state: State, - ) -> Result> { - self.execute(RunSeed::fresh( - state, - vec![Activation::node(self.entry.clone())], - Some(thread_id.into()), - )) - .await - } - - /// Runs one threaded execution with an explicit recursive-agent binding. - pub async fn run_with_thread_agent_binding( - &self, - thread_id: impl Into, - state: State, - binding: crate::subagent_node::AgentInvocationBinding, - ) -> Result> { - self.execute( - RunSeed::fresh( - state, - vec![Activation::node(self.entry.clone())], - Some(thread_id.into()), - ) - .with_binding(binding), - ) - .await - } - - /// Runs the graph under a thread id with one or more external inputs in the - /// first superstep, persisting checkpoints at every boundary when a - /// checkpointer is configured. - pub async fn run_with_thread_inputs( - &self, - thread_id: impl Into, - state: State, - inputs: impl IntoIterator, - ) -> Result> { - let active = self.initial_inputs(inputs)?; - self.execute(RunSeed::fresh(state, active, Some(thread_id.into()))) - .await - } - - /// Resumes an interrupted run from its latest checkpoint, re-running the - /// interrupted node(s) with the resume value supplied by `command`. - /// - /// Requires a checkpointer and an existing checkpoint for the thread; - /// otherwise returns [`TinyAgentsError::Resume`]. - pub async fn resume( - &self, - thread_id: impl Into, - command: Command, - ) -> Result> { - self.resume_from(thread_id, ResumeTarget::Latest, command) - .await - } - - /// Resumes an interrupted run with a host-bound recursive-agent binding. - /// - /// Like [`Self::resume`], this reloads the latest checkpoint for `thread_id`. - /// The binding is scoped solely to the resumed execution and is propagated - /// to every resumed node and nested subgraph; it is never retained by this - /// reusable graph value. - pub async fn resume_with_agent_binding( - &self, - thread_id: impl Into, - command: Command, - binding: crate::subagent_node::AgentInvocationBinding, - ) -> Result> { - self.resume_from_with_agent_binding(thread_id, ResumeTarget::Latest, command, binding) - .await - } - - /// Retries a failed run from its latest (failure-boundary) checkpoint, - /// re-running the node that failed and the not-yet-run tail of that step. - /// - /// This is the resume counterpart for the *failure* path (as opposed to a - /// human interrupt): after a node handler aborts a checkpointed run — a - /// transient outage that outlived the node-retry policy, or a hard crash — - /// the run leaves a resumable checkpoint (see - /// [`CompiledGraph::with_node_retry`]). Calling `retry` re-runs exactly what - /// did not complete, carrying no resume value. It is shorthand for - /// [`CompiledGraph::resume`] with an empty [`Command`]. - /// - /// To continue on *user feedback* instead of a bare retry, first inspect the - /// committed state with - /// [`get_state`](CompiledGraph::get_state), edit it with - /// [`update_state`](CompiledGraph::update_state), then call `retry` (or - /// `resume`) — the edited state is what the re-run sees. - pub async fn retry(&self, thread_id: impl Into) -> Result> { - self.resume_from(thread_id, ResumeTarget::Latest, Command::new()) - .await - } - - /// Retries a failed run with a host-bound recursive-agent binding. - /// - /// This is the binding-aware counterpart to [`Self::retry`]. The supplied - /// binding is available only to this retry and any descendants it spawns. - pub async fn retry_with_agent_binding( - &self, - thread_id: impl Into, - binding: crate::subagent_node::AgentInvocationBinding, - ) -> Result> { - self.resume_from_with_agent_binding( - thread_id, - ResumeTarget::Latest, - Command::new(), - binding, - ) - .await - } - - /// Resumes a run from a specific checkpoint (time-travel resume). - /// - /// [`ResumeTarget::Latest`] behaves exactly like [`CompiledGraph::resume`]; - /// [`ResumeTarget::Checkpoint`] replays forward from an older checkpoint's - /// config — re-running its pending nodes (and applying `command`'s resume - /// value to any interrupted node) without mutating the original record. The - /// addressed checkpoint is read-only; the replay appends new boundary - /// checkpoints to the thread rather than rewriting history. - /// - /// Requires a checkpointer and a matching checkpoint with pending nodes; - /// otherwise returns [`TinyAgentsError::Resume`]. - pub async fn resume_from( - &self, - thread_id: impl Into, - target: ResumeTarget, - command: Command, - ) -> Result> { - 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. - /// - /// This is the binding-aware counterpart to [`Self::resume_from`]. It is - /// useful when a durable continuation reaches a - /// [`SubAgentNode`](crate::SubAgentNode) after an interrupt or retry. - /// The binding remains execution-scoped, including for resumed nested - /// subgraphs, and is not stored on [`CompiledGraph`](crate::CompiledGraph). - pub async fn resume_from_with_agent_binding( - &self, - thread_id: impl Into, - target: ResumeTarget, - command: Command, - binding: crate::subagent_node::AgentInvocationBinding, - ) -> Result> { - self.resume_from_inner( - thread_id.into(), - target, - command, - Some(binding), - RunOptions::default(), - ) - .await - } - - fn initial_inputs( - &self, - inputs: impl IntoIterator, - ) -> Result> { - let mut active = Vec::new(); - for input in inputs { - let node = if input.node.as_str() == START { - self.entry.clone() - } else if input.node.as_str() == END { - return Err(TinyAgentsError::Graph( - "graph input cannot target END".to_string(), - )); - } else { - if !self.nodes.contains_key(&input.node) { - return Err(TinyAgentsError::MissingNode(input.node.to_string())); - } - input.node - }; - active.push(Activation { - node, - send_arg: input.payload.map(Arc::new), - task_id: TaskId::from(String::new()), - }); - } - if active.is_empty() { - return Err(TinyAgentsError::Validation( - "run_with_inputs requires at least one input".to_string(), - )); - } - Ok(active) - } - - // ---- State inspection & time travel ------------------------------------ - - /// Returns the configured checkpointer or a [`TinyAgentsError::Checkpoint`] - /// when inspection is attempted on a graph without durability. - pub(super) async fn execute( - &self, - 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. - 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.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`, - /// wrapping any existing sink as the live downstream. Returns a plain clone - /// when no journal is configured. - fn clone_with_journal_sink(&self, run_id: &RunId, thread_id: &Option) -> Self { - let Some(journal) = &self.journal else { - return self.clone(); - }; - let mut sink = crate::observability::JournalGraphSink::new( - journal.clone(), - run_id.clone(), - self.graph_id.clone(), - ) - .with_namespace(self.namespace.clone()) - .with_thread(thread_id.clone()); - if let Some(inner) = &self.event_sink { - sink = sink.with_inner(inner.clone()); - } - let mut this = self.clone(); - this.event_sink = Some(Arc::new(sink)); - this - } - - /// 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, - seed: RunSeed, - ) -> Result> { - let RunSeed { - mut state, - active: initial_active, - thread_id, - resume_map, - barriers: initial_barriers, - parent: initial_parent, - binding, - resume_seed, - options, - .. - } = seed; - - 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 }; - - // Record the run as live before the first superstep is scheduled. - 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() { - // 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; - } - // 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; - } - - 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 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`). - state = match self.apply_updates(state, step_run.updates) { - Ok(state) => state, - Err(err) => return self.fail_and_return(&mut ctx, err).await, - }; - - // 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, - }; - - // 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 - .handle_failure_boundary(&mut ctx, sb, &state, fail) - .await; - } - if !step_run.interrupted.is_empty() { - return self - .handle_interrupt_boundary(&mut ctx, sb, state, step_run.interrupted) - .await; - } - - active = match self.advance(&mut ctx, sb, &state).await { - Ok(next) => next, - Err(err) => return self.fail_and_return(&mut ctx, err).await, - }; - } - - Ok(self.finish_run(&mut ctx, state).await) - } - - /// 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. - /// - /// 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, - runner: &StepRunner<'_, State, Update>, - ctx: &mut RunCtx<'_, State, Update>, - active: &[Activation], - state: &State, - step: usize, - ) -> Result>> { - let Some(token) = ctx.cancellation.clone() else { - return runner.run_step(ctx, active, state, step).await.map(Some); - }; - if token.is_cancelled() { - return Ok(None); - } - tokio::select! { - biased; - _ = token.cancelled() => Ok(None), - result = runner.run_step(ctx, active, state, step) => result.map(Some), - } - } - - /// 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, - 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)); - } - // 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 - ))); - } - } - // 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)?; - } - 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)); - } - } - ctx.emit(GraphEvent::StepStarted { - step: ctx.steps, - active: activation_nodes(active), - }); - Ok(ctx.steps) - } - - /// 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, - 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, - }); - - 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: ctx.steps, - interrupts: Vec::new(), - status, - checkpoint_id: ctx.last_checkpoint.clone(), - drained: false, - } - } -} diff --git a/crates/tinyagents-graph/src/compiled/interrupt_selectors_tests.rs b/crates/tinyagents-graph/src/compiled/interrupt_selectors_tests.rs deleted file mode 100644 index 51aad28d5..000000000 --- a/crates/tinyagents-graph/src/compiled/interrupt_selectors_tests.rs +++ /dev/null @@ -1,505 +0,0 @@ -//! 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 deleted file mode 100644 index 1df50a35b..000000000 --- a/crates/tinyagents-graph/src/compiled/mod.rs +++ /dev/null @@ -1,712 +0,0 @@ -//! Superstep executor for the durable graph. -//! -//! This is the engine that makes the recursive runtime durable: it drives a -//! [`CompiledGraph`] in checkpointed supersteps, and because a node handler may -//! recurse into another compiled graph (a subgraph) or a sub-agent, every level -//! of that recursion is observed through the same step/boundary/checkpoint -//! discipline — child runs roll their state, events, and interrupts up through -//! the parent's reducer and checkpointer. -//! -//! The executor runs in supersteps. Each step: take the active node set, run -//! each active node against the committed state snapshot, collect updates / -//! commands / interrupts, apply the reducer at the step boundary, persist a -//! checkpoint at the boundary (when a checkpointer is configured), then select -//! the next active set. The loop stops when the active set empties, every -//! branch reaches [`END`], an interrupt pauses the run, or the recursion limit -//! is hit (a deterministic [`TinyAgentsError::RecursionLimit`]). -//! -//! By default execution is sequential within a step. When the graph is compiled -//! with [`crate::GraphBuilder::with_parallel`], a step with more than one -//! active node runs every branch concurrently via -//! [`futures::future::join_all`], yet the data flow — snapshot reads, boundary -//! reducer application, boundary checkpointing — is identical: each branch reads -//! the same committed snapshot (its own clone), and results are folded into the -//! reducer in deterministic active-set order at the step boundary, so the merged -//! state is reproducible regardless of which branch finishes first. -//! -//! ## Concurrency and interrupt semantics -//! -//! - 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: 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; 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 -//! -//! Two opt-in mechanisms make a run durable under transient failure and -//! restartable after a hard one: -//! -//! - **Node retry.** With -//! [`CompiledGraph::with_node_retry`], a node whose handler fails with a -//! [retryable][tinyagents_harness::retry::is_retryable] error (a model or tool -//! error — the transient class) is re-run from its start up to the policy's -//! attempt cap, emitting -//! [`GraphEvent::NodeRetryScheduled`] -//! and sleeping the opt-in backoff between attempts. A single network blip is -//! absorbed without touching the run. -//! - **Resumable failure.** When a handler fails beyond the retry budget (or the -//! error is non-retryable), the executor does not discard the step. On a -//! checkpointed thread it folds the branches that already completed into -//! committed state and persists a failure-boundary checkpoint whose -//! `next_nodes` schedule the failed node (and the not-yet-run tail) for a -//! later [`CompiledGraph::resume`]/[`CompiledGraph::retry`], with the error -//! and failed node stamped into the checkpoint metadata. The run then reports -//! `Failed` (carrying that checkpoint id) and returns the error. A caller can -//! restart it as-is, or continue on operator feedback by editing state with -//! [`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, DrainHandle, DrainSignal, GraphExecution, GraphInput, ResumeTarget, RunOptions, - StateSnapshot, -}; - -pub(crate) use types::AsyncCheckpointWrites; - -use std::collections::{BTreeMap, HashMap, HashSet}; -use std::sync::Arc; -use std::time::{Duration, SystemTime}; - -use crate::builder::{ - BarrierRelief, Branch, BuilderNode, END, ForkId, IdleClock, NodeContext, NodeFuture, - NodeHandler, NodeMeta, NodePolicy, START, UpdateCodec, -}; -use crate::checkpoint::{ - BarrierArrivals, Checkpoint, CheckpointConfig, CheckpointTuple, Checkpointer, DurabilityMode, - PendingActivation, -}; -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, GraphEventEnvelope, GraphEventSink}; -use crate::{Result, TinyAgentsError}; -use tinyagents_harness::ids::{ - CheckpointId, ExecutionStatus, GraphId, InterruptId, NodeId, RunId, TaskId, ThreadId, -}; -use tinyagents_harness::retry::is_retryable; - -/// Allocates a fresh checkpoint id (string form) that is collision-free across -/// process restarts. -/// -/// Delegates to [`tinyagents_harness::ids::new_checkpoint_id`]: a resumed thread -/// restarted in a new process must never re-mint a checkpoint id it already -/// used, or the lineage map (`prune`) and time-travel resume corrupt. The bare -/// process-local counter this used to build ids from restarted at `0` every -/// process and did exactly that. -fn next_checkpoint_id() -> String { - tinyagents_harness::ids::new_checkpoint_id() - .as_str() - .to_string() -} - -/// Projects a loaded [`CheckpointTuple`] onto a [`StateSnapshot`] for the -/// state-inspection API. The checkpoint's listing metadata is derived through -/// [`Checkpoint::to_metadata`](crate::Checkpoint::to_metadata) so a -/// snapshot's `metadata` always matches what `Checkpointer::list` reports. -fn snapshot_from_tuple(tuple: CheckpointTuple) -> StateSnapshot { - let CheckpointTuple { - config, - checkpoint, - parent_config, - .. - } = tuple; - let metadata = checkpoint.to_metadata(); - // `checkpoint` was normalized on read, so `.tasks` is the single source - // of truth here regardless of the stored record's original format - // version. - let next_nodes: Vec = checkpoint.tasks.iter().map(|t| t.node.clone()).collect(); - StateSnapshot { - values: checkpoint.state, - tasks: next_nodes.clone(), - next_nodes, - config, - metadata, - parent_config, - pending_interrupts: checkpoint.interrupts, - } -} - -/// 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 { - /// Active-set index of the branch whose handler ultimately failed (after any - /// retries). The executor derives the failed node, the completed lower-index - /// branches (whose successors it schedules) and the pending tail (which it - /// re-runs) from this index against the step's active set — preserving each - /// pending branch's [`Send`] argument. - failed_index: usize, - /// The escalated error. - error: TinyAgentsError, -} - -/// One scheduled activation in the active set: the node plus an optional -/// per-invocation [`Send`] argument delivered via [`NodeContext::send_arg`]. -/// -/// Plain edge/`goto`/conditional activations carry `send_arg == None`; a -/// [`crate::Send`] packet carries `Some(arg)`. Multiple activations may -/// target the same node within a step (map-reduce fanout), so the active set is -/// a `Vec` of these rather than a deduplicated node set. -#[derive(Clone)] -struct Activation { - node: NodeId, - /// `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 { - fn node(node: NodeId) -> Self { - Self { - node, - send_arg: None, - task_id: TaskId::from(String::new()), - } - } -} - -impl From<&Activation> for PendingActivation { - fn from(a: &Activation) -> Self { - PendingActivation { - node: a.node.clone(), - send_arg: a.send_arg.clone(), - task_id: a.task_id.clone(), - } - } -} - -impl From<&PendingActivation> for Activation { - fn from(p: &PendingActivation) -> Self { - Activation { - node: p.node.clone(), - send_arg: p.send_arg.clone(), - task_id: p.task_id.clone(), - } - } -} - -/// Projects the live barrier-arrival map onto its serializable checkpoint form. -fn barriers_to_persisted(map: &HashMap>) -> Vec { - map.iter() - .map(|(node, arrived)| BarrierArrivals { - node: node.clone(), - arrived: arrived.iter().cloned().collect(), - }) - .collect() -} - -/// Rebuilds the live barrier-arrival map from a checkpoint's persisted form. -fn barriers_from_persisted(persisted: &[BarrierArrivals]) -> HashMap> { - persisted - .iter() - .map(|b| (b.node.clone(), b.arrived.iter().cloned().collect())) - .collect() -} - -/// Projects the live per-node visit-count map onto a JSON object for -/// embedding into a checkpoint's `metadata` (see the I3 finding in -/// `docs/runtime-comparison/code-review-graph.md`: without this, per-node -/// visit caps reset on every resume instead of bounding the whole thread's -/// lifetime). `NodeId` is not itself a valid `serde_json` map key type, so -/// this builds the object directly rather than serializing the `HashMap`. -fn node_visits_to_json(node_visits: &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 { - active.iter().map(|a| a.node.clone()).collect() -} - -/// The nodes of *this* graph that paused at an interrupt boundary — the ones a -/// resume value belongs to. -/// -/// Read from the `interrupted_nodes` metadata the interrupt boundary stamps, -/// which is the local activation rather than -/// [`Interrupt::node`](crate::command::Interrupt::node): a subgraph node -/// re-emits its *child's* interrupt, so the recorded interrupt can name a node -/// that does not exist in this graph. Falls back to the recorded interrupts for -/// checkpoints written before that metadata existed — intersected with -/// `pending`, because a legacy checkpoint of a re-emitted child interrupt names -/// a node this graph never schedules. An empty result tells the caller to fan -/// the resume value across the pending set, which is what those checkpoints got -/// before the metadata existed. -fn interrupted_nodes(checkpoint: &Checkpoint, pending: &[Activation]) -> Vec { - let stamped = checkpoint - .metadata - .get("interrupted_nodes") - .and_then(serde_json::Value::as_array) - .map(|nodes| { - nodes - .iter() - .filter_map(serde_json::Value::as_str) - .map(NodeId::from) - .collect::>() - }) - .unwrap_or_default(); - if stamped.is_empty() { - let scheduled: HashSet<&NodeId> = pending.iter().map(|a| &a.node).collect(); - return checkpoint - .interrupts - .iter() - .map(|i| i.node.clone()) - .filter(|node| scheduled.contains(node)) - .collect(); - } - stamped -} - -impl CompiledGraph { - /// Internal constructor used by the builder. - #[allow(clippy::too_many_arguments)] - #[allow(clippy::too_many_arguments)] - pub(crate) fn from_parts( - graph_id: GraphId, - name: Option, - nodes: HashMap>, - edges: HashMap>, - branches: HashMap>, - command_nodes: HashSet, - waiting: HashMap>, - entry: NodeId, - reducer: Arc>, - recursion_limit: usize, - parallel: bool, - max_concurrency: Option, - node_timeout: Option, - node_meta: HashMap, - barrier_reliefs: Vec, - ) -> Self { - Self { - graph_id, - name, - nodes: Arc::new(nodes), - edges: Arc::new(edges), - branches: Arc::new(branches), - command_nodes: Arc::new(command_nodes), - waiting: Arc::new(waiting), - barrier_reliefs: Arc::new(barrier_reliefs), - node_meta: Arc::new(node_meta), - entry, - reducer, - recursion_limit, - recursion_policy: crate::recursion::RecursionPolicy::default(), - recursion_frames: Vec::new(), - recursion_node: None, - checkpointer: None, - event_sink: None, - journal: None, - status_store: None, - namespace: Vec::new(), - parallel, - max_concurrency, - node_timeout, - 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 - } - - /// The optional human-readable graph name, if one was set via - /// [`GraphBuilder::with_name`](crate::GraphBuilder::with_name). - pub fn name(&self) -> Option<&str> { - self.name.as_deref() - } - - /// The checkpoint namespace (empty for top-level graphs). - pub fn namespace(&self) -> &[String] { - &self.namespace - } - - /// Attaches a checkpointer, enabling durability, interrupts, and resume. - pub fn with_checkpointer(mut self, checkpointer: Arc>) -> Self - where - State: Send + Sync + 'static, - { - self.checkpointer = Some(checkpointer); - self - } - - /// Attaches an event sink for low-level streaming/observability. - pub fn with_event_sink(mut self, sink: Arc) -> Self { - self.event_sink = Some(sink); - self - } - - /// Sets the [`DurabilityMode`] that governs when boundary checkpoints are - /// persisted. - /// - /// The default is [`DurabilityMode::Sync`] (persist before the next step). - /// [`DurabilityMode::Async`] hands non-terminal boundary writes to spawned - /// background tasks so checkpoint I/O stays off the superstep critical - /// path; a failed background write fails the run at the next durability - /// boundary, and every in-flight write is awaited at the terminal / - /// interrupt boundary so the run result reflects persistence failures (see - /// [`DurabilityMode::Async`] for the full semantics). - /// [`DurabilityMode::Exit`] persists only the terminal checkpoint (and any - /// interrupt boundary, which is required for resume), skipping - /// intermediate boundaries. - pub fn with_durability(mut self, durability: DurabilityMode) -> Self { - self.durability = durability; - self - } - - /// Sets the per-node [`RetryPolicy`](tinyagents_harness::retry::RetryPolicy) applied around every node handler. - /// - /// Opt-in network resilience for the graph: when a node handler fails with a - /// [retryable][tinyagents_harness::retry::is_retryable] error (a model or tool - /// error — the transient class), the executor re-runs the node from its - /// start up to the policy's attempt cap, emitting a - /// [`GraphEvent::NodeRetryScheduled`](crate::stream::GraphEvent::NodeRetryScheduled) - /// before each retry. Backoff between attempts is slept on only when the - /// policy opts in via - /// [`RetryPolicy::with_backoff_sleep`](tinyagents_harness::retry::RetryPolicy::with_backoff_sleep). - /// - /// Non-retryable errors, and retryable errors once attempts are exhausted, - /// escalate — and, on a checkpointed thread, leave a resumable - /// failure-boundary checkpoint (see [`CompiledGraph::resume`]) so the run can - /// be restarted or continued rather than lost. Without a policy (the - /// default) the first node error aborts the run immediately. - pub fn with_node_retry(mut self, policy: tinyagents_harness::retry::RetryPolicy) -> Self { - self.node_retry = Some(policy); - 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. - /// - /// Unlike wrapping [`run`](Self::run) in an external - /// [`tokio::time::timeout`] — which aborts mid-super-step and cannot leave a - /// clean checkpoint — this stops the run *between* super-steps: when the - /// elapsed run time first reaches `deadline`, the run fails with - /// [`TinyAgentsError::Timeout`] and (on a checkpointed thread) the last - /// committed boundary checkpoint stays intact, so the run can be resumed or - /// inspected rather than lost. - /// - /// The deadline bounds *scheduling*, not a single in-flight node: a - /// long-running node still runs to completion within its super-step (bound - /// it independently with a per-node timeout). `None` (the default) imposes - /// no deadline. - pub fn with_run_deadline(mut self, deadline: std::time::Duration) -> Self { - self.run_deadline = Some(deadline); - self - } - - /// Sets the checkpoint namespace (used by subgraph wrappers). - pub fn with_namespace(mut self, namespace: Vec) -> Self { - self.namespace = namespace; - self - } - - /// Sets the [`RecursionPolicy`] enforced while this graph runs. - /// - /// The policy bounds three independently-tracked recursion dimensions: - /// run-tree depth (`max_depth`), per-node activations within a run - /// (`max_visits_per_node`), and total super-steps per run - /// (`max_total_steps`). The effective per-run step cap is the smaller of - /// the policy's `max_total_steps` and the builder's recursion limit, so - /// configuring a policy never *loosens* an existing limit. - pub fn with_recursion_policy(mut self, policy: RecursionPolicy) -> Self { - self.recursion_policy = policy; - self - } - - /// Seeds the inherited recursion frames of an enclosing run. - /// - /// A subgraph or sub-agent wrapper passes the parent run's frame stack so - /// this run extends the parent's recursion tree (its root frame's `depth` - /// and `parent` continue from the caller) rather than starting a fresh tree - /// at depth zero. Top-level graphs leave this empty. - pub fn with_recursion_frames(mut self, frames: Vec) -> Self { - self.recursion_frames = frames; - self - } - - /// Sets the hosting node id used as this run's root recursion-frame node. - /// - /// A subgraph wrapper sets this to the embedding node id so the child run's - /// frame (and the parent/child [`RunTree`](crate::RunTree)) names the - /// node that ran the embedded graph. Top-level graphs leave this unset. - pub fn with_recursion_node(mut self, node: NodeId) -> Self { - self.recursion_node = Some(node); - self - } - - /// Attaches a durable event journal. Every emitted [`GraphEvent`] is wrapped - /// into a [`crate::observability::GraphObservation`] (stamped with the - /// run's lineage, the graph's checkpoint namespace, and the run id) and - /// appended for offset-addressable replay. Opt-in; default off. - pub fn with_event_journal( - mut self, - journal: Arc, - ) -> Self { - self.journal = Some(journal); - self - } - - /// Attaches a run-status store. The executor writes a compact - /// [`GraphRunStatus`] at every lifecycle boundary (start, terminal, - /// interrupt, failure) so observers can poll run state. Opt-in; default off. - pub fn with_status_store( - mut self, - status_store: Arc, - ) -> Self { - self.status_store = Some(status_store); - self - } - - /// 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)] -#[path = "drain_tests.rs"] -mod drain_test; -#[cfg(test)] -#[path = "durable_task_tests.rs"] -mod durable_task_test; -#[cfg(test)] -#[path = "durable_tests.rs"] -mod durable_test; -#[cfg(test)] -#[path = "interrupt_selectors_tests.rs"] -mod interrupt_selectors_test; -#[cfg(test)] -#[path = "policy_tests.rs"] -mod policy_test; -#[cfg(test)] -#[path = "resume_tests.rs"] -mod resume_test; -#[cfg(test)] -#[path = "mod_tests.rs"] -mod test; diff --git a/crates/tinyagents-graph/src/compiled/mod_tests.rs b/crates/tinyagents-graph/src/compiled/mod_tests.rs deleted file mode 100644 index 124289884..000000000 --- a/crates/tinyagents-graph/src/compiled/mod_tests.rs +++ /dev/null @@ -1,4886 +0,0 @@ -//! Unit tests for the superstep executor: sequential and parallel runs, -//! reducer fan-in ordering, conditional/command routing, checkpoint -//! persistence, interrupt/resume, and recursion-limit enforcement. -// The deprecated `SubAgentNode` keeps its behaviour and its coverage. -#![allow(deprecated)] - -use super::*; -use crate::builder::{GraphBuilder, GraphDefaults, NodeContext, Route}; -use crate::checkpoint::{Checkpointer, InMemoryCheckpointer}; -use crate::command::{Command, Interrupt, NodeResult, Send}; -use crate::reducer::ClosureStateReducer; -use crate::stream::{CollectingSink, GraphEvent}; -use async_trait::async_trait; -use serde_json::json; -use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering as AtomicOrdering}; -use std::sync::{Arc, Mutex}; -use std::time::Duration; -use tinyagents_harness::cancel::CancellationToken; -use tinyagents_harness::events::{AgentEvent, EventSink, RecordingListener}; -use tinyagents_harness::ids::ExecutionStatus; -use tinyagents_harness::retry::RetryPolicy; - -#[derive(Clone, Debug, PartialEq)] -struct Counter { - value: i32, - log: Vec, -} - -/// A host invoker whose recorded request makes a continuation's live binding -/// observable without putting host capabilities in durable graph state. -#[derive(Clone, Default)] -struct BindingRecordingInvoker(Arc>>); - -#[async_trait] -impl crate::AgentInvoker for BindingRecordingInvoker { - async fn invoke( - &self, - request: crate::AgentInvocation, - ) -> crate::Result { - self.0.lock().unwrap().push(request.clone()); - request.events.emit(AgentEvent::StateUpdate); - Ok(crate::SubAgentOutput { - text: request.input.prompt, - ..Default::default() - }) - } -} - -struct BindingFailingInvoker; - -#[async_trait] -impl crate::AgentInvoker for BindingFailingInvoker { - async fn invoke( - &self, - _request: crate::AgentInvocation, - ) -> crate::Result { - Err(TinyAgentsError::Model("continuation failure".to_string())) - } -} - -fn agent_binding(invoker: Arc) -> crate::AgentInvocationBinding { - crate::AgentInvocationBinding::new(invoker, EventSink::new(), CancellationToken::new()) -} - -/// Builds a graph whose nodes return partial `i32` updates merged by a custom -/// reducer that adds to `value` and records a log entry. -fn adding_graph() -> CompiledGraph { - GraphBuilder::::new() - .set_reducer(ClosureStateReducer::new(|mut s: Counter, u: i32| { - s.value += u; - s.log.push(format!("+{u}")); - Ok(s) - })) - .add_node("inc", |_s: Counter, _c: NodeContext| async move { - Ok(NodeResult::Update(1)) - }) - .add_node("double", |s: Counter, _c: NodeContext| async move { - Ok(NodeResult::Update(s.value)) - }) - .set_entry("inc") - .add_edge("inc", "double") - .set_finish("double") - .compile() - .unwrap() -} - -#[tokio::test] -async fn partial_updates_and_reducer() { - let graph = adding_graph(); - let run = graph - .run(Counter { - value: 0, - log: vec![], - }) - .await - .unwrap(); - // inc -> value=1 ; double -> +1 (value snapshot) -> value=2 - assert_eq!(run.state.value, 2); - assert_eq!(run.state.log, vec!["+1", "+1"]); - assert_eq!(run.steps, 2); - assert_eq!( - run.visited - .iter() - .map(|n| n.to_string()) - .collect::>(), - vec!["inc", "double"] - ); - assert_eq!(run.status.status, ExecutionStatus::Completed); - assert!(!run.is_interrupted()); -} - -#[tokio::test] -async fn conditional_routing_selects_branch() { - let graph = GraphBuilder::::overwrite() - .add_node("start", |_s, _c: NodeContext| async move { - Ok(NodeResult::Update(0)) - }) - .add_node("even", |_s, _c: NodeContext| async move { - Ok(NodeResult::Update(100)) - }) - .add_node("odd", |_s, _c: NodeContext| async move { - Ok(NodeResult::Update(200)) - }) - .set_entry("start") - .add_conditional_edges( - "start", - |s: &i32| { - if *s % 2 == 0 { - "even".to_string() - } else { - "odd".to_string() - } - }, - [("even", "even"), ("odd", "odd")], - ) - .set_finish("even") - .set_finish("odd") - .compile() - .unwrap(); - - let run = graph.run(0).await.unwrap(); - 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() - .add_node("router", |_s, _c: NodeContext| async move { - Ok(NodeResult::Command( - Command::update(5).with_goto(["target"]), - )) - }) - .add_node("target", |s, _c: NodeContext| async move { - Ok(NodeResult::Update(s + 1)) - }) - .set_entry("router") - .mark_command_routing("router") - .set_finish("target") - .compile() - .unwrap(); - - let run = graph.run(0).await.unwrap(); - assert_eq!(run.state, 6); -} - -#[tokio::test] -async fn command_goto_rejects_unknown_target_immediately() { - let graph = GraphBuilder::::overwrite() - .add_node("router", |_s, _c: NodeContext| async move { - Ok(NodeResult::Command(Command::goto(["missing"]))) - }) - .set_entry("router") - .mark_command_routing("router") - .compile() - .unwrap(); - - let err = graph.run(0).await.unwrap_err(); - match err { - TinyAgentsError::MissingNode(node) => assert_eq!(node, "missing"), - other => panic!("expected MissingNode, got {other:?}"), - } -} - -#[tokio::test] -async fn command_goto_rejects_start_target() { - let graph = GraphBuilder::::overwrite() - .add_node("router", |_s, _c: NodeContext| async move { - Ok(NodeResult::Command(Command::goto(["__start__"]))) - }) - .set_entry("router") - .mark_command_routing("router") - .compile() - .unwrap(); - - let err = graph.run(0).await.unwrap_err(); - assert!(matches!(err, TinyAgentsError::Graph(_)), "got {err:?}"); - assert!(err.to_string().contains("START"), "{err}"); -} - -#[tokio::test] -async fn invalid_command_goto_is_not_persisted_as_next_node() { - let cp = Arc::new(InMemoryCheckpointer::::new()); - let graph = GraphBuilder::::overwrite() - .add_node("router", |_s, _c: NodeContext| async move { - Ok(NodeResult::Command( - Command::update(1).with_goto(["missing"]), - )) - }) - .set_entry("router") - .mark_command_routing("router") - .compile() - .unwrap() - .with_checkpointer(cp.clone()); - - let err = graph.run_with_thread("bad-goto", 0).await.unwrap_err(); - assert!( - matches!(err, TinyAgentsError::MissingNode(_)), - "got {err:?}" - ); - assert_eq!( - cp.count("bad-goto"), - 0, - "invalid runtime route must fail before boundary checkpoint persistence" - ); -} - -#[tokio::test] -async fn recursion_limit_is_deterministic() { - let graph = GraphBuilder::::overwrite() - .with_recursion_limit(3) - .add_node("loop", |s, _c: NodeContext| async move { - Ok(NodeResult::Update(s + 1)) - }) - .set_entry("loop") - .add_edge("loop", "loop") - .compile() - .unwrap(); - - let err = graph.run(0).await.unwrap_err(); - assert!(matches!(err, TinyAgentsError::RecursionLimit(3))); -} - -#[tokio::test] -async fn superstep_count_matches_path_length() { - let graph = GraphBuilder::::overwrite() - .add_node("a", |s, _c: NodeContext| async move { - Ok(NodeResult::Update(s + 1)) - }) - .add_node("b", |s, _c: NodeContext| async move { - Ok(NodeResult::Update(s + 1)) - }) - .add_node("c", |s, _c: NodeContext| async move { - Ok(NodeResult::Update(s + 1)) - }) - .set_entry("a") - .add_edge("a", "b") - .add_edge("b", "c") - .set_finish("c") - .compile() - .unwrap(); - let run = graph.run(0).await.unwrap(); - assert_eq!(run.steps, 3); - assert_eq!(run.state, 3); -} - -#[tokio::test] -async fn checkpoints_persist_at_boundaries() { - let cp = Arc::new(InMemoryCheckpointer::::new()); - let graph = GraphBuilder::::overwrite() - .add_node("a", |s, _c: NodeContext| async move { - Ok(NodeResult::Update(s + 1)) - }) - .add_node("b", |s, _c: NodeContext| async move { - Ok(NodeResult::Update(s + 1)) - }) - .set_entry("a") - .add_edge("a", "b") - .set_finish("b") - .compile() - .unwrap() - .with_checkpointer(cp.clone()); - - let run = graph.run_with_thread("t1", 0).await.unwrap(); - assert_eq!(run.state, 2); - assert!(run.checkpoint_id.is_some()); - - // one checkpoint per superstep boundary - let list = cp.list("t1").await.unwrap(); - assert_eq!(list.len(), 2); - // lineage is chained - assert!(list[0].parent_checkpoint_id.is_none()); - assert_eq!( - list[1].parent_checkpoint_id.as_deref(), - Some(list[0].checkpoint_id.as_str()) - ); -} - -#[tokio::test] -async fn exit_durability_persists_only_terminal_checkpoint() { - use crate::checkpoint::DurabilityMode; - - let cp = Arc::new(InMemoryCheckpointer::::new()); - let graph = GraphBuilder::::overwrite() - .add_node("a", |s, _c: NodeContext| async move { - Ok(NodeResult::Update(s + 1)) - }) - .add_node("b", |s, _c: NodeContext| async move { - Ok(NodeResult::Update(s + 1)) - }) - .set_entry("a") - .add_edge("a", "b") - .set_finish("b") - .compile() - .unwrap() - .with_checkpointer(cp.clone()) - .with_durability(DurabilityMode::Exit); - - let run = graph.run_with_thread("t1", 0).await.unwrap(); - assert_eq!(run.state, 2); - // Only the terminal boundary is persisted under Exit durability. - assert_eq!(cp.count("t1"), 1); - assert!(run.checkpoint_id.is_some()); - let list = cp.list("t1").await.unwrap(); - assert_eq!(list.len(), 1); - // The single record is the terminal boundary: no pending next nodes. - assert!(list[0].next_nodes.is_empty()); -} - -#[tokio::test] -async fn interrupt_then_resume_reruns_node() { - let cp = Arc::new(InMemoryCheckpointer::::new()); - let graph = GraphBuilder::::overwrite() - .add_node("approve", |s, ctx: NodeContext| async move { - match ctx.resume { - // resumed: apply the approved increment - Some(value) => { - let bump = value.get("bump").and_then(|v| v.as_i64()).unwrap_or(0) as i32; - Ok(NodeResult::Update(s + bump)) - } - // first run: pause for human approval - 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() - .with_checkpointer(cp.clone()); - - // first run pauses - 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); - - // resume re-runs the interrupted node with the resume value - let resumed = graph - .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 resume_emits_restore_not_save_for_the_loaded_checkpoint() { - // Resuming loads a checkpoint; that read must surface as CheckpointRestored, - // never CheckpointSaved (which would inflate persisted-checkpoint counts). - let cp = Arc::new(InMemoryCheckpointer::::new()); - let sink = Arc::new(CollectingSink::new()); - let graph = GraphBuilder::::overwrite() - .add_node("approve", |s, ctx: NodeContext| async move { - match ctx.resume { - Some(_) => Ok(NodeResult::Update(s + 1)), - None => Ok(NodeResult::Interrupt(Interrupt::new("approve", json!({})))), - } - }) - .set_entry("approve") - .set_finish("approve") - .compile() - .unwrap() - .with_checkpointer(cp.clone()) - .with_event_sink(sink.clone()); - - let paused = graph.run_with_thread("t", 0).await.unwrap(); - let loaded = paused - .checkpoint_id - .clone() - .expect("interrupt persisted a checkpoint"); - - // Only inspect events emitted during the resume (the initial run genuinely - // saved the interrupt checkpoint). - let before = sink.events().len(); - graph - .resume("t", Command::resume(json!(null))) - .await - .unwrap(); - let resume_events = sink.events(); - let resume_events = &resume_events[before..]; - - assert!( - resume_events.iter().any(|e| matches!( - e, - GraphEvent::CheckpointRestored { checkpoint_id } if *checkpoint_id == loaded - )), - "resume must emit CheckpointRestored for the loaded checkpoint" - ); - assert!( - !resume_events.iter().any(|e| matches!( - e, - GraphEvent::CheckpointSaved { checkpoint_id, .. } if *checkpoint_id == loaded - )), - "loading a checkpoint on resume must not re-emit it as saved" - ); -} - -#[tokio::test] -async fn resume_preserves_parent_checkpoint_lineage() { - // A run that boundary-checkpoints, interrupts, then resumes to completion - // must keep a single connected lineage: the first post-resume checkpoint - // chains onto the loaded one instead of orphaning the pre-interrupt - // history. Without it, get_state_history stops at the resume point and - // prune deletes the ancestors it should protect. - let cp = Arc::new(InMemoryCheckpointer::::new()); - let graph = GraphBuilder::::overwrite() - .add_node("start", |s, _c: NodeContext| async move { - Ok(NodeResult::Update(s)) - }) - .add_node("approve", |s, ctx: NodeContext| async move { - match ctx.resume { - Some(_) => Ok(NodeResult::Update(s + 1)), - None => Ok(NodeResult::Interrupt(Interrupt::new("approve", json!({})))), - } - }) - .add_node("done", |s, _c: NodeContext| async move { - Ok(NodeResult::Update(s)) - }) - .set_entry("start") - .add_edge("start", "approve") - .add_edge("approve", "done") - .set_finish("done") - .compile() - .unwrap() - .with_checkpointer(cp.clone()); - - let paused = graph.run_with_thread("hitl", 10).await.unwrap(); - assert!(paused.is_interrupted()); - let resumed = graph - .resume("hitl", Command::resume(json!(null))) - .await - .unwrap(); - assert!(!resumed.is_interrupted()); - - // Four boundary checkpoints: start, approve(interrupt), approve(resumed), - // done — all reachable through the parent lineage from the latest. - let history = graph.get_state_history("hitl", None).await.unwrap(); - assert_eq!( - history.len(), - 4, - "full lineage must walk past the resume point; got steps {:?}", - history.iter().map(|s| s.metadata.step).collect::>() - ); - // Connected chain: exactly one root, every parent present. - let ids: std::collections::HashSet<&str> = history - .iter() - .map(|s| s.metadata.checkpoint_id.as_str()) - .collect(); - let roots = history - .iter() - .filter(|s| s.metadata.parent_checkpoint_id.is_none()) - .count(); - assert_eq!(roots, 1, "a connected lineage has exactly one root"); - for s in &history { - if let Some(parent) = &s.metadata.parent_checkpoint_id { - assert!( - ids.contains(parent.as_str()), - "parent `{parent}` must be present in the walked history" - ); - } - } - - // Prune protects the ancestor chain of the retained window: keeping the - // latest still keeps the pre-interrupt checkpoints it depends on. - cp.prune("hitl", 1).await.unwrap(); - assert_eq!( - cp.list("hitl").await.unwrap().len(), - 4, - "prune must retain the full ancestor chain across the resume boundary" - ); -} - -#[tokio::test] -async fn interrupt_without_checkpointer_errors_instead_of_pausing() { - let graph = GraphBuilder::::overwrite() - .add_node("approve", |_s, _ctx: NodeContext| async move { - Ok(NodeResult::Interrupt(Interrupt::new( - "approve", - json!({ "ask": "approve?" }), - ))) - }) - .set_entry("approve") - .set_finish("approve") - .compile() - .unwrap(); - - let err = graph.run_with_thread("hitl", 10).await.unwrap_err(); - assert!(matches!(err, TinyAgentsError::Resume(_)), "got {err:?}"); - assert!(err.to_string().contains("checkpointer"), "{err}"); -} - -#[tokio::test] -async fn interrupt_without_thread_errors_instead_of_pausing() { - let graph = GraphBuilder::::overwrite() - .add_node("approve", |_s, _ctx: NodeContext| async move { - Ok(NodeResult::Interrupt(Interrupt::new( - "approve", - json!({ "ask": "approve?" }), - ))) - }) - .set_entry("approve") - .set_finish("approve") - .compile() - .unwrap() - .with_checkpointer(Arc::new(InMemoryCheckpointer::::new())); - - let err = graph.run(10).await.unwrap_err(); - assert!(matches!(err, TinyAgentsError::Resume(_)), "got {err:?}"); - assert!(err.to_string().contains("thread id"), "{err}"); -} - -#[tokio::test] -async fn resume_without_checkpointer_errors() { - let graph = GraphBuilder::::overwrite() - .add_node("a", |s, _c: NodeContext| async move { - Ok(NodeResult::Update(s)) - }) - .set_entry("a") - .set_finish("a") - .compile() - .unwrap(); - let err = graph - .resume("t", Command::resume(json!(null))) - .await - .unwrap_err(); - assert!(matches!(err, TinyAgentsError::Resume(_))); -} - -#[tokio::test] -async fn events_are_emitted() { - let sink = Arc::new(CollectingSink::new()); - let graph = GraphBuilder::::overwrite() - .add_node("a", |s, _c: NodeContext| async move { - Ok(NodeResult::Update(s + 1)) - }) - .set_entry("a") - .set_finish("a") - .compile() - .unwrap() - .with_event_sink(sink.clone()); - - graph.run(1).await.unwrap(); - let events = sink.events(); - assert!( - events - .iter() - .any(|e| matches!(e, GraphEvent::StepStarted { .. })) - ); - assert!( - events - .iter() - .any(|e| matches!(e, GraphEvent::NodeCompleted { .. })) - ); - assert!( - events - .iter() - .any(|e| matches!(e, GraphEvent::StepCompleted { .. })) - ); -} - -// --- State inspection & time travel ---------------------------------------- - -/// A linear `a -> b -> c` counter graph (each node `+1`) wired to `cp`, used by -/// the inspection/time-travel tests. -fn chain_graph(cp: Arc>) -> CompiledGraph { - GraphBuilder::::overwrite() - .add_node("a", |s, _c: NodeContext| async move { - Ok(NodeResult::Update(s + 1)) - }) - .add_node("b", |s, _c: NodeContext| async move { - Ok(NodeResult::Update(s + 1)) - }) - .add_node("c", |s, _c: NodeContext| async move { - Ok(NodeResult::Update(s + 1)) - }) - .set_entry("a") - .add_edge("a", "b") - .add_edge("b", "c") - .set_finish("c") - .compile() - .unwrap() - .with_checkpointer(cp) -} - -#[tokio::test] -async fn get_state_and_history_walk_the_lineage() { - use crate::CheckpointSource; - - let cp = Arc::new(InMemoryCheckpointer::::new()); - let graph = chain_graph(cp.clone()); - graph.run_with_thread("t", 0).await.unwrap(); - - // Latest snapshot is the terminal boundary: state 3, no pending nodes. - let latest = graph.get_state("t", None).await.unwrap().unwrap(); - assert_eq!(latest.values, 3); - assert!(latest.next_nodes.is_empty()); - assert_eq!(latest.metadata.source, CheckpointSource::Loop); - - // History is newest-first along the parent chain: 3 boundaries. - let history = graph.get_state_history("t", None).await.unwrap(); - assert_eq!( - history.iter().map(|s| s.values).collect::>(), - vec![3, 2, 1] - ); - // The oldest snapshot has no parent; younger ones chain to their parent. - assert!(history.last().unwrap().parent_config.is_none()); - assert_eq!( - history[0].parent_config.as_ref().unwrap().checkpoint_id, - history[1].config.checkpoint_id, - ); - - // limit caps to the most recent snapshots. - let limited = graph.get_state_history("t", Some(2)).await.unwrap(); - assert_eq!(limited.len(), 2); - assert_eq!(limited[0].values, 3); - - // Unknown thread / missing checkpointer behave as documented. - assert!(graph.get_state("missing", None).await.unwrap().is_none()); -} - -#[tokio::test] -async fn update_state_goes_through_the_reducer() { - use crate::CheckpointSource; - - let cp = Arc::new(InMemoryCheckpointer::::new()); - let graph = adding_graph().with_checkpointer(cp.clone()); - graph - .run_with_thread( - "t", - Counter { - value: 0, - log: vec![], - }, - ) - .await - .unwrap(); - - // Manual write: the reducer adds 10 and records a log entry (proving it is - // not a raw overwrite). - let config = graph.update_state("t", 10, None).await.unwrap(); - let snap = graph - .get_state("t", Some(&config.checkpoint_id.unwrap())) - .await - .unwrap() - .unwrap(); - assert_eq!(snap.values.value, 12); - assert_eq!(snap.values.log, vec!["+1", "+1", "+10"]); - assert_eq!(snap.metadata.source, CheckpointSource::Update); - - // Attributing to a missing node is rejected. - let err = graph - .update_state("t", 1, Some("nope".into())) - .await - .unwrap_err(); - assert!(matches!(err, TinyAgentsError::MissingNode(_))); -} - -#[tokio::test] -async fn update_state_as_node_sets_successor_pending_nodes() { - let cp = Arc::new(InMemoryCheckpointer::::new()); - let graph = chain_graph(cp.clone()); - graph.run_with_thread("t", 0).await.unwrap(); - - // Attribute a write to `a`: the new checkpoint's pending nodes become a's - // successor (`b`), so a resume continues from there. - let config = graph.update_state("t", 5, Some("a".into())).await.unwrap(); - let snap = graph - .get_state("t", Some(&config.checkpoint_id.unwrap())) - .await - .unwrap() - .unwrap(); - assert_eq!( - snap.next_nodes - .iter() - .map(|n| n.to_string()) - .collect::>(), - vec!["b".to_string()] - ); -} - -#[tokio::test] -async fn update_state_as_command_node_is_rejected() { - // A command node routes dynamically, so it has no static successors. Using - // it as `as_node` would persist an empty `next_nodes` and silently render - // the thread non-resumable; the write must be rejected instead. - let cp = Arc::new(InMemoryCheckpointer::::new()); - let graph = GraphBuilder::::overwrite() - .add_node("router", |_s, _c: NodeContext| async move { - Ok(NodeResult::Command( - Command::update(5).with_goto(["target"]), - )) - }) - .add_node("target", |s, _c: NodeContext| async move { - Ok(NodeResult::Update(s + 1)) - }) - .set_entry("router") - .mark_command_routing("router") - .set_finish("target") - .compile() - .unwrap() - .with_checkpointer(cp); - graph.run_with_thread("t", 0).await.unwrap(); - - let err = graph - .update_state("t", 1, Some("router".into())) - .await - .unwrap_err(); - assert!(matches!(err, TinyAgentsError::Graph(_)), "got {err:?}"); - assert!(err.to_string().contains("non-resumable"), "{err}"); - - // A plain node is still accepted. - graph - .update_state("t", 1, Some("target".into())) - .await - .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; - - let cp = Arc::new(InMemoryCheckpointer::::new()); - let graph = chain_graph(cp.clone()); - graph.run_with_thread("t", 0).await.unwrap(); - let before = cp.count("t"); - - let last = graph - .bulk_update_state("t", [(10, None), (100, None)]) - .await - .unwrap(); - // Two new update checkpoints were appended. - assert_eq!(cp.count("t"), before + 2); - let snap = graph - .get_state("t", Some(&last.checkpoint_id.unwrap())) - .await - .unwrap() - .unwrap(); - // overwrite reducer: 3 -> 10 -> 100 (last write wins each step). - assert_eq!(snap.values, 100); - assert_eq!(snap.metadata.source, CheckpointSource::Update); - - // Empty bulk is rejected (no resulting config). - let err = graph.bulk_update_state("t", []).await.unwrap_err(); - assert!(matches!(err, TinyAgentsError::Checkpoint(_))); -} - -#[tokio::test] -async fn fork_state_does_not_mutate_source() { - use crate::CheckpointSource; - - let cp = Arc::new(InMemoryCheckpointer::::new()); - let graph = chain_graph(cp.clone()); - graph.run_with_thread("src", 0).await.unwrap(); - let src_before = cp.count("src"); - let src_latest = graph.get_state("src", None).await.unwrap().unwrap(); - - let forked = graph.fork_state("src", None, "dst").await.unwrap(); - // Source thread is untouched: same count, same latest state/source. - assert_eq!(cp.count("src"), src_before); - let src_after = graph.get_state("src", None).await.unwrap().unwrap(); - assert_eq!(src_after.values, src_latest.values); - assert_eq!(src_after.metadata.source, src_latest.metadata.source); - - // Target carries the forked state as a fresh root (no parent), source=fork. - let dst = graph - .get_state("dst", Some(&forked.checkpoint_id.unwrap())) - .await - .unwrap() - .unwrap(); - assert_eq!(dst.values, src_latest.values); - assert_eq!(dst.metadata.source, CheckpointSource::Fork); - assert!(dst.parent_config.is_none()); -} - -#[tokio::test] -async fn resume_from_older_checkpoint_replays_forward() { - let cp = Arc::new(InMemoryCheckpointer::::new()); - let graph = chain_graph(cp.clone()); - graph.run_with_thread("t", 0).await.unwrap(); - - // The first boundary (after `a`) has state 1 and pending node `b`. - let list = cp.list("t").await.unwrap(); - let after_a = &list[0]; - assert_eq!( - after_a - .next_nodes - .iter() - .map(|n| n.to_string()) - .collect::>(), - vec!["b".to_string()] - ); - - // Time-travel resume from that older checkpoint replays b -> c forward. - let replayed = graph - .resume_from( - "t", - ResumeTarget::Checkpoint(after_a.checkpoint_id.clone()), - Command::new(), - ) - .await - .unwrap(); - assert!(!replayed.is_interrupted()); - assert_eq!(replayed.state, 3); - - // Resuming an unknown checkpoint id errors. - let err = graph - .resume_from("t", ResumeTarget::Checkpoint("nope".into()), Command::new()) - .await - .unwrap_err(); - assert!(matches!(err, TinyAgentsError::Resume(_))); -} - -#[tokio::test] -async fn resume_from_with_agent_binding_keeps_host_capabilities_live_only() { - // A bound run pauses before delegation. Its checkpoint must be enough to - // resume graph state, but must never retain the invoker, event sink, or - // cancellation handle that happened to start the run. - let checkpointer = Arc::new(InMemoryCheckpointer::::new()); - let graph = GraphBuilder::::overwrite() - .add_node("gate", |state: String, ctx: NodeContext| async move { - if state == "unbound" || ctx.resume.is_some() { - Ok(NodeResult::Update(state)) - } else { - Ok(NodeResult::Interrupt(Interrupt::new( - "gate", - json!({ "ask": "continue?" }), - ))) - } - }) - .add_node( - "delegate", - crate::subagent_node(crate::SubAgentNode::from_fns( - "researcher", - |state: &String| crate::SubAgentInput::prompt(state.clone()), - |output: crate::SubAgentOutput| output.text, - )), - ) - .set_entry("gate") - .add_edge("gate", "delegate") - .set_finish("delegate") - .compile() - .unwrap() - .with_checkpointer(checkpointer.clone()); - - let initial_invoker = Arc::new(BindingRecordingInvoker::default()); - let paused = graph - .run_with_thread_agent_binding( - "resume-binding", - "question".to_string(), - agent_binding(initial_invoker), - ) - .await - .unwrap(); - assert!(paused.is_interrupted()); - - let checkpoint = checkpointer - .get("resume-binding", None) - .await - .unwrap() - .expect("interrupt is checkpointed"); - let interrupted_checkpoint_id = checkpoint.checkpoint_id.clone(); - let checkpoint_json = serde_json::to_value(checkpoint).unwrap(); - assert!(checkpoint_json.get("agent_binding").is_none()); - assert!( - !checkpoint_json - .to_string() - .contains("AgentInvocationBinding"), - "a durable checkpoint must contain no live invocation capability" - ); - - // Reloading this bound checkpoint without a fresh binding reaches the - // SubAgentNode but must fail closed. Retrying the resulting failure - // boundary without a binding must do the same; neither continuation may - // recover capabilities from the original bound execution. - let unbound_resume = graph - .resume("resume-binding", Command::resume(json!("approved"))) - .await - .unwrap_err(); - assert!(matches!(unbound_resume, TinyAgentsError::Capability(_))); - assert!( - unbound_resume - .to_string() - .contains("sub-agent `researcher`"), - "unbound resume must fail at SubAgentNode: {unbound_resume}" - ); - - let unbound_retry = graph.retry("resume-binding").await.unwrap_err(); - assert!(matches!(unbound_retry, TinyAgentsError::Capability(_))); - assert!( - unbound_retry.to_string().contains("sub-agent `researcher`"), - "unbound retry must fail at SubAgentNode: {unbound_retry}" - ); - - let latest_after_unbound_failures = checkpointer - .get("resume-binding", None) - .await - .unwrap() - .expect("unbound failures are checkpointed") - .checkpoint_id; - assert_ne!( - latest_after_unbound_failures, interrupted_checkpoint_id, - "the selected interrupt checkpoint must no longer be latest" - ); - - let invoker = Arc::new(BindingRecordingInvoker::default()); - let events = EventSink::new(); - let listener = Arc::new(RecordingListener::new()); - events.subscribe(listener.clone()); - let cancellation = CancellationToken::new(); - let resumed = graph - .resume_from_with_agent_binding( - "resume-binding", - ResumeTarget::Checkpoint(interrupted_checkpoint_id), - Command::resume(json!("approved")), - crate::AgentInvocationBinding::new(invoker.clone(), events, cancellation.clone()), - ) - .await - .unwrap(); - - assert_eq!(resumed.state, "question"); - let request = invoker.0.lock().unwrap().pop().expect("delegate ran"); - assert_eq!(request.parent_run_id, resumed.run_id); - assert_eq!(request.root_run_id, resumed.root_run_id); - let request_cancellation = request.cancellation.expect("binding supplies cancellation"); - assert!( - !request_cancellation.is_cancelled(), - "the captured request must initially observe the supplied live token" - ); - cancellation.cancel(); - assert!( - request_cancellation.is_cancelled(), - "cancelling the supplied token after invocation must reach the captured request" - ); - assert_eq!(listener.len(), 1); -} - -#[tokio::test] -async fn retry_with_agent_binding_replaces_failed_run_capabilities() { - let checkpointer = Arc::new(InMemoryCheckpointer::::new()); - let graph = GraphBuilder::::overwrite() - .add_node( - "delegate", - crate::subagent_node(crate::SubAgentNode::from_fns( - "researcher", - |state: &String| crate::SubAgentInput::prompt(state.clone()), - |output: crate::SubAgentOutput| output.text, - )), - ) - .set_entry("delegate") - .set_finish("delegate") - .compile() - .unwrap() - .with_checkpointer(checkpointer); - - let failed = graph - .run_with_thread_agent_binding( - "retry-binding", - "question".to_string(), - agent_binding(Arc::new(BindingFailingInvoker)), - ) - .await - .unwrap_err(); - assert!(matches!(failed, TinyAgentsError::Model(_))); - - // Give every capability a distinct observable behavior: the replacement - // invoker records the retry and its event sink has this listener only. The - // live cancellation handle is checked after the invocation below, proving - // the request received this exact replacement token. - let replacement = Arc::new(BindingRecordingInvoker::default()); - let events = EventSink::new(); - let listener = Arc::new(RecordingListener::new()); - events.subscribe(listener.clone()); - let cancellation = CancellationToken::new(); - let retried = graph - .retry_with_agent_binding( - "retry-binding", - crate::AgentInvocationBinding::new(replacement.clone(), events, cancellation.clone()), - ) - .await - .unwrap(); - - assert_eq!(retried.state, "question"); - let request = replacement - .0 - .lock() - .unwrap() - .pop() - .expect("retry delegated"); - assert_eq!(request.parent_run_id, retried.run_id); - assert_eq!(request.root_run_id, retried.root_run_id); - let request_cancellation = request - .cancellation - .expect("retry binding supplies a cancellation token"); - assert!( - !request_cancellation.is_cancelled(), - "the captured retry request must initially observe a live token" - ); - cancellation.cancel(); - assert!( - request_cancellation.is_cancelled(), - "cancelling the supplied token after invocation must reach the captured retry request" - ); - assert_eq!( - listener.len(), - 1, - "retry must forward the supplied event sink to the replacement invoker" - ); -} - -// --- Parallel (fan-out / fan-in) execution --------------------------------- - -#[derive(Clone, Debug, Default, PartialEq)] -struct Fan { - /// Values contributed by branches, in reducer-application order. - values: Vec, - /// Fork branch indices observed by branches, in reducer-application order. - forks: Vec, - /// Sum a downstream join node computed over the merged `values`. - joined_sum: Option, -} - -#[derive(Clone, Debug)] -enum FanUpdate { - Branch { value: i32, fork: usize }, - Join(i32), -} - -/// Shared instrumentation proving how many branches were in flight at once. -#[derive(Clone)] -struct Concurrency { - inflight: Arc, - max: Arc, -} - -impl Concurrency { - fn new() -> Self { - Self { - inflight: Arc::new(AtomicUsize::new(0)), - max: Arc::new(AtomicUsize::new(0)), - } - } - - fn max_observed(&self) -> usize { - self.max.load(AtomicOrdering::SeqCst) - } - - async fn track(&self, sleep: Duration, value: T) -> T { - let now = self.inflight.fetch_add(1, AtomicOrdering::SeqCst) + 1; - self.max.fetch_max(now, AtomicOrdering::SeqCst); - tokio::time::sleep(sleep).await; - self.inflight.fetch_sub(1, AtomicOrdering::SeqCst); - value - } -} - -/// Builds a fan-out/fan-in graph: `super` routes to three branches that each -/// contribute a value and their fork index; all three converge on `join`, which -/// observes the merged state. `parallel` toggles concurrent branch execution. -/// Branch sleeps are deliberately reversed (a shortest, c longest) so reducer -/// ordering cannot accidentally match completion order. -fn fanout_graph(parallel: bool, conc: Concurrency) -> CompiledGraph { - let (c_a, c_b, c_c) = (conc.clone(), conc.clone(), conc); - GraphBuilder::::new() - .with_parallel(parallel) - .set_reducer(ClosureStateReducer::new(|mut s: Fan, u: FanUpdate| { - match u { - FanUpdate::Branch { value, fork } => { - s.values.push(value); - s.forks.push(fork); - } - FanUpdate::Join(sum) => s.joined_sum = Some(sum), - } - Ok(s) - })) - .add_node("super", |_s: Fan, _c: NodeContext| async move { - Ok(NodeResult::Command( - Command::default().with_goto(["a", "b", "c"]), - )) - }) - .add_node("a", move |_s: Fan, c: NodeContext| { - let conc = c_a.clone(); - let fork = c - .fork - .as_ref() - .map(|f| f.branch_index) - .unwrap_or(usize::MAX); - async move { - Ok(NodeResult::Update( - conc.track( - Duration::from_millis(20), - FanUpdate::Branch { value: 1, fork }, - ) - .await, - )) - } - }) - .add_node("b", move |_s: Fan, c: NodeContext| { - let conc = c_b.clone(); - let fork = c - .fork - .as_ref() - .map(|f| f.branch_index) - .unwrap_or(usize::MAX); - async move { - Ok(NodeResult::Update( - conc.track( - Duration::from_millis(60), - FanUpdate::Branch { value: 2, fork }, - ) - .await, - )) - } - }) - .add_node("c", move |_s: Fan, c: NodeContext| { - let conc = c_c.clone(); - let fork = c - .fork - .as_ref() - .map(|f| f.branch_index) - .unwrap_or(usize::MAX); - async move { - Ok(NodeResult::Update( - conc.track( - Duration::from_millis(100), - FanUpdate::Branch { value: 4, fork }, - ) - .await, - )) - } - }) - .add_node("join", |s: Fan, _c: NodeContext| async move { - Ok(NodeResult::Update(FanUpdate::Join(s.values.iter().sum()))) - }) - .set_entry("super") - .mark_command_routing("super") - .add_edge("a", "join") - .add_edge("b", "join") - .add_edge("c", "join") - .set_finish("join") - .compile() - .unwrap() -} - -#[tokio::test] -async fn parallel_runs_branches_concurrently_and_merges() { - let conc = Concurrency::new(); - let graph = fanout_graph(true, conc.clone()); - let run = graph.run(Fan::default()).await.unwrap(); - - // All three branches ran at the same time. - assert_eq!(conc.max_observed(), 3); - // Reducer merged every branch's contribution. - assert_eq!(run.state.values, vec![1, 2, 4]); - // Fork indices are deterministic active-set positions, not completion order. - assert_eq!(run.state.forks, vec![0, 1, 2]); - // Downstream join observed the merged state. - assert_eq!(run.state.joined_sum, Some(7)); - // super | (a,b,c) | join == 3 supersteps. - assert_eq!(run.steps, 3); -} - -#[tokio::test] -async fn sequential_mode_runs_one_branch_at_a_time() { - let conc = Concurrency::new(); - let graph = fanout_graph(false, conc.clone()); - let run = graph.run(Fan::default()).await.unwrap(); - - // Never more than one branch in flight in sequential mode. - assert_eq!(conc.max_observed(), 1); - // Same deterministic merge as the parallel run. - assert_eq!(run.state.values, vec![1, 2, 4]); - assert_eq!(run.state.joined_sum, Some(7)); - // Sequential branches get no fork identity. - assert_eq!(run.state.forks, vec![usize::MAX, usize::MAX, usize::MAX]); - assert_eq!(run.steps, 3); -} - -#[tokio::test] -async fn parallel_merge_is_reproducible() { - // Run the same parallel fan-out repeatedly; the merged order must be stable - // regardless of which branch's sleep finishes first. - for _ in 0..5 { - let graph = fanout_graph(true, Concurrency::new()); - let run = graph.run(Fan::default()).await.unwrap(); - assert_eq!(run.state.values, vec![1, 2, 4]); - assert_eq!(run.state.forks, vec![0, 1, 2]); - assert_eq!(run.state.joined_sum, Some(7)); - } -} - -#[tokio::test] -async fn recursion_limit_is_deterministic_in_parallel() { - // A self-looping fan-out in parallel mode must still hit the recursion limit - // deterministically at the configured number of supersteps. - let graph = GraphBuilder::::overwrite() - .with_parallel(true) - .with_recursion_limit(3) - .add_node("loop", |s, _c: NodeContext| async move { - Ok(NodeResult::Command( - Command::update(s + 1).with_goto(["loop"]), - )) - }) - .set_entry("loop") - .mark_command_routing("loop") - .compile() - .unwrap(); - - let err = graph.run(0).await.unwrap_err(); - assert!(matches!(err, TinyAgentsError::RecursionLimit(3))); -} - -#[tokio::test] -async fn parallel_interrupt_pauses_at_lowest_index_branch() { - // When a parallel branch interrupts, the step pauses; the interrupted branch - // and every later active node become the checkpoint's pending nodes, while - // lower-index successful branches' updates are still applied. - let cp = Arc::new(InMemoryCheckpointer::::new()); - let graph = GraphBuilder::::new() - .with_parallel(true) - .set_reducer(ClosureStateReducer::new(|mut s: Fan, u: FanUpdate| { - if let FanUpdate::Branch { value, fork } = u { - s.values.push(value); - s.forks.push(fork); - } - Ok(s) - })) - .add_node("super", |_s: Fan, _c: NodeContext| async move { - Ok(NodeResult::Command( - Command::default().with_goto(["a", "b"]), - )) - }) - .add_node("a", |_s: Fan, _c: NodeContext| async move { - Ok(NodeResult::Update(FanUpdate::Branch { value: 1, fork: 0 })) - }) - .add_node("b", |_s: Fan, _c: NodeContext| async move { - Ok(NodeResult::Interrupt(Interrupt::new("b", json!({})))) - }) - .set_entry("super") - .mark_command_routing("super") - .set_finish("a") - .set_finish("b") - .compile() - .unwrap() - .with_checkpointer(cp.clone()); - - let paused = graph.run_with_thread("fan", Fan::default()).await.unwrap(); - assert!(paused.is_interrupted()); - // Lower-index branch `a` committed before the pause. - assert_eq!(paused.state.values, vec![1]); - // The interrupting branch `b` is the head of the pending set. - assert_eq!( - paused.status.active_nodes.first().map(|n| n.to_string()), - Some("b".to_string()) - ); -} - -#[tokio::test] -async fn parallel_interrupt_schedules_completed_branch_successors() { - // Parallel [a, b]: a routes to successor `x` and completes; b interrupts. - // After resume, x (a's successor) must still run — its scheduling used to be - // dropped at the interrupt boundary, so x silently never executed. - let cp = Arc::new(InMemoryCheckpointer::::new()); - let graph = GraphBuilder::::new() - .with_parallel(true) - .set_reducer(ClosureStateReducer::new(|mut s: Counter, u: i32| { - s.value += u; - s.log.push(format!("+{u}")); - Ok(s) - })) - .add_node("super", |_s: Counter, _c: NodeContext| async move { - Ok(NodeResult::Command( - Command::default().with_goto(["a", "b"]), - )) - }) - .add_node("a", |_s: Counter, _c: NodeContext| async move { - Ok(NodeResult::Update(1)) - }) - .add_node("b", |_s: Counter, c: NodeContext| async move { - match c.resume { - Some(_) => Ok(NodeResult::Update(100)), - None => Ok(NodeResult::Interrupt(Interrupt::new("b", json!({})))), - } - }) - .add_node("x", |_s: Counter, _c: NodeContext| async move { - Ok(NodeResult::Update(10)) - }) - .set_entry("super") - .mark_command_routing("super") - .add_edge("a", "x") - .set_finish("b") - .set_finish("x") - .compile() - .unwrap() - .with_checkpointer(cp.clone()); - - let paused = graph - .run_with_thread( - "t", - Counter { - value: 0, - log: vec![], - }, - ) - .await - .unwrap(); - assert!(paused.is_interrupted()); - assert_eq!(paused.state.value, 1, "branch a committed before the pause"); - - let done = graph - .resume("t", Command::resume(json!(null))) - .await - .unwrap(); - assert!( - done.visited.iter().any(|n| n.as_str() == "x"), - "a's successor x must run after resume" - ); - // 1 (a) + 100 (b resume) + 10 (x) — every scheduled branch ran once. - 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 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!("+{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 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!({})))) - } - } - } - }) - .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( - "t-higher-index", - Counter { - value: 0, - log: vec![], - }, - ) - .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"); - - let done = graph - .resume("t-higher-index", Command::resume(json!(null))) - .await - .unwrap(); - 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 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| { - 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", |_s: Counter, c: NodeContext| async move { - match c.resume { - Some(_) => Ok(NodeResult::Update(2)), - None => Ok(NodeResult::Interrupt(Interrupt::new("lo", json!({})))), - } - }) - // `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"]))) - }) - .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("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("hi") - .set_finish("lo") - .set_finish("x") - .set_finish("y") - .compile() - .unwrap() - .with_checkpointer(cp.clone()); - - let paused = graph - .run_with_thread( - "t-carried-goto", - Counter { - value: 0, - log: vec![], - }, - ) - .await - .unwrap(); - assert!(paused.is_interrupted()); - assert_eq!( - paused.state.value, 20, - "hi's update committed before the pause" - ); - - let done = graph - .resume("t-carried-goto", Command::resume(json!(null))) - .await - .unwrap(); - - assert_eq!( - x_calls.load(AtomicOrdering::SeqCst), - 1, - "hi's persisted goto(\"x\") must run exactly once after resume" - ); - 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 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![], - }) - .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!( - !resumed.is_interrupted(), - "the resume must carry `lo` past its interrupt check, not pause it again" - ); - - // 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!( - 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 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("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 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_checkpointer(cp.clone()); - - 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() - .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 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(); - 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" - ); -} - -/// 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 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!("w:{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("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!({})))); - } - Ok(NodeResult::Update(arg)) - }) - .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()); - - // 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"]); -} - -#[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 - .visited - .iter() - .filter(|n| n.as_str() == "worker") - .count(); - assert_eq!(worker_runs, 3); - let mut log = run.state.log.clone(); - log.sort(); - assert_eq!(log, vec!["worker:10", "worker:20", "worker:30"]); -} - -#[tokio::test] -async fn repeated_send_activations_keep_distinct_commands() { - // Regression: two `Send` activations of the *same* node each return a - // distinct `Command::goto`. A node-keyed goto map let the second clobber - // the first, so both branches routed to the survivor's target (and one - // sink was dropped). Keyed per activation, each keeps its own routing. - let graph = GraphBuilder::::new() - .set_reducer(ClosureStateReducer::new(|mut s: Counter, u: i32| { - s.value += u; - s.log.push(format!("n:{u}")); - Ok(s) - })) - .with_parallel(true) - .add_node("dispatch", |_s: Counter, _c: NodeContext| async move { - Ok(NodeResult::Command(Command::send([ - Send::new("worker", json!(1)), - Send::new("worker", json!(2)), - ]))) - }) - // Each worker routes to a different sink based on its own send arg. - .add_node("worker", |_s: Counter, c: NodeContext| async move { - let arg = c.send_arg.expect("worker carries a send arg"); - let target = if arg.as_i64() == Some(1) { - "sink_a" - } else { - "sink_b" - }; - Ok(NodeResult::Command(Command::new().with_goto([target]))) - }) - .add_node("sink_a", |_s: Counter, _c: NodeContext| async move { - Ok(NodeResult::Update(10)) - }) - .add_node("sink_b", |_s: Counter, _c: NodeContext| async move { - Ok(NodeResult::Update(20)) - }) - .mark_command_routing("dispatch") - .mark_command_routing("worker") - .set_entry("dispatch") - .set_finish("sink_a") - .set_finish("sink_b") - .compile() - .unwrap(); - - let run = graph - .run(Counter { - value: 0, - log: vec![], - }) - .await - .unwrap(); - - // Both sinks must have run — one per activation's own goto. - assert!( - run.visited.iter().any(|n| n.as_str() == "sink_a"), - "sink_a (worker arg 1's target) must run" - ); - assert!( - run.visited.iter().any(|n| n.as_str() == "sink_b"), - "sink_b (worker arg 2's target) must run" - ); - assert_eq!(run.state.value, 30, "both sinks contributed (10 + 20)"); -} - -#[tokio::test] -async fn run_with_inputs_seeds_start_and_peer_node() { - let graph = GraphBuilder::::new() - .set_reducer(ClosureStateReducer::new(|mut s: Counter, u: String| { - s.log.push(u); - Ok(s) - })) - .with_parallel(true) - .add_node("user_loop", |_s: Counter, c: NodeContext| async move { - let input = c - .send_arg - .expect("start input should be delivered to entry node") - .as_str() - .expect("user payload is a string") - .to_string(); - Ok(NodeResult::Update(format!("user:{input}"))) - }) - .add_node("tool_loop", |_s: Counter, c: NodeContext| async move { - let tool = c - .send_arg - .expect("tool input should be delivered to peer node") - .get("tool") - .and_then(|v| v.as_str()) - .expect("tool payload names the tool") - .to_string(); - Ok(NodeResult::Update(format!("tool:{tool}"))) - }) - .set_entry("user_loop") - .set_finish("user_loop") - .set_finish("tool_loop") - .compile() - .unwrap(); - - let run = graph - .run_with_inputs( - Counter { - value: 0, - log: vec![], - }, - [ - GraphInput::start(json!("hello")), - GraphInput::new("tool_loop", json!({ "tool": "search" })), - ], - ) - .await - .unwrap(); - - assert_eq!(run.steps, 1); - assert_eq!(run.state.log, vec!["user:hello", "tool:search"]); - assert_eq!( - run.visited.iter().map(|n| n.as_str()).collect::>(), - vec!["user_loop", "tool_loop"] - ); -} - -#[tokio::test] -async fn run_with_inputs_preserves_repeated_inputs_to_same_node() { - let graph = GraphBuilder::::new() - .set_reducer(ClosureStateReducer::new(|mut s: Counter, u: String| { - s.log.push(u); - Ok(s) - })) - .add_node("worker", |_s: Counter, c: NodeContext| async move { - let item = c - .send_arg - .expect("external input should carry an item") - .as_i64() - .expect("item payload is an integer"); - Ok(NodeResult::Update(format!("item:{item}"))) - }) - .set_entry("worker") - .set_finish("worker") - .compile() - .unwrap(); - - let run = graph - .run_with_inputs( - Counter { - value: 0, - log: vec![], - }, - [ - GraphInput::new("worker", json!(1)), - GraphInput::new("worker", json!(2)), - GraphInput::new("worker", json!(3)), - ], - ) - .await - .unwrap(); - - assert_eq!(run.steps, 1); - assert_eq!(run.state.log, vec!["item:1", "item:2", "item:3"]); - assert_eq!( - run.visited - .iter() - .filter(|node| node.as_str() == "worker") - .count(), - 3 - ); -} - -/// A node with normal `goto` (no `send_arg`) gets `None`, while the same node -/// reached via `Send` gets the packet's argument — proving the two coexist. -#[tokio::test] -async fn goto_activation_has_no_send_arg() { - let graph = GraphBuilder::::new() - .set_reducer(ClosureStateReducer::new(|mut s: Counter, u: i32| { - s.value += u; - Ok(s) - })) - .add_node("start", |_s: Counter, _c: NodeContext| async move { - Ok(NodeResult::Command(Command::goto(["sink"]))) - }) - .add_node("sink", |_s: Counter, c: NodeContext| async move { - // Plain goto activation: no per-invocation argument. - assert!(c.send_arg.is_none()); - Ok(NodeResult::Update(1)) - }) - .mark_command_routing("start") - .set_entry("start") - .set_finish("sink") - .compile() - .unwrap(); - let run = graph - .run(Counter { - value: 0, - log: vec![], - }) - .await - .unwrap(); - assert_eq!(run.state.value, 1); -} - -/// A user route enum with `Display` can label conditional edges directly -/// (typed routes), and the [`Route`] newtype is accepted interchangeably. -#[tokio::test] -async fn typed_enum_conditional_route() { - #[derive(Clone, Copy)] - enum Decision { - Approve, - Reject, - } - impl std::fmt::Display for Decision { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - match self { - Decision::Approve => f.write_str("approve"), - Decision::Reject => f.write_str("reject"), - } - } - } - - let graph = GraphBuilder::::overwrite() - .add_node("gate", |s: i32, _c: NodeContext| async move { - Ok(NodeResult::Update(s)) - }) - .add_node("approved", |_s: i32, _c: NodeContext| async move { - Ok(NodeResult::Update(100)) - }) - .add_node("rejected", |_s: i32, _c: NodeContext| async move { - Ok(NodeResult::Update(-1)) - }) - .set_entry("gate") - // Router returns the enum directly (impl ToString); the route table is - // keyed by the enum variant and the `Route` newtype interchangeably. - .add_conditional_edges( - "gate", - |s: &i32| { - if *s > 0 { - Decision::Approve - } else { - Decision::Reject - } - }, - [ - (Route::new(Decision::Approve), "approved"), - (Route::new(Decision::Reject), "rejected"), - ], - ) - .set_finish("approved") - .set_finish("rejected") - .compile() - .unwrap(); - - assert_eq!(graph.run(5).await.unwrap().state, 100); - assert_eq!(graph.run(-3).await.unwrap().state, -1); -} - -/// A barrier/waiting node activates exactly once, only after *all* of its -/// registered predecessors have completed — even when they finish in different -/// supersteps. -#[tokio::test] -async fn waiting_edge_barrier_joins_staggered_predecessors() { - let graph = GraphBuilder::::new() - .set_reducer(ClosureStateReducer::new(|mut s: Counter, u: i32| { - s.value += u; - Ok(s) - })) - .add_node("start", |_s: Counter, _c: NodeContext| async move { - // Fan out to a fast predecessor and a one-hop chain. - Ok(NodeResult::Command(Command::goto(["p1", "inter"]))) - }) - .add_node("p1", |_s: Counter, _c: NodeContext| async move { - Ok(NodeResult::Update(1)) - }) - .add_node("inter", |_s: Counter, _c: NodeContext| async move { - Ok(NodeResult::Update(1)) - }) - .add_node("p2", |_s: Counter, _c: NodeContext| async move { - Ok(NodeResult::Update(1)) - }) - .add_node("join", |_s: Counter, _c: NodeContext| async move { - Ok(NodeResult::Update(10)) - }) - .mark_command_routing("start") - .set_entry("start") - // p1 completes in step 2; p2 only after inter (step 3). The barrier - // holds `join` until both have arrived. - .add_waiting_edge("p1", "join") - .add_edge("inter", "p2") - .add_waiting_edge("p2", "join") - .set_finish("join") - .compile() - .unwrap(); - - let run = graph - .run(Counter { - value: 0, - log: vec![], - }) - .await - .unwrap(); - - // join ran exactly once (not once per predecessor arrival). - let join_runs = run.visited.iter().filter(|n| n.as_str() == "join").count(); - assert_eq!(join_runs, 1); - // p1 + inter + p2 + join = 1 + 1 + 1 + 10. - assert_eq!(run.state.value, 13); - // join is the last node visited, proving it waited for both branches. - assert_eq!(run.visited.last().unwrap().as_str(), "join"); -} - -/// `add_sequence` is sugar for a chain of direct edges. -#[tokio::test] -async fn add_sequence_chains_direct_edges() { - 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("a", |_s: Counter, _c: NodeContext| async move { - Ok(NodeResult::Update(1)) - }) - .add_node("b", |_s: Counter, _c: NodeContext| async move { - Ok(NodeResult::Update(1)) - }) - .add_node("c", |_s: Counter, _c: NodeContext| async move { - Ok(NodeResult::Update(1)) - }) - .set_entry("a") - .add_sequence(["a", "b", "c"]) - .set_finish("c") - .compile() - .unwrap(); - - let run = graph - .run(Counter { - value: 0, - log: vec![], - }) - .await - .unwrap(); - assert_eq!(run.state.value, 3); - assert_eq!( - run.visited - .iter() - .map(|n| n.to_string()) - .collect::>(), - vec!["a", "b", "c"] - ); -} - -/// `with_max_concurrency` (via `set_defaults`) bounds the number of node -/// handlers in flight at once within a parallel superstep. -#[tokio::test(flavor = "multi_thread", worker_threads = 4)] -async fn max_concurrency_bounds_in_flight_branches() { - let in_flight = Arc::new(AtomicUsize::new(0)); - let max_seen = Arc::new(AtomicUsize::new(0)); - - let worker_in_flight = in_flight.clone(); - let worker_max = max_seen.clone(); - - let graph = GraphBuilder::::new() - .set_reducer(ClosureStateReducer::new(|mut s: Counter, u: i32| { - s.value += u; - Ok(s) - })) - .set_defaults(GraphDefaults { - parallel: Some(true), - max_concurrency: Some(2), - ..Default::default() - }) - .add_node("dispatch", |_s: Counter, _c: NodeContext| async move { - Ok(NodeResult::Command(Command::send([ - Send::new("worker", json!(1)), - Send::new("worker", json!(1)), - Send::new("worker", json!(1)), - Send::new("worker", json!(1)), - ]))) - }) - .add_node("worker", move |_s: Counter, _c: NodeContext| { - let in_flight = worker_in_flight.clone(); - let max_seen = worker_max.clone(); - async move { - let now = in_flight.fetch_add(1, AtomicOrdering::SeqCst) + 1; - max_seen.fetch_max(now, AtomicOrdering::SeqCst); - tokio::time::sleep(Duration::from_millis(30)).await; - in_flight.fetch_sub(1, AtomicOrdering::SeqCst); - Ok(NodeResult::Update(1)) - } - }) - .mark_command_routing("dispatch") - .set_entry("dispatch") - .set_finish("worker") - .compile() - .unwrap(); - - let run = graph - .run(Counter { - value: 0, - log: vec![], - }) - .await - .unwrap(); - - // All four workers ran and contributed. - assert_eq!(run.state.value, 4); - // Never more than the configured bound of 2 in flight simultaneously. - assert!( - max_seen.load(AtomicOrdering::SeqCst) <= 2, - "max in-flight {} exceeded bound", - max_seen.load(AtomicOrdering::SeqCst) - ); - // And concurrency actually happened (a chunk of 2 overlapped). - assert_eq!(max_seen.load(AtomicOrdering::SeqCst), 2); -} - -/// With `max_concurrency`, the executor uses a rolling `buffered(limit)` window -/// rather than fixed `join_all` chunks, so a slow branch does not head-of-line -/// block later branches: a new branch starts as soon as any in-flight one -/// finishes. A fixed-chunk executor would run the long branch's chunk to -/// completion before starting the next chunk, so the long branch would overlap -/// at most its single chunk-mate. -#[tokio::test(flavor = "multi_thread", worker_threads = 4)] -async fn max_concurrency_uses_rolling_window_not_chunks() { - // A shared flag marks the long branch as running; short branches count how - // many of them start while the long branch is still in flight. - let long_running = Arc::new(AtomicBool::new(false)); - let overlapped_with_long = Arc::new(AtomicUsize::new(0)); - - let w_long = long_running.clone(); - let w_overlap = overlapped_with_long.clone(); - - let graph = GraphBuilder::::new() - .set_reducer(ClosureStateReducer::new(|mut s: Counter, u: i32| { - s.value += u; - Ok(s) - })) - .set_defaults(GraphDefaults { - parallel: Some(true), - max_concurrency: Some(2), - ..Default::default() - }) - .add_node("dispatch", |_s: Counter, _c: NodeContext| async move { - // One long branch (arg 100) plus three short branches (arg 5). - Ok(NodeResult::Command(Command::send([ - Send::new("worker", json!(100)), - Send::new("worker", json!(5)), - Send::new("worker", json!(5)), - Send::new("worker", json!(5)), - ]))) - }) - .add_node("worker", move |_s: Counter, c: NodeContext| { - let long_running = w_long.clone(); - let overlapped = w_overlap.clone(); - async move { - let ms = c.send_arg.and_then(|v| v.as_u64()).unwrap_or(0); - if ms >= 50 { - // The long branch: flag itself running for its whole life. - long_running.store(true, AtomicOrdering::SeqCst); - tokio::time::sleep(Duration::from_millis(ms)).await; - long_running.store(false, AtomicOrdering::SeqCst); - } else { - // A short branch: did it get to start while the long branch - // was still running? Only possible with a rolling window. - if long_running.load(AtomicOrdering::SeqCst) { - overlapped.fetch_add(1, AtomicOrdering::SeqCst); - } - tokio::time::sleep(Duration::from_millis(ms)).await; - } - Ok(NodeResult::Update(1)) - } - }) - .mark_command_routing("dispatch") - .set_entry("dispatch") - .set_finish("worker") - .compile() - .unwrap(); - - let run = graph - .run(Counter { - value: 0, - log: vec![], - }) - .await - .unwrap(); - - assert_eq!(run.state.value, 4, "all four workers ran"); - // With a rolling window the two short branches that start after the initial - // pair (slot freed as each short one finishes) run while the long branch is - // still going. A fixed-chunk executor would finish the long branch's chunk - // first, so at most one short branch could overlap it. - assert!( - overlapped_with_long.load(AtomicOrdering::SeqCst) >= 2, - "expected the rolling window to overlap the long branch with >=2 short \ - branches, saw {}", - overlapped_with_long.load(AtomicOrdering::SeqCst) - ); -} - -/// A per-node default timeout fails the run with [`TinyAgentsError::Timeout`] -/// when a handler does not resolve in time. -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn node_timeout_fails_slow_handler() { - let graph = GraphBuilder::::overwrite() - .with_node_timeout(Duration::from_millis(20)) - .add_node("slow", |s: i32, _c: NodeContext| async move { - tokio::time::sleep(Duration::from_millis(200)).await; - Ok(NodeResult::Update(s)) - }) - .set_entry("slow") - .set_finish("slow") - .compile() - .unwrap(); - - let err = graph.run(0).await.unwrap_err(); - assert!(matches!(err, TinyAgentsError::Timeout(_))); -} - -// ── Whole-run wall-clock deadline ──────────────────────────────────────────── - -/// A per-run deadline stops the run *between* super-steps once the elapsed run -/// time reaches it, surfacing [`TinyAgentsError::Timeout`]. -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn run_deadline_stops_between_supersteps() { - let graph = GraphBuilder::::overwrite() - .add_node("a", |s: i32, _c: NodeContext| async move { - tokio::time::sleep(Duration::from_millis(40)).await; - Ok(NodeResult::Update(s + 1)) - }) - .add_node("b", |s: i32, _c: NodeContext| async move { - Ok(NodeResult::Update(s + 1)) - }) - .set_entry("a") - .add_edge("a", "b") - .set_finish("b") - .compile() - .unwrap() - .with_run_deadline(Duration::from_millis(20)); - - // The first boundary (elapsed ~0) admits node `a`; the next boundary - // (elapsed ~40ms ≥ 20ms) trips the deadline before `b` ever runs. - let err = graph.run(0).await.unwrap_err(); - assert!(matches!(err, TinyAgentsError::Timeout(_)), "got {err:?}"); -} - -/// A run that finishes within its deadline is unaffected — no false trip. -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn run_deadline_allows_a_run_that_finishes_in_time() { - let graph = GraphBuilder::::overwrite() - .add_node("a", |s: i32, _c: NodeContext| async move { - Ok(NodeResult::Update(s + 1)) - }) - .add_node("b", |s: i32, _c: NodeContext| async move { - Ok(NodeResult::Update(s + 1)) - }) - .set_entry("a") - .add_edge("a", "b") - .set_finish("b") - .compile() - .unwrap() - .with_run_deadline(Duration::from_secs(30)); - - let run = graph.run(0).await.unwrap(); - assert_eq!(run.state, 2); -} - -/// On a checkpointed thread, a deadline trip leaves the last committed boundary -/// checkpoint intact — so the run can be resumed to completion rather than lost -/// (the durability win over an external `tokio::time::timeout`). -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn run_deadline_leaves_last_checkpoint_resumable() { - let cp = Arc::new(InMemoryCheckpointer::::new()); - let topology = || { - GraphBuilder::::overwrite() - .add_node("a", |s: i32, _c: NodeContext| async move { - tokio::time::sleep(Duration::from_millis(40)).await; - Ok(NodeResult::Update(s + 1)) - }) - .add_node("b", |s: i32, _c: NodeContext| async move { - Ok(NodeResult::Update(s + 1)) - }) - .add_node("c", |s: i32, _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() - }; - - // Trips after `a`'s boundary (state=1, next=[b]) but before `b` runs. - let deadlined = topology() - .with_checkpointer(cp.clone()) - .with_run_deadline(Duration::from_millis(20)); - let err = deadlined.run_with_thread("t", 0).await.unwrap_err(); - assert!(matches!(err, TinyAgentsError::Timeout(_)), "got {err:?}"); - - // The boundary checkpoint from the completed super-step survived intact. - let list = cp.list("t").await.unwrap(); - assert!( - !list.is_empty(), - "the pre-deadline boundary checkpoint is intact" - ); - - // Resuming (no deadline) continues from that checkpoint to completion. - let resumed = topology().with_checkpointer(cp.clone()); - let run = resumed - .resume("t", Command::resume(json!(null))) - .await - .unwrap(); - assert_eq!(run.state, 3, "resume ran the remaining super-steps b and c"); -} - -// ── Network resilience: node retry + resumable failures ────────────────────── - -/// A single-node graph whose handler fails (with a retryable model error) the -/// first `fail_times` invocations, then succeeds with `+1`. The shared counter -/// lets a test observe how many attempts were made. -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() -} - -#[tokio::test] -async fn node_retry_recovers_transient_failure() { - let attempts = Arc::new(AtomicUsize::new(0)); - let sink = Arc::new(CollectingSink::new()); - // Fail twice, succeed on the third attempt; the policy allows 1 try + 3 - // retries, so recovery is within budget. - let graph = flaky_graph(2, attempts.clone()) - .with_node_retry(RetryPolicy::default().with_max_attempts(4)) - .with_event_sink(sink.clone()); - - let run = graph.run(10).await.unwrap(); - assert_eq!(run.state, 11); - assert_eq!(run.status.status, ExecutionStatus::Completed); - assert_eq!(attempts.load(AtomicOrdering::SeqCst), 3); - - let retries = sink - .events() - .into_iter() - .filter(|e| matches!(e, GraphEvent::NodeRetryScheduled { .. })) - .count(); - assert_eq!(retries, 2, "one retry per transient failure"); -} - -#[tokio::test] -async fn exhausted_retries_leave_a_resumable_failure_checkpoint() { - let cp = Arc::new(InMemoryCheckpointer::::new()); - let attempts = Arc::new(AtomicUsize::new(0)); - // Fail the first 3 invocations. With 1 try + 1 retry the first run exhausts - // its budget (2 attempts: n=0,1) and aborts, leaving a resumable checkpoint. - let graph = flaky_graph(3, attempts.clone()) - .with_node_retry(RetryPolicy::default().with_max_attempts(2)) - .with_checkpointer(cp.clone()); - - let err = graph.run_with_thread("net", 100).await.unwrap_err(); - assert!(matches!(err, TinyAgentsError::Model(_)), "got {err:?}"); - assert_eq!(attempts.load(AtomicOrdering::SeqCst), 2, "1 try + 1 retry"); - - // The failure boundary is durable: the checkpoint schedules the failed node - // for re-run at the first superstep. - let list = cp.list("net").await.unwrap(); - let last = list.last().expect("a failure checkpoint was persisted"); - assert_eq!(last.next_nodes, vec![NodeId::from("flaky")]); - let snapshot = graph.get_state("net", None).await.unwrap().unwrap(); - assert_eq!( - snapshot.metadata.step, 1, - "failure boundary is at the first superstep" - ); - - // Retry: attempt n=2 fails, retry n=3 (3<3 is false) succeeds — the run - // completes without losing the earlier progress. - let resumed = graph.retry("net").await.unwrap(); - assert_eq!( - resumed.state, 101, - "resume re-runs the failed node to success" - ); - assert_eq!(resumed.status.status, ExecutionStatus::Completed); - assert_eq!(attempts.load(AtomicOrdering::SeqCst), 4); -} - -#[tokio::test] -async fn node_failure_without_retry_policy_is_resumable() { - // No node-retry policy configured: the first failure aborts the run, but a - // checkpointed thread still leaves a resumable failure boundary (the - // "resumable abort" default) rather than losing the run. - let cp = Arc::new(InMemoryCheckpointer::::new()); - let attempts = Arc::new(AtomicUsize::new(0)); - let graph = flaky_graph(1, attempts.clone()).with_checkpointer(cp.clone()); - - let err = graph.run_with_thread("once", 7).await.unwrap_err(); - assert!(matches!(err, TinyAgentsError::Model(_)), "got {err:?}"); - assert_eq!( - attempts.load(AtomicOrdering::SeqCst), - 1, - "no retries attempted" - ); - - // Failed status carries the resumable checkpoint id. - let status = graph.get_state("once", None).await.unwrap().unwrap(); - assert_eq!(status.next_nodes, vec![NodeId::from("flaky")]); - - // The transient condition has cleared; retry completes the run. - let resumed = graph.retry("once").await.unwrap(); - assert_eq!(resumed.state, 8); - assert_eq!(resumed.status.status, ExecutionStatus::Completed); -} - -#[tokio::test] -async fn edit_state_then_retry_uses_the_edited_state() { - // User-feedback continuation: after a failure, the operator edits committed - // state via update_state, then retries; the re-run sees the edited value. - let cp = Arc::new(InMemoryCheckpointer::::new()); - let attempts = Arc::new(AtomicUsize::new(0)); - let graph = flaky_graph(1, attempts.clone()).with_checkpointer(cp.clone()); - - graph.run_with_thread("feedback", 0).await.unwrap_err(); - - // Operator bumps the committed state by +40, inheriting the failure - // boundary's pending nodes (`flaky`), then retries. The node adds +1 to - // whatever state it now sees. - graph.update_state("feedback", 40, None).await.unwrap(); - let resumed = graph.retry("feedback").await.unwrap(); - assert_eq!( - resumed.state, 41, - "retry runs against the edited state (40) + 1" - ); - assert_eq!(resumed.status.status, ExecutionStatus::Completed); -} - -#[tokio::test] -async fn parallel_partial_progress_is_preserved_on_failure() { - // A parallel step where one branch succeeds and a lower/higher-index branch - // fails: the successful branch's update is folded into committed state and - // the failure checkpoint schedules only the failed branch for re-run. - let cp = Arc::new(InMemoryCheckpointer::::new()); - let attempts = Arc::new(AtomicUsize::new(0)); - let graph = GraphBuilder::::new() - .with_parallel(true) - .set_reducer(ClosureStateReducer::new(|s: i32, u: i32| Ok(s + u))) - .add_node("seed", |_s, _c: NodeContext| async move { - Ok(NodeResult::Command(Command::goto(["ok", "flaky"]))) - }) - .add_node("ok", |_s, _c: NodeContext| async move { - Ok(NodeResult::Update(1)) - }) - .add_node("flaky", move |_s, _c: NodeContext| { - let attempts = attempts.clone(); - async move { - // Fail on the first invocation, succeed on resume. - if attempts.fetch_add(1, AtomicOrdering::SeqCst) == 0 { - Err(TinyAgentsError::Model("branch blip".into())) - } else { - Ok(NodeResult::Update(10)) - } - } - }) - .set_entry("seed") - .mark_command_routing("seed") - .set_finish("ok") - .set_finish("flaky") - .compile() - .unwrap() - .with_checkpointer(cp.clone()); - - // First run: seed fans out; the "ok" branch commits +1, "flaky" aborts. - graph.run_with_thread("fanout", 0).await.unwrap_err(); - let snapshot = graph.get_state("fanout", None).await.unwrap().unwrap(); - assert_eq!( - snapshot.values, 1, - "the successful branch's +1 is preserved" - ); - assert!( - snapshot.next_nodes.contains(&NodeId::from("flaky")), - "only the failed branch is scheduled for re-run: {:?}", - snapshot.next_nodes - ); - - // Resume: the flaky branch now succeeds (+10) without re-running "ok". - let resumed = graph.retry("fanout").await.unwrap(); - assert_eq!(resumed.state, 11, "1 (preserved) + 10 (re-run branch)"); - assert_eq!(resumed.status.status, ExecutionStatus::Completed); -} - -// --------------------------------------------------------------------------- -// DurabilityMode::Async -// --------------------------------------------------------------------------- - -/// Delegating checkpointer whose `put` sleeps first, then records completion, -/// so tests can observe whether the executor awaited the write inline. -struct SlowCheckpointer { - inner: Arc>, - delay: Duration, - completed_puts: Arc, -} - -#[async_trait::async_trait] -impl Checkpointer for SlowCheckpointer { - async fn put( - &self, - checkpoint: crate::checkpoint::Checkpoint, - ) -> tinyagents_harness::error::Result { - tokio::time::sleep(self.delay).await; - let id = self.inner.put(checkpoint).await?; - self.completed_puts.fetch_add(1, AtomicOrdering::SeqCst); - Ok(id) - } - - async fn get( - &self, - thread_id: &str, - checkpoint_id: Option<&str>, - ) -> tinyagents_harness::error::Result>> { - self.inner.get(thread_id, checkpoint_id).await - } - - async fn list( - &self, - thread_id: &str, - ) -> tinyagents_harness::error::Result> { - self.inner.list(thread_id).await - } - - async fn list_threads(&self) -> tinyagents_harness::error::Result> { - self.inner.list_threads().await - } - - async fn delete_thread(&self, thread_id: &str) -> tinyagents_harness::error::Result<()> { - self.inner.delete_thread(thread_id).await - } - - async fn delete_checkpoints( - &self, - thread_id: &str, - ids: &[String], - ) -> tinyagents_harness::error::Result { - self.inner.delete_checkpoints(thread_id, ids).await - } -} - -/// Delegating checkpointer that fails every *non-terminal* boundary `put` -/// (records with pending next nodes), simulating a broken store while the -/// terminal write still succeeds. -struct FailNonTerminalCheckpointer { - inner: Arc>, -} - -#[async_trait::async_trait] -impl Checkpointer for FailNonTerminalCheckpointer { - async fn put( - &self, - checkpoint: crate::checkpoint::Checkpoint, - ) -> tinyagents_harness::error::Result { - if !checkpoint.tasks.is_empty() { - return Err(tinyagents_harness::error::TinyAgentsError::Checkpoint( - "injected background write failure".to_string(), - )); - } - self.inner.put(checkpoint).await - } - - async fn get( - &self, - thread_id: &str, - checkpoint_id: Option<&str>, - ) -> tinyagents_harness::error::Result>> { - self.inner.get(thread_id, checkpoint_id).await - } - - async fn list( - &self, - thread_id: &str, - ) -> tinyagents_harness::error::Result> { - self.inner.list(thread_id).await - } - - async fn list_threads(&self) -> tinyagents_harness::error::Result> { - self.inner.list_threads().await - } - - async fn delete_thread(&self, thread_id: &str) -> tinyagents_harness::error::Result<()> { - self.inner.delete_thread(thread_id).await - } - - async fn delete_checkpoints( - &self, - thread_id: &str, - ids: &[String], - ) -> tinyagents_harness::error::Result { - self.inner.delete_checkpoints(thread_id, ids).await - } -} - -fn two_step_graph() -> crate::GraphBuilder { - GraphBuilder::::overwrite() - .add_node("a", |s, _c: NodeContext| async move { - Ok(NodeResult::Update(s + 1)) - }) - .add_node("b", |s, _c: NodeContext| async move { - Ok(NodeResult::Update(s + 1)) - }) - .set_entry("a") - .add_edge("a", "b") - .set_finish("b") -} - -#[tokio::test] -async fn async_durability_persists_every_boundary_with_intact_lineage() { - use crate::checkpoint::DurabilityMode; - - let cp = Arc::new(InMemoryCheckpointer::::new()); - let graph = two_step_graph() - .compile() - .unwrap() - .with_checkpointer(cp.clone()) - .with_durability(DurabilityMode::Async); - - let run = graph.run_with_thread("t-async", 0).await.unwrap(); - assert_eq!(run.state, 2); - // Both boundaries are durable by run end: the run drained its background - // writes before writing the terminal checkpoint. - let list = cp.list("t-async").await.unwrap(); - assert_eq!(list.len(), 2); - // Lineage stays chained even though the first write ran in the background - // (its id was minted before the write was handed off). - assert_eq!( - list[1].parent_checkpoint_id.as_deref(), - Some(list[0].checkpoint_id.as_str()) - ); - assert_eq!( - run.checkpoint_id.as_ref().map(|id| id.as_str()), - Some(list[1].checkpoint_id.as_str()) - ); -} - -#[tokio::test] -async fn async_durability_does_not_await_non_terminal_writes_inline() { - use crate::checkpoint::DurabilityMode; - - let completed_puts = Arc::new(AtomicUsize::new(0)); - let observed_at_b = Arc::new(AtomicUsize::new(usize::MAX)); - let cp = Arc::new(SlowCheckpointer { - inner: Arc::new(InMemoryCheckpointer::new()), - delay: Duration::from_millis(100), - completed_puts: completed_puts.clone(), - }); - - let puts_for_b = completed_puts.clone(); - let seen = observed_at_b.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 puts = puts_for_b.clone(); - let seen = seen.clone(); - async move { - // Record how many checkpoint writes had *completed* when this - // node ran. Under Sync durability the step-1 boundary write - // (100ms) would have finished first; under Async it is still - // in flight. - seen.store(puts.load(AtomicOrdering::SeqCst), AtomicOrdering::SeqCst); - Ok(NodeResult::Update(s + 1)) - } - }) - .set_entry("a") - .add_edge("a", "b") - .set_finish("b") - .compile() - .unwrap() - .with_checkpointer(cp) - .with_durability(DurabilityMode::Async); - - let run = graph.run_with_thread("t-async-slow", 0).await.unwrap(); - assert_eq!(run.state, 2); - assert_eq!( - observed_at_b.load(AtomicOrdering::SeqCst), - 0, - "node b must start while the step-1 boundary write is still in flight" - ); - // ...but by run end every write has been drained and is durable. - assert_eq!(completed_puts.load(AtomicOrdering::SeqCst), 2); -} - -#[tokio::test] -async fn async_durability_surfaces_background_write_failure_in_run_result() { - use crate::checkpoint::DurabilityMode; - - let cp = Arc::new(FailNonTerminalCheckpointer { - inner: Arc::new(InMemoryCheckpointer::new()), - }); - let graph = two_step_graph() - .compile() - .unwrap() - .with_checkpointer(cp) - .with_durability(DurabilityMode::Async); - - // The step-1 boundary write fails in the background; the run must not - // report success — the failure surfaces at the next durability boundary - // or, at the latest, at the terminal drain. - let err = graph - .run_with_thread("t-async-fail", 0) - .await - .expect_err("a lost background checkpoint must fail the run"); - assert!( - err.to_string() - .contains("injected background write failure"), - "unexpected error: {err}" - ); -} - -// --------------------------------------------------------------------------- -// Regression: resume value targeting, barrier-gated manual writes, async drains -// --------------------------------------------------------------------------- - -#[tokio::test] -async fn resume_value_reaches_only_the_interrupted_node() { - // Parallel [a, b]: a routes to successor `x` and completes; b interrupts. - // The interrupt boundary schedules both `x` and `b`, but only `b` actually - // interrupted — `x` has never run, so it must observe `ctx.resume == None` - // on its first activation. Fanning the resume value across the whole - // pending set used to drive `x` down its "already approved" arm. - let cp = Arc::new(InMemoryCheckpointer::::new()); - let graph = GraphBuilder::::new() - .with_parallel(true) - .set_reducer(ClosureStateReducer::new(|mut s: Counter, u: i32| { - s.value += u; - s.log.push(format!("+{u}")); - Ok(s) - })) - .add_node("super", |_s: Counter, _c: NodeContext| async move { - Ok(NodeResult::Command( - Command::default().with_goto(["a", "b"]), - )) - }) - .add_node("a", |_s: Counter, _c: NodeContext| async move { - Ok(NodeResult::Update(1)) - }) - .add_node("b", |_s: Counter, c: NodeContext| async move { - match c.resume { - Some(_) => Ok(NodeResult::Update(100)), - None => Ok(NodeResult::Interrupt(Interrupt::new("b", json!({})))), - } - }) - // 10 on a first (unresumed) activation, 1000 if it wrongly sees a - // resume value it never asked for. - .add_node("x", |_s: Counter, c: NodeContext| async move { - Ok(NodeResult::Update(if c.resume.is_some() { - 1000 - } else { - 10 - })) - }) - .set_entry("super") - .mark_command_routing("super") - .add_edge("a", "x") - .set_finish("b") - .set_finish("x") - .compile() - .unwrap() - .with_checkpointer(cp.clone()); - - let paused = graph - .run_with_thread( - "t-resume-scope", - Counter { - value: 0, - log: vec![], - }, - ) - .await - .unwrap(); - assert!(paused.is_interrupted()); - - let done = graph - .resume("t-resume-scope", Command::resume(json!({"approved": true}))) - .await - .unwrap(); - // 1 (a) + 100 (b, resumed) + 10 (x, first activation, no resume value). - assert_eq!( - done.state.value, 111, - "the resume value must reach only the node that interrupted" - ); -} - -#[tokio::test] -async fn attributed_update_does_not_fire_an_unsatisfied_barrier() { - // Diamond `super -> {b, c} -> merge` joined by waiting edges. `c` interrupts - // on its first activation, so the pause leaves `c` pending with only `b` - // arrived at the barrier. A manual write attributed to `b` must not - // schedule `merge` (that would run the join without `c`'s contribution), - // and the recorded arrival must let `merge` fire once `c` completes. - let cp = Arc::new(InMemoryCheckpointer::::new()); - let interrupted = Arc::new(AtomicBool::new(false)); - let once = interrupted.clone(); - let graph = GraphBuilder::::new() - .with_parallel(true) - .set_reducer(ClosureStateReducer::new(|mut s: Counter, u: i32| { - s.value += u; - s.log.push(format!("+{u}")); - Ok(s) - })) - .add_node("super", |_s: Counter, _c: NodeContext| async move { - Ok(NodeResult::Command( - Command::default().with_goto(["b", "c"]), - )) - }) - .add_node("b", |_s: Counter, _c: NodeContext| async move { - Ok(NodeResult::Update(1)) - }) - .add_node("c", move |_s: Counter, _c: NodeContext| { - let once = once.clone(); - async move { - if once.swap(true, AtomicOrdering::SeqCst) { - Ok(NodeResult::Update(2)) - } else { - Ok(NodeResult::Interrupt(Interrupt::new("c", json!({})))) - } - } - }) - .add_node("merge", |_s: Counter, _c: NodeContext| async move { - Ok(NodeResult::Update(100)) - }) - .set_entry("super") - .mark_command_routing("super") - .add_waiting_edge("b", "merge") - .add_waiting_edge("c", "merge") - .set_finish("merge") - .compile() - .unwrap() - .with_checkpointer(cp.clone()); - - let paused = graph - .run_with_thread( - "t-barrier-update", - Counter { - value: 0, - log: vec![], - }, - ) - .await - .unwrap(); - assert!(paused.is_interrupted()); - - // Operator edits state and attributes the write to `b`, the barrier - // predecessor that already ran. - graph - .update_state("t-barrier-update", 10, Some(NodeId::from("b"))) - .await - .unwrap(); - let written = cp.get("t-barrier-update", None).await.unwrap().unwrap(); - 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"), - "an unsatisfied barrier must not be scheduled by an attributed write" - ); - assert!( - written_next_nodes.iter().any(|n| n.as_str() == "c"), - "the still-pending barrier predecessor must stay scheduled" - ); - - let done = graph.retry("t-barrier-update").await.unwrap(); - assert!( - done.visited.iter().any(|n| n.as_str() == "merge"), - "merge must fire once the remaining predecessor arrives" - ); - // 1 (b) + 10 (manual write) + 2 (c) + 100 (merge). - assert_eq!(done.state.value, 113); -} - -/// Fan-out graph used by the attributed-write scheduling tests: `super` forks -/// into `b -> x -> y` and `c`, where `c` interrupts on its first activation. -/// The pause therefore leaves a checkpoint with two independent pending -/// branches (`x` from the completed `b`, and the still-interrupted `c`). -fn forked_interrupt_graph( - cp: Arc>, - interrupted: 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(["b", "c"]), - )) - }) - .add_node("b", |_s: Counter, _c: NodeContext| async move { - Ok(NodeResult::Update(1)) - }) - .add_node("x", |_s: Counter, _c: NodeContext| async move { - Ok(NodeResult::Update(20)) - }) - .add_node("y", |_s: Counter, _c: NodeContext| async move { - Ok(NodeResult::Update(40)) - }) - .add_node("c", move |_s: Counter, _c: NodeContext| { - let once = interrupted.clone(); - async move { - if once.swap(true, AtomicOrdering::SeqCst) { - Ok(NodeResult::Update(2)) - } else { - Ok(NodeResult::Interrupt(Interrupt::new("c", json!({})))) - } - } - }) - .set_entry("super") - .mark_command_routing("super") - .add_sequence(["b", "x", "y"]) - .set_finish("y") - .set_finish("c") - .compile() - .unwrap() - .with_checkpointer(cp) -} - -#[tokio::test] -async fn attributed_update_keeps_other_pending_branches_scheduled() { - // `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))); - - let paused = graph - .run_with_thread( - "t-fork-update", - Counter { - value: 0, - log: vec![], - }, - ) - .await - .unwrap(); - assert!(paused.is_interrupted()); - - let before = cp.get("t-fork-update", None).await.unwrap().unwrap(); - 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("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() == "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:?}" - ); - assert!( - !written_next_nodes.iter().any(|n| n.as_str() == "b"), - "the attributed node itself is completed, not pending: {written_next_nodes:?}" - ); - - let done = graph.retry("t-fork-update").await.unwrap(); - assert!( - done.visited.iter().any(|n| n.as_str() == "c"), - "the dropped branch must still run, visited {:?}", - done.visited - ); - 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] -async fn attributed_update_to_sink_node_keeps_other_pending_branches() { - // `c` is terminal, so an attributed write to it schedules nothing of its - // own. Replacing the pending set with that empty routing would drop the - // sibling `x` branch *and* leave a checkpoint with nothing to resume. - let cp = Arc::new(InMemoryCheckpointer::::new()); - let graph = forked_interrupt_graph(cp.clone(), Arc::new(AtomicBool::new(false))); - - let paused = graph - .run_with_thread( - "t-fork-sink", - Counter { - value: 0, - log: vec![], - }, - ) - .await - .unwrap(); - assert!(paused.is_interrupted()); - - graph - .update_state("t-fork-sink", 10, Some(NodeId::from("c"))) - .await - .unwrap(); - let written = cp.get("t-fork-sink", None).await.unwrap().unwrap(); - assert_eq!( - written - .tasks - .iter() - .map(|t| t.node.to_string()) - .collect::>(), - vec!["x".to_string()], - "the sibling branch must survive an attributed write to a sink node" - ); - - let done = graph.retry("t-fork-sink").await.unwrap(); - // 1 (b) + 10 (manual write) + 20 (x) + 40 (y); `c` is attributed as done. - assert_eq!(done.state.value, 71); -} - -#[tokio::test] -async fn attributed_update_preserves_pending_send_args_of_other_branches() { - // 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) - .set_reducer(ClosureStateReducer::new(|mut s: Counter, u: i32| { - s.value += u; - s.log.push(format!("w:{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("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!({})))); - } - Ok(NodeResult::Update(arg)) - }) - .add_node("side", |_s: Counter, _c: NodeContext| async move { - Ok(NodeResult::Update(0)) - }) - .add_node("tail", |_s: Counter, _c: NodeContext| async move { - Ok(NodeResult::Update(0)) - }) - .set_entry("dispatch") - .mark_command_routing("dispatch") - .add_edge("side", "tail") - .set_finish("worker") - .set_finish("tail") - .compile() - .unwrap() - .with_checkpointer(cp.clone()); - - let paused = graph - .run_with_thread( - "t-send-update", - Counter { - value: 0, - log: vec![], - }, - ) - .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" - ); - - 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.tasks.clone(); - let args: Vec = pending - .iter() - .filter(|a| a.node.as_str() == "worker") - .map(|a| { - a.send_arg - .as_ref() - .expect("pending Send activations keep their arg") - .as_i64() - .unwrap() - }) - .collect(); - 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 it" - ); -} - -#[tokio::test] -async fn async_durability_drains_background_writes_on_abort() { - // The recursion-limit abort returns `Err` mid-run. Any in-flight background - // checkpoint write must be settled first: dropping the tracker would detach - // the tasks, discarding their outcome and racing a caller that immediately - // retries the thread. - use crate::checkpoint::DurabilityMode; - - let completed_puts = Arc::new(AtomicUsize::new(0)); - let cp = Arc::new(SlowCheckpointer { - inner: Arc::new(InMemoryCheckpointer::new()), - delay: Duration::from_millis(50), - completed_puts: completed_puts.clone(), - }); - let graph = GraphBuilder::::overwrite() - .add_node("a", |s, _c: NodeContext| async move { - Ok(NodeResult::Update(s + 1)) - }) - .set_entry("a") - .add_edge("a", "a") - .with_recursion_limit(3) - .compile() - .unwrap() - .with_checkpointer(cp) - .with_durability(DurabilityMode::Async); - - let err = graph - .run_with_thread("t-async-abort", 0) - .await - .expect_err("the run must abort at the recursion limit"); - assert!(matches!(err, TinyAgentsError::RecursionLimit(3))); - assert_eq!( - completed_puts.load(AtomicOrdering::SeqCst), - 3, - "every background boundary write must be settled before the abort returns" - ); -} - -/// Delegating checkpointer that makes the *first* `put` slow and every later -/// one instant, so an unserialized background writer would append boundary 2 -/// before boundary 1. -struct FirstPutSlowCheckpointer { - inner: Arc>, - started: AtomicUsize, -} - -#[async_trait::async_trait] -impl Checkpointer for FirstPutSlowCheckpointer { - async fn put( - &self, - checkpoint: crate::checkpoint::Checkpoint, - ) -> tinyagents_harness::error::Result { - if self.started.fetch_add(1, AtomicOrdering::SeqCst) == 0 { - tokio::time::sleep(Duration::from_millis(100)).await; - } - self.inner.put(checkpoint).await - } - - async fn get( - &self, - thread_id: &str, - checkpoint_id: Option<&str>, - ) -> tinyagents_harness::error::Result>> { - self.inner.get(thread_id, checkpoint_id).await - } - - async fn list( - &self, - thread_id: &str, - ) -> tinyagents_harness::error::Result> { - self.inner.list(thread_id).await - } - - async fn list_threads(&self) -> tinyagents_harness::error::Result> { - self.inner.list_threads().await - } - - async fn delete_thread(&self, thread_id: &str) -> tinyagents_harness::error::Result<()> { - self.inner.delete_thread(thread_id).await - } - - async fn delete_checkpoints( - &self, - thread_id: &str, - ids: &[String], - ) -> tinyagents_harness::error::Result { - self.inner.delete_checkpoints(thread_id, ids).await - } -} - -#[tokio::test] -async fn async_durability_appends_boundaries_in_order() { - // Every bundled backend defines a thread's "latest" checkpoint by insertion - // order, so background writes must land in boundary order even when an - // earlier `put` is slower than a later one. - use crate::checkpoint::DurabilityMode; - - let cp = Arc::new(FirstPutSlowCheckpointer { - inner: Arc::new(InMemoryCheckpointer::new()), - started: AtomicUsize::new(0), - }); - let graph = GraphBuilder::::overwrite() - .add_node("a", |s, _c: NodeContext| async move { - Ok(NodeResult::Update(s + 1)) - }) - .add_node("b", |s, _c: NodeContext| async move { - Ok(NodeResult::Update(s + 1)) - }) - .add_node("c", |s, _c: NodeContext| async move { - Ok(NodeResult::Update(s + 1)) - }) - .set_entry("a") - .add_edge("a", "b") - .add_edge("b", "c") - .set_finish("c") - .compile() - .unwrap() - .with_checkpointer(cp.clone()) - .with_durability(DurabilityMode::Async); - - let run = graph.run_with_thread("t-async-order", 0).await.unwrap(); - assert_eq!(run.state, 3); - // Insertion order is listing order: each record's parent must be the one - // appended just before it. - let list = cp.list("t-async-order").await.unwrap(); - assert_eq!(list.len(), 3); - for pair in list.windows(2) { - assert_eq!( - pair[1].parent_checkpoint_id.as_deref(), - Some(pair[0].checkpoint_id.as_str()), - "boundary writes landed out of order" - ); - } -} - -#[tokio::test] -async fn legacy_interrupt_checkpoint_without_stamped_nodes_still_resumes() { - // Checkpoints written before the `interrupted_nodes` metadata existed carry - // only `Interrupt::node`, which for a re-emitted child interrupt (what a - // subgraph node does) names a node this graph never schedules. The resume - // value must still reach the paused node — falling back to the pending set — - // rather than being addressed to a node that does not exist here, which - // would re-run the paused node unresumed and interrupt forever. - let cp = Arc::new(InMemoryCheckpointer::::new()); - let graph = GraphBuilder::::new() - .set_reducer(ClosureStateReducer::new(|mut s: Counter, u: i32| { - s.value += u; - s.log.push(format!("+{u}")); - Ok(s) - })) - .add_node("gate", |_s: Counter, c: NodeContext| async move { - match c.resume { - // The interrupt names a foreign node, as a subgraph node's - // re-emitted child interrupt does. - None => Ok(NodeResult::Interrupt(Interrupt::new( - "child-gate", - json!({}), - ))), - Some(_) => Ok(NodeResult::Update(5)), - } - }) - .set_entry("gate") - .set_finish("gate") - .compile() - .unwrap() - .with_checkpointer(cp.clone()); - - let paused = graph - .run_with_thread( - "t-legacy-resume", - Counter { - value: 0, - log: vec![], - }, - ) - .await - .unwrap(); - assert!(paused.is_interrupted()); - - // Age the boundary checkpoint into its pre-upgrade shape. - let mut legacy = cp.get("t-legacy-resume", None).await.unwrap().unwrap(); - legacy - .metadata - .as_object_mut() - .unwrap() - .remove("interrupted_nodes"); - cp.put(legacy).await.unwrap(); - - let done = graph - .resume("t-legacy-resume", Command::resume(json!("go"))) - .await - .unwrap(); - assert!( - !done.is_interrupted(), - "a legacy interrupt checkpoint must still deliver the resume value" - ); - assert_eq!(done.state.value, 5); -} - -/// Delegating checkpointer whose first `put` fails *slowly* (so the next -/// boundary's write is already chained behind it) and records every attempt. -struct FirstPutFailsSlowlyCheckpointer { - inner: Arc>, - attempts: Arc, -} - -#[async_trait::async_trait] -impl Checkpointer for FirstPutFailsSlowlyCheckpointer { - async fn put( - &self, - checkpoint: crate::checkpoint::Checkpoint, - ) -> tinyagents_harness::error::Result { - if self.attempts.fetch_add(1, AtomicOrdering::SeqCst) == 0 { - tokio::time::sleep(Duration::from_millis(100)).await; - return Err(tinyagents_harness::error::TinyAgentsError::Checkpoint( - "injected background write failure".to_string(), - )); - } - self.inner.put(checkpoint).await - } - - async fn get( - &self, - thread_id: &str, - checkpoint_id: Option<&str>, - ) -> tinyagents_harness::error::Result>> { - self.inner.get(thread_id, checkpoint_id).await - } - - async fn list( - &self, - thread_id: &str, - ) -> tinyagents_harness::error::Result> { - self.inner.list(thread_id).await - } - - async fn list_threads(&self) -> tinyagents_harness::error::Result> { - self.inner.list_threads().await - } - - async fn delete_thread(&self, thread_id: &str) -> tinyagents_harness::error::Result<()> { - self.inner.delete_thread(thread_id).await - } - - async fn delete_checkpoints( - &self, - thread_id: &str, - ids: &[String], - ) -> tinyagents_harness::error::Result { - self.inner.delete_checkpoints(thread_id, ids).await - } -} - -#[tokio::test] -async fn async_durability_skips_a_write_whose_predecessor_failed() { - // Writing boundary N+1 after boundary N failed would durably append a - // record whose `parent_checkpoint_id` points at something that never - // persisted. The chained write must skip its own `put` and report the - // failure that broke the lineage. - use crate::checkpoint::DurabilityMode; - - let attempts = Arc::new(AtomicUsize::new(0)); - let cp = Arc::new(FirstPutFailsSlowlyCheckpointer { - inner: Arc::new(InMemoryCheckpointer::new()), - attempts: attempts.clone(), - }); - let graph = GraphBuilder::::overwrite() - .add_node("a", |s, _c: NodeContext| async move { - Ok(NodeResult::Update(s + 1)) - }) - .add_node("b", |s, _c: NodeContext| async move { - Ok(NodeResult::Update(s + 1)) - }) - .add_node("c", |s, _c: NodeContext| async move { - Ok(NodeResult::Update(s + 1)) - }) - .set_entry("a") - .add_edge("a", "b") - .add_edge("b", "c") - .set_finish("c") - .compile() - .unwrap() - .with_checkpointer(cp.clone()) - .with_durability(DurabilityMode::Async); - - let err = graph - .run_with_thread("t-async-orphan", 0) - .await - .expect_err("a lost background checkpoint must fail the run"); - assert!( - err.to_string() - .contains("injected background write failure"), - "unexpected error: {err}" - ); - assert_eq!( - attempts.load(AtomicOrdering::SeqCst), - 1, - "the write chained behind a failed one must not be attempted" - ); - assert!( - cp.list("t-async-orphan").await.unwrap().is_empty(), - "no orphaned checkpoint may be appended after a broken lineage" - ); -} - -/// 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/policy_tests.rs b/crates/tinyagents-graph/src/compiled/policy_tests.rs deleted file mode 100644 index 42b504c16..000000000 --- a/crates/tinyagents-graph/src/compiled/policy_tests.rs +++ /dev/null @@ -1,556 +0,0 @@ -//! 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 deleted file mode 100644 index d856543fe..000000000 --- a/crates/tinyagents-graph/src/compiled/resume.rs +++ /dev/null @@ -1,401 +0,0 @@ -//! 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/resume_tests.rs b/crates/tinyagents-graph/src/compiled/resume_tests.rs deleted file mode 100644 index 2535de452..000000000 --- a/crates/tinyagents-graph/src/compiled/resume_tests.rs +++ /dev/null @@ -1,121 +0,0 @@ -//! Resume honours completion markers: a pending task that the checkpoint's -//! write ledger says already ran is skipped, unless every pending task is -//! marked, in which case all of them re-run rather than strand the thread. - -use std::sync::Arc; -use std::sync::atomic::{AtomicUsize, Ordering}; - -use serde_json::json; - -use crate::builder::{GraphBuilder, NodeContext}; -use crate::checkpoint::{CheckpointConfig, Checkpointer, InMemoryCheckpointer, PendingWrite}; -use crate::command::{Command, Interrupt, NodeResult}; -use crate::reducer::ClosureStateReducer; - -/// The graph, its checkpointer and the `p`/`q` run counters. -type Fixture = ( - crate::CompiledGraph, - Arc>, - Arc, - Arc, -); - -/// `p` and `q` both pause on their first run; on resume each adds its value. -fn two_pausing_nodes() -> Fixture { - let cp = Arc::new(InMemoryCheckpointer::::new()); - let p_runs = Arc::new(AtomicUsize::new(0)); - let q_runs = Arc::new(AtomicUsize::new(0)); - let (p, q) = (p_runs.clone(), q_runs.clone()); - let graph = GraphBuilder::::new() - .with_parallel(true) - .set_reducer(ClosureStateReducer::new(|s: i32, u: i32| Ok(s + u))) - .add_node("split", |_s: i32, _c: NodeContext| async move { - Ok(NodeResult::Command( - Command::default().with_goto(["p", "q"]), - )) - }) - .add_node("p", move |_s: i32, c: NodeContext| { - let runs = p.clone(); - async move { - runs.fetch_add(1, Ordering::SeqCst); - match c.resume { - Some(_) => Ok(NodeResult::Update(1)), - None => Ok(NodeResult::Interrupt(Interrupt::new("p", json!({})))), - } - } - }) - .add_node("q", move |_s: i32, c: NodeContext| { - let runs = q.clone(); - async move { - runs.fetch_add(1, Ordering::SeqCst); - match c.resume { - Some(_) => Ok(NodeResult::Update(10)), - None => Ok(NodeResult::Interrupt(Interrupt::new("q", json!({})))), - } - } - }) - .set_entry("split") - .mark_command_routing("split") - .set_finish("p") - .set_finish("q") - .compile() - .unwrap() - .with_checkpointer(cp.clone()); - (graph, cp, p_runs, q_runs) -} - -/// Pauses on thread `t` and records a completion marker for each of `marked`. -async fn pause_and_mark( - graph: &crate::CompiledGraph, - cp: &InMemoryCheckpointer, - marked: &[&str], -) { - let paused = graph.run_with_thread("t", 0).await.unwrap(); - assert!(paused.is_interrupted()); - let checkpoint = cp.get("t", None).await.unwrap().unwrap(); - assert_eq!(checkpoint.tasks.len(), 2, "both p and q are pending"); - let writes: Vec = checkpoint - .tasks - .iter() - .filter(|task| marked.contains(&task.node.as_str())) - .map(|task| PendingWrite::completion_marker(task.node.clone(), task.task_id.clone())) - .collect(); - let config = CheckpointConfig { - thread_id: "t".into(), - checkpoint_id: Some(checkpoint.checkpoint_id.clone()), - namespace: Vec::new(), - }; - cp.put_writes(&config, &writes).await.unwrap(); -} - -#[tokio::test] -async fn resume_skips_a_pending_task_with_a_completion_marker() { - let (graph, cp, p_runs, q_runs) = two_pausing_nodes(); - pause_and_mark(&graph, &cp, &["p"]).await; - - let done = graph - .resume("t", Command::resume(json!(null))) - .await - .unwrap(); - assert_eq!( - p_runs.load(Ordering::SeqCst), - 1, - "p is marked done: not re-run" - ); - assert_eq!(q_runs.load(Ordering::SeqCst), 2, "q pauses, then resumes"); - assert_eq!(done.state, 10); -} - -#[tokio::test] -async fn resume_reruns_everything_when_every_pending_task_is_marked() { - let (graph, cp, p_runs, q_runs) = two_pausing_nodes(); - pause_and_mark(&graph, &cp, &["p", "q"]).await; - - let done = graph - .resume("t", Command::resume(json!(null))) - .await - .unwrap(); - assert_eq!(p_runs.load(Ordering::SeqCst), 2); - assert_eq!(q_runs.load(Ordering::SeqCst), 2); - assert_eq!(done.state, 11); -} diff --git a/crates/tinyagents-graph/src/compiled/routing.rs b/crates/tinyagents-graph/src/compiled/routing.rs deleted file mode 100644 index a3111cad4..000000000 --- a/crates/tinyagents-graph/src/compiled/routing.rs +++ /dev/null @@ -1,290 +0,0 @@ -//! Routing: resolving a completed step's active set into the next -//! superstep's activations (`goto`/conditional branches, [`Send`] -//! fanout, and interrupt-durability preconditions). -//! -//! Split out of `compiled/mod.rs`; see that module's doc comment for the -//! executor's overall design. - -use super::*; - -impl CompiledGraph -where - State: Clone + Send + Sync + 'static, - Update: Send + 'static, -{ - /// Resolves the next superstep's activation set from a completed step's - /// results: routes each completed activation (`self.route`), applies - /// barrier gating for waiting/fan-in nodes, and runs the mixed-fan-in - /// barrier-relief pass so a barrier fed by a conditional predecessor - /// doesn't deadlock when that branch wasn't taken. - /// - /// `goto_map` carries per-activation-index [`Command::goto`] targets - /// (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, - 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 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 (orig_index, activation) in completed.iter() { - let node_id = &activation.node; - 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_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) { - let arrived = barrier_arrivals.entry(tnode.clone()).or_default(); - arrived.insert(node_id.clone()); - if !required.is_subset(arrived) { - continue; - } - barrier_arrivals.remove(&tnode); - } - // `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().map(Arc::new); - if send_arg.is_some() { - next.push(Activation { - node: tnode, - send_arg, - task_id: TaskId::from(String::new()), - }); - } else if next_seen.insert(tnode.clone()) { - next.push(Activation { - node: tnode, - send_arg: None, - task_id: TaskId::from(String::new()), - }); - } - } - } - - // Mixed fan-in barrier relief: a waiting/barrier node normally - // activates only once every registered predecessor has arrived. When - // one of those predecessors is reachable only via a conditional - // branch, and the *taken* branch does not lead toward it, that - // predecessor will never arrive on its own — register a phantom - // arrival on its behalf so the barrier can still clear on the - // predecessors that actually ran, instead of deadlocking forever. - // - // The check is keyed on `source`'s *resolved routing target* this - // step (via `reaches_deterministically`), not on whether - // `relief_node` is freshly scheduled into `next` this same - // superstep. A same-superstep check is wrong for a multi-hop - // conditional predecessor (`source --branch--> x --main--> - // relief_node`, `x` a plain pass-through): `relief_node` would not - // yet be in `next` on the step `source` itself completes — even on - // the *taken* branch, where `relief_node` WILL run once `x` does — - // so a same-superstep check would fire a premature phantom arrival - // and let the barrier clear before the real predecessor's data - // commits, reintroducing the exact data-loss bug this primitive - // exists to prevent. Resolving `source`'s actual target and walking - // it forward through deterministic (non-branching) edges is correct - // for both direct and multi-hop cases because it answers "was the - // branch leading to `relief_node` taken", not "did `relief_node` - // happen to run in lockstep with `source`". - for relief in self.barrier_reliefs.iter() { - let source_indices: Vec = completed - .iter() - .enumerate() - .filter(|(_, (_, activation))| activation.node == relief.source) - .map(|(index, _)| index) - .collect(); - if source_indices.is_empty() { - continue; - } - let branch_taken = source_indices.iter().any(|index| { - resolved[*index].iter().any(|target| { - self.reaches_deterministically( - target.node(), - &relief.relief_node, - &relief.barrier_node, - ) - }) - }); - // A relief_node freshly scheduled into `next` this step (a - // direct, single-hop predecessor) is also proof the branch was - // taken; kept as a defensive fallback alongside the resolved- - // target check above. - if branch_taken || next_seen.contains(&relief.relief_node) { - continue; - } - let Some(required) = self.waiting.get(&relief.barrier_node) else { - continue; - }; - let arrived = barrier_arrivals - .entry(relief.barrier_node.clone()) - .or_default(); - arrived.insert(relief.relief_node.clone()); - if !required.is_subset(arrived) { - continue; - } - barrier_arrivals.remove(&relief.barrier_node); - if next_seen.insert(relief.barrier_node.clone()) { - next.push(Activation { - node: relief.barrier_node.clone(), - send_arg: None, - task_id: TaskId::from(String::new()), - }); - } - } - - Ok(next) - } - - /// Whether `to` is reachable from `from` by following only deterministic - /// static routing — plain/waiting edges (`self.edges`), which resolve to - /// exactly one successor with no runtime decision — without ever - /// expanding through `stop`. - /// - /// This is what makes barrier-relief evaluation correct for a - /// multi-hop conditional predecessor: once a brancher's routing decision - /// for this step is resolved to a concrete target, whether that target - /// eventually leads to a barrier's conditional predecessor is a static - /// property of the compiled topology for any chain of plain pass-through - /// nodes — it does not depend on when each hop happens to run. A further - /// conditional/command node along the way (no `self.edges` entry) is a - /// second runtime decision this walk cannot resolve ahead of time, so it - /// conservatively reports unreachable there (falling back to the - /// same-superstep check). - fn reaches_deterministically(&self, from: &NodeId, to: &NodeId, stop: &NodeId) -> bool { - if from == to { - return true; - } - // 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(current) = stack.pop() { - if !seen.insert(current) { - continue; - } - 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); - } - } - false - } - - /// Resolves the next routing targets for `node_id`. - /// - /// Command `goto` (which may include [`Send`] packets) wins over static and - /// conditional edges; edge/conditional targets are plain node activations. - /// - /// `goto` carries this specific activation's [`Command::goto`] targets (when - /// it returned a command), passed per-activation rather than looked up by - /// node id so repeated activations of one node never share routing. - pub(super) fn route( - &self, - node_id: &NodeId, - goto: Option<&[RouteTarget]>, - state: &State, - ) -> Result> { - if let Some(targets) = goto { - self.validate_route_targets(node_id, targets)?; - return Ok(targets.to_vec()); - } - 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).to_string(); - let target = branch.routes.get(&route).cloned().ok_or_else(|| { - TinyAgentsError::MissingRoute { - node: node_id.to_string(), - route, - } - })?; - return Ok(vec![RouteTarget::Node(target)]); - } - // Sink: no outgoing routing, the branch ends here. - Ok(Vec::new()) - } - - fn validate_route_targets(&self, node_id: &NodeId, targets: &[RouteTarget]) -> Result<()> { - for target in targets { - let target_node = target.node(); - if target_node.as_str() == END { - continue; - } - if target_node.as_str() == START { - return Err(TinyAgentsError::Graph(format!( - "command goto from node `{node_id}` cannot target START" - ))); - } - if !self.nodes.contains_key(target_node) { - return Err(TinyAgentsError::MissingNode(target_node.to_string())); - } - } - Ok(()) - } - - /// Checks the precondition for emitting an [`Interrupt`]: a checkpointer - /// and a thread id must both be present, since an interrupt is only - /// resumable if it can be persisted and later addressed by thread. - pub(super) fn require_interrupt_durability(&self, thread_id: &Option) -> Result<()> { - if self.checkpointer.is_none() { - return Err(TinyAgentsError::Resume( - "interrupt emitted without a configured checkpointer".to_string(), - )); - } - if thread_id.is_none() { - return Err(TinyAgentsError::Resume( - "interrupt emitted without a thread id".to_string(), - )); - } - Ok(()) - } -} diff --git a/crates/tinyagents-graph/src/compiled/run_ctx.rs b/crates/tinyagents-graph/src/compiled/run_ctx.rs deleted file mode 100644 index 9e9ce6fa3..000000000 --- a/crates/tinyagents-graph/src/compiled/run_ctx.rs +++ /dev/null @@ -1,567 +0,0 @@ -//! 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 deleted file mode 100644 index 5c48f9b68..000000000 --- a/crates/tinyagents-graph/src/compiled/state_api.rs +++ /dev/null @@ -1,418 +0,0 @@ -//! State inspection and manual state-write API (`get_state`, -//! `get_state_history`, `update_state`, `bulk_update_state`, `fork_state`). -//! -//! Split out of `compiled/mod.rs`; see that module's doc comment for the -//! executor's overall durability design. - -use super::*; - -impl CompiledGraph -where - State: Clone + Send + Sync + 'static, - Update: Send + 'static, -{ - fn require_checkpointer(&self) -> Result<&Arc>> { - self.checkpointer - .as_ref() - .ok_or_else(|| TinyAgentsError::Checkpoint("no checkpointer configured".to_string())) - } - - /// Builds a [`CheckpointConfig`] addressing `checkpoint_id` (or the latest - /// when `None`) under this graph's namespace. - fn config_for(&self, thread_id: &str, checkpoint_id: Option<&str>) -> CheckpointConfig { - CheckpointConfig { - thread_id: thread_id.to_string(), - checkpoint_id: checkpoint_id.map(str::to_string), - namespace: self.namespace.clone(), - } - } - pub async fn get_state( - &self, - thread_id: &str, - checkpoint_id: Option<&str>, - ) -> Result>> { - let checkpointer = self.require_checkpointer()?; - let config = self.config_for(thread_id, checkpoint_id); - Ok(checkpointer - .get_tuple(config) - .await? - .map(snapshot_from_tuple)) - } - - /// Returns a thread's state history newest-first, walking the - /// `parent_checkpoint_id` lineage from the latest checkpoint backwards. - /// - /// `limit` caps the number of snapshots returned (the most recent ones). - /// Requires a configured checkpointer. - pub async fn get_state_history( - &self, - thread_id: &str, - limit: Option, - ) -> Result>> { - let checkpointer = self.require_checkpointer()?; - // Delegate the parent-lineage walk to the checkpointer so backends that - // would otherwise re-read the whole thread per hop (the file backend) can - // read it once and walk in memory (O(H) instead of O(H²)). - let tuples = checkpointer - .state_history(thread_id, &self.namespace, limit) - .await?; - Ok(tuples.into_iter().map(snapshot_from_tuple).collect()) - } - - /// Applies a manual state write to a thread, producing a new checkpoint with - /// source `update`. - /// - /// The write is a genuine graph write: `update` is folded through the same - /// [`StateReducer`](crate::StateReducer) the executor uses, on top of - /// the thread's latest committed state. When `as_node` is supplied it must - /// name a real node (else [`TinyAgentsError::MissingNode`]); the write is - /// attributed to that node, which is treated as having just completed: it - /// leaves the pending set and its routing successors are *merged into* the - /// base checkpoint's remaining pending work (so a subsequent resume - /// continues from after the attributed node without dropping the branches - /// it never touched). Sibling branches keep their `Send` args, and a - /// successor that is already pending is not scheduled twice. When the - /// attributed node has several pending `Send` activations, the write - /// completes all of them at once — a manual write cannot name which packet - /// it stands for — and the successor is scheduled once. A command node - /// cannot be used as `as_node` (it routes dynamically and has no static - /// successors); doing so returns [`TinyAgentsError::Graph`] rather than - /// silently producing a non-resumable checkpoint. A successor reached by a - /// waiting edge is barrier-gated exactly as it would be during a run: it is - /// scheduled only once every required predecessor has arrived (the write - /// records the attributed node's arrival), so a manual write can never fire - /// a join ahead of a still-pending branch — and because the other pending - /// predecessors are retained, they still run and clear the join. With - /// `as_node == None` the latest pending node set is preserved. Requires a - /// configured checkpointer and an existing checkpoint for the thread. - pub async fn update_state( - &self, - thread_id: &str, - update: Update, - as_node: Option, - ) -> Result { - let checkpointer = self.require_checkpointer()?; - if let Some(node) = &as_node { - if !self.nodes.contains_key(node) { - return Err(TinyAgentsError::MissingNode(node.to_string())); - } - // A command node routes dynamically (via the [`Command`] it returns - // at runtime), so it has no static successors to schedule here. - // Attributing a manual write to one would persist an empty - // `next_nodes` and silently render the thread non-resumable, so - // reject it at write time instead. - if self.command_nodes.contains(node) { - return Err(TinyAgentsError::Graph(format!( - "cannot update state as node `{node}`: it routes dynamically \ - via Command and has no static successors, so the resulting \ - checkpoint would be non-resumable" - ))); - } - } - - let base = checkpointer - .get_scoped(thread_id, None, &self.namespace) - .await? - .ok_or_else(|| { - TinyAgentsError::Checkpoint(format!( - "cannot update state: no checkpoint exists for thread `{thread_id}`" - )) - })?; - 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 (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 and any - // carried-forward completions' successors, merged into the base - // checkpoint's still-pending work. - // - // `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; - } - 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()), - }); - } - } - 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 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?; - 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) - } - - /// Applies a sequence of manual writes as successive `update` checkpoints, - /// returning the config of the last one written. - /// - /// Each `(update, as_node)` pair is applied with [`CompiledGraph::update_state`] - /// in order, so every step layers on the previous one's committed state and - /// produces its own checkpoint. Returns [`TinyAgentsError::Checkpoint`] when - /// the iterator is empty (there is no resulting config to return). - pub async fn bulk_update_state( - &self, - thread_id: &str, - updates: impl IntoIterator)>, - ) -> Result { - let mut last: Option = None; - for (update, as_node) in updates { - last = Some(self.update_state(thread_id, update, as_node).await?); - } - last.ok_or_else(|| { - TinyAgentsError::Checkpoint("bulk_update_state received no updates".to_string()) - }) - } - - /// Forks a checkpoint into a new thread, producing a fresh root checkpoint - /// with source `fork`. - /// - /// Copies the addressed source checkpoint's committed state, pending nodes, - /// completed tasks, pending writes, and interrupts into `target_thread` under - /// a brand-new checkpoint id with no parent (the root of the new thread). The - /// source record is read with `get` and never mutated, so time-travel forks - /// are non-destructive. With `source_checkpoint_id == None` the source - /// thread's latest checkpoint is forked. Requires a configured checkpointer. - pub async fn fork_state( - &self, - source_thread: &str, - source_checkpoint_id: Option<&str>, - target_thread: &str, - ) -> Result { - let checkpointer = self.require_checkpointer()?; - let source = checkpointer - .get_scoped(source_thread, source_checkpoint_id, &self.namespace) - .await? - .ok_or_else(|| { - TinyAgentsError::Checkpoint(format!( - "cannot fork: no checkpoint found for thread `{source_thread}`" - )) - })?; - let step = source.to_metadata().step; - let checkpoint_id = next_checkpoint_id(); - let config = self.config_for(target_thread, Some(&checkpoint_id)); - // `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_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 deleted file mode 100644 index 47d40c782..000000000 --- a/crates/tinyagents-graph/src/compiled/step.rs +++ /dev/null @@ -1,1006 +0,0 @@ -//! 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/types.rs b/crates/tinyagents-graph/src/compiled/types.rs deleted file mode 100644 index cf4bd88ca..000000000 --- a/crates/tinyagents-graph/src/compiled/types.rs +++ /dev/null @@ -1,588 +0,0 @@ -//! Compiled graph and execution-result types. -//! -//! [`CompiledGraph`] is the frozen, runnable output of -//! [`GraphBuilder::compile`](crate::builder::GraphBuilder::compile); its -//! behavior (the superstep loop, routing, state-inspection API) lives in the -//! sibling `executor.rs`, `routing.rs`, and `state_api.rs` files, which all -//! `impl` against the type defined here. [`GraphExecution`] and -//! [`StateSnapshot`] are the two result shapes those entry points return — -//! a finished/paused run and a point-in-time checkpoint view, respectively. - -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, NodePolicy, UpdateCodec}; -use crate::checkpoint::{CheckpointConfig, CheckpointMetadata, Checkpointer, DurabilityMode}; -use crate::command::Interrupt; -use crate::observability::{GraphEventJournal, GraphStatusStore}; -use crate::recursion::{ChildRun, RunTree}; -use crate::reducer::StateReducer; -use crate::status::GraphRunStatus; -use crate::stream::GraphEventSink; -use tinyagents_harness::ids::{CheckpointId, GraphId, NodeId, RunId}; - -/// An immutable, validated graph ready to execute. -/// -/// Cheap to clone (all heavy fields are `Arc`-shared) and safe to run -/// concurrently. Construct one with [`crate::GraphBuilder::compile`]. -pub struct CompiledGraph { - pub(crate) graph_id: GraphId, - /// Optional human-readable graph name surfaced by the topology export. - pub(crate) name: Option, - pub(crate) nodes: Arc>>, - pub(crate) edges: Arc>>, - pub(crate) branches: Arc>>, - pub(crate) command_nodes: Arc>, - /// Barrier/waiting edges: target -> the predecessor set that must all - /// complete (across steps) before the target activates. - pub(crate) waiting: Arc>>, - /// Mixed fan-in barrier relief registrations; see [`BarrierRelief`]. - pub(crate) barrier_reliefs: Arc>, - /// Behavior-free per-node markers/metadata surfaced by the topology export. - pub(crate) node_meta: Arc>, - pub(crate) entry: NodeId, - pub(crate) reducer: Arc>, - pub(crate) recursion_limit: usize, - /// Explicit recursion caps (run-tree depth, per-node visits, total steps) - /// enforced by the executor; configured via - /// [`CompiledGraph::with_recursion_policy`](crate::CompiledGraph::with_recursion_policy). - pub(crate) recursion_policy: crate::recursion::RecursionPolicy, - /// Inherited recursion frames of an enclosing run (empty for a top-level - /// graph). A subgraph/sub-agent wrapper seeds these so a nested run extends - /// the parent's recursion stack rather than starting a fresh tree. - pub(crate) recursion_frames: Vec, - /// The hosting node this graph runs under when embedded as a subgraph; used - /// as the `node_id` of this run's root recursion frame so the run tree names - /// the embedding node. `None` for a top-level run. - pub(crate) recursion_node: Option, - pub(crate) checkpointer: Option>>, - pub(crate) event_sink: Option>, - /// Optional durable observation journal (opt-in via - /// [`CompiledGraph::with_event_journal`]). - pub(crate) journal: Option>, - /// Optional run-status surface (opt-in via - /// [`CompiledGraph::with_status_store`]). - pub(crate) status_store: Option>, - pub(crate) namespace: Vec, - /// When true, the active node set of a superstep is executed concurrently. - pub(crate) parallel: bool, - /// Upper bound on concurrently-running branches per step (`None` = unbounded). - pub(crate) max_concurrency: Option, - /// Default per-node handler timeout (`None` = no timeout). - pub(crate) node_timeout: Option, - /// Optional whole-run wall-clock deadline (`None` = no deadline). Checked at - /// every super-step boundary: when the elapsed run time reaches it the run - /// stops *between* super-steps with a [`TinyAgentsError::Timeout`], leaving - /// the last committed boundary checkpoint intact and resumable. Configured - /// via [`CompiledGraph::with_run_deadline`](crate::CompiledGraph::with_run_deadline). - pub(crate) run_deadline: Option, - /// When checkpoints are persisted relative to execution (default - /// [`DurabilityMode::Sync`]). - pub(crate) durability: DurabilityMode, - /// Optional per-node retry policy applied around each node handler. When - /// set, a handler that fails with a - /// [retryable][tinyagents_harness::retry::is_retryable] error is re-run from its - /// start (with opt-in backoff) up to the policy's attempt cap before the - /// failure escalates. `None` (default) preserves the original - /// 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 { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("CompiledGraph") - .field("graph_id", &self.graph_id) - .field("nodes", &self.nodes.len()) - .field("entry", &self.entry) - .field("recursion_limit", &self.recursion_limit) - .field("namespace", &self.namespace) - .field("parallel", &self.parallel) - .finish_non_exhaustive() - } -} - -impl Clone for CompiledGraph { - fn clone(&self) -> Self { - Self { - graph_id: self.graph_id.clone(), - name: self.name.clone(), - nodes: self.nodes.clone(), - edges: self.edges.clone(), - branches: self.branches.clone(), - command_nodes: self.command_nodes.clone(), - waiting: self.waiting.clone(), - barrier_reliefs: self.barrier_reliefs.clone(), - node_meta: self.node_meta.clone(), - entry: self.entry.clone(), - reducer: self.reducer.clone(), - recursion_limit: self.recursion_limit, - recursion_policy: self.recursion_policy, - recursion_frames: self.recursion_frames.clone(), - recursion_node: self.recursion_node.clone(), - checkpointer: self.checkpointer.clone(), - event_sink: self.event_sink.clone(), - journal: self.journal.clone(), - status_store: self.status_store.clone(), - namespace: self.namespace.clone(), - parallel: self.parallel, - max_concurrency: self.max_concurrency, - node_timeout: self.node_timeout, - 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(), - } - } -} - -/// Tracks checkpoint writes handed to background tasks under -/// [`DurabilityMode::Async`]. -/// -/// # Failure semantics -/// -/// Async durability is fire-and-forget for *latency*, never for *outcome*: a -/// background `put` error is recorded here rather than dropped, and the -/// executor surfaces it at the next durability boundary -/// ([`Self::take_failure`]) or — at the latest — when the run drains all -/// in-flight writes at its terminal/interrupt boundary ([`Self::drain`]). A -/// run therefore never reports success while one of its checkpoints silently -/// failed to persist. -/// -/// # Ordering -/// -/// Writes are *chained*, not merely spawned: each background task awaits the -/// previously spawned one before performing its own `put`. Every bundled -/// backend defines a thread's "latest" checkpoint by insertion order, so -/// letting boundary N+1 land before boundary N would make a concurrent reader -/// (or a `retry` racing a straggler append) observe a stale record as latest. -#[derive(Default)] -pub(crate) struct AsyncCheckpointWrites { - /// In-flight (or finished-but-unharvested) background writes, in the order - /// they were spawned. Because each spawn chains onto its predecessor, this - /// holds at most the tail of the chain. - handles: Vec>>, -} - -impl AsyncCheckpointWrites { - /// Spawns `write` so it runs only after every previously spawned write has - /// settled, keeping appends in boundary order. - /// - /// The spawned task reports the first error along the chain (in spawn - /// order), so a predecessor's failure is never lost by being folded behind - /// a later success. Awaiting a panicked predecessor yields a join error - /// rather than deadlocking the chain. - pub(crate) fn spawn_ordered(&mut self, runtime: &tokio::runtime::Handle, write: F) - where - F: std::future::Future> - + Send - + 'static, - { - let previous = self.handles.pop(); - self.handles.push(runtime.spawn(async move { - let prior = match previous { - Some(handle) => Self::harvest(handle.await), - None => None, - }; - // A predecessor that never persisted would leave this record's - // `parent_checkpoint_id` dangling, so skip the write entirely and - // report the failure that broke the lineage. - match prior { - Some(err) => Err(err), - None => write.await, - } - })); - } - - /// Harvests writes that have already finished, without blocking on the - /// ones still in flight. Returns the first error found (a failed `put` or - /// a panicked task), or `None` when every finished write succeeded. - pub(crate) async fn take_failure( - &mut self, - ) -> Option { - let mut failure = None; - let mut remaining = Vec::with_capacity(self.handles.len()); - for handle in self.handles.drain(..) { - if !handle.is_finished() { - remaining.push(handle); - continue; - } - // Finished, so this await is immediate. - if let Some(err) = Self::harvest(handle.await) - && failure.is_none() - { - failure = Some(err); - } - } - self.handles = remaining; - failure - } - - /// Awaits every in-flight write. Returns the first error (spawn order), - /// after all writes have settled. This is the "final await at run end": - /// terminal, interrupt, and failure boundaries drain before returning so - /// the run result reflects persistence failures. - pub(crate) async fn drain(&mut self) -> tinyagents_harness::error::Result<()> { - let mut failure = None; - for handle in self.handles.drain(..) { - if let Some(err) = Self::harvest(handle.await) - && failure.is_none() - { - failure = Some(err); - } - } - match failure { - Some(err) => Err(err), - None => Ok(()), - } - } - - /// Maps one settled join result onto its error, if any. - fn harvest( - joined: std::result::Result< - tinyagents_harness::error::Result, - tokio::task::JoinError, - >, - ) -> Option { - match joined { - Ok(Ok(_)) => None, - Ok(Err(err)) => Some(err), - Err(join_err) => Some(tinyagents_harness::error::TinyAgentsError::Checkpoint( - format!("background checkpoint write task failed: {join_err}"), - )), - } - } -} - -/// The result of a durable graph run. -/// -/// Carries the final committed state, the visited node history, the superstep -/// count, any pending interrupts, the run status snapshot, and the latest -/// checkpoint id (when checkpointing was enabled). -#[derive(Clone, Debug)] -pub struct GraphExecution { - /// Final committed state. - pub state: State, - /// This run's own id. - pub run_id: RunId, - /// The graph that produced this run. - pub graph_id: GraphId, - /// The root run id of the recursion tree (equals `run_id` for a top-level - /// run; the shared ancestor for a subgraph/sub-agent child run). - pub root_run_id: RunId, - /// The enclosing run's id, when this run was spawned as a child. - pub parent_run_id: Option, - /// Child runs spawned from subgraph nodes during this run, in completion - /// order, keyed by the embedding node. - pub child_runs: Vec, - /// Ordered list of executed nodes (may repeat across supersteps). - pub visited: Vec, - /// Number of supersteps executed. - pub steps: usize, - /// Interrupts that paused the run, if any. - pub interrupts: Vec, - /// Compact status snapshot at the final boundary. - 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. -/// -/// Most runs should target [`START`] via [`GraphInput::start`], which delivers -/// the payload to the graph's compiled entry node. More complex orchestrations -/// can seed additional nodes in the same first superstep, letting independent -/// loops receive user input, tool results, webhook payloads, or other external -/// events without forcing every input through one entry handler. -#[derive(Clone, Debug, PartialEq)] -pub struct GraphInput { - /// Target node for this input. `__start__` is accepted as the virtual entry - /// target and is resolved to the graph's compiled entry node at run time. - pub node: NodeId, - /// Optional per-input payload delivered as - /// [`NodeContext::send_arg`](crate::NodeContext::send_arg). - pub payload: Option, -} - -impl GraphInput { - /// Delivers `payload` to the graph's compiled entry node. - pub fn start(payload: serde_json::Value) -> Self { - Self { - node: NodeId::from(START), - payload: Some(payload), - } - } - - /// Delivers `payload` directly to `node` in the first superstep. - pub fn new(node: impl Into, payload: serde_json::Value) -> Self { - Self { - node: node.into(), - payload: Some(payload), - } - } - - /// Activates `node` without a per-input payload. - pub fn node(node: impl Into) -> Self { - Self { - node: node.into(), - payload: None, - } - } -} - -impl GraphExecution { - /// Returns true when the run paused on an interrupt rather than completing. - pub fn is_interrupted(&self) -> bool { - !self.interrupts.is_empty() - } - - /// Builds the parent/child run-lineage view for this run. - /// - /// This is the run-id counterpart to the live recursion stack: it reports - /// this run's id, the shared root, the enclosing parent (when this was a - /// child run), and every child run spawned from a subgraph node. - pub fn run_tree(&self) -> RunTree { - RunTree { - run_id: self.run_id.clone(), - root_run_id: self.root_run_id.clone(), - parent_run_id: self.parent_run_id.clone(), - children: self.child_runs.clone(), - } - } -} - -/// A point-in-time view of a thread's checkpointed state, returned by -/// [`CompiledGraph::get_state`](crate::CompiledGraph::get_state) and -/// [`CompiledGraph::get_state_history`](crate::CompiledGraph::get_state_history). -/// -/// This is the state-inspection / time-travel surface: it bundles the committed -/// channel values at a checkpoint with everything a caller needs to reason about -/// or resume from that point — the next nodes that would run, the config that -/// addresses this snapshot, the config of its parent (for walking the lineage), -/// the checkpoint metadata (source/step/run id), and any pending interrupts. -#[derive(Clone, Debug)] -pub struct StateSnapshot { - /// Committed graph state (channel values) at the snapshot's checkpoint. - pub values: State, - /// Nodes that would run if execution resumed from this snapshot. - pub next_nodes: Vec, - /// Tasks scheduled for the next superstep (mirrors `next_nodes`). - pub tasks: Vec, - /// Config that addresses this snapshot's checkpoint. - pub config: CheckpointConfig, - /// Lightweight metadata for the snapshot's checkpoint (source, step, run id). - pub metadata: CheckpointMetadata, - /// Config addressing the parent checkpoint, when one exists. - pub parent_config: Option, - /// Interrupts that paused the run at this checkpoint, if any. - 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 -/// for [`ResumeTarget::Latest`]; [`CompiledGraph::resume_from`](crate::CompiledGraph::resume_from) -/// accepts an explicit target so a run can be replayed from an older checkpoint -/// config (time travel) without mutating the original record. -#[derive(Clone, Debug, PartialEq, Eq)] -pub enum ResumeTarget { - /// Resume from the thread's most recent checkpoint. - Latest, - /// Resume from a specific checkpoint id within the thread. - Checkpoint(String), -} diff --git a/crates/tinyagents-graph/src/dag/README.md b/crates/tinyagents-graph/src/dag/README.md deleted file mode 100644 index 03feaef9a..000000000 --- a/crates/tinyagents-graph/src/dag/README.md +++ /dev/null @@ -1,64 +0,0 @@ -# graph::dag - -Structural validation of a dependency DAG — duplicate ids, dangling edges, and -cycles — over a borrowed node view, independent of any one caller's node -shape or of the graph runtime. - -Multiple graph surfaces (task boards, plan steps, and host-defined workflows) -keep re-deriving the same -question: given nodes that each name the nodes they must run after, is the -result a well-formed DAG? Rather than each caller reimplementing Kahn's -algorithm against its own struct, this module takes a borrowed [`DagNode`] -view that any caller can project its own type into for the duration of one -validation call. It touches no graph runtime state and is usable before -anything is compiled or executed. - -## Public surface - -- `DagNode<'a>` — a borrowed `id` plus the `depends_on` ids it names (edges - read as `depends_on -> id`: a node runs after everything it names). - `DagNode::new(id, depends_on)` builds one from an id and any iterator of - dependency ids. -- `DagIssue` — a structural problem: `DuplicateNode { id }`, - `UnknownDependency { node, depends_on }`, or `Cycle` (a self-edge counts as - a cycle). Deliberately narrow: domain rules a host layers on top (e.g. "a - phase must name at least one agent") are not this module's concern. -- `has_cycle(nodes) -> bool` — the cycle question alone, for a caller that - already validated ids/edges its own way. -- `validate_dag(nodes) -> Vec` — duplicates, dangling edges, and - cycles in one pass, returning every issue found (not just the first) so a - definition can be shown with all its problems at once. An empty result - means the input is a well-formed DAG. - -## Files - -| File | Role | -| --- | --- | -| `types.rs` | `DagNode`, `DagIssue`. | -| `mod.rs` | `has_cycle`, `validate_dag` (Kahn's algorithm, O(V + E)). | -| `mod_tests.rs` | Unit tests (acyclic acceptance, cycles, self-edges, dangling edges, the duplicate-id false-cycle guard). | - -## Semantics worth knowing - -- **Edges pointing at unknown ids are ignored by the cycle check** and - reported separately as `DagIssue::UnknownDependency` — a dangling edge - cannot close a loop, and letting it participate would turn one mistake - into two errors. -- **Duplicate ids never trip a false cycle.** The cycle graph is built from - exactly one declaration per id (first declaration wins); reachability is - compared against the unique-node count, not the input length. -- **A self-edge is a cycle.** A host that wants to surface self-dependency as - its own diagnostic must check for it before calling. -- A caller that wants to add dependency edges to a node id that already - exists must merge those edges into that node's single `DagNode` - declaration — a second `DagNode` with the same id has its edges silently - dropped by the dedupe (first declaration wins), so a real cycle it would - have introduced can go undetected. This is unreachable for a caller that - only ever mints fresh ids, but is a live footgun for one that revisits an - existing id. - -## How it fits together - -`graph::dag` has no dependency on the rest of `graph`; it is a leaf utility -re-exported at the crate root for graph modules and hosts that need to project -their own node types into `DagNode` views before compiling or scheduling them. diff --git a/crates/tinyagents-graph/src/dag/mod.rs b/crates/tinyagents-graph/src/dag/mod.rs deleted file mode 100644 index 36427f008..000000000 --- a/crates/tinyagents-graph/src/dag/mod.rs +++ /dev/null @@ -1,165 +0,0 @@ -//! Structural validation of a dependency DAG. -//! -//! Orchestration surfaces keep re-deriving the same graph question: given nodes -//! that each name the nodes they must run after, are the ids unique, do all the -//! edges land, and is the whole thing acyclic? Workflow phases, team task -//! boards and plan steps all ask it of different node shapes, so this module -//! takes a borrowed [`DagNode`] view rather than any one caller's struct. -//! -//! Two entry points: -//! -//! - [`has_cycle`] — the cycle question alone, for a caller that already -//! validated ids and edges its own way (a scheduler admitting one new node -//! at a time, say, where a dangling edge on an *existing* node is not the new -//! node's fault). -//! - [`validate_dag`] — duplicates, dangling edges and cycles in one pass, -//! returning every issue rather than the first, so a definition can be shown -//! with all of its problems at once. -//! -//! # Semantics worth knowing -//! -//! - **Edges pointing at unknown ids are ignored by the cycle check.** They are -//! reported as [`DagIssue::UnknownDependency`] instead. A dangling edge -//! cannot close a loop, and letting it participate would turn one mistake -//! into two errors. -//! - **Duplicate ids never trip a false cycle.** The graph is keyed by unique -//! id, and reachability is compared against the unique-node count rather than -//! the input length — otherwise a repeated id (already reported as -//! [`DagIssue::DuplicateNode`]) would look like an unreachable node and read -//! as a cycle. -//! - **A self-edge is a cycle.** `a` depending on `a` yields -//! [`DagIssue::Cycle`]. Callers that surface self-dependency as its own -//! diagnostic should test for it before calling. -//! -//! The algorithm is Kahn's: count indegrees, drain the zero-indegree frontier, -//! and compare the number of nodes drained against the number of nodes. Runs in -//! O(V + E), allocates only the working maps, and touches no graph runtime -//! state — this module is pure structure and is usable before anything is -//! compiled or executed. -//! -//! ``` -//! use tinyagents_graph::dag::{DagIssue, DagNode, validate_dag}; -//! -//! let nodes = vec![ -//! DagNode::new("plan", []), -//! DagNode::new("build", ["plan"]), -//! DagNode::new("ship", ["build"]), -//! ]; -//! assert!(validate_dag(&nodes).is_empty()); -//! -//! let looped = vec![DagNode::new("a", ["b"]), DagNode::new("b", ["a"])]; -//! assert_eq!(validate_dag(&looped), vec![DagIssue::Cycle]); -//! ``` - -use std::collections::{HashMap, HashSet, VecDeque}; - -mod types; - -#[cfg(test)] -#[path = "mod_tests.rs"] -mod test; - -pub use types::{DagIssue, DagNode}; - -/// Reports whether the dependency edges contain a cycle. -/// -/// Edges naming an id that is not a declared node are ignored (see the module -/// docs). Duplicate ids are collapsed to one node (first declaration wins), -/// so they cannot produce a false positive — but a caller that wants to add -/// edges to a node id that already exists must merge them into that one -/// declaration; a second `DagNode` with the same id has its edges silently -/// dropped instead of contributing to the cycle check (see the comment at -/// the dedupe below). -pub fn has_cycle(nodes: &[DagNode<'_>]) -> bool { - let ids: HashSet<&str> = nodes.iter().map(|n| n.id).collect(); - let mut indegree: HashMap<&str, usize> = ids.iter().map(|&id| (id, 0)).collect(); - let mut adjacency: HashMap<&str, Vec<&str>> = HashMap::new(); - - // A repeated id is reported separately as `DagIssue::DuplicateNode`; the - // cycle graph must still be built from exactly one declaration per id, or - // two conflicting declarations of the same id (e.g. `a[]` and `a[b]`) get - // merged into edges that no single declaration actually names, which can - // fabricate a cycle that would not exist under either declaration alone. - // - // Consequence for callers: a caller that wants to add dependency edges to - // a node that already exists in `nodes` must merge those edges into that - // node's single `DagNode` declaration, not append a second `DagNode` with - // the same id and the new edges — the second declaration's edges are - // silently dropped here (first declaration wins), so a real cycle they - // would have introduced goes undetected. This is unreachable in a caller - // that only ever mints fresh, never-reused ids (e.g. a fresh UUID per - // node), but is a live footgun for any caller that revisits an existing - // id. - let mut processed: HashSet<&str> = HashSet::with_capacity(ids.len()); - for node in nodes { - if !processed.insert(node.id) { - continue; - } - for &dep in &node.depends_on { - if ids.contains(dep) { - // Edge dep -> node: `node` runs after `dep`. - adjacency.entry(dep).or_default().push(node.id); - *indegree.entry(node.id).or_insert(0) += 1; - } - } - } - - let mut queue: VecDeque<&str> = indegree - .iter() - .filter(|&(_, &d)| d == 0) - .map(|(&n, _)| n) - .collect(); - let mut visited = 0usize; - while let Some(node) = queue.pop_front() { - visited += 1; - if let Some(children) = adjacency.get(node) { - for &child in children { - let entry = indegree.get_mut(child).expect("child in indegree"); - *entry -= 1; - if *entry == 0 { - queue.push_back(child); - } - } - } - } - - // Compare against the unique-id count, not `nodes.len()`: the graph is keyed - // by unique id, so duplicates (reported separately) must not read as - // unreachable nodes. - visited != indegree.len() -} - -/// Validates ids, edges and acyclicity in one pass, returning every issue found. -/// -/// Issues are ordered: duplicate ids in input order, then dangling edges in -/// input order, then a single [`DagIssue::Cycle`] if the resolvable edges close -/// a loop. An empty result means the input is a well-formed DAG. -pub fn validate_dag(nodes: &[DagNode<'_>]) -> Vec { - let mut issues = Vec::new(); - - let mut seen: HashSet<&str> = HashSet::new(); - for node in nodes { - if !seen.insert(node.id) { - issues.push(DagIssue::DuplicateNode { - id: node.id.to_string(), - }); - } - } - - for node in nodes { - for &dep in &node.depends_on { - if !seen.contains(dep) { - issues.push(DagIssue::UnknownDependency { - node: node.id.to_string(), - depends_on: dep.to_string(), - }); - } - } - } - - if has_cycle(nodes) { - issues.push(DagIssue::Cycle); - } - - issues -} diff --git a/crates/tinyagents-graph/src/dag/mod_tests.rs b/crates/tinyagents-graph/src/dag/mod_tests.rs deleted file mode 100644 index 14691c65e..000000000 --- a/crates/tinyagents-graph/src/dag/mod_tests.rs +++ /dev/null @@ -1,140 +0,0 @@ -//! Unit tests for dependency-DAG validation: acyclic acceptance, cycle and -//! self-edge detection, dangling edges, and the duplicate-id false-cycle guard. - -use super::*; - -#[test] -fn linear_chain_is_acyclic() { - let nodes = vec![ - DagNode::new("a", []), - DagNode::new("b", ["a"]), - DagNode::new("c", ["b"]), - ]; - assert!(!has_cycle(&nodes)); - assert!(validate_dag(&nodes).is_empty()); -} - -#[test] -fn diamond_is_acyclic() { - let nodes = vec![ - DagNode::new("root", []), - DagNode::new("left", ["root"]), - DagNode::new("right", ["root"]), - DagNode::new("join", ["left", "right"]), - ]; - assert!(!has_cycle(&nodes)); - assert!(validate_dag(&nodes).is_empty()); -} - -#[test] -fn empty_graph_is_acyclic() { - assert!(!has_cycle(&[])); - assert!(validate_dag(&[]).is_empty()); -} - -#[test] -fn two_node_cycle_is_detected() { - let nodes = vec![DagNode::new("a", ["b"]), DagNode::new("b", ["a"])]; - assert!(has_cycle(&nodes)); - assert_eq!(validate_dag(&nodes), vec![DagIssue::Cycle]); -} - -#[test] -fn longer_cycle_is_detected() { - let nodes = vec![ - DagNode::new("a", ["c"]), - DagNode::new("b", ["a"]), - DagNode::new("c", ["b"]), - ]; - assert!(has_cycle(&nodes)); -} - -#[test] -fn self_edge_is_a_cycle() { - let nodes = vec![DagNode::new("a", ["a"])]; - assert!(has_cycle(&nodes)); - assert_eq!(validate_dag(&nodes), vec![DagIssue::Cycle]); -} - -#[test] -fn cycle_off_to_the_side_is_still_detected() { - let nodes = vec![ - DagNode::new("start", []), - DagNode::new("x", ["y"]), - DagNode::new("y", ["x"]), - ]; - assert!(has_cycle(&nodes)); -} - -#[test] -fn dangling_edges_are_reported_and_ignored_by_the_cycle_check() { - let nodes = vec![DagNode::new("only", ["ghost"])]; - assert!(!has_cycle(&nodes)); - assert_eq!( - validate_dag(&nodes), - vec![DagIssue::UnknownDependency { - node: "only".to_string(), - depends_on: "ghost".to_string(), - }] - ); -} - -#[test] -fn duplicate_ids_are_reported_without_a_false_cycle() { - let nodes = vec![DagNode::new("a", []), DagNode::new("a", [])]; - assert!(!has_cycle(&nodes)); - assert_eq!( - validate_dag(&nodes), - vec![DagIssue::DuplicateNode { - id: "a".to_string() - }] - ); -} - -#[test] -fn conflicting_duplicate_declarations_do_not_fabricate_a_cycle() { - // `a` is declared twice with conflicting dependency lists ([] and [b]), - // and `b` depends on `a`. Merging both declarations of `a` would add an - // edge b -> a (from the second `a[b]`) alongside a -> b (from `b[a]`), - // fabricating a cycle that neither declaration of `a` alone produces: - // the first (`a[]`) yields a -> b only, and the second (`a[b]`) yields - // b -> a only. Only one declaration per id may contribute edges. - let nodes = vec![ - DagNode::new("a", []), - DagNode::new("a", ["b"]), - DagNode::new("b", ["a"]), - ]; - assert!(!has_cycle(&nodes)); - assert_eq!( - validate_dag(&nodes), - vec![DagIssue::DuplicateNode { - id: "a".to_string() - }] - ); -} - -#[test] -fn all_issue_kinds_are_reported_together() { - let nodes = vec![ - DagNode::new("a", ["b"]), - DagNode::new("b", ["a"]), - DagNode::new("b", ["ghost"]), - ]; - let issues = validate_dag(&nodes); - assert!(issues.contains(&DagIssue::DuplicateNode { - id: "b".to_string() - })); - assert!(issues.contains(&DagIssue::UnknownDependency { - node: "b".to_string(), - depends_on: "ghost".to_string(), - })); - assert!(issues.contains(&DagIssue::Cycle)); -} - -#[test] -fn node_view_borrows_owned_dependency_ids() { - let deps = ["a".to_string(), "b".to_string()]; - let node = DagNode::new("c", deps.iter().map(String::as_str)); - assert_eq!(node.id, "c"); - assert_eq!(node.depends_on, vec!["a", "b"]); -} diff --git a/crates/tinyagents-graph/src/dag/types.rs b/crates/tinyagents-graph/src/dag/types.rs deleted file mode 100644 index 66a680ca4..000000000 --- a/crates/tinyagents-graph/src/dag/types.rs +++ /dev/null @@ -1,67 +0,0 @@ -//! Types for structural validation of a dependency DAG. -//! -//! A [`DagNode`] is a borrowed view of one node — its id plus the ids it -//! declares a dependency on — so a host can validate its own node shape without -//! this crate learning anything about that shape. - -/// One node of a dependency graph: an id and the ids it depends on. -/// -/// Borrowed on purpose: callers hold their own node structs (workflow phases, -/// team tasks, plan steps) and project them into this view for the duration of -/// a validation call. Edges are read as `depends_on -> id`, i.e. a node runs -/// after everything it names. -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct DagNode<'a> { - /// This node's id. Ids are compared as strings and are expected to be - /// unique; a repeat is reported as [`DagIssue::DuplicateNode`]. - pub id: &'a str, - /// Ids this node depends on. An entry naming no declared node is reported - /// as [`DagIssue::UnknownDependency`] and is ignored by the cycle check. - pub depends_on: Vec<&'a str>, -} - -impl<'a> DagNode<'a> { - /// Builds a node view from an id and any iterator of dependency ids. - /// - /// ``` - /// use tinyagents_graph::dag::DagNode; - /// - /// let deps = vec!["a".to_string(), "b".to_string()]; - /// let node = DagNode::new("c", deps.iter().map(String::as_str)); - /// assert_eq!(node.depends_on, vec!["a", "b"]); - /// ``` - pub fn new(id: &'a str, depends_on: impl IntoIterator) -> Self { - Self { - id, - depends_on: depends_on.into_iter().collect(), - } - } -} - -/// A structural problem found in a dependency graph. -/// -/// Deliberately small: these are the graph-shaped facts only. Domain rules a -/// host layers on top — "a phase must name at least one agent", "a task may not -/// depend on itself by name" — stay with the host, which knows what to call -/// them. -#[derive(Debug, Clone, PartialEq, Eq)] -pub enum DagIssue { - /// Two nodes share the same id. - DuplicateNode { - /// The repeated id. - id: String, - }, - /// A node named a dependency that is not a declared node. - UnknownDependency { - /// The node declaring the edge. - node: String, - /// The id it named. - depends_on: String, - }, - /// The dependency edges contain at least one cycle. - /// - /// A self-edge (`a` depends on `a`) is a cycle and is reported this way; a - /// host that wants to distinguish it should check for that case itself - /// before calling. - Cycle, -} diff --git a/crates/tinyagents-graph/src/delegation/README.md b/crates/tinyagents-graph/src/delegation/README.md deleted file mode 100644 index bc5cafb2f..000000000 --- a/crates/tinyagents-graph/src/delegation/README.md +++ /dev/null @@ -1,102 +0,0 @@ -# `graph::delegation` - -A multi-stage sub-agent delegation pipeline expressed as a durable state graph. - -```text - plan ─▶ execute ─▶ review ──approved / budget spent──▶ [approval] ─▶ finalize ─▶ END - ▲ │ - └──────revise───────┘ -``` - -The `approval` node exists only when the run is human-gated -(`DelegationConfig::require_review_approval`). - -## Why it is here - -Chaining sub-agents ad hoc gives you no revision budget, no resume after a -crash, and no place to put a human in the loop. This module is that pipeline -written once against the graph runtime, so the routing, the bounded revision -loop, the checkpoint classification and the durable approval pause are one -reviewed implementation rather than a per-host reinvention. - -## Public surface - -| Item | Role | -| --- | --- | -| `run_delegation` | Run to completion, return the terminal `DelegationState`. The convenience wrapper for the non-gated shape, where the graph never interrupts. | -| `run_delegation_durable` | Run, reporting via `DelegationOutcome::pending` whether it finalized or parked on a human-approval interrupt. | -| `run_or_resume_delegation` | Classify the thread's latest checkpoint, then resume / return / start fresh accordingly. | -| `resume_delegation` | Deliver an approver's decision to a parked run. | -| `deny_decision` | The canonical deny value, for resume-on-timeout. | -| `delegation_graph_topology` | Structure-only export (nodes, edges, routing) for inspection. | - -## The stage worker is injected - -Every entry point takes a `run_stage` closure: - -```rust,ignore -Fn(DelegationStage, DelegationState) -> impl Future> -``` - -This module owns *when* a stage runs and *where its result routes*; it owns -nothing about *how* a stage runs. A host passes a closure dispatching each -`DelegationStage` to its own sub-agent runner; tests pass a deterministic mock, -which is what makes the orchestration mechanics unit-testable without a model. - -The same reasoning applies to observability: `DelegationConfig::event_sink` is -an optional host-supplied `GraphEventSink`. The module never reaches for a -host-specific logger. - -## `DelegationState` is an on-disk format - -This is the constraint to read before editing anything in `types.rs`. - -A checkpoint written by one release is decoded by a later one. So the serde -representation of `DelegationState` — field names, `#[serde(...)]` attributes, -defaults, and the resulting JSON — is a compatibility contract with every -installation that has a run in flight. - -`CURRENT_SCHEMA_VERSION` is what makes that contract enforceable rather than -merely documented: - -- a fresh run stamps its state with the current version; -- pre-versioned checkpoints decode to `0` via `#[serde(default)]`; -- `run_or_resume_delegation` **expires** any checkpoint below the current - version instead of resuming or returning it. - -That last point closes a gap a decode failure alone cannot: a stale record whose -fields still happen to decode (an empty `executions` list, say) is structurally -readable but semantically wrong for the current graph, and would otherwise be -resumed as if it were current. - -`mod_tests.rs` pins the exact serialized JSON of a fully-populated state, of the -default state, and of a pre-versioned record. A field addition that drifts the -shape fails those tests rather than reaching a user's disk unnoticed. Changing -the shape deliberately means bumping `CURRENT_SCHEMA_VERSION` in the same -change. - -## Bounds and failure handling - -- **Revision budget.** `max_revisions` caps reviewer-requested revisions; at the - cap the review is force-approved so a never-satisfied reviewer still - terminates. -- **Recursion policy.** A `RecursionPolicy` bounds visits per node and total - steps as a backstop to the in-state counter — the loop terminates even if the - counter logic is wrong. -- **Retry.** `max_attempts(1)` preserves single-attempt node semantics; the seam - is wired so raising it is a one-line change. -- **Cancellation** is checked at each node boundary and routes to `finalize`, - so a cancelled run still produces a state rather than an error. -- **Checkpoint read errors are split by kind.** A *decode* failure expires the - checkpoint and starts fresh; an *operational* failure (SQLite busy, I/O, a - poisoned lock) is propagated. Treating the second as the first would silently - discard valid durable work on a transient fault. - -## Durable pause vs. interactive pause - -The `approval` node is a **checkpointed graph interrupt**: the pause lives -entirely in the checkpointer and survives a process restart, and only -`resume_delegation` releases it. That is a durability mechanism for the pause — -it grants no approval authority and bypasses no host security boundary. A host's -own in-memory, TTL-bounded prompt for a live turn is a different mechanism and -is unaffected by anything here. diff --git a/crates/tinyagents-graph/src/delegation/graph.rs b/crates/tinyagents-graph/src/delegation/graph.rs deleted file mode 100644 index 40bcb8b37..000000000 --- a/crates/tinyagents-graph/src/delegation/graph.rs +++ /dev/null @@ -1,332 +0,0 @@ -//! Structure of the delegation graph: nodes, routing, the state reducer, and -//! the recursion/retry policy that bounds the `execute` ⇄ `review` loop. - -use std::future::Future; - -use serde_json::json; - -use super::run::decision_is_approve; -use super::types::{ - DelegationConfig, DelegationStage, DelegationStageOutput, DelegationState, DelegationUpdate, - StepRecord, -}; -use crate::export::GraphTopology; -use crate::recursion::RecursionPolicy; -use crate::{ - ClosureStateReducer, Command, CompiledGraph, END, GraphBuilder, Interrupt, NodeContext, - NodeResult, -}; -use tinyagents_harness::cancel::CancellationToken; -use tinyagents_harness::retry::RetryPolicy; - -/// Build (but do not run) the delegation `CompiledGraph`. Shared by -/// [`run_delegation`] and [`delegation_graph_topology`] so the graph's structure -/// has one definition. -pub(super) fn build_delegation_graph( - max_revisions: usize, - cancel: CancellationToken, - require_review_approval: bool, - run_stage: F, -) -> Result, String> -where - F: Fn(DelegationStage, DelegationState) -> Fut + Clone + Send + Sync + 'static, - Fut: Future> + Send + 'static, -{ - let mut builder = GraphBuilder::::new().set_reducer( - ClosureStateReducer::new(|mut s: DelegationState, u: DelegationUpdate| { - match u { - DelegationUpdate::Plan(p) => s.plan = Some(p), - DelegationUpdate::Execution { prompt, result } => { - let index = s.executions.len(); - s.executions.push(StepRecord { - index, - prompt, - result, - }); - } - DelegationUpdate::Review { note, approved } => { - s.reviews.push(note); - s.approved = approved; - if !approved { - s.revisions += 1; - } - } - DelegationUpdate::HumanDecision { approved } => { - s.human_approved = Some(approved); - s.denied = !approved; - // A denial overrides the reviewer's in-graph approval: the - // human gate is the final authority on whether the result - // may finalize. - if !approved { - s.approved = false; - } - } - DelegationUpdate::Final(f) => s.final_output = Some(f), - DelegationUpdate::Cancelled => s.cancelled = true, - } - Ok(s) - }), - ); - - // plan: produce the plan, then route to execute (or finalize if cancelled). - let run_plan = run_stage.clone(); - let cancel_plan = cancel.clone(); - builder = builder.add_node("plan", move |s: DelegationState, _c: NodeContext| { - let run_plan = run_plan.clone(); - let cancel = cancel_plan.clone(); - async move { - if cancel.is_cancelled() { - return Ok(NodeResult::Command( - Command::default() - .with_update(DelegationUpdate::Cancelled) - .with_goto(["finalize"]), - )); - } - let out = run_plan(DelegationStage::Plan, s) - .await - .map_err(to_node_err)?; - Ok(NodeResult::Command( - Command::default() - .with_update(DelegationUpdate::Plan(out.text)) - .with_goto(["execute"]), - )) - } - }); - - // execute: run the plan; route to review. - let run_exec = run_stage.clone(); - let cancel_exec = cancel.clone(); - builder = builder.add_node("execute", move |s: DelegationState, _c: NodeContext| { - let run_exec = run_exec.clone(); - let cancel = cancel_exec.clone(); - async move { - if cancel.is_cancelled() { - return Ok(NodeResult::Command( - Command::default() - .with_update(DelegationUpdate::Cancelled) - .with_goto(["finalize"]), - )); - } - let out = run_exec(DelegationStage::Execute, s) - .await - .map_err(to_node_err)?; - Ok(NodeResult::Command( - Command::default() - .with_update(DelegationUpdate::Execution { - prompt: out.prompt.unwrap_or_default(), - result: out.text, - }) - .with_goto(["review"]), - )) - } - }); - - // review: approve (→ finalize) or request a revision (→ execute), bounded by - // `max_revisions` so a never-approving reviewer still terminates. - let run_review = run_stage.clone(); - let cancel_review = cancel.clone(); - builder = builder.add_node("review", move |s: DelegationState, _c: NodeContext| { - let run_review = run_review.clone(); - let cancel = cancel_review.clone(); - async move { - if cancel.is_cancelled() { - return Ok(NodeResult::Command( - Command::default() - .with_update(DelegationUpdate::Cancelled) - .with_goto(["finalize"]), - )); - } - let revisions = s.revisions; - let out = run_review(DelegationStage::Review, s) - .await - .map_err(to_node_err)?; - // Approve when the reviewer is satisfied OR the revision budget is spent. - let approved = out.approved || revisions >= max_revisions; - // An approved result routes to the durable human-approval gate when - // the run is human-gated; otherwise it finalizes directly. A - // not-approved result always loops back to `execute` for a revision. - let next = if !approved { - "execute" - } else if require_review_approval { - "approval" - } else { - "finalize" - }; - Ok(NodeResult::Command( - Command::default() - .with_update(DelegationUpdate::Review { - note: out.text, - approved, - }) - .with_goto([next]), - )) - } - }); - - // approval (only when human-gated): a durable human-in-the-loop pause. - // - // First entry (`ctx.resume` is `None`): emit `NodeResult::Interrupt`. The - // executor persists a boundary checkpoint (Sync durability) and returns - // control to the caller — the pause now survives a process restart. Nothing - // finalizes until a resume arrives. - // - // Re-entry (`ctx.resume` is `Some(decision)`): the approver's decision was - // delivered via `Command { resume: .. }`. Apply approve/deny and route to - // `finalize` (deny is honoured there as a blocked/denied result). This is a - // durability mechanism for the PAUSE only — it grants no new approval - // authority and never bypasses the security/approval boundary. - // - // Durable-vs-chat boundary: this pause is a *checkpointed graph interrupt*, - // distinct from the interactive chat-turn approval gate (10-min TTL steering - // pause via `ApprovalRequestCard`), which parks a live in-memory chat turn - // and is deliberately left untouched. - if require_review_approval { - let cancel_approval = cancel.clone(); - builder = builder.add_node("approval", move |s: DelegationState, ctx: NodeContext| { - let cancel = cancel_approval.clone(); - async move { - // Cancellation must be honoured here exactly like every other - // node: without this check, a cancellation that arrives while - // a gated run is waiting for approval (or during the - // preceding review worker) is invisible to this boundary. A - // decisionless retry would then interrupt again indefinitely - // (nothing ever routes it to `finalize`), while an approving - // resume would finalize successfully despite the run having - // been cancelled. - if cancel.is_cancelled() { - return Ok(NodeResult::Command( - Command::default() - .with_update(DelegationUpdate::Cancelled) - .with_goto(["finalize"]), - )); - } - match ctx.resume { - None => { - let payload = json!({ - "kind": "delegation_review", - "reviews": s.reviews, - "executions": s.executions_texts(), - "revisions": s.revisions, - }); - tracing::info!( - revisions = s.revisions, - "[interrupt] delegation review reached durable human-approval gate; pausing" - ); - Ok(NodeResult::Interrupt(Interrupt::with_id( - "delegation-review-approval", - "approval", - payload, - ))) - } - Some(decision) => { - let approved = decision_is_approve(&decision); - tracing::info!( - approved, - "[interrupt] delegation review resumed with human decision" - ); - Ok(NodeResult::Command( - Command::default() - .with_update(DelegationUpdate::HumanDecision { approved }) - .with_goto(["finalize"]), - )) - } - } - } - }); - } - - // finalize: synthesize the final output from the accumulated state, then end. - builder = builder.add_node( - "finalize", - move |s: DelegationState, _c: NodeContext| async move { - let summary = s - .executions - .last() - .map(|r| r.result.clone()) - .unwrap_or_else(|| "".to_string()); - let final_text = if s.cancelled { - format!("cancelled after {} execution(s)", s.executions.len()) - } else if s.denied { - // `s.denied` is set only by `DelegationUpdate::HumanDecision` - // (the durable human-approval gate) — the in-graph reviewer - // never sets it: a not-approved review routes back to - // `execute` and only increments `revisions`. Attribute the - // denial to the actor that can actually produce it. - format!( - "denied by approver after {} execution(s)", - s.executions.len() - ) - } else { - summary - }; - Ok(NodeResult::Command( - Command::default() - .with_update(DelegationUpdate::Final(final_text)) - .with_goto([END]), - )) - }, - ); - - builder = builder - .set_entry("plan") - .mark_command_routing("plan") - .mark_command_routing("execute") - .mark_command_routing("review") - .mark_command_routing("finalize"); - - if require_review_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 - .compile() - .map_err(|e| format!("delegation graph compile failed: {e}"))? - // Bound the execute⇄review loop as a backstop to the in-state counter: - // each of execute/review may be visited at most max_revisions + 1 times. - .with_recursion_policy(RecursionPolicy { - max_visits_per_node: Some(max_revisions + 2), - max_total_steps: (max_revisions + 1) * 4 + 8, - ..RecursionPolicy::default() - }) - // Adapter-first landing of the crate-native per-node RetryPolicy - // (tinyagents 1.5.0 `CompiledGraph::with_node_retry`). Conservative: - // `max_attempts(1)` preserves today's single-attempt semantics exactly - // (no bespoke retry glue existed here) and backoff sleeping stays off - // (the default), so a transient node-handler failure surfaces as it does - // today. This wires the seam so raising the attempt cap / enabling - // backoff is a one-line, gated follow-up rather than a rewrite. - .with_node_retry(RetryPolicy::default().with_max_attempts(1)); - - Ok(graph) -} - -/// Structure-only [`GraphTopology`] of the delegation graph for debug / -/// inspection (issue #4249, Phase 4). Built with a no-op stub stage worker — -/// the topology exposes only node names, edges, and routing, never closure -/// bodies. -pub fn delegation_graph_topology() -> Result { - let graph = build_delegation_graph( - DelegationConfig::default().max_revisions, - CancellationToken::new(), - // Topology export uses the non-gated shape (the four revision-loop - // nodes); the durable `approval` interrupt node is additive and only - // present when a run is human-gated. - false, - |_stage, _state| async { Ok(DelegationStageOutput::done("")) }, - )?; - Ok(graph.topology()) -} - -/// Map an injected-stage error string into a graph node error. -fn to_node_err(e: String) -> crate::TinyAgentsError { - crate::TinyAgentsError::Model(e) -} diff --git a/crates/tinyagents-graph/src/delegation/mod.rs b/crates/tinyagents-graph/src/delegation/mod.rs deleted file mode 100644 index c23a1622b..000000000 --- a/crates/tinyagents-graph/src/delegation/mod.rs +++ /dev/null @@ -1,74 +0,0 @@ -//! Multi-stage sub-agent delegation as a durable, resumable state graph. -//! -//! Where a single harness call drives *one* agent turn, this module composes -//! several sub-agent stages into a checkpointed state machine — the graph-native -//! alternative to ad-hoc sub-agent chaining: -//! -//! ```text -//! plan ─▶ execute ─▶ review ──approved/maxed──▶ finalize ─▶ END -//! ▲ │ -//! └─────revise────────┘ -//! ``` -//! -//! It exercises most of the graph layer: -//! - **conditional routing** — `review` returns a [`Command`](crate::Command) -//! routing to `execute` (revise) or `finalize` (done) based on the stage result; -//! - **recursion bounds** — a [`RecursionPolicy`](crate::RecursionPolicy) -//! caps the `execute` ⇄ `review` loop as a backstop to the in-state -//! `revisions` counter; -//! - **durable checkpoint/resume** — an optional -//! [`Checkpointer`](crate::Checkpointer) persists the typed -//! [`DelegationState`] at every super-step boundary, so a crashed or paused -//! run resumes from its last node; -//! - **human-in-the-loop** — with [`DelegationConfig::require_review_approval`] -//! an approved review parks on a durable -//! [`Interrupt`](crate::Interrupt) that survives a process restart and -//! is released by [`resume_delegation`]; -//! - **cooperative cancellation** — a -//! [`CancellationToken`](tinyagents_harness::cancel::CancellationToken) short-circuits the pipeline -//! to `finalize` at the next node boundary. -//! -//! # The per-stage worker is injected -//! -//! [`run_delegation`] and friends take a `run_stage` closure, so this module owns -//! the orchestration mechanics and nothing about *how* a stage runs. A host -//! passes a closure dispatching each [`DelegationStage`] to its own sub-agent -//! runner; tests pass a deterministic mock. -//! -//! # [`DelegationState`] is an on-disk format -//! -//! Checkpointed state is read back by a later process — potentially a later -//! release. [`CURRENT_SCHEMA_VERSION`] stamps fresh runs and -//! [`run_or_resume_delegation`] expires any checkpoint below it rather than -//! misreading a stale shape. Changing a field name, a `#[serde(...)]` attribute, -//! or a default breaks resume for existing installations; `test.rs` pins the -//! exact serialized JSON so such a change fails loudly. See `README.md`. -//! -//! # Entry points -//! -//! - [`run_delegation`] — run to completion, return the terminal state. The -//! convenience wrapper for the non-gated shape. -//! - [`run_delegation_durable`] — run, reporting whether it finalized or parked -//! on a human-approval interrupt. -//! - [`run_or_resume_delegation`] — resume a thread's last checkpoint boundary -//! when one is live and compatible, else start fresh. -//! - [`resume_delegation`] — deliver an approver's decision to a parked run. -//! - [`delegation_graph_topology`] — structure-only export for inspection. - -mod graph; -mod run; -mod types; - -pub use graph::delegation_graph_topology; -pub use run::{ - deny_decision, resume_delegation, run_delegation, run_delegation_durable, - run_or_resume_delegation, -}; -pub use types::{ - CURRENT_SCHEMA_VERSION, DelegationConfig, DelegationOutcome, DelegationStage, - DelegationStageOutput, DelegationState, PendingApproval, StepRecord, -}; - -#[cfg(test)] -#[path = "mod_tests.rs"] -mod test; diff --git a/crates/tinyagents-graph/src/delegation/mod_tests.rs b/crates/tinyagents-graph/src/delegation/mod_tests.rs deleted file mode 100644 index 434b5ed32..000000000 --- a/crates/tinyagents-graph/src/delegation/mod_tests.rs +++ /dev/null @@ -1,1272 +0,0 @@ -//! Unit tests for the delegation graph: the revision loop and its budget, -//! cancellation, durable checkpoint/resume classification, the human-approval -//! interrupt, and the persisted shape of [`DelegationState`]. - -use std::future::Future; -use std::sync::atomic::{AtomicUsize, Ordering}; -use std::sync::{Arc, Mutex}; -use tinystoragedrivers_testkit::MemoryBackend; - -use serde::{Deserialize, Serialize}; -use serde_json::json; - -use super::run::{decision_is_approve, is_incompatible_checkpoint_error}; -use super::*; -use crate::Interrupt; -use crate::checkpoint::{Checkpoint, Checkpointer, PendingActivation}; -use tinyagents_harness::cancel::CancellationToken; - -/// A reviewer that rejects the first `reject_first` executions, then approves, -/// driving the execute⇄review revision loop. -fn flow_runner( - reject_first: usize, -) -> impl Fn( - DelegationStage, - DelegationState, -) -> std::pin::Pin> + Send>> -+ Clone -+ Send -+ Sync -+ 'static { - let reviews = Arc::new(AtomicUsize::new(0)); - move |stage, _state| { - let reviews = reviews.clone(); - Box::pin(async move { - match stage { - DelegationStage::Plan => Ok(DelegationStageOutput::done("PLAN")), - DelegationStage::Execute => Ok(DelegationStageOutput::done("EXEC")), - DelegationStage::Review => { - let n = reviews.fetch_add(1, Ordering::SeqCst); - Ok(DelegationStageOutput { - text: format!("review-{n}"), - approved: n >= reject_first, - prompt: None, - }) - } - } - }) - } -} - -#[tokio::test] -async fn approves_first_pass_no_revision() { - let state = run_delegation(DelegationConfig::default(), flow_runner(0)) - .await - .expect("runs"); - assert_eq!(state.plan.as_deref(), Some("PLAN")); - assert_eq!(state.executions.len(), 1, "one execution, no revision"); - assert_eq!(state.executions[0].index, 0); - assert_eq!(state.executions[0].result, "EXEC"); - assert_eq!(state.reviews.len(), 1); - assert_eq!(state.revisions, 0); - assert!(state.approved); - assert_eq!(state.final_output.as_deref(), Some("EXEC")); -} - -#[tokio::test] -async fn revises_then_approves() { - // Reject the first review → one revision (a second execute+review). - let state = run_delegation(DelegationConfig::default(), flow_runner(1)) - .await - .expect("runs"); - assert_eq!(state.executions.len(), 2, "initial + one revised execution"); - assert_eq!(state.reviews.len(), 2); - assert_eq!(state.revisions, 1); - assert!(state.approved); -} - -#[tokio::test] -async fn revision_budget_caps_a_never_approving_reviewer() { - // Reviewer never approves on its own; the max_revisions cap forces finalize. - let config = DelegationConfig { - max_revisions: 2, - ..DelegationConfig::default() - }; - let state = run_delegation(config, flow_runner(999)) - .await - .expect("runs"); - // revisions counted: 1st review (rev 1), 2nd review (rev 2), 3rd review - // hits revisions>=2 → forced approve. So 3 executions, 3 reviews. - assert_eq!(state.revisions, 2, "stops at the revision budget"); - assert!(state.approved, "forced-approved at the cap"); - assert_eq!(state.executions.len(), 3); -} - -#[tokio::test] -async fn cancellation_short_circuits_to_finalize() { - let cancel = CancellationToken::new(); - cancel.cancel(); - let ran = Arc::new(Mutex::new(Vec::::new())); - let ran2 = ran.clone(); - let runner = move |stage: DelegationStage, _s: DelegationState| { - let ran = ran2.clone(); - Box::pin(async move { - ran.lock().unwrap().push(stage); - Ok::<_, String>(DelegationStageOutput::done("X")) - }) as std::pin::Pin + Send>> - }; - let config = DelegationConfig { - cancel, - ..DelegationConfig::default() - }; - let state = run_delegation(config, runner).await.expect("runs"); - assert!(state.cancelled, "state flagged cancelled"); - assert!(state.final_output.is_some()); - assert!( - ran.lock().unwrap().is_empty(), - "no stage worker ran once cancelled at the plan boundary" - ); -} - -#[tokio::test] -async fn human_gated_run_parks_on_interrupt_then_resume_approves() { - let checkpoint_backend = MemoryBackend::new(); - let cp: Arc> = - Arc::new(crate::test_support::dir_checkpointer(&checkpoint_backend)); - let make_config = || DelegationConfig { - require_review_approval: true, - checkpointer: Some(cp.clone()), - thread_id: Some("hg-approve".to_string()), - ..DelegationConfig::default() - }; - - // First pass: reviewer approves on the first review, so the run reaches - // the durable human-approval gate and parks on an interrupt. - let outcome = run_delegation_durable(make_config(), flow_runner(0)) - .await - .expect("runs"); - 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" - ); - - // Reopen through a fresh graph and checkpointer handle over the same - // explicitly owned backend, then re-enter via Command::resume. - let resumed = resume_delegation(make_config(), json!("approve_once"), flow_runner(0)) - .await - .expect("resumes"); - assert!(resumed.pending.is_none(), "resume clears the pause"); - assert_eq!(resumed.state.human_approved, Some(true)); - assert!(!resumed.state.denied); - assert!( - resumed.state.final_output.is_some(), - "resumes from checkpoint to finalize" - ); -} - -#[tokio::test] -async fn concurrent_calls_for_the_same_thread_never_overlap() { - let checkpoint_backend = MemoryBackend::new(); - // Two concurrent `run_or_resume_delegation` calls for the SAME stable - // thread_id must not both read the checkpoint and dispatch independently: - // without a per-thread lock held across classification and the selected - // run/resume operation, both can observe the same starting point and run - // the same pending stage concurrently, duplicating execute-stage - // external side effects and producing competing checkpoint histories. - let concurrent = Arc::new(AtomicUsize::new(0)); - let max_concurrent = Arc::new(AtomicUsize::new(0)); - - let make_runner = || { - let concurrent = concurrent.clone(); - let max_concurrent = max_concurrent.clone(); - move |stage: DelegationStage, _s: DelegationState| { - let concurrent = concurrent.clone(); - let max_concurrent = max_concurrent.clone(); - Box::pin(async move { - let now = concurrent.fetch_add(1, Ordering::SeqCst) + 1; - max_concurrent.fetch_max(now, Ordering::SeqCst); - // Widen the race window: without the per-thread lock, the - // second call's stages would start executing while the - // first call is still mid-flight. - tokio::time::sleep(std::time::Duration::from_millis(20)).await; - concurrent.fetch_sub(1, Ordering::SeqCst); - let text = match stage { - DelegationStage::Review => "approve".to_string(), - _ => "ok".to_string(), - }; - Ok::<_, String>(DelegationStageOutput::done(text)) - }) as std::pin::Pin + Send>> - } - }; - - let cp: Arc> = - Arc::new(crate::test_support::dir_checkpointer(&checkpoint_backend)); - let make_config = || DelegationConfig { - checkpointer: Some(cp.clone()), - thread_id: Some("racing-thread".to_string()), - ..DelegationConfig::default() - }; - - let (r1, r2) = tokio::join!( - run_or_resume_delegation(make_config(), make_runner()), - run_or_resume_delegation(make_config(), make_runner()), - ); - r1.expect("first call runs"); - r2.expect("second call runs"); - - assert_eq!( - max_concurrent.load(Ordering::SeqCst), - 1, - "the per-thread lock must serialize the two calls' stage execution; \ - observed overlapping stage invocations" - ); -} - -#[tokio::test] -async fn run_delegation_rejects_require_review_approval() { - let checkpoint_backend = MemoryBackend::new(); - // `run_delegation` documents itself as the non-gated convenience wrapper - // that always returns a terminal state. With the gate on, a legitimately - // configured durable run can still return at the approval interrupt with - // `final_output` unset; discarding `DelegationOutcome::pending` (as - // `.state` alone does) would silently hand back an unfinished state with - // no signal that approval is pending. Reject the misuse instead of - // discarding it. - let cp: Arc> = - Arc::new(crate::test_support::dir_checkpointer(&checkpoint_backend)); - let config = DelegationConfig { - require_review_approval: true, - checkpointer: Some(cp), - thread_id: Some("hg-wrong-wrapper".to_string()), - ..DelegationConfig::default() - }; - let err = run_delegation(config, flow_runner(0)) - .await - .expect_err("must reject require_review_approval=true"); - assert!( - err.contains("require_review_approval"), - "error should name the misuse: {err}" - ); -} - -#[tokio::test] -async fn human_gated_without_checkpointer_or_thread_id_is_rejected() { - let checkpoint_backend = MemoryBackend::new(); - // `require_review_approval` documents that it needs both a checkpointer - // and a thread_id (interrupts require durability). Without this guard the - // run would still park on the approval interrupt with nothing persisted, - // filling `PendingApproval::thread_id` with an empty string and stranding - // the pause forever (`resume_delegation` has no checkpoint/thread to - // resume from). Both misconfigurations must be rejected up front. - let neither = DelegationConfig { - require_review_approval: true, - ..DelegationConfig::default() - }; - let err = run_delegation_durable(neither, flow_runner(0)) - .await - .expect_err("must reject: no checkpointer, no thread_id"); - assert!( - err.contains("require_review_approval"), - "error should name the misconfigured field: {err}" - ); - - let cp: Arc> = - Arc::new(crate::test_support::dir_checkpointer(&checkpoint_backend)); - let checkpointer_only = DelegationConfig { - require_review_approval: true, - checkpointer: Some(cp), - thread_id: None, - ..DelegationConfig::default() - }; - run_delegation_durable(checkpointer_only, flow_runner(0)) - .await - .expect_err("must reject: checkpointer without a thread_id"); - - let thread_only = DelegationConfig { - require_review_approval: true, - checkpointer: None, - thread_id: Some("hg-no-cp".to_string()), - ..DelegationConfig::default() - }; - run_delegation_durable(thread_only, flow_runner(0)) - .await - .expect_err("must reject: thread_id without a checkpointer"); -} - -#[tokio::test] -async fn ttl_expiry_resume_with_deny_blocks_the_result() { - let checkpoint_backend = MemoryBackend::new(); - let cp: Arc> = - Arc::new(crate::test_support::dir_checkpointer(&checkpoint_backend)); - let make_config = || DelegationConfig { - require_review_approval: true, - checkpointer: Some(cp.clone()), - thread_id: Some("hg-deny".to_string()), - ..DelegationConfig::default() - }; - - let outcome = run_delegation_durable(make_config(), flow_runner(0)) - .await - .expect("runs"); - assert!(outcome.pending.is_some(), "parks awaiting approval"); - - // TTL expiry → resume-with-deny preserves the timeout-deny behavior. - let resumed = resume_delegation(make_config(), deny_decision(), flow_runner(0)) - .await - .expect("resumes"); - assert_eq!(resumed.state.human_approved, Some(false)); - assert!(resumed.state.denied, "deny is honoured as a blocked result"); - assert!( - !resumed.state.approved, - "human deny overrides the reviewer's in-graph approval" - ); - assert!( - resumed - .state - .final_output - .as_deref() - .unwrap_or_default() - .contains("denied") - ); -} - -#[tokio::test] -async fn durable_checkpointer_persists_thread_state() { - let checkpoint_backend = MemoryBackend::new(); - let cp: Arc> = - Arc::new(crate::test_support::dir_checkpointer(&checkpoint_backend)); - let config = DelegationConfig { - checkpointer: Some(cp.clone()), - thread_id: Some("run-1".to_string()), - ..DelegationConfig::default() - }; - let state = run_delegation(config, flow_runner(1)).await.expect("runs"); - assert!(state.approved); - // The checkpointer recorded the run under its thread id. - let threads = cp.list_threads().await.expect("list threads"); - assert!( - threads.iter().any(|t| t == "run-1"), - "thread persisted, saw {threads:?}" - ); - let checkpoints = cp.list("run-1").await.expect("list checkpoints"); - assert!( - !checkpoints.is_empty(), - "at least one super-step boundary checkpoint persisted" - ); -} - -/// The boxed-future type every inline test runner returns. -type BoxedStageFut = - std::pin::Pin> + Send>>; - -#[tokio::test] -async fn execution_records_capture_index_and_prompt() { - // A worker that surfaces an execute prompt; assert the per-step record - // carries index + prompt + result (issue #3884). - let runner = move |stage: DelegationStage, _s: DelegationState| { - Box::pin(async move { - match stage { - DelegationStage::Plan => Ok(DelegationStageOutput::done("PLAN")), - DelegationStage::Execute => Ok(DelegationStageOutput { - text: "EXEC".to_string(), - approved: true, - prompt: Some("EXEC-PROMPT".to_string()), - }), - DelegationStage::Review => Ok(DelegationStageOutput { - text: "APPROVE".to_string(), - approved: true, - prompt: None, - }), - } - }) as BoxedStageFut - }; - let state = run_delegation(DelegationConfig::default(), runner) - .await - .expect("runs"); - assert_eq!(state.executions.len(), 1); - assert_eq!(state.executions[0].index, 0); - assert_eq!(state.executions[0].result, "EXEC"); - assert_eq!(state.executions[0].prompt, "EXEC-PROMPT"); -} - -#[test] -fn legacy_string_executions_do_not_deserialize_into_step_records() { - // A pre-#3884 checkpoint stored `executions` as a flat string array. It - // must fail to load into the new `Vec` — which is exactly - // what makes `run_or_resume_delegation` expire the stale checkpoint - // instead of misreading it. - let legacy = r#"{"plan":"P","executions":["raw a","raw b"],"reviews":[],"revisions":0,"approved":true,"final_output":null,"cancelled":false}"#; - assert!(serde_json::from_str::(legacy).is_err()); -} - -#[test] -fn schema_version_defaults_to_zero_and_step_records_round_trip() { - // A pre-versioned record (no `schema_version`) loads as version 0. - let unversioned = r#"{"plan":null,"executions":[],"reviews":[],"revisions":0,"approved":false,"final_output":null,"cancelled":false}"#; - let s: DelegationState = serde_json::from_str(unversioned).expect("loads"); - assert_eq!(s.schema_version, 0); - - // A fresh run stamps the current version and round-trips step records. - let mut fresh = DelegationState::new_run(); - assert_eq!(fresh.schema_version, CURRENT_SCHEMA_VERSION); - fresh.executions.push(StepRecord { - index: 0, - prompt: "P".to_string(), - result: "R".to_string(), - }); - let json = serde_json::to_string(&fresh).expect("serializes"); - let back: DelegationState = serde_json::from_str(&json).expect("round-trips"); - assert_eq!(back.schema_version, CURRENT_SCHEMA_VERSION); - assert_eq!(back.executions[0].result, "R"); - assert_eq!(back.executions[0].prompt, "P"); -} - -#[tokio::test] -async fn run_or_resume_starts_fresh_without_a_checkpoint() { - let checkpoint_backend = MemoryBackend::new(); - let cp: Arc> = - Arc::new(crate::test_support::dir_checkpointer(&checkpoint_backend)); - let config = DelegationConfig { - checkpointer: Some(cp), - thread_id: Some("fresh-1".to_string()), - ..DelegationConfig::default() - }; - let outcome = run_or_resume_delegation(config, flow_runner(0)) - .await - .expect("runs fresh"); - assert!(outcome.state.final_output.is_some()); - assert_eq!(outcome.state.executions.len(), 1); - assert_eq!(outcome.state.schema_version, CURRENT_SCHEMA_VERSION); -} - -#[tokio::test] -async fn run_or_resume_continues_from_last_boundary_after_a_crash() { - let checkpoint_backend = MemoryBackend::new(); - let cp: Arc> = - Arc::new(crate::test_support::dir_checkpointer(&checkpoint_backend)); - // Count stage invocations across BOTH the crashed run and the resume, to - // prove plan/execute are NOT re-run on resume. - let plan_runs = Arc::new(AtomicUsize::new(0)); - let exec_runs = Arc::new(AtomicUsize::new(0)); - let review_runs = Arc::new(AtomicUsize::new(0)); - let make_runner = |crash_first_review: bool| { - let plan_runs = plan_runs.clone(); - let exec_runs = exec_runs.clone(); - let review_runs = review_runs.clone(); - move |stage: DelegationStage, _s: DelegationState| { - let plan_runs = plan_runs.clone(); - let exec_runs = exec_runs.clone(); - let review_runs = review_runs.clone(); - Box::pin(async move { - match stage { - DelegationStage::Plan => { - plan_runs.fetch_add(1, Ordering::SeqCst); - Ok(DelegationStageOutput::done("PLAN")) - } - DelegationStage::Execute => { - exec_runs.fetch_add(1, Ordering::SeqCst); - Ok(DelegationStageOutput { - text: "EXEC".to_string(), - approved: true, - prompt: Some("EXEC-PROMPT".to_string()), - }) - } - DelegationStage::Review => { - let n = review_runs.fetch_add(1, Ordering::SeqCst); - if crash_first_review && n == 0 { - Err("simulated crash during review".to_string()) - } else { - Ok(DelegationStageOutput { - text: "APPROVE".to_string(), - approved: true, - prompt: None, - }) - } - } - } - }) as BoxedStageFut - } - }; - let make_config = || DelegationConfig { - checkpointer: Some(cp.clone()), - thread_id: Some("resume-1".to_string()), - ..DelegationConfig::default() - }; - - // Run 1: crashes during the first review. plan + execute completed and - // were checkpointed at their super-step boundaries; the run returns Err. - let crashed = run_or_resume_delegation(make_config(), make_runner(true)).await; - assert!(crashed.is_err(), "first run crashes at review"); - assert_eq!(plan_runs.load(Ordering::SeqCst), 1); - assert_eq!(exec_runs.load(Ordering::SeqCst), 1); - - // Run 2 (resume): a resumable checkpoint exists → re-enter at the pending - // node (review), NOT plan; plan/execute must not run again. - let resumed = run_or_resume_delegation(make_config(), make_runner(false)) - .await - .expect("resumes"); - assert!( - resumed.state.final_output.is_some(), - "resumed run finalizes" - ); - assert_eq!( - plan_runs.load(Ordering::SeqCst), - 1, - "plan not re-run on resume" - ); - assert_eq!( - exec_runs.load(Ordering::SeqCst), - 1, - "execute not re-run on resume" - ); - assert_eq!( - resumed.state.executions.len(), - 1, - "the pre-crash execution survived; no duplicate" - ); - assert_eq!(resumed.state.executions[0].result, "EXEC"); -} - -#[tokio::test] -async fn run_or_resume_is_idempotent_on_a_finalized_thread() { - let checkpoint_backend = MemoryBackend::new(); - let cp: Arc> = - Arc::new(crate::test_support::dir_checkpointer(&checkpoint_backend)); - let stage_runs = Arc::new(AtomicUsize::new(0)); - let make_runner = || { - let stage_runs = stage_runs.clone(); - move |stage: DelegationStage, _s: DelegationState| { - let stage_runs = stage_runs.clone(); - Box::pin(async move { - stage_runs.fetch_add(1, Ordering::SeqCst); - match stage { - DelegationStage::Review => Ok(DelegationStageOutput { - text: "APPROVE".to_string(), - approved: true, - prompt: None, - }), - _ => Ok(DelegationStageOutput::done("X")), - } - }) as BoxedStageFut - } - }; - let make_config = || DelegationConfig { - checkpointer: Some(cp.clone()), - thread_id: Some("done-1".to_string()), - ..DelegationConfig::default() - }; - - let first = run_or_resume_delegation(make_config(), make_runner()) - .await - .expect("runs"); - assert!(first.state.final_output.is_some()); - let after_first = stage_runs.load(Ordering::SeqCst); - assert!(after_first > 0, "the first run invoked stage workers"); - - // Re-invoke the SAME thread: a terminal checkpoint → return the finalized - // state without re-running any stage worker. - let second = run_or_resume_delegation(make_config(), make_runner()) - .await - .expect("idempotent"); - assert!(second.state.final_output.is_some()); - assert_eq!( - stage_runs.load(Ordering::SeqCst), - after_first, - "no stage worker re-ran on a finalized thread" - ); -} - -#[tokio::test] -async fn incompatible_checkpoint_expires_to_a_fresh_run() { - let checkpoint_backend = MemoryBackend::new(); - // A pre-#3884 checkpoint (executions as `Vec`) left in the store - // must not crash a resume: `run_or_resume_delegation` expires it and runs - // fresh, never panicking or surfacing the decode error. - #[derive(Clone, Debug, Serialize, Deserialize)] - struct LegacyState { - plan: Option, - executions: Vec, - reviews: Vec, - revisions: usize, - approved: bool, - final_output: Option, - cancelled: bool, - } - // Seed the store as the OLD state type under the thread. - let legacy_cp: crate::checkpoint::DriverCheckpointer = - crate::test_support::dir_checkpointer(&checkpoint_backend); - let legacy = Checkpoint::new( - LegacyState { - plan: Some("old".to_string()), - executions: vec!["a".to_string(), "b".to_string()], - reviews: vec![], - revisions: 0, - approved: true, - final_output: None, - cancelled: false, - }, - 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 - // undecodable record is expired and a fresh run completes. - let cp: Arc> = Arc::new( - crate::test_support::dir_checkpointer::(&checkpoint_backend), - ); - let config = DelegationConfig { - checkpointer: Some(cp), - thread_id: Some("legacy-1".to_string()), - ..DelegationConfig::default() - }; - let outcome = run_or_resume_delegation(config, flow_runner(0)) - .await - .expect("expires stale checkpoint and runs fresh"); - assert!(outcome.state.final_output.is_some(), "fresh run completed"); - assert_eq!(outcome.state.executions.len(), 1); - assert_eq!(outcome.state.executions[0].result, "EXEC"); -} - -#[tokio::test] -async fn checkpoint_below_current_schema_version_expires_to_fresh_run() { - let checkpoint_backend = MemoryBackend::new(); - // A decodable but OLD-schema checkpoint (schema_version defaults to 0 — - // e.g. a pre-#3884 record whose executions happened to be empty) must be - // expired, not resumed: `schema_version` is a real guard, not just a doc. - let seed: crate::checkpoint::DriverCheckpointer = - crate::test_support::dir_checkpointer(&checkpoint_backend); - let checkpoint = Checkpoint::new( - DelegationState { - plan: Some("stale".to_string()), - ..Default::default() - }, - 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" - ); - seed.put(checkpoint) - .await - .expect("seed old-schema checkpoint"); - - let cp: Arc> = Arc::new( - crate::test_support::dir_checkpointer::(&checkpoint_backend), - ); - let config = DelegationConfig { - checkpointer: Some(cp), - thread_id: Some("old-schema".to_string()), - ..DelegationConfig::default() - }; - let outcome = run_or_resume_delegation(config, flow_runner(0)) - .await - .expect("expires + fresh"); - assert!(outcome.state.final_output.is_some(), "fresh run completed"); - assert_eq!( - outcome.state.schema_version, CURRENT_SCHEMA_VERSION, - "fresh run stamped the current version" - ); - assert_eq!( - outcome.state.plan.as_deref(), - Some("PLAN"), - "re-planned from scratch, not resumed with the stale plan" - ); -} - -#[tokio::test] -async fn checkpoint_above_current_schema_version_also_expires_to_fresh_run() { - let checkpoint_backend = MemoryBackend::new(); - // A checkpoint stamped with a schema version NEWER than this binary - // understands (e.g. written by a newer binary during a rollback, or in a - // mixed-version deployment) must be expired exactly like an older one, not - // resumed. Serde can decode a newer record just fine by ignoring fields it - // doesn't recognize, so an inequality check — not merely `<` — is what - // catches this: resuming it under this binary's older semantics would be - // silently wrong rather than loudly incompatible. - let seed: crate::checkpoint::DriverCheckpointer = - crate::test_support::dir_checkpointer(&checkpoint_backend); - let checkpoint = Checkpoint::new( - DelegationState { - plan: Some("from-the-future".to_string()), - schema_version: CURRENT_SCHEMA_VERSION + 1, - ..Default::default() - }, - 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"); - - let cp: Arc> = Arc::new( - crate::test_support::dir_checkpointer::(&checkpoint_backend), - ); - let config = DelegationConfig { - checkpointer: Some(cp), - thread_id: Some("future-schema".to_string()), - ..DelegationConfig::default() - }; - let outcome = run_or_resume_delegation(config, flow_runner(0)) - .await - .expect("expires + fresh"); - assert!(outcome.state.final_output.is_some(), "fresh run completed"); - assert_eq!( - outcome.state.schema_version, CURRENT_SCHEMA_VERSION, - "fresh run stamped the current version, not the future one" - ); - assert_eq!( - outcome.state.plan.as_deref(), - Some("PLAN"), - "re-planned from scratch, not resumed with the future-schema plan" - ); -} - -#[test] -fn decision_is_approve_rejects_unrecognized_string_prefixes() { - // Object-shaped decisions must be checked against the SAME allowlist as - // string-shaped ones, not a `starts_with("approve")` prefix check: an - // unvalidated JSON value like `{"decision":"approve_not_authorized"}` - // must not release the durable human-approval gate just because it - // begins with "approve". - assert!(!decision_is_approve( - &json!({"decision": "approve_not_authorized"}) - )); - assert!(!decision_is_approve( - &json!({"decision": "approved_by_nobody"}) - )); - assert!(!decision_is_approve(&json!("approve_not_authorized"))); - - // The real allowlisted spellings still approve, in both shapes. - assert!(decision_is_approve(&json!("approve_once"))); - assert!(decision_is_approve(&json!("approve_always_for_tool"))); - assert!(decision_is_approve(&json!({"decision": "approve_once"}))); - assert!(decision_is_approve(&json!({"approved": true}))); - assert!(decision_is_approve(&json!(true))); - - // Denials and unrecognized shapes stay denials. - assert!(!decision_is_approve(&json!("deny"))); - assert!(!decision_is_approve(&json!(false))); - assert!(!decision_is_approve(&json!(null))); -} - -#[tokio::test] -async fn cancelled_checkpoint_still_scheduling_finalize_is_resumed_not_terminal() { - let checkpoint_backend = MemoryBackend::new(); - // Every cancellation route goes through `goto finalize` (see - // `graph.rs`), so a checkpoint can be `cancelled == true` with - // `final_output == None` and `next_nodes == ["finalize"]` if the process - // stopped after the cancellation update was checkpointed but before - // `finalize` ran. Treating `cancelled` as its own terminal signal - // classified that live boundary as terminal and returned it unfinished, - // never producing the cancellation summary. - let seed: crate::checkpoint::DriverCheckpointer = - crate::test_support::dir_checkpointer(&checkpoint_backend); - let checkpoint = Checkpoint::new( - DelegationState { - plan: Some("PLAN".to_string()), - cancelled: true, - schema_version: CURRENT_SCHEMA_VERSION, - ..Default::default() - }, - 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"); - - let cp: Arc> = Arc::new( - crate::test_support::dir_checkpointer::(&checkpoint_backend), - ); - let config = DelegationConfig { - checkpointer: Some(cp), - thread_id: Some("cancelled-mid-flight".to_string()), - ..DelegationConfig::default() - }; - let outcome = run_or_resume_delegation(config, flow_runner(0)) - .await - .expect("resumes to finalize"); - assert!( - outcome.state.final_output.is_some(), - "must resume through finalize and produce a cancellation summary, not \ - return the unfinished checkpoint verbatim" - ); - assert!(outcome.state.cancelled); -} - -#[tokio::test] -async fn cancellation_at_the_approval_node_routes_to_finalize() { - let checkpoint_backend = MemoryBackend::new(); - // Cancellation arriving while a gated run is parked at the approval - // interrupt (or during the preceding review worker) must be honoured at - // that boundary exactly like every other node. Without the check, a - // decisionless retry interrupts again indefinitely, and an approving - // resume finalizes successfully despite the cancellation. - let cancel = CancellationToken::new(); - let cp: Arc> = - Arc::new(crate::test_support::dir_checkpointer(&checkpoint_backend)); - let config = DelegationConfig { - require_review_approval: true, - checkpointer: Some(cp.clone()), - thread_id: Some("cancel-at-approval".to_string()), - cancel: cancel.clone(), - ..DelegationConfig::default() - }; - let outcome = run_delegation_durable(config, flow_runner(0)) - .await - .expect("parks on approval"); - outcome.pending.expect("parked on the approval interrupt"); - assert!( - outcome.state.final_output.is_none(), - "not finalized yet — still awaiting approval" - ); - - // Cancel while parked, then resume with an approval decision: the - // cancellation must win, not the approval. - cancel.cancel(); - let resumed_config = DelegationConfig { - require_review_approval: true, - checkpointer: Some(cp), - thread_id: Some("cancel-at-approval".to_string()), - cancel, - ..DelegationConfig::default() - }; - let resumed = resume_delegation(resumed_config, json!("approve_once"), flow_runner(0)) - .await - .expect("resumes"); - assert!(resumed.pending.is_none(), "resume clears the pause"); - assert!( - resumed.state.cancelled, - "cancellation must be honoured at the approval boundary, not silently \ - overridden by an approving resume" - ); - assert!( - resumed.state.final_output.is_some(), - "cancellation routes to finalize, producing a cancellation summary" - ); -} - -#[tokio::test] -async fn concurrent_resume_delegation_calls_for_the_same_thread_never_overlap() { - let checkpoint_backend = MemoryBackend::new(); - // The approval interrupt's resume path routes straight from `approval` to - // `finalize` WITHOUT calling `run_stage` at all (`finalize` only reads - // accumulated state), so a `run_stage`-based concurrency probe cannot - // observe the lock during `resume_delegation` — it only ever sees the - // SETUP `run_delegation_durable` call's stage execution. Instrument the - // checkpointer's `get` instead, since that IS on `resume_delegation`'s - // critical path (the schema-mismatch check reads the checkpoint before - // dispatching), and is where two racing resumes would actually overlap - // without the per-thread lock. - struct BlockingGetCheckpointer { - inner: Arc>, - concurrent: Arc, - max_concurrent: Arc, - } - - #[async_trait::async_trait] - impl Checkpointer for BlockingGetCheckpointer { - async fn put( - &self, - checkpoint: Checkpoint, - ) -> crate::Result { - self.inner.put(checkpoint).await - } - - async fn get( - &self, - thread_id: &str, - checkpoint_id: Option<&str>, - ) -> crate::Result>> { - let now = self.concurrent.fetch_add(1, Ordering::SeqCst) + 1; - self.max_concurrent.fetch_max(now, Ordering::SeqCst); - tokio::time::sleep(std::time::Duration::from_millis(20)).await; - let result = self.inner.get(thread_id, checkpoint_id).await; - self.concurrent.fetch_sub(1, Ordering::SeqCst); - result - } - - async fn list( - &self, - thread_id: &str, - ) -> crate::Result> { - self.inner.list(thread_id).await - } - - async fn list_threads(&self) -> crate::Result> { - self.inner.list_threads().await - } - - async fn delete_thread(&self, thread_id: &str) -> crate::Result<()> { - self.inner.delete_thread(thread_id).await - } - - async fn delete_checkpoints( - &self, - thread_id: &str, - ids: &[String], - ) -> crate::Result { - self.inner.delete_checkpoints(thread_id, ids).await - } - } - - let real_cp: Arc> = - Arc::new(crate::test_support::dir_checkpointer(&checkpoint_backend)); - - // Park on the approval interrupt first, through the real checkpointer. - let setup_config = DelegationConfig { - require_review_approval: true, - checkpointer: Some(real_cp.clone()), - thread_id: Some("racing-resume-thread".to_string()), - ..DelegationConfig::default() - }; - run_delegation_durable(setup_config, flow_runner(0)) - .await - .expect("parks on approval") - .pending - .expect("parked on the approval interrupt"); - - let concurrent = Arc::new(AtomicUsize::new(0)); - let max_concurrent = Arc::new(AtomicUsize::new(0)); - let wrapped: Arc> = Arc::new(BlockingGetCheckpointer { - inner: real_cp, - concurrent: concurrent.clone(), - max_concurrent: max_concurrent.clone(), - }); - let make_config = || DelegationConfig { - require_review_approval: true, - checkpointer: Some(wrapped.clone()), - thread_id: Some("racing-resume-thread".to_string()), - ..DelegationConfig::default() - }; - - let (r1, r2) = tokio::join!( - resume_delegation(make_config(), json!("approve_once"), flow_runner(0)), - resume_delegation(make_config(), json!("approve_once"), flow_runner(0)), - ); - // Exactly one of the two racing resumes may succeed in applying the - // decision; a second resume of an already-finalized thread is expected to - // either no-op or surface a graph-level error — either is acceptable, but - // overlapping `get` calls across the two resumes is not: that is exactly - // the race the per-thread lock exists to prevent. - assert!(r1.is_ok() || r2.is_ok(), "at least one resume must succeed"); - assert_eq!( - max_concurrent.load(Ordering::SeqCst), - 1, - "the per-thread lock must serialize the two resumes' checkpoint reads; \ - observed overlapping `get` calls" - ); -} -#[tokio::test] -async fn resume_delegation_rejects_a_schema_mismatched_checkpoint() { - let checkpoint_backend = MemoryBackend::new(); - // `resume_delegation` is a public entry point a host calls directly with - // an approver's decision — it bypasses `run_or_resume_delegation`'s - // match arms entirely, so `CompiledGraph::resume` would otherwise apply - // the decision to a schema-mismatched checkpoint with no check at all. - // During a rollback/mixed-version deployment, serde can decode a newer - // additive checkpoint by ignoring unknown fields, letting an older - // binary consume the approval under outdated semantics. - let seed: crate::checkpoint::DriverCheckpointer = - crate::test_support::dir_checkpointer(&checkpoint_backend); - let checkpoint = Checkpoint::new( - DelegationState { - plan: Some("PLAN".to_string()), - schema_version: CURRENT_SCHEMA_VERSION + 1, - ..Default::default() - }, - vec![PendingActivation { - node: tinyagents_harness::ids::NodeId::from("approval"), - send_arg: None, - task_id: tinyagents_harness::ids::TaskId::from(String::new()), - }], - ) - .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"); - - let cp: Arc> = Arc::new( - crate::test_support::dir_checkpointer::(&checkpoint_backend), - ); - let config = DelegationConfig { - require_review_approval: true, - checkpointer: Some(cp), - thread_id: Some("resume-future-schema".to_string()), - ..DelegationConfig::default() - }; - let err = resume_delegation(config, json!("approve_once"), flow_runner(0)) - .await - .expect_err("must refuse to resume a schema-mismatched checkpoint"); - assert!( - err.contains("schema_version"), - "error should name the mismatch: {err}" - ); -} - -#[test] -fn incompatible_checkpoint_error_matches_schema_not_corrupt_or_operational() { - use crate::TinyAgentsError; - // Positively identified schema mismatch → safe to expire. - assert!(is_incompatible_checkpoint_error( - &TinyAgentsError::Checkpoint( - "sqlite checkpointer: decode [schema] record: invalid type: string".to_string() - ) - )); - assert!(is_incompatible_checkpoint_error( - &TinyAgentsError::Checkpoint( - "sqlite checkpointer: decode [schema] next_nodes: missing field `foo`".to_string() - ) - )); - // Ambiguous/likely corruption must NOT be treated as a safe-to-expire - // schema mismatch — deleting it would discard the only evidence of the - // corruption and silently restart the delegation from `plan`, potentially - // repeating already-completed execute-stage side effects. - assert!(!is_incompatible_checkpoint_error( - &TinyAgentsError::Checkpoint( - "sqlite checkpointer: decode [corrupt] record: EOF while parsing a value".to_string() - ) - )); - assert!(!is_incompatible_checkpoint_error( - &TinyAgentsError::Checkpoint( - "file checkpointer: decode [corrupt] record: expected `,` or `}`".to_string() - ) - )); - // Operational failures must NOT be treated as incompatible (they must - // propagate, not silently restart durable work). - assert!(!is_incompatible_checkpoint_error( - &TinyAgentsError::Checkpoint( - "sqlite checkpointer: query latest checkpoint: database is locked".to_string() - ) - )); - assert!(!is_incompatible_checkpoint_error( - &TinyAgentsError::Checkpoint("sqlite checkpointer: connection lock poisoned".to_string()) - )); - assert!(!is_incompatible_checkpoint_error(&TinyAgentsError::Resume( - "no checkpoint".to_string() - ))); -} - -#[test] -fn decode_json_err_only_tags_schema_for_the_record_context() { - use crate::checkpoint::decode_json_err; - - // `Category::Data` errors decoding checkpoint-internal metadata - // ("namespace", "next_nodes", "header", "write payload") must NEVER be - // tagged `[schema]` — those fields have no versioned shape of their own, - // so a Data-category failure there is corruption by construction, not a - // caller's schema evolution. - let data_err = serde_json::from_str::("\"not a number\"").unwrap_err(); - assert_eq!(data_err.classify(), serde_json::error::Category::Data); - for what in ["namespace", "next_nodes", "header", "write payload"] { - let wrapped = decode_json_err("sqlite checkpointer", what, { - serde_json::from_str::("\"not a number\"").unwrap_err() - }); - assert!( - format!("{wrapped}").contains(&format!("decode [corrupt] {what}")), - "non-record contexts must always tag [corrupt], even for a \ - Data-category error: {wrapped}" - ); - } - let _ = data_err; - - // The "record" context is the only one eligible for [schema]. - let wrapped = decode_json_err( - "sqlite checkpointer", - "record", - serde_json::from_str::("\"not a number\"").unwrap_err(), - ); - assert!(format!("{wrapped}").contains("decode [schema] record")); -} - -#[test] -fn decode_json_err_classifies_data_errors_as_schema_and_others_as_corrupt() { - use crate::checkpoint::decode_json_err; - - // `Category::Data`: syntactically valid JSON that doesn't match the - // target type — the shape a legacy or newer schema produces. - let data_err = serde_json::from_str::("\"not a number\"").unwrap_err(); - assert_eq!(data_err.classify(), serde_json::error::Category::Data); - let wrapped = decode_json_err("sqlite checkpointer", "record", data_err); - assert!( - format!("{wrapped}").contains("decode [schema] record"), - "data-category decode errors must be tagged [schema]: {wrapped}" - ); - - // `Category::Eof`: the bytes were truncated — real corruption, not a - // schema difference. - let eof_err = serde_json::from_str::("{\"a\":").unwrap_err(); - assert_eq!(eof_err.classify(), serde_json::error::Category::Eof); - let wrapped = decode_json_err("file checkpointer", "record", eof_err); - assert!( - format!("{wrapped}").contains("decode [corrupt] record"), - "eof-category decode errors must be tagged [corrupt]: {wrapped}" - ); - - // `Category::Syntax`: malformed bytes — also corruption, not a schema - // difference. - let syntax_err = serde_json::from_str::("{not json}").unwrap_err(); - assert_eq!(syntax_err.classify(), serde_json::error::Category::Syntax); - let wrapped = decode_json_err("sqlite checkpointer", "record", syntax_err); - assert!( - format!("{wrapped}").contains("decode [corrupt] record"), - "syntax-category decode errors must be tagged [corrupt]: {wrapped}" - ); -} - -#[tokio::test] -async fn terminal_checkpoint_with_a_pending_interrupt_surfaces_it() { - let checkpoint_backend = MemoryBackend::new(); - // A terminal-classified checkpoint that still carries an interrupt must - // surface it, not silently drop `pending`. (The live routing never - // produces this shape; the terminal branch is defensive.) - let seed: crate::checkpoint::DriverCheckpointer = - crate::test_support::dir_checkpointer(&checkpoint_backend); - let mut state = DelegationState::new_run(); - state.final_output = Some("done".to_string()); - 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" }), - )]) - .with_metadata(json!({})); - seed.put(checkpoint).await.expect("seed terminal+interrupt"); - - let cp: Arc> = Arc::new( - crate::test_support::dir_checkpointer::(&checkpoint_backend), - ); - let config = DelegationConfig { - checkpointer: Some(cp), - thread_id: Some("terminal-interrupt".to_string()), - ..DelegationConfig::default() - }; - let outcome = run_or_resume_delegation(config, flow_runner(0)) - .await - .expect("terminal"); - let pending = outcome - .pending - .expect("the carried interrupt is surfaced, not dropped"); - assert_eq!(pending.node, "approval"); - assert_eq!(pending.interrupt_id, "intr-1"); - assert_eq!(outcome.state.final_output.as_deref(), Some("done")); -} - -/// Pins the exact serialized shape of a fully-populated [`DelegationState`]. -/// -/// `DelegationState` is an **on-disk checkpoint format**: installations carry -/// persisted state written by an earlier release, and a resume decodes it with -/// whatever the current binary declares. A renamed field, a changed -/// `#[serde(...)]` attribute, a new field without `#[serde(default)]`, or a -/// reordering that changes the emitted JSON silently breaks resume for those -/// installations — the checkpoint decodes into something subtly different, or -/// fails to decode at all and is expired, discarding in-flight work. -/// -/// The literal below is the byte-for-byte output of the implementation this -/// module was moved from, captured before the move. It is not merely a snapshot -/// of current behaviour: it is the compatibility contract. A change here must be -/// deliberate and paired with a [`CURRENT_SCHEMA_VERSION`] bump so -/// `run_or_resume_delegation` expires older checkpoints rather than misreading -/// them. -#[test] -fn serialized_state_shape_is_pinned() { - let state = DelegationState { - plan: Some("PLAN".into()), - executions: vec![ - StepRecord { - index: 0, - prompt: "P0".into(), - result: "R0".into(), - }, - StepRecord { - index: 1, - prompt: "P1".into(), - result: "R1".into(), - }, - ], - reviews: vec!["review-0".into(), "review-1".into()], - revisions: 1, - approved: true, - final_output: Some("FINAL".into()), - cancelled: false, - human_approved: Some(true), - denied: false, - schema_version: CURRENT_SCHEMA_VERSION, - }; - - let json = serde_json::to_string(&state).expect("serializes"); - assert_eq!( - json, - r#"{"plan":"PLAN","executions":[{"index":0,"prompt":"P0","result":"R0"},{"index":1,"prompt":"P1","result":"R1"}],"reviews":["review-0","review-1"],"revisions":1,"approved":true,"final_output":"FINAL","cancelled":false,"human_approved":true,"denied":false,"schema_version":1}"#, - "DelegationState's serialized shape changed; see this test's doc comment" - ); - - // And it decodes back to an equal value, so the pin covers both directions. - let back: DelegationState = serde_json::from_str(&json).expect("round-trips"); - assert_eq!(back.plan, state.plan); - assert_eq!(back.executions, state.executions); - assert_eq!(back.reviews, state.reviews); - assert_eq!(back.revisions, state.revisions); - assert_eq!(back.approved, state.approved); - assert_eq!(back.final_output, state.final_output); - assert_eq!(back.cancelled, state.cancelled); - assert_eq!(back.human_approved, state.human_approved); - assert_eq!(back.denied, state.denied); - assert_eq!(back.schema_version, state.schema_version); -} - -/// A pre-versioned checkpoint — one written before `schema_version`, -/// `human_approved` and `denied` existed — must still decode, taking the -/// documented defaults. This is the half of the contract that lets -/// `run_or_resume_delegation` *classify* a stale checkpoint (version `0`) rather -/// than failing to read it at all. -#[test] -fn pre_versioned_state_decodes_with_documented_defaults() { - let legacy = r#"{"plan":"PLAN","executions":[],"reviews":[],"revisions":0,"approved":false,"final_output":null,"cancelled":false}"#; - let state: DelegationState = serde_json::from_str(legacy).expect("decodes"); - assert_eq!( - state.schema_version, 0, - "defaults below CURRENT, so it expires" - ); - assert_eq!(state.human_approved, None); - assert!(!state.denied); - assert!(state.schema_version < CURRENT_SCHEMA_VERSION); -} - -/// The default state (a fresh, unstarted run) also has a pinned shape — it is -/// what the very first checkpoint of a run serializes from. -#[test] -fn default_state_shape_is_pinned() { - let json = serde_json::to_string(&DelegationState::default()).expect("serializes"); - assert_eq!( - json, - r#"{"plan":null,"executions":[],"reviews":[],"revisions":0,"approved":false,"final_output":null,"cancelled":false,"human_approved":null,"denied":false,"schema_version":0}"# - ); -} diff --git a/crates/tinyagents-graph/src/delegation/run.rs b/crates/tinyagents-graph/src/delegation/run.rs deleted file mode 100644 index 05943be92..000000000 --- a/crates/tinyagents-graph/src/delegation/run.rs +++ /dev/null @@ -1,514 +0,0 @@ -//! Driving the delegation graph: fresh runs, durable resumes, and the -//! checkpoint classification that decides between the two. - -use std::future::Future; -use std::sync::{Arc, OnceLock}; - -use serde_json::{Value, json}; -use tokio::sync::Mutex; - -use super::graph::build_delegation_graph; -use super::types::{ - CURRENT_SCHEMA_VERSION, DelegationConfig, DelegationOutcome, DelegationStage, - DelegationStageOutput, DelegationState, DelegationUpdate, PendingApproval, -}; -use crate::checkpoint::{Checkpoint, Checkpointer}; -use crate::execution_locks::RunLockMap; -use crate::{Command, END}; - -/// Serializes [`run_or_resume_delegation`]'s classify-then-dispatch critical -/// section per thread (see the lock acquisition there for why). Same pattern -/// as `graph::goals::store::thread_lock` / `graph::todos::store::thread_lock` -/// — a weak-value map so an idle thread's mutex is reclaimed instead of -/// leaking for the life of the process. -fn thread_lock(thread_id: &str) -> Arc> { - static LOCKS: OnceLock = OnceLock::new(); - LOCKS.get_or_init(RunLockMap::default).lock_for(thread_id) -} - -/// Run the plan→execute⇄review→finalize delegation graph, invoking `run_stage` -/// for each stage. Returns the final [`DelegationState`]. -/// -/// `run_stage` is the seam to the agent harness: production passes a closure that -/// dispatches each [`DelegationStage`] to `run_subagent`; tests pass a mock. -/// -/// This is the non-gated convenience wrapper: with the default config -/// (`require_review_approval = false`) the graph never interrupts, so the -/// returned state is always terminal. -/// -/// Rejects `require_review_approval = true` rather than silently discarding -/// [`DelegationOutcome::pending`]: with the gate on, a legitimately -/// configured run can still return at the durable approval interrupt with -/// `final_output` unset. Returning only `.state` here would hand the caller -/// an unfinished state with no indication that approval is pending — the -/// caller loses the one signal that would tell it to call -/// [`resume_delegation`]. Use [`run_delegation_durable`] (or -/// [`run_or_resume_delegation`]) when the gate is on, and read -/// `DelegationOutcome::pending`. -pub async fn run_delegation( - config: DelegationConfig, - run_stage: F, -) -> Result -where - F: Fn(DelegationStage, DelegationState) -> Fut + Clone + Send + Sync + 'static, - Fut: Future> + Send + 'static, -{ - if config.require_review_approval { - return Err( - "run_delegation does not support require_review_approval=true (it would \ - silently discard a pending durable approval interrupt); call \ - run_delegation_durable or run_or_resume_delegation instead and read \ - DelegationOutcome::pending" - .to_string(), - ); - } - Ok(run_delegation_durable(config, run_stage).await?.state) -} - -/// Run the delegation graph and report whether it finalized or parked on a -/// durable human-approval interrupt. -/// -/// When [`DelegationConfig::require_review_approval`] is set and the reviewer -/// approves, the `approval` node emits [`NodeResult::Interrupt`](crate::NodeResult::Interrupt); the executor -/// persists a checkpoint (Sync durability — the crate default) and returns -/// control here with the interrupt in [`DelegationOutcome::pending`]. Deliver the -/// approver's decision later with [`resume_delegation`] — it may run after a -/// process restart, since the pause lives entirely in the checkpointer. -pub async fn run_delegation_durable( - config: DelegationConfig, - run_stage: F, -) -> Result -where - F: Fn(DelegationStage, DelegationState) -> Fut + Clone + Send + Sync + 'static, - Fut: Future> + Send + 'static, -{ - // `require_review_approval` documents that it needs a checkpointer + - // thread_id (interrupts require durability). Without this check either - // may be absent and the graph still parks on the approval interrupt: no - // checkpoint is written, `into_outcome` fills `PendingApproval::thread_id` - // with an empty string, and `resume_delegation` then fails with - // "delegation resume requires a thread_id"/"a checkpointer" — the pause - // can never be released and the delegated work is stranded. Reject the - // misconfiguration up front instead. - if config.require_review_approval - && (config.checkpointer.is_none() || config.thread_id.is_none()) - { - return Err( - "delegation human-approval gate (require_review_approval) requires both a \ - checkpointer and a thread_id" - .to_string(), - ); - } - - let thread_id = config.thread_id.clone(); - let mut graph = build_delegation_graph( - config.max_revisions, - config.cancel.clone(), - config.require_review_approval, - run_stage, - )?; - if let Some(sink) = config.event_sink.clone() { - graph = graph.with_event_sink(sink); - } - - if let Some(cp) = config.checkpointer { - graph = graph.with_checkpointer(cp); - } - - tracing::info!( - max_revisions = config.max_revisions, - durable = thread_id.is_some(), - human_gated = config.require_review_approval, - "[delegation] running sub-agent delegation graph" - ); - - let initial = DelegationState::new_run(); - let execution = match thread_id.clone() { - Some(tid) => graph.run_with_thread(tid, initial).await, - None => graph.run(initial).await, - } - .map_err(|e| format!("delegation graph run failed: {e}"))?; - - Ok(into_outcome(execution, thread_id)) -} - -/// Resume a delegation graph parked on a durable human-approval interrupt, -/// delivering the approver's `decision` through `Command { resume: .. }`. -/// -/// The graph is rebuilt (its node closures are not serializable — only the typed -/// state is checkpointed) with the same checkpointer + `thread_id`, then -/// re-entered at the interrupted node via [`CompiledGraph::resume`](crate::CompiledGraph::resume) (the -/// `ResumeTarget::Latest` checkpoint). `decision` maps to approve/deny via -/// `decision_is_approve`, so passing the approval RPC's `ApprovalDecision` -/// (serialized with its stable `as_str()` wire value — `approve_once` / -/// `approve_always_for_tool` / `deny`) routes the existing decision contract -/// into the resume **without changing that contract**. -/// -/// TTL expiry → resume-with-deny: call this with [`deny_decision`] to preserve -/// the existing timeout-deny behavior for a pause that was never answered. -pub async fn resume_delegation( - config: DelegationConfig, - decision: Value, - run_stage: F, -) -> Result -where - F: Fn(DelegationStage, DelegationState) -> Fut + Clone + Send + Sync + 'static, - Fut: Future> + Send + 'static, -{ - // This is the public entry point a host calls directly with an approver's - // decision — it does NOT go through `run_or_resume_delegation`, so it must - // provide its own two protections that function otherwise only gets: - // - // 1. Per-thread serialization: two concurrent approval callbacks for the - // SAME paused thread (a retried callback racing with a deny, say) must - // not both load the same interrupt checkpoint and independently - // finalize opposite decisions, appending competing histories. Uses the - // SAME `thread_lock` map as `run_or_resume_delegation`, so the two - // entry points serialize against each other too, not just themselves. - // 2. Schema-version validation: `CompiledGraph::resume` (inside - // `resume_graph`) applies the decision to whatever checkpoint it - // finds, with no schema check of its own — that check lives in - // `run_or_resume_delegation`'s match arms, which this path bypasses. - // During a rollback or mixed-version deployment, serde can decode a - // checkpoint written by a newer binary by ignoring its added fields, - // letting an older binary consume the approval and finalize the state - // under outdated semantics. Reject rather than guess. - let thread_id = config - .thread_id - .clone() - .ok_or_else(|| "delegation resume requires a thread_id".to_string())?; - let cp = config - .checkpointer - .clone() - .ok_or_else(|| "delegation resume requires a checkpointer".to_string())?; - - let lock = thread_lock(&thread_id); - let _guard = lock.lock().await; - - if let Some(checkpoint) = cp - .get(thread_id.as_str(), None) - .await - .map_err(|e| format!("delegation checkpoint read failed for thread {thread_id}: {e}"))? - && checkpoint.state.schema_version != CURRENT_SCHEMA_VERSION - { - return Err(format!( - "delegation resume refused for thread {thread_id}: checkpoint schema_version {} does not match this binary's schema_version {} — the approval decision was NOT applied", - checkpoint.state.schema_version, CURRENT_SCHEMA_VERSION - )); - } - - let approved = decision_is_approve(&decision); - tracing::info!( - approved, - "[interrupt] resuming durable delegation graph with approval decision" - ); - let command = Command { - resume: Some(decision), - ..Command::default() - }; - resume_graph(config, command, run_stage).await -} - -/// Run the delegation graph, resuming from the last checkpoint boundary when the -/// configured thread has a live, compatible, non-terminal checkpoint, else -/// starting fresh (issue #3884 — node-level checkpoint & resume). -/// -/// Classifies the thread's latest checkpoint and routes accordingly: -/// - **resumable** (a crash/failure left a mid-run boundary) → re-run only the -/// not-yet-completed nodes from that boundary via [`CompiledGraph::resume`](crate::CompiledGraph::resume) -/// with an empty command — never restarting from `plan`, and never re-running -/// an already-completed step (its [`StepRecord`](crate::StepRecord) is restored from the state); -/// - **terminal** (already finalized/cancelled) → return the stored final state -/// without re-running (idempotent re-invocation of a stable thread); -/// - **absent** (no checkpoint) → a fresh durable run; -/// - **incompatible** (an undecodable record — e.g. a pre-#3884 `Vec` -/// `executions` shape) → log, best-effort prune, and start fresh. The decode -/// error is swallowed into a fresh-run decision, never propagated, never a panic. -/// -/// Callers that mint a unique `thread_id` per run (today's default) always take -/// the fresh path unchanged; the resume paths activate only when a caller reuses -/// a stable `thread_id`, so this is byte-compatible for existing callers while -/// wiring the resume seam that #3881 (plan-edit resume) builds on. -pub async fn run_or_resume_delegation( - config: DelegationConfig, - run_stage: F, -) -> Result -where - F: Fn(DelegationStage, DelegationState) -> Fut + Clone + Send + Sync + 'static, - Fut: Future> + Send + 'static, -{ - // Without a checkpointer + thread there is nothing to resume from. - let (Some(cp), Some(tid)) = (config.checkpointer.clone(), config.thread_id.clone()) else { - return run_delegation_durable(config, run_stage).await; - }; - - // Serialize checkpoint classification + the selected run/resume operation - // per thread. Without this, two concurrent callers for the same stable - // `thread_id` can both read the same checkpoint (or both observe none) - // before either starts executing — neither this wrapper nor the - // `CompiledGraph` run/resume entry points hold a per-thread lock across - // that gap on their own. Both would then run the same pending stage, - // duplicating execute-stage external side effects and producing - // competing checkpoint histories. The lock is held for the whole - // classify-then-dispatch critical section below, released on return. - let lock = thread_lock(&tid); - let _guard = lock.lock().await; - - match cp.get(tid.as_str(), None).await { - // A checkpoint written under a schema version other than the current - // one is expired rather than resumed/returned — its semantics may not - // match this binary's graph. This is what makes `schema_version` an - // actual guard, not just documentation, and closes the empty- - // `executions` gap that a decode failure alone cannot catch (e.g. a - // pre-#3884 record whose `executions` happened to be empty and so - // still decoded into `Vec`). - // - // A version *greater* than `CURRENT_SCHEMA_VERSION` is rejected here - // too, not only a lesser one: during a rollback or a mixed-version - // deployment, serde can decode a checkpoint written by a newer binary - // by silently ignoring its added fields, and this older binary would - // then resume or return that state under outdated semantics. The - // 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 => { - tracing::warn!( - thread_id = %tid, - schema_version = checkpoint.state.schema_version, - current = CURRENT_SCHEMA_VERSION, - "[delegation] checkpoint schema version does not match the current binary; pruning and starting fresh" - ); - prune_thread(cp.as_ref(), &tid).await; - run_delegation_durable(config, run_stage).await - } - Ok(Some(checkpoint)) if checkpoint_is_resumable(&checkpoint) => { - tracing::info!( - thread_id = %tid, - "[delegation] resuming durable delegation from its last checkpoint boundary" - ); - // A crash/failure resume carries no decision value — an empty command - // simply re-runs the pending node(s) (the crate's `retry` semantics). - resume_graph(config, Command::default(), run_stage).await - } - Ok(Some(checkpoint)) => { - // Terminal: return the finalized state without re-running. Defensive: - // if this checkpoint still carried an unconsumed interrupt, surface it - // instead of silently dropping it. The current routing never produces - // this (an interrupt boundary schedules its node and is classified - // resumable above), but a future schedule could, and a dropped pause - // would strand a run. - let pending = checkpoint.interrupts.first().map(|i| PendingApproval { - interrupt_id: i.id.clone(), - node: i.node.as_str().to_string(), - payload: i.payload.clone(), - thread_id: tid.clone(), - }); - if pending.is_some() { - tracing::warn!( - thread_id = %tid, - "[delegation] terminal-classified checkpoint carried a pending interrupt; surfacing it" - ); - } else { - tracing::info!( - thread_id = %tid, - "[delegation] thread already terminal; returning finalized state without re-running" - ); - } - Ok(DelegationOutcome { - state: checkpoint.state, - pending, - }) - } - Ok(None) => { - tracing::debug!( - thread_id = %tid, - "[delegation] no checkpoint for thread; starting a fresh durable run" - ); - run_delegation_durable(config, run_stage).await - } - // Only a *decode / shape-incompatibility* read error expires the - // checkpoint. An operational error (SQLite busy / I/O / poisoned lock) - // 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) => { - tracing::warn!( - thread_id = %tid, - error = %e, - "[delegation] undecodable/incompatible checkpoint; pruning and starting fresh" - ); - prune_thread(cp.as_ref(), &tid).await; - run_delegation_durable(config, run_stage).await - } - Err(e) => { - tracing::error!( - thread_id = %tid, - error = %e, - "[delegation] checkpoint read failed (operational); not restarting — propagating error" - ); - Err(format!( - "delegation checkpoint read failed for thread {tid}: {e}" - )) - } - } -} - -/// Best-effort prune of a dead/expired checkpoint thread so it is not re-probed -/// 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 { - tracing::debug!( - thread_id = %thread_id, - error = %e, - "[delegation] could not prune checkpoint thread (non-fatal)" - ); - } -} - -/// Whether a `Checkpointer::get` error is a *positively identified* schema -/// mismatch (safe to expire the checkpoint) rather than an operational failure -/// (SQLite busy / I/O / poisoned lock) or ambiguous/likely corruption — either -/// of which must propagate rather than silently discard the record. -/// -/// Both backends (`checkpoint::file`, `checkpoint::sqlite`) route every decode -/// failure through `checkpoint::decode_json_err`, which classifies the -/// underlying `serde_json::Error` and tags the message `"decode [schema] …"` -/// when the JSON parsed but didn't match the target type (a missing field, a -/// renamed variant — exactly what a legacy or newer schema produces), or -/// `"decode [corrupt] …"` when the bytes themselves were malformed or -/// truncated (a write torn by a crash, disk corruption). Only the former is -/// pruned here: matching on a bare `"decode"` substring previously expired a -/// corrupted record exactly as readily as an outdated one, deleting the only -/// evidence of the corruption and restarting the delegation from `plan`, -/// potentially repeating already-completed execute-stage side effects. -pub(super) fn is_incompatible_checkpoint_error(e: &crate::TinyAgentsError) -> bool { - matches!(e, crate::TinyAgentsError::Checkpoint(msg) if msg.contains("decode [schema]")) -} - -/// Whether a loaded checkpoint still has work to resume: a non-finalized run -/// that still schedules a real (non-`END`) node. -/// -/// Cancellation alone does NOT mean finalized: every cancellation route goes -/// through `goto finalize` (see `graph.rs`), so a checkpoint can be -/// `cancelled == true` with `final_output == None` and `next_nodes == -/// ["finalize"]` if the process stopped after the cancellation update was -/// checkpointed but before `finalize` ran. Treating `cancelled` as an -/// independent terminal signal classified that live boundary as terminal, so -/// `run_or_resume_delegation` returned the run unfinished and never produced -/// the cancellation summary. `final_output` — the one field `finalize` -/// itself sets — is the actual terminal signal; the schedule is what decides -/// whether there is still work to resume. -fn checkpoint_is_resumable(checkpoint: &Checkpoint) -> bool { - if checkpoint.state.final_output.is_some() { - return false; - } - // `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 -/// the typed state is checkpointed) with the same checkpointer + `thread_id`, and -/// re-enter it at the latest checkpoint's pending node(s) via -/// [`CompiledGraph::resume`]. `command` carries an approver's decision for the -/// human-approval interrupt path, or is empty for a plain crash/failure resume. -async fn resume_graph( - config: DelegationConfig, - command: Command, - run_stage: F, -) -> Result -where - F: Fn(DelegationStage, DelegationState) -> Fut + Clone + Send + Sync + 'static, - Fut: Future> + Send + 'static, -{ - let thread_id = config - .thread_id - .clone() - .ok_or_else(|| "delegation resume requires a thread_id".to_string())?; - let cp = config - .checkpointer - .clone() - .ok_or_else(|| "delegation resume requires a checkpointer".to_string())?; - - let mut graph = build_delegation_graph( - config.max_revisions, - config.cancel.clone(), - config.require_review_approval, - run_stage, - )? - .with_checkpointer(cp); - if let Some(sink) = config.event_sink.clone() { - graph = graph.with_event_sink(sink); - } - - let execution = graph - .resume(thread_id.clone(), command) - .await - .map_err(|e| format!("delegation graph resume failed: {e}"))?; - - Ok(into_outcome(execution, Some(thread_id))) -} - -/// The canonical deny decision used for TTL-expiry resume (resume-with-deny), -/// serialized to the approval RPC's stable `deny` wire value. -pub fn deny_decision() -> Value { - json!("deny") -} - -/// Fold a finished/paused graph execution into a [`DelegationOutcome`], -/// surfacing the first pending interrupt (if the run parked on one). -fn into_outcome( - execution: crate::GraphExecution, - thread_id: Option, -) -> DelegationOutcome { - let pending = execution.interrupts.first().map(|i| { - tracing::info!( - interrupt_id = %i.id, - node = %i.node.as_str(), - "[interrupt] delegation run parked on durable human-approval interrupt" - ); - PendingApproval { - interrupt_id: i.id.clone(), - node: i.node.as_str().to_string(), - payload: i.payload.clone(), - thread_id: thread_id.clone().unwrap_or_default(), - } - }); - DelegationOutcome { - state: execution.state, - pending, - } -} - -/// Map an approval decision value onto approve/deny. Accepts the approval RPC's -/// stable string forms (`approve_once`, `approve_always_for_tool`, `deny`), a -/// bare bool, or an object carrying `approved`/`decision` — so the existing -/// decision contract routes into `Command::resume` unchanged. -pub(super) fn decision_is_approve(decision: &Value) -> bool { - /// Recognised string spellings for an approval decision. Also used for - /// the object-shaped `{"decision": ..}` form below: a prefix check there - /// (`starts_with("approve")`) previously treated any unvalidated string - /// beginning with "approve" (e.g. an attacker- or bug-supplied - /// `"approve_not_authorized"`) as approval, releasing the durable - /// human-approval gate for a decision value nothing actually approved. - const APPROVE_STRINGS: &[&str] = &[ - "approve_once", - "approve_always_for_tool", - "approve", - "approved", - ]; - - match decision { - Value::Bool(b) => *b, - Value::String(s) => APPROVE_STRINGS.contains(&s.as_str()), - Value::Object(m) => { - if let Some(b) = m.get("approved").and_then(Value::as_bool) { - return b; - } - m.get("decision") - .and_then(Value::as_str) - .map(|d| APPROVE_STRINGS.contains(&d)) - .unwrap_or(false) - } - _ => false, - } -} diff --git a/crates/tinyagents-graph/src/delegation/types.rs b/crates/tinyagents-graph/src/delegation/types.rs deleted file mode 100644 index b6fc7d043..000000000 --- a/crates/tinyagents-graph/src/delegation/types.rs +++ /dev/null @@ -1,232 +0,0 @@ -//! Types threaded through the delegation graph. -//! -//! [`DelegationState`] is the durable one: it is what a [`Checkpointer`] -//! persists, so its serde representation is an on-disk format. See the -//! `schema_version` field and the module README before changing any field name, -//! `#[serde(...)]` attribute, or default. - -use std::sync::Arc; - -use serde::{Deserialize, Serialize}; -use serde_json::Value; - -use crate::checkpoint::Checkpointer; -use crate::stream::GraphEventSink; -use tinyagents_harness::cancel::CancellationToken; - -/// Which stage a delegation node is asking the injected worker to run. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum DelegationStage { - /// Produce a plan for the task. - Plan, - /// Execute the current plan (re-run on revision). - Execute, - /// Review the latest execution; may approve or request a revision. - Review, -} - -/// What an injected stage worker returns. -#[derive(Debug, Clone)] -pub struct DelegationStageOutput { - /// The stage's textual output (plan text, execution result, or review note). - pub text: String, - /// Only meaningful for [`DelegationStage::Review`]: `true` approves the - /// execution and ends the loop; `false` requests another revision. - pub approved: bool, - /// The exact prompt handed to this stage's worker, when it surfaces one. - /// Persisted into [`StepRecord::prompt`] for per-step provenance (read only - /// for the execute stage; ignored elsewhere). `None` when the worker does not - /// surface a prompt — e.g. the deterministic test mock. - pub prompt: Option, -} - -impl DelegationStageOutput { - /// A plain non-review stage output (the `approved` flag is unused and no - /// prompt is surfaced). - pub fn done(text: impl Into) -> Self { - Self { - text: text.into(), - approved: true, - prompt: None, - } - } -} - -/// Current on-disk schema version for a checkpointed [`DelegationState`]. Bumped -/// only on a breaking state-shape change — introduced with the -/// `executions: Vec` → `Vec` migration (issue #3884). -/// Pre-versioned records deserialize to `0` via `#[serde(default)]`, so a resume -/// can tell a stale checkpoint from a current one. -pub const CURRENT_SCHEMA_VERSION: u32 = 1; - -/// One completed execute-stage pass, recorded durably so a resumed run knows -/// exactly how far it got and can render/finalize per step rather than from a -/// flat text log. Replaces the former `executions: Vec` (issue #3884). -#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] -pub struct StepRecord { - /// 0-based execute pass: `0` is the first execution, `n` the n-th revision. - pub index: usize, - /// The exact prompt handed to the execute sub-agent for this pass — per-step - /// provenance and the seam a later plan-edit slice (#3881) diffs against. - /// Empty when the worker did not surface one. - #[serde(default)] - pub prompt: String, - /// The sub-agent's result text — the value the former `Vec` entry held. - pub result: String, -} - -/// Typed working state threaded through (and checkpointed across) the delegation -/// graph. Serde-serializable so a [`Checkpointer`] can persist and restore it. -#[derive(Debug, Clone, Default, Serialize, Deserialize)] -pub struct DelegationState { - /// The plan produced by the `plan` stage. - pub plan: Option, - /// One record per execution pass (the first plus each revision), typed so a - /// resumed run can render/finalize per step (widened from `Vec`, - /// issue #3884). - pub executions: Vec, - /// One entry per review pass. - pub reviews: Vec, - /// Number of revisions the reviewer requested (loops back to `execute`). - pub revisions: usize, - /// Set once the reviewer approves or the revision cap is hit. - pub approved: bool, - /// The final synthesized output (set by `finalize`). - pub final_output: Option, - /// Set when the run short-circuited because its token was cancelled. - pub cancelled: bool, - /// The durable human-approval decision, once a resume delivers one: - /// `Some(true)` = approved, `Some(false)` = denied, `None` = not gated / - /// still awaiting. Only meaningful when `require_review_approval` is set. - #[serde(default)] - pub human_approved: Option, - /// Set when the durable human-approval gate denied the delegated result - /// (deny semantics: block the action, finalize as denied). - #[serde(default)] - pub denied: bool, - /// On-disk schema version, stamped [`CURRENT_SCHEMA_VERSION`] on a fresh run - /// and defaulting to `0` for pre-versioned checkpoints. - /// [`run_or_resume_delegation`](crate::run_or_resume_delegation) expires any checkpoint whose version is below - /// `CURRENT_SCHEMA_VERSION` (and any that fails to deserialize) instead of - /// resuming or returning it — so a shape change that stays structurally - /// decodable is still not misread. - #[serde(default)] - pub schema_version: u32, -} - -impl DelegationState { - /// A fresh run's initial state, stamped with the current schema version so - /// its checkpoints are self-identifying. - pub(super) fn new_run() -> Self { - Self { - schema_version: CURRENT_SCHEMA_VERSION, - ..Self::default() - } - } - - /// The latest execution result text, if any — the projection the review - /// prompt and the finalize summary read. - pub fn last_result(&self) -> Option<&str> { - self.executions.last().map(|r| r.result.as_str()) - } - - /// The execution result texts in order — the flat projection used for the - /// durable approval-interrupt payload, kept `Vec<&str>` so that wire shape - /// is unchanged from the pre-#3884 `Vec`. - pub fn executions_texts(&self) -> Vec<&str> { - self.executions.iter().map(|r| r.result.as_str()).collect() - } -} - -/// Reducer updates emitted by the delegation nodes. -pub(crate) enum DelegationUpdate { - Plan(String), - Execution { - prompt: String, - result: String, - }, - Review { - note: String, - approved: bool, - }, - /// A durable human-approval decision delivered by a resume command. - HumanDecision { - approved: bool, - }, - Final(String), - Cancelled, -} - -/// Configuration for a delegation run. -pub struct DelegationConfig { - /// Upper bound on reviewer-requested revisions before forcing `finalize`. - pub max_revisions: usize, - /// Optional durable checkpointer (e.g. a `DriverCheckpointer`). When set with a - /// `thread_id`, the run persists its state at every super-step boundary. - pub checkpointer: Option>>, - /// Thread id for checkpoint keying; required for the checkpointer to persist. - pub thread_id: Option, - /// Cooperative cancellation; checked at each node boundary. - pub cancel: CancellationToken, - /// When set, an approved review does not finalize directly: the run reaches - /// a durable **human-approval** interrupt (`NodeResult::Interrupt`) that is - /// persisted via the checkpointer (Sync durability) and survives a process - /// restart. The pause is only released by [`resume_delegation`](crate::resume_delegation) carrying the - /// approver's decision. Requires `checkpointer` + `thread_id` (interrupts - /// require durability). - /// - /// This is the **durable** approval boundary — distinct from the interactive - /// chat-turn approval gate (the 10-min TTL steering pause surfaced via - /// `ApprovalRequestCard`), which parks a live chat turn in memory and is left - /// exactly as-is. Durable graphs pause by checkpoint; chat turns pause by - /// steering. See the `approval` node below. - pub require_review_approval: bool, - /// Optional observability sink for the graph's run/node/route events. - /// - /// Host-supplied and `None` by default: this crate emits its own `tracing` - /// diagnostics for the graph runtime, and a host that wants delegation runs - /// journalled into its own observability surface attaches a sink here - /// rather than the module reaching for a host-specific logger. - pub event_sink: Option>, -} - -impl Default for DelegationConfig { - fn default() -> Self { - Self { - max_revisions: 2, - checkpointer: None, - thread_id: None, - cancel: CancellationToken::new(), - require_review_approval: false, - event_sink: None, - } - } -} - -/// A durable human-approval pause the delegation graph is parked on. -/// -/// Produced when a run reaches the `approval` interrupt (see -/// [`DelegationConfig::require_review_approval`]). The pause is already -/// persisted as a checkpoint keyed by `thread_id`; the approver's decision is -/// delivered later via [`resume_delegation`](crate::resume_delegation), which survives a process restart. -#[derive(Debug, Clone)] -pub struct PendingApproval { - /// Stable id of the emitted interrupt (matches a resume value to this pause). - pub interrupt_id: String, - /// The node that emitted the interrupt (always `"approval"` here). - pub node: String, - /// Approval-request payload presented to the approver (review notes, etc.). - pub payload: Value, - /// Thread id the paused graph is checkpointed under; the resume key. - pub thread_id: String, -} - -/// Outcome of a durable delegation run or resume. -#[derive(Debug, Clone)] -pub struct DelegationOutcome { - /// The latest committed [`DelegationState`] at the run/resume boundary. - pub state: DelegationState, - /// `Some` when the run is parked on a durable human-approval interrupt; - /// `None` when the run reached a terminal (finalized) boundary. - pub pending: Option, -} diff --git a/crates/tinyagents-graph/src/execution_locks.rs b/crates/tinyagents-graph/src/execution_locks.rs deleted file mode 100644 index 84e9da165..000000000 --- a/crates/tinyagents-graph/src/execution_locks.rs +++ /dev/null @@ -1,33 +0,0 @@ -//! Private per-thread lock state for the graph's executor and delegation runner. -//! -//! Reusable lock maps live in `tinyagents-harness-tools`; these two graph -//! entry points retain only the private serialization state they need until -//! the graph runtime is retired. - -use std::collections::HashMap; -use std::sync::{Arc, Mutex as StdMutex, Weak}; - -use tokio::sync::Mutex; - -#[derive(Default)] -pub(super) struct RunLockMap { - locks: StdMutex>>>, -} - -impl RunLockMap { - pub(super) fn lock_for(&self, key: &str) -> Arc> { - let mut locks = self - .locks - .lock() - .unwrap_or_else(|_| panic!("graph run lock map poisoned")); - if let Some(lock) = locks.get(key).and_then(Weak::upgrade) { - return lock; - } - let lock = Arc::new(Mutex::new(())); - locks.insert(key.to_owned(), Arc::downgrade(&lock)); - if locks.len() >= 16 { - locks.retain(|_, weak| weak.strong_count() > 0); - } - lock - } -} diff --git a/crates/tinyagents-graph/src/export/README.md b/crates/tinyagents-graph/src/export/README.md deleted file mode 100644 index 2d7501fed..000000000 --- a/crates/tinyagents-graph/src/export/README.md +++ /dev/null @@ -1,75 +0,0 @@ -# graph::export - -Graph export and visualization — the introspection surface that lets a -recursive harness read back the shape of any graph, including one a model -authored or assembled at runtime. - -Graphs in this runtime can be built by hand (`GraphBuilder`) or assembled at -runtime. This module reduces both sources to one inspectable, behavior-free -description — `GraphTopology` — so a graph can be diffed, snapshotted in -tests, or drawn for a human reviewing what an agent just constructed. It -implements the spec's "graph serialization to JSON" and "Mermaid export" -future features (see `docs/modules/graph/visualization-testkit.md`). - -## Where topology comes from - -`GraphTopology` can be extracted from two sources, both yielding the same -shape so visualization and test snapshots share one truth: - -- `crate::CompiledGraph::topology` — a validated, frozen graph. -- `crate::GraphBuilder::topology` — a graph still under construction (entry - may be unresolved). - -None of these expose runnable behavior (handler/router closures, reducers); -only structure is captured, which is what makes `GraphTopology` cheap to -clone, `serde`-serializable, and safe to compare across builds. - -## Public surface - -- `GraphTopology` — the top-level, `serde`-serializable structural - description of a graph: nodes, edges, conditional routes, barrier edges, - channel bindings, a policy summary, and a validation report. All - collections are stored in stable sorted order so exports are - deterministic regardless of `HashMap` iteration order. -- `NodeInfo` — one node's structural metadata (kind, routing flags, `goto` - hints, metadata, derived `NodePolicySummary`). -- `EdgeInfo` / `ConditionalEdgeInfo` / `RouteInfo` — direct and - router-resolved edges. -- `WaitingEdgeInfo` — a barrier/fan-in join: `target` activates only once - every node in `predecessors` has completed. -- `ChannelInfo` — a state-channel-to-reducer binding (empty for compiled - whole-state graphs). -- `GraphPolicySummary` / `NodePolicySummary` — derived, graph- and - node-level execution policy summaries computed from the topology itself. -- `ValidationReport` — structural errors (dangling entry/edge/route/barrier - targets) and non-fatal warnings (unreachable or dead-end nodes) computed - over a topology. -- `to_json` / `from_json` — pretty JSON round-trip for `GraphTopology`. -- `to_mermaid` — deterministic [Mermaid](https://mermaid.js.org/) `flowchart` - rendering of a topology. - -## Files - -| File | Role | -| --- | --- | -| `types.rs` | The `GraphTopology` shape and its constituent structs; no behavior, `serde`-only. | -| `mod.rs` | Topology extraction (`build_topology`, `node_parts`), structural `validate`ation, JSON round-trip, and Mermaid rendering. | -| `mod_tests.rs` | Unit tests covering topology extraction from built/compiled graphs, JSON round-tripping, Mermaid output, and validation errors/warnings. | - -## Operational constraints - -- A `GraphTopology` never carries runnable behavior — no handler closures, - router closures, or reducers. Anything derived from a topology (JSON, - Mermaid, diffs) is therefore safe to log, snapshot in tests, or show to an - end user without leaking implementation closures. -- All topology collections are sorted before being handed back - (`build_topology` normalizes ordering), so two extractions of the same - logical graph produce byte-identical JSON and Mermaid output — this is - what makes snapshot testing viable. -- `ValidationReport` is computed structurally from the topology itself, not - re-derived from the live graph; a topology extracted from a `CompiledGraph` - is typically clean (compilation already validated it), while a topology - extracted from an in-progress `GraphBuilder` may surface real - in-progress issues (e.g. an unresolved entry). -- `START`/`END` are virtual boundaries and are never treated as undeclared - nodes by `validate`, even though they don't appear in `GraphTopology::nodes`. diff --git a/crates/tinyagents-graph/src/export/mod.rs b/crates/tinyagents-graph/src/export/mod.rs deleted file mode 100644 index 3750811cc..000000000 --- a/crates/tinyagents-graph/src/export/mod.rs +++ /dev/null @@ -1,696 +0,0 @@ -//! Graph export and visualization — the introspection surface that lets a -//! recursive harness read back the shape of any graph, including ones a model -//! authored or assembled at runtime. -//! -//! Because graphs in this runtime can be built by hand or assembled at runtime, -//! it matters that both reduce to one inspectable description. Export captures that -//! description as a behavior-free [`GraphTopology`] — never the runnable -//! handler/router closures — so a graph can be diffed, snapshotted in tests, or -//! drawn for a human reviewing what an agent just constructed. -//! -//! This module turns a graph's structure into portable artifacts: a -//! `serde`-serializable [`GraphTopology`], a pretty JSON document, and a -//! deterministic [Mermaid](https://mermaid.js.org/) `flowchart`. It implements -//! the spec's "graph serialization to JSON" and "Mermaid export" future -//! features (see `docs/modules/graph/visualization-testkit.md`). -//! -//! Topology can be extracted from two sources, both yielding the same -//! [`GraphTopology`] shape so visualization and test snapshots share one truth: -//! -//! - [`crate::CompiledGraph::topology`] — a validated, frozen graph. -//! - [`crate::GraphBuilder::topology`] — a graph still under -//! construction (entry may be unresolved). -//! -//! None of these expose runnable behavior (handler/router closures, reducers); -//! only structure is captured. -//! -//! ``` -//! use tinyagents_graph::{GraphBuilder, NodeResult, START, END}; -//! use tinyagents_graph::export::{to_json, to_mermaid}; -//! -//! let graph = GraphBuilder::::overwrite() -//! .add_node("a", |s, _| async move { Ok(NodeResult::Update(s + 1)) }) -//! .add_node("b", |s, _| async move { Ok(NodeResult::Update(s + 1)) }) -//! .add_edge(START, "a") -//! .add_edge("a", "b") -//! .add_edge("b", END) -//! .compile() -//! .unwrap(); -//! -//! let topology = graph.topology(); -//! let json = to_json(&topology); -//! let mermaid = to_mermaid(&topology); -//! assert!(mermaid.contains("flowchart TD")); -//! ``` - -mod report; -mod types; - -pub use report::GraphTopologyReport; -pub use types::{ - ChannelInfo, ConditionalEdgeInfo, EdgeInfo, GraphPolicySummary, GraphTopology, NodeInfo, - NodePolicySummary, RouteInfo, ValidationReport, WaitingEdgeInfo, -}; - -use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet, VecDeque}; - -use crate::Result; -use crate::builder::Branch; -use crate::builder::{END, GraphBuilder, START}; -use crate::compiled::CompiledGraph; -use tinyagents_harness::ids::NodeId; - -/// A behavior-free description of one node, fed into [`build_topology`]. -struct NodePart { - id: String, - kind: Option, - command_routing: bool, - subgraph: bool, - interrupt: bool, - deferred: bool, - command_destinations: Vec, - metadata: BTreeMap, -} - -/// Inputs to [`build_topology`]: a behavior-free view of one graph's structure. -struct TopologyParts { - graph_id: String, - name: Option, - recursion_limit: usize, - parallel: bool, - max_concurrency: Option, - node_timeout_ms: Option, - /// Every declared node. - nodes: Vec, - /// Raw edges, including the synthetic `START`/`END` edges and any edges that - /// are also registered as barrier/waiting edges. - edges: Vec<(String, String)>, - /// `(from, [(label, target)])` conditional routes. - conditional: Vec<(String, Vec<(String, String)>)>, - /// `(target, [predecessors])` barrier/waiting (fan-in) edges. - waiting: Vec<(String, Vec)>, - /// `(channel, reducer)` bindings. - channels: Vec<(String, String)>, -} - -/// Folds raw structural parts into a normalized, deterministically-ordered -/// [`GraphTopology`]. Synthetic `START`/`END` edges are lifted into -/// [`GraphTopology::entry`] and [`GraphTopology::finish_nodes`]; barrier edges -/// are lifted out of the direct-edge set into [`GraphTopology::waiting_edges`]. -fn build_topology(parts: TopologyParts) -> GraphTopology { - let TopologyParts { - graph_id, - name, - recursion_limit, - parallel, - max_concurrency, - node_timeout_ms, - nodes, - edges, - conditional, - waiting, - channels, - } = parts; - - // Normalize waiting edges first so direct-edge folding can skip them. - let mut waiting_edges: Vec = waiting - .into_iter() - .map(|(target, mut predecessors)| { - predecessors.sort(); - predecessors.dedup(); - WaitingEdgeInfo { - target, - predecessors, - } - }) - .collect(); - waiting_edges.sort(); - let waiting_pairs: BTreeSet<(String, String)> = waiting_edges - .iter() - .flat_map(|w| { - w.predecessors - .iter() - .map(move |p| (p.clone(), w.target.clone())) - }) - .collect(); - let barrier_targets: BTreeSet = - waiting_edges.iter().map(|w| w.target.clone()).collect(); - - let mut entry: Option = None; - let mut direct: Vec = Vec::new(); - let mut finish_nodes: Vec = Vec::new(); - - for (from, to) in edges { - if from == START { - entry = Some(to); - } else if to == END { - finish_nodes.push(from); - } else if waiting_pairs.contains(&(from.clone(), to.clone())) { - // Represented as a waiting edge, not a direct one. - } else { - direct.push(EdgeInfo { from, to }); - } - } - - let mut conditional_edges: Vec = conditional - .into_iter() - .map(|(from, routes)| { - let mut routes: Vec = routes - .into_iter() - .map(|(label, target)| RouteInfo { label, target }) - .collect(); - routes.sort(); - ConditionalEdgeInfo { from, routes } - }) - .collect(); - - let channels: Vec = channels - .into_iter() - .map(|(name, reducer)| ChannelInfo { name, reducer }) - .collect(); - - // Stable ordering so exports are deterministic. - direct.sort(); - conditional_edges.sort_by(|a, b| a.from.cmp(&b.from)); - finish_nodes.sort(); - finish_nodes.dedup(); - - // Index sets used to derive per-node policy summaries. - let conditional_set: BTreeSet<&str> = - conditional_edges.iter().map(|c| c.from.as_str()).collect(); - let direct_from_set: BTreeSet<&str> = direct.iter().map(|e| e.from.as_str()).collect(); - let finish_set: BTreeSet<&str> = finish_nodes.iter().map(String::as_str).collect(); - - let mut node_infos: Vec = nodes - .into_iter() - .map(|n| { - let routing = if n.command_routing { - "command" - } else if conditional_set.contains(n.id.as_str()) { - "conditional" - } else if finish_set.contains(n.id.as_str()) { - "terminal" - } else if direct_from_set.contains(n.id.as_str()) { - "static" - } else { - "unrouted" - }; - let policy = NodePolicySummary { - routing: routing.to_string(), - barrier: barrier_targets.contains(&n.id), - interrupt: n.interrupt, - deferred: n.deferred, - subgraph: n.subgraph, - }; - let mut command_destinations = n.command_destinations; - command_destinations.sort(); - command_destinations.dedup(); - NodeInfo { - id: n.id, - kind: n.kind, - command_routing: n.command_routing, - subgraph: n.subgraph, - interrupt: n.interrupt, - deferred: n.deferred, - command_destinations, - metadata: n.metadata, - policy, - } - }) - .collect(); - node_infos.sort_by(|a, b| a.id.cmp(&b.id)); - - let policy = GraphPolicySummary { - recursion_limit, - parallel, - max_concurrency, - node_timeout_ms, - }; - - let validation = validate( - &node_infos, - entry.as_deref(), - &direct, - &conditional_edges, - &waiting_edges, - &finish_nodes, - ); - - GraphTopology { - graph_id, - name, - entry, - recursion_limit, - parallel, - nodes: node_infos, - edges: direct, - conditional_edges, - waiting_edges, - finish_nodes, - channels, - policy, - validation, - } -} - -/// Computes a structural [`ValidationReport`] over the assembled topology. -/// -/// Errors flag references to undeclared nodes (a dangling entry, edge, route, -/// or barrier target) which would make the graph invalid; warnings flag -/// non-fatal observations (unreachable nodes, `unrouted` dead-ends). The virtual -/// `START`/`END` boundaries are never treated as undeclared. -fn validate( - nodes: &[NodeInfo], - entry: Option<&str>, - direct: &[EdgeInfo], - conditional: &[ConditionalEdgeInfo], - waiting: &[WaitingEdgeInfo], - finish_nodes: &[String], -) -> ValidationReport { - let declared: BTreeSet<&str> = nodes.iter().map(|n| n.id.as_str()).collect(); - let known = |id: &str| id == START || id == END || declared.contains(id); - - let mut errors: Vec = Vec::new(); - let mut warnings: Vec = Vec::new(); - - match entry { - None => errors.push("graph has no entry node (no START edge)".to_string()), - Some(e) if !known(e) => { - errors.push(format!("entry node `{e}` is not declared")); - } - Some(_) => {} - } - - for edge in direct { - if !known(&edge.from) { - errors.push(format!("edge source `{}` is not declared", edge.from)); - } - if !known(&edge.to) { - errors.push(format!("edge target `{}` is not declared", edge.to)); - } - } - for cond in conditional { - if !known(&cond.from) { - errors.push(format!( - "conditional source `{}` is not declared", - cond.from - )); - } - for route in &cond.routes { - if !known(&route.target) { - errors.push(format!( - "conditional route `{}` of `{}` targets undeclared node `{}`", - route.label, cond.from, route.target - )); - } - } - } - for w in waiting { - if !known(&w.target) { - errors.push(format!("barrier target `{}` is not declared", w.target)); - } - for pred in &w.predecessors { - if !known(pred) { - errors.push(format!( - "barrier predecessor `{pred}` of `{}` is not declared", - w.target - )); - } - } - } - for node in nodes { - for dest in &node.command_destinations { - if !known(dest) { - errors.push(format!( - "command destination `{dest}` of `{}` is not declared", - node.id - )); - } - } - } - for f in finish_nodes { - if !known(f) { - errors.push(format!("finish node `{f}` is not declared")); - } - } - - // Reachability from the entry across every successor relation. - if let Some(entry) = entry { - let mut successors: BTreeMap<&str, BTreeSet<&str>> = BTreeMap::new(); - for edge in direct { - successors.entry(&edge.from).or_default().insert(&edge.to); - } - for cond in conditional { - let entry_set = successors.entry(&cond.from).or_default(); - for route in &cond.routes { - entry_set.insert(&route.target); - } - } - for w in waiting { - for pred in &w.predecessors { - successors.entry(pred).or_default().insert(&w.target); - } - } - for node in nodes { - for dest in &node.command_destinations { - successors.entry(&node.id).or_default().insert(dest); - } - } - - let mut seen: BTreeSet<&str> = BTreeSet::new(); - let mut queue: VecDeque<&str> = VecDeque::new(); - seen.insert(entry); - queue.push_back(entry); - while let Some(node) = queue.pop_front() { - if let Some(next) = successors.get(node) { - for &n in next { - if seen.insert(n) { - queue.push_back(n); - } - } - } - } - for node in nodes { - if !seen.contains(node.id.as_str()) { - warnings.push(format!("node `{}` is unreachable from the entry", node.id)); - } - } - } - - // Dead-end (unrouted) nodes that are not terminal. - for node in nodes { - if node.policy.routing == "unrouted" { - warnings.push(format!( - "node `{}` has no outgoing route and is not terminal", - node.id - )); - } - } - - errors.sort(); - errors.dedup(); - warnings.sort(); - warnings.dedup(); - - ValidationReport { - ok: errors.is_empty(), - errors, - warnings, - } -} - -/// Builds the per-node [`NodePart`] list from a declared-node iterator, the -/// command-routing set, and the behavior-free node-metadata map. -fn node_parts<'a, I>( - ids: I, - is_command: impl Fn(&tinyagents_harness::ids::NodeId) -> bool, - meta: &std::collections::HashMap, -) -> Vec -where - I: IntoIterator, -{ - ids.into_iter() - .map(|id| { - let m = meta.get(id); - NodePart { - id: id.to_string(), - kind: m.and_then(|m| m.kind.clone()), - command_routing: is_command(id), - subgraph: m.is_some_and(|m| m.subgraph), - interrupt: m.is_some_and(|m| m.interrupt), - deferred: m.is_some_and(|m| m.deferred), - command_destinations: m - .map(|m| { - m.command_destinations - .iter() - .map(|d| d.to_string()) - .collect() - }) - .unwrap_or_default(), - metadata: m.map(|m| m.metadata.clone()).unwrap_or_default(), - } - }) - .collect() -} - -/// Flattens a direct-edge map into `(from, to)` string pairs. -fn edge_pairs(edges: &HashMap>) -> Vec<(String, String)> { - edges - .iter() - .flat_map(|(from, targets)| { - targets - .iter() - .map(move |to| (from.to_string(), to.to_string())) - }) - .collect() -} - -/// Flattens the branch map into `(from, [(label, target)])` conditional routes. -fn conditional_routes( - branches: &HashMap>, -) -> Vec<(String, Vec<(String, String)>)> { - branches - .iter() - .map(|(from, branch)| { - let routes = branch - .routes - .iter() - .map(|(label, target)| (label.clone(), target.to_string())) - .collect(); - (from.to_string(), routes) - }) - .collect() -} - -/// Flattens the barrier map into `(target, [predecessors])` fan-in edges. -fn waiting_sets(waiting: &HashMap>) -> Vec<(String, Vec)> { - waiting - .iter() - .map(|(target, preds)| { - ( - target.to_string(), - preds.iter().map(ToString::to_string).collect(), - ) - }) - .collect() -} - -impl CompiledGraph { - /// Extracts a behavior-free [`GraphTopology`] describing this graph's - /// structure (id, nodes, direct edges, conditional routes, entry, finish - /// nodes). Node handler and router closures are never exposed. - pub fn topology(&self) -> GraphTopology { - let nodes = node_parts( - self.nodes.keys(), - |id| self.command_nodes.contains(id), - &self.node_meta, - ); - let edges = edge_pairs(&self.edges); - let conditional = conditional_routes(&self.branches); - let waiting = waiting_sets(&self.waiting); - build_topology(TopologyParts { - graph_id: self.graph_id().to_string(), - name: self.name().map(str::to_string), - recursion_limit: self.recursion_limit, - parallel: self.parallel, - max_concurrency: self.max_concurrency, - node_timeout_ms: self.node_timeout.map(|d| d.as_millis()), - nodes, - edges, - conditional, - waiting, - channels: Vec::new(), - }) - } -} - -impl GraphBuilder { - /// Extracts a [`GraphTopology`] from an in-progress builder. - /// - /// Unlike [`CompiledGraph::topology`], the builder has not been validated: - /// the entry may be unresolved (`None`) if `START` has no edge yet, and - /// dangling targets are reported as-is. - pub fn topology(&self) -> GraphTopology { - let nodes = node_parts( - self.nodes.keys(), - |id| self.command_nodes.contains(id), - &self.node_meta, - ); - let edges = edge_pairs(&self.edges); - let conditional = conditional_routes(&self.branches); - let waiting = waiting_sets(&self.waiting); - build_topology(TopologyParts { - graph_id: self.graph_id.to_string(), - name: self.name.clone(), - recursion_limit: self.recursion_limit, - parallel: self.parallel, - max_concurrency: self.max_concurrency, - node_timeout_ms: self.node_timeout.map(|d| d.as_millis()), - nodes, - edges, - conditional, - waiting, - channels: Vec::new(), - }) - } -} - -/// Serializes a [`GraphTopology`] to a pretty-printed JSON document. -pub fn to_json(topology: &GraphTopology) -> String { - // GraphTopology is a plain data struct of strings/bools/usize; serialization - // is infallible in practice, but fall back to an empty object defensively. - serde_json::to_string_pretty(topology).unwrap_or_else(|_| "{}".to_string()) -} - -/// Deserializes a [`GraphTopology`] from a JSON document produced by -/// [`to_json`]. Returns [`tinyagents_harness::error::TinyAgentsError::Serialization`] on -/// malformed input. -pub fn from_json(json: &str) -> Result { - Ok(serde_json::from_str(json)?) -} - -/// Renders a [`GraphTopology`] as a Mermaid `flowchart TD`. -/// -/// The synthetic `START` and `END` boundaries are drawn as stadium nodes, -/// direct edges use `-->`, and conditional routes are labeled `-- label -->`. -/// Output is deterministic: node declarations and edges follow the topology's -/// already-sorted ordering, so the same graph always renders identically. -pub fn to_mermaid(topology: &GraphTopology) -> String { - let mut out = String::from("flowchart TD\n"); - - // Boundary nodes. - out.push_str(" START([START])\n"); - out.push_str(" END([END])\n"); - - // Node declarations with a quoted label so ids with reserved characters - // still render. Subgraph nodes use the Mermaid subroutine shape (`[[...]]`) - // so they read as embedded graphs; everything else is a plain box. - for node in &topology.nodes { - let id = mermaid_id(&node.id); - let label = escape_label(&node.id); - if node.subgraph { - out.push_str(&format!(" {id}[[\"{label}\"]]\n")); - } else { - out.push_str(&format!(" {id}[\"{label}\"]\n")); - } - } - - // Marker classes. Emitted only when at least one node carries the marker so - // the output stays minimal and deterministic. - emit_marker_class(&mut out, topology, "subgraph", |n| n.subgraph); - emit_marker_class(&mut out, topology, "interrupt", |n| n.interrupt); - emit_marker_class(&mut out, topology, "deferred", |n| n.deferred); - - out.push('\n'); - - // Entry edge. - if let Some(entry) = &topology.entry { - out.push_str(&format!(" START --> {}\n", mermaid_ref(entry))); - } - - // Direct edges. - for edge in &topology.edges { - out.push_str(&format!( - " {} --> {}\n", - mermaid_ref(&edge.from), - mermaid_ref(&edge.to) - )); - } - - // Conditional edges, labeled. - for cond in &topology.conditional_edges { - for route in &cond.routes { - out.push_str(&format!( - " {} -- {} --> {}\n", - mermaid_ref(&cond.from), - escape_label(&route.label), - mermaid_ref(&route.target) - )); - } - } - - // Barrier/waiting (fan-in) edges drawn as dotted, labeled joins. - for w in &topology.waiting_edges { - for pred in &w.predecessors { - out.push_str(&format!( - " {} -. barrier .-> {}\n", - mermaid_ref(pred), - mermaid_ref(&w.target) - )); - } - } - - // Command `goto` destination hints drawn as dotted, labeled edges. - for node in &topology.nodes { - for dest in &node.command_destinations { - out.push_str(&format!( - " {} -. goto .-> {}\n", - mermaid_ref(&node.id), - mermaid_ref(dest) - )); - } - } - - // Finish edges. - for node in &topology.finish_nodes { - out.push_str(&format!(" {} --> END\n", mermaid_ref(node))); - } - - out -} - -/// Emits a Mermaid `classDef` plus `class` assignments for every node matching -/// `pred`, in the topology's already-sorted node order. Nothing is emitted when -/// no node matches, keeping the diagram minimal and deterministic. -fn emit_marker_class( - out: &mut String, - topology: &GraphTopology, - class: &str, - pred: impl Fn(&NodeInfo) -> bool, -) { - let matching: Vec<&NodeInfo> = topology.nodes.iter().filter(|n| pred(n)).collect(); - if matching.is_empty() { - return; - } - out.push_str(&format!(" classDef {class} stroke-dasharray: 4 2;\n")); - for node in matching { - out.push_str(&format!(" class {} {class}\n", mermaid_id(&node.id))); - } -} - -/// Maps a node id to its Mermaid reference token. The reserved `START`/`END` -/// boundaries map to the literal boundary nodes; everything else is sanitized. -fn mermaid_ref(id: &str) -> String { - if id == START || id == "START" { - "START".to_string() - } else if id == END || id == "END" { - "END".to_string() - } else { - mermaid_id(id) - } -} - -/// Produces a Mermaid-safe identifier from an arbitrary node id by replacing -/// any non-alphanumeric/underscore character with `_` and prefixing `n_` so the -/// result never collides with the reserved `START`/`END` tokens. -fn mermaid_id(id: &str) -> String { - let mut sanitized = String::with_capacity(id.len() + 2); - sanitized.push_str("n_"); - for ch in id.chars() { - if ch.is_ascii_alphanumeric() || ch == '_' { - sanitized.push(ch); - } else { - sanitized.push('_'); - } - } - sanitized -} - -/// Escapes a label for use inside a Mermaid quoted string or edge label. -fn escape_label(label: &str) -> String { - label.replace('"', """) -} - -#[cfg(test)] -#[path = "mod_tests.rs"] -mod test; diff --git a/crates/tinyagents-graph/src/export/mod_tests.rs b/crates/tinyagents-graph/src/export/mod_tests.rs deleted file mode 100644 index 063daa0d1..000000000 --- a/crates/tinyagents-graph/src/export/mod_tests.rs +++ /dev/null @@ -1,289 +0,0 @@ -//! Unit tests for graph export / visualization. - -use crate::builder::{END, GraphBuilder, START}; -use crate::command::NodeResult; -use crate::export::{from_json, to_json, to_mermaid}; - -/// Builds a small branching graph: START -> route -> {even,odd} with `route` -/// conditionally selecting a successor, both successors finishing. -fn sample_graph() -> crate::CompiledGraph { - GraphBuilder::::overwrite() - .with_graph_id("demo") - .add_node("route", |s, _| async move { Ok(NodeResult::Update(s)) }) - .add_node("even", |s, _| async move { Ok(NodeResult::Update(s + 1)) }) - .add_node("odd", |s, _| async move { Ok(NodeResult::Update(s - 1)) }) - .add_edge(START, "route") - .add_conditional_edges( - "route", - |s: &i64| { - if s % 2 == 0 { - "even".to_string() - } else { - "odd".to_string() - } - }, - [("even", "even"), ("odd", "odd")], - ) - .add_edge("even", END) - .add_edge("odd", END) - .compile() - .expect("compiles") -} - -#[test] -fn extracts_topology_structure() { - let topology = sample_graph().topology(); - - assert_eq!(topology.graph_id, "demo"); - assert_eq!(topology.entry.as_deref(), Some("route")); - - // Nodes are sorted by id. - let ids: Vec<&str> = topology.nodes.iter().map(|n| n.id.as_str()).collect(); - assert_eq!(ids, vec!["even", "odd", "route"]); - - // `route` has a conditional edge with both labeled routes (sorted). - assert_eq!(topology.conditional_edges.len(), 1); - let cond = &topology.conditional_edges[0]; - assert_eq!(cond.from, "route"); - let routes: Vec<(&str, &str)> = cond - .routes - .iter() - .map(|r| (r.label.as_str(), r.target.as_str())) - .collect(); - assert_eq!(routes, vec![("even", "even"), ("odd", "odd")]); - - // Both `even` and `odd` are finish nodes; START/END edges are lifted out. - assert_eq!(topology.finish_nodes, vec!["even", "odd"]); - assert!(topology.edges.is_empty()); -} - -#[test] -fn json_round_trips() { - let topology = sample_graph().topology(); - let json = to_json(&topology); - assert!(json.contains("\"graph_id\": \"demo\"")); - - let restored = from_json(&json).expect("valid json"); - assert_eq!(restored, topology); -} - -#[test] -fn from_json_rejects_malformed_input() { - assert!(from_json("not json").is_err()); -} - -#[test] -fn mermaid_contains_nodes_edges_and_labels() { - let topology = sample_graph().topology(); - let mermaid = to_mermaid(&topology); - - assert!(mermaid.starts_with("flowchart TD\n")); - assert!(mermaid.contains("START([START])")); - assert!(mermaid.contains("END([END])")); - // Node declarations carry the original label. - assert!(mermaid.contains("n_route[\"route\"]")); - // Entry edge. - assert!(mermaid.contains("START --> n_route")); - // Conditional routes are labeled. - assert!(mermaid.contains("n_route -- even --> n_even")); - assert!(mermaid.contains("n_route -- odd --> n_odd")); - // Finish edges. - assert!(mermaid.contains("n_even --> END")); - assert!(mermaid.contains("n_odd --> END")); -} - -#[test] -fn mermaid_is_deterministic() { - let a = to_mermaid(&sample_graph().topology()); - let b = to_mermaid(&sample_graph().topology()); - assert_eq!(a, b); -} - -#[test] -fn direct_edges_are_captured() { - let graph = GraphBuilder::::overwrite() - .add_node("a", |s, _| async move { Ok(NodeResult::Update(s)) }) - .add_node("b", |s, _| async move { Ok(NodeResult::Update(s)) }) - .add_edge(START, "a") - .add_edge("a", "b") - .add_edge("b", END) - .compile() - .expect("compiles"); - - let topology = graph.topology(); - assert_eq!(topology.edges.len(), 1); - assert_eq!(topology.edges[0].from, "a"); - assert_eq!(topology.edges[0].to, "b"); - assert_eq!(topology.finish_nodes, vec!["b"]); -} - -#[test] -fn builder_topology_works_before_compile() { - let builder = GraphBuilder::::overwrite() - .with_graph_id("wip") - .add_node("a", |s, _| async move { Ok(NodeResult::Update(s)) }) - .add_edge(START, "a"); - - let topology = builder.topology(); - assert_eq!(topology.graph_id, "wip"); - assert_eq!(topology.entry.as_deref(), Some("a")); - assert_eq!(topology.nodes.len(), 1); -} - -/// A graph exercising every marker kind: a conditional route, a barrier -/// (waiting/fan-in) edge, and a subgraph node, plus an interrupt and a deferred -/// marker. The structured export must surface them all, and Mermaid must render. -#[test] -fn exports_all_marker_kinds_and_renders_mermaid() { - let graph = GraphBuilder::::overwrite() - .with_graph_id("markers") - .with_name("Markers Demo") - .add_node("route", |s, _| async move { Ok(NodeResult::Update(s)) }) - .add_node("left", |s, _| async move { Ok(NodeResult::Update(s + 1)) }) - .add_node("right", |s, _| async move { Ok(NodeResult::Update(s - 1)) }) - .add_node("join", |s, _| async move { Ok(NodeResult::Update(s)) }) - .add_conditional_edges( - "route", - |s: &i64| { - if *s >= 0 { - "left".to_string() - } else { - "right".to_string() - } - }, - [("left", "left"), ("right", "right")], - ) - .add_waiting_edge("left", "join") - .add_waiting_edge("right", "join") - .add_edge(START, "route") - .add_edge("join", END) - .mark_subgraph("left") - .mark_interrupt("route") - .mark_deferred("join") - .with_node_kind("left", "subgraph") - .with_node_metadata("route", "owner", "triage") - .compile() - .expect("compiles"); - - let topology = graph.topology(); - - // Graph id + name. - assert_eq!(topology.graph_id, "markers"); - assert_eq!(topology.name.as_deref(), Some("Markers Demo")); - - // Graph-level policy summary. - assert_eq!(topology.policy.recursion_limit, 50); - assert!(!topology.policy.parallel); - - // Conditional route labels. - assert_eq!(topology.conditional_edges.len(), 1); - assert_eq!(topology.conditional_edges[0].from, "route"); - - // Barrier / waiting edges: `join` waits on both branches; they are lifted - // out of the direct edge set. - assert_eq!(topology.waiting_edges.len(), 1); - assert_eq!(topology.waiting_edges[0].target, "join"); - assert_eq!( - topology.waiting_edges[0].predecessors, - vec!["left", "right"] - ); - assert!(topology.edges.is_empty()); - - // Per-node markers + policy summaries. - let route = topology.nodes.iter().find(|n| n.id == "route").unwrap(); - assert!(route.interrupt); - assert_eq!(route.policy.routing, "conditional"); - assert!(route.policy.interrupt); - assert_eq!( - route.metadata.get("owner").map(String::as_str), - Some("triage") - ); - - let left = topology.nodes.iter().find(|n| n.id == "left").unwrap(); - assert!(left.subgraph); - assert!(left.policy.subgraph); - assert_eq!(left.kind.as_deref(), Some("subgraph")); - - let join = topology.nodes.iter().find(|n| n.id == "join").unwrap(); - assert!(join.deferred); - assert!(join.policy.barrier); - assert!(join.policy.deferred); - assert_eq!(join.policy.routing, "terminal"); - - // Validation is clean for this well-formed graph. - assert!(topology.validation.ok); - assert!(topology.validation.errors.is_empty()); - - // JSON round-trips with all the new fields. - let restored = from_json(&to_json(&topology)).expect("round trip"); - assert_eq!(restored, topology); - - // Mermaid renders every marker kind. - let mermaid = to_mermaid(&topology); - assert!(mermaid.starts_with("flowchart TD\n")); - // Subgraph node uses the subroutine shape and a class. - assert!(mermaid.contains("n_left[[\"left\"]]")); - assert!(mermaid.contains("class n_left subgraph")); - // Interrupt + deferred classes. - assert!(mermaid.contains("class n_route interrupt")); - assert!(mermaid.contains("class n_join deferred")); - // Conditional labels and barrier joins. - assert!(mermaid.contains("n_route -- left --> n_left")); - assert!(mermaid.contains("n_left -. barrier .-> n_join")); - assert!(mermaid.contains("n_right -. barrier .-> n_join")); - assert!(mermaid.contains("n_join --> END")); -} - -/// Command-routing nodes can declare `goto` destination hints, which the export -/// surfaces structurally and as dotted Mermaid edges. -#[test] -fn command_destination_hints_are_exported() { - let graph = GraphBuilder::::overwrite() - .add_node("dispatch", |s, _| async move { Ok(NodeResult::Update(s)) }) - .add_node("a", |s, _| async move { Ok(NodeResult::Update(s)) }) - .add_node("b", |s, _| async move { Ok(NodeResult::Update(s)) }) - .add_edge(START, "dispatch") - .with_command_destinations("dispatch", ["a", "b"]) - .add_edge("a", END) - .add_edge("b", END) - .compile() - .expect("compiles"); - - let topology = graph.topology(); - let dispatch = topology.nodes.iter().find(|n| n.id == "dispatch").unwrap(); - assert!(dispatch.command_routing); - assert_eq!(dispatch.command_destinations, vec!["a", "b"]); - assert_eq!(dispatch.policy.routing, "command"); - - let mermaid = to_mermaid(&topology); - assert!(mermaid.contains("n_dispatch -. goto .-> n_a")); - assert!(mermaid.contains("n_dispatch -. goto .-> n_b")); -} - -/// A builder-stage topology with a dangling conditional target and an -/// unreachable node surfaces a validation error and warning respectively. -#[test] -fn validation_flags_dangling_and_unreachable() { - let builder = GraphBuilder::::overwrite() - .add_node("a", |s, _| async move { Ok(NodeResult::Update(s)) }) - .add_node("orphan", |s, _| async move { Ok(NodeResult::Update(s)) }) - .add_edge(START, "a") - .add_conditional_edges("a", |_: &i64| "x".to_string(), [("x", "missing")]); - - let topology = builder.topology(); - assert!(!topology.validation.ok); - assert!( - topology - .validation - .errors - .iter() - .any(|e| e.contains("missing")) - ); - assert!( - topology - .validation - .warnings - .iter() - .any(|w| w.contains("orphan") && w.contains("unreachable")) - ); -} diff --git a/crates/tinyagents-graph/src/export/report.rs b/crates/tinyagents-graph/src/export/report.rs deleted file mode 100644 index cfeb5e929..000000000 --- a/crates/tinyagents-graph/src/export/report.rs +++ /dev/null @@ -1,39 +0,0 @@ -//! Rendered topology reports for debug / inspection UIs. - -use super::{GraphTopology, to_json, to_mermaid}; - -/// A rendered topology for one named graph: Mermaid and JSON renderings plus -/// the structural validation outcome. Behaviour-free; never closure bodies. -#[derive(Debug, Clone)] -pub struct GraphTopologyReport { - /// Stable graph label (e.g. `"agent_teams:member"`). - pub name: &'static str, - /// Mermaid `flowchart TD` rendering. - pub mermaid: String, - /// Pretty-printed JSON of the full topology. - pub json: String, - /// `true` when the structural validation found no errors. - pub ok: bool, - /// Structural defects (missing nodes, unreachable routes, ...). - pub errors: Vec, - /// Non-fatal observations. - pub warnings: Vec, -} - -impl GraphTopologyReport { - /// Render `topology` under the stable label `name`. - pub fn describe(name: &'static str, topology: &GraphTopology) -> Self { - Self { - name, - mermaid: to_mermaid(topology), - json: to_json(topology), - ok: topology.validation.ok, - errors: topology.validation.errors.clone(), - warnings: topology.validation.warnings.clone(), - } - } -} - -#[cfg(test)] -#[path = "report_tests.rs"] -mod tests; diff --git a/crates/tinyagents-graph/src/export/report_tests.rs b/crates/tinyagents-graph/src/export/report_tests.rs deleted file mode 100644 index 271b2bc48..000000000 --- a/crates/tinyagents-graph/src/export/report_tests.rs +++ /dev/null @@ -1,12 +0,0 @@ -use super::*; - -#[test] -fn describe_renders_mermaid_json_and_validation() { - let topology = crate::delegation::delegation_graph_topology().expect("delegation topology"); - let report = GraphTopologyReport::describe("delegation", &topology); - assert_eq!(report.name, "delegation"); - assert!(report.ok, "{:?}", report.errors); - assert!(report.mermaid.contains("flowchart")); - assert!(!report.json.is_empty()); - assert_eq!(report.errors, topology.validation.errors); -} diff --git a/crates/tinyagents-graph/src/export/types.rs b/crates/tinyagents-graph/src/export/types.rs deleted file mode 100644 index bbbf61191..000000000 --- a/crates/tinyagents-graph/src/export/types.rs +++ /dev/null @@ -1,208 +0,0 @@ -//! Serializable graph-topology types for export and visualization — the single, -//! behavior-free description every graph in the recursive runtime collapses to. -//! -//! Hand-built and runtime-assembled graphs extract to the same -//! [`GraphTopology`] shape, so one type backs JSON export, Mermaid rendering, -//! and test snapshots regardless of how the graph came to exist. -//! -//! These types describe the *structure* of a graph — its nodes, edges, and -//! conditional routes — without referencing any runnable node behavior (handler -//! closures, router closures, reducers). That makes a [`GraphTopology`] cheap to -//! clone, `serde`-serializable, and safe to snapshot in tests or render as a -//! diagram. Extract one from a compiled or built graph with -//! [`crate::CompiledGraph::topology`] / -//! [`crate::GraphBuilder::topology`]. - -use std::collections::BTreeMap; - -use serde::{Deserialize, Serialize}; - -/// A serializable, behavior-free description of a graph's structure. -/// -/// The topology is the single source of truth shared by JSON export -/// ([`crate::export::to_json`]) and Mermaid rendering -/// ([`crate::export::to_mermaid`]). All collections are stored in a -/// stable, sorted order so exports are deterministic regardless of the -/// underlying `HashMap` iteration order. -#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] -pub struct GraphTopology { - /// The graph identifier. - pub graph_id: String, - /// The optional human-readable graph name (descriptive; `graph_id` remains - /// the stable identifier). - #[serde(default, skip_serializing_if = "Option::is_none")] - pub name: Option, - /// The entry node (the target of the virtual `START` node), if known. - pub entry: Option, - /// Maximum number of supersteps the graph may execute (0 when unknown). - pub recursion_limit: usize, - /// Whether the active node set of a superstep runs concurrently. - pub parallel: bool, - /// All declared nodes, sorted by id. - pub nodes: Vec, - /// Direct (unconditional) edges between nodes, sorted by `(from, to)`. - /// Excludes the synthetic `START`/`END` edges, which are represented by - /// [`Self::entry`] and [`Self::finish_nodes`]. - pub edges: Vec, - /// Conditional (router-driven) edges, sorted by source node id. - pub conditional_edges: Vec, - /// Barrier / waiting (fan-in) edges, sorted by target node id. A target's - /// predecessor set must *all* complete before it activates. - #[serde(default, skip_serializing_if = "Vec::is_empty")] - pub waiting_edges: Vec, - /// Nodes that route directly to the virtual `END` node, sorted. - pub finish_nodes: Vec, - /// State channel / reducer bindings, when available (populated from a - /// declarative state-channel binding; empty for whole-state graphs). - pub channels: Vec, - /// Graph-level execution policy summary (recursion limit, concurrency, - /// per-node timeout). - #[serde(default)] - pub policy: GraphPolicySummary, - /// Structural validation report computed over this topology. - #[serde(default)] - pub validation: ValidationReport, -} - -/// A single graph node's structural metadata. -#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] -pub struct NodeInfo { - /// The node id. - pub id: String, - /// The node kind, when supplied by the graph builder. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub kind: Option, - /// True when this node routes exclusively via a `Command` `goto` rather - /// than static or conditional edges. - #[serde(default, skip_serializing_if = "is_false")] - pub command_routing: bool, - /// True when this node embeds and runs a child graph (a subgraph node). - #[serde(default, skip_serializing_if = "is_false")] - pub subgraph: bool, - /// True when this node is an interrupt point that can pause the run. - #[serde(default, skip_serializing_if = "is_false")] - pub interrupt: bool, - /// True when this node is a deferred join (activates after the frontier - /// drains). - #[serde(default, skip_serializing_if = "is_false")] - pub deferred: bool, - /// Declared `goto` destination hints for a command-routing node, sorted. - /// Advisory only — the runtime resolves the real target at runtime. - #[serde(default, skip_serializing_if = "Vec::is_empty")] - pub command_destinations: Vec, - /// Free-form, sorted key/value annotations carried from the builder. - #[serde(default, skip_serializing_if = "BTreeMap::is_empty")] - pub metadata: BTreeMap, - /// A compact, derived summary of this node's routing/behavior policy. - #[serde(default)] - pub policy: NodePolicySummary, -} - -/// A barrier / waiting (fan-in) edge: `target` only activates once every node -/// in `predecessors` has completed, possibly across supersteps. -#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -pub struct WaitingEdgeInfo { - /// The join node that waits on its predecessors. - pub target: String, - /// The predecessor nodes that must all complete first, sorted. - pub predecessors: Vec, -} - -/// Graph-level execution policy, summarized for export. -#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)] -pub struct GraphPolicySummary { - /// Maximum number of supersteps the graph may execute (0 when unknown). - pub recursion_limit: usize, - /// Whether the active node set of a superstep runs concurrently. - pub parallel: bool, - /// Upper bound on branches run concurrently per step (`None` = unbounded). - #[serde(default, skip_serializing_if = "Option::is_none")] - pub max_concurrency: Option, - /// Default per-node handler timeout in milliseconds (`None` = no timeout). - #[serde(default, skip_serializing_if = "Option::is_none")] - pub node_timeout_ms: Option, -} - -/// A derived, per-node summary of how a node routes and what role it plays. -/// -/// Every field is computed from the topology itself, so it stays in sync with -/// the structural fields and is safe to snapshot in tests. -#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)] -pub struct NodePolicySummary { - /// The node's routing discipline: one of `static`, `conditional`, - /// `command`, `terminal`, or `unrouted`. - pub routing: String, - /// True when the node is a barrier/waiting (fan-in) join target. - #[serde(default, skip_serializing_if = "is_false")] - pub barrier: bool, - /// True when the node is an interrupt point. - #[serde(default, skip_serializing_if = "is_false")] - pub interrupt: bool, - /// True when the node is a deferred join. - #[serde(default, skip_serializing_if = "is_false")] - pub deferred: bool, - /// True when the node embeds a subgraph. - #[serde(default, skip_serializing_if = "is_false")] - pub subgraph: bool, -} - -/// A structural validation report computed over a [`GraphTopology`]. -/// -/// `errors` are structural defects that would make the graph invalid (a missing -/// entry, a dangling edge/route/barrier target). `warnings` are non-fatal -/// observations (unreachable nodes, dead-end nodes). A compiled graph is already -/// validated, so its report is typically clean; a builder-stage topology may -/// surface in-progress issues. -#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)] -pub struct ValidationReport { - /// True when there are no errors. - pub ok: bool, - /// Structural defects, sorted and deduplicated. - #[serde(default, skip_serializing_if = "Vec::is_empty")] - pub errors: Vec, - /// Non-fatal observations, sorted and deduplicated. - #[serde(default, skip_serializing_if = "Vec::is_empty")] - pub warnings: Vec, -} - -/// A direct, unconditional edge from one node to another. -#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -pub struct EdgeInfo { - /// The source node id. - pub from: String, - /// The target node id. - pub to: String, -} - -/// Conditional routing out of a node: a set of labeled targets resolved at -/// runtime by the node's router. -#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] -pub struct ConditionalEdgeInfo { - /// The source node id. - pub from: String, - /// The labeled routes, sorted by label. - pub routes: Vec, -} - -/// One labeled branch of a conditional edge. -#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -pub struct RouteInfo { - /// The route label returned by the router. - pub label: String, - /// The target node id (may be the virtual `END`). - pub target: String, -} - -/// A state-channel-to-reducer binding. -#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] -pub struct ChannelInfo { - /// The channel name. - pub name: String, - /// The reducer reference bound to the channel. - pub reducer: String, -} - -/// Serde helper: skip serializing `false` booleans. -fn is_false(value: &bool) -> bool { - !*value -} diff --git a/crates/tinyagents-graph/src/lib.rs b/crates/tinyagents-graph/src/lib.rs deleted file mode 100644 index 2e1e7cb94..000000000 --- a/crates/tinyagents-graph/src/lib.rs +++ /dev/null @@ -1,121 +0,0 @@ -//! TinyAgents graph runtime. -//! -//! The graph module is TinyAgents' durable workflow runtime (LangGraph-style) -//! and one of the load-bearing surfaces of the crate's recursive language-model -//! architecture: because a node can embed another compiled graph -//! ([`subgraph`]) or invoke a sub-agent, **graphs run graphs** and orchestration -//! recurses while every step stays typed, checkpointed, and observable. A -//! workflow driven from a host orchestrator lowers into exactly these same -//! types. -//! -//! The pieces: partial updates and reducers ([`reducer`]), commands and -//! interrupts ([`command`]), a builder/compile contract ([`builder`]), a -//! superstep executor ([`compiled`]), checkpointing ([`checkpoint`]), -//! streaming/events ([`stream`]), run-status snapshots ([`status`]), graph -//! export/visualization ([`export`]), dependency-DAG validation ([`dag`]), -//! subgraph embedding ([`subgraph`]), and checkpointed workflow execution. -//! -//! Each concern lives in its own submodule with `types.rs` (definitions), -//! `mod.rs` (implementations), and `test.rs` (unit tests). - -pub mod agent_loop; -pub mod builder; -pub mod cache; -pub mod channel; -pub mod checkpoint; -pub mod command; -pub mod compiled; -pub mod dag; -pub mod delegation; -mod execution_locks; -pub mod export; -pub mod observability; -pub mod orchestration; -pub mod parallel; -pub mod recursion; -pub mod reducer; -pub mod status; -pub mod storage; -pub mod stream; -pub mod subagent_node; -pub mod subgraph; -pub mod testkit; - -pub use tinyagents_harness::error::{Result, TinyAgentsError}; - -#[cfg(test)] -mod test_support; - -// --- Durable execution model --- -pub use builder::{ - END, ForkId, GraphBuilder, GraphDefaults, IdleClock, NodeCachePolicy, NodeContext, NodeFuture, - NodeHandler, NodePolicy, Route, RouterFn, START, -}; -pub use cache::{DriverTaskCache, InMemoryTaskCache, TaskCache, TaskCacheKey}; -pub use channel::{ - Barrier, BinaryAggregate, Channel, ChannelSet, ChannelState, ChannelUpdate, ChannelWrite, - Delta, Ephemeral, LastValue, Messages, NamedBarrier, ReducerRegistry, Topic, Untracked, -}; -pub use checkpoint::DriverCheckpointer; -pub use checkpoint::{ - BarrierArrivals, CHECKPOINT_FORMAT_VERSION, Checkpoint, CheckpointConfig, CheckpointMetadata, - CheckpointSource, CheckpointTuple, Checkpointer, CompletedTask, DurabilityMode, - InMemoryCheckpointer, PendingActivation, PendingWrite, -}; -pub use command::{Command, Interrupt, NodeResult, RouteTarget, Send}; -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, - DelegationStage, DelegationStageOutput, DelegationState, PendingApproval, StepRecord, - delegation_graph_topology, deny_decision, resume_delegation, run_delegation, - run_delegation_durable, run_or_resume_delegation, -}; -pub use export::{ - ChannelInfo, ConditionalEdgeInfo, EdgeInfo, GraphPolicySummary, GraphTopology, - GraphTopologyReport, NodeInfo, NodePolicySummary, RouteInfo, ValidationReport, WaitingEdgeInfo, - from_json, to_json, to_mermaid, -}; -pub use observability::{ - GraphEventJournal, GraphHealthSummary, GraphLangfuseExporter, GraphLatencyMetrics, - GraphNodeHealth, GraphNodeLatency, GraphObservation, GraphStatusStore, GraphStepLatency, - GraphTracingSink, InMemoryGraphEventJournal, InMemoryGraphStatusStore, JournalGraphSink, - SpanMetadataFn, StoreGraphEventJournal, build_graph_langfuse_batch, -}; -pub use orchestration::{ - CancelledDetachedTask, DetachedTaskRegistry, DetachedTaskRegistryError, DetachedTaskSnapshot, - DetachedTaskWaitOutcome, DriverTaskStore, InMemoryTaskStore, OrchestrationControlOutcome, - OrchestrationTaskFilter, OrchestrationTaskKind, OrchestrationTaskRecord, - OrchestrationTaskResult, OrchestrationTaskSpec, OrchestrationTaskStatus, SteeringRegistry, - TaskStore, -}; -pub use parallel::{FailurePolicy, ItemOutcome, ParallelOptions, ParallelOutcome, map_reduce}; -pub use recursion::{ - ChildRun, ChildRunSink, RecursionFrame, RecursionPolicy, RecursionStack, RunTree, -}; -pub use reducer::{ - AppendReducer, ClosureReducer, ClosureStateReducer, MaxReducer, MinReducer, OverwriteReducer, - OverwriteStateReducer, Reducer, SetUnionReducer, StateReducer, -}; -pub use status::GraphRunStatus; -pub use storage::{ - collections, import_legacy_leases, import_legacy_task_cache, open_storage, storage_schema, -}; -pub use stream::{ - CollectingSink, GraphEvent, GraphEventEnvelope, GraphEventSink, NoopSink, StreamMode, - StreamProjection, project_graph_event, -}; -#[allow(deprecated)] -pub use subagent_node::{ - AgentInvocation, AgentInvocationBinding, AgentInvoker, InputMapper, OutputMapper, - SubAgentBudget, SubAgentInput, SubAgentNode, SubAgentOutput, SubAgentPolicy, subagent_node, -}; -pub use subgraph::{adapter_subgraph_node, shared_subgraph_node}; -pub use testkit::{ - GraphAssertions, GraphEventRecorder, GraphRun, RetryCountingNode, StreamCollector, - assert_graph, failing_node, fanout_node, interrupting_node, noop_node, run_recorded, - scripted_route_node, scripted_update_node, subagent_fake_node, subgraph_test_node, -}; diff --git a/crates/tinyagents-graph/src/observability/README.md b/crates/tinyagents-graph/src/observability/README.md deleted file mode 100644 index e00ce4fea..000000000 --- a/crates/tinyagents-graph/src/observability/README.md +++ /dev/null @@ -1,91 +0,0 @@ -# graph::observability - -Durable observability for the graph runtime — journals, status stores, and -the journaling event sink. - -The live `graph::stream` layer emits transient `GraphEvent`s into an -in-process `GraphEventSink`. This module makes that history **durable and -correlatable** so a UI, supervisor, or test can reconstruct a recursive graph -run tree after the fact. - -`CompiledGraph` wires into this module through builder-style -`with_status_store` and `with_event_journal`; both are opt-in and default off -so existing runs are unchanged. - -## Public surface - -- `GraphObservation` — a durable envelope pairing a `GraphEvent` with its run - lineage (`run_id` / `parent_run_id` / `root_run_id`), `graph_id`, - `checkpoint_id`, subgraph `namespace`, `step`, `offset`, and timestamp. This - is the unit everything else in the module is built from. -- `GraphEventJournal` (trait) — an append-only, offset-addressable journal of - observations. - - `InMemoryGraphEventJournal` — in-process implementation for tests. - - `StoreGraphEventJournal` — store-backed implementation; stream key is - the run id, so a `harness::store::AppendStore` durably persists a run's - observations. -- `GraphStatusStore` (trait) — a compact "what is running now?" surface over - `graph::GraphRunStatus`. - - `InMemoryGraphStatusStore` — in-process implementation with a - `thread_id → run ids` index so `list_by_thread` is proportional to that - thread's runs. Unbounded by default; `with_max_runs(n)` caps retained - runs, evicting the oldest terminal run first (oldest live run when none - is terminal). -- `JournalGraphSink` — a `GraphEventSink` that wraps each emitted event into a - `GraphObservation` and appends it to a journal, optionally also forwarding - to a live `inner` sink (`with_lineage`, `with_thread`, `with_namespace`, - `with_inner` builders). -- `GraphLatencyMetrics` — per-step/per-node timing rollups derived from a - run's observations (`from_observations`, `average_step_ms`, - `average_node_ms`). -- `GraphHealthSummary` — per-node success/failure counts derived from - observations — node-level **tool health** telemetry (`from_observations`, - `from_status`). -- `GraphLangfuseExporter` (`langfuse/`) — exports a run's observations to - Langfuse, turning supersteps and nodes into timed spans (failures promoted - to `ERROR`) and attaching the health summary to the trace. It emits a - structural graph-run span (`{trace}:run:{run_id}`) that steps, subgraphs, and - point events nest under. It shares the harness `LangfuseClient` transport and - defaults its `traceId` to the run's `root_run_id`; because the graph-run span - uses the same id scheme the harness exporter parents its agent run spans to, - a graph run and the agent/tool runs its nodes spawn nest into one trace tree. - `send_observations_otlp` / `build_otlp_requests` export the same trace to - Langfuse v4 over OTLP (trace attributes on every span, deterministic span - ids); the batch methods remain for legacy-mode deployments. - -## Persistence bridge - -`JournalGraphSink` bridges the synchronous `GraphEventSink::emit` hook to the -async journal API through a background `AppendWorker`: `emit` hands the -observation to a bounded channel drained on a dedicated thread, so it never -blocks the executor on I/O. Persistence is **best-effort**: a full queue drops -(and counts) the overflow, backend errors are counted and reported through -`tracing` (rate-limited to one report per failure run plus a reminder every five -minutes — see the harness `observability` README for the full error policy) -rather than propagated, and neither aborts the run. The executor calls -`GraphEventSink::flush` -after the terminal run event (and callers can call it directly) to block until -the durable log has caught up. Do not rely on the sink for delivery guarantees -stronger than "usually persisted, never run-blocking." - -## Files - -| File | Role | -| --- | --- | -| `types.rs` | Every public type: `GraphObservation`, journal/status-store traits and in-memory impls, `JournalGraphSink`, latency/health rollups. | -| `mod.rs` | Behavioral code: rollup computation, journal/store/sink impls. | -| `langfuse/` | `GraphLangfuseExporter` and its span-construction logic. | -| `mod_tests.rs` | Unit tests (journal round-trips, sink lineage, latency/health rollups). | - -## Operational constraints - -- `StoreGraphEventJournal` keys the stream by run id; reusing a run id across - logically distinct runs interleaves their observations. -- The health summary's failure counts are node-scoped, not step-scoped — a - node retried and eventually succeeded still contributes its earlier - failures to `GraphHealthSummary`, by design (it is a *tool health* signal, - not a final-status signal). -- The Langfuse exporter shares transport code with - `harness::observability::LangfuseClient` (via `harness::observability`'s - crate-visible `clean_nulls`/`iso_ms` helpers); keep timestamp/null-pruning - behavior in sync between the two if either changes. diff --git a/crates/tinyagents-graph/src/observability/langfuse.rs b/crates/tinyagents-graph/src/observability/langfuse.rs deleted file mode 100644 index 7edaa9484..000000000 --- a/crates/tinyagents-graph/src/observability/langfuse.rs +++ /dev/null @@ -1,660 +0,0 @@ -//! Langfuse ingestion exporter for durable graph observations. -//! -//! Where [`tinyagents_harness::observability::LangfuseClient`] exports an *agent* -//! run's observations (model generations, tool calls), this exporter turns a -//! *graph* run's durable [`GraphObservation`] stream into a Langfuse trace: each -//! superstep and node handler becomes a timed span, node/subgraph failures are -//! promoted to `ERROR` level, and per-node **health telemetry** rides along on -//! the trace metadata. -//! -//! # Unified traces across the graph and its agents -//! -//! Graph nodes frequently delegate to agents and tools (see -//! [`crate::SubAgentNode`]). Those child agent runs share the graph run's -//! `root_run_id`, and both exporters default their Langfuse `traceId` to that -//! root run id. Sending a graph run's observations with this exporter **and** -//! the child agent runs' observations with the harness -//! [`LangfuseClient`](tinyagents_harness::observability::LangfuseClient) therefore -//! lands every graph step, node, model generation, and tool call under one -//! Langfuse trace — full end-to-end telemetry including tool health. -//! -//! # Langfuse v4 -//! -//! Langfuse v4 ingests traces only over OTLP. [`GraphLangfuseExporter::send_observations_otlp`] -//! converts the batch with the harness `trace_export::ingestion_otlp` -//! projection, which copies trace-level attributes onto every span and keeps -//! the deterministic observation ids as span ids, so the unified-trace -//! nesting above holds on v4 too. The batch methods remain for inspection and -//! for Langfuse deployments still on the legacy path. -//! -//! The exporter is pull-based and best-effort by construction: build the batch -//! from a completed (or in-progress) observation slice read back from a -//! [`GraphEventJournal`](crate::GraphEventJournal), then send it. It does -//! not sit on the hot path of graph execution. - -use std::collections::HashMap; -use std::collections::VecDeque; -use std::sync::Arc; - -use serde_json::{Map, Value, json}; - -use crate::observability::{GraphHealthSummary, GraphObservation}; -use crate::stream::GraphEvent; -use tinyagents_harness::error::{Result, TinyAgentsError}; -use tinyagents_harness::ids::NodeId; -use tinyagents_harness::observability::trace_export::ingestion_otlp::ingestion_batch_otlp_requests; -use tinyagents_harness::observability::{LangfuseClient, LangfuseTraceConfig, clean_nulls, iso_ms}; - -/// A FIFO queue of pending span starts, each `(observation index, ts_ms)`. -/// Duplicate span identities (a node re-run in a later step, say) queue in -/// arrival order and pair with their terminals FIFO. -type StartQueue = VecDeque<(usize, u64)>; - -/// A host-supplied hook that contributes extra metadata to every span the -/// exporter emits. Called with the span's coordinate observation; returning -/// `None` (or an empty map) contributes nothing. Host keys are merged **over** -/// the exporter's built-in coordinate keys, so a host can also override them. -pub type SpanMetadataFn = dyn Fn(&GraphObservation) -> Option> + Send + Sync; - -mod batch; -pub use batch::build_graph_langfuse_batch; - -/// Async Langfuse exporter for durable graph observations. -/// -/// Wraps a shared [`LangfuseClient`] for transport (auth, endpoint -/// normalization, `207 Multi-Status` handling) and adds graph-aware payload -/// construction on top. -#[derive(Clone)] -pub struct GraphLangfuseExporter { - client: LangfuseClient, - /// Optional host metadata injector, merged into every span's metadata. - span_metadata_fn: Option>, -} - -impl std::fmt::Debug for GraphLangfuseExporter { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("GraphLangfuseExporter") - .field("client", &self.client) - .field( - "span_metadata_fn", - &self.span_metadata_fn.as_ref().map(|_| "Fn(..)"), - ) - .finish() - } -} - -impl GraphLangfuseExporter { - /// Wraps an existing [`LangfuseClient`] as a graph exporter. - pub fn new(client: LangfuseClient) -> Self { - Self { - client, - span_metadata_fn: None, - } - } - - /// Builds an exporter from the same environment variables as - /// [`LangfuseClient::from_env`]. - pub fn from_env() -> Result { - Ok(Self::new(LangfuseClient::from_env()?)) - } - - /// Installs a host metadata injector: `f` is called once per emitted span - /// with the span's coordinate observation, and any returned keys are merged - /// into that span's metadata (host keys win on collision). Default: none. - pub fn with_span_metadata_fn( - mut self, - f: impl Fn(&GraphObservation) -> Option> + Send + Sync + 'static, - ) -> Self { - self.span_metadata_fn = Some(Arc::new(f)); - self - } - - /// Returns the underlying transport client. - pub fn client(&self) -> &LangfuseClient { - &self.client - } - - /// Returns the normalized ingestion endpoint. - pub fn endpoint(&self) -> &str { - self.client.endpoint() - } - - /// Builds a Langfuse ingestion payload for `observations` without sending it. - /// - /// The batch always begins with a `trace-create` event, followed by a - /// timed `span-create` for every completed (or still-running) superstep, - /// node, and subgraph, then an `event-create` for every remaining - /// observation (routes, checkpoints, interrupts, custom writes, run - /// lifecycle). Node health telemetry is attached to the trace metadata. - pub fn build_ingestion_batch( - &self, - trace: LangfuseTraceConfig, - observations: &[GraphObservation], - ) -> Result { - build_graph_langfuse_batch(trace, observations, self.span_metadata_fn.as_deref()) - } - - /// Builds the Langfuse v4 OTLP/HTTP JSON requests for `observations` - /// without sending them: the trace of [`Self::build_ingestion_batch`], - /// with trace name, session, tags, metadata (including node health) and - /// release copied onto every span, under resource `service.name = - /// service` and scope `.graph`. - pub fn build_otlp_requests( - &self, - trace: LangfuseTraceConfig, - observations: &[GraphObservation], - service: &str, - ) -> Result> { - let payload = self.build_ingestion_batch(trace, observations)?; - ingestion_batch_otlp_requests(&payload, service, &format!("{service}.graph")) - } - - /// Sends `observations` to Langfuse v4 over OTLP - /// ([`Self::build_otlp_requests`] through - /// [`LangfuseClient::send_otlp`]), stopping at the first failed request. - /// Prefer this over [`Self::send_observations`]: Langfuse v4 no longer - /// accepts traces on the legacy batch path. - pub async fn send_observations_otlp( - &self, - trace: LangfuseTraceConfig, - observations: &[GraphObservation], - service: &str, - ) -> Result> { - let requests = self.build_otlp_requests(trace, observations, service)?; - let mut responses = Vec::with_capacity(requests.len()); - for request in requests { - responses.push(self.client.send_otlp(request).await?); - } - Ok(responses) - } - - /// Sends `observations` as one legacy Langfuse ingestion batch. - /// - /// Langfuse Cloud removes this path for traces on 2026-11-16; use - /// [`Self::send_observations_otlp`]. - pub async fn send_observations( - &self, - trace: LangfuseTraceConfig, - observations: &[GraphObservation], - ) -> Result { - let payload = self.build_ingestion_batch(trace, observations)?; - self.client.send_batch(payload).await - } -} - -/// The stable Langfuse observation id for a run's span, shared with the harness -/// exporter (`{trace_id}:run:{run_id}`) so a graph run and the harness agent -/// runs its nodes spawn resolve to one nested tree under the same trace. -fn run_span_id(trace_id: &str, run_id: &str) -> String { - format!("{trace_id}:run:{run_id}") -} - -/// Builds the structural `span-create` for the graph run, parented to the trace. -/// -/// The span brackets the run: it starts at the first observation and ends at the -/// last (or at a terminal `RunCompleted`/`RunFailed` when present). It is a plain -/// structural anchor — run failure is still surfaced by the dedicated -/// `run.failed` point event — so its only job is to give steps, subgraphs, and -/// any child agent runs a single parent to nest under. -fn graph_run_span( - trace_id: &str, - run_span_id: &str, - observations: &[GraphObservation], - first: &GraphObservation, -) -> Value { - let start_ts = first.ts_ms; - // Prefer a terminal run event's timestamp; else close at the last observation. - let end_ts = observations - .iter() - .rev() - .find(|o| { - matches!( - o.event, - GraphEvent::RunCompleted { .. } | GraphEvent::RunFailed { .. } - ) - }) - .or_else(|| observations.last()) - .map(|o| o.ts_ms) - .unwrap_or(start_ts); - json!({ - "id": run_span_id, - "timestamp": iso_ms(end_ts), - "type": "span-create", - "body": clean_nulls(json!({ - "id": run_span_id, - "traceId": trace_id, - "name": first.graph_id.as_str(), - "startTime": iso_ms(start_ts), - "endTime": iso_ms(end_ts), - "metadata": json!({ - "run_id": first.run_id.as_str(), - "root_run_id": first.root_run_id.as_str(), - "graph_id": first.graph_id.as_str(), - }), - })), - }) -} - -/// Resolves the Langfuse trace id: the configured id when set, else the first -/// observation's root run id so it aligns with the harness agent exporter. -fn resolve_trace_id(trace: &LangfuseTraceConfig, observations: &[GraphObservation]) -> String { - if let Some(id) = &trace.trace_id - && !id.trim().is_empty() - { - return id.clone(); - } - observations[0].root_run_id.as_str().to_string() -} - -/// Builds the `trace-create` batch event, defaulting the trace name to the -/// graph id and the session to the run's thread, and folding node-health -/// telemetry plus graph coordinates into the trace metadata. -fn trace_create( - trace_id: &str, - trace_ts: &str, - trace: &LangfuseTraceConfig, - first: &GraphObservation, - health: &GraphHealthSummary, -) -> Value { - let name = trace - .name - .clone() - .unwrap_or_else(|| first.graph_id.as_str().to_string()); - let session_id = trace - .session_id - .clone() - .or_else(|| first.thread_id.as_ref().map(|t| t.as_str().to_string())); - let mut metadata = json!({ - "graph_id": first.graph_id.as_str(), - "root_run_id": first.root_run_id.as_str(), - "health": health, - }); - if !first.namespace.is_empty() - && let Value::Object(map) = &mut metadata - { - map.insert("namespace".to_string(), json!(first.namespace)); - } - if let (Value::Object(dst), Value::Object(extra)) = (&mut metadata, &trace.metadata) { - for (k, v) in extra { - dst.insert(k.clone(), v.clone()); - } - } - - json!({ - "id": format!("{trace_id}:trace"), - "timestamp": trace_ts, - "type": "trace-create", - "body": clean_nulls(json!({ - "id": trace_id, - "timestamp": trace_ts, - "name": name, - "userId": trace.user_id, - "sessionId": session_id, - "environment": trace.environment, - "release": trace.release, - "version": trace.version, - "tags": if trace.tags.is_empty() { Value::Null } else { json!(trace.tags) }, - "metadata": metadata, - })), - }) -} - -/// Emits a timed `span-create` for every step, node, and subgraph by pairing -/// each start observation with its terminal one (FIFO for duplicate keys). -/// Both the start and terminal indices are marked `consumed` so they are not -/// re-emitted as point events. Unpaired starts become open spans (start only). -fn push_span_events( - trace_id: &str, - run_parent: &str, - observations: &[GraphObservation], - consumed: &mut [bool], - batch: &mut Vec, - injector: Option<&SpanMetadataFn>, -) { - // Pending starts keyed by their span identity, each holding (index, ts). - let mut step_starts: HashMap = HashMap::new(); - let mut node_starts: HashMap<(NodeId, usize), StartQueue> = HashMap::new(); - let mut subgraph_starts: HashMap<(NodeId, Vec), StartQueue> = HashMap::new(); - - for (idx, obs) in observations.iter().enumerate() { - match &obs.event { - GraphEvent::StepStarted { step, .. } => { - step_starts - .entry(*step) - .or_default() - .push_back((idx, obs.ts_ms)); - consumed[idx] = true; - } - GraphEvent::StepCompleted { step } => { - if let Some((start_idx, start_ts)) = pop(&mut step_starts, step) { - consumed[idx] = true; - batch.push(step_span( - trace_id, - run_parent, - *step, - start_ts, - Some(obs.ts_ms), - obs, - &observations[start_idx], - injector, - )); - } - } - GraphEvent::NodeStarted { node, step } => { - node_starts - .entry((node.clone(), *step)) - .or_default() - .push_back((idx, obs.ts_ms)); - consumed[idx] = true; - } - GraphEvent::NodeCompleted { node, step } => { - if let Some((_, start_ts)) = pop(&mut node_starts, &(node.clone(), *step)) { - consumed[idx] = true; - batch.push(node_span( - trace_id, - node, - *step, - start_ts, - Some(obs.ts_ms), - None, - obs, - injector, - )); - } - } - GraphEvent::NodeFailed { node, step, error } => { - if let Some((_, start_ts)) = pop(&mut node_starts, &(node.clone(), *step)) { - consumed[idx] = true; - batch.push(node_span( - trace_id, - node, - *step, - start_ts, - Some(obs.ts_ms), - Some(error.as_str()), - obs, - injector, - )); - } - } - GraphEvent::SubgraphStarted { node, namespace } => { - subgraph_starts - .entry((node.clone(), namespace.clone())) - .or_default() - .push_back((idx, obs.ts_ms)); - consumed[idx] = true; - } - GraphEvent::SubgraphCompleted { node, namespace } => { - if let Some((_, start_ts)) = - pop(&mut subgraph_starts, &(node.clone(), namespace.clone())) - { - consumed[idx] = true; - batch.push(subgraph_span( - trace_id, - run_parent, - node, - namespace, - start_ts, - Some(obs.ts_ms), - obs, - injector, - )); - } - } - _ => {} - } - } - - // Any still-open starts become spans with a start time but no end time. - for (step, mut queue) in step_starts { - while let Some((start_idx, start_ts)) = queue.pop_front() { - let obs = &observations[start_idx]; - batch.push(step_span( - trace_id, run_parent, step, start_ts, None, obs, obs, injector, - )); - } - } - for ((node, step), mut queue) in node_starts { - while let Some((start_idx, start_ts)) = queue.pop_front() { - batch.push(node_span( - trace_id, - &node, - step, - start_ts, - None, - None, - &observations[start_idx], - injector, - )); - } - } - for ((node, namespace), mut queue) in subgraph_starts { - while let Some((start_idx, start_ts)) = queue.pop_front() { - batch.push(subgraph_span( - trace_id, - run_parent, - &node, - &namespace, - start_ts, - None, - &observations[start_idx], - injector, - )); - } - } -} - -/// Emits an `event-create` for every observation not already represented by a -/// span, mapping `run.failed` to `ERROR` level with the rendered error. -fn push_point_events( - trace_id: &str, - run_parent: &str, - observations: &[GraphObservation], - consumed: &[bool], - batch: &mut Vec, -) { - for (idx, obs) in observations.iter().enumerate() { - if consumed[idx] { - continue; - } - // The trace itself represents the run start; skip the duplicate event. - if matches!(obs.event, GraphEvent::RunStarted { .. }) { - continue; - } - let ts = iso_ms(obs.ts_ms); - let (level, status) = match &obs.event { - GraphEvent::RunFailed { error, .. } => (Some("ERROR"), Some(error.clone())), - _ => (None, None), - }; - batch.push(json!({ - "id": obs.event_id.as_str(), - "timestamp": ts, - "type": "event-create", - "body": clean_nulls(json!({ - "id": obs.event_id.as_str(), - "traceId": trace_id, - "parentObservationId": run_parent, - "name": obs.event.kind(), - "startTime": ts, - "level": level, - "statusMessage": status, - "metadata": span_metadata(obs, None, None, None), - })), - })); - } -} - -/// Builds a `span-create` for a superstep, parented directly to the trace. -#[allow(clippy::too_many_arguments)] -fn step_span( - trace_id: &str, - run_parent: &str, - step: usize, - start_ts: u64, - end_ts: Option, - terminal: &GraphObservation, - start: &GraphObservation, - injector: Option<&SpanMetadataFn>, -) -> Value { - let metadata = span_metadata(start, None, Some(step), injector); - span_event( - trace_id, - &format!("{trace_id}:step:{step}"), - Some(run_parent.to_string()), - &format!("step {step}"), - start_ts, - end_ts, - None, - terminal, - metadata, - ) -} - -/// Builds a `span-create` for a node handler, parented to its superstep span. -/// -/// Node spans carry the Langfuse **Agent Graph view** keys `langgraph_node` -/// (the node id) and `langgraph_step` (the superstep index) in metadata, so -/// Langfuse can lay the trace out as a graph. -#[allow(clippy::too_many_arguments)] -fn node_span( - trace_id: &str, - node: &NodeId, - step: usize, - start_ts: u64, - end_ts: Option, - error: Option<&str>, - terminal: &GraphObservation, - injector: Option<&SpanMetadataFn>, -) -> Value { - let metadata = span_metadata(terminal, Some(node.as_str()), Some(step), injector); - span_event( - trace_id, - &format!("{trace_id}:node:{}:{step}", node.as_str()), - Some(format!("{trace_id}:step:{step}")), - node.as_str(), - start_ts, - end_ts, - error, - terminal, - metadata, - ) -} - -/// Builds a `span-create` for an embedded subgraph, parented to the trace. -#[allow(clippy::too_many_arguments)] -fn subgraph_span( - trace_id: &str, - run_parent: &str, - node: &NodeId, - namespace: &[String], - start_ts: u64, - end_ts: Option, - terminal: &GraphObservation, - injector: Option<&SpanMetadataFn>, -) -> Value { - let metadata = span_metadata(terminal, None, Some(terminal.step), injector); - span_event( - trace_id, - &format!("{trace_id}:subgraph:{}", namespace.join("/")), - Some(run_parent.to_string()), - &format!("subgraph {}", node.as_str()), - start_ts, - end_ts, - None, - terminal, - metadata, - ) -} - -/// Shared `span-create` builder. `error` promotes the span to `ERROR` level. -/// The batch item id comes from the terminal observation so it is unique; the -/// caller supplies the fully-built span `metadata`. -#[allow(clippy::too_many_arguments)] -fn span_event( - trace_id: &str, - span_id: &str, - parent: Option, - name: &str, - start_ts: u64, - end_ts: Option, - error: Option<&str>, - terminal: &GraphObservation, - metadata: Value, -) -> Value { - let start_iso = iso_ms(start_ts); - let end_iso = end_ts.map(iso_ms); - let (level, status) = match error { - Some(err) => (Some("ERROR"), Some(err.to_string())), - None => (None, None), - }; - json!({ - "id": terminal.event_id.as_str(), - "timestamp": end_iso.clone().unwrap_or_else(|| start_iso.clone()), - "type": "span-create", - "body": clean_nulls(json!({ - "id": span_id, - "traceId": trace_id, - "parentObservationId": parent, - "name": name, - "startTime": start_iso, - "endTime": end_iso, - "level": level, - "statusMessage": status, - "metadata": metadata, - })), - }) -} - -/// Extracts the correlation coordinates every span/event carries in metadata, -/// stamping the Langfuse Agent-Graph-view keys (`langgraph_node`, -/// `langgraph_step`) when the caller supplies them and merging any -/// host-injected keys last (host keys win on collision). -fn span_metadata( - obs: &GraphObservation, - langgraph_node: Option<&str>, - langgraph_step: Option, - injector: Option<&SpanMetadataFn>, -) -> Value { - let mut metadata = json!({ - "run_id": obs.run_id.as_str(), - "root_run_id": obs.root_run_id.as_str(), - "parent_run_id": obs.parent_run_id.as_ref().map(|id| id.as_str()), - "graph_id": obs.graph_id.as_str(), - "checkpoint_id": obs.checkpoint_id.as_ref().map(|id| id.as_str()), - "namespace": if obs.namespace.is_empty() { Value::Null } else { json!(obs.namespace) }, - "step": obs.step, - "offset": obs.offset, - "event": obs.event, - }); - if let Value::Object(map) = &mut metadata { - if let Some(node) = langgraph_node { - map.insert("langgraph_node".to_string(), json!(node)); - } - if let Some(step) = langgraph_step { - map.insert("langgraph_step".to_string(), json!(step)); - } - if let Some(extra) = injector.and_then(|f| f(obs)) { - for (k, v) in extra { - map.insert(k, v); - } - } - } - metadata -} - -/// Pops the FIFO-oldest pending start for `key`, cleaning up empty queues. -fn pop( - starts: &mut HashMap, - key: &K, -) -> Option<(usize, u64)> { - let queue = starts.get_mut(key)?; - let popped = queue.pop_front(); - if queue.is_empty() { - starts.remove(key); - } - popped -} - -#[cfg(test)] -#[path = "langfuse/langfuse_tests.rs"] -mod test; diff --git a/crates/tinyagents-graph/src/observability/langfuse/batch.rs b/crates/tinyagents-graph/src/observability/langfuse/batch.rs deleted file mode 100644 index 9f9a6424c..000000000 --- a/crates/tinyagents-graph/src/observability/langfuse/batch.rs +++ /dev/null @@ -1,56 +0,0 @@ -//! Transport-independent Langfuse graph payload construction. - -use super::*; - -/// Build one graph ingestion batch without constructing a transport client. -/// -/// Use this when the host supplies authentication, endpoint selection, or an -/// export sink separately. The result has the same trace, structural run span, -/// graph spans, point events, and health metadata as -/// [`GraphLangfuseExporter::build_ingestion_batch`]. No environment variables, -/// credentials, or network access are required. -/// -/// `span_metadata_fn` receives each span's coordinate observation. Its keys -/// override the built-in coordinate metadata, matching the exporter's hook. -/// An empty observation slice returns [`TinyAgentsError::Validation`]. -pub fn build_graph_langfuse_batch( - trace: LangfuseTraceConfig, - observations: &[GraphObservation], - span_metadata_fn: Option<&SpanMetadataFn>, -) -> Result { - if observations.is_empty() { - return Err(TinyAgentsError::Validation( - "at least one observation is required".to_string(), - )); - } - let trace_id = resolve_trace_id(&trace, observations); - let first = &observations[0]; - let trace_ts = iso_ms(first.ts_ms); - let health = GraphHealthSummary::from_observations(observations); - - let mut batch = Vec::with_capacity(observations.len() + 2); - batch.push(trace_create(&trace_id, &trace_ts, &trace, first, &health)); - - // A structural span for the graph run itself, sitting directly under the - // trace. Every step, subgraph, and point event nests under it, and — key - // for a unified trace — its id follows the same `{trace}:run:{run_id}` - // scheme the harness exporter parents its run spans to. A sub-agent that - // a graph node spawns carries `parent_run_id == graph run id`, so its - // harness-exported run span nests directly under this graph-run span - // rather than floating at the trace root. - let run_parent = run_span_id(&trace_id, first.root_run_id.as_str()); - batch.push(graph_run_span(&trace_id, &run_parent, observations, first)); - - let mut consumed = vec![false; observations.len()]; - push_span_events( - &trace_id, - &run_parent, - observations, - &mut consumed, - &mut batch, - span_metadata_fn, - ); - push_point_events(&trace_id, &run_parent, observations, &consumed, &mut batch); - - Ok(json!({ "batch": batch })) -} diff --git a/crates/tinyagents-graph/src/observability/langfuse/langfuse_tests.rs b/crates/tinyagents-graph/src/observability/langfuse/langfuse_tests.rs deleted file mode 100644 index 4b960213c..000000000 --- a/crates/tinyagents-graph/src/observability/langfuse/langfuse_tests.rs +++ /dev/null @@ -1,518 +0,0 @@ -//! Unit tests for the graph Langfuse exporter: batch shape, timed spans, -//! failure-level promotion, health telemetry, and trace-id alignment. - -use super::*; -use crate::observability::GraphObservation; -use crate::stream::GraphEvent; -use tinyagents_harness::ids::{CheckpointId, EventId, GraphId, NodeId, RunId, ThreadId}; - -/// `2024-01-01T00:00:00.000Z` in epoch millis; test deltas add onto this so ISO -/// assertions read as `…:00.0NNZ`. -const BASE: u64 = 1_704_067_200_000; - -/// Builds an observation `delta_ms` after [`BASE`] with the given -/// offset/step/event under a fixed run. -fn obs(offset: u64, step: usize, delta_ms: u64, event: GraphEvent) -> GraphObservation { - let ts_ms = BASE + delta_ms; - GraphObservation { - event_id: EventId::new(format!("evt-{offset}")), - run_id: RunId::new("run-1"), - root_run_id: RunId::new("root-1"), - parent_run_id: None, - thread_id: Some(ThreadId::new("thread-1")), - graph_id: GraphId::new("demo-graph"), - checkpoint_id: None, - namespace: Vec::new(), - step, - offset, - ts_ms, - event, - } -} - -/// A representative run: start → step → node a (ok) → node b (fail) → step end -/// → run fail. -fn sample_run() -> Vec { - let a = NodeId::new("a"); - let b = NodeId::new("b"); - vec![ - obs( - 0, - 0, - 1_000, - GraphEvent::RunStarted { - run_id: RunId::new("run-1"), - }, - ), - obs( - 1, - 1, - 1_010, - GraphEvent::StepStarted { - step: 1, - active: vec![a.clone(), b.clone()], - }, - ), - obs( - 2, - 1, - 1_020, - GraphEvent::NodeStarted { - node: a.clone(), - step: 1, - }, - ), - obs( - 3, - 1, - 1_045, - GraphEvent::NodeCompleted { - node: a.clone(), - step: 1, - }, - ), - obs( - 4, - 1, - 1_050, - GraphEvent::NodeStarted { - node: b.clone(), - step: 1, - }, - ), - obs( - 5, - 1, - 1_080, - GraphEvent::NodeFailed { - node: b.clone(), - step: 1, - error: "boom".to_string(), - }, - ), - obs(6, 1, 1_090, GraphEvent::StepCompleted { step: 1 }), - obs( - 7, - 1, - 1_100, - GraphEvent::RunFailed { - run_id: RunId::new("run-1"), - error: "boom".to_string(), - }, - ), - ] -} - -fn exporter() -> GraphLangfuseExporter { - GraphLangfuseExporter::new( - LangfuseClient::proxy("https://backend.test/telemetry/langfuse/ingestion", "tok").unwrap(), - ) -} - -/// Finds the first batch item of the given `type` whose body `id` matches. -fn find<'a>(batch: &'a [Value], ty: &str, id: &str) -> Option<&'a Value> { - batch - .iter() - .find(|e| e["type"] == ty && e["body"]["id"] == id) -} - -#[test] -fn empty_observations_are_rejected() { - let err = exporter() - .build_ingestion_batch(LangfuseTraceConfig::default(), &[]) - .unwrap_err(); - assert!(matches!(err, TinyAgentsError::Validation(_))); -} - -#[test] -fn trace_defaults_to_root_run_and_graph_id() { - let batch = exporter() - .build_ingestion_batch(LangfuseTraceConfig::default(), &sample_run()) - .unwrap(); - let events = batch["batch"].as_array().unwrap(); - - let trace = &events[0]; - assert_eq!(trace["type"], "trace-create"); - // Trace id aligns with the harness exporter (root run id) so agent tool - // spans land under the same trace. - assert_eq!(trace["body"]["id"], "root-1"); - assert_eq!(trace["body"]["name"], "demo-graph"); - assert_eq!(trace["body"]["sessionId"], "thread-1"); -} - -#[test] -fn config_overrides_trace_id_name_and_session() { - let batch = exporter() - .build_ingestion_batch( - LangfuseTraceConfig { - trace_id: Some("custom-trace".to_string()), - name: Some("My Graph".to_string()), - session_id: Some("sess-9".to_string()), - user_id: Some("user-1".to_string()), - ..Default::default() - }, - &sample_run(), - ) - .unwrap(); - let events = batch["batch"].as_array().unwrap(); - let trace = &events[0]; - assert_eq!(trace["body"]["id"], "custom-trace"); - assert_eq!(trace["body"]["name"], "My Graph"); - assert_eq!(trace["body"]["sessionId"], "sess-9"); - assert_eq!(trace["body"]["userId"], "user-1"); - // Spans reference the overridden trace id. - let events: &[Value] = events; - assert!(find(events, "span-create", "custom-trace:step:1").is_some()); -} - -#[test] -fn steps_and_nodes_become_timed_spans() { - let batch = exporter() - .build_ingestion_batch(LangfuseTraceConfig::default(), &sample_run()) - .unwrap(); - let events = batch["batch"].as_array().unwrap(); - - // The graph run itself is a structural span under the trace, named for the - // graph, that every step nests beneath. Its id uses the shared - // `{trace}:run:{run}` scheme so harness sub-agent runs (parent_run_id = the - // graph run) nest under it too, unifying graph + agent into one trace tree. - let run_span = find(events, "span-create", "root-1:run:root-1").expect("graph run span"); - assert_eq!(run_span["body"]["name"], "demo-graph"); - assert!(run_span["body"].get("parentObservationId").is_none()); - assert_eq!(run_span["body"]["startTime"], "2024-01-01T00:00:01.000Z"); - - // Superstep span with real start/end times. - let step = find(events, "span-create", "root-1:step:1").expect("step span"); - assert_eq!(step["body"]["name"], "step 1"); - assert_eq!(step["body"]["startTime"], "2024-01-01T00:00:01.010Z"); - assert_eq!(step["body"]["endTime"], "2024-01-01T00:00:01.090Z"); - // Steps parent to the graph-run span, not straight to the trace. - assert_eq!(step["body"]["parentObservationId"], "root-1:run:root-1"); - - // Node a completed cleanly and is parented to its step span. - let node_a = find(events, "span-create", "root-1:node:a:1").expect("node a span"); - assert_eq!(node_a["body"]["name"], "a"); - assert_eq!(node_a["body"]["parentObservationId"], "root-1:step:1"); - assert!(node_a["body"].get("level").is_none()); - assert_eq!(node_a["body"]["endTime"], "2024-01-01T00:00:01.045Z"); -} - -#[test] -fn node_failure_is_promoted_to_error_span() { - let batch = exporter() - .build_ingestion_batch(LangfuseTraceConfig::default(), &sample_run()) - .unwrap(); - let events = batch["batch"].as_array().unwrap(); - - let node_b = find(events, "span-create", "root-1:node:b:1").expect("node b span"); - assert_eq!(node_b["body"]["level"], "ERROR"); - assert_eq!(node_b["body"]["statusMessage"], "boom"); - assert_eq!(node_b["body"]["endTime"], "2024-01-01T00:00:01.080Z"); -} - -#[test] -fn run_failed_is_an_error_event_and_run_started_is_not_duplicated() { - let batch = exporter() - .build_ingestion_batch(LangfuseTraceConfig::default(), &sample_run()) - .unwrap(); - let events = batch["batch"].as_array().unwrap(); - - // RunStarted is represented by the trace, not a separate event. - assert!( - !events - .iter() - .any(|e| e["body"]["name"] == "run.started" && e["type"] == "event-create") - ); - - let failed = events - .iter() - .find(|e| e["body"]["name"] == "run.failed") - .expect("run.failed event"); - assert_eq!(failed["type"], "event-create"); - assert_eq!(failed["body"]["level"], "ERROR"); - assert_eq!(failed["body"]["statusMessage"], "boom"); -} - -#[test] -fn health_telemetry_rides_on_trace_metadata() { - let batch = exporter() - .build_ingestion_batch(LangfuseTraceConfig::default(), &sample_run()) - .unwrap(); - let events = batch["batch"].as_array().unwrap(); - let health = &events[0]["body"]["metadata"]["health"]; - - assert_eq!(health["total_started"], 2); - assert_eq!(health["total_completed"], 1); - assert_eq!(health["total_failed"], 1); - assert_eq!(health["run_failed"], true); - let nodes = health["nodes"].as_array().unwrap(); - assert_eq!(nodes.len(), 2); - assert_eq!(nodes[0]["node"], "a"); - assert_eq!(nodes[1]["node"], "b"); - assert_eq!(nodes[1]["failed"], 1); -} - -#[test] -fn node_spans_carry_langgraph_agent_graph_view_keys() { - let batch = exporter() - .build_ingestion_batch(LangfuseTraceConfig::default(), &sample_run()) - .unwrap(); - let events = batch["batch"].as_array().unwrap(); - - // Node spans carry both Agent-Graph-view keys: the node id and superstep. - let node_a = find(events, "span-create", "root-1:node:a:1").expect("node a span"); - assert_eq!(node_a["body"]["metadata"]["langgraph_node"], "a"); - assert_eq!(node_a["body"]["metadata"]["langgraph_step"], 1); - let node_b = find(events, "span-create", "root-1:node:b:1").expect("node b span"); - assert_eq!(node_b["body"]["metadata"]["langgraph_node"], "b"); - assert_eq!(node_b["body"]["metadata"]["langgraph_step"], 1); - - // Step spans carry the superstep index but no node id. - let step = find(events, "span-create", "root-1:step:1").expect("step span"); - assert_eq!(step["body"]["metadata"]["langgraph_step"], 1); - assert!(step["body"]["metadata"].get("langgraph_node").is_none()); - - // Existing coordinate keys are preserved alongside the new ones. - assert_eq!(node_a["body"]["metadata"]["run_id"], "run-1"); - assert_eq!(node_a["body"]["metadata"]["graph_id"], "demo-graph"); -} - -#[test] -fn subgraph_spans_carry_langgraph_step() { - let sub = NodeId::new("sub"); - let ns = vec!["sub".to_string()]; - let observations = vec![ - obs( - 0, - 2, - 1_000, - GraphEvent::SubgraphStarted { - node: sub.clone(), - namespace: ns.clone(), - }, - ), - obs( - 1, - 2, - 1_050, - GraphEvent::SubgraphCompleted { - node: sub.clone(), - namespace: ns.clone(), - }, - ), - ]; - let batch = exporter() - .build_ingestion_batch(LangfuseTraceConfig::default(), &observations) - .unwrap(); - let events = batch["batch"].as_array().unwrap(); - let span = find(events, "span-create", "root-1:subgraph:sub").expect("subgraph span"); - assert_eq!(span["body"]["metadata"]["langgraph_step"], 2); - assert!(span["body"]["metadata"].get("langgraph_node").is_none()); -} - -#[test] -fn host_span_metadata_injector_merges_and_wins_on_collision() { - let exporter = exporter().with_span_metadata_fn(|obs| { - let mut extra = serde_json::Map::new(); - extra.insert("flow_id".to_string(), json!("flow-42")); - // Host keys win over built-ins on collision. - extra.insert("graph_id".to_string(), json!("host-override")); - // Prove the injector sees the observation it decorates. - extra.insert("seen_step".to_string(), json!(obs.step)); - Some(extra) - }); - let batch = exporter - .build_ingestion_batch(LangfuseTraceConfig::default(), &sample_run()) - .unwrap(); - let events = batch["batch"].as_array().unwrap(); - - let node_a = find(events, "span-create", "root-1:node:a:1").expect("node a span"); - assert_eq!(node_a["body"]["metadata"]["flow_id"], "flow-42"); - assert_eq!(node_a["body"]["metadata"]["graph_id"], "host-override"); - assert_eq!(node_a["body"]["metadata"]["seen_step"], 1); - // Built-in Agent-Graph-view keys survive the merge. - assert_eq!(node_a["body"]["metadata"]["langgraph_node"], "a"); - - let step = find(events, "span-create", "root-1:step:1").expect("step span"); - assert_eq!(step["body"]["metadata"]["flow_id"], "flow-42"); - - // Point events (e.g. run.failed) are untouched by the injector. - let failed = events - .iter() - .find(|e| e["body"]["name"] == "run.failed") - .expect("run.failed event"); - assert!(failed["body"]["metadata"].get("flow_id").is_none()); -} - -#[test] -fn open_node_span_has_start_but_no_end() { - let a = NodeId::new("a"); - let observations = vec![ - obs( - 0, - 1, - 1_000, - GraphEvent::StepStarted { - step: 1, - active: vec![a.clone()], - }, - ), - obs( - 1, - 1, - 1_010, - GraphEvent::NodeStarted { - node: a.clone(), - step: 1, - }, - ), - ]; - let batch = exporter() - .build_ingestion_batch(LangfuseTraceConfig::default(), &observations) - .unwrap(); - let events = batch["batch"].as_array().unwrap(); - - let node = find(events, "span-create", "root-1:node:a:1").expect("open node span"); - assert_eq!(node["body"]["startTime"], "2024-01-01T00:00:01.010Z"); - assert!(node["body"].get("endTime").is_none()); -} - -#[test] -fn checkpoint_events_carry_coordinates_in_metadata() { - let mut observations = sample_run(); - observations.push(GraphObservation { - checkpoint_id: Some(CheckpointId::new("ckpt-7")), - ..obs( - 8, - 1, - 1_110, - GraphEvent::CheckpointSaved { - checkpoint_id: CheckpointId::new("ckpt-7"), - step: Some(8), - }, - ) - }); - let batch = exporter() - .build_ingestion_batch(LangfuseTraceConfig::default(), &observations) - .unwrap(); - let events = batch["batch"].as_array().unwrap(); - - let ckpt = events - .iter() - .find(|e| e["body"]["name"] == "checkpoint.saved") - .expect("checkpoint event"); - assert_eq!(ckpt["body"]["metadata"]["checkpoint_id"], "ckpt-7"); - assert_eq!(ckpt["body"]["metadata"]["step"], 1); -} - -#[test] -fn transport_free_batch_preserves_graph_coordinates_and_host_metadata() { - let observations = sample_run(); - let metadata = |_observation: &GraphObservation| { - Some(Map::from_iter([ - ("flow_id".to_owned(), json!("flow-9")), - ("langgraph_node".to_owned(), json!("host-node")), - ])) - }; - let trace = LangfuseTraceConfig { - trace_id: Some("flow:9:run-1".to_owned()), - session_id: Some("flow:9".to_owned()), - metadata: json!({"source": "flows"}), - ..Default::default() - }; - let payload = build_graph_langfuse_batch(trace, &observations, Some(&metadata)).unwrap(); - let batch = payload["batch"].as_array().unwrap(); - assert_eq!(batch[0]["type"], "trace-create"); - assert_eq!(batch[0]["body"]["id"], "flow:9:run-1"); - assert_eq!(batch[0]["body"]["sessionId"], "flow:9"); - assert_eq!(batch[0]["body"]["metadata"]["source"], "flows"); - let node = find(batch, "span-create", "flow:9:run-1:node:a:1").unwrap(); - assert_eq!(node["body"]["metadata"]["flow_id"], "flow-9"); - assert_eq!(node["body"]["metadata"]["langgraph_node"], "host-node"); - assert_eq!(node["body"]["metadata"]["langgraph_step"], 1); - assert_eq!(node["body"]["parentObservationId"], "flow:9:run-1:step:1"); -} - -#[test] -fn transport_free_batch_rejects_an_empty_journal() { - assert!(matches!( - build_graph_langfuse_batch(LangfuseTraceConfig::default(), &[], None), - Err(TinyAgentsError::Validation(message)) if message == "at least one observation is required" - )); -} - -#[test] -fn transport_free_batch_uses_root_identity_without_an_explicit_trace() { - let payload = - build_graph_langfuse_batch(LangfuseTraceConfig::default(), &sample_run(), None).unwrap(); - let batch = payload["batch"].as_array().unwrap(); - assert_eq!(batch[0]["body"]["id"], "root-1"); - assert!(find(batch, "span-create", "root-1:run:root-1").is_some()); - assert!(find(batch, "span-create", "root-1:node:b:1").is_some()); -} - -/// The string value of OTLP span attribute `key`, if present. -fn otlp_attr<'a>(span: &'a Value, key: &str) -> Option<&'a str> { - span["attributes"] - .as_array()? - .iter() - .find(|item| item["key"] == key)?["value"]["stringValue"] - .as_str() -} - -#[test] -fn otlp_export_is_a_v4_trace_with_trace_attributes_on_every_span() { - let trace = LangfuseTraceConfig { - name: Some("flow.run:demo".into()), - session_id: Some("flow:demo".into()), - release: Some("1.0.0".into()), - tags: vec!["run:flow".into()], - ..Default::default() - }; - let requests = exporter() - .build_otlp_requests(trace, &sample_run(), "openhuman") - .unwrap(); - assert_eq!(requests.len(), 1); - let scope = &requests[0]["resourceSpans"][0]["scopeSpans"][0]; - assert_eq!(scope["scope"]["name"], "openhuman.graph"); - let spans = scope["spans"].as_array().unwrap(); - for span in spans { - assert_eq!( - otlp_attr(span, "langfuse.trace.name"), - Some("flow.run:demo") - ); - assert_eq!(otlp_attr(span, "langfuse.session.id"), Some("flow:demo")); - assert_eq!(otlp_attr(span, "langfuse.release"), Some("1.0.0")); - assert_eq!( - otlp_attr(span, "langfuse.trace.metadata.graph_id"), - Some("demo-graph") - ); - assert!(otlp_attr(span, "langfuse.trace.metadata.health").is_some()); - } - - // One root: the structural graph-run span, everything else nests below it. - let roots: Vec<_> = spans - .iter() - .filter(|s| s.get("parentSpanId").is_none()) - .collect(); - assert_eq!(roots.len(), 1); - assert_eq!(roots[0]["name"], "demo-graph"); - - let by_name = |name: &str| spans.iter().find(|s| s["name"] == name).unwrap(); - let (step, node_b) = (by_name("step 1"), by_name("b")); - assert_eq!(step["parentSpanId"], roots[0]["spanId"]); - assert_eq!(node_b["parentSpanId"], step["spanId"]); - assert_eq!( - otlp_attr(node_b, "langfuse.observation.level"), - Some("ERROR") - ); - assert_eq!( - otlp_attr(node_b, "langfuse.observation.metadata.langgraph_node"), - Some("b") - ); - let failed = by_name("run.failed"); - assert_eq!( - otlp_attr(failed, "langfuse.observation.type"), - Some("event") - ); -} diff --git a/crates/tinyagents-graph/src/observability/mod.rs b/crates/tinyagents-graph/src/observability/mod.rs deleted file mode 100644 index 4ad4c10dc..000000000 --- a/crates/tinyagents-graph/src/observability/mod.rs +++ /dev/null @@ -1,628 +0,0 @@ -//! Durable observability for the graph runtime — journals, status stores, and -//! the journaling event sink. -//! -//! The live [`crate::stream`] layer emits transient [`GraphEvent`]s into -//! an in-process [`GraphEventSink`]. This module makes that history **durable -//! and correlatable** so a UI, supervisor, or test can reconstruct a recursive -//! graph run tree after the fact: -//! -//! - [`GraphObservation`] — a durable envelope pairing an event with its run -//! lineage (`run_id` / `parent_run_id` / `root_run_id`), `graph_id`, -//! `checkpoint_id`, subgraph `namespace`, `step`, `offset`, and timestamp. -//! - [`GraphEventJournal`] — an append-only, offset-addressable journal of -//! observations, with an [`InMemoryGraphEventJournal`] and a store-backed -//! [`StoreGraphEventJournal`] (stream key = run id). -//! - [`GraphStatusStore`] — a compact "what is running now?" surface over -//! [`crate::GraphRunStatus`], with an [`InMemoryGraphStatusStore`]. -//! - [`JournalGraphSink`] — a [`GraphEventSink`] that wraps each emitted event -//! into a [`GraphObservation`] and appends it to a journal, optionally also -//! forwarding to a live `inner` sink. -//! - [`GraphLatencyMetrics`] / [`GraphHealthSummary`] — rollups derived from a -//! run's observations: per-step/per-node timings, and per-node -//! success/failure counts (node-level **tool health** telemetry). -//! - [`GraphLangfuseExporter`] — exports a run's observations to Langfuse, -//! turning supersteps and nodes into timed spans (failures promoted to -//! `ERROR`) and attaching the health summary to the trace. It shares the -//! harness [`LangfuseClient`](tinyagents_harness::observability::LangfuseClient) -//! transport and defaults its `traceId` to the run's `root_run_id`, so a -//! graph run and the agent/tool runs its nodes spawn land under one trace. -//! -//! [`CompiledGraph`](crate::CompiledGraph) can be wired to write to a -//! status store and a journal through its builder-style -//! [`with_status_store`](crate::CompiledGraph::with_status_store) and -//! [`with_event_journal`](crate::CompiledGraph::with_event_journal) -//! methods; both are opt-in and default off so existing runs are unchanged. -//! -//! The journaling sink bridges the synchronous [`GraphEventSink::emit`] hook to -//! the async journal API through a background drain (an -//! [`AppendWorker`](tinyagents_harness::observability)): `emit` never blocks the -//! executor on I/O and persistence is best-effort (a full bounded queue drops -//! rather than stalls; backend errors are reported, not propagated). The -//! executor calls [`GraphEventSink::flush`] after the terminal run event so a -//! caller reading the journal right after the run returns sees a complete log. - -mod langfuse; -mod tracing_sink; -mod types; - -pub use langfuse::{GraphLangfuseExporter, SpanMetadataFn, build_graph_langfuse_batch}; -pub use tracing_sink::GraphTracingSink; -pub use types::*; - -use std::collections::{HashMap, VecDeque}; -use std::sync::Arc; -use std::sync::atomic::Ordering; -use std::time::SystemTime; - -use async_trait::async_trait; - -use crate::status::GraphRunStatus; -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}; -use tinyagents_harness::store::AppendStore; - -// --------------------------------------------------------------------------- -// GraphLatencyMetrics -// --------------------------------------------------------------------------- - -impl GraphLatencyMetrics { - /// Builds latency rollups from durable observations for one graph run. - /// - /// Incomplete steps or node executions are ignored because there is no - /// terminal timestamp to measure against. Duplicate node activations for the - /// same `(node, step)` are paired in FIFO order. - pub fn from_observations(observations: &[GraphObservation]) -> Self { - let mut metrics = Self::default(); - let mut run_start: Option = None; - let mut step_starts: HashMap = HashMap::new(); - let mut node_starts: HashMap<(NodeId, usize), VecDeque> = HashMap::new(); - - for obs in observations { - match &obs.event { - GraphEvent::RunStarted { .. } if run_start.is_none() => { - run_start = Some(obs.ts_ms); - } - GraphEvent::RunStarted { .. } => {} - GraphEvent::RunCompleted { .. } | GraphEvent::RunFailed { .. } => { - if metrics.run_elapsed_ms.is_none() - && let Some(start) = run_start - { - metrics.run_elapsed_ms = Some(obs.ts_ms.saturating_sub(start)); - } - } - GraphEvent::StepStarted { step, .. } => { - step_starts.insert(*step, obs.ts_ms); - } - GraphEvent::StepCompleted { step } => { - if let Some(start) = step_starts.remove(step) { - metrics.record_step(GraphStepLatency { - step: *step, - elapsed_ms: obs.ts_ms.saturating_sub(start), - }); - } - } - GraphEvent::NodeStarted { node, step } => { - node_starts - .entry((node.clone(), *step)) - .or_default() - .push_back(obs.ts_ms); - } - GraphEvent::NodeCompleted { node, step } => { - if let Some(start) = pop_node_start(&mut node_starts, node, *step) { - metrics.record_node(GraphNodeLatency { - node: node.clone(), - step: *step, - elapsed_ms: obs.ts_ms.saturating_sub(start), - failed: false, - }); - } - } - GraphEvent::NodeFailed { node, step, .. } => { - if let Some(start) = pop_node_start(&mut node_starts, node, *step) { - metrics.record_node(GraphNodeLatency { - node: node.clone(), - step: *step, - elapsed_ms: obs.ts_ms.saturating_sub(start), - failed: true, - }); - } - } - _ => {} - } - } - - metrics - } - - /// Builds a run-level latency summary from a compact status snapshot. - /// - /// Status snapshots do not contain per-step or per-node timings, but they - /// do carry started/updated/ended timestamps for end-to-end elapsed time. - pub fn from_status(status: &GraphRunStatus) -> Self { - let end = status.ended_at.unwrap_or(status.updated_at); - Self { - run_elapsed_ms: duration_ms(status.started_at, end), - ..Self::default() - } - } - - /// Average completed-step latency. - pub fn average_step_ms(&self) -> Option { - average(self.total_step_ms, self.steps.len()) - } - - /// Average completed-node latency. - pub fn average_node_ms(&self) -> Option { - average(self.total_node_ms, self.nodes.len()) - } - - fn record_step(&mut self, latency: GraphStepLatency) { - self.total_step_ms = self.total_step_ms.saturating_add(latency.elapsed_ms); - self.max_step_ms = self.max_step_ms.max(latency.elapsed_ms); - self.steps.push(latency); - } - - fn record_node(&mut self, latency: GraphNodeLatency) { - self.total_node_ms = self.total_node_ms.saturating_add(latency.elapsed_ms); - self.max_node_ms = self.max_node_ms.max(latency.elapsed_ms); - self.nodes.push(latency); - } -} - -// --------------------------------------------------------------------------- -// GraphHealthSummary -// --------------------------------------------------------------------------- - -impl GraphHealthSummary { - /// Builds a node/tool health rollup from durable observations for one run. - /// - /// Counts every `node.started`, `node.completed`, and `node.failed` - /// observation per node, plus whether the run itself failed. Per-node - /// entries are sorted by node id so the summary is deterministic. - pub fn from_observations(observations: &[GraphObservation]) -> Self { - let mut per_node: HashMap = HashMap::new(); - let mut summary = Self::default(); - - for obs in observations { - match &obs.event { - GraphEvent::NodeStarted { node, .. } => { - entry_for(&mut per_node, node).started += 1; - summary.total_started += 1; - } - GraphEvent::NodeCompleted { node, .. } => { - entry_for(&mut per_node, node).completed += 1; - summary.total_completed += 1; - } - GraphEvent::NodeFailed { node, .. } => { - entry_for(&mut per_node, node).failed += 1; - summary.total_failed += 1; - } - GraphEvent::RunFailed { .. } => { - summary.run_failed = true; - } - _ => {} - } - } - - summary.nodes = per_node.into_values().collect(); - summary - .nodes - .sort_by(|a, b| a.node.as_str().cmp(b.node.as_str())); - summary - } -} - -/// Returns the mutable health entry for `node`, inserting a zeroed one keyed by -/// the node id if absent. -fn entry_for<'a>( - per_node: &'a mut HashMap, - node: &NodeId, -) -> &'a mut GraphNodeHealth { - per_node - .entry(node.clone()) - .or_insert_with(|| GraphNodeHealth { - node: node.clone(), - started: 0, - completed: 0, - failed: 0, - }) -} - -// --------------------------------------------------------------------------- -// InMemoryGraphEventJournal -// --------------------------------------------------------------------------- - -impl InMemoryGraphEventJournal { - /// Creates a new, empty in-memory journal. - pub fn new() -> Self { - Self::default() - } - - /// Returns the number of observations stored for `run_id`. - pub fn len(&self, run_id: &str) -> usize { - self.runs - .lock() - .expect("InMemoryGraphEventJournal lock poisoned") - .get(run_id) - .map(|v| v.len()) - .unwrap_or(0) - } - - /// Returns `true` when no observations are stored for `run_id`. - pub fn is_empty(&self, run_id: &str) -> bool { - self.len(run_id) == 0 - } -} - -#[async_trait] -impl GraphEventJournal for InMemoryGraphEventJournal { - async fn append(&self, obs: GraphObservation) -> Result { - let mut runs = self - .runs - .lock() - .map_err(|e| poisoned("InMemoryGraphEventJournal", e))?; - let entries = runs.entry(obs.run_id.as_str().to_string()).or_default(); - let offset = entries.len() as u64; - entries.push(obs); - Ok(offset) - } - - async fn read_from(&self, run_id: &str, offset: u64) -> Result> { - let runs = self - .runs - .lock() - .map_err(|e| poisoned("InMemoryGraphEventJournal", e))?; - let Some(entries) = runs.get(run_id) else { - return Ok(Vec::new()); - }; - Ok(entries.iter().skip(offset as usize).cloned().collect()) - } -} - -// --------------------------------------------------------------------------- -// StoreGraphEventJournal -// --------------------------------------------------------------------------- - -impl StoreGraphEventJournal { - /// Wraps `store` as a graph event journal whose stream key is the run id. - pub fn new(store: A) -> Self { - Self { store } - } - - /// Returns a reference to the backing store. - pub fn store(&self) -> &A { - &self.store - } -} - -#[async_trait] -impl GraphEventJournal for StoreGraphEventJournal { - async fn append(&self, obs: GraphObservation) -> Result { - let stream = obs.run_id.as_str().to_string(); - let value = serde_json::to_value(&obs)?; - self.store.append(&stream, value).await - } - - async fn read_from(&self, run_id: &str, offset: u64) -> Result> { - let raw = self.store.read_from(run_id, offset).await?; - let mut out = Vec::with_capacity(raw.len()); - for (_offset, value) in raw { - out.push(serde_json::from_value(value)?); - } - Ok(out) - } -} - -// --------------------------------------------------------------------------- -// InMemoryGraphStatusStore -// --------------------------------------------------------------------------- - -impl InMemoryGraphStatusStore { - /// Creates a new, empty, **unbounded** in-memory status store. - pub fn new() -> Self { - Self::default() - } - - /// Caps the store at `max` distinct runs. - /// - /// Once a `put_status` for a *new* run pushes the store past the cap, the - /// oldest **terminal** run (completed / failed / cancelled) is evicted - /// first; when every retained run is still live, the oldest run overall is - /// evicted. Overwriting an already-recorded run never triggers eviction. A - /// `max` of `0` retains nothing. The default (via [`Self::new`] / - /// [`Default`]) is unbounded. - pub fn with_max_runs(mut self, max: usize) -> Self { - self.max_runs = Some(max); - self - } - - /// Returns the number of distinct runs with a recorded status. - pub fn len(&self) -> usize { - self.state - .lock() - .expect("InMemoryGraphStatusStore lock poisoned") - .statuses - .len() - } - - /// Returns `true` when no statuses have been recorded. - pub fn is_empty(&self) -> bool { - self.len() == 0 - } -} - -impl StatusStoreState { - /// Removes `run_id` from the thread index, dropping the thread's bucket - /// when it becomes empty. - fn unindex_thread(&mut self, thread_id: &str, run_id: &str) { - if let Some(runs) = self.by_thread.get_mut(thread_id) { - runs.retain(|id| id != run_id); - if runs.is_empty() { - self.by_thread.remove(thread_id); - } - } - } - - /// Evicts one run: the oldest terminal run if any, otherwise the oldest - /// run overall. No-op when the store is empty. - fn evict_one(&mut self) { - let idx = self - .order - .iter() - .position(|id| self.statuses.get(id).is_none_or(|s| s.is_terminal())) - .unwrap_or(0); - let Some(run_id) = self.order.remove(idx) else { - return; - }; - if let Some(evicted) = self.statuses.remove(&run_id) - && let Some(thread_id) = evicted.thread_id.as_ref() - { - self.unindex_thread(thread_id.as_str(), &run_id); - } - } -} - -#[async_trait] -impl GraphStatusStore for InMemoryGraphStatusStore { - async fn put_status(&self, status: GraphRunStatus) -> Result<()> { - let mut state = self - .state - .lock() - .map_err(|e| poisoned("InMemoryGraphStatusStore", e))?; - let run_id = status.run_id.as_str().to_string(); - let thread_id = status.thread_id.as_ref().map(|t| t.as_str().to_string()); - - match state.statuses.insert(run_id.clone(), status) { - Some(previous) => { - // Overwrite: keep the thread index coherent if the thread - // changed (rare, but cheap to handle). - let previous_thread = previous.thread_id.as_ref().map(|t| t.as_str().to_string()); - if previous_thread != thread_id { - if let Some(old) = previous_thread { - state.unindex_thread(&old, &run_id); - } - if let Some(new) = thread_id { - state.by_thread.entry(new).or_default().push(run_id); - } - } - } - None => { - // First status for this run: index it and enforce the cap. - if let Some(thread) = thread_id { - state - .by_thread - .entry(thread) - .or_default() - .push(run_id.clone()); - } - state.order.push_back(run_id); - if let Some(max) = self.max_runs { - while state.statuses.len() > max { - state.evict_one(); - } - } - } - } - Ok(()) - } - - async fn get_status(&self, run_id: &str) -> Result> { - let state = self - .state - .lock() - .map_err(|e| poisoned("InMemoryGraphStatusStore", e))?; - Ok(state.statuses.get(run_id).cloned()) - } - - async fn list_by_thread(&self, thread_id: &str) -> Result> { - let state = self - .state - .lock() - .map_err(|e| poisoned("InMemoryGraphStatusStore", e))?; - Ok(state - .by_thread - .get(thread_id) - .map(|runs| { - runs.iter() - .filter_map(|id| state.statuses.get(id).cloned()) - .collect() - }) - .unwrap_or_default()) - } -} - -// --------------------------------------------------------------------------- -// JournalGraphSink -// --------------------------------------------------------------------------- - -impl JournalGraphSink { - /// Builds a journal sink for `run_id` of `graph_id`. `root_run_id` defaults - /// to `run_id` (a top-level run) and the namespace is empty; use the - /// builder methods to set a parent, root, thread, namespace, or downstream - /// sink. - pub fn new(journal: Arc, run_id: RunId, graph_id: GraphId) -> Self { - let worker = Arc::new(AppendWorker::spawn( - "graph-journal-sink", - DEFAULT_DRAIN_CAPACITY, - move |obs: GraphObservation| { - let journal = Arc::clone(&journal); - async move { journal.append(obs).await.map(|_| ()) } - }, - )); - Self { - worker, - inner: None, - root_run_id: run_id.clone(), - run_id, - parent_run_id: None, - thread_id: None, - graph_id, - namespace: Vec::new(), - offset: Arc::new(std::sync::atomic::AtomicU64::new(0)), - step: Arc::new(std::sync::atomic::AtomicU64::new(0)), - } - } - - /// Sets the parent and root run ids stamped onto every observation. - pub fn with_lineage(mut self, parent_run_id: Option, root_run_id: RunId) -> Self { - self.parent_run_id = parent_run_id; - self.root_run_id = root_run_id; - self - } - - /// Sets the thread id stamped onto every observation. - pub fn with_thread(mut self, thread_id: Option) -> Self { - self.thread_id = thread_id; - self - } - - /// Sets the checkpoint namespace stamped onto every observation. - /// - /// A subgraph sink is given the child namespace here so its observations - /// carry the nested path. - pub fn with_namespace(mut self, namespace: Vec) -> Self { - self.namespace = namespace; - self - } - - /// Forwards every event to `inner` in addition to journaling it, so one - /// configured sink can both persist and broadcast. - pub fn with_inner(mut self, inner: Arc) -> Self { - self.inner = Some(inner); - self - } - - /// Builds the durable observation for `event`, advancing the offset and the - /// latest-step trackers. - fn observe(&self, event: &GraphEvent) -> GraphObservation { - let offset = self.offset.fetch_add(1, Ordering::Relaxed); - // Track the latest superstep so events without a step (route, run - // lifecycle) still carry the step they happened during. - let step = match event.step() { - Some(step) => { - self.step.store(step as u64, Ordering::Relaxed); - step - } - None => self.step.load(Ordering::Relaxed) as usize, - }; - GraphObservation { - event_id: EventId::new(format!("{}-{offset}", self.run_id.as_str())), - run_id: self.run_id.clone(), - root_run_id: self.root_run_id.clone(), - parent_run_id: self.parent_run_id.clone(), - thread_id: self.thread_id.clone(), - graph_id: self.graph_id.clone(), - checkpoint_id: checkpoint_of(event), - namespace: self.namespace.clone(), - step, - offset, - ts_ms: now_ms(), - event: event.clone(), - } - } -} - -impl GraphEventSink for JournalGraphSink { - 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(envelope); - } - } - - fn flush(&self) { - self.worker.flush(); - if let Some(inner) = &self.inner { - inner.flush(); - } - } -} - -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()), - _ => None, - } -} - -// --------------------------------------------------------------------------- -// Shared helpers -// --------------------------------------------------------------------------- - -fn pop_node_start( - starts: &mut HashMap<(NodeId, usize), VecDeque>, - node: &NodeId, - step: usize, -) -> Option { - let key = (node.clone(), step); - let queue = starts.get_mut(&key)?; - let start = queue.pop_front(); - if queue.is_empty() { - starts.remove(&key); - } - start -} - -fn average(total: u64, count: usize) -> Option { - (count > 0).then_some(total / count as u64) -} - -fn duration_ms(start: SystemTime, end: SystemTime) -> Option { - end.duration_since(start) - .ok() - .map(|duration| duration.as_millis() as u64) -} - -/// Builds a uniform poisoned-lock validation error for the in-memory backends. -fn poisoned( - what: &str, - err: E, -) -> tinyagents_harness::error::TinyAgentsError { - tinyagents_harness::error::TinyAgentsError::Validation(format!("{what} lock poisoned: {err}")) -} - -#[cfg(test)] -#[path = "mod_tests.rs"] -mod test; diff --git a/crates/tinyagents-graph/src/observability/mod_tests.rs b/crates/tinyagents-graph/src/observability/mod_tests.rs deleted file mode 100644 index 7220e19c4..000000000 --- a/crates/tinyagents-graph/src/observability/mod_tests.rs +++ /dev/null @@ -1,602 +0,0 @@ -//! Unit tests for the graph durable observability layer: journaling sinks, -//! offset-addressable replay, status-store lifecycle, store-backed journals, and -//! namespaced subgraph observations. - -use super::*; -use crate::builder::{GraphBuilder, NodeContext}; -use crate::command::NodeResult; -use crate::compiled::CompiledGraph; -use crate::stream::{CollectingSink, GraphEvent, GraphEventSink}; -use std::sync::Arc; -use tinyagents_harness::error::TinyAgentsError; -use tinyagents_harness::ids::{ExecutionStatus, GraphId, NodeId, RunId}; -use tinyagents_harness::store::InMemoryAppendStore; - -/// A two-node line graph over `i32` with overwrite semantics: `a -> b`. -fn line_graph() -> CompiledGraph { - GraphBuilder::::overwrite() - .add_node("a", |_s, _c: NodeContext| async move { - Ok(NodeResult::Update(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() -} - -/// A single-node graph whose node always errors. -fn failing_graph() -> CompiledGraph { - GraphBuilder::::overwrite() - .add_node("boom", |_s, _c: NodeContext| async move { - Err::, _>(TinyAgentsError::Validation("boom".to_string())) - }) - .set_entry("boom") - .set_finish("boom") - .compile() - .unwrap() -} - -fn graph_obs(run: &str, offset: u64, step: usize, event: GraphEvent) -> GraphObservation { - GraphObservation { - event_id: tinyagents_harness::ids::EventId::new(format!("g-evt-{offset}")), - run_id: RunId::new(run), - root_run_id: RunId::new(run), - parent_run_id: None, - thread_id: None, - graph_id: GraphId::new("graph-latency"), - checkpoint_id: None, - namespace: Vec::new(), - step, - offset, - ts_ms: 1_000 + offset, - event, - } -} - -#[test] -fn graph_latency_metrics_include_run_steps_and_nodes() { - let run_id = RunId::new("run-latency"); - let node_a = NodeId::new("a"); - let node_b = NodeId::new("b"); - let observations = vec![ - graph_obs( - run_id.as_str(), - 0, - 0, - GraphEvent::RunStarted { - run_id: run_id.clone(), - }, - ), - graph_obs( - run_id.as_str(), - 10, - 1, - GraphEvent::StepStarted { - step: 1, - active: vec![node_a.clone(), node_b.clone()], - }, - ), - graph_obs( - run_id.as_str(), - 20, - 1, - GraphEvent::NodeStarted { - node: node_a.clone(), - step: 1, - }, - ), - graph_obs( - run_id.as_str(), - 45, - 1, - GraphEvent::NodeCompleted { - node: node_a.clone(), - step: 1, - }, - ), - graph_obs( - run_id.as_str(), - 50, - 1, - GraphEvent::NodeStarted { - node: node_b.clone(), - step: 1, - }, - ), - graph_obs( - run_id.as_str(), - 80, - 1, - GraphEvent::NodeFailed { - node: node_b.clone(), - step: 1, - error: "boom".to_string(), - }, - ), - graph_obs( - run_id.as_str(), - 90, - 1, - GraphEvent::StepCompleted { step: 1 }, - ), - graph_obs( - run_id.as_str(), - 100, - 1, - GraphEvent::RunFailed { - run_id: run_id.clone(), - error: "boom".to_string(), - }, - ), - ]; - - let metrics = GraphLatencyMetrics::from_observations(&observations); - assert_eq!(metrics.run_elapsed_ms, Some(100)); - assert_eq!(metrics.steps.len(), 1); - assert_eq!(metrics.steps[0].step, 1); - assert_eq!(metrics.steps[0].elapsed_ms, 80); - assert_eq!(metrics.total_step_ms, 80); - assert_eq!(metrics.average_step_ms(), Some(80)); - - assert_eq!(metrics.nodes.len(), 2); - assert_eq!(metrics.nodes[0].node, node_a); - assert_eq!(metrics.nodes[0].elapsed_ms, 25); - assert!(!metrics.nodes[0].failed); - assert_eq!(metrics.nodes[1].node, node_b); - assert_eq!(metrics.nodes[1].elapsed_ms, 30); - assert!(metrics.nodes[1].failed); - assert_eq!(metrics.total_node_ms, 55); - assert_eq!(metrics.max_node_ms, 30); - assert_eq!(metrics.average_node_ms(), Some(27)); -} - -#[test] -fn graph_health_summary_counts_node_outcomes() { - let run = "run-health"; - let a = NodeId::new("a"); - let b = NodeId::new("b"); - let observations = vec![ - graph_obs( - run, - 0, - 0, - GraphEvent::RunStarted { - run_id: RunId::new(run), - }, - ), - graph_obs( - run, - 1, - 1, - GraphEvent::NodeStarted { - node: a.clone(), - step: 1, - }, - ), - graph_obs( - run, - 2, - 1, - GraphEvent::NodeCompleted { - node: a.clone(), - step: 1, - }, - ), - graph_obs( - run, - 3, - 1, - GraphEvent::NodeStarted { - node: b.clone(), - step: 1, - }, - ), - graph_obs( - run, - 4, - 1, - GraphEvent::NodeFailed { - node: b.clone(), - step: 1, - error: "boom".to_string(), - }, - ), - graph_obs( - run, - 5, - 1, - GraphEvent::RunFailed { - run_id: RunId::new(run), - error: "boom".to_string(), - }, - ), - ]; - - let health = GraphHealthSummary::from_observations(&observations); - assert_eq!(health.total_started, 2); - assert_eq!(health.total_completed, 1); - assert_eq!(health.total_failed, 1); - assert!(health.run_failed); - assert!(!health.is_healthy()); - assert_eq!(health.failure_rate(), 0.5); - - // Sorted by node id, with per-node health surfaced. - assert_eq!(health.nodes.len(), 2); - assert_eq!(health.nodes[0].node, a); - assert!(health.nodes[0].is_healthy()); - assert_eq!(health.nodes[1].node, b); - assert!(!health.nodes[1].is_healthy()); - assert_eq!(health.nodes[1].failure_rate(), 1.0); - - let unhealthy: Vec<_> = health.unhealthy_nodes().map(|n| n.node.clone()).collect(); - assert_eq!(unhealthy, vec![b]); -} - -#[test] -fn healthy_run_reports_no_failures() { - let run = "run-ok"; - let a = NodeId::new("a"); - let observations = vec![ - graph_obs( - run, - 0, - 1, - GraphEvent::NodeStarted { - node: a.clone(), - step: 1, - }, - ), - graph_obs( - run, - 1, - 1, - GraphEvent::NodeCompleted { - node: a.clone(), - step: 1, - }, - ), - graph_obs( - run, - 2, - 1, - GraphEvent::RunCompleted { - run_id: RunId::new(run), - steps: 1, - }, - ), - ]; - - let health = GraphHealthSummary::from_observations(&observations); - assert!(health.is_healthy()); - assert_eq!(health.failure_rate(), 0.0); - assert_eq!(health.unhealthy_nodes().count(), 0); -} - -#[tokio::test] -async fn run_with_journal_records_replayable_observations() { - let journal = Arc::new(InMemoryGraphEventJournal::new()); - let graph = line_graph().with_event_journal(journal.clone()); - - let run = graph.run(0).await.unwrap(); - let run_id = run.status.run_id.as_str().to_string(); - - let all = journal.read_from(&run_id, 0).await.unwrap(); - assert!( - all.len() >= 3, - "expected several observations, got {}", - all.len() - ); - - // Offsets are dense and monotonic from 0, and lineage/graph are stamped. - for (i, obs) in all.iter().enumerate() { - assert_eq!(obs.offset, i as u64); - assert_eq!(obs.run_id.as_str(), run_id); - assert_eq!(obs.root_run_id.as_str(), run_id); - assert_eq!(obs.graph_id, *graph.graph_id()); - } - - // The run lifecycle bookends the stream. - assert!(matches!( - all.first().unwrap().event, - GraphEvent::RunStarted { .. } - )); - assert!( - all.iter() - .any(|o| matches!(o.event, GraphEvent::RunCompleted { steps: 2, .. })) - ); - - // Replay from a mid-stream offset returns exactly the tail. - let tail = journal.read_from(&run_id, 2).await.unwrap(); - assert_eq!(tail.len(), all.len() - 2); - assert_eq!(tail.first().unwrap().offset, 2); - - // Reading an unknown run is empty, not an error. - assert!(journal.read_from("nope", 0).await.unwrap().is_empty()); -} - -#[tokio::test] -async fn status_store_reflects_run_lifecycle() { - let store = Arc::new(InMemoryGraphStatusStore::new()); - let graph = line_graph().with_status_store(store.clone()); - - let run = graph.run_with_thread("t-1", 0).await.unwrap(); - let run_id = run.status.run_id.as_str().to_string(); - - let status = store.get_status(&run_id).await.unwrap().unwrap(); - assert_eq!(status.status, ExecutionStatus::Completed); - assert_eq!(status.current_step, 2); - assert!(status.ended_at.is_some()); - assert!(status.error.is_none()); - - // Indexed by thread for supervisor queries. - let by_thread = store.list_by_thread("t-1").await.unwrap(); - assert_eq!(by_thread.len(), 1); - assert_eq!(by_thread[0].run_id.as_str(), run_id); -} - -#[tokio::test] -async fn status_store_records_failed_runs() { - let store = Arc::new(InMemoryGraphStatusStore::new()); - let graph = failing_graph().with_status_store(store.clone()); - - let err = graph.run_with_thread("t-fail", 0).await.unwrap_err(); - assert!(matches!(err, TinyAgentsError::Validation(_))); - - let failed = store.list_by_thread("t-fail").await.unwrap(); - assert_eq!(failed.len(), 1); - assert_eq!(failed[0].status, ExecutionStatus::Failed); - assert!(failed[0].error.as_deref().unwrap().contains("boom")); - assert!(failed[0].ended_at.is_some()); -} - -#[tokio::test] -async fn failed_run_journals_run_failed_event() { - let journal = Arc::new(InMemoryGraphEventJournal::new()); - let store = Arc::new(InMemoryGraphStatusStore::new()); - let graph = failing_graph() - .with_event_journal(journal.clone()) - .with_status_store(store.clone()); - - let _ = graph.run_with_thread("t-boom", 0).await.unwrap_err(); - - // Recover the executor-assigned run id via the thread-indexed status store. - let status = store.list_by_thread("t-boom").await.unwrap(); - let run_id = status[0].run_id.as_str().to_string(); - - let obs = journal.read_from(&run_id, 0).await.unwrap(); - assert!(matches!( - obs.first().unwrap().event, - GraphEvent::RunStarted { .. } - )); - assert!(obs.iter().any( - |o| matches!(&o.event, GraphEvent::RunFailed { error, .. } if error.contains("boom")) - )); -} - -#[tokio::test] -async fn namespaced_subgraph_observations_carry_child_namespace() { - let journal = Arc::new(InMemoryGraphEventJournal::new()); - let child_ns = vec!["parent".to_string(), "child".to_string()]; - let graph = line_graph() - .with_namespace(child_ns.clone()) - .with_event_journal(journal.clone()); - - let run = graph.run(0).await.unwrap(); - let obs = journal - .read_from(run.status.run_id.as_str(), 0) - .await - .unwrap(); - - assert!(!obs.is_empty()); - assert!( - obs.iter().all(|o| o.namespace == child_ns), - "every observation should carry the child namespace" - ); -} - -#[tokio::test] -async fn journal_sink_used_directly_forwards_to_inner() { - // Drive the public JournalGraphSink as the live event sink: it both journals - // observations under its configured run id and forwards to an inner sink. - let journal = Arc::new(InMemoryGraphEventJournal::new()); - let collector = Arc::new(CollectingSink::new()); - let sink = JournalGraphSink::new(journal.clone(), RunId::new("fixed-run"), GraphId::new("g")) - .with_inner(collector.clone()); - - sink.emit(GraphEventEnvelope::for_test(GraphEvent::StepStarted { - step: 1, - active: Vec::new(), - })); - sink.emit(GraphEventEnvelope::for_test(GraphEvent::RouteSelected { - node: "a".into(), - target: "b".into(), - })); - sink.emit(GraphEventEnvelope::for_test(GraphEvent::StepCompleted { - step: 1, - })); - - // Forwarded to the live sink. - assert_eq!(collector.len(), 3); - - // Persistence is asynchronous; block until the durable log catches up. - sink.flush(); - - // Journaled with dense offsets; the step is carried forward onto the - // step-less RouteSelected event. - let obs = journal.read_from("fixed-run", 0).await.unwrap(); - assert_eq!(obs.len(), 3); - assert_eq!(obs[0].step, 1); - assert_eq!(obs[1].step, 1); // RouteSelected inherits the last seen step - assert_eq!(obs[2].offset, 2); -} - -#[tokio::test] -async fn store_backed_journal_round_trips() { - let store = StoreGraphEventJournal::new(InMemoryAppendStore::new()); - let obs = GraphObservation { - event_id: tinyagents_harness::ids::EventId::new("e0"), - run_id: RunId::new("r0"), - root_run_id: RunId::new("r0"), - parent_run_id: None, - thread_id: None, - graph_id: GraphId::new("g0"), - checkpoint_id: None, - namespace: vec!["ns".to_string()], - step: 1, - offset: 0, - ts_ms: 42, - event: GraphEvent::RunStarted { - run_id: RunId::new("r0"), - }, - }; - let off = store.append(obs.clone()).await.unwrap(); - assert_eq!(off, 0); - - let read = store.read_from("r0", 0).await.unwrap(); - assert_eq!(read.len(), 1); - assert_eq!(read[0], obs); -} - -#[test] -fn graph_event_kind_and_step_are_stable() { - assert_eq!( - GraphEvent::RunStarted { - run_id: RunId::new("r") - } - .kind(), - "run.started" - ); - assert_eq!( - GraphEvent::SubgraphStarted { - node: "n".into(), - namespace: vec![] - } - .kind(), - "subgraph.started" - ); - let forked = GraphEvent::ContextForked { - node: "n".into(), - fork: 0, - step: 3, - }; - assert_eq!(forked.kind(), "context.forked"); - assert_eq!(forked.step(), Some(3)); - assert_eq!(GraphEvent::RecursionDepthChanged { depth: 2 }.step(), None); -} - -// ── InMemoryGraphStatusStore indexing and retention ─────────────────────────── - -/// Builds a status for `run_id` on `thread_id` in the given lifecycle state. -fn status_on_thread(run_id: &str, thread_id: &str, status: ExecutionStatus) -> GraphRunStatus { - let mut s = GraphRunStatus::new(RunId::new(run_id), GraphId::new("g"), status); - s.thread_id = Some(tinyagents_harness::ids::ThreadId::new(thread_id)); - s -} - -#[tokio::test] -async fn status_store_thread_index_tracks_overwrites() { - let store = InMemoryGraphStatusStore::new(); - store - .put_status(status_on_thread("r-1", "t-a", ExecutionStatus::Running)) - .await - .unwrap(); - store - .put_status(status_on_thread("r-2", "t-a", ExecutionStatus::Running)) - .await - .unwrap(); - store - .put_status(status_on_thread("r-3", "t-b", ExecutionStatus::Running)) - .await - .unwrap(); - - assert_eq!(store.list_by_thread("t-a").await.unwrap().len(), 2); - assert_eq!(store.list_by_thread("t-b").await.unwrap().len(), 1); - assert!(store.list_by_thread("t-missing").await.unwrap().is_empty()); - - // Overwriting a run on the same thread does not duplicate the index entry. - store - .put_status(status_on_thread("r-1", "t-a", ExecutionStatus::Completed)) - .await - .unwrap(); - assert_eq!(store.list_by_thread("t-a").await.unwrap().len(), 2); - - // Re-homing a run to a different thread moves it in the index. - store - .put_status(status_on_thread("r-2", "t-b", ExecutionStatus::Running)) - .await - .unwrap(); - assert_eq!(store.list_by_thread("t-a").await.unwrap().len(), 1); - assert_eq!(store.list_by_thread("t-b").await.unwrap().len(), 2); - assert_eq!(store.len(), 3); -} - -#[tokio::test] -async fn status_store_cap_evicts_oldest_terminal_first() { - let store = InMemoryGraphStatusStore::new().with_max_runs(2); - store - .put_status(status_on_thread("r-live", "t", ExecutionStatus::Running)) - .await - .unwrap(); - store - .put_status(status_on_thread("r-done", "t", ExecutionStatus::Completed)) - .await - .unwrap(); - // Third run exceeds the cap: the terminal `r-done` goes first even though - // `r-live` is older. - store - .put_status(status_on_thread("r-new", "t", ExecutionStatus::Running)) - .await - .unwrap(); - - assert_eq!(store.len(), 2); - assert!(store.get_status("r-done").await.unwrap().is_none()); - assert!(store.get_status("r-live").await.unwrap().is_some()); - assert!(store.get_status("r-new").await.unwrap().is_some()); - - // The thread index no longer serves the evicted run. - let by_thread = store.list_by_thread("t").await.unwrap(); - assert_eq!(by_thread.len(), 2); - assert!(by_thread.iter().all(|s| s.run_id.as_str() != "r-done")); -} - -#[tokio::test] -async fn status_store_cap_falls_back_to_oldest_live_run() { - let store = InMemoryGraphStatusStore::new().with_max_runs(2); - for id in ["r-1", "r-2", "r-3"] { - store - .put_status(status_on_thread(id, "t", ExecutionStatus::Running)) - .await - .unwrap(); - } - // No terminal run to prefer, so the oldest live run is evicted. - assert_eq!(store.len(), 2); - assert!(store.get_status("r-1").await.unwrap().is_none()); - assert!(store.get_status("r-2").await.unwrap().is_some()); - assert!(store.get_status("r-3").await.unwrap().is_some()); - assert_eq!(store.list_by_thread("t").await.unwrap().len(), 2); -} - -#[tokio::test] -async fn status_store_overwrite_never_evicts() { - let store = InMemoryGraphStatusStore::new().with_max_runs(2); - store - .put_status(status_on_thread("r-1", "t", ExecutionStatus::Running)) - .await - .unwrap(); - store - .put_status(status_on_thread("r-2", "t", ExecutionStatus::Running)) - .await - .unwrap(); - // Updating an existing run at capacity is not an insertion. - store - .put_status(status_on_thread("r-1", "t", ExecutionStatus::Completed)) - .await - .unwrap(); - assert_eq!(store.len(), 2); - assert!(store.get_status("r-1").await.unwrap().is_some()); - assert!(store.get_status("r-2").await.unwrap().is_some()); -} diff --git a/crates/tinyagents-graph/src/observability/tracing_sink.rs b/crates/tinyagents-graph/src/observability/tracing_sink.rs deleted file mode 100644 index 6ff56e1a4..000000000 --- a/crates/tinyagents-graph/src/observability/tracing_sink.rs +++ /dev/null @@ -1,70 +0,0 @@ -//! [`GraphTracingSink`] — mirrors the graph executor's lifecycle -//! stream onto `tracing` diagnostics. - -use std::sync::Arc; -use std::sync::atomic::Ordering; - -use crate::stream::{GraphEvent, GraphEventEnvelope, GraphEventSink}; - -/// A [`GraphEventSink`] that mirrors the graph executor's lifecycle -/// stream onto `tracing` diagnostics — an observability journal for -/// graph runs. Node/step/run/route transitions land as -/// grep-friendly `[graph]` lines tagged with `label`; the running event count is -/// exposed for tests. Hosts share one per graph family (fan-out, -/// delegation, …). -pub struct GraphTracingSink { - label: String, - count: Arc, -} - -impl GraphTracingSink { - /// Build a sink tagging its lines with `label` (e.g. `"delegation:graph"`). - /// Accepts both string literals and runtime-built labels. - pub fn new(label: impl Into) -> Self { - Self { - label: label.into(), - count: Arc::new(std::sync::atomic::AtomicUsize::new(0)), - } - } - - /// Shared counter of events observed. - pub fn counter(&self) -> Arc { - self.count.clone() - } -} - -impl GraphEventSink for GraphTracingSink { - fn emit(&self, envelope: GraphEventEnvelope) { - let event = envelope.event; - self.count.fetch_add(1, Ordering::Relaxed); - let label = self.label.as_str(); - match &event { - GraphEvent::RunStarted { run_id } => { - tracing::debug!(label, ?run_id, "[graph] run started") - } - GraphEvent::RunCompleted { steps, .. } => { - tracing::debug!(label, steps, "[graph] run completed") - } - GraphEvent::RunFailed { error, .. } => { - tracing::warn!(label, %error, "[graph] run failed") - } - GraphEvent::NodeStarted { node, step } => { - tracing::debug!(label, ?node, step, "[graph] node started") - } - GraphEvent::NodeCompleted { node, step } => { - tracing::debug!(label, ?node, step, "[graph] node completed") - } - GraphEvent::NodeFailed { node, error, .. } => { - tracing::warn!(label, ?node, %error, "[graph] node failed") - } - GraphEvent::RouteSelected { node, target } => { - tracing::trace!(label, ?node, ?target, "[graph] route selected") - } - _ => tracing::trace!(label, "[graph] event"), - } - } -} - -#[cfg(test)] -#[path = "tracing_sink_tests.rs"] -mod tests; diff --git a/crates/tinyagents-graph/src/observability/tracing_sink_tests.rs b/crates/tinyagents-graph/src/observability/tracing_sink_tests.rs deleted file mode 100644 index 43bc21710..000000000 --- a/crates/tinyagents-graph/src/observability/tracing_sink_tests.rs +++ /dev/null @@ -1,26 +0,0 @@ -use super::*; -use std::sync::atomic::Ordering; -use tinyagents_harness::ids::{NodeId, RunId}; - -#[test] -fn every_emitted_event_is_counted_including_unmatched_variants() { - let sink = GraphTracingSink::new("test:graph"); - let count = sink.counter(); - sink.emit(GraphEventEnvelope::for_test(GraphEvent::RunStarted { - run_id: RunId::new("r"), - })); - sink.emit(GraphEventEnvelope::for_test(GraphEvent::NodeFailed { - node: NodeId::new("n"), - step: 1, - error: "boom".into(), - })); - sink.emit(GraphEventEnvelope::for_test(GraphEvent::StepCompleted { - step: 1, - })); - assert_eq!(count.load(Ordering::Relaxed), 3); -} - -#[test] -fn runtime_built_labels_are_accepted() { - let _ = GraphTracingSink::new(format!("team:{}", 7)); -} diff --git a/crates/tinyagents-graph/src/observability/types.rs b/crates/tinyagents-graph/src/observability/types.rs deleted file mode 100644 index 0e75c2b6a..000000000 --- a/crates/tinyagents-graph/src/observability/types.rs +++ /dev/null @@ -1,397 +0,0 @@ -//! Type definitions for the graph durable observability layer. -//! -//! These types add **durability** to the live [`crate::stream`] event -//! vocabulary: an envelope ([`GraphObservation`]) that carries a run's lineage, -//! checkpoint coordinates, namespace, and a timestamp so a single -//! [`GraphEvent`] can be journaled, replayed, and correlated across a recursive -//! graph run tree; pluggable journal and status traits; and a -//! [`JournalGraphSink`] that wraps emitted events into observations. -//! -//! All public items here are re-exported through [`super`]. Trait -//! implementations, sink logic, and tests live in the sibling `mod.rs` and -//! `test.rs` files. - -use std::collections::HashMap; -use std::sync::atomic::AtomicU64; -use std::sync::{Arc, Mutex}; - -use async_trait::async_trait; -use serde::{Deserialize, Serialize}; - -use crate::status::GraphRunStatus; -use crate::stream::{GraphEvent, GraphEventSink}; -use tinyagents_harness::error::Result; -use tinyagents_harness::ids::{CheckpointId, EventId, GraphId, NodeId, RunId, ThreadId}; -use tinyagents_harness::observability::AppendWorker; -use tinyagents_harness::store::AppendStore; - -// --------------------------------------------------------------------------- -// GraphObservation -// --------------------------------------------------------------------------- - -/// A durable observability envelope around a [`GraphEvent`]. -/// -/// Where a raw [`GraphEvent`] is the transient, in-process signal the executor -/// emits at each boundary, a `GraphObservation` adds everything a durable -/// journal or external trace needs to correlate the event without an in-memory -/// broadcast: the run's `run_id`, its `parent_run_id` / `root_run_id` lineage, -/// the owning `graph_id`, the latest `checkpoint_id`, the subgraph `namespace`, -/// the superstep `step`, the stream `offset`, and a wall-clock `ts_ms`. -#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] -pub struct GraphObservation { - /// Stable, unique identifier for this observation. - pub event_id: EventId, - - /// The run that emitted the event. - pub run_id: RunId, - - /// Root ancestor run, equal to `run_id` for top-level runs. - pub root_run_id: RunId, - - /// Parent run id when this run was spawned by another run (a subgraph or - /// sub-agent). `None` for top-level runs. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub parent_run_id: Option, - - /// The thread the run executes under, when checkpointing is enabled. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub thread_id: Option, - - /// The graph this run belongs to. - pub graph_id: GraphId, - - /// The latest checkpoint id at the time the event was emitted, if any. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub checkpoint_id: Option, - - /// The checkpoint namespace (the child path for nested subgraph runs). - #[serde(default, skip_serializing_if = "Vec::is_empty")] - pub namespace: Vec, - - /// The superstep number the event is associated with (`0` before the run - /// has entered a superstep). - pub step: usize, - - /// Monotonic position of the observation within its run's stream. - pub offset: u64, - - /// Wall-clock time the observation was created, in Unix-epoch milliseconds. - pub ts_ms: u64, - - /// The typed event payload. - pub event: GraphEvent, -} - -// --------------------------------------------------------------------------- -// Graph latency metrics -// --------------------------------------------------------------------------- - -/// Latency for one graph superstep. -/// -/// Derived by correlating `step.started` and `step.completed` observations with -/// the same `step` value. -#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] -pub struct GraphStepLatency { - /// 1-based graph superstep. - pub step: usize, - - /// Wall-clock elapsed time between step start and completion. - pub elapsed_ms: u64, -} - -/// Latency for one node handler execution. -/// -/// Derived by correlating `node.started` with either `node.completed` or -/// `node.failed` for the same node and step. The `failed` flag distinguishes -/// successful and failed terminal observations while keeping both in the same -/// latency rollup. -#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] -pub struct GraphNodeLatency { - /// Node whose handler ran. - pub node: NodeId, - - /// 1-based graph superstep. - pub step: usize, - - /// Wall-clock elapsed time between node start and terminal observation. - pub elapsed_ms: u64, - - /// True when the node ended with `node.failed`. - pub failed: bool, -} - -/// Summarized latency metrics for a single graph run. -#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)] -pub struct GraphLatencyMetrics { - /// End-to-end graph run latency, when both run start and terminal events - /// exist. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub run_elapsed_ms: Option, - - /// Per-superstep latencies in completion order. - #[serde(default)] - pub steps: Vec, - - /// Per-node handler latencies in terminal-event order. - #[serde(default)] - pub nodes: Vec, - - /// Sum of completed superstep latency. - pub total_step_ms: u64, - - /// Slowest completed superstep. - pub max_step_ms: u64, - - /// Sum of completed node-handler latency. - pub total_node_ms: u64, - - /// Slowest completed node handler. - pub max_node_ms: u64, -} - -// --------------------------------------------------------------------------- -// Graph health telemetry -// --------------------------------------------------------------------------- - -/// Health rollup for a single graph node across one run. -/// -/// A graph node is the graph's unit of work — frequently a delegated agent or -/// tool call (see [`crate::SubAgentNode`]) — so per-node success/failure -/// counts double as **tool health** telemetry for the graph. Counts are derived -/// from durable [`GraphObservation`]s by correlating `node.started` with -/// `node.completed` / `node.failed`. -#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] -pub struct GraphNodeHealth { - /// The node these counts describe. - pub node: NodeId, - - /// Number of `node.started` observations seen. - pub started: u64, - - /// Number of `node.completed` observations seen. - pub completed: u64, - - /// Number of `node.failed` observations seen. - pub failed: u64, -} - -impl GraphNodeHealth { - /// Terminal attempts (`completed + failed`) — activations with a known - /// outcome. `started` may exceed this while a node is still running. - pub fn attempts(&self) -> u64 { - self.completed.saturating_add(self.failed) - } - - /// Fraction of terminal attempts that failed, in `0.0..=1.0`. Returns `0.0` - /// when the node has no terminal attempts yet. - pub fn failure_rate(&self) -> f64 { - let attempts = self.attempts(); - if attempts == 0 { - 0.0 - } else { - self.failed as f64 / attempts as f64 - } - } - - /// True when the node has never failed. - pub fn is_healthy(&self) -> bool { - self.failed == 0 - } -} - -/// Aggregate node/tool health for a single graph run. -/// -/// Built from the same durable observation stream that feeds -/// [`GraphLatencyMetrics`], this is the compact "is anything unhealthy?" surface -/// a supervisor or dashboard reads, and the telemetry the Langfuse exporter -/// attaches to a trace. Per-node entries are sorted by node id for stable -/// output. -#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)] -pub struct GraphHealthSummary { - /// Per-node health entries, sorted by node id. - #[serde(default)] - pub nodes: Vec, - - /// Total `node.started` observations across all nodes. - pub total_started: u64, - - /// Total `node.completed` observations across all nodes. - pub total_completed: u64, - - /// Total `node.failed` observations across all nodes. - pub total_failed: u64, - - /// True when the run itself emitted `run.failed`. - pub run_failed: bool, -} - -impl GraphHealthSummary { - /// Terminal node attempts (`total_completed + total_failed`). - pub fn total_attempts(&self) -> u64 { - self.total_completed.saturating_add(self.total_failed) - } - - /// Fraction of terminal node attempts that failed, in `0.0..=1.0`. - pub fn failure_rate(&self) -> f64 { - let attempts = self.total_attempts(); - if attempts == 0 { - 0.0 - } else { - self.total_failed as f64 / attempts as f64 - } - } - - /// True when no node failed and the run did not fail. - pub fn is_healthy(&self) -> bool { - self.total_failed == 0 && !self.run_failed - } - - /// The nodes that failed at least once, in node-id order. - pub fn unhealthy_nodes(&self) -> impl Iterator { - self.nodes.iter().filter(|n| !n.is_healthy()) - } -} - -// --------------------------------------------------------------------------- -// GraphEventJournal -// --------------------------------------------------------------------------- - -/// A durable, append-only journal of [`GraphObservation`]s keyed by run id. -/// -/// Journals decouple durable replay from live broadcast: a UI or supervisor can -/// attach after a graph run has started and reconstruct history by reading from -/// a known offset rather than relying on having subscribed to an in-memory -/// [`GraphEventSink`]. -#[async_trait] -pub trait GraphEventJournal: Send + Sync { - /// Appends `obs` to the journal and returns the offset it was stored at - /// within its run's stream. - async fn append(&self, obs: GraphObservation) -> Result; - - /// Returns every observation for `run_id` whose stream offset is `>= - /// offset`, in offset order. Reading from `0` replays the whole run; - /// reading an unknown run returns an empty `Vec`. - async fn read_from(&self, run_id: &str, offset: u64) -> Result>; -} - -/// In-memory [`GraphEventJournal`] backed by a per-run `Vec`. -/// -/// Cheaply clonable through an inner [`Arc`]; clones share the same streams. -/// There is no durability — entries are lost when the last clone drops. -#[derive(Clone, Debug, Default)] -pub struct InMemoryGraphEventJournal { - /// `run_id → ordered observations`. - pub(crate) runs: Arc>>>, -} - -/// [`GraphEventJournal`] backed by any [`AppendStore`]. -/// -/// Each run's observations are appended to the store under a stream named by -/// the run id, so `read_from` resumes from a durable offset. Pair with -/// [`tinyagents_harness::store::DriverAppendStore`] for a durable journal or -/// [`tinyagents_harness::store::InMemoryAppendStore`] for deterministic tests. -#[derive(Clone, Debug)] -pub struct StoreGraphEventJournal { - /// The backing append store; stream key is the run id. - pub(crate) store: A, -} - -// --------------------------------------------------------------------------- -// GraphStatusStore -// --------------------------------------------------------------------------- - -/// A readable status surface for graph runs. -/// -/// Status records are overwritten by `run_id` ("what is running now?") in -/// contrast to the append-only journal ("what happened?"). A -/// [`GraphRunStatus`] is compact: ids, step, active nodes, pending interrupts, -/// and timestamps — never full graph state, which belongs to the checkpointer. -#[async_trait] -pub trait GraphStatusStore: Send + Sync { - /// Inserts or overwrites the status for its `run_id`. - async fn put_status(&self, status: GraphRunStatus) -> Result<()>; - - /// Returns the latest status for `run_id`, or `None` if unknown. - async fn get_status(&self, run_id: &str) -> Result>; - - /// Returns all known statuses whose `thread_id` matches `thread_id`, in - /// unspecified order. - async fn list_by_thread(&self, thread_id: &str) -> Result>; -} - -/// In-memory [`GraphStatusStore`] backed by a `run_id → status` map plus a -/// `thread_id → run ids` secondary index, so -/// [`list_by_thread`](GraphStatusStore::list_by_thread) is proportional to -/// that thread's runs rather than every recorded run. -/// -/// # Retention -/// Unbounded by default. Long-lived processes should opt into a run cap with -/// [`InMemoryGraphStatusStore::with_max_runs`]: once the number of distinct -/// runs exceeds the cap, the store evicts the oldest **terminal** run first -/// (completed / failed / cancelled), falling back to the oldest run overall -/// when every run is still live. -/// -/// Cheaply clonable through an inner [`Arc`]; clones share the same state -/// *and* the same cap. -#[derive(Clone, Debug, Default)] -pub struct InMemoryGraphStatusStore { - /// Shared map + indexes, kept coherent under one lock. - pub(crate) state: Arc>, - /// Maximum retained runs; `None` means unbounded (default). - pub(crate) max_runs: Option, -} - -/// Internal state of [`InMemoryGraphStatusStore`]: the primary map plus the -/// indexes that keep `list_by_thread` fast and eviction ordered. -#[derive(Debug, Default)] -pub(crate) struct StatusStoreState { - /// `run_id → latest status`. - pub(crate) statuses: HashMap, - /// `thread_id → run ids` secondary index for `list_by_thread`. - pub(crate) by_thread: HashMap>, - /// Run ids in first-insertion order, driving oldest-first eviction. - pub(crate) order: std::collections::VecDeque, -} - -// --------------------------------------------------------------------------- -// JournalGraphSink -// --------------------------------------------------------------------------- - -/// A [`GraphEventSink`] that wraps each emitted [`GraphEvent`] into a durable -/// [`GraphObservation`] and appends it to a [`GraphEventJournal`]. -/// -/// The sink is configured with the emitting run's lineage and checkpoint -/// coordinates; each received event is stamped with a monotonically increasing -/// `offset`, the latest observed `step`, and the configured `namespace`. The -/// observation is then handed to a background [`AppendWorker`] that persists it -/// off the executor thread (best-effort — a full bounded queue drops rather -/// than stalling the run, and append errors are reported, not propagated). -/// [`crate::stream::GraphEventSink::flush`] blocks until the durable log -/// has caught up; the executor calls it after the terminal run event. -/// -/// An optional `inner` sink lets the journal sink also forward each event to a -/// live transport (for example a [`crate::stream::CollectingSink`]) so a -/// single configured sink can both persist and broadcast. -#[derive(Clone)] -pub struct JournalGraphSink { - /// Background drain that persists observations without blocking the run. - pub(crate) worker: Arc>, - /// Optional downstream sink that also receives every event. - pub(crate) inner: Option>, - /// The run that owns events delivered to this sink. - pub(crate) run_id: RunId, - /// Root run id stamped onto every observation. - pub(crate) root_run_id: RunId, - /// Parent run id stamped onto every observation. - pub(crate) parent_run_id: Option, - /// Thread id stamped onto every observation. - pub(crate) thread_id: Option, - /// The graph this run belongs to. - pub(crate) graph_id: GraphId, - /// The checkpoint namespace stamped onto every observation. - pub(crate) namespace: Vec, - /// Monotonic per-run observation offset counter. - pub(crate) offset: Arc, - /// The latest observed superstep, used to stamp events with no step. - pub(crate) step: Arc, -} diff --git a/crates/tinyagents-graph/src/orchestration.rs b/crates/tinyagents-graph/src/orchestration.rs deleted file mode 100644 index 3bcfd54d6..000000000 --- a/crates/tinyagents-graph/src/orchestration.rs +++ /dev/null @@ -1,13 +0,0 @@ -//! Compatibility shim: the detached-task machinery moved to `tinyagents-tasks`. -//! -//! The detached-task runtime types still provided by `tinyagents-tasks` are -//! re-exported here so their existing `tinyagents_graph::orchestration::*` -//! paths keep compiling. Removed orchestration tool adapters and recovery -//! helpers are not part of this compatibility surface. Prefer -//! `tinyagents_tasks` in new code. - -pub use tinyagents_tasks::*; - -#[cfg(test)] -#[path = "orchestration_tests.rs"] -mod tests; diff --git a/crates/tinyagents-graph/src/orchestration_tests.rs b/crates/tinyagents-graph/src/orchestration_tests.rs deleted file mode 100644 index 92c739d19..000000000 --- a/crates/tinyagents-graph/src/orchestration_tests.rs +++ /dev/null @@ -1,31 +0,0 @@ -//! The old graph paths must keep resolving after the move to `tinyagents-tasks`. - -use crate::orchestration::{ - DetachedTaskRegistry, DriverTaskStore, InMemoryTaskStore, OrchestrationTaskStatus, - SteeringRegistry, TaskStore, TaskStoreRegistry, reconcile_orphaned_tasks, -}; - -#[test] -fn old_graph_paths_are_the_tasks_crate_types() { - fn same(_: std::marker::PhantomData, _: std::marker::PhantomData) {} - same( - std::marker::PhantomData::, - std::marker::PhantomData::, - ); - same( - std::marker::PhantomData::, - std::marker::PhantomData::, - ); - same( - std::marker::PhantomData::, - std::marker::PhantomData::, - ); - let _ = std::marker::PhantomData::<( - InMemoryTaskStore, - DriverTaskStore, - TaskStoreRegistry, - DetachedTaskRegistry<(), ()>, - Box, - )>; - let _ = reconcile_orphaned_tasks; -} diff --git a/crates/tinyagents-graph/src/parallel/README.md b/crates/tinyagents-graph/src/parallel/README.md deleted file mode 100644 index b8b722ae5..000000000 --- a/crates/tinyagents-graph/src/parallel/README.md +++ /dev/null @@ -1,65 +0,0 @@ -# graph::parallel - -Ordered, bounded-concurrency parallel map/reduce with a configurable failure -policy, plus (in `claims/`) the shared-workspace claim arbitration that -decides when fanned-out workers are safe to run concurrently at all. - -Graph `Send` fanout is the low-level primitive for parallel supersteps, but -callers frequently want a *reusable* "run these N items concurrently and -reduce the results" helper independent of the graph executor: deterministic -input-order results, a concurrency cap, per-item success/failure isolation, -and a policy for what to do when some items fail (fail-fast, collect-all, -quorum, best-effort). That is what [`map_reduce`] provides. - -## Public surface - -- `map_reduce(items, options, f) -> Result>` — runs `f` - over `items` with bounded concurrency (`ParallelOptions::max_concurrency`, - `0` = unbounded), collecting per-item outcomes in **input order** regardless - of completion order, and applying `options.failure_policy`. -- `ParallelOptions` — builder-style options: `with_max_concurrency`, - `with_failure_policy`, `with_item_timeout`, `with_total_timeout`, - `with_cancellation`. -- `FailurePolicy` — `FailFast` (return the first input-order error, cancel the - rest), `CollectAll` (default; always `Ok`, per-item outcomes), `Quorum(n)` - (error unless at least `n` items succeed), `BestEffort` (always `Ok`, keep - only successes). -- `ItemOutcome` / `ParallelOutcome` — per-item and aggregate results, - with `success_count`/`failure_count`/`successes`/`into_successes` helpers. -- `claims` (re-exported): `WorkspaceClaim`, `DispatchPlan`, `DispatchMode`, - `ClaimConflict`, `ClaimPathError`, `parse_relative_claim_paths`, - `paths_overlap`, `writes_shared_workspace`, - `plan_shared_workspace_dispatch` — see `claims/README.md`. - -## Files - -| File | Role | -| --- | --- | -| `types.rs` | `FailurePolicy`, `ParallelOptions`, `ItemOutcome`, `ParallelOutcome`. | -| `mod.rs` | `map_reduce`: the bounded-concurrency, input-order-preserving driver. | -| `mod_tests.rs` | Unit tests (ordering under out-of-order completion, each failure policy, timeouts, cancellation). | -| `claims/` | Shared-workspace claim arbitration (own `README.md`). | - -## Operational constraints - -- `map_reduce` re-orders by input index after collection; it does **not** - stream partial results — callers needing incremental output should not use - this helper. -- `FailFast` cancels remaining in-flight work by dropping the underlying - stream once every item with a smaller input index has resolved, not merely - on the first error observed — see the `fail_fast_error` bookkeeping in - `mod.rs` for why the first-completed error is not necessarily the one - returned. -- `item_timeout` and `total_timeout` are independent: an item exceeding - `item_timeout` becomes a per-item failure handled by `FailurePolicy`; the - whole call exceeding `total_timeout` aborts unconditionally with - `TinyAgentsError::Timeout`, even under `CollectAll`/`BestEffort`. - -## Relation to neighbours - -`map_reduce` answers "run these N items concurrently"; `claims::plan_shared_workspace_dispatch` -(see `claims/README.md`) answers the orthogonal question of whether running -fanned-out workers concurrently is *safe* against a shared filesystem root. -Neither module depends on the other at the type level — a caller composes -them: plan dispatch first, then drive the resulting parallel/serial groups -through `map_reduce` (or an equivalent scheduler). diff --git a/crates/tinyagents-graph/src/parallel/claims/README.md b/crates/tinyagents-graph/src/parallel/claims/README.md deleted file mode 100644 index 207bd4fc9..000000000 --- a/crates/tinyagents-graph/src/parallel/claims/README.md +++ /dev/null @@ -1,86 +0,0 @@ -# graph::parallel::claims - -Shared-workspace claim arbitration for parallel agent fan-out — a decision -module with **no I/O and no scheduling**: given a set of `WorkspaceClaim`s it -returns a `DispatchPlan` saying which workers may run concurrently, which must -be serialized, and which cannot be scheduled at all. - -`parallel::map_reduce` answers "run these N items with bounded concurrency." -It cannot answer whether running them concurrently is *safe*. When fanned-out -workers share one filesystem root, two of them mutating the same file at the -same time corrupts it silently — there is no error, just a plausible-looking -result built on a torn tree. This module owns that decision, and only that -decision. - -## The rule - -A worker is safe to run in parallel when it either has its own root -(`isolated`) or never writes (`!writes`). A worker that writes a *shared* root -must declare the paths it owns; it is then serialized, and its claim is -checked against every claim already granted. Claims are granted in **input -order, first-writer-wins**, which makes the plan a pure function of the input -regardless of which worker happens to finish first — the same request always -yields the same rejection. - -## What stays with the caller - -The crate reports a `ClaimConflict` as data; it never decides what a host does -about one. Whether an unbounded write is a hard rejection or a warning, and -what sentence the user reads, are product decisions. Likewise the syntax a -claim arrives in: `parse_relative_claim_paths` takes the *body* of a claim -list, not whatever prefix or key a host wraps it in. - -## Public surface - -- `WorkspaceClaim` — one worker's declared relationship to the shared - workspace (`isolated`, `read_only`, `writing` constructors). -- `plan_shared_workspace_dispatch(claims) -> DispatchPlan` — the planner. -- `DispatchPlan` — `modes` (index-aligned `Option`, `None` = - rejected) plus `conflicts` (index-keyed rejections); `has_serial_work`, - `parallel_indices`, `serial_indices`. -- `DispatchMode` — `Parallel` / `Serial`. -- `ClaimConflict` — `UnboundedWrite` (writer with no declared paths) / - `Overlap` (claim collides with an earlier one); `worker_id()` accessor. -- `ClaimPathError` — `Absolute` / `Escaping`, returned by - `parse_relative_claim_paths`. -- `parse_relative_claim_paths(spec) -> Result, ClaimPathError>` — - parses a comma/newline-separated, optionally bulleted claim list into safe, - sorted, deduplicated relative paths. -- `paths_overlap(left, right) -> bool` — component-wise (not textual) overlap - check: `src/a` and `src/ab` are distinct, `src/a` and `src/a/inner.rs` - collide. -- `writes_shared_workspace(effects: &ToolSideEffects) -> bool` — derives a - "does this tool write?" claim input from tool side-effect metadata. - -## Files - -| File | Role | -| --- | --- | -| `types.rs` | `WorkspaceClaim`, `ClaimPathError`, `ClaimConflict`, `DispatchMode`, `DispatchPlan`. | -| `mod.rs` | `parse_relative_claim_paths`, `paths_overlap`, `writes_shared_workspace`, `plan_shared_workspace_dispatch`. | -| `mod_tests.rs` | Unit tests (path parsing/safety, overlap semantics, planner ordering and conflict reporting). | - -## Operational constraints - -- What a host *does* with a `ClaimConflict` (hard rejection vs. warning, and - what a user reads) is a product decision this crate deliberately does not - make — it only reports the conflict as data. -- `parse_relative_claim_paths` rejects absolute paths and anything containing - `..` or a platform path prefix, because either would let a claim reach - outside the shared root it is meant to partition. This is a security - boundary, not a convenience check — do not bypass it by constructing - `WorkspaceClaim::writing` paths from unvalidated input. -- Claim order matters: `plan_shared_workspace_dispatch` grants first-writer- - wins in the order `claims` is given, so callers that need deterministic, - reproducible plans must pass claims in a stable order. -- This module performs no I/O and spawns nothing; it is pure decision logic - the caller (e.g. `parallel::map_reduce`) uses to decide dispatch, not an - executor. - -## Relation to neighbours - -Sits alongside `parallel::map_reduce` (parent module) as the "is this safe?" -half of parallel fan-out; `map_reduce` handles "run it," this module handles -"may it run concurrently." Consumed by hosts that fan out agent/tool workers -over a shared filesystem root before dispatching them through `map_reduce` or -an equivalent scheduler. diff --git a/crates/tinyagents-graph/src/parallel/claims/mod.rs b/crates/tinyagents-graph/src/parallel/claims/mod.rs deleted file mode 100644 index 74e1c4484..000000000 --- a/crates/tinyagents-graph/src/parallel/claims/mod.rs +++ /dev/null @@ -1,167 +0,0 @@ -//! Shared-workspace claim arbitration for parallel agent fan-out. -//! -//! [`map_reduce`](super::map_reduce) answers "run these N items with bounded -//! concurrency". It cannot answer whether running them concurrently is *safe*. -//! When fanned-out workers share one filesystem root, two of them mutating the -//! same file at the same time corrupts it silently — there is no error, just a -//! plausible-looking result built on a torn tree. -//! -//! This module owns that decision, and only that decision. Given a set of -//! [`WorkspaceClaim`]s it returns a [`DispatchPlan`] saying which workers may -//! run concurrently, which must be serialized, and which cannot be scheduled at -//! all. It performs no I/O, spawns nothing, and never renders a message. -//! -//! # The rule -//! -//! A worker is safe to run in parallel when it either has its own root -//! (`isolated`) or never writes (`!writes`). A worker that writes a *shared* -//! root must declare the paths it owns; it is then serialized, and its claim is -//! checked against every claim already granted. -//! -//! Claims are granted in **input order, first-writer-wins**. That ordering is -//! load-bearing: it makes the plan a pure function of the input, independent of -//! how quickly any worker happens to finish, so the same request always yields -//! the same rejection. -//! -//! # What stays with the caller -//! -//! The crate reports a [`ClaimConflict`] as data; it never decides what a host -//! does about one. Whether an unbounded write is a hard rejection or a warning, -//! and what sentence the user reads, are product decisions. Likewise the syntax -//! a claim arrives in: [`parse_relative_claim_paths`] takes the *body* of a -//! claim list, not whatever prefix or key a host wraps it in. - -mod types; - -pub use types::{ClaimConflict, ClaimPathError, DispatchMode, DispatchPlan, WorkspaceClaim}; - -use std::path::{Component, Path, PathBuf}; - -use tinytools::ToolSideEffects; - -/// Parses a claim list into safe, sorted, deduplicated relative paths. -/// -/// Entries are separated by commas or newlines and may carry a leading `-` or -/// `*` bullet. Empty entries are skipped, so a trailing separator is harmless. -/// -/// # Errors -/// -/// Returns [`ClaimPathError::Absolute`] for an absolute path and -/// [`ClaimPathError::Escaping`] for one containing `..` or a platform prefix — -/// either would let a claim reach outside the shared root it is meant to -/// partition. -pub fn parse_relative_claim_paths(spec: &str) -> Result, ClaimPathError> { - let mut paths = Vec::new(); - for raw in spec.split([',', '\n']) { - let trimmed = raw.trim().trim_start_matches(['-', '*']).trim(); - if trimmed.is_empty() { - continue; - } - let path = PathBuf::from(trimmed); - if path.is_absolute() { - return Err(ClaimPathError::Absolute { - raw: trimmed.to_string(), - }); - } - if path - .components() - .any(|component| matches!(component, Component::ParentDir | Component::Prefix(_))) - { - return Err(ClaimPathError::Escaping { - raw: trimmed.to_string(), - }); - } - paths.push(path); - } - paths.sort(); - paths.dedup(); - Ok(paths) -} - -/// `true` when two claimed paths cannot be held by different workers at once. -/// -/// Comparison is **component-wise**, not textual: `src/a` and `src/ab` are -/// distinct, while `src/a` and `src/a/inner.rs` collide because one contains the -/// other. -pub fn paths_overlap(left: &Path, right: &Path) -> bool { - left == right || left.starts_with(right) || right.starts_with(left) -} - -/// `true` when a tool's declared effects can mutate a shared workspace. -/// -/// Reads the declaration rather than a name or a host permission enum, so any -/// host can derive a claim from the tool metadata it already publishes. -/// `read_only` wins over the mutating flags: a tool that declares both is taken -/// at its most conservative word. -pub fn writes_shared_workspace(effects: &ToolSideEffects) -> bool { - if effects.read_only { - return false; - } - effects.writes_files || effects.installs_dependencies || effects.destructive -} - -/// Decides which workers may run concurrently against one shared workspace. -/// -/// Claims are considered in input order; each writing worker's paths are checked -/// against those already granted, and on the first overlap that worker is -/// rejected (the earlier holder keeps its claim). Rejected workers get `None` in -/// [`DispatchPlan::modes`] and an entry in [`DispatchPlan::conflicts`]. -pub fn plan_shared_workspace_dispatch(claims: &[WorkspaceClaim]) -> DispatchPlan { - let mut plan = DispatchPlan { - modes: Vec::with_capacity(claims.len()), - conflicts: Vec::new(), - }; - // Granted (path, worker_id) pairs, accumulated in input order. - let mut granted: Vec<(PathBuf, String)> = Vec::new(); - - for (index, claim) in claims.iter().enumerate() { - if claim.isolated || !claim.writes { - plan.modes.push(Some(DispatchMode::Parallel)); - continue; - } - - if claim.paths.is_empty() { - plan.modes.push(None); - plan.conflicts.push(( - index, - ClaimConflict::UnboundedWrite { - worker_id: claim.worker_id.clone(), - }, - )); - continue; - } - - let overlap = claim.paths.iter().find_map(|path| { - granted - .iter() - .find(|(claimed, _)| paths_overlap(path, claimed)) - .map(|(claimed, holder)| (claimed.clone(), holder.clone())) - }); - - match overlap { - Some((path, other_worker_id)) => { - plan.modes.push(None); - plan.conflicts.push(( - index, - ClaimConflict::Overlap { - worker_id: claim.worker_id.clone(), - other_worker_id, - path, - }, - )); - } - None => { - for path in &claim.paths { - granted.push((path.clone(), claim.worker_id.clone())); - } - plan.modes.push(Some(DispatchMode::Serial)); - } - } - } - - plan -} - -#[cfg(test)] -#[path = "mod_tests.rs"] -mod test; diff --git a/crates/tinyagents-graph/src/parallel/claims/mod_tests.rs b/crates/tinyagents-graph/src/parallel/claims/mod_tests.rs deleted file mode 100644 index ff12004e7..000000000 --- a/crates/tinyagents-graph/src/parallel/claims/mod_tests.rs +++ /dev/null @@ -1,328 +0,0 @@ -//! Unit tests for shared-workspace claim arbitration. - -use super::*; - -fn p(s: &str) -> PathBuf { - PathBuf::from(s) -} - -// --- parse_relative_claim_paths ----------------------------------------- - -#[test] -fn parses_comma_and_newline_separated_entries() { - let parsed = parse_relative_claim_paths("src/a.rs, src/b.rs\nsrc/c.rs").expect("parses"); - assert_eq!(parsed, vec![p("src/a.rs"), p("src/b.rs"), p("src/c.rs")]); -} - -#[test] -fn strips_bullet_markers_and_surrounding_space() { - let parsed = parse_relative_claim_paths(" - src/a.rs \n * src/b.rs").expect("parses"); - assert_eq!(parsed, vec![p("src/a.rs"), p("src/b.rs")]); -} - -#[test] -fn sorts_and_deduplicates() { - let parsed = parse_relative_claim_paths("src/b.rs, src/a.rs, src/b.rs").expect("parses"); - assert_eq!(parsed, vec![p("src/a.rs"), p("src/b.rs")]); -} - -#[test] -fn skips_empty_entries_so_trailing_separators_are_harmless() { - let parsed = parse_relative_claim_paths("src/a.rs,,\n \n").expect("parses"); - assert_eq!(parsed, vec![p("src/a.rs")]); -} - -#[test] -fn empty_spec_yields_no_paths() { - assert_eq!( - parse_relative_claim_paths("").expect("parses"), - Vec::::new() - ); -} - -#[test] -fn rejects_absolute_paths() { - let err = parse_relative_claim_paths("/etc/passwd").expect_err("absolute is rejected"); - assert_eq!( - err, - ClaimPathError::Absolute { - raw: "/etc/passwd".to_string() - } - ); -} - -#[test] -fn rejects_parent_dir_escapes() { - let err = parse_relative_claim_paths("../outside.rs").expect_err("escape is rejected"); - assert_eq!( - err, - ClaimPathError::Escaping { - raw: "../outside.rs".to_string() - } - ); -} - -#[test] -fn rejects_escape_hidden_mid_path() { - let err = parse_relative_claim_paths("src/../../outside.rs").expect_err("escape is rejected"); - assert!(matches!(err, ClaimPathError::Escaping { .. })); -} - -// --- paths_overlap ------------------------------------------------------- - -#[test] -fn identical_paths_overlap() { - assert!(paths_overlap(&p("src/a.rs"), &p("src/a.rs"))); -} - -#[test] -fn a_directory_overlaps_a_file_beneath_it() { - assert!(paths_overlap(&p("src"), &p("src/a.rs"))); - assert!(paths_overlap(&p("src/a.rs"), &p("src"))); -} - -#[test] -fn sibling_paths_sharing_a_textual_prefix_do_not_overlap() { - // The classic bug: `starts_with` on strings would call these an overlap. - // `Path::starts_with` is component-wise, and this pins that. - assert!(!paths_overlap(&p("src/a"), &p("src/ab"))); - assert!(!paths_overlap(&p("src/ab"), &p("src/a"))); - assert!(!paths_overlap(&p("a.rs"), &p("ab.rs"))); -} - -#[test] -fn unrelated_paths_do_not_overlap() { - assert!(!paths_overlap(&p("src/a.rs"), &p("docs/b.md"))); -} - -// --- writes_shared_workspace -------------------------------------------- - -#[test] -fn read_only_never_writes_even_when_other_flags_are_set() { - let effects = ToolSideEffects { - read_only: true, - writes_files: true, - destructive: true, - ..Default::default() - }; - assert!(!writes_shared_workspace(&effects)); -} - -#[test] -fn file_writes_installs_and_destructive_all_count_as_writing() { - for effects in [ - ToolSideEffects { - writes_files: true, - ..Default::default() - }, - ToolSideEffects { - installs_dependencies: true, - ..Default::default() - }, - ToolSideEffects { - destructive: true, - ..Default::default() - }, - ] { - assert!( - writes_shared_workspace(&effects), - "{effects:?} should count as writing" - ); - } -} - -#[test] -fn network_and_payment_alone_do_not_touch_the_workspace() { - let effects = ToolSideEffects { - network: true, - payment: true, - external_service: true, - ..Default::default() - }; - assert!(!writes_shared_workspace(&effects)); -} - -#[test] -fn default_effects_do_not_write() { - assert!(!writes_shared_workspace(&ToolSideEffects::default())); -} - -// --- plan_shared_workspace_dispatch ------------------------------------- - -#[test] -fn isolated_and_read_only_workers_all_run_in_parallel() { - let plan = plan_shared_workspace_dispatch(&[ - WorkspaceClaim::isolated("w1"), - WorkspaceClaim::read_only("w2"), - ]); - assert_eq!( - plan.modes, - vec![Some(DispatchMode::Parallel), Some(DispatchMode::Parallel)] - ); - assert!(plan.conflicts.is_empty()); - assert!(!plan.has_serial_work()); - assert_eq!(plan.parallel_indices(), vec![0, 1]); - assert!(plan.serial_indices().is_empty()); -} - -#[test] -fn an_isolated_worker_that_writes_still_runs_in_parallel() { - // Isolation short-circuits the write check: its root is its own. - let mut claim = WorkspaceClaim::isolated("w1"); - claim.writes = true; - let plan = plan_shared_workspace_dispatch(&[claim]); - assert_eq!(plan.modes, vec![Some(DispatchMode::Parallel)]); - assert!(plan.conflicts.is_empty()); -} - -#[test] -fn a_shared_writer_with_a_claim_is_serialized() { - let plan = - plan_shared_workspace_dispatch(&[WorkspaceClaim::writing("w1", vec![p("src/a.rs")])]); - assert_eq!(plan.modes, vec![Some(DispatchMode::Serial)]); - assert!(plan.conflicts.is_empty()); - assert!(plan.has_serial_work()); - assert_eq!(plan.serial_indices(), vec![0]); -} - -#[test] -fn a_shared_writer_without_a_claim_is_rejected_as_unbounded() { - let plan = plan_shared_workspace_dispatch(&[WorkspaceClaim::writing("w1", vec![])]); - assert_eq!(plan.modes, vec![None]); - assert_eq!( - plan.conflicts, - vec![( - 0, - ClaimConflict::UnboundedWrite { - worker_id: "w1".to_string() - } - )] - ); -} - -#[test] -fn two_writers_with_disjoint_claims_are_both_serialized() { - let plan = plan_shared_workspace_dispatch(&[ - WorkspaceClaim::writing("w1", vec![p("src/a.rs")]), - WorkspaceClaim::writing("w2", vec![p("src/b.rs")]), - ]); - assert_eq!( - plan.modes, - vec![Some(DispatchMode::Serial), Some(DispatchMode::Serial)] - ); - assert!(plan.conflicts.is_empty()); - assert_eq!(plan.serial_indices(), vec![0, 1]); -} - -#[test] -fn on_overlap_the_later_worker_is_rejected_and_the_earlier_keeps_its_claim() { - let plan = plan_shared_workspace_dispatch(&[ - WorkspaceClaim::writing("w1", vec![p("src/a.rs")]), - WorkspaceClaim::writing("w2", vec![p("src/a.rs")]), - ]); - assert_eq!(plan.modes, vec![Some(DispatchMode::Serial), None]); - assert_eq!( - plan.conflicts, - vec![( - 1, - ClaimConflict::Overlap { - worker_id: "w2".to_string(), - other_worker_id: "w1".to_string(), - path: p("src/a.rs"), - } - )] - ); -} - -#[test] -fn overlap_is_detected_through_directory_containment() { - let plan = plan_shared_workspace_dispatch(&[ - WorkspaceClaim::writing("w1", vec![p("src")]), - WorkspaceClaim::writing("w2", vec![p("src/a.rs")]), - ]); - assert_eq!(plan.modes, vec![Some(DispatchMode::Serial), None]); - assert_eq!(plan.conflicts.len(), 1); -} - -#[test] -fn a_rejected_worker_does_not_reserve_its_paths_for_later_workers() { - // w2 is rejected against w1, so w2's *other* path stays free for w3. - let plan = plan_shared_workspace_dispatch(&[ - WorkspaceClaim::writing("w1", vec![p("src/a.rs")]), - WorkspaceClaim::writing("w2", vec![p("src/a.rs"), p("src/z.rs")]), - WorkspaceClaim::writing("w3", vec![p("src/z.rs")]), - ]); - assert_eq!( - plan.modes, - vec![Some(DispatchMode::Serial), None, Some(DispatchMode::Serial)] - ); - assert_eq!(plan.conflicts.len(), 1); - assert_eq!(plan.conflicts[0].0, 1); -} - -#[test] -fn rejection_follows_input_order_not_claim_content() { - // Swapping the inputs swaps who is rejected — the plan is a pure function - // of input order, which is what makes it reproducible across runs. - let forward = plan_shared_workspace_dispatch(&[ - WorkspaceClaim::writing("first", vec![p("shared.rs")]), - WorkspaceClaim::writing("second", vec![p("shared.rs")]), - ]); - let reversed = plan_shared_workspace_dispatch(&[ - WorkspaceClaim::writing("second", vec![p("shared.rs")]), - WorkspaceClaim::writing("first", vec![p("shared.rs")]), - ]); - assert_eq!(forward.conflicts[0].1.worker_id(), "second"); - assert_eq!(reversed.conflicts[0].1.worker_id(), "first"); -} - -#[test] -fn mixed_batches_keep_every_index_aligned_with_the_input() { - let plan = plan_shared_workspace_dispatch(&[ - WorkspaceClaim::isolated("isolated"), - WorkspaceClaim::read_only("reader"), - WorkspaceClaim::writing("writer", vec![p("src/a.rs")]), - WorkspaceClaim::writing("clasher", vec![p("src/a.rs")]), - WorkspaceClaim::writing("unbounded", vec![]), - ]); - assert_eq!( - plan.modes, - vec![ - Some(DispatchMode::Parallel), - Some(DispatchMode::Parallel), - Some(DispatchMode::Serial), - None, - None, - ] - ); - assert_eq!(plan.parallel_indices(), vec![0, 1]); - assert_eq!(plan.serial_indices(), vec![2]); - assert_eq!( - plan.conflicts - .iter() - .map(|(index, _)| *index) - .collect::>(), - vec![3, 4] - ); -} - -#[test] -fn an_empty_batch_plans_nothing() { - let plan = plan_shared_workspace_dispatch(&[]); - assert_eq!(plan, DispatchPlan::default()); - assert!(!plan.has_serial_work()); -} - -#[test] -fn claim_conflict_reports_the_worker_it_rejected() { - let unbounded = ClaimConflict::UnboundedWrite { - worker_id: "w1".to_string(), - }; - assert_eq!(unbounded.worker_id(), "w1"); - let overlap = ClaimConflict::Overlap { - worker_id: "w2".to_string(), - other_worker_id: "w1".to_string(), - path: p("src/a.rs"), - }; - assert_eq!(overlap.worker_id(), "w2"); -} diff --git a/crates/tinyagents-graph/src/parallel/claims/types.rs b/crates/tinyagents-graph/src/parallel/claims/types.rs deleted file mode 100644 index 34dbd7c1f..000000000 --- a/crates/tinyagents-graph/src/parallel/claims/types.rs +++ /dev/null @@ -1,159 +0,0 @@ -//! Type definitions for shared-workspace claim arbitration. - -use std::path::PathBuf; - -/// One worker's declared relationship to a fan-out's shared workspace. -/// -/// A worker is safe to run concurrently with its siblings when it either has a -/// workspace of its own ([`isolated`](Self::isolated)) or never mutates the one -/// it shares ([`writes`](Self::writes) is `false`). Anything else needs an -/// explicit, disjoint set of owned [`paths`](Self::paths) before it can be -/// scheduled at all. -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct WorkspaceClaim { - /// Caller-chosen identifier, echoed back in conflicts so the caller can - /// phrase its own diagnostics. - pub worker_id: String, - /// The worker runs in its own isolated root and can never collide. - pub isolated: bool, - /// The worker can mutate files in whatever root it runs in. - pub writes: bool, - /// Relative paths this worker declares exclusive ownership of. Empty means - /// "unbounded" — a writing, non-isolated worker with no claim yields - /// [`ClaimConflict::UnboundedWrite`]. - pub paths: Vec, -} - -impl WorkspaceClaim { - /// A worker with its own workspace root. Never conflicts, always parallel. - pub fn isolated(worker_id: impl Into) -> Self { - Self { - worker_id: worker_id.into(), - isolated: true, - writes: false, - paths: Vec::new(), - } - } - - /// A worker sharing the workspace but never mutating it. - pub fn read_only(worker_id: impl Into) -> Self { - Self { - worker_id: worker_id.into(), - isolated: false, - writes: false, - paths: Vec::new(), - } - } - - /// A worker sharing the workspace and mutating the given owned paths. - /// - /// An empty `paths` is accepted here and reported as - /// [`ClaimConflict::UnboundedWrite`] at planning time, so callers can build - /// claims uniformly and let the planner do the rejecting. - pub fn writing(worker_id: impl Into, paths: Vec) -> Self { - Self { - worker_id: worker_id.into(), - isolated: false, - writes: true, - paths, - } - } -} - -/// Why a claim string could not be parsed into safe relative paths. -/// -/// Carries the offending input rather than a rendered sentence: the phrasing a -/// host shows its users is the host's, not this crate's. -#[derive(Debug, Clone, PartialEq, Eq)] -pub enum ClaimPathError { - /// The path was absolute; claims are always relative to the shared root. - Absolute { - /// The offending entry, as written. - raw: String, - }, - /// The path escaped the shared root via `..` or a platform path prefix. - Escaping { - /// The offending entry, as written. - raw: String, - }, -} - -/// Why a worker cannot be scheduled against the shared workspace. -#[derive(Debug, Clone, PartialEq, Eq)] -pub enum ClaimConflict { - /// The worker mutates a shared workspace without declaring what it owns. - UnboundedWrite { - /// The worker that could not be scheduled. - worker_id: String, - }, - /// The worker's claim overlaps one already granted to an earlier worker. - Overlap { - /// The worker that could not be scheduled. - worker_id: String, - /// The earlier worker holding the conflicting claim. - other_worker_id: String, - /// The path both workers claim. - path: PathBuf, - }, -} - -impl ClaimConflict { - /// The worker this conflict rejected. - pub fn worker_id(&self) -> &str { - match self { - Self::UnboundedWrite { worker_id } => worker_id, - Self::Overlap { worker_id, .. } => worker_id, - } - } -} - -/// How a worker must be dispatched relative to its siblings. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum DispatchMode { - /// Safe to run concurrently with every other worker. - Parallel, - /// Mutates the shared workspace; must not overlap another writer in time. - Serial, -} - -/// Which workers may run concurrently, in input order. -/// -/// [`modes`](Self::modes) is index-aligned with the planner's input. Rejected -/// workers appear in [`conflicts`](Self::conflicts) keyed by that same index and -/// are **not** represented in `modes`, so the two are read together. -#[derive(Debug, Clone, PartialEq, Eq, Default)] -pub struct DispatchPlan { - /// Per-worker dispatch mode, indexed by input position. `None` marks a - /// worker that was rejected and must not run. - pub modes: Vec>, - /// Rejections, paired with the input index of the worker they rejected. - pub conflicts: Vec<(usize, ClaimConflict)>, -} - -impl DispatchPlan { - /// `true` when at least one schedulable worker must run serially. - pub fn has_serial_work(&self) -> bool { - self.modes - .iter() - .any(|mode| matches!(mode, Some(DispatchMode::Serial))) - } - - /// Input indices of the workers that may run concurrently. - pub fn parallel_indices(&self) -> Vec { - self.indices_matching(DispatchMode::Parallel) - } - - /// Input indices of the workers that must run one at a time, in input order. - pub fn serial_indices(&self) -> Vec { - self.indices_matching(DispatchMode::Serial) - } - - /// Input indices of workers scheduled in `wanted` mode. - fn indices_matching(&self, wanted: DispatchMode) -> Vec { - self.modes - .iter() - .enumerate() - .filter_map(|(index, mode)| (*mode == Some(wanted)).then_some(index)) - .collect() - } -} diff --git a/crates/tinyagents-graph/src/parallel/mod.rs b/crates/tinyagents-graph/src/parallel/mod.rs deleted file mode 100644 index 0173c0b38..000000000 --- a/crates/tinyagents-graph/src/parallel/mod.rs +++ /dev/null @@ -1,180 +0,0 @@ -//! Ordered, bounded-concurrency parallel map/reduce with a configurable failure -//! policy. -//! -//! Graph `Send` fanout is the low-level primitive for parallel supersteps, but -//! callers frequently want a *reusable* "run these N items concurrently and -//! reduce the results" helper with deterministic input-order results, a -//! concurrency cap, per-item success/failure isolation, and a policy for what to -//! do when some items fail (fail-fast, collect-all, quorum, best-effort). That -//! is what [`map_reduce`] provides, independent of the graph executor. -//! -//! Results are always returned in **input order** even though items complete out -//! of order: concurrency is driven by -//! [`buffer_unordered`](futures::stream::StreamExt::buffer_unordered), which -//! yields items as they finish, and each future carries its input index so the -//! completed results are re-sorted into input-order slots before being returned. - -mod claims; -mod types; - -pub use claims::{ - ClaimConflict, ClaimPathError, DispatchMode, DispatchPlan, WorkspaceClaim, - parse_relative_claim_paths, paths_overlap, plan_shared_workspace_dispatch, - writes_shared_workspace, -}; -pub use types::{FailurePolicy, ItemOutcome, ParallelOptions, ParallelOutcome}; - -use futures::stream::StreamExt; -use std::future::Future; - -use tinyagents_harness::error::{Result, TinyAgentsError}; - -/// Runs `f` over `items` concurrently (bounded by -/// [`ParallelOptions::max_concurrency`]), collecting per-item outcomes in input -/// order and applying the configured [`FailurePolicy`]. -/// -/// `f` receives each item's input index and the item, and returns a future -/// yielding `Result`. The mapping is deterministic in input order regardless -/// of completion order. -/// -/// # Errors -/// -/// - [`FailurePolicy::FailFast`] returns the first item error (in input order), -/// cancelling remaining in-flight work. -/// - [`FailurePolicy::Quorum`] returns [`TinyAgentsError::Graph`] when fewer than -/// the required number of items succeeded. -/// - [`FailurePolicy::CollectAll`] and [`FailurePolicy::BestEffort`] never error. -pub async fn map_reduce( - items: Vec, - options: ParallelOptions, - f: F, -) -> Result> -where - F: Fn(usize, I) -> Fut, - Fut: Future>, -{ - let total = items.len(); - let concurrency = if options.max_concurrency == 0 { - total.max(1) - } else { - options.max_concurrency - }; - - let item_timeout = options.item_timeout; - // Each future carries its input index so results can be re-ordered. - // `buffer_unordered` bounds concurrency to `concurrency` and yields items as - // they complete; we re-order into `slots` by index. Dropping the stream on a - // fail-fast break cancels any remaining in-flight work. - let mut stream = futures::stream::iter(items.into_iter().enumerate().map(|(index, item)| { - let fut = f(index, item); - async move { - // A per-item timeout turns a slow item into a recoverable failure - // rather than stalling the whole batch. - let result = match item_timeout { - Some(limit) => match tokio::time::timeout(limit, fut).await { - Ok(r) => r, - Err(_) => Err(TinyAgentsError::Timeout(format!( - "parallel item {index} exceeded {} ms", - limit.as_millis() - ))), - }, - None => fut.await, - }; - (index, result) - } - })) - .buffer_unordered(concurrency); - - let mut slots: Vec>> = (0..total).map(|_| None).collect(); - // Under FailFast we return the first failure in *input* order, not the first - // to complete (items finish out of order under `buffer_unordered`). Track the - // lowest-index error seen; once every earlier item has also resolved, no - // smaller-index error can still appear, so that error is final and we can - // drop the stream to cancel the remaining in-flight work. - let mut fail_fast_error: Option<(usize, TinyAgentsError)> = None; - - // The collection loop, wrapped so it can be raced against an overall timeout - // and a cancellation token without duplicating the drain logic. - let collect = async { - while let Some((index, result)) = stream.next().await { - match result { - Ok(value) => slots[index] = Some(Ok(value)), - Err(err) => { - if options.failure_policy == FailurePolicy::FailFast { - // Mark this index resolved so the "all earlier items - // done" check below can see it. - slots[index] = Some(Err(String::new())); - if fail_fast_error - .as_ref() - .is_none_or(|(seen, _)| index < *seen) - { - fail_fast_error = Some((index, err)); - } - let min_index = fail_fast_error.as_ref().map(|(i, _)| *i).unwrap_or(index); - if slots[..min_index].iter().all(Option::is_some) { - break; // dropping the stream cancels remaining work - } - continue; - } - slots[index] = Some(Err(err.to_string())); - } - } - } - }; - - let cancelled = async { - match &options.cancellation { - Some(token) => token.cancelled().await, - None => std::future::pending::<()>().await, - } - }; - - // Race collection against cancellation, all under an optional total timeout. - let raced = async { - tokio::select! { - biased; - _ = cancelled => Err(TinyAgentsError::Cancelled), - () = collect => Ok(()), - } - }; - match options.total_timeout { - Some(limit) => match tokio::time::timeout(limit, raced).await { - Ok(inner) => inner?, - Err(_) => { - return Err(TinyAgentsError::Timeout(format!( - "parallel map/reduce exceeded {} ms", - limit.as_millis() - ))); - } - }, - None => raced.await?, - } - - if let Some((_, err)) = fail_fast_error { - return Err(err); - } - - let mut outcomes: Vec> = Vec::with_capacity(total); - for (index, slot) in slots.into_iter().enumerate() { - if let Some(result) = slot { - outcomes.push(ItemOutcome { index, result }); - } - } - - let success_count = outcomes.iter().filter(|o| o.is_ok()).count(); - - match options.failure_policy { - FailurePolicy::Quorum(required) if success_count < required => Err(TinyAgentsError::Graph( - format!("parallel quorum not met: {success_count}/{required} items succeeded"), - )), - FailurePolicy::BestEffort => { - outcomes.retain(|o| o.is_ok()); - Ok(ParallelOutcome { outcomes }) - } - _ => Ok(ParallelOutcome { outcomes }), - } -} - -#[cfg(test)] -#[path = "mod_tests.rs"] -mod test; diff --git a/crates/tinyagents-graph/src/parallel/mod_tests.rs b/crates/tinyagents-graph/src/parallel/mod_tests.rs deleted file mode 100644 index cdaa3e0ce..000000000 --- a/crates/tinyagents-graph/src/parallel/mod_tests.rs +++ /dev/null @@ -1,218 +0,0 @@ -//! Tests for the parallel map/reduce helper. - -use super::*; -use tinyagents_harness::error::TinyAgentsError; - -#[tokio::test] -async fn preserves_input_order_despite_completion_order() { - // Items complete out of order (larger sleeps first) but results stay ordered. - let out = map_reduce( - vec![3u64, 1, 2], - ParallelOptions::default(), - |_i, n| async move { - tokio::time::sleep(std::time::Duration::from_millis(n)).await; - Ok::<_, TinyAgentsError>(n * 10) - }, - ) - .await - .unwrap(); - - let values: Vec = out.into_successes(); - assert_eq!(values, vec![30, 10, 20]); -} - -#[tokio::test] -async fn collect_all_records_per_item_failure() { - let out = map_reduce( - vec![1, 2, 3, 4], - ParallelOptions::default().with_failure_policy(FailurePolicy::CollectAll), - |_i, n| async move { - if n % 2 == 0 { - Err(TinyAgentsError::Graph(format!("even {n}"))) - } else { - Ok(n) - } - }, - ) - .await - .unwrap(); - - assert_eq!(out.success_count(), 2); - assert_eq!(out.failure_count(), 2); - // Outcomes preserve input order and index. - assert_eq!(out.outcomes[0].index, 0); - assert!(out.outcomes[1].result.is_err()); -} - -#[tokio::test] -async fn fail_fast_returns_first_error() { - let result = map_reduce( - vec![1, 2, 3], - ParallelOptions::default().with_failure_policy(FailurePolicy::FailFast), - |_i, n| async move { - if n == 2 { - Err(TinyAgentsError::Graph("boom".to_string())) - } else { - Ok(n) - } - }, - ) - .await; - assert!(matches!(result, Err(TinyAgentsError::Graph(_)))); -} - -#[tokio::test] -async fn fail_fast_returns_first_error_in_input_order_not_completion_order() { - // Two items fail. The later-index one (index 2) completes almost - // immediately, while the earlier-index one (index 0) fails after a delay. - // FailFast must return the earlier item's error (input order), not the - // first-to-complete error. - let result = map_reduce( - vec![0usize, 1, 2], - ParallelOptions::default().with_failure_policy(FailurePolicy::FailFast), - |i, _n| async move { - match i { - 0 => { - tokio::time::sleep(std::time::Duration::from_millis(30)).await; - Err(TinyAgentsError::Graph("first".to_string())) - } - 2 => Err(TinyAgentsError::Graph("third".to_string())), - _ => Ok(i), - } - }, - ) - .await; - match result { - Err(TinyAgentsError::Graph(msg)) => assert_eq!(msg, "first"), - other => panic!("expected the input-order-first error, got {other:?}"), - } -} - -#[tokio::test] -async fn quorum_requires_minimum_successes() { - let opts = ParallelOptions::default().with_failure_policy(FailurePolicy::Quorum(3)); - let items = vec![1, 2, 3, 4]; - let f = |_i: usize, n: i32| async move { - if n <= 2 { - Ok(n) - } else { - Err(TinyAgentsError::Graph("nope".to_string())) - } - }; - // Only 2 succeed, quorum of 3 not met. - let result = map_reduce(items.clone(), opts, f).await; - assert!(matches!(result, Err(TinyAgentsError::Graph(_)))); - - // Quorum of 2 is met. - let opts2 = ParallelOptions::default().with_failure_policy(FailurePolicy::Quorum(2)); - let ok = map_reduce(items, opts2, |_i, n| async move { - if n <= 2 { - Ok(n) - } else { - Err(TinyAgentsError::Graph("nope".to_string())) - } - }) - .await - .unwrap(); - assert_eq!(ok.success_count(), 2); -} - -#[tokio::test] -async fn best_effort_keeps_only_successes() { - let out = map_reduce( - vec![1, 2, 3], - ParallelOptions::default().with_failure_policy(FailurePolicy::BestEffort), - |_i, n| async move { - if n == 2 { - Err(TinyAgentsError::Graph("skip".to_string())) - } else { - Ok(n) - } - }, - ) - .await - .unwrap(); - assert_eq!(out.into_successes(), vec![1, 3]); -} - -#[tokio::test] -async fn concurrency_cap_bounds_simultaneous_work() { - use std::sync::Arc; - use std::sync::atomic::{AtomicUsize, Ordering}; - - let live = Arc::new(AtomicUsize::new(0)); - let peak = Arc::new(AtomicUsize::new(0)); - let opts = ParallelOptions::default().with_max_concurrency(2); - - map_reduce(vec![0u64; 6], opts, |_i, _n| { - let live = live.clone(); - let peak = peak.clone(); - async move { - let now = live.fetch_add(1, Ordering::SeqCst) + 1; - peak.fetch_max(now, Ordering::SeqCst); - tokio::time::sleep(std::time::Duration::from_millis(5)).await; - live.fetch_sub(1, Ordering::SeqCst); - Ok::<_, TinyAgentsError>(()) - } - }) - .await - .unwrap(); - - assert!(peak.load(Ordering::SeqCst) <= 2, "concurrency exceeded cap"); -} - -#[tokio::test(start_paused = true)] -async fn item_timeout_fails_only_the_slow_item() { - // Item 1 hangs; a per-item timeout turns it into a recoverable failure - // while the fast items still succeed (CollectAll). - let out = map_reduce( - vec![0u64, 3_600_000, 0], - ParallelOptions::default() - .with_failure_policy(FailurePolicy::CollectAll) - .with_item_timeout(std::time::Duration::from_millis(50)), - |_i, ms| async move { - tokio::time::sleep(std::time::Duration::from_millis(ms)).await; - Ok::<_, TinyAgentsError>(ms) - }, - ) - .await - .unwrap(); - - assert_eq!(out.success_count(), 2); - assert_eq!(out.failure_count(), 1); - assert!(!out.outcomes[1].is_ok(), "the hanging item timed out"); -} - -#[tokio::test(start_paused = true)] -async fn total_timeout_aborts_the_batch() { - let err = map_reduce( - vec![3_600_000u64], - ParallelOptions::default().with_total_timeout(std::time::Duration::from_millis(20)), - |_i, ms| async move { - tokio::time::sleep(std::time::Duration::from_millis(ms)).await; - Ok::<_, TinyAgentsError>(ms) - }, - ) - .await - .expect_err("total timeout should abort"); - assert!(matches!(err, TinyAgentsError::Timeout(_))); -} - -#[tokio::test] -async fn cancellation_token_stops_the_batch() { - use tinyagents_harness::cancel::CancellationToken; - - let token = CancellationToken::new(); - token.cancel(); - let err = map_reduce( - vec![0u64], - ParallelOptions::default().with_cancellation(token), - |_i, _n| async move { - tokio::time::sleep(std::time::Duration::from_secs(3600)).await; - Ok::<_, TinyAgentsError>(0u64) - }, - ) - .await - .expect_err("a pre-cancelled token should abort"); - assert!(matches!(err, TinyAgentsError::Cancelled)); -} diff --git a/crates/tinyagents-graph/src/parallel/types.rs b/crates/tinyagents-graph/src/parallel/types.rs deleted file mode 100644 index d8c6eae8b..000000000 --- a/crates/tinyagents-graph/src/parallel/types.rs +++ /dev/null @@ -1,129 +0,0 @@ -//! Types for the ordered, bounded parallel map/reduce helper. - -/// What [`map_reduce`](super::map_reduce) does when some items fail. -#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] -pub enum FailurePolicy { - /// Stop at the first failing item (in input order) and return its error; - /// remaining in-flight work is dropped/cancelled. - FailFast, - /// Run every item to completion and always return `Ok`, recording per-item - /// success or failure for the caller to inspect. - #[default] - CollectAll, - /// Run every item to completion; return `Ok` when at least `n` items - /// succeeded, otherwise return an error. - Quorum(usize), - /// Run every item to completion and always return `Ok`, silently keeping - /// only the successful outputs. - BestEffort, -} - -/// Options controlling parallel execution. -#[derive(Clone, Debug, Default)] -pub struct ParallelOptions { - /// Maximum number of items to run concurrently. `0` means unbounded. - pub max_concurrency: usize, - /// How to react to per-item failures. - pub failure_policy: FailurePolicy, - /// Per-item wall-clock timeout. An item exceeding it fails with a timeout - /// message (then handled by the [`FailurePolicy`]); `None` means no per-item - /// bound. - pub item_timeout: Option, - /// Overall wall-clock timeout for the whole map/reduce. On elapse the call - /// returns [`TinyAgentsError::Timeout`][tinyagents_harness::error::TinyAgentsError::Timeout] - /// and remaining in-flight work is dropped; `None` means no overall bound. - pub total_timeout: Option, - /// Cooperative cancellation token: when cancelled, the call stops collecting - /// and returns [`TinyAgentsError::Cancelled`][tinyagents_harness::error::TinyAgentsError::Cancelled], - /// dropping remaining in-flight work. - pub cancellation: Option, -} - -impl ParallelOptions { - /// Bounds concurrency to `n` simultaneous items. - pub fn with_max_concurrency(mut self, n: usize) -> Self { - self.max_concurrency = n; - self - } - - /// Sets the failure policy. - pub fn with_failure_policy(mut self, policy: FailurePolicy) -> Self { - self.failure_policy = policy; - self - } - - /// Bounds each item to `timeout` of wall-clock execution. - pub fn with_item_timeout(mut self, timeout: std::time::Duration) -> Self { - self.item_timeout = Some(timeout); - self - } - - /// Bounds the whole map/reduce to `timeout` of wall-clock execution. - pub fn with_total_timeout(mut self, timeout: std::time::Duration) -> Self { - self.total_timeout = Some(timeout); - self - } - - /// Attaches a cooperative cancellation token. - pub fn with_cancellation( - mut self, - token: tinyagents_harness::cancel::CancellationToken, - ) -> Self { - self.cancellation = Some(token); - self - } -} - -/// The outcome of one item in a parallel map/reduce, tagged with its input -/// index so callers can correlate results back to inputs. -#[derive(Clone, Debug, PartialEq, Eq)] -pub struct ItemOutcome { - /// The input index (0-based) this outcome corresponds to. - pub index: usize, - /// `Ok(value)` on success, `Err(message)` on failure. - pub result: std::result::Result, -} - -impl ItemOutcome { - /// Returns `true` when the item succeeded. - pub fn is_ok(&self) -> bool { - self.result.is_ok() - } -} - -/// The collected result of a parallel map/reduce, preserving input order. -#[derive(Clone, Debug, PartialEq, Eq)] -pub struct ParallelOutcome { - /// Per-item outcomes, in input order. Under [`FailurePolicy::BestEffort`] - /// only successful items are present. - pub outcomes: Vec>, -} - -impl ParallelOutcome { - /// Number of successful items. - pub fn success_count(&self) -> usize { - self.outcomes.iter().filter(|o| o.is_ok()).count() - } - - /// Number of failed items. - pub fn failure_count(&self) -> usize { - self.outcomes.iter().filter(|o| !o.is_ok()).count() - } - - /// Borrows every successful value, in input order. - pub fn successes(&self) -> Vec<&T> { - self.outcomes - .iter() - .filter_map(|o| o.result.as_ref().ok()) - .collect() - } - - /// Consumes the outcome and returns only the successful values, in input - /// order. - pub fn into_successes(self) -> Vec { - self.outcomes - .into_iter() - .filter_map(|o| o.result.ok()) - .collect() - } -} diff --git a/crates/tinyagents-graph/src/recursion/README.md b/crates/tinyagents-graph/src/recursion/README.md deleted file mode 100644 index 1956a4123..000000000 --- a/crates/tinyagents-graph/src/recursion/README.md +++ /dev/null @@ -1,85 +0,0 @@ -# graph::recursion - -Recursion policy and depth tracking: the bound-and-observe contract for -nested graph/subgraph/sub-agent recursion. - -The graph runtime allows recursion — a graph can embed itself as a subgraph, a -subgraph can call back into its parent, an agent node can call another agent -that re-enters the same graph, a router can loop nodes until state converges, -and a `Send` fanout can schedule many recursive child tasks — but only under -explicit, enforced limits. This module is that contract; `compiled::executor` -is the sole enforcer, building one `RecursionStack` per run, pushing/popping -frames around every nested call, and stamping the live stack into checkpoint -metadata and `GraphEvent::RecursionDepthChanged`. - -## Public surface - -- `RecursionFrame` — one level of the run tree: `graph_id`, `node_id` (the - hosting node, `None` for a root frame), `run_id`, `task_id`, `namespace`, - `depth`, and `parent` (the enclosing frame's run id). Serialized into - checkpoint metadata so a UI can render nested runs without replaying event - logs. -- `RecursionPolicy` — the configured caps, tracked independently: - `max_depth` (run-tree depth, default 25), `max_visits_per_node` (optional - cap on activations of one node in a run, default unbounded), - `max_total_steps` (supersteps per run, default 1000). `Default` gives the - conservative defaults above. -- `RecursionStack` — the live frame stack for one executing run, paired with - the policy that bounds it: - - `new(policy)` / `with_frames(frames, policy)` (the latter seeds inherited - parent frames for a subgraph/sub-agent child run). - - `push`/`pop` — symmetric frame push/pop; `push` enforces `max_depth` - (`TinyAgentsError::SubAgentDepth`) and leaves the frame unpushed on - failure. - - `depth()` / `frames()` — current depth and the frame list, root-first. - - `check_total_steps(steps)` — enforces `max_total_steps` - (`TinyAgentsError::RecursionLimit`). - - `record_node_visit(counts, node)` — increments and enforces - `max_visits_per_node` (`TinyAgentsError::NodeVisitLimit`) when configured. -- `ChildRun` — a reference to a child run spawned from a subgraph/sub-agent - node: `node`, `graph_id`, `run_id`, `root_run_id`, and rolled-up `usage` - (only populated for sub-agent children; subgraph children track usage - through their own model calls). -- `ChildRunSink` — a thread-safe, per-run collector the executor hands to node - contexts so a subgraph/sub-agent node can report the `ChildRun` it spawned; - `record`/`drain` (drained once per superstep boundary). -- `RunTree` — a flat, after-the-fact parent/child lineage view derived from a - completed run: `run_id`, `root_run_id`, `parent_run_id`, `children`; - `is_root()` reports whether `parent_run_id` is `None`. The run-id - counterpart to the live `RecursionStack`. - -## Files - -| File | Role | -| --- | --- | -| `types.rs` | All type definitions and their `impl` blocks. | -| `mod.rs` | Module doc only (re-exports; no additional behavior). | -| `mod_tests.rs` | Unit tests for the stack contract and its enforcement inside the executor. | - -## How it fits together - -- `builder::types::GraphBuilder`/`compiled::types::CompiledGraph` carry - `recursion_policy`, `recursion_frames` (inherited from an enclosing run), - and `recursion_node` (the hosting node when embedded); `subgraph` and - `subagent_node` are what seed `recursion_frames` so a nested run extends the - parent's tree instead of starting fresh. -- `compiled::executor` enforces the three caps with distinct errors - (`SubAgentDepth`, `NodeVisitLimit`, `RecursionLimit`), records the current - stack into checkpoint metadata under a `recursion` array, and emits - `GraphEvent::RecursionDepthChanged` on depth changes. -- `compiled::types::GraphExecution::run_tree()` builds a `RunTree` from the - run's own ids plus its accumulated `child_runs` — the completed-run - counterpart to the live `RecursionStack`. -- `status::GraphRunStatus` carries `root_run_id`/`parent_run_id` in the same - shape, so the live-status and completed-run views of the recursion tree - stay consistent. - -## Operational constraints - -- `RecursionStack::push`/`pop` must stay balanced around every nested call — - an executor path that pushes without a matching pop (e.g. on an early - return from an error branch) leaves the stack permanently deeper than the - actual call depth for the rest of that run. -- `max_visits_per_node` counts are per-run, not persisted across a - resume/retry from a checkpoint — a long-running resumed thread does not - inherit visit counts from before the checkpoint boundary. diff --git a/crates/tinyagents-graph/src/recursion/mod.rs b/crates/tinyagents-graph/src/recursion/mod.rs deleted file mode 100644 index f9e42a834..000000000 --- a/crates/tinyagents-graph/src/recursion/mod.rs +++ /dev/null @@ -1,35 +0,0 @@ -//! Recursion policy and depth tracking for the durable graph runtime. -//! -//! The graph allows recursion — a graph can invoke itself as a subgraph, a -//! subgraph can call back into its parent, an agent node can call another agent -//! that re-enters the same graph, a router can loop nodes until state -//! converges, and a `Send` fanout can schedule many recursive child tasks — but -//! only with explicit limits and tracking. -//! -//! This module provides that contract: -//! -//! - [`RecursionFrame`] names one level of the run tree (graph, node, run, task, -//! namespace, depth, parent run). -//! - [`RecursionPolicy`] holds the caps: `max_depth` (run-tree depth), -//! `max_visits_per_node` (node-loop recursion), and `max_total_steps` -//! (super-steps per run). -//! - [`RecursionStack`] is the live frame stack: every graph/subgraph/sub-agent -//! call pushes a frame (enforcing `max_depth`) and every return pops it, so -//! the stack stays symmetric and always describes the path from the root run -//! to the current one. -//! -//! The executor (`CompiledGraph::execute`) builds one stack per -//! run, tracks graph-call depth separately from node-loop visits, enforces the -//! caps with clear recursion errors ([`crate::TinyAgentsError::SubAgentDepth`], -//! [`crate::TinyAgentsError::NodeVisitLimit`], -//! [`crate::TinyAgentsError::RecursionLimit`]), records the current stack into -//! checkpoint metadata under a `recursion` array, and emits the live depth on -//! [`crate::GraphEvent::RecursionDepthChanged`]. - -mod types; - -pub use types::{ChildRun, ChildRunSink, RecursionFrame, RecursionPolicy, RecursionStack, RunTree}; - -#[cfg(test)] -#[path = "mod_tests.rs"] -mod test; diff --git a/crates/tinyagents-graph/src/recursion/mod_tests.rs b/crates/tinyagents-graph/src/recursion/mod_tests.rs deleted file mode 100644 index 4bce2ad35..000000000 --- a/crates/tinyagents-graph/src/recursion/mod_tests.rs +++ /dev/null @@ -1,210 +0,0 @@ -//! Unit tests for recursion policy, frames, and depth tracking — both the -//! standalone [`RecursionStack`] contract and its enforcement inside the -//! superstep executor. - -use super::*; -use crate::TinyAgentsError; -use crate::builder::{GraphBuilder, NodeContext}; -use crate::checkpoint::{Checkpointer, InMemoryCheckpointer}; -use crate::command::NodeResult; -use crate::stream::{CollectingSink, GraphEvent}; -use std::collections::HashMap; -use std::sync::Arc; -use tinyagents_harness::ids::{GraphId, NodeId, RunId}; - -fn frame(run: &str, depth: usize, parent: Option<&str>) -> RecursionFrame { - RecursionFrame { - graph_id: GraphId::new("g"), - node_id: None, - run_id: RunId::new(run), - task_id: None, - namespace: Vec::new(), - depth, - parent: parent.map(RunId::new), - } -} - -#[test] -fn default_policy_is_conservative() { - let p = RecursionPolicy::default(); - assert_eq!(p.max_depth, 25); - assert_eq!(p.max_visits_per_node, None); - assert_eq!(p.max_total_steps, 1000); -} - -#[test] -fn frames_push_and_pop_are_symmetric() { - let mut stack = RecursionStack::new(RecursionPolicy::default()); - assert_eq!(stack.depth(), 0); - stack.push(frame("a", 0, None)).unwrap(); - stack.push(frame("b", 1, Some("a"))).unwrap(); - stack.push(frame("c", 2, Some("b"))).unwrap(); - assert_eq!(stack.depth(), 3); - assert_eq!(stack.frames().len(), 3); - - assert_eq!(stack.pop().unwrap().run_id, RunId::new("c")); - assert_eq!(stack.pop().unwrap().run_id, RunId::new("b")); - assert_eq!(stack.pop().unwrap().run_id, RunId::new("a")); - assert_eq!(stack.depth(), 0); - assert!(stack.pop().is_none()); -} - -#[test] -fn depth_limit_trips_on_push() { - let policy = RecursionPolicy { - max_depth: 2, - max_visits_per_node: None, - max_total_steps: 1000, - }; - let mut stack = RecursionStack::new(policy); - stack.push(frame("a", 0, None)).unwrap(); - stack.push(frame("b", 1, Some("a"))).unwrap(); - let err = stack.push(frame("c", 2, Some("b"))).unwrap_err(); - assert!(matches!(err, TinyAgentsError::SubAgentDepth(2))); - // The rejected frame is not retained, so the stack stays consistent. - assert_eq!(stack.depth(), 2); -} - -#[test] -fn with_frames_seeds_inherited_depth() { - let stack = - RecursionStack::with_frames(vec![frame("root", 0, None)], RecursionPolicy::default()); - assert_eq!(stack.depth(), 1); - assert_eq!(stack.frames()[0].run_id, RunId::new("root")); -} - -#[test] -fn node_visit_limit_trips() { - let policy = RecursionPolicy { - max_depth: 25, - max_visits_per_node: Some(2), - max_total_steps: 1000, - }; - let stack = RecursionStack::new(policy); - let mut counts: HashMap = HashMap::new(); - let node = NodeId::new("loop"); - stack.record_node_visit(&mut counts, &node).unwrap(); - stack.record_node_visit(&mut counts, &node).unwrap(); - let err = stack.record_node_visit(&mut counts, &node).unwrap_err(); - assert!(matches!( - err, - TinyAgentsError::NodeVisitLimit { limit: 2, .. } - )); -} - -#[test] -fn node_visit_unbounded_when_unset() { - let stack = RecursionStack::new(RecursionPolicy::default()); - let mut counts: HashMap = HashMap::new(); - let node = NodeId::new("loop"); - for _ in 0..100 { - stack.record_node_visit(&mut counts, &node).unwrap(); - } -} - -#[test] -fn total_steps_check() { - let policy = RecursionPolicy { - max_depth: 25, - max_visits_per_node: None, - max_total_steps: 3, - }; - let stack = RecursionStack::new(policy); - assert!(stack.check_total_steps(2).is_ok()); - let err = stack.check_total_steps(3).unwrap_err(); - assert!(matches!(err, TinyAgentsError::RecursionLimit(3))); -} - -// ---- Executor wiring ------------------------------------------------------- - -#[tokio::test] -async fn policy_max_total_steps_trips_in_executor() { - let graph = GraphBuilder::::overwrite() - .add_node("loop", |s, _c: NodeContext| async move { - Ok(NodeResult::Update(s + 1)) - }) - .set_entry("loop") - .add_edge("loop", "loop") - .compile() - .unwrap() - .with_recursion_policy(RecursionPolicy { - max_depth: 25, - max_visits_per_node: None, - max_total_steps: 3, - }); - - let err = graph.run(0).await.unwrap_err(); - assert!(matches!(err, TinyAgentsError::RecursionLimit(3))); -} - -#[tokio::test] -async fn policy_node_visit_limit_trips_in_executor() { - let graph = GraphBuilder::::overwrite() - .add_node("loop", |s, _c: NodeContext| async move { - Ok(NodeResult::Update(s + 1)) - }) - .set_entry("loop") - .add_edge("loop", "loop") - .compile() - .unwrap() - .with_recursion_policy(RecursionPolicy { - max_depth: 25, - max_visits_per_node: Some(2), - max_total_steps: 1000, - }); - - let err = graph.run(0).await.unwrap_err(); - assert!(matches!( - err, - TinyAgentsError::NodeVisitLimit { limit: 2, .. } - )); -} - -#[tokio::test] -async fn run_emits_recursion_depth_event() { - let sink = Arc::new(CollectingSink::default()); - let graph = GraphBuilder::::overwrite() - .add_node("only", |s, _c: NodeContext| async move { - Ok(NodeResult::Update(s + 1)) - }) - .set_entry("only") - .set_finish("only") - .compile() - .unwrap() - .with_event_sink(sink.clone()); - - graph.run(0).await.unwrap(); - let saw_depth = sink - .events() - .iter() - .any(|e| matches!(e, GraphEvent::RecursionDepthChanged { depth: 1 })); - assert!(saw_depth, "expected a RecursionDepthChanged {{ depth: 1 }}"); -} - -#[tokio::test] -async fn checkpoint_metadata_records_recursion_stack() { - let checkpointer = Arc::new(InMemoryCheckpointer::::new()); - let graph = GraphBuilder::::overwrite() - .add_node("a", |s, _c: NodeContext| async move { - Ok(NodeResult::Update(s + 1)) - }) - .set_entry("a") - .set_finish("a") - .compile() - .unwrap() - .with_checkpointer(checkpointer.clone()); - - graph.run_with_thread("t1", 0).await.unwrap(); - - let checkpoint = checkpointer.get("t1", None).await.unwrap().unwrap(); - let recursion = checkpoint - .metadata - .get("recursion") - .expect("metadata carries a recursion array"); - let frames = recursion.as_array().expect("recursion is an array"); - assert_eq!(frames.len(), 1, "one frame for the top-level run"); - assert_eq!(frames[0]["depth"], 0); - assert!(frames[0]["run_id"].is_string()); - // The original `source` field is preserved alongside the new array. - assert_eq!(checkpoint.metadata["source"], "loop"); -} diff --git a/crates/tinyagents-graph/src/recursion/types.rs b/crates/tinyagents-graph/src/recursion/types.rs deleted file mode 100644 index 5c8fb8945..000000000 --- a/crates/tinyagents-graph/src/recursion/types.rs +++ /dev/null @@ -1,273 +0,0 @@ -//! Recursion policy and depth-tracking types. -//! -//! See the module [`mod`](super) docs for how these are wired into the -//! executor. This file holds the plain data definitions: a [`RecursionFrame`] -//! (one level of the run tree), a [`RecursionPolicy`] (the configured caps), and -//! a [`RecursionStack`] (the live frame stack plus enforcement helpers). - -use std::sync::{Arc, Mutex}; - -use serde::{Deserialize, Serialize}; - -use crate::{Result, TinyAgentsError}; -use tinyagents_harness::ids::{CheckpointId, GraphId, NodeId, RunId, TaskId}; - -/// One level of the graph/subgraph/sub-agent recursion tree. -/// -/// A frame is pushed for every recursive call (a graph invoking a subgraph, a -/// subgraph invoking a sub-agent, a router looping a graph back into itself) and -/// popped on return, so the live [`RecursionStack`] always describes the path -/// from the root run down to the currently-executing run. Frames are serialized -/// into checkpoint metadata so a UI can render nested runs without -/// reconstructing the tree from event logs. -#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] -pub struct RecursionFrame { - /// The graph this frame is executing. - pub graph_id: GraphId, - /// The hosting node, when this frame was entered from a parent node (the - /// root frame of a top-level run has `None`). - pub node_id: Option, - /// The run id of this frame's execution. - pub run_id: RunId, - /// The scheduled task this frame descends from, when graph-backed (for - /// example a `Send` fanout task). - pub task_id: Option, - /// The checkpoint namespace of this frame's graph (empty for top-level). - pub namespace: Vec, - /// Zero-based depth of this frame in the recursion tree (the root is `0`). - pub depth: usize, - /// The run id of the enclosing frame, when this frame is a child run. - pub parent: Option, -} - -/// Explicit limits that bound recursive graph execution. -/// -/// The three caps are tracked independently so graph-call recursion (nested -/// runs), node-loop recursion (a router cycling the same node), and total work -/// (super-steps) can each be reasoned about and surfaced separately: -/// -/// - `max_depth` bounds the run-tree depth ([`RecursionStack::push`] enforces -/// it, failing with [`TinyAgentsError::SubAgentDepth`]). -/// - `max_visits_per_node` optionally bounds how many times any single node may -/// be activated within one run (failing with -/// [`TinyAgentsError::NodeVisitLimit`]). -/// - `max_total_steps` bounds the number of super-steps a single run may execute -/// (failing with [`TinyAgentsError::RecursionLimit`]). -#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] -pub struct RecursionPolicy { - /// Maximum run-tree depth (root run included). Reaching it on a push fails. - pub max_depth: usize, - /// Optional cap on activations of a single node within one run. - pub max_visits_per_node: Option, - /// Maximum number of super-steps a single run may execute. - pub max_total_steps: usize, -} - -impl Default for RecursionPolicy { - /// A conservative default: 25 levels of run-tree depth, no per-node visit - /// cap, and 1000 super-steps per run. - fn default() -> Self { - Self { - max_depth: 25, - max_visits_per_node: None, - max_total_steps: 1000, - } - } -} - -/// The live stack of [`RecursionFrame`]s for an executing run tree, paired with -/// the [`RecursionPolicy`] that bounds it. -/// -/// Every graph/subgraph/sub-agent call [`push`](RecursionStack::push)es a frame -/// (enforcing `max_depth`) and every return [`pop`](RecursionStack::pop)s it, so -/// the stack stays symmetric. The executor builds one stack per run from the -/// inherited parent frames, pushes this run's frame, and consults it when -/// persisting checkpoints and enforcing the step cap. -#[derive(Clone, Debug, Default)] -pub struct RecursionStack { - frames: Vec, - policy: RecursionPolicy, -} - -impl RecursionStack { - /// Creates an empty stack governed by `policy`. - pub fn new(policy: RecursionPolicy) -> Self { - Self { - frames: Vec::new(), - policy, - } - } - - /// Creates a stack pre-seeded with the inherited `frames` of an enclosing - /// run, governed by `policy`. Used when a subgraph/sub-agent run continues - /// the parent's recursion tree. - pub fn with_frames(frames: Vec, policy: RecursionPolicy) -> Self { - Self { frames, policy } - } - - /// The policy bounding this stack. - pub fn policy(&self) -> &RecursionPolicy { - &self.policy - } - - /// The current recursion depth (number of frames on the stack). - pub fn depth(&self) -> usize { - self.frames.len() - } - - /// The frames currently on the stack, root-first. - pub fn frames(&self) -> &[RecursionFrame] { - &self.frames - } - - /// Pushes a recursion frame, enforcing [`RecursionPolicy::max_depth`]. - /// - /// Returns [`TinyAgentsError::SubAgentDepth`] when the push would make the - /// stack deeper than `max_depth` allows; the frame is **not** pushed in that - /// case, so a caller that recovers keeps a consistent stack. - pub fn push(&mut self, frame: RecursionFrame) -> Result<()> { - if self.frames.len() + 1 > self.policy.max_depth { - return Err(TinyAgentsError::SubAgentDepth(self.policy.max_depth)); - } - self.frames.push(frame); - Ok(()) - } - - /// Pops the deepest frame, returning it (or `None` when empty). - pub fn pop(&mut self) -> Option { - self.frames.pop() - } - - /// Enforces [`RecursionPolicy::max_total_steps`] against an executed-step - /// count, returning [`TinyAgentsError::RecursionLimit`] when the run has - /// reached the cap. - pub fn check_total_steps(&self, steps: usize) -> Result<()> { - if steps >= self.policy.max_total_steps { - return Err(TinyAgentsError::RecursionLimit(self.policy.max_total_steps)); - } - Ok(()) - } - - /// Records one activation of `node` against `counts` and enforces - /// [`RecursionPolicy::max_visits_per_node`] when configured. - /// - /// Increments the per-node counter, returning - /// [`TinyAgentsError::NodeVisitLimit`] when the node has now been activated - /// more times than the policy allows. - pub fn record_node_visit( - &self, - counts: &mut std::collections::HashMap, - node: &NodeId, - ) -> Result<()> { - let Some(max) = self.policy.max_visits_per_node else { - return Ok(()); - }; - let count = counts.entry(node.clone()).or_insert(0); - *count += 1; - if *count > max { - return Err(TinyAgentsError::NodeVisitLimit { - node: node.to_string(), - limit: max, - }); - } - Ok(()) - } -} - -/// A reference to a child run spawned from a node within an enclosing run. -/// -/// Recorded when a subgraph node embeds and runs a [`CompiledGraph`](crate::CompiledGraph): -/// the child gets its own [`run_id`](ChildRun::run_id) while preserving the -/// enclosing run's [`root_run_id`](ChildRun::root_run_id), so a caller can -/// reconstruct the parent/child run lineage after a run completes. Child runs -/// are surfaced on [`GraphExecution`](crate::GraphExecution) and embedded -/// into the parent checkpoint metadata under a `child_runs` array. -#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] -pub struct ChildRun { - /// The embedding node that ran the child graph. - pub node: NodeId, - /// The graph id of the embedded child. - pub graph_id: GraphId, - /// The child run's own (distinct) run id. - pub run_id: RunId, - /// The shared root run id of the whole recursion tree (preserved from the - /// enclosing run). - pub root_run_id: RunId, - /// Token usage rolled up from the child run, when the child reported any. - /// - /// Subgraph children leave this at the default (their usage is tracked by - /// their own model calls); a [`mod@crate::subagent_node`] sub-agent child - /// folds the delegated harness agent's [`UsageTotals`](tinyinference_llm::usage::UsageTotals) here so it is visible - /// 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 -/// subgraph node can report the [`ChildRun`] it spawned back to the enclosing -/// run. -/// -/// One sink is created per run; the executor drains it at each superstep -/// boundary to embed that step's child runs into the boundary checkpoint and to -/// accumulate them onto the final [`GraphExecution`](crate::GraphExecution). -#[derive(Clone, Debug, Default)] -pub struct ChildRunSink { - inner: Arc>>, -} - -impl ChildRunSink { - /// Creates an empty sink. - pub fn new() -> Self { - Self::default() - } - - /// Records a spawned child run. - pub fn record(&self, child: ChildRun) { - if let Ok(mut guard) = self.inner.lock() { - guard.push(child); - } - } - - /// Removes and returns every child run recorded since the last drain. - pub fn drain(&self) -> Vec { - match self.inner.lock() { - Ok(mut guard) => std::mem::take(&mut *guard), - Err(_) => Vec::new(), - } - } -} - -/// A flat parent/child run-lineage view derived from a completed -/// [`GraphExecution`](crate::GraphExecution). -/// -/// `RunTree` is the run-id counterpart to a [`RecursionStack`]: where the stack -/// describes the live path while a run executes, a `RunTree` is the after-the-fact -/// summary a caller reads to see this run's id, the shared root, the enclosing -/// parent (when this was itself a child run), and every child run it spawned. -#[derive(Clone, Debug, PartialEq, Eq)] -pub struct RunTree { - /// This run's id. - pub run_id: RunId, - /// The root run id of the whole recursion tree (equals `run_id` for a - /// top-level run). - pub root_run_id: RunId, - /// The enclosing run's id, when this run was spawned as a child. - pub parent_run_id: Option, - /// The child runs spawned from nodes within this run. - pub children: Vec, -} - -impl RunTree { - /// True when this run is the root of its recursion tree. - pub fn is_root(&self) -> bool { - self.parent_run_id.is_none() - } -} diff --git a/crates/tinyagents-graph/src/reducer/README.md b/crates/tinyagents-graph/src/reducer/README.md deleted file mode 100644 index a6f5bb8ea..000000000 --- a/crates/tinyagents-graph/src/reducer/README.md +++ /dev/null @@ -1,79 +0,0 @@ -# graph::reducer - -The deterministic fan-in point of the durable executor: how updates from one -or more active nodes in a superstep are folded back into committed state. - -When a superstep runs several active nodes — including concurrent branches -(`GraphBuilder::with_parallel`) or a node that merged results from a subgraph -or sub-agent — `compiled::executor` folds their `NodeResult` updates through -the configured reducer in deterministic active-set order at the step -boundary, so the merged state is reproducible regardless of which branch -finished first (see `compiled::README.md`'s "Sequential vs. parallel steps"). -This module supplies the trait definitions and the built-in reducer set; -`channel` builds a richer per-field reducer on top of the same idea for -channel-style state. - -## Public surface - -### Traits - -- `Reducer` — merges two values of the same channel type: - `reduce(&self, current: T, update: T) -> Result`. Used for channel-style - state where each key has its own merge policy (see `channel`). -- `StateReducer` — merges a partial `Update` into the whole - `State`: `apply(&self, state: State, update: Update) -> Result`. This - is the contract `CompiledGraph` executes against; `GraphBuilder::set_reducer` - configures it, and `GraphBuilder::overwrite` is shorthand for the - whole-state-overwrite default. - -### Built-in `Reducer` markers - -- `OverwriteReducer` — last-value semantics (`update` replaces `current`). -- `AppendReducer` — appends the update vector onto the current vector. -- `SetUnionReducer` — unions the update vector into the current vector, - skipping duplicates and preserving first-seen order (requires `T: Eq + Hash - + Clone`). -- `MinReducer` / `MaxReducer` — keeps the smaller/larger of the two values - (requires `T: PartialOrd`). -- `ClosureReducer` — a custom binary reducer backed by an `Fn(T, T) -> - Result` closure; construct with `ClosureReducer::new`. - -### `StateReducer` implementations - -- `OverwriteStateReducer` — overwrites whole state with the update; the - default for `State == Update` graphs (`GraphBuilder::overwrite`). -- `ClosureStateReducer` — a custom state reducer backed by - an `Fn(State, Update) -> Result` closure; construct with - `ClosureStateReducer::new`. - -## Files - -| File | Role | -| --- | --- | -| `types.rs` | `Reducer`, `StateReducer` traits; the built-in marker structs. | -| `mod.rs` | `impl` blocks wiring each marker/closure type to its trait. | -| `mod_tests.rs` | Unit tests for every built-in reducer and the closure-backed ones. | - -## How it fits together - -- `builder::GraphBuilder::set_reducer` accepts any `StateReducer`; `compile()` stores it as `Arc>` on the - frozen `CompiledGraph`. -- `compiled::executor` is the sole caller of `StateReducer::apply`, invoked - once per completed activation's update at each superstep boundary, in - active-set order. -- `channel` is a parallel, richer reducer surface: it applies a per-field - `Reducer` to each named channel and detects same-step concurrent writes - a plain `StateReducer` would silently overwrite — see `channel/README.md`. - -## Operational constraints - -- A `Reducer`/`StateReducer` implementation must be a pure, order-sensitive - merge over exactly the two values it is given — the executor relies on - applying updates in deterministic active-set order to make parallel-step - results reproducible; a reducer with side effects or hidden ordering - dependencies breaks that guarantee. -- `SetUnionReducer`'s duplicate detection and `Min`/`MaxReducer`'s ordering - both depend on the element type's `Eq`/`Hash`/`PartialOrd` impls matching - the domain's notion of equality/ordering — a type with a `PartialOrd` that - disagrees with intended domain order will silently keep the "wrong" value. diff --git a/crates/tinyagents-graph/src/reducer/mod.rs b/crates/tinyagents-graph/src/reducer/mod.rs deleted file mode 100644 index ac3acf3e0..000000000 --- a/crates/tinyagents-graph/src/reducer/mod.rs +++ /dev/null @@ -1,135 +0,0 @@ -//! Reducer trait implementations. -//! -//! Reducers are the deterministic fan-in point of the recursive runtime: when a -//! superstep runs several active nodes — including parallel branches or a node -//! that merged results from a subgraph or sub-agent — the executor folds their -//! updates back into one committed state through a reducer, in deterministic -//! active-set order, so merged state is reproducible regardless of completion -//! order. This module supplies the built-in channel/state reducers and the -//! closure-backed escape hatches; see `types` for the trait and marker -//! definitions. - -mod types; - -pub use types::{ - AppendReducer, ClosureReducer, ClosureStateReducer, MaxReducer, MinReducer, OverwriteReducer, - OverwriteStateReducer, Reducer, SetUnionReducer, StateReducer, -}; - -use std::collections::HashSet; -use std::marker::PhantomData; - -use crate::Result; - -impl Reducer for OverwriteReducer -where - T: Send + Sync, -{ - fn reduce(&self, _current: T, update: T) -> Result { - Ok(update) - } -} - -impl Reducer> for AppendReducer -where - T: Send + Sync, -{ - fn reduce(&self, mut current: Vec, mut update: Vec) -> Result> { - current.append(&mut update); - Ok(current) - } -} - -impl Reducer> for SetUnionReducer -where - T: Eq + std::hash::Hash + Clone + Send + Sync, -{ - fn reduce(&self, current: Vec, update: Vec) -> Result> { - let mut seen: HashSet = current.iter().cloned().collect(); - let mut out = current; - for item in update { - if seen.insert(item.clone()) { - out.push(item); - } - } - Ok(out) - } -} - -impl Reducer for MinReducer -where - T: PartialOrd + Send + Sync, -{ - fn reduce(&self, current: T, update: T) -> Result { - Ok(if update < current { update } else { current }) - } -} - -impl Reducer for MaxReducer -where - T: PartialOrd + Send + Sync, -{ - fn reduce(&self, current: T, update: T) -> Result { - Ok(if update > current { update } else { current }) - } -} - -impl ClosureReducer -where - F: Fn(T, T) -> Result + Send + Sync, -{ - /// Creates a custom reducer from a binary merge closure. - pub fn new(f: F) -> Self { - Self { - f, - _marker: PhantomData, - } - } -} - -impl Reducer for ClosureReducer -where - T: Send + Sync, - F: Fn(T, T) -> Result + Send + Sync, -{ - fn reduce(&self, current: T, update: T) -> Result { - (self.f)(current, update) - } -} - -impl StateReducer for OverwriteStateReducer -where - State: Send + Sync, -{ - fn apply(&self, _state: State, update: State) -> Result { - Ok(update) - } -} - -impl ClosureStateReducer -where - F: Fn(State, Update) -> Result + Send + Sync, -{ - /// Creates a custom state reducer from an `apply` closure. - pub fn new(f: F) -> Self { - Self { - f, - _marker: PhantomData, - } - } -} - -impl StateReducer for ClosureStateReducer -where - State: Send + Sync, - Update: Send + Sync, - F: Fn(State, Update) -> Result + Send + Sync, -{ - fn apply(&self, state: State, update: Update) -> Result { - (self.f)(state, update) - } -} - -#[cfg(test)] -#[path = "mod_tests.rs"] -mod test; diff --git a/crates/tinyagents-graph/src/reducer/mod_tests.rs b/crates/tinyagents-graph/src/reducer/mod_tests.rs deleted file mode 100644 index 40e45c0cf..000000000 --- a/crates/tinyagents-graph/src/reducer/mod_tests.rs +++ /dev/null @@ -1,56 +0,0 @@ -//! Unit tests for the built-in reducers (overwrite, append, set-union, min, -//! max) and the closure-backed channel/state reducers. - -use super::*; - -#[test] -fn overwrite_keeps_update() { - let r = OverwriteReducer; - assert_eq!(Reducer::reduce(&r, 1, 2).unwrap(), 2); -} - -#[test] -fn append_concatenates() { - let r = AppendReducer; - assert_eq!(r.reduce(vec![1, 2], vec![3, 4]).unwrap(), vec![1, 2, 3, 4]); -} - -#[test] -fn set_union_dedups() { - let r = SetUnionReducer; - assert_eq!(r.reduce(vec![1, 2], vec![2, 3]).unwrap(), vec![1, 2, 3]); -} - -#[test] -fn min_and_max() { - assert_eq!(Reducer::reduce(&MinReducer, 5, 3).unwrap(), 3); - assert_eq!(Reducer::reduce(&MinReducer, 2, 9).unwrap(), 2); - assert_eq!(Reducer::reduce(&MaxReducer, 5, 3).unwrap(), 5); - assert_eq!(Reducer::reduce(&MaxReducer, 2, 9).unwrap(), 9); -} - -#[test] -fn closure_reducer_runs() { - let r = ClosureReducer::new(|a: i32, b: i32| Ok(a + b)); - assert_eq!(r.reduce(2, 3).unwrap(), 5); -} - -#[test] -fn closure_state_reducer_applies_partial() { - #[derive(Clone, PartialEq, Debug)] - struct S { - total: i32, - } - let r = ClosureStateReducer::new(|mut s: S, u: i32| { - s.total += u; - Ok(s) - }); - let s = r.apply(S { total: 1 }, 4).unwrap(); - assert_eq!(s.total, 5); -} - -#[test] -fn overwrite_state_reducer() { - let r = OverwriteStateReducer; - assert_eq!(StateReducer::apply(&r, 1, 2).unwrap(), 2); -} diff --git a/crates/tinyagents-graph/src/reducer/types.rs b/crates/tinyagents-graph/src/reducer/types.rs deleted file mode 100644 index 9971e3f7d..000000000 --- a/crates/tinyagents-graph/src/reducer/types.rs +++ /dev/null @@ -1,65 +0,0 @@ -//! Reducer traits and built-in reducer markers. -//! -//! Reducers define how concurrent or sequential writes merge into channel -//! values at superstep boundaries. Two traits are provided: -//! -//! - [`Reducer`]: merges two values of the same channel type. Used for -//! channel-style state where each key has its own merge policy. -//! - [`StateReducer`]: merges a partial `Update` into the whole -//! `State`. This is the milestone-1 contract used by the executor: whole-state -//! updates are acceptable, while the partial-update path enables typed, -//! channel-like merging without rewriting node code. - -use std::marker::PhantomData; - -use crate::Result; - -/// Merges two values of the same channel type. -pub trait Reducer: Send + Sync { - /// Merges `update` into `current`, producing the new channel value. - fn reduce(&self, current: T, update: T) -> Result; -} - -/// Merges a partial `Update` into the whole graph `State`. -pub trait StateReducer: Send + Sync { - /// Applies `update` to `state`, producing the new state. - fn apply(&self, state: State, update: Update) -> Result; -} - -/// Overwrites the current value with the update (last-value semantics). -#[derive(Clone, Copy, Debug, Default)] -pub struct OverwriteReducer; - -/// Appends the update vector onto the current vector. -#[derive(Clone, Copy, Debug, Default)] -pub struct AppendReducer; - -/// Unions the update vector into the current vector, skipping duplicates and -/// preserving first-seen order. -#[derive(Clone, Copy, Debug, Default)] -pub struct SetUnionReducer; - -/// Keeps the smaller of the two values. -#[derive(Clone, Copy, Debug, Default)] -pub struct MinReducer; - -/// Keeps the larger of the two values. -#[derive(Clone, Copy, Debug, Default)] -pub struct MaxReducer; - -/// A custom binary [`Reducer`] backed by a closure. -pub struct ClosureReducer { - pub(crate) f: F, - pub(crate) _marker: PhantomData T>, -} - -/// A [`StateReducer`] that overwrites whole state with the update (the default -/// for `State == Update` graphs). -#[derive(Clone, Copy, Debug, Default)] -pub struct OverwriteStateReducer; - -/// A custom [`StateReducer`] backed by a closure. -pub struct ClosureStateReducer { - pub(crate) f: F, - pub(crate) _marker: PhantomData State>, -} diff --git a/crates/tinyagents-graph/src/status/README.md b/crates/tinyagents-graph/src/status/README.md deleted file mode 100644 index dbb085823..000000000 --- a/crates/tinyagents-graph/src/status/README.md +++ /dev/null @@ -1,46 +0,0 @@ -# graph::status - -`GraphRunStatus` — a compact, cheaply-pollable summary of a graph run, -distinct from a checkpoint. - -Where a `checkpoint::Checkpoint` preserves the *resumable state* of a run, a -`GraphRunStatus` is a lightweight record an observer can read to answer "is -this run active?", "which node is executing?", "which interrupt is waiting?" -— without deserializing full graph state. Because each status carries -`root_run_id` / `parent_run_id`, these records form the same run-tree shape -as `recursion::RunTree`, connecting a parent run to every subgraph/sub-agent -run it spawns, so an observer can roll up progress, interrupts, and errors -across levels of recursion. - -## Public surface - -- `GraphRunStatus` — the record itself: run/root/parent run ids, thread id, - graph id, latest checkpoint id and namespace, coarse `ExecutionStatus`, - current superstep, active nodes, pending interrupts, last event id, - timestamps (`started_at`/`updated_at`/`ended_at`), and a rendered error - summary for failed runs. - - `GraphRunStatus::new(run_id, graph_id, status)` — a fresh status for a - top-level run with no recorded progress (`root_run_id` mirrors `run_id`, - zeroed step/nodes/interrupts). - - `GraphRunStatus::is_terminal()` — true for `Completed`, `Failed`, or - `Cancelled`. - -## Files - -| File | Role | -| --- | --- | -| `types.rs` | The `GraphRunStatus` struct definition. | -| `mod.rs` | `GraphRunStatus::new` and `::is_terminal`. | -| `mod_tests.rs` | Unit tests for the constructor's defaults and terminal-state detection. | - -## How it fits together - -- `compiled::types::GraphExecution` carries a `status: GraphRunStatus` - snapshot at the final boundary of a run. -- `observability::GraphStatusStore` (opt-in via - `CompiledGraph::with_status_store`) persists `GraphRunStatus` updates as a - run progresses, so an external poller doesn't need to replay the event - journal to answer "is this run still running?". -- This module has no dependency on `checkpoint`, `stream`, or `observability` - — it defines only the record shape; the executor and observability layer - are what populate and persist it. diff --git a/crates/tinyagents-graph/src/status/mod.rs b/crates/tinyagents-graph/src/status/mod.rs deleted file mode 100644 index 369122c9e..000000000 --- a/crates/tinyagents-graph/src/status/mod.rs +++ /dev/null @@ -1,62 +0,0 @@ -//! Graph run status snapshots — the lightweight observability surface for the -//! tree of recursive runs. -//! -//! Where checkpoints persist *resumable state*, a [`GraphRunStatus`] is a -//! cheap, readable summary an observer can poll to answer "is this run active?", -//! "which node is executing?", and "which interrupt is waiting?". Because each -//! status carries `root_run_id` / `parent_run_id`, the records form the run -//! tree that links a parent graph to the subgraphs and sub-agents it recurses -//! into, so progress, interrupts, and errors can be rolled up across levels. -//! -//! See `types` for the [`GraphRunStatus`] definition. - -mod types; - -pub use types::GraphRunStatus; - -use std::time::SystemTime; - -use tinyagents_harness::ids::{ExecutionStatus, GraphId, RunId}; - -impl GraphRunStatus { - /// Creates a fresh status for a top-level run with no recorded progress yet. - pub fn new(run_id: RunId, graph_id: GraphId, status: ExecutionStatus) -> Self { - let now = SystemTime::now(); - Self { - root_run_id: run_id.clone(), - run_id, - parent_run_id: None, - thread_id: None, - graph_id, - checkpoint_id: None, - checkpoint_namespace: Vec::new(), - status, - current_step: 0, - active_nodes: Vec::new(), - pending_interrupts: Vec::new(), - last_event_id: None, - started_at: now, - updated_at: now, - ended_at: None, - error: None, - } - } - - /// 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::Drained - ) - } -} - -#[cfg(test)] -#[path = "mod_tests.rs"] -mod test; diff --git a/crates/tinyagents-graph/src/status/mod_tests.rs b/crates/tinyagents-graph/src/status/mod_tests.rs deleted file mode 100644 index 171298cb1..000000000 --- a/crates/tinyagents-graph/src/status/mod_tests.rs +++ /dev/null @@ -1,31 +0,0 @@ -//! Unit tests for [`GraphRunStatus`]: the constructor's defaults (root run id -//! mirrors the run id, zeroed progress) and terminal-state detection across the -//! lifecycle. - -use super::*; -use tinyagents_harness::ids::{ExecutionStatus, GraphId, RunId}; - -#[test] -fn new_status_defaults() { - let s = GraphRunStatus::new( - RunId::from("r1"), - GraphId::from("g1"), - ExecutionStatus::Pending, - ); - assert_eq!(s.run_id, s.root_run_id); - assert_eq!(s.current_step, 0); - assert!(s.active_nodes.is_empty()); - assert!(!s.is_terminal()); -} - -#[test] -fn terminal_detection() { - let mut s = GraphRunStatus::new( - RunId::from("r1"), - GraphId::from("g1"), - ExecutionStatus::Running, - ); - assert!(!s.is_terminal()); - s.status = ExecutionStatus::Completed; - assert!(s.is_terminal()); -} diff --git a/crates/tinyagents-graph/src/status/types.rs b/crates/tinyagents-graph/src/status/types.rs deleted file mode 100644 index 8262e831c..000000000 --- a/crates/tinyagents-graph/src/status/types.rs +++ /dev/null @@ -1,55 +0,0 @@ -//! Compact, observable graph run status records. -//! -//! These records are the nodes of the recursive run tree: `root_run_id` and -//! `parent_run_id` connect a top-level run to every subgraph/sub-agent run it -//! spawns, while the remaining fields capture live execution (current step, -//! active nodes, pending interrupts) without deserializing full graph state. - -use std::time::SystemTime; - -use tinyagents_harness::ids::{ - CheckpointId, EventId, ExecutionStatus, GraphId, InterruptId, NodeId, RunId, ThreadId, -}; - -/// A compact, readable summary of a graph run at an execution boundary. -/// -/// Status records are **not** checkpoints. A [`crate::Checkpoint`] -/// preserves the resumable state of a run; a `GraphRunStatus` summarizes live -/// and recent execution so observers (UIs, supervisors, tests) can answer "is -/// this run active?", "which node is executing?", and "which interrupt is -/// waiting?" without deserializing full graph state. -#[derive(Clone, Debug)] -pub struct GraphRunStatus { - /// This run's id. - pub run_id: RunId, - /// The root run id of the run tree (equal to `run_id` for top-level runs). - pub root_run_id: RunId, - /// The parent run id when this run is a subgraph/sub-agent child. - pub parent_run_id: Option, - /// The thread id when checkpointing is enabled. - pub thread_id: Option, - /// The graph this run belongs to. - pub graph_id: GraphId, - /// The latest checkpoint id, if any. - pub checkpoint_id: Option, - /// The checkpoint namespace (for nested subgraph runs). - pub checkpoint_namespace: Vec, - /// Coarse lifecycle status. - pub status: ExecutionStatus, - /// The current (or final) superstep number. - pub current_step: usize, - /// Nodes active at this boundary. - pub active_nodes: Vec, - /// Interrupts awaiting a resume command. - pub pending_interrupts: Vec, - /// The last emitted event id, if event ids are tracked. - pub last_event_id: Option, - /// When the run started. - pub started_at: SystemTime, - /// When this status was last updated. - pub updated_at: SystemTime, - /// When the run reached a terminal state, if it has. - pub ended_at: Option, - /// A rendered error summary for failed runs. - pub error: Option, -} diff --git a/crates/tinyagents-graph/src/storage/import.rs b/crates/tinyagents-graph/src/storage/import.rs deleted file mode 100644 index 37ad24940..000000000 --- a/crates/tinyagents-graph/src/storage/import.rs +++ /dev/null @@ -1,213 +0,0 @@ -//! Moving v1 data into `tsd`-owned collections. -//! -//! A document step may only write its own owner's collections, and the graph's -//! leases and task cache now live in `tsd.leases` and `tsd.kv`, which `tsd` -//! owns. So the move is an importer the host calls once per scope after -//! `Registry::migrate`, in the style of a legacy-file importer. Each records -//! completion in `tsd.migrations` (`Migrations::done(owner, step)`), so a -//! second call does nothing. - -use serde_json::Value; -use tinyagents_harness::ids::now_ms; -use tinyagents_storage_kit::{ - CollectionSpec, Cursor, ErrorKind, Filter, ID_FIELD, Key, Leases, MIGRATIONS_COLLECTION, - Migrations, Precondition, Query, Result, ScopedStorage, Sort, WriteOptions, primitives, -}; - -use super::steps::LEGACY; -use super::{KV_OWNER, LEASE_OWNER, task_cache_ns}; -use crate::storage::types::lease_key; - -/// Rows read per page while importing. -const PAGE: usize = 200; - -/// Every row of the legacy collection `collection`, oldest id first. -async fn legacy_rows( - storage: &ScopedStorage, - collection: &str, -) -> Result>> { - let docs = storage.documents(); - docs.ensure_collection(&CollectionSpec::new(collection)) - .await?; - let mut rows = Vec::new(); - let mut cursor: Option = None; - loop { - let mut query = Query::filter(Filter::default()) - .sort(Sort::asc(ID_FIELD)) - .limit(PAGE); - if let Some(after) = cursor.take() { - query = query.after(after); - } - let page = docs.query(collection, &query).await?; - rows.extend(page.items); - match page.next { - Some(next) => cursor = Some(next), - None => return Ok(rows), - } - } -} - -/// The thread a v1 lease id belongs to: ids were `len:thread`. -fn v1_lease_thread(id: &str) -> Option<&str> { - let (len, thread) = id.split_once(':')?; - (len.parse::().ok()? == thread.len()).then_some(thread) -} - -/// Carry the live execution leases of the v1 checkpointers -/// (`_leases`: `{owner, expires_at_ms}`) into `tsd.leases` for the -/// scope `storage`, keeping the holder and the time left, then remove the -/// legacy collections. A lease that has already lapsed is dropped. Returns how -/// many leases were carried over. -/// -/// A completed import is skipped; a missing collection counts as empty. -/// -/// # Errors -/// -/// The first storage error; the import resumes on the next call because the -/// legacy collections are removed only after every lease was carried over. -pub async fn import_legacy_leases(storage: &ScopedStorage) -> Result { - tinyagents_storage_kit::with_migration_lease(storage, "graph", |lease| async move { - import_legacy_leases_with_lease(storage, &lease).await - }) - .await -} - -pub(crate) async fn import_legacy_leases_with_lease( - storage: &ScopedStorage, - lease: &tinyagents_storage_kit::MigrationLease, -) -> Result { - storage - .documents() - .ensure_collection(&CollectionSpec::new(MIGRATIONS_COLLECTION)) - .await?; - let migrations = Migrations::new(storage); - if migrations.done(LEASE_OWNER, "lease_import").await? { - tracing::debug!("[graph-storage] lease import already done"); - return Ok(0); - } - storage - .documents() - .ensure_collection(&primitives::leases_collection_spec()) - .await?; - let leases = Leases::new(storage.documents()); - let mut carried = 0_u64; - for legacy in &LEGACY { - for row in legacy_rows(storage, legacy.leases).await? { - let (Some(holder), Some(expires)) = ( - row.doc.get("owner").and_then(Value::as_str), - row.doc.get("expires_at_ms").and_then(Value::as_i64), - ) else { - continue; - }; - let Some(thread) = v1_lease_thread(&row.id) else { - continue; - }; - let now = now_ms() as i64; - if expires <= now { - continue; - } - lease.renew().await?; - let held = leases - .acquire( - LEASE_OWNER, - &lease_key(legacy.ns, thread), - holder, - expires - now, - ) - .await?; - carried += u64::from(held.is_some()); - } - } - for legacy in &LEGACY { - lease.renew().await?; - storage - .documents() - .delete_where(legacy.leases, &Filter::default()) - .await?; - } - lease.mark_done(LEASE_OWNER, "lease_import").await?; - tracing::debug!(carried, "[graph-storage] lease import complete"); - Ok(carried) -} - -/// Carry the v1 node-result cache (`graph_task_cache`: -/// `{graph, node, hash, value, expiresAtMs}`) into `tsd.kv` under owner -/// `graph` for the scope `storage`, moving the deadline into the entry's -/// `expires_ms` metadata, then remove the legacy collection. Expired rows are -/// dropped. Returns how many entries were written. -/// -/// A completed import is skipped; an entry already present in `tsd.kv` is kept. -/// -/// # Errors -/// -/// The first storage error; the import resumes on the next call. -pub async fn import_legacy_task_cache(storage: &ScopedStorage) -> Result { - tinyagents_storage_kit::with_migration_lease(storage, "graph", |lease| async move { - import_legacy_task_cache_with_lease(storage, &lease).await - }) - .await -} - -pub(crate) async fn import_legacy_task_cache_with_lease( - storage: &ScopedStorage, - lease: &tinyagents_storage_kit::MigrationLease, -) -> Result { - const LEGACY_COLLECTION: &str = "graph_task_cache"; - const STEP: &str = "kv_import_graph_task_cache"; - storage - .documents() - .ensure_collection(&CollectionSpec::new(MIGRATIONS_COLLECTION)) - .await?; - let migrations = Migrations::new(storage); - if migrations.done(KV_OWNER, STEP).await? { - tracing::debug!("[graph-storage] task cache import already done"); - return Ok(0); - } - storage - .documents() - .ensure_collection(&primitives::kv_collection_spec()) - .await?; - let now = now_ms() as i64; - let mut written = 0_u64; - for row in legacy_rows(storage, LEGACY_COLLECTION).await? { - let (Some(graph), Some(node), Some(hash), Some(value)) = ( - row.doc.get("graph").and_then(Value::as_str), - row.doc.get("node").and_then(Value::as_str), - row.doc.get("hash").and_then(Value::as_str), - row.doc.get("value"), - ) else { - continue; - }; - let expires = row.doc.get("expiresAtMs").and_then(Value::as_i64); - if expires.is_some_and(|at| at <= now) { - continue; - } - let ns = task_cache_ns(super::DEFAULT_CACHE_NAMESPACE, graph); - let key = Key::of([node, hash])?.as_str().to_owned(); - let id = Key::of([KV_OWNER, &ns, &key])?.as_str().to_owned(); - let mut options = WriteOptions::new(); - if let Some(at) = expires { - options = options.expires_ms(at); - } - let body = serde_json::json!({"owner": KV_OWNER, "ns": ns, "key": key, "value": value}); - lease.renew().await?; - match storage - .documents() - .put_with(primitives::KV, &id, body, Precondition::Absent, options) - .await - { - Ok(_) => written += 1, - Err(error) - if matches!(error.kind(), ErrorKind::AlreadyExists | ErrorKind::Conflict) => {} - Err(error) => return Err(error), - } - } - lease.renew().await?; - storage - .documents() - .delete_where(LEGACY_COLLECTION, &Filter::default()) - .await?; - lease.mark_done(KV_OWNER, STEP).await?; - tracing::debug!(written, "[graph-storage] task cache import complete"); - Ok(written) -} diff --git a/crates/tinyagents-graph/src/storage/mod.rs b/crates/tinyagents-graph/src/storage/mod.rs deleted file mode 100644 index b910eeeb9..000000000 --- a/crates/tinyagents-graph/src/storage/mod.rs +++ /dev/null @@ -1,155 +0,0 @@ -//! The graph crate's v2 storage declarations: its `storage_schema()`, the -//! stored bodies of the one graph checkpointer and the steps and importers -//! that bring the v1 collections forward. - -mod import; -mod steps; -mod types; - -use tinyagents_storage_kit::{ - CollectionSpec, IndexSpec, MigrateOptions, Owner, Registry, Schema, ScopeClass, ScopedStorage, - primitives, -}; - -pub use import::{import_legacy_leases, import_legacy_task_cache}; -pub(crate) use import::{import_legacy_leases_with_lease, import_legacy_task_cache_with_lease}; -pub(crate) use types::{ - CheckpointBody, ThreadBody, WritesBody, checkpoint_id_key, lease_key, subgraph_key, - thread_id_key, writes_id_key, -}; - -/// The body version of every `graph.*` collection. -/// -/// A collection takes one v1 -> v2 step per legacy prefix (see -/// [`storage_schema`]) and a step chain cannot fork, so the chain runs 1 -> 2 -/// -> 3 -> 4 and the declared version is the end of it. The body shape itself -/// has not changed since the first step. -pub(crate) const SCHEMA_VERSION: u32 = 4; - -/// Stored checkpoints, one document per checkpoint. -pub(crate) const CHECKPOINTS: &str = "graph.checkpoints"; - -/// The per-thread insertion counters. -pub(crate) const THREADS: &str = "graph.threads"; - -/// The merged pending writes of a checkpoint. -pub(crate) const WRITES: &str = "graph.writes"; - -/// The namespace [`DriverCheckpointer::new`](crate::checkpoint::DriverCheckpointer::new) uses. -pub(crate) const DEFAULT_NAMESPACE: &str = "graph"; - -/// The `tsd.leases` owner of thread execution leases. -pub(crate) const LEASE_OWNER: &str = "graph"; - -/// The `tsd.kv` owner of the node-result cache. -pub(crate) const KV_OWNER: &str = "graph"; - -/// The `tsd.kv` namespace of the node-result cache when none is given. -pub(crate) const DEFAULT_CACHE_NAMESPACE: &str = "task_cache"; - -/// The `tsd.kv` namespace of `graph`'s node results in the cache `base`. -pub(crate) fn task_cache_ns(base: &str, graph: &str) -> String { - format!("{base}:{graph}") -} - -/// `graph.checkpoints`: every read filters on the namespace and thread. -pub(crate) fn checkpoints_spec() -> CollectionSpec { - CollectionSpec::new(CHECKPOINTS) - .index(IndexSpec::new("by_thread_seq", ["ns", "thread", "seq"])) - .index(IndexSpec::new( - "by_thread_id", - ["ns", "thread", "checkpoint_id"], - )) - .index(IndexSpec::new( - "by_thread_subgraph", - ["ns", "thread", "subgraph", "seq"], - )) -} - -/// `graph.threads`. -pub(crate) fn threads_spec() -> CollectionSpec { - CollectionSpec::new(THREADS).index(IndexSpec::new("by_ns", ["ns"])) -} - -/// `graph.writes`. -pub(crate) fn writes_spec() -> CollectionSpec { - CollectionSpec::new(WRITES).index(IndexSpec::new("by_thread", ["ns", "thread"])) -} - -/// `tsd.leases`, which the checkpointer claims through `Leases`. -pub(crate) fn leases_spec() -> CollectionSpec { - CollectionSpec::new(primitives::LEASES) -} - -/// The graph crate's storage declaration for its checkpointer and the steps -/// that bring the v1 names and shapes forward. [`open_storage`] runs this -/// owner-local migration against the already-bound scope; the checkpointer's -/// `DriverCheckpointer::open` calls it before returning the handle. -/// -/// | v1 | v2 | -/// |---|---| -/// | `graph_{checkpoints,threads,writes}` | `graph.{checkpoints,threads,writes}`, `ns` = `graph` | -/// | `delegation_graph_{checkpoints,threads,writes}` | the same collections, `ns` = `delegation` | -/// | `flows_graph_{checkpoints,threads,writes}` | the same collections, `ns` = `flows` | -/// | `_leases` | `tsd.leases`, owner `graph` | -/// | `graph_task_cache` | `tsd.kv`, owner `graph` | -/// -/// The old prefix becomes the `ns` body field and is part of every id and -/// index, so two namespaces never collide. The checkpoint record stays a typed -/// field (`checkpoint`), not a JSON string. A collection gets one step per -/// prefix, so the steps chain through versions 1 to 4. -/// -/// `tsd.leases` and `tsd.kv` belong to `tsd`, so a document step cannot write -/// them: [`open_storage`] calls [`import_legacy_leases`] and -/// [`import_legacy_task_cache`] after the steps. The goal and todo lists use the -/// harness `Store`, which lives in `tsd.kv` under the store's own owner. -#[must_use] -pub fn storage_schema() -> Schema { - let owner = Owner::new("graph").expect("`graph` is a valid owner"); - let mut schema = Schema::new(owner) - .collection_with(SCHEMA_VERSION, checkpoints_spec(), ScopeClass::Agent) - .collection_with(SCHEMA_VERSION, threads_spec(), ScopeClass::Agent) - .collection_with(SCHEMA_VERSION, writes_spec(), ScopeClass::Agent); - for step in steps::all() { - schema = schema.step(step); - } - schema -} - -/// Collections owned by graph; callers declare these on their chosen backend. -#[must_use] -pub fn collections() -> Vec { - storage_schema().collections() -} - -/// Open graph storage for one scope, applying graph-owned migrations and -/// importing the graph's legacy cache and lease records before use. After the -/// first success for a scope, later calls in this process skip the work. -pub async fn open_storage( - storage: ScopedStorage, -) -> tinyagents_storage_kit::error::Result { - let owned_storage = storage.clone(); - tinyagents_storage_kit::with_migration_lease_once(&storage, "graph", move |lease| async move { - let storage = owned_storage; - let registry = Registry::compose([storage_schema()])?; - for spec in collections() { - storage.documents().ensure_collection(&spec).await?; - } - for spec in [ - primitives::leases_collection_spec(), - primitives::kv_collection_spec(), - ] { - storage.documents().ensure_collection(&spec).await?; - } - registry.migrate(&storage, &MigrateOptions::new()).await?; - import_legacy_leases_with_lease(&storage, &lease).await?; - import_legacy_task_cache_with_lease(&storage, &lease).await?; - Ok(()) - }) - .await?; - Ok(storage) -} - -#[cfg(test)] -#[path = "mod_tests.rs"] -mod tests; diff --git a/crates/tinyagents-graph/src/storage/mod_tests.rs b/crates/tinyagents-graph/src/storage/mod_tests.rs deleted file mode 100644 index 4ea6b0873..000000000 --- a/crates/tinyagents-graph/src/storage/mod_tests.rs +++ /dev/null @@ -1,356 +0,0 @@ -use super::*; -use crate::cache::{DriverTaskCache, TaskCache, TaskCacheKey}; -use crate::checkpoint::{Checkpoint, Checkpointer, DriverCheckpointer}; -use serde_json::{Value, json}; -use std::sync::Arc; -use std::time::Duration; -use tinyagents_harness::ids::{GraphId, NodeId, now_ms}; -use tinyagents_storage_kit::{ - DocumentStoreExt, MigrateOptions, Precondition, Query, Registry, Scope, ScopeClass, - ScopedStorage, StorageBackend, -}; -use tinystoragedrivers_testkit::MemoryBackend; - -fn scoped() -> ScopedStorage { - MemoryBackend::new() - .for_scope(&Scope::new("local").unwrap()) - .unwrap() -} - -async fn put_v1(storage: &ScopedStorage, collection: &str, id: &str, body: Value) { - storage - .documents() - .put(collection, id, body, Precondition::None) - .await - .unwrap(); -} - -async fn migrate(storage: &ScopedStorage) { - let registry = Registry::compose([storage_schema()]).unwrap(); - for item in storage_schema().items { - if let tinyagents_storage_kit::Item::Collection { spec, .. } = item { - storage.documents().ensure_collection(&spec).await.unwrap(); - } - } - for spec in [ - tinyagents_storage_kit::primitives::leases_collection_spec(), - tinyagents_storage_kit::primitives::kv_collection_spec(), - ] { - storage.documents().ensure_collection(&spec).await.unwrap(); - } - registry - .migrate(storage, &MigrateOptions::new()) - .await - .unwrap(); -} - -/// A v1 checkpoint row as `DriverCheckpointer` wrote it under `_checkpoints`. -async fn v1_checkpoint( - storage: &ScopedStorage, - prefix: &str, - thread: &str, - seq: u64, - id: &str, - state: i32, -) { - let checkpoint = Checkpoint::new(state, Vec::new()) - .with_thread_id(thread.to_string()) - .with_checkpoint_id(id.to_string()); - put_v1( - storage, - &format!("{prefix}_checkpoints"), - &format!("{}:{thread}/20:{seq:020}", thread.len()), - json!({ - "thread": thread, "namespace": "", "seq": seq, "checkpoint_id": id, - "record": serde_json::to_value(&checkpoint).unwrap(), - }), - ) - .await; - put_v1( - storage, - &format!("{prefix}_threads"), - &format!("{}:{thread}", thread.len()), - json!({"thread": thread, "next": seq + 1}), - ) - .await; -} - -fn checkpointer(storage: &ScopedStorage, ns: &str) -> DriverCheckpointer { - DriverCheckpointer::with_namespace(Arc::clone(storage.documents()), ns) -} - -#[test] -fn the_schema_composes_and_names_the_one_checkpointer() { - let schema = storage_schema(); - assert_eq!(schema.owner.as_str(), "graph"); - let names: Vec<&str> = schema.items.iter().map(|item| item.name()).collect(); - assert_eq!( - names, - ["graph.checkpoints", "graph.threads", "graph.writes"] - ); - assert!( - schema - .items - .iter() - .all(|item| item.scope() == ScopeClass::Agent) - ); - Registry::compose([schema]).unwrap(); -} - -#[tokio::test] -async fn each_v1_prefix_becomes_its_own_namespace_of_the_same_collections() { - let storage = scoped(); - // The same thread id under all three prefixes must not collide. - v1_checkpoint(&storage, "graph", "t", 0, "g1", 1).await; - v1_checkpoint(&storage, "delegation_graph", "t", 0, "d1", 2).await; - v1_checkpoint(&storage, "flows_graph", "t", 0, "f1", 3).await; - migrate(&storage).await; - - for (ns, id, state) in [ - ("graph", "g1", 1), - ("delegation", "d1", 2), - ("flows", "f1", 3), - ] { - let cp = checkpointer(&storage, ns); - let found = cp.get("t", None).await.unwrap().expect(ns); - assert_eq!( - (found.checkpoint_id.as_str(), found.state), - (id, state), - "{ns}" - ); - assert_eq!(cp.list_threads().await.unwrap(), ["t"], "{ns}"); - // The counter survived: the next checkpoint continues after the migrated one. - let next = Checkpoint::new(9, Vec::new()) - .with_thread_id("t".to_string()) - .with_checkpoint_id("next".to_string()); - cp.put(next).await.unwrap(); - let history = cp.get_thread("t").await.unwrap(); - assert_eq!(history.len(), 2, "{ns}"); - assert_eq!(history[1].checkpoint_id, "next", "{ns}"); - } - - let row = storage - .documents() - .query_all("graph.checkpoints", &Query::all()) - .await - .unwrap() - .into_iter() - .find(|row| row.doc["ns"] == "delegation" && row.doc["checkpoint_id"] == "d1") - .unwrap(); - assert_eq!(row.doc["thread"], "t"); - assert_eq!(row.doc["seq"], 0); - assert!( - row.doc["checkpoint"].is_object(), - "typed, not a JSON string" - ); - assert!(row.doc.get("record").is_none()); - // A row carries the version its step ends at: one step per prefix. - assert_eq!(row.meta.schema, 3); -} - -#[tokio::test] -async fn v1_pending_writes_follow_their_checkpoint_into_the_namespace() { - let storage = scoped(); - put_v1( - &storage, - "delegation_graph_writes", - "3:t/4:::c1", - json!({"thread": "t", "namespace": ["sub"], "checkpoint_id": "c1", - "writes": [{"node": "n", "task_id": "task", "idx": 0, "channel": "out", "payload": "x"}]}), - ) - .await; - migrate(&storage).await; - let cp = checkpointer(&storage, "delegation"); - let config = crate::checkpoint::CheckpointConfig { - thread_id: "t".to_string(), - checkpoint_id: Some("c1".to_string()), - namespace: vec!["sub".to_string()], - }; - let writes = cp.get_writes(&config).await.unwrap(); - assert_eq!(writes.len(), 1); - assert_eq!(writes[0].payload, json!("x")); - assert!( - checkpointer(&storage, "graph") - .get_writes(&config) - .await - .unwrap() - .is_empty() - ); -} - -#[tokio::test] -async fn migrating_twice_changes_nothing() { - let storage = scoped(); - v1_checkpoint(&storage, "graph", "t", 0, "c1", 1).await; - migrate(&storage).await; - migrate(&storage).await; - let rows = storage - .documents() - .query_all("graph.checkpoints", &Query::all()) - .await - .unwrap(); - assert_eq!(rows.len(), 1); -} - -#[tokio::test] -async fn a_live_v1_lease_is_carried_over_and_an_expired_one_dropped() { - let storage = scoped(); - let now = now_ms(); - put_v1( - &storage, - "delegation_graph_leases", - "live", - json!({"owner": "worker-a", "expires_at_ms": now + 60_000}), - ) - .await; - // Ids are the v1 `key([thread])`: length-prefixed. - put_v1( - &storage, - "delegation_graph_leases", - "1:a", - json!({"owner": "worker-a", "expires_at_ms": now + 60_000}), - ) - .await; - put_v1( - &storage, - "flows_graph_leases", - "1:b", - json!({"owner": "worker-b", "expires_at_ms": now.saturating_sub(1)}), - ) - .await; - assert_eq!(import_legacy_leases(&storage).await.unwrap(), 1); - - let delegation = checkpointer(&storage, "delegation"); - assert!( - !delegation - .try_claim("a", "other", Duration::from_secs(5)) - .await - .unwrap() - ); - assert!( - delegation - .renew("a", "worker-a", Duration::from_secs(5)) - .await - .unwrap() - ); - let flows = checkpointer(&storage, "flows"); - assert!( - flows - .try_claim("b", "other", Duration::from_secs(5)) - .await - .unwrap() - ); - // The legacy collections are gone and the import is recorded. - assert!( - storage - .documents() - .query_all("delegation_graph_leases", &Query::all()) - .await - .unwrap() - .is_empty() - ); - assert_eq!(import_legacy_leases(&storage).await.unwrap(), 0); -} - -#[tokio::test] -async fn v1_task_cache_rows_move_into_the_graph_kv_namespace() { - let storage = scoped(); - let now = now_ms(); - put_v1( - &storage, - "graph_task_cache", - "h1", - json!({"graph": "g", "node": "n", "hash": "abc", "value": {"v": 1}, "expiresAtMs": null}), - ) - .await; - put_v1( - &storage, - "graph_task_cache", - "h2", - json!({"graph": "g", "node": "n", "hash": "old", "value": 2, "expiresAtMs": now.saturating_sub(5)}), - ) - .await; - put_v1( - &storage, - "graph_task_cache", - "h3", - json!({"graph": "g", "node": "n", "hash": "soon", "value": 3, "expiresAtMs": now + 60_000}), - ) - .await; - assert_eq!(import_legacy_task_cache(&storage).await.unwrap(), 2); - let cache = DriverTaskCache::new(Arc::clone(storage.documents())); - let key = |hash: &str| TaskCacheKey::new(GraphId::new("g"), NodeId::from("n"), hash); - assert_eq!(cache.get(&key("abc")).await.unwrap(), Some(json!({"v": 1}))); - assert_eq!(cache.get(&key("soon")).await.unwrap(), Some(json!(3))); - assert_eq!(cache.get(&key("old")).await.unwrap(), None); - let soon = storage - .documents() - .query_all("tsd.kv", &Query::all()) - .await - .unwrap() - .into_iter() - .find(|row| row.doc["key"].as_str().is_some_and(|k| k.ends_with("soon"))) - .unwrap(); - assert_eq!(soon.doc["owner"], "graph"); - assert!( - soon.meta.expires_ms.is_some(), - "the deadline is metadata now" - ); - assert_eq!(import_legacy_task_cache(&storage).await.unwrap(), 0); -} - -#[tokio::test] -async fn open_runs_the_graph_migration_so_a_v1_checkpoint_is_visible() { - let storage = scoped(); - v1_checkpoint(&storage, "delegation_graph", "t", 0, "d1", 2).await; - // The sync constructor over an unmigrated scope does not see v1 rows. - assert!( - checkpointer(&storage, "delegation") - .get("t", None) - .await - .unwrap() - .is_none() - ); - - let cp = DriverCheckpointer::::open(storage.clone(), "delegation") - .await - .unwrap(); - let found = cp.get("t", None).await.unwrap().expect("migrated on open"); - assert_eq!((found.checkpoint_id.as_str(), found.state), ("d1", 2)); -} - -#[tokio::test] -async fn a_second_open_skips_the_migration_lease_and_changes_nothing() { - let storage = scoped(); - v1_checkpoint(&storage, "delegation_graph", "t", 0, "d1", 2).await; - DriverCheckpointer::::open(storage.clone(), "delegation") - .await - .unwrap(); - // Another process now holds the graph's migration lease: a second open - // that took it again would wait for that lease to expire. - tinyagents_storage_kit::primitives::Leases::new(storage.documents()) - .acquire_with(tinyagents_storage_kit::primitives::Acquire { - owner: "graph", - key: "owner_migrations", - holder: "someone-else", - ttl_ms: 600_000, - endpoint: None, - held_epoch: None, - }) - .await - .unwrap(); - let cp = tokio::time::timeout( - Duration::from_secs(5), - DriverCheckpointer::::open(storage.clone(), "delegation"), - ) - .await - .expect("the second open must not wait for the migration lease") - .unwrap(); - assert_eq!(cp.get_thread("t").await.unwrap().len(), 1); - let rows = storage - .documents() - .query_all("graph.checkpoints", &Query::all()) - .await - .unwrap(); - assert_eq!(rows.len(), 1); -} diff --git a/crates/tinyagents-graph/src/storage/steps.rs b/crates/tinyagents-graph/src/storage/steps.rs deleted file mode 100644 index da4bd2bb9..000000000 --- a/crates/tinyagents-graph/src/storage/steps.rs +++ /dev/null @@ -1,199 +0,0 @@ -//! The v1 -> v2 migration steps of [`super::storage_schema`]. Each is a pure -//! function of its input, so a resumed run produces the same output. - -use serde_json::Value; -use tinyagents_storage_kit::{SourceRecord, Step, TargetRecord}; - -use super::types::{ - CheckpointBody, ThreadBody, WritesBody, checkpoint_id_key, subgraph_key, thread_id_key, - writes_id_key, -}; -use super::{CHECKPOINTS, THREADS, WRITES}; -use crate::checkpoint::PendingWrite; - -/// One v1 family: the namespace it becomes and its three legacy collections. -pub(super) struct Legacy { - pub(super) ns: &'static str, - pub(super) checkpoints: &'static str, - pub(super) threads: &'static str, - pub(super) writes: &'static str, - pub(super) leases: &'static str, -} - -/// Every prefix `DriverCheckpointer` wrote under: its default (`graph`), the -/// delegation graph's (`delegation_graph`) and workflows' (`flows_graph`). -pub(super) const LEGACY: [Legacy; 3] = [ - Legacy { - ns: "graph", - checkpoints: "graph_checkpoints", - threads: "graph_threads", - writes: "graph_writes", - leases: "graph_leases", - }, - Legacy { - ns: "delegation", - checkpoints: "delegation_graph_checkpoints", - threads: "delegation_graph_threads", - writes: "delegation_graph_writes", - leases: "delegation_graph_leases", - }, - Legacy { - ns: "flows", - checkpoints: "flows_graph_checkpoints", - threads: "flows_graph_threads", - writes: "flows_graph_writes", - leases: "flows_graph_leases", - }, -]; - -/// The step ids per legacy family, in [`LEGACY`] order. -const CHECKPOINT_STEPS: [&str; 3] = [ - "checkpoints_graph_v2", - "checkpoints_delegation_v2", - "checkpoints_flows_v2", -]; -const THREAD_STEPS: [&str; 3] = [ - "threads_graph_v2", - "threads_delegation_v2", - "threads_flows_v2", -]; -const WRITE_STEPS: [&str; 3] = ["writes_graph_v2", "writes_delegation_v2", "writes_flows_v2"]; - -/// The steps of the three collections. A collection takes one step per legacy -/// prefix and a step chain cannot fork, so the chain runs -/// `1 -> 2 -> 3 -> 4`: the declared version counts the prefixes, not four -/// reshapes of the body. -pub(super) fn all() -> Vec { - let mut steps = Vec::new(); - for (index, legacy) in LEGACY.iter().enumerate() { - let from = 1 + u32::try_from(index).unwrap_or(0); - steps.push(checkpoints(CHECKPOINT_STEPS[index], from, legacy)); - } - for (index, legacy) in LEGACY.iter().enumerate() { - let from = 1 + u32::try_from(index).unwrap_or(0); - steps.push(threads(THREAD_STEPS[index], from, legacy)); - } - for (index, legacy) in LEGACY.iter().enumerate() { - let from = 1 + u32::try_from(index).unwrap_or(0); - steps.push(writes(WRITE_STEPS[index], from, legacy)); - } - steps -} - -/// A target that keeps the time the record was first written, when known. -fn target( - collection: &str, - id: String, - body: &impl serde::Serialize, - created_ms: i64, -) -> Option { - let record = TargetRecord::new(collection, id, serde_json::to_value(body).ok()?); - Some(if created_ms > 0 { - record.created_ms(created_ms) - } else { - record - }) -} - -fn str_field(body: &Value, field: &str) -> Option { - body.get(field)?.as_str().map(str::to_owned) -} - -/// The subgraph namespace of a v1 row: an array of segments. -fn namespace_of(value: Option<&Value>) -> Vec { - value - .and_then(Value::as_array) - .map(|parts| { - parts - .iter() - .filter_map(|part| part.as_str().map(str::to_owned)) - .collect() - }) - .unwrap_or_default() -} - -fn checkpoints(id: &'static str, from: u32, legacy: &'static Legacy) -> Step { - Step::new(id, CHECKPOINTS, from, from + 1, move |batch| { - Ok(batch - .iter() - .filter_map(|source| checkpoint_row(legacy.ns, source)) - .collect()) - }) - .source(legacy.checkpoints) -} - -fn checkpoint_row(ns: &str, source: &SourceRecord) -> Option { - let thread = str_field(&source.body, "thread")?; - let record = source.body.get("record")?.clone(); - if !record.is_object() { - return None; - } - let seq = source.body.get("seq")?.as_u64()?; - let body = CheckpointBody { - ns: ns.to_owned(), - subgraph: subgraph_key(&namespace_of(record.get("namespace"))), - checkpoint_id: str_field(&source.body, "checkpoint_id")?, - thread: thread.clone(), - seq, - checkpoint: record.clone(), - }; - let written = record - .get("created_at") - .and_then(Value::as_i64) - .filter(|ms| *ms > 0) - .unwrap_or(source.meta.created_ms); - target( - CHECKPOINTS, - checkpoint_id_key(ns, &thread, seq), - &body, - written, - ) -} - -fn threads(id: &'static str, from: u32, legacy: &'static Legacy) -> Step { - Step::new(id, THREADS, from, from + 1, move |batch| { - Ok(batch - .iter() - .filter_map(|source| { - let thread = str_field(&source.body, "thread")?; - let body = ThreadBody { - ns: legacy.ns.to_owned(), - next: source.body.get("next")?.as_u64()?, - thread: thread.clone(), - }; - target( - THREADS, - thread_id_key(legacy.ns, &thread), - &body, - source.meta.created_ms, - ) - }) - .collect()) - }) - .source(legacy.threads) -} - -fn writes(id: &'static str, from: u32, legacy: &'static Legacy) -> Step { - Step::new(id, WRITES, from, from + 1, move |batch| { - Ok(batch - .iter() - .filter_map(|source| { - let thread = str_field(&source.body, "thread")?; - let checkpoint_id = str_field(&source.body, "checkpoint_id")?; - let subgraph = subgraph_key(&namespace_of(source.body.get("namespace"))); - let writes: Vec = - serde_json::from_value(source.body.get("writes")?.clone()).ok()?; - let id = writes_id_key(legacy.ns, &thread, &subgraph, &checkpoint_id); - let body = WritesBody { - ns: legacy.ns.to_owned(), - thread, - subgraph, - checkpoint_id, - writes, - }; - target(WRITES, id, &body, source.meta.created_ms) - }) - .collect()) - }) - .source(legacy.writes) -} diff --git a/crates/tinyagents-graph/src/storage/types.rs b/crates/tinyagents-graph/src/storage/types.rs deleted file mode 100644 index 829292c71..000000000 --- a/crates/tinyagents-graph/src/storage/types.rs +++ /dev/null @@ -1,78 +0,0 @@ -//! The stored bodies of the one graph checkpointer (spec schema-v2 §5). - -use serde::{Deserialize, Serialize}; -use serde_json::Value; -use tinyagents_storage_kit::Key; - -use crate::checkpoint::PendingWrite; - -/// The body of one `graph.checkpoints` document. -/// -/// `checkpoint` is the opaque checkpoint payload, kept as a typed field rather -/// than a JSON string; everything else is an indexed projection. -#[derive(Debug, Clone, Serialize, Deserialize)] -pub(crate) struct CheckpointBody { - pub(crate) ns: String, - pub(crate) thread: String, - pub(crate) subgraph: String, - pub(crate) seq: u64, - pub(crate) checkpoint_id: String, - pub(crate) checkpoint: Value, -} - -/// The body of one `graph.threads` document: the next insertion sequence. -#[derive(Debug, Clone, Serialize, Deserialize)] -pub(crate) struct ThreadBody { - pub(crate) ns: String, - pub(crate) thread: String, - pub(crate) next: u64, -} - -/// The body of one `graph.writes` document. -#[derive(Debug, Clone, Serialize, Deserialize)] -pub(crate) struct WritesBody { - pub(crate) ns: String, - pub(crate) thread: String, - pub(crate) subgraph: String, - pub(crate) checkpoint_id: String, - pub(crate) writes: Vec, -} - -/// A subgraph namespace as one string, injectively: the root is empty and any -/// other namespace is `Key::of(["n", parts...])`, so `["a", "b"]`, `["a/b"]` -/// and `[""]` never meet. -pub(crate) fn subgraph_key(namespace: &[String]) -> String { - if namespace.is_empty() { - return String::new(); - } - Key::of(std::iter::once("n").chain(namespace.iter().map(String::as_str))) - .map(|key| key.as_str().to_owned()) - .unwrap_or_default() -} - -/// The id of the checkpoint inserted at `seq` for `thread` in `ns`. -pub(crate) fn checkpoint_id_key(ns: &str, thread: &str, seq: u64) -> String { - key(&[ns, thread, &format!("{seq:020}")]) -} - -/// The id of `thread`'s counter document in `ns`. -pub(crate) fn thread_id_key(ns: &str, thread: &str) -> String { - key(&[ns, thread]) -} - -/// The id of the pending writes of `checkpoint_id` in `subgraph`. -pub(crate) fn writes_id_key(ns: &str, thread: &str, subgraph: &str, checkpoint_id: &str) -> String { - key(&[ns, thread, subgraph, checkpoint_id]) -} - -/// The lease key of `thread` in `ns` (owner `graph`). -pub(crate) fn lease_key(ns: &str, thread: &str) -> String { - key(&[ns, thread]) -} - -/// `Key::of(parts)` as text; the parts are never empty, so it cannot fail. -fn key(parts: &[&str]) -> String { - Key::of(parts) - .map(|key| key.as_str().to_owned()) - .unwrap_or_default() -} diff --git a/crates/tinyagents-graph/src/stream/README.md b/crates/tinyagents-graph/src/stream/README.md deleted file mode 100644 index 2a186030c..000000000 --- a/crates/tinyagents-graph/src/stream/README.md +++ /dev/null @@ -1,64 +0,0 @@ -# graph::stream - -The live, in-process event surface: pluggable sinks and the wire vocabulary -(`GraphEvent`, `StreamMode`) the durable executor uses to narrate its own -execution as it happens. - -The executor emits a `GraphEvent` at every meaningful boundary — run -start/end, superstep start/end, node lifecycle, retries, routing decisions, -checkpoints, interrupts, subgraph enter/exit, forked branches, recursion depth -changes, and arbitrary custom writes. Routing those events into a -`GraphEventSink` is what lets a UI or an enclosing graph watch a -subgraph or sub-agent execute in real time; because every event carries its -node/step, nested runs' streams can be merged and attributed back up the run -tree. `observability` is what makes this same event vocabulary durable -(journaled, replayable) rather than merely live — see -`observability/README.md`. - -## Public surface - -- `GraphEventSink` (trait) — `emit(&self, event: GraphEvent)` receives one - event and must not block the executor; `flush(&self)` (default no-op) lets - an async-persisting sink let callers wait for the durable log to catch up - (the executor calls it after a terminal run event). -- `NoopSink` — drops every event. -- `CollectingSink` — records every event (in emission order) for inspection - in tests and UIs: `new()`, `events()`, `len()`, `is_empty()`. -- `GraphEvent` — the event enum itself (see `types.rs` for every variant); - `kind()` returns a stable dot-separated name (e.g. `"node.completed"`) for - logging/filtering; `step()` returns the associated superstep number when the - variant carries one. -- `StreamMode` — LangGraph-style selection of which projection of the stream a - caller wants: `Values`, `Updates`, `Messages`, `Debug`, `Interrupts`, - `Custom`. The milestone executor exposes this as a plain selection enum; - richer typed `StreamPart` projection is noted as future work. - -## Files - -| File | Role | -| --- | --- | -| `types.rs` | `GraphEvent` (every emitted variant), `StreamMode`. | -| `mod.rs` | `GraphEventSink` trait, `NoopSink`, `CollectingSink`. | -| `mod_tests.rs` | Unit tests for the two built-in sinks. | - -## How it fits together - -- `CompiledGraph::with_event_sink` attaches a `GraphEventSink`; the executor - (`compiled::executor`, `compiled::routing`) calls `self.emit(...)` - throughout the superstep loop. -- `GraphEvent` is `serde`-(de)serializable so a single event can be wrapped - into a durable `observability::GraphObservation` envelope, journaled, and - replayed — the live stream and the durable journal share this same event - shape rather than maintaining parallel vocabularies. -- `GraphEvent::step()` is what stamps the `step` field of a - `GraphObservation` when journaling. - -## Operational constraints - -- A `GraphEventSink::emit` implementation must not block the executor thread - — a sink that persists asynchronously should hand off work (e.g. to a - channel or background task) and implement `flush` to let callers wait for - it to settle, rather than doing the write inline. -- `CollectingSink` is unbounded and in-memory; it is meant for tests and - short-lived UIs, not as a durability mechanism for long-running production - runs. diff --git a/crates/tinyagents-graph/src/stream/mod.rs b/crates/tinyagents-graph/src/stream/mod.rs deleted file mode 100644 index fa93bc449..000000000 --- a/crates/tinyagents-graph/src/stream/mod.rs +++ /dev/null @@ -1,106 +0,0 @@ -//! Pluggable event sinks for graph streaming and observability — how a -//! recursive run makes its progress visible to the outside (and to a parent -//! run). -//! -//! The durable executor emits a stream of [`GraphEvent`]s as it walks -//! supersteps, schedules tasks, updates state, saves checkpoints, and raises -//! interrupts. Routing those events into a [`GraphEventSink`] is what lets a -//! UI or enclosing graph watch a subgraph or sub-agent execute in -//! real time; because every event is tagged with its node and step, the streams -//! of nested runs can be merged and attributed back up the run tree. -//! -//! See `types` for the event and stream-mode definitions. The executor emits -//! [`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 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 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. - /// - /// Sinks that persist asynchronously (off the executor thread) override this - /// so callers can guarantee the durable log has caught up — for example - /// before reading a journal back. The executor calls it after a terminal - /// run event. The default is a no-op for synchronous/in-memory sinks. - fn flush(&self) {} -} - -/// A sink that drops every event. -#[derive(Clone, Copy, Debug, Default)] -pub struct NoopSink; - -impl GraphEventSink for NoopSink { - 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>>, -} - -impl CollectingSink { - /// Creates an empty collecting sink. - pub fn new() -> Self { - Self::default() - } - - /// 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() - } - - /// Returns the number of recorded events. - pub fn len(&self) -> usize { - self.events.lock().map(|g| g.len()).unwrap_or(0) - } - - /// Returns true when no events were recorded. - pub fn is_empty(&self) -> bool { - self.len() == 0 - } -} - -impl GraphEventSink for CollectingSink { - fn emit(&self, envelope: GraphEventEnvelope) { - if let Ok(mut guard) = self.events.lock() { - guard.push(envelope); - } - } -} - -#[cfg(test)] -#[path = "mod_tests.rs"] -mod test; diff --git a/crates/tinyagents-graph/src/stream/mod_tests.rs b/crates/tinyagents-graph/src/stream/mod_tests.rs deleted file mode 100644 index 847434fe4..000000000 --- a/crates/tinyagents-graph/src/stream/mod_tests.rs +++ /dev/null @@ -1,31 +0,0 @@ -//! Unit tests for the event sinks: [`CollectingSink`] records emitted events in -//! order for inspection, and [`NoopSink`] silently drops them. - -use super::*; -use tinyagents_harness::ids::NodeId; - -#[test] -fn collecting_sink_records_events() { - let sink = CollectingSink::new(); - assert!(sink.is_empty()); - sink.emit(GraphEventEnvelope::for_test(GraphEvent::StepStarted { - step: 1, - active: vec![NodeId::from("a")], - })); - sink.emit(GraphEventEnvelope::for_test(GraphEvent::StepCompleted { - step: 1, - })); - assert_eq!(sink.len(), 2); - assert!(matches!( - sink.events()[0], - GraphEvent::StepStarted { step: 1, .. } - )); -} - -#[test] -fn noop_sink_drops_events() { - let sink = NoopSink; - sink.emit(GraphEventEnvelope::for_test(GraphEvent::StepCompleted { - step: 1, - })); -} diff --git a/crates/tinyagents-graph/src/stream/project.rs b/crates/tinyagents-graph/src/stream/project.rs deleted file mode 100644 index cf2ea8420..000000000 --- a/crates/tinyagents-graph/src/stream/project.rs +++ /dev/null @@ -1,303 +0,0 @@ -//! [`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)] -#[path = "project/project_tests.rs"] -mod test; diff --git a/crates/tinyagents-graph/src/stream/project/project_tests.rs b/crates/tinyagents-graph/src/stream/project/project_tests.rs deleted file mode 100644 index 79e30653d..000000000 --- a/crates/tinyagents-graph/src/stream/project/project_tests.rs +++ /dev/null @@ -1,178 +0,0 @@ -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 { - parent_call_id: None, - call_id: CallId::from("call-1".to_string()), - tool_name: "search".into(), - input: None, - }); - projection.fold_agent_event(&AgentEvent::ToolCompleted { - parent_call_id: None, - 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 { - parent_call_id: None, - 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 { - parent_call_id: None, - call_id: CallId::from("call-1".to_string()), - tool_name: "search".into(), - input: None, - }); - 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/types.rs b/crates/tinyagents-graph/src/stream/types.rs deleted file mode 100644 index c27456799..000000000 --- a/crates/tinyagents-graph/src/stream/types.rs +++ /dev/null @@ -1,386 +0,0 @@ -//! Low-level graph events and high-level stream modes — the wire vocabulary the -//! recursive executor uses to narrate its own execution. -//! -//! [`GraphEvent`] is the fine-grained, per-node/per-step lifecycle signal the -//! durable executor emits at every boundary; [`StreamMode`] is the LangGraph- -//! style selection of *which projection* of that stream a caller wants (full -//! values, per-node updates, model messages, debug detail, interrupts, or -//! custom node writes). Together they let observers — including a parent run -//! consuming a subgraph — follow a run without inspecting its internal state. - -use serde::{Deserialize, Serialize}; - -use crate::command::Interrupt; -use tinyagents_harness::ids::{CheckpointId, NodeId, RunId, TaskId}; - -/// A low-level graph lifecycle event emitted through a [`super::GraphEventSink`]. -/// -/// These are the durable-executor analogues of the observability event model in -/// the graph spec, reduced to the set the milestone executor actually emits. -/// -/// The variants are serde-serializable so a single event can be wrapped into a -/// durable [`crate::observability::GraphObservation`] envelope, journaled, -/// and replayed. -#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] -pub enum GraphEvent { - /// The run began (emitted once before the first superstep). - RunStarted { - /// The run that started. - run_id: RunId, - }, - /// The run finished successfully. - RunCompleted { - /// The run that completed. - run_id: RunId, - /// Total supersteps executed. - steps: usize, - }, - /// The run aborted with an error. - RunFailed { - /// The run that failed. - run_id: RunId, - /// 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. - step: usize, - /// Nodes scheduled to run this step. - active: Vec, - }, - /// A superstep finished and its boundary work (reducer, checkpoint) ran. - StepCompleted { - /// 1-based step number. - step: usize, - }, - /// A task was scheduled for a node in the active set. - TaskScheduled { - /// Target node. - node: NodeId, - /// 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. - node: NodeId, - /// Step number. - step: usize, - }, - /// A node handler completed successfully. - NodeCompleted { - /// Node id. - node: NodeId, - /// Step number. - step: usize, - }, - /// A node handler returned an error. - NodeFailed { - /// Node id. - node: NodeId, - /// Step number. - step: usize, - /// Rendered error. - error: String, - }, - /// A node handler failed with a retryable error and a retry was scheduled - /// under the graph's node-retry policy. Emitted before the (opt-in) backoff - /// wait and the re-run of the node from its start. - NodeRetryScheduled { - /// Node id. - node: NodeId, - /// Step number. - step: usize, - /// The 1-based retry attempt about to be made. - attempt: usize, - }, - /// A node produced a state update applied at the boundary. - StateUpdated { - /// Node id. - node: NodeId, - /// Step number. - step: usize, - }, - /// A route was selected for a node. - RouteSelected { - /// Source node. - node: NodeId, - /// Selected next node. - target: NodeId, - }, - /// A checkpoint was persisted at a superstep boundary. - 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 { - /// The checkpoint id that was loaded. - checkpoint_id: CheckpointId, - }, - /// A node emitted an interrupt and the run paused. - InterruptEmitted { - /// The emitted interrupt. - interrupt: Interrupt, - }, - /// An embedded subgraph began executing under a child namespace. - SubgraphStarted { - /// The parent node hosting the subgraph. - node: NodeId, - /// The child checkpoint namespace. - namespace: Vec, - }, - /// An embedded subgraph finished executing. - SubgraphCompleted { - /// The parent node hosting the subgraph. - node: NodeId, - /// The child checkpoint namespace. - namespace: Vec, - }, - /// A parallel superstep forked an execution branch for a node. - ContextForked { - /// The node whose branch was forked. - node: NodeId, - /// The branch (fork) index within the active set. - fork: usize, - /// Step number. - step: usize, - }, - /// The effective recursion/namespace depth changed. - RecursionDepthChanged { - /// The new depth (number of enclosing namespaces). - depth: usize, - }, - /// An arbitrary user-defined event written from inside a node. - Custom { - /// A stable name for the custom event. - name: String, - /// Free-form structured payload. - data: serde_json::Value, - }, -} - -impl GraphEvent { - /// Returns a stable, dot-separated string that names the kind of event. - /// - /// The returned string is a static literal, suitable for logging, - /// filtering, and serde-independent routing. Examples: `"run.started"`, - /// `"step.started"`, `"node.completed"`. - pub fn kind(&self) -> &'static str { - match self { - 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", - GraphEvent::NodeRetryScheduled { .. } => "node.retry_scheduled", - GraphEvent::StateUpdated { .. } => "state.updated", - GraphEvent::RouteSelected { .. } => "route.selected", - GraphEvent::CheckpointSaved { .. } => "checkpoint.saved", - GraphEvent::CheckpointRestored { .. } => "checkpoint.restored", - GraphEvent::InterruptEmitted { .. } => "interrupt.emitted", - GraphEvent::SubgraphStarted { .. } => "subgraph.started", - GraphEvent::SubgraphCompleted { .. } => "subgraph.completed", - GraphEvent::ContextForked { .. } => "context.forked", - GraphEvent::RecursionDepthChanged { .. } => "recursion.depth_changed", - GraphEvent::Custom { .. } => "custom", - } - } - - /// Returns the superstep number this event is associated with, when the - /// variant carries one. Used to stamp the `step` field of a durable - /// [`crate::observability::GraphObservation`]. - pub fn step(&self) -> Option { - match self { - 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, .. } | GraphEvent::RunDrained { steps, .. } => { - Some(*steps) - } - _ => None, - } - } -} - -/// High-level projection modes for a graph run stream. -/// -/// 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, - /// Per-node/per-task state updates. - Updates, - /// Harness message or token deltas from model nodes. - Messages, - /// Checkpoints plus task internals. - Debug, - /// Pending interrupts only. - 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/README.md b/crates/tinyagents-graph/src/subagent_node/README.md deleted file mode 100644 index 24586c0fe..000000000 --- a/crates/tinyagents-graph/src/subagent_node/README.md +++ /dev/null @@ -1,95 +0,0 @@ -# graph::subagent_node - -> **Deprecated.** `SubAgentNode` and `subagent_node` are marked -> `#[deprecated(note = "use tinyagents_orchestration SubagentDriver")]`. Graph -> cannot depend on orchestration, so the node cannot adapt the driver; its -> behaviour is unchanged. Use `SubagentDriver` for new work. - -Sub-agent nodes — the graph node that delegates to a harness *agent* (a -model-driven agent loop) invoked through an explicit host capability. - -Where `graph::subgraph` embeds an entire `CompiledGraph` as a node (graph -recursion), this module embeds a *harness agent* as a node (graph-to-agent -recursion): a graph step hands its work to a host-selected, -independently-observable agent and folds the agent's answer back into the -parent graph state. A graph never synthesizes a harness `RunContext` itself; -it only ever calls out through a host-bound `AgentInvoker`, so a graph without -one fails closed rather than silently fabricating agent behavior. - -## Public surface - -- `SubAgentNode` — binds an agent name to an `InputMapper` - (parent `State` → `SubAgentInput`), an `OutputMapper` (`SubAgentOutput` → - parent `Update`), and a `SubAgentPolicy`. `::new` / `::from_fns` construct - it; `.with_policy` sets the policy. -- `subagent_node(node) -> Handler` — lowers a `SubAgentNode` - into an ordinary graph node handler: resolves the agent's `AgentInvoker` - from the node context, projects state into `SubAgentInput`, mints a child - `run_id` parented to the enclosing run, runs under the policy - (timeout/retry/budget), records the child run (with its usage) onto the - parent's execution rollup, forwards the child's events, and folds the - output into a parent update. -- `AgentInvoker` (trait) — the object-safe host boundary for graph-to-agent - recursion; implementations must dispatch through the same host entry point - used for a top-level harness run and create the child via - `RunContext::child`. -- `AgentInvocation` — an explicit request for a host-owned recursive agent - invocation (agent id, mapped input, graph/node identity, parent/root run - ids, forwarded event sink and cancellation). -- `AgentInvocationBinding` — the atomic, execution-scoped capability supplied - to one graph run (`invoker` + forwarded `events` + `cancellation`); never - stored on a reusable `CompiledGraph`. -- `SubAgentInput` / `SubAgentOutput` — the structured carriers crossing the - graph↔agent boundary; `SubAgentOutput` also carries `UsageTotals` and call - counts so usage rolls up and a budget can be enforced. -- `SubAgentPolicy` — timeout/retry/budget policy for one invocation; defers to - the harness `RetryPolicy` for backoff. Default is a conservative *single - attempt, no timeout, no budget* so a node never silently re-runs a - non-idempotent agent. -- `SubAgentBudget` — optional caps on model/tool calls and input/output tokens, - checked *after* the child run returns; violating one fails the node with - `TinyAgentsError::LimitExceeded`. `apply_call_caps` tightens a `RunConfig` - for in-run enforcement; `max_cost` is carried but not enforced here - (`to_budget_limits` hands it to the harness `BudgetMiddleware`). - `SubAgentPolicy::retry_after_tool_calls` (default `false`) is honoured by the - orchestration paths; the graph node cannot see whether a failed attempt ran - tools. The policy types are re-exported from - `tinyagents_orchestration::subagent`. -- `InputMapper` / `OutputMapper` — type aliases for the parent↔ - child mapping closures. - -## Files - -| File | Role | -| --- | --- | -| `types.rs` | `AgentInvocation`, `AgentInvocationBinding`, `AgentInvoker`, `SubAgentInput`, `SubAgentOutput`, `SubAgentBudget`, `SubAgentPolicy`, `InputMapper`, `OutputMapper`, `SubAgentNode`. | -| `mod.rs` | `SubAgentNode` builders, `subagent_node` (the handler lowering), retry/timeout/budget execution (`run_with_policy`), and child-run recording. | -| `mod_tests.rs` | Unit tests (delegation, mapper wiring, retry/timeout behavior, budget enforcement, child-run recording). | - -## Operational constraints - -- `AgentInvocationBinding` is execution-scoped: it is supplied per graph run - (via `NodeContext::agent_binding`), never cached on a `CompiledGraph`. - Cloning it is only meant for descendants of the same execution tree — do - not reuse one binding across unrelated top-level executions, or their - cancellation/event signals bleed into each other. -- A node without an attached `AgentInvocationBinding` fails immediately with - `TinyAgentsError::Capability` rather than falling back to any default - behavior — this is deliberate fail-closed design, not a bug to work around - by fabricating a binding. -- `SubAgentPolicy`'s default retry policy allows exactly one attempt; a caller - that wants retries must opt in explicitly with `with_retry`, since retrying - a non-idempotent agent call by default would be unsafe. -- The minted child `run_id` always preserves the run tree's `root_run_id` and - is parented to the enclosing graph run, so observability/tracing (see - `graph::observability`) can nest a sub-agent's model/tool calls under the - graph run that spawned it. - -## Relation to neighbours - -Structural counterpart to `graph::subgraph` (graph-to-graph recursion) and -`graph::parallel` (fan-out without an agent boundary). Child-run bookkeeping -here mirrors what `graph::subgraph` records via `crate::recursion::ChildRun`, -so both recursion styles show up uniformly in `GraphExecution::child_runs` / -`run_tree`. `graph::testkit::subagent_fake_node` fakes this module's -child-run recording for tests without a live agent or registry. diff --git a/crates/tinyagents-graph/src/subagent_node/mod.rs b/crates/tinyagents-graph/src/subagent_node/mod.rs deleted file mode 100644 index 9ce99778a..000000000 --- a/crates/tinyagents-graph/src/subagent_node/mod.rs +++ /dev/null @@ -1,204 +0,0 @@ -//! Sub-agent nodes — the graph node that delegates to a harness *agent* (a -//! model-driven agent loop) invoked through an explicit host capability. -//! -//! Where [`crate::subgraph`] embeds an entire [`CompiledGraph`](crate::CompiledGraph) as a -//! node, this module embeds a *harness agent* as a node: a graph step hands its -//! work to a host-selected, independently-observable agent and folds the agent's -//! answer back into the parent graph state. -//! -//! The pieces: -//! -//! - [`SubAgentNode`] binds an agent `ComponentId` to an [`InputMapper`] -//! (parent `State` → [`SubAgentInput`]), an [`OutputMapper`] -//! ([`SubAgentOutput`] → parent `Update`), and a [`SubAgentPolicy`]. -//! - [`subagent_node`] lowers a [`SubAgentNode`] into an ordinary graph node -//! `Handler`: it obtains the carried [`AgentInvoker`], creates a distinct -//! child `run_id` that preserves the run tree's `root_run_id` and is parented -//! to the enclosing graph run, applies timeout/retry/budget policy, maps the -//! child output into the parent update, records the child run (with its usage) -//! onto the parent execution rollup, and forwards the child run's harness -//! events onto the host-provided event sink. -//! -//! **Deprecated** (see [`SubAgentNode`]): use `tinyagents_orchestration`'s -//! `SubagentDriver`. The graph crate cannot depend on orchestration, so the -//! node keeps its own behaviour rather than adapting the driver. -//! -//! See `types` for the data definitions and `test.rs` for focused tests. - -mod types; - -use std::future::Future; -use std::pin::Pin; -use std::sync::Arc; -use std::time::Duration; - -pub use types::*; - -use crate::builder::NodeContext; -use crate::command::NodeResult; -use crate::recursion::ChildRun; -use crate::{Result, TinyAgentsError}; -use tinyagents_harness::ids::next_seq; -use tinyagents_harness::ids::{GraphId, RunId}; - -type Handler = Box< - dyn Fn(S, NodeContext) -> Pin>> + Send>> - + Send - + Sync, ->; - -#[allow(deprecated)] -impl SubAgentNode { - /// Builds a sub-agent node delegating to the registered agent named `agent`, - /// with the given parent↔child mappers and a default [`SubAgentPolicy`]. - pub fn new( - agent: impl Into, - input_mapper: InputMapper, - output_mapper: OutputMapper, - ) -> Self { - Self { - agent: agent.into(), - input_mapper, - output_mapper, - policy: SubAgentPolicy::default(), - } - } - - /// Builds a sub-agent node from plain closures (a convenience over - /// [`SubAgentNode::new`] that wraps them as [`InputMapper`]/[`OutputMapper`]). - pub fn from_fns(agent: impl Into, input: I, output: O) -> Self - where - I: Fn(&State) -> SubAgentInput + Send + Sync + 'static, - O: Fn(SubAgentOutput) -> Update + Send + Sync + 'static, - { - Self::new(agent, Arc::new(input), Arc::new(output)) - } - - /// Sets the invocation policy, returning `self` for chaining. - pub fn with_policy(mut self, policy: SubAgentPolicy) -> Self { - self.policy = policy; - self - } -} - -/// Lowers a [`SubAgentNode`] plus a host-bound `invoker` into a graph node -/// `Handler`. -/// -/// At each activation the handler: -/// -/// 1. resolves [`SubAgentNode::agent`] against `registry` (failing with -/// [`TinyAgentsError::Capability`] when the name is not a registered agent), -/// 2. projects the committed `state` into a [`SubAgentInput`], -/// 3. mints a distinct child `run_id` that preserves the run tree's -/// `root_run_id` and is parented to the enclosing graph run, -/// 4. runs the agent under the node's [`SubAgentPolicy`] (timeout/retry), then -/// enforces the work budget, -/// 5. records the child run — with its rolled-up [`UsageTotals`](tinyinference_llm::usage::UsageTotals) — onto the -/// enclosing run's child-run sink, and -/// 6. folds the [`SubAgentOutput`] into a parent `Update` via the output mapper. -#[deprecated(note = "use tinyagents_orchestration SubagentDriver")] -#[allow(deprecated)] -pub fn subagent_node(node: SubAgentNode) -> Handler -where - State: Clone + Send + Sync + 'static, - Update: Send + 'static, -{ - let node = Arc::new(node); - Box::new(move |state: State, ctx: NodeContext| { - let node = node.clone(); - Box::pin(async move { - let input = (node.input_mapper)(&state); - let binding = ctx.agent_binding.clone().ok_or_else(|| { - TinyAgentsError::Capability(format!( - "sub-agent `{}` requires an execution-scoped AgentInvocationBinding", - node.agent - )) - })?; - let output = run_with_policy(&binding, &node.agent, input, &ctx, &node.policy).await?; - - record_child_run(&ctx, &node.agent, &output); - - let update = (node.output_mapper)(output); - Ok(NodeResult::Update(update)) - }) - }) -} - -/// Runs `agent` under `policy`: applies the per-attempt timeout, retries -/// transient failures per the retry policy, then enforces the work budget. -async fn run_with_policy( - binding: &AgentInvocationBinding, - agent_id: &str, - input: SubAgentInput, - ctx: &NodeContext, - policy: &SubAgentPolicy, -) -> Result { - let mut attempt = 0; - loop { - let fut = binding.invoker.invoke(AgentInvocation { - agent_id: agent_id.to_string(), - input: input.clone(), - graph_id: ctx.graph_id.clone(), - node_id: ctx.node_id.clone(), - parent_run_id: ctx.run_id.clone(), - root_run_id: ctx - .root_run_id - .clone() - .unwrap_or_else(|| ctx.run_id.clone()), - events: binding.events.clone(), - cancellation: Some(binding.cancellation.clone()), - }); - let result = match policy.timeout { - Some(timeout) => match tokio::time::timeout(timeout, fut).await { - Ok(result) => result, - Err(_) => Err(TinyAgentsError::Timeout(format!( - "sub-agent `{}` timed out after {timeout:?}", - agent_id - ))), - }, - None => fut.await, - }; - - match result { - Ok(output) => return policy.budget.check(&output, agent_id).map(|()| output), - Err(err) => { - if policy.retry.should_retry(attempt) - && tinyagents_harness::retry::is_retryable(&err) - { - attempt += 1; - let backoff = policy.retry.backoff_for_attempt(attempt); - if backoff > Duration::ZERO { - tokio::time::sleep(backoff).await; - } - continue; - } - return Err(err); - } - } - } -} - -/// Mints a distinct child run id (preserving the run tree root, parented to the -/// enclosing graph run) and records it — with the child's rolled-up usage — -/// onto the enclosing run's child-run sink, when one is attached. -fn record_child_run(ctx: &NodeContext, agent: &str, output: &SubAgentOutput) { - let Some(sink) = &ctx.child_runs else { - return; - }; - let root_run_id = ctx - .root_run_id - .clone() - .unwrap_or_else(|| ctx.run_id.clone()); - sink.record(ChildRun { - node: ctx.node_id.clone(), - graph_id: GraphId::new(format!("agent:{agent}")), - run_id: RunId::new(format!("subagent-{}", next_seq())), - root_run_id, - usage: output.usage, - checkpoint_id: None, - }); -} - -#[cfg(test)] -#[path = "mod_tests.rs"] -mod test; diff --git a/crates/tinyagents-graph/src/subagent_node/mod_tests.rs b/crates/tinyagents-graph/src/subagent_node/mod_tests.rs deleted file mode 100644 index 529d650da..000000000 --- a/crates/tinyagents-graph/src/subagent_node/mod_tests.rs +++ /dev/null @@ -1,314 +0,0 @@ -//! Tests for host-driven graph-to-agent delegation. -// The deprecated `SubAgentNode` keeps its behaviour and its coverage. -#![allow(deprecated)] - -use std::sync::{Arc, Mutex}; - -use async_trait::async_trait; - -use super::*; -use crate::builder::GraphBuilder; -use tinyagents_harness::cancel::CancellationToken; -use tinyagents_harness::events::{AgentEvent, EventSink, RecordingListener}; - -#[derive(Clone, Default)] -struct RecordingInvoker { - requests: Arc>>, -} - -impl RecordingInvoker { - fn requests(&self) -> Vec { - self.requests.lock().unwrap().clone() - } -} - -#[async_trait] -impl AgentInvoker for RecordingInvoker { - async fn invoke(&self, request: AgentInvocation) -> crate::Result { - self.requests.lock().unwrap().push(request.clone()); - request.events.emit(AgentEvent::StateUpdate); - Ok(SubAgentOutput { - text: format!("done:{}", request.input.prompt), - model_calls: 1, - ..SubAgentOutput::default() - }) - } -} - -struct FailingInvoker; - -#[async_trait] -impl AgentInvoker for FailingInvoker { - async fn invoke(&self, _request: AgentInvocation) -> crate::Result { - Err(crate::TinyAgentsError::Model( - "temporary failure".to_string(), - )) - } -} - -fn graph() -> crate::CompiledGraph { - GraphBuilder::::overwrite() - .add_node( - "delegate", - subagent_node(SubAgentNode::from_fns( - "researcher", - |state: &String| SubAgentInput::prompt(state.clone()), - |output: SubAgentOutput| output.text, - )), - ) - .set_entry("delegate") - .set_finish("delegate") - .compile() - .unwrap() -} - -fn binding(invoker: Arc) -> AgentInvocationBinding { - AgentInvocationBinding::new(invoker, EventSink::new(), CancellationToken::new()) -} - -#[tokio::test] -async fn delegation_uses_carried_invoker_and_preserves_graph_lineage() { - let invoker = Arc::new(RecordingInvoker::default()); - let graph = graph(); - - let run = graph - .run_with_agent_binding("question".to_string(), binding(invoker.clone())) - .await - .unwrap(); - - assert_eq!(run.state, "done:question"); - let requests = invoker.requests(); - assert_eq!(requests.len(), 1); - let request = &requests[0]; - assert_eq!(request.agent_id, "researcher"); - assert_eq!(request.input.prompt, "question"); - assert_eq!(request.parent_run_id, run.run_id); - assert_eq!(request.root_run_id, run.root_run_id); - assert_eq!(request.node_id.as_str(), "delegate"); - assert_eq!(request.graph_id, run.graph_id); - assert!(request.cancellation.is_some()); -} - -#[tokio::test] -async fn concurrent_sibling_graph_runs_have_isolated_parent_identity() { - let left_invoker = Arc::new(RecordingInvoker::default()); - let right_invoker = Arc::new(RecordingInvoker::default()); - let graph = graph(); - let left_cancellation = CancellationToken::new(); - left_cancellation.cancel(); - let right_cancellation = CancellationToken::new(); - - let (left, right) = tokio::join!( - graph.run_with_agent_binding( - "left".to_string(), - AgentInvocationBinding::new(left_invoker.clone(), EventSink::new(), left_cancellation,) - ), - graph.run_with_agent_binding( - "right".to_string(), - AgentInvocationBinding::new( - right_invoker.clone(), - EventSink::new(), - right_cancellation, - ) - ) - ); - let left = left.unwrap(); - let right = right.unwrap(); - let left_request = left_invoker.requests().pop().unwrap(); - let right_request = right_invoker.requests().pop().unwrap(); - assert_ne!(left.run_id, right.run_id); - assert_eq!(left_request.input.prompt, "left"); - assert_eq!(left_request.parent_run_id, left.run_id); - assert_eq!(left_request.root_run_id, left.root_run_id); - assert!(left_request.cancellation.unwrap().is_cancelled()); - assert_eq!(right_request.input.prompt, "right"); - assert_eq!(right_request.parent_run_id, right.run_id); - assert_eq!(right_request.root_run_id, right.root_run_id); - assert!(!right_request.cancellation.unwrap().is_cancelled()); -} - -#[tokio::test] -async fn missing_host_invoker_is_an_explicit_capability_error() { - let graph = GraphBuilder::::overwrite() - .add_node( - "delegate", - subagent_node(SubAgentNode::from_fns( - "researcher", - |state: &String| SubAgentInput::prompt(state.clone()), - |output: SubAgentOutput| output.text, - )), - ) - .set_entry("delegate") - .set_finish("delegate") - .compile() - .unwrap(); - - let error = graph.run("question".to_string()).await.unwrap_err(); - assert!(matches!(error, crate::TinyAgentsError::Capability(_))); -} - -#[tokio::test] -async fn resume_binding_reaches_a_later_subagent_with_host_signals() { - // The initial invocation pauses before reaching the sub-agent, so the - // continuation must supply a fresh execution-scoped binding rather than - // relying on mutable state retained by the compiled graph. - let checkpointer = Arc::new(crate::checkpoint::InMemoryCheckpointer::::new()); - let graph = GraphBuilder::::overwrite() - .add_node( - "gate", - |state: String, ctx: crate::builder::NodeContext| async move { - if ctx.resume.is_some() { - Ok(crate::command::NodeResult::Update(state)) - } else { - Ok(crate::command::NodeResult::Interrupt( - crate::command::Interrupt::new("gate", serde_json::json!({ "ask": "go?" })), - )) - } - }, - ) - .add_node( - "delegate", - subagent_node(SubAgentNode::from_fns( - "researcher", - |state: &String| SubAgentInput::prompt(state.clone()), - |output: SubAgentOutput| output.text, - )), - ) - .set_entry("gate") - .add_edge("gate", "delegate") - .set_finish("delegate") - .compile() - .unwrap() - .with_checkpointer(checkpointer); - - let paused = graph - .run_with_thread("resume", "question".to_string()) - .await - .unwrap(); - assert!(paused.is_interrupted()); - - let invoker = Arc::new(RecordingInvoker::default()); - let events = EventSink::new(); - let listener = Arc::new(RecordingListener::new()); - events.subscribe(listener.clone()); - let cancellation = CancellationToken::new(); - cancellation.cancel(); - - let resumed = graph - .resume_with_agent_binding( - "resume", - crate::command::Command::resume(serde_json::json!("approved")), - AgentInvocationBinding::new(invoker.clone(), events, cancellation), - ) - .await - .unwrap(); - - assert_eq!(resumed.state, "done:question"); - let request = invoker.requests().pop().expect("resumed sub-agent ran"); - assert_eq!(request.parent_run_id, resumed.run_id); - assert_eq!(request.root_run_id, resumed.root_run_id); - assert!( - request - .cancellation - .expect("binding has cancellation") - .is_cancelled() - ); - assert_eq!( - listener.len(), - 1, - "the resumed sub-agent used the supplied event sink" - ); -} - -#[tokio::test] -async fn retry_binding_is_fresh_and_reaches_the_failed_subagent() { - // A durable failure from one execution must not capture that execution's - // host capability. Retrying with a different binding reaches the node and - // succeeds through the new invoker. - let checkpointer = Arc::new(crate::checkpoint::InMemoryCheckpointer::::new()); - let graph = graph().with_checkpointer(checkpointer); - - let failed = graph - .run_with_thread_agent_binding( - "retry", - "question".to_string(), - binding(Arc::new(FailingInvoker)), - ) - .await - .unwrap_err(); - assert!(matches!(failed, crate::TinyAgentsError::Model(_))); - - let replacement = Arc::new(RecordingInvoker::default()); - let retried = graph - .retry_with_agent_binding("retry", binding(replacement.clone())) - .await - .unwrap(); - - assert_eq!(retried.state, "done:question"); - let request = replacement - .requests() - .pop() - .expect("retry used replacement binding"); - assert_eq!(request.parent_run_id, retried.run_id); - assert_eq!(request.root_run_id, retried.root_run_id); -} - -fn usage_output(input: u64, output: u64) -> SubAgentOutput { - let mut out = SubAgentOutput::default(); - out.usage.usage.input_tokens = input; - out.usage.usage.output_tokens = output; - out -} - -#[test] -fn token_budget_is_enforced_after_the_child_returns() { - let budget = SubAgentBudget::unlimited() - .with_max_input_tokens(100) - .with_max_output_tokens(50); - assert!(budget.check(&usage_output(100, 50), "a").is_ok()); - let err = budget.check(&usage_output(101, 0), "a").unwrap_err(); - assert!(matches!(err, TinyAgentsError::LimitExceeded(m) if m.contains("input-token"))); - let err = budget.check(&usage_output(0, 51), "a").unwrap_err(); - assert!(matches!(err, TinyAgentsError::LimitExceeded(m) if m.contains("output-token"))); -} - -#[test] -fn budget_maps_onto_the_harness_budget_limits_including_cost() { - let limits = SubAgentBudget::unlimited() - .with_max_input_tokens(7) - .with_max_output_tokens(8) - .with_max_cost(1.5) - .to_budget_limits(); - assert_eq!(limits.max_input_tokens, Some(7)); - assert_eq!(limits.max_output_tokens, Some(8)); - assert_eq!(limits.max_cost, Some(1.5)); -} - -#[test] -fn call_caps_only_ever_tighten_a_run_config() { - let budget = SubAgentBudget { - max_model_calls: Some(3), - max_tool_calls: Some(9), - ..SubAgentBudget::unlimited() - }; - let mut loose = tinyagents_harness::context::RunConfig::new("r"); - budget.apply_call_caps(&mut loose); - assert_eq!(loose.max_model_calls, Some(3)); - assert_eq!(loose.max_tool_calls, Some(9)); - let mut tight = tinyagents_harness::context::RunConfig::new("r") - .with_max_model_calls(2) - .with_max_tool_calls(20); - budget.apply_call_caps(&mut tight); - assert_eq!(tight.max_model_calls, Some(2), "a tighter cap is kept"); - assert_eq!(tight.max_tool_calls, Some(9)); -} - -#[test] -fn policy_defaults_to_no_retry_after_tool_calls() { - assert!(!SubAgentPolicy::default().retry_after_tool_calls); - assert!( - SubAgentPolicy::default() - .with_retry_after_tool_calls(true) - .retry_after_tool_calls - ); -} diff --git a/crates/tinyagents-graph/src/subagent_node/types.rs b/crates/tinyagents-graph/src/subagent_node/types.rs deleted file mode 100644 index 90702581c..000000000 --- a/crates/tinyagents-graph/src/subagent_node/types.rs +++ /dev/null @@ -1,382 +0,0 @@ -//! Type definitions for [`SubAgentNode`](super::SubAgentNode) — the graph node -//! that delegates through a host-bound [`AgentInvoker`]. -//! -//! See the module [`mod`](super) docs for how these are wired into a node -//! handler. This file holds the host invocation trait ([`AgentInvoker`]), the -//! structured input/output carriers -//! ([`SubAgentInput`]/[`SubAgentOutput`]), the per-call policy -//! ([`SubAgentPolicy`]/[`SubAgentBudget`]), the parent↔child mapping aliases -//! ([`InputMapper`]/[`OutputMapper`]), and the [`SubAgentNode`] descriptor. - -use std::sync::Arc; -use std::time::Duration; - -use async_trait::async_trait; -use serde_json::Value; - -use crate::Result; -use tinyagents_harness::cancel::CancellationToken; -use tinyagents_harness::context::RunConfig; -use tinyagents_harness::events::EventSink; -use tinyagents_harness::ids::{GraphId, NodeId, RunId}; -use tinyagents_harness::middleware::BudgetLimits; -use tinyagents_harness::retry::RetryPolicy; -use tinyinference_llm::usage::UsageTotals; - -/// An explicit request for a host-owned recursive agent invocation. -/// -/// A graph never synthesizes a harness context. The host binds an invoker to a -/// graph entry point and uses the request's parent graph identity to invoke its -/// actual parent [`RunContext`](tinyagents_harness::context::RunContext), which -/// in turn creates the child through `RunContext::child`. -#[derive(Clone)] -pub struct AgentInvocation { - /// Registered agent definition id. - pub agent_id: String, - /// Mapped child input. - pub input: SubAgentInput, - /// Graph containing the delegating node. - pub graph_id: GraphId, - /// Node issuing the delegation. - pub node_id: NodeId, - /// Immediate parent graph run. - pub parent_run_id: RunId, - /// Root graph run shared by all descendants. - pub root_run_id: RunId, - /// Parent's event sink, forwarded by the host into the child context. - pub events: EventSink, - /// Parent cancellation signal, forwarded by the host into the child context. - pub cancellation: Option, -} - -/// Atomic, execution-scoped host capability for recursive agent invocation. -/// -/// This value is supplied to one graph run; it is never stored on a reusable -/// [`CompiledGraph`](crate::CompiledGraph). Cloning it is only for descendants -/// of that same execution tree, which preserves one parent invocation context -/// while preventing separate top-level executions from bleeding signals. -#[derive(Clone)] -pub struct AgentInvocationBinding { - /// Host entry point bound to the parent invocation context. - pub invoker: Arc, - /// Parent event sink forwarded to every descendant request. - pub events: EventSink, - /// Parent cancellation signal forwarded to every descendant request. - pub cancellation: CancellationToken, -} - -impl AgentInvocationBinding { - /// Creates the complete binding required for graph-to-agent recursion. - #[must_use] - pub fn new( - invoker: Arc, - events: EventSink, - cancellation: CancellationToken, - ) -> Self { - Self { - invoker, - events, - cancellation, - } - } -} - -/// Object-safe host boundary for graph-to-agent recursion. -/// -/// Implementations must be explicitly bound by the host to an owned or -/// `Arc`-backed parent invocation context. They must dispatch through the same -/// host entry point used for a top-level harness run and create the child with -/// `RunContext::child`. This avoids globals, task locals, downcasts, and -/// fabricated `Default` state. A graph without an invoker fails closed. -#[async_trait] -pub trait AgentInvoker: Send + Sync { - /// Invokes a requested agent using the host-bound parent context. - async fn invoke(&self, request: AgentInvocation) -> Result; -} - -/// The structured input a [`SubAgentNode`] hands to a delegated agent. -/// -/// `prompt` is the user-facing text the child run is seeded with; `data` is an -/// optional structured payload an agent (or an adapter) may consult. -#[derive(Clone, Debug, Default, PartialEq, Eq)] -pub struct SubAgentInput { - /// The user prompt the child run is seeded with. - pub prompt: String, - /// Optional structured side-channel payload. - pub data: Option, -} - -impl SubAgentInput { - /// Builds an input carrying just a `prompt`. - pub fn prompt(prompt: impl Into) -> Self { - Self { - prompt: prompt.into(), - data: None, - } - } - - /// Attaches a structured `data` payload, returning `self` for chaining. - pub fn with_data(mut self, data: Value) -> Self { - self.data = Some(data); - self - } -} - -/// The structured output a delegated agent returns to its [`SubAgentNode`]. -/// -/// Besides the final `text` and any parsed `structured` value, the output -/// carries the child run's [`UsageTotals`] and call counts so the node can roll -/// the child's usage into the parent execution and a budget can be enforced. -#[derive(Clone, Debug, Default, PartialEq, Eq)] -pub struct SubAgentOutput { - /// The child run's final assistant text. - pub text: String, - /// Parsed structured output, when the child run produced one. - pub structured: Option, - /// Cumulative token usage across the child run's model calls. - pub usage: UsageTotals, - /// Number of model calls the child run dispatched. - pub model_calls: usize, - /// Number of tool invocations the child run executed. - pub tool_calls: usize, -} - -/// An optional cap on the work a single sub-agent invocation may perform. -/// -/// What is enforced, and where: -/// -/// | field | enforcement | -/// | --- | --- | -/// | `max_model_calls`, `max_tool_calls` | after the child returns here ([`Self::check`]); *during* the run wherever a [`RunConfig`] is available ([`Self::apply_call_caps`]) | -/// | `max_input_tokens`, `max_output_tokens` | after the child returns ([`Self::check`]) from the reported usage; a run that overshoots is reported, not interrupted | -/// | `max_cost` | **not enforced** here: cost needs a pricing table this layer does not own. [`Self::to_budget_limits`] hands the same caps to the harness `BudgetMiddleware`, which does enforce them during the run | -/// -/// A run that exceeds a checked cap fails with -/// [`TinyAgentsError::LimitExceeded`](crate::TinyAgentsError::LimitExceeded). -/// The default ([`SubAgentBudget::unlimited`]) imposes no cap. -#[derive(Clone, Copy, Debug, Default, PartialEq)] -pub struct SubAgentBudget { - /// Maximum model calls the child run may make (`None` = unbounded). - pub max_model_calls: Option, - /// Maximum tool calls the child run may make (`None` = unbounded). - pub max_tool_calls: Option, - /// Maximum cumulative input tokens (`None` = unbounded). Checked after the - /// child returns. - pub max_input_tokens: Option, - /// Maximum cumulative output tokens (`None` = unbounded). Checked after - /// the child returns. - pub max_output_tokens: Option, - /// Maximum cumulative estimated cost in the pricing table's currency - /// (`None` = unbounded). Not enforced by [`Self::check`]; see the type - /// docs and [`Self::to_budget_limits`]. - pub max_cost: Option, -} - -impl SubAgentBudget { - /// A budget that imposes no cap. - pub fn unlimited() -> Self { - Self::default() - } - - /// Caps model calls, returning `self` for chaining. - pub fn with_max_model_calls(mut self, max: usize) -> Self { - self.max_model_calls = Some(max); - self - } - - /// Caps tool calls, returning `self` for chaining. - pub fn with_max_tool_calls(mut self, max: usize) -> Self { - self.max_tool_calls = Some(max); - self - } - - /// Caps cumulative input tokens, returning `self` for chaining. - pub fn with_max_input_tokens(mut self, max: u64) -> Self { - self.max_input_tokens = Some(max); - self - } - - /// Caps cumulative output tokens, returning `self` for chaining. - pub fn with_max_output_tokens(mut self, max: u64) -> Self { - self.max_output_tokens = Some(max); - self - } - - /// Caps cumulative estimated cost, returning `self` for chaining. See the - /// type docs: this is carried, not enforced, by this layer. - pub fn with_max_cost(mut self, max: f64) -> Self { - self.max_cost = Some(max); - self - } - - /// Tightens `config`'s call caps to this budget. A cap the config already - /// sets lower is kept: a budget can narrow a run, never widen it. - pub fn apply_call_caps(&self, config: &mut RunConfig) { - if let Some(max) = self.max_model_calls { - config.max_model_calls = Some(config.max_model_calls.map_or(max, |set| set.min(max))); - } - if let Some(max) = self.max_tool_calls { - config.max_tool_calls = Some(config.max_tool_calls.map_or(max, |set| set.min(max))); - } - } - - /// The token and cost caps as the harness's in-run - /// [`BudgetLimits`], for a host to attach as a `BudgetMiddleware` on the - /// child harness — the only place a cost cap can be enforced mid-run. - pub fn to_budget_limits(&self) -> BudgetLimits { - BudgetLimits { - max_input_tokens: self.max_input_tokens, - max_output_tokens: self.max_output_tokens, - max_cost: self.max_cost, - ..BudgetLimits::default() - } - } - - /// Returns `Ok(())` when `output` is within budget, else - /// [`TinyAgentsError::LimitExceeded`](crate::TinyAgentsError::LimitExceeded). - pub fn check(&self, output: &SubAgentOutput, agent: &str) -> Result<()> { - if let Some(max) = self.max_model_calls - && output.model_calls > max - { - return Err(crate::TinyAgentsError::LimitExceeded(format!( - "sub-agent `{agent}` exceeded model-call budget: {} > {max}", - output.model_calls - ))); - } - if let Some(max) = self.max_tool_calls - && output.tool_calls > max - { - return Err(crate::TinyAgentsError::LimitExceeded(format!( - "sub-agent `{agent}` exceeded tool-call budget: {} > {max}", - output.tool_calls - ))); - } - if let Some(max) = self.max_input_tokens - && output.usage.usage.input_tokens > max - { - return Err(crate::TinyAgentsError::LimitExceeded(format!( - "sub-agent `{agent}` exceeded input-token budget: {} > {max}", - output.usage.usage.input_tokens - ))); - } - if let Some(max) = self.max_output_tokens - && output.usage.usage.output_tokens > max - { - return Err(crate::TinyAgentsError::LimitExceeded(format!( - "sub-agent `{agent}` exceeded output-token budget: {} > {max}", - output.usage.usage.output_tokens - ))); - } - Ok(()) - } -} - -/// Timeout, retry, and budget policy applied around a sub-agent invocation. -/// -/// This is a thin graph-local struct that defers to the harness -/// [`RetryPolicy`] for retry/backoff and reuses the harness usage accounting for -/// budgeting, so a graph node and an in-harness sub-agent share the same -/// resilience semantics. The default is a *single attempt, no timeout, no -/// budget* — deliberately conservative so a node never silently re-runs a -/// non-idempotent agent. -#[derive(Clone, Debug)] -pub struct SubAgentPolicy { - /// Optional wall-clock timeout for a single attempt. On elapse the node - /// fails with [`TinyAgentsError::Timeout`](crate::TinyAgentsError::Timeout). - pub timeout: Option, - /// Retry/backoff policy applied across attempts. - pub retry: RetryPolicy, - /// Optional cap on the work the child run may perform. - pub budget: SubAgentBudget, - /// Whether a failed attempt may be retried after the child already ran - /// tools. Defaults to `false`: tools can have side effects, so a retry is - /// only safe when the failed attempt provably ran none. - /// - /// The graph node cannot observe whether a failed [`AgentInvoker`] call - /// ran tools (an error carries no run), so it retries exactly as the - /// `retry` policy says; the orchestration paths (`SubagentDriver`, - /// `SubAgentTool`) can, and honour this flag. - pub retry_after_tool_calls: bool, -} - -impl Default for SubAgentPolicy { - fn default() -> Self { - Self { - timeout: None, - retry: RetryPolicy::default().with_max_attempts(1), - budget: SubAgentBudget::unlimited(), - retry_after_tool_calls: false, - } - } -} - -impl SubAgentPolicy { - /// Sets the per-attempt timeout, returning `self` for chaining. - pub fn with_timeout(mut self, timeout: Duration) -> Self { - self.timeout = Some(timeout); - self - } - - /// Sets the retry policy, returning `self` for chaining. - pub fn with_retry(mut self, retry: RetryPolicy) -> Self { - self.retry = retry; - self - } - - /// Sets the work budget, returning `self` for chaining. - pub fn with_budget(mut self, budget: SubAgentBudget) -> Self { - self.budget = budget; - self - } - - /// Allows retrying a failed attempt even after the child ran tools, - /// returning `self` for chaining. Only for idempotent tool sets. - pub fn with_retry_after_tool_calls(mut self, allow: bool) -> Self { - self.retry_after_tool_calls = allow; - self - } -} - -/// Maps parent graph `State` into the [`SubAgentInput`] a delegated agent runs -/// over. -pub type InputMapper = Arc SubAgentInput + Send + Sync>; - -/// Maps a delegated agent's [`SubAgentOutput`] into the parent graph `Update` -/// merged through the graph reducer. -pub type OutputMapper = Arc Update + Send + Sync>; - -/// A graph node that delegates to a harness agent resolved by name. -/// -/// A `SubAgentNode` binds an agent `ComponentId` (resolved against a -/// capability registry at run time) to a pair of mappers and a -/// [`SubAgentPolicy`]. Lower it into a graph node handler with -/// [`subagent_node`](super::subagent_node). -/// -/// Deprecated: `tinyagents-graph` cannot depend on `tinyagents-orchestration` -/// (the dependency runs the other way), so this node cannot be an adapter over -/// `SubagentDriver`. Run agent steps through `SubagentDriver` instead; it owns -/// spawn admission, result policy, role and the typed outcome. The type keeps -/// its behaviour unchanged. -#[deprecated(note = "use tinyagents_orchestration SubagentDriver")] -pub struct SubAgentNode { - /// The registered agent name to resolve and delegate to. - pub agent: String, - /// Projects parent state into the child input. - pub input_mapper: InputMapper, - /// Folds the child output back into a parent update. - pub output_mapper: OutputMapper, - /// Timeout/retry/budget policy applied around the invocation. - pub policy: SubAgentPolicy, -} - -#[allow(deprecated)] -impl Clone for SubAgentNode { - fn clone(&self) -> Self { - Self { - agent: self.agent.clone(), - input_mapper: self.input_mapper.clone(), - output_mapper: self.output_mapper.clone(), - policy: self.policy.clone(), - } - } -} diff --git a/crates/tinyagents-graph/src/subgraph/README.md b/crates/tinyagents-graph/src/subgraph/README.md deleted file mode 100644 index 970140b52..000000000 --- a/crates/tinyagents-graph/src/subgraph/README.md +++ /dev/null @@ -1,83 +0,0 @@ -# graph::subgraph - -Subgraph node adapters — the graph-level recursion surface where a graph runs -another graph. - -This is the structural counterpart to harness sub-agents (a model calling a -model): here an entire `CompiledGraph` is embedded *as a node* inside a parent -graph, so "graphs that run graphs" is just an ordinary node handler. Each -embedding extends the child's checkpoint namespace with the embedding node id, -which keeps every level of a recursively nested run durable and -collision-free, and the executor's recursion limit bounds how deep that -nesting can go. - -## Embedding modes - -- **`shared_subgraph_node(child)`** — parent and child share the same - `State`/`Update` channel (`Update == State`). The child runs over the - parent's state as passed to the node, and its final state becomes the - parent update. -- **`adapter_subgraph_node(child, to_child, from_child)`** — parent and child - use different state shapes. `to_child: Fn(&P) -> C` projects the parent - state into the child's input; `from_child: Fn(&P, C) -> PU` folds the - child's final state back into a parent update. - -Both wrap a `CompiledGraph` into a node handler usable with -`GraphBuilder::add_node`. - -## Interrupt propagation - -A child that pauses on an interrupt must surface that interrupt to the -parent rather than have its partial state treated as a completed output — -both adapters check `execution.is_interrupted()` and return -`NodeResult::Interrupt(..)` instead of folding the paused child's state -through `from_child` (or returning it directly, for the shared-state case). - -When a parent is resumed through one of `CompiledGraph`'s binding-aware -continuation APIs, both adapters forward that same live -`AgentInvocationBinding` into every resumed child branch. Bindings remain -execution-only data — they are not checkpointed — so a resumed child that -reaches a `SubAgentNode` without one fails closed. - -## Namespace and recursion bookkeeping - -Internal helpers (not part of the public surface, but load-bearing for -correctness): - -- `namespaced(child, ctx)` — clones `child` and extends its checkpoint - namespace with the embedding node id, preventing parent/child checkpoint - collisions when both share a `thread_id`. -- `child_for(child, ctx)` — prepares an embedded child for a run: applies - `namespaced`, seeds it with the enclosing run's live recursion frames (so - the child run extends the parent's recursion tree rather than starting a - fresh one), and records the embedding node so the child's root frame names - it. - -## Public surface - -- `shared_subgraph_node(child: CompiledGraph) -> Handler` -- `adapter_subgraph_node(child, to_child, from_child) -> Handler` - -`types.rs` is documentation-only — the conceptual overview above — because the -adapter constructors return closures rather than named types; there is no -separate public type to document. - -## Files - -| File | Role | -| --- | --- | -| `types.rs` | Documentation-only: the two embedding modes, conceptually. | -| `mod.rs` | `shared_subgraph_node`, `adapter_subgraph_node`, and the namespace/recursion-frame plumbing that makes nested checkpoints and recursion trees work. | -| `mod_tests.rs` | Unit tests (shared vs. adapter state mapping, checkpoint namespace isolation, interrupt propagation, nested recursion limits). | - -## Operational constraints - -- A parent and its embedded subgraph may share a `thread_id`, but they never - share a checkpoint namespace — the embedding always appends the node id. - A caller inspecting checkpoints directly (bypassing the node handler) must - account for this namespace suffix or it will not find the child's records - under the plain parent namespace. -- Deep subgraph nesting is bounded by the executor's recursion limit, which is - seeded from the parent's live recursion frames on each embedding — a graph - that embeds itself (directly or through a cycle of subgraphs) will hit - `TinyAgentsError::RecursionLimit` rather than recurse unbounded. diff --git a/crates/tinyagents-graph/src/subgraph/mod.rs b/crates/tinyagents-graph/src/subgraph/mod.rs deleted file mode 100644 index fd689e256..000000000 --- a/crates/tinyagents-graph/src/subgraph/mod.rs +++ /dev/null @@ -1,322 +0,0 @@ -//! Subgraph node adapters — the graph-level recursion surface where a graph -//! runs another graph. -//! -//! This is the structural counterpart to harness sub-agents (a model calling a -//! model): here an entire [`CompiledGraph`] is embedded *as a node* inside a -//! parent graph, so "graphs that run graphs" is just an ordinary node handler. -//! Each embedding extends the child's checkpoint namespace with the embedding -//! node id, which keeps every level of a recursively -//! nested run durable and collision-free, and the executor's recursion limit -//! bounds how deep that nesting can go. -//! -//! See `types` for the conceptual overview of the two embedding modes. The -//! functions here wrap a [`CompiledGraph`] into a node handler usable with -//! [`crate::GraphBuilder::add_node`]: -//! -//! - [`shared_subgraph_node`] — parent and child share one state channel. -//! - [`adapter_subgraph_node`] — parent and child use different state shapes, -//! bridged by `to_child` / `from_child` mappings. - -mod types; - -use std::future::Future; -use std::pin::Pin; - -use crate::Result; -use crate::builder::NodeContext; -use crate::command::{Command, NodeResult}; -use crate::compiled::{CompiledGraph, GraphExecution}; -use crate::recursion::ChildRun; - -type Handler = Box< - dyn Fn(S, NodeContext) -> Pin>> + Send>> - + Send - + Sync, ->; - -/// Embeds `child` as a shared-state subgraph node. -/// -/// The child uses the same `State` channel as the parent (so `Update == State`) -/// and runs over the parent state passed to the node. Its final state becomes -/// the parent update. The child's checkpoint namespace is extended with the -/// embedding node id. -pub fn shared_subgraph_node(child: CompiledGraph) -> Handler -where - State: Clone + Send + Sync + 'static, -{ - Box::new(move |state: State, ctx: NodeContext| { - let child = child_for(&child, &ctx); - let thread_id = ctx.thread_id.clone(); - let resume = ctx.resume.clone(); - let binding = ctx.agent_binding.clone(); - let recorder = ChildRunRecorder::new(&ctx); - Box::pin(async move { - let execution = drive_child(child, thread_id, state, resume, binding).await?; - recorder.record(&execution); - // A child that paused on an interrupt must surface it to the parent - // rather than have its partial state treated as a completed output. - if execution.is_interrupted() { - return Ok(NodeResult::Interrupt(child_interrupt(execution))); - } - Ok(NodeResult::Update(execution.state)) - }) - }) -} - -/// Embeds `child` as an adapter subgraph node mapping between parent and child -/// state shapes. -/// -/// `to_child` projects the parent state `P` into the child input `C`; -/// `from_child` folds the child's final state back into a parent update `PU`. -pub fn adapter_subgraph_node( - child: CompiledGraph, - to_child: ToChild, - from_child: FromChild, -) -> Handler -where - P: Clone + Send + Sync + 'static, - PU: Send + 'static, - C: Clone + Send + Sync + 'static, - CU: Send + 'static, - ToChild: Fn(&P) -> C + Send + Sync + Clone + 'static, - FromChild: Fn(&P, C) -> PU + Send + Sync + Clone + 'static, -{ - Box::new(move |state: P, ctx: NodeContext| { - let child = child_for(&child, &ctx); - let thread_id = ctx.thread_id.clone(); - let resume = ctx.resume.clone(); - let binding = ctx.agent_binding.clone(); - let recorder = ChildRunRecorder::new(&ctx); - let to_child = to_child.clone(); - let from_child = from_child.clone(); - Box::pin(async move { - let child_input = to_child(&state); - let execution = drive_child(child, thread_id, child_input, resume, binding).await?; - recorder.record(&execution); - // Propagate a child interrupt to the parent instead of folding a - // paused child's partial state through `from_child`. - if execution.is_interrupted() { - return Ok(NodeResult::Interrupt(child_interrupt(execution))); - } - let update = from_child(&state, execution.state); - Ok(NodeResult::Update(update)) - }) - }) -} - -/// 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()); - 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 -/// namespace with the embedding node id (so nested checkpoints never collide), -/// seeds it with the enclosing run's live recursion frames (so the child run -/// extends the parent's recursion tree rather than starting a fresh one), and -/// records the embedding node so the child's root frame names it. -fn child_for(child: &CompiledGraph, ctx: &NodeContext) -> CompiledGraph { - namespaced(child, ctx) - .with_recursion_frames(ctx.recursion_frames.clone()) - .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 -/// resumed activation (the parent was resumed and delivered a value to this -/// node) the child is *resumed* from its own checkpoint with that value, so a -/// subgraph interrupt is reachable through the parent's resume API rather than -/// 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, - state: S, - resume: Option, - binding: Option, -) -> Result> -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 - .run_with_thread_agent_binding(thread_id, state, binding) - .await - } - (None, _, Some(binding)) => child.run_with_agent_binding(state, binding).await, - (Some(thread_id), Some(value), Some(binding)) => { - child - .resume_with_agent_binding(thread_id, Command::resume(value), binding) - .await - } - (Some(thread_id), Some(value), None) => { - child.resume(thread_id, Command::resume(value)).await - } - (Some(thread_id), None, None) => child.run_with_thread(thread_id, state).await, - (None, _, None) => child.run(state).await, - } -} - -/// Extracts the child's paused interrupt so the parent node can re-emit it. -fn child_interrupt(mut execution: GraphExecution) -> crate::command::Interrupt { - execution - .interrupts - .drain(..) - .next() - .expect("an interrupted execution carries at least one interrupt") -} - -/// Captures the enclosing run's child-run sink and lineage so a subgraph node can -/// report the child run it spawned (its distinct run id sharing the parent's -/// root) back to the executor after the embedded graph returns. -struct ChildRunRecorder { - node: tinyagents_harness::ids::NodeId, - sink: Option, -} - -impl ChildRunRecorder { - fn new(ctx: &NodeContext) -> Self { - Self { - node: ctx.node_id.clone(), - sink: ctx.child_runs.clone(), - } - } - - /// Records the embedded run's id (keyed by the embedding node) into the - /// enclosing run's sink, when one is attached. - fn record(&self, execution: &GraphExecution) { - if let Some(sink) = &self.sink { - sink.record(ChildRun { - node: self.node.clone(), - graph_id: execution.graph_id.clone(), - 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(), - }); - } - } -} - -#[cfg(test)] -#[path = "mod_tests.rs"] -mod test; diff --git a/crates/tinyagents-graph/src/subgraph/mod_tests.rs b/crates/tinyagents-graph/src/subgraph/mod_tests.rs deleted file mode 100644 index cc444df01..000000000 --- a/crates/tinyagents-graph/src/subgraph/mod_tests.rs +++ /dev/null @@ -1,1210 +0,0 @@ -//! Unit tests for the subgraph adapters: shared-state embedding, adapter -//! embedding (including folding child output together with parent context), and -//! checkpoint-namespace isolation — verifying that recursively nested and -//! sibling embeddings accumulate distinct namespaces and never collide on -//! checkpoint ids when sharing one checkpointer and thread. -// The deprecated `SubAgentNode` keeps its behaviour and its coverage. -#![allow(deprecated)] - -use std::collections::HashSet; -use std::sync::Arc; - -use super::*; -use crate::builder::{GraphBuilder, NodeContext}; -use crate::checkpoint::{Checkpointer, InMemoryCheckpointer}; -use crate::command::NodeResult; -use crate::reducer::ClosureStateReducer; -use async_trait::async_trait; -use tinyagents_harness::ids::{NodeId, RunId}; - -#[derive(Default)] -struct NestedRecordingInvoker(std::sync::Mutex>); - -#[async_trait] -impl crate::subagent_node::AgentInvoker for NestedRecordingInvoker { - async fn invoke( - &self, - request: crate::subagent_node::AgentInvocation, - ) -> crate::Result { - self.0.lock().unwrap().push(request.clone()); - request - .events - .emit(tinyagents_harness::events::AgentEvent::StateUpdate); - Ok(crate::subagent_node::SubAgentOutput { - text: request.input.prompt, - ..Default::default() - }) - } -} - -struct NestedFailingInvoker; - -#[async_trait] -impl crate::subagent_node::AgentInvoker for NestedFailingInvoker { - async fn invoke( - &self, - _request: crate::subagent_node::AgentInvocation, - ) -> crate::Result { - Err(crate::TinyAgentsError::Model( - "nested continuation failure".to_string(), - )) - } -} - -fn delegating_child( - checkpointer: Arc>, -) -> CompiledGraph { - GraphBuilder::::overwrite() - .add_node( - "delegate", - crate::subagent_node::subagent_node(crate::subagent_node::SubAgentNode::from_fns( - "researcher", - |state: &String| crate::subagent_node::SubAgentInput::prompt(state.clone()), - |output: crate::subagent_node::SubAgentOutput| output.text, - )), - ) - .set_entry("delegate") - .set_finish("delegate") - .compile() - .unwrap() - .with_checkpointer(checkpointer) -} - -fn nested_binding( - invoker: Arc, -) -> crate::subagent_node::AgentInvocationBinding { - crate::subagent_node::AgentInvocationBinding::new( - invoker, - tinyagents_harness::events::EventSink::new(), - tinyagents_harness::cancel::CancellationToken::new(), - ) -} - -/// 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), - run_id: RunId::new("run-test"), - thread_id: None, - step: 1, - resume: None, - fork: None, - send_arg: None, - root_run_id: None, - 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(), - } -} - -/// A small child graph (shared state `i32`) that adds 10. -fn child_add_ten() -> CompiledGraph { - GraphBuilder::::overwrite() - .add_node("add", |s: i32, _c: NodeContext| async move { - Ok(NodeResult::Update(s + 10)) - }) - .set_entry("add") - .set_finish("add") - .compile() - .unwrap() -} - -#[tokio::test] -async fn shared_state_subgraph() { - let child = child_add_ten(); - let parent = GraphBuilder::::overwrite() - .add_node("pre", |s: i32, _c: NodeContext| async move { - Ok(NodeResult::Update(s + 1)) - }) - .add_node("child", shared_subgraph_node(child)) - .set_entry("pre") - .add_edge("pre", "child") - .set_finish("child") - .compile() - .unwrap(); - - // 0 -> pre(+1) -> child(+10) = 11 - let run = parent.run(0).await.unwrap(); - assert_eq!(run.state, 11); -} - -#[tokio::test] -async fn embedded_graph_propagates_the_host_agent_invoker() { - let invoker = Arc::new(NestedRecordingInvoker::default()); - let child = GraphBuilder::::overwrite() - .add_node( - "delegate", - crate::subagent_node::subagent_node(crate::subagent_node::SubAgentNode::from_fns( - "researcher", - |state: &String| crate::subagent_node::SubAgentInput::prompt(state.clone()), - |out: crate::subagent_node::SubAgentOutput| out.text, - )), - ) - .set_entry("delegate") - .set_finish("delegate") - .compile() - .unwrap(); - let parent = GraphBuilder::::overwrite() - .add_node("child", shared_subgraph_node(child)) - .set_entry("child") - .set_finish("child") - .compile() - .unwrap(); - - let run = parent - .run_with_agent_binding( - "nested".to_string(), - crate::subagent_node::AgentInvocationBinding::new( - invoker.clone(), - tinyagents_harness::events::EventSink::new(), - tinyagents_harness::cancel::CancellationToken::new(), - ), - ) - .await - .unwrap(); - let requests = invoker.0.lock().unwrap(); - assert_eq!(requests.len(), 1); - assert_eq!(requests[0].root_run_id, run.root_run_id); - assert_ne!(requests[0].parent_run_id, run.run_id); - assert_eq!(requests[0].input.prompt, "nested"); -} - -#[derive(Clone, Debug, PartialEq)] -struct ParentState { - name: String, - score: i32, -} - -#[tokio::test] -async fn adapter_subgraph_maps_state() { - // child works on a bare i32 score - let child = child_add_ten(); - - let parent = GraphBuilder::::new() - .set_reducer(ClosureStateReducer::new(|_old, new: ParentState| Ok(new))) - .add_node( - "score", - adapter_subgraph_node( - child, - // project parent -> child input - |p: &ParentState| p.score, - // fold child output -> parent update - |p: &ParentState, child_score: i32| ParentState { - name: p.name.clone(), - score: child_score, - }, - ), - ) - .set_entry("score") - .set_finish("score") - .compile() - .unwrap(); - - let run = parent - .run(ParentState { - name: "alice".to_string(), - score: 5, - }) - .await - .unwrap(); - assert_eq!(run.state.name, "alice"); - assert_eq!(run.state.score, 15); -} - -#[tokio::test] -async fn adapter_folds_child_output_with_parent_context() { - // `from_child` receives BOTH the original parent state and the child output, - // so it can combine them rather than just replacing parent fields. - let child = child_add_ten(); - - let parent = GraphBuilder::::new() - .set_reducer(ClosureStateReducer::new(|_old, new: ParentState| Ok(new))) - .add_node( - "score", - adapter_subgraph_node( - child, - |p: &ParentState| p.score, - |p: &ParentState, child_score: i32| ParentState { - name: format!("{}-scored", p.name), - // combine parent's own score (5) with the child output (15). - score: p.score + child_score, - }, - ), - ) - .set_entry("score") - .set_finish("score") - .compile() - .unwrap(); - - let run = parent - .run(ParentState { - name: "bob".to_string(), - score: 5, - }) - .await - .unwrap(); - // child: 5 + 10 = 15; from_child uses both args: 5 + 15 = 20. - assert_eq!(run.state.score, 20); - assert_eq!(run.state.name, "bob-scored"); -} - -#[test] -fn namespaced_clone_appends_embedding_node_id() { - let child = child_add_ten(); - assert!(child.namespace().is_empty()); - let scoped = namespaced(&child, &ctx_for("embed")); - assert_eq!(scoped.namespace(), &["embed".to_string()]); -} - -#[test] -fn nested_namespaces_accumulate_and_stay_distinct() { - let child = child_add_ten(); - // Two levels of embedding accumulate node ids in order. - let outer = namespaced(&child, &ctx_for("outer")); - let inner = namespaced(&outer, &ctx_for("inner")); - assert_eq!( - inner.namespace(), - &["outer".to_string(), "inner".to_string()] - ); - // Siblings embedded under different node ids get distinct namespaces, so - // their checkpoints can never collide. - let sibling = namespaced(&child, &ctx_for("other")); - assert_ne!(outer.namespace(), sibling.namespace()); -} - -#[tokio::test] -async fn namespaced_children_persist_under_isolated_namespaces() { - // A single compiled child embedded under two different node ids, sharing one - // checkpointer and thread: every checkpoint is tagged with its own namespace - // and keeps a globally-unique id, so the two embeddings never collide. - let ckpt = Arc::new(InMemoryCheckpointer::::new()); - let child = child_add_ten().with_checkpointer(ckpt.clone()); - - let branch_a = namespaced(&child, &ctx_for("branch_a")); - let branch_b = namespaced(&child, &ctx_for("branch_b")); - - branch_a.run_with_thread("t", 0).await.unwrap(); - branch_b.run_with_thread("t", 1).await.unwrap(); - - let list = ckpt.list("t").await.unwrap(); - assert_eq!(list.len(), 2); - - // Checkpoint ids are unique (no collision). - let ids: HashSet<&str> = list.iter().map(|m| m.checkpoint_id.as_str()).collect(); - assert_eq!(ids.len(), 2); - - // Each embedding's checkpoint carries its own namespace. - assert!( - list.iter() - .any(|m| m.namespace == vec!["branch_a".to_string()]) - ); - assert!( - list.iter() - .any(|m| m.namespace == vec!["branch_b".to_string()]) - ); -} - -#[tokio::test] -async fn embedded_child_persists_under_parent_thread_and_child_namespace() { - let ckpt = Arc::new(InMemoryCheckpointer::::new()); - let child = child_add_ten().with_checkpointer(ckpt.clone()); - let parent = GraphBuilder::::overwrite() - .add_node("child", shared_subgraph_node(child.clone())) - .set_entry("child") - .set_finish("child") - .compile() - .unwrap() - .with_checkpointer(ckpt.clone()); - - let run = parent.run_with_thread("t", 0).await.unwrap(); - assert_eq!(run.state, 10); - - let list = ckpt.list("t").await.unwrap(); - assert_eq!(list.len(), 2); - assert!(list.iter().any(|m| m.namespace.is_empty())); - let child_meta = list - .iter() - .find(|m| m.namespace == vec!["child".to_string()]) - .expect("child checkpoint is stored under the embedding namespace"); - - let child_scoped = namespaced(&child, &ctx_for("child")); - let child_state = child_scoped - .get_state("t", Some(&child_meta.checkpoint_id)) - .await - .unwrap() - .expect("child checkpoint can be loaded from the parent thread"); - assert_eq!(child_state.values, 10); - assert_eq!(child_state.next_nodes, Vec::::new()); -} - -/// A child graph that pauses on an interrupt until resumed. -fn child_interrupts() -> CompiledGraph { - GraphBuilder::::overwrite() - .add_node("gate", |s: i32, c: NodeContext| async move { - match c.resume { - Some(_) => Ok(NodeResult::Update(s + 10)), - None => Ok(NodeResult::Interrupt(crate::command::Interrupt::new( - "gate", - serde_json::json!({ "ask": "ok?" }), - ))), - } - }) - .set_entry("gate") - .set_finish("gate") - .compile() - .unwrap() -} - -#[tokio::test] -async fn subgraph_interrupt_propagates_to_parent() { - // A child graph that interrupts must pause the *parent* run — not be - // swallowed as a completed output with the child's partial state. - let ckpt = Arc::new(InMemoryCheckpointer::::new()); - let child = child_interrupts().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 run = parent.run_with_thread("t", 0).await.unwrap(); - assert!( - run.is_interrupted(), - "a child interrupt must pause the parent instead of being swallowed" - ); - assert_eq!(run.interrupts.len(), 1, "the child's interrupt is surfaced"); - assert_eq!( - run.interrupts[0].node.as_str(), - "gate", - "the propagated interrupt names the child's paused node" - ); -} - -#[tokio::test] -async fn subgraph_interrupt_resumes_child_from_its_own_checkpoint() { - // Resuming the parent must resume the *child* from its namespaced - // checkpoint (not re-run it and not load the parent's checkpoint), so the - // whole nested run completes. Exercises namespace-aware resume end to end. - let ckpt = Arc::new(InMemoryCheckpointer::::new()); - let child = child_interrupts().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 paused = parent.run_with_thread("t", 0).await.unwrap(); - assert!(paused.is_interrupted()); - - let done = parent - .resume( - "t", - crate::command::Command::resume(serde_json::json!("go")), - ) - .await - .unwrap(); - assert!( - !done.is_interrupted(), - "resuming the parent must drive the child to completion" - ); - // Child added 10 to the initial 0. - assert_eq!(done.state, 10); -} - -#[tokio::test] -async fn resumed_subgraph_passes_the_supplied_binding_to_its_subagent() { - // The child pauses before its SubAgentNode. The parent resume must route - // the fresh binding into the child resume (rather than dropping it at the - // subgraph boundary), where it reaches the durable continuation. - let ckpt = Arc::new(InMemoryCheckpointer::::new()); - let child = GraphBuilder::::overwrite() - .add_node("gate", |state: String, ctx: NodeContext| async move { - if ctx.resume.is_some() { - Ok(NodeResult::Update(state)) - } else { - Ok(NodeResult::Interrupt(crate::command::Interrupt::new( - "gate", - serde_json::json!({ "ask": "continue?" }), - ))) - } - }) - .add_node( - "delegate", - crate::subagent_node::subagent_node(crate::subagent_node::SubAgentNode::from_fns( - "researcher", - |state: &String| crate::subagent_node::SubAgentInput::prompt(state.clone()), - |output: crate::subagent_node::SubAgentOutput| output.text, - )), - ) - .set_entry("gate") - .add_edge("gate", "delegate") - .set_finish("delegate") - .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); - - let paused = parent - .run_with_thread("nested-resume", "question".to_string()) - .await - .unwrap(); - assert!(paused.is_interrupted()); - - let invoker = Arc::new(NestedRecordingInvoker::default()); - let events = tinyagents_harness::events::EventSink::new(); - let listener = Arc::new(tinyagents_harness::events::RecordingListener::new()); - events.subscribe(listener.clone()); - let cancellation = tinyagents_harness::cancel::CancellationToken::new(); - cancellation.cancel(); - let resumed = parent - .resume_with_agent_binding( - "nested-resume", - crate::command::Command::resume(serde_json::json!("go")), - crate::subagent_node::AgentInvocationBinding::new( - invoker.clone(), - events, - cancellation, - ), - ) - .await - .unwrap(); - - assert_eq!(resumed.state, "question"); - assert_eq!(resumed.child_runs.len(), 1); - let request = invoker.0.lock().unwrap().pop().expect("child delegated"); - assert_eq!(request.parent_run_id, resumed.child_runs[0].run_id); - assert_eq!(request.root_run_id, resumed.root_run_id); - assert!( - request - .cancellation - .expect("binding has cancellation") - .is_cancelled() - ); - assert_eq!( - listener.len(), - 1, - "sub-agent event used the resumed binding sink" - ); -} - -#[tokio::test] -async fn resumed_adapter_subgraph_passes_the_supplied_binding_to_its_subagent() { - // Exercise the adapter route separately: it has its own child-driving - // closure, so a shared-state test alone cannot prove a bound resume is not - // dropped while mapping parent state into child state and back. - let ckpt = Arc::new(InMemoryCheckpointer::::new()); - let child = GraphBuilder::::overwrite() - .add_node("gate", |state: String, ctx: NodeContext| async move { - if ctx.resume.is_some() { - Ok(NodeResult::Update(state)) - } else { - Ok(NodeResult::Interrupt(crate::command::Interrupt::new( - "gate", - serde_json::json!({ "ask": "continue?" }), - ))) - } - }) - .add_node( - "delegate", - crate::subagent_node::subagent_node(crate::subagent_node::SubAgentNode::from_fns( - "researcher", - |state: &String| crate::subagent_node::SubAgentInput::prompt(state.clone()), - |output: crate::subagent_node::SubAgentOutput| output.text, - )), - ) - .set_entry("gate") - .add_edge("gate", "delegate") - .set_finish("delegate") - .compile() - .unwrap() - .with_checkpointer(ckpt.clone()); - let parent = GraphBuilder::::overwrite() - .add_node( - "child", - adapter_subgraph_node( - child, - |state: &String| state.clone(), - |_parent, child| child, - ), - ) - .set_entry("child") - .set_finish("child") - .compile() - .unwrap() - .with_checkpointer(ckpt); - - assert!( - parent - .run_with_thread("adapter-resume", "question".to_string()) - .await - .unwrap() - .is_interrupted() - ); - - let invoker = Arc::new(NestedRecordingInvoker::default()); - let events = tinyagents_harness::events::EventSink::new(); - let listener = Arc::new(tinyagents_harness::events::RecordingListener::new()); - events.subscribe(listener.clone()); - let cancellation = tinyagents_harness::cancel::CancellationToken::new(); - cancellation.cancel(); - let resumed = parent - .resume_with_agent_binding( - "adapter-resume", - crate::command::Command::resume(serde_json::json!("go")), - crate::subagent_node::AgentInvocationBinding::new( - invoker.clone(), - events, - cancellation, - ), - ) - .await - .unwrap(); - - assert_eq!(resumed.state, "question"); - let request = invoker.0.lock().unwrap().pop().expect("child delegated"); - assert_eq!(request.parent_run_id, resumed.child_runs[0].run_id); - assert_eq!(request.root_run_id, resumed.root_run_id); - assert!(request.cancellation.unwrap().is_cancelled()); - assert_eq!(listener.len(), 1); -} - -#[tokio::test] -async fn binding_reaches_a_grandchild_that_itself_resumes() { - // The grandchild interrupts first. Resuming the root must therefore carry - // one binding through *two* resumed drive_child branches, preserving the - // same event sink and cancellation token at the eventual SubAgentNode. - let ckpt = Arc::new(InMemoryCheckpointer::::new()); - let grandchild = GraphBuilder::::overwrite() - .add_node("gate", |state: String, ctx: NodeContext| async move { - if ctx.resume.is_some() { - Ok(NodeResult::Update(state)) - } else { - Ok(NodeResult::Interrupt(crate::command::Interrupt::new( - "gate", - serde_json::json!({ "ask": "continue?" }), - ))) - } - }) - .add_node( - "delegate", - crate::subagent_node::subagent_node(crate::subagent_node::SubAgentNode::from_fns( - "researcher", - |state: &String| crate::subagent_node::SubAgentInput::prompt(state.clone()), - |output: crate::subagent_node::SubAgentOutput| output.text, - )), - ) - .set_entry("gate") - .add_edge("gate", "delegate") - .set_finish("delegate") - .compile() - .unwrap() - .with_checkpointer(ckpt.clone()); - let child = GraphBuilder::::overwrite() - .add_node("grandchild", shared_subgraph_node(grandchild)) - .set_entry("grandchild") - .set_finish("grandchild") - .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); - - assert!( - parent - .run_with_thread("grandchild-resume", "question".to_string()) - .await - .unwrap() - .is_interrupted() - ); - - let invoker = Arc::new(NestedRecordingInvoker::default()); - let events = tinyagents_harness::events::EventSink::new(); - let listener = Arc::new(tinyagents_harness::events::RecordingListener::new()); - events.subscribe(listener.clone()); - let cancellation = tinyagents_harness::cancel::CancellationToken::new(); - cancellation.cancel(); - let resumed = parent - .resume_with_agent_binding( - "grandchild-resume", - crate::command::Command::resume(serde_json::json!("go")), - crate::subagent_node::AgentInvocationBinding::new( - invoker.clone(), - events, - cancellation, - ), - ) - .await - .unwrap(); - - assert_eq!(resumed.state, "question"); - let request = invoker - .0 - .lock() - .unwrap() - .pop() - .expect("grandchild delegated"); - assert_eq!(request.root_run_id, resumed.root_run_id); - assert!(request.cancellation.unwrap().is_cancelled()); - assert_eq!(listener.len(), 1); -} - -#[tokio::test] -async fn retrying_shared_subgraph_passes_fresh_binding_to_failed_child() { - // A failed child leaves the parent node pending. Retrying that parent with - // a replacement binding drives the shared-state child through the fresh - // threaded binding branch, rather than reviving the failed invoker. - let ckpt = Arc::new(InMemoryCheckpointer::::new()); - let child = delegating_child(ckpt.clone()); - let parent = GraphBuilder::::overwrite() - .add_node("child", shared_subgraph_node(child)) - .set_entry("child") - .set_finish("child") - .compile() - .unwrap() - .with_checkpointer(ckpt); - - let failed = parent - .run_with_thread_agent_binding( - "shared-retry", - "question".to_string(), - nested_binding(Arc::new(NestedFailingInvoker)), - ) - .await - .unwrap_err(); - assert!(matches!(failed, crate::TinyAgentsError::Model(_))); - - let replacement = Arc::new(NestedRecordingInvoker::default()); - let retried = parent - .retry_with_agent_binding("shared-retry", nested_binding(replacement.clone())) - .await - .unwrap(); - - assert_eq!(retried.state, "question"); - let request = replacement - .0 - .lock() - .unwrap() - .pop() - .expect("shared child used the retry binding"); - assert_eq!(request.parent_run_id, retried.child_runs[0].run_id); - assert_eq!(request.root_run_id, retried.root_run_id); -} - -#[tokio::test] -async fn retrying_adapter_subgraph_passes_fresh_binding_to_failed_child() { - // This covers the adapter's distinct child-driving closure. Mapping parent - // state into and out of the child must not drop the retry binding. - let ckpt = Arc::new(InMemoryCheckpointer::::new()); - let child = delegating_child(ckpt.clone()); - let parent = GraphBuilder::::overwrite() - .add_node( - "child", - adapter_subgraph_node( - child, - |state: &String| state.clone(), - |_parent, child| child, - ), - ) - .set_entry("child") - .set_finish("child") - .compile() - .unwrap() - .with_checkpointer(ckpt); - - let failed = parent - .run_with_thread_agent_binding( - "adapter-retry", - "question".to_string(), - nested_binding(Arc::new(NestedFailingInvoker)), - ) - .await - .unwrap_err(); - assert!(matches!(failed, crate::TinyAgentsError::Model(_))); - - let replacement = Arc::new(NestedRecordingInvoker::default()); - let retried = parent - .retry_with_agent_binding("adapter-retry", nested_binding(replacement.clone())) - .await - .unwrap(); - - assert_eq!(retried.state, "question"); - let request = replacement - .0 - .lock() - .unwrap() - .pop() - .expect("adapter child used the retry binding"); - assert_eq!(request.parent_run_id, retried.child_runs[0].run_id); - assert_eq!(request.root_run_id, retried.root_run_id); -} - -#[tokio::test] -async fn retrying_nested_shared_subgraphs_rebinds_the_failed_grandchild() { - // Retrying from the root traverses two independent shared-subgraph - // drive_child calls. The grandchild failure must be repaired solely by the - // fresh root binding, with no capability retained in either checkpoint. - let ckpt = Arc::new(InMemoryCheckpointer::::new()); - let grandchild = delegating_child(ckpt.clone()); - let child = GraphBuilder::::overwrite() - .add_node("grandchild", shared_subgraph_node(grandchild)) - .set_entry("grandchild") - .set_finish("grandchild") - .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); - - let failed = parent - .run_with_thread_agent_binding( - "grandchild-retry", - "question".to_string(), - nested_binding(Arc::new(NestedFailingInvoker)), - ) - .await - .unwrap_err(); - assert!(matches!(failed, crate::TinyAgentsError::Model(_))); - - let replacement = Arc::new(NestedRecordingInvoker::default()); - let retried = parent - .retry_with_agent_binding("grandchild-retry", nested_binding(replacement.clone())) - .await - .unwrap(); - - assert_eq!(retried.state, "question"); - let request = replacement - .0 - .lock() - .unwrap() - .pop() - .expect("grandchild used the retry binding"); - assert_eq!(request.root_run_id, retried.root_run_id); - assert_eq!(retried.child_runs.len(), 1); -} - -#[tokio::test] -async fn subgraph_child_run_distinct_and_shares_root() { - // A parent embedding one child: the parent run records exactly one child run - // whose run id differs from the parent's, and whose root run id equals the - // parent's (the child preserves the root of the recursion tree). - let child = child_add_ten(); - let parent = GraphBuilder::::overwrite() - .add_node("child", shared_subgraph_node(child)) - .set_entry("child") - .set_finish("child") - .compile() - .unwrap(); - - let run = parent.run(0).await.unwrap(); - assert_eq!(run.state, 10); - - // Exactly one child run, keyed by the embedding node. - assert_eq!(run.child_runs.len(), 1); - let child_run = &run.child_runs[0]; - assert_eq!(child_run.node.as_str(), "child"); - // Distinct child run id, shared root. - assert_ne!(child_run.run_id, run.run_id); - assert_eq!(child_run.root_run_id, run.run_id); - assert_eq!(run.root_run_id, run.run_id); - assert!(run.parent_run_id.is_none()); - - // The run tree mirrors the execution's lineage. - let tree = run.run_tree(); - assert!(tree.is_root()); - assert_eq!(tree.children.len(), 1); - assert_eq!(tree.children[0].run_id, child_run.run_id); -} - -#[tokio::test] -async fn nested_subgraphs_produce_distinct_ids_sharing_one_root() { - // grandchild adds 10; child embeds grandchild then is itself embedded in the - // parent, so the run tree is three deep: parent -> child -> grandchild. - let grandchild = child_add_ten(); - let child = GraphBuilder::::overwrite() - .add_node("grandchild", shared_subgraph_node(grandchild)) - .set_entry("grandchild") - .set_finish("grandchild") - .compile() - .unwrap(); - let parent = GraphBuilder::::overwrite() - .add_node("child", shared_subgraph_node(child)) - .set_entry("child") - .set_finish("child") - .compile() - .unwrap(); - - let run = parent.run(0).await.unwrap(); - // 0 -> child -> grandchild(+10) = 10 - assert_eq!(run.state, 10); - - // Parent records the child run; that child's run shares the parent's root. - assert_eq!(run.child_runs.len(), 1); - let child_run = &run.child_runs[0]; - assert_eq!(child_run.node.as_str(), "child"); - assert_eq!(child_run.root_run_id, run.run_id); - - // All three run ids are distinct (parent, child, grandchild). The grandchild - // is recorded on the child run's own child_runs, but the top-level parent - // sees only its direct child; we assert direct-child distinctness here. - assert_ne!(child_run.run_id, run.run_id); -} - -#[tokio::test] -async fn parent_frames_balanced_after_subgraph_returns() { - // Two sibling subgraph nodes run in sequence under one parent. Because each - // child runs on its own seeded recursion stack, the parent's depth is never - // mutated by a child returning, so both child runs see depth-consistent - // lineage: each is a direct child of the parent and shares its root. - let parent = GraphBuilder::::overwrite() - .add_node("a", shared_subgraph_node(child_add_ten())) - .add_node("b", shared_subgraph_node(child_add_ten())) - .set_entry("a") - .add_edge("a", "b") - .set_finish("b") - .compile() - .unwrap(); - - let run = parent.run(0).await.unwrap(); - // 0 -> a(+10) -> b(+10) = 20 - assert_eq!(run.state, 20); - - assert_eq!(run.child_runs.len(), 2); - let nodes: Vec<&str> = run.child_runs.iter().map(|c| c.node.as_str()).collect(); - assert_eq!(nodes, vec!["a", "b"]); - // Both children share the parent's root and have distinct run ids. - for c in &run.child_runs { - assert_eq!(c.root_run_id, run.run_id); - assert_ne!(c.run_id, run.run_id); - } - assert_ne!(run.child_runs[0].run_id, run.child_runs[1].run_id); -} - -#[tokio::test] -async fn child_runs_recorded_in_checkpoint_metadata() { - // With a checkpointer + thread, the boundary checkpoint that committed the - // subgraph node carries a `child_runs` array in its metadata keyed by node. - let ckpt = Arc::new(InMemoryCheckpointer::::new()); - let parent = GraphBuilder::::overwrite() - .add_node("child", shared_subgraph_node(child_add_ten())) - .set_entry("child") - .set_finish("child") - .compile() - .unwrap() - .with_checkpointer(ckpt.clone()); - - let run = parent.run_with_thread("t", 0).await.unwrap(); - assert_eq!(run.child_runs.len(), 1); - - // Walk every persisted checkpoint's raw metadata for the `child_runs` array - // naming the embedding node. - let list = ckpt.list("t").await.unwrap(); - let mut found = false; - for meta in &list { - let checkpoint = ckpt - .get("t", Some(&meta.checkpoint_id)) - .await - .unwrap() - .unwrap(); - if checkpoint - .metadata - .get("child_runs") - .and_then(|v| v.as_array()) - .is_some_and(|arr| { - arr.iter() - .any(|c| c.get("node").and_then(|n| n.as_str()) == Some("child")) - }) - { - found = true; - break; - } - } - 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/subgraph/types.rs b/crates/tinyagents-graph/src/subgraph/types.rs deleted file mode 100644 index 21f8419b6..000000000 --- a/crates/tinyagents-graph/src/subgraph/types.rs +++ /dev/null @@ -1,20 +0,0 @@ -//! Subgraph node adapters — conceptual overview of the graph-level recursion -//! modes. -//! -//! Embedding a [`crate::CompiledGraph`] as a node is how this runtime -//! expresses recursion at the graph layer (the analogue of a model calling a -//! model in the harness). A child graph can be embedded in a parent in two -//! modes: -//! -//! - **shared-state**: parent and child share the same `State`/`Update` -//! channel. The child runs over the parent's state and its final state is -//! returned as the parent update. -//! - **adapter**: parent and child use different state shapes. A `to_child` -//! mapping projects parent state into the child input, and a `from_child` -//! mapping folds the child's final state back into a parent update. -//! -//! Both adapters append the embedding node id to the child's checkpoint -//! namespace so parent and child checkpoint ids never collide. - -// This module is documentation-only for types; the adapter constructors live in -// `mod.rs` because they return closures rather than named types. diff --git a/crates/tinyagents-graph/src/test_support.rs b/crates/tinyagents-graph/src/test_support.rs deleted file mode 100644 index c281bcf7e..000000000 --- a/crates/tinyagents-graph/src/test_support.rs +++ /dev/null @@ -1,14 +0,0 @@ -//! Shared testkit fixtures for storage-port conformance and reopen tests. - -use tinyagents_storage_kit::{Scope, StorageBackend}; -use tinystoragedrivers_testkit::MemoryBackend; - -use crate::checkpoint::DriverCheckpointer; - -/// A checkpointer over the explicit testkit backend owned by one test. -pub(crate) fn dir_checkpointer(backend: &MemoryBackend) -> DriverCheckpointer { - let scoped = backend - .for_scope(&Scope::new("local").unwrap()) - .expect("local scope on a single-operator database"); - DriverCheckpointer::new(std::sync::Arc::clone(scoped.documents())) -} diff --git a/crates/tinyagents-graph/src/testkit/README.md b/crates/tinyagents-graph/src/testkit/README.md deleted file mode 100644 index f435d648a..000000000 --- a/crates/tinyagents-graph/src/testkit/README.md +++ /dev/null @@ -1,106 +0,0 @@ -# graph::testkit - -Graph-test building blocks — deterministic node doubles, an event recorder, a -stream projector, a fluent run-assertion builder, and storage conformance -suites for the durable graph runtime. - -This is the *graph-level* counterpart to `harness::testkit` (model/tool -doubles + trajectories): here the units under test are graph nodes and -supersteps, so the doubles are node handlers and the assertions read a run's -export/event/checkpoint truth. Because a node can recurse into a -`graph::subgraph` or a `graph::subagent_node`, the same recorder that observes -a top-level run also captures the events and child-run rollups of the nested -runs it spawns, so recursion stays observable in tests. - -## Node doubles - -Each returns a closure ready for `GraphBuilder::add_node`: - -| Helper | Behavior | -|--------|----------| -| `noop_node` | Routes onward with no state update | -| `scripted_update_node` | Emits queued updates (saturating the last) | -| `scripted_route_node` | Emits queued `goto` route-sets | -| `fanout_node` | Emits one `Send` per arg (fanout) | -| `failing_node` | Always returns an error | -| `RetryCountingNode` | Counts activations, fails the first N | -| `interrupting_node` | Interrupts until resumed, then updates | -| `subgraph_test_node` | Embeds a child graph (shared state) | -| `subagent_fake_node` | Records a child run + updates (fake sub-agent) | - -## Observation & assertions - -- `GraphEventRecorder` — captures the `GraphEvent` stream from a run (`.sink()` - to wire it, `.events()` / `.kinds()` to read back, `.collector()` to - project). -- `StreamCollector` — projects a recorded event list into test-friendly views: - `node_order()`, `updates()`, `routes()`, `interrupts()`, - `checkpoint_count()`, `custom()`. -- `GraphRun` — bundles a `GraphExecution` with its recorded - events and checkpoint history; built with `::new(execution)` + - `.with_events(..)` / `.with_history(..)`. The single `GraphRun` is the test - truth — execution, events, and checkpoints all read from it. -- `run_recorded(..)` — runs a graph with the recorder wired and returns the - bundled `GraphRun`. -- `assert_graph(&run) -> GraphAssertions<'_, State>` — opens a fluent - assertion builder: `.visited(..)`, `.routed(from, to)`, - `.checkpoint_count(n)`, `.state_history(f)`, `.checkpoint(f)`, - `.completed()`, `.interrupted()` — each panics with a descriptive message on - failure and returns `&Self` for chaining. - -## Storage conformance suites (`conformance.rs`) - -Durable graph stores are hard to migrate safely without a shared contract: -two backends that both implement a trait should behave identically. These -functions encode that contract once so any backend — built-in or a -caller-supplied adapter — can be certified by running the same assertions: - -- `checkpointer_contract(cp)` — basic `Checkpointer` contract: put/get - (latest and specific), duplicate-id resolution, insertion-order listing, - unknown-thread misses, namespace-scoped lookup, `list_threads`, - `delete_thread`, and `prune`. -- `checkpointer_concurrent_contract(cp: Arc)` — concurrent-write contract: - many tasks putting distinct checkpoints on one shared instance must not lose - writes. -- `checkpointer_writes_contract(cp)` — pending-writes ledger contract: data - writes are idempotent by `(task_id, idx)`, control-plane writes upsert, - writes are namespace-isolated, and `delete_checkpoints`/`delete_thread` - remove the writes they own. -- `checkpointer_lineage_contract(cp)` — lineage and thread-operation contract: - `get_tuple` addressing (including parent config), `state_history` - ordering/limit/cycle-termination, `copy_thread` semantics (including - rejecting a non-empty target), and per-namespace `prune`. -- `taskstore_contract(store)` — basic CRUD/lifecycle contract for a - `tinyagents_tasks::TaskStore` implementation. -- `taskstore_concurrent_contract(store: Arc)` — concurrent-access contract - (no lost updates, no corrupted records under concurrent writers). -- `taskstore_replay_contract(reopen)` — replay contract: a store reopened via - `reopen` must replay to the same state. A test that uses memory should pass - another handle to the same backend; it verifies driver reopening within the - test, not persistence across a process restart. - -Each function **panics** with a descriptive message on the first violation — -call them from a `#[tokio::test]` / `#[test]` in the crate implementing a new -backend. - -## Files - -| File | Role | -| --- | --- | -| `types.rs` | `GraphEventRecorder`, `StreamCollector`, `GraphRun`, `GraphAssertions`, `RetryCountingNode`. | -| `mod.rs` | Node doubles, `run_recorded`, `assert_graph`. | -| `conformance.rs` | `checkpointer_contract`, `checkpointer_concurrent_contract`, `checkpointer_writes_contract`, `checkpointer_lineage_contract`, `taskstore_contract`, `taskstore_concurrent_contract`, `taskstore_replay_contract`. | -| `mod_tests.rs` | Unit tests for the testkit itself (each double, recorder, assertion). | - -## Operational constraints - -- These helpers are intended for tests; none of them are optimized for - production workloads (e.g. `GraphEventRecorder` buffers every event - unbounded in memory). -- `RetryCountingNode` and `interrupting_node` hold interior mutable counters — - construct a fresh instance per test rather than sharing one across - independent test cases, or activation counts will leak between them. -- Conformance suites assume the store under test starts empty (or, for - `taskstore_replay_contract`, that `reopen` yields an equivalent fresh - handle to the same underlying storage) — seed/clean state accordingly in - the calling test. diff --git a/crates/tinyagents-graph/src/testkit/conformance.rs b/crates/tinyagents-graph/src/testkit/conformance.rs deleted file mode 100644 index 4e45b54eb..000000000 --- a/crates/tinyagents-graph/src/testkit/conformance.rs +++ /dev/null @@ -1,821 +0,0 @@ -//! Reusable storage **conformance** (contract) suites. -//! -//! Durable graph stores are hard to migrate safely without a shared contract: -//! two backends that both implement a trait should behave identically. These -//! functions encode that contract once so any backend — the built-in ones or a -//! caller-supplied adapter — can be certified by running the same assertions. -//! -//! Each function panics with a descriptive message on the first violation, so -//! call them from a `#[tokio::test]` / `#[test]`. - -use crate::checkpoint::{Checkpoint, Checkpointer, PendingActivation}; -use crate::orchestration::{ - OrchestrationTaskFilter, OrchestrationTaskKind, OrchestrationTaskResult, - OrchestrationTaskStatus, TaskStore, -}; -use tinyagents_harness::ids::{NodeId, TaskId}; - -fn contract_checkpoint( - thread: &str, - id: &str, - parent: Option<&str>, - step: usize, -) -> Checkpoint { - 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`. -/// -/// Covers put/get (latest and specific), insertion-order listing, unknown-thread -/// misses, `list_threads`, `delete_thread`, and `prune`. Any backend that -/// passes this behaves interchangeably for the durable graph runtime. -pub async fn checkpointer_contract(cp: C) -where - C: Checkpointer, -{ - cp.put(contract_checkpoint("t1", "c1", None, 1)) - .await - .expect("put c1"); - cp.put(contract_checkpoint("t1", "c2", Some("c1"), 2)) - .await - .expect("put c2"); - - // Latest wins. - let latest = cp.get("t1", None).await.expect("get latest").expect("some"); - assert_eq!(latest.checkpoint_id, "c2", "latest checkpoint id"); - assert_eq!(latest.state, 2, "latest state"); - - // Specific lookup. - let first = cp - .get("t1", Some("c1")) - .await - .expect("get specific") - .expect("some"); - assert_eq!(first.checkpoint_id, "c1", "specific checkpoint id"); - - // Duplicate-id lookup: every backend resolves a re-used checkpoint id to - // the *last* written record (matching the append-only history and - // `get(None)`), so the three backends stay interchangeable. - cp.put(contract_checkpoint("dup", "d1", None, 1)) - .await - .expect("put dup first"); - cp.put(contract_checkpoint("dup", "d1", None, 9)) - .await - .expect("put dup second"); - let dup = cp - .get("dup", Some("d1")) - .await - .expect("get dup") - .expect("some"); - assert_eq!( - dup.state, 9, - "duplicate checkpoint id resolves to the last written record" - ); - - // Unknown thread / checkpoint. - assert!(cp.get("nope", None).await.expect("get miss").is_none()); - assert!( - cp.get("t1", Some("missing")) - .await - .expect("get miss") - .is_none() - ); - - // Namespace-scoped lookup: checkpoints sharing a thread but differing in - // namespace never resolve to each other. `get_scoped(None, ns)` returns the - // latest in that namespace, not the global latest for the thread. - let mut root = contract_checkpoint("ns", "root1", None, 1); - root.namespace = vec![]; - cp.put(root).await.expect("put root ns"); - let mut child = contract_checkpoint("ns", "child1", None, 5); - child.namespace = vec!["child".to_string()]; - cp.put(child).await.expect("put child ns"); - // The global latest is the child checkpoint, but the root namespace still - // resolves to root1 (and vice versa). - let scoped_root = cp - .get_scoped("ns", None, &[]) - .await - .expect("scoped root") - .expect("some"); - assert_eq!( - scoped_root.checkpoint_id, "root1", - "root namespace ignores the child-namespace checkpoint" - ); - let scoped_child = cp - .get_scoped("ns", None, &["child".to_string()]) - .await - .expect("scoped child") - .expect("some"); - assert_eq!( - scoped_child.checkpoint_id, "child1", - "child namespace resolves to its own checkpoint" - ); - - // Listing preserves insertion order and projects lineage. - let list = cp.list("t1").await.expect("list"); - assert_eq!(list.len(), 2, "listed count"); - assert_eq!(list[0].checkpoint_id, "c1", "list order[0]"); - assert_eq!( - list[1].parent_checkpoint_id.as_deref(), - Some("c1"), - "list lineage" - ); - - // Thread listing includes the written thread. - let threads = cp.list_threads().await.expect("list_threads"); - assert!( - threads.iter().any(|t| t == "t1"), - "list_threads contains t1" - ); - - // Prune keeps a window (and the ancestor chain of what it keeps). - let second_thread = "t2"; - for i in 1..=3 { - let parent = (i > 1).then(|| format!("c{}", i - 1)); - cp.put(contract_checkpoint( - second_thread, - &format!("c{i}"), - parent.as_deref(), - i, - )) - .await - .expect("put for prune"); - } - cp.prune(second_thread, 1).await.expect("prune"); - let pruned = cp.list(second_thread).await.expect("list pruned"); - assert!(!pruned.is_empty(), "prune keeps at least the window"); - - // Delete removes the thread entirely. - cp.delete_thread("t1").await.expect("delete_thread"); - assert!( - cp.get("t1", None) - .await - .expect("get after delete") - .is_none() - ); -} - -/// Runs the [`TaskStore`] contract against `store`. -/// -/// Covers the full lifecycle state machine (pending → running → complete / -/// fail / timeout), cooperative cancellation, kill, deadline updates, filtering -/// by status, and terminal-transition rejection. Any backend that passes this -/// behaves interchangeably for orchestration. -pub fn taskstore_contract(store: S) -where - S: TaskStore, -{ - let spec = |id: &str| { - crate::orchestration::OrchestrationTaskSpec::new( - id, - OrchestrationTaskKind::SubAgent { - agent: "worker".into(), - }, - ) - }; - - // Happy path: pending → running → completed. - let happy = TaskId::new("happy"); - let record = store.insert(spec("happy")).expect("insert"); - assert_eq!(record.status, OrchestrationTaskStatus::Pending); - assert_eq!( - store.mark_running(&happy).expect("running").status, - OrchestrationTaskStatus::Running - ); - let done = store - .complete(&happy, OrchestrationTaskResult::text("ok")) - .expect("complete"); - assert_eq!(done.status, OrchestrationTaskStatus::Completed); - assert!(done.ended_at.is_some(), "completed sets ended_at"); - - // Terminal tasks reject further control. - assert!( - store.request_cancel(&happy).is_err(), - "terminal task rejects cancel" - ); - - // Failure path. - let bad = TaskId::new("bad"); - store.insert(spec("bad")).expect("insert bad"); - assert_eq!( - store.fail(&bad, "boom".into()).expect("fail").status, - OrchestrationTaskStatus::Failed - ); - - // Timeout path. - let slow = TaskId::new("slow"); - store.insert(spec("slow")).expect("insert slow"); - assert_eq!( - store - .timeout(&slow, "deadline".into()) - .expect("timeout") - .status, - OrchestrationTaskStatus::TimedOut - ); - - // Cooperative cancellation: request then mark. - let cancelable = TaskId::new("cancelable"); - store.insert(spec("cancelable")).expect("insert cancelable"); - store.request_cancel(&cancelable).expect("request_cancel"); - assert_eq!( - store.get(&cancelable).expect("get").status, - OrchestrationTaskStatus::CancelRequested - ); - assert_eq!( - store.mark_cancelled(&cancelable).expect("mark").status, - OrchestrationTaskStatus::Cancelled - ); - - // Kill abandons a live task. - let doomed = TaskId::new("doomed"); - store.insert(spec("doomed")).expect("insert doomed"); - store.kill(&doomed).expect("kill"); - assert_eq!( - store.get(&doomed).expect("get").status, - OrchestrationTaskStatus::Abandoned - ); - - // Deadline update on a live task. - let timed = TaskId::new("timed"); - store.insert(spec("timed")).expect("insert timed"); - let updated = store.set_timeout_ms(&timed, 1234).expect("set_timeout_ms"); - assert_eq!(updated.spec.timeout_ms, Some(1234)); - - // Filtering by status returns only matches. - let completed = store.list(OrchestrationTaskFilter { - status: Some(OrchestrationTaskStatus::Completed), - ..OrchestrationTaskFilter::default() - }); - assert_eq!(completed.len(), 1, "one completed task"); - assert_eq!(completed[0].spec.task_id.as_str(), "happy"); -} - -/// Concurrent-write contract for a [`Checkpointer`]: many tasks put distinct -/// checkpoints on one shared instance concurrently; every put must be durable -/// and retrievable with no lost writes. -pub async fn checkpointer_concurrent_contract(cp: std::sync::Arc) -where - C: Checkpointer + 'static, -{ - const WRITERS: usize = 8; - let mut handles = Vec::new(); - for w in 0..WRITERS { - let cp = cp.clone(); - handles.push(tokio::spawn(async move { - cp.put(contract_checkpoint("shared", &format!("c{w}"), None, w)) - .await - .expect("concurrent put"); - })); - } - for h in handles { - h.await.expect("writer task joins"); - } - - // Every concurrently-written checkpoint is retrievable by id. - for w in 0..WRITERS { - let got = cp.get("shared", Some(&format!("c{w}"))).await.expect("get"); - assert!(got.is_some(), "checkpoint c{w} survived concurrent writes"); - } -} - -/// Concurrent-write contract: many threads insert and advance distinct tasks -/// against one shared store; every write must land exactly once with no lost -/// updates. Any backend that passes is safe to share across a run tree. -pub fn taskstore_concurrent_contract(store: std::sync::Arc) -where - S: TaskStore + 'static, -{ - const WRITERS: usize = 8; - std::thread::scope(|scope| { - for w in 0..WRITERS { - let store = store.clone(); - scope.spawn(move || { - let id = TaskId::new(format!("task-{w}")); - store - .insert(crate::orchestration::OrchestrationTaskSpec::new( - id.as_str(), - OrchestrationTaskKind::SubAgent { - agent: "worker".into(), - }, - )) - .expect("concurrent insert"); - store.mark_running(&id).expect("concurrent mark_running"); - }); - } - }); - - let all = store.list(OrchestrationTaskFilter::default()); - assert_eq!(all.len(), WRITERS, "every concurrent insert must land once"); - assert!( - all.iter() - .all(|r| r.status == OrchestrationTaskStatus::Running), - "every concurrently-advanced task reached Running" - ); -} - -/// Durable replay contract: state written through one handle must survive -/// re-opening the backing store through `reopen`. Only meaningful for durable -/// backends (call it for those; an in-memory store would drop state on reopen). -pub fn taskstore_replay_contract(reopen: F) -where - S: TaskStore, - F: Fn() -> S, -{ - let id = TaskId::new("survivor"); - let spec = || { - crate::orchestration::OrchestrationTaskSpec::new( - "survivor", - OrchestrationTaskKind::SubAgent { - agent: "worker".into(), - }, - ) - }; - { - let store = reopen(); - store.insert(spec()).expect("insert"); - store.mark_running(&id).expect("running"); - store - .complete(&id, OrchestrationTaskResult::text("done")) - .expect("complete"); - } - // Re-open: the terminal state and full transition history are reconstructed. - let reopened = reopen(); - let record = reopened.get(&id).expect("task survives reopen"); - assert_eq!(record.status, OrchestrationTaskStatus::Completed); - assert_eq!( - reopened.history(&id).len(), - 3, - "pending → running → completed history replays" - ); -} - -// ── Checkpointer: lineage, thread ops and pending writes ───────────────────── -// -// The original `checkpointer_contract` covered put/get/list/list_threads/ -// delete_thread/prune and nothing else — no `get_tuple`, no `state_history`, no -// `copy_thread`, no `delete_checkpoints`, and nothing at all about pending -// writes. That gap is not incidental: it is precisely why a dead -// `pending_writes` field, an unguarded lineage walk and a `copy_thread` that -// silently corrupts a non-empty target all survived in-tree. These suites port -// the invariants from LangGraph's `checkpoint-conformance` spec. - -use crate::checkpoint::{CheckpointConfig, PendingWrite, WRITES_IDX_RESUME}; - -/// Builds a checkpoint in `namespace` rather than at the root. -fn scoped_checkpoint( - thread: &str, - id: &str, - parent: Option<&str>, - step: usize, - namespace: &[&str], -) -> Checkpoint { - let mut checkpoint = contract_checkpoint(thread, id, parent, step); - checkpoint.namespace = namespace.iter().map(|s| s.to_string()).collect(); - checkpoint -} - -fn config(thread: &str, id: &str, namespace: &[&str]) -> CheckpointConfig { - CheckpointConfig { - thread_id: thread.to_string(), - checkpoint_id: Some(id.to_string()), - namespace: namespace.iter().map(|s| s.to_string()).collect(), - } -} - -/// Runs the **pending-writes** contract against `cp`. -/// -/// Pins the partial-failure protocol: idempotent data writes, upserting -/// control-plane writes, namespace isolation, an empty ledger for a fresh -/// checkpoint, and removal along with the checkpoint/thread they belong to. -/// -/// A backend that does not implement the protocol at all (the trait's default -/// no-op bodies) will fail the very first assertion — which is the point: the -/// suite certifies the protocol, it does not quietly accept its absence. -pub async fn checkpointer_writes_contract(cp: C) -where - C: Checkpointer, -{ - let thread = "writes"; - cp.put(contract_checkpoint(thread, "w1", None, 1)) - .await - .expect("put w1"); - - // A brand-new checkpoint carries no writes. - let fresh = cp - .get_tuple(config(thread, "w1", &[])) - .await - .expect("get_tuple") - .expect("some"); - assert!( - fresh.pending_writes.is_empty(), - "a newly written checkpoint starts with an empty pending-writes ledger" - ); - - // Data writes land and read back in order. - let writes = vec![ - PendingWrite::data("n", "task-a", 0, "out", serde_json::json!("first")), - PendingWrite::data("n", "task-a", 1, "out", serde_json::json!("second")), - ]; - cp.put_writes(&config(thread, "w1", &[]), &writes) - .await - .expect("put_writes"); - let stored = cp - .get_writes(&config(thread, "w1", &[])) - .await - .expect("get_writes"); - assert_eq!(stored.len(), 2, "both data writes stored"); - - // Re-putting the same `(task_id, idx)` data writes must not duplicate them. - cp.put_writes(&config(thread, "w1", &[]), &writes) - .await - .expect("duplicate put_writes"); - let stored = cp - .get_writes(&config(thread, "w1", &[])) - .await - .expect("get_writes after duplicate"); - assert_eq!( - stored.len(), - 2, - "a duplicate (task_id, idx) data write is ignored, not appended" - ); - assert_eq!( - stored[0].payload, - serde_json::json!("first"), - "an ignored duplicate leaves the original value in place" - ); - - // A control-plane write (negative idx) upserts: exactly one row, latest value. - for value in ["v1", "v2"] { - cp.put_writes( - &config(thread, "w1", &[]), - &[PendingWrite::data( - "n", - "task-a", - WRITES_IDX_RESUME, - "__resume__", - serde_json::json!(value), - )], - ) - .await - .expect("put control-plane write"); - } - let stored = cp - .get_writes(&config(thread, "w1", &[])) - .await - .expect("get_writes after control-plane"); - let control: Vec<&PendingWrite> = stored.iter().filter(|w| w.is_control_plane()).collect(); - assert_eq!( - control.len(), - 1, - "a control-plane write stores exactly one row (upsert, not append)" - ); - assert_eq!( - control[0].payload, - serde_json::json!("v2"), - "the control-plane upsert keeps the latest value" - ); - - // `get_tuple` surfaces the ledger. - let tuple = cp - .get_tuple(config(thread, "w1", &[])) - .await - .expect("get_tuple") - .expect("some"); - assert_eq!( - tuple.pending_writes.len(), - stored.len(), - "get_tuple surfaces the persisted pending writes" - ); - - // Namespace isolation: same thread, same checkpoint id, same channel, but a - // different namespace is a different ledger. - cp.put(scoped_checkpoint(thread, "w1", None, 1, &["child"])) - .await - .expect("put child-namespace w1"); - cp.put_writes( - &config(thread, "w1", &["child"]), - &[PendingWrite::data( - "n", - "task-a", - 0, - "out", - serde_json::json!("child"), - )], - ) - .await - .expect("put child-namespace writes"); - let child = cp - .get_writes(&config(thread, "w1", &["child"])) - .await - .expect("child writes"); - assert_eq!(child.len(), 1, "the child namespace has its own ledger"); - assert_eq!(child[0].payload, serde_json::json!("child")); - let root = cp - .get_writes(&config(thread, "w1", &[])) - .await - .expect("root writes"); - assert_eq!( - root.len(), - stored.len(), - "writing the child namespace did not touch the root ledger" - ); - - // `delete_checkpoints` takes the checkpoint's writes with it. - cp.put(contract_checkpoint(thread, "w2", Some("w1"), 2)) - .await - .expect("put w2"); - cp.put_writes( - &config(thread, "w2", &[]), - &[PendingWrite::completion_marker("n", "task-b")], - ) - .await - .expect("put w2 writes"); - cp.delete_checkpoints(thread, &["w2".to_string()]) - .await - .expect("delete_checkpoints"); - assert!( - cp.get_writes(&config(thread, "w2", &[])) - .await - .expect("writes after delete_checkpoints") - .is_empty(), - "deleting a checkpoint deletes its writes" - ); - - // `delete_thread` clears every namespace's ledger, not just the root. - cp.delete_thread(thread).await.expect("delete_thread"); - for namespace in [&[][..], &["child"][..]] { - assert!( - cp.get_writes(&config(thread, "w1", namespace)) - .await - .expect("writes after delete_thread") - .is_empty(), - "delete_thread removes the writes of every namespace ({namespace:?})" - ); - } -} - -/// Runs the **lineage and thread-operation** contract against `cp`. -/// -/// Covers `get_tuple` addressing, `state_history` ordering/limit/cycle safety, -/// `copy_thread` semantics (including the empty-target rule), and per-namespace -/// pruning. -pub async fn checkpointer_lineage_contract(cp: C) -where - C: Checkpointer, -{ - let thread = "lineage"; - for i in 1..=3 { - let parent = (i > 1).then(|| format!("c{}", i - 1)); - cp.put(contract_checkpoint( - thread, - &format!("c{i}"), - parent.as_deref(), - i, - )) - .await - .expect("put lineage"); - } - // An embedded subgraph shares the thread but keeps its own lineage. - cp.put(scoped_checkpoint(thread, "s1", None, 1, &["child"])) - .await - .expect("put child s1"); - cp.put(scoped_checkpoint(thread, "s2", Some("s1"), 2, &["child"])) - .await - .expect("put child s2"); - - // `get_tuple` resolves the addressing config and the parent config. - let tuple = cp - .get_tuple(config(thread, "c2", &[])) - .await - .expect("get_tuple") - .expect("some"); - assert_eq!(tuple.config.checkpoint_id.as_deref(), Some("c2")); - assert_eq!( - tuple - .parent_config - .as_ref() - .and_then(|c| c.checkpoint_id.as_deref()), - Some("c1"), - "get_tuple carries the parent's config" - ); - // A `None` id addresses the latest checkpoint *in the namespace*. - let latest = cp - .get_tuple(CheckpointConfig::latest(thread)) - .await - .expect("get_tuple latest") - .expect("some"); - assert_eq!( - latest.checkpoint.checkpoint_id, "c3", - "an unaddressed get_tuple resolves the latest checkpoint in the namespace, \ - not the thread's global latest" - ); - - // `state_history` is strictly newest-first along the parent spine, scoped - // to the namespace, and `limit` caps from the newest end. - let history = cp - .state_history(thread, &[], None) - .await - .expect("state_history"); - let ids: Vec<&str> = history - .iter() - .map(|t| t.checkpoint.checkpoint_id.as_str()) - .collect(); - assert_eq!(ids, vec!["c3", "c2", "c1"], "state_history is newest-first"); - let limited = cp - .state_history(thread, &[], Some(2)) - .await - .expect("state_history limited"); - assert_eq!( - limited - .iter() - .map(|t| t.checkpoint.checkpoint_id.as_str()) - .collect::>(), - vec!["c3", "c2"], - "limit keeps the most recent hops" - ); - let child_history = cp - .state_history(thread, &["child".to_string()], None) - .await - .expect("child state_history"); - assert_eq!( - child_history - .iter() - .map(|t| t.checkpoint.checkpoint_id.as_str()) - .collect::>(), - vec!["s2", "s1"], - "state_history is namespace-scoped" - ); - - // A self-referential parent pointer must terminate the walk rather than - // spin forever. `parent_checkpoint_id` is caller data, not a structurally - // guaranteed DAG edge. - let cyclic_thread = "cycle"; - let mut cyclic = contract_checkpoint(cyclic_thread, "loop1", None, 1); - cyclic.parent_checkpoint_id = Some("loop1".to_string()); - cp.put(cyclic).await.expect("put cyclic"); - let walked = cp - .state_history(cyclic_thread, &[], None) - .await - .expect("state_history over a cycle terminates"); - assert_eq!( - walked.len(), - 1, - "a lineage cycle is visited once, not forever" - ); - - // `copy_thread` preserves order, ids and namespaces, and leaves the source - // untouched. - cp.copy_thread(thread, "lineage-copy") - .await - .expect("copy_thread"); - let source_after: Vec = cp - .list(thread) - .await - .expect("list source") - .into_iter() - .map(|m| m.checkpoint_id) - .collect(); - let copied = cp.list("lineage-copy").await.expect("list copy"); - assert_eq!( - source_after.len(), - 5, - "copy_thread leaves the source thread unchanged" - ); - assert_eq!( - copied - .iter() - .map(|m| m.checkpoint_id.clone()) - .collect::>(), - source_after, - "copy_thread preserves checkpoint ids and their order" - ); - assert!( - copied - .iter() - .any(|m| m.namespace == vec!["child".to_string()]), - "copy_thread preserves namespaces" - ); - - // Copying onto a non-empty thread is rejected: reusing ids across two - // lineages in one thread is silent corruption, not a merge. - let err = cp.copy_thread(thread, "lineage-copy").await; - assert!( - err.is_err(), - "copy_thread into a non-empty target must be rejected, not interleaved" - ); - - // `prune` keeps the latest per namespace (plus ancestors), so an embedded - // subgraph's lineage survives a prune driven by the parent's recency. - cp.prune(thread, 1).await.expect("prune"); - let kept = cp.list(thread).await.expect("list pruned"); - let namespaces: Vec<&Vec> = kept.iter().map(|m| &m.namespace).collect(); - assert!( - namespaces.iter().any(|ns| ns.is_empty()), - "prune keeps the root namespace's latest" - ); - assert!( - namespaces - .iter() - .any(|ns| ns.as_slice() == ["child".to_string()]), - "prune keeps the latest of EVERY namespace, not a thread-wide window" - ); - - // A pruned-but-kept checkpoint keeps its writes. - let survivor = kept - .iter() - .find(|m| m.namespace.is_empty()) - .expect("a root-namespace survivor"); - let survivor_config = config(thread, &survivor.checkpoint_id, &[]); - cp.put_writes( - &survivor_config, - &[PendingWrite::completion_marker("n", "task-keep")], - ) - .await - .expect("put survivor writes"); - cp.prune(thread, 1).await.expect("prune again"); - assert_eq!( - cp.get_writes(&survivor_config) - .await - .expect("survivor writes") - .len(), - 1, - "prune preserves the writes of the checkpoints it keeps" - ); -} - -/// Runs the [`TaskCache`](crate::TaskCache) contract against `cache`. -/// -/// Covers miss-then-hit round trips, per-key isolation, overwrite, TTL expiry -/// (a lapsed entry reads as a miss and stays gone), and `clear` dropping one -/// graph's entries without touching another's. Takes a short real sleep for the -/// TTL check. -pub async fn task_cache_contract(cache: C) -where - C: crate::TaskCache, -{ - use crate::TaskCacheKey; - use serde_json::json; - use std::time::Duration; - use tinyagents_harness::ids::GraphId; - - let key = |graph: &str, node: &str, hash: &str| { - TaskCacheKey::new(GraphId::new(graph), NodeId::from(node), hash) - }; - let k = key("g", "n", "h1"); - assert_eq!(cache.get(&k).await.unwrap(), None, "empty cache misses"); - cache.put(&k, json!({"v": [1, 2]}), None).await.unwrap(); - assert_eq!(cache.get(&k).await.unwrap(), Some(json!({"v": [1, 2]}))); - for other in [ - key("g", "n", "h2"), - key("g", "m", "h1"), - key("x", "n", "h1"), - ] { - assert_eq!(cache.get(&other).await.unwrap(), None, "keys are distinct"); - } - cache.put(&k, json!("new"), None).await.unwrap(); - assert_eq!( - cache.get(&k).await.unwrap(), - Some(json!("new")), - "overwrite" - ); - - let short = key("g", "n", "short"); - cache - .put(&short, json!(1), Some(Duration::from_secs(2))) - .await - .unwrap(); - assert_eq!(cache.get(&short).await.unwrap(), Some(json!(1))); - tokio::time::sleep(Duration::from_millis(2_100)).await; - assert_eq!(cache.get(&short).await.unwrap(), None, "expired after TTL"); - assert_eq!(cache.get(&short).await.unwrap(), None, "stays expired"); - cache.put(&short, json!(2), None).await.unwrap(); - assert_eq!( - cache.get(&short).await.unwrap(), - Some(json!(2)), - "re-put revives an expired key" - ); - - cache - .put(&key("other", "n", "h"), json!(9), None) - .await - .unwrap(); - cache.clear(&GraphId::new("g")).await.unwrap(); - assert_eq!(cache.get(&k).await.unwrap(), None, "clear drops the graph"); - assert_eq!(cache.get(&short).await.unwrap(), None); - assert_eq!( - cache.get(&key("other", "n", "h")).await.unwrap(), - Some(json!(9)), - "clear leaves other graphs" - ); -} diff --git a/crates/tinyagents-graph/src/testkit/mod.rs b/crates/tinyagents-graph/src/testkit/mod.rs deleted file mode 100644 index a2027bcc5..000000000 --- a/crates/tinyagents-graph/src/testkit/mod.rs +++ /dev/null @@ -1,641 +0,0 @@ -//! Graph-test building blocks — deterministic node doubles, an event recorder, a -//! stream projector, and a fluent run-assertion builder for the durable graph -//! runtime. -//! -//! This is the *graph-level* counterpart to the harness -//! [`testkit`](tinyagents_harness::testkit) (model/tool doubles + trajectories): -//! here the units under test are graph nodes and supersteps, so the doubles are -//! node handlers and the assertions read a run's export/event/checkpoint truth. -//! Because a node can recurse into a [subgraph](crate::subgraph) or a -//! [sub-agent](mod@crate::subagent_node), the same recorder that observes a -//! top-level run also captures the events and child-run rollups of the nested -//! runs it spawns, so recursion stays observable in tests. -//! -//! # Node doubles -//! -//! Each returns a closure ready for -//! [`GraphBuilder::add_node`](crate::GraphBuilder::add_node): -//! -//! | Helper | Behavior | -//! |--------|----------| -//! | [`noop_node`] | Routes onward with no state update | -//! | [`scripted_update_node`] | Emits queued updates (saturating the last) | -//! | [`scripted_route_node`] | Emits queued `goto` route-sets | -//! | [`fanout_node`] | Emits one [`Send`](crate::Send) per arg (fanout) | -//! | [`failing_node`] | Always returns an error | -//! | [`RetryCountingNode`] | Counts activations, fails the first N | -//! | [`interrupting_node`] | Interrupts until resumed, then updates | -//! | [`subgraph_test_node`] | Embeds a child graph (shared state) | -//! | [`subagent_fake_node`] | Records a child run + updates (fake sub-agent) | -//! -//! # Observation & assertions -//! -//! [`GraphEventRecorder`] captures the [`GraphEvent`] stream; [`StreamCollector`] -//! projects it; [`run_recorded`] runs a graph with the recorder wired and -//! bundles the result, recorded events, and checkpoint history into a -//! [`GraphRun`]; [`assert_graph`] opens a fluent [`GraphAssertions`] over it. The -//! single [`GraphRun`] is the test truth — execution, events, and checkpoints -//! all read from it. - -pub mod conformance; -mod types; - -pub use types::{ - GraphAssertions, GraphEventRecorder, GraphRun, RetryCountingNode, StreamCollector, -}; - -use std::sync::Arc; -use std::sync::atomic::{AtomicUsize, Ordering}; - -use serde_json::Value; - -use crate::builder::{NodeContext, NodeFuture}; -use crate::command::{Command, Interrupt, NodeResult, Send as SendPacket}; -use crate::compiled::{CompiledGraph, GraphExecution, StateSnapshot}; -use crate::recursion::ChildRun; -use crate::stream::{CollectingSink, GraphEvent, GraphEventSink}; -use crate::{Result, TinyAgentsError}; -use tinyagents_harness::ids::{GraphId, NodeId, RunId}; -use tinyinference_llm::usage::UsageTotals; - -// --------------------------------------------------------------------------- -// Node doubles -// --------------------------------------------------------------------------- - -/// A node that performs no state update and routes onward via its static or -/// conditional edges (emits an empty [`Command`]). -pub fn noop_node() --> impl Fn(State, NodeContext) -> NodeFuture + Send + Sync + 'static -where - State: Send + 'static, - Update: Send + 'static, -{ - move |_state, _ctx| -> NodeFuture { - Box::pin(async move { Ok(NodeResult::Command(Command::new())) }) - } -} - -/// A node that emits a fixed queue of updates, one per activation. -/// -/// The first activation returns the first update, the second the second, and so -/// on; once the queue is drained every later activation re-emits the *last* -/// update (so the node stays well-defined inside a loop). An empty queue makes -/// every activation fail with [`TinyAgentsError::Graph`]. -pub fn scripted_update_node( - updates: impl IntoIterator, -) -> impl Fn(State, NodeContext) -> NodeFuture + Send + Sync + 'static -where - State: Send + 'static, - Update: Clone + Send + Sync + 'static, -{ - let updates: Arc> = Arc::new(updates.into_iter().collect()); - let idx = Arc::new(AtomicUsize::new(0)); - move |_state, _ctx| -> NodeFuture { - let updates = updates.clone(); - let idx = idx.clone(); - Box::pin(async move { - if updates.is_empty() { - return Err(TinyAgentsError::Graph( - "scripted_update_node has no scripted updates".to_string(), - )); - } - let i = idx.fetch_add(1, Ordering::Relaxed).min(updates.len() - 1); - Ok(NodeResult::Update(updates[i].clone())) - }) - } -} - -/// A node that routes via a fixed queue of `goto` target-sets, one set per -/// activation (saturating the last set once drained). -/// -/// Each item is the set of nodes to activate next; route to -/// [`END`](crate::END) to end a branch. An empty queue makes every -/// activation fail with [`TinyAgentsError::Graph`]. -pub fn scripted_route_node( - routes: I, -) -> impl Fn(State, NodeContext) -> NodeFuture + Send + Sync + 'static -where - State: Send + 'static, - Update: Send + 'static, - I: IntoIterator, - R: IntoIterator, - N: Into, -{ - let routes: Arc>> = Arc::new( - routes - .into_iter() - .map(|r| r.into_iter().map(Into::into).collect()) - .collect(), - ); - let idx = Arc::new(AtomicUsize::new(0)); - move |_state, _ctx| -> NodeFuture { - let routes = routes.clone(); - let idx = idx.clone(); - Box::pin(async move { - if routes.is_empty() { - return Err(TinyAgentsError::Graph( - "scripted_route_node has no scripted routes".to_string(), - )); - } - let i = idx.fetch_add(1, Ordering::Relaxed).min(routes.len() - 1); - Ok(NodeResult::Command(Command::goto(routes[i].clone()))) - }) - } -} - -/// A node that fans out to `target` once per `arg`, delivering each arg through -/// [`NodeContext::send_arg`](crate::NodeContext::send_arg). -/// -/// This is the map-reduce / parallel-tool primitive: it emits one -/// [`Send`](crate::Send) per argument, so `target` is scheduled once for -/// each work item with its own per-invocation input. -pub fn fanout_node( - target: impl Into, - args: impl IntoIterator, -) -> impl Fn(State, NodeContext) -> NodeFuture + Send + Sync + 'static -where - State: Send + 'static, - Update: Send + 'static, -{ - let target = target.into(); - let args: Arc> = Arc::new(args.into_iter().collect()); - move |_state, _ctx| -> NodeFuture { - let target = target.clone(); - let args = args.clone(); - Box::pin(async move { - let sends: Vec = args - .iter() - .map(|a| SendPacket::new(target.clone(), a.clone())) - .collect(); - Ok(NodeResult::Command(Command::send(sends))) - }) - } -} - -/// A node that always fails with [`TinyAgentsError::Graph`] carrying `message`. -pub fn failing_node( - message: impl Into, -) -> impl Fn(State, NodeContext) -> NodeFuture + Send + Sync + 'static -where - State: Send + 'static, - Update: Send + 'static, -{ - let message = message.into(); - move |_state, _ctx| -> NodeFuture { - let message = message.clone(); - Box::pin(async move { Err(TinyAgentsError::Graph(message)) }) - } -} - -/// A node that interrupts (pausing the run for human input) until a resume -/// value arrives, then emits `on_resume`. -/// -/// On an activation with no resume value it returns -/// [`NodeResult::Interrupt`] carrying -/// `payload`; on a resumed activation (a non-empty -/// [`NodeContext::resume`](crate::NodeContext)) it returns -/// `NodeResult::Update(on_resume)`. Requires a checkpointer to actually pause -/// and resume. -pub fn interrupting_node( - payload: Value, - on_resume: Update, -) -> impl Fn(State, NodeContext) -> NodeFuture + Send + Sync + 'static -where - State: Send + 'static, - Update: Clone + Send + Sync + 'static, -{ - move |_state, ctx: NodeContext| -> NodeFuture { - let payload = payload.clone(); - let on_resume = on_resume.clone(); - Box::pin(async move { - match ctx.resume { - Some(_) => Ok(NodeResult::Update(on_resume)), - None => Ok(NodeResult::Interrupt(Interrupt::new( - ctx.node_id.clone(), - payload, - ))), - } - }) - } -} - -/// Embeds `child` as a shared-state subgraph node (a thin wrapper over -/// [`shared_subgraph_node`](crate::shared_subgraph_node)). -/// -/// The child runs over the parent state and its final state becomes the parent -/// update, recording a [`ChildRun`] onto the enclosing run so the subgraph is -/// visible on the parent [`GraphExecution::child_runs`]. -pub fn subgraph_test_node( - child: CompiledGraph, -) -> Box NodeFuture + Send + Sync> -where - State: Clone + Send + Sync + 'static, -{ - crate::subgraph::shared_subgraph_node(child) -} - -/// A fake sub-agent node: records a [`ChildRun`] (with `usage`) onto the -/// enclosing run's child-run sink and emits `update`. -/// -/// This mimics what [`subagent_node`](crate::subagent_node()) records, -/// without needing a registry or a live agent, so tests can assert the -/// parent-run child rollup ([`GraphExecution::child_runs`] / -/// [`run_tree`](crate::GraphExecution::run_tree)) deterministically. The -/// child run preserves the enclosing run's `root_run_id`. -pub fn subagent_fake_node( - agent: impl Into, - update: Update, - usage: UsageTotals, -) -> impl Fn(State, NodeContext) -> NodeFuture + Send + Sync + 'static -where - State: Send + 'static, - Update: Clone + Send + Sync + 'static, -{ - let agent = agent.into(); - move |_state, ctx: NodeContext| -> NodeFuture { - let agent = agent.clone(); - let update = update.clone(); - Box::pin(async move { - if let Some(sink) = &ctx.child_runs { - let root_run_id = ctx - .root_run_id - .clone() - .unwrap_or_else(|| ctx.run_id.clone()); - sink.record(ChildRun { - node: ctx.node_id.clone(), - graph_id: GraphId::new(format!("agent:{agent}")), - run_id: RunId::new(format!( - "subagent-fake-{}", - tinyagents_harness::ids::next_seq() - )), - root_run_id, - usage, - checkpoint_id: None, - }); - } - Ok(NodeResult::Update(update)) - }) - } -} - -impl RetryCountingNode { - /// Creates a counter whose nodes fail their first `fail_times` activations - /// and succeed afterwards. - pub fn new(fail_times: usize) -> Self { - Self { - attempts: Arc::new(AtomicUsize::new(0)), - fail_times, - } - } - - /// The number of activations recorded so far across every handler this - /// counter produced. - pub fn attempts(&self) -> usize { - self.attempts.load(Ordering::Relaxed) - } - - /// Produces a node handler sharing this counter: the first `fail_times` - /// activations return [`TinyAgentsError::Graph`]; later ones succeed with - /// `success`. - pub fn handler( - &self, - success: Update, - ) -> impl Fn(State, NodeContext) -> NodeFuture + Send + Sync + 'static - where - State: Send + 'static, - Update: Clone + Send + Sync + 'static, - { - let attempts = self.attempts.clone(); - let fail_times = self.fail_times; - move |_state, _ctx| -> NodeFuture { - let attempts = attempts.clone(); - let success = success.clone(); - Box::pin(async move { - let n = attempts.fetch_add(1, Ordering::Relaxed) + 1; - if n <= fail_times { - Err(TinyAgentsError::Graph(format!( - "retry_counting_node: attempt {n} of {fail_times} failing" - ))) - } else { - Ok(NodeResult::Update(success)) - } - }) - } - } -} - -// --------------------------------------------------------------------------- -// GraphEventRecorder -// --------------------------------------------------------------------------- - -impl GraphEventRecorder { - /// Creates an empty recorder. - pub fn new() -> Self { - Self { - sink: CollectingSink::new(), - } - } - - /// An [`Arc`]-wrapped event sink to hand to - /// [`CompiledGraph::with_event_sink`](crate::CompiledGraph::with_event_sink). - pub fn sink(&self) -> Arc { - Arc::new(self.sink.clone()) - } - - /// A snapshot of the recorded events, in emission order. - pub fn events(&self) -> Vec { - self.sink.events() - } - - /// The `kind()` string of each recorded event, in emission order. - pub fn kinds(&self) -> Vec { - self.sink - .events() - .iter() - .map(|e| e.kind().to_string()) - .collect() - } - - /// A [`StreamCollector`] over the events recorded so far. - pub fn collector(&self) -> StreamCollector { - StreamCollector::new(self.sink.events()) - } -} - -// --------------------------------------------------------------------------- -// StreamCollector -// --------------------------------------------------------------------------- - -impl StreamCollector { - /// Wraps an owned event sequence for projection. - pub fn new(events: Vec) -> Self { - Self { events } - } - - /// The recorded events, unprojected. - pub fn events(&self) -> &[GraphEvent] { - &self.events - } - - /// The executed-node order (one entry per - /// [`GraphEvent::NodeCompleted`](crate::GraphEvent)). - pub fn node_order(&self) -> Vec { - self.events - .iter() - .filter_map(|e| match e { - GraphEvent::NodeCompleted { node, .. } => Some(node.clone()), - _ => None, - }) - .collect() - } - - /// The nodes that wrote state, in order (one entry per `StateUpdated`). - pub fn updates(&self) -> Vec { - self.events - .iter() - .filter_map(|e| match e { - GraphEvent::StateUpdated { node, .. } => Some(node.clone()), - _ => None, - }) - .collect() - } - - /// The `(from, to)` routes the executor selected, in order. - pub fn routes(&self) -> Vec<(NodeId, NodeId)> { - self.events - .iter() - .filter_map(|e| match e { - GraphEvent::RouteSelected { node, target } => Some((node.clone(), target.clone())), - _ => None, - }) - .collect() - } - - /// The interrupts emitted during the run, in order. - pub fn interrupts(&self) -> Vec { - self.events - .iter() - .filter_map(|e| match e { - GraphEvent::InterruptEmitted { interrupt } => Some(interrupt.clone()), - _ => None, - }) - .collect() - } - - /// The number of persisted checkpoints (one per `CheckpointSaved`). - pub fn checkpoint_count(&self) -> usize { - self.events - .iter() - .filter(|e| matches!(e, GraphEvent::CheckpointSaved { .. })) - .count() - } - - /// The custom node writes, as `(name, data)` pairs in order. - pub fn custom(&self) -> Vec<(String, Value)> { - self.events - .iter() - .filter_map(|e| match e { - GraphEvent::Custom { name, data } => Some((name.clone(), data.clone())), - _ => None, - }) - .collect() - } -} - -// --------------------------------------------------------------------------- -// GraphRun + run_recorded -// --------------------------------------------------------------------------- - -impl GraphRun { - /// Bundles an execution with no recorded events or checkpoint history. - pub fn new(execution: GraphExecution) -> Self { - Self { - execution, - events: Vec::new(), - history: Vec::new(), - } - } - - /// Attaches a recorded event stream. - pub fn with_events(mut self, events: Vec) -> Self { - self.events = events; - self - } - - /// Attaches a checkpoint history (newest-first). - pub fn with_history(mut self, history: Vec>) -> Self { - self.history = history; - self - } - - /// The recorded events as a [`StreamCollector`]. - pub fn collector(&self) -> StreamCollector { - StreamCollector::new(self.events.clone()) - } -} - -/// Runs `graph` to completion with an event recorder wired in and bundles the -/// result into a [`GraphRun`]. -/// -/// When `thread` is `Some`, the run executes under that thread id (so a -/// configured checkpointer persists boundary checkpoints) and the thread's -/// checkpoint history is collected into [`GraphRun::history`]; when `None`, the -/// run executes without a thread (no checkpoints, empty history). The graph is -/// cloned so the caller's instance is untouched. -pub async fn run_recorded( - graph: &CompiledGraph, - thread: Option<&str>, - state: State, -) -> Result> -where - State: Clone + Send + Sync + 'static, - Update: Send + 'static, -{ - let recorder = GraphEventRecorder::new(); - let graph = graph.clone().with_event_sink(recorder.sink()); - let execution = match thread { - Some(thread) => graph.run_with_thread(thread, state).await?, - None => graph.run(state).await?, - }; - let history = match thread { - Some(thread) => graph - .get_state_history(thread, None) - .await - .unwrap_or_default(), - None => Vec::new(), - }; - Ok(GraphRun { - execution, - events: recorder.events(), - history, - }) -} - -// --------------------------------------------------------------------------- -// assert_graph -// --------------------------------------------------------------------------- - -/// Opens a fluent [`GraphAssertions`] over a [`GraphRun`]. -/// -/// ```ignore -/// assert_graph(&run) -/// .visited(["agent", "tools", "agent"]) -/// .routed("agent", "tools") -/// .checkpoint_count(3) -/// .completed(); -/// ``` -pub fn assert_graph(run: &GraphRun) -> GraphAssertions<'_, State> { - GraphAssertions { run } -} - -impl GraphAssertions<'_, State> { - /// Asserts the run visited exactly `expected`, in order (repeats included). - pub fn visited(&self, expected: I) -> &Self - where - I: IntoIterator, - N: Into, - { - let expected: Vec = expected.into_iter().map(Into::into).collect(); - assert_eq!( - self.run.execution.visited, expected, - "assert_graph: expected visited {expected:?} but run visited {:?}", - self.run.execution.visited - ); - self - } - - /// Asserts the executor selected a route from `from` to `to`. - /// - /// Reads the recorded `RouteSelected` events when present; when no events - /// were recorded it falls back to adjacency in the visited sequence (`to` - /// immediately follows `from`). - pub fn routed(&self, from: impl Into, to: impl Into) -> &Self { - let from = from.into(); - let to = to.into(); - let by_event = self.run.events.iter().any(|e| { - matches!( - e, - GraphEvent::RouteSelected { node, target } - if *node == from && *target == to - ) - }); - let by_visited = || { - self.run - .execution - .visited - .windows(2) - .any(|w| w[0] == from && w[1] == to) - }; - assert!( - by_event || (self.run.events.is_empty() && by_visited()), - "assert_graph: expected a route from `{from}` to `{to}` but none was found" - ); - self - } - - /// Asserts the run persisted exactly `n` checkpoints. - /// - /// Counts the recorded `CheckpointSaved` events when present, else falls - /// back to the collected checkpoint-history length. - pub fn checkpoint_count(&self, n: usize) -> &Self { - let count = if self.run.events.is_empty() { - self.run.history.len() - } else { - self.run.collector().checkpoint_count() - }; - assert_eq!( - count, n, - "assert_graph: expected {n} checkpoint(s) but found {count}" - ); - self - } - - /// Asserts against the thread's checkpoint history (newest-first) via a - /// caller-supplied predicate closure. - pub fn state_history(&self, f: impl FnOnce(&[StateSnapshot])) -> &Self { - f(&self.run.history); - self - } - - /// Asserts against the latest checkpoint snapshot via a caller-supplied - /// closure. Panics when the run has no checkpoint history. - pub fn checkpoint(&self, f: impl FnOnce(&StateSnapshot)) -> &Self { - let latest = self - .run - .history - .first() - .expect("assert_graph: expected a checkpoint but the run history is empty"); - f(latest); - self - } - - /// Asserts the run completed (no pending interrupts and a `Completed` - /// status). - pub fn completed(&self) -> &Self { - assert!( - !self.run.execution.is_interrupted(), - "assert_graph: expected the run to complete but it was interrupted: {:?}", - self.run.execution.interrupts - ); - assert_eq!( - self.run.execution.status.status, - tinyagents_harness::ids::ExecutionStatus::Completed, - "assert_graph: expected a Completed status but found {:?}", - self.run.execution.status.status - ); - self - } - - /// Asserts the run paused on an interrupt rather than completing. - pub fn interrupted(&self) -> &Self { - assert!( - self.run.execution.is_interrupted(), - "assert_graph: expected the run to be interrupted but it completed" - ); - self - } -} - -#[cfg(test)] -#[path = "mod_tests.rs"] -mod test; diff --git a/crates/tinyagents-graph/src/testkit/mod_tests.rs b/crates/tinyagents-graph/src/testkit/mod_tests.rs deleted file mode 100644 index 213de276c..000000000 --- a/crates/tinyagents-graph/src/testkit/mod_tests.rs +++ /dev/null @@ -1,327 +0,0 @@ -//! Unit tests for the graph testkit: every node double and every assertion. - -use std::sync::Arc; - -use serde_json::json; - -use super::*; -use crate::builder::{END, GraphBuilder}; -use crate::checkpoint::InMemoryCheckpointer; -use crate::command::{Command, NodeResult}; -use crate::reducer::ClosureStateReducer; -use tinyagents_harness::ids::NodeId; -use tinyinference_llm::usage::UsageTotals; - -fn ids(values: &[&str]) -> Vec { - values.iter().map(|v| NodeId::from(*v)).collect() -} - -// ── noop_node ──────────────────────────────────────────────────────────────── - -#[tokio::test] -async fn noop_node_passes_through_without_updating_state() { - let graph = GraphBuilder::::overwrite() - .add_node("a", noop_node()) - .set_entry("a") - .set_finish("a") - .compile() - .unwrap(); - - let run = run_recorded(&graph, None, 7).await.unwrap(); - assert_eq!(run.execution.state, 7); - assert_graph(&run).visited(["a"]).completed(); -} - -// ── scripted_update_node ───────────────────────────────────────────────────── - -#[tokio::test] -async fn scripted_update_node_emits_queued_updates_in_order() { - let graph = GraphBuilder::::overwrite() - .add_node("a", scripted_update_node([10])) - .add_node("b", scripted_update_node([20])) - .set_entry("a") - .add_edge("a", "b") - .set_finish("b") - .compile() - .unwrap(); - - let run = run_recorded(&graph, None, 0).await.unwrap(); - assert_eq!(run.execution.state, 20); - assert_graph(&run).visited(["a", "b"]).completed(); -} - -#[tokio::test] -async fn scripted_update_node_saturates_last_update_in_a_loop() { - // The "loop" node re-emits its last scripted update on the third visit. - let graph = GraphBuilder::::overwrite() - .add_node("loop", scripted_update_node([1, 2])) - .add_node( - "gate", - scripted_route_node(vec![vec!["loop"], vec!["loop"], vec![END]]), - ) - .set_entry("loop") - .add_edge("loop", "gate") - .compile() - .unwrap(); - - let run = run_recorded(&graph, None, 0).await.unwrap(); - // loop visited 3 times: updates 1, 2, then saturated 2. - assert_eq!(run.execution.state, 2); - assert_graph(&run).visited(["loop", "gate", "loop", "gate", "loop", "gate"]); -} - -#[tokio::test] -async fn scripted_update_node_without_updates_fails() { - let graph = GraphBuilder::::overwrite() - .add_node("a", scripted_update_node(Vec::::new())) - .set_entry("a") - .set_finish("a") - .compile() - .unwrap(); - - assert!(graph.run(0).await.is_err()); -} - -// ── scripted_route_node ────────────────────────────────────────────────────── - -#[tokio::test] -async fn scripted_route_node_routes_to_queued_targets() { - let graph = GraphBuilder::::overwrite() - .add_node("a", scripted_route_node(vec![vec!["b"], vec![END]])) - .add_node("b", noop_node()) - .set_entry("a") - .add_edge("b", "a") - .compile() - .unwrap(); - - let run = run_recorded(&graph, None, 0).await.unwrap(); - assert_graph(&run) - .visited(["a", "b", "a"]) - .routed("a", "b") - .routed("b", "a"); -} - -// ── fanout_node ────────────────────────────────────────────────────────────── - -#[tokio::test] -async fn fanout_node_schedules_target_once_per_arg() { - // Sum reducer so each worker's send_arg folds into the final state. - let graph = GraphBuilder::::new() - .set_reducer(ClosureStateReducer::new(|s: i64, u: i64| Ok(s + u))) - .add_node("fan", fanout_node("work", [json!(1), json!(2), json!(3)])) - .add_node("work", |_s, ctx| async move { - let n = ctx.send_arg.and_then(|v| v.as_i64()).unwrap_or(0); - Ok(NodeResult::Update(n)) - }) - .set_entry("fan") - .set_finish("work") - .compile() - .unwrap(); - - let run = run_recorded(&graph, None, 0).await.unwrap(); - // fan + three work activations. - let work_visits = run - .execution - .visited - .iter() - .filter(|n| n.as_str() == "work") - .count(); - assert_eq!(work_visits, 3); - assert_eq!(run.execution.state, 6); -} - -// ── failing_node ───────────────────────────────────────────────────────────── - -#[tokio::test] -async fn failing_node_aborts_the_run() { - let graph = GraphBuilder::::overwrite() - .add_node("boom", failing_node("kaboom")) - .set_entry("boom") - .set_finish("boom") - .compile() - .unwrap(); - - let err = graph.run(0).await.unwrap_err(); - assert!(err.to_string().contains("kaboom"), "got: {err}"); -} - -// ── RetryCountingNode ──────────────────────────────────────────────────────── - -#[tokio::test] -async fn retry_counting_node_counts_and_eventually_succeeds() { - let counter = RetryCountingNode::new(0); - let graph = GraphBuilder::::overwrite() - .add_node("a", counter.handler(99)) - .set_entry("a") - .set_finish("a") - .compile() - .unwrap(); - - let run = run_recorded(&graph, None, 0).await.unwrap(); - assert_eq!(run.execution.state, 99); - assert_eq!(counter.attempts(), 1); -} - -#[tokio::test] -async fn retry_counting_node_fails_first_attempts() { - let counter = RetryCountingNode::new(1); - let graph = GraphBuilder::::overwrite() - .add_node("a", counter.handler(99)) - .set_entry("a") - .set_finish("a") - .compile() - .unwrap(); - - // The single failing attempt aborts the run. - assert!(graph.run(0).await.is_err()); - assert_eq!(counter.attempts(), 1); -} - -// ── interrupting_node ──────────────────────────────────────────────────────── - -#[tokio::test] -async fn interrupting_node_pauses_then_resumes() { - let cp = Arc::new(InMemoryCheckpointer::::new()); - let graph = GraphBuilder::::overwrite() - .add_node("hitl", interrupting_node(json!("approve?"), 42)) - .set_entry("hitl") - .set_finish("hitl") - .compile() - .unwrap() - .with_checkpointer(cp.clone()); - - let paused = graph.run_with_thread("t1", 0).await.unwrap(); - assert_graph(&GraphRun::new(paused)).interrupted(); - - let resumed = graph - .resume("t1", Command::resume(json!(true))) - .await - .unwrap(); - let final_state = resumed.state; - assert_graph(&GraphRun::new(resumed)).completed(); - assert_eq!(final_state, 42); -} - -// ── subgraph_test_node ─────────────────────────────────────────────────────── - -#[tokio::test] -async fn subgraph_test_node_embeds_a_child_graph() { - let child = GraphBuilder::::overwrite() - .add_node( - "double", - |s, _| async move { Ok(NodeResult::Update(s * 2)) }, - ) - .set_entry("double") - .set_finish("double") - .compile() - .unwrap(); - - let parent = GraphBuilder::::overwrite() - .add_node("sub", subgraph_test_node(child)) - .set_entry("sub") - .set_finish("sub") - .compile() - .unwrap(); - - let run = run_recorded(&parent, None, 5).await.unwrap(); - assert_eq!(run.execution.state, 10); - assert_graph(&run).visited(["sub"]).completed(); - assert_eq!(run.execution.child_runs.len(), 1); - assert_eq!(run.execution.child_runs[0].node.as_str(), "sub"); -} - -// ── subagent_fake_node ─────────────────────────────────────────────────────── - -#[tokio::test] -async fn subagent_fake_node_records_a_child_run() { - let usage = UsageTotals::default(); - let graph = GraphBuilder::::overwrite() - .add_node("ask", subagent_fake_node("helper", 1, usage)) - .set_entry("ask") - .set_finish("ask") - .compile() - .unwrap(); - - let run = run_recorded(&graph, None, 0).await.unwrap(); - assert_eq!(run.execution.child_runs.len(), 1); - let child = &run.execution.child_runs[0]; - assert_eq!(child.node.as_str(), "ask"); - assert_eq!(child.graph_id.as_str(), "agent:helper"); - // The child preserves the parent run as the recursion-tree root. - assert_eq!(child.root_run_id, run.execution.run_id); -} - -// ── GraphEventRecorder + StreamCollector ───────────────────────────────────── - -#[tokio::test] -async fn event_recorder_captures_events_and_collector_projects_them() { - let recorder = GraphEventRecorder::new(); - let graph = GraphBuilder::::overwrite() - .add_node("a", scripted_update_node([1])) - .add_node("b", noop_node()) - .set_entry("a") - .add_edge("a", "b") - .set_finish("b") - .compile() - .unwrap() - .with_event_sink(recorder.sink()); - - graph.run(0).await.unwrap(); - - let kinds = recorder.kinds(); - assert!(kinds.contains(&"run.started".to_string())); - assert!(kinds.contains(&"run.completed".to_string())); - - let collector = recorder.collector(); - assert_eq!(collector.node_order(), ids(&["a", "b"])); - assert_eq!(collector.updates(), ids(&["a"])); - assert_eq!( - collector.routes(), - vec![(NodeId::from("a"), NodeId::from("b"))] - ); - assert!(collector.interrupts().is_empty()); -} - -// ── assert_graph: the exact doc example, plus checkpoint assertions ────────── - -#[tokio::test] -async fn assert_graph_doc_example_with_checkpoints() { - let cp = Arc::new(InMemoryCheckpointer::::new()); - let graph = GraphBuilder::::overwrite() - .add_node("agent", scripted_route_node(vec![vec!["tools"], vec![END]])) - .add_node("tools", noop_node()) - .set_entry("agent") - .add_edge("tools", "agent") - .compile() - .unwrap() - .with_checkpointer(cp); - - let run = run_recorded(&graph, Some("t1"), 0).await.unwrap(); - - assert_graph(&run) - .visited(["agent", "tools", "agent"]) - .routed("agent", "tools") - .checkpoint_count(3) - .completed(); - - // state_history + checkpoint assertions read the durable snapshots. - assert_graph(&run) - .state_history(|history| assert_eq!(history.len(), 3)) - .checkpoint(|latest| assert_eq!(latest.values, 0)); -} - -#[tokio::test] -async fn routed_falls_back_to_visited_adjacency_without_events() { - // A GraphRun assembled by hand (no events) routes off visited adjacency. - let graph = GraphBuilder::::overwrite() - .add_node("a", scripted_update_node([1])) - .add_node("b", noop_node()) - .set_entry("a") - .add_edge("a", "b") - .set_finish("b") - .compile() - .unwrap(); - let execution = graph.run(0).await.unwrap(); - let run = GraphRun::new(execution); - assert_graph(&run).routed("a", "b"); -} diff --git a/crates/tinyagents-graph/src/testkit/types.rs b/crates/tinyagents-graph/src/testkit/types.rs deleted file mode 100644 index 2cc8647c7..000000000 --- a/crates/tinyagents-graph/src/testkit/types.rs +++ /dev/null @@ -1,87 +0,0 @@ -//! Data types for the graph testkit. -//! -//! See the module [`mod`](super) docs for the building blocks these back: the -//! event recorder ([`GraphEventRecorder`]), the projection collector -//! ([`StreamCollector`]), the bundled run-under-test ([`GraphRun`]), the fluent -//! assertion builder ([`GraphAssertions`]), and the stateful retry-counting node -//! ([`RetryCountingNode`]). This file holds the plain definitions; the impls and -//! the node-helper free functions live in `mod.rs`. - -use std::sync::Arc; -use std::sync::atomic::AtomicUsize; - -use crate::compiled::{GraphExecution, StateSnapshot}; -use crate::stream::{CollectingSink, GraphEvent}; - -/// Captures the durable executor's [`GraphEvent`] stream for inspection. -/// -/// Wraps a [`CollectingSink`]; hand its [`sink`](GraphEventRecorder::sink) to -/// [`CompiledGraph::with_event_sink`](crate::CompiledGraph::with_event_sink) -/// (or use [`run_recorded`](super::run_recorded), which wires it for you) so a -/// run's events are recorded, then read them back as raw events, `kind` strings, -/// or a [`StreamCollector`] of higher-level projections. -#[derive(Clone, Default)] -pub struct GraphEventRecorder { - pub(crate) sink: CollectingSink, -} - -/// A higher-level projection over a recorded [`GraphEvent`] stream. -/// -/// Where [`GraphEventRecorder`] hands back raw events, `StreamCollector` slices -/// them into the views a test usually wants: the executed-node order, the nodes -/// that wrote state, the `(from, to)` routes the executor selected, the emitted -/// interrupts, the persisted-checkpoint count, and any custom node writes. -#[derive(Clone, Default)] -pub struct StreamCollector { - pub(crate) events: Vec, -} - -/// A completed (or interrupted) run bundled with everything the assertions need. -/// -/// [`assert_graph`](super::assert_graph) reads its truth from one `GraphRun`: -/// the [`GraphExecution`] (visited nodes, steps, status, interrupts), the -/// recorded [`GraphEvent`] stream (real route selections and checkpoint saves), -/// and the thread's checkpoint `history` (newest-first, as -/// [`get_state_history`](crate::CompiledGraph::get_state_history) returns -/// it). Build one with [`run_recorded`](super::run_recorded), or assemble it by -/// hand via [`GraphRun::new`] and the `with_*` builders. -pub struct GraphRun { - /// The execution result (final state, visited nodes, steps, status). - pub execution: GraphExecution, - /// The recorded low-level graph events, in emission order. - pub events: Vec, - /// The thread's checkpoint history, newest-first (empty for a non-durable - /// run, or one executed without a thread id). - pub history: Vec>, -} - -/// A fluent, panic-on-failure assertion builder over a [`GraphRun`]. -/// -/// Returned by [`assert_graph`](super::assert_graph). Every method asserts and -/// returns `&Self` so checks chain: -/// -/// ```ignore -/// assert_graph(&run) -/// .visited(["agent", "tools", "agent"]) -/// .routed("agent", "tools") -/// .checkpoint_count(3) -/// .completed(); -/// ``` -pub struct GraphAssertions<'a, State> { - pub(crate) run: &'a GraphRun, -} - -/// A stateful node double that counts its activations and fails the first -/// `fail_times` of them. -/// -/// Each [`handler`](RetryCountingNode::handler) it produces shares the same -/// activation counter, so a test can run the graph and then read -/// [`attempts`](RetryCountingNode::attempts) to assert how many times the node -/// ran. The first `fail_times` activations return -/// [`TinyAgentsError::Graph`](crate::TinyAgentsError::Graph); subsequent ones -/// succeed with the configured update. -#[derive(Clone)] -pub struct RetryCountingNode { - pub(crate) attempts: Arc, - pub(crate) fail_times: usize, -} diff --git a/crates/tinyagents-harness-tools/Cargo.toml b/crates/tinyagents-harness-tools/Cargo.toml index 13f733fe1..73828e783 100644 --- a/crates/tinyagents-harness-tools/Cargo.toml +++ b/crates/tinyagents-harness-tools/Cargo.toml @@ -12,15 +12,33 @@ anyhow = { workspace = true } async-trait = { workspace = true } chrono = { workspace = true } chrono-tz = "0.10" +base64 = { version = "0.23", optional = true } +futures-util = { version = "0.3", default-features = false, optional = true } +fs2 = { version = "0.4", optional = true } +glob = { version = "0.3", optional = true } +libc = { version = "0.2", optional = true } +log = { version = "0.4", optional = true } +parking_lot = "0.12" +regex = { version = "1", optional = true } +reqwest = { workspace = true, features = ["stream", "multipart"], optional = true } +sha2 = { version = "0.11", optional = true } serde = { workspace = true } serde_json = { workspace = true } -tinyagents-harness = { path = "../tinyagents-harness", version = "2.11.0", default-features = false } -tinyagents-tasks = { path = "../tinyagents-tasks", version = "2.11.0", default-features = false, optional = true } +tinyagents-harness = { path = "../tinyagents-harness", version = "3.0.0", default-features = false } +tinyagents-tasks = { path = "../tinyagents-tasks", version = "3.0.0", default-features = false, optional = true } tinyinference-llm = { path = "../../vendor/tinyinference/crates/tinyinference-llm", version = "0.5.0", optional = true } tinytools = { path = "../../vendor/tinytools/crates/tinytools", version = "0.6.0" } -tinyagents-orchestration = { path = "../tinyagents-orchestration", version = "2.11.0", default-features = false, optional = true } +tinyagents-orchestration = { path = "../tinyagents-orchestration", version = "3.0.0", default-features = false, optional = true } tokio = { workspace = true, features = ["sync", "time", "macros", "rt", "rt-multi-thread"] } tracing = { workspace = true } +uuid = { workspace = true } +walkdir = { version = "2", optional = true } + +[target.'cfg(unix)'.dependencies] +rustix = { workspace = true, optional = true, features = ["process"] } + +[target.'cfg(not(unix))'.dependencies] +rustix = { version = "1", default-features = false, optional = true, features = ["fs", "process"] } [features] default = ["todos", "goals", "skills"] @@ -29,6 +47,11 @@ goals = [] skills = [] state-block = ["goals", "todos", "dep:tinyinference-llm"] subagents = ["dep:tinyagents-orchestration", "dep:tinyagents-tasks", "dep:tinyinference-llm"] +fs = ["dep:base64", "dep:fs2", "dep:glob", "dep:libc", "dep:log", "dep:regex", "dep:walkdir", "tokio/fs", "tokio/io-util"] +net = ["dep:futures-util", "dep:log", "dep:reqwest", "dep:sha2", "tokio/fs", "tokio/io-util"] +exec = ["dep:rustix", "tokio/fs", "tokio/io-util", "tokio/process"] +media = [] +plan-review = [] [dev-dependencies] tempfile = { workspace = true } diff --git a/crates/tinyagents-harness-tools/README.md b/crates/tinyagents-harness-tools/README.md index 7e9780957..c8bbc1000 100644 --- a/crates/tinyagents-harness-tools/README.md +++ b/crates/tinyagents-harness-tools/README.md @@ -13,6 +13,13 @@ clarification, plan-exit, per-call binding, and thread-lock support are available without optional features. The `state-block` feature enables the run-state middleware and its required `goals` and `todos` modules. +The opt-in `fs`, `net`, and `exec` features expose the standard tool families +under `tinyagents_harness_tools::{fs,net,exec}`. The `media` feature exposes +the generic artifact-filing wrapper at `tinyagents_harness_tools::media`, and +`plan-review` exposes the host-backed review tool at +`tinyagents_harness_tools::plan_review`. These features do not change the +default set. + ## Public modules and tools - `goals` provides `GoalStore`, `GoalTool`, observer notifications, atomic @@ -24,10 +31,21 @@ run-state middleware and its required `goals` and `todos` modules. goal, plan, and tool-call facts in each model request. - `skills` provides `UseSkillTool`, `PackCatalog`, `PackRegistryHandle`, and `UseSkillDispatch` for binding a selected tool from a named pack. + Hosts can recover the shared registry from an erased registered tool with + `skills::pack_registry_handle`; the generic `ToolCallBinding` still pins each + call's admission and execution target. - `subagents` provides `DriverSubAgentTool`, `SubAgentJobsTool`, and `SubAgentMessageTool` over the registry owned by `tinyagents_orchestration::subagent::SubAgentJobRegistry`. - `thread_locks` provides the public `ThreadLockMap` concurrency primitive. +- `fs` provides filesystem tools and file-state operations in `fs`. +- `net` provides network tools in `net`. +- `exec` provides `ExecTool`, `InstallTool`, process timeout utilities, and + host-supplied execution ports in `exec`. +- `media` provides `MediaArtifactTool` and its `ArtifactFiling` port in + `media`. +- `plan-review` provides `RequestPlanReviewTool` and its `PlanReviewPort` in + `plan_review`. - The crate root exports `CurrentTimeTool`, `ResolveTimeTool`, `AskClarificationTool`, `PlanExitTool`, `WaitTool`, `WaitLoopTool`, and the `per_call` binding adapter. @@ -35,3 +53,31 @@ run-state middleware and its required `goals` and `todos` modules. Subagent registry policy stays in `tinyagents-orchestration`; this crate owns the reusable tool declarations. The harness owns tool registration, admission, schema preparation, discovery, and dispatch infrastructure. + +### Host-backed tools + +`ExecTool` and `InstallTool` receive an `ExecGate` and an `ExecBackend` from the +host. The gate authorizes the parsed request before the backend can spawn a +process; the permit carries the backend selection, capability, working +directory, and environment chosen for that call. `NativeExecBackend::new` +selects a shell flavor, but does not choose whether native execution is safe. +An embedding host wires the pair with +`ExecTool::new(Arc, Arc)` and supplies the same +host-owned gate/backend ports to `InstallTool`. +Hosts own approval, security policy, sandbox selection, and audit policy and +may instead provide a sandbox or remote backend. Keep those decisions in the +host adapter. + +For turns that can be cancelled, hosts should create a +`exec::timeout::ProcessCleanup`, run the turn future inside `scope`, drop that +future when cancelling, and await `wait` before acknowledging cancellation. +The native backend registers its child waiter with the active scope; the wait +ensures the process and captured output pipes are finished first. Work handed +to a host-owned task must explicitly inherit cleanup with +`ProcessCleanup::hold` and drop the hold after stopping. + +`MediaArtifactTool` takes an `ArtifactFiling` implementation. The host owns +artifact persistence, file-root checks, and blob lifecycle; the generic wrapper +only maps returned artifact records and rewrites their paths. Likewise, +`RequestPlanReviewTool` uses `PlanReviewPort` to request a host review and +returns its resolution without owning a UI or approval policy. diff --git a/crates/tinyagents-harness-tools/src/exec/detect_tools/mod.rs b/crates/tinyagents-harness-tools/src/exec/detect_tools/mod.rs new file mode 100644 index 000000000..df237c2b7 --- /dev/null +++ b/crates/tinyagents-harness-tools/src/exec/detect_tools/mod.rs @@ -0,0 +1,184 @@ +//! Tool: `detect_tools` — report which developer toolchains are installed on PATH. +//! +//! Lets the agent ground its plans in what the host actually has rather than +//! assuming. Read-only: it only scans `$PATH` for executables (no subprocesses, +//! no writes), so it is safe in every access mode. + +use async_trait::async_trait; +use serde_json::json; +use std::path::PathBuf; +use tinytools::{PermissionLevel, Tool, ToolResult}; + +/// Common developer tools probed when the caller doesn't specify a list. +const DEFAULT_CANDIDATES: &[&str] = &[ + "node", "npm", "npx", "pnpm", "yarn", "bun", "deno", "python3", "python", "pip3", "pip", "uv", + "pipx", "cargo", "rustc", "go", "gcc", "cc", "clang", "make", "git", "gh", "docker", "podman", + "kubectl", "rg", "jq", "fd", "curl", "wget", +]; + +/// Read-only tool reporting which developer toolchains are on `PATH`. +#[derive(Debug)] +pub struct DetectToolsTool; + +impl DetectToolsTool { + #[must_use] + /// Create the tool. + pub fn new() -> Self { + Self + } +} + +impl Default for DetectToolsTool { + fn default() -> Self { + Self::new() + } +} + +/// Fallback executable extensions when Windows has no `PATHEXT` set. +const DEFAULT_PATHEXT: &str = ".EXE;.CMD;.BAT"; + +/// Locate `name` on `$PATH`, honoring `PATHEXT` on Windows. Returns the first +/// matching executable path, or `None` if not found. +#[must_use] +pub fn find_on_path(name: &str) -> Option { + let path = std::env::var_os("PATH")?; + let pathext = cfg!(windows) + .then(|| std::env::var("PATHEXT").unwrap_or_else(|_| DEFAULT_PATHEXT.to_string())); + let file_names = candidate_file_names(name, pathext.as_deref()); + for dir in std::env::split_paths(&path) { + for file_name in &file_names { + let candidate = dir.join(file_name); + if is_executable_file(&candidate) { + return Some(candidate); + } + } + } + None +} + +/// Whether `path` names a regular file the current process can execute. +fn is_executable_file(path: &std::path::Path) -> bool { + if !path.is_file() { + return false; + } + // On Unix, checking mode bits alone ignores the current process's + // effective credentials and supplementary groups. + #[cfg(unix)] + { + rustix::fs::accessat( + rustix::fs::CWD, + path.as_os_str().as_encoded_bytes(), + rustix::fs::Access::EXEC_OK, + rustix::fs::AtFlags::EACCESS, + ) + .is_ok() + } + #[cfg(not(unix))] + { + true + } +} + +/// The file names to probe in each `PATH` directory for `name`. +/// +/// `pathext` is the Windows `PATHEXT` list (`;`-separated), or `None` on +/// platforms that run a bare file name. A name that already ends in one of +/// those extensions (compared case-insensitively, as Windows does) is probed +/// unchanged; any other name is probed once per extension. +fn candidate_file_names(name: &str, pathext: Option<&str>) -> Vec { + let extensions: Vec<&str> = pathext + .into_iter() + .flat_map(|list| list.split(';')) + .filter(|ext| !ext.is_empty()) + .collect(); + let lower_name = name.to_ascii_lowercase(); + let has_extension = extensions + .iter() + .any(|ext| lower_name.ends_with(&ext.to_ascii_lowercase())); + if extensions.is_empty() || has_extension { + return vec![name.to_string()]; + } + extensions + .iter() + .map(|ext| format!("{name}{ext}")) + .collect() +} + +#[async_trait] +impl Tool for DetectToolsTool { + fn name(&self) -> &'static str { + "detect_tools" + } + + fn description(&self) -> &'static str { + "Detect which developer tools / language runtimes are installed on the host PATH \ + (e.g. node, python3, cargo, docker, git, rg). Use this before assuming a tool \ + exists or before proposing to install one. Read-only — scans PATH only." + } + + fn parameters_schema(&self) -> serde_json::Value { + json!({ + "type": "object", + "properties": { + "tools": { + "type": "array", + "items": { "type": "string" }, + "description": "Optional list of tool names to probe. If omitted, a default \ + catalog of common developer tools is probed." + } + } + }) + } + + fn permission_level(&self) -> PermissionLevel { + PermissionLevel::ReadOnly + } + + async fn execute(&self, args: serde_json::Value) -> anyhow::Result { + let requested: Vec = args + .get("tools") + .and_then(|v| v.as_array()) + .map(|a| { + a.iter() + .filter_map(|x| x.as_str().map(str::to_string)) + .collect() + }) + .unwrap_or_default(); + let candidates: Vec = if requested.is_empty() { + DEFAULT_CANDIDATES + .iter() + .map(|s| (*s).to_string()) + .collect() + } else { + requested + }; + + let mut available = Vec::new(); + let mut missing = Vec::new(); + for name in &candidates { + match find_on_path(name) { + Some(p) => available.push(json!({ + "name": name, + "path": p.to_string_lossy(), + })), + None => missing.push(name.clone()), + } + } + + tracing::debug!( + probed = candidates.len(), + available = available.len(), + "[detect_tools] PATH scan complete" + ); + let payload = json!({ + "available": available, + "missing": missing, + "probed": candidates.len(), + }); + Ok(ToolResult::success(serde_json::to_string_pretty(&payload)?)) + } +} + +#[cfg(test)] +#[path = "mod_tests.rs"] +mod tests; diff --git a/crates/tinyagents-harness-tools/src/exec/detect_tools/mod_tests.rs b/crates/tinyagents-harness-tools/src/exec/detect_tools/mod_tests.rs new file mode 100644 index 000000000..f88a39926 --- /dev/null +++ b/crates/tinyagents-harness-tools/src/exec/detect_tools/mod_tests.rs @@ -0,0 +1,125 @@ +//! Unit tests for `PATH` probing, `PATHEXT` candidate names, and the +//! `detect_tools` tool's payload. + +use super::*; + +#[cfg(unix)] +#[test] +fn executable_lookup_requires_current_process_access() -> anyhow::Result<()> { + use std::os::unix::fs::PermissionsExt; + + let dir = std::env::temp_dir().join(format!("tinytools-exec-check-{}", std::process::id())); + std::fs::create_dir_all(&dir)?; + let file = dir.join("candidate"); + std::fs::write(&file, "binary")?; + std::fs::set_permissions(&file, std::fs::Permissions::from_mode(0o600))?; + assert!(!is_executable_file(&file)); + std::fs::set_permissions(&file, std::fs::Permissions::from_mode(0o700))?; + assert!(is_executable_file(&file)); + assert!(!is_executable_file(&dir)); + std::fs::remove_dir_all(dir)?; + Ok(()) +} + +#[test] +fn name_and_permission() { + let tool = DetectToolsTool::new(); + assert_eq!(tool.name(), "detect_tools"); + assert_eq!(tool.permission_level(), PermissionLevel::ReadOnly); +} + +#[test] +fn default_and_metadata_contracts_are_available() { + let tool = ::default(); + assert!(tool.description().contains("PATH")); + assert_eq!( + tool.parameters_schema()["properties"]["tools"]["type"], + "array" + ); +} + +#[tokio::test] +async fn non_string_tool_names_fall_back_to_the_default_catalog() -> anyhow::Result<()> { + let result = DetectToolsTool::new() + .execute(json!({"tools": [null, 3]})) + .await?; + let payload: serde_json::Value = serde_json::from_str(&result.output())?; + assert_eq!(payload["probed"], super::DEFAULT_CANDIDATES.len()); + Ok(()) +} + +#[tokio::test] +async fn missing_tool_reported_missing() -> anyhow::Result<()> { + let tool = DetectToolsTool::new(); + let result = tool + .execute(json!({ "tools": ["definitely_not_a_real_binary_xyz_123"] })) + .await?; + assert!(!result.is_error); + let payload: serde_json::Value = serde_json::from_str(&result.output())?; + assert_eq!(payload["probed"], 1); + assert_eq!(payload["available"], json!([])); + assert_eq!( + payload["missing"], + json!(["definitely_not_a_real_binary_xyz_123"]) + ); + Ok(()) +} + +#[tokio::test] +async fn available_plus_missing_equals_probed() -> anyhow::Result<()> { + let tool = DetectToolsTool::new(); + let result = tool + .execute(json!({ "tools": ["sh", "definitely_not_a_real_binary_xyz_123"] })) + .await?; + let payload: serde_json::Value = serde_json::from_str(&result.output())?; + let avail = payload["available"].as_array().map_or(0, Vec::len); + let miss = payload["missing"].as_array().map_or(0, Vec::len); + assert_eq!(avail + miss, 2); + Ok(()) +} + +#[test] +fn bare_name_is_probed_unchanged_without_pathext() { + assert_eq!(candidate_file_names("git", None), vec!["git".to_string()]); +} + +#[test] +fn extensionless_name_gets_each_pathext_extension() { + assert_eq!( + candidate_file_names("git", Some(".EXE;.CMD")), + vec!["git.EXE".to_string(), "git.CMD".to_string()] + ); +} + +#[test] +fn name_already_carrying_a_pathext_extension_is_probed_unchanged() { + // Windows must probe `git.exe`, not `git.exe.EXE`. + assert_eq!( + candidate_file_names("git.exe", Some(".EXE;.CMD;.BAT")), + vec!["git.exe".to_string()] + ); + assert_eq!( + candidate_file_names("build.Cmd", Some(".EXE;.CMD")), + vec!["build.Cmd".to_string()] + ); +} + +#[test] +fn name_with_a_non_pathext_extension_still_gets_pathext_appended() { + assert_eq!( + candidate_file_names("python3.11", Some(".EXE")), + vec!["python3.11.EXE".to_string()] + ); +} + +#[test] +fn empty_pathext_entries_are_ignored() { + assert_eq!( + candidate_file_names("git", Some(".EXE;;")), + vec!["git.EXE".to_string()] + ); + assert_eq!( + candidate_file_names("git", Some("")), + vec!["git".to_string()] + ); +} diff --git a/crates/tinyagents-harness-tools/src/exec/install_schema.json b/crates/tinyagents-harness-tools/src/exec/install_schema.json new file mode 100644 index 000000000..b685b63f3 --- /dev/null +++ b/crates/tinyagents-harness-tools/src/exec/install_schema.json @@ -0,0 +1,8 @@ +{ + "type": "object", + "properties": { + "package": { "type": "string", "description": "Package name to install (e.g. 'ripgrep', 'jq', '@scope/cli')." }, + "manager": { "type": "string", "enum": ["apt", "apt-get", "dnf", "yum", "pacman", "apk", "brew", "winget", "pipx", "npm", "cargo"], "description": "Package manager to use. If omitted, the host's system package manager is auto-detected." } + }, + "required": ["package"] +} diff --git a/crates/tinyagents-harness-tools/src/exec/install_tool.rs b/crates/tinyagents-harness-tools/src/exec/install_tool.rs new file mode 100644 index 000000000..21b8f92e1 --- /dev/null +++ b/crates/tinyagents-harness-tools/src/exec/install_tool.rs @@ -0,0 +1,280 @@ +use std::{ + ffi::OsString, + sync::Arc, + time::{Duration, Instant}, +}; + +use async_trait::async_trait; +use serde_json::Value; +use tinytools::{PermissionLevel, Tool, ToolCallOptions, ToolResult, ToolRunContext}; + +use super::{ + ExecAuditEvent, ExecAuditOutcome, ExecAuthorizationRequest, ExecBackend, ExecBackendOutcome, + ExecGate, ExecOperation, ExecRequest, detect_tools::find_on_path, +}; + +const INSTALL_TIMEOUT: Duration = Duration::from_secs(300); +const MAX_OUTPUT_BYTES: usize = 1_048_576; + +pub struct InstallTool { + gate: Arc, + backend: Arc, +} + +impl InstallTool { + #[must_use] + pub fn new(gate: Arc, backend: Arc) -> Self { + Self { gate, backend } + } + + async fn execute_with_context_and_path_lookup( + &self, + args: Value, + _options: ToolCallOptions, + context: Option<&dyn ToolRunContext>, + on_path: impl FnMut(&str) -> bool + Send, + ) -> anyhow::Result { + execute_install(self, args, context, on_path).await + } +} + +fn manager_argv(manager: &str) -> Option<(&'static str, &'static [&'static str])> { + Some(match manager { + "apt" | "apt-get" => ("apt-get", &["install", "-y"]), + "dnf" => ("dnf", &["install", "-y"]), + "yum" => ("yum", &["install", "-y"]), + "pacman" => ("pacman", &["-S", "--noconfirm"]), + "apk" => ("apk", &["add"]), + "brew" => ("brew", &["install"]), + "winget" => ("winget", &["install", "-e", "--id"]), + "pipx" => ("pipx", &["install"]), + "npm" => ("npm", &["install", "-g"]), + "cargo" => ("cargo", &["install"]), + _ => return None, + }) +} + +fn resolve_manager( + requested: Option<&str>, + mut on_path: impl FnMut(&str) -> bool, +) -> Result<(String, &'static str, &'static [&'static str]), String> { + let manager = match requested { + Some(manager) => manager.to_owned(), + None => ["apt-get", "dnf", "yum", "pacman", "apk", "brew", "winget"] + .into_iter() + .find(|name| on_path(name)) + .unwrap_or("") + .to_owned(), + }; + if manager.is_empty() { + return Err("No supported package manager detected on PATH. Pass 'manager' explicitly (e.g. brew, pipx, npm, cargo).".to_owned()); + } + let Some((program, prefix)) = manager_argv(&manager) else { + return Err(format!("Unsupported package manager '{manager}'")); + }; + if !on_path(program) { + return Err(format!( + "Package manager '{program}' is not installed on PATH" + )); + } + Ok((manager, program, prefix)) +} + +fn valid_package(package: &str) -> bool { + !package.is_empty() + && package.len() <= 200 + && package.chars().all(|c| { + c.is_ascii_alphanumeric() || matches!(c, '.' | '_' | '+' | '-' | '@' | '/' | ':') + }) +} + +fn truncate_stream(bytes: &[u8], label: &str) -> String { + let text = String::from_utf8_lossy(bytes); + if text.len() <= MAX_OUTPUT_BYTES { + return text.into_owned(); + } + let mut end = MAX_OUTPUT_BYTES; + while !text.is_char_boundary(end) { + end -= 1; + } + format!("{}\n... [{label} truncated at 1 MiB]", &text[..end]) +} + +#[async_trait] +impl Tool for InstallTool { + fn name(&self) -> &str { + "install_tool" + } + fn description(&self) -> &str { + "Install an OS or language package on the host (apt/dnf/pacman/apk/brew/winget/pipx/npm -g/cargo). HIGH IMPACT: host policy must require explicit user approval. Prefer detect_tools first; only install what's actually missing." + } + fn parameters_schema(&self) -> Value { + crate::static_schema!(include_str!("install_schema.json")) + } + fn permission_level(&self) -> PermissionLevel { + PermissionLevel::Dangerous + } + fn external_effect(&self) -> bool { + true + } + fn external_effect_with_args(&self, args: &Value) -> bool { + self.gate.external_effect_with_args(args) + } + async fn prepare_authorization(&self, args: &Value) -> anyhow::Result> { + self.gate.prepare_authorization(args).await + } + async fn execute(&self, args: Value) -> anyhow::Result { + self.execute_with_context(args, ToolCallOptions::default(), None) + .await + } + async fn execute_with_context( + &self, + args: Value, + options: ToolCallOptions, + context: Option<&dyn ToolRunContext>, + ) -> anyhow::Result { + self.execute_with_context_and_path_lookup(args, options, context, |name| { + find_on_path(name).is_some() + }) + .await + } +} + +async fn execute_install( + tool: &InstallTool, + args: Value, + context: Option<&dyn ToolRunContext>, + on_path: impl FnMut(&str) -> bool + Send, +) -> anyhow::Result { + let Some(package) = args.get("package").and_then(Value::as_str).map(str::trim) else { + return Ok(ToolResult::error("Missing 'package' parameter")); + }; + if !valid_package(package) { + return Ok(ToolResult::error(format!( + "Invalid package name '{package}': only letters, digits and ._+-@/: are allowed" + ))); + } + let requested_manager = args.get("manager").and_then(Value::as_str); + let authorization_operation = ExecOperation::Install { + package: package.to_owned(), + // Keep automatic selection opaque until the gate authorizes the install. + manager: requested_manager.unwrap_or("auto").to_owned(), + }; + let authorization = ExecAuthorizationRequest { + operation: authorization_operation.clone(), + timeout: Some(INSTALL_TIMEOUT), + }; + let started = Instant::now(); + let approval_required = tool.gate.external_effect_with_args(&args); + let permit = match tool.gate.authorize(&authorization, context).await { + Ok(permit) => permit, + Err(denied) => { + tool.gate.record(ExecAuditEvent { + operation: authorization_operation, + approval_required, + authorized: false, + allowed: false, + backend_selection: None, + process_attempted: false, + outcome: ExecAuditOutcome::GateDenied, + duration: started.elapsed(), + }); + return Ok(ToolResult::error(denied.0)); + } + }; + let backend_selection = permit.backend_selection.clone(); + let (manager, program, prefix) = match resolve_manager(requested_manager, on_path) { + Ok(resolved) => resolved, + Err(error) => { + tool.gate.record(ExecAuditEvent { + operation: authorization_operation, + approval_required, + authorized: true, + allowed: false, + backend_selection: Some(backend_selection.clone()), + process_attempted: false, + outcome: ExecAuditOutcome::InstallManagerResolutionFailed, + duration: started.elapsed(), + }); + return Ok(ToolResult::error(error)); + } + }; + let operation = ExecOperation::Install { + package: package.to_owned(), + manager: manager.clone(), + }; + let mut argv = prefix.iter().map(OsString::from).collect::>(); + argv.push(package.into()); + let request = ExecRequest { + program: program.into(), + args: argv, + action_dir: permit.action_dir, + environment: permit.environment, + timeout: Some(INSTALL_TIMEOUT), + shell_command: None, + backend_selection: permit.backend_selection, + backend_capability: permit.backend_capability, + }; + let outcome = tool.backend.execute(request, context, None).await; + let (allowed, attempted, audit, result) = match outcome { + ExecBackendOutcome::Refused { kind, reason } => ( + false, + false, + ExecAuditOutcome::BackendRefused { kind }, + ToolResult::error(reason), + ), + ExecBackendOutcome::LaunchFailed { error } => ( + true, + true, + ExecAuditOutcome::LaunchFailed, + ToolResult::error(format!("Failed to launch {program}: {error}")), + ), + ExecBackendOutcome::Completed(output) if output.timed_out => ( + true, + true, + ExecAuditOutcome::Completed { + success: false, + timed_out: true, + }, + ToolResult::error("Install timed out after 300s and was killed"), + ), + ExecBackendOutcome::Completed(output) => { + let stdout = truncate_stream(&output.stdout, "stdout"); + let stderr = truncate_stream(&output.stderr, "stderr"); + let result = if output.exit_code == Some(0) { + ToolResult::success(format!("Installed '{package}' via {program}.\n{stdout}")) + } else { + ToolResult::error(format!( + "{program} failed to install '{package}' (exit {:?}).\n{}", + output.exit_code, + if stderr.is_empty() { stdout } else { stderr } + )) + }; + let success = !result.is_error; + ( + true, + true, + ExecAuditOutcome::Completed { + success, + timed_out: false, + }, + result, + ) + } + }; + tool.gate.record(ExecAuditEvent { + operation, + approval_required, + authorized: true, + allowed, + backend_selection: Some(backend_selection), + process_attempted: attempted, + outcome: audit, + duration: started.elapsed(), + }); + Ok(result) +} + +#[cfg(test)] +#[path = "install_tool_tests.rs"] +mod tests; diff --git a/crates/tinyagents-harness-tools/src/exec/install_tool_tests.rs b/crates/tinyagents-harness-tools/src/exec/install_tool_tests.rs new file mode 100644 index 000000000..0fa8e1b32 --- /dev/null +++ b/crates/tinyagents-harness-tools/src/exec/install_tool_tests.rs @@ -0,0 +1,357 @@ +use super::*; +use crate::exec::{ + ExecBackendCapability, ExecDenied, ExecOutput, ExecPermit, ExecProgress, ShellFlavor, +}; +use serde_json::{Value, json}; +use std::{ + collections::BTreeMap, + ffi::OsString, + path::PathBuf, + sync::{ + Mutex, + atomic::{AtomicUsize, Ordering}, + }, +}; + +#[derive(Default)] +struct State { + denied: bool, + events: Mutex>, + backend_calls: AtomicUsize, + authorization_calls: AtomicUsize, + authorization_request: Mutex>, + manager_lookup_calls: AtomicUsize, + request: Mutex>, + backend_outcome: Mutex>, +} +struct Gate(Arc); + +#[async_trait] +impl ExecGate for Gate { + async fn prepare_authorization(&self, args: &Value) -> anyhow::Result> { + Ok(Some(json!({"prepared_package":args.get("package")}))) + } + fn external_effect_with_args(&self, args: &Value) -> bool { + args.get("gate_effect") + .and_then(Value::as_bool) + .unwrap_or(false) + } + async fn authorize( + &self, + request: &ExecAuthorizationRequest, + _: Option<&dyn ToolRunContext>, + ) -> Result { + self.0.authorization_calls.fetch_add(1, Ordering::SeqCst); + *self.0.authorization_request.lock().unwrap() = Some(request.clone()); + if self.0.denied { + return Err(ExecDenied("installation denied".into())); + } + Ok(ExecPermit { + action_dir: PathBuf::from("/tmp"), + environment: BTreeMap::new(), + backend_selection: "test".into(), + backend_capability: ExecBackendCapability::new(()), + }) + } + fn record(&self, event: ExecAuditEvent) { + self.0.events.lock().unwrap().push(event); + } +} +struct Backend(Arc); + +#[async_trait] +impl ExecBackend for Backend { + fn shell_flavor(&self) -> ShellFlavor { + ShellFlavor::Posix + } + async fn execute( + &self, + request: ExecRequest, + _: Option<&dyn ToolRunContext>, + _: Option, + ) -> ExecBackendOutcome { + self.0.backend_calls.fetch_add(1, Ordering::SeqCst); + *self.0.request.lock().unwrap() = Some(request); + if let Some(outcome) = self.0.backend_outcome.lock().unwrap().take() { + return outcome; + } + ExecBackendOutcome::Completed(ExecOutput { + stdout: b"ok".to_vec(), + stderr: Vec::new(), + exit_code: Some(0), + timed_out: false, + }) + } +} + +fn tool(denied: bool) -> (InstallTool, Arc) { + let state = Arc::new(State { + denied, + ..State::default() + }); + ( + InstallTool::new( + Arc::new(Gate(state.clone())), + Arc::new(Backend(state.clone())), + ), + state, + ) +} + +#[test] +fn package_validation_rejects_shell_injection_and_oversized_input() { + assert!(!valid_package("pkg; touch /tmp/pwned")); + assert!(!valid_package(&"a".repeat(201))); + for package in ["ripgrep", "@scope/cli", "lib:feature", "a+b-c_1.0"] { + assert!(valid_package(package)); + } +} + +#[tokio::test] +async fn install_delegates_authorization_preparation_and_argument_effect_to_gate() { + let (tool, _) = tool(false); + let args = json!({"package":"ripgrep","gate_effect":false}); + assert_eq!( + tool.prepare_authorization(&args).await.unwrap(), + Some(json!({"prepared_package":"ripgrep"})) + ); + assert!(!tool.external_effect_with_args(&args)); + assert!(tool.external_effect_with_args(&json!({"package":"ripgrep","gate_effect":true}))); +} + +#[test] +fn every_supported_manager_keeps_its_exact_argument_prefix() { + let cases = [ + ("apt", "apt-get", &["install", "-y"][..]), + ("apt-get", "apt-get", &["install", "-y"][..]), + ("dnf", "dnf", &["install", "-y"][..]), + ("yum", "yum", &["install", "-y"][..]), + ("pacman", "pacman", &["-S", "--noconfirm"][..]), + ("apk", "apk", &["add"][..]), + ("brew", "brew", &["install"][..]), + ("winget", "winget", &["install", "-e", "--id"][..]), + ("pipx", "pipx", &["install"][..]), + ("npm", "npm", &["install", "-g"][..]), + ("cargo", "cargo", &["install"][..]), + ]; + for (manager, program, prefix) in cases { + let (actual_program, actual_prefix) = manager_argv(manager).unwrap(); + assert_eq!(actual_program, program); + assert_eq!(actual_prefix, prefix); + } + assert_eq!(manager_argv("unknown"), None); +} + +#[test] +fn missing_path_and_unsupported_managers_keep_exact_errors() { + assert_eq!( + resolve_manager(Some("cargo"), |_| false).unwrap_err(), + "Package manager 'cargo' is not installed on PATH" + ); + assert_eq!( + resolve_manager(None, |_| false).unwrap_err(), + "No supported package manager detected on PATH. Pass 'manager' explicitly (e.g. brew, pipx, npm, cargo)." + ); + assert_eq!( + resolve_manager(Some("unknown"), |_| true).unwrap_err(), + "Unsupported package manager 'unknown'" + ); +} + +#[test] +fn package_output_truncation_is_per_stream_and_utf8_safe() { + let text = format!("{}é", "a".repeat(MAX_OUTPUT_BYTES - 1)); + assert_eq!(text.len(), MAX_OUTPUT_BYTES + 1); + let truncated = truncate_stream(text.as_bytes(), "stdout"); + assert!(truncated.ends_with("\n... [stdout truncated at 1 MiB]")); + assert!(truncated.is_char_boundary(truncated.len())); + assert!(!truncated[..MAX_OUTPUT_BYTES].ends_with('�')); +} + +#[tokio::test] +async fn invalid_package_never_reaches_gate_or_backend() { + let (tool, state) = tool(false); + let result = tool + .execute(json!({"package":"pkg; touch /tmp/pwned","manager":"cargo"})) + .await + .unwrap(); + assert!(result.is_error); + assert_eq!(state.backend_calls.load(Ordering::SeqCst), 0); + assert!(state.events.lock().unwrap().is_empty()); +} + +#[tokio::test] +async fn denied_install_never_calls_backend_and_records_denial() { + let (tool, state) = tool(true); + let result = tool + .execute(json!({"package":"ripgrep","manager":"cargo"})) + .await + .unwrap(); + assert!(result.is_error); + assert_eq!(result.text(), "installation denied"); + assert_eq!(state.backend_calls.load(Ordering::SeqCst), 0); + let events = state.events.lock().unwrap(); + assert_eq!(events.len(), 1); + assert_eq!(events[0].outcome, ExecAuditOutcome::GateDenied); + assert!(!events[0].process_attempted); +} + +#[tokio::test] +async fn denied_install_is_authorized_before_manager_resolution() { + for (args, expected_manager) in [ + (json!({"package":"ripgrep"}), "auto"), + ( + json!({"package":"ripgrep","manager":"unsupported"}), + "unsupported", + ), + ] { + let (tool, state) = tool(true); + let result = tool + .execute_with_context_and_path_lookup(args, ToolCallOptions::default(), None, |_| { + state.manager_lookup_calls.fetch_add(1, Ordering::SeqCst); + false + }) + .await + .unwrap(); + + assert_eq!(result.text(), "installation denied"); + assert_eq!(state.authorization_calls.load(Ordering::SeqCst), 1); + assert_eq!(state.manager_lookup_calls.load(Ordering::SeqCst), 0); + assert_eq!( + state + .authorization_request + .lock() + .unwrap() + .as_ref() + .unwrap() + .operation, + ExecOperation::Install { + package: "ripgrep".to_owned(), + manager: expected_manager.to_owned(), + } + ); + let events = state.events.lock().unwrap(); + assert_eq!(events.len(), 1); + assert_eq!(events[0].outcome, ExecAuditOutcome::GateDenied); + assert!(!events[0].process_attempted); + } +} + +#[tokio::test] +async fn authorized_install_manager_resolution_failure_is_audited_without_execution() { + let cases = [ + ( + json!({"package":"ripgrep"}), + "auto", + "No supported package manager detected on PATH. Pass 'manager' explicitly (e.g. brew, pipx, npm, cargo).", + ), + ( + json!({"package":"ripgrep","manager":"unsupported"}), + "unsupported", + "Unsupported package manager 'unsupported'", + ), + ]; + + for (args, expected_manager, expected_error) in cases { + let (tool, state) = tool(false); + let lookup_state = state.clone(); + let result = tool + .execute_with_context_and_path_lookup( + args, + ToolCallOptions::default(), + None, + move |_| { + lookup_state + .manager_lookup_calls + .fetch_add(1, Ordering::SeqCst); + false + }, + ) + .await + .unwrap(); + + assert_eq!(result.text(), expected_error); + assert_eq!(state.authorization_calls.load(Ordering::SeqCst), 1); + assert_eq!(state.backend_calls.load(Ordering::SeqCst), 0); + let events = state.events.lock().unwrap(); + assert_eq!(events.len(), 1); + assert!(events[0].authorized); + assert!(!events[0].allowed); + assert!(!events[0].process_attempted); + assert_eq!( + events[0].outcome, + ExecAuditOutcome::InstallManagerResolutionFailed + ); + assert_eq!(events[0].backend_selection.as_deref(), Some("test")); + assert_eq!( + events[0].operation, + ExecOperation::Install { + package: "ripgrep".to_owned(), + manager: expected_manager.to_owned(), + } + ); + } +} + +#[tokio::test] +async fn install_uses_exact_manager_argv_and_no_shell_command() { + let (tool, state) = tool(false); + let result = tool + .execute(json!({"package":"@scope/cli","manager":"cargo"})) + .await + .unwrap(); + assert!(!result.is_error); + let request = state.request.lock().unwrap(); + let request = request.as_ref().unwrap(); + assert_eq!(request.program, OsString::from("cargo")); + assert_eq!( + request.args, + vec![OsString::from("install"), OsString::from("@scope/cli")] + ); + assert_eq!(request.shell_command, None); + assert_eq!(request.timeout, Some(INSTALL_TIMEOUT)); + assert_eq!(state.events.lock().unwrap().len(), 1); +} + +#[tokio::test] +async fn install_maps_launch_failure_nonzero_exit_and_timeout_exactly() { + let (tool, state) = tool(false); + *state.backend_outcome.lock().unwrap() = Some(ExecBackendOutcome::LaunchFailed { + error: "permission denied".into(), + }); + assert_eq!( + tool.execute(json!({"package":"ripgrep","manager":"cargo"})) + .await + .unwrap() + .text(), + "Failed to launch cargo: permission denied" + ); + + *state.backend_outcome.lock().unwrap() = Some(ExecBackendOutcome::Completed(ExecOutput { + stdout: b"informational".to_vec(), + stderr: b"compile failed".to_vec(), + exit_code: Some(17), + timed_out: false, + })); + assert_eq!( + tool.execute(json!({"package":"ripgrep","manager":"cargo"})) + .await + .unwrap() + .text(), + "cargo failed to install 'ripgrep' (exit Some(17)).\ncompile failed" + ); + + *state.backend_outcome.lock().unwrap() = Some(ExecBackendOutcome::Completed(ExecOutput { + stdout: b"partial".to_vec(), + stderr: b"still working".to_vec(), + exit_code: None, + timed_out: true, + })); + assert_eq!( + tool.execute(json!({"package":"ripgrep","manager":"cargo"})) + .await + .unwrap() + .text(), + "Install timed out after 300s and was killed" + ); +} diff --git a/crates/tinyagents-harness-tools/src/exec/mod.rs b/crates/tinyagents-harness-tools/src/exec/mod.rs new file mode 100644 index 000000000..36bdb734b --- /dev/null +++ b/crates/tinyagents-harness-tools/src/exec/mod.rs @@ -0,0 +1,25 @@ +//! Process utilities and gated execution tools. + +pub mod detect_tools; +pub mod install_tool; +pub mod native_backend; +pub mod ports; +pub mod progress; +mod shell_description; +mod shell_support; +pub mod timeout; + +pub use install_tool::InstallTool; +pub use native_backend::{NativeExecBackend, NativeExecCapability}; +pub use ports::{ + ExecAuditEvent, ExecAuditOutcome, ExecAuthorizationRequest, ExecBackend, ExecBackendCapability, + ExecBackendOutcome, ExecDenied, ExecGate, ExecOperation, ExecOutput, ExecPermit, + ExecRefusalKind, ExecRequest, ShellFlavor, +}; +pub use progress::ExecProgress; +mod tool; +pub use tool::ExecTool; + +#[cfg(test)] +#[path = "ports_tests.rs"] +mod ports_tests; diff --git a/crates/tinyagents-harness-tools/src/exec/native_backend.rs b/crates/tinyagents-harness-tools/src/exec/native_backend.rs new file mode 100644 index 000000000..8bb93b25c --- /dev/null +++ b/crates/tinyagents-harness-tools/src/exec/native_backend.rs @@ -0,0 +1,129 @@ +//! Native process backend for hosts that explicitly select local execution. + +use tinytools::ToolRunContext; + +use super::{ + ExecBackend, ExecBackendOutcome, ExecOutput, ExecProgress, ExecRequest, ShellFlavor, + timeout::{ + CapturedCommandResult, output_or_kill_captured_isolated, + output_unbounded_with_progress_isolated, + }, +}; + +/// Capability token required by [`NativeExecBackend`] before it may spawn. +#[derive(Debug, Clone, Copy, Default)] +pub struct NativeExecCapability; + +/// Runs an authorized request as a local child process. +/// +/// Hosts opt into this backend by constructing it with the selected +/// [`ShellFlavor`] and providing it to an [`ExecTool`](super::ExecTool) through +/// the execution ports. The backend only checks the opaque native capability +/// supplied by the host; it does not implement approval, command policy, +/// sandbox selection, or audit policy. +/// +/// A host that acknowledges cancellation must retain a +/// [`super::timeout::ProcessCleanup`] for the turn, run the turn inside +/// [`ProcessCleanup::scope`](super::timeout::ProcessCleanup::scope), drop the +/// scoped turn future on cancellation, and await +/// [`ProcessCleanup::wait`](super::timeout::ProcessCleanup::wait) before +/// acknowledging cleanup. This lets the backend's child waiter finish and +/// close captured output pipes before the host reports that cancellation is +/// complete. +pub struct NativeExecBackend { + flavor: ShellFlavor, +} + +#[cfg(test)] +#[path = "native_backend_tests.rs"] +mod tests; + +impl NativeExecBackend { + /// Selects the shell syntax used by this local backend. + /// + /// The host remains responsible for authorizing a request and attaching a + /// [`NativeExecCapability`] only when local execution is permitted. + #[must_use] + pub fn new(flavor: ShellFlavor) -> Self { + Self { flavor } + } +} + +#[async_trait::async_trait] +impl ExecBackend for NativeExecBackend { + fn shell_flavor(&self) -> ShellFlavor { + self.flavor + } + + async fn execute( + &self, + request: ExecRequest, + _context: Option<&dyn ToolRunContext>, + progress: Option, + ) -> ExecBackendOutcome { + use std::process::Stdio; + if request.backend_selection != "native" + || request + .backend_capability + .downcast_ref::() + .is_none() + { + return ExecBackendOutcome::Refused { + kind: super::ExecRefusalKind::Preflight, + reason: "native execution was not selected or authorized".into(), + }; + } + if request.program.is_empty() { + return ExecBackendOutcome::Refused { + kind: super::ExecRefusalKind::Preflight, + reason: "program is empty".into(), + }; + } + let mut command = tokio::process::Command::new(&request.program); + command + .args(&request.args) + .current_dir(&request.action_dir) + .env_clear() + .envs(request.environment) + .stdin(Stdio::null()); + let run = async { + match request.timeout { + Some(timeout) => { + match output_or_kill_captured_isolated(&mut command, timeout, progress).await { + CapturedCommandResult::Completed(Ok(output)) => Ok(ExecOutput { + stdout: output.stdout, + stderr: output.stderr, + exit_code: output.status.code(), + timed_out: false, + }), + CapturedCommandResult::Completed(Err(error)) => Err(error), + CapturedCommandResult::TimedOut(output) => Ok(ExecOutput { + stdout: output + .as_ref() + .map_or_else(Vec::new, |output| output.stdout.clone()), + stderr: output + .as_ref() + .map_or_else(Vec::new, |output| output.stderr.clone()), + exit_code: output.and_then(|output| output.status.code()), + timed_out: true, + }), + } + } + None => output_unbounded_with_progress_isolated(&mut command, progress) + .await + .map(|output| ExecOutput { + stdout: output.stdout, + stderr: output.stderr, + exit_code: output.status.code(), + timed_out: false, + }), + } + }; + match run.await { + Ok(output) => ExecBackendOutcome::Completed(output), + Err(error) => ExecBackendOutcome::LaunchFailed { + error: error.to_string(), + }, + } + } +} diff --git a/crates/tinyagents-harness-tools/src/exec/native_backend_tests.rs b/crates/tinyagents-harness-tools/src/exec/native_backend_tests.rs new file mode 100644 index 000000000..071ef88b9 --- /dev/null +++ b/crates/tinyagents-harness-tools/src/exec/native_backend_tests.rs @@ -0,0 +1,234 @@ +use super::*; +use crate::exec::{ExecBackendCapability, ExecRequest}; +use std::{collections::BTreeMap, ffi::OsString, path::PathBuf, time::Duration}; + +fn request( + program: &str, + args: &[&str], + timeout: Option, + capability: bool, +) -> ExecRequest { + ExecRequest { + program: OsString::from(program), + args: args.iter().map(OsString::from).collect(), + action_dir: PathBuf::from("."), + environment: BTreeMap::new(), + timeout, + shell_command: None, + backend_selection: "native".into(), + backend_capability: if capability { + ExecBackendCapability::new(NativeExecCapability) + } else { + ExecBackendCapability::new(()) + }, + } +} + +#[cfg(target_os = "linux")] +fn process_is_gone_or_zombie(pid: i32) -> bool { + let Ok(stat) = std::fs::read_to_string(format!("/proc/{pid}/stat")) else { + return true; + }; + stat.rsplit_once(')') + .and_then(|(_, tail)| tail.split_whitespace().next()) + .is_some_and(|state| state == "Z") +} + +#[cfg(target_os = "linux")] +fn shell_quote(path: &std::path::Path) -> String { + format!("'{}'", path.display().to_string().replace('\'', "'\\''")) +} + +#[cfg(unix)] +#[tokio::test] +async fn native_backend_preserves_argv_and_captures_both_streams() { + use crate::exec::{ExecBackend, ExecBackendOutcome, ShellFlavor}; + let backend = NativeExecBackend::new(ShellFlavor::Posix); + let outcome = backend + .execute( + request( + "sh", + &["-c", "printf 'hello'; printf 'warning' >&2"], + None, + true, + ), + None, + None, + ) + .await; + let ExecBackendOutcome::Completed(output) = outcome else { + panic!("expected process output") + }; + assert_eq!(output.stdout, b"hello"); + assert_eq!(output.stderr, b"warning"); + assert_eq!(output.exit_code, Some(0)); + assert!(!output.timed_out); +} + +#[tokio::test] +async fn native_backend_refuses_without_the_host_capability() { + use crate::exec::{ExecBackend, ExecBackendOutcome, ExecRefusalKind, ShellFlavor}; + let backend = NativeExecBackend::new(ShellFlavor::Posix); + let outcome = backend + .execute( + request( + "sh", + &["-c", "touch /tmp/tinyagents-should-not-run"], + None, + false, + ), + None, + None, + ) + .await; + assert!(matches!( + outcome, + ExecBackendOutcome::Refused { + kind: ExecRefusalKind::Preflight, + .. + } + )); +} + +#[cfg(unix)] +#[tokio::test] +async fn native_backend_does_not_reintroduce_ambient_turn_environment() { + use crate::exec::{ExecBackend, ExecBackendOutcome, ShellFlavor, timeout::CommandEnvironment}; + let backend = NativeExecBackend::new(ShellFlavor::Posix); + let authorized = + CommandEnvironment::new([("DENIED_FROM_TURN".to_owned(), "secret".to_owned())]); + let outcome=authorized.scope(backend.execute(request("sh", &["-c", "if [ -z \"${DENIED_FROM_TURN:-}\" ]; then printf clean; else printf leaked; fi"],None,true),None,None)).await; + let ExecBackendOutcome::Completed(output) = outcome else { + panic!("expected process output") + }; + assert_eq!(output.stdout, b"clean"); +} + +#[cfg(unix)] +#[tokio::test] +async fn native_backend_timeout_kills_the_process_group_before_returning() { + use crate::exec::{ExecBackend, ExecBackendOutcome, ShellFlavor}; + let backend = NativeExecBackend::new(ShellFlavor::Posix); + let outcome = backend + .execute( + request( + "sh", + &["-c", "printf partial; sleep 10"], + Some(Duration::from_millis(40)), + true, + ), + None, + None, + ) + .await; + let ExecBackendOutcome::Completed(output) = outcome else { + panic!("expected timeout result") + }; + assert!(output.timed_out); + assert_eq!(output.stdout, b"partial"); +} + +#[cfg(target_os = "linux")] +#[tokio::test] +async fn native_timeout_returns_after_descendant_and_pipes_are_cleaned() { + use crate::exec::{ExecBackend, ExecBackendOutcome, ShellFlavor}; + let dir = tempfile::tempdir().unwrap(); + let pidfile = dir.path().join("child.pid"); + let command = format!( + "sleep 30 & child=$!; echo $child > {}; printf captured; wait", + shell_quote(&pidfile), + ); + let backend = NativeExecBackend::new(ShellFlavor::Posix); + let cleanup = crate::exec::timeout::ProcessCleanup::default(); + let task_cleanup = cleanup.clone(); + let task = tokio::spawn(async move { + task_cleanup + .scope(backend.execute( + request("sh", &["-c", &command], Some(Duration::from_secs(5)), true), + None, + None, + )) + .await + }); + tokio::time::timeout(Duration::from_secs(3), async { + loop { + if std::fs::read_to_string(&pidfile) + .ok() + .and_then(|pid| pid.trim().parse::().ok()) + .is_some() + { + break; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("native descendant started before timeout"); + let outcome = tokio::time::timeout(Duration::from_secs(8), task) + .await + .expect("native timeout returned") + .expect("native process task completed"); + cleanup.wait().await; + let ExecBackendOutcome::Completed(output) = outcome else { + panic!("expected timeout completion") + }; + assert!(output.timed_out); + assert_eq!(output.stdout, b"captured"); + let pid: i32 = std::fs::read_to_string(pidfile) + .unwrap() + .trim() + .parse() + .unwrap(); + assert!( + process_is_gone_or_zombie(pid), + "descendant {pid} was still executing after timeout returned" + ); +} + +#[cfg(target_os = "linux")] +#[tokio::test] +async fn cancellation_waits_for_native_process_tree_reaping() { + use crate::exec::{ExecBackend, ShellFlavor, timeout::ProcessCleanup}; + let dir = tempfile::tempdir().unwrap(); + let pidfile = dir.path().join("child.pid"); + let command = format!( + "sleep 30 & child=$!; echo $child > {}; wait", + shell_quote(&pidfile), + ); + let backend = NativeExecBackend::new(ShellFlavor::Posix); + let cleanup = ProcessCleanup::default(); + let turn_cleanup = cleanup.clone(); + let turn = tokio::spawn(async move { + turn_cleanup + .scope(backend.execute(request("sh", &["-c", &command], None, true), None, None)) + .await + }); + tokio::time::timeout(Duration::from_secs(3), async { + loop { + if std::fs::read_to_string(&pidfile) + .ok() + .and_then(|pid| pid.trim().parse::().ok()) + .is_some() + { + break; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("native child began before cancellation"); + turn.abort(); + let _ = turn.await; + tokio::time::timeout(Duration::from_secs(3), cleanup.wait()) + .await + .expect("native process reaped before cancellation acknowledgment"); + let pid: i32 = std::fs::read_to_string(pidfile) + .unwrap() + .trim() + .parse() + .unwrap(); + assert!( + process_is_gone_or_zombie(pid), + "descendant {pid} was still executing after cleanup acknowledgment" + ); +} diff --git a/crates/tinyagents-harness-tools/src/exec/ports.rs b/crates/tinyagents-harness-tools/src/exec/ports.rs new file mode 100644 index 000000000..e7cb9165a --- /dev/null +++ b/crates/tinyagents-harness-tools/src/exec/ports.rs @@ -0,0 +1,151 @@ +use std::{ + any::Any, collections::BTreeMap, ffi::OsString, path::PathBuf, sync::Arc, time::Duration, +}; + +use async_trait::async_trait; +use serde_json::Value; +use tinytools::ToolRunContext; + +use super::ExecProgress; + +/// A concrete operation presented to the host's execution policy. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum ExecOperation { + Shell { + command: String, + category: Option, + }, + Install { + package: String, + manager: String, + }, +} + +#[derive(Debug, Clone)] +pub struct ExecAuthorizationRequest { + pub operation: ExecOperation, + pub timeout: Option, +} + +/// Opaque host-owned data selected during authorization. +#[derive(Clone)] +pub struct ExecBackendCapability(Arc); + +impl ExecBackendCapability { + #[must_use] + pub fn new(value: T) -> Self { + Self(Arc::new(value)) + } + + #[must_use] + pub fn downcast_ref(&self) -> Option<&T> { + self.0.as_ref().downcast_ref() + } +} + +#[derive(Clone)] +pub struct ExecPermit { + pub action_dir: PathBuf, + pub environment: BTreeMap, + pub backend_selection: String, + pub backend_capability: ExecBackendCapability, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ExecDenied(pub String); + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum ExecRefusalKind { + BackendUnavailable, + Configuration, + Preflight, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum ExecAuditOutcome { + GateDenied, + /// The host authorized an install, but the requested or automatic package + /// manager could not be resolved. No process launch was attempted. + InstallManagerResolutionFailed, + BackendRefused { + kind: ExecRefusalKind, + }, + LaunchFailed, + Completed { + success: bool, + timed_out: bool, + }, +} + +#[derive(Debug, Clone)] +pub struct ExecAuditEvent { + pub operation: ExecOperation, + pub approval_required: bool, + pub authorized: bool, + pub allowed: bool, + pub backend_selection: Option, + pub process_attempted: bool, + pub outcome: ExecAuditOutcome, + pub duration: Duration, +} + +#[derive(Clone)] +pub struct ExecRequest { + pub program: OsString, + pub args: Vec, + pub action_dir: PathBuf, + pub environment: BTreeMap, + pub timeout: Option, + pub shell_command: Option, + pub backend_selection: String, + pub backend_capability: ExecBackendCapability, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ExecOutput { + pub stdout: Vec, + pub stderr: Vec, + pub exit_code: Option, + pub timed_out: bool, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum ExecBackendOutcome { + Refused { + kind: ExecRefusalKind, + reason: String, + }, + LaunchFailed { + error: String, + }, + Completed(ExecOutput), +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ShellFlavor { + Posix, + Cmd, +} + +#[async_trait] +pub trait ExecGate: Send + Sync { + async fn prepare_authorization(&self, args: &Value) -> anyhow::Result>; + fn external_effect_with_args(&self, args: &Value) -> bool; + async fn authorize( + &self, + request: &ExecAuthorizationRequest, + context: Option<&dyn ToolRunContext>, + ) -> Result; + fn record(&self, event: ExecAuditEvent); +} + +#[async_trait] +pub trait ExecBackend: Send + Sync { + fn shell_flavor(&self) -> ShellFlavor; + async fn execute( + &self, + request: ExecRequest, + context: Option<&dyn ToolRunContext>, + progress: Option, + ) -> ExecBackendOutcome; +} diff --git a/crates/tinyagents-harness-tools/src/exec/ports_tests.rs b/crates/tinyagents-harness-tools/src/exec/ports_tests.rs new file mode 100644 index 000000000..0fdc69b79 --- /dev/null +++ b/crates/tinyagents-harness-tools/src/exec/ports_tests.rs @@ -0,0 +1,14 @@ +use super::ports::{ExecBackend, ExecGate}; +use std::sync::Arc; + +#[test] +fn host_ports_are_object_safe_and_thread_safe() { + fn assert_send_sync() {} + fn assert_gate_object(_: Option>) {} + fn assert_backend_object(_: Option>) {} + + assert_send_sync::>(); + assert_send_sync::>(); + assert_gate_object(None); + assert_backend_object(None); +} diff --git a/crates/tinyagents-harness-tools/src/exec/progress.rs b/crates/tinyagents-harness-tools/src/exec/progress.rs new file mode 100644 index 000000000..a7e3f4a97 --- /dev/null +++ b/crates/tinyagents-harness-tools/src/exec/progress.rs @@ -0,0 +1,150 @@ +//! Progress for a running shell call. +//! +//! A shell call can legitimately run for many minutes (a build, a test suite, +//! a download), and until it exits the turn shows nothing. While it runs, +//! `with_exec_progress` reports a status line through the harness's +//! [`ToolRunContext::report_progress`] seam every [`HEARTBEAT_INTERVAL`]: the +//! elapsed time, how much output the command has written, and its last output +//! line. The agent loop turns each report into `AgentEvent::ToolProgressDetail` +//! for the call, which the web channel projects as `tool_progress`. +//! +//! The output figures come from [`ExecProgress`], which +//! `shell_exec::collect_process` feeds with every chunk it reads from the +//! process. Paths that do not stream output (the Docker backend) still report +//! the elapsed time. + +use std::sync::{Arc, Mutex}; +use std::time::Duration; + +use tinytools::{ToolProgress, ToolRunContext}; + +/// How often a running shell call reports that it is still working. +pub const HEARTBEAT_INTERVAL: Duration = Duration::from_secs(15); + +/// Longest last-output line carried in a status line, in characters. +const LAST_LINE_CHARS: usize = 120; + +/// Most characters of an unfinished output line kept until its terminator +/// arrives. Only the first [`LAST_LINE_CHARS`] are ever shown, so this bounds +/// memory without changing what a status line says. +const PENDING_LINE_CHARS: usize = 1_024; + +#[derive(Debug, Default)] +struct Meter { + bytes: u64, + /// The last complete, non-empty output line. + last_line: Option, + /// The start of the line still being written. Output arrives in chunks + /// that can split a line anywhere, so a line only counts once its + /// terminator arrives; otherwise a status line could quote half of one. + pending: String, + pending_chars: usize, +} + +/// Output written so far by the shell call running on this task. Cheap to +/// clone; every clone sees the same figures. +#[derive(Clone, Debug, Default)] +pub struct ExecProgress(Arc>); + +impl ExecProgress { + /// Fold one chunk of stdout or stderr into the figures. `\n` and `\r` + /// both end a line, so a progress bar redrawn with `\r` reports its + /// latest state. + pub fn record(&self, chunk: &[u8]) { + let mut meter = self.0.lock().unwrap_or_else(|p| p.into_inner()); + meter.bytes = meter.bytes.saturating_add(chunk.len() as u64); + let text = String::from_utf8_lossy(chunk); + for (index, piece) in text.split(['\n', '\r']).enumerate() { + if index > 0 { + // A terminator preceded this piece: the pending line is done. + let line = std::mem::take(&mut meter.pending); + meter.pending_chars = 0; + let line = line.trim(); + if !line.is_empty() { + meter.last_line = Some(line.chars().take(LAST_LINE_CHARS).collect()); + } + } + let room = PENDING_LINE_CHARS - meter.pending_chars; + let kept: String = piece.chars().take(room).collect(); + meter.pending_chars += kept.chars().count(); + meter.pending.push_str(&kept); + } + } + + pub fn snapshot(&self) -> (u64, Option) { + let meter = self.0.lock().unwrap_or_else(|p| p.into_inner()); + (meter.bytes, meter.last_line.clone()) + } +} + +/// The status line for a call `elapsed` into its run. +pub(crate) fn progress_message(elapsed: Duration, bytes: u64, last_line: Option<&str>) -> String { + let secs = elapsed.as_secs(); + let elapsed = if secs >= 60 { + format!("{}m {:02}s", secs / 60, secs % 60) + } else { + format!("{secs}s") + }; + if bytes == 0 { + return format!("running {elapsed} · no output yet"); + } + let size = if bytes >= 1024 * 1024 { + format!("{:.1} MB", bytes as f64 / (1024.0 * 1024.0)) + } else if bytes >= 1024 { + format!("{:.1} KB", bytes as f64 / 1024.0) + } else { + format!("{bytes} B") + }; + match last_line { + Some(line) => format!("running {elapsed} · {size} output · last line: {line}"), + None => format!("running {elapsed} · {size} output"), + } +} + +/// Run the shell call built by `build`, reporting its progress through +/// `context` every `interval` until it finishes. `build` receives the call's +/// [`ExecProgress`] and hands it down to the process runner, which feeds it; +/// without a context there is nowhere to report, so it receives `None` and the +/// call runs as is. +pub(crate) async fn with_exec_progress( + context: Option<&dyn ToolRunContext>, + interval: Duration, + build: B, +) -> F::Output +where + B: FnOnce(Option) -> F, + F: std::future::Future, +{ + let Some(context) = context else { + return build(None).await; + }; + let meter = ExecProgress::default(); + let fut = build(Some(meter.clone())); + let started = tokio::time::Instant::now(); + // Boxed: the shell call's future is large, and this wrapper would + // otherwise hold it inline a second time on the tool's stack. + let mut work = Box::pin(fut); + let mut ticker = tokio::time::interval_at(tokio::time::Instant::now() + interval, interval); + loop { + tokio::select! { + output = &mut work => return output, + _ = ticker.tick() => { + let (bytes, last_line) = meter.snapshot(); + let elapsed = started.elapsed(); + // The status line can quote command output, so only its + // content-free figures are logged. + tracing::debug!( + elapsed_secs = elapsed.as_secs(), + output_bytes = bytes, + "[shell] reporting progress" + ); + let message = progress_message(elapsed, bytes, last_line.as_deref()); + context.report_progress(ToolProgress::message(message)); + } + } + } +} + +#[cfg(test)] +#[path = "progress_tests.rs"] +mod tests; diff --git a/crates/tinyagents-harness-tools/src/exec/progress_tests.rs b/crates/tinyagents-harness-tools/src/exec/progress_tests.rs new file mode 100644 index 000000000..881cb4de7 --- /dev/null +++ b/crates/tinyagents-harness-tools/src/exec/progress_tests.rs @@ -0,0 +1,116 @@ +use super::*; + +/// A context that records what the tool reported. +#[derive(Default)] +struct Recorder { + updates: Mutex>, +} + +impl ToolRunContext for Recorder { + fn report_progress(&self, update: ToolProgress) { + self.updates + .lock() + .unwrap() + .push(update.message.unwrap_or_default()); + } +} + +#[test] +fn message_names_elapsed_time_output_size_and_last_line() { + assert_eq!( + progress_message(Duration::from_secs(45), 0, None), + "running 45s · no output yet" + ); + assert_eq!( + progress_message( + Duration::from_secs(252), + 39_117, + Some("Compiling serde v1.0") + ), + "running 4m 12s · 38.2 KB output · last line: Compiling serde v1.0" + ); + assert_eq!( + progress_message(Duration::from_secs(61), 3 * 1024 * 1024, None), + "running 1m 01s · 3.0 MB output" + ); +} + +#[test] +fn meter_counts_bytes_and_keeps_the_last_non_empty_line() { + let meter = ExecProgress::default(); + meter.record(b"one\ntwo\n\n"); + meter.record(b" three \n"); + meter.record(b"\n"); + assert_eq!(meter.snapshot(), (20, Some("three".to_string()))); +} + +#[test] +fn meter_truncates_a_long_last_line() { + let meter = ExecProgress::default(); + meter.record(format!("{}\n", "x".repeat(5_000)).as_bytes()); + let (_, line) = meter.snapshot(); + assert_eq!(line.unwrap().chars().count(), LAST_LINE_CHARS); +} + +/// A line split across chunks is reported whole, never from its second half. +#[test] +fn meter_joins_a_line_split_across_chunks() { + let meter = ExecProgress::default(); + meter.record(b"Compil"); + assert_eq!( + meter.snapshot().1, + None, + "an unfinished line is not reported" + ); + meter.record(b"ing serde v1.0\nDown"); + assert_eq!(meter.snapshot().1.as_deref(), Some("Compiling serde v1.0")); + meter.record(b"loading\n"); + assert_eq!(meter.snapshot().1.as_deref(), Some("Downloading")); +} + +/// A progress bar redrawn with carriage returns reports its latest state. +#[test] +fn meter_treats_a_carriage_return_as_a_line_end() { + let meter = ExecProgress::default(); + meter.record(b"\r 10%\r 55%\r"); + assert_eq!(meter.snapshot().1.as_deref(), Some("55%")); +} + +/// A 40s command that writes output reports twice, with the output it has +/// written so far, through the meter handed to the call. +#[tokio::test(start_paused = true)] +async fn a_long_call_reports_its_output_as_it_runs() { + let recorder = Recorder::default(); + let work = |meter: Option| async move { + let meter = meter.expect("meter handed to the call"); + meter.record(b"Compiling a\n"); + tokio::time::sleep(Duration::from_secs(20)).await; + meter.record(b"Compiling b\n"); + tokio::time::sleep(Duration::from_secs(20)).await; + "done" + }; + let output = with_exec_progress(Some(&recorder), HEARTBEAT_INTERVAL, work).await; + assert_eq!(output, "done"); + let updates = recorder.updates.lock().unwrap().clone(); + assert_eq!( + updates, + vec![ + "running 15s · 12 B output · last line: Compiling a".to_string(), + "running 30s · 24 B output · last line: Compiling b".to_string(), + ] + ); +} + +#[tokio::test(start_paused = true)] +async fn a_short_call_reports_nothing() { + let recorder = Recorder::default(); + let output = with_exec_progress(Some(&recorder), HEARTBEAT_INTERVAL, |_| async { 7 }).await; + assert_eq!(output, 7); + assert!(recorder.updates.lock().unwrap().is_empty()); +} + +#[tokio::test] +async fn without_a_context_no_meter_is_handed_out() { + let seen = with_exec_progress(None, HEARTBEAT_INTERVAL, |meter| async move { meter }).await; + assert!(seen.is_none()); +} diff --git a/crates/tinyagents-harness-tools/src/exec/shell_description.rs b/crates/tinyagents-harness-tools/src/exec/shell_description.rs new file mode 100644 index 000000000..f4e24a33a --- /dev/null +++ b/crates/tinyagents-harness-tools/src/exec/shell_description.rs @@ -0,0 +1,20 @@ +use super::ShellFlavor; + +pub(super) fn shell_description(flavor: ShellFlavor) -> &'static str { + match flavor { + ShellFlavor::Cmd => { + "Execute a command under Windows cmd.exe (`cmd /C`), not a POSIX shell: run code, \ + manipulate workspace files, or launch applications (`start \"\" music:`). Use cmd \ + syntax: `cd` with no argument prints the current directory (no `pwd`), `dir` lists \ + (no `ls`), `type` prints a file (no `cat`), `findstr` searches (no `grep`), \ + `del`/`copy`/`move` replace rm/cp/mv, variables are `%VAR%` (not `$VAR`), and \ + commands chain with `&&` (`;` is not a separator). For PowerShell run \ + `powershell -NoProfile -Command \"...\"`. Only stdout/stderr comes back, so a script \ + that computes silently or only writes a file returns nothing — print what you need, \ + or read the file afterwards." + } + ShellFlavor::Posix => { + "Execute a shell command: run code, manipulate workspace files, or launch applications (`open -a Music`, `xdg-open music://`). Only stdout/stderr comes back, so a script that computes silently or only writes a file returns nothing — print what you need, or read the file afterwards." + } + } +} diff --git a/crates/tinyagents-harness-tools/src/exec/shell_schema_cmd.json b/crates/tinyagents-harness-tools/src/exec/shell_schema_cmd.json new file mode 100644 index 000000000..43f925e3f --- /dev/null +++ b/crates/tinyagents-harness-tools/src/exec/shell_schema_cmd.json @@ -0,0 +1,9 @@ +{ + "type": "object", + "properties": { + "command": { "type": "string", "description": "The cmd.exe command to execute (cmd syntax, not POSIX)" }, + "category": { "type": "string", "enum": ["read", "write", "network", "install", "destructive"], "description": "Optional self-declared risk; can only raise the approval requirement." }, + "timeout_secs": { "type": "integer", "minimum": 0, "maximum": 3600, "description": "Optional wall-clock timeout (seconds, up to 3600) for this command before it is killed. Use a larger value for long-running work (builds, test suites, solvers). Omitted or 0 runs without a deadline." } + }, + "required": ["command"] +} diff --git a/crates/tinyagents-harness-tools/src/exec/shell_schema_posix.json b/crates/tinyagents-harness-tools/src/exec/shell_schema_posix.json new file mode 100644 index 000000000..147cfb718 --- /dev/null +++ b/crates/tinyagents-harness-tools/src/exec/shell_schema_posix.json @@ -0,0 +1,9 @@ +{ + "type": "object", + "properties": { + "command": { "type": "string", "description": "The shell command to execute" }, + "category": { "type": "string", "enum": ["read", "write", "network", "install", "destructive"], "description": "Optional self-declared risk; can only raise the approval requirement." }, + "timeout_secs": { "type": "integer", "minimum": 0, "maximum": 3600, "description": "Optional wall-clock timeout (seconds, up to 3600) for this command before it is killed. Use a larger value for long-running work (builds, test suites, solvers). Omitted or 0 runs without a deadline." } + }, + "required": ["command"] +} diff --git a/crates/tinyagents-harness-tools/src/exec/shell_support.rs b/crates/tinyagents-harness-tools/src/exec/shell_support.rs new file mode 100644 index 000000000..cfe90c0f5 --- /dev/null +++ b/crates/tinyagents-harness-tools/src/exec/shell_support.rs @@ -0,0 +1,94 @@ +use std::{collections::BTreeMap, ffi::OsString, time::Duration}; + +use tinytools::{ToolContent, ToolResult}; + +use super::ExecOutput; + +pub(super) const MAX_OUTPUT_BYTES: usize = 1_048_576; +pub(super) const COMMAND_TIMEOUT_EXIT_CODE: i32 = 124; +pub(super) const COMMAND_TIMEOUT_NOTE: &str = "[exit 124: the command's own `timeout` limit expired; the exec tool's timeout_secs did not fire]"; +const OUTPUT_TRUNCATED: &str = "\n... [output truncated at 1MB]"; +const STDERR_TRUNCATED: &str = "\n... [stderr truncated at 1MB]"; + +pub(super) fn tool_timeout_message(secs: u64) -> String { + format!( + "Command timed out after {secs}s and was killed: the exec tool's timeout_secs limit fired." + ) +} + +pub(super) const PYTHON_UTF8_DEFAULTS: &[(&str, &str)] = + &[("PYTHONUTF8", "1"), ("PYTHONIOENCODING", "utf-8")]; + +pub(super) fn python_utf8_env( + lookup: impl Fn(&str) -> Option, +) -> Vec<(&'static str, OsString)> { + PYTHON_UTF8_DEFAULTS + .iter() + .map(|(name, default)| (*name, lookup(name).unwrap_or_else(|| (*default).into()))) + .collect() +} + +pub(super) fn command_failure(code: Option, stdout: &str, stderr: &str) -> ToolResult { + let mut failure = tinytools::command_failure(code, stdout, stderr); + if code == Some(COMMAND_TIMEOUT_EXIT_CODE) { + tracing::debug!("[exec] exit 124: attributing it to the command's own timeout"); + failure.content.push(ToolContent::Text { + text: COMMAND_TIMEOUT_NOTE.to_string(), + }); + } + failure +} + +fn bounded_stream(bytes: &[u8], marker: &str) -> String { + let text = String::from_utf8_lossy(bytes); + if text.len() <= MAX_OUTPUT_BYTES { + return text.into_owned(); + } + let mut end = MAX_OUTPUT_BYTES; + while !text.is_char_boundary(end) { + end -= 1; + } + format!("{}{marker}", &text[..end]) +} + +pub(super) fn exec_result(output: ExecOutput, explicit_timeout: Option) -> ToolResult { + let stdout = bounded_stream(&output.stdout, OUTPUT_TRUNCATED); + let stderr = bounded_stream(&output.stderr, STDERR_TRUNCATED); + if output.timed_out { + let secs = explicit_timeout.map_or(0, |timeout| timeout.as_secs()); + let mut message = tool_timeout_message(secs); + if !stdout.is_empty() { + message.push_str("\n[stdout]\n"); + message.push_str(&stdout); + } + if !stderr.is_empty() { + message.push_str("\n[stderr]\n"); + message.push_str(&stderr); + } + return ToolResult::error(message); + } + if output.exit_code == Some(0) { + if stderr.is_empty() { + ToolResult::success(stdout) + } else { + ToolResult::success(format!("{stdout}\n[stderr]\n{stderr}")) + } + } else { + command_failure(output.exit_code, &stdout, &stderr) + } +} + +pub(super) fn python_environment( + mut environment: BTreeMap, +) -> BTreeMap { + for (name, value) in + python_utf8_env(|name| environment.get(std::ffi::OsStr::new(name)).cloned()) + { + environment.insert(name.into(), value); + } + environment +} + +#[cfg(test)] +#[path = "shell_support_tests.rs"] +mod tests; diff --git a/crates/tinyagents-harness-tools/src/exec/shell_support_tests.rs b/crates/tinyagents-harness-tools/src/exec/shell_support_tests.rs new file mode 100644 index 000000000..d5bce0813 --- /dev/null +++ b/crates/tinyagents-harness-tools/src/exec/shell_support_tests.rs @@ -0,0 +1,106 @@ +use super::*; +use crate::exec::ExecOutput; +use std::ffi::OsString; + +fn output(stdout: Vec, stderr: Vec, exit_code: Option, timed_out: bool) -> ExecOutput { + ExecOutput { + stdout, + stderr, + exit_code, + timed_out, + } +} + +#[test] +fn explicit_timeout_preserves_captured_streams() { + let result = exec_result( + output( + b"partial stdout".to_vec(), + b"partial stderr".to_vec(), + None, + true, + ), + Some(Duration::from_secs(17)), + ); + assert!(result.is_error); + let text = result.text(); + assert!(text.contains("Command timed out after 17s")); + assert!(text.contains("[stdout]\npartial stdout")); + assert!(text.contains("[stderr]\npartial stderr")); +} + +#[test] +fn successful_stderr_is_preserved() { + let result = exec_result( + output(b"out".to_vec(), b"warning".to_vec(), Some(0), false), + None, + ); + assert!(!result.is_error); + assert_eq!(result.text(), "out\n[stderr]\nwarning"); +} + +#[test] +fn command_failure_preserves_exit_and_both_streams() { + let result = exec_result( + output(b"out".to_vec(), b"err".to_vec(), Some(7), false), + None, + ); + assert!(result.is_error); + let text = result.text(); + assert!(text.contains("exit code 7")); + assert!(text.contains("out")); + assert!(text.contains("err")); +} + +#[test] +fn exit_124_identifies_the_command_timeout() { + let result = exec_result(output(b"done".to_vec(), Vec::new(), Some(124), false), None); + assert!(result.is_error); + assert!(result.text().contains(COMMAND_TIMEOUT_NOTE)); +} + +#[test] +fn each_stream_is_truncated_at_a_utf8_boundary() { + let mut stdout = vec![b'a'; MAX_OUTPUT_BYTES - 1]; + stdout.extend_from_slice("é".as_bytes()); + let mut stderr = stdout.clone(); + stderr.extend_from_slice(b"extra"); + let result = exec_result(output(stdout, stderr, Some(0), false), None); + assert!(!result.is_error); + let text = result.text(); + assert!(text.contains(OUTPUT_TRUNCATED)); + assert!(text.contains(STDERR_TRUNCATED)); + assert!(!text.contains('�')); +} + +#[test] +fn python_utf8_defaults_keep_explicit_values() { + let env = python_environment(BTreeMap::from([ + (OsString::from("PYTHONUTF8"), OsString::from("0")), + (OsString::from("KEEP"), OsString::from("value")), + ])); + assert_eq!( + env.get(std::ffi::OsStr::new("PYTHONUTF8")), + Some(&OsString::from("0")) + ); + assert_eq!( + env.get(std::ffi::OsStr::new("PYTHONIOENCODING")), + Some(&OsString::from("utf-8")) + ); + assert_eq!( + env.get(std::ffi::OsStr::new("KEEP")), + Some(&OsString::from("value")) + ); +} + +#[test] +fn python_utf8_lookup_fills_only_missing_defaults() { + let env = python_utf8_env(|name| (name == "PYTHONUTF8").then(|| OsString::from("0"))); + assert_eq!( + env, + vec![ + ("PYTHONUTF8", OsString::from("0")), + ("PYTHONIOENCODING", OsString::from("utf-8")), + ] + ); +} diff --git a/crates/tinyagents-harness-tools/src/exec/timeout/command_environment.rs b/crates/tinyagents-harness-tools/src/exec/timeout/command_environment.rs new file mode 100644 index 000000000..92bcb2620 --- /dev/null +++ b/crates/tinyagents-harness-tools/src/exec/timeout/command_environment.rs @@ -0,0 +1,70 @@ +//! Child-process environment carried by one dispatch task. + +use std::{collections::BTreeMap, future::Future, sync::Arc}; + +tokio::task_local! { static ACTIVE: CommandEnvironment; } + +/// A complete, turn-owned environment for builtin tool subprocesses. +/// This never mutates the host process environment. Host-spawned tasks must +/// explicitly inherit its scope. Interpreter pools cannot accept scoped work. +#[derive(Clone, Debug)] +pub struct CommandEnvironment(Arc>); + +impl CommandEnvironment { + /// Own the exact variables permitted in child processes. + #[must_use] + pub fn new(env: impl IntoIterator) -> Self { + Self(Arc::new(env.into_iter().collect())) + } + /// Whether this task requires its own child environment. + #[must_use] + pub fn is_active() -> bool { + ACTIVE.try_with(|_| ()).is_ok() + } + /// Run a future using this environment for builtin subprocesses. + pub async fn scope(&self, future: F) -> F::Output { + ACTIVE.scope(self.clone(), future).await + } + /// Read a functional child variable from the scoped map, falling back to + /// the process only when no map was supplied. + #[must_use] + pub fn var_os(name: &str) -> Option { + ACTIVE + .try_with(|environment| environment.0.get(name).map(Into::into)) + .unwrap_or_else(|_| std::env::var_os(name)) + } + /// UTF-8 counterpart used by managed interpreter command builders. + pub fn var(name: &str) -> Result { + ACTIVE + .try_with(|environment| { + environment + .0 + .get(name) + .cloned() + .ok_or(std::env::VarError::NotPresent) + }) + .unwrap_or_else(|_| std::env::var(name)) + } + pub(super) fn apply(command: &mut tokio::process::Command) { + let _ = ACTIVE.try_with(|environment| { + // Explicit runtime/security additions (managed PATH, scratch TMP, + // Git policy) remain authoritative over the base environment. + let explicit: Vec<_> = command + .as_std() + .get_envs() + .map(|(key, value)| (key.to_owned(), value.map(ToOwned::to_owned))) + .collect(); + command.env_clear().envs(environment.0.iter()); + for (key, value) in explicit { + match value { + Some(value) => { + command.env(key, value); + } + None => { + command.env_remove(key); + } + } + } + }); + } +} diff --git a/crates/tinyagents-harness-tools/src/exec/timeout/group_exit.rs b/crates/tinyagents-harness-tools/src/exec/timeout/group_exit.rs new file mode 100644 index 000000000..a30cc7e1c --- /dev/null +++ b/crates/tinyagents-harness-tools/src/exec/timeout/group_exit.rs @@ -0,0 +1,59 @@ +//! Wait for signalled Linux process-group members to stop executing. +//! +//! The caller keeps the group leader unreaped during this check, reserving +//! its PID and preventing a new process group from reusing the identifier. + +/// Wait until the reserved group contains no executing Linux processes. +#[cfg(target_os = "linux")] +pub(super) async fn wait(pid: u32) -> std::io::Result<()> { + loop { + let active = tokio::task::spawn_blocking(move || active_members(pid)) + .await + .map_err(std::io::Error::other)??; + if !active { + return Ok(()); + } + tokio::time::sleep(std::time::Duration::from_millis(1)).await; + } +} + +/// Other platforms retain their existing direct-child acknowledgement. +#[cfg(not(target_os = "linux"))] +pub(super) async fn wait(_: u32) -> std::io::Result<()> { + Ok(()) +} + +/// Inspect kernel process states without reaping the reserved group leader. +#[cfg(target_os = "linux")] +fn active_members(group: u32) -> std::io::Result { + for entry in std::fs::read_dir("/proc")? { + let entry = entry?; + if entry.file_name().to_string_lossy().parse::().is_err() { + continue; + } + // Process entries can disappear or become unreadable between + // enumeration and reading. + if let Ok(stat) = std::fs::read_to_string(entry.path().join("stat")) { + if is_active_member(&stat, group) { + return Ok(true); + } + } + } + Ok(false) +} + +/// Zombies and dead tasks cannot execute further side effects. +#[cfg(target_os = "linux")] +pub(super) fn is_active_member(stat: &str, group: u32) -> bool { + let Some((_, fields)) = stat.rsplit_once(") ") else { + return false; + }; + let mut fields = fields.split_whitespace(); + let state = fields.next(); + let process_group = fields.nth(1).and_then(|value| value.parse::().ok()); + process_group == Some(group) && !matches!(state, Some("Z" | "X" | "x")) +} + +#[cfg(all(test, target_os = "linux"))] +#[path = "group_exit_tests.rs"] +mod tests; diff --git a/crates/tinyagents-harness-tools/src/exec/timeout/group_exit_tests.rs b/crates/tinyagents-harness-tools/src/exec/timeout/group_exit_tests.rs new file mode 100644 index 000000000..9b20c5bf8 --- /dev/null +++ b/crates/tinyagents-harness-tools/src/exec/timeout/group_exit_tests.rs @@ -0,0 +1,44 @@ +#![allow(clippy::unwrap_used, clippy::expect_used)] +use super::*; + +#[test] +fn exited_members_do_not_block_cleanup_but_running_members_do() { + for state in ["R", "S", "D", "T", "I"] { + assert!(is_active_member( + &format!("12 (name with ) spaces) {state} 1 42 42"), + 42 + )); + } + for state in ["Z", "X", "x"] { + assert!(!is_active_member( + &format!("12 (sleep) {state} 1 42 42"), + 42 + )); + } + assert!(!is_active_member("12 (sleep) R 1 43 43", 42)); + assert!(!is_active_member("malformed", 42)); +} + +#[tokio::test] +async fn wait_returns_immediately_for_a_group_with_no_members() { + // PID 0x7fff_fff0 is above any pid_max, so no process group matches. + wait(0x7fff_fff0).await.unwrap(); + assert!(!active_members(0x7fff_fff0).unwrap()); +} + +#[tokio::test] +async fn wait_polls_until_the_group_has_stopped() { + let mut cmd = tokio::process::Command::new("sh"); + cmd.arg("-c").arg("sleep 0.3"); + cmd.process_group(0); + let mut child = cmd.spawn().unwrap(); + let pid = child.id().unwrap(); + assert!(active_members(pid).unwrap()); + // The group leader stays unreaped, so `wait` only returns once the + // member has exited (zombie) after polling at least once. + tokio::time::timeout(std::time::Duration::from_secs(10), wait(pid)) + .await + .expect("group stops") + .unwrap(); + child.wait().await.unwrap(); +} diff --git a/crates/tinyagents-harness-tools/src/exec/timeout/mod.rs b/crates/tinyagents-harness-tools/src/exec/timeout/mod.rs new file mode 100644 index 000000000..924f38bb7 --- /dev/null +++ b/crates/tinyagents-harness-tools/src/exec/timeout/mod.rs @@ -0,0 +1,41 @@ +//! Tool-execution timeouts and child-process management. +//! +//! - [`settings`]: bounds, resolution and the process-global inherited timeout. +//! The host reads its config and env, then calls [`configure`]. +//! - [`output_or_kill`] / [`output_unbounded`] / [`output_with_input`]: capture a +//! command's output, killing its whole process group on a deadline or drop. +//! - [`ProcessCleanup`]: lets an embedder await reaping after cancelling a turn; +//! [`CleanupHold`] extends that wait to work a host tool hands off to its own task. +//! - [`CommandEnvironment`]: a turn-owned child environment, never the host's. + +#![allow( + clippy::missing_errors_doc, + clippy::ignored_unit_patterns, + clippy::struct_field_names, + clippy::collapsible_if +)] + +mod command_environment; +mod group_exit; +mod output; +mod process_cleanup; +mod settings; + +pub use command_environment::CommandEnvironment; +pub use output::{ + CapturedCommandResult, kill_process_group, output_or_kill, output_or_kill_captured_isolated, + output_or_kill_with_progress, output_or_kill_with_progress_isolated, output_unbounded, + output_unbounded_with_progress, output_unbounded_with_progress_isolated, output_with_input, + output_with_input_and_progress, own_process_group, set_task_wrapper, +}; +pub use process_cleanup::{CleanupHold, ProcessCleanup}; +pub use settings::{ + DEFAULT_TIMEOUT_SECS, MAX_TIMEOUT_SECS, MIN_TIMEOUT_SECS, ResolvedTimeout, + SANDBOX_UNBOUNDED_CAP_SECS, TOOL_TIMEOUT_GRACE_SECS, TimeoutSettings, configure, + explicit_call_timeout_duration, explicit_call_timeout_secs, override_from, + parse_tool_timeout_secs, resolve_effective, settings, +}; + +#[cfg(test)] +#[path = "mod_tests.rs"] +mod tests; diff --git a/crates/tinyagents-harness-tools/src/exec/timeout/mod_tests.rs b/crates/tinyagents-harness-tools/src/exec/timeout/mod_tests.rs new file mode 100644 index 000000000..04af0b0d2 --- /dev/null +++ b/crates/tinyagents-harness-tools/src/exec/timeout/mod_tests.rs @@ -0,0 +1,425 @@ +#![allow( + clippy::expect_used, + clippy::panic, + clippy::unwrap_used, + clippy::cast_sign_loss +)] +use super::*; +use std::time::Duration; + +use tinytools::ToolTimeout; + +fn sh(script: &str) -> tokio::process::Command { + let mut cmd = tokio::process::Command::new("sh"); + cmd.arg("-c").arg(script); + cmd +} + +#[cfg(unix)] +fn process_exists(pid: i32) -> bool { + rustix::process::Pid::from_raw(pid) + .is_some_and(|pid| rustix::process::test_kill_process(pid).is_ok()) +} + +#[test] +fn configure_pushes_the_inherited_timeout_into_the_global() { + let global = configure(TimeoutSettings::with_inherited_secs(77)); + assert_eq!(global.inherited_timeout(), Duration::from_secs(77)); + configure(TimeoutSettings::with_inherited_secs(88)); + assert_eq!(settings().inherited_timeout(), Duration::from_secs(88)); + assert_eq!( + settings().resolve(ToolTimeout::Inherit).deadline, + Some(Duration::from_secs(88)) + ); + assert_eq!(settings().resolve(ToolTimeout::Unbounded).deadline, None); +} + +#[test] +fn default_when_env_missing() { + assert_eq!(parse_tool_timeout_secs(None), DEFAULT_TIMEOUT_SECS); +} + +#[test] +fn default_when_value_not_numeric() { + assert_eq!( + parse_tool_timeout_secs(Some("not-a-number")), + DEFAULT_TIMEOUT_SECS + ); + assert_eq!(parse_tool_timeout_secs(Some("")), DEFAULT_TIMEOUT_SECS); + assert_eq!(parse_tool_timeout_secs(Some("12x")), DEFAULT_TIMEOUT_SECS); +} + +#[test] +fn default_when_value_zero() { + // 0 seconds would disable the timeout — reject and fall back. + assert_eq!(parse_tool_timeout_secs(Some("0")), DEFAULT_TIMEOUT_SECS); +} + +#[test] +fn default_when_value_above_max() { + assert_eq!(parse_tool_timeout_secs(Some("3601")), DEFAULT_TIMEOUT_SECS); + assert_eq!( + parse_tool_timeout_secs(Some("99999999999")), + DEFAULT_TIMEOUT_SECS + ); +} + +#[test] +fn default_when_value_negative_or_signed() { + // Negative values fail u64 parse and fall back to default. + assert_eq!(parse_tool_timeout_secs(Some("-5")), DEFAULT_TIMEOUT_SECS); +} + +#[test] +fn accepts_valid_values_at_boundaries() { + assert_eq!(parse_tool_timeout_secs(Some("1")), MIN_TIMEOUT_SECS); + assert_eq!(parse_tool_timeout_secs(Some("3600")), MAX_TIMEOUT_SECS); +} + +#[test] +fn accepts_valid_midrange_value() { + assert_eq!(parse_tool_timeout_secs(Some("300")), 300); +} + +#[test] +fn env_override_takes_precedence_over_config() { + // When the env var holds a valid value it wins over the config value. + assert_eq!(resolve_effective(300, Some("600")), 600); +} + +#[test] +fn config_value_used_when_env_absent_or_invalid() { + // No env → config drives the effective value (bounded). + assert_eq!(resolve_effective(300, None), 300); + // Present-but-invalid env (non-numeric / 0 / out of range) is ignored, + // so the config value still applies. + assert_eq!(resolve_effective(300, Some("nonsense")), 300); + assert_eq!(resolve_effective(300, Some("0")), 300); + assert_eq!(resolve_effective(300, Some("4000")), 300); +} + +#[test] +fn config_value_is_bounded() { + // An out-of-range config value falls back to the default rather than + // being applied verbatim. + assert_eq!(resolve_effective(0, None), DEFAULT_TIMEOUT_SECS); + assert_eq!(resolve_effective(99_999, None), DEFAULT_TIMEOUT_SECS); +} + +#[test] +fn explicit_call_timeout_unbounded_when_absent_or_disabled() { + // No request, or an explicit 0, means "run unbounded" (None). + assert_eq!(explicit_call_timeout_secs(None, MAX_TIMEOUT_SECS), None); + assert_eq!(explicit_call_timeout_secs(Some(0), MAX_TIMEOUT_SECS), None); +} + +#[test] +fn explicit_call_timeout_enforces_and_clamps_request() { + // An in-range request is enforced verbatim. + assert_eq!( + explicit_call_timeout_secs(Some(900), MAX_TIMEOUT_SECS), + Some(900) + ); + // Below the floor clamps up to MIN. + assert_eq!( + explicit_call_timeout_secs(Some(0), MAX_TIMEOUT_SECS), + None, + "0 disables rather than clamping to MIN" + ); + // Above the cap clamps down to the cap. + assert_eq!( + explicit_call_timeout_secs(Some(99_999), MAX_TIMEOUT_SECS), + Some(MAX_TIMEOUT_SECS) + ); + // A tool with a tighter own ceiling (node/npm at 1800) clamps to it. + assert_eq!(explicit_call_timeout_secs(Some(99_999), 1800), Some(1800)); + assert_eq!(explicit_call_timeout_secs(Some(600), 1800), Some(600)); +} + +#[test] +fn explicit_call_timeout_duration_matches_secs() { + assert_eq!( + explicit_call_timeout_duration(Some(900), MAX_TIMEOUT_SECS), + Some(Duration::from_secs(900)) + ); + assert_eq!(explicit_call_timeout_duration(None, MAX_TIMEOUT_SECS), None); +} + +#[test] +fn override_from_rejects_invalid() { + assert_eq!(override_from(None), None); + assert_eq!(override_from(Some("")), None); + assert_eq!(override_from(Some("0")), None); + assert_eq!(override_from(Some("abc")), None); + assert_eq!(override_from(Some("3601")), None); + assert_eq!(override_from(Some("120")), Some(120)); +} + +/// Table of `(policy, inherited config secs, expected (deadline_ms, budget_secs))`. +/// Pins the deadline resolution. +#[test] +fn resolve_tool_deadline_table() { + use tinytools::ToolTimeout; + let ms = |m: u64| Some(Duration::from_millis(m)); + let cases: Vec<(ToolTimeout, u64, (Option, u64))> = vec![ + // Inherit: the global config value, no grace. + (ToolTimeout::Inherit, 120, (ms(120_000), 120)), + (ToolTimeout::Inherit, 1, (ms(1_000), 1)), + (ToolTimeout::Inherit, 3600, (ms(3_600_000), 3600)), + // Unbounded: no deadline, zero budget. + (ToolTimeout::Unbounded, 120, (None, 0)), + // Millis: rounded UP to whole seconds, clamped, +5s grace on deadline. + (ToolTimeout::Millis(0), 120, (ms(6_000), 1)), + (ToolTimeout::Millis(1), 120, (ms(6_000), 1)), + (ToolTimeout::Millis(1_000), 120, (ms(6_000), 1)), + (ToolTimeout::Millis(1_001), 120, (ms(7_000), 2)), + (ToolTimeout::Millis(1_500), 120, (ms(7_000), 2)), + (ToolTimeout::Millis(30_000), 120, (ms(35_000), 30)), + (ToolTimeout::Millis(3_600_000), 120, (ms(3_605_000), 3600)), + (ToolTimeout::Millis(3_600_001), 120, (ms(3_605_000), 3600)), + (ToolTimeout::Millis(u64::MAX), 120, (ms(3_605_000), 3600)), + ]; + for (policy, inherited_secs, expected) in cases { + // Resolve against the pure helper so the process-global stays untouched. + let r = TimeoutSettings::with_inherited_secs(inherited_secs).resolve(policy); + let got = (r.deadline, r.budget_ms / 1000); + assert_eq!( + got, expected, + "policy {policy:?} inherited {inherited_secs}s" + ); + } +} + +/// Env / config precedence table for the effective inherited timeout. +#[test] +fn resolve_effective_table() { + let cases: &[(u64, Option<&str>, u64)] = &[ + (300, None, 300), + (300, Some("600"), 600), + (300, Some("1"), 1), + (300, Some("3600"), 3600), + (300, Some("3601"), 300), + (300, Some("0"), 300), + (300, Some("-1"), 300), + (300, Some(""), 300), + (300, Some("abc"), 300), + (0, None, DEFAULT_TIMEOUT_SECS), + (0, Some("45"), 45), + (3601, None, DEFAULT_TIMEOUT_SECS), + (u64::MAX, None, DEFAULT_TIMEOUT_SECS), + (1, None, 1), + (3600, None, 3600), + ]; + for &(config, env, expected) in cases { + assert_eq!( + resolve_effective(config, env), + expected, + "config {config} env {env:?}" + ); + } +} + +/// The deadline must take the command's whole process group with it. A shell +/// pipeline is grandchildren of the tool's child, and a `grep -rl … /` the tool +/// had reported as killed ran on for half an hour at a full core. +#[cfg(unix)] +#[tokio::test] +async fn a_timed_out_command_takes_its_whole_process_group_with_it() { + let dir = tempfile::tempdir() + .unwrap() + .keep() + .join(format!("oh-pgkill-{}", std::process::id())); + std::fs::create_dir_all(&dir).unwrap(); + let pidfile = dir.join("grandchild.pid"); + // `sleep 30 &` is a grandchild in the shell's own process group; `wait` + // keeps the shell alive past the deadline. + let mut cmd = sh(&format!("sleep 30 & echo $! > {}; wait", pidfile.display())); + let result = output_or_kill(&mut cmd, Duration::from_millis(700)).await; + assert!(result.is_err(), "the deadline must fire on a 30s sleep"); + + let grandchild: i32 = loop { + if let Ok(text) = std::fs::read_to_string(&pidfile) { + if let Ok(pid) = text.trim().parse() { + break pid; + } + } + tokio::time::sleep(Duration::from_millis(20)).await; + }; + let mut gone = false; + for _ in 0..100 { + // kill(pid, 0) succeeds while the process (or its zombie) exists. + if !process_exists(grandchild) { + gone = true; + break; + } + // A container's PID 1 may leave a killed orphan unreaped. It cannot + // execute side effects; use the same kernel-state rule as cleanup. + #[cfg(target_os = "linux")] + if let Ok(stat) = std::fs::read_to_string(format!("/proc/{grandchild}/stat")) { + let group = rustix::process::Pid::from_raw(grandchild) + .and_then(|pid| rustix::process::getpgid(Some(pid)).ok()); + if let Some(group) = group.map(|g| g.as_raw_nonzero().get() as u32) { + if !group_exit::is_active_member(&stat, group) { + gone = true; + break; + } + } + } + tokio::time::sleep(Duration::from_millis(20)).await; + } + let _ = std::fs::remove_dir_all(&dir); + assert!(gone, "grandchild {grandchild} survived the deadline"); +} + +/// Without a deadline the child still dies with the handle: a cancelled tool +/// future must not leave the command running either. +#[cfg(unix)] +#[tokio::test] +async fn a_shell_family_child_dies_with_a_dropped_future() { + let dir = tempfile::tempdir() + .unwrap() + .keep() + .join(format!("oh-dropkill-{}", std::process::id())); + std::fs::create_dir_all(&dir).unwrap(); + let pidfile = dir.join("child.pid"); + let mut cmd = sh(&format!("echo $$ > {}; sleep 30", pidfile.display())); + // `timeout` takes the future by value and drops it when the deadline + // passes -- exactly the abandoned-future path a cancelled tool call takes. + let _ = tokio::time::timeout(Duration::from_millis(300), output_unbounded(&mut cmd)).await; + let shell: i32 = loop { + if let Ok(text) = std::fs::read_to_string(&pidfile) { + if let Ok(pid) = text.trim().parse() { + break pid; + } + } + tokio::time::sleep(Duration::from_millis(20)).await; + }; + let mut gone = false; + for _ in 0..100 { + if !process_exists(shell) { + gone = true; + break; + } + tokio::time::sleep(Duration::from_millis(20)).await; + } + let _ = std::fs::remove_dir_all(&dir); + assert!(gone, "shell {shell} survived its dropped future"); +} + +/// A reaped scope: after the turn future is dropped, `ProcessCleanup::wait` +/// returns only once the sleeping child group is gone. +#[cfg(target_os = "linux")] +#[tokio::test] +async fn process_cleanup_waits_for_a_cancelled_sleeping_group() { + let cleanup = ProcessCleanup::default(); + let turn = { + let cleanup = cleanup.clone(); + tokio::spawn(async move { + cleanup + .scope(async { + let mut cmd = sh("sleep 30 & wait"); + let _ = output_unbounded(&mut cmd).await; + }) + .await; + }) + }; + tokio::time::sleep(Duration::from_millis(300)).await; + turn.abort(); + let _ = turn.await; + tokio::time::timeout(Duration::from_secs(10), cleanup.wait()) + .await + .expect("cleanup completes once the group is reaped"); +} + +#[tokio::test] +async fn command_environment_replaces_the_child_environment() { + let env = CommandEnvironment::new([ + ("PATH".to_string(), "/usr/bin:/bin".to_string()), + ("ONLY_THIS".to_string(), "1".to_string()), + ]); + let out = env + .scope(async { + let mut cmd = sh("echo ${ONLY_THIS:-missing}-${HOME:-nohome}"); + output_unbounded(&mut cmd).await.unwrap() + }) + .await; + assert_eq!(String::from_utf8_lossy(&out.stdout).trim(), "1-nohome"); +} + +#[tokio::test] +async fn command_environment_scopes_variable_lookup() { + assert!(!CommandEnvironment::is_active()); + let env = CommandEnvironment::new([("SCOPED_ONLY".to_string(), "yes".to_string())]); + env.scope(async { + assert!(CommandEnvironment::is_active()); + assert_eq!( + CommandEnvironment::var_os("SCOPED_ONLY").unwrap(), + std::ffi::OsString::from("yes") + ); + assert_eq!(CommandEnvironment::var("SCOPED_ONLY").unwrap(), "yes"); + assert!(CommandEnvironment::var_os("PATH").is_none()); + assert!(CommandEnvironment::var("PATH").is_err()); + }) + .await; + assert!(CommandEnvironment::var_os("PATH").is_some()); + assert!(CommandEnvironment::var("PATH").is_ok()); +} + +#[tokio::test] +async fn explicit_command_env_overrides_and_removals_survive_the_scoped_base() { + let env = CommandEnvironment::new([ + ("PATH".to_string(), "/usr/bin:/bin".to_string()), + ("BASE".to_string(), "base".to_string()), + ("DROP".to_string(), "x".to_string()), + ]); + let out = env + .scope(async { + let mut cmd = sh("echo $BASE-$EXTRA-${DROP:-dropped}"); + cmd.env("BASE", "override") + .env("EXTRA", "e") + .env_remove("DROP"); + output_unbounded(&mut cmd).await.unwrap() + }) + .await; + assert_eq!( + String::from_utf8_lossy(&out.stdout).trim(), + "override-e-dropped" + ); +} + +#[tokio::test] +async fn process_cleanup_reports_active_only_inside_its_scope() { + assert!(!ProcessCleanup::is_active()); + let cleanup = ProcessCleanup::default(); + cleanup + .scope(async { + assert!(ProcessCleanup::is_active()); + }) + .await; + cleanup.wait().await; +} + +/// A hold taken inside a turn's scope keeps `wait` pending until dropped, so +/// detached host work can finish its own cleanup before cancellation is +/// acknowledged. A hold taken outside any scope registers nothing. +#[tokio::test] +async fn a_cleanup_hold_keeps_wait_pending_until_dropped() { + let cleanup = ProcessCleanup::default(); + let hold = cleanup.scope(async { ProcessCleanup::hold() }).await; + assert!( + tokio::time::timeout(Duration::from_millis(50), cleanup.wait()) + .await + .is_err(), + "wait must not return while the hold is alive" + ); + drop(hold); + tokio::time::timeout(Duration::from_secs(1), cleanup.wait()) + .await + .expect("wait returns once the hold is dropped"); + + let unscoped = ProcessCleanup::default(); + let _outside = ProcessCleanup::hold(); + tokio::time::timeout(Duration::from_secs(1), unscoped.wait()) + .await + .expect("a hold outside any scope registers nothing"); +} diff --git a/crates/tinyagents-harness-tools/src/exec/timeout/output.rs b/crates/tinyagents-harness-tools/src/exec/timeout/output.rs new file mode 100644 index 000000000..6a215c49c --- /dev/null +++ b/crates/tinyagents-harness-tools/src/exec/timeout/output.rs @@ -0,0 +1,387 @@ +//! Child-process management: bounded capture, process-group kill, reaping. + +use std::future::Future; +use std::pin::Pin; +use std::sync::{Arc, Mutex, OnceLock}; +use std::time::Duration; + +use super::ProcessCleanup; +use super::command_environment::CommandEnvironment; +use super::group_exit; +use super::process_cleanup::Reaped; +use crate::exec::ExecProgress; + +type BoxTask = Pin> + Send>>; + +static TASK_WRAPPER: OnceLock BoxTask> = OnceLock::new(); + +/// Register a wrapper applied to every internal waiter task before it is +/// spawned, so a host can carry its ambient task-locals (request context, +/// identity) into it. First registration wins; returns whether it took effect. +pub fn set_task_wrapper(wrapper: fn(BoxTask) -> BoxTask) -> bool { + TASK_WRAPPER.set(wrapper).is_ok() +} + +fn spawn_wrapped( + fut: impl Future> + Send + 'static, +) -> tokio::task::JoinHandle> { + let fut: BoxTask = Box::pin(fut); + tokio::spawn(match TASK_WRAPPER.get() { + Some(wrap) => wrap(fut), + None => fut, + }) +} + +/// Run `cmd` to completion with its output captured, or kill it -- and every +/// process it started -- when `deadline` passes. +/// +/// `tokio::time::timeout(deadline, cmd.output())` only abandons the future. +/// Without `kill_on_drop` the child is not even signalled, and a shell's +/// pipeline (`grep -r … | head`) is a set of grandchildren that no kill of the +/// direct child reaches anyway. One `grep -rl … /` the shell tool had reported +/// as "timed out after 600s and was killed" ran on for half an hour at a full +/// core inside a two-CPU container, starving the agent that had started it. +/// +/// So the child is spawned as the leader of its own process group and the +/// whole group is signalled when the deadline fires. The result has the same +/// shape as `timeout(deadline, cmd.output())`, so a caller's match arms do not +/// change. Stdio is set the way `output()` sets it: stdin closed, both output +/// streams captured. +pub async fn output_or_kill( + cmd: &mut tokio::process::Command, + deadline: Duration, +) -> Result, tokio::time::error::Elapsed> { + output_or_kill_with_progress(cmd, deadline, None).await +} + +/// `output_or_kill` with streamed bytes sent to the supplied progress meter. +pub async fn output_or_kill_with_progress( + cmd: &mut tokio::process::Command, + deadline: Duration, + progress: Option, +) -> Result, tokio::time::error::Elapsed> { + output_or_kill_with_progress_mode(cmd, deadline, progress, true).await +} + +/// Like [`output_or_kill_with_progress`] without applying task-local ambient +/// variables to the child process. +pub async fn output_or_kill_with_progress_isolated( + cmd: &mut tokio::process::Command, + deadline: Duration, + progress: Option, +) -> Result, tokio::time::error::Elapsed> { + output_or_kill_with_progress_mode(cmd, deadline, progress, false).await +} + +/// Process result that retains bytes collected before a deadline killed it. +pub enum CapturedCommandResult { + Completed(std::io::Result), + /// `None` means the deadline fired before the waiter captured output. + TimedOut(Option), +} + +/// Run with a deadline and return the captured streams after cleanup on timeout. +pub async fn output_or_kill_captured_isolated( + cmd: &mut tokio::process::Command, + deadline: Duration, + progress: Option, +) -> CapturedCommandResult { + let cleanup = ProcessCleanup::default(); + let captured = Arc::new(Mutex::new(None)); + match cleanup + .scope(tokio::time::timeout( + deadline, + output_with_input_and_progress_mode( + cmd, + None, + progress, + false, + Some(Arc::clone(&captured)), + ), + )) + .await + { + Ok(result) => CapturedCommandResult::Completed(result), + Err(_) => { + cleanup.wait().await; + let output = captured + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .take(); + CapturedCommandResult::TimedOut(output) + } + } +} + +async fn output_or_kill_with_progress_mode( + cmd: &mut tokio::process::Command, + deadline: Duration, + progress: Option, + apply_task_environment: bool, +) -> Result, tokio::time::error::Elapsed> { + let cleanup = ProcessCleanup::default(); + match cleanup + .scope(tokio::time::timeout( + deadline, + output_unbounded_with_progress_mode(cmd, progress, apply_task_environment), + )) + .await + { + Ok(result) => Ok(result), + Err(elapsed) => { + cleanup.wait().await; + Err(elapsed) + } + } +} + +/// Capture a command with no deadline, killing its process group if the +/// future is dropped. The child is reaped by an owned waiter even when the +/// caller cancels; a scoped [`super::ProcessCleanup`] can await that waiter. +pub async fn output_unbounded( + cmd: &mut tokio::process::Command, +) -> std::io::Result { + output_with_input_and_progress(cmd, None, None).await +} + +/// Capture a command with no deadline while reporting output chunks. +pub async fn output_unbounded_with_progress( + cmd: &mut tokio::process::Command, + progress: Option, +) -> std::io::Result { + output_unbounded_with_progress_mode(cmd, progress, true).await +} + +/// Like [`output_unbounded_with_progress`] without applying task-local ambient +/// variables to the child process. +pub async fn output_unbounded_with_progress_isolated( + cmd: &mut tokio::process::Command, + progress: Option, +) -> std::io::Result { + output_unbounded_with_progress_mode(cmd, progress, false).await +} + +async fn output_unbounded_with_progress_mode( + cmd: &mut tokio::process::Command, + progress: Option, + apply_task_environment: bool, +) -> std::io::Result { + output_with_input_and_progress_mode(cmd, None, progress, apply_task_environment, None).await +} + +/// Capture a command, optionally supplying stdin, in the current cleanup scope. +pub async fn output_with_input( + cmd: &mut tokio::process::Command, + input: Option>, +) -> std::io::Result { + output_with_input_and_progress(cmd, input, None).await +} + +/// Capture a command, optionally supplying stdin, and stream output bytes to +/// the execution progress meter. +pub async fn output_with_input_and_progress( + cmd: &mut tokio::process::Command, + input: Option>, + progress: Option, +) -> std::io::Result { + output_with_input_and_progress_mode(cmd, input, progress, true, None).await +} + +async fn output_with_input_and_progress_mode( + cmd: &mut tokio::process::Command, + input: Option>, + progress: Option, + apply_task_environment: bool, + captured: Option>>>, +) -> std::io::Result { + use std::process::Stdio; + cmd.stdin(if input.is_some() { + Stdio::piped() + } else { + Stdio::null() + }) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .kill_on_drop(true); + if apply_task_environment { + CommandEnvironment::apply(cmd); + } + own_process_group(cmd.as_std_mut()); + let mut child = cmd.spawn()?; + let stdin = child.stdin.take(); + let reaped = Reaped::register(); + let (cancel, cancellation) = tokio::sync::watch::channel(false); + let waiter = spawn_wrapped(async move { + let _reaped = reaped; + let write_input = async move { + if let (Some(mut stdin), Some(input)) = (stdin, input) { + use tokio::io::AsyncWriteExt; + // An early child exit closes stdin. Its exit status/output is + // authoritative; a broken pipe must not hide it. + let _ = stdin.write_all(&input).await; + } + }; + let (_, output) = tokio::join!( + write_input, + collect_command_output(child, cancellation, progress) + ); + if let (Some(captured), Ok(output)) = (captured, output.as_ref()) { + *captured + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) = Some(std::process::Output { + status: output.status, + stdout: output.stdout.clone(), + stderr: output.stderr.clone(), + }); + } + output + }); + let _cancel_on_drop = CancelOnDrop(cancel); + waiter.await.map_err(std::io::Error::other)? +} + +// The caller signals only the owned waiter. It never retains a PID after +// that waiter reaps the child, so late future drops cannot kill a reused PID. +struct CancelOnDrop(tokio::sync::watch::Sender); + +impl Drop for CancelOnDrop { + fn drop(&mut self) { + self.0.send_replace(true); + } +} + +struct CommandGroup(Option); + +impl CommandGroup { + fn kill(&self) { + if let Some(pid) = self.0 { + kill_process_group(pid); + } + } +} + +impl Drop for CommandGroup { + fn drop(&mut self) { + self.kill(); + } +} + +async fn collect_command_output( + mut child: tokio::process::Child, + mut cancellation: tokio::sync::watch::Receiver, + progress: Option, +) -> std::io::Result { + let mut group = CommandGroup(child.id()); + let mut stdout = child + .stdout + .take() + .ok_or_else(|| std::io::Error::other("command stdout is not piped"))?; + let mut stderr = child + .stderr + .take() + .ok_or_else(|| std::io::Error::other("command stderr is not piped"))?; + let mut stdout_bytes = Vec::new(); + let mut stderr_bytes = Vec::new(); + { + let drain = async { + tokio::try_join!( + read_stream(&mut stdout, &mut stdout_bytes, progress.as_ref()), + read_stream(&mut stderr, &mut stderr_bytes, progress.as_ref()), + ) + }; + tokio::pin!(drain); + tokio::select! { + biased; + _ = async { let _ = cancellation.wait_for(|cancelled| *cancelled).await; } => { + // The leader has not been reaped, even if it already exited. + // Its PID cannot be reused while signalling this group. + group.kill(); + child.start_kill()?; + if let Ok(result) = tokio::time::timeout(Duration::from_secs(2), drain).await { + result?; + } + } + result = &mut drain => { result?; } + } + } + let status = tokio::select! { + biased; + _ = async { let _ = cancellation.wait_for(|cancelled| *cancelled).await; } => { + group.kill(); + child.start_kill()?; + // Pipe EOF and the shell's exit can precede a descendant's exit. + // Keep the leader unreaped while checking the reserved group. + let stopped = match group.0 { + Some(pid) => group_exit::wait(pid).await, + None => Ok(()), + }; + let status = child.wait().await?; + // Disarm before propagating a check error after reaping: its PID + // could now be reused, including while unwinding this function. + group.0 = None; + stopped?; + status + } + result = child.wait() => result?, + }; + // No await between reaping and disarming. Only this waiter owns the PID. + group.0 = None; + Ok(std::process::Output { + status, + stdout: stdout_bytes, + stderr: stderr_bytes, + }) +} + +async fn read_stream( + reader: &mut (impl tokio::io::AsyncRead + Unpin), + output: &mut Vec, + progress: Option<&ExecProgress>, +) -> std::io::Result<()> { + use tokio::io::AsyncReadExt; + let mut chunk = [0; 8_192]; + loop { + let read = reader.read(&mut chunk).await?; + if read == 0 { + return Ok(()); + } + output.extend_from_slice(&chunk[..read]); + if let Some(progress) = progress { + progress.record(&chunk[..read]); + } + } +} + +/// Make `cmd` the leader of a new process group when it is spawned, so that +/// [`kill_process_group`] can reach everything it starts. A no-op off Unix. +pub fn own_process_group(cmd: &mut std::process::Command) { + #[cfg(unix)] + { + use std::os::unix::process::CommandExt; + cmd.process_group(0); + } + #[cfg(not(unix))] + { + let _ = cmd; + } +} + +/// Send SIGKILL to the process group led by `pid` -- a child spawned via +/// [`own_process_group`] and everything it started. Off Unix the direct child +/// is what `kill_on_drop` reaches and no group exists to signal. +pub fn kill_process_group(pid: u32) { + #[cfg(unix)] + { + use rustix::process::{Pid, Signal}; + let Some(pid) = i32::try_from(pid).ok().and_then(Pid::from_raw) else { + return; + }; + // A signal to a process group this process created; the kernel + // validates the target. A failure (group already gone) is expected. + let _ = rustix::process::kill_process_group(pid, Signal::KILL); + } + #[cfg(not(unix))] + { + let _ = pid; + } +} diff --git a/crates/tinyagents-harness-tools/src/exec/timeout/process_cleanup.rs b/crates/tinyagents-harness-tools/src/exec/timeout/process_cleanup.rs new file mode 100644 index 000000000..845f11c93 --- /dev/null +++ b/crates/tinyagents-harness-tools/src/exec/timeout/process_cleanup.rs @@ -0,0 +1,94 @@ +//! Track command reaping for an embedder that awaits turn cancellation. + +use std::future::Future; +use std::sync::{Arc, Mutex}; + +use tokio::sync::watch; + +tokio::task_local! { + static ACTIVE: Vec; +} + +/// Command waiters belonging to one turn. Clone before scoping the turn; +/// after dropping the turn future, [`wait`](Self::wait) awaits their reaping. +/// Tasks spawned by a host tool must explicitly inherit this scope. +#[derive(Clone, Default, Debug)] +pub struct ProcessCleanup(Arc>>>); + +impl ProcessCleanup { + /// Whether the current task must acknowledge owned subprocess cleanup. + /// Interpreter pools have no per-job cancellation acknowledgement and + /// therefore cannot accept work from this scope. + #[must_use] + pub fn is_active() -> bool { + ACTIVE.try_with(|_| ()).is_ok() + } + + /// Run a future with command waiters registered to this turn. + pub async fn scope(&self, future: impl Future) -> T { + let mut scopes = ACTIVE.try_with(Clone::clone).unwrap_or_default(); + scopes.push(self.clone()); + ACTIVE.scope(scopes, future).await + } + + /// Register work that the current turn's [`wait`](Self::wait) must outlast + /// until the returned hold is dropped: the explicit inheritance a host tool + /// needs when it hands a command to a task or service of its own. Take it + /// inside the scope (before spawning), move it into the detached work, and + /// drop it once that work has stopped. Outside any scope it holds nothing. + pub fn hold() -> CleanupHold { + CleanupHold { + _reaped: Reaped::register(), + } + } + + /// Wait for every registered command to exit and its output pipes to close. + /// Call only after the scoped future has finished or been dropped, so no + /// further commands can be registered by it. + pub async fn wait(&self) { + let waiters = self + .0 + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .clone(); + for mut waiter in waiters { + let _ = waiter.wait_for(|done| *done).await; + } + } +} + +/// Keeps the turn's cleanup waiting while alive. See [`ProcessCleanup::hold`]. +#[must_use = "the cleanup stops waiting as soon as the hold is dropped"] +pub struct CleanupHold { + _reaped: Reaped, +} + +impl std::fmt::Debug for CleanupHold { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("CleanupHold").finish_non_exhaustive() + } +} + +pub(super) struct Reaped(watch::Sender); + +impl Reaped { + pub(super) fn register() -> Self { + let (done, waiter) = watch::channel(false); + let _ = ACTIVE.try_with(|scopes| { + for cleanup in scopes { + cleanup + .0 + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .push(waiter.clone()); + } + }); + Self(done) + } +} + +impl Drop for Reaped { + fn drop(&mut self) { + self.0.send_replace(true); + } +} diff --git a/crates/tinyagents-harness-tools/src/exec/timeout/settings.rs b/crates/tinyagents-harness-tools/src/exec/timeout/settings.rs new file mode 100644 index 000000000..6fd88abfe --- /dev/null +++ b/crates/tinyagents-harness-tools/src/exec/timeout/settings.rs @@ -0,0 +1,164 @@ +//! Wall-clock bounds for tool execution: pure resolution, no host policy. +//! +//! The host decides where the number comes from (config, an operator env var) +//! and pushes it in through [`configure`]; this module only bounds, resolves +//! and stores it. The inherited value lives in an atomic so a later push +//! reaches every reader on its next call. + +use std::sync::OnceLock; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::time::Duration; + +use tinytools::ToolTimeout; + +/// Default tool-execution timeout in seconds when nothing else is configured. +pub const DEFAULT_TIMEOUT_SECS: u64 = 120; +/// Smallest accepted timeout. `0` would disable the timeout entirely, so it is +/// rejected and falls back to the default. +pub const MIN_TIMEOUT_SECS: u64 = 1; +/// Largest accepted timeout (1 hour). +pub const MAX_TIMEOUT_SECS: u64 = 3600; +/// Effective-unbounded cap (24h) for sandbox backends, which require a finite +/// deadline. +pub const SANDBOX_UNBOUNDED_CAP_SECS: u64 = 86_400; +/// Slack added on top of an explicit per-call budget before the hard deadline +/// fires, so a tool finishing right at its budget is not killed by jitter. +pub const TOOL_TIMEOUT_GRACE_SECS: u64 = 5; + +/// A resolved per-call deadline. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct ResolvedTimeout { + /// Hard deadline, or `None` to run unbounded. + pub deadline: Option, + /// The user-facing budget in milliseconds (`0` when unbounded). + pub budget_ms: u64, +} + +/// Bounds plus the inherited timeout, with the inherited value atomic. +#[derive(Debug)] +pub struct TimeoutSettings { + inherited_ms: AtomicU64, + min_ms: u64, + max_ms: u64, + grace_ms: u64, +} + +impl TimeoutSettings { + /// Settings with explicit bounds, all in milliseconds. + #[must_use] + pub fn new(inherited_ms: u64, min_ms: u64, max_ms: u64, grace_ms: u64) -> Self { + Self { + inherited_ms: AtomicU64::new(inherited_ms), + min_ms, + max_ms, + grace_ms, + } + } + + /// Settings with the crate's default bounds, inheriting `inherited_secs`. + #[must_use] + pub fn with_inherited_secs(inherited_secs: u64) -> Self { + Self::new( + inherited_secs.saturating_mul(1000), + MIN_TIMEOUT_SECS * 1000, + MAX_TIMEOUT_SECS * 1000, + TOOL_TIMEOUT_GRACE_SECS * 1000, + ) + } + + /// The timeout `Inherit` tools get. + #[must_use] + pub fn inherited_timeout(&self) -> Duration { + Duration::from_millis(self.inherited_ms.load(Ordering::Relaxed)) + } + + /// Replace the inherited timeout; readers see it on their next call. + pub fn set_inherited_timeout_ms(&self, ms: u64) { + self.inherited_ms.store(ms, Ordering::Relaxed); + } + + /// Resolve a tool's policy. An explicit millisecond request is rounded up + /// to whole seconds, clamped to the bounds, and padded with the grace. + #[must_use] + pub fn resolve(&self, policy: ToolTimeout) -> ResolvedTimeout { + match policy { + ToolTimeout::Inherit => { + let ms = self.inherited_ms.load(Ordering::Relaxed); + ResolvedTimeout { + deadline: Some(Duration::from_millis(ms)), + budget_ms: ms, + } + } + ToolTimeout::Unbounded => ResolvedTimeout { + deadline: None, + budget_ms: 0, + }, + ToolTimeout::Millis(req) => { + let rounded = req.saturating_add(999) / 1000 * 1000; + let budget = rounded.clamp(self.min_ms, self.max_ms); + ResolvedTimeout { + deadline: Some(Duration::from_millis(budget.saturating_add(self.grace_ms))), + budget_ms: budget, + } + } + } + } +} + +static SETTINGS: OnceLock = OnceLock::new(); + +/// Install the process-global settings. The first call installs `settings`; +/// later calls only push its inherited timeout into the installed one, so +/// handles obtained earlier stay live. Returns the global. +pub fn configure(settings: TimeoutSettings) -> &'static TimeoutSettings { + let ms = settings.inherited_ms.load(Ordering::Relaxed); + let global = SETTINGS.get_or_init(|| settings); + global.set_inherited_timeout_ms(ms); + tracing::debug!("[tool_timeout] configured inherited timeout {ms}ms"); + global +} + +/// The process-global settings, seeded with [`DEFAULT_TIMEOUT_SECS`] when the +/// host never called [`configure`]. +#[must_use] +pub fn settings() -> &'static TimeoutSettings { + SETTINGS.get_or_init(|| TimeoutSettings::with_inherited_secs(DEFAULT_TIMEOUT_SECS)) +} + +/// Parse a raw value into a bounded timeout; anything missing, non-numeric or +/// outside `MIN_TIMEOUT_SECS..=MAX_TIMEOUT_SECS` yields [`DEFAULT_TIMEOUT_SECS`]. +#[must_use] +pub fn parse_tool_timeout_secs(raw: Option<&str>) -> u64 { + override_from(raw).unwrap_or(DEFAULT_TIMEOUT_SECS) +} + +/// A raw operator override, if it is a valid in-range number of seconds. +#[must_use] +pub fn override_from(raw: Option<&str>) -> Option { + raw.and_then(|s| s.parse::().ok()) + .filter(|&n| (MIN_TIMEOUT_SECS..=MAX_TIMEOUT_SECS).contains(&n)) +} + +/// Env override wins when valid; otherwise the (bounded) config value applies. +#[must_use] +pub fn resolve_effective(config_secs: u64, env_raw: Option<&str>) -> u64 { + override_from(env_raw) + .unwrap_or_else(|| parse_tool_timeout_secs(Some(&config_secs.to_string()))) +} + +/// Resolve an explicit per-call request for an otherwise unbounded tool. +/// `None`/`Some(0)` run unbounded; otherwise clamp to `MIN..=cap`. +#[must_use] +pub fn explicit_call_timeout_secs(requested: Option, cap: u64) -> Option { + let cap = cap.clamp(MIN_TIMEOUT_SECS, MAX_TIMEOUT_SECS); + match requested { + None | Some(0) => None, + Some(n) => Some(n.clamp(MIN_TIMEOUT_SECS, cap)), + } +} + +/// [`explicit_call_timeout_secs`] as a [`Duration`]. +#[must_use] +pub fn explicit_call_timeout_duration(requested: Option, cap: u64) -> Option { + explicit_call_timeout_secs(requested, cap).map(Duration::from_secs) +} diff --git a/crates/tinyagents-harness-tools/src/exec/tool.rs b/crates/tinyagents-harness-tools/src/exec/tool.rs new file mode 100644 index 000000000..7108179bd --- /dev/null +++ b/crates/tinyagents-harness-tools/src/exec/tool.rs @@ -0,0 +1,190 @@ +use std::{ffi::OsString, sync::Arc, time::Instant}; + +use async_trait::async_trait; +use serde_json::Value; +use tinytools::{PermissionLevel, Tool, ToolCallOptions, ToolResult, ToolRunContext, ToolTimeout}; + +use super::{ + ExecAuditEvent, ExecAuditOutcome, ExecAuthorizationRequest, ExecBackend, ExecBackendOutcome, + ExecDenied, ExecGate, ExecOperation, ExecRequest, ShellFlavor, + progress::{HEARTBEAT_INTERVAL, with_exec_progress}, + shell_description::shell_description, + shell_support::{exec_result, python_environment}, + timeout::{MAX_TIMEOUT_SECS, explicit_call_timeout_duration}, +}; + +/// Generic shell execution tool. Host policy and process selection are injected. +pub struct ExecTool { + gate: Arc, + backend: Arc, +} + +impl ExecTool { + /// Creates an exec tool with its host authorization gate and selected backend. + #[must_use] + pub fn new(gate: Arc, backend: Arc) -> Self { + Self { gate, backend } + } + + fn operation(&self, args: &Value) -> Result { + let Some(command) = args.get("command").and_then(Value::as_str) else { + return Err("Missing required command string.".to_owned()); + }; + Ok(ExecOperation::Shell { + command: command.to_owned(), + category: args + .get("category") + .and_then(Value::as_str) + .map(ToOwned::to_owned), + }) + } +} + +#[async_trait] +impl Tool for ExecTool { + fn name(&self) -> &str { + "exec" + } + + fn description(&self) -> &str { + shell_description(self.backend.shell_flavor()) + } + + fn parameters_schema(&self) -> Value { + match self.backend.shell_flavor() { + ShellFlavor::Posix => crate::static_schema!(include_str!("shell_schema_posix.json")), + ShellFlavor::Cmd => crate::static_schema!(include_str!("shell_schema_cmd.json")), + } + } + + fn max_result_size_chars(&self) -> Option { + Some(30_000) + } + + fn timeout_policy(&self, args: &Value) -> ToolTimeout { + match args.get("timeout_secs").and_then(Value::as_u64) { + None | Some(0) => ToolTimeout::Unbounded, + Some(secs) => ToolTimeout::Millis(secs.saturating_mul(1_000)), + } + } + + fn permission_level(&self) -> PermissionLevel { + PermissionLevel::Execute + } + + fn external_effect_with_args(&self, args: &Value) -> bool { + self.gate.external_effect_with_args(args) + } + + async fn prepare_authorization(&self, args: &Value) -> anyhow::Result> { + self.gate.prepare_authorization(args).await + } + + async fn execute(&self, args: Value) -> anyhow::Result { + self.execute_with_context(args, ToolCallOptions::default(), None) + .await + } + + async fn execute_with_context( + &self, + args: Value, + _options: ToolCallOptions, + context: Option<&dyn ToolRunContext>, + ) -> anyhow::Result { + let operation = match self.operation(&args) { + Ok(operation) => operation, + Err(message) => return Ok(ToolResult::error(message)), + }; + let timeout_secs = args.get("timeout_secs").and_then(Value::as_u64); + let timeout = explicit_call_timeout_duration(timeout_secs, MAX_TIMEOUT_SECS); + let authorization = ExecAuthorizationRequest { + operation: operation.clone(), + timeout, + }; + let approval_required = self.gate.external_effect_with_args(&args); + let started = Instant::now(); + let permit = match self.gate.authorize(&authorization, context).await { + Ok(permit) => permit, + Err(ExecDenied(reason)) => { + self.gate.record(ExecAuditEvent { + operation, + approval_required, + authorized: false, + allowed: false, + backend_selection: None, + process_attempted: false, + outcome: ExecAuditOutcome::GateDenied, + duration: started.elapsed(), + }); + return Ok(ToolResult::error(reason)); + } + }; + + let flavor = self.backend.shell_flavor(); + let command = match &authorization.operation { + ExecOperation::Shell { command, .. } => command.clone(), + ExecOperation::Install { .. } => unreachable!("ExecTool only handles shell operations"), + }; + let (program, args): (OsString, Vec) = match flavor { + ShellFlavor::Posix => ("sh".into(), vec!["-c".into(), command.clone().into()]), + ShellFlavor::Cmd => ("cmd.exe".into(), vec!["/C".into(), command.clone().into()]), + }; + let backend_selection = permit.backend_selection.clone(); + let request = ExecRequest { + program, + args, + action_dir: permit.action_dir, + environment: python_environment(permit.environment), + timeout, + shell_command: Some(command), + backend_selection: permit.backend_selection, + backend_capability: permit.backend_capability, + }; + let backend = Arc::clone(&self.backend); + let outcome = with_exec_progress(context, HEARTBEAT_INTERVAL, |progress| async move { + backend.execute(request, context, progress).await + }) + .await; + + let (allowed, process_attempted, audit_outcome, result) = match outcome { + ExecBackendOutcome::Refused { kind, reason } => ( + false, + false, + ExecAuditOutcome::BackendRefused { kind }, + ToolResult::error(reason), + ), + ExecBackendOutcome::LaunchFailed { error } => ( + true, + true, + ExecAuditOutcome::LaunchFailed, + ToolResult::error(format!("Sandbox execution failed: {error}")), + ), + ExecBackendOutcome::Completed(output) => { + let timed_out = output.timed_out; + let result = exec_result(output, timeout); + let success = !result.is_error; + ( + true, + true, + ExecAuditOutcome::Completed { success, timed_out }, + result, + ) + } + }; + self.gate.record(ExecAuditEvent { + operation, + approval_required, + authorized: true, + allowed, + backend_selection: Some(backend_selection), + process_attempted, + outcome: audit_outcome, + duration: started.elapsed(), + }); + Ok(result) + } +} + +#[cfg(test)] +#[path = "tool_tests.rs"] +mod tests; diff --git a/crates/tinyagents-harness-tools/src/exec/tool_tests.rs b/crates/tinyagents-harness-tools/src/exec/tool_tests.rs new file mode 100644 index 000000000..84e495e16 --- /dev/null +++ b/crates/tinyagents-harness-tools/src/exec/tool_tests.rs @@ -0,0 +1,380 @@ +use super::*; +use crate::exec::{ + ExecAuditEvent, ExecAuditOutcome, ExecAuthorizationRequest, ExecBackend, ExecBackendOutcome, + ExecDenied, ExecGate, ExecOutput, ExecPermit, ExecProgress, ExecRefusalKind, ExecRequest, + ShellFlavor, +}; +use serde_json::json; +use std::{ + collections::BTreeMap, + ffi::OsString, + path::PathBuf, + sync::{ + Mutex, + atomic::{AtomicUsize, Ordering}, + }, + time::Duration, +}; + +#[derive(Default)] +struct GateState { + denied: bool, + events: Mutex>, + requests: Mutex>, +} + +struct TestGate(Arc); + +#[async_trait] +impl ExecGate for TestGate { + async fn prepare_authorization(&self, args: &Value) -> anyhow::Result> { + Ok(Some(json!({"prepared": args.get("command")}))) + } + + fn external_effect_with_args(&self, args: &Value) -> bool { + args.get("command") + .and_then(Value::as_str) + .is_some_and(|command| !command.is_empty()) + } + + async fn authorize( + &self, + request: &ExecAuthorizationRequest, + _context: Option<&dyn tinytools::ToolRunContext>, + ) -> Result { + self.0.requests.lock().unwrap().push(request.clone()); + if self.0.denied { + return Err(ExecDenied("blocked by test policy".to_owned())); + } + Ok(ExecPermit { + action_dir: PathBuf::from("/authorized/action"), + environment: BTreeMap::from([ + (OsString::from("KEEP"), OsString::from("provided")), + (OsString::from("PYTHONUTF8"), OsString::from("0")), + ]), + backend_selection: "native-test".to_owned(), + backend_capability: crate::exec::ExecBackendCapability::new("opaque-capability"), + }) + } + + fn record(&self, event: ExecAuditEvent) { + self.0.events.lock().unwrap().push(event); + } +} + +struct TestBackend { + flavor: ShellFlavor, + outcome: Mutex, + calls: AtomicUsize, + requests: Mutex>, +} + +impl TestBackend { + fn new(flavor: ShellFlavor, outcome: ExecBackendOutcome) -> Self { + Self { + flavor, + outcome: Mutex::new(outcome), + calls: AtomicUsize::new(0), + requests: Mutex::new(Vec::new()), + } + } +} + +#[async_trait] +impl ExecBackend for TestBackend { + fn shell_flavor(&self) -> ShellFlavor { + self.flavor + } + + async fn execute( + &self, + request: ExecRequest, + _context: Option<&dyn tinytools::ToolRunContext>, + _progress: Option, + ) -> ExecBackendOutcome { + self.calls.fetch_add(1, Ordering::SeqCst); + self.requests.lock().unwrap().push(request); + self.outcome.lock().unwrap().clone() + } +} + +fn harness( + denied: bool, + flavor: ShellFlavor, + outcome: ExecBackendOutcome, +) -> (ExecTool, Arc, Arc) { + let state = Arc::new(GateState { + denied, + ..GateState::default() + }); + let backend = Arc::new(TestBackend::new(flavor, outcome)); + let tool = ExecTool::new(Arc::new(TestGate(Arc::clone(&state))), backend.clone()); + (tool, state, backend) +} + +fn completed( + stdout: &[u8], + stderr: &[u8], + code: Option, + timed_out: bool, +) -> ExecBackendOutcome { + ExecBackendOutcome::Completed(ExecOutput { + stdout: stdout.to_vec(), + stderr: stderr.to_vec(), + exit_code: code, + timed_out, + }) +} + +#[tokio::test] +async fn a_denied_gate_never_calls_the_backend() { + let (tool, state, backend) = harness( + true, + ShellFlavor::Posix, + completed(b"unexpected", b"", Some(0), false), + ); + let result = tool + .execute(json!({"command":"echo should-not-run"})) + .await + .unwrap(); + assert!(result.is_error); + assert_eq!(result.text(), "blocked by test policy"); + assert_eq!(backend.calls.load(Ordering::SeqCst), 0); + assert!(backend.requests.lock().unwrap().is_empty()); + let events = state.events.lock().unwrap(); + assert_eq!(events.len(), 1); + assert_eq!(events[0].outcome, ExecAuditOutcome::GateDenied); + assert!(!events[0].authorized); + assert!(!events[0].allowed); + assert!(!events[0].process_attempted); +} + +#[tokio::test] +async fn authorized_shell_uses_the_gate_permit_and_records_success() { + let (tool, state, backend) = harness( + false, + ShellFlavor::Posix, + completed(b"done", b"", Some(0), false), + ); + let result = tool + .execute(json!({"command":"echo done", "category":"write", "timeout_secs":12})) + .await + .unwrap(); + assert!(!result.is_error); + assert_eq!(result.text(), "done"); + assert_eq!(backend.calls.load(Ordering::SeqCst), 1); + let requests = backend.requests.lock().unwrap(); + assert_eq!(requests.len(), 1); + assert_eq!(requests[0].program, OsString::from("sh")); + assert_eq!( + requests[0].args, + vec![OsString::from("-c"), OsString::from("echo done")] + ); + assert_eq!(requests[0].action_dir, PathBuf::from("/authorized/action")); + assert_eq!(requests[0].timeout, Some(Duration::from_secs(12))); + assert_eq!(requests[0].shell_command.as_deref(), Some("echo done")); + assert_eq!(requests[0].backend_selection, "native-test"); + assert_eq!( + requests[0].environment.get(std::ffi::OsStr::new("KEEP")), + Some(&OsString::from("provided")) + ); + assert_eq!( + requests[0] + .environment + .get(std::ffi::OsStr::new("PYTHONUTF8")), + Some(&OsString::from("0")) + ); + assert_eq!( + requests[0] + .environment + .get(std::ffi::OsStr::new("PYTHONIOENCODING")), + Some(&OsString::from("utf-8")) + ); + let events = state.events.lock().unwrap(); + assert_eq!(events.len(), 1); + assert!(events[0].approval_required); + assert!(events[0].authorized); + assert!(events[0].allowed); + assert!(events[0].process_attempted); + assert_eq!(events[0].backend_selection.as_deref(), Some("native-test")); + assert_eq!( + events[0].outcome, + ExecAuditOutcome::Completed { + success: true, + timed_out: false + } + ); +} + +#[tokio::test] +async fn cmd_flavor_uses_cmd_argv_and_matching_description() { + let (tool, _, backend) = harness( + false, + ShellFlavor::Cmd, + completed(b"ok", b"", Some(0), false), + ); + assert!(tool.description().contains("Windows cmd.exe")); + assert_eq!( + tool.parameters_schema()["properties"]["command"]["description"], + "The cmd.exe command to execute (cmd syntax, not POSIX)" + ); + tool.execute(json!({"command":"echo ok"})).await.unwrap(); + let requests = backend.requests.lock().unwrap(); + assert_eq!(requests[0].program, OsString::from("cmd.exe")); + assert_eq!( + requests[0].args, + vec![OsString::from("/C"), OsString::from("echo ok")] + ); +} + +#[tokio::test] +async fn successful_stderr_is_returned_and_audited_as_success() { + let (tool, state, _) = harness( + false, + ShellFlavor::Posix, + completed(b"out", b"warning", Some(0), false), + ); + let result = tool.execute(json!({"command":"run"})).await.unwrap(); + assert!(!result.is_error); + assert_eq!(result.text(), "out\n[stderr]\nwarning"); + assert_eq!( + state.events.lock().unwrap()[0].outcome, + ExecAuditOutcome::Completed { + success: true, + timed_out: false + } + ); +} + +#[tokio::test] +async fn command_failure_and_timeout_are_audited_as_failed_completions() { + for (output, expected) in [ + ( + completed(b"out", b"err", Some(7), false), + ExecAuditOutcome::Completed { + success: false, + timed_out: false, + }, + ), + ( + completed(b"partial", b"", None, true), + ExecAuditOutcome::Completed { + success: false, + timed_out: true, + }, + ), + ] { + let (tool, state, _) = harness(false, ShellFlavor::Posix, output); + let result = tool + .execute(json!({"command":"run", "timeout_secs":5})) + .await + .unwrap(); + assert!(result.is_error); + assert_eq!(state.events.lock().unwrap()[0].outcome, expected); + } +} + +#[tokio::test] +async fn backend_refusal_is_not_reported_as_allowed_execution() { + let (tool, state, backend) = harness( + false, + ShellFlavor::Posix, + ExecBackendOutcome::Refused { + kind: ExecRefusalKind::Preflight, + reason: "sandbox unavailable".into(), + }, + ); + let result = tool.execute(json!({"command":"run"})).await.unwrap(); + assert!(result.is_error); + assert_eq!(backend.calls.load(Ordering::SeqCst), 1); + let event = &state.events.lock().unwrap()[0]; + assert!(event.authorized); + assert!(!event.allowed); + assert!(!event.process_attempted); + assert_eq!( + event.outcome, + ExecAuditOutcome::BackendRefused { + kind: ExecRefusalKind::Preflight + } + ); +} + +#[tokio::test] +async fn launch_failure_is_audited_after_process_attempt() { + let (tool, state, _) = harness( + false, + ShellFlavor::Posix, + ExecBackendOutcome::LaunchFailed { + error: "spawn denied".into(), + }, + ); + let result = tool.execute(json!({"command":"run"})).await.unwrap(); + assert!(result.is_error); + assert!( + result + .text() + .contains("Sandbox execution failed: spawn denied") + ); + let event = &state.events.lock().unwrap()[0]; + assert!(event.authorized); + assert!(event.allowed); + assert!(event.process_attempted); + assert_eq!(event.outcome, ExecAuditOutcome::LaunchFailed); +} + +#[tokio::test] +async fn authorization_preparation_and_external_effect_classification_delegate_to_gate() { + let (tool, _, _) = harness( + false, + ShellFlavor::Posix, + completed(b"", b"", Some(0), false), + ); + let args = json!({"command":"echo"}); + assert_eq!( + tool.prepare_authorization(&args).await.unwrap(), + Some(json!({"prepared":"echo"})) + ); + assert!(tool.external_effect_with_args(&args)); + assert!(!tool.external_effect_with_args(&json!({"command":""}))); +} + +#[tokio::test] +async fn missing_command_fails_without_authorization_or_audit() { + let (tool, state, backend) = harness( + false, + ShellFlavor::Posix, + completed(b"", b"", Some(0), false), + ); + assert!(tool.execute(json!({})).await.unwrap().is_error); + assert!(state.requests.lock().unwrap().is_empty()); + assert!(state.events.lock().unwrap().is_empty()); + assert_eq!(backend.calls.load(Ordering::SeqCst), 0); +} + +#[test] +fn tool_metadata_and_timeout_policy_are_stable() { + let (tool, _, _) = harness( + false, + ShellFlavor::Posix, + completed(b"", b"", Some(0), false), + ); + assert_eq!(tool.name(), "exec"); + assert_eq!(tool.permission_level(), PermissionLevel::Execute); + assert_eq!(tool.max_result_size_chars(), Some(30_000)); + assert_eq!(tool.timeout_policy(&json!({})), ToolTimeout::Unbounded); + assert_eq!( + tool.timeout_policy(&json!({"timeout_secs":0})), + ToolTimeout::Unbounded + ); + assert_eq!( + tool.timeout_policy(&json!({"timeout_secs":12})), + ToolTimeout::Millis(12_000) + ); + let schema = tool.parameters_schema(); + assert_eq!(schema["required"], json!(["command"])); + assert_eq!(schema["properties"]["timeout_secs"]["maximum"], 3600); + assert_eq!( + schema["properties"]["category"]["enum"], + json!(["read", "write", "network", "install", "destructive"]) + ); +} diff --git a/crates/tinyagents-harness-tools/src/fs/file_state/agent_context.rs b/crates/tinyagents-harness-tools/src/fs/file_state/agent_context.rs new file mode 100644 index 000000000..1624adfbe --- /dev/null +++ b/crates/tinyagents-harness-tools/src/fs/file_state/agent_context.rs @@ -0,0 +1,24 @@ +//! Task-local carrier for the currently-executing agent's identity so +//! file tools can attribute reads/writes without widening the `Tool` trait. +//! +//! Set by the agent harness around tool execution; tools read via [`current_file_state_agent_id`]. + +tokio::task_local! { + static FILE_STATE_AGENT_ID: String; +} + +/// Returns the current agent's identity for file-state tracking, if set. +/// +/// Returns `None` outside an agent turn (CLI, JSON-RPC direct, unit tests). +#[must_use] +pub fn current_file_state_agent_id() -> Option { + FILE_STATE_AGENT_ID.try_with(std::clone::Clone::clone).ok() +} + +/// Run `future` with `agent_id` installed as the file-state identity. +pub async fn with_file_state_agent_id(agent_id: String, future: F) -> R +where + F: std::future::Future, +{ + FILE_STATE_AGENT_ID.scope(agent_id, future).await +} diff --git a/crates/tinyagents-harness-tools/src/fs/file_state/file_state_tests.rs b/crates/tinyagents-harness-tools/src/fs/file_state/file_state_tests.rs new file mode 100644 index 000000000..1cebb2bdb --- /dev/null +++ b/crates/tinyagents-harness-tools/src/fs/file_state/file_state_tests.rs @@ -0,0 +1,353 @@ +//! Tests for file-state scoping, tracking, staleness checks, and path locks. + +use super::*; + +#[tokio::test] +async fn returns_none_outside_scope() { + assert_eq!(current_file_state_agent_id(), None); +} + +#[tokio::test] +async fn installs_and_reads_agent_id() { + let observed = + with_file_state_agent_id("agent-1".into(), async { current_file_state_agent_id() }).await; + assert_eq!(observed, Some("agent-1".to_string())); +} + +#[tokio::test] +async fn does_not_leak_across_scopes() { + with_file_state_agent_id("agent-1".into(), async { + assert_eq!(current_file_state_agent_id(), Some("agent-1".to_string())); + }) + .await; + assert_eq!(current_file_state_agent_id(), None); +} + +#[tokio::test] +async fn nested_scope_overrides_outer() { + with_file_state_agent_id("parent".into(), async { + assert_eq!(current_file_state_agent_id(), Some("parent".to_string())); + with_file_state_agent_id("child".into(), async { + assert_eq!(current_file_state_agent_id(), Some("child".to_string())); + }) + .await; + assert_eq!(current_file_state_agent_id(), Some("parent".to_string())); + }) + .await; +} +// Tests for read/write tracking, staleness checks, and path locks. +use std::path::{Path, PathBuf}; +use std::sync::Arc; +use std::time::{Duration, Instant, SystemTime}; +use tokio::sync::Mutex; + +fn fresh_coordinator() -> Arc { + Arc::new(FileStateCoordinator::new()) +} + +#[test] +fn record_and_check_no_staleness() { + let coord = fresh_coordinator(); + let path = PathBuf::from("/tmp/test/a.txt"); + coord.reads.write().insert( + ("agent-a".to_string(), path.clone()), + ReadStamp { + mtime: SystemTime::now(), + timestamp: Instant::now(), + partial: false, + }, + ); + let partial = coord + .reads + .read() + .get(&("agent-a".to_string(), path.clone())) + .map(|rs| rs.partial); + assert_eq!(partial, Some(false)); + assert!(coord.writes.read().get(&path).is_none()); +} + +#[test] +fn detect_sibling_write_staleness() { + let coord = fresh_coordinator(); + let path = PathBuf::from("/tmp/test/b.txt"); + coord.record_read( + "agent-a", + path.clone(), + SystemTime::now(), + false, + Instant::now(), + ); + std::thread::sleep(Duration::from_millis(5)); + coord.record_write("agent-b", path.clone()); + let stale = coord.stale_reads_for_parent("agent-a"); + assert_eq!(stale, vec![path]); +} + +#[test] +fn own_write_does_not_trigger_staleness() { + let coord = fresh_coordinator(); + let path = PathBuf::from("/tmp/test/c.txt"); + coord.record_read( + "agent-a", + path.clone(), + SystemTime::now(), + false, + Instant::now(), + ); + std::thread::sleep(Duration::from_millis(5)); + coord.record_write("agent-a", path.clone()); + let stale = coord.stale_reads_for_parent("agent-a"); + assert_eq!(stale.len(), 0); + assert_eq!(coord.check_stale_read("agent-a", &path), None); +} + +#[test] +fn partial_read_detected() { + let coord = fresh_coordinator(); + let path = PathBuf::from("/tmp/test/d.txt"); + coord.reads.write().insert( + ("agent-a".to_string(), path.clone()), + ReadStamp { + mtime: SystemTime::now(), + timestamp: Instant::now(), + partial: true, + }, + ); + let partial = coord + .reads + .read() + .get(&("agent-a".to_string(), path.clone())) + .map(|rs| rs.partial); + assert_eq!(partial, Some(true)); +} + +#[test] +fn parent_stale_files_detects_child_writes() { + let coord = fresh_coordinator(); + let path = PathBuf::from("/tmp/test/e.txt"); + coord.record_read( + "parent", + path.clone(), + SystemTime::now(), + false, + Instant::now(), + ); + std::thread::sleep(Duration::from_millis(5)); + coord.record_write("child-1", path.clone()); + assert_eq!(coord.stale_reads_for_parent("parent"), vec![path.clone()]); + assert_eq!( + coord.parent_stale_files("parent", &["child-1".to_string()]), + vec![path] + ); + assert_eq!( + coord + .parent_stale_files("parent", &["someone-else".to_string()]) + .len(), + 0 + ); +} + +#[test] +fn paths_written_by_collects_correctly() { + let coord = fresh_coordinator(); + let p1 = PathBuf::from("/tmp/test/f1.txt"); + let p2 = PathBuf::from("/tmp/test/f2.txt"); + coord.record_write("child-1", p1.clone()); + coord.record_write("child-2", p2); + let result = coord.paths_written_by(&["child-1".to_string()]); + assert_eq!(result.len(), 1); + assert_eq!(result.get("child-1"), Some(&vec![p1])); +} + +#[tokio::test] +async fn path_lock_serialises_access() { + let coord = fresh_coordinator(); + let path = PathBuf::from("/tmp/test/lock.txt"); + let mutex = { + let mut locks = coord.path_locks.write(); + locks + .entry(path.clone()) + .or_insert_with(|| Arc::new(Mutex::new(()))) + .clone() + }; + + let guard = mutex.lock().await; + assert!(mutex.try_lock().is_err()); + drop(guard); + assert!(mutex.try_lock().is_ok()); +} + +#[tokio::test] +async fn global_api_tracks_reads_writes_and_locks() -> anyhow::Result<()> { + use super::{ + acquire_path_lock, check_partial_read, check_stale_read, init_global, parent_stale_files, + record_read, record_write, try_global, + }; + + init_global(false); + init_global(true); + assert!(try_global().is_some()); + + let path = PathBuf::from("/tmp/test/global-flow.txt"); + record_read( + "reader", + path.clone(), + SystemTime::now(), + true, + Instant::now(), + ); + assert!(check_partial_read("reader", &path).is_some()); + assert!(check_stale_read("reader", &path).is_none()); + + record_read( + "reader", + path.clone(), + SystemTime::now(), + false, + Instant::now(), + ); + assert!(check_partial_read("reader", &path).is_none()); + + std::thread::sleep(Duration::from_millis(5)); + record_write("writer", path.clone()); + let msg = check_stale_read("reader", &path) + .ok_or_else(|| anyhow::anyhow!("expected a stale read after the sibling write"))?; + assert!(msg.contains("writer")); + assert_eq!( + parent_stale_files("reader", &["writer".to_string()]), + vec![path.clone()] + ); + assert!(check_stale_read("writer", &path).is_none()); + + let guard = acquire_path_lock(&path).await; + assert!(guard.is_some()); + Ok(()) +} + +#[test] +fn paths_written_by_keeps_a_path_after_another_agent_overwrites_it() { + let coord = fresh_coordinator(); + let path = PathBuf::from("/tmp/test/history-shared.txt"); + coord.record_write("child-1", path.clone()); + coord.record_write("child-2", path.clone()); + + let result = coord.paths_written_by(&["child-1".to_string(), "child-2".to_string()]); + assert_eq!(result.get("child-1"), Some(&vec![path.clone()])); + assert_eq!(result.get("child-2"), Some(&vec![path])); +} + +#[test] +fn sibling_write_during_an_in_flight_read_is_reported_stale() { + let coord = fresh_coordinator(); + let path = PathBuf::from("/tmp/test/in-flight.txt"); + let read_started = Instant::now(); + std::thread::sleep(Duration::from_millis(2)); + // The sibling write lands after the reader opened the file but before + // the reader got round to recording the read. + coord.record_write("sibling", path.clone()); + std::thread::sleep(Duration::from_millis(2)); + coord.record_read( + "reader", + path.clone(), + SystemTime::now(), + false, + read_started, + ); + + assert_eq!(coord.stale_reads_for_parent("reader"), vec![path]); +} + +// ── A later writer must not mask an earlier one ───────────── + +/// Parent reads `path`, then child-1 and child-2 write it in that order. +fn parent_read_then_two_child_writes(coord: &FileStateCoordinator, path: &Path) { + coord.record_read( + "parent", + path.to_path_buf(), + SystemTime::now(), + false, + Instant::now(), + ); + std::thread::sleep(Duration::from_millis(2)); + coord.record_write("child-1", path.to_path_buf()); + std::thread::sleep(Duration::from_millis(2)); + coord.record_write("child-2", path.to_path_buf()); +} + +#[test] +fn parent_stale_files_reports_an_earlier_child_write_masked_by_a_later_one() { + let coord = fresh_coordinator(); + let path = PathBuf::from("/tmp/test/masked-child.txt"); + parent_read_then_two_child_writes(&coord, &path); + + assert_eq!( + coord.parent_stale_files("parent", &["child-1".to_string()]), + vec![path.clone()] + ); + assert_eq!( + coord.parent_stale_files("parent", &["child-2".to_string()]), + vec![path] + ); +} + +#[test] +fn stale_reads_for_parent_reports_a_path_written_by_two_children() { + let coord = fresh_coordinator(); + let path = PathBuf::from("/tmp/test/two-children.txt"); + parent_read_then_two_child_writes(&coord, &path); + + assert_eq!(coord.stale_reads_for_parent("parent"), vec![path]); +} + +#[test] +fn own_later_write_does_not_mask_a_sibling_write_during_an_in_flight_read() { + // The reader opens the file, a sibling writes it, the reader then writes + // it through another tool, and only afterwards records the read that + // started before the sibling's write. The latest writer is the reader + // itself, but the content it holds may predate the sibling's change. + let coord = fresh_coordinator(); + let path = PathBuf::from("/tmp/test/own-write-masks.txt"); + let read_started = Instant::now(); + std::thread::sleep(Duration::from_millis(2)); + coord.record_write("sibling", path.clone()); + std::thread::sleep(Duration::from_millis(2)); + coord.record_write("reader", path.clone()); + coord.record_read( + "reader", + path.clone(), + SystemTime::now(), + false, + read_started, + ); + + assert_eq!(coord.stale_reads_for_parent("reader"), vec![path.clone()]); + let msg = coord.check_stale_read("reader", &path); + assert!( + msg.as_deref().is_some_and(|m| m.contains("'sibling'")), + "got: {msg:?}" + ); +} + +#[test] +fn a_write_before_the_read_is_not_stale() { + let coord = fresh_coordinator(); + let path = PathBuf::from("/tmp/test/write-then-read.txt"); + coord.record_write("child-1", path.clone()); + std::thread::sleep(Duration::from_millis(2)); + coord.record_read( + "parent", + path.clone(), + SystemTime::now(), + false, + Instant::now(), + ); + + assert_eq!(coord.stale_reads_for_parent("parent").len(), 0); + assert_eq!( + coord + .parent_stale_files("parent", &["child-1".to_string()]) + .len(), + 0 + ); + assert_eq!(coord.check_stale_read("parent", &path), None); +} diff --git a/crates/tinyagents-harness-tools/src/fs/file_state/mod.rs b/crates/tinyagents-harness-tools/src/fs/file_state/mod.rs new file mode 100644 index 000000000..0035da878 --- /dev/null +++ b/crates/tinyagents-harness-tools/src/fs/file_state/mod.rs @@ -0,0 +1,32 @@ +//! Process-wide file state coordinator for cross-agent staleness detection. +//! +//! Parallel subagents and worker threads share a workspace. Without +//! coordination one worker can read a file, a sibling can edit it, and +//! the first worker can later write based on stale content. This module +//! tracks per-agent read stamps and, for every path, each agent's latest +//! write stamp so that write tools can detect the conflict and return a +//! model-facing error requiring the agent to re-read. A read counts as stale +//! when *any* other agent wrote after it, not only the most recent writer. +//! +//! A read tool captures `Instant::now()` *before* it opens the file and hands +//! that stamp to [`record_read`] afterwards, so a sibling write racing the +//! read's I/O is ordered after the read and still reported stale. +//! +//! The guard is opt-in for the process: the host calls [`init_global`] with +//! `enabled` (typically derived from its own configuration or environment). +//! Until it does, or when it passes `false`, every operation is a no-op. + +mod agent_context; +mod ops; +mod types; + +pub use agent_context::{current_file_state_agent_id, with_file_state_agent_id}; +pub use ops::{ + acquire_path_lock, check_partial_read, check_stale_read, init_global, parent_stale_files, + record_read, record_write, try_global, +}; +pub use types::{FileStateCoordinator, ReadStamp}; + +#[cfg(test)] +#[path = "file_state_tests.rs"] +mod tests; diff --git a/crates/tinyagents-harness-tools/src/fs/file_state/ops.rs b/crates/tinyagents-harness-tools/src/fs/file_state/ops.rs new file mode 100644 index 000000000..efa52b368 --- /dev/null +++ b/crates/tinyagents-harness-tools/src/fs/file_state/ops.rs @@ -0,0 +1,212 @@ +//! Operational API for the file state coordinator. + +use std::path::{Path, PathBuf}; +use std::sync::{Arc, OnceLock}; +use std::time::{Instant, SystemTime}; +use tokio::sync::{Mutex, OwnedMutexGuard}; + +use super::types::{FileStateCoordinator, ReadStamp, writers_after_read}; + +// ── Singleton ──────────────────────────────────────────────────────────── + +static GLOBAL: OnceLock> = OnceLock::new(); + +/// Initialise the process-global coordinator when `enabled`. Safe to call +/// multiple times; only the first enabling call wins. Passing `false` leaves the +/// guard off, so [`try_global`] keeps returning `None` and every tracking +/// operation stays a no-op. +pub fn init_global(enabled: bool) { + if !enabled { + tracing::debug!("[file_state] guard disabled by host"); + return; + } + let _ = GLOBAL.set(Arc::new(FileStateCoordinator::new())); + tracing::debug!("[file_state] coordinator initialised"); +} + +/// Returns the global coordinator, or `None` when disabled / not yet initialised. +pub fn try_global() -> Option> { + GLOBAL.get().cloned() +} + +// ── Read tracking ──────────────────────────────────────────────────────── + +/// Record that `agent_id` read `resolved_path`. +/// +/// `read_started` must be captured with [`Instant::now`] *before* the file +/// is opened, not after its contents were read. A sibling write that lands +/// while the read is in flight is then ordered after the read and reported +/// stale by [`check_stale_read`], instead of being silently absorbed. +/// +/// ``` +/// use std::path::PathBuf; +/// use std::time::{Instant, SystemTime}; +/// use tinyagents_harness_tools::fs::file_state::record_read; +/// +/// let path = PathBuf::from("/workspace/notes.txt"); +/// let read_started = Instant::now(); +/// // ... open and read the file, then stat it for its mtime ... +/// record_read("agent-1", path, SystemTime::now(), false, read_started); +/// ``` +pub fn record_read( + agent_id: &str, + resolved_path: PathBuf, + mtime: SystemTime, + partial: bool, + read_started: Instant, +) { + let Some(coord) = try_global() else { return }; + coord.record_read(agent_id, resolved_path, mtime, partial, read_started); +} + +impl FileStateCoordinator { + /// Record a read on this coordinator; [`record_read`] delegates here. + pub(crate) fn record_read( + &self, + agent_id: &str, + resolved_path: PathBuf, + mtime: SystemTime, + partial: bool, + read_started: Instant, + ) { + tracing::trace!( + agent = agent_id, + path = %resolved_path.display(), + partial, + "[file_state] record_read" + ); + self.reads.write().insert( + (agent_id.to_string(), resolved_path), + ReadStamp { + mtime, + timestamp: read_started, + partial, + }, + ); + } +} + +// ── Write tracking ─────────────────────────────────────────────────────── + +/// Record that `agent_id` wrote `resolved_path`. +pub fn record_write(agent_id: &str, resolved_path: PathBuf) { + let Some(coord) = try_global() else { return }; + coord.record_write(agent_id, resolved_path); +} + +impl FileStateCoordinator { + /// Record a write on this coordinator; [`record_write`] delegates here. + pub(crate) fn record_write(&self, agent_id: &str, resolved_path: PathBuf) { + tracing::trace!( + agent = agent_id, + path = %resolved_path.display(), + "[file_state] record_write" + ); + let now = Instant::now(); + self.writes + .write() + .entry(resolved_path.clone()) + .or_default() + .insert(agent_id.to_string(), now); + // Also update this agent's own read stamp so its own subsequent + // writes don't trigger self-staleness. + self.reads.write().insert( + (agent_id.to_string(), resolved_path), + ReadStamp { + mtime: SystemTime::now(), + timestamp: now, + partial: false, + }, + ); + } +} + +// ── Staleness checks ───────────────────────────────────────────────────── + +/// Check whether `agent_id`'s view of `resolved_path` is stale because +/// another agent wrote to it after this agent's last read. Returns an +/// error message when stale, `None` when safe. +#[must_use] +pub fn check_stale_read(agent_id: &str, resolved_path: &Path) -> Option { + try_global()?.check_stale_read(agent_id, resolved_path) +} + +impl FileStateCoordinator { + /// [`check_stale_read`] against this coordinator. + pub(crate) fn check_stale_read(&self, agent_id: &str, resolved_path: &Path) -> Option { + let reads = self.reads.read(); + let writes = self.writes.read(); + let read_key = (agent_id.to_string(), resolved_path.to_path_buf()); + let read_stamp = reads.get(&read_key)?; + let path_writers = writes.get(resolved_path)?; + // Name the most recent of the writers that landed after the read; + // on a tie, the alphabetically first, so the message is deterministic. + let (latest, _) = writers_after_read(path_writers, agent_id, read_stamp.timestamp) + .max_by_key(|&(name, written_at)| (written_at, std::cmp::Reverse(name)))?; + let display_path = resolved_path.display(); + Some(format!( + "Stale read: file '{display_path}' was modified by agent '{latest}' after your last \ + read. Re-read the file before editing." + )) + } +} + +/// Check whether `agent_id`'s last read of `resolved_path` was partial. +/// Returns an error message when partial, `None` when safe. +#[must_use] +pub fn check_partial_read(agent_id: &str, resolved_path: &Path) -> Option { + let coord = try_global()?; + let reads = coord.reads.read(); + let read_key = (agent_id.to_string(), resolved_path.to_path_buf()); + let read_stamp = reads.get(&read_key)?; + if read_stamp.partial { + let display_path = resolved_path.display(); + Some(format!( + "Partial read: your last read of '{display_path}' was partial (paginated). \ + Perform a full read before overwriting." + )) + } else { + None + } +} + +// ── Path locking ───────────────────────────────────────────────────────── + +/// Acquire an async lock on `resolved_path` for a read-modify-write +/// section. Returns an `OwnedMutexGuard` that releases when dropped. +/// Returns `None` when the coordinator is disabled. +pub async fn acquire_path_lock(resolved_path: &Path) -> Option> { + let coord = try_global()?; + let mutex = { + let mut locks = coord.path_locks.write(); + locks + .entry(resolved_path.to_path_buf()) + .or_insert_with(|| Arc::new(Mutex::new(()))) + .clone() + }; + Some(mutex.lock_owned().await) +} + +// ── Parent reminder ────────────────────────────────────────────────────── + +/// Return resolved paths that `parent_agent_id` had previously read but +/// were subsequently written by any agent in `child_agent_ids`. +#[must_use] +pub fn parent_stale_files(parent_agent_id: &str, child_agent_ids: &[String]) -> Vec { + try_global().map_or_else(Vec::new, |coord| { + coord.parent_stale_files(parent_agent_id, child_agent_ids) + }) +} + +impl FileStateCoordinator { + /// [`parent_stale_files`] against this coordinator. + pub(crate) fn parent_stale_files( + &self, + parent_agent_id: &str, + child_agent_ids: &[String], + ) -> Vec { + self.stale_reads(parent_agent_id, |writer| { + child_agent_ids.iter().any(|child| child == writer) + }) + } +} diff --git a/crates/tinyagents-harness-tools/src/fs/file_state/types.rs b/crates/tinyagents-harness-tools/src/fs/file_state/types.rs new file mode 100644 index 000000000..28a4fa902 --- /dev/null +++ b/crates/tinyagents-harness-tools/src/fs/file_state/types.rs @@ -0,0 +1,122 @@ +//! Core types for the file state coordinator. + +use parking_lot::RwLock; +use std::collections::HashMap; +use std::path::PathBuf; +use std::sync::Arc; +use std::time::{Instant, SystemTime}; +use tokio::sync::Mutex; + +/// Snapshot of a single file read by an agent. +#[derive(Debug, Clone)] +pub struct ReadStamp { + /// Filesystem mtime at the moment of the read. + pub mtime: SystemTime, + /// Monotonic clock timestamp taken just before the read's I/O began, so + /// any write that lands during the read orders after it. + pub timestamp: Instant, + /// Whether the read was partial (paginated / offset+limit). + pub partial: bool, +} + +/// For one path: each agent that wrote it, mapped to the monotonic instant +/// of that agent's latest write. +pub(crate) type PathWriters = HashMap; + +/// Process-global coordinator that tracks file reads and writes across +/// all agents in the process. Thread-safe via `RwLock`. +#[derive(Debug)] +pub struct FileStateCoordinator { + /// Per-agent, per-resolved-path read stamps. + /// Key: `(agent_id, canonical_path)`. + pub(crate) reads: RwLock>, + + /// Per-resolved-path writers, each with its own latest write instant. + /// A later write by one agent never erases another agent's entry, so a + /// staleness check sees every writer, not just the most recent. + pub(crate) writes: RwLock>, + + /// Per-resolved-path async mutex for serialising read-modify-write + /// sections (used by `edit` and `apply_patch`). + pub(crate) path_locks: RwLock>>>, +} + +impl Default for FileStateCoordinator { + fn default() -> Self { + Self::new() + } +} + +impl FileStateCoordinator { + #[must_use] + /// Create an empty coordinator. + pub fn new() -> Self { + Self { + reads: RwLock::new(HashMap::new()), + writes: RwLock::new(HashMap::new()), + path_locks: RwLock::new(HashMap::new()), + } + } + + /// Return the set of resolved paths that `parent_agent_id` has read + /// but were subsequently written by any other agent. + #[must_use] + pub fn stale_reads_for_parent(&self, parent_agent_id: &str) -> Vec { + self.stale_reads(parent_agent_id, |_| true) + } + + /// Paths `reader` read that some other agent accepted by `counts` wrote + /// after that read. Sorted. + pub(crate) fn stale_reads(&self, reader: &str, counts: impl Fn(&str) -> bool) -> Vec { + let reads = self.reads.read(); + let writes = self.writes.read(); + let mut stale: Vec = reads + .iter() + .filter(|((agent_id, _), _)| agent_id == reader) + .filter(|((_, path), read_stamp)| { + writes.get(path).is_some_and(|writers| { + writers_after_read(writers, reader, read_stamp.timestamp) + .any(|(writer, _)| counts(writer)) + }) + }) + .map(|((_, path), _)| path.clone()) + .collect(); + stale.sort(); + stale + } + + /// Collect every path written by each agent in `agent_ids`, keyed by + /// agent. A path appears under every listed agent that ever wrote it, + /// even when a different agent wrote it afterwards. Paths are sorted; + /// agents with no writes are absent. + #[must_use] + pub fn paths_written_by(&self, agent_ids: &[String]) -> HashMap> { + let writes = self.writes.read(); + let mut result: HashMap> = HashMap::new(); + for (path, writers) in writes.iter() { + for agent_id in agent_ids.iter().filter(|id| writers.contains_key(*id)) { + result + .entry(agent_id.clone()) + .or_default() + .push(path.clone()); + } + } + for paths in result.values_mut() { + paths.sort(); + } + result + } +} + +/// The agents other than `reader` whose latest write to a path came after +/// `read_at`, each with that write's instant. +pub(crate) fn writers_after_read<'a>( + writers: &'a PathWriters, + reader: &'a str, + read_at: Instant, +) -> impl Iterator + 'a { + writers + .iter() + .filter(move |(writer, written_at)| writer.as_str() != reader && **written_at > read_at) + .map(|(writer, written_at)| (writer.as_str(), *written_at)) +} diff --git a/crates/tinyagents-harness-tools/src/fs/filesystem/apply_patch/mod.rs b/crates/tinyagents-harness-tools/src/fs/filesystem/apply_patch/mod.rs new file mode 100644 index 000000000..e6510c8d0 --- /dev/null +++ b/crates/tinyagents-harness-tools/src/fs/filesystem/apply_patch/mod.rs @@ -0,0 +1,474 @@ +//! `apply_patch` — atomic multi-edit across one or more files. +//! +//! Coding-harness baseline tool (issue #1205). Takes an array of +//! `{path, old_string, new_string}` edits and applies them atomically: +//! every edit is validated up front (path, exact-match, uniqueness) +//! before any file is written. If any edit fails validation, no files +//! are touched. +//! +//! **Creating a file** is an empty `old_string` against a path that does not +//! exist yet; `new_string` becomes the whole contents. Before that existed the +//! tool could only edit, and an agent asked to produce a document had no route +//! at all: it would create a placeholder with `shell` purely so there was +//! something to patch (the life-scenario benchmark caught two 1-byte files +//! containing `x`). An empty `old_string` against a file that *does* exist is +//! still an error — "replace nothing" is ambiguous, not a create. + +use super::gate::{FsGate, gate_for_context}; +use crate::fs::file_state; +use async_trait::async_trait; +use serde_json::json; +use std::collections::HashMap; +use std::path::PathBuf; +use std::sync::Arc; +use tinytools::ToolRunContext; +use tinytools::{PermissionLevel, Tool, ToolCallOptions, ToolResult}; + +const MAX_FILE_BYTES: u64 = 5 * 1024 * 1024; +const MAX_EDITS: usize = 50; + +/// Apply a multi-edit patch to workspace files atomically. +#[derive(Debug)] +pub struct ApplyPatchTool { + gate: Arc, +} + +impl ApplyPatchTool { + /// Create the tool. + #[must_use] + pub fn new(gate: Arc) -> Self { + Self { gate } + } +} + +#[async_trait] +impl Tool for ApplyPatchTool { + fn name(&self) -> &'static str { + "apply_patch" + } + + fn description(&self) -> &'static str { + "Apply a batch of exact-string edits across one or more files atomically. \ + All edits are validated before any are written; validation failure rolls \ + back the whole batch. Each edit is `{path, old_string, new_string, replace_all?}`; \ + a top-level `path` is the default for edits that omit their own. \ + To CREATE a new file, pass an empty `old_string` with the full contents \ + as `new_string`; the path must not already exist." + } + + fn parameters_schema(&self) -> serde_json::Value { + json!({ + "type": "object", + "properties": { + "path": { + "type": "string", + "description": "Default file for edits that omit their own `path`." + }, + "edits": { + "type": "array", + "description": "Ordered list of edits.", + "items": { + "type": "object", + "properties": { + "path": { + "type": "string", + "description": "File to edit; defaults to the top-level `path`." + }, + "old_string": { + "type": "string", + "description": "Exact text to replace. Empty means CREATE: the path must not exist and `new_string` becomes the whole file." + }, + "new_string": { "type": "string" }, + "replace_all": { "type": "boolean", "default": false } + }, + "required": ["old_string", "new_string"] + } + } + }, + "required": ["edits"] + }) + } + + fn permission_level(&self) -> PermissionLevel { + PermissionLevel::Write + } + + /// `apply_patch` modifies existing files → in ask-before-edit it routes + /// through the human approval gate; in Full it runs; read-only is blocked + /// in `execute`. + fn external_effect_with_args(&self, _args: &serde_json::Value) -> bool { + self.gate.write_needs_approval() + } + + async fn execute(&self, args: serde_json::Value) -> anyhow::Result { + self.execute_in_context(args, None).await + } + + async fn execute_with_context( + &self, + args: serde_json::Value, + _options: ToolCallOptions, + context: Option<&dyn ToolRunContext>, + ) -> anyhow::Result { + self.execute_in_context(args, context).await + } +} + +impl ApplyPatchTool { + async fn execute_in_context( + &self, + args: serde_json::Value, + context: Option<&dyn ToolRunContext>, + ) -> anyhow::Result { + let edits = args + .get("edits") + .and_then(|v| v.as_array()) + .ok_or_else(|| anyhow::anyhow!("Missing 'edits' array"))?; + if edits.is_empty() { + return Ok(ToolResult::error("`edits` array is empty")); + } + if edits.len() > MAX_EDITS { + return Ok(ToolResult::error(format!( + "Too many edits: {} (max {MAX_EDITS})", + edits.len() + ))); + } + + if !self.gate.can_act() { + return Ok(ToolResult::error( + "[policy-blocked] Action blocked: autonomy is read-only", + )); + } + if self.gate.is_rate_limited() { + return Ok(ToolResult::error( + "Rate limit exceeded: too many actions in the last hour", + )); + } + if !self.gate.record_action() { + return Ok(ToolResult::error( + "Rate limit exceeded: action budget exhausted", + )); + } + + let path_policy = gate_for_context(&self.gate, context, "apply_patch"); + + // "One file, several edits" is a natural call shape, and models write + // it with the path once at the top level (4 of 10 calls in one run; + // the per-edit `path` requirement rejected every one of them and the + // run halted on the fourth). The top-level path is the default; an + // edit's own path still wins. + let default_path = match args.get("path") { + Some(value) => Some( + value + .as_str() + .ok_or_else(|| anyhow::anyhow!("top-level `path` must be a string"))?, + ), + None => None, + }; + + // Parse + group edits by file. + let mut parsed: Vec = Vec::with_capacity(edits.len()); + for (i, raw) in edits.iter().enumerate() { + let path = match raw.get("path") { + Some(value) => value.as_str().ok_or_else(|| { + anyhow::anyhow!("edit[{i}]: `path` must be a string") + })?, + None => default_path.ok_or_else(|| { + anyhow::anyhow!( + "edit[{i}]: missing `path` (give each edit a `path`, or one top-level `path` for all edits)" + ) + })?, + }; + let old_string = raw + .get("old_string") + .and_then(|v| v.as_str()) + .ok_or_else(|| anyhow::anyhow!("edit[{i}]: missing `old_string`"))?; + let new_string = raw + .get("new_string") + .and_then(|v| v.as_str()) + .ok_or_else(|| anyhow::anyhow!("edit[{i}]: missing `new_string`"))?; + let replace_all = raw + .get("replace_all") + .and_then(serde_json::Value::as_bool) + .unwrap_or(false); + + if !path_policy.is_path_string_allowed(path) { + return Ok(ToolResult::error(format!( + "edit[{i}]: path not allowed: {path}" + ))); + } + // An empty `old_string` is a create, and only against a path that + // does not exist yet. Resolved the same way every other edit is — + // joined onto `action_dir` — so "exists" means the same thing here + // as it does in the apply loop below. + let exists = path_policy.action_dir().join(path).exists() + || (std::path::Path::new(path).is_absolute() + && std::path::Path::new(path).exists()); + let create = old_string.is_empty(); + if create && exists { + return Ok(ToolResult::error(format!( + "edit[{i}]: `old_string` must not be empty for an existing file ({path}). \ + Pass the exact text to replace, or write to a new path to create a file." + ))); + } + parsed.push(ParsedEdit { + index: i, + path: path.to_string(), + old_string: old_string.to_string(), + new_string: new_string.to_string(), + replace_all, + create, + }); + } + + // Acquire per-path locks for all unique paths before any reads. + let unique_paths: Vec = { + let mut seen = std::collections::HashSet::new(); + parsed + .iter() + .filter_map(|e| { + if seen.insert(e.path.clone()) { + Some(e.path.clone()) + } else { + None + } + }) + .collect() + }; + let mut _path_guards = Vec::new(); + for p in &unique_paths { + let full = path_policy.action_dir().join(p); + if let Ok(resolved) = tokio::fs::canonicalize(&full).await + && let Some(guard) = file_state::acquire_path_lock(&resolved).await + { + _path_guards.push(guard); + } + } + + // File-state guard: reject edits based on stale or partial reads. + if let Some(agent_id) = file_state::current_file_state_agent_id() { + for p in &unique_paths { + let full = path_policy.action_dir().join(p); + if let Ok(resolved) = tokio::fs::canonicalize(&full).await { + if let Some(msg) = file_state::check_stale_read(&agent_id, &resolved) { + tracing::debug!( + agent = %agent_id, + path = %resolved.display(), + "[file_state] apply_patch blocked: stale read" + ); + return Ok(ToolResult::error(msg)); + } + if let Some(msg) = file_state::check_partial_read(&agent_id, &resolved) { + tracing::debug!( + agent = %agent_id, + path = %resolved.display(), + "[file_state] apply_patch blocked: partial read" + ); + return Ok(ToolResult::error(msg)); + } + } + } + } + + // Resolve paths + load file contents (once per file). Apply edits in + // memory; if any edit fails, return without writing. + let mut buffers: HashMap = HashMap::new(); + for edit in &parsed { + if !buffers.contains_key(&edit.path) { + let full = path_policy.action_dir().join(&edit.path); + + // Symlink check must happen on the *unresolved* path — + // canonicalize resolves symlinks, so a check after that + // point would never see the link. + if let Ok(meta) = tokio::fs::symlink_metadata(&full).await + && meta.file_type().is_symlink() + { + return Ok(ToolResult::error(format!( + "edit[{}]: refusing to edit through symlink", + edit.index + ))); + } + + // Security check: validate path string, resolve symlinks, confirm + // workspace containment. A create has no file to canonicalize, + // so it resolves through the parent — the same gate + // `file_write` uses, which walks up to the deepest existing + // ancestor and still checks it for symlink escapes. + let resolved = if edit.create { + match path_policy.validate_parent_path(&edit.path).await { + Ok(p) => p, + Err(msg) => { + return Ok(ToolResult::error(format!("edit[{}]: {msg}", edit.index))); + } + } + } else { + match path_policy.validate_path(&edit.path).await { + Ok(p) => p, + Err(msg) => { + return Ok(ToolResult::error(format!("edit[{}]: {msg}", edit.index))); + } + } + }; + if edit.create { + if let Some(parent) = resolved.parent() + && let Err(e) = tokio::fs::create_dir_all(parent).await + { + return Ok(ToolResult::error(format!( + "edit[{}]: failed to create parent of {}: {e}", + edit.index, edit.path + ))); + } + buffers.insert( + edit.path.clone(), + FileBuffer { + resolved, + original: None, + contents: edit.new_string.clone(), + edit_count: 1, + }, + ); + continue; + } + if let Ok(meta) = tokio::fs::metadata(&resolved).await + && meta.len() > MAX_FILE_BYTES + { + return Ok(ToolResult::error(format!( + "edit[{}]: file too large ({} bytes)", + edit.index, + meta.len() + ))); + } + let contents = match tokio::fs::read_to_string(&resolved).await { + Ok(c) => c, + Err(e) => { + return Ok(ToolResult::error(format!( + "edit[{}]: failed to read {}: {e}", + edit.index, edit.path + ))); + } + }; + buffers.insert( + edit.path.clone(), + FileBuffer { + resolved, + original: Some(contents.clone()), + contents, + edit_count: 0, + }, + ); + } + + if edit.create { + // A second create for the same path in one batch. The first one + // already populated the buffer; a duplicate would silently + // discard one of the two bodies, so say so. + return Ok(ToolResult::error(format!( + "edit[{}]: {} is created earlier in this batch; only one create per path", + edit.index, edit.path + ))); + } + + let buf = buffers.get_mut(&edit.path).unwrap(); + let count = buf.contents.matches(&edit.old_string).count(); + if count == 0 { + return Ok(ToolResult::error(format!( + "edit[{}]: `old_string` not found in {}", + edit.index, edit.path + ))); + } + if count > 1 && !edit.replace_all { + return Ok(ToolResult::error(format!( + "edit[{}]: `old_string` matches {count} times in {}; pass `replace_all`", + edit.index, edit.path + ))); + } + buf.contents = if edit.replace_all { + buf.contents.replace(&edit.old_string, &edit.new_string) + } else { + buf.contents.replacen(&edit.old_string, &edit.new_string, 1) + }; + buf.edit_count += count; + } + + // Best-effort atomic write across files. We cannot get true + // multi-file atomicity without filesystem-level transactions, + // but if the i-th write fails we attempt to restore originals + // for the i-1 already-written files from the in-memory snapshot. + let mut summary: Vec = Vec::new(); + let mut written: Vec<&FileBuffer> = Vec::new(); + for (path, buf) in &buffers { + if let Err(e) = tokio::fs::write(&buf.resolved, &buf.contents).await { + let restore_errors = restore_originals(&written).await; + let suffix = if restore_errors.is_empty() { + "; previously-written files restored from snapshot".to_string() + } else { + format!("; restore failed for: {}", restore_errors.join(", ")) + }; + return Ok(ToolResult::error(format!( + "Failed to write {path}: {e}{suffix}" + ))); + } + written.push(buf); + summary.push(match buf.original { + None => format!("{path}: created ({} bytes)", buf.contents.len()), + Some(_) => format!("{path}: {} replacement(s)", buf.edit_count), + }); + } + // Record writes in the file-state coordinator. + if let Some(agent_id) = file_state::current_file_state_agent_id() { + for buf in buffers.values() { + file_state::record_write(&agent_id, buf.resolved.clone()); + } + } + + summary.sort(); + Ok(ToolResult::success(format!( + "Applied {} edit(s) across {} file(s)\n{}", + parsed.len(), + buffers.len(), + summary.join("\n") + ))) + } +} + +async fn restore_originals(written: &[&FileBuffer]) -> Vec { + let mut errors = Vec::new(); + for buf in written { + let result = match &buf.original { + Some(original) => tokio::fs::write(&buf.resolved, original).await, + // The batch created this file, so "restore" means remove it — + // otherwise a failed multi-file patch leaves a half-written new + // file behind and the caller cannot tell it apart from a success. + None => tokio::fs::remove_file(&buf.resolved).await, + }; + if let Err(e) = result { + errors.push(format!("{}: {e}", buf.resolved.display())); + } + } + errors +} + +struct ParsedEdit { + index: usize, + path: String, + old_string: String, + new_string: String, + replace_all: bool, + /// Empty `old_string` against a path that does not exist: `new_string` + /// becomes the whole file. + create: bool, +} + +struct FileBuffer { + resolved: PathBuf, + /// Snapshot of the file's contents as we first read them. + /// Used to restore on a partial-write failure. `None` for a file this batch + /// is creating — there is nothing to restore *to*, so rollback deletes it. + original: Option, + contents: String, + edit_count: usize, +} + +#[cfg(test)] +#[path = "mod_tests.rs"] +mod tests; diff --git a/crates/tinyagents-harness-tools/src/fs/filesystem/apply_patch/mod_tests.rs b/crates/tinyagents-harness-tools/src/fs/filesystem/apply_patch/mod_tests.rs new file mode 100644 index 000000000..769662715 --- /dev/null +++ b/crates/tinyagents-harness-tools/src/fs/filesystem/apply_patch/mod_tests.rs @@ -0,0 +1,648 @@ +//! Behavior tests for the `apply_patch` tool, driven through a fake [`FsGate`]. + +#![allow(clippy::expect_used, clippy::panic, clippy::unwrap_used)] + +use super::*; +use crate::fs::filesystem::test_support::TestGate; + +fn test_security(workspace: std::path::PathBuf) -> Arc { + TestGate::supervised(workspace) +} + +#[test] +fn apply_patch_name() { + let tool = ApplyPatchTool::new(test_security(std::env::temp_dir())); + assert_eq!(tool.name(), "apply_patch"); +} + +#[tokio::test] +async fn apply_patch_applies_multiple_edits() { + let dir = std::env::temp_dir().join("openhuman_test_patch_multi"); + let _ = tokio::fs::remove_dir_all(&dir).await; + tokio::fs::create_dir_all(&dir).await.unwrap(); + tokio::fs::write(dir.join("a.txt"), "alpha\nbravo") + .await + .unwrap(); + tokio::fs::write(dir.join("b.txt"), "one two") + .await + .unwrap(); + + let tool = ApplyPatchTool::new(test_security(dir.clone())); + let result = tool + .execute(json!({ + "edits": [ + { "path": "a.txt", "old_string": "alpha", "new_string": "ALPHA" }, + { "path": "b.txt", "old_string": "two", "new_string": "TWO" } + ] + })) + .await + .unwrap(); + assert!(!result.is_error, "{}", result.output()); + let a = tokio::fs::read_to_string(dir.join("a.txt")).await.unwrap(); + let b = tokio::fs::read_to_string(dir.join("b.txt")).await.unwrap(); + assert_eq!(a, "ALPHA\nbravo"); + assert_eq!(b, "one TWO"); + + let _ = tokio::fs::remove_dir_all(&dir).await; +} + +#[tokio::test] +async fn apply_patch_atomic_on_validation_failure() { + let dir = std::env::temp_dir().join("openhuman_test_patch_atomic"); + let _ = tokio::fs::remove_dir_all(&dir).await; + tokio::fs::create_dir_all(&dir).await.unwrap(); + tokio::fs::write(dir.join("a.txt"), "alpha").await.unwrap(); + tokio::fs::write(dir.join("b.txt"), "bravo").await.unwrap(); + + let tool = ApplyPatchTool::new(test_security(dir.clone())); + // Second edit will fail (no match) — first must NOT be applied. + let result = tool + .execute(json!({ + "edits": [ + { "path": "a.txt", "old_string": "alpha", "new_string": "ALPHA" }, + { "path": "b.txt", "old_string": "missing", "new_string": "x" } + ] + })) + .await + .unwrap(); + assert!(result.is_error); + let a = tokio::fs::read_to_string(dir.join("a.txt")).await.unwrap(); + assert_eq!(a, "alpha", "atomic: first edit must not be persisted"); + + let _ = tokio::fs::remove_dir_all(&dir).await; +} + +#[tokio::test] +async fn apply_patch_chained_edits_same_file() { + let dir = std::env::temp_dir().join("openhuman_test_patch_chain"); + let _ = tokio::fs::remove_dir_all(&dir).await; + tokio::fs::create_dir_all(&dir).await.unwrap(); + tokio::fs::write(dir.join("a.txt"), "one two three") + .await + .unwrap(); + + let tool = ApplyPatchTool::new(test_security(dir.clone())); + let result = tool + .execute(json!({ + "edits": [ + { "path": "a.txt", "old_string": "one", "new_string": "ONE" }, + { "path": "a.txt", "old_string": "two", "new_string": "TWO" } + ] + })) + .await + .unwrap(); + assert!(!result.is_error, "{}", result.output()); + let updated = tokio::fs::read_to_string(dir.join("a.txt")).await.unwrap(); + assert_eq!(updated, "ONE TWO three"); + + let _ = tokio::fs::remove_dir_all(&dir).await; +} + +#[tokio::test] +async fn apply_patch_rejects_empty_edits() { + let dir = std::env::temp_dir().join("openhuman_test_patch_empty"); + let _ = tokio::fs::remove_dir_all(&dir).await; + tokio::fs::create_dir_all(&dir).await.unwrap(); + + let tool = ApplyPatchTool::new(test_security(dir.clone())); + let result = tool.execute(json!({"edits": []})).await.unwrap(); + assert!(result.is_error); + + let _ = tokio::fs::remove_dir_all(&dir).await; +} + +#[tokio::test] +async fn apply_patch_rejects_traversal() { + let tool = ApplyPatchTool::new(test_security(std::env::temp_dir())); + let result = tool + .execute(json!({ + "edits": [ + { "path": "../etc/passwd", "old_string": "x", "new_string": "y" } + ] + })) + .await + .unwrap(); + assert!(result.is_error); + assert!(result.output().contains("not allowed")); +} + +// -- create mode -------------------------------------------------------------- +// +// Empty `old_string` + a path that does not exist = create. Before this the +// tool could only edit, so an agent asked to produce a document had no route: +// the life-scenario `meal-plan` run left two 1-byte files containing `x`, +// placeholders it made with `shell` purely so `apply_patch` had something to +// patch. + +#[tokio::test] +async fn apply_patch_creates_a_new_file_from_an_empty_old_string() { + let dir = std::env::temp_dir().join("openhuman_test_patch_create"); + let _ = tokio::fs::remove_dir_all(&dir).await; + tokio::fs::create_dir_all(&dir).await.unwrap(); + + let tool = ApplyPatchTool::new(test_security(dir.clone())); + let result = tool + .execute(json!({ + "edits": [ + { "path": "out/plan.md", "old_string": "", "new_string": "# Plan\nday one\n" } + ] + })) + .await + .unwrap(); + assert!(!result.is_error, "{}", result.output()); + let written = tokio::fs::read_to_string(dir.join("out/plan.md")) + .await + .expect("the file (and its parent) must have been created"); + assert_eq!(written, "# Plan\nday one\n"); + assert!(result.output().contains("created"), "{}", result.output()); +} + +#[tokio::test] +async fn apply_patch_takes_a_top_level_path_as_the_default_for_every_edit() { + // The shape a model writes for "one file, several edits": the path once, + // at the top level. Each such call used to be rejected for a missing + // `edits[0].path`, and four in a row halted a run with 26 minutes left. + let dir = std::env::temp_dir().join("openhuman_test_patch_top_level_path"); + let _ = tokio::fs::remove_dir_all(&dir).await; + tokio::fs::create_dir_all(&dir).await.unwrap(); + tokio::fs::write(dir.join("a.c"), "int x = 1;\nint y = 2;\n") + .await + .unwrap(); + + let tool = ApplyPatchTool::new(test_security(dir.clone())); + let result = tool + .execute(json!({ + "path": "a.c", + "edits": [ + { "old_string": "int x = 1;", "new_string": "int x = 10;" }, + { "old_string": "int y = 2;", "new_string": "int y = 20;" } + ] + })) + .await + .unwrap(); + assert!(!result.is_error, "{}", result.output()); + assert_eq!( + tokio::fs::read_to_string(dir.join("a.c")).await.unwrap(), + "int x = 10;\nint y = 20;\n" + ); +} + +#[tokio::test] +async fn apply_patch_rejects_malformed_edit_path_instead_of_using_default() { + let dir = std::env::temp_dir().join("openhuman_test_patch_malformed_path"); + let _ = tokio::fs::remove_dir_all(&dir).await; + tokio::fs::create_dir_all(&dir).await.unwrap(); + tokio::fs::write(dir.join("default.txt"), "original") + .await + .unwrap(); + + let tool = ApplyPatchTool::new(test_security(dir.clone())); + let result = tool + .execute(json!({ + "path": "default.txt", + "edits": [{ "path": null, "old_string": "original", "new_string": "changed" }] + })) + .await; + + assert!(result.is_err()); + assert!( + result + .unwrap_err() + .to_string() + .contains("`path` must be a string") + ); + assert_eq!( + tokio::fs::read_to_string(dir.join("default.txt")) + .await + .unwrap(), + "original" + ); + let _ = tokio::fs::remove_dir_all(&dir).await; +} + +#[tokio::test] +async fn apply_patch_rejects_a_malformed_top_level_path() { + let dir = std::env::temp_dir().join("openhuman_test_patch_malformed_top_level_path"); + let _ = tokio::fs::remove_dir_all(&dir).await; + tokio::fs::create_dir_all(&dir).await.unwrap(); + tokio::fs::write(dir.join("a.txt"), "original") + .await + .unwrap(); + + let tool = ApplyPatchTool::new(test_security(dir.clone())); + let result = tool + .execute(json!({ + "path": 123, + "edits": [{ "path": "a.txt", "old_string": "original", "new_string": "changed" }] + })) + .await; + + assert!(result.is_err()); + assert!( + result + .unwrap_err() + .to_string() + .contains("top-level `path` must be a string") + ); + assert_eq!( + tokio::fs::read_to_string(dir.join("a.txt")).await.unwrap(), + "original" + ); + let _ = tokio::fs::remove_dir_all(&dir).await; +} + +#[tokio::test] +async fn apply_patch_lets_an_edits_own_path_win_over_the_top_level_one() { + let dir = std::env::temp_dir().join("openhuman_test_patch_top_level_path_override"); + let _ = tokio::fs::remove_dir_all(&dir).await; + tokio::fs::create_dir_all(&dir).await.unwrap(); + tokio::fs::write(dir.join("a.txt"), "alpha").await.unwrap(); + tokio::fs::write(dir.join("b.txt"), "bravo").await.unwrap(); + + let tool = ApplyPatchTool::new(test_security(dir.clone())); + let result = tool + .execute(json!({ + "path": "a.txt", + "edits": [ + { "old_string": "alpha", "new_string": "ALPHA" }, + { "path": "b.txt", "old_string": "bravo", "new_string": "BRAVO" } + ] + })) + .await + .unwrap(); + assert!(!result.is_error, "{}", result.output()); + assert_eq!( + tokio::fs::read_to_string(dir.join("a.txt")).await.unwrap(), + "ALPHA" + ); + assert_eq!( + tokio::fs::read_to_string(dir.join("b.txt")).await.unwrap(), + "BRAVO" + ); +} + +#[tokio::test] +async fn apply_patch_names_both_ways_to_give_a_path_when_neither_is_given() { + let dir = std::env::temp_dir().join("openhuman_test_patch_no_path_anywhere"); + let _ = tokio::fs::remove_dir_all(&dir).await; + tokio::fs::create_dir_all(&dir).await.unwrap(); + + let tool = ApplyPatchTool::new(test_security(dir.clone())); + let err = tool + .execute(json!({ + "edits": [ { "old_string": "a", "new_string": "b" } ] + })) + .await + .expect_err("no path anywhere is an argument error"); + let text = err.to_string(); + assert!(text.contains("edit[0]: missing `path`"), "{text}"); + assert!(text.contains("one top-level `path`"), "{text}"); +} + +#[tokio::test] +async fn apply_patch_refuses_an_empty_old_string_on_an_existing_file() { + let dir = std::env::temp_dir().join("openhuman_test_patch_create_existing"); + let _ = tokio::fs::remove_dir_all(&dir).await; + tokio::fs::create_dir_all(&dir).await.unwrap(); + tokio::fs::write(dir.join("a.txt"), "alpha").await.unwrap(); + + let tool = ApplyPatchTool::new(test_security(dir.clone())); + let result = tool + .execute(json!({ + "edits": [{ "path": "a.txt", "old_string": "", "new_string": "overwritten" }] + })) + .await + .unwrap(); + assert!(result.is_error, "{}", result.output()); + // Unchanged — "replace nothing" must never become "replace everything". + let a = tokio::fs::read_to_string(dir.join("a.txt")).await.unwrap(); + assert_eq!(a, "alpha"); +} + +#[tokio::test] +async fn apply_patch_mixes_a_create_and_an_edit_in_one_batch() { + let dir = std::env::temp_dir().join("openhuman_test_patch_create_mixed"); + let _ = tokio::fs::remove_dir_all(&dir).await; + tokio::fs::create_dir_all(&dir).await.unwrap(); + tokio::fs::write(dir.join("a.txt"), "alpha").await.unwrap(); + + let tool = ApplyPatchTool::new(test_security(dir.clone())); + let result = tool + .execute(json!({ + "edits": [ + { "path": "a.txt", "old_string": "alpha", "new_string": "ALPHA" }, + { "path": "b.txt", "old_string": "", "new_string": "bravo" } + ] + })) + .await + .unwrap(); + assert!(!result.is_error, "{}", result.output()); + assert_eq!( + tokio::fs::read_to_string(dir.join("a.txt")).await.unwrap(), + "ALPHA" + ); + assert_eq!( + tokio::fs::read_to_string(dir.join("b.txt")).await.unwrap(), + "bravo" + ); +} + +#[tokio::test] +async fn apply_patch_refuses_two_creates_for_the_same_path() { + let dir = std::env::temp_dir().join("openhuman_test_patch_create_dup"); + let _ = tokio::fs::remove_dir_all(&dir).await; + tokio::fs::create_dir_all(&dir).await.unwrap(); + + let tool = ApplyPatchTool::new(test_security(dir.clone())); + let result = tool + .execute(json!({ + "edits": [ + { "path": "c.txt", "old_string": "", "new_string": "first" }, + { "path": "c.txt", "old_string": "", "new_string": "second" } + ] + })) + .await + .unwrap(); + assert!(result.is_error, "{}", result.output()); + assert!( + !dir.join("c.txt").exists(), + "a rejected batch must write nothing" + ); +} + +#[tokio::test] +async fn apply_patch_create_still_obeys_the_path_policy() { + let dir = std::env::temp_dir().join("openhuman_test_patch_create_escape"); + let _ = tokio::fs::remove_dir_all(&dir).await; + tokio::fs::create_dir_all(&dir).await.unwrap(); + + let tool = ApplyPatchTool::new(test_security(dir.clone())); + let result = tool + .execute(json!({ + "edits": [{ "path": "../escaped.md", "old_string": "", "new_string": "nope" }] + })) + .await + .unwrap(); + assert!(result.is_error, "{}", result.output()); +} + +// ── Coverage of limits, budgets, guards, symlinks and write failures ──── + +use crate::fs::filesystem::test_support::{AutonomyLevel, RacyGate, WorkspaceContext}; + +fn one_edit(path: &str, old: &str, new: &str) -> serde_json::Value { + json!({"edits": [{"path": path, "old_string": old, "new_string": new}]}) +} + +#[tokio::test] +async fn apply_patch_rejects_too_many_edits() { + let dir = tempfile::tempdir().unwrap(); + let edits: Vec<_> = (0..=MAX_EDITS) + .map(|_| json!({"path": "a.txt", "old_string": "a", "new_string": "b"})) + .collect(); + let tool = ApplyPatchTool::new(test_security(dir.path().to_path_buf())); + let result = tool.execute(json!({"edits": edits})).await.unwrap(); + assert!(result.is_error); + assert!(result.output().contains("Too many edits")); +} + +#[tokio::test] +async fn apply_patch_enforces_autonomy_and_budgets() { + let dir = tempfile::tempdir().unwrap(); + std::fs::write(dir.path().join("a.txt"), "x").unwrap(); + let args = one_edit("a.txt", "x", "y"); + + let read_only = ApplyPatchTool::new(TestGate::with( + dir.path().to_path_buf(), + AutonomyLevel::ReadOnly, + 100, + )); + let result = read_only.execute(args.clone()).await.unwrap(); + assert!(result.output().contains("autonomy is read-only")); + + let limited = ApplyPatchTool::new(TestGate::with( + dir.path().to_path_buf(), + AutonomyLevel::Supervised, + 0, + )); + let result = limited.execute(args.clone()).await.unwrap(); + assert!( + result + .output() + .contains("too many actions in the last hour") + ); + + let racy = ApplyPatchTool::new(Arc::new(RacyGate(test_security(dir.path().to_path_buf())))); + let result = racy.execute(args).await.unwrap(); + assert!(result.output().contains("action budget exhausted")); +} + +#[tokio::test] +async fn apply_patch_resolves_an_absolute_create_target() { + let dir = tempfile::tempdir().unwrap(); + let target = dir.path().join("abs_new.txt"); + let tool = ApplyPatchTool::new(test_security(dir.path().to_path_buf())); + let result = tool + .execute(one_edit(target.to_str().unwrap(), "", "created")) + .await + .unwrap(); + assert!(!result.is_error, "{}", result.output()); + assert_eq!(std::fs::read_to_string(target).unwrap(), "created"); +} + +#[cfg(unix)] +#[tokio::test] +async fn apply_patch_refuses_to_edit_through_a_symlink() { + let dir = tempfile::tempdir().unwrap(); + std::fs::write(dir.path().join("real.txt"), "x").unwrap(); + std::os::unix::fs::symlink(dir.path().join("real.txt"), dir.path().join("link.txt")).unwrap(); + let tool = ApplyPatchTool::new(test_security(dir.path().to_path_buf())); + let result = tool.execute(one_edit("link.txt", "x", "y")).await.unwrap(); + assert!(result.is_error); + assert!(result.output().contains("refusing to edit through symlink")); +} + +#[tokio::test] +async fn apply_patch_reports_unresolvable_and_unreadable_targets() { + let dir = tempfile::tempdir().unwrap(); + let tool = ApplyPatchTool::new(test_security(dir.path().to_path_buf())); + let result = tool + .execute(one_edit("absent.txt", "x", "y")) + .await + .unwrap(); + assert!(result.is_error); + assert!(result.output().starts_with("edit[0]: ")); + assert!(result.output().contains("Failed to resolve")); + + std::fs::write(dir.path().join("bin.dat"), [0xff_u8, 0xfe]).unwrap(); + let result = tool.execute(one_edit("bin.dat", "x", "y")).await.unwrap(); + assert!(result.output().contains("failed to read bin.dat")); +} + +#[tokio::test] +async fn apply_patch_reports_a_create_whose_parent_cannot_be_made() { + let dir = tempfile::tempdir().unwrap(); + std::fs::write(dir.path().join("file.txt"), "x").unwrap(); + let tool = ApplyPatchTool::new(test_security(dir.path().to_path_buf())); + let result = tool + .execute(one_edit("file.txt/child.txt", "", "body")) + .await + .unwrap(); + assert!(result.is_error); + assert!( + result + .output() + .contains("failed to create parent of file.txt/child.txt") + ); +} + +#[tokio::test] +async fn apply_patch_refuses_an_oversized_file() { + let dir = tempfile::tempdir().unwrap(); + let file = std::fs::File::create(dir.path().join("big.txt")).unwrap(); + file.set_len(MAX_FILE_BYTES + 1).unwrap(); + let tool = ApplyPatchTool::new(test_security(dir.path().to_path_buf())); + let result = tool.execute(one_edit("big.txt", "x", "y")).await.unwrap(); + assert!(result.output().contains("file too large")); +} + +#[tokio::test] +async fn apply_patch_needs_replace_all_for_repeated_matches() { + let dir = tempfile::tempdir().unwrap(); + std::fs::write(dir.path().join("r.txt"), "a a a").unwrap(); + let tool = ApplyPatchTool::new(test_security(dir.path().to_path_buf())); + let result = tool.execute(one_edit("r.txt", "a", "b")).await.unwrap(); + assert!(result.is_error); + assert!(result.output().contains("matches 3 times")); + assert_eq!( + std::fs::read_to_string(dir.path().join("r.txt")).unwrap(), + "a a a" + ); + + let result = tool + .execute(json!({"edits": [ + {"path": "r.txt", "old_string": "a", "new_string": "b", "replace_all": true} + ]})) + .await + .unwrap(); + assert!(!result.is_error, "{}", result.output()); + assert_eq!( + std::fs::read_to_string(dir.path().join("r.txt")).unwrap(), + "b b b" + ); +} + +#[tokio::test] +async fn apply_patch_uses_the_context_workspace_when_one_is_threaded() { + let home = tempfile::tempdir().unwrap(); + let isolated = tempfile::tempdir().unwrap(); + std::fs::write(isolated.path().join("w.txt"), "x").unwrap(); + let tool = ApplyPatchTool::new(test_security(home.path().to_path_buf())); + let context = WorkspaceContext::at(isolated.path()); + let result = tool + .execute_with_context( + one_edit("w.txt", "x", "y"), + ToolCallOptions::default(), + Some(&context), + ) + .await + .unwrap(); + assert!(!result.is_error, "{}", result.output()); + assert_eq!( + std::fs::read_to_string(isolated.path().join("w.txt")).unwrap(), + "y" + ); +} + +#[tokio::test] +async fn apply_patch_honours_the_file_state_guard_and_records_writes() { + crate::fs::file_state::init_global(true); + let dir = tempfile::tempdir().unwrap(); + let target = dir.path().canonicalize().unwrap().join("g.txt"); + std::fs::write(&target, "x").unwrap(); + let tool = ApplyPatchTool::new(test_security(dir.path().to_path_buf())); + let agent = format!("ap-agent-{}", dir.path().display()); + let other = format!("ap-other-{}", dir.path().display()); + let run = |agent: String| { + crate::fs::file_state::with_file_state_agent_id( + agent, + tool.execute(one_edit("g.txt", "x", "y")), + ) + }; + + crate::fs::file_state::record_read( + &agent, + target.clone(), + std::time::SystemTime::now(), + true, + std::time::Instant::now(), + ); + assert!( + run(agent.clone()) + .await + .unwrap() + .output() + .contains("Partial read") + ); + + crate::fs::file_state::record_read( + &agent, + target.clone(), + std::time::SystemTime::now(), + false, + std::time::Instant::now(), + ); + crate::fs::file_state::record_write(&other, target.clone()); + assert!( + run(agent.clone()) + .await + .unwrap() + .output() + .contains("Stale read") + ); + + crate::fs::file_state::record_read( + &agent, + target.clone(), + std::time::SystemTime::now(), + false, + std::time::Instant::now(), + ); + let result = run(agent.clone()).await.unwrap(); + assert!(!result.is_error, "{}", result.output()); + assert!(crate::fs::file_state::check_stale_read(&agent, &target).is_none()); +} + +/// A create whose file name exceeds the OS limit passes every check and then +/// fails at write time. Whichever order the batch writes in, no file from the +/// failed batch may be left changed, and the error must say it was restored. +#[tokio::test] +async fn apply_patch_restores_earlier_writes_when_a_later_write_fails() { + let long_name = "n".repeat(300); + let mut saw_restore = false; + for _ in 0..40 { + let dir = tempfile::tempdir().unwrap(); + std::fs::write(dir.path().join("keep.txt"), "orig").unwrap(); + let tool = ApplyPatchTool::new(test_security(dir.path().to_path_buf())); + let result = tool + .execute(json!({"edits": [ + {"path": "keep.txt", "old_string": "orig", "new_string": "changed"}, + {"path": "fresh.txt", "old_string": "", "new_string": "new"}, + {"path": long_name, "old_string": "", "new_string": "boom"}, + ]})) + .await + .unwrap(); + assert!(result.is_error, "{}", result.output()); + assert!(result.output().contains("Failed to write")); + assert_eq!( + std::fs::read_to_string(dir.path().join("keep.txt")).unwrap(), + "orig" + ); + assert!(!dir.path().join("fresh.txt").exists()); + saw_restore |= result.output().contains("restored from snapshot"); + } + assert!( + saw_restore, + "no iteration wrote a file before the failing one" + ); +} diff --git a/crates/tinyagents-harness-tools/src/fs/filesystem/contract_tests.rs b/crates/tinyagents-harness-tools/src/fs/filesystem/contract_tests.rs new file mode 100644 index 000000000..e312a0017 --- /dev/null +++ b/crates/tinyagents-harness-tools/src/fs/filesystem/contract_tests.rs @@ -0,0 +1,152 @@ +//! Wire-contract fixtures: the name, description, permission level, exposure +//! and JSON Schema each filesystem tool advertises to the model. +//! +//! These are what a model is prompted with, and what existing transcripts were +//! produced against, so they are pinned against literal fixture files rather +//! than re-derived from the code under test. A change here is a change to the +//! tool's public contract and must be made on purpose. + +#![allow(clippy::expect_used, clippy::panic, clippy::unwrap_used)] + +use std::sync::Arc; + +use serde_json::{Value, json}; +use tinytools::Tool; + +use super::*; +use crate::fs::filesystem::test_support::{AutonomyLevel, TestGate}; + +fn tools() -> Vec<(&'static str, Box)> { + let gate = TestGate::supervised(std::env::temp_dir()); + let workspace = std::env::temp_dir(); + vec![ + ( + "apply_patch", + Box::new(ApplyPatchTool::new(gate.clone())) as Box, + ), + ("csv_export", Box::new(CsvExportTool::new(gate.clone()))), + ("edit_file", Box::new(EditFileTool::new(gate.clone()))), + ("file_read", Box::new(FileReadTool::new(gate.clone()))), + ("file_write", Box::new(FileWriteTool::new(gate.clone()))), + ( + "git_operations", + Box::new(GitOperationsTool::new(gate.clone(), workspace.clone())), + ), + ("glob_search", Box::new(GlobTool::new(gate.clone()))), + ("grep", Box::new(GrepTool::new(gate.clone()))), + ("image_info", Box::new(ImageInfoTool::new(gate.clone()))), + ("list_files", Box::new(ListFilesTool::new(gate))), + ("read_diff", Box::new(ReadDiffTool::new(workspace.clone()))), + ( + "run_linter", + Box::new(RunLinterTool::new(workspace.clone())), + ), + ("run_tests", Box::new(RunTestsTool::new(workspace.clone()))), + ( + "update_memory_md", + Box::new(UpdateMemoryMdTool::new(workspace.clone())), + ), + ( + "read_workspace_state", + Box::new(WorkspaceStateTool::new(workspace)), + ), + ] +} + +fn describe(tool: &dyn Tool) -> Value { + json!({ + "name": tool.name(), + "description": tool.description(), + "permission_level": format!("{:?}", tool.permission_level()), + "exposure": format!("{:?}", tool.exposure()), + "schema": tool.parameters_schema(), + }) +} + +fn fixture(name: &str) -> Value { + let raw = match name { + "apply_patch" => include_str!("fixtures/apply_patch.json"), + "csv_export" => include_str!("fixtures/csv_export.json"), + "edit_file" => include_str!("fixtures/edit_file.json"), + "file_read" => include_str!("fixtures/file_read.json"), + "file_write" => include_str!("fixtures/file_write.json"), + "git_operations" => include_str!("fixtures/git_operations.json"), + "glob_search" => include_str!("fixtures/glob_search.json"), + "grep" => include_str!("fixtures/grep.json"), + "image_info" => include_str!("fixtures/image_info.json"), + "list_files" => include_str!("fixtures/list_files.json"), + "read_diff" => include_str!("fixtures/read_diff.json"), + "run_linter" => include_str!("fixtures/run_linter.json"), + "run_tests" => include_str!("fixtures/run_tests.json"), + "update_memory_md" => include_str!("fixtures/update_memory_md.json"), + "read_workspace_state" => include_str!("fixtures/read_workspace_state.json"), + other => panic!("no fixture for {other}"), + }; + serde_json::from_str(raw).expect("fixture is valid JSON") +} + +#[test] +fn every_tool_advertises_exactly_its_pinned_contract() { + for (fixture_name, tool) in tools() { + assert_eq!( + describe(tool.as_ref()), + fixture(fixture_name), + "contract drift in `{fixture_name}`" + ); + } +} + +#[test] +fn tool_names_are_unique() { + let mut names: Vec = tools().iter().map(|(_, t)| t.name().to_string()).collect(); + names.sort(); + let before = names.len(); + names.dedup(); + assert_eq!(before, names.len()); +} + +#[test] +fn a_gate_can_be_shared_across_tools() { + let gate: Arc = TestGate::supervised(std::env::temp_dir()); + let _read = FileReadTool::new(Arc::clone(&gate)); + let _write = FileWriteTool::new(gate); +} + +fn gate_at(autonomy: AutonomyLevel) -> Arc { + TestGate::with(std::env::temp_dir(), autonomy, 1_000_000) +} + +#[test] +fn write_tools_route_through_approval_only_when_the_gate_asks() { + let existing = tempfile::tempdir().unwrap(); + std::fs::write(existing.path().join("f.txt"), "x").unwrap(); + let existing_arg = json!({"path": existing.path().join("f.txt").to_string_lossy()}); + let new_arg = json!({"path": existing.path().join("missing.txt").to_string_lossy()}); + + for (autonomy, prompts) in [ + (AutonomyLevel::Supervised, true), + (AutonomyLevel::Full, false), + (AutonomyLevel::ReadOnly, false), + ] { + let gate = gate_at(autonomy); + assert_eq!( + EditFileTool::new(gate.clone()).external_effect_with_args(&json!({})), + prompts + ); + assert_eq!( + ApplyPatchTool::new(gate.clone()).external_effect_with_args(&json!({})), + prompts + ); + let write = FileWriteTool::new(gate.clone()); + // Overwriting prompts, creating does not, an unknown path fails safe. + assert_eq!(write.external_effect_with_args(&existing_arg), prompts); + assert!(!write.external_effect_with_args(&new_arg)); + assert_eq!(write.external_effect_with_args(&json!({})), prompts); + let git = GitOperationsTool::new(gate, std::env::temp_dir()); + assert_eq!( + git.external_effect_with_args(&json!({"operation": "commit"})), + prompts + ); + assert!(!git.external_effect_with_args(&json!({"operation": "status"}))); + } +} diff --git a/crates/tinyagents-harness-tools/src/fs/filesystem/csv_export/extra_tests.rs b/crates/tinyagents-harness-tools/src/fs/filesystem/csv_export/extra_tests.rs new file mode 100644 index 000000000..9c94ce112 --- /dev/null +++ b/crates/tinyagents-harness-tools/src/fs/filesystem/csv_export/extra_tests.rs @@ -0,0 +1,161 @@ +//! Refusal, size-label and failure-path tests for `csv_export`. + +#![allow(clippy::expect_used, clippy::panic, clippy::unwrap_used)] + +use super::*; +use crate::fs::filesystem::test_support::{AutonomyLevel, TestGate, WrapGate}; +use tempfile::TempDir; + +fn tool_in(dir: &TempDir) -> CsvExportTool { + CsvExportTool::new(TestGate::supervised(dir.path().to_path_buf())) +} + +fn data(v: &serde_json::Value) -> String { + serde_json::to_string(v).unwrap() +} + +#[tokio::test] +async fn read_only_autonomy_is_refused() { + let dir = TempDir::new().unwrap(); + let gate = TestGate::with(dir.path().to_path_buf(), AutonomyLevel::ReadOnly, 100); + let r = CsvExportTool::new(gate) + .execute(json!({"data": data(&json!([{"a": 1}])), "filename": "x.csv"})) + .await + .unwrap(); + assert!(r.is_error); + assert!(r.output().contains("autonomy is read-only")); +} + +#[tokio::test] +async fn exhausted_hourly_limit_is_refused() { + let dir = TempDir::new().unwrap(); + let gate = TestGate::with(dir.path().to_path_buf(), AutonomyLevel::Supervised, 0); + let r = CsvExportTool::new(gate) + .execute(json!({"data": data(&json!([{"a": 1}])), "filename": "x.csv"})) + .await + .unwrap(); + assert!(r.output().contains("too many actions in the last hour")); +} + +#[tokio::test] +async fn exhausted_action_budget_is_refused_before_writing() { + let dir = TempDir::new().unwrap(); + let gate = WrapGate::new(TestGate::supervised(dir.path().to_path_buf())) + .no_budget() + .arc(); + let r = CsvExportTool::new(gate) + .execute(json!({"data": data(&json!([{"a": 1}])), "filename": "x.csv"})) + .await + .unwrap(); + assert!(r.output().contains("action budget exhausted")); + assert!(!dir.path().join("exports/x.csv").exists()); +} + +#[tokio::test] +async fn invalid_json_is_reported() { + let dir = TempDir::new().unwrap(); + let r = tool_in(&dir) + .execute(json!({"data": "{not json", "filename": "x.csv"})) + .await + .unwrap(); + assert!(r.is_error); + assert!(r.output().contains("Failed to parse data as JSON")); +} + +#[tokio::test] +async fn empty_array_is_reported() { + let dir = TempDir::new().unwrap(); + let r = tool_in(&dir) + .execute(json!({"data": "[]", "filename": "x.csv"})) + .await + .unwrap(); + assert!(r.output().contains("nothing to export")); +} + +#[tokio::test] +async fn a_path_the_gate_refuses_is_reported() { + let dir = TempDir::new().unwrap(); + let r = tool_in(&dir) + .execute(json!({"data": data(&json!([{"a": 1}])), "filename": "../escape.csv"})) + .await + .unwrap(); + assert!(r.is_error); + assert!(r.output().contains("Path not allowed")); +} + +#[cfg(unix)] +#[tokio::test] +async fn a_symlinked_target_is_refused() { + let dir = TempDir::new().unwrap(); + std::fs::create_dir(dir.path().join("exports")).unwrap(); + std::fs::write(dir.path().join("real.txt"), "keep").unwrap(); + std::os::unix::fs::symlink( + dir.path().join("real.txt"), + dir.path().join("exports/link.csv"), + ) + .unwrap(); + let r = tool_in(&dir) + .execute(json!({"data": data(&json!([{"a": 1}])), "filename": "link.csv"})) + .await + .unwrap(); + assert!(r.is_error); + assert!(r.output().contains("Refusing to write through symlink")); + assert_eq!( + std::fs::read_to_string(dir.path().join("real.txt")).unwrap(), + "keep" + ); +} + +#[tokio::test] +async fn a_directory_in_the_way_reports_the_write_failure() { + let dir = TempDir::new().unwrap(); + std::fs::create_dir_all(dir.path().join("exports/out.csv")).unwrap(); + let r = tool_in(&dir) + .execute(json!({"data": data(&json!([{"a": 1}])), "filename": "out.csv"})) + .await + .unwrap(); + assert!(r.is_error); + assert!(r.output().contains("Failed to write CSV file")); +} + +#[tokio::test] +async fn size_labels_scale_to_kb_and_mb() { + let dir = TempDir::new().unwrap(); + let kb_rows: Vec<_> = (0..200) + .map(|i| json!({"n": format!("row-{i:04}-padding")})) + .collect(); + let r = tool_in(&dir) + .execute(json!({"data": data(&json!(kb_rows)), "filename": "kb.csv"})) + .await + .unwrap(); + assert!(r.output().contains(" KB)"), "{}", r.output()); + + let big = "x".repeat(1024 * 1024 + 10); + let r = tool_in(&dir) + .execute(json!({"data": data(&json!([{"blob": big}])), "filename": "mb.csv"})) + .await + .unwrap(); + assert!(r.output().contains(" MB)"), "{}", r.output()); +} + +#[test] +fn cells_and_columns_cover_bools_and_non_object_rows() { + assert_eq!(value_to_cell(&json!(true)), "true"); + assert_eq!(value_to_cell(&json!(null)), ""); + assert_eq!(value_to_cell(&json!([1, 2])), "[1,2]"); + assert_eq!(resolve_columns(&[json!(1), json!(2)], None).len(), 0); +} + +#[tokio::test] +async fn execute_with_context_delegates() { + let dir = TempDir::new().unwrap(); + let r = tool_in(&dir) + .execute_with_context( + json!({"data": data(&json!([{"ok": true}])), "filename": "c.csv"}), + ToolCallOptions::default(), + None, + ) + .await + .unwrap(); + assert!(!r.is_error, "{}", r.output()); +} diff --git a/crates/tinyagents-harness-tools/src/fs/filesystem/csv_export/mod.rs b/crates/tinyagents-harness-tools/src/fs/filesystem/csv_export/mod.rs new file mode 100644 index 000000000..da2a648d0 --- /dev/null +++ b/crates/tinyagents-harness-tools/src/fs/filesystem/csv_export/mod.rs @@ -0,0 +1,270 @@ +use super::gate::{FsGate, gate_for_context}; +use async_trait::async_trait; +use serde_json::json; +use std::sync::Arc; +use tinytools::ToolRunContext; +use tinytools::{Tool, ToolCallOptions, ToolResult}; + +/// Export structured data (JSON array of objects) as a CSV file to the workspace. +#[derive(Debug)] +pub struct CsvExportTool { + gate: Arc, +} + +impl CsvExportTool { + /// Create the tool. + #[must_use] + pub fn new(gate: Arc) -> Self { + Self { gate } + } +} + +/// Escape a value for inclusion in a CSV cell. Wraps the value in +/// double-quotes when it contains commas, quotes, or newlines. Embedded +/// double-quotes are escaped by doubling them per RFC 4180. +fn csv_escape(value: &str) -> String { + if value.contains(',') || value.contains('"') || value.contains('\n') || value.contains('\r') { + let escaped = value.replace('"', "\"\""); + format!("\"{escaped}\"") + } else { + value.to_string() + } +} + +/// Convert a `serde_json::Value` into a plain string suitable for a CSV +/// cell. Objects and arrays are serialised as compact JSON; booleans and +/// numbers use their natural representation; nulls become the empty +/// string. +fn value_to_cell(v: &serde_json::Value) -> String { + match v { + serde_json::Value::Null => String::new(), + serde_json::Value::String(s) => s.clone(), + serde_json::Value::Bool(b) => b.to_string(), + serde_json::Value::Number(n) => n.to_string(), + // Nested objects/arrays → compact JSON string + other => other.to_string(), + } +} + +/// Collect column headers from a JSON array. If `columns` is provided, +/// use those in order. Otherwise, collect all keys from the first object +/// in the array (sorted alphabetically — `serde_json` uses `BTreeMap` by +/// default). Callers who need a specific column order should pass the +/// `columns` parameter. +fn resolve_columns(items: &[serde_json::Value], columns: Option<&[String]>) -> Vec { + if let Some(cols) = columns { + return cols.to_vec(); + } + // Collect keys from the first object. + if let Some(first) = items.first() + && let Some(obj) = first.as_object() + { + return obj.keys().cloned().collect(); + } + Vec::new() +} + +/// Render a JSON array of objects into a CSV string. +fn render_csv(items: &[serde_json::Value], columns: &[String]) -> String { + let mut buf = String::new(); + + // Header row + let header: Vec = columns.iter().map(|c| csv_escape(c)).collect(); + buf.push_str(&header.join(",")); + buf.push('\n'); + + // Data rows + for item in items { + let row: Vec = columns + .iter() + .map(|col| { + let cell_value = item.get(col).map(value_to_cell).unwrap_or_default(); + csv_escape(&cell_value) + }) + .collect(); + buf.push_str(&row.join(",")); + buf.push('\n'); + } + + buf +} + +#[async_trait] +impl Tool for CsvExportTool { + fn name(&self) -> &'static str { + "csv_export" + } + + fn description(&self) -> &'static str { + "Export structured data (JSON array of objects) as a CSV file to the workspace. \ + Returns the file path. Use when the user wants raw tabular data from a tool \ + result that's too large to include inline." + } + + fn parameters_schema(&self) -> serde_json::Value { + json!({ + "type": "object", + "properties": { + "data": { + "type": "string", + "description": "JSON string containing an array of objects to export. Each object becomes a row; keys become column headers." + }, + "filename": { + "type": "string", + "description": "Output filename (without path). Will be written to workspace/exports/. Example: 'github-issues-2026-04-16.csv'" + }, + "columns": { + "type": "array", + "items": { "type": "string" }, + "description": "Optional ordered list of column names to include. If omitted, all keys from the first object are used as headers." + } + }, + "required": ["data", "filename"] + }) + } + + fn permission_level(&self) -> tinytools::PermissionLevel { + tinytools::PermissionLevel::Write + } + + async fn execute(&self, args: serde_json::Value) -> anyhow::Result { + self.execute_in_context(args, None).await + } + + async fn execute_with_context( + &self, + args: serde_json::Value, + _options: ToolCallOptions, + context: Option<&dyn ToolRunContext>, + ) -> anyhow::Result { + self.execute_in_context(args, context).await + } +} + +impl CsvExportTool { + async fn execute_in_context( + &self, + args: serde_json::Value, + context: Option<&dyn ToolRunContext>, + ) -> anyhow::Result { + let data_str = args + .get("data") + .and_then(|v| v.as_str()) + .ok_or_else(|| anyhow::anyhow!("Missing 'data' parameter"))?; + + let filename = args + .get("filename") + .and_then(|v| v.as_str()) + .ok_or_else(|| anyhow::anyhow!("Missing 'filename' parameter"))?; + + let columns: Option> = args.get("columns").and_then(|v| { + v.as_array().map(|arr| { + arr.iter() + .filter_map(|item| item.as_str().map(String::from)) + .collect() + }) + }); + + // Security: check write permission + if !self.gate.can_act() { + return Ok(ToolResult::error( + "[policy-blocked] Action blocked: autonomy is read-only", + )); + } + + if self.gate.is_rate_limited() { + return Ok(ToolResult::error( + "Rate limit exceeded: too many actions in the last hour", + )); + } + + // Parse the JSON data + let parsed: serde_json::Value = match serde_json::from_str(data_str) { + Ok(v) => v, + Err(e) => { + return Ok(ToolResult::error(format!( + "Failed to parse data as JSON: {e}" + ))); + } + }; + + let items = match parsed.as_array() { + Some(arr) => arr, + None => { + return Ok(ToolResult::error( + "Data must be a JSON array of objects, but got a non-array value", + )); + } + }; + + if items.is_empty() { + return Ok(ToolResult::error("Data array is empty — nothing to export")); + } + + // Resolve columns and render CSV + let cols = resolve_columns(items, columns.as_deref()); + let csv_content = render_csv(items, &cols); + let csv_bytes = csv_content.len(); + + // Validate the relative path + let relative_path = format!("exports/{filename}"); + + let path_policy = gate_for_context(&self.gate, context, "csv_export"); + + // Security check first: validate path string, resolve symlinks, confirm workspace + // containment. validate_parent_path walks up to the deepest existing ancestor so + // it does not require the exports/ directory to exist yet. + let resolved_target = match path_policy.validate_parent_path(&relative_path).await { + Ok(p) => p, + Err(msg) => return Ok(ToolResult::error(msg)), + }; + + // Create exports/ directory only at the validated, resolved location. + if let Some(resolved_parent) = resolved_target.parent() { + tokio::fs::create_dir_all(resolved_parent).await?; + } + + // If the target already exists and is a symlink, refuse to follow it + if let Ok(meta) = tokio::fs::symlink_metadata(&resolved_target).await + && meta.file_type().is_symlink() + { + return Ok(ToolResult::error(format!( + "Refusing to write through symlink: {}", + resolved_target.display() + ))); + } + + if !self.gate.record_action() { + return Ok(ToolResult::error( + "Rate limit exceeded: action budget exhausted", + )); + } + + // Write the CSV file + match tokio::fs::write(&resolved_target, &csv_content).await { + Ok(()) => { + let size_display = if csv_bytes >= 1024 * 1024 { + format!("{:.1} MB", csv_bytes as f64 / (1024.0 * 1024.0)) + } else if csv_bytes >= 1024 { + format!("{:.1} KB", csv_bytes as f64 / 1024.0) + } else { + format!("{csv_bytes} bytes") + }; + + Ok(ToolResult::success(format!( + "Exported {} rows to {relative_path} ({size_display})", + items.len() + ))) + } + Err(e) => Ok(ToolResult::error(format!("Failed to write CSV file: {e}"))), + } + } +} + +#[cfg(test)] +#[path = "mod_tests.rs"] +mod tests; + +#[cfg(test)] +#[path = "extra_tests.rs"] +mod extra_tests; diff --git a/crates/tinyagents-harness-tools/src/fs/filesystem/csv_export/mod_tests.rs b/crates/tinyagents-harness-tools/src/fs/filesystem/csv_export/mod_tests.rs new file mode 100644 index 000000000..ef818bada --- /dev/null +++ b/crates/tinyagents-harness-tools/src/fs/filesystem/csv_export/mod_tests.rs @@ -0,0 +1,234 @@ +//! Behavior tests for the `csv_export` tool, driven through a fake [`FsGate`]. + +#![allow(clippy::expect_used, clippy::panic, clippy::unwrap_used)] + +use super::*; +use crate::fs::filesystem::test_support::TestGate; + +fn test_security(workspace: std::path::PathBuf) -> Arc { + TestGate::supervised(workspace) +} + +#[test] +fn csv_export_name() { + let tool = CsvExportTool::new(test_security(std::env::temp_dir())); + assert_eq!(tool.name(), "csv_export"); +} + +#[test] +fn csv_export_schema_has_required_fields() { + let tool = CsvExportTool::new(test_security(std::env::temp_dir())); + let schema = tool.parameters_schema(); + assert!(schema["properties"]["data"].is_object()); + assert!(schema["properties"]["filename"].is_object()); + assert!(schema["properties"]["columns"].is_object()); + let required = schema["required"].as_array().unwrap(); + assert!(required.contains(&json!("data"))); + assert!(required.contains(&json!("filename"))); +} + +#[tokio::test] +async fn csv_export_formats_simple_array() { + let dir = std::env::temp_dir().join("openhuman_test_csv_export_simple"); + let _ = tokio::fs::remove_dir_all(&dir).await; + tokio::fs::create_dir_all(&dir).await.unwrap(); + + let tool = CsvExportTool::new(test_security(dir.clone())); + let data = serde_json::to_string(&json!([ + {"name": "Alice", "age": 30, "city": "NYC"}, + {"name": "Bob", "age": 25, "city": "LA"}, + {"name": "Carol", "age": 35, "city": "Chicago"} + ])) + .unwrap(); + + let result = tool + .execute(json!({ + "data": data, + "filename": "people.csv" + })) + .await + .unwrap(); + + assert!(!result.is_error, "unexpected error: {}", result.output()); + assert!(result.output().contains("3 rows")); + + let content = tokio::fs::read_to_string(dir.join("exports/people.csv")) + .await + .unwrap(); + let lines: Vec<&str> = content.trim().lines().collect(); + assert_eq!(lines.len(), 4, "header + 3 data rows"); + + // Header should contain the keys from the first object + let header = lines[0]; + assert!(header.contains("name")); + assert!(header.contains("age")); + assert!(header.contains("city")); + + // Data rows should contain values + assert!(lines[1].contains("Alice")); + assert!(lines[2].contains("Bob")); + assert!(lines[3].contains("Carol")); + + let _ = tokio::fs::remove_dir_all(&dir).await; +} + +#[tokio::test] +async fn csv_export_handles_missing_keys() { + let dir = std::env::temp_dir().join("openhuman_test_csv_export_missing_keys"); + let _ = tokio::fs::remove_dir_all(&dir).await; + tokio::fs::create_dir_all(&dir).await.unwrap(); + + let tool = CsvExportTool::new(test_security(dir.clone())); + let data = serde_json::to_string(&json!([ + {"name": "Alice", "age": 30, "city": "NYC"}, + {"name": "Bob"}, + {"name": "Carol", "city": "Chicago"} + ])) + .unwrap(); + + let result = tool + .execute(json!({ + "data": data, + "filename": "sparse.csv", + "columns": ["name", "age", "city"] + })) + .await + .unwrap(); + + assert!(!result.is_error, "unexpected error: {}", result.output()); + + let content = tokio::fs::read_to_string(dir.join("exports/sparse.csv")) + .await + .unwrap(); + let lines: Vec<&str> = content.trim().lines().collect(); + assert_eq!(lines.len(), 4); + + // Bob's row should have empty cells for age and city + let bob_row = lines[2]; + let bob_cells: Vec<&str> = bob_row.split(',').collect(); + assert_eq!(bob_cells.len(), 3, "Bob row should have 3 cells"); + assert_eq!(bob_cells[0], "Bob"); + assert_eq!(bob_cells[1], "", "missing age should be empty"); + assert_eq!(bob_cells[2], "", "missing city should be empty"); + + let _ = tokio::fs::remove_dir_all(&dir).await; +} + +#[tokio::test] +async fn csv_export_respects_column_order() { + let dir = std::env::temp_dir().join("openhuman_test_csv_export_column_order"); + let _ = tokio::fs::remove_dir_all(&dir).await; + tokio::fs::create_dir_all(&dir).await.unwrap(); + + let tool = CsvExportTool::new(test_security(dir.clone())); + let data = serde_json::to_string(&json!([ + {"name": "Alice", "age": 30, "city": "NYC"}, + {"name": "Bob", "age": 25, "city": "LA"} + ])) + .unwrap(); + + let result = tool + .execute(json!({ + "data": data, + "filename": "ordered.csv", + "columns": ["city", "name", "age"] + })) + .await + .unwrap(); + + assert!(!result.is_error, "unexpected error: {}", result.output()); + + let content = tokio::fs::read_to_string(dir.join("exports/ordered.csv")) + .await + .unwrap(); + let lines: Vec<&str> = content.trim().lines().collect(); + assert_eq!( + lines[0], "city,name,age", + "header must follow requested column order" + ); + assert_eq!(lines[1], "NYC,Alice,30"); + assert_eq!(lines[2], "LA,Bob,25"); + + let _ = tokio::fs::remove_dir_all(&dir).await; +} + +#[tokio::test] +async fn csv_export_rejects_non_array_input() { + let dir = std::env::temp_dir().join("openhuman_test_csv_export_non_array"); + let _ = tokio::fs::remove_dir_all(&dir).await; + tokio::fs::create_dir_all(&dir).await.unwrap(); + + let tool = CsvExportTool::new(test_security(dir.clone())); + let data = serde_json::to_string(&json!({"not": "an array"})).unwrap(); + + let result = tool + .execute(json!({ + "data": data, + "filename": "bad.csv" + })) + .await + .unwrap(); + + assert!(result.is_error); + assert!( + result.output().contains("non-array"), + "error should mention non-array, got: {}", + result.output() + ); + + let _ = tokio::fs::remove_dir_all(&dir).await; +} + +#[tokio::test] +async fn csv_export_handles_nested_values() { + let dir = std::env::temp_dir().join("openhuman_test_csv_export_nested"); + let _ = tokio::fs::remove_dir_all(&dir).await; + tokio::fs::create_dir_all(&dir).await.unwrap(); + + let tool = CsvExportTool::new(test_security(dir.clone())); + let data = serde_json::to_string(&json!([ + { + "name": "Alice", + "tags": ["admin", "dev"], + "meta": {"role": "lead", "level": 5} + }, + { + "name": "Bob", + "tags": [], + "meta": null + } + ])) + .unwrap(); + + let result = tool + .execute(json!({ + "data": data, + "filename": "nested.csv", + "columns": ["name", "tags", "meta"] + })) + .await + .unwrap(); + + assert!(!result.is_error, "unexpected error: {}", result.output()); + + let content = tokio::fs::read_to_string(dir.join("exports/nested.csv")) + .await + .unwrap(); + let lines: Vec<&str> = content.trim().lines().collect(); + assert_eq!(lines.len(), 3, "header + 2 data rows"); + + // Alice's tags should be serialized as a JSON string (in quotes because it contains commas) + let alice_row = lines[1]; + assert!(alice_row.contains("Alice")); + // The JSON array should be serialized as a string and quoted + assert!( + alice_row.contains(r#"[""admin"",""dev""]"#), + "nested arrays should be JSON-serialized in CSV: {alice_row}" + ); + + // Bob's meta is null → empty cell + let bob_row = lines[2]; + assert!(bob_row.contains("Bob")); + + let _ = tokio::fs::remove_dir_all(&dir).await; +} diff --git a/crates/tinyagents-harness-tools/src/fs/filesystem/edit_file/mod.rs b/crates/tinyagents-harness-tools/src/fs/filesystem/edit_file/mod.rs new file mode 100644 index 000000000..2336c989d --- /dev/null +++ b/crates/tinyagents-harness-tools/src/fs/filesystem/edit_file/mod.rs @@ -0,0 +1,242 @@ +//! `edit` — string-replace edit on a single file. +//! +//! Coding-harness baseline tool (issue #1205). Models the +//! Anthropic/Claude-Code `Edit` semantics: exact-match `old_string` → +//! `new_string` substitution. By default, `old_string` MUST match +//! exactly once in the file (so the model can't accidentally edit +//! every match). Set `replace_all` to override. + +use super::gate::{FsGate, gate_for_context}; +use crate::fs::file_state; +use async_trait::async_trait; +use serde_json::json; +use std::sync::Arc; +use tinytools::ToolRunContext; +use tinytools::{PermissionLevel, Tool, ToolCallOptions, ToolResult}; + +const MAX_FILE_BYTES: u64 = 5 * 1024 * 1024; + +/// Replace an exact string in a workspace file. +#[derive(Debug)] +pub struct EditFileTool { + gate: Arc, + sink: Arc, +} + +impl EditFileTool { + /// Create the tool. + #[must_use] + pub fn new(gate: Arc) -> Self { + Self { + gate, + sink: super::file_sink::os_sink(), + } + } + + /// Sends this tool's writes somewhere other than the OS. + #[cfg(test)] + pub(super) fn with_sink(mut self, sink: Arc) -> Self { + self.sink = sink; + self + } +} + +#[async_trait] +impl Tool for EditFileTool { + fn exposure(&self) -> tinytools::ToolExposure { + tinytools::ToolExposure::Deferred + } + + fn name(&self) -> &'static str { + "edit" + } + + fn description(&self) -> &'static str { + "Edit a file by exact string replacement. By default `old_string` must \ + match exactly once. Set `replace_all` to true to replace every match." + } + + fn parameters_schema(&self) -> serde_json::Value { + json!({ + "type": "object", + "properties": { + "path": { "type": "string", "description": "Workspace-relative file path." }, + "old_string": { "type": "string", "description": "Text to find. Must match exactly." }, + "new_string": { "type": "string", "description": "Replacement text." }, + "replace_all": { + "type": "boolean", + "description": "If true, replace every occurrence (default false).", + "default": false + } + }, + "required": ["path", "old_string", "new_string"] + }) + } + + fn permission_level(&self) -> PermissionLevel { + PermissionLevel::Write + } + + /// `edit` always modifies an **existing** file → in ask-before-edit it + /// routes through the human approval gate; in Full it runs; read-only is + /// blocked in `execute`. + fn external_effect_with_args(&self, _args: &serde_json::Value) -> bool { + self.gate.write_needs_approval() + } + + async fn execute(&self, args: serde_json::Value) -> anyhow::Result { + self.execute_in_context(args, None).await + } + + async fn execute_with_context( + &self, + args: serde_json::Value, + _options: ToolCallOptions, + context: Option<&dyn ToolRunContext>, + ) -> anyhow::Result { + self.execute_in_context(args, context).await + } +} + +impl EditFileTool { + async fn execute_in_context( + &self, + args: serde_json::Value, + context: Option<&dyn ToolRunContext>, + ) -> anyhow::Result { + let path = args + .get("path") + .and_then(|v| v.as_str()) + .ok_or_else(|| anyhow::anyhow!("Missing 'path' parameter"))?; + let old_string = args + .get("old_string") + .and_then(|v| v.as_str()) + .ok_or_else(|| anyhow::anyhow!("Missing 'old_string' parameter"))?; + let new_string = args + .get("new_string") + .and_then(|v| v.as_str()) + .ok_or_else(|| anyhow::anyhow!("Missing 'new_string' parameter"))?; + let replace_all = args + .get("replace_all") + .and_then(serde_json::Value::as_bool) + .unwrap_or(false); + + if old_string.is_empty() { + return Ok(ToolResult::error("`old_string` must not be empty")); + } + if old_string == new_string { + return Ok(ToolResult::error( + "`old_string` and `new_string` are identical — nothing to do", + )); + } + + if !self.gate.can_act() { + return Ok(ToolResult::error( + "[policy-blocked] Action blocked: autonomy is read-only", + )); + } + if self.gate.is_rate_limited() { + return Ok(ToolResult::error( + "Rate limit exceeded: too many actions in the last hour", + )); + } + if !self.gate.record_action() { + return Ok(ToolResult::error( + "Rate limit exceeded: action budget exhausted", + )); + } + + let path_policy = gate_for_context(&self.gate, context, "edit"); + let full = path_policy.action_dir().join(path); + + // Symlink check must happen on the *unresolved* path — + // `canonicalize` resolves symlinks, so checking after that point + // would always see the link's final target. + if let Ok(meta) = tokio::fs::symlink_metadata(&full).await + && meta.file_type().is_symlink() + { + return Ok(ToolResult::error(format!( + "Refusing to edit through symlink: {}", + full.display() + ))); + } + + // Security check: validate path string, resolve symlinks, confirm workspace containment. + let resolved = match path_policy.validate_path(path).await { + Ok(p) => p, + Err(msg) => return Ok(ToolResult::error(msg)), + }; + + if let Ok(meta) = tokio::fs::metadata(&resolved).await + && meta.len() > MAX_FILE_BYTES + { + return Ok(ToolResult::error(format!( + "File too large: {} bytes (limit: {MAX_FILE_BYTES} bytes)", + meta.len() + ))); + } + + // Acquire per-path lock for the read-modify-write section. + let _path_guard = file_state::acquire_path_lock(&resolved).await; + + // File-state guard: reject edits based on stale or partial reads. + if let Some(agent_id) = file_state::current_file_state_agent_id() { + if let Some(msg) = file_state::check_stale_read(&agent_id, &resolved) { + tracing::debug!( + agent = %agent_id, + path = %resolved.display(), + "[file_state] edit blocked: stale read" + ); + return Ok(ToolResult::error(msg)); + } + if let Some(msg) = file_state::check_partial_read(&agent_id, &resolved) { + tracing::debug!( + agent = %agent_id, + path = %resolved.display(), + "[file_state] edit blocked: partial read" + ); + return Ok(ToolResult::error(msg)); + } + } + + let contents = match tokio::fs::read_to_string(&resolved).await { + Ok(c) => c, + Err(e) => return Ok(ToolResult::error(format!("Failed to read file: {e}"))), + }; + + let count = contents.matches(old_string).count(); + if count == 0 { + return Ok(ToolResult::error(format!( + "`old_string` not found in {path}" + ))); + } + if count > 1 && !replace_all { + return Ok(ToolResult::error(format!( + "`old_string` matches {count} times in {path}; pass `replace_all: true` or \ + expand `old_string` so it is unique" + ))); + } + + let updated = if replace_all { + contents.replace(old_string, new_string) + } else { + contents.replacen(old_string, new_string, 1) + }; + + match self.sink.write(&resolved, updated.as_bytes()).await { + Ok(()) => { + if let Some(agent_id) = file_state::current_file_state_agent_id() { + file_state::record_write(&agent_id, resolved); + } + Ok(ToolResult::success(format!( + "Edited {path}: {count} replacement(s)" + ))) + } + Err(e) => Ok(ToolResult::error(format!("Failed to write file: {e}"))), + } + } +} + +#[cfg(test)] +#[path = "mod_tests.rs"] +mod tests; diff --git a/crates/tinyagents-harness-tools/src/fs/filesystem/edit_file/mod_tests.rs b/crates/tinyagents-harness-tools/src/fs/filesystem/edit_file/mod_tests.rs new file mode 100644 index 000000000..8afffed3a --- /dev/null +++ b/crates/tinyagents-harness-tools/src/fs/filesystem/edit_file/mod_tests.rs @@ -0,0 +1,316 @@ +//! Behavior tests for the `edit_file` tool, driven through a fake [`FsGate`]. + +#![allow(clippy::expect_used, clippy::panic, clippy::unwrap_used)] + +use super::*; +use crate::fs::filesystem::test_support::{AutonomyLevel, TestGate}; + +fn test_security(workspace: std::path::PathBuf) -> Arc { + TestGate::supervised(workspace) +} + +fn test_security_readonly(workspace: std::path::PathBuf) -> Arc { + TestGate::with(workspace, AutonomyLevel::ReadOnly, 1_000_000) +} + +#[test] +fn edit_name() { + let tool = EditFileTool::new(test_security(std::env::temp_dir())); + assert_eq!(tool.name(), "edit"); +} + +#[tokio::test] +async fn edit_replaces_unique_match() { + let dir = std::env::temp_dir().join("openhuman_test_edit_unique"); + let _ = tokio::fs::remove_dir_all(&dir).await; + tokio::fs::create_dir_all(&dir).await.unwrap(); + tokio::fs::write(dir.join("f.txt"), "alpha bravo") + .await + .unwrap(); + + let tool = EditFileTool::new(test_security(dir.clone())); + let result = tool + .execute(json!({"path": "f.txt", "old_string": "bravo", "new_string": "charlie"})) + .await + .unwrap(); + assert!(!result.is_error, "{}", result.output()); + let updated = tokio::fs::read_to_string(dir.join("f.txt")).await.unwrap(); + assert_eq!(updated, "alpha charlie"); + + let _ = tokio::fs::remove_dir_all(&dir).await; +} + +#[tokio::test] +async fn edit_rejects_ambiguous_match() { + let dir = std::env::temp_dir().join("openhuman_test_edit_ambig"); + let _ = tokio::fs::remove_dir_all(&dir).await; + tokio::fs::create_dir_all(&dir).await.unwrap(); + tokio::fs::write(dir.join("f.txt"), "x x x").await.unwrap(); + + let tool = EditFileTool::new(test_security(dir.clone())); + let result = tool + .execute(json!({"path": "f.txt", "old_string": "x", "new_string": "y"})) + .await + .unwrap(); + assert!(result.is_error); + assert!(result.output().contains("matches 3 times")); + + let _ = tokio::fs::remove_dir_all(&dir).await; +} + +#[tokio::test] +async fn edit_replace_all() { + let dir = std::env::temp_dir().join("openhuman_test_edit_all"); + let _ = tokio::fs::remove_dir_all(&dir).await; + tokio::fs::create_dir_all(&dir).await.unwrap(); + tokio::fs::write(dir.join("f.txt"), "x x x").await.unwrap(); + + let tool = EditFileTool::new(test_security(dir.clone())); + let result = tool + .execute( + json!({"path": "f.txt", "old_string": "x", "new_string": "y", "replace_all": true}), + ) + .await + .unwrap(); + assert!(!result.is_error); + let updated = tokio::fs::read_to_string(dir.join("f.txt")).await.unwrap(); + assert_eq!(updated, "y y y"); + + let _ = tokio::fs::remove_dir_all(&dir).await; +} + +#[tokio::test] +async fn edit_no_match() { + let dir = std::env::temp_dir().join("openhuman_test_edit_nomatch"); + let _ = tokio::fs::remove_dir_all(&dir).await; + tokio::fs::create_dir_all(&dir).await.unwrap(); + tokio::fs::write(dir.join("f.txt"), "alpha").await.unwrap(); + + let tool = EditFileTool::new(test_security(dir.clone())); + let result = tool + .execute(json!({"path": "f.txt", "old_string": "zulu", "new_string": "x"})) + .await + .unwrap(); + assert!(result.is_error); + assert!(result.output().contains("not found")); + + let _ = tokio::fs::remove_dir_all(&dir).await; +} + +#[tokio::test] +async fn edit_blocks_readonly_mode() { + let dir = std::env::temp_dir().join("openhuman_test_edit_ro"); + let _ = tokio::fs::remove_dir_all(&dir).await; + tokio::fs::create_dir_all(&dir).await.unwrap(); + tokio::fs::write(dir.join("f.txt"), "abc").await.unwrap(); + + let tool = EditFileTool::new(test_security_readonly(dir.clone())); + let result = tool + .execute(json!({"path": "f.txt", "old_string": "abc", "new_string": "xyz"})) + .await + .unwrap(); + assert!(result.is_error); + assert!(result.output().contains("read-only")); + + let _ = tokio::fs::remove_dir_all(&dir).await; +} + +#[tokio::test] +async fn edit_rejects_empty_old_string() { + let dir = std::env::temp_dir().join("openhuman_test_edit_empty_old"); + let _ = tokio::fs::remove_dir_all(&dir).await; + tokio::fs::create_dir_all(&dir).await.unwrap(); + tokio::fs::write(dir.join("f.txt"), "abc").await.unwrap(); + + let tool = EditFileTool::new(test_security(dir.clone())); + let result = tool + .execute(json!({"path": "f.txt", "old_string": "", "new_string": "x"})) + .await + .unwrap(); + assert!(result.is_error); + + let _ = tokio::fs::remove_dir_all(&dir).await; +} + +#[tokio::test] +async fn edit_rejects_identical_strings() { + let dir = std::env::temp_dir().join("openhuman_test_edit_same"); + let _ = tokio::fs::remove_dir_all(&dir).await; + tokio::fs::create_dir_all(&dir).await.unwrap(); + tokio::fs::write(dir.join("f.txt"), "abc").await.unwrap(); + + let tool = EditFileTool::new(test_security(dir.clone())); + let result = tool + .execute(json!({"path": "f.txt", "old_string": "abc", "new_string": "abc"})) + .await + .unwrap(); + assert!(result.is_error); + assert!(result.output().contains("identical")); + + let _ = tokio::fs::remove_dir_all(&dir).await; +} + +#[tokio::test] +async fn edit_reports_an_os_write_failure_rather_than_a_silent_success() { + let dir = std::env::temp_dir().join("openhuman_test_edit_os_write_fail"); + let _ = tokio::fs::remove_dir_all(&dir).await; + tokio::fs::create_dir_all(&dir).await.unwrap(); + let file = dir.join("f.txt"); + tokio::fs::write(&file, "abc").await.unwrap(); + + let tool = EditFileTool::new(test_security(dir.clone())).with_sink(Arc::new( + crate::fs::filesystem::file_sink::RefusingSink(std::io::ErrorKind::PermissionDenied), + )); + let result = tool + .execute(json!({"path": "f.txt", "old_string": "abc", "new_string": "xyz"})) + .await + .unwrap(); + + assert!( + result.is_error, + "a refused write must surface as an error, not a fabricated success" + ); + assert!(result.output().contains("Failed to write file")); + assert_eq!( + tokio::fs::read_to_string(&file).await.unwrap(), + "abc", + "the file must be left as it was when the write did not happen" + ); + + let _ = tokio::fs::remove_dir_all(&dir).await; +} + +// ── Coverage of the budget, symlink, size, read and guard branches ────── + +use crate::fs::filesystem::test_support::{RacyGate, WorkspaceContext}; + +fn edit_args(path: &str) -> serde_json::Value { + json!({"path": path, "old_string": "aaa", "new_string": "bbb"}) +} + +#[tokio::test] +async fn edit_reports_rate_limits() { + let dir = tempfile::tempdir().unwrap(); + std::fs::write(dir.path().join("f.txt"), "aaa").unwrap(); + + let limited = EditFileTool::new(TestGate::with( + dir.path().to_path_buf(), + AutonomyLevel::Supervised, + 0, + )); + let result = limited.execute(edit_args("f.txt")).await.unwrap(); + assert!( + result + .output() + .contains("too many actions in the last hour") + ); + + let racy = EditFileTool::new(Arc::new(RacyGate(test_security(dir.path().to_path_buf())))); + let result = racy.execute(edit_args("f.txt")).await.unwrap(); + assert!(result.output().contains("action budget exhausted")); +} + +#[cfg(unix)] +#[tokio::test] +async fn edit_refuses_to_edit_through_a_symlink() { + let dir = tempfile::tempdir().unwrap(); + std::fs::write(dir.path().join("real.txt"), "aaa").unwrap(); + std::os::unix::fs::symlink(dir.path().join("real.txt"), dir.path().join("link.txt")).unwrap(); + let tool = EditFileTool::new(test_security(dir.path().to_path_buf())); + let result = tool.execute(edit_args("link.txt")).await.unwrap(); + assert!(result.is_error); + assert!(result.output().contains("Refusing to edit through symlink")); +} + +#[tokio::test] +async fn edit_reports_a_missing_file_and_unreadable_contents() { + let dir = tempfile::tempdir().unwrap(); + let tool = EditFileTool::new(test_security(dir.path().to_path_buf())); + let result = tool.execute(edit_args("absent.txt")).await.unwrap(); + assert!(result.is_error); + assert!(result.output().contains("Failed to resolve")); + + std::fs::write(dir.path().join("bin.dat"), [0xff_u8, 0xfe, 0xfd]).unwrap(); + let result = tool.execute(edit_args("bin.dat")).await.unwrap(); + assert!(result.is_error); + assert!(result.output().contains("Failed to read file")); +} + +#[tokio::test] +async fn edit_refuses_an_oversized_file() { + let dir = tempfile::tempdir().unwrap(); + let file = std::fs::File::create(dir.path().join("big.txt")).unwrap(); + file.set_len(MAX_FILE_BYTES + 1).unwrap(); + let tool = EditFileTool::new(test_security(dir.path().to_path_buf())); + let result = tool.execute(edit_args("big.txt")).await.unwrap(); + assert!(result.is_error); + assert!(result.output().contains("File too large")); +} + +#[tokio::test] +async fn edit_uses_the_context_workspace_when_one_is_threaded() { + let home = tempfile::tempdir().unwrap(); + let isolated = tempfile::tempdir().unwrap(); + std::fs::write(isolated.path().join("w.txt"), "aaa").unwrap(); + let tool = EditFileTool::new(test_security(home.path().to_path_buf())); + let context = WorkspaceContext::at(isolated.path()); + let result = tool + .execute_with_context( + edit_args("w.txt"), + ToolCallOptions::default(), + Some(&context), + ) + .await + .unwrap(); + assert!(!result.is_error, "{}", result.output()); + assert_eq!( + std::fs::read_to_string(isolated.path().join("w.txt")).unwrap(), + "bbb" + ); +} + +#[tokio::test] +async fn edit_honours_the_file_state_guard_and_records_its_write() { + crate::fs::file_state::init_global(true); + let dir = tempfile::tempdir().unwrap(); + let target = dir.path().canonicalize().unwrap().join("g.txt"); + std::fs::write(&target, "aaa").unwrap(); + let tool = EditFileTool::new(test_security(dir.path().to_path_buf())); + let agent = format!("ed-agent-{}", dir.path().display()); + let other = format!("ed-other-{}", dir.path().display()); + let run = |agent: String| { + crate::fs::file_state::with_file_state_agent_id(agent, tool.execute(edit_args("g.txt"))) + }; + + crate::fs::file_state::record_read( + &agent, + target.clone(), + std::time::SystemTime::now(), + true, + std::time::Instant::now(), + ); + let result = run(agent.clone()).await.unwrap(); + assert!(result.output().contains("Partial read")); + + crate::fs::file_state::record_read( + &agent, + target.clone(), + std::time::SystemTime::now(), + false, + std::time::Instant::now(), + ); + crate::fs::file_state::record_write(&other, target.clone()); + let result = run(agent.clone()).await.unwrap(); + assert!(result.output().contains("Stale read")); + + crate::fs::file_state::record_read( + &agent, + target.clone(), + std::time::SystemTime::now(), + false, + std::time::Instant::now(), + ); + let result = run(agent.clone()).await.unwrap(); + assert!(!result.is_error, "{}", result.output()); + assert!(crate::fs::file_state::check_stale_read(&agent, &target).is_none()); +} diff --git a/crates/tinyagents-harness-tools/src/fs/filesystem/file_read/extra_tests.rs b/crates/tinyagents-harness-tools/src/fs/filesystem/file_read/extra_tests.rs new file mode 100644 index 000000000..e8409a737 --- /dev/null +++ b/crates/tinyagents-harness-tools/src/fs/filesystem/file_read/extra_tests.rs @@ -0,0 +1,96 @@ +//! Budget, metadata-failure, file-state and paging-marker tests for `file_read`. + +#![allow(clippy::expect_used, clippy::panic, clippy::unwrap_used)] + +use super::*; +use crate::fs::filesystem::test_support::{TestGate, WorkspaceContext, WrapGate}; +use tempfile::TempDir; + +#[test] +fn reads_are_concurrency_safe() { + let dir = TempDir::new().unwrap(); + let tool = FileReadTool::new(TestGate::supervised(dir.path().to_path_buf())); + assert!(tool.is_concurrency_safe(&json!({}))); +} + +#[tokio::test] +async fn exhausted_action_budget_is_refused() { + let dir = TempDir::new().unwrap(); + std::fs::write(dir.path().join("a.txt"), "x").unwrap(); + let gate = WrapGate::new(TestGate::supervised(dir.path().to_path_buf())) + .no_budget() + .arc(); + let r = FileReadTool::new(gate) + .execute(json!({"path": "a.txt"})) + .await + .unwrap(); + assert!(r.output().contains("action budget exhausted")); +} + +#[tokio::test] +async fn missing_metadata_is_reported() { + let dir = TempDir::new().unwrap(); + let gate = WrapGate::new(TestGate::supervised(dir.path().to_path_buf())) + .fixed_path(dir.path().join("vanished.txt")) + .arc(); + let r = FileReadTool::new(gate) + .execute(json!({"path": "vanished.txt"})) + .await + .unwrap(); + assert!(r.is_error); + assert!(r.output().contains("Failed to read file metadata")); +} + +#[tokio::test] +async fn a_read_under_an_agent_id_records_file_state() { + let dir = TempDir::new().unwrap(); + std::fs::write(dir.path().join("a.txt"), "tracked").unwrap(); + let tool = FileReadTool::new(TestGate::supervised(dir.path().to_path_buf())); + let r = crate::fs::file_state::with_file_state_agent_id( + "agent-file-read-extra".to_string(), + tool.execute(json!({"path": "a.txt"})), + ) + .await + .unwrap(); + assert_eq!(r.output(), "tracked"); +} + +#[tokio::test] +async fn execute_with_context_delegates() { + let dir = TempDir::new().unwrap(); + std::fs::write(dir.path().join("a.txt"), "ctx").unwrap(); + let tool = FileReadTool::new(TestGate::supervised(dir.path().to_path_buf())); + let ctx = WorkspaceContext::at(dir.path()); + let r = tool + .execute_with_context( + json!({"path": "a.txt"}), + ToolCallOptions::default(), + Some(&ctx), + ) + .await + .unwrap(); + assert_eq!(r.output(), "ctx"); +} + +#[tokio::test] +async fn an_oversized_path_drops_out_of_the_continuation_marker() { + let dir = TempDir::new().unwrap(); + let file = dir.path().join("big.txt"); + std::fs::write(&file, "a".repeat(40 * 1024)).unwrap(); + let gate = WrapGate::new(TestGate::supervised(dir.path().to_path_buf())) + .fixed_path(file) + .arc(); + let long_path = "p".repeat(20 * 1024); + let r = FileReadTool::new(gate) + .execute(json!({"path": long_path})) + .await + .unwrap(); + assert!(!r.is_error); + let out = r.output(); + assert!( + out.contains("continue with file_read at \"offset\":"), + "marker missing" + ); + assert!(!out.contains("ppppp")); + assert!(out.len() <= 16 * 1024); +} diff --git a/crates/tinyagents-harness-tools/src/fs/filesystem/file_read/mod.rs b/crates/tinyagents-harness-tools/src/fs/filesystem/file_read/mod.rs new file mode 100644 index 000000000..637e987e9 --- /dev/null +++ b/crates/tinyagents-harness-tools/src/fs/filesystem/file_read/mod.rs @@ -0,0 +1,240 @@ +use super::gate::{FsGate, gate_for_context}; +use crate::fs::file_state; +use async_trait::async_trait; +use serde_json::json; +use std::sync::Arc; +use tinytools::ToolRunContext; +use tinytools::{Tool, ToolCallOptions, ToolResult}; + +const MAX_FILE_SIZE_BYTES: u64 = 10 * 1024 * 1024; +/// Keep a page below the harness's 16 KiB per-result ceiling so the continuation +/// marker survives the result middleware. Returning the whole remainder and +/// relying on that middleware to truncate it loses the next offset and makes a +/// model repeat the same read forever. +const MAX_PAGE_BYTES: usize = 12 * 1024; +const MAX_TOOL_OUTPUT_BYTES: usize = 16 * 1024; + +/// Read file contents with path sandboxing +#[derive(Debug)] +pub struct FileReadTool { + gate: Arc, +} + +impl FileReadTool { + /// Largest file this tool will open. Anything that must stay readable + /// through it (a persisted tool-result artifact) has to fit. + pub const MAX_FILE_SIZE_BYTES: u64 = MAX_FILE_SIZE_BYTES; + + /// Create the tool. + #[must_use] + pub fn new(gate: Arc) -> Self { + Self { gate } + } +} + +#[async_trait] +impl Tool for FileReadTool { + fn name(&self) -> &'static str { + "file_read" + } + + fn description(&self) -> &'static str { + "Read the contents of a file in your working directory (the action sandbox). \ + Relative paths resolve against that directory; paths outside it are blocked. \ + To read a file written by `shell`, confirm its location with `pwd` and use the \ + same relative path. Long files are returned one page at a time; continue with \ + the exact `offset` reported at the end of the page." + } + + fn parameters_schema(&self) -> serde_json::Value { + json!({ + "type": "object", + "properties": { + "path": { + "type": "string", + "description": "Relative path to the file within the workspace" + }, + "offset": { + "type": "integer", + "minimum": 0, + "description": "Byte offset to start reading from. Use the offset a paged read reports to continue it." + } + }, + "required": ["path"] + }) + } + + /// Pure read — safe to fan out across parallel `file_read` calls. + fn is_concurrency_safe(&self, _args: &serde_json::Value) -> bool { + true + } + + async fn execute(&self, args: serde_json::Value) -> anyhow::Result { + self.execute_in_context(args, None).await + } + + async fn execute_with_context( + &self, + args: serde_json::Value, + _options: ToolCallOptions, + context: Option<&dyn ToolRunContext>, + ) -> anyhow::Result { + self.execute_in_context(args, context).await + } +} + +impl FileReadTool { + async fn execute_in_context( + &self, + args: serde_json::Value, + context: Option<&dyn ToolRunContext>, + ) -> anyhow::Result { + let path = args + .get("path") + .and_then(|v| v.as_str()) + .ok_or_else(|| anyhow::anyhow!("Missing 'path' parameter"))?; + + if self.gate.is_rate_limited() { + return Ok(ToolResult::error( + "Rate limit exceeded: too many actions in the last hour", + )); + } + + // Record action BEFORE validation so that every non-trivially-rejected + // request consumes rate limit budget. This prevents attackers from probing + // path existence (via canonicalize errors) without rate limit cost. + if !self.gate.record_action() { + return Ok(ToolResult::error( + "Rate limit exceeded: action budget exhausted", + )); + } + + // Security check: validate path string, resolve symlinks, confirm workspace containment. + let path_policy = gate_for_context(&self.gate, context, "file_read"); + let resolved_path = match path_policy.validate_path(path).await { + Ok(p) => p, + Err(msg) => return Ok(ToolResult::error(msg)), + }; + + // Captured before any read I/O: `file_state` compares it against + // sibling agents' writes to decide whether this read may already be stale. + let read_started = std::time::Instant::now(); + + // Check file size AFTER canonicalization to prevent TOCTOU symlink bypass + match tokio::fs::metadata(&resolved_path).await { + Ok(meta) => { + if meta.len() > MAX_FILE_SIZE_BYTES { + return Ok(ToolResult::error(format!( + "File too large: {} bytes (limit: {MAX_FILE_SIZE_BYTES} bytes)", + meta.len() + ))); + } + } + Err(e) => { + return Ok(ToolResult::error(format!( + "Failed to read file metadata: {e}" + ))); + } + } + + match tokio::fs::read_to_string(&resolved_path).await { + Ok(contents) => { + if let Some(agent_id) = file_state::current_file_state_agent_id() { + let mtime = tokio::fs::metadata(&resolved_path) + .await + .ok() + .and_then(|m| m.modified().ok()) + .unwrap_or(std::time::SystemTime::UNIX_EPOCH); + file_state::record_read(&agent_id, resolved_path, mtime, false, read_started); + } + // An absent or null offset reads from the start; anything else + // that is not a non-negative integer is rejected, so a + // malformed continuation never silently re-reads the file from + // byte 0. + let offset = match args.get("offset") { + None | Some(serde_json::Value::Null) => 0, + Some(value) => match value.as_u64().and_then(|o| usize::try_from(o).ok()) { + Some(offset) => offset, + None => { + return Ok(ToolResult::error(format!( + "offset must be a non-negative integer byte offset, got {value}" + ))); + } + }, + }; + if offset > contents.len() { + return Ok(ToolResult::error(format!( + "offset {offset} is past the end of the file ({} bytes)", + contents.len() + ))); + } + // Rejected rather than snapped to a boundary: a caller paging + // from `offset` computes the next offset from where it asked to + // start, so a silent move backwards would skip or repeat bytes. + if !contents.is_char_boundary(offset) { + return Ok(ToolResult::error(format!( + "offset {offset} falls inside a multi-byte character; continue from the offset a paged read reports" + ))); + } + Ok(ToolResult::success(page_contents(&contents, path, offset))) + } + Err(e) => Ok(ToolResult::error(format!("Failed to read file: {e}"))), + } + } +} + +fn page_contents(contents: &str, path: &str, offset: usize) -> String { + let initial_end = offset.saturating_add(MAX_PAGE_BYTES).min(contents.len()); + if initial_end == contents.len() { + return contents[offset..].to_string(); + } + + let path_json = + serde_json::to_string(path).unwrap_or_else(|_| "\"\"".to_string()); + let marker_with_path = |end: usize| { + format!( + "\n\n[file_read page: bytes {offset}..{end} of {}; continue with file_read \ + {{\"path\":{path_json},\"offset\":{end}}}]", + contents.len() + ) + }; + let marker_without_path = |end: usize| { + format!( + "\n\n[file_read page: bytes {offset}..{end} of {}; continue with file_read at \ + \"offset\":{end}]", + contents.len() + ) + }; + + // Size against the longest offset this page can report. If the escaped path + // consumes too much of the middleware budget, omit it: the caller already + // has the path and the continuation offset is the irreplaceable part. + let longest_with_path = marker_with_path(initial_end); + let include_path = longest_with_path.len() + 4 <= MAX_TOOL_OUTPUT_BYTES; + let marker_len = if include_path { + longest_with_path.len() + } else { + marker_without_path(initial_end).len() + }; + let content_budget = MAX_PAGE_BYTES.min(MAX_TOOL_OUTPUT_BYTES.saturating_sub(marker_len)); + let mut end = offset.saturating_add(content_budget).min(contents.len()); + while end > offset && !contents.is_char_boundary(end) { + end -= 1; + } + let marker = if include_path { + marker_with_path(end) + } else { + marker_without_path(end) + }; + let mut page = contents[offset..end].to_string(); + page.push_str(&marker); + page +} + +#[cfg(test)] +#[path = "mod_tests.rs"] +mod tests; + +#[cfg(test)] +#[path = "extra_tests.rs"] +mod extra_tests; diff --git a/crates/tinyagents-harness-tools/src/fs/filesystem/file_read/mod_tests.rs b/crates/tinyagents-harness-tools/src/fs/filesystem/file_read/mod_tests.rs new file mode 100644 index 000000000..7571f21c2 --- /dev/null +++ b/crates/tinyagents-harness-tools/src/fs/filesystem/file_read/mod_tests.rs @@ -0,0 +1,431 @@ +//! Behavior tests for the `file_read` tool, driven through a fake [`FsGate`]. + +#![allow(clippy::expect_used, clippy::panic, clippy::unwrap_used)] + +use super::*; +use crate::fs::filesystem::test_support::{AutonomyLevel, TestGate}; + +fn test_security(workspace: std::path::PathBuf) -> Arc { + TestGate::supervised(workspace) +} + +fn test_security_with( + workspace: std::path::PathBuf, + autonomy: AutonomyLevel, + max_actions_per_hour: usize, +) -> Arc { + TestGate::with(workspace, autonomy, max_actions_per_hour) +} + +#[test] +fn file_read_name() { + let tool = FileReadTool::new(test_security(std::env::temp_dir())); + assert_eq!(tool.name(), "file_read"); +} + +#[test] +fn file_read_schema_has_path() { + let tool = FileReadTool::new(test_security(std::env::temp_dir())); + let schema = tool.parameters_schema(); + assert!(schema["properties"]["path"].is_object()); + assert!( + schema["required"] + .as_array() + .unwrap() + .contains(&json!("path")) + ); +} + +#[tokio::test] +async fn file_read_existing_file() { + let dir = std::env::temp_dir().join("openhuman_test_file_read"); + let _ = tokio::fs::remove_dir_all(&dir).await; + tokio::fs::create_dir_all(&dir).await.unwrap(); + tokio::fs::write(dir.join("test.txt"), "hello world") + .await + .unwrap(); + + let tool = FileReadTool::new(test_security(dir.clone())); + let result = tool.execute(json!({"path": "test.txt"})).await.unwrap(); + assert!(!result.is_error); + assert_eq!(result.output(), "hello world"); + assert!(!result.is_error); + + let _ = tokio::fs::remove_dir_all(&dir).await; +} + +#[tokio::test] +async fn file_read_nonexistent_file() { + let dir = std::env::temp_dir().join("openhuman_test_file_read_missing"); + let _ = tokio::fs::remove_dir_all(&dir).await; + tokio::fs::create_dir_all(&dir).await.unwrap(); + + let tool = FileReadTool::new(test_security(dir.clone())); + let result = tool.execute(json!({"path": "nope.txt"})).await.unwrap(); + assert!(result.is_error); + assert!(&result.output().contains("Failed to resolve")); + + let _ = tokio::fs::remove_dir_all(&dir).await; +} + +#[tokio::test] +async fn file_read_blocks_path_traversal() { + let dir = std::env::temp_dir().join("openhuman_test_file_read_traversal"); + let _ = tokio::fs::remove_dir_all(&dir).await; + tokio::fs::create_dir_all(&dir).await.unwrap(); + + let tool = FileReadTool::new(test_security(dir.clone())); + let result = tool + .execute(json!({"path": "../../../etc/passwd"})) + .await + .unwrap(); + assert!(result.is_error); + assert!(&result.output().contains("not allowed")); + + let _ = tokio::fs::remove_dir_all(&dir).await; +} + +#[tokio::test] +async fn file_read_blocks_absolute_path() { + let tool = FileReadTool::new(test_security(std::env::temp_dir())); + let result = tool.execute(json!({"path": "/etc/passwd"})).await.unwrap(); + assert!(result.is_error); + assert!(&result.output().contains("not allowed")); +} + +#[tokio::test] +async fn file_read_blocks_when_rate_limited() { + let dir = std::env::temp_dir().join("openhuman_test_file_read_rate_limited"); + let _ = tokio::fs::remove_dir_all(&dir).await; + tokio::fs::create_dir_all(&dir).await.unwrap(); + tokio::fs::write(dir.join("test.txt"), "hello world") + .await + .unwrap(); + + let tool = FileReadTool::new(test_security_with( + dir.clone(), + AutonomyLevel::Supervised, + 0, + )); + let result = tool.execute(json!({"path": "test.txt"})).await.unwrap(); + + assert!(result.is_error); + assert!(result.output().contains("Rate limit exceeded")); + + let _ = tokio::fs::remove_dir_all(&dir).await; +} + +#[tokio::test] +async fn file_read_allows_readonly_mode() { + let dir = std::env::temp_dir().join("openhuman_test_file_read_readonly"); + let _ = tokio::fs::remove_dir_all(&dir).await; + tokio::fs::create_dir_all(&dir).await.unwrap(); + tokio::fs::write(dir.join("test.txt"), "readonly ok") + .await + .unwrap(); + + let tool = FileReadTool::new(test_security_with(dir.clone(), AutonomyLevel::ReadOnly, 20)); + let result = tool.execute(json!({"path": "test.txt"})).await.unwrap(); + + assert!(!result.is_error); + assert_eq!(result.output(), "readonly ok"); + + let _ = tokio::fs::remove_dir_all(&dir).await; +} + +#[tokio::test] +async fn file_read_missing_path_param() { + let tool = FileReadTool::new(test_security(std::env::temp_dir())); + let result = tool.execute(json!({})).await; + assert!(result.is_err()); +} + +#[tokio::test] +async fn file_read_empty_file() { + let dir = std::env::temp_dir().join("openhuman_test_file_read_empty"); + let _ = tokio::fs::remove_dir_all(&dir).await; + tokio::fs::create_dir_all(&dir).await.unwrap(); + tokio::fs::write(dir.join("empty.txt"), "").await.unwrap(); + + let tool = FileReadTool::new(test_security(dir.clone())); + let result = tool.execute(json!({"path": "empty.txt"})).await.unwrap(); + assert!(!result.is_error); + assert_eq!(result.output(), ""); + + let _ = tokio::fs::remove_dir_all(&dir).await; +} + +#[tokio::test] +async fn file_read_nested_path() { + let dir = std::env::temp_dir().join("openhuman_test_file_read_nested"); + let _ = tokio::fs::remove_dir_all(&dir).await; + tokio::fs::create_dir_all(dir.join("sub/dir")) + .await + .unwrap(); + tokio::fs::write(dir.join("sub/dir/deep.txt"), "deep content") + .await + .unwrap(); + + let tool = FileReadTool::new(test_security(dir.clone())); + let result = tool + .execute(json!({"path": "sub/dir/deep.txt"})) + .await + .unwrap(); + assert!(!result.is_error); + assert_eq!(result.output(), "deep content"); + + let _ = tokio::fs::remove_dir_all(&dir).await; +} + +#[cfg(unix)] +#[tokio::test] +async fn file_read_blocks_symlink_escape() { + use std::os::unix::fs::symlink; + + let root = std::env::temp_dir().join("openhuman_test_file_read_symlink_escape"); + let workspace = root.join("workspace"); + let outside = root.join("outside"); + + let _ = tokio::fs::remove_dir_all(&root).await; + tokio::fs::create_dir_all(&workspace).await.unwrap(); + tokio::fs::create_dir_all(&outside).await.unwrap(); + + tokio::fs::write(outside.join("secret.txt"), "outside workspace") + .await + .unwrap(); + + symlink(outside.join("secret.txt"), workspace.join("escape.txt")).unwrap(); + + let tool = FileReadTool::new(test_security(workspace.clone())); + let result = tool.execute(json!({"path": "escape.txt"})).await.unwrap(); + + assert!(result.is_error); + // After the symlink-safe canonical check landed in + // SecurityPolicy::is_path_allowed (#1927), the policy layer blocks + // the escape before file_read's own resolved-path check runs — the + // error becomes "Path not allowed by security policy". Either + // message signals defense-in-depth worked. + let out = result.output(); + assert!( + out.contains("escapes workspace") || out.contains("not allowed"), + "expected escape/not-allowed error, got: {out}" + ); + + let _ = tokio::fs::remove_dir_all(&root).await; +} + +#[tokio::test] +async fn file_read_nonexistent_consumes_rate_limit_budget() { + let dir = std::env::temp_dir().join("openhuman_test_file_read_probe"); + let _ = tokio::fs::remove_dir_all(&dir).await; + tokio::fs::create_dir_all(&dir).await.unwrap(); + + // Allow only 2 actions total + let tool = FileReadTool::new(test_security_with( + dir.clone(), + AutonomyLevel::Supervised, + 2, + )); + + // Both reads fail (file doesn't exist) but should consume budget + let r1 = tool.execute(json!({"path": "nope1.txt"})).await.unwrap(); + assert!(r1.is_error); + assert!(r1.output().contains("Failed to resolve")); + + let r2 = tool.execute(json!({"path": "nope2.txt"})).await.unwrap(); + assert!(r2.is_error); + assert!(r2.output().contains("Failed to resolve")); + + // Third attempt should be rate limited even though file doesn't exist + let r3 = tool.execute(json!({"path": "nope3.txt"})).await.unwrap(); + assert!(r3.is_error); + let r3_output = r3.output(); + assert!( + r3_output.contains("Rate limit"), + "Expected rate limit error, got: {r3_output:?}" + ); + + let _ = tokio::fs::remove_dir_all(&dir).await; +} + +#[tokio::test] +async fn file_read_rejects_oversized_file() { + let dir = std::env::temp_dir().join("openhuman_test_file_read_large"); + let _ = tokio::fs::remove_dir_all(&dir).await; + tokio::fs::create_dir_all(&dir).await.unwrap(); + + // Create a file just over 10 MB + let big = vec![b'x'; 10 * 1024 * 1024 + 1]; + tokio::fs::write(dir.join("huge.bin"), &big).await.unwrap(); + + let tool = FileReadTool::new(test_security(dir.clone())); + let result = tool.execute(json!({"path": "huge.bin"})).await.unwrap(); + assert!(result.is_error); + assert!(&result.output().contains("File too large")); + + let _ = tokio::fs::remove_dir_all(&dir).await; +} + +#[tokio::test] +async fn file_read_reports_invalid_utf8_as_an_error_not_a_panic() { + let dir = std::env::temp_dir().join("openhuman_test_file_read_non_utf8"); + let _ = tokio::fs::remove_dir_all(&dir).await; + tokio::fs::create_dir_all(&dir).await.unwrap(); + + // Passes the size cap, fails `read_to_string`: 0x80 is a lone UTF-8 + // continuation byte with no lead byte, invalid at any position. + tokio::fs::write(dir.join("binary.dat"), [0xFFu8, 0xFE, 0x00, 0x80]) + .await + .unwrap(); + + let tool = FileReadTool::new(test_security(dir.clone())); + let result = tool.execute(json!({"path": "binary.dat"})).await.unwrap(); + assert!( + result.is_error, + "invalid UTF-8 must surface as a tool error, not succeed with replaced/garbled text" + ); + assert!(&result.output().contains("Failed to read file")); + + let _ = tokio::fs::remove_dir_all(&dir).await; +} + +#[tokio::test] +async fn file_read_offset_continues_from_the_given_byte() { + let dir = tempfile::tempdir().unwrap(); + tokio::fs::write(dir.path().join("paged.txt"), "hello world") + .await + .unwrap(); + let tool = FileReadTool::new(test_security(dir.path().to_path_buf())); + + let page = tool + .execute(json!({"path": "paged.txt", "offset": 6})) + .await + .unwrap(); + assert!(!page.is_error, "{}", page.output()); + assert_eq!( + page.output(), + "world", + "a read with an offset must start at that byte" + ); + + let past_end = tool + .execute(json!({"path": "paged.txt", "offset": 99})) + .await + .unwrap(); + assert!( + past_end.is_error, + "an offset past the end must be reported, not returned as an empty read" + ); +} + +#[tokio::test] +async fn file_read_pages_long_files_and_reports_a_working_continuation() { + let dir = tempfile::tempdir().unwrap(); + let contents = format!("{}éTAIL", "x".repeat(MAX_PAGE_BYTES - 1)); + tokio::fs::write(dir.path().join("long.txt"), &contents) + .await + .unwrap(); + let tool = FileReadTool::new(test_security(dir.path().to_path_buf())); + + let first = tool.execute(json!({"path": "long.txt"})).await.unwrap(); + assert!(!first.is_error, "{}", first.output()); + assert!( + first.output().contains("continue with file_read"), + "a truncated page must tell the caller exactly how to continue: {}", + first.output() + ); + assert!( + first.output().len() < 16 * 1024, + "the page and marker must survive the harness's per-result budget" + ); + let marker = first + .output() + .split("\"offset\":") + .nth(1) + .and_then(|tail| tail.split('}').next()) + .and_then(|value| value.parse::().ok()) + .expect("continuation offset in marker"); + assert!(contents.is_char_boundary(marker)); + + let second = tool + .execute(json!({"path": "long.txt", "offset": marker})) + .await + .unwrap(); + assert!(!second.is_error, "{}", second.output()); + assert_eq!( + format!("{}{}", &first.output()[..marker], second.output()), + contents + ); +} + +#[test] +fn file_read_reserves_the_continuation_marker_for_a_long_escaped_path() { + let contents = "x".repeat(MAX_PAGE_BYTES * 2); + let path = format!("{}quoted\\\"name.txt", "nested\\\\".repeat(900)); + let page = page_contents(&contents, &path, 0); + + assert!(page.len() <= MAX_TOOL_OUTPUT_BYTES, "{}", page.len()); + assert!(page.contains("continue with file_read")); + assert!(page.contains("\"offset\":")); + assert!( + !page.contains(&path), + "an oversized escaped path should be omitted so the offset survives" + ); +} + +#[tokio::test] +async fn file_read_rejects_an_offset_inside_a_multibyte_character() { + let dir = tempfile::tempdir().unwrap(); + // "é" is two bytes, so byte 1 is inside it. + tokio::fs::write(dir.path().join("utf8.txt"), "é world") + .await + .unwrap(); + let tool = FileReadTool::new(test_security(dir.path().to_path_buf())); + + let inside = tool + .execute(json!({"path": "utf8.txt", "offset": 1})) + .await + .unwrap(); + assert!( + inside.is_error, + "an offset inside a character must be rejected, not moved, or paging skips bytes: {}", + inside.output() + ); + + let boundary = tool + .execute(json!({"path": "utf8.txt", "offset": 2})) + .await + .unwrap(); + assert_eq!(boundary.output(), " world"); +} + +#[tokio::test] +async fn file_read_rejects_an_offset_that_is_not_a_non_negative_integer() { + let dir = tempfile::tempdir().unwrap(); + tokio::fs::write(dir.path().join("offsets.txt"), "hello world") + .await + .unwrap(); + let tool = FileReadTool::new(test_security(dir.path().to_path_buf())); + + for bad in [json!(-5), json!("6"), json!(1.5)] { + let result = tool + .execute(json!({"path": "offsets.txt", "offset": bad.clone()})) + .await + .unwrap(); + assert!( + result.is_error, + "offset {bad} must be rejected, not treated as 0: got {:?}", + result.output() + ); + } + + let null_offset = tool + .execute(json!({"path": "offsets.txt", "offset": null})) + .await + .unwrap(); + assert_eq!( + null_offset.output(), + "hello world", + "a null offset reads from the start" + ); +} diff --git a/crates/tinyagents-harness-tools/src/fs/filesystem/file_sink.rs b/crates/tinyagents-harness-tools/src/fs/filesystem/file_sink.rs new file mode 100644 index 000000000..a50dcf201 --- /dev/null +++ b/crates/tinyagents-harness-tools/src/fs/filesystem/file_sink.rs @@ -0,0 +1,59 @@ +//! The filesystem write the tools go through. +//! +//! Every tool that reports "the write failed" has a branch that only runs when +//! the OS refuses, and that branch is the one worth asserting: it decides +//! whether a failed edit is reported honestly or fabricated as a success. +//! +//! Provoking a real refusal from a test is where this gets awkward. A +//! mode-bit denial is a discretionary check, and a root test runner carries +//! `CAP_DAC_OVERRIDE` and writes through it — so a fixture built on +//! permissions succeeds under CI, and the case either fails on a premise that +//! never held or is skipped and asserts nothing at all. Both read as coverage. +//! +//! So the write is a seam instead. Production installs the OS; a test installs +//! one that refuses, and the reporting branch becomes ordinary to assert on +//! any platform and as any user. + +use std::path::Path; +use std::sync::Arc; + +use async_trait::async_trait; + +/// Where a tool's file writes land. +#[async_trait] +pub(super) trait FileSink: std::fmt::Debug + Send + Sync { + async fn write(&self, path: &Path, contents: &[u8]) -> std::io::Result<()>; +} + +/// The real filesystem. +#[derive(Debug)] +pub(super) struct OsFileSink; + +#[async_trait] +impl FileSink for OsFileSink { + async fn write(&self, path: &Path, contents: &[u8]) -> std::io::Result<()> { + tokio::fs::write(path, contents).await + } +} + +/// The sink a tool uses unless a caller installs another one. +pub(super) fn os_sink() -> Arc { + Arc::new(OsFileSink) +} + +/// A sink that refuses every write with `kind`, for asserting the branch a +/// tool takes when the filesystem says no. +#[cfg(test)] +#[derive(Debug)] +pub(super) struct RefusingSink(pub std::io::ErrorKind); + +#[cfg(test)] +#[async_trait] +impl FileSink for RefusingSink { + async fn write(&self, path: &Path, _contents: &[u8]) -> std::io::Result<()> { + Err(std::io::Error::new( + self.0, + format!("refused by the test sink: {}", path.display()), + )) + } +} diff --git a/crates/tinyagents-harness-tools/src/fs/filesystem/file_write/mod.rs b/crates/tinyagents-harness-tools/src/fs/filesystem/file_write/mod.rs new file mode 100644 index 000000000..6a8947a50 --- /dev/null +++ b/crates/tinyagents-harness-tools/src/fs/filesystem/file_write/mod.rs @@ -0,0 +1,234 @@ +use super::gate::{FsGate, gate_for_context}; +use crate::fs::file_state; +use async_trait::async_trait; +use serde_json::json; +use std::sync::Arc; +use tinytools::ToolRunContext; +use tinytools::{PermissionLevel, Tool, ToolCallOptions, ToolResult}; + +const MAX_CONTENT_BYTES: usize = 5 * 1024 * 1024; + +/// Write file contents with path sandboxing +#[derive(Debug)] +pub struct FileWriteTool { + gate: Arc, + approval_workspace_root: Option, + sink: Arc, +} + +impl FileWriteTool { + /// Create the tool. + #[must_use] + pub fn new(gate: Arc) -> Self { + Self { + gate, + approval_workspace_root: None, + sink: super::file_sink::os_sink(), + } + } + + /// Sends this tool's writes somewhere other than the OS. + #[cfg(test)] + pub(super) fn with_sink(mut self, sink: Arc) -> Self { + self.sink = sink; + self + } + + /// Use the same effective workspace root for approval routing that tool + /// execution receives through its `ToolExecutionContext`. + pub fn with_approval_workspace_root( + gate: Arc, + approval_workspace_root: std::path::PathBuf, + ) -> Self { + Self { + gate, + approval_workspace_root: Some(approval_workspace_root), + sink: super::file_sink::os_sink(), + } + } +} + +#[async_trait] +impl Tool for FileWriteTool { + fn name(&self) -> &'static str { + "file_write" + } + + fn description(&self) -> &'static str { + "Write contents to a file in your working directory (the action sandbox). \ + Relative paths resolve against that directory; writes outside it are blocked. \ + Reference the file later by the same relative path so `file_read` resolves to it." + } + + fn parameters_schema(&self) -> serde_json::Value { + json!({ + "type": "object", + "properties": { + "path": { + "type": "string", + "description": "Relative path to the file within the workspace" + }, + "content": { + "type": "string", + "description": "Content to write to the file" + } + }, + "required": ["path", "content"] + }) + } + + fn permission_level(&self) -> PermissionLevel { + PermissionLevel::Write + } + + /// "Ask before edit": modifying an **existing** file routes through the + /// human approval gate in ask-before-edit mode; creating a **new** file is + /// free. In Full neither prompts; in read-only `execute` blocks via + /// `can_act()`. The existence probe is best-effort (relative to the + /// workspace); when the path can't be resolved we fail safe and prompt. + /// + /// The probe runs at gate-routing time, microseconds before `execute()` in + /// the same sequential turn. A create→edit flip in that window would require + /// an external process to win the race — outside this gate's threat model, + /// which governs the agent, not concurrent writers — and `execute()` still + /// enforces workspace containment + symlink refusal regardless, so a write + /// that slips through as "create" cannot escape the sandbox. We therefore + /// do not re-probe in `execute()`. + fn external_effect_with_args(&self, args: &serde_json::Value) -> bool { + if !self.gate.write_needs_approval() { + return false; // Full (allow) or read-only (blocked in execute) + } + let Some(path) = args.get("path").and_then(|v| v.as_str()) else { + return true; // unknown path → prompt (fail safe) + }; + let target = if std::path::Path::new(path).is_absolute() { + std::path::PathBuf::from(path) + } else { + self.approval_workspace_root + .as_deref() + .unwrap_or_else(|| self.gate.action_dir()) + .join(path) + }; + // Sync `stat` — intentionally blocking, since the `Tool` trait makes + // this method sync. Fast for local paths; would only need + // `block_in_place` if a remote/slow filesystem is ever supported here. + target.exists() // exists = edit → prompt; new = create → free + } + + async fn execute(&self, args: serde_json::Value) -> anyhow::Result { + self.execute_in_context(args, None).await + } + + async fn execute_with_context( + &self, + args: serde_json::Value, + _options: ToolCallOptions, + context: Option<&dyn ToolRunContext>, + ) -> anyhow::Result { + self.execute_in_context(args, context).await + } +} + +impl FileWriteTool { + async fn execute_in_context( + &self, + args: serde_json::Value, + context: Option<&dyn ToolRunContext>, + ) -> anyhow::Result { + let path = args + .get("path") + .and_then(|v| v.as_str()) + .ok_or_else(|| anyhow::anyhow!("Missing 'path' parameter"))?; + + let content = args + .get("content") + .and_then(|v| v.as_str()) + .ok_or_else(|| anyhow::anyhow!("Missing 'content' parameter"))?; + + if !self.gate.can_act() { + return Ok(ToolResult::error( + "[policy-blocked] Action blocked: autonomy is read-only", + )); + } + + if content.len() > MAX_CONTENT_BYTES { + return Ok(ToolResult::error(format!( + "Content too large: {} bytes (limit: {MAX_CONTENT_BYTES} bytes)", + content.len() + ))); + } + + if self.gate.is_rate_limited() { + return Ok(ToolResult::error( + "Rate limit exceeded: too many actions in the last hour", + )); + } + + // Security check first: validate path string, resolve symlinks, confirm workspace + // containment. validate_parent_path walks up to the deepest existing ancestor so + // it does not require the parent directory to exist yet. + let path_policy = gate_for_context(&self.gate, context, "file_write"); + let resolved_target = match path_policy.validate_parent_path(path).await { + Ok(p) => p, + Err(msg) => return Ok(ToolResult::error(msg)), + }; + + // Create parent directory only at the validated, resolved location. + if let Some(resolved_parent) = resolved_target.parent() { + tokio::fs::create_dir_all(resolved_parent).await?; + } + + // If the target already exists and is a symlink, refuse to follow it + if let Ok(meta) = tokio::fs::symlink_metadata(&resolved_target).await + && meta.file_type().is_symlink() + { + return Ok(ToolResult::error(format!( + "Refusing to write through symlink: {}", + resolved_target.display() + ))); + } + + if !self.gate.record_action() { + return Ok(ToolResult::error( + "Rate limit exceeded: action budget exhausted", + )); + } + + // File-state guard: reject writes based on stale or partial reads. + if let Some(agent_id) = file_state::current_file_state_agent_id() { + if let Some(msg) = file_state::check_stale_read(&agent_id, &resolved_target) { + tracing::debug!( + agent = %agent_id, + path = %resolved_target.display(), + "[file_state] file_write blocked: stale read" + ); + return Ok(ToolResult::error(msg)); + } + if let Some(msg) = file_state::check_partial_read(&agent_id, &resolved_target) { + tracing::debug!( + agent = %agent_id, + path = %resolved_target.display(), + "[file_state] file_write blocked: partial read" + ); + return Ok(ToolResult::error(msg)); + } + } + + match self.sink.write(&resolved_target, content.as_bytes()).await { + Ok(()) => { + if let Some(agent_id) = file_state::current_file_state_agent_id() { + file_state::record_write(&agent_id, resolved_target); + } + Ok(ToolResult::success(format!( + "Written {} bytes to {path}", + content.len() + ))) + } + Err(e) => Ok(ToolResult::error(format!("Failed to write file: {e}"))), + } + } +} + +#[cfg(test)] +#[path = "mod_tests.rs"] +mod tests; diff --git a/crates/tinyagents-harness-tools/src/fs/filesystem/file_write/mod_tests.rs b/crates/tinyagents-harness-tools/src/fs/filesystem/file_write/mod_tests.rs new file mode 100644 index 000000000..edcd9717d --- /dev/null +++ b/crates/tinyagents-harness-tools/src/fs/filesystem/file_write/mod_tests.rs @@ -0,0 +1,502 @@ +//! Behavior tests for the `file_write` tool, driven through a fake [`FsGate`]. + +#![allow(clippy::expect_used, clippy::panic, clippy::unwrap_used)] + +use super::*; +use crate::fs::filesystem::test_support::{AutonomyLevel, TestGate}; + +fn test_security(workspace: std::path::PathBuf) -> Arc { + TestGate::supervised(workspace) +} + +fn test_security_with( + workspace: std::path::PathBuf, + autonomy: AutonomyLevel, + max_actions_per_hour: usize, +) -> Arc { + TestGate::with(workspace, autonomy, max_actions_per_hour) +} + +#[test] +fn file_write_name() { + let tool = FileWriteTool::new(test_security(std::env::temp_dir())); + assert_eq!(tool.name(), "file_write"); +} + +#[test] +fn file_write_schema_has_path_and_content() { + let tool = FileWriteTool::new(test_security(std::env::temp_dir())); + let schema = tool.parameters_schema(); + assert!(schema["properties"]["path"].is_object()); + assert!(schema["properties"]["content"].is_object()); + let required = schema["required"].as_array().unwrap(); + assert!(required.contains(&json!("path"))); + assert!(required.contains(&json!("content"))); +} + +#[test] +fn approval_probe_uses_effective_workspace_root() { + let root = tempfile::tempdir().expect("root"); + let profile_workspace = root.path().join("profiles/alice"); + std::fs::create_dir_all(&profile_workspace).expect("profile workspace"); + std::fs::write(profile_workspace.join("notes.md"), "existing").expect("seed file"); + + let tool = FileWriteTool::with_approval_workspace_root( + test_security(root.path().to_path_buf()), + profile_workspace, + ); + + assert!(tool.external_effect_with_args(&json!({ + "path": "notes.md", + "content": "updated" + }))); +} + +#[tokio::test] +async fn file_write_creates_file() { + let dir = std::env::temp_dir().join("openhuman_test_file_write"); + let _ = tokio::fs::remove_dir_all(&dir).await; + tokio::fs::create_dir_all(&dir).await.unwrap(); + + let tool = FileWriteTool::new(test_security(dir.clone())); + let result = tool + .execute(json!({"path": "out.txt", "content": "written!"})) + .await + .unwrap(); + assert!(!result.is_error); + assert!(result.output().contains("8 bytes")); + + let content = tokio::fs::read_to_string(dir.join("out.txt")) + .await + .unwrap(); + assert_eq!(content, "written!"); + + let _ = tokio::fs::remove_dir_all(&dir).await; +} + +#[tokio::test] +async fn file_write_creates_parent_dirs() { + let dir = std::env::temp_dir().join("openhuman_test_file_write_nested"); + let _ = tokio::fs::remove_dir_all(&dir).await; + tokio::fs::create_dir_all(&dir).await.unwrap(); + + let tool = FileWriteTool::new(test_security(dir.clone())); + let result = tool + .execute(json!({"path": "a/b/c/deep.txt", "content": "deep"})) + .await + .unwrap(); + assert!(!result.is_error); + + let content = tokio::fs::read_to_string(dir.join("a/b/c/deep.txt")) + .await + .unwrap(); + assert_eq!(content, "deep"); + + let _ = tokio::fs::remove_dir_all(&dir).await; +} + +#[tokio::test] +async fn file_write_overwrites_existing() { + let dir = std::env::temp_dir().join("openhuman_test_file_write_overwrite"); + let _ = tokio::fs::remove_dir_all(&dir).await; + tokio::fs::create_dir_all(&dir).await.unwrap(); + tokio::fs::write(dir.join("exist.txt"), "old") + .await + .unwrap(); + + let tool = FileWriteTool::new(test_security(dir.clone())); + let result = tool + .execute(json!({"path": "exist.txt", "content": "new"})) + .await + .unwrap(); + assert!(!result.is_error); + + let content = tokio::fs::read_to_string(dir.join("exist.txt")) + .await + .unwrap(); + assert_eq!(content, "new"); + + let _ = tokio::fs::remove_dir_all(&dir).await; +} + +#[tokio::test] +async fn file_write_blocks_path_traversal() { + let dir = std::env::temp_dir().join("openhuman_test_file_write_traversal"); + let _ = tokio::fs::remove_dir_all(&dir).await; + tokio::fs::create_dir_all(&dir).await.unwrap(); + + let tool = FileWriteTool::new(test_security(dir.clone())); + let result = tool + .execute(json!({"path": "../../etc/evil", "content": "bad"})) + .await + .unwrap(); + assert!(result.is_error); + assert!(&result.output().contains("not allowed")); + + let _ = tokio::fs::remove_dir_all(&dir).await; +} + +#[tokio::test] +async fn file_write_blocks_absolute_path() { + let tool = FileWriteTool::new(test_security(std::env::temp_dir())); + let result = tool + .execute(json!({"path": "/etc/evil", "content": "bad"})) + .await + .unwrap(); + assert!(result.is_error); + assert!(&result.output().contains("not allowed")); +} + +#[tokio::test] +async fn file_write_missing_path_param() { + let tool = FileWriteTool::new(test_security(std::env::temp_dir())); + let result = tool.execute(json!({"content": "data"})).await; + assert!(result.is_err()); +} + +#[tokio::test] +async fn file_write_missing_content_param() { + let tool = FileWriteTool::new(test_security(std::env::temp_dir())); + let result = tool.execute(json!({"path": "file.txt"})).await; + assert!(result.is_err()); +} + +#[tokio::test] +async fn file_write_empty_content() { + let dir = std::env::temp_dir().join("openhuman_test_file_write_empty"); + let _ = tokio::fs::remove_dir_all(&dir).await; + tokio::fs::create_dir_all(&dir).await.unwrap(); + + let tool = FileWriteTool::new(test_security(dir.clone())); + let result = tool + .execute(json!({"path": "empty.txt", "content": ""})) + .await + .unwrap(); + assert!(!result.is_error); + assert!(result.output().contains("0 bytes")); + + let _ = tokio::fs::remove_dir_all(&dir).await; +} + +#[cfg(unix)] +#[tokio::test] +async fn file_write_blocks_symlink_escape() { + use std::os::unix::fs::symlink; + + let root = std::env::temp_dir().join("openhuman_test_file_write_symlink_escape"); + let workspace = root.join("workspace"); + let outside = root.join("outside"); + + let _ = tokio::fs::remove_dir_all(&root).await; + tokio::fs::create_dir_all(&workspace).await.unwrap(); + tokio::fs::create_dir_all(&outside).await.unwrap(); + + symlink(&outside, workspace.join("escape_dir")).unwrap(); + + let tool = FileWriteTool::new(test_security(workspace.clone())); + let result = tool + .execute(json!({"path": "escape_dir/hijack.txt", "content": "bad"})) + .await + .unwrap(); + + assert!(result.is_error); + // SecurityPolicy now blocks symlink escapes at the is_path_allowed + // layer (#1927) — error becomes "Path not allowed by security + // policy" rather than the deeper "escapes workspace" message. + let out = result.output(); + assert!( + out.contains("escapes workspace") || out.contains("not allowed"), + "expected escape/not-allowed error, got: {out}" + ); + assert!(!outside.join("hijack.txt").exists()); + + let _ = tokio::fs::remove_dir_all(&root).await; +} + +#[tokio::test] +async fn file_write_blocks_readonly_mode() { + let dir = std::env::temp_dir().join("openhuman_test_file_write_readonly"); + let _ = tokio::fs::remove_dir_all(&dir).await; + tokio::fs::create_dir_all(&dir).await.unwrap(); + + let tool = FileWriteTool::new(test_security_with(dir.clone(), AutonomyLevel::ReadOnly, 20)); + let result = tool + .execute(json!({"path": "out.txt", "content": "should-block"})) + .await + .unwrap(); + + assert!(result.is_error); + assert!(result.output().contains("read-only")); + // The readonly block must carry the hard-reject marker so the agent + // harness recognizes it and halts on a verbatim repeat instead of + // grinding. Ties this tool's literal to the marker const — the + // const→detector half lives in `RepeatedToolFailureMiddleware` and is + // covered by its tests (`crates/openhuman-core/src/agent/tinyagents/middleware.rs`). + assert!( + result + .output() + .contains(crate::fs::filesystem::test_support::POLICY_BLOCKED_MARKER), + "file_write readonly block must carry the hard-reject marker: {}", + result.output() + ); + assert!(!dir.join("out.txt").exists()); + + let _ = tokio::fs::remove_dir_all(&dir).await; +} + +#[tokio::test] +async fn file_write_blocks_when_rate_limited() { + let dir = std::env::temp_dir().join("openhuman_test_file_write_rate_limited"); + let _ = tokio::fs::remove_dir_all(&dir).await; + tokio::fs::create_dir_all(&dir).await.unwrap(); + + let tool = FileWriteTool::new(test_security_with( + dir.clone(), + AutonomyLevel::Supervised, + 0, + )); + let result = tool + .execute(json!({"path": "out.txt", "content": "should-block"})) + .await + .unwrap(); + + assert!(result.is_error); + assert!(result.output().contains("Rate limit exceeded")); + assert!(!dir.join("out.txt").exists()); + + let _ = tokio::fs::remove_dir_all(&dir).await; +} + +// ── §5.1 TOCTOU / symlink file write protection tests ──── + +#[cfg(unix)] +#[tokio::test] +async fn file_write_blocks_symlink_target_file() { + use std::os::unix::fs::symlink; + + let root = std::env::temp_dir().join("openhuman_test_file_write_symlink_target"); + let workspace = root.join("workspace"); + let outside = root.join("outside"); + + let _ = tokio::fs::remove_dir_all(&root).await; + tokio::fs::create_dir_all(&workspace).await.unwrap(); + tokio::fs::create_dir_all(&outside).await.unwrap(); + + // Create a file outside and symlink to it inside workspace + tokio::fs::write(outside.join("target.txt"), "original") + .await + .unwrap(); + symlink(outside.join("target.txt"), workspace.join("linked.txt")).unwrap(); + + let tool = FileWriteTool::new(test_security(workspace.clone())); + let result = tool + .execute(json!({"path": "linked.txt", "content": "overwritten"})) + .await + .unwrap(); + + assert!(result.is_error, "writing through symlink must be blocked"); + // The symlink-safe is_path_allowed check (#1927) blocks at the + // policy layer before the tool's own symlink-target detection + // runs; accept either error message. + let out = result.output(); + assert!( + out.contains("symlink") || out.contains("not allowed"), + "error should mention symlink or policy block, got: {out}" + ); + + // Verify original file was not modified + let content = tokio::fs::read_to_string(outside.join("target.txt")) + .await + .unwrap(); + assert_eq!(content, "original", "original file must not be modified"); + + let _ = tokio::fs::remove_dir_all(&root).await; +} + +#[tokio::test] +async fn file_write_blocks_null_byte_in_path() { + let dir = std::env::temp_dir().join("openhuman_test_file_write_null"); + let _ = tokio::fs::remove_dir_all(&dir).await; + tokio::fs::create_dir_all(&dir).await.unwrap(); + + let tool = FileWriteTool::new(test_security(dir.clone())); + let result = tool + .execute(json!({"path": "file\u{0000}.txt", "content": "bad"})) + .await + .unwrap(); + assert!(result.is_error, "paths with null bytes must be blocked"); + + let _ = tokio::fs::remove_dir_all(&dir).await; +} + +/// `file_write` enforces the same 5MB content-size cap as `edit_file` and +/// `grep`'s file caps (`file_read` allows up to 10MB for reads). +#[tokio::test] +async fn file_write_reports_a_refused_write_rather_than_a_silent_success() { + let dir = std::env::temp_dir().join("openhuman_test_file_write_refused"); + let _ = tokio::fs::remove_dir_all(&dir).await; + tokio::fs::create_dir_all(&dir).await.unwrap(); + + let tool = FileWriteTool::new(test_security(dir.clone())).with_sink(Arc::new( + crate::fs::filesystem::file_sink::RefusingSink(std::io::ErrorKind::PermissionDenied), + )); + let result = tool + .execute(json!({"path": "f.txt", "content": "abc"})) + .await + .unwrap(); + + assert!( + result.is_error, + "a refused write must surface as an error, not a fabricated success" + ); + assert!( + !dir.join("f.txt").exists(), + "and nothing may be left behind claiming the write happened" + ); + + let _ = tokio::fs::remove_dir_all(&dir).await; +} + +#[tokio::test] +async fn file_write_refuses_an_oversized_content_write() { + let dir = std::env::temp_dir().join("openhuman_test_file_write_oversized"); + let _ = tokio::fs::remove_dir_all(&dir).await; + tokio::fs::create_dir_all(&dir).await.unwrap(); + + let tool = FileWriteTool::new(test_security(dir.clone())); + let content = "x".repeat(64 * 1024 * 1024); + let result = tool + .execute(json!({"path": "huge.txt", "content": content})) + .await + .unwrap(); + assert!( + result.is_error, + "a 64MB write must be refused by a stated ceiling, not written: {}", + result.output() + ); + + let _ = tokio::fs::remove_dir_all(&dir).await; +} + +// ── Coverage of the guard, symlink and budget branches ────────────────── + +use crate::fs::filesystem::test_support::RacyGate; + +#[tokio::test] +async fn file_write_refuses_a_racing_budget() { + let dir = tempfile::tempdir().unwrap(); + let tool = FileWriteTool::new(Arc::new(RacyGate(test_security(dir.path().to_path_buf())))); + let result = tool + .execute(json!({"path": "a.txt", "content": "x"})) + .await + .unwrap(); + assert!(result.is_error); + assert!(result.output().contains("action budget exhausted")); +} + +#[cfg(unix)] +#[tokio::test] +async fn file_write_refuses_to_write_through_an_in_workspace_symlink() { + let dir = tempfile::tempdir().unwrap(); + std::fs::write(dir.path().join("real.txt"), "real").unwrap(); + std::os::unix::fs::symlink(dir.path().join("real.txt"), dir.path().join("link.txt")).unwrap(); + let tool = FileWriteTool::new(test_security(dir.path().to_path_buf())); + let result = tool + .execute(json!({"path": "link.txt", "content": "x"})) + .await + .unwrap(); + assert!(result.is_error); + assert!( + result + .output() + .contains("Refusing to write through symlink") + ); + assert_eq!( + std::fs::read_to_string(dir.path().join("real.txt")).unwrap(), + "real" + ); +} + +#[tokio::test] +async fn file_write_honours_the_file_state_guard_and_records_its_write() { + crate::fs::file_state::init_global(true); + let dir = tempfile::tempdir().unwrap(); + let target = dir.path().canonicalize().unwrap().join("g.txt"); + std::fs::write(&target, "old").unwrap(); + let tool = FileWriteTool::new(test_security(dir.path().to_path_buf())); + let agent = format!("fw-agent-{}", dir.path().display()); + let other = format!("fw-other-{}", dir.path().display()); + + // Partial read blocks the overwrite. + crate::fs::file_state::record_read( + &agent, + target.clone(), + std::time::SystemTime::now(), + true, + std::time::Instant::now(), + ); + let result = crate::fs::file_state::with_file_state_agent_id( + agent.clone(), + tool.execute(json!({"path": "g.txt", "content": "new"})), + ) + .await + .unwrap(); + assert!(result.is_error); + assert!(result.output().contains("Partial read")); + + // A sibling's later write makes the read stale. + crate::fs::file_state::record_read( + &agent, + target.clone(), + std::time::SystemTime::now(), + false, + std::time::Instant::now(), + ); + crate::fs::file_state::record_write(&other, target.clone()); + let result = crate::fs::file_state::with_file_state_agent_id( + agent.clone(), + tool.execute(json!({"path": "g.txt", "content": "new"})), + ) + .await + .unwrap(); + assert!(result.is_error); + assert!(result.output().contains("Stale read")); + + // With a fresh full read the write goes through and is recorded. + crate::fs::file_state::record_read( + &agent, + target.clone(), + std::time::SystemTime::now(), + false, + std::time::Instant::now(), + ); + let result = crate::fs::file_state::with_file_state_agent_id( + agent.clone(), + tool.execute(json!({"path": "g.txt", "content": "new"})), + ) + .await + .unwrap(); + assert!(!result.is_error, "{}", result.output()); + assert_eq!(std::fs::read_to_string(&target).unwrap(), "new"); + assert!(crate::fs::file_state::check_stale_read(&agent, &target).is_none()); +} + +#[tokio::test] +async fn file_write_uses_the_context_workspace_when_one_is_threaded() { + use crate::fs::filesystem::test_support::WorkspaceContext; + let home = tempfile::tempdir().unwrap(); + let isolated = tempfile::tempdir().unwrap(); + let tool = FileWriteTool::new(test_security(home.path().to_path_buf())); + let context = WorkspaceContext::at(isolated.path()); + let result = tool + .execute_with_context( + json!({"path": "in_ws.txt", "content": "hi"}), + ToolCallOptions::default(), + Some(&context), + ) + .await + .unwrap(); + assert!(!result.is_error, "{}", result.output()); + assert!(isolated.path().join("in_ws.txt").exists()); + assert!(!home.path().join("in_ws.txt").exists()); +} diff --git a/crates/tinyagents-harness-tools/src/fs/filesystem/fixtures/apply_patch.json b/crates/tinyagents-harness-tools/src/fs/filesystem/fixtures/apply_patch.json new file mode 100644 index 000000000..7b98f41ec --- /dev/null +++ b/crates/tinyagents-harness-tools/src/fs/filesystem/fixtures/apply_patch.json @@ -0,0 +1,46 @@ +{ + "description": "Apply a batch of exact-string edits across one or more files atomically. All edits are validated before any are written; validation failure rolls back the whole batch. Each edit is `{path, old_string, new_string, replace_all?}`; a top-level `path` is the default for edits that omit their own. To CREATE a new file, pass an empty `old_string` with the full contents as `new_string`; the path must not already exist.", + "exposure": "Direct", + "name": "apply_patch", + "permission_level": "Write", + "schema": { + "properties": { + "path": { + "description": "Default file for edits that omit their own `path`.", + "type": "string" + }, + "edits": { + "description": "Ordered list of edits.", + "items": { + "properties": { + "new_string": { + "type": "string" + }, + "old_string": { + "description": "Exact text to replace. Empty means CREATE: the path must not exist and `new_string` becomes the whole file.", + "type": "string" + }, + "path": { + "description": "File to edit; defaults to the top-level `path`.", + "type": "string" + }, + "replace_all": { + "default": false, + "type": "boolean" + } + }, + "required": [ + "old_string", + "new_string" + ], + "type": "object" + }, + "type": "array" + } + }, + "required": [ + "edits" + ], + "type": "object" + } +} diff --git a/crates/tinyagents-harness-tools/src/fs/filesystem/fixtures/csv_export.json b/crates/tinyagents-harness-tools/src/fs/filesystem/fixtures/csv_export.json new file mode 100644 index 000000000..dd4e4d5a0 --- /dev/null +++ b/crates/tinyagents-harness-tools/src/fs/filesystem/fixtures/csv_export.json @@ -0,0 +1,30 @@ +{ + "description": "Export structured data (JSON array of objects) as a CSV file to the workspace. Returns the file path. Use when the user wants raw tabular data from a tool result that's too large to include inline.", + "exposure": "Direct", + "name": "csv_export", + "permission_level": "Write", + "schema": { + "properties": { + "columns": { + "description": "Optional ordered list of column names to include. If omitted, all keys from the first object are used as headers.", + "items": { + "type": "string" + }, + "type": "array" + }, + "data": { + "description": "JSON string containing an array of objects to export. Each object becomes a row; keys become column headers.", + "type": "string" + }, + "filename": { + "description": "Output filename (without path). Will be written to workspace/exports/. Example: 'github-issues-2026-04-16.csv'", + "type": "string" + } + }, + "required": [ + "data", + "filename" + ], + "type": "object" + } +} diff --git a/crates/tinyagents-harness-tools/src/fs/filesystem/fixtures/edit_file.json b/crates/tinyagents-harness-tools/src/fs/filesystem/fixtures/edit_file.json new file mode 100644 index 000000000..b6d958175 --- /dev/null +++ b/crates/tinyagents-harness-tools/src/fs/filesystem/fixtures/edit_file.json @@ -0,0 +1,33 @@ +{ + "description": "Edit a file by exact string replacement. By default `old_string` must match exactly once. Set `replace_all` to true to replace every match.", + "exposure": "Deferred", + "name": "edit", + "permission_level": "Write", + "schema": { + "properties": { + "new_string": { + "description": "Replacement text.", + "type": "string" + }, + "old_string": { + "description": "Text to find. Must match exactly.", + "type": "string" + }, + "path": { + "description": "Workspace-relative file path.", + "type": "string" + }, + "replace_all": { + "default": false, + "description": "If true, replace every occurrence (default false).", + "type": "boolean" + } + }, + "required": [ + "path", + "old_string", + "new_string" + ], + "type": "object" + } +} diff --git a/crates/tinyagents-harness-tools/src/fs/filesystem/fixtures/file_read.json b/crates/tinyagents-harness-tools/src/fs/filesystem/fixtures/file_read.json new file mode 100644 index 000000000..fa9744fd3 --- /dev/null +++ b/crates/tinyagents-harness-tools/src/fs/filesystem/fixtures/file_read.json @@ -0,0 +1,23 @@ +{ + "description": "Read the contents of a file in your working directory (the action sandbox). Relative paths resolve against that directory; paths outside it are blocked. To read a file written by `shell`, confirm its location with `pwd` and use the same relative path. Long files are returned one page at a time; continue with the exact `offset` reported at the end of the page.", + "exposure": "Direct", + "name": "file_read", + "permission_level": "ReadOnly", + "schema": { + "properties": { + "offset": { + "description": "Byte offset to start reading from. Use the offset a paged read reports to continue it.", + "minimum": 0, + "type": "integer" + }, + "path": { + "description": "Relative path to the file within the workspace", + "type": "string" + } + }, + "required": [ + "path" + ], + "type": "object" + } +} diff --git a/crates/tinyagents-harness-tools/src/fs/filesystem/fixtures/file_write.json b/crates/tinyagents-harness-tools/src/fs/filesystem/fixtures/file_write.json new file mode 100644 index 000000000..8621d0bd7 --- /dev/null +++ b/crates/tinyagents-harness-tools/src/fs/filesystem/fixtures/file_write.json @@ -0,0 +1,23 @@ +{ + "description": "Write contents to a file in your working directory (the action sandbox). Relative paths resolve against that directory; writes outside it are blocked. Reference the file later by the same relative path so `file_read` resolves to it.", + "exposure": "Direct", + "name": "file_write", + "permission_level": "Write", + "schema": { + "properties": { + "content": { + "description": "Content to write to the file", + "type": "string" + }, + "path": { + "description": "Relative path to the file within the workspace", + "type": "string" + } + }, + "required": [ + "path", + "content" + ], + "type": "object" + } +} diff --git a/crates/tinyagents-harness-tools/src/fs/filesystem/fixtures/git_operations.json b/crates/tinyagents-harness-tools/src/fs/filesystem/fixtures/git_operations.json new file mode 100644 index 000000000..302536c1b --- /dev/null +++ b/crates/tinyagents-harness-tools/src/fs/filesystem/fixtures/git_operations.json @@ -0,0 +1,66 @@ +{ + "description": "Perform structured Git operations (status, diff, log, branch, commit, add, checkout, stash). Provides parsed JSON output and integrates with security policy for autonomy controls.", + "exposure": "Direct", + "name": "git_operations", + "permission_level": "ReadOnly", + "schema": { + "properties": { + "action": { + "description": "Stash action (for 'stash' operation)", + "enum": [ + "push", + "pop", + "list", + "drop" + ], + "type": "string" + }, + "branch": { + "description": "Branch name (for 'checkout' operation)", + "type": "string" + }, + "cached": { + "description": "Show staged changes (for 'diff' operation)", + "type": "boolean" + }, + "files": { + "description": "File or path to diff (for 'diff' operation, default: '.')", + "type": "string" + }, + "index": { + "description": "Stash index (for 'stash' with 'drop' action)", + "type": "integer" + }, + "limit": { + "description": "Number of log entries (for 'log' operation, default: 10)", + "type": "integer" + }, + "message": { + "description": "Commit message (for 'commit' operation)", + "type": "string" + }, + "operation": { + "description": "Git operation to perform", + "enum": [ + "status", + "diff", + "log", + "branch", + "commit", + "add", + "checkout", + "stash" + ], + "type": "string" + }, + "paths": { + "description": "File paths to stage (for 'add' operation)", + "type": "string" + } + }, + "required": [ + "operation" + ], + "type": "object" + } +} diff --git a/crates/tinyagents-harness-tools/src/fs/filesystem/fixtures/glob_search.json b/crates/tinyagents-harness-tools/src/fs/filesystem/fixtures/glob_search.json new file mode 100644 index 000000000..926937d5b --- /dev/null +++ b/crates/tinyagents-harness-tools/src/fs/filesystem/fixtures/glob_search.json @@ -0,0 +1,27 @@ +{ + "description": "Find files matching a glob pattern (e.g. `src/**/*.rs`), searching the action sandbox unless `path` says otherwise. Returns paths newest-first, ready to pass straight to `file_read` / `grep`.", + "exposure": "Direct", + "name": "glob", + "permission_level": "ReadOnly", + "schema": { + "properties": { + "max_results": { + "description": "Cap on returned paths (default 500).", + "minimum": 1, + "type": "integer" + }, + "path": { + "description": "Directory to search within. Defaults to the working directory (action sandbox). Must be a location you're allowed to read; relative paths resolve against the working directory.", + "type": "string" + }, + "pattern": { + "description": "Glob pattern, e.g. `**/*.rs` or `src/**/*.{ts,tsx}` (single brace expansion not supported — list patterns separately). Matched against paths relative to the search directory.", + "type": "string" + } + }, + "required": [ + "pattern" + ], + "type": "object" + } +} diff --git a/crates/tinyagents-harness-tools/src/fs/filesystem/fixtures/grep.json b/crates/tinyagents-harness-tools/src/fs/filesystem/fixtures/grep.json new file mode 100644 index 000000000..039c98fd6 --- /dev/null +++ b/crates/tinyagents-harness-tools/src/fs/filesystem/fixtures/grep.json @@ -0,0 +1,33 @@ +{ + "description": "Search file contents in the workspace with a regular expression. Returns up to `max_matches` matches (default 200) as `path:line:text` lines.", + "exposure": "Direct", + "name": "grep", + "permission_level": "ReadOnly", + "schema": { + "properties": { + "case_insensitive": { + "default": false, + "description": "If true, compile the regex with the `i` flag.", + "type": "boolean" + }, + "max_matches": { + "description": "Cap on matches returned (default 200).", + "minimum": 1, + "type": "integer" + }, + "path": { + "default": ".", + "description": "Optional sub-path to restrict the search to (relative to workspace).", + "type": "string" + }, + "pattern": { + "description": "Regular expression (Rust `regex` crate syntax).", + "type": "string" + } + }, + "required": [ + "pattern" + ], + "type": "object" + } +} diff --git a/crates/tinyagents-harness-tools/src/fs/filesystem/fixtures/image_info.json b/crates/tinyagents-harness-tools/src/fs/filesystem/fixtures/image_info.json new file mode 100644 index 000000000..3125c3d19 --- /dev/null +++ b/crates/tinyagents-harness-tools/src/fs/filesystem/fixtures/image_info.json @@ -0,0 +1,22 @@ +{ + "description": "Read image file metadata (format, dimensions, size) and optionally return base64-encoded data.", + "exposure": "Direct", + "name": "image_info", + "permission_level": "ReadOnly", + "schema": { + "properties": { + "include_base64": { + "description": "Include base64-encoded image data in output (default: false)", + "type": "boolean" + }, + "path": { + "description": "Path to the image file (absolute or relative to workspace)", + "type": "string" + } + }, + "required": [ + "path" + ], + "type": "object" + } +} diff --git a/crates/tinyagents-harness-tools/src/fs/filesystem/fixtures/list_files.json b/crates/tinyagents-harness-tools/src/fs/filesystem/fixtures/list_files.json new file mode 100644 index 000000000..98415d486 --- /dev/null +++ b/crates/tinyagents-harness-tools/src/fs/filesystem/fixtures/list_files.json @@ -0,0 +1,16 @@ +{ + "description": "List entries in a workspace directory (non-recursive). Each line is `\t` where kind is `dir`, `file`, or `link`.", + "exposure": "Direct", + "name": "list", + "permission_level": "ReadOnly", + "schema": { + "properties": { + "path": { + "default": ".", + "description": "Directory path relative to the workspace (default `.`).", + "type": "string" + } + }, + "type": "object" + } +} diff --git a/crates/tinyagents-harness-tools/src/fs/filesystem/fixtures/read_diff.json b/crates/tinyagents-harness-tools/src/fs/filesystem/fixtures/read_diff.json new file mode 100644 index 000000000..f32196a0d --- /dev/null +++ b/crates/tinyagents-harness-tools/src/fs/filesystem/fixtures/read_diff.json @@ -0,0 +1,23 @@ +{ + "description": "Get the git diff of current changes. Can diff staged, unstaged, or against a specific base branch/commit. Returns file paths and hunks.", + "exposure": "Deferred", + "name": "read_diff", + "permission_level": "ReadOnly", + "schema": { + "properties": { + "base": { + "description": "Base ref to diff against (e.g. 'main', 'HEAD~3'). Default: unstaged changes.", + "type": "string" + }, + "path_filter": { + "description": "Limit diff to a specific path or glob.", + "type": "string" + }, + "staged": { + "description": "Show staged changes only (--cached). Default: false.", + "type": "boolean" + } + }, + "type": "object" + } +} diff --git a/crates/tinyagents-harness-tools/src/fs/filesystem/fixtures/read_workspace_state.json b/crates/tinyagents-harness-tools/src/fs/filesystem/fixtures/read_workspace_state.json new file mode 100644 index 000000000..962bb1bb1 --- /dev/null +++ b/crates/tinyagents-harness-tools/src/fs/filesystem/fixtures/read_workspace_state.json @@ -0,0 +1,21 @@ +{ + "description": "Get a read-only overview of the workspace: git status (modified/untracked files), recent commits, and top-level directory structure. Useful for understanding the current project state before planning tasks.", + "exposure": "Direct", + "name": "read_workspace_state", + "permission_level": "ReadOnly", + "schema": { + "properties": { + "include_tree": { + "default": true, + "description": "Include top-level directory tree (default: true).", + "type": "boolean" + }, + "recent_commits": { + "default": 5, + "description": "Number of recent commits to show (default: 5).", + "type": "integer" + } + }, + "type": "object" + } +} diff --git a/crates/tinyagents-harness-tools/src/fs/filesystem/fixtures/run_linter.json b/crates/tinyagents-harness-tools/src/fs/filesystem/fixtures/run_linter.json new file mode 100644 index 000000000..4c16bbdfa --- /dev/null +++ b/crates/tinyagents-harness-tools/src/fs/filesystem/fixtures/run_linter.json @@ -0,0 +1,25 @@ +{ + "description": "Run linting tools on the codebase. Supports 'clippy' for Rust and 'eslint' for TypeScript/JavaScript. Returns warnings and errors.", + "exposure": "Deferred", + "name": "run_linter", + "permission_level": "Execute", + "schema": { + "properties": { + "linter": { + "default": "auto", + "description": "Which linter to run. 'auto' detects from project files.", + "enum": [ + "clippy", + "eslint", + "auto" + ], + "type": "string" + }, + "path": { + "description": "Limit linting to a specific path.", + "type": "string" + } + }, + "type": "object" + } +} diff --git a/crates/tinyagents-harness-tools/src/fs/filesystem/fixtures/run_tests.json b/crates/tinyagents-harness-tools/src/fs/filesystem/fixtures/run_tests.json new file mode 100644 index 000000000..f0f5dc9f1 --- /dev/null +++ b/crates/tinyagents-harness-tools/src/fs/filesystem/fixtures/run_tests.json @@ -0,0 +1,30 @@ +{ + "description": "Run the project test suite. Supports 'cargo_test' for Rust and 'vitest' for TypeScript/JavaScript. Returns pass/fail results with output.", + "exposure": "Deferred", + "name": "run_tests", + "permission_level": "Execute", + "schema": { + "properties": { + "filter": { + "description": "Filter to run specific tests (e.g. test name or module).", + "type": "string" + }, + "runner": { + "default": "auto", + "description": "Which test runner to use. 'auto' detects from project files.", + "enum": [ + "cargo_test", + "vitest", + "auto" + ], + "type": "string" + }, + "timeout_secs": { + "default": 120, + "description": "Timeout in seconds (default: 120).", + "type": "integer" + } + }, + "type": "object" + } +} diff --git a/crates/tinyagents-harness-tools/src/fs/filesystem/fixtures/update_memory_md.json b/crates/tinyagents-harness-tools/src/fs/filesystem/fixtures/update_memory_md.json new file mode 100644 index 000000000..16237185a --- /dev/null +++ b/crates/tinyagents-harness-tools/src/fs/filesystem/fixtures/update_memory_md.json @@ -0,0 +1,40 @@ +{ + "description": "Append or update sections in MEMORY.md or SKILL.md workspace files. Use 'append' to add new notes at the end, or 'replace_section' to overwrite the body under a named '## Section' heading.", + "exposure": "Direct", + "name": "update_memory_md", + "permission_level": "Write", + "schema": { + "properties": { + "action": { + "description": "'append' adds content at the end; 'replace_section' replaces the body of the named section.", + "enum": [ + "append", + "replace_section" + ], + "type": "string" + }, + "content": { + "description": "The markdown text to write.", + "type": "string" + }, + "file": { + "description": "Which workspace markdown file to modify.", + "enum": [ + "MEMORY.md", + "SKILL.md" + ], + "type": "string" + }, + "section_title": { + "description": "Required for 'replace_section': the heading text (without '## ').", + "type": "string" + } + }, + "required": [ + "file", + "action", + "content" + ], + "type": "object" + } +} diff --git a/crates/tinyagents-harness-tools/src/fs/filesystem/gate.rs b/crates/tinyagents-harness-tools/src/fs/filesystem/gate.rs new file mode 100644 index 000000000..ca9965e94 --- /dev/null +++ b/crates/tinyagents-harness-tools/src/fs/filesystem/gate.rs @@ -0,0 +1,129 @@ +//! The host's say over every filesystem tool call. +//! +//! The tools in this module read, write, search and patch files, but they do +//! not decide what they are *allowed* to touch. That decision (an autonomy +//! level, a workspace boundary, an approval prompt, an hourly action budget, +//! credential stores that are never reachable) belongs to the host whose threat +//! model and configuration it depends on. A tool asks an [`FsGate`] instead. +//! +//! The trait is exactly the set of questions the tools put to a policy, and +//! nothing wider: it carries no decision *types* of its own, only booleans and +//! resolved paths, so a host maps its own policy onto it without translating a +//! vocabulary. + +use std::path::{Path, PathBuf}; +use std::sync::Arc; + +use async_trait::async_trait; +use tinytools::ToolRunContext; + +/// The host policy a filesystem tool consults before it acts. +/// +/// Implementations must be cheap to call: the tools ask on every invocation. +/// All methods take `&self`; an implementation that counts actions keeps that +/// state behind interior mutability. +#[async_trait] +pub trait FsGate: std::fmt::Debug + Send + Sync { + /// Whether the host permits the agent to change anything at all right now. + /// + /// `false` blocks every mutating tool before it touches the disk. + fn can_act(&self) -> bool; + + /// Whether the host is in a strictly read-only mode. + /// + /// Asked by `git_operations`, on top of [`can_act`](Self::can_act), before + /// a write operation runs. + fn is_read_only(&self) -> bool; + + /// Whether the action budget is already spent, without consuming any of it. + fn is_rate_limited(&self) -> bool; + + /// Consume one unit of the action budget. + /// + /// Returns `false` when the call is over budget and must be refused. + fn record_action(&self) -> bool; + + /// Whether a write must be confirmed by a human before it runs. + /// + /// Drives the tools' `external_effect_with_args` flag, which routes the + /// call through the host's approval flow. + fn write_needs_approval(&self) -> bool; + + /// The directory relative paths resolve against. + fn action_dir(&self) -> &Path; + + /// Cheap, string-level check of whether `path` may be used at all. + /// + /// Used to filter search hits without touching the disk. + fn is_path_string_allowed(&self, path: &str) -> bool; + + /// Batch form of [`is_path_string_allowed`](Self::is_path_string_allowed). + /// + /// Answers are in request order and the result has the same length as + /// `paths`. The default asks the single-path question once per entry, so + /// every existing gate behaves exactly as before. A host that can answer a + /// whole batch more cheaply (one round trip to a policy module instead of + /// one per path) overrides this. + /// + /// A caller must treat an answer whose length differs from `paths` as + /// every path refused. + fn paths_string_allowed(&self, paths: &[String]) -> Vec { + paths + .iter() + .map(|path| self.is_path_string_allowed(path)) + .collect() + } + + /// Resolve an existing path and confirm the host allows it. + /// + /// # Errors + /// + /// Returns the message shown to the model when the path is refused or + /// cannot be resolved. + async fn validate_path(&self, path: &str) -> Result; + + /// Resolve a path whose *parent* must exist but whose target may not, and + /// confirm the host allows it. Used for writes that create files. + /// + /// # Errors + /// + /// Returns the message shown to the model when the path is refused or + /// cannot be resolved. + async fn validate_parent_path(&self, path: &str) -> Result; + + /// A gate for the same policy, rooted at `root`. + /// + /// The run's isolated workspace becomes both the directory relative paths + /// resolve against and a readable-and-writable root. The grant is + /// per-call: it must not mutate `self`, so concurrent turns cannot race + /// each other, and it must not widen whatever the host treats as always + /// forbidden. `root` originates from trusted in-process code (the session + /// builder or a sub-agent runner), never from model-supplied text. + fn scoped_to_workspace(&self, root: &Path) -> Arc; +} + +/// Resolve the gate a tool should validate paths against for this call. +/// +/// When the harness threaded a workspace descriptor into the call (an +/// edit-capable worker running in an isolated checkout), the returned gate is +/// scoped to that workspace root; otherwise it is `gate` itself. Counting and +/// approval questions ([`FsGate::can_act`], [`FsGate::record_action`], ...) +/// are still asked of the unscoped gate. +pub(crate) fn gate_for_context( + gate: &Arc, + context: Option<&dyn ToolRunContext>, + tool: &str, +) -> Arc { + match context.and_then(ToolRunContext::workspace) { + Some(workspace) => { + tracing::debug!( + tool, + workspace_root = %workspace.root.display(), + policy_id = %workspace.policy_id, + "[tools:filesystem] granting TinyAgents workspace descriptor as action dir + trusted root" + ); + gate.scoped_to_workspace(&workspace.root) + } + None => Arc::clone(gate), + } +} diff --git a/crates/tinyagents-harness-tools/src/fs/filesystem/gate_tests.rs b/crates/tinyagents-harness-tools/src/fs/filesystem/gate_tests.rs new file mode 100644 index 000000000..727a3b049 --- /dev/null +++ b/crates/tinyagents-harness-tools/src/fs/filesystem/gate_tests.rs @@ -0,0 +1,95 @@ +//! Tests for how a call's workspace descriptor scopes the gate. + +#![allow(clippy::unwrap_used)] + +use std::sync::Arc; + +use super::gate::{FsGate, gate_for_context}; +use super::test_support::{TestGate, WorkspaceContext}; + +#[tokio::test] +async fn without_a_descriptor_the_same_gate_is_used() { + let dir = tempfile::tempdir().unwrap(); + let gate: Arc = TestGate::supervised(dir.path().to_path_buf()); + let resolved = gate_for_context(&gate, None, "file_read"); + assert!(Arc::ptr_eq(&gate, &resolved)); +} + +#[tokio::test] +async fn a_descriptor_roots_the_gate_at_the_isolated_workspace() { + let home = tempfile::tempdir().unwrap(); + let isolated = tempfile::tempdir().unwrap(); + std::fs::write(isolated.path().join("seen.txt"), "x").unwrap(); + let gate: Arc = TestGate::supervised(home.path().to_path_buf()); + + let context = WorkspaceContext::at(isolated.path()); + let scoped = gate_for_context(&gate, Some(&context), "file_read"); + + assert_eq!(scoped.action_dir(), isolated.path()); + let resolved = scoped.validate_path("seen.txt").await.unwrap(); + assert_eq!( + resolved, + isolated.path().join("seen.txt").canonicalize().unwrap() + ); + // The original gate keeps its own root and refuses the isolated file. + assert!(gate.validate_path("seen.txt").await.is_err()); +} + +#[tokio::test] +async fn the_fake_gates_delegate_every_question_to_the_gate_they_wrap() { + use super::test_support::{AutonomyLevel, RacyGate, WrapGate}; + + let dir = tempfile::tempdir().unwrap(); + std::fs::write(dir.path().join("f.txt"), "x").unwrap(); + let inner = TestGate::with(dir.path().to_path_buf(), AutonomyLevel::Supervised, 5); + + let racy = RacyGate(inner.clone()); + let wrap = WrapGate::new(inner.clone()).arc(); + let gates: [&dyn FsGate; 2] = [&racy, wrap.as_ref()]; + for gate in gates { + assert!(gate.can_act()); + assert!(!gate.is_read_only()); + assert!(gate.write_needs_approval()); + assert_eq!(gate.action_dir(), dir.path()); + assert!(gate.is_path_string_allowed("f.txt")); + assert!(gate.validate_path("f.txt").await.is_ok()); + assert!(gate.validate_parent_path("new.txt").await.is_ok()); + let scoped = gate.scoped_to_workspace(dir.path()); + assert_eq!(scoped.action_dir(), dir.path()); + } +} + +#[tokio::test] +async fn the_fake_gate_refuses_what_a_real_policy_would() { + let dir = tempfile::tempdir().unwrap(); + let gate = TestGate::supervised(dir.path().to_path_buf()); + + // No parent / no file name to resolve. + assert!(gate.validate_parent_path("/").await.is_err()); + // Nothing on disk to resolve, even at the filesystem root. + let far = TestGate::supervised(std::path::PathBuf::from("/definitely/not/a/real/root")); + assert!(!far.is_path_string_allowed("x.txt")); + assert!(far.validate_parent_path("x.txt").await.is_err()); + // A NUL byte is never a path. + assert!(!gate.is_path_string_allowed("a\0b")); +} + +#[test] +fn the_default_batch_answer_maps_the_single_path_check_in_order() { + let dir = tempfile::tempdir().unwrap(); + let gate = TestGate::supervised(dir.path().to_path_buf()); + let paths = vec![ + "ok.txt".to_string(), + "../escape.txt".to_string(), + "a\0b".to_string(), + "sub/ok.rs".to_string(), + ]; + let batch = gate.paths_string_allowed(&paths); + let single: Vec = paths + .iter() + .map(|p| gate.is_path_string_allowed(p)) + .collect(); + assert_eq!(batch, single); + assert_eq!(batch, vec![true, false, false, true]); + assert_eq!(gate.paths_string_allowed(&[]), Vec::::new()); +} diff --git a/crates/tinyagents-harness-tools/src/fs/filesystem/git_operations/config.rs b/crates/tinyagents-harness-tools/src/fs/filesystem/git_operations/config.rs new file mode 100644 index 000000000..ebc2b8424 --- /dev/null +++ b/crates/tinyagents-harness-tools/src/fs/filesystem/git_operations/config.rs @@ -0,0 +1,376 @@ +//! Ambient-git-config hardening for [`super::git_operations`]. +//! +//! Split out of `git_operations.rs` for the Rust layout gate. Cohesive on its +//! own: the allow-list, the neutralised keys, and the two helpers that build a +//! hardened `git` invocation are one policy, and reviewing that policy is +//! easier when it is not interleaved with the command implementations. + +use std::path::Path; + +/// Repository config keys this tool will run `git` under. +/// +/// This is an allowlist, and the direction is the whole point. Several git +/// config keys name a command that git then executes — `core.fsmonitor` is +/// run by `git status` (used by the `status` operation here), `core.sshCommand` +/// by anything that reaches a remote, `diff.external` by `diff`, `core.pager` +/// and `core.editor` by the commands that use them, and the `filter.*.process` +/// / `*.clean` / `*.smudge` and `diff.*.textconv` families by content +/// operations. Enumerating *those* and clearing them is the obvious fix and it +/// is the wrong shape: the list is only correct until git adds a key, and a +/// denylist that has gone stale reads as protection while providing none. +/// +/// An allowlist ages the other way. A key nobody here has heard of is refused, +/// so a new git release makes this tool fail closed and loud rather than +/// silently regain the hole. +/// +/// The entries are `section.key`, lowercased, with any subsection elided — +/// `remote.origin.url` is checked as `remote.url`. +pub(crate) const ALLOWED_REPO_CONFIG: &[&str] = &[ + // What `git init` and `git clone` write, and nothing else. + // + // `core.worktree` is deliberately absent, unlike the read-only sibling + // this list started from. It redirects the working-tree root Git + // operates against — including a linked-worktree-shaped redirect this + // tool does not otherwise offer — and this tool runs *write* operations + // (`checkout`, `add`, `commit`, `stash`) whose target directory is + // supposed to be `action_dir`/the resolved workspace, never something a + // repository's own config gets to name. Worktree isolation is already + // handled through `WorkspaceDescriptor` (`effective_action_dir_for_context`), + // a trusted, in-process mechanism — nothing here needs the config key. + "core.repositoryformatversion", + "core.filemode", + "core.bare", + "core.logallrefupdates", + "core.ignorecase", + "core.precomposeunicode", + "core.symlinks", + "remote.url", + "remote.fetch", + "remote.pushurl", + "remote.mirror", + "branch.remote", + "branch.merge", + "branch.rebase", + "submodule.active", + "submodule.url", + "user.name", + "user.email", + "pull.rebase", + "push.default", + "init.defaultbranch", + // Inert settings ordinary repositories carry that a first-draft allowlist + // would refuse, making the tool useless on a large class of real + // workspaces. Each is a value git *interprets*; none names a program git + // runs. + "core.autocrlf", + "core.eol", + "core.untrackedcache", + "core.longpaths", + "core.fscache", + "core.hidedotfiles", + "core.sparsecheckout", + "core.sparsecheckoutcone", + "commit.gpgsign", + "tag.gpgsign", + "remote.tagopt", + "remote.prune", + "remote.partialclonefilter", + "remote.promisor", + "branch.vscodemerge", + "gc.auto", + "fetch.prune", + // SHA-256 repositories and worktree-scoped config. Only these two + // `extensions.*` keys — the namespace as a whole is where git puts + // repository-format switches, and a blanket allow would admit whatever it + // adds next. + "extensions.objectformat", + "extensions.worktreeconfig", + // `filter..required` is a boolean. The driver's actual programs — + // `clean`, `smudge`, `process` — are NOT here and must not be; see the LFS + // note on `NEUTRALISED_CONFIG`. + "filter.required", + "lfs.repositoryformatversion", +]; + +/// Command-valued keys cleared on the command line as a second layer. +/// +/// Command-line `-c` outranks every config file, so this genuinely +/// neutralises these keys even when a repository sets them — and it does so +/// at the moment `git` actually runs, not at the moment +/// [`first_disallowed_repo_config_key`] happened to inspect the config a +/// moment earlier. That distinction matters: the inspection and the real +/// command are two separate `git` invocations, so a key set in the gap +/// between them (a second concurrent writer, or a worktree-scoped write) is +/// invisible to the inspection but still reaches here — where it is +/// neutralised regardless of timing. It is a denylist and therefore cannot be +/// the *only* guarantee — [`ALLOWED_REPO_CONFIG`] is the one that fails +/// closed on a key nobody anticipated — but this layer is what actually holds +/// under a race, not the inspection step. +/// +/// `credential.helper` is command-valued — a value beginning `!` is run as a +/// shell command — so it belongs nowhere near [`ALLOWED_REPO_CONFIG`] despite +/// reading like a mere preference; it is refused there instead, since none of +/// this tool's operations reach a remote and so have no legitimate use for it +/// to preserve. +/// +/// A `git lfs install` clone is refused, deliberately: `filter.lfs.clean`, +/// `.smudge` and `.process` each name a program, so an LFS working copy +/// cannot be read by this tool. That is the fail-closed answer and it is the +/// intended one. Only `filter..required`, a boolean, is allowed. +/// +/// `core.hooksPath` and `commit.gpgSign` are handled separately, in +/// [`hardened_git`] — see there for why. +pub(crate) const NEUTRALISED_CONFIG: &[&str] = &[ + "core.fsmonitor=", + "core.sshCommand=", + "core.pager=cat", + "core.editor=false", + // NOT `diff.external=`. An empty value does not disable an external diff — + // git tries to *execute* the empty string and the whole command dies with + // `error: cannot run : No such file or directory` / `fatal: external diff + // died`, so every `diff` operation failed rather than being hardened. + // Suppression belongs on the command instead: `git diff --no-ext-diff`, + // which ignores `diff.external` however the repository set it. Verified + // both ways against a repo with `diff.external=/bin/false`: plain `diff` + // dies, `--no-ext-diff` prints the patch. + "sequence.editor=false", + "uploadpack.packObjectsHook=", +]; + +/// The entries of [`NEUTRALISED_CONFIG`] that are also safe to force on an +/// *arbitrary* `git` invocation, plus `commit.gpgSign`. +/// +/// The `shell` tool runs git too, and since the `files` tool pack withheld +/// `git_operations` it runs most of it. Nothing there controls argv, so the +/// `-c` layer [`hardened_git`] builds is unavailable; the same overrides ride +/// `GIT_CONFIG_PARAMETERS` instead, which git documents as the mechanism `-c` +/// itself uses to reach child processes, at the same precedence. +/// +/// **This is the denylist half of the policy, not the fail-closed half.** +/// [`ALLOWED_REPO_CONFIG`] cannot come with it: inspecting the repository +/// config before every shell command would cost a `git` round trip on commands +/// that have nothing to do with git, and would refuse to run in an ordinary +/// LFS checkout. So this narrows the surface; it does not close it the way +/// `git_operations` does. +/// +/// Three entries of [`NEUTRALISED_CONFIG`] are deliberately **not** here: +/// +/// - `core.sshCommand=`. It is the same trap the `diff.external` note above +/// records: an empty value is not "use the default", it is a command git +/// tries to execute. `git_operations` gets away with listing it because none +/// of its operations reach a remote; a shell `git fetch` / `git push` does, +/// and would die on every one. Excluded rather than risked. +/// - `diff.external`. Unsuppressable by value for the same reason, and the fix +/// `git_operations` uses — `--no-ext-diff` — is an argv change this path +/// cannot make. A repository that sets it keeps it, and that is a known gap. +/// - `core.hooksPath`. There is no portable inert value for it, and overriding +/// it would change ordinary shell behavior by suppressing repository hooks. +/// Shell-invoked git therefore keeps the repository's configured hooks. +pub(crate) const SHELL_NEUTRALISED_CONFIG: &[&str] = &[ + "core.fsmonitor=", + "core.pager=cat", + "core.editor=false", + "sequence.editor=false", + "uploadpack.packObjectsHook=", + // Not in `NEUTRALISED_CONFIG` — `hardened_git` adds it separately, for the + // reason written there: a repository must not get to spend the host's + // signing key on a commit nobody asked to sign. + "commit.gpgSign=false", +]; + +/// Environment overrides applied to arbitrary commands launched by the shell +/// tool so any nested `git` process receives the safe subset above. +/// +/// Preserve ambient `GIT_CONFIG_PARAMETERS` exactly, including non-UTF-8 +/// bytes, because git uses it to propagate a parent's `-c` settings. The +/// safety overrides are appended so they win when the same key is repeated. +#[must_use] +pub fn shell_git_env() -> std::collections::HashMap { + let mut parameters = std::env::var_os("GIT_CONFIG_PARAMETERS").unwrap_or_default(); + for entry in SHELL_NEUTRALISED_CONFIG { + let Some((key, value)) = entry.split_once('=') else { + continue; + }; + if !parameters.is_empty() { + parameters.push(" "); + } + parameters.push(format!("'{key}'='{value}'")); + } + + std::collections::HashMap::from([( + std::ffi::OsString::from("GIT_CONFIG_PARAMETERS"), + parameters, + )]) +} + +/// A path git will read as an empty config file. +/// +/// `GIT_CONFIG_GLOBAL` must name something readable-and-empty rather than be +/// unset — unsetting it lets git fall back to `~/.gitconfig`, which is the +/// thing being suppressed. `/dev/null` is not a path on Windows; `NUL` is. +pub(crate) const NULL_CONFIG_PATH: &str = if cfg!(windows) { "NUL" } else { "/dev/null" }; + +/// Close the system/global config files and the command-valued `GIT_*` env +/// vars on a `git` invocation, without touching anything about how it reads +/// the repository's own config. +/// +/// `GIT_CONFIG_NOSYSTEM` and `GIT_CONFIG_GLOBAL` close the system and global +/// config files. Note what they do *not* close: the repository's own local +/// and worktree-scoped config, which is what an agent-writable workspace +/// actually lets an attacker author. That is handled by +/// [`first_disallowed_repo_config_key`] — a separate step, precisely because +/// [`hardened_git`]'s `-c` layer below must not be present while that step is +/// reading what the repository itself set (see its own doc comment). +pub(crate) fn suppress_ambient_git_config( + cmd: &mut tokio::process::Command, +) -> &mut tokio::process::Command { + cmd.env("GIT_CONFIG_NOSYSTEM", "1") + .env("GIT_CONFIG_GLOBAL", NULL_CONFIG_PATH) + // A parent `git -c` invocation propagates its command-line settings + // through these variables. They outrank every config file, so leaving + // them intact would let an ambient `core.hooksPath` or any other + // command-valued setting bypass both the file suppression and the + // repository allowlist below. + .env_remove("GIT_CONFIG_PARAMETERS") + .env_remove("GIT_CONFIG_COUNT") + // `git` consults these before it reads any config file. + .env_remove("GIT_EXTERNAL_DIFF") + .env_remove("GIT_PAGER") + .env_remove("GIT_EDITOR") + .env_remove("GIT_SSH") + .env_remove("GIT_SSH_COMMAND") + .env_remove("GIT_ALTERNATE_OBJECT_DIRECTORIES") + .env("GIT_TERMINAL_PROMPT", "0") +} + +/// Build the `git` invocation actually used to run a requested operation. +/// +/// Layers [`suppress_ambient_git_config`] under the [`NEUTRALISED_CONFIG`] +/// `-c` overrides, which outrank every config file the repository itself +/// could set — including one set in the gap between +/// [`first_disallowed_repo_config_key`]'s inspection and this invocation. +/// +/// Two more `-c` overrides are added here rather than in the static +/// [`NEUTRALISED_CONFIG`] list, because their safe value is not a fixed +/// literal: +/// +/// - `core.hooksPath=` pointed at [`NULL_CONFIG_PATH`]. `commit` and +/// `checkout` are the two operations this tool exposes that run hooks +/// (`pre-commit`/`commit-msg`/`post-commit`, `post-checkout`), and a +/// repository-writable `core.hooksPath` naming a directory with an +/// executable `pre-commit` in it is exactly the shape of the worktree-scoped +/// bypass this hardening exists to close — verified directly: pointing +/// `core.hooksPath` at a directory containing a `pre-commit` that touches a +/// marker file, then running with `-c core.hooksPath=`, leaves +/// the marker untouched. A previous version of this comment claimed there +/// was no portable value that meant "nowhere"; there is — the same one +/// [`suppress_ambient_git_config`] already uses for `GIT_CONFIG_GLOBAL`, +/// since a location with nothing at it is exactly what git needs it to be. +/// - `commit.gpgSign=false`. `commit.gpgsign` is on [`ALLOWED_REPO_CONFIG`] as +/// an ordinary boolean, but a repository could still set it to force every +/// commit through this tool to be GPG-signed — with whatever real signing +/// key the *host* happens to have configured, silently attributing a +/// cryptographic signature to a commit the operator did not ask this tool to +/// sign. Overriding it here removes that decision from the repository +/// entirely, the same way `core.editor=false` removes commit message +/// editing from it in [`NEUTRALISED_CONFIG`] above. +pub(crate) fn hardened_git(dir: &Path) -> tokio::process::Command { + let mut cmd = tokio::process::Command::new("git"); + suppress_ambient_git_config(&mut cmd).current_dir(dir); + for kv in NEUTRALISED_CONFIG { + cmd.arg("-c").arg(kv); + } + cmd.arg("-c") + .arg(format!("core.hooksPath={NULL_CONFIG_PATH}")) + .arg("-c") + .arg("commit.gpgSign=false"); + cmd +} + +/// Normalise a `git config --list` key to the `section.key` form +/// [`ALLOWED_REPO_CONFIG`] uses, dropping any subsection. +/// +/// `remote.origin.url` → `remote.url`; `core.filemode` → `core.filemode`. A +/// subsection may itself contain dots (`includeIf.gitdir:~/x.y/.path`), so the +/// first and last components are the reliable ones. +pub(crate) fn normalise_config_key(key: &str) -> String { + let key = key.to_ascii_lowercase(); + match (key.find('.'), key.rfind('.')) { + (Some(first), Some(last)) if first != last => { + format!("{}.{}", &key[..first], &key[last + 1..]) + } + _ => key, + } +} + +/// Returns the first repository config key at `dir` that is not on +/// [`ALLOWED_REPO_CONFIG`], or `None` if every key it sets is recognised. +/// +/// Deliberately **not** `--local`: when `extensions.worktreeConfig` is set — +/// itself on [`ALLOWED_REPO_CONFIG`] as an ordinary, non-command-valued +/// setting — git additionally reads `config.worktree`, a second file `--local` +/// does not cover. A `core.hooksPath` set with `git config --worktree ...` +/// lands there, is invisible to `--local`, and is exactly the kind of key this +/// check exists to catch: `commit` runs the hook it names. Bare `--list` +/// returns the same merged view `git` itself consults — local *and* +/// worktree-scoped — with system and global excluded by +/// [`suppress_ambient_git_config`] instead of by a location flag. Verified +/// directly: `git config --worktree core.hooksPath ...` followed by +/// `--local --null` omits it; the same followed by a bare `--list --null` +/// (system/global suppressed) reports it. +/// +/// This step must run through [`suppress_ambient_git_config`] alone, **not** +/// [`hardened_git`]: the latter's `-c` layer would inject the very keys this +/// check inspects for (`core.fsmonitor=`, `core.pager=cat`, …), none of which +/// are themselves on [`ALLOWED_REPO_CONFIG`], and every invocation would +/// refuse itself. Reading the config this way also does not consult +/// `core.fsmonitor` or spawn a pager when its output is captured, so this +/// inspection step does not have the property it is checking for. +/// +/// A non-zero exit here is treated as a refusal, not as "nothing to +/// distrust": by the time this runs, the caller (`execute_in_context`) has +/// already confirmed `dir` — or one of its parents — contains a `.git`, so +/// `git config --list` failing means the config could not be read, not that +/// there is none. Proceeding to run the real command against config this step +/// never actually inspected would defeat the point of inspecting it first. +/// The refusal text for a repository whose config carries `key`. +/// +/// Lives here so the two paths that produce it — the guard inside +/// `run_git_command_in` and the repository probe in `execute_in_context`, +/// which reaches this conclusion without ever getting as far as the guard — +/// cannot word the same refusal differently. +pub(crate) fn disallowed_config_refusal(dir: &Path, key: &str) -> String { + format!( + "refusing to run git in {}: its repository config sets `{key}`, which is \ + not on the allowlist of configuration this tool will run under. \ + Several git config keys name a command git then executes, and this \ + directory is agent-writable, so unrecognised configuration is treated \ + as untrusted rather than honoured.", + dir.display() + ) +} + +pub(crate) async fn first_disallowed_repo_config_key(dir: &Path) -> anyhow::Result> { + let mut cmd = tokio::process::Command::new("git"); + suppress_ambient_git_config(&mut cmd).current_dir(dir); + let output = cmd.args(["config", "--list", "--null"]).output().await?; + + if !output.status.success() { + anyhow::bail!( + "refusing to run git in {}: could not inspect its repository config \ + ({}), so whether it is safe to run under could not be determined", + dir.display(), + String::from_utf8_lossy(&output.stderr).trim() + ); + } + + let listing = String::from_utf8_lossy(&output.stdout); + for entry in listing.split('\0').filter(|e| !e.is_empty()) { + // `--null` separates entries with NUL and key from value with LF. + let key = entry.split('\n').next().unwrap_or(entry); + if !ALLOWED_REPO_CONFIG.contains(&normalise_config_key(key).as_str()) { + return Ok(Some(key.to_string())); + } + } + Ok(None) +} diff --git a/crates/tinyagents-harness-tools/src/fs/filesystem/git_operations/config_tests.rs b/crates/tinyagents-harness-tools/src/fs/filesystem/git_operations/config_tests.rs new file mode 100644 index 000000000..c63cbe778 --- /dev/null +++ b/crates/tinyagents-harness-tools/src/fs/filesystem/git_operations/config_tests.rs @@ -0,0 +1,671 @@ +//! Repository-config hardening (#5494) — the tests for +//! [`super::config`]. +//! +//! Split out of `test.rs` for the Rust layout gate, along the +//! same seam as the module split: these exercise the allow-list and the +//! hardened `git` invocation rather than the command surface. +//! +//! Attached to `git_operations` (not to the config module) so `use super::*` +//! still resolves `GitOperationsTool` the way it did before the split, and so +//! the fixtures in the sibling test module stay reachable. + +#![allow( + clippy::expect_used, + clippy::panic, + clippy::unwrap_used, + clippy::items_after_statements, + clippy::single_match_else, + clippy::single_match +)] + +use super::config::{ + NEUTRALISED_CONFIG, SHELL_NEUTRALISED_CONFIG, normalise_config_key, shell_git_env, +}; +// The fixtures stay in `test.rs` and are shared rather than +// duplicated: both modules are children of `git_operations`, so `pub(super)` +// there makes them reachable here. +use super::tests::{error_text, hermetic, init_git_repo, test_tool}; +use super::*; +use tempfile::TempDir; + +// ── run_git_command_in: repository config hardening (issue #5494) ───────── +// +// `run_git_command_in` backs every operation this tool exposes, including +// `status`, which — like `read_workspace_state`'s `run_git` before #5493 — +// invokes `core.fsmonitor` from the repository's own `.git/config`. That file +// lives in `action_dir`, which `file_write` and `git_operations` itself +// (`add`, `commit`, `checkout`) can write to, so it is attacker-controlled +// input, not trusted configuration. + +/// Write a `core.fsmonitor` hook into `dir`'s repository config that creates a +/// marker file when git runs it, and return the marker's path. +/// +/// Runs the hook once up front and asserts the marker appears, then removes +/// it — so a later absent marker means the hardening refused the hook, not +/// that the hook itself was silently broken (e.g. by `{:?}`-escaping a path +/// the shell would quote differently than Rust's `Debug` does). +#[cfg(unix)] +fn plant_fsmonitor_hook(dir: &std::path::Path) -> std::path::PathBuf { + let hook = dir.join("hook.sh"); + let marker = dir.join("COMMAND_RAN"); + std::fs::write( + &hook, + format!("#!/bin/sh\ntouch {:?}\nexit 1\n", marker.to_string_lossy()), + ) + .unwrap(); + use std::os::unix::fs::PermissionsExt; + std::fs::set_permissions(&hook, std::fs::Permissions::from_mode(0o755)).unwrap(); + + // Running a newly-written file directly can race the overlay filesystem in + // CI with ETXTBSY. Invoke the same script through the shell instead; git + // executes the configured hook through its interpreter as well. + std::process::Command::new("sh") + .arg(&hook) + .status() + .unwrap(); + assert!(marker.exists(), "the planted hook does not run at all"); + std::fs::remove_file(&marker).unwrap(); + + // Written with `git config` rather than by appending to the file: + // appending only lands in `[core]` while `[core]` happens to be the last + // section, which is true of a fresh `git init` and is not a property + // worth depending on. + let ok = hermetic( + std::process::Command::new("git") + .args(["config", "core.fsmonitor"]) + .arg(&hook) + .current_dir(dir), + ) + .status() + .unwrap() + .success(); + assert!(ok, "failed to plant the hook in the repository config"); + marker +} + +/// Set a repository config key with `git config`, asserting it took. +fn set_config(dir: &std::path::Path, key: &str, value: &str) { + let ok = hermetic( + std::process::Command::new("git") + .args(["config", key, value]) + .current_dir(dir), + ) + .status() + .unwrap() + .success(); + assert!(ok, "failed to set {key} in the test workspace"); +} + +/// Issue #5494. `git status` executes the command named by the workspace's +/// own repository config unless `run_git_command_in` refuses to run under it. +/// Revert the hardening and this test fails by finding the marker — verified, +/// not assumed. +#[cfg(unix)] +#[tokio::test] +async fn repository_config_naming_a_command_does_not_get_to_run_it() { + let tmp = TempDir::new().unwrap(); + init_git_repo(tmp.path()); + let marker = plant_fsmonitor_hook(tmp.path()); + + let tool = test_tool(tmp.path()); + let result = tool.execute(json!({"operation": "status"})).await; + let msg = error_text(&result); + + assert!( + !marker.exists(), + "`git status` executed the command named by the workspace's own \ + repository config — this tool is a code-execution primitive" + ); + assert!( + msg.contains("fsmonitor"), + "the refusal should name the key that caused it, got: {msg}" + ); +} + +/// The allowlist has to leave an ordinary repository working, or the fix is +/// just a different way of breaking the tool. +#[tokio::test] +async fn an_ordinary_repository_still_reports_status() { + let tmp = TempDir::new().unwrap(); + init_git_repo(tmp.path()); + std::fs::write(tmp.path().join("tracked.txt"), "hi").unwrap(); + + let tool = test_tool(tmp.path()); + let result = tool.execute(json!({"operation": "status"})).await.unwrap(); + + assert!(!result.is_error, "got: {}", result.output()); + assert!( + result.output().contains("tracked.txt"), + "a plain `git init` workspace must still report status, got: {}", + result.output() + ); +} + +/// A first-draft allowlist that refused any repository carrying an ordinary +/// setting like `core.autocrlf` would report nothing useful for a large class +/// of real workspaces. +#[tokio::test] +async fn an_inert_setting_an_ordinary_repository_carries_is_allowed() { + let tmp = TempDir::new().unwrap(); + init_git_repo(tmp.path()); + set_config(tmp.path(), "core.autocrlf", "input"); + set_config(tmp.path(), "gc.auto", "0"); + set_config(tmp.path(), "remote.origin.prune", "true"); + + let tool = test_tool(tmp.path()); + let result = tool.execute(json!({"operation": "status"})).await.unwrap(); + + assert!( + !result.is_error && !result.output().contains("not on the allowlist"), + "an ordinary repository must still report status, got: {}", + result.output() + ); +} + +/// The other half of the same question, and the answer is the opposite one. +/// `filter.lfs.clean` names a program, so an LFS working copy is refused — +/// fail-closed, and intended rather than an oversight. +#[tokio::test] +async fn an_lfs_clone_is_refused_because_its_filter_names_a_program() { + let tmp = TempDir::new().unwrap(); + init_git_repo(tmp.path()); + // What `git lfs install` writes. `required` is inert and allowed; the + // three programs are not. + set_config(tmp.path(), "filter.lfs.required", "true"); + set_config(tmp.path(), "filter.lfs.clean", "git-lfs clean -- %f"); + + let tool = test_tool(tmp.path()); + let result = tool.execute(json!({"operation": "status"})).await; + let msg = error_text(&result); + + assert!( + msg.contains("filter.lfs.clean"), + "the refusal must name the key that caused it, got: {msg}" + ); +} + +/// `credential.helper` reads like a preference and is command-valued: a value +/// beginning `!` is run as a shell command. It must be refused however inert +/// it reads. +#[tokio::test] +async fn credential_helper_is_refused_despite_looking_like_a_preference() { + let tmp = TempDir::new().unwrap(); + init_git_repo(tmp.path()); + set_config(tmp.path(), "credential.helper", "!echo pwned"); + + let tool = test_tool(tmp.path()); + let result = tool.execute(json!({"operation": "status"})).await; + let msg = error_text(&result); + + assert!( + msg.contains("credential.helper"), + "a command-valued key must be refused however inert it reads, got: {msg}" + ); +} + +/// The refusal must hold for a write operation too, not just `status` — the +/// same repository config runs under `commit`/`add`/`checkout`/`stash`. +#[tokio::test] +async fn refusal_also_covers_write_operations() { + let tmp = TempDir::new().unwrap(); + init_git_repo(tmp.path()); + set_config(tmp.path(), "credential.helper", "!echo pwned"); + + let tool = test_tool(tmp.path()); + let result = tool + .execute(json!({"operation": "commit", "message": "test"})) + .await + .unwrap(); + + assert!( + result.is_error && result.output().contains("credential.helper"), + "write operations must be refused under untrusted repo config too, got: {}", + result.output() + ); +} + +/// `core.worktree` redirects the working-tree root every write operation +/// here (`checkout`, `add`, `commit`, `stash`) targets. Left on the +/// allowlist, a repository config could point that root outside +/// `action_dir` and turn a supposedly sandboxed write into one against an +/// arbitrary directory. Nothing this tool does needs the key — worktree +/// isolation goes through `WorkspaceDescriptor` instead. +#[tokio::test] +async fn core_worktree_is_refused_because_it_can_redirect_writes_outside_the_sandbox() { + let tmp = TempDir::new().unwrap(); + let elsewhere = TempDir::new().unwrap(); + init_git_repo(tmp.path()); + set_config( + tmp.path(), + "core.worktree", + &elsewhere.path().to_string_lossy(), + ); + + let tool = test_tool(tmp.path()); + let result = tool.execute(json!({"operation": "status"})).await; + let msg = error_text(&result); + + assert!( + msg.contains("core.worktree"), + "the refusal must name the key that caused it, got: {msg}" + ); +} + +/// `extensions.worktreeConfig` is itself allowlisted as an ordinary setting, +/// but turning it on makes git additionally read `config.worktree` — a +/// second file `--local` alone does not see. A `core.hooksPath` set there is +/// invisible to a `--local`-only inspection and would still run on the next +/// `commit`. The inspection step must read the same merged view git does. +#[tokio::test] +async fn a_hookspath_hidden_in_worktree_scoped_config_is_still_refused() { + let tmp = TempDir::new().unwrap(); + init_git_repo(tmp.path()); + set_config(tmp.path(), "extensions.worktreeConfig", "true"); + // What `git config --worktree core.hooksPath ` writes; `set_config` + // only reaches `--local`, so this key is planted directly the same way + // the production inspection step reads it — via a real `git config + // --worktree` invocation — to prove the bypass is closed, not just that + // `set_config` happens to skip it. + let ok = hermetic( + std::process::Command::new("git") + .args(["config", "--worktree", "core.hooksPath"]) + .arg(tmp.path()) + .current_dir(tmp.path()), + ) + .status() + .unwrap() + .success(); + assert!(ok, "failed to set core.hooksPath in worktree-scoped config"); + + let tool = test_tool(tmp.path()); + let result = tool.execute(json!({"operation": "status"})).await; + let msg = error_text(&result); + + assert!( + msg.contains("core.hookspath"), + "a hookspath hidden in worktree-scoped config must still be refused, got: {msg}" + ); +} + +/// The allowlist inspection and the real command are two separate `git` +/// invocations, so a config change landing in the gap between them would be +/// invisible to the first and still reach the second. This test calls +/// `hardened_git` directly — skipping `first_disallowed_repo_config_key` +/// entirely, standing in for that gap — to prove the second invocation does +/// not depend on the first having caught anything: `core.hooksPath` is +/// neutralised at the point of execution regardless of what any inspection +/// saw or missed. +#[cfg(unix)] +#[tokio::test] +async fn hardened_git_neutralises_hookspath_even_if_the_allowlist_check_never_ran() { + let tmp = TempDir::new().unwrap(); + init_git_repo(tmp.path()); + set_config(tmp.path(), "user.email", "test@example.com"); + set_config(tmp.path(), "user.name", "Test"); + + let hooks_dir = tmp.path().join("evil-hooks"); + std::fs::create_dir(&hooks_dir).unwrap(); + let marker = tmp.path().join("HOOK_RAN"); + std::fs::write( + hooks_dir.join("pre-commit"), + format!("#!/bin/sh\ntouch {:?}\nexit 0\n", marker.to_string_lossy()), + ) + .unwrap(); + { + use std::os::unix::fs::PermissionsExt; + std::fs::set_permissions( + hooks_dir.join("pre-commit"), + std::fs::Permissions::from_mode(0o755), + ) + .unwrap(); + } + set_config(tmp.path(), "core.hooksPath", &hooks_dir.to_string_lossy()); + + std::fs::write(tmp.path().join("f.txt"), "hi").unwrap(); + let staged = hermetic( + std::process::Command::new("git") + .args(["add", "f.txt"]) + .current_dir(tmp.path()), + ) + .status() + .unwrap() + .success(); + assert!(staged, "failed to stage the test file"); + + let output = super::hardened_git(tmp.path()) + .args(["commit", "-m", "msg"]) + .output() + .await + .unwrap(); + + assert!( + output.status.success(), + "commit should succeed: {}", + String::from_utf8_lossy(&output.stderr) + ); + assert!( + !marker.exists(), + "hardened_git ran the repository-configured pre-commit hook — the \ + override that is supposed to hold even without the allowlist check \ + did not" + ); +} + +/// `commit.gpgsign` is on `ALLOWED_REPO_CONFIG` as an ordinary boolean, but +/// left un-neutralised it would let a repository force every commit through +/// this tool to be signed. `output.status.success()` alone does not prove +/// that: a repository that also configures a *working* `gpg.program` would +/// make a signed commit succeed too, so this plants a fake one that always +/// signs successfully and then inspects the commit object itself for a +/// `gpgsig` header — the only assertion that actually distinguishes "signing +/// was skipped" from "signing was attempted and happened to work". +#[cfg(unix)] +#[tokio::test] +async fn hardened_git_neutralises_forced_commit_signing() { + use std::os::unix::fs::PermissionsExt; + + let tmp = TempDir::new().unwrap(); + init_git_repo(tmp.path()); + set_config(tmp.path(), "user.email", "test@example.com"); + set_config(tmp.path(), "user.name", "Test"); + set_config(tmp.path(), "commit.gpgsign", "true"); + + // A `gpg.program` that always "signs" successfully, so a regression here + // fails by finding a signature, not by the commit merely erroring out — + // the same distinction CodeRabbit's review raised. + let fake_gpg = tmp.path().join("fake-gpg.sh"); + std::fs::write( + &fake_gpg, + "#!/bin/sh\n\ + printf '%s\\n' '[GNUPG:] BEGIN_SIGNING H10' >&2\n\ + cat >/dev/null\n\ + printf -- '-----BEGIN PGP SIGNATURE-----\\n\\nZmFrZQ==\\n-----END PGP SIGNATURE-----\\n'\n\ + printf '%s\\n' '[GNUPG:] SIG_CREATED D 1 10 00 0 0123456789ABCDEF' >&2\n", + ) + .unwrap(); + std::fs::set_permissions(&fake_gpg, std::fs::Permissions::from_mode(0o755)).unwrap(); + set_config(tmp.path(), "gpg.program", &fake_gpg.to_string_lossy()); + + std::fs::write(tmp.path().join("f.txt"), "hi").unwrap(); + let staged = hermetic( + std::process::Command::new("git") + .args(["add", "f.txt"]) + .current_dir(tmp.path()), + ) + .status() + .unwrap() + .success(); + assert!(staged, "failed to stage the test file"); + + let output = super::hardened_git(tmp.path()) + .args(["commit", "-m", "msg"]) + .output() + .await + .unwrap(); + assert!( + output.status.success(), + "commit failed: {}", + String::from_utf8_lossy(&output.stderr) + ); + + let commit_object = super::hardened_git(tmp.path()) + .args(["cat-file", "-p", "HEAD"]) + .output() + .await + .unwrap(); + assert!(commit_object.status.success()); + let commit_object = String::from_utf8_lossy(&commit_object.stdout); + assert!( + !commit_object.lines().any(|l| l.starts_with("gpgsig ")), + "commit.gpgsign=true must not be honoured, but HEAD carries a \ + signature: {commit_object}" + ); +} + +/// The config-inspection step must fail closed: if `git config --list +/// --local` cannot be read, that is not the same as "nothing to distrust", +/// and running the real command anyway would skip the check entirely. +#[cfg(unix)] +#[tokio::test] +async fn unreadable_repo_config_fails_closed_rather_than_running_anyway() { + use std::os::unix::fs::PermissionsExt; + + let tmp = TempDir::new().unwrap(); + init_git_repo(tmp.path()); + let config_path = tmp.path().join(".git").join("config"); + std::fs::set_permissions(&config_path, std::fs::Permissions::from_mode(0o000)).unwrap(); + + // Root (and some CI containers run as root) ignores this permission bit + // entirely, which would make the assertion below meaningless rather than + // wrong. Detect that up front instead of failing on an unrelated cause. + let permission_enforced = std::fs::File::open(&config_path).is_err(); + + let result = if permission_enforced { + let tool = test_tool(tmp.path()); + Some(tool.execute(json!({"operation": "status"})).await) + } else { + None + }; + + // Restore permissions before the TempDir is dropped, so cleanup doesn't + // fail on an unreadable file. + std::fs::set_permissions(&config_path, std::fs::Permissions::from_mode(0o644)).unwrap(); + + let Some(result) = result else { + eprintln!( + "skipping: file permissions are not enforced against this process (running as root?)" + ); + return; + }; + let msg = error_text(&result); + + assert!( + msg.contains("could not inspect its repository config"), + "an unreadable repo config must refuse, not silently proceed, got: {msg}" + ); +} + +#[test] +fn a_subsection_is_elided_so_one_entry_covers_every_remote() { + assert_eq!(normalise_config_key("remote.origin.url"), "remote.url"); + assert_eq!(normalise_config_key("remote.a.b.c.url"), "remote.url"); + assert_eq!(normalise_config_key("core.fileMode"), "core.filemode"); + assert_eq!(normalise_config_key("core.fsmonitor"), "core.fsmonitor"); + // The subsection itself contains dots; the first and last components + // remain the reliable ones. + assert_eq!( + normalise_config_key("includeIf.gitdir:~/x.y/.path"), + "includeif.path" + ); + // A key with no dot at all is returned unchanged rather than panicking. + assert_eq!(normalise_config_key("bare"), "bare"); +} + +// ── External diff suppression ───────────────────────────────────────────── + +/// An ordinary repository's `diff` must produce its patch. +/// +/// That reads like it could not possibly regress, and it did. Suppression of +/// an external diff was attempted with `-c diff.external=` in +/// `NEUTRALISED_CONFIG`, and an empty value does not disable one — git tries +/// to *execute* the empty string, so **every** diff died with +/// `error: cannot run : No such file or directory` / +/// `fatal: external diff died`, on every repository, hostile or not. The +/// hardening removed the operation instead of hardening it. `--no-ext-diff` +/// on the diff command is the real suppression. +/// +/// This lives in the lib suite deliberately. The integration test that caught +/// it sits in `raw_coverage_all`, and a change to `git_operations.rs` maps to +/// the `openhuman::tools` libtest filter — which never selects that target. So +/// the lane the change picks did not run the test covering the change, and the +/// breakage sat until an unrelated PR happened to select the other filter. +/// A test here runs whenever this file is touched. +/// +/// A repository that *sets* `diff.external` is a separate matter and is +/// already refused before reaching the invocation: the key is on neither +/// allowlist, so `first_disallowed_repo_config_key` rejects the repository +/// outright. `--no-ext-diff` is the second layer, covering the gap between +/// that inspection and the command. +#[tokio::test] +async fn an_ordinary_repository_still_produces_a_diff() { + let tmp = TempDir::new().unwrap(); + init_git_repo(tmp.path()); + + // A committer identity has to be set in the repository itself. `hermetic` + // closes the global and system config, which is the point of it — so on a + // CI container with no identity of its own `git commit` fails with + // "Author identity unknown". Both keys are on `ALLOWED_REPO_CONFIG`, so + // setting them does not trip the repository-config refusal. + set_config(tmp.path(), "user.email", "test@example.invalid"); + set_config(tmp.path(), "user.name", "Test"); + + let tracked = tmp.path().join("tracked.txt"); + std::fs::write(&tracked, "first\n").unwrap(); + // `hermetic` closes the ambient git config the way the fixtures do — + // without it a developer's global `commit.gpgsign` or hooks path can fail + // this commit for reasons unrelated to the test. + // Slices, not arrays: `["add", "tracked.txt"]` and `["commit", "-m", "one"]` + // are `[&str; 2]` and `[&str; 3]`, which are different types and cannot + // share an array literal. + for args in [&["add", "tracked.txt"][..], &["commit", "-m", "one"][..]] { + let ok = hermetic( + std::process::Command::new("git") + .args(args) + .current_dir(tmp.path()), + ) + .status() + .unwrap() + .success(); + assert!(ok, "failed to run `git {}` in the test workspace", args[0]); + } + std::fs::write(&tracked, "first\nsecond\n").unwrap(); + + let tool = test_tool(tmp.path()); + let result = tool + .execute(json!({"operation": "diff", "files": "tracked.txt"})) + .await + .expect("a diff on a plain repository must not error out"); + + assert!(!result.is_error, "got: {}", result.output()); + assert!( + result.output().contains("second"), + "the added line must appear in the patch: {}", + result.output() + ); + assert!( + !result.output().contains("external diff died"), + "git must not be handed an empty `diff.external` to execute: {}", + result.output() + ); +} + +/// The shell subset must stay a subset, so widening the policy is a decision. +/// +/// `SHELL_NEUTRALISED_CONFIG` exists because the `shell` tool cannot use the +/// `-c` layer `hardened_git` builds. The risk is drift in one direction: +/// someone adds a newly-dangerous key to `NEUTRALISED_CONFIG`, the shell path +/// silently keeps running without it, and the gap is invisible. This test +/// cannot force the addition — some entries genuinely must not cross, and the +/// doc comment says which and why — but it does force every shell entry to be +/// one the main policy already recognises. +#[test] +fn the_shell_subset_never_invents_a_key_the_main_policy_does_not_have() { + for entry in SHELL_NEUTRALISED_CONFIG { + let known = NEUTRALISED_CONFIG.contains(entry) + // `hardened_git` adds this one as an argument rather than through + // the static list; it is still the same policy. + || *entry == "commit.gpgSign=false"; + assert!( + known, + "`{entry}` is neutralised for the shell but is not part of the \ + git_operations policy it is supposed to mirror" + ); + } +} + +/// The three deliberate exclusions must stay excluded, with their reasons. +/// +/// Each would break ordinary work rather than harden it: an empty +/// `core.sshCommand` is a command git tries to execute (the same trap the +/// `diff.external` note records), while overriding `core.hooksPath` would +/// suppress repository hooks. Re-adding any of them looks like tightening +/// security and is actually an outage, so it should be a conscious edit here +/// first. +#[test] +fn the_shell_subset_keeps_its_documented_exclusions() { + for excluded in ["core.sshCommand=", "diff.external=", "core.hooksPath="] { + assert!( + !SHELL_NEUTRALISED_CONFIG.contains(&excluded), + "`{excluded}` is excluded from the shell subset on purpose — see \ + the doc comment before changing this" + ); + } +} + +#[test] +fn shell_git_environment_forces_the_safe_config_subset() { + let env = shell_git_env(); + let parameters = env + .get(std::ffi::OsStr::new("GIT_CONFIG_PARAMETERS")) + .expect("the shell git environment sets GIT_CONFIG_PARAMETERS") + .to_string_lossy(); + + for expected in SHELL_NEUTRALISED_CONFIG { + let (key, value) = expected.split_once('=').unwrap(); + assert!( + parameters.contains(&format!("'{key}'='{value}'")), + "`{expected}` is missing from GIT_CONFIG_PARAMETERS: {parameters}" + ); + } +} + +/// Exercise the environment through Git itself rather than only checking its +/// serialization. A syntactically plausible `GIT_CONFIG_PARAMETERS` value is +/// not useful if Git rejects it, and `core.fsmonitor=` must actually override +/// the command-valued setting in the repository config. +#[cfg(unix)] +#[test] +fn shell_git_environment_is_accepted_and_neutralises_repository_fsmonitor() { + let tmp = TempDir::new().unwrap(); + init_git_repo(tmp.path()); + let marker = plant_fsmonitor_hook(tmp.path()); + + // Prove this Git version and repository really do invoke the planted + // fsmonitor hook. Without this control, an absent marker below could be a + // false positive caused by Git never consulting the setting. + let plain = hermetic( + std::process::Command::new("git") + .arg("status") + .current_dir(tmp.path()), + ) + .output() + .expect("failed to run the unhardened control command"); + assert!( + plain.status.success(), + "the unhardened control command failed: {}", + String::from_utf8_lossy(&plain.stderr) + ); + assert!( + marker.exists(), + "the unhardened control did not execute the planted fsmonitor hook" + ); + std::fs::remove_file(&marker).unwrap(); + + let env = shell_git_env(); + let mut hardened = std::process::Command::new("git"); + hardened.arg("status").current_dir(tmp.path()); + hermetic(&mut hardened).envs(env); + let output = hardened + .output() + .expect("failed to run git with the shell hardening environment"); + + assert!( + output.status.success(), + "Git rejected the shell hardening environment: {}", + String::from_utf8_lossy(&output.stderr) + ); + assert!( + !marker.exists(), + "Git executed repository core.fsmonitor despite the shell hardening environment" + ); +} diff --git a/crates/tinyagents-harness-tools/src/fs/filesystem/git_operations/mod.rs b/crates/tinyagents-harness-tools/src/fs/filesystem/git_operations/mod.rs new file mode 100644 index 000000000..03d97810f --- /dev/null +++ b/crates/tinyagents-harness-tools/src/fs/filesystem/git_operations/mod.rs @@ -0,0 +1,656 @@ +use super::gate::FsGate; +use async_trait::async_trait; +use serde_json::json; +use std::path::{Path, PathBuf}; +use std::sync::Arc; +use tinytools::ToolRunContext; +use tinytools::{Tool, ToolCallOptions, ToolResult}; + +pub(crate) mod config; +mod render; + +pub use self::config::shell_git_env; +use self::config::{disallowed_config_refusal, first_disallowed_repo_config_key, hardened_git}; +use self::render::{render_branch_markdown, render_log_markdown, render_status_markdown}; + +/// Git operations tool for structured repository management. +/// Provides safe, parsed git operations with JSON output. +#[derive(Debug)] +pub struct GitOperationsTool { + gate: Arc, + action_dir: PathBuf, +} + +impl GitOperationsTool { + /// Create the tool. + #[must_use] + pub fn new(gate: Arc, action_dir: PathBuf) -> Self { + Self { gate, action_dir } + } + + /// Resolve the working directory for git operations. + /// + /// Returns the per-worker git-worktree checkout when the tinyagents harness + /// threaded a [`WorkspaceDescriptor`] into this call's + /// [`ToolExecutionContext`] — an edit-capable worker running with + /// `isolation = "worktree"`, whose isolated worktree root is carried on the + /// run context (`RunContext::with_workspace`) and surfaced per tool call via + /// `ToolExecutionContext::from_run_context`. Otherwise falls back to the + /// tool's configured `action_dir`, which preserves the non-isolated + /// behaviour exactly. See #3376, #4249 (08.5). + fn effective_action_dir_for_context(&self, context: Option<&dyn ToolRunContext>) -> PathBuf { + if let Some(workspace) = context.and_then(|ctx| ctx.workspace()) { + tracing::debug!( + workspace_root = %workspace.root.display(), + policy_id = %workspace.policy_id, + "[git_operations] using TinyAgents workspace descriptor as action dir" + ); + return workspace.root.clone(); + } + self.action_dir.clone() + } + + /// Sanitize git arguments to prevent injection attacks + fn sanitize_git_args(&self, args: &str) -> anyhow::Result> { + let mut result = Vec::new(); + for arg in args.split_whitespace() { + // Block dangerous git options that could lead to command injection + let arg_lower = arg.to_lowercase(); + if arg_lower.starts_with("--exec=") + || arg_lower.starts_with("--upload-pack=") + || arg_lower.starts_with("--receive-pack=") + || arg_lower.starts_with("--pager=") + || arg_lower.starts_with("--editor=") + || arg_lower == "--no-verify" + || arg_lower.contains("$(") + || arg_lower.contains('`') + || arg.contains('|') + || arg.contains(';') + || arg.contains('>') + { + anyhow::bail!("Blocked potentially dangerous git argument: {arg}"); + } + // Block `-c` config injection (exact match or `-c=...` prefix). + // This must not false-positive on `--cached` or `-cached`. + if arg_lower == "-c" || arg_lower.starts_with("-c=") { + anyhow::bail!("Blocked potentially dangerous git argument: {arg}"); + } + result.push(arg.to_string()); + } + Ok(result) + } + + /// Check if an operation requires write access + fn requires_write_access(&self, operation: &str) -> bool { + matches!( + operation, + "commit" | "add" | "checkout" | "stash" | "reset" | "revert" + ) + } + + /// Check if an operation is read-only + #[allow(dead_code)] // exercised by the tests; kept with `requires_write_access` + fn is_read_only(&self, operation: &str) -> bool { + matches!( + operation, + "status" | "diff" | "log" | "show" | "branch" | "rev-parse" + ) + } + + async fn run_git_command_in(&self, cwd: &Path, args: &[&str]) -> anyhow::Result { + if let Some(key) = first_disallowed_repo_config_key(cwd).await? { + tracing::debug!( + "[git_operations] refusing to run git: dir={}, disallowed_config_key={key}", + cwd.display() + ); + anyhow::bail!("{}", disallowed_config_refusal(cwd, &key)) + } + + let output = hardened_git(cwd).args(args).output().await?; + + if !output.status.success() { + let stderr = String::from_utf8_lossy(&output.stderr); + anyhow::bail!("Git command failed: {stderr}"); + } + + Ok(String::from_utf8_lossy(&output.stdout).to_string()) + } + + async fn git_status(&self, cwd: &Path, _args: serde_json::Value) -> anyhow::Result { + let output = self + .run_git_command_in(cwd, &["status", "--porcelain=2", "--branch"]) + .await?; + + // Parse git status output into structured format + let mut result = serde_json::Map::new(); + let mut branch = String::new(); + let mut staged = Vec::new(); + let mut unstaged = Vec::new(); + let mut untracked = Vec::new(); + + for line in output.lines() { + if line.starts_with("# branch.head ") { + branch = line.trim_start_matches("# branch.head ").to_string(); + } else if let Some(rest) = line.strip_prefix("1 ") { + // Ordinary changed entry + let mut parts = rest.splitn(3, ' '); + if let (Some(staging), Some(path)) = (parts.next(), parts.next()) + && !staging.is_empty() + { + let status_char = staging.chars().next().unwrap_or(' '); + if status_char != '.' && status_char != ' ' { + staged.push(json!({"path": path, "status": status_char})); + } + let status_char = staging.chars().nth(1).unwrap_or(' '); + if status_char != '.' && status_char != ' ' { + unstaged.push(json!({"path": path, "status": status_char})); + } + } + } else if let Some(rest) = line.strip_prefix("? ") { + untracked.push(rest.to_string()); + } + } + + result.insert("branch".to_string(), json!(branch)); + result.insert("staged".to_string(), json!(staged)); + result.insert("unstaged".to_string(), json!(unstaged)); + result.insert("untracked".to_string(), json!(untracked)); + result.insert( + "clean".to_string(), + json!(staged.is_empty() && unstaged.is_empty() && untracked.is_empty()), + ); + + let mut tr = ToolResult::success(serde_json::to_string_pretty(&result).unwrap_or_default()); + tr.markdown_formatted = Some(render_status_markdown(&result)); + Ok(tr) + } + + async fn git_diff(&self, cwd: &Path, args: serde_json::Value) -> anyhow::Result { + let files = args.get("files").and_then(|v| v.as_str()).unwrap_or("."); + let cached = args + .get("cached") + .and_then(serde_json::Value::as_bool) + .unwrap_or(false); + + // Validate files argument against injection patterns + self.sanitize_git_args(files)?; + + // `--no-ext-diff` is where external-diff suppression has to live: it is + // a diff-command flag, and the `-c diff.external=` that used to stand in + // for it made git exec the empty string instead of disabling anything. + // + // `--no-textconv` is a SEPARATE mechanism and needs its own flag. + // `--no-ext-diff` covers `diff.external` and `diff..command`; + // it does not touch `diff..textconv`, which a `.gitattributes` + // line (`*.bin diff=evil`) can select and which git then EXECUTES to + // render a binary file as text. Verified against a scratch repo: with + // `--no-ext-diff` alone the textconv script ran; adding + // `--no-textconv` it did not. `hardened_git`'s `-c` list cannot close + // this — driver names are arbitrary, so there is no finite set of keys + // to neutralise, and the suppression has to be a command flag. + let mut git_args = vec!["diff", "--no-ext-diff", "--no-textconv", "--unified=3"]; + if cached { + git_args.push("--cached"); + } + git_args.push("--"); + git_args.push(files); + + let output = self.run_git_command_in(cwd, &git_args).await?; + + // Parse diff into structured hunks + let mut result = serde_json::Map::new(); + let mut hunks = Vec::new(); + let mut current_file = String::new(); + let mut current_hunk = serde_json::Map::new(); + let mut lines = Vec::new(); + + for line in output.lines() { + if line.starts_with("diff --git ") { + if !lines.is_empty() { + current_hunk.insert("lines".to_string(), json!(lines)); + if !current_hunk.is_empty() { + hunks.push(serde_json::Value::Object(current_hunk.clone())); + } + lines = Vec::new(); + current_hunk = serde_json::Map::new(); + } + let parts: Vec<&str> = line.split_whitespace().collect(); + if parts.len() >= 4 { + current_file = parts[3].trim_start_matches("b/").to_string(); + current_hunk.insert("file".to_string(), json!(current_file)); + } + } else if line.starts_with("@@ ") { + if !lines.is_empty() { + current_hunk.insert("lines".to_string(), json!(lines)); + if !current_hunk.is_empty() { + hunks.push(serde_json::Value::Object(current_hunk.clone())); + } + lines = Vec::new(); + current_hunk = serde_json::Map::new(); + current_hunk.insert("file".to_string(), json!(current_file)); + } + current_hunk.insert("header".to_string(), json!(line)); + } else if !line.is_empty() { + lines.push(json!({ + "text": line, + "type": if line.starts_with('+') { "add" } + else if line.starts_with('-') { "delete" } + else { "context" } + })); + } + } + + if !lines.is_empty() { + current_hunk.insert("lines".to_string(), json!(lines)); + if !current_hunk.is_empty() { + hunks.push(serde_json::Value::Object(current_hunk)); + } + } + + result.insert("hunks".to_string(), json!(hunks)); + result.insert("file_count".to_string(), json!(hunks.len())); + + Ok(ToolResult::success( + serde_json::to_string_pretty(&result).unwrap_or_default(), + )) + } + + async fn git_log(&self, cwd: &Path, args: serde_json::Value) -> anyhow::Result { + let limit_raw = args + .get("limit") + .and_then(serde_json::Value::as_u64) + .unwrap_or(10); + let limit = usize::try_from(limit_raw).unwrap_or(usize::MAX).min(1000); + let limit_str = limit.to_string(); + + let output = self + .run_git_command_in( + cwd, + &[ + "log", + &format!("-{limit_str}"), + "--pretty=format:%H|%an|%ae|%ad|%s", + "--date=iso", + ], + ) + .await?; + + let mut commits = Vec::new(); + + for line in output.lines() { + let parts: Vec<&str> = line.split('|').collect(); + if parts.len() >= 5 { + commits.push(json!({ + "hash": parts[0], + "author": parts[1], + "email": parts[2], + "date": parts[3], + "message": parts[4] + })); + } + } + + let mut tr = ToolResult::success( + serde_json::to_string_pretty(&json!({ "commits": commits })).unwrap_or_default(), + ); + tr.markdown_formatted = Some(render_log_markdown(&commits)); + Ok(tr) + } + + async fn git_branch(&self, cwd: &Path, _args: serde_json::Value) -> anyhow::Result { + let output = self + .run_git_command_in(cwd, &["branch", "--format=%(refname:short)|%(HEAD)"]) + .await?; + + let mut branches = Vec::new(); + let mut current = String::new(); + + for line in output.lines() { + if let Some((name, head)) = line.split_once('|') { + let is_current = head == "*"; + if is_current { + current = name.to_string(); + } + branches.push(json!({ + "name": name, + "current": is_current + })); + } + } + + let mut tr = ToolResult::success( + serde_json::to_string_pretty(&json!({ + "current": current, + "branches": branches + })) + .unwrap_or_default(), + ); + tr.markdown_formatted = Some(render_branch_markdown(¤t, &branches)); + Ok(tr) + } + + fn truncate_commit_message(message: &str) -> String { + if message.chars().count() > 2000 { + format!("{}...", message.chars().take(1997).collect::()) + } else { + message.to_string() + } + } + + async fn git_commit(&self, cwd: &Path, args: serde_json::Value) -> anyhow::Result { + let message = args + .get("message") + .and_then(|v| v.as_str()) + .ok_or_else(|| anyhow::anyhow!("Missing 'message' parameter"))?; + + // Sanitize commit message + let sanitized = message + .lines() + .map(str::trim) + .filter(|l| !l.is_empty()) + .collect::>() + .join("\n"); + + if sanitized.is_empty() { + anyhow::bail!("Commit message cannot be empty"); + } + + // Limit message length + let message = Self::truncate_commit_message(&sanitized); + + let output = self + .run_git_command_in(cwd, &["commit", "-m", &message]) + .await; + + match output { + Ok(_) => Ok(ToolResult::success(format!("Committed: {message}"))), + Err(e) => Ok(ToolResult::error(format!("Commit failed: {e}"))), + } + } + + async fn git_add(&self, cwd: &Path, args: serde_json::Value) -> anyhow::Result { + let paths = args + .get("paths") + .and_then(|v| v.as_str()) + .ok_or_else(|| anyhow::anyhow!("Missing 'paths' parameter"))?; + + // Validate paths against injection patterns + self.sanitize_git_args(paths)?; + + let output = self.run_git_command_in(cwd, &["add", "--", paths]).await; + + match output { + Ok(_) => Ok(ToolResult::success(format!("Staged: {paths}"))), + Err(e) => Ok(ToolResult::error(format!("Add failed: {e}"))), + } + } + + async fn git_checkout( + &self, + cwd: &Path, + args: serde_json::Value, + ) -> anyhow::Result { + let branch = args + .get("branch") + .and_then(|v| v.as_str()) + .ok_or_else(|| anyhow::anyhow!("Missing 'branch' parameter"))?; + + // Sanitize branch name + let sanitized = self.sanitize_git_args(branch)?; + + if sanitized.is_empty() || sanitized.len() > 1 { + anyhow::bail!("Invalid branch specification"); + } + + let branch_name = &sanitized[0]; + + // Block dangerous branch names + if branch_name.contains('@') || branch_name.contains('^') || branch_name.contains('~') { + anyhow::bail!("Branch name contains invalid characters"); + } + + let output = self + .run_git_command_in(cwd, &["checkout", branch_name]) + .await; + + match output { + Ok(_) => Ok(ToolResult::success(format!( + "Switched to branch: {branch_name}" + ))), + Err(e) => Ok(ToolResult::error(format!("Checkout failed: {e}"))), + } + } + + async fn git_stash(&self, cwd: &Path, args: serde_json::Value) -> anyhow::Result { + let action = args + .get("action") + .and_then(|v| v.as_str()) + .unwrap_or("push"); + + let output = match action { + "push" | "save" => { + self.run_git_command_in(cwd, &["stash", "push", "-m", "auto-stash"]) + .await + } + "pop" => self.run_git_command_in(cwd, &["stash", "pop"]).await, + "list" => self.run_git_command_in(cwd, &["stash", "list"]).await, + "drop" => { + let index_raw = args + .get("index") + .and_then(serde_json::Value::as_u64) + .unwrap_or(0); + let index = i32::try_from(index_raw) + .map_err(|_| anyhow::anyhow!("stash index too large: {index_raw}"))?; + self.run_git_command_in(cwd, &["stash", "drop", &format!("stash@{{{index}}}")]) + .await + } + _ => anyhow::bail!("Unknown stash action: {action}. Use: push, pop, list, drop"), + }; + + match output { + Ok(out) => Ok(ToolResult::success(out)), + Err(e) => Ok(ToolResult::error(format!("Stash {action} failed: {e}"))), + } + } +} + +#[async_trait] +impl Tool for GitOperationsTool { + fn name(&self) -> &'static str { + "git_operations" + } + + fn description(&self) -> &'static str { + "Perform structured Git operations (status, diff, log, branch, commit, add, checkout, stash). Provides parsed JSON output and integrates with security policy for autonomy controls." + } + + fn parameters_schema(&self) -> serde_json::Value { + json!({ + "type": "object", + "properties": { + "operation": { + "type": "string", + "enum": ["status", "diff", "log", "branch", "commit", "add", "checkout", "stash"], + "description": "Git operation to perform" + }, + "message": { + "type": "string", + "description": "Commit message (for 'commit' operation)" + }, + "paths": { + "type": "string", + "description": "File paths to stage (for 'add' operation)" + }, + "branch": { + "type": "string", + "description": "Branch name (for 'checkout' operation)" + }, + "files": { + "type": "string", + "description": "File or path to diff (for 'diff' operation, default: '.')" + }, + "cached": { + "type": "boolean", + "description": "Show staged changes (for 'diff' operation)" + }, + "limit": { + "type": "integer", + "description": "Number of log entries (for 'log' operation, default: 10)" + }, + "action": { + "type": "string", + "enum": ["push", "pop", "list", "drop"], + "description": "Stash action (for 'stash' operation)" + }, + "index": { + "type": "integer", + "description": "Stash index (for 'stash' with 'drop' action)" + } + }, + "required": ["operation"] + }) + } + + fn supports_markdown(&self) -> bool { + true + } + + /// Write git operations (commit/add/checkout/stash/…) route through the + /// human approval gate in ask-before-edit; read operations (status/diff/ + /// log/branch) never prompt. In Full, writes run; read-only is blocked in + /// `execute` via the existing `can_act()` / autonomy check. + fn external_effect_with_args(&self, args: &serde_json::Value) -> bool { + let operation = args.get("operation").and_then(|v| v.as_str()).unwrap_or(""); + self.requires_write_access(operation) && self.gate.write_needs_approval() + } + + async fn execute_with_options( + &self, + args: serde_json::Value, + _options: ToolCallOptions, + ) -> anyhow::Result { + // git_operations always populates `markdown_formatted` for the + // structured sub-operations (status/diff/log/branch). The harness + // picks it up when `prefer_markdown` is on; the JSON content + // block is preserved for callers that want the raw structure. + self.execute_in_context(args, None).await + } + + async fn execute_with_context( + &self, + args: serde_json::Value, + _options: ToolCallOptions, + context: Option<&dyn ToolRunContext>, + ) -> anyhow::Result { + self.execute_in_context(args, context).await + } + + async fn execute(&self, args: serde_json::Value) -> anyhow::Result { + self.execute_in_context(args, None).await + } +} + +impl GitOperationsTool { + async fn execute_in_context( + &self, + args: serde_json::Value, + context: Option<&dyn ToolRunContext>, + ) -> anyhow::Result { + let operation = match args.get("operation").and_then(|v| v.as_str()) { + Some(op) => op, + None => { + return Ok(ToolResult::error("Missing 'operation' parameter")); + } + }; + + let effective_dir = self.effective_action_dir_for_context(context); + // Validate the repository instead of trusting the presence of a `.git` + // path. This handles linked-worktree gitfiles and rejects malformed + // ancestor markers that Git itself cannot open. + // + // Probed with `hardened_git` directly rather than through + // `run_git_command_in`. That helper refuses up front when the + // repository config carries a disallowed key, and routing the probe + // through it made the refusal indistinguishable from "there is no + // repository here": `is_ok_and` discarded the error, so a repo setting + // `core.fsmonitor` / `credential.helper` / `core.worktree` was reported + // as **not a git repository** and the message naming the key was lost + // (#5494). `first_disallowed_repo_config_key` even documents the + // ordering this restores — it expects the caller to have established + // that `dir` is a repository before it runs. + // + // Safe to probe unguarded: `hardened_git` already neutralises every + // config key that names a command, so this invocation cannot be turned + // into an execution primitive, and the real command below still goes + // through the full allow-list check in `run_git_command_in`. + let is_worktree = hardened_git(&effective_dir) + .args(["rev-parse", "--is-inside-work-tree"]) + .output() + .await + .is_ok_and(|output| { + output.status.success() && String::from_utf8_lossy(&output.stdout).trim() == "true" + }); + if !is_worktree { + // The probe can fail *because* the repository is misconfigured: a + // `core.worktree` redirect git cannot follow, or a config file it + // cannot read. Ask the guard before concluding there is nothing + // here, so a refusal is reported as a refusal — "Not in a git + // repository" is both false and un-actionable for those (#5494). + return Ok( + match first_disallowed_repo_config_key(&effective_dir).await { + Ok(Some(key)) => { + ToolResult::error(disallowed_config_refusal(&effective_dir, &key)) + } + Err(error) => ToolResult::error(error.to_string()), + Ok(None) => ToolResult::error("Not in a git repository"), + }, + ); + } + + // Check autonomy level for write operations + if self.requires_write_access(operation) { + if !self.gate.can_act() { + return Ok(ToolResult::error( + "[policy-blocked] Action blocked: git write operations require higher autonomy level", + )); + } + + if self.gate.is_read_only() { + return Ok(ToolResult::error( + "[policy-blocked] Action blocked: read-only mode", + )); + } + } + + // Record action for rate limiting + if !self.gate.record_action() { + return Ok(ToolResult::error("Action blocked: rate limit exceeded")); + } + + // Execute the requested operation + match operation { + "status" => self.git_status(&effective_dir, args).await, + "diff" => self.git_diff(&effective_dir, args).await, + "log" => self.git_log(&effective_dir, args).await, + "branch" => self.git_branch(&effective_dir, args).await, + "commit" => self.git_commit(&effective_dir, args).await, + "add" => self.git_add(&effective_dir, args).await, + "checkout" => self.git_checkout(&effective_dir, args).await, + "stash" => self.git_stash(&effective_dir, args).await, + _ => Ok(ToolResult::error(format!("Unknown operation: {operation}"))), + } + } +} + +#[cfg(test)] +#[path = "config_tests.rs"] +mod config_tests; +#[cfg(test)] +#[path = "ops_tests.rs"] +mod ops_tests; +#[cfg(test)] +#[path = "render_tests.rs"] +mod render_tests; +#[cfg(test)] +#[path = "mod_tests.rs"] +mod tests; diff --git a/crates/tinyagents-harness-tools/src/fs/filesystem/git_operations/mod_tests.rs b/crates/tinyagents-harness-tools/src/fs/filesystem/git_operations/mod_tests.rs new file mode 100644 index 000000000..92f09a069 --- /dev/null +++ b/crates/tinyagents-harness-tools/src/fs/filesystem/git_operations/mod_tests.rs @@ -0,0 +1,455 @@ +//! Behavior tests for the `git_operations` tool, driven through a fake [`FsGate`]. + +#![allow(clippy::expect_used, clippy::panic, clippy::unwrap_used)] + +use super::config::NULL_CONFIG_PATH; +use super::*; +use crate::fs::filesystem::test_support::{AutonomyLevel, TestGate, WorkspaceContext}; +use tempfile::TempDir; + +pub(super) fn test_tool(dir: &std::path::Path) -> GitOperationsTool { + let security = TestGate::with(dir.to_path_buf(), AutonomyLevel::Supervised, 1_000_000); + GitOperationsTool::new(security, dir.to_path_buf()) +} + +#[test] +fn sanitize_git_blocks_injection() { + let tmp = TempDir::new().unwrap(); + let tool = test_tool(tmp.path()); + + // Should block dangerous arguments + assert!(tool.sanitize_git_args("--exec=rm -rf /").is_err()); + assert!(tool.sanitize_git_args("$(echo pwned)").is_err()); + assert!(tool.sanitize_git_args("`malicious`").is_err()); + assert!(tool.sanitize_git_args("arg | cat").is_err()); + assert!(tool.sanitize_git_args("arg; rm file").is_err()); +} + +#[test] +fn sanitize_git_blocks_pager_editor_injection() { + let tmp = TempDir::new().unwrap(); + let tool = test_tool(tmp.path()); + + assert!(tool.sanitize_git_args("--pager=less").is_err()); + assert!(tool.sanitize_git_args("--editor=vim").is_err()); +} + +#[test] +fn sanitize_git_blocks_config_injection() { + let tmp = TempDir::new().unwrap(); + let tool = test_tool(tmp.path()); + + // Exact `-c` flag (config injection) + assert!(tool.sanitize_git_args("-c core.sshCommand=evil").is_err()); + assert!(tool.sanitize_git_args("-c=core.pager=less").is_err()); +} + +#[test] +fn sanitize_git_blocks_no_verify() { + let tmp = TempDir::new().unwrap(); + let tool = test_tool(tmp.path()); + + assert!(tool.sanitize_git_args("--no-verify").is_err()); +} + +#[test] +fn sanitize_git_blocks_redirect_in_args() { + let tmp = TempDir::new().unwrap(); + let tool = test_tool(tmp.path()); + + assert!(tool.sanitize_git_args("file.txt > /tmp/out").is_err()); +} + +#[test] +fn sanitize_git_cached_not_blocked() { + let tmp = TempDir::new().unwrap(); + let tool = test_tool(tmp.path()); + + // --cached must NOT be blocked by the `-c` check + assert!(tool.sanitize_git_args("--cached").is_ok()); + // Other safe flags starting with -c prefix + assert!(tool.sanitize_git_args("-cached").is_ok()); +} + +#[test] +fn sanitize_git_allows_safe() { + let tmp = TempDir::new().unwrap(); + let tool = test_tool(tmp.path()); + + // Should allow safe arguments + assert!(tool.sanitize_git_args("main").is_ok()); + assert!(tool.sanitize_git_args("feature/test-branch").is_ok()); + assert!(tool.sanitize_git_args("--cached").is_ok()); + assert!(tool.sanitize_git_args("src/main.rs").is_ok()); + assert!(tool.sanitize_git_args(".").is_ok()); +} + +/// Parity guard for the worktree-isolation action-dir override (#3376, +/// #4249 08.5). A worktree-isolated worker's git operation MUST resolve its CWD +/// from the carried `WorkspaceDescriptor` (the isolated worktree), never the +/// tool's configured `action_dir`. WITHOUT a descriptor it falls back to +/// `action_dir` — the non-isolated path, byte-identical to before. This encodes +/// the behaviour the deleted `worktree_context.rs` task-local used to provide. +#[test] +fn git_resolves_cwd_from_workspace_descriptor() { + let action_tmp = TempDir::new().unwrap(); + let worktree_tmp = TempDir::new().unwrap(); + let tool = test_tool(action_tmp.path()); + + // WITH a descriptor → the worktree root wins. + let tool_ctx = WorkspaceContext::at(worktree_tmp.path()); + assert_eq!( + tool.effective_action_dir_for_context(Some(&tool_ctx)), + worktree_tmp.path().to_path_buf(), + "git with a WorkspaceDescriptor must resolve CWD to the worktree root" + ); + + // WITHOUT a descriptor → configured action_dir (non-isolated parity). + assert_eq!( + tool.effective_action_dir_for_context(None), + action_tmp.path().to_path_buf(), + "git with no descriptor must fall back to the configured action_dir" + ); +} + +#[test] +fn requires_write_detection() { + let tmp = TempDir::new().unwrap(); + let tool = test_tool(tmp.path()); + + assert!(tool.requires_write_access("commit")); + assert!(tool.requires_write_access("add")); + assert!(tool.requires_write_access("checkout")); + + assert!(!tool.requires_write_access("status")); + assert!(!tool.requires_write_access("diff")); + assert!(!tool.requires_write_access("log")); +} + +#[test] +fn branch_is_not_write_gated() { + let tmp = TempDir::new().unwrap(); + let tool = test_tool(tmp.path()); + + // Branch listing is read-only; it must not require write access + assert!(!tool.requires_write_access("branch")); + assert!(tool.is_read_only("branch")); +} + +#[test] +fn is_read_only_detection() { + let tmp = TempDir::new().unwrap(); + let tool = test_tool(tmp.path()); + + assert!(tool.is_read_only("status")); + assert!(tool.is_read_only("diff")); + assert!(tool.is_read_only("log")); + assert!(tool.is_read_only("branch")); + + assert!(!tool.is_read_only("commit")); + assert!(!tool.is_read_only("add")); +} + +#[tokio::test] +async fn blocks_readonly_mode_for_write_ops() { + let tmp = TempDir::new().unwrap(); + init_git_repo(tmp.path()); + + let security = TestGate::with(tmp.path().to_path_buf(), AutonomyLevel::ReadOnly, 1_000_000); + let tool = GitOperationsTool::new(security, tmp.path().to_path_buf()); + + let result = tool + .execute(json!({"operation": "commit", "message": "test"})) + .await + .unwrap(); + assert!(result.is_error); + // can_act() returns false for ReadOnly, so we get the "higher autonomy level" message + assert!(result.output().contains("higher autonomy")); +} + +#[tokio::test] +async fn allows_branch_listing_in_readonly_mode() { + let tmp = TempDir::new().unwrap(); + init_git_repo(tmp.path()); + + let security = TestGate::with(tmp.path().to_path_buf(), AutonomyLevel::ReadOnly, 1_000_000); + let tool = GitOperationsTool::new(security, tmp.path().to_path_buf()); + + let result = tool.execute(json!({"operation": "branch"})).await.unwrap(); + // Branch listing must not be blocked by read-only autonomy + let error_msg = result.output(); + assert!( + !error_msg.contains("read-only") && !error_msg.contains("higher autonomy"), + "branch listing should not be blocked in read-only mode, got: {error_msg}" + ); +} + +#[tokio::test] +async fn allows_readonly_ops_in_readonly_mode() { + let tmp = TempDir::new().unwrap(); + let security = TestGate::with(tmp.path().to_path_buf(), AutonomyLevel::ReadOnly, 1_000_000); + let tool = GitOperationsTool::new(security, tmp.path().to_path_buf()); + + // This will fail because there's no git repo, but it shouldn't be blocked by autonomy + let result = tool.execute(json!({"operation": "status"})).await.unwrap(); + // The error should be about git (not about autonomy/read-only mode) + assert!(result.is_error, "Expected failure due to missing git repo"); + let error_msg = result.output(); + assert!( + !error_msg.is_empty(), + "Expected a git-related error message" + ); + assert!( + !error_msg.contains("read-only") && !error_msg.contains("autonomy"), + "Error should be about git, not about autonomy restrictions: {error_msg}" + ); +} + +#[tokio::test] +async fn rejects_missing_operation() { + let tmp = TempDir::new().unwrap(); + let tool = test_tool(tmp.path()); + + let result = tool.execute(json!({})).await.unwrap(); + assert!(result.is_error); + assert!(result.output().contains("Missing 'operation'")); +} + +#[tokio::test] +async fn rejects_unknown_operation() { + let tmp = TempDir::new().unwrap(); + init_git_repo(tmp.path()); + + let tool = test_tool(tmp.path()); + + let result = tool.execute(json!({"operation": "push"})).await.unwrap(); + assert!(result.is_error); + assert!(result.output().contains("Unknown operation")); +} + +#[test] +fn truncates_multibyte_commit_message_without_panicking() { + let long = "🦀".repeat(2500); + let truncated = GitOperationsTool::truncate_commit_message(&long); + + assert_eq!(truncated.chars().count(), 2000); +} + +// ── truncate_commit_message: short messages pass through unchanged ───────── + +#[test] +fn truncate_short_message_unchanged() { + let msg = "Fix the bug"; + assert_eq!(GitOperationsTool::truncate_commit_message(msg), msg); +} + +#[test] +fn truncate_exact_2000_chars_unchanged() { + let msg = "a".repeat(2000); + let result = GitOperationsTool::truncate_commit_message(&msg); + assert_eq!(result.chars().count(), 2000); + assert!(!result.ends_with("...")); +} + +#[test] +fn truncate_2001_chars_adds_ellipsis() { + let msg = "a".repeat(2001); + let result = GitOperationsTool::truncate_commit_message(&msg); + assert!(result.ends_with("...")); + assert_eq!(result.chars().count(), 2000); +} + +// ── sanitize_git_args: allow leading dash that is not -c ───────────────── + +#[test] +fn sanitize_git_allows_other_flags() { + let tmp = TempDir::new().unwrap(); + let tool = test_tool(tmp.path()); + assert!(tool.sanitize_git_args("--follow").is_ok()); + assert!(tool.sanitize_git_args("-p").is_ok()); + assert!(tool.sanitize_git_args("-n5").is_ok()); +} + +// ── requires_write_access completeness ──────────────────────────────────── + +#[test] +fn requires_write_access_covers_all_write_ops() { + let tmp = TempDir::new().unwrap(); + let tool = test_tool(tmp.path()); + for op in ["commit", "add", "checkout", "stash", "reset", "revert"] { + assert!( + tool.requires_write_access(op), + "'{op}' should require write access" + ); + } +} + +// ── schema validation ───────────────────────────────────────────────────── + +#[test] +fn schema_has_required_operation() { + let tmp = TempDir::new().unwrap(); + let tool = test_tool(tmp.path()); + let schema = tool.parameters_schema(); + let required = schema["required"].as_array().unwrap(); + assert!( + required.contains(&serde_json::json!("operation")), + "schema required should include 'operation'" + ); +} + +#[test] +fn schema_enumerates_operations() { + let tmp = TempDir::new().unwrap(); + let tool = test_tool(tmp.path()); + let schema = tool.parameters_schema(); + let ops = schema["properties"]["operation"]["enum"] + .as_array() + .unwrap(); + let op_names: Vec<&str> = ops.iter().map(|v| v.as_str().unwrap()).collect(); + for expected in [ + "status", "diff", "log", "branch", "commit", "add", "checkout", "stash", + ] { + assert!( + op_names.contains(&expected), + "schema should include '{expected}'" + ); + } +} + +// ── git_operations tool name / description ──────────────────────────────── + +#[test] +fn tool_name_and_description() { + let tmp = TempDir::new().unwrap(); + let tool = test_tool(tmp.path()); + assert_eq!(tool.name(), "git_operations"); + assert_ne!(tool.description().len(), 0); + assert!(tool.description().contains("Git")); +} + +// ── not_in_git_repo returns error (covers the git-repo check) ───────────── + +#[tokio::test] +async fn not_in_git_repo_returns_error() { + let tmp = TempDir::new().unwrap(); + // Do NOT init a git repo + let tool = test_tool(tmp.path()); + let result = tool.execute(json!({"operation": "status"})).await.unwrap(); + assert!(result.is_error); + assert!(result.output().contains("Not in a git repository")); +} + +/// Suppress the developer's own system/global git config on a raw +/// `std::process::Command`, so a machine-local `init.templateDir` or similar +/// cannot write extra keys into a test repository's `.git/config` and make +/// these tests depend on ambient environment. Mirrors [`hardened_git`]'s two +/// env vars; the production code under test applies its own suppression when +/// it later reads this same config, so this only affects setup. +pub(super) fn hermetic(cmd: &mut std::process::Command) -> &mut std::process::Command { + cmd.env("GIT_CONFIG_NOSYSTEM", "1") + .env("GIT_CONFIG_GLOBAL", NULL_CONFIG_PATH) + .env_remove("GIT_CONFIG_PARAMETERS") + .env_remove("GIT_CONFIG_COUNT") +} + +/// Initialise a git repo at `path` and fail the test if `git init` +/// itself didn't succeed (so we don't misread later assertion failures +/// as product bugs when the real problem is a missing/broken git). +pub(super) fn init_git_repo(path: &std::path::Path) { + let output = hermetic( + std::process::Command::new("git") + .args(["init"]) + .current_dir(path), + ) + .output() + .expect("failed to spawn `git init`"); + assert!( + output.status.success(), + "`git init` failed: {}", + String::from_utf8_lossy(&output.stderr) + ); +} + +/// Extract the error text from a Result — whether the +/// failure came through `Err(anyhow::Error)` or `Ok(ToolResult::error)`. +pub(super) fn error_text(result: &anyhow::Result) -> String { + match result { + Ok(r) => { + assert!(r.is_error, "expected a tool-error ToolResult"); + r.output().clone() + } + Err(e) => e.to_string(), + } +} + +// ── stash: unknown action returns error ──────────────────────────────────── + +#[tokio::test] +async fn stash_unknown_action_returns_error() { + let tmp = TempDir::new().unwrap(); + init_git_repo(tmp.path()); + + let tool = test_tool(tmp.path()); + let result = tool + .execute(json!({"operation": "stash", "action": "squash"})) + .await; + let msg = error_text(&result); + assert!( + msg.contains("Unknown stash action"), + "expected 'Unknown stash action' in error, got: {msg}" + ); +} + +// ── checkout: dangerous characters ──────────────────────────────────────── + +#[tokio::test] +async fn checkout_rejects_dangerous_branch_names() { + let tmp = TempDir::new().unwrap(); + init_git_repo(tmp.path()); + + let tool = test_tool(tmp.path()); + + for dangerous in ["main@{1}", "HEAD^", "v1~2"] { + let result = tool + .execute(json!({"operation": "checkout", "branch": dangerous})) + .await; + let msg = error_text(&result); + assert!( + msg.contains("invalid characters") || msg.contains("Invalid branch"), + "expected a dangerous-branch rejection for '{dangerous}', got: {msg}" + ); + } +} + +// ── commit: missing message ──────────────────────────────────────────────── + +#[tokio::test] +async fn commit_missing_message_returns_error() { + let tmp = TempDir::new().unwrap(); + init_git_repo(tmp.path()); + + let tool = test_tool(tmp.path()); + let result = tool.execute(json!({"operation": "commit"})).await; + let msg = error_text(&result); + assert!( + msg.contains("Missing 'message' parameter"), + "expected missing-message error, got: {msg}" + ); +} + +// ── add: missing paths ───────────────────────────────────────────────────── + +#[tokio::test] +async fn add_missing_paths_returns_error() { + let tmp = TempDir::new().unwrap(); + init_git_repo(tmp.path()); + + let tool = test_tool(tmp.path()); + let result = tool.execute(json!({"operation": "add"})).await; + let msg = error_text(&result); + assert!( + msg.contains("Missing 'paths' parameter"), + "expected missing-paths error, got: {msg}" + ); +} diff --git a/crates/tinyagents-harness-tools/src/fs/filesystem/git_operations/ops_tests.rs b/crates/tinyagents-harness-tools/src/fs/filesystem/git_operations/ops_tests.rs new file mode 100644 index 000000000..e9a8ff799 --- /dev/null +++ b/crates/tinyagents-harness-tools/src/fs/filesystem/git_operations/ops_tests.rs @@ -0,0 +1,368 @@ +//! End-to-end tests of the `git_operations` sub-operations against real, +//! hermetic repositories: parsing, success paths, failure paths and gate +//! refusals. + +#![allow(clippy::expect_used, clippy::panic, clippy::unwrap_used)] + +use super::tests::{hermetic, init_git_repo, test_tool}; +use super::*; +use crate::fs::filesystem::test_support::{AutonomyLevel, TestGate, WorkspaceContext}; +use tempfile::TempDir; + +fn git(dir: &std::path::Path, args: &[&str]) { + let out = hermetic( + std::process::Command::new("git") + .args(args) + .current_dir(dir), + ) + .output() + .expect("spawn git"); + assert!( + out.status.success(), + "git {args:?} failed: {}", + String::from_utf8_lossy(&out.stderr) + ); +} + +/// A repo with one commit of `a.txt` and identity configured locally. +fn repo_with_commit() -> TempDir { + let tmp = TempDir::new().unwrap(); + init_git_repo(tmp.path()); + git(tmp.path(), &["config", "user.email", "t@example.invalid"]); + git(tmp.path(), &["config", "user.name", "Tester"]); + git(tmp.path(), &["checkout", "-b", "main"]); + std::fs::write(tmp.path().join("a.txt"), "one\ntwo\nthree\n").unwrap(); + git(tmp.path(), &["add", "a.txt"]); + git(tmp.path(), &["commit", "-m", "first commit"]); + tmp +} + +async fn run(tool: &GitOperationsTool, args: serde_json::Value) -> ToolResult { + tool.execute(args).await.unwrap() +} + +#[test] +fn supports_markdown_is_advertised() { + let tmp = TempDir::new().unwrap(); + assert!(test_tool(tmp.path()).supports_markdown()); +} + +#[tokio::test] +async fn status_parses_staged_unstaged_and_untracked() { + let tmp = repo_with_commit(); + std::fs::write(tmp.path().join("a.txt"), "one\nTWO\nthree\n").unwrap(); + git(tmp.path(), &["add", "a.txt"]); + std::fs::write(tmp.path().join("a.txt"), "changed again\n").unwrap(); + std::fs::write(tmp.path().join("new.txt"), "n").unwrap(); + let tool = test_tool(tmp.path()); + + let result = run(&tool, json!({"operation": "status"})).await; + assert!(!result.is_error, "{}", result.output()); + let parsed: serde_json::Value = serde_json::from_str(&result.output()).unwrap(); + assert_ne!(parsed["staged"].as_array().unwrap().len(), 0); + assert_ne!(parsed["unstaged"].as_array().unwrap().len(), 0); + assert!( + parsed["untracked"] + .as_array() + .unwrap() + .contains(&json!("new.txt")) + ); + assert!(result.markdown_formatted.is_some()); +} + +#[tokio::test] +async fn diff_reports_hunks_for_unstaged_and_cached_changes() { + let tmp = repo_with_commit(); + std::fs::write(tmp.path().join("a.txt"), "one\nTWO\nthree\n").unwrap(); + std::fs::write(tmp.path().join("b.txt"), "x\n").unwrap(); + git(tmp.path(), &["add", "b.txt"]); + let tool = test_tool(tmp.path()); + + let unstaged = run(&tool, json!({"operation": "diff"})).await; + let parsed: serde_json::Value = serde_json::from_str(&unstaged.output()).unwrap(); + assert!(parsed["file_count"].as_u64().unwrap() >= 1); + assert_eq!(parsed["hunks"][0]["file"], "a.txt"); + let kinds: Vec<&str> = parsed["hunks"] + .as_array() + .unwrap() + .iter() + .flat_map(|h| h["lines"].as_array().unwrap().iter()) + .map(|l| l["type"].as_str().unwrap()) + .collect(); + assert!(kinds.contains(&"add") && kinds.contains(&"delete") && kinds.contains(&"context")); + + let cached = run(&tool, json!({"operation": "diff", "cached": true})).await; + let parsed: serde_json::Value = serde_json::from_str(&cached.output()).unwrap(); + assert_eq!(parsed["hunks"][0]["file"], "b.txt"); +} + +#[tokio::test] +async fn diff_with_several_files_and_hunks_splits_them() { + let tmp = repo_with_commit(); + let long: String = (0..40).fold(String::new(), |acc, i| acc + &format!("line{i}\n")); + std::fs::write(tmp.path().join("big.txt"), &long).unwrap(); + git(tmp.path(), &["add", "big.txt"]); + git(tmp.path(), &["commit", "-m", "big"]); + let edited = long + .replace("line1\n", "LINE1\n") + .replace("line38\n", "LINE38\n"); + std::fs::write(tmp.path().join("big.txt"), edited).unwrap(); + std::fs::write(tmp.path().join("a.txt"), "different\n").unwrap(); + let tool = test_tool(tmp.path()); + + let result = run(&tool, json!({"operation": "diff"})).await; + let parsed: serde_json::Value = serde_json::from_str(&result.output()).unwrap(); + assert!(parsed["file_count"].as_u64().unwrap() >= 3); + assert!( + parsed["hunks"] + .as_array() + .unwrap() + .iter() + .any(|h| h["header"].is_string()) + ); +} + +#[tokio::test] +async fn diff_rejects_injected_file_arguments() { + let tmp = repo_with_commit(); + let tool = test_tool(tmp.path()); + let result = tool + .execute(json!({"operation": "diff", "files": "--exec=evil"})) + .await; + assert!(result.is_err() || result.unwrap().is_error); +} + +#[tokio::test] +async fn log_lists_commits_with_markdown() { + let tmp = repo_with_commit(); + let tool = test_tool(tmp.path()); + let result = run(&tool, json!({"operation": "log", "limit": 5})).await; + assert!(!result.is_error, "{}", result.output()); + let parsed: serde_json::Value = serde_json::from_str(&result.output()).unwrap(); + assert_eq!(parsed["commits"][0]["message"], "first commit"); + assert_eq!(parsed["commits"][0]["author"], "Tester"); + assert!(result.markdown_formatted.unwrap().contains("first commit")); +} + +#[tokio::test] +async fn branch_lists_the_current_branch() { + let tmp = repo_with_commit(); + git(tmp.path(), &["branch", "side"]); + let tool = test_tool(tmp.path()); + let result = run(&tool, json!({"operation": "branch"})).await; + let parsed: serde_json::Value = serde_json::from_str(&result.output()).unwrap(); + assert_eq!(parsed["current"], "main"); + assert_eq!(parsed["branches"].as_array().unwrap().len(), 2); + assert!(result.markdown_formatted.unwrap().contains("← current")); +} + +#[tokio::test] +async fn add_then_commit_succeeds() { + let tmp = repo_with_commit(); + std::fs::write(tmp.path().join("c.txt"), "c").unwrap(); + let tool = test_tool(tmp.path()); + + let added = run(&tool, json!({"operation": "add", "paths": "c.txt"})).await; + assert_eq!(added.output(), "Staged: c.txt"); + + let committed = run( + &tool, + json!({"operation": "commit", "message": " add c \n\n more "}), + ) + .await; + assert!(!committed.is_error, "{}", committed.output()); + assert_eq!(committed.output(), "Committed: add c\nmore"); +} + +#[tokio::test] +async fn add_of_a_missing_path_reports_failure() { + let tmp = repo_with_commit(); + let tool = test_tool(tmp.path()); + let result = run(&tool, json!({"operation": "add", "paths": "nope.txt"})).await; + assert!(result.is_error); + assert!(result.output().starts_with("Add failed:")); +} + +#[tokio::test] +async fn commit_with_nothing_staged_reports_failure() { + let tmp = repo_with_commit(); + let tool = test_tool(tmp.path()); + let result = run(&tool, json!({"operation": "commit", "message": "empty"})).await; + assert!(result.is_error); + assert!(result.output().starts_with("Commit failed:")); +} + +#[tokio::test] +async fn commit_with_blank_message_is_an_error() { + let tmp = repo_with_commit(); + let tool = test_tool(tmp.path()); + let result = tool + .execute(json!({"operation": "commit", "message": " \n "})) + .await; + assert!(result.unwrap_err().to_string().contains("cannot be empty")); +} + +#[test] +fn overlong_commit_messages_are_truncated() { + let long = "x".repeat(2500); + let out = GitOperationsTool::truncate_commit_message(&long); + assert_eq!(out.chars().count(), 2000); + assert!(out.ends_with("...")); + assert_eq!(GitOperationsTool::truncate_commit_message("short"), "short"); +} + +#[tokio::test] +async fn checkout_switches_and_reports_failures() { + let tmp = repo_with_commit(); + git(tmp.path(), &["branch", "side"]); + let tool = test_tool(tmp.path()); + + let ok = run(&tool, json!({"operation": "checkout", "branch": "side"})).await; + assert_eq!(ok.output(), "Switched to branch: side"); + + let missing = run(&tool, json!({"operation": "checkout", "branch": "ghost"})).await; + assert!(missing.is_error); + assert!(missing.output().starts_with("Checkout failed:")); + + let invalid = tool + .execute(json!({"operation": "checkout", "branch": "--evil"})) + .await; + assert!(invalid.is_err() || invalid.unwrap().is_error); +} + +#[tokio::test] +async fn stash_push_list_pop_and_drop() { + let tmp = repo_with_commit(); + let tool = test_tool(tmp.path()); + std::fs::write(tmp.path().join("a.txt"), "dirty\n").unwrap(); + + let push = run(&tool, json!({"operation": "stash", "action": "push"})).await; + assert!(!push.is_error, "{}", push.output()); + let list = run(&tool, json!({"operation": "stash", "action": "list"})).await; + assert!(list.output().contains("auto-stash")); + let pop = run(&tool, json!({"operation": "stash", "action": "pop"})).await; + assert!(!pop.is_error, "{}", pop.output()); + + git(tmp.path(), &["stash", "push", "-m", "again"]); + let drop = run( + &tool, + json!({"operation": "stash", "action": "drop", "index": 0}), + ) + .await; + assert!(!drop.is_error, "{}", drop.output()); + + let bad = run( + &tool, + json!({"operation": "stash", "action": "drop", "index": 7}), + ) + .await; + assert!(bad.is_error); + assert!(bad.output().starts_with("Stash drop failed:")); + + let huge = tool + .execute(json!({"operation": "stash", "action": "drop", "index": 4_000_000_000u64})) + .await; + assert!(huge.unwrap_err().to_string().contains("too large")); +} + +#[tokio::test] +async fn pop_with_no_stash_reports_failure() { + let tmp = repo_with_commit(); + let tool = test_tool(tmp.path()); + let result = run(&tool, json!({"operation": "stash", "action": "pop"})).await; + assert!(result.is_error); + assert!(result.output().starts_with("Stash pop failed:")); +} + +#[tokio::test] +async fn write_operations_are_refused_when_the_gate_is_read_only() { + let tmp = repo_with_commit(); + let gate = TestGate::with(tmp.path().to_path_buf(), AutonomyLevel::ReadOnly, 1_000_000); + let tool = GitOperationsTool::new(gate, tmp.path().to_path_buf()); + let result = run(&tool, json!({"operation": "add", "paths": "a.txt"})).await; + assert!(result.is_error); + assert!(result.output().contains("higher autonomy level")); +} + +/// A gate that permits acting but reports strict read-only mode. +#[derive(Debug)] +struct ReadOnlyButActing(std::path::PathBuf); + +#[async_trait] +impl FsGate for ReadOnlyButActing { + fn can_act(&self) -> bool { + true + } + fn is_read_only(&self) -> bool { + true + } + fn is_rate_limited(&self) -> bool { + false + } + fn record_action(&self) -> bool { + true + } + fn write_needs_approval(&self) -> bool { + false + } + fn action_dir(&self) -> &Path { + &self.0 + } + fn is_path_string_allowed(&self, _path: &str) -> bool { + true + } + async fn validate_path(&self, path: &str) -> Result { + Ok(PathBuf::from(path)) + } + async fn validate_parent_path(&self, path: &str) -> Result { + Ok(PathBuf::from(path)) + } + fn scoped_to_workspace(&self, root: &Path) -> Arc { + Arc::new(Self(root.to_path_buf())) + } +} + +#[tokio::test] +async fn read_only_mode_blocks_writes_even_when_the_gate_allows_acting() { + let tmp = repo_with_commit(); + let gate: Arc = Arc::new(ReadOnlyButActing(tmp.path().to_path_buf())); + let tool = GitOperationsTool::new(gate, tmp.path().to_path_buf()); + let result = run(&tool, json!({"operation": "add", "paths": "a.txt"})).await; + assert_eq!( + result.output(), + "[policy-blocked] Action blocked: read-only mode" + ); +} + +#[tokio::test] +async fn an_exhausted_action_budget_blocks_the_operation() { + let tmp = repo_with_commit(); + let gate = TestGate::with(tmp.path().to_path_buf(), AutonomyLevel::Supervised, 0); + let tool = GitOperationsTool::new(gate, tmp.path().to_path_buf()); + let result = run(&tool, json!({"operation": "status"})).await; + assert_eq!(result.output(), "Action blocked: rate limit exceeded"); +} + +#[tokio::test] +async fn the_execute_wrappers_agree_and_honour_a_workspace_descriptor() { + let unrelated = TempDir::new().unwrap(); + let repo = repo_with_commit(); + let tool = test_tool(unrelated.path()); + let context = WorkspaceContext::at(repo.path()); + + let with_options = tool + .execute_with_options(json!({"operation": "branch"}), ToolCallOptions::default()) + .await + .unwrap(); + assert!(with_options.is_error, "the unrelated dir is not a repo"); + + let with_context = tool + .execute_with_context( + json!({"operation": "branch"}), + ToolCallOptions::default(), + Some(&context), + ) + .await + .unwrap(); + assert!(!with_context.is_error, "{}", with_context.output()); +} diff --git a/crates/tinyagents-harness-tools/src/fs/filesystem/git_operations/render.rs b/crates/tinyagents-harness-tools/src/fs/filesystem/git_operations/render.rs new file mode 100644 index 000000000..124c355a1 --- /dev/null +++ b/crates/tinyagents-harness-tools/src/fs/filesystem/git_operations/render.rs @@ -0,0 +1,91 @@ +//! Markdown renderers for [`super::git_operations`] results. +//! +//! Split out of `git_operations.rs` for the Rust layout gate. Pure formatting: +//! these take already-parsed JSON and produce the human-readable summary, with +//! no git invocation of their own. + +pub(super) fn render_status_markdown( + result: &serde_json::Map, +) -> String { + let mut out = String::new(); + if let Some(branch) = result.get("branch").and_then(|v| v.as_str()) { + out.push_str(&format!("**branch**: `{branch}`\n")); + } + let clean = result + .get("clean") + .and_then(serde_json::Value::as_bool) + .unwrap_or(false); + if clean { + out.push_str("_Working tree clean._\n"); + return out; + } + let push_section = |out: &mut String, label: &str, items: Option<&Vec>| { + if let Some(items) = items + && !items.is_empty() + { + out.push_str(&format!("\n**{label}** ({})\n", items.len())); + for it in items { + if let (Some(p), Some(s)) = ( + it.get("path").and_then(|v| v.as_str()), + it.get("status").and_then(|v| v.as_str()), + ) { + out.push_str(&format!("- `{s}` {p}\n")); + } + } + } + }; + push_section( + &mut out, + "staged", + result.get("staged").and_then(|v| v.as_array()), + ); + push_section( + &mut out, + "unstaged", + result.get("unstaged").and_then(|v| v.as_array()), + ); + if let Some(items) = result.get("untracked").and_then(|v| v.as_array()) + && !items.is_empty() + { + out.push_str(&format!("\n**untracked** ({})\n", items.len())); + for it in items { + if let Some(p) = it.as_str() { + out.push_str(&format!("- {p}\n")); + } + } + } + out +} + +pub(super) fn render_log_markdown(commits: &[serde_json::Value]) -> String { + if commits.is_empty() { + return "_No commits._".to_string(); + } + let mut out = format!("# Commits ({})\n", commits.len()); + for c in commits { + let hash = c.get("hash").and_then(|v| v.as_str()).unwrap_or(""); + let short = hash.get(..hash.len().min(8)).unwrap_or(hash); + let author = c.get("author").and_then(|v| v.as_str()).unwrap_or(""); + let date = c.get("date").and_then(|v| v.as_str()).unwrap_or(""); + let msg = c.get("message").and_then(|v| v.as_str()).unwrap_or(""); + out.push_str(&format!("- `{short}` {msg} _(by {author}, {date})_\n")); + } + out +} + +pub(super) fn render_branch_markdown(current: &str, branches: &[serde_json::Value]) -> String { + let mut out = format!("**current**: `{current}`\n\n## Branches\n"); + for b in branches { + let name = b.get("name").and_then(|v| v.as_str()).unwrap_or(""); + let cur = b + .get("current") + .and_then(serde_json::Value::as_bool) + .unwrap_or(false); + if cur { + out.push_str(&format!("- **{name}** ← current\n")); + } else { + out.push_str(&format!("- {name}\n")); + } + } + out +} diff --git a/crates/tinyagents-harness-tools/src/fs/filesystem/git_operations/render_tests.rs b/crates/tinyagents-harness-tools/src/fs/filesystem/git_operations/render_tests.rs new file mode 100644 index 000000000..f41c5077b --- /dev/null +++ b/crates/tinyagents-harness-tools/src/fs/filesystem/git_operations/render_tests.rs @@ -0,0 +1,83 @@ +//! Tests for the markdown renderers: literal JSON in, literal markdown out. + +#![allow(clippy::unwrap_used)] + +use super::render::{render_branch_markdown, render_log_markdown, render_status_markdown}; +use serde_json::{Value, json}; + +fn status(value: &Value) -> String { + render_status_markdown(value.as_object().unwrap()) +} + +#[test] +fn clean_status_stops_after_the_branch_line() { + let md = + status(&json!({"branch": "main", "clean": true, "staged": [{"path": "a", "status": "M"}]})); + assert_eq!(md, "**branch**: `main`\n_Working tree clean._\n"); +} + +#[test] +fn dirty_status_lists_each_section_with_counts() { + let md = status(&json!({ + "branch": "feat", + "clean": false, + "staged": [{"path": "a.rs", "status": "M"}, {"path": "b.rs", "status": "A"}], + "unstaged": [{"path": "c.rs", "status": "D"}], + "untracked": ["new.txt"], + })); + assert_eq!( + md, + "**branch**: `feat`\n\ + \n**staged** (2)\n- `M` a.rs\n- `A` b.rs\n\ + \n**unstaged** (1)\n- `D` c.rs\n\ + \n**untracked** (1)\n- new.txt\n" + ); +} + +#[test] +fn status_skips_empty_sections_and_malformed_entries() { + let md = status(&json!({ + "staged": [], + "unstaged": [{"path": "only-path"}, {"status": "M"}], + "untracked": [42], + })); + assert_eq!(md, "\n**unstaged** (2)\n\n**untracked** (1)\n"); +} + +#[test] +fn status_without_a_branch_or_clean_flag_renders_nothing() { + assert_eq!(status(&json!({})), ""); +} + +#[test] +fn log_renders_short_hashes_and_defaults_missing_fields() { + assert_eq!(render_log_markdown(&[]), "_No commits._"); + let md = render_log_markdown(&[ + json!({"hash": "0123456789abcdef", "author": "Ann", "date": "2026-01-02", "message": "fix"}), + json!({"hash": "abc"}), + json!({}), + ]); + assert_eq!( + md, + "# Commits (3)\n\ + - `01234567` fix _(by Ann, 2026-01-02)_\n\ + - `abc` _(by , )_\n\ + - `` _(by , )_\n" + ); +} + +#[test] +fn branches_mark_the_current_one() { + let md = render_branch_markdown( + "main", + &[ + json!({"name": "main", "current": true}), + json!({"name": "dev"}), + json!({}), + ], + ); + assert_eq!( + md, + "**current**: `main`\n\n## Branches\n- **main** ← current\n- dev\n- \n" + ); +} diff --git a/crates/tinyagents-harness-tools/src/fs/filesystem/glob_search/extra_tests.rs b/crates/tinyagents-harness-tools/src/fs/filesystem/glob_search/extra_tests.rs new file mode 100644 index 000000000..108328cc8 --- /dev/null +++ b/crates/tinyagents-harness-tools/src/fs/filesystem/glob_search/extra_tests.rs @@ -0,0 +1,153 @@ +//! Budget, root-validation, truncation and policy-filter tests for `glob`. + +#![allow(clippy::expect_used, clippy::panic, clippy::unwrap_used)] + +use super::*; +use crate::fs::filesystem::test_support::{TestGate, WorkspaceContext, WrapGate}; +use tempfile::TempDir; + +#[test] +fn glob_is_concurrency_safe() { + let dir = TempDir::new().unwrap(); + let tool = GlobTool::new(TestGate::supervised(dir.path().to_path_buf())); + assert!(tool.is_concurrency_safe(&json!({}))); +} + +#[tokio::test] +async fn rate_limit_and_exhausted_budget_are_refused() { + let dir = TempDir::new().unwrap(); + let limited = GlobTool::new(TestGate::with( + dir.path().to_path_buf(), + crate::fs::filesystem::test_support::AutonomyLevel::Supervised, + 0, + )); + let r = limited.execute(json!({"pattern": "*"})).await.unwrap(); + assert!(r.output().contains("too many actions in the last hour")); + + let gate = WrapGate::new(TestGate::supervised(dir.path().to_path_buf())) + .no_budget() + .arc(); + let r = GlobTool::new(gate) + .execute(json!({"pattern": "*"})) + .await + .unwrap(); + assert!(r.output().contains("action budget exhausted")); +} + +#[tokio::test] +async fn a_root_that_vanishes_after_validation_is_reported() { + let dir = TempDir::new().unwrap(); + let gate = WrapGate::new(TestGate::supervised(dir.path().to_path_buf())) + .fixed_path(dir.path().join("gone")) + .arc(); + let r = GlobTool::new(gate) + .execute(json!({"pattern": "*"})) + .await + .unwrap(); + assert!(r.is_error); + assert!(r.output().contains("is not accessible")); +} + +#[tokio::test] +async fn a_missing_action_dir_falls_back_to_its_raw_path() { + let dir = TempDir::new().unwrap(); + std::fs::write(dir.path().join("a.txt"), "").unwrap(); + let gate = WrapGate::new(TestGate::supervised(dir.path().to_path_buf())) + .action_dir_override(dir.path().join("does-not-exist")) + .fixed_path(dir.path().to_path_buf()) + .arc(); + let r = GlobTool::new(gate) + .execute(json!({"pattern": "*.txt"})) + .await + .unwrap(); + assert!(!r.is_error, "{}", r.output()); + assert!(r.output().contains("a.txt")); +} + +#[tokio::test] +async fn results_report_truncation() { + let dir = TempDir::new().unwrap(); + std::fs::write(dir.path().join("a.txt"), "").unwrap(); + std::fs::write(dir.path().join("b.txt"), "").unwrap(); + let tool = GlobTool::new(TestGate::supervised(dir.path().to_path_buf())); + let ctx = WorkspaceContext::at(dir.path()); + let r = tool + .execute_with_context( + json!({"pattern": "*.txt", "max_results": 1}), + ToolCallOptions::default(), + Some(&ctx), + ) + .await + .unwrap(); + assert!( + r.output().starts_with("1 match(es) (truncated at 1)"), + "{}", + r.output() + ); +} + +#[cfg(unix)] +#[tokio::test] +async fn hits_the_gate_refuses_are_filtered_out() { + let dir = TempDir::new().unwrap(); + let outside = TempDir::new().unwrap(); + std::fs::write(outside.path().join("secret.txt"), "s").unwrap(); + std::fs::write(dir.path().join("ok.txt"), "").unwrap(); + std::os::unix::fs::symlink( + outside.path().join("secret.txt"), + dir.path().join("link.txt"), + ) + .unwrap(); + let tool = GlobTool::new(TestGate::supervised(dir.path().to_path_buf())); + let r = tool.execute(json!({"pattern": "*.txt"})).await.unwrap(); + assert!(r.output().contains("ok.txt")); + assert!(!r.output().contains("link.txt"), "{}", r.output()); +} + +fn seed_many(dir: &std::path::Path, n: usize) { + std::fs::create_dir_all(dir.join("sub")).unwrap(); + for i in 0..n { + std::fs::write(dir.join(format!("sub/f{i}.txt")), "x").unwrap(); + } +} + +#[tokio::test] +async fn glob_asks_the_gate_once_per_search_however_many_hits() { + let dir = TempDir::new().unwrap(); + seed_many(dir.path(), 40); + let wrap = WrapGate::new(TestGate::supervised(dir.path().to_path_buf())); + let calls = wrap.batch_calls(); + let tool = GlobTool::new(wrap.arc()); + let r = tool.execute(json!({"pattern": "**/*.txt"})).await.unwrap(); + assert!(!r.is_error, "{}", r.output()); + assert!(r.output().starts_with("40 match(es)"), "{}", r.output()); + assert_eq!(calls.load(std::sync::atomic::Ordering::SeqCst), 1); +} + +#[tokio::test] +async fn glob_refuses_every_hit_when_the_batch_answer_has_the_wrong_length() { + let dir = TempDir::new().unwrap(); + seed_many(dir.path(), 5); + let wrap = WrapGate::new(TestGate::supervised(dir.path().to_path_buf())).wrong_length_batch(); + let tool = GlobTool::new(wrap.arc()); + let r = tool.execute(json!({"pattern": "**/*.txt"})).await.unwrap(); + assert!(!r.is_error, "{}", r.output()); + assert_eq!(r.output().trim(), "0 match(es)"); +} + +#[tokio::test] +async fn glob_batch_filtering_matches_the_per_hit_verdicts() { + let dir = TempDir::new().unwrap(); + seed_many(dir.path(), 6); + let gate = TestGate::supervised(dir.path().to_path_buf()); + let tool = GlobTool::new(gate.clone()); + let r = tool.execute(json!({"pattern": "**/*.txt"})).await.unwrap(); + let out = r.output(); + let mut got: Vec<&str> = out.lines().skip(1).collect(); + got.sort_unstable(); + let mut want: Vec = (0..6).map(|i| format!("sub/f{i}.txt")).collect(); + want.retain(|p| gate.is_path_string_allowed(p)); + want.sort_unstable(); + assert_eq!(got, want); + assert_eq!(got.len(), 6); +} diff --git a/crates/tinyagents-harness-tools/src/fs/filesystem/glob_search/mod.rs b/crates/tinyagents-harness-tools/src/fs/filesystem/glob_search/mod.rs new file mode 100644 index 000000000..803285b52 --- /dev/null +++ b/crates/tinyagents-harness-tools/src/fs/filesystem/glob_search/mod.rs @@ -0,0 +1,344 @@ +//! `glob` — find files by glob pattern. +//! +//! Coding-harness baseline tool (issue #1205): pure file discovery +//! by pattern (e.g. `src/**/*.rs`). +//! +//! ## Path-root contract (issue #3357) +//! +//! `glob` resolves its search root through the **same** seam the reader tools +//! use — [`FsGate::validate_path`] — and filters every hit through +//! [`FsGate::paths_string_allowed`] (one batch call per search). This guarantees the core +//! invariant a discovery tool must hold: **every path `glob` returns is a path +//! `file_read`/`grep`/`list` can open, and nothing else.** +//! +//! Previously `glob` walked `workspace_dir` (internal product state) and +//! returned paths relative to it, while the readers resolve relative paths +//! against `action_dir`. The two roots normally differ, so a path `glob` +//! returned was meaningless to the readers — every follow-up read failed with +//! "No such file or directory (os error 2)" (#3357). Walking `workspace_dir` +//! also let the agent enumerate internal state (memory DBs, sessions, tokens). +//! +//! The optional `path` argument lets the agent search any location it is +//! allowed to read (the action sandbox by default, or a granted `trusted_root`) +//! — it is resolved through `validate_path`, so it can never widen access +//! beyond what the readers already permit. + +use super::gate::{FsGate, gate_for_context}; +use async_trait::async_trait; +use glob::Pattern; +use serde_json::json; +use std::path::Path; +use std::sync::Arc; +use tinytools::ToolRunContext; +use tinytools::{PermissionLevel, Tool, ToolCallOptions, ToolResult}; +use walkdir::WalkDir; + +const DEFAULT_MAX_RESULTS: usize = 500; + +/// Find files by glob pattern, newest first. +#[derive(Debug)] +pub struct GlobTool { + gate: Arc, +} + +impl GlobTool { + /// Create the tool. + #[must_use] + pub fn new(gate: Arc) -> Self { + Self { gate } + } +} + +#[async_trait] +impl Tool for GlobTool { + fn name(&self) -> &'static str { + "glob" + } + + fn description(&self) -> &'static str { + "Find files matching a glob pattern (e.g. `src/**/*.rs`), searching the action sandbox unless `path` says otherwise. Returns paths newest-first, ready to pass straight to `file_read` / `grep`." + } + + fn parameters_schema(&self) -> serde_json::Value { + json!({ + "type": "object", + "properties": { + "pattern": { + "type": "string", + "description": "Glob pattern, e.g. `**/*.rs` or `src/**/*.{ts,tsx}` (single brace expansion not supported — list patterns separately). Matched against paths relative to the search directory." + }, + "path": { + "type": "string", + "description": "Directory to search within. Defaults to the working directory (action sandbox). Must be a location you're allowed to read; relative paths resolve against the working directory." + }, + "max_results": { + "type": "integer", + "description": "Cap on returned paths (default 500).", + "minimum": 1 + } + }, + "required": ["pattern"] + }) + } + + fn permission_level(&self) -> PermissionLevel { + PermissionLevel::ReadOnly + } + + /// Pure read — safe to fan out across parallel `glob` calls. + fn is_concurrency_safe(&self, _args: &serde_json::Value) -> bool { + true + } + + async fn execute(&self, args: serde_json::Value) -> anyhow::Result { + self.execute_in_context(args, None).await + } + + async fn execute_with_context( + &self, + args: serde_json::Value, + _options: ToolCallOptions, + context: Option<&dyn ToolRunContext>, + ) -> anyhow::Result { + self.execute_in_context(args, context).await + } +} + +impl GlobTool { + async fn execute_in_context( + &self, + args: serde_json::Value, + context: Option<&dyn ToolRunContext>, + ) -> anyhow::Result { + let pattern_str = args + .get("pattern") + .and_then(|v| v.as_str()) + .ok_or_else(|| anyhow::anyhow!("Missing 'pattern' parameter"))?; + // Default the search root to "." — which `validate_path` resolves to the + // action sandbox (action_dir), the same root file_read/grep/list use. + let search_path = args + .get("path") + .and_then(|v| v.as_str()) + .map(str::trim) + .filter(|s| !s.is_empty()) + .unwrap_or("."); + let max_results = args + .get("max_results") + .and_then(serde_json::Value::as_u64) + .map_or(DEFAULT_MAX_RESULTS, |n| (n as usize).max(1)); + + if self.gate.is_rate_limited() { + return Ok(ToolResult::error( + "Rate limit exceeded: too many actions in the last hour", + )); + } + if !self.gate.record_action() { + return Ok(ToolResult::error( + "Rate limit exceeded: action budget exhausted", + )); + } + + let pattern = match Pattern::new(pattern_str) { + Ok(p) => p, + Err(e) => return Ok(ToolResult::error(format!("Invalid glob pattern: {e}"))), + }; + + let path_policy = gate_for_context(&self.gate, context, "glob"); + + log::debug!( + "[tools:glob] search start: pattern='{pattern_str}' path='{search_path}' max_results={max_results}" + ); + + // Resolve + authorize the search root through the readers' validation + // seam. A disallowed or missing root yields a clear, path-naming error + // instead of the opaque "No such file or directory" the old workspace_dir + // root produced (#3357). + let base = match path_policy.validate_path(search_path).await { + Ok(p) => p, + Err(e) => { + log::debug!("[tools:glob] search path rejected: path='{search_path}' error={e}"); + return Ok(ToolResult::error(format!( + "glob search path '{search_path}' is not accessible: {e}" + ))); + } + }; + // `validate_path` can authorize a *file* as well as a directory. Walking a + // file root yields the file itself, whose path relative to `base` is empty, + // so the pattern never matches and the agent gets a misleading "0 match(es)" + // instead of a clear error. Reject non-directory roots explicitly. + match tokio::fs::metadata(&base).await { + Ok(meta) if meta.is_dir() => {} + Ok(_) => { + log::debug!("[tools:glob] search path is not a directory: path='{search_path}'"); + return Ok(ToolResult::error(format!( + "glob search path '{search_path}' is not a directory" + ))); + } + Err(e) => { + log::debug!( + "[tools:glob] search path not accessible: path='{search_path}' error={e}" + ); + return Ok(ToolResult::error(format!( + "glob search path '{search_path}' is not accessible: {e}" + ))); + } + } + log::debug!( + "[tools:glob] resolved search root: '{}' (action_dir='{}')", + base.display(), + path_policy.action_dir().display() + ); + + // Canonical action sandbox, used to decide whether a hit is rendered + // relative (inside the sandbox → directly file_read-able by relative + // path) or absolute (in a granted trusted root → file_read-able as-is). + // `base` is already canonical (validate_path canonicalizes), so if this + // fallback fires the two roots become asymmetric and strip_prefix below + // may miss for an in-sandbox hit, rendering it absolute. Harmless: the + // absolute path is still readable and the per-hit filter still applies. + let action_root = tokio::fs::canonicalize(path_policy.action_dir()) + .await + .unwrap_or_else(|e| { + log::trace!( + "[tools:glob] action_dir canonicalize fallback: path='{}' error={e}", + path_policy.action_dir().display() + ); + path_policy.action_dir().to_path_buf() + }); + + let result = tokio::task::spawn_blocking(move || { + collect_matches( + &base, + &action_root, + path_policy.as_ref(), + &pattern, + max_results, + ) + }) + .await + .map_err(|e| anyhow::anyhow!("scan task failed: {e}"))?; + + let (paths, truncated) = result; + log::debug!( + "[tools:glob] scan complete: {} match(es) truncated={}", + paths.len(), + truncated + ); + let header = if truncated { + format!("{} match(es) (truncated at {max_results})", paths.len()) + } else { + format!("{} match(es)", paths.len()) + }; + + let mut body = String::with_capacity(paths.len() * 32 + header.len() + 1); + body.push_str(&header); + for p in paths { + body.push('\n'); + body.push_str(&p); + } + Ok(ToolResult::success(body)) + } +} + +/// Walk `base`, returning glob matches as paths the reader tools can open. +/// +/// - The pattern is matched against each file's path **relative to `base`** +/// (the searched directory), the conventional glob behavior. +/// - The returned string is **relative to `action_root`** when the hit lives +/// inside the action sandbox (so `file_read("src/a.rs")` works), otherwise +/// the **absolute** path (so `file_read` resolves it as-is via a trusted +/// root). +/// - Every hit is filtered through [`FsGate::paths_string_allowed`] (one batch +/// call per search) on that exact returned string, so `glob` never surfaces a path the readers +/// would reject (internal state, forbidden trees, symlink escapes). +/// +/// `max_results` bounds the returned list, not the walk: the full tree under +/// `base` is always traversed and only the output is truncated (pre-existing +/// behavior — the cap caps results, not work). +fn collect_matches( + base: &Path, + action_root: &Path, + gate: &dyn FsGate, + pattern: &Pattern, + max_results: usize, +) -> (Vec, bool) { + let mut candidates: Vec<(std::time::SystemTime, String)> = Vec::new(); + + for entry in WalkDir::new(base) + .follow_links(false) + .into_iter() + .filter_entry(|e| !is_skipped(e.file_name().to_string_lossy().as_ref())) + .filter_map(std::result::Result::ok) + { + if !entry.file_type().is_file() { + continue; + } + let abs = entry.path(); + + // Match the pattern against the path relative to the searched root. + let rel_to_base = match abs.strip_prefix(base) { + Ok(p) => p.to_string_lossy().replace('\\', "/"), + Err(_) => continue, + }; + if !pattern.matches(&rel_to_base) { + continue; + } + + // Render the path the way the agent should hand it to file_read: + // sandbox-relative when inside action_dir, absolute otherwise. + let display = match abs.strip_prefix(action_root) { + Ok(rel) => rel.to_string_lossy().replace('\\', "/"), + Err(_) => abs.to_string_lossy().to_string(), + }; + + let mtime = entry + .metadata() + .ok() + .and_then(|m| m.modified().ok()) + .unwrap_or(std::time::SystemTime::UNIX_EPOCH); + candidates.push((mtime, display)); + } + + // Fail-closed: only surface paths the readers would also accept. The whole + // walk is asked in ONE batch so a host that answers over a bus pays a + // single round trip per search, not one per hit. + let displays: Vec = candidates.iter().map(|(_, p)| p.clone()).collect(); + let verdicts = gate.paths_string_allowed(&displays); + if verdicts.len() != candidates.len() { + log::warn!( + "[tools:glob] gate answered {} verdict(s) for {} path(s); refusing every hit", + verdicts.len(), + candidates.len() + ); + candidates.clear(); + } + let mut hits: Vec<(std::time::SystemTime, String)> = Vec::with_capacity(candidates.len()); + for ((mtime, display), allowed) in candidates.into_iter().zip(verdicts) { + if allowed { + hits.push((mtime, display)); + } else { + log::trace!("[tools:glob] path filtered by policy: '{display}'"); + } + } + + // Newest first. + hits.sort_by_key(|item| std::cmp::Reverse(item.0)); + let truncated = hits.len() > max_results; + let paths: Vec = hits.into_iter().take(max_results).map(|(_, p)| p).collect(); + (paths, truncated) +} + +fn is_skipped(name: &str) -> bool { + matches!( + name, + ".git" | "node_modules" | "target" | ".next" | "dist" | "build" | ".cache" + ) +} + +#[cfg(test)] +#[path = "mod_tests.rs"] +mod tests; + +#[cfg(test)] +#[path = "extra_tests.rs"] +mod extra_tests; diff --git a/crates/tinyagents-harness-tools/src/fs/filesystem/glob_search/mod_tests.rs b/crates/tinyagents-harness-tools/src/fs/filesystem/glob_search/mod_tests.rs new file mode 100644 index 000000000..d8a712a9a --- /dev/null +++ b/crates/tinyagents-harness-tools/src/fs/filesystem/glob_search/mod_tests.rs @@ -0,0 +1,264 @@ +//! Behavior tests for the `glob` tool, driven through a fake [`FsGate`]. + +#![allow(clippy::expect_used, clippy::panic, clippy::unwrap_used)] + +use super::*; +use crate::fs::filesystem::test_support::TestGate; + +fn test_security(workspace: std::path::PathBuf) -> Arc { + TestGate::supervised(workspace) +} + +/// Gate with an action sandbox and extra granted roots: the production shape, +/// and the configuration that surfaced #3357. +fn test_security_split( + action_dir: std::path::PathBuf, + _workspace_dir: std::path::PathBuf, + extra_roots: Vec, +) -> Arc { + TestGate::split(action_dir, extra_roots) +} + +#[test] +fn glob_name() { + let tool = GlobTool::new(test_security(std::env::temp_dir())); + assert_eq!(tool.name(), "glob"); +} + +#[tokio::test] +async fn glob_matches_extension() { + let dir = std::env::temp_dir().join("openhuman_test_glob_ext"); + let _ = tokio::fs::remove_dir_all(&dir).await; + tokio::fs::create_dir_all(dir.join("src/sub")) + .await + .unwrap(); + tokio::fs::write(dir.join("src/a.rs"), "// a") + .await + .unwrap(); + tokio::fs::write(dir.join("src/sub/b.rs"), "// b") + .await + .unwrap(); + tokio::fs::write(dir.join("src/c.txt"), "c").await.unwrap(); + + let tool = GlobTool::new(test_security(dir.clone())); + let result = tool.execute(json!({"pattern": "**/*.rs"})).await.unwrap(); + assert!(!result.is_error); + let output = result.output(); + assert!(output.contains("src/a.rs")); + assert!(output.contains("src/sub/b.rs")); + assert!(!output.contains("c.txt")); + + let _ = tokio::fs::remove_dir_all(&dir).await; +} + +#[tokio::test] +async fn glob_invalid_pattern() { + let dir = std::env::temp_dir().join("openhuman_test_glob_invalid"); + let _ = tokio::fs::remove_dir_all(&dir).await; + tokio::fs::create_dir_all(&dir).await.unwrap(); + let tool = GlobTool::new(test_security(dir.clone())); + let result = tool.execute(json!({"pattern": "**["})).await.unwrap(); + assert!(result.is_error); + assert!(result.output().contains("Invalid glob")); + let _ = tokio::fs::remove_dir_all(&dir).await; +} + +#[tokio::test] +async fn glob_skips_node_modules() { + let dir = std::env::temp_dir().join("openhuman_test_glob_skip"); + let _ = tokio::fs::remove_dir_all(&dir).await; + tokio::fs::create_dir_all(dir.join("node_modules")) + .await + .unwrap(); + tokio::fs::write(dir.join("node_modules/lib.js"), "") + .await + .unwrap(); + tokio::fs::write(dir.join("app.js"), "").await.unwrap(); + + let tool = GlobTool::new(test_security(dir.clone())); + let result = tool.execute(json!({"pattern": "**/*.js"})).await.unwrap(); + let output = result.output(); + assert!(output.contains("app.js")); + assert!(!output.contains("node_modules")); + + let _ = tokio::fs::remove_dir_all(&dir).await; +} + +/// Regression for #3357: glob roots at `action_dir` (so its hits are readable), +/// and never surfaces files living under the internal `workspace_dir`. +#[tokio::test] +async fn glob_roots_at_action_dir_and_excludes_workspace() { + let root = std::env::temp_dir().join("openhuman_test_glob_split"); + let action = root.join("action"); + let workspace = root.join("workspace"); + let _ = tokio::fs::remove_dir_all(&root).await; + tokio::fs::create_dir_all(&action).await.unwrap(); + tokio::fs::create_dir_all(&workspace).await.unwrap(); + tokio::fs::write(action.join("keep.txt"), "keep") + .await + .unwrap(); + tokio::fs::write(workspace.join("secret.txt"), "secret") + .await + .unwrap(); + + let security = test_security_split(action.clone(), workspace.clone(), vec![]); + let tool = GlobTool::new(security.clone()); + let result = tool.execute(json!({"pattern": "**/*.txt"})).await.unwrap(); + assert!(!result.is_error, "{}", result.output()); + let output = result.output(); + // Action-sandbox file is found, rendered relative... + assert!(output.contains("keep.txt"), "missing keep.txt: {output}"); + // ...and the internal workspace file is NOT enumerated. + assert!( + !output.contains("secret.txt"), + "leaked workspace file: {output}" + ); + + // The glob hit must be directly readable by the reader tools. + assert!( + security.validate_path("keep.txt").await.is_ok(), + "glob hit not resolvable by validate_path" + ); + + let _ = tokio::fs::remove_dir_all(&root).await; +} + +/// "As wide as the readers": glob can search any granted trusted root via +/// the `path` arg, returning absolute paths that the readers accept as-is. +#[tokio::test] +async fn glob_searches_named_trusted_root() { + let root = std::env::temp_dir().join("openhuman_test_glob_trusted"); + let action = root.join("action"); + let workspace = root.join("workspace"); + let granted = root.join("granted"); + let _ = tokio::fs::remove_dir_all(&root).await; + tokio::fs::create_dir_all(&action).await.unwrap(); + tokio::fs::create_dir_all(&workspace).await.unwrap(); + tokio::fs::create_dir_all(&granted).await.unwrap(); + tokio::fs::write(granted.join("data.csv"), "x") + .await + .unwrap(); + + let security = test_security_split(action.clone(), workspace.clone(), vec![granted.clone()]); + let tool = GlobTool::new(security.clone()); + let granted_abs = granted.to_string_lossy().to_string(); + let result = tool + .execute(json!({"pattern": "**/*.csv", "path": granted_abs})) + .await + .unwrap(); + assert!(!result.is_error, "{}", result.output()); + let output = result.output(); + assert!(output.contains("data.csv"), "missing data.csv: {output}"); + // Rendered absolute (outside the action sandbox). + assert!( + output.contains(&granted.to_string_lossy().to_string()), + "expected absolute path: {output}" + ); + + let _ = tokio::fs::remove_dir_all(&root).await; +} + +/// A search root the policy disallows yields a clear error, not ENOENT. +#[tokio::test] +async fn glob_rejects_disallowed_search_path() { + let root = std::env::temp_dir().join("openhuman_test_glob_reject"); + let action = root.join("action"); + let workspace = root.join("workspace"); + let _ = tokio::fs::remove_dir_all(&root).await; + tokio::fs::create_dir_all(&action).await.unwrap(); + tokio::fs::create_dir_all(&workspace).await.unwrap(); + + let security = test_security_split(action.clone(), workspace.clone(), vec![]); + let tool = GlobTool::new(security); + // Absolute path outside any trusted root, under workspace_only. + let result = tool + .execute(json!({"pattern": "**/*", "path": "/etc"})) + .await + .unwrap(); + assert!(result.is_error); + assert!( + result.output().contains("not accessible"), + "{}", + result.output() + ); + + let _ = tokio::fs::remove_dir_all(&root).await; +} + +/// Security backstop: a symlink inside the sandbox pointing OUTSIDE it must +/// never leak the target. Two layers cover this and this test exercises +/// both: the walk runs with `follow_links(false)`, so the symlink entry is +/// dropped by the `is_file()` gate (a symlink is not a regular file) and is +/// never descended; and `is_path_string_allowed` is the fail-closed per-hit +/// backstop — assert directly that it rejects the escape's resolved string, +/// since that is the check `collect_matches` leans on if the walk gate is +/// ever bypassed. A legitimate in-sandbox file is still found. +#[cfg(unix)] +#[tokio::test] +async fn glob_does_not_leak_symlink_escape() { + let root = std::env::temp_dir().join("openhuman_test_glob_symlink"); + let action = root.join("action"); + let workspace = root.join("workspace"); + let outside = root.join("outside"); + let _ = tokio::fs::remove_dir_all(&root).await; + tokio::fs::create_dir_all(action.join("sub")).await.unwrap(); + tokio::fs::create_dir_all(&workspace).await.unwrap(); + tokio::fs::create_dir_all(&outside).await.unwrap(); + tokio::fs::write(action.join("sub/ok.txt"), "ok") + .await + .unwrap(); + tokio::fs::write(outside.join("secret.txt"), "secret") + .await + .unwrap(); + // Symlink inside the sandbox pointing at the outside (untrusted) tree. + std::os::unix::fs::symlink(&outside, action.join("escape")).unwrap(); + + // `outside` is deliberately NOT registered as a trusted root. + let security = test_security_split(action.clone(), workspace.clone(), vec![]); + let tool = GlobTool::new(security.clone()); + let result = tool.execute(json!({"pattern": "**/*.txt"})).await.unwrap(); + assert!(!result.is_error, "{}", result.output()); + let output = result.output(); + // The in-sandbox file is found, the escape target is not enumerated. + assert!( + output.contains("ok.txt"), + "missing in-sandbox file: {output}" + ); + assert!( + !output.contains("secret"), + "leaked symlink-escape target: {output}" + ); + // The fail-closed backstop, exercised directly: the resolved escape + // path is rejected as a string regardless of how it was reached. + let escaped = outside.join("secret.txt").to_string_lossy().to_string(); + assert!( + !security.is_path_string_allowed(&escaped), + "policy must reject the escape target string: {escaped}" + ); + + let _ = tokio::fs::remove_dir_all(&root).await; +} + +/// A `path` that resolves to a *file* (not a directory) yields a clear +/// "not a directory" error rather than a misleading "0 match(es)". +#[tokio::test] +async fn glob_rejects_file_search_path() { + let dir = std::env::temp_dir().join("openhuman_test_glob_file_root"); + let _ = tokio::fs::remove_dir_all(&dir).await; + tokio::fs::create_dir_all(&dir).await.unwrap(); + tokio::fs::write(dir.join("file.txt"), "x").await.unwrap(); + + let tool = GlobTool::new(test_security(dir.clone())); + let result = tool + .execute(json!({"pattern": "**/*", "path": "file.txt"})) + .await + .unwrap(); + assert!(result.is_error, "{}", result.output()); + assert!( + result.output().contains("not a directory"), + "{}", + result.output() + ); + + let _ = tokio::fs::remove_dir_all(&dir).await; +} diff --git a/crates/tinyagents-harness-tools/src/fs/filesystem/grep/extra_tests.rs b/crates/tinyagents-harness-tools/src/fs/filesystem/grep/extra_tests.rs new file mode 100644 index 000000000..a920df2b4 --- /dev/null +++ b/crates/tinyagents-harness-tools/src/fs/filesystem/grep/extra_tests.rs @@ -0,0 +1,54 @@ +//! Budget, delegation and unreadable-file tests for `grep`. + +#![allow(clippy::expect_used, clippy::panic, clippy::unwrap_used)] + +use super::*; +use crate::fs::filesystem::test_support::{TestGate, WorkspaceContext, WrapGate}; +use tempfile::TempDir; + +#[test] +fn grep_is_concurrency_safe() { + let dir = TempDir::new().unwrap(); + let tool = GrepTool::new(TestGate::supervised(dir.path().to_path_buf())); + assert!(tool.is_concurrency_safe(&json!({}))); +} + +#[tokio::test] +async fn rate_limit_and_exhausted_budget_are_refused() { + let dir = TempDir::new().unwrap(); + let limited = GrepTool::new(TestGate::with( + dir.path().to_path_buf(), + crate::fs::filesystem::test_support::AutonomyLevel::Supervised, + 0, + )); + let r = limited.execute(json!({"pattern": "x"})).await.unwrap(); + assert!(r.output().contains("too many actions in the last hour")); + + let gate = WrapGate::new(TestGate::supervised(dir.path().to_path_buf())) + .no_budget() + .arc(); + let r = GrepTool::new(gate) + .execute(json!({"pattern": "x"})) + .await + .unwrap(); + assert!(r.output().contains("action budget exhausted")); +} + +#[tokio::test] +async fn non_utf8_files_are_skipped_and_others_still_match() { + let dir = TempDir::new().unwrap(); + std::fs::write(dir.path().join("bin.dat"), [0xff, 0xfe, 0x00, 0xc3]).unwrap(); + std::fs::write(dir.path().join("t.txt"), "needle here\n").unwrap(); + let tool = GrepTool::new(TestGate::supervised(dir.path().to_path_buf())); + let ctx = WorkspaceContext::at(dir.path()); + let r = tool + .execute_with_context( + json!({"pattern": "needle"}), + ToolCallOptions::default(), + Some(&ctx), + ) + .await + .unwrap(); + assert!(r.output().contains("t.txt"), "{}", r.output()); + assert!(!r.output().contains("bin.dat")); +} diff --git a/crates/tinyagents-harness-tools/src/fs/filesystem/grep/mod.rs b/crates/tinyagents-harness-tools/src/fs/filesystem/grep/mod.rs new file mode 100644 index 000000000..d84bf7768 --- /dev/null +++ b/crates/tinyagents-harness-tools/src/fs/filesystem/grep/mod.rs @@ -0,0 +1,241 @@ +//! `grep` — regex search across files in the workspace. +//! +//! Coding-harness baseline tool (issue #1205): a first-class +//! file-navigation primitive that lets the agent search for a regex +//! across the workspace without falling through to `shell`. Uses the +//! same path-sandboxing + rate-limiting as `file_read`. + +use super::gate::{FsGate, gate_for_context}; +use async_trait::async_trait; +use regex::Regex; +use serde_json::json; +use std::path::Path; +use std::sync::Arc; +use tinytools::ToolRunContext; +use tinytools::{PermissionLevel, Tool, ToolCallOptions, ToolResult}; +use walkdir::WalkDir; + +const DEFAULT_MAX_MATCHES: usize = 200; +const MAX_LINE_BYTES: usize = 2_000; +const MAX_FILE_BYTES: u64 = 5 * 1024 * 1024; + +/// Search file contents with a regular expression. +#[derive(Debug)] +pub struct GrepTool { + gate: Arc, +} + +impl GrepTool { + /// Create the tool. + #[must_use] + pub fn new(gate: Arc) -> Self { + Self { gate } + } +} + +#[async_trait] +impl Tool for GrepTool { + fn name(&self) -> &'static str { + "grep" + } + + fn description(&self) -> &'static str { + "Search file contents in the workspace with a regular expression. \ + Returns up to `max_matches` matches (default 200) as `path:line:text` lines." + } + + fn parameters_schema(&self) -> serde_json::Value { + json!({ + "type": "object", + "properties": { + "pattern": { + "type": "string", + "description": "Regular expression (Rust `regex` crate syntax)." + }, + "path": { + "type": "string", + "description": "Optional sub-path to restrict the search to (relative to workspace).", + "default": "." + }, + "max_matches": { + "type": "integer", + "description": "Cap on matches returned (default 200).", + "minimum": 1 + }, + "case_insensitive": { + "type": "boolean", + "description": "If true, compile the regex with the `i` flag.", + "default": false + } + }, + "required": ["pattern"] + }) + } + + fn permission_level(&self) -> PermissionLevel { + PermissionLevel::ReadOnly + } + + /// Pure read — safe to fan out across parallel `grep` calls. + fn is_concurrency_safe(&self, _args: &serde_json::Value) -> bool { + true + } + + async fn execute(&self, args: serde_json::Value) -> anyhow::Result { + self.execute_in_context(args, None).await + } + + async fn execute_with_context( + &self, + args: serde_json::Value, + _options: ToolCallOptions, + context: Option<&dyn ToolRunContext>, + ) -> anyhow::Result { + self.execute_in_context(args, context).await + } +} + +impl GrepTool { + async fn execute_in_context( + &self, + args: serde_json::Value, + context: Option<&dyn ToolRunContext>, + ) -> anyhow::Result { + let pattern = args + .get("pattern") + .and_then(|v| v.as_str()) + .ok_or_else(|| anyhow::anyhow!("Missing 'pattern' parameter"))?; + let sub_path = args.get("path").and_then(|v| v.as_str()).unwrap_or("."); + let max_matches = args + .get("max_matches") + .and_then(serde_json::Value::as_u64) + .map_or(DEFAULT_MAX_MATCHES, |n| (n as usize).max(1)); + let case_insensitive = args + .get("case_insensitive") + .and_then(serde_json::Value::as_bool) + .unwrap_or(false); + + if self.gate.is_rate_limited() { + return Ok(ToolResult::error( + "Rate limit exceeded: too many actions in the last hour", + )); + } + if !self.gate.record_action() { + return Ok(ToolResult::error( + "Rate limit exceeded: action budget exhausted", + )); + } + + let regex = match build_regex(pattern, case_insensitive) { + Ok(r) => r, + Err(e) => return Ok(ToolResult::error(format!("Invalid regex: {e}"))), + }; + + let path_policy = gate_for_context(&self.gate, context, "grep"); + + // Security check: validate path string, resolve symlinks, confirm workspace containment. + let resolved_root = match path_policy.validate_path(sub_path).await { + Ok(p) => p, + Err(msg) => return Ok(ToolResult::error(msg)), + }; + + let workspace = path_policy.action_dir().to_path_buf(); + let result = tokio::task::spawn_blocking(move || { + scan_for_matches(&resolved_root, &workspace, ®ex, max_matches) + }) + .await + .map_err(|e| anyhow::anyhow!("scan task failed: {e}"))?; + + let (matches, scanned, truncated) = result; + let header = if truncated { + format!( + "{} match(es) (truncated at {max_matches}); scanned {scanned} file(s)", + matches.len() + ) + } else { + format!("{} match(es); scanned {scanned} file(s)", matches.len()) + }; + + let mut body = String::with_capacity(matches.len() * 80 + header.len() + 1); + body.push_str(&header); + for m in &matches { + body.push('\n'); + body.push_str(m); + } + Ok(ToolResult::success(body)) + } +} + +fn build_regex(pattern: &str, case_insensitive: bool) -> Result { + if case_insensitive { + regex::RegexBuilder::new(pattern) + .case_insensitive(true) + .build() + } else { + Regex::new(pattern) + } +} + +fn scan_for_matches( + root: &Path, + workspace: &Path, + regex: &Regex, + max_matches: usize, +) -> (Vec, usize, bool) { + let mut matches: Vec = Vec::new(); + let mut scanned = 0usize; + let mut truncated = false; + + 'outer: for entry in WalkDir::new(root) + .follow_links(false) + .into_iter() + .filter_entry(|e| !is_skipped(e.file_name().to_string_lossy().as_ref())) + .filter_map(std::result::Result::ok) + { + if !entry.file_type().is_file() { + continue; + } + let path = entry.path(); + let Ok(meta) = entry.metadata() else { continue }; + if meta.len() > MAX_FILE_BYTES { + continue; + } + let Ok(contents) = std::fs::read_to_string(path) else { + continue; + }; + scanned += 1; + let rel = path.strip_prefix(workspace).unwrap_or(path); + for (lineno, line) in contents.lines().enumerate() { + if regex.is_match(line) { + // `MAX_LINE_BYTES` is a BYTE budget — use the byte-aware + // helper. The earlier migration to `truncate_with_suffix` + // mis-typed this as a char budget; for multi-byte text + // (CJK / emoji) the rendered line could balloon to ~3× + // the intended cap. Per CodeRabbit critical review on + // PR #1549. + let display_line = super::text::truncate_at_byte_boundary(line, MAX_LINE_BYTES); + matches.push(format!("{}:{}:{}", rel.display(), lineno + 1, display_line)); + if matches.len() >= max_matches { + truncated = true; + break 'outer; + } + } + } + } + (matches, scanned, truncated) +} + +fn is_skipped(name: &str) -> bool { + matches!( + name, + ".git" | "node_modules" | "target" | ".next" | "dist" | "build" | ".cache" + ) +} + +#[cfg(test)] +#[path = "mod_tests.rs"] +mod tests; + +#[cfg(test)] +#[path = "extra_tests.rs"] +mod extra_tests; diff --git a/crates/tinyagents-harness-tools/src/fs/filesystem/grep/mod_tests.rs b/crates/tinyagents-harness-tools/src/fs/filesystem/grep/mod_tests.rs new file mode 100644 index 000000000..bec0087c3 --- /dev/null +++ b/crates/tinyagents-harness-tools/src/fs/filesystem/grep/mod_tests.rs @@ -0,0 +1,184 @@ +//! Behavior tests for the `grep` tool, driven through a fake [`FsGate`]. + +#![allow(clippy::expect_used, clippy::panic, clippy::unwrap_used)] + +use super::*; +use crate::fs::filesystem::test_support::TestGate; + +fn test_security(workspace: std::path::PathBuf) -> Arc { + TestGate::supervised(workspace) +} + +#[test] +fn grep_name_and_schema() { + let tool = GrepTool::new(test_security(std::env::temp_dir())); + assert_eq!(tool.name(), "grep"); + let schema = tool.parameters_schema(); + assert!(schema["properties"]["pattern"].is_object()); + assert!( + schema["required"] + .as_array() + .unwrap() + .contains(&json!("pattern")) + ); +} + +#[tokio::test] +async fn grep_finds_matches() { + let dir = std::env::temp_dir().join("openhuman_test_grep_finds"); + let _ = tokio::fs::remove_dir_all(&dir).await; + tokio::fs::create_dir_all(&dir).await.unwrap(); + tokio::fs::write(dir.join("a.txt"), "alpha\nbravo\ncharlie") + .await + .unwrap(); + tokio::fs::write(dir.join("b.txt"), "alpha2").await.unwrap(); + + let tool = GrepTool::new(test_security(dir.clone())); + let result = tool.execute(json!({"pattern": "^alpha"})).await.unwrap(); + assert!(!result.is_error); + let output = result.output(); + assert!(output.contains("a.txt:1:alpha")); + assert!(output.contains("b.txt:1:alpha2")); + assert!(!output.contains("bravo")); + + let _ = tokio::fs::remove_dir_all(&dir).await; +} + +#[tokio::test] +async fn grep_invalid_regex() { + let dir = std::env::temp_dir().join("openhuman_test_grep_invalid"); + let _ = tokio::fs::remove_dir_all(&dir).await; + tokio::fs::create_dir_all(&dir).await.unwrap(); + + let tool = GrepTool::new(test_security(dir.clone())); + let result = tool + .execute(json!({"pattern": "([unclosed"})) + .await + .unwrap(); + assert!(result.is_error); + assert!(result.output().contains("Invalid regex")); + + let _ = tokio::fs::remove_dir_all(&dir).await; +} + +#[tokio::test] +async fn grep_case_insensitive() { + let dir = std::env::temp_dir().join("openhuman_test_grep_ci"); + let _ = tokio::fs::remove_dir_all(&dir).await; + tokio::fs::create_dir_all(&dir).await.unwrap(); + tokio::fs::write(dir.join("c.txt"), "Hello World") + .await + .unwrap(); + + let tool = GrepTool::new(test_security(dir.clone())); + let result = tool + .execute(json!({"pattern": "hello", "case_insensitive": true})) + .await + .unwrap(); + assert!(!result.is_error); + assert!(result.output().contains("c.txt:1:Hello World")); + + let _ = tokio::fs::remove_dir_all(&dir).await; +} + +#[tokio::test] +async fn grep_blocks_path_traversal() { + let tool = GrepTool::new(test_security(std::env::temp_dir())); + let result = tool + .execute(json!({"pattern": ".", "path": "../.."})) + .await + .unwrap(); + assert!(result.is_error); + assert!(result.output().contains("not allowed")); +} + +#[tokio::test] +async fn grep_skips_node_modules_and_git() { + let dir = std::env::temp_dir().join("openhuman_test_grep_skip"); + let _ = tokio::fs::remove_dir_all(&dir).await; + tokio::fs::create_dir_all(dir.join("node_modules")) + .await + .unwrap(); + tokio::fs::create_dir_all(dir.join(".git")).await.unwrap(); + tokio::fs::write(dir.join("node_modules/x.txt"), "needle") + .await + .unwrap(); + tokio::fs::write(dir.join(".git/x.txt"), "needle") + .await + .unwrap(); + tokio::fs::write(dir.join("real.txt"), "needle") + .await + .unwrap(); + + let tool = GrepTool::new(test_security(dir.clone())); + let result = tool.execute(json!({"pattern": "needle"})).await.unwrap(); + assert!(!result.is_error); + let output = result.output(); + assert!(output.contains("real.txt")); + assert!(!output.contains("node_modules")); + assert!(!output.contains(".git")); + + let _ = tokio::fs::remove_dir_all(&dir).await; +} + +#[tokio::test] +async fn grep_respects_max_matches() { + let dir = std::env::temp_dir().join("openhuman_test_grep_max"); + let _ = tokio::fs::remove_dir_all(&dir).await; + tokio::fs::create_dir_all(&dir).await.unwrap(); + let mut text = String::new(); + for _ in 0..50 { + text.push_str("hit\n"); + } + tokio::fs::write(dir.join("many.txt"), text).await.unwrap(); + + let tool = GrepTool::new(test_security(dir.clone())); + let result = tool + .execute(json!({"pattern": "hit", "max_matches": 5})) + .await + .unwrap(); + assert!(!result.is_error); + assert!(result.output().contains("truncated")); + + let _ = tokio::fs::remove_dir_all(&dir).await; +} + +/// A file over `MAX_FILE_BYTES` is silently excluded rather than read — this +/// pins that the exclusion is at least honest about it: the file does not +/// count toward `scanned`, so the report never claims to have searched bytes +/// it skipped. It does not surface the file was too large; a caller who +/// expected a match in it sees only a lower `scanned` count with no reason. +#[tokio::test] +async fn grep_excludes_an_oversized_file_and_does_not_count_it_scanned() { + let dir = std::env::temp_dir().join("openhuman_test_grep_oversized"); + let _ = tokio::fs::remove_dir_all(&dir).await; + tokio::fs::create_dir_all(&dir).await.unwrap(); + + let small = format!("{}\nneedle\n", "line ".repeat(10)); + tokio::fs::write(dir.join("small.txt"), &small) + .await + .unwrap(); + let huge = "needle\n".repeat(1_000_000); + tokio::fs::write(dir.join("huge.txt"), &huge).await.unwrap(); + + let tool = GrepTool::new(test_security(dir.clone())); + let result = tool.execute(json!({"pattern": "needle"})).await.unwrap(); + assert!(!result.is_error, "{}", result.output()); + assert!( + result.output().contains("small.txt"), + "the small file's match must still be reported: {}", + result.output() + ); + assert!( + !result.output().contains("huge.txt"), + "the oversized file must not be scanned at all: {}", + result.output() + ); + assert!( + result.output().contains("scanned 1 file"), + "the scanned count must not include the file it skipped for size: {}", + result.output() + ); + + let _ = tokio::fs::remove_dir_all(&dir).await; +} diff --git a/crates/tinyagents-harness-tools/src/fs/filesystem/image_info/mod.rs b/crates/tinyagents-harness-tools/src/fs/filesystem/image_info/mod.rs new file mode 100644 index 000000000..da647744f --- /dev/null +++ b/crates/tinyagents-harness-tools/src/fs/filesystem/image_info/mod.rs @@ -0,0 +1,211 @@ +use super::gate::FsGate; +use async_trait::async_trait; +use serde_json::json; +use std::fmt::Write; +use std::sync::Arc; +use tinytools::{Tool, ToolResult}; + +/// Maximum file size we will read and base64-encode (5 MB). +const MAX_IMAGE_BYTES: u64 = 5_242_880; + +/// Tool to read image metadata and optionally return base64-encoded data. +/// +/// Since providers are currently text-only, this tool extracts what it can +/// (file size, format, dimensions from header bytes) and provides base64 +/// data for future multimodal provider support. +#[derive(Debug)] +pub struct ImageInfoTool { + gate: Arc, +} + +impl ImageInfoTool { + /// An `image_info` tool that reads only what `gate` allows. + pub fn new(gate: Arc) -> Self { + Self { gate } + } + + /// Detect image format from first few bytes (magic numbers). + fn detect_format(bytes: &[u8]) -> &'static str { + if bytes.len() < 4 { + return "unknown"; + } + if bytes.starts_with(b"\x89PNG") { + "png" + } else if bytes.starts_with(b"\xFF\xD8\xFF") { + "jpeg" + } else if bytes.starts_with(b"GIF8") { + "gif" + } else if bytes.starts_with(b"RIFF") && bytes.len() >= 12 && &bytes[8..12] == b"WEBP" { + "webp" + } else if bytes.starts_with(b"BM") { + "bmp" + } else { + "unknown" + } + } + + /// Try to extract dimensions from image header bytes. + /// Returns (width, height) if detectable. + fn extract_dimensions(bytes: &[u8], format: &str) -> Option<(u32, u32)> { + match format { + "png" => { + // PNG IHDR chunk: bytes 16-19 = width, 20-23 = height (big-endian) + if bytes.len() >= 24 { + let w = u32::from_be_bytes([bytes[16], bytes[17], bytes[18], bytes[19]]); + let h = u32::from_be_bytes([bytes[20], bytes[21], bytes[22], bytes[23]]); + Some((w, h)) + } else { + None + } + } + "gif" => { + // GIF: bytes 6-7 = width, 8-9 = height (little-endian) + if bytes.len() >= 10 { + let w = u32::from(u16::from_le_bytes([bytes[6], bytes[7]])); + let h = u32::from(u16::from_le_bytes([bytes[8], bytes[9]])); + Some((w, h)) + } else { + None + } + } + "bmp" => { + // BMP: bytes 18-21 = width, 22-25 = height (little-endian, signed) + if bytes.len() >= 26 { + let w = u32::from_le_bytes([bytes[18], bytes[19], bytes[20], bytes[21]]); + let h_raw = i32::from_le_bytes([bytes[22], bytes[23], bytes[24], bytes[25]]); + let h = h_raw.unsigned_abs(); + Some((w, h)) + } else { + None + } + } + "jpeg" => Self::jpeg_dimensions(bytes), + _ => None, + } + } + + /// Parse JPEG SOF markers to extract dimensions. + fn jpeg_dimensions(bytes: &[u8]) -> Option<(u32, u32)> { + let mut i = 2; // skip SOI marker + while i + 1 < bytes.len() { + if bytes[i] != 0xFF { + return None; + } + let marker = bytes[i + 1]; + i += 2; + + // SOF0..SOF3 markers contain dimensions + if (0xC0..=0xC3).contains(&marker) { + if i + 7 <= bytes.len() { + let h = u32::from(u16::from_be_bytes([bytes[i + 3], bytes[i + 4]])); + let w = u32::from(u16::from_be_bytes([bytes[i + 5], bytes[i + 6]])); + return Some((w, h)); + } + return None; + } + + // Skip this segment + if i + 1 < bytes.len() { + let seg_len = u16::from_be_bytes([bytes[i], bytes[i + 1]]) as usize; + if seg_len < 2 { + return None; // Malformed segment (valid segments have length >= 2) + } + i += seg_len; + } else { + return None; + } + } + None + } +} + +#[async_trait] +impl Tool for ImageInfoTool { + fn name(&self) -> &'static str { + "image_info" + } + + fn description(&self) -> &'static str { + "Read image file metadata (format, dimensions, size) and optionally return base64-encoded data." + } + + fn parameters_schema(&self) -> serde_json::Value { + json!({ + "type": "object", + "properties": { + "path": { + "type": "string", + "description": "Path to the image file (absolute or relative to workspace)" + }, + "include_base64": { + "type": "boolean", + "description": "Include base64-encoded image data in output (default: false)" + } + }, + "required": ["path"] + }) + } + + async fn execute(&self, args: serde_json::Value) -> anyhow::Result { + let path_str = args + .get("path") + .and_then(|v| v.as_str()) + .ok_or_else(|| anyhow::anyhow!("Missing 'path' parameter"))?; + + let include_base64 = args + .get("include_base64") + .and_then(serde_json::Value::as_bool) + .unwrap_or(false); + + // Security check: validate path string, resolve symlinks, confirm workspace containment. + let resolved = match self.gate.validate_path(path_str).await { + Ok(p) => p, + Err(msg) => return Ok(ToolResult::error(msg)), + }; + + let metadata = tokio::fs::metadata(&resolved) + .await + .map_err(|e| anyhow::anyhow!("Failed to read file metadata: {e}"))?; + + let file_size = metadata.len(); + + if file_size > MAX_IMAGE_BYTES { + return Ok(ToolResult::error(format!( + "Image too large: {file_size} bytes (max {MAX_IMAGE_BYTES} bytes)" + ))); + } + + let bytes = tokio::fs::read(&resolved) + .await + .map_err(|e| anyhow::anyhow!("Failed to read image file: {e}"))?; + + let format = Self::detect_format(&bytes); + let dimensions = Self::extract_dimensions(&bytes, format); + + let mut output = format!("File: {path_str}\nFormat: {format}\nSize: {file_size} bytes"); + + if let Some((w, h)) = dimensions { + let _ = write!(output, "\nDimensions: {w}x{h}"); + } + + if include_base64 { + use base64::Engine; + let encoded = base64::engine::general_purpose::STANDARD.encode(&bytes); + let mime = match format { + "png" => "image/png", + "jpeg" => "image/jpeg", + "gif" => "image/gif", + "webp" => "image/webp", + "bmp" => "image/bmp", + _ => "application/octet-stream", + }; + let _ = write!(output, "\ndata:{mime};base64,{encoded}"); + } + + Ok(ToolResult::success(output)) + } +} + +#[cfg(test)] +#[path = "mod_tests.rs"] +mod tests; diff --git a/crates/tinyagents-harness-tools/src/fs/filesystem/image_info/mod_tests.rs b/crates/tinyagents-harness-tools/src/fs/filesystem/image_info/mod_tests.rs new file mode 100644 index 000000000..40245aa3d --- /dev/null +++ b/crates/tinyagents-harness-tools/src/fs/filesystem/image_info/mod_tests.rs @@ -0,0 +1,342 @@ +#![allow(clippy::expect_used, clippy::unwrap_used, clippy::panic)] + +use super::*; +use crate::fs::filesystem::test_support::{AutonomyLevel, TestGate}; + +fn test_security() -> Arc { + TestGate::with(std::env::temp_dir(), AutonomyLevel::Full, 1_000_000) +} + +#[test] +fn image_info_tool_name() { + let tool = ImageInfoTool::new(test_security()); + assert_eq!(tool.name(), "image_info"); +} + +#[test] +fn image_info_tool_description() { + let tool = ImageInfoTool::new(test_security()); + assert_ne!(tool.description().len(), 0); + assert!(tool.description().contains("image")); +} + +#[test] +fn image_info_tool_schema() { + let tool = ImageInfoTool::new(test_security()); + let schema = tool.parameters_schema(); + assert!(schema["properties"]["path"].is_object()); + assert!(schema["properties"]["include_base64"].is_object()); + let required = schema["required"].as_array().unwrap(); + assert!(required.contains(&json!("path"))); +} + +#[test] +fn image_info_tool_spec() { + let tool = ImageInfoTool::new(test_security()); + let spec = tool.spec(); + assert_eq!(spec.name, "image_info"); + assert!(spec.parameters.is_object()); +} + +// ── Format detection ──────────────────────────────────────── + +#[test] +fn detect_png() { + let bytes = b"\x89PNG\r\n\x1a\n"; + assert_eq!(ImageInfoTool::detect_format(bytes), "png"); +} + +#[test] +fn detect_jpeg() { + let bytes = b"\xFF\xD8\xFF\xE0"; + assert_eq!(ImageInfoTool::detect_format(bytes), "jpeg"); +} + +#[test] +fn detect_gif() { + let bytes = b"GIF89a"; + assert_eq!(ImageInfoTool::detect_format(bytes), "gif"); +} + +#[test] +fn detect_webp() { + let bytes = b"RIFF\x00\x00\x00\x00WEBP"; + assert_eq!(ImageInfoTool::detect_format(bytes), "webp"); +} + +#[test] +fn detect_bmp() { + let bytes = b"BM\x00\x00"; + assert_eq!(ImageInfoTool::detect_format(bytes), "bmp"); +} + +#[test] +fn detect_unknown_short() { + let bytes = b"\x00\x01"; + assert_eq!(ImageInfoTool::detect_format(bytes), "unknown"); +} + +#[test] +fn detect_unknown_garbage() { + let bytes = b"this is not an image"; + assert_eq!(ImageInfoTool::detect_format(bytes), "unknown"); +} + +// ── Dimension extraction ──────────────────────────────────── + +#[test] +fn png_dimensions() { + // Minimal PNG IHDR: 8-byte signature + 4-byte length + 4-byte IHDR + 4-byte width + 4-byte height + let mut bytes = vec![ + 0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A, // PNG signature + 0x00, 0x00, 0x00, 0x0D, // IHDR length + 0x49, 0x48, 0x44, 0x52, // "IHDR" + 0x00, 0x00, 0x03, 0x20, // width: 800 + 0x00, 0x00, 0x02, 0x58, // height: 600 + ]; + bytes.extend_from_slice(&[0u8; 10]); // padding + let dims = ImageInfoTool::extract_dimensions(&bytes, "png"); + assert_eq!(dims, Some((800, 600))); +} + +#[test] +fn gif_dimensions() { + let bytes = [ + 0x47, 0x49, 0x46, 0x38, 0x39, 0x61, // GIF89a + 0x40, 0x01, // width: 320 (LE) + 0xF0, 0x00, // height: 240 (LE) + ]; + let dims = ImageInfoTool::extract_dimensions(&bytes, "gif"); + assert_eq!(dims, Some((320, 240))); +} + +#[test] +fn bmp_dimensions() { + let mut bytes = vec![0u8; 26]; + bytes[0] = b'B'; + bytes[1] = b'M'; + // width at offset 18 (LE): 1024 + bytes[18] = 0x00; + bytes[19] = 0x04; + bytes[20] = 0x00; + bytes[21] = 0x00; + // height at offset 22 (LE): 768 + bytes[22] = 0x00; + bytes[23] = 0x03; + bytes[24] = 0x00; + bytes[25] = 0x00; + let dims = ImageInfoTool::extract_dimensions(&bytes, "bmp"); + assert_eq!(dims, Some((1024, 768))); +} + +#[test] +fn jpeg_dimensions() { + // Minimal JPEG-like byte sequence with SOF0 marker + let mut bytes: Vec = vec![ + 0xFF, 0xD8, // SOI + 0xFF, 0xE0, // APP0 marker + 0x00, 0x10, // APP0 length = 16 + ]; + bytes.extend_from_slice(&[0u8; 14]); // APP0 payload + bytes.extend_from_slice(&[ + 0xFF, 0xC0, // SOF0 marker + 0x00, 0x11, // SOF0 length + 0x08, // precision + 0x01, 0xE0, // height: 480 + 0x02, 0x80, // width: 640 + ]); + let dims = ImageInfoTool::extract_dimensions(&bytes, "jpeg"); + assert_eq!(dims, Some((640, 480))); +} + +#[test] +fn jpeg_malformed_zero_length_segment() { + // Zero-length segment should return None instead of looping forever + let bytes: Vec = vec![ + 0xFF, 0xD8, // SOI + 0xFF, 0xE0, // APP0 marker + 0x00, 0x00, // length = 0 (malformed) + ]; + let dims = ImageInfoTool::extract_dimensions(&bytes, "jpeg"); + assert!(dims.is_none()); +} + +#[test] +fn unknown_format_no_dimensions() { + let bytes = b"random data here"; + let dims = ImageInfoTool::extract_dimensions(bytes, "unknown"); + assert!(dims.is_none()); +} + +// ── Execute tests ─────────────────────────────────────────── + +#[tokio::test] +async fn execute_missing_path() { + let tool = ImageInfoTool::new(test_security()); + let result = tool.execute(json!({})).await; + assert!(result.is_err()); +} + +#[tokio::test] +async fn execute_nonexistent_file() { + let tool = ImageInfoTool::new(test_security()); + let result = tool + .execute(json!({"path": "nonexistent_image_xyz.png"})) + .await + .unwrap(); + assert!(result.is_error); + assert!( + result.output().contains("not allowed") || result.output().contains("Failed to resolve"), + "unexpected error: {}", + result.output() + ); +} + +#[tokio::test] +async fn execute_real_file() { + // Create a minimal valid PNG + let dir = std::env::temp_dir().join("tinytools_image_info_test"); + let _ = tokio::fs::create_dir_all(&dir).await; + let png_path = dir.join("test.png"); + + // Minimal 1x1 red PNG (67 bytes) + let png_bytes: Vec = vec![ + 0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A, // signature + 0x00, 0x00, 0x00, 0x0D, // IHDR length + 0x49, 0x48, 0x44, 0x52, // IHDR + 0x00, 0x00, 0x00, 0x01, // width: 1 + 0x00, 0x00, 0x00, 0x01, // height: 1 + 0x08, 0x02, 0x00, 0x00, 0x00, // bit depth, color type, etc. + 0x90, 0x77, 0x53, 0xDE, // CRC + 0x00, 0x00, 0x00, 0x0C, // IDAT length + 0x49, 0x44, 0x41, 0x54, // IDAT + 0x08, 0xD7, 0x63, 0xF8, 0xCF, 0xC0, 0x00, 0x00, 0x00, 0x02, 0x00, 0x01, 0xE2, 0x21, 0xBC, + 0x33, // CRC + 0x00, 0x00, 0x00, 0x00, // IEND length + 0x49, 0x45, 0x4E, 0x44, // IEND + 0xAE, 0x42, 0x60, 0x82, // CRC + ]; + tokio::fs::write(&png_path, &png_bytes).await.unwrap(); + + let tool = ImageInfoTool::new(test_security()); + let result = tool + .execute(json!({"path": png_path.to_string_lossy()})) + .await + .unwrap(); + assert!(!result.is_error); + assert!(result.output().contains("Format: png")); + assert!(result.output().contains("Dimensions: 1x1")); + assert!(!result.output().contains("data:")); + + // Clean up + let _ = tokio::fs::remove_dir_all(&dir).await; +} + +/// A decompression-bomb-shaped PNG (a header claiming an enormous canvas, +/// backed by a file well under `MAX_IMAGE_BYTES`) must not blow up memory or +/// hang: `extract_dimensions` reads four fixed IHDR bytes and never inflates +/// pixel data, so the claimed 60000x60000 canvas is read as fast as any other +/// header and no allocation is proportional to it. +#[tokio::test] +async fn execute_reports_a_bomb_shaped_header_without_decoding_pixels() { + let dir = std::env::temp_dir().join("tinytools_image_info_bomb"); + let _ = tokio::fs::create_dir_all(&dir).await; + let png_path = dir.join("bomb.png"); + + let mut bytes = vec![ + 0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A, // signature + 0x00, 0x00, 0x00, 0x0D, // IHDR length + 0x49, 0x48, 0x44, 0x52, // "IHDR" + 0x00, 0x00, 0xEA, 0x60, // width: 60000 + 0x00, 0x00, 0xEA, 0x60, // height: 60000 + ]; + bytes.extend_from_slice(&[0u8; 10]); // pad past the 24 bytes IHDR parsing reads + tokio::fs::write(&png_path, &bytes).await.unwrap(); + + let tool = ImageInfoTool::new(test_security()); + let result = tokio::time::timeout( + std::time::Duration::from_secs(2), + tool.execute(json!({"path": png_path.to_string_lossy()})), + ) + .await + .expect("image-info execution must settle within two seconds") + .unwrap(); + assert!(!result.is_error, "{}", result.output()); + assert!(result.output().contains("Dimensions: 60000x60000")); + + let _ = tokio::fs::remove_dir_all(&dir).await; +} + +#[tokio::test] +async fn execute_with_base64() { + let dir = std::env::temp_dir().join("tinytools_image_info_b64"); + let _ = tokio::fs::create_dir_all(&dir).await; + let png_path = dir.join("test_b64.png"); + + // Minimal 1x1 PNG + let png_bytes: Vec = vec![ + 0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A, 0x00, 0x00, 0x00, 0x0D, 0x49, 0x48, 0x44, + 0x52, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x01, 0x08, 0x02, 0x00, 0x00, 0x00, 0x90, + 0x77, 0x53, 0xDE, 0x00, 0x00, 0x00, 0x0C, 0x49, 0x44, 0x41, 0x54, 0x08, 0xD7, 0x63, 0xF8, + 0xCF, 0xC0, 0x00, 0x00, 0x00, 0x02, 0x00, 0x01, 0xE2, 0x21, 0xBC, 0x33, 0x00, 0x00, 0x00, + 0x00, 0x49, 0x45, 0x4E, 0x44, 0xAE, 0x42, 0x60, 0x82, + ]; + tokio::fs::write(&png_path, &png_bytes).await.unwrap(); + + let tool = ImageInfoTool::new(test_security()); + let result = tool + .execute(json!({"path": png_path.to_string_lossy(), "include_base64": true})) + .await + .unwrap(); + assert!(!result.is_error); + assert!(result.output().contains("data:image/png;base64,")); + + let _ = tokio::fs::remove_dir_all(&dir).await; +} + +#[test] +fn truncated_headers_yield_no_dimensions() { + for format in ["png", "gif", "bmp"] { + assert_eq!( + ImageInfoTool::extract_dimensions(&[0u8; 5], format), + None, + "{format}" + ); + } +} + +#[test] +fn jpeg_dimension_scan_rejects_bad_structure() { + // A non-0xFF byte where a marker should be. + assert_eq!( + ImageInfoTool::extract_dimensions(&[0xFF, 0xD8, 0x00, 0x00], "jpeg"), + None + ); + // SOF marker with too few bytes after it. + assert_eq!( + ImageInfoTool::extract_dimensions(&[0xFF, 0xD8, 0xFF, 0xC0, 0, 0], "jpeg"), + None + ); + // A skippable marker with no length bytes left. + assert_eq!( + ImageInfoTool::extract_dimensions(&[0xFF, 0xD8, 0xFF, 0xE0], "jpeg"), + None + ); + // A well-formed segment that runs out before any SOF marker. + assert_eq!( + ImageInfoTool::extract_dimensions(&[0xFF, 0xD8, 0xFF, 0xE0, 0x00, 0x02], "jpeg"), + None + ); +} + +#[tokio::test] +async fn execute_on_a_directory_reports_a_read_failure() { + let dir = std::env::temp_dir().join("tinytools_image_info_dir_test"); + tokio::fs::create_dir_all(&dir).await.unwrap(); + let tool = ImageInfoTool::new(test_security()); + let err = tool + .execute(json!({"path": dir.to_string_lossy()})) + .await + .unwrap_err(); + assert!(err.to_string().contains("Failed to read image file")); +} diff --git a/crates/tinyagents-harness-tools/src/fs/filesystem/list_files/extra_tests.rs b/crates/tinyagents-harness-tools/src/fs/filesystem/list_files/extra_tests.rs new file mode 100644 index 000000000..32dafc69b --- /dev/null +++ b/crates/tinyagents-harness-tools/src/fs/filesystem/list_files/extra_tests.rs @@ -0,0 +1,52 @@ +//! Budget, delegation and entry-cap tests for `list_files`. + +#![allow(clippy::expect_used, clippy::panic, clippy::unwrap_used)] + +use super::*; +use crate::fs::filesystem::test_support::{TestGate, WorkspaceContext, WrapGate}; +use tempfile::TempDir; + +#[tokio::test] +async fn rate_limit_and_exhausted_budget_are_refused() { + let dir = TempDir::new().unwrap(); + let limited = ListFilesTool::new(TestGate::with( + dir.path().to_path_buf(), + crate::fs::filesystem::test_support::AutonomyLevel::Supervised, + 0, + )); + let r = limited.execute(json!({})).await.unwrap(); + assert!(r.output().contains("too many actions in the last hour")); + + let gate = WrapGate::new(TestGate::supervised(dir.path().to_path_buf())) + .no_budget() + .arc(); + let r = ListFilesTool::new(gate).execute(json!({})).await.unwrap(); + assert!(r.output().contains("action budget exhausted")); +} + +#[tokio::test] +async fn execute_with_context_delegates() { + let dir = TempDir::new().unwrap(); + std::fs::write(dir.path().join("f.txt"), "").unwrap(); + let tool = ListFilesTool::new(TestGate::supervised(dir.path().to_path_buf())); + let ctx = WorkspaceContext::at(dir.path()); + let r = tool + .execute_with_context(json!({}), ToolCallOptions::default(), Some(&ctx)) + .await + .unwrap(); + assert!(r.output().contains("file\tf.txt")); +} + +#[tokio::test] +async fn listing_stops_at_the_entry_cap() { + let dir = TempDir::new().unwrap(); + for i in 0..(MAX_ENTRIES + 5) { + std::fs::write(dir.path().join(format!("f{i:05}")), "").unwrap(); + } + let tool = ListFilesTool::new(TestGate::supervised(dir.path().to_path_buf())); + let r = tool.execute(json!({})).await.unwrap(); + assert!( + r.output() + .starts_with(&format!("{MAX_ENTRIES} entr(ies) in .")) + ); +} diff --git a/crates/tinyagents-harness-tools/src/fs/filesystem/list_files/mod.rs b/crates/tinyagents-harness-tools/src/fs/filesystem/list_files/mod.rs new file mode 100644 index 000000000..dd07adb98 --- /dev/null +++ b/crates/tinyagents-harness-tools/src/fs/filesystem/list_files/mod.rs @@ -0,0 +1,147 @@ +//! `list` — directory listing. +//! +//! Coding-harness baseline tool (issue #1205): non-recursive directory +//! listing keyed by a workspace-relative path. Distinguishes files, +//! directories, and symlinks. Path sandboxing matches `file_read`. + +use super::gate::{FsGate, gate_for_context}; +use async_trait::async_trait; +use serde_json::json; +use std::sync::Arc; +use tinytools::ToolRunContext; +use tinytools::{PermissionLevel, Tool, ToolCallOptions, ToolResult}; + +const MAX_ENTRIES: usize = 1_000; + +/// List the entries of a workspace directory. +#[derive(Debug)] +pub struct ListFilesTool { + gate: Arc, +} + +impl ListFilesTool { + /// Create the tool. + #[must_use] + pub fn new(gate: Arc) -> Self { + Self { gate } + } +} + +#[async_trait] +impl Tool for ListFilesTool { + fn name(&self) -> &'static str { + "list" + } + + fn description(&self) -> &'static str { + "List entries in a workspace directory (non-recursive). Each line is \ + `\t` where kind is `dir`, `file`, or `link`." + } + + fn parameters_schema(&self) -> serde_json::Value { + json!({ + "type": "object", + "properties": { + "path": { + "type": "string", + "description": "Directory path relative to the workspace (default `.`).", + "default": "." + } + } + }) + } + + fn permission_level(&self) -> PermissionLevel { + PermissionLevel::ReadOnly + } + + async fn execute(&self, args: serde_json::Value) -> anyhow::Result { + self.execute_in_context(args, None).await + } + + async fn execute_with_context( + &self, + args: serde_json::Value, + _options: ToolCallOptions, + context: Option<&dyn ToolRunContext>, + ) -> anyhow::Result { + self.execute_in_context(args, context).await + } +} + +impl ListFilesTool { + async fn execute_in_context( + &self, + args: serde_json::Value, + context: Option<&dyn ToolRunContext>, + ) -> anyhow::Result { + let path = args.get("path").and_then(|v| v.as_str()).unwrap_or("."); + + if self.gate.is_rate_limited() { + return Ok(ToolResult::error( + "Rate limit exceeded: too many actions in the last hour", + )); + } + if !self.gate.record_action() { + return Ok(ToolResult::error( + "Rate limit exceeded: action budget exhausted", + )); + } + + // Security check: validate path string, resolve symlinks, confirm workspace containment. + let path_policy = gate_for_context(&self.gate, context, "list"); + let resolved = match path_policy.validate_path(path).await { + Ok(p) => p, + Err(msg) => return Ok(ToolResult::error(msg)), + }; + + let mut read = match tokio::fs::read_dir(&resolved).await { + Ok(r) => r, + Err(e) => return Ok(ToolResult::error(format!("Failed to read directory: {e}"))), + }; + + let mut entries: Vec<(String, String)> = Vec::new(); + loop { + match read.next_entry().await { + Ok(Some(entry)) => { + let name = entry.file_name().to_string_lossy().into_owned(); + let kind = match entry.file_type().await { + Ok(t) if t.is_symlink() => "link", + Ok(t) if t.is_dir() => "dir", + Ok(_) => "file", + Err(_) => "unknown", + }; + entries.push((kind.to_string(), name)); + if entries.len() >= MAX_ENTRIES { + break; + } + } + Ok(None) => break, + Err(e) => { + return Ok(ToolResult::error(format!( + "Failed to enumerate directory: {e}" + ))); + } + } + } + + entries.sort_by(|a, b| a.1.cmp(&b.1)); + + let mut body = format!("{} entr(ies) in {path}", entries.len()); + for (kind, name) in entries { + body.push('\n'); + body.push_str(&kind); + body.push('\t'); + body.push_str(&name); + } + Ok(ToolResult::success(body)) + } +} + +#[cfg(test)] +#[path = "mod_tests.rs"] +mod tests; + +#[cfg(test)] +#[path = "extra_tests.rs"] +mod extra_tests; diff --git a/crates/tinyagents-harness-tools/src/fs/filesystem/list_files/mod_tests.rs b/crates/tinyagents-harness-tools/src/fs/filesystem/list_files/mod_tests.rs new file mode 100644 index 000000000..0d6c88032 --- /dev/null +++ b/crates/tinyagents-harness-tools/src/fs/filesystem/list_files/mod_tests.rs @@ -0,0 +1,76 @@ +//! Behavior tests for the `list_files` tool, driven through a fake [`FsGate`]. + +#![allow(clippy::expect_used, clippy::panic, clippy::unwrap_used)] + +use super::*; +use crate::fs::filesystem::test_support::TestGate; + +fn test_security(workspace: std::path::PathBuf) -> Arc { + TestGate::supervised(workspace) +} + +#[test] +fn list_name() { + let tool = ListFilesTool::new(test_security(std::env::temp_dir())); + assert_eq!(tool.name(), "list"); +} + +#[tokio::test] +async fn list_lists_files_and_dirs() { + let dir = std::env::temp_dir().join("openhuman_test_list"); + let _ = tokio::fs::remove_dir_all(&dir).await; + tokio::fs::create_dir_all(dir.join("sub")).await.unwrap(); + tokio::fs::write(dir.join("a.txt"), "x").await.unwrap(); + + let tool = ListFilesTool::new(test_security(dir.clone())); + let result = tool.execute(json!({})).await.unwrap(); + assert!(!result.is_error); + let output = result.output(); + assert!(output.contains("file\ta.txt")); + assert!(output.contains("dir\tsub")); + + let _ = tokio::fs::remove_dir_all(&dir).await; +} + +#[tokio::test] +async fn list_blocks_path_traversal() { + let tool = ListFilesTool::new(test_security(std::env::temp_dir())); + let result = tool.execute(json!({"path": "../../etc"})).await.unwrap(); + assert!(result.is_error); + assert!(result.output().contains("not allowed")); +} + +#[tokio::test] +async fn list_missing_dir() { + let dir = std::env::temp_dir().join("openhuman_test_list_missing"); + let _ = tokio::fs::remove_dir_all(&dir).await; + tokio::fs::create_dir_all(&dir).await.unwrap(); + let tool = ListFilesTool::new(test_security(dir.clone())); + let result = tool.execute(json!({"path": "nope"})).await.unwrap(); + assert!(result.is_error); + assert!(result.output().contains("Failed to resolve")); + let _ = tokio::fs::remove_dir_all(&dir).await; +} + +#[tokio::test] +async fn list_a_plain_file_reports_the_read_dir_failure() { + // Distinct from `list_missing_dir`: the path resolves (the file exists), + // so `validate_path` succeeds; the failure is `read_dir` itself refusing + // a non-directory (ENOTDIR), one step later in the same function. + let dir = std::env::temp_dir().join("openhuman_test_list_not_a_dir"); + let _ = tokio::fs::remove_dir_all(&dir).await; + tokio::fs::create_dir_all(&dir).await.unwrap(); + tokio::fs::write(dir.join("plain.txt"), "not a directory") + .await + .unwrap(); + + let tool = ListFilesTool::new(test_security(dir.clone())); + let result = tool.execute(json!({"path": "plain.txt"})).await.unwrap(); + assert!( + result.is_error, + "listing a plain file must be a refusal, not an empty-looking success" + ); + assert!(result.output().contains("Failed to read directory")); + + let _ = tokio::fs::remove_dir_all(&dir).await; +} diff --git a/crates/tinyagents-harness-tools/src/fs/filesystem/mod.rs b/crates/tinyagents-harness-tools/src/fs/filesystem/mod.rs new file mode 100644 index 000000000..2dba8e04c --- /dev/null +++ b/crates/tinyagents-harness-tools/src/fs/filesystem/mod.rs @@ -0,0 +1,73 @@ +//! Filesystem tools: read/write/edit/patch, grep/glob/list, diffs, CSV export, +//! git operations, linter and test runners, the workspace memory index, a +//! read-only workspace overview and image metadata. +//! +//! Every tool that can touch a path takes an `Arc`: the tool +//! performs the I/O, the host's gate decides whether it is allowed. The gate +//! trait ([`FsGate`]) is the whole seam; nothing here knows what an autonomy +//! level, an approval prompt or a workspace boundary is. + +// These tools moved here verbatim from a host crate; their behavior, messages +// and control flow are pinned by tests and by the fixtures in `fixtures/`, so +// the purely stylistic pedantic lints below are allowed rather than reshaping +// working code (long `execute` bodies, `usize as f64` size labels, and so on). +#![allow( + clippy::too_many_lines, + clippy::cast_precision_loss, + clippy::cast_possible_truncation, + clippy::similar_names, + clippy::used_underscore_binding, + clippy::manual_let_else, + clippy::format_push_string, + clippy::unwrap_used, + clippy::unused_self, + clippy::unnecessary_debug_formatting +)] + +mod apply_patch; +mod csv_export; +mod edit_file; +mod file_read; +mod file_sink; +mod file_write; +mod gate; +mod git_operations; +mod glob_search; +mod grep; +mod image_info; +mod list_files; +mod read_diff; +mod read_only; +mod run_linter; +mod run_tests; +mod text; +mod update_memory_md; +mod workspace_state; + +#[cfg(test)] +#[path = "gate_tests.rs"] +mod gate_tests; +#[cfg(test)] +mod test_support; + +pub use apply_patch::ApplyPatchTool; +pub use csv_export::CsvExportTool; +pub use edit_file::EditFileTool; +pub use file_read::FileReadTool; +pub use file_write::FileWriteTool; +pub use gate::FsGate; +pub use git_operations::{GitOperationsTool, shell_git_env}; +pub use glob_search::GlobTool; +pub use grep::GrepTool; +pub use image_info::ImageInfoTool; +pub use list_files::ListFilesTool; +pub use read_diff::ReadDiffTool; +pub use read_only::{is_read_only_call, is_read_only_command}; +pub use run_linter::RunLinterTool; +pub use run_tests::RunTestsTool; +pub use update_memory_md::UpdateMemoryMdTool; +pub use workspace_state::WorkspaceStateTool; + +#[cfg(test)] +#[path = "contract_tests.rs"] +mod contract_tests; diff --git a/crates/tinyagents-harness-tools/src/fs/filesystem/read_diff/mod.rs b/crates/tinyagents-harness-tools/src/fs/filesystem/read_diff/mod.rs new file mode 100644 index 000000000..d7f5a5c09 --- /dev/null +++ b/crates/tinyagents-harness-tools/src/fs/filesystem/read_diff/mod.rs @@ -0,0 +1,139 @@ +//! Tool: `read_diff` — structured git diff output for the Critic archetype. + +use async_trait::async_trait; +use serde_json::json; +use std::path::PathBuf; +use tinytools::ToolRunContext; +use tinytools::{PermissionLevel, Tool, ToolCallOptions, ToolResult}; + +/// Returns `git diff` output in a structured format. +#[derive(Debug)] +pub struct ReadDiffTool { + workspace_dir: PathBuf, +} + +impl ReadDiffTool { + #[must_use] + /// Create the tool. + pub fn new(workspace_dir: PathBuf) -> Self { + Self { workspace_dir } + } + + fn workspace_dir_for_context(&self, context: Option<&dyn ToolRunContext>) -> PathBuf { + if let Some(workspace) = context.and_then(|ctx| ctx.workspace()) { + tracing::debug!( + workspace_root = %workspace.root.display(), + policy_id = %workspace.policy_id, + "[read_diff] using TinyAgents workspace descriptor as workspace dir" + ); + return workspace.root.clone(); + } + self.workspace_dir.clone() + } +} + +#[async_trait] +impl Tool for ReadDiffTool { + fn exposure(&self) -> tinytools::ToolExposure { + tinytools::ToolExposure::Deferred + } + + fn name(&self) -> &'static str { + "read_diff" + } + + fn description(&self) -> &'static str { + "Get the git diff of current changes. Can diff staged, unstaged, or against a \ + specific base branch/commit. Returns file paths and hunks." + } + + fn parameters_schema(&self) -> serde_json::Value { + json!({ + "type": "object", + "properties": { + "base": { + "type": "string", + "description": "Base ref to diff against (e.g. 'main', 'HEAD~3'). Default: unstaged changes." + }, + "staged": { + "type": "boolean", + "description": "Show staged changes only (--cached). Default: false." + }, + "path_filter": { + "type": "string", + "description": "Limit diff to a specific path or glob." + } + } + }) + } + + fn permission_level(&self) -> PermissionLevel { + PermissionLevel::ReadOnly + } + + async fn execute(&self, args: serde_json::Value) -> anyhow::Result { + self.execute_with_context(args, ToolCallOptions::default(), None) + .await + } + + async fn execute_with_context( + &self, + args: serde_json::Value, + _options: ToolCallOptions, + context: Option<&dyn ToolRunContext>, + ) -> anyhow::Result { + let workspace_dir = self.workspace_dir_for_context(context); + let base = args.get("base").and_then(|v| v.as_str()); + let staged = args + .get("staged") + .and_then(serde_json::Value::as_bool) + .unwrap_or(false); + let path_filter = args.get("path_filter").and_then(|v| v.as_str()); + + let mut git_args = vec!["diff", "--stat", "-p"]; + + if staged { + git_args.push("--cached"); + } + + let base_str = base.map(std::string::ToString::to_string); + if let Some(ref bs) = base_str { + git_args.push(bs); + } + + if let Some(pf) = path_filter { + git_args.push("--"); + git_args.push(pf); + } + + tracing::debug!( + workspace = %workspace_dir.display(), + ?git_args, + "[read_diff] running git diff" + ); + + let output = tokio::process::Command::new("git") + .args(&git_args) + .current_dir(&workspace_dir) + .output() + .await?; + + if output.status.success() { + let diff = String::from_utf8_lossy(&output.stdout); + tracing::debug!("[read_diff] success, diff length={}", diff.len()); + if diff.trim().is_empty() { + Ok(ToolResult::success("No changes found.")) + } else { + Ok(ToolResult::success(diff.to_string())) + } + } else { + let stderr = String::from_utf8_lossy(&output.stderr); + tracing::debug!("[read_diff] failed: {stderr}"); + Ok(ToolResult::error(stderr.to_string())) + } + } +} + +#[cfg(test)] +#[path = "mod_tests.rs"] +mod tests; diff --git a/crates/tinyagents-harness-tools/src/fs/filesystem/read_diff/mod_tests.rs b/crates/tinyagents-harness-tools/src/fs/filesystem/read_diff/mod_tests.rs new file mode 100644 index 000000000..a767621de --- /dev/null +++ b/crates/tinyagents-harness-tools/src/fs/filesystem/read_diff/mod_tests.rs @@ -0,0 +1,144 @@ +#![allow(clippy::expect_used, clippy::panic, clippy::unwrap_used)] + +use super::*; +use serde_json::json; +use tempfile::TempDir; + +fn make_tool(dir: &TempDir) -> ReadDiffTool { + ReadDiffTool::new(dir.path().to_path_buf()) +} + +#[test] +fn name_is_correct() { + let tmp = TempDir::new().unwrap(); + assert_eq!(make_tool(&tmp).name(), "read_diff"); +} + +#[test] +fn description_is_non_empty() { + let tmp = TempDir::new().unwrap(); + assert_ne!(make_tool(&tmp).description().len(), 0); +} + +#[test] +fn schema_is_object_type() { + let tmp = TempDir::new().unwrap(); + let schema = make_tool(&tmp).parameters_schema(); + assert_eq!(schema["type"], "object"); +} + +#[test] +fn permission_level_is_read_only() { + let tmp = TempDir::new().unwrap(); + assert_eq!( + make_tool(&tmp).permission_level(), + PermissionLevel::ReadOnly + ); +} + +#[tokio::test] +async fn execute_returns_error_for_non_git_dir() { + let tmp = TempDir::new().unwrap(); + let result = make_tool(&tmp).execute(json!({})).await.unwrap(); + // Non-git dir: git will fail, tool returns error + assert!(result.is_error); +} + +#[tokio::test] +async fn execute_no_changes_in_clean_git_repo() { + let tmp = TempDir::new().unwrap(); + // Init a git repo and make an initial commit so there's nothing to diff + let _ = std::process::Command::new("git") + .args(["init"]) + .current_dir(tmp.path()) + .output(); + let _ = std::process::Command::new("git") + .args(["commit", "--allow-empty", "-m", "init"]) + .current_dir(tmp.path()) + .env("GIT_AUTHOR_NAME", "test") + .env("GIT_AUTHOR_EMAIL", "t@t.com") + .env("GIT_COMMITTER_NAME", "test") + .env("GIT_COMMITTER_EMAIL", "t@t.com") + .output(); + let result = make_tool(&tmp).execute(json!({})).await.unwrap(); + assert!(!result.is_error); + assert!(result.output().contains("No changes found.")); +} + +fn git_in(dir: &std::path::Path, args: &[&str]) { + let out = std::process::Command::new("git") + .args(["-c", "user.email=t@example.invalid", "-c", "user.name=T"]) + .args(args) + .current_dir(dir) + .env("GIT_CONFIG_NOSYSTEM", "1") + .env("GIT_CONFIG_GLOBAL", "/dev/null") + .output() + .unwrap(); + assert!( + out.status.success(), + "{}", + String::from_utf8_lossy(&out.stderr) + ); +} + +fn repo_with_change() -> TempDir { + let tmp = TempDir::new().unwrap(); + git_in(tmp.path(), &["init"]); + std::fs::write(tmp.path().join("a.txt"), "one\n").unwrap(); + std::fs::write(tmp.path().join("b.txt"), "one\n").unwrap(); + git_in(tmp.path(), &["add", "."]); + git_in(tmp.path(), &["commit", "-m", "init"]); + std::fs::write(tmp.path().join("a.txt"), "two\n").unwrap(); + std::fs::write(tmp.path().join("b.txt"), "two\n").unwrap(); + tmp +} + +#[tokio::test] +async fn path_filter_limits_the_diff() { + let tmp = repo_with_change(); + let result = make_tool(&tmp) + .execute(json!({"path_filter": "a.txt"})) + .await + .unwrap(); + assert!(!result.is_error); + assert!(result.output().contains("a.txt")); + assert!(!result.output().contains("b.txt")); +} + +#[tokio::test] +async fn base_ref_and_staged_flags_are_passed_to_git() { + let tmp = repo_with_change(); + git_in(tmp.path(), &["add", "a.txt"]); + let staged = make_tool(&tmp) + .execute(json!({"staged": true})) + .await + .unwrap(); + assert!(staged.output().contains("a.txt")); + assert!(!staged.output().contains("b.txt")); + + let against_head = make_tool(&tmp) + .execute(json!({"base": "HEAD"})) + .await + .unwrap(); + assert!(against_head.output().contains("b.txt")); + + let bad_base = make_tool(&tmp) + .execute(json!({"base": "no-such-ref"})) + .await + .unwrap(); + assert!(bad_base.is_error); +} + +#[tokio::test] +async fn a_workspace_descriptor_overrides_the_configured_directory() { + use crate::fs::filesystem::test_support::WorkspaceContext; + let elsewhere = TempDir::new().unwrap(); + let repo = repo_with_change(); + let context = WorkspaceContext::at(repo.path()); + let result = make_tool(&elsewhere) + .execute_with_context(json!({}), ToolCallOptions::default(), Some(&context)) + .await + .unwrap(); + assert!(!result.is_error, "{}", result.output()); + assert!(result.output().contains("a.txt")); +} diff --git a/crates/tinyagents-harness-tools/src/fs/filesystem/read_only.rs b/crates/tinyagents-harness-tools/src/fs/filesystem/read_only.rs new file mode 100644 index 000000000..a008504b4 --- /dev/null +++ b/crates/tinyagents-harness-tools/src/fs/filesystem/read_only.rs @@ -0,0 +1,92 @@ +//! Which tool calls only read files the agent can read again. +//! +//! A host skips an LLM summary of tool output for these. A summary of source code is a paraphrase, and an agent about to +//! edit a file needs its exact text: on live coding runs every summarized +//! `cat` of a source file was followed by `juice_retrieve` of the whole +//! original or a fresh `sed -n` read of the file. The summary bought nothing +//! and cost a model call plus up to `TinyJuice`'s 8-second summary timeout, +//! which it hit on most attempts. +//! +//! The test is deliberately narrow: a file-reading tool, or a shell command +//! made only of `cd` and plain reads (`cat`, `head`, `tail`, `nl`, `sed` +//! without `-i`, `echo`/`printf` separators), with no redirects, heredocs, +//! substitutions or other programs. Anything else (a build, a test run, a +//! search) keeps the normal pipeline: its output is not a file on disk. + +use serde_json::Value; + +/// Tools whose whole job is returning a file's contents. +const FILE_READ_TOOLS: &[&str] = &["file_read", "read_file", "read", "view_file"]; + +/// Shell tools, by the names hosts and common harness adapters use. +const SHELL_TOOLS: &[&str] = &["shell", "bash", "exec", "run_command", "terminal"]; + +/// Programs that print a file (`sed` only without in-place editing). +const READERS: &[&str] = &["cat", "head", "tail", "nl", "sed"]; + +/// Programs allowed beside the readers: navigation and separators. +const NEUTRAL: &[&str] = &["cd", "echo", "printf", "true", "pwd"]; + +/// Whether this call only prints files the agent can re-read (see the module +/// docs). `use_skill` is followed into the tool it wraps. +#[must_use] +pub fn is_read_only_call(tool_name: &str, args: &Value) -> bool { + let (name, args) = if tool_name == "use_skill" { + match (args.get("tool").and_then(Value::as_str), args.get("args")) { + (Some(inner), Some(inner_args)) => (inner, inner_args), + _ => return false, + } + } else { + (tool_name, args) + }; + if FILE_READ_TOOLS.contains(&name) { + return true; + } + if !SHELL_TOOLS.contains(&name) { + return false; + } + args.get("command") + .or_else(|| args.get("cmd")) + .and_then(Value::as_str) + .is_some_and(is_read_only_command) +} + +/// Whether a shell command line only prints files: every segment (split on +/// `&&`, `||`, `;`, `|` and newlines) runs a reader or a neutral program, at +/// least one runs a reader, and nothing redirects, substitutes or feeds a +/// heredoc. +#[must_use] +pub fn is_read_only_command(command: &str) -> bool { + if command.contains(['>', '<', '`', '$']) { + return false; + } + let mut reads = false; + for segment in command + .split(['\n', ';', '|', '&']) + .map(str::trim) + .filter(|s| !s.is_empty()) + { + let mut words = segment.split_whitespace(); + let Some(program) = words.next() else { + continue; + }; + let program = program.rsplit('/').next().unwrap_or(program); + if program == "sed" + && segment + .split_whitespace() + .any(|w| w.starts_with("-i") || w == "--in-place" || w.starts_with("--in-place=")) + { + return false; + } + if READERS.contains(&program) { + reads = true; + } else if !NEUTRAL.contains(&program) { + return false; + } + } + reads +} + +#[cfg(test)] +#[path = "read_only_tests.rs"] +mod tests; diff --git a/crates/tinyagents-harness-tools/src/fs/filesystem/read_only_tests.rs b/crates/tinyagents-harness-tools/src/fs/filesystem/read_only_tests.rs new file mode 100644 index 000000000..ce034c70f --- /dev/null +++ b/crates/tinyagents-harness-tools/src/fs/filesystem/read_only_tests.rs @@ -0,0 +1,58 @@ +use super::*; +use serde_json::json; + +#[test] +fn plain_reads_of_files_are_file_reads() { + for command in [ + "cat vm/vmExprFunction.go", + "cd /app && cat vm/vmExprFunction.go", + "cd /app && cat -n parser/lexer.go", + "cd /app && sed -n '120,300p' vm/vmExprFunction.go", + "cd /app && cat parser/Makefile && echo \"=== parser.go.y ===\" && cat parser/parser.go.y", + "cd /app && sed -n '125,150p' ast/expr.go && echo \"=== vmExprFunction.go ===\" && cat vm/vmExprFunction.go", + "head -50 README.md; tail -20 CHANGELOG.md", + "cat -n src/lib.rs | sed -n '1,80p'", + ] { + assert!(is_read_only_command(command), "{command}"); + } +} + +#[test] +fn anything_that_runs_writes_or_searches_is_not() { + for command in [ + "go test ./vm/...", + "cd /app && git diff", + "grep -rn FuncExpr .", + "cat foo.go > bar.go", + "sed -i 's/a/b/' foo.go", + "sed --in-place 's/a/b/' foo.go", + "sed --in-place=.bak -n 's/old/new/p' file", + "cat <<'EOF' > x.go\npackage x\nEOF", + "cat $(git ls-files)", + "cat README.md && printf '%s\\n' \"$BUILD_LOG\"", + "cd /app && cat a.go && go build ./...", + "echo hello", + "cd /app", + ] { + assert!(!is_read_only_command(command), "{command}"); + } +} + +#[test] +fn file_read_tools_and_wrapped_calls_are_recognised() { + assert!(is_read_only_call("file_read", &json!({"path": "a.rs"}))); + assert!(is_read_only_call("shell", &json!({"command": "cat a.rs"}))); + assert!(!is_read_only_call( + "shell", + &json!({"command": "cargo test"}) + )); + assert!(!is_read_only_call( + "web_fetch", + &json!({"url": "https://example.com"}) + )); + assert!(is_read_only_call( + "use_skill", + &json!({"tool": "shell", "args": {"command": "cat a.rs"}}) + )); + assert!(!is_read_only_call("use_skill", &json!({"tool": "shell"}))); +} diff --git a/crates/tinyagents-harness-tools/src/fs/filesystem/run_linter/mod.rs b/crates/tinyagents-harness-tools/src/fs/filesystem/run_linter/mod.rs new file mode 100644 index 000000000..4ad182d6b --- /dev/null +++ b/crates/tinyagents-harness-tools/src/fs/filesystem/run_linter/mod.rs @@ -0,0 +1,159 @@ +//! Tool: `run_linter` — run linting tools for the Critic archetype. + +use async_trait::async_trait; +use serde_json::json; +use std::path::PathBuf; +use tinytools::ToolRunContext; +use tinytools::{PermissionLevel, Tool, ToolCallOptions, ToolResult}; + +/// Runs linters (cargo clippy, eslint) and returns structured findings. +#[derive(Debug)] +pub struct RunLinterTool { + workspace_dir: PathBuf, +} + +impl RunLinterTool { + #[must_use] + /// Create the tool. + pub fn new(workspace_dir: PathBuf) -> Self { + Self { workspace_dir } + } + + fn workspace_dir_for_context(&self, context: Option<&dyn ToolRunContext>) -> PathBuf { + if let Some(workspace) = context.and_then(|ctx| ctx.workspace()) { + tracing::debug!( + workspace_root = %workspace.root.display(), + policy_id = %workspace.policy_id, + "[run_linter] using TinyAgents workspace descriptor as workspace dir" + ); + return workspace.root.clone(); + } + self.workspace_dir.clone() + } +} + +#[async_trait] +impl Tool for RunLinterTool { + fn exposure(&self) -> tinytools::ToolExposure { + tinytools::ToolExposure::Deferred + } + + fn name(&self) -> &'static str { + "run_linter" + } + + fn description(&self) -> &'static str { + "Run linting tools on the codebase. Supports 'clippy' for Rust and 'eslint' for \ + TypeScript/JavaScript. Returns warnings and errors." + } + + fn parameters_schema(&self) -> serde_json::Value { + json!({ + "type": "object", + "properties": { + "linter": { + "type": "string", + "enum": ["clippy", "eslint", "auto"], + "description": "Which linter to run. 'auto' detects from project files.", + "default": "auto" + }, + "path": { + "type": "string", + "description": "Limit linting to a specific path." + } + } + }) + } + + fn permission_level(&self) -> PermissionLevel { + PermissionLevel::Execute + } + + async fn execute(&self, args: serde_json::Value) -> anyhow::Result { + self.execute_with_context(args, ToolCallOptions::default(), None) + .await + } + + async fn execute_with_context( + &self, + args: serde_json::Value, + _options: ToolCallOptions, + context: Option<&dyn ToolRunContext>, + ) -> anyhow::Result { + let workspace_dir = self.workspace_dir_for_context(context); + let linter = args + .get("linter") + .and_then(|v| v.as_str()) + .unwrap_or("auto"); + + let linter = if linter == "auto" { + if workspace_dir.join("Cargo.toml").exists() { + "clippy" + } else if workspace_dir.join("package.json").exists() { + "eslint" + } else { + return Ok(ToolResult::error( + "Could not detect project type for linting.", + )); + } + } else { + linter + }; + + let output = match linter { + "clippy" => { + tokio::process::Command::new("cargo") + .args([ + "clippy", + "--message-format=short", + "--", + "-W", + "clippy::all", + ]) + .current_dir(&workspace_dir) + .output() + .await? + } + "eslint" => { + let path = args.get("path").and_then(|v| v.as_str()).unwrap_or("."); + if path.starts_with('/') || path.contains("..") { + return Ok(ToolResult::error( + "path must be a relative path within the workspace \ + (no absolute paths or '..')", + )); + } + tokio::process::Command::new("npx") + .args(["eslint", "--format", "compact", path]) + .current_dir(&workspace_dir) + .output() + .await? + } + other => { + return Ok(ToolResult::error(format!("Unknown linter: {other}"))); + } + }; + + let stdout = String::from_utf8_lossy(&output.stdout); + let stderr = String::from_utf8_lossy(&output.stderr); + + let combined = if stdout.is_empty() { + stderr.to_string() + } else { + format!("{stdout}\n{stderr}") + }; + + if output.status.success() { + Ok(ToolResult::success(combined)) + } else { + Ok(ToolResult::error(format!( + "Linter exited with code {:?}\n{}", + output.status.code(), + combined + ))) + } + } +} + +#[cfg(test)] +#[path = "mod_tests.rs"] +mod tests; diff --git a/crates/tinyagents-harness-tools/src/fs/filesystem/run_linter/mod_tests.rs b/crates/tinyagents-harness-tools/src/fs/filesystem/run_linter/mod_tests.rs new file mode 100644 index 000000000..d382a40bb --- /dev/null +++ b/crates/tinyagents-harness-tools/src/fs/filesystem/run_linter/mod_tests.rs @@ -0,0 +1,97 @@ +#![allow(clippy::expect_used, clippy::panic, clippy::unwrap_used)] + +use super::*; +use serde_json::json; +use tempfile::TempDir; + +fn make_tool(dir: &TempDir) -> RunLinterTool { + RunLinterTool::new(dir.path().to_path_buf()) +} + +#[test] +fn name_is_correct() { + let tmp = TempDir::new().unwrap(); + assert_eq!(make_tool(&tmp).name(), "run_linter"); +} + +#[test] +fn description_is_non_empty() { + let tmp = TempDir::new().unwrap(); + assert_ne!(make_tool(&tmp).description().len(), 0); +} + +#[test] +fn schema_is_object_type() { + let tmp = TempDir::new().unwrap(); + let schema = make_tool(&tmp).parameters_schema(); + assert_eq!(schema["type"], "object"); +} + +#[test] +fn permission_level_is_execute() { + let tmp = TempDir::new().unwrap(); + assert_eq!(make_tool(&tmp).permission_level(), PermissionLevel::Execute); +} + +#[tokio::test] +async fn auto_returns_error_when_no_project_files() { + let tmp = TempDir::new().unwrap(); + let result = make_tool(&tmp) + .execute(json!({"linter": "auto"})) + .await + .unwrap(); + assert!(result.is_error); + assert!(result.output().contains("Could not detect project type")); +} + +#[tokio::test] +async fn unknown_linter_returns_error() { + let tmp = TempDir::new().unwrap(); + let result = make_tool(&tmp) + .execute(json!({"linter": "rubocop"})) + .await + .unwrap(); + assert!(result.is_error); + assert!(result.output().contains("Unknown linter")); +} + +#[tokio::test] +async fn eslint_rejects_absolute_path() { + let tmp = TempDir::new().unwrap(); + // Create a package.json so linter resolves to eslint + std::fs::write(tmp.path().join("package.json"), "{}").unwrap(); + let result = make_tool(&tmp) + .execute(json!({"linter": "eslint", "path": "/etc/passwd"})) + .await + .unwrap(); + assert!(result.is_error); + assert!(result.output().contains("relative path")); +} + +#[tokio::test] +async fn eslint_rejects_path_traversal() { + let tmp = TempDir::new().unwrap(); + std::fs::write(tmp.path().join("package.json"), "{}").unwrap(); + let result = make_tool(&tmp) + .execute(json!({"linter": "eslint", "path": "../secret"})) + .await + .unwrap(); + assert!(result.is_error); + assert!(result.output().contains("relative path")); +} + +#[tokio::test] +async fn a_workspace_descriptor_overrides_the_configured_directory() { + use crate::fs::filesystem::test_support::WorkspaceContext; + let configured = TempDir::new().unwrap(); + let isolated = TempDir::new().unwrap(); + let context = WorkspaceContext::at(isolated.path()); + // Neither directory has project files; the answer only proves the call + // ran, and the descriptor branch is what resolved the directory. + let result = make_tool(&configured) + .execute_with_context(json!({}), ToolCallOptions::default(), Some(&context)) + .await + .unwrap(); + assert!(result.is_error); + assert!(result.output().contains("Could not detect project type")); +} diff --git a/crates/tinyagents-harness-tools/src/fs/filesystem/run_tests/mod.rs b/crates/tinyagents-harness-tools/src/fs/filesystem/run_tests/mod.rs new file mode 100644 index 000000000..9f4164b0b --- /dev/null +++ b/crates/tinyagents-harness-tools/src/fs/filesystem/run_tests/mod.rs @@ -0,0 +1,200 @@ +//! Tool: `run_tests` — run test suites for the Critic archetype. + +use async_trait::async_trait; +use serde_json::json; +use std::path::PathBuf; +use tinytools::ToolRunContext; +use tinytools::{PermissionLevel, Tool, ToolCallOptions, ToolResult}; + +/// Runs test suites (cargo test, vitest) and returns pass/fail with output. +#[derive(Debug)] +pub struct RunTestsTool { + workspace_dir: PathBuf, +} + +impl RunTestsTool { + #[must_use] + /// Create the tool. + pub fn new(workspace_dir: PathBuf) -> Self { + Self { workspace_dir } + } + + fn workspace_dir_for_context(&self, context: Option<&dyn ToolRunContext>) -> PathBuf { + if let Some(workspace) = context.and_then(|ctx| ctx.workspace()) { + tracing::debug!( + workspace_root = %workspace.root.display(), + policy_id = %workspace.policy_id, + "[run_tests] using TinyAgents workspace descriptor as workspace dir" + ); + return workspace.root.clone(); + } + self.workspace_dir.clone() + } +} + +#[async_trait] +impl Tool for RunTestsTool { + fn exposure(&self) -> tinytools::ToolExposure { + tinytools::ToolExposure::Deferred + } + + fn name(&self) -> &'static str { + "run_tests" + } + + fn description(&self) -> &'static str { + "Run the project test suite. Supports 'cargo_test' for Rust and 'vitest' for \ + TypeScript/JavaScript. Returns pass/fail results with output." + } + + fn parameters_schema(&self) -> serde_json::Value { + json!({ + "type": "object", + "properties": { + "runner": { + "type": "string", + "enum": ["cargo_test", "vitest", "auto"], + "description": "Which test runner to use. 'auto' detects from project files.", + "default": "auto" + }, + "filter": { + "type": "string", + "description": "Filter to run specific tests (e.g. test name or module)." + }, + "timeout_secs": { + "type": "integer", + "description": "Timeout in seconds (default: 120).", + "default": 120 + } + } + }) + } + + fn permission_level(&self) -> PermissionLevel { + PermissionLevel::Execute + } + + async fn execute(&self, args: serde_json::Value) -> anyhow::Result { + self.execute_with_context(args, ToolCallOptions::default(), None) + .await + } + + async fn execute_with_context( + &self, + args: serde_json::Value, + _options: ToolCallOptions, + context: Option<&dyn ToolRunContext>, + ) -> anyhow::Result { + let workspace_dir = self.workspace_dir_for_context(context); + let runner = args + .get("runner") + .and_then(|v| v.as_str()) + .unwrap_or("auto"); + let filter = args.get("filter").and_then(|v| v.as_str()); + let timeout_secs = args + .get("timeout_secs") + .and_then(serde_json::Value::as_u64) + .unwrap_or(120); + + let runner = if runner == "auto" { + if workspace_dir.join("Cargo.toml").exists() { + "cargo_test" + } else if workspace_dir.join("package.json").exists() { + "vitest" + } else { + return Ok(ToolResult::error( + "Could not detect project type for testing.", + )); + } + } else { + runner + }; + + let mut cmd = match runner { + "cargo_test" => { + let mut c = tokio::process::Command::new("cargo"); + c.arg("test"); + if let Some(f) = filter { + c.arg(f); + } + c + } + "vitest" => { + let mut c = tokio::process::Command::new("npx"); + c.args(["vitest", "run"]); + if let Some(f) = filter { + c.arg(f); + } + c + } + other => { + return Ok(ToolResult::error(format!("Unknown test runner: {other}"))); + } + }; + + tracing::debug!( + workspace = %workspace_dir.display(), + "[run_tests] runner={runner}, filter={filter:?}, timeout={timeout_secs}s" + ); + + cmd.current_dir(&workspace_dir); + cmd.kill_on_drop(true); + + let output = + match tokio::time::timeout(std::time::Duration::from_secs(timeout_secs), cmd.output()) + .await + { + Ok(Ok(output)) => output, + Ok(Err(e)) => { + return Ok(ToolResult::error(format!( + "failed to spawn test runner: {e}" + ))); + } + Err(_) => { + return Ok(ToolResult::error(format!( + "test execution timed out after {timeout_secs}s" + ))); + } + }; + + let stdout = String::from_utf8_lossy(&output.stdout); + let stderr = String::from_utf8_lossy(&output.stderr); + let combined = format!("{stdout}\n{stderr}"); + + // Truncate on a safe UTF-8 char boundary. + let truncated = if combined.len() > 8000 { + let safe_end = combined + .char_indices() + .take_while(|(i, _)| *i <= 8000) + .last() + .map_or(0, |(i, c)| i + c.len_utf8()); + format!( + "{}...\n[truncated, {} total chars]", + &combined[..safe_end], + combined.len() + ) + } else { + combined + }; + + tracing::debug!( + "[run_tests] exit_code={:?}, output_len={}", + output.status.code(), + truncated.len() + ); + + if output.status.success() { + Ok(ToolResult::success(truncated)) + } else { + Ok(ToolResult::error(format!( + "Tests exited with code {:?}\n\n{}", + output.status.code(), + truncated + ))) + } + } +} + +#[cfg(test)] +#[path = "mod_tests.rs"] +mod tests; diff --git a/crates/tinyagents-harness-tools/src/fs/filesystem/run_tests/mod_tests.rs b/crates/tinyagents-harness-tools/src/fs/filesystem/run_tests/mod_tests.rs new file mode 100644 index 000000000..edfcdaca1 --- /dev/null +++ b/crates/tinyagents-harness-tools/src/fs/filesystem/run_tests/mod_tests.rs @@ -0,0 +1,164 @@ +//! Behavior tests for the `run_tests` tool. + +#![allow(clippy::expect_used, clippy::panic, clippy::unwrap_used)] + +use super::*; +use crate::fs::filesystem::test_support::WorkspaceContext; +use serde_json::json; +use tempfile::TempDir; + +fn make_tool(dir: &TempDir) -> RunTestsTool { + RunTestsTool::new(dir.path().to_path_buf()) +} + +/// A throwaway cargo package whose single test body is `test_body`. +fn cargo_project(test_body: &str) -> TempDir { + let tmp = TempDir::new().unwrap(); + std::fs::write( + tmp.path().join("Cargo.toml"), + "[package]\nname = \"fixture\"\nversion = \"0.0.0\"\nedition = \"2021\"\n\n[workspace]\n", + ) + .unwrap(); + std::fs::create_dir(tmp.path().join("src")).unwrap(); + std::fs::write( + tmp.path().join("src/lib.rs"), + format!("#[test]\nfn fixture_test() {{\n{test_body}\n}}\n"), + ) + .unwrap(); + tmp +} + +#[test] +fn name_and_description_are_pinned() { + let tmp = TempDir::new().unwrap(); + let tool = make_tool(&tmp); + assert_eq!(tool.name(), "run_tests"); + assert_eq!( + tool.description(), + "Run the project test suite. Supports 'cargo_test' for Rust and 'vitest' for \ + TypeScript/JavaScript. Returns pass/fail results with output." + ); +} + +#[test] +fn schema_is_pinned() { + let tmp = TempDir::new().unwrap(); + assert_eq!( + make_tool(&tmp).parameters_schema(), + json!({ + "type": "object", + "properties": { + "runner": { + "type": "string", + "enum": ["cargo_test", "vitest", "auto"], + "description": "Which test runner to use. 'auto' detects from project files.", + "default": "auto" + }, + "filter": { + "type": "string", + "description": "Filter to run specific tests (e.g. test name or module)." + }, + "timeout_secs": { + "type": "integer", + "description": "Timeout in seconds (default: 120).", + "default": 120 + } + } + }) + ); +} + +#[test] +fn permission_level_and_exposure() { + let tmp = TempDir::new().unwrap(); + let tool = make_tool(&tmp); + assert_eq!(tool.permission_level(), PermissionLevel::Execute); + assert_eq!(tool.exposure(), tinytools::ToolExposure::Deferred); +} + +#[tokio::test] +async fn auto_returns_error_when_no_project_files() { + let tmp = TempDir::new().unwrap(); + let result = make_tool(&tmp).execute(json!({})).await.unwrap(); + assert!(result.is_error); + assert!( + result + .output() + .contains("Could not detect project type for testing.") + ); +} + +#[tokio::test] +async fn unknown_runner_returns_error() { + let tmp = TempDir::new().unwrap(); + let result = make_tool(&tmp) + .execute(json!({"runner": "pytest"})) + .await + .unwrap(); + assert!(result.is_error); + assert_eq!(result.output(), "Unknown test runner: pytest"); +} + +#[tokio::test] +async fn auto_detects_a_node_project_as_vitest() { + let tmp = TempDir::new().unwrap(); + std::fs::write(tmp.path().join("package.json"), "{}").unwrap(); + // The runner may or may not be installed; either way the tool must have + // picked `vitest` rather than reporting an undetectable project. + let result = make_tool(&tmp) + .execute(json!({"timeout_secs": 1})) + .await + .unwrap(); + assert!(!result.output().contains("Could not detect project type")); +} + +#[tokio::test] +async fn passing_cargo_project_reports_success() { + let tmp = cargo_project("assert_eq!(1 + 1, 2);"); + let result = make_tool(&tmp) + .execute(json!({"runner": "cargo_test", "filter": "fixture_test"})) + .await + .unwrap(); + assert!(!result.is_error, "{}", result.output()); + assert!(result.output().contains("test result: ok")); +} + +#[tokio::test] +async fn failing_cargo_project_reports_the_exit_code_and_truncates_output() { + let tmp = cargo_project("panic!(\"{}\", \"x\".repeat(20_000));"); + let result = make_tool(&tmp).execute(json!({})).await.unwrap(); + assert!(result.is_error); + let output = result.output(); + assert!(output.starts_with("Tests exited with code Some(101)")); + assert!(output.contains("[truncated, ")); +} + +#[tokio::test] +async fn a_timeout_is_reported() { + let tmp = cargo_project("assert!(true);"); + let result = make_tool(&tmp) + .execute(json!({"runner": "cargo_test", "timeout_secs": 0})) + .await + .unwrap(); + assert!(result.is_error); + assert_eq!(result.output(), "test execution timed out after 0s"); +} + +#[tokio::test] +async fn a_workspace_descriptor_overrides_the_configured_directory() { + let configured = TempDir::new().unwrap(); + let isolated = TempDir::new().unwrap(); + std::fs::write(isolated.path().join("package.json"), "{}").unwrap(); + let context = WorkspaceContext::at(isolated.path()); + // The configured directory has no project files; only the isolated + // workspace does, so detection succeeding proves the override was used. + let result = make_tool(&configured) + .execute_with_context( + json!({"timeout_secs": 1}), + ToolCallOptions::default(), + Some(&context), + ) + .await + .unwrap(); + assert!(!result.output().contains("Could not detect project type")); +} diff --git a/crates/tinyagents-harness-tools/src/fs/filesystem/test_support.rs b/crates/tinyagents-harness-tools/src/fs/filesystem/test_support.rs new file mode 100644 index 000000000..1a3ffb58f --- /dev/null +++ b/crates/tinyagents-harness-tools/src/fs/filesystem/test_support.rs @@ -0,0 +1,385 @@ +//! A small, self-contained [`FsGate`] for the filesystem tools' unit tests. +//! +//! It models just enough of a host policy to exercise the tools: an autonomy +//! tier, an action budget, an approval flag, and a set of allowed roots that +//! paths must resolve into. It is *not* a copy of any host's policy; the +//! host's own semantics are tested in the host, through its adapter. + +#![allow( + clippy::expect_used, + clippy::unwrap_used, + missing_docs, + unreachable_pub +)] + +use std::path::{Component, Path, PathBuf}; +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; + +use async_trait::async_trait; + +use super::gate::FsGate; + +/// Prefix the fake puts on refusals a real host would mark as policy-blocked. +pub const POLICY_BLOCKED_MARKER: &str = "[policy-blocked]"; + +/// Autonomy tier of the fake gate. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum AutonomyLevel { + /// Nothing that changes state may run. + ReadOnly, + /// Writes ask for approval. + Supervised, + /// Writes run unprompted. + Full, +} + +/// A fake host policy. +#[derive(Debug)] +pub struct TestGate { + autonomy: AutonomyLevel, + action_dir: PathBuf, + roots: Vec, + max_actions: usize, + used: AtomicUsize, +} + +impl TestGate { + /// Supervised gate whose only root is `dir`. + pub fn supervised(dir: PathBuf) -> Arc { + Self::with(dir, AutonomyLevel::Supervised, 1_000_000) + } + + /// Gate over `dir` with an explicit tier and action budget. + pub fn with(dir: PathBuf, autonomy: AutonomyLevel, max_actions: usize) -> Arc { + Arc::new(Self { + autonomy, + roots: vec![dir.clone()], + action_dir: dir, + max_actions, + used: AtomicUsize::new(0), + }) + } + + /// Supervised gate resolving relative paths against `action_dir`, with + /// `extra_roots` also readable and writable. + pub fn split(action_dir: PathBuf, extra_roots: Vec) -> Arc { + let mut roots = vec![action_dir.clone()]; + roots.extend(extra_roots); + Arc::new(Self { + autonomy: AutonomyLevel::Supervised, + roots, + action_dir, + max_actions: 1_000_000, + used: AtomicUsize::new(0), + }) + } + + fn full_path(&self, path: &str) -> PathBuf { + if Path::new(path).is_absolute() { + PathBuf::from(path) + } else { + self.action_dir.join(path) + } + } + + /// Canonicalize the deepest existing ancestor of `full` and test it + /// against the roots, so a symlink out of a root is caught. + fn resolved_inside_roots(&self, full: &Path) -> bool { + let mut existing = full.to_path_buf(); + while !existing.exists() { + match existing.parent() { + Some(parent) => existing = parent.to_path_buf(), + None => return false, + } + } + let Ok(canonical) = existing.canonicalize() else { + return false; + }; + self.roots.iter().any(|root| { + root.canonicalize() + .is_ok_and(|root| canonical.starts_with(root)) + }) + } + + fn blocked(path: &str) -> String { + format!( + "{POLICY_BLOCKED_MARKER} Path not allowed by security policy: {path}. Do not retry \ + this path; use an allowed location (the workspace or a granted folder)." + ) + } +} + +#[async_trait] +impl FsGate for TestGate { + fn can_act(&self) -> bool { + self.autonomy != AutonomyLevel::ReadOnly + } + + fn is_read_only(&self) -> bool { + self.autonomy == AutonomyLevel::ReadOnly + } + + fn is_rate_limited(&self) -> bool { + self.used.load(Ordering::SeqCst) >= self.max_actions + } + + fn record_action(&self) -> bool { + self.used.fetch_add(1, Ordering::SeqCst) < self.max_actions + } + + fn write_needs_approval(&self) -> bool { + self.autonomy == AutonomyLevel::Supervised + } + + fn action_dir(&self) -> &Path { + &self.action_dir + } + + fn is_path_string_allowed(&self, path: &str) -> bool { + if path.contains('\0') { + return false; + } + let candidate = Path::new(path); + if candidate + .components() + .any(|component| component == Component::ParentDir) + { + return false; + } + self.resolved_inside_roots(&self.full_path(path)) + } + + async fn validate_path(&self, path: &str) -> Result { + if !self.is_path_string_allowed(path) { + return Err(Self::blocked(path)); + } + tokio::fs::canonicalize(self.full_path(path)) + .await + .map_err(|e| format!("Failed to resolve path '{path}': {e}")) + } + + async fn validate_parent_path(&self, path: &str) -> Result { + if !self.is_path_string_allowed(path) { + return Err(Self::blocked(path)); + } + let full = self.full_path(path); + let parent = full + .parent() + .ok_or_else(|| format!("Invalid path (no parent): {path}"))?; + let file_name = full + .file_name() + .ok_or_else(|| format!("Invalid path (no filename): {path}"))?; + let mut existing = parent.to_path_buf(); + while !existing.exists() { + match existing.parent() { + Some(next) => existing = next.to_path_buf(), + None => break, + } + } + let canonical = tokio::fs::canonicalize(&existing) + .await + .map_err(|e| format!("Failed to resolve parent of '{path}': {e}"))?; + let suffix = parent.strip_prefix(&existing).unwrap_or(Path::new("")); + Ok(canonical.join(suffix).join(file_name)) + } + + fn scoped_to_workspace(&self, root: &Path) -> Arc { + let mut roots = self.roots.clone(); + roots.push(root.to_path_buf()); + Arc::new(Self { + autonomy: self.autonomy, + action_dir: root.to_path_buf(), + roots, + max_actions: self.max_actions, + used: AtomicUsize::new(self.used.load(Ordering::SeqCst)), + }) + } +} + +/// A run context that carries an isolated workspace, as a harness threads one +/// into a worktree-isolated worker. +#[derive(Debug)] +pub struct WorkspaceContext(pub tinytools::WorkspaceDescriptor); + +impl WorkspaceContext { + /// Context whose workspace is rooted at `root`. + pub fn at(root: &Path) -> Self { + Self( + tinytools::WorkspaceDescriptor::new(root.to_path_buf()).with_policy_id("test-worktree"), + ) + } +} + +impl tinytools::ToolRunContext for WorkspaceContext { + fn workspace(&self) -> Option<&tinytools::WorkspaceDescriptor> { + Some(&self.0) + } +} + +/// A gate whose budget check passes but whose `record_action` refuses, the +/// interleaving where another call spends the last unit between the two. +#[derive(Debug)] +pub struct RacyGate(pub Arc); + +#[async_trait] +impl FsGate for RacyGate { + fn can_act(&self) -> bool { + self.0.can_act() + } + fn is_read_only(&self) -> bool { + self.0.is_read_only() + } + fn is_rate_limited(&self) -> bool { + false + } + fn record_action(&self) -> bool { + false + } + fn write_needs_approval(&self) -> bool { + self.0.write_needs_approval() + } + fn action_dir(&self) -> &Path { + self.0.action_dir() + } + fn is_path_string_allowed(&self, path: &str) -> bool { + self.0.is_path_string_allowed(path) + } + async fn validate_path(&self, path: &str) -> Result { + self.0.validate_path(path).await + } + async fn validate_parent_path(&self, path: &str) -> Result { + self.0.validate_parent_path(path).await + } + fn scoped_to_workspace(&self, root: &Path) -> Arc { + self.0.scoped_to_workspace(root) + } +} + +/// A [`TestGate`] wrapper that can force conditions the plain fake cannot +/// reach: a budget that reports "not limited" yet refuses to record, a path +/// resolver pinned to one answer, or an action directory that does not exist. +#[derive(Debug)] +pub struct WrapGate { + inner: Arc, + no_budget: bool, + fixed_path: Option, + action_dir: Option, + batch_calls: Arc, + wrong_length_batch: bool, +} + +impl WrapGate { + /// Wrap `inner` without overriding anything yet. + pub fn new(inner: Arc) -> Self { + Self { + inner, + no_budget: false, + fixed_path: None, + action_dir: None, + batch_calls: Arc::new(AtomicUsize::new(0)), + wrong_length_batch: false, + } + } + + /// Answer every batch with one verdict too few. + #[must_use] + pub fn wrong_length_batch(mut self) -> Self { + self.wrong_length_batch = true; + self + } + + /// Counter of `paths_string_allowed` calls this gate has received. + pub fn batch_calls(&self) -> Arc { + Arc::clone(&self.batch_calls) + } + + /// `is_rate_limited` stays `false` but `record_action` returns `false`. + #[must_use] + pub fn no_budget(mut self) -> Self { + self.no_budget = true; + self + } + + /// Both validators answer `Ok(path)` regardless of input. + #[must_use] + pub fn fixed_path(mut self, path: PathBuf) -> Self { + self.fixed_path = Some(path); + self + } + + /// Report `dir` as the action directory. + #[must_use] + pub fn action_dir_override(mut self, dir: PathBuf) -> Self { + self.action_dir = Some(dir); + self + } + + /// Erase into the trait object the tools take. + pub fn arc(self) -> Arc { + Arc::new(self) + } +} + +#[async_trait] +impl FsGate for WrapGate { + fn can_act(&self) -> bool { + self.inner.can_act() + } + + fn is_read_only(&self) -> bool { + self.inner.is_read_only() + } + + fn is_rate_limited(&self) -> bool { + !self.no_budget && self.inner.is_rate_limited() + } + + fn record_action(&self) -> bool { + !self.no_budget && self.inner.record_action() + } + + fn write_needs_approval(&self) -> bool { + self.inner.write_needs_approval() + } + + fn action_dir(&self) -> &Path { + self.action_dir + .as_deref() + .unwrap_or_else(|| self.inner.action_dir()) + } + + fn is_path_string_allowed(&self, path: &str) -> bool { + self.inner.is_path_string_allowed(path) + } + + fn paths_string_allowed(&self, paths: &[String]) -> Vec { + self.batch_calls.fetch_add(1, Ordering::SeqCst); + let mut answers: Vec = paths + .iter() + .map(|path| self.inner.is_path_string_allowed(path)) + .collect(); + if self.wrong_length_batch { + answers.pop(); + } + answers + } + + async fn validate_path(&self, path: &str) -> Result { + match &self.fixed_path { + Some(fixed) => Ok(fixed.clone()), + None => self.inner.validate_path(path).await, + } + } + + async fn validate_parent_path(&self, path: &str) -> Result { + match &self.fixed_path { + Some(fixed) => Ok(fixed.clone()), + None => self.inner.validate_parent_path(path).await, + } + } + + fn scoped_to_workspace(&self, root: &Path) -> Arc { + self.inner.scoped_to_workspace(root) + } +} diff --git a/crates/tinyagents-harness-tools/src/fs/filesystem/text.rs b/crates/tinyagents-harness-tools/src/fs/filesystem/text.rs new file mode 100644 index 000000000..ef5018fe7 --- /dev/null +++ b/crates/tinyagents-harness-tools/src/fs/filesystem/text.rs @@ -0,0 +1,24 @@ +//! Small text helpers shared by the filesystem tools. + +/// Truncate `s` to at most `max_bytes` bytes, appending a single-character +/// ellipsis `…` (3 bytes) if truncated. The returned string's total byte +/// length never exceeds `max_bytes`. +pub(super) fn truncate_at_byte_boundary(s: &str, max_bytes: usize) -> String { + if s.len() <= max_bytes { + return s.to_string(); + } + let ellipsis = "…"; + let ellipsis_len = ellipsis.len(); + if max_bytes < ellipsis_len { + return String::new(); + } + let mut end = max_bytes - ellipsis_len; + while end > 0 && !s.is_char_boundary(end) { + end -= 1; + } + format!("{}{}", &s[..end], ellipsis) +} + +#[cfg(test)] +#[path = "text_tests.rs"] +mod tests; diff --git a/crates/tinyagents-harness-tools/src/fs/filesystem/text_tests.rs b/crates/tinyagents-harness-tools/src/fs/filesystem/text_tests.rs new file mode 100644 index 000000000..abd85f2fb --- /dev/null +++ b/crates/tinyagents-harness-tools/src/fs/filesystem/text_tests.rs @@ -0,0 +1,25 @@ +use super::truncate_at_byte_boundary; + +#[test] +fn short_strings_pass_through() { + assert_eq!(truncate_at_byte_boundary("abc", 3), "abc"); +} + +#[test] +fn long_strings_end_in_an_ellipsis_within_budget() { + let out = truncate_at_byte_boundary("abcdefghij", 6); + assert_eq!(out, "abc…"); + assert!(out.len() <= 6); +} + +#[test] +fn never_splits_a_multibyte_character() { + // 'é' is two bytes; a budget landing mid-character backs up. + let out = truncate_at_byte_boundary("ééééé", 6); + assert_eq!(out, "é…"); +} + +#[test] +fn a_budget_smaller_than_the_ellipsis_yields_nothing() { + assert_eq!(truncate_at_byte_boundary("abcdef", 2), ""); +} diff --git a/crates/tinyagents-harness-tools/src/fs/filesystem/update_memory_md/mod.rs b/crates/tinyagents-harness-tools/src/fs/filesystem/update_memory_md/mod.rs new file mode 100644 index 000000000..f0627eaa0 --- /dev/null +++ b/crates/tinyagents-harness-tools/src/fs/filesystem/update_memory_md/mod.rs @@ -0,0 +1,413 @@ +//! Tool: `update_memory_md` — append or update sections in MEMORY.md or SKILL.md. + +use async_trait::async_trait; +use serde_json::json; +use std::collections::HashMap; +use std::path::{Path, PathBuf}; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::{Arc, LazyLock, Mutex}; +use tinytools::ToolRunContext; +use tinytools::{PermissionLevel, Tool, ToolCallOptions, ToolResult}; + +/// Allowed workspace markdown files this tool may modify. +const ALLOWED_FILES: &[&str] = &["MEMORY.md", "SKILL.md"]; + +/// Process-global registry of per-workspace write locks (#4458). +/// +/// `update_memory_md` performs a read-modify-write on `MEMORY.md`/`SKILL.md`. +/// The per-run `MemoryProtocolTracker` has zero cross-run awareness, so two +/// concurrent runs (parallel forks, cron) racing the same workspace file would +/// otherwise clobber each other's append. We serialize every write to a given +/// workspace directory through a shared async mutex, keyed by the canonicalized +/// (falling back to raw) workspace path, so concurrent index updates queue +/// instead of racing. Combined with the temp-file + atomic-rename write below, +/// a killed process can never leave a truncated file. +static WORKSPACE_WRITE_LOCKS: LazyLock>>>> = + LazyLock::new(|| Mutex::new(HashMap::new())); + +/// Monotonic counter for temp-file uniqueness within a process. +static TEMP_FILE_SEQ: AtomicU64 = AtomicU64::new(0); + +/// Return (creating if needed) the shared async write lock for `workspace_dir`. +/// +/// The lock key is the canonicalized workspace path when it resolves (so two +/// spellings of the same directory share one lock), else the raw path. +fn workspace_write_lock(workspace_dir: &Path) -> Arc> { + let key = workspace_dir + .canonicalize() + .unwrap_or_else(|_| workspace_dir.to_path_buf()); + let mut map = WORKSPACE_WRITE_LOCKS + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + Arc::clone( + map.entry(key) + .or_insert_with(|| Arc::new(tokio::sync::Mutex::new(()))), + ) +} + +/// Acquire an **inter-process** advisory write lock for `workspace_dir` (#4458). +/// +/// [`WORKSPACE_WRITE_LOCKS`] only serializes writers inside this OS process, but +/// cron launches work via separate `tokio::process::Command` subprocesses that +/// don't share that mutex — so two cron runs could still clobber the same +/// `MEMORY.md` mid read-modify-write. This takes an `fs2` exclusive `flock` on a +/// sentinel `.memory-write.lock` file in the workspace; the returned `File` +/// holds the lock until it is dropped (end of the write). `flock` acquisition +/// blocks, so it runs on a blocking thread. +async fn acquire_cross_process_write_lock(workspace_dir: &Path) -> anyhow::Result { + let lock_path = workspace_dir.join(".memory-write.lock"); + tokio::task::spawn_blocking(move || { + use fs2::FileExt; + if let Some(parent) = lock_path.parent() { + let _ = std::fs::create_dir_all(parent); + } + // Reject a symlinked lock file: a project-controlled `.memory-write.lock` + // symlink (including a dangling one) could otherwise redirect this + // create/open/lock to a path OUTSIDE the already-containment-checked + // workspace, bypassing the symlink hardening applied to MEMORY.md / + // SKILL.md. If it exists it must be a regular file. + if let Ok(meta) = std::fs::symlink_metadata(&lock_path) + && meta.file_type().is_symlink() + { + return Err(anyhow::anyhow!( + "workspace lock file {lock_path:?} is a symlink; refusing to follow it" + )); + } + let mut opts = std::fs::OpenOptions::new(); + opts.create(true).write(true).truncate(false); + #[cfg(unix)] + { + // O_NOFOLLOW closes the TOCTOU window: if a symlink is swapped in + // after the check above, the open fails (ELOOP) rather than follows. + use std::os::unix::fs::OpenOptionsExt; + opts.custom_flags(libc::O_NOFOLLOW); + } + let file = opts + .open(&lock_path) + .map_err(|e| anyhow::anyhow!("open workspace lock file {lock_path:?}: {e}"))?; + file.lock_exclusive() + .map_err(|e| anyhow::anyhow!("acquire workspace write flock: {e}"))?; + Ok::(file) + }) + .await + .map_err(|e| anyhow::anyhow!("workspace lock task join failed: {e}"))? +} + +/// Atomically replace `path`'s contents with `content`. +/// +/// Writes to a sibling temp file in the same directory (so the rename stays on +/// one filesystem and is atomic) and `rename`s it over the target. A crash +/// mid-write leaves either the old file or the complete new file — never a +/// half-written truncation. Callers MUST hold the per-workspace write lock so +/// the read-modify-write is serialized end-to-end. +async fn atomic_write(path: &Path, file: &str, content: &str) -> anyhow::Result<()> { + let dir = path + .parent() + .ok_or_else(|| anyhow::anyhow!("target path has no parent directory"))?; + let seq = TEMP_FILE_SEQ.fetch_add(1, Ordering::Relaxed); + let tmp_name = format!(".{file}.{}.{seq}.tmp", std::process::id()); + let tmp_path = dir.join(tmp_name); + + tracing::debug!( + tmp = %tmp_path.display(), + target = %path.display(), + bytes = content.len(), + "[update_memory_md] atomic write: staging temp file" + ); + + if let Err(e) = tokio::fs::write(&tmp_path, content).await { + // Clean up a partially-written temp file so a failed stage doesn't + // litter the workspace (CodeRabbit: temp not cleaned on initial write). + let _ = tokio::fs::remove_file(&tmp_path).await; + return Err(anyhow::anyhow!("Failed to stage temp file for {file}: {e}")); + } + + if let Err(e) = tokio::fs::rename(&tmp_path, path).await { + // Best-effort cleanup so a failed rename doesn't litter the workspace. + let _ = tokio::fs::remove_file(&tmp_path).await; + return Err(anyhow::anyhow!("Failed to atomically write {file}: {e}")); + } + + tracing::debug!( + target = %path.display(), + "[update_memory_md] atomic write: rename committed" + ); + Ok(()) +} + +/// Appends or replaces a named section in MEMORY.md or SKILL.md. +/// +/// Supports two actions: +/// - `append`: adds `content` to the end of the file. +/// - `replace_section`: locates the first `## {section_title}` heading and +/// replaces the body (lines until the next `##` heading or EOF) with `content`. +#[derive(Debug)] +pub struct UpdateMemoryMdTool { + workspace_dir: PathBuf, +} + +impl UpdateMemoryMdTool { + #[must_use] + /// Create the tool. + pub fn new(workspace_dir: PathBuf) -> Self { + Self { workspace_dir } + } + + fn workspace_dir_for_context(&self, context: Option<&dyn ToolRunContext>) -> PathBuf { + if let Some(workspace) = context.and_then(|ctx| ctx.workspace()) { + tracing::debug!( + workspace_root = %workspace.root.display(), + policy_id = %workspace.policy_id, + "[update_memory_md] using TinyAgents workspace descriptor as workspace dir" + ); + return workspace.root.clone(); + } + self.workspace_dir.clone() + } +} + +#[async_trait] +impl Tool for UpdateMemoryMdTool { + fn name(&self) -> &'static str { + "update_memory_md" + } + + fn description(&self) -> &'static str { + "Append or update sections in MEMORY.md or SKILL.md workspace files. \ + Use 'append' to add new notes at the end, or 'replace_section' to \ + overwrite the body under a named '## Section' heading." + } + + fn parameters_schema(&self) -> serde_json::Value { + json!({ + "type": "object", + "required": ["file", "action", "content"], + "properties": { + "file": { + "type": "string", + "enum": ["MEMORY.md", "SKILL.md"], + "description": "Which workspace markdown file to modify." + }, + "action": { + "type": "string", + "enum": ["append", "replace_section"], + "description": "'append' adds content at the end; \ + 'replace_section' replaces the body of the named section." + }, + "section_title": { + "type": "string", + "description": "Required for 'replace_section': the heading text (without '## ')." + }, + "content": { + "type": "string", + "description": "The markdown text to write." + } + } + }) + } + + fn permission_level(&self) -> PermissionLevel { + PermissionLevel::Write + } + + async fn execute(&self, args: serde_json::Value) -> anyhow::Result { + self.execute_with_context(args, ToolCallOptions::default(), None) + .await + } + + async fn execute_with_context( + &self, + args: serde_json::Value, + _options: ToolCallOptions, + context: Option<&dyn ToolRunContext>, + ) -> anyhow::Result { + let workspace_dir = self.workspace_dir_for_context(context); + let file = args + .get("file") + .and_then(|v| v.as_str()) + .ok_or_else(|| anyhow::anyhow!("Missing 'file' parameter"))?; + + let action = args + .get("action") + .and_then(|v| v.as_str()) + .ok_or_else(|| anyhow::anyhow!("Missing 'action' parameter"))?; + + let content = args + .get("content") + .and_then(|v| v.as_str()) + .ok_or_else(|| anyhow::anyhow!("Missing 'content' parameter"))?; + + // Guard: only allow MEMORY.md and SKILL.md. + if !ALLOWED_FILES.contains(&file) { + return Ok(ToolResult::error(format!( + "File '{file}' is not allowed. Permitted files: MEMORY.md, SKILL.md" + ))); + } + + let target_path = workspace_dir.join(file); + + // Prevent symlink-based workspace escape. + let workspace_canon = self + .workspace_dir_for_context(context) + .canonicalize() + .map_err(|e| anyhow::anyhow!("Failed to canonicalize workspace: {e}"))?; + // Check parent dir exists and canonicalize to detect symlinks. + let parent = target_path.parent().unwrap_or(&workspace_dir); + let parent_canon = parent + .canonicalize() + .unwrap_or_else(|_| parent.to_path_buf()); + if !parent_canon.starts_with(&workspace_canon) { + return Ok(ToolResult::error(format!( + "File path '{file}' resolves outside workspace" + ))); + } + + tracing::debug!("[update_memory_md] action={action} file={file} path={target_path:?}"); + + // #4458: serialize the whole read-modify-write against concurrent runs + // targeting the same workspace. The guard is held across read + atomic + // write so no interleaving append can be lost. + let lock = workspace_write_lock(&workspace_dir); + let _guard = lock.lock().await; + // Also take a cross-process advisory lock so cron subprocesses (which + // don't share the in-process mutex above) can't clobber the same file + // mid-RMW. Held across read + atomic write; released on drop. + let _file_lock = acquire_cross_process_write_lock(&workspace_dir).await?; + tracing::debug!( + workspace = %workspace_dir.display(), + "[update_memory_md] acquired per-workspace write lock (in-process + cross-process flock)" + ); + + match action { + "append" => self.do_append(&target_path, file, content).await, + "replace_section" => { + let section_title = args + .get("section_title") + .and_then(|v| v.as_str()) + .ok_or_else(|| { + anyhow::anyhow!("'section_title' is required for 'replace_section' action") + })?; + self.do_replace_section(&target_path, file, section_title, content) + .await + } + other => Ok(ToolResult::error(format!( + "Unknown action '{other}'. Use 'append' or 'replace_section'." + ))), + } + } +} + +impl UpdateMemoryMdTool { + /// Append `content` to the end of `path`, creating the file if it does not exist. + async fn do_append( + &self, + path: &std::path::Path, + file: &str, + content: &str, + ) -> anyhow::Result { + // Read existing content (empty string if file not found). + let existing = read_or_empty(path).await?; + + let separator = if existing.is_empty() || existing.ends_with('\n') { + "" + } else { + "\n" + }; + let new_content = format!("{existing}{separator}{content}\n"); + + atomic_write(path, file, &new_content).await?; + + let bytes = new_content.len(); + tracing::info!( + "[update_memory_md] appended {} bytes to {file}", + content.len() + ); + + Ok(ToolResult::success(format!( + "Appended {} bytes to {file} ({bytes} bytes total).", + content.len() + ))) + } + + /// Replace the body of the section headed `## {section_title}` in `path`. + /// + /// If the section is not found it is appended as a new section at the end. + async fn do_replace_section( + &self, + path: &std::path::Path, + file: &str, + section_title: &str, + content: &str, + ) -> anyhow::Result { + let existing = read_or_empty(path).await?; + let heading = format!("## {section_title}"); + + let lines: Vec<&str> = existing.lines().collect(); + let section_start = lines.iter().position(|l| l.trim() == heading.as_str()); + + let new_file_content = if let Some(start_idx) = section_start { + // Find where the next ## heading begins (or end of file). + let body_start = start_idx + 1; + let next_heading = lines[body_start..] + .iter() + .position(|l| l.starts_with("## ")) + .map(|rel| body_start + rel); + + let before: String = lines[..=start_idx].join("\n"); + let after: String = match next_heading { + Some(end_idx) => { + let tail = lines[end_idx..].join("\n"); + format!("\n{tail}") + } + None => String::new(), + }; + + // Ensure content is separated from the heading by a blank line. + let body = if content.trim().is_empty() { + String::new() + } else { + format!("\n{content}") + }; + + format!("{before}{body}{after}\n") + } else { + // Section not found — append it. + tracing::debug!( + "[update_memory_md] section '{section_title}' not found in {file}, appending" + ); + let separator = if existing.is_empty() || existing.ends_with('\n') { + "" + } else { + "\n" + }; + format!("{existing}{separator}{heading}\n{content}\n") + }; + + atomic_write(path, file, &new_file_content).await?; + + tracing::info!( + "[update_memory_md] replaced section '{}' in {file} ({} bytes written)", + section_title, + new_file_content.len() + ); + + Ok(ToolResult::success(format!( + "Section '{}' updated in {file} ({} bytes).", + section_title, + new_file_content.len() + ))) + } +} + +/// Read file to string, returning an empty string when the file does not exist. +async fn read_or_empty(path: &std::path::Path) -> anyhow::Result { + match tokio::fs::read_to_string(path).await { + Ok(s) => Ok(s), + Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(String::new()), + Err(e) => Err(anyhow::anyhow!("Failed to read {}: {e}", path.display())), + } +} + +#[cfg(test)] +#[path = "mod_tests.rs"] +mod tests; diff --git a/crates/tinyagents-harness-tools/src/fs/filesystem/update_memory_md/mod_tests.rs b/crates/tinyagents-harness-tools/src/fs/filesystem/update_memory_md/mod_tests.rs new file mode 100644 index 000000000..a577b85be --- /dev/null +++ b/crates/tinyagents-harness-tools/src/fs/filesystem/update_memory_md/mod_tests.rs @@ -0,0 +1,357 @@ +#![allow( + clippy::expect_used, + clippy::panic, + clippy::unwrap_used, + clippy::items_after_statements, + clippy::single_match_else, + clippy::single_match +)] + +use super::*; + +fn make_tool(dir: &std::path::Path) -> UpdateMemoryMdTool { + UpdateMemoryMdTool::new(dir.to_path_buf()) +} + +#[tokio::test] +async fn append_creates_file_if_missing() { + let dir = tempfile::tempdir().unwrap(); + let tool = make_tool(dir.path()); + let result = tool + .execute(json!({ + "file": "MEMORY.md", + "action": "append", + "content": "first note" + })) + .await + .unwrap(); + assert!(!result.is_error, "{:?}", result.output()); + let text = std::fs::read_to_string(dir.path().join("MEMORY.md")).unwrap(); + assert!(text.contains("first note")); +} + +#[tokio::test] +async fn append_adds_to_existing() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("MEMORY.md"); + std::fs::write(&path, "existing\n").unwrap(); + let tool = make_tool(dir.path()); + tool.execute(json!({ + "file": "MEMORY.md", + "action": "append", + "content": "second note" + })) + .await + .unwrap(); + let text = std::fs::read_to_string(&path).unwrap(); + assert!(text.contains("existing")); + assert!(text.contains("second note")); +} + +#[tokio::test] +async fn replace_section_overwrites_body() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("MEMORY.md"); + std::fs::write(&path, "## Lessons\nold body\n## Other\nkept\n").unwrap(); + let tool = make_tool(dir.path()); + tool.execute(json!({ + "file": "MEMORY.md", + "action": "replace_section", + "section_title": "Lessons", + "content": "new body" + })) + .await + .unwrap(); + let text = std::fs::read_to_string(&path).unwrap(); + assert!(text.contains("new body"), "new body missing: {text}"); + assert!( + !text.contains("old body"), + "old body should be gone: {text}" + ); + assert!(text.contains("## Other"), "other section missing: {text}"); + assert!(text.contains("kept"), "other section body missing: {text}"); +} + +#[tokio::test] +async fn replace_section_appends_when_not_found() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("SKILL.md"); + std::fs::write(&path, "# Header\n").unwrap(); + let tool = make_tool(dir.path()); + tool.execute(json!({ + "file": "SKILL.md", + "action": "replace_section", + "section_title": "New Section", + "content": "brand new" + })) + .await + .unwrap(); + let text = std::fs::read_to_string(&path).unwrap(); + assert!(text.contains("## New Section"), "heading missing: {text}"); + assert!(text.contains("brand new"), "content missing: {text}"); +} + +#[tokio::test] +async fn replace_section_with_empty_content() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("MEMORY.md"); + std::fs::write(&path, "## Notes\nold stuff\n## End\ndone\n").unwrap(); + let tool = make_tool(dir.path()); + tool.execute(json!({ + "file": "MEMORY.md", + "action": "replace_section", + "section_title": "Notes", + "content": "" + })) + .await + .unwrap(); + let text = std::fs::read_to_string(&path).unwrap(); + assert!( + !text.contains("old stuff"), + "old body should be gone: {text}" + ); + assert!(text.contains("## End"), "other section missing: {text}"); +} + +#[tokio::test] +async fn append_to_empty_memory_file() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("MEMORY.md"); + std::fs::write(&path, "").unwrap(); + let tool = make_tool(dir.path()); + let result = tool + .execute(json!({ + "file": "MEMORY.md", + "action": "append", + "content": "first line" + })) + .await + .unwrap(); + assert!(!result.is_error, "unexpected error: {}", result.output()); + let text = std::fs::read_to_string(&path).unwrap(); + assert!(text.contains("first line")); +} + +#[tokio::test] +async fn replace_section_creates_memory_file_if_missing() { + let dir = tempfile::tempdir().unwrap(); + let tool = make_tool(dir.path()); + let result = tool + .execute(json!({ + "file": "MEMORY.md", + "action": "replace_section", + "section_title": "First", + "content": "hello" + })) + .await + .unwrap(); + assert!(!result.is_error, "unexpected error: {}", result.output()); + let text = std::fs::read_to_string(dir.path().join("MEMORY.md")).unwrap(); + assert!(text.contains("## First")); + assert!(text.contains("hello")); +} + +#[tokio::test] +async fn rejects_unknown_action() { + let dir = tempfile::tempdir().unwrap(); + let tool = make_tool(dir.path()); + let result = tool + .execute(json!({ + "file": "MEMORY.md", + "action": "delete_all", + "content": "x" + })) + .await + .unwrap(); + assert!(result.is_error); +} + +#[tokio::test] +async fn replace_section_missing_section_title_errors() { + let dir = tempfile::tempdir().unwrap(); + let tool = make_tool(dir.path()); + let result = tool + .execute(json!({ + "file": "MEMORY.md", + "action": "replace_section", + "content": "x" + })) + .await; + // May return Err or Ok with is_error + match result { + Ok(r) => assert!(r.is_error), + Err(_) => {} // also acceptable + } +} + +#[test] +fn tool_name_and_description() { + let dir = tempfile::tempdir().unwrap(); + let tool = make_tool(dir.path()); + assert_eq!(tool.name(), "update_memory_md"); + assert_ne!(tool.description().len(), 0); +} + +#[test] +fn parameters_schema_has_required_fields() { + let dir = tempfile::tempdir().unwrap(); + let tool = make_tool(dir.path()); + let schema = tool.parameters_schema(); + let required = schema["required"].as_array().unwrap(); + assert!(required.contains(&json!("file"))); + assert!(required.contains(&json!("action"))); +} + +#[tokio::test] +async fn rejects_disallowed_file() { + let dir = tempfile::tempdir().unwrap(); + let tool = make_tool(dir.path()); + let result = tool + .execute(json!({ + "file": "../../etc/passwd", + "action": "append", + "content": "evil" + })) + .await + .unwrap(); + assert!(result.is_error); + assert!(result.output().contains("not allowed")); +} + +#[tokio::test] +async fn missing_parameters_are_errors() { + let dir = tempfile::tempdir().unwrap(); + let tool = make_tool(dir.path()); + for (args, needle) in [ + (json!({"action": "append", "content": "x"}), "'file'"), + (json!({"file": "MEMORY.md", "content": "x"}), "'action'"), + ( + json!({"file": "MEMORY.md", "action": "append"}), + "'content'", + ), + ] { + let err = tool.execute(args).await.unwrap_err(); + assert!(err.to_string().contains(needle), "{err}"); + } +} + +#[tokio::test] +async fn a_missing_workspace_fails_to_canonicalize() { + let dir = tempfile::tempdir().unwrap(); + let tool = make_tool(&dir.path().join("does-not-exist")); + let err = tool + .execute(json!({"file": "MEMORY.md", "action": "append", "content": "x"})) + .await + .unwrap_err(); + assert!(err.to_string().contains("Failed to canonicalize workspace")); +} + +#[tokio::test] +async fn an_unreadable_target_is_reported() { + let dir = tempfile::tempdir().unwrap(); + // A directory named MEMORY.md cannot be read as text (and is not NotFound). + std::fs::create_dir(dir.path().join("MEMORY.md")).unwrap(); + let err = make_tool(dir.path()) + .execute(json!({"file": "MEMORY.md", "action": "append", "content": "x"})) + .await + .unwrap_err(); + assert!(err.to_string().contains("Failed to read"), "{err}"); +} + +#[tokio::test] +async fn atomic_write_reports_staging_and_rename_failures() { + let dir = tempfile::tempdir().unwrap(); + let missing_dir_target = dir.path().join("nope").join("MEMORY.md"); + let err = atomic_write(&missing_dir_target, "MEMORY.md", "x") + .await + .unwrap_err(); + assert!(err.to_string().contains("Failed to stage temp file")); + + // Renaming a file over a directory fails after the temp file was staged. + let dir_target = dir.path().join("MEMORY.md"); + std::fs::create_dir(&dir_target).unwrap(); + let err = atomic_write(&dir_target, "MEMORY.md", "x") + .await + .unwrap_err(); + assert!(err.to_string().contains("Failed to atomically write")); + let leftovers: Vec<_> = std::fs::read_dir(dir.path()) + .unwrap() + .filter_map(Result::ok) + .filter(|e| e.file_name().to_string_lossy().ends_with(".tmp")) + .collect(); + assert!(leftovers.is_empty()); +} + +#[cfg(unix)] +#[tokio::test] +async fn a_symlinked_lock_file_is_refused() { + let dir = tempfile::tempdir().unwrap(); + let outside = tempfile::tempdir().unwrap(); + std::os::unix::fs::symlink( + outside.path().join("target"), + dir.path().join(".memory-write.lock"), + ) + .unwrap(); + let err = make_tool(dir.path()) + .execute(json!({"file": "MEMORY.md", "action": "append", "content": "x"})) + .await + .unwrap_err(); + assert!(err.to_string().contains("is a symlink"), "{err}"); +} + +#[tokio::test] +async fn a_workspace_descriptor_overrides_the_configured_directory() { + use crate::fs::filesystem::test_support::WorkspaceContext; + let configured = tempfile::tempdir().unwrap(); + let isolated = tempfile::tempdir().unwrap(); + let context = WorkspaceContext::at(isolated.path()); + let result = make_tool(configured.path()) + .execute_with_context( + json!({"file": "MEMORY.md", "action": "append", "content": "hi"}), + ToolCallOptions::default(), + Some(&context), + ) + .await + .unwrap(); + assert!(!result.is_error); + assert!(isolated.path().join("MEMORY.md").exists()); + assert!(!configured.path().join("MEMORY.md").exists()); +} + +#[tokio::test] +async fn append_separates_from_content_lacking_a_trailing_newline() { + let dir = tempfile::tempdir().unwrap(); + std::fs::write(dir.path().join("MEMORY.md"), "no newline").unwrap(); + make_tool(dir.path()) + .execute(json!({"file": "MEMORY.md", "action": "append", "content": "next"})) + .await + .unwrap(); + assert_eq!( + std::fs::read_to_string(dir.path().join("MEMORY.md")).unwrap(), + "no newline\nnext\n" + ); +} + +#[tokio::test] +async fn replace_section_before_another_heading_keeps_the_tail() { + let dir = tempfile::tempdir().unwrap(); + std::fs::write(dir.path().join("MEMORY.md"), "## A\nold\n## B\nkeep\n").unwrap(); + make_tool(dir.path()) + .execute(json!({"file": "MEMORY.md", "action": "replace_section", "section_title": "A", "content": "new"})) + .await + .unwrap(); + let text = std::fs::read_to_string(dir.path().join("MEMORY.md")).unwrap(); + assert!(text.contains("## A\nnew\n## B\nkeep"), "{text}"); +} + +#[tokio::test] +async fn replace_section_appends_after_a_file_without_trailing_newline() { + let dir = tempfile::tempdir().unwrap(); + std::fs::write(dir.path().join("MEMORY.md"), "intro").unwrap(); + make_tool(dir.path()) + .execute(json!({"file": "MEMORY.md", "action": "replace_section", "section_title": "New", "content": "body"})) + .await + .unwrap(); + let text = std::fs::read_to_string(dir.path().join("MEMORY.md")).unwrap(); + assert_eq!(text, "intro\n## New\nbody\n"); +} diff --git a/crates/tinyagents-harness-tools/src/fs/filesystem/workspace_state/mod.rs b/crates/tinyagents-harness-tools/src/fs/filesystem/workspace_state/mod.rs new file mode 100644 index 000000000..166ae4994 --- /dev/null +++ b/crates/tinyagents-harness-tools/src/fs/filesystem/workspace_state/mod.rs @@ -0,0 +1,158 @@ +//! Tool: `read_workspace_state` — read-only workspace overview for Orchestrator/Planner. + +use super::git_operations::config::{ + disallowed_config_refusal, first_disallowed_repo_config_key, hardened_git, +}; +use async_trait::async_trait; +use serde_json::json; +use std::path::PathBuf; +use tinytools::{PermissionLevel, Tool, ToolResult}; + +/// Returns a summary of the workspace: git status, file tree, recent commits. +#[derive(Debug)] +pub struct WorkspaceStateTool { + workspace_dir: PathBuf, +} + +impl WorkspaceStateTool { + /// A `read_workspace_state` tool reporting on `workspace_dir`. + #[must_use] + pub fn new(workspace_dir: PathBuf) -> Self { + Self { workspace_dir } + } +} + +#[async_trait] +impl Tool for WorkspaceStateTool { + fn name(&self) -> &'static str { + "read_workspace_state" + } + + fn description(&self) -> &'static str { + "Get a read-only overview of the workspace: git status (modified/untracked files), \ + recent commits, and top-level directory structure. Useful for understanding the \ + current project state before planning tasks." + } + + fn parameters_schema(&self) -> serde_json::Value { + json!({ + "type": "object", + "properties": { + "include_tree": { + "type": "boolean", + "description": "Include top-level directory tree (default: true).", + "default": true + }, + "recent_commits": { + "type": "integer", + "description": "Number of recent commits to show (default: 5).", + "default": 5 + } + } + }) + } + + fn permission_level(&self) -> PermissionLevel { + PermissionLevel::ReadOnly + } + + async fn execute(&self, args: serde_json::Value) -> anyhow::Result { + let include_tree = args + .get("include_tree") + .and_then(serde_json::Value::as_bool) + .unwrap_or(true); + let recent_commits = args + .get("recent_commits") + .and_then(serde_json::Value::as_u64) + .unwrap_or(5) as usize; + + tracing::debug!( + "[workspace_state] dir={}, include_tree={include_tree}, recent_commits={recent_commits}", + self.workspace_dir.display() + ); + + let mut output = String::new(); + let dir = &self.workspace_dir; + + // Git status + output.push_str("## Git Status\n"); + match run_git(dir, &["status", "--porcelain"]).await { + Ok(status) if status.trim().is_empty() => { + output.push_str("Clean working tree.\n"); + } + Ok(status) => { + output.push_str(&status); + } + Err(e) => { + output.push_str(&format!("(not a git repo or error: {e})\n")); + } + } + + // Recent commits + output.push_str(&format!("\n## Recent Commits (last {recent_commits})\n")); + let log_arg = format!("-{recent_commits}"); + match run_git(dir, &["log", &log_arg, "--oneline", "--no-decorate"]).await { + Ok(log) => output.push_str(&log), + Err(e) => output.push_str(&format!("(error: {e})\n")), + } + + // Directory tree (top-level only) + if include_tree { + output.push_str("\n## Directory Tree (top-level)\n"); + match tokio::fs::read_dir(dir).await { + Ok(mut entries) => { + let mut names = Vec::new(); + while let Ok(Some(entry)) = entries.next_entry().await { + let name = entry.file_name().to_string_lossy().to_string(); + if !name.starts_with('.') { + let suffix = if entry.file_type().await.is_ok_and(|ft| ft.is_dir()) { + "/" + } else { + "" + }; + names.push(format!("{name}{suffix}")); + } + } + names.sort(); + for name in &names { + output.push_str(&format!(" {name}\n")); + } + } + Err(e) => output.push_str(&format!("(error reading dir: {e})\n")), + } + } + + tracing::debug!("[workspace_state] output length={}", output.len()); + Ok(ToolResult::success(output)) + } +} + +async fn run_git(dir: &std::path::Path, args: &[&str]) -> anyhow::Result { + if let Some(key) = first_disallowed_repo_config_key(dir).await? { + // The refusal is otherwise only visible folded into the tool's own + // output, which is not greppable when an operator is asking why a + // workspace stopped reporting. Correlation fields: the directory and + // the key that caused it. + tracing::debug!( + "[workspace_state] refusing to run git: dir={}, disallowed_config_key={key}", + dir.display() + ); + anyhow::bail!("{}", disallowed_config_refusal(dir, &key)) + } + + let output = hardened_git(dir).args(args).output().await?; + + if output.status.success() { + Ok(String::from_utf8_lossy(&output.stdout).to_string()) + } else { + anyhow::bail!( + "git {} failed: {}", + args.join(" "), + String::from_utf8_lossy(&output.stderr) + ) + } +} + +#[cfg(test)] +#[path = "mod_tests.rs"] +mod tests; diff --git a/crates/tinyagents-harness-tools/src/fs/filesystem/workspace_state/mod_tests.rs b/crates/tinyagents-harness-tools/src/fs/filesystem/workspace_state/mod_tests.rs new file mode 100644 index 000000000..128632401 --- /dev/null +++ b/crates/tinyagents-harness-tools/src/fs/filesystem/workspace_state/mod_tests.rs @@ -0,0 +1,289 @@ +use super::*; +use crate::fs::filesystem::git_operations::config::normalise_config_key; +use serde_json::json; +use tempfile::TempDir; + +fn make_tool(dir: &TempDir) -> WorkspaceStateTool { + WorkspaceStateTool::new(dir.path().to_path_buf()) +} + +#[test] +fn name_is_correct() { + let tmp = TempDir::new().unwrap(); + assert_eq!(make_tool(&tmp).name(), "read_workspace_state"); +} + +#[test] +fn description_is_non_empty() { + let tmp = TempDir::new().unwrap(); + assert_ne!(make_tool(&tmp).description().len(), 0); +} + +#[test] +fn schema_is_object_type() { + let tmp = TempDir::new().unwrap(); + let schema = make_tool(&tmp).parameters_schema(); + assert_eq!(schema["type"], "object"); +} + +#[test] +fn permission_level_is_read_only() { + let tmp = TempDir::new().unwrap(); + assert_eq!( + make_tool(&tmp).permission_level(), + PermissionLevel::ReadOnly + ); +} + +#[tokio::test] +async fn output_contains_git_status_section() { + let tmp = TempDir::new().unwrap(); + let result = make_tool(&tmp).execute(json!({})).await.unwrap(); + assert!(!result.is_error); + assert!(result.output().contains("Git Status")); +} + +#[tokio::test] +async fn include_tree_false_omits_directory_tree() { + let tmp = TempDir::new().unwrap(); + let result = make_tool(&tmp) + .execute(json!({"include_tree": false})) + .await + .unwrap(); + assert!(!result.is_error); + assert!(!result.output().contains("Directory Tree")); +} + +/// Write a `core.fsmonitor` hook into `dir`'s repository config that +/// creates a marker file when git runs it, and return the marker's path. +/// +/// `git status` invokes `core.fsmonitor`, and the value is read from the +/// repository's own `.git/config` — a file inside the workspace, which is +/// where `file_write` puts things. The hook exits non-zero so git falls +/// back to a normal scan and `status` still succeeds; the point of the test +/// is whether the command ran at all, not what it returned. +fn plant_fsmonitor_hook(dir: &TempDir) -> std::path::PathBuf { + let hook = dir.path().join("hook.sh"); + let marker = dir.path().join("COMMAND_RAN"); + std::fs::write( + &hook, + format!("#!/bin/sh\ntouch {:?}\nexit 1\n", marker.to_string_lossy()), + ) + .unwrap(); + #[cfg(unix)] + { + use std::os::unix::fs::PermissionsExt; + std::fs::set_permissions(&hook, std::fs::Permissions::from_mode(0o755)).unwrap(); + } + // Written with `git config` rather than by appending to the file: + // appending only lands in `[core]` while `[core]` happens to be the + // last section, which is true of a fresh `git init` and is not a + // property worth depending on. + let ok = std::process::Command::new("git") + .args(["config", "core.fsmonitor"]) + .arg(&hook) + .current_dir(dir.path()) + .status() + .unwrap() + .success(); + assert!(ok, "failed to plant the hook in the repository config"); + marker +} + +fn git_init(dir: &TempDir) { + // `.status().unwrap()` only unwraps the *spawn*: a non-zero exit would + // pass silently here and surface later as a confusing assertion about + // status output. Match `plant_fsmonitor_hook`, which already checks. + let ok = std::process::Command::new("git") + .args(["init", "-q", "."]) + .current_dir(dir.path()) + .status() + .unwrap() + .success(); + assert!(ok, "git init failed in the test workspace"); +} + +/// Set a repository config key with `git config`, asserting it took. +fn set_config(dir: &TempDir, key: &str, value: &str) { + let ok = std::process::Command::new("git") + .args(["config", key, value]) + .current_dir(dir.path()) + .status() + .unwrap() + .success(); + assert!(ok, "failed to set {key} in the test workspace"); +} + +/// Issue #459. The workspace this tool reads is the same directory the +/// agent's own file tools write to, so its `.git/config` is attacker- +/// controlled input. `core.fsmonitor` names a command that `git status` +/// executes. +/// +/// Revert either layer of the fix in `run_git` and this test fails by +/// finding the marker — verified, not assumed. +#[cfg(unix)] +#[tokio::test] +async fn repository_config_naming_a_command_does_not_get_to_run_it() { + let tmp = TempDir::new().unwrap(); + git_init(&tmp); + let marker = plant_fsmonitor_hook(&tmp); + + let result = make_tool(&tmp).execute(json!({})).await.unwrap(); + + assert!( + !marker.exists(), + "`git status` executed the command named by the workspace's own \ + repository config — this tool is a code-execution primitive" + ); + assert!( + result.output().contains("fsmonitor"), + "the refusal should name the key that caused it, got: {}", + result.output() + ); +} + +/// The allowlist has to leave an ordinary repository working, or the fix is +/// just a different way of breaking the tool. +#[cfg(unix)] +#[tokio::test] +async fn an_ordinary_repository_still_reports_status_and_log() { + let tmp = TempDir::new().unwrap(); + git_init(&tmp); + std::fs::write(tmp.path().join("tracked.txt"), "hi").unwrap(); + + let result = make_tool(&tmp).execute(json!({})).await.unwrap(); + let out = result.output(); + + assert!( + out.contains("tracked.txt"), + "a plain `git init` workspace must still report status, got: {out}" + ); + assert!( + !out.contains("not on the allowlist"), + "`git init`'s own config must not trip the allowlist, got: {out}" + ); +} + +/// The first draft of the allowlist refused any repository carrying an +/// ordinary setting like `core.autocrlf`, which is most of them on Windows +/// and many elsewhere — the tool would have reported nothing useful for a +/// large class of real workspaces. Raised by `CodeRabbit` on the PR. +#[cfg(unix)] +#[tokio::test] +async fn an_inert_setting_an_ordinary_repository_carries_is_allowed() { + let tmp = TempDir::new().unwrap(); + git_init(&tmp); + set_config(&tmp, "core.autocrlf", "input"); + set_config(&tmp, "gc.auto", "0"); + set_config(&tmp, "remote.origin.prune", "true"); + std::fs::write(tmp.path().join("tracked.txt"), "hi").unwrap(); + + let out = make_tool(&tmp).execute(json!({})).await.unwrap().output(); + + assert!( + out.contains("tracked.txt"), + "an ordinary repository must still report status, got: {out}" + ); + assert!(!out.contains("not on the allowlist"), "got: {out}"); +} + +/// The other half of the same question, and the answer is the opposite one. +/// `filter.lfs.clean` names a program, so an LFS working copy is refused — +/// fail-closed, and **intended** rather than an oversight. Pinned so that +/// widening the allowlist for ergonomics cannot quietly admit it. +#[cfg(unix)] +#[tokio::test] +async fn an_lfs_clone_is_refused_because_its_filter_names_a_program() { + let tmp = TempDir::new().unwrap(); + git_init(&tmp); + // What `git lfs install` writes. `required` is inert and allowed; the + // three programs are not. + set_config(&tmp, "filter.lfs.required", "true"); + set_config(&tmp, "filter.lfs.clean", "git-lfs clean -- %f"); + + let out = make_tool(&tmp).execute(json!({})).await.unwrap().output(); + + assert!( + out.contains("filter.lfs.clean"), + "the refusal must name the key that caused it, got: {out}" + ); +} + +/// `credential.helper` reads like a preference and is command-valued: a +/// value beginning `!` is run as a shell command. `CodeRabbit`'s review +/// listed it among the inert keys to allow; it is not one, and allowlisting +/// it would have reopened the hole this PR closes. +#[cfg(unix)] +#[tokio::test] +async fn credential_helper_is_refused_despite_looking_like_a_preference() { + let tmp = TempDir::new().unwrap(); + git_init(&tmp); + set_config(&tmp, "credential.helper", "!echo pwned"); + + let out = make_tool(&tmp).execute(json!({})).await.unwrap().output(); + + assert!( + out.contains("credential.helper"), + "a command-valued key must be refused however inert it reads, got: {out}" + ); +} + +#[test] +fn a_subsection_is_elided_so_one_entry_covers_every_remote() { + assert_eq!(normalise_config_key("remote.origin.url"), "remote.url"); + assert_eq!(normalise_config_key("remote.a.b.c.url"), "remote.url"); + assert_eq!(normalise_config_key("core.fileMode"), "core.filemode"); + assert_eq!(normalise_config_key("core.fsmonitor"), "core.fsmonitor"); +} + +#[tokio::test] +async fn lists_non_hidden_files_in_tree() { + let tmp = TempDir::new().unwrap(); + std::fs::write(tmp.path().join("readme.txt"), "hi").unwrap(); + std::fs::write(tmp.path().join(".hidden"), "skip").unwrap(); + let result = make_tool(&tmp) + .execute(json!({"include_tree": true, "recent_commits": 0})) + .await + .unwrap(); + assert!(!result.is_error); + let out = result.output(); + assert!(out.contains("readme.txt")); + assert!(!out.contains(".hidden")); +} + +/// The hardening now comes from `git_operations`' policy, which is stricter +/// than the copy this tool used to carry: `core.worktree` redirects the tree +/// git operates on and is not on the allowlist. +#[cfg(unix)] +#[tokio::test] +async fn core_worktree_is_refused_like_any_other_unrecognised_key() { + let tmp = TempDir::new().unwrap(); + git_init(&tmp); + set_config(&tmp, "core.worktree", "/tmp"); + + let out = make_tool(&tmp).execute(json!({})).await.unwrap().output(); + + assert!(out.contains("core.worktree"), "got: {out}"); + assert!(out.contains("not on the allowlist"), "got: {out}"); +} + +/// `git config --local` never shows a worktree-scoped key; the shared check +/// reads the merged view, so a hook path set there is still caught. +#[cfg(unix)] +#[tokio::test] +async fn a_worktree_scoped_command_key_is_refused() { + let tmp = TempDir::new().unwrap(); + git_init(&tmp); + set_config(&tmp, "extensions.worktreeConfig", "true"); + let ok = std::process::Command::new("git") + .args(["config", "--worktree", "core.fsmonitor", "/bin/true"]) + .current_dir(tmp.path()) + .status() + .unwrap() + .success(); + assert!(ok, "failed to set a worktree-scoped key"); + + let out = make_tool(&tmp).execute(json!({})).await.unwrap().output(); + + assert!(out.contains("fsmonitor"), "got: {out}"); +} diff --git a/crates/tinyagents-harness-tools/src/fs/mod.rs b/crates/tinyagents-harness-tools/src/fs/mod.rs new file mode 100644 index 000000000..901582d21 --- /dev/null +++ b/crates/tinyagents-harness-tools/src/fs/mod.rs @@ -0,0 +1,4 @@ +//! Filesystem tools and file-state coordination. + +pub mod file_state; +pub mod filesystem; diff --git a/crates/tinyagents-harness-tools/src/goals/mod.rs b/crates/tinyagents-harness-tools/src/goals/mod.rs index 609890417..2df0ea9f3 100644 --- a/crates/tinyagents-harness-tools/src/goals/mod.rs +++ b/crates/tinyagents-harness-tools/src/goals/mod.rs @@ -11,7 +11,7 @@ //! //! It is the harness-level counterpart of OpenHuman's `thread_goals`, minus the //! app-specific coupling (event bus, RPC envelopes, heartbeat scheduler): the -//! primitive is provider-neutral and drives off the graph runtime. +//! primitive is provider-neutral and works with the harness tools runtime. pub mod budget; mod observer; diff --git a/crates/tinyagents-harness-tools/src/goals/store.rs b/crates/tinyagents-harness-tools/src/goals/store.rs index f0028b3cf..4bd7a60a4 100644 --- a/crates/tinyagents-harness-tools/src/goals/store.rs +++ b/crates/tinyagents-harness-tools/src/goals/store.rs @@ -9,7 +9,7 @@ //! so every mutation runs `load → mutate → put` under a **per-thread async //! mutex** (`thread_lock`) — the process-local analogue of OpenHuman's //! file-rename atomicity. Inside that lock the `goal_id` compare-and-set guard -//! (see [`account_turn_atomic`]) still rejects +//! (see `account_turn_atomic`) still rejects //! stale accounting from a replaced goal. //! //! # Single-process only diff --git a/crates/tinyagents-harness-tools/src/lib.rs b/crates/tinyagents-harness-tools/src/lib.rs index 8961b7772..348ef5d6c 100644 --- a/crates/tinyagents-harness-tools/src/lib.rs +++ b/crates/tinyagents-harness-tools/src/lib.rs @@ -9,26 +9,37 @@ mod ask_clarification; mod canonical_json; pub mod per_call; mod plan_exit; +pub mod run_mode; +pub mod schema; +pub mod scratchpad; +#[cfg(feature = "state-block")] +pub mod state_block; pub mod thread_locks; mod time; mod time_parse; mod wait; +#[cfg(feature = "exec")] +pub mod exec; +#[cfg(feature = "fs")] +pub mod fs; #[cfg(feature = "goals")] pub mod goals; -#[cfg(feature = "todos")] -pub mod scratchpad; +#[cfg(feature = "media")] +pub mod media; +#[cfg(feature = "net")] +pub mod net; +#[cfg(feature = "plan-review")] +pub mod plan_review; #[cfg(feature = "skills")] pub mod skills; -#[cfg(feature = "state-block")] -pub mod state_block; #[cfg(feature = "subagents")] pub mod subagents; #[cfg(feature = "todos")] pub mod todos; pub use ask_clarification::AskClarificationTool; -pub use plan_exit::{PLAN_EXIT_MARKER, PlanExitModeLookup, PlanExitTool}; +pub use plan_exit::{PLAN_EXIT_MARKER, ParkedPlanReviewLookup, PlanExitHost, PlanExitTool}; pub use time::{CurrentTimeTool, ResolveTimeTool, register_time_tools, time_tools}; pub use wait::{WaitLoopTool, WaitTool}; diff --git a/crates/tinyagents-harness-tools/src/media/artifact_filing.rs b/crates/tinyagents-harness-tools/src/media/artifact_filing.rs new file mode 100644 index 000000000..02267655d --- /dev/null +++ b/crates/tinyagents-harness-tools/src/media/artifact_filing.rs @@ -0,0 +1,33 @@ +use async_trait::async_trait; +use std::path::PathBuf; +use tinytools::ToolRunContext; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum MediaArtifactKind { + Image, + Video, +} +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct GeneratedArtifact { + pub path: PathBuf, + pub artifact_kind: MediaArtifactKind, + pub kind: String, + pub media_type: Option, + pub bytes: Option, + pub title: String, + pub tool_call_id: Option, +} +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct FiledArtifact { + pub id: String, + pub path: String, +} + +#[async_trait] +pub trait ArtifactFiling: Send + Sync { + async fn file( + &self, + generated: GeneratedArtifact, + context: Option<&dyn ToolRunContext>, + ) -> Result; +} diff --git a/crates/tinyagents-harness-tools/src/media/artifact_tool.rs b/crates/tinyagents-harness-tools/src/media/artifact_tool.rs new file mode 100644 index 000000000..b44804012 --- /dev/null +++ b/crates/tinyagents-harness-tools/src/media/artifact_tool.rs @@ -0,0 +1,293 @@ +use super::{ArtifactFiling, GeneratedArtifact, MediaArtifactKind}; +use async_trait::async_trait; +use serde_json::{Value, json}; +use std::time::Duration; +use std::{path::PathBuf, sync::Arc}; +use tinytools::{ + PermissionLevel, Tool, ToolCallOptions, ToolCategory, ToolContent, ToolPolicy, ToolProgress, + ToolResult, ToolRunContext, ToolScope, ToolTimeout, +}; + +#[derive(Debug, Clone, PartialEq)] +pub struct MediaJobProgress { + pub id: String, + pub state: String, + pub outputs: usize, + pub output_indices: Vec, + pub cost_usd: Option, + pub error: Option, +} +pub type MediaProgressFn = Arc; +pub type MediaInnerFactory = + Arc Fn(MediaProgressFn, Option<&'a dyn ToolRunContext>) -> T + Send + Sync>; +/// Host-supplied extraction of its provider-assigned call identifier. +pub type ToolCallIdExtractor = + Arc Fn(Option<&'a dyn ToolRunContext>) -> Option + Send + Sync>; + +fn canonical_tool_call_id(context: Option<&dyn ToolRunContext>) -> Option { + context + .and_then(ToolRunContext::host_extension) + .and_then(|extension| { + extension.downcast_ref::() + }) + .map(|context| context.call_id.to_string()) +} + +fn progress_label(kind: &MediaArtifactKind) -> &'static str { + match kind { + MediaArtifactKind::Image => "Generating image", + MediaArtifactKind::Video => "Generating video", + } +} + +fn heartbeat_message(label: &str, elapsed: Duration, job_state: Option<&str>) -> String { + let seconds = elapsed.as_secs(); + let elapsed = if seconds >= 60 { + format!("{}m {:02}s", seconds / 60, seconds % 60) + } else { + format!("{seconds}s") + }; + match job_state { + Some(state) => format!("{label} — {elapsed} elapsed (job {state})"), + None => format!("{label} — {elapsed} elapsed"), + } +} + +pub struct MediaArtifactTool { + inner: T, + kind: MediaArtifactKind, + per_call: Option>, + call_id: Option, + filing: Arc, +} +impl MediaArtifactTool { + #[must_use] + pub fn new(inner: T, kind: MediaArtifactKind, filing: Arc) -> Self { + Self { + inner, + kind, + per_call: None, + call_id: None, + filing, + } + } + #[must_use] + pub fn with_per_call_inner(mut self, factory: MediaInnerFactory) -> Self { + self.per_call = Some(factory); + self + } + /// Let a host extract its call id from its own erased context extension. + #[must_use] + pub fn with_tool_call_id(mut self, extract: ToolCallIdExtractor) -> Self { + self.call_id = Some(extract); + self + } + fn title(args: &Value, index: usize, total: usize) -> String { + let base = args + .get("prompt") + .and_then(Value::as_str) + .map(str::trim) + .filter(|s| !s.is_empty()) + .map(|s| s.chars().take(60).collect()) + .unwrap_or_else(|| "Generated media".into()); + if total > 1 { + format!("{base} ({}/{total})", index + 1) + } else { + base + } + } + + async fn run_with_heartbeat>>( + &self, + context: Option<&dyn ToolRunContext>, + status: Arc>>, + run: F, + ) -> anyhow::Result { + let Some(context) = context else { + return run.await; + }; + let label = progress_label(&self.kind); + let start = tokio::time::Instant::now(); + let mut run = Box::pin(run); + let mut ticker = tokio::time::interval_at( + tokio::time::Instant::now() + Duration::from_secs(15), + Duration::from_secs(15), + ); + loop { + tokio::select! { + result=&mut run => return result, + _=ticker.tick()=>{ + let state=status.lock().unwrap_or_else(std::sync::PoisonError::into_inner).clone(); + let message=heartbeat_message(label,start.elapsed(),state.as_ref().map(|state|state.state.as_str())); + context.report_progress(ToolProgress::message(message)); + } + } + } + } + async fn file_entries( + &self, + mut result: ToolResult, + args: &Value, + context: Option<&dyn ToolRunContext>, + call_id: Option<&str>, + ) -> ToolResult { + let Some(data) = result.content.iter_mut().find_map(|block| { + if let ToolContent::Json { data } = block { + Some(data) + } else { + None + } + }) else { + return result; + }; + let Some(entries) = data.get_mut("artifacts").and_then(Value::as_array_mut) else { + return result; + }; + let total = entries.len(); + let mut replacements = Vec::new(); + for (index, entry) in entries.iter_mut().enumerate() { + let Some(path) = entry.get("path").and_then(Value::as_str).map(PathBuf::from) else { + continue; + }; + let generated = GeneratedArtifact { + path: path.clone(), + artifact_kind: self.kind.clone(), + kind: entry + .get("type") + .or_else(|| entry.get("kind")) + .and_then(Value::as_str) + .unwrap_or("") + .to_owned(), + media_type: entry + .get("media_type") + .or_else(|| entry.get("format")) + .and_then(Value::as_str) + .map(ToOwned::to_owned), + bytes: entry.get("bytes").and_then(Value::as_u64), + title: Self::title(args, index, total), + tool_call_id: call_id.map(ToOwned::to_owned), + }; + match self.filing.file(generated, context).await { + Ok(filed) => { + if let Some(obj) = entry.as_object_mut() { + obj.insert("path".into(), json!(filed.path)); + obj.insert("artifact_id".into(), json!(filed.id)); + } + replacements.push((path.to_string_lossy().into_owned(), filed.path)); + } + Err(error) => { + if let Some(obj) = entry.as_object_mut() { + obj.insert("artifact_error".into(), json!(error)); + } + } + } + } + if let Some(markdown) = result.markdown_formatted.as_mut() { + for (source, destination) in replacements { + *markdown = markdown.replace(&source, &destination); + } + } + result + } +} + +#[async_trait] +impl Tool for MediaArtifactTool { + fn name(&self) -> &str { + self.inner.name() + } + fn description(&self) -> &str { + self.inner.description() + } + fn parameters_schema(&self) -> Value { + self.inner.parameters_schema() + } + fn policy(&self) -> ToolPolicy { + self.inner.policy() + } + fn permission_level(&self) -> PermissionLevel { + self.inner.permission_level() + } + fn permission_level_with_args(&self, a: &Value) -> PermissionLevel { + self.inner.permission_level_with_args(a) + } + fn scope(&self) -> ToolScope { + self.inner.scope() + } + fn category(&self) -> ToolCategory { + self.inner.category() + } + fn external_effect(&self) -> bool { + self.inner.external_effect() + } + fn external_effect_with_args(&self, a: &Value) -> bool { + self.inner.external_effect_with_args(a) + } + fn timeout_policy(&self, a: &Value) -> ToolTimeout { + self.inner.timeout_policy(a) + } + fn supports_markdown(&self) -> bool { + self.inner.supports_markdown() + } + fn max_result_size_chars(&self) -> Option { + self.inner.max_result_size_chars() + } + async fn execute(&self, args: Value) -> anyhow::Result { + self.execute_with_context(args, ToolCallOptions::default(), None) + .await + } + async fn execute_with_options( + &self, + args: Value, + options: ToolCallOptions, + ) -> anyhow::Result { + self.execute_with_context(args, options, None).await + } + async fn execute_with_context( + &self, + args: Value, + options: ToolCallOptions, + context: Option<&dyn ToolRunContext>, + ) -> anyhow::Result { + let status = Arc::new(std::sync::Mutex::new(None)); + let callback_status = Arc::clone(&status); + let progress: MediaProgressFn = Arc::new(move |update| { + *callback_status + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) = Some(update.clone()); + }); + let inner = self + .per_call + .as_ref() + .map(|factory| factory(progress, context)); + let run = async { + match inner { + Some(tool) => { + tool.execute_with_context(args.clone(), options, context) + .await + } + None => { + self.inner + .execute_with_context(args.clone(), options, context) + .await + } + } + }; + let result = self.run_with_heartbeat(context, status, run).await?; + if result.is_error { + return Ok(result); + } + let call_id = self + .call_id + .as_ref() + .and_then(|extract| extract(context)) + .or_else(|| canonical_tool_call_id(context)); + Ok(self + .file_entries(result, &args, context, call_id.as_deref()) + .await) + } +} + +#[cfg(test)] +#[path = "artifact_tool_tests.rs"] +mod tests; diff --git a/crates/tinyagents-harness-tools/src/media/artifact_tool_tests.rs b/crates/tinyagents-harness-tools/src/media/artifact_tool_tests.rs new file mode 100644 index 000000000..4600ca825 --- /dev/null +++ b/crates/tinyagents-harness-tools/src/media/artifact_tool_tests.rs @@ -0,0 +1,709 @@ +use super::*; +use crate::media::FiledArtifact; +use async_trait::async_trait; +use serde_json::{Value, json}; +use std::{ + path::PathBuf, + sync::{Arc, Mutex}, +}; +use tinytools::{Tool, ToolCallOptions, ToolContent, ToolResult, ToolRunContext}; +use tinytools::{ToolAccess, ToolPolicy, ToolScope, ToolSideEffects, WorkspaceAccess}; + +struct Inner; +#[async_trait] +impl Tool for Inner { + fn name(&self) -> &str { + "inner" + } + fn description(&self) -> &str { + "stable" + } + fn parameters_schema(&self) -> Value { + json!({"type":"object"}) + } + fn max_result_size_chars(&self) -> Option { + Some(1234) + } + async fn execute(&self, _: Value) -> anyhow::Result { + Ok(ToolResult::success_with_markdown( + json!({"artifacts":[{"type":"image","path":"blob:key","bytes":4}]}), + "saved blob:key", + )) + } +} +#[derive(Default)] +struct Filer(Mutex>); +#[async_trait] +impl ArtifactFiling for Filer { + async fn file( + &self, + g: GeneratedArtifact, + _: Option<&dyn ToolRunContext>, + ) -> Result { + self.0.lock().unwrap().push(g); + Ok(FiledArtifact { + id: "id1".into(), + path: "/files/image.png".into(), + }) + } +} + +struct CallIdentity(&'static str); +struct Context { + identity: CallIdentity, + progress: Mutex>, +} +impl ToolRunContext for Context { + fn host_extension(&self) -> Option<&(dyn std::any::Any + Send + Sync)> { + Some(&self.identity) + } + fn report_progress(&self, update: tinytools::ToolProgress) { + self.progress + .lock() + .unwrap() + .push(update.message.unwrap_or_default()); + } +} + +struct RoutingFiler; +#[async_trait] +impl ArtifactFiling for RoutingFiler { + async fn file( + &self, + generated: GeneratedArtifact, + context: Option<&dyn ToolRunContext>, + ) -> Result { + let identity = context + .and_then(|context| context.host_extension()) + .and_then(|extension| extension.downcast_ref::()) + .ok_or("missing call identity")?; + assert_eq!(generated.tool_call_id, Some(identity.0.to_owned())); + let filename = generated + .path + .file_name() + .and_then(|name| name.to_str()) + .ok_or("missing generated filename")?; + let filename = filename.strip_prefix("blob:").unwrap_or(filename); + Ok(FiledArtifact { + id: format!("{}-{filename}", identity.0), + path: format!("/files/{}-{filename}.png", identity.0), + }) + } +} + +struct OneFails; +#[async_trait] +impl ArtifactFiling for OneFails { + async fn file( + &self, + generated: GeneratedArtifact, + _: Option<&dyn ToolRunContext>, + ) -> Result { + if generated.path == std::path::Path::new("blob:bad") { + return Err("filing failed".into()); + } + Ok(FiledArtifact { + id: "good-id".into(), + path: "/files/good.png".into(), + }) + } +} + +struct MultiInner; +#[async_trait] +impl Tool for MultiInner { + fn name(&self) -> &str { + "multi" + } + fn description(&self) -> &str { + "multi description" + } + fn parameters_schema(&self) -> Value { + json!({"type":"object","properties":{"prompt":{"type":"string"}}}) + } + fn category(&self) -> ToolCategory { + ToolCategory::Workflow + } + fn external_effect(&self) -> bool { + true + } + fn permission_level(&self) -> PermissionLevel { + PermissionLevel::Dangerous + } + fn timeout_policy(&self, _: &Value) -> ToolTimeout { + ToolTimeout::Unbounded + } + async fn execute(&self, _: Value) -> anyhow::Result { + Ok(ToolResult::success_with_markdown( + json!({"artifacts":[ + {"type":"image","path":"blob:good","bytes":3}, + {"type":"image","path":"blob:bad","bytes":4}, + {"type":"image","path":"blob:last","bytes":5} + ],"other":"retained"}), + "good blob:good; bad blob:bad; last blob:last", + )) + } +} + +#[derive(Clone, Copy)] +struct MetadataInner { + temporary: bool, +} +impl MetadataInner { + fn text(&self, retained: &'static str, temporary: &'static str) -> &'static str { + if self.temporary { temporary } else { retained } + } +} +#[async_trait] +impl Tool for MetadataInner { + fn name(&self) -> &str { + self.text("retained_media", "temporary_media") + } + fn description(&self) -> &str { + self.text("retained description", "temporary description") + } + fn parameters_schema(&self) -> Value { + json!({"type":"object","properties":{"source":{"const":self.text("retained","temporary")}}}) + } + fn policy(&self) -> ToolPolicy { + ToolPolicy::classified() + .with_side_effects(ToolSideEffects { + network: !self.temporary, + writes_files: self.temporary, + ..ToolSideEffects::default() + }) + .with_access(ToolAccess { + workspace: WorkspaceAccess::Scoped, + trusted_roots: vec![self.text("/retained", "/temporary").into()], + approval_required: !self.temporary, + ..ToolAccess::default() + }) + } + fn permission_level(&self) -> PermissionLevel { + if self.temporary { + PermissionLevel::Write + } else { + PermissionLevel::ReadOnly + } + } + fn permission_level_with_args(&self, args: &Value) -> PermissionLevel { + if args["danger"] == true { + PermissionLevel::Dangerous + } else { + self.permission_level() + } + } + fn scope(&self) -> ToolScope { + if self.temporary { + ToolScope::AgentOnly + } else { + ToolScope::CliRpcOnly + } + } + fn category(&self) -> ToolCategory { + if self.temporary { + ToolCategory::System + } else { + ToolCategory::Workflow + } + } + fn external_effect(&self) -> bool { + self.temporary + } + fn external_effect_with_args(&self, args: &Value) -> bool { + args["send"] == true + } + fn max_result_size_chars(&self) -> Option { + Some(if self.temporary { 22 } else { 11 }) + } + fn timeout_policy(&self, _: &Value) -> ToolTimeout { + ToolTimeout::Millis(if self.temporary { 22_000 } else { 11_000 }) + } + fn supports_markdown(&self) -> bool { + !self.temporary + } + async fn execute(&self, _: Value) -> anyhow::Result { + Ok(ToolResult::success_with_markdown( + json!({"artifacts":[ + {"type":"image","path":"blob:first"}, + {"type":"video","path":"blob:second"} + ]}), + "blob:first then blob:second", + )) + } +} + +async fn assert_retained_metadata(factory: bool) { + let retained = MetadataInner { temporary: false }; + let filer = Arc::new(Filer::default()); + let mut tool = MediaArtifactTool::new(retained, MediaArtifactKind::Image, filer.clone()); + if factory { + tool = tool.with_per_call_inner(Arc::new(|_, _| MetadataInner { temporary: true })); + } + let assert_metadata = |tool: &MediaArtifactTool| { + assert_eq!(tool.name(), "retained_media"); + assert_eq!(tool.description(), "retained description"); + assert_eq!( + tool.parameters_schema(), + json!({"type":"object","properties":{"source":{"const":"retained"}}}) + ); + assert_eq!(tool.policy(), retained.policy()); + assert_eq!(tool.scope(), ToolScope::CliRpcOnly); + assert_eq!(tool.category(), ToolCategory::Workflow); + assert_eq!(tool.permission_level(), PermissionLevel::ReadOnly); + assert_eq!( + tool.permission_level_with_args(&json!({"danger":true})), + PermissionLevel::Dangerous + ); + assert!(!tool.external_effect()); + assert!(tool.external_effect_with_args(&json!({"send":true}))); + assert_eq!(tool.max_result_size_chars(), Some(11)); + assert_eq!(tool.timeout_policy(&json!({})), ToolTimeout::Millis(11_000)); + assert!(tool.supports_markdown()); + }; + assert_metadata(&tool); + let long_prompt = "x".repeat(70); + tool.execute(json!({"prompt":long_prompt})).await.unwrap(); + assert_metadata(&tool); + let filed = filer.0.lock().unwrap(); + assert_eq!(filed.len(), 2); + assert_eq!(filed[0].title, format!("{} (1/2)", "x".repeat(60))); + assert_eq!(filed[1].title, format!("{} (2/2)", "x".repeat(60))); + assert_eq!(filed[0].artifact_kind, MediaArtifactKind::Image); + assert_eq!(filed[1].artifact_kind, MediaArtifactKind::Image); +} + +#[tokio::test] +async fn retained_metadata_survives_execution_with_and_without_per_call_factory() { + assert_retained_metadata(false).await; + assert_retained_metadata(true).await; +} + +struct ConcurrentInner { + identity: &'static str, + progress: MediaProgressFn, + barrier: Arc, +} +#[async_trait] +impl Tool for ConcurrentInner { + fn name(&self) -> &str { + "concurrent" + } + fn description(&self) -> &str { + "per-call concurrent tool" + } + fn parameters_schema(&self) -> Value { + json!({"type":"object"}) + } + async fn execute(&self, _: Value) -> anyhow::Result { + (self.progress)(&MediaJobProgress { + id: self.identity.into(), + state: format!("{}-job", self.identity), + outputs: 2, + output_indices: vec![0, 1], + cost_usd: None, + error: None, + }); + self.barrier.wait().await; + tokio::time::sleep(Duration::from_secs(20)).await; + Ok(ToolResult::success_with_markdown( + json!({"artifacts":[ + {"type":"image","path":format!("blob:{}-one", self.identity)}, + {"type":"image","path":format!("blob:{}-two", self.identity)} + ]}), + format!("blob:{}-one and blob:{}-two", self.identity, self.identity), + )) + } +} + +struct ProgressInner(MediaProgressFn); +#[async_trait] +impl Tool for ProgressInner { + fn name(&self) -> &str { + "progress" + } + fn description(&self) -> &str { + "progress" + } + fn parameters_schema(&self) -> Value { + json!({"type":"object"}) + } + async fn execute(&self, _: Value) -> anyhow::Result { + (self.0)(&MediaJobProgress { + id: "job-1".into(), + state: "in_progress".into(), + outputs: 0, + output_indices: Vec::new(), + cost_usd: None, + error: None, + }); + tokio::time::sleep(Duration::from_secs(20)).await; + Ok(ToolResult::success("done")) + } +} + +struct ProgressArtifactInner(MediaProgressFn); +#[async_trait] +impl Tool for ProgressArtifactInner { + fn name(&self) -> &str { + "progress_artifact" + } + fn description(&self) -> &str { + "progress artifact" + } + fn parameters_schema(&self) -> Value { + json!({"type":"object"}) + } + async fn execute(&self, _: Value) -> anyhow::Result { + (self.0)(&MediaJobProgress { + id: "job-kind".into(), + state: "running".into(), + outputs: 1, + output_indices: vec![0], + cost_usd: None, + error: None, + }); + tokio::time::sleep(Duration::from_secs(20)).await; + Ok(ToolResult::success_with_markdown( + json!({"artifacts":[{"type":"media","path":"blob:kind"}]}), + "blob:kind", + )) + } +} + +struct ErrorInner; +#[async_trait] +impl Tool for ErrorInner { + fn name(&self) -> &str { + "error" + } + fn description(&self) -> &str { + "error" + } + fn parameters_schema(&self) -> Value { + json!({"type":"object"}) + } + async fn execute(&self, _: Value) -> anyhow::Result { + Ok(ToolResult::error("provider failed")) + } +} + +struct UnknownShapeInner; +#[async_trait] +impl Tool for UnknownShapeInner { + fn name(&self) -> &str { + "unknown_shape" + } + fn description(&self) -> &str { + "unknown" + } + fn parameters_schema(&self) -> Value { + json!({"type":"object"}) + } + async fn execute(&self, _: Value) -> anyhow::Result { + Ok(ToolResult::success_with_markdown( + json!({"result":"provider-specific","items":[{"url":"https://example.test/media"}]}), + "provider-specific markdown", + )) + } +} + +#[derive(Default)] +struct ProgressRecorder(Mutex>); +impl ToolRunContext for ProgressRecorder { + fn report_progress(&self, update: tinytools::ToolProgress) { + self.0 + .lock() + .unwrap() + .push(update.message.unwrap_or_default()); + } +} + +#[tokio::test] +async fn wrapper_keeps_metadata_and_files_each_artifact() { + let filer = Arc::new(Filer::default()); + let tool = MediaArtifactTool::new(Inner, MediaArtifactKind::Video, filer.clone()); + assert_eq!(tool.name(), "inner"); + assert_eq!(tool.description(), "stable"); + assert_eq!(tool.parameters_schema(), json!({"type":"object"})); + assert_eq!(tool.max_result_size_chars(), Some(1234)); + let result = tool.execute(json!({"prompt":" a title "})).await.unwrap(); + let ToolContent::Json { data } = &result.content[0] else { + panic!("expected JSON content") + }; + assert_eq!(data["artifacts"][0]["path"], "/files/image.png"); + assert_eq!(data["artifacts"][0]["artifact_id"], "id1"); + assert_eq!( + result.markdown_formatted.as_deref(), + Some("saved /files/image.png") + ); + let filed = filer.0.lock().unwrap(); + assert_eq!(filed[0].artifact_kind, MediaArtifactKind::Video); + assert_eq!(filed[0].title, "a title"); + assert_eq!(filed[0].path, PathBuf::from("blob:key")); +} + +#[tokio::test(start_paused = true)] +async fn concurrent_calls_keep_each_invocations_context_and_rewrite_separate_results() { + let barrier = Arc::new(tokio::sync::Barrier::new(2)); + let factory_barrier = Arc::clone(&barrier); + let tool = MediaArtifactTool::new( + ConcurrentInner { + identity: "unused", + progress: Arc::new(|_| {}), + barrier: Arc::clone(&barrier), + }, + MediaArtifactKind::Image, + Arc::new(RoutingFiler), + ) + .with_per_call_inner(Arc::new(move |progress, context| { + let identity = context + .and_then(|context| context.host_extension()) + .and_then(|extension| extension.downcast_ref::()) + .expect("factory receives each invocation context") + .0; + ConcurrentInner { + identity, + progress, + barrier: Arc::clone(&factory_barrier), + } + })) + .with_tool_call_id(Arc::new(|context| { + context + .and_then(|context| context.host_extension()) + .and_then(|extension| extension.downcast_ref::()) + .map(|identity| identity.0.to_owned()) + })); + let first = Context { + identity: CallIdentity("first"), + progress: Mutex::new(Vec::new()), + }; + let second = Context { + identity: CallIdentity("second"), + progress: Mutex::new(Vec::new()), + }; + let (left, right) = tokio::time::timeout(Duration::from_secs(60), async { + tokio::join!( + tool.execute_with_context( + json!({"prompt":"first"}), + ToolCallOptions::default(), + Some(&first) + ), + tool.execute_with_context( + json!({"prompt":"second"}), + ToolCallOptions::default(), + Some(&second) + ), + ) + }) + .await + .expect("both factory calls must reach the overlap barrier"); + for (result, identity) in [(left.unwrap(), "first"), (right.unwrap(), "second")] { + let ToolContent::Json { data } = &result.content[0] else { + panic!("expected JSON result") + }; + assert_eq!( + data["artifacts"][0]["artifact_id"], + format!("{identity}-{identity}-one") + ); + assert_eq!( + data["artifacts"][0]["path"], + format!("/files/{identity}-{identity}-one.png") + ); + assert_eq!( + data["artifacts"][1]["artifact_id"], + format!("{identity}-{identity}-two") + ); + assert_eq!( + data["artifacts"][1]["path"], + format!("/files/{identity}-{identity}-two.png") + ); + assert_eq!( + result.markdown_formatted.as_deref(), + Some( + format!( + "/files/{identity}-{identity}-one.png and /files/{identity}-{identity}-two.png" + ) + .as_str() + ) + ); + } + assert_eq!( + *first.progress.lock().unwrap(), + vec!["Generating image — 15s elapsed (job first-job)"] + ); + assert_eq!( + *second.progress.lock().unwrap(), + vec!["Generating image — 15s elapsed (job second-job)"] + ); +} + +#[tokio::test] +async fn generated_artifact_uses_canonical_harness_call_id_when_context_has_one() { + use tinyagents_harness::{ + context::{RunConfig, RunContext}, + ids::CallId, + tool::ToolExecutionContext, + }; + let context = ToolExecutionContext::from_run_context( + &RunContext::new(RunConfig::new("media-test"), ()), + CallId::new("provider-call-42"), + ); + let filer = Arc::new(Filer::default()); + let tool = MediaArtifactTool::new(Inner, MediaArtifactKind::Image, filer.clone()); + tool.execute_with_context( + json!({"prompt":"call id"}), + ToolCallOptions::default(), + Some(&context), + ) + .await + .unwrap(); + assert_eq!( + filer.0.lock().unwrap()[0].tool_call_id.as_deref(), + Some("provider-call-42") + ); +} + +#[tokio::test] +async fn per_call_inner_preserves_stable_metadata_and_files_multiple_entries_independently() { + let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let factory_calls = Arc::clone(&calls); + let tool = MediaArtifactTool::new(MultiInner, MediaArtifactKind::Image, Arc::new(OneFails)) + .with_per_call_inner(Arc::new(move |_progress, _context| { + factory_calls.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + MultiInner + })); + assert_eq!(tool.name(), "multi"); + assert_eq!(tool.description(), "multi description"); + assert_eq!(tool.category(), ToolCategory::Workflow); + assert_eq!(tool.permission_level(), PermissionLevel::Dangerous); + assert!(tool.external_effect()); + assert_eq!(tool.timeout_policy(&json!({})), ToolTimeout::Unbounded); + let result = tool + .execute(json!({"prompt":format!(" {} ", "x".repeat(70))})) + .await + .unwrap(); + assert_eq!(calls.load(std::sync::atomic::Ordering::SeqCst), 1); + let ToolContent::Json { data } = &result.content[0] else { + panic!("expected JSON") + }; + assert_eq!(data["other"], "retained"); + assert_eq!(data["artifacts"][0]["path"], "/files/good.png"); + assert_eq!(data["artifacts"][0]["artifact_id"], "good-id"); + assert_eq!(data["artifacts"][1]["path"], "blob:bad"); + assert_eq!(data["artifacts"][1]["artifact_error"], "filing failed"); + assert_eq!(data["artifacts"][2]["path"], "/files/good.png"); + assert_eq!(data["artifacts"][2]["artifact_id"], "good-id"); + assert_eq!( + result.markdown_formatted.as_deref(), + Some("good /files/good.png; bad blob:bad; last /files/good.png") + ); +} + +#[test] +fn configured_media_kinds_keep_their_exact_heartbeat_labels() { + assert_eq!( + progress_label(&MediaArtifactKind::Image), + "Generating image" + ); + assert_eq!( + progress_label(&MediaArtifactKind::Video), + "Generating video" + ); +} + +#[test] +fn heartbeat_keeps_elapsed_duration_when_job_state_is_present() { + assert_eq!( + heartbeat_message( + "Generating video", + std::time::Duration::from_secs(17), + Some("processing") + ), + "Generating video — 17s elapsed (job processing)" + ); + assert_eq!( + heartbeat_message( + "Generating image", + std::time::Duration::from_secs(125), + None + ), + "Generating image — 2m 05s elapsed" + ); +} + +#[tokio::test(start_paused = true)] +async fn media_heartbeat_reports_per_call_job_state() { + let recorder = ProgressRecorder::default(); + let tool = MediaArtifactTool::new( + ProgressInner(Arc::new(|_| {})), + MediaArtifactKind::Video, + Arc::new(Filer::default()), + ) + .with_per_call_inner(Arc::new(|progress, _context| ProgressInner(progress))); + let result = tool + .execute_with_context(json!({}), ToolCallOptions::default(), Some(&recorder)) + .await + .unwrap(); + assert_eq!(result.text(), "done"); + assert_eq!( + *recorder.0.lock().unwrap(), + vec!["Generating video — 15s elapsed (job in_progress)"], + ); +} + +#[tokio::test(start_paused = true)] +async fn image_and_video_factory_calls_keep_configured_kind_and_heartbeat() { + for (kind, label) in [ + (MediaArtifactKind::Image, "Generating image"), + (MediaArtifactKind::Video, "Generating video"), + ] { + let recorder = ProgressRecorder::default(); + let filer = Arc::new(Filer::default()); + let tool = MediaArtifactTool::new( + ProgressArtifactInner(Arc::new(|_| {})), + kind.clone(), + filer.clone(), + ) + .with_per_call_inner(Arc::new(|progress, _context| { + ProgressArtifactInner(progress) + })); + let result = tool + .execute_with_context(json!({}), ToolCallOptions::default(), Some(&recorder)) + .await + .unwrap(); + let ToolContent::Json { data } = &result.content[0] else { + panic!("expected JSON result") + }; + assert_eq!(data["artifacts"][0]["path"], "/files/image.png"); + assert_eq!( + *recorder.0.lock().unwrap(), + vec![format!("{label} — 15s elapsed (job running)")] + ); + assert_eq!(filer.0.lock().unwrap()[0].artifact_kind, kind); + } +} + +#[tokio::test] +async fn errors_and_unknown_output_shapes_are_preserved_without_filing() { + let filer = Arc::new(Filer::default()); + let error_tool = MediaArtifactTool::new(ErrorInner, MediaArtifactKind::Image, filer.clone()); + let error = error_tool.execute(json!({})).await.unwrap(); + assert!(error.is_error); + assert_eq!(error.text(), "provider failed"); + let unknown_tool = + MediaArtifactTool::new(UnknownShapeInner, MediaArtifactKind::Image, filer.clone()); + let unknown = unknown_tool.execute(json!({})).await.unwrap(); + let ToolContent::Json { data } = &unknown.content[0] else { + panic!("expected JSON") + }; + assert_eq!( + data, + &json!({"result":"provider-specific","items":[{"url":"https://example.test/media"}]}) + ); + assert_eq!( + unknown.markdown_formatted.as_deref(), + Some("provider-specific markdown") + ); + assert!(filer.0.lock().unwrap().is_empty()); +} diff --git a/crates/tinyagents-harness-tools/src/media/mod.rs b/crates/tinyagents-harness-tools/src/media/mod.rs new file mode 100644 index 000000000..2a71865a2 --- /dev/null +++ b/crates/tinyagents-harness-tools/src/media/mod.rs @@ -0,0 +1,7 @@ +//! Host-neutral wrappers for tools that generate media artifacts. +mod artifact_filing; +mod artifact_tool; +pub use artifact_filing::{ArtifactFiling, FiledArtifact, GeneratedArtifact, MediaArtifactKind}; +pub use artifact_tool::{ + MediaArtifactTool, MediaInnerFactory, MediaJobProgress, MediaProgressFn, ToolCallIdExtractor, +}; diff --git a/crates/tinyagents-harness-tools/src/net/mod.rs b/crates/tinyagents-harness-tools/src/net/mod.rs new file mode 100644 index 000000000..ae8de9195 --- /dev/null +++ b/crates/tinyagents-harness-tools/src/net/mod.rs @@ -0,0 +1,3 @@ +//! Network tools behind host-provided egress gates. + +pub mod network; diff --git a/crates/tinyagents-harness-tools/src/net/network/contract_tests.rs b/crates/tinyagents-harness-tools/src/net/network/contract_tests.rs new file mode 100644 index 000000000..33ac31f00 --- /dev/null +++ b/crates/tinyagents-harness-tools/src/net/network/contract_tests.rs @@ -0,0 +1,126 @@ +//! Wire-contract fixtures: the name, description, permission level, exposure +//! and JSON Schema each network tool advertises to the model. +//! +//! These are what a model is prompted with, and what existing transcripts were +//! produced against, so they are pinned against literal fixture files rather +//! than re-derived from the code under test. A change here is a change to the +//! tool's public contract and must be made on purpose. + +#![allow(clippy::expect_used, clippy::panic, clippy::unwrap_used)] + +use std::path::PathBuf; +use std::sync::Arc; + +use serde_json::{Value, json}; +use tinytools::Tool; + +use super::test_support::{AutonomyLevel, DEFAULT_LIMITS, TestHtml, TestNetGate}; +use super::*; + +fn tools() -> Vec<(&'static str, Box)> { + let gate = TestNetGate::supervised(); + vec![ + ( + "curl", + Box::new(CurlTool::new( + gate.clone(), + vec![], + PathBuf::from("."), + "downloads".into(), + 1024, + 30, + )) as Box, + ), + ( + "http_request", + Box::new(HttpRequestTool::new( + gate.clone(), + vec![], + 1024, + 30, + DEFAULT_LIMITS, + )), + ), + ( + "pushover", + Box::new(PushoverTool::new(gate.clone(), PathBuf::from("."))), + ), + ( + "web_fetch", + Box::new(WebFetchTool::new( + gate, + vec![], + None, + None, + DEFAULT_LIMITS, + Arc::new(TestHtml), + )), + ), + ] +} + +fn describe(tool: &dyn Tool) -> Value { + json!({ + "name": tool.name(), + "description": tool.description(), + "permission_level": format!("{:?}", tool.permission_level()), + "exposure": format!("{:?}", tool.exposure()), + "schema": tool.parameters_schema(), + }) +} + +fn fixture(name: &str) -> Value { + let raw = match name { + "curl" => include_str!("fixtures/curl.json"), + "http_request" => include_str!("fixtures/http_request.json"), + "pushover" => include_str!("fixtures/pushover.json"), + "web_fetch" => include_str!("fixtures/web_fetch.json"), + other => panic!("no fixture for {other}"), + }; + serde_json::from_str(raw).expect("fixture is valid JSON") +} + +#[test] +fn every_tool_advertises_exactly_its_pinned_contract() { + for (fixture_name, tool) in tools() { + assert_eq!( + describe(tool.as_ref()), + fixture(fixture_name), + "contract drift in `{fixture_name}`" + ); + } +} + +#[test] +fn web_fetch_with_a_host_property_extends_only_the_properties() { + let tool = WebFetchTool::new( + TestNetGate::supervised(), + vec![], + None, + None, + DEFAULT_LIMITS, + Arc::new(TestHtml), + ) + .with_schema_property("summary_focus", json!({"type": "string"})); + let mut expected = fixture("web_fetch"); + expected["schema"]["properties"]["summary_focus"] = json!({"type": "string"}); + assert_eq!(describe(&tool), expected); +} + +#[test] +fn only_the_approval_gated_tools_route_through_approval() { + let args = json!({"url": "https://example.com"}); + for (autonomy, prompts) in [ + (AutonomyLevel::Supervised, true), + (AutonomyLevel::Full, false), + (AutonomyLevel::ReadOnly, false), + ] { + let gate = TestNetGate::with(autonomy, 100); + let curl = CurlTool::new(gate.clone(), vec![], ".".into(), "d".into(), 1, 1); + let http = HttpRequestTool::new(gate.clone(), vec![], 1, 1, DEFAULT_LIMITS); + let fetch = WebFetchTool::new(gate, vec![], None, None, DEFAULT_LIMITS, Arc::new(TestHtml)); + assert_eq!(curl.external_effect_with_args(&args), prompts); + assert_eq!(http.external_effect_with_args(&args), prompts); + assert!(!fetch.external_effect_with_args(&args)); + } +} diff --git a/crates/tinyagents-harness-tools/src/net/network/curl.rs b/crates/tinyagents-harness-tools/src/net/network/curl.rs new file mode 100644 index 000000000..b03dfe654 --- /dev/null +++ b/crates/tinyagents-harness-tools/src/net/network/curl.rs @@ -0,0 +1,429 @@ +//! `curl` — download files from the web to a path under the workspace. +//! +//! Distinct from `http_request`: instead of returning the body inline +//! (size-capped), `curl` streams to disk with a hard byte ceiling. Same +//! SSRF/allowlist guards (shared via `url_guard`), shares +//! `http_request.allowed_domains` so there is one allowlist to reason +//! about. + +use super::gate::{NetGate, host_of}; +use tinyagents_harness::url_guard::{ValidatedUrl, normalize_allowed_domains}; + +use super::gate::validate_destination; +use async_trait::async_trait; +use futures_util::StreamExt; +use sha2::{Digest, Sha256}; +use std::fmt::Write as _; +use std::path::{Component, Path, PathBuf}; +use std::sync::Arc; +use std::time::Duration; +use tinytools::{PermissionLevel, Tool, ToolResult}; +use tokio::fs; +use tokio::io::AsyncWriteExt; + +/// Downloads a file from an http(s) URL into the workspace. +#[derive(Debug)] +pub struct CurlTool { + gate: Arc, + allowed_domains: Vec, + workspace_dir: PathBuf, + dest_subdir: String, + max_download_bytes: u64, + timeout_secs: u64, +} + +impl CurlTool { + /// A `curl` tool over `gate`, downloading under `/`. + pub fn new( + gate: Arc, + allowed_domains: Vec, + workspace_dir: PathBuf, + dest_subdir: String, + max_download_bytes: u64, + timeout_secs: u64, + ) -> Self { + Self { + gate, + allowed_domains: normalize_allowed_domains(allowed_domains), + workspace_dir, + dest_subdir: sanitize_dest_subdir(&dest_subdir), + max_download_bytes, + timeout_secs, + } + } + + /// Resolve a user-supplied dest path to an absolute path inside + /// `/`. Rejects absolute paths, `..` + /// segments, and any other escape attempts. + fn resolve_dest(&self, dest: &str) -> anyhow::Result { + let trimmed = dest.trim(); + if trimmed.is_empty() { + anyhow::bail!("dest_path cannot be empty"); + } + + let p = Path::new(trimmed); + if p.is_absolute() { + anyhow::bail!("dest_path must be relative — got absolute path"); + } + + for component in p.components() { + match component { + Component::Normal(_) | Component::CurDir => {} + Component::ParentDir => { + anyhow::bail!("dest_path may not contain '..'"); + } + Component::Prefix(_) | Component::RootDir => { + anyhow::bail!("dest_path must be relative"); + } + } + } + + let root = self.workspace_dir.join(&self.dest_subdir); + let resolved = root.join(p); + + // Belt-and-braces: ensure the resolved path still lives under root. + // Lexical check is sufficient because we already rejected `..`. + if !resolved.starts_with(&root) { + anyhow::bail!("dest_path resolves outside the downloads root"); + } + + Ok(resolved) + } + + async fn validate_url( + &self, + raw_url: &str, + headers: &[(String, String)], + ) -> anyhow::Result { + validate_destination( + &*self.gate, + raw_url, + &self.allowed_domains, + "GET", + headers, + None, + ) + .await + } + + fn default_filename_from_url(url: &str) -> String { + let after_scheme = url.split_once("://").map_or(url, |(_, rest)| rest); + let path_part = after_scheme.split_once('/').map_or("", |(_, p)| p); + let last = path_part + .split('?') + .next() + .unwrap_or("") + .rsplit('/') + .next() + .unwrap_or(""); + let cleaned: String = last + .chars() + .filter(|c| c.is_alphanumeric() || matches!(c, '.' | '-' | '_')) + .collect(); + if cleaned.is_empty() { + "download.bin".into() + } else { + cleaned + } + } +} + +#[async_trait] +impl Tool for CurlTool { + fn exposure(&self) -> tinytools::ToolExposure { + tinytools::ToolExposure::Deferred + } + + fn name(&self) -> &'static str { + "curl" + } + + fn description(&self) -> &'static str { + "Download a file from an http(s) URL into the workspace. The body is streamed to disk \ + with a hard byte ceiling. Same allowlist as `http_request`. Returns the saved path, \ + bytes written, content-type, and SHA-256 of the file." + } + + fn parameters_schema(&self) -> serde_json::Value { + serde_json::json!({ + "type": "object", + "properties": { + "url": { + "type": "string", + "description": "HTTP or HTTPS URL of the file to download" + }, + "dest_path": { + "type": "string", + "description": "Destination path relative to the downloads root inside the workspace. No '..' or absolute paths. If omitted, the filename is inferred from the URL." + }, + "headers": { + "type": "object", + "description": "Optional HTTP headers (e.g. {\"Authorization\": \"Bearer …\"})", + "default": {} + } + }, + "required": ["url"] + }) + } + + /// Downloading from the network is the always-ask `Network` bucket — it + /// prompts the human in both ask-before-edit and Full; read-only is blocked + /// in `execute`. + fn external_effect_with_args(&self, _args: &serde_json::Value) -> bool { + self.gate.network_needs_approval() + } + + fn permission_level(&self) -> PermissionLevel { + PermissionLevel::Write + } + + async fn execute(&self, args: serde_json::Value) -> anyhow::Result { + let url = args + .get("url") + .and_then(|v| v.as_str()) + .ok_or_else(|| anyhow::anyhow!("Missing 'url' parameter"))?; + + let dest_arg = args.get("dest_path").and_then(|v| v.as_str()); + let headers_val = args + .get("headers") + .cloned() + .unwrap_or_else(|| serde_json::json!({})); + + if !self.gate.can_act() { + tracing::debug!(target: "[curl]", url = %url, "blocked: autonomy read-only"); + return Ok(ToolResult::error( + "[policy-blocked] Action blocked: autonomy is read-only", + )); + } + if !self.gate.record_action() { + tracing::debug!(target: "[curl]", url = %url, "blocked: rate limit"); + return Ok(ToolResult::error("Action blocked: rate limit exceeded")); + } + + // Local-only enforcement (privacy epic S7, #4441): mirror the read-only + // `can_act()` deny above — under LocalOnly, refuse the download before + // URL validation / DNS so nothing leaves the device. + { + let host = host_of(url); + if let Some(msg) = self.gate.local_only_block(&host) { + // Log only the host, never the full URL: a raw URL can carry + // secrets in its query string (pre-signed links, tokens). The + // gate helper already logs `desc.service` (host only) too. + tracing::debug!(target: "[curl]", host = %host, "blocked: local-only privacy mode"); + return Ok(ToolResult::error(msg)); + } + } + + let request_headers: Vec<(String, String)> = headers_val + .as_object() + .map(|obj| { + obj.iter() + .filter_map(|(k, v)| v.as_str().map(|s| (k.clone(), s.to_string()))) + .collect() + }) + .unwrap_or_default(); + let target = match self.validate_url(url, &request_headers).await { + Ok(v) => v, + Err(e) => { + tracing::debug!(target: "[curl]", url = %url, reason = %e, "url validation failed"); + return Ok(ToolResult::error(e.to_string())); + } + }; + + let url = target.url.clone(); + + // Egress spine (privacy epic S2/S7, #4436/#4441): a curl download + // contacts an external host — disclose the destination (and that custom + // headers ride along) before the request. Enforcement already ran at the + // top of `execute`; this is the observe-only disclosure for permitted + // downloads. + { + let host = host_of(&url); + let has_headers = headers_val.as_object().is_some_and(|h| !h.is_empty()); + self.gate.disclose(&host, false, has_headers); + } + + let dest = match dest_arg { + Some(d) => d.to_string(), + None => Self::default_filename_from_url(&url), + }; + let dest_path = match self.resolve_dest(&dest) { + Ok(p) => p, + Err(e) => { + tracing::debug!(target: "[curl]", url = %url, dest = %dest, reason = %e, "dest_path rejected"); + return Ok(ToolResult::error(e.to_string())); + } + }; + + if let Some(parent) = dest_path.parent() + && let Err(e) = fs::create_dir_all(parent).await + { + tracing::error!(target: "[curl]", url = %url, dest = %dest_path.display(), reason = %e, "create_dir_all failed"); + return Ok(ToolResult::error(format!( + "Failed to create destination directory: {e}" + ))); + } + + let builder = reqwest::Client::builder() + .timeout(Duration::from_secs(self.timeout_secs)) + .connect_timeout(Duration::from_secs(10)) + .redirect(reqwest::redirect::Policy::none()); + let builder = self.gate.prepare_client("tool.curl", builder); + // Connect only to the vetted addresses; the URL keeps its host name. + let builder = builder.resolve_to_addrs(&target.host, &target.addrs); + let client = match builder.build() { + Ok(c) => c, + Err(e) => { + tracing::error!(target: "[curl]", reason = %e, "HTTP client build failed"); + return Ok(ToolResult::error(format!("HTTP client build failed: {e}"))); + } + }; + + let mut request = client.get(&url); + if let Some(obj) = headers_val.as_object() { + for (k, v) in obj { + if let Some(s) = v.as_str() { + request = request.header(k, s); + } + } + } + + tracing::debug!(target: "[curl]", url = %url, dest = %dest_path.display(), "starting download"); + + let response = match request.send().await { + Ok(r) => r, + Err(e) => { + tracing::error!(target: "[curl]", url = %url, reason = %e, "request send failed"); + return Ok(ToolResult::error(format!("Request failed: {e}"))); + } + }; + + let status = response.status(); + if !status.is_success() { + tracing::debug!(target: "[curl]", url = %url, status = %status.as_u16(), "non-success HTTP status"); + return Ok(ToolResult::error(format!( + "HTTP {} {}", + status.as_u16(), + status.canonical_reason().unwrap_or("Unknown") + ))); + } + + let content_type = response + .headers() + .get(reqwest::header::CONTENT_TYPE) + .and_then(|v| v.to_str().ok()) + .unwrap_or("application/octet-stream") + .to_string(); + + let mut file = match fs::File::create(&dest_path).await { + Ok(f) => f, + Err(e) => { + tracing::error!(target: "[curl]", dest = %dest_path.display(), reason = %e, "fs::File::create failed"); + return Ok(ToolResult::error(format!( + "Failed to create destination file: {e}" + ))); + } + }; + + let mut hasher = Sha256::new(); + let mut bytes_written: u64 = 0; + let mut stream = response.bytes_stream(); + + while let Some(chunk) = stream.next().await { + let chunk = match chunk { + Ok(c) => c, + Err(e) => { + drop(file); + if let Err(rm) = fs::remove_file(&dest_path).await { + tracing::debug!(target: "[curl]", dest = %dest_path.display(), reason = %rm, "cleanup remove_file failed"); + } + tracing::error!(target: "[curl]", url = %url, bytes_written, reason = %e, "stream error"); + return Ok(ToolResult::error(format!("Stream error: {e}"))); + } + }; + + if bytes_written.saturating_add(chunk.len() as u64) > self.max_download_bytes { + let _ = file.flush().await; + drop(file); + if let Err(rm) = fs::remove_file(&dest_path).await { + tracing::debug!(target: "[curl]", dest = %dest_path.display(), reason = %rm, "cleanup remove_file failed"); + } + tracing::error!(target: "[curl]", url = %url, bytes_written, max = self.max_download_bytes, "size cap exceeded — download aborted"); + return Ok(ToolResult::error(format!( + "Download exceeded max_download_bytes ({} bytes)", + self.max_download_bytes + ))); + } + + if let Err(e) = file.write_all(&chunk).await { + drop(file); + if let Err(rm) = fs::remove_file(&dest_path).await { + tracing::debug!(target: "[curl]", dest = %dest_path.display(), reason = %rm, "cleanup remove_file failed"); + } + tracing::error!(target: "[curl]", dest = %dest_path.display(), bytes_written, reason = %e, "write_all failed"); + return Ok(ToolResult::error(format!("Write failed: {e}"))); + } + hasher.update(&chunk); + bytes_written += chunk.len() as u64; + } + + if let Err(e) = file.flush().await { + drop(file); + if let Err(rm) = fs::remove_file(&dest_path).await { + tracing::debug!(target: "[curl]", dest = %dest_path.display(), reason = %rm, "cleanup remove_file failed"); + } + tracing::error!(target: "[curl]", dest = %dest_path.display(), bytes_written, reason = %e, "flush failed"); + return Ok(ToolResult::error(format!("Flush failed: {e}"))); + } + + let mut sha256 = String::with_capacity(64); + for byte in hasher.finalize() { + let _ = write!(sha256, "{byte:02x}"); + } + + tracing::debug!( + target: "[curl]", + url = %url, + dest = %dest_path.display(), + bytes = bytes_written, + content_type = %content_type, + sha256 = %sha256, + "download complete" + ); + + let payload = serde_json::json!({ + "path": dest_path.display().to_string(), + "bytes_written": bytes_written, + "content_type": content_type, + "sha256": sha256, + }); + Ok(ToolResult::success(payload.to_string())) + } +} + +/// Sanitize the configured `dest_subdir` so a malicious or misconfigured +/// `[curl].dest_subdir` cannot escape the workspace via absolute paths +/// or `..` segments. Drops disallowed components rather than panicking; +/// falls back to `"downloads"` if everything is filtered out. +fn sanitize_dest_subdir(raw: &str) -> String { + let trimmed = raw.trim(); + if trimmed.is_empty() { + return "downloads".into(); + } + let p = Path::new(trimmed); + let mut buf = PathBuf::new(); + for component in p.components() { + // Drop everything else: absolute roots, prefixes, parent dirs, cur dirs. + if let Component::Normal(c) = component { + buf.push(c); + } + } + if buf.as_os_str().is_empty() { + return "downloads".into(); + } + buf.to_string_lossy().into_owned() +} + +#[cfg(test)] +#[path = "curl_tests.rs"] +mod tests; diff --git a/crates/tinyagents-harness-tools/src/net/network/curl_tests.rs b/crates/tinyagents-harness-tools/src/net/network/curl_tests.rs new file mode 100644 index 000000000..c15a1d95a --- /dev/null +++ b/crates/tinyagents-harness-tools/src/net/network/curl_tests.rs @@ -0,0 +1,324 @@ +use super::*; +use crate::net::network::test_support::{AutonomyLevel, TestNetGate}; +use tempfile::TempDir; +fn slash_norm(s: String) -> String { + s.replace('\\', "/") +} + +fn tool(tmp: &TempDir, allow: Vec<&str>) -> CurlTool { + CurlTool::new( + TestNetGate::supervised(), + allow.into_iter().map(String::from).collect(), + tmp.path().to_path_buf(), + "downloads".into(), + 1024 * 1024, + 30, + ) +} + +#[test] +fn sanitize_dest_subdir_strips_absolute_paths() { + assert_eq!( + slash_norm(sanitize_dest_subdir("/etc/passwd")), + "etc/passwd" + ); + assert_eq!(sanitize_dest_subdir("//foo"), "foo"); +} + +#[test] +fn sanitize_dest_subdir_strips_parent_segments() { + assert_eq!(sanitize_dest_subdir("../../etc"), "etc"); + assert_eq!(slash_norm(sanitize_dest_subdir("a/../b")), "a/b"); +} + +#[test] +fn sanitize_dest_subdir_falls_back_to_downloads() { + assert_eq!(sanitize_dest_subdir(""), "downloads"); + assert_eq!(sanitize_dest_subdir(" "), "downloads"); + assert_eq!(sanitize_dest_subdir(".."), "downloads"); + assert_eq!(sanitize_dest_subdir("/"), "downloads"); +} + +#[test] +fn sanitize_dest_subdir_keeps_normal_paths() { + assert_eq!(sanitize_dest_subdir("downloads"), "downloads"); + assert_eq!( + slash_norm(sanitize_dest_subdir("artifacts/build")), + "artifacts/build" + ); +} + +#[test] +fn new_sanitizes_malicious_dest_subdir() { + let tmp = TempDir::new().unwrap(); + let t = CurlTool::new( + TestNetGate::supervised(), + vec!["example.com".into()], + tmp.path().to_path_buf(), + "../../etc".into(), + 1024, + 30, + ); + let resolved = t.resolve_dest("file.txt").unwrap(); + // Sanitizer reduced "../../etc" to "etc"; resolution must stay under workspace. + assert!(resolved.starts_with(tmp.path().join("etc"))); + assert!(resolved.starts_with(tmp.path())); +} + +#[test] +fn resolve_dest_normal() { + let tmp = TempDir::new().unwrap(); + let t = tool(&tmp, vec!["example.com"]); + let p = t.resolve_dest("foo/bar.txt").unwrap(); + assert!(p.starts_with(tmp.path().join("downloads"))); + assert!(p.ends_with("foo/bar.txt")); +} + +#[test] +fn resolve_dest_rejects_absolute() { + let tmp = TempDir::new().unwrap(); + let t = tool(&tmp, vec!["example.com"]); + let err = t.resolve_dest("/etc/passwd").unwrap_err().to_string(); + assert!(err.contains("relative")); +} + +#[test] +fn resolve_dest_rejects_parent_dir() { + let tmp = TempDir::new().unwrap(); + let t = tool(&tmp, vec!["example.com"]); + let err = t.resolve_dest("../etc/passwd").unwrap_err().to_string(); + assert!(err.contains("..")); +} + +#[test] +fn resolve_dest_rejects_nested_parent_dir() { + let tmp = TempDir::new().unwrap(); + let t = tool(&tmp, vec!["example.com"]); + let err = t.resolve_dest("a/../../b").unwrap_err().to_string(); + assert!(err.contains("..")); +} + +#[test] +fn resolve_dest_rejects_empty() { + let tmp = TempDir::new().unwrap(); + let t = tool(&tmp, vec!["example.com"]); + assert!(t.resolve_dest("").is_err()); + assert!(t.resolve_dest(" ").is_err()); +} + +#[tokio::test] +async fn validate_url_rejects_disallowed_domain() { + let tmp = TempDir::new().unwrap(); + let t = tool(&tmp, vec!["example.com"]); + let err = t + .validate_url("https://evil.test/archive.tar.gz", &[]) + .await + .unwrap_err() + .to_string(); + assert!(err.contains("allowed websites")); +} + +#[test] +fn default_filename_from_url_basic() { + assert_eq!( + CurlTool::default_filename_from_url("https://example.com/foo/bar.zip"), + "bar.zip" + ); +} + +#[test] +fn default_filename_from_url_query_stripped() { + assert_eq!( + CurlTool::default_filename_from_url("https://example.com/file.tar.gz?token=x"), + "file.tar.gz" + ); +} + +#[test] +fn default_filename_from_url_root_falls_back() { + assert_eq!( + CurlTool::default_filename_from_url("https://example.com/"), + "download.bin" + ); +} + +#[tokio::test] +async fn execute_blocks_when_rate_limited() { + let tmp = TempDir::new().unwrap(); + let t = CurlTool::new( + TestNetGate::with(AutonomyLevel::Supervised, 0), + vec!["example.com".into()], + tmp.path().into(), + "downloads".into(), + 1024, + 30, + ); + let result = t + .execute(serde_json::json!({"url": "https://example.com/x"})) + .await + .unwrap(); + assert!(result.is_error); + assert!(result.output().contains("rate limit")); +} + +#[tokio::test] +async fn execute_blocked_under_local_only_privacy_mode() { + // Privacy epic S7 (#4441): under LocalOnly the download is refused with a + // `[policy-blocked]` result before URL validation / network. + let tmp = TempDir::new().unwrap(); + let t = CurlTool::new( + TestNetGate::local_only(), + vec!["example.com".into()], + tmp.path().to_path_buf(), + "downloads".into(), + 1024 * 1024, + 30, + ); + let result = t + .execute(serde_json::json!({"url": "https://example.com/x"})) + .await + .unwrap(); + assert!(result.is_error); + assert!( + result.output().contains("[policy-blocked]"), + "got: {}", + result.output() + ); + assert!( + result.output().contains("Local-only"), + "got: {}", + result.output() + ); +} + +/// Live integration smoke: downloads example.com (a tiny, stable +/// public page). Gated behind `OPENHUMAN_CURL_LIVE_TEST=1` so CI / +/// offline runs don't depend on the network. +#[tokio::test] +async fn live_download_example_com() { + if std::env::var("OPENHUMAN_CURL_LIVE_TEST").ok().as_deref() != Some("1") { + return; + } + let tmp = TempDir::new().unwrap(); + let t = tool(&tmp, vec!["example.com"]); + let result = t + .execute(serde_json::json!({ + "url": "https://example.com/", + "dest_path": "example.html" + })) + .await + .unwrap(); + assert!(!result.is_error, "live curl errored: {}", result.output()); + let payload: serde_json::Value = serde_json::from_str(&result.output()).unwrap(); + let bytes = payload["bytes_written"].as_u64().unwrap(); + assert!(bytes > 100, "unexpectedly small download: {bytes} bytes"); + let path = payload["path"].as_str().unwrap(); + let content = std::fs::read_to_string(path).unwrap(); + assert!(content.to_lowercase().contains("example domain")); +} + +#[tokio::test] +async fn execute_rejects_allowlist_miss() { + let tmp = TempDir::new().unwrap(); + let t = tool(&tmp, vec!["example.com"]); + let result = t + .execute(serde_json::json!({"url": "https://other.example.org/x"})) + .await + .unwrap(); + assert!(result.is_error); + assert!(result.output().contains("allowed websites")); +} + +/// The downloads root cannot be created (a plain file already occupies its +/// path) — `create_dir_all` must be reported as a tool error, never a panic, +/// and no network request is attempted before that check runs. The URL uses +/// a public IPv4 literal (never actually contacted — `create_dir_all` fails +/// and returns before any HTTP client is built) so `validate_url_with_dns_check` +/// takes its IP-literal short-circuit and this test needs no live DNS. +#[tokio::test] +async fn execute_reports_a_create_dir_all_failure_rather_than_panicking() { + let tmp = TempDir::new().unwrap(); + std::fs::write(tmp.path().join("downloads"), b"not a directory").unwrap(); + let t = tool(&tmp, vec!["8.8.8.8"]); + let result = t + .execute(serde_json::json!({"url": "https://8.8.8.8/x", "dest_path": "sub/file.txt"})) + .await + .unwrap(); + assert!( + result.is_error, + "an unwritable downloads root must be reported, not panic: {}", + result.output() + ); + assert!( + result + .output() + .contains("Failed to create destination directory"), + "must fail at the create_dir_all step specifically, not e.g. DNS resolution: {}", + result.output() + ); +} + +#[tokio::test] +async fn execute_blocks_readonly_mode() { + let tmp = TempDir::new().unwrap(); + let t = CurlTool::new( + TestNetGate::with(AutonomyLevel::ReadOnly, 100), + vec!["example.com".into()], + tmp.path().into(), + "downloads".into(), + 1024, + 30, + ); + let result = t + .execute(serde_json::json!({"url": "https://example.com/x"})) + .await + .unwrap(); + assert!(result.is_error); + assert!(result.output().contains("read-only")); +} + +#[test] +fn approval_follows_the_gate() { + let tmp = TempDir::new().unwrap(); + let args = serde_json::json!({"url": "https://example.com/x"}); + for (autonomy, prompts) in [ + (AutonomyLevel::Supervised, true), + (AutonomyLevel::Full, false), + ] { + let t = CurlTool::new( + TestNetGate::with(autonomy, 100), + vec!["example.com".into()], + tmp.path().into(), + "downloads".into(), + 1024, + 30, + ); + assert_eq!(t.external_effect_with_args(&args), prompts); + } +} + +#[tokio::test] +async fn a_download_discloses_its_destination_and_headers_before_the_request() { + let tmp = TempDir::new().unwrap(); + let gate = TestNetGate::supervised(); + // A public IPv4 literal short-circuits DNS; `create_dir_all` then fails + // (a file occupies the downloads root), so nothing is ever contacted. + std::fs::write(tmp.path().join("downloads"), b"not a directory").unwrap(); + let t = CurlTool::new( + gate.clone(), + vec!["8.8.8.8".into()], + tmp.path().into(), + "downloads".into(), + 1024, + 30, + ); + let result = t + .execute(serde_json::json!({ + "url": "https://8.8.8.8/x", + "headers": {"Authorization": "Bearer t"} + })) + .await + .unwrap(); + assert!(result.is_error); + assert_eq!(gate.disclosed(), vec![("8.8.8.8".to_string(), false, true)]); +} diff --git a/crates/tinyagents-harness-tools/src/net/network/fixtures/curl.json b/crates/tinyagents-harness-tools/src/net/network/fixtures/curl.json new file mode 100644 index 000000000..0cfc5856f --- /dev/null +++ b/crates/tinyagents-harness-tools/src/net/network/fixtures/curl.json @@ -0,0 +1,27 @@ +{ + "description": "Download a file from an http(s) URL into the workspace. The body is streamed to disk with a hard byte ceiling. Same allowlist as `http_request`. Returns the saved path, bytes written, content-type, and SHA-256 of the file.", + "exposure": "Deferred", + "name": "curl", + "permission_level": "Write", + "schema": { + "properties": { + "dest_path": { + "description": "Destination path relative to the downloads root inside the workspace. No '..' or absolute paths. If omitted, the filename is inferred from the URL.", + "type": "string" + }, + "headers": { + "default": {}, + "description": "Optional HTTP headers (e.g. {\"Authorization\": \"Bearer …\"})", + "type": "object" + }, + "url": { + "description": "HTTP or HTTPS URL of the file to download", + "type": "string" + } + }, + "required": [ + "url" + ], + "type": "object" + } +} diff --git a/crates/tinyagents-harness-tools/src/net/network/fixtures/http_request.json b/crates/tinyagents-harness-tools/src/net/network/fixtures/http_request.json new file mode 100644 index 000000000..c09b9e97d --- /dev/null +++ b/crates/tinyagents-harness-tools/src/net/network/fixtures/http_request.json @@ -0,0 +1,32 @@ +{ + "description": "Make HTTP requests to external APIs. Supports GET, POST, PUT, DELETE, PATCH, HEAD, OPTIONS methods. Security constraints: allowlist-only domains, no local/private hosts, configurable timeout and response size limits.", + "exposure": "Direct", + "name": "http_request", + "permission_level": "Write", + "schema": { + "properties": { + "body": { + "description": "Optional request body (for POST, PUT, PATCH requests)", + "type": "string" + }, + "headers": { + "default": {}, + "description": "Optional HTTP headers as key-value pairs (e.g., {\"Authorization\": \"Bearer token\", \"Content-Type\": \"application/json\"})", + "type": "object" + }, + "method": { + "default": "GET", + "description": "HTTP method (GET, POST, PUT, DELETE, PATCH, HEAD, OPTIONS)", + "type": "string" + }, + "url": { + "description": "HTTP or HTTPS URL to request", + "type": "string" + } + }, + "required": [ + "url" + ], + "type": "object" + } +} diff --git a/crates/tinyagents-harness-tools/src/net/network/fixtures/pushover.json b/crates/tinyagents-harness-tools/src/net/network/fixtures/pushover.json new file mode 100644 index 000000000..7cea38c00 --- /dev/null +++ b/crates/tinyagents-harness-tools/src/net/network/fixtures/pushover.json @@ -0,0 +1,30 @@ +{ + "description": "Send a Pushover notification to your device. Requires PUSHOVER_TOKEN and PUSHOVER_USER_KEY in .env file.", + "exposure": "Direct", + "name": "pushover", + "permission_level": "ReadOnly", + "schema": { + "properties": { + "message": { + "description": "The notification message to send", + "type": "string" + }, + "priority": { + "description": "Message priority: -2 (lowest/silent), -1 (low/no sound), 0 (normal), 1 (high), 2 (emergency/repeating)", + "type": "integer" + }, + "sound": { + "description": "Notification sound override (e.g., 'pushover', 'bike', 'bugle', 'cashregister', etc.)", + "type": "string" + }, + "title": { + "description": "Optional notification title", + "type": "string" + } + }, + "required": [ + "message" + ], + "type": "object" + } +} diff --git a/crates/tinyagents-harness-tools/src/net/network/fixtures/web_fetch.json b/crates/tinyagents-harness-tools/src/net/network/fixtures/web_fetch.json new file mode 100644 index 000000000..a715af69d --- /dev/null +++ b/crates/tinyagents-harness-tools/src/net/network/fixtures/web_fetch.json @@ -0,0 +1,27 @@ +{ + "description": "GET a URL and read the page. HTML returns as Markdown (links kept, scripts dropped); `raw: true` for the body as sent. For POST or custom headers use `http_request`.", + "exposure": "Direct", + "name": "web_fetch", + "permission_level": "ReadOnly", + "schema": { + "properties": { + "max_bytes": { + "description": "Cap the returned text at this many bytes (default 1_000_000). Bounds the OUTPUT — the extracted markdown, or the raw body with raw:true — never the markup the extractor reads, so lowering it cannot cost you content the page actually had.", + "minimum": 1, + "type": "integer" + }, + "raw": { + "description": "Return the body as sent.", + "type": "boolean" + }, + "url": { + "description": "Absolute http(s) URL.", + "type": "string" + } + }, + "required": [ + "url" + ], + "type": "object" + } +} diff --git a/crates/tinyagents-harness-tools/src/net/network/gate.rs b/crates/tinyagents-harness-tools/src/net/network/gate.rs new file mode 100644 index 000000000..8a280a2c0 --- /dev/null +++ b/crates/tinyagents-harness-tools/src/net/network/gate.rs @@ -0,0 +1,223 @@ +//! The host's say over every network tool call. +//! +//! The tools in this module reach the internet, but they do not decide whether +//! they are *allowed* to. Whether the agent may act at all, how many actions it +//! has left this hour, whether an outbound request needs a human's approval, +//! whether a privacy mode forbids leaving the device, and how the process-wide +//! proxy is applied to a client all belong to the host whose threat model and +//! configuration they depend on. A tool asks a [`NetGate`] instead. +//! +//! Like [`crate::fs::filesystem::FsGate`], the trait carries no decision *types* of +//! its own, only booleans, strings and the `reqwest` builder the tool was +//! already holding, so a host maps its own policy onto it without translating +//! a vocabulary. + +/// What these tools call themselves on the wire. +/// +/// reqwest sends no `User-Agent` unless told to, and a missing one is not a +/// cosmetic omission. GitHub's REST API refuses the request outright: +/// +/// ```text +/// 403 Request forbidden by administrative rules. +/// Please make sure your request has a User-Agent header +/// ``` +/// +/// Reproducible on demand — the same URL in the same second answers 403 with +/// no header and 200 with one — so every `api.github.com` call through these +/// tools failed, always, and the 403 was then read as a credentials problem. +/// Several other APIs require one too, and anonymous traffic is the first a +/// rate limiter penalises. +/// +/// Identifying rather than disguised: a server that wants to throttle or block +/// this traffic should be able to name it. +pub(super) const USER_AGENT: &str = concat!("tinytools/", env!("CARGO_PKG_VERSION")); + +/// The host policy a network tool consults before it acts. +/// +/// Implementations must be cheap to call: the tools ask on every invocation. +/// All methods take `&self`; an implementation that counts actions keeps that +/// state behind interior mutability. +pub trait NetGate: std::fmt::Debug + Send + Sync { + /// Whether the host permits the agent to act at all right now. + /// + /// `false` blocks the tool before it touches the network. + fn can_act(&self) -> bool; + + /// Whether the action budget is already spent, without consuming any of it. + fn is_rate_limited(&self) -> bool; + + /// Consume one unit of the action budget. + /// + /// Returns `false` when the call is over budget and must be refused. + fn record_action(&self) -> bool; + + /// Whether an outbound request must be confirmed by a human before it runs. + /// + /// Drives the tools' `external_effect_with_args` flag, which routes the + /// call through the host's approval flow. + fn network_needs_approval(&self) -> bool; + + /// The refusal message when the host's privacy mode forbids contacting + /// `host`, or `None` when the request may proceed. + /// + /// Asked before URL validation and DNS, so a refused request leaks nothing, + /// not even a lookup. + fn local_only_block(&self, host: &str) -> Option; + + /// Tell the host that a request is about to leave the device for `host`. + /// + /// Observe-only: called after enforcement has already let the request + /// through. `has_body` is true when a request body rides along and + /// `has_headers` when caller-supplied headers do. + fn disclose(&self, host: &str, has_body: bool, has_headers: bool); + + /// Apply the host's process-wide proxy settings for `service` to a client + /// builder the tool has already configured. + fn prepare_client( + &self, + service: &str, + builder: reqwest::ClientBuilder, + ) -> reqwest::ClientBuilder; + + /// A ready-made client for `service` with the given total and connect + /// timeouts, proxied and TLS-configured the way the host does it. + fn timeout_client( + &self, + service: &str, + timeout_secs: u64, + connect_timeout_secs: u64, + ) -> reqwest::Client; + + /// Let the host authorize the destination of a request and say which socket + /// addresses the connection may use. + /// + /// Called after the lexical checks (scheme, allowlist, literal private + /// host) and before any DNS or connection. The answer decides how the tool + /// proceeds: + /// + /// * `None` (the default): the host takes no part, and the tool resolves and + /// vets the host itself ([`validate_url_with_dns_check`]). + /// * `Some(Ok(addrs))`: the host authorized the destination. The tool + /// connects only to `addrs` (no second DNS lookup) while keeping the URL's + /// host for the `Host` header, SNI and certificate checks. Addresses that + /// are private or local are refused regardless, so a host cannot widen the + /// tool beyond the public internet. + /// * `Some(Err(reason))`: the host refused; the tool returns `reason`. + /// + /// A host that already holds a decision engine for destinations (a security + /// module that resolves, applies an address floor and pins) answers here so + /// there is one URL guard, not one per tool. + /// + /// [`validate_url_with_dns_check`]: tinyagents_harness::url_guard::validate_url_with_dns_check + fn authorize_destination<'a>( + &'a self, + destination: &'a Destination<'a>, + ) -> DestinationFuture<'a> { + let _ = destination; + Box::pin(std::future::ready(None)) + } +} + +/// A request a tool is about to make, as the host sees it when it authorizes +/// the destination. Borrowed from the call; nothing here is retained. +#[derive(Debug, Clone, Copy)] +pub struct Destination<'a> { + /// The URL that passed the lexical checks. + pub url: &'a str, + /// Its host (a name or an IPv4 literal). + pub host: &'a str, + /// Its port (the scheme default when the URL names none). + pub port: u16, + /// HTTP method. + pub method: &'a str, + /// Caller-supplied headers. + pub headers: &'a [(String, String)], + /// Request body, if any. + pub body: Option<&'a str>, +} + +/// The host's answer to [`NetGate::authorize_destination`]; see there. +pub type DestinationFuture<'a> = std::pin::Pin< + Box< + dyn std::future::Future, String>>> + + Send + + 'a, + >, +>; + +/// Lexically validate `raw_url`, then resolve it: through the gate when the +/// host authorizes destinations, otherwise with the shared DNS check. The +/// returned addresses are the only ones the tool may connect to. +pub(super) async fn validate_destination( + gate: &dyn NetGate, + raw_url: &str, + allowed_domains: &[String], + method: &str, + headers: &[(String, String)], + body: Option<&str>, +) -> anyhow::Result { + use tinyagents_harness::url_guard::{ + ValidatedUrl, extract_host, extract_port, is_private_or_local_host, validate_url, + validate_url_with_dns_check, + }; + let url = validate_url(raw_url, allowed_domains)?; + let host = extract_host(&url)?; + let port = extract_port(&url)?; + let destination = Destination { + url: &url, + host: &host, + port, + method, + headers, + body, + }; + match gate.authorize_destination(&destination).await { + None => validate_url_with_dns_check(raw_url, allowed_domains).await, + Some(Err(reason)) => anyhow::bail!("{reason}"), + Some(Ok(addrs)) => { + if addrs.is_empty() { + anyhow::bail!("the host authorized '{host}' without an address"); + } + if let Some(blocked) = addrs + .iter() + .find(|addr| is_private_or_local_host(&addr.ip().to_string())) + { + anyhow::bail!( + "Blocked local/private address {} for '{host}'", + blocked.ip() + ); + } + Ok(ValidatedUrl { url, host, addrs }) + } + } +} + +/// The host of `url`, or `"unknown"` when it does not parse or has none. +/// +/// The value handed to [`NetGate::local_only_block`] and +/// [`NetGate::disclose`]. +pub(crate) fn host_of(url: &str) -> String { + reqwest::Url::parse(url) + .ok() + .and_then(|u| u.host_str().map(str::to_string)) + .unwrap_or_else(|| "unknown".to_string()) +} + +/// Response-size and timeout limits a tool falls back to when it is +/// constructed with a `0` (or absent) value. +/// +/// A `0` byte cap truncates every body to nothing and a `0` second timeout +/// fails every request instantly, so neither is ever a meaningful limit; it +/// only shows up when a stale configuration slips through. The host owns what +/// the sensible default is and passes it in. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct HttpLimits { + /// Response-body cap, in bytes. + pub max_response_size: usize, + /// Whole-request timeout, in seconds. + pub timeout_secs: u64, +} + +#[cfg(test)] +#[path = "gate_tests.rs"] +mod tests; diff --git a/crates/tinyagents-harness-tools/src/net/network/gate_tests.rs b/crates/tinyagents-harness-tools/src/net/network/gate_tests.rs new file mode 100644 index 000000000..1ac7a718e --- /dev/null +++ b/crates/tinyagents-harness-tools/src/net/network/gate_tests.rs @@ -0,0 +1,183 @@ +use super::*; +use std::net::SocketAddr; +use std::sync::Mutex; + +/// What the gate was asked: url, host, port, method, headers, body. +type Asked = ( + String, + String, + u16, + String, + Vec<(String, String)>, + Option, +); + +/// A gate that answers `authorize_destination` with a canned result and keeps +/// what it was asked. +#[derive(Debug)] +struct Authorizing { + answer: Option, String>>, + seen: Mutex>, +} + +impl Authorizing { + fn new(answer: Option, String>>) -> Self { + Self { + answer, + seen: Mutex::new(Vec::new()), + } + } +} + +impl NetGate for Authorizing { + fn can_act(&self) -> bool { + true + } + fn is_rate_limited(&self) -> bool { + false + } + fn record_action(&self) -> bool { + true + } + fn network_needs_approval(&self) -> bool { + false + } + fn local_only_block(&self, _host: &str) -> Option { + None + } + fn disclose(&self, _host: &str, _has_body: bool, _has_headers: bool) {} + fn prepare_client( + &self, + _service: &str, + builder: reqwest::ClientBuilder, + ) -> reqwest::ClientBuilder { + builder + } + fn timeout_client( + &self, + _service: &str, + _timeout_secs: u64, + _connect_timeout_secs: u64, + ) -> reqwest::Client { + reqwest::Client::new() + } + fn authorize_destination<'a>( + &'a self, + destination: &'a Destination<'a>, + ) -> DestinationFuture<'a> { + if let Ok(mut seen) = self.seen.lock() { + seen.push(( + destination.url.to_string(), + destination.host.to_string(), + destination.port, + destination.method.to_string(), + destination.headers.to_vec(), + destination.body.map(str::to_string), + )); + } + Box::pin(std::future::ready(self.answer.clone())) + } +} + +fn public() -> SocketAddr { + "93.184.216.34:443".parse().unwrap() +} + +#[tokio::test] +async fn a_gate_that_takes_no_part_leaves_the_tools_own_check_in_force() { + // The trait default is `None`, so the shared DNS check runs; an IP literal + // needs no lookup and is returned as the single pinned address. + let gate = Authorizing::new(None); + let target = validate_destination(&gate, "http://93.184.216.34/x", &[], "GET", &[], None) + .await + .unwrap(); + assert_eq!(target.addrs, vec!["93.184.216.34:80".parse().unwrap()]); + let blocked = validate_destination(&gate, "http://169.254.169.254/", &[], "GET", &[], None) + .await + .unwrap_err(); + assert!(blocked.to_string().contains("Blocked"), "{blocked}"); +} + +#[tokio::test] +async fn an_authorizing_gate_supplies_the_addresses_and_sees_the_request() { + let gate = Authorizing::new(Some(Ok(vec![public()]))); + let headers = vec![("X-A".to_string(), "1".to_string())]; + let target = validate_destination( + &gate, + "https://example.test/path?q=1", + &[], + "POST", + &headers, + Some("payload"), + ) + .await + .unwrap(); + assert_eq!(target.url, "https://example.test/path?q=1"); + assert_eq!(target.host, "example.test"); + assert_eq!(target.addrs, vec![public()]); + let seen = gate.seen.lock().unwrap(); + assert_eq!( + seen[0], + ( + "https://example.test/path?q=1".to_string(), + "example.test".to_string(), + 443, + "POST".to_string(), + headers, + Some("payload".to_string()) + ) + ); +} + +#[tokio::test] +async fn the_gates_refusal_reason_reaches_the_caller() { + let gate = Authorizing::new(Some(Err("egress denied: Privacy".to_string()))); + let error = validate_destination(&gate, "https://example.test/", &[], "GET", &[], None) + .await + .unwrap_err(); + assert_eq!(error.to_string(), "egress denied: Privacy"); +} + +#[tokio::test] +async fn lexical_checks_run_before_the_gate_is_asked() { + let gate = Authorizing::new(Some(Ok(vec![public()]))); + for url in [ + "ftp://example.test/", + "http://127.0.0.1/", + "http://localhost/", + ] { + assert!( + validate_destination(&gate, url, &[], "GET", &[], None) + .await + .is_err(), + "{url}" + ); + } + let allowed = vec!["allowed.test".to_string()]; + assert!( + validate_destination(&gate, "https://other.test/", &allowed, "GET", &[], None) + .await + .is_err() + ); + assert!( + gate.seen.lock().unwrap().is_empty(), + "the gate must not see a URL the tool already refused" + ); +} + +#[tokio::test] +async fn a_gate_cannot_widen_the_tool_to_private_or_empty_answers() { + for bad in [ + vec!["127.0.0.1:443".parse().unwrap()], + vec!["10.0.0.5:443".parse().unwrap()], + vec![public(), "169.254.169.254:443".parse().unwrap()], + vec![], + ] { + let gate = Authorizing::new(Some(Ok(bad))); + assert!( + validate_destination(&gate, "https://example.test/", &[], "GET", &[], None) + .await + .is_err() + ); + } +} diff --git a/crates/tinyagents-harness-tools/src/net/network/http_request.rs b/crates/tinyagents-harness-tools/src/net/network/http_request.rs new file mode 100644 index 000000000..7e7ef1003 --- /dev/null +++ b/crates/tinyagents-harness-tools/src/net/network/http_request.rs @@ -0,0 +1,483 @@ +use super::gate::{HttpLimits, NetGate, host_of}; +use async_trait::async_trait; +use serde_json::json; +use std::sync::Arc; +use std::time::Duration; +use tinytools::{PermissionLevel, Tool, ToolResult}; + +use tinyagents_harness::url_guard::{ValidatedUrl, normalize_allowed_domains}; + +use super::gate::validate_destination; + +/// What a [`PaymentHook`] hands back to retry a `402 Payment Required`. +pub struct PaymentAttempt { + /// Extra request headers that carry the payment (added to the original set). + pub headers: Vec<(String, String)>, + /// Called once with how the retried request came out, so the host can + /// settle whatever ledger entry it opened when it produced `headers`. + pub settle: Box, +} + +impl std::fmt::Debug for PaymentAttempt { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("PaymentAttempt") + .field("headers", &self.headers.len()) + .finish_non_exhaustive() + } +} + +/// How the request retried with payment headers came out. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct PaymentOutcome { + /// HTTP status of the retried response. + pub status: u16, + /// Whether that status is a success. + pub success: bool, + /// The retried response's `PAYMENT-RESPONSE` header, verbatim, when it had one. + pub payment_response: Option, +} + +/// A host's answer to a `402 Payment Required` that names a payment scheme. +/// +/// The tool sees the 402, asks the hook to pay, and retries once with the +/// headers it returns. What paying means (a wallet, a ledger, a budget) is +/// entirely the host's. +#[async_trait] +pub trait PaymentHook: std::fmt::Debug + Send + Sync { + /// Produce the payment headers for a 402 from `url`. + /// + /// `response_headers` are the 402 response's headers, which carry the + /// payment requirements. + /// + /// # Errors + /// + /// The message returned to the model as the tool error. + async fn pay( + &self, + url: &str, + response_headers: &reqwest::header::HeaderMap, + ) -> Result; +} + +/// HTTP request tool for API interactions. +/// Supports GET, POST, PUT, DELETE methods with configurable security. +#[derive(Debug)] +pub struct HttpRequestTool { + gate: Arc, + allowed_domains: Vec, + max_response_size: usize, + timeout_secs: u64, + payment: Option>, +} + +impl HttpRequestTool { + /// An `http_request` tool over `gate`. + /// + /// A `0` `max_response_size` or `timeout_secs` falls back to `defaults`. + pub fn new( + gate: Arc, + allowed_domains: Vec, + max_response_size: usize, + timeout_secs: u64, + defaults: HttpLimits, + ) -> Self { + // Treat `0` as "use default": a 0-byte cap or 0-second timeout is never + // a meaningful limit, only a footgun (see migration 5→6). Pull the + // fallbacks from `HttpRequestConfig::default()` so the tool, the schema + // default, and the migration share one source and can't drift. A `0` + // here means a stale/invalid config slipped past the migration, so + // surface it with a stable, grep-friendly, non-sensitive log line. + let max_response_size = if max_response_size == 0 { + log::warn!( + "[tool.http_request] coercing invalid limit field=max_response_size \ + from=0 to={} (stale/invalid config — see migration 5→6)", + defaults.max_response_size + ); + defaults.max_response_size + } else { + max_response_size + }; + let timeout_secs = if timeout_secs == 0 { + log::warn!( + "[tool.http_request] coercing invalid limit field=timeout_secs \ + from=0 to={} (stale/invalid config — see migration 5→6)", + defaults.timeout_secs + ); + defaults.timeout_secs + } else { + timeout_secs + }; + Self { + gate, + allowed_domains: normalize_allowed_domains(allowed_domains), + max_response_size, + timeout_secs, + payment: None, + } + } + + /// Answer a `402 Payment Required` through `hook` instead of returning it unpaid. + #[must_use] + pub fn with_payment_hook(mut self, hook: Arc) -> Self { + self.payment = Some(hook); + self + } + + async fn validate_url( + &self, + raw_url: &str, + method: &str, + headers: &[(String, String)], + body: Option<&str>, + ) -> anyhow::Result { + validate_destination( + &*self.gate, + raw_url, + &self.allowed_domains, + method, + headers, + body, + ) + .await + } + + fn validate_method(&self, method: &str) -> anyhow::Result { + match method.to_uppercase().as_str() { + "GET" => Ok(reqwest::Method::GET), + "POST" => Ok(reqwest::Method::POST), + "PUT" => Ok(reqwest::Method::PUT), + "DELETE" => Ok(reqwest::Method::DELETE), + "PATCH" => Ok(reqwest::Method::PATCH), + "HEAD" => Ok(reqwest::Method::HEAD), + "OPTIONS" => Ok(reqwest::Method::OPTIONS), + _ => anyhow::bail!( + "Unsupported HTTP method: {method}. Supported: GET, POST, PUT, DELETE, PATCH, HEAD, OPTIONS" + ), + } + } + + fn parse_headers(&self, headers: &serde_json::Value) -> Vec<(String, String)> { + let mut result = Vec::new(); + if let Some(obj) = headers.as_object() { + for (key, value) in obj { + if let Some(str_val) = value.as_str() { + result.push((key.clone(), str_val.to_string())); + } + } + } + result + } + + #[allow(dead_code)] + fn redact_headers_for_display(headers: &[(String, String)]) -> Vec<(String, String)> { + headers + .iter() + .map(|(key, value)| { + let lower = key.to_lowercase(); + let is_sensitive = lower.contains("authorization") + || lower.contains("api-key") + || lower.contains("apikey") + || lower.contains("token") + || lower.contains("secret"); + if is_sensitive { + (key.clone(), "***REDACTED***".into()) + } else { + (key.clone(), value.clone()) + } + }) + .collect() + } + + fn is_safe_response_header(name: &str) -> bool { + matches!( + name.to_ascii_lowercase().as_str(), + "content-type" + | "content-length" + | "retry-after" + | "x-ratelimit-limit" + | "x-ratelimit-remaining" + | "x-ratelimit-reset" + ) + } + + async fn execute_request( + &self, + target: &ValidatedUrl, + method: reqwest::Method, + headers: Vec<(String, String)>, + body: Option<&str>, + ) -> anyhow::Result { + let builder = reqwest::Client::builder() + .timeout(Duration::from_secs(self.timeout_secs)) + .connect_timeout(Duration::from_secs(10)) + .user_agent(super::gate::USER_AGENT) + .redirect(reqwest::redirect::Policy::none()); + let builder = self.gate.prepare_client("tool.http_request", builder); + // Connect only to the vetted addresses; the URL keeps its host name so + // `Host`, SNI and certificate checks are unchanged. + let builder = builder.resolve_to_addrs(&target.host, &target.addrs); + let client = builder.build()?; + + let mut request = client.request(method, &target.url); + + for (key, value) in headers { + request = request.header(&key, &value); + } + + if let Some(body_str) = body { + request = request.body(body_str.to_string()); + } + + Ok(request.send().await?) + } + + /// Retry a `402 Payment Required` once with the headers `hook` produced. + async fn handle_payment_required( + &self, + hook: &Arc, + initial_response: reqwest::Response, + target: &ValidatedUrl, + method: reqwest::Method, + headers: Vec<(String, String)>, + body: Option<&str>, + ) -> Result { + log::debug!( + "[tool.http_request] 402 received with PAYMENT-REQUIRED, attempting x402 payment for {}", + target.url + ); + + let attempt = hook.pay(&target.url, initial_response.headers()).await?; + + let mut retry_headers = headers; + retry_headers.extend(attempt.headers); + + let response = self + .execute_request(target, method, retry_headers, body) + .await + .map_err(|e| format!("x402 retry request failed: {e}"))?; + + (attempt.settle)(PaymentOutcome { + status: response.status().as_u16(), + success: response.status().is_success(), + payment_response: response + .headers() + .get("PAYMENT-RESPONSE") + .and_then(|v| v.to_str().ok()) + .map(str::to_string), + }); + + Ok(response) + } + + async fn format_response(&self, response: reqwest::Response) -> anyhow::Result { + let status = response.status(); + let status_code = status.as_u16(); + + // Name *and* value. This printed the name twice — `{}: {:?}` over + // `(k.as_str(), k.as_str())` — so every line read + // `x-ratelimit-remaining: "x-ratelimit-remaining"` and the header block + // carried no information at all. The redaction beside it only makes + // sense if values were meant to be shown, which is the reading taken + // here: rate-limit and retry-after headers are exactly what a caller + // needs on the failures this block renders. + let headers_text = response + .headers() + .iter() + .map(|(name, value)| { + // Response headers are untrusted and service-specific headers + // can carry credentials under arbitrary names. Only expose + // values from this small set of useful diagnostic headers. + if Self::is_safe_response_header(name.as_str()) { + format!( + "{}: {}", + name.as_str(), + value.to_str().unwrap_or("") + ) + } else { + format!("{}: ***REDACTED***", name.as_str()) + } + }) + .collect::>() + .join(", "); + + let response_text = match response.text().await { + Ok(text) => self.truncate_response(&text), + Err(e) => format!("[Failed to read response body: {e}]"), + }; + + let output = format!( + "Status: {} {}\nResponse Headers: {}\n\nResponse Body:\n{}", + status_code, + status.canonical_reason().unwrap_or("Unknown"), + headers_text, + response_text + ); + + if status.is_success() { + Ok(ToolResult::success(output)) + } else { + // The same `output` on both arms. This returned the bare string + // `HTTP 403` — twenty characters — having already built the status + // line, headers and body and then dropped them, so the one part of + // the response that said *why* never reached the caller. A real + // case: GitHub's 403 names the missing `User-Agent` header outright, + // with a documentation link, and all of it was discarded. + Ok(ToolResult::error(output)) + } + } + + fn truncate_response(&self, text: &str) -> String { + if text.len() > self.max_response_size { + let mut truncated = text + .chars() + .take(self.max_response_size) + .collect::(); + truncated.push_str("\n\n... [Response truncated due to size limit] ..."); + truncated + } else { + text.to_string() + } + } +} + +#[async_trait] +impl Tool for HttpRequestTool { + fn name(&self) -> &'static str { + "http_request" + } + + fn description(&self) -> &'static str { + "Make HTTP requests to external APIs. Supports GET, POST, PUT, DELETE, PATCH, HEAD, OPTIONS methods. \ + Security constraints: allowlist-only domains, no local/private hosts, configurable timeout and response size limits." + } + + fn parameters_schema(&self) -> serde_json::Value { + json!({ + "type": "object", + "properties": { + "url": { + "type": "string", + "description": "HTTP or HTTPS URL to request" + }, + "method": { + "type": "string", + "description": "HTTP method (GET, POST, PUT, DELETE, PATCH, HEAD, OPTIONS)", + "default": "GET" + }, + "headers": { + "type": "object", + "description": "Optional HTTP headers as key-value pairs (e.g., {\"Authorization\": \"Bearer token\", \"Content-Type\": \"application/json\"})", + "default": {} + }, + "body": { + "type": "string", + "description": "Optional request body (for POST, PUT, PATCH requests)" + } + }, + "required": ["url"] + }) + } + + /// Rich HTTP semantics (methods, headers, request bodies, and x402 retry) + /// are the same Network-class risk as `curl`: read-only autonomy is blocked + /// in `execute`, and supervised/full tiers route through ApprovalGate. + fn external_effect_with_args(&self, _args: &serde_json::Value) -> bool { + self.gate.network_needs_approval() + } + + fn permission_level(&self) -> PermissionLevel { + PermissionLevel::Write + } + + async fn execute(&self, args: serde_json::Value) -> anyhow::Result { + let url = args + .get("url") + .and_then(|v| v.as_str()) + .ok_or_else(|| anyhow::anyhow!("Missing 'url' parameter"))?; + + let method_str = args.get("method").and_then(|v| v.as_str()).unwrap_or("GET"); + let headers_val = args.get("headers").cloned().unwrap_or(json!({})); + let body = args.get("body").and_then(|v| v.as_str()); + + if !self.gate.can_act() { + return Ok(ToolResult::error( + "[policy-blocked] Action blocked: autonomy is read-only", + )); + } + + if !self.gate.record_action() { + return Ok(ToolResult::error("Action blocked: rate limit exceeded")); + } + + // Local-only enforcement (privacy epic S7, #4441): mirror the read-only + // `can_act()` deny above — under LocalOnly, refuse the outbound request + // before URL validation / DNS so nothing (not even a DNS lookup for the + // host) leaves the device. The post-validation `emit_external_transfer` + // below stays the S2 disclosure point for permitted requests. + { + let host = host_of(url); + if let Some(msg) = self.gate.local_only_block(&host) { + return Ok(ToolResult::error(msg)); + } + } + + let request_headers = self.parse_headers(&headers_val); + let target = match self + .validate_url(url, method_str, &request_headers, body) + .await + { + Ok(v) => v, + Err(e) => return Ok(ToolResult::error(e.to_string())), + }; + let url = target.url.clone(); + + // Egress spine (privacy epic S2, #4436): an agent-driven HTTP request to + // an allowlisted host leaves the device — disclose the destination and + // everything that rides with it (body + custom headers) before the + // round-trip. + { + let has_headers = headers_val.as_object().is_some_and(|h| !h.is_empty()); + self.gate + .disclose(&host_of(&url), body.is_some(), has_headers); + } + + let method = match self.validate_method(method_str) { + Ok(m) => m, + Err(e) => return Ok(ToolResult::error(e.to_string())), + }; + + let response = match self + .execute_request(&target, method.clone(), request_headers.clone(), body) + .await + { + Ok(r) => r, + Err(e) => return Ok(ToolResult::error(format!("HTTP request failed: {e}"))), + }; + + // Payment-required: if the server returns 402 with a PAYMENT-REQUIRED + // header and the host installed a hook, attempt to pay and retry. With + // no hook a 402 passes through unpaid. + let response = match &self.payment { + Some(hook) + if response.status() == reqwest::StatusCode::PAYMENT_REQUIRED + && (response.headers().get("PAYMENT-REQUIRED").is_some() + || response.headers().get("X-PAYMENT-REQUIRED").is_some()) => + { + match self + .handle_payment_required(hook, response, &target, method, request_headers, body) + .await + { + Ok(paid_response) => paid_response, + Err(msg) => return Ok(ToolResult::error(msg)), + } + } + _ => response, + }; + + self.format_response(response).await + } +} + +#[cfg(test)] +#[path = "http_request_tests.rs"] +mod tests; diff --git a/crates/tinyagents-harness-tools/src/net/network/http_request_tests.rs b/crates/tinyagents-harness-tools/src/net/network/http_request_tests.rs new file mode 100644 index 000000000..eb0130779 --- /dev/null +++ b/crates/tinyagents-harness-tools/src/net/network/http_request_tests.rs @@ -0,0 +1,672 @@ +use super::*; +use crate::net::network::test_support::{ + AutonomyLevel, DEFAULT_LIMITS, TestNetGate, pinned_to_local, +}; + +fn test_tool(allowed_domains: Vec<&str>) -> HttpRequestTool { + let security = TestNetGate::supervised(); + HttpRequestTool::new( + security, + allowed_domains.into_iter().map(String::from).collect(), + 1_000_000, + 30, + DEFAULT_LIMITS, + ) +} + +#[test] +fn zero_limits_fall_back_to_defaults() { + // Stale configs (or a bad write) can pass 0; a 0-second timeout fails + // every request and a 0-byte cap truncates every body. The constructor + // must coerce both to the module defaults — never let 0 reach reqwest. + let security = TestNetGate::supervised(); + let tool = HttpRequestTool::new( + security, + vec!["example.com".to_string()], + 0, + 0, + DEFAULT_LIMITS, + ); + let defaults = DEFAULT_LIMITS; + assert_eq!(tool.max_response_size, defaults.max_response_size); + assert_eq!(tool.timeout_secs, defaults.timeout_secs); + assert_ne!(tool.timeout_secs, 0); + assert_ne!(tool.max_response_size, 0); +} + +#[test] +fn nonzero_limits_are_preserved() { + let security = TestNetGate::supervised(); + let tool = HttpRequestTool::new( + security, + vec!["example.com".to_string()], + 2048, + 12, + DEFAULT_LIMITS, + ); + assert_eq!(tool.max_response_size, 2048); + assert_eq!(tool.timeout_secs, 12); +} + +#[test] +fn validate_accepts_valid_methods() { + let tool = test_tool(vec!["example.com"]); + assert!(tool.validate_method("GET").is_ok()); + assert!(tool.validate_method("POST").is_ok()); + assert!(tool.validate_method("PUT").is_ok()); + assert!(tool.validate_method("DELETE").is_ok()); + assert!(tool.validate_method("PATCH").is_ok()); + assert!(tool.validate_method("HEAD").is_ok()); + assert!(tool.validate_method("OPTIONS").is_ok()); +} + +#[test] +fn validate_rejects_invalid_method() { + let tool = test_tool(vec!["example.com"]); + let err = tool.validate_method("INVALID").unwrap_err().to_string(); + assert!(err.contains("Unsupported HTTP method")); +} + +#[tokio::test] +async fn validate_url_rejects_disallowed_domain() { + let tool = test_tool(vec!["example.com"]); + let err = tool + .validate_url("https://evil.test/path", "GET", &[], None) + .await + .unwrap_err() + .to_string(); + assert!(err.contains("allowed websites")); +} + +#[tokio::test] +async fn execute_blocks_readonly_mode() { + let security = TestNetGate::with(AutonomyLevel::ReadOnly, 100); + let tool = HttpRequestTool::new( + security, + vec!["example.com".into()], + 1_000_000, + 30, + DEFAULT_LIMITS, + ); + let result = tool + .execute(json!({"url": "https://example.com"})) + .await + .unwrap(); + assert!(result.is_error); + assert!(result.output().contains("read-only")); +} + +#[tokio::test] +async fn execute_blocks_when_rate_limited() { + let security = TestNetGate::with(AutonomyLevel::Supervised, 0); + let tool = HttpRequestTool::new( + security, + vec!["example.com".into()], + 1_000_000, + 30, + DEFAULT_LIMITS, + ); + let result = tool + .execute(json!({"url": "https://example.com"})) + .await + .unwrap(); + assert!(result.is_error); + assert!(result.output().contains("rate limit")); +} + +#[tokio::test] +async fn execute_blocked_under_local_only_privacy_mode() { + // Privacy epic S7 (#4441): under LocalOnly the request is refused with a + // `[policy-blocked]` result before URL validation / network. + let tool = HttpRequestTool::new( + TestNetGate::local_only(), + vec!["example.com".into()], + 1_000_000, + 30, + DEFAULT_LIMITS, + ); + let result = tool + .execute(json!({"url": "https://example.com"})) + .await + .unwrap(); + assert!(result.is_error); + assert!( + result.output().contains("[policy-blocked]"), + "got: {}", + result.output() + ); + assert!( + result.output().contains("Local-only"), + "got: {}", + result.output() + ); +} + +#[test] +fn truncate_response_within_limit() { + let tool = test_tool(vec!["example.com"]); + let text = "hello world"; + assert_eq!(tool.truncate_response(text), "hello world"); +} + +#[test] +fn truncate_response_over_limit() { + let tool = HttpRequestTool::new( + TestNetGate::supervised(), + vec!["example.com".into()], + 10, + 30, + DEFAULT_LIMITS, + ); + let text = "hello world this is long"; + let truncated = tool.truncate_response(text); + assert!(truncated.len() <= 10 + 60); + assert!(truncated.contains("[Response truncated")); +} + +#[test] +fn parse_headers_preserves_original_values() { + let tool = test_tool(vec!["example.com"]); + let headers = json!({ + "Authorization": "Bearer secret", + "Content-Type": "application/json", + "X-API-Key": "my-key" + }); + let parsed = tool.parse_headers(&headers); + assert_eq!(parsed.len(), 3); + assert!( + parsed + .iter() + .any(|(k, v)| k == "Authorization" && v == "Bearer secret") + ); + assert!( + parsed + .iter() + .any(|(k, v)| k == "X-API-Key" && v == "my-key") + ); + assert!( + parsed + .iter() + .any(|(k, v)| k == "Content-Type" && v == "application/json") + ); +} + +#[test] +fn redact_headers_for_display_redacts_sensitive() { + let headers = vec![ + ("Authorization".into(), "Bearer secret".into()), + ("Content-Type".into(), "application/json".into()), + ("X-API-Key".into(), "my-key".into()), + ("X-Secret-Token".into(), "tok-123".into()), + ]; + let redacted = HttpRequestTool::redact_headers_for_display(&headers); + assert_eq!(redacted.len(), 4); + assert!( + redacted + .iter() + .any(|(k, v)| k == "Authorization" && v == "***REDACTED***") + ); + assert!( + redacted + .iter() + .any(|(k, v)| k == "X-API-Key" && v == "***REDACTED***") + ); + assert!( + redacted + .iter() + .any(|(k, v)| k == "X-Secret-Token" && v == "***REDACTED***") + ); + assert!( + redacted + .iter() + .any(|(k, v)| k == "Content-Type" && v == "application/json") + ); +} + +#[test] +fn redact_headers_does_not_alter_original() { + let headers = vec![("Authorization".into(), "Bearer real-token".into())]; + let _ = HttpRequestTool::redact_headers_for_display(&headers); + assert_eq!(headers[0].1, "Bearer real-token"); +} + +#[test] +fn only_explicitly_safe_response_header_values_are_displayed() { + for safe in [ + "content-type", + "content-length", + "retry-after", + "x-ratelimit-remaining", + ] { + assert!(HttpRequestTool::is_safe_response_header(safe), "{safe}"); + assert!(HttpRequestTool::is_safe_response_header( + &safe.to_ascii_uppercase() + )); + } + + for sensitive in [ + "set-cookie", + "www-authenticate", + "proxy-authenticate", + "authentication-info", + "x-api-key", + "x-service-token", + ] { + assert!( + !HttpRequestTool::is_safe_response_header(sensitive), + "{sensitive} must be redacted" + ); + } +} + +#[test] +fn redirect_policy_is_none() { + let tool = test_tool(vec!["example.com"]); + assert_eq!(tool.name(), "http_request"); +} + +#[test] +fn supervised_http_request_is_external_effect_for_approval_gate() { + let tool = test_tool(vec!["example.com"]); + assert_eq!(tool.permission_level(), PermissionLevel::Write); + assert!(tool.external_effect_with_args(&json!({ + "url": "https://example.com/api", + "method": "POST", + "headers": { "Authorization": "Bearer token" }, + "body": "{}" + }))); +} + +#[test] +fn readonly_http_request_is_not_external_effect_because_execute_blocks() { + let security = TestNetGate::with(AutonomyLevel::ReadOnly, 100); + let tool = HttpRequestTool::new( + security, + vec!["example.com".into()], + 1_000_000, + 30, + DEFAULT_LIMITS, + ); + assert!(!tool.external_effect_with_args(&json!({ + "url": "https://example.com/api", + "method": "GET" + }))); +} + +#[tokio::test] +async fn a_request_discloses_body_and_header_presence_to_the_gate() { + // A public IPv4 literal skips DNS; the disallowed method fails after + // disclosure, so nothing is ever contacted. + let gate = TestNetGate::supervised(); + let tool = HttpRequestTool::new( + gate.clone(), + vec!["8.8.8.8".into()], + 1_000_000, + 30, + DEFAULT_LIMITS, + ); + let result = tool + .execute(json!({ + "url": "https://8.8.8.8/x", + "method": "TRACE", + "headers": {"Authorization": "Bearer t"}, + "body": "{}" + })) + .await + .unwrap(); + assert!(result.is_error); + assert!(result.output().contains("Unsupported HTTP method")); + assert_eq!(gate.disclosed(), vec![("8.8.8.8".to_string(), true, true)]); +} + +#[tokio::test] +async fn a_bare_request_discloses_neither_body_nor_headers() { + let gate = TestNetGate::supervised(); + let tool = HttpRequestTool::new( + gate.clone(), + vec!["8.8.8.8".into()], + 1_000_000, + 30, + DEFAULT_LIMITS, + ); + let _ = tool + .execute(json!({"url": "https://8.8.8.8/x", "method": "TRACE"})) + .await + .unwrap(); + assert_eq!( + gate.disclosed(), + vec![("8.8.8.8".to_string(), false, false)] + ); +} + +#[test] +fn the_host_of_an_unparseable_url_is_unknown() { + assert_eq!(super::super::gate::host_of("not a url"), "unknown"); + assert_eq!( + super::super::gate::host_of("https://api.example.com/v1"), + "api.example.com" + ); +} + +/// Serves one canned HTTP response per accepted connection, in order, and +/// records each request's raw head. Returns the bound address. +async fn serve( + responses: Vec, +) -> (std::net::SocketAddr, Arc>>) { + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let seen = Arc::new(std::sync::Mutex::new(Vec::new())); + let log = Arc::clone(&seen); + tokio::spawn(async move { + for response in responses { + let (mut socket, _) = listener.accept().await.unwrap(); + let mut buf = vec![0u8; 8192]; + let n = socket.read(&mut buf).await.unwrap(); + log.lock() + .unwrap() + .push(String::from_utf8_lossy(&buf[..n]).to_string()); + socket.write_all(response.as_bytes()).await.unwrap(); + let _ = socket.shutdown().await; + } + }); + (addr, seen) +} + +#[derive(Debug)] +struct RecordingHook { + settled: Arc>>, + fail: bool, +} + +#[async_trait] +impl PaymentHook for RecordingHook { + async fn pay( + &self, + _url: &str, + response_headers: &reqwest::header::HeaderMap, + ) -> Result { + if self.fail { + return Err("x402 payment failed: no wallet".into()); + } + assert!(response_headers.get("PAYMENT-REQUIRED").is_some()); + let settled = Arc::clone(&self.settled); + Ok(PaymentAttempt { + headers: vec![("PAYMENT-SIGNATURE".into(), "sig".into())], + settle: Box::new(move |outcome| settled.lock().unwrap().push(outcome)), + }) + } +} + +const PAYMENT_REQUIRED: &str = "HTTP/1.1 402 Payment Required\r\nPAYMENT-REQUIRED: abc\r\nContent-Length: 0\r\nConnection: close\r\n\r\n"; + +// Loopback is refused by the SSRF guard, so these drive the private request +// path directly rather than `execute`. +#[tokio::test] +async fn a_402_is_retried_once_with_the_hooks_headers_and_settled() { + let ok = "HTTP/1.1 200 OK\r\nPAYMENT-RESPONSE: resp\r\nContent-Length: 2\r\nConnection: close\r\n\r\nok"; + let (addr, seen) = serve(vec![PAYMENT_REQUIRED.to_string(), ok.to_string()]).await; + let settled = Arc::new(std::sync::Mutex::new(Vec::new())); + let hook: Arc = Arc::new(RecordingHook { + settled: Arc::clone(&settled), + fail: false, + }); + let tool = test_tool(vec![]); + let url = format!("http://{addr}/paid"); + let target = pinned_to_local(&url); + let first = tool + .execute_request( + &target, + reqwest::Method::GET, + vec![("X-A".into(), "1".into())], + None, + ) + .await + .unwrap(); + assert_eq!(first.status(), reqwest::StatusCode::PAYMENT_REQUIRED); + + let paid = tool + .handle_payment_required( + &hook, + first, + &target, + reqwest::Method::GET, + vec![("X-A".into(), "1".into())], + None, + ) + .await + .unwrap(); + assert!(paid.status().is_success()); + + let requests = seen.lock().unwrap().clone(); + assert_eq!(requests.len(), 2); + assert!( + !requests[0] + .to_ascii_lowercase() + .contains("payment-signature") + ); + assert!( + requests[1] + .to_ascii_lowercase() + .contains("payment-signature: sig") + ); + assert!(requests[1].to_ascii_lowercase().contains("x-a: 1")); + assert_eq!( + *settled.lock().unwrap(), + vec![PaymentOutcome { + status: 200, + success: true, + payment_response: Some("resp".into()), + }] + ); +} + +#[tokio::test] +async fn a_hook_failure_is_returned_as_the_error_without_a_retry() { + let (addr, seen) = serve(vec![PAYMENT_REQUIRED.to_string()]).await; + let hook: Arc = Arc::new(RecordingHook { + settled: Arc::default(), + fail: true, + }); + let tool = test_tool(vec![]); + let url = format!("http://{addr}/paid"); + let target = pinned_to_local(&url); + let first = tool + .execute_request(&target, reqwest::Method::GET, vec![], None) + .await + .unwrap(); + let err = tool + .handle_payment_required(&hook, first, &target, reqwest::Method::GET, vec![], None) + .await + .unwrap_err(); + assert_eq!(err, "x402 payment failed: no wallet"); + assert_eq!(seen.lock().unwrap().len(), 1); +} + +#[tokio::test] +async fn a_body_is_sent_and_a_success_response_is_formatted_with_redacted_cookies() { + let ok = "HTTP/1.1 200 OK\r\nSet-Cookie: session=abc\r\nWWW-Authenticate: Bearer secret-challenge\r\nX-Api-Key: api-secret\r\nX-Service-Token: service-secret\r\nX-Other: 1\r\nContent-Length: 5\r\nConnection: close\r\n\r\nhello"; + let (addr, seen) = serve(vec![ok.to_string()]).await; + let tool = test_tool(vec![]); + let response = tool + .execute_request( + &pinned_to_local(&format!("http://{addr}/post")), + reqwest::Method::POST, + vec![], + Some("payload"), + ) + .await + .unwrap(); + let result = tool.format_response(response).await.unwrap(); + assert!(!result.is_error); + let text = result.text(); + assert!(text.contains("hello"), "{text}"); + assert!(text.contains("set-cookie: ***REDACTED***"), "{text}"); + assert!(!text.contains("session=abc"), "{text}"); + for secret in ["secret-challenge", "api-secret", "service-secret"] { + assert!(!text.contains(secret), "response leaked {secret}: {text}"); + } + for header in ["www-authenticate", "x-api-key", "x-service-token"] { + assert!( + text.to_ascii_lowercase() + .contains(&format!("{header}: ***redacted***")), + "{text}" + ); + } + assert!(seen.lock().unwrap()[0].ends_with("payload")); +} + +#[tokio::test] +async fn a_non_success_response_is_reported_as_an_error() { + // The failure now carries the whole response rather than just its code. + // + // It used to answer `HTTP 500` —, and in a measured case `HTTP 403`, + // twenty characters, after building the status line, headers and body and + // then dropping them. The discarded body was the only thing that said what + // was wrong: GitHub's 403 names the missing `User-Agent` header outright. + let body = "Request forbidden by administrative rules. \ + Please make sure your request has a User-Agent header"; + let (addr, _) = serve(vec![format!( + "HTTP/1.1 403 Forbidden\r\nContent-Length: {}\r\nX-RateLimit-Remaining: 59\r\nConnection: close\r\n\r\n{}", + body.len(), + body + )]) + .await; + let tool = test_tool(vec![]); + let response = tool + .execute_request( + &pinned_to_local(&format!("http://{addr}/")), + reqwest::Method::GET, + vec![], + None, + ) + .await + .unwrap(); + let result = tool.format_response(response).await.unwrap(); + assert!(result.is_error, "a 403 is still a failure"); + let text = result.text(); + // Anchored: a caller classifying the failure reads the status from the head + // of the text, and an unreadable one gets guessed at. + assert!(text.starts_with("Status: 403"), "{text}"); + // The reason the server gave. + assert!(text.contains("User-Agent header"), "{text}"); + // Header *values*, not the name printed twice. `x-ratelimit-remaining` is + // exactly what a caller wants on the failures this block renders. + let lower = text.to_ascii_lowercase(); + assert!(lower.contains("x-ratelimit-remaining: 59"), "{text}"); +} + +#[tokio::test] +async fn an_unreadable_body_is_reported_inline() { + // Promise more bytes than are sent, then close: reading the body fails. + let (addr, _) = serve(vec![ + "HTTP/1.1 200 OK\r\nContent-Length: 100\r\nConnection: close\r\n\r\nshort".to_string(), + ]) + .await; + let tool = test_tool(vec![]); + let response = tool + .execute_request( + &pinned_to_local(&format!("http://{addr}/")), + reqwest::Method::GET, + vec![], + None, + ) + .await + .unwrap(); + let result = tool.format_response(response).await.unwrap(); + assert!(result.text().contains("Failed to read response body")); +} + +#[tokio::test] +async fn a_missing_url_is_an_error() { + let tool = test_tool(vec![]); + let err = tool.execute(json!({})).await.unwrap_err(); + assert!(err.to_string().contains("Missing 'url'")); +} + +#[test] +fn a_payment_hook_can_be_installed_and_attempts_debug_without_headers() { + let hook: Arc = Arc::new(RecordingHook { + settled: Arc::default(), + fail: false, + }); + let tool = test_tool(vec![]).with_payment_hook(hook); + assert!(tool.payment.is_some()); + let attempt = PaymentAttempt { + headers: vec![("PAYMENT-SIGNATURE".into(), "secret".into())], + settle: Box::new(|_| {}), + }; + let shown = format!("{attempt:?}"); + assert!(shown.contains("PaymentAttempt") && !shown.contains("secret")); +} + +#[test] +fn the_test_gate_builds_a_client_with_the_requested_timeouts() { + let gate = TestNetGate::supervised(); + let _client = gate.timeout_client("svc", 5, 2); +} + +/// Both network tools must name themselves on the wire. +/// +/// Not cosmetic: GitHub's REST API answers 403 to a request with no +/// `User-Agent`, so every `api.github.com` call through these tools failed +/// until this was set. `serve` records the raw request, which is the only way +/// to assert an outgoing header actually left the process. +#[tokio::test] +async fn an_outgoing_request_carries_a_user_agent() -> anyhow::Result<()> { + let (addr, seen) = serve(vec![ + "HTTP/1.1 200 OK\r\nContent-Length: 2\r\nConnection: close\r\n\r\nok".to_string(), + ]) + .await; + let tool = test_tool(vec![]); + let _ = tool + .execute_request( + &pinned_to_local(&format!("http://{addr}/")), + reqwest::Method::GET, + vec![], + None, + ) + .await?; + + let request = seen + .lock() + .map_err(|error| anyhow::anyhow!("request log mutex poisoned: {error}"))?[0] + .clone(); + let lower = request.to_ascii_lowercase(); + assert!( + lower.contains("user-agent:"), + "no User-Agent was sent:\n{request}" + ); + assert!( + lower.contains("user-agent: tinytools/"), + "the header must identify this crate:\n{request}" + ); + Ok(()) +} + +/// The connection goes to the vetted address, not to whatever the host name +/// resolves to a second time: the name below does not resolve at all, and the +/// request still reaches the pinned socket with the name in `Host`. +#[tokio::test] +async fn the_request_connects_to_the_pinned_address_and_keeps_the_host_name() { + let ok = "HTTP/1.1 200 OK\r\nContent-Length: 2\r\nConnection: close\r\n\r\nok"; + let (addr, seen) = serve(vec![ok.to_string()]).await; + let tool = test_tool(vec![]); + let target = tinyagents_harness::url_guard::ValidatedUrl { + url: format!("http://pinned.invalid:{}/x", addr.port()), + host: "pinned.invalid".to_string(), + addrs: vec![addr], + }; + let response = tool + .execute_request(&target, reqwest::Method::GET, vec![], None) + .await + .unwrap(); + assert_eq!(response.status(), reqwest::StatusCode::OK); + let request = seen.lock().unwrap().join(""); + assert!( + request + .to_ascii_lowercase() + .contains(&format!("host: pinned.invalid:{}", addr.port())), + "{request}" + ); +} diff --git a/crates/tinyagents-harness-tools/src/net/network/mod.rs b/crates/tinyagents-harness-tools/src/net/network/mod.rs new file mode 100644 index 000000000..7544fd874 --- /dev/null +++ b/crates/tinyagents-harness-tools/src/net/network/mod.rs @@ -0,0 +1,50 @@ +//! Network tools: `http_request`, `web_fetch`, `curl` and `pushover`. +//! +//! Every tool takes an `Arc`: the tool performs the I/O, the +//! host's gate decides whether it may (autonomy, action budget, approval, +//! privacy mode) and how the process-wide proxy is applied. The gate trait +//! ([`NetGate`]) is the whole seam; nothing here knows what an autonomy level, +//! an approval prompt or an egress descriptor is. URL allowlisting and SSRF +//! checks are [`tinyagents_harness::url_guard`]'s. +//! +//! Two more seams keep host-specific behavior out of the tools: a +//! [`PaymentHook`] answers a `402 Payment Required` for `http_request`, and an +//! [`HtmlExtractor`] turns pages into Markdown for `web_fetch`. + +// These tools moved here verbatim from a host crate; their behavior, messages +// and control flow are pinned by tests and by the fixtures in `fixtures/`, so +// the purely stylistic pedantic lints below are allowed rather than reshaping +// working code (long `execute` bodies, `usize as f64` size labels, and so on). +#![allow( + clippy::too_many_lines, + clippy::cast_precision_loss, + clippy::cast_possible_truncation, + clippy::similar_names, + clippy::used_underscore_binding, + clippy::manual_let_else, + clippy::format_push_string, + clippy::unwrap_used, + clippy::unused_self, + clippy::unused_async, + clippy::needless_pass_by_value, + clippy::doc_markdown, + clippy::unnecessary_wraps +)] + +mod curl; +mod gate; +mod http_request; +mod pushover; +pub mod web_fetch; + +#[cfg(test)] +#[path = "contract_tests.rs"] +mod contract_tests; +#[cfg(test)] +mod test_support; + +pub use curl::CurlTool; +pub use gate::{Destination, DestinationFuture, HttpLimits, NetGate}; +pub use http_request::{HttpRequestTool, PaymentAttempt, PaymentHook, PaymentOutcome}; +pub use pushover::PushoverTool; +pub use web_fetch::{AsyncHtmlExtractor, HtmlExtractor, WebFetchTool}; diff --git a/crates/tinyagents-harness-tools/src/net/network/pushover.rs b/crates/tinyagents-harness-tools/src/net/network/pushover.rs new file mode 100644 index 000000000..defd9bf27 --- /dev/null +++ b/crates/tinyagents-harness-tools/src/net/network/pushover.rs @@ -0,0 +1,217 @@ +use super::gate::NetGate; +use async_trait::async_trait; +use serde_json::json; +use std::path::PathBuf; +use std::sync::Arc; +use tinytools::{Tool, ToolResult}; + +const PUSHOVER_API_URL: &str = "https://api.pushover.net/1/messages.json"; +const PUSHOVER_REQUEST_TIMEOUT_SECS: u64 = 15; + +/// Sends a Pushover notification using credentials from the workspace `.env`. +#[derive(Debug)] +pub struct PushoverTool { + gate: Arc, + workspace_dir: PathBuf, +} + +impl PushoverTool { + /// A `pushover` tool over `gate`, reading `/.env`. + pub fn new(gate: Arc, workspace_dir: PathBuf) -> Self { + Self { + gate, + workspace_dir, + } + } + + fn parse_env_value(raw: &str) -> String { + let raw = raw.trim(); + + let unquoted = if raw.len() >= 2 + && ((raw.starts_with('"') && raw.ends_with('"')) + || (raw.starts_with('\'') && raw.ends_with('\''))) + { + &raw[1..raw.len() - 1] + } else { + raw + }; + + // Keep support for inline comments in unquoted values: + // KEY=value # comment + unquoted.split_once(" #").map_or_else( + || unquoted.trim().to_string(), + |(value, _)| value.trim().to_string(), + ) + } + + async fn get_credentials(&self) -> anyhow::Result<(String, String)> { + let env_path = self.workspace_dir.join(".env"); + let content = tokio::fs::read_to_string(&env_path) + .await + .map_err(|e| anyhow::anyhow!("Failed to read {}: {}", env_path.display(), e))?; + + let mut token = None; + let mut user_key = None; + + for line in content.lines() { + let line = line.trim(); + if line.starts_with('#') || line.is_empty() { + continue; + } + let line = line.strip_prefix("export ").map_or(line, str::trim); + if let Some((key, value)) = line.split_once('=') { + let key = key.trim(); + let value = Self::parse_env_value(value); + + if key.eq_ignore_ascii_case("PUSHOVER_TOKEN") { + token = Some(value); + } else if key.eq_ignore_ascii_case("PUSHOVER_USER_KEY") { + user_key = Some(value); + } + } + } + + let token = token.ok_or_else(|| anyhow::anyhow!("PUSHOVER_TOKEN not found in .env"))?; + let user_key = + user_key.ok_or_else(|| anyhow::anyhow!("PUSHOVER_USER_KEY not found in .env"))?; + + Ok((token, user_key)) + } +} + +#[async_trait] +impl Tool for PushoverTool { + fn name(&self) -> &'static str { + "pushover" + } + + fn description(&self) -> &'static str { + "Send a Pushover notification to your device. Requires PUSHOVER_TOKEN and PUSHOVER_USER_KEY in .env file." + } + + fn parameters_schema(&self) -> serde_json::Value { + json!({ + "type": "object", + "properties": { + "message": { + "type": "string", + "description": "The notification message to send" + }, + "title": { + "type": "string", + "description": "Optional notification title" + }, + "priority": { + "type": "integer", + "description": "Message priority: -2 (lowest/silent), -1 (low/no sound), 0 (normal), 1 (high), 2 (emergency/repeating)" + }, + "sound": { + "type": "string", + "description": "Notification sound override (e.g., 'pushover', 'bike', 'bugle', 'cashregister', etc.)" + } + }, + "required": ["message"] + }) + } + + // Intentionally NOT marked external_effect=true in v1. + // + // `execute()` below already enforces local policy via + // `gate.can_act()` (read-only autonomy block) and + // `gate.record_action()` (rate limit). The gate runs + // BEFORE `execute()`, so flipping external_effect would prompt + // the user even for calls that the policy will immediately + // reject — a dead-end approval flow. Leaving this tool at the + // trait default (false) means the local policy stays the + // single approval surface until the gate learns to consult + // policy before parking. + // + // Follow-up #1339-v3: thread the SecurityPolicy into the gate + // pre-check so external_effect can stay true here without + // generating ghost approvals. + + async fn execute(&self, args: serde_json::Value) -> anyhow::Result { + if !self.gate.can_act() { + return Ok(ToolResult::error( + "[policy-blocked] Action blocked: autonomy is read-only", + )); + } + + if !self.gate.record_action() { + return Ok(ToolResult::error("Action blocked: rate limit exceeded")); + } + + let message = args + .get("message") + .and_then(|v| v.as_str()) + .map(str::trim) + .filter(|v| !v.is_empty()) + .ok_or_else(|| anyhow::anyhow!("Missing 'message' parameter"))? + .to_string(); + + let title = args.get("title").and_then(|v| v.as_str()).map(String::from); + + let priority = match args.get("priority").and_then(serde_json::Value::as_i64) { + Some(value) if (-2..=2).contains(&value) => Some(value), + Some(value) => { + return Ok(ToolResult::error(format!( + "Invalid 'priority': {value}. Expected integer in range -2..=2" + ))); + } + None => None, + }; + + let sound = args.get("sound").and_then(|v| v.as_str()).map(String::from); + + let (token, user_key) = self.get_credentials().await?; + + let mut form = reqwest::multipart::Form::new() + .text("token", token) + .text("user", user_key) + .text("message", message); + + if let Some(title) = title { + form = form.text("title", title); + } + + if let Some(priority) = priority { + form = form.text("priority", priority.to_string()); + } + + if let Some(sound) = sound { + form = form.text("sound", sound); + } + + let client = self + .gate + .timeout_client("tool.pushover", PUSHOVER_REQUEST_TIMEOUT_SECS, 10); + let response = client.post(PUSHOVER_API_URL).multipart(form).send().await?; + + let status = response.status(); + let body = response.text().await.unwrap_or_default(); + + if !status.is_success() { + return Ok(ToolResult::error(format!( + "Pushover API returned status {status}" + ))); + } + + let api_status = serde_json::from_str::(&body) + .ok() + .and_then(|json| json.get("status").and_then(serde_json::Value::as_i64)); + + if api_status == Some(1) { + Ok(ToolResult::success(format!( + "Pushover notification sent successfully. Response: {body}" + ))) + } else { + Ok(ToolResult::error( + "Pushover API returned an application-level error", + )) + } + } +} + +#[cfg(test)] +#[path = "pushover_tests.rs"] +mod tests; diff --git a/crates/tinyagents-harness-tools/src/net/network/pushover_tests.rs b/crates/tinyagents-harness-tools/src/net/network/pushover_tests.rs new file mode 100644 index 000000000..a3aca412e --- /dev/null +++ b/crates/tinyagents-harness-tools/src/net/network/pushover_tests.rs @@ -0,0 +1,209 @@ +use super::*; +use crate::net::network::test_support::{AutonomyLevel, TestNetGate}; +use std::fs; +use tempfile::TempDir; + +fn test_security(level: AutonomyLevel, max_actions_per_hour: usize) -> Arc { + TestNetGate::with(level, max_actions_per_hour) +} + +#[test] +fn pushover_tool_name() { + let tool = PushoverTool::new( + test_security(AutonomyLevel::Full, 100), + PathBuf::from("/tmp"), + ); + assert_eq!(tool.name(), "pushover"); +} + +#[test] +fn pushover_tool_description() { + let tool = PushoverTool::new( + test_security(AutonomyLevel::Full, 100), + PathBuf::from("/tmp"), + ); + assert_ne!(tool.description().len(), 0); +} + +#[test] +fn pushover_tool_has_parameters_schema() { + let tool = PushoverTool::new( + test_security(AutonomyLevel::Full, 100), + PathBuf::from("/tmp"), + ); + let schema = tool.parameters_schema(); + assert_eq!(schema["type"], "object"); + assert!(schema["properties"].get("message").is_some()); +} + +#[test] +fn pushover_tool_requires_message() { + let tool = PushoverTool::new( + test_security(AutonomyLevel::Full, 100), + PathBuf::from("/tmp"), + ); + let schema = tool.parameters_schema(); + let required = schema["required"].as_array().unwrap(); + assert!(required.contains(&serde_json::Value::String("message".to_string()))); +} + +#[tokio::test] +async fn credentials_parsed_from_env_file() { + let tmp = TempDir::new().unwrap(); + let env_path = tmp.path().join(".env"); + fs::write( + &env_path, + "PUSHOVER_TOKEN=testtoken123\nPUSHOVER_USER_KEY=userkey456\n", + ) + .unwrap(); + + let tool = PushoverTool::new( + test_security(AutonomyLevel::Full, 100), + tmp.path().to_path_buf(), + ); + let result = tool.get_credentials().await; + + assert!(result.is_ok()); + let (token, user_key) = result.unwrap(); + assert_eq!(token, "testtoken123"); + assert_eq!(user_key, "userkey456"); +} + +#[tokio::test] +async fn credentials_fail_without_env_file() { + let tmp = TempDir::new().unwrap(); + let tool = PushoverTool::new( + test_security(AutonomyLevel::Full, 100), + tmp.path().to_path_buf(), + ); + let result = tool.get_credentials().await; + + assert!(result.is_err()); +} + +#[tokio::test] +async fn credentials_fail_without_token() { + let tmp = TempDir::new().unwrap(); + let env_path = tmp.path().join(".env"); + fs::write(&env_path, "PUSHOVER_USER_KEY=userkey456\n").unwrap(); + + let tool = PushoverTool::new( + test_security(AutonomyLevel::Full, 100), + tmp.path().to_path_buf(), + ); + let result = tool.get_credentials().await; + + assert!(result.is_err()); +} + +#[tokio::test] +async fn credentials_fail_without_user_key() { + let tmp = TempDir::new().unwrap(); + let env_path = tmp.path().join(".env"); + fs::write(&env_path, "PUSHOVER_TOKEN=testtoken123\n").unwrap(); + + let tool = PushoverTool::new( + test_security(AutonomyLevel::Full, 100), + tmp.path().to_path_buf(), + ); + let result = tool.get_credentials().await; + + assert!(result.is_err()); +} + +#[tokio::test] +async fn credentials_ignore_comments() { + let tmp = TempDir::new().unwrap(); + let env_path = tmp.path().join(".env"); + fs::write(&env_path, "# This is a comment\nPUSHOVER_TOKEN=realtoken\n# Another comment\nPUSHOVER_USER_KEY=realuser\n").unwrap(); + + let tool = PushoverTool::new( + test_security(AutonomyLevel::Full, 100), + tmp.path().to_path_buf(), + ); + let result = tool.get_credentials().await; + + assert!(result.is_ok()); + let (token, user_key) = result.unwrap(); + assert_eq!(token, "realtoken"); + assert_eq!(user_key, "realuser"); +} + +#[test] +fn pushover_tool_supports_priority() { + let tool = PushoverTool::new( + test_security(AutonomyLevel::Full, 100), + PathBuf::from("/tmp"), + ); + let schema = tool.parameters_schema(); + assert!(schema["properties"].get("priority").is_some()); +} + +#[test] +fn pushover_tool_supports_sound() { + let tool = PushoverTool::new( + test_security(AutonomyLevel::Full, 100), + PathBuf::from("/tmp"), + ); + let schema = tool.parameters_schema(); + assert!(schema["properties"].get("sound").is_some()); +} + +#[tokio::test] +async fn credentials_support_export_and_quoted_values() { + let tmp = TempDir::new().unwrap(); + let env_path = tmp.path().join(".env"); + fs::write( + &env_path, + "export PUSHOVER_TOKEN=\"quotedtoken\"\nPUSHOVER_USER_KEY='quoteduser'\n", + ) + .unwrap(); + + let tool = PushoverTool::new( + test_security(AutonomyLevel::Full, 100), + tmp.path().to_path_buf(), + ); + let result = tool.get_credentials().await; + + assert!(result.is_ok()); + let (token, user_key) = result.unwrap(); + assert_eq!(token, "quotedtoken"); + assert_eq!(user_key, "quoteduser"); +} + +#[tokio::test] +async fn execute_blocks_readonly_mode() { + let tool = PushoverTool::new( + test_security(AutonomyLevel::ReadOnly, 100), + PathBuf::from("/tmp"), + ); + + let result = tool.execute(json!({"message": "hello"})).await.unwrap(); + assert!(result.is_error); + assert!(result.output().contains("read-only")); +} + +#[tokio::test] +async fn execute_blocks_rate_limit() { + let tool = PushoverTool::new(test_security(AutonomyLevel::Full, 0), PathBuf::from("/tmp")); + + let result = tool.execute(json!({"message": "hello"})).await.unwrap(); + assert!(result.is_error); + assert!(result.output().contains("rate limit")); +} + +#[tokio::test] +async fn execute_rejects_priority_out_of_range() { + let tool = PushoverTool::new( + test_security(AutonomyLevel::Full, 100), + PathBuf::from("/tmp"), + ); + + let result = tool + .execute(json!({"message": "hello", "priority": 5})) + .await + .unwrap(); + + assert!(result.is_error); + assert!(result.output().contains("-2..=2")); +} diff --git a/crates/tinyagents-harness-tools/src/net/network/test_support.rs b/crates/tinyagents-harness-tools/src/net/network/test_support.rs new file mode 100644 index 000000000..bb5a52598 --- /dev/null +++ b/crates/tinyagents-harness-tools/src/net/network/test_support.rs @@ -0,0 +1,176 @@ +//! A small, self-contained [`NetGate`] for the network tools' unit tests. +//! +//! It models just enough of a host policy to exercise the tools: an autonomy +//! tier, an action budget, an approval flag, a local-only switch and a record +//! of what the tool disclosed. It is *not* a copy of any host's policy; the +//! host's own semantics are tested in the host, through its adapter. + +#![allow( + clippy::expect_used, + clippy::unwrap_used, + missing_docs, + unreachable_pub +)] + +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::{Arc, Mutex}; + +use super::gate::{HttpLimits, NetGate}; +use super::web_fetch::HtmlExtractor; + +/// Prefix the fake puts on refusals a real host would mark as policy-blocked. +pub const POLICY_BLOCKED_MARKER: &str = "[policy-blocked]"; + +/// Autonomy tier of the fake gate. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum AutonomyLevel { + /// Nothing may run. + ReadOnly, + /// Outbound requests ask for approval. + Supervised, + /// Outbound requests run unprompted. + Full, +} + +/// The defaults the tests hand to constructors that take [`HttpLimits`]. +pub const DEFAULT_LIMITS: HttpLimits = HttpLimits { + max_response_size: 1_000_000, + timeout_secs: 30, +}; + +/// A fake host policy. +#[derive(Debug)] +pub struct TestNetGate { + autonomy: AutonomyLevel, + max_actions: usize, + used: AtomicUsize, + local_only: bool, + disclosed: Mutex>, +} + +impl TestNetGate { + /// Supervised gate with an effectively unlimited budget. + pub fn supervised() -> Arc { + Self::with(AutonomyLevel::Supervised, 1_000_000) + } + + /// Gate with an explicit tier and action budget. + pub fn with(autonomy: AutonomyLevel, max_actions: usize) -> Arc { + Arc::new(Self { + autonomy, + max_actions, + used: AtomicUsize::new(0), + local_only: false, + disclosed: Mutex::new(Vec::new()), + }) + } + + /// Supervised gate whose privacy mode refuses every outbound request. + pub fn local_only() -> Arc { + Arc::new(Self { + autonomy: AutonomyLevel::Supervised, + max_actions: 1_000_000, + used: AtomicUsize::new(0), + local_only: true, + disclosed: Mutex::new(Vec::new()), + }) + } + + /// Everything the tool has disclosed so far: `(host, has_body, has_headers)`. + pub fn disclosed(&self) -> Vec<(String, bool, bool)> { + self.disclosed.lock().unwrap().clone() + } +} + +impl NetGate for TestNetGate { + fn can_act(&self) -> bool { + self.autonomy != AutonomyLevel::ReadOnly + } + + fn is_rate_limited(&self) -> bool { + self.used.load(Ordering::SeqCst) >= self.max_actions + } + + fn record_action(&self) -> bool { + self.used.fetch_add(1, Ordering::SeqCst) < self.max_actions + } + + fn network_needs_approval(&self) -> bool { + self.autonomy == AutonomyLevel::Supervised + } + + fn local_only_block(&self, host: &str) -> Option { + self.local_only.then(|| { + format!("{POLICY_BLOCKED_MARKER} Local-only privacy mode blocks contacting {host}.") + }) + } + + fn disclose(&self, host: &str, has_body: bool, has_headers: bool) { + self.disclosed + .lock() + .unwrap() + .push((host.to_string(), has_body, has_headers)); + } + + fn prepare_client( + &self, + _service: &str, + builder: reqwest::ClientBuilder, + ) -> reqwest::ClientBuilder { + builder + } + + fn timeout_client( + &self, + _service: &str, + timeout_secs: u64, + connect_timeout_secs: u64, + ) -> reqwest::Client { + reqwest::Client::builder() + .timeout(std::time::Duration::from_secs(timeout_secs)) + .connect_timeout(std::time::Duration::from_secs(connect_timeout_secs)) + .build() + .unwrap_or_default() + } +} + +/// An [`HtmlExtractor`] that recognises ` bool { + let head = body.trim_start().to_ascii_lowercase(); + head.starts_with(" String { + let mut out = String::new(); + let mut in_tag = false; + for ch in html.chars() { + match ch { + '<' => in_tag = true, + '>' => in_tag = false, + c if !in_tag => out.push(c), + _ => {} + } + } + out + } +} + +/// A [`ValidatedUrl`] for a loopback test server: the host is the IP literal in +/// `url`, pinned to that one socket address, as the SSRF guard would return it +/// for a public literal. +pub fn pinned_to_local(url: &str) -> tinyagents_harness::url_guard::ValidatedUrl { + let parsed = reqwest::Url::parse(url).expect("test URL parses"); + let host = parsed.host_str().expect("test URL has a host").to_string(); + let ip: std::net::IpAddr = host.parse().expect("test URL host is an IP literal"); + let port = parsed.port_or_known_default().expect("test URL has a port"); + tinyagents_harness::url_guard::ValidatedUrl { + url: url.to_string(), + host, + addrs: vec![std::net::SocketAddr::new(ip, port)], + } +} diff --git a/crates/tinyagents-harness-tools/src/net/network/web_fetch.rs b/crates/tinyagents-harness-tools/src/net/network/web_fetch.rs new file mode 100644 index 000000000..57644e5e6 --- /dev/null +++ b/crates/tinyagents-harness-tools/src/net/network/web_fetch.rs @@ -0,0 +1,674 @@ +//! `web_fetch` — fetch a URL and return its text body. +//! +//! Coding-harness baseline tool (issue #1205). Distinct from +//! `http_request` (full method/header surface) and `curl` (writes to +//! disk). `web_fetch` is the single-purpose "GET and read" primitive +//! the agent reaches for when researching: returns the response body +//! as text, capped, with a tiny preamble (status + final URL). +//! +//! A 4xx/5xx response is a failed fetch: it comes back as an error result +//! (`is_error`) naming the status and a short body excerpt, so a host that +//! budgets or retries on tool errors sees blocked and rate-limited pages for +//! what they are. 3xx responses are not followed and stay successful reports. + +use super::gate::{HttpLimits, NetGate, USER_AGENT, host_of}; +use tinyagents_harness::url_guard::{ValidatedUrl, normalize_allowed_domains}; + +use super::gate::validate_destination; +use async_trait::async_trait; +use serde_json::json; +use std::sync::Arc; +use std::time::Duration; +use tinytools::{PermissionLevel, Tool, ToolResult}; + +/// How `web_fetch` turns an HTML page into something a model should read. +/// +/// Content transforms are somebody else's specialty (the host wires in +/// whichever engine it has), so the tool asks instead of owning one. +pub trait HtmlExtractor: std::fmt::Debug + Send + Sync { + /// Whether `body`, which arrived with no usable `Content-Type`, is HTML. + fn looks_like_html(&self, body: &str) -> bool; + + /// Convert an HTML document to Markdown (links kept, scripts dropped). + fn to_markdown(&self, html: &str) -> String; +} + +/// An asynchronous, fallible HTML transform supplied by the host. +/// +/// Use this interface when extraction runs in a separate service or module. +/// Failures propagate to the fetch caller; the tool does not retry locally. +#[async_trait] +pub trait AsyncHtmlExtractor: std::fmt::Debug + Send + Sync { + /// Whether a response without a usable content type contains HTML. + /// + /// # Errors + /// Returns the host's detection failure. + async fn looks_like_html(&self, body: &str) -> anyhow::Result; + + /// Convert HTML to Markdown, preserving links and dropping scripts. + /// + /// # Errors + /// Returns the host's extraction failure. + async fn to_markdown(&self, html: &str) -> anyhow::Result; +} + +#[derive(Debug)] +enum HtmlProvider { + Local(Arc), + Async { + extractor: Arc, + timeout: Duration, + }, +} + +impl HtmlProvider { + async fn looks_like_html(&self, body: &str) -> anyhow::Result { + match self { + Self::Local(html) => Ok(html.looks_like_html(body)), + Self::Async { extractor, timeout } => { + tokio::time::timeout(*timeout, extractor.looks_like_html(body)) + .await + .map_err(|_| anyhow::anyhow!("HTML detection timed out"))? + } + } + } + + async fn to_markdown(&self, body: &str) -> anyhow::Result { + match self { + Self::Local(html) => Ok(html.to_markdown(body)), + Self::Async { extractor, timeout } => { + tokio::time::timeout(*timeout, extractor.to_markdown(body)) + .await + .map_err(|_| anyhow::anyhow!("HTML extraction timed out"))? + } + } + } +} + +/// Fetches a URL and returns its text body. +#[derive(Debug)] +pub struct WebFetchTool { + gate: Arc, + allowed_domains: Vec, + max_bytes: usize, + timeout_secs: u64, + html: HtmlProvider, + extra_properties: Vec<(String, serde_json::Value)>, +} + +impl WebFetchTool { + /// A `web_fetch` tool over `gate`. + /// + /// Both `None` and `Some(0)` for the limits mean "use `defaults`". + pub fn new( + gate: Arc, + allowed_domains: Vec, + max_bytes: Option, + timeout_secs: Option, + defaults: HttpLimits, + html: Arc, + ) -> Self { + Self::with_provider( + gate, + allowed_domains, + max_bytes, + timeout_secs, + defaults, + HtmlProvider::Local(html), + ) + } + + /// A `web_fetch` tool with asynchronous HTML detection and extraction. + /// + /// Limits and raw-response handling match [`Self::new`]. Extraction errors + /// propagate without invoking another provider. Both asynchronous detection + /// and conversion are bounded by the configured request timeout. + #[must_use] + pub fn new_async( + gate: Arc, + allowed_domains: Vec, + max_bytes: Option, + timeout_secs: Option, + defaults: HttpLimits, + html: Arc, + ) -> Self { + let resolved_timeout = match timeout_secs { + Some(0) => { + log::warn!( + "[tool.web_fetch] coercing invalid limit field=timeout_secs from=0 to={} \ + (stale/invalid config — see migration 5→6)", + defaults.timeout_secs + ); + defaults.timeout_secs + } + Some(seconds) => seconds, + None => defaults.timeout_secs, + }; + Self::with_provider( + gate, + allowed_domains, + max_bytes, + Some(resolved_timeout), + defaults, + HtmlProvider::Async { + extractor: html, + timeout: Duration::from_secs(resolved_timeout), + }, + ) + } + + fn with_provider( + gate: Arc, + allowed_domains: Vec, + max_bytes: Option, + timeout_secs: Option, + defaults: HttpLimits, + html: HtmlProvider, + ) -> Self { + // Treat both `None` and `Some(0)` as "use default": callers wire these + // from `[http_request]`, and a 0-byte cap truncates every body to + // nothing while a 0-second timeout fails every request instantly. + // Stale-zero configs are repaired on load (migration 5→6); this clamp + // is the always-on guard at the point of use. The fallbacks come from + // the host's `defaults` so the tool shares one source with the schema + + // migration (no cross-layer drift). `Some(0)` is a genuine + // misconfiguration, so log it (grep-friendly, no payload); a bare + // `None` is a normal "use default" call and stays quiet. + let max_bytes = match max_bytes { + Some(0) => { + log::warn!( + "[tool.web_fetch] coercing invalid limit field=max_bytes \ + from=0 to={} (stale/invalid config — see migration 5→6)", + defaults.max_response_size + ); + defaults.max_response_size + } + Some(n) => n, + None => defaults.max_response_size, + }; + let timeout_secs = match timeout_secs { + Some(0) => { + log::warn!( + "[tool.web_fetch] coercing invalid limit field=timeout_secs \ + from=0 to={} (stale/invalid config — see migration 5→6)", + defaults.timeout_secs + ); + defaults.timeout_secs + } + Some(n) => n, + None => defaults.timeout_secs, + }; + Self { + gate, + allowed_domains: normalize_allowed_domains(allowed_domains), + max_bytes, + timeout_secs, + html, + extra_properties: Vec::new(), + } + } + + /// Add a property to the advertised parameter schema (not `required`). + /// + /// The host uses this to opt the tool into arguments it interprets itself + /// before the call reaches the tool. + #[must_use] + pub fn with_schema_property(mut self, name: &str, schema: serde_json::Value) -> Self { + self.extra_properties.push((name.to_string(), schema)); + self + } +} + +#[async_trait] +impl Tool for WebFetchTool { + fn name(&self) -> &'static str { + "web_fetch" + } + + fn description(&self) -> &'static str { + "GET a URL and read the page. HTML returns as Markdown (links kept, \ + scripts dropped); `raw: true` for the body as sent. For POST or \ + custom headers use `http_request`." + } + + fn parameters_schema(&self) -> serde_json::Value { + let mut schema = json!({ + "type": "object", + "properties": { + "url": { "type": "string", "description": "Absolute http(s) URL." }, + "max_bytes": { + "type": "integer", + "description": "Cap the returned text at this many bytes \ + (default 1_000_000). Bounds the OUTPUT — the extracted \ + markdown, or the raw body with raw:true — never the markup \ + the extractor reads, so lowering it cannot cost you content \ + the page actually had.", + "minimum": 1 + }, + "raw": { + "type": "boolean", + "description": "Return the body as sent." + } + }, + "required": ["url"] + }); + if let Some(properties) = schema["properties"].as_object_mut() { + for (name, property) in &self.extra_properties { + properties.insert(name.clone(), property.clone()); + } + } + schema + } + + fn permission_level(&self) -> PermissionLevel { + PermissionLevel::ReadOnly + } + + /// Idempotent GET — safe to fan out across parallel `web_fetch` + /// calls. Targets that throttle aggressively are the user's + /// concern; we don't try to second-guess at the tool layer. + fn is_concurrency_safe(&self, _args: &serde_json::Value) -> bool { + true + } + + /// How much of a page reaches the model in one result. + /// + /// Extraction does most of the work — after HTML→Markdown a long + /// documentation page is usually a few thousand chars — so this + /// bites only on genuinely large documents. What it no longer does + /// is throw the remainder away: `ToolOutputMiddleware` spills the + /// full extracted page to an artifact and returns the `file_read` + /// call that pages it, which is what this comment used to + /// recommend while the cap itself disabled the affordance. + /// + /// 24k rather than the old 50k because the point of reference + /// moved: 50k was a bound on raw markup, this is clean Markdown. + /// Hermes budgets 15,000 chars of extracted text for the same job; + /// Codex caps every tool result at ~10,000 tokens. + fn max_result_size_chars(&self) -> Option { + Some(24_000) + } + + async fn execute(&self, args: serde_json::Value) -> anyhow::Result { + let raw_url = args + .get("url") + .and_then(|v| v.as_str()) + .ok_or_else(|| anyhow::anyhow!("Missing 'url' parameter"))?; + let max_bytes = args + .get("max_bytes") + .and_then(serde_json::Value::as_u64) + .map_or(self.max_bytes, |n| (n as usize).max(1)); + let raw_requested = args + .get("raw") + .and_then(serde_json::Value::as_bool) + .unwrap_or(false); + + if self.gate.is_rate_limited() { + return Ok(ToolResult::error( + "Rate limit exceeded: too many actions in the last hour", + )); + } + if !self.gate.record_action() { + return Ok(ToolResult::error( + "Rate limit exceeded: action budget exhausted", + )); + } + + // Local-only enforcement (privacy epic S7, #4441): refuse the fetch under + // LocalOnly before URL validation / DNS. The post-validation + // `emit_external_transfer` below stays the S2 disclosure point. + { + let host = host_of(raw_url); + if let Some(msg) = self.gate.local_only_block(&host) { + return Ok(ToolResult::error(msg)); + } + } + + let target = match validate_destination( + &*self.gate, + raw_url, + &self.allowed_domains, + "GET", + &[], + None, + ) + .await + { + Ok(u) => u, + Err(e) => return Ok(ToolResult::error(format!("URL rejected: {e}"))), + }; + let url = target.url.clone(); + + // Egress spine (privacy epic S2, #4436): disclose the fetch destination + // before contacting the host. + self.gate.disclose(&host_of(&url), false, false); + + self.fetch_validated(&target, max_bytes, raw_requested) + .await + } +} + +impl WebFetchTool { + /// Issue the GET for a URL that already passed the gate and the SSRF + /// guard, and render the response for the model. + async fn fetch_validated( + &self, + target: &ValidatedUrl, + max_bytes: usize, + raw_requested: bool, + ) -> anyhow::Result { + // Disable automatic redirect following: reqwest follows up to 10 + // redirects by default, and a redirect target may be on a host + // outside the allowed-domains list. We surface 3xx responses to + // the caller so they can decide whether to refetch the new URL. + let client = match reqwest::Client::builder() + .timeout(Duration::from_secs(self.timeout_secs)) + .user_agent(USER_AGENT) + .redirect(reqwest::redirect::Policy::none()) + // Connect only to the vetted addresses; the URL keeps its host name. + .resolve_to_addrs(&target.host, &target.addrs) + .build() + { + Ok(c) => c, + Err(e) => return Ok(ToolResult::error(format!("Failed to build client: {e}"))), + }; + + let url = target.url.as_str(); + let resp = match client.get(url).send().await { + Ok(r) => r, + Err(e) => return Ok(ToolResult::error(format!("Request failed: {e}"))), + }; + let status = resp.status(); + let final_url = resp.url().to_string(); + let location = resp + .headers() + .get(reqwest::header::LOCATION) + .and_then(|v| v.to_str().ok()) + .map(str::to_string); + let content_type = resp + .headers() + .get(reqwest::header::CONTENT_TYPE) + .and_then(|v| v.to_str().ok()) + .map(str::to_string); + let retry_after = resp + .headers() + .get(reqwest::header::RETRY_AFTER) + .and_then(|v| v.to_str().ok()) + .map(str::to_string); + let body = match resp.text().await { + Ok(b) => b, + Err(e) => return Ok(ToolResult::error(format!("Failed to read body: {e}"))), + }; + + if let Some(loc) = &location + && status.is_redirection() + { + return Ok(ToolResult::success(format!( + "status={} url={} location={loc}\n[redirect not followed — re-call web_fetch with the location URL if it's an allowed domain]", + status.as_u16(), + final_url + ))); + } + + // A 4xx/5xx is a failed fetch, not a page: returning it as a success + // let a blocked or rate-limited site count as research done. (3xx is + // handled above and stays a successful "not followed" report.) + if status.is_client_error() || status.is_server_error() { + let host = host_of(&final_url); + log::debug!( + "[tool.web_fetch] http error status={} host={host} retry_after_present={}", + status.as_u16(), + retry_after.is_some() + ); + let excerpt = + error_body_excerpt(&self.html, &body, content_type.as_deref(), raw_requested) + .await?; + return Ok(ToolResult::error(http_error_message( + status, + &host, + retry_after.as_deref(), + &excerpt, + ))); + } + + let downloaded = body.len(); + + // Markdown by default. A page's prose is a small fraction of its + // bytes; handing the raw document to the model (and to the payload + // summarizer behind it) is how one research turn came to cost + // 1,083,069 input tokens. The host's `HtmlExtractor` owns + // content transforms. + let converted = + !raw_requested && is_html(&self.html, &body, content_type.as_deref()).await?; + let rendered = render_body(&self.html, body, converted, max_bytes).await?; + + let extracted = rendered.extracted; + let mut header = format!("status={} url={final_url}", status.as_u16()); + if converted { + header.push_str(" content=markdown"); + } + // `output_capped_at`, not the old `download_capped_at`: nothing here + // ever capped a download — `resp.text()` above materialises the whole + // body regardless — and naming it that sent a reader looking in the + // wrong place for the content that went missing. + if rendered.output_capped { + header.push_str(&format!(" output_capped_at={max_bytes}B")); + } + append_markup_truncation_header(&mut header, &rendered); + if converted && extracted < downloaded { + header.push_str(&format!(" extracted={extracted}B_of_{downloaded}B")); + } + header.push('\n'); + let content = rendered.content; + + // Full extracted content. Bounding it — the head/tail window, the + // spill to an artifact and the paging handle — belongs to + // `ToolOutputMiddleware`, which applies one rule to every tool. + // Codex enforces exactly this invariant at a single chokepoint + // (`context_manager/history.rs`), which is why a tool there cannot + // leak an unbounded payload however it misbehaves. + Ok(ToolResult::success(format!("{header}{content}"))) + } +} + +/// How much of an error response's body is quoted back to the model. +const ERROR_EXCERPT_CHARS: usize = 300; + +/// The model-facing text for a 4xx/5xx response: what happened, why it +/// matters, and what to do next, plus a short excerpt of the body. +fn http_error_message( + status: reqwest::StatusCode, + host: &str, + retry_after: Option<&str>, + excerpt: &str, +) -> String { + let code = status.as_u16(); + let reason = status.canonical_reason().unwrap_or("Unknown Status"); + let mut msg = format!("HTTP {code} {reason} from {host}; "); + match code { + 429 => { + msg.push_str("the site is rate limiting requests."); + if let Some(wait) = retry_after.map(str::trim).filter(|w| !w.is_empty()) { + msg.push_str(&format!(" Retry-After: {wait}.")); + } + msg.push_str(" Try another source, or retry later."); + } + 401 | 403 => msg.push_str("the site refused the request. Try another source."), + 404 | 410 => { + msg.push_str( + "the page does not exist at this URL. Check the URL or try another source.", + ); + } + 500..=599 => { + msg.push_str( + "the server failed to handle the request. Retry later or try another source.", + ); + } + _ => msg.push_str("the server rejected the request. Try another source."), + } + if !excerpt.is_empty() { + msg.push_str("\nResponse excerpt: "); + msg.push_str(excerpt); + } + msg +} + +/// A short, single-line, text-only excerpt of an error response body. +/// +/// HTML goes through the host extractor (unless the caller asked for `raw`) +/// so the model reads the page's words rather than its markup. +async fn error_body_excerpt( + extractor: &HtmlProvider, + body: &str, + content_type: Option<&str>, + raw_requested: bool, +) -> anyhow::Result { + let (bounded_body, _) = cap_extractor_input(body); + let text = if !raw_requested && is_html(extractor, bounded_body, content_type).await? { + extractor.to_markdown(bounded_body).await? + } else { + body.to_string() + }; + let collapsed = text.split_whitespace().collect::>().join(" "); + Ok(match collapsed.char_indices().nth(ERROR_EXCERPT_CHARS) { + Some((cut, _)) => format!("{}...", &collapsed[..cut]), + None => collapsed, + }) +} + +/// Raw markup handed to the HTML extractor, at most. +/// +/// Not a byte budget — the caller's `max_bytes` owns that, applied to the +/// output. This exists only so a pathological document cannot cost unbounded +/// CPU in `to_markdown`, a cost the old pre-truncation ordering hid by +/// accident. 8 MiB is ~19x the largest real page measured here (a 434 KB +/// client-rendered spreadsheet) and ~8x the default `max_response_size`, so no +/// realistic page reaches it. +const EXTRACTOR_INPUT_CEILING: usize = 8 * 1024 * 1024; + +/// A body turned into what the caller reads. +struct RenderedBody { + /// The text to return, bounded by the caller's `max_bytes`. + content: String, + /// Length *before* that bound, so the header's ratio describes the + /// extraction rather than the truncation. + extracted: usize, + /// Whether `content` was cut to fit `max_bytes`. + output_capped: bool, + /// Whether the markup was cut before the extractor saw it, which only + /// happens past [`EXTRACTOR_INPUT_CEILING`]. + markup_truncated: bool, +} + +/// Add the extractor input ceiling to a fetch header when markup was cut. +fn append_markup_truncation_header(header: &mut String, rendered: &RenderedBody) { + if rendered.markup_truncated { + header.push_str(&format!(" markup_truncated_at={EXTRACTOR_INPUT_CEILING}B")); + } +} + +/// Convert, **then** bound. +/// +/// The order is the whole of this function. It used to be the other way round, +/// and the cap was therefore destroying the thing it was meant to measure: a +/// 432,864-byte client-rendered page fetched with `max_bytes: 50000` was cut at +/// byte 50,000 — mid-tag, mid-DOM — and the readability pass, handed that +/// wreckage, recovered only the ``. The result was +/// `extracted=48B_of_432864B`, reported as `status=200`. The same URL with no +/// `max_bytes` yields `extracted=37243B_of_433638B`: the real document. +/// Identical tool, identical extractor, 776x the content, one parameter. +/// +/// So a caller setting a sensible cost bound silently lost the page, and the +/// knob it would then reach for — raising the cap — was the right knob turned +/// too timidly, which is the worst case for learning anything from the failure. +/// +/// Bounding the output is also what the schema promises, and it costs no +/// memory: the caller has already materialised the whole body. Only the +/// extractor's input needs a ceiling, and that is for CPU. +async fn render_body( + html: &HtmlProvider, + body: String, + converted: bool, + max_bytes: usize, +) -> anyhow::Result<RenderedBody> { + let (markup_truncated, body) = if converted { + let (body, truncated) = cap_extractor_input(&body); + (truncated, body.to_string()) + } else { + (false, body) + }; + let full = if converted { + html.to_markdown(&body).await? + } else { + body + }; + let extracted = full.len(); + let (content, output_capped) = if extracted > max_bytes { + let cut = floor_char_boundary(&full, max_bytes); + (full[..cut].to_string(), true) + } else { + (full, false) + }; + Ok(RenderedBody { + content, + extracted, + output_capped, + markup_truncated, + }) +} + +/// Bound every body before it crosses into a host supplied extractor. +fn cap_extractor_input(body: &str) -> (&str, bool) { + if body.len() > EXTRACTOR_INPUT_CEILING { + let cut = floor_char_boundary(body, EXTRACTOR_INPUT_CEILING); + (&body[..cut], true) + } else { + (body, false) + } +} + +/// The largest index at or below `index` that is a char boundary of `s`. +fn floor_char_boundary(s: &str, index: usize) -> usize { + if index >= s.len() { + return s.len(); + } + let mut end = index; + while end > 0 && !s.is_char_boundary(end) { + end -= 1; + } + end +} + +/// Is this HTML? The server's own `Content-Type` is authoritative when it +/// says so; otherwise fall back to the host's content detection, which +/// already distinguishes HTML from JSON, diffs and code. +async fn is_html( + extractor: &HtmlProvider, + body: &str, + content_type: Option<&str>, +) -> anyhow::Result<bool> { + if let Some(ct) = content_type { + let ct = ct.to_ascii_lowercase(); + let mime = ct.split(';').next().unwrap_or("").trim().to_string(); + // An explicit non-HTML type is a statement, not a guess: a JSON API + // that happens to embed markup must come back verbatim. + if !mime.is_empty() && mime != "text/html" && mime != "application/xhtml+xml" { + return Ok(false); + } + if !mime.is_empty() { + return Ok(true); + } + } + let (body, _) = cap_extractor_input(body); + extractor.looks_like_html(body).await +} + +/// The HTTP status of a `web_fetch` error result. This remains available from +/// the producer module, while the single parser lives in `tinytools::status` +/// so lightweight hosts can consume the result contract without depending on +/// `tinyagents-harness-tools`. +pub use tinytools::status::web_fetch_status as status_of; + +#[cfg(test)] +#[path = "web_fetch_tests.rs"] +mod tests; diff --git a/crates/tinyagents-harness-tools/src/net/network/web_fetch_tests.rs b/crates/tinyagents-harness-tools/src/net/network/web_fetch_tests.rs new file mode 100644 index 000000000..bccbb89c6 --- /dev/null +++ b/crates/tinyagents-harness-tools/src/net/network/web_fetch_tests.rs @@ -0,0 +1,819 @@ +#![allow(clippy::expect_used, clippy::unwrap_used, clippy::panic)] + +use super::*; +use crate::net::network::test_support::{DEFAULT_LIMITS, TestHtml, TestNetGate, pinned_to_local}; +use std::sync::Mutex; + +fn test_security() -> Arc<TestNetGate> { + TestNetGate::supervised() +} + +fn fetch( + gate: Arc<TestNetGate>, + allowed: Vec<String>, + max: Option<usize>, + timeout: Option<u64>, +) -> WebFetchTool { + WebFetchTool::new( + gate, + allowed, + max, + timeout, + DEFAULT_LIMITS, + Arc::new(TestHtml), + ) +} + +#[test] +fn web_fetch_name_and_schema() { + let tool = fetch(test_security(), vec!["example.com".into()], None, None); + assert_eq!(tool.name(), "web_fetch"); + let schema = tool.parameters_schema(); + assert!(schema["properties"]["url"].is_object()); + assert!( + schema["required"] + .as_array() + .unwrap() + .contains(&json!("url")) + ); +} + +#[test] +fn a_host_can_opt_the_schema_into_extra_optional_arguments() { + let tool = fetch(test_security(), vec!["example.com".into()], None, None).with_schema_property( + "summary_focus", + json!({"type": "string", "description": "What you need."}), + ); + let schema = tool.parameters_schema(); + assert_eq!(schema["properties"]["summary_focus"]["type"], "string"); + assert!( + !schema["required"] + .as_array() + .unwrap() + .contains(&json!("summary_focus")) + ); + // Without the opt-in the schema is exactly the base one. + let plain = fetch(test_security(), vec![], None, None).parameters_schema(); + assert!(plain["properties"].get("summary_focus").is_none()); +} + +#[test] +fn zero_and_none_limits_fall_back_to_defaults() { + // Callers wire these from `[http_request]`; a stale `Some(0)` is a + // 0-byte cap (empty bodies) and a 0-second timeout (instant failure). + // Both `None` and `Some(0)` must coerce to the shared schema defaults. + let defaults = DEFAULT_LIMITS; + let from_zero = fetch( + test_security(), + vec!["example.com".into()], + Some(0), + Some(0), + ); + assert_eq!(from_zero.max_bytes, defaults.max_response_size); + assert_eq!(from_zero.timeout_secs, defaults.timeout_secs); + assert_ne!(from_zero.timeout_secs, 0); + assert_ne!(from_zero.max_bytes, 0); + + let from_none = fetch(test_security(), vec!["example.com".into()], None, None); + assert_eq!(from_none.max_bytes, defaults.max_response_size); + assert_eq!(from_none.timeout_secs, defaults.timeout_secs); +} + +#[test] +fn nonzero_limits_are_preserved() { + let tool = fetch( + test_security(), + vec!["example.com".into()], + Some(4096), + Some(15), + ); + assert_eq!(tool.max_bytes, 4096); + assert_eq!(tool.timeout_secs, 15); +} + +#[tokio::test] +async fn web_fetch_rejects_disallowed_domain() { + let tool = fetch(test_security(), vec!["example.com".into()], None, None); + let result = tool + .execute(json!({ "url": "https://evil.test/path" })) + .await + .unwrap(); + assert!(result.is_error); + assert!(result.output().contains("URL rejected")); +} + +#[tokio::test] +async fn web_fetch_rejects_invalid_url() { + let tool = fetch(test_security(), vec!["example.com".into()], None, None); + let result = tool.execute(json!({ "url": "not-a-url" })).await.unwrap(); + assert!(result.is_error); +} + +#[tokio::test] +async fn web_fetch_blocked_under_local_only_privacy_mode() { + // Privacy epic S7 (#4441): under LocalOnly the fetch is refused with a + // `[policy-blocked]` result before any URL validation / network. + let tool = fetch( + TestNetGate::local_only(), + vec!["example.com".into()], + None, + None, + ); + let result = tool + .execute(json!({ "url": "https://example.com/data" })) + .await + .unwrap(); + assert!(result.is_error); + assert!( + result.output().contains("[policy-blocked]"), + "got: {}", + result.output() + ); + assert!( + result.output().contains("Local-only"), + "got: {}", + result.output() + ); +} + +#[tokio::test] +async fn test_web_fetch_truncation_utf8() { + // Mock body with multi-byte char exactly at budget + let body = "Hello 🦀 World"; // 🦀 is at index 6-9 + let max_bytes = 8; + // Should truncate at index 6 + let cut = floor_char_boundary(body, max_bytes); + assert_eq!(cut, 6); + assert_eq!(&body[..cut], "Hello "); +} + +// --- content extraction ---------------------------------------------------- + +async fn html(body: &str, content_type: Option<&str>) -> bool { + is_html(&HtmlProvider::Local(Arc::new(TestHtml)), body, content_type) + .await + .unwrap() +} + +#[tokio::test] +async fn an_explicit_html_content_type_selects_markdown_conversion() { + assert!(html("<p>hi</p>", Some("text/html; charset=utf-8")).await); + assert!(html("<p>hi</p>", Some("application/xhtml+xml")).await); +} + +#[tokio::test] +async fn an_explicit_non_html_content_type_is_taken_at_its_word() { + // A JSON API that happens to quote markup must come back verbatim — + // the server said what it sent, so we don't second-guess it by sniffing. + let body = r#"{"html": "<div><p>one</p><span>two</span><a href="/x">two</a></div>"}"#; + assert!(!html(body, Some("application/json")).await); + assert!(!html("<p>x</p>", Some("text/plain")).await); +} + +#[tokio::test] +async fn a_missing_content_type_falls_back_to_content_detection() { + assert!(html("<!DOCTYPE html><html><body><p>hi</p></body></html>", None).await); + assert!(!html("# Just a README\n\nSome prose.\n", None).await); +} + +#[tokio::test] +async fn an_empty_content_type_does_not_veto_detection() { + assert!(html("<!DOCTYPE html><html><body>x</body></html>", Some("")).await); +} + +#[test] +fn the_schema_offers_the_raw_escape_hatch() { + let tool = fetch(test_security(), vec![], None, None); + let schema = tool.parameters_schema(); + assert!( + schema["properties"].get("raw").is_some(), + "a caller must be able to opt out of conversion: {schema}" + ); + assert!( + tool.description().contains("Markdown"), + "the model needs to know what it will get back: {}", + tool.description() + ); +} + +#[test] +fn the_declared_cap_is_sized_for_extracted_markdown_not_raw_markup() { + let tool = fetch(test_security(), vec![], None, None); + let cap = tool + .max_result_size_chars() + .expect("web_fetch declares a cap"); + assert!( + (8_000..=32_000).contains(&cap), + "cap should sit in the same range as Hermes (15k chars) and Codex \ + (~10k tokens) budget for one result, got {cap}" + ); +} + +#[tokio::test] +async fn html_is_converted_through_the_host_extractor_only_when_it_is_html() { + let body = "<!DOCTYPE html><html><body><p>hi</p></body></html>"; + assert!(html(body, None).await); + assert_eq!(TestHtml.to_markdown(body), "hi"); +} + +#[tokio::test] +async fn execute_blocks_when_rate_limited() { + let tool = fetch( + TestNetGate::with( + crate::net::network::test_support::AutonomyLevel::Supervised, + 0, + ), + vec!["example.com".into()], + None, + None, + ); + let result = tool + .execute(json!({ "url": "https://example.com/data" })) + .await + .unwrap(); + assert!(result.is_error); + assert!(result.output().contains("Rate limit exceeded")); +} + +// --- HTTP error statuses --------------------------------------------------- +// +// Loopback is refused by the SSRF guard, so these drive `fetch_validated` +// (everything after validation) against a one-shot local server. + +/// Serves one canned response and returns the base URL. +async fn serve_once(response: &str) -> String { + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let response = response.to_string(); + tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.unwrap(); + let mut buf = vec![0u8; 8192]; + let _ = socket.read(&mut buf).await.unwrap(); + socket.write_all(response.as_bytes()).await.unwrap(); + let _ = socket.shutdown().await; + }); + format!("http://{addr}/page") +} + +#[tokio::test] +async fn an_outgoing_request_carries_a_user_agent() { + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let seen = Arc::new(std::sync::Mutex::new(String::new())); + let request_log = Arc::clone(&seen); + let server = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.unwrap(); + let mut request = Vec::new(); + let mut buf = [0u8; 1024]; + while !request.windows(4).any(|window| window == b"\r\n\r\n") { + let n = socket.read(&mut buf).await.unwrap(); + assert!( + n != 0, + "peer closed before sending the complete HTTP headers" + ); + request.extend_from_slice(&buf[..n]); + } + *request_log.lock().unwrap() = String::from_utf8_lossy(&request).to_string(); + socket + .write_all(b"HTTP/1.1 200 OK\r\nContent-Length: 2\r\nConnection: close\r\n\r\nok") + .await + .unwrap(); + }); + + let tool = fetch(test_security(), vec![], None, None); + let result = tool + .fetch_validated( + &pinned_to_local(&format!("http://{addr}/")), + 1_000_000, + false, + ) + .await + .unwrap(); + assert!(!result.is_error, "got: {}", result.output()); + server.await.unwrap(); + + let request = seen.lock().unwrap().to_ascii_lowercase(); + let user_agent = request + .lines() + .find_map(|line| line.strip_prefix("user-agent: ")); + assert_eq!( + user_agent, + Some(concat!("tinytools/", env!("CARGO_PKG_VERSION"))) + ); +} + +fn http_response(status_line: &str, headers: &str, body: &str) -> String { + format!( + "HTTP/1.1 {status_line}\r\n{headers}Content-Length: {}\r\nConnection: close\r\n\r\n{body}", + body.len() + ) +} + +async fn fetch_canned(response: String) -> ToolResult { + let url = serve_once(&response).await; + fetch(test_security(), vec![], None, None) + .fetch_validated(&pinned_to_local(&url), 1_000_000, false) + .await + .unwrap() +} + +#[tokio::test] +async fn a_200_is_a_successful_result_with_the_status_header() { + let result = fetch_canned(http_response( + "200 OK", + "Content-Type: text/plain\r\n", + "hello", + )) + .await; + assert!(!result.is_error, "got: {}", result.output()); + assert!(result.output().starts_with("status=200 url=")); + assert!(result.output().ends_with("hello")); +} + +#[tokio::test] +async fn a_403_is_an_error_naming_the_status_and_suggesting_another_source() { + let result = fetch_canned(http_response( + "403 Forbidden", + "Content-Type: text/plain\r\n", + "Access denied by bot protection", + )) + .await; + assert!(result.is_error, "got: {}", result.output()); + let out = result.output(); + assert!( + out.contains("HTTP 403 Forbidden from 127.0.0.1"), + "got: {out}" + ); + assert!(out.contains("refused the request"), "got: {out}"); + assert!(out.contains("another source"), "got: {out}"); + assert!( + out.contains("Access denied by bot protection"), + "got: {out}" + ); +} + +#[tokio::test] +async fn a_429_is_an_error_that_reports_rate_limiting_and_retry_after() { + let result = fetch_canned(http_response( + "429 Too Many Requests", + "Retry-After: 120\r\n", + "", + )) + .await; + assert!(result.is_error, "got: {}", result.output()); + let out = result.output(); + assert!(out.contains("HTTP 429 Too Many Requests"), "got: {out}"); + assert!(out.contains("rate limit"), "got: {out}"); + assert!(out.contains("Retry-After: 120"), "got: {out}"); +} + +#[tokio::test] +async fn a_429_without_retry_after_still_reports_rate_limiting() { + let result = fetch_canned(http_response("429 Too Many Requests", "", "")).await; + assert!(result.is_error); + assert!(result.output().contains("rate limit")); + assert!(!result.output().contains("Retry-After")); +} + +#[tokio::test] +async fn a_404_is_an_error() { + let result = fetch_canned(http_response( + "404 Not Found", + "Content-Type: text/html\r\n", + "<!DOCTYPE html><html><body><p>No such page</p></body></html>", + )) + .await; + assert!(result.is_error, "got: {}", result.output()); + let out = result.output(); + assert!(out.contains("HTTP 404 Not Found"), "got: {out}"); + // The excerpt is the page's text, not its markup. + assert!(out.contains("No such page"), "got: {out}"); + assert!(!out.contains("<p>"), "got: {out}"); +} + +#[tokio::test] +async fn a_5xx_is_an_error() { + let result = fetch_canned(http_response("503 Service Unavailable", "", "")).await; + assert!(result.is_error); + assert!(result.output().contains("HTTP 503 Service Unavailable")); +} + +#[tokio::test] +async fn an_error_body_excerpt_is_short() { + let long = "x".repeat(5_000); + let result = fetch_canned(http_response( + "500 Internal Server Error", + "Content-Type: text/plain\r\n", + &long, + )) + .await; + assert!(result.is_error); + assert!( + result.output().len() < 1_000, + "excerpt must be bounded, got {} bytes", + result.output().len() + ); +} + +#[tokio::test] +async fn a_redirect_is_still_reported_as_a_successful_result() { + let result = fetch_canned(http_response( + "301 Moved Permanently", + "Location: https://example.com/new\r\n", + "", + )) + .await; + assert!(!result.is_error, "got: {}", result.output()); + assert!(result.output().contains("status=301")); + assert!(result.output().contains("location=https://example.com/new")); +} + +// --- The cap bounds the output, not the extractor's input ------------------ +// +// `TestHtml::to_markdown` strips tags, so a document whose prose sits past the +// cap offset is the shape that discriminates: cutting the markup first loses +// the prose outright, cutting the rendered text keeps it. + +/// Markup whose readable text sits behind `filler` bytes of attribute. +fn page_with_prose_after(filler: usize) -> String { + format!( + "<!DOCTYPE html><html><head><title>T\ +

the prose that matters

", + "f".repeat(filler) + ) +} + +#[tokio::test] +async fn the_cap_applies_to_the_extracted_text_not_the_markup() { + // The regression this function exists for. A 432,864-byte page fetched at + // `max_bytes: 50000` used to come back as `extracted=48B_of_432864B` — the + // whole document reduced to its `` — because the markup was cut + // mid-DOM before the extractor ever saw it. + let body = page_with_prose_after(4_000); + assert!(body.len() > 1_000, "the prose must sit past the cap"); + + let rendered = render_body(&HtmlProvider::Local(Arc::new(TestHtml)), body, true, 1_000) + .await + .unwrap(); + assert!( + rendered.content.contains("the prose that matters"), + "cutting the markup first would have lost this: {:?}", + rendered.content + ); +} + +#[test] +fn cutting_the_markup_first_really_does_destroy_the_extraction() { + // The counterfactual, so the test above cannot pass vacuously: the old + // ordering, performed by hand, loses the prose from the same document. + let body = page_with_prose_after(4_000); + let truncated = &body[..1_000]; + assert!( + !TestHtml + .to_markdown(truncated) + .contains("the prose that matters"), + "a pre-extraction cut is what destroyed the content" + ); +} + +#[tokio::test] +async fn the_reported_length_is_the_extraction_not_the_truncation() { + // `extracted` is pre-cap on purpose: the header's `extracted=XB_of_YB` + // ratio describes how much of the document the extractor found, which is + // the number that reveals a collapse. Measuring post-cap would report the + // cap back to the caller as if it were the page. + let body = page_with_prose_after(4_000); + let full = TestHtml.to_markdown(&body); + let rendered = render_body(&HtmlProvider::Local(Arc::new(TestHtml)), body, true, 4) + .await + .unwrap(); + assert_eq!(rendered.extracted, full.len()); + assert!(rendered.output_capped); + assert!(rendered.content.len() <= 4); +} + +#[tokio::test] +async fn an_output_within_the_cap_is_returned_whole_and_unflagged() { + let body = page_with_prose_after(16); + let rendered = render_body( + &HtmlProvider::Local(Arc::new(TestHtml)), + body, + true, + 1_000_000, + ) + .await + .unwrap(); + assert!(!rendered.output_capped); + assert!(!rendered.markup_truncated); + assert!(rendered.content.contains("the prose that matters")); +} + +#[tokio::test] +async fn markup_truncated_input_is_reported_in_the_fetch_header() { + // Drive the body across the extractor's independent input ceiling. The + // large attribute keeps extracted text small, while proving that the + // extractor input itself was bounded and reported to the caller. + let body = format!( + "<!DOCTYPE html><html><body><div data-pad=\"{}\"></div><p>visible</p></body></html>", + "x".repeat(EXTRACTOR_INPUT_CEILING) + ); + assert!(body.len() > EXTRACTOR_INPUT_CEILING); + let rendered = render_body(&HtmlProvider::Local(Arc::new(TestHtml)), body, true, 1_000) + .await + .unwrap(); + assert!(rendered.markup_truncated); + let mut output = "status=200 url=https://example.com content=markdown".to_string(); + append_markup_truncation_header(&mut output, &rendered); + assert!( + output.contains(&format!("markup_truncated_at={EXTRACTOR_INPUT_CEILING}B")), + "header should disclose the extractor input ceiling: {output}" + ); +} + +#[tokio::test] +async fn raw_output_is_bounded_by_the_same_cap() { + // With `raw: true` there is no extraction, so the cap applies to the body + // itself — the one case where cutting the input and cutting the output are + // the same act. + let rendered = render_body( + &HtmlProvider::Local(Arc::new(TestHtml)), + "abcdefghij".to_string(), + false, + 4, + ) + .await + .unwrap(); + assert_eq!(rendered.content, "abcd"); + assert!(rendered.output_capped); + assert_eq!(rendered.extracted, 10); +} + +#[test] +fn the_markup_ceiling_is_far_above_any_real_page() { + // Sized against the measured worst case, not picked round. A ceiling near + // the old default would reintroduce the bug for ordinary documents. + const { + assert!(EXTRACTOR_INPUT_CEILING >= 8 * 1024 * 1024); + assert!(EXTRACTOR_INPUT_CEILING > 433_638 * 10); + } +} + +#[test] +fn the_schema_says_which_side_of_the_extractor_it_bounds() { + // The description is a caller's only account of what the knob does, and + // the old wording ("Truncate body at this many bytes") is what made + // lowering it look free. + let tool = fetch(test_security(), vec![], None, None); + let schema = tool.parameters_schema(); + let desc = schema["properties"]["max_bytes"]["description"] + .as_str() + .expect("max_bytes documents itself"); + assert!(desc.contains("OUTPUT"), "{desc}"); + assert!(desc.contains("never the markup"), "{desc}"); +} + +#[derive(Debug)] +struct RemoteHtml { + fail_detection: bool, + fail_extraction: bool, +} + +#[async_trait] +impl AsyncHtmlExtractor for RemoteHtml { + async fn looks_like_html(&self, body: &str) -> anyhow::Result<bool> { + tokio::task::yield_now().await; + anyhow::ensure!(!self.fail_detection, "remote detection unavailable"); + Ok(TestHtml.looks_like_html(body)) + } + + async fn to_markdown(&self, body: &str) -> anyhow::Result<String> { + tokio::task::yield_now().await; + anyhow::ensure!(!self.fail_extraction, "remote extraction unavailable"); + Ok(TestHtml.to_markdown(body)) + } +} + +fn remote_fetch(fail_detection: bool, fail_extraction: bool) -> WebFetchTool { + WebFetchTool::new_async( + test_security(), + vec![], + None, + None, + DEFAULT_LIMITS, + Arc::new(RemoteHtml { + fail_detection, + fail_extraction, + }), + ) +} + +#[tokio::test] +async fn remote_extraction_reads_before_applying_the_output_cap() { + let tool = remote_fetch(false, false); + let body = page_with_prose_after(4_000); + let url = serve_once(&http_response("200 OK", "", &body)).await; + let result = tool + .fetch_validated(&pinned_to_local(&url), 1_000, false) + .await + .unwrap(); + assert!(result.output().contains("the prose that matters")); + assert!(result.output().contains("content=markdown")); +} + +#[tokio::test] +async fn remote_detection_failure_propagates_without_returning_markup() { + let tool = remote_fetch(true, false); + let url = serve_once(&http_response("200 OK", "", "<html>private</html>")).await; + let error = tool + .fetch_validated(&pinned_to_local(&url), 1_000, false) + .await + .unwrap_err(); + assert_eq!(error.to_string(), "remote detection unavailable"); +} + +#[tokio::test] +async fn remote_extraction_failure_propagates_for_success_and_error_responses() { + let tool = remote_fetch(false, true); + for status in ["200 OK", "403 Forbidden"] { + let url = serve_once(&http_response( + status, + "Content-Type: text/html\r\n", + "<p>private</p>", + )) + .await; + let error = tool + .fetch_validated(&pinned_to_local(&url), 1_000, false) + .await + .unwrap_err(); + assert_eq!(error.to_string(), "remote extraction unavailable"); + } +} + +#[tokio::test] +async fn raw_and_explicit_non_html_responses_do_not_call_the_remote_provider() { + let tool = remote_fetch(true, true); + for (headers, raw) in [("Content-Type: application/json\r\n", false), ("", true)] { + let url = serve_once(&http_response("200 OK", headers, "<p>raw</p>")).await; + let result = tool + .fetch_validated(&pinned_to_local(&url), 1_000, raw) + .await + .unwrap(); + assert!(result.output().ends_with("<p>raw</p>")); + assert!(!result.output().contains("content=markdown")); + } +} + +#[derive(Debug)] +struct PendingHtml; + +#[async_trait] +impl AsyncHtmlExtractor for PendingHtml { + async fn looks_like_html(&self, _body: &str) -> anyhow::Result<bool> { + std::future::pending().await + } + + async fn to_markdown(&self, _body: &str) -> anyhow::Result<String> { + std::future::pending().await + } +} + +#[tokio::test(start_paused = true)] +async fn pending_remote_detection_and_extraction_are_bounded() { + let tool = WebFetchTool::new_async( + test_security(), + vec![], + None, + Some(2), + DEFAULT_LIMITS, + Arc::new(PendingHtml), + ); + let error = is_html(&tool.html, "<html>body</html>", None) + .await + .unwrap_err(); + assert_eq!(error.to_string(), "HTML detection timed out"); + let error = render_body(&tool.html, "<p>body</p>".into(), true, 1_000) + .await + .err() + .unwrap(); + assert_eq!(error.to_string(), "HTML extraction timed out"); + // Error-response excerpts are governed by the same provider deadline. + let error = error_body_excerpt(&tool.html, "<p>body</p>", Some("text/html"), false) + .await + .unwrap_err(); + assert_eq!(error.to_string(), "HTML extraction timed out"); +} + +#[derive(Debug)] +struct InputLengthHtml(Arc<Mutex<Vec<usize>>>); + +#[async_trait] +impl AsyncHtmlExtractor for InputLengthHtml { + async fn looks_like_html(&self, body: &str) -> anyhow::Result<bool> { + self.0.lock().unwrap().push(body.len()); + Ok(true) + } + + async fn to_markdown(&self, body: &str) -> anyhow::Result<String> { + self.0.lock().unwrap().push(body.len()); + Ok(String::new()) + } +} + +#[tokio::test] +async fn every_remote_extractor_call_receives_at_most_the_input_ceiling() { + let lengths = Arc::new(Mutex::new(Vec::new())); + let extractor = HtmlProvider::Async { + extractor: Arc::new(InputLengthHtml(lengths.clone())), + timeout: Duration::from_secs(1), + }; + let body = "x".repeat(EXTRACTOR_INPUT_CEILING + 1); + + // Missing Content-Type takes the remote detection path. + assert!(is_html(&extractor, &body, None).await.unwrap()); + // Error bodies can call detection and extraction directly. + error_body_excerpt(&extractor, &body, None, false) + .await + .unwrap(); + // An explicit HTML error response skips detection but still caps extraction. + error_body_excerpt(&extractor, &body, Some("text/html"), false) + .await + .unwrap(); + + assert_eq!( + *lengths.lock().unwrap(), + vec![ + EXTRACTOR_INPUT_CEILING, + EXTRACTOR_INPUT_CEILING, + EXTRACTOR_INPUT_CEILING, + EXTRACTOR_INPUT_CEILING, + ] + ); +} + +#[test] +fn remote_provider_timeout_uses_the_same_zero_and_none_defaults_as_http() { + for configured in [Some(0), None, Some(7)] { + let tool = WebFetchTool::new_async( + test_security(), + vec![], + None, + configured, + DEFAULT_LIMITS, + Arc::new(PendingHtml), + ); + let HtmlProvider::Async { timeout, .. } = tool.html else { + panic!("expected remote provider") + }; + assert_eq!(timeout, Duration::from_secs(tool.timeout_secs)); + } +} + +#[test] +fn status_of_reads_the_leading_web_fetch_error_shape() { + assert_eq!( + status_of("HTTP 403 Forbidden from example.com; try another source\nbody"), + Some(403) + ); + assert_eq!( + status_of(" HTTP 503 Service Unavailable from a.b; retry later"), + Some(503) + ); +} + +#[test] +fn status_of_rejects_quoted_bare_or_non_error_statuses() { + assert_eq!(status_of("403 Forbidden"), None); + assert_eq!(status_of("HTTP 403"), None); + assert_eq!(status_of("HTTP 200 OK from example.com; fine"), None); + assert_eq!(status_of("HTTP 99 x from h; g"), None); + assert_eq!(status_of("HTTP 404 Not Found from two words; g"), None); + assert_eq!( + status_of("error\nHTTP 403 Forbidden from example.com; quoted"), + None + ); +} + +/// As for `http_request`: the fetch connects to the vetted address and keeps +/// the host name, rather than resolving the name again. +#[tokio::test] +async fn the_fetch_connects_to_the_pinned_address() { + let url = serve_once(&http_response( + "200 OK", + "Content-Type: text/plain\r\n", + "pinned", + )) + .await; + let port = reqwest::Url::parse(&url).unwrap().port().unwrap(); + let target = ValidatedUrl { + url: format!("http://pinned.invalid:{port}/page"), + host: "pinned.invalid".to_string(), + addrs: vec![format!("127.0.0.1:{port}").parse().unwrap()], + }; + let result = fetch(test_security(), vec![], None, None) + .fetch_validated(&target, 1_000_000, false) + .await + .unwrap(); + assert!(!result.is_error, "{result:?}"); +} diff --git a/crates/tinyagents-harness-tools/src/plan_exit.rs b/crates/tinyagents-harness-tools/src/plan_exit.rs index b0e8af9db..45300b2ec 100644 --- a/crates/tinyagents-harness-tools/src/plan_exit.rs +++ b/crates/tinyagents-harness-tools/src/plan_exit.rs @@ -1,40 +1,58 @@ //! Host-configured plan-to-build handoff tool. -use std::sync::Arc; +use std::{hash::Hash, sync::Arc}; +use crate::run_mode::RunModeStore; use async_trait::async_trait; use serde_json::json; +use tinyagents_harness::middleware::library::RunMode; use tinytools::{PermissionLevel, Tool, ToolCallOptions, ToolResult, ToolRunContext}; -use tinyagents_harness::middleware::library::{RunMode, RunModeHandle}; - /// Stable marker emitted for a plan-to-build handoff. pub const PLAN_EXIT_MARKER: &str = "[plan_exit]"; -/// Host callback that resolves the live mode handle only when this call may -/// leave plan mode. Hosts can keep tenant lookup and approval checks here. -pub type PlanExitModeLookup = dyn Fn(&dyn ToolRunContext) -> Option<RunModeHandle> + Send + Sync; +/// Host-owned state operations required for a generic plan exit. +pub trait PlanExitHost: Send + Sync { + type Scope: Clone + Eq + Hash + Send + Sync + 'static; + fn scope_and_thread(&self, context: &dyn ToolRunContext) -> (Self::Scope, String); +} + +/// Optional host-supplied lookup for a pending review in the resolved scope. +pub trait ParkedPlanReviewLookup<S>: Send + Sync { + fn has_parked_review(&self, scope: &S, thread_id: &str) -> bool; +} /// Tool that returns a plan and switches its host-selected run mode to Build. -pub struct PlanExitTool { - lookup: Arc<PlanExitModeLookup>, +pub struct PlanExitTool<S> { + host: Arc<dyn PlanExitHost<Scope = S>>, + run_modes: Arc<dyn RunModeStore<S>>, + reviews: Option<Arc<dyn ParkedPlanReviewLookup<S>>>, } -impl PlanExitTool { +impl<S: Clone + Eq + Hash + Send + Sync + 'static> PlanExitTool<S> { /// Stable tool name advertised to the model and host registries. pub const NAME: &'static str = "plan_exit"; /// Stable tool description advertised to the model and host registries. pub const DESCRIPTION: &'static str = "Exit plan mode and hand off the plan to execution. Call this once the plan is complete — the wrapped harness will switch to build mode (when wired). The `plan` argument is the user-facing plan text that downstream agents will execute against."; - /// Creates a tool using the host's run-mode lookup and policy checks. - pub fn new(lookup: Arc<PlanExitModeLookup>) -> Self { - Self { lookup } + /// Creates a tool whose parked-review check and mode transition are + /// delegated to the host's generic stores. + pub fn new( + host: Arc<dyn PlanExitHost<Scope = S>>, + run_modes: Arc<dyn RunModeStore<S>>, + reviews: Option<Arc<dyn ParkedPlanReviewLookup<S>>>, + ) -> Self { + Self { + host, + run_modes, + reviews, + } } } #[async_trait] -impl Tool for PlanExitTool { +impl<S: Clone + Eq + Hash + Send + Sync + 'static> Tool for PlanExitTool<S> { fn name(&self) -> &str { Self::NAME } @@ -77,10 +95,18 @@ impl Tool for PlanExitTool { return Ok(ToolResult::error("`plan` must not be empty")); } - if let Some(context) = context - && let Some(mode) = (self.lookup)(context) - { - mode.set(RunMode::Build); + if let Some(context) = context { + let (scope, thread_id) = self.host.scope_and_thread(context); + if self + .reviews + .as_ref() + .is_some_and(|reviews| reviews.has_parked_review(&scope, &thread_id)) + { + return Ok(ToolResult::error( + "A plan review is still pending for this thread", + )); + } + self.run_modes.set_mode(&scope, &thread_id, RunMode::Build); } Ok(ToolResult::success(format!("{PLAN_EXIT_MARKER}\n{plan}"))) diff --git a/crates/tinyagents-harness-tools/src/plan_exit_tests.rs b/crates/tinyagents-harness-tools/src/plan_exit_tests.rs index c058e4e5e..b93f13bc2 100644 --- a/crates/tinyagents-harness-tools/src/plan_exit_tests.rs +++ b/crates/tinyagents-harness-tools/src/plan_exit_tests.rs @@ -1,48 +1,148 @@ use super::*; - +#[cfg(feature = "plan-review")] +use crate::plan_review::PlanReviewStore; +use std::sync::{Arc, Mutex}; use tinyagents_harness::{ context::{RunConfig, RunContext}, ids::CallId, + middleware::library::{RunMode, RunModeHandle}, tool::ToolExecutionContext, }; +struct Host; +impl PlanExitHost for Host { + type Scope = u8; + fn scope_and_thread(&self, _: &dyn tinytools::ToolRunContext) -> (u8, String) { + (1, "thread".into()) + } +} +#[derive(Default)] +struct Changes(Mutex<Vec<RunMode>>); +impl crate::run_mode::ModeChangeNotifier<u8> for Changes { + fn changed(&self, _: &u8, _: &str, mode: RunMode) { + self.0.lock().unwrap().push(mode); + } +} +fn tool( + mode: RunMode, + reviews: Option<Arc<dyn ParkedPlanReviewLookup<u8>>>, + changes: Arc<Changes>, +) -> (PlanExitTool<u8>, RunModeHandle) { + let modes = Arc::new(crate::run_mode::InMemoryRunModeStore::new(changes)); + modes.set_mode(&1, "thread", mode); + let handle = modes.handle_for_thread(&1, "thread"); + (PlanExitTool::new(Arc::new(Host), modes, reviews), handle) +} +#[cfg(feature = "plan-review")] +#[derive(Default)] +struct NoopReviews(tokio::sync::Notify); +#[cfg(feature = "plan-review")] +impl crate::plan_review::PlanReviewNotifier<u8> for NoopReviews { + fn requested(&self, _: &crate::plan_review::ParkedReview<u8>) { + self.0.notify_one(); + } + fn decided( + &self, + _: &crate::plan_review::ParkedReview<u8>, + _: &crate::plan_review::PlanReviewResolution, + _: crate::plan_review::PlanReviewReason, + ) { + } +} + #[tokio::test] -async fn plan_exit_switches_the_host_resolved_mode_and_keeps_the_marker() { - let mode = RunModeHandle::new(RunMode::Plan); - let lookup_mode = mode.clone(); - let tool = PlanExitTool::new(Arc::new(move |_| Some(lookup_mode.clone()))); +async fn plan_exit_switches_mode_through_notifying_store_and_keeps_marker() { + let changes = Arc::new(Changes::default()); + let (tool, mode) = tool(RunMode::Plan, None, changes.clone()); let parent = RunContext::new(RunConfig::new("plan-run"), ()); let context = ToolExecutionContext::from_run_context(&parent, CallId::new("call-1")); - let result = tool .execute_with_context( - json!({"plan": " step one "}), + json!({"plan":" step one "}), ToolCallOptions::default(), Some(&context), ) .await - .expect("execute"); - + .unwrap(); assert_eq!(mode.get(), RunMode::Build); + assert_eq!( + *changes.0.lock().unwrap(), + vec![RunMode::Plan, RunMode::Build] + ); assert_eq!(result.text(), format!("{PLAN_EXIT_MARKER}\nstep one")); } #[tokio::test] -async fn plan_exit_leaves_mode_unchanged_when_host_withholds_the_handle() { - let mode = RunModeHandle::new(RunMode::Plan); - let tool = PlanExitTool::new(Arc::new(|_| None)); +#[cfg(feature = "plan-review")] +async fn plan_exit_refuses_while_review_is_parked_and_leaves_mode_unchanged() { + let notifier = Arc::new(NoopReviews::default()); + let reviews = Arc::new(PlanReviewStore::new(notifier.clone())); + let waiter_store = reviews.clone(); + let parked = tokio::spawn(async move { + waiter_store + .request_review( + 1, + "thread".into(), + None, + crate::plan_review::PlanReviewRequest { + summary: "plan".into(), + steps: vec![], + tool_call_id: None, + }, + ) + .await + }); + notifier.0.notified().await; + let changes = Arc::new(Changes::default()); + let review_lookup: Arc<dyn ParkedPlanReviewLookup<u8>> = reviews.clone(); + let (tool, mode) = tool(RunMode::Plan, Some(review_lookup), changes.clone()); let parent = RunContext::new(RunConfig::new("plan-run"), ()); let context = ToolExecutionContext::from_run_context(&parent, CallId::new("call-2")); - let result = tool .execute_with_context( - json!({"plan": "approved later"}), + json!({"plan":"step"}), ToolCallOptions::default(), Some(&context), ) .await - .expect("execute"); + .unwrap(); + assert!(result.is_error); + assert_eq!(mode.get(), RunMode::Plan); + assert_eq!(changes.0.lock().unwrap().last(), Some(&RunMode::Plan)); + let review = reviews.parked_review_for_thread(&1, "thread").unwrap(); + reviews.decide( + &1, + &review.request_id, + crate::plan_review::PlanReviewResolution::Approve, + ); + let _ = parked.await; +} + +#[tokio::test] +async fn plan_exit_rejects_empty_plan() { + let (tool, _) = tool(RunMode::Plan, None, Arc::new(Changes::default())); + let result = tool.execute(json!({"plan":" \n "})).await.unwrap(); + assert!(result.is_error); + assert!(result.text().contains("must not be empty")); +} +#[tokio::test] +async fn plan_exit_without_context_returns_plan_without_changing_mode() { + let changes = Arc::new(Changes::default()); + let (tool, mode) = tool(RunMode::Plan, None, changes.clone()); + let before = changes.0.lock().unwrap().clone(); + let result = tool.execute(json!({"plan":"step"})).await.unwrap(); + assert!(!result.is_error); + assert_eq!(result.text(), format!("{PLAN_EXIT_MARKER}\nstep")); assert_eq!(mode.get(), RunMode::Plan); - assert_eq!(result.text(), format!("{PLAN_EXIT_MARKER}\napproved later")); + assert_eq!(*changes.0.lock().unwrap(), before); +} + +#[test] +fn plan_exit_keeps_tool_metadata() { + use tinytools::{PermissionLevel, Tool}; + + let (tool, _) = tool(RunMode::Build, None, Arc::new(Changes::default())); + assert_eq!(tool.name(), "plan_exit"); + assert_eq!(tool.permission_level(), PermissionLevel::None); } diff --git a/crates/tinyagents-harness-tools/src/plan_review/README.md b/crates/tinyagents-harness-tools/src/plan_review/README.md new file mode 100644 index 000000000..c93d87aed --- /dev/null +++ b/crates/tinyagents-harness-tools/src/plan_review/README.md @@ -0,0 +1,21 @@ +# Plan review + +`RequestPlanReviewTool` formats the model-facing review request and preserves +the established result text. `PlanReviewHost` supplies the interactive +surface, thread identity, run mode, and tool-call id, then routes an interactive +request into the host's scope-aware `PlanReviewStore`. A missing surface uses +the existing non-interactive auto-approval response; calls outside plan mode +return the not-applicable response. + +`PlanReviewStore<S>` treats `S` as an opaque scope key. It keeps a snapshot and +waiter for each pending request, indexes the newest request per scoped thread, +and rejects after its ten-minute default TTL. A timed out request is reported +as expired; a cancelled waiter or withdrawn sender is rejected and reported as +cancelled. `PlanReviewNotifier` receives requested and terminal decisions so +the host can publish its product events without putting product policy in the +generic store. + +The store is in-memory because a parked turn cannot resume after process exit. +Hosts own identity resolution, persistence, and event payload mapping. See +[`plan_exit`](../plan_exit.rs) and [`run_mode`](../run_mode.rs) for the paired +mode transition contracts. diff --git a/crates/tinyagents-harness-tools/src/plan_review/mod.rs b/crates/tinyagents-harness-tools/src/plan_review/mod.rs new file mode 100644 index 000000000..d36618785 --- /dev/null +++ b/crates/tinyagents-harness-tools/src/plan_review/mod.rs @@ -0,0 +1,17 @@ +//! Generic plan review tool and its host-supplied review port. +mod port; +mod store; +mod tool; +pub mod types; +pub use port::{ + InteractiveReviewSurface, PlanReviewHost, PlanReviewPort, PlanReviewRequest, + PlanReviewResolution, +}; +pub use store::{ParkedReview, PlanReviewNotifier, PlanReviewReason, PlanReviewStore}; +pub use tool::RequestPlanReviewTool; +#[cfg(test)] +#[path = "store_tests.rs"] +mod store_tests; +#[cfg(test)] +#[path = "tool_tests.rs"] +mod tests; diff --git a/crates/tinyagents-harness-tools/src/plan_review/port.rs b/crates/tinyagents-harness-tools/src/plan_review/port.rs new file mode 100644 index 000000000..37c1c5ee5 --- /dev/null +++ b/crates/tinyagents-harness-tools/src/plan_review/port.rs @@ -0,0 +1,44 @@ +pub use super::types::PlanReviewResolution; +use async_trait::async_trait; +use tinyagents_harness::middleware::library::RunMode; +use tinytools::ToolRunContext; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct PlanReviewRequest { + pub summary: String, + pub steps: Vec<String>, + pub tool_call_id: Option<String>, +} + +/// Opaque marker that the host has an interactive surface for this request. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct InteractiveReviewSurface; + +#[async_trait] +pub trait PlanReviewHost: Send + Sync { + /// `None` selects the established non-interactive auto-approval path. + fn interactive_surface( + &self, + _context: Option<&dyn ToolRunContext>, + ) -> Option<InteractiveReviewSurface> { + Some(InteractiveReviewSurface) + } + fn thread_id(&self, _context: Option<&dyn ToolRunContext>) -> Option<String> { + None + } + fn run_mode(&self, _context: Option<&dyn ToolRunContext>, _thread_id: Option<&str>) -> RunMode { + RunMode::Plan + } + fn tool_call_id(&self, _context: Option<&dyn ToolRunContext>) -> Option<String> { + None + } + async fn request_review( + &self, + request: PlanReviewRequest, + context: Option<&dyn ToolRunContext>, + ) -> Result<PlanReviewResolution, String>; +} + +/// Compatibility name retained from the initial generic port introduced in +/// the harness-tools crate. +pub use PlanReviewHost as PlanReviewPort; diff --git a/crates/tinyagents-harness-tools/src/plan_review/store.rs b/crates/tinyagents-harness-tools/src/plan_review/store.rs new file mode 100644 index 000000000..64c6f468b --- /dev/null +++ b/crates/tinyagents-harness-tools/src/plan_review/store.rs @@ -0,0 +1,269 @@ +//! In-memory, fail-closed state for human plan reviews. +use super::PlanReviewResolution; +use parking_lot::Mutex; +use std::{collections::HashMap, hash::Hash, sync::Arc, time::Duration}; +use tokio::sync::oneshot; +use uuid::Uuid; + +const DEFAULT_TTL: Duration = Duration::from_secs(600); + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum PlanReviewReason { + Approved, + Rejected, + Revised, + Expired, + Cancelled, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ParkedReview<S> { + pub request_id: String, + pub scope: S, + pub thread_id: String, + pub client_id: Option<String>, + pub summary: String, + pub steps: Vec<String>, + pub tool_call_id: Option<String>, + pub expires_at: String, +} + +pub trait PlanReviewNotifier<S>: Send + Sync { + fn requested(&self, review: &ParkedReview<S>); + fn decided( + &self, + review: &ParkedReview<S>, + resolution: &PlanReviewResolution, + reason: PlanReviewReason, + ); +} + +struct Entry<S> { + snapshot: ParkedReview<S>, + sender: oneshot::Sender<PlanReviewResolution>, +} +struct State<S> { + entries: HashMap<String, Entry<S>>, + newest: HashMap<(S, String), String>, +} + +/// A scoped collection of parked reviews. Scope keys are supplied by the host; +/// this crate never assigns meaning to tenant or user identity. +pub struct PlanReviewStore<S> { + ttl: Duration, + state: Mutex<State<S>>, + notifier: Arc<dyn PlanReviewNotifier<S>>, +} + +impl<S> crate::plan_exit::ParkedPlanReviewLookup<S> for PlanReviewStore<S> +where + S: Clone + Eq + Hash + Send + Sync + 'static, +{ + fn has_parked_review(&self, scope: &S, thread_id: &str) -> bool { + PlanReviewStore::has_parked_review(self, scope, thread_id) + } +} + +impl<S> PlanReviewStore<S> +where + S: Clone + Eq + Hash + Send + Sync + 'static, +{ + pub fn new(notifier: Arc<dyn PlanReviewNotifier<S>>) -> Self { + Self::with_ttl(DEFAULT_TTL, notifier) + } + pub fn with_ttl(ttl: Duration, notifier: Arc<dyn PlanReviewNotifier<S>>) -> Self { + Self { + ttl, + state: Mutex::new(State { + entries: HashMap::new(), + newest: HashMap::new(), + }), + notifier, + } + } + + pub async fn request_review( + &self, + scope: S, + thread_id: String, + client_id: Option<String>, + request: super::PlanReviewRequest, + ) -> PlanReviewResolution { + let id = format!("plan-{}", Uuid::new_v4()); + let expires_at = (chrono::Utc::now() + + chrono::Duration::from_std(self.ttl).unwrap_or_default()) + .to_rfc3339(); + let snapshot = ParkedReview { + request_id: id.clone(), + scope: scope.clone(), + thread_id: thread_id.clone(), + client_id, + summary: request.summary, + steps: request.steps, + tool_call_id: request.tool_call_id, + expires_at, + }; + let (tx, rx) = oneshot::channel(); + let superseded = { + let mut state = self.state.lock(); + let old = state + .newest + .insert((scope.clone(), thread_id.clone()), id.clone()); + let superseded = old.and_then(|old_id| state.entries.remove(&old_id)); + state.entries.insert( + id.clone(), + Entry { + snapshot: snapshot.clone(), + sender: tx, + }, + ); + superseded + }; + drop(superseded); + let mut guard = ParkGuard { + store: self, + snapshot: snapshot.clone(), + armed: true, + }; + self.notifier.requested(&snapshot); + let (resolution, reason) = match tokio::time::timeout(self.ttl, rx).await { + Ok(Ok(value)) => { + let reason = reason_for(&value); + (value, reason) + } + Ok(Err(_)) => (PlanReviewResolution::Reject, PlanReviewReason::Cancelled), + Err(_) => (PlanReviewResolution::Reject, PlanReviewReason::Expired), + }; + guard.armed = false; + self.remove(&snapshot); + self.notifier.decided(&snapshot, &resolution, reason); + resolution + } + + pub fn decide(&self, scope: &S, request_id: &str, resolution: PlanReviewResolution) -> bool { + let mut state = self.state.lock(); + if !state + .entries + .get(request_id) + .is_some_and(|entry| entry.snapshot.scope.eq(scope)) + { + return false; + } + if state + .newest + .get(&( + scope.clone(), + state.entries[request_id].snapshot.thread_id.clone(), + )) + .is_none_or(|id| id != request_id) + { + return false; + } + let entry = state + .entries + .remove(request_id) + .expect("entry checked above"); + let key = ( + entry.snapshot.scope.clone(), + entry.snapshot.thread_id.clone(), + ); + if state.newest.get(&key).is_some_and(|id| id == request_id) { + state.newest.remove(&key); + } + drop(state); + entry.sender.send(resolution).is_ok() + } + /// Drops a pending sender, waking its waiter into the fail-closed cancelled + /// path. Hosts use this when the review surface is withdrawn. + pub fn cancel(&self, scope: &S, request_id: &str) -> bool { + let mut state = self.state.lock(); + if !state + .entries + .get(request_id) + .is_some_and(|entry| entry.snapshot.scope.eq(scope)) + { + return false; + } + let entry = state + .entries + .remove(request_id) + .expect("entry checked above"); + let key = ( + entry.snapshot.scope.clone(), + entry.snapshot.thread_id.clone(), + ); + if state.newest.get(&key).is_some_and(|id| id == request_id) { + state.newest.remove(&key); + } + drop(entry); + true + } + pub fn decide_by_thread( + &self, + scope: &S, + thread_id: &str, + resolution: PlanReviewResolution, + ) -> bool { + let id: Option<String> = self + .state + .lock() + .newest + .get(&(scope.clone(), thread_id.to_owned())) + .cloned(); + id.is_some_and(|id| self.decide(scope, &id, resolution)) + } + pub fn parked_review_for_thread(&self, scope: &S, thread_id: &str) -> Option<ParkedReview<S>> { + let state = self.state.lock(); + let id = state.newest.get(&(scope.clone(), thread_id.to_owned()))?; + state.entries.get(id).map(|entry| entry.snapshot.clone()) + } + pub fn has_parked_review(&self, scope: &S, thread_id: &str) -> bool { + self.parked_review_for_thread(scope, thread_id).is_some() + } + fn remove(&self, snapshot: &ParkedReview<S>) { + let mut state = self.state.lock(); + state.entries.remove(&snapshot.request_id); + let key = (snapshot.scope.clone(), snapshot.thread_id.clone()); + if state + .newest + .get(&key) + .is_some_and(|id| id == &snapshot.request_id) + { + state.newest.remove(&key); + } + } +} + +fn reason_for(resolution: &PlanReviewResolution) -> PlanReviewReason { + match resolution { + PlanReviewResolution::Approve | PlanReviewResolution::NonInteractiveApprove => { + PlanReviewReason::Approved + } + PlanReviewResolution::Reject => PlanReviewReason::Rejected, + PlanReviewResolution::Revise { .. } => PlanReviewReason::Revised, + PlanReviewResolution::NotApplicable => PlanReviewReason::Cancelled, + } +} +struct ParkGuard<'a, S> +where + S: Clone + Eq + Hash + Send + Sync + 'static, +{ + store: &'a PlanReviewStore<S>, + snapshot: ParkedReview<S>, + armed: bool, +} +impl<S> Drop for ParkGuard<'_, S> +where + S: Clone + Eq + Hash + Send + Sync + 'static, +{ + fn drop(&mut self) { + if self.armed { + self.store.remove(&self.snapshot); + self.store.notifier.decided( + &self.snapshot, + &PlanReviewResolution::Reject, + PlanReviewReason::Cancelled, + ); + } + } +} diff --git a/crates/tinyagents-harness-tools/src/plan_review/store_tests.rs b/crates/tinyagents-harness-tools/src/plan_review/store_tests.rs new file mode 100644 index 000000000..55b45fa4b --- /dev/null +++ b/crates/tinyagents-harness-tools/src/plan_review/store_tests.rs @@ -0,0 +1,198 @@ +use super::*; +use crate::plan_review::PlanReviewRequest; +use std::sync::Mutex as StdMutex; +use std::{sync::Arc, time::Duration}; + +#[derive(Default)] +struct Events { + decisions: StdMutex<Vec<(String, PlanReviewReason)>>, + requested: tokio::sync::Notify, +} +impl PlanReviewNotifier<u8> for Events { + fn requested(&self, _: &ParkedReview<u8>) { + self.requested.notify_one(); + } + fn decided( + &self, + review: &ParkedReview<u8>, + _: &PlanReviewResolution, + reason: PlanReviewReason, + ) { + self.decisions + .lock() + .unwrap() + .push((review.request_id.clone(), reason)); + } +} +fn store(ttl: Duration) -> (Arc<PlanReviewStore<u8>>, Arc<Events>) { + let events = Arc::new(Events::default()); + ( + Arc::new(PlanReviewStore::with_ttl(ttl, events.clone())), + events, + ) +} +fn request( + store: Arc<PlanReviewStore<u8>>, + scope: u8, + thread: &str, +) -> tokio::task::JoinHandle<PlanReviewResolution> { + let thread = thread.to_owned(); + tokio::spawn(async move { + store + .request_review( + scope, + thread, + None, + PlanReviewRequest { + summary: "plan".into(), + steps: vec!["one".into()], + tool_call_id: None, + }, + ) + .await + }) +} + +#[tokio::test(start_paused = true)] +async fn timeout_rejects_and_reports_expired() { + let (store, events) = store(Duration::from_secs(600)); + let task = request(store.clone(), 1, "t"); + events.requested.notified().await; + assert!(store.parked_review_for_thread(&1, "t").is_some()); + tokio::time::advance(Duration::from_secs(600)).await; + assert_eq!(task.await.unwrap(), PlanReviewResolution::Reject); + let decisions = events.decisions.lock().unwrap(); + assert_eq!(decisions.len(), 1); + assert_eq!(decisions[0].1, PlanReviewReason::Expired); + drop(decisions); + assert!(!store.has_parked_review(&1, "t")); +} + +#[tokio::test] +async fn approve_resolves_the_parked_turn() { + let (store, events) = store(Duration::from_secs(60)); + let task = request(store.clone(), 1, "approve"); + events.requested.notified().await; + assert!(store.decide_by_thread(&1, "approve", PlanReviewResolution::Approve)); + assert_eq!(task.await.unwrap(), PlanReviewResolution::Approve); + let decisions = events.decisions.lock().unwrap(); + assert_eq!(decisions.len(), 1); + assert_eq!(decisions[0].1, PlanReviewReason::Approved); +} + +#[tokio::test] +async fn revise_returns_feedback_to_the_parked_turn() { + let (store, events) = store(Duration::from_secs(60)); + let task = request(store.clone(), 1, "revise"); + events.requested.notified().await; + let resolution = PlanReviewResolution::Revise { + feedback: "add tests".into(), + }; + assert!(store.decide_by_thread(&1, "revise", resolution.clone())); + assert_eq!(task.await.unwrap(), resolution); + let decisions = events.decisions.lock().unwrap(); + assert_eq!(decisions.len(), 1); + assert_eq!(decisions[0].1, PlanReviewReason::Revised); +} + +#[test] +fn unknown_request_cannot_be_decided() { + let (store, _) = store(Duration::from_secs(60)); + assert!(!store.decide(&1, "missing", PlanReviewResolution::Approve)); + assert!(!store.decide_by_thread(&1, "missing", PlanReviewResolution::Approve)); +} + +#[tokio::test] +async fn parked_snapshot_replays_routing_call_id_and_expiry() { + let (store, events) = store(Duration::from_secs(60)); + let parked_store = store.clone(); + let task = tokio::spawn(async move { + parked_store + .request_review( + 1, + "replay".into(), + Some("client".into()), + PlanReviewRequest { + summary: "plan".into(), + steps: vec!["step".into()], + tool_call_id: Some("call".into()), + }, + ) + .await + }); + events.requested.notified().await; + let snapshot = store.parked_review_for_thread(&1, "replay").unwrap(); + assert_eq!(snapshot.client_id.as_deref(), Some("client")); + assert_eq!(snapshot.tool_call_id.as_deref(), Some("call")); + assert_eq!(snapshot.steps, vec!["step"]); + assert!(!snapshot.expires_at.is_empty()); + assert!(store.decide(&1, &snapshot.request_id, PlanReviewResolution::Approve)); + assert_eq!(task.await.unwrap(), PlanReviewResolution::Approve); + assert!(store.parked_review_for_thread(&1, "replay").is_none()); +} + +#[test] +fn no_snapshot_exists_for_an_unparked_thread() { + let (store, _) = store(Duration::from_secs(60)); + assert!(store.parked_review_for_thread(&1, "missing").is_none()); +} + +#[tokio::test] +async fn dropped_park_reports_cancelled_and_cleans_snapshot() { + let (store, events) = store(Duration::from_secs(600)); + let task = request(store.clone(), 1, "t"); + events.requested.notified().await; + task.abort(); + let _ = task.await; + assert!(!store.has_parked_review(&1, "t")); + let decisions = events.decisions.lock().unwrap(); + assert_eq!(decisions.len(), 1); + assert_eq!(decisions[0].1, PlanReviewReason::Cancelled); +} + +#[tokio::test] +async fn dropped_sender_is_cancelled_not_expired() { + let (store, events) = store(Duration::from_secs(600)); + let task = request(store.clone(), 1, "t"); + events.requested.notified().await; + let id = store.parked_review_for_thread(&1, "t").unwrap().request_id; + assert!(store.cancel(&1, &id)); + assert_eq!(task.await.unwrap(), PlanReviewResolution::Reject); + let decisions = events.decisions.lock().unwrap(); + assert_eq!(decisions.len(), 1); + assert_eq!(decisions[0].1, PlanReviewReason::Cancelled); +} + +#[tokio::test] +async fn decisions_are_scope_isolated_and_newest_review_is_replayed() { + let (store, events) = store(Duration::from_secs(60)); + let one = request(store.clone(), 1, "same"); + events.requested.notified().await; + let old = store.parked_review_for_thread(&1, "same").unwrap(); + let two = request(store.clone(), 1, "same"); + events.requested.notified().await; + let newest = store.parked_review_for_thread(&1, "same").unwrap(); + assert_ne!(old.request_id, newest.request_id); + assert!(store.parked_review_for_thread(&2, "same").is_none()); + assert!(!store.decide(&2, &newest.request_id, PlanReviewResolution::Approve)); + assert!(store.decide_by_thread(&1, "same", PlanReviewResolution::Approve)); + assert_eq!(two.await.unwrap(), PlanReviewResolution::Approve); + assert!(!store.decide(&1, &old.request_id, PlanReviewResolution::Reject)); + assert_eq!(one.await.unwrap(), PlanReviewResolution::Reject); + let decisions = events.decisions.lock().unwrap(); + assert_eq!(decisions.len(), 2); + assert_eq!( + decisions + .iter() + .filter(|(_, reason)| *reason == PlanReviewReason::Cancelled) + .count(), + 1 + ); + assert_eq!( + decisions + .iter() + .filter(|(_, reason)| *reason == PlanReviewReason::Approved) + .count(), + 1 + ); +} diff --git a/crates/tinyagents-harness-tools/src/plan_review/tool.rs b/crates/tinyagents-harness-tools/src/plan_review/tool.rs new file mode 100644 index 000000000..31272ee58 --- /dev/null +++ b/crates/tinyagents-harness-tools/src/plan_review/tool.rs @@ -0,0 +1,107 @@ +use super::{PlanReviewHost, PlanReviewRequest, PlanReviewResolution}; +use async_trait::async_trait; +use serde_json::{Value, json}; +use std::sync::Arc; +use tinyagents_harness::middleware::library::RunMode; +use tinytools::{PermissionLevel, Tool, ToolCallOptions, ToolResult, ToolRunContext, ToolTimeout}; + +pub struct RequestPlanReviewTool { + port: Option<Arc<dyn PlanReviewHost>>, +} +impl RequestPlanReviewTool { + #[must_use] + pub fn new(port: Option<Arc<dyn PlanReviewHost>>) -> Self { + Self { port } + } +} + +#[async_trait] +impl Tool for RequestPlanReviewTool { + fn name(&self) -> &str { + "request_plan_review" + } + fn description(&self) -> &str { + "Pause the turn so the user can approve a thread-scoped plan before you execute it. Blocks until they decide, then returns approved, rejected, or revise with their feedback. Non-interactive policy is supplied by the host." + } + fn parameters_schema(&self) -> Value { + json!({"type":"object","properties":{"summary":{"type":"string","description":"One-line description of the plan being reviewed."},"steps":{"type":"array","description":"Ordered plan steps shown to the user for review.","items":{"type":"string"}}},"required":["summary"]}) + } + fn permission_level(&self) -> PermissionLevel { + PermissionLevel::None + } + fn external_effect(&self) -> bool { + false + } + fn timeout_policy(&self, _: &Value) -> ToolTimeout { + ToolTimeout::Unbounded + } + async fn execute(&self, args: Value) -> anyhow::Result<ToolResult> { + self.execute_with_context(args, ToolCallOptions::default(), None) + .await + } + async fn execute_with_context( + &self, + args: Value, + _: ToolCallOptions, + context: Option<&dyn ToolRunContext>, + ) -> anyhow::Result<ToolResult> { + let Some(port) = &self.port else { + return Ok(ToolResult::error( + "Plan review is unavailable: no review port is installed.", + )); + }; + let summary = args + .get("summary") + .and_then(Value::as_str) + .unwrap_or("") + .trim() + .to_owned(); + let steps = args + .get("steps") + .and_then(Value::as_array) + .map(|values| { + values + .iter() + .filter_map(Value::as_str) + .map(str::trim) + .filter(|s| !s.is_empty()) + .map(ToOwned::to_owned) + .collect() + }) + .unwrap_or_default(); + let Some(_) = port.interactive_surface(context) else { + return Ok(ToolResult::success( + "approved: non-interactive turn (no review surface) — proceed with the plan.", + )); + }; + let thread_id = port.thread_id(context); + if port.run_mode(context, thread_id.as_deref()) != RunMode::Plan { + return Ok(ToolResult::success( + "not applicable: this thread is not in plan mode, so there is no plan to review. Do not call `request_plan_review` again this turn.", + )); + } + let tool_call_id = port.tool_call_id(context); + let resolution = match port + .request_review( + PlanReviewRequest { + summary, + steps, + tool_call_id, + }, + context, + ) + .await + { + Ok(value) => value, + Err(error) => return Ok(ToolResult::error(format!("Plan review failed: {error}"))), + }; + let result = match resolution { + PlanReviewResolution::Approve => "approved: the user approved the plan — proceed and execute it now.".to_owned(), + PlanReviewResolution::NonInteractiveApprove => "approved: non-interactive turn (no review surface) — proceed with the plan.".to_owned(), + PlanReviewResolution::NotApplicable => "not applicable: this thread is not in plan mode, so there is no plan to review. Do not call `request_plan_review` again this turn.".to_owned(), + PlanReviewResolution::Reject => "rejected: the user rejected the plan — do NOT execute it. Briefly ask what they would like to do instead.".to_owned(), + PlanReviewResolution::Revise { feedback } => format!("revise: the user requested changes before executing. Their feedback:\n{feedback}\nRevise the plan accordingly, then call `request_plan_review` again before executing."), + }; + Ok(ToolResult::success(result)) + } +} diff --git a/crates/tinyagents-harness-tools/src/plan_review/tool_tests.rs b/crates/tinyagents-harness-tools/src/plan_review/tool_tests.rs new file mode 100644 index 000000000..f5d099338 --- /dev/null +++ b/crates/tinyagents-harness-tools/src/plan_review/tool_tests.rs @@ -0,0 +1,263 @@ +use super::*; +use async_trait::async_trait; +use serde_json::json; +use std::sync::Arc; +use std::sync::Mutex; +use tinyagents_harness::middleware::library::RunMode; +use tinytools::{Tool, ToolRunContext}; + +struct Port { + resolution: PlanReviewResolution, + request: Mutex<Option<PlanReviewRequest>>, +} +struct FailingPort; +struct NonInteractive; +struct BuildMode; +#[async_trait] +impl PlanReviewHost for NonInteractive { + fn interactive_surface( + &self, + _: Option<&dyn ToolRunContext>, + ) -> Option<InteractiveReviewSurface> { + None + } + async fn request_review( + &self, + _: PlanReviewRequest, + _: Option<&dyn ToolRunContext>, + ) -> Result<PlanReviewResolution, String> { + panic!("non-interactive requests are auto-approved before parking") + } +} +#[async_trait] +impl PlanReviewHost for BuildMode { + fn run_mode(&self, _: Option<&dyn ToolRunContext>, _: Option<&str>) -> RunMode { + RunMode::Build + } + async fn request_review( + &self, + _: PlanReviewRequest, + _: Option<&dyn ToolRunContext>, + ) -> Result<PlanReviewResolution, String> { + panic!("non-plan requests are not parked") + } +} +#[async_trait] +impl PlanReviewHost for FailingPort { + async fn request_review( + &self, + _: PlanReviewRequest, + _: Option<&dyn ToolRunContext>, + ) -> Result<PlanReviewResolution, String> { + Err("review surface unavailable".into()) + } +} +#[async_trait] +impl PlanReviewHost for Port { + fn tool_call_id(&self, _: Option<&dyn ToolRunContext>) -> Option<String> { + Some("call-1".into()) + } + async fn request_review( + &self, + request: PlanReviewRequest, + _: Option<&dyn ToolRunContext>, + ) -> Result<PlanReviewResolution, String> { + *self.request.lock().unwrap() = Some(request); + Ok(self.resolution.clone()) + } +} + +#[tokio::test] +async fn missing_port_fails_closed() { + let result = RequestPlanReviewTool::new(None) + .execute(json!({"summary":"plan"})) + .await + .unwrap(); + assert!(result.is_error); + assert!(result.text().contains("no review port")); +} + +#[tokio::test] +async fn missing_interactive_surface_auto_approves() { + let result = RequestPlanReviewTool::new(Some(Arc::new(NonInteractive))) + .execute(json!({"summary":"plan"})) + .await + .unwrap(); + assert_eq!( + result.text(), + "approved: non-interactive turn (no review surface) — proceed with the plan." + ); +} + +#[tokio::test] +async fn build_mode_returns_not_applicable_without_parking() { + let result = RequestPlanReviewTool::new(Some(Arc::new(BuildMode))) + .execute(json!({"summary":"plan"})) + .await + .unwrap(); + assert!(result.text().starts_with("not applicable:")); +} + +#[tokio::test] +async fn failing_port_returns_error_without_approval() { + let result = RequestPlanReviewTool::new(Some(Arc::new(FailingPort))) + .execute(json!({"summary":"plan"})) + .await + .unwrap(); + assert!(result.is_error); + assert_eq!( + result.text(), + "Plan review failed: review surface unavailable" + ); +} + +#[tokio::test] +async fn request_is_trimmed_and_resolution_text_is_preserved() { + let port = Arc::new(Port { + resolution: PlanReviewResolution::Revise { + feedback: "add tests".into(), + }, + request: Mutex::new(None), + }); + let tool = RequestPlanReviewTool::new(Some(port.clone())); + let result = tool + .execute(json!({"summary":" plan ","steps":[" first "," ",2,"second"]})) + .await + .unwrap(); + assert_eq!( + result.text(), + "revise: the user requested changes before executing. Their feedback:\nadd tests\nRevise the plan accordingly, then call `request_plan_review` again before executing." + ); + assert_eq!( + *port.request.lock().unwrap(), + Some(PlanReviewRequest { + summary: "plan".into(), + steps: vec!["first".into(), "second".into()], + tool_call_id: Some("call-1".into()), + }) + ); +} + +#[tokio::test] +async fn each_resolution_has_its_distinct_result() { + let cases = [ + ( + PlanReviewResolution::Approve, + "approved: the user approved the plan — proceed and execute it now.", + ), + ( + PlanReviewResolution::NonInteractiveApprove, + "approved: non-interactive turn (no review surface) — proceed with the plan.", + ), + ( + PlanReviewResolution::NotApplicable, + "not applicable: this thread is not in plan mode, so there is no plan to review. Do not call `request_plan_review` again this turn.", + ), + ( + PlanReviewResolution::Reject, + "rejected: the user rejected the plan — do NOT execute it. Briefly ask what they would like to do instead.", + ), + ]; + for (resolution, expected) in cases { + let tool = RequestPlanReviewTool::new(Some(Arc::new(Port { + resolution, + request: Mutex::new(None), + }))); + assert_eq!( + tool.execute(json!({"summary":"x"})).await.unwrap().text(), + expected + ); + } +} + +#[tokio::test] +async fn interactive_tool_parks_in_the_scoped_store_until_decided() { + use crate::plan_review::{ParkedReview, PlanReviewNotifier, PlanReviewReason, PlanReviewStore}; + use std::time::Duration; + + type DecisionEvent = ( + String, + PlanReviewResolution, + Option<String>, + PlanReviewReason, + ); + #[derive(Default)] + struct StoreEvents { + requested: tokio::sync::Notify, + decided: Mutex<Vec<DecisionEvent>>, + } + impl PlanReviewNotifier<u8> for StoreEvents { + fn requested(&self, _: &ParkedReview<u8>) { + self.requested.notify_one(); + } + fn decided( + &self, + review: &ParkedReview<u8>, + resolution: &PlanReviewResolution, + reason: PlanReviewReason, + ) { + self.decided.lock().unwrap().push(( + review.thread_id.clone(), + resolution.clone(), + review.tool_call_id.clone(), + reason, + )); + } + } + struct StorePort(Arc<PlanReviewStore<u8>>); + #[async_trait] + impl PlanReviewHost for StorePort { + fn thread_id(&self, _: Option<&dyn ToolRunContext>) -> Option<String> { + Some("thread-7".into()) + } + fn tool_call_id(&self, _: Option<&dyn ToolRunContext>) -> Option<String> { + Some("call-9".into()) + } + async fn request_review( + &self, + request: PlanReviewRequest, + _: Option<&dyn ToolRunContext>, + ) -> Result<PlanReviewResolution, String> { + Ok(self + .0 + .request_review(42, "thread-7".into(), Some("client-3".into()), request) + .await) + } + } + + let events = Arc::new(StoreEvents::default()); + let store = Arc::new(PlanReviewStore::with_ttl( + Duration::from_secs(60), + events.clone(), + )); + let tool = RequestPlanReviewTool::new(Some(Arc::new(StorePort(store.clone())))); + let pending = tool.execute(json!({"summary":"ship it","steps":["build"]})); + tokio::pin!(pending); + assert!( + tokio::time::timeout(Duration::ZERO, &mut pending) + .await + .is_err() + ); + events.requested.notified().await; + let parked = store.parked_review_for_thread(&42, "thread-7").unwrap(); + assert_eq!(parked.client_id.as_deref(), Some("client-3")); + assert_eq!(parked.tool_call_id.as_deref(), Some("call-9")); + assert_eq!(parked.summary, "ship it"); + assert_eq!(parked.steps, vec!["build"]); + assert!(store.decide(&42, &parked.request_id, PlanReviewResolution::Approve)); + let result = pending.await.unwrap(); + assert_eq!( + result.text(), + "approved: the user approved the plan — proceed and execute it now." + ); + assert!(!store.has_parked_review(&42, "thread-7")); + assert_eq!( + *events.decided.lock().unwrap(), + vec![( + "thread-7".into(), + PlanReviewResolution::Approve, + Some("call-9".into()), + PlanReviewReason::Approved, + )] + ); +} diff --git a/crates/tinyagents-harness-tools/src/plan_review/types.rs b/crates/tinyagents-harness-tools/src/plan_review/types.rs new file mode 100644 index 000000000..af06c2f41 --- /dev/null +++ b/crates/tinyagents-harness-tools/src/plan_review/types.rs @@ -0,0 +1,39 @@ +//! Wire-compatible plan-review decisions and internal tool outcomes. + +use serde::{Deserialize, Serialize}; + +/// The user's decision on a parked plan, plus internal tool outcomes preserved +/// for the generic plan-review tool's established result behavior. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(tag = "decision", rename_all = "snake_case")] +pub enum PlanReviewResolution { + /// Run the plan as proposed. + Approve, + /// Internal outcome when no interactive review surface exists. + #[serde(skip)] + NonInteractiveApprove, + /// Internal outcome when the current run is not in plan mode. + #[serde(skip)] + NotApplicable, + /// Drop the plan. + Reject, + /// Revise the plan using user feedback. + Revise { feedback: String }, +} + +impl PlanReviewResolution { + /// Stable decision label for events and logs; revise feedback is omitted. + pub fn as_str(&self) -> &'static str { + match self { + Self::Approve => "approve", + Self::Reject => "reject", + Self::Revise { .. } => "revise", + Self::NonInteractiveApprove => "non_interactive_approve", + Self::NotApplicable => "not_applicable", + } + } +} + +#[cfg(test)] +#[path = "types_tests.rs"] +mod tests; diff --git a/crates/tinyagents-harness-tools/src/plan_review/types_tests.rs b/crates/tinyagents-harness-tools/src/plan_review/types_tests.rs new file mode 100644 index 000000000..5236c485f --- /dev/null +++ b/crates/tinyagents-harness-tools/src/plan_review/types_tests.rs @@ -0,0 +1,45 @@ +use super::*; + +#[test] +fn decisions_keep_the_original_tagged_json_wire_form() { + let fixtures = [ + (PlanReviewResolution::Approve, r#"{"decision":"approve"}"#), + (PlanReviewResolution::Reject, r#"{"decision":"reject"}"#), + ( + PlanReviewResolution::Revise { + feedback: "add tests".into(), + }, + r#"{"decision":"revise","feedback":"add tests"}"#, + ), + ]; + for (resolution, expected) in fixtures { + let value = serde_json::to_string(&resolution).unwrap(); + assert_eq!(value, expected); + assert_eq!( + serde_json::from_str::<PlanReviewResolution>(expected).unwrap(), + resolution + ); + } +} + +#[test] +fn internal_tool_outcomes_are_not_client_wire_decisions() { + assert!(serde_json::to_string(&PlanReviewResolution::NonInteractiveApprove).is_err()); + assert!(serde_json::to_string(&PlanReviewResolution::NotApplicable).is_err()); + assert!( + serde_json::from_str::<PlanReviewResolution>(r#"{"decision":"non_interactive_approve"}"#) + .is_err() + ); + assert!( + serde_json::from_str::<PlanReviewResolution>(r#"{"decision":"not_applicable"}"#).is_err() + ); + assert_eq!(PlanReviewResolution::Approve.as_str(), "approve"); + assert_eq!(PlanReviewResolution::Reject.as_str(), "reject"); + assert_eq!( + PlanReviewResolution::Revise { + feedback: "private feedback".into() + } + .as_str(), + "revise" + ); +} diff --git a/crates/tinyagents-harness-tools/src/run_mode.rs b/crates/tinyagents-harness-tools/src/run_mode.rs new file mode 100644 index 000000000..5bb49e05c --- /dev/null +++ b/crates/tinyagents-harness-tools/src/run_mode.rs @@ -0,0 +1,83 @@ +//! Reusable per-thread run-mode storage and change notification. +use parking_lot::Mutex; +use std::{collections::HashMap, hash::Hash, sync::Arc}; +use tinyagents_harness::middleware::library::{RunMode, RunModeHandle}; + +pub trait ModeChangeNotifier<S>: Send + Sync { + fn changed(&self, scope: &S, thread_id: &str, mode: RunMode); +} +pub trait RunModeStore<S>: Send + Sync { + fn handle_for_thread(&self, scope: &S, thread_id: &str) -> RunModeHandle; + fn get_mode(&self, scope: &S, thread_id: &str) -> RunMode; + fn set_mode(&self, scope: &S, thread_id: &str, mode: RunMode); +} + +pub struct InMemoryRunModeStore<S> { + handles: Mutex<HashMap<(S, String), RunModeHandle>>, + update: Mutex<()>, + notifier: Arc<dyn ModeChangeNotifier<S>>, +} +impl<S: Clone + Eq + Hash> InMemoryRunModeStore<S> { + pub fn new(notifier: Arc<dyn ModeChangeNotifier<S>>) -> Self { + Self { + handles: Mutex::new(HashMap::new()), + update: Mutex::new(()), + notifier, + } + } +} +impl<S: Clone + Eq + Hash + Send + Sync + 'static> RunModeStore<S> for InMemoryRunModeStore<S> { + fn handle_for_thread(&self, scope: &S, thread_id: &str) -> RunModeHandle { + self.handles + .lock() + .entry((scope.clone(), thread_id.to_owned())) + .or_default() + .clone() + } + fn get_mode(&self, scope: &S, thread_id: &str) -> RunMode { + self.handles + .lock() + .get(&(scope.clone(), thread_id.to_owned())) + .map_or(RunMode::Build, RunModeHandle::get) + } + fn set_mode(&self, scope: &S, thread_id: &str, mode: RunMode) { + let changed = { + let _update = self.update.lock(); + let key = (scope.clone(), thread_id.to_owned()); + let mut handles = self.handles.lock(); + match handles.get(&key) { + Some(handle) if mode == handle.get() => false, + Some(handle) => { + handle.set(mode); + true + } + None if mode == RunMode::Build => false, + None => { + handles.insert(key, RunModeHandle::new(mode)); + true + } + } + }; + if changed { + self.notifier.changed(scope, thread_id, mode); + } + } +} + +pub fn mode_label(mode: RunMode) -> &'static str { + match mode { + RunMode::Build => "build", + RunMode::Plan => "plan", + } +} +pub fn parse_mode_label(label: &str) -> Option<RunMode> { + match label { + "build" => Some(RunMode::Build), + "plan" => Some(RunMode::Plan), + _ => None, + } +} + +#[cfg(test)] +#[path = "run_mode_tests.rs"] +mod tests; diff --git a/crates/tinyagents-harness-tools/src/run_mode_tests.rs b/crates/tinyagents-harness-tools/src/run_mode_tests.rs new file mode 100644 index 000000000..e7d13204a --- /dev/null +++ b/crates/tinyagents-harness-tools/src/run_mode_tests.rs @@ -0,0 +1,80 @@ +use super::*; +use std::sync::Mutex as StdMutex; + +#[derive(Default)] +struct Events(StdMutex<Vec<(u8, String, RunMode)>>); +impl ModeChangeNotifier<u8> for Events { + fn changed(&self, scope: &u8, thread: &str, mode: RunMode) { + self.0 + .lock() + .unwrap() + .push((*scope, thread.to_owned(), mode)); + } +} + +#[test] +fn modes_are_scoped_reused_and_notified() { + let events = Arc::new(Events::default()); + let store = InMemoryRunModeStore::new(events.clone()); + let first = store.handle_for_thread(&1, "same"); + let second = store.handle_for_thread(&1, "same"); + assert_eq!(first.get(), RunMode::Build); + store.set_mode(&1, "same", RunMode::Plan); + assert_eq!(second.get(), RunMode::Plan); + assert_eq!(store.get_mode(&2, "same"), RunMode::Build); + assert_eq!( + *events.0.lock().unwrap(), + vec![(1, "same".into(), RunMode::Plan)] + ); +} + +#[test] +fn run_mode_labels_round_trip() { + assert_eq!(mode_label(RunMode::Plan), "plan"); + assert_eq!(mode_label(RunMode::Build), "build"); + assert_eq!(parse_mode_label("plan"), Some(RunMode::Plan)); + assert_eq!(parse_mode_label("build"), Some(RunMode::Build)); + assert_eq!(parse_mode_label("other"), None); +} + +#[test] +fn unchanged_mode_does_not_notify_and_unknown_reads_do_not_create_handles() { + let events = Arc::new(Events::default()); + let store = InMemoryRunModeStore::new(events.clone()); + + assert_eq!(store.get_mode(&1, "unseen"), RunMode::Build); + assert!(store.handles.lock().is_empty()); + store.set_mode(&1, "unseen", RunMode::Build); + assert!(events.0.lock().unwrap().is_empty()); + assert!(store.handles.lock().is_empty()); +} + +struct ReentrantNotifier { + store: std::sync::OnceLock<std::sync::Weak<InMemoryRunModeStore<u8>>>, + calls: std::sync::atomic::AtomicUsize, +} + +impl ModeChangeNotifier<u8> for ReentrantNotifier { + fn changed(&self, scope: &u8, thread: &str, mode: RunMode) { + if self.calls.fetch_add(1, std::sync::atomic::Ordering::SeqCst) == 0 { + self.store + .get() + .and_then(std::sync::Weak::upgrade) + .unwrap() + .set_mode(scope, thread, mode); + } + } +} + +#[test] +fn reentrant_same_mode_change_notifies_once() { + let notifier = Arc::new(ReentrantNotifier { + store: std::sync::OnceLock::new(), + calls: std::sync::atomic::AtomicUsize::new(0), + }); + let store = Arc::new(InMemoryRunModeStore::new(notifier.clone())); + notifier.store.set(Arc::downgrade(&store)).unwrap(); + + store.set_mode(&1, "thread", RunMode::Plan); + assert_eq!(notifier.calls.load(std::sync::atomic::Ordering::SeqCst), 1); +} diff --git a/crates/tinyagents-harness-tools/src/schema.rs b/crates/tinyagents-harness-tools/src/schema.rs new file mode 100644 index 000000000..a56684191 --- /dev/null +++ b/crates/tinyagents-harness-tools/src/schema.rs @@ -0,0 +1,25 @@ +//! Parse-once utility for embedded JSON tool schemas. + +use serde_json::Value; + +/// Parse an embedded schema and identify the source location on malformed JSON. +pub fn parse_embedded(site: &str, source: &str) -> Value { + serde_json::from_str(source) + .unwrap_or_else(|error| panic!("embedded tool schema at {site} is not valid JSON: {error}")) +} + +/// Parse an embedded JSON schema once and return a clone at each call site. +#[macro_export] +macro_rules! static_schema { + ($source:expr) => {{ + static SCHEMA: ::std::sync::LazyLock<::serde_json::Value> = + ::std::sync::LazyLock::new(|| { + $crate::schema::parse_embedded(concat!(file!(), ":", line!()), $source) + }); + ::std::clone::Clone::clone(&*SCHEMA) + }}; +} + +#[cfg(test)] +#[path = "schema_tests.rs"] +mod tests; diff --git a/crates/tinyagents-harness-tools/src/schema_tests.rs b/crates/tinyagents-harness-tools/src/schema_tests.rs new file mode 100644 index 000000000..f2a0c7d5c --- /dev/null +++ b/crates/tinyagents-harness-tools/src/schema_tests.rs @@ -0,0 +1,16 @@ +use super::*; + +#[test] +fn embedded_schema_parses_and_returns_independent_values() { + let mut first = crate::static_schema!(r#"{"type":"object","required":["command"]}"#); + let second = crate::static_schema!(r#"{"type":"object","required":["command"]}"#); + first["type"] = "changed".into(); + assert_eq!(second["type"], "object"); + assert_eq!(second["required"][0], "command"); +} + +#[test] +#[should_panic(expected = "embedded tool schema at test-site is not valid JSON")] +fn malformed_embedded_schema_names_the_call_site() { + let _ = parse_embedded("test-site", "not json"); +} diff --git a/crates/tinyagents-harness-tools/src/skills/mod.rs b/crates/tinyagents-harness-tools/src/skills/mod.rs index 0cc58f12e..e558acab3 100644 --- a/crates/tinyagents-harness-tools/src/skills/mod.rs +++ b/crates/tinyagents-harness-tools/src/skills/mod.rs @@ -36,7 +36,7 @@ pub use handle::PackRegistryHandle; pub use render::{ NoSuchPackTool, named_tool, render_pack_filtered, route_sentence, scope_use_skill_spec, }; -pub use tool::{USE_SKILL, UseSkillTool}; +pub use tool::{USE_SKILL, UseSkillTool, pack_registry_handle}; pub(crate) use types::PinnedPackCall; pub use types::{PackSelection, ToolPack}; diff --git a/crates/tinyagents-harness-tools/src/skills/mod_permissions_tests.rs b/crates/tinyagents-harness-tools/src/skills/mod_permissions_tests.rs index c83dc461a..a9607f337 100644 --- a/crates/tinyagents-harness-tools/src/skills/mod_permissions_tests.rs +++ b/crates/tinyagents-harness-tools/src/skills/mod_permissions_tests.rs @@ -232,19 +232,17 @@ async fn read_only_permission_ceiling_refuses_a_packed_delete_before_execution() handle.bind(Arc::downgrade(&tools)); let dispatcher = Arc::new(UseSkillTool::new(handle, CATALOG)); let mut harness: AgentHarness<(), ()> = AgentHarness::new(); - harness - .register_model( - "mock", - Arc::new(MockModel::with_responses(vec![ - testkit::tool_call_response(ToolCall::new( - "packed-delete", - USE_SKILL, - json!({"skill": "code", "tool": "delete_source"}), - )), - testkit::text_response("finished"), - ])), - ) - .register_tool(dispatcher); + harness.with_model_resolver(Arc::new(tinyagents_harness::host::FixedModelResolver::new( + Arc::new(MockModel::with_responses(vec![ + testkit::tool_call_response(ToolCall::new( + "packed-delete", + USE_SKILL, + json!({"skill": "code", "tool": "delete_source"}), + )), + testkit::text_response("finished"), + ])), + ))); + harness.register_tool(dispatcher); harness.with_policy(tinyagents_harness::runtime::RunPolicy { tool_rules: tinyagents_harness::tool::ToolRulePolicy::default() .with_permission_ceiling(PermissionLevel::ReadOnly), diff --git a/crates/tinyagents-harness-tools/src/skills/mod_tests.rs b/crates/tinyagents-harness-tools/src/skills/mod_tests.rs index a6cf0fb9f..6a2cc86f6 100644 --- a/crates/tinyagents-harness-tools/src/skills/mod_tests.rs +++ b/crates/tinyagents-harness-tools/src/skills/mod_tests.rs @@ -54,6 +54,7 @@ struct FakeTool { external: bool, timeout: ToolTimeout, prepares_authorization: bool, + binding: Option<ToolCallBinding>, } impl FakeTool { @@ -64,8 +65,14 @@ impl FakeTool { external: false, timeout: ToolTimeout::Inherit, prepares_authorization: false, + binding: None, } } + + fn with_binding(mut self, binding: ToolCallBinding) -> Self { + self.binding = Some(binding); + self + } } #[async_trait] @@ -96,6 +103,11 @@ impl Tool for FakeTool { fn timeout_policy(&self, _args: &Value) -> ToolTimeout { self.timeout } + fn host_extension(&self) -> Option<&(dyn std::any::Any + Send + Sync)> { + self.binding + .as_ref() + .map(|binding| binding as &(dyn std::any::Any + Send + Sync)) + } } /// A durable registry holding `tools` plus a bound `use_skill`. @@ -262,6 +274,32 @@ async fn authorization_and_execution_use_the_same_packed_tool_after_rebind() { assert_eq!(result.text(), "authorized"); } +#[tokio::test] +async fn erased_use_skill_exposes_its_shared_registry_handle() { + let handle = PackRegistryHandle::default(); + let tool: Arc<dyn Tool> = Arc::new(UseSkillTool::new(handle.clone(), CATALOG)); + let recovered = pack_registry_handle(tool.as_ref()).expect("registered use_skill handle"); + let registry = Arc::new(vec![Box::new(FakeTool::plain("a_one")) as Box<dyn Tool>]); + recovered.bind(Arc::downgrade(®istry)); + + let result = tool + .execute(json!({"skill": "alpha", "tool": "a_one", "args": {}})) + .await + .unwrap(); + assert_eq!(result.text(), "a_one:{}"); +} + +#[test] +fn unrelated_tool_has_no_pack_registry_handle() { + let tool = FakeTool::plain("anything_else"); + assert!(pack_registry_handle(&tool).is_none()); + + let binding = ToolCallBinding::new(|_, _| unreachable!("metadata lookup does not bind")) + .with_metadata(String::from("unrelated metadata")); + let tool = FakeTool::plain("another_tool").with_binding(binding); + assert!(pack_registry_handle(&tool).is_none()); +} + #[tokio::test] async fn catalogs_with_the_same_start_pointer_but_different_lengths_conflict() { let handle = PackRegistryHandle::default(); @@ -379,13 +417,15 @@ async fn registered_pack_dispatch_runs_through_the_harness_with_admitted_target_ )); let mut harness: AgentHarness<(), String> = AgentHarness::new(); harness - .register_model( - "mock", - Arc::new(MockModel::with_responses(vec![ - response, - testkit::text_response("finished"), - ])), - ) + .with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named( + "mock", + Arc::new(MockModel::with_responses(vec![ + response, + testkit::text_response("finished"), + ])), + ), + )) .register_tool_dispatch(dispatch); let parent = tinyagents_harness::context::RunContext::new( tinyagents_harness::context::RunConfig::new("outer-run"), @@ -429,17 +469,19 @@ async fn harness_rule_denial_blocks_the_pinned_inner_tool_before_typed_dispatch( }); let mut harness: AgentHarness<(), String> = AgentHarness::new(); harness - .register_model( - "mock", - Arc::new(MockModel::with_responses(vec![ - testkit::tool_call_response(ToolCall::new( - "denied-call", - USE_SKILL, - json!({"skill":"alpha", "tool":"a_one", "args":{}}), - )), - testkit::text_response("finished"), - ])), - ) + .with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named( + "mock", + Arc::new(MockModel::with_responses(vec![ + testkit::tool_call_response(ToolCall::new( + "denied-call", + USE_SKILL, + json!({"skill":"alpha", "tool":"a_one", "args":{}}), + )), + testkit::text_response("finished"), + ])), + ), + )) .register_tool_dispatch(Arc::new(UseSkillDispatch::new(outer, selector))); let policy = tinyagents_harness::runtime::RunPolicy { tool_rules: tinyagents_harness::tool::ToolRulePolicy::new(tinytools::ToolRuleSet::single( @@ -577,7 +619,8 @@ async fn an_unbound_handle_reports_the_registry_unavailable() { /// must be reachable, and rebinding it must repoint the handle. #[tokio::test] async fn a_tool_in_the_synthesised_set_is_reachable_and_rebinding_repoints() { - let (durable, handle) = bound_with_handle(vec![]); + let (durable, _) = bound_with_handle(vec![]); + let handle = pack_registry_handle(find(&durable, USE_SKILL)).expect("registered handle"); let first: Arc<Vec<Box<dyn Tool>>> = Arc::new(vec![Box::new(FakeTool::plain("a_one"))]); handle.bind_synthesized(Arc::downgrade(&first)); let use_skill = find(&durable, USE_SKILL); @@ -625,8 +668,8 @@ fn resolve_registry_for_enforces_pack_membership() { #[tokio::test] async fn the_filtered_listing_omits_tools_the_session_cannot_call() { - let (_tools, handle) = - bound_with_handle(vec![FakeTool::plain("a_one"), FakeTool::plain("a_two")]); + let (tools, _) = bound_with_handle(vec![FakeTool::plain("a_one"), FakeTool::plain("a_two")]); + let handle = pack_registry_handle(find(&tools, USE_SKILL)).expect("registered handle"); let listing = render_pack_filtered(&CATALOG, "alpha", &handle, &|name| name == "a_one", "") .expect("one callable tool"); assert!(listing.contains("## `a_one`")); @@ -669,9 +712,13 @@ fn no_such_pack_tool_renders_each_shape() { } #[test] -fn route_sentence_prefers_a_callable_tool_and_falls_back_to_owners() { +fn route_sentence_keeps_owner_identity_when_a_delegate_tool_is_callable() { assert_eq!( route_sentence(&["a".to_string(), "b".to_string()], &["o"]), + "Call `a` or `b` to delegate; these tools are owned by `o`. The delegate runs them directly. Do not retry this skill." + ); + assert_eq!( + route_sentence(&["a".to_string(), "b".to_string()], &[]), "Call `a` or `b` instead — that agent owns these tools and runs them directly. Do not retry this skill." ); assert_eq!( @@ -767,7 +814,9 @@ async fn generic_binding_applies_tool_rules_to_the_selected_packed_tool() { ["a_one"], )); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("scripted", model); + harness.with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named("scripted", model), + )); harness.register_tool(Arc::new(UseSkillTool::new(handle, CATALOG))); harness.with_policy(tinyagents_harness::runtime::RunPolicy { tool_rules: tinyagents_harness::tool::ToolRulePolicy::new(rules), diff --git a/crates/tinyagents-harness-tools/src/skills/render.rs b/crates/tinyagents-harness-tools/src/skills/render.rs index 85f2537f7..8c2537cee 100644 --- a/crates/tinyagents-harness-tools/src/skills/render.rs +++ b/crates/tinyagents-harness-tools/src/skills/render.rs @@ -154,8 +154,9 @@ impl NoSuchPackTool<'_> { /// `callable_delegates` are delegation tool names the caller has already /// confirmed this session can invoke — naming the *tool* rather than the agent /// is the difference between guidance and an instruction, and a model left to -/// guess the call retries. When none can be reached the owning agents are named -/// instead: strictly worse, but still better than a bare denial. +/// guess the call retries. When owners are supplied, keep that ownership +/// metadata alongside the callable route. When none can be reached the owning +/// agents are named instead: strictly worse, but still better than a bare denial. pub fn route_sentence(callable_delegates: &[String], owners: &[&str]) -> String { if !callable_delegates.is_empty() { let names = callable_delegates @@ -163,6 +164,16 @@ pub fn route_sentence(callable_delegates: &[String], owners: &[&str]) -> String .map(|t| format!("`{t}`")) .collect::<Vec<_>>() .join(" or "); + if !owners.is_empty() { + let owner_names = owners + .iter() + .map(|owner| format!("`{owner}`")) + .collect::<Vec<_>>() + .join(" or "); + return format!( + "Call {names} to delegate; these tools are owned by {owner_names}. The delegate runs them directly. Do not retry this skill." + ); + } return format!( "Call {names} instead — that agent owns these tools and runs them directly. \ Do not retry this skill." diff --git a/crates/tinyagents-harness-tools/src/skills/tool.rs b/crates/tinyagents-harness-tools/src/skills/tool.rs index 091148823..47d1ca272 100644 --- a/crates/tinyagents-harness-tools/src/skills/tool.rs +++ b/crates/tinyagents-harness-tools/src/skills/tool.rs @@ -250,7 +250,8 @@ impl UseSkillTool { ); let binding_handle = handle.clone(); let binding = - ToolCallBinding::new(move |outer, args| binding_handle.bind_call(outer, args)); + ToolCallBinding::new(move |outer, args| binding_handle.bind_call(outer, args)) + .with_metadata(handle.clone()); Self { handle, catalog, @@ -270,6 +271,18 @@ impl UseSkillTool { } } +/// Recovers the shared registry handle from an erased registered `use_skill` tool. +/// +/// The handle lets a host bind durable and synthesized registries and render +/// their filtered pack listings. Calls still pass through the generic +/// [`ToolCallBinding`] so admission and execution use the same pinned target. +pub fn pack_registry_handle(tool: &dyn Tool) -> Option<PackRegistryHandle> { + tool.host_extension() + .and_then(|extension| extension.downcast_ref::<ToolCallBinding>()) + .and_then(|binding| binding.metadata::<PackRegistryHandle>()) + .cloned() +} + #[async_trait] impl Tool for UseSkillTool { fn name(&self) -> &str { diff --git a/crates/tinyagents-harness-tools/src/subagents/driver_tool.rs b/crates/tinyagents-harness-tools/src/subagents/driver_tool.rs index 9a618a778..78e44db44 100644 --- a/crates/tinyagents-harness-tools/src/subagents/driver_tool.rs +++ b/crates/tinyagents-harness-tools/src/subagents/driver_tool.rs @@ -1,8 +1,9 @@ -//! `SubAgentTool` adapter for the host-planned [`SubagentDriver`] lifecycle. +//! Async tool adapter for the host-planned [`SubagentDriver`] lifecycle. //! //! The request factory is host supplied: it owns authorization, child context -//! construction, durable task identity, and opaque planner options. The driver -//! then remains the only path that plans, executes, and persists the run. +//! construction, durable task identity, and opaque planner options. Every +//! accepted call starts a background job and returns its status immediately; +//! waiting and completion delivery remain separate host operations. use std::future::Future; use std::marker::PhantomData; @@ -17,9 +18,9 @@ use tinyagents_harness::tool::ToolDispatch; use tinytools::{PermissionLevel, Tool, ToolCategory, ToolResult, ToolScope, ToolTimeout}; use tinyagents_orchestration::subagent::DriverJobRegistration; -use tinyagents_orchestration::subagent::InlineJobGuard; +use tinyagents_orchestration::subagent::JobSettleGuard; use tinyagents_orchestration::subagent::{ - SubAgentJobRegistry, SubagentDriver, SubagentError, SubagentOutcomeKind, SubagentRequest, + SubAgentJobRegistry, SubagentDriver, SubagentError, SubagentRequest, }; fn status_tool_result( @@ -59,10 +60,9 @@ type OwnerScope = dyn Fn(u64) -> String + Send + Sync; /// A typed-parent dispatcher that runs through [`SubagentDriver`]. /// /// Hosts construct the validated request with `request_factory`; this keeps -/// authorization and execution policy in the host planner/executor. A call -/// requesting background mode is detached only when `delivery_target` returns -/// a target. Otherwise the child runs inline and the result is returned to the -/// caller. +/// authorization and execution policy in the host planner/executor. Every call +/// starts a detached job and returns its task reference immediately. The +/// optional delivery target controls whether host completion callbacks run. pub struct DriverSubAgentTool<State, Ctx: Send + 'static, H: Send + 'static> { name: String, description: String, @@ -105,14 +105,9 @@ impl<State, Ctx: Send + 'static, H: Send + 'static> DriverSubAgentTool<State, Ct parameters: json!({ "type": "object", "required": ["input"], - "properties": { - "input": { "type": "string" }, - "mode": { - "type": "string", - "enum": ["background", "inline"], - "description": "Background returns a task id when a delivery target exists; otherwise it runs inline." - } - } + "properties": { + "input": { "type": "string" } + } }), driver, request_factory: Arc::new(request_factory), @@ -231,28 +226,12 @@ impl<State, Ctx: Send + 'static, H: Send + 'static> DriverSubAgentTool<State, Ct self } - fn requested_mode( - args: &Value, - ) -> Result<tinyagents_orchestration::subagent::SubAgentMode, String> { - match args.get("mode") { - None | Some(Value::Null) => { - Ok(tinyagents_orchestration::subagent::SubAgentMode::Background) - } - Some(Value::String(mode)) if mode == "background" => { - Ok(tinyagents_orchestration::subagent::SubAgentMode::Background) - } - Some(Value::String(mode)) if mode == "inline" => { - Ok(tinyagents_orchestration::subagent::SubAgentMode::Inline) - } - Some(_) => Err("`mode` must be \"background\" or \"inline\"".into()), - } - } - async fn execute_request( &self, mut request: SubagentRequest<Ctx, H>, target: Option<String>, owner: String, + parent_run_id: tinyagents_harness::ids::RunId, ) -> anyhow::Result<ToolResult> where Ctx: Send + 'static, @@ -261,16 +240,20 @@ impl<State, Ctx: Send + 'static, H: Send + 'static> DriverSubAgentTool<State, Ct let task_id = request.task_id().to_owned(); let bridged_job = if let Some(jobs) = &self.jobs { let continuing = request.resume().is_some(); - let (id, steering, run_queue, cancellation, created) = - jobs.register_driver_job(DriverJobRegistration { - key: request.task_key().clone(), - agent: request.target().unwrap_or(&self.name).to_owned(), - owner: owner.clone(), - cancellation: request.run_context().cancellation.clone(), - steering: request.run_context().steering.clone(), - run_queue: request.run_context().run_queue.clone(), - continuing, - })?; + let (id, steering, run_queue, cancellation, created) = jobs + .register_driver_job_with_parent_run( + DriverJobRegistration { + key: request.task_key().clone(), + agent: request.target().unwrap_or(&self.name).to_owned(), + owner: owner.clone(), + workspace: request.run_context().workspace.clone(), + cancellation: request.run_context().cancellation.clone(), + steering: request.run_context().steering.clone(), + run_queue: request.run_context().run_queue.clone(), + continuing, + }, + parent_run_id.clone(), + )?; if created { jobs.mark_running(&id); } @@ -286,14 +269,11 @@ impl<State, Ctx: Send + 'static, H: Send + 'static> DriverSubAgentTool<State, Ct let job_guard = bridged_job .as_ref() .filter(|(_, _, created)| *created) - .map(|(jobs, id, _)| InlineJobGuard::new(jobs.clone(), id.clone())); + .map(|(jobs, id, _)| JobSettleGuard::new(jobs.clone(), id.clone())); if let Some((jobs, id, false)) = &bridged_job { - let job = if target.is_some() { - jobs.get(id.as_str()) - .ok_or_else(|| anyhow::anyhow!("registered subagent job disappeared"))? - } else { - jobs.wait_owned(id.as_str(), &owner).await? - }; + let job = jobs + .get(id.as_str()) + .ok_or_else(|| anyhow::anyhow!("registered subagent job disappeared"))?; let status = match job.status { tinyagents_orchestration::subagent::SubAgentJobStatus::Queued => "queued", tinyagents_orchestration::subagent::SubAgentJobStatus::Running => "running", @@ -312,97 +292,68 @@ impl<State, Ctx: Send + 'static, H: Send + 'static> DriverSubAgentTool<State, Ct job.error, )); } - if let Some(target) = target { - let cancellation = request.run_context().cancellation.clone(); - let driver = Arc::clone(&self.driver); - let completion = self.completion.clone(); - let async_completion = self.async_completion.clone(); - let detached_task_id = task_id.clone(); - let completion_target = target.clone(); - let bridge = bridged_job.clone(); - let mut job_guard = job_guard; - tokio::spawn(async move { - match driver.run(request, cancellation).await { - Ok(result) => { - if result.disposition - != tinyagents_orchestration::subagent::SubagentPersistenceDisposition::ObserverCancelled - && let Some(completion) = async_completion - { - completion(completion_target.clone(), Ok(result.clone())).await; - } - if let Some((jobs, id, _)) = &bridge { - jobs.mark_driver_result(id, Ok(result.outcome.clone())); - } - if let Some(guard) = &mut job_guard { - guard.disarm(); - } - if result.should_emit_host_effects() - && let Some(completion) = completion - { - completion(&completion_target, Ok(result.outcome)); - } - } - Err(error) => { - let error = error.to_string(); - if let Some(completion) = async_completion { - completion(completion_target.clone(), Err(error.clone())).await; - } - if let Some((jobs, id, _)) = &bridge { - jobs.mark_driver_result(id, Err(error.clone())); - } - if let Some(guard) = &mut job_guard { - guard.disarm(); - } - if let Some(completion) = completion { - completion(&completion_target, Err(error.clone())); - } - tracing::error!(task_id = %detached_task_id, error, "driver-backed subagent failed"); - } - } - }); - return Ok(ToolResult::json(json!({ - "task_id": public_task_id, - "status": "running", - "delivery_target": target, - }))); - } let cancellation = request.run_context().cancellation.clone(); + let driver = Arc::clone(&self.driver); + let completion = self.completion.clone(); + let async_completion = self.async_completion.clone(); + let detached_task_id = task_id.clone(); + let bridge = bridged_job.clone(); + let callback_target = target.clone(); let mut job_guard = job_guard; - match self.driver.run(request, cancellation).await { - Ok(result) => { - let outcome = result.outcome; - if let Some((jobs, id, _)) = &bridged_job { - jobs.mark_driver_result(id, Ok(outcome.clone())); + tokio::spawn(async move { + let attempt_bridge = bridge.clone(); + let attempt_observer = move |run_id: &tinyagents_harness::ids::RunId| { + if let Some((jobs, id, _)) = &attempt_bridge { + jobs.set_attempt_run_id(id, run_id.as_str()); } - if let Some(guard) = &mut job_guard { - guard.disarm(); - } - let (status, error) = match &outcome.status { - SubagentOutcomeKind::Completed => ("completed", None), - SubagentOutcomeKind::AwaitingInput(_) => ("awaiting_input", None), - SubagentOutcomeKind::Incomplete(incomplete) => { - ("incomplete", Some(incomplete.reason.clone())) + }; + match driver + .run_with_attempt_observer(request, cancellation, attempt_observer) + .await + { + Ok(result) => { + if result.disposition + != tinyagents_orchestration::subagent::SubagentPersistenceDisposition::ObserverCancelled + && let (Some(target), Some(callback)) = (&callback_target, &async_completion) + { + callback(target.clone(), Ok(result.clone())).await; + } + if let Some((jobs, id, _)) = &bridge { + jobs.mark_driver_result(id, Ok(result.outcome.clone())); + } + if let Some(guard) = &mut job_guard { + guard.disarm(); + } + if result.should_emit_host_effects() + && let (Some(target), Some(callback)) = (&callback_target, &completion) + { + callback(target, Ok(result.outcome)); } - SubagentOutcomeKind::Cancelled => ("cancelled", Some("cancelled".into())), - }; - Ok(status_tool_result( - public_task_id, - status, - Some(outcome.output), - error, - )) - } - Err(error) => { - if let Some((jobs, id, _)) = &bridged_job { - jobs.mark_driver_result(id, Err(error.to_string())); } - if let Some(guard) = &mut job_guard { - guard.disarm(); + Err(error) => { + let error = error.to_string(); + if let (Some(target), Some(callback)) = (&callback_target, &async_completion) { + callback(target.clone(), Err(error.clone())).await; + } + if let Some((jobs, id, _)) = &bridge { + jobs.mark_driver_result(id, Err(error.clone())); + } + if let Some(guard) = &mut job_guard { + guard.disarm(); + } + if let (Some(target), Some(callback)) = (&callback_target, &completion) { + callback(target, Err(error.clone())); + } + tracing::error!(task_id = %detached_task_id, error, "driver-backed subagent failed"); } - Ok(ToolResult::error(format!("subagent driver: {error}"))) } + }); + let mut response = json!({ "task_id": public_task_id, "status": "running" }); + if let Some(target) = target { + response["delivery_target"] = Value::String(target); } + Ok(ToolResult::json(response)) } } @@ -437,21 +388,18 @@ where _options: tinytools::ToolCallOptions, parent: &RunContext<Ctx>, ) -> anyhow::Result<ToolResult> { - let requested = match Self::requested_mode(&arguments) { - Ok(mode) => mode, - Err(error) => return Ok(ToolResult::error(error)), - }; let request = match (self.request_factory)(state, parent, &arguments, &call_id) { Ok(request) => request, Err(error) => return Ok(ToolResult::error(format!("subagent request: {error}"))), }; - let target = if requested == tinyagents_orchestration::subagent::SubAgentMode::Background { - (self.delivery_target)(&request) - } else { - None - }; - self.execute_request(request, target, (self.owner_scope)(parent.instance_id())) - .await + let target = (self.delivery_target)(&request); + self.execute_request( + request, + target, + (self.owner_scope)(parent.instance_id()), + parent.run_id().clone(), + ) + .await } } diff --git a/crates/tinyagents-harness-tools/src/subagents/driver_tool_tests.rs b/crates/tinyagents-harness-tools/src/subagents/driver_tool_tests.rs index 7767e57ef..2662bedbf 100644 --- a/crates/tinyagents-harness-tools/src/subagents/driver_tool_tests.rs +++ b/crates/tinyagents-harness-tools/src/subagents/driver_tool_tests.rs @@ -45,6 +45,101 @@ impl ChatModel<()> for CloseCorrectionModel { struct Planner; +struct RecordingRequestPlanner { + requests: Arc<std::sync::Mutex<Vec<(SubagentTaskKey, String)>>>, +} + +struct ReplacingPlanner; + +struct ReplacingRetryPlanner; + +#[async_trait] +impl SubagentPlanner<(), ()> for ReplacingPlanner { + async fn prepare( + &self, + request: SubagentRequest<(), ()>, + ) -> std::result::Result<PreparedSubagent<()>, SubagentError> { + let parts = request.into_parts(); + Ok(PreparedSubagent::new( + parts.task_key.task_id, + "worker", + vec![Message::user(parts.input)], + ToolSnapshot::new(vec![]).unwrap(), + RunContext::new(RunConfig::new("planner-replaced-actual"), ()), + )) + } +} + +#[async_trait] +impl SubagentPlanner<(), ()> for ReplacingRetryPlanner { + async fn prepare( + &self, + request: SubagentRequest<(), ()>, + ) -> std::result::Result<PreparedSubagent<()>, SubagentError> { + let parts = request.into_parts(); + let retry_parent = parts + .run_context + .child(RunConfig::new("retry-template"), ()) + .map_err(|error| SubagentError::InvalidRequest(error.to_string()))?; + let retry_policy = SubAgentPolicy::default().with_retry( + tinyagents_harness::retry::RetryPolicy::default() + .with_max_attempts(2) + .with_backoff_sleep(false), + ); + Ok(PreparedSubagent::new( + parts.task_key.task_id, + "worker", + vec![Message::user(parts.input)], + ToolSnapshot::new(vec![]).unwrap(), + RunContext::new(RunConfig::new("planner-replaced-actual"), ()), + ) + .with_policy(retry_policy) + .with_retry_context(Arc::new(move |attempt| { + retry_parent + .child(RunConfig::new(format!("custom-retry-{attempt}")), ()) + .map_err(|error| SubagentError::InvalidRequest(error.to_string())) + }))) + } +} + +struct RetryThenNestedExecutor { + nested_tool: Arc<dyn ToolDispatch<(), ()>>, + attempts: std::sync::atomic::AtomicUsize, +} + +#[async_trait] +impl SubagentExecutor<()> for RetryThenNestedExecutor { + async fn execute( + &self, + execution: SubagentExecution<()>, + ) -> std::result::Result<SubagentOutcome, SubagentError> { + if self + .attempts + .fetch_add(1, std::sync::atomic::Ordering::SeqCst) + == 0 + { + return Err(SubagentError::Transient { + message: "retry before nested spawn".into(), + tools_ran: false, + }); + } + ToolDispatch::<(), ()>::execute( + self.nested_tool.as_ref(), + &(), + CallId::new("nested-from-custom-retry"), + json!({"input":"nested child"}), + tinytools::ToolCallOptions::default(), + &execution.prepared.run_context, + ) + .await + .map_err(|error| SubagentError::InvalidRequest(error.to_string()))?; + Ok(SubagentOutcome::completed( + execution.prepared.task_id, + "retried and spawned nested child", + )) + } +} + #[async_trait] impl SubagentPlanner<(), ()> for Planner { async fn prepare( @@ -62,6 +157,43 @@ impl SubagentPlanner<(), ()> for Planner { } } +#[async_trait] +impl SubagentPlanner<(), ()> for RecordingRequestPlanner { + async fn prepare( + &self, + request: SubagentRequest<(), ()>, + ) -> std::result::Result<PreparedSubagent<()>, SubagentError> { + let parts = request.into_parts(); + self.requests + .lock() + .unwrap() + .push((parts.task_key.clone(), parts.input.clone())); + Ok(PreparedSubagent::new( + parts.task_key.task_id, + "worker", + vec![Message::user(parts.input)], + ToolSnapshot::new(vec![]).unwrap(), + parts.run_context, + )) + } +} + +struct TimeoutPlanner; + +#[async_trait] +impl SubagentPlanner<(), ()> for TimeoutPlanner { + async fn prepare( + &self, + request: SubagentRequest<(), ()>, + ) -> Result<PreparedSubagent<()>, SubagentError> { + Planner.prepare(request).await.map(|prepared| { + prepared.with_policy( + SubAgentPolicy::default().with_timeout(std::time::Duration::from_millis(10)), + ) + }) + } +} + struct Executor; #[async_trait] @@ -231,6 +363,52 @@ fn selector_test_session() -> DurableSubagentSession { } } +#[tokio::test] +async fn awaiting_user_queries_preserve_the_durable_task_key() { + let mut session = selector_test_session(); + session.task_key = "analysis-follow-up".into(); + session.current_task_id = Some("task-follow-up".into()); + session.status = DurableSubagentStatus::AwaitingUser; + session.latest_history = Some(vec![Message::user("Which environment?")]); + let session_id = session.subagent_session_id.clone(); + let store = Arc::new(SingleSessionStore(Arc::new(std::sync::Mutex::new(session)))); + let parent = RunContext::new(RunConfig::new("parent"), ()); + let jobs = SubAgentJobsTool::new(SubAgentJobRegistry::new()) + .with_sessions(store, |_| "durable-parent".into()); + + let individual = ToolDispatch::<(), ()>::execute( + &jobs, + &(), + CallId::new("await-one"), + json!({"action":"await","job_id":session_id}), + tinytools::ToolCallOptions::default(), + &parent, + ) + .await + .unwrap(); + let individual: Value = serde_json::from_str(&individual.output()).unwrap(); + let individual_envelope = individual["awaiting_user_envelope"].as_str().unwrap(); + let individual_envelope = SubagentAwaitingUser::parse(individual_envelope).unwrap(); + assert_eq!(individual_envelope.task_key, "analysis-follow-up"); + + let collected = ToolDispatch::<(), ()>::execute( + &jobs, + &(), + CallId::new("await-all"), + json!({"action":"await"}), + tinytools::ToolCallOptions::default(), + &parent, + ) + .await + .unwrap(); + let collected: Value = serde_json::from_str(&collected.output()).unwrap(); + let collected_envelope = collected["awaiting_user"][0]["awaiting_user_envelope"] + .as_str() + .unwrap(); + let collected_envelope = SubagentAwaitingUser::parse(collected_envelope).unwrap(); + assert_eq!(collected_envelope.task_key, "analysis-follow-up"); +} + #[tokio::test] async fn model_can_correct_foreign_durable_close_and_continue_to_owned_close() { let session = Arc::new(std::sync::Mutex::new(selector_test_session())); @@ -241,7 +419,9 @@ async fn model_can_correct_foreign_durable_close_and_continue_to_owned_close() { .recover_selector_errors(); let model = Arc::new(CloseCorrectionModel(std::sync::atomic::AtomicUsize::new(0))); let mut harness = AgentHarness::new(); - harness.register_model("script", model.clone()); + harness.with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named("script", model.clone()), + )); harness.register_tool_dispatch(Arc::new(jobs)); let run = harness @@ -396,6 +576,7 @@ async fn fixed_query_alias_ignores_action_override_and_keeps_host_declaration() key: SubagentTaskKey::from_context(&parent, "fixed-query", None), agent: "worker".into(), owner: parent.instance_id().to_string(), + workspace: None, cancellation: CancellationToken::new(), steering: None, run_queue: None, @@ -448,6 +629,7 @@ async fn message_alias_routes_unique_contents_to_queue_lanes_or_default_steering key: SubagentTaskKey::from_context(&parent, "message-routing", None), agent: "worker".into(), owner: parent.instance_id().to_string(), + workspace: None, cancellation: CancellationToken::new(), steering: Some(steering), run_queue: Some(queue), @@ -560,7 +742,7 @@ async fn bounded_await_observes_terminal_state_after_deadline_before_same_future let task_key = SubagentTaskKey::from_context(&parent, "bounded-await-race", None); let request_key = task_key.clone(); let tool = DriverSubAgentTool::new( - "spawn_async_subagent", + "spawn_subagent", "Run the worker in the background.", driver, move |_: &(), parent: &RunContext<()>, args: &Value, _: &CallId| { @@ -582,14 +764,14 @@ async fn bounded_await_observes_terminal_state_after_deadline_before_same_future .with_parameters(json!({ "type":"object", "required":["agent_id","prompt"], - "properties":{"agent_id":{"type":"string"},"prompt":{"type":"string"},"mode":{"type":"string"}} + "properties":{"agent_id":{"type":"string"},"prompt":{"type":"string"}} })) .with_job_registry(jobs.clone()); let spawned = ToolDispatch::<(), ()>::execute( &tool, &(), CallId::new("spawn-bounded-await-race"), - json!({"agent_id":"worker","prompt":"finish at the deadline","mode":"background"}), + json!({"agent_id":"worker","prompt":"finish at the deadline"}), tinytools::ToolCallOptions::default(), &parent, ) @@ -936,8 +1118,72 @@ impl SubagentPersistence for TerminalReplayPersistence { } } +#[tokio::test(start_paused = true)] +async fn detached_driver_timeout_is_recorded_as_incomplete_job_status() { + let jobs = SubAgentJobRegistry::new(); + let started = Arc::new(tokio::sync::Notify::new()); + let release = Arc::new(tokio::sync::Semaphore::new(0)); + let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let driver = Arc::new( + SubagentDriver::new(SubagentCapabilities { + planner: Some(Arc::new(TimeoutPlanner)), + executor: Some(Arc::new(BlockingExecutor { + started: started.clone(), + release, + calls, + })), + persistence: Some(Arc::new(Persistence)), + }) + .unwrap(), + ); + let tool = DriverSubAgentTool::new( + "spawn_subagent", + "worker description", + driver, + |_: &(), parent: &RunContext<()>, args: &Value, _: &CallId| { + let child = parent + .child(RunConfig::new("timeout-child"), ()) + .map_err(|error| SubagentError::InvalidRequest(error.to_string()))?; + SubagentRequest::fresh_from_parent( + parent, + child, + "timeout-task", + (), + args["prompt"].as_str().unwrap_or_default(), + None, + ) + }, + |_| None, + ) + .with_job_registry(jobs.clone()); + let parent = RunContext::new(RunConfig::new("timeout-parent"), ()); + let spawned = ToolDispatch::<(), ()>::execute( + &tool, + &(), + CallId::new("timeout-spawn"), + json!({"prompt":"work"}), + tinytools::ToolCallOptions::default(), + &parent, + ) + .await + .unwrap(); + let value: Value = serde_json::from_str(&spawned.output()).unwrap(); + started.notified().await; + tokio::time::advance(std::time::Duration::from_millis(11)).await; + let job = jobs + .wait_owned( + value["task_id"].as_str().unwrap(), + parent.instance_id().to_string(), + ) + .await + .unwrap(); + assert_eq!(job.status, SubAgentJobStatus::Incomplete); + assert_eq!(job.incomplete_kind, Some(IncompleteKind::Timeout)); +} + #[tokio::test] -async fn background_request_without_delivery_target_runs_inline() { +async fn background_request_without_delivery_target_spawns_and_can_be_waited_separately() { + let jobs = SubAgentJobRegistry::new(); let tool = DriverSubAgentTool::new( "worker", "worker description", @@ -956,14 +1202,15 @@ async fn background_request_without_delivery_target_runs_inline() { ) }, |_request| None, - ); + ) + .with_job_registry(jobs.clone()); let parent = RunContext::new(RunConfig::new("parent"), ()); let result = ToolDispatch::<(), ()>::execute( &tool, &(), CallId::new("call-1"), - json!({"input":"do work", "mode":"background"}), + json!({"input":"do work"}), tinytools::ToolCallOptions::default(), &parent, ) @@ -972,8 +1219,16 @@ async fn background_request_without_delivery_target_runs_inline() { assert!(!result.is_error); let value: Value = serde_json::from_str(&result.output()).unwrap(); - assert_eq!(value["status"], "completed"); - assert_eq!(value["output"], "driver result"); + assert_eq!(value["status"], "running"); + let job = jobs + .wait_owned( + value["task_id"].as_str().unwrap(), + &parent.instance_id().to_string(), + ) + .await + .unwrap(); + assert_eq!(job.status, SubAgentJobStatus::Completed); + assert_eq!(job.output.as_deref(), Some("driver result")); } #[tokio::test] @@ -1003,7 +1258,7 @@ async fn background_request_with_delivery_target_returns_task_reference() { &tool, &(), CallId::new("call-2"), - json!({"input":"do work", "mode":"background"}), + json!({"input":"do work"}), tinytools::ToolCallOptions::default(), &parent, ) @@ -1014,7 +1269,6 @@ async fn background_request_with_delivery_target_returns_task_reference() { let value: Value = serde_json::from_str(&result.output()).unwrap(); assert_eq!(value["status"], "running"); assert_eq!(value["delivery_target"], "parent-thread"); - assert_eq!(value["task_id"], "worker-task"); } #[tokio::test(flavor = "current_thread")] @@ -1062,7 +1316,7 @@ async fn awaited_background_completion_commits_before_jobs_become_terminal() { &tool, &(), CallId::new("call-await-completion"), - json!({"input":"do work", "mode":"background"}), + json!({"input":"do work"}), tinytools::ToolCallOptions::default(), &parent, ) @@ -1144,7 +1398,7 @@ async fn awaited_completion_repairs_replayed_terminal_before_jobs_become_termina &tool, &(), CallId::new("call-replayed-completion"), - json!({"input":"do work", "mode":"background"}), + json!({"input":"do work"}), tinytools::ToolCallOptions::default(), &parent, ) @@ -1183,7 +1437,7 @@ fn driver_subagent_tool_preserves_host_name_description_and_schema() { let schema = json!({ "type":"object", "required":["agent_id","prompt"], - "properties":{"agent_id":{"type":"string"},"prompt":{"type":"string"},"mode":{"type":"string"}} + "properties":{"agent_id":{"type":"string"},"prompt":{"type":"string"}} }); let tool = DriverSubAgentTool::new( "delegate_to", @@ -1216,10 +1470,22 @@ fn driver_subagent_tool_preserves_host_name_description_and_schema() { #[tokio::test] async fn driver_subagent_tool_routes_execution_through_the_driver() { + let jobs = SubAgentJobRegistry::new(); + let requests = Arc::new(std::sync::Mutex::new(Vec::new())); + let driver = Arc::new( + SubagentDriver::new(SubagentCapabilities { + planner: Some(Arc::new(RecordingRequestPlanner { + requests: requests.clone(), + })), + executor: Some(Arc::new(Executor)), + persistence: Some(Arc::new(Persistence)), + }) + .unwrap(), + ); let tool = DriverSubAgentTool::new( "subagent", "worker description", - driver(), + driver, |_: &(), parent: &RunContext<()>, args: &Value, _: &CallId| { let child = parent .child(RunConfig::new("worker-child"), ()) @@ -1234,24 +1500,204 @@ async fn driver_subagent_tool_routes_execution_through_the_driver() { ) }, |_| None, - ); + ) + .with_job_registry(jobs.clone()); let parent = RunContext::new(RunConfig::new("parent"), ()); let result = ToolDispatch::<(), ()>::execute( &tool, &(), CallId::new("driver-route"), - json!({"input":"do work", "mode":"inline"}), + json!({"input":"do work"}), + tinytools::ToolCallOptions::default(), + &parent, + ) + .await + .unwrap(); + assert!(!result.is_error); + let response: Value = serde_json::from_str(&result.output()).unwrap(); + assert_eq!(response["status"], "running"); + let task_id = response["task_id"].as_str().expect("spawned job id"); + assert!(!task_id.is_empty()); + assert!(response.get("output").is_none()); + + let job = jobs + .wait_owned(task_id, parent.instance_id().to_string()) + .await + .unwrap(); + assert_eq!(job.status, SubAgentJobStatus::Completed); + assert_eq!(job.agent, "subagent"); + assert_eq!(job.output.as_deref(), Some("driver result")); + + { + let requests = requests.lock().unwrap(); + assert_eq!(requests.len(), 1); + let (request_key, request_input) = &requests[0]; + assert_eq!(request_key.task_id, "driver-task"); + assert_eq!(request_key.parent_run_id, parent.run_id().as_str()); + assert_eq!(request_input, "do work"); + } + + let settled = jobs + .settle_root_run_descendants( + parent.run_id(), + parent.instance_id().to_string(), + tokio::time::Instant::now() + std::time::Duration::from_secs(30), + &CancellationToken::new(), + ) + .await + .unwrap(); + assert_eq!(settled.len(), 1); + assert_eq!(settled[0].status, SubAgentJobStatus::Completed); +} + +#[tokio::test] +async fn driver_job_records_the_planners_actual_run_context_for_descendant_settlement() { + let jobs = SubAgentJobRegistry::new(); + let driver = Arc::new( + SubagentDriver::new(SubagentCapabilities { + planner: Some(Arc::new(ReplacingPlanner)), + executor: Some(Arc::new(Executor)), + persistence: Some(Arc::new(Persistence)), + }) + .unwrap(), + ); + let tool = DriverSubAgentTool::new( + "subagent", + "worker description", + driver, + |_: &(), parent: &RunContext<()>, _: &Value, _: &CallId| { + let child = parent + .child(RunConfig::new("worker-request-child"), ()) + .map_err(|error| SubagentError::InvalidRequest(error.to_string()))?; + SubagentRequest::fresh_from_parent( + parent, + child, + "planner-replaced-task", + (), + "input", + None, + ) + }, + |_| None, + ) + .with_job_registry(jobs.clone()); + let parent = RunContext::new(RunConfig::new("planner-replaced-parent"), ()); + ToolDispatch::<(), ()>::execute( + &tool, + &(), + CallId::new("planner-replaced-driver"), + json!({"input":"ignored"}), tinytools::ToolCallOptions::default(), &parent, ) .await .unwrap(); + + let settled = jobs + .settle_root_run_descendants( + parent.run_id(), + parent.instance_id().to_string(), + tokio::time::Instant::now() + std::time::Duration::from_secs(30), + &CancellationToken::new(), + ) + .await + .unwrap(); + assert_eq!(settled.len(), 1); assert_eq!( - serde_json::from_str::<Value>(&result.output()).unwrap()["output"], - "driver result" + settled[0].subagent_run_id.as_deref(), + Some("planner-replaced-actual") ); } +#[tokio::test] +async fn driver_job_tracks_planner_replacement_custom_retry_and_nested_child_lineage() { + let jobs = SubAgentJobRegistry::new(); + let nested_driver = driver(); + let nested_tool: Arc<dyn ToolDispatch<(), ()>> = Arc::new( + DriverSubAgentTool::new( + "nested_subagent", + "Run a nested worker.", + nested_driver, + |_: &(), parent: &RunContext<()>, _: &Value, _: &CallId| { + let child = parent + .child(RunConfig::new("nested-worker-child"), ()) + .map_err(|error| SubagentError::InvalidRequest(error.to_string()))?; + SubagentRequest::fresh_from_parent( + parent, + child, + "nested-task", + (), + "nested input", + None, + ) + }, + |_| None, + ) + .with_job_registry(jobs.clone()) + .with_owner_scope(Arc::new(|_| "owner".to_owned())), + ); + let driver = Arc::new( + SubagentDriver::new(SubagentCapabilities { + planner: Some(Arc::new(ReplacingRetryPlanner)), + executor: Some(Arc::new(RetryThenNestedExecutor { + nested_tool, + attempts: std::sync::atomic::AtomicUsize::new(0), + })), + persistence: Some(Arc::new(Persistence)), + }) + .unwrap(), + ); + let tool = DriverSubAgentTool::new( + "subagent", + "worker description", + driver, + |_: &(), parent: &RunContext<()>, _: &Value, _: &CallId| { + let child = parent + .child(RunConfig::new("worker-request-child"), ()) + .map_err(|error| SubagentError::InvalidRequest(error.to_string()))?; + SubagentRequest::fresh_from_parent( + parent, + child, + "retry-nested-task", + (), + "input", + None, + ) + }, + |_| None, + ) + .with_job_registry(jobs.clone()) + .with_owner_scope(Arc::new(|_| "owner".to_owned())); + let parent = RunContext::new(RunConfig::new("retry-nested-parent"), ()); + ToolDispatch::<(), ()>::execute( + &tool, + &(), + CallId::new("retry-nested-driver"), + json!({"input":"ignored"}), + tinytools::ToolCallOptions::default(), + &parent, + ) + .await + .unwrap(); + + let settled = jobs + .settle_root_run_descendants( + parent.run_id(), + "owner", + tokio::time::Instant::now() + std::time::Duration::from_secs(30), + &CancellationToken::new(), + ) + .await + .unwrap(); + assert_eq!(settled.len(), 2, "all registered jobs: {:?}", jobs.list()); + assert!( + settled + .iter() + .all(|job| job.status == SubAgentJobStatus::Completed) + ); + assert!(settled.iter().any(|job| job.agent == "nested_subagent")); +} + #[tokio::test] async fn driver_subagent_tool_routes_detached_completion_to_host_callback() { let (tx, rx) = tokio::sync::oneshot::channel(); @@ -1287,7 +1733,7 @@ async fn driver_subagent_tool_routes_detached_completion_to_host_callback() { &tool, &(), CallId::new("driver-background"), - json!({"input":"do work", "mode":"background"}), + json!({"input":"do work"}), tinytools::ToolCallOptions::default(), &parent, ) @@ -1354,10 +1800,10 @@ async fn driver_alias_registry_controls_the_same_scoped_job_end_to_end() { let schema = json!({ "type":"object", "required":["agent_id","prompt"], - "properties":{"agent_id":{"type":"string"},"prompt":{"type":"string"},"mode":{"type":"string"}} + "properties":{"agent_id":{"type":"string"},"prompt":{"type":"string"}} }); let tool = DriverSubAgentTool::new( - "spawn_async_subagent", + "spawn_subagent", "Run the allowed worker in the background.", driver, |_: &(), parent: &RunContext<()>, args: &Value, _: &CallId| { @@ -1398,7 +1844,7 @@ async fn driver_alias_registry_controls_the_same_scoped_job_end_to_end() { &tool, &(), CallId::new("spawn"), - json!({"agent_id":"worker","prompt":"do work","mode":"background"}), + json!({"agent_id":"worker","prompt":"do work"}), tinytools::ToolCallOptions::default(), &parent, ) @@ -1412,7 +1858,7 @@ async fn driver_alias_registry_controls_the_same_scoped_job_end_to_end() { &tool, &(), CallId::new("spawn-again"), - json!({"agent_id":"worker","prompt":"do work","mode":"background"}), + json!({"agent_id":"worker","prompt":"do work"}), tinytools::ToolCallOptions::default(), &parent, ) @@ -1550,7 +1996,7 @@ async fn driver_alias_registry_controls_the_same_scoped_job_end_to_end() { &tool, &(), CallId::new("spawn-foreign"), - json!({"agent_id":"worker","prompt":"do work","mode":"background"}), + json!({"agent_id":"worker","prompt":"do work"}), tinytools::ToolCallOptions::default(), &foreign, ) @@ -1678,7 +2124,7 @@ async fn driver_alias_registry_controls_the_same_scoped_job_end_to_end() { &tool, &(), CallId::new("spawn-terminal-replay"), - json!({"agent_id":"worker","prompt":"do work","mode":"background"}), + json!({"agent_id":"worker","prompt":"do work"}), tinytools::ToolCallOptions::default(), &parent, ) @@ -1723,7 +2169,7 @@ async fn jobs_tool_treats_a_null_action_as_query() { &tool, &(), CallId::new("spawn-null-action"), - json!({"input":"work", "mode":"background"}), + json!({"input":"work"}), tinytools::ToolCallOptions::default(), &parent, ) @@ -1777,7 +2223,7 @@ async fn oversized_message_request_ids_are_rejected_by_the_moved_shell() { &tool, &(), CallId::new("spawn-oversized-request"), - json!({"input":"work", "mode":"background"}), + json!({"input":"work"}), tinytools::ToolCallOptions::default(), &parent, ) @@ -1869,7 +2315,7 @@ async fn driver_alias_pause_stays_nonterminal_and_continues_on_the_same_job() { "properties":{"agent_id":{"type":"string"},"prompt":{"type":"string"},"continue":{"type":"boolean"}} }); let tool = DriverSubAgentTool::new( - "spawn_async_subagent", + "spawn_subagent", "worker description", driver, move |_: &(), parent: &RunContext<()>, args: &Value, _: &CallId| { @@ -1898,7 +2344,7 @@ async fn driver_alias_pause_stays_nonterminal_and_continues_on_the_same_job() { .with_target("worker"); Ok(request) }, - |_| Some("parent-thread".into()), + |_| None, ) .with_parameters(schema) .with_job_registry(jobs.clone()) @@ -1908,7 +2354,7 @@ async fn driver_alias_pause_stays_nonterminal_and_continues_on_the_same_job() { &tool, &(), CallId::new("spawn-pausing-child"), - json!({"agent_id":"worker","prompt":"wait","mode":"background"}), + json!({"agent_id":"worker","prompt":"wait"}), tinytools::ToolCallOptions::default(), &parent, ) @@ -1965,7 +2411,7 @@ async fn driver_alias_pause_stays_nonterminal_and_continues_on_the_same_job() { &tool, &(), CallId::new("continue-child"), - json!({"agent_id":"worker","prompt":"answer","continue":true,"mode":"background"}), + json!({"agent_id":"worker","prompt":"answer","continue":true}), tinytools::ToolCallOptions::default(), &next_turn, ) @@ -2007,7 +2453,11 @@ async fn cancel_driver_job_by_full_key_checks_owner_and_waits_for_settlement() { let key = SubagentTaskKey::from_context(&parent, "cancel-task", None); let request_key = key.clone(); let owner_scope: Arc<dyn Fn(u64) -> String + Send + Sync> = Arc::new(|_| "parent-a".into()); - let tool = DriverSubAgentTool::new("spawn_async_subagent", "worker description", driver, move |_: &(), parent: &RunContext<()>, args: &Value, _: &CallId| { + let tool = DriverSubAgentTool::new( + "spawn_subagent", + "worker description", + driver, + move |_: &(), parent: &RunContext<()>, args: &Value, _: &CallId| { let child = parent .child(RunConfig::new("child-run"), ()) .map_err(|error| SubagentError::InvalidRequest(error.to_string()))?; @@ -2020,16 +2470,21 @@ async fn cancel_driver_job_by_full_key_checks_owner_and_waits_for_settlement() { None, ) .map(|request| request.with_target("worker")) - }, |_| Some("thread-a".into())).with_parameters(json!({ + }, + |_| Some("thread-a".into()), + ) + .with_parameters(json!({ "type":"object", "required":["agent_id","prompt"], - "properties":{"agent_id":{"type":"string"},"prompt":{"type":"string"},"mode":{"type":"string"}} - })).with_job_registry(jobs.clone()).with_owner_scope(owner_scope); + "properties":{"agent_id":{"type":"string"},"prompt":{"type":"string"}} + })) + .with_job_registry(jobs.clone()) + .with_owner_scope(owner_scope); let spawned = ToolDispatch::<(), ()>::execute( &tool, &(), CallId::new("spawn-cancellable"), - json!({"agent_id":"worker","prompt":"do work","mode":"background"}), + json!({"agent_id":"worker","prompt":"do work"}), tinytools::ToolCallOptions::default(), &parent, ) @@ -2111,7 +2566,7 @@ async fn parent_cancellation_cascades_to_running_driver_jobs() { &tool, &(), CallId::new("spawn-parent-cancel"), - json!({"input":"work", "mode":"background"}), + json!({"input":"work"}), tinytools::ToolCallOptions::default(), &parent, ) @@ -2174,7 +2629,7 @@ async fn cancelling_one_driver_job_leaves_its_sibling_running() { &tool, &(), CallId::new(call_id), - json!({"task_id":task_id,"input":"work","mode":"background"}), + json!({"task_id":task_id,"input":"work"}), tinytools::ToolCallOptions::default(), &parent, ) @@ -2262,7 +2717,7 @@ fn inline_blocking_alias( } #[tokio::test(start_paused = true)] -async fn concurrent_inline_driver_alias_coalesces_and_awaits_shared_result() { +async fn concurrent_async_driver_aliases_coalesce_without_waiting_for_completion() { let jobs = SubAgentJobRegistry::new(); let started = Arc::new(tokio::sync::Notify::new()); let release = Arc::new(tokio::sync::Semaphore::new(0)); @@ -2282,34 +2737,29 @@ async fn concurrent_inline_driver_alias_coalesces_and_awaits_shared_result() { &parent, ) }; - let mut leader = Box::pin(dispatch("inline-leader")); - tokio::select! { - biased; - result = &mut leader => panic!("leader completed before being released: {result:?}"), - _ = started.notified() => {}, - } - let mut follower = Box::pin(dispatch("inline-follower")); - tokio::select! { - biased; - result = &mut follower => panic!("coalesced inline call returned before the shared result: {result:?}"), - _ = tokio::task::yield_now() => {}, - } + let leader = dispatch("async-leader").await.unwrap(); + let leader_value: Value = serde_json::from_str(&leader.output()).unwrap(); + assert_eq!(leader_value["status"], "running"); + started.notified().await; + let follower = dispatch("async-follower").await.unwrap(); + let follower_value: Value = serde_json::from_str(&follower.output()).unwrap(); + assert_eq!(follower_value["status"], "running"); + assert_eq!(follower_value["task_id"], leader_value["task_id"]); assert_eq!(calls.load(std::sync::atomic::Ordering::SeqCst), 1); release.add_permits(1); - let leader = leader.await.unwrap(); - let follower = follower.await.unwrap(); - for result in [leader, follower] { - assert!(!result.is_error); - assert_eq!( - serde_json::from_str::<Value>(&result.output()).unwrap()["status"], - "completed" - ); - } + let settled = jobs + .wait_owned( + leader_value["task_id"].as_str().unwrap(), + &parent.instance_id().to_string(), + ) + .await + .unwrap(); + assert_eq!(settled.status, SubAgentJobStatus::Completed); assert_eq!(calls.load(std::sync::atomic::Ordering::SeqCst), 1); } #[tokio::test(flavor = "current_thread", start_paused = true)] -async fn immediate_inline_replay_waits_for_background_winner_and_emits_one_completion() { +async fn immediate_async_replay_returns_job_snapshot_and_emits_one_completion() { let jobs = SubAgentJobRegistry::new(); let started = Arc::new(tokio::sync::Notify::new()); let release = Arc::new(tokio::sync::Semaphore::new(0)); @@ -2322,7 +2772,7 @@ async fn immediate_inline_replay_waits_for_background_winner_and_emits_one_compl "required":["prompt"], "properties":{"prompt":{"type":"string"}} }); - let make_tool = |background: bool| { + let make_tool = || { let request = |_: &(), parent: &RunContext<()>, args: &Value, _: &CallId| { let child = parent .child(RunConfig::new("mixed-child"), ()) @@ -2337,16 +2787,8 @@ async fn immediate_inline_replay_waits_for_background_winner_and_emits_one_compl ) .map(|request| request.with_target("worker")) }; - let name = if background { - "spawn_async_subagent" - } else { - "spawn_subagent" - }; - let delivery = if background { - Some("parent-thread".to_owned()) - } else { - None - }; + let name = "spawn_subagent"; + let delivery = Some("parent-thread".to_owned()); let tool = DriverSubAgentTool::new( name, "worker description", @@ -2356,27 +2798,22 @@ async fn immediate_inline_replay_waits_for_background_winner_and_emits_one_compl ) .with_parameters(schema.clone()) .with_job_registry(jobs.clone()); - if background { - let completions = completions.clone(); - let completion_notified = completion_notified.clone(); - tool.with_completion_handler(move |_, result| { - assert!(result.is_ok()); - completions.fetch_add(1, std::sync::atomic::Ordering::SeqCst); - completion_notified.notify_one(); - }) - } else { - tool - } + let completions = completions.clone(); + let completion_notified = completion_notified.clone(); + tool.with_completion_handler(move |_, result| { + assert!(result.is_ok()); + completions.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + completion_notified.notify_one(); + }) }; - let background = make_tool(true); - let inline = make_tool(false); + let tool = make_tool(); let parent = RunContext::new(RunConfig::new("mixed-parent"), ()); let spawned = ToolDispatch::<(), ()>::execute( - &background, + &tool, &(), CallId::new("background-first"), - json!({"prompt":"same request","mode":"background"}), + json!({"prompt":"same request"}), tinytools::ToolCallOptions::default(), &parent, ) @@ -2387,28 +2824,22 @@ async fn immediate_inline_replay_waits_for_background_winner_and_emits_one_compl .unwrap() .to_owned(); - let inline_result = ToolDispatch::<(), ()>::execute( - &inline, + started.notified().await; + let replay = ToolDispatch::<(), ()>::execute( + &tool, &(), - CallId::new("inline-immediately-after-background"), - json!({"prompt":"same request","mode":"inline"}), + CallId::new("async-replay"), + json!({"prompt":"same request"}), tinytools::ToolCallOptions::default(), &parent, - ); - tokio::pin!(inline_result); - tokio::select! { - biased; - result = &mut inline_result => panic!("inline alias returned before background winner finished: {result:?}"), - _ = started.notified() => {}, - } + ) + .await + .unwrap(); + let replay_value: Value = serde_json::from_str(&replay.output()).unwrap(); + assert_eq!(replay_value["status"], "running"); + assert_eq!(replay_value["task_id"], task_id); assert_eq!(calls.load(std::sync::atomic::Ordering::SeqCst), 1); release.add_permits(1); - let inline_result = inline_result.await.unwrap(); - assert!(!inline_result.is_error); - assert_eq!( - serde_json::from_str::<Value>(&inline_result.output()).unwrap()["status"], - "completed" - ); completion_notified.notified().await; assert_eq!(completions.load(std::sync::atomic::Ordering::SeqCst), 1); assert_eq!( @@ -2418,7 +2849,7 @@ async fn immediate_inline_replay_waits_for_background_winner_and_emits_one_compl } #[tokio::test(start_paused = true)] -async fn dropping_inline_driver_alias_marks_its_registered_job_cancelled() { +async fn dropping_caller_does_not_cancel_async_job_but_owner_can_cancel_it() { let jobs = SubAgentJobRegistry::new(); let started = Arc::new(tokio::sync::Notify::new()); let release = Arc::new(tokio::sync::Semaphore::new(0)); @@ -2428,30 +2859,27 @@ async fn dropping_inline_driver_alias_marks_its_registered_job_cancelled() { blocking_driver(started.clone(), release, calls), ); let parent = RunContext::new(RunConfig::new("drop-parent"), ()); - let mut execution = Box::pin(ToolDispatch::<(), ()>::execute( + let result = ToolDispatch::<(), ()>::execute( &tool, &(), - CallId::new("drop-inline"), - json!({"prompt":"inline"}), + CallId::new("spawn-before-caller-drop"), + json!({"prompt":"async"}), tinytools::ToolCallOptions::default(), &parent, - )); - tokio::select! { - biased; - result = &mut execution => panic!("child completed before drop: {result:?}"), - _ = started.notified() => {}, - } - drop(execution); + ) + .await + .unwrap(); + let value: Value = serde_json::from_str(&result.output()).unwrap(); + started.notified().await; + let job_id = value["task_id"].as_str().unwrap(); + assert_eq!(jobs.get(job_id).unwrap().status, SubAgentJobStatus::Running); + jobs.cancel_owned(job_id, parent.instance_id().to_string()) + .unwrap(); let job = jobs - .list() - .into_iter() - .next() - .expect("registered job remains visible"); + .wait_owned(job_id, parent.instance_id().to_string()) + .await + .unwrap(); assert_eq!(job.status, SubAgentJobStatus::Cancelled); - assert_eq!( - jobs.get(job.id.as_str()).unwrap().status, - SubAgentJobStatus::Cancelled - ); } #[tokio::test] @@ -2466,7 +2894,7 @@ async fn panicking_background_driver_alias_settles_and_replays_as_tool_error() { .unwrap(), ); let tool = DriverSubAgentTool::new( - "spawn_async_subagent", + "spawn_subagent", "worker description", driver, |_: &(), parent: &RunContext<()>, args: &Value, _: &CallId| { @@ -2496,7 +2924,7 @@ async fn panicking_background_driver_alias_settles_and_replays_as_tool_error() { &tool, &(), CallId::new("panic-spawn"), - json!({"prompt":"panic","mode":"background"}), + json!({"prompt":"panic"}), tinytools::ToolCallOptions::default(), &parent, ) @@ -2533,7 +2961,7 @@ async fn panicking_background_driver_alias_settles_and_replays_as_tool_error() { &tool, &(), CallId::new("panic-replay"), - json!({"prompt":"panic","mode":"background"}), + json!({"prompt":"panic"}), tinytools::ToolCallOptions::default(), &parent, ) @@ -2556,6 +2984,7 @@ async fn cancelling_paused_driver_job_settles_before_wait_returns() { key: SubagentTaskKey::from_context(&parent, "paused-cancel", None), agent: "worker".into(), owner: parent.instance_id().to_string(), + workspace: None, cancellation: CancellationToken::new(), steering: None, run_queue: None, diff --git a/crates/tinyagents-harness-tools/src/subagents/jobs.rs b/crates/tinyagents-harness-tools/src/subagents/jobs.rs index 2f8731a81..34d21124d 100644 --- a/crates/tinyagents-harness-tools/src/subagents/jobs.rs +++ b/crates/tinyagents-harness-tools/src/subagents/jobs.rs @@ -528,6 +528,7 @@ impl<State: Send + Sync, Ctx: Send + Sync> ToolDispatch<State, Ctx> for SubAgent .unwrap_or_default(); let envelope = SubagentAwaitingUser { task_id: session.current_task_id.clone().unwrap_or_default(), + task_key: session.task_key.clone(), agent_id: session.agent_id.clone(), worker_thread_id: session.worker_thread_id.clone(), question, @@ -683,6 +684,7 @@ impl<State: Send + Sync, Ctx: Send + Sync> ToolDispatch<State, Ctx> for SubAgent .unwrap_or_default(); let envelope = SubagentAwaitingUser { task_id: session.current_task_id.clone().unwrap_or_default(), + task_key: session.task_key.clone(), agent_id: session.agent_id.clone(), worker_thread_id: session.worker_thread_id.clone(), question, diff --git a/crates/tinyagents-harness-tools/src/thread_locks/README.md b/crates/tinyagents-harness-tools/src/thread_locks/README.md index 185e472fd..88598a13f 100644 --- a/crates/tinyagents-harness-tools/src/thread_locks/README.md +++ b/crates/tinyagents-harness-tools/src/thread_locks/README.md @@ -23,7 +23,7 @@ ever seen. ## Public surface This module is `pub(crate)` — it has no crate-external API. Within -`tinyagents-graph`: +the harness tools runtime: - `ThreadLockMap` — the weak-value lock map. - `ThreadLockMap::new(what: &'static str)` — creates an empty map; `what` diff --git a/crates/tinyagents-harness-tools/src/todos/README.md b/crates/tinyagents-harness-tools/src/todos/README.md index c3a2069cd..0fd378c4e 100644 --- a/crates/tinyagents-harness-tools/src/todos/README.md +++ b/crates/tinyagents-harness-tools/src/todos/README.md @@ -106,7 +106,7 @@ inert while the thread has no plan with open steps. ```rust,ignore use std::sync::Arc; -use tinyagents_graph::{TodoItem, TodoStatus, TodoTool, todo_store}; +use tinyagents_harness_tools::todos::{TodoItem, TodoStatus, TodoTool, todo_store}; use tinyagents_harness::store::{InMemoryStore, Store}; let store: Arc<dyn Store> = Arc::new(InMemoryStore::default()); diff --git a/crates/tinyagents-harness-tools/src/todos/mod.rs b/crates/tinyagents-harness-tools/src/todos/mod.rs index 1310d1859..a1589e759 100644 --- a/crates/tinyagents-harness-tools/src/todos/mod.rs +++ b/crates/tinyagents-harness-tools/src/todos/mod.rs @@ -3,10 +3,10 @@ //! The todo list holds the concrete steps the run is working through: an //! ordered list of [`TodoItem`]s, each `pending`, `in_progress`, //! `completed` or `cancelled`. This module owns the data model and markdown -//! rendering ([`types`]), the step order (`sequence`: one step at a time, -//! top to bottom), harness-[`Store`](tinyagents_harness::store::Store)-backed -//! persistence that keeps it ([`store`]), the model-facing tool ([`tool`]), -//! and the middleware that holds a run to its plan ([`middleware`]). +//! rendering (`types`), the step order (`sequence`: one step at a time, +//! top to bottom), harness `Store`-backed persistence that keeps it (`store`), +//! the model-facing tool (`tool`), and the middleware that holds a run to its +//! plan (`middleware`). //! //! The model writes the whole list once, then ticks steps off with //! `{"complete": N}`; nothing else hangs off an item — no ids, no approvals, diff --git a/crates/tinyagents-harness-tools/src/todos/store.rs b/crates/tinyagents-harness-tools/src/todos/store.rs index 96fa4ab8b..3a165edf6 100644 --- a/crates/tinyagents-harness-tools/src/todos/store.rs +++ b/crates/tinyagents-harness-tools/src/todos/store.rs @@ -1,10 +1,10 @@ //! CRUD for the per-thread todo list, on the harness -//! [`Store`](tinyagents_harness::store::Store). +//! `Store`. //! //! Each thread's list is a single serialized [`TodoList`] value under the //! [`TODOS_NAMESPACE`] namespace, keyed by the hex-encoded thread id. Every //! mutation runs `load → mutate → normalise → put` under a **per-thread async -//! mutex** ([`thread_lock`]) so the read-modify-write is atomic within the +//! mutex** (`thread_lock`) so the read-modify-write is atomic within the //! process (the same single-process caveat as //! [`GoalStore`](crate::goals::GoalStore)). //! diff --git a/crates/tinyagents-harness/Cargo.toml b/crates/tinyagents-harness/Cargo.toml index 2deb92fc0..f707c2f29 100644 --- a/crates/tinyagents-harness/Cargo.toml +++ b/crates/tinyagents-harness/Cargo.toml @@ -29,7 +29,7 @@ serde_json = { workspace = true } sha2 = "0.11" thiserror = "2" tracing = { workspace = true } -tinyagents-definition = { path = "../tinyagents-definition", version = "2.11.0" } +tinyagents-definition = { path = "../tinyagents-definition", version = "3.0.0" } tinytools-agent = { path = "../../vendor/tinytools/crates/tinytools-agent", version = "0.6.0", default-features = false } tinyinference-llm = { path = "../../vendor/tinyinference/crates/tinyinference-llm", version = "0.5.0" } tinyinference-providers = { path = "../../vendor/tinyinference/crates/tinyinference-providers", version = "0.5.0" } @@ -38,8 +38,6 @@ tinyinference-embeddings = { path = "../../vendor/tinyinference/crates/tinyinfer tinyinference-image = { path = "../../vendor/tinyinference/crates/tinyinference-image", version = "0.5.0", optional = true } tinyinference-video = { path = "../../vendor/tinyinference/crates/tinyinference-video", version = "0.5.0", optional = true } tinytools = { path = "../../vendor/tinytools/crates/tinytools", version = "0.6.0" } -# Shared URL admission for opt-in remote multimodal attachments. -tinytools-std = { path = "../../vendor/tinytools/crates/tinytools-std", version = "0.6.0", optional = true } tokio = { workspace = true, features = ["sync", "time", "macros", "rt", "rt-multi-thread", "fs", "io-util", "process"] } # Storage ports (no database client): the `storage-drivers` adapters put # the harness `Store` / `AppendStore` on any tinystoragedrivers backend. @@ -66,7 +64,7 @@ sqlite = [] # ports, so a host's chosen backend (SQLite, MongoDB, memory) holds harness # data. storage-drivers = [] -multimodal = ["dep:flate2", "dep:reqwest", "dep:tar", "dep:tinytools-std", "dep:zip"] +multimodal = ["dep:flate2", "dep:reqwest", "dep:tar", "dep:zip"] # Lossless PNG re-compression (`multimodal::optimize_png_lossless`) through # `oxipng`. Off by default: without it the helper returns `None` and callers # keep the original bytes, and `oxipng`/`libdeflater` (a C build) are not linked. @@ -86,6 +84,7 @@ tracing = ["tinytools-agent/tracing"] [dev-dependencies] png = "0.17" +toml = "0.8" tinystoragedrivers-testkit = { workspace = true } tokio = { workspace = true, default-features = true, features = ["macros", "rt-multi-thread", "time", "test-util"] } diff --git a/crates/tinyagents-harness/src/README.md b/crates/tinyagents-harness/src/README.md index d2c4ce51f..99c1972ec 100644 --- a/crates/tinyagents-harness/src/README.md +++ b/crates/tinyagents-harness/src/README.md @@ -18,37 +18,33 @@ the module map below. | Module | Concern | | --- | --- | -| `agent_loop` | The default model-tool-model agent loop — one model call driven to completion, plus tool dispatch. | +| `agent_loop` | The default model-tool-model agent loop — one model call driven to completion, tool dispatch, and caller-consumable event streaming. | | `artifacts` | Filesystem offload for oversized worker artifacts on long-horizon runs. | | `blocking` (private) | `run_blocking`, the helper that moves synchronous file/DB work off the tokio runtime. | | `cache` | Prompt, response, and layout caches for repeated/replayed requests. | | `cancel` | Cooperative, runtime-agnostic run cancellation (`CancellationToken`). See [`cancel/README.md`](cancel/README.md). | -| `capability` | `Capability` bundles (instructions, toolset, middleware, model defaults, exposure, `defer_loading`) composed as one named unit, plus `CapabilityToolSet` / `LoadCapabilityTool` for on-demand loading. | | `config` | Crate-owned session configuration, decoupled from any host's config schema. | | `context` | `RunContext` — the unit of recursion: run configuration and runtime context threaded through every nested layer. | | `cost` | Additive cost accounting (`CostTotals`) that rolls a child run's cost up into its parent. | | `error` | The crate-wide `TinyAgentsError` type and `Result` alias every fallible surface funnels through. | | `events` | The typed observability layer (`AgentEvent`, sinks, listeners) — the live in-process event surface. | | `finish_reason` (private) | Shared helper that recognises provider spellings of an output-cap stop (`length`, `max_tokens`, `MAX_TOKENS`). | -| `handoff` | Progressive-disclosure handoff cache for oversized tool results, keeping bloated payloads out of sub-agent history. | | `host` | Host capability traits — the seams a host implements to supply product-specific behavior to a generic runtime. | | `ids` | Identifier newtypes (`RunId`, `CallId`, …) and lifecycle enums used to correlate a recursive run tree. | | `limits` | Run-scoped limit enforcement (model/tool call caps, wall clock) that keeps recursion bounded. | | `media` (feature `media`) | `GenerateImageTool` / `GenerateVideoTool`, tools over the `tinyinference-image` / `tinyinference-video` generators. | | `middleware` | The middleware stack wrapping every level of the recursion identically. See [`middleware/README.md`](middleware/README.md). | -| `model_registry` | Runtime-owned executable model registry, name resolution, and fallback ordering. | +| `host` | The resolver port and fixed/prebuilt implementations hosts use to supply models for execution. | | `multimodal` (feature `multimodal`) | Attachment resolution for `[IMAGE:…]` / `[FILE:…]` markers into model-readable bytes. | | `no_progress` | Detects a stuck turn (identical failing/successful tool calls) and escalates through a nudge/halt ladder. See [`no_progress/README.md`](no_progress/README.md). | | `observability` | Durable observability — journals, status stores, sinks — making the live event stream persistent. See [`observability/README.md`](observability/README.md). | | `prompt` | Prompt assembly — templates and `PromptBuilder` turning runtime values into the final request. | | `providers` | Model adapters whose behavior depends on TinyAgents-specific prompt dialects (e.g. Claude Code/Agent SDK). | -| `retriever` | Provider-neutral retrieval contracts (`Retriever`) for injecting ranked context into a prompt. See [`retriever/README.md`](retriever/README.md). | | `retry` | Retry/backoff, model fallback, and rate-limiting policies applied uniformly to every model call. See [`retry/README.md`](retry/README.md). | | `run_queue` | A generic multi-lane FIFO queue (steer/followup/collect) for messages arriving during an active run. See [`run_queue/README.md`](run_queue/README.md). | | `runtime` | The harness runtime facade (`AgentHarness`) and invocation-local runtime wiring. See [`runtime/README.md`](runtime/README.md). | | `steering` | Policy-checked, observable orchestrator → sub-agent steering commands. See [`steering/README.md`](steering/README.md). | | `store` | Long-term key-value and append-only stream storage backends (`Store`, `AppendStore`, `NamespacedStore`). See [`store/README.md`](store/README.md). | -| `stream` | Higher-level streaming projections (`StreamChunk`, `StreamSink`) from the raw event stream. See [`stream/README.md`](stream/README.md). | | `structured` | Structured (typed, JSON-schema-validated) output extraction from a model call. | | `summarization` | Explicit message trimming, summarization, and compression policies (the harness's answer to context rot). | | `terminal` | `TerminalOutcome`, the typed reason a run ended (completed, halted, limit reached, timeout phase, ...). | @@ -62,20 +58,17 @@ the module map below. ## How the pieces fit together 1. **Configure**: a host builds a `RunContext` (`context`) carrying its - `StoreRegistry` (`store`), `ModelRegistry` (`model_registry`), - `ToolRegistry` (`tool`), limits (`limits`), cancellation token (`cancel`), + `StoreRegistry` (`store`), model resolver (`host`), `ToolRegistry` (`tool`), limits (`limits`), cancellation token (`cancel`), and event sink (`events`). 2. **Run**: `agent_loop` drives the model ⇄ tool cycle for one turn, - consulting `retry`/`model_registry` on provider failure, `no_progress` and - `summarization` to keep the loop productive and in-budget, `handoff` to - keep oversized tool results out of history, and `middleware` around every + consulting `retry` and the configured resolver on provider failure, `no_progress` and + `summarization` to keep the loop productive and in-budget, and `middleware` around every step. 3. **Recurse**: a tool can wrap a whole nested `AgentHarness` invocation - (`SubAgent` / `SubAgentTool` in `tinyagents-orchestration`), so the same + (`DriverSubAgentTool` in `tinyagents-harness-tools`), so the same loop runs again one level deeper, tracked by `ids`/`limits`/`cost` roll-up. -4. **Observe**: every step emits an `AgentEvent` (`events`); `stream` projects - those onto consumer-facing `StreamChunk`s, and `observability` durably - journals them. +4. **Observe**: every step emits an `AgentEvent` (`events`); `agent_loop` can + stream those events to callers, and `observability` durably journals them. 5. **Extend**: host-owned `AgentMiddleware` wraps the complete run to load and save memory, prepare and clean up workspaces, or attach other product policy. The harness keeps only the execution seam. diff --git a/crates/tinyagents-harness/src/agent_loop/README.md b/crates/tinyagents-harness/src/agent_loop/README.md index e9bc801fd..913078cbb 100644 --- a/crates/tinyagents-harness/src/agent_loop/README.md +++ b/crates/tinyagents-harness/src/agent_loop/README.md @@ -4,7 +4,7 @@ The default model-tool-model agent loop: the innermost turn of the recursive harness. This loop is where one model call is driven to completion. Because a whole -harness can be exposed as a tool (`tinyagents_orchestration::subagent::SubAgentTool`), the +harness can be exposed as a tool (`tinyagents_harness_tools::subagents::DriverSubAgentTool`), the tools this loop executes may themselves be other agents — "a model calling a model" is just this loop nested inside one of its own tool calls. Each invocation runs inside a `RunContext` that tracks recursion depth, fans diff --git a/crates/tinyagents-harness/src/agent_loop/deferred_tests.rs b/crates/tinyagents-harness/src/agent_loop/deferred_tests.rs index eb4dfd1eb..33085211b 100644 --- a/crates/tinyagents-harness/src/agent_loop/deferred_tests.rs +++ b/crates/tinyagents-harness/src/agent_loop/deferred_tests.rs @@ -127,13 +127,13 @@ fn tool_result_text(messages: &[Message], call_id: &str) -> Option<String> { 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( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "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()); @@ -192,13 +192,13 @@ async fn deferred_run( crate::middleware::AgentRun, ) { let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "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()); @@ -332,7 +332,7 @@ async fn resume_refuses_an_incomplete_resolution_and_names_the_missing_ids() { #[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( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "mock", Arc::new(MockModel::with_responses(vec![ response( @@ -345,7 +345,7 @@ async fn external_tool_call_is_deferred_and_its_host_result_is_injected_on_resum ), 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(), @@ -404,13 +404,13 @@ impl crate::tool::DeferredToolHandler for ApproveAllHandler { async fn inline_handler_resolves_deferrals_without_surfacing_them() { let recorder = EventRecorder::new(); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "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")); @@ -464,10 +464,10 @@ impl crate::tool::DeferredToolHandler for SilentHandler { #[tokio::test] async fn inline_handler_that_leaves_calls_unresolved_fails_the_run() { let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "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)); @@ -509,7 +509,7 @@ impl Tool for SelfDeferringTool { async fn tool_raising_approval_required_defers_with_its_metadata() { let recorder = EventRecorder::new(); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "mock", Arc::new(MockModel::with_responses(vec![response( vec![ToolCall::new( @@ -519,7 +519,7 @@ async fn tool_raising_approval_required_defers_with_its_metadata() { )], "", )])), - ); + ))); harness.register_tool(Arc::new(SelfDeferringTool)); let ctx = RunContext::new(RunConfig::new("exec-defer"), ()).with_events(recorder.sink()); @@ -561,13 +561,13 @@ 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( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "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"); @@ -613,13 +613,13 @@ async fn human_approval_middleware_deny_outcome_answers_the_model_without_runnin use crate::middleware::library::{ApprovalOutcome, HumanApprovalMiddleware}; let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "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")); @@ -682,13 +682,13 @@ fn finish_only_batch() -> ModelResponse { fn terminate_harness(first: ModelResponse) -> AgentHarness<()> { let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "mock", Arc::new(MockModel::with_responses(vec![ first, response(Vec::new(), "model saw the results"), ])), - ); + ))); harness.register_tool(RecordingTool::plain("lookup", "found")); harness.register_tool(Arc::new(FinishTool)); harness diff --git a/crates/tinyagents-harness/src/agent_loop/entry.rs b/crates/tinyagents-harness/src/agent_loop/entry.rs index 10fc52ff3..b5152fbb7 100644 --- a/crates/tinyagents-harness/src/agent_loop/entry.rs +++ b/crates/tinyagents-harness/src/agent_loop/entry.rs @@ -23,33 +23,9 @@ impl<State: Send + Sync, Ctx: Send + Sync> AgentBaseCall<State, Ctx> ) -> BoxAgentFuture<'a> { Box::pin(async move { ctx.streaming = request.streaming; - 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 - } - } + self.harness + .run_loop(state, ctx, run, status, request.input, request.streaming) + .await }) } } diff --git a/crates/tinyagents-harness/src/agent_loop/lifecycle.rs b/crates/tinyagents-harness/src/agent_loop/lifecycle.rs index 5a72db8ea..0d1df4db3 100644 --- a/crates/tinyagents-harness/src/agent_loop/lifecycle.rs +++ b/crates/tinyagents-harness/src/agent_loop/lifecycle.rs @@ -29,10 +29,6 @@ pub(crate) struct TurnTracker { turn: u32, /// The open turn and the transcript index it started at. open: Option<(u32, usize)>, - /// Whether the seed has been fixed. A context that never went through - /// [`TurnTracker::new`] (a graph run entered mid-flight) is seeded on first - /// use by [`TurnTracker::ensure_seeded`]. - seeded: bool, } pub(crate) fn role_of(message: &Message) -> &'static str { @@ -54,39 +50,6 @@ impl TurnTracker { seed_len, turn: 0, open: None, - seeded: true, - } - } - - /// Continues numbering from `completed_turns` (a resumed run in a fresh - /// runtime starts its tracker at zero) and, when `open_from` is given and no - /// turn is open, re-opens the in-flight turn that began at that transcript - /// index, without announcing it again. - pub(crate) fn adopt(&mut self, completed_turns: u32, open_from: Option<usize>) { - self.turn = self.turn.max(completed_turns); - if let Some(start) = open_from - && self.open.is_none() - { - tracing::debug!( - target: "tinyagents::agent_loop", - turn = self.turn, - start, - "[agent_loop] re-opened the in-flight turn after a resume" - ); - self.open = Some((self.turn, start)); - } - } - - /// Treats the first `len` messages as seed input if no seed was fixed yet. - /// A no-op once seeded, so repeated node entries never swallow appends. - pub(crate) fn ensure_seeded(&mut self, len: usize) { - if !self.seeded { - tracing::debug!( - target: "tinyagents::agent_loop", - seed_len = len, - "[agent_loop] lifecycle tracker seeded on first use" - ); - *self = Self::new(len); } } @@ -244,15 +207,3 @@ impl<Ctx> RunContext<Ctx> { self.turns.rebase(&self.events, new_len, reason); } } - -impl<Ctx> RunContext<Ctx> { - /// See [`TurnTracker::adopt`]. - pub(crate) fn adopt_turn_state(&mut self, completed_turns: u32, open_from: Option<usize>) { - self.turns.adopt(completed_turns, open_from); - } - - /// Fixes the lifecycle seed at `len` messages unless one is already set. - pub(crate) fn ensure_turn_tracker_seeded(&mut self, len: usize) { - self.turns.ensure_seeded(len); - } -} diff --git a/crates/tinyagents-harness/src/agent_loop/lifecycle_tests.rs b/crates/tinyagents-harness/src/agent_loop/lifecycle_tests.rs index 7149755ad..1037d12cf 100644 --- a/crates/tinyagents-harness/src/agent_loop/lifecycle_tests.rs +++ b/crates/tinyagents-harness/src/agent_loop/lifecycle_tests.rs @@ -80,7 +80,10 @@ fn tool_turn(ids: &[&str]) -> ModelResponse { fn harness(responses: Vec<ModelResponse>, capture: PayloadCapture) -> AgentHarness<()> { let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", Arc::new(ScriptedModel::new(responses))); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + Arc::new(ScriptedModel::new(responses)), + ))); harness.register_tool(Arc::new(EchoTool)); harness.with_policy(RunPolicy { capture, @@ -189,7 +192,10 @@ async fn message_payloads_follow_the_capture_policy() { #[tokio::test] async fn a_failed_turn_is_still_closed() { let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", Arc::new(FailingModel)); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + Arc::new(FailingModel), + ))); let recorder = EventRecorder::new(); let ctx = RunContext::new(RunConfig::new("fail"), ()).with_events(recorder.sink()); let partial = harness @@ -483,7 +489,7 @@ async fn a_mixed_structured_turn_closes_before_queued_messages_are_drained() { "", ); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "mock", Arc::new(ToolStructuredScript { profile: tinyinference_llm::model::ModelProfile { @@ -494,7 +500,7 @@ async fn a_mixed_structured_turn_closes_before_queued_messages_are_drained() { }, responses: std::sync::Mutex::new(vec![mixed, last].into()), }), - ); + ))); harness.register_tool(Arc::new(EchoTool)); harness.with_policy(RunPolicy { end_strategy: crate::runtime::EndStrategy::Exhaustive, @@ -560,12 +566,12 @@ async fn messages_appended_by_after_agent_middleware_are_announced() { } } let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "mock", Arc::new(crate::testkit::ScriptedModel::new(vec![ crate::testkit::text_response("done"), ])), - ); + ))); harness.push_middleware(Arc::new(Appender)); let recorder = EventRecorder::new(); let ctx = RunContext::new(RunConfig::new("aa-append"), ()).with_events(recorder.sink()); diff --git a/crates/tinyagents-harness/src/agent_loop/mod.rs b/crates/tinyagents-harness/src/agent_loop/mod.rs index 1a572edbc..41389b8a6 100644 --- a/crates/tinyagents-harness/src/agent_loop/mod.rs +++ b/crates/tinyagents-harness/src/agent_loop/mod.rs @@ -3,7 +3,7 @@ //! This loop is the innermost turn of the recursive runtime: it is where one //! model call is driven to completion, and because a //! whole harness can be exposed as a tool -//! (`SubAgentTool` in `tinyagents-orchestration`), the very tools this loop +//! (the driver-backed spawn tool in `tinyagents-harness-tools`), the tools this loop //! executes may themselves be other agents — so "a model calling a model" is //! just this loop nested inside one of its own tool calls. Each invocation runs //! inside a [`RunContext`] that tracks recursion depth, fans usage/cost up to a @@ -89,6 +89,7 @@ //! provider integration retries after a genuine, growing delay while unit tests //! stay sleep-free. +pub mod phases; mod types; pub use types::*; @@ -97,6 +98,7 @@ use std::future::Future; use std::sync::Arc; use std::time::Duration; +use crate::agent_loop::ModelBinding; use crate::cache::{ResponseCache, cache_key}; use crate::context::{LoopTarget, MiddlewareControl, RunConfig, RunContext}; use crate::error::{Result, TinyAgentsError}; @@ -106,7 +108,6 @@ use crate::middleware::{ AgentBaseCall, AgentRun, BoxAgentFuture, BoxModelFuture, BoxToolFuture, ModelBaseCall, ToolBaseCall, }; -use crate::model_registry::{ResolvedModelBinding, model_eligible}; use crate::runtime::{AgentHarness, EndStrategy, InvalidArgsPolicy, UnknownToolPolicy}; use crate::structured::{StructuredExtractor, StructuredStrategy}; use crate::terminal::{TerminalClass, TerminalOutcome, TimeoutPhase}; @@ -125,11 +126,10 @@ mod handoff_transform; mod host_budget; mod lifecycle; mod mixed_turn; -mod model_call; +pub(crate) mod model_call; mod model_switch; mod model_turn; mod nested; -pub mod phases; mod response_recovery; mod run_loop; pub(crate) mod stream; diff --git a/crates/tinyagents-harness/src/agent_loop/mod_tests.rs b/crates/tinyagents-harness/src/agent_loop/mod_tests.rs index e512d2988..61f801b0a 100644 --- a/crates/tinyagents-harness/src/agent_loop/mod_tests.rs +++ b/crates/tinyagents-harness/src/agent_loop/mod_tests.rs @@ -776,7 +776,10 @@ async fn agent_middleware_owns_context_and_observes_failed_partial_runs() { }); let mut successful: AgentHarness<()> = AgentHarness::new(); - successful.register_model("mock", Arc::new(MockModel::constant("done"))); + successful.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + Arc::new(MockModel::constant("done")), + ))); successful.push_agent_middleware(middleware.clone()); let run = successful .invoke_default(&(), vec![Message::user("go")]) @@ -785,12 +788,12 @@ async fn agent_middleware_owns_context_and_observes_failed_partial_runs() { assert_eq!(run.messages[0].text(), "remembered by the host"); let mut failing: AgentHarness<()> = AgentHarness::new(); - failing.register_model( + failing.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "mock", Arc::new(FailingModel { attempts: Mutex::new(0), }), - ); + ))); failing.push_agent_middleware(middleware); assert!( failing @@ -805,13 +808,13 @@ async fn agent_middleware_owns_context_and_observes_failed_partial_runs() { #[tokio::test] async fn wrap_middleware_fires_around_model_and_tool_calls() { let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "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"))); harness.push_model_middleware(Arc::new(StampModelWrap)); harness.push_tool_middleware(Arc::new(StampToolWrap)); @@ -838,7 +841,10 @@ async fn wrap_model_short_circuit_skips_provider() { let model = Arc::new(FailingModel { attempts: Mutex::new(0), }); - harness.register_model("mock", model.clone()); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + model.clone(), + ))); harness.push_model_middleware(Arc::new(ShortCircuitModelWrap)); let run = harness @@ -853,7 +859,10 @@ async fn wrap_model_short_circuit_skips_provider() { #[tokio::test] async fn single_model_call_no_tools() { let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", Arc::new(MockModel::constant("hello there"))); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + Arc::new(MockModel::constant("hello there")), + ))); let run = harness .invoke_default(&(), vec![Message::user("hi")]) @@ -875,7 +884,10 @@ async fn empty_response_fails_the_run_when_guard_enabled() { // when the guard is enabled — it fails with a typed `EmptyResponse` so the // caller can re-prompt instead of silently succeeding on empty content. let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", Arc::new(MockModel::constant(""))); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + Arc::new(MockModel::constant("")), + ))); harness.with_policy(RunPolicy { error_on_empty_response: true, ..RunPolicy::default() @@ -896,7 +908,10 @@ async fn empty_response_terminates_normally_when_guard_disabled() { // The guard is opt-in: with the default policy an empty completion still // terminates the run with a blank final answer (preserved behavior). let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", Arc::new(MockModel::constant(""))); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + Arc::new(MockModel::constant("")), + ))); let run = harness .invoke_default(&(), vec![Message::user("hi")]) @@ -913,7 +928,10 @@ async fn reasoning_only_stop_retries_once_when_enabled() { text_response("recovered", 4, 3), ])); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", Arc::clone(&model) as _); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + Arc::clone(&model) as _, + ))); harness.with_policy(RunPolicy { empty_response_retries: 1, ..RunPolicy::default() @@ -948,7 +966,10 @@ async fn reasoning_only_stop_keeps_default_blank_final() { reasoning_only_stop_response(), ])); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", Arc::clone(&model) as _); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + Arc::clone(&model) as _, + ))); let run = harness .invoke_default(&(), vec![Message::user("hi")]) @@ -965,7 +986,10 @@ async fn reasoning_only_stop_retry_exhaustion_respects_empty_guard() { reasoning_only_stop_response(), ])); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", Arc::clone(&model) as _); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + Arc::clone(&model) as _, + ))); harness.with_policy(RunPolicy { empty_response_retries: 1, error_on_empty_response: true, @@ -990,7 +1014,10 @@ async fn empty_response_retry_budget_resets_after_a_tool_turn() { ])); let mut harness: AgentHarness<()> = AgentHarness::new(); harness - .register_model("mock", Arc::clone(&model) as _) + .with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + Arc::clone(&model) as _, + ))) .register_tool(Arc::new(FakeTool::new("fake", "tool output"))); harness.with_policy(RunPolicy { empty_response_retries: 1, @@ -1014,7 +1041,10 @@ async fn empty_response_retry_does_not_replace_an_explicit_continuation() { text_response("done", 4, 3), ])); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", Arc::clone(&model) as _); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + Arc::clone(&model) as _, + ))); harness.with_policy(RunPolicy { empty_response_retries: 1, ..RunPolicy::default() @@ -1042,7 +1072,10 @@ async fn empty_response_retry_does_not_replay_a_cached_blank() { text_response("recovered", 4, 3), ])); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", Arc::clone(&model) as _); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + Arc::clone(&model) as _, + ))); harness.with_response_cache(Arc::new(InMemoryResponseCache::new())); harness.with_policy(RunPolicy { empty_response_retries: 1, @@ -1074,7 +1107,10 @@ async fn empty_response_retry_ignores_a_blank_cached_before_opt_in() { text_response("recovered", 4, 3), ])); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", Arc::clone(&model) as _); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + Arc::clone(&model) as _, + ))); harness.with_response_cache(Arc::new(InMemoryResponseCache::new())); let input = vec![Message::user("same request")]; @@ -1107,7 +1143,10 @@ async fn truncated_empty_retry_does_not_replay_a_cached_blank() { text_response("recovered", 4, 3), ])); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", Arc::clone(&model) as _); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + Arc::clone(&model) as _, + ))); harness.with_policy(RunPolicy { // The first retry keeps its cap; only the second may grow it. ..RunPolicy::default() @@ -1142,7 +1181,10 @@ async fn truncated_empty_response_retries_then_succeeds() { text_response("recovered", 4, 3), ])); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", Arc::clone(&model) as _); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + Arc::clone(&model) as _, + ))); use crate::testkit::EventRecorder; let recorder = EventRecorder::new(); @@ -1196,7 +1238,10 @@ async fn truncated_empty_boost_does_not_leak_into_later_turns() { ])); let mut harness: AgentHarness<()> = AgentHarness::new(); harness - .register_model("mock", Arc::clone(&model) as _) + .with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + Arc::clone(&model) as _, + ))) .register_tool(Arc::new(FakeTool::new("fake", "tool output"))); harness.with_policy(RunPolicy { // This test pins the cap ladder; keep reasoning on so the ladder runs. @@ -1237,7 +1282,10 @@ async fn truncated_empty_retry_stops_at_the_cap_ceiling_and_nudges() { text_response("recovered", 4, 3), ])); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", Arc::clone(&model) as _); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + Arc::clone(&model) as _, + ))); harness.with_policy(RunPolicy { // This test pins the cap ladder; keep reasoning on so the ladder runs. truncated_empty_retries: 4, @@ -1344,7 +1392,9 @@ async fn truncated_empty_retry_yields_to_the_clock() { calls: Mutex::new(0), }); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", model as _); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", model as _, + ))); harness.with_policy(RunPolicy { limits: RunLimits::default().with_max_wall_clock_ms(Some(1_000)), truncated_empty_nudges: 1, @@ -1405,7 +1455,10 @@ async fn truncated_empty_nudges_repeat_while_the_clock_allows_with_a_halving_cap text_response("recovered", 4, 3), ])); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", Arc::clone(&model) as _); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + Arc::clone(&model) as _, + ))); harness.with_policy(RunPolicy { // This test pins the cap ladder; keep reasoning on so the ladder runs. limits: RunLimits::default().with_max_wall_clock_ms(Some(60_000)), @@ -1456,7 +1509,10 @@ async fn truncated_empty_nudges_stop_at_their_limit_even_with_clock_left() { let responses: Vec<ModelResponse> = (0..12).map(|_| truncated_empty_response(2048)).collect(); let model = Arc::new(crate::testkit::ScriptedModel::new(responses)); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", Arc::clone(&model) as _); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + Arc::clone(&model) as _, + ))); harness.with_policy(RunPolicy { limits: RunLimits::default().with_max_wall_clock_ms(Some(60_000)), ..RunPolicy::default() @@ -1484,7 +1540,10 @@ async fn truncated_empty_nudges_do_not_repeat_without_a_clock() { let responses: Vec<ModelResponse> = (0..6).map(|_| truncated_empty_response(2048)).collect(); let model = Arc::new(crate::testkit::ScriptedModel::new(responses)); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", Arc::clone(&model) as _); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + Arc::clone(&model) as _, + ))); let ctx = RunContext::new( RunConfig::new("truncated-no-clock").with_max_turn_output_tokens(2048), (), @@ -1512,7 +1571,10 @@ async fn truncated_empty_retry_budget_is_restored_for_a_later_turn() { ])); let mut harness: AgentHarness<()> = AgentHarness::new(); harness - .register_model("mock", Arc::clone(&model) as _) + .with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + Arc::clone(&model) as _, + ))) .register_tool(Arc::new(FakeTool::new("fake", "tool output"))); let run = harness @@ -1534,7 +1596,10 @@ async fn truncated_empty_retry_budget_stays_unset_when_request_had_none() { text_response("recovered", 4, 3), ])); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", Arc::clone(&model) as _); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + Arc::clone(&model) as _, + ))); let run = harness .invoke_default(&(), vec![Message::user("hi")]) @@ -1561,7 +1626,10 @@ async fn truncated_empty_retries_exhausted_returns_blank_by_default() { truncated_empty_response(4096), ])); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", Arc::clone(&model) as _); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + Arc::clone(&model) as _, + ))); let ctx = RunContext::new( RunConfig::new("truncated-exhausted").with_max_turn_output_tokens(2048), @@ -1595,7 +1663,10 @@ async fn truncated_empty_nudge_after_spent_retries_reaches_the_tool_call() { ])); let mut harness: AgentHarness<()> = AgentHarness::new(); harness - .register_model("mock", Arc::clone(&model) as _) + .with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + Arc::clone(&model) as _, + ))) .register_tool(Arc::new(FakeTool::new("fake", "tool output"))); let ctx = RunContext::new( @@ -1640,7 +1711,10 @@ async fn truncated_empty_nudge_asks_for_an_answer_when_no_tool_is_callable() { text_response("short answer", 4, 3), ])); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", Arc::clone(&model) as _); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + Arc::clone(&model) as _, + ))); let run = harness .invoke_default(&(), vec![Message::user("hi")]) @@ -1670,7 +1744,10 @@ async fn truncated_empty_nudge_is_skipped_when_no_model_call_remains() { truncated_empty_response(4096), ])); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", Arc::clone(&model) as _); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + Arc::clone(&model) as _, + ))); harness.with_policy(RunPolicy { limits: RunLimits::default().with_max_model_calls(2), ..RunPolicy::default() @@ -1693,7 +1770,10 @@ async fn truncated_empty_retry_is_skipped_when_no_model_call_remains() { truncated_empty_response(2048), ])); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", Arc::clone(&model) as _); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + Arc::clone(&model) as _, + ))); harness.with_policy(RunPolicy { limits: RunLimits::default().with_max_model_calls(1), ..RunPolicy::default() @@ -1717,7 +1797,10 @@ async fn truncated_empty_nudge_disabled_by_zero_policy() { truncated_empty_response(4096), ])); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", Arc::clone(&model) as _); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + Arc::clone(&model) as _, + ))); harness.with_policy(RunPolicy { truncated_empty_nudges: 0, ..RunPolicy::default() @@ -1742,7 +1825,10 @@ async fn truncated_empty_retries_exhausted_errors_when_guard_enabled() { truncated_empty_response(4096), ])); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", Arc::clone(&model) as _); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + Arc::clone(&model) as _, + ))); harness.with_policy(RunPolicy { error_on_empty_response: true, ..RunPolicy::default() @@ -1767,7 +1853,10 @@ async fn truncated_empty_retry_disabled_by_zero_policy_even_with_clock() { text_response("unexpected recovery", 4, 3), ])); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", Arc::clone(&model) as _); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + Arc::clone(&model) as _, + ))); harness.with_policy(RunPolicy { truncated_empty_retries: 0, truncated_empty_nudges: 0, @@ -1791,7 +1880,10 @@ async fn length_finish_with_text_is_not_treated_as_truncated_empty() { partial.finish_reason = Some("length".to_string()); let model = Arc::new(crate::testkit::ScriptedModel::new(vec![partial])); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", Arc::clone(&model) as _); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + Arc::clone(&model) as _, + ))); let run = harness .invoke_default(&(), vec![Message::user("hi")]) @@ -1847,7 +1939,10 @@ async fn length_truncation_rejects_only_the_last_native_call() { ])); let mut harness: AgentHarness<()> = AgentHarness::new(); harness - .register_model("mock", Arc::clone(&model) as _) + .with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + Arc::clone(&model) as _, + ))) .register_tool(Arc::clone(&tool) as _); let recorder = crate::testkit::EventRecorder::new(); @@ -1897,7 +1992,10 @@ async fn length_truncated_tool_calls_run_when_the_guard_is_disabled() { ])); let mut harness: AgentHarness<()> = AgentHarness::new(); harness - .register_model("mock", Arc::clone(&model) as _) + .with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + Arc::clone(&model) as _, + ))) .register_tool(Arc::clone(&tool) as _); harness.with_policy(RunPolicy { reject_truncated_tool_calls: false, @@ -1925,7 +2023,10 @@ async fn max_tokens_finish_reason_counts_as_a_length_stop() { ])); let mut harness: AgentHarness<()> = AgentHarness::new(); harness - .register_model("mock", Arc::clone(&model) as _) + .with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + Arc::clone(&model) as _, + ))) .register_tool(Arc::clone(&tool) as _); harness @@ -1954,7 +2055,10 @@ async fn consecutive_length_truncated_tool_turns_are_bounded_and_raise_the_token ])); let mut harness: AgentHarness<()> = AgentHarness::new(); harness - .register_model("mock", Arc::clone(&model) as _) + .with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + Arc::clone(&model) as _, + ))) .register_tool(Arc::clone(&tool) as _); let ctx = RunContext::new( @@ -1994,7 +2098,10 @@ async fn truncated_tool_call_budget_resets_after_a_clean_tool_turn() { ])); let mut harness: AgentHarness<()> = AgentHarness::new(); harness - .register_model("mock", Arc::clone(&model) as _) + .with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + Arc::clone(&model) as _, + ))) .register_tool(Arc::clone(&tool) as _); let ctx = RunContext::new( @@ -2042,7 +2149,10 @@ async fn a_text_recovered_last_call_is_rejected_under_a_length_stop() { ])); let mut harness: AgentHarness<()> = AgentHarness::new(); harness - .register_model("mock", Arc::clone(&model) as _) + .with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + Arc::clone(&model) as _, + ))) .register_tool(Arc::clone(&tool) as _); harness.with_policy(RunPolicy { text_dialect_recovery: crate::runtime::TextDialectRecovery::On, @@ -2084,7 +2194,10 @@ async fn a_forced_text_dialect_call_cut_by_a_length_stop_is_rejected() { ])); let mut harness: AgentHarness<()> = AgentHarness::new(); harness - .register_model("mock", Arc::clone(&model) as _) + .with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + Arc::clone(&model) as _, + ))) .register_tool(Arc::clone(&tool) as _); harness.with_policy(RunPolicy { tool_dialect: crate::config::ToolDispatcher::Xml, @@ -2112,7 +2225,10 @@ async fn duplicate_provider_ids_reject_the_positional_last_call() { ])); let mut harness: AgentHarness<()> = AgentHarness::new(); harness - .register_model("mock", Arc::clone(&model) as _) + .with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + Arc::clone(&model) as _, + ))) .register_tool(Arc::clone(&tool) as _); let run = harness @@ -2140,7 +2256,10 @@ async fn an_empty_max_tokens_reply_takes_the_truncated_empty_retry_path() { text_response("recovered", 4, 3), ])); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", Arc::clone(&model) as _); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + Arc::clone(&model) as _, + ))); let ctx = RunContext::new( RunConfig::new("max-tokens-empty").with_max_turn_output_tokens(2048), @@ -2180,7 +2299,10 @@ async fn an_invalid_call_is_rejected_not_repaired_under_a_length_stop() { ])); let mut harness: AgentHarness<()> = AgentHarness::new(); harness - .register_model("mock", Arc::clone(&model) as _) + .with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + Arc::clone(&model) as _, + ))) .register_tool(Arc::clone(&tool) as _); harness @@ -2285,7 +2407,10 @@ fn harness_with_terminating_tools( ) -> (AgentHarness<()>, Arc<crate::testkit::ScriptedModel>) { let model = Arc::new(crate::testkit::ScriptedModel::new(responses)); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", Arc::clone(&model) as _); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + Arc::clone(&model) as _, + ))); for tool in tools { harness.register_tool(Arc::new(tool)); } @@ -2428,13 +2553,13 @@ async fn return_direct_in_a_mixed_batch_still_ends_the_run() { #[tokio::test] async fn model_requests_tool_then_finishes() { let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "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 run = harness @@ -2461,13 +2586,13 @@ async fn model_requests_tool_then_finishes() { #[tokio::test] async fn dynamic_toolset_change_folds_into_the_leading_system_message_by_default() { let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "mock", Arc::new(MockModel::with_responses(vec![ tool_call_response("call-1", "search", json!({"q": "x"})), text_response("done", 4, 2), ])), - ); + ))); let search: Arc<dyn Tool> = Arc::new(FakeTool::new("search", "search-output")); let browse: Arc<dyn Tool> = Arc::new(FakeTool::new("browse", "browse-output")); let toolset = Arc::new(DynamicToolSet { @@ -2531,7 +2656,7 @@ async fn dynamic_toolset_change_appends_exactly_one_patch_when_the_profile_allow mid_conversation_system_messages: true, ..ModelProfile::default() }; - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "mock", Arc::new(ProfiledModel { inner: MockModel::with_responses(vec![ @@ -2540,7 +2665,7 @@ async fn dynamic_toolset_change_appends_exactly_one_patch_when_the_profile_allow ]), profile, }), - ); + ))); let search: Arc<dyn Tool> = Arc::new(FakeTool::new("search", "search-output")); let browse: Arc<dyn Tool> = Arc::new(FakeTool::new("browse", "browse-output")); let toolset = Arc::new(DynamicToolSet { @@ -2598,147 +2723,6 @@ async fn dynamic_toolset_change_appends_exactly_one_patch_when_the_profile_allow /// 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<dyn Tool>, - } - #[async_trait] - impl crate::tool::toolset::ToolSet<(), ()> for SingleToolSet { - async fn tools(&self, _ctx: &RunContext<()>) -> Result<Vec<Arc<dyn Tool>>> { - Ok(vec![self.tool.clone()]) - } - async fn call( - &self, - name: &str, - args: serde_json::Value, - _ctx: &RunContext<()>, - ) -> Result<ToolResult> { - 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<_>>(), - 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(""); @@ -2748,13 +2732,13 @@ async fn after_tool_receives_distinct_identity_for_same_named_calls() { ]; let seen = Arc::new(Mutex::new(Vec::new())); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "mock", Arc::new(MockModel::with_responses(vec![ parallel_calls, text_response("done", 1, 1), ])), - ); + ))); harness.register_tool(Arc::new(FakeTool::new("lookup", "tool-output"))); harness.push_middleware(Arc::new(ToolInvocationRecorder { seen: Arc::clone(&seen), @@ -2778,10 +2762,10 @@ async fn after_tool_receives_distinct_identity_for_same_named_calls() { async fn max_model_calls_limit_triggers_limit_exceeded() { let mut harness: AgentHarness<()> = AgentHarness::new(); // A model that always asks for the tool -> the loop would never stop. - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "mock", Arc::new(MockModel::with_tool_call("spin", json!({}))), - ); + ))); harness.register_tool(Arc::new(FakeTool::new("spin", "again"))); harness.with_policy(RunPolicy { limits: RunLimits::default().with_max_model_calls(1), @@ -2808,10 +2792,10 @@ async fn policy_model_call_limit_above_run_config_default_is_honored() { // survives past 25 calls and, once it does trip, reports the limit that // actually applies. let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "mock", Arc::new(MockModel::with_tool_call("spin", json!({}))), - ); + ))); harness.register_tool(Arc::new(FakeTool::new("spin", "again"))); harness.with_policy(RunPolicy { limits: RunLimits::default() @@ -2837,10 +2821,10 @@ async fn policy_model_call_limit_above_run_config_default_is_honored() { #[tokio::test] async fn max_tool_calls_limit_triggers_limit_exceeded() { let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "mock", Arc::new(MockModel::with_tool_call("spin", json!({}))), - ); + ))); harness.register_tool(Arc::new(FakeTool::new("spin", "again"))); harness.with_policy(RunPolicy { limits: RunLimits::default() @@ -2863,7 +2847,10 @@ async fn max_tool_calls_limit_triggers_limit_exceeded() { async fn before_model_middleware_mutates_request() { let mut harness: AgentHarness<()> = AgentHarness::new(); // Echo returns the last user message; the middleware injects one. - harness.register_model("mock", Arc::new(MockModel::echo())); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + Arc::new(MockModel::echo()), + ))); harness.push_middleware(Arc::new(InjectMiddleware { text: "injected" })); let run = harness @@ -2902,7 +2889,10 @@ async fn prefix_mutating_middleware_refreshes_provider_cache_key() { let model = Arc::new(crate::testkit::ScriptedModel::replies(vec!["done"])); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", model.clone()); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + model.clone(), + ))); harness.with_policy(RunPolicy { cache: CachePolicy { protect_prompt_prefix: true, @@ -2963,7 +2953,10 @@ async fn prefix_mutating_wrap_middleware_refreshes_provider_cache_key() { let model = Arc::new(crate::testkit::ScriptedModel::replies(vec!["done"])); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", model.clone()); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + model.clone(), + ))); harness.with_policy(RunPolicy { cache: CachePolicy { protect_prompt_prefix: true, @@ -3039,7 +3032,10 @@ async fn middleware_owned_cache_segments_are_preserved_and_fingerprinted() { let model = Arc::new(crate::testkit::ScriptedModel::replies(vec!["done"])); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", model.clone()); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + model.clone(), + ))); harness.with_policy(RunPolicy { cache: CachePolicy { protect_prompt_prefix: true, @@ -3068,13 +3064,13 @@ async fn middleware_owned_cache_segments_are_preserved_and_fingerprinted() { #[tokio::test] async fn usage_accumulates_across_calls() { let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "mock", Arc::new(MockModel::with_responses(vec![ tool_call_response("call-1", "lookup", json!({})), text_response("done", 4, 2), ])), - ); + ))); harness.register_tool(Arc::new(FakeTool::new("lookup", "out"))); let run = harness @@ -3093,10 +3089,10 @@ async fn usage_accumulates_across_calls() { #[tokio::test] async fn tool_not_found_errors_under_the_fail_policy() { let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "mock", Arc::new(MockModel::with_tool_call("missing", json!({}))), - ); + ))); // No tool registered. harness.with_policy(RunPolicy { unknown_tool: UnknownToolPolicy::Fail, @@ -3116,13 +3112,13 @@ async fn tool_not_found_errors_under_the_fail_policy() { #[tokio::test] async fn unknown_tool_return_tool_error_recovers() { let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "mock", Arc::new(MockModel::with_responses(vec![ tool_call_response("call-1", "missing", json!({})), text_response("recovered", 1, 1), ])), - ); + ))); harness.with_policy(RunPolicy { unknown_tool: UnknownToolPolicy::ReturnToolError, ..RunPolicy::default() @@ -3149,13 +3145,13 @@ async fn unknown_tool_return_tool_error_recovers() { async fn unknown_tool_rewrite_retargets_to_real_tool() { let lookup = Arc::new(FakeTool::new("lookup", "out")); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "mock", Arc::new(MockModel::with_responses(vec![ tool_call_response("call-1", "missing", json!({})), text_response("done", 1, 1), ])), - ); + ))); harness.register_tool(lookup.clone()); harness.with_policy(RunPolicy { unknown_tool: UnknownToolPolicy::Rewrite { @@ -3179,13 +3175,13 @@ async fn unknown_tool_rewrite_retargets_to_real_tool() { #[tokio::test] async fn invalid_tool_arguments_fail_before_tool_execution_under_the_fail_policy() { let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "mock", Arc::new(MockModel::with_tool_call( "strict_lookup", json!({ "query": 42, "extra": true }), )), - ); + ))); let calls = Arc::new(Mutex::new(0)); harness.register_tool(Arc::new(StrictLookupTool { calls: Arc::clone(&calls), @@ -3211,7 +3207,7 @@ async fn invalid_tool_arguments_fail_before_tool_execution_under_the_fail_policy #[tokio::test] async fn invalid_tool_arguments_return_tool_error_recovers() { let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "mock", Arc::new(MockModel::with_responses(vec![ // `query` is required to be a string; 42 violates the schema, so @@ -3219,7 +3215,7 @@ async fn invalid_tool_arguments_return_tool_error_recovers() { tool_call_response("call-1", "strict_lookup", json!({ "query": 42 })), text_response("recovered", 1, 1), ])), - ); + ))); let calls = Arc::new(Mutex::new(0)); harness.register_tool(Arc::new(StrictLookupTool { calls: Arc::clone(&calls), @@ -3255,7 +3251,7 @@ async fn invalid_tool_arguments_return_tool_error_recovers() { #[tokio::test] async fn normalized_json_string_arguments_execute_registered_tool() { let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "mock", Arc::new(MockModel::with_responses(vec![ tool_call_response( @@ -3265,7 +3261,7 @@ async fn normalized_json_string_arguments_execute_registered_tool() { ), text_response("done", 1, 1), ])), - ); + ))); let calls = Arc::new(Mutex::new(0)); harness.register_tool(Arc::new(StrictLookupTool { calls: Arc::clone(&calls), @@ -3287,13 +3283,13 @@ async fn normalized_json_string_arguments_execute_registered_tool() { #[tokio::test] async fn normalized_non_object_executes_tool_without_required_fields() { let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "mock", Arc::new(MockModel::with_responses(vec![ tool_call_response("call-1", "permissive", json!(null)), text_response("done", 1, 1), ])), - ); + ))); let tool = Arc::new(FakeTool::new("permissive", "ok")); harness.register_tool(tool.clone()); harness.with_policy(RunPolicy { @@ -3320,13 +3316,13 @@ async fn normalization_preserves_a_decoded_but_schema_invalid_scalar() { // replaced with `{}`, letting the tool execute with fabricated empty // arguments instead of surfacing the model's real type mismatch. let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "mock", Arc::new(MockModel::with_responses(vec![ tool_call_response("call-1", "permissive", json!("true")), text_response("recovered", 1, 1), ])), - ); + ))); let tool = Arc::new(FakeTool::new("permissive", "ok")); harness.register_tool(tool.clone()); harness.with_policy(RunPolicy { @@ -3359,13 +3355,13 @@ async fn normalization_preserves_a_decoded_but_schema_invalid_scalar() { #[tokio::test] async fn normalization_preserves_valid_primitive_arguments() { let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "mock", Arc::new(MockModel::with_responses(vec![ tool_call_response("call-1", "string_echo", json!("hello")), text_response("done", 1, 1), ])), - ); + ))); let calls = Arc::new(Mutex::new(0)); harness.register_tool(Arc::new(StringEchoTool { calls: Arc::clone(&calls), @@ -3393,13 +3389,13 @@ async fn normalization_preserves_valid_primitive_arguments() { #[tokio::test] async fn normalization_decodes_required_only_object_schema() { let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "mock", Arc::new(MockModel::with_responses(vec![ tool_call_response("call-1", "required_only", json!("{\"query\":\"rust\"}")), text_response("done", 1, 1), ])), - ); + ))); let calls = Arc::new(Mutex::new(0)); harness.register_tool(Arc::new(RequiredOnlyTool { calls: Arc::clone(&calls), @@ -3421,13 +3417,13 @@ async fn normalization_decodes_required_only_object_schema() { #[tokio::test] async fn normalization_decodes_object_valued_enum_schema() { let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "mock", Arc::new(MockModel::with_responses(vec![ tool_call_response("call-1", "object_enum", json!("{\"query\":\"rust\"}")), text_response("done", 1, 1), ])), - ); + ))); let calls = Arc::new(Mutex::new(0)); harness.register_tool(Arc::new(ObjectEnumTool { calls: Arc::clone(&calls), @@ -3449,13 +3445,13 @@ async fn normalization_decodes_object_valued_enum_schema() { #[tokio::test] async fn normalization_preserves_decoded_array_union_arguments() { let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "mock", Arc::new(MockModel::with_responses(vec![ tool_call_response("call-1", "object_array", json!("[1,2]")), text_response("done", 1, 1), ])), - ); + ))); let calls = Arc::new(Mutex::new(0)); harness.register_tool(Arc::new(ObjectArrayTool { calls: Arc::clone(&calls), @@ -3483,7 +3479,7 @@ async fn normalization_preserves_decoded_array_union_arguments() { #[tokio::test] async fn normalization_preserves_decoded_invalid_object_for_validation() { let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "mock", Arc::new(MockModel::with_responses(vec![ tool_call_response( @@ -3493,7 +3489,7 @@ async fn normalization_preserves_decoded_invalid_object_for_validation() { ), text_response("recovered", 1, 1), ])), - ); + ))); let calls = Arc::new(Mutex::new(0)); harness.register_tool(Arc::new(OptionalStrictObjectTool { calls: Arc::clone(&calls), @@ -3533,13 +3529,13 @@ async fn normalization_preserves_decoded_invalid_object_for_validation() { #[tokio::test] async fn normalization_preserves_direct_invalid_object_for_validation() { let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "mock", Arc::new(MockModel::with_responses(vec![ tool_call_response("call-1", "optional_strict_object", json!({"extra": true})), text_response("recovered", 1, 1), ])), - ); + ))); let calls = Arc::new(Mutex::new(0)); harness.register_tool(Arc::new(OptionalStrictObjectTool { calls: Arc::clone(&calls), @@ -3567,13 +3563,13 @@ async fn normalization_preserves_direct_invalid_object_for_validation() { #[tokio::test] async fn normalization_preserves_required_field_validation_errors() { let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "mock", Arc::new(MockModel::with_responses(vec![ tool_call_response("call-1", "strict_lookup", json!(null)), text_response("recovered", 1, 1), ])), - ); + ))); let calls = Arc::new(Mutex::new(0)); harness.register_tool(Arc::new(StrictLookupTool { calls: Arc::clone(&calls), @@ -3602,13 +3598,13 @@ async fn malformed_tool_arguments_recover_as_error_tool_result() { // well-formed args, not unparseable ones), proving the leniency is // unconditional. let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "mock", Arc::new(MockModel::with_responses(vec![ invalid_tool_call_response("call-x", "strict_lookup", "{\"query\":"), text_response("recovered", 1, 1), ])), - ); + ))); let calls = Arc::new(Mutex::new(0)); harness.register_tool(Arc::new(StrictLookupTool { calls: Arc::clone(&calls), @@ -3654,7 +3650,7 @@ async fn provider_invalid_arguments_recoverable_by_relaxed_json_are_repaired_and let tool = Arc::new(crate::testkit::FakeTool::returning("lookup", "found it")); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "mock", Arc::new(MockModel::with_responses(vec![ // Unquoted object key: `tinytools_agent::repair::json::recover_object` @@ -3662,7 +3658,7 @@ async fn provider_invalid_arguments_recoverable_by_relaxed_json_are_repaired_and invalid_tool_call_response("call-x", "lookup", "{query:\"weather\"}"), text_response("found it", 1, 1), ])), - ); + ))); harness.register_tool(tool.clone()); let recorder = EventRecorder::new(); @@ -3700,10 +3696,10 @@ async fn provider_invalid_arguments_recoverable_by_relaxed_json_are_repaired_and #[tokio::test] async fn structured_output_is_extracted() { let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "mock", Arc::new(MockModel::constant(r#"{"value":"hi","score":42}"#)), - ); + ))); harness.with_policy(RunPolicy { default_response_format: Some(ResponseFormat::json_schema( "answer", @@ -3727,10 +3723,10 @@ async fn auto_format_uses_provider_schema_for_native_model() { // MockModel advertises a permissive profile (native structured output), so // `Auto` resolves to provider-native schema mode and parses the JSON text. let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "mock", Arc::new(MockModel::constant(r#"{"value":"native","score":1}"#)), - ); + ))); harness.with_policy(RunPolicy { default_response_format: Some(ResponseFormat::auto("answer", json!({"type": "object"}))), ..RunPolicy::default() @@ -3751,7 +3747,10 @@ async fn auto_format_uses_tool_call_for_non_native_model() { // fallback; the structured value is read from the tool-call arguments and // the artificial tool call is treated as the final response. let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("tool", Arc::new(ToolStructuredModel::new())); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "tool", + Arc::new(ToolStructuredModel::new()), + ))); harness.with_policy(RunPolicy { default_response_format: Some(ResponseFormat::auto("answer", json!({"type": "object"}))), ..RunPolicy::default() @@ -3803,13 +3802,13 @@ impl ChatModel<()> for CutThenWholeStructuredModel { #[tokio::test] async fn a_length_cut_structured_output_call_is_retried_not_extracted() { let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "tool", Arc::new(CutThenWholeStructuredModel { profile: ToolStructuredModel::new().profile, calls: std::sync::atomic::AtomicUsize::new(0), }), - ); + ))); harness.with_policy(RunPolicy { default_response_format: Some(ResponseFormat::auto("answer", json!({"type": "object"}))), ..RunPolicy::default() @@ -3838,7 +3837,10 @@ async fn pformat_dialect_recovers_the_structured_output_fallback_tool() { let model = Arc::new(PFormatStructuredModel::new()); let mut harness: AgentHarness<()> = AgentHarness::new(); harness - .register_model("mock", model.clone()) + .with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + model.clone(), + ))) .with_policy(RunPolicy { tool_dialect: crate::config::ToolDispatcher::Pformat, default_response_format: Some(ResponseFormat::auto( @@ -3900,8 +3902,10 @@ async fn native_tool_dispatcher_requires_tool_calling_capability() { }); let mut harness: AgentHarness<()> = AgentHarness::new(); harness - .register_model("mock", incapable.clone()) - .set_default_model("mock") + .with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + incapable.clone(), + ))) .register_tool(Arc::new(FakeTool::new("lookup", "tool-output"))) .with_policy(RunPolicy { tool_dialect: crate::config::ToolDispatcher::Native, @@ -3965,8 +3969,10 @@ async fn native_tool_dispatcher_gates_on_the_post_middleware_tool_set() { }); let mut harness: AgentHarness<()> = AgentHarness::new(); harness - .register_model("mock", incapable.clone()) - .set_default_model("mock") + .with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + incapable.clone(), + ))) .push_middleware(Arc::new(InjectToolMiddleware)) .with_policy(RunPolicy { tool_dialect: crate::config::ToolDispatcher::Native, @@ -4008,8 +4014,10 @@ async fn native_tool_dispatcher_gates_on_auto_structured_output_with_no_ordinary }); let mut harness: AgentHarness<()> = AgentHarness::new(); harness - .register_model("mock", incapable.clone()) - .set_default_model("mock") + .with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + incapable.clone(), + ))) .with_policy(RunPolicy { tool_dialect: crate::config::ToolDispatcher::Native, default_response_format: Some(ResponseFormat::auto( @@ -4068,7 +4076,10 @@ async fn retry_backoff_sleeps_the_documented_schedule() { timestamps: Mutex::new(Vec::new()), }); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("flaky", model.clone()); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "flaky", + model.clone(), + ))); harness.with_policy(RunPolicy { retry: policy, ..RunPolicy::default() @@ -4103,8 +4114,14 @@ async fn retry_then_fallback_succeeds() { let failing = Arc::new(FailingModel { attempts: Mutex::new(0), }); - harness.register_model("primary", failing.clone()); - harness.register_model("backup", Arc::new(MockModel::constant("recovered"))); + let backup = Arc::new(MockModel::constant("recovered")); + harness.with_model_resolver(Arc::new( + crate::host::PrebuiltModelResolver::new( + failing.clone(), + [("backup".to_string(), backup as Arc<dyn ChatModel<()>>)], + ) + .with_primary_name("primary"), + )); harness.with_policy(RunPolicy { // 2 attempts on primary, then fall back to backup. retry: RetryPolicy::default().with_max_attempts(2), @@ -4131,8 +4148,16 @@ async fn provider_budget_refusal_does_not_call_model_fallback() { "unexpected fallback", ])); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("primary", primary.clone()); - harness.register_model("fallback", fallback.clone()); + harness.with_model_resolver(Arc::new( + crate::host::PrebuiltModelResolver::new( + primary.clone(), + [( + "fallback".to_string(), + fallback.clone() as Arc<dyn ChatModel<()>>, + )], + ) + .with_primary_name("primary"), + )); harness.with_policy(RunPolicy { fallback: Some(FallbackPolicy::new(["primary", "fallback"])), ..RunPolicy::default() @@ -4158,7 +4183,10 @@ async fn run_limits_max_retries_per_call_caps_a_looser_retry_policy() { let failing = Arc::new(FailingModel { attempts: Mutex::new(0), }); - harness.register_model("primary", failing.clone()); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "primary", + failing.clone(), + ))); harness.with_policy(RunPolicy { retry: RetryPolicy::default().with_max_attempts(5), limits: RunLimits::default().with_max_retries_per_call(1), @@ -4182,8 +4210,16 @@ async fn a_provider_budget_refusal_stops_retries_and_model_fallback() { "unexpected fallback", ])); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("primary", primary.clone()); - harness.register_model("fallback", fallback.clone()); + harness.with_model_resolver(Arc::new( + crate::host::PrebuiltModelResolver::new( + primary.clone(), + [( + "fallback".to_string(), + fallback.clone() as Arc<dyn ChatModel<()>>, + )], + ) + .with_primary_name("primary"), + )); harness.with_policy(RunPolicy { retry: RetryPolicy::default().with_max_attempts(4), fallback: Some(FallbackPolicy::new(["primary", "fallback"])), @@ -4213,7 +4249,10 @@ async fn retry_middleware_and_run_policy_retry_do_not_multiply_attempts() { let failing = Arc::new(FailingModel { attempts: Mutex::new(0), }); - harness.register_model("primary", failing.clone()); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "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( @@ -4245,7 +4284,10 @@ async fn provider_error_401_is_not_retried() { status: 401, attempts: Mutex::new(0), }); - harness.register_model("primary", model.clone()); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "primary", + model.clone(), + ))); harness.with_policy(RunPolicy { retry: RetryPolicy::default().with_max_attempts(5), ..RunPolicy::default() @@ -4269,7 +4311,10 @@ async fn provider_error_429_is_retried_up_to_max_attempts() { status: 429, attempts: Mutex::new(0), }); - harness.register_model("primary", model.clone()); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "primary", + model.clone(), + ))); harness.with_policy(RunPolicy { retry: RetryPolicy::default().with_max_attempts(3), ..RunPolicy::default() @@ -4286,12 +4331,12 @@ async fn provider_error_429_is_retried_up_to_max_attempts() { #[tokio::test] async fn non_retryable_or_exhausted_without_fallback_errors() { let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "primary", Arc::new(FailingModel { attempts: Mutex::new(0), }), - ); + ))); harness.with_policy(RunPolicy { retry: RetryPolicy::default().with_max_attempts(1), ..RunPolicy::default() @@ -4318,8 +4363,16 @@ async fn fallback_chain_with_repeated_model_name_terminates() { let backup = Arc::new(FailingModel { attempts: Mutex::new(0), }); - harness.register_model("primary", primary.clone()); - harness.register_model("backup", backup.clone()); + harness.with_model_resolver(Arc::new( + crate::host::PrebuiltModelResolver::new( + primary.clone(), + [( + "backup".to_string(), + backup.clone() as Arc<dyn ChatModel<()>>, + )], + ) + .with_primary_name("primary"), + )); harness.with_policy(RunPolicy { retry: RetryPolicy::default().with_max_attempts(1), fallback: Some(FallbackPolicy::new(["primary", "backup", "primary"])), @@ -4452,9 +4505,22 @@ async fn runtime_fallback_skips_capability_ineligible_candidate() { }); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("primary", primary.clone()); - harness.register_model("backup", backup.clone()); - harness.register_model("tertiary", tertiary.clone()); + harness.with_model_resolver(Arc::new( + crate::host::PrebuiltModelResolver::new( + primary.clone(), + [ + ( + "backup".to_string(), + backup.clone() as Arc<dyn ChatModel<()>>, + ), + ( + "tertiary".to_string(), + tertiary.clone() as Arc<dyn ChatModel<()>>, + ), + ], + ) + .with_primary_name("primary"), + )); harness.push_middleware(Arc::new(RequireCapsMiddleware { model: "primary", caps: CapabilitySet { @@ -4519,7 +4585,10 @@ async fn native_tool_calling_model_does_not_execute_quoted_text_dialect_markup() }); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("native", model.clone()); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "native", + model.clone(), + ))); harness.register_tool(tool.clone()); let run = harness @@ -4545,7 +4614,10 @@ async fn invoke_with_status_reports_completed() { use crate::ids::{ExecutionStatus, HarnessPhase}; let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", Arc::new(MockModel::constant("ok"))); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + Arc::new(MockModel::constant("ok")), + ))); let result = harness .invoke_with_status(&(), (), RunConfig::new("run-x"), vec![Message::user("hi")]) @@ -4688,10 +4760,10 @@ async fn invoke_streaming_fires_on_model_delta_per_delta_and_accumulates() { let reasonings = Arc::new(Mutex::new(Vec::new())); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "stream", Arc::new(StreamingMock::from_text_chunks(["Hel", "lo, ", "world"])), - ); + ))); harness.push_middleware(Arc::new(DeltaRecorder { count: count.clone(), texts: texts.clone(), @@ -4734,10 +4806,10 @@ async fn streamed_repetitive_narration_stops_the_live_model_call() { }) .collect::<String>(); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "stream", Arc::new(StreamingMock::from_text_chunks([repeated])), - ); + ))); let error = harness .invoke_streaming( @@ -4766,7 +4838,10 @@ async fn streaming_delta_transform_controls_final_run_and_cached_response() { ModelStreamItem::Completed(ModelResponse::assistant("raw-secret")), ])); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("stream", model.clone()); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "stream", + model.clone(), + ))); harness.with_response_cache(Arc::new(InMemoryResponseCache::new())); harness.push_middleware(Arc::new(SecretRedactingDelta)); @@ -4830,7 +4905,7 @@ async fn streaming_turn_keeps_a_signed_thinking_signature_ahead_of_a_tool_call() .push(ToolCall::new("call-1", "lookup", json!({}))); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "stream", Arc::new(StreamingMock::new(vec![ ModelStreamItem::Started, @@ -4843,7 +4918,7 @@ async fn streaming_turn_keeps_a_signed_thinking_signature_ahead_of_a_tool_call() }), ModelStreamItem::Completed(terminal), ])), - ); + ))); harness.register_tool(tool.clone()); // Cap the run at one model call: the mock always replays the same @@ -4897,7 +4972,7 @@ async fn streaming_middleware_can_suppress_a_standalone_tool_delta() { .tool_calls .push(ToolCall::new("blocked-call", "blocked", json!({}))); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "stream", Arc::new(StreamingMock::new(vec![ ModelStreamItem::Started, @@ -4909,7 +4984,7 @@ async fn streaming_middleware_can_suppress_a_standalone_tool_delta() { }), ModelStreamItem::Completed(terminal), ])), - ); + ))); harness.register_tool(tool.clone()); harness.push_middleware(Arc::new(SuppressToolDelta)); @@ -4943,7 +5018,7 @@ async fn streaming_tool_delta_transform_controls_terminal_dispatch() { .tool_calls .push(ToolCall::new("raw-call", "blocked", json!({"raw": true}))); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "stream", Arc::new(StreamingMock::new(vec![ ModelStreamItem::Started, @@ -4955,7 +5030,7 @@ async fn streaming_tool_delta_transform_controls_terminal_dispatch() { }), ModelStreamItem::Completed(terminal), ])), - ); + ))); harness.register_tool(blocked.clone()); harness.register_tool(safe.clone()); harness.push_middleware(Arc::new(RewriteToolDelta)); @@ -4984,7 +5059,7 @@ async fn invoke_streaming_forwards_reasoning_deltas_to_middleware_and_events() { let reasonings = Arc::new(Mutex::new(Vec::new())); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "stream", Arc::new(StreamingMock::new(vec![ ModelStreamItem::Started, @@ -4992,7 +5067,7 @@ async fn invoke_streaming_forwards_reasoning_deltas_to_middleware_and_events() { ModelStreamItem::MessageDelta(MessageDelta::text("answer")), ModelStreamItem::Completed(ModelResponse::assistant("answer")), ])), - ); + ))); harness.push_middleware(Arc::new(DeltaRecorder { count: count.clone(), texts: texts.clone(), @@ -5034,10 +5109,10 @@ async fn invoke_streaming_emits_model_delta_events() { use crate::testkit::{EventRecorder, StreamingMock}; let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "stream", Arc::new(StreamingMock::from_text_chunks(["a", "b"])), - ); + ))); let recorder = EventRecorder::new(); let ctx = RunContext::new(RunConfig::new("stream-run"), ()).with_events(recorder.sink()); @@ -5119,13 +5194,13 @@ impl Tool for CancelOnCallTool { async fn token_cancelled_before_run_yields_cancelled() { let invocations = Arc::new(Mutex::new(0usize)); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "mock", Arc::new(CountingToolModel { name: "cancel_me", invocations: invocations.clone(), }), - ); + ))); // Pre-cancel the token before the run starts. let token = CancellationToken::new(); @@ -5149,13 +5224,13 @@ async fn cancelled_mid_run_stops_before_next_model_call() { let token = CancellationToken::new(); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "mock", Arc::new(CountingToolModel { name: "cancel_me", invocations: invocations.clone(), }), - ); + ))); harness.register_tool(Arc::new(CancelOnCallTool { token: token.clone(), })); @@ -5201,12 +5276,12 @@ async fn cancelled_during_unary_model_call_drops_the_in_flight_call() { let started = Arc::new(tokio::sync::Notify::new()); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "mock", Arc::new(BlockForeverModel { started: started.clone(), }), - ); + ))); let ctx = RunContext::new(RunConfig::new("cancel-unary"), ()).with_cancellation(token.clone()); @@ -5241,10 +5316,10 @@ async fn slow_model_call_is_timed_out_by_remaining_budget() { // (20ms), so the per-call timeout must interrupt it mid-flight and surface a // `Timeout` error rather than waiting for the model to return. let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "slow", Arc::new(SlowModel::new(Duration::from_millis(200), "too late")), - ); + ))); let config = RunConfig::new("timeout-run").with_timeout_ms(20); let err = harness @@ -5260,7 +5335,10 @@ async fn fast_model_call_succeeds_under_same_budget() { // Control: under the same small timeout a fast model completes well within // the budget, proving the timeout only fires on genuinely slow calls. let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("fast", Arc::new(MockModel::constant("done"))); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "fast", + Arc::new(MockModel::constant("done")), + ))); let config = RunConfig::new("fast-run").with_timeout_ms(20); let run = harness @@ -5280,10 +5358,10 @@ async fn slow_streaming_model_call_is_timed_out() { // The streaming path must enforce the same per-call budget: the default // `stream` impl delegates to `invoke`, whose sleep exceeds the 20ms budget. let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "slow", Arc::new(SlowModel::new(Duration::from_millis(200), "too late")), - ); + ))); let config = RunConfig::new("timeout-stream-run").with_timeout_ms(20); let err = harness @@ -5304,10 +5382,10 @@ async fn slow_tool_call_is_timed_out_by_remaining_budget() { // budget (20ms), so the same per-call timeout used for model calls must // interrupt it too. let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "mock", Arc::new(MockModel::with_tool_call("slow", json!({}))), - ); + ))); harness.register_tool(Arc::new(SlowTool { delay: Duration::from_millis(200), })); @@ -5335,10 +5413,10 @@ async fn per_model_call_ceiling_times_out_a_slow_call_with_run_time_left() { // 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( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "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() @@ -5365,6 +5443,7 @@ async fn per_model_call_ceiling_times_out_a_slow_call_with_run_time_left() { async fn per_model_call_ceiling_consults_the_fallback_chain_instead_of_aborting() { use std::time::Duration; + use crate::host::PrebuiltModelResolver; use crate::testkit::{ScriptedModel, SlowModel}; // Same setup as `per_model_call_ceiling_times_out_a_slow_call_with_run_time_left`, @@ -5378,12 +5457,15 @@ async fn per_model_call_ceiling_consults_the_fallback_chain_instead_of_aborting( 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", slow.clone()); - harness.register_model("fallback", fallback.clone()); - harness.with_policy(RunPolicy { - fallback: Some(FallbackPolicy::new(["fallback"])), - ..RunPolicy::default() - }); + let resolver = PrebuiltModelResolver::new( + slow.clone(), + [( + "fallback".to_string(), + fallback.clone() as Arc<dyn ChatModel<()>>, + )], + ) + .with_primary_name("slow"); + harness.with_model_resolver(Arc::new(resolver)); harness.with_policy(RunPolicy { limits: RunLimits::default().with_max_model_call_ms(Some(20)), @@ -5408,6 +5490,11 @@ async fn per_model_call_ceiling_consults_the_fallback_chain_instead_of_aborting( 1, "the fallback chain must actually have been consulted and called" ); + assert_eq!( + *slow.calls.lock().unwrap(), + 1, + "the primary must time out first" + ); } #[tokio::test] @@ -5419,10 +5506,10 @@ async fn per_model_call_ceiling_bounds_calls_without_any_run_deadline() { // 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( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "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() @@ -5454,10 +5541,10 @@ async fn run_remainder_bounds_a_model_call_when_tighter_than_the_ceiling() { // deadline (20ms): the tighter source wins, and the error blames the // remaining wall-clock budget. let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "slow", Arc::new(SlowModel::new(Duration::from_millis(200), "too late")), - ); + ))); harness.with_policy(RunPolicy { limits: RunLimits::default().with_max_model_call_ms(Some(10_000)), ..RunPolicy::default() @@ -5486,13 +5573,13 @@ async fn per_model_call_ceiling_does_not_bound_tool_calls() { // entire child run), so the run completes despite the tool outliving the // model-call ceiling many times over. let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "mock", Arc::new(MockModel::with_responses(vec![ tool_call_response("call-1", "slow", json!({})), text_response("done", 0, 0), ])), - ); + ))); harness.register_tool(Arc::new(SlowTool { delay: Duration::from_millis(100), })); @@ -5517,13 +5604,13 @@ async fn inherited_tool_timeout_is_enforced_without_a_run_deadline() { let settings = ToolTimeoutSettings::new(20, 1, 1_000, 0); let mut harness: AgentHarness<()> = AgentHarness::new(); harness.with_tool_timeout_settings(settings.clone()); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "mock", Arc::new(MockModel::with_responses(vec![ tool_call_response("call-1", "policy_slow", json!({})), text_response("recovered", 2, 1), ])), - ); + ))); harness.register_tool(Arc::new(PolicySlowTool { name: "policy_slow", delay: Duration::from_millis(200), @@ -5554,13 +5641,13 @@ async fn tool_timeout_unwinds_wrap_middleware() { let mut harness: AgentHarness<()> = AgentHarness::new(); harness.with_tool_timeout_settings(ToolTimeoutSettings::new(10, 1, 1_000, 0)); harness.push_tool_middleware(Arc::new(StampToolWrap)); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "mock", Arc::new(MockModel::with_responses(vec![ tool_call_response("call-1", "policy_slow", json!({})), text_response("recovered", 2, 1), ])), - ); + ))); harness.register_tool(Arc::new(PolicySlowTool { name: "policy_slow", delay: Duration::from_millis(200), @@ -5584,13 +5671,13 @@ async fn tool_timeout_resolves_after_wrap_middleware_mutation() { let mut harness: AgentHarness<()> = AgentHarness::new(); harness.with_tool_timeout_settings(ToolTimeoutSettings::new(10, 1, 1_000, 0)); harness.push_tool_middleware(Arc::new(UnboundToolWrap)); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "mock", Arc::new(MockModel::with_responses(vec![ tool_call_response("call-1", "policy_slow", json!({})), text_response("done", 2, 1), ])), - ); + ))); harness.register_tool(Arc::new(PolicySlowTool { name: "policy_slow", delay: Duration::from_millis(30), @@ -5612,13 +5699,13 @@ async fn parallel_tool_timeouts_return_ordered_recoverable_errors() { let mut harness: AgentHarness<()> = AgentHarness::new(); harness.with_tool_timeout_settings(ToolTimeoutSettings::new(10, 1, 1_000, 0)); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "mock", Arc::new(MockModel::with_responses(vec![ multi_tool_call_response(vec![("call-a", "slow_a"), ("call-b", "slow_b")]), text_response("recovered", 2, 1), ])), - ); + ))); for name in ["slow_a", "slow_b"] { harness.register_tool(Arc::new(PolicySlowTool { name, @@ -5654,7 +5741,10 @@ async fn response_cache_serves_repeated_request_without_calling_model() { let cache = Arc::new(InMemoryResponseCache::new()); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", model.clone()); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + model.clone(), + ))); harness.with_response_cache(cache.clone()); // First run: cache miss, model invoked once. @@ -5715,7 +5805,10 @@ async fn multi_turn_request_with_prior_assistant_turn_is_not_cached() { ])); let cache = Arc::new(InMemoryResponseCache::new()); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", model.clone()); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + model.clone(), + ))); harness.with_response_cache(cache.clone()); let convo = vec![ @@ -5745,7 +5838,10 @@ async fn no_cache_attached_invokes_model_each_run() { // Control: without a cache the model is invoked on every run. let model = Arc::new(MockModel::echo()); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", model.clone()); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + model.clone(), + ))); harness .invoke_default(&(), vec![Message::user("hello")]) @@ -5794,7 +5890,10 @@ async fn request_cache_policy_overrides_run_policy_to_disable_caching() { let model = Arc::new(MockModel::echo()); let cache = Arc::new(InMemoryResponseCache::new()); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", model.clone()); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + model.clone(), + ))); harness.with_response_cache(cache.clone()); harness.push_middleware(Arc::new(DisableCaching)); @@ -5847,7 +5946,10 @@ async fn hook_failure_delivers_on_error_exactly_once() { // compensating did it twice for one failure. let calls = Arc::new(Mutex::new(0usize)); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", Arc::new(MockModel::constant("hi"))); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + Arc::new(MockModel::constant("hi")), + ))); harness.push_middleware(Arc::new(FailingHookMiddleware { on_error_calls: calls.clone(), })); @@ -5887,10 +5989,10 @@ async fn middleware_control_stops_loop_with_final_response() { let mut harness: AgentHarness<()> = AgentHarness::new(); // The model asks for a tool; without control the loop would execute it and // continue, but the control outcome stops the run first. - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "mock", Arc::new(MockModel::with_tool_call("lookup", json!({}))), - ); + ))); harness.register_tool(Arc::new(FakeTool::new("lookup", "out"))); harness.push_middleware(Arc::new(EarlyStopMiddleware)); @@ -5949,7 +6051,10 @@ impl Middleware<(), ()> for InterruptMiddleware { #[tokio::test] async fn middleware_control_can_interrupt_run() { let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", Arc::new(MockModel::constant("hi"))); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + Arc::new(MockModel::constant("hi")), + ))); harness.push_middleware(Arc::new(InterruptMiddleware)); let err = harness @@ -6005,7 +6110,10 @@ async fn tool_schemas_are_stable_across_turns() { }; let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", Arc::new(model)); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + Arc::new(model), + ))); harness.register_tool(Arc::new(FakeTool::new("spin", "again"))); let run = harness @@ -6059,7 +6167,10 @@ async fn skipped_model_override_emits_diagnostic_event() { use crate::testkit::EventRecorder; let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", Arc::new(MockModel::constant("ok"))); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + Arc::new(MockModel::constant("ok")), + ))); harness.push_middleware(Arc::new(OverrideModelMiddleware("missing-model"))); let recorder = EventRecorder::new(); @@ -6084,11 +6195,19 @@ async fn skipped_model_override_emits_diagnostic_event() { /// An override that resolution honors must not emit the diagnostic. #[tokio::test] async fn honored_model_override_emits_no_diagnostic_event() { + use crate::host::PrebuiltModelResolver; use crate::testkit::EventRecorder; let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", Arc::new(MockModel::constant("default"))); - harness.register_model("special", Arc::new(MockModel::constant("special answer"))); + let resolver = PrebuiltModelResolver::new( + Arc::new(MockModel::constant("default")), + [( + "special".to_string(), + Arc::new(MockModel::constant("special answer")) as Arc<dyn ChatModel<()>>, + )], + ) + .with_primary_name("mock"); + harness.with_model_resolver(Arc::new(resolver)); harness.push_middleware(Arc::new(OverrideModelMiddleware("special"))); let recorder = EventRecorder::new(); @@ -6181,7 +6300,7 @@ async fn concurrent_tool_failure_fails_every_started_sibling_before_returning() use crate::testkit::EventRecorder; let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "mock", Arc::new(MockModel::with_responses(vec![multi_tool_call_response( // "boom" fails fast and comes first in call order, so the fold @@ -6190,7 +6309,7 @@ async fn concurrent_tool_failure_fails_every_started_sibling_before_returning() // 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", @@ -6300,13 +6419,13 @@ fn probe_pair( #[tokio::test] async fn independent_tool_calls_in_one_turn_run_concurrently() { let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "mock", Arc::new(MockModel::with_responses(vec![ multi_tool_call_response(vec![("call-a", "alpha"), ("call-b", "beta")]), text_response("done", 4, 2), ])), - ); + ))); let max_seen = probe_pair(&mut harness, (80, 80)); let run = harness @@ -6330,7 +6449,7 @@ async fn max_tool_concurrency_bounds_how_many_tools_run_at_once() { // 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( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "mock", Arc::new(MockModel::with_responses(vec![ multi_tool_call_response(vec![ @@ -6341,7 +6460,7 @@ async fn max_tool_concurrency_bounds_how_many_tools_run_at_once() { ]), 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"] { @@ -6374,13 +6493,13 @@ async fn max_tool_concurrency_bounds_how_many_tools_run_at_once() { #[tokio::test] async fn parallel_tool_results_keep_original_call_order_and_ids() { let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "mock", Arc::new(MockModel::with_responses(vec![ multi_tool_call_response(vec![("call-a", "alpha"), ("call-b", "beta")]), text_response("done", 4, 2), ])), - ); + ))); // alpha finishes *after* beta; the transcript must still list alpha first. let _ = probe_pair(&mut harness, (120, 0)); @@ -6406,13 +6525,13 @@ async fn parallel_tool_results_keep_original_call_order_and_ids() { #[tokio::test] async fn tool_wrap_middleware_no_longer_forces_serial_execution() { let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "mock", Arc::new(MockModel::with_responses(vec![ multi_tool_call_response(vec![("call-a", "alpha"), ("call-b", "beta")]), text_response("done", 4, 2), ])), - ); + ))); let max_seen = probe_pair(&mut harness, (40, 40)); // The wrap onion runs inside each concurrent call (`wrap_tool` takes a // shared `&RunContext`), so a wrap does not force serial execution. @@ -6437,13 +6556,13 @@ async fn tool_wrap_middleware_no_longer_forces_serial_execution() { #[tokio::test] async fn unknown_tool_recovery_keeps_its_slot_in_a_parallel_turn() { let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "mock", Arc::new(MockModel::with_responses(vec![ multi_tool_call_response(vec![("call-x", "missing"), ("call-b", "beta")]), text_response("done", 4, 2), ])), - ); + ))); let _ = probe_pair(&mut harness, (0, 0)); harness.with_policy(RunPolicy { unknown_tool: UnknownToolPolicy::ReturnToolError, @@ -6506,7 +6625,10 @@ async fn collect_stream(items: Vec<AgentStreamItem>) -> (Vec<AgentEvent>, AgentS #[tokio::test] async fn invoke_stream_yields_events_then_completed() { let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", Arc::new(MockModel::constant("hello there"))); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + Arc::new(MockModel::constant("hello there")), + ))); let items: Vec<AgentStreamItem> = harness .invoke_stream( @@ -6546,7 +6668,10 @@ async fn invoke_stream_yields_events_then_completed() { #[tokio::test] async fn invoke_stream_in_context_preserves_caller_context() { let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", Arc::new(MockModel::constant("hello there"))); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + Arc::new(MockModel::constant("hello there")), + ))); let ctx = RunContext::new( RunConfig::new("caller-run").with_thread("caller-thread"), (), @@ -6581,7 +6706,10 @@ async fn invoke_stream_in_context_preserves_caller_context() { #[tokio::test] async fn invoke_stream_in_context_unsubscribes_channel_listener() { let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", Arc::new(MockModel::constant("hello there"))); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + Arc::new(MockModel::constant("hello there")), + ))); let events = EventSink::new(); let ctx = RunContext::new(RunConfig::new("shared-events-run"), ()).with_events(events.clone()); @@ -6601,7 +6729,10 @@ fn invoke_stream_in_context_stream_is_send() { fn assert_send<T: Send>(_value: T) {} let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", Arc::new(MockModel::constant("hello there"))); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + Arc::new(MockModel::constant("hello there")), + ))); let ctx = RunContext::new(RunConfig::new("send-stream-run"), ()); assert_send(harness.invoke_stream_in_context(&(), ctx, vec![Message::user("hi")])); @@ -6610,13 +6741,13 @@ fn invoke_stream_in_context_stream_is_send() { #[tokio::test] async fn invoke_stream_surfaces_tool_lifecycle() { let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "mock", Arc::new(MockModel::with_responses(vec![ tool_call_response("c1", "spin", json!({})), text_response("done", 5, 2), ])), - ); + ))); harness.register_tool(Arc::new(FakeTool::new("spin", "again"))); let items: Vec<AgentStreamItem> = harness @@ -6655,10 +6786,10 @@ async fn invoke_stream_surfaces_tool_lifecycle() { async fn invoke_stream_yields_failed_terminal_on_error() { let mut harness: AgentHarness<()> = AgentHarness::new(); // A model that always asks for the tool, capped so the loop trips the limit. - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "mock", Arc::new(MockModel::with_tool_call("spin", json!({}))), - ); + ))); harness.register_tool(Arc::new(FakeTool::new("spin", "again"))); harness.with_policy(RunPolicy { limits: RunLimits::default().with_max_model_calls(1), @@ -6696,7 +6827,10 @@ async fn invoke_stream_scripted_incremental_deltas_surface_reasoning() { ModelStreamItem::Completed(ModelResponse::assistant("hello")), ]; let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", Arc::new(MockModel::streaming_script(script))); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + Arc::new(MockModel::streaming_script(script)), + ))); let items: Vec<AgentStreamItem> = harness .invoke_stream( @@ -6768,13 +6902,13 @@ async fn tool_completed_event_carries_outcome() { } let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "mock", Arc::new(MockModel::with_responses(vec![ tool_call_response("c1", "boom", json!({})), text_response("done", 1, 1), ])), - ); + ))); harness.register_tool(Arc::new(FailTool)); let recorder = Arc::new(RecordingListener::new()); @@ -6818,13 +6952,13 @@ async fn tool_started_event_carries_input_when_capture_enabled() { // completion, so `ToolStarted` must carry the same captured input the // policy already puts on `ToolCompleted`. let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "mock", Arc::new(MockModel::with_responses(vec![ tool_call_response("c1", "echo", json!({ "q": "weather" })), text_response("done", 1, 1), ])), - ); + ))); harness.register_tool(Arc::new(FakeTool::new("echo", "ok"))); harness.with_policy(RunPolicy { capture: PayloadCapture::all(), @@ -6860,13 +6994,13 @@ async fn tool_started_event_has_no_input_when_capture_disabled() { // Default policy is payload-free: `ToolStarted.input` stays `None` so no // tool argument is captured unless the host opts in. let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "mock", Arc::new(MockModel::with_responses(vec![ tool_call_response("c1", "echo", json!({ "q": "weather" })), text_response("done", 1, 1), ])), - ); + ))); harness.register_tool(Arc::new(FakeTool::new("echo", "ok"))); let recorder = Arc::new(RecordingListener::new()); @@ -6906,13 +7040,13 @@ fn continuing(text: &str, nudge: &str) -> ModelResponse { #[tokio::test] async fn continue_turn_keeps_the_floor_and_appends_the_nudge() { let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "mock", Arc::new(MockModel::with_responses(vec![ continuing("let me look that up", "(next)"), ModelResponse::assistant("found it".to_string()), ])), - ); + ))); let run = harness .invoke_default(&(), vec![Message::user("hi")]) @@ -6934,13 +7068,13 @@ async fn continue_turn_keeps_the_floor_and_appends_the_nudge() { #[tokio::test] async fn no_continue_turn_ends_on_the_first_tool_less_reply() { let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "mock", Arc::new(MockModel::with_responses(vec![ ModelResponse::assistant("done".to_string()), ModelResponse::assistant("never reached".to_string()), ])), - ); + ))); let run = harness .invoke_default(&(), vec![Message::user("hi")]) @@ -6956,7 +7090,7 @@ async fn no_continue_turn_ends_on_the_first_tool_less_reply() { #[tokio::test] async fn an_endless_continue_is_bounded_by_max_model_calls() { let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "mock", Arc::new(MockModel::with_responses(vec![ continuing("still going", "(next)"), @@ -6965,7 +7099,7 @@ async fn an_endless_continue_is_bounded_by_max_model_calls() { continuing("still going", "(next)"), continuing("still going", "(next)"), ])), - ); + ))); let config = RunConfig::new("continue-cap").with_max_model_calls(3); let err = harness @@ -7047,7 +7181,7 @@ impl Tool for ArgumentRecordingTool { #[tokio::test] async fn echoed_schema_arguments_are_unwrapped_before_validation() { let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "mock", Arc::new(MockModel::with_responses(vec![ tool_call_response( @@ -7063,7 +7197,7 @@ async fn echoed_schema_arguments_are_unwrapped_before_validation() { ), text_response("done", 1, 1), ])), - ); + ))); let seen = Arc::new(Mutex::new(Vec::new())); harness.register_tool(Arc::new(ArgumentRecordingTool { seen: Arc::clone(&seen), @@ -7104,13 +7238,13 @@ async fn wrapped_arguments_are_unwrapped_from_every_known_envelope_key() { json!({ "input": { "query": "rust" } }), ] { let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "mock", Arc::new(MockModel::with_responses(vec![ tool_call_response("call-1", "strict_lookup", envelope.clone()), text_response("done", 1, 1), ])), - ); + ))); let seen = Arc::new(Mutex::new(Vec::new())); harness.register_tool(Arc::new(ArgumentRecordingTool { seen: Arc::clone(&seen), @@ -7136,7 +7270,7 @@ async fn wrapped_arguments_are_unwrapped_from_every_known_envelope_key() { #[tokio::test] async fn echo_unwrap_leaves_a_tool_that_really_takes_properties_alone() { let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "mock", Arc::new(MockModel::with_responses(vec![ tool_call_response( @@ -7146,7 +7280,7 @@ async fn echo_unwrap_leaves_a_tool_that_really_takes_properties_alone() { ), text_response("done", 1, 1), ])), - ); + ))); let calls = Arc::new(Mutex::new(0)); harness.register_tool(Arc::new(SchemaShapedTool { calls: Arc::clone(&calls), @@ -7168,7 +7302,7 @@ async fn echo_unwrap_leaves_a_tool_that_really_takes_properties_alone() { #[tokio::test] async fn echo_unwrap_is_skipped_when_the_inner_value_is_still_invalid() { let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "mock", Arc::new(MockModel::with_responses(vec![ tool_call_response( @@ -7184,7 +7318,7 @@ async fn echo_unwrap_is_skipped_when_the_inner_value_is_still_invalid() { ), text_response("done", 1, 1), ])), - ); + ))); let calls = Arc::new(Mutex::new(0)); harness.register_tool(Arc::new(StrictLookupTool { calls: Arc::clone(&calls), @@ -7247,7 +7381,10 @@ async fn resolved_profile_schema_transform_is_applied_to_tool_schemas() { }), ); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", model.clone()); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + model.clone(), + ))); harness.register_tool(Arc::new(DefsTool)); harness @@ -7283,7 +7420,10 @@ async fn resolved_profile_schema_transform_is_applied_to_the_structured_output_s }), ); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", model.clone()); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + model.clone(), + ))); harness.with_policy(RunPolicy { default_response_format: Some(ResponseFormat::auto( "answer", @@ -7328,7 +7468,10 @@ async fn default_structured_mode_prompted_injects_schema_into_the_system_segment }), ); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", model.clone()); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + model.clone(), + ))); harness.with_policy(RunPolicy { default_response_format: Some(ResponseFormat::auto( "answer", @@ -7407,7 +7550,10 @@ async fn thinking_level_map_resolves_a_named_reasoning_effort() { }), ); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", model.clone()); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + model.clone(), + ))); harness.push_middleware(Arc::new(RequestReasoningEffort(ReasoningEffort::High))); harness @@ -7446,7 +7592,10 @@ async fn thinking_level_map_does_not_override_an_explicit_budget() { }), ); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", model.clone()); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + model.clone(), + ))); struct ExplicitBudget; #[async_trait] @@ -7498,7 +7647,10 @@ async fn thinking_tags_are_split_out_of_a_unary_response_into_a_thinking_block() }), ); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", model.clone()); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + model.clone(), + ))); let run = harness .invoke_default(&(), vec![Message::user("go")]) @@ -7527,7 +7679,10 @@ async fn a_model_without_thinking_tags_configured_leaves_text_untouched() { "<think>not a reasoning tag here</think>plain text", )])); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", model.clone()); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + model.clone(), + ))); let run = harness .invoke_default(&(), vec![Message::user("go")]) @@ -7556,7 +7711,9 @@ async fn streaming_ignores_leading_whitespace_on_the_first_text_delta_only() { }), ); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("stream", model); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "stream", model, + ))); let run = harness .invoke_streaming( @@ -7577,7 +7734,9 @@ async fn streaming_without_the_profile_flag_keeps_leading_whitespace() { let model = Arc::new(StreamingMock::from_text_chunks([" Hello", ", world"])); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("stream", model); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "stream", model, + ))); let run = harness .invoke_streaming( @@ -7749,13 +7908,13 @@ mod tool_effects_test { #[tokio::test] async fn agent_loop_writes_started_then_completed_around_a_tool_call() { let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "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()); @@ -7778,14 +7937,14 @@ mod tool_effects_test { #[tokio::test] async fn crash_between_started_and_settled_leaves_an_unresolved_row() { let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "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(), @@ -7945,14 +8104,14 @@ mod tool_effects_test { #[tokio::test] async fn ledger_started_failure_aborts_the_run_under_the_default_policy() { let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "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 { @@ -7972,13 +8131,13 @@ mod tool_effects_test { #[tokio::test] async fn ledger_started_failure_is_ignored_under_the_continue_policy() { let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "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()); @@ -8040,13 +8199,13 @@ mod tool_effects_test { async fn resume_deferred_settles_the_deferred_row_completed_without_a_synthesized_crash_answer() { let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "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, @@ -8097,14 +8256,14 @@ mod tool_effects_test { #[tokio::test] async fn resume_deferred_settles_the_deferred_row_failed_when_the_retry_errors() { let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "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, @@ -8139,12 +8298,12 @@ mod tool_effects_test { #[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( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "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", @@ -8259,7 +8418,10 @@ async fn tiered_system_messages_become_one_cacheable_segment_each() { // fingerprint must cover both, in order. let model = Arc::new(crate::testkit::ScriptedModel::replies(vec!["done"])); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", model.clone()); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + model.clone(), + ))); harness.with_policy(RunPolicy { cache: CachePolicy { protect_prompt_prefix: true, @@ -9036,7 +9198,10 @@ async fn run_policy_default_reasoning_reaches_every_request() { let model = Arc::new(ScriptedModel::replies(vec!["done"])); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", model.clone()); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + model.clone(), + ))); harness.with_policy(RunPolicy { default_reasoning: Some(ReasoningConfig::effort(ReasoningEffort::Low)), ..RunPolicy::default() @@ -9077,7 +9242,10 @@ async fn run_policy_default_reasoning_goes_through_the_profile_map() { }), ); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", model.clone()); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + model.clone(), + ))); harness.with_policy(RunPolicy { default_reasoning: Some(ReasoningConfig::effort(ReasoningEffort::High)), ..RunPolicy::default() @@ -9104,7 +9272,10 @@ async fn request_reasoning_wins_over_the_run_policy_default() { let model = Arc::new(ScriptedModel::replies(vec!["done"])); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", model.clone()); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + model.clone(), + ))); harness.with_policy(RunPolicy { default_reasoning: Some(ReasoningConfig::effort(ReasoningEffort::Low)), ..RunPolicy::default() @@ -9138,7 +9309,10 @@ async fn thread_reasoning_snapshot_flows_through_policy_and_request_override_win let model = Arc::new(ScriptedModel::replies(vec!["first", "second", "third"])); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", model.clone()); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + model.clone(), + ))); harness.with_policy(RunPolicy { default_reasoning: turn_snapshot, ..RunPolicy::default() @@ -9304,7 +9478,9 @@ fn structured_mixed_harness( ) -> AgentHarness<()> { let mut harness: AgentHarness<()> = AgentHarness::new(); harness - .register_model("mock", model) + .with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", model, + ))) .register_tool(tool) .with_policy(RunPolicy { end_strategy, @@ -9381,13 +9557,13 @@ async fn graceful_mixed_turn_without_truncation_still_finishes_in_one_call() { #[tokio::test] async fn dynamic_toolset_fold_is_reported_as_a_transcript_rewrite() { let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "mock", Arc::new(MockModel::with_responses(vec![ tool_call_response("call-1", "search", json!({"q": "x"})), text_response("done", 4, 2), ])), - ); + ))); let search: Arc<dyn Tool> = Arc::new(FakeTool::new("search", "search-output")); let browse: Arc<dyn Tool> = Arc::new(FakeTool::new("browse", "browse-output")); let toolset = Arc::new(DynamicToolSet { @@ -9427,7 +9603,10 @@ async fn a_dead_call_retry_preserves_configured_reasoning_effort() { ])); let mut harness: AgentHarness<()> = AgentHarness::new(); harness - .register_model("mock", Arc::clone(&model) as _) + .with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + Arc::clone(&model) as _, + ))) .register_tool(Arc::new(FakeTool::new("fake", "tool output"))); harness.with_policy(RunPolicy { default_reasoning: Some(ReasoningConfig::effort(ReasoningEffort::High)), diff --git a/crates/tinyagents-harness/src/agent_loop/model_call.rs b/crates/tinyagents-harness/src/agent_loop/model_call.rs index 6099f6a04..5e02d8b7e 100644 --- a/crates/tinyagents-harness/src/agent_loop/model_call.rs +++ b/crates/tinyagents-harness/src/agent_loop/model_call.rs @@ -22,6 +22,71 @@ use crate::no_progress::StreamTextStallDetector; use crate::retry::{FailoverDecision, FailoverReason, FailoverState, decide}; use tinyinference_llm::cache::CachePolicy; +pub(crate) fn local_resolve_request(request: &ModelRequest) -> crate::host::ModelResolveRequest { + let mut resolve = crate::host::ModelResolveRequest::new("local"); + if let Some(pin) = request.model.as_deref() { + resolve = resolve.with_model_pin(pin); + } else if let Some(hint) = request.model_hints.first() { + resolve = resolve.with_model_hint(&hint.model); + } + if let Some(required) = request.required_capabilities.clone() { + resolve = resolve.with_required_capabilities(required); + } + resolve +} + +pub(crate) fn preview_local_model_profile<State: Send + Sync>( + resolver: Option<&Arc<dyn crate::host::ModelResolver<State>>>, + fixed_model: Option<&Arc<dyn ChatModel<State>>>, + request: &ModelRequest, +) -> Option<tinyinference_llm::model::ModelProfile> { + if let Some(resolver) = resolver { + let resolve = local_resolve_request(request); + let model = resolver.preview(&resolve)?; + if model_is_eligible(model.as_ref(), request.required_capabilities.as_ref()) { + return model.profile().cloned(); + } + if request.model.is_some() || !request.model_hints.is_empty() { + let mut default = crate::host::ModelResolveRequest::new("local"); + if let Some(required) = request.required_capabilities.clone() { + default = default.with_required_capabilities(required); + } + return resolver + .preview(&default) + .filter(|model| { + model_is_eligible(model.as_ref(), request.required_capabilities.as_ref()) + }) + .and_then(|model| model.profile().cloned()); + } + return None; + } + fixed_model + .filter(|model| model_is_eligible(model.as_ref(), request.required_capabilities.as_ref())) + .and_then(|model| model.profile().cloned()) +} + +pub(super) fn model_is_eligible<State: Send + Sync>( + model: &dyn ChatModel<State>, + required: Option<&tinyinference_llm::model::CapabilitySet>, +) -> bool { + model.profile().is_none_or(|profile| { + profile.is_usable() && required.is_none_or(|required| profile.satisfies(required)) + }) +} + +fn ensure_model_eligible<State: Send + Sync>( + model: &dyn ChatModel<State>, + request: &ModelRequest, +) -> Result<()> { + if model_is_eligible(model, request.required_capabilities.as_ref()) { + Ok(()) + } else { + Err(TinyAgentsError::ModelNotFound( + request.model.clone().unwrap_or_else(|| "<default>".into()), + )) + } +} + /// Converts a configured stream window to a [`Duration`]. `None` and a zero /// window both mean "no bound": a zero-length window would fail every stream /// before it could deliver anything. @@ -30,28 +95,9 @@ fn positive_window(ms: Option<u64>) -> Option<Duration> { } impl<State: Send + Sync, Ctx: Send + Sync> AgentHarness<State, Ctx> { - /// Previews which model `request` would reach through the local - /// [`crate::model_registry::ModelRegistry`] and returns its capability - /// profile, without dispatching anything. - /// - /// A pure registry lookup (no network call), so it is cheap enough to run - /// before every `before_model` pass and before every tool-change patch. - /// `None` when nothing resolves or the model advertises no profile. A - /// Hosted runs use their resolver for the profile shown to middleware; - /// this local preview is only their initial fallback. - pub(super) fn preview_model_profile( - &self, - request: &ModelRequest, - ) -> Option<tinyinference_llm::model::ModelProfile> { - self.models - .resolve_request(request, None, None) - .and_then(|binding| binding.model.profile().cloned()) - } - /// Resolves the model binding through the host's routing authority when /// this run is a hosted invocation; returns `None` for a plain SDK run so - /// the caller falls through to local [`crate::model_registry::ModelRegistry`] - /// resolution instead. + /// the caller falls through to the configured local resolver instead. /// /// The first call for a run (`ctx.depth() == 0`) is flagged /// [`as_team_lead`][crate::host::ModelResolveRequest::as_team_lead] so the @@ -68,7 +114,7 @@ impl<State: Send + Sync, Ctx: Send + Sync> AgentHarness<State, Ctx> { &self, ctx: &RunContext<Ctx>, request: &ModelRequest, - ) -> Result<Option<ResolvedModelBinding<State>>> { + ) -> Result<Option<ModelBinding<State>>> { let Some(host_run) = crate::runtime::host_invocation_binding::<State, Ctx>(ctx)? else { return Ok(None); }; @@ -123,7 +169,7 @@ impl<State: Send + Sync, Ctx: Send + Sync> AgentHarness<State, Ctx> { .profile() .and_then(|profile| profile.model.clone()) .unwrap_or_else(|| format!("host:{}", host_run.agent_id)); - Ok(Some(ResolvedModelBinding { + Ok(Some(ModelBinding { resolved: ResolvedModel { name, requested: resolve.model_pin, @@ -136,6 +182,69 @@ impl<State: Send + Sync, Ctx: Send + Sync> AgentHarness<State, Ctx> { model, })) } + + pub(crate) async fn resolve_local_model( + &self, + ctx: &RunContext<Ctx>, + request: &ModelRequest, + ) -> Result<ModelBinding<State>> { + let resolve = local_resolve_request(request); + let (model, selected_name) = if let Some(resolver) = self.model_resolver.as_ref() { + let (budget, bound) = self.model_call_budget(ctx); + let mut model = ctx + .bounded(budget, resolver.resolve(&resolve), || { + format!( + "local model resolution for run `{}` exceeded its {bound}", + ctx.run_id() + ) + }) + .await?; + let mut selected_name = resolver.selected_name(&resolve); + if !model_is_eligible(model.as_ref(), request.required_capabilities.as_ref()) + && (request.model.is_some() || !request.model_hints.is_empty()) + { + let mut default_request = crate::host::ModelResolveRequest::new("local"); + if let Some(required) = request.required_capabilities.clone() { + default_request = default_request.with_required_capabilities(required); + } + let (budget, bound) = self.model_call_budget(ctx); + model = ctx + .bounded(budget, resolver.resolve(&default_request), || { + format!( + "local default model resolution for run `{}` exceeded its {bound}", + ctx.run_id() + ) + }) + .await?; + selected_name = resolver.selected_name(&default_request); + } + (model, selected_name) + } else if let Some(model) = self.fixed_model.clone() { + let selected_name = model.profile().and_then(|profile| profile.model.clone()); + (model, selected_name) + } else { + return Err(TinyAgentsError::ModelNotFound( + request.model.clone().unwrap_or_else(|| "<default>".into()), + )); + }; + ensure_model_eligible(model.as_ref(), request)?; + let name = selected_name + .or_else(|| model.profile().and_then(|profile| profile.model.clone())) + .or_else(|| request.model.clone()) + .unwrap_or_else(|| "default".into()); + Ok(ModelBinding { + resolved: ResolvedModel { + name: name.clone(), + requested: request.model.clone(), + source: if request.model.as_deref() == Some(name.as_str()) { + ModelResolutionSource::RequestOverride + } else { + ModelResolutionSource::RegistryDefault + }, + }, + model, + }) + } /// Invokes a model, consulting the local response cache around the /// retry/fallback path. /// @@ -185,7 +294,7 @@ impl<State: Send + Sync, Ctx: Send + Sync> AgentHarness<State, Ctx> { ctx: &mut RunContext<Ctx>, request: &ModelRequest, call_id: &CallId, - binding: ResolvedModelBinding<State>, + binding: ModelBinding<State>, shape: &super::dialect::CallShape, ) -> Result<ModelResponse> { let streaming = shape.streaming; @@ -610,7 +719,7 @@ impl<State: Send + Sync, Ctx: Send + Sync> AgentHarness<State, Ctx> { ctx: &mut RunContext<Ctx>, request: &ModelRequest, call_id: &CallId, - binding: ResolvedModelBinding<State>, + binding: ModelBinding<State>, shape: &super::dialect::CallShape, ) -> Result<ModelResponse> { let streaming = shape.streaming; @@ -637,8 +746,9 @@ impl<State: Send + Sync, Ctx: Send + Sync> AgentHarness<State, Ctx> { // anyway — failing again is better than refusing to try. if ctx.limits.is_model_skipped(¤t_name) && crate::runtime::host_invocation_binding::<State, Ctx>(ctx)?.is_none() - && let Some((name, next_model)) = - self.select_fallback(ctx, request, ¤t_name, &mut visited, None) + && let Some((name, next_model)) = self + .select_fallback(ctx, request, ¤t_name, &mut visited, None) + .await? { tracing::debug!( call_id = %call_id.as_str(), @@ -901,13 +1011,9 @@ impl<State: Send + Sync, Ctx: Send + Sync> AgentHarness<State, Ctx> { .and_then(|profile| profile.max_input_tokens) .unwrap_or(u64::MAX) }); - let selected = self.select_fallback( - ctx, - request, - ¤t_name, - &mut visited, - larger_than, - ); + let selected = self + .select_fallback(ctx, request, ¤t_name, &mut visited, larger_than) + .await?; let decision = decide( reason, FailoverState { @@ -967,18 +1073,18 @@ impl<State: Send + Sync, Ctx: Send + Sync> AgentHarness<State, Ctx> { /// the same gate here a primary failure could silently fall back to a model /// that can't call tools, lacks vision, or has a smaller context window /// (issue #4641). `allow_retired` is `false` to match - /// `ModelRegistry::resolve_request`. Every skipped candidate is recorded as + /// hosted model selection. Every skipped candidate is recorded as /// visited and surfaced as [`AgentEvent::FallbackSkipped`]. With /// `larger_than = Some(n)` a candidate must also declare a /// `max_input_tokens` strictly greater than `n` (context-overflow failover). - fn select_fallback( + async fn select_fallback( &self, ctx: &mut RunContext<Ctx>, request: &ModelRequest, cursor: &str, visited: &mut std::collections::HashSet<String>, larger_than: Option<u64>, - ) -> Option<(String, Arc<dyn ChatModel<State>>)> { + ) -> Result<Option<(String, Arc<dyn ChatModel<State>>)>> { let required = request.required_capabilities.as_ref(); // A steered model outside the chain falls back through the whole // chain (the original primary's), not from a position it lacks. @@ -998,8 +1104,37 @@ impl<State: Send + Sync, Ctx: Send + Sync> AgentHarness<State, Ctx> { }) .map(str::to_owned) .filter(|name| !visited.contains(name)); - let (name, next_model) = - next.and_then(|name| self.models.get(&name).map(|m| (name, m)))?; + let Some(name) = next else { + return Ok(None); + }; + let host = crate::runtime::host_invocation_binding::<State, Ctx>(ctx)?; + let resolver = host + .as_ref() + .map(|binding| &binding.host.models) + .or(self.model_resolver.as_ref()); + let next_model = if let Some(resolver) = resolver { + let (budget, bound) = self.model_call_budget(ctx); + ctx.bounded(budget, resolver.resolve_named(&name), || { + format!( + "fallback model resolution for run `{}` exceeded its {bound}", + ctx.run_id() + ) + }) + .await? + } else { + self.fixed_model.clone().filter(|model| { + model.profile().and_then(|profile| profile.model.as_deref()) + == Some(name.as_str()) + }) + }; + let Some(next_model) = next_model else { + visited.insert(name.clone()); + ctx.emit(AgentEvent::FallbackSkipped { + model: name.clone(), + }); + cursor = name; + continue; + }; let written_off = ctx.limits.is_model_skipped(&name); // `larger_than` (context-overflow failover): the candidate must // *declare* a window strictly above the given one. @@ -1009,7 +1144,7 @@ impl<State: Send + Sync, Ctx: Send + Sync> AgentHarness<State, Ctx> { .and_then(|profile| profile.max_input_tokens) .is_none_or(|window| window <= floor) }); - if written_off || too_small || !model_eligible(next_model.as_ref(), required, false) { + if written_off || too_small || !model_is_eligible(next_model.as_ref(), required) { visited.insert(name.clone()); ctx.emit(AgentEvent::FallbackSkipped { model: name.clone(), @@ -1017,7 +1152,7 @@ impl<State: Send + Sync, Ctx: Send + Sync> AgentHarness<State, Ctx> { cursor = name; continue; } - return Some((name, next_model)); + return Ok(Some((name, next_model))); } } @@ -1835,8 +1970,8 @@ impl<State: Send + Sync, Ctx: Send + Sync> ModelCallBase<'_, State, Ctx> { &self, ctx: &mut RunContext<Ctx>, request: &ModelRequest, - ) -> Result<ResolvedModelBinding<State>> { - let captured = || ResolvedModelBinding { + ) -> Result<ModelBinding<State>> { + let captured = || ModelBinding { resolved: self.resolved.clone(), model: Arc::clone(&self.model), }; @@ -1864,11 +1999,58 @@ impl<State: Send + Sync, Ctx: Send + Sync> ModelCallBase<'_, State, Ctx> { if requested == self.resolved.name { return Ok(captured()); } - match self.harness.models.resolve_request(request, None, None) { - Some(binding) - if binding.resolved.source == ModelResolutionSource::RequestOverride - && binding.resolved.name == requested => + let binding = if let Some(resolver) = self.harness.model_resolver.as_ref() { + let (budget, bound) = self.harness.model_call_budget(ctx); + let named = match ctx + .bounded(budget, resolver.resolve_named(requested), || { + format!( + "model override resolution for run `{}` exceeded its {bound}", + ctx.run_id() + ) + }) + .await { + Ok(model) => model, + Err( + error @ (TinyAgentsError::Cancelled + | TinyAgentsError::Timeout(_) + | TinyAgentsError::CallTimeout(_)), + ) => return Err(error), + Err(_) => None, + }; + named + .filter(|model| { + model_is_eligible(model.as_ref(), request.required_capabilities.as_ref()) + }) + .map(|model| ModelBinding { + resolved: ResolvedModel { + name: requested.to_owned(), + requested: Some(requested.to_owned()), + source: ModelResolutionSource::RequestOverride, + }, + model, + }) + } else { + self.harness + .fixed_model + .clone() + .filter(|model| { + model.profile().and_then(|profile| profile.model.as_deref()) == Some(requested) + }) + .filter(|model| { + model_is_eligible(model.as_ref(), request.required_capabilities.as_ref()) + }) + .map(|model| ModelBinding { + resolved: ResolvedModel { + name: requested.to_owned(), + requested: Some(requested.to_owned()), + source: ModelResolutionSource::RequestOverride, + }, + model, + }) + }; + match binding { + Some(binding) => { tracing::debug!( call_id = %self.call_id.as_str(), from = %self.resolved.name, @@ -1877,7 +2059,7 @@ impl<State: Send + Sync, Ctx: Send + Sync> ModelCallBase<'_, State, Ctx> { ); Ok(binding) } - _ => { + None => { tracing::warn!( call_id = %self.call_id.as_str(), requested = %requested, diff --git a/crates/tinyagents-harness/src/agent_loop/model_call_failover_tests.rs b/crates/tinyagents-harness/src/agent_loop/model_call_failover_tests.rs index fc239849a..2119121f8 100644 --- a/crates/tinyagents-harness/src/agent_loop/model_call_failover_tests.rs +++ b/crates/tinyagents-harness/src/agent_loop/model_call_failover_tests.rs @@ -111,8 +111,15 @@ fn harness_with_primary( max_attempts: usize, ) -> AgentHarness<()> { let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("primary", primary); - harness.register_model("backup", backup.clone()); + let resolver = crate::host::PrebuiltModelResolver::new( + primary, + [( + "backup".to_string(), + backup.clone() as Arc<dyn ChatModel<()>>, + )], + ) + .with_primary_name("primary"); + harness.with_model_resolver(Arc::new(resolver)); harness.with_policy(RunPolicy { retry: RetryPolicy::default() .with_max_attempts(max_attempts) @@ -129,8 +136,15 @@ fn harness( max_attempts: usize, ) -> AgentHarness<()> { let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("primary", primary.clone()); - harness.register_model("backup", backup.clone()); + let resolver = crate::host::PrebuiltModelResolver::new( + primary.clone(), + [( + "backup".to_string(), + backup.clone() as Arc<dyn ChatModel<()>>, + )], + ) + .with_primary_name("primary"); + harness.with_model_resolver(Arc::new(resolver)); harness.with_policy(RunPolicy { retry: RetryPolicy::default() .with_max_attempts(max_attempts) @@ -267,8 +281,15 @@ impl ChatModel<()> for Windowed { fn windowed_harness(primary: &Arc<Windowed>, sibling: &Arc<Windowed>) -> AgentHarness<()> { let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("primary", primary.clone()); - harness.register_model("sibling", sibling.clone()); + let resolver = crate::host::PrebuiltModelResolver::new( + primary.clone(), + [( + "sibling".to_string(), + sibling.clone() as Arc<dyn ChatModel<()>>, + )], + ) + .with_primary_name("primary"); + harness.with_model_resolver(Arc::new(resolver)); harness.with_policy(RunPolicy { retry: RetryPolicy::default() .with_max_attempts(2) @@ -383,7 +404,10 @@ async fn plain_auth_failure_is_retried_on_the_next_call() { fn single_model_harness(model: &Arc<ScriptedOutcomes>) -> AgentHarness<()> { let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("primary", model.clone()); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "primary", + model.clone(), + ))); harness.with_policy(RunPolicy { retry: RetryPolicy::default() .with_max_attempts(1) @@ -432,9 +456,15 @@ async fn a_fallback_chosen_by_the_skip_hint_is_not_tried_twice() { let b = ScriptedOutcomes::failing(400, "bad request", false); let c = ScriptedOutcomes::failing(400, "bad request", false); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("a", a.clone()); - harness.register_model("b", b.clone()); - harness.register_model("c", c.clone()); + let resolver = crate::host::PrebuiltModelResolver::new( + a.clone(), + [ + ("b".to_string(), b.clone() as Arc<dyn ChatModel<()>>), + ("c".to_string(), c.clone() as Arc<dyn ChatModel<()>>), + ], + ) + .with_primary_name("a"); + harness.with_model_resolver(Arc::new(resolver)); harness.with_policy(RunPolicy { retry: RetryPolicy::default() .with_max_attempts(1) @@ -485,12 +515,12 @@ async fn a_cancel_during_a_retry_backoff_stops_without_waiting_it_out() { for middleware in [false, true] { let token = crate::cancel::CancellationToken::new(); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "flaky", Arc::new(CancelThenFail { token: token.clone(), }), - ); + ))); if middleware { harness.push_model_middleware(Arc::new( crate::middleware::library::RetryMiddleware::new(backoff.clone()), diff --git a/crates/tinyagents-harness/src/agent_loop/model_call_reasoning_policy_tests.rs b/crates/tinyagents-harness/src/agent_loop/model_call_reasoning_policy_tests.rs index 07c1ec27e..36214258c 100644 --- a/crates/tinyagents-harness/src/agent_loop/model_call_reasoning_policy_tests.rs +++ b/crates/tinyagents-harness/src/agent_loop/model_call_reasoning_policy_tests.rs @@ -103,7 +103,10 @@ async fn reasoning_past_the_requested_budget_finishes_without_client_interruptio "complete answer", )])); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", Arc::clone(&model) as _); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + Arc::clone(&model) as _, + ))); harness.with_policy(RunPolicy { default_reasoning: Some(ReasoningConfig { effort: Some(ReasoningEffort::High), diff --git a/crates/tinyagents-harness/src/agent_loop/model_profile_preview_tests.rs b/crates/tinyagents-harness/src/agent_loop/model_profile_preview_tests.rs index b4586a1cd..fa58a0a94 100644 --- a/crates/tinyagents-harness/src/agent_loop/model_profile_preview_tests.rs +++ b/crates/tinyagents-harness/src/agent_loop/model_profile_preview_tests.rs @@ -7,15 +7,44 @@ use std::sync::{Arc, Mutex}; use async_trait::async_trait; +use crate::host::{ModelResolveRequest, ModelResolver}; use crate::middleware::Middleware; use crate::testkit::ScriptedModel; -use tinyinference_llm::model::ModelProfile; +use tinyinference_llm::model::{ChatModel, ModelProfile, ModelRequest}; /// Records the profile each `before_model` call saw. struct SeenProfiles(Arc<Mutex<Vec<Option<ModelProfile>>>>); struct SelectModel(&'static str); +struct DynamicResolver { + primary: Arc<dyn ChatModel<()>>, + secondary: Arc<dyn ChatModel<()>>, + pins: Arc<Mutex<Vec<Option<String>>>>, +} + +#[async_trait] +impl ModelResolver<()> for DynamicResolver { + async fn resolve(&self, request: &ModelResolveRequest) -> Result<Arc<dyn ChatModel<()>>> { + let pin = request.model_pin.clone(); + self.pins.lock().unwrap().push(pin.clone()); + Ok(if pin.as_deref() == Some("second") { + self.secondary.clone() + } else { + self.primary.clone() + }) + } + + fn selected_name(&self, request: &ModelResolveRequest) -> Option<String> { + Some( + request + .model_pin + .clone() + .unwrap_or_else(|| "first".to_owned()), + ) + } +} + #[async_trait] impl Middleware<()> for SelectModel { fn name(&self) -> &str { @@ -60,7 +89,9 @@ async fn before_model_sees_the_target_model_profile() { let model = Arc::new(ScriptedModel::replies(vec!["done"]).with_profile(profile.clone())); let seen = Arc::new(Mutex::new(Vec::new())); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", model); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", model, + ))); harness.push_middleware(Arc::new(SeenProfiles(seen.clone()))); harness @@ -74,6 +105,7 @@ async fn before_model_sees_the_target_model_profile() { #[tokio::test] async fn later_middleware_sees_model_selected_by_earlier_middleware() { + use crate::host::{ModelResolver, PrebuiltModelResolver}; let first_profile = ModelProfile { model: Some("first".into()), hoists_system_messages: true, @@ -86,14 +118,15 @@ async fn later_middleware_sees_model_selected_by_earlier_middleware() { }; let seen = Arc::new(Mutex::new(Vec::new())); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( - "first", - Arc::new(ScriptedModel::replies(vec!["wrong"]).with_profile(first_profile)), - ); - harness.register_model( - "second", - Arc::new(ScriptedModel::replies(vec!["done"]).with_profile(second_profile.clone())), + let first: Arc<dyn ChatModel<()>> = + Arc::new(ScriptedModel::replies(vec!["wrong"]).with_profile(first_profile)); + let second: Arc<dyn ChatModel<()>> = + Arc::new(ScriptedModel::replies(vec!["done"]).with_profile(second_profile.clone())); + let resolver: Arc<dyn ModelResolver<()>> = Arc::new( + PrebuiltModelResolver::new(first, [("second".to_owned(), second)]) + .with_primary_name("first"), ); + harness.with_model_resolver(resolver); harness.push_middleware(Arc::new(SelectModel("second"))); harness.push_middleware(Arc::new(SeenProfiles(seen.clone()))); @@ -104,3 +137,82 @@ async fn later_middleware_sees_model_selected_by_earlier_middleware() { assert_eq!(seen.lock().unwrap().as_slice(), [Some(second_profile)]); } + +#[tokio::test] +async fn later_middleware_sees_dynamically_resolved_model_without_preview() { + let first_profile = ModelProfile { + model: Some("first".into()), + hoists_system_messages: false, + ..ModelProfile::default() + }; + let second_profile = ModelProfile { + model: Some("second".into()), + hoists_system_messages: true, + ..ModelProfile::default() + }; + let primary: Arc<dyn ChatModel<()>> = + Arc::new(ScriptedModel::replies(vec!["wrong"]).with_profile(first_profile)); + let secondary: Arc<dyn ChatModel<()>> = + Arc::new(ScriptedModel::replies(vec!["done"]).with_profile(second_profile.clone())); + + let pins = Arc::new(Mutex::new(Vec::new())); + let resolver = Arc::new(DynamicResolver { + primary, + secondary, + pins: pins.clone(), + }); + let seen = Arc::new(Mutex::new(Vec::new())); + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness.with_model_resolver(resolver); + harness.push_middleware(Arc::new(SelectModel("second"))); + harness.push_middleware(Arc::new(SeenProfiles(seen.clone()))); + + let run = harness + .invoke_default(&(), vec![Message::user("go")]) + .await + .expect("run succeeds using the selected route"); + assert_eq!(run.text().as_deref(), Some("done")); + assert_eq!(seen.lock().unwrap().as_slice(), [Some(second_profile)]); + assert_eq!( + pins.lock().unwrap().as_slice(), + [None, Some("second".to_owned())] + ); +} + +#[test] +fn profile_preview_uses_the_eligible_default_for_an_ineligible_pin() { + use crate::host::{ModelResolver, PrebuiltModelResolver}; + use tinyinference_llm::model::{CapabilitySet, ModelRequest}; + + let primary: Arc<dyn ChatModel<()>> = Arc::new( + ScriptedModel::replies(vec!["primary"]).with_profile(ModelProfile { + model: Some("primary".into()), + tool_calling: true, + ..ModelProfile::default() + }), + ); + let pinned: Arc<dyn ChatModel<()>> = Arc::new( + ScriptedModel::replies(vec!["pinned"]).with_profile(ModelProfile { + model: Some("pinned".into()), + tool_calling: false, + ..ModelProfile::default() + }), + ); + let resolver: Arc<dyn ModelResolver<()>> = Arc::new( + PrebuiltModelResolver::new(primary, [("pinned".to_owned(), pinned)]) + .with_primary_name("primary"), + ); + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness.with_model_resolver(resolver); + + let request = ModelRequest::default() + .with_model("pinned") + .with_required_capabilities(CapabilitySet { + tool_calling: true, + ..CapabilitySet::default() + }); + let profile = harness + .preview_model_profile_for_request(&request) + .expect("eligible primary profile is previewed"); + assert_eq!(profile.model.as_deref(), Some("primary")); +} diff --git a/crates/tinyagents-harness/src/agent_loop/model_switch.rs b/crates/tinyagents-harness/src/agent_loop/model_switch.rs index 350efcd46..eb0947a5c 100644 --- a/crates/tinyagents-harness/src/agent_loop/model_switch.rs +++ b/crates/tinyagents-harness/src/agent_loop/model_switch.rs @@ -13,6 +13,7 @@ //! up front instead, through the existing //! [`AgentEvent::ModelOverrideSkipped`] diagnostic. +use super::model_call::model_is_eligible; use super::*; impl<State: Send + Sync, Ctx: Send + Sync> AgentHarness<State, Ctx> { @@ -49,30 +50,55 @@ impl<State: Send + Sync, Ctx: Send + Sync> AgentHarness<State, Ctx> { /// call; a rejection puts it back when `request.model` still carries the /// rejected name, so that name never reaches resolution or an adapter that /// honours `request.model`. A model a middleware picked is left alone. - pub(super) fn apply_steered_model_switch( + pub(super) async fn apply_steered_model_switch( &self, ctx: &mut RunContext<Ctx>, request: &mut ModelRequest, model_before_switch: &Option<String>, - ) { + cached_candidate: Option<Arc<dyn ChatModel<State>>>, + ) -> Result<()> { let Some(handle) = ctx.steering.clone() else { - return; + return Ok(()); }; let Some(requested) = handle.model_override() else { - return; + return Ok(()); }; let hosted = !matches!( crate::runtime::host_invocation_binding::<State, Ctx>(ctx), Ok(None) ); - let eligible = !hosted - && self.models.get(&requested).is_some_and(|model| { - model_eligible( - model.as_ref(), - request.required_capabilities.as_ref(), - false, - ) - }); + let candidate = if hosted { + None + } else if cached_candidate.is_some() { + cached_candidate + } else if let Some(resolver) = self.model_resolver.as_ref() { + let (budget, bound) = self.model_call_budget(ctx); + match ctx + .bounded(budget, resolver.resolve_named(&requested), || { + format!( + "steered model resolution for run `{}` exceeded its {bound}", + ctx.run_id() + ) + }) + .await + { + Ok(candidate) => candidate, + Err( + error @ (TinyAgentsError::Cancelled + | TinyAgentsError::Timeout(_) + | TinyAgentsError::CallTimeout(_)), + ) => return Err(error), + Err(_) => None, + } + } else { + self.fixed_model.clone().filter(|model| { + model.profile().and_then(|profile| profile.model.as_deref()) + == Some(requested.as_str()) + }) + }; + let eligible = candidate.as_ref().is_some_and(|model| { + model_is_eligible(model.as_ref(), request.required_capabilities.as_ref()) + }); if eligible { if request.model.as_deref() != Some(requested.as_str()) { match request.model.as_deref() { @@ -92,15 +118,28 @@ impl<State: Send + Sync, Ctx: Send + Sync> AgentHarness<State, Ctx> { } request.model = Some(requested); } - return; + return Ok(()); } if request.model.as_deref() == Some(requested.as_str()) { request.model = model_before_switch.clone(); } + let resolve = super::model_call::local_resolve_request(request); let resolved = self - .models - .resolve_request(request, None, None) - .map_or_else(|| "<default>".to_string(), |binding| binding.resolved.name); + .model_resolver + .as_ref() + .and_then(|resolver| resolver.selected_name(&resolve)) + .or_else(|| { + self.model_resolver + .as_ref() + .and_then(|resolver| resolver.preview(&resolve)) + .and_then(|model| model.profile().and_then(|profile| profile.model.clone())) + }) + .or_else(|| { + self.fixed_model + .as_ref() + .and_then(|model| model.profile().and_then(|profile| profile.model.clone())) + }) + .unwrap_or_else(|| "<default>".to_string()); tracing::warn!( target: "tinyagents::steering", run_id = %ctx.run_id(), @@ -123,6 +162,7 @@ impl<State: Send + Sync, Ctx: Send + Sync> AgentHarness<State, Ctx> { accepted: false, }); } + Ok(()) } /// Reports a steered switch as applied (`Steered { accepted: true }`, once diff --git a/crates/tinyagents-harness/src/agent_loop/model_switch_tests.rs b/crates/tinyagents-harness/src/agent_loop/model_switch_tests.rs index b5e89d8f5..30fdb0648 100644 --- a/crates/tinyagents-harness/src/agent_loop/model_switch_tests.rs +++ b/crates/tinyagents-harness/src/agent_loop/model_switch_tests.rs @@ -4,6 +4,7 @@ //! fallback chain) belongs to the new model. use std::collections::VecDeque; +use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{Arc, Mutex}; use async_trait::async_trait; @@ -126,12 +127,83 @@ fn two_models( ) -> (AgentHarness<()>, Arc<ScriptedModel>, Arc<ScriptedModel>) { let (a, b) = (Arc::new(a), Arc::new(b)); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("a", a.clone()); - harness.register_model("b", b.clone()); + let resolver = crate::host::PrebuiltModelResolver::new( + a.clone(), + [("b".to_string(), b.clone() as Arc<dyn ChatModel<()>>)], + ) + .with_primary_name("a"); + harness.with_model_resolver(Arc::new(resolver)); harness.register_tool(Arc::new(FakeTool::returning("lookup", "ok"))); (harness, a, b) } +struct CountingResolver { + inner: crate::host::PrebuiltModelResolver<()>, + named_calls: Arc<AtomicUsize>, +} + +#[async_trait] +impl crate::host::ModelResolver<()> for CountingResolver { + async fn resolve( + &self, + request: &crate::host::ModelResolveRequest, + ) -> crate::error::Result<Arc<dyn ChatModel<()>>> { + self.inner.resolve(request).await + } + + async fn resolve_named( + &self, + name: &str, + ) -> crate::error::Result<Option<Arc<dyn ChatModel<()>>>> { + self.named_calls.fetch_add(1, Ordering::SeqCst); + self.inner.resolve_named(name).await + } + + fn preview( + &self, + request: &crate::host::ModelResolveRequest, + ) -> Option<Arc<dyn ChatModel<()>>> { + self.inner.preview(request) + } + + fn selected_name(&self, request: &crate::host::ModelResolveRequest) -> Option<String> { + self.inner.selected_name(request) + } +} + +#[tokio::test] +async fn sticky_switch_reuses_the_model_binding_for_unchanged_inputs() { + let a = Arc::new(ScriptedModel::replies(vec!["from a"]).with_profile(profile("openai", "a"))); + let b = Arc::new(ScriptedModel::replies(vec!["from b"]).with_profile(profile("openai", "b"))); + let inner = crate::host::PrebuiltModelResolver::new( + a, + [("b".to_owned(), b.clone() as Arc<dyn ChatModel<()>>)], + ) + .with_primary_name("a"); + let named_calls = Arc::new(AtomicUsize::new(0)); + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness.with_model_resolver(Arc::new(CountingResolver { + inner, + named_calls: named_calls.clone(), + })); + + let handle = SteeringHandle::allow_all_with_model_switch(); + handle.send(switch("b")); + let recorder = EventRecorder::new(); + let run = harness + .invoke_in_context(&(), context(&handle, &recorder), vec![Message::user("hi")]) + .await + .expect("the switched model runs"); + + assert_eq!(run.text(), Some("from b".to_owned())); + assert_eq!(b.requests().len(), 1); + assert_eq!( + named_calls.load(Ordering::SeqCst), + 0, + "the eligible binding resolved for this request is reused by both switch checks" + ); +} + /// Answers from a script and, on its first call, queues `command` on the /// run's steering handle: an orchestrator reacting mid-run. struct SteersOnFirstCall { @@ -213,8 +285,12 @@ async fn mid_run_switch_applies_to_the_next_call_and_stays_sticky() { ModelResponse::assistant("done"), ])); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("a", a); - harness.register_model("b", b.clone()); + let resolver = crate::host::PrebuiltModelResolver::new( + a.clone(), + [("b".to_string(), b.clone() as Arc<dyn ChatModel<()>>)], + ) + .with_primary_name("a"); + harness.with_model_resolver(Arc::new(resolver)); harness.register_tool(Arc::new(FakeTool::returning("lookup", "ok"))); let recorder = EventRecorder::new(); @@ -449,8 +525,14 @@ fn failover_harness( models: Vec<(&str, Arc<dyn ChatModel<()>>)>, ) -> AgentHarness<()> { let mut harness: AgentHarness<()> = AgentHarness::new(); - for (name, model) in models { - harness.register_model(name, model); + let mut models = models.into_iter(); + if let Some((primary_name, primary)) = models.next() { + let resolver = crate::host::PrebuiltModelResolver::new( + primary, + models.map(|(name, model)| (name.to_string(), model)), + ) + .with_primary_name(primary_name); + harness.with_model_resolver(Arc::new(resolver)); } harness.with_policy(RunPolicy { retry: RetryPolicy::default() @@ -655,11 +737,20 @@ async fn rejection_after_middleware_restores_request_model_and_reports_once() { #[tokio::test] async fn switch_wins_over_a_model_chosen_by_before_model_middleware() { - let (mut harness, _, b) = two_models( - ScriptedModel::replies(vec!["from a"]), - ScriptedModel::replies(vec!["from b"]), - ); - harness.register_model("c", Arc::new(ScriptedModel::replies(vec!["from c"]))); + let a = Arc::new(ScriptedModel::replies(vec!["from a"])); + let b = Arc::new(ScriptedModel::replies(vec!["from b"])); + let c = Arc::new(ScriptedModel::replies(vec!["from c"])); + let resolver = crate::host::PrebuiltModelResolver::new( + a.clone(), + [ + ("b".to_string(), b.clone() as Arc<dyn ChatModel<()>>), + ("c".to_string(), c.clone() as Arc<dyn ChatModel<()>>), + ], + ) + .with_primary_name("a"); + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness.with_model_resolver(Arc::new(resolver)); + harness.register_tool(Arc::new(FakeTool::returning("lookup", "ok"))); harness.push_middleware(Arc::new(Tweak { model: Some("c"), tool_calling: false, @@ -676,6 +767,10 @@ async fn switch_wins_over_a_model_chosen_by_before_model_middleware() { assert_eq!(run.text(), Some("from b".to_string())); assert_eq!(model_started(&recorder), vec!["b"]); assert_eq!(b.requests().len(), 1); + assert!( + c.requests().is_empty(), + "the switch wins over the available middleware route" + ); } #[tokio::test] @@ -687,7 +782,15 @@ async fn rejected_switch_keeps_the_model_a_middleware_picked() { let c = Arc::new( ScriptedModel::replies(vec!["from c"]).with_profile(tool_capable("openai", "gpt-5-mini")), ); - harness.register_model("c", c.clone()); + let resolver = crate::host::PrebuiltModelResolver::new( + a.clone(), + [ + ("b".to_string(), b.clone() as Arc<dyn ChatModel<()>>), + ("c".to_string(), c.clone() as Arc<dyn ChatModel<()>>), + ], + ) + .with_primary_name("a"); + harness.with_model_resolver(Arc::new(resolver)); harness.push_middleware(Arc::new(Tweak { model: Some("c"), tool_calling: true, diff --git a/crates/tinyagents-harness/src/agent_loop/nested_tests.rs b/crates/tinyagents-harness/src/agent_loop/nested_tests.rs index 37e62ce22..3120df518 100644 --- a/crates/tinyagents-harness/src/agent_loop/nested_tests.rs +++ b/crates/tinyagents-harness/src/agent_loop/nested_tests.rs @@ -237,13 +237,13 @@ fn parent_call(id: &str, tool: &str) -> ToolCall { /// A harness whose model makes `calls` in one turn, then says "done". fn harness_with(calls: Vec<ToolCall>, limits: RunLimits) -> AgentHarness<()> { let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "mock", Arc::new(MockModel::with_responses(vec![ response(calls), text_response("done"), ])), - ); + ))); harness.with_policy(RunPolicy { limits, ..RunPolicy::default() diff --git a/crates/tinyagents-harness/src/agent_loop/phases.rs b/crates/tinyagents-harness/src/agent_loop/phases.rs index 9cb7e53c2..049e89ba8 100644 --- a/crates/tinyagents-harness/src/agent_loop/phases.rs +++ b/crates/tinyagents-harness/src/agent_loop/phases.rs @@ -1,50 +1,4 @@ -//! 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}; +//! Reusable agent-loop execution helpers. use crate::context::RunContext; use crate::error::Result; @@ -52,140 +6,19 @@ 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<StructuredPlan>, -} - -/// 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)] +/// The messages and tool names produced by executing one batch of calls. +#[derive(Clone, Debug)] pub struct ToolBatchOutcome { - /// The tool-result messages to append to the transcript, in call order. + /// Tool-result messages appended in call order. pub results: Vec<Message>, - /// Names of calls that reached a tool executor, in execution order (see - /// [`AgentRun::executed_tools`]). + /// Calls that reached an executor, in execution order. pub executed_tools: Vec<String>, } -/// 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<serde_json::Value>, - /// Which schema variant matched, under - /// [`crate::structured::StructuredStrategy::ToolCallUnion`]. `None` for - /// every other strategy. - pub structured_variant: Option<String>, -} - -/// 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<State: Send + Sync, Ctx: Send + Sync>: 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, Ctx>, - state: &State, - ctx: &mut RunContext<Ctx>, - run: &mut AgentRun, - status: &mut HarnessRunStatus, - input: Vec<Message>, - 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_with_promotions`][super::AgentHarness::execute_tools_with_promotions] -/// 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_with_promotions` 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. +/// Executes a batch with the same admission, middleware, and result-folding +/// pipeline used by the standard agent loop. pub async fn execute_tool_batch<State: Send + Sync, Ctx: Send + Sync>( harness: &AgentHarness<State, Ctx>, state: &State, @@ -213,68 +46,3 @@ pub async fn execute_tool_batch<State: Send + Sync, Ctx: Send + Sync>( executed_tools: run.executed_tools[executed_before..].to_vec(), }) } - -// ── Lifecycle events for alternate drivers ───────────────────────────────── - -/// Fixes the lifecycle seed for a driver that enters a run mid-flight (a node -/// activation of the compiled-graph loop, possibly resumed from a checkpoint): -/// the first `transcript_len` messages are input and are never announced. -/// A no-op once a seed exists, so it is safe to call on every node entry. -pub fn lifecycle_seed<Ctx>(ctx: &mut RunContext<Ctx>, transcript_len: usize) { - ctx.ensure_turn_tracker_seeded(transcript_len); -} - -/// Announces transcript appends not yet announced (`MessageAppended`), using -/// the harness's payload-capture policy. Mirrors the direct loop's flush -/// points; tool calls nested inside a tool never reach the transcript, so they -/// are never announced. -pub fn lifecycle_flush<State: Send + Sync, Ctx: Send + Sync>( - harness: &AgentHarness<State, Ctx>, - ctx: &mut RunContext<Ctx>, - messages: &[Message], -) { - ctx.flush_transcript(harness.policy().capture, messages); -} - -/// Opens the next turn (`TurnStarted`), first announcing pending appends and -/// closing any turn still open. Call right before the model call. -pub fn lifecycle_start_turn<State: Send + Sync, Ctx: Send + Sync>( - harness: &AgentHarness<State, Ctx>, - ctx: &mut RunContext<Ctx>, - messages: &[Message], -) -> u32 { - ctx.start_turn(harness.policy().capture, messages) -} - -/// Announces pending appends and closes the open turn (`TurnCompleted`), if -/// any. Idempotent: closing with no open turn only flushes. -pub fn lifecycle_close_turn<State: Send + Sync, Ctx: Send + Sync>( - harness: &AgentHarness<State, Ctx>, - ctx: &mut RunContext<Ctx>, - messages: &[Message], -) { - ctx.close_turn(harness.policy().capture, messages); -} - -/// Reports that the transcript was truncated to `new_len` messages, emitting -/// `MessageRetracted` for each announced message removed (highest index first). -/// A driver whose node discards its state (an interrupt re-runs the node from -/// its entry state) calls this so a mirror stays consistent with the transcript -/// that will actually be kept. -pub fn lifecycle_retract<Ctx>(ctx: &mut RunContext<Ctx>, new_len: usize) { - ctx.retract_transcript(new_len); -} - -/// Re-aligns the lifecycle tracker with a run entered mid-flight (resumed from a -/// checkpoint in a fresh runtime, whose tracker starts at zero): turn numbering -/// continues from `completed_turns`, and when `open_turn_start` is `Some(index)` -/// and no turn is open, the in-flight turn that began at that transcript index -/// is re-opened (not announced again) so the next close reports it. Never moves -/// numbering backwards, so it is safe to call on every node entry. -pub fn lifecycle_resume<Ctx>( - ctx: &mut RunContext<Ctx>, - completed_turns: u32, - open_turn_start: Option<usize>, -) { - ctx.adopt_turn_state(completed_turns, open_turn_start); -} diff --git a/crates/tinyagents-harness/src/agent_loop/rich_tool_tests.rs b/crates/tinyagents-harness/src/agent_loop/rich_tool_tests.rs index 9f28147af..d81766119 100644 --- a/crates/tinyagents-harness/src/agent_loop/rich_tool_tests.rs +++ b/crates/tinyagents-harness/src/agent_loop/rich_tool_tests.rs @@ -167,7 +167,10 @@ struct Fixture { fn fixture(responses: Vec<ModelResponse>) -> Fixture { let model = Arc::new(ScriptedModel::new(responses)); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", Arc::clone(&model) as _); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + Arc::clone(&model) as _, + ))); Fixture { harness, model, diff --git a/crates/tinyagents-harness/src/agent_loop/run_depth_tests.rs b/crates/tinyagents-harness/src/agent_loop/run_depth_tests.rs index f347f1d00..173965af0 100644 --- a/crates/tinyagents-harness/src/agent_loop/run_depth_tests.rs +++ b/crates/tinyagents-harness/src/agent_loop/run_depth_tests.rs @@ -31,7 +31,10 @@ impl Middleware<(), ()> for RecordDepth { async fn driver_scopes_the_real_configured_depth_for_unary_and_streaming_runs() { let recorded = Arc::new(std::sync::Mutex::new(Vec::new())); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", Arc::new(MockModel::constant("done"))); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + Arc::new(MockModel::constant("done")), + ))); harness.push_middleware(Arc::new(RecordDepth(recorded.clone()))); let incoming = RunDepth { depth: 1, diff --git a/crates/tinyagents-harness/src/agent_loop/run_loop.rs b/crates/tinyagents-harness/src/agent_loop/run_loop.rs index 3a53971b6..261d4640e 100644 --- a/crates/tinyagents-harness/src/agent_loop/run_loop.rs +++ b/crates/tinyagents-harness/src/agent_loop/run_loop.rs @@ -47,8 +47,13 @@ impl<State: Send + Sync, Ctx: Send + Sync> AgentHarness<State, Ctx> { // 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); + run.collected.extend( + queue + .take_admitted(crate::run_queue::QueueLane::Collect, self.policy.queue_mode) + .await + .into_iter() + .map(|item| item.value.clone()), + ); } let exit = match outcome { @@ -374,7 +379,7 @@ impl<State: Send + Sync, Ctx: Send + Sync> AgentHarness<State, Ctx> { Option<String>, Vec<tinyinference_llm::model::ModelHint>, Option<tinyinference_llm::model::CapabilitySet>, - crate::model_registry::ResolvedModelBinding<State>, + crate::agent_loop::ModelBinding<State>, )>, )); // A pending steered model switch decides the model this call @@ -395,7 +400,20 @@ impl<State: Send + Sync, Ctx: Send + Sync> AgentHarness<State, Ctx> { )); profile } else { - self.preview_model_profile(&patch_request) + match self.resolve_local_model(ctx, &patch_request).await { + Ok(binding) => { + let profile = binding.model.profile().cloned(); + *resolution_cache.lock().unwrap() = Some(( + patch_request.model.clone(), + patch_request.model_hints.clone(), + patch_request.required_capabilities.clone(), + binding, + )); + profile + } + Err(TinyAgentsError::ModelNotFound(_)) => None, + Err(error) => return Err(error), + } }; // B6: re-consult the toolset chain and declare any live change as a @@ -429,7 +447,24 @@ impl<State: Send + Sync, Ctx: Send + Sync> AgentHarness<State, Ctx> { // handoff/budget/dialect decisions below all see the new model. // Re-applied after `before_model` (see below). let model_before_switch = request.model.clone(); - self.apply_steered_model_switch(ctx, &mut request, &model_before_switch); + let switch_candidate = self.steered_model(ctx).and_then(|requested| { + resolution_cache + .lock() + .unwrap() + .as_ref() + .and_then(|(_, _, _, binding)| { + (binding.resolved.name == requested + && binding.resolved.requested.as_deref() == Some(requested.as_str())) + .then(|| binding.model.clone()) + }) + }); + self.apply_steered_model_switch( + ctx, + &mut request, + &model_before_switch, + switch_candidate, + ) + .await?; // Known tool requirements must shape the hosted profile seen by // middleware. The later gate below still catches tools added by @@ -483,7 +518,16 @@ impl<State: Send + Sync, Ctx: Send + Sync> AgentHarness<State, Ctx> { Some((key.0, key.1, key.2, binding)); Ok(profile) } else { - Ok(harness.preview_model_profile(request)) + match harness.resolve_local_model(ctx, request).await { + Ok(binding) => { + let profile = binding.model.profile().cloned(); + *resolution_cache.lock().unwrap() = + Some((key.0, key.1, key.2, binding)); + Ok(profile) + } + Err(TinyAgentsError::ModelNotFound(_)) => Ok(None), + Err(error) => Err(error), + } } }) }, @@ -510,7 +554,24 @@ impl<State: Send + Sync, Ctx: Send + Sync> AgentHarness<State, Ctx> { // Middleware may have chosen another model or added capability // requirements; the steered switch is re-validated and wins. - self.apply_steered_model_switch(ctx, &mut request, &model_before_switch); + let switch_candidate = self.steered_model(ctx).and_then(|requested| { + resolution_cache + .lock() + .unwrap() + .as_ref() + .and_then(|(_, _, _, binding)| { + (binding.resolved.name == requested + && binding.resolved.requested.as_deref() == Some(requested.as_str())) + .then(|| binding.model.clone()) + }) + }); + self.apply_steered_model_switch( + ctx, + &mut request, + &model_before_switch, + switch_candidate, + ) + .await?; // Safe checkpoint: a control requested from `before_model_control` // (for example `BudgetMiddleware` finding the budget already @@ -553,16 +614,7 @@ impl<State: Send + Sync, Ctx: Send + Sync> AgentHarness<State, Ctx> { let binding = if let Some(binding) = hosted_binding { binding } else { - self.models - .resolve_request(&request, None, None) - .ok_or_else(|| { - TinyAgentsError::ModelNotFound( - request - .model - .clone() - .unwrap_or_else(|| "<default>".to_string()), - ) - })? + self.resolve_local_model(ctx, &request).await? }; ctx.model_profile = binding.model.profile().cloned(); crate::middleware::library::rehome_ephemeral_system_instructions( diff --git a/crates/tinyagents-harness/src/agent_loop/run_loop_dead_call_tests.rs b/crates/tinyagents-harness/src/agent_loop/run_loop_dead_call_tests.rs index 02defe23f..99a8661fe 100644 --- a/crates/tinyagents-harness/src/agent_loop/run_loop_dead_call_tests.rs +++ b/crates/tinyagents-harness/src/agent_loop/run_loop_dead_call_tests.rs @@ -27,7 +27,10 @@ async fn efforts(reasoning_off: bool) -> Vec<Option<ReasoningEffort>> { crate::testkit::text_response("done").with_usage(Usage::new(4, 3)), ])); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", Arc::clone(&model) as _); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + Arc::clone(&model) as _, + ))); harness.with_policy(RunPolicy { default_reasoning: Some(ReasoningConfig::effort(ReasoningEffort::High)), truncated_empty_reasoning_off: reasoning_off, @@ -86,7 +89,10 @@ async fn efforts_across_steps(sticky_after: u32) -> Vec<Option<ReasoningEffort>> crate::testkit::text_response("done").with_usage(Usage::new(4, 3)), ])); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", Arc::clone(&model) as _); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + Arc::clone(&model) as _, + ))); harness.register_tool(Arc::new(crate::testkit::FakeTool::new("probe"))); harness.with_policy(RunPolicy { default_reasoning: Some(ReasoningConfig::effort(ReasoningEffort::High)), @@ -132,7 +138,10 @@ async fn constant_cap_retries_reach_the_provider_instead_of_replaying_cached_bla crate::testkit::text_response("recovered"), ])); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", model.clone()); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + Arc::clone(&model) as _, + ))); harness.with_response_cache(Arc::new(crate::cache::InMemoryResponseCache::new())); harness.with_policy(RunPolicy { truncated_empty_constant_cap: true, @@ -162,7 +171,10 @@ async fn truncated_empty_recovery_bypasses_a_blank_cached_before_retries_were_en let input = vec![Message::user("same question")]; let seed_model = Arc::new(ScriptedModel::new(vec![dead(2048)])); let mut seed: AgentHarness<()> = AgentHarness::new(); - seed.register_model("mock", seed_model.clone()); + seed.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + Arc::clone(&seed_model) as _, + ))); seed.with_response_cache(cache.clone()); seed.with_policy(RunPolicy { truncated_empty_retries: 0, @@ -185,7 +197,10 @@ async fn truncated_empty_recovery_bypasses_a_blank_cached_before_retries_were_en "fresh answer", )])); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", model.clone()); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + Arc::clone(&model) as _, + ))); harness.with_response_cache(cache); harness.with_policy(RunPolicy { truncated_empty_constant_cap: true, @@ -246,7 +261,10 @@ async fn only_the_tool_calls_of_a_dead_call_retry_are_marked_as_such() { ])); let flags = Arc::new(RetryFlags::default()); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", Arc::clone(&model) as _); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + Arc::clone(&model) as _, + ))); harness.register_tool(Arc::new(crate::testkit::FakeTool::new("probe"))); harness.push_middleware(flags.clone()); harness.with_policy(RunPolicy { diff --git a/crates/tinyagents-harness/src/agent_loop/run_loop_recovery_tests.rs b/crates/tinyagents-harness/src/agent_loop/run_loop_recovery_tests.rs index 84397ec47..9fe80a9a3 100644 --- a/crates/tinyagents-harness/src/agent_loop/run_loop_recovery_tests.rs +++ b/crates/tinyagents-harness/src/agent_loop/run_loop_recovery_tests.rs @@ -183,7 +183,10 @@ async fn text_dialect_recovery_policy_on_stream_and_unary() { }); } let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", Arc::new(model)); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + Arc::new(model), + ))); harness.register_tool(tool.clone()); harness.with_policy(RunPolicy { text_dialect_recovery: policy, diff --git a/crates/tinyagents-harness/src/agent_loop/run_loop_withheld_tests.rs b/crates/tinyagents-harness/src/agent_loop/run_loop_withheld_tests.rs index 49b93850f..ff372adbd 100644 --- a/crates/tinyagents-harness/src/agent_loop/run_loop_withheld_tests.rs +++ b/crates/tinyagents-harness/src/agent_loop/run_loop_withheld_tests.rs @@ -73,7 +73,9 @@ fn harness_with( max_calls: usize, ) -> AgentHarness<()> { let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", model); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", model, + ))); harness.with_policy(RunPolicy { limits: RunLimits::default().with_max_model_calls(max_calls), ..RunPolicy::default() diff --git a/crates/tinyagents-harness/src/agent_loop/run_queue_tests.rs b/crates/tinyagents-harness/src/agent_loop/run_queue_tests.rs index b0103a949..cbc8b6dec 100644 --- a/crates/tinyagents-harness/src/agent_loop/run_queue_tests.rs +++ b/crates/tinyagents-harness/src/agent_loop/run_queue_tests.rs @@ -157,7 +157,10 @@ struct Fixture { fn fixture(responses: Vec<ModelResponse>, 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_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + Arc::clone(&model) as _, + ))); harness.with_policy(RunPolicy { queue_mode: mode, ..RunPolicy::default() @@ -399,6 +402,49 @@ async fn collect_lane_lands_on_the_run_and_never_reaches_the_model() { assert_eq!(fx.queue.status().await.collects, 0); } +#[tokio::test] +async fn collect_lane_rejects_revoked_guarded_messages_before_returning_the_run() { + use crate::run_queue::{QueueAdmission, QueueAdmissionError, QueueAdmissionPermit}; + use std::sync::atomic::{AtomicBool, Ordering}; + + #[derive(Debug)] + struct Revocable(AtomicBool); + impl QueueAdmission for Revocable { + fn acquire(&self) -> Result<Arc<dyn QueueAdmissionPermit>, QueueAdmissionError> { + self.0 + .load(Ordering::SeqCst) + .then(|| Arc::new(()) as Arc<dyn QueueAdmissionPermit>) + .ok_or(QueueAdmissionError::Rejected) + } + } + + let fx = fixture(vec![final_turn("done")], QueueMode::All); + let admission = Arc::new(Revocable(AtomicBool::new(true))); + fx.queue + .push_guarded( + QueueLane::Collect, + Message::user("revoked observation"), + admission.clone(), + ) + .await; + fx.queue + .push(QueueLane::Collect, Message::user("ordinary observation")) + .await; + admission.0.store(false, Ordering::SeqCst); + + let run = fx + .harness + .invoke_in_context(&(), fx.ctx("collect-revoked"), vec![Message::user("go")]) + .await + .expect("run succeeds"); + + assert_eq!( + run.collected.iter().map(Message::text).collect::<Vec<_>>(), + vec!["ordinary observation"] + ); + assert_eq!(fx.queue.status().await.collects, 0); +} + // ── Boundary semantics ────────────────────────────────────────────────────── #[tokio::test] diff --git a/crates/tinyagents-harness/src/agent_loop/stream_idle_timeout_tests.rs b/crates/tinyagents-harness/src/agent_loop/stream_idle_timeout_tests.rs index 48d189e77..a1b56ccf8 100644 --- a/crates/tinyagents-harness/src/agent_loop/stream_idle_timeout_tests.rs +++ b/crates/tinyagents-harness/src/agent_loop/stream_idle_timeout_tests.rs @@ -143,7 +143,9 @@ fn harness_with( attempts: usize, ) -> AgentHarness<()> { let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("primary", primary); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "primary", primary, + ))); harness.with_policy(RunPolicy { limits, retry: RetryPolicy::default() @@ -163,12 +165,18 @@ fn harness_with_chain( attempts: usize, ) -> AgentHarness<()> { let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("primary", primary); - let mut models = vec!["primary".to_string()]; - for (name, model) in others { - harness.register_model(name, model); - models.push(name.to_string()); - } + let models = std::iter::once("primary".to_string()) + .chain(others.iter().map(|(name, _)| (*name).to_string())) + .collect::<Vec<_>>(); + harness.with_model_resolver(Arc::new( + crate::host::PrebuiltModelResolver::new( + primary, + others + .into_iter() + .map(|(name, model)| (name.to_string(), model)), + ) + .with_primary_name("primary"), + )); harness.with_policy(RunPolicy { limits, retry: RetryPolicy::default() @@ -812,10 +820,10 @@ async fn non_streaming_calls_ignore_the_idle_timeout() { // A buffered call has no inter-event gaps to measure: a 300s call under a // 1s idle timeout must still complete. let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "slow", Arc::new(SlowModel::new(Duration::from_secs(300), "done")), - ); + ))); harness.with_policy(RunPolicy { limits: RunLimits::default() .with_stream_idle_timeout_ms(Some(1_000)) diff --git a/crates/tinyagents-harness/src/agent_loop/terminal_outcome_tests.rs b/crates/tinyagents-harness/src/agent_loop/terminal_outcome_tests.rs index a4f1ff37d..f380a6450 100644 --- a/crates/tinyagents-harness/src/agent_loop/terminal_outcome_tests.rs +++ b/crates/tinyagents-harness/src/agent_loop/terminal_outcome_tests.rs @@ -64,7 +64,9 @@ impl ChatModel<()> for PendingModel { fn harness_with(model: Arc<dyn ChatModel<()>>) -> AgentHarness<()> { let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", model); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", model, + ))); harness } diff --git a/crates/tinyagents-harness/src/agent_loop/tool_rules_tests.rs b/crates/tinyagents-harness/src/agent_loop/tool_rules_tests.rs index 492538300..82ac6e125 100644 --- a/crates/tinyagents-harness/src/agent_loop/tool_rules_tests.rs +++ b/crates/tinyagents-harness/src/agent_loop/tool_rules_tests.rs @@ -127,7 +127,7 @@ fn rules(value: Value) -> ToolRules { fn harness_with(model: Arc<ScriptedModel>, policy: ToolRulePolicy) -> AgentHarness<()> { let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("scripted", model as _); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::new(model as _))); harness.with_policy(RunPolicy { tool_rules: policy, ..RunPolicy::default() @@ -271,14 +271,9 @@ async fn indirect_target_resolution_obeys_run_cancellation() { let cancellation = crate::cancel::CancellationToken::new(); let ctx = RunContext::new(RunConfig::new("cancel-target-resolution"), ()) .with_cancellation(cancellation.clone()); - harness.register_model( - "scripted", - Arc::new(ScriptedModel::new(vec![calls(vec![( - "target", - "blocked_target", - json!({}), - )])])) as _, - ); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::new(Arc::new( + ScriptedModel::new(vec![calls(vec![("target", "blocked_target", json!({}))])]), + ) as _))); let mut run = Box::pin(harness.invoke_in_context(&(), ctx, vec![Message::user("go")])); tokio::select! { @@ -311,14 +306,9 @@ async fn indirect_target_resolution_obeys_run_deadline() { ] }))), ..RunPolicy::default() }); - harness.register_model( - "scripted", - Arc::new(ScriptedModel::new(vec![calls(vec![( - "target", - "blocked_target", - json!({}), - )])])) as _, - ); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::new(Arc::new( + ScriptedModel::new(vec![calls(vec![("target", "blocked_target", json!({}))])]), + ) as _))); let ctx = RunContext::new( RunConfig::new("deadline-target-resolution").with_timeout_ms(500), (), @@ -471,7 +461,7 @@ async fn normalized_arguments_are_checked_against_the_rules_again() { ] }))); let execute = RuleTool::dispatcher("execute"); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("scripted", model as _); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::new(model as _))); harness.with_policy(RunPolicy { invalid_args: crate::runtime::InvalidArgsPolicy::NormalizeThenReturnToolError, tool_rules: policy, @@ -615,7 +605,7 @@ async fn a_rewrite_target_carries_its_own_approval_rule() { ModelResponse::assistant("never reached"), ])); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("scripted", model as _); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::new(model as _))); harness.with_policy(RunPolicy { unknown_tool: UnknownToolPolicy::Rewrite { tool_name: "send".to_string(), diff --git a/crates/tinyagents-harness/src/agent_loop/tool_timing_tests.rs b/crates/tinyagents-harness/src/agent_loop/tool_timing_tests.rs index fc46c060c..b88c24965 100644 --- a/crates/tinyagents-harness/src/agent_loop/tool_timing_tests.rs +++ b/crates/tinyagents-harness/src/agent_loop/tool_timing_tests.rs @@ -84,7 +84,10 @@ fn tool_call_then_text() -> MockModel { async fn tool_row_text(policy: RunPolicy) -> String { let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", Arc::new(tool_call_then_text())); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + Arc::new(tool_call_then_text()), + ))); harness.register_tool(Arc::new(EchoTool)); harness.with_policy(policy); let run = harness diff --git a/crates/tinyagents-harness/src/agent_loop/tools_args_recovery_tests.rs b/crates/tinyagents-harness/src/agent_loop/tools_args_recovery_tests.rs index b9d6dd3d3..65ebe05fb 100644 --- a/crates/tinyagents-harness/src/agent_loop/tools_args_recovery_tests.rs +++ b/crates/tinyagents-harness/src/agent_loop/tools_args_recovery_tests.rs @@ -64,7 +64,9 @@ async fn run_one(name: &str, arguments: Value) -> Outcome { text_response("done"), ])); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", model as _); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", model as _, + ))); harness.register_tool(Arc::new(SkillTool { received: Arc::clone(&received), })); diff --git a/crates/tinyagents-harness/src/agent_loop/tools_progress_tests.rs b/crates/tinyagents-harness/src/agent_loop/tools_progress_tests.rs index 08bc908be..20dfaca05 100644 --- a/crates/tinyagents-harness/src/agent_loop/tools_progress_tests.rs +++ b/crates/tinyagents-harness/src/agent_loop/tools_progress_tests.rs @@ -112,7 +112,9 @@ struct Fixture { fn fixture(first_turn: ModelResponse) -> Fixture { let model = Arc::new(ScriptedModel::new(vec![first_turn, response(vec![], "ok")])); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", model as _); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", model as _, + ))); Fixture { harness, recorder: EventRecorder::new(), diff --git a/crates/tinyagents-harness/src/agent_loop/truncated_tool_call_tests.rs b/crates/tinyagents-harness/src/agent_loop/truncated_tool_call_tests.rs index ac8a844c7..ebfee4264 100644 --- a/crates/tinyagents-harness/src/agent_loop/truncated_tool_call_tests.rs +++ b/crates/tinyagents-harness/src/agent_loop/truncated_tool_call_tests.rs @@ -30,7 +30,10 @@ fn harness(responses: Vec<ModelResponse>, retries: u32) -> (AgentHarness<()>, Ar let other = Arc::new(FakeTool::returning("write_file", "written")); let mut harness: AgentHarness<()> = AgentHarness::new(); harness - .register_model("mock", Arc::new(ScriptedModel::new(responses)) as _) + .with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + Arc::new(ScriptedModel::new(responses)) as _, + ))) .register_tool(Arc::clone(&tool) as _) .register_tool(other as _); harness.with_policy(RunPolicy { diff --git a/crates/tinyagents-harness/src/agent_loop/turn_control.rs b/crates/tinyagents-harness/src/agent_loop/turn_control.rs index 9803f8079..af48f0efe 100644 --- a/crates/tinyagents-harness/src/agent_loop/turn_control.rs +++ b/crates/tinyagents-harness/src/agent_loop/turn_control.rs @@ -24,7 +24,7 @@ impl<State: Send + Sync, Ctx: Send + Sync> AgentHarness<State, Ctx> { let Some(queue) = ctx.run_queue.clone() else { return false; }; - let items = queue.take(lane, self.policy.queue_mode).await; + let items = queue.take_admitted(lane, self.policy.queue_mode).await; if items.is_empty() { return false; } @@ -38,12 +38,13 @@ impl<State: Send + Sync, Ctx: Send + Sync> AgentHarness<State, Ctx> { Message::Tool(_) => self.policy.capture.tool_io, _ => self.policy.capture.model_io, }; - let captured: Vec<serde_json::Value> = if items.iter().any(is_captured) { + let captured: Vec<serde_json::Value> = if items.iter().any(|item| is_captured(&item.value)) + { items .iter() - .map(|message| { - if is_captured(message) { - super::lifecycle::to_value_logged(message) + .map(|item| { + if is_captured(&item.value) { + super::lifecycle::to_value_logged(&item.value) } else { serde_json::Value::Null } @@ -52,7 +53,7 @@ impl<State: Send + Sync, Ctx: Send + Sync> AgentHarness<State, Ctx> { } else { Vec::new() }; - messages.extend(items); + messages.extend(items.iter().map(|item| item.value.clone())); let record = ctx.emit(AgentEvent::QueuedMessageApplied { lane, count, diff --git a/crates/tinyagents-harness/src/agent_loop/types.rs b/crates/tinyagents-harness/src/agent_loop/types.rs index cc8b91d77..e92c79788 100644 --- a/crates/tinyagents-harness/src/agent_loop/types.rs +++ b/crates/tinyagents-harness/src/agent_loop/types.rs @@ -14,9 +14,18 @@ use crate::middleware::AgentRun; use crate::steering::PauseState; use std::collections::{BTreeMap, BTreeSet}; +use std::sync::Arc; +use tinyinference_llm::model::{ChatModel, ResolvedModel}; use super::*; +/// A resolved executable model and the metadata produced for this call. +/// Private to the loop; selection authority remains with `ModelResolver`. +pub(crate) struct ModelBinding<State: Send + Sync> { + pub(crate) resolved: ResolvedModel, + pub(crate) model: Arc<dyn ChatModel<State>>, +} + /// The result of an agent-loop invocation that keeps the partial run even when /// the loop fails. /// diff --git a/crates/tinyagents-harness/src/agent_loop/wrap_concurrency_tests.rs b/crates/tinyagents-harness/src/agent_loop/wrap_concurrency_tests.rs index d8dc6ee5b..99083efb7 100644 --- a/crates/tinyagents-harness/src/agent_loop/wrap_concurrency_tests.rs +++ b/crates/tinyagents-harness/src/agent_loop/wrap_concurrency_tests.rs @@ -165,13 +165,13 @@ struct Rig { /// Three sleeping tools (alpha slowest, gamma fastest) behind one probe wrap. fn rig(mode: Mode, concurrent_safe: bool) -> Rig { let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "mock", Arc::new(MockModel::with_responses(vec![ batch_response(), crate::testkit::text_response("done").with_usage(Usage::new(4, 2)), ])), - ); + ))); let ran = Arc::new(Mutex::new(Vec::new())); let active = Arc::new(std::sync::atomic::AtomicUsize::new(0)); let max_seen = Arc::new(std::sync::atomic::AtomicUsize::new(0)); diff --git a/crates/tinyagents-harness/src/capability/mod.rs b/crates/tinyagents-harness/src/capability/mod.rs deleted file mode 100644 index 0893e423d..000000000 --- a/crates/tinyagents-harness/src/capability/mod.rs +++ /dev/null @@ -1,190 +0,0 @@ -//! 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<State, Ctx>` composes [`crate::tool::toolset::ToolSet<State, -//! Ctx>`] and [`crate::middleware::Middleware<State, Ctx>`] 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<State>` 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<dyn Any>` 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)] -#[path = "mod_tests.rs"] -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<State: Send + Sync, Ctx: Send + Sync> Capability<State, Ctx> { - /// 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<String>) -> 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<String>) -> 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<dyn ToolSet<State, Ctx>>) -> 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<dyn Middleware<State, Ctx>>) -> 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 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<Self> { - 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/mod_tests.rs b/crates/tinyagents-harness/src/capability/mod_tests.rs deleted file mode 100644 index 8a1ad4eaf..000000000 --- a/crates/tinyagents-harness/src/capability/mod_tests.rs +++ /dev/null @@ -1,265 +0,0 @@ -//! 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<String>) -> 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<ToolResult> { - Ok(ToolResult::success(format!("{}-result", self.name))) - } -} - -/// A single-tool `ToolSet` wrapping one [`StubTool`]. -struct StubToolSet { - tool: Arc<dyn Tool>, -} - -impl StubToolSet { - fn new(name: impl Into<String>) -> Self { - Self { - tool: Arc::new(StubTool::new(name)), - } - } -} - -#[async_trait] -impl ToolSet<(), ()> for StubToolSet { - async fn tools(&self, _ctx: &RunContext<()>) -> Result<Vec<Arc<dyn Tool>>> { - Ok(vec![self.tool.clone()]) - } - - async fn call( - &self, - name: &str, - args: serde_json::Value, - _ctx: &RunContext<()>, - ) -> Result<ToolResult> { - 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<String> = 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 deleted file mode 100644 index 90702ebba..000000000 --- a/crates/tinyagents-harness/src/capability/types.rs +++ /dev/null @@ -1,450 +0,0 @@ -//! 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<ResponseFormat>, - /// Overrides [`crate::runtime::RunPolicy::fallback`] when set. - pub fallback: Option<FallbackPolicy>, -} - -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: the unit an `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<State, Ctx>`] and -/// [`crate::tool::toolset::ToolSet<State, Ctx>`] — 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<dyn Any>` for exactly that reason). -pub struct Capability<State: Send + Sync, Ctx: Send + Sync> { - /// 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<String>, - /// The tools this capability contributes, if any. - pub toolset: Option<Arc<dyn ToolSet<State, Ctx>>>, - /// Middleware appended to the harness's stack when this capability is - /// installed. - pub middleware: Vec<Arc<dyn Middleware<State, Ctx>>>, - /// Model-request defaults applied to the harness's [`crate::runtime::RunPolicy`] - /// when this capability is installed. - pub model_defaults: Option<ModelRequestDefaults>, - /// 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<State: Send + Sync, Ctx: Send + Sync> Clone for Capability<State, Ctx> { - /// 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<State: Send + Sync, Ctx: Send + Sync> std::fmt::Debug for Capability<State, Ctx> { - /// 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<String>, - #[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<ModelDefaultsSpec>, -} - -/// 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<ExposureSpec> for ToolExposure { - fn from(value: ExposureSpec) -> Self { - match value { - ExposureSpec::Direct => ToolExposure::Direct, - ExposureSpec::Deferred => ToolExposure::Deferred, - ExposureSpec::Hidden => ToolExposure::Hidden, - } - } -} - -impl From<ToolExposure> 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<ResponseFormat>, - #[serde(default, skip_serializing_if = "Vec::is_empty")] - pub(crate) fallback_models: Vec<String>, -} - -/// [`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<RwLock<HashSet<String>>>`) 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<State: Send + Sync, Ctx: Send + Sync> { - pub(crate) capabilities: Vec<Capability<State, Ctx>>, - pub(crate) loaded: Arc<RwLock<HashSet<String>>>, - pub(crate) load_tool: Option<Arc<LoadCapabilityTool>>, -} - -impl<State: Send + Sync, Ctx: Send + Sync> CapabilityToolSet<State, Ctx> { - /// 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<Capability<State, Ctx>>) -> 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<String> { - let mut names: Vec<String> = self - .loaded - .read() - .expect("capability load state lock poisoned") - .iter() - .cloned() - .collect(); - names.sort(); - names - } - - fn is_loaded(&self, capability: &Capability<State, Ctx>, loaded: &HashSet<String>) -> 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<Arc<dyn Tool>> { - self.load_tool - .as_ref() - .map(|tool| tool.clone() as Arc<dyn Tool>) - } -} - -#[async_trait] -impl<State: Send + Sync, Ctx: Send + Sync> ToolSet<State, Ctx> for CapabilityToolSet<State, Ctx> { - async fn tools(&self, ctx: &RunContext<Ctx>) -> Result<Vec<Arc<dyn Tool>>> { - 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<dyn Tool>, - ); - } - } - } - if let Some(load_tool) = &self.load_tool { - out.push(load_tool.clone() as Arc<dyn Tool>); - } - Ok(out) - } - - async fn call(&self, name: &str, args: Value, ctx: &RunContext<Ctx>) -> Result<ToolResult> { - 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<String> { - 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<Ctx>) -> 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<String>, - pub(crate) loaded: Arc<RwLock<HashSet<String>>>, -} - -impl LoadCapabilityTool { - pub(crate) fn new(available: Vec<String>, loaded: Arc<RwLock<HashSet<String>>>) -> 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<ToolResult> { - 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/close.rs b/crates/tinyagents-harness/src/close.rs index 73f29c2d2..37e8ac04b 100644 --- a/crates/tinyagents-harness/src/close.rs +++ b/crates/tinyagents-harness/src/close.rs @@ -9,13 +9,380 @@ //! observe completed responses while preserving the legacy entry points and //! configuration structs. +use futures::StreamExt; use tinyinference_llm::message::Message; use tinyinference_llm::model::{ - ChatModel, ModelRequest, ModelResponse, ModelStreamItem, ToolChoice, + ChatModel, ModelRequest, ModelResponse, ModelStream, ModelStreamItem, ToolChoice, }; use tinyinference_llm::usage::Usage; use tinytools_agent::dialect::{DialectResponse, ToolDialect}; +/// Host copy used by deterministic close and repair paths. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct CloseCopy { + pub checkpoint_template: &'static str, + pub final_summary_template: &'static str, + pub empty_checkpoint_results: &'static str, + pub empty_final_summary: &'static str, + pub repair_template: &'static str, + pub max_iteration_instruction: &'static str, + pub final_write_instruction: &'static str, + pub final_answer_instruction: &'static str, + pub truncated_answer_instruction: &'static str, + pub stop_note_preamble: &'static str, + pub no_tool_calls: &'static str, + pub repair_no_reply: &'static str, + pub repair_quoted_instruction: &'static str, + pub repair_unverified: &'static str, + pub instruction_needles: &'static [&'static str], + pub minimum_quoted_words: usize, + pub verification_prompt_template: &'static str, +} + +const ENGLISH_FINAL_ANSWER_NEEDLE: &str = "You have finished using tools for this turn but have not yet written a reply to the user. Tools are no longer available and nothing more will run this turn, so do not call any tools and do not describe steps you are about to take. Write a self-contained final message that reports what actually happened: what you found, changed or established, grounded in the tool results above and the tool records below. If the request was not completed, say so and give the reason from the failing tool's own error message, keeping any link it includes. Do not state anything the tool records contradict. If nothing conclusive resulted, say so plainly.\n\nThese directions are addressed to you and are not part of the conversation. Do not quote or restate them, in whole or in part. Do not list or describe the tools available to you. Do not narrate your deliberation: no thinking aloud, no correcting yourself mid-reply, no weighing what to do. Write only the message the user will read."; +const ENGLISH_TRUNCATED_ANSWER_NEEDLE: &str = "Your last reply ran out of output tokens while you were still reasoning, so this turn ended before you made your next tool call or wrote a reply to the user. The work is unfinished: say so plainly and name what is still left to do. Tools are no longer available and nothing more will run this turn, so do not call any tools and do not describe steps you are about to take. Write a self-contained final message that reports what actually happened: what you found, changed or established, grounded in the tool results above and the tool records below. If the request was not completed, say so and give the reason from the failing tool's own error message, keeping any link it includes. Do not state anything the tool records contradict. If nothing conclusive resulted, say so plainly.\n\nThese directions are addressed to you and are not part of the conversation. Do not quote or restate them, in whole or in part. Do not list or describe the tools available to you. Do not narrate your deliberation: no thinking aloud, no correcting yourself mid-reply, no weighing what to do. Write only the message the user will read."; +const ENGLISH_ITERATION_NEEDLE: &str = "You have reached the maximum number of tool calls allowed for this single turn, so you cannot call any more tools right now. Do not attempt another tool call.\n\nFirst, report the substance of what this turn produced, grounded in the tool results above: what you found, what you changed, or what you established — whichever this task was. Be concrete about it. If you were gathering information, give the information itself rather than saying that you gathered it; if you were changing something, say what is now different. Include your most recent tool calls: they completed, and this is your only opportunity to use them.\n\nThen close with a brief **Still to do** line naming what remains, so the user can ask you to carry on.\n\nLet the length follow the substance — neither a one-line status nor a raw dump of tool output. If the results support no conclusion yet, say that plainly and name what is missing."; +const ENGLISH_FINAL_WRITE_NEEDLE: &str = "This is the last call on which you can use a tool, and the only tools left are the ones that write files. Gathering is over — anything you have not found by now will not be found in this turn.\n\nIf this task asked you to produce a file (a report, a guide, a document) and you have not written it yet, write it now with file_write, from what is already in the results above, and mark inside it any part you could not confirm.\n\nIf this task is a code change, apply as much of the change as you can right now as real edits to the source files: apply_patch for targeted edits, or file_write to rewrite a file whose full content you have. A partial set of real edits is worth more than any description of them. Do not write notes, a plan, a summary, a TODO list or any other document into the project in place of the change. If you cannot edit a file exactly, leave it untouched and name it in your reply.\n\nIf there is nothing to write — the task asked only for an answer, or the work is already written — then do not call a tool. Answer instead, and you will be asked to conclude next."; + +/// English copy preserving OpenHuman's close and checkpoint wording. +pub const ENGLISH_CLOSE_COPY: CloseCopy = CloseCopy { + checkpoint_template: "I reached the tool-call limit for this turn ({steps} steps), so I paused here before I could write up what I found.\n\n**Raw results from this turn**\n{results}\n**Next steps:** I'll continue from here — just reply (e.g. \"continue\") and I'll pick up where I left off.", + final_summary_template: "I finished this turn without writing up a result. Here is what each tool call returned:\n{results}\nTell me how you'd like to proceed.", + empty_checkpoint_results: "\n- (no tools completed yet)\n", + empty_final_summary: "I finished this turn but produced no result to report.", + repair_template: "{repair_reason} Write the message again.\n\n{instruction}", + max_iteration_instruction: "You have reached the maximum number of tool calls allowed for this single turn, so you cannot call any more tools right now. Do not attempt another tool call.\n\nFirst, report the substance of what this turn produced, grounded in the tool results above: what you found, what you changed, or what you established — whichever this task was. Be concrete about it. If you were gathering information, give the information itself rather than saying that you gathered it; if you were changing something, say what is now different. Include your most recent tool calls: they completed, and this is your only opportunity to use them.\n\nThen close with a brief **Still to do** line naming what remains, so the user can ask you to carry on.\n\nLet the length follow the substance — neither a one-line status nor a raw dump of tool output. If the results support no conclusion yet, say that plainly and name what is missing.", + final_write_instruction: "This is the last call on which you can use a tool, and the only tools left are the ones that write files. Gathering is over — anything you have not found by now will not be found in this turn.\n\nIf this task asked you to produce a file (a report, a guide, a document) and you have not written it yet, write it now with file_write, from what is already in the results above, and mark inside it any part you could not confirm.\n\nIf this task is a code change, apply as much of the change as you can right now as real edits to the source files: apply_patch for targeted edits, or file_write to rewrite a file whose full content you have. A partial set of real edits is worth more than any description of them. Do not write notes, a plan, a summary, a TODO list or any other document into the project in place of the change. If you cannot edit a file exactly, leave it untouched and name it in your reply.\n\nIf there is nothing to write — the task asked only for an answer, or the work is already written — then do not call a tool. Answer instead, and you will be asked to conclude next.", + final_answer_instruction: "You have finished using tools for this turn but have not yet written a reply to the user. Tools are no longer available and nothing more will run this turn, so do not call any tools and do not describe steps you are about to take. Write a self-contained final message that reports what actually happened: what you found, changed or established, grounded in the tool results above and the tool records below. If the request was not completed, say so and give the reason from the failing tool's own error message, keeping any link it includes. Do not state anything the tool records contradict. If nothing conclusive resulted, say so plainly.\n\nThese directions are addressed to you and are not part of the conversation. Do not quote or restate them, in whole or in part. Do not list or describe the tools available to you. Do not narrate your deliberation: no thinking aloud, no correcting yourself mid-reply, no weighing what to do. Write only the message the user will read.", + truncated_answer_instruction: "Your last reply ran out of output tokens while you were still reasoning, so this turn ended before you made your next tool call or wrote a reply to the user. The work is unfinished: say so plainly and name what is still left to do. Tools are no longer available and nothing more will run this turn, so do not call any tools and do not describe steps you are about to take. Write a self-contained final message that reports what actually happened: what you found, changed or established, grounded in the tool results above and the tool records below. If the request was not completed, say so and give the reason from the failing tool's own error message, keeping any link it includes. Do not state anything the tool records contradict. If nothing conclusive resulted, say so plainly.\n\nThese directions are addressed to you and are not part of the conversation. Do not quote or restate them, in whole or in part. Do not list or describe the tools available to you. Do not narrate your deliberation: no thinking aloud, no correcting yourself mid-reply, no weighing what to do. Write only the message the user will read.", + stop_note_preamble: "The harness stopped this turn early because its tool calls stopped making progress: they kept failing, or kept repeating the same step. Its stop note is written for you, not for the user, so explain it in your own words rather than repeating it, and describe each call as the tool records show it:", + no_tool_calls: "(no tool calls completed)", + repair_no_reply: "Your previous reply was empty or tried to call a tool. No tools will run: write the message itself.", + repair_quoted_instruction: "Your previous reply repeated these directions, or the stop note, back to the user. That text is addressed to you alone and must never appear in the message. Do not reproduce any part of it, do not refer to it, and do not explain that you were given it.", + repair_unverified: "Your previous reply did not pass the check that it reports what actually happened, grounded in the records below.", + instruction_needles: &[ + ENGLISH_FINAL_ANSWER_NEEDLE, + ENGLISH_TRUNCATED_ANSWER_NEEDLE, + ENGLISH_ITERATION_NEEDLE, + ENGLISH_FINAL_WRITE_NEEDLE, + "The harness stopped this turn early because its tool calls stopped making progress: they kept failing, or kept repeating the same step. Its stop note is written for you, not for the user, so explain it in your own words rather than repeating it, and describe each call as the tool records show it:", + ], + minimum_quoted_words: 10, + verification_prompt_template: "You are checking a reply before it is shown to a user. Below are the user's request, the records of the tool calls made while handling it, and the reply.\n\nAnswer REJECT if any of these is true:\n1. The reply only says what the assistant will do or is about to do, instead of reporting what happened.\n2. The reply states something the tool records contradict, for example that something does not exist or did not work when a record shows it succeeded, or that something succeeded when its record shows it failed.\n3. The request was not completed, a failed record gives the reason, and the reply does not pass that reason on.\n4. The reply narrates the assistant's own deliberation — thinking aloud, correcting itself part way through, or weighing what to do — rather than stating the outcome. Reporting what a tool call returned is not deliberation.\n5. The reply recites or enumerates the tools the assistant has, instead of reporting what this turn's tool calls produced. Naming a tool the reply actually used, or answering a request that asked what the assistant can do, is not a violation.\n\nOtherwise answer ACCEPT. Reply with the single word ACCEPT or REJECT.\n\n<user_request>\n{user_request}\n</user_request>\n\n<tool_records>\n{records}\n</tool_records>\n\n<reply>\n{reply}\n</reply>", +}; + +/// Substitute only tokens present in the template; inserted content stays literal. +pub(crate) fn interpolate_once(template: &str, fields: &[(&str, &str)]) -> String { + let mut output = String::with_capacity(template.len()); + let mut remaining = template; + while let Some((index, key, value)) = fields + .iter() + .filter_map(|(key, value)| remaining.find(key).map(|index| (index, *key, *value))) + .min_by_key(|(index, _, _)| *index) + { + output.push_str(&remaining[..index]); + output.push_str(value); + remaining = &remaining[index + key.len()..]; + } + output.push_str(remaining); + output +} + +/// Build the checkpoint prose around already-rendered results. +pub fn checkpoint_with_copy( + copy: &CloseCopy, + steps: usize, + rendered_results: &str, + has_results: bool, +) -> String { + interpolate_once( + copy.checkpoint_template, + &[ + ("{steps}", &steps.to_string()), + ( + "{results}", + if has_results { + rendered_results + } else { + copy.empty_checkpoint_results + }, + ), + ], + ) +} + +/// Build the fallback close prose around already-rendered results. +pub fn final_summary_with_copy( + copy: &CloseCopy, + rendered_results: &str, + has_results: bool, +) -> String { + if !has_results { + return copy.empty_final_summary.to_owned(); + } + interpolate_once( + copy.final_summary_template, + &[("{results}", rendered_results)], + ) +} + +/// Build the evidence-grounded closing instruction using host copy. +pub fn final_answer_instruction( + copy: &CloseCopy, + stop_reason: Option<&str>, + truncated: bool, + records: &str, +) -> String { + use crate::middleware::wrap_harness_instruction; + let mut directive = String::new(); + if let Some(reason) = stop_reason { + directive.push_str(copy.stop_note_preamble); + directive.push_str("\n<stop_note>\n"); + directive.push_str(reason.trim()); + directive.push_str("\n</stop_note>\n\n"); + } + directive.push_str(if truncated && stop_reason.is_none() { + copy.truncated_answer_instruction + } else { + copy.final_answer_instruction + }); + let mut out = wrap_harness_instruction(&directive); + out.push_str("\n\n<tool_records>\n"); + out.push_str(if records.trim().is_empty() { + copy.no_tool_calls + } else { + records.trim() + }); + out.push_str("\n</tool_records>"); + out +} + +/// Build the verifier prompt without rescanning inserted user or tool content. +pub fn verification_prompt( + copy: &CloseCopy, + request: &str, + records: &str, + reply: &str, + request_char_cap: usize, +) -> String { + let records = if records.trim().is_empty() { + copy.no_tool_calls + } else { + records.trim() + }; + interpolate_once( + copy.verification_prompt_template, + &[ + ("{user_request}", &truncate_chars(request, request_char_cap)), + ("{records}", records), + ("{reply}", reply.trim()), + ], + ) +} + +/// Derive normalized quote-detection phrases from a copy bundle. +pub fn instruction_needles(copy: &CloseCopy, stop_reason: Option<&str>) -> Vec<String> { + let mut sources = copy.instruction_needles.to_vec(); + if let Some(reason) = stop_reason { + sources.push(reason); + } + harness_instruction_needles(&sources, copy.minimum_quoted_words) +} + +/// Build the single corrective prompt for a rejected close response. +pub fn repair_with_copy(copy: &CloseCopy, instruction: &str, violation: CloseViolation) -> String { + let reason = match violation { + CloseViolation::NoReply => copy.repair_no_reply, + CloseViolation::QuotedInstruction => copy.repair_quoted_instruction, + CloseViolation::Unverified => copy.repair_unverified, + }; + interpolate_once( + copy.repair_template, + &[("{repair_reason}", reason), ("{instruction}", instruction)], + ) +} + +/// Host-provided model builder for evidence-grounded close calls. +pub trait CloseModelSource: Send + Sync { + /// Build the selected model for one close or repair request. + fn build_summarizer( + &self, + model: &str, + temperature: f64, + thread_id: Option<&str>, + ) -> Result<std::sync::Arc<dyn ChatModel<()>>, String>; +} + +/// Close model that resolves its host model source for each request. +pub struct CloseSourceModel<'a> { + pub source: &'a dyn CloseModelSource, + pub model: &'a str, + pub temperature: f64, + pub thread_id: Option<&'a str>, + pub trace: Option<&'a std::sync::Arc<crate::observability::SupplementalCallTrace>>, +} + +#[async_trait::async_trait] +impl ChatModel<()> for CloseSourceModel<'_> { + async fn invoke( + &self, + state: &(), + mut request: ModelRequest, + ) -> tinyinference_llm::Result<ModelResponse> { + let model = self + .source + .build_summarizer(self.model, self.temperature, self.thread_id) + .map_err(tinyinference_llm::Error::Model)?; + let mut call = self.trace.map(|trace| { + crate::observability::SupplementalCallGuard::start(trace, self.model, &mut request) + }); + match model.invoke(state, request).await { + Ok(response) => { + if let Some(call) = &mut call { + call.complete(&response); + } + Ok(response) + } + Err(error) => Err(error), + } + } + + async fn stream( + &self, + state: &(), + mut request: ModelRequest, + ) -> tinyinference_llm::Result<ModelStream> { + let model = self + .source + .build_summarizer(self.model, self.temperature, self.thread_id) + .map_err(tinyinference_llm::Error::Model)?; + let call = self.trace.map(|trace| { + crate::observability::SupplementalCallGuard::start(trace, self.model, &mut request) + }); + let correlation = request.correlation.clone(); + let stream = model.stream(state, request).await?; + let mut metadata = stream.metadata().clone(); + if metadata.correlation.is_none() { + metadata.correlation = correlation; + } + let traced = futures::stream::unfold( + (stream, call, false), + |(mut stream, mut call, terminal)| async move { + if terminal { + return None; + } + match stream.next().await { + Some(item) => { + let terminal = matches!( + &item, + ModelStreamItem::Completed(_) + | ModelStreamItem::Failed(_) + | ModelStreamItem::ProviderFailed(_) + ); + match &item { + ModelStreamItem::Completed(response) => { + if let Some(call) = &mut call { + call.complete(response); + } + } + ModelStreamItem::Failed(_) | ModelStreamItem::ProviderFailed(_) => { + call.take(); + } + _ => {} + } + Some((item, (stream, call, terminal))) + } + None => { + // EOF is terminal even when the consumer retains the + // stream after observing `None`. + call.take(); + None + } + } + }, + ); + Ok(ModelStream::new(Box::pin(traced)).with_metadata(metadata)) + } +} + +/// Usage and context-window accounting for supplemental close requests. +#[derive(Debug, Clone, Default, PartialEq)] +pub struct CloseUsage { + pub model_calls: usize, + pub input_tokens: u64, + pub output_tokens: u64, + pub cached_input_tokens: u64, + pub last_call_input_tokens: u64, + pub last_call_output_tokens: u64, + pub cost: crate::cost::CostTally, + pub pricing_model: String, +} + +impl CloseUsage { + /// Empty accounting for supplemental calls made on `model`. + pub fn for_model(model: &str) -> Self { + Self { + pricing_model: model.to_owned(), + ..Self::default() + } + } + + /// Record one attempted request and its provider usage when available. + pub fn record( + &mut self, + _kind: CloseRequestKind, + usage: Option<&Usage>, + measures_context: bool, + ) { + self.model_calls += 1; + if let Some(usage) = usage { + if measures_context { + self.last_call_input_tokens = usage.input_tokens; + self.last_call_output_tokens = usage.output_tokens; + } + self.input_tokens += usage.input_tokens; + self.output_tokens += usage.output_tokens; + self.cached_input_tokens += usage.cache_read_tokens; + } + } + + /// Record a close response, excluding verifier calls from final context. + pub fn record_close_call(&mut self, kind: CloseRequestKind, response: Option<&ModelResponse>) { + self.record( + kind, + response + .and_then(|response| response.usage.or(response.message.usage)) + .as_ref(), + kind != CloseRequestKind::Verification, + ); + } + + /// Record a close request and let the host attach pricing from its billing + /// and catalog adapters without importing host types into the harness. + pub fn record_close_call_with_cost( + &mut self, + kind: CloseRequestKind, + response: Option<&ModelResponse>, + cost: impl FnOnce(Option<&ModelResponse>) -> Option<(String, crate::cost::CallCost)>, + ) { + self.record_close_call(kind, response); + if let Some((pricing_model, call_cost)) = cost(response) { + self.pricing_model = pricing_model; + self.cost.add(call_cost); + } + } +} + +/// Invoke a close builder only when the supplied stop reason is classified. +pub fn classified_halt_close<T>( + stop_reason: Option<&str>, + is_classified_halt: impl FnOnce(&str) -> bool, + build_close: impl FnOnce() -> T, +) -> Option<T> { + let reason = stop_reason?; + is_classified_halt(reason).then(build_close) +} + /// One completed tool call captured by the host before transcript compaction. #[derive(Debug, Clone, PartialEq, Eq)] pub struct CloseToolResult { @@ -239,9 +606,9 @@ pub fn deterministic_checkpoint( results: &[CloseToolResult], total_budget: usize, ) -> String { - template - .replace("{steps}", &steps.to_string()) - .replace("{results}", &render_tool_results(results, total_budget)) + let steps = steps.to_string(); + let rendered = render_tool_results(results, total_budget); + interpolate_once(template, &[("{steps}", &steps), ("{results}", &rendered)]) } /// Build the deterministic summary used when a normal close is rejected or @@ -256,9 +623,14 @@ pub fn deterministic_final_summary( stop_reason: Option<&str>, total_budget: usize, ) -> String { - template - .replace("{stop_reason}", stop_reason.unwrap_or("")) - .replace("{results}", &render_tool_results(results, total_budget)) + let rendered = render_tool_results(results, total_budget); + interpolate_once( + template, + &[ + ("{stop_reason}", stop_reason.unwrap_or("")), + ("{results}", &rendered), + ], + ) } /// Whether a candidate repeats any harness-only phrase supplied by the host. diff --git a/crates/tinyagents-harness/src/close_tests.rs b/crates/tinyagents-harness/src/close_tests.rs index 11bdea72f..b0cbaf89b 100644 --- a/crates/tinyagents-harness/src/close_tests.rs +++ b/crates/tinyagents-harness/src/close_tests.rs @@ -1,5 +1,491 @@ use super::*; +#[test] +fn deterministic_templates_do_not_rescan_inserted_placeholder_text() { + let checkpoint = deterministic_checkpoint( + "{steps}|{results}", + 3, + &[CloseToolResult { + name: "tool_{steps}".to_owned(), + success: true, + content: "result {results}".to_owned(), + }], + 100, + ); + assert_eq!(checkpoint, "3|- `tool_{steps}` (ok): result {results}"); + + let summary = deterministic_final_summary( + "{stop_reason}|{results}", + &[CloseToolResult { + name: "tool_{results}".to_owned(), + success: false, + content: "failed {stop_reason}".to_owned(), + }], + Some("halt {results}"), + 100, + ); + assert_eq!( + summary, + "halt {results}|- `tool_{results}` (failed): failed {stop_reason}" + ); +} + +#[test] +fn close_copy_prompts_preserve_placeholder_shaped_content() { + let copy = CloseCopy { + verification_prompt_template: "request={user_request};records={records};reply={reply}", + repair_template: "{repair_reason}|{instruction}", + ..ENGLISH_CLOSE_COPY + }; + assert_eq!( + verification_prompt( + ©, + "Explain {records}", + "record {reply}", + "reply {user_request}", + 200 + ), + "request=Explain {records};records=record {reply};reply=reply {user_request}" + ); + assert_eq!( + repair_with_copy(©, "try {repair_reason}", CloseViolation::NoReply), + "Your previous reply was empty or tried to call a tool. No tools will run: write the message itself.|try {repair_reason}" + ); +} + +#[test] +fn close_copy_renderers_preserve_custom_literals_and_all_repair_reasons() { + let copy = CloseCopy { + checkpoint_template: "checkpoint {steps}: {results}", + final_summary_template: "summary: {results}", + empty_checkpoint_results: "empty checkpoint", + empty_final_summary: "empty summary", + repair_template: "reason={repair_reason};instruction={instruction}", + final_answer_instruction: "answer copy", + truncated_answer_instruction: "truncated copy", + stop_note_preamble: "stop preamble", + no_tool_calls: "no calls", + repair_no_reply: "repair empty", + repair_quoted_instruction: "repair quote", + repair_unverified: "repair verification", + verification_prompt_template: "request={user_request}|records={records}|reply={reply}", + instruction_needles: &["custom needle with enough words to match"], + minimum_quoted_words: 2, + ..ENGLISH_CLOSE_COPY + }; + assert_eq!( + checkpoint_with_copy(©, 4, "body", true), + "checkpoint 4: body" + ); + assert_eq!( + checkpoint_with_copy(©, 4, "", false), + "checkpoint 4: empty checkpoint" + ); + assert_eq!( + final_summary_with_copy(©, "body", true), + "summary: body" + ); + assert_eq!(final_summary_with_copy(©, "", false), "empty summary"); + assert_eq!( + repair_with_copy(©, "write {repair_reason}", CloseViolation::NoReply), + "reason=repair empty;instruction=write {repair_reason}" + ); + assert_eq!( + repair_with_copy(©, "instruction", CloseViolation::QuotedInstruction), + "reason=repair quote;instruction=instruction" + ); + assert_eq!( + repair_with_copy(©, "instruction", CloseViolation::Unverified), + "reason=repair verification;instruction=instruction" + ); + let answer = final_answer_instruction(©, Some("halt"), false, "records"); + assert!(answer.contains("stop preamble\n<stop_note>\nhalt\n</stop_note>")); + assert!(answer.contains("answer copy")); + let truncated = final_answer_instruction(©, None, true, " "); + assert!(truncated.contains("truncated copy")); + assert!(truncated.contains("<tool_records>\nno calls\n</tool_records>")); + assert_eq!( + instruction_needles(©, None), + ["custom needle with enough words to match"] + ); +} + +#[test] +fn verification_prompt_does_not_rescan_placeholder_shaped_content() { + let copy = CloseCopy { + verification_prompt_template: "request=<{user_request}> records=<{records}> reply=<{reply}>", + ..ENGLISH_CLOSE_COPY + }; + assert_eq!( + verification_prompt( + ©, + "Explain {records}", + "record contains {reply} and {user_request}", + "reply contains {records}", + 200, + ), + "request=<Explain {records}> records=<record contains {reply} and {user_request}> reply=<reply contains {records}>" + ); +} + +#[test] +fn final_summary_copy_keeps_each_result_status_and_body_byte_for_byte() { + let rendered = "\n- `install_item` — failed\n > no direct download\n"; + assert_eq!( + final_summary_with_copy(&ENGLISH_CLOSE_COPY, rendered, true), + "I finished this turn without writing up a result. Here is what each tool call returned:\n\n- `install_item` — failed\n > no direct download\n\nTell me how you'd like to proceed." + ); +} + +#[test] +fn english_close_copy_frames_records_and_keeps_write_and_truncation_rules() { + let records = "\n- `read_file` — ok\n > configuration found\n"; + let instruction = final_answer_instruction(&ENGLISH_CLOSE_COPY, None, false, records); + let frame_end = instruction + .find("</harness_instruction>") + .expect("directive is framed"); + let records_at = instruction + .find("<tool_records>") + .expect("records included"); + assert!(frame_end < records_at); + assert!(instruction.contains("Do not quote or restate them")); + assert!(instruction.contains("Do not list or describe the tools available to you")); + assert!(instruction.contains("Do not narrate your deliberation")); + assert!(instruction.contains(" > configuration found")); + + let halted = final_answer_instruction(&ENGLISH_CLOSE_COPY, Some("stop note"), true, records); + assert!(halted.contains("<stop_note>\nstop note\n</stop_note>")); + assert!(halted.contains(ENGLISH_CLOSE_COPY.final_answer_instruction)); + let truncated = final_answer_instruction(&ENGLISH_CLOSE_COPY, None, true, records); + assert!(truncated.contains("ran out of output tokens")); + assert!(!truncated.contains("You have finished using tools")); + + let write = ENGLISH_CLOSE_COPY.final_write_instruction.to_lowercase(); + assert!(write.contains("apply_patch") && write.contains("file_write")); + assert!(write.contains("source files") && write.contains("partial set of real edits")); + assert!(write.contains("asked you to produce a file")); + for forbidden in ["notes", "plan", "summary"] { + assert!(write.contains(forbidden)); + } +} + +#[test] +fn english_close_guard_rejects_the_instruction_and_stop_note_it_received() { + let stop_note = "The stop note says to stop now before the work is complete."; + let needles = instruction_needles(&ENGLISH_CLOSE_COPY, Some(stop_note)); + let needle_refs = needles.iter().map(String::as_str).collect::<Vec<_>>(); + let directive = ENGLISH_CLOSE_COPY + .final_answer_instruction + .split('.') + .next() + .unwrap(); + assert!(quotes_harness_instruction_with_minimum( + &format!("{directive}. Here is the result."), + &needle_refs, + 1 + )); + assert!(quotes_harness_instruction_with_minimum( + stop_note, + &needle_refs, + 1 + )); + assert!(!quotes_harness_instruction_with_minimum( + "I used `read_file` and `write_file` to update the config.", + &needle_refs, + 1 + )); +} + +#[test] +fn english_verification_prompt_uses_no_tool_fallback_for_empty_records() { + let expected = "You are checking a reply before it is shown to a user. Below are the user's request, the records of the tool calls made while handling it, and the reply.\n\nAnswer REJECT if any of these is true:\n1. The reply only says what the assistant will do or is about to do, instead of reporting what happened.\n2. The reply states something the tool records contradict, for example that something does not exist or did not work when a record shows it succeeded, or that something succeeded when its record shows it failed.\n3. The request was not completed, a failed record gives the reason, and the reply does not pass that reason on.\n4. The reply narrates the assistant's own deliberation — thinking aloud, correcting itself part way through, or weighing what to do — rather than stating the outcome. Reporting what a tool call returned is not deliberation.\n5. The reply recites or enumerates the tools the assistant has, instead of reporting what this turn's tool calls produced. Naming a tool the reply actually used, or answering a request that asked what the assistant can do, is not a violation.\n\nOtherwise answer ACCEPT. Reply with the single word ACCEPT or REJECT.\n\n<user_request>\nstatus?\n</user_request>\n\n<tool_records>\n(no tool calls completed)\n</tool_records>\n\n<reply>\nNothing ran.\n</reply>"; + assert_eq!( + verification_prompt(&ENGLISH_CLOSE_COPY, "status?", "", "Nothing ran.", 200), + expected + ); + assert_eq!( + verification_prompt( + &ENGLISH_CLOSE_COPY, + "status?", + " \n\t ", + "Nothing ran.", + 200 + ), + expected + ); +} + +#[test] +fn close_copy_builders_match_english_checkpoint_and_empty_summary() { + assert_eq!( + checkpoint_with_copy( + &ENGLISH_CLOSE_COPY, + 2, + "- `read_file` — ok\n > data\n", + true + ), + "I reached the tool-call limit for this turn (2 steps), so I paused here before I could write up what I found.\n\n**Raw results from this turn**\n- `read_file` — ok\n > data\n\n**Next steps:** I'll continue from here — just reply (e.g. \"continue\") and I'll pick up where I left off." + ); + assert_eq!( + final_summary_with_copy(&ENGLISH_CLOSE_COPY, "", false), + ENGLISH_CLOSE_COPY.empty_final_summary + ); + assert_eq!( + checkpoint_with_copy(&ENGLISH_CLOSE_COPY, 25, "", false), + "I reached the tool-call limit for this turn (25 steps), so I paused here before I could write up what I found.\n\n**Raw results from this turn**\n\n- (no tools completed yet)\n\n**Next steps:** I'll continue from here — just reply (e.g. \"continue\") and I'll pick up where I left off." + ); + let needles = instruction_needles(&ENGLISH_CLOSE_COPY, None); + let needle_refs = needles.iter().map(String::as_str).collect::<Vec<_>>(); + assert!(quotes_harness_instruction_with_minimum( + ENGLISH_CLOSE_COPY.final_write_instruction, + &needle_refs, + 1 + )); +} + +#[test] +fn result_renderer_accepts_the_host_block_format_and_charges_rendered_bytes() { + let results = [ + CloseToolResult { + name: "read_file".into(), + success: true, + content: "first line\nsecond line".into(), + }, + CloseToolResult { + name: "write_file".into(), + success: false, + content: "permission denied".into(), + }, + ]; + let rendered = render_tool_results_with_options( + &results, + CloseResultRenderOptions { + per_result_cap: 100, + truncation_marker: "…", + total_budget: 200, + render_block: |result: &CloseToolResult| { + let status = if result.success { "ok" } else { "failed" }; + let body = result + .content + .lines() + .map(|line| format!(" > {line}\n")) + .collect::<String>(); + format!("- `{}` — {status}\n{body}", result.name) + }, + omitted_notice: |count| { + if count == 0 { + String::new() + } else { + format!( + "\n_({count} earlier tool result(s) omitted for length — the most recent are shown.)_\n" + ) + } + }, + }, + ); + assert_eq!( + rendered, + "- `read_file` — ok\n > first line\n > second line\n- `write_file` — failed\n > permission denied\n" + ); + assert_eq!( + checkpoint_with_copy(&ENGLISH_CLOSE_COPY, 3, &rendered, true,), + "I reached the tool-call limit for this turn (3 steps), so I paused here before I could write up what I found.\n\n**Raw results from this turn**\n- `read_file` — ok\n > first line\n > second line\n- `write_file` — failed\n > permission denied\n\n**Next steps:** I'll continue from here — just reply (e.g. \"continue\") and I'll pick up where I left off." + ); +} + +#[test] +fn host_block_golden_keeps_its_leading_newline_and_multiline_output() { + let rendered = render_tool_results_with_options( + &[CloseToolResult { + name: "read_file".into(), + success: true, + content: "first line\nsecond line".into(), + }], + CloseResultRenderOptions { + per_result_cap: 800, + truncation_marker: "… [truncated]", + total_budget: 4_000, + render_block: |result: &CloseToolResult| { + let status = if result.success { "ok" } else { "failed" }; + let mut block = format!("\n- `{}` — {status}\n", result.name); + for line in result.content.lines() { + block.push_str(" > "); + block.push_str(line); + block.push('\n'); + } + block + }, + omitted_notice: |_| String::new(), + }, + ); + assert_eq!( + rendered, + "\n- `read_file` — ok\n > first line\n > second line\n" + ); +} + +#[test] +fn newline_heavy_results_are_budgeted_after_host_block_rendering() { + let body = (0..200).map(|_| "y").collect::<Vec<_>>().join("\n"); + let results = (0..20) + .map(|index| CloseToolResult { + name: format!("list_issues_{index}"), + success: true, + content: body.clone(), + }) + .collect::<Vec<_>>(); + let rendered = render_tool_results_with_options( + &results, + CloseResultRenderOptions { + per_result_cap: 800, + truncation_marker: "… [truncated]", + total_budget: 4_000, + render_block: |result: &CloseToolResult| { + let mut block = format!("\n- `{}` — ok\n", result.name); + for line in result.content.lines() { + block.push_str(" > "); + block.push_str(line); + block.push('\n'); + } + block + }, + omitted_notice: |count| { + if count == 0 { + String::new() + } else { + format!( + "\n_({count} earlier tool result(s) omitted for length — the most recent are shown.)_\n" + ) + } + }, + }, + ); + assert!(rendered.chars().count() <= 4_000 + 200); + assert!(rendered.contains("earlier tool result(s) omitted for length")); + assert!(rendered.contains("`list_issues_19` — ok")); + assert!(rendered.contains(" > y\n")); +} + +#[test] +fn close_usage_counts_attempts_and_excludes_verifier_from_context_tokens() { + let mut usage = CloseUsage::for_model("close-model"); + let generation = Usage { + input_tokens: 11, + output_tokens: 4, + cache_read_tokens: 3, + ..Default::default() + }; + usage.record(CloseRequestKind::Generation, Some(&generation), true); + let verification = Usage { + input_tokens: 50, + output_tokens: 2, + cache_read_tokens: 0, + ..Default::default() + }; + usage.record(CloseRequestKind::Verification, Some(&verification), false); + usage.record(CloseRequestKind::RequiredOutputRepair, None, true); + assert_eq!(usage.model_calls, 3); + assert_eq!( + ( + usage.input_tokens, + usage.output_tokens, + usage.cached_input_tokens + ), + (61, 6, 3) + ); + assert_eq!( + (usage.last_call_input_tokens, usage.last_call_output_tokens), + (11, 4) + ); + assert_eq!(usage.pricing_model, "close-model"); +} + +#[test] +fn close_usage_accepts_host_billing_without_host_types() { + let mut usage = CloseUsage::for_model("selected-model"); + let mut response = ModelResponse::assistant("closed"); + response.usage = Some(Usage { + input_tokens: 17, + output_tokens: 5, + ..Default::default() + }); + + usage.record_close_call_with_cost(CloseRequestKind::Generation, Some(&response), |response| { + assert!(response.is_some()); + Some(( + "actual-route-model".into(), + crate::cost::CallCost::Charged(0.007), + )) + }); + usage.record_close_call_with_cost(CloseRequestKind::RequiredOutputRepair, None, |response| { + assert!(response.is_none()); + None + }); + + assert_eq!(usage.model_calls, 2); + assert_eq!((usage.input_tokens, usage.output_tokens), (17, 5)); + assert_eq!(usage.pricing_model, "actual-route-model"); + assert_eq!(usage.cost.usd(), Some(0.007)); +} + +#[test] +fn classified_halt_close_builds_only_for_classified_reason() { + let mut built = false; + let result = classified_halt_close( + Some("halt"), + |reason| reason == "halt", + || { + built = true; + 42 + }, + ); + assert_eq!(result, Some(42)); + assert!(built); + assert_eq!( + classified_halt_close::<()>(Some("ordinary"), |_| false, || ()), + None + ); + assert_eq!( + classified_halt_close::<()>(None, |_| panic!("no reason"), || panic!("no close")), + None + ); +} + +#[test] +fn classified_halt_renders_partial_work_without_provider_usage() { + let results = [ + CloseToolResult { + name: "list_issues".into(), + success: true, + content: "private success payload".into(), + }, + CloseToolResult { + name: "fetch_pr".into(), + success: false, + content: "permission denied".into(), + }, + ]; + let close = classified_halt_close( + Some("Stopping after 2 attempts: classified failure"), + |reason| reason.starts_with("Stopping after "), + || { + let groups = crate::turn_stop::group_stop_results(&results, str::to_owned); + let output = crate::turn_stop::render_stop_summary_with_copy( + "a step was blocked", + &groups, + &crate::turn_stop::ENGLISH_STOP_COPY, + ); + (output, CloseUsage::default()) + }, + ) + .unwrap(); + assert!(close.0.contains("`list_issues` succeeded")); + assert!(!close.0.contains("private success payload")); + assert!(close.0.contains("`fetch_pr` failed: permission denied")); + assert_eq!(close.1.model_calls, 0); + assert_eq!(close.1.cost.usd(), Some(0.0)); +} + use std::collections::VecDeque; use std::sync::{Arc, Mutex}; @@ -80,6 +566,369 @@ fn scripted_model(answers: &[&str]) -> ScriptedModel { } } +struct StaticCloseSource(Arc<dyn ChatModel<()>>); + +impl CloseModelSource for StaticCloseSource { + fn build_summarizer( + &self, + _model: &str, + _temperature: f64, + _thread_id: Option<&str>, + ) -> Result<Arc<dyn ChatModel<()>>, String> { + Ok(Arc::clone(&self.0)) + } +} + +struct RecordingCloseSource { + model: Arc<dyn ChatModel<()>>, + builds: Mutex<Vec<(String, f64, Option<String>)>>, +} + +impl CloseModelSource for RecordingCloseSource { + fn build_summarizer( + &self, + model: &str, + temperature: f64, + thread_id: Option<&str>, + ) -> Result<Arc<dyn ChatModel<()>>, String> { + self.builds.lock().unwrap().push(( + model.to_owned(), + temperature, + thread_id.map(str::to_owned), + )); + Ok(Arc::clone(&self.model)) + } +} + +struct EmptyCloseEnricher; + +impl crate::observability::SupplementalEnricher for EmptyCloseEnricher { + fn enrich( + &self, + _model: &str, + _response: &ModelResponse, + ) -> (Option<String>, crate::events::GenerationEnrichment) { + (None, Default::default()) + } +} + +struct InvokeFailureModel; + +#[async_trait] +impl ChatModel<()> for InvokeFailureModel { + async fn invoke(&self, _: &(), _: ModelRequest) -> tinyinference_llm::Result<ModelResponse> { + Err(tinyinference_llm::Error::Model( + "synthetic provider failure".into(), + )) + } +} + +#[tokio::test] +async fn close_source_model_stamps_request_and_traces_completion() { + use crate::events::{EventListener, EventSink, RecordingListener}; + use crate::ids::RunId; + + let inner = Arc::new(scripted_model(&["answer"])); + let source = RecordingCloseSource { + model: inner.clone(), + builds: Mutex::new(Vec::new()), + }; + let events = EventSink::new(); + let listener = Arc::new(RecordingListener::new()); + events.subscribe(listener.clone() as Arc<dyn EventListener>); + let trace = crate::observability::SupplementalCallTrace::new( + events, + RunId::new("run-close"), + Arc::new(EmptyCloseEnricher), + ); + let close = CloseSourceModel { + source: &source, + model: "selected", + temperature: 0.2, + thread_id: Some("thread"), + trace: Some(&trace), + }; + close + .invoke(&(), ModelRequest::new(Vec::new())) + .await + .unwrap(); + + assert_eq!( + source.builds.lock().unwrap().as_slice(), + [("selected".to_owned(), 0.2, Some("thread".to_owned()))] + ); + + let requests = inner.requests.lock().unwrap(); + let correlation = requests[0].correlation.as_ref().unwrap(); + assert_eq!(correlation.run_id, "run-close"); + assert!(correlation.call_id.ends_with("-host-supplemental-1")); + let events = listener.events(); + assert_eq!(events.len(), 3); + assert!(matches!( + events[0].event, + crate::events::AgentEvent::ModelStarted { .. } + )); + assert!(matches!( + events[1].event, + crate::events::AgentEvent::UsageRecorded { .. } + )); + assert!(matches!( + events[2].event, + crate::events::AgentEvent::ModelCompleted { .. } + )); +} + +#[tokio::test] +async fn close_source_model_records_one_failure_when_stream_ends_without_terminal_item() { + use crate::events::{AgentEvent, EventListener, EventSink, RecordingListener}; + use crate::ids::RunId; + use futures::StreamExt; + + let inner = Arc::new(StreamScriptedModel { + responses: Mutex::new(VecDeque::from([Vec::new()])), + requests: Mutex::new(Vec::new()), + }); + let source = StaticCloseSource(inner); + let events = EventSink::new(); + let listener = Arc::new(RecordingListener::new()); + events.subscribe(listener.clone() as Arc<dyn EventListener>); + let trace = crate::observability::SupplementalCallTrace::new( + events, + RunId::new("run-close"), + Arc::new(EmptyCloseEnricher), + ); + let close = CloseSourceModel { + source: &source, + model: "selected", + temperature: 0.2, + thread_id: None, + trace: Some(&trace), + }; + let mut stream = close.stream(&(), ModelRequest::default()).await.unwrap(); + assert!(stream.next().await.is_none()); + let events_at_eof = listener.events(); + assert_eq!(events_at_eof.len(), 2); + assert!(matches!( + events_at_eof[1].event, + AgentEvent::ModelFailed { .. } + )); + drop(stream); + let events = listener.events(); + assert_eq!(events.len(), 2); + assert!(matches!(events[0].event, AgentEvent::ModelStarted { .. })); + assert!(matches!(events[1].event, AgentEvent::ModelFailed { .. })); +} + +#[tokio::test] +async fn close_source_model_records_failure_when_consumer_drops_stream_early() { + use crate::events::{AgentEvent, EventListener, EventSink, RecordingListener}; + use crate::ids::RunId; + use futures::StreamExt; + + let inner = Arc::new(StreamScriptedModel { + responses: Mutex::new(VecDeque::from([vec![ModelStreamItem::Started]])), + requests: Mutex::new(Vec::new()), + }); + let source = StaticCloseSource(inner); + let events = EventSink::new(); + let listener = Arc::new(RecordingListener::new()); + events.subscribe(listener.clone() as Arc<dyn EventListener>); + let trace = crate::observability::SupplementalCallTrace::new( + events, + RunId::new("run-close-drop"), + Arc::new(EmptyCloseEnricher), + ); + let close = CloseSourceModel { + source: &source, + model: "selected", + temperature: 0.2, + thread_id: None, + trace: Some(&trace), + }; + let mut stream = close.stream(&(), ModelRequest::default()).await.unwrap(); + assert!(matches!( + stream.next().await, + Some(ModelStreamItem::Started) + )); + drop(stream); + let events = listener.events(); + assert_eq!(events.len(), 2); + assert!(matches!(events[0].event, AgentEvent::ModelStarted { .. })); + assert!(matches!(events[1].event, AgentEvent::ModelFailed { .. })); +} + +#[tokio::test] +async fn close_source_model_completes_stream_trace_once() { + use crate::events::{AgentEvent, EventListener, EventSink, RecordingListener}; + use crate::ids::RunId; + use futures::StreamExt; + + let mut response = ModelResponse::assistant("answer"); + response.usage = Some(Usage::new(2, 1)); + let inner = Arc::new(StreamScriptedModel { + responses: Mutex::new(VecDeque::from([vec![ModelStreamItem::Completed(response)]])), + requests: Mutex::new(Vec::new()), + }); + let source = StaticCloseSource(inner); + let events = EventSink::new(); + let listener = Arc::new(RecordingListener::new()); + events.subscribe(listener.clone() as Arc<dyn EventListener>); + let trace = crate::observability::SupplementalCallTrace::new( + events, + RunId::new("run-close-stream"), + Arc::new(EmptyCloseEnricher), + ); + let close = CloseSourceModel { + source: &source, + model: "selected", + temperature: 0.2, + thread_id: None, + trace: Some(&trace), + }; + let mut stream = close.stream(&(), ModelRequest::default()).await.unwrap(); + assert!(matches!( + stream.next().await, + Some(ModelStreamItem::Completed(_)) + )); + drop(stream); + let events = listener.events(); + assert_eq!(events.len(), 3); + assert!(matches!(events[0].event, AgentEvent::ModelStarted { .. })); + assert!(matches!(events[1].event, AgentEvent::UsageRecorded { .. })); + assert!(matches!(events[2].event, AgentEvent::ModelCompleted { .. })); +} + +struct RouteRecordingEnricher( + Arc<Mutex<Vec<Option<tinyinference_llm::model::ResolvedModelRoute>>>>, +); + +impl crate::observability::SupplementalEnricher for RouteRecordingEnricher { + fn enrich( + &self, + model: &str, + response: &ModelResponse, + ) -> (Option<String>, crate::events::GenerationEnrichment) { + self.0.lock().unwrap().push(response.resolved_route.clone()); + (Some(model.to_owned()), Default::default()) + } +} + +struct RoutedStreamModel(tinyinference_llm::model::ResolvedModelRoute); + +#[async_trait] +impl ChatModel<()> for RoutedStreamModel { + async fn invoke(&self, _: &(), _: ModelRequest) -> tinyinference_llm::Result<ModelResponse> { + Err(tinyinference_llm::Error::Unsupported( + "stream-only test model".into(), + )) + } + + async fn stream( + &self, + _: &(), + _: ModelRequest, + ) -> tinyinference_llm::Result<tinyinference_llm::model::ModelStream> { + Ok( + tinyinference_llm::model::ModelStream::new(Box::pin(futures::stream::iter([ + ModelStreamItem::Completed(ModelResponse::assistant("answer")), + ]))) + .with_resolved_route(self.0.clone()), + ) + } +} + +#[tokio::test] +async fn close_source_model_enriches_the_fully_projected_stream_response() { + use crate::events::{AgentEvent, EventListener, EventSink, RecordingListener}; + use crate::ids::RunId; + use futures::StreamExt; + use tinyinference_llm::model::ResolvedModelRoute; + + let route = ResolvedModelRoute::new("provider-x", "actual-model", "fallback-route"); + let source = StaticCloseSource(Arc::new(RoutedStreamModel(route.clone()))); + let events = EventSink::new(); + let listener = Arc::new(RecordingListener::new()); + events.subscribe(listener.clone() as Arc<dyn EventListener>); + let observed_routes = Arc::new(Mutex::new(Vec::new())); + let trace = crate::observability::SupplementalCallTrace::new( + events, + RunId::new("run-close-route"), + Arc::new(RouteRecordingEnricher(observed_routes.clone())), + ); + let close = CloseSourceModel { + source: &source, + model: "selected-model", + temperature: 0.2, + thread_id: None, + trace: Some(&trace), + }; + let mut stream = close.stream(&(), ModelRequest::default()).await.unwrap(); + assert_eq!(stream.metadata().resolved_route.as_ref(), Some(&route)); + let response = match stream.next().await { + Some(ModelStreamItem::Completed(response)) => response, + other => panic!("expected completed response, got {other:?}"), + }; + assert_eq!(response.resolved_route.as_ref(), Some(&route)); + assert_eq!(observed_routes.lock().unwrap().as_slice(), [Some(route)]); + assert!(stream.next().await.is_none()); + let events = listener.events(); + assert!(matches!(events[0].event, AgentEvent::ModelStarted { .. })); + assert!(matches!(events[1].event, AgentEvent::ModelCompleted { .. })); +} + +#[tokio::test] +async fn close_source_model_records_invoke_and_stream_failures_once() { + use crate::events::{AgentEvent, EventListener, EventSink, RecordingListener}; + use crate::ids::RunId; + use futures::StreamExt; + + let events = EventSink::new(); + let listener = Arc::new(RecordingListener::new()); + events.subscribe(listener.clone() as Arc<dyn EventListener>); + let trace = crate::observability::SupplementalCallTrace::new( + events, + RunId::new("run-close-failures"), + Arc::new(EmptyCloseEnricher), + ); + let source = StaticCloseSource(Arc::new(InvokeFailureModel)); + let close = CloseSourceModel { + source: &source, + model: "selected", + temperature: 0.0, + thread_id: None, + trace: Some(&trace), + }; + assert!(close.invoke(&(), ModelRequest::default()).await.is_err()); + let events = listener.events(); + assert_eq!(events.len(), 2); + assert!(matches!(events[0].event, AgentEvent::ModelStarted { .. })); + assert!(matches!(events[1].event, AgentEvent::ModelFailed { .. })); + + let stream_source = StaticCloseSource(Arc::new(StreamScriptedModel { + responses: Mutex::new(VecDeque::from([vec![ModelStreamItem::Failed( + "synthetic stream failure".into(), + )]])), + requests: Mutex::new(Vec::new()), + })); + let stream_close = CloseSourceModel { + source: &stream_source, + ..close + }; + let mut stream = stream_close + .stream(&(), ModelRequest::default()) + .await + .unwrap(); + assert!(matches!( + stream.next().await, + Some(ModelStreamItem::Failed(_)) + )); + drop(stream); + let events = listener.events(); + assert_eq!(events.len(), 4); + assert!(matches!(events[2].event, AgentEvent::ModelStarted { .. })); + assert!(matches!(events[3].event, AgentEvent::ModelFailed { .. })); +} + #[test] fn legacy_and_options_close_futures_remain_send() { fn assert_send<T: Send>(_: T) {} diff --git a/crates/tinyagents-harness/src/config/context.rs b/crates/tinyagents-harness/src/config/context.rs new file mode 100644 index 000000000..d6786992f --- /dev/null +++ b/crates/tinyagents-harness/src/config/context.rs @@ -0,0 +1,133 @@ +use serde::{Deserialize, Serialize}; + +/// Default number of recent tool results preserved by microcompaction. +pub const DEFAULT_KEEP_RECENT_TOOL_RESULTS: usize = 5; + +/// Generic token and result controls used by context middleware. +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(default)] +pub struct ContextConfig { + /// Master switch for context-reduction middleware. + #[serde(default = "default_enabled")] + pub enabled: bool, + /// Enables clearing older tool-result payloads before summarization. + #[serde(default = "default_true")] + pub microcompact_enabled: bool, + /// Enables automatic summarization when context approaches its limit. + #[serde(default = "default_true")] + pub autocompact_enabled: bool, + /// Number of recent tool results left untouched by microcompaction. + #[serde(default = "default_microcompact_keep_recent")] + pub microcompact_keep_recent: usize, + /// Maximum bytes retained from one tool result; zero disables the cap. + #[serde(default = "default_tool_result_budget_bytes")] + pub tool_result_budget_bytes: usize, + /// Payload size at which a host may start a summary operation. + #[serde( + default = "default_summarizer_payload_threshold_tokens", + alias = "summarizer_payload_threshold_bytes" + )] + pub summarizer_payload_threshold_tokens: usize, + /// Maximum payload size passed to a summary operation. + #[serde( + default = "default_summarizer_max_payload_tokens", + alias = "summarizer_max_payload_bytes" + )] + pub summarizer_max_payload_tokens: usize, + /// Optional model hint for hosts that provide summarization. + #[serde(default)] + pub summarizer_model: Option<String>, + /// Requests markdown tool-result rendering where a tool supports it. + #[serde(default = "default_true")] + pub prefer_markdown_tool_output: bool, + /// Enables tool-output compaction before the byte-budget cap. + #[serde(default = "default_true")] + pub compaction_enabled: bool, + /// Optional absolute context-token threshold; zero means no override. + #[serde(default)] + pub compaction_trigger_tokens: Option<u64>, + /// Checkpoint strategy selected when context compaction runs. + #[serde(default)] + pub compaction_strategy: CompactionStrategy, +} + +/// Normalized compaction options carried by a turn. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub struct CompactionSettings { + /// Normalized absolute trigger override. + pub trigger_tokens: Option<u64>, + /// Checkpoint strategy. + pub strategy: CompactionStrategy, +} + +/// How context compaction creates a checkpoint. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum CompactionStrategy { + /// Preserve typed task state. + #[default] + TaskState, + /// Produce a free-form summary. + Summary, +} + +impl ContextConfig { + pub fn compaction_settings(&self) -> CompactionSettings { + CompactionSettings { + trigger_tokens: self.compaction_trigger_tokens.filter(|value| *value > 0), + strategy: self.compaction_strategy, + } + } +} + +impl CompactionStrategy { + pub fn parse(raw: &str) -> Option<Self> { + match raw.trim().to_ascii_lowercase().replace('-', "_").as_str() { + "task_state" | "typed" | "s4" => Some(Self::TaskState), + "summary" | "free_form" | "s0" => Some(Self::Summary), + _ => None, + } + } +} + +fn default_enabled() -> bool { + true +} +fn default_true() -> bool { + true +} +fn default_microcompact_keep_recent() -> usize { + DEFAULT_KEEP_RECENT_TOOL_RESULTS +} +fn default_tool_result_budget_bytes() -> usize { + crate::config::TurnConfig::default().tool_result_budget_bytes +} +fn default_summarizer_payload_threshold_tokens() -> usize { + 4_000 +} +fn default_summarizer_max_payload_tokens() -> usize { + 64_000 +} + +impl Default for ContextConfig { + fn default() -> Self { + Self { + enabled: default_enabled(), + microcompact_enabled: default_true(), + autocompact_enabled: default_true(), + microcompact_keep_recent: default_microcompact_keep_recent(), + tool_result_budget_bytes: default_tool_result_budget_bytes(), + summarizer_payload_threshold_tokens: default_summarizer_payload_threshold_tokens(), + summarizer_max_payload_tokens: default_summarizer_max_payload_tokens(), + summarizer_model: None, + prefer_markdown_tool_output: default_true(), + compaction_enabled: default_true(), + compaction_trigger_tokens: None, + compaction_strategy: CompactionStrategy::default(), + } + } +} + +#[cfg(test)] +#[path = "context_tests.rs"] +mod tests; diff --git a/crates/tinyagents-harness/src/config/context_tests.rs b/crates/tinyagents-harness/src/config/context_tests.rs new file mode 100644 index 000000000..4bc7cf506 --- /dev/null +++ b/crates/tinyagents-harness/src/config/context_tests.rs @@ -0,0 +1,75 @@ +use super::*; + +#[test] +fn context_config_preserves_defaults_aliases_and_strategy_spellings() { + let default = ContextConfig::default(); + let parsed: ContextConfig = toml::from_str("").unwrap(); + assert_eq!( + serde_json::to_value(parsed).unwrap(), + serde_json::to_value(&default).unwrap() + ); + + let canonical = r#" +enabled = false +microcompact_enabled = false +autocompact_enabled = true +microcompact_keep_recent = 9 +tool_result_budget_bytes = 1234 +summarizer_payload_threshold_tokens = 4000 +summarizer_max_payload_tokens = 64000 +summarizer_model = "test-model" +prefer_markdown_tool_output = false +compaction_enabled = true +compaction_trigger_tokens = 123 +compaction_strategy = "summary" +"#; + let value: ContextConfig = toml::from_str(canonical).unwrap(); + assert!(!value.enabled); + assert!(!value.microcompact_enabled); + assert!(value.autocompact_enabled); + assert_eq!(value.microcompact_keep_recent, 9); + assert_eq!(value.tool_result_budget_bytes, 1234); + assert_eq!(value.summarizer_payload_threshold_tokens, 4000); + assert_eq!(value.summarizer_max_payload_tokens, 64000); + assert_eq!(value.summarizer_model.as_deref(), Some("test-model")); + assert!(!value.prefer_markdown_tool_output); + assert!(value.compaction_enabled); + assert_eq!(value.compaction_trigger_tokens, Some(123)); + assert_eq!(value.compaction_settings().trigger_tokens, Some(123)); + assert_eq!( + value.compaction_settings().strategy, + CompactionStrategy::Summary + ); + + let aliases: ContextConfig = + toml::from_str("summarizer_payload_threshold_bytes=41\nsummarizer_max_payload_bytes=42") + .unwrap(); + assert_eq!(aliases.summarizer_payload_threshold_tokens, 41); + assert_eq!(aliases.summarizer_max_payload_tokens, 42); + for spelling in ["task_state", "typed", "s4"] { + assert_eq!( + CompactionStrategy::parse(spelling), + Some(CompactionStrategy::TaskState) + ); + } + for spelling in ["summary", "free_form", "s0"] { + assert_eq!( + CompactionStrategy::parse(spelling), + Some(CompactionStrategy::Summary) + ); + } + assert_eq!( + CompactionStrategy::parse(" TASK-STATE "), + Some(CompactionStrategy::TaskState) + ); + assert_eq!(CompactionStrategy::parse("unknown"), None); + let mut zero = value.clone(); + zero.compaction_trigger_tokens = Some(0); + assert_eq!(zero.compaction_settings().trigger_tokens, None); + let encoded = toml::to_string(&default).unwrap(); + let round_trip: ContextConfig = toml::from_str(&encoded).unwrap(); + assert_eq!( + serde_json::to_value(round_trip).unwrap(), + serde_json::to_value(default).unwrap() + ); +} diff --git a/crates/tinyagents-harness/src/config/mod.rs b/crates/tinyagents-harness/src/config/mod.rs index c251df6aa..f0861827c 100644 --- a/crates/tinyagents-harness/src/config/mod.rs +++ b/crates/tinyagents-harness/src/config/mod.rs @@ -28,9 +28,11 @@ //! (memory, security, budget, progress) are trait objects supplied separately; //! nothing in this module is a behaviour seam. +mod context; mod required_output; mod types; +pub use context::*; pub use required_output::{ find_required_block, output_satisfies_contract, repair_instruction, synthesize_block, }; diff --git a/crates/tinyagents-harness/src/context/types.rs b/crates/tinyagents-harness/src/context/types.rs index f751246a9..2024e77db 100644 --- a/crates/tinyagents-harness/src/context/types.rs +++ b/crates/tinyagents-harness/src/context/types.rs @@ -245,7 +245,7 @@ impl std::fmt::Debug for StateUpdate { /// channel. /// /// This is the harness-native complement to the graph -/// `Command`/`Interrupt` vocabulary (in `tinyagents-graph`): the agent loop drains any requested control at its safe +/// control requests: the agent loop drains any requested control at its safe /// checkpoints (after each model response) and acts on it, so behaviors like /// "stop after an early-exit tool" or "pause on budget" no longer need a /// bespoke side channel. Requests are visible via diff --git a/crates/tinyagents-harness/src/cost/dispatch_admission_tests.rs b/crates/tinyagents-harness/src/cost/dispatch_admission_tests.rs index 3046ba071..288b6fa0e 100644 --- a/crates/tinyagents-harness/src/cost/dispatch_admission_tests.rs +++ b/crates/tinyagents-harness/src/cost/dispatch_admission_tests.rs @@ -521,7 +521,10 @@ async fn agent_turn_preserves_typed_admission_refusal_from_chat_model() { options(Some(admission.clone()), None, &counts), ); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("controlled", Arc::new(controlled)); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "controlled", + Arc::new(controlled), + ))); let error = harness .invoke( @@ -557,7 +560,10 @@ async fn concurrent_agent_turns_keep_typed_admission_refusals_call_local() { options(Some(admission.clone()), None, &counts), ); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("controlled", Arc::new(controlled)); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "controlled", + Arc::new(controlled), + ))); let harness = Arc::new(harness); let invoke = |harness: Arc<AgentHarness<()>>| async move { harness @@ -618,7 +624,10 @@ async fn admitted_agent_turn_still_completes_normally() { options(Some(admission.clone()), None, &counts), ); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("controlled", Arc::new(controlled)); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "controlled", + Arc::new(controlled), + ))); harness .invoke( @@ -644,7 +653,10 @@ async fn streaming_agent_turn_preserves_typed_admission_refusal() { options(Some(admission.clone()), None, &counts), ); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("controlled", Arc::new(controlled)); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "controlled", + Arc::new(controlled), + ))); let error = harness .invoke_streaming( @@ -706,7 +718,10 @@ async fn assert_openai_retry_rechecks_committed_admission(with_money_budget: boo options(Some(admission.clone()), money_budget, &counts), ); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("controlled", Arc::new(controlled)); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "controlled", + Arc::new(controlled), + ))); let server_permit_state = admission.permits.clone(); let server_listener = listener.clone(); let (first_response_sent, first_response_received) = tokio::sync::oneshot::channel(); @@ -814,13 +829,13 @@ impl ChatModel<()> for LazyTypedRefusalModel { #[tokio::test] async fn turn_keeps_typed_dispatch_outcome_scope_through_lazy_stream_poll() { let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "controlled", Arc::new(LazyTypedRefusalModel { reason: DispatchRefusalReason::CallLimit, barrier: None, }), - ); + ))); let error = harness .invoke_streaming( @@ -843,21 +858,21 @@ async fn turn_keeps_typed_dispatch_outcome_scope_through_lazy_stream_poll() { async fn concurrent_lazy_turns_keep_distinct_typed_refusals_call_local() { let barrier = Arc::new(tokio::sync::Barrier::new(2)); let mut first_harness: AgentHarness<()> = AgentHarness::new(); - first_harness.register_model( + first_harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "controlled", Arc::new(LazyTypedRefusalModel { reason: DispatchRefusalReason::CallLimit, barrier: Some(barrier.clone()), }), - ); + ))); let mut second_harness: AgentHarness<()> = AgentHarness::new(); - second_harness.register_model( + second_harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "controlled", Arc::new(LazyTypedRefusalModel { reason: DispatchRefusalReason::Cancelled, barrier: Some(barrier), }), - ); + ))); let first = tokio::spawn(async move { first_harness diff --git a/crates/tinyagents-harness/src/error.rs b/crates/tinyagents-harness/src/error.rs index 2047fd851..4a7f5a442 100644 --- a/crates/tinyagents-harness/src/error.rs +++ b/crates/tinyagents-harness/src/error.rs @@ -1,9 +1,8 @@ //! Crate-wide error type and `Result` alias. //! -//! Every fallible surface of the recursive runtime — graph execution, the -//! harness agent loop, sub-agent recursion, and -//! registry binding — funnels through [`TinyAgentsError`] so failures from a -//! deeply nested child run roll up to the caller through one uniform type. +//! Fallible harness, sub-agent, orchestration-task, and registry operations +//! use [`TinyAgentsError`] so failures from a nested child run roll up through +//! one uniform type. //! Downstream code should prefer the [`Result`] alias exported here. use thiserror::Error; @@ -14,54 +13,49 @@ pub type Result<T> = std::result::Result<T, TinyAgentsError>; /// The single error type returned by every fallible TinyAgents operation. /// -/// Variants are grouped by the surface that raises them: graph construction and -/// execution, model/tool invocation, run limits and policy, graph durability, -/// and graph execution. +/// Variants cover model/tool invocation, harness control, run limits and +/// policy, orchestration tasks, and persistence. Legacy graph error variants +/// remain for public API compatibility; this workspace no longer ships a graph +/// execution runtime. #[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. + /// Legacy graph configuration error retained for API compatibility. The + /// current workspace has no graph runtime that emits it. #[error("graph start node is not configured")] MissingStart, - /// An edge, route, or run referenced a node name that is not present in the - /// graph. The payload is the missing node name. + /// Legacy graph configuration error retained for API compatibility. The + /// payload is the missing node name. #[error("node `{0}` does not exist")] MissingNode(String), - /// An edge declares a destination node that does not exist. The payload is - /// the missing target name. + /// Legacy graph configuration error retained for API compatibility. The + /// payload is the missing target name. #[error("edge points to missing node `{0}`")] MissingEdgeTarget(String), - /// A conditional router returned a `route` label that is not wired to any - /// destination from `node`. + /// Legacy graph configuration error retained for API compatibility. The + /// current workspace has no graph router that emits it. #[error("conditional route `{route}` from node `{node}` does not exist")] MissingRoute { node: String, route: String }, - /// Graph execution performed more super-steps than the configured recursion - /// limit allows (typically an unintended cycle). The payload is the limit - /// that was hit. Contrast with [`TinyAgentsError::SubAgentDepth`], which - /// counts nested run-tree levels rather than super-steps. + /// Legacy graph recursion error retained for API compatibility. The + /// payload is the configured limit; current sub-agent depth failures use + /// [`TinyAgentsError::SubAgentDepth`]. #[error("graph exceeded the recursion limit of {0} steps")] RecursionLimit(usize), /// A sub-agent invocation would exceed the configured maximum recursion /// depth. The payload is the `max_depth` cap that was reached. /// - /// This is distinct from [`TinyAgentsError::RecursionLimit`], which counts - /// graph *super-steps*; `SubAgentDepth` counts nested run-tree *levels* - /// (parent → child → grandchild …) so the two limits can be reasoned about - /// and surfaced independently. + /// Nested run-tree levels (parent → child → grandchild …) are counted + /// independently from the legacy [`TinyAgentsError::RecursionLimit`]. #[error("sub-agent recursion exceeded the maximum depth of {0}")] SubAgentDepth(usize), - /// A single graph node was activated more times within one run than the - /// `RecursionPolicy`'s (in `tinyagents-graph`) `max_visits_per_node` allows (an - /// unbounded node-loop). This is node-loop recursion, tracked separately - /// from [`TinyAgentsError::RecursionLimit`] (total super-steps) and - /// [`TinyAgentsError::SubAgentDepth`] (run-tree depth). + /// Legacy graph node-visit limit retained for API compatibility. The + /// current workspace has no graph executor that emits it. #[error("node `{node}` exceeded its visit limit of {limit}")] NodeVisitLimit { /// The node that was over-visited. @@ -299,30 +293,30 @@ pub enum TinyAgentsError { #[error("embedding error: {0}")] Embedding(String), - // --- graph durability errors --- - /// Generic graph runtime error. + // --- legacy graph errors and current orchestration errors --- + /// Generic error variant retained for compatibility and used by the + /// orchestration-task and completion-router APIs. It does not imply that a + /// graph execution runtime is present. #[error("graph error: {0}")] Graph(String), - /// Execution was interrupted (human-in-the-loop / external approval). + /// A harness run was interrupted for human input or external approval. + /// The legacy display text is retained for compatibility. #[error("graph interrupted at node `{node}`: {message}")] Interrupted { node: String, message: String }, - /// Two or more concurrent branches in a single superstep wrote the same - /// non-aggregate channel (for example a `LastValue` channel, in - /// `tinyagents-graph`), so the merge cannot pick a single deterministic winner. Use an - /// aggregate channel (one whose `Channel::allows_concurrent` is `true`, such as - /// `Topic` or `BinaryAggregate`) when fan-out branches must - /// write the same key. The payload describes the offending channel. + /// Legacy graph state-merge error retained for API compatibility. The + /// current workspace has no graph executor that emits it. #[error("invalid concurrent update: {0}")] InvalidConcurrentUpdate(String), - /// A checkpoint could not be written, read, or located. + /// Legacy execution-checkpoint error retained for API compatibility. + /// Current durable conversation persistence is owned by `tinyagents-session`. #[error("checkpoint error: {0}")] Checkpoint(String), - /// Resume was requested but checkpointing was not configured or no - /// checkpoint was found. + /// Legacy resume error retained for API compatibility. Session resume is + /// addressed through `tinyagents-session` records. #[error("cannot resume: {0}")] Resume(String), @@ -349,11 +343,9 @@ pub enum TinyAgentsError { /// writing a backing database for the session store and run ledger /// (`tinyagents-session`). /// - /// Distinct from [`TinyAgentsError::Checkpoint`], which covers graph - /// checkpoint durability: a session-store failure means run *history* could - /// not be recorded or queried, while a checkpoint failure means a run - /// cannot be resumed. The payload carries the operation context and the - /// underlying driver message. + /// The payload carries the operation context and the underlying driver + /// message. Legacy [`TinyAgentsError::Checkpoint`] remains only for public + /// compatibility; session persistence failures use this variant. #[error("storage error: {0}")] Storage(String), } diff --git a/crates/tinyagents-harness/src/events/types.rs b/crates/tinyagents-harness/src/events/types.rs index 619c23bc9..106bbe457 100644 --- a/crates/tinyagents-harness/src/events/types.rs +++ b/crates/tinyagents-harness/src/events/types.rs @@ -658,7 +658,7 @@ pub enum AgentEvent { /// An existing sub-agent was *reused* for a follow-up turn rather than /// reconstructed, carrying the prior conversation context forward. /// - /// Emitted by `SubAgentSession` in `tinyagents-orchestration` on every send + /// Emitted by the host session lifecycle on every continuation /// after the first (i.e. `turn >= 1`), so post-completion reuse — the /// orchestrator → sub-agent → human input → *same* sub-agent pattern — is /// visible in the event stream and distinguishable from a fresh diff --git a/crates/tinyagents-harness/src/handoff.rs b/crates/tinyagents-harness/src/handoff.rs deleted file mode 100644 index 2be55fd47..000000000 --- a/crates/tinyagents-harness/src/handoff.rs +++ /dev/null @@ -1,364 +0,0 @@ -//! Progressive-disclosure handoff cache for oversized tool results. -//! -//! Sub-agents regularly call tools -//! that return megabyte-scale payloads — `GMAIL_LIST_MESSAGES`, -//! `NOTION_GET_PAGE`, `GOOGLEDRIVE_LIST_FILES`. The default behaviour pushes -//! that raw blob into the sub-agent's history as a tool-result message, and -//! the NEXT iteration ships the bloated history back to the provider where -//! it hits the model's context-length ceiling. -//! -//! Progressive disclosure fixes this: when a tool returns too much data we -//! stash the full payload here, replace it in history with a short -//! placeholder (size + preview + `result_id` + how to query it), and expose -//! an extraction tool — host-side, since a tool is host vocabulary — that the -//! sub-agent can call with a targeted query. The extractor only runs when -//! the sub-agent actually asks for a narrower view. -//! -//! This module owns: -//! * the thresholds and limits (token cut-off, preview size, max entries); -//! * 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; - -// ── Tunables ─────────────────────────────────────────────────────────────── - -/// Token threshold above which a tool result is routed to the handoff -/// cache instead of being pushed into history raw. Token count is -/// estimated at ~4 chars/token. -/// -/// Set at `50_000` so the clean Gmail / Notion envelopes emitted by provider -/// post-processing fit through unchanged for normal workloads — only -/// genuinely oversized results (bulk fetches, raw thread dumps) are routed -/// through the `extract_from_result` path. -pub const HANDOFF_OVERSIZE_THRESHOLD_TOKENS: usize = 50_000; - -/// Characters of the raw payload to surface in the placeholder preview. -/// Enough for the sub-agent to recognise the shape (JSON keys, first -/// record) and often small enough to answer trivial questions without a -/// follow-up `extract_from_result` call. -pub const HANDOFF_PREVIEW_CHARS: usize = 1500; - -/// Maximum entries per session. Bounded to keep memory use predictable on -/// long-running sub-agents that might call many large tools. When over -/// capacity we evict the oldest entry (FIFO); callers see "no cached -/// result" for evicted ids and can either re-run the tool or ask the -/// 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<String>, -} - -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 -/// the top of `run_typed_mode` and shared (via `Arc`) with both the inner -/// tool-call loop (writes) and the `extract_from_result` tool (reads). -#[derive(Default)] -pub struct ResultHandoffCache { - inner: StdMutex<HandoffInner>, -} - -#[derive(Default)] -struct HandoffInner { - /// FIFO of inserted ids, used for eviction. - order: Vec<String>, - /// Content by id. - entries: HashMap<String, CachedResult>, - /// Monotonic counter for id generation within the session. - next_id: u64, -} - -/// One stashed tool result: the raw content plus the name of the tool that -/// produced it, returned by [`ResultHandoffCache::get`] and stored on -/// [`ResultHandoffCache::store`]. -pub struct CachedResult { - /// Name of the tool whose result was stashed. - pub tool_name: String, - /// The (cleaned) payload the sub-agent can query with - /// `extract_from_result`. - pub content: String, -} - -impl ResultHandoffCache { - /// Creates an empty cache. - pub fn new() -> Self { - Self::default() - } - - /// Stash a payload and return a stable, short, grep-friendly id. - pub fn store(&self, tool_name: String, content: String) -> String { - let mut g = match self.inner.lock() { - Ok(g) => g, - Err(poisoned) => poisoned.into_inner(), - }; - g.next_id = g.next_id.saturating_add(1); - let id = format!("res_{:x}", g.next_id); - g.order.push(id.clone()); - g.entries - .insert(id.clone(), CachedResult { tool_name, content }); - while g.order.len() > HANDOFF_MAX_ENTRIES { - let evicted = g.order.remove(0); - g.entries.remove(&evicted); - } - id - } - - /// Looks up a previously stashed result by id, returning `None` if it was - /// never stored or has since been evicted (FIFO, past - /// [`HANDOFF_MAX_ENTRIES`]). - pub fn get(&self, result_id: &str) -> Option<CachedResult> { - let g = self.inner.lock().ok()?; - g.entries.get(result_id).map(|r| CachedResult { - tool_name: r.tool_name.clone(), - content: r.content.clone(), - }) - } -} - -/// Apply the progressive-disclosure handoff to a tool result. If a cache is -/// present and the cleaned result is large enough, stash the raw payload and -/// substitute a short placeholder the sub-agent can drill into with -/// `extract_from_result`. Errors and already-extracted output pass through -/// unchanged. -/// -/// `threshold_tokens` is the size above which a payload is stashed rather than -/// inlined. It is a parameter rather than a constant read here because a host -/// legitimately varies it — by agent, by model context window, or to exercise -/// 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 == 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 { - tracing::debug!( - tool = %tool_name, - before_bytes = pre_len, - after_bytes = cleaned.len(), - saved_pct = ((pre_len - cleaned.len()) * 100) / pre_len.max(1), - "[handoff] cleaned tool output (stripped markup/data-uris/whitespace)" - ); - } - cleaned - }; - 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(config, tool_name, &id, &cleaned); - tracing::info!( - task_id = %task_id, - agent_id = %agent_id, - tool = %tool_name, - raw_tokens = tokens, - raw_bytes = cleaned.len(), - threshold_tokens, - result_id = %id, - "[handoff] stashed oversized tool output; substituted placeholder into history" - ); - placeholder - } else { - cleaned - } -} - -// ── Placeholder renderer ─────────────────────────────────────────────────── - -/// 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 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( - 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\ - {extractor}(result_id=\"{result_id}\", query=\"<specific question>\")\n\ - Good queries name the exact fields/identifiers you need \ - (e.g. \"subject and sender of the 5 most recent messages\"). \ - Tool: {tool_name}", - raw_bytes = raw.len(), - preview_chars = preview.chars().count(), - ) -} - -// ── Content hygiene helpers (used by the extract path) ───────────────────── - -use regex::Regex; -use std::sync::LazyLock; - -/// Strip common noise from tool outputs before they're stashed or chunked. -/// -/// Agent tools frequently return raw HTML email bodies, inline SVG, base64 -/// data URIs, CSS/JS blocks, and collapsed whitespace — all of which bloat -/// the handoff cache and waste summarizer context on tokens that carry -/// zero semantic value for most extraction queries. Cleaning before the -/// oversize check means (a) some payloads drop below threshold entirely -/// and skip the extract pipeline, (b) chunked payloads fit more real -/// content per chunk, and (c) summarizers see clean text instead of -/// parsing around markup. -pub fn clean_tool_output(content: &str) -> String { - static SCRIPT_RE: LazyLock<Regex> = - LazyLock::new(|| Regex::new(r"(?is)<script\b[^>]*>.*?</script\s*>").unwrap()); - static STYLE_RE: LazyLock<Regex> = - LazyLock::new(|| Regex::new(r"(?is)<style\b[^>]*>.*?</style\s*>").unwrap()); - static SVG_RE: LazyLock<Regex> = - LazyLock::new(|| Regex::new(r"(?is)<svg\b[^>]*>.*?</svg\s*>").unwrap()); - static HTML_COMMENT_RE: LazyLock<Regex> = - LazyLock::new(|| Regex::new(r"(?s)<!--.*?-->").unwrap()); - static DATA_URI_RE: LazyLock<Regex> = - LazyLock::new(|| Regex::new(r"(?i)data:[a-z0-9.+\-/]+;base64,[A-Za-z0-9+/=]+").unwrap()); - static HTML_TAG_RE: LazyLock<Regex> = LazyLock::new(|| Regex::new(r"<[^>]+>").unwrap()); - static WS_RE: LazyLock<Regex> = LazyLock::new(|| Regex::new(r"[ \t\f\v]+").unwrap()); - static BLANK_LINE_RE: LazyLock<Regex> = LazyLock::new(|| Regex::new(r"\n{3,}").unwrap()); - - let cleaned = SCRIPT_RE.replace_all(content, ""); - let cleaned = STYLE_RE.replace_all(&cleaned, ""); - let cleaned = SVG_RE.replace_all(&cleaned, "[svg]"); - let cleaned = HTML_COMMENT_RE.replace_all(&cleaned, ""); - let cleaned = DATA_URI_RE.replace_all(&cleaned, "[data-uri]"); - let cleaned = HTML_TAG_RE.replace_all(&cleaned, ""); - let cleaned = WS_RE.replace_all(&cleaned, " "); - let cleaned = BLANK_LINE_RE.replace_all(&cleaned, "\n\n"); - cleaned.trim().to_string() -} - -/// Split `content` into chunks no larger than `budget` bytes, breaking -/// at natural boundaries (blank lines, then single newlines) so the -/// extraction LLM rarely sees a structure torn mid-record. Falls back to -/// char-safe slicing for pathological single-line inputs. -pub fn chunk_content(content: &str, budget: usize) -> Vec<String> { - if content.len() <= budget { - return vec![content.to_string()]; - } - - let mut chunks: Vec<String> = Vec::new(); - let mut current = String::with_capacity(budget.min(content.len())); - - let flush = |current: &mut String, chunks: &mut Vec<String>| { - if !current.is_empty() { - chunks.push(std::mem::take(current)); - } - }; - - for line in content.lines() { - let projected = current.len() + line.len() + 1; - if projected > budget && !current.is_empty() { - flush(&mut current, &mut chunks); - } - if line.len() > budget { - // Single line exceeds budget (e.g. JSON with no formatting). - // Emit any pending content, then slice the line at char - // boundaries so we don't panic on multi-byte chars. - flush(&mut current, &mut chunks); - let mut remaining = line; - while !remaining.is_empty() { - let mut cut = budget.min(remaining.len()); - while cut > 0 && !remaining.is_char_boundary(cut) { - cut -= 1; - } - if cut == 0 { - // Degenerate case — shouldn't happen for normal - // text. Take the entire remaining line to avoid - // an infinite loop. - chunks.push(remaining.to_string()); - break; - } - chunks.push(remaining[..cut].to_string()); - remaining = &remaining[cut..]; - } - } else { - current.push_str(line); - current.push('\n'); - } - } - flush(&mut current, &mut chunks); - - chunks -} - -#[cfg(test)] -#[path = "handoff_tests.rs"] -mod tests; diff --git a/crates/tinyagents-harness/src/handoff_tests.rs b/crates/tinyagents-harness/src/handoff_tests.rs deleted file mode 100644 index ac6b3a4a4..000000000 --- a/crates/tinyagents-harness/src/handoff_tests.rs +++ /dev/null @@ -1,284 +0,0 @@ -//! Tests for the progressive-disclosure handoff cache. -//! -//! The module had none where it came from — it was covered only indirectly, -//! through the sub-agent runner's integration tests. These pin the behaviour -//! that a caller actually depends on, and in particular the cases where a -//! regression would be silent rather than loud. - -use super::*; - -fn cache() -> ResultHandoffCache { - ResultHandoffCache::new() -} - -/// A payload comfortably over any threshold used here. -fn big(n: usize) -> String { - "x".repeat(n) -} - -// ── The cache ───────────────────────────────────────────────────────────────── - -#[test] -fn a_stored_payload_round_trips_by_id() { - let c = cache(); - let id = c.store("gmail_list".to_string(), "payload".to_string()); - let got = c.get(&id).expect("stored payload is retrievable"); - assert_eq!(got.content, "payload"); - assert_eq!(got.tool_name, "gmail_list"); -} - -#[test] -fn ids_are_unique_across_stores() { - // Ids are handed to a model and used as lookup keys. A collision would - // silently serve one tool's payload in answer to another's query. - let c = cache(); - let ids: Vec<String> = (0..5) - .map(|i| c.store(format!("tool{i}"), format!("body{i}"))) - .collect(); - let mut sorted = ids.clone(); - sorted.sort(); - sorted.dedup(); - assert_eq!(sorted.len(), ids.len(), "duplicate result ids: {ids:?}"); -} - -#[test] -fn an_unknown_id_is_none_rather_than_a_panic() { - // The id comes back through a model, so it can be anything at all. - assert!(cache().get("res_deadbeef").is_none()); - assert!(cache().get("").is_none()); -} - -#[test] -fn eviction_is_fifo_and_bounded() { - // The cache is per-spawn and unbounded growth is a leak, but evicting the - // WRONG end would drop the payload a model is most likely to ask about. - let c = cache(); - let ids: Vec<String> = (0..HANDOFF_MAX_ENTRIES + 2) - .map(|i| c.store("t".to_string(), format!("body{i}"))) - .collect(); - - assert!(c.get(&ids[0]).is_none(), "oldest entry should be evicted"); - assert!(c.get(&ids[1]).is_none(), "second-oldest should be evicted"); - let newest = ids.last().expect("at least one id"); - assert!(c.get(newest).is_some(), "newest entry must survive"); -} - -// ── apply_handoff ───────────────────────────────────────────────────────────── - -#[test] -fn a_small_result_passes_through_untouched_and_is_not_cached() { - let c = cache(); - 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(), - "nothing should have been stashed for a small result" - ); -} - -#[test] -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, - &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")); - assert!(out.contains("extract_from_result")); - - // The placeholder must name an id that actually resolves, or the model is - // told to call a tool that cannot find anything. - let id = out - .split("result_id=\"") - .nth(1) - .and_then(|s| s.split('"').next()) - .expect("placeholder carries a result_id"); - assert_eq!( - c.get(id).expect("the advertised id resolves").content, - raw, - "full fidelity must be preserved in the cache" - ); -} - -#[test] -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(), - &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); -} - -#[test] -fn an_error_result_passes_through_however_large() { - // Errors are diagnostic text the agent must read directly. Stashing one - // 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, - &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] -fn an_extraction_result_is_never_re_stashed() { - // The extractor answers a query against a stashed payload. Handing its - // answer back through the same path would stash the answer and hand the - // model another placeholder — a loop that never converges. - let c = cache(); - let raw = big(8_000); - let out = apply_handoff( - &c, - &HandoffConfig::default(), - "extract_from_result", - "task", - "agent", - raw.clone(), - 1, - ); - assert_eq!(out, raw); -} - -// ── The placeholder ─────────────────────────────────────────────────────────── - -#[test] -fn the_placeholder_reports_size_and_previews_the_head() { - let raw = format!("HEAD-MARKER{}", big(5_000)); - let text = build_handoff_placeholder(&HandoffConfig::default(), "gmail_list", "res_1", &raw); - - assert!(text.contains("res_1")); - assert!(text.contains("gmail_list")); - assert!( - text.contains(&raw.len().to_string()), - "raw byte count shown" - ); - assert!( - text.contains("HEAD-MARKER"), - "the preview must come from the start of the payload" - ); - assert!( - text.len() < raw.len(), - "a placeholder larger than the payload defeats the purpose" - ); -} - -#[test] -fn a_short_payload_previews_whole_without_padding() { - 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. - assert!(text.contains("first 4 chars")); -} - -#[test] -fn the_preview_is_capped_for_a_large_payload() { - let raw = big(HANDOFF_PREVIEW_CHARS * 4); - let text = build_handoff_placeholder(&HandoffConfig::default(), "t", "res_1", &raw); - assert!(text.contains(&format!("first {HANDOFF_PREVIEW_CHARS} chars"))); -} - -#[test] -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(&HandoffConfig::default(), "t", "res_1", &raw); - assert!(text.is_char_boundary(text.len())); -} diff --git a/crates/tinyagents-harness/src/host/model_resolver.rs b/crates/tinyagents-harness/src/host/model_resolver.rs index ad1fd5a15..de30aba15 100644 --- a/crates/tinyagents-harness/src/host/model_resolver.rs +++ b/crates/tinyagents-harness/src/host/model_resolver.rs @@ -228,6 +228,28 @@ pub trait ModelResolver<State: Send + Sync>: Send + Sync { /// should therefore be cheap and must not assume the answer is memoized by /// the caller. async fn resolve(&self, req: &ModelResolveRequest) -> Result<Arc<dyn ChatModel<State>>>; + + /// Resolves an exact configured name without falling back to a default. + /// Used by the generic runtime to walk an explicit fallback chain. Hosts + /// with dynamic routing may leave this unsupported; named fallbacks then + /// cannot be selected through that resolver. + async fn resolve_named(&self, _name: &str) -> Result<Option<Arc<dyn ChatModel<State>>>> { + Ok(None) + } + + /// Returns a cheap synchronous preview for middleware that shapes a + /// request from the target model profile. Network-backed resolvers should + /// leave this unavailable and use their normal resolution path. + fn preview(&self, _req: &ModelResolveRequest) -> Option<Arc<dyn ChatModel<State>>> { + None + } + + /// The selected configured name, when the resolver can identify it + /// without changing its answer. This carries observability and fallback + /// identity alongside the executable model. + fn selected_name(&self, _req: &ModelResolveRequest) -> Option<String> { + None + } } // ── FixedModelResolver ──────────────────────────────────────────────────────── @@ -242,12 +264,21 @@ pub trait ModelResolver<State: Send + Sync>: Send + Sync { /// about product tiers. pub struct FixedModelResolver<State: Send + Sync> { model: Arc<dyn ChatModel<State>>, + name: Option<String>, } impl<State: Send + Sync> FixedModelResolver<State> { /// Wraps `model` as the answer to every resolution request. pub fn new(model: Arc<dyn ChatModel<State>>) -> Self { - Self { model } + Self { model, name: None } + } + + /// Wraps one model while preserving its configured runtime name. + pub fn named(name: impl Into<String>, model: Arc<dyn ChatModel<State>>) -> Self { + Self { + model, + name: Some(name.into()), + } } /// The wrapped model. @@ -262,6 +293,7 @@ impl<State: Send + Sync> Clone for FixedModelResolver<State> { fn clone(&self) -> Self { Self { model: Arc::clone(&self.model), + name: self.name.clone(), } } } @@ -278,6 +310,28 @@ impl<State: Send + Sync> ModelResolver<State> for FixedModelResolver<State> { async fn resolve(&self, _req: &ModelResolveRequest) -> Result<Arc<dyn ChatModel<State>>> { Ok(Arc::clone(&self.model)) } + + async fn resolve_named(&self, name: &str) -> Result<Option<Arc<dyn ChatModel<State>>>> { + let matches = self.name.as_deref() == Some(name) + || self + .model + .profile() + .and_then(|profile| profile.model.as_deref()) + == Some(name); + Ok(matches.then(|| Arc::clone(&self.model))) + } + + fn preview(&self, _req: &ModelResolveRequest) -> Option<Arc<dyn ChatModel<State>>> { + Some(Arc::clone(&self.model)) + } + + fn selected_name(&self, _req: &ModelResolveRequest) -> Option<String> { + self.name.clone().or_else(|| { + self.model + .profile() + .and_then(|profile| profile.model.clone()) + }) + } } /// Resolves from models that a host has already constructed for one turn. @@ -289,6 +343,7 @@ impl<State: Send + Sync> ModelResolver<State> for FixedModelResolver<State> { /// credentials, capabilities, or provider state. pub struct PrebuiltModelResolver<State: Send + Sync> { primary: Arc<dyn ChatModel<State>>, + primary_name: Option<String>, lead: Arc<dyn ChatModel<State>>, routes: std::collections::HashMap<String, Arc<dyn ChatModel<State>>>, } @@ -297,6 +352,7 @@ impl<State: Send + Sync> Clone for PrebuiltModelResolver<State> { fn clone(&self) -> Self { Self { primary: Arc::clone(&self.primary), + primary_name: self.primary_name.clone(), lead: Arc::clone(&self.lead), routes: self.routes.clone(), } @@ -312,6 +368,7 @@ impl<State: Send + Sync> PrebuiltModelResolver<State> { Self { lead: Arc::clone(&primary), primary, + primary_name: None, routes: routes.into_iter().collect(), } } @@ -321,6 +378,12 @@ impl<State: Send + Sync> PrebuiltModelResolver<State> { self.lead = lead; self } + + /// Gives the primary model its exact configured name for fallback lookup. + pub fn with_primary_name(mut self, name: impl Into<String>) -> Self { + self.primary_name = Some(name.into()); + self + } } impl<State: Send + Sync> std::fmt::Debug for PrebuiltModelResolver<State> { @@ -350,6 +413,50 @@ impl<State: Send + Sync> ModelResolver<State> for PrebuiltModelResolver<State> { .cloned() .unwrap_or_else(|| Arc::clone(&self.primary))) } + + async fn resolve_named(&self, name: &str) -> Result<Option<Arc<dyn ChatModel<State>>>> { + if self.primary_name.as_deref() == Some(name) { + return Ok(Some(Arc::clone(&self.primary))); + } + Ok(self.routes.get(name).cloned()) + } + + fn preview(&self, req: &ModelResolveRequest) -> Option<Arc<dyn ChatModel<State>>> { + if req.is_team_lead { + return Some(Arc::clone(&self.lead)); + } + let hinted_route = req + .model_hint + .as_deref() + .map(str::trim) + .filter(|hint| !hint.is_empty()) + .map(|hint| format!("hint:{hint}")); + let route = req.model_pin().or(hinted_route.as_deref()); + Some( + route + .and_then(|name| self.routes.get(name)) + .cloned() + .unwrap_or_else(|| Arc::clone(&self.primary)), + ) + } + + fn selected_name(&self, req: &ModelResolveRequest) -> Option<String> { + let hinted_route = req + .model_hint + .as_deref() + .map(str::trim) + .filter(|hint| !hint.is_empty()) + .map(|hint| format!("hint:{hint}")); + let route = req.model_pin().or(hinted_route.as_deref()); + if let Some(name) = route.filter(|name| self.routes.contains_key(*name)) { + return Some(name.to_owned()); + } + self.primary_name.clone().or_else(|| { + self.primary + .profile() + .and_then(|profile| profile.model.clone()) + }) + } } #[cfg(test)] diff --git a/crates/tinyagents-harness/src/host/model_resolver_tests.rs b/crates/tinyagents-harness/src/host/model_resolver_tests.rs index a7c1830b7..c84826b62 100644 --- a/crates/tinyagents-harness/src/host/model_resolver_tests.rs +++ b/crates/tinyagents-harness/src/host/model_resolver_tests.rs @@ -201,6 +201,35 @@ async fn fixed_resolver_ignores_every_routing_field() { ); } +#[tokio::test] +async fn prebuilt_named_lookup_does_not_substitute_the_primary_for_unknown_names() { + let resolver = PrebuiltModelResolver::new( + Arc::new(EchoModel("primary")), + [( + "backup".to_string(), + Arc::new(EchoModel("backup")) as Arc<dyn ChatModel<()>>, + )], + ) + .with_primary_name("primary"); + + let backup = resolver + .resolve_named("backup") + .await + .expect("lookup succeeds") + .expect("registered named route exists"); + assert!(Arc::ptr_eq( + &backup, + resolver.routes.get("backup").expect("backup route") + )); + assert!( + resolver + .resolve_named("unregistered") + .await + .expect("unknown names are a normal miss") + .is_none() + ); +} + #[tokio::test] async fn resolved_model_shares_the_wrapped_allocation() { // Resolution must hand back a handle to the same long-lived client, not diff --git a/crates/tinyagents-harness/src/ids/README.md b/crates/tinyagents-harness/src/ids/README.md index 5333cf542..b3a352e30 100644 --- a/crates/tinyagents-harness/src/ids/README.md +++ b/crates/tinyagents-harness/src/ids/README.md @@ -5,8 +5,8 @@ dependency-free id-allocation primitives used to mint them. ## Why this exists -Every long-lived harness concept (runs, threads, calls, events, graph nodes, -checkpoints, interrupts, sessions, cells) is keyed by a small, cheap-to-clone +Every long-lived harness concept (runs, threads, calls, events, orchestration +task labels, checkpoints, interrupts, sessions, cells) is keyed by a small, cheap-to-clone string newtype instead of a bare `String`. Wrapping the raw string prevents accidentally passing, say, a `ThreadId` where a `RunId` is expected, and gives every id a consistent `new`/`as_str`/`Display`/`From<&str>`/`From<String>` @@ -14,8 +14,8 @@ surface generated by one macro (`impl_string_id!` in `mod.rs`). These ids are also what makes recursion observable: a `RunId` names one run, and pairing it with the `root_run_id`/`parent_run_id` recorded in -`crate::events::HarnessRunStatus` lets a child run (a sub-agent or sub-graph -invocation) be traced back up to the top-level run that spawned it. +`crate::events::HarnessRunStatus` lets a child run (for example a sub-agent +invocation) be traced back to the top-level run that spawned it. ## Public surface @@ -25,8 +25,7 @@ invocation) be traced back up to the top-level run that spawned it. `Display`, and `From<&str>`/`From<String>` via `impl_string_id!` in `mod.rs`. - [`ExecutionStatus`], [`HarnessPhase`] (`types.rs`) — coarse lifecycle enums - shared by direct model calls, agent loops, and graph-node invocations, both - serde-encoded `snake_case`. + shared by direct model calls and agent loops, both serde-encoded `snake_case`. - Allocators (`mod.rs`): [`new_run_id`], [`new_checkpoint_id`], [`new_session_id`], [`new_cell_id`], [`new_call_id`] — mint fresh, process-unique (and, for run/checkpoint ids, restart-collision-free) ids. @@ -49,11 +48,10 @@ invocation) be traced back up to the top-level run that spawned it. `process_nonce` (seeded once from wall-clock nanoseconds) gives collision-freedom across a process restart. An id minted from `next_seq` alone would restart at the same values after every process restart. -- `new_checkpoint_id`'s restart-safety is load-bearing, not cosmetic: - checkpoint ids key the parent-lineage spine that graph pruning, history - listing, and `ResumeTarget::Checkpoint` all walk. A duplicate checkpoint id - across a restart can delete a live record's ancestor or resume the wrong - checkpoint. +- `new_checkpoint_id` mixes the per-process nonce into ids persisted by + adapters, avoiding duplicate minted values after a restart. This is an + opaque compatibility id; durable conversation identity and resume are owned + by `tinyagents-session` and use `SessionRef`. - `now_ms` is the crate's single millisecond-timestamp convention (epoch, saturating to `0` before the epoch). Reuse it rather than re-deriving a timestamp locally, so every timestamped record in the crate agrees on the diff --git a/crates/tinyagents-harness/src/ids/mod.rs b/crates/tinyagents-harness/src/ids/mod.rs index 932532521..99f41d895 100644 --- a/crates/tinyagents-harness/src/ids/mod.rs +++ b/crates/tinyagents-harness/src/ids/mod.rs @@ -3,7 +3,7 @@ //! These ids are the keys that make recursion observable and correlatable: a //! [`RunId`] names one run, and pairing it with the `root_run_id` / //! `parent_run_id` recorded in [`crate::events::HarnessRunStatus`] lets -//! a child run (a sub-agent or sub-graph invocation) be traced back up to the +//! a child run (for example, a sub-agent invocation) be traced back up to the //! top-level run that spawned it. A `From`/`Display`/`as_str` surface is //! generated for each newtype by a single macro so the ids stay cheap to clone, //! log, and serialize. @@ -88,13 +88,12 @@ pub fn next_seq() -> u64 { /// (in practice) distinct across restarts. /// /// The bare monotonic [`next_seq`] restarts at `0` in every new process, so ids -/// built from it alone (`ckpt-0`, `ckpt-1`, …) *collide* across a restart — a -/// resumed thread would re-mint checkpoint ids it already used, corrupting the -/// parent-lineage map and time-travel resume. Mixing in this nonce makes ids -/// collision-free across restarts while [`next_seq`] keeps them ordered within -/// a process. Seeded once from the wall clock (nanoseconds since the epoch); -/// only ever used as an opaque uniqueness component, never parsed or compared -/// for time. +/// built from it alone could repeat after a restart. Mixing in this nonce +/// avoids collisions for adapters that persist generated ids, while +/// [`next_seq`] keeps them ordered within a process. Session resume itself is +/// owned by `tinyagents-session` and does not rely on this nonce. The nonce is +/// seeded once from wall-clock nanoseconds; it is only used as an opaque +/// uniqueness component, never parsed or compared for time. pub fn process_nonce() -> u64 { static NONCE: OnceLock<u64> = OnceLock::new(); *NONCE.get_or_init(|| { @@ -128,10 +127,8 @@ pub fn new_run_id() -> RunId { /// Allocates a fresh [`CheckpointId`] of the form `ckpt-<nonce>-<n>`, /// collision-free across process restarts (see [`process_nonce`]). /// -/// Restart-safety is essential here: checkpoint ids are the keys of the -/// parent-lineage spine that `prune`, `get_state_history`, and -/// `ResumeTarget::Checkpoint` all walk, so a duplicate id across a restart can -/// delete a live record's ancestor or resume the wrong checkpoint. +/// The process nonce avoids collisions when adapters persist checkpoint ids +/// across restarts. Session resume itself is handled by `tinyagents-session`. pub fn new_checkpoint_id() -> CheckpointId { CheckpointId(format!("ckpt-{}-{}", process_nonce(), next_seq())) } diff --git a/crates/tinyagents-harness/src/ids/types.rs b/crates/tinyagents-harness/src/ids/types.rs index 003915242..ac7cb483b 100644 --- a/crates/tinyagents-harness/src/ids/types.rs +++ b/crates/tinyagents-harness/src/ids/types.rs @@ -1,17 +1,16 @@ //! Identifier newtypes and lifecycle enums for the harness. //! -//! Every long-lived harness concept (runs, threads, calls, events, graph -//! nodes, checkpoints, interrupts) is keyed by a small, cheap-to-clone string +//! Every long-lived harness concept (runs, threads, calls, and events) is keyed by a small, cheap-to-clone string //! newtype. Wrapping the raw `String` in a distinct type prevents accidentally //! passing, say, a [`ThreadId`] where a [`RunId`] is expected. use serde::{Deserialize, Serialize}; /// Identifies a single harness run (one model call, one agent loop, or one -/// graph-node invocation of the harness). +/// nested harness invocation). /// -/// A `RunId` is the unit of recursion in TinyAgents: when a run spawns a -/// sub-agent or sub-graph, the child gets its own `RunId` while +/// A `RunId` identifies a unit of agent work: when a run spawns a +/// sub-agent, the child gets its own `RunId` while /// [`crate::events::HarnessRunStatus`] records that child's /// `parent_run_id` (the spawning run) and `root_run_id` (the top-level /// ancestor), so the full recursion tree is reconstructable from these ids. @@ -38,21 +37,22 @@ pub struct EventId(pub(crate) String); #[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)] pub struct ComponentId(pub(crate) String); -/// Identifies a state graph. +/// Identifies a graph-kind orchestration task. This opaque type is retained for +/// task descriptor compatibility; the workspace no longer provides a graph +/// execution runtime. #[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)] pub struct GraphId(pub(crate) String); -/// Identifies a node within a state graph. +/// Identifies an optional node label carried by an orchestration task record. #[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)] pub struct NodeId(pub(crate) String); -/// Identifies a unit of scheduled graph work (a node activation or a `Send` -/// fanout task) within a run. +/// Identifies a managed orchestration task, including sub-agent jobs and +/// persisted task records. /// -/// A `TaskId` distinguishes the individual recursive tasks a graph schedules — -/// for example each child task produced by a `Send` fanout — so a -/// `RecursionFrame` (in `tinyagents-graph`) can name the exact task a nested -/// call descends from, independent of the [`NodeId`] it ran. +/// `tinyagents-tasks` and `tinyagents-orchestration` use it to correlate task +/// lifecycle, steering, completion, and recovery records. It is independent of +/// the optional [`NodeId`] metadata on a task. #[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)] pub struct TaskId(pub(crate) String); @@ -61,7 +61,7 @@ pub struct TaskId(pub(crate) String); /// A `SessionId` names one host-driven session: a persistent namespace and /// capability boundary that a (possibly model-driven) orchestrator drives one /// cell at a time. Pairing it with a [`RunId`] lets session events and the child -/// model/tool/graph runs a session spawns be correlated back to the session +/// model/tool runs a session spawns be correlated back to the session /// that issued them. #[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)] pub struct SessionId(pub(crate) String); @@ -74,7 +74,8 @@ pub struct SessionId(pub(crate) String); #[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)] pub struct CellId(pub(crate) String); -/// Identifies a persisted graph checkpoint. +/// Identifies a persisted execution checkpoint. Retained for compatibility; +/// durable conversation identity and resume are owned by `tinyagents-session`. #[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)] pub struct CheckpointId(pub(crate) String); @@ -82,8 +83,7 @@ pub struct CheckpointId(pub(crate) String); #[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)] pub struct InterruptId(pub(crate) String); -/// Coarse lifecycle status shared by direct model calls, agent loops, and -/// graph-node harness invocations. +/// Coarse lifecycle status shared by direct model calls and agent loops. #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)] #[serde(rename_all = "snake_case")] pub enum ExecutionStatus { @@ -99,10 +99,7 @@ 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. + /// Stopped at a safe boundary after the active work completed. Drained, } diff --git a/crates/tinyagents-harness/src/lib.rs b/crates/tinyagents-harness/src/lib.rs index 25c79ae64..fb2624639 100644 --- a/crates/tinyagents-harness/src/lib.rs +++ b/crates/tinyagents-harness/src/lib.rs @@ -13,7 +13,8 @@ //! # Cargo features //! //! - `sqlite` — retained as a no-op for compatibility; durable stores are the `Driver*` types over `tinystoragedrivers`. -//! - `multimodal` — image/audio/binary content resolution ([`multimodal`]), +//! - `multimodal` — image/audio/binary content resolution (see the `multimodal` +//! module when the feature is enabled), //! pulling in `reqwest` and `flate2`. //! - `claude-code` — the Claude Code CLI and Claude Agent SDK provider //! adapters under [`providers`]. @@ -25,27 +26,10 @@ //! `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`] is exported for a host to build on, but [`agent_loop`] does -//! not call into it on its own — it is 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`]. -//! //! Persisting a thread's transcript across runs lives in the separate //! `tinyagents-session` crate: 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`, @@ -62,10 +46,12 @@ pub mod artifacts; pub mod cache; pub mod cancel; mod canonical_json; -pub mod capability; pub mod catalog; pub mod close; pub mod config; +pub use config::{ + CompactionSettings, CompactionStrategy, ContextConfig, DEFAULT_KEEP_RECENT_TOOL_RESULTS, +}; pub mod context; pub mod cost; pub mod dialect; @@ -73,14 +59,12 @@ pub mod driver_ids; pub mod error; pub mod events; mod finish_reason; -pub mod handoff; pub mod host; pub mod ids; pub mod limits; #[cfg(feature = "media")] pub mod media; pub mod middleware; -pub mod model_registry; #[cfg(feature = "multimodal")] pub mod multimodal; pub mod no_progress; @@ -89,14 +73,12 @@ pub mod oneshot; pub mod progress; pub mod prompt; pub mod providers; -pub mod retriever; pub mod retry; pub mod run_queue; pub mod runtime; pub mod steering; pub mod storage; pub mod store; -pub mod stream; pub mod structured; pub mod summarization; pub mod terminal; @@ -113,6 +95,9 @@ pub mod turn_run_error; pub mod turn_run_finalize; pub mod turn_runner; pub mod turn_stop; +/// Shared URL admission rules for outbound requests, including SSRF checks +/// and DNS address pinning. +pub mod url_guard; pub mod workspace; #[cfg(feature = "media")] @@ -131,14 +116,9 @@ 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::*; -pub use model_registry::{ModelRegistry, ModelSelection, ResolvedModelBinding}; pub use no_progress::{ ArgumentChurnDetector, CallGate, ClassifiedFailure, ClassifiedFailureTracker, DEFAULT_IDENTICAL_HALT_THRESHOLD, DEFAULT_REPEAT_CALL_THRESHOLD, diff --git a/crates/tinyagents-harness/src/media/mod_tests.rs b/crates/tinyagents-harness/src/media/mod_tests.rs index bf203beac..ab4da3205 100644 --- a/crates/tinyagents-harness/src/media/mod_tests.rs +++ b/crates/tinyagents-harness/src/media/mod_tests.rs @@ -465,13 +465,13 @@ async fn video_timeout_stays_resumable_through_agent_harness() { ); let mut harness: AgentHarness<()> = AgentHarness::new(); harness.with_tool_timeout_settings(ToolTimeoutSettings::new(1, 1, 1, 0)); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "mock", Arc::new(MockModel::with_responses(vec![ tool_call_response("call-1", GENERATE_VIDEO_TOOL_NAME, json!({ "prompt": "x" })), text_response("recovered"), ])), - ); + ))); harness.register_tool(Arc::new(video)); let run = harness diff --git a/crates/tinyagents-harness/src/middleware/library/context_loop_tests.rs b/crates/tinyagents-harness/src/middleware/library/context_loop_tests.rs index 8981e9a7d..e4e47a9c2 100644 --- a/crates/tinyagents-harness/src/middleware/library/context_loop_tests.rs +++ b/crates/tinyagents-harness/src/middleware/library/context_loop_tests.rs @@ -102,7 +102,10 @@ fn fixture_with( Box::new(summarizer), ))); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", Arc::clone(&model) as _); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + Arc::clone(&model) as _, + ))); harness.register_tool(Arc::new(FakeTool::returning("read", big_output()))); harness.push_middleware(mw.clone()); harness.push_model_middleware(mw.clone()); @@ -398,7 +401,10 @@ async fn a_provider_overflow_compacts_once_and_retries() { Box::new(summarizer), )); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", model.clone() as _); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + model.clone() as _, + ))); harness.register_tool(Arc::new(FakeTool::returning("read", big_output()))); harness.push_middleware(mw.clone()); harness.push_model_middleware(mw.clone()); diff --git a/crates/tinyagents-harness/src/middleware/library/context_task_state_tests.rs b/crates/tinyagents-harness/src/middleware/library/context_task_state_tests.rs index dfd805f60..dd5503f2a 100644 --- a/crates/tinyagents-harness/src/middleware/library/context_task_state_tests.rs +++ b/crates/tinyagents-harness/src/middleware/library/context_task_state_tests.rs @@ -57,7 +57,10 @@ async fn task_state_checkpoints_carry_the_task_and_files_across_compactions() { .with_keep_recent_tokens(400), ); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", Arc::clone(&model) as _); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + Arc::clone(&model) as _, + ))); harness.register_tool(Arc::new(FakeTool::returning("shell", "r".repeat(800)))); harness.push_middleware(mw.clone()); harness.push_model_middleware(mw.clone()); diff --git a/crates/tinyagents-harness/src/middleware/library/ephemeral_instructions/mod_tests.rs b/crates/tinyagents-harness/src/middleware/library/ephemeral_instructions/mod_tests.rs index 11f0faaa2..f9d3d355e 100644 --- a/crates/tinyagents-harness/src/middleware/library/ephemeral_instructions/mod_tests.rs +++ b/crates/tinyagents-harness/src/middleware/library/ephemeral_instructions/mod_tests.rs @@ -244,13 +244,13 @@ async fn native_tool_corrections_reach_the_next_request_without_entering_history let queue = EphemeralInstructionQueue::new(); let captured = Arc::new(Mutex::new(Vec::new())); let mut harness: crate::runtime::AgentHarness<()> = crate::runtime::AgentHarness::new(); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "mock", Arc::new(MockModel::with_responses(vec![ response, ModelResponse::assistant("done"), ])), - ); + ))); harness.register_tool(Arc::new(ReadTool)); harness.push_middleware(Arc::new(NativeProducer(queue.clone()))); harness.push_middleware(Arc::new(queue.middleware())); @@ -309,13 +309,13 @@ async fn model_control_jump_keeps_corrections_for_the_next_dispatched_request() let queue = EphemeralInstructionQueue::new(); let captured = Arc::new(Mutex::new(Vec::new())); let mut harness: crate::runtime::AgentHarness<()> = crate::runtime::AgentHarness::new(); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "mock", Arc::new(MockModel::with_responses(vec![ response, ModelResponse::assistant("done"), ])), - ); + ))); harness.register_tool(Arc::new(ReadTool)); harness.push_middleware(Arc::new(NativeProducer(queue.clone()))); harness.push_middleware(Arc::new(SkipSecondRequest( diff --git a/crates/tinyagents-harness/src/middleware/library/mod_tests.rs b/crates/tinyagents-harness/src/middleware/library/mod_tests.rs index e5143efde..0e8e8fe60 100644 --- a/crates/tinyagents-harness/src/middleware/library/mod_tests.rs +++ b/crates/tinyagents-harness/src/middleware/library/mod_tests.rs @@ -2024,6 +2024,24 @@ async fn plan_mode_allowlist_keeps_specific_tools_available() { .expect("allowlisted tool stays executable in plan mode"); } +#[test] +fn common_plan_mode_allowlist_covers_review_exit_and_host_planning_tools() { + let names = plan_mode_allowlist(); + for required in [ + "plan_exit", + "request_plan_review", + "todo", + "goal_get", + "goal_set", + "goal_complete", + "goal_pause", + "goal_resume", + "goal_clear", + ] { + assert!(names.contains(&required), "missing {required}"); + } +} + #[tokio::test] async fn plan_mode_treats_unclassified_tools_as_side_effecting() { // Fail-closed: a tool with no policy entry at all is hidden/denied in diff --git a/crates/tinyagents-harness/src/middleware/library/repeat_escalation_loop_tests.rs b/crates/tinyagents-harness/src/middleware/library/repeat_escalation_loop_tests.rs index 2595e8062..7b5ddeff4 100644 --- a/crates/tinyagents-harness/src/middleware/library/repeat_escalation_loop_tests.rs +++ b/crates/tinyagents-harness/src/middleware/library/repeat_escalation_loop_tests.rs @@ -66,12 +66,12 @@ async fn blocked_calls_never_execute_and_the_second_block_pauses_the_run() { runs: Mutex::new(0), }); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "mock", Arc::new(MockModel::with_responses( (0..12).map(repeat_call).collect::<Vec<_>>(), )), - ); + ))); harness.register_tool(tool.clone()); harness.push_middleware(Arc::new(RepeatProgressMiddleware::new( steering.clone(), @@ -137,12 +137,12 @@ async fn a_later_registered_after_tool_sees_the_marker_on_refused_calls() { runs: Mutex::new(0), }); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "mock", Arc::new(MockModel::with_responses( (0..12).map(repeat_call).collect::<Vec<_>>(), )), - ); + ))); harness.register_tool(tool); harness.push_middleware(Arc::new(RepeatProgressMiddleware::new( steering.clone(), @@ -171,12 +171,12 @@ async fn a_guard_halt_is_reported_as_a_halted_terminal_outcome() { let steering = SteeringHandle::allow_all(); let summary: HaltSummarySlot = Arc::new(Mutex::new(None)); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "mock", Arc::new(MockModel::with_responses( (0..12).map(repeat_call).collect::<Vec<_>>(), )), - ); + ))); harness.register_tool(Arc::new(CountingTool { runs: Mutex::new(0), })); diff --git a/crates/tinyagents-harness/src/middleware/library/types.rs b/crates/tinyagents-harness/src/middleware/library/types.rs index 7da300f24..9352e3d42 100644 --- a/crates/tinyagents-harness/src/middleware/library/types.rs +++ b/crates/tinyagents-harness/src/middleware/library/types.rs @@ -725,3 +725,19 @@ pub fn plan_mode_middleware( ) -> PlanModeMiddleware { PlanModeMiddleware::new(mode, policies) } + +/// Product-level tool names commonly required for plan-mode control. +/// Hosts may extend this list for their own safe planning tools. +pub const fn plan_mode_allowlist() -> &'static [&'static str] { + &[ + "plan_exit", + "request_plan_review", + "todo", + "goal_get", + "goal_set", + "goal_complete", + "goal_pause", + "goal_resume", + "goal_clear", + ] +} diff --git a/crates/tinyagents-harness/src/middleware/library/unmet_deliverable_tests.rs b/crates/tinyagents-harness/src/middleware/library/unmet_deliverable_tests.rs index bd4d37f2b..09ffedfc8 100644 --- a/crates/tinyagents-harness/src/middleware/library/unmet_deliverable_tests.rs +++ b/crates/tinyagents-harness/src/middleware/library/unmet_deliverable_tests.rs @@ -136,7 +136,7 @@ async fn an_unwritten_deliverable_holds_the_answer_once_and_permits_a_fix() { let request = format!("Do the work and write it to {}", missing.display()); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "mock", Arc::new(ScriptedModel::new(vec![ tool_round("c0", "writer"), @@ -145,7 +145,7 @@ async fn an_unwritten_deliverable_holds_the_answer_once_and_permits_a_fix() { tinyinference_llm::model::ModelResponse::assistant("written and answered".to_string()), tinyinference_llm::model::ModelResponse::assistant("must not be reached".to_string()), ])), - ); + ))); harness.register_tool(Arc::new(FakeTool::returning("writer", "ok"))); harness.with_policy(RunPolicy { limits: RunLimits::default() @@ -177,14 +177,14 @@ async fn an_unwritten_deliverable_holds_the_answer_once_and_permits_a_fix() { #[tokio::test] async fn a_request_naming_no_missing_file_is_left_alone() { let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "mock", Arc::new(ScriptedModel::new(vec![ tool_round("c0", "writer"), tinyinference_llm::model::ModelResponse::assistant("done".to_string()), tinyinference_llm::model::ModelResponse::assistant("must not be reached".to_string()), ])), - ); + ))); harness.register_tool(Arc::new(FakeTool::returning("writer", "ok"))); harness.with_policy(RunPolicy { limits: RunLimits::default().with_max_model_calls(20), @@ -213,7 +213,9 @@ async fn install_holds_a_final_answer_while_the_named_file_is_missing() { ModelResponse::assistant("written and answered".to_string()), ])); let mut root: AgentHarness<()> = AgentHarness::new(); - root.register_model("mock", script); + root.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", script, + ))); root.register_tool(Arc::new(FakeTool::returning("writer", "ok"))); root.with_policy(RunPolicy { limits: RunLimits::default() @@ -244,14 +246,14 @@ async fn candidates_come_from_the_current_turns_request() { let _ = std::fs::remove_file(&old); let _ = std::fs::remove_file(&new); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "mock", Arc::new(ScriptedModel::new(vec![ tool_round("c0", "writer"), tinyinference_llm::model::ModelResponse::assistant("first answer".to_string()), tinyinference_llm::model::ModelResponse::assistant("second answer".to_string()), ])), - ); + ))); harness.register_tool(Arc::new(FakeTool::returning("writer", "ok"))); harness.with_policy(RunPolicy { limits: RunLimits::default() @@ -418,14 +420,14 @@ async fn the_runs_state_is_dropped_when_the_turn_ends() { let _ = std::fs::remove_file(&missing); let middleware = Arc::new(UnmetDeliverableMiddleware::new(None)); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "mock", Arc::new(ScriptedModel::new(vec![ tool_round("c0", "writer"), tinyinference_llm::model::ModelResponse::assistant("first answer".to_string()), tinyinference_llm::model::ModelResponse::assistant("second answer".to_string()), ])), - ); + ))); harness.register_tool(Arc::new(FakeTool::returning("writer", "ok"))); harness.with_policy(RunPolicy { limits: RunLimits::default() diff --git a/crates/tinyagents-harness/src/middleware/library/verify_before_finish/middleware_tests.rs b/crates/tinyagents-harness/src/middleware/library/verify_before_finish/middleware_tests.rs index 56035f980..4fe85184f 100644 --- a/crates/tinyagents-harness/src/middleware/library/verify_before_finish/middleware_tests.rs +++ b/crates/tinyagents-harness/src/middleware/library/verify_before_finish/middleware_tests.rs @@ -62,7 +62,10 @@ async fn drive( ) -> (AgentRun, Vec<tinyinference_llm::model::ModelRequest>) { let model = Arc::new(ScriptedModel::new(responses)); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", model.clone()); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + model.clone(), + ))); harness.register_tool(Arc::new(FakeTool::returning("lookup", "found"))); harness.register_tool(Arc::new(FakeTool::returning("todo", "listed"))); harness.push_middleware(Arc::new(mw)); diff --git a/crates/tinyagents-harness/src/middleware/library/workload_tests.rs b/crates/tinyagents-harness/src/middleware/library/workload_tests.rs index 4f980df5b..90bee6683 100644 --- a/crates/tinyagents-harness/src/middleware/library/workload_tests.rs +++ b/crates/tinyagents-harness/src/middleware/library/workload_tests.rs @@ -693,15 +693,13 @@ async fn resolved_route_middleware_observes_stream_folded_terminal_response() { |context: &RouteContext, route| *context.0.lock().unwrap() = Some(route), ); let mut harness: AgentHarness<(), RouteContext> = AgentHarness::new(); - harness - .register_model( - "stream", - Arc::new(RouteRecordingModel::new( - Arc::new(ScriptedModel::replies(vec!["streamed"])), - route.clone(), - )), - ) - .set_default_model("stream"); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "stream", + Arc::new(RouteRecordingModel::new( + Arc::new(ScriptedModel::replies(vec!["streamed"])), + route.clone(), + )) as _, + ))); harness.push_model_middleware(Arc::new(middleware)); let context = RouteContext::default(); let slot = context.0.clone(); diff --git a/crates/tinyagents-harness/src/model_registry/README.md b/crates/tinyagents-harness/src/model_registry/README.md deleted file mode 100644 index 0f299d401..000000000 --- a/crates/tinyagents-harness/src/model_registry/README.md +++ /dev/null @@ -1,57 +0,0 @@ -# harness::model_registry - -Runtime-owned executable model registry and selection policy. - -## Why this exists - -Provider-neutral inference types and calls live in `tinyinference`; this -module owns only the harness-side concerns: registering executable models -under runtime names, and deterministically picking one for a given call. - -## Public surface - -- [`ModelRegistry<State>`] (`types.rs` + `mod.rs`) — name-keyed registry of - `Arc<dyn ChatModel<State>>` handles plus one designated default (the first - model registered, unless overridden with `set_default`). `register`, - `get`, `default_model`, `default_name`, `names` are the basic accessors; - `resolve`/`resolve_request` are the selection entry points. -- [`ModelSelection`] (`types.rs`) — the input policy for one resolution: - explicit override, reusable previous selection, priority-ordered runtime - hints, agent-level default, and a capability requirement, each field - corresponding to one precedence tier `resolve` walks. -- [`ResolvedModelBinding<State>`] (`types.rs`) — the resolution result: the - executable `model` handle plus `resolved` - (`tinyinference_llm::model::ResolvedModel`) metadata recording which name - and source won. -- [`model_eligible`] (`mod.rs`, crate-private) — the shared eligibility check - every candidate in `resolve` passes through: capability satisfaction plus, - unless `allow_retired` is set, usability. - -## Files - -| File | Role | -| ---------- | ---------------------------------------------------------------------- | -| `types.rs` | `ModelSelection`, `ModelRegistry`, `ResolvedModelBinding` definitions (plus their `Debug` impls). | -| `mod.rs` | `ModelRegistry` methods, the `resolve` selection algorithm, `model_eligible`/`model_satisfies`/`binding` helpers. | -| `mod_tests.rs` | Unit tests for registration, default handling, and each resolution precedence tier. | - -## Operational constraints - -- `resolve`'s precedence order is a contract, not an implementation detail: - explicit request override → reusable previous selection → priority-sorted - hints → agent default → registry-wide default. A source that names an - unregistered or ineligible model is skipped, not treated as a hard failure - — resolution only returns `None` once every tier is exhausted. Any change - to this order is a behavior change, not a docs change. -- Hint ordering is a **stable** sort: equal-priority hints keep the caller's - original order rather than an arbitrary one. Preserve that stability if the - sort is ever touched. -- `allow_retired` only ever *widens* eligibility (letting a retired model be - selected); it must never be used to route fresh, non-pinned selections to a - retired model — that path exists specifically for finishing an in-flight - conversation already pinned to one. -- `ModelRegistry::register` sets the default implicitly on the *first* - registration only; subsequent registrations never change it unless - `set_default` is called explicitly. A test or caller relying on - registration order to pick the default should register the intended - default first. diff --git a/crates/tinyagents-harness/src/model_registry/mod.rs b/crates/tinyagents-harness/src/model_registry/mod.rs deleted file mode 100644 index 75f5fdd53..000000000 --- a/crates/tinyagents-harness/src/model_registry/mod.rs +++ /dev/null @@ -1,237 +0,0 @@ -//! Runtime-owned executable model registry and selection policy. -//! -//! Provider-neutral inference types and calls live in `tinyinference`; this -//! module owns only name registration, fallback ordering, and runtime defaults. -//! -//! [`ModelRegistry`] maps runtime model names to executable -//! `Arc<dyn ChatModel<State>>` handles and holds one designated default. -//! [`ModelRegistry::resolve`] is the single selection algorithm: given a -//! [`ModelSelection`] it tries, in strict precedence order, an explicit -//! per-request override, a reusable previous selection, priority-sorted -//! runtime hints, the agent's own default, and finally the registry-wide -//! default — returning the first candidate that is both registered and -//! `model_eligible`. [`ModelRegistry::resolve_request`] is the common-case -//! entry point that builds a [`ModelSelection`] from one -//! [`ModelRequest`]. - -mod types; - -use std::sync::Arc; - -use tinyinference_llm::model::{ - CapabilitySet, ChatModel, ModelProfile, ModelRequest, ModelResolutionSource, ResolvedModel, -}; - -pub use types::*; - -impl<State: Send + Sync> ModelRegistry<State> { - /// Creates an empty registry. - pub fn new() -> Self { - Self { - models: std::collections::HashMap::new(), - default: None, - } - } - - /// Registers `model` under `name`. - pub fn register( - &mut self, - name: impl Into<String>, - model: Arc<dyn ChatModel<State>>, - ) -> &mut Self { - let name = name.into(); - if self.default.is_none() { - self.default = Some(name.clone()); - } - self.models.insert(name, model); - self - } - - /// Sets the runtime default model name. - pub fn set_default(&mut self, name: impl Into<String>) -> &mut Self { - self.default = Some(name.into()); - self - } - - /// Looks up a model by runtime name. - pub fn get(&self, name: &str) -> Option<Arc<dyn ChatModel<State>>> { - self.models.get(name).cloned() - } - - /// Returns the configured default model. - pub fn default_model(&self) -> Option<Arc<dyn ChatModel<State>>> { - self.default.as_deref().and_then(|name| self.get(name)) - } - - /// Returns the configured default name. - pub fn default_name(&self) -> Option<&str> { - self.default.as_deref() - } - - /// Returns registered names in sorted order. - pub fn names(&self) -> Vec<String> { - let mut names: Vec<String> = self.models.keys().cloned().collect(); - names.sort(); - names - } - - /// Resolves a model using override, previous state, hints, and defaults. - /// - /// Tries each source in strict precedence order and returns the first - /// candidate that is both registered and `model_eligible` (capability - /// match, and usable unless `selection.allow_retired`): - /// - /// 1. [`ModelSelection::requested`] — an explicit per-request override. - /// 2. [`ModelSelection::previous`], only when - /// [`ModelSelection::reuse_previous`] is set — a durable prior - /// selection the caller wants to keep using. - /// 3. [`ModelSelection::hints`], sorted by descending - /// [`ModelHint::priority`][tinyinference_llm::model::ModelHint::priority] - /// with ties broken by original order (stable) — runtime routing hints. - /// 4. [`ModelSelection::agent_default`] — the agent definition's own - /// default. - /// 5. The registry-wide default ([`ModelRegistry::default_name`]). - /// - /// A source that names a model this registry does not have, or that fails - /// the eligibility check, is skipped rather than treated as a hard - /// failure — resolution keeps falling through to the next source. - /// Returns `None` only when every source is exhausted. - pub fn resolve(&self, selection: ModelSelection) -> Option<ResolvedModelBinding<State>> { - let required = selection.required_capabilities.as_ref(); - let allow_retired = selection.allow_retired; - if let Some(requested) = selection.requested - && let Some(model) = self.get(&requested) - && model_eligible(model.as_ref(), required, allow_retired) - { - return Some(binding( - model, - requested.clone(), - Some(requested), - ModelResolutionSource::RequestOverride, - )); - } - if selection.reuse_previous - && let Some(previous) = selection.previous - && let Some(model) = self.get(&previous.name) - && model_eligible(model.as_ref(), required, allow_retired) - { - return Some(binding( - model, - previous.name, - previous.requested, - ModelResolutionSource::StateReuse, - )); - } - // Higher priority first; equal priorities keep their original - // (caller-supplied) order rather than an arbitrary sort order. - let mut hints: Vec<(usize, _)> = selection.hints.into_iter().enumerate().collect(); - hints.sort_by(|(left_index, left), (right_index, right)| { - right - .priority - .cmp(&left.priority) - .then_with(|| left_index.cmp(right_index)) - }); - for (_, hint) in hints { - if let Some(model) = self.get(&hint.model) - && model_eligible(model.as_ref(), required, allow_retired) - { - return Some(binding( - model, - hint.model.clone(), - Some(hint.model), - ModelResolutionSource::Hint, - )); - } - } - if let Some(agent_default) = selection.agent_default - && let Some(model) = self.get(&agent_default) - && model_eligible(model.as_ref(), required, allow_retired) - { - return Some(binding( - model, - agent_default.clone(), - Some(agent_default), - ModelResolutionSource::AgentDefault, - )); - } - let name = self.default_name()?.to_string(); - self.default_model() - .filter(|model| model_eligible(model.as_ref(), required, allow_retired)) - .map(|model| binding(model, name, None, ModelResolutionSource::RegistryDefault)) - } - - /// Resolves a model from one direct TinyInference request. - pub fn resolve_request( - &self, - request: &ModelRequest, - agent_default: Option<&str>, - previous: Option<ResolvedModel>, - ) -> Option<ResolvedModelBinding<State>> { - self.resolve(ModelSelection { - requested: request.model.clone(), - previous, - reuse_previous: request.reuse_previous_model, - hints: request.model_hints.clone(), - agent_default: agent_default.map(ToOwned::to_owned), - required_capabilities: request.required_capabilities.clone(), - allow_retired: false, - }) - } -} - -fn binding<State: Send + Sync>( - model: Arc<dyn ChatModel<State>>, - name: String, - requested: Option<String>, - source: ModelResolutionSource, -) -> ResolvedModelBinding<State> { - ResolvedModelBinding { - resolved: ResolvedModel { - name, - requested, - source, - }, - model, - } -} - -fn model_satisfies<State: Send + Sync>( - model: &dyn ChatModel<State>, - required: Option<&CapabilitySet>, -) -> bool { - match required { - None => true, - Some(required) if required == &CapabilitySet::default() => true, - Some(required) => model - .profile() - .is_some_and(|profile| profile.satisfies(required)), - } -} - -/// Whether `model` may be selected for a call requiring `required` capabilities. -/// -/// Composes two independent checks: capability satisfaction (via -/// [`ModelProfile::satisfies`]) and, unless `allow_retired` is set, -/// usability (a model with no declared profile is always usable; one with a -/// profile must pass [`ModelProfile::is_usable`]). `allow_retired` exists so -/// callers can explicitly resolve a retired model (e.g. to finish an -/// in-flight conversation pinned to it) without opening retired models to -/// fresh selection generally. -pub(crate) fn model_eligible<State: Send + Sync>( - model: &dyn ChatModel<State>, - required: Option<&CapabilitySet>, - allow_retired: bool, -) -> bool { - model_satisfies(model, required) - && (allow_retired || model.profile().is_none_or(ModelProfile::is_usable)) -} - -impl<State: Send + Sync> Default for ModelRegistry<State> { - fn default() -> Self { - Self::new() - } -} - -#[cfg(test)] -#[path = "mod_tests.rs"] -mod test; diff --git a/crates/tinyagents-harness/src/model_registry/mod_tests.rs b/crates/tinyagents-harness/src/model_registry/mod_tests.rs deleted file mode 100644 index f9a46eb35..000000000 --- a/crates/tinyagents-harness/src/model_registry/mod_tests.rs +++ /dev/null @@ -1,27 +0,0 @@ -//! Model registry unit tests. - -use async_trait::async_trait; -use tinyinference_llm::model::{ChatModel, ModelRequest, ModelResponse}; - -use super::*; - -struct StaticModel; - -#[async_trait] -impl ChatModel<()> for StaticModel { - async fn invoke( - &self, - _state: &(), - _request: ModelRequest, - ) -> tinyinference_llm::Result<ModelResponse> { - Ok(ModelResponse::assistant("ok")) - } -} - -#[test] -fn first_registration_becomes_default() { - let mut registry = ModelRegistry::new(); - registry.register("default", Arc::new(StaticModel)); - assert_eq!(registry.default_name(), Some("default")); - assert_eq!(registry.names(), vec!["default"]); -} diff --git a/crates/tinyagents-harness/src/model_registry/types.rs b/crates/tinyagents-harness/src/model_registry/types.rs deleted file mode 100644 index eb58d3c21..000000000 --- a/crates/tinyagents-harness/src/model_registry/types.rs +++ /dev/null @@ -1,75 +0,0 @@ -//! Runtime-owned model registration and selection types. -//! -//! See `mod.rs` for [`ModelRegistry::resolve`], the algorithm that consumes -//! [`ModelSelection`] and returns a [`ResolvedModelBinding`]. - -use std::collections::HashMap; -use std::sync::Arc; - -use tinyinference_llm::model::{CapabilitySet, ChatModel, ModelHint, ResolvedModel}; - -/// Input policy for resolving one registered model. -/// -/// Field order below matches the precedence [`ModelRegistry::resolve`] -/// applies: `requested` beats `previous` (when `reuse_previous`) beats -/// `hints` (priority-sorted) beats `agent_default` beats the registry-wide -/// default. -#[derive(Clone, Debug, Default, PartialEq)] -pub struct ModelSelection { - /// Explicit per-request model override; wins over every other source - /// when the named model is registered and eligible. - pub requested: Option<String>, - /// A durable prior selection (name, and how it was originally resolved). - /// Only consulted when [`Self::reuse_previous`] is set. - pub previous: Option<ResolvedModel>, - /// Whether [`Self::previous`] may be reused for this resolution. - pub reuse_previous: bool, - /// Ordered runtime hints, consulted in descending - /// [`ModelHint::priority`] order (original order breaks ties). - pub hints: Vec<ModelHint>, - /// The agent definition's own default model name, consulted after every - /// hint has been tried and rejected. - pub agent_default: Option<String>, - /// Capabilities a resolved model must satisfy, checked against every - /// candidate regardless of which source it came from. - pub required_capabilities: Option<CapabilitySet>, - /// Whether a retired model may still be selected (e.g. to finish an - /// in-flight conversation already pinned to it) rather than being - /// treated as ineligible. - pub allow_retired: bool, -} - -/// Name-keyed runtime registry of executable models. -pub struct ModelRegistry<State: Send + Sync> { - pub(crate) models: HashMap<String, Arc<dyn ChatModel<State>>>, - pub(crate) default: Option<String>, -} - -impl<State: Send + Sync> std::fmt::Debug for ModelRegistry<State> { - fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - let mut names: Vec<&str> = self.models.keys().map(String::as_str).collect(); - names.sort_unstable(); - formatter - .debug_struct("ModelRegistry") - .field("models", &names) - .field("default", &self.default) - .finish() - } -} - -/// Executable binding plus durable selection metadata. -pub struct ResolvedModelBinding<State: Send + Sync> { - /// Selected-model metadata. - pub resolved: ResolvedModel, - /// Executable model handle. - pub model: Arc<dyn ChatModel<State>>, -} - -impl<State: Send + Sync> std::fmt::Debug for ResolvedModelBinding<State> { - fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - formatter - .debug_struct("ResolvedModelBinding") - .field("resolved", &self.resolved) - .finish_non_exhaustive() - } -} diff --git a/crates/tinyagents-harness/src/multimodal/resolve.rs b/crates/tinyagents-harness/src/multimodal/resolve.rs index 8cbfe4878..ab5aea974 100644 --- a/crates/tinyagents-harness/src/multimodal/resolve.rs +++ b/crates/tinyagents-harness/src/multimodal/resolve.rs @@ -94,7 +94,7 @@ impl TextExtractor for NoTextExtractor { /// be used here: reqwest does not let a request override its DNS or redirect /// policy, so that client could reach a different address after validation. async fn guarded_remote_get(source: &str) -> std::result::Result<reqwest::Response, String> { - let validated = tinytools_std::url_guard::validate_url_with_dns_check(source, &[]) + let validated = crate::url_guard::validate_url_with_dns_check(source, &[]) .await .map_err(|error| error.to_string())?; let client = guarded_remote_client(&validated).map_err(|error| error.to_string())?; @@ -105,9 +105,7 @@ async fn guarded_remote_get(source: &str) -> std::result::Result<reqwest::Respon .map_err(|error| error.to_string()) } -fn guarded_remote_client( - validated: &tinytools_std::url_guard::ValidatedUrl, -) -> reqwest::Result<Client> { +fn guarded_remote_client(validated: &crate::url_guard::ValidatedUrl) -> reqwest::Result<Client> { Client::builder() .no_proxy() .redirect(reqwest::redirect::Policy::none()) diff --git a/crates/tinyagents-harness/src/multimodal/resolve_tests.rs b/crates/tinyagents-harness/src/multimodal/resolve_tests.rs index 8cb458730..b96828312 100644 --- a/crates/tinyagents-harness/src/multimodal/resolve_tests.rs +++ b/crates/tinyagents-harness/src/multimodal/resolve_tests.rs @@ -107,7 +107,7 @@ async fn guarded_client_does_not_follow_redirects() { let address = listener.local_addr().unwrap(); // The synthetic validated address tests the transport policy in isolation. // Production obtains this struct only from the DNS guard, which rejects it. - let validated = tinytools_std::url_guard::ValidatedUrl { + let validated = crate::url_guard::ValidatedUrl { url: format!("http://example.com:{}/image.png", address.port()), host: "example.com".to_string(), addrs: vec![address], diff --git a/crates/tinyagents-harness/src/observability/README.md b/crates/tinyagents-harness/src/observability/README.md index 7a587d728..c41a54893 100644 --- a/crates/tinyagents-harness/src/observability/README.md +++ b/crates/tinyagents-harness/src/observability/README.md @@ -53,12 +53,18 @@ surface and operational constraints. - `ProcessProfile` — the serializable wall-time, CPU utilization, and baseline/peak RSS report returned by `ProcessProfiler::finish`. - `LangfuseAuth`, `LangfuseClient`, `LangfuseTraceConfig` (re-exported from the - private `langfuse` submodule) — the Langfuse exporter used to ship harness - (and, via shared helpers, graph) traces to Langfuse. The exporter emits a + private `langfuse` submodule) — shared Langfuse credentials, HTTP transport, + trace config, and the exporter for durable harness observations. The agent + exporter emits a per-run span (`{trace}:run:{run_id}`, `agent`/`sub-agent`) and nests every generation, tool span, and event under it, so a run — and sub-agent recursion across separately-exported batches — surfaces as a tree, matching LangChain's callback run hierarchy. +- `spans::Span`, `spans::SpanState`, and `spans::SpanExporter` (behind the + `langfuse` feature) — a host-neutral model for already-mapped run, step, + generation, tool, and event observations. It builds the Langfuse ingestion + shape and delegates v4 OTLP conversion and transport to the shared helpers + and client. Hosts retain event mapping, identity, and content-capture policy. - `LangfuseScore` / `LangfuseScoreValue` — an evaluation score (numeric, categorical, or boolean; trace- or observation-scoped) attached to an exported trace via `LangfuseClient::create_score`, mirroring Langfuse's `createScore`. @@ -105,7 +111,8 @@ silently excluded from the rollup rather than reported with a bogus duration. | `types.rs` | Every public type: `AgentObservation`, journal/status-store traits and in-memory impls, sink structs, latency types. | | `profile.rs` | Scoped process CPU/RSS sampling for stress tests and host-specific workloads. | | `mod.rs` | Behavioral code: latency rollups, journal/store/sink impls. | -| `langfuse/` | `LangfuseClient` and payload helpers (`clean_nulls`, `iso_ms`) shared with `graph::observability::langfuse`, split into `mod.rs` (impl + helpers), `types.rs` (`LangfuseAuth`/`LangfuseClient`/`LangfuseTraceConfig`), and `mod_tests.rs`. | +| `langfuse/` | `LangfuseClient` and payload helpers (`clean_nulls`, `iso_ms`), split into `mod.rs` (impl + helpers), `types.rs` (`LangfuseAuth`/`LangfuseClient`/`LangfuseTraceConfig`), and `mod_tests.rs`. | +| `spans.rs`, `spans/` | Host-neutral span records and Langfuse batch/OTLP exporter; module README documents lifecycle and shape. | | `mod_tests.rs` | Unit tests (journal round-trips, redaction, latency rollups, sink fan-out). | ## Operational constraints diff --git a/crates/tinyagents-harness/src/observability/mod.rs b/crates/tinyagents-harness/src/observability/mod.rs index b8841a44e..a56c09478 100644 --- a/crates/tinyagents-harness/src/observability/mod.rs +++ b/crates/tinyagents-harness/src/observability/mod.rs @@ -28,6 +28,7 @@ mod journal_entry; #[cfg(feature = "langfuse")] mod langfuse; mod profile; +mod supplemental; mod types; mod worker; @@ -39,6 +40,8 @@ pub use physical_dispatch::{ }; pub mod reaper; pub mod replay; +#[cfg(feature = "langfuse")] +pub mod spans; pub mod status; pub(crate) mod status_record; pub mod trace_export; @@ -52,6 +55,7 @@ pub use langfuse::{ LangfuseAuth, LangfuseClient, LangfuseScore, LangfuseScoreValue, LangfuseTraceConfig, }; pub use profile::{ProcessProfile, ProcessProfiler, ProcessSnapshot}; +pub use supplemental::{SupplementalCallGuard, SupplementalCallTrace, SupplementalEnricher}; // 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")] diff --git a/crates/tinyagents-harness/src/observability/spans.rs b/crates/tinyagents-harness/src/observability/spans.rs new file mode 100644 index 000000000..d87e38cc1 --- /dev/null +++ b/crates/tinyagents-harness/src/observability/spans.rs @@ -0,0 +1,303 @@ +//! Host-neutral span records and Langfuse exporters. +//! +//! Hosts translate their own run, step, and span observations into [`Span`] +//! records. This module owns the common Langfuse envelope and OTLP conversion; +//! it does not collect events, choose capture policy, or own a transport. + +use std::collections::BTreeMap; + +use serde::{Deserialize, Serialize}; +use serde_json::{Value, json}; + +use crate::error::Result; +use crate::observability::trace_export::ingestion_otlp::ingestion_batch_otlp_requests; +use crate::observability::{LangfuseClient, LangfuseTraceConfig, clean_nulls, iso_ms}; + +/// Completion state for a span-like observation. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum SpanState { + /// The observation has no explicit completion result. + #[default] + Unset, + /// The observation completed successfully. + Ok, + /// The observation completed with an error. + Error, +} + +/// One host-mapped run, step, span, generation, or point event. +/// +/// `span_id` is the stable observation identity used for parent links. A host +/// may supply a separate `event_id` for Langfuse's ingestion-envelope id. +/// `run_id` and `step_id` are correlation coordinates; the exporter preserves +/// them in metadata without interpreting host-specific values. +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct Span { + /// Stable span or observation identity. + pub span_id: String, + /// Stable run identity that produced this observation. + pub run_id: String, + /// Root ancestor run id, when the host exports nested runs under one trace. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub root_run_id: Option<String>, + /// Optional host-defined step identity. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub step_id: Option<String>, + /// Optional parent observation identity. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub parent_span_id: Option<String>, + /// Optional distinct event id for the ingestion envelope. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub event_id: Option<String>, + /// Human-readable operation name. + pub name: String, + /// Langfuse ingestion type (`span-create`, `event-create`, or `generation-create`). + #[serde(default = "default_kind")] + pub kind: String, + /// Start time in Unix epoch milliseconds. + pub start_ms: u64, + /// End time in Unix epoch milliseconds; absent while still running. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub end_ms: Option<u64>, + /// Completion state. + #[serde(default)] + pub status: SpanState, + /// Optional status detail, commonly an error message. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub status_message: Option<String>, + /// Optional input payload, subject to the host's capture policy. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub input: Option<Value>, + /// Optional output payload, subject to the host's capture policy. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub output: Option<Value>, + /// Optional model token/cost usage details. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub usage: Option<Value>, + /// Additional host-owned observation metadata. + #[serde(default)] + pub metadata: BTreeMap<String, Value>, +} + +fn default_kind() -> String { + "span-create".to_string() +} + +impl Span { + /// Creates a timed span with an explicit stable id and parent. + pub fn new( + span_id: impl Into<String>, + run_id: impl Into<String>, + name: impl Into<String>, + start_ms: u64, + ) -> Self { + Self { + span_id: span_id.into(), + run_id: run_id.into(), + root_run_id: None, + step_id: None, + parent_span_id: None, + event_id: None, + name: name.into(), + kind: default_kind(), + start_ms, + end_ms: None, + status: SpanState::Unset, + status_message: None, + input: None, + output: None, + usage: None, + metadata: BTreeMap::new(), + } + } + + /// Sets the step coordinate. + pub fn with_step_id(mut self, step_id: impl Into<String>) -> Self { + self.step_id = Some(step_id.into()); + self + } + + /// Sets the root run identity used when deriving a default trace id. + pub fn with_root_run_id(mut self, root_run_id: impl Into<String>) -> Self { + self.root_run_id = Some(root_run_id.into()); + self + } + + /// Sets the parent span identity. + pub fn with_parent(mut self, parent_span_id: impl Into<String>) -> Self { + self.parent_span_id = Some(parent_span_id.into()); + self + } + + /// Sets the event id used by the Langfuse ingestion envelope. + pub fn with_event_id(mut self, event_id: impl Into<String>) -> Self { + self.event_id = Some(event_id.into()); + self + } + + /// Sets the end time and completion state. + pub fn finish(mut self, end_ms: u64, status: SpanState) -> Self { + self.end_ms = Some(end_ms); + self.status = status; + self + } + + /// Sets the input and output payloads. + pub fn with_io(mut self, input: Option<Value>, output: Option<Value>) -> Self { + self.input = input; + self.output = output; + self + } + + /// Sets model usage details. + pub fn with_usage(mut self, usage: Value) -> Self { + self.usage = Some(usage); + self + } + + /// Sets an explanatory status message, usually for a failed observation. + pub fn with_status_message(mut self, message: impl Into<String>) -> Self { + self.status_message = Some(message.into()); + self + } +} + +/// Builds a Langfuse ingestion payload from host-mapped spans. +/// +/// The first span's `run_id` supplies the default trace id. A trace-create +/// record is followed by one observation per span, preserving input order. +/// Empty input is rejected because a trace cannot be identified safely. +pub fn build_ingestion_batch(trace: LangfuseTraceConfig, spans: &[Span]) -> Result<Value> { + let first = spans.first().ok_or_else(|| { + crate::error::TinyAgentsError::Validation("at least one span is required".to_string()) + })?; + let trace_id = trace + .trace_id + .as_deref() + .filter(|id| !id.trim().is_empty()) + .or(first.root_run_id.as_deref()) + .unwrap_or(&first.run_id); + let trace_ts = iso_ms(first.start_ms); + let trace_event = json!({ + "id": format!("{trace_id}:trace"), + "timestamp": trace_ts, + "type": "trace-create", + "body": clean_nulls(json!({ + "id": trace_id, + "timestamp": trace_ts, + "name": trace.name, + "userId": trace.user_id, + "sessionId": trace.session_id, + "environment": trace.environment, + "release": trace.release, + "version": trace.version, + "tags": if trace.tags.is_empty() { Value::Null } else { json!(trace.tags) }, + "metadata": trace.metadata, + })), + }); + let mut batch = Vec::with_capacity(spans.len() + 1); + batch.push(trace_event); + for span in spans { + let start = iso_ms(span.start_ms); + let end = span.end_ms.map(iso_ms); + let event_id = span.event_id.as_deref().unwrap_or(&span.span_id); + let level = matches!(span.status, SpanState::Error).then_some("ERROR"); + let status_message = span.status_message.as_deref(); + let mut metadata = span.metadata.clone(); + metadata.insert("run_id".to_string(), json!(span.run_id)); + if let Some(root_run_id) = &span.root_run_id { + metadata.insert("root_run_id".to_string(), json!(root_run_id)); + } + if let Some(step_id) = &span.step_id { + metadata.insert("step_id".to_string(), json!(step_id)); + } + let body = clean_nulls(json!({ + "id": span.span_id, + "traceId": trace_id, + "parentObservationId": span.parent_span_id, + "name": span.name, + "startTime": start, + "endTime": end, + "level": level, + "statusMessage": status_message, + "input": span.input, + "output": span.output, + "usageDetails": span.usage, + "metadata": metadata, + })); + batch.push(json!({ + "id": event_id, + "timestamp": end.as_deref().unwrap_or(&start), + "type": span.kind, + "body": body, + })); + } + Ok(json!({ "batch": batch })) +} + +/// Langfuse exporter over the shared client; host identity and capture policy +/// remain with the caller that builds [`Span`] values. +#[derive(Clone, Debug)] +pub struct SpanExporter { + client: LangfuseClient, +} + +impl SpanExporter { + /// Wraps the shared Langfuse transport client. + pub fn new(client: LangfuseClient) -> Self { + Self { client } + } + + /// Returns the underlying client. + pub fn client(&self) -> &LangfuseClient { + &self.client + } + + /// Builds a legacy Langfuse ingestion batch without sending it. + pub fn build_ingestion_batch( + &self, + trace: LangfuseTraceConfig, + spans: &[Span], + ) -> Result<Value> { + build_ingestion_batch(trace, spans) + } + + /// Builds bounded Langfuse v4 OTLP requests without sending them. + pub fn build_otlp_requests( + &self, + trace: LangfuseTraceConfig, + spans: &[Span], + service: &str, + scope: &str, + ) -> Result<Vec<Value>> { + let payload = build_ingestion_batch(trace, spans)?; + ingestion_batch_otlp_requests(&payload, service, scope) + } + + /// Sends spans through Langfuse v4 OTLP, stopping on the first failure. + pub async fn send_otlp( + &self, + trace: LangfuseTraceConfig, + spans: &[Span], + service: &str, + scope: &str, + ) -> Result<Vec<Value>> { + let mut responses = Vec::new(); + for request in self.build_otlp_requests(trace, spans, service, scope)? { + responses.push(self.client.send_otlp(request).await?); + } + Ok(responses) + } + + /// Sends spans through the legacy Langfuse ingestion API. + pub async fn send(&self, trace: LangfuseTraceConfig, spans: &[Span]) -> Result<Value> { + self.client + .send_batch(build_ingestion_batch(trace, spans)?) + .await + } +} + +#[cfg(test)] +#[path = "spans_tests.rs"] +mod tests; diff --git a/crates/tinyagents-harness/src/observability/spans/README.md b/crates/tinyagents-harness/src/observability/spans/README.md new file mode 100644 index 000000000..335d0990f --- /dev/null +++ b/crates/tinyagents-harness/src/observability/spans/README.md @@ -0,0 +1,34 @@ +# harness::observability::spans + +This module is the host-neutral input model and Langfuse exporter for mapped +run, step, span, generation, and event observations. Hosts retain ownership of +collection, identity, content redaction/capture, and mapping their native events +into [`Span`](../spans.rs). The exporter owns the stable Langfuse JSON envelope +and delegates OTLP projection, request sizing, credentials, HTTP status and +partial-failure handling to the existing observability helpers and +[`LangfuseClient`](../langfuse/README.md). + +## Lifecycle and errors + +Export is pull-based: build a batch from the spans currently available and send +it at the host's chosen boundary. The exporter has no background queue, flush, +or close lifecycle. Ingestion sends one batch; OTLP sends bounded requests in +order and stops on the first error. A transport error is returned to the host. +The existing client reports failed HTTP responses, legacy `207` partial +failures, and OTLP `partialSuccess.rejectedSpans` through the shared error +path. + +## Shape + +`build_ingestion_batch` places one `trace-create` envelope first and preserves +the caller's span order afterward. Each observation keeps its supplied stable +id, parent id, timing, status, input/output, usage details, and metadata. Trace +ids default to `root_run_id`, then `run_id`, on the first span. Hosts should +pass a root span when they need the full run boundary represented in the +Langfuse tree. + +`build_otlp_requests` projects that same batch with +`trace_export::ingestion_otlp`; callers choose both the resource service name +and instrumentation scope. The projection and client preserve Langfuse v4 +parentage, payload, usage, trace attributes, request bounds, and failure +semantics. diff --git a/crates/tinyagents-harness/src/observability/spans_tests.rs b/crates/tinyagents-harness/src/observability/spans_tests.rs new file mode 100644 index 000000000..6251fa1bf --- /dev/null +++ b/crates/tinyagents-harness/src/observability/spans_tests.rs @@ -0,0 +1,165 @@ +use super::*; +use crate::observability::trace_export::ingestion_otlp::ingestion_batch_otlp_requests; + +const BASE: u64 = 1_704_067_200_000; + +fn span(id: &str, name: &str, start: u64, end: u64) -> Span { + Span::new(id, "run-1", name, BASE + start).finish(BASE + end, SpanState::Ok) +} + +#[test] +fn builds_nested_three_step_trace_with_payloads_and_usage() { + let spans = vec![ + Span::new("trace:run:run-1", "run-1", "demo", BASE) + .finish(BASE + 50, SpanState::Ok) + .with_io( + Some(json!({"input":"question"})), + Some(json!({"output":"answer"})), + ), + span("trace:step:1", "step 1", 1, 10).with_parent("trace:run:run-1"), + Span::new("trace:model:1", "run-1", "model call", BASE + 2) + .with_parent("trace:step:1") + .with_step_id("1") + .with_io( + Some(json!([{"role":"user","content":"question"}])), + Some(json!({"content":"thinking"})), + ) + .with_usage(json!({"input":12,"output":4,"total":16})) + .finish(BASE + 8, SpanState::Ok), + span("trace:step:2", "step 2", 11, 25).with_parent("trace:run:run-1"), + Span::new("trace:tool:2", "run-1", "search", BASE + 12) + .with_parent("trace:step:2") + .with_step_id("2") + .with_io( + Some(json!({"query":"question"})), + Some(json!({"result":"found"})), + ) + .finish(BASE + 20, SpanState::Ok), + span("trace:step:3", "step 3", 26, 50).with_parent("trace:run:run-1"), + Span::new("trace:model:3", "run-1", "final answer", BASE + 27) + .with_parent("trace:step:3") + .with_step_id("3") + .with_usage(json!({"input":20,"output":5,"total":25})) + .finish(BASE + 49, SpanState::Ok), + ]; + let batch = build_ingestion_batch( + LangfuseTraceConfig { + trace_id: Some("trace".to_string()), + name: Some("demo".to_string()), + session_id: Some("thread-1".to_string()), + ..Default::default() + }, + &spans, + ) + .unwrap(); + let events = batch["batch"].as_array().unwrap(); + assert_eq!(events.len(), 8); + assert_eq!(events[0]["type"], "trace-create"); + assert_eq!(events[1]["body"]["parentObservationId"], Value::Null); + assert_eq!(events[2]["body"]["parentObservationId"], "trace:run:run-1"); + assert_eq!(events[3]["body"]["parentObservationId"], "trace:step:1"); + assert_eq!(events[3]["body"]["input"][0]["content"], "question"); + assert_eq!(events[3]["body"]["usageDetails"]["total"], 16); + assert_eq!(events[7]["body"]["parentObservationId"], "trace:step:3"); + assert_eq!(events[7]["body"]["startTime"], "2024-01-01T00:00:00.027Z"); +} + +#[test] +fn builds_a_nested_run_and_three_step_span_batch() { + let spans = vec![ + Span::new("root-1:run:root-1", "run-1", "demo-graph", BASE + 1_000) + .with_event_id("evt-0") + .finish(BASE + 1_090, SpanState::Ok) + .with_io(None, None), + span("root-1:step:1", "step 1", 1_010, 1_030) + .with_parent("root-1:run:root-1") + .with_event_id("evt-3"), + span("root-1:node:a:1", "a", 1_015, 1_025) + .with_parent("root-1:step:1") + .with_event_id("evt-2"), + span("root-1:step:2", "step 2", 1_040, 1_060) + .with_parent("root-1:run:root-1") + .with_event_id("evt-6"), + span("root-1:node:b:2", "b", 1_045, 1_055) + .with_parent("root-1:step:2") + .with_event_id("evt-5"), + span("root-1:step:3", "step 3", 1_070, 1_090) + .with_parent("root-1:run:root-1") + .with_event_id("evt-9"), + span("root-1:node:c:3", "c", 1_075, 1_085) + .with_parent("root-1:step:3") + .with_event_id("evt-8"), + ]; + let trace = LangfuseTraceConfig { + trace_id: Some("root-1".to_string()), + name: Some("demo-graph".to_string()), + session_id: Some("thread-1".to_string()), + metadata: json!({"graph_id":"demo-graph","root_run_id":"root-1","health":{}}), + ..Default::default() + }; + let batch = build_ingestion_batch(trace, &spans).unwrap(); + let events = batch["batch"].as_array().unwrap(); + assert_eq!(events[0]["body"]["id"], "root-1"); + assert_eq!(events[0]["body"]["sessionId"], "thread-1"); + assert_eq!(events[1]["id"], "evt-0"); + assert_eq!(events[1]["body"]["id"], "root-1:run:root-1"); + assert_eq!( + events[2]["body"]["parentObservationId"], + "root-1:run:root-1" + ); + assert_eq!(events[3]["body"]["parentObservationId"], "root-1:step:1"); + assert_eq!(events[6]["body"]["id"], "root-1:step:3"); +} + +#[test] +fn failed_and_open_spans_keep_status_and_timing_shape() { + let spans = vec![ + Span::new("failed", "run-1", "operation", BASE) + .finish(BASE + 4, SpanState::Error) + .with_status_message("failed operation"), + Span::new("open", "run-1", "running", BASE + 5), + Span::new("ok-note", "run-1", "completed", BASE + 6) + .finish(BASE + 7, SpanState::Ok) + .with_status_message("completed with note"), + Span::new("open-note", "run-1", "still running", BASE + 8) + .with_status_message("in progress"), + ]; + let batch = build_ingestion_batch(LangfuseTraceConfig::default(), &spans).unwrap(); + let events = batch["batch"].as_array().unwrap(); + assert_eq!(events[1]["body"]["level"], "ERROR"); + assert_eq!(events[1]["body"]["statusMessage"], "failed operation"); + assert!(events[2]["body"].get("endTime").is_none()); + assert_eq!(events[3]["body"]["statusMessage"], "completed with note"); + assert!(events[3]["body"].get("level").is_none()); + assert_eq!(events[4]["body"]["statusMessage"], "in progress"); + assert!(events[4]["body"].get("level").is_none()); + assert!(events[4]["body"].get("endTime").is_none()); +} + +#[test] +fn otlp_conversion_keeps_parentage_payload_and_usage() { + let spans = vec![ + Span::new("root", "run-1", "root", BASE).finish(BASE + 10, SpanState::Ok), + Span::new("model", "run-1", "model", BASE + 1) + .with_parent("root") + .with_io( + Some(json!({"prompt":"hi"})), + Some(json!({"answer":"hello"})), + ) + .with_usage(json!({"input":2,"output":1})) + .finish(BASE + 9, SpanState::Ok), + ]; + let batch = build_ingestion_batch(LangfuseTraceConfig::default(), &spans).unwrap(); + let requests = + ingestion_batch_otlp_requests(&batch, "test-service", "test-service.spans").unwrap(); + let encoded = requests[0].to_string(); + assert!(encoded.contains("langfuse.observation.usage_details")); + assert!(encoded.contains("langfuse.observation.input")); + assert!(encoded.contains("langfuse.observation.output")); + assert!(encoded.contains("parentSpanId")); +} + +#[test] +fn rejects_empty_span_batch() { + assert!(build_ingestion_batch(LangfuseTraceConfig::default(), &[]).is_err()); +} diff --git a/crates/tinyagents-harness/src/observability/supplemental.rs b/crates/tinyagents-harness/src/observability/supplemental.rs new file mode 100644 index 000000000..6be6f8e92 --- /dev/null +++ b/crates/tinyagents-harness/src/observability/supplemental.rs @@ -0,0 +1,172 @@ +//! Event lifecycle for host model calls performed outside the agent loop. + +use std::sync::{ + Arc, Mutex, + atomic::{AtomicBool, AtomicUsize, Ordering}, +}; + +use tinyinference_llm::model::{ModelCallCorrelation, ModelRequest, ModelResponse}; + +use crate::{ + events::{AgentEvent, EventSink, GenerationEnrichment}, + ids::{CallId, RunId, now_ms}, +}; + +/// Host enrichment for an out-of-loop model call. +pub trait SupplementalEnricher: Send + Sync { + /// Return the actual model label and any validated provider pricing data. + fn enrich( + &self, + model: &str, + response: &ModelResponse, + ) -> (Option<String>, GenerationEnrichment); +} + +/// Correlated trace state shared by supplemental calls in one run. +pub struct SupplementalCallTrace { + events: EventSink, + run_id: RunId, + next_call: AtomicUsize, + actual_model: Mutex<Option<String>>, + cache_hit: AtomicBool, + enricher: Arc<dyn SupplementalEnricher>, +} + +impl SupplementalCallTrace { + /// Creates an event trace; call ids begin at one and remain unique per trace. + pub fn new( + events: EventSink, + run_id: RunId, + enricher: Arc<dyn SupplementalEnricher>, + ) -> Arc<Self> { + Arc::new(Self { + events, + run_id, + next_call: AtomicUsize::new(1), + actual_model: Mutex::new(None), + cache_hit: AtomicBool::new(false), + enricher, + }) + } + + /// The owning run identifier. + pub fn run_id(&self) -> &RunId { + &self.run_id + } + + /// The actual model from the latest completed supplemental call. + pub fn actual_model(&self) -> Option<String> { + self.actual_model + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .clone() + } + + /// Whether the latest completed supplemental call came from cache. + pub fn is_cache_hit(&self) -> bool { + self.cache_hit.load(Ordering::SeqCst) + } +} + +/// One in-flight supplemental call; dropping it records a single failure. +pub struct SupplementalCallGuard { + trace: Arc<SupplementalCallTrace>, + call_id: CallId, + model: String, + started_at_ms: u64, + input: serde_json::Value, + completed: bool, +} + +impl SupplementalCallGuard { + /// Starts a correlated call and emits its `ModelStarted` event. + pub fn start( + trace: &Arc<SupplementalCallTrace>, + model: &str, + request: &mut ModelRequest, + ) -> Self { + let n = trace.next_call.fetch_add(1, Ordering::Relaxed); + let call_id = CallId::new(format!("{}-host-supplemental-{n}", trace.run_id.as_str())); + let started_at_ms = now_ms(); + request.correlation = Some(ModelCallCorrelation::new( + trace.run_id.as_str(), + call_id.as_str(), + )); + trace.events.emit(AgentEvent::ModelStarted { + call_id: call_id.clone(), + model: model.to_owned(), + }); + Self { + trace: Arc::clone(trace), + call_id, + model: model.to_owned(), + started_at_ms, + input: serde_json::to_value(&request.messages).unwrap_or(serde_json::Value::Null), + completed: false, + } + } + + /// Records completion, usage and host enrichment exactly once. + pub fn complete(&mut self, response: &ModelResponse) { + let mut usage = response.usage.or(response.message.usage); + if response.served_from_cache { + self.trace.events.emit(AgentEvent::ModelResponseCached { + call_id: self.call_id.clone(), + }); + if let Some(usage) = &mut usage { + usage.charged_amount = None; + } + } + self.trace + .cache_hit + .store(response.served_from_cache, Ordering::SeqCst); + if let Some(usage) = usage { + self.trace.events.emit(AgentEvent::UsageRecorded { usage }); + } + self.trace.events.emit(AgentEvent::ModelCompleted { + call_id: self.call_id.clone(), + started_at_ms: Some(self.started_at_ms), + usage, + input: Some(self.input.clone()), + output: Some( + serde_json::to_value(&response.message).unwrap_or(serde_json::Value::Null), + ), + }); + let (actual_model, enrichment) = self.trace.enricher.enrich(&self.model, response); + if let Some(actual_model) = actual_model { + *self + .trace + .actual_model + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) = Some(actual_model); + } + if enrichment.provider_id.is_some() + || enrichment.pricing.is_some() + || enrichment.estimated_cost_usd.is_some() + { + self.trace.events.emit(AgentEvent::GenerationEnriched { + call_id: self.call_id.clone(), + enrichment, + }); + } + self.completed = true; + } +} + +impl Drop for SupplementalCallGuard { + fn drop(&mut self) { + if !self.completed { + self.trace.events.emit(AgentEvent::ModelFailed { + call_id: self.call_id.clone(), + model: self.model.clone(), + started_at_ms: Some(self.started_at_ms), + attempts: Some(1), + error: "supplemental model call failed".to_owned(), + }); + } + } +} + +#[cfg(test)] +#[path = "supplemental_tests.rs"] +mod tests; diff --git a/crates/tinyagents-harness/src/observability/supplemental_tests.rs b/crates/tinyagents-harness/src/observability/supplemental_tests.rs new file mode 100644 index 000000000..48709f83e --- /dev/null +++ b/crates/tinyagents-harness/src/observability/supplemental_tests.rs @@ -0,0 +1,110 @@ +use super::*; +use crate::events::{AgentEvent, EventListener, EventSink, RecordingListener}; +use crate::ids::RunId; +use tinyinference_llm::{ + model::{ModelRequest, ModelResponse}, + usage::{ChargedAmount, Usage}, +}; + +struct Enricher; + +impl SupplementalEnricher for Enricher { + fn enrich( + &self, + model: &str, + _response: &ModelResponse, + ) -> (Option<String>, GenerationEnrichment) { + ( + Some(format!("actual-{model}")), + GenerationEnrichment { + provider_id: Some("provider".into()), + ..Default::default() + }, + ) + } +} + +fn setup() -> (Arc<SupplementalCallTrace>, Arc<RecordingListener>) { + let events = EventSink::new(); + let listener = Arc::new(RecordingListener::new()); + events.subscribe(listener.clone() as Arc<dyn EventListener>); + ( + SupplementalCallTrace::new(events, RunId::new("run-1"), Arc::new(Enricher)), + listener, + ) +} + +#[test] +fn completed_supplemental_call_is_correlated_and_enriched_once() { + let (trace, listener) = setup(); + let mut request = + ModelRequest::new(vec![tinyinference_llm::message::Message::user("question")]); + let mut call = SupplementalCallGuard::start(&trace, "selected", &mut request); + let correlation = request.correlation.expect("correlation is attached"); + assert_eq!(correlation.run_id, "run-1"); + assert!(correlation.call_id.ends_with("-host-supplemental-1")); + call.complete(&ModelResponse::assistant("answer")); + drop(call); + + let events = listener + .events() + .into_iter() + .map(|record| record.event) + .collect::<Vec<_>>(); + assert!(matches!( + &events[..], + [ + AgentEvent::ModelStarted { .. }, + AgentEvent::ModelCompleted { .. }, + AgentEvent::GenerationEnriched { .. } + ] + )); + assert_eq!(trace.actual_model().as_deref(), Some("actual-selected")); + assert!(!trace.is_cache_hit()); +} + +#[test] +fn dropping_incomplete_supplemental_call_emits_one_failure() { + let (trace, listener) = setup(); + let mut request = ModelRequest::default(); + drop(SupplementalCallGuard::start( + &trace, + "selected", + &mut request, + )); + let events = listener.events(); + assert_eq!(events.len(), 2); + assert!(matches!(events[0].event, AgentEvent::ModelStarted { .. })); + assert!(matches!(events[1].event, AgentEvent::ModelFailed { .. })); +} + +#[test] +fn cached_completion_clears_charge_before_usage_event() { + let (trace, listener) = setup(); + let mut request = ModelRequest::default(); + let mut call = SupplementalCallGuard::start(&trace, "selected", &mut request); + let mut response = ModelResponse::assistant("answer"); + response.served_from_cache = true; + response.usage = Some(Usage { + charged_amount: Some(ChargedAmount::usd_micros(7)), + ..Default::default() + }); + call.complete(&response); + + let events = listener + .events() + .into_iter() + .map(|record| record.event) + .collect::<Vec<_>>(); + assert_eq!(events.len(), 5); + assert!(matches!(events[0], AgentEvent::ModelStarted { .. })); + assert!(matches!(events[1], AgentEvent::ModelResponseCached { .. })); + assert!( + matches!(&events[2], AgentEvent::UsageRecorded { usage } if usage.charged_amount.is_none()) + ); + assert!( + matches!(&events[3], AgentEvent::ModelCompleted { usage: Some(usage), .. } if usage.charged_amount.is_none()) + ); + assert!(matches!(events[4], AgentEvent::GenerationEnriched { .. })); + assert!(trace.is_cache_hit()); +} diff --git a/crates/tinyagents-harness/src/prompt/README.md b/crates/tinyagents-harness/src/prompt/README.md index 39ce3c878..8a3fa32e5 100644 --- a/crates/tinyagents-harness/src/prompt/README.md +++ b/crates/tinyagents-harness/src/prompt/README.md @@ -42,8 +42,8 @@ KV-cache/prompt-cache keeps hitting. reports byte offsets for reusable prefix boundaries. Pass the result to `PromptBuilder::push_tiered_system_prompt` to mark stable/context segments cacheable and the volatile tail non-cacheable. -- Free rendering helpers: `render_heading`, `render_optional_section`, - `render_tool_catalogue`, `render_retrieved_documents`. +- Free rendering helpers: `render_heading`, `render_optional_section`, and + `render_tool_catalogue`. - Model-guidance helpers: `needs_execution_discipline`, `execution_discipline_for`, and `execution_discipline_for_profile` select the exported `EXECUTION_DISCIPLINE` block by model family, falling back from an @@ -98,6 +98,3 @@ KV-cache/prompt-cache keeps hitting. with `cache_segments` populated; downstream, `crate::cache` and `crate::middleware`'s `PromptCacheGuardMiddleware` read those segments to detect prefix invalidation across calls. -- `render_retrieved_documents` accepts `crate::retriever::RetrievedDocument`, - connecting this module to the harness's retrieval surface without a hard - dependency in the other direction. diff --git a/crates/tinyagents-harness/src/prompt/mod.rs b/crates/tinyagents-harness/src/prompt/mod.rs index 070e0c034..3eeed1029 100644 --- a/crates/tinyagents-harness/src/prompt/mod.rs +++ b/crates/tinyagents-harness/src/prompt/mod.rs @@ -285,20 +285,6 @@ pub fn render_tool_catalogue(rows: &[(String, String)]) -> String { .join("\n") } -/// Renders generic retrieved documents in ranked order for context composition. -pub fn render_retrieved_documents(documents: &[crate::retriever::RetrievedDocument]) -> String { - documents - .iter() - .map(|document| { - format!( - "[{} score={}]\n{}", - document.id, document.score, document.content - ) - }) - .collect::<Vec<_>>() - .join("\n\n") -} - fn utf8_prefix(text: &str, max_bytes: usize) -> usize { if text.len() <= max_bytes { return text.len(); diff --git a/crates/tinyagents-harness/src/retriever/README.md b/crates/tinyagents-harness/src/retriever/README.md deleted file mode 100644 index 45d42663f..000000000 --- a/crates/tinyagents-harness/src/retriever/README.md +++ /dev/null @@ -1,44 +0,0 @@ -# harness::retriever - -Provider-neutral retrieval contracts used by agent context composition. - -## Why this exists - -The harness does not construct embedding providers, own credentials, or -decide memory/source scope — those are host concerns, typically backed -directly by `tinyinference_embeddings` or a host memory adapter. What the -harness needs is a narrow, stable seam: given a query, get back ranked -documents, and render them into a prompt section. This module is exactly that -seam and nothing more. - -## Public surface - -- [`Retriever`] — the provider- and storage-neutral async trait. A single - method, `retrieve(RetrievalRequest) -> Result<Vec<RetrievedDocument>>`, - returning at most `request.limit` documents in descending relevance order. -- [`RetrievalRequest`] — query text, result limit, opaque caller metadata, and - a [`crate::CancellationToken`] for cooperative cancellation of expensive - retrieval work. -- [`RetrievedDocument`] — one ranked result: id, model-facing content, score, - and source metadata (for citations / host projection). -- [`compose_retrieval_context`] — the one entry point most callers need: drives - a `Retriever` and renders the result as a [`crate::prompt::PromptSection`] - ready to insert into an agent's composed prompt. - -## Files - -| File | Role | -| --- | --- | -| `mod.rs` | `compose_retrieval_context`, the sole free function tying the trait to prompt composition. | -| `types.rs` | `RetrievalRequest`, `RetrievedDocument`, and the `Retriever` trait. | -| `mod_tests.rs` | A recording fake `Retriever` exercising cancellation, ordering, and metadata round-tripping. | - -## Operational constraints - -- Authorization and source scope are the host's responsibility *before* a - `RetrievalRequest` is constructed — the trait itself performs no filtering. -- `compose_retrieval_context` preserves whatever order the `Retriever` - returns; it does not re-rank or deduplicate. -- Cancellation is cooperative: an implementation should check - `RetrievalRequest::cancellation` and return - [`crate::TinyAgentsError::Cancelled`] promptly rather than ignoring it. diff --git a/crates/tinyagents-harness/src/retriever/mod.rs b/crates/tinyagents-harness/src/retriever/mod.rs deleted file mode 100644 index 69c141448..000000000 --- a/crates/tinyagents-harness/src/retriever/mod.rs +++ /dev/null @@ -1,37 +0,0 @@ -//! Generic retrieval contracts used by agent context composition. -//! -//! The harness does not construct embedding providers or assign memory scope. -//! Those are host concerns. It provides only the narrow request/result/trait -//! seam needed to inject ranked context into a prompt. Embedding callers use -//! the direct `tinyinference_embeddings` 0.3 types; no local embedding facade -//! is introduced here. -//! -//! [`compose_retrieval_context`] is the only entry point most callers need: it -//! drives a caller-supplied [`Retriever`] and renders the result as a -//! [`crate::prompt::PromptSection`] ready to insert into an agent's composed -//! prompt (see [`crate::prompt`]). - -mod types; - -pub use types::*; - -/// Retrieves ranked context and returns one caller-insertable prompt section. -/// -/// The caller owns source policy and where the returned section is placed; this -/// function only preserves retriever order and uses the request's cancellation -/// and limit unchanged. -pub async fn compose_retrieval_context( - retriever: &dyn Retriever, - request: RetrievalRequest, - section_name: impl Into<String>, -) -> crate::Result<crate::prompt::PromptSection> { - let documents = retriever.retrieve(request).await?; - Ok(crate::prompt::PromptSection::new( - section_name, - crate::prompt::render_retrieved_documents(&documents), - )) -} - -#[cfg(test)] -#[path = "mod_tests.rs"] -mod test; diff --git a/crates/tinyagents-harness/src/retriever/mod_tests.rs b/crates/tinyagents-harness/src/retriever/mod_tests.rs deleted file mode 100644 index d62363eb8..000000000 --- a/crates/tinyagents-harness/src/retriever/mod_tests.rs +++ /dev/null @@ -1,75 +0,0 @@ -//! Tests for the retrieval contract: a fake [`Retriever`] verifying -//! cancellation propagation, request/limit/metadata round-tripping, and -//! [`compose_retrieval_context`] rendering documents in retriever order. - -use std::sync::{Arc, Mutex}; - -use async_trait::async_trait; -use serde_json::json; - -use super::*; - -#[derive(Default)] -struct RecordingRetriever { - requests: Arc<Mutex<Vec<RetrievalRequest>>>, -} - -#[async_trait] -impl Retriever for RecordingRetriever { - async fn retrieve(&self, request: RetrievalRequest) -> crate::Result<Vec<RetrievedDocument>> { - if request.cancellation.is_cancelled() { - return Err(crate::TinyAgentsError::Cancelled); - } - self.requests.lock().unwrap().push(request.clone()); - Ok(vec![ - RetrievedDocument::new("first", "first result", 0.9) - .with_metadata(json!({"source": "memory"})), - RetrievedDocument::new("second", "second result", 0.2), - ] - .into_iter() - .take(request.limit) - .collect()) - } -} - -#[tokio::test] -async fn retrieval_contract_propagates_cancellation() { - let retriever = RecordingRetriever::default(); - let cancellation = crate::CancellationToken::new(); - cancellation.cancel(); - let result = retriever - .retrieve(RetrievalRequest::new("query", 2).with_cancellation(cancellation)) - .await; - assert!(matches!(result, Err(crate::TinyAgentsError::Cancelled))); -} - -#[tokio::test] -async fn context_composition_preserves_retriever_order_scores_and_metadata() { - let retriever = RecordingRetriever::default(); - let section = - compose_retrieval_context(&retriever, RetrievalRequest::new("query", 2), "retrieved") - .await - .unwrap(); - assert_eq!(section.name, "retrieved"); - assert_eq!( - section.content, - "[first score=0.9]\nfirst result\n\n[second score=0.2]\nsecond result" - ); -} - -#[tokio::test] -async fn retrieval_contract_preserves_limit_score_metadata_and_request_shape() { - let retriever = RecordingRetriever::default(); - let request = - RetrievalRequest::new("find the answer", 1).with_metadata(json!({"collection": "scoped"})); - - let docs = retriever.retrieve(request).await.unwrap(); - - assert_eq!(docs.len(), 1); - assert_eq!(docs[0].id, "first"); - assert_eq!(docs[0].score, 0.9); - assert_eq!(docs[0].metadata, json!({"source": "memory"})); - let recorded = retriever.requests.lock().unwrap(); - assert_eq!(recorded[0].query, "find the answer"); - assert_eq!(recorded[0].metadata, json!({"collection": "scoped"})); -} diff --git a/crates/tinyagents-harness/src/retriever/types.rs b/crates/tinyagents-harness/src/retriever/types.rs deleted file mode 100644 index 2cd893254..000000000 --- a/crates/tinyagents-harness/src/retriever/types.rs +++ /dev/null @@ -1,92 +0,0 @@ -//! Provider-neutral retrieval contracts. -//! -//! Embedding provider construction belongs to the host or directly to -//! `tinyinference-embeddings`. This module only names the data an agent needs -//! after a host has decided what it is permitted to retrieve. - -use async_trait::async_trait; -use serde_json::Value; - -use crate::CancellationToken; -use crate::Result; - -/// A retrieval request supplied by context composition. -#[derive(Clone, Debug)] -pub struct RetrievalRequest { - /// Query text. Hosts are responsible for authorization and source scope - /// before creating this request. - pub query: String, - /// Maximum number of ranked documents to return. - pub limit: usize, - /// Optional caller metadata, preserved for a host adapter but never - /// interpreted by the generic contract. - pub metadata: Value, - /// Cooperative cancellation signal for expensive retrieval work. - pub cancellation: CancellationToken, -} - -impl RetrievalRequest { - /// Creates a request with no host-specific metadata. - pub fn new(query: impl Into<String>, limit: usize) -> Self { - Self { - query: query.into(), - limit, - metadata: Value::Null, - cancellation: CancellationToken::new(), - } - } - - /// Attaches opaque host metadata. - pub fn with_metadata(mut self, metadata: Value) -> Self { - self.metadata = metadata; - self - } - - /// Uses the caller's cancellation tree for retrieval work. - pub fn with_cancellation(mut self, cancellation: CancellationToken) -> Self { - self.cancellation = cancellation; - self - } -} - -/// One ranked document returned by a [`Retriever`]. -#[derive(Clone, Debug, PartialEq)] -pub struct RetrievedDocument { - /// Stable document identifier, used for host-owned deduplication. - pub id: String, - /// Model-facing document text. The host decides redaction and byte caps. - pub content: String, - /// Relevance score in the retriever's documented ranking scale. - pub score: f32, - /// Source metadata retained for citations and host projection. - pub metadata: Value, -} - -impl RetrievedDocument { - /// Creates a retrieval result with no source metadata. - pub fn new(id: impl Into<String>, content: impl Into<String>, score: f32) -> Self { - Self { - id: id.into(), - content: content.into(), - score, - metadata: Value::Null, - } - } - - /// Attaches source metadata. - pub fn with_metadata(mut self, metadata: Value) -> Self { - self.metadata = metadata; - self - } -} - -/// Provider- and storage-neutral ranked retrieval interface. -/// -/// An implementation may use `tinyinference_embeddings` directly, a remote -/// vector store, or a host memory adapter. Credentials, provider settings, -/// file access, and collection scope deliberately remain outside this trait. -#[async_trait] -pub trait Retriever: Send + Sync { - /// Returns at most `request.limit` documents in descending relevance order. - async fn retrieve(&self, request: RetrievalRequest) -> Result<Vec<RetrievedDocument>>; -} diff --git a/crates/tinyagents-harness/src/run_queue/forwarder_tests.rs b/crates/tinyagents-harness/src/run_queue/forwarder_tests.rs index 38c977611..05f79221b 100644 --- a/crates/tinyagents-harness/src/run_queue/forwarder_tests.rs +++ b/crates/tinyagents-harness/src/run_queue/forwarder_tests.rs @@ -227,9 +227,10 @@ async fn collect_reaches_the_next_model_boundary_as_additional_context() { forward_collects(&queue, &handle, "thread-test", &sink).await; let model = Arc::new(ScriptedModel::replies(vec!["done"])); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness - .register_model("scripted", model.clone()) - .set_default_model("scripted"); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "scripted", + model.clone(), + ))); harness .invoke_in_context( &(), @@ -431,6 +432,61 @@ async fn native_message_guard_requeues_full_residual_and_leaves_non_user_items() )); } +#[tokio::test] +async fn native_forwarder_requeues_residual_without_stripping_guarded_queue_items() { + use crate::run_queue::{QueueAdmission, QueueAdmissionError, QueueAdmissionPermit}; + use std::sync::atomic::{AtomicBool, Ordering}; + + #[derive(Debug)] + struct Admission(AtomicBool); + impl QueueAdmission for Admission { + fn acquire(&self) -> Result<Arc<dyn QueueAdmissionPermit>, QueueAdmissionError> { + self.0 + .load(Ordering::SeqCst) + .then(|| Arc::new(()) as Arc<dyn QueueAdmissionPermit>) + .ok_or(QueueAdmissionError::Rejected) + } + } + + let queue = Arc::new(RunQueue::<TaMessage>::new()); + let admission = Arc::new(Admission(AtomicBool::new(true))); + queue + .push_guarded( + QueueLane::Steer, + TaMessage::user("lifecycle guarded"), + admission, + ) + .await; + queue + .push(QueueLane::Steer, TaMessage::user("forwarded then restored")) + .await; + let handle = SteeringHandle::allow_all(); + let (sink, _) = recorder(); + let guard = MessageSteeringForwarderGuard::new( + handle.clone(), + Some(queue.clone()), + None, + "native".into(), + sink.clone(), + ); + + forward_message_steers(&queue, &handle, "native", &sink).await; + drop(guard); + + let mut restored = queue + .take_admitted(QueueLane::Steer, crate::run_queue::QueueMode::All) + .await + .into_iter() + .map(|item| item.value.text()) + .collect::<Vec<_>>(); + restored.sort_unstable(); + assert_eq!( + restored, + vec!["forwarded then restored", "lifecycle guarded"], + "legacy forwarding may process ordinary items but must retain guarded payloads and their admission" + ); +} + #[tokio::test] async fn native_message_residual_round_trips_without_empty_prefix_blocks() { let queue = Arc::new(RunQueue::<TaMessage>::new()); diff --git a/crates/tinyagents-harness/src/run_queue/mod.rs b/crates/tinyagents-harness/src/run_queue/mod.rs index 9d52d304e..2b26192c3 100644 --- a/crates/tinyagents-harness/src/run_queue/mod.rs +++ b/crates/tinyagents-harness/src/run_queue/mod.rs @@ -44,7 +44,10 @@ pub use forwarder::{ forward_message_steers, forward_steers, lane_prefix, }; -pub use types::{QueueLane, QueueMode, QueueStatus, RunQueueHandle}; +pub use types::{ + QueueAdmission, QueueAdmissionError, QueueAdmissionPermit, QueueLane, QueueMode, QueueStatus, + RunQueueHandle, +}; /// Thread-safe FIFO queue split into steer, follow-up, and collect lanes. /// @@ -63,11 +66,51 @@ pub struct RunQueue<T> { struct RunQueueInner<T> { /// Immediate steering messages, consumed as soon as the loop reaches a /// safe iteration boundary. - steers: Vec<T>, + steers: Vec<Queued<T>>, /// Deferred follow-up work, consumed once the current run settles. - followups: Vec<T>, + followups: Vec<Queued<T>>, /// Collected context (e.g. observations) accumulated for later use. - collects: Vec<T>, + collects: Vec<Queued<T>>, +} + +#[derive(Debug)] +struct Queued<T> { + value: T, + admission: Option<std::sync::Arc<dyn QueueAdmission>>, +} + +enum AdmissionDecision<T> { + Accepted(AdmittedQueueItem<T>), + Retry(Queued<T>), + Rejected, +} + +impl<T> Queued<T> { + fn admitted(value: T) -> Self { + Self { + value, + admission: None, + } + } + + fn guarded(value: T, admission: std::sync::Arc<dyn QueueAdmission>) -> Self { + Self { + value, + admission: Some(admission), + } + } + + fn is_guarded(&self) -> bool { + self.admission.is_some() + } +} + +/// One item transferred from a run queue with its lifecycle permit retained +/// until the caller completes the queue-to-turn handoff. +pub struct AdmittedQueueItem<T> { + /// The queued payload. + pub value: T, + _permit: Option<std::sync::Arc<dyn types::QueueAdmissionPermit>>, } impl<T> RunQueue<T> { @@ -97,6 +140,160 @@ impl<T> RunQueue<T> { /// Appends `item` to `lane` without awaiting: the synchronous form of /// [`Self::push`], for `Drop` impls and other non-async callers. pub fn push_now(&self, lane: QueueLane, item: T) { + self.push_entry_now(lane, Queued::admitted(item)); + } + + /// Appends `item` with a retained capability checked when the agent loop + /// consumes it. Rejected items are discarded without entering the model + /// transcript. + pub async fn push_guarded( + &self, + lane: QueueLane, + item: T, + admission: std::sync::Arc<dyn QueueAdmission>, + ) { + self.push_guarded_now(lane, item, admission); + } + + /// Synchronous form of [`Self::push_guarded`]. + pub fn push_guarded_now( + &self, + lane: QueueLane, + item: T, + admission: std::sync::Arc<dyn QueueAdmission>, + ) { + self.push_entry_now(lane, Queued::guarded(item, admission)); + } + + /// Takes items with their lifecycle permits. The permits are acquired + /// after releasing the queue mutex, so durable admission never blocks + /// queue producers while performing storage I/O. + pub async fn take_admitted( + &self, + lane: QueueLane, + mode: QueueMode, + ) -> Vec<AdmittedQueueItem<T>> { + if mode == QueueMode::OneAtATime { + let mut permits = std::collections::HashMap::new(); + let count = { + let inner = self.lock(); + match lane { + QueueLane::Steer => inner.steers.len(), + QueueLane::Followup => inner.followups.len(), + QueueLane::Collect => inner.collects.len(), + } + }; + let mut deferred = Vec::new(); + for _ in 0..count { + let item = { + let mut inner = self.lock(); + let items = match lane { + QueueLane::Steer => &mut inner.steers, + QueueLane::Followup => &mut inner.followups, + QueueLane::Collect => &mut inner.collects, + }; + if items.is_empty() { + break; + } + items.remove(0) + }; + match Self::admit_item(item, &mut permits) { + AdmissionDecision::Accepted(item) => { + self.requeue_front_now(lane, deferred); + return vec![item]; + } + AdmissionDecision::Retry(item) => deferred.push(item), + AdmissionDecision::Rejected => {} + } + } + self.requeue_front_now(lane, deferred); + return Vec::new(); + } + let queued = { + let mut inner = self.lock(); + let items = match lane { + QueueLane::Steer => &mut inner.steers, + QueueLane::Followup => &mut inner.followups, + QueueLane::Collect => &mut inner.collects, + }; + std::mem::take(items) + }; + let mut permits = std::collections::HashMap::new(); + let mut admitted = Vec::new(); + let mut deferred = Vec::new(); + for item in queued { + match Self::admit_item(item, &mut permits) { + AdmissionDecision::Accepted(item) => admitted.push(item), + AdmissionDecision::Retry(item) => deferred.push(item), + AdmissionDecision::Rejected => {} + } + } + self.requeue_front_now(lane, deferred); + admitted + } + + fn admit_item( + item: Queued<T>, + permits: &mut std::collections::HashMap< + String, + std::sync::Arc<dyn types::QueueAdmissionPermit>, + >, + ) -> AdmissionDecision<T> { + let Queued { value, admission } = item; + let permit = match admission.as_ref() { + Some(capability) => { + let key = capability.batch_key(); + if let Some(key) = key.as_ref() { + match permits.get(key) { + Some(permit) => Some(std::sync::Arc::clone(permit)), + None => { + let permit = match capability.acquire_for(Some(key)) { + Ok(permit) => permit, + Err(types::QueueAdmissionError::Retry) => { + return AdmissionDecision::Retry(Queued { value, admission }); + } + Err(types::QueueAdmissionError::Rejected) => { + return AdmissionDecision::Rejected; + } + }; + permits.insert(key.clone(), std::sync::Arc::clone(&permit)); + Some(permit) + } + } + } else { + match capability.acquire_for(None) { + Ok(permit) => Some(permit), + Err(types::QueueAdmissionError::Retry) => { + return AdmissionDecision::Retry(Queued { value, admission }); + } + Err(types::QueueAdmissionError::Rejected) => { + return AdmissionDecision::Rejected; + } + } + } + } + None => None, + }; + AdmissionDecision::Accepted(AdmittedQueueItem { + value, + _permit: permit, + }) + } + + fn requeue_front_now(&self, lane: QueueLane, items: Vec<Queued<T>>) { + if items.is_empty() { + return; + } + let mut inner = self.lock(); + let target = match lane { + QueueLane::Steer => &mut inner.steers, + QueueLane::Followup => &mut inner.followups, + QueueLane::Collect => &mut inner.collects, + }; + target.splice(0..0, items); + } + + fn push_entry_now(&self, lane: QueueLane, item: Queued<T>) { let mut inner = self.lock(); match lane { QueueLane::Steer => inner.steers.push(item), @@ -118,7 +315,7 @@ impl<T> RunQueue<T> { QueueLane::Followup => &mut inner.followups, QueueLane::Collect => &mut inner.collects, }; - target.splice(0..0, items); + target.splice(0..0, items.into_iter().map(Queued::admitted)); } /// Drains one lane in FIFO order. @@ -130,11 +327,20 @@ impl<T> RunQueue<T> { /// of [`Self::drain`]. pub fn drain_now(&self, lane: QueueLane) -> Vec<T> { let mut inner = self.lock(); - match lane { + let items = match lane { QueueLane::Steer => std::mem::take(&mut inner.steers), QueueLane::Followup => std::mem::take(&mut inner.followups), QueueLane::Collect => std::mem::take(&mut inner.collects), - } + }; + let (unguarded, guarded): (Vec<_>, Vec<_>) = + items.into_iter().partition(|item| !item.is_guarded()); + let target = match lane { + QueueLane::Steer => &mut inner.steers, + QueueLane::Followup => &mut inner.followups, + QueueLane::Collect => &mut inner.collects, + }; + *target = guarded; + unguarded.into_iter().map(|item| item.value).collect() } /// Drains matching items from one lane in FIFO order and leaves all other @@ -151,8 +357,12 @@ impl<T> RunQueue<T> { let mut drained = Vec::new(); let mut remaining = Vec::with_capacity(items.len()); for item in std::mem::take(items) { - if pred(&item) { - drained.push(item); + if item.is_guarded() { + remaining.push(item); + continue; + } + if pred(&item.value) { + drained.push(item.value); } else { remaining.push(item); } @@ -174,11 +384,10 @@ impl<T> RunQueue<T> { QueueLane::Followup => &mut inner.followups, QueueLane::Collect => &mut inner.collects, }; - if items.is_empty() { - Vec::new() - } else { - vec![items.remove(0)] - } + let Some(index) = items.iter().position(|item| !item.is_guarded()) else { + return Vec::new(); + }; + vec![items.remove(index).value] } } } @@ -225,22 +434,19 @@ impl<T> RunQueue<T> { inner .steers .iter() - .cloned() - .map(|item| (QueueLane::Steer, item)), + .map(|item| (QueueLane::Steer, item.value.clone())), ); items.extend( inner .followups .iter() - .cloned() - .map(|item| (QueueLane::Followup, item)), + .map(|item| (QueueLane::Followup, item.value.clone())), ); items.extend( inner .collects .iter() - .cloned() - .map(|item| (QueueLane::Collect, item)), + .map(|item| (QueueLane::Collect, item.value.clone())), ); items } @@ -254,9 +460,9 @@ impl<T> RunQueue<T> { pub async fn remove_where(&self, pred: impl Fn(&T) -> bool) -> usize { let mut inner = self.lock(); let before = inner.steers.len() + inner.followups.len() + inner.collects.len(); - inner.steers.retain(|item| !pred(item)); - inner.followups.retain(|item| !pred(item)); - inner.collects.retain(|item| !pred(item)); + inner.steers.retain(|item| !pred(&item.value)); + inner.followups.retain(|item| !pred(&item.value)); + inner.collects.retain(|item| !pred(&item.value)); let after = inner.steers.len() + inner.followups.len() + inner.collects.len(); before - after } diff --git a/crates/tinyagents-harness/src/run_queue/mod_tests.rs b/crates/tinyagents-harness/src/run_queue/mod_tests.rs index d27735230..079989555 100644 --- a/crates/tinyagents-harness/src/run_queue/mod_tests.rs +++ b/crates/tinyagents-harness/src/run_queue/mod_tests.rs @@ -2,6 +2,58 @@ //! snapshots, `clear`, and lane independence. use super::*; +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; + +#[derive(Debug)] +struct MutableAdmission(AtomicBool); + +impl QueueAdmission for MutableAdmission { + fn acquire(&self) -> Result<Arc<dyn QueueAdmissionPermit>, QueueAdmissionError> { + self.0 + .load(Ordering::SeqCst) + .then(|| Arc::new(()) as Arc<dyn QueueAdmissionPermit>) + .ok_or(QueueAdmissionError::Rejected) + } +} + +#[derive(Debug)] +struct RetryOnce(AtomicBool); + +impl QueueAdmission for RetryOnce { + fn acquire(&self) -> Result<Arc<dyn QueueAdmissionPermit>, QueueAdmissionError> { + if self.0.swap(false, Ordering::SeqCst) { + Err(QueueAdmissionError::Retry) + } else { + Ok(Arc::new(())) + } + } +} + +#[tokio::test] +async fn transient_admission_result_retains_item_for_the_next_boundary() { + let queue = RunQueue::new(); + queue + .push_guarded( + QueueLane::Followup, + "recover after migration", + Arc::new(RetryOnce(AtomicBool::new(true))), + ) + .await; + + assert!( + queue + .take_admitted(QueueLane::Followup, QueueMode::All) + .await + .is_empty() + ); + assert_eq!(queue.status().await.followups, 1); + let retried = queue + .take_admitted(QueueLane::Followup, QueueMode::All) + .await; + assert_eq!(retried.len(), 1); + assert_eq!(retried[0].value, "recover after migration"); +} #[tokio::test] async fn new_queue_is_empty() { @@ -81,6 +133,53 @@ async fn take_one_at_a_time_pops_only_the_oldest_item() { ); } +#[tokio::test] +async fn take_discards_items_whose_retained_admission_was_revoked() { + let queue = RunQueue::new(); + let admission = Arc::new(MutableAdmission(AtomicBool::new(true))); + queue + .push_guarded(QueueLane::Followup, "child result", admission.clone()) + .await; + queue.push(QueueLane::Followup, "ordinary follow-up").await; + admission.0.store(false, Ordering::SeqCst); + + assert_eq!( + queue.take(QueueLane::Followup, QueueMode::All).await, + vec!["ordinary follow-up"], + "legacy extraction must not transfer a guarded payload without its permit" + ); + assert_eq!( + queue + .take_admitted(QueueLane::Followup, QueueMode::All) + .await + .into_iter() + .map(|item| item.value) + .collect::<Vec<_>>(), + Vec::<&str>::new(), + "a revoked item is rejected at the lifecycle admission boundary" + ); +} + +#[tokio::test] +async fn one_at_a_time_skips_a_revoked_oldest_item_without_blocking_the_lane() { + let queue = RunQueue::new(); + let admission = Arc::new(MutableAdmission(AtomicBool::new(false))); + queue + .push_guarded(QueueLane::Steer, "revoked", admission) + .await; + queue.push(QueueLane::Steer, "current").await; + + assert_eq!( + queue + .take_admitted(QueueLane::Steer, QueueMode::OneAtATime) + .await + .into_iter() + .map(|item| item.value) + .collect::<Vec<_>>(), + vec!["current"] + ); +} + #[tokio::test] async fn take_all_drains_the_whole_lane() { let queue = RunQueue::new(); diff --git a/crates/tinyagents-harness/src/run_queue/types.rs b/crates/tinyagents-harness/src/run_queue/types.rs index d9b5be69e..f45eeb6db 100644 --- a/crates/tinyagents-harness/src/run_queue/types.rs +++ b/crates/tinyagents-harness/src/run_queue/types.rs @@ -15,6 +15,47 @@ use tinyinference_llm::message::Message; /// [`Message::user`] (or [`Message::system`]) — the host chooses the role. pub type RunQueueHandle = Arc<super::RunQueue<Message>>; +/// A retained capability checked when a queued item is consumed by an agent +/// turn. Hosts can attach durable lifecycle admission to messages without +/// changing the transcript payload. +pub trait QueueAdmission: std::fmt::Debug + Send + Sync { + /// Acquires a short-lived permit that linearizes this item with lifecycle + /// closure. Implementations must return an error when the lifecycle no + /// longer admits the item. The permit is held only while the harness + /// transfers the item into a turn, never across a provider await. + fn acquire(&self) -> Result<std::sync::Arc<dyn QueueAdmissionPermit>, QueueAdmissionError>; + + /// Acquires the identity returned by the queue's immediately preceding + /// `batch_key` call. Implementations with replaceable authority should + /// reject a changed identity so a permit is never cached under a stale key. + fn acquire_for( + &self, + _expected_key: Option<&str>, + ) -> Result<std::sync::Arc<dyn QueueAdmissionPermit>, QueueAdmissionError> { + self.acquire() + } + + /// Stable identity for capabilities that may share one batch permit. + /// `None` means each item requires an independent acquisition. + fn batch_key(&self) -> Option<String> { + None + } +} + +/// Owned authority held across the queue-to-turn transfer. +pub trait QueueAdmissionPermit: Send + Sync {} + +impl QueueAdmissionPermit for () {} + +/// Why a queued item could not acquire its admission capability. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum QueueAdmissionError { + /// A binding is changing; retain the item for a later queue boundary. + Retry, + /// The lifecycle permanently revoked this item; discard it. + Rejected, +} + /// A queue lane consumed by the agent runtime. #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] #[serde(rename_all = "snake_case")] diff --git a/crates/tinyagents-harness/src/runtime/README.md b/crates/tinyagents-harness/src/runtime/README.md index 98b3cd95f..c2f4f762e 100644 --- a/crates/tinyagents-harness/src/runtime/README.md +++ b/crates/tinyagents-harness/src/runtime/README.md @@ -5,8 +5,8 @@ invocation entry points that sit on top of the default agent loop (`crate::agent_loop`). `AgentHarness<State, Ctx>` is the durable, reusable object a caller builds -once — it owns the model registry, tool registry, middleware stack, and -run policy. The agent loop itself (model ⇄ tool) is implemented as inherent +once — it owns the configured model resolver, tool registry, middleware stack, +and run policy. The agent loop itself (model ⇄ tool) is implemented as inherent methods on `AgentHarness` in the sibling `agent_loop` module; this module owns construction/registration (`mod.rs`), the type definitions (`types.rs`), and the separate host-driven invocation surface (`agent.rs`) that wires a @@ -33,7 +33,7 @@ enforcement — only how the turn gets *set up* differs. | Item | File | Role | | --- | --- | --- | -| `AgentHarness<State, Ctx>` | `types.rs` (impl in `mod.rs`, `agent.rs`) | The builder/facade: registries, middleware, policy, response cache. | +| `AgentHarness<State, Ctx>` | `types.rs` (impl in `mod.rs`, `agent.rs`) | The builder/facade: model resolver, tools, middleware, policy, response cache. | | `RunPolicy` | `types.rs` | Declarative cross-cutting policy: limits, unknown-tool/invalid-args handling, retry, fallback, response format, payload capture, cache policy, empty-response handling. | | `UnknownToolPolicy`, `InvalidArgsPolicy` | `types.rs` | How the loop recovers from a hallucinated tool name or schema-invalid arguments. | | `PayloadCapture` | `types.rs` | Opt-in capture of model/tool payloads onto observability events. | @@ -44,7 +44,7 @@ enforcement — only how the turn gets *set up* differs. | File | Role | | --- | --- | -| `mod.rs` | Builder/registration/accessor methods on `AgentHarness` (`register_model`, `push_agent_middleware`, `push_middleware`, `with_response_cache`, …). | +| `mod.rs` | Builder/accessor methods on `AgentHarness` (`with_model_resolver`, `with_model`, `register_tool`, `push_agent_middleware`, `push_middleware`, `with_response_cache`, …). | | `types.rs` | Public type definitions: `RunPolicy` and its sub-policies, `PayloadCapture`, `AgentHarness`'s fields, `InvocationRuntime`, and the crate-private `HostInvocationBinding`. | | `agent.rs` | Host-driven invocation: definition resolution, security screening, context composition, memory/experience recall, the `invoke_agent*` entry points, and post-turn memory/learning/experience finalization. | | `mod_tests.rs` | Tests for `AgentHarness` construction/registration and `RunPolicy` defaults. | diff --git a/crates/tinyagents-harness/src/runtime/mod.rs b/crates/tinyagents-harness/src/runtime/mod.rs index 180f4d6b5..a81c6a86b 100644 --- a/crates/tinyagents-harness/src/runtime/mod.rs +++ b/crates/tinyagents-harness/src/runtime/mod.rs @@ -40,12 +40,11 @@ pub use types::*; use std::sync::Arc; use crate::cache::ResponseCache; +use crate::host::ModelResolver; use crate::middleware::{ AgentMiddleware, Middleware, MiddlewareStack, ModelMiddleware, ToolMiddleware, }; -use crate::model_registry::ModelRegistry; use crate::tool::{ToolDispatch, ToolRegistry, ToolTimeoutSettings}; -use tinyinference_llm::model::ChatModel; use tinytools::Tool; /// Installs a host-selected recursion ceiling without widening the active @@ -60,11 +59,40 @@ pub fn install_depth<State: Send + Sync, Ctx: Send + Sync>( } impl<State: Send + Sync, Ctx: Send + Sync> AgentHarness<State, Ctx> { + /// Resolves one direct SDK model request through this harness's configured + /// model port. Hosted invocations use their invocation-local resolver in + /// the agent loop instead. + pub async fn resolve_model_for_request( + &self, + ctx: &crate::context::RunContext<Ctx>, + request: &tinyinference_llm::model::ModelRequest, + ) -> crate::error::Result<( + Option<String>, + Arc<dyn tinyinference_llm::model::ChatModel<State>>, + )> { + let binding = self.resolve_local_model(ctx, request).await?; + Ok((Some(binding.resolved.name), binding.model)) + } + + /// Returns the profile a direct SDK request would reach without dispatching + /// a model call. Hosts with dynamic routing can decline synchronous preview. + pub fn preview_model_profile_for_request( + &self, + request: &tinyinference_llm::model::ModelRequest, + ) -> Option<tinyinference_llm::model::ModelProfile> { + crate::agent_loop::model_call::preview_local_model_profile( + self.model_resolver.as_ref(), + self.fixed_model.as_ref(), + request, + ) + } + /// Creates an empty harness with default policy and no models, tools, or /// middleware registered. pub fn new() -> Self { Self { - models: ModelRegistry::new(), + model_resolver: None, + fixed_model: None, tools: ToolRegistry::new(), middleware: MiddlewareStack::new(), policy: RunPolicy::default(), @@ -73,27 +101,24 @@ impl<State: Send + Sync, Ctx: Send + Sync> AgentHarness<State, Ctx> { output_validator: None, deferred_tool_handler: None, toolset: None, - capabilities: Vec::new(), - capability_base_toolset: None, - loop_driver: None, } } - /// Registers a model under `name`. The first registered model becomes the - /// registry default unless one is already set. Returns `&mut Self` for - /// chaining. - pub fn register_model( - &mut self, - name: impl Into<String>, - model: Arc<dyn ChatModel<State>>, - ) -> &mut Self { - self.models.register(name, model); + /// Supplies the resolver used for direct SDK turns. Hosted invocations + /// continue to use their invocation-local host resolver. + pub fn with_model_resolver(&mut self, resolver: Arc<dyn ModelResolver<State>>) -> &mut Self { + self.model_resolver = Some(resolver); self } - /// Sets the default model name used when a request specifies no override. - pub fn set_default_model(&mut self, name: impl Into<String>) -> &mut Self { - self.models.set_default(name); + /// Supplies a single model directly when the state type carries a + /// non-static borrow and therefore cannot be used through the resolver + /// trait object. + pub fn with_model( + &mut self, + model: Arc<dyn tinyinference_llm::model::ChatModel<State>>, + ) -> &mut Self { + self.fixed_model = Some(model); self } @@ -293,93 +318,6 @@ impl<State: Send + Sync, Ctx: Send + Sync> AgentHarness<State, Ctx> { 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<State, Ctx>, - ) -> &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<dyn crate::tool::toolset::ToolSet<State, Ctx>> = - 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<State> { - &self.models - } - /// Returns a reference to the tool registry. pub fn tools(&self) -> &ToolRegistry<State, Ctx> { &self.tools @@ -394,27 +332,6 @@ impl<State: Send + Sync, Ctx: Send + Sync> AgentHarness<State, Ctx> { 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<dyn crate::agent_loop::phases::LoopDriver<State, Ctx>>, - ) -> &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<dyn crate::agent_loop::phases::LoopDriver<State, Ctx>>> { - self.loop_driver.as_ref() - } } impl<State: Send + Sync, Ctx: Send + Sync> Default for AgentHarness<State, Ctx> { diff --git a/crates/tinyagents-harness/src/runtime/mod_tests.rs b/crates/tinyagents-harness/src/runtime/mod_tests.rs index 81aa8b9c1..572001d32 100644 --- a/crates/tinyagents-harness/src/runtime/mod_tests.rs +++ b/crates/tinyagents-harness/src/runtime/mod_tests.rs @@ -20,7 +20,7 @@ use crate::middleware::{LoggingMiddleware, Middleware, ModelFallbackMiddleware}; use crate::retry::{FallbackPolicy, RetryPolicy}; use crate::runtime::{AgentHarness, AgentInvocation, AgentTurnRequest, RunPolicy}; use crate::testkit::ScriptedModel; -use futures::StreamExt; +use futures::{FutureExt, StreamExt}; use sha2::Digest; use tinyagents_definition::{ AgentDefinition, DefinitionRegistry, DefinitionRegistryError, InMemoryDefinitionRegistry, @@ -58,6 +58,11 @@ struct PendingResolver { started: Arc<tokio::sync::Notify>, } +struct PendingNamedFallbackResolver { + primary: Arc<dyn ChatModel<()>>, + named_started: Arc<tokio::sync::Notify>, +} + struct FirstThenPendingResolver { initial: Arc<dyn ChatModel<()>>, started: Arc<tokio::sync::Notify>, @@ -561,6 +566,33 @@ impl crate::host::ModelResolver<()> for PendingResolver { } } +#[async_trait] +impl crate::host::ModelResolver<()> for PendingNamedFallbackResolver { + async fn resolve( + &self, + _request: &crate::host::ModelResolveRequest, + ) -> crate::error::Result<Arc<dyn ChatModel<()>>> { + Ok(self.primary.clone()) + } + + async fn resolve_named( + &self, + _name: &str, + ) -> crate::error::Result<Option<Arc<dyn ChatModel<()>>>> { + self.named_started.notify_one(); + std::future::pending().await + } + + fn selected_name(&self, request: &crate::host::ModelResolveRequest) -> Option<String> { + Some( + request + .model_pin + .clone() + .unwrap_or_else(|| "primary".to_owned()), + ) + } +} + #[async_trait] impl crate::host::ModelResolver<()> for FirstThenPendingResolver { async fn resolve( @@ -966,23 +998,12 @@ impl Tool for InjectedArgumentTool { #[test] fn new_harness_is_empty_with_default_policy() { let harness: AgentHarness<()> = AgentHarness::new(); - assert!(harness.models().default_name().is_none()); + assert!(harness.model_resolver.is_none()); assert_eq!(harness.tools().names().len(), 0); assert!(harness.middleware().is_empty()); assert_eq!(harness.policy(), &RunPolicy::default()); } -#[test] -fn register_first_model_becomes_default() { - let mut harness: AgentHarness<()> = AgentHarness::new(); - harness - .register_model("a", Arc::new(MockModel::constant("a"))) - .register_model("b", Arc::new(MockModel::constant("b"))); - assert_eq!(harness.models().default_name(), Some("a")); - harness.set_default_model("b"); - assert_eq!(harness.models().default_name(), Some("b")); -} - #[test] fn register_tool_and_push_middleware() { let mut harness: AgentHarness<()> = AgentHarness::new(); @@ -1027,7 +1048,7 @@ fn install_depth_only_narrows_the_existing_limit() { #[test] fn default_matches_new() { let harness: AgentHarness<()> = AgentHarness::default(); - assert!(harness.models().default_name().is_none()); + assert!(harness.model_resolver.is_none()); } #[tokio::test] @@ -1081,8 +1102,8 @@ async fn host_driven_turn_resolves_and_composes_without_touching_explicit_sdk_de ); yield_until(|| progress.len() == 2).await; assert!( - harness.models().default_name().is_none(), - "host resolution does not mutate SDK model defaults" + harness.model_resolver.is_none(), + "host resolution does not mutate the SDK resolver" ); } @@ -1231,6 +1252,79 @@ async fn initial_host_model_resolution_is_cancelled_while_the_resolver_is_pendin assert_eq!(error.timeout_bound, None); } +#[tokio::test(start_paused = true)] +async fn pending_named_fallback_resolution_obeys_the_model_call_deadline() { + let started = Arc::new(tokio::sync::Notify::new()); + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness.with_model_resolver(Arc::new(PendingNamedFallbackResolver { + primary: Arc::new(RetryableFailingModel), + named_started: started.clone(), + })); + harness.with_policy(RunPolicy { + limits: RunLimits::default().with_max_model_call_ms(Some(10)), + retry: RetryPolicy::default() + .with_max_attempts(1) + .with_backoff_sleep(false), + fallback: Some(FallbackPolicy::new(["primary", "backup"])), + ..RunPolicy::default() + }); + + let error = harness + .invoke_default(&(), vec![tinyinference_llm::message::Message::user("go")]) + .await + .expect_err("the named fallback resolver remains pending"); + assert!( + matches!( + error, + crate::error::TinyAgentsError::Timeout(_) + | crate::error::TinyAgentsError::CallTimeout(_) + ), + "got {error:?}" + ); + assert!( + started.notified().now_or_never().is_some(), + "fallback lookup was entered" + ); +} + +#[tokio::test] +async fn pending_named_fallback_resolution_is_cancelled_with_the_run() { + let started = Arc::new(tokio::sync::Notify::new()); + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness.with_model_resolver(Arc::new(PendingNamedFallbackResolver { + primary: Arc::new(RetryableFailingModel), + named_started: started.clone(), + })); + harness.with_policy(RunPolicy { + retry: RetryPolicy::default() + .with_max_attempts(1) + .with_backoff_sleep(false), + fallback: Some(FallbackPolicy::new(["primary", "backup"])), + ..RunPolicy::default() + }); + let token = crate::CancellationToken::new(); + let invocation = harness.invoke_in_context( + &(), + RunContext::new(RunConfig::new("named-fallback-cancel"), ()) + .with_cancellation(token.clone()), + vec![tinyinference_llm::message::Message::user("go")], + ); + tokio::pin!(invocation); + let error = tokio::select! { + _ = started.notified() => { + token.cancel(); + tokio::time::timeout(Duration::from_secs(1), &mut invocation) + .await.expect("cancellation drops the pending resolver") + .expect_err("cancelled fallback resolution cannot succeed") + } + result = &mut invocation => panic!("pending named resolver unexpectedly finished: {result:?}"), + }; + assert!( + matches!(error, crate::error::TinyAgentsError::Cancelled), + "got {error:?}" + ); +} + #[tokio::test] async fn policy_only_deadline_bounds_initial_host_resolution_with_a_timeout_error() { let started = Arc::new(tokio::sync::Notify::new()); @@ -1526,7 +1620,9 @@ async fn hosted_model_fallback_rebinds_through_the_host_resolver_not_the_local_r resolver.clone(), ); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("host-backup", Arc::new(MockModel::constant("local bypass"))); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::new(Arc::new( + MockModel::constant("local bypass"), + )))); harness.push_middleware(Arc::new(FallbackGuidance)); harness.push_model_middleware(Arc::new(ModelFallbackMiddleware::new(["host-backup"]))); @@ -3105,7 +3201,7 @@ async fn retryable_classifier_changes_the_model_visible_result_without_redispatc async fn explicit_model_paths_accept_borrowed_state() { async fn exercise<'state>(state: &'state str) { let mut harness: AgentHarness<&'state str> = AgentHarness::new(); - harness.register_model("borrowed", Arc::new(BorrowedStateModel)); + harness.with_model(Arc::new(BorrowedStateModel)); let unary = harness .invoke( diff --git a/crates/tinyagents-harness/src/runtime/types.rs b/crates/tinyagents-harness/src/runtime/types.rs index 2fe88cc03..f41038a3d 100644 --- a/crates/tinyagents-harness/src/runtime/types.rs +++ b/crates/tinyagents-harness/src/runtime/types.rs @@ -23,9 +23,9 @@ use std::sync::Arc; use crate::cache::ResponseCache; use crate::host::HostCapabilities; +use crate::host::ModelResolver; 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}; @@ -531,20 +531,6 @@ pub struct RunPolicy { /// 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, /// Whether each executed tool's result row ends with how long the call /// took, as a trailing `[took 12.3s]` line the model reads. /// @@ -557,20 +543,6 @@ pub struct RunPolicy { pub tool_result_durations: bool, } -/// 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 { @@ -738,7 +710,6 @@ impl Default for RunPolicy { end_strategy: EndStrategy::default(), structured_strategy_override: None, queue_mode: QueueMode::default(), - execution: LoopExecution::default(), tool_result_durations: false, } } @@ -763,15 +734,18 @@ impl Default for RunPolicy { /// ``` /// use std::sync::Arc; /// use tinyinference_llm::providers::MockModel; +/// use tinyagents_harness::host::FixedModelResolver; /// use tinyagents_harness::runtime::AgentHarness; /// /// let mut harness: AgentHarness<()> = AgentHarness::new(); -/// harness.register_model("mock", Arc::new(MockModel::constant("hello"))); -/// assert_eq!(harness.models().default_name(), Some("mock")); +/// harness.with_model_resolver(Arc::new(FixedModelResolver::named("mock", Arc::new(MockModel::constant("hello"))))); /// ``` pub struct AgentHarness<State: Send + Sync, Ctx: Send + Sync = ()> { - /// Name-keyed registry of chat models with an optional default. - pub(crate) models: ModelRegistry<State>, + /// Host-supplied resolver for local SDK turns. + pub(crate) model_resolver: Option<Arc<dyn ModelResolver<State>>>, + /// One directly supplied model, supporting borrowed state without + /// imposing an object lifetime on the routing trait. + pub(crate) fixed_model: Option<Arc<dyn tinyinference_llm::model::ChatModel<State>>>, /// Name-keyed registry of tools exposed to the model. pub(crate) tools: ToolRegistry<State, Ctx>, /// Ordered middleware stack wrapping agent, model, and tool execution. @@ -808,23 +782,6 @@ pub struct AgentHarness<State: Send + Sync, Ctx: Send + Sync = ()> { /// [`AgentHarness::with_toolset`]. See that method's doc comment for /// exactly which turn behavior this changes. pub(crate) toolset: Option<Arc<dyn crate::tool::toolset::ToolSet<State, Ctx>>>, - /// 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<crate::capability::Capability<State, Ctx>>, - /// 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<Arc<dyn crate::tool::toolset::ToolSet<State, Ctx>>>, - /// 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<Arc<dyn crate::agent_loop::phases::LoopDriver<State, Ctx>>>, } /// The non-serializable mechanics selected for one hosted invocation. diff --git a/crates/tinyagents-harness/src/steering/README.md b/crates/tinyagents-harness/src/steering/README.md index 11538ebe2..2d29a6b4e 100644 --- a/crates/tinyagents-harness/src/steering/README.md +++ b/crates/tinyagents-harness/src/steering/README.md @@ -4,7 +4,7 @@ Policy-checked, observable orchestrator → sub-agent steering: how a *parent* in the run tree exerts typed control over a *child* it is currently running, without killing or restarting it. -This is the mid-run counterpart to `tinyagents_orchestration::subagent::SubAgentSession` reuse +This is the mid-run counterpart to `durable child-session reuse through the host-supplied driver lifecycle (which resumes a *completed* child) — together they cover both ways an orchestrator keeps a sub-agent "in play". An orchestrating agent, a human UI, a graph supervisor, or a test harness attaches a `SteeringHandle` to a run's diff --git a/crates/tinyagents-harness/src/steering/mod.rs b/crates/tinyagents-harness/src/steering/mod.rs index 54f80360e..9a171e26b 100644 --- a/crates/tinyagents-harness/src/steering/mod.rs +++ b/crates/tinyagents-harness/src/steering/mod.rs @@ -4,7 +4,7 @@ //! typed control over a *child* it is currently running: an orchestrating agent //! (or human, or graph supervisor) injects [`SteeringCommand`]s into a live //! sub-agent loop without killing or restarting it. It is the mid-run -//! counterpart to `SubAgentSession` reuse in `tinyagents-orchestration` (which +//! counterpart to durable child-session reuse in `tinyagents-orchestration` (which //! resumes a *completed* child) — together they cover both ways an orchestrator //! keeps a sub-agent "in play". //! diff --git a/crates/tinyagents-harness/src/steering/mod_tests.rs b/crates/tinyagents-harness/src/steering/mod_tests.rs index 82fb45e81..64784f612 100644 --- a/crates/tinyagents-harness/src/steering/mod_tests.rs +++ b/crates/tinyagents-harness/src/steering/mod_tests.rs @@ -274,7 +274,10 @@ async fn orchestrator_injects_message_mid_run_next_model_call_sees_it() { }); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", model.clone()); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + model.clone(), + ))); harness.register_tool(Arc::new(NoopTool)); let ctx: RunContext = RunContext::new(RunConfig::new("run-inject"), ()) @@ -318,7 +321,10 @@ async fn cancel_terminates_the_run() { handle.send(SteeringCommand::Cancel); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", Arc::new(MockModel::constant("never reached"))); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + Arc::new(MockModel::constant("never reached")), + ))); let ctx: RunContext = RunContext::new(RunConfig::new("run-cancel"), ()) .with_events(recorder.sink()) @@ -349,7 +355,10 @@ async fn disallowed_command_is_skipped_and_the_run_still_completes() { handle.send(SteeringCommand::InjectMessage(Message::user("nope"))); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", Arc::new(MockModel::constant("reached"))); + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( + "mock", + Arc::new(MockModel::constant("reached")), + ))); let ctx: RunContext = RunContext::new(RunConfig::new("run-reject"), ()) .with_events(recorder.sink()) diff --git a/crates/tinyagents-harness/src/steering/types.rs b/crates/tinyagents-harness/src/steering/types.rs index fdd4a77b0..3979d5194 100644 --- a/crates/tinyagents-harness/src/steering/types.rs +++ b/crates/tinyagents-harness/src/steering/types.rs @@ -56,8 +56,8 @@ pub enum SteeringTarget { #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] #[serde(rename_all = "snake_case", tag = "command")] #[non_exhaustive] -// Keep `InjectMessage(Message)` source-compatible; boxing it only to silence -// this layout lint would change the public Rust constructor's payload type. +// Keep `InjectMessage(Message)` source-compatible: boxing it would break the +// public command API, and commands are short-lived control messages. #[allow(clippy::large_enum_variant)] pub enum SteeringCommand { /// Cooperatively pause the run: the loop stops issuing further model and diff --git a/crates/tinyagents-harness/src/store/README.md b/crates/tinyagents-harness/src/store/README.md index 4582a29e7..ad5e8cdf0 100644 --- a/crates/tinyagents-harness/src/store/README.md +++ b/crates/tinyagents-harness/src/store/README.md @@ -11,9 +11,7 @@ nested executions read and write the same namespaced values, so a deeply nested call can persist a result a sibling or a later turn picks up. It is the harness-side persistence layer for runtime data — events, model and tool call records, message history, artifacts, memory — and is intentionally -separate from graph checkpointing (which belongs to `tinyagents-graph`'s -`checkpoint` module) and from prompt/model context assembly (`prompt`, -`context`). +separate from prompt/model context assembly (`prompt`, `context`). ## Public surface diff --git a/crates/tinyagents-harness/src/store/conformance.rs b/crates/tinyagents-harness/src/store/conformance.rs index 42ae06107..588623015 100644 --- a/crates/tinyagents-harness/src/store/conformance.rs +++ b/crates/tinyagents-harness/src/store/conformance.rs @@ -1,9 +1,8 @@ //! 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 +//! Mirrors 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]`. diff --git a/crates/tinyagents-harness/src/store/mod.rs b/crates/tinyagents-harness/src/store/mod.rs index 105eac34b..deb0afef6 100644 --- a/crates/tinyagents-harness/src/store/mod.rs +++ b/crates/tinyagents-harness/src/store/mod.rs @@ -4,14 +4,12 @@ //! that outlives any single run: parent and child runs, sub-agents, and //! nested executions read and write the same namespaced values, so a //! deeply nested call can persist a result that a sibling or a later turn picks -//! up. It is the harness-side persistence layer, distinct from graph -//! checkpointing. +//! up. //! //! The store is the persistence layer for harness runtime data: events, model //! call records, tool call records, message history, artifacts, and memory. It -//! is intentionally separate from graph checkpointing (which belongs to the -//! graph module) and from prompt/model context assembly (which belongs to the -//! model and prompt modules). +//! is intentionally separate from prompt/model context assembly (which belongs +//! to the model and prompt modules). //! //! # Primary types //! - [`Store`] — the core async trait every backend implements. diff --git a/crates/tinyagents-harness/src/stream/README.md b/crates/tinyagents-harness/src/stream/README.md deleted file mode 100644 index 80172a351..000000000 --- a/crates/tinyagents-harness/src/stream/README.md +++ /dev/null @@ -1,56 +0,0 @@ -# harness::stream - -Higher-level streaming projections for the harness — the typed lens an -observer uses to watch a run (and, transitively, the sub-agents and -sub-graphs it spawns) unfold in real time. - -## Why this exists - -State snapshots, diffs, model deltas, debug traces, and interrupts from a run -are projected as filtered [`StreamChunk`]s so a parent driving the run can -consume only the categories it cares about, -instead of every caller re-implementing delta reassembly and mode filtering -over raw [`crate::events::AgentEvent`]s. - -The chunk *types* (`types.rs`) are independent of `crate::events`; only the -projection (`project.rs`) depends on it, so callers that never touch events -pay nothing for the coupling. - -## Public surface - -- [`StreamMode`] — selects which chunk categories a consumer wants: `Values`, - `Updates`, `Messages`, `Debug`, `Interrupts`, `Custom`. -- [`StreamChunk`] — the typed union of chunk categories, adjacently tagged - for serialization (`{"type": …, "content": …}`). `mode()` returns the - `StreamMode` that gates a given chunk. -- [`StreamSink`] — a synchronous, mode-filtered buffer. `new`/`all` construct - it; `push`/`push_event` submit chunks/events; `drain`/`peek` read the - buffer; `enable`/`disable` adjust the active mode set at runtime. -- [`stream`] — a synchronous helper filtering an already-collected slice of - chunks by a set of modes (tests, post-processing). -- [`project_event`] / [`project_event_for_modes`] / [`projected_mode`] — the - `AgentEvent` → `StreamChunk` projection. `project_event_for_modes` is the - one a streaming run loop should call per event: it skips - clone/serialization work for any mode nobody subscribed to. - -## Files - -| File | Role | -| --- | --- | -| `mod.rs` | `StreamSink` method implementations and the standalone `stream()` filter helper. | -| `project.rs` | The `AgentEvent` → `StreamChunk` projection: `project_event`, `project_event_for_modes`, `projected_mode`, and the routing table in its module docs. | -| `types.rs` | `StreamMode`, `StreamChunk`, `StreamSink` type definitions. | -| `mod_tests.rs` | `StreamSink` filtering/push/drain/peek/enable/disable, `StreamChunk::mode` mapping, and the `stream()` helper. | - -## Operational constraints - -- Each `AgentEvent` projects to **at most one** `StreamChunk`, so a consumer - subscribed to several modes never receives the same event twice in two - shapes (see the routing table in `project.rs`). -- `StreamMode::Values` is **never** produced by the projection — a full state - snapshot is graph state, which the event stream does not carry; the graph - runtime pushes `StreamChunk::Values` itself via `StreamSink::push`. -- `StreamMode::Custom` is likewise never produced by projection — by - definition it is the caller's own extension channel, pushed directly. -- `StreamSink` is synchronous and single-threaded (`RefCell`-backed buffer); - wrap it in `Arc<Mutex<_>>` to share across threads. diff --git a/crates/tinyagents-harness/src/stream/frame.rs b/crates/tinyagents-harness/src/stream/frame.rs deleted file mode 100644 index 86759d09d..000000000 --- a/crates/tinyagents-harness/src/stream/frame.rs +++ /dev/null @@ -1,439 +0,0 @@ -//! 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<String>, - }, - /// 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<AssistantMessage>, - /// The stop/finish reason reported before the failure, when known. - #[serde(default, skip_serializing_if = "Option::is_none")] - stop_reason: Option<String>, - }, -} - -// --------------------------------------------------------------------------- -// 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<AssistantFrame>, - /// Per-block running tool-argument JSON and fragment count since the - /// last checkpoint, keyed by block index. - tool_progress: BTreeMap<usize, (String, usize)>, -} - -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<AssistantFrame> { - 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<AssistantFrame> { - 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), - ProviderExtension { - block_type: String, - }, - ToolCall { - id: Option<String>, - name: Option<String>, - 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<ContentBlock>, - /// 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<ToolCall>, - /// Most recent usage value seen. - pub usage: Option<Usage>, - /// Present once a terminal [`AssistantFrame::Completed`] or - /// [`AssistantFrame::Failed`] frame has been folded in. - pub terminal: Option<PartialTerminal>, -} - -/// 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<String>, - }, - /// 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<AssistantMessage>, - /// The stop/finish reason reported before the failure, when known. - stop_reason: Option<String>, - }, -} - -/// 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<usize, OpenBlock> = BTreeMap::new(); - let mut closed: BTreeMap<usize, ContentBlock> = 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::ProviderExtension { block_type } => OpenBlock::ProviderExtension { - block_type: block_type.clone(), - }, - 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::ProviderExtension { block_type } => { - format!("provider extension: {block_type}") - } - OpenBlock::ToolCall { json_so_far, .. } => json_so_far, - }; - (index, text) - }) - .collect(); - - PartialAssistantMessage { - content, - open_blocks, - tool_calls, - usage, - terminal, - } -} - -#[cfg(test)] -#[path = "frame/frame_tests.rs"] -mod test; diff --git a/crates/tinyagents-harness/src/stream/frame/frame_tests.rs b/crates/tinyagents-harness/src/stream/frame/frame_tests.rs deleted file mode 100644 index 6f185929d..000000000 --- a/crates/tinyagents-harness/src/stream/frame/frame_tests.rs +++ /dev/null @@ -1,308 +0,0 @@ -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<ModelStreamItem> { - 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())] - ); -} - -#[test] -fn provider_extension_frames_retain_the_completed_content_block() { - let extension = json!({"type": "future_block", "opaque": true}); - let partial = reduce_frames(&[ - AssistantFrame::BlockStart { - index: 0, - kind: BlockKind::ProviderExtension { - block_type: "future_block".into(), - }, - }, - AssistantFrame::BlockEnd { - index: 0, - block: ContentBlock::ProviderExtension(extension.clone()), - }, - ]); - - assert_eq!( - partial.content, - vec![ContentBlock::ProviderExtension(extension)] - ); - assert!(partial.open_blocks.is_empty()); -} diff --git a/crates/tinyagents-harness/src/stream/mod.rs b/crates/tinyagents-harness/src/stream/mod.rs deleted file mode 100644 index 8375bac48..000000000 --- a/crates/tinyagents-harness/src/stream/mod.rs +++ /dev/null @@ -1,187 +0,0 @@ -//! Higher-level streaming projections for the harness. -//! -//! In the recursive architecture this is the typed lens an observer uses to -//! watch work unfold across a nested run tree: state snapshots, diffs, model -//! deltas, debug traces, and interrupts from a run — and, transitively, from -//! the sub-agents and sub-graphs it spawns — are projected as filtered -//! [`StreamChunk`]s so a parent driving the run can -//! consume only the categories it cares about. -//! -//! This module provides: -//! -//! - [`StreamMode`] — an enum selecting which chunk categories a consumer -//! wants to receive. -//! - [`StreamChunk`] — a typed union of all chunk categories (state snapshots, -//! diffs, model deltas, debug output, interrupts, and custom extensions). -//! - [`StreamSink`] — a synchronous, mode-filtered buffer for [`StreamChunk`]s. -//! - [`stream`] — a convenience helper that filters a slice of chunks by a -//! set of modes and returns the matching subset. -//! -//! - [`project_event`] / [`project_event_for_modes`] / [`projected_mode`] — the -//! [`AgentEvent`] → [`StreamChunk`] projection (see the `project` module docs -//! for the routing table). -//! -//! The chunk *types* are independent of `crate::events`; only the -//! projection depends on it. Callers that never touch events still pay nothing -//! for the coupling. -//! -//! # Wiring note for the agent loop (wave 2) -//! -//! `invoke_stream` should build a `StreamSink` from the caller's requested -//! [`StreamMode`]s and call [`StreamSink::push_event`] (or -//! [`project_event_for_modes`] directly) for every emitted [`AgentEvent`], -//! instead of handing raw events to the caller. [`StreamMode::Values`] is the -//! 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::*; - -use std::collections::HashSet; - -use crate::events::AgentEvent; - -// --------------------------------------------------------------------------- -// StreamSink impls -// --------------------------------------------------------------------------- - -impl StreamSink { - /// Creates a sink that accepts chunks for the given modes. - /// - /// Pass an empty slice to create a sink that discards every chunk (useful - /// as a no-op sink in tests). - pub fn new(modes: impl IntoIterator<Item = StreamMode>) -> Self { - Self { - active_modes: modes.into_iter().collect(), - buffer: std::cell::RefCell::new(Vec::new()), - } - } - - /// Creates a sink that accepts **all** chunk modes. - pub fn all() -> Self { - Self::new([ - StreamMode::Values, - StreamMode::Updates, - StreamMode::Messages, - StreamMode::Debug, - StreamMode::Interrupts, - StreamMode::Custom, - ]) - } - - /// Returns `true` when the given mode is active on this sink. - pub fn is_active(&self, mode: StreamMode) -> bool { - self.active_modes.contains(&mode) - } - - /// Returns a reference to the set of active modes. - pub fn active_modes(&self) -> &HashSet<StreamMode> { - &self.active_modes - } - - /// Adds a mode to the active set. - pub fn enable(&mut self, mode: StreamMode) { - self.active_modes.insert(mode); - } - - /// Removes a mode from the active set. Chunks of that mode already in the - /// buffer are not removed. - pub fn disable(&mut self, mode: StreamMode) { - self.active_modes.remove(&mode); - } - - /// Buffers a chunk if its mode is active; silently discards it otherwise. - pub fn push(&self, chunk: StreamChunk) { - if self.active_modes.contains(&chunk.mode()) { - self.buffer.borrow_mut().push(chunk); - } - } - - /// Projects `event` and buffers the resulting chunk when its mode is - /// active. Returns `true` when a chunk was buffered. - /// - /// The one-line adapter between the event bus and this sink: a streaming - /// run loop calls this for every emitted [`AgentEvent`] and the sink's - /// active-mode set does the filtering. Projection is skipped entirely for - /// inactive modes (see [`project_event_for_modes`]). - pub fn push_event(&self, event: &AgentEvent) -> bool { - let modes: Vec<StreamMode> = self.active_modes.iter().copied().collect(); - match project_event_for_modes(event, &modes) { - Some(chunk) => { - self.buffer.borrow_mut().push(chunk); - true - } - None => false, - } - } - - /// Returns all buffered chunks in push order and clears the buffer. - pub fn drain(&self) -> Vec<StreamChunk> { - self.buffer.borrow_mut().drain(..).collect() - } - - /// Returns the number of buffered chunks. - pub fn len(&self) -> usize { - self.buffer.borrow().len() - } - - /// Returns `true` when no chunks are buffered. - pub fn is_empty(&self) -> bool { - self.buffer.borrow().is_empty() - } - - /// Peeks at the buffered chunks without consuming them. - pub fn peek(&self) -> Vec<StreamChunk> { - self.buffer.borrow().clone() - } -} - -// --------------------------------------------------------------------------- -// stream() helper -// --------------------------------------------------------------------------- - -/// Filters a collection of [`StreamChunk`]s to those matching a set of -/// [`StreamMode`]s and returns the matching chunks in input order. -/// -/// This is a synchronous, allocation-based helper for contexts where the full -/// chunk list is already collected (for example in tests or post-processing). -/// It does not require an async runtime. -/// -/// # Example -/// -/// ``` -/// use tinyagents_harness::stream::{stream, StreamChunk, StreamMode}; -/// use tinyinference_llm::message::MessageDelta; -/// -/// let chunks = vec![ -/// StreamChunk::Message(MessageDelta::text("hi")), -/// StreamChunk::Debug("trace".into()), -/// StreamChunk::Updates(serde_json::json!({"key": "value"})), -/// ]; -/// -/// let messages = stream(&chunks, &[StreamMode::Messages]); -/// assert_eq!(messages.len(), 1); -/// -/// let all = stream(&chunks, &[StreamMode::Messages, StreamMode::Debug, StreamMode::Updates]); -/// assert_eq!(all.len(), 3); -/// ``` -pub fn stream(chunks: &[StreamChunk], modes: &[StreamMode]) -> Vec<StreamChunk> { - let mode_set: HashSet<StreamMode> = modes.iter().copied().collect(); - chunks - .iter() - .filter(|c| mode_set.contains(&c.mode())) - .cloned() - .collect() -} - -#[cfg(test)] -#[path = "mod_tests.rs"] -mod test; diff --git a/crates/tinyagents-harness/src/stream/mod_tests.rs b/crates/tinyagents-harness/src/stream/mod_tests.rs deleted file mode 100644 index 3a3b34c0c..000000000 --- a/crates/tinyagents-harness/src/stream/mod_tests.rs +++ /dev/null @@ -1,332 +0,0 @@ -//! Tests for the harness streaming projections. -//! -//! Cover [`StreamSink`] mode filtering, push/drain/peek, enable/disable of -//! active modes, the all-modes and empty-modes sinks, [`StreamChunk::mode`] -//! variant mapping, and the standalone [`stream`] filter helper. - -use serde_json::json; - -use crate::stream::{StreamChunk, StreamMode, StreamSink, stream}; -use tinyinference_llm::message::MessageDelta; - -#[test] -fn smoke_sink_filters_by_mode() { - let sink = StreamSink::new([StreamMode::Messages]); - - sink.push(StreamChunk::Message(MessageDelta { - text: "hello".into(), - reasoning: String::new(), - tool_call: None, - })); - // Debug chunk should be discarded (mode not active). - sink.push(StreamChunk::Debug("internal note".into())); - - assert_eq!(sink.len(), 1); - - let chunks = sink.drain(); - assert_eq!(chunks.len(), 1); - assert!(matches!(chunks[0], StreamChunk::Message(_))); - - // Buffer is cleared after drain. - assert!(sink.is_empty()); -} - -#[test] -fn smoke_sink_all_accepts_every_mode() { - let sink = StreamSink::all(); - - sink.push(StreamChunk::Values(json!({"state": 1}))); - sink.push(StreamChunk::Updates(json!({"delta": 2}))); - sink.push(StreamChunk::Message(MessageDelta::default())); - sink.push(StreamChunk::Debug("trace".into())); - sink.push(StreamChunk::Interrupt(json!({"kind": "approval"}))); - sink.push(StreamChunk::Custom(json!({"ext": true}))); - - assert_eq!(sink.drain().len(), 6); -} - -#[test] -fn smoke_stream_helper_filters() { - let chunks = vec![ - StreamChunk::Message(MessageDelta { - text: "tok".into(), - reasoning: String::new(), - tool_call: None, - }), - StreamChunk::Debug("trace".into()), - StreamChunk::Values(json!(null)), - ]; - - let msgs = stream(&chunks, &[StreamMode::Messages]); - assert_eq!(msgs.len(), 1); - - let two = stream(&chunks, &[StreamMode::Messages, StreamMode::Debug]); - assert_eq!(two.len(), 2); - - let none = stream(&chunks, &[]); - assert!(none.is_empty()); -} - -#[test] -fn smoke_chunk_mode_matches_variant() { - assert_eq!(StreamChunk::Values(json!(null)).mode(), StreamMode::Values); - assert_eq!( - StreamChunk::Updates(json!(null)).mode(), - StreamMode::Updates - ); - assert_eq!( - StreamChunk::Message(MessageDelta::default()).mode(), - StreamMode::Messages - ); - assert_eq!(StreamChunk::Debug("x".into()).mode(), StreamMode::Debug); - assert_eq!( - StreamChunk::Interrupt(json!(null)).mode(), - StreamMode::Interrupts - ); - assert_eq!(StreamChunk::Custom(json!(null)).mode(), StreamMode::Custom); -} - -#[test] -fn sink_enable_disable_and_is_active() { - let mut sink = StreamSink::new([StreamMode::Messages]); - assert!(sink.is_active(StreamMode::Messages)); - assert!(!sink.is_active(StreamMode::Debug)); - - // Enabling Debug starts accepting debug chunks. - sink.enable(StreamMode::Debug); - assert!(sink.is_active(StreamMode::Debug)); - sink.push(StreamChunk::Debug("now kept".into())); - assert_eq!(sink.len(), 1); - - // Disabling Messages discards subsequent message chunks but keeps buffered. - sink.disable(StreamMode::Messages); - assert!(!sink.is_active(StreamMode::Messages)); - sink.push(StreamChunk::Message(MessageDelta::default())); - assert_eq!(sink.len(), 1); -} - -#[test] -fn sink_active_modes_returns_set() { - let sink = StreamSink::new([StreamMode::Values, StreamMode::Custom]); - let modes = sink.active_modes(); - assert_eq!(modes.len(), 2); - assert!(modes.contains(&StreamMode::Values)); - assert!(modes.contains(&StreamMode::Custom)); -} - -#[test] -fn sink_empty_modes_discards_everything() { - let sink = StreamSink::new([]); - sink.push(StreamChunk::Debug("x".into())); - sink.push(StreamChunk::Values(json!(1))); - assert!(sink.is_empty()); - assert_eq!(sink.len(), 0); -} - -#[test] -fn sink_peek_does_not_consume() { - let sink = StreamSink::all(); - sink.push(StreamChunk::Debug("a".into())); - sink.push(StreamChunk::Debug("b".into())); - - let peeked = sink.peek(); - assert_eq!(peeked.len(), 2); - // Peek leaves the buffer intact. - assert_eq!(sink.len(), 2); - // Drain still returns the same chunks. - assert_eq!(sink.drain().len(), 2); - assert!(sink.is_empty()); -} - -/// Round-trips a [`StreamChunk`] through JSON and asserts it deserializes back -/// to an equal value, proving every variant survives serde (including the -/// scalar `Value` payloads an internally tagged enum could not encode). -fn roundtrip_chunk(chunk: StreamChunk) { - let value = serde_json::to_value(&chunk).expect("serialize StreamChunk"); - let back: StreamChunk = serde_json::from_value(value).expect("deserialize StreamChunk"); - assert_eq!(chunk, back); -} - -#[test] -fn stream_chunk_roundtrips_every_variant() { - // Scalar / null / array Values are exactly what internal tagging corrupted. - roundtrip_chunk(StreamChunk::Values(json!(null))); - roundtrip_chunk(StreamChunk::Values(json!(42))); - roundtrip_chunk(StreamChunk::Values(json!({ "state": 1 }))); - roundtrip_chunk(StreamChunk::Updates(json!([1, 2, 3]))); - roundtrip_chunk(StreamChunk::Message(MessageDelta::text("hi"))); - roundtrip_chunk(StreamChunk::Debug("trace".into())); - roundtrip_chunk(StreamChunk::Interrupt(json!({ "kind": "approval" }))); - roundtrip_chunk(StreamChunk::Custom(json!(true))); -} - -#[test] -fn stream_chunk_null_values_does_not_corrupt_to_empty_object() { - let value = serde_json::to_value(StreamChunk::Values(json!(null))).unwrap(); - assert_eq!(value["type"], json!("values")); - assert_eq!(value["content"], json!(null)); -} - -// --------------------------------------------------------------------------- -// C6: AgentEvent -> StreamChunk projection -// --------------------------------------------------------------------------- - -mod project { - use crate::events::AgentEvent; - use crate::ids::{CallId, RunId}; - use crate::stream::{ - StreamChunk, StreamMode, StreamSink, project_event, project_event_for_modes, projected_mode, - }; - use tinyinference_llm::message::MessageDelta; - - fn delta_event() -> AgentEvent { - AgentEvent::ModelDelta { - run_id: RunId::new("r1"), - call_id: CallId::new("c1"), - delta: MessageDelta::text("hello"), - } - } - - #[test] - fn model_deltas_project_onto_messages_mode() { - // Before this projection existed, `StreamMode` / `StreamChunk` / - // `StreamSink` were referenced NOWHERE outside their own module, and - // every caller re-implemented delta reassembly against raw events. - let chunk = project_event(&delta_event()).expect("a delta must project"); - assert_eq!(chunk.mode(), StreamMode::Messages); - assert_eq!(chunk, StreamChunk::Message(MessageDelta::text("hello"))); - } - - #[test] - fn an_interrupting_control_is_the_producer_of_stream_chunk_interrupt() { - // `StreamChunk::Interrupt` was defined but never constructed anywhere - // in the crate. - let event = AgentEvent::ControlApplied { - control: "interrupt".into(), - detail: "approval_node: needs sign-off".into(), - }; - let chunk = project_event(&event).expect("must project"); - assert_eq!(chunk.mode(), StreamMode::Interrupts); - match chunk { - StreamChunk::Interrupt(value) => { - assert_eq!(value["control"], "interrupt"); - assert_eq!(value["detail"], "approval_node: needs sign-off"); - } - other => panic!("expected an Interrupt chunk, got {other:?}"), - } - } - - #[test] - fn a_non_interrupting_control_stays_in_the_debug_channel() { - let event = AgentEvent::ControlApplied { - control: "stop_with_final".into(), - detail: "done".into(), - }; - assert_eq!(projected_mode(&event), Some(StreamMode::Debug)); - } - - #[test] - fn state_updates_project_onto_updates_mode() { - let chunk = project_event(&AgentEvent::StateUpdate).expect("must project"); - assert_eq!(chunk.mode(), StreamMode::Updates); - } - - #[test] - fn every_other_event_falls_through_to_debug_and_keeps_its_payload() { - let event = AgentEvent::ToolStarted { - parent_call_id: None, - call_id: CallId::new("c9"), - tool_name: "search".into(), - input: None, - }; - let chunk = project_event(&event).expect("must project"); - assert_eq!(chunk.mode(), StreamMode::Debug); - match chunk { - // The typed payload survives, not just the kind string. - StreamChunk::Debug(text) => { - assert!(text.contains("tool_started"), "got {text}"); - assert!(text.contains("search"), "got {text}"); - } - other => panic!("expected a Debug chunk, got {other:?}"), - } - } - - #[test] - fn stream_closed_is_a_terminator_not_content() { - assert!(project_event(&AgentEvent::StreamClosed).is_none()); - assert!(projected_mode(&AgentEvent::StreamClosed).is_none()); - assert!(project_event_for_modes(&AgentEvent::StreamClosed, &[StreamMode::Debug]).is_none()); - } - - #[test] - fn one_event_never_projects_into_two_modes() { - // A consumer subscribed to several modes must not see the same event - // twice in two shapes. - for event in [ - delta_event(), - AgentEvent::StateUpdate, - AgentEvent::ControlApplied { - control: "interrupt".into(), - detail: "d".into(), - }, - AgentEvent::MemoryLoaded, - ] { - let all = [ - StreamMode::Values, - StreamMode::Updates, - StreamMode::Messages, - StreamMode::Debug, - StreamMode::Interrupts, - StreamMode::Custom, - ]; - let hits = all - .iter() - .filter(|mode| project_event_for_modes(&event, &[**mode]).is_some()) - .count(); - assert_eq!(hits, 1, "{} projected into {hits} modes", event.kind()); - } - } - - #[test] - fn mode_filtering_happens_producer_side() { - let event = delta_event(); - assert!(project_event_for_modes(&event, &[StreamMode::Messages]).is_some()); - assert!(project_event_for_modes(&event, &[StreamMode::Debug]).is_none()); - assert!(project_event_for_modes(&event, &[]).is_none()); - // Multiplexed mode sets work the way LangGraph's do. - assert!( - project_event_for_modes(&event, &[StreamMode::Debug, StreamMode::Messages]).is_some() - ); - } - - #[test] - fn push_event_bridges_the_event_bus_into_a_sink() { - let sink = StreamSink::new([StreamMode::Messages]); - assert!(sink.push_event(&delta_event())); - assert!(!sink.push_event(&AgentEvent::StateUpdate)); - assert!(!sink.push_event(&AgentEvent::StreamClosed)); - - let chunks = sink.drain(); - assert_eq!(chunks.len(), 1); - assert_eq!(chunks[0].mode(), StreamMode::Messages); - } - - #[test] - fn values_and_custom_are_never_produced_by_the_projection() { - // Documented gap: a full state snapshot is graph state, which the event - // stream does not carry, and Custom is the caller's own channel. - for event in [ - delta_event(), - AgentEvent::StateUpdate, - AgentEvent::MemorySaved, - AgentEvent::RunCompleted { - run_id: RunId::new("r1"), - outcome: None, - }, - ] { - let mode = projected_mode(&event); - assert_ne!(mode, Some(StreamMode::Values)); - assert_ne!(mode, Some(StreamMode::Custom)); - } - } -} diff --git a/crates/tinyagents-harness/src/stream/project.rs b/crates/tinyagents-harness/src/stream/project.rs deleted file mode 100644 index 4c7fa9aba..000000000 --- a/crates/tinyagents-harness/src/stream/project.rs +++ /dev/null @@ -1,156 +0,0 @@ -//! Projection from the raw [`AgentEvent`] stream onto typed [`StreamChunk`]s. -//! -//! This is the missing half of the streaming surface. [`StreamMode`], -//! [`StreamChunk`] and [`StreamSink`] were fully defined and tested but used -//! nowhere else in the crate — callers got raw `AgentEvent`s and each -//! re-implemented delta reassembly and filtering. This module supplies the one -//! function that turns an event into the chunk a consumer actually asked for. -//! -//! # Shape (ported from LangGraph) -//! -//! LangGraph multiplexes several *stream modes* over one producer: the producer -//! filters by the requested mode set (`_loop.py`'s `stream_modes` check) and the -//! consumer shapes the result (`main.py`'s `stream()` fan-out). The same split -//! applies here: -//! -//! - **Producer-side filtering** — [`project_event_for_modes`] returns `None` -//! for an event whose chunk is not in the requested mode set, so nothing is -//! allocated or cloned for a mode nobody is listening to. -//! - **Consumer-side shaping** — [`project_event`] gives the full projection -//! when the caller wants to route chunks itself. -//! -//! # One event, at most one chunk -//! -//! Each [`AgentEvent`] projects to **at most one** [`StreamChunk`], so a -//! consumer subscribed to several modes never receives the same event twice in -//! two shapes. The routing table: -//! -//! | Event | Chunk | Mode | -//! |---|---|---| -//! | [`AgentEvent::ModelDelta`] | [`StreamChunk::Message`] | [`StreamMode::Messages`] | -//! | [`AgentEvent::StateUpdate`] | [`StreamChunk::Updates`] | [`StreamMode::Updates`] | -//! | [`AgentEvent::ControlApplied`] with an interrupting control | [`StreamChunk::Interrupt`] | [`StreamMode::Interrupts`] | -//! | [`AgentEvent::StreamClosed`] | — (not projected) | — | -//! | everything else | [`StreamChunk::Debug`] | [`StreamMode::Debug`] | -//! -//! [`StreamMode::Values`] is deliberately **not** produced here: a full state -//! snapshot is graph state, which the event stream does not carry. The graph -//! runtime pushes [`StreamChunk::Values`] itself. -//! -//! [`StreamMode::Custom`] is likewise never produced — by definition it is the -//! caller's own extension channel. - -use crate::events::AgentEvent; - -use super::{StreamChunk, StreamMode}; - -/// Control kinds from -/// [`MiddlewareControl::kind`][crate::context::MiddlewareControl::kind] -/// that mean "the run paused and is waiting for something external", and so -/// belong on [`StreamMode::Interrupts`] rather than in the debug firehose. -const INTERRUPTING_CONTROLS: [&str; 1] = ["interrupt"]; - -/// Projects one [`AgentEvent`] onto the [`StreamChunk`] a consumer sees. -/// -/// Returns `None` for an event with no meaningful chunk representation — today -/// only [`AgentEvent::StreamClosed`], which is a stream terminator rather than -/// content (consumers detect end-of-stream by the stream ending). -/// -/// # Example -/// -/// ``` -/// use tinyagents_harness::events::AgentEvent; -/// use tinyagents_harness::ids::{CallId, RunId}; -/// use tinyinference_llm::message::MessageDelta; -/// use tinyagents_harness::stream::{project_event, StreamChunk}; -/// -/// let event = AgentEvent::ModelDelta { -/// run_id: RunId::new("r1"), -/// call_id: CallId::new("c1"), -/// delta: MessageDelta::text("hi"), -/// }; -/// assert!(matches!(project_event(&event), Some(StreamChunk::Message(_)))); -/// ``` -pub fn project_event(event: &AgentEvent) -> Option<StreamChunk> { - match event { - // The whole point of `messages` mode: raw token/tool-call fragments, - // already reassembled by the provider adapter. - AgentEvent::ModelDelta { delta, .. } => Some(StreamChunk::Message(delta.clone())), - - // `StateUpdate` is payload-free by design (see its variant docs), so the - // chunk carries the fact of the update and its provenance rather than a - // diff the event never had. - AgentEvent::StateUpdate => Some(StreamChunk::Updates(serde_json::json!({ - "kind": event.kind(), - }))), - - // A human-in-the-loop pause. This is the only producer of - // `StreamChunk::Interrupt` — the variant existed but nothing ever - // constructed it. - AgentEvent::ControlApplied { control, detail } - if INTERRUPTING_CONTROLS.contains(&control.as_str()) => - { - Some(StreamChunk::Interrupt(serde_json::json!({ - "control": control, - "detail": detail, - }))) - } - - // A terminator, not content. - AgentEvent::StreamClosed => None, - - // Everything else is diagnostic. Rendered as the event's own JSON so a - // debug consumer keeps the full typed payload rather than a lossy - // summary; falls back to the stable kind string if serialization ever - // fails (it cannot for the current variants, but a `Debug` chunk must - // never be the thing that panics a run). - other => Some(StreamChunk::Debug( - serde_json::to_string(other).unwrap_or_else(|_| other.kind().to_string()), - )), - } -} - -/// Producer-side filtered projection: like [`project_event`], but returns -/// `None` when the resulting chunk's mode is not in `modes`. -/// -/// This is the function a streaming run loop should call per event — it skips -/// the clone/serialization for any mode nobody subscribed to, which matters -/// because the [`StreamMode::Debug`] catch-all serializes every event. -/// -/// # Example -/// -/// ``` -/// use tinyagents_harness::events::AgentEvent; -/// use tinyagents_harness::stream::{project_event_for_modes, StreamMode}; -/// -/// let event = AgentEvent::StateUpdate; -/// assert!(project_event_for_modes(&event, &[StreamMode::Updates]).is_some()); -/// assert!(project_event_for_modes(&event, &[StreamMode::Messages]).is_none()); -/// ``` -pub fn project_event_for_modes(event: &AgentEvent, modes: &[StreamMode]) -> Option<StreamChunk> { - let mode = projected_mode(event)?; - if !modes.contains(&mode) { - return None; - } - project_event(event) -} - -/// Returns the [`StreamMode`] an event projects onto, without building the -/// chunk. -/// -/// Cheap enough to call in a hot loop, and it is what makes -/// [`project_event_for_modes`] able to skip work rather than build-then-discard. -/// `None` mirrors [`project_event`] returning `None`. -pub fn projected_mode(event: &AgentEvent) -> Option<StreamMode> { - match event { - AgentEvent::ModelDelta { .. } => Some(StreamMode::Messages), - AgentEvent::StateUpdate => Some(StreamMode::Updates), - AgentEvent::ControlApplied { control, .. } - if INTERRUPTING_CONTROLS.contains(&control.as_str()) => - { - Some(StreamMode::Interrupts) - } - AgentEvent::StreamClosed => None, - _ => Some(StreamMode::Debug), - } -} diff --git a/crates/tinyagents-harness/src/stream/types.rs b/crates/tinyagents-harness/src/stream/types.rs deleted file mode 100644 index c6508d5df..000000000 --- a/crates/tinyagents-harness/src/stream/types.rs +++ /dev/null @@ -1,140 +0,0 @@ -//! Type definitions for the harness streaming surface. -//! -//! All structs, enums, and traits in this module form the public API of -//! `crate::stream`. Implementations, free functions, and tests live -//! in the sibling `mod.rs` and `test.rs` files. -//! -//! The stream module is intentionally independent of `crate::events` -//! so it can be used without the full observability stack. - -use std::collections::HashSet; - -use serde::{Deserialize, Serialize}; - -use tinyinference_llm::message::MessageDelta; - -// --------------------------------------------------------------------------- -// StreamMode -// --------------------------------------------------------------------------- - -/// Selects which categories of [`StreamChunk`]s a consumer wants to receive. -/// -/// Consumers subscribe to one or more modes; a [`StreamSink`] only buffers -/// chunks that match the active mode set. Using a narrow mode set reduces -/// allocation and processing in tight loops. -#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum StreamMode { - /// Raw state value snapshots (full key/value state blobs). - Values, - /// Incremental state diffs (only changed keys). - Updates, - /// Model token and message deltas. - Messages, - /// Verbose debug output: intermediate steps, routing decisions, phase - /// transitions, and internal notes. - Debug, - /// Human-in-the-loop interrupt notifications. - Interrupts, - /// Caller-defined extension chunks passed through without filtering. - Custom, -} - -// --------------------------------------------------------------------------- -// StreamChunk -// --------------------------------------------------------------------------- - -/// A single item produced by a streaming harness or graph run. -/// -/// Each variant corresponds to a [`StreamMode`]: a chunk is buffered by a -/// [`StreamSink`] only when its matching mode is active. -/// -/// All variants derive `Clone` so chunks can be fanned out to multiple -/// consumers. -/// -/// # Serialization -/// -/// This enum is **adjacently tagged** (`{"type": …, "content": …}`) rather than -/// internally tagged. Most variants wrap a bare [`serde_json::Value`] or -/// [`String`]; an internally tagged enum cannot serialize a scalar or array -/// payload (serde errors, or corrupts `Values(json!(null))` into `{}`), so -/// adjacent tagging is required for every variant to round-trip. -#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] -#[serde(rename_all = "snake_case", tag = "type", content = "content")] -pub enum StreamChunk { - /// A complete state value snapshot (corresponds to [`StreamMode::Values`]). - Values(serde_json::Value), - - /// An incremental state diff (corresponds to [`StreamMode::Updates`]). - Updates(serde_json::Value), - - /// An incremental model response fragment (corresponds to - /// [`StreamMode::Messages`]). - Message(MessageDelta), - - /// A verbose debug string (corresponds to [`StreamMode::Debug`]). - Debug(String), - - /// A human-in-the-loop interrupt payload (corresponds to - /// [`StreamMode::Interrupts`]). - /// - /// The inner value carries interrupt metadata (kind, resume node, etc.) - /// as untyped JSON so the interrupt shape can vary without a hard coupling - /// to a specific interrupt type. - Interrupt(serde_json::Value), - - /// A caller-defined chunk that bypasses mode filtering and is always - /// buffered when [`StreamMode::Custom`] is active. - Custom(serde_json::Value), -} - -impl StreamChunk { - /// Returns the [`StreamMode`] that gates this chunk variant. - /// - /// Use this to implement custom routing or logging without matching on the - /// full enum. - pub fn mode(&self) -> StreamMode { - match self { - StreamChunk::Values(_) => StreamMode::Values, - StreamChunk::Updates(_) => StreamMode::Updates, - StreamChunk::Message(_) => StreamMode::Messages, - StreamChunk::Debug(_) => StreamMode::Debug, - StreamChunk::Interrupt(_) => StreamMode::Interrupts, - StreamChunk::Custom(_) => StreamMode::Custom, - } - } -} - -// --------------------------------------------------------------------------- -// StreamSink -// --------------------------------------------------------------------------- - -/// An in-process buffer for [`StreamChunk`]s filtered by an active set of -/// [`StreamMode`]s. -/// -/// Producers call [`StreamSink::push`] to submit chunks; the sink silently -/// discards chunks whose mode is not in the active set. Consumers call -/// [`StreamSink::drain`] to retrieve and clear all buffered chunks. -/// -/// `StreamSink` is deliberately synchronous and single-threaded. Wrap it in -/// `Arc<Mutex<…>>` when sharing across threads. -/// -/// # Example -/// -/// ``` -/// use tinyagents_harness::stream::{StreamChunk, StreamMode, StreamSink}; -/// use tinyinference_llm::message::MessageDelta; -/// -/// let sink = StreamSink::new(vec![StreamMode::Messages]); -/// sink.push(StreamChunk::Message(MessageDelta::text("hello"))); -/// sink.push(StreamChunk::Debug("ignored".into())); -/// -/// let chunks = sink.drain(); -/// assert_eq!(chunks.len(), 1); -/// ``` -pub struct StreamSink { - /// The set of modes whose chunks are accepted by this sink. - pub(crate) active_modes: HashSet<StreamMode>, - /// Buffered chunks, in push order. - pub(crate) buffer: std::cell::RefCell<Vec<StreamChunk>>, -} diff --git a/crates/tinyagents-harness/src/summarization/payload_summarizer.rs b/crates/tinyagents-harness/src/summarization/payload_summarizer.rs index 887231657..3dab31c62 100644 --- a/crates/tinyagents-harness/src/summarization/payload_summarizer.rs +++ b/crates/tinyagents-harness/src/summarization/payload_summarizer.rs @@ -85,9 +85,8 @@ pub async fn run_unary_summary<Ctx: Send + Sync + 'static>( let mut harness: AgentHarness<(), Ctx> = AgentHarness::new(); harness.with_policy(policy); - harness - .register_model(model_name, model) - .set_default_model(model_name); + let _ = model_name; + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::new(model))); let run = harness .invoke_in_context( &(), diff --git a/crates/tinyagents-harness/src/tool/binding.rs b/crates/tinyagents-harness/src/tool/binding.rs index 544c5bb3c..bf549869a 100644 --- a/crates/tinyagents-harness/src/tool/binding.rs +++ b/crates/tinyagents-harness/src/tool/binding.rs @@ -27,6 +27,7 @@ type Binder = dyn Fn(Arc<dyn Tool>, &Value) -> Arc<dyn BoundToolCall> + Send + S /// implementation-specific routing model. pub struct ToolCallBinding { bind: Arc<Binder>, + metadata: Option<Arc<dyn Any + Send + Sync>>, } impl ToolCallBinding { @@ -36,9 +37,25 @@ impl ToolCallBinding { ) -> Self { Self { bind: Arc::new(bind), + metadata: None, } } + /// Associates implementation-owned registration metadata with this binding. + /// + /// Metadata describes the registered tool and is not the call-scoped target + /// returned by [`Self::bind`]. Admission and execution must continue to use + /// the bound call so both observe the same captured target. + pub fn with_metadata<T: Any + Send + Sync>(mut self, metadata: T) -> Self { + self.metadata = Some(Arc::new(metadata)); + self + } + + /// Returns registration metadata when it has the requested concrete type. + pub fn metadata<T: Any + Send + Sync>(&self) -> Option<&T> { + self.metadata.as_ref()?.downcast_ref() + } + /// Captures the indirect target selected for these arguments. pub fn bind(&self, outer: Arc<dyn Tool>, arguments: &Value) -> Arc<dyn BoundToolCall> { (self.bind)(outer, arguments) diff --git a/crates/tinyagents-harness/src/tool/binding_tests.rs b/crates/tinyagents-harness/src/tool/binding_tests.rs index 4e37ca2d8..174929dc8 100644 --- a/crates/tinyagents-harness/src/tool/binding_tests.rs +++ b/crates/tinyagents-harness/src/tool/binding_tests.rs @@ -99,3 +99,28 @@ fn canonical_adapter_forwards_tool_extension_to_the_call_binding_contract() { assert_eq!(target.name(), "selected"); assert_eq!(effective_args, &args); } + +#[test] +fn binding_exposes_typed_registration_metadata_without_changing_call_binding() { + let binding = ToolCallBinding::new(|outer, args| { + Arc::new(BoundProbe { + outer: outer.clone(), + target: Arc::new(Target), + args: args.clone(), + }) + }) + .with_metadata(String::from("registry metadata")); + + assert_eq!( + binding.metadata::<String>().map(String::as_str), + Some("registry metadata") + ); + assert!(binding.metadata::<usize>().is_none()); + + let outer: Arc<dyn Tool> = Arc::new(Target); + let arguments = json!({"value": 7}); + let bound = binding.bind(outer.clone(), &arguments); + let (target, effective_args) = bound.admission_target().unwrap(); + assert_eq!(target.name(), "selected"); + assert_eq!(effective_args, &arguments); +} diff --git a/crates/tinyagents-harness/src/tool/mod.rs b/crates/tinyagents-harness/src/tool/mod.rs index b317505dc..1174da21f 100644 --- a/crates/tinyagents-harness/src/tool/mod.rs +++ b/crates/tinyagents-harness/src/tool/mod.rs @@ -50,7 +50,9 @@ pub use shared::{ CanonicalSharedToolAdapter, EarlyExit, EarlyExitHook, ToolExecScope, ToolExecution, }; pub use signature::*; -pub use snapshot::{ToolSnapshot, ToolSnapshotError}; +pub use snapshot::{ + ToolSnapshot, ToolSnapshotError, ToolSnapshotExtension, ToolSnapshotExtensions, +}; pub use timeout::*; /// Name of the generic skill disclosure and indirect-call tool. diff --git a/crates/tinyagents-harness/src/tool/select/README.md b/crates/tinyagents-harness/src/tool/select/README.md index 08bb14417..0405d6be5 100644 --- a/crates/tinyagents-harness/src/tool/select/README.md +++ b/crates/tinyagents-harness/src/tool/select/README.md @@ -61,6 +61,4 @@ select::*`) and has no dependency on the rest of `tool/` or on `tinytools` beyond the `SelectableTool` view a caller constructs. Its intended consumer is a host that binds a large third-party tool catalogue (e.g. a sub-agent scoped to hundreds of external actions) and needs to narrow it before -advertising it to a model; the `tinyagents-graph` crate's todo-dispatch -module (`crates/tinyagents-graph/src/todos/dispatch/`) is the current -caller. +advertising it to a model. diff --git a/crates/tinyagents-harness/src/tool/shared/mod_permissions_tests.rs b/crates/tinyagents-harness/src/tool/shared/mod_permissions_tests.rs index 68aae95d0..44d86a07e 100644 --- a/crates/tinyagents-harness/src/tool/shared/mod_permissions_tests.rs +++ b/crates/tinyagents-harness/src/tool/shared/mod_permissions_tests.rs @@ -122,14 +122,14 @@ async fn shared_target_finish_activity_distinguishes_preview_from_write() { )]; call.finish_reason = Some("tool_calls".into()); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( + harness.with_model_resolver(Arc::new(crate::host::FixedModelResolver::named( "mock", Arc::new(ScriptedModel::new(vec![ call, ModelResponse::assistant("draft"), ModelResponse::assistant("checked"), - ])), - ); + ])) as _, + ))); harness.register_tool(Arc::new(metadata_adapter(PermissionLevel::ReadOnly, false))); harness.with_policy(RunPolicy { limits: RunLimits::default().with_max_model_calls(10), diff --git a/crates/tinyagents-harness/src/tool/snapshot.rs b/crates/tinyagents-harness/src/tool/snapshot.rs index 5d2649cb0..4761d7a91 100644 --- a/crates/tinyagents-harness/src/tool/snapshot.rs +++ b/crates/tinyagents-harness/src/tool/snapshot.rs @@ -1,4 +1,9 @@ -use std::collections::BTreeMap; +use std::{ + any::TypeId, + collections::{BTreeMap, HashMap}, + fmt, + sync::Arc, +}; use tinytools::ToolSpec; @@ -19,14 +24,57 @@ pub enum ToolSnapshotError { Decode(String), } +/// Type-erased host-owned state carried alongside a tool snapshot. +/// +/// Extensions are retained by in-memory snapshot copies and are excluded from +/// the serialized tool declarations. +pub trait ToolSnapshotExtension: std::any::Any + fmt::Debug + Send + Sync { + fn as_any(&self) -> &dyn std::any::Any; +} + +impl<T: std::any::Any + fmt::Debug + Send + Sync> ToolSnapshotExtension for T { + fn as_any(&self) -> &dyn std::any::Any { + self + } +} + +#[derive(Clone, Default)] +pub struct ToolSnapshotExtensions { + values: HashMap<TypeId, Arc<dyn ToolSnapshotExtension>>, +} + +impl fmt::Debug for ToolSnapshotExtensions { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("ToolSnapshotExtensions") + .field("len", &self.values.len()) + .finish() + } +} + +impl ToolSnapshotExtensions { + pub fn insert<T: ToolSnapshotExtension>(&mut self, value: T) { + self.values.insert(TypeId::of::<T>(), Arc::new(value)); + } + + pub fn get<T: ToolSnapshotExtension>(&self) -> Option<&T> { + self.values + .get(&TypeId::of::<T>())? + .as_ref() + .as_any() + .downcast_ref::<T>() + } +} + /// An immutable, model-visible tool declaration set for one session turn. /// /// It records declarations only; choosing which tools are permitted and wiring -/// their executors remains a host/driver responsibility. +/// their executors remains a host/driver responsibility. Typed extensions stay +/// in memory and do not appear in the serialized transcript value. #[derive(Clone, Debug, Default)] pub struct ToolSnapshot { specs: Vec<ToolSpec>, exact: bool, + extensions: ToolSnapshotExtensions, } impl ToolSnapshot { @@ -50,6 +98,7 @@ impl ToolSnapshot { Ok(Self { specs: names.into_values().collect(), exact: false, + extensions: ToolSnapshotExtensions::default(), }) } @@ -69,6 +118,17 @@ impl ToolSnapshot { &self.specs } + /// Attaches an immutable, host-defined value keyed by its Rust type. + pub fn with_extension<T: ToolSnapshotExtension>(mut self, extension: T) -> Self { + self.extensions.insert(extension); + self + } + + /// Reads a host-defined extension without serializing it into tool specs. + pub fn extension<T: ToolSnapshotExtension>(&self) -> Option<&T> { + self.extensions.get::<T>() + } + /// A narrowed copy holding only the declarations `keep` accepts, with the /// one-off ([`Self::exact`]) marker preserved. The source is untouched. pub fn retaining(&self, keep: impl Fn(&ToolSpec) -> bool) -> Self { @@ -80,6 +140,7 @@ impl ToolSnapshot { .cloned() .collect(), exact: self.exact, + extensions: self.extensions.clone(), } } @@ -103,6 +164,9 @@ impl ToolSnapshot { &self, recorded: &ToolSnapshot, ) -> Result<(Self, usize), ToolSnapshotError> { + if self.exact { + return Ok((self.clone(), 0)); + } let retained: Vec<ToolSpec> = recorded .specs .iter() @@ -113,7 +177,8 @@ impl ToolSnapshot { return Ok((self.clone(), 0)); } let count = retained.len(); - let merged = Self::new(self.specs.iter().cloned().chain(retained).collect())?; + let mut merged = Self::new(self.specs.iter().cloned().chain(retained).collect())?; + merged.extensions = self.extensions.clone(); Ok((merged, count)) } } diff --git a/crates/tinyagents-harness/src/tool/snapshot_tests.rs b/crates/tinyagents-harness/src/tool/snapshot_tests.rs index b1244ebf1..69af29053 100644 --- a/crates/tinyagents-harness/src/tool/snapshot_tests.rs +++ b/crates/tinyagents-harness/src/tool/snapshot_tests.rs @@ -37,3 +37,55 @@ fn snapshot_json_round_trip_preserves_exact_flag_semantics() { assert_eq!(restored.to_json(), snapshot.to_json()); assert!(!restored.is_exact()); } + +#[derive(Clone, Debug, PartialEq, Eq)] +struct HostAlpha(u8); +#[derive(Clone, Debug, PartialEq, Eq)] +struct HostBeta(&'static str); + +#[test] +fn typed_extensions_survive_snapshot_copies_without_changing_json() { + let original_json = serde_json::json!([]); + let snapshot = ToolSnapshot::from_json(&original_json) + .unwrap() + .with_extension(HostAlpha(7)) + .with_extension(HostBeta("host-owned")); + assert_eq!(snapshot.extension::<HostAlpha>(), Some(&HostAlpha(7))); + assert_eq!( + snapshot.extension::<HostBeta>(), + Some(&HostBeta("host-owned")) + ); + assert_eq!(snapshot.to_json(), original_json); + + let clone = snapshot.clone(); + let narrowed = snapshot.retaining(|_| true); + let (retained, count) = snapshot.with_retained(&ToolSnapshot::default()).unwrap(); + assert_eq!(count, 0); + for copy in [&clone, &narrowed, &retained] { + assert_eq!(copy.extension::<HostAlpha>(), Some(&HostAlpha(7))); + assert_eq!(copy.extension::<HostBeta>(), Some(&HostBeta("host-owned"))); + assert_eq!(copy.to_json(), original_json); + } + let replaced = clone.with_extension(HostAlpha(8)); + assert_eq!(replaced.extension::<HostAlpha>(), Some(&HostAlpha(8))); + assert_eq!(snapshot.extension::<HostAlpha>(), Some(&HostAlpha(7))); +} + +#[test] +fn exact_snapshot_never_retains_recorded_declarations() { + let exact = ToolSnapshot::new(vec![spec("live", "live")]) + .unwrap() + .exact(); + let recorded = ToolSnapshot::new(vec![spec("recorded", "recorded")]).unwrap(); + let (merged, count) = exact.with_retained(&recorded).unwrap(); + assert_eq!(count, 0); + assert!(merged.is_exact()); + assert_eq!( + merged + .specs() + .iter() + .map(|tool| tool.name.as_str()) + .collect::<Vec<_>>(), + ["live"] + ); +} diff --git a/crates/tinyagents-harness/src/tool/toolset/mod.rs b/crates/tinyagents-harness/src/tool/toolset/mod.rs index e408aa196..f70c6e439 100644 --- a/crates/tinyagents-harness/src/tool/toolset/mod.rs +++ b/crates/tinyagents-harness/src/tool/toolset/mod.rs @@ -261,12 +261,6 @@ pub(crate) struct OverrideTool { pub(crate) description: Option<String>, pub(crate) parameters: Option<Value>, pub(crate) policy_transform: Option<Arc<dyn Fn(ToolPolicy) -> ToolPolicy + Send + Sync>>, - /// Overrides [`Tool::exposure`] instead of delegating to `inner`. Used by - /// [`crate::capability::CapabilityToolSet`] (gap G3) to apply a - /// [`crate::capability::Capability::exposure`] uniformly to every tool a - /// capability's inner [`ToolSet`] contributes, regardless of what each - /// individual tool declares for itself. - pub(crate) exposure_override: Option<ToolExposure>, } impl OverrideTool { @@ -277,17 +271,9 @@ impl OverrideTool { description: None, parameters: None, policy_transform: None, - exposure_override: None, } } - /// Overrides [`Tool::exposure`] with `exposure` instead of delegating to - /// `inner`. - pub(crate) fn with_exposure(mut self, exposure: ToolExposure) -> Self { - self.exposure_override = Some(exposure); - self - } - pub(crate) fn with_name(mut self, name: impl Into<String>) -> Self { self.name = Some(name.into()); self @@ -382,8 +368,7 @@ impl Tool for OverrideTool { } fn exposure(&self) -> ToolExposure { - self.exposure_override - .unwrap_or_else(|| self.inner.exposure()) + self.inner.exposure() } // A renamed or prefixed tool is still the same tool to the host's tool diff --git a/crates/tinyagents-harness/src/turn_outcome.rs b/crates/tinyagents-harness/src/turn_outcome.rs index 78c1cf879..644620c06 100644 --- a/crates/tinyagents-harness/src/turn_outcome.rs +++ b/crates/tinyagents-harness/src/turn_outcome.rs @@ -9,6 +9,67 @@ use std::sync::{Arc, Mutex}; use tinyinference_llm::model::ModelResponse; +/// Event-only terminal token for hosts that finish post-run close work before +/// committing durable session status. +#[derive(Clone)] +pub struct DeferredTurnTerminal { + pub run_id: crate::ids::RunId, + pub events: crate::events::EventSink, +} + +impl std::fmt::Debug for DeferredTurnTerminal { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("DeferredTurnTerminal") + .finish_non_exhaustive() + } +} + +impl DeferredTurnTerminal { + /// Emit the optional stop and successful terminal events in order. + pub fn finish(self, stop: Option<crate::turn_stop::TurnStop>, turn: u32) { + if let Some(stop) = stop { + self.events + .emit(crate::events::AgentEvent::TurnStopped { turn, stop }); + } + self.events.emit(crate::events::AgentEvent::RunFinalized { + run_id: self.run_id, + }); + } + + /// Emit content-free cancellation and finalization events. + pub fn finish_cancelled(self) { + use crate::terminal::{TerminalOutcome, TerminalReason}; + self.events.emit(crate::events::AgentEvent::RunFailed { + run_id: self.run_id.clone(), + error: "turn cancelled after harness completion".into(), + outcome: Some(TerminalOutcome::new( + TerminalReason::Cancelled, + "turn cancelled after harness completion", + )), + }); + self.events.emit(crate::events::AgentEvent::RunFinalized { + run_id: self.run_id, + }); + } + + /// Emit a generic failure and finalization without leaking runtime error text. + pub fn finish_failed(self) { + use crate::terminal::{TerminalOutcome, TerminalReason}; + self.events.emit(crate::events::AgentEvent::RunFailed { + run_id: self.run_id.clone(), + error: "host turn failed after harness completion".into(), + outcome: Some(TerminalOutcome::new( + TerminalReason::Internal, + "host turn failed after harness completion", + )), + }); + self.events.emit(crate::events::AgentEvent::RunFinalized { + run_id: self.run_id, + }); + } +} + /// Whether a run's final response is a reply that ran out of output tokens /// before producing anything: `finish_reason = length`, no visible text and no /// tool call. @@ -27,7 +88,7 @@ pub fn ended_out_of_output_budget(final_response: Option<&ModelResponse>) -> boo /// One tool call's execution outcome, captured at the tool boundary before the /// harness folds the result into a message that drops its failure flag. -#[derive(Debug, Clone)] +#[derive(Debug, Clone, PartialEq)] pub struct ToolCallOutcome { /// Provider call id. pub call_id: String, diff --git a/crates/tinyagents-harness/src/turn_outcome_tests.rs b/crates/tinyagents-harness/src/turn_outcome_tests.rs index d296fc7ea..d28462180 100644 --- a/crates/tinyagents-harness/src/turn_outcome_tests.rs +++ b/crates/tinyagents-harness/src/turn_outcome_tests.rs @@ -1,4 +1,108 @@ use super::*; + +#[test] +fn deferred_terminal_emits_stop_before_finalization() { + use crate::events::{AgentEvent, EventListener, EventSink, RecordingListener}; + use crate::ids::RunId; + let events = EventSink::new(); + let listener = Arc::new(RecordingListener::new()); + events.subscribe(listener.clone() as Arc<dyn EventListener>); + DeferredTurnTerminal { + run_id: RunId::new("run"), + events, + } + .finish(Some(crate::turn_stop::TurnStop::wind_down()), 4); + let events = listener.events(); + assert_eq!(events.len(), 2); + assert!(matches!( + events[0].event, + AgentEvent::TurnStopped { turn: 4, .. } + )); + assert!(matches!(events[1].event, AgentEvent::RunFinalized { .. })); +} + +#[test] +fn deferred_terminal_failure_emits_failure_before_finalization() { + use crate::events::{AgentEvent, EventListener, EventSink, RecordingListener}; + use crate::ids::RunId; + let events = EventSink::new(); + let listener = Arc::new(RecordingListener::new()); + events.subscribe(listener.clone() as Arc<dyn EventListener>); + DeferredTurnTerminal { + run_id: RunId::new("run"), + events, + } + .finish_cancelled(); + let events = listener.events(); + assert_eq!(events.len(), 2); + assert!( + matches!(&events[0].event, AgentEvent::RunFailed { error, .. } if error == "turn cancelled after harness completion") + ); + assert!(matches!(events[1].event, AgentEvent::RunFinalized { .. })); +} + +#[test] +fn deferred_terminal_failure_uses_content_free_error_and_outcome() { + use crate::events::{AgentEvent, EventListener, EventSink, RecordingListener}; + use crate::ids::RunId; + use crate::terminal::{TerminalOutcome, TerminalReason}; + let events = EventSink::new(); + let listener = Arc::new(RecordingListener::new()); + events.subscribe(listener.clone() as Arc<dyn EventListener>); + DeferredTurnTerminal { + run_id: RunId::new("run-failed"), + events, + } + .finish_failed(); + let events = listener.events(); + assert_eq!(events.len(), 2); + assert!(matches!( + &events[0].event, + AgentEvent::RunFailed { run_id, error, outcome: Some(outcome) } + if run_id.as_str() == "run-failed" + && error == "host turn failed after harness completion" + && outcome == &TerminalOutcome::new( + TerminalReason::Internal, + "host turn failed after harness completion" + ) + )); + assert!(matches!( + &events[1].event, + AgentEvent::RunFinalized { run_id } if run_id.as_str() == "run-failed" + )); +} + +#[test] +fn tool_call_outcome_equality_covers_every_observed_field() { + let original = ToolCallOutcome { + call_id: "call-1".into(), + name: "read_file".into(), + arguments: serde_json::json!({"path": "notes.txt"}), + success: false, + content: "denied".into(), + duration_ms: 17, + }; + assert_eq!(original, original.clone()); + + let mut changed = original.clone(); + changed.call_id.push_str("-changed"); + assert_ne!(original, changed); + let mut changed = original.clone(); + changed.name.push_str("_changed"); + assert_ne!(original, changed); + let mut changed = original.clone(); + changed.arguments["path"] = serde_json::json!("other.txt"); + assert_ne!(original, changed); + let mut changed = original.clone(); + changed.success = true; + assert_ne!(original, changed); + let mut changed = original.clone(); + changed.content.push_str(" changed"); + assert_ne!(original, changed); + let mut changed = original.clone(); + changed.duration_ms += 1; + assert_ne!(original, changed); +} use tinyinference_llm::message::{AssistantMessage, ContentBlock}; use tinyinference_llm::tool::ToolCall; diff --git a/crates/tinyagents-harness/src/turn_stop.rs b/crates/tinyagents-harness/src/turn_stop.rs index 03c76597e..d720d3c95 100644 --- a/crates/tinyagents-harness/src/turn_stop.rs +++ b/crates/tinyagents-harness/src/turn_stop.rs @@ -14,10 +14,134 @@ //! content-free by construction: a typed breaker cause carries only stable //! class and operation identifiers, never the quoted tool output or target. +use crate::close::{CloseToolResult, interpolate_once}; use serde::de::Error as _; use serde::ser::SerializeStruct; use serde::{Deserialize, Deserializer, Serialize, Serializer}; +/// Host-provided prose and line templates for deterministic stop summaries. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct StopCopy { + pub opening_template: &'static str, + pub what_happened_heading: &'static str, + pub success_template: &'static str, + pub failure_template: &'static str, + pub hidden_calls_template: &'static str, + pub proceed_prompt: &'static str, +} + +/// English renderer copy for stopped-turn summaries. +pub const ENGLISH_STOP_COPY: StopCopy = StopCopy { + opening_template: "I stopped this turn early because {reason}.\n", + what_happened_heading: "\n**What happened**\n", + success_template: "- `{name}` succeeded{count}", + failure_template: "- `{name}` failed{count}{error}", + hidden_calls_template: "- …and {count} more tool call(s)", + proceed_prompt: "\nTell me how you'd like to proceed, or ask me to try again.", +}; + +/// Summary row grouped by operation and scrubbed result text. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct StopSummaryGroup { + pub name: String, + pub success: bool, + pub error: Option<String>, + pub count: usize, +} + +/// Group repeated results after the host has scrubbed failure text. +pub fn group_stop_results<F>( + results: &[CloseToolResult], + mut scrubbed_error_line: F, +) -> Vec<StopSummaryGroup> +where + F: FnMut(&str) -> String, +{ + let mut groups: Vec<StopSummaryGroup> = Vec::new(); + for result in results { + let error = if result.success { + None + } else { + Some(scrubbed_error_line(&result.content)).filter(|line| !line.is_empty()) + }; + if let Some(group) = groups + .iter_mut() + .find(|g| g.name == result.name && g.success == result.success && g.error == error) + { + group.count += 1; + } else { + groups.push(StopSummaryGroup { + name: result.name.clone(), + success: result.success, + error, + count: 1, + }); + } + } + groups +} + +/// Render a bounded stop summary with caller-selected product copy. +pub fn render_stop_summary_with_copy( + reason: &str, + groups: &[StopSummaryGroup], + copy: &StopCopy, +) -> String { + use std::fmt::Write as _; + let mut out = interpolate_once(copy.opening_template, &[("{reason}", reason)]); + if !groups.is_empty() { + out.push_str(copy.what_happened_heading); + let shown_from = groups.len().saturating_sub(8); + for group in groups.iter().skip(shown_from) { + let name = truncate_stop_name(&group.name, 60); + let count = if group.count == 1 { + String::new() + } else { + format!(" {} times", group.count) + }; + let line = if group.success { + interpolate_once( + copy.success_template, + &[("{name}", &name), ("{count}", &count)], + ) + } else { + let error = group + .error + .as_ref() + .map(|e| format!(": {e}")) + .unwrap_or_default(); + interpolate_once( + copy.failure_template, + &[("{name}", &name), ("{count}", &count), ("{error}", &error)], + ) + }; + let _ = writeln!(out, "{line}"); + } + let hidden: usize = groups + .iter() + .take(shown_from) + .map(|group| group.count) + .sum(); + if hidden > 0 { + let hidden = hidden.to_string(); + let _ = writeln!( + out, + "{}", + interpolate_once(copy.hidden_calls_template, &[("{count}", &hidden)]) + ); + } + } + out.push_str(copy.proceed_prompt); + out +} + +fn truncate_stop_name(text: &str, max_chars: usize) -> String { + match text.char_indices().nth(max_chars) { + Some((idx, _)) => format!("{}...", text[..idx].trim_end()), + None => text.to_owned(), + } +} + /// Which mechanism stopped the turn. #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] #[serde(rename_all = "snake_case")] diff --git a/crates/tinyagents-harness/src/turn_stop_tests.rs b/crates/tinyagents-harness/src/turn_stop_tests.rs index 849efb371..662326e4f 100644 --- a/crates/tinyagents-harness/src/turn_stop_tests.rs +++ b/crates/tinyagents-harness/src/turn_stop_tests.rs @@ -1,5 +1,148 @@ use super::*; +#[test] +fn stop_summary_with_copy_preserves_placeholder_shaped_values() { + let groups = [StopSummaryGroup { + name: "tool_{count}_{error}".to_owned(), + success: false, + error: Some("failed {name}".to_owned()), + count: 3, + }]; + let copy = StopCopy { + opening_template: "reason: {reason}\n", + failure_template: "- `{name}` failed{count}{error}", + ..ENGLISH_STOP_COPY + }; + assert_eq!( + render_stop_summary_with_copy("reason {count}", &groups, ©), + "reason: reason {count}\n\n**What happened**\n- `tool_{count}_{error}` failed 3 times: failed {name}\n\nTell me how you'd like to proceed, or ask me to try again." + ); +} + +#[test] +fn stop_copy_keeps_transient_summary_golden() { + let groups = [ + StopSummaryGroup { + name: "get_profile".into(), + success: true, + error: None, + count: 1, + }, + StopSummaryGroup { + name: "use_skill".into(), + success: false, + error: Some("tool `use_skill` timed out after 30000 ms".into()), + count: 3, + }, + ]; + assert_eq!( + render_stop_summary_with_copy( + "a service I needed kept timing out or was temporarily unavailable", + &groups, + &ENGLISH_STOP_COPY + ), + "I stopped this turn early because a service I needed kept timing out or was temporarily unavailable.\n\n**What happened**\n- `get_profile` succeeded\n- `use_skill` failed 3 times: tool `use_skill` timed out after 30000 ms\n\nTell me how you'd like to proceed, or ask me to try again." + ); +} + +#[test] +fn stop_summary_groups_scrubbed_failures_and_keeps_the_newest_eight_rows() { + let mut results = vec![CloseToolResult { + name: "oldest".into(), + success: true, + content: "ok".into(), + }]; + for index in 0..8 { + results.push(CloseToolResult { + name: format!("tool-{index}"), + success: false, + content: format!("secret-{index}"), + }); + } + results.extend([ + CloseToolResult { + name: "repeat".into(), + success: false, + content: "credential: hidden".into(), + }, + CloseToolResult { + name: "repeat".into(), + success: false, + content: "credential: hidden".into(), + }, + ]); + let groups = group_stop_results(&results, |content| { + content + .replace("credential: hidden", "credential redacted") + .replace("secret-", "redacted-") + }); + assert_eq!(groups.last().unwrap().count, 2); + assert_eq!( + groups.last().unwrap().error.as_deref(), + Some("credential redacted") + ); + let output = render_stop_summary_with_copy("reason", &groups, &ENGLISH_STOP_COPY); + assert!(!output.contains("oldest")); + assert!(!output.contains("secret-")); + assert!(output.contains("- …and 2 more tool call(s)")); + assert!(output.contains("- `repeat` failed 2 times: credential redacted")); +} + +#[test] +fn stop_summary_reproduces_the_production_repeated_timeout_golden() { + let results = [ + CloseToolResult { + name: "get_profile".into(), + success: true, + content: "profile loaded".into(), + }, + CloseToolResult { + name: "use_skill".into(), + success: false, + content: "tool `use_skill` timed out after 30000 ms".into(), + }, + CloseToolResult { + name: "use_skill".into(), + success: false, + content: "tool `use_skill` timed out after 30000 ms".into(), + }, + CloseToolResult { + name: "use_skill".into(), + success: false, + content: "tool `use_skill` timed out after 30000 ms".into(), + }, + ]; + let groups = group_stop_results(&results, str::to_owned); + assert_eq!( + render_stop_summary_with_copy( + "a service I needed kept timing out or was temporarily unavailable", + &groups, + &ENGLISH_STOP_COPY, + ), + "I stopped this turn early because a service I needed kept timing out or was temporarily unavailable.\n\n**What happened**\n- `get_profile` succeeded\n- `use_skill` failed 3 times: tool `use_skill` timed out after 30000 ms\n\nTell me how you'd like to proceed, or ask me to try again." + ); +} + +#[test] +fn successful_repeats_are_reported_without_a_failure_claim() { + let results = (0..3) + .map(|_| CloseToolResult { + name: "list_items".into(), + success: true, + content: "three items".into(), + }) + .collect::<Vec<_>>(); + let groups = group_stop_results(&results, str::to_owned); + let output = + render_stop_summary_with_copy("the same step repeated", &groups, &ENGLISH_STOP_COPY); + assert_eq!( + output, + "I stopped this turn early because the same step repeated.\n\n**What happened**\n- `list_items` succeeded 3 times\n\nTell me how you'd like to proceed, or ask me to try again." + ); + assert!(!output.to_lowercase().contains("failed")); + assert!(!output.contains("three items")); +} + fn fallback(note: &str) -> &'static str { if note.contains("in a row failed") { "turn_stop.varied_failures" diff --git a/crates/tinyagents-harness/src/url_guard.rs b/crates/tinyagents-harness/src/url_guard.rs new file mode 100644 index 000000000..610eaf592 --- /dev/null +++ b/crates/tinyagents-harness/src/url_guard.rs @@ -0,0 +1,497 @@ +//! Shared URL validation + SSRF guards for outbound network tools. +//! +//! Used by `http_request`, `curl`, and any future tool that takes a +//! user-supplied URL. Two allowlist modes: +//! +//! - **Open allowlist** (`allowed_domains` is empty): any public non-private +//! host is permitted. All SSRF guards still apply (loopback / RFC1918 / +//! link-local / multicast / documentation / shared-address / +//! IPv4-mapped and transition IPv6, `localhost` / `*.localhost` / `*.local`). +//! - **Strict allowlist** (`allowed_domains` is non-empty): only the listed +//! domains and their subdomains are permitted. +//! +//! Both modes enforce: http(s) only, no whitespace, no userinfo, no IPv6 hosts, +//! no backslash anywhere, and no percent-encoding in the host — the last two +//! because a WHATWG parser would read them as a different host than the one +//! checked here. +//! +//! **Alternate IP notations** (octal, hex, decimal): Rust's `IpAddr::parse` +//! rejects them so they are treated as plain hostnames. In strict-allowlist +//! mode they are rejected by the domain check. In open-allowlist mode they +//! pass `validate_url` but are caught by `validate_url_with_dns_check` +//! because they fail real-world DNS resolution. +//! +//! ## DNS Rebinding +//! +//! Hostname validation alone is insufficient: an attacker can register a +//! domain that alternates DNS responses between a public IP (passing the +//! allowlist) and a private IP (e.g. 127.0.0.1). +//! [`crate::url_guard::validate_url_with_dns_check`] resolves the hostname, vets every +//! resolved IP, and returns them in a [`crate::url_guard::ValidatedUrl`]. That closes the gap +//! **only if the caller connects to [`crate::url_guard::ValidatedUrl::addrs`]** — for example +//! via `reqwest::ClientBuilder::resolve_to_addrs` — rather than letting its +//! HTTP client resolve the hostname a second time. This crate carries no +//! HTTP client, so the pinning is the caller's responsibility. This utility +//! deliberately does not own an HTTP client or host network policy. + +use std::future::Future; +use std::net::{IpAddr, SocketAddr, ToSocketAddrs}; + +/// Validate a URL against the allowlist + SSRF rules. Returns the +/// original URL on success. +/// +/// # Errors +/// +/// Fails when the URL is empty, contains whitespace or a backslash, is not +/// `http(s)`, has percent-encoding in its host, names a local/private host, +/// or (in strict mode) is outside the allowlist. +pub fn validate_url(raw_url: &str, allowed_domains: &[String]) -> anyhow::Result<String> { + let url = raw_url.trim(); + + if url.is_empty() { + anyhow::bail!("URL cannot be empty"); + } + + if url.chars().any(char::is_whitespace) { + anyhow::bail!("URL cannot contain whitespace"); + } + + if !url.starts_with("http://") && !url.starts_with("https://") { + anyhow::bail!("Only http:// and https:// URLs are allowed"); + } + + reject_backslash(url)?; + + let host = extract_host(url)?; + + if is_private_or_local_host(&host) { + tracing::debug!( + "[url_guard] ssrf block: host={host} mode={}", + if allowed_domains.is_empty() { + "open" + } else { + "strict" + } + ); + anyhow::bail!("Blocked local/private host: {host}"); + } + + // Empty allowed_domains = open mode: any public non-private host is + // permitted (same as ["*"]). This ensures the http_request tool works + // out of the box regardless of whether the user configured an explicit + // domain list, and keeps web-fetch consistent across routing paths. + // A non-empty list = strict mode: only listed domains pass. (#2700) + if !allowed_domains.is_empty() && !host_matches_allowlist(&host, allowed_domains) { + tracing::debug!( + "[url_guard] strict-allowlist rejection: host={host} allowed={allowed_domains:?}" + ); + anyhow::bail!( + "I'm not allowed to open '{host}' — it isn't in your allowed websites. \ + Add it (or turn on \"Allow all sites\") under \ + Settings → Advanced → Search engine → Allowed websites, then ask me again." + ); + } + + tracing::debug!( + "[url_guard] validate_url ok: host={host} mode={}", + if allowed_domains.is_empty() { + "open" + } else { + "strict" + } + ); + + Ok(url.to_string()) +} + +/// A URL that passed [`validate_url_with_dns_check`], together with the +/// exact socket addresses that were vetted. +/// +/// The addresses are the point: DNS can answer differently the next time it +/// is asked, so a client that re-resolves `host` may connect somewhere that +/// was never checked. Pin the connection to [`addrs`](Self::addrs) instead — +/// for example with `reqwest::ClientBuilder::resolve_to_addrs(&host, &addrs)` +/// — and keep `url` unchanged so TLS SNI and the `Host` header still name +/// `host`. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ValidatedUrl { + /// The validated URL, trimmed, otherwise exactly as supplied. + pub url: String, + /// The lowercase host the URL names (a hostname or an IP literal). + pub host: String, + /// Every address `host` resolved to, each paired with the URL's port; + /// all are public. For an IP-literal host this is that single address. + pub addrs: Vec<SocketAddr>, +} + +impl ValidatedUrl { + /// The validated URL, retaining its hostname as the request authority. + #[must_use] + pub fn url(&self) -> &str { + &self.url + } + + /// Vetted connection destinations to pin the HTTP client to. + #[must_use] + pub fn addresses(&self) -> &[SocketAddr] { + &self.addrs + } +} + +/// Like [`validate_url`] but also resolves the hostname via DNS and +/// verifies that none of the resolved IPs are private/local. +/// +/// This only defends against DNS rebinding — an attacker's domain answering +/// with a public IP here and 127.0.0.1 at request time — when the caller +/// connects to the returned [`ValidatedUrl::addrs`] rather than resolving +/// the hostname again. Callers should use this function instead of +/// `validate_url` in all paths that make outbound HTTP requests, and pin +/// the connection as described on [`ValidatedUrl`]. +/// +/// # Errors +/// +/// Everything [`validate_url`] rejects, plus DNS failure, an empty answer, or +/// any resolved address that is private/local. +pub async fn validate_url_with_dns_check( + raw_url: &str, + allowed_domains: &[String], +) -> anyhow::Result<ValidatedUrl> { + validate_url_with_dns_check_with_resolver(raw_url, allowed_domains, resolve_host_ips).await +} + +async fn validate_url_with_dns_check_with_resolver<F, Fut>( + raw_url: &str, + allowed_domains: &[String], + resolver: F, +) -> anyhow::Result<ValidatedUrl> +where + F: FnOnce(String, u16) -> Fut, + Fut: Future<Output = anyhow::Result<Vec<IpAddr>>>, +{ + let url = validate_url(raw_url, allowed_domains)?; + + let host = extract_host(&url)?; + + let port = extract_port(&url)?; + + // If the host is already a valid IP literal, `is_private_or_local_host` + // has already checked it above. We only need DNS resolution for hostnames. + if let Ok(ip) = host.parse::<IpAddr>() { + return Ok(ValidatedUrl { + url, + host, + addrs: vec![SocketAddr::new(ip, port)], + }); + } + + tracing::debug!("[url_guard] resolving DNS for host={host} port={port}"); + let addrs = resolver(host.clone(), port).await?; + + if addrs.is_empty() { + anyhow::bail!("DNS resolution returned no addresses for '{host}'"); + } + + tracing::debug!("[url_guard] DNS resolved host={host} addrs={}", addrs.len()); + + for addr in &addrs { + let ip_str = addr.to_string(); + if is_private_or_local_host(&ip_str) { + tracing::debug!("[url_guard] DNS rebinding blocked host={host} resolved_ip={ip_str}"); + anyhow::bail!( + "DNS rebinding blocked: '{host}' resolved to private/local address {ip_str}" + ); + } + } + + Ok(ValidatedUrl { + url, + host, + addrs: addrs + .into_iter() + .map(|ip| SocketAddr::new(ip, port)) + .collect(), + }) +} + +async fn resolve_host_ips(host: String, port: u16) -> anyhow::Result<Vec<IpAddr>> { + let log_host = host.clone(); + tokio::task::spawn_blocking(move || { + (host.as_str(), port) + .to_socket_addrs() + .map_err(|e| { + tracing::debug!( + "[url_guard] DNS resolution failed host={host} port={port} error={e}" + ); + anyhow::anyhow!("DNS resolution failed for '{host}': {e}") + }) + .map(|iter| iter.map(|addr| addr.ip()).collect()) + }) + .await + .map_err(|e| { + tracing::debug!( + "[url_guard] DNS resolution task failed host={log_host} port={port} error={e}" + ); + anyhow::anyhow!("DNS resolution task failed for '{log_host}': {e}") + })? +} + +#[must_use] +/// Normalise an allowlist: strip scheme/path, lowercase, drop invalid entries and +/// duplicates. An empty result means open mode. +pub fn normalize_allowed_domains(domains: Vec<String>) -> Vec<String> { + if domains.is_empty() { + return Vec::new(); + } + let mut normalized = domains + .into_iter() + .filter_map(|d| normalize_domain(&d)) + .collect::<Vec<_>>(); + normalized.sort_unstable(); + normalized.dedup(); + if normalized.is_empty() { + // All entries were malformed (whitespace-only, scheme-only, etc.) and + // filtered out. Returning empty would silently enter open mode; instead + // return a sentinel that keeps the tool in strict mode and rejects every + // URL — fail-closed on misconfiguration. (#2738) + tracing::warn!( + "[url_guard] all configured allowed_domains entries are invalid — \ + treating as misconfigured allowlist (fail-closed)" + ); + return vec!["<misconfigured-allowlist>".to_string()]; + } + normalized +} + +/// Normalise one allowlist entry to a bare lowercase host, or `None` if invalid. +pub fn normalize_domain(raw: &str) -> Option<String> { + let mut d = raw.trim().to_lowercase(); + if d.is_empty() { + return None; + } + + if let Some(stripped) = d.strip_prefix("https://") { + d = stripped.to_string(); + } else if let Some(stripped) = d.strip_prefix("http://") { + d = stripped.to_string(); + } + + if let Some((host, _)) = d.split_once('/') { + d = host.to_string(); + } + + d = d.trim_start_matches('.').trim_end_matches('.').to_string(); + + if let Some((host, _)) = d.split_once(':') { + d = host.to_string(); + } + + if d.is_empty() || d.chars().any(char::is_whitespace) { + return None; + } + + Some(d) +} + +/// Refuse a URL containing `\` anywhere. WHATWG URL parsers (browsers, +/// `reqwest`'s `url` crate) treat `\` as a path separator in `http(s)` URLs, +/// so `http://127.0.0.1\.example.com/` names `127.0.0.1` to a real client +/// while a naive split on `/` would see a subdomain of `example.com`. +fn reject_backslash(url: &str) -> anyhow::Result<()> { + if url.contains('\\') { + anyhow::bail!("URL cannot contain a backslash"); + } + Ok(()) +} + +/// Split an `http(s)` URL into whether it is plain `http` and its authority +/// (`host[:port]`), refusing inputs a WHATWG parser would read differently. +fn split_authority(url: &str) -> anyhow::Result<(bool, &str)> { + let (is_http, rest) = if let Some(rest) = url.strip_prefix("http://") { + (true, rest) + } else if let Some(rest) = url.strip_prefix("https://") { + (false, rest) + } else { + anyhow::bail!("Only http:// and https:// URLs are allowed"); + }; + + reject_backslash(url)?; + + let authority = rest + .split(['/', '?', '#']) + .next() + .ok_or_else(|| anyhow::anyhow!("Invalid URL"))?; + + // WHATWG percent-decodes the host, so `%31%32%37.0.0.1` is 127.0.0.1 on + // the wire while the literal text matches neither the SSRF checks nor + // the allowlist. + if authority.contains('%') { + anyhow::bail!("URL host cannot contain percent-encoded characters"); + } + + if authority.starts_with('[') { + anyhow::bail!("IPv6 hosts are not supported in http_request"); + } + + Ok((is_http, authority)) +} + +/// Extract the host part of an `http(s)` URL. +/// +/// # Errors +/// +/// Fails on a missing/empty host, userinfo, an IPv6 literal, a backslash +/// anywhere in the URL, or percent-encoding in the authority. +pub fn extract_host(url: &str) -> anyhow::Result<String> { + let (_, authority) = split_authority(url)?; + + if authority.is_empty() { + anyhow::bail!("URL must include a host"); + } + + if authority.contains('@') { + anyhow::bail!("URL userinfo is not allowed"); + } + + let host = authority + .split(':') + .next() + .unwrap_or_default() + .trim() + .trim_end_matches('.') + .to_lowercase(); + + if host.is_empty() { + anyhow::bail!("URL must include a valid host"); + } + + Ok(host) +} + +/// Extract the explicit or scheme-default port of an `http(s)` URL. +/// +/// # Errors +/// +/// Fails when the URL has no valid port, is an IPv6 literal, contains a +/// backslash, or has percent-encoding in the authority. +pub fn extract_port(url: &str) -> anyhow::Result<u16> { + let (is_http, authority) = split_authority(url)?; + + if let Some((_, port)) = authority.rsplit_once(':') { + if port.is_empty() || !port.chars().all(|ch| ch.is_ascii_digit()) { + anyhow::bail!("URL port must be numeric"); + } + return port + .parse::<u16>() + .map_err(|_| anyhow::anyhow!("URL port is out of range")); + } + + Ok(if is_http { 80 } else { 443 }) +} + +#[must_use] +/// Whether `host` equals, or is a subdomain of, an allowlist entry. +pub fn host_matches_allowlist(host: &str, allowed_domains: &[String]) -> bool { + allowed_domains.iter().any(|domain| { + // `"*"` is the explicit allow-all wildcard (the "Allow all sites" + // toggle), mirroring the browser tool. Local/private hosts are still + // rejected upstream by `is_private_or_local_host`, so a wildcard only + // opens *public* hosts, never the loopback/RFC1918 SSRF surface. + domain == "*" + || host == domain + || host + .strip_suffix(domain) + .is_some_and(|prefix| prefix.ends_with('.')) + }) +} + +#[must_use] +/// Whether `host` is a local name or resolves lexically to a non-global address. +pub fn is_private_or_local_host(host: &str) -> bool { + let unbracketed = host + .strip_prefix('[') + .and_then(|h| h.strip_suffix(']')) + .unwrap_or(host); + let bare = unbracketed.strip_suffix('.').unwrap_or(unbracketed); + + let lower = bare.to_ascii_lowercase(); + + let has_local_tld = lower + .rsplit('.') + .next() + .is_some_and(|label| label == "local"); + + if lower == "localhost" || lower.ends_with(".localhost") || has_local_tld { + return true; + } + + if let Ok(ip) = bare.parse::<std::net::IpAddr>() { + return match ip { + std::net::IpAddr::V4(v4) => is_non_global_v4(v4), + std::net::IpAddr::V6(v6) => is_non_global_v6(v6), + }; + } + + false +} + +#[must_use] +/// Whether an IPv4 address is non-global (loopback, private, link-local, ...). +pub fn is_non_global_v4(v4: std::net::Ipv4Addr) -> bool { + let [a, b, c, _] = v4.octets(); + v4.is_loopback() + || v4.is_private() + || v4.is_link_local() + || v4.is_unspecified() + || v4.is_broadcast() + || v4.is_multicast() + || (a == 100 && (64..=127).contains(&b)) + || a >= 240 + || (a == 192 && b == 0 && c == 0) + || (a == 192 && b == 88 && c == 99) + || (a == 198 && b == 51 && c == 100) + || (a == 203 && b == 0 && c == 113) + || (a == 198 && (18..=19).contains(&b)) + // 0.0.0.0/8 — "this network" (RFC 1122 §3.2.1.3). `is_unspecified()` only + // covers 0.0.0.0 itself, but the whole /8 routes to the local host on + // Linux. Carried over from the `ops_install` copy this replaces. + || a == 0 +} + +/// Whether an IPv6 address is non-global, including private IPv4 destinations +/// embedded in transition addresses. +pub fn is_non_global_v6(v6: std::net::Ipv6Addr) -> bool { + let segs = v6.segments(); + let embedded_v4 = |hi: u16, lo: u16| { + let [first, second] = hi.to_be_bytes(); + let [third, fourth] = lo.to_be_bytes(); + is_non_global_v4(std::net::Ipv4Addr::new(first, second, third, fourth)) + }; + v6.is_loopback() + || v6.is_unspecified() + || v6.is_multicast() + || (segs[0] & 0xfe00) == 0xfc00 + || (segs[0] & 0xffc0) == 0xfe80 + || (segs[0] == 0x2001 && segs[1] == 0x0db8) + || (segs[0] == 0x0100 && segs[1] == 0 && segs[2] == 0 && segs[3] <= 1) + || (segs[0] == 0x2001 && segs[1] == 0x0002 && segs[2] == 0) + // Local-use translation (RFC 8215) and the well-known NAT64 prefix + // can embed addresses that translate to private IPv4 destinations. + || (segs[0] == 0x0064 && segs[1] == 0xff9b && segs[2] == 1) + || (segs[0] == 0x0064 + && segs[1] == 0xff9b + && segs[2..6] == [0; 4] + && embedded_v4(segs[6], segs[7])) + // 6to4 carries the destination IPv4 address immediately after 2002::/16. + || (segs[0] == 0x2002 && embedded_v4(segs[1], segs[2])) + // Teredo carries a server IPv4 address, but the destination is the + // client's XOR-obfuscated IPv4 address in the last two segments. + || (segs[..2] == [0x2001, 0] && embedded_v4(!segs[6], !segs[7])) + || (segs[0] & 0xfff0) == 0x3ff0 + || segs[0] == 0x5f00 + // `to_ipv4` covers both mapped and deprecated compatible addresses. + || v6.to_ipv4().is_some_and(is_non_global_v4) +} + +#[cfg(test)] +#[path = "url_guard_tests.rs"] +mod tests; diff --git a/crates/tinyagents-harness/src/url_guard_tests.rs b/crates/tinyagents-harness/src/url_guard_tests.rs new file mode 100644 index 000000000..ec2c82004 --- /dev/null +++ b/crates/tinyagents-harness/src/url_guard_tests.rs @@ -0,0 +1,770 @@ +//! Unit tests for URL validation, SSRF host classification, and the DNS +//! check's vetted addresses. + +use super::*; + +/// Turns an expected rejection into its message, and an unexpected success +/// into a test failure, without `unwrap_err`. +trait Rejection { + fn rejection(self) -> anyhow::Result<String>; +} + +impl<T: std::fmt::Debug> Rejection for anyhow::Result<T> { + fn rejection(self) -> anyhow::Result<String> { + match self { + Ok(value) => anyhow::bail!("expected a rejection, got {value:?}"), + Err(err) => Ok(err.to_string()), + } + } +} + +#[test] +fn normalize_domain_strips_scheme_path_and_case() { + let got = normalize_domain(" HTTPS://Docs.Example.com/path "); + assert_eq!(got.as_deref(), Some("docs.example.com")); +} + +#[test] +fn normalizes_http_domains_and_rejects_empty_hosts() { + assert_eq!( + normalize_domain("http://Example.com:8080/path"), + Some("example.com".into()) + ); + assert_eq!(normalize_domain("https://"), None); + assert!(extract_host("http:///path").is_err()); + assert!(extract_host("http://:80/path").is_err()); +} + +#[test] +fn rejects_malformed_ports() { + assert!(extract_port("http://example.com:abc").is_err()); + assert!(extract_port("http://example.com:65536").is_err()); +} + +#[tokio::test] +async fn system_dns_resolves_numeric_loopback_without_external_network() -> anyhow::Result<()> { + let resolved = super::resolve_host_ips("127.0.0.1".to_string(), 80).await?; + assert_eq!(resolved, vec!["127.0.0.1".parse::<std::net::IpAddr>()?]); + assert!(super::resolve_host_ips(String::new(), 80).await.is_err()); + Ok(()) +} + +#[test] +fn normalize_allowed_domains_deduplicates() { + let got = normalize_allowed_domains(vec![ + "example.com".into(), + "EXAMPLE.COM".into(), + "https://example.com/".into(), + ]); + assert_eq!(got, vec!["example.com".to_string()]); +} + +#[test] +fn validate_accepts_exact_domain() -> anyhow::Result<()> { + let allow = vec!["example.com".to_string()]; + let got = validate_url("https://example.com/docs", &allow)?; + assert_eq!(got, "https://example.com/docs"); + Ok(()) +} + +#[test] +fn validate_accepts_http() { + let allow = vec!["example.com".to_string()]; + assert!(validate_url("http://example.com", &allow).is_ok()); +} + +#[test] +fn validate_accepts_subdomain() { + let allow = vec!["example.com".to_string()]; + assert!(validate_url("https://api.example.com/v1", &allow).is_ok()); +} + +#[test] +fn validate_rejects_allowlist_miss() -> anyhow::Result<()> { + let allow = vec!["example.com".to_string()]; + let err = validate_url("https://google.com", &allow).rejection()?; + assert!(err.contains("allowed websites")); + Ok(()) +} + +#[test] +fn validate_wildcard_allows_any_public_host() { + let allow = vec!["*".to_string()]; + assert!(validate_url("https://example.com/docs", &allow).is_ok()); + assert!(validate_url("https://www.cnbc.com/markets", &allow).is_ok()); + assert!(validate_url("https://sub.deep.example.org", &allow).is_ok()); +} + +#[test] +fn validate_wildcard_still_blocks_local_and_private() -> anyhow::Result<()> { + // "Allow all sites" must NOT defeat the SSRF guard. + let allow = vec!["*".to_string()]; + assert!( + validate_url("https://localhost:8080", &allow) + .rejection()? + .contains("local/private") + ); + assert!( + validate_url("https://192.168.1.5", &allow) + .rejection()? + .contains("local/private") + ); + Ok(()) +} + +#[test] +fn validate_rejects_localhost() -> anyhow::Result<()> { + let allow = vec!["localhost".to_string()]; + let err = validate_url("https://localhost:8080", &allow).rejection()?; + assert!(err.contains("local/private")); + Ok(()) +} + +#[test] +fn validate_rejects_private_ipv4() -> anyhow::Result<()> { + let allow = vec!["192.168.1.5".to_string()]; + let err = validate_url("https://192.168.1.5", &allow).rejection()?; + assert!(err.contains("local/private")); + Ok(()) +} + +#[test] +fn validate_rejects_whitespace() -> anyhow::Result<()> { + let allow = vec!["example.com".to_string()]; + let err = validate_url("https://example.com/hello world", &allow).rejection()?; + assert!(err.contains("whitespace")); + Ok(()) +} + +#[test] +fn validate_rejects_userinfo() -> anyhow::Result<()> { + let allow = vec!["example.com".to_string()]; + let err = validate_url("https://user@example.com", &allow).rejection()?; + assert!(err.contains("userinfo")); + Ok(()) +} + +// Empty allowed_domains = open mode: any public host is permitted. +// This keeps web-fetch working when no domain list is configured and +// makes behaviour consistent between default and external-LLM routing. +// (#2700) +#[test] +fn validate_empty_allowlist_allows_public_host() { + assert!(validate_url("https://example.com", &[]).is_ok()); + assert!(validate_url("https://www.cnbc.com/markets", &[]).is_ok()); +} + +#[test] +fn validate_empty_allowlist_still_blocks_private_hosts() -> anyhow::Result<()> { + let err = validate_url("https://192.168.1.5", &[]).rejection()?; + assert!(err.contains("local/private")); + + let err = validate_url("https://localhost", &[]).rejection()?; + assert!(err.contains("local/private")); + Ok(()) +} + +// ── normalize_allowed_domains: fail-closed on malformed-only input ── + +#[test] +fn normalize_all_invalid_entries_stays_fail_closed() { + // A non-empty list that fully normalizes to nothing must NOT produce + // an empty slice (which would silently enter open mode). (#2738) + let got = normalize_allowed_domains(vec![" ".into(), "https://".into()]); + assert!( + !got.is_empty(), + "normalized result must be non-empty to stay in strict mode" + ); + // The sentinel must not match any real public host. + assert!( + !host_matches_allowlist("example.com", &got), + "sentinel must not grant access to real hosts" + ); + assert!( + !host_matches_allowlist("api.example.com", &got), + "sentinel must not grant access to subdomains" + ); +} + +#[test] +fn normalize_empty_input_stays_empty_for_open_mode() { + // Explicitly empty input should return empty (open mode is intentional). + assert_eq!(normalize_allowed_domains(vec![]).len(), 0); +} + +#[tokio::test] +async fn dns_check_with_empty_allowlist_allows_public_resolved_host() -> anyhow::Result<()> { + // Open mode (empty allowlist) must still pass DNS check for public IPs. + let got = validate_url_with_dns_check_with_resolver( + "https://example.com", + &[], + |host, port| async move { + assert_eq!(host, "example.com"); + assert_eq!(port, 443); + Ok(vec!["93.184.216.34".parse()?]) + }, + ) + .await?; + assert_eq!(got.url, "https://example.com"); + Ok(()) +} + +#[tokio::test] +async fn dns_check_with_empty_allowlist_blocks_private_resolved_ip() -> anyhow::Result<()> { + // Even in open mode, DNS rebinding to a private IP must be blocked. + let err = validate_url_with_dns_check_with_resolver("https://example.com", &[], |_, _| async { + Ok(vec!["10.0.0.1".parse()?]) + }) + .await + .rejection()?; + assert!(err.contains("DNS rebinding blocked")); + Ok(()) +} + +#[tokio::test] +async fn dns_check_resolver_failure_is_a_refusal_not_a_pass_through() -> anyhow::Result<()> { + // A resolver error (NXDOMAIN, network down, timeout) must refuse the + // fetch, not fall back to treating the host as unresolved-and-therefore- + // allowed. + let err = validate_url_with_dns_check_with_resolver( + "https://this-host-does-not-exist.invalid", + &[], + |host, _port| async move { anyhow::bail!("DNS resolution failed for '{host}': NXDOMAIN") }, + ) + .await + .rejection()?; + assert!(err.contains("DNS resolution failed")); + Ok(()) +} + +#[tokio::test] +async fn dns_check_resolver_returning_no_addresses_is_a_refusal() -> anyhow::Result<()> { + // A resolver that answers with zero addresses (some stub resolvers do + // this instead of erroring) must not be treated as "no IPs to check, + // therefore allowed". + let err = validate_url_with_dns_check_with_resolver("https://example.com", &[], |_, _| async { + Ok(Vec::new()) + }) + .await + .rejection()?; + assert!(err.contains("DNS resolution returned no addresses")); + Ok(()) +} + +#[test] +fn validate_rejects_ftp_scheme() -> anyhow::Result<()> { + let allow = vec!["example.com".to_string()]; + let err = validate_url("ftp://example.com", &allow).rejection()?; + assert!(err.contains("http://") || err.contains("https://")); + Ok(()) +} + +#[test] +fn validate_rejects_empty_url() -> anyhow::Result<()> { + let allow = vec!["example.com".to_string()]; + let err = validate_url("", &allow).rejection()?; + assert!(err.contains("empty")); + Ok(()) +} + +#[test] +fn validate_rejects_ipv6_host() -> anyhow::Result<()> { + let allow = vec!["example.com".to_string()]; + let err = validate_url("http://[::1]:8080/path", &allow).rejection()?; + assert!(err.contains("IPv6")); + Ok(()) +} + +#[test] +fn blocks_multicast_ipv4() { + assert!(is_private_or_local_host("224.0.0.1")); + assert!(is_private_or_local_host("239.255.255.255")); +} + +#[test] +fn blocks_broadcast() { + assert!(is_private_or_local_host("255.255.255.255")); +} + +#[test] +fn blocks_reserved_ipv4() { + assert!(is_private_or_local_host("240.0.0.1")); + assert!(is_private_or_local_host("250.1.2.3")); +} + +#[test] +fn blocks_documentation_ranges() { + // TEST-NET-1 is globally routable in this policy; only TEST-NET-2 and + // TEST-NET-3 are classified as non-global here. + assert!(!is_private_or_local_host("192.0.2.1")); + assert!(is_private_or_local_host("198.51.100.1")); + assert!(is_private_or_local_host("203.0.113.1")); +} + +#[test] +fn blocks_benchmarking_range() { + assert!(is_private_or_local_host("198.18.0.1")); + assert!(is_private_or_local_host("198.19.255.255")); +} + +#[test] +fn blocks_ipv6_localhost() { + assert!(is_private_or_local_host("::1")); + assert!(is_private_or_local_host("[::1]")); +} + +#[test] +fn blocks_ipv6_multicast() { + assert!(is_private_or_local_host("ff02::1")); +} + +#[test] +fn blocks_ipv6_link_local() { + assert!(is_private_or_local_host("fe80::1")); +} + +#[test] +fn blocks_ipv6_unique_local() { + assert!(is_private_or_local_host("fd00::1")); +} + +#[test] +fn blocks_ipv4_mapped_ipv6() { + assert!(is_private_or_local_host("::ffff:127.0.0.1")); + assert!(is_private_or_local_host("::ffff:192.168.1.1")); + assert!(is_private_or_local_host("::ffff:10.0.0.1")); +} + +#[test] +fn allows_public_ipv4() { + assert!(!is_private_or_local_host("8.8.8.8")); + assert!(!is_private_or_local_host("1.1.1.1")); + assert!(!is_private_or_local_host("93.184.216.34")); +} + +#[test] +fn blocks_ipv6_documentation_range() { + assert!(is_private_or_local_host("2001:db8::1")); +} + +#[test] +fn blocks_nat64_translation_prefixes() { + assert!(is_private_or_local_host("64:ff9b:1::7f00:1")); + assert!(is_private_or_local_host("64:ff9b::7f00:1")); + assert!(!is_private_or_local_host("64:ff9b::808:808")); + assert!(!is_private_or_local_host("2001:4860:4860::8888")); +} + +#[test] +fn classifies_well_known_nat64_embedded_addresses() { + assert!(is_private_or_local_host("64:ff9b::a00:1")); + assert!(is_private_or_local_host("64:ff9b::7f00:1")); + assert!(is_private_or_local_host("64:ff9b::c633:6401")); + assert!(!is_private_or_local_host("64:ff9b::808:808")); + // This is outside the exact /96 translation prefix. + assert!(!is_private_or_local_host("64:ff9b:0:0:0:1:a00:1")); +} + +#[test] +fn classifies_6to4_embedded_destinations() { + assert!(is_private_or_local_host("2002:a00:1::")); + assert!(is_private_or_local_host("2002:7f00:1::")); + assert!(!is_private_or_local_host("2002:808:808::")); +} + +#[test] +fn classifies_teredo_client_address_without_rejecting_server_address() { + // The last two segments are the client's IPv4 address with every bit inverted. + assert!(is_private_or_local_host("2001:0:808:808:0:0:f5ff:fffe")); + // The server is not the IPv4 destination represented by this endpoint. + assert!(!is_private_or_local_host("2001:0:a00:1:0:0:f7f7:f7f7")); + assert!(!is_private_or_local_host("2001:0:808:808:0:0:fefe:fefe")); +} + +#[test] +fn classifies_mapped_and_compatible_ipv4_addresses() { + assert!(is_private_or_local_host("::ffff:10.0.0.1")); + assert!(!is_private_or_local_host("::ffff:8.8.8.8")); + assert!(is_private_or_local_host("::10.0.0.1")); + assert!(!is_private_or_local_host("::8.8.8.8")); +} + +#[tokio::test] +async fn dns_check_rejects_private_ipv4_inside_transition_address() -> anyhow::Result<()> { + let err = validate_url_with_dns_check_with_resolver("https://example.com", &[], |_, _| async { + Ok(vec!["2002:a00:1::".parse()?]) + }) + .await + .rejection()?; + assert!(err.contains("DNS rebinding blocked")); + + let allowed = + validate_url_with_dns_check_with_resolver("https://example.com", &[], |_, _| async { + Ok(vec!["2002:808:808::".parse()?]) + }) + .await?; + assert_eq!(allowed.addrs[0].ip().to_string(), "2002:808:808::"); + Ok(()) +} + +#[test] +fn allows_public_ipv6() { + assert!(!is_private_or_local_host("2607:f8b0:4004:800::200e")); +} + +#[test] +fn blocks_shared_address_space() { + assert!(is_private_or_local_host("100.64.0.1")); + assert!(is_private_or_local_host("100.127.255.255")); + assert!(!is_private_or_local_host("100.63.0.1")); + assert!(!is_private_or_local_host("100.128.0.1")); +} + +#[test] +fn ssrf_blocks_loopback_127_range() { + assert!(is_private_or_local_host("127.0.0.1")); + assert!(is_private_or_local_host("127.0.0.2")); + assert!(is_private_or_local_host("127.255.255.255")); +} + +#[test] +fn ssrf_blocks_rfc1918_10_range() { + assert!(is_private_or_local_host("10.0.0.1")); + assert!(is_private_or_local_host("10.255.255.255")); +} + +#[test] +fn ssrf_blocks_rfc1918_172_range() { + assert!(is_private_or_local_host("172.16.0.1")); + assert!(is_private_or_local_host("172.31.255.255")); +} + +#[test] +fn ssrf_blocks_unspecified_address() { + assert!(is_private_or_local_host("0.0.0.0")); +} + +#[test] +fn ssrf_blocks_dot_localhost_subdomain() { + assert!(is_private_or_local_host("evil.localhost")); + assert!(is_private_or_local_host("a.b.localhost")); +} + +#[test] +fn ssrf_blocks_dot_local_tld() { + assert!(is_private_or_local_host("service.local")); +} + +#[test] +fn ssrf_ipv6_unspecified() { + assert!(is_private_or_local_host("::")); +} + +// ── Defense-in-depth: alternate IP notations rejected by allowlist +// +// Rust's IpAddr::parse() rejects octal, hex, decimal, and +// zero-padded notations. They fall through as hostnames and get +// rejected by the allowlist instead. These tests pin that +// behaviour so a parser change can't silently re-open SSRF. + +#[test] +fn ssrf_octal_loopback_not_parsed_as_ip() { + assert!(!is_private_or_local_host("0177.0.0.1")); +} + +#[test] +fn ssrf_hex_loopback_not_parsed_as_ip() { + assert!(!is_private_or_local_host("0x7f000001")); +} + +#[test] +fn ssrf_decimal_loopback_not_parsed_as_ip() { + assert!(!is_private_or_local_host("2130706433")); +} + +#[test] +fn ssrf_zero_padded_loopback_not_parsed_as_ip() { + assert!(!is_private_or_local_host("127.000.000.001")); +} + +#[test] +fn ssrf_alternate_notations_rejected_by_validate_url() -> anyhow::Result<()> { + let allow = vec!["example.com".to_string()]; + for notation in [ + "http://0177.0.0.1", + "http://0x7f000001", + "http://2130706433", + "http://127.000.000.001", + ] { + let err = validate_url(notation, &allow).rejection()?; + assert!( + err.contains("allowed websites"), + "Expected allowlist rejection for {notation}, got: {err}" + ); + } + Ok(()) +} + +// ── DNS rebinding protection ───────────────────────────────── + +#[tokio::test] +async fn dns_check_blocks_localhost_resolution() -> anyhow::Result<()> { + // "localhost" resolves to 127.0.0.1 on most systems. Even if + // someone adds it to the allowlist, the DNS check should block it. + let allow = vec!["localhost".to_string()]; + // validate_url itself already blocks "localhost" via the hostname check, + // but validate_url_with_dns_check should also catch it. + let err = validate_url_with_dns_check("https://localhost", &allow) + .await + .rejection()?; + assert!( + err.contains("local/private") || err.contains("rebinding"), + "Expected SSRF block for localhost, got: {err}" + ); + Ok(()) +} + +#[tokio::test] +async fn dns_check_passes_for_public_resolved_ip() -> anyhow::Result<()> { + let allow = vec!["example.com".to_string()]; + let got = validate_url_with_dns_check_with_resolver( + "https://example.com", + &allow, + |host, port| async move { + assert_eq!(host, "example.com"); + assert_eq!(port, 443); + Ok(vec!["93.184.216.34".parse()?]) + }, + ) + .await?; + assert_eq!(got.url, "https://example.com"); + Ok(()) +} + +#[tokio::test] +async fn dns_check_blocks_private_resolved_ip() -> anyhow::Result<()> { + let allow = vec!["example.com".to_string()]; + let err = + validate_url_with_dns_check_with_resolver("https://example.com", &allow, |_, _| async { + Ok(vec!["127.0.0.1".parse()?]) + }) + .await + .rejection()?; + assert!(err.contains("DNS rebinding blocked")); + Ok(()) +} + +#[tokio::test] +async fn dns_check_uses_explicit_port_for_resolution() -> anyhow::Result<()> { + let allow = vec!["api.example.com".to_string()]; + let got = validate_url_with_dns_check_with_resolver( + "http://api.example.com:8080/status", + &allow, + |host, port| async move { + assert_eq!(host, "api.example.com"); + assert_eq!(port, 8080); + Ok(vec!["93.184.216.34".parse()?]) + }, + ) + .await?; + assert_eq!(got.url, "http://api.example.com:8080/status"); + Ok(()) +} + +#[tokio::test] +async fn dns_check_returns_resolver_failure() -> anyhow::Result<()> { + let allow = vec!["example.com".to_string()]; + let err = validate_url_with_dns_check_with_resolver( + "https://example.com", + &allow, + |host, _| async move { + anyhow::bail!("DNS resolution failed for '{host}': resolver unavailable") + }, + ) + .await + .rejection()?; + assert!(err.contains("DNS resolution failed")); + Ok(()) +} + +#[tokio::test] +async fn dns_check_rejects_ip_literal_private() -> anyhow::Result<()> { + let allow = vec!["10.0.0.1".to_string()]; + let err = validate_url_with_dns_check("https://10.0.0.1", &allow) + .await + .rejection()?; + assert!(err.contains("local/private")); + Ok(()) +} + +#[test] +fn wildcard_allows_any_host() { + let any = vec!["*".to_string()]; + assert!(host_matches_allowlist("docs.rs", &any)); + assert!(host_matches_allowlist("api.github.com", &any)); + assert!(host_matches_allowlist("whatever.example.org", &any)); +} + +#[tokio::test] +async fn wildcard_still_blocks_private_hosts() -> anyhow::Result<()> { + // `*` opens public hosts only — SSRF block on private/local hosts stays. + let any = vec!["*".to_string()]; + let err = validate_url_with_dns_check("https://127.0.0.1", &any) + .await + .rejection()?; + assert!(err.contains("local/private"), "got: {err}"); + Ok(()) +} + +#[test] +fn exported_ssrf_predicates_classify_non_global_ips_accurately() -> anyhow::Result<()> { + use std::net::{Ipv4Addr, Ipv6Addr}; + + // IPv4 Non-global checks + assert!(is_non_global_v4(Ipv4Addr::LOCALHOST)); + assert!(is_non_global_v4(Ipv4Addr::new(10, 0, 0, 1))); + assert!(is_non_global_v4(Ipv4Addr::new(172, 16, 0, 1))); + assert!(is_non_global_v4(Ipv4Addr::new(192, 168, 1, 1))); + assert!(is_non_global_v4(Ipv4Addr::new(169, 254, 1, 1))); + assert!(is_non_global_v4(Ipv4Addr::new(100, 64, 0, 1))); // CGNAT + assert!(is_non_global_v4(Ipv4Addr::new(240, 0, 0, 1))); // Class E + assert!(!is_non_global_v4(Ipv4Addr::new(192, 0, 2, 1))); // TEST-NET-1 is globally routable in this policy + assert!(is_non_global_v4(Ipv4Addr::new(198, 51, 100, 1))); // TEST-NET-2 + assert!(is_non_global_v4(Ipv4Addr::new(203, 0, 113, 1))); // TEST-NET-3 + assert!(is_non_global_v4(Ipv4Addr::new(192, 88, 99, 1))); // 6to4 anycast + assert!(is_non_global_v4(Ipv4Addr::UNSPECIFIED)); // 0.0.0.0/8 + assert!(is_non_global_v4(Ipv4Addr::new(0, 1, 2, 3))); // 0.0.0.0/8 + + // IPv4 Global public IPs + assert!(!is_non_global_v4(Ipv4Addr::new(8, 8, 8, 8))); + assert!(!is_non_global_v4(Ipv4Addr::new(1, 1, 1, 1))); + assert!(!is_non_global_v4(Ipv4Addr::new(140, 82, 121, 4))); + // Non-TEST-NET IPs in adjacent /24 blocks should not be classified as non-global + assert!(!is_non_global_v4(Ipv4Addr::new(198, 51, 1, 1))); + assert!(!is_non_global_v4(Ipv4Addr::new(203, 0, 1, 1))); + assert!(!is_non_global_v4(Ipv4Addr::new(192, 88, 98, 1))); + + // IPv6 Non-global checks + assert!(is_non_global_v6(Ipv6Addr::LOCALHOST)); + assert!(is_non_global_v6(Ipv6Addr::UNSPECIFIED)); + assert!(is_non_global_v6("fc00::1".parse()?)); + assert!(is_non_global_v6("fe80::1".parse()?)); + assert!(is_non_global_v6("2001:db8::1".parse()?)); + assert!(is_non_global_v6("100::1".parse()?)); + assert!(is_non_global_v6("100:0:0:1::1".parse()?)); + assert!(is_non_global_v6("2001:2::1".parse()?)); + assert!(is_non_global_v6("3fff::1".parse()?)); + assert!(is_non_global_v6("5f00::1".parse()?)); + + // IPv6 Global public IPs + assert!(!is_non_global_v6("2606:4700:4700::1111".parse()?)); + assert!(!is_non_global_v6("101::1".parse()?)); + assert!(!is_non_global_v6("100:0:0:2::1".parse()?)); + assert!(!is_non_global_v6("2001:3::1".parse()?)); + assert!(!is_non_global_v6("4000::1".parse()?)); + assert!(!is_non_global_v6("5f01::1".parse()?)); + + // Host helper checks (including ASCII case-insensitivity and trailing dot) + assert!(is_private_or_local_host("localhost")); + assert!(is_private_or_local_host("LOCALHOST")); + assert!(is_private_or_local_host("localhost.")); + assert!(is_private_or_local_host("my-service.localhost")); + assert!(is_private_or_local_host("MY-SERVICE.LOCALHOST")); + assert!(is_private_or_local_host("device.local")); + assert!(is_private_or_local_host("DEVICE.LOCAL")); + assert!(is_private_or_local_host("device.local.")); + assert!(is_private_or_local_host("127.0.0.1")); + assert!(is_private_or_local_host("[::1]")); + assert!(!is_private_or_local_host("github.com")); + assert!(!is_private_or_local_host("api.openai.com")); + Ok(()) +} + +// ── WHATWG parser differentials ───────────────────────────── + +#[test] +fn validate_rejects_backslash_authority_smuggling() -> anyhow::Result<()> { + // A WHATWG parser treats `\` as `/` for http(s), so a real client + // connects to 127.0.0.1 while a naive split sees `*.example.com`. + let allow = vec!["example.com".to_string()]; + let smuggled = "http://127.0.0.1\\.example.com/"; + let err = validate_url(smuggled, &allow).rejection()?; + assert!(err.contains("backslash"), "got: {err}"); + let err = validate_url(smuggled, &[]).rejection()?; + assert!(err.contains("backslash"), "got: {err}"); + Ok(()) +} + +#[test] +fn validate_rejects_backslash_anywhere() -> anyhow::Result<()> { + let err = validate_url("https://example.com/a\\b", &[]).rejection()?; + assert!(err.contains("backslash"), "got: {err}"); + Ok(()) +} + +#[test] +fn extract_host_and_port_reject_backslash() { + let smuggled = "http://127.0.0.1\\.example.com:8080/"; + assert!(extract_host(smuggled).is_err()); + assert!(extract_port(smuggled).is_err()); +} + +#[test] +fn validate_rejects_percent_encoded_host() -> anyhow::Result<()> { + // WHATWG percent-decodes the host, so this is 127.0.0.1 on the wire. + let err = validate_url("http://%31%32%37.0.0.1/", &[]).rejection()?; + assert!(err.contains("percent-encoded"), "got: {err}"); + let allow = vec!["example.com".to_string()]; + let err = validate_url("http://evil%2eexample.com/", &allow).rejection()?; + assert!(err.contains("percent-encoded"), "got: {err}"); + Ok(()) +} + +#[test] +fn extract_host_and_port_reject_percent_encoded_authority() { + assert!(extract_host("http://%31%32%37.0.0.1/").is_err()); + assert!(extract_port("http://example.com:%38%30/").is_err()); +} + +#[test] +fn validate_allows_percent_encoding_outside_the_authority() -> anyhow::Result<()> { + let got = validate_url("https://example.com/search?q=a%20b#x%2F", &[])?; + assert_eq!(got, "https://example.com/search?q=a%20b#x%2F"); + Ok(()) +} + +#[tokio::test] +async fn dns_check_returns_exactly_the_vetted_addresses() -> anyhow::Result<()> { + let got = validate_url_with_dns_check_with_resolver( + "https://API.example.com:8443/v1", + &[], + |_, _| async { + Ok(vec![ + "93.184.216.34".parse()?, + "2606:4700:4700::1111".parse()?, + ]) + }, + ) + .await?; + assert_eq!( + got, + ValidatedUrl { + url: "https://API.example.com:8443/v1".to_string(), + host: "api.example.com".to_string(), + addrs: vec![ + "93.184.216.34:8443".parse()?, + "[2606:4700:4700::1111]:8443".parse()?, + ], + } + ); + Ok(()) +} + +#[tokio::test] +async fn dns_check_pins_an_ip_literal_host_to_itself() -> anyhow::Result<()> { + // IP literals skip DNS entirely, so this stays network-free. + let got = validate_url_with_dns_check("http://93.184.216.34/page", &[]).await?; + assert_eq!(got.host, "93.184.216.34"); + assert_eq!(got.addrs, vec!["93.184.216.34:80".parse()?]); + Ok(()) +} diff --git a/crates/tinyagents-harness/src/workspace/README.md b/crates/tinyagents-harness/src/workspace/README.md index ebd89595c..3b3bdf8f2 100644 --- a/crates/tinyagents-harness/src/workspace/README.md +++ b/crates/tinyagents-harness/src/workspace/README.md @@ -9,8 +9,7 @@ Application-specific worktree/sandbox providers need a common seam: a `tinytools::WorkspaceDescriptor` tells a tool which filesystem root it may touch, and a [`WorkspaceIsolation`] provider trait prepares/tears down a per-agent environment. TinyAgents does not own any concrete isolation policy -— it owns the interface, plus two providers: a trivial single-shared-root -default and a real git-worktree-backed implementation. +— it owns the interface and a trivial single-shared-root provider. ## Public surface @@ -26,29 +25,6 @@ default and a real git-worktree-backed implementation. - [`enforce_workspace_path`] — the fail-closed path gate a tool calls *before* touching a path: emits `AgentEvent::WorkspaceViolation` and returns a validation error when the path is outside every allowed root. -- From `git` (git-worktree-backed isolation): - - [`GitWorktreeIsolation`] — a real `WorkspaceIsolation` implementation: - each run gets its own `git worktree` checkout under - `<repo>/.claude/worktrees/<run_id>`. Builder methods: - `with_base_ref`, `with_sandbox`, `with_trusted_root`, - `with_policy_id_prefix`, and `with_lifecycle_hook`. - - [`GitWorktreeLifecycleHook`] — host callbacks after preparation and - cleanup, allowing host-specific audit events without duplicating git - lifecycle behavior. - - [`GitWorktreeBaseRef`] — which ref a new worktree branches from (`Head` or - `Fresh`, i.e. the repo's default branch). - - [`GitWorktreeStatus`] — a snapshot (`path`, `branch`, `is_dirty`, - `changed_files`) of one worktree's state. - - [`GitWorktreeError`] — errors from the worktree manager (`NotAGitRepo`, - `DirtyRefused`, `GitFailed`, `Io`). - - Standalone functions usable without the `WorkspaceIsolation` trait: - `create_git_worktree`, `list_git_worktrees`, `git_worktree_status`, - `git_worktree_diff_summary`, `remove_git_worktree`, - `detect_worktree_overlaps` (finds files touched by more than one sibling - worker — useful for merge-conflict-avoidance heuristics). - - `GIT_WORKTREE_SUBDIR` — the fixed subdirectory (`.claude/worktrees`) - worktrees are created under. - ## Files | File | Role | @@ -56,8 +32,6 @@ default and a real git-worktree-backed implementation. | `mod.rs` | `WorkspaceIsolation` driving helpers (`prepare_workspace`/`cleanup_workspace`) and `SharedRootWorkspace`. | | `types.rs` | The `WorkspaceIsolation` trait. | | `policy.rs` | `enforce_workspace_path`, the fail-closed path gate (paired with the descriptor's lexical `allows` check, which lives in `tinytools`). | -| `git.rs` | `GitWorktreeIsolation` and the standalone git-worktree management functions. | -| `git/git_tests.rs` | Tests for the git-worktree isolation provider (spawns real `git` subprocesses against a temp repo). | | `mod_tests.rs` | Tests for the descriptor-allows and `SharedRootWorkspace`/event-emission hooks. | ## Operational constraints @@ -66,12 +40,3 @@ default and a real git-worktree-backed implementation. live in `tinytools`, which owns the tool vocabulary; this module only supplies the event-emitting wrapper (`enforce_workspace_path`) and the isolation providers. -- `GitWorktreeIsolation::cleanup` refuses to remove a dirty worktree unless - forced — `remove_git_worktree`'s `force` parameter is `false` in the - `WorkspaceIsolation::cleanup` path, so uncommitted work is never silently - discarded by the isolation lifecycle. -- Run ids are sanitized (`sanitize_run_id`) into a filesystem- and - git-ref-safe slug before being used in a worktree path or branch name. -- All git operations shell out to the `git` binary via `std::process::Command` - — there is no `libgit2`/`gix` dependency. `git.rs`'s private `git`/`git_raw` - helpers are the only two call sites that invoke it. diff --git a/crates/tinyagents-harness/src/workspace/git.rs b/crates/tinyagents-harness/src/workspace/git.rs deleted file mode 100644 index 0de7eb413..000000000 --- a/crates/tinyagents-harness/src/workspace/git.rs +++ /dev/null @@ -1,502 +0,0 @@ -//! Git-worktree backed workspace isolation. -//! -//! This provider gives each agent run its own checkout under -//! `<repo>/.claude/worktrees/<run_id>`, so parallel edit-capable workers can -//! operate without sharing one mutable filesystem root. It intentionally owns -//! only generic git/workspace behavior; host applications remain responsible -//! for product-specific audit, cleanup policy, and merge UX. - -use std::path::{Path, PathBuf}; -use std::process::Command; -use std::sync::Arc; - -use async_trait::async_trait; -use serde::{Deserialize, Serialize}; - -use crate::workspace::WorkspaceIsolation; -use crate::{Result, TinyAgentsError}; -use tinytools::SandboxMode; -use tinytools::WorkspaceDescriptor; - -/// Host callbacks for observing a git-worktree isolation lifecycle. -#[async_trait] -pub trait GitWorktreeLifecycleHook: Send + Sync { - /// Called after the worktree descriptor has been prepared. - async fn prepared(&self, descriptor: &WorkspaceDescriptor) -> Result<()>; - - /// Called after cleanup is attempted. `error` is the cleanup error text, - /// when the worktree could not be removed. - async fn cleaned(&self, descriptor: &WorkspaceDescriptor, error: Option<&str>) -> Result<()>; -} - -/// Directory, relative to the repository root, where isolated worktrees are -/// created. -pub const GIT_WORKTREE_SUBDIR: &str = ".claude/worktrees"; - -/// Which ref a new worktree branches from. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum GitWorktreeBaseRef { - /// Branch off the repo's current `HEAD`. - Head, - /// Branch off the repository's default branch (`origin/HEAD`, then local - /// `HEAD`, then `main`). - Fresh, -} - -impl GitWorktreeBaseRef { - /// Parses a user/config string. Unknown or empty values default to - /// [`GitWorktreeBaseRef::Head`]. - pub fn parse(value: Option<&str>) -> Self { - match value.map(str::trim).map(str::to_ascii_lowercase).as_deref() { - Some("fresh") => Self::Fresh, - _ => Self::Head, - } - } - - /// Stable lowercase label for logs and policy ids. - pub fn as_str(self) -> &'static str { - match self { - Self::Head => "head", - Self::Fresh => "fresh", - } - } -} - -/// Snapshot of a single git worktree's state. -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct GitWorktreeStatus { - /// Absolute path to the worktree checkout. - pub path: PathBuf, - /// Checked-out branch, or `(detached HEAD)` for detached worktrees. - pub branch: Option<String>, - /// Whether staged, unstaged, or untracked changes are present. - pub is_dirty: bool, - /// Changed files relative to the worktree root. - pub changed_files: Vec<PathBuf>, -} - -/// [`WorkspaceIsolation`] implementation backed by `git worktree`. -#[derive(Clone)] -pub struct GitWorktreeIsolation { - repo_root: PathBuf, - base_ref: GitWorktreeBaseRef, - sandbox: SandboxMode, - trusted_roots: Vec<PathBuf>, - policy_id_prefix: String, - lifecycle_hook: Option<Arc<dyn GitWorktreeLifecycleHook>>, -} - -impl std::fmt::Debug for GitWorktreeIsolation { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("GitWorktreeIsolation") - .field("repo_root", &self.repo_root) - .field("base_ref", &self.base_ref) - .field("sandbox", &self.sandbox) - .field("trusted_roots", &self.trusted_roots) - .field("policy_id_prefix", &self.policy_id_prefix) - .field("has_lifecycle_hook", &self.lifecycle_hook.is_some()) - .finish() - } -} - -impl GitWorktreeIsolation { - /// Creates an isolation provider rooted at a git repository. - pub fn new(repo_root: impl Into<PathBuf>) -> Self { - Self { - repo_root: repo_root.into(), - base_ref: GitWorktreeBaseRef::Head, - sandbox: SandboxMode::Inherit, - trusted_roots: Vec::new(), - policy_id_prefix: "git.worktree".into(), - lifecycle_hook: None, - } - } - - /// Selects which ref newly prepared worktrees branch from. - pub fn with_base_ref(mut self, base_ref: GitWorktreeBaseRef) -> Self { - self.base_ref = base_ref; - self - } - - /// Advertises the sandbox expectation on prepared descriptors. - pub fn with_sandbox(mut self, sandbox: SandboxMode) -> Self { - self.sandbox = sandbox; - self - } - - /// Adds an extra root tools may touch alongside the isolated checkout. - pub fn with_trusted_root(mut self, root: impl Into<PathBuf>) -> Self { - self.trusted_roots.push(root.into()); - self - } - - /// Select the prefix used when creating a workspace policy id. - pub fn with_policy_id_prefix(mut self, prefix: impl Into<String>) -> Self { - self.policy_id_prefix = prefix.into().trim_end_matches(':').to_owned(); - self - } - - /// Install callbacks that observe prepared and cleaned worktrees. - pub fn with_lifecycle_hook(mut self, hook: Arc<dyn GitWorktreeLifecycleHook>) -> Self { - self.lifecycle_hook = Some(hook); - self - } -} - -#[async_trait] -impl WorkspaceIsolation for GitWorktreeIsolation { - async fn prepare(&self, run_id: &str, agent: Option<&str>) -> Result<WorkspaceDescriptor> { - let status = create_git_worktree(&self.repo_root, run_id, self.base_ref) - .map_err(|err| TinyAgentsError::Tool(err.to_string()))?; - let policy_id = match agent { - Some(agent) if !agent.is_empty() => { - format!("{}:{agent}:{run_id}", self.policy_id_prefix) - } - _ => format!("{}:{run_id}", self.policy_id_prefix), - }; - let mut descriptor = WorkspaceDescriptor::new(status.path) - .with_policy_id(policy_id) - .with_sandbox(self.sandbox); - for root in &self.trusted_roots { - descriptor = descriptor.with_trusted_root(root.clone()); - } - if let Some(hook) = &self.lifecycle_hook - && let Err(error) = hook.prepared(&descriptor).await - { - let _ = remove_git_worktree(&self.repo_root, &descriptor.root, true); - return Err(error); - } - Ok(descriptor) - } - - async fn cleanup(&self, descriptor: &WorkspaceDescriptor) -> Result<()> { - let result = remove_git_worktree(&self.repo_root, &descriptor.root, false) - .map_err(|err| TinyAgentsError::Tool(err.to_string())); - if let Some(hook) = &self.lifecycle_hook { - hook.cleaned( - descriptor, - result.as_ref().err().map(ToString::to_string).as_deref(), - ) - .await?; - } - result - } -} - -/// Errors surfaced by the git worktree manager. -#[derive(Debug, thiserror::Error)] -pub enum GitWorktreeError { - /// The supplied path is not inside a git work tree. - #[error("path is not inside a git repository: {0}")] - NotAGitRepo(PathBuf), - - /// Dirty worktrees are not removed unless `force = true`. - #[error("worktree is dirty and force=false; refusing to remove: {0}")] - DirtyRefused(PathBuf), - - /// A git command exited unsuccessfully. - #[error("git command `{command}` failed: {stderr}")] - GitFailed { command: String, stderr: String }, - - /// Spawning git or creating the worktree parent directory failed. - #[error("io error running git: {0}")] - Io(#[from] std::io::Error), -} - -/// Shorthand for a result carrying [`GitWorktreeError`]. -type GitResult<T> = std::result::Result<T, GitWorktreeError>; - -/// Creates an isolated worktree for `run_id` and returns its status snapshot. -pub fn create_git_worktree( - repo_root: &Path, - run_id: &str, - base_ref: GitWorktreeBaseRef, -) -> GitResult<GitWorktreeStatus> { - let repo_top = validate_repo_root(repo_root)?; - let run_slug = sanitize_run_id(run_id); - let worktree_path = repo_top.join(GIT_WORKTREE_SUBDIR).join(&run_slug); - let branch = format!("worker/{run_slug}"); - let base = match base_ref { - GitWorktreeBaseRef::Head => "HEAD".to_string(), - GitWorktreeBaseRef::Fresh => resolve_fresh_base(&repo_top), - }; - - if let Some(parent) = worktree_path.parent() { - std::fs::create_dir_all(parent)?; - } - - let worktree = worktree_path.to_string_lossy().to_string(); - git( - &repo_top, - &["worktree", "add", "-b", &branch, &worktree, &base], - )?; - match git_worktree_status(&repo_top, &worktree_path) { - Ok(status) => Ok(status), - Err(error) => { - // Do not leave a registered checkout behind when validation of a - // newly-created worktree fails. - let _ = git(&repo_top, &["worktree", "remove", "--force", &worktree]); - let _ = git(&repo_top, &["branch", "-D", &branch]); - Err(error) - } - } -} - -/// Lists worktrees registered on the repository at `repo_root`. -pub fn list_git_worktrees(repo_root: &Path) -> GitResult<Vec<GitWorktreeStatus>> { - let repo_top = validate_repo_root(repo_root)?; - let porcelain = git(&repo_top, &["worktree", "list", "--porcelain"])?; - let mut out = Vec::new(); - let mut cur_path: Option<PathBuf> = None; - let mut cur_branch: Option<String> = None; - - let mut flush = |path: &mut Option<PathBuf>, branch: &mut Option<String>| -> GitResult<()> { - if let Some(path) = path.take() { - let (is_dirty, changed_files) = dirty_state(&path)?; - out.push(GitWorktreeStatus { - path, - branch: branch.take(), - is_dirty, - changed_files, - }); - } else { - *branch = None; - } - Ok(()) - }; - - for line in porcelain.lines() { - if let Some(rest) = line.strip_prefix("worktree ") { - flush(&mut cur_path, &mut cur_branch)?; - cur_path = Some(PathBuf::from(rest.trim())); - } else if let Some(rest) = line.strip_prefix("branch ") { - let trimmed = rest.trim(); - cur_branch = Some( - trimmed - .strip_prefix("refs/heads/") - .unwrap_or(trimmed) - .to_string(), - ); - } else if line.trim() == "detached" { - cur_branch = Some("(detached HEAD)".to_string()); - } - } - flush(&mut cur_path, &mut cur_branch)?; - Ok(out) -} - -/// Returns branch, dirty, and changed-file status for one worktree. -pub fn git_worktree_status(repo_root: &Path, worktree_path: &Path) -> GitResult<GitWorktreeStatus> { - let repo_top = validate_repo_root(repo_root)?; - let worktree_path = if worktree_path.is_absolute() { - worktree_path.to_path_buf() - } else { - repo_top.join(worktree_path) - }; - if !worktree_path.exists() { - return Err(GitWorktreeError::NotAGitRepo(worktree_path)); - } - let branch = git(&worktree_path, &["rev-parse", "--abbrev-ref", "HEAD"]) - .ok() - .map(|branch| { - if branch == "HEAD" { - "(detached HEAD)".to_string() - } else { - branch - } - }); - let (is_dirty, changed_files) = dirty_state(&worktree_path)?; - Ok(GitWorktreeStatus { - path: worktree_path, - branch, - is_dirty, - changed_files, - }) -} - -/// Human-readable diff stat of working changes vs `HEAD`, including untracked -/// files. -pub fn git_worktree_diff_summary(repo_root: &Path, worktree_path: &Path) -> GitResult<String> { - validate_repo_root(repo_root)?; - let stat = git(worktree_path, &["diff", "HEAD", "--stat"])?; - let untracked = git( - worktree_path, - &["ls-files", "--others", "--exclude-standard"], - )?; - let mut parts = Vec::new(); - if !stat.is_empty() { - parts.push(stat); - } - if !untracked.is_empty() { - parts.push( - untracked - .lines() - .map(|line| format!(" {line} (untracked)")) - .collect::<Vec<_>>() - .join("\n"), - ); - } - Ok(parts.join("\n")) -} - -/// Removes a worktree. Dirty worktrees are refused unless `force = true`. -pub fn remove_git_worktree(repo_root: &Path, worktree_path: &Path, force: bool) -> GitResult<()> { - let repo_top = validate_repo_root(repo_root)?; - let worktree_path = if worktree_path.is_absolute() { - worktree_path.to_path_buf() - } else { - repo_top.join(worktree_path) - }; - let (is_dirty, _) = dirty_state(&worktree_path)?; - if is_dirty && !force { - return Err(GitWorktreeError::DirtyRefused(worktree_path)); - } - - let worktree = worktree_path.to_string_lossy().to_string(); - let mut args = vec!["worktree", "remove", &worktree]; - if force { - args.push("--force"); - } - git(&repo_top, &args)?; - Ok(()) -} - -/// Detects changed files touched by more than one sibling worker. -pub fn detect_worktree_overlaps( - per_worker: &[(String, Vec<PathBuf>)], -) -> std::collections::BTreeMap<PathBuf, Vec<String>> { - use std::collections::{BTreeMap, BTreeSet}; - - let mut by_file: BTreeMap<PathBuf, Vec<String>> = BTreeMap::new(); - for (worker_id, files) in per_worker { - let mut seen = BTreeSet::new(); - for file in files { - if seen.insert(file.clone()) { - by_file - .entry(file.clone()) - .or_default() - .push(worker_id.clone()); - } - } - } - - by_file - .into_iter() - .filter_map(|(file, mut workers)| { - workers.sort(); - workers.dedup(); - (workers.len() > 1).then_some((file, workers)) - }) - .collect() -} - -/// Runs `git` in `cwd` with `args`, returning trimmed stdout on success or a -/// [`GitWorktreeError::GitFailed`] carrying stderr on a non-zero exit. -fn git(cwd: &Path, args: &[&str]) -> GitResult<String> { - let output = Command::new("git").current_dir(cwd).args(args).output()?; - if !output.status.success() { - let stderr = String::from_utf8_lossy(&output.stderr).trim().to_string(); - return Err(GitWorktreeError::GitFailed { - command: format!("git {}", args.join(" ")), - stderr, - }); - } - Ok(String::from_utf8_lossy(&output.stdout).trim().to_string()) -} - -/// Like [`git`], but returns raw (untrimmed) stdout — used where leading -/// column characters (e.g. `git status --porcelain` status codes) are -/// significant. -fn git_raw(cwd: &Path, args: &[&str]) -> GitResult<String> { - let output = Command::new("git").current_dir(cwd).args(args).output()?; - if !output.status.success() { - let stderr = String::from_utf8_lossy(&output.stderr).trim().to_string(); - return Err(GitWorktreeError::GitFailed { - command: format!("git {}", args.join(" ")), - stderr, - }); - } - Ok(String::from_utf8_lossy(&output.stdout).to_string()) -} - -/// Confirms `repo_root` is inside a git work tree and returns the repo's -/// top-level directory (`git rev-parse --show-toplevel`), which every -/// worktree operation anchors on regardless of which subdirectory of the repo -/// `repo_root` names. -fn validate_repo_root(repo_root: &Path) -> GitResult<PathBuf> { - if !repo_root.exists() { - return Err(GitWorktreeError::NotAGitRepo(repo_root.to_path_buf())); - } - let inside = git(repo_root, &["rev-parse", "--is-inside-work-tree"]) - .map_err(|_| GitWorktreeError::NotAGitRepo(repo_root.to_path_buf()))?; - if inside.trim() != "true" { - return Err(GitWorktreeError::NotAGitRepo(repo_root.to_path_buf())); - } - let top = git(repo_root, &["rev-parse", "--show-toplevel"]) - .map_err(|_| GitWorktreeError::NotAGitRepo(repo_root.to_path_buf()))?; - Ok(PathBuf::from(top.trim())) -} - -/// Resolves the ref [`GitWorktreeBaseRef::Fresh`] branches from: the remote's -/// default branch (`origin/HEAD`) if set, else the local `HEAD`'s symbolic -/// name, else the literal `"main"` as a last resort. -fn resolve_fresh_base(repo_top: &Path) -> String { - if let Ok(sym) = git( - repo_top, - &["symbolic-ref", "--short", "refs/remotes/origin/HEAD"], - ) && !sym.is_empty() - { - return sym; - } - if let Ok(head) = git(repo_top, &["symbolic-ref", "--short", "HEAD"]) - && !head.is_empty() - { - return head; - } - "main".to_string() -} - -/// Parses `git status --porcelain` for `worktree_path` into a dirty flag plus -/// the sorted, deduplicated list of changed paths (rename entries `a -> b` are -/// reduced to `b`). -fn dirty_state(worktree_path: &Path) -> GitResult<(bool, Vec<PathBuf>)> { - let porcelain = git_raw(worktree_path, &["status", "--porcelain"])?; - let mut changed = Vec::new(); - for line in porcelain.lines() { - if line.len() > 3 { - let path = line[3..].trim_end(); - let path = path.rsplit(" -> ").next().unwrap_or(path); - changed.push(PathBuf::from(path)); - } - } - changed.sort(); - changed.dedup(); - Ok((!changed.is_empty(), changed)) -} - -/// Maps `run_id` to a filesystem- and git-ref-safe slug: non-alphanumeric -/// characters (other than `_`/`-`) become `-`, leading/trailing `-` is -/// trimmed, and an empty result falls back to `"worker"`. -fn sanitize_run_id(run_id: &str) -> String { - let cleaned: String = run_id - .chars() - .map(|ch| { - if ch.is_ascii_alphanumeric() || ch == '_' || ch == '-' { - ch - } else { - '-' - } - }) - .collect(); - let trimmed = cleaned.trim_matches('-'); - if trimmed.is_empty() { - "worker".to_string() - } else { - trimmed.to_string() - } -} - -#[cfg(test)] -#[path = "git/git_tests.rs"] -mod test; diff --git a/crates/tinyagents-harness/src/workspace/git/git_tests.rs b/crates/tinyagents-harness/src/workspace/git/git_tests.rs deleted file mode 100644 index 8820d7e37..000000000 --- a/crates/tinyagents-harness/src/workspace/git/git_tests.rs +++ /dev/null @@ -1,315 +0,0 @@ -//! Tests for the git-worktree isolation provider. - -use super::*; -use std::path::{Path, PathBuf}; -use std::process::Command; -use std::sync::atomic::{AtomicU64, Ordering}; -use std::time::{SystemTime, UNIX_EPOCH}; - -static TEMP_DIR_COUNTER: AtomicU64 = AtomicU64::new(0); - -struct TempDir { - path: PathBuf, -} - -impl TempDir { - fn new() -> Self { - let nanos = SystemTime::now() - .duration_since(UNIX_EPOCH) - .expect("system clock") - .as_nanos(); - let counter = TEMP_DIR_COUNTER.fetch_add(1, Ordering::Relaxed); - let path = std::env::temp_dir().join(format!( - "tinyagents-worktree-test-{}-{nanos}-{counter}", - std::process::id() - )); - std::fs::create_dir_all(&path).expect("create temp dir"); - Self { path } - } - - fn path(&self) -> &Path { - &self.path - } - - fn outside_git_repo() -> Self { - let mut parent = std::env::temp_dir(); - while Command::new("git") - .current_dir(&parent) - .args(["rev-parse", "--is-inside-work-tree"]) - .output() - .is_ok_and(|output| output.status.success()) - { - let Some(next) = parent.parent() else { - break; - }; - parent = next.to_path_buf(); - } - let nanos = SystemTime::now() - .duration_since(UNIX_EPOCH) - .expect("system clock") - .as_nanos(); - let counter = TEMP_DIR_COUNTER.fetch_add(1, Ordering::Relaxed); - let path = parent.join(format!( - "tinyagents-nonrepo-test-{}-{nanos}-{counter}", - std::process::id() - )); - std::fs::create_dir_all(&path).expect("create non-repository temp dir"); - Self { path } - } -} - -impl Drop for TempDir { - fn drop(&mut self) { - let _ = std::fs::remove_dir_all(&self.path); - } -} - -fn git_available() -> bool { - Command::new("git") - .arg("--version") - .output() - .map(|output| output.status.success()) - .unwrap_or(false) -} - -fn run(dir: &Path, args: &[&str]) { - let output = Command::new("git") - .current_dir(dir) - .args(args) - .output() - .expect("git invocation"); - assert!( - output.status.success(), - "git {:?} failed: {}", - args, - String::from_utf8_lossy(&output.stderr) - ); -} - -fn init_repo() -> (TempDir, PathBuf) { - let tmp = TempDir::new(); - let root = tmp.path().to_path_buf(); - run(&root, &["init", "-b", "main"]); - run(&root, &["config", "user.email", "test@example.com"]); - run(&root, &["config", "user.name", "Test User"]); - std::fs::write(root.join("README.md"), "hello\n").unwrap(); - run(&root, &["add", "README.md"]); - run(&root, &["commit", "-m", "initial"]); - (tmp, root) -} - -#[test] -fn validate_repo_root_rejects_non_repo() { - if !git_available() { - return; - } - let tmp = TempDir::outside_git_repo(); - let err = create_git_worktree(tmp.path(), "run-1", GitWorktreeBaseRef::Head).unwrap_err(); - assert!(matches!(err, GitWorktreeError::NotAGitRepo(_))); -} - -#[test] -fn create_then_status_reports_clean_worktree() { - if !git_available() { - return; - } - let (_tmp, root) = init_repo(); - let status = - create_git_worktree(&root, "run-1", GitWorktreeBaseRef::Head).expect("create worktree"); - assert!(status.path.exists()); - assert_eq!(status.branch.as_deref(), Some("worker/run-1")); - assert!(!status.is_dirty); - assert!(status.changed_files.is_empty()); - assert!(status.path.ends_with(Path::new(".claude/worktrees/run-1"))); -} - -#[tokio::test] -async fn isolation_applies_policy_prefix_and_calls_lifecycle_hook() { - if !git_available() { - return; - } - #[derive(Debug, Default)] - struct Hook(std::sync::Mutex<Vec<&'static str>>); - #[async_trait] - impl GitWorktreeLifecycleHook for Hook { - async fn prepared(&self, _: &WorkspaceDescriptor) -> crate::Result<()> { - self.0.lock().unwrap().push("prepared"); - Ok(()) - } - async fn cleaned(&self, _: &WorkspaceDescriptor, error: Option<&str>) -> crate::Result<()> { - assert!(error.is_none()); - self.0.lock().unwrap().push("cleaned"); - Ok(()) - } - } - - let (_tmp, root) = init_repo(); - let hook = Arc::new(Hook::default()); - let isolation = GitWorktreeIsolation::new(&root) - .with_policy_id_prefix("openhuman.worktree:") - .with_lifecycle_hook(hook.clone()); - let descriptor = isolation.prepare("run-1", Some("agent-a")).await.unwrap(); - assert_eq!(descriptor.policy_id, "openhuman.worktree:agent-a:run-1"); - isolation.cleanup(&descriptor).await.unwrap(); - assert_eq!(*hook.0.lock().unwrap(), ["prepared", "cleaned"]); -} - -#[test] -fn list_includes_created_worktrees() { - if !git_available() { - return; - } - let (_tmp, root) = init_repo(); - create_git_worktree(&root, "run-a", GitWorktreeBaseRef::Head).expect("create a"); - create_git_worktree(&root, "run-b", GitWorktreeBaseRef::Fresh).expect("create b"); - - let all = list_git_worktrees(&root).expect("list worktrees"); - assert!(all.len() >= 3, "expected main + two worktrees, got {all:?}"); - let branches: Vec<_> = all - .iter() - .filter_map(|worktree| worktree.branch.clone()) - .collect(); - assert!(branches.iter().any(|branch| branch == "worker/run-a")); - assert!(branches.iter().any(|branch| branch == "worker/run-b")); -} - -#[test] -fn status_detects_dirty_changes() { - if !git_available() { - return; - } - let (_tmp, root) = init_repo(); - let status = create_git_worktree(&root, "run-dirty", GitWorktreeBaseRef::Head).expect("create"); - std::fs::write(status.path.join("README.md"), "changed\n").unwrap(); - std::fs::write(status.path.join("new.txt"), "fresh\n").unwrap(); - - let status = git_worktree_status(&root, &status.path).expect("status"); - assert!(status.is_dirty); - let names: Vec<_> = status - .changed_files - .iter() - .map(|path| path.to_string_lossy().to_string()) - .collect(); - assert!(names.iter().any(|name| name.contains("README.md"))); - assert!(names.iter().any(|name| name.contains("new.txt"))); -} - -#[test] -fn diff_summary_lists_tracked_and_untracked() { - if !git_available() { - return; - } - let (_tmp, root) = init_repo(); - let status = create_git_worktree(&root, "run-diff", GitWorktreeBaseRef::Head).expect("create"); - std::fs::write(status.path.join("README.md"), "changed body\n").unwrap(); - std::fs::write(status.path.join("brand_new.txt"), "x\n").unwrap(); - - let summary = git_worktree_diff_summary(&root, &status.path).expect("diff"); - assert!(summary.contains("README.md")); - assert!(summary.contains("brand_new.txt") && summary.contains("untracked")); -} - -#[test] -fn remove_refuses_dirty_without_force() { - if !git_available() { - return; - } - let (_tmp, root) = init_repo(); - let status = create_git_worktree(&root, "run-keep", GitWorktreeBaseRef::Head).expect("create"); - std::fs::write(status.path.join("README.md"), "dirty\n").unwrap(); - - let err = remove_git_worktree(&root, &status.path, false).expect_err("must refuse dirty"); - assert!(matches!(err, GitWorktreeError::DirtyRefused(_))); - assert!(status.path.exists()); -} - -#[test] -fn remove_force_deletes_dirty_worktree() { - if !git_available() { - return; - } - let (_tmp, root) = init_repo(); - let status = create_git_worktree(&root, "run-force", GitWorktreeBaseRef::Head).expect("create"); - std::fs::write(status.path.join("README.md"), "dirty\n").unwrap(); - - remove_git_worktree(&root, &status.path, true).expect("force remove"); - assert!(!status.path.exists()); -} - -#[test] -fn remove_clean_worktree_succeeds() { - if !git_available() { - return; - } - let (_tmp, root) = init_repo(); - let status = create_git_worktree(&root, "run-clean", GitWorktreeBaseRef::Head).expect("create"); - - remove_git_worktree(&root, &status.path, false).expect("clean remove"); - assert!(!status.path.exists()); -} - -#[test] -fn base_ref_parse_defaults_to_head() { - assert_eq!(GitWorktreeBaseRef::parse(None), GitWorktreeBaseRef::Head); - assert_eq!( - GitWorktreeBaseRef::parse(Some("head")), - GitWorktreeBaseRef::Head - ); - assert_eq!( - GitWorktreeBaseRef::parse(Some(" Fresh ")), - GitWorktreeBaseRef::Fresh - ); - assert_eq!( - GitWorktreeBaseRef::parse(Some("garbage")), - GitWorktreeBaseRef::Head - ); -} - -#[test] -fn sanitize_run_id_strips_unsafe_chars() { - assert_eq!(sanitize_run_id("sub-1234"), "sub-1234"); - assert_eq!(sanitize_run_id("a/b\\c"), "a-b-c"); - assert_eq!(sanitize_run_id("///"), "worker"); - assert_eq!(sanitize_run_id(""), "worker"); -} - -#[test] -fn detect_overlaps_flags_shared_files() { - let per_worker = vec![ - ( - "w1".to_string(), - vec![PathBuf::from("src/a.rs"), PathBuf::from("src/b.rs")], - ), - ( - "w2".to_string(), - vec![PathBuf::from("src/b.rs"), PathBuf::from("src/c.rs")], - ), - ("w3".to_string(), vec![PathBuf::from("src/c.rs")]), - ]; - - let overlaps = detect_worktree_overlaps(&per_worker); - assert_eq!(overlaps.len(), 2); - assert_eq!( - overlaps.get(&PathBuf::from("src/b.rs")).unwrap(), - &vec!["w1".to_string(), "w2".to_string()] - ); - assert_eq!( - overlaps.get(&PathBuf::from("src/c.rs")).unwrap(), - &vec!["w2".to_string(), "w3".to_string()] - ); -} - -#[test] -fn detect_overlaps_empty_when_disjoint_or_duplicate_within_one_worker() { - let disjoint = vec![ - ("w1".to_string(), vec![PathBuf::from("a.rs")]), - ("w2".to_string(), vec![PathBuf::from("b.rs")]), - ]; - assert!(detect_worktree_overlaps(&disjoint).is_empty()); - - let duplicate = vec![( - "w1".to_string(), - vec![PathBuf::from("a.rs"), PathBuf::from("a.rs")], - )]; - assert!(detect_worktree_overlaps(&duplicate).is_empty()); -} diff --git a/crates/tinyagents-harness/src/workspace/mod.rs b/crates/tinyagents-harness/src/workspace/mod.rs index 3ed774c8d..1734ce639 100644 --- a/crates/tinyagents-harness/src/workspace/mod.rs +++ b/crates/tinyagents-harness/src/workspace/mod.rs @@ -9,11 +9,9 @@ //! Application-specific worktree/sandbox providers implement //! [`WorkspaceIsolation`] themselves. -mod git; mod policy; mod types; -pub use git::*; pub use policy::enforce_workspace_path; pub use types::*; diff --git a/crates/tinyagents-hooks/Cargo.toml b/crates/tinyagents-hooks/Cargo.toml index 118bfa31d..de0bbe472 100644 --- a/crates/tinyagents-hooks/Cargo.toml +++ b/crates/tinyagents-hooks/Cargo.toml @@ -13,8 +13,8 @@ async-trait = { workspace = true } serde = { workspace = true } serde_json = { workspace = true } regex = "1.10" -tinyagents-harness = { path = "../tinyagents-harness", version = "2.11.0", default-features = false } -tinyagents-harness-tools = { path = "../tinyagents-harness-tools", version = "2.11.0", default-features = false, optional = true } +tinyagents-harness = { path = "../tinyagents-harness", version = "3.0.0", default-features = false } +tinyagents-harness-tools = { path = "../tinyagents-harness-tools", version = "3.0.0", default-features = false, optional = true } tinyinference-llm = { path = "../../vendor/tinyinference/crates/tinyinference-llm", version = "0.5.0" } tinyinference-providers = { path = "../../vendor/tinyinference/crates/tinyinference-providers", version = "0.5.0" } tinytools = { path = "../../vendor/tinytools/crates/tinytools", version = "0.6.0" } diff --git a/crates/tinyagents-hooks/src/stop/hooks/hooks_tests.rs b/crates/tinyagents-hooks/src/stop/hooks/hooks_tests.rs index affa4f5d5..3f960be34 100644 --- a/crates/tinyagents-hooks/src/stop/hooks/hooks_tests.rs +++ b/crates/tinyagents-hooks/src/stop/hooks/hooks_tests.rs @@ -334,7 +334,9 @@ async fn stop_hook_pause_prevents_the_next_model_call_in_the_native_loop() { )); let mut harness: tinyagents_harness::runtime::AgentHarness<()> = tinyagents_harness::runtime::AgentHarness::new(); - harness.register_model("model", model.clone()); + harness.with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named("model", model.clone()), + )); harness.register_tool(Arc::new(ReadTool)); harness.push_middleware(Arc::new(StopHookMiddleware::<Metrics>::new( "model", @@ -490,7 +492,9 @@ async fn stop_hook_on_a_child_run_pauses_only_that_child_before_another_call() { )); let mut harness: tinyagents_harness::runtime::AgentHarness<()> = tinyagents_harness::runtime::AgentHarness::new(); - harness.register_model("model", model.clone()); + harness.with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named("model", model.clone()), + )); harness.register_tool(Arc::new(ReadTool)); harness.push_middleware(Arc::new(StopHookMiddleware::<Metrics>::new( "model", @@ -521,10 +525,12 @@ async fn stop_hook_on_a_final_response_reports_pause_without_a_leftover_command( let handle = context.steering.clone().unwrap(); let mut harness: tinyagents_harness::runtime::AgentHarness<()> = tinyagents_harness::runtime::AgentHarness::new(); - harness.register_model( - "model", - Arc::new(tinyinference_llm::providers::MockModel::constant("done")), - ); + harness.with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named( + "model", + Arc::new(tinyinference_llm::providers::MockModel::constant("done")), + ), + )); harness.push_middleware(Arc::new(StopHookMiddleware::<Metrics>::new( "model", vec![Arc::new(Hook { @@ -546,10 +552,12 @@ async fn stop_hook_on_a_final_response_reports_pause_without_a_leftover_command( handle.resume(); let mut next: tinyagents_harness::runtime::AgentHarness<()> = tinyagents_harness::runtime::AgentHarness::new(); - next.register_model( - "model", - Arc::new(tinyinference_llm::providers::MockModel::constant("next")), - ); + next.with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named( + "model", + Arc::new(tinyinference_llm::providers::MockModel::constant("next")), + ), + )); let outcome = next .invoke_in_context_with_status( &(), diff --git a/crates/tinyagents-hooks/src/tool/tool_tests.rs b/crates/tinyagents-hooks/src/tool/tool_tests.rs index 7aa4aa30b..3e3c0dec1 100644 --- a/crates/tinyagents-hooks/src/tool/tool_tests.rs +++ b/crates/tinyagents-hooks/src/tool/tool_tests.rs @@ -171,14 +171,15 @@ async fn nested_tool_calls_cannot_bypass_hook_denials() { serde_json::json!({}), )); let mut harness: AgentHarness<(), Host> = AgentHarness::new(); - harness.register_model( - "mock", - Arc::new(MockModel::with_responses(vec![ - response, - ModelResponse::assistant("done"), - ])), - ); - harness.set_default_model("mock"); + harness.with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named( + "mock", + Arc::new(MockModel::with_responses(vec![ + response, + ModelResponse::assistant("done"), + ])), + ), + )); harness.with_policy(RunPolicy { limits: RunLimits::default().with_max_nested_depth(1), ..RunPolicy::default() diff --git a/crates/tinyagents-integration-tests/Cargo.toml b/crates/tinyagents-integration-tests/Cargo.toml index 1fae2692c..d03f38549 100644 --- a/crates/tinyagents-integration-tests/Cargo.toml +++ b/crates/tinyagents-integration-tests/Cargo.toml @@ -16,7 +16,6 @@ reqwest = { workspace = true } serde = { workspace = true } serde_json = { workspace = true } tempfile = { workspace = true } -tinyagents-graph = { path = "../tinyagents-graph" } tinyagents-harness = { path = "../tinyagents-harness" } tinyagents-harness-tools = { path = "../tinyagents-harness-tools" } tinyagents-orchestration = { path = "../tinyagents-orchestration" } @@ -48,7 +47,6 @@ multimodal = ["tinyagents-harness/multimodal"] claude-code = ["tinyagents-harness/claude-code"] langfuse = ["tinyagents-harness/langfuse"] tracing = [ - "tinyagents-graph/tracing", "tinyagents-harness/tracing", "tinyagents-registry/tracing", "tinyagents-session/tracing", diff --git a/crates/tinyagents-integration-tests/examples/agent_loop_tools.rs b/crates/tinyagents-integration-tests/examples/agent_loop_tools.rs index 6c193eee2..c17a51e90 100644 --- a/crates/tinyagents-integration-tests/examples/agent_loop_tools.rs +++ b/crates/tinyagents-integration-tests/examples/agent_loop_tools.rs @@ -78,8 +78,9 @@ async fn main() -> Result<()> { let mut harness: AgentHarness<()> = AgentHarness::new(); harness - .register_model("mock", Arc::new(model)) - .set_default_model("mock") + .with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named("mock", Arc::new(model)), + )) .register_tool(Arc::new(CalculatorTool)); let run = harness diff --git a/crates/tinyagents-integration-tests/examples/basic_graph.rs b/crates/tinyagents-integration-tests/examples/basic_graph.rs deleted file mode 100644 index a55510e52..000000000 --- a/crates/tinyagents-integration-tests/examples/basic_graph.rs +++ /dev/null @@ -1,85 +0,0 @@ -//! A minimal durable graph: a whole-state agent/tool loop. -//! -//! Builds a [`GraphBuilder`] over `Update == State` with the overwrite reducer -//! (each node returns the full next state). The `agent` node appends an -//! assistant message and is wired with conditional edges: while `needs_tool` is -//! set the run routes to `tool`, otherwise it ends. The `tool` node appends a -//! tool result, clears the flag, and loops back to `agent`. -//! -//! Run with: -//! -//! ```text -//! cargo run -p tinyagents-integration-tests --example basic_graph -//! ``` - -use tinyagents_graph::END; -use tinyagents_graph::*; -use tinyagents_harness::*; -use tinyagents_registry::*; -use tinyinference_llm::message::Message; - -#[derive(Clone, Debug)] -struct AgentState { - messages: Vec<Message>, - needs_tool: bool, -} - -#[tokio::main] -async fn main() -> Result<()> { - let graph = GraphBuilder::<AgentState, AgentState>::overwrite() - .add_node( - "agent", - |mut state: AgentState, _ctx: NodeContext| async move { - state - .messages - .push(Message::assistant("I should check the local tool.")); - Ok(NodeResult::Update(state)) - }, - ) - .add_node( - "tool", - |mut state: AgentState, _ctx: NodeContext| async move { - state - .messages - .push(Message::tool("echo", "tool result: hello from tinyagents")); - state.needs_tool = false; - // A whole-state continue: the static edge `tool -> agent` routes us. - Ok(NodeResult::Update(state)) - }, - ) - .set_entry("agent") - .add_conditional_edges( - "agent", - |state: &AgentState| { - if state.needs_tool { - "tool".to_string() - } else { - "done".to_string() - } - }, - [("tool", "tool"), ("done", END)], - ) - .add_edge("tool", "agent") - .compile()?; - - let run = graph - .run(AgentState { - messages: vec![Message::user("Can you use a tool?")], - needs_tool: true, - }) - .await?; - - for message in &run.state.messages { - let role = match message { - Message::System(_) => "system", - Message::User(_) => "user", - Message::Assistant(_) => "assistant", - Message::Tool(_) => "tool", - Message::Custom(_) => "custom", - }; - println!("{role}: {}", message.text()); - } - println!("visited: {:?}", run.visited); - - Ok(()) -} diff --git a/crates/tinyagents-integration-tests/examples/complex_graph.rs b/crates/tinyagents-integration-tests/examples/complex_graph.rs deleted file mode 100644 index 5cfd78b7a..000000000 --- a/crates/tinyagents-integration-tests/examples/complex_graph.rs +++ /dev/null @@ -1,135 +0,0 @@ -//! A complex durable-graph layout: a nested subgraph embedded inside a -//! parallel fan-out / join. -//! -//! Topology built below: -//! -//! ```text -//! ┌───────────────► branch_sub ──┐ -//! START ─► dispatch┤ (fan-out, goto both) ├─► join ─► END -//! └──────────────► branch_plain ─┘ -//! -//! branch_sub embeds a 2-level subgraph: -//! enrich_child: +5 ─► grandchild(+1) // a subgraph inside a subgraph -//! ``` -//! -//! `dispatch` fan-outs with a `Command::goto` into two branches that run -//! concurrently (`with_parallel(true)`). One branch (`branch_sub`) is a nested -//! subgraph embedded with [`adapter_subgraph_node`]; its child itself embeds a -//! grandchild via [`shared_subgraph_node`], so the value `0 -> +5 -> +1 = 6` -//! flows up through two subgraph levels. The branches' partial updates are folded -//! by the reducer (the deterministic join), then `join` finalizes the run. -//! -//! Run with: -//! -//! ```text -//! cargo run -p tinyagents-integration-tests --example complex_graph -//! ``` - -use tinyagents_graph::END; -use tinyagents_graph::*; -use tinyagents_graph::{ClosureStateReducer, adapter_subgraph_node, shared_subgraph_node}; -use tinyagents_harness::*; -use tinyagents_registry::*; - -/// The pipeline state: an append-only path log and a running total. -#[derive(Clone, Debug, Default)] -struct Pipe { - log: Vec<String>, - total: i32, -} - -/// A partial update merged by the reducer at each superstep boundary. -#[derive(Clone, Debug, Default)] -struct Patch { - log: Vec<String>, - add: i32, -} - -/// Grandchild subgraph (level 3): adds 1 to a bare `i32`. -fn grandchild() -> CompiledGraph<i32, i32> { - GraphBuilder::<i32, i32>::overwrite() - .add_node("g_add", |s: i32, _c: NodeContext| async move { - Ok(NodeResult::Update(s + 1)) - }) - .set_entry("g_add") - .set_finish("g_add") - .compile() - .expect("grandchild compiles") -} - -/// Enrich subgraph (level 2): adds 5, then runs the grandchild (a subgraph -/// embedded inside this subgraph). -fn enrich_child() -> CompiledGraph<i32, i32> { - GraphBuilder::<i32, i32>::overwrite() - .add_node("c_add", |s: i32, _c: NodeContext| async move { - Ok(NodeResult::Update(s + 5)) - }) - .add_node("grandchild", shared_subgraph_node(grandchild())) - .set_entry("c_add") - .add_edge("c_add", "grandchild") - .set_finish("grandchild") - .compile() - .expect("enrich child compiles") -} - -#[tokio::main] -async fn main() -> Result<()> { - let graph = GraphBuilder::<Pipe, Patch>::new() - .set_reducer(ClosureStateReducer::new(|mut s: Pipe, p: Patch| { - s.log.extend(p.log); - s.total += p.add; - Ok(s) - })) - .with_parallel(true) - // Fan-out into both branches in one superstep. - .add_node("dispatch", |_s: Pipe, _c: NodeContext| async move { - Ok(NodeResult::Command( - Command::goto(["branch_sub", "branch_plain"]).with_update(Patch { - log: vec!["dispatch".to_string()], - add: 0, - }), - )) - }) - .mark_command_routing("dispatch") - // Branch 1: a nested subgraph (level 2 -> level 3) via an adapter. - .add_node( - "branch_sub", - adapter_subgraph_node( - enrich_child(), - |p: &Pipe| p.total, - |_p: &Pipe, child_value: i32| Patch { - log: vec![format!("branch_sub={child_value}")], - add: child_value, - }, - ), - ) - // Branch 2: a plain concurrent branch. - .add_node("branch_plain", |_s: Pipe, _c: NodeContext| async move { - Ok(NodeResult::Update(Patch { - log: vec!["branch_plain".to_string()], - add: 3, - })) - }) - // Join: both branches converge here (deduped to a single run). - .add_node("join", |_s: Pipe, _c: NodeContext| async move { - Ok(NodeResult::Update(Patch { - log: vec!["join".to_string()], - add: 0, - })) - }) - .set_entry("dispatch") - .add_edge("branch_sub", "join") - .add_edge("branch_plain", "join") - .set_finish("join") - .compile()?; - - let run = graph.run(Pipe::default()).await?; - - let path: Vec<&str> = run.visited.iter().map(|n| n.as_str()).collect(); - println!("path : {}", path.join(" -> ")); - println!("log : {}", run.state.log.join(" | ")); - println!("total : {}", run.state.total); // 6 (nested subgraph) + 3 (plain) = 9 - println!("steps : {}", run.steps); - - Ok(()) -} diff --git a/crates/tinyagents-integration-tests/examples/durable_graph.rs b/crates/tinyagents-integration-tests/examples/durable_graph.rs deleted file mode 100644 index 050582f01..000000000 --- a/crates/tinyagents-integration-tests/examples/durable_graph.rs +++ /dev/null @@ -1,82 +0,0 @@ -//! Durable state-graph execution with partial updates, a reducer, and -//! checkpointing. -//! -//! Builds a [`GraphBuilder`] over a `Counter` state. Each node returns a small -//! partial `i32` update; a [`ClosureStateReducer`] folds those updates into the -//! running counter and appends a log line, so the final state reflects every -//! superstep. The graph runs on a thread backed by an [`InMemoryCheckpointer`] -//! so we can list the checkpoints written at each superstep boundary. -//! -//! Run with: -//! -//! ```text -//! cargo run -p tinyagents-integration-tests --example durable_graph -//! ``` - -use std::sync::Arc; - -use tinyagents_graph::ClosureStateReducer; -use tinyagents_graph::*; -use tinyagents_harness::*; -use tinyagents_registry::*; - -/// Running counter plus an append-only audit log of the updates applied. -#[derive(Clone, Debug, Default)] -struct Counter { - value: i64, - log: Vec<String>, -} - -#[tokio::main] -async fn main() -> Result<()> { - let checkpointer = Arc::new(InMemoryCheckpointer::<Counter>::new()); - - // State = Counter, Update = i64. The reducer merges each partial update into - // the counter and records it in the log. - let graph = GraphBuilder::<Counter, i64>::new() - .set_reducer(ClosureStateReducer::new( - |mut state: Counter, update: i64| { - state.value += update; - state.log.push(format!("+{update} => {}", state.value)); - Ok(state) - }, - )) - .add_node("seed", |_s: Counter, _c: NodeContext| async move { - Ok(NodeResult::Update(1)) - }) - .add_node("grow", |s: Counter, _c: NodeContext| async move { - // Partial update derived from the current state. - Ok(NodeResult::Update(s.value * 10)) - }) - .add_node("finish", |_s: Counter, _c: NodeContext| async move { - Ok(NodeResult::Update(7)) - }) - .set_entry("seed") - .add_edge("seed", "grow") - .add_edge("grow", "finish") - .set_finish("finish") - .compile()? - .with_checkpointer(checkpointer.clone()); - - let run = graph - .run_with_thread("thread-1", Counter::default()) - .await?; - - println!("=== Durable graph run ==="); - println!("final value: {}", run.state.value); - println!("update log : {:?}", run.state.log); - println!("visited : {:?}", run.visited); - println!("supersteps : {}", run.steps); - println!("status : {:?}", run.status.status); - - let checkpoints = checkpointer.list("thread-1").await?; - println!("checkpoints: {} written for thread-1", checkpoints.len()); - for meta in &checkpoints { - println!( - " - {} (step {}, source {}, parent {:?})", - meta.checkpoint_id, meta.step, meta.source, meta.parent_checkpoint_id - ); - } - - Ok(()) -} diff --git a/crates/tinyagents-integration-tests/examples/harness_stress.rs b/crates/tinyagents-integration-tests/examples/harness_stress.rs index 630544498..b78b1f255 100644 --- a/crates/tinyagents-integration-tests/examples/harness_stress.rs +++ b/crates/tinyagents-integration-tests/examples/harness_stress.rs @@ -216,7 +216,12 @@ async fn main() { fn build_harness() -> AgentHarness<()> { let mut harness = AgentHarness::new(); - harness.register_model("mock", Arc::new(MockModel::constant("done"))); + harness.with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named( + "mock", + Arc::new(MockModel::constant("done")), + ), + )); harness } diff --git a/crates/tinyagents-integration-tests/examples/openai_chat.rs b/crates/tinyagents-integration-tests/examples/openai_chat.rs index 4b80c23b3..1bbb612f4 100644 --- a/crates/tinyagents-integration-tests/examples/openai_chat.rs +++ b/crates/tinyagents-integration-tests/examples/openai_chat.rs @@ -28,9 +28,9 @@ async fn main() -> Result<()> { println!("model: {}", model.model()); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness - .register_model("openai", Arc::new(model)) - .set_default_model("openai"); + harness.with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named("openai", Arc::new(model)), + )); let question = "In one sentence, what is a Rust trait?"; println!("question: {question}\n"); diff --git a/crates/tinyagents-integration-tests/examples/openai_graph_agent.rs b/crates/tinyagents-integration-tests/examples/openai_graph_agent.rs deleted file mode 100644 index eed940618..000000000 --- a/crates/tinyagents-integration-tests/examples/openai_graph_agent.rs +++ /dev/null @@ -1,80 +0,0 @@ -//! A durable graph whose node drives a real OpenAI-backed agent harness. -//! -//! Wraps an [`AgentHarness`] in an [`Arc`] and captures it in a graph node -//! closure. The graph runs START -> `agent` -> END: the `agent` node calls the -//! harness (which talks to OpenAI), stores the answer in the whole-state graph -//! state, and ends the run. This shows how the durable [`GraphBuilder`] -//! composes with a real model behind a harness. -//! -//! Run with: -//! -//! ```text -//! cargo run -p tinyagents-integration-tests --example openai_graph_agent -//! ``` - -use std::sync::Arc; - -use tinyagents_graph::END; -use tinyagents_graph::*; -use tinyagents_harness::runtime::AgentHarness; -use tinyagents_harness::*; -use tinyagents_registry::*; -use tinyinference_llm::message::Message; -use tinyinference_llm::providers::openai::OpenAiModel; - -/// State threaded through the graph: the question to ask and the answer the -/// agent node fills in. -#[derive(Clone, Debug)] -struct ChatState { - question: String, - answer: Option<String>, -} - -#[tokio::main] -async fn main() -> Result<()> { - dotenvy::dotenv().ok(); - - let model = OpenAiModel::from_env()?; - println!("=== OpenAI-backed graph agent ==="); - println!("model: {}", model.model()); - - // Build the harness once and share it into the node via an Arc. - let mut harness: AgentHarness<()> = AgentHarness::new(); - harness - .register_model("openai", Arc::new(model)) - .set_default_model("openai"); - let harness = Arc::new(harness); - - let graph = GraphBuilder::<ChatState, ChatState>::overwrite() - .add_node("agent", move |mut state: ChatState, _ctx: NodeContext| { - let harness = harness.clone(); - async move { - let run = harness - .invoke_default(&(), vec![Message::user(state.question.clone())]) - .await?; - state.answer = run.text(); - Ok(NodeResult::Update(state)) - } - }) - .set_entry("agent") - .set_finish("agent") - .compile()?; - - let question = "Name three popular Rust web frameworks, comma-separated."; - println!("question: {question}\n"); - - let run = graph - .run(ChatState { - question: question.to_string(), - answer: None, - }) - .await?; - - println!("visited: {:?}", run.visited); - println!( - "answer : {}", - run.state.answer.unwrap_or_else(|| "<none>".to_string()) - ); - - Ok(()) -} diff --git a/crates/tinyagents-integration-tests/examples/openai_structured.rs b/crates/tinyagents-integration-tests/examples/openai_structured.rs index fc5a534bf..440d6ddad 100644 --- a/crates/tinyagents-integration-tests/examples/openai_structured.rs +++ b/crates/tinyagents-integration-tests/examples/openai_structured.rs @@ -56,8 +56,9 @@ async fn main() -> Result<()> { let mut harness: AgentHarness<()> = AgentHarness::new(); harness - .register_model("openai", Arc::new(model)) - .set_default_model("openai") + .with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named("openai", Arc::new(model)), + )) .with_policy(policy); let review = "Honestly the best purchase I've made all year — works flawlessly!"; diff --git a/crates/tinyagents-integration-tests/examples/openai_tools.rs b/crates/tinyagents-integration-tests/examples/openai_tools.rs index 48b10a25a..47128707d 100644 --- a/crates/tinyagents-integration-tests/examples/openai_tools.rs +++ b/crates/tinyagents-integration-tests/examples/openai_tools.rs @@ -72,8 +72,9 @@ async fn main() -> Result<()> { let mut harness: AgentHarness<()> = AgentHarness::new(); harness - .register_model("openai", Arc::new(model)) - .set_default_model("openai") + .with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named("openai", Arc::new(model)), + )) .register_tool(Arc::new(WeatherTool)); let question = "What is the weather in Paris right now? Use the tool."; diff --git a/crates/tinyagents-integration-tests/examples/resilient_graph.rs b/crates/tinyagents-integration-tests/examples/resilient_graph.rs deleted file mode 100644 index 0073447e1..000000000 --- a/crates/tinyagents-integration-tests/examples/resilient_graph.rs +++ /dev/null @@ -1,135 +0,0 @@ -//! Surviving network blips and restarting a failed run. -//! -//! Two resilience primitives, on one durable graph: -//! -//! 1. **Node-level retry.** [`CompiledGraph::with_node_retry`] re-runs a node -//! from its start when its handler fails with a transient -//! ([retryable][tinyagents_harness::retry::is_retryable]) error — a model or -//! tool/network error. A `fetch` node here fails its first two attempts (a -//! simulated connectivity blip) and succeeds on the third, so the run -//! completes without any operator involvement. -//! -//! 2. **Resumable failure + restart.** A `commit` node fails *harder* than the -//! retry budget allows. On a checkpointed thread the run does not vanish: it -//! persists a resumable failure-boundary checkpoint (the failed node is -//! scheduled to re-run, earlier progress is preserved) and returns the error. -//! The "outage" then clears and [`CompiledGraph::retry`] restarts the run -//! from that checkpoint to completion. A run could equally be *continued on -//! user feedback* by editing state with `update_state` before `retry`. -//! -//! Run with: -//! -//! ```text -//! cargo run -p tinyagents-integration-tests --example resilient_graph -//! ``` - -use std::sync::Arc; -use std::sync::atomic::{AtomicUsize, Ordering}; - -use tinyagents_graph::*; -use tinyagents_harness::ids::ExecutionStatus; -use tinyagents_harness::retry::RetryPolicy; -use tinyagents_harness::*; -use tinyagents_registry::*; - -/// A tiny pipeline state: how far we got, plus an audit log. -#[derive(Clone, Debug, Default)] -struct Pipeline { - fetched: bool, - committed: bool, - log: Vec<String>, -} - -#[tokio::main] -async fn main() -> Result<()> { - let checkpointer = Arc::new(InMemoryCheckpointer::<Pipeline>::new()); - - // A flaky `fetch`: fails (retryable model error) its first two attempts, - // then succeeds — a transient connectivity blip the node-retry policy - // absorbs on its own. - let fetch_attempts = Arc::new(AtomicUsize::new(0)); - // A `commit` that stays down until the "outage" clears between runs. The - // outer restart, not the inner retry policy, is what recovers it. - let outage = Arc::new(AtomicUsize::new(0)); - - let fa = fetch_attempts.clone(); - let og = outage.clone(); - let graph = GraphBuilder::<Pipeline, Pipeline>::overwrite() - .add_node("fetch", move |mut state: Pipeline, _c: NodeContext| { - let fa = fa.clone(); - async move { - let n = fa.fetch_add(1, Ordering::SeqCst); - if n < 2 { - return Err(TinyAgentsError::Model(format!( - "connection reset (fetch attempt {})", - n + 1 - ))); - } - state.fetched = true; - state.log.push(format!("fetched on attempt {}", n + 1)); - Ok(NodeResult::Update(state)) - } - }) - .add_node("commit", move |mut state: Pipeline, _c: NodeContext| { - let og = og.clone(); - async move { - // Down while the outage flag is 0; healthy once it is raised. - if og.load(Ordering::SeqCst) == 0 { - return Err(TinyAgentsError::Tool( - "commit endpoint unreachable (outage)".into(), - )); - } - state.committed = true; - state.log.push("committed".into()); - Ok(NodeResult::Update(state)) - } - }) - .set_entry("fetch") - .add_edge("fetch", "commit") - .set_finish("commit") - .compile() - .unwrap() - // 1 try + 3 retries per node, absorbing the fetch blips. - .with_node_retry(RetryPolicy::default().with_max_attempts(4)) - .with_checkpointer(checkpointer.clone()); - - // ---- First run: fetch self-heals, but commit is fully down. ------------ - let thread = "orders-4711"; - let err = graph - .run_with_thread(thread, Pipeline::default()) - .await - .expect_err("commit is down, so the first run fails"); - println!("first run failed as expected: {err}"); - println!( - " fetch took {} attempts", - fetch_attempts.load(Ordering::SeqCst) - ); - - // The failure left a resumable checkpoint: `fetch`'s progress is committed - // and `commit` is scheduled to re-run. - let snapshot = graph - .get_state(thread, None) - .await? - .expect("a failure-boundary checkpoint exists"); - println!( - " durable state: fetched={}, committed={}", - snapshot.values.fetched, snapshot.values.committed - ); - println!(" will re-run on resume: {:?}", snapshot.next_nodes); - println!(" checkpoints on thread: {}", checkpointer.count(thread)); - - // ---- The outage clears; restart the run from the checkpoint. ----------- - outage.store(1, Ordering::SeqCst); - let resumed = graph.retry(thread).await?; - assert_eq!(resumed.status.status, ExecutionStatus::Completed); - println!("\nrestarted to completion:"); - println!( - " fetched={}, committed={}", - resumed.state.fetched, resumed.state.committed - ); - for line in &resumed.state.log { - println!(" - {line}"); - } - - Ok(()) -} diff --git a/crates/tinyagents-integration-tests/examples/subconscious_loop/README.md b/crates/tinyagents-integration-tests/examples/subconscious_loop/README.md deleted file mode 100644 index 3da7b485b..000000000 --- a/crates/tinyagents-integration-tests/examples/subconscious_loop/README.md +++ /dev/null @@ -1,159 +0,0 @@ -# Subconscious Loop Example - -This folder contains a runnable, fully offline example of an autonomous -closed-loop agent harness with a dedicated subconscious layer. It translates the -LangGraph-style reference design into TinyAgents' typed graph runtime while -keeping every step deterministic enough for normal `cargo test` coverage. - -Run it with: - -```text -cargo run -p tinyagents-integration-tests --example subconscious_loop -``` - -Run the integration coverage with: - -```text -cargo test --test e2e_subconscious_loop_example -``` - -## Files - -- `main.rs`: small executable entrypoint that seeds an initial channel packet, - runs the graph, and prints the final state. -- `autonomous_loop.rs`: the reusable graph, state types, node implementations, - routing functions, reducer, mock retrieval, diff summarization, and - subconscious steering logic. -- `../../tests/e2e_subconscious_loop_example.rs`: integration tests that import - the same implementation file used by the example binary, so documented and - tested behavior stay aligned. - -## Architecture - -The example models three cognitive tiers: - -- Quick layer: `frontend_agent` turns raw channel payloads into macro - instructions, then later compiles the final channel response. -- Reasoning layer: `agent_execution` retrieves from long-term memory, simulates - sub-agent execution, performs semantic extraction, emits a sequential - world-state diff, and decides whether to force a subconscious evaluation. -- Subconscious layer: `subconscious_eval` consumes gated world summaries and - returns a short steering directive, then resets the trigger state to avoid an - infinite evaluation loop. - -The graph nodes are: - -```text -channel_ingestion - -> frontend_agent - -> agent_execution - -> summarization_gate - -> frontend_agent - -> context_manager_hook - -> subconscious_eval | END -``` - -The second `frontend_agent` pass is intentional. The first pass defers work to -the Reasoning layer by writing `agent_instructions`; the second pass sees -`agent_reply` and writes `channel_response`. - -## State Model - -`SystemState` is the committed graph state. It contains the surface channel -fields, operational memory fields, routing flags, and an `event_log` used by the -example output and tests. - -Important fields: - -- `messages`: channel-visible messages. -- `agent_instructions`: macro instruction produced by the Quick layer. -- `agent_reply`: Reasoning-layer result consumed by the Quick layer. -- `semantic_history`: entity-centric extracted traces from execution. -- `sequential_diffs`: small world-state mutations accumulated across cycles. -- `gated_world_summary`: consolidated package forwarded to the subconscious - layer. -- `context_utilization`: simulated context pressure. -- `trigger_subconscious`: event escalation flag set by the Reasoning layer. -- `cron_due`: scheduled consolidation flag. -- `long_term_memory`: mock vector database records. -- `retrieved_context`: memory traces retrieved for the current run. -- `subconscious_steering`: final steering directive. - -Nodes do not overwrite `SystemState` directly. They return a `StatePatch`, and -`apply_patch` merges partial updates through `ClosureStateReducer`. This mirrors -the durable graph contract: nodes emit small updates, and the reducer is the -single deterministic fan-in point. - -## Memory Lifecycle - -The Reasoning layer performs a mock retrieval step against `long_term_memory`. -The implementation uses string matching rather than embeddings so the example -has no network calls, API keys, or nondeterministic model output. - -When `context_utilization` reaches the eviction threshold, the -`context_manager_hook` moves semantic traces into `long_term_memory`, replaces -the active `semantic_history` with an eviction marker, and lowers utilization. -This demonstrates the bi-directional memory loop: - -1. retrieve historical traces before execution; -2. execute and extract semantic history; -3. evict dense traces into long-term memory under pressure; -4. retrieve them again on future runs. - -## Summarization Gate - -The graph does not forward every diff to the subconscious layer. The -`summarization_gate` forwards only when either condition is true: - -- at least three sequential diffs are queued; -- `trigger_subconscious` is set by the Reasoning layer. - -Forwarding converts the queued `WorldDiff` values into one `GatedWorldSummary` -containing a macro trend, critical event count, and total magnitude. After -forwarding, the diff queue is cleared. If the threshold is not met, the gate -holds the diffs and records that decision in `event_log`. - -## Hybrid Triggers - -The example supports both event and cron-style triggers: - -- Event trigger: `agent_execution` sets `trigger_subconscious` when the - instruction contains critical anomaly terms or when context pressure is very - high. -- Cron trigger: callers can seed `SystemState::cron_due = true` before running - the graph to force a scheduled subconscious review. - -`subconscious_eval` always clears both `trigger_subconscious` and `cron_due`. -That reset is part of the circuit breaker that prevents repeated asynchronous -evaluation loops. - -## Offline Simulation Boundaries - -This example intentionally avoids real LLM and vector database calls. It is a -reference architecture sample, not a hosted provider integration. The simulated -parts are isolated behind small functions: - -- `retrieve_context`: stand-in for vector DB retrieval. -- `agent_execution_node`: stand-in for Reasoning LLM orchestration and - sub-agent spawning. -- `summarize_diffs`: stand-in for the summarization gate model. -- `subconscious_eval_node`: stand-in for deep reflection and steering. - -A production version would replace those functions with provider-backed model -calls and a real `Retriever`, while keeping the same graph topology, state -fields, routing flags, and reducer shape. - -## Test Coverage - -The integration test covers the main architectural guardrails: - -- normal execution retrieves memory and holds a single diff below the gate - threshold; -- critical Reasoning escalation forces the summarization gate and routes through - `subconscious_eval`; -- sequential diff accumulation triggers the subconscious layer without an - anomaly; -- context pressure evicts semantic history into long-term memory; -- the public `run_subconscious_loop` helper used by `main.rs` stays executable. - -These tests are offline and run as part of the default suite. diff --git a/crates/tinyagents-integration-tests/examples/subconscious_loop/autonomous_loop.rs b/crates/tinyagents-integration-tests/examples/subconscious_loop/autonomous_loop.rs deleted file mode 100644 index 7677122df..000000000 --- a/crates/tinyagents-integration-tests/examples/subconscious_loop/autonomous_loop.rs +++ /dev/null @@ -1,496 +0,0 @@ -//! A fully offline, testable version of the autonomous closed-loop harness. -//! -//! The example maps the Python/LangGraph reference architecture onto -//! TinyAgents' typed graph runtime: -//! -//! - channel ingestion and front-end response management run as lightweight -//! graph nodes; -//! - the reasoning node performs deterministic retrieval, sub-agent simulation, -//! semantic extraction, sequential diff generation, and event escalation; -//! - the summarization gate compresses world-state diffs before the deep layer; -//! - the context manager evicts semantic history into a mock vector store when -//! utilization crosses a threshold; -//! - the subconscious node consumes gated summaries and emits a short steering -//! directive while resetting escalation state. -//! -//! The implementation uses deterministic functions instead of live LLM calls so -//! `cargo run -p tinyagents-integration-tests --example subconscious_loop` and the integration tests stay -//! offline and reproducible. - -use tinyagents_graph::ClosureStateReducer; -use tinyagents_graph::END; -use tinyagents_graph::*; -use tinyagents_harness::*; -use tinyagents_registry::*; - -const DIFF_GATE_MIN_SEQUENCE: usize = 3; -const CONTEXT_EVICTION_THRESHOLD: f32 = 0.85; - -/// One visible channel message in the surface layer. -#[derive(Clone, Debug, Default, PartialEq, Eq)] -pub struct ChannelMessage { - pub role: String, - pub content: String, -} - -/// A compact world-state mutation produced by the reasoning layer. -#[derive(Clone, Debug, PartialEq, Eq)] -pub struct WorldDiff { - pub event: String, - pub magnitude: i32, - pub status: String, -} - -impl WorldDiff { - pub fn new(event: impl Into<String>, magnitude: i32, status: impl Into<String>) -> Self { - Self { - event: event.into(), - magnitude, - status: status.into(), - } - } -} - -/// Dense package forwarded to the subconscious layer. -#[derive(Clone, Debug, PartialEq, Eq)] -pub struct GatedWorldSummary { - pub macro_trend: String, - pub critical_events: usize, - pub total_magnitude: i32, -} - -/// A mock long-term memory record. In production this would be a vector DB row. -#[derive(Clone, Debug, PartialEq, Eq)] -pub struct MemoryRecord { - pub key: String, - pub text: String, -} - -impl MemoryRecord { - pub fn new(key: impl Into<String>, text: impl Into<String>) -> Self { - Self { - key: key.into(), - text: text.into(), - } - } -} - -/// The state carried through the closed-loop graph. -#[derive(Clone, Debug, PartialEq)] -pub struct SystemState { - pub messages: Vec<ChannelMessage>, - pub channel_source: String, - pub raw_channel_payload: String, - pub agent_instructions: String, - pub agent_reply: String, - pub channel_response: String, - pub subconscious_steering: String, - pub semantic_history: Vec<String>, - pub sequential_diffs: Vec<WorldDiff>, - pub gated_world_summary: Option<GatedWorldSummary>, - pub context_utilization: f32, - pub trigger_subconscious: bool, - pub cron_due: bool, - pub long_term_memory: Vec<MemoryRecord>, - pub retrieved_context: Vec<String>, - pub event_log: Vec<String>, -} - -impl SystemState { - /// Builds an initial channel packet with a seeded memory corpus. - pub fn new(channel_source: impl Into<String>, raw_payload: impl Into<String>) -> Self { - Self { - messages: Vec::new(), - channel_source: channel_source.into(), - raw_channel_payload: raw_payload.into(), - agent_instructions: String::new(), - agent_reply: String::new(), - channel_response: String::new(), - subconscious_steering: String::new(), - semantic_history: Vec::new(), - sequential_diffs: Vec::new(), - gated_world_summary: None, - context_utilization: 0.1, - trigger_subconscious: false, - cron_due: false, - long_term_memory: vec![ - MemoryRecord::new( - "resource", - "Historical trace: resource reallocation occurred in cycle 42.", - ), - MemoryRecord::new( - "subagent", - "Historical trace: repeated sub-agent failures respond well to lower temperature.", - ), - ], - retrieved_context: Vec::new(), - event_log: Vec::new(), - } - } -} - -/// Partial updates merged at graph step boundaries. -#[derive(Clone, Debug, Default)] -pub struct StatePatch { - pub messages: Vec<ChannelMessage>, - pub agent_instructions: Option<String>, - pub agent_reply: Option<String>, - pub channel_response: Option<String>, - pub subconscious_steering: Option<String>, - pub append_semantic_history: Vec<String>, - pub replace_semantic_history: Option<Vec<String>>, - pub append_sequential_diffs: Vec<WorldDiff>, - pub clear_sequential_diffs: bool, - pub set_gated_world_summary: Option<GatedWorldSummary>, - pub clear_gated_world_summary: bool, - pub context_utilization: Option<f32>, - pub context_utilization_delta: f32, - pub trigger_subconscious: Option<bool>, - pub cron_due: Option<bool>, - pub append_long_term_memory: Vec<MemoryRecord>, - pub append_retrieved_context: Vec<String>, - pub append_event_log: Vec<String>, -} - -impl StatePatch { - fn event(event: impl Into<String>) -> Self { - Self { - append_event_log: vec![event.into()], - ..Self::default() - } - } -} - -/// Builds the closed-loop graph. -pub fn build_subconscious_loop_graph() -> Result<CompiledGraph<SystemState, StatePatch>> { - GraphBuilder::<SystemState, StatePatch>::new() - .set_reducer(ClosureStateReducer::new(apply_patch)) - .add_node("channel_ingestion", channel_ingestion_node) - .add_node("frontend_agent", frontend_agent_node) - .add_node("agent_execution", agent_execution_node) - .add_node("summarization_gate", summarization_gate_node) - .add_node("context_manager_hook", context_manager_hook_node) - .add_node("subconscious_eval", subconscious_eval_node) - .set_entry("channel_ingestion") - .add_edge("channel_ingestion", "frontend_agent") - .add_conditional_edges( - "frontend_agent", - automated_loop_router, - [ - ("agent_execution", "agent_execution"), - ("context_manager_hook", "context_manager_hook"), - ], - ) - .add_edge("agent_execution", "summarization_gate") - .add_edge("summarization_gate", "frontend_agent") - .add_conditional_edges( - "context_manager_hook", - subconscious_router, - [ - ("subconscious_eval", "subconscious_eval"), - ("END", "__end__"), - ], - ) - .set_finish("subconscious_eval") - .with_node_kind("channel_ingestion", "channel") - .with_node_kind("frontend_agent", "quick_llm") - .with_node_kind("agent_execution", "reasoning_llm") - .with_node_kind("summarization_gate", "compression_gate") - .with_node_kind("context_manager_hook", "context_manager") - .with_node_kind("subconscious_eval", "subconscious_llm") - .compile() -} - -/// Runs the example graph from an initial state. -pub async fn run_subconscious_loop(initial: SystemState) -> Result<SystemState> { - let graph = build_subconscious_loop_graph()?; - Ok(graph.run(initial).await?.state) -} - -fn apply_patch(mut state: SystemState, patch: StatePatch) -> Result<SystemState> { - state.messages.extend(patch.messages); - - if let Some(value) = patch.agent_instructions { - state.agent_instructions = value; - } - if let Some(value) = patch.agent_reply { - state.agent_reply = value; - } - if let Some(value) = patch.channel_response { - state.channel_response = value; - } - if let Some(value) = patch.subconscious_steering { - state.subconscious_steering = value; - } - - if let Some(history) = patch.replace_semantic_history { - state.semantic_history = history; - } - state.semantic_history.extend(patch.append_semantic_history); - - if patch.clear_sequential_diffs { - state.sequential_diffs.clear(); - } - state.sequential_diffs.extend(patch.append_sequential_diffs); - - if patch.clear_gated_world_summary { - state.gated_world_summary = None; - } - if let Some(summary) = patch.set_gated_world_summary { - state.gated_world_summary = Some(summary); - } - - if let Some(value) = patch.context_utilization { - state.context_utilization = value.clamp(0.0, 1.0); - } - if patch.context_utilization_delta != 0.0 { - state.context_utilization = - (state.context_utilization + patch.context_utilization_delta).clamp(0.0, 1.0); - } - - if let Some(value) = patch.trigger_subconscious { - state.trigger_subconscious = value; - } - if let Some(value) = patch.cron_due { - state.cron_due = value; - } - - state.long_term_memory.extend(patch.append_long_term_memory); - state - .retrieved_context - .extend(patch.append_retrieved_context); - state.event_log.extend(patch.append_event_log); - - Ok(state) -} - -fn channel_ingestion_node(state: SystemState, _ctx: NodeContext) -> impl FutureNode { - async move { - Ok(NodeResult::Update(StatePatch { - messages: vec![ChannelMessage { - role: "user".to_string(), - content: state.raw_channel_payload.clone(), - }], - append_event_log: vec![format!( - "channel_ingestion: packet from {}", - state.channel_source - )], - ..StatePatch::default() - })) - } -} - -fn frontend_agent_node(state: SystemState, _ctx: NodeContext) -> impl FutureNode { - async move { - if state.agent_reply.is_empty() { - Ok(NodeResult::Update(StatePatch { - agent_instructions: Some(format!( - "AUTONOMOUS_EXECUTE: {}", - state.raw_channel_payload - )), - append_event_log: vec!["frontend_agent: deferred macro instruction".to_string()], - ..StatePatch::default() - })) - } else { - Ok(NodeResult::Update(StatePatch { - channel_response: Some(format!("Completed. Output: {}", state.agent_reply)), - append_event_log: vec!["frontend_agent: compiled final response".to_string()], - ..StatePatch::default() - })) - } - } -} - -fn agent_execution_node(state: SystemState, _ctx: NodeContext) -> impl FutureNode { - async move { - let instructions = state.agent_instructions.clone(); - let retrieved = retrieve_context(&state.long_term_memory, &instructions); - let anomaly = contains_any( - &instructions, - &["critical", "cascade", "failure", "anomaly"], - ); - let magnitude = if anomaly { - 91 - } else if contains_any(&instructions, &["resource", "matrix", "allocation"]) { - 14 - } else { - 4 - }; - let status = if anomaly { "degraded" } else { "stable" }; - let semantic_trace = format!( - "[Semantic Extraction] entity=sub_agents event=resource_shift magnitude={magnitude}% status={status} error_code={}", - if anomaly { "SUBAGENT_CASCADE" } else { "NONE" } - ); - let force_subconscious = anomaly || state.context_utilization >= 0.9; - - Ok(NodeResult::Update(StatePatch { - agent_reply: Some(format!( - "Pipeline mutations executed via sub-agents; {} memory traces retrieved.", - retrieved.len() - )), - append_semantic_history: vec![semantic_trace], - append_sequential_diffs: vec![WorldDiff::new("resource_shift", magnitude, status)], - trigger_subconscious: Some(force_subconscious), - context_utilization_delta: if anomaly { 0.2 } else { 0.05 }, - append_retrieved_context: retrieved, - append_event_log: vec![ - "agent_execution: retrieved long-term context".to_string(), - format!("agent_execution: emitted sequential diff status={status}"), - ], - ..StatePatch::default() - })) - } -} - -fn summarization_gate_node(state: SystemState, _ctx: NodeContext) -> impl FutureNode { - async move { - let should_forward = - state.sequential_diffs.len() >= DIFF_GATE_MIN_SEQUENCE || state.trigger_subconscious; - - if !should_forward { - return Ok(NodeResult::Update(StatePatch::event( - "summarization_gate: threshold not met; holding diffs", - ))); - } - - Ok(NodeResult::Update(StatePatch { - set_gated_world_summary: Some(summarize_diffs(&state.sequential_diffs)?), - clear_sequential_diffs: true, - append_event_log: vec![ - "summarization_gate: forwarded consolidated world summary".to_string(), - ], - ..StatePatch::default() - })) - } -} - -fn context_manager_hook_node(state: SystemState, _ctx: NodeContext) -> impl FutureNode { - async move { - if state.context_utilization < CONTEXT_EVICTION_THRESHOLD { - return Ok(NodeResult::Update(StatePatch::event( - "context_manager_hook: utilization below eviction threshold", - ))); - } - - let records = state - .semantic_history - .iter() - .enumerate() - .map(|(idx, trace)| MemoryRecord::new(format!("evicted-trace-{idx}"), trace.clone())) - .collect(); - - Ok(NodeResult::Update(StatePatch { - replace_semantic_history: Some(vec![ - "--- semantic history evicted to vector DB ---".to_string(), - ]), - append_long_term_memory: records, - context_utilization: Some(0.2), - append_event_log: vec![ - "context_manager_hook: evicted semantic history to long-term memory".to_string(), - ], - ..StatePatch::default() - })) - } -} - -fn subconscious_eval_node(state: SystemState, _ctx: NodeContext) -> impl FutureNode { - async move { - let summary = state - .gated_world_summary - .clone() - .unwrap_or(GatedWorldSummary { - macro_trend: "Cron review found no gated world-state package.".to_string(), - critical_events: 0, - total_magnitude: 0, - }); - - let directive = if summary.critical_events > 0 || summary.total_magnitude >= 50 { - "STEERING_DIRECTIVE: High mutability detected. Lower sub-agent temperature and require retrieval before spawning." - } else { - "STEERING_DIRECTIVE: System stable. Preserve current execution policy." - }; - - Ok(NodeResult::Update(StatePatch { - subconscious_steering: Some(directive.to_string()), - trigger_subconscious: Some(false), - cron_due: Some(false), - clear_gated_world_summary: true, - append_event_log: vec![format!( - "subconscious_eval: digested summary '{}'", - summary.macro_trend - )], - ..StatePatch::default() - })) - } -} - -fn automated_loop_router(state: &SystemState) -> &'static str { - if state.channel_response.is_empty() { - "agent_execution" - } else { - "context_manager_hook" - } -} - -fn subconscious_router(state: &SystemState) -> &'static str { - if state.trigger_subconscious || state.cron_due || state.gated_world_summary.is_some() { - "subconscious_eval" - } else { - "END" - } -} - -fn retrieve_context(memory: &[MemoryRecord], query: &str) -> Vec<String> { - let query = query.to_ascii_lowercase(); - memory - .iter() - .filter(|record| { - query.contains(&record.key.to_ascii_lowercase()) - || record - .key - .split(['-', '_']) - .any(|term| !term.is_empty() && query.contains(term)) - }) - .map(|record| record.text.clone()) - .collect() -} - -fn summarize_diffs(diffs: &[WorldDiff]) -> Result<GatedWorldSummary> { - if diffs.is_empty() { - return Err(TinyAgentsError::Validation( - "summarization gate received no diffs".to_string(), - )); - } - - let critical_events = diffs - .iter() - .filter(|diff| diff.status != "stable" || diff.magnitude >= 50) - .count(); - let total_magnitude = diffs.iter().map(|diff| diff.magnitude).sum(); - - Ok(GatedWorldSummary { - macro_trend: format!( - "Aggregated {} operational shifts; {} critical events; total magnitude {}%.", - diffs.len(), - critical_events, - total_magnitude - ), - critical_events, - total_magnitude, - }) -} - -fn contains_any(haystack: &str, needles: &[&str]) -> bool { - let haystack = haystack.to_ascii_lowercase(); - needles.iter().any(|needle| haystack.contains(needle)) -} - -trait FutureNode: - std::future::Future<Output = Result<NodeResult<StatePatch>>> + std::marker::Send + 'static -{ -} - -impl<T> FutureNode for T where - T: std::future::Future<Output = Result<NodeResult<StatePatch>>> + std::marker::Send + 'static -{ -} diff --git a/crates/tinyagents-integration-tests/examples/subconscious_loop/main.rs b/crates/tinyagents-integration-tests/examples/subconscious_loop/main.rs deleted file mode 100644 index 37f486abf..000000000 --- a/crates/tinyagents-integration-tests/examples/subconscious_loop/main.rs +++ /dev/null @@ -1,33 +0,0 @@ -//! Runs the autonomous subconscious-loop graph example. -//! -//! ```text -//! cargo run -p tinyagents-integration-tests --example subconscious_loop -//! ``` - -mod autonomous_loop; - -use autonomous_loop::{SystemState, run_subconscious_loop}; -use tinyagents_harness::Result; - -#[tokio::main] -async fn main() -> Result<()> { - let final_state = run_subconscious_loop(SystemState::new( - "telegram", - "Reallocate the resource matrix and report whether the system remains stable.", - )) - .await?; - - println!("=== Autonomous closed-loop subconscious example ==="); - println!("channel response : {}", final_state.channel_response); - println!("retrieved context: {:?}", final_state.retrieved_context); - println!("semantic history : {:?}", final_state.semantic_history); - println!("sequential diffs : {:?}", final_state.sequential_diffs); - println!("steering : {}", final_state.subconscious_steering); - println!("context usage : {:.2}", final_state.context_utilization); - println!("visited events :"); - for event in &final_state.event_log { - println!(" - {event}"); - } - - Ok(()) -} diff --git a/crates/tinyagents-integration-tests/tests/common/storage.rs b/crates/tinyagents-integration-tests/tests/common/storage.rs index 8a41afe54..fb8bfe22f 100644 --- a/crates/tinyagents-integration-tests/tests/common/storage.rs +++ b/crates/tinyagents-integration-tests/tests/common/storage.rs @@ -5,9 +5,8 @@ use std::path::Path; use std::sync::Arc; -use tinyagents_graph::checkpoint::DriverCheckpointer; -use tinyagents_session::port::AgentStorageFactory; -use tinyagents_storage_kit::{DocumentStore, Scope, ScopedStorage, StreamStore}; +use tinyagents_session::port::{AgentStorageFactory, ConditionedAgentStorageFactory}; +use tinyagents_storage_kit::{DocumentStore, Scope, ScopedStorage, StreamStore, WriteCondition}; use tinystoragedrivers_api::{Storage, StorageBackend}; use tinystoragedrivers_document::sqlite::{SqliteDocuments, SqliteStreams}; use tinystoragedrivers_testkit::MemoryBackend; @@ -25,12 +24,28 @@ pub fn memory() -> ScopedStorage { /// exercise durable document/stream state only. pub fn sqlite(dir: &Path, name: &str) -> ScopedStorage { let scope = Scope::new("integration-sqlite").unwrap(); + sqlite_storage(dir, name, &scope, None) +} + +fn sqlite_storage( + dir: &Path, + name: &str, + scope: &Scope, + condition: Option<&WriteCondition>, +) -> ScopedStorage { let connection = tinystoragedrivers_connect::sqlite::SqliteConnection::open(dir.join(name)).unwrap(); - let documents = - Arc::new(SqliteDocuments::new(connection.clone(), scope.clone())) as Arc<dyn DocumentStore>; + let documents = match condition { + Some(condition) => Arc::new(SqliteDocuments::conditioned( + connection.clone(), + scope.clone(), + condition.clone(), + )) as Arc<dyn DocumentStore>, + None => Arc::new(SqliteDocuments::new(connection.clone(), scope.clone())) + as Arc<dyn DocumentStore>, + }; let streams = Arc::new(SqliteStreams::new(connection, scope.clone())) as Arc<dyn StreamStore>; - let memory = MemoryBackend::new().for_scope(&scope).unwrap(); + let memory = MemoryBackend::new().for_scope(scope).unwrap(); let storage = Storage::new( documents, streams, @@ -39,7 +54,7 @@ pub fn sqlite(dir: &Path, name: &str) -> ScopedStorage { memory.queues().clone(), memory.secrets().clone(), ); - ScopedStorage::new(storage, scope) + ScopedStorage::new(storage, scope.clone()) } /// The document port of `storage`. @@ -52,7 +67,24 @@ pub fn fixed_factory(storage: ScopedStorage) -> AgentStorageFactory { Arc::new(move |_| Ok(storage.clone())) } -/// A durable checkpointer over `checkpoints.sqlite` in `dir`. -pub fn checkpointer<State>(dir: &Path) -> DriverCheckpointer<State> { - DriverCheckpointer::new(docs(&sqlite(dir, "checkpoints.sqlite"))) +/// Host factories over one SQLite scope, including the native conditioned +/// opener required for durable transcript-generation writes. +pub fn conditioned_sqlite_factories( + dir: &Path, + name: &str, +) -> (AgentStorageFactory, ConditionedAgentStorageFactory) { + let directory = dir.to_path_buf(); + let name = name.to_owned(); + let scope = Scope::new("integration-conditioned-sqlite").unwrap(); + let plain = { + let directory = directory.clone(); + let name = name.clone(); + let scope = scope.clone(); + Arc::new(move |_: &str| Ok(sqlite_storage(&directory, &name, &scope, None))) + as AgentStorageFactory + }; + let conditioned = Arc::new(move |_: &str, condition: &WriteCondition| { + Ok(sqlite_storage(&directory, &name, &scope, Some(condition))) + }) as ConditionedAgentStorageFactory; + (plain, conditioned) } diff --git a/crates/tinyagents-integration-tests/tests/conformance.rs b/crates/tinyagents-integration-tests/tests/conformance.rs deleted file mode 100644 index 6cd9f3823..000000000 --- a/crates/tinyagents-integration-tests/tests/conformance.rs +++ /dev/null @@ -1,97 +0,0 @@ -//! Storage conformance: the same contract suite applied to every built-in -//! backend, proving they behave interchangeably (gap #17). - -mod common; - -use common::storage; -use tinyagents_graph::checkpoint::{DriverCheckpointer, InMemoryCheckpointer}; -use tinyagents_graph::orchestration::{DriverTaskStore, InMemoryTaskStore}; -use tinyagents_graph::testkit::conformance::{ - checkpointer_concurrent_contract, checkpointer_contract, task_cache_contract, - taskstore_concurrent_contract, taskstore_contract, taskstore_replay_contract, -}; -use tinyagents_graph::{DriverTaskCache, InMemoryTaskCache}; - -#[tokio::test] -async fn in_memory_checkpointer_satisfies_contract() { - checkpointer_contract(InMemoryCheckpointer::<i32>::new()).await; -} - -#[tokio::test] -async fn memory_driver_checkpointer_satisfies_contract() { - checkpointer_contract(DriverCheckpointer::<i32>::new(storage::docs( - &storage::memory(), - ))) - .await; -} - -#[tokio::test] -async fn sqlite_driver_checkpointer_satisfies_contract() { - let dir = tempfile::tempdir().unwrap(); - checkpointer_contract(storage::checkpointer::<i32>(dir.path())).await; -} - -#[test] -fn in_memory_task_store_satisfies_contract() { - taskstore_contract(InMemoryTaskStore::new()); -} - -#[test] -fn memory_driver_task_store_satisfies_contract() { - taskstore_contract(DriverTaskStore::new(storage::docs(&storage::memory())).unwrap()); -} - -#[test] -fn sqlite_driver_task_store_satisfies_contract() { - let dir = tempfile::tempdir().unwrap(); - let store = - DriverTaskStore::new(storage::docs(&storage::sqlite(dir.path(), "tasks.sqlite"))).unwrap(); - taskstore_contract(store); -} - -#[tokio::test] -async fn in_memory_checkpointer_handles_concurrent_writes() { - checkpointer_concurrent_contract(std::sync::Arc::new(InMemoryCheckpointer::<i32>::new())).await; -} - -#[tokio::test] -async fn sqlite_driver_checkpointer_handles_concurrent_writes() { - let dir = tempfile::tempdir().unwrap(); - checkpointer_concurrent_contract(std::sync::Arc::new(storage::checkpointer::<i32>( - dir.path(), - ))) - .await; -} - -#[test] -fn in_memory_task_store_handles_concurrent_writes() { - taskstore_concurrent_contract(std::sync::Arc::new(InMemoryTaskStore::new())); -} - -#[test] -fn sqlite_driver_task_store_handles_concurrent_writes() { - let dir = tempfile::tempdir().unwrap(); - let store = - DriverTaskStore::new(storage::docs(&storage::sqlite(dir.path(), "tasks.sqlite"))).unwrap(); - taskstore_concurrent_contract(std::sync::Arc::new(store)); -} - -#[test] -fn sqlite_driver_task_store_replays_after_restart() { - let dir = tempfile::tempdir().unwrap(); - taskstore_replay_contract(|| { - DriverTaskStore::new(storage::docs(&storage::sqlite(dir.path(), "tasks.sqlite"))).unwrap() - }); -} - -#[tokio::test] -async fn task_caches_satisfy_the_contract() { - task_cache_contract(InMemoryTaskCache::new()).await; - task_cache_contract(DriverTaskCache::new(storage::docs(&storage::memory()))).await; - let dir = tempfile::tempdir().unwrap(); - task_cache_contract(DriverTaskCache::new(storage::docs(&storage::sqlite( - dir.path(), - "cache.sqlite", - )))) - .await; -} diff --git a/crates/tinyagents-integration-tests/tests/e2e_agent_graph.rs b/crates/tinyagents-integration-tests/tests/e2e_agent_graph.rs deleted file mode 100644 index 34ab1aead..000000000 --- a/crates/tinyagents-integration-tests/tests/e2e_agent_graph.rs +++ /dev/null @@ -1,138 +0,0 @@ -//! TRUE end-to-end: a durable graph node that drives a full [`AgentHarness`] -//! agent run. -//! -//! This composes three subsystems that the per-module integration tests -//! exercise in isolation: -//! -//! - the **graph** executor (`GraphBuilder` / `NodeResult`), -//! - the **harness** agent loop (model → tool → model), and -//! - the **testkit** (`ScriptedModel`, `FakeTool`, `EventRecorder`, -//! `Trajectory`) plus a real **middleware** (`UsageAccountingMiddleware`). -//! -//! A single graph node owns an `Arc<AgentHarness>` and invokes it inside a -//! caller-supplied `RunContext` whose `EventSink` is wired to a shared -//! `EventRecorder`. After the graph reaches its terminal node we assert on the -//! accumulated state, the recorded event trajectory (the tool really ran), and -//! the usage totals folded by the middleware — never on model prose. - -use std::sync::Arc; - -use serde_json::json; - -use tinyagents_graph::*; -use tinyagents_harness::context::{RunConfig, RunContext}; -use tinyagents_harness::middleware::UsageAccountingMiddleware; -use tinyagents_harness::runtime::AgentHarness; -use tinyagents_harness::testkit::{EventRecorder, FakeTool, ScriptedModel, Trajectory}; -use tinyagents_harness::*; -use tinyagents_registry::*; -use tinyinference_llm::message::{AssistantMessage, ContentBlock, Message}; -use tinyinference_llm::model::ModelResponse; -use tinyinference_llm::tool::ToolCall; -use tinyinference_llm::usage::Usage; - -/// State threaded through the graph: the question to answer and the agent's -/// final text once the `agent` node has run. -#[derive(Clone, Debug, Default)] -struct AgentGraphState { - question: String, - answer: Option<String>, -} - -/// Builds an assistant response that asks for a single tool call. -fn tool_call_response(id: &str, name: &str, arguments: serde_json::Value) -> ModelResponse { - tinyagents_harness::testkit::tool_call_response(ToolCall::new(id, name, arguments)) - .with_usage(Usage::new(6, 2)) -} - -/// Builds a plain-text final assistant response. -fn text_response(text: &str, input: u64, output: u64) -> ModelResponse { - tinyagents_harness::testkit::text_response(text).with_usage(Usage::new(input, output)) -} - -#[tokio::test] -async fn graph_node_drives_harness_agent_with_tool_loop() { - // Shared recorder: the node feeds the agent run's events here so the test - // can reconstruct a Trajectory after the graph completes. - let recorder = Arc::new(EventRecorder::new()); - - // Usage accounting middleware folds every model response's usage into a - // running total we read back after the run. - let usage_mw = Arc::new(UsageAccountingMiddleware::new()); - - // Build the harness: script a tool call then a final answer, register a - // real tool, and attach the usage middleware. - let mut harness: AgentHarness<()> = AgentHarness::new(); - harness - .register_model( - "scripted", - Arc::new(ScriptedModel::new(vec![ - tool_call_response("call-1", "lookup_user", json!({ "user_id": "u-1" })), - text_response("resolved via lookup", 4, 3), - ])), - ) - .set_default_model("scripted") - .register_tool(Arc::new(FakeTool::returning("lookup_user", "Ada Lovelace"))) - .push_middleware(usage_mw.clone()); - let harness = Arc::new(harness); - - // A single graph node that runs the agent and folds the answer into the - // whole-state graph state. - let node_harness = harness.clone(); - let node_recorder = recorder.clone(); - let graph = GraphBuilder::<AgentGraphState, AgentGraphState>::overwrite() - .add_node( - "agent", - move |mut state: AgentGraphState, _ctx: NodeContext| { - let harness = node_harness.clone(); - let recorder = node_recorder.clone(); - async move { - // Wire the agent run's events into the shared recorder via a - // caller-supplied RunContext. - let ctx: RunContext<()> = RunContext::new(RunConfig::new("agent-node-run"), ()) - .with_events(recorder.sink()); - - let run = harness - .invoke_in_context(&(), ctx, vec![Message::user(state.question.clone())]) - .await?; - - state.answer = run.text(); - Ok(NodeResult::Update(state)) - } - }, - ) - .set_entry("agent") - .set_finish("agent") - .compile() - .expect("graph compiles"); - - let run = graph - .run(AgentGraphState { - question: "who is user u-1?".to_string(), - answer: None, - }) - .await - .expect("graph runs to completion"); - - // Graph-level: the terminal node ran and the agent's answer landed in state. - let visited: Vec<String> = run.visited.iter().map(ToString::to_string).collect(); - assert_eq!(visited, vec!["agent"]); - assert_eq!(run.state.answer.as_deref(), Some("resolved via lookup")); - - // Harness-level via the trajectory: the tool actually fired, between two - // model calls, and the run completed. - let traj = Trajectory::from_events(recorder.events()); - traj.assert_tool_called("lookup_user"); - assert_eq!(traj.tool_call_count("lookup_user"), 1); - traj.assert_model_called_times(2); - traj.assert_completed(); - traj.assert_order(&["model.started", "lookup_user", "model.completed"]) - .expect("tool runs between the two model calls"); - - // Middleware-level: usage was folded across both model calls. - let totals = usage_mw.totals(); - assert_eq!(totals.calls, 2); - assert_eq!(totals.usage.input_tokens, 6 + 4); - assert_eq!(totals.usage.output_tokens, 2 + 3); - assert_eq!(totals.usage.total_tokens, 15); -} diff --git a/crates/tinyagents-integration-tests/tests/e2e_budget.rs b/crates/tinyagents-integration-tests/tests/e2e_budget.rs index 0e68e7c84..d061a8319 100644 --- a/crates/tinyagents-integration-tests/tests/e2e_budget.rs +++ b/crates/tinyagents-integration-tests/tests/e2e_budget.rs @@ -123,8 +123,9 @@ async fn token_budget_blocks_multi_call_run() { let mut harness: AgentHarness<()> = AgentHarness::new(); harness - .register_model("mock", Arc::new(model)) - .set_default_model("mock") + .with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named("mock", Arc::new(model)), + )) .register_tool(Arc::new(FakeTool::returning("noop", "ok"))) .push_middleware(Arc::new(mw)); @@ -183,8 +184,9 @@ async fn shared_tracker_rolls_up_and_blocks_across_runs() { ]); let mut harness_a: AgentHarness<()> = AgentHarness::new(); harness_a - .register_model("mock", Arc::new(model_a)) - .set_default_model("mock") + .with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named("mock", Arc::new(model_a)), + )) .register_tool(Arc::new(FakeTool::returning("noop", "ok"))) .push_middleware(Arc::new( BudgetMiddleware::new(limits).with_tracker(tracker.clone()), @@ -213,8 +215,9 @@ async fn shared_tracker_rolls_up_and_blocks_across_runs() { ]); let mut harness_b: AgentHarness<()> = AgentHarness::new(); harness_b - .register_model("mock", Arc::new(model_b)) - .set_default_model("mock") + .with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named("mock", Arc::new(model_b)), + )) .register_tool(Arc::new(FakeTool::returning("noop", "ok"))) .push_middleware(Arc::new( BudgetMiddleware::new(limits).with_tracker(tracker.clone()), @@ -363,8 +366,9 @@ async fn below_threshold_run_completes_without_exceeding() { let mut harness: AgentHarness<()> = AgentHarness::new(); harness - .register_model("mock", Arc::new(model)) - .set_default_model("mock") + .with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named("mock", Arc::new(model)), + )) .register_tool(Arc::new(FakeTool::returning("noop", "ok"))) .push_middleware(Arc::new(mw)); @@ -415,8 +419,9 @@ async fn budget_middleware_does_not_double_record_usage() { }); let mut harness: AgentHarness<()> = AgentHarness::new(); harness - .register_model("mock", Arc::new(model)) - .set_default_model("mock") + .with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named("mock", Arc::new(model)), + )) .register_tool(Arc::new(FakeTool::returning("noop", "ok"))) .push_middleware(Arc::new(mw)); @@ -451,8 +456,12 @@ async fn input_reservation_blocks_oversized_call_in_a_live_run() { let mut harness: AgentHarness<()> = AgentHarness::new(); harness - .register_model("mock", Arc::new(MockModel::constant("unused"))) - .set_default_model("mock") + .with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named( + "mock", + Arc::new(MockModel::constant("unused")), + ), + )) .push_middleware(Arc::new(mw)); // A long prompt estimates well over the 5-token input reservation budget. @@ -492,11 +501,12 @@ async fn fitting_call_emits_reserved_and_reconciled_events() { let mut harness: AgentHarness<()> = AgentHarness::new(); harness - .register_model( - "mock", - Arc::new(MockModel::with_responses(vec![text_response("done", 3, 1)])), - ) - .set_default_model("mock") + .with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named( + "mock", + Arc::new(MockModel::with_responses(vec![text_response("done", 3, 1)])), + ), + )) .push_middleware(Arc::new(mw)); let ctx = RunContext::new(RunConfig::new("reserve-ok"), ()).with_events(recorder.sink()); diff --git a/crates/tinyagents-integration-tests/tests/e2e_complex_graph.rs b/crates/tinyagents-integration-tests/tests/e2e_complex_graph.rs deleted file mode 100644 index 8bd6684c1..000000000 --- a/crates/tinyagents-integration-tests/tests/e2e_complex_graph.rs +++ /dev/null @@ -1,351 +0,0 @@ -//! End-to-end coverage for COMPLEX durable-graph workflows. -//! -//! These offline tests exercise the harder compositions the durable graph -//! runtime is meant to support: -//! -//! - **Graph + subgraph nesting** using BOTH embedding modes in one parent: -//! [`shared_subgraph_node`] (parent and child share the `State` channel) and -//! [`adapter_subgraph_node`] (parent state is projected into a differently -//! shaped child and folded back). -//! - **Deep nesting** (a subgraph that itself embeds a subgraph, >= 2 levels) -//! running to completion with correctly merged state, while child checkpoints -//! stay isolated and never collide with the parent's checkpoint thread. -//! - **A complex layout**: a `Command::goto` fan-out into a parallel branch set -//! (`with_parallel`), a join node, a conditional route, and a loop-back edge -//! bounded deterministically by the recursion limit. -//! -//! All scenarios are deterministic and run without network access. - -use std::sync::Arc; - -use tinyagents_graph::END; -use tinyagents_graph::*; -use tinyagents_graph::{ClosureStateReducer, adapter_subgraph_node, shared_subgraph_node}; -use tinyagents_harness::*; -use tinyagents_registry::*; - -// --------------------------------------------------------------------------- -// Scenario 1: graph + subgraph nesting via BOTH shared and adapter embedding. -// --------------------------------------------------------------------------- - -/// A document the parent graph carries end to end. -#[derive(Clone, Debug, Default, PartialEq, Eq)] -struct Doc { - text: String, - word_count: usize, - sentiment: i32, - log: Vec<String>, -} - -/// A shared-state child subgraph that normalizes a [`Doc`] in two steps. It -/// uses the same `State`/`Update` channel as the parent (whole-state overwrite), -/// so it can be embedded with [`shared_subgraph_node`]. -fn normalize_child() -> CompiledGraph<Doc, Doc> { - GraphBuilder::<Doc, Doc>::overwrite() - .add_node("upper", |mut d: Doc, _c: NodeContext| async move { - d.text = d.text.to_uppercase(); - d.log.push("upper".to_string()); - Ok(NodeResult::Update(d)) - }) - .add_node("count", |mut d: Doc, _c: NodeContext| async move { - d.word_count = d.text.split_whitespace().count(); - d.log.push("count".to_string()); - Ok(NodeResult::Update(d)) - }) - .set_entry("upper") - .add_edge("upper", "count") - .set_finish("count") - .compile() - .expect("normalize child compiles") -} - -/// A differently-shaped child subgraph that scores a bare `i32`: doubles it then -/// adds one. It is embedded via [`adapter_subgraph_node`] because its state shape -/// (`i32`) differs from the parent's (`Doc`). -fn score_child() -> CompiledGraph<i32, i32> { - GraphBuilder::<i32, i32>::overwrite() - .add_node("double", |s: i32, _c: NodeContext| async move { - Ok(NodeResult::Update(s * 2)) - }) - .add_node("plus_one", |s: i32, _c: NodeContext| async move { - Ok(NodeResult::Update(s + 1)) - }) - .set_entry("double") - .add_edge("double", "plus_one") - .set_finish("plus_one") - .compile() - .expect("score child compiles") -} - -#[tokio::test] -async fn parent_nests_shared_and_adapter_subgraphs() { - let parent = GraphBuilder::<Doc, Doc>::overwrite() - // Shared-state subgraph: runs over the whole parent Doc. - .add_node("normalize", shared_subgraph_node(normalize_child())) - // Adapter subgraph: project Doc -> i32, fold the i32 score back into Doc. - .add_node( - "score", - adapter_subgraph_node( - score_child(), - |d: &Doc| d.text.chars().count() as i32, - |d: &Doc, child_score: i32| { - let mut d = d.clone(); - d.sentiment = child_score; - d.log.push("scored".to_string()); - d - }, - ), - ) - .set_entry("normalize") - .add_edge("normalize", "score") - .set_finish("score") - .compile() - .expect("parent compiles"); - - let run = parent - .run(Doc { - text: "hello world".to_string(), - ..Doc::default() - }) - .await - .expect("run succeeds"); - - // The shared child uppercased the text and counted words. - assert_eq!(run.state.text, "HELLO WORLD"); - assert_eq!(run.state.word_count, 2); - // The adapter child scored chars("HELLO WORLD") = 11 -> *2 + 1 = 23. - assert_eq!(run.state.sentiment, 23); - // Both children's work is reflected in the parent's merged log. - assert_eq!(run.state.log, vec!["upper", "count", "scored"]); - // The PARENT only sees its own two nodes; subgraph internals are private. - let visited: Vec<&str> = run.visited.iter().map(|n| n.as_str()).collect(); - assert_eq!(visited, vec!["normalize", "score"]); -} - -// --------------------------------------------------------------------------- -// Scenario 2: deep nesting (subgraph inside a subgraph) + checkpoint isolation. -// --------------------------------------------------------------------------- - -#[tokio::test] -async fn deep_nested_subgraphs_merge_state_without_checkpoint_collision() { - // Level 3 (innermost): +100, with its own checkpointer. - let inner_ckpt = Arc::new(InMemoryCheckpointer::<i32>::new()); - let inner = GraphBuilder::<i32, i32>::overwrite() - .add_node("add100", |s: i32, _c: NodeContext| async move { - Ok(NodeResult::Update(s + 100)) - }) - .set_entry("add100") - .set_finish("add100") - .compile() - .expect("inner compiles") - .with_checkpointer(inner_ckpt.clone()); - - // Level 2 (middle): +10, then run the inner subgraph. - let mid_ckpt = Arc::new(InMemoryCheckpointer::<i32>::new()); - let middle = GraphBuilder::<i32, i32>::overwrite() - .add_node("add10", |s: i32, _c: NodeContext| async move { - Ok(NodeResult::Update(s + 10)) - }) - .add_node("inner", shared_subgraph_node(inner)) - .set_entry("add10") - .add_edge("add10", "inner") - .set_finish("inner") - .compile() - .expect("middle compiles") - .with_checkpointer(mid_ckpt.clone()); - - // Level 1 (parent/top): +1, then run the middle subgraph. - let parent_ckpt = Arc::new(InMemoryCheckpointer::<i32>::new()); - let parent = GraphBuilder::<i32, i32>::overwrite() - .add_node("add1", |s: i32, _c: NodeContext| async move { - Ok(NodeResult::Update(s + 1)) - }) - .add_node("middle", shared_subgraph_node(middle)) - .set_entry("add1") - .add_edge("add1", "middle") - .set_finish("middle") - .compile() - .expect("parent compiles") - .with_checkpointer(parent_ckpt.clone()); - - let run = parent - .run_with_thread("deep", 0) - .await - .expect("threaded run succeeds"); - - // 0 +1 (parent) +10 (middle) +100 (inner) = 111, merged across all levels. - assert_eq!(run.state, 111); - assert!(!run.is_interrupted()); - - // The top-level run persists one checkpoint per superstep boundary (add1, - // middle), all under the empty top-level namespace. - let list = parent_ckpt.list("deep").await.expect("list succeeds"); - assert_eq!(list.len(), 2); - assert!(list.iter().all(|m| m.namespace.is_empty())); - // Checkpoint ids are unique within the thread (parent chained correctly). - assert_eq!(list[0].parent_checkpoint_id, None); - assert_eq!( - list[1].parent_checkpoint_id.as_deref(), - Some(list[0].checkpoint_id.as_str()) - ); - - // Embedded subgraphs inherit the parent thread id, but their checkpoints - // stay isolated by checkpointer and embedding namespace. - let middle_list = mid_ckpt.list("deep").await.expect("middle list succeeds"); - assert_eq!(middle_list.len(), 2); - assert!( - middle_list - .iter() - .all(|m| m.namespace == vec!["middle".to_string()]) - ); - - let inner_list = inner_ckpt.list("deep").await.expect("inner list succeeds"); - assert_eq!(inner_list.len(), 1); - assert!( - inner_list - .iter() - .all(|m| m.namespace == vec!["inner".to_string()]) - ); -} - -// --------------------------------------------------------------------------- -// Scenario 3: complex layout — fan-out -> parallel branches -> join, plus a -// conditional loop-back bounded by the recursion limit. -// --------------------------------------------------------------------------- - -/// The pipeline state: an append-only audit log, two accumulators written by -/// the parallel branches, and a round counter incremented at each join. -#[derive(Clone, Debug, Default, PartialEq, Eq)] -struct Pipe { - log: Vec<String>, - a: i32, - b: i32, - rounds: i32, -} - -/// A partial update folded by the reducer. Because the parallel branches each -/// write a different field, the reducer is the deterministic fan-in/join: every -/// branch's patch is applied in stable active-set order. -#[derive(Clone, Debug, Default)] -struct Patch { - log: Vec<String>, - add_a: i32, - add_b: i32, - inc_round: bool, -} - -fn pipe_reducer() --> ClosureStateReducer<Pipe, Patch, impl Fn(Pipe, Patch) -> tinyagents_harness::Result<Pipe>> { - ClosureStateReducer::new(|mut s: Pipe, p: Patch| { - s.log.extend(p.log); - s.a += p.add_a; - s.b += p.add_b; - if p.inc_round { - s.rounds += 1; - } - Ok(s) - }) -} - -/// Builds the fan-out/parallel/join/loop-back graph. `max_rounds` bounds the -/// loop via the join's conditional router; `recursion_limit` is the hard ceiling. -fn complex_graph(max_rounds: i32, recursion_limit: usize) -> CompiledGraph<Pipe, Patch> { - GraphBuilder::<Pipe, Patch>::new() - .set_reducer(pipe_reducer()) - .with_parallel(true) - .with_recursion_limit(recursion_limit) - // Fan-out: route to BOTH parallel branches via an explicit goto command. - .add_node("dispatch", |_s: Pipe, _c: NodeContext| async move { - Ok(NodeResult::Command( - tinyagents_graph::Command::goto(["work_a", "work_b"]).with_update(Patch { - log: vec!["dispatch".to_string()], - ..Patch::default() - }), - )) - }) - .mark_command_routing("dispatch") - // Two branches that run concurrently in one superstep. - .add_node("work_a", |_s: Pipe, _c: NodeContext| async move { - Ok(NodeResult::Update(Patch { - log: vec!["a".to_string()], - add_a: 1, - ..Patch::default() - })) - }) - .add_node("work_b", |_s: Pipe, _c: NodeContext| async move { - Ok(NodeResult::Update(Patch { - log: vec!["b".to_string()], - add_b: 10, - ..Patch::default() - })) - }) - // Join: both branches' static edges converge here (deduped to one run). - .add_node("join", |_s: Pipe, _c: NodeContext| async move { - Ok(NodeResult::Update(Patch { - log: vec!["join".to_string()], - inc_round: true, - ..Patch::default() - })) - }) - .set_entry("dispatch") - .add_edge("work_a", "join") - .add_edge("work_b", "join") - // Conditional loop-back: keep looping until we hit `max_rounds`. - .add_conditional_edges( - "join", - move |s: &Pipe| { - if s.rounds < max_rounds { - "again".to_string() - } else { - "done".to_string() - } - }, - [("again", "dispatch"), ("done", tinyagents_graph::END)], - ) - .compile() - .expect("complex graph compiles") -} - -#[tokio::test] -async fn fan_out_parallel_join_with_bounded_loop_back() { - // Loop runs exactly two rounds, then routes to END. - let graph = complex_graph(2, 50); - let run = graph.run(Pipe::default()).await.expect("run succeeds"); - - // Two rounds of (a += 1, b += 10); two joins increment rounds to 2. - assert_eq!(run.state.a, 2); - assert_eq!(run.state.b, 20); - assert_eq!(run.state.rounds, 2); - assert_eq!( - run.state.log, - vec!["dispatch", "a", "b", "join", "dispatch", "a", "b", "join"] - ); - - // Visited order is deterministic: parallel branches fold in active-set index - // order (work_a before work_b) regardless of completion order. - let visited: Vec<&str> = run.visited.iter().map(|n| n.as_str()).collect(); - assert_eq!( - visited, - vec![ - "dispatch", "work_a", "work_b", "join", "dispatch", "work_a", "work_b", "join", - ] - ); - - // dispatch(1) -> {work_a,work_b}(2) -> join(3) -> dispatch(4) -> ... -> join(6). - assert_eq!(run.steps, 6); -} - -#[tokio::test] -async fn unbounded_loop_back_hits_recursion_limit_deterministically() { - // `max_rounds` larger than the recursion limit allows means the router never - // routes to END; the run must terminate on the deterministic ceiling. - let graph = complex_graph(1_000, 4); - let err = graph - .run(Pipe::default()) - .await - .expect_err("the loop never reaches `done`"); - assert!( - matches!(err, TinyAgentsError::RecursionLimit(4)), - "expected RecursionLimit(4), got {err:?}" - ); -} diff --git a/crates/tinyagents-integration-tests/tests/e2e_control_and_steer.rs b/crates/tinyagents-integration-tests/tests/e2e_control_and_steer.rs index bba35b24c..bde8ab9cd 100644 --- a/crates/tinyagents-integration-tests/tests/e2e_control_and_steer.rs +++ b/crates/tinyagents-integration-tests/tests/e2e_control_and_steer.rs @@ -6,7 +6,6 @@ use std::sync::Arc; use async_trait::async_trait; use serde_json::json; -use tinyagents_graph::*; use tinyagents_harness::context::{MiddlewareControl, RunConfig, RunContext}; use tinyagents_harness::events::AgentEvent; use tinyagents_harness::middleware::Middleware; @@ -112,7 +111,12 @@ impl Middleware<(), ()> for TwoControlMiddleware { async fn higher_precedence_control_wins_and_is_journaled() { for interrupt_first in [true, false] { let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", Arc::new(MockModel::constant("hi"))); + harness.with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named( + "mock", + Arc::new(MockModel::constant("hi")), + ), + )); harness.push_middleware(Arc::new(TwoControlMiddleware { interrupt_first })); let recorder = EventRecorder::new(); @@ -150,10 +154,12 @@ async fn stop_with_final_control_ends_run_before_tools() { // The model asks for a tool it would normally execute and loop on, but the // control middleware stops the run at the checkpoint after the model call. let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( - "mock", - Arc::new(MockModel::with_tool_call("lookup", json!({}))), - ); + harness.with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named( + "mock", + Arc::new(MockModel::with_tool_call("lookup", json!({}))), + ), + )); harness.register_tool(Arc::new(FakeTool::returning("lookup", "out"))); harness.push_middleware(Arc::new(StopControlMiddleware)); @@ -191,7 +197,12 @@ async fn stop_with_final_control_ends_run_before_tools() { #[tokio::test] async fn interrupt_control_surfaces_as_interrupted_error() { let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", Arc::new(MockModel::constant("hi"))); + harness.with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named( + "mock", + Arc::new(MockModel::constant("hi")), + ), + )); harness.push_middleware(Arc::new(InterruptControlMiddleware)); let err = harness @@ -211,7 +222,12 @@ async fn interrupt_control_surfaces_as_interrupted_error() { #[tokio::test] async fn preloaded_cancel_steering_cancels_a_real_run() { let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", Arc::new(MockModel::constant("done"))); + harness.with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named( + "mock", + Arc::new(MockModel::constant("done")), + ), + )); // Attach a steering handle already carrying a Cancel command; the agent loop // drains it at the checkpoint before the first model call and ends the run. diff --git a/crates/tinyagents-integration-tests/tests/e2e_durable_interrupt.rs b/crates/tinyagents-integration-tests/tests/e2e_durable_interrupt.rs deleted file mode 100644 index ee4689efa..000000000 --- a/crates/tinyagents-integration-tests/tests/e2e_durable_interrupt.rs +++ /dev/null @@ -1,110 +0,0 @@ -//! TRUE end-to-end: a durable [`CompiledGraph`] with reducer-merged partial -//! updates, an [`InMemoryCheckpointer`], and a human-in-the-loop -//! interrupt/resume round trip spanning a thread. -//! -//! This composes the **durable graph builder/executor**, the **reducer** -//! (partial-update merge), the **checkpointer** (durability), and the -//! **command/interrupt** model. An early node does work, a middle node -//! interrupts for approval, and a later node performs more work after resume. -//! We assert the run pauses, a checkpoint is persisted, and the resumed final -//! state reflects updates from *both* sides of the interrupt. - -use std::sync::Arc; - -use serde_json::json; - -use tinyagents_graph::ClosureStateReducer; -use tinyagents_graph::*; -use tinyagents_harness::*; -use tinyagents_registry::*; - -/// Audit log of the actions applied through the reducer. Each node contributes -/// a partial `String` update merged onto the running list. -#[derive(Clone, Debug, Default, PartialEq, Eq)] -struct Workflow { - actions: Vec<String>, -} - -#[tokio::test] -async fn durable_graph_interrupts_then_resumes_with_full_state() { - let checkpointer = Arc::new(InMemoryCheckpointer::<Workflow>::new()); - - let graph = GraphBuilder::<Workflow, String>::new() - .set_reducer(ClosureStateReducer::new( - |mut state: Workflow, update: String| { - state.actions.push(update); - Ok(state) - }, - )) - // Early node: real pre-interrupt work. - .add_node("prep", |_s: Workflow, _c: NodeContext| async move { - Ok(NodeResult::Update("prep".to_string())) - }) - // Middle node: pause for human approval, then apply the approval. - .add_node("gate", |_s: Workflow, ctx: NodeContext| async move { - match ctx.resume { - Some(value) => { - let who = value - .get("approved_by") - .and_then(|v| v.as_str()) - .unwrap_or("unknown"); - Ok(NodeResult::Update(format!("approved:{who}"))) - } - None => Ok(NodeResult::Interrupt(Interrupt::new( - "gate", - json!({ "ask": "approve this workflow?" }), - ))), - } - }) - // Later node: post-interrupt work. - .add_node("finalize", |_s: Workflow, _c: NodeContext| async move { - Ok(NodeResult::Update("finalize".to_string())) - }) - .set_entry("prep") - .add_edge("prep", "gate") - .add_edge("gate", "finalize") - .set_finish("finalize") - .compile() - .expect("graph compiles") - .with_checkpointer(checkpointer.clone()); - - // First pass: prep runs, gate interrupts before finalize. - let paused = graph - .run_with_thread("wf-1", Workflow::default()) - .await - .expect("first pass succeeds"); - - assert!(paused.is_interrupted(), "the run should pause at the gate"); - assert_eq!(paused.interrupts.len(), 1); - assert_eq!(paused.interrupts[0].node.as_str(), "gate"); - // Pre-interrupt work is already reflected in the paused state. - assert_eq!(paused.state.actions, vec!["prep".to_string()]); - - // Durability: at least one checkpoint was persisted for the thread. - let checkpoints = checkpointer.list("wf-1").await.expect("list succeeds"); - assert!( - !checkpoints.is_empty(), - "an interrupt must persist a checkpoint" - ); - - // Resume: inject the approval, gate completes, finalize runs. - let resumed = graph - .resume("wf-1", Command::resume(json!({ "approved_by": "ada" }))) - .await - .expect("resume succeeds"); - - assert!(!resumed.is_interrupted(), "resumed run must finish"); - // Final state reflects BOTH pre- and post-interrupt updates. - assert_eq!( - resumed.state.actions, - vec![ - "prep".to_string(), - "approved:ada".to_string(), - "finalize".to_string(), - ] - ); - assert!( - resumed.steps >= 1, - "the resumed pass advanced at least once" - ); -} diff --git a/crates/tinyagents-integration-tests/tests/e2e_embeddings.rs b/crates/tinyagents-integration-tests/tests/e2e_embeddings.rs index 0078a6a70..f0ed614be 100644 --- a/crates/tinyagents-integration-tests/tests/e2e_embeddings.rs +++ b/crates/tinyagents-integration-tests/tests/e2e_embeddings.rs @@ -11,7 +11,6 @@ use std::sync::Arc; use serde_json::json; -use tinyagents_graph::*; use tinyagents_harness::*; use tinyagents_registry::*; use tinyinference_embeddings::{InMemoryVectorStore, MockEmbeddingModel, Retriever}; diff --git a/crates/tinyagents-integration-tests/tests/e2e_graph_support_contracts.rs b/crates/tinyagents-integration-tests/tests/e2e_graph_support_contracts.rs deleted file mode 100644 index 749e6f26d..000000000 --- a/crates/tinyagents-integration-tests/tests/e2e_graph_support_contracts.rs +++ /dev/null @@ -1,290 +0,0 @@ -//! End-to-end contracts for graph support surfaces: file-backed checkpoints and -//! the public graph testkit helpers. - -use std::time::{SystemTime, UNIX_EPOCH}; - -use serde_json::json; - -use tinyagents_graph::testkit::{ - GraphEventRecorder, RetryCountingNode, assert_graph, fanout_node, interrupting_node, noop_node, - run_recorded, scripted_route_node, scripted_update_node, subagent_fake_node, - subgraph_test_node, -}; -use tinyagents_graph::*; -use tinyagents_harness::ids::ExecutionStatus; -mod common; -use common::storage; -use tinyagents_harness::*; -use tinyagents_registry::*; -use tinyinference_llm::usage::UsageTotals; - -#[tokio::test] -async fn driver_checkpointer_persists_lists_copies_prunes_and_deletes_threads() { - let dir = tempfile::tempdir().unwrap(); - let checkpointer = std::sync::Arc::new(storage::checkpointer::<i32>(dir.path())); - - let graph = GraphBuilder::<i32, i32>::overwrite() - .add_node("a", |s, _c: NodeContext| async move { - Ok(NodeResult::Update(s + 1)) - }) - .add_node("b", |s, _c: NodeContext| async move { - Ok(NodeResult::Update(s + 1)) - }) - .set_entry("a") - .add_edge("a", "b") - .set_finish("b") - .compile() - .expect("graph compiles") - .with_checkpointer(checkpointer.clone()); - - let run = graph - .run_with_thread("thread/with space", 0) - .await - .expect("threaded file-backed run succeeds"); - assert_eq!(run.state, 2); - assert_eq!(run.status.status, ExecutionStatus::Completed); - - let list = checkpointer.list("thread/with space").await.expect("list"); - assert_eq!(list.len(), 2); - assert_eq!(list[0].source, CheckpointSource::Loop); - assert_eq!( - CheckpointSource::parse("fork"), - Some(CheckpointSource::Fork) - ); - assert_eq!(CheckpointSource::Update.to_string(), "update"); - - let latest = checkpointer - .get("thread/with space", None) - .await - .expect("get latest") - .expect("latest checkpoint"); - assert_eq!(latest.state, 2); - assert_eq!( - latest.parent_checkpoint_id.as_deref(), - Some(list[0].checkpoint_id.as_str()) - ); - - let tuple = checkpointer - .get_tuple(CheckpointConfig::latest("thread/with space")) - .await - .expect("tuple") - .expect("tuple present"); - assert_eq!(tuple.config.thread_id, "thread/with space"); - assert!(tuple.parent_config.is_some()); - - checkpointer - .copy_thread("thread/with space", "copy-target") - .await - .expect("copy thread"); - let copied = checkpointer.list("copy-target").await.expect("copied list"); - assert_eq!(copied.len(), 2); - assert!(copied.iter().all(|m| m.thread_id == "copy-target")); - - let threads = checkpointer.list_threads().await.expect("threads"); - assert!(threads.contains(&"thread/with space".to_string())); - assert!(threads.contains(&"copy-target".to_string())); - - let removed = checkpointer - .delete_checkpoints("copy-target", &[copied[0].checkpoint_id.clone()]) - .await - .expect("delete one"); - assert_eq!(removed, 1); - assert_eq!(checkpointer.list("copy-target").await.unwrap().len(), 1); - - let pruned = checkpointer - .prune("thread/with space", 1) - .await - .expect("prune keeps latest and ancestors"); - assert_eq!(pruned, 0); - - checkpointer - .delete_thread("copy-target") - .await - .expect("delete thread"); - assert!(checkpointer.list("copy-target").await.unwrap().is_empty()); -} - -#[tokio::test] -async fn driver_checkpointer_state_history_walks_lineage_newest_first() { - // Walks the parent lineage and returns snapshots newest-first with the - // lineage spine intact. - let dir = tempfile::tempdir().unwrap(); - let checkpointer = std::sync::Arc::new(storage::checkpointer::<i32>(dir.path())); - let graph = GraphBuilder::<i32, i32>::overwrite() - .add_node("a", |s, _c: NodeContext| async move { - Ok(NodeResult::Update(s + 1)) - }) - .add_node("b", |s, _c: NodeContext| async move { - Ok(NodeResult::Update(s + 1)) - }) - .set_entry("a") - .add_edge("a", "b") - .set_finish("b") - .compile() - .expect("graph compiles") - .with_checkpointer(checkpointer.clone()); - - graph - .run_with_thread("hist-thread", 0) - .await - .expect("threaded run succeeds"); - - // Full history is newest-first, and each snapshot's parent addresses the - // next (older) snapshot — the lineage spine walked in memory. - let history = graph - .get_state_history("hist-thread", None) - .await - .expect("state history"); - assert_eq!(history.len(), 2); - assert_eq!(history[0].values, 2); - assert_eq!(history[1].values, 1); - assert_eq!( - history[0] - .parent_config - .as_ref() - .and_then(|c| c.checkpoint_id.as_deref()), - history[1].config.checkpoint_id.as_deref(), - "newest checkpoint's parent is the older checkpoint" - ); - - // `limit` caps to the most recent snapshots. - let recent = graph - .get_state_history("hist-thread", Some(1)) - .await - .expect("limited history"); - assert_eq!(recent.len(), 1); - assert_eq!(recent[0].values, 2); -} - -#[tokio::test] -async fn graph_testkit_helpers_drive_recorded_graphs_and_assertions() { - let checkpointer = std::sync::Arc::new(InMemoryCheckpointer::<i32>::new()); - let graph = GraphBuilder::<i32, i32>::overwrite() - .add_node("seed", scripted_update_node([1])) - .add_node("route", scripted_route_node([["fanout"]])) - .add_node("fanout", fanout_node("worker", [json!(2), json!(3)])) - .add_node("worker", |s: i32, ctx: NodeContext| async move { - let arg = ctx.send_arg.and_then(|v| v.as_i64()).unwrap_or(0) as i32; - Ok(NodeResult::Update(s + arg)) - }) - .add_node("agent", subagent_fake_node("fake", 10, UsageTotals::new())) - .add_node("noop", noop_node()) - .set_entry("seed") - .add_edge("seed", "route") - .mark_command_routing("route") - .mark_command_routing("fanout") - .add_waiting_edge("worker", "agent") - .add_edge("agent", "noop") - .set_finish("noop") - .compile() - .expect("graph compiles") - .with_checkpointer(checkpointer); - - let recorded = run_recorded(&graph, Some("graph-testkit"), 0) - .await - .expect("recorded run"); - assert_eq!(recorded.execution.state, 10); - assert_eq!(recorded.execution.child_runs.len(), 1); - assert_graph(&recorded) - .visited([ - "seed", "route", "fanout", "worker", "worker", "agent", "noop", - ]) - .routed("seed", "route") - .checkpoint_count(6) - .state_history(|history| assert_eq!(history.len(), 6)) - .checkpoint(|latest| { - assert_eq!(latest.values, 10); - assert!(latest.next_nodes.is_empty()); - }) - .completed(); - - let collector = recorded.collector(); - assert!(collector.events().len() >= 5); - assert!( - collector - .node_order() - .iter() - .any(|n| n.as_str() == "worker") - ); - assert!(collector.updates().iter().any(|n| n.as_str() == "agent")); - assert!( - collector - .routes() - .iter() - .any(|(from, to)| { from.as_str() == "seed" && to.as_str() == "route" }) - ); - assert_eq!(collector.checkpoint_count(), 6); - assert!(collector.interrupts().is_empty()); - assert!(collector.custom().is_empty()); - - let manual_recorder = GraphEventRecorder::new(); - let manual_graph = GraphBuilder::<i32, i32>::overwrite() - .add_node("one", scripted_update_node([1])) - .set_entry("one") - .set_finish("one") - .compile() - .unwrap() - .with_event_sink(manual_recorder.sink()); - manual_graph.run(0).await.unwrap(); - assert!( - manual_recorder - .kinds() - .contains(&"node.completed".to_string()) - ); - assert!(!manual_recorder.collector().events().is_empty()); -} - -#[tokio::test] -async fn graph_testkit_interrupt_retry_subgraph_and_failure_paths_are_public_contracts() { - let retry = RetryCountingNode::new(1); - let retry_graph = GraphBuilder::<i32, i32>::overwrite() - .add_node("retry", retry.handler(7)) - .set_entry("retry") - .set_finish("retry") - .compile() - .unwrap(); - let err = retry_graph.run(0).await.unwrap_err(); - assert!(matches!(err, TinyAgentsError::Graph(_))); - assert_eq!(retry.attempts(), 1); - - let retry_success = GraphBuilder::<i32, i32>::overwrite() - .add_node("retry", retry.handler(7)) - .set_entry("retry") - .set_finish("retry") - .compile() - .unwrap(); - assert_eq!(retry_success.run(0).await.unwrap().state, 7); - assert_eq!(retry.attempts(), 2); - - let child = GraphBuilder::<i32, i32>::overwrite() - .add_node("add", |s, _c: NodeContext| async move { - Ok(NodeResult::Update(s + 5)) - }) - .set_entry("add") - .set_finish("add") - .compile() - .unwrap(); - let parent = GraphBuilder::<i32, i32>::overwrite() - .add_node("child", subgraph_test_node(child)) - .set_entry("child") - .set_finish("child") - .compile() - .unwrap(); - assert_eq!(parent.run(1).await.unwrap().state, 6); - - let interrupting = GraphBuilder::<i32, i32>::overwrite() - .add_node("hitl", interrupting_node(json!({ "ask": true }), 11)) - .set_entry("hitl") - .set_finish("hitl") - .compile() - .unwrap() - .with_checkpointer(std::sync::Arc::new(InMemoryCheckpointer::<i32>::new())) - .with_durability(DurabilityMode::Exit); - let paused = interrupting.run_with_thread("hitl", 0).await.unwrap(); - assert_graph(&tinyagents_graph::testkit::GraphRun::new(paused)).interrupted(); - let resumed = interrupting - .resume("hitl", Command::resume(json!(true))) - .await - .unwrap(); - assert_eq!(resumed.state, 11); -} diff --git a/crates/tinyagents-integration-tests/tests/e2e_graph_todos.rs b/crates/tinyagents-integration-tests/tests/e2e_graph_todos.rs index 9458b7127..f60b17fe2 100644 --- a/crates/tinyagents-integration-tests/tests/e2e_graph_todos.rs +++ b/crates/tinyagents-integration-tests/tests/e2e_graph_todos.rs @@ -34,31 +34,32 @@ async fn model_drives_the_todo_tool_and_list_persists_to_the_thread() { let mut harness: AgentHarness<()> = AgentHarness::new(); harness - .register_model( - "mock", - Arc::new(MockModel::with_responses(vec![ - // First turn: write the plan via the `todo` tool. - tool_call_response( - "call-1", - "todo", - json!({ "todos": [ + .with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named( + "mock", + Arc::new(MockModel::with_responses(vec![ + // First turn: write the plan via the `todo` tool. + tool_call_response( + "call-1", + "todo", + json!({ "todos": [ { "content": "Write the integration test", "status": "in_progress" }, { "content": "Run it", "status": "pending" } ] }), - ), - // Second turn: tick the first one off and start the next. - tool_call_response( - "call-2", - "todo", - json!({ "todos": [ + ), + // Second turn: tick the first one off and start the next. + tool_call_response( + "call-2", + "todo", + json!({ "todos": [ { "content": "Write the integration test", "status": "completed" }, { "content": "Run it", "status": "in_progress" } ] }), - ), - text_response("done"), - ])), - ) - .set_default_model("mock") + ), + text_response("done"), + ])), + ), + )) .register_tool(Arc::new(TodoTool::new(store.clone()))); let run = harness diff --git a/crates/tinyagents-integration-tests/tests/e2e_harness_mock_capabilities.rs b/crates/tinyagents-integration-tests/tests/e2e_harness_mock_capabilities.rs index 8505d37a1..c7b612018 100644 --- a/crates/tinyagents-integration-tests/tests/e2e_harness_mock_capabilities.rs +++ b/crates/tinyagents-integration-tests/tests/e2e_harness_mock_capabilities.rs @@ -14,9 +14,6 @@ use tinyagents_harness::context::{RunConfig, RunContext}; use tinyagents_harness::middleware::{LoggingMiddleware, MicrocompactMiddleware, Middleware}; use tinyagents_harness::runtime::{AgentHarness, RunPolicy}; use tinyagents_harness::testkit::{FakeTool, ScriptedModel}; -use tinyagents_orchestration::subagent::{ - ChildDataPolicy, SubAgent, SubAgentSession, SubAgentTool, -}; use tinyinference_llm::message::{AssistantMessage, ContentBlock, Message}; use tinyinference_llm::model::{ModelRequest, ModelResponse}; use tinyinference_llm::tool::ToolCall; @@ -83,8 +80,9 @@ async fn mock_provider_reuses_a_cached_single_turn_response() { let cache = Arc::new(InMemoryResponseCache::new()); let mut harness: AgentHarness<()> = AgentHarness::new(); harness - .register_model("mock", model.clone()) - .set_default_model("mock") + .with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named("mock", model.clone()), + )) .with_response_cache(cache.clone()); let input = vec![Message::system("Be concise."), Message::user("hello")]; @@ -121,8 +119,9 @@ async fn mock_provider_receives_one_stable_prompt_prefix_key_across_a_tool_loop( policy.cache.protect_prompt_prefix = true; let mut harness: AgentHarness<()> = AgentHarness::new(); harness - .register_model("mock", model.clone()) - .set_default_model("mock") + .with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named("mock", model.clone()), + )) .register_tool(Arc::new(FakeTool::returning("lookup", "result"))) .with_policy(policy); @@ -169,15 +168,16 @@ async fn tool_loop_compacts_old_results_and_runs_all_lifecycle_hooks() { let hooks = Arc::new(LoggingMiddleware::new()); let mut harness: AgentHarness<()> = AgentHarness::new(); harness - .register_model( - "mock", - Arc::new(ScriptedModel::new(vec![ - tool_turn(vec![ToolCall::new("one", "lookup", json!({"n": 1}))]), - tool_turn(vec![ToolCall::new("two", "lookup", json!({"n": 2}))]), - text_turn("complete"), - ])), - ) - .set_default_model("mock") + .with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named( + "mock", + Arc::new(ScriptedModel::new(vec![ + tool_turn(vec![ToolCall::new("one", "lookup", json!({"n": 1}))]), + tool_turn(vec![ToolCall::new("two", "lookup", json!({"n": 2}))]), + text_turn("complete"), + ])), + ), + )) .register_tool(Arc::new(FakeTool::returning("lookup", "large result"))) .push_middleware(Arc::new(MicrocompactMiddleware::new(1, "[compacted]"))) .push_middleware(probe.clone()) @@ -206,96 +206,42 @@ async fn tool_loop_compacts_old_results_and_runs_all_lifecycle_hooks() { assert_eq!(counts.on_error, 0); } +/// This deterministic matrix exercises both the no-tool and regular-tool paths. #[tokio::test] -async fn subagent_session_keeps_context_across_multiple_turns() { - let child_model = Arc::new(ScriptedModel::replies(vec![ - "first answer", - "second answer", - ])); - let mut child_harness: AgentHarness<()> = AgentHarness::new(); - child_harness.register_model("child", child_model.clone()); - let child = Arc::new(SubAgent::new( - "researcher", - "researches", - Arc::new(child_harness), - )); - - let mut session = SubAgentSession::new(child); - session - .send(&(), (), vec![Message::user("first question")]) - .await - .expect("first child turn"); - let second = session - .send(&(), (), vec![Message::user("follow up")]) - .await - .expect("reused child turn"); - - assert_eq!(session.turns(), 2); - assert_eq!(second.text(), Some("second answer".to_string())); - let second_request = &child_model.requests()[1]; - let texts: Vec<String> = second_request.messages.iter().map(Message::text).collect(); - assert!(texts.contains(&"first question".to_string())); - assert!(texts.contains(&"first answer".to_string())); - assert!(texts.contains(&"follow up".to_string())); -} - -/// This is intentionally deterministic rather than random: every combination -/// is reproducible, prints its flags on failure, and still exercises the same -/// branching surface a seed-based fuzz test would explore. -#[tokio::test] -async fn mock_fuzz_matrix_composes_regular_and_subagent_tools() { +async fn mock_matrix_composes_regular_tools() { for use_lookup in [false, true] { - for use_subagent in [false, true] { - let label = format!("lookup={use_lookup}, subagent={use_subagent}"); - let mut harness: AgentHarness<()> = AgentHarness::new(); - let mut calls = Vec::new(); - if use_lookup { - harness.register_tool(Arc::new(FakeTool::returning("lookup", "lookup result"))); - calls.push(ToolCall::new("lookup-call", "lookup", json!({"q": "x"}))); - } - if use_subagent { - let mut child_harness: AgentHarness<()> = AgentHarness::new(); - child_harness.register_model( - "child", - Arc::new(ScriptedModel::replies(vec!["child result"])), - ); - let child = Arc::new(SubAgent::new( - "delegate", - "delegates", - Arc::new(child_harness), - )); - harness.register_tool_dispatch(Arc::new(SubAgentTool::new( - child, - ChildDataPolicy::new(|state: &()| *state), - ))); - calls.push(ToolCall::new( - "child-call", - "delegate", - json!({"input": "work"}), - )); - } - let responses = if calls.is_empty() { - vec![text_turn("final")] - } else { - vec![tool_turn(calls), text_turn("final")] - }; - harness - .register_model("parent", Arc::new(ScriptedModel::new(responses))) - .set_default_model("parent") - .with_policy(RunPolicy::default()); - - let run = harness - .invoke_default(&(), vec![Message::user(format!("run {label}"))]) - .await - .unwrap_or_else(|error| panic!("scenario {label} failed: {error:?}")); - let expected_tools = usize::from(use_lookup) + usize::from(use_subagent); - assert_eq!(run.tool_calls, expected_tools, "scenario {label}"); - assert_eq!( - run.model_calls, - if expected_tools == 0 { 1 } else { 2 }, - "scenario {label}" - ); - assert_eq!(run.text(), Some("final".to_string()), "scenario {label}"); + let label = format!("lookup={use_lookup}"); + let mut harness: AgentHarness<()> = AgentHarness::new(); + let mut calls = Vec::new(); + if use_lookup { + harness.register_tool(Arc::new(FakeTool::returning("lookup", "lookup result"))); + calls.push(ToolCall::new("lookup-call", "lookup", json!({"q": "x"}))); } + let responses = if calls.is_empty() { + vec![text_turn("final")] + } else { + vec![tool_turn(calls), text_turn("final")] + }; + harness + .with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named( + "parent", + Arc::new(ScriptedModel::new(responses)), + ), + )) + .with_policy(RunPolicy::default()); + + let run = harness + .invoke_default(&(), vec![Message::user(format!("run {label}"))]) + .await + .unwrap_or_else(|error| panic!("scenario {label} failed: {error:?}")); + let expected_tools = usize::from(use_lookup); + assert_eq!(run.tool_calls, expected_tools, "scenario {label}"); + assert_eq!( + run.model_calls, + if expected_tools == 0 { 1 } else { 2 }, + "scenario {label}" + ); + assert_eq!(run.text(), Some("final".to_string()), "scenario {label}"); } } diff --git a/crates/tinyagents-integration-tests/tests/e2e_harness_provider_contracts.rs b/crates/tinyagents-integration-tests/tests/e2e_harness_provider_contracts.rs index ba46770e8..124baf96d 100644 --- a/crates/tinyagents-integration-tests/tests/e2e_harness_provider_contracts.rs +++ b/crates/tinyagents-integration-tests/tests/e2e_harness_provider_contracts.rs @@ -4,8 +4,6 @@ use futures::stream; use serde::Deserialize; use serde_json::json; use tinyagents_harness::cost::{CostTotals, ModelPricing, estimate_cost}; -use tinyagents_harness::model_registry::{ModelRegistry, ModelSelection}; -use tinyagents_harness::stream::{StreamChunk, StreamMode, StreamSink, stream as filter_stream}; use tinyagents_harness::structured::{ StructuredExtractor, StructuredStrategy, response_format_for_strategy, }; @@ -182,51 +180,15 @@ async fn model_request_response_registry_and_stream_contracts_are_stable() { assert_eq!(response.message.usage.unwrap().effective_total(), 5); assert_eq!(response.resolved_model, Some(resolved)); - let mut registry: ModelRegistry<()> = ModelRegistry::new(); - registry - .register("default", Arc::new(MockModel::constant("default"))) - .register("direct", Arc::new(MockModel::constant("direct"))) - .register("hinted", Arc::new(MockModel::constant("hinted"))); - registry.set_default("default"); - - let binding = registry - .resolve_request(&request, Some("default"), None) - .unwrap(); - assert_eq!(binding.resolved.name, "direct"); + assert_eq!(request.model.as_deref(), Some("direct")); assert_eq!( - binding.resolved.source, - ModelResolutionSource::RequestOverride - ); - assert_eq!(registry.names(), vec!["default", "direct", "hinted"]); - - let fallback = registry - .resolve(ModelSelection { - hints: vec![ - ModelHint { - model: "missing".into(), - priority: 100, - reason: None, - }, - ModelHint { - model: "hinted".into(), - priority: 1, - reason: None, - }, - ], - ..ModelSelection::default() - }) - .unwrap(); - assert_eq!(fallback.resolved.name, "hinted"); - assert_eq!(fallback.resolved.source, ModelResolutionSource::Hint); - - let unsatisfied = registry.resolve(ModelSelection { - required_capabilities: Some(CapabilitySet { - min_input_tokens: Some(u64::MAX), + request.required_capabilities, + Some(CapabilitySet { + tool_calling: true, + streaming: true, ..CapabilitySet::default() - }), - ..ModelSelection::default() - }); - assert!(unsatisfied.is_none()); + }) + ); let mut accumulator = StreamAccumulator::new(); accumulator.push(&ModelStreamItem::Started); @@ -319,44 +281,6 @@ fn model_profiles_stream_chunks_usage_and_cost_are_additive() { ..CapabilitySet::default() })); - let chunks = vec![ - StreamChunk::Values(json!({ "state": 1 })), - StreamChunk::Updates(json!({ "state": 2 })), - StreamChunk::Message(MessageDelta { - text: "delta".into(), - reasoning: String::new(), - tool_call: Some(ToolDelta { - call_id: "tool-1".into(), - content: "partial".into(), - tool_name: None, - ..Default::default() - }), - }), - StreamChunk::Debug("trace".into()), - StreamChunk::Interrupt(json!({ "kind": "input" })), - StreamChunk::Custom(json!({ "x": true })), - ]; - assert_eq!(chunks[0].mode(), StreamMode::Values); - assert_eq!( - filter_stream(&chunks, &[StreamMode::Messages, StreamMode::Custom]), - vec![chunks[2].clone(), chunks[5].clone()] - ); - - let mut sink = StreamSink::new([StreamMode::Messages]); - assert!(sink.is_active(StreamMode::Messages)); - assert!(!sink.is_active(StreamMode::Debug)); - sink.push(chunks[2].clone()); - sink.push(chunks[3].clone()); - assert_eq!(sink.len(), 1); - sink.enable(StreamMode::Debug); - sink.push(chunks[3].clone()); - sink.disable(StreamMode::Messages); - sink.push(chunks[2].clone()); - assert_eq!(sink.peek(), vec![chunks[2].clone(), chunks[3].clone()]); - assert_eq!(sink.drain().len(), 2); - assert!(sink.is_empty()); - assert_eq!(StreamSink::all().active_modes().len(), 6); - let mut usage = Usage::new(10, 5); usage.cache_read_tokens = 3; usage.cache_creation_tokens = 2; diff --git a/crates/tinyagents-integration-tests/tests/e2e_middleware.rs b/crates/tinyagents-integration-tests/tests/e2e_middleware.rs index 088dd0223..8ba573681 100644 --- a/crates/tinyagents-integration-tests/tests/e2e_middleware.rs +++ b/crates/tinyagents-integration-tests/tests/e2e_middleware.rs @@ -211,8 +211,9 @@ async fn stacked_middleware_recovers_redacts_and_traces() { let mut harness: AgentHarness<()> = AgentHarness::new(); harness - .register_model("flaky", Arc::new(model)) - .set_default_model("flaky") + .with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named("flaky", Arc::new(model)), + )) .register_tool(Arc::new(FakeTool::returning("search", "tool-output"))) .push_model_middleware(Arc::new(RetryMiddleware::new(RetryPolicy::default()))) .push_middleware(Arc::new(ToolAllowlistMiddleware::new(["search"]))) @@ -296,8 +297,9 @@ async fn stacked_middleware_rejects_disallowed_tool() { let mut harness: AgentHarness<()> = AgentHarness::new(); harness - .register_model("scripted", Arc::new(scripted)) - .set_default_model("scripted") + .with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named("scripted", Arc::new(scripted)), + )) // Register the tool so the *only* reason for failure is the allowlist, // not a missing-tool error. .register_tool(danger.clone()) @@ -347,8 +349,9 @@ async fn invalid_tool_arguments_are_rejected_before_execution() { let mut harness: AgentHarness<()> = AgentHarness::new(); harness - .register_model("scripted", Arc::new(scripted)) - .set_default_model("scripted") + .with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named("scripted", Arc::new(scripted)), + )) .register_tool(Arc::new(StrictLookupTool { calls: Arc::clone(&calls), })); diff --git a/crates/tinyagents-integration-tests/tests/e2e_misc_public_helpers.rs b/crates/tinyagents-integration-tests/tests/e2e_misc_public_helpers.rs index cc015f63f..27425b90a 100644 --- a/crates/tinyagents-integration-tests/tests/e2e_misc_public_helpers.rs +++ b/crates/tinyagents-integration-tests/tests/e2e_misc_public_helpers.rs @@ -1,268 +1,12 @@ -use std::sync::Arc; use std::thread; use std::time::Duration; use serde_json::json; -use tinyagents_graph::observability::{ - GraphEventJournal, GraphStatusStore, InMemoryGraphEventJournal, InMemoryGraphStatusStore, - JournalGraphSink, StoreGraphEventJournal, -}; -use tinyagents_graph::reducer::{ - AppendReducer, ClosureReducer, ClosureStateReducer, MaxReducer, MinReducer, OverwriteReducer, - OverwriteStateReducer, Reducer, SetUnionReducer, StateReducer, -}; -use tinyagents_graph::status::GraphRunStatus; -use tinyagents_graph::stream::{ - CollectingSink, GraphEvent, GraphEventEnvelope, GraphEventSink, NoopSink, StreamMode, -}; -use tinyagents_graph::{Command, Interrupt, NodeResult, Send}; -use tinyagents_harness::ids::{ - CheckpointId, ExecutionStatus, GraphId, InterruptId, NodeId, RunId, ThreadId, new_call_id, - new_cell_id, new_session_id, next_seq, -}; use tinyagents_harness::limits::{LimitTracker, RunLimits}; -use tinyagents_harness::store::{AppendStore, InMemoryAppendStore}; +use tinyagents_harness::{new_call_id, new_cell_id, new_session_id, next_seq}; use tinyinference_llm::tool::{ToolCall, ToolFormat, ToolSchema}; use tinytools::ToolResult; -#[tokio::test] -async fn graph_reducers_streams_observability_and_status_helpers_work() { - assert_eq!(OverwriteReducer.reduce(1, 2).unwrap(), 2); - assert_eq!( - AppendReducer.reduce(vec![1], vec![2, 3]).unwrap(), - vec![1, 2, 3] - ); - assert_eq!( - SetUnionReducer.reduce(vec![1, 2], vec![2, 3]).unwrap(), - vec![1, 2, 3] - ); - assert_eq!(MinReducer.reduce(10, 2).unwrap(), 2); - assert_eq!(MaxReducer.reduce(10, 20).unwrap(), 20); - assert_eq!( - ClosureReducer::new(|a: i32, b: i32| Ok(a * b)) - .reduce(3, 4) - .unwrap(), - 12 - ); - assert_eq!(OverwriteStateReducer.apply("old", "new").unwrap(), "new"); - assert_eq!( - ClosureStateReducer::new(|state: i32, update: i32| Ok(state + update)) - .apply(2, 5) - .unwrap(), - 7 - ); - - let send = Send { - node: NodeId::new("worker"), - arg: json!({ "item": 1 }), - }; - let command = Command::send([send.clone()]) - .with_sends([send]) - .with_goto(["next"]) - .with_update(3) - .with_resume(json!({ "resume": true })); - assert_eq!(command.update, Some(3)); - assert_eq!(command.goto.len(), 3); - assert_eq!(command.resume.as_ref().unwrap()["resume"], true); - assert_eq!(Command::<i32>::update(9).update, Some(9)); - assert_eq!( - Command::<i32>::resume(json!("ok")).resume, - Some(json!("ok")) - ); - let interrupt = Interrupt::new("node-a", json!({ "question": "continue?" })); - assert!(interrupt.id.starts_with("interrupt-node-a-")); - let explicit = Interrupt::with_id("int-1", "node-b", json!(1)); - assert_eq!(explicit.id, "int-1"); - - let events = vec![ - GraphEvent::RunStarted { - run_id: RunId::new("run"), - }, - GraphEvent::RunCompleted { - run_id: RunId::new("run"), - steps: 2, - }, - GraphEvent::RunFailed { - run_id: RunId::new("run"), - error: "bad".into(), - }, - GraphEvent::StepStarted { - step: 1, - active: vec![NodeId::new("a")], - }, - GraphEvent::StepCompleted { step: 1 }, - GraphEvent::TaskScheduled { - node: NodeId::new("a"), - step: 1, - }, - GraphEvent::NodeStarted { - node: NodeId::new("a"), - step: 1, - }, - GraphEvent::NodeCompleted { - node: NodeId::new("a"), - step: 1, - }, - GraphEvent::NodeFailed { - node: NodeId::new("a"), - step: 1, - error: "bad".into(), - }, - GraphEvent::StateUpdated { - node: NodeId::new("a"), - step: 1, - }, - GraphEvent::RouteSelected { - node: NodeId::new("a"), - target: NodeId::new("b"), - }, - GraphEvent::CheckpointSaved { - checkpoint_id: CheckpointId::new("cp-1"), - step: Some(1), - }, - GraphEvent::InterruptEmitted { - interrupt: explicit.clone(), - }, - GraphEvent::SubgraphStarted { - node: NodeId::new("a"), - namespace: vec!["child".into()], - }, - GraphEvent::SubgraphCompleted { - node: NodeId::new("a"), - namespace: vec!["child".into()], - }, - GraphEvent::ContextForked { - node: NodeId::new("a"), - fork: 0, - step: 1, - }, - GraphEvent::RecursionDepthChanged { depth: 2 }, - GraphEvent::Custom { - name: "custom".into(), - data: json!({ "x": 1 }), - }, - ]; - assert_eq!(events[0].kind(), "run.started"); - assert_eq!(events[3].step(), Some(1)); - assert_eq!(events[1].step(), Some(2)); - assert_eq!(events[10].step(), None); - assert_ne!(StreamMode::Values, StreamMode::Debug); - - let envelope = |event: GraphEvent| GraphEventEnvelope { - run_id: RunId::new("run-g"), - task_id: None, - ns: Vec::new(), - seq: 0, - event, - }; - - let sink = CollectingSink::new(); - assert!(sink.is_empty()); - for event in events.clone() { - sink.emit(envelope(event)); - } - assert_eq!(sink.len(), events.len()); - assert_eq!(sink.events()[0].kind(), "run.started"); - NoopSink.emit(envelope(GraphEvent::Custom { - name: "drop".into(), - data: json!(null), - })); - - let journal = Arc::new(InMemoryGraphEventJournal::new()); - assert!(journal.is_empty("run-g")); - let journal_sink = JournalGraphSink::new( - journal.clone(), - RunId::new("run-g"), - GraphId::new("graph-g"), - ) - .with_lineage(Some(RunId::new("parent")), RunId::new("root")) - .with_thread(Some(ThreadId::new("thread"))) - .with_namespace(vec!["child".into()]) - .with_inner(Arc::new(sink.clone())); - journal_sink.emit(envelope(GraphEvent::StepStarted { - step: 3, - active: vec![NodeId::new("a")], - })); - journal_sink.emit(envelope(GraphEvent::CheckpointSaved { - checkpoint_id: CheckpointId::new("cp-3"), - step: Some(3), - })); - // Persistence is asynchronous; block until the durable log catches up. - journal_sink.flush(); - assert_eq!(journal.len("run-g"), 2); - let observations = journal.read_from("run-g", 0).await.unwrap(); - assert_eq!(observations[0].step, 3); - assert_eq!(observations[1].step, 3); - assert_eq!( - observations[1].checkpoint_id.as_ref().unwrap().as_str(), - "cp-3" - ); - assert_eq!( - observations[0].parent_run_id.as_ref().unwrap().as_str(), - "parent" - ); - assert_eq!( - observations[0].thread_id.as_ref().unwrap().as_str(), - "thread" - ); - assert_eq!(observations[0].namespace, vec!["child"]); - - let append_store = InMemoryAppendStore::new(); - let store_journal = StoreGraphEventJournal::new(append_store.clone()); - assert_eq!( - store_journal.append(observations[0].clone()).await.unwrap(), - 0 - ); - assert_eq!(store_journal.store().len("run-g").await.unwrap(), 1); - assert_eq!( - store_journal.read_from("run-g", 0).await.unwrap()[0] - .event - .kind(), - "step.started" - ); - - let status_store = InMemoryGraphStatusStore::new(); - assert!(status_store.is_empty()); - let mut status = GraphRunStatus::new( - RunId::new("run-status"), - GraphId::new("graph-status"), - ExecutionStatus::Running, - ); - status.thread_id = Some(ThreadId::new("thread")); - status.parent_run_id = Some(RunId::new("parent")); - status.current_step = 4; - status.active_nodes = vec![NodeId::new("a")]; - status.pending_interrupts = vec![InterruptId::new("int")]; - assert!(!status.is_terminal()); - status_store.put_status(status.clone()).await.unwrap(); - assert_eq!(status_store.len(), 1); - assert_eq!( - status_store - .get_status("run-status") - .await - .unwrap() - .unwrap() - .current_step, - 4 - ); - assert_eq!( - status_store.list_by_thread("thread").await.unwrap().len(), - 1 - ); - assert!( - GraphRunStatus::new( - RunId::new("done"), - GraphId::new("g"), - ExecutionStatus::Completed - ) - .is_terminal() - ); - - let _ = NodeResult::<i32>::Update(1); - let _ = NodeResult::<i32>::Command(Command::default()); - let _ = NodeResult::<i32>::Interrupt(interrupt); -} - #[test] fn tool_schema_limits_and_ids_cover_public_helpers() { let schema = ToolSchema::new( diff --git a/crates/tinyagents-integration-tests/tests/e2e_observability.rs b/crates/tinyagents-integration-tests/tests/e2e_observability.rs index 020aeff31..9f764863e 100644 --- a/crates/tinyagents-integration-tests/tests/e2e_observability.rs +++ b/crates/tinyagents-integration-tests/tests/e2e_observability.rs @@ -9,11 +9,6 @@ //! that the run's completion is observable both through a wired //! [`HarnessStatusStore`] and through the journaled events. //! -//! 2. A graph run with a [`JournalGraphSink`] (wired through -//! [`CompiledGraph::with_event_journal`]). The test asserts the -//! [`GraphObservation`]s are replayable from offset `0` and carry their -//! run/step coordinates. -//! //! All assertions are structural (event kinds, offsets, ids, redaction) — never //! model prose — so the test is deterministic. @@ -21,7 +16,6 @@ use std::sync::Arc; use serde_json::json; -use tinyagents_graph::*; use tinyagents_harness::context::{RunConfig, RunContext}; use tinyagents_harness::events::{AgentEvent, RecordingListener}; use tinyagents_harness::ids::{ExecutionStatus, RunId}; @@ -86,14 +80,15 @@ async fn harness_run_journals_redacted_replayable_events() { // varied lifecycle: run/model/tool boundaries plus completion. let mut harness: AgentHarness<()> = AgentHarness::new(); harness - .register_model( - MODEL_NAME, - Arc::new(MockModel::with_responses(vec![ - tool_call_response("call-1", "lookup", json!({ "q": "x" })), - text_response("done"), - ])), - ) - .set_default_model(MODEL_NAME) + .with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named( + MODEL_NAME, + Arc::new(MockModel::with_responses(vec![ + tool_call_response("call-1", "lookup", json!({ "q": "x" })), + text_response("done"), + ])), + ), + )) .register_tool(Arc::new(FakeTool::returning("lookup", "tool-output"))); // Durable journal + a plain recorder (live broadcast) to cross-check. @@ -204,14 +199,15 @@ async fn harness_run_with_capture_exports_generation_and_tool_io() { // the tool-create Input/Output panels — the fix for tinyhumansai/tinyagents#6. let mut harness: AgentHarness<()> = AgentHarness::new(); harness - .register_model( - "capture-model", - Arc::new(MockModel::with_responses(vec![ - tool_call_response("call-1", "lookup", json!({ "q": "weather" })), - text_response("all done"), - ])), - ) - .set_default_model("capture-model") + .with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named( + "capture-model", + Arc::new(MockModel::with_responses(vec![ + tool_call_response("call-1", "lookup", json!({ "q": "weather" })), + text_response("all done"), + ])), + ), + )) .register_tool(Arc::new(FakeTool::returning("lookup", "tool-output"))) .with_policy(RunPolicy { capture: PayloadCapture::all(), @@ -304,149 +300,3 @@ async fn harness_run_with_capture_exports_generation_and_tool_io() { "every generation nests under the run span" ); } - -/// A two-node line graph over `i32` with overwrite semantics: `a -> b`. -fn line_graph() -> tinyagents_graph::CompiledGraph<i32, i32> { - GraphBuilder::<i32, i32>::overwrite() - .add_node("a", |_s, _c: NodeContext| async move { - Ok(NodeResult::Update(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() - .expect("graph compiles") -} - -#[tokio::test] -async fn graph_run_journals_replayable_observations_with_coords() { - let journal = Arc::new(InMemoryGraphEventJournal::new()); - // `with_event_journal` wires a JournalGraphSink as the run's event sink. - let graph = line_graph().with_event_journal(journal.clone()); - - let run = graph.run(0).await.expect("graph run succeeds"); - let run_id = run.status.run_id.as_str().to_string(); - - let obs = journal - .read_from(&run_id, 0) - .await - .expect("replay graph journal"); - assert!( - obs.len() >= 3, - "expected several graph observations, got {}", - obs.len() - ); - - // Dense, monotonic offsets from 0; run/graph coordinates are stamped. - for (i, o) in obs.iter().enumerate() { - assert_eq!(o.offset, i as u64, "dense offsets from 0"); - assert_eq!(o.run_id.as_str(), run_id); - assert_eq!(o.root_run_id.as_str(), run_id); - assert_eq!(&o.graph_id, graph.graph_id()); - } - - // Step coordinates are carried: the executor runs at least one superstep. - assert!( - obs.iter().any(|o| o.step >= 1), - "observations should carry superstep coordinates" - ); - - // The run lifecycle bookends the stream. - assert!( - matches!(obs.first().unwrap().event, GraphEvent::RunStarted { .. }), - "first observation is RunStarted" - ); - assert!( - obs.iter() - .any(|o| matches!(o.event, GraphEvent::RunCompleted { .. })), - "completion is observable from the graph journal" - ); - - // Replay from a mid-stream offset returns exactly the tail. - let tail = journal.read_from(&run_id, 2).await.expect("replay tail"); - assert_eq!(tail.len(), obs.len() - 2); - assert_eq!(tail.first().unwrap().offset, 2); - - // Reading an unknown run is empty, not an error. - assert!( - journal.read_from("nope", 0).await.unwrap().is_empty(), - "unknown run replays empty" - ); -} - -#[tokio::test] -async fn graph_observations_export_to_langfuse_trace_and_spans() { - // Run a real graph, journal it, then feed the durable observations through - // the graph Langfuse exporter and assert the offline batch structure. - let journal = Arc::new(InMemoryGraphEventJournal::new()); - let graph = line_graph().with_event_journal(journal.clone()); - - let run = graph.run(0).await.expect("graph run succeeds"); - let run_id = run.status.run_id.as_str().to_string(); - let observations = journal - .read_from(&run_id, 0) - .await - .expect("replay graph journal"); - - // The exporter reuses the harness Langfuse transport (proxy mode here). - let exporter = GraphLangfuseExporter::new( - LangfuseClient::proxy("https://backend.test", "tok").expect("client"), - ); - let batch = exporter - .build_ingestion_batch(LangfuseTraceConfig::default(), &observations) - .expect("build batch"); - let events = batch["batch"].as_array().expect("batch array"); - - // The trace is created first and its id defaults to the run's root run id, - // aligning with the harness agent exporter for a unified trace. - assert_eq!(events[0]["type"], "trace-create"); - assert_eq!(events[0]["body"]["id"], run_id); - - // Node health telemetry rides on the trace: two nodes ran and completed. - let health = &events[0]["body"]["metadata"]["health"]; - assert_eq!(health["total_completed"], 2); - assert_eq!(health["total_failed"], 0); - assert_eq!(health["run_failed"], false); - - // Each node handler becomes a span, parented to its superstep span. - let node_a = format!("{run_id}:node:a:1"); - let node_b = format!("{run_id}:node:b:2"); - let has_node_span = |id: &str| { - events - .iter() - .any(|e| e["type"] == "span-create" && e["body"]["id"] == id) - }; - assert!(has_node_span(&node_a), "node a span present"); - assert!(has_node_span(&node_b), "node b span present"); - - // No node failed, so no span carries an ERROR level. - assert!( - !events - .iter() - .any(|e| e["type"] == "span-create" && e["body"]["level"] == "ERROR"), - "healthy run has no ERROR spans" - ); - - // The graph run is a structural span under the trace, and every step nests - // beneath it. Its id follows the `{trace}:run:{run}` scheme the harness - // exporter parents its agent run spans to, so a graph run and the sub-agent - // runs its nodes spawn reconstruct a single nested tree under one trace. - let graph_run_span = format!("{run_id}:run:{run_id}"); - let run_span = events - .iter() - .find(|e| e["body"]["id"] == graph_run_span) - .expect("graph run span"); - assert_eq!(run_span["type"], "span-create"); - assert!(run_span["body"].get("parentObservationId").is_none()); - let step_1 = events - .iter() - .find(|e| e["body"]["id"] == format!("{run_id}:step:1")) - .expect("step 1 span"); - assert_eq!(step_1["body"]["parentObservationId"], graph_run_span); - // A harness agent run spawned by a graph node carries parent_run_id == the - // graph run id, so its harness run span resolves to exactly this parent id - // and nests under the graph run rather than floating at the trace root. -} diff --git a/crates/tinyagents-integration-tests/tests/e2e_parallel_agents.rs b/crates/tinyagents-integration-tests/tests/e2e_parallel_agents.rs deleted file mode 100644 index 3bdfcbaba..000000000 --- a/crates/tinyagents-integration-tests/tests/e2e_parallel_agents.rs +++ /dev/null @@ -1,285 +0,0 @@ -//! TRUE end-to-end, fully OFFLINE: parallel branch execution + forking on a -//! durable [`CompiledGraph`]. -//! -//! This composes the **durable graph builder/executor** (`with_parallel`), -//! the **command** model (`Command::goto` fan-out), the **reducer** -//! (concurrent partial-update merge), and the **fork identity** carried on -//! [`NodeContext`]. No model or network is involved, so this test always runs -//! and must pass, with or without `OPENAI_API_KEY`. -//! -//! Topology built by [`fanout_graph`]: -//! -//! ```text -//! ┌──> worker-a ──┐ -//! START -> dispatch ──> worker-b ──> aggregate -> END -//! └──> worker-c ──┘ -//! └──> worker-d ──┘ -//! ``` -//! -//! `dispatch` fans out to four worker branches via a single -//! `Command::goto(["worker-a", .., "worker-d"])`. Each worker contributes a -//! partial update (its produced value, the fork branch index it observed, and -//! its own name) that the reducer folds into shared channels at the superstep -//! boundary. The downstream `aggregate` node then reads the *merged* state and -//! computes a derived total, proving fan-in visibility. - -use std::sync::Arc; -use std::sync::atomic::{AtomicUsize, Ordering}; -use std::time::Duration; - -use tinyagents_graph::ClosureStateReducer; -use tinyagents_graph::END; -use tinyagents_graph::*; -use tinyagents_harness::*; -use tinyagents_registry::*; - -/// Committed graph state. Every field is filled by branches merging through the -/// reducer; `total` is derived downstream from the merged `values`. -#[derive(Clone, Debug, Default, PartialEq, Eq)] -struct AgentState { - /// Values produced by each worker branch, in reducer-application order. - values: Vec<i32>, - /// The `branch_index` each worker observed on its `NodeContext::fork`. - /// `usize::MAX` is the sentinel for "no fork identity" (sequential mode). - forks: Vec<usize>, - /// Names of the workers that contributed, in reducer-application order. - workers: Vec<String>, - /// Sum computed downstream over the merged `values` (fan-in proof). - total: Option<i32>, -} - -/// Partial-update type merged through the [`ClosureStateReducer`]. -#[derive(Clone, Debug)] -enum AgentUpdate { - /// One worker branch's contribution. - Work { - value: i32, - fork: usize, - worker: String, - }, - /// The downstream aggregate's derived total. - Total(i32), -} - -/// Shared instrumentation that records the maximum number of branches that were -/// ever in flight simultaneously. In parallel mode this should reach the branch -/// count; in sequential mode it must never exceed one. -#[derive(Clone)] -struct Inflight { - current: Arc<AtomicUsize>, - max: Arc<AtomicUsize>, -} - -impl Inflight { - fn new() -> Self { - Self { - current: Arc::new(AtomicUsize::new(0)), - max: Arc::new(AtomicUsize::new(0)), - } - } - - fn max_observed(&self) -> usize { - self.max.load(Ordering::SeqCst) - } - - /// Marks a branch as in flight for `sleep`, updating the high-water mark, - /// then yields `value`. - async fn run<T>(&self, sleep: Duration, value: T) -> T { - let now = self.current.fetch_add(1, Ordering::SeqCst) + 1; - self.max.fetch_max(now, Ordering::SeqCst); - tokio::time::sleep(sleep).await; - self.current.fetch_sub(1, Ordering::SeqCst); - value - } -} - -/// Reads the fork branch index from a [`NodeContext`], using `usize::MAX` as -/// the "no fork identity" sentinel (sequential mode / single-node steps). -fn fork_index(ctx: &NodeContext) -> usize { - ctx.fork - .as_ref() - .map(|f| f.branch_index) - .unwrap_or(usize::MAX) -} - -/// Builds the fan-out/fan-in agent graph. `parallel` toggles concurrent branch -/// execution; `inflight` instruments concurrency. Worker sleeps are reversed -/// (a longest, d shortest) so reducer-merge order can never accidentally match -/// completion order — any deterministic ordering must come from the active-set -/// position, not from which branch finishes first. -fn fanout_graph(parallel: bool, inflight: Inflight) -> CompiledGraph<AgentState, AgentUpdate> { - // Each worker: (name, produced value, sleep millis). - let workers = [ - ("worker-a", 1, 80u64), - ("worker-b", 2, 60), - ("worker-c", 4, 40), - ("worker-d", 8, 20), - ]; - - let mut builder = GraphBuilder::<AgentState, AgentUpdate>::new() - .with_parallel(parallel) - .set_reducer(ClosureStateReducer::new( - |mut s: AgentState, u: AgentUpdate| { - match u { - AgentUpdate::Work { - value, - fork, - worker, - } => { - s.values.push(value); - s.forks.push(fork); - s.workers.push(worker); - } - AgentUpdate::Total(t) => s.total = Some(t), - } - Ok(s) - }, - )) - // dispatch fans out to all workers in one command. - .add_node("dispatch", |_s: AgentState, _c: NodeContext| async move { - Ok(NodeResult::Command(Command::default().with_goto([ - "worker-a", "worker-b", "worker-c", "worker-d", - ]))) - }); - - for (name, value, sleep_ms) in workers { - let inflight = inflight.clone(); - builder = builder.add_node(name, move |_s: AgentState, c: NodeContext| { - let inflight = inflight.clone(); - let fork = fork_index(&c); - async move { - let update = AgentUpdate::Work { - value, - fork, - worker: name.to_string(), - }; - Ok(NodeResult::Update( - inflight.run(Duration::from_millis(sleep_ms), update).await, - )) - } - }); - } - - builder = builder - // aggregate reads the merged state after fan-in. - .add_node("aggregate", |s: AgentState, _c: NodeContext| async move { - Ok(NodeResult::Update(AgentUpdate::Total( - s.values.iter().sum(), - ))) - }) - .set_entry("dispatch") - .mark_command_routing("dispatch") - .add_edge("worker-a", "aggregate") - .add_edge("worker-b", "aggregate") - .add_edge("worker-c", "aggregate") - .add_edge("worker-d", "aggregate") - .set_finish("aggregate"); - - builder.compile().expect("graph compiles") -} - -#[tokio::test] -async fn parallel_fanout_merges_all_branches_and_downstream_sees_them() { - let inflight = Inflight::new(); - let graph = fanout_graph(true, inflight.clone()); - let run = graph.run(AgentState::default()).await.unwrap(); - - // All four worker branches were in flight at the same time. - assert_eq!( - inflight.max_observed(), - 4, - "parallel mode should run every branch concurrently" - ); - - // The reducer folded in every branch's contribution, in deterministic - // active-set order (goto order), NOT completion order. - assert_eq!(run.state.values, vec![1, 2, 4, 8]); - assert_eq!( - run.state.workers, - vec!["worker-a", "worker-b", "worker-c", "worker-d"], - ); - - // Each fork saw its stable branch index from the active set. - assert_eq!(run.state.forks, vec![0, 1, 2, 3]); - - // The downstream aggregate node observed the fully merged state. - assert_eq!(run.state.total, Some(15)); - - // dispatch | (a,b,c,d) | aggregate == 3 supersteps. - assert_eq!(run.steps, 3); - - // Every node was visited; workers executed within a single superstep. - for node in [ - "dispatch", - "worker-a", - "worker-b", - "worker-c", - "worker-d", - "aggregate", - ] { - assert!( - run.visited.iter().any(|n| n.as_str() == node), - "expected `{node}` in visited history {:?}", - run.visited, - ); - } -} - -#[tokio::test] -async fn sequential_mode_runs_one_branch_at_a_time_without_fork_identity() { - let inflight = Inflight::new(); - let graph = fanout_graph(false, inflight.clone()); - let run = graph.run(AgentState::default()).await.unwrap(); - - // Never more than one branch in flight in sequential mode. - assert_eq!( - inflight.max_observed(), - 1, - "sequential mode must serialize branches" - ); - - // Sequential branches receive no fork identity. - assert_eq!( - run.state.forks, - vec![usize::MAX, usize::MAX, usize::MAX, usize::MAX], - ); -} - -#[tokio::test] -async fn parallel_and_sequential_reach_the_same_final_state() { - // Final committed state (modulo fork identity) must be identical whether - // branches run concurrently or one-at-a-time. This is the core determinism - // guarantee of the executor: parallelism is an execution strategy, not a - // semantic change. - let parallel = fanout_graph(true, Inflight::new()) - .run(AgentState::default()) - .await - .unwrap(); - let sequential = fanout_graph(false, Inflight::new()) - .run(AgentState::default()) - .await - .unwrap(); - - assert_eq!(parallel.state.values, sequential.state.values); - assert_eq!(parallel.state.workers, sequential.state.workers); - assert_eq!(parallel.state.total, sequential.state.total); - assert_eq!(parallel.steps, sequential.steps); - - // The only intentional difference is fork identity. - assert_ne!(parallel.state.forks, sequential.state.forks); -} - -#[tokio::test] -async fn parallel_merge_order_is_reproducible_across_runs() { - // Re-running the same concurrent fan-out must always produce the same merge - // order, regardless of which branch's sleep finishes first. - for _ in 0..5 { - let run = fanout_graph(true, Inflight::new()) - .run(AgentState::default()) - .await - .unwrap(); - assert_eq!(run.state.values, vec![1, 2, 4, 8]); - assert_eq!(run.state.forks, vec![0, 1, 2, 3]); - assert_eq!(run.state.total, Some(15)); - } -} diff --git a/crates/tinyagents-integration-tests/tests/e2e_parallel_and_journals.rs b/crates/tinyagents-integration-tests/tests/e2e_parallel_and_journals.rs deleted file mode 100644 index fa426564f..000000000 --- a/crates/tinyagents-integration-tests/tests/e2e_parallel_and_journals.rs +++ /dev/null @@ -1,500 +0,0 @@ -//! End-to-end coverage for two independent harness building blocks: -//! -//! - **Part A** exercises the ordered, bounded [`map_reduce`] helper driving -//! *real* async agent work: each mapped item spins up (or shares) an -//! [`AgentHarness`] and performs a genuine `invoke_default` model call. We -//! assert input-order preservation, skewed-completion ordering, every failure -//! policy, and concurrency bounding. -//! - **Part B** feeds a durable [`InMemoryEventJournal`] from a real recorded -//! harness run and exercises windowed / filtered reads, then validates run -//! lineage queries over an [`InMemoryStatusStore`]. - -use std::sync::Arc; -use std::sync::atomic::{AtomicUsize, Ordering}; -use std::time::Duration; - -use tinyagents_graph::*; -use tinyagents_harness::context::{RunConfig, RunContext}; -use tinyagents_harness::events::HarnessRunStatus; -use tinyagents_harness::ids::{ComponentId, EventId, HarnessPhase, RunId, ThreadId}; -use tinyagents_harness::runtime::AgentHarness; -use tinyagents_harness::testkit::EventRecorder; -use tinyagents_harness::*; -use tinyagents_registry::*; -use tinyinference_llm::message::Message; -use tinyinference_llm::providers::MockModel; - -// --------------------------------------------------------------------------- -// Helpers -// --------------------------------------------------------------------------- - -/// Builds a harness whose default model always replies with `reply`. -fn constant_harness(reply: &str) -> AgentHarness<()> { - let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("m", Arc::new(MockModel::constant(reply))); - harness -} - -// --------------------------------------------------------------------------- -// Part A — parallel map/reduce over real harness invocations -// --------------------------------------------------------------------------- - -/// Test 1: N=5 real harness invocations run concurrently through `map_reduce` -/// (CollectAll). All five succeed and the outcomes stay in input order, each -/// carrying the index of the input that produced it. -#[tokio::test] -async fn map_reduce_runs_real_harness_invocations_in_input_order() { - let harness = Arc::new(constant_harness("reply")); - let items: Vec<usize> = (0..5).collect(); - - let opts = ParallelOptions::default().with_failure_policy(FailurePolicy::CollectAll); - - let outcome = map_reduce(items, opts, move |index, item| { - let harness = harness.clone(); - async move { - // Distinct, index-derived prompt so each unit of work is unique. - let prompt = format!("question #{item}"); - let run = harness - .invoke_default(&(), vec![Message::user(prompt)]) - .await?; - let text = run.text().unwrap_or_default(); - Ok::<_, TinyAgentsError>(format!("{index}:{text}")) - } - }) - .await - .expect("map_reduce should not error under CollectAll"); - - assert_eq!(outcome.success_count(), 5); - assert_eq!(outcome.failure_count(), 0); - - // Outcomes preserve input order: outcome[i].index == i and the value was - // produced by the i-th closure invocation. - for (i, item) in outcome.outcomes.iter().enumerate() { - assert_eq!(item.index, i, "outcome index must match input position"); - assert!(item.is_ok()); - assert_eq!(item.result.as_ref().unwrap(), &format!("{i}:reply")); - } - - // The borrowed successes are likewise in input order. - let successes = outcome.successes(); - assert_eq!(successes.len(), 5); - for (i, text) in successes.iter().enumerate() { - assert_eq!(*text, &format!("{i}:reply")); - } -} - -/// Test 2: items sleep for descending durations (so completion order is the -/// reverse of input order), yet results still land in input order. -#[tokio::test] -async fn map_reduce_preserves_input_order_under_skewed_completion() { - let harness = Arc::new(constant_harness("done")); - let items: Vec<usize> = (0..5).collect(); - let n = items.len(); - - let outcome = map_reduce(items, ParallelOptions::default(), move |index, item| { - let harness = harness.clone(); - async move { - // Earlier items sleep longest, so they complete last. - let millis = ((n - index) as u64) * 20; - tokio::time::sleep(Duration::from_millis(millis)).await; - let run = harness - .invoke_default(&(), vec![Message::user(format!("item {item}"))]) - .await?; - Ok::<_, TinyAgentsError>(format!("{index}:{}", run.text().unwrap_or_default())) - } - }) - .await - .expect("map_reduce should succeed"); - - assert_eq!(outcome.success_count(), 5); - let ordered = outcome.into_successes(); - let expected: Vec<String> = (0..5).map(|i| format!("{i}:done")).collect(); - assert_eq!( - ordered, expected, - "results must follow input order, not completion order" - ); -} - -/// Test 3: each failure policy behaves as documented when roughly half the -/// items fail. Odd indices fail; even indices succeed (3 successes of 5). -#[tokio::test] -async fn map_reduce_honors_each_failure_policy() { - // Closure factory: even index succeeds, odd index fails. - fn work( - harness: Arc<AgentHarness<()>>, - ) -> impl Fn( - usize, - usize, - ) -> std::pin::Pin< - Box< - dyn std::future::Future<Output = tinyagents_harness::Result<String>> - + std::marker::Send, - >, - > { - move |index, _item| { - let harness = harness.clone(); - Box::pin(async move { - if index % 2 == 1 { - return Err(TinyAgentsError::Graph(format!("item {index} failed"))); - } - let run = harness - .invoke_default(&(), vec![Message::user(format!("ok {index}"))]) - .await?; - Ok::<_, TinyAgentsError>(format!("{index}:{}", run.text().unwrap_or_default())) - }) - } - } - - let harness = Arc::new(constant_harness("ok")); - - // CollectAll: never errors; records per-item successes and failures. - let collect = map_reduce( - (0..5).collect::<Vec<_>>(), - ParallelOptions::default().with_failure_policy(FailurePolicy::CollectAll), - work(harness.clone()), - ) - .await - .expect("CollectAll never errors"); - assert_eq!(collect.success_count(), 3); - assert_eq!(collect.failure_count(), 2); - // Failures are recorded, not raised, and remain index-tagged in order. - for item in &collect.outcomes { - assert_eq!(item.is_ok(), item.index % 2 == 0); - } - - // FailFast: returns the first error (in input order → index 1). - let fail_fast = map_reduce( - (0..5).collect::<Vec<_>>(), - ParallelOptions::default().with_failure_policy(FailurePolicy::FailFast), - work(harness.clone()), - ) - .await; - match fail_fast { - Err(TinyAgentsError::Graph(msg)) => assert_eq!(msg, "item 1 failed"), - other => panic!("FailFast should return the first item error, got {other:?}"), - } - - // Quorum below the success count errors. - let quorum_high = map_reduce( - (0..5).collect::<Vec<_>>(), - ParallelOptions::default().with_failure_policy(FailurePolicy::Quorum(4)), - work(harness.clone()), - ) - .await; - assert!( - matches!(quorum_high, Err(TinyAgentsError::Graph(_))), - "Quorum(4) with only 3 successes must error" - ); - - // Quorum at/below the success count is Ok. - let quorum_ok = map_reduce( - (0..5).collect::<Vec<_>>(), - ParallelOptions::default().with_failure_policy(FailurePolicy::Quorum(3)), - work(harness.clone()), - ) - .await - .expect("Quorum(3) with 3 successes must be Ok"); - assert_eq!(quorum_ok.success_count(), 3); - - // BestEffort: Ok, keeping only the successful outputs. - let best = map_reduce( - (0..5).collect::<Vec<_>>(), - ParallelOptions::default().with_failure_policy(FailurePolicy::BestEffort), - work(harness), - ) - .await - .expect("BestEffort never errors"); - assert_eq!(best.success_count(), 3); - assert_eq!( - best.failure_count(), - 0, - "BestEffort drops failures entirely" - ); - let kept = best.into_successes(); - assert_eq!( - kept, - vec!["0:ok".to_string(), "2:ok".to_string(), "4:ok".to_string()] - ); -} - -/// Test 4: `with_max_concurrency(2)` bounds the number of simultaneously -/// executing closures. A live counter tracks in-flight work and a peak counter -/// records the high-water mark, which must never exceed 2. -#[tokio::test] -async fn map_reduce_bounds_concurrency() { - let harness = Arc::new(constant_harness("x")); - let live = Arc::new(AtomicUsize::new(0)); - let peak = Arc::new(AtomicUsize::new(0)); - - let opts = ParallelOptions::default().with_max_concurrency(2); - - let live_c = live.clone(); - let peak_c = peak.clone(); - let outcome = map_reduce((0..8).collect::<Vec<_>>(), opts, move |index, _item| { - let harness = harness.clone(); - let live = live_c.clone(); - let peak = peak_c.clone(); - async move { - let now = live.fetch_add(1, Ordering::SeqCst) + 1; - peak.fetch_max(now, Ordering::SeqCst); - // Hold the slot long enough for overlap to be observable. - tokio::time::sleep(Duration::from_millis(15)).await; - let run = harness - .invoke_default(&(), vec![Message::user(format!("n{index}"))]) - .await?; - live.fetch_sub(1, Ordering::SeqCst); - Ok::<_, TinyAgentsError>(run.text().unwrap_or_default()) - } - }) - .await - .expect("bounded map_reduce should succeed"); - - assert_eq!(outcome.success_count(), 8); - let observed_peak = peak.load(Ordering::SeqCst); - assert!( - observed_peak <= 2, - "peak concurrency {observed_peak} must not exceed the configured bound of 2" - ); - assert!( - observed_peak >= 1, - "at least one closure must have run (peak was {observed_peak})" - ); -} - -/// Test 4b: a per-item timeout converts only the slow item into a recoverable -/// failure while the fast items still perform real harness work and succeed -/// (CollectAll). Time is virtual (`start_paused`) so the batch is deterministic. -#[tokio::test(start_paused = true)] -async fn map_reduce_item_timeout_fails_only_the_slow_item() { - let harness = Arc::new(constant_harness("ok")); - // The middle item sleeps for an hour; the per-item timeout is 50ms. - let out = map_reduce( - vec![0u64, 3_600_000, 0], - ParallelOptions::default() - .with_failure_policy(FailurePolicy::CollectAll) - .with_item_timeout(Duration::from_millis(50)), - move |index, ms| { - let harness = harness.clone(); - async move { - tokio::time::sleep(Duration::from_millis(ms)).await; - let run = harness - .invoke_default(&(), vec![Message::user(format!("n{index}"))]) - .await?; - Ok::<_, TinyAgentsError>(run.text().unwrap_or_default()) - } - }, - ) - .await - .expect("CollectAll never errors even when an item times out"); - - assert_eq!(out.success_count(), 2, "the two fast items complete"); - assert_eq!(out.failure_count(), 1, "the hanging item times out"); - assert!(!out.outcomes[1].is_ok(), "index 1 is the timed-out item"); - // The per-item outcome carries the stringified Timeout error. - let message = out.outcomes[1] - .result - .as_ref() - .expect_err("the slow item is a failure"); - assert!( - message.contains("timed out"), - "the slow item's error should describe a timeout, got {message:?}" - ); -} - -/// Test 4c: a total (batch-wide) timeout aborts the whole `map_reduce` and -/// surfaces [`TinyAgentsError::Timeout`]. -#[tokio::test(start_paused = true)] -async fn map_reduce_total_timeout_aborts_the_batch() { - let err = map_reduce( - vec![3_600_000u64], - ParallelOptions::default().with_total_timeout(Duration::from_millis(20)), - |_index, ms| async move { - tokio::time::sleep(Duration::from_millis(ms)).await; - Ok::<_, TinyAgentsError>(ms) - }, - ) - .await - .expect_err("the total timeout must abort the batch"); - assert!( - matches!(err, TinyAgentsError::Timeout(_)), - "expected Timeout, got {err:?}" - ); -} - -/// Test 4d: a pre-cancelled [`CancellationToken`] aborts the batch cooperatively -/// and surfaces [`TinyAgentsError::Cancelled`]. -#[tokio::test] -async fn map_reduce_cancellation_token_stops_the_batch() { - let token = CancellationToken::new(); - token.cancel(); - - let err = map_reduce( - vec![0u64], - ParallelOptions::default().with_cancellation(token), - |_index, _n| async move { - tokio::time::sleep(Duration::from_secs(3_600)).await; - Ok::<_, TinyAgentsError>(0u64) - }, - ) - .await - .expect_err("a pre-cancelled token must abort the batch"); - assert!( - matches!(err, TinyAgentsError::Cancelled), - "expected Cancelled, got {err:?}" - ); -} - -// --------------------------------------------------------------------------- -// Part B — journals fed from a real run, plus status lineage -// --------------------------------------------------------------------------- - -/// Test 5: run a real harness with an [`EventRecorder`], wrap each recorded -/// [`AgentEvent`] into an [`AgentObservation`] (increasing offset, shared -/// run id) and append them to an [`InMemoryEventJournal`]. Then exercise -/// full, windowed, and kind-filtered reads. -#[tokio::test] -async fn journal_windowed_and_filtered_reads_from_real_run() { - let harness = constant_harness("hello journal"); - let recorder = EventRecorder::new(); - - let run_id_str = "journal-run"; - let ctx: RunContext<()> = - RunContext::new(RunConfig::new(run_id_str), ()).with_events(recorder.sink()); - harness - .invoke_in_context(&(), ctx, vec![Message::user("hi")]) - .await - .expect("real harness run should succeed"); - - let events = recorder.events(); - assert!( - events.len() >= 2, - "a real run should emit at least run.started and run.completed" - ); - - // Feed the durable journal from the real recorded events. - let journal = InMemoryEventJournal::default(); - let run_id = RunId::new(run_id_str); - for (offset, event) in events.iter().cloned().enumerate() { - let obs = AgentObservation { - event_id: EventId::new(format!("e{offset}")), - run_id: run_id.clone(), - parent_run_id: None, - root_run_id: run_id.clone(), - offset: offset as u64, - ts_ms: offset as u64, - event, - }; - journal.append(obs).await.expect("append should succeed"); - } - - // read_from(0) replays the entire run. - let all = journal - .read_from(run_id_str, 0) - .await - .expect("read_from should succeed"); - assert_eq!(all.len(), events.len()); - - // read_window truncates to at most `limit`. - let window = journal - .read_window(run_id_str, 0, 2) - .await - .expect("read_window should succeed"); - assert!(window.len() <= 2, "window must honor the limit"); - assert_eq!(window.len(), 2.min(events.len())); - - // read_filtered returns only the requested kinds. A real run always emits - // run.started and run.completed. - let kinds = ["run.started", "run.completed"]; - let filtered = journal - .read_filtered(run_id_str, 0, &kinds) - .await - .expect("read_filtered should succeed"); - assert!( - !filtered.is_empty(), - "run.started/run.completed should be present in a real run" - ); - for obs in &filtered { - assert!( - obs.event.kind().starts_with("run."), - "filtered observation kind {:?} should start with run.", - obs.event.kind() - ); - assert!(kinds.contains(&obs.event.kind())); - } - // Both terminal-lifecycle kinds should be captured. - assert!(filtered.iter().any(|o| o.event.kind() == "run.started")); - assert!(filtered.iter().any(|o| o.event.kind() == "run.completed")); - - // An unknown run reads back empty. - let missing = journal - .read_from("no-such-run", 0) - .await - .expect("reading an unknown run is Ok"); - assert!(missing.is_empty()); -} - -/// Test 6: build a small run tree in an [`InMemoryStatusStore`] and validate -/// lineage queries: `list_by_root` returns every descendant of a root, while -/// `list_active` excludes terminal (completed) runs. -#[tokio::test] -async fn status_lineage_by_root_and_active() { - let store = InMemoryStatusStore::default(); - let component = ComponentId::new("agent_loop"); - let root = RunId::new("root"); - - // Parent (the root itself), running. - let mut parent = HarnessRunStatus::new(root.clone(), component.clone()) - .with_thread(ThreadId::new("thread-1")); - parent.mark_running(HarnessPhase::Model); - - // Child under the root, running. - let mut child = HarnessRunStatus::new(RunId::new("child"), component.clone()) - .with_parent(root.clone(), root.clone()); - child.mark_running(HarnessPhase::Tools); - - // Completed sibling under the same root. - let mut sibling = HarnessRunStatus::new(RunId::new("sibling"), component.clone()) - .with_parent(root.clone(), root.clone()); - sibling.mark_completed(); - - // Unrelated run under a different root. - let mut unrelated = HarnessRunStatus::new(RunId::new("other"), component.clone()); - unrelated.mark_running(HarnessPhase::Model); - - for status in [parent, child, sibling, unrelated] { - store.put_status(status).await.expect("put_status"); - } - - // list_by_root returns the 3 descendants (root, child, sibling). - let mut lineage = store.list_by_root("root").await.expect("list_by_root"); - lineage.sort_by(|a, b| a.run_id.as_str().cmp(b.run_id.as_str())); - let lineage_ids: Vec<&str> = lineage.iter().map(|s| s.run_id.as_str()).collect(); - assert_eq!(lineage_ids, vec!["child", "root", "sibling"]); - assert!( - !lineage_ids.contains(&"other"), - "the unrelated run must not appear under this root" - ); - - // list_active excludes the completed sibling but keeps the running runs. - let active = store.list_active().await.expect("list_active"); - let active_ids: Vec<&str> = active.iter().map(|s| s.run_id.as_str()).collect(); - assert!(active_ids.contains(&"root")); - assert!(active_ids.contains(&"child")); - assert!(active_ids.contains(&"other")); - assert!( - !active_ids.contains(&"sibling"), - "the completed sibling must be excluded from active runs" - ); - - // list_by_thread finds the thread-tagged root. - let by_thread = store - .list_by_thread("thread-1") - .await - .expect("list_by_thread"); - assert_eq!(by_thread.len(), 1); - assert_eq!(by_thread[0].run_id.as_str(), "root"); - - // get_status round-trips a known run. - let got = store.get_status("child").await.expect("get_status"); - assert!(got.is_some()); -} diff --git a/crates/tinyagents-integration-tests/tests/e2e_prompt_cache_kv.rs b/crates/tinyagents-integration-tests/tests/e2e_prompt_cache_kv.rs index 03ff1461d..0e35053fe 100644 --- a/crates/tinyagents-integration-tests/tests/e2e_prompt_cache_kv.rs +++ b/crates/tinyagents-integration-tests/tests/e2e_prompt_cache_kv.rs @@ -142,7 +142,9 @@ impl Tool for LookupTool { async fn harness_tool_turns_reuse_the_provider_kv_prompt_prefix() { let server = Arc::new(KvCacheMockServer::default()); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock-kv", server.clone()); + harness.with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named("mock-kv", server.clone()), + )); harness.register_tool(Arc::new(LookupTool)); let mut policy = RunPolicy::default(); policy.cache.protect_prompt_prefix = true; diff --git a/crates/tinyagents-integration-tests/tests/e2e_public_api_contracts.rs b/crates/tinyagents-integration-tests/tests/e2e_public_api_contracts.rs index d7dbb099f..184de3f2f 100644 --- a/crates/tinyagents-integration-tests/tests/e2e_public_api_contracts.rs +++ b/crates/tinyagents-integration-tests/tests/e2e_public_api_contracts.rs @@ -1,5 +1,4 @@ -//! End-to-end contract coverage for public support APIs that recursive graph -//! and harness runs depend on. +//! End-to-end contract coverage for public harness support APIs. //! //! These tests intentionally sit in `tests/` rather than module-local unit tests //! so the e2e coverage pass exercises the same public surface downstream users @@ -13,10 +12,6 @@ use std::time::{Duration, Instant}; use serde_json::{Map, Value, json}; -use tinyagents_graph::{ - Barrier, BinaryAggregate, ChannelSet, ChannelState, ChannelUpdate, Delta, Ephemeral, LastValue, - Messages, NamedBarrier, Topic, Untracked, -}; use tinyagents_harness::TinyAgentsError; use tinyagents_harness::cache::{ CacheLayoutEvent, InMemoryResponseCache, PromptCacheLayout, ResponseCache, cache_key, @@ -185,99 +180,6 @@ async fn stores_round_trip_across_ephemeral_and_file_backends() { let _ = std::fs::remove_dir_all(&dir); } -#[tokio::test] -async fn channel_state_contracts_handle_merges_barriers_and_conflicts() { - let mut set = ChannelSet::new() - .with_channel("last", LastValue) - .with_channel("topic", Topic) - .with_channel("delta", Delta) - .with_channel("messages", Messages) - .with_channel("ephemeral", Ephemeral) - .with_channel("scratch", Untracked) - .with_channel("barrier", Barrier::new(2)) - .with_channel("named", NamedBarrier::new(["left", "right"])) - .with_channel( - "sum", - BinaryAggregate::new(|current, incoming| { - Ok(Value::from( - current.as_i64().unwrap_or(0) + incoming.as_i64().unwrap_or(0), - )) - }), - ); - - assert!(set.contains("last")); - assert!(!set.allows_concurrent("last").unwrap()); - assert!(set.allows_concurrent("topic").unwrap()); - assert!(!set.is_ready("barrier").unwrap()); - - set.apply_update("last", json!("a")).unwrap(); - set.apply_update("topic", json!("first")).unwrap(); - set.apply_update("topic", json!(["second", "third"])) - .unwrap(); - set.apply_update("delta", json!(2)).unwrap(); - set.apply_update("delta", json!(3)).unwrap(); - set.apply_update("messages", json!({ "id": "m1", "text": "old" })) - .unwrap(); - set.apply_update("messages", json!({ "id": "m1", "text": "new" })) - .unwrap(); - set.apply_update("ephemeral", json!("one-shot")).unwrap(); - set.apply_update("scratch", json!("hidden")).unwrap(); - set.apply_update("barrier", json!("left")).unwrap(); - set.apply_update("barrier", json!("right")).unwrap(); - set.apply_update("named", json!({ "left": true })).unwrap(); - set.apply_update("named", json!({ "right": true })).unwrap(); - set.apply_update("sum", json!(5)).unwrap(); - set.apply_update("sum", json!(7)).unwrap(); - - assert_eq!(set.get("delta"), Some(&json!(5))); - assert_eq!(set.get("sum"), Some(&json!(12))); - assert!(set.is_ready("barrier").unwrap()); - assert!(set.is_ready("named").unwrap()); - assert!(!set.snapshot().contains_key("scratch")); - - let state = ChannelState::new() - .with_channel("last", LastValue) - .with_channel("topic", Topic) - .with_channel("delta", Delta) - .with_channel("messages", Messages) - .with_channel("ephemeral", Ephemeral) - .with_channel("scratch", Untracked) - .with_channel("barrier", Barrier::new(2)) - .with_channel("named", NamedBarrier::new(["left", "right"])) - .with_channel( - "sum", - BinaryAggregate::new(|current, incoming| { - Ok(Value::from( - current.as_i64().unwrap_or(0) + incoming.as_i64().unwrap_or(0), - )) - }), - ); - let state = state - .merge( - ChannelUpdate::new() - .set("last", json!("step-1")) - .set("ephemeral", json!("flash")) - .at_step(1), - ) - .unwrap(); - assert_eq!(state.channels().get("ephemeral"), Some(&json!("flash"))); - let state = state - .merge(ChannelUpdate::new().set("last", json!("step-2")).at_step(2)) - .unwrap(); - assert!(state.channels().get("ephemeral").is_none()); - - let conflict = state - .clone() - .merge(ChannelUpdate::new().set("last", json!("a")).at_step(3)) - .unwrap() - .merge(ChannelUpdate::new().set("last", json!("b")).at_step(3)) - .unwrap_err(); - assert!(matches!( - conflict, - TinyAgentsError::InvalidConcurrentUpdate(_) - )); -} - #[tokio::test] async fn retry_rate_limit_and_summarization_contracts_are_deterministic() { let retry = RetryPolicy::default() diff --git a/crates/tinyagents-integration-tests/tests/e2e_reasoning_and_selection.rs b/crates/tinyagents-integration-tests/tests/e2e_reasoning_and_selection.rs index c41881dd2..8c3514d39 100644 --- a/crates/tinyagents-integration-tests/tests/e2e_reasoning_and_selection.rs +++ b/crates/tinyagents-integration-tests/tests/e2e_reasoning_and_selection.rs @@ -21,13 +21,11 @@ use futures::StreamExt; use tinyagents_harness::context::{RunConfig, RunContext}; use tinyagents_harness::events::AgentEvent; use tinyagents_harness::middleware::{ContextualToolSelectionMiddleware, ToolSelectionContext}; -use tinyagents_harness::model_registry::{ModelRegistry, ModelSelection}; use tinyagents_harness::runtime::AgentHarness; use tinyagents_harness::testkit::{EventRecorder, FakeTool, ScriptedModel, StreamingMock}; use tinyinference_llm::message::{AssistantMessage, ContentBlock, Message, MessageDelta}; use tinyinference_llm::model::{ - ChatModel, ModelProfile, ModelRequest, ModelResolutionSource, ModelResponse, ModelStatus, - ModelStreamItem, StreamAccumulator, + ChatModel, ModelProfile, ModelRequest, ModelResponse, ModelStreamItem, StreamAccumulator, }; use tinyinference_llm::tool::ToolSchema; use tinyinference_llm::usage::Usage; @@ -166,7 +164,9 @@ async fn exposed_tools_at_depth( let scripted = Arc::new(ScriptedModel::new(vec![text_response("done")])); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", scripted.clone()); + harness.with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named("mock", scripted.clone()), + )); harness.register_tool(Arc::new(FakeTool::returning("a", "ok"))); harness.register_tool(Arc::new(FakeTool::returning("b", "ok"))); harness.register_tool(Arc::new(FakeTool::returning("c", "ok"))); @@ -244,7 +244,9 @@ async fn inheriting_narrows_child_and_emits_tools_filtered_event() { let scripted = Arc::new(ScriptedModel::new(vec![text_response("done")])); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", scripted.clone()); + harness.with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named("mock", scripted.clone()), + )); harness.register_tool(Arc::new(FakeTool::returning("a", "ok"))); harness.register_tool(Arc::new(FakeTool::returning("b", "ok"))); harness.register_tool(Arc::new(FakeTool::returning("c", "ok"))); @@ -295,122 +297,3 @@ async fn inheriting_narrows_child_and_emits_tools_filtered_event() { "the exposure decision is journaled under the tool.filtered kind" ); } - -// --------------------------------------------------------------------------- -// Part C — retired-model resolution gating -// --------------------------------------------------------------------------- - -/// A minimal [`ChatModel`] that advertises a caller-supplied [`ModelProfile`], -/// so resolution can observe its lifecycle [`ModelStatus`]. -struct ProfiledModel { - profile: ModelProfile, -} - -impl ProfiledModel { - fn with_status(status: ModelStatus) -> Self { - Self { - profile: ModelProfile { - status, - ..ModelProfile::permissive() - }, - } - } -} - -#[async_trait] -impl ChatModel<()> for ProfiledModel { - fn profile(&self) -> Option<&ModelProfile> { - Some(&self.profile) - } - - async fn invoke( - &self, - _state: &(), - _request: ModelRequest, - ) -> tinyinference_llm::Result<ModelResponse> { - Ok(text_response("ok")) - } -} - -/// A model whose profile reports [`ModelStatus::Retired`] is skipped by -/// resolution even on an explicit override (it falls back to the live default), -/// and opting in via `allow_retired` re-admits it as the request override. -#[tokio::test] -async fn retired_override_falls_back_to_live_default_unless_allowed() { - let mut registry: ModelRegistry<()> = ModelRegistry::new(); - // "live" registers first, so it is the registry default. - registry - .register( - "live", - Arc::new(ProfiledModel::with_status(ModelStatus::Stable)), - ) - .register( - "retired", - Arc::new(ProfiledModel::with_status(ModelStatus::Retired)), - ); - - // A live model resolves via an explicit override. - let live = registry - .resolve(ModelSelection { - requested: Some("live".into()), - ..ModelSelection::default() - }) - .expect("a live model resolves"); - assert_eq!(live.resolved.name, "live"); - assert_eq!(live.resolved.source, ModelResolutionSource::RequestOverride); - - // A retired override is skipped (fail closed) and resolution falls through to - // the live registry default rather than selecting the retired model. - let fell_back = registry - .resolve(ModelSelection { - requested: Some("retired".into()), - ..ModelSelection::default() - }) - .expect("resolution falls back to the live default"); - assert_eq!( - fell_back.resolved.name, "live", - "a retired override must not be selected; it falls back to the live default" - ); - assert_eq!( - fell_back.resolved.source, - ModelResolutionSource::RegistryDefault - ); - - // `allow_retired = true` re-admits the retired model as the request override. - let readmitted = registry - .resolve(ModelSelection { - requested: Some("retired".into()), - allow_retired: true, - ..ModelSelection::default() - }) - .expect("allow_retired re-admits the retired model"); - assert_eq!(readmitted.resolved.name, "retired"); - assert_eq!( - readmitted.resolved.source, - ModelResolutionSource::RequestOverride - ); -} - -/// When every candidate is retired, default resolution yields nothing; -/// `allow_retired` makes the retired default resolvable again. -#[tokio::test] -async fn all_retired_registry_resolves_to_none_until_allowed() { - let mut registry: ModelRegistry<()> = ModelRegistry::new(); - registry.register( - "only_retired", - Arc::new(ProfiledModel::with_status(ModelStatus::Retired)), - ); - - assert!( - registry.resolve(ModelSelection::default()).is_none(), - "a registry with only a retired default resolves to None" - ); - - let allowed = registry - .resolve(ModelSelection { - allow_retired: true, - ..ModelSelection::default() - }) - .expect("allow_retired admits the retired default"); - assert_eq!(allowed.resolved.name, "only_retired"); -} diff --git a/crates/tinyagents-integration-tests/tests/e2e_registry_observability_contracts.rs b/crates/tinyagents-integration-tests/tests/e2e_registry_observability_contracts.rs index b49558159..3c6cace08 100644 --- a/crates/tinyagents-integration-tests/tests/e2e_registry_observability_contracts.rs +++ b/crates/tinyagents-integration-tests/tests/e2e_registry_observability_contracts.rs @@ -94,10 +94,8 @@ async fn capability_registry_resolves_aliases_and_hands_off_runtime_registries() assert!(registry.has(ComponentKind::Graph, "nb")); - let mut model_registry = registry.to_model_registry(); - model_registry.set_default("fast"); - assert_eq!(model_registry.names(), vec!["fast", "primary"]); - assert!(model_registry.get("fast").is_some()); + assert!(registry.model("primary").is_some()); + assert!(registry.model("fast").is_some()); let tool_registry = registry.to_tool_registry::<()>(); assert_eq!(tool_registry.names(), vec!["lookup"]); diff --git a/crates/tinyagents-integration-tests/tests/e2e_steering.rs b/crates/tinyagents-integration-tests/tests/e2e_steering.rs index 17ca9b320..a216103ec 100644 --- a/crates/tinyagents-integration-tests/tests/e2e_steering.rs +++ b/crates/tinyagents-integration-tests/tests/e2e_steering.rs @@ -7,7 +7,6 @@ use std::sync::Arc; -use tinyagents_graph::*; use tinyagents_harness::context::{RunConfig, RunContext}; use tinyagents_harness::runtime::AgentHarness; use tinyagents_harness::testkit::{EventRecorder, ScriptedModel, Trajectory}; @@ -24,8 +23,9 @@ use tinyinference_llm::message::Message; fn harness_with_scripted(replies: Vec<&str>) -> (AgentHarness<(), ()>, Arc<ScriptedModel>) { let model = Arc::new(ScriptedModel::replies(replies)); let mut harness = AgentHarness::<(), ()>::new(); - harness.register_model("scripted", model.clone()); - harness.set_default_model("scripted"); + harness.with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named("scripted", model.clone()), + )); (harness, model) } diff --git a/crates/tinyagents-integration-tests/tests/e2e_stream_error_code.rs b/crates/tinyagents-integration-tests/tests/e2e_stream_error_code.rs index 7edd2f6ff..a55b32fac 100644 --- a/crates/tinyagents-integration-tests/tests/e2e_stream_error_code.rs +++ b/crates/tinyagents-integration-tests/tests/e2e_stream_error_code.rs @@ -89,9 +89,9 @@ async fn streamed_run_against(server: &StreamErrorServer) -> tinyagents_harness: .with_base_url(&server.base_url) .with_model("test-model"); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness - .register_model("remote", Arc::new(model)) - .set_default_model("remote"); + harness.with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named("remote", Arc::new(model)), + )); harness .invoke_streaming_default(&(), vec![Message::user("go")]) .await diff --git a/crates/tinyagents-integration-tests/tests/e2e_streaming_cancel.rs b/crates/tinyagents-integration-tests/tests/e2e_streaming_cancel.rs index 8daab2848..782104e9c 100644 --- a/crates/tinyagents-integration-tests/tests/e2e_streaming_cancel.rs +++ b/crates/tinyagents-integration-tests/tests/e2e_streaming_cancel.rs @@ -50,10 +50,12 @@ async fn streaming_invoke_fires_on_model_delta_and_accumulates_final_response() let texts = Arc::new(Mutex::new(Vec::new())); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( - "stream", - Arc::new(StreamingMock::from_text_chunks(["Hel", "lo, ", "world"])), - ); + harness.with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named( + "stream", + Arc::new(StreamingMock::from_text_chunks(["Hel", "lo, ", "world"])), + ), + )); harness.push_middleware(Arc::new(DeltaRecorder { texts: texts.clone(), })); @@ -89,10 +91,12 @@ async fn streaming_invoke_fires_on_model_delta_and_accumulates_final_response() #[tokio::test] async fn pre_cancelled_token_yields_cancelled_on_streaming_path() { let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( - "stream", - Arc::new(StreamingMock::from_text_chunks(["never", "delivered"])), - ); + harness.with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named( + "stream", + Arc::new(StreamingMock::from_text_chunks(["never", "delivered"])), + ), + )); // Pre-cancel the token before the run starts. let token = CancellationToken::new(); diff --git a/crates/tinyagents-integration-tests/tests/e2e_subconscious_loop_example.rs b/crates/tinyagents-integration-tests/tests/e2e_subconscious_loop_example.rs deleted file mode 100644 index 4577ec02b..000000000 --- a/crates/tinyagents-integration-tests/tests/e2e_subconscious_loop_example.rs +++ /dev/null @@ -1,119 +0,0 @@ -//! Integration coverage for the runnable subconscious-loop example. - -#[path = "../examples/subconscious_loop/autonomous_loop.rs"] -mod autonomous_loop; - -use autonomous_loop::{ - SystemState, WorldDiff, build_subconscious_loop_graph, run_subconscious_loop, -}; - -#[tokio::test] -async fn normal_loop_retrieves_memory_and_holds_single_diff() { - let graph = build_subconscious_loop_graph().expect("graph compiles"); - let run = graph - .run(SystemState::new( - "web", - "Shift the resource matrix by a small amount.", - )) - .await - .expect("graph runs"); - - let visited: Vec<_> = run.visited.iter().map(|node| node.as_str()).collect(); - assert_eq!( - visited, - vec![ - "channel_ingestion", - "frontend_agent", - "agent_execution", - "summarization_gate", - "frontend_agent", - "context_manager_hook", - ] - ); - assert!(run.state.channel_response.contains("Completed. Output")); - assert_eq!(run.state.retrieved_context.len(), 1); - assert_eq!(run.state.sequential_diffs.len(), 1); - assert!(run.state.gated_world_summary.is_none()); - assert!(!run.state.trigger_subconscious); - assert!(run.state.subconscious_steering.is_empty()); - - let helper_state = run_subconscious_loop(SystemState::new("web", "Routine helper smoke.")) - .await - .expect("helper runs"); - assert!(helper_state.channel_response.contains("Completed. Output")); -} - -#[tokio::test] -async fn reasoning_escalation_forces_gate_and_resets_trigger() { - let graph = build_subconscious_loop_graph().expect("graph compiles"); - let run = graph - .run(SystemState::new( - "telegram", - "CRITICAL cascading sub-agent failure during resource reallocation.", - )) - .await - .expect("graph runs"); - - let visited: Vec<_> = run.visited.iter().map(|node| node.as_str()).collect(); - assert!(visited.contains(&"subconscious_eval")); - assert!(run.state.sequential_diffs.is_empty()); - assert!(run.state.gated_world_summary.is_none()); - assert!(!run.state.trigger_subconscious); - assert!( - run.state - .subconscious_steering - .contains("Lower sub-agent temperature") - ); - assert!( - run.state - .semantic_history - .iter() - .any(|trace| trace.contains("error_code=SUBAGENT_CASCADE")) - ); -} - -#[tokio::test] -async fn sequential_diff_threshold_routes_through_subconscious() { - let graph = build_subconscious_loop_graph().expect("graph compiles"); - let mut state = SystemState::new("web", "Shift the resource matrix again."); - state.sequential_diffs = vec![ - WorldDiff::new("resource_shift", 9, "stable"), - WorldDiff::new("resource_shift", 11, "stable"), - ]; - - let run = graph.run(state).await.expect("graph runs"); - - let visited: Vec<_> = run.visited.iter().map(|node| node.as_str()).collect(); - assert!(visited.contains(&"subconscious_eval")); - assert!(run.state.sequential_diffs.is_empty()); - assert!(run.state.gated_world_summary.is_none()); - assert_eq!( - run.state.subconscious_steering, - "STEERING_DIRECTIVE: System stable. Preserve current execution policy." - ); - assert!( - run.state - .event_log - .iter() - .any(|entry| entry.contains("Aggregated 3 operational shifts")) - ); -} - -#[tokio::test] -async fn context_manager_evicts_semantic_history_to_long_term_memory() { - let graph = build_subconscious_loop_graph().expect("graph compiles"); - let mut state = SystemState::new("web", "Perform routine work."); - state.context_utilization = 0.84; - let initial_memory_count = state.long_term_memory.len(); - - let run = graph.run(state).await.expect("graph runs"); - - let visited: Vec<_> = run.visited.iter().map(|node| node.as_str()).collect(); - assert!(!visited.contains(&"subconscious_eval")); - assert_eq!(run.state.context_utilization, 0.2); - assert_eq!( - run.state.semantic_history, - vec!["--- semantic history evicted to vector DB ---".to_string()] - ); - assert!(run.state.long_term_memory.len() > initial_memory_count); -} diff --git a/crates/tinyagents-integration-tests/tests/e2e_tool_dialects.rs b/crates/tinyagents-integration-tests/tests/e2e_tool_dialects.rs index 26fa75cbd..649113e37 100644 --- a/crates/tinyagents-integration-tests/tests/e2e_tool_dialects.rs +++ b/crates/tinyagents-integration-tests/tests/e2e_tool_dialects.rs @@ -15,6 +15,7 @@ use serde_json::json; use tinyagents_harness::config::ToolDispatcher; use tinyagents_harness::context::RunContext; use tinyagents_harness::events::{AgentEvent, RecordingListener}; +use tinyagents_harness::host::FixedModelResolver; use tinyagents_harness::middleware::{ Middleware, MiddlewareModelOutcome, ModelHandler, ModelMiddleware, }; @@ -101,8 +102,9 @@ fn harness_with( ) -> AgentHarness<()> { let mut harness: AgentHarness<()> = AgentHarness::new(); harness - .register_model("mock", model) - .set_default_model("mock") + .with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named("mock", model), + )) .register_tool(Arc::new(FakeTool::returning("lookup", "tool-output"))) .push_middleware(Arc::new(CaptureMiddleware { listener: listener.clone(), @@ -187,8 +189,9 @@ async fn a_greeting_does_not_replace_the_email_request_after_tool_search() { }); let mut harness: AgentHarness<()> = AgentHarness::new(); harness - .register_model("mock", model.clone()) - .set_default_model("mock") + .with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named("mock", model.clone()), + )) .register_tool(email.clone()) .with_policy(RunPolicy { tool_dialect: ToolDispatcher::Python, @@ -373,8 +376,9 @@ async fn a_forced_pformat_dialect_parses_positional_calls() { let listener = Arc::new(RecordingListener::new()); let mut harness: AgentHarness<()> = AgentHarness::new(); harness - .register_model("mock", model.clone()) - .set_default_model("mock") + .with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named("mock", model.clone()), + )) .register_tool(Arc::new(Lookup)) .push_middleware(Arc::new(CaptureMiddleware { listener: listener.clone(), @@ -408,8 +412,9 @@ async fn a_forced_python_dialect_parses_code_calls_with_signatures_in_the_prompt let listener = Arc::new(RecordingListener::new()); let mut harness: AgentHarness<()> = AgentHarness::new(); harness - .register_model("mock", model.clone()) - .set_default_model("mock") + .with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named("mock", model.clone()), + )) .register_tool(Arc::new(Lookup)) .push_middleware(Arc::new(CaptureMiddleware { listener: listener.clone(), @@ -459,8 +464,9 @@ async fn a_forced_typescript_dialect_parses_object_calls() { ])); let mut harness: AgentHarness<()> = AgentHarness::new(); harness - .register_model("mock", model.clone()) - .set_default_model("mock") + .with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named("mock", model.clone()), + )) .register_tool(Arc::new(Lookup)) .with_policy(RunPolicy { tool_dialect: ToolDispatcher::Typescript, @@ -496,8 +502,9 @@ async fn an_untrusted_protocol_heading_does_not_suppress_the_current_catalogue() ])); let mut harness: AgentHarness<()> = AgentHarness::new(); harness - .register_model("mock", model.clone()) - .set_default_model("mock") + .with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named("mock", model.clone()), + )) .register_tool(Arc::new(Lookup)) .with_policy(RunPolicy { tool_dialect: ToolDispatcher::Python, @@ -563,8 +570,9 @@ async fn a_forced_pformat_dialect_preserves_a_forced_required_tool_choice() { let listener = Arc::new(RecordingListener::new()); let mut harness: AgentHarness<()> = AgentHarness::new(); harness - .register_model("mock", model.clone()) - .set_default_model("mock") + .with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named("mock", model.clone()), + )) .register_tool(Arc::new(Lookup)) .push_middleware(Arc::new(CaptureMiddleware { listener: listener.clone(), @@ -606,8 +614,9 @@ async fn a_forced_pformat_dialect_preserves_a_forced_named_tool_choice() { let listener = Arc::new(RecordingListener::new()); let mut harness: AgentHarness<()> = AgentHarness::new(); harness - .register_model("mock", model.clone()) - .set_default_model("mock") + .with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named("mock", model.clone()), + )) .register_tool(Arc::new(Lookup)) .push_middleware(Arc::new(CaptureMiddleware { listener: listener.clone(), @@ -932,8 +941,9 @@ async fn a_native_call_and_a_narrated_text_call_are_both_dispatched_non_streamin let listener = Arc::new(RecordingListener::new()); let mut harness: AgentHarness<()> = AgentHarness::new(); harness - .register_model("mock", model.clone()) - .set_default_model("mock") + .with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named("mock", model.clone()), + )) .register_tool(Arc::new(FakeTool::returning("lookup", "lookup-output"))) .register_tool(Arc::new(FakeTool::returning("second", "second-output"))) .push_middleware(Arc::new(CaptureMiddleware { @@ -983,8 +993,9 @@ async fn a_native_call_and_a_narrated_text_call_are_both_dispatched_when_streame let listener = Arc::new(RecordingListener::new()); let mut harness: AgentHarness<()> = AgentHarness::new(); harness - .register_model("mock", model.clone()) - .set_default_model("mock") + .with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named("mock", model.clone()), + )) .register_tool(Arc::new(FakeTool::returning("lookup", "lookup-output"))) .register_tool(Arc::new(FakeTool::returning("second", "second-output"))) .push_middleware(Arc::new(CaptureMiddleware { @@ -1263,8 +1274,9 @@ async fn dropped_call_nudge_budget_resets_after_a_mixed_structured_and_tool_turn let listener = Arc::new(RecordingListener::new()); let mut harness: AgentHarness<()> = AgentHarness::new(); harness - .register_model("mock", model.clone()) - .set_default_model("mock") + .with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named("mock", model.clone()), + )) .register_tool(Arc::new(FakeTool::returning("lookup", "tool-output"))) .push_middleware(Arc::new(CaptureMiddleware { listener: listener.clone(), @@ -1324,8 +1336,9 @@ async fn auto_dialect_falls_back_to_xml_for_a_model_that_cannot_make_native_tool let listener = Arc::new(RecordingListener::new()); let mut harness: AgentHarness<()> = AgentHarness::new(); harness - .register_model("mock", model.clone()) - .set_default_model("mock") + .with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named("mock", model.clone()), + )) .register_tool(Arc::new(FakeTool::returning("lookup", "tool-output"))) .push_middleware(Arc::new(CaptureMiddleware { listener: listener.clone(), @@ -1820,8 +1833,12 @@ async fn python_dialect_dispatch(text: &str) -> (Vec<serde_json::Value>, Vec<ser let search_calls = Arc::new(Mutex::new(Vec::new())); let mut harness: AgentHarness<()> = AgentHarness::new(); harness - .register_model("mock", Arc::new(ScriptedModel::replies(vec![text, "done"]))) - .set_default_model("mock") + .with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named( + "mock", + Arc::new(ScriptedModel::replies(vec![text, "done"])), + ), + )) .register_tool(Arc::new(RecordingTool { name: "todo", schema: json!({ @@ -1913,8 +1930,7 @@ async fn a_streamed_native_delta_keeps_the_narrated_call() { let model = Arc::new(StreamingMock::new(items)); let mut harness: AgentHarness<()> = AgentHarness::new(); harness - .register_model("mock", model.clone()) - .set_default_model("mock") + .with_model_resolver(Arc::new(FixedModelResolver::new(model.clone()))) .register_tool(Arc::new(FakeTool::returning("lookup", "lookup-output"))) .register_tool(Arc::new(FakeTool::returning("second", "second-output"))) .with_policy(RunPolicy { diff --git a/crates/tinyagents-integration-tests/tests/e2e_tool_policy.rs b/crates/tinyagents-integration-tests/tests/e2e_tool_policy.rs index afa08db0c..bd93b706e 100644 --- a/crates/tinyagents-integration-tests/tests/e2e_tool_policy.rs +++ b/crates/tinyagents-integration-tests/tests/e2e_tool_policy.rs @@ -71,7 +71,9 @@ async fn strict_policy_exposes_only_the_classified_safe_tool() { "finished", )])); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", scripted.clone()); + harness.with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named("mock", scripted.clone()), + )); harness.register_tool(Arc::new(SafeTool)); harness.register_tool(Arc::new(PaymentTool)); harness.push_middleware(Arc::new(ToolPolicyMiddleware::strict( diff --git a/crates/tinyagents-integration-tests/tests/e2e_unknown_tool_policy.rs b/crates/tinyagents-integration-tests/tests/e2e_unknown_tool_policy.rs index 6f75740f1..e5acf7d81 100644 --- a/crates/tinyagents-integration-tests/tests/e2e_unknown_tool_policy.rs +++ b/crates/tinyagents-integration-tests/tests/e2e_unknown_tool_policy.rs @@ -62,10 +62,12 @@ fn single_unknown_tool_event(events: &[AgentEvent]) -> (String, String) { #[tokio::test] async fn fail_policy_errors_on_unregistered_tool() { let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( - "mock", - Arc::new(MockModel::with_tool_call("missing", json!({}))), - ); + harness.with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named( + "mock", + Arc::new(MockModel::with_tool_call("missing", json!({}))), + ), + )); // `ReturnToolError` is the default now (a hallucinated tool name is a // routine model mistake), so a hard stop has to be asked for. harness.with_policy(RunPolicy { @@ -89,13 +91,15 @@ async fn fail_policy_errors_on_unregistered_tool() { #[tokio::test] async fn return_tool_error_recovers_and_emits_event() { let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( - "mock", - Arc::new(MockModel::with_responses(vec![ - tool_call_response("c1", "missing", json!({})), - text_response("recovered"), - ])), - ); + harness.with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named( + "mock", + Arc::new(MockModel::with_responses(vec![ + tool_call_response("c1", "missing", json!({})), + text_response("recovered"), + ])), + ), + )); harness.with_policy(RunPolicy { unknown_tool: UnknownToolPolicy::ReturnToolError, ..RunPolicy::default() @@ -131,13 +135,15 @@ async fn return_tool_error_preserves_original_arguments() { let original = json!({ "query": "weather", "n": 3 }); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( - "mock", - Arc::new(MockModel::with_responses(vec![ - tool_call_response("c1", "search", original.clone()), - text_response("recovered"), - ])), - ); + harness.with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named( + "mock", + Arc::new(MockModel::with_responses(vec![ + tool_call_response("c1", "search", original.clone()), + text_response("recovered"), + ])), + ), + )); harness.with_policy(RunPolicy { unknown_tool: UnknownToolPolicy::ReturnToolError, ..RunPolicy::default() @@ -173,13 +179,15 @@ async fn rewrite_retargets_to_real_tool() { let fake = Arc::new(FakeTool::returning("lookup", "out")); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( - "mock", - Arc::new(MockModel::with_responses(vec![ - tool_call_response("c1", "missing", json!({})), - text_response("done"), - ])), - ); + harness.with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named( + "mock", + Arc::new(MockModel::with_responses(vec![ + tool_call_response("c1", "missing", json!({})), + text_response("done"), + ])), + ), + )); harness.register_tool(fake.clone()); harness.with_policy(RunPolicy { unknown_tool: UnknownToolPolicy::Rewrite { @@ -210,13 +218,15 @@ async fn rewrite_retargets_to_real_tool() { #[tokio::test] async fn rewrite_to_missing_target_falls_back_to_tool_error() { let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( - "mock", - Arc::new(MockModel::with_responses(vec![ - tool_call_response("c1", "missing", json!({})), - text_response("recovered"), - ])), - ); + harness.with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named( + "mock", + Arc::new(MockModel::with_responses(vec![ + tool_call_response("c1", "missing", json!({})), + text_response("recovered"), + ])), + ), + )); // Rewrite target "nope" is itself unregistered → fall back to tool-error. harness.with_policy(RunPolicy { unknown_tool: UnknownToolPolicy::Rewrite { @@ -259,10 +269,12 @@ async fn recovery_is_bounded_by_tool_call_limit() { // MockModel::with_tool_call repeats the same unknown call on every // invocation, so without a bound the loop would spin forever. let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( - "mock", - Arc::new(MockModel::with_tool_call("missing", json!({}))), - ); + harness.with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named( + "mock", + Arc::new(MockModel::with_tool_call("missing", json!({}))), + ), + )); harness.with_policy(RunPolicy { unknown_tool: UnknownToolPolicy::ReturnToolError, limits: RunLimits::default().with_max_tool_calls(3), diff --git a/crates/tinyagents-integration-tests/tests/feature_graph_fanout.rs b/crates/tinyagents-integration-tests/tests/feature_graph_fanout.rs deleted file mode 100644 index 6cdc69d1c..000000000 --- a/crates/tinyagents-integration-tests/tests/feature_graph_fanout.rs +++ /dev/null @@ -1,143 +0,0 @@ -//! Feature tests for graph-level `Send` fan-out joined through a state reducer. -//! -//! Where `tests/e2e_complex_graph.rs` fans out with a static `goto` to two -//! *named* branch nodes, this covers the dynamic `Send` primitive: one worker -//! node scheduled once per argument, each activation receiving its own -//! `send_arg`, with the per-branch updates folded back into committed state by -//! the reducer in deterministic branch order — under both sequential and -//! bounded-parallel execution — and (in parallel mode) a distinct fork identity -//! per branch. - -use std::sync::{Arc, Mutex}; - -use serde_json::json; - -use tinyagents_graph::*; -use tinyagents_harness::*; -use tinyagents_registry::*; - -/// Committed state: the append-only list of squared work items. -#[derive(Clone, Debug, Default, PartialEq, Eq)] -struct Squares { - values: Vec<i64>, -} - -/// Records each worker's `(fork branch index, arg)`; the fork is `None` in -/// sequential mode and `Some(_)` under parallel execution. -type ForkLog = Arc<Mutex<Vec<(Option<usize>, i64)>>>; - -/// Builds a `spread -> square (xN) -> collect` graph. `spread` fans out one -/// `Send` per argument; each `square` activation squares its own `send_arg` and -/// records the fork it ran on. `collect` is a join that runs once after every -/// branch has folded in. -fn fanout_graph( - parallel: bool, - max_concurrency: Option<usize>, - forks: ForkLog, -) -> tinyagents_graph::CompiledGraph<Squares, i64> { - let mut builder = GraphBuilder::<Squares, i64>::new() - .set_reducer(ClosureStateReducer::new(|mut s: Squares, u: i64| { - s.values.push(u); - Ok(s) - })) - .with_parallel(parallel); - if let Some(cap) = max_concurrency { - builder = builder.with_max_concurrency(cap); - } - builder - .add_node( - "spread", - fanout_node("square", [json!(1), json!(2), json!(3), json!(4)]), - ) - .add_node("square", move |_s: Squares, ctx: NodeContext| { - let forks = forks.clone(); - async move { - let arg = ctx - .send_arg - .as_ref() - .and_then(|v| v.as_i64()) - .expect("each worker carries a send arg"); - let branch = ctx.fork.as_ref().map(|f| f.branch_index); - forks.lock().unwrap().push((branch, arg)); - Ok(NodeResult::Update(arg * arg)) - } - }) - .add_node("collect", tinyagents_graph::noop_node()) - .set_entry("spread") - .add_edge("square", "collect") - .set_finish("collect") - .compile() - .expect("fan-out graph compiles") -} - -#[tokio::test] -async fn send_fanout_folds_each_branch_update_through_the_reducer() { - let forks: ForkLog = Arc::new(Mutex::new(Vec::new())); - let run = run_recorded( - &fanout_graph(false, None, forks.clone()), - None, - Squares::default(), - ) - .await - .expect("run succeeds"); - - // The reducer folds the four branch updates in deterministic branch order. - assert_eq!(run.execution.state.values, vec![1, 4, 9, 16]); - - // The worker was scheduled once per argument, each with its own send arg. - let square_visits = run - .execution - .visited - .iter() - .filter(|n| n.as_str() == "square") - .count(); - assert_eq!(square_visits, 4); - - let mut args: Vec<i64> = forks.lock().unwrap().iter().map(|(_, a)| *a).collect(); - args.sort_unstable(); - assert_eq!(args, vec![1, 2, 3, 4]); -} - -#[tokio::test] -async fn each_parallel_branch_receives_a_distinct_fork_identity() { - let forks: ForkLog = Arc::new(Mutex::new(Vec::new())); - let _ = run_recorded( - &fanout_graph(true, None, forks.clone()), - None, - Squares::default(), - ) - .await - .expect("run succeeds"); - - let mut recorded = forks.lock().unwrap().clone(); - recorded.sort_by_key(|(branch, _)| branch.unwrap_or(usize::MAX)); - // Four distinct branch indices (0..4), each paired with its own send arg. - assert_eq!( - recorded, - vec![(Some(0), 1), (Some(1), 2), (Some(2), 3), (Some(3), 4)] - ); -} - -#[tokio::test] -async fn parallel_fanout_is_deterministic_under_a_bounded_concurrency_cap() { - let forks: ForkLog = Arc::new(Mutex::new(Vec::new())); - // Run in parallel with at most two branches in flight at once. - let graph = fanout_graph(true, Some(2), forks.clone()); - - let run = run_recorded(&graph, None, Squares::default()) - .await - .expect("parallel run succeeds"); - - // Reducer application stays in stable branch order regardless of the - // parallel completion order, so the joined state matches the sequential run. - assert_eq!(run.execution.state.values, vec![1, 4, 9, 16]); - - // The join runs exactly once after all branches drained. - let collect_visits = run - .execution - .visited - .iter() - .filter(|n| n.as_str() == "collect") - .count(); - assert_eq!(collect_visits, 1); -} diff --git a/crates/tinyagents-integration-tests/tests/feature_graph_goal_budget.rs b/crates/tinyagents-integration-tests/tests/feature_graph_goal_budget.rs deleted file mode 100644 index 7848eb325..000000000 --- a/crates/tinyagents-integration-tests/tests/feature_graph_goal_budget.rs +++ /dev/null @@ -1,44 +0,0 @@ -//! Public coverage for goal accounting and budget guard preparation. - -use std::sync::Arc; -use std::time::Duration; - -use tinyagents_harness::store::{InMemoryStore, Store}; -use tinyagents_harness_tools::goals::{BudgetVerdict, GoalStore, ThreadGoalStatus}; - -#[tokio::test] -async fn prepared_budget_guard_tracks_atomic_accounting_and_limited_status() { - let store: Arc<dyn Store> = Arc::new(InMemoryStore::default()); - let goals = GoalStore::new(store); - goals.set("thread", "finish", Some(10)).await.unwrap(); - let prepared = goals.prepare_turn("thread").await.unwrap(); - let guard = prepared.budget_guard.unwrap(); - assert_eq!( - guard.check(&goals.backing_store(), 9).await.unwrap(), - BudgetVerdict::Continue - ); - - let update = goals - .account_turn("thread", 6, 0, Duration::from_millis(500)) - .await - .unwrap() - .unwrap(); - assert_eq!(update.tokens_used, 6); - assert_eq!(update.status, ThreadGoalStatus::Active); - assert!( - guard - .check(&goals.backing_store(), 4) - .await - .unwrap() - .is_stop() - ); - - let limited = goals - .account_turn("thread", 4, 0, Duration::from_millis(500)) - .await - .unwrap() - .unwrap(); - assert_eq!(limited.status, ThreadGoalStatus::BudgetLimited); - assert!(limited.over_budget()); - assert_eq!(limited.time_used_subsec_nanos, 0); -} diff --git a/crates/tinyagents-integration-tests/tests/feature_graph_parallel.rs b/crates/tinyagents-integration-tests/tests/feature_graph_parallel.rs deleted file mode 100644 index 65d799205..000000000 --- a/crates/tinyagents-integration-tests/tests/feature_graph_parallel.rs +++ /dev/null @@ -1,204 +0,0 @@ -//! Feature tests for `graph::parallel::map_reduce` — the reusable, bounded -//! concurrency map/reduce helper independent of the graph executor. -//! -//! This helper had no integration coverage: these tests pin down input-order -//! result determinism despite out-of-order completion, the concurrency cap, the -//! per-item timeout, and every [`FailurePolicy`] variant. - -use std::sync::Arc; -use std::sync::atomic::{AtomicUsize, Ordering}; -use std::time::Duration; - -use tinyagents_graph::*; -use tinyagents_harness::*; -use tinyagents_registry::*; - -#[tokio::test] -async fn results_are_returned_in_input_order_despite_out_of_order_completion() { - // Larger inputs finish *sooner* (shorter sleep), so completion order is the - // reverse of input order; the outcome must still be input-ordered. - let items = vec![1_u64, 2, 3, 4, 5]; - let outcome = map_reduce( - items, - ParallelOptions::default().with_failure_policy(FailurePolicy::CollectAll), - |_index, item| async move { - let delay = 50u64.saturating_sub(item * 5); - tokio::time::sleep(Duration::from_millis(delay)).await; - Ok::<u64, TinyAgentsError>(item * item) - }, - ) - .await - .expect("collect-all never errors"); - - let indices: Vec<usize> = outcome.outcomes.iter().map(|o| o.index).collect(); - assert_eq!(indices, vec![0, 1, 2, 3, 4]); - assert_eq!(outcome.into_successes(), vec![1, 4, 9, 16, 25]); -} - -#[tokio::test] -async fn bounded_concurrency_never_exceeds_the_configured_cap() { - let in_flight = Arc::new(AtomicUsize::new(0)); - let peak = Arc::new(AtomicUsize::new(0)); - - let items: Vec<u64> = (0..8).collect(); - let outcome = map_reduce( - items, - ParallelOptions::default().with_max_concurrency(3), - |_index, _item| { - let in_flight = in_flight.clone(); - let peak = peak.clone(); - async move { - let now = in_flight.fetch_add(1, Ordering::SeqCst) + 1; - peak.fetch_max(now, Ordering::SeqCst); - tokio::time::sleep(Duration::from_millis(10)).await; - in_flight.fetch_sub(1, Ordering::SeqCst); - Ok::<u64, TinyAgentsError>(0) - } - }, - ) - .await - .expect("collect-all never errors"); - - assert_eq!(outcome.success_count(), 8); - assert!( - peak.load(Ordering::SeqCst) <= 3, - "observed {} concurrent items, cap was 3", - peak.load(Ordering::SeqCst) - ); -} - -#[tokio::test] -async fn fail_fast_returns_the_first_error_in_input_order() { - let attempts = Arc::new(AtomicUsize::new(0)); - let items = vec![1_u64, 2, 3, 4]; - let err = map_reduce( - items, - ParallelOptions::default() - .with_max_concurrency(1) - .with_failure_policy(FailurePolicy::FailFast), - { - let attempts = attempts.clone(); - move |_index, item| { - let attempts = attempts.clone(); - async move { - attempts.fetch_add(1, Ordering::SeqCst); - if item == 2 { - Err(TinyAgentsError::Graph("item two blew up".to_string())) - } else { - Ok(item) - } - } - } - }, - ) - .await - .expect_err("fail-fast surfaces the first item error"); - - assert!(err.to_string().contains("item two blew up"), "got: {err}"); - // With concurrency 1 the batch stops after the failing item, never reaching 4. - assert!(attempts.load(Ordering::SeqCst) <= 2); -} - -#[tokio::test] -async fn collect_all_records_each_success_and_failure() { - let items = vec![1_u64, 2, 3, 4, 5]; - let outcome = map_reduce( - items, - ParallelOptions::default().with_failure_policy(FailurePolicy::CollectAll), - |_index, item| async move { - if item % 2 == 0 { - Err(TinyAgentsError::Graph(format!("even {item}"))) - } else { - Ok(item) - } - }, - ) - .await - .expect("collect-all always returns Ok"); - - assert_eq!(outcome.outcomes.len(), 5); - assert_eq!(outcome.success_count(), 3); - assert_eq!(outcome.failure_count(), 2); - assert_eq!(outcome.successes(), vec![&1, &3, &5]); - // Failures preserve their input index for correlation. - let failed_indices: Vec<usize> = outcome - .outcomes - .iter() - .filter(|o| !o.is_ok()) - .map(|o| o.index) - .collect(); - assert_eq!(failed_indices, vec![1, 3]); -} - -#[tokio::test] -async fn quorum_succeeds_at_threshold_and_errors_below_it() { - let run = |threshold: usize, fail_from: u64| async move { - let items = vec![1_u64, 2, 3, 4]; - map_reduce( - items, - ParallelOptions::default().with_failure_policy(FailurePolicy::Quorum(threshold)), - move |_index, item| async move { - if item >= fail_from { - Err(TinyAgentsError::Graph(format!("fail {item}"))) - } else { - Ok(item) - } - }, - ) - .await - }; - - // 3 of 4 succeed (items 1,2,3); a quorum of 3 is met. - let met = run(3, 4).await.expect("quorum of 3 is satisfied"); - assert_eq!(met.success_count(), 3); - - // Only 1 of 4 succeeds (item 1); a quorum of 3 is not met. - let unmet = run(3, 2).await.expect_err("quorum of 3 is not satisfied"); - assert!(matches!(unmet, TinyAgentsError::Graph(_)), "got: {unmet:?}"); -} - -#[tokio::test] -async fn best_effort_keeps_only_the_successful_outputs() { - let items = vec![1_u64, 2, 3, 4, 5, 6]; - let outcome = map_reduce( - items, - ParallelOptions::default().with_failure_policy(FailurePolicy::BestEffort), - |_index, item| async move { - if item % 3 == 0 { - Err(TinyAgentsError::Graph(format!("mult of three {item}"))) - } else { - Ok(item) - } - }, - ) - .await - .expect("best-effort always returns Ok"); - - // Only successes are retained, in input order; failures are dropped entirely. - assert_eq!(outcome.failure_count(), 0); - assert_eq!(outcome.into_successes(), vec![1, 2, 4, 5]); -} - -#[tokio::test] -async fn per_item_timeout_turns_a_slow_item_into_a_recoverable_failure() { - let items = vec![10_u64, 200, 30]; - let outcome = map_reduce( - items, - ParallelOptions::default() - .with_failure_policy(FailurePolicy::CollectAll) - .with_item_timeout(Duration::from_millis(60)), - |_index, item| async move { - tokio::time::sleep(Duration::from_millis(item)).await; - Ok::<u64, TinyAgentsError>(item) - }, - ) - .await - .expect("collect-all absorbs the timeout as a per-item failure"); - - // The 200ms item exceeds the 60ms bound and is recorded as a failure; the - // fast items still succeed. - assert_eq!(outcome.success_count(), 2); - assert_eq!(outcome.failure_count(), 1); - let failed = outcome.outcomes.iter().find(|o| !o.is_ok()).unwrap(); - assert_eq!(failed.index, 1); -} diff --git a/crates/tinyagents-integration-tests/tests/feature_graph_routing.rs b/crates/tinyagents-integration-tests/tests/feature_graph_routing.rs deleted file mode 100644 index 73385cb51..000000000 --- a/crates/tinyagents-integration-tests/tests/feature_graph_routing.rs +++ /dev/null @@ -1,178 +0,0 @@ -//! Feature tests for the graph runtime's routing surface. -//! -//! These exercise the *executed* routing paths — conditional edges evaluated -//! against committed state at the superstep boundary and dynamic `Command` -//! `goto` routing — rather than the topology export. Complements -//! `tests/graph_durable.rs` (which only checks that an invalid command target -//! fails) and `tests/e2e_graph_export.rs` (which renders conditional edges -//! without running them). - -use tinyagents_graph::END; -use tinyagents_graph::*; -use tinyagents_harness::*; -use tinyagents_registry::*; - -/// A whole-state document that records the path of nodes it flowed through. -#[derive(Clone, Debug, Default, PartialEq, Eq)] -struct Doc { - n: i64, - path: Vec<String>, -} - -impl Doc { - fn seed(n: i64) -> Self { - Doc { - n, - path: Vec::new(), - } - } - - fn stamped(mut self, label: &str) -> Self { - self.path.push(label.to_string()); - self - } -} - -/// A classify node whose conditional router picks a terminal branch by parity, -/// routing the `zero` case straight to `END`. -fn parity_graph() -> tinyagents_graph::CompiledGraph<Doc, Doc> { - GraphBuilder::<Doc, Doc>::overwrite() - .add_node("classify", |s: Doc, _c: NodeContext| async move { - Ok(NodeResult::Update(s.stamped("classify"))) - }) - .add_node("even", |s: Doc, _c: NodeContext| async move { - Ok(NodeResult::Update(s.stamped("even"))) - }) - .add_node("odd", |s: Doc, _c: NodeContext| async move { - Ok(NodeResult::Update(s.stamped("odd"))) - }) - .set_entry("classify") - .add_conditional_edges( - "classify", - |d: &Doc| { - if d.n == 0 { - "zero" - } else if d.n % 2 == 0 { - "even" - } else { - "odd" - } - }, - [("even", "even"), ("odd", "odd"), ("zero", END)], - ) - .set_finish("even") - .set_finish("odd") - .compile() - .expect("parity graph compiles") -} - -#[tokio::test] -async fn conditional_router_selects_even_branch_by_committed_state() { - let run = run_recorded(&parity_graph(), None, Doc::seed(4)) - .await - .expect("run succeeds"); - - assert_graph(&run) - .visited(["classify", "even"]) - .routed("classify", "even") - .completed(); - assert_eq!(run.execution.state.path, vec!["classify", "even"]); -} - -#[tokio::test] -async fn conditional_router_selects_odd_branch_by_committed_state() { - let run = run_recorded(&parity_graph(), None, Doc::seed(7)) - .await - .expect("run succeeds"); - - assert_graph(&run) - .visited(["classify", "odd"]) - .routed("classify", "odd") - .completed(); - assert_eq!(run.execution.state.path, vec!["classify", "odd"]); -} - -#[tokio::test] -async fn conditional_router_can_route_directly_to_end() { - let run = run_recorded(&parity_graph(), None, Doc::seed(0)) - .await - .expect("run succeeds"); - - // The `zero` label resolves to END, so the run terminates at `classify` - // without visiting a branch node. - assert_graph(&run).visited(["classify"]).completed(); - assert_eq!(run.execution.state.path, vec!["classify"]); -} - -#[tokio::test] -async fn conditional_router_is_re_evaluated_each_loop_iteration() { - // A self-loop whose router terminates only once the counter crosses a - // threshold — the router runs at every superstep boundary. - let graph = GraphBuilder::<i64, i64>::overwrite() - .with_recursion_limit(16) - .add_node("inc", |s: i64, _c: NodeContext| async move { - Ok(NodeResult::Update(s + 1)) - }) - .set_entry("inc") - .add_conditional_edges( - "inc", - |n: &i64| if *n >= 3 { "done" } else { "again" }, - [("again", "inc"), ("done", END)], - ) - .compile() - .expect("loop graph compiles"); - - let run = run_recorded(&graph, None, 0).await.expect("run succeeds"); - - assert_graph(&run) - .visited(["inc", "inc", "inc"]) - .completed(); - assert_eq!(run.execution.state, 3); - assert_eq!(run.execution.steps, 3); -} - -#[tokio::test] -async fn command_goto_routes_dynamically_to_the_chosen_target() { - // A command-routing node that resolves its successor at runtime from the - // committed state, choosing between two declared destinations. - fn threshold_graph() -> tinyagents_graph::CompiledGraph<Doc, Doc> { - GraphBuilder::<Doc, Doc>::overwrite() - .add_node("route", |s: Doc, _c: NodeContext| async move { - let target = if s.n >= 10 { "high" } else { "low" }; - Ok(NodeResult::Command( - Command::goto([target]).with_update(s.stamped("route")), - )) - }) - .add_node("low", |s: Doc, _c: NodeContext| async move { - Ok(NodeResult::Update(s.stamped("low"))) - }) - .add_node("high", |s: Doc, _c: NodeContext| async move { - Ok(NodeResult::Update(s.stamped("high"))) - }) - .set_entry("route") - .with_command_destinations("route", ["low", "high"]) - .set_finish("low") - .set_finish("high") - .compile() - .expect("threshold graph compiles") - } - - let graph = threshold_graph(); - - let low_run = run_recorded(&graph, None, Doc::seed(3)) - .await - .expect("low run succeeds"); - assert_graph(&low_run) - .visited(["route", "low"]) - .routed("route", "low") - .completed(); - - let high_run = run_recorded(&graph, None, Doc::seed(42)) - .await - .expect("high run succeeds"); - assert_graph(&high_run) - .visited(["route", "high"]) - .routed("route", "high") - .completed(); - assert_eq!(high_run.execution.state.path, vec!["route", "high"]); -} diff --git a/crates/tinyagents-integration-tests/tests/feature_graph_streaming.rs b/crates/tinyagents-integration-tests/tests/feature_graph_streaming.rs deleted file mode 100644 index d351722d3..000000000 --- a/crates/tinyagents-integration-tests/tests/feature_graph_streaming.rs +++ /dev/null @@ -1,171 +0,0 @@ -//! Feature tests for the graph event stream and run-status snapshot. -//! -//! Exercises the live [`GraphEvent`] surface projected through -//! [`StreamCollector`] (node order, state writes, selected routes, checkpoint -//! boundaries, custom writes) and the compact [`GraphRunStatus`] snapshot on a -//! completed run. Complements `tests/e2e_graph_support_contracts.rs`, which -//! smoke-tests the testkit helpers, by asserting the concrete event lifecycle -//! and ordering. - -use std::sync::Arc; - -use tinyagents_graph::END; -use tinyagents_graph::*; -use tinyagents_harness::ids::ExecutionStatus; -use tinyagents_harness::*; -use tinyagents_registry::*; - -/// A linear two-node graph plus a conditional edge, so the stream carries both -/// static and conditional `RouteSelected` events. -fn pipeline() -> tinyagents_graph::CompiledGraph<i64, i64> { - GraphBuilder::<i64, i64>::overwrite() - .add_node("seed", |s: i64, _c: NodeContext| async move { - Ok(NodeResult::Update(s + 1)) - }) - .add_node("double", |s: i64, _c: NodeContext| async move { - Ok(NodeResult::Update(s * 2)) - }) - .add_node("finish", |s: i64, _c: NodeContext| async move { - Ok(NodeResult::Update(s + 100)) - }) - .set_entry("seed") - .add_edge("seed", "double") - .add_conditional_edges( - "double", - |n: &i64| if *n > 0 { "go" } else { "stop" }, - [("go", "finish"), ("stop", END)], - ) - .set_finish("finish") - .compile() - .expect("pipeline compiles") -} - -#[tokio::test] -async fn event_stream_reports_the_full_run_lifecycle() { - let run = run_recorded(&pipeline(), None, 0) - .await - .expect("run succeeds"); - - let kinds = run.events.iter().map(|e| e.kind()).collect::<Vec<_>>(); - - // The run brackets every step with start/complete, and every executed node - // reports started/completed with a state write. - assert_eq!(kinds.first(), Some(&"run.started")); - assert_eq!(kinds.last(), Some(&"run.completed")); - for expected in [ - "step.started", - "node.started", - "node.completed", - "state.updated", - "route.selected", - "step.completed", - ] { - assert!( - kinds.contains(&expected), - "expected a `{expected}` event, saw {kinds:?}" - ); - } -} - -#[tokio::test] -async fn stream_collector_projects_node_order_updates_and_routes() { - let run = run_recorded(&pipeline(), None, 0) - .await - .expect("run succeeds"); - let collector = StreamCollector::new(run.events.clone()); - - let order: Vec<String> = collector - .node_order() - .iter() - .map(|n| n.as_str().to_string()) - .collect(); - assert_eq!(order, vec!["seed", "double", "finish"]); - - let updates: Vec<String> = collector - .updates() - .iter() - .map(|n| n.as_str().to_string()) - .collect(); - assert_eq!(updates, vec!["seed", "double", "finish"]); - - // Both the static edge (seed -> double) and the conditional edge - // (double -> finish) surface as selected routes; the final -> END edge does - // not emit a route. - let routes: Vec<(String, String)> = collector - .routes() - .iter() - .map(|(f, t)| (f.as_str().to_string(), t.as_str().to_string())) - .collect(); - assert!(routes.contains(&("seed".to_string(), "double".to_string()))); - assert!(routes.contains(&("double".to_string(), "finish".to_string()))); - - // 0 -> +1 -> *2 -> +100. - assert_eq!(run.execution.state, 102); -} - -#[tokio::test] -async fn checkpoint_saved_events_track_every_superstep_boundary() { - let graph = pipeline().with_checkpointer(Arc::new(InMemoryCheckpointer::<i64>::new())); - - let run = run_recorded(&graph, Some("stream-thread"), 0) - .await - .expect("threaded run succeeds"); - - let collector = StreamCollector::new(run.events.clone()); - // One checkpoint per executed superstep (seed, double, finish). - assert_eq!(collector.checkpoint_count(), 3); - assert_eq!(run.execution.steps, 3); - // The recorded checkpoint history mirrors the boundary count. - assert_eq!(run.history.len(), 3); -} - -#[tokio::test] -async fn run_status_snapshot_reflects_a_completed_run() { - let run = run_recorded(&pipeline(), None, 0) - .await - .expect("run succeeds"); - let status = &run.execution.status; - - assert_eq!(status.status, ExecutionStatus::Completed); - assert!(status.is_terminal()); - assert_eq!(status.current_step, run.execution.steps); - assert!( - status.active_nodes.is_empty(), - "a completed run has no active nodes, saw {:?}", - status.active_nodes - ); - assert!(status.pending_interrupts.is_empty()); - assert!( - status.ended_at.is_some(), - "a terminal run records an end time" - ); - assert!(status.error.is_none()); -} - -#[tokio::test] -async fn a_failed_node_emits_run_failed_and_a_failed_status() { - let graph = GraphBuilder::<i64, i64>::overwrite() - .add_node("boom", tinyagents_graph::failing_node("kaboom")) - .set_entry("boom") - .set_finish("boom") - .compile() - .expect("graph compiles"); - - // Wire the recorder manually so we can inspect events even though the run - // returns an error. - let recorder = tinyagents_graph::GraphEventRecorder::new(); - let graph = graph.with_event_sink(recorder.sink()); - let err = graph.run(0).await.expect_err("failing node aborts the run"); - assert!(err.to_string().contains("kaboom"), "got: {err}"); - - let kinds = recorder.kinds(); - assert!(kinds.contains(&"node.failed".to_string()), "saw {kinds:?}"); - assert!(kinds.contains(&"run.failed".to_string()), "saw {kinds:?}"); - // No run.completed event is emitted for a failed run. - assert!( - !recorder - .events() - .iter() - .any(|e| matches!(e, GraphEvent::RunCompleted { .. })) - ); -} diff --git a/crates/tinyagents-integration-tests/tests/feature_harness_agent_loop.rs b/crates/tinyagents-integration-tests/tests/feature_harness_agent_loop.rs index 34cb09a82..94d2e0d1b 100644 --- a/crates/tinyagents-integration-tests/tests/feature_harness_agent_loop.rs +++ b/crates/tinyagents-integration-tests/tests/feature_harness_agent_loop.rs @@ -88,18 +88,19 @@ async fn parallel_tool_calls_in_one_turn_all_execute() { let mut harness: AgentHarness<()> = AgentHarness::new(); harness - .register_model( - "mock", - Arc::new(MockModel::with_responses(vec![ - multi_tool_call_response(vec![ - ToolCall::new("c1", "alpha", json!({ "n": 1 })), - ToolCall::new("c2", "beta", json!({ "n": 2 })), - ToolCall::new("c3", "alpha", json!({ "n": 3 })), - ]), - text_response("all done"), - ])), - ) - .set_default_model("mock") + .with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named( + "mock", + Arc::new(MockModel::with_responses(vec![ + multi_tool_call_response(vec![ + ToolCall::new("c1", "alpha", json!({ "n": 1 })), + ToolCall::new("c2", "beta", json!({ "n": 2 })), + ToolCall::new("c3", "alpha", json!({ "n": 3 })), + ]), + text_response("all done"), + ])), + ), + )) .register_tool(Arc::new(FakeTool::returning("alpha", "A"))) .register_tool(Arc::new(FakeTool::returning("beta", "B"))); @@ -137,11 +138,12 @@ async fn max_tool_calls_limit_returns_limit_exceeded() { let mut harness: AgentHarness<()> = AgentHarness::new(); // A model that always requests a tool would loop forever without a cap. harness - .register_model( - "mock", - Arc::new(MockModel::with_tool_call("spin", json!({}))), - ) - .set_default_model("mock") + .with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named( + "mock", + Arc::new(MockModel::with_tool_call("spin", json!({}))), + ), + )) .register_tool(Arc::new(FakeTool::returning("spin", "again"))) .with_policy(RunPolicy { // Generous model budget so the *tool* cap is what trips. @@ -171,14 +173,15 @@ async fn tool_round_trip_emits_lifecycle_events_in_order() { let mut harness: AgentHarness<()> = AgentHarness::new(); harness - .register_model( - "mock", - Arc::new(MockModel::with_responses(vec![ - tool_call_response("c1", "lookup", json!({})), - text_response("final"), - ])), - ) - .set_default_model("mock") + .with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named( + "mock", + Arc::new(MockModel::with_responses(vec![ + tool_call_response("c1", "lookup", json!({})), + text_response("final"), + ])), + ), + )) .register_tool(Arc::new(FakeTool::returning("lookup", "out"))); let ctx = RunContext::new(RunConfig::new("ordering"), ()).with_events(recorder.sink()); @@ -211,14 +214,15 @@ async fn tool_round_trip_emits_lifecycle_events_in_order() { async fn invoke_with_status_reports_completed_run() { let mut harness: AgentHarness<()> = AgentHarness::new(); harness - .register_model( - "mock", - Arc::new(MockModel::with_responses(vec![ - tool_call_response("c1", "lookup", json!({})), - text_response("done"), - ])), - ) - .set_default_model("mock") + .with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named( + "mock", + Arc::new(MockModel::with_responses(vec![ + tool_call_response("c1", "lookup", json!({})), + text_response("done"), + ])), + ), + )) .register_tool(Arc::new(FakeTool::returning("lookup", "out"))); let result = harness @@ -252,14 +256,15 @@ async fn tools_are_dispatched_by_name() { let mut harness: AgentHarness<()> = AgentHarness::new(); harness - .register_model( - "mock", - Arc::new(MockModel::with_responses(vec![ - tool_call_response("c1", "beta", json!({})), - text_response("ok"), - ])), - ) - .set_default_model("mock") + .with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named( + "mock", + Arc::new(MockModel::with_responses(vec![ + tool_call_response("c1", "beta", json!({})), + text_response("ok"), + ])), + ), + )) .register_tool(alpha.clone()) .register_tool(beta.clone()); diff --git a/crates/tinyagents-integration-tests/tests/feature_harness_structured.rs b/crates/tinyagents-integration-tests/tests/feature_harness_structured.rs index 7d6e34dee..8029d2dec 100644 --- a/crates/tinyagents-integration-tests/tests/feature_harness_structured.rs +++ b/crates/tinyagents-integration-tests/tests/feature_harness_structured.rs @@ -291,8 +291,9 @@ async fn harness_run_surfaces_structured_value_after_tool_round_trip() { let mut harness: AgentHarness<()> = AgentHarness::new(); harness - .register_model("mock", Arc::new(scripted)) - .set_default_model("mock") + .with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named("mock", Arc::new(scripted)), + )) .register_tool(Arc::new(tinyagents_harness::testkit::FakeTool::returning( "lookup", "looked-up", @@ -331,8 +332,12 @@ async fn harness_run_fails_when_final_text_is_not_valid_structured_output() { // structured output: extraction must fail the run closed. let mut harness: AgentHarness<()> = AgentHarness::new(); harness - .register_model("mock", Arc::new(MockModel::constant("just some prose"))) - .set_default_model("mock") + .with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named( + "mock", + Arc::new(MockModel::constant("just some prose")), + ), + )) .with_policy(RunPolicy { default_response_format: Some(ResponseFormat::json_schema("answer", object_schema())), ..RunPolicy::default() diff --git a/crates/tinyagents-integration-tests/tests/graph_durable.rs b/crates/tinyagents-integration-tests/tests/graph_durable.rs deleted file mode 100644 index 4c997e6a2..000000000 --- a/crates/tinyagents-integration-tests/tests/graph_durable.rs +++ /dev/null @@ -1,264 +0,0 @@ -//! End-to-end coverage for the durable state-graph runtime. -//! -//! Covers reducer-applied partial updates, the deterministic recursion-limit -//! error, checkpoint put/get/list via [`InMemoryCheckpointer`], and an -//! interrupt-then-resume round trip. - -use std::sync::Arc; - -use serde_json::json; - -use tinyagents_graph::ClosureStateReducer; -use tinyagents_graph::shared_subgraph_node; -use tinyagents_graph::*; -use tinyagents_harness::*; -use tinyagents_registry::*; - -/// Running counter plus an audit log, used to prove partial updates are merged -/// by the reducer rather than overwriting whole state. -#[derive(Clone, Debug, Default, PartialEq, Eq)] -struct Counter { - value: i64, - log: Vec<String>, -} - -#[tokio::test] -async fn reducer_merges_partial_updates_into_final_state() { - let graph = GraphBuilder::<Counter, i64>::new() - .set_reducer(ClosureStateReducer::new( - |mut state: Counter, update: i64| { - state.value += update; - state.log.push(format!("+{update}")); - Ok(state) - }, - )) - .add_node("seed", |_s: Counter, _c: NodeContext| async move { - Ok(NodeResult::Update(1)) - }) - .add_node("grow", |s: Counter, _c: NodeContext| async move { - Ok(NodeResult::Update(s.value * 10)) - }) - .add_node("finish", |_s: Counter, _c: NodeContext| async move { - Ok(NodeResult::Update(7)) - }) - .set_entry("seed") - .add_edge("seed", "grow") - .add_edge("grow", "finish") - .set_finish("finish") - .compile() - .expect("graph compiles"); - - let run = graph.run(Counter::default()).await.expect("run succeeds"); - - // 0 +1 => 1, then +10 => 11, then +7 => 18. - assert_eq!(run.state.value, 18); - assert_eq!(run.state.log, vec!["+1", "+10", "+7"]); - let visited: Vec<&str> = run.visited.iter().map(|n| n.as_str()).collect(); - assert_eq!(visited, vec!["seed", "grow", "finish"]); - assert_eq!(run.steps, 3); -} - -#[tokio::test] -async fn recursion_limit_is_deterministic() { - let graph = GraphBuilder::<i64, i64>::overwrite() - .with_recursion_limit(3) - .add_node("loop", |s: i64, _c: NodeContext| async move { - Ok(NodeResult::Update(s + 1)) - }) - .set_entry("loop") - .add_edge("loop", "loop") - .compile() - .expect("graph compiles"); - - let err = graph.run(0).await.expect_err("the loop never terminates"); - assert!( - matches!(err, TinyAgentsError::RecursionLimit(3)), - "got {err:?}" - ); -} - -#[tokio::test] -async fn checkpointer_put_get_list_round_trip() { - let checkpointer = Arc::new(InMemoryCheckpointer::<i64>::new()); - - let graph = GraphBuilder::<i64, i64>::overwrite() - .add_node("a", |s: i64, _c: NodeContext| async move { - Ok(NodeResult::Update(s + 1)) - }) - .add_node("b", |s: i64, _c: NodeContext| async move { - Ok(NodeResult::Update(s + 1)) - }) - .set_entry("a") - .add_edge("a", "b") - .set_finish("b") - .compile() - .expect("graph compiles") - .with_checkpointer(checkpointer.clone()); - - let run = graph - .run_with_thread("t1", 0) - .await - .expect("threaded run succeeds"); - assert_eq!(run.state, 2); - assert!(run.checkpoint_id.is_some()); - - // One checkpoint per superstep boundary, chained parent -> child. - let list = checkpointer.list("t1").await.expect("list succeeds"); - assert_eq!(list.len(), 2); - assert_eq!(list[0].parent_checkpoint_id, None); - assert_eq!( - list[1].parent_checkpoint_id.as_deref(), - Some(list[0].checkpoint_id.as_str()) - ); - - // `get` with None returns the latest committed checkpoint. - let latest = checkpointer - .get("t1", None) - .await - .expect("get latest succeeds") - .expect("a latest checkpoint exists"); - assert_eq!(latest.state, 2); - assert_eq!(latest.thread_id, "t1"); - - // `get` for an unknown thread is empty, not an error. - let missing = checkpointer - .get("does-not-exist", None) - .await - .expect("get is infallible"); - assert!(missing.is_none()); -} - -#[tokio::test] -async fn invalid_command_goto_does_not_write_poisoned_checkpoint() { - let checkpointer = Arc::new(InMemoryCheckpointer::<i64>::new()); - let graph = GraphBuilder::<i64, i64>::overwrite() - .add_node("router", |_s: i64, _c: NodeContext| async move { - Ok(NodeResult::Command( - Command::update(1).with_goto(["missing"]), - )) - }) - .set_entry("router") - .mark_command_routing("router") - .compile() - .expect("graph compiles") - .with_checkpointer(checkpointer.clone()); - - let err = graph - .run_with_thread("bad-route", 0) - .await - .expect_err("invalid runtime command target fails immediately"); - assert!( - matches!(err, TinyAgentsError::MissingNode(_)), - "got {err:?}" - ); - assert!( - checkpointer - .get("bad-route", None) - .await - .expect("get succeeds") - .is_none(), - "invalid command route should not persist a checkpoint with missing next nodes" - ); -} - -#[tokio::test] -async fn subgraph_child_persists_under_parent_thread_namespace() { - let checkpointer = Arc::new(InMemoryCheckpointer::<i64>::new()); - let child = GraphBuilder::<i64, i64>::overwrite() - .add_node("add", |s: i64, _c: NodeContext| async move { - Ok(NodeResult::Update(s + 10)) - }) - .set_entry("add") - .set_finish("add") - .compile() - .expect("child graph compiles") - .with_checkpointer(checkpointer.clone()); - - let parent = GraphBuilder::<i64, i64>::overwrite() - .add_node("child", shared_subgraph_node(child)) - .set_entry("child") - .set_finish("child") - .compile() - .expect("parent graph compiles") - .with_checkpointer(checkpointer.clone()); - - let run = parent - .run_with_thread("subgraph-thread", 0) - .await - .expect("threaded parent run succeeds"); - assert_eq!(run.state, 10); - - let checkpoints = checkpointer - .list("subgraph-thread") - .await - .expect("list succeeds"); - assert_eq!(checkpoints.len(), 2); - assert!(checkpoints.iter().any(|c| c.namespace.is_empty())); - assert!( - checkpoints - .iter() - .any(|c| c.namespace == vec!["child".to_string()]), - "embedded child checkpoint should use the parent thread and child namespace" - ); -} - -#[tokio::test] -async fn interrupt_then_resume_yields_resumed_result() { - let graph = GraphBuilder::<i64, i64>::overwrite() - .add_node("approve", |s: i64, ctx: NodeContext| async move { - match ctx.resume { - Some(value) => { - let bump = value.get("bump").and_then(|v| v.as_i64()).unwrap_or(0); - Ok(NodeResult::Update(s + bump)) - } - None => Ok(NodeResult::Interrupt(Interrupt::new( - "approve", - json!({ "ask": "approve?" }), - ))), - } - }) - .set_entry("approve") - .set_finish("approve") - .compile() - .expect("graph compiles") - .with_checkpointer(Arc::new(InMemoryCheckpointer::<i64>::new())); - - // First pass pauses at the interrupt. - let paused = graph - .run_with_thread("hitl", 10) - .await - .expect("first pass succeeds"); - assert!(paused.is_interrupted()); - assert_eq!(paused.interrupts.len(), 1); - - // Resuming injects the approval value and finishes: 10 + 5 = 15. - let resumed = graph - .resume("hitl", Command::resume(json!({ "bump": 5 }))) - .await - .expect("resume succeeds"); - assert!(!resumed.is_interrupted()); - assert_eq!(resumed.state, 15); -} - -#[tokio::test] -async fn interrupt_requires_resumable_durability() { - let graph = GraphBuilder::<i64, i64>::overwrite() - .add_node("approve", |_s: i64, _c: NodeContext| async move { - Ok(NodeResult::Interrupt(Interrupt::new( - "approve", - json!({ "ask": "approve?" }), - ))) - }) - .set_entry("approve") - .set_finish("approve") - .compile() - .expect("graph compiles") - .with_checkpointer(Arc::new(InMemoryCheckpointer::<i64>::new())); - - let err = graph - .run(10) - .await - .expect_err("interrupt without a thread id is not resumable"); - assert!(matches!(err, TinyAgentsError::Resume(_)), "got {err:?}"); - assert!(err.to_string().contains("thread id"), "{err}"); -} diff --git a/crates/tinyagents-integration-tests/tests/graph_lifecycle_events.rs b/crates/tinyagents-integration-tests/tests/graph_lifecycle_events.rs deleted file mode 100644 index e4cf79e08..000000000 --- a/crates/tinyagents-integration-tests/tests/graph_lifecycle_events.rs +++ /dev/null @@ -1,604 +0,0 @@ -//! The graph loop driver announces the same turn/message lifecycle events as -//! the harness's direct loop (`TurnStarted`, `TurnCompleted`, -//! `MessageAppended`), once each, and never for nested tool calls. - -use std::sync::Arc; - -use async_trait::async_trait; -use serde_json::{Value, json}; - -use tinyagents_graph::agent_loop::{AgentLoopGraphExt, GraphLoopDriver}; -use tinyagents_harness::context::{RunConfig, RunContext}; -use tinyagents_harness::events::AgentEvent; -use tinyagents_harness::limits::RunLimits; -use tinyagents_harness::runtime::{AgentHarness, LoopExecution, RunPolicy}; -use tinyagents_harness::steering::{SteeringCommand, SteeringHandle}; -use tinyagents_harness::testkit::{EventRecorder, FakeTool}; -use tinyagents_harness::tool::ToolExecutionContext; -use tinyinference_llm::message::{AssistantMessage, Message}; -use tinyinference_llm::model::ModelResponse; -use tinyinference_llm::providers::MockModel; -use tinyinference_llm::tool::ToolCall; -use tinytools::{Tool, ToolResult}; - -fn tool_call_response(id: &str, name: &str) -> ModelResponse { - ModelResponse { - message: AssistantMessage { - id: Some(format!("msg-{id}")), - content: Vec::new(), - tool_calls: vec![ToolCall::new(id, name, json!({}))], - usage: None, - origin: None, - }, - finish_reason: Some("tool_calls".to_string()), - ..ModelResponse::assistant("") - } -} - -fn harness_for( - execution: LoopExecution, - responses: Vec<ModelResponse>, - limits: RunLimits, -) -> AgentHarness<()> { - let mut harness: AgentHarness<()> = AgentHarness::new(); - harness - .register_model("mock", Arc::new(MockModel::with_responses(responses))) - .set_default_model("mock"); - if matches!(execution, LoopExecution::Graph) { - harness.with_loop_driver(Arc::new(GraphLoopDriver::new())); - } - harness.with_policy(RunPolicy { - execution, - limits, - ..RunPolicy::default() - }); - harness -} - -/// The lifecycle events of a run, flattened to comparable strings. -fn lifecycle(events: &[AgentEvent]) -> Vec<String> { - events - .iter() - .filter_map(|event| match event { - AgentEvent::TurnStarted { turn } => Some(format!("turn.started:{turn}")), - AgentEvent::TurnCompleted { - turn, - tool_call_ids, - .. - } => Some(format!( - "turn.completed:{turn}:{}", - tool_call_ids - .iter() - .map(ToString::to_string) - .collect::<Vec<_>>() - .join(",") - )), - AgentEvent::MessageAppended { - role, - index, - call_id, - .. - } => Some(format!( - "append:{index}:{role}:{}", - call_id - .as_ref() - .map(ToString::to_string) - .unwrap_or_default() - )), - _ => None, - }) - .collect() -} - -async fn run_with( - execution: LoopExecution, - responses: Vec<ModelResponse>, - input: Vec<Message>, - steering: Option<SteeringHandle>, -) -> Vec<String> { - let mut harness = harness_for(execution, responses, RunLimits::default()); - harness.register_tool(Arc::new(FakeTool::returning("lookup", "tool-output"))); - let recorder = EventRecorder::new(); - let mut ctx = RunContext::new(RunConfig::new("lifecycle"), ()).with_events(recorder.sink()); - if let Some(steering) = steering { - ctx = ctx.with_steering(steering); - } - harness - .invoke_in_context(&(), ctx, input) - .await - .expect("run completes"); - lifecycle(&recorder.events()) -} - -#[tokio::test] -async fn a_plain_run_announces_the_same_lifecycle_as_the_direct_loop() { - let script = || vec![ModelResponse::assistant("done")]; - let input = || vec![Message::user("hi")]; - let direct = run_with(LoopExecution::Direct, script(), input(), None).await; - let graph = run_with(LoopExecution::Graph, script(), input(), None).await; - assert!(!direct.is_empty(), "the direct loop emits lifecycle events"); - assert_eq!(graph, direct); -} - -#[tokio::test] -async fn a_tool_run_announces_each_message_once_and_closes_each_turn() { - let script = || { - vec![ - tool_call_response("call-1", "lookup"), - ModelResponse::assistant("done"), - ] - }; - let input = || vec![Message::user("look it up")]; - let direct = run_with(LoopExecution::Direct, script(), input(), None).await; - let graph = run_with(LoopExecution::Graph, script(), input(), None).await; - assert_eq!(graph, direct); - - // Input is never announced; every later index appears exactly once. - let mut indices: Vec<&String> = graph.iter().filter(|e| e.starts_with("append:")).collect(); - let total = indices.len(); - indices.dedup(); - assert_eq!( - indices.len(), - total, - "no duplicate MessageAppended: {graph:?}" - ); - assert!( - !graph.iter().any(|e| e.starts_with("append:0:")), - "{graph:?}" - ); - assert!( - graph.iter().any(|e| e == "turn.completed:1:call-1"), - "{graph:?}" - ); -} - -#[tokio::test] -async fn a_steering_injected_message_is_announced_like_the_direct_loop() { - let run = |execution| async move { - let steering = SteeringHandle::allow_all(); - steering.send(SteeringCommand::InjectMessage(Message::user("extra"))); - run_with( - execution, - vec![ModelResponse::assistant("done")], - vec![Message::user("hello")], - Some(steering), - ) - .await - }; - assert_eq!( - run(LoopExecution::Graph).await, - run(LoopExecution::Direct).await - ); -} - -/// A tool that calls another tool through the harness. -struct Caller; - -#[async_trait] -impl Tool for Caller { - fn name(&self) -> &str { - "caller" - } - fn description(&self) -> &str { - "calls another tool" - } - fn parameters_schema(&self) -> Value { - json!({"type": "object"}) - } - async fn execute(&self, _arguments: Value) -> anyhow::Result<ToolResult> { - unreachable!("dispatched through execute_with_context") - } - async fn execute_with_context( - &self, - _arguments: Value, - _options: tinytools::ToolCallOptions, - context: Option<&dyn tinytools::ToolRunContext>, - ) -> anyhow::Result<ToolResult> { - let harness = context - .and_then(tinytools::ToolRunContext::host_extension) - .and_then(|any| any.downcast_ref::<ToolExecutionContext>()) - .expect("the harness installs its context") - .clone(); - harness.call_tool("lookup", json!({})).await?; - Ok(ToolResult::success("caller-out")) - } -} - -#[tokio::test] -async fn nested_tool_calls_do_not_emit_message_appended_on_the_graph_driver() { - let mut rows = Vec::new(); - for execution in [LoopExecution::Direct, LoopExecution::Graph] { - let mut harness = harness_for( - execution, - vec![ - tool_call_response("p1", "caller"), - ModelResponse::assistant("done"), - ], - RunLimits::default().with_max_nested_depth(3), - ); - harness.register_tool(Arc::new(FakeTool::returning("lookup", "leaf-out"))); - harness.register_tool(Arc::new(Caller)); - let recorder = EventRecorder::new(); - let ctx = RunContext::new(RunConfig::new("nested"), ()).with_events(recorder.sink()); - harness - .invoke_in_context(&(), ctx, vec![Message::user("go")]) - .await - .expect("run completes"); - let tool_rows: Vec<String> = lifecycle(&recorder.events()) - .into_iter() - .filter(|e| e.contains(":tool:")) - .collect(); - assert_eq!(tool_rows.len(), 1, "{execution:?}: {tool_rows:?}"); - assert!(tool_rows[0].ends_with(":tool:p1"), "{tool_rows:?}"); - rows.push(tool_rows); - } - assert_eq!(rows[0], rows[1]); -} - -#[tokio::test] -async fn iter_stepping_announces_appends_once_and_never_the_input() { - let mut harness: AgentHarness<()> = AgentHarness::new(); - harness - .register_model( - "mock", - Arc::new(MockModel::with_responses(vec![ - tool_call_response("c1", "lookup"), - ModelResponse::assistant("done"), - ])), - ) - .set_default_model("mock") - .register_tool(Arc::new(FakeTool::returning("lookup", "out"))); - let harness = Arc::new(harness); - let recorder = EventRecorder::new(); - let ctx = RunContext::new(RunConfig::new("iter-lifecycle"), ()).with_events(recorder.sink()); - let mut iter = harness - .iter(Arc::new(()), ctx, vec![Message::user("go")]) - .expect("iter starts"); - iter.run_to_end().await.expect("finishes"); - - let events = lifecycle(&recorder.events()); - assert!( - !events.iter().any(|e| e.starts_with("append:0:")), - "{events:?}" - ); - let appended: Vec<&String> = events.iter().filter(|e| e.starts_with("append:")).collect(); - assert_eq!(appended.len(), 3, "assistant, tool, assistant: {events:?}"); - assert_eq!( - events - .iter() - .filter(|e| e.starts_with("turn.started")) - .count(), - events - .iter() - .filter(|e| e.starts_with("turn.completed")) - .count(), - "{events:?}" - ); -} - -/// Requests an approval interrupt on the first `after_model` hook only. -struct PauseOnce(std::sync::atomic::AtomicBool); - -#[async_trait] -impl tinyagents_harness::middleware::Middleware<(), ()> for PauseOnce { - fn name(&self) -> &str { - "pause_once" - } - - async fn after_model( - &self, - ctx: &mut RunContext<()>, - _state: &(), - _response: &mut ModelResponse, - ) -> tinyagents_harness::Result<()> { - if self.0.swap(false, std::sync::atomic::Ordering::SeqCst) { - ctx.request_control(tinyagents_harness::context::MiddlewareControl::Interrupt { - node: "review".into(), - message: "needs approval".into(), - }); - } - Ok(()) - } -} - -/// An interrupted node discards its state and re-runs on resume, in a fresh -/// runtime (a restart). Lifecycle events must stay consistent: the input is -/// never announced and a message index is never announced twice without a -/// retraction in between. -#[tokio::test] -async fn a_checkpoint_resume_in_a_fresh_runtime_neither_reannounces_input_nor_duplicates() { - use tinyagents_graph::InMemoryCheckpointer; - use tinyagents_graph::agent_loop::{LoopRuntime, LoopState, compile_loop}; - - let build = |armed: bool| { - let mut harness: AgentHarness<()> = AgentHarness::new(); - harness - .register_model( - "mock", - Arc::new(MockModel::with_responses(vec![ - tool_call_response("c1", "lookup"), - ModelResponse::assistant("done"), - ])), - ) - .set_default_model("mock") - .register_tool(Arc::new(FakeTool::returning("lookup", "out"))) - .push_middleware(Arc::new(PauseOnce(std::sync::atomic::AtomicBool::new( - armed, - )))); - Arc::new(harness) - }; - let checkpointer = Arc::new(InMemoryCheckpointer::<LoopState>::default()); - let recorder = EventRecorder::new(); - let graph_for = |harness: Arc<AgentHarness<()>>| { - let ctx = - RunContext::new(RunConfig::new("resume-lifecycle"), ()).with_events(recorder.sink()); - let rt = Arc::new(LoopRuntime::for_run(harness, Arc::new(()), ctx)); - compile_loop(rt) - .expect("compiles") - .with_checkpointer(checkpointer.clone()) - }; - - let first = graph_for(build(true)) - .run_with_thread("t", LoopState::seed(vec![Message::user("go")])) - .await - .expect("first leg reaches the interrupt"); - assert_eq!(first.interrupts.len(), 1); - - // A different harness and runtime resumes from the checkpoint. The pause - // already fired, so the second leg runs to the end. - let second = build(false); - let resumed = graph_for(second) - .resume( - "t", - tinyagents_graph::Command { - update: None, - goto: Vec::new(), - resume: Some(json!({ "approved": true })), - resume_by_task: Default::default(), - }, - ) - .await - .expect("resume completes"); - assert!(resumed.state.finished); - - let events = lifecycle(&recorder.events()); - assert!( - !events.iter().any(|e| e.starts_with("append:0:")), - "{events:?}" - ); - let mut live = std::collections::BTreeSet::new(); - for event in recorder.events() { - match event { - AgentEvent::MessageAppended { index, .. } => { - assert!( - live.insert(index), - "index {index} announced twice: {events:?}" - ); - } - AgentEvent::MessageRetracted { index } => { - live.remove(&index); - } - _ => {} - } - } -} - -/// A serial batch whose second call trips the tool cap with an error: the -/// first call's result is already on the transcript and must be announced and -/// counted by the closing `TurnCompleted`, as in the direct loop. -#[tokio::test] -async fn a_partially_executed_tool_batch_is_announced_when_the_batch_errors() { - let mut traces = Vec::new(); - for execution in [LoopExecution::Direct, LoopExecution::Graph] { - let mut response = tool_call_response("a", "lookup"); - response - .message - .tool_calls - .push(ToolCall::new("b", "lookup", json!({}))); - let mut harness = harness_for( - execution, - vec![response, ModelResponse::assistant("done")], - RunLimits::default().with_max_tool_calls(1), - ); - harness.register_tool(Arc::new(FakeTool::returning("lookup", "out"))); - let recorder = EventRecorder::new(); - let ctx = RunContext::new(RunConfig::new("partial"), ()).with_events(recorder.sink()); - let result = harness - .invoke_in_context(&(), ctx, vec![Message::user("go")]) - .await; - assert!(result.is_err(), "{execution:?}: the cap errors the run"); - traces.push(lifecycle(&recorder.events())); - } - assert!( - traces[0].iter().any(|e| e == "append:2:tool:a"), - "{:?}", - traces[0] - ); - assert_eq!(traces[1], traces[0]); -} - -/// Approval interrupt raised after the tool batch: the batch's messages and -/// its turn close are not left announced for a state the graph discards. -#[tokio::test] -async fn an_interrupt_after_the_tool_batch_retracts_instead_of_closing_the_turn() { - use tinyagents_graph::InMemoryCheckpointer; - use tinyagents_graph::agent_loop::{LoopRuntime, LoopState, compile_loop}; - use tinyagents_harness::context::MiddlewareControl; - - struct PauseAfterTools(std::sync::atomic::AtomicBool); - - #[async_trait] - impl tinyagents_harness::middleware::Middleware<(), ()> for PauseAfterTools { - fn name(&self) -> &str { - "pause_after_tools" - } - async fn after_tool( - &self, - ctx: &mut RunContext<()>, - _state: &(), - _invocation: &tinyagents_harness::middleware::ToolInvocationIdentity, - _result: &mut ToolResult, - ) -> tinyagents_harness::Result<()> { - if self.0.swap(false, std::sync::atomic::Ordering::SeqCst) { - ctx.request_control(MiddlewareControl::Interrupt { - node: "review".into(), - message: "needs approval".into(), - }); - } - Ok(()) - } - } - - let mut harness: AgentHarness<()> = AgentHarness::new(); - harness - .register_model( - "mock", - Arc::new(MockModel::with_responses(vec![ - tool_call_response("c1", "lookup"), - ModelResponse::assistant("done"), - ])), - ) - .set_default_model("mock") - .register_tool(Arc::new(FakeTool::returning("lookup", "out"))) - .push_middleware(Arc::new(PauseAfterTools( - std::sync::atomic::AtomicBool::new(true), - ))); - let recorder = EventRecorder::new(); - let ctx = RunContext::new(RunConfig::new("pause-tools"), ()).with_events(recorder.sink()); - let rt = Arc::new(LoopRuntime::for_run(Arc::new(harness), Arc::new(()), ctx)); - let checkpointer = Arc::new(InMemoryCheckpointer::<LoopState>::default()); - let graph = compile_loop(rt) - .expect("compiles") - .with_checkpointer(checkpointer.clone()); - let first = graph - .run_with_thread("t", LoopState::seed(vec![Message::user("go")])) - .await - .expect("reaches the interrupt"); - assert_eq!(first.interrupts.len(), 1); - - let events = lifecycle(&recorder.events()); - assert!( - !events.iter().any(|e| e.starts_with("turn.completed:1:c1")), - "a turn holding discarded tool results must not be reported complete: {events:?}" - ); - // Net of retractions, the discarded tool message (index 2) is not live. - let mut live = std::collections::BTreeSet::new(); - for event in recorder.events() { - match event { - AgentEvent::MessageAppended { index, .. } => { - live.insert(index); - } - AgentEvent::MessageRetracted { index } => { - live.remove(&index); - } - _ => {} - } - } - assert!( - !live.contains(&2), - "discarded tool message left live: {events:?}" - ); - - // A fresh runtime (a restart) resumes from the checkpoint: the re-run tools - // node announces its results and closes the turn it belongs to, with the - // turn numbering continuing from the checkpoint. - let mut second: AgentHarness<()> = AgentHarness::new(); - second - .register_model( - "mock", - Arc::new(MockModel::with_responses(vec![ModelResponse::assistant( - "done", - )])), - ) - .set_default_model("mock") - .register_tool(Arc::new(FakeTool::returning("lookup", "out"))); - let rt2 = Arc::new(LoopRuntime::for_run( - Arc::new(second), - Arc::new(()), - RunContext::new(RunConfig::new("pause-tools"), ()).with_events(recorder.sink()), - )); - let resumed = compile_loop(rt2) - .expect("compiles") - .with_checkpointer(checkpointer) - .resume( - "t", - tinyagents_graph::Command { - update: None, - goto: Vec::new(), - resume: Some(json!({ "approved": true })), - resume_by_task: Default::default(), - }, - ) - .await - .expect("resume completes"); - assert!(resumed.state.finished); - let events = lifecycle(&recorder.events()); - let count = |prefix: &str| events.iter().filter(|e| e.starts_with(prefix)).count(); - assert_eq!(count("turn.completed:1:c1"), 1, "{events:?}"); - assert_eq!(count("turn.started:1"), 1, "{events:?}"); - assert_eq!( - count("turn.started:2"), - 1, - "numbering continues: {events:?}" - ); - assert_eq!(count("turn.started"), count("turn.completed"), "{events:?}"); -} - -/// `LoopIter` and `compile_loop` propagate a node error directly, with no -/// driver epilogue: the nodes themselves close the turn they opened and -/// announce what had been appended. -#[tokio::test] -async fn a_node_error_in_the_iterator_closes_its_turn_and_announces_partial_results() { - // Provider error in the model node. - let mut harness: AgentHarness<()> = AgentHarness::new(); - harness - .register_model("mock", Arc::new(FailingModel)) - .set_default_model("mock"); - let recorder = EventRecorder::new(); - let ctx = RunContext::new(RunConfig::new("model-error"), ()).with_events(recorder.sink()); - let mut iter = Arc::new(harness) - .iter(Arc::new(()), ctx, vec![Message::user("go")]) - .expect("iter starts"); - assert!(iter.run_to_end().await.is_err()); - let events = lifecycle(&recorder.events()); - let count = |prefix: &str| events.iter().filter(|e| e.starts_with(prefix)).count(); - assert_eq!(count("turn.started"), count("turn.completed"), "{events:?}"); - - // A serial batch whose second call trips the tool cap. - let mut response = tool_call_response("a", "lookup"); - response - .message - .tool_calls - .push(ToolCall::new("b", "lookup", json!({}))); - let mut harness = harness_for( - LoopExecution::Direct, - vec![response, ModelResponse::assistant("done")], - RunLimits::default().with_max_tool_calls(1), - ); - harness.register_tool(Arc::new(FakeTool::returning("lookup", "out"))); - let recorder = EventRecorder::new(); - let ctx = RunContext::new(RunConfig::new("batch-error"), ()).with_events(recorder.sink()); - let mut iter = Arc::new(harness) - .iter(Arc::new(()), ctx, vec![Message::user("go")]) - .expect("iter starts"); - assert!(iter.run_to_end().await.is_err()); - let events = lifecycle(&recorder.events()); - assert!(events.iter().any(|e| e == "append:2:tool:a"), "{events:?}"); - assert!( - events.iter().any(|e| e == "turn.completed:1:a"), - "{events:?}" - ); -} - -struct FailingModel; - -#[async_trait] -impl tinyinference_llm::model::ChatModel<()> for FailingModel { - async fn invoke( - &self, - _: &(), - _: tinyinference_llm::model::ModelRequest, - ) -> tinyinference_llm::Result<ModelResponse> { - Err(tinyinference_llm::Error::Model("boom".to_string())) - } -} diff --git a/crates/tinyagents-integration-tests/tests/harness_agent_loop.rs b/crates/tinyagents-integration-tests/tests/harness_agent_loop.rs index 21e970ed6..ee96a8f17 100644 --- a/crates/tinyagents-integration-tests/tests/harness_agent_loop.rs +++ b/crates/tinyagents-integration-tests/tests/harness_agent_loop.rs @@ -23,8 +23,7 @@ use tinyagents_harness::runtime::{AgentHarness, RunPolicy}; use tinyagents_harness::testkit::{FakeTool, Trajectory}; use tinyinference_llm::message::{AssistantMessage, ContentBlock, Message}; use tinyinference_llm::model::{ - CapabilitySet, ChatModel, ModelHint, ModelProfile, ModelRequest, ModelResolutionSource, - ModelResponse, + CapabilitySet, ChatModel, ModelProfile, ModelRequest, ModelResolutionSource, ModelResponse, }; use tinyinference_llm::providers::MockModel; use tinyinference_llm::tool::ToolCall; @@ -105,11 +104,7 @@ impl Middleware<(), ()> for RequireJsonSchemaMiddleware { json_schema: true, ..CapabilitySet::default() }); - request.model_hints.push(ModelHint { - model: "capable".to_string(), - priority: 1, - reason: Some("needs json schema".to_string()), - }); + request.model = Some("capable".to_string()); Ok(()) } } @@ -120,8 +115,12 @@ async fn single_turn_response_completes_with_one_model_call() { let mut harness: AgentHarness<()> = AgentHarness::new(); harness - .register_model("mock", Arc::new(MockModel::constant("hello there"))) - .set_default_model("mock") + .with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named( + "mock", + Arc::new(MockModel::constant("hello there")), + ), + )) .push_middleware(Arc::new(CaptureMiddleware { listener: listener.clone(), })); @@ -142,33 +141,34 @@ async fn single_turn_response_completes_with_one_model_call() { } #[tokio::test] -async fn required_capabilities_select_capable_model() { +async fn explicit_model_override_uses_the_named_capable_route() { let mut harness: AgentHarness<()> = AgentHarness::new(); harness - .register_model( - "incapable", - Arc::new(ProfiledIntegrationModel { - profile: ModelProfile::default(), - text: "incapable", - }), - ) - .set_default_model("incapable") - .register_model( - "capable", - Arc::new(ProfiledIntegrationModel { - profile: ModelProfile { - json_schema: true, - ..ModelProfile::default() - }, - text: "capable", - }), - ) + .with_model_resolver(Arc::new( + tinyagents_harness::host::PrebuiltModelResolver::new( + Arc::new(ProfiledIntegrationModel { + profile: ModelProfile::default(), + text: "incapable", + }), + [( + "capable".to_string(), + Arc::new(ProfiledIntegrationModel { + profile: ModelProfile { + json_schema: true, + ..ModelProfile::default() + }, + text: "capable", + }) as Arc<dyn ChatModel<()>>, + )], + ) + .with_primary_name("incapable"), + )) .push_middleware(Arc::new(RequireJsonSchemaMiddleware)); let run = harness .invoke_default(&(), vec![Message::user("pick a model")]) .await - .expect("capable hinted model satisfies required capabilities"); + .expect("named route satisfies the request's required capabilities"); assert_eq!(run.text(), Some("capable".to_string())); let resolved = run @@ -177,7 +177,7 @@ async fn required_capabilities_select_capable_model() { .resolved_model .expect("resolved model metadata"); assert_eq!(resolved.name, "capable"); - assert_eq!(resolved.source, ModelResolutionSource::Hint); + assert_eq!(resolved.source, ModelResolutionSource::RequestOverride); } #[tokio::test] @@ -186,14 +186,15 @@ async fn multi_step_tool_loop_calls_tool_then_finishes() { 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), - ])), - ) - .set_default_model("mock") + .with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named( + "mock", + Arc::new(MockModel::with_responses(vec![ + tool_call_response("call-1", "lookup", json!({ "q": "x" })), + text_response("done", 4, 2), + ])), + ), + )) .register_tool(Arc::new(FakeTool::returning("lookup", "tool-output"))) .push_middleware(Arc::new(CaptureMiddleware { listener: listener.clone(), @@ -223,11 +224,12 @@ async fn max_model_calls_limit_returns_limit_exceeded() { let mut harness: AgentHarness<()> = AgentHarness::new(); // A model that always asks for a tool would loop forever without a cap. harness - .register_model( - "mock", - Arc::new(MockModel::with_tool_call("spin", json!({}))), - ) - .set_default_model("mock") + .with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named( + "mock", + Arc::new(MockModel::with_tool_call("spin", json!({}))), + ), + )) .register_tool(Arc::new(FakeTool::returning("spin", "again"))) .with_policy(RunPolicy { limits: RunLimits::default().with_max_model_calls(1), @@ -249,14 +251,15 @@ async fn max_model_calls_limit_returns_limit_exceeded() { async fn usage_accumulates_across_model_calls() { 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), - ])), - ) - .set_default_model("mock") + .with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named( + "mock", + Arc::new(MockModel::with_responses(vec![ + tool_call_response("call-1", "lookup", json!({})), + text_response("done", 4, 2), + ])), + ), + )) .register_tool(Arc::new(FakeTool::returning("lookup", "out"))); let run = harness diff --git a/crates/tinyagents-integration-tests/tests/harness_stress.rs b/crates/tinyagents-integration-tests/tests/harness_stress.rs index a21dbdfd6..66ada1b17 100644 --- a/crates/tinyagents-integration-tests/tests/harness_stress.rs +++ b/crates/tinyagents-integration-tests/tests/harness_stress.rs @@ -17,7 +17,12 @@ use tinyinference_llm::providers::MockModel; fn harness() -> Arc<AgentHarness<()>> { let mut harness = AgentHarness::new(); - harness.register_model("mock", Arc::new(MockModel::constant("done"))); + harness.with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named( + "mock", + Arc::new(MockModel::constant("done")), + ), + )); Arc::new(harness) } diff --git a/crates/tinyagents-integration-tests/tests/live_cache.rs b/crates/tinyagents-integration-tests/tests/live_cache.rs index 43003c7a6..a1e9a29e4 100644 --- a/crates/tinyagents-integration-tests/tests/live_cache.rs +++ b/crates/tinyagents-integration-tests/tests/live_cache.rs @@ -82,7 +82,9 @@ async fn live_openai_response_cache_hits_on_repeated_question() { let cache = Arc::new(InMemoryResponseCache::new()); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("openai", counting); + harness.with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named("openai", counting), + )); harness.with_response_cache(cache); // A short, deterministic question. The harness builds the `ModelRequest` diff --git a/crates/tinyagents-integration-tests/tests/live_checkpoint_resume.rs b/crates/tinyagents-integration-tests/tests/live_checkpoint_resume.rs deleted file mode 100644 index 6c77a0d78..000000000 --- a/crates/tinyagents-integration-tests/tests/live_checkpoint_resume.rs +++ /dev/null @@ -1,158 +0,0 @@ -//! LIVE end-to-end: a durable [`CompiledGraph`] whose single node drives a real -//! OpenAI [`AgentHarness`] — but only *after* a human-in-the-loop interrupt. -//! -//! This composes the **durable graph runtime** (builder/executor + reducer), the -//! **checkpointer** (durability), the **command/interrupt** model, and a real -//! network-backed **harness** run. The flow deliberately straddles the model -//! call with a checkpoint boundary: -//! -//! 1. First pass: the node sees no `resume` value and immediately interrupts -//! *before* making any model call. The runtime persists a checkpoint. -//! 2. Resume: a [`Command::resume`] re-enters the node, which now performs the -//! real OpenAI model call and folds the answer into state. -//! -//! We assert *structurally* — the run pauses, a checkpoint persists, the resumed -//! run completes, and a model call actually happened — never on LLM prose. -//! -//! # Skips gracefully -//! -//! This test is `#[ignore]`d and only runs opted in via -//! `tests/common/live.rs::require_live`, so `cargo test` passes with no key -//! configured and never dials a real provider by accident. - -mod common; - -#[tokio::test] -#[ignore = "network: set TINYAGENTS_LIVE=1 and run with --ignored"] -async fn live_durable_graph_checkpoints_then_resumes_across_model_call() { - use std::sync::Arc; - - use serde_json::json; - - use tinyagents_graph::*; - use tinyagents_harness::runtime::AgentHarness; - use tinyagents_harness::*; - use tinyagents_registry::*; - use tinyinference_llm::message::Message; - use tinyinference_llm::providers::openai::OpenAiModel; - - if !common::live::require_live(&["OPENAI_API_KEY"]) { - return; - } - - /// State threaded through the durable graph: the question to ask, plus the - /// answer + model-call count once the node has driven the real harness. - #[derive(Clone, Debug, Default)] - struct QaState { - question: String, - answer: Option<String>, - model_calls: usize, - } - - // A real OpenAI-backed harness; the graph node owns a shared handle to it. - let mut harness: AgentHarness<()> = AgentHarness::new(); - harness - .register_model( - "openai", - Arc::new(OpenAiModel::from_env().expect("OPENAI_API_KEY present")), - ) - .set_default_model("openai"); - let harness = Arc::new(harness); - - let checkpointer = Arc::new(InMemoryCheckpointer::<QaState>::new()); - - let node_harness = harness.clone(); - let graph = GraphBuilder::<QaState, QaState>::overwrite() - .with_graph_id("live-checkpoint-resume") - .add_node("ask", move |mut state: QaState, ctx: NodeContext| { - let harness = node_harness.clone(); - async move { - match ctx.resume { - // First pass: pause *before* the model call so the checkpoint - // boundary lands squarely in front of the network request. - None => Ok(NodeResult::Interrupt(Interrupt::new( - "ask", - json!({ "ask": "approve calling the model?" }), - ))), - // Resume pass: perform the real OpenAI call and fold the - // result into state. - Some(_) => { - let run = harness - .invoke_default(&(), vec![Message::user(state.question.clone())]) - .await?; - state.answer = run.text(); - state.model_calls = run.model_calls; - Ok(NodeResult::Update(state)) - } - } - } - }) - .set_entry("ask") - .set_finish("ask") - .compile() - .expect("durable graph compiles") - .with_checkpointer(checkpointer.clone()); - - // First pass: the node interrupts before touching the network. - let paused = graph - .run_with_thread( - "ckpt-thread", - QaState { - question: "Reply with a single short greeting.".to_string(), - answer: None, - model_calls: 0, - }, - ) - .await - .expect("first pass succeeds"); - - assert!( - paused.is_interrupted(), - "the run should pause before the model call" - ); - assert_eq!(paused.interrupts.len(), 1); - assert_eq!(paused.interrupts[0].node.as_str(), "ask"); - // No model call has happened yet on the paused side of the interrupt. - assert!( - paused.state.answer.is_none(), - "no answer should exist before resume" - ); - - // Durability: the interrupt persisted at least one checkpoint for the thread. - let checkpoints = checkpointer - .list("ckpt-thread") - .await - .expect("list succeeds"); - assert!( - !checkpoints.is_empty(), - "an interrupt must persist a checkpoint" - ); - - // Resume: re-enter the node, which now makes the real OpenAI model call. - let resumed = graph - .resume("ckpt-thread", Command::resume(json!({ "approved": true }))) - .await - .expect("resume succeeds"); - - // Structural assertions only — never assert on exact LLM prose. - assert!(!resumed.is_interrupted(), "resumed run must complete"); - assert!( - resumed.state.model_calls >= 1, - "the resumed pass should make at least one model call" - ); - let answer = resumed.state.answer.unwrap_or_default(); - assert!( - !answer.trim().is_empty(), - "the resumed run should fold a non-empty model answer into state" - ); - - // The resumed run committed at least as many checkpoints as the paused one. - let after_resume = checkpointer - .list("ckpt-thread") - .await - .expect("list succeeds"); - assert!( - after_resume.len() >= checkpoints.len(), - "resuming should not drop persisted checkpoints" - ); -} diff --git a/crates/tinyagents-integration-tests/tests/live_local_models.rs b/crates/tinyagents-integration-tests/tests/live_local_models.rs index 8af87cb67..376d49f34 100644 --- a/crates/tinyagents-integration-tests/tests/live_local_models.rs +++ b/crates/tinyagents-integration-tests/tests/live_local_models.rs @@ -649,8 +649,12 @@ async fn local_runtimes_complete_a_full_tool_loop() { let mut harness: AgentHarness<()> = AgentHarness::new(); harness.register_tool(tool.clone()); harness - .register_model("local", Arc::new(runtime.model())) - .set_default_model("local") + .with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named( + "local", + Arc::new(runtime.model()), + ), + )) .with_policy(local_run_policy()) .push_middleware(Arc::new(ForceFirstToolCall::new())); diff --git a/crates/tinyagents-integration-tests/tests/live_sdk_gaps.rs b/crates/tinyagents-integration-tests/tests/live_sdk_gaps.rs index 016a819ff..4bc45af2f 100644 --- a/crates/tinyagents-integration-tests/tests/live_sdk_gaps.rs +++ b/crates/tinyagents-integration-tests/tests/live_sdk_gaps.rs @@ -44,7 +44,9 @@ async fn live_budget_blocks_second_call() { // 1) Sanity: a normal harness run succeeds against the real model. { let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("openai", model.clone()); + harness.with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named("openai", model.clone()), + )); let run = harness .invoke_default(&(), vec![Message::user("Reply with exactly: hi")]) .await @@ -68,7 +70,9 @@ async fn live_budget_blocks_second_call() { .with_tracker(tracker); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("openai", model); + harness.with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named("openai", model), + )); harness.push_middleware(Arc::new(budget)); let err = harness @@ -143,7 +147,9 @@ async fn live_tool_policy_exposes_classified_tool() { Arc::new(OpenAiModel::from_env().expect("OPENAI_API_KEY present")); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("openai", model); + harness.with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named("openai", model), + )); harness.register_tool(Arc::new(AddTool { calls: calls.clone(), })); diff --git a/crates/tinyagents-integration-tests/tests/live_steering.rs b/crates/tinyagents-integration-tests/tests/live_steering.rs deleted file mode 100644 index 5d9bceeae..000000000 --- a/crates/tinyagents-integration-tests/tests/live_steering.rs +++ /dev/null @@ -1,61 +0,0 @@ -//! Live steering test against the real OpenAI API. -//! -//! This test is `#[ignore]`d and only runs opted in via -//! `tests/common/live.rs::require_live`, so `cargo test` stays green without -//! credentials and never dials a real provider by accident. Run it for real -//! with: -//! -//! ```text -//! TINYAGENTS_LIVE=1 cargo test --test live_steering -- --ignored --nocapture -//! ``` - -mod common; - -use std::sync::Arc; - -use tinyagents_graph::*; -use tinyagents_harness::context::{RunConfig, RunContext}; -use tinyagents_harness::runtime::AgentHarness; -use tinyagents_harness::*; -use tinyagents_registry::*; -use tinyinference_llm::message::Message; -use tinyinference_llm::providers::openai::OpenAiModel; - -#[tokio::test] -#[ignore = "network: set TINYAGENTS_LIVE=1 and run with --ignored"] -async fn orchestrator_steers_a_real_openai_run() { - if !common::live::require_live(&["OPENAI_API_KEY"]) { - return; - } - - let model = OpenAiModel::from_env().expect("OpenAiModel from env"); - let mut harness = AgentHarness::<(), ()>::new(); - harness.register_model("openai", Arc::new(model)); - harness.set_default_model("openai"); - - // The orchestrator injects an overriding instruction before the model call. - let steering = SteeringHandle::allow_all(); - steering.send(SteeringCommand::Redirect { - instruction: "Ignore the question's topic. Reply with EXACTLY the single \ - word BANANA in uppercase and nothing else." - .into(), - }); - - let ctx = RunContext::new(RunConfig::new("live-steer"), ()).with_steering(steering); - - let run = harness - .invoke_in_context( - &(), - ctx, - vec![Message::user("Tell me about the history of Rome.")], - ) - .await - .expect("live steered run completes"); - - let answer = run.text().unwrap_or_default(); - eprintln!("steered answer: {answer:?}"); - assert!( - answer.to_uppercase().contains("BANANA"), - "the steered instruction should dominate the response, got {answer:?}" - ); -} diff --git a/crates/tinyagents-integration-tests/tests/live_subagents.rs b/crates/tinyagents-integration-tests/tests/live_subagents.rs index 90b1d08e6..cba16c172 100644 --- a/crates/tinyagents-integration-tests/tests/live_subagents.rs +++ b/crates/tinyagents-integration-tests/tests/live_subagents.rs @@ -1,11 +1,9 @@ //! LIVE end-to-end: a real OpenAI parent agent composes the answer of a real //! OpenAI child sub-agent. //! -//! This is the network-backed sibling of `e2e_subagents.rs`. The parent -//! [`AgentHarness`] is given a [`SubAgentTool`] wrapping a child harness; both -//! are driven by a real [`OpenAiModel`]. We assert *structurally* (the parent -//! finished, produced non-empty text, and the sub-agent tool ran) rather than -//! on the exact model prose. +//! This is the network-backed sibling of the mock-backed subagent coverage. +//! It uses the public planner/executor lifecycle and `DriverSubAgentTool`, +//! then waits for the detached child through `SubAgentJobsTool`. //! //! # Skips gracefully //! @@ -18,73 +16,210 @@ mod common; #[tokio::test] #[ignore = "network: set TINYAGENTS_LIVE=1 and run with --ignored"] async fn live_openai_parent_composes_child_subagent() { - use std::sync::Arc; + use std::collections::HashMap; + use std::sync::{Arc, Mutex}; + use async_trait::async_trait; + use serde_json::{Value, json}; use tinyagents_harness::context::{RunConfig, RunContext}; use tinyagents_harness::runtime::AgentHarness; use tinyagents_harness::testkit::{EventRecorder, Trajectory}; - use tinyagents_harness_tools::subagents::{SubAgentJobsTool, SubAgentMessageTool}; - use tinyagents_orchestration::subagent::{ChildDataPolicy, SubAgent, SubAgentTool}; + use tinyagents_harness::tool::ToolSnapshot; + use tinyagents_harness_tools::subagents::{DriverSubAgentTool, SubAgentJobsTool}; + use tinyagents_orchestration::subagent::*; use tinyinference_llm::message::Message; + use tinyinference_llm::model::ChatModel; use tinyinference_llm::providers::openai::OpenAiModel; if !common::live::require_live(&["OPENAI_API_KEY"]) { return; } - // Child agent: a focused "math expert" sub-agent driven by a real model. - let mut child: AgentHarness<()> = AgentHarness::new(); - child - .register_model( - "openai", - Arc::new(OpenAiModel::from_env().expect("OPENAI_API_KEY present")), - ) - .set_default_model("openai"); + #[derive(Default)] + struct MemoryPersistence { + terminals: Mutex<HashMap<SubagentTaskKey, SubagentOutcome>>, + pauses: Mutex<HashMap<SubagentTaskKey, SubagentOutcome>>, + } - let subagent = Arc::new( - SubAgent::new( - "math_expert", - "Computes precise arithmetic answers. Pass the arithmetic question as `input`.", - Arc::new(child), - ) - .with_system_prompt( - "You are a meticulous arithmetic engine. Reply with only the numeric answer.", - ), + #[async_trait] + impl SubagentPersistence for MemoryPersistence { + async fn load_terminal( + &self, + key: &SubagentTaskKey, + ) -> Result<Option<SubagentOutcome>, SubagentError> { + Ok(self.terminals.lock().unwrap().get(key).cloned()) + } + + async fn load( + &self, + key: &SubagentTaskKey, + ) -> Result<Option<SubagentResume>, SubagentError> { + Ok(self + .pauses + .lock() + .unwrap() + .get(key) + .and_then(|outcome| match &outcome.status { + SubagentOutcomeKind::AwaitingInput(pause) => Some(pause.resume.clone()), + _ => None, + })) + } + + async fn load_pause( + &self, + key: &SubagentTaskKey, + ) -> Result<Option<SubagentOutcome>, SubagentError> { + Ok(self.pauses.lock().unwrap().get(key).cloned()) + } + + async fn save_pause( + &self, + pause: PersistedSubagentPause, + ) -> Result<SubagentPausePersistenceDisposition, SubagentError> { + self.pauses.lock().unwrap().insert(pause.key, pause.outcome); + Ok(SubagentPausePersistenceDisposition::Inserted) + } + + async fn record_terminal( + &self, + key: &SubagentTaskKey, + outcome: &SubagentOutcome, + _: Option<&SubagentResume>, + ) -> Result<SubagentTerminalPersistenceDisposition, SubagentError> { + let mut terminals = self.terminals.lock().unwrap(); + if terminals.contains_key(key) { + return Ok(SubagentTerminalPersistenceDisposition::Existing); + } + self.pauses.lock().unwrap().remove(key); + terminals.insert(key.clone(), outcome.clone()); + Ok(SubagentTerminalPersistenceDisposition::Inserted) + } + } + + struct Planner; + + #[async_trait] + impl SubagentPlanner for Planner { + async fn prepare( + &self, + request: SubagentRequest, + ) -> Result<PreparedSubagent, SubagentError> { + let parts = request.into_parts(); + Ok(PreparedSubagent::new( + parts.task_key.task_id, + "math_expert", + vec![ + Message::system( + "You are a meticulous arithmetic engine. Reply with only the numeric answer.", + ), + Message::user(parts.input), + ], + ToolSnapshot::new(vec![]).expect("empty child tool snapshot is valid"), + parts.run_context, + )) + } + } + + struct Executor { + child: Arc<AgentHarness<()>>, + } + + #[async_trait] + impl SubagentExecutor for Executor { + async fn execute( + &self, + execution: SubagentExecution, + ) -> Result<SubagentOutcome, SubagentError> { + let task_id = execution.prepared.task_id; + let run = self + .child + .invoke_in_context( + &(), + execution.prepared.run_context, + execution.prepared.input, + ) + .await + .map_err(|error| SubagentError::Execution(error.to_string()))?; + Ok(SubagentOutcome::completed( + task_id, + run.text().unwrap_or_default(), + )) + } + } + + let model: Arc<dyn ChatModel<()>> = + Arc::new(OpenAiModel::from_env().expect("OPENAI_API_KEY present")); + + let mut child = AgentHarness::new(); + child.with_model(model.clone()); + let child = Arc::new(child); + + let driver = Arc::new( + SubagentDriver::new(SubagentCapabilities { + planner: Some(Arc::new(Planner)), + executor: Some(Arc::new(Executor { child })), + persistence: Some(Arc::new(MemoryPersistence::default())), + }) + .expect("all subagent lifecycle capabilities are configured"), ); - let tool = Arc::new(SubAgentTool::new( - subagent, - ChildDataPolicy::new(|parent: &()| *parent), - )); - let jobs = tool.job_registry().clone(); - - // Parent agent: also a real model, equipped with the sub-agent as a tool. - let mut parent: AgentHarness<()> = AgentHarness::new(); - parent.register_tool_dispatch(tool); - parent.register_tool_dispatch(Arc::new(SubAgentJobsTool::new(jobs.clone()))); - parent.register_tool_dispatch(Arc::new(SubAgentMessageTool::new(jobs.clone()))); - parent - .register_model( - "openai", - Arc::new(OpenAiModel::from_env().expect("OPENAI_API_KEY present")), - ) - .set_default_model("openai"); + let jobs = SubAgentJobRegistry::new(); + let request_jobs = jobs.clone(); + let spawn = DriverSubAgentTool::new( + "spawn_subagent", + "Start a named expert agent on a task and return its task id.", + driver, + move |_: &(), + parent: &RunContext<()>, + args: &Value, + call_id: &tinyagents_harness::ids::CallId| { + let child = parent + .child( + RunConfig::new(format!("live-child-{}", call_id.as_str())), + (), + ) + .map_err(|error| SubagentError::InvalidRequest(error.to_string()))?; + SubagentRequest::fresh_from_parent( + parent, + child, + format!("math-expert-{}", call_id.as_str()), + (), + args.get("prompt") + .and_then(Value::as_str) + .unwrap_or_default(), + None, + ) + .map(|request| request.with_target("math_expert")) + }, + |_| None, + ) + .with_parameters(json!({ + "type": "object", + "required": ["agent_id", "prompt"], + "properties": { + "agent_id": {"type": "string", "enum": ["math_expert"]}, + "prompt": {"type": "string"} + } + })) + .with_job_registry(request_jobs.clone()); + + let mut parent = AgentHarness::new(); + parent.register_tool_dispatch(Arc::new(spawn)); + parent.register_tool_dispatch(Arc::new(SubAgentJobsTool::new(request_jobs.clone()))); + parent.with_model(model); let recorder = EventRecorder::new(); let ctx = RunContext::new(RunConfig::new("live-parent-run"), ()).with_events(recorder.sink()); - let run = parent .invoke_in_context( &(), ctx, vec![Message::user( - "Start math_expert to compute 17 * 23. Use subagent_jobs with its job_id until \ - it completes, then state the result in a short sentence.", + "Use spawn_subagent to ask math_expert to compute 17 * 23. Then use subagent_jobs with action=await and the returned task_id to wait for the child. State the numeric result in a short sentence.", )], ) .await .expect("live parent run succeeds"); - // Structural assertions only — never assert on exact LLM prose. let final_text = run.text().unwrap_or_default(); assert!( !final_text.trim().is_empty(), @@ -95,27 +230,24 @@ async fn live_openai_parent_composes_child_subagent() { "parent should make at least one model call" ); - let traj = Trajectory::from_events(recorder.events()); - traj.assert_completed(); - // The model is strongly steered to call the sub-agent; if it did, the - // trajectory records it as a `math_expert` tool invocation. - if traj.tool_was_called("math_expert") { - assert!( - run.tool_calls >= 1, - "a recorded sub-agent tool call should be reflected in run.tool_calls" - ); - let job = tokio::time::timeout(std::time::Duration::from_secs(120), async { - loop { - if let Some(job) = jobs.list().into_iter().next() - && job.status.is_terminal() - { - return job; - } - tokio::time::sleep(std::time::Duration::from_millis(25)).await; + let trajectory = Trajectory::from_events(recorder.events()); + trajectory.assert_completed(); + assert!( + trajectory.tool_was_called("spawn_subagent"), + "parent should start the real child through DriverSubAgentTool" + ); + + let job = tokio::time::timeout(std::time::Duration::from_secs(120), async { + loop { + if let Some(job) = jobs.list().into_iter().next() + && job.status.is_terminal() + { + return job; } - }) - .await - .expect("live math subagent job finishes within two minutes"); - assert!(job.output.is_some(), "completed math job has output"); - } + tokio::time::sleep(std::time::Duration::from_millis(25)).await; + } + }) + .await + .expect("live math subagent job finishes within two minutes"); + assert!(job.output.is_some(), "completed math job has output"); } diff --git a/crates/tinyagents-integration-tests/tests/live_tool_deferral.rs b/crates/tinyagents-integration-tests/tests/live_tool_deferral.rs index 1ce8ab69a..58580fe30 100644 --- a/crates/tinyagents-integration-tests/tests/live_tool_deferral.rs +++ b/crates/tinyagents-integration-tests/tests/live_tool_deferral.rs @@ -307,8 +307,9 @@ async fn run_once( let mut harness: AgentHarness<()> = AgentHarness::new(); harness - .register_model("live", model.clone()) - .set_default_model("live") + .with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named("live", model.clone()), + )) .register_tool(Arc::new(ReadNote)) .push_middleware(Arc::new(Capture { listener: listener.clone(), @@ -507,8 +508,9 @@ async fn live_search_promotes_typed_schema_before_the_real_call() { }); let mut harness: AgentHarness<()> = AgentHarness::new(); harness - .register_model("live", model.clone()) - .set_default_model("live") + .with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named("live", model.clone()), + )) .register_tool(Arc::new(LiveTool { name: "stock_quote", description: "Fetch the trading price of a stock ticker.", diff --git a/crates/tinyagents-integration-tests/tests/loop_as_graph.rs b/crates/tinyagents-integration-tests/tests/loop_as_graph.rs deleted file mode 100644 index 5d0a5c2b2..000000000 --- a/crates/tinyagents-integration-tests/tests/loop_as_graph.rs +++ /dev/null @@ -1,871 +0,0 @@ -//! Equivalence coverage for A5: the compiled-graph rendition of the agent -//! loop (`tinyagents_graph::agent_loop`) against the harness's built-in -//! direct loop. -//! -//! Each scenario below runs the identical scripted model/harness setup -//! twice — once with `RunPolicy::execution = LoopExecution::Direct` (the -//! default), once with `LoopExecution::Graph` plus -//! `AgentHarness::with_loop_driver(Arc::new(GraphLoopDriver::new()))` — and -//! asserts the two runs agree on transcript, structured output, and usage, -//! and that the direct run's `AgentEvent` kind sequence appears (in order, -//! extra graph events allowed) within the graph run's sequence. See -//! `docs/modules/harness/state-graph.md` for the documented scope of the -//! graph rendition (it is a subset of the direct loop's behavior, not a -//! byte-for-byte reimplementation). - -use std::sync::Arc; - -use serde_json::json; - -use tinyagents_graph::InMemoryCheckpointer; -use tinyagents_graph::agent_loop::{ - AgentLoopGraphExt, GraphLoopDriver, LoopRuntime, LoopState, compile_loop, node, -}; - -mod common; -use common::storage; -use tinyagents_harness::TinyAgentsError; -use tinyagents_harness::context::{RunConfig, RunContext}; -use tinyagents_harness::runtime::{AgentHarness, LoopExecution, RunPolicy}; -use tinyagents_harness::steering::{SteeringCommand, SteeringHandle}; -use tinyagents_harness::testkit::{EventRecorder, FakeTool}; -use tinyinference_llm::message::{AssistantMessage, Message}; -use tinyinference_llm::model::{ModelResponse, ResponseFormat}; -use tinyinference_llm::providers::MockModel; -use tinyinference_llm::tool::ToolCall; -use tinyinference_llm::usage::Usage; - -fn tool_call_response(id: &str, name: &str, arguments: serde_json::Value) -> ModelResponse { - ModelResponse { - message: AssistantMessage { - id: Some(format!("msg-{id}")), - 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()), - ..ModelResponse::assistant("") - } -} - -/// Builds a harness for `execution`, registering `model` as the default. -fn harness_for(execution: LoopExecution, model: Arc<MockModel>) -> AgentHarness<()> { - let mut harness: AgentHarness<()> = AgentHarness::new(); - harness - .register_model("mock", model) - .set_default_model("mock"); - if matches!(execution, LoopExecution::Graph) { - harness.with_loop_driver(Arc::new(GraphLoopDriver::new())); - } - harness.with_policy(RunPolicy { - execution, - ..RunPolicy::default() - }); - harness -} - -/// Asserts `expected` appears, in order, as a (not necessarily contiguous) -/// subsequence of `actual` — the "same kind sequence, extra graph events -/// allowed" contract. -/// -/// The turn/message lifecycle events (`turn.*`, `message.appended`) are part of -/// the contract: the graph driver announces them like the direct loop does, and -/// `graph_lifecycle_events.rs` pins their exact shape. -fn assert_kinds_subsequence(expected: &[String], actual: &[String]) { - let mut cursor = 0; - for kind in expected { - let Some(offset) = actual[cursor..].iter().position(|k| k == kind) else { - panic!( - "expected event kind `{kind}` not found (in order) in graph run's kinds: \ - {actual:?}; direct run's kinds were: {expected:?}" - ); - }; - cursor += offset + 1; - } -} - -// ── Scenario 1: tool call ─────────────────────────────────────────────────── - -#[tokio::test] -async fn tool_call_scenario_matches_direct_and_graph() { - let mut direct_run = None; - let mut direct_kinds = None; - let mut graph_run = None; - let mut graph_kinds = None; - - for execution in [LoopExecution::Direct, LoopExecution::Graph] { - let model = Arc::new(MockModel::with_responses(vec![ - tool_call_response("call-1", "lookup", json!({ "q": "x" })), - ModelResponse::assistant("done"), - ])); - let mut harness = harness_for(execution, model); - harness.register_tool(Arc::new(FakeTool::returning("lookup", "tool-output"))); - - let recorder = EventRecorder::new(); - let ctx = RunContext::new(RunConfig::new("tool-call"), ()).with_events(recorder.sink()); - let run = harness - .invoke_in_context(&(), ctx, vec![Message::user("look something up")]) - .await - .expect("run completes"); - - match execution { - LoopExecution::Direct => { - direct_run = Some(run); - direct_kinds = Some(recorder.kinds()); - } - LoopExecution::Graph => { - graph_run = Some(run); - graph_kinds = Some(recorder.kinds()); - } - } - } - - let (direct_run, graph_run) = (direct_run.unwrap(), graph_run.unwrap()); - assert_eq!(direct_run.model_calls, graph_run.model_calls); - assert_eq!(direct_run.tool_calls, graph_run.tool_calls); - assert_eq!(direct_run.executed_tools, graph_run.executed_tools); - assert_eq!(direct_run.text(), graph_run.text()); - assert_eq!(direct_run.usage, graph_run.usage); - assert_kinds_subsequence(&direct_kinds.unwrap(), &graph_kinds.unwrap()); -} - -// ── Scenario 2: structured output ─────────────────────────────────────────── - -#[tokio::test] -async fn structured_output_scenario_matches_direct_and_graph() { - let schema = json!({ - "type": "object", - "properties": { "answer": { "type": "string" } }, - "required": ["answer"], - }); - - let mut direct_run = None; - let mut graph_run = None; - - for execution in [LoopExecution::Direct, LoopExecution::Graph] { - let model = Arc::new(MockModel::with_responses(vec![ModelResponse::assistant( - r#"{"answer":"42"}"#, - )])); - let mut harness = harness_for(execution, model); - let mut policy = harness.policy().clone(); - policy.default_response_format = Some(ResponseFormat::auto("answer", schema.clone())); - harness.with_policy(policy); - - let run = harness - .invoke_default(&(), vec![Message::user("what is the answer")]) - .await - .expect("run completes"); - - match execution { - LoopExecution::Direct => direct_run = Some(run), - LoopExecution::Graph => graph_run = Some(run), - } - } - - let (direct_run, graph_run) = (direct_run.unwrap(), graph_run.unwrap()); - assert_eq!(direct_run.structured, graph_run.structured); - assert_eq!(direct_run.structured, Some(json!({ "answer": "42" }))); -} - -// ── Scenario 3: limit stop ────────────────────────────────────────────────── - -#[tokio::test] -async fn limit_stop_scenario_matches_direct_and_graph() { - let mut direct_err = None; - let mut graph_err = None; - - for execution in [LoopExecution::Direct, LoopExecution::Graph] { - // Always requests a tool call, so the loop never finishes on its own - // and must hit the model-call cap. - let model = Arc::new(MockModel::with_tool_call("lookup", json!({}))); - let mut harness = harness_for(execution, model); - harness.register_tool(Arc::new(FakeTool::returning("lookup", "tool-output"))); - - let ctx = RunContext::new(RunConfig::new("limit-stop").with_max_model_calls(2), ()); - let err = harness - .invoke_in_context(&(), ctx, vec![Message::user("loop forever")]) - .await - .expect_err("run hits the model-call cap"); - - match execution { - LoopExecution::Direct => direct_err = Some(err), - LoopExecution::Graph => graph_err = Some(err), - } - } - - let direct_err = direct_err.unwrap(); - let graph_err = graph_err.unwrap(); - assert!( - matches!(direct_err, TinyAgentsError::LimitExceeded(_)), - "direct: {direct_err:?}" - ); - assert!( - matches!(graph_err, TinyAgentsError::LimitExceeded(_)), - "graph: {graph_err:?}" - ); -} - -// ── Scenario 4: approval interrupt ────────────────────────────────────────── - -mod interrupt_middleware { - use async_trait::async_trait; - use tinyagents_harness::context::{MiddlewareControl, RunContext}; - use tinyagents_harness::middleware::Middleware; - use tinyinference_llm::model::ModelResponse; - - pub struct RequireApproval; - - #[async_trait] - impl Middleware<(), ()> for RequireApproval { - fn name(&self) -> &str { - "require_approval" - } - - async fn after_model( - &self, - ctx: &mut RunContext<()>, - _state: &(), - _response: &mut ModelResponse, - ) -> tinyagents_harness::Result<()> { - ctx.request_control(MiddlewareControl::Interrupt { - node: "review".into(), - message: "needs approval".into(), - }); - Ok(()) - } - } -} - -#[tokio::test] -async fn approval_interrupt_scenario_matches_direct_and_graph() { - for execution in [LoopExecution::Direct, LoopExecution::Graph] { - let model = Arc::new(MockModel::constant("hi")); - let mut harness = harness_for(execution, model); - harness.push_middleware(Arc::new(interrupt_middleware::RequireApproval)); - - let err = harness - .invoke_default(&(), vec![Message::user("do the risky thing")]) - .await - .expect_err("both engines surface MiddlewareControl::Interrupt as an error"); - - match err { - TinyAgentsError::Interrupted { node, message } => { - assert_eq!(node, "review"); - assert_eq!(message, "needs approval"); - } - other => panic!("expected Interrupted, got {other:?}"), - } - } -} - -// ── Scenario 5: steering inject ───────────────────────────────────────────── - -#[tokio::test] -async fn steering_inject_scenario_matches_direct_and_graph() { - let mut direct_messages = None; - let mut graph_messages = None; - - for execution in [LoopExecution::Direct, LoopExecution::Graph] { - let model = Arc::new(MockModel::with_responses(vec![ModelResponse::assistant( - "done", - )])); - let harness = harness_for(execution, model.clone()); - - let steering = SteeringHandle::allow_all(); - steering.send(SteeringCommand::InjectMessage(Message::user( - "ORCHESTRATOR: answer in French.", - ))); - let ctx = RunContext::new(RunConfig::new("steer-inject"), ()).with_steering(steering); - - let run = harness - .invoke_in_context(&(), ctx, vec![Message::user("hello")]) - .await - .expect("run completes"); - assert_eq!(run.model_calls, 1); - - let injected = run - .messages - .iter() - .any(|message| message.text().contains("ORCHESTRATOR")); - assert!( - injected, - "the injected steering message must reach the transcript" - ); - - match execution { - LoopExecution::Direct => direct_messages = Some(run.messages), - LoopExecution::Graph => graph_messages = Some(run.messages), - } - } - - assert_eq!(direct_messages, graph_messages); -} - -// ── Scenario 6: output retry ──────────────────────────────────────────────── - -#[tokio::test] -async fn output_retry_scenario_matches_direct_and_graph() { - let schema = json!({ - "type": "object", - "properties": { "answer": { "type": "string" } }, - "required": ["answer"], - }); - - let mut direct_run = None; - let mut graph_run = None; - - for execution in [LoopExecution::Direct, LoopExecution::Graph] { - // First reply is not valid JSON for the schema; second reply repairs it. - let model = Arc::new(MockModel::with_responses(vec![ - ModelResponse::assistant("not json"), - ModelResponse::assistant(r#"{"answer":"fixed"}"#), - ])); - let mut harness = harness_for(execution, model); - let mut policy = harness.policy().clone(); - policy.default_response_format = Some(ResponseFormat::auto("answer", schema.clone())); - harness.with_policy(policy); - - let run = harness - .invoke_default(&(), vec![Message::user("what is the answer")]) - .await - .expect("run completes after one repair turn"); - - match execution { - LoopExecution::Direct => direct_run = Some(run), - LoopExecution::Graph => graph_run = Some(run), - } - } - - let (direct_run, graph_run) = (direct_run.unwrap(), graph_run.unwrap()); - assert_eq!(direct_run.structured, graph_run.structured); - assert_eq!(direct_run.structured, Some(json!({ "answer": "fixed" }))); - assert_eq!(direct_run.model_calls, graph_run.model_calls); - assert_eq!(direct_run.model_calls, 2); -} - -// ── Checkpoint + resume across an interrupt ───────────────────────────────── - -mod pause_middleware { - use async_trait::async_trait; - use tinyagents_harness::context::{MiddlewareControl, RunContext}; - use tinyagents_harness::middleware::Middleware; - use tinyinference_llm::model::ModelResponse; - - /// Requests a graph-level interrupt (an approval gate) on the first - /// `after_model` hook only. - pub struct PauseOnce(pub std::sync::atomic::AtomicBool); - - impl PauseOnce { - pub fn new() -> Self { - Self(std::sync::atomic::AtomicBool::new(true)) - } - } - - #[async_trait] - impl Middleware<(), ()> for PauseOnce { - fn name(&self) -> &str { - "pause_once" - } - - async fn after_model( - &self, - ctx: &mut RunContext<()>, - _state: &(), - _response: &mut ModelResponse, - ) -> tinyagents_harness::Result<()> { - if self.0.swap(false, std::sync::atomic::Ordering::SeqCst) { - ctx.request_control(MiddlewareControl::Interrupt { - node: "review".into(), - message: "needs approval".into(), - }); - } - Ok(()) - } - } -} - -/// Drives `compile_loop`'s real `CompiledGraph` directly (not through -/// `AgentHarness::invoke`/`GraphLoopDriver`), which is what actually gets -/// checkpoint/resume: A5 item 2's "approvals surfacing as graph interrupts". -async fn checkpoint_resume_with<C>(checkpointer: Arc<C>) -where - C: tinyagents_graph::checkpoint::Checkpointer<tinyagents_graph::agent_loop::LoopState> - + 'static, -{ - let model = Arc::new(MockModel::with_responses(vec![ - tool_call_response("call-1", "lookup", json!({})), - ModelResponse::assistant("done"), - ])); - let mut harness: AgentHarness<()> = AgentHarness::new(); - harness - .register_model("mock", model) - .set_default_model("mock") - .register_tool(Arc::new(FakeTool::returning("lookup", "tool-output"))) - .push_middleware(Arc::new(pause_middleware::PauseOnce::new())); - let harness = Arc::new(harness); - - let rt = Arc::new(LoopRuntime::for_run( - harness.clone(), - Arc::new(()), - RunContext::new(RunConfig::new("checkpoint-resume"), ()), - )); - let graph = compile_loop(rt) - .expect("graph compiles") - .with_checkpointer(checkpointer); - - let execution = graph - .run_with_thread( - "checkpoint-resume-thread", - LoopState::seed(vec![Message::user("look something up")]), - ) - .await - .expect("first leg completes to the interrupt"); - assert_eq!( - execution.interrupts.len(), - 1, - "run paused at the approval gate" - ); - assert!(!execution.state.finished); - - let resumed = graph - .resume( - "checkpoint-resume-thread", - tinyagents_graph::Command { - update: None, - goto: Vec::new(), - resume: Some(json!({ "approved": true })), - resume_by_task: Default::default(), - }, - ) - .await - .expect("resume drives the run to completion"); - assert!(resumed.state.finished); - assert_eq!(resumed.state.final_text.as_deref(), Some("done")); -} - -#[tokio::test] -async fn checkpoint_resume_across_interrupt_in_memory() { - checkpoint_resume_with(Arc::new(InMemoryCheckpointer::new())).await; -} - -#[tokio::test] -async fn checkpoint_resume_across_interrupt_sqlite() { - let dir = tempdir(); - checkpoint_resume_with(Arc::new(storage::checkpointer(dir.path()))).await; -} - -/// Minimal temp-dir helper (avoids pulling in the `tempfile` crate just for -/// this one test). -fn tempdir() -> TempDir { - let path = std::env::temp_dir().join(format!( - "loop_as_graph-{}-{}", - std::process::id(), - tinyagents_harness::ids::now_ms() - )); - std::fs::create_dir_all(&path).expect("create temp dir"); - TempDir(path) -} - -struct TempDir(std::path::PathBuf); - -impl TempDir { - fn path(&self) -> &std::path::Path { - &self.0 - } -} - -impl Drop for TempDir { - fn drop(&mut self) { - let _ = std::fs::remove_dir_all(&self.0); - } -} - -// ── `iter()` stepping with `override_next` ────────────────────────────────── - -#[tokio::test] -async fn iter_steps_node_by_node_and_honors_override_next() { - let model = Arc::new(MockModel::with_responses(vec![ - tool_call_response("call-1", "lookup", json!({})), - ModelResponse::assistant("done"), - ])); - let mut harness: AgentHarness<()> = AgentHarness::new(); - harness - .register_model("mock", model) - .set_default_model("mock") - .register_tool(Arc::new(FakeTool::returning("lookup", "tool-output"))); - let harness = Arc::new(harness); - - let ctx = RunContext::new(RunConfig::new("iter-stepping"), ()); - let mut iter = harness - .iter(Arc::new(()), ctx, vec![Message::user("look something up")]) - .expect("iter starts"); - - let step = iter.next().await.expect("plan step").expect("not finished"); - assert_eq!(step.node, node::PLAN); - assert_eq!(step.next.as_deref(), Some(node::MODEL)); - - let step = iter - .next() - .await - .expect("model step") - .expect("not finished"); - assert_eq!(step.node, node::MODEL); - assert_eq!(step.next.as_deref(), Some(node::TOOLS)); - - // Redirect the very next activation back to `plan` instead of the - // naturally-routed `tools` — exercising `override_next` — then let the - // (now unoverridden) routing carry the run to completion. - iter.override_next(node::PLAN); - let step = iter - .next() - .await - .expect("overridden step") - .expect("not finished"); - assert_eq!(step.node, node::PLAN); - assert_eq!(step.next.as_deref(), Some(node::MODEL)); - - let state = iter.run_to_end().await.expect("run finishes"); - assert!(state.finished); -} - -// ── Model profile preview (#6962) ─────────────────────────────────────────── - -/// Both loops expose the target model's profile to `before_model` middleware -/// on `RunContext::model_profile`, so middleware can avoid shapes the model -/// handles badly (a new system message on a model that hoists them). -#[tokio::test] -async fn before_model_sees_the_model_profile_in_direct_and_graph() { - use async_trait::async_trait; - use std::sync::Mutex; - use tinyagents_harness::middleware::Middleware; - use tinyinference_llm::model::{ChatModel, ModelProfile, ModelRequest}; - - struct SeenProfiles(Arc<Mutex<Vec<Option<ModelProfile>>>>); - - #[async_trait] - impl Middleware<(), ()> for SeenProfiles { - fn name(&self) -> &str { - "seen-profiles" - } - - async fn before_model( - &self, - ctx: &mut RunContext<()>, - _state: &(), - _request: &mut ModelRequest, - ) -> tinyagents_harness::Result<()> { - self.0.lock().unwrap().push(ctx.model_profile.clone()); - Ok(()) - } - } - - for execution in [LoopExecution::Direct, LoopExecution::Graph] { - let model = Arc::new(MockModel::with_responses(vec![ModelResponse::assistant( - "done", - )])); - let expected = ChatModel::<()>::profile(model.as_ref()).cloned(); - assert!(expected.is_some(), "the mock advertises a profile"); - let mut harness = harness_for(execution, model); - let seen = Arc::new(Mutex::new(Vec::new())); - harness.push_middleware(Arc::new(SeenProfiles(seen.clone()))); - - let ctx = RunContext::new(RunConfig::new("profile-preview"), ()); - harness - .invoke_in_context(&(), ctx, vec![Message::user("hi")]) - .await - .expect("run completes"); - - assert_eq!( - seen.lock().unwrap().as_slice(), - [expected], - "{execution:?} loop must expose the target profile to before_model" - ); - } -} - -/// A hint added by `before_model` must invalidate the graph loop's initial -/// resolution, because hints participate in registry selection just like an -/// explicit model name. -#[tokio::test] -async fn graph_reresolves_when_before_model_adds_a_model_hint() { - use async_trait::async_trait; - use tinyagents_harness::middleware::Middleware; - use tinyinference_llm::model::{ModelHint, ModelRequest}; - - struct SelectHint; - - #[async_trait] - impl Middleware<(), ()> for SelectHint { - fn name(&self) -> &str { - "select-hint" - } - - async fn before_model( - &self, - _ctx: &mut RunContext<()>, - _state: &(), - request: &mut ModelRequest, - ) -> tinyagents_harness::Result<()> { - request.model_hints.push(ModelHint { - model: "hinted".to_string(), - priority: 100, - reason: Some("test route".to_string()), - }); - Ok(()) - } - } - - let default = Arc::new(MockModel::with_responses(vec![ModelResponse::assistant( - "default", - )])); - let hinted = Arc::new(MockModel::with_responses(vec![ModelResponse::assistant( - "hinted", - )])); - let mut harness = harness_for(LoopExecution::Graph, default); - harness - .register_model("hinted", hinted) - .set_default_model("mock") - .push_middleware(Arc::new(SelectHint)); - - let run = harness - .invoke_default(&(), vec![Message::user("route me")]) - .await - .expect("run completes"); - - assert_eq!(run.text(), Some("hinted".to_string())); -} - -// ── Terminal outcome parity ───────────────────────────────────────────────── - -/// A `StopWithPartial` cap must be reported identically by both engines: a -/// limit outcome, not a plain completion. -#[tokio::test] -async fn model_cap_stop_reports_the_same_terminal_outcome_in_both_engines() { - use tinyagents_harness::events::LimitKind; - use tinyagents_harness::limits::{LimitBehavior, RunLimits}; - use tinyagents_harness::terminal::TerminalReason; - - for execution in [LoopExecution::Direct, LoopExecution::Graph] { - let model = Arc::new(MockModel::with_tool_call("spin", serde_json::json!({}))); - let mut harness = harness_for(execution, model); - harness.register_tool(Arc::new(tinyagents_harness::testkit::FakeTool::returning( - "spin", "again", - ))); - harness.with_policy(RunPolicy { - execution, - limits: RunLimits::default() - .with_max_model_calls(2) - .with_behavior(LimitBehavior::StopWithPartial), - ..RunPolicy::default() - }); - let run = harness - .invoke_default(&(), vec![Message::user("go")]) - .await - .expect("StopWithPartial completes the run"); - let outcome = run.terminal.expect("terminal outcome"); - assert_eq!( - outcome.reason, - TerminalReason::LimitReached(Some(LimitKind::ModelCalls)), - "{execution:?}" - ); - } -} - -/// The graph engine honors `StopWithPartial` for the tool-call cap (the direct -/// loop surfaces it as a `LimitExceeded` error carrying `ToolCalls`); the stop -/// must carry `ToolCalls`, not an untyped limit. -#[tokio::test] -async fn graph_tool_cap_stop_reports_a_typed_tool_calls_limit() { - use tinyagents_harness::events::LimitKind; - use tinyagents_harness::limits::{LimitBehavior, RunLimits}; - use tinyagents_harness::terminal::TerminalReason; - - for execution in [LoopExecution::Graph] { - let model = Arc::new(MockModel::with_tool_call("spin", serde_json::json!({}))); - let mut harness = harness_for(execution, model); - harness.register_tool(Arc::new(tinyagents_harness::testkit::FakeTool::returning( - "spin", "again", - ))); - harness.with_policy(RunPolicy { - execution, - limits: RunLimits::default() - .with_max_tool_calls(1) - .with_behavior(LimitBehavior::StopWithPartial), - ..RunPolicy::default() - }); - let run = harness - .invoke_default(&(), vec![Message::user("go")]) - .await - .expect("StopWithPartial completes the run"); - let outcome = run.terminal.expect("terminal outcome"); - assert_eq!( - outcome.reason, - TerminalReason::LimitReached(Some(LimitKind::ToolCalls)), - "{execution:?}" - ); - assert!( - outcome.message.contains("tool_calls"), - "{execution:?}: {outcome:?}" - ); - } -} - -struct BoomModel; - -#[async_trait::async_trait] -impl tinyinference_llm::model::ChatModel<()> for BoomModel { - async fn invoke( - &self, - _: &(), - _: tinyinference_llm::model::ModelRequest, - ) -> tinyinference_llm::Result<ModelResponse> { - Err(tinyinference_llm::Error::Model("model boom".into())) - } -} - -struct FailingAfterAgent; - -#[async_trait::async_trait] -impl tinyagents_harness::middleware::Middleware<(), ()> for FailingAfterAgent { - fn name(&self) -> &str { - "failing_after_agent" - } - - async fn after_agent( - &self, - _ctx: &mut RunContext<()>, - _state: &(), - _run: &mut tinyagents_harness::middleware::AgentRun, - ) -> tinyagents_harness::Result<()> { - Err(TinyAgentsError::Middleware("cleanup boom".into())) - } -} - -/// When the run already failed, a failing `after_agent` hook must not replace -/// the originating error, in either engine. -#[tokio::test] -async fn a_failing_after_agent_keeps_the_original_error_in_both_engines() { - for execution in [LoopExecution::Direct, LoopExecution::Graph] { - let mut harness: AgentHarness<()> = AgentHarness::new(); - harness - .register_model("mock", Arc::new(BoomModel)) - .set_default_model("mock"); - if matches!(execution, LoopExecution::Graph) { - harness.with_loop_driver(Arc::new(GraphLoopDriver::new())); - } - harness.with_policy(RunPolicy { - execution, - ..RunPolicy::default() - }); - harness.push_middleware(Arc::new(FailingAfterAgent)); - let error = harness - .invoke_default(&(), vec![Message::user("go")]) - .await - .expect_err("the run fails"); - assert!( - error.to_string().contains("model boom"), - "{execution:?}: {error}" - ); - } -} - -/// A failing `after_agent` hook on an otherwise successful run is the surfaced -/// error, so the partial run's terminal outcome must be that failure rather -/// than the stale completion recorded before the hook ran. -#[tokio::test] -async fn a_failing_after_agent_on_a_successful_run_records_a_failure_outcome() { - use tinyagents_harness::terminal::{TerminalClass, TerminalReason}; - - for execution in [LoopExecution::Direct, LoopExecution::Graph] { - let model = Arc::new(MockModel::with_responses(vec![ModelResponse::assistant( - "fine", - )])); - let mut harness = harness_for(execution, model); - harness.push_middleware(Arc::new(FailingAfterAgent)); - let ctx = RunContext::new(tinyagents_harness::context::RunConfig::new("aa"), ()); - let partial = harness - .invoke_in_context_collecting_partial(&(), ctx, vec![Message::user("go")]) - .await; - let error = partial.error.expect("the hook error surfaces"); - assert!(error.to_string().contains("cleanup boom"), "{execution:?}"); - let outcome = partial.run.terminal.expect("terminal outcome"); - assert_ne!(outcome.reason, TerminalReason::Completed, "{execution:?}"); - assert_eq!(outcome.class, TerminalClass::Failure, "{execution:?}"); - } -} - -/// Under `LimitBehavior::Error` the graph engine's tool-cap failure still -/// carries the concrete `ToolCalls` kind on the partial run's outcome. -#[tokio::test] -async fn graph_tool_cap_error_carries_the_tool_calls_kind() { - use tinyagents_harness::events::LimitKind; - use tinyagents_harness::limits::RunLimits; - use tinyagents_harness::terminal::TerminalReason; - - let model = Arc::new(MockModel::with_tool_call("spin", serde_json::json!({}))); - let mut harness = harness_for(LoopExecution::Graph, model); - harness.register_tool(Arc::new(tinyagents_harness::testkit::FakeTool::returning( - "spin", "again", - ))); - harness.with_policy(RunPolicy { - execution: LoopExecution::Graph, - limits: RunLimits::default().with_max_tool_calls(1), - ..RunPolicy::default() - }); - let ctx = RunContext::new(tinyagents_harness::context::RunConfig::new("cap"), ()); - let partial = harness - .invoke_in_context_collecting_partial(&(), ctx, vec![Message::user("go")]) - .await; - assert!(partial.error.is_some()); - assert_eq!( - partial.run.terminal.expect("outcome").reason, - TerminalReason::LimitReached(Some(LimitKind::ToolCalls)) - ); -} - -struct LimitFromHook; - -#[async_trait::async_trait] -impl tinyagents_harness::middleware::Middleware<(), ()> for LimitFromHook { - fn name(&self) -> &str { - "limit_from_hook" - } - - async fn before_tool( - &self, - _ctx: &mut RunContext<()>, - _state: &(), - _call: &mut tinyinference_llm::tool::ToolCall, - ) -> tinyagents_harness::Result<()> { - Err(TinyAgentsError::LimitExceeded("policy budget".into())) - } -} - -/// A `LimitExceeded` raised by middleware is not the tool-call cap, so under -/// `StopWithPartial` it must fail the run in both engines instead of becoming -/// a partial stop labeled `ToolCalls`. -#[tokio::test] -async fn a_middleware_limit_error_is_not_a_tool_cap_partial_stop() { - use tinyagents_harness::limits::{LimitBehavior, RunLimits}; - - for execution in [LoopExecution::Direct, LoopExecution::Graph] { - let model = Arc::new(MockModel::with_tool_call("spin", serde_json::json!({}))); - let mut harness = harness_for(execution, model); - harness.register_tool(Arc::new(tinyagents_harness::testkit::FakeTool::returning( - "spin", "again", - ))); - harness.push_middleware(Arc::new(LimitFromHook)); - harness.with_policy(RunPolicy { - execution, - limits: RunLimits::default().with_behavior(LimitBehavior::StopWithPartial), - ..RunPolicy::default() - }); - let result = harness.invoke_default(&(), vec![Message::user("go")]).await; - assert!( - matches!(result, Err(TinyAgentsError::LimitExceeded(_))), - "{execution:?}: {result:?}" - ); - } -} diff --git a/crates/tinyagents-integration-tests/tests/persistence_conformance.rs b/crates/tinyagents-integration-tests/tests/persistence_conformance.rs deleted file mode 100644 index c31654fe6..000000000 --- a/crates/tinyagents-integration-tests/tests/persistence_conformance.rs +++ /dev/null @@ -1,53 +0,0 @@ -//! Checkpointer conformance for the lineage and pending-writes contracts. -//! -//! The pre-existing `tests/conformance.rs` covers put/get/list/delete_thread/ -//! prune. These two suites cover what it did not — `get_tuple`, -//! `state_history`, `copy_thread`, `delete_checkpoints` and the whole -//! pending-writes protocol — against every bundled backend, so a defect in -//! any one of them cannot hide behind the others. - -mod common; - -use common::storage; -use tinyagents_graph::checkpoint::{DriverCheckpointer, InMemoryCheckpointer}; -use tinyagents_graph::testkit::conformance::{ - checkpointer_lineage_contract, checkpointer_writes_contract, -}; - -#[tokio::test] -async fn in_memory_checkpointer_satisfies_the_writes_contract() { - checkpointer_writes_contract(InMemoryCheckpointer::<i32>::new()).await; -} - -#[tokio::test] -async fn in_memory_checkpointer_satisfies_the_lineage_contract() { - checkpointer_lineage_contract(InMemoryCheckpointer::<i32>::new()).await; -} - -#[tokio::test] -async fn memory_driver_checkpointer_satisfies_the_writes_contract() { - checkpointer_writes_contract(DriverCheckpointer::<i32>::new(storage::docs( - &storage::memory(), - ))) - .await; -} - -#[tokio::test] -async fn memory_driver_checkpointer_satisfies_the_lineage_contract() { - checkpointer_lineage_contract(DriverCheckpointer::<i32>::new(storage::docs( - &storage::memory(), - ))) - .await; -} - -#[tokio::test] -async fn sqlite_driver_checkpointer_satisfies_the_writes_contract() { - let dir = tempfile::tempdir().unwrap(); - checkpointer_writes_contract(storage::checkpointer::<i32>(dir.path())).await; -} - -#[tokio::test] -async fn sqlite_driver_checkpointer_satisfies_the_lineage_contract() { - let dir = tempfile::tempdir().unwrap(); - checkpointer_lineage_contract(storage::checkpointer::<i32>(dir.path())).await; -} diff --git a/crates/tinyagents-integration-tests/tests/persistence_store.rs b/crates/tinyagents-integration-tests/tests/persistence_store.rs deleted file mode 100644 index 1f7422043..000000000 --- a/crates/tinyagents-integration-tests/tests/persistence_store.rs +++ /dev/null @@ -1,140 +0,0 @@ -//! Regression tests for the driver-backed checkpointer and append store: the -//! identity questions the old file backends once got wrong (thread ids that -//! differ only by case, the empty thread id, offsets handed out by several -//! handles on one store) must stay right on every backend. - -mod common; - -use std::sync::Arc; - -use common::storage; -use tinyagents_graph::checkpoint::{Checkpoint, Checkpointer, PendingActivation}; -use tinyagents_harness::ids::NodeId; -use tinyagents_harness::store::{AppendStore, DriverAppendStore}; - -fn checkpoint(thread: &str, id: &str) -> Checkpoint<i32> { - Checkpoint::new( - 1, - 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_metadata(serde_json::json!({ "source": "loop", "step": 1 })) -} - -/// Thread ids that differ only by case must not share state: on a -/// case-insensitive filesystem `Alice` and `alice` were the same file, which -/// silently merged two unrelated lineages. -#[tokio::test] -async fn thread_ids_differing_only_by_case_do_not_share_state() { - let dir = tempfile::tempdir().unwrap(); - let cp = storage::checkpointer::<i32>(dir.path()); - cp.put(checkpoint("Alice", "upper")).await.unwrap(); - cp.put(checkpoint("alice", "lower")).await.unwrap(); - - let upper = cp.list("Alice").await.unwrap(); - let lower = cp.list("alice").await.unwrap(); - assert_eq!(upper.len(), 1, "`Alice` holds only its own checkpoint"); - assert_eq!(lower.len(), 1, "`alice` holds only its own checkpoint"); - assert_eq!(upper[0].checkpoint_id, "upper"); - assert_eq!(lower[0].checkpoint_id, "lower"); -} - -/// The empty thread id is a valid thread and must be visible to listing. -#[tokio::test] -async fn the_empty_thread_id_is_visible_to_list_threads() { - let dir = tempfile::tempdir().unwrap(); - let cp = storage::checkpointer::<i32>(dir.path()); - cp.put(checkpoint("", "c1")).await.unwrap(); - let threads = cp.list_threads().await.unwrap(); - assert!( - threads.iter().any(|t| t.is_empty()), - "list_threads reports the empty thread id: {threads:?}" - ); -} - -/// Separators, dots and non-ASCII are ordinary characters in a thread id. -#[tokio::test] -async fn thread_ids_with_separators_and_unicode_round_trip() { - let dir = tempfile::tempdir().unwrap(); - let cp = storage::checkpointer::<i32>(dir.path()); - let ids = ["a/b", "a\\b", "..", "con", "naïve ☃", "with space", "x.y.z"]; - for (n, id) in ids.iter().enumerate() { - cp.put(checkpoint(id, &format!("c{n}"))).await.unwrap(); - } - let mut listed = cp.list_threads().await.unwrap(); - listed.sort(); - let mut expected: Vec<String> = ids.iter().map(|s| s.to_string()).collect(); - expected.sort(); - assert_eq!(listed, expected); - for (n, id) in ids.iter().enumerate() { - let checkpoints = cp.list(id).await.unwrap(); - assert_eq!(checkpoints.len(), 1, "{id}"); - assert_eq!(checkpoints[0].checkpoint_id, format!("c{n}")); - } -} - -/// One undecodable record must not break listing for every other thread. -#[tokio::test] -async fn one_undecodable_record_does_not_break_other_threads() { - use tinystoragedrivers_api::DocumentStore; - use tinystoragedrivers_types::Precondition; - let dir = tempfile::tempdir().unwrap(); - let scoped = storage::sqlite(dir.path(), "checkpoints.sqlite"); - let docs = storage::docs(&scoped); - let cp = tinyagents_graph::checkpoint::DriverCheckpointer::<i32>::new(docs.clone()); - cp.put(checkpoint("good", "c1")).await.unwrap(); - docs.put( - "graph_checkpoints", - "poison", - serde_json::json!({"thread": "bad", "namespace": "", "seq": 0, "checkpoint_id": "x", "record": {"nope": true}}), - Precondition::None, - ) - .await - .unwrap(); - let threads = cp.list_threads().await.expect("listing survives"); - assert!(threads.iter().any(|t| t == "good")); - assert_eq!(cp.list("good").await.unwrap().len(), 1); -} - -/// Two store handles over one database must keep handing out fresh offsets: -/// a per-instance counter would re-issue offsets the first had already used. -#[tokio::test] -async fn offsets_stay_unique_across_store_instances() { - let dir = tempfile::tempdir().unwrap(); - let streams = || Arc::clone(storage::sqlite(dir.path(), "journal.sqlite").streams()); - let first = DriverAppendStore::new(streams()); - let second = DriverAppendStore::new(streams()); - - let mut offsets = Vec::new(); - offsets.push(first.append("s", serde_json::json!(0)).await.unwrap()); - offsets.push(second.append("s", serde_json::json!(1)).await.unwrap()); - offsets.push(first.append("s", serde_json::json!(2)).await.unwrap()); - offsets.push(second.append("s", serde_json::json!(3)).await.unwrap()); - - assert_eq!( - offsets, - vec![0, 1, 2, 3], - "a second instance continues the stream instead of restarting it" - ); - assert_eq!(second.len("s").await.unwrap(), 4); -} - -/// `read_from` resolves by offset, not by position. -#[tokio::test] -async fn read_from_resolves_by_offset_not_position() { - let store = DriverAppendStore::new(Arc::clone(storage::memory().streams())); - for n in 0..5 { - store.append("s", serde_json::json!(n)).await.unwrap(); - } - let tail = store.read_from("s", 3).await.unwrap(); - assert_eq!( - tail, - vec![(3, serde_json::json!(3)), (4, serde_json::json!(4))] - ); - assert!(store.read_from("s", 5).await.unwrap().is_empty()); -} diff --git a/crates/tinyagents-integration-tests/tests/runtime_primitives_resilience.rs b/crates/tinyagents-integration-tests/tests/runtime_primitives_resilience.rs index 96671fc6c..021986788 100644 --- a/crates/tinyagents-integration-tests/tests/runtime_primitives_resilience.rs +++ b/crates/tinyagents-integration-tests/tests/runtime_primitives_resilience.rs @@ -12,8 +12,7 @@ //! - LOOP-6 — `Started` events had no terminal partner on error paths. //! - LOOP-8 — steering applied part of a rejected batch, and a pause was unresumable. //! - LOOP-9 — a panicking listener wedged the event sink forever. -//! - C6 — `StreamChunk` had no producer. - +//! use std::sync::Arc; use std::time::Duration; @@ -32,8 +31,7 @@ use tinyagents_harness::steering::{ SteeringCommand, SteeringCommandKind, SteeringHandle, SteeringOutcome, SteeringPolicy, apply_pending_steering, }; -use tinyagents_harness::stream::{StreamChunk, StreamMode, StreamSink, project_event}; -use tinyinference_llm::message::{Message, MessageDelta}; +use tinyinference_llm::message::Message; // ── LOOP-2: jitter must never collapse the backoff ─────────────────────────── @@ -333,37 +331,6 @@ fn loop8_a_pause_is_latched_and_resumable_across_checkpoints() { assert!(handle.pause_state().is_none()); } -// ── C6: the stream projection ──────────────────────────────────────────────── - -#[test] -fn c6_events_project_onto_stream_chunks() { - let delta = AgentEvent::ModelDelta { - run_id: RunId::new("r1"), - call_id: CallId::new("c1"), - delta: MessageDelta::text("hi"), - }; - assert!(matches!( - project_event(&delta), - Some(StreamChunk::Message(_)) - )); - - // The Interrupt variant finally has a producer. - let interrupt = AgentEvent::ControlApplied { - control: "interrupt".into(), - detail: "needs approval".into(), - }; - assert!(matches!( - project_event(&interrupt), - Some(StreamChunk::Interrupt(_)) - )); - - // Mode filtering happens producer-side. - let sink = StreamSink::new([StreamMode::Interrupts]); - assert!(!sink.push_event(&delta)); - assert!(sink.push_event(&interrupt)); - assert_eq!(sink.drain().len(), 1); -} - // ── C5: the no-progress tracker is drivable from a hook ────────────────────── #[test] diff --git a/crates/tinyagents-integration-tests/tests/session_conformance.rs b/crates/tinyagents-integration-tests/tests/session_conformance.rs index 06a53bd30..3bcb4c0ae 100644 --- a/crates/tinyagents-integration-tests/tests/session_conformance.rs +++ b/crates/tinyagents-integration-tests/tests/session_conformance.rs @@ -5,7 +5,9 @@ //! defect in one cannot hide behind another. The run ledger is certified on the //! memory and sqlite drivers. `TranscriptHistory` genuinely has two //! implementations — the production driver-backed one and the in-memory test -//! double — and both are certified here. +//! double — and both are certified here. The driver-backed transcript suite +//! uses SQLite because generation writes require native lease-conditioned +//! document handles. mod common; @@ -57,7 +59,12 @@ fn contract_meta() -> TranscriptMeta { #[test] fn driver_transcript_history_satisfies_the_conformance_suite() { use tinyagents_session::port::{DriverSessionStores, SessionStoreProvider}; - let stores = DriverSessionStores::new(storage::fixed_factory(storage::memory())).unwrap(); + let dir = tempfile::tempdir().unwrap(); + let (factory, conditioned_factory) = + storage::conditioned_sqlite_factories(dir.path(), "transcripts.sqlite"); + let stores = DriverSessionStores::new(factory) + .unwrap() + .with_conditioned_storage_factory(conditioned_factory); let history = stores .for_agent("contract") .transcripts diff --git a/crates/tinyagents-integration-tests/tests/tool_deferral.rs b/crates/tinyagents-integration-tests/tests/tool_deferral.rs index c21d3c123..2137dcbce 100644 --- a/crates/tinyagents-integration-tests/tests/tool_deferral.rs +++ b/crates/tinyagents-integration-tests/tests/tool_deferral.rs @@ -241,8 +241,9 @@ async fn deferred_tool_is_promoted_after_search_and_restored_on_resume() { let mut harness: AgentHarness<()> = AgentHarness::new(); harness - .register_model("mock", model.clone()) - .set_default_model("mock") + .with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named("mock", model.clone()), + )) .register_tool(Arc::new(FakeTool::returning("read_file", "contents"))) .register_tool(deferred.clone()) .register_tool(unrelated.clone()) @@ -301,8 +302,9 @@ async fn deferred_tool_is_promoted_after_search_and_restored_on_resume() { let resumed_model = RecordingModel::new(vec![text("resumed")]); let mut resumed_harness: AgentHarness<()> = AgentHarness::new(); resumed_harness - .register_model("mock", resumed_model.clone()) - .set_default_model("mock") + .with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named("mock", resumed_model.clone()), + )) .register_tool(Arc::new(FakeTool::returning("read_file", "contents"))) .register_tool(Arc::new(ChangedDeferredTool)); resumed_harness.register_tool(unrelated); @@ -402,8 +404,9 @@ async fn disabled_discovery_drops_the_bridge_and_keeps_direct_calls_working() { let mut harness: AgentHarness<()> = AgentHarness::new(); harness - .register_model("mock", model.clone()) - .set_default_model("mock") + .with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named("mock", model.clone()), + )) .register_tool(Arc::new(FakeTool::returning("read_file", "contents"))) .register_tool(deferred.clone()) .with_policy(RunPolicy { @@ -438,8 +441,9 @@ async fn no_deferred_tools_means_no_bridge() { let model = RecordingModel::new(vec![text("done")]); let mut harness: AgentHarness<()> = AgentHarness::new(); harness - .register_model("mock", model.clone()) - .set_default_model("mock") + .with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named("mock", model.clone()), + )) .register_tool(Arc::new(FakeTool::returning("read_file", "contents"))); harness .invoke_default(&(), vec![Message::user("hi")]) @@ -468,8 +472,9 @@ async fn tool_searched_query_is_payload_free_by_default() { let mut harness: AgentHarness<()> = AgentHarness::new(); harness - .register_model("mock", model.clone()) - .set_default_model("mock") + .with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named("mock", model.clone()), + )) .register_tool(deferred) .push_middleware(Arc::new(CaptureMiddleware { listener: listener.clone(), @@ -510,8 +515,9 @@ async fn host_registered_tool_search_wins_over_the_intrinsic_bridge() { let mut harness: AgentHarness<()> = AgentHarness::new(); harness - .register_model("mock", model.clone()) - .set_default_model("mock") + .with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named("mock", model.clone()), + )) .register_tool(deferred.clone()) .register_tool(host_search.clone()); @@ -545,8 +551,9 @@ async fn unknown_tool_corrective_does_not_advertise_a_host_registered_tool_searc let mut harness: AgentHarness<()> = AgentHarness::new(); harness - .register_model("mock", model.clone()) - .set_default_model("mock") + .with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named("mock", model.clone()), + )) .register_tool(deferred) .register_tool(host_search) .with_policy(RunPolicy { @@ -598,8 +605,9 @@ async fn hidden_registration_suppresses_the_matching_bridge_schema() { let mut harness: AgentHarness<()> = AgentHarness::new(); harness - .register_model("mock", model.clone()) - .set_default_model("mock") + .with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named("mock", model.clone()), + )) .register_tool(deferred.clone()) .register_tool(hidden_search.clone()); @@ -634,8 +642,9 @@ async fn tool_schemas_projection_applies_to_wire_and_catalog() { let mut harness: AgentHarness<()> = AgentHarness::new(); harness - .register_model("mock", model.clone()) - .set_default_model("mock") + .with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named("mock", model.clone()), + )) .register_tool(direct) .register_tool(deferred) .register_tool(ExposedTool::new( @@ -711,8 +720,9 @@ async fn a_deferred_event_is_emitted_only_for_a_deferred_tool_called_by_name() { let mut harness: AgentHarness<()> = AgentHarness::new(); harness - .register_model("mock", model.clone()) - .set_default_model("mock") + .with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named("mock", model.clone()), + )) .register_tool(direct.clone()) .register_tool(deferred.clone()) .register_tool(hidden.clone()) diff --git a/crates/tinyagents-integration-tests/tests/verify_before_finish_execution.rs b/crates/tinyagents-integration-tests/tests/verify_before_finish_execution.rs index 8275d54df..092e308b5 100644 --- a/crates/tinyagents-integration-tests/tests/verify_before_finish_execution.rs +++ b/crates/tinyagents-integration-tests/tests/verify_before_finish_execution.rs @@ -9,6 +9,7 @@ use tinytools::{PermissionLevel, Tool, ToolResult}; use tinyagents_harness::Result; use tinyagents_harness::context::RunContext; +use tinyagents_harness::host::FixedModelResolver; use tinyagents_harness::limits::RunLimits; use tinyagents_harness::middleware::{ AgentRun, FinishActivity, Middleware, MiddlewareToolOutcome, ToolHandler, @@ -130,7 +131,7 @@ fn rig(calls: Vec<ToolCall>) -> Rig { ]); handle.bind(Arc::downgrade(&packed)); let mut harness = AgentHarness::new(); - harness.register_model("mock", model); + harness.with_model_resolver(Arc::new(FixedModelResolver::new(model))); harness.register_tool(Arc::new(UseSkillTool::new( handle, PackCatalog::new(PACKS, "[missing]"), diff --git a/crates/tinyagents-integration-tests/tests/wave2_cache_key_scope.rs b/crates/tinyagents-integration-tests/tests/wave2_cache_key_scope.rs index 8e0cb5d44..48b092c6e 100644 --- a/crates/tinyagents-integration-tests/tests/wave2_cache_key_scope.rs +++ b/crates/tinyagents-integration-tests/tests/wave2_cache_key_scope.rs @@ -63,25 +63,29 @@ async fn shared_cache_does_not_cross_serve_between_two_providers() { let local_calls = Arc::new(AtomicUsize::new(0)); let mut hosted: AgentHarness<()> = AgentHarness::new(); - hosted.register_model( - "chat", - Arc::new(IdentifiedModel { - identity: "openai|gpt-5|https://api.openai.com/v1||abc".to_string(), - answer: "hosted answer".to_string(), - calls: hosted_calls.clone(), - }), - ); + hosted.with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named( + "chat", + Arc::new(IdentifiedModel { + identity: "openai|gpt-5|https://api.openai.com/v1||abc".to_string(), + answer: "hosted answer".to_string(), + calls: hosted_calls.clone(), + }), + ), + )); hosted.with_response_cache(cache.clone()); let mut local: AgentHarness<()> = AgentHarness::new(); - local.register_model( - "chat", - Arc::new(IdentifiedModel { - identity: "ollama|llama3.2|http://localhost:11434/v1||no-credential".to_string(), - answer: "local answer".to_string(), - calls: local_calls.clone(), - }), - ); + local.with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named( + "chat", + Arc::new(IdentifiedModel { + identity: "ollama|llama3.2|http://localhost:11434/v1||no-credential".to_string(), + answer: "local answer".to_string(), + calls: local_calls.clone(), + }), + ), + )); local.with_response_cache(cache.clone()); let hosted_run = hosted @@ -112,14 +116,16 @@ async fn identical_identity_still_shares_the_cache() { let build = |calls: Arc<AtomicUsize>| { let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( - "chat", - Arc::new(IdentifiedModel { - identity: identity.to_string(), - answer: "same answer".to_string(), - calls, - }), - ); + harness.with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named( + "chat", + Arc::new(IdentifiedModel { + identity: identity.to_string(), + answer: "same answer".to_string(), + calls, + }), + ), + )); harness.with_response_cache(cache.clone()); harness }; diff --git a/crates/tinyagents-integration-tests/tests/wave2_cache_layout.rs b/crates/tinyagents-integration-tests/tests/wave2_cache_layout.rs index 86b9fd61e..755349437 100644 --- a/crates/tinyagents-integration-tests/tests/wave2_cache_layout.rs +++ b/crates/tinyagents-integration-tests/tests/wave2_cache_layout.rs @@ -378,7 +378,9 @@ mod run_policy_breakpoints { seen: Mutex::new(Vec::new()), }); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("rec", model.clone()); + harness.with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named("rec", model.clone()), + )); harness.push_middleware(Arc::new(DeclareSystemPrefix)); harness.with_policy(RunPolicy { cache: policy, diff --git a/crates/tinyagents-integration-tests/tests/wave2_cache_loop.rs b/crates/tinyagents-integration-tests/tests/wave2_cache_loop.rs index 04f8013c9..6697b10f2 100644 --- a/crates/tinyagents-integration-tests/tests/wave2_cache_loop.rs +++ b/crates/tinyagents-integration-tests/tests/wave2_cache_loop.rs @@ -11,7 +11,6 @@ use std::sync::atomic::{AtomicUsize, Ordering}; use async_trait::async_trait; -use tinyagents_graph::*; use tinyagents_harness::cache::{InMemoryResponseCache, ResponseCache}; use tinyagents_harness::context::{RunConfig, RunContext}; use tinyagents_harness::middleware::ModelFallbackMiddleware; @@ -124,22 +123,22 @@ async fn a_fallback_answer_is_never_cached_under_the_primary_key() { let cache = Arc::new(InMemoryResponseCache::new()); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( - "primary", - Arc::new(AlwaysFailing { - identity: "primary-identity", - calls: primary_calls.clone(), - }), - ); - harness.register_model( - "backup", - Arc::new(FixedModel::new( - "backup-identity", - "backup answer", - backup_calls.clone(), - )), - ); - harness.set_default_model("primary"); + let primary = Arc::new(AlwaysFailing { + identity: "primary-identity", + calls: primary_calls.clone(), + }); + let backup = Arc::new(FixedModel::new( + "backup-identity", + "backup answer", + backup_calls.clone(), + )); + harness.with_model_resolver(Arc::new( + tinyagents_harness::host::PrebuiltModelResolver::new( + primary as Arc<dyn ChatModel<()>>, + [("backup".to_string(), backup as Arc<dyn ChatModel<()>>)], + ) + .with_primary_name("primary"), + )); harness.with_response_cache(cache.clone()); harness.with_policy(RunPolicy { retry: RetryPolicy::default().with_max_attempts(1), @@ -184,10 +183,12 @@ async fn a_cache_hit_is_marked_served_from_cache() { let calls = Arc::new(AtomicUsize::new(0)); let cache = Arc::new(InMemoryResponseCache::new()); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( - "chat", - Arc::new(FixedModel::new("id", "answer", calls.clone())), - ); + harness.with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named( + "chat", + Arc::new(FixedModel::new("id", "answer", calls.clone())), + ), + )); harness.with_response_cache(cache); let cold = harness @@ -225,10 +226,12 @@ async fn a_streaming_cache_hit_still_emits_deltas() { let calls = Arc::new(AtomicUsize::new(0)); let cache = Arc::new(InMemoryResponseCache::new()); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( - "chat", - Arc::new(FixedModel::new("id", "streamed answer", calls.clone())), - ); + harness.with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named( + "chat", + Arc::new(FixedModel::new("id", "streamed answer", calls.clone())), + ), + )); harness.with_response_cache(cache); let cold_events = EventRecorder::new(); @@ -280,10 +283,12 @@ async fn a_streaming_run_does_not_reuse_a_unary_cache_entry() { let calls = Arc::new(AtomicUsize::new(0)); let cache = Arc::new(InMemoryResponseCache::new()); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( - "chat", - Arc::new(FixedModel::new("id", "answer", calls.clone())), - ); + harness.with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named( + "chat", + Arc::new(FixedModel::new("id", "answer", calls.clone())), + ), + )); harness.with_response_cache(cache); harness @@ -314,10 +319,12 @@ async fn a_failing_cache_never_fails_the_run() { let calls = Arc::new(AtomicUsize::new(0)); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( - "chat", - Arc::new(FixedModel::new("id", "answer", calls.clone())), - ); + harness.with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named( + "chat", + Arc::new(FixedModel::new("id", "answer", calls.clone())), + ), + )); harness.with_response_cache(Arc::new(BrokenCache { gets: gets.clone(), puts: puts.clone(), @@ -347,22 +354,22 @@ async fn model_fallback_middleware_actually_switches_models() { let backup_calls = Arc::new(AtomicUsize::new(0)); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( - "primary", - Arc::new(AlwaysFailing { - identity: "primary-identity", - calls: primary_calls.clone(), - }), - ); - harness.register_model( - "backup", - Arc::new(FixedModel::new( - "backup-identity", - "backup answer", - backup_calls.clone(), - )), - ); - harness.set_default_model("primary"); + let primary = Arc::new(AlwaysFailing { + identity: "primary-identity", + calls: primary_calls.clone(), + }); + let backup = Arc::new(FixedModel::new( + "backup-identity", + "backup answer", + backup_calls.clone(), + )); + harness.with_model_resolver(Arc::new( + tinyagents_harness::host::PrebuiltModelResolver::new( + primary as Arc<dyn ChatModel<()>>, + [("backup".to_string(), backup as Arc<dyn ChatModel<()>>)], + ) + .with_primary_name("primary"), + )); // No harness-level `FallbackPolicy` — the switch must come from the wrap // middleware alone, which steers by mutating `request.model`. harness.with_policy(RunPolicy { @@ -396,21 +403,22 @@ async fn model_fallback_middleware_actually_switches_models() { #[tokio::test] async fn an_unresolvable_wrap_override_keeps_the_resolved_binding() { - // Fail-closed: naming a model the registry cannot resolve must not silently + // Fail-closed: naming a model the resolver cannot resolve must not silently // substitute a different one — it keeps the resolved binding and makes the // skip observable, matching what `run_loop` already does for a pre-wrap // override. let primary_calls = Arc::new(AtomicUsize::new(0)); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( - "primary", - Arc::new(AlwaysFailing { - identity: "primary-identity", - calls: primary_calls.clone(), - }), - ); - harness.set_default_model("primary"); + harness.with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named( + "primary", + Arc::new(AlwaysFailing { + identity: "primary-identity", + calls: primary_calls.clone(), + }), + ), + )); harness.with_policy(RunPolicy { retry: RetryPolicy::default() .with_max_attempts(1) diff --git a/crates/tinyagents-integration-tests/tests/wave2_loop_cache_accounting.rs b/crates/tinyagents-integration-tests/tests/wave2_loop_cache_accounting.rs index bfa82ba34..0cf575783 100644 --- a/crates/tinyagents-integration-tests/tests/wave2_loop_cache_accounting.rs +++ b/crates/tinyagents-integration-tests/tests/wave2_loop_cache_accounting.rs @@ -35,7 +35,12 @@ impl ChatModel<()> for FlaggedModel { async fn run_usage(served_from_cache: bool) -> tinyinference_llm::usage::UsageTotals { let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("m", Arc::new(FlaggedModel { served_from_cache })); + harness.with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named( + "m", + Arc::new(FlaggedModel { served_from_cache }), + ), + )); harness .invoke_default(&(), vec![Message::user("go")]) .await @@ -73,12 +78,14 @@ async fn a_cache_served_response_does_not_consume_the_budget() { let tracker = budget.tracker(); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( - "m", - Arc::new(FlaggedModel { - served_from_cache: true, - }), - ); + harness.with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named( + "m", + Arc::new(FlaggedModel { + served_from_cache: true, + }), + ), + )); harness.push_middleware(Arc::new(budget)); for _ in 0..5 { diff --git a/crates/tinyagents-integration-tests/tests/wave2_loop_control.rs b/crates/tinyagents-integration-tests/tests/wave2_loop_control.rs index ed43f0e2e..66c3c5136 100644 --- a/crates/tinyagents-integration-tests/tests/wave2_loop_control.rs +++ b/crates/tinyagents-integration-tests/tests/wave2_loop_control.rs @@ -51,7 +51,9 @@ impl Middleware<(), ()> for StopAfterToolMiddleware { fn spinning_harness() -> (AgentHarness<()>, Arc<MockModel>) { let mut harness: AgentHarness<()> = AgentHarness::new(); let model = Arc::new(MockModel::with_tool_call("spin", json!({}))); - harness.register_model("mock", model.clone()); + harness.with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named("mock", model.clone()), + )); harness.register_tool(Arc::new(FakeTool::returning("spin", "again"))); (harness, model) } @@ -143,7 +145,12 @@ async fn a_steering_pause_is_distinguishable_from_a_clean_finish() { #[tokio::test] async fn an_unsteered_run_still_reports_completed() { let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", Arc::new(MockModel::constant("done"))); + harness.with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named( + "mock", + Arc::new(MockModel::constant("done")), + ), + )); let result = harness .invoke_with_status(&(), (), RunConfig::new("plain"), vec![Message::user("go")]) diff --git a/crates/tinyagents-integration-tests/tests/wave2_loop_limits.rs b/crates/tinyagents-integration-tests/tests/wave2_loop_limits.rs index 4b1834182..893c7408c 100644 --- a/crates/tinyagents-integration-tests/tests/wave2_loop_limits.rs +++ b/crates/tinyagents-integration-tests/tests/wave2_loop_limits.rs @@ -25,7 +25,9 @@ use tinyinference_llm::providers::MockModel; fn spinning_harness() -> (AgentHarness<()>, Arc<MockModel>) { let mut harness: AgentHarness<()> = AgentHarness::new(); let model = Arc::new(MockModel::with_tool_call("spin", json!({}))); - harness.register_model("mock", model.clone()); + harness.with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named("mock", model.clone()), + )); harness.register_tool(Arc::new(FakeTool::returning("spin", "again"))); (harness, model) } diff --git a/crates/tinyagents-integration-tests/tests/wave2_loop_recovery.rs b/crates/tinyagents-integration-tests/tests/wave2_loop_recovery.rs index 0c5c1151a..b4d1023f0 100644 --- a/crates/tinyagents-integration-tests/tests/wave2_loop_recovery.rs +++ b/crates/tinyagents-integration-tests/tests/wave2_loop_recovery.rs @@ -102,10 +102,12 @@ impl Tool for WeatherTool { #[tokio::test] async fn an_unknown_tool_name_recovers_by_default() { let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( - "mock", - Arc::new(TwoTurnModel::new(ToolCall::new("c1", "nope", json!({})))), - ); + harness.with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named( + "mock", + Arc::new(TwoTurnModel::new(ToolCall::new("c1", "nope", json!({})))), + ), + )); harness.register_tool(Arc::new(WeatherTool)); let run = harness @@ -131,14 +133,16 @@ async fn an_unknown_tool_name_recovers_by_default() { #[tokio::test] async fn invalid_tool_arguments_recover_by_default() { let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( - "mock", - Arc::new(TwoTurnModel::new(ToolCall::new( - "c1", - "weather", - json!({ "city": 5 }), - ))), - ); + harness.with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named( + "mock", + Arc::new(TwoTurnModel::new(ToolCall::new( + "c1", + "weather", + json!({ "city": 5 }), + ))), + ), + )); harness.register_tool(Arc::new(WeatherTool)); let run = harness @@ -160,10 +164,12 @@ async fn invalid_tool_arguments_recover_by_default() { #[tokio::test] async fn the_fail_variants_are_still_available() { let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( - "mock", - Arc::new(TwoTurnModel::new(ToolCall::new("c1", "nope", json!({})))), - ); + harness.with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named( + "mock", + Arc::new(TwoTurnModel::new(ToolCall::new("c1", "nope", json!({})))), + ), + )); harness.register_tool(Arc::new(WeatherTool)); harness.with_policy(RunPolicy { unknown_tool: UnknownToolPolicy::Fail, diff --git a/crates/tinyagents-integration-tests/tests/wave2_loop_structured.rs b/crates/tinyagents-integration-tests/tests/wave2_loop_structured.rs index c3a254099..eedcb337d 100644 --- a/crates/tinyagents-integration-tests/tests/wave2_loop_structured.rs +++ b/crates/tinyagents-integration-tests/tests/wave2_loop_structured.rs @@ -106,7 +106,9 @@ async fn structured_tool_strategy_does_not_force_the_schema_tool_when_real_tools ])])); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("rec", model.clone()); + harness.with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named("rec", model.clone()), + )); harness.register_tool(Arc::new(FakeTool::returning("search", "hits"))); harness.with_policy(RunPolicy { default_response_format: Some(tinyinference_llm::model::ResponseFormat::auto( @@ -149,7 +151,9 @@ async fn structured_tool_strategy_still_forces_the_schema_tool_with_no_real_tool ])])); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("rec", model.clone()); + harness.with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named("rec", model.clone()), + )); harness.with_policy(RunPolicy { default_response_format: Some(tinyinference_llm::model::ResponseFormat::auto( "result", @@ -194,7 +198,9 @@ async fn a_structured_hit_does_not_discard_sibling_real_tool_calls() { let search = Arc::new(FakeTool::returning("search", "hits")); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("rec", model.clone()); + harness.with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named("rec", model.clone()), + )); harness.register_tool(search.clone()); harness.with_policy(RunPolicy { default_response_format: Some(tinyinference_llm::model::ResponseFormat::auto( @@ -230,7 +236,9 @@ async fn a_schema_name_colliding_with_a_registered_tool_fails_closed() { let model = Arc::new(RecordingModel::new(vec![ModelResponse::assistant("hi")])); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("rec", model); + harness.with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named("rec", model), + )); harness.register_tool(Arc::new(FakeTool::returning("result", "hits"))); harness.with_policy(RunPolicy { default_response_format: Some(tinyinference_llm::model::ResponseFormat::auto( @@ -294,7 +302,9 @@ async fn a_schema_name_colliding_with_the_discovery_bridge_fails_closed() { let model = Arc::new(RecordingModel::new(vec![ModelResponse::assistant("hi")])); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("rec", model); + harness.with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named("rec", model), + )); harness.register_tool(Arc::new(DeferredStockQuote)); harness.with_policy(RunPolicy { default_response_format: Some(tinyinference_llm::model::ResponseFormat::auto( diff --git a/crates/tinyagents-integration-tests/tests/wave2_tools_execution.rs b/crates/tinyagents-integration-tests/tests/wave2_tools_execution.rs index 6da2ed26d..ca7f0eee7 100644 --- a/crates/tinyagents-integration-tests/tests/wave2_tools_execution.rs +++ b/crates/tinyagents-integration-tests/tests/wave2_tools_execution.rs @@ -217,13 +217,15 @@ fn tool_message_ids(messages: &[Message]) -> Vec<String> { #[tokio::test] async fn tool_result_call_id_is_overwritten_with_the_admitted_id() { let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( - "mock", - Arc::new(MockModel::with_responses(vec![ - tool_calls_response(vec![ToolCall::new("call_abc", "wrong_id", json!({}))]), - text_response("done"), - ])), - ); + harness.with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named( + "mock", + Arc::new(MockModel::with_responses(vec![ + tool_calls_response(vec![ToolCall::new("call_abc", "wrong_id", json!({}))]), + text_response("done"), + ])), + ), + )); harness.register_tool(Arc::new(WrongCallIdTool)); let run = harness @@ -244,14 +246,16 @@ async fn tool_result_call_id_is_overwritten_with_the_admitted_id() { async fn concurrent_admission_failure_emits_no_tool_started() { let calls = Arc::new(AtomicUsize::new(0)); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( - "mock", - Arc::new(MockModel::with_responses(vec![tool_calls_response(vec![ - ToolCall::new("c1", "echo_a", json!({})), - ToolCall::new("c2", "echo_b", json!({})), - ToolCall::new("c3", "echo_c", json!({})), - ])])), - ); + harness.with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named( + "mock", + Arc::new(MockModel::with_responses(vec![tool_calls_response(vec![ + ToolCall::new("c1", "echo_a", json!({})), + ToolCall::new("c2", "echo_b", json!({})), + ToolCall::new("c3", "echo_c", json!({})), + ])])), + ), + )); for name in ["echo_a", "echo_b", "echo_c"] { harness.register_tool(Arc::new(EchoTool { name: name.to_string(), @@ -294,13 +298,15 @@ async fn concurrent_admission_failure_emits_no_tool_started() { async fn serial_tool_error_becomes_a_model_visible_result() { let calls = Arc::new(AtomicUsize::new(0)); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( - "mock", - Arc::new(MockModel::with_responses(vec![ - tool_calls_response(vec![ToolCall::new("c1", "flaky", json!({}))]), - text_response("recovered"), - ])), - ); + harness.with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named( + "mock", + Arc::new(MockModel::with_responses(vec![ + tool_calls_response(vec![ToolCall::new("c1", "flaky", json!({}))]), + text_response("recovered"), + ])), + ), + )); harness.register_tool(Arc::new(ErringTool { name: "flaky".into(), recoverable: true, @@ -332,16 +338,18 @@ async fn serial_tool_error_becomes_a_model_visible_result() { async fn concurrent_tool_error_becomes_a_model_visible_result() { let calls = Arc::new(AtomicUsize::new(0)); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( - "mock", - Arc::new(MockModel::with_responses(vec![ - tool_calls_response(vec![ - ToolCall::new("c1", "flaky", json!({})), - ToolCall::new("c2", "echo_a", json!({})), - ]), - text_response("recovered"), - ])), - ); + harness.with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named( + "mock", + Arc::new(MockModel::with_responses(vec![ + tool_calls_response(vec![ + ToolCall::new("c1", "flaky", json!({})), + ToolCall::new("c2", "echo_a", json!({})), + ]), + text_response("recovered"), + ])), + ), + )); harness.register_tool(Arc::new(ErringTool { name: "flaky".into(), recoverable: true, @@ -370,12 +378,14 @@ async fn concurrent_tool_error_becomes_a_model_visible_result() { async fn fatal_tool_error_emits_tool_failed_and_clears_active_calls() { let calls = Arc::new(AtomicUsize::new(0)); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( - "mock", - Arc::new(MockModel::with_responses(vec![tool_calls_response(vec![ - ToolCall::new("c1", "fatal", json!({})), - ])])), - ); + harness.with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named( + "mock", + Arc::new(MockModel::with_responses(vec![tool_calls_response(vec![ + ToolCall::new("c1", "fatal", json!({})), + ])])), + ), + )); harness.register_tool(Arc::new(ErringTool { name: "fatal".into(), recoverable: false, @@ -422,17 +432,19 @@ async fn fatal_tool_error_emits_tool_failed_and_clears_active_calls() { async fn duplicate_call_ids_do_not_clear_each_others_active_entry() { let calls = Arc::new(AtomicUsize::new(0)); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( - "mock", - Arc::new(MockModel::with_responses(vec![ - // A provider that reuses one id across two calls in the same turn. - tool_calls_response(vec![ - ToolCall::new("dup", "echo_a", json!({})), - ToolCall::new("dup", "flaky", json!({})), - ]), - text_response("done"), - ])), - ); + harness.with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named( + "mock", + Arc::new(MockModel::with_responses(vec![ + // A provider that reuses one id across two calls in the same turn. + tool_calls_response(vec![ + ToolCall::new("dup", "echo_a", json!({})), + ToolCall::new("dup", "flaky", json!({})), + ]), + text_response("done"), + ])), + ), + )); harness.register_tool(Arc::new(EchoTool { name: "echo_a".into(), calls: calls.clone(), @@ -477,13 +489,15 @@ async fn duplicate_call_ids_do_not_clear_each_others_active_entry() { #[tokio::test] async fn unknown_tool_recovery_emits_started_and_completed() { let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( - "mock", - Arc::new(MockModel::with_responses(vec![ - tool_calls_response(vec![ToolCall::new("c1", "missing", json!({}))]), - text_response("recovered"), - ])), - ); + harness.with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named( + "mock", + Arc::new(MockModel::with_responses(vec![ + tool_calls_response(vec![ToolCall::new("c1", "missing", json!({}))]), + text_response("recovered"), + ])), + ), + )); harness.with_policy(RunPolicy { unknown_tool: UnknownToolPolicy::ReturnToolError, ..RunPolicy::default() @@ -521,12 +535,14 @@ async fn before_tool_rejection_surfaces_as_a_middleware_error() { // error, is what surfaces while budget remains. let calls = Arc::new(AtomicUsize::new(0)); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( - "mock", - Arc::new(MockModel::with_responses(vec![tool_calls_response(vec![ - ToolCall::new("c1", "echo_a", json!({})), - ])])), - ); + harness.with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named( + "mock", + Arc::new(MockModel::with_responses(vec![tool_calls_response(vec![ + ToolCall::new("c1", "echo_a", json!({})), + ])])), + ), + )); harness.register_tool(Arc::new(EchoTool { name: "echo_a".into(), calls: calls.clone(), @@ -551,18 +567,20 @@ async fn before_tool_rejection_surfaces_as_a_middleware_error() { async fn injected_arguments_are_stripped_before_validation() { let seen = Arc::new(std::sync::Mutex::new(Vec::new())); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( - "mock", - Arc::new(MockModel::with_responses(vec![ - tool_calls_response(vec![ToolCall::new( - "c1", - "injected", - // The model forges the hidden key it was never shown. - json!({ "query": "hi", "thread_id": "forged" }), - )]), - text_response("done"), - ])), - ); + harness.with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named( + "mock", + Arc::new(MockModel::with_responses(vec![ + tool_calls_response(vec![ToolCall::new( + "c1", + "injected", + // The model forges the hidden key it was never shown. + json!({ "query": "hi", "thread_id": "forged" }), + )]), + text_response("done"), + ])), + ), + )); harness.register_tool(Arc::new(InjectedArgTool { seen: seen.clone() })); harness diff --git a/crates/tinyagents-integration-tests/tests/wave3_control_outcomes.rs b/crates/tinyagents-integration-tests/tests/wave3_control_outcomes.rs index bd2682357..8b603d73c 100644 --- a/crates/tinyagents-integration-tests/tests/wave3_control_outcomes.rs +++ b/crates/tinyagents-integration-tests/tests/wave3_control_outcomes.rs @@ -58,7 +58,9 @@ impl Middleware<()> for StopBeforeNCalls { async fn before_model_jump_to_end_stops_the_loop_gracefully() { let mut harness: AgentHarness<()> = AgentHarness::new(); let model = Arc::new(MockModel::with_tool_call("spin", json!({}))); - harness.register_model("mock", model.clone()); + harness.with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named("mock", model.clone()), + )); harness.register_tool(Arc::new(FakeTool::returning("spin", "again"))); harness.push_middleware(Arc::new(StopBeforeNCalls { limit: 1, @@ -117,7 +119,9 @@ async fn jump_to_model_skips_the_turns_tool_calls() { // the test bounds the run with max_model_calls and asserts the tool // itself never actually ran on the skipped turn. let model = Arc::new(MockModel::with_tool_call("spin", json!({}))); - harness.register_model("mock", model.clone()); + harness.with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named("mock", model.clone()), + )); let tool = Arc::new(FakeTool::returning("spin", "ran")); harness.register_tool(tool.clone()); harness.push_middleware(Arc::new(SkipToolsOnce { @@ -178,7 +182,9 @@ impl Middleware<()> for InterruptAfterTool { async fn after_tool_control_interrupt_surfaces_as_interrupted() { let mut harness: AgentHarness<()> = AgentHarness::new(); let model = Arc::new(MockModel::with_tool_call("spin", json!({}))); - harness.register_model("mock", model.clone()); + harness.with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named("mock", model.clone()), + )); harness.register_tool(Arc::new(FakeTool::returning("spin", "again"))); harness.push_middleware(Arc::new(InterruptAfterTool)); @@ -253,7 +259,12 @@ impl Middleware<()> for Winner { async fn first_non_continue_control_wins_and_non_observers_after_it_are_skipped() { let ran = Arc::new(Mutex::new(Vec::new())); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", Arc::new(MockModel::constant("done"))); + harness.with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named( + "mock", + Arc::new(MockModel::constant("done")), + ), + )); harness.push_middleware(Arc::new(Winner { ran: ran.clone() })); harness.push_middleware(Arc::new(Recorder { label: "non-observer", @@ -308,7 +319,9 @@ impl Tool for ReturnDirectTool { async fn tool_return_direct_ends_the_loop_with_the_tools_own_output() { let mut harness: AgentHarness<()> = AgentHarness::new(); let model = Arc::new(MockModel::with_tool_call("finalize", json!({}))); - harness.register_model("mock", model.clone()); + harness.with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named("mock", model.clone()), + )); harness.register_tool(Arc::new(ReturnDirectTool)); let run = harness @@ -353,7 +366,9 @@ impl Tool for GotoModelTool { async fn tool_goto_model_is_honored_via_middleware_control() { let mut harness: AgentHarness<()> = AgentHarness::new(); let model = Arc::new(MockModel::constant("done")); - harness.register_model("mock", model.clone()); + harness.with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named("mock", model.clone()), + )); harness.register_tool(Arc::new(GotoModelTool)); // The model itself never calls `reroute` (MockModel::constant produces @@ -396,7 +411,9 @@ impl Middleware<()> for StopAfterFirstTurn { async fn should_stop_after_turn_ends_the_run_once_a_turn_executed_a_tool() { let mut harness: AgentHarness<()> = AgentHarness::new(); let model = Arc::new(MockModel::with_tool_call("spin", json!({}))); - harness.register_model("mock", model.clone()); + harness.with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named("mock", model.clone()), + )); harness.register_tool(Arc::new(FakeTool::returning("spin", "again"))); harness.push_middleware(Arc::new(StopAfterFirstTurn { seen_turn: Mutex::new(false), diff --git a/crates/tinyagents-integration-tests/tests/wave3_deferred_tools.rs b/crates/tinyagents-integration-tests/tests/wave3_deferred_tools.rs index 8b0803a78..b8a8302e2 100644 --- a/crates/tinyagents-integration-tests/tests/wave3_deferred_tools.rs +++ b/crates/tinyagents-integration-tests/tests/wave3_deferred_tools.rs @@ -76,19 +76,21 @@ fn assistant(tool_calls: Vec<ToolCall>, text: &str) -> ModelResponse { fn harness(delete: Arc<DeleteTool>) -> AgentHarness<()> { let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( - "mock", - Arc::new(MockModel::with_responses(vec![ - assistant( - vec![ - ToolCall::new("call-delete", "delete", json!({"path": "/tmp/x"})), - ToolCall::new("call-click", "browser_click", json!({"x": 1})), - ], - "", - ), - assistant(Vec::new(), "done"), - ])), - ); + harness.with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named( + "mock", + Arc::new(MockModel::with_responses(vec![ + assistant( + vec![ + ToolCall::new("call-delete", "delete", json!({"path": "/tmp/x"})), + ToolCall::new("call-click", "browser_click", json!({"x": 1})), + ], + "", + ), + assistant(Vec::new(), "done"), + ])), + ), + )); harness.register_tool(delete); harness.register_external_tool(ToolSchema { name: "browser_click".into(), @@ -146,13 +148,15 @@ async fn deferred_run_survives_a_json_round_trip_and_resumes_in_a_fresh_harness( }); let mut second_harness = harness(delete2.clone()); // The resumed leg only needs the final answer from the model. - second_harness.register_model( - "mock", - Arc::new(MockModel::with_responses(vec![assistant( - Vec::new(), - "done", - )])), - ); + second_harness.with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named( + "mock", + Arc::new(MockModel::with_responses(vec![assistant( + Vec::new(), + "done", + )])), + ), + )); let run = second_harness .resume_deferred( &(), diff --git a/crates/tinyagents-integration-tests/tests/wave3_output_retry.rs b/crates/tinyagents-integration-tests/tests/wave3_output_retry.rs index 5dbd3597e..1eb1af138 100644 --- a/crates/tinyagents-integration-tests/tests/wave3_output_retry.rs +++ b/crates/tinyagents-integration-tests/tests/wave3_output_retry.rs @@ -74,8 +74,9 @@ async fn validator_rejects_once_then_accepts() { let mut harness: AgentHarness<()> = AgentHarness::new(); harness - .register_model("mock", Arc::new(scripted)) - .set_default_model("mock") + .with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named("mock", Arc::new(scripted)), + )) .with_output_validator(Arc::new(MinScore { threshold: 50 })) .with_policy(RunPolicy { default_response_format: Some(ResponseFormat::json_schema("answer", object_schema())), @@ -142,8 +143,9 @@ async fn exhausting_output_retries_fails_the_run() { let mut harness: AgentHarness<()> = AgentHarness::new(); harness - .register_model("mock", Arc::new(scripted)) - .set_default_model("mock") + .with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named("mock", Arc::new(scripted)), + )) .with_output_validator(validator.clone()) .with_policy(RunPolicy { default_response_format: Some(ResponseFormat::json_schema("answer", object_schema())), @@ -174,8 +176,12 @@ async fn exhausting_output_retries_fails_the_run() { async fn structured_as_deserializes_the_typed_output() { let mut harness: AgentHarness<()> = AgentHarness::new(); harness - .register_model("mock", Arc::new(MockModel::constant(r#"{"score":7}"#))) - .set_default_model("mock") + .with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named( + "mock", + Arc::new(MockModel::constant(r#"{"score":7}"#)), + ), + )) .with_policy(RunPolicy { default_response_format: Some(ResponseFormat::json_schema("answer", object_schema())), ..RunPolicy::default() @@ -193,7 +199,12 @@ async fn structured_as_deserializes_the_typed_output() { #[tokio::test] async fn structured_as_errors_when_the_run_produced_no_structured_output() { let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", Arc::new(MockModel::constant("plain text"))); + harness.with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named( + "mock", + Arc::new(MockModel::constant("plain text")), + ), + )); let run = harness .invoke_default(&(), vec![Message::user("hi")]) diff --git a/crates/tinyagents-integration-tests/tests/wave3_structured_modes.rs b/crates/tinyagents-integration-tests/tests/wave3_structured_modes.rs index 19bd7fb85..2ee46c4c2 100644 --- a/crates/tinyagents-integration-tests/tests/wave3_structured_modes.rs +++ b/crates/tinyagents-integration-tests/tests/wave3_structured_modes.rs @@ -146,7 +146,9 @@ async fn graceful_runs_the_tool_then_finishes_with_the_first_answer() { let mut harness: AgentHarness<()> = AgentHarness::new(); harness - .register_model("mock", model.clone()) + .with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named("mock", model.clone()), + )) .register_tool(search.clone()) .with_policy(RunPolicy { default_response_format: Some(ResponseFormat::auto("result", schema())), @@ -181,7 +183,9 @@ async fn early_finishes_immediately_and_never_runs_the_accompanying_tool() { let mut harness: AgentHarness<()> = AgentHarness::new(); harness - .register_model("mock", model.clone()) + .with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named("mock", model.clone()), + )) .register_tool(search.clone()) .with_policy(RunPolicy { default_response_format: Some(ResponseFormat::auto("result", schema())), @@ -212,7 +216,9 @@ async fn exhaustive_ignores_the_first_output_tool_and_waits_for_a_clean_turn() { let mut harness: AgentHarness<()> = AgentHarness::new(); harness - .register_model("mock", model.clone()) + .with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named("mock", model.clone()), + )) .register_tool(search.clone()) .with_policy(RunPolicy { default_response_format: Some(ResponseFormat::auto("result", schema())), @@ -250,7 +256,9 @@ async fn prompted_mode_injects_the_schema_into_the_system_prompt_and_extracts_fr let mut harness: AgentHarness<()> = AgentHarness::new(); harness - .register_model("mock", model.clone()) + .with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named("mock", model.clone()), + )) .with_policy(RunPolicy { default_response_format: Some(ResponseFormat::auto("result", schema())), structured_strategy_override: Some(StructuredStrategyOverride::Prompted { @@ -291,7 +299,9 @@ async fn tool_call_union_records_which_variant_matched() { let mut harness: AgentHarness<()> = AgentHarness::new(); harness - .register_model("mock", model.clone()) + .with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named("mock", model.clone()), + )) .with_policy(RunPolicy { default_response_format: Some(ResponseFormat::auto( "result", diff --git a/crates/tinyagents-live/Cargo.toml b/crates/tinyagents-live/Cargo.toml index 52428e540..ae1394552 100644 --- a/crates/tinyagents-live/Cargo.toml +++ b/crates/tinyagents-live/Cargo.toml @@ -13,7 +13,7 @@ description = "Runs live (realtime voice) agent sessions with harness tools and tinyliveagents = { path = "../../vendor/tinyliveagents/crates/tinyliveagents", version = "0.1.0" } # Tool calls run through the harness's own admission / middleware / fold # pipeline (`agent_loop::phases::execute_tool_batch`). -tinyagents-harness = { path = "../tinyagents-harness", version = "2.11.0" } +tinyagents-harness = { path = "../tinyagents-harness", version = "3.0.0" } serde_json = { workspace = true } tokio = { workspace = true, features = ["sync", "time"] } tracing = { workspace = true } @@ -23,7 +23,7 @@ async-trait = { workspace = true } anyhow = { workspace = true } tokio = { workspace = true, features = ["sync", "time", "macros", "rt-multi-thread"] } # The example builds a harness with the testkit's FakeTool. -tinyagents-harness = { path = "../tinyagents-harness", version = "2.11.0" } +tinyagents-harness = { path = "../tinyagents-harness", version = "3.0.0" } [lints] workspace = true diff --git a/crates/tinyagents-orchestration/Cargo.toml b/crates/tinyagents-orchestration/Cargo.toml index e7092dd35..1733c8594 100644 --- a/crates/tinyagents-orchestration/Cargo.toml +++ b/crates/tinyagents-orchestration/Cargo.toml @@ -9,8 +9,7 @@ rust-version.workspace = true description = "Host-neutral subagent invocation and lifecycle orchestration." [dependencies] -tinyagents-definition = { path = "../tinyagents-definition", version = "2.11.0" } -anyhow = { workspace = true } +tinyagents-definition = { path = "../tinyagents-definition", version = "3.0.0" } async-trait = { workspace = true } chrono = { workspace = true } serde = { workspace = true } @@ -18,10 +17,9 @@ serde_json = { workspace = true } sha2 = "0.11" thiserror = "2" parking_lot = "0.12" -tinyagents-graph = { path = "../tinyagents-graph", version = "2.11.0", default-features = false } -tinyagents-tasks = { path = "../tinyagents-tasks", version = "2.11.0", default-features = false } -tinyagents-harness = { path = "../tinyagents-harness", version = "2.11.0", default-features = false } -tinyagents-session = { path = "../tinyagents-session", version = "2.11.0", default-features = false } +tinyagents-tasks = { path = "../tinyagents-tasks", version = "3.0.0", default-features = false } +tinyagents-harness = { path = "../tinyagents-harness", version = "3.0.0", default-features = false } +tinyagents-session = { path = "../tinyagents-session", version = "3.0.0", default-features = false } tinyinference-llm = { path = "../../vendor/tinyinference/crates/tinyinference-llm", version = "0.5.0" } tinytools = { path = "../../vendor/tinytools/crates/tinytools", version = "0.6.0" } tokio = { workspace = true, features = ["sync", "rt", "macros", "time"] } @@ -41,12 +39,14 @@ tracing = [ [dev-dependencies] tinystoragedrivers-testkit = { workspace = true } +tinystoragedrivers-connect = { workspace = true } +tinystoragedrivers-document = { workspace = true } dotenvy = "0.15" futures = { workspace = true } tempfile = "3" -tinyagents-definition = { path = "../tinyagents-definition", version = "2.11.0" } -tinyagents-graph = { path = "../tinyagents-graph", version = "2.11.0" } -tinyagents-registry = { path = "../tinyagents-registry", version = "2.11.0" } +tinyagents-definition = { path = "../tinyagents-definition", version = "3.0.0" } +tinyagents-registry = { path = "../tinyagents-registry", version = "3.0.0" } +tinyagents-runtime = { path = "../tinyagents-runtime", version = "3.0.0" } tokio = { workspace = true, default-features = true, features = ["macros", "rt-multi-thread", "time", "test-util"] } [lints] diff --git a/crates/tinyagents-orchestration/examples/orchestrator_subagents.rs b/crates/tinyagents-orchestration/examples/orchestrator_subagents.rs deleted file mode 100644 index 713bf66fa..000000000 --- a/crates/tinyagents-orchestration/examples/orchestrator_subagents.rs +++ /dev/null @@ -1,299 +0,0 @@ -//! Flagship registry showcase: an orchestrator that *designs* which sub-agents -//! to call at runtime by resolving them **by name** from a -//! [`CapabilityRegistry`], then runs them and composes their results. -//! -//! The pattern this example demonstrates is the whole point of the named -//! capability registry: -//! -//! 1. **Register named capabilities.** Several specialized sub-agents -//! (`researcher`, `coder`, `summarizer`), each a [`SubAgent`] over an -//! [`OpenAiModel`] with a distinct system prompt, are wrapped as -//! [`SubAgentTool`]s and registered in a [`CapabilityRegistry`] under their -//! names. -//! 2. **Discover.** The set of available capabilities (names + descriptions) is -//! read back out of the registry — nothing is hard-coded into the planner. -//! 3. **Design.** An *orchestrator* agent is given the task plus the discovered -//! menu and decides, via structured output, **which** registered sub-agents -//! to invoke. It returns a list of names. -//! 4. **Bind at runtime.** Each chosen name is resolved from the registry with -//! [`CapabilityRegistry::tool`] and spawned in parallel. Each spawn returns -//! a job id; the registry is queried until the results can be composed. -//! -//! Capabilities are therefore *named, discovered, and bound at runtime* — the -//! orchestrator never holds a direct handle to any sub-agent; it only knows -//! their names and looks them up in the registry when it decides to use them. -//! -//! Run with: -//! -//! ```text -//! cargo run -p tinyagents-integration-tests --example orchestrator_subagents -//! ``` - -use std::collections::HashMap; -use std::sync::Arc; - -use futures::future::join_all; -use serde_json::{Value, json}; - -use tinyagents_graph::*; -use tinyagents_harness::context::{RunConfig, RunContext}; -use tinyagents_harness::middleware::AgentRun; -use tinyagents_harness::runtime::{AgentHarness, RunPolicy}; -use tinyagents_harness::tool::ToolDispatch; -use tinyagents_orchestration::subagent::{ - ChildDataPolicy, SubAgent, SubAgentJobRegistry, SubAgentTool, -}; -use tinyagents_registry::*; -use tinyinference_llm::message::Message; -use tinyinference_llm::model::{ChatModel, ResponseFormat}; -use tinyinference_llm::providers::openai::OpenAiModel; - -/// A specialized sub-agent's static identity. -struct AgentSpec { - name: &'static str, - description: &'static str, - system_prompt: &'static str, -} - -/// The three specialist sub-agents this orchestrator can choose from. -const SPECIALISTS: &[AgentSpec] = &[ - AgentSpec { - name: "researcher", - description: "Gathers and explains factual background on a topic.", - system_prompt: "You are a meticulous researcher. Answer with a few concise, factual \ - bullet points of background relevant to the user's task. No code.", - }, - AgentSpec { - name: "coder", - description: "Writes small, focused code snippets and explains them.", - system_prompt: "You are a senior Rust engineer. When asked, produce a short, correct \ - code snippet with a one-line explanation. Keep it minimal.", - }, - AgentSpec { - name: "summarizer", - description: "Condenses material into a short, plain-language summary.", - system_prompt: "You are an editor. Produce a crisp, plain-language summary in 2-3 \ - sentences. No jargon.", - }, -]; - -/// Builds a [`SubAgentTool`] wrapping a [`SubAgent`] over the shared OpenAI -/// model with the spec's distinct system prompt. -fn build_specialist( - spec: &AgentSpec, - model: Arc<dyn ChatModel<()>>, - jobs: SubAgentJobRegistry, -) -> SubAgentTool<()> { - let mut harness: AgentHarness<()> = AgentHarness::new(); - harness - .register_model("model", model) - .set_default_model("model"); - let subagent = SubAgent::new(spec.name, spec.description, Arc::new(harness)) - .with_system_prompt(spec.system_prompt); - SubAgentTool::new( - Arc::new(subagent), - ChildDataPolicy::new(|parent: &()| *parent), - ) - .with_job_registry(jobs) -} - -/// Reads the `{ "agents": [..] }` selection out of an [`AgentRun`], preferring -/// the extracted structured output and falling back to parsing the raw text. -fn parse_selection(run: &AgentRun) -> Vec<String> { - let value: Value = run - .structured - .clone() - .or_else(|| run.text().and_then(|t| serde_json::from_str(&t).ok())) - .unwrap_or(Value::Null); - value - .get("agents") - .and_then(Value::as_array) - .map(|arr| { - arr.iter() - .filter_map(|v| v.as_str().map(str::to_owned)) - .collect() - }) - .unwrap_or_default() -} - -#[tokio::main] -async fn main() -> Result<()> { - dotenvy::dotenv().ok(); - - let model: Arc<dyn ChatModel<()>> = Arc::new(OpenAiModel::from_env()?); - println!("=== Orchestrator designs sub-agents via the registry ==="); - - // 1. Register every specialist sub-agent by name in the capability registry. - let mut registry: CapabilityRegistry<()> = CapabilityRegistry::new(); - let mut dispatches: HashMap<String, Arc<SubAgentTool<()>>> = HashMap::new(); - let jobs = SubAgentJobRegistry::new(); - for spec in SPECIALISTS { - let dispatch = Arc::new(build_specialist(spec, model.clone(), jobs.clone())); - registry.register_tool(dispatch.tool())?; - dispatches.insert(spec.name.to_owned(), dispatch); - } - - // 2. Discover the available capabilities straight out of the registry. - let menu: Vec<(String, String)> = registry - .names(ComponentKind::Tool) - .into_iter() - .map(|name| { - let desc = registry - .tool(&name) - .map(|t| t.description().to_owned()) - .unwrap_or_default(); - (name, desc) - }) - .collect(); - let menu_text = menu - .iter() - .map(|(name, desc)| format!("- {name}: {desc}")) - .collect::<Vec<_>>() - .join("\n"); - let available: Vec<String> = menu.iter().map(|(n, _)| n.clone()).collect(); - println!("registered capabilities:\n{menu_text}\n"); - - // 3. Design: the orchestrator picks which named sub-agents to use, returning - // a structured `{ "agents": [..] }` constrained to the registered names. - let task = "Explain what a Rust trait object is and give a tiny code example, \ - then summarize it for a beginner."; - println!("task: {task}\n"); - - let selection_schema = ResponseFormat::json_schema( - "agent_selection", - json!({ - "type": "object", - "properties": { - "agents": { - "type": "array", - "items": { "type": "string", "enum": available }, - "description": "Names of the sub-agents to invoke for this task." - } - }, - "required": ["agents"], - "additionalProperties": false - }), - ); - - let mut orchestrator: AgentHarness<()> = AgentHarness::new(); - orchestrator - .register_model("model", model.clone()) - .set_default_model("model") - .with_policy(RunPolicy { - default_response_format: Some(selection_schema), - ..RunPolicy::default() - }); - - let plan = orchestrator - .invoke_default( - &(), - vec![ - Message::system(format!( - "You are an orchestrator. You have these named sub-agents available:\n\ - {menu_text}\n\n\ - Choose the minimal subset whose combined skills solve the user's task. \ - Respond ONLY with the requested JSON object listing the sub-agent names." - )), - Message::user(task), - ], - ) - .await?; - - let mut chosen = parse_selection(&plan); - chosen.retain(|name| registry.has(ComponentKind::Tool, name)); - if chosen.is_empty() { - // Defensive fallback so the showcase always does *something* useful. - chosen = available.clone(); - } - println!("orchestrator chose: {chosen:?}\n"); - - // 4. Bind at runtime: resolve each chosen name and spawn its sub-agent job. - let runs = chosen.iter().enumerate().map(|(i, name)| { - let dispatch = dispatches - .get(name) - .cloned() - .expect("chosen name resolves in the registry"); - let name = name.clone(); - async move { - let parent = RunContext::new(RunConfig::new(format!("orchestrator-{i}")), ()); - let result = dispatch - .invoke_in_parent_context( - &(), - json!({ "input": task }), - tinytools::ToolCallOptions::default(), - &parent, - ) - .await?; - let job_id = serde_json::from_str::<Value>(&result.output())? - .get("job_id") - .and_then(Value::as_str) - .expect("spawn result contains a job id") - .to_owned(); - Ok::<(String, String), tinyagents_harness::TinyAgentsError>((name, job_id)) - } - }); - let spawned: Vec<(String, String)> = join_all(runs).await.into_iter().collect::<Result<_>>()?; - let outputs = tokio::time::timeout(std::time::Duration::from_secs(120), async { - loop { - let completed = spawned - .iter() - .filter_map(|(name, job_id)| { - let job = jobs.get(job_id)?; - job.status.is_terminal().then(|| { - ( - name.clone(), - job.output - .or(job.error) - .unwrap_or_else(|| "no output".into()), - ) - }) - }) - .collect::<Vec<_>>(); - if completed.len() == spawned.len() { - return completed; - } - tokio::time::sleep(std::time::Duration::from_millis(25)).await; - } - }) - .await - .map_err(|_| { - tinyagents_harness::TinyAgentsError::Validation( - "subagent jobs did not finish within two minutes".into(), - ) - })?; - - for (name, text) in &outputs { - println!("── {name} ──\n{text}\n"); - } - - // 5. Compose: synthesize the sub-agent outputs into one final answer. - let composed_context = outputs - .iter() - .map(|(name, text)| format!("[{name}]\n{text}")) - .collect::<Vec<_>>() - .join("\n\n"); - - let mut composer: AgentHarness<()> = AgentHarness::new(); - composer - .register_model("model", model) - .set_default_model("model"); - let composed = composer - .invoke_default( - &(), - vec![ - Message::system( - "Combine the labeled sub-agent outputs below into one coherent answer to \ - the user's task. Preserve any code. Do not mention the sub-agents.", - ), - Message::user(format!( - "Task: {task}\n\nSub-agent outputs:\n{composed_context}" - )), - ], - ) - .await?; - - println!("=== composed answer ==="); - println!("{}", composed.text().unwrap_or_default()); - - Ok(()) -} diff --git a/crates/tinyagents-orchestration/src/README.md b/crates/tinyagents-orchestration/src/README.md index 81bbf3d36..cad78af68 100644 --- a/crates/tinyagents-orchestration/src/README.md +++ b/crates/tinyagents-orchestration/src/README.md @@ -1,127 +1,41 @@ -# `tinyagents-orchestration` — subagent orchestration +# `tinyagents-orchestration` — subagent lifecycle -This crate owns subagent orchestration policy and job lifecycle. Generic tool -shells for its registry live in `tinyagents-harness-tools`, which depends on -this crate only under its `subagents` feature. +This crate owns the host-neutral policy and lifecycle for asynchronous +subagents. Hosts provide planning, execution, persistence, authorization, and +product effects. ## Public surface -- `run_progress_ledger::RunProgressLedger` records one foreground turn's - progress events and terminal status through the host-installed - `RunLedgerStore`; `ProgressLedgerContext` carries parent-call correlation to - the host's projection. Storage failures do not interrupt the live turn, and - terminal transitions are retried during `finish()` before interruption - fallback. See the [module guide](run_progress_ledger/README.md) for its - design and operational constraints. - -- `subagent::SubAgent` runs an `AgentHarness` as a child with inherited - lineage, cancellation, events, host authority, and recursion limits. -- `subagent::SubAgentTool` starts a child asynchronously through typed - parent-context tool dispatch and immediately returns a stable job id. Hosts - Hosts that use the prepared lifecycle construct - `tinyagents_harness_tools::subagents::DriverSubAgentTool` directly. -- `subagent::SubAgentJobRegistry` records queued, running, completed, failed, - and cancelled jobs on top of `tinyagents_tasks::DetachedTaskRegistry`. - `tinyagents_harness_tools::subagents::SubAgentJobsTool` queries them and - `SubAgentMessageTool` sends messages to live children. -- `tinyagents_harness_tools::subagents::SubAgentJobsTool::with_sessions` overlays scoped durable sessions - on list/query/close/await; `with_task_store` adds parent-run task records. - `tinyagents_harness_tools::subagents::SubAgentMessageTool::with_sessions` accepts durable session ids for live - steering. The host supplies parent-scope mapping closures so durable stores - do not decide authorization. -- `subagent::MultiTargetDelegationTool` exposes a configured target enum and - structured task envelope. Its `blocking` field maps to the standard `mode` - field, and exposure can be set to `Hidden` for host-only tools. -- `subagent::DelegationTool` runs the graph plan/execute/review cycle through - an injected stage worker. Its default keeps the one-shot execution and - standalone `APPROVE` parser; `with_resumable_execution` opts checkpointed - threads into classify-and-resume behavior. Hosts can retain their worker's - approval decision with `ReviewApprovalPolicy::WorkerOutput`, project the - complete state, pending approval, or graph error into a `ToolResult` through - `with_result_projection`, and preserve declaration metadata with - `with_metadata`. -- `parallel::ParallelFanoutTool` composes host task parsing, admission policy, - asynchronous preflight, branch-context creation, and execution with ordered - results and explicit rejection records. -- `subagent::DriverSubagentSessionStore` persists reusable child-session - identity and history on the host document-store port. Hosts can continue by - durable id, close sessions, and prevent stale task completions from replacing - a newer assignment. -- `subagent::SubAgentSession` reuses one child and its transcript across turns. -- `subagent::SpawnPolicy` / `SpawnAdmission` bound child fan-out (per-parent - live cap, per-root total budget, target allowlist) with atomic reservation; - see `subagent/README.md`. - `subagent::SubagentDriver` coordinates durable resume, preparation, - execution, pause, and terminal persistence through host-supplied traits. - It is the single engine for an agent step: team member steps and workflow - agent children run through it too (`subagent::run_agent_step`), so they get - the same spawn admission, policy, result policy and typed outcome. -- `steering::{steer_detached, steer_detached_with_request_id}` sends lane-checked - steering to a live child handle, or delegates it to the host's - `SteerFallback` when the handle is absent or the retained task has finished. - `SteeringRegistry` remains in `tinyagents-tasks`, where - `DetachedTaskRegistry` receives it at construction. `steering::cancel_for_thread` - cancels registered children belonging to a parent thread. - -The direct invocation implementation and its tool-focused tests live in -`subagent/invocation/`. Durable lifecycle files live beside it in `subagent/`. -End-to-end and live subagent tests live in the crate-level `tests/` directory, -and `tests/live_orchestrator_subagents.rs` exercises network-backed job-based -delegation. + execution, pauses, retries, and terminal persistence through host-supplied + traits. +- `subagent::SubAgentJobRegistry` records spawned jobs and supports status, + cancellation, steering, messages, and descendant settlement. The + `DriverSubAgentTool` in `tinyagents-harness-tools` starts jobs and returns a + task reference immediately; registry wait and detached wait APIs are + separate host-facing operations. +- Detached lifecycle helpers in `subagent::detached` support completion + routing and settlement, including `wait_all_detached_with_owners` for + headless hosts. +- `subagent::SpawnPolicy` / `SpawnAdmission` bound child fan-out, and + `SubAgentPolicy`, `SubAgentBudget`, and `SubAgentOutput` define execution + policy and accounting. +- `DriverSubagentSessionStore` persists reusable child-session identity and + history through a host document-store port. +- `steering` provides lane-checked control of live child runs. ## Status vocabularies -`OrchestrationTaskStatus` (crate `tinyagents-tasks`) is the one durable -lifecycle status for a subagent or task run. The other enums stay because their -serialized forms are persisted or consumed elsewhere, but each is expressed in -terms of it through `From` / `TryFrom` impls (`NoEquivalentStatus` is the error -of the fallible ones). The single mapping table, with the lossy cases, is in the -`tinyagents-tasks` README. `DetachedSubagentStatus` stays the live type that -carries a payload (output, question, error). - -| Type | Crate | Role | -| --- | --- | --- | -| `OrchestrationTaskStatus` | tasks | Canonical durable status. | -| `DetachedSubagentStatus` | orchestration (`subagent::detached`) | Live status published to waiters, with payload. Not serialized. | -| `SubAgentJobStatus` | orchestration (`subagent::invocation`) | Status on `SubAgentJob` snapshots. | -| `SubagentOutcomeKind` | orchestration (`subagent`) | What a driver-level run returned, with its pause/incomplete payload. | -| `AgentRunStatus` | session (`run_ledger`) | Durable ledger row status. | -| `TranscriptSubagentStatus` | session (`transcript::view`) | Display projection read back from a transcript. | -| `CompletionStatus` | tasks (`completions`) | How a finished child ended, for the completion router. | - -The old `SubagentStatus` names remain as `#[deprecated]` type aliases at their -original paths. `status::task_status_to_run_status` and -`status::run_status_to_task_status` are `#[deprecated]` in favour of -`AgentRunStatus::from` / `OrchestrationTaskStatus::from`. No wire format -changed; golden serde tests pin each one. - -Dependency direction decides where an impl lives: ledger and transcript-view -conversions are in `tinyagents-session`, `CompletionStatus` in -`tinyagents-tasks`, and the job, detached and outcome ones beside their types -in this crate. `to_task_status()` / `to_run_status()` on the job and detached -types are kept as thin wrappers over the `From` impls. +`OrchestrationTaskStatus` in `tinyagents-tasks` is the canonical durable task +status. `DetachedSubagentStatus`, `SubAgentJobStatus`, `SubagentOutcomeKind`, +`AgentRunStatus`, transcript status, and completion status retain their +role-specific payloads and serialized forms, with conversions documented in +`tinyagents-tasks`. ## Boundaries -The crate owns subagent invocation and lifecycle (`subagent/`), durable team -coordination (`teams/`: `TeamService`, `TeamLedger`, `run_member_graph`) and -durable phase-DAG workflows (`workflow/`: `WorkflowEngine`, `WorkflowStore`). -Generic graph execution remains in `tinyagents-graph`, while provider calls, -tool dispatch mechanics, events, and run contexts remain in -`tinyagents-harness`. This crate composes those primitives into child-agent -behavior. - -`subagent::SubagentDriver` is the single agent-step engine (#349). Team member -steps (`teams`) and workflow child steps (`workflow`) do not hand-roll their -own lifecycle; both run through `subagent::run_agent_step`, which adapts a -host-supplied worker onto the driver, so every agent step gets the same spawn -admission, policy, result policy and typed outcome. Workflow steps are the one -exception for structured output: `workflow::run_child_step` clears -`result_policy.schema` before calling `run_agent_step` and validates the -structured value itself (`workflow/child_step.rs`). See `teams/README.md` and -`workflow/README.md`. - -Hosts remain responsible for agent definitions, model selection, credentials, -workspace policy, durable storage implementations, and authorization. Hosted -child invocations always re-enter through the exact capability bundle carried -by their parent context. +This crate owns async subagent lifecycle only. Generic graph execution stays +in `tinyagents-orchestration`; tool dispatch, events, and run contexts stay in +`tinyagents-harness`. Hosts remain responsible for agent definitions, model +selection, credentials, workspace policy, storage implementations, and +authorization. diff --git a/crates/tinyagents-orchestration/src/lib.rs b/crates/tinyagents-orchestration/src/lib.rs index 32a7900e2..ae0cc682f 100644 --- a/crates/tinyagents-orchestration/src/lib.rs +++ b/crates/tinyagents-orchestration/src/lib.rs @@ -1,30 +1,21 @@ //! Host-neutral subagent invocation and lifecycle orchestration. //! -//! This crate owns TinyAgents' subagent-facing surfaces: direct child-agent -//! invocation, the typed tool adapter, reusable child sessions, and the -//! durable lifecycle driver. Hosts supply persistence and concrete execution; -//! policy, credentials, progress, and RPC remain host concerns. +//! This crate owns TinyAgents' host-neutral async subagent lifecycle driver, +//! detached jobs, and completion handling. Hosts supply persistence and +//! concrete execution; policy, credentials, progress, and RPC remain host +//! concerns. //! //! `OrchestrationTaskStatus` is the canonical status; the job, detached, outcome, //! run-ledger and completion statuses convert to and from it with `From` / -//! `TryFrom` (see the [`status`] module docs). +//! `TryFrom` (see the `tinyagents-tasks` status mapping docs). //! //! Dependency direction is deliberately one way: -//! `orchestration -> {harness, runtime}`. The lower-level crates never -//! depend on this composition layer. +//! `orchestration -> {harness, session, tasks, storage ports}`. Lower-level +//! crates never depend on this composition layer. -pub mod parallel; -pub mod run_progress_ledger; -pub mod status; pub mod steering; pub mod subagent; -pub mod teams; -pub mod topology; -pub mod workflow; #[cfg(test)] #[path = "lib_boundary_tests.rs"] mod boundary_tests; - -#[cfg(test)] -mod test_support; diff --git a/crates/tinyagents-orchestration/src/parallel/collect.rs b/crates/tinyagents-orchestration/src/parallel/collect.rs deleted file mode 100644 index bc8f1efef..000000000 --- a/crates/tinyagents-orchestration/src/parallel/collect.rs +++ /dev/null @@ -1,70 +0,0 @@ -//! Aggregating branch results and detecting cross-branch overlap. - -use std::path::PathBuf; - -use serde::Serialize; - -/// What the collector needs to know about one finished branch. -pub trait BranchReport { - fn branch_id(&self) -> &str; - fn succeeded(&self) -> bool; - /// Files (relative to the branch root) the branch changed. - fn changed_files(&self) -> &[String]; -} - -/// A file touched by more than one branch. -#[derive(Debug, Clone, PartialEq, Eq, Serialize)] -pub struct OverlapWarning { - pub file: String, - pub workers: Vec<String>, -} - -/// Aggregate of a finished batch, results in input order. -#[derive(Debug, Clone)] -pub struct Collected<R> { - pub results: Vec<R>, - pub failures: usize, - pub overlap_warnings: Vec<OverlapWarning>, -} - -impl<R: BranchReport> Collected<R> { - pub fn new(results: Vec<R>) -> Self { - let failures = results.iter().filter(|r| !r.succeeded()).count(); - let overlap_warnings = overlap_warnings(&results); - Self { - results, - failures, - overlap_warnings, - } - } -} - -impl<R> Collected<R> { - pub fn total(&self) -> usize { - self.results.len() - } - pub fn succeeded(&self) -> usize { - self.results.len().saturating_sub(self.failures) - } -} - -/// Files changed by more than one branch. -pub fn overlap_warnings<R: BranchReport>(results: &[R]) -> Vec<OverlapWarning> { - let per_worker: Vec<(String, Vec<PathBuf>)> = results - .iter() - .filter(|r| !r.changed_files().is_empty()) - .map(|r| { - ( - r.branch_id().to_string(), - r.changed_files().iter().map(PathBuf::from).collect(), - ) - }) - .collect(); - tinyagents_harness::workspace::detect_worktree_overlaps(&per_worker) - .into_iter() - .map(|(file, workers)| OverlapWarning { - file: file.to_string_lossy().into_owned(), - workers, - }) - .collect() -} diff --git a/crates/tinyagents-orchestration/src/parallel/hooks.rs b/crates/tinyagents-orchestration/src/parallel/hooks.rs deleted file mode 100644 index c9b9d1f64..000000000 --- a/crates/tinyagents-orchestration/src/parallel/hooks.rs +++ /dev/null @@ -1,87 +0,0 @@ -//! Host seams for a parallel run: progress, approval and configured -//! concurrency. All defaults are no-ops so a host opts in to only what it needs. - -use std::future::Future; - -use tinyagents_harness::{CancellationToken, Result}; - -use super::staging::Worker; -use super::workers::run; - -/// Callbacks a host attaches to a parallel run. Methods are synchronous and -/// must not block; they are invoked on the branch's task. -pub trait RunHooks<P, R>: Send + Sync { - /// Approval gate, consulted just before a worker starts. `Err(reason)` - /// skips execution and the worker's result is built by the `denied` - /// closure given to [`run_with_hooks`]. - fn approve(&self, _index: usize, _worker: &Worker<P>) -> std::result::Result<(), String> { - Ok(()) - } - /// Progress: a worker is about to execute (after approval). - fn on_start(&self, _index: usize, _worker: &Worker<P>) {} - /// Progress: a worker finished (including a denied one). - fn on_finish(&self, _index: usize, _result: &R) {} -} - -/// Clamp a host-configured concurrency to the batch size. `configured == 0` -/// means "one slot per worker"; a non-zero value is never larger than `n`. -pub fn resolve_concurrency(configured: usize, n: usize) -> usize { - let n = n.max(1); - if configured == 0 { - n - } else { - configured.min(n) - } -} - -/// [`run`] with host hooks. Results stay in input order. A worker the hook -/// denies is not executed; `denied(worker, reason)` supplies its result. -pub async fn run_with_hooks<P, C, R, B, F, Fut, D>( - workers: Vec<Worker<P>>, - configured_concurrency: usize, - cancel: CancellationToken, - hooks: &dyn RunHooks<P, R>, - branch_ctx: B, - exec: F, - denied: D, -) -> Result<Vec<R>> -where - B: Fn(usize) -> C, - F: Fn(Worker<P>, C) -> Fut, - Fut: Future<Output = R>, - D: Fn(Worker<P>, String) -> R, -{ - let concurrency = resolve_concurrency(configured_concurrency, workers.len()); - // The branch context is built only after approval, so a denied worker has - // no setup side effects. Hooks run on - // the branch's own task, so `on_finish` fires as each worker completes. - run( - workers, - concurrency, - cancel, - |i| i, - |w, idx| { - // Exactly one of `ready` (denied) or `fut` (approved) is set. - let (ready, fut) = match hooks.approve(idx, &w) { - Ok(()) => { - hooks.on_start(idx, &w); - (None, Some(exec(w, branch_ctx(idx)))) - } - Err(reason) => { - tracing::debug!(index = idx, "[parallel] worker denied by approval hook"); - (Some(denied(w, reason)), None) - } - }; - async move { - let r = match (ready, fut) { - (Some(r), _) => r, - (None, Some(fut)) => fut.await, - (None, None) => unreachable!("approved or denied"), - }; - hooks.on_finish(idx, &r); - r - } - }, - ) - .await -} diff --git a/crates/tinyagents-orchestration/src/parallel/hooks_tests.rs b/crates/tinyagents-orchestration/src/parallel/hooks_tests.rs deleted file mode 100644 index b79736cc4..000000000 --- a/crates/tinyagents-orchestration/src/parallel/hooks_tests.rs +++ /dev/null @@ -1,97 +0,0 @@ -use super::*; -use std::sync::Mutex; -use tinyagents_harness::CancellationToken; - -#[derive(Clone, Debug)] -struct P; - -struct Policy; -impl AdmissionPolicy<P> for Policy { - fn admit(&self, _i: usize, _t: &Task<P>) -> AdmissionVerdict { - AdmissionVerdict::ReadOnly - } -} - -fn workers(ids: &[&str]) -> Vec<Worker<P>> { - Plan::new( - ids.iter() - .map(|id| Task { - id: (*id).into(), - prompt: "p".into(), - ownership: None, - payload: P, - }) - .collect(), - ) - .stage(&Policy) - .into_iter() - .map(|p| match p { - Preflight::Prepared(w) => w, - Preflight::Rejected(_) => panic!("rejected"), - }) - .collect() -} - -#[derive(Default)] -struct Rec { - deny: Option<&'static str>, - started: Mutex<Vec<usize>>, - finished: Mutex<Vec<String>>, -} -impl RunHooks<P, String> for Rec { - fn approve(&self, _i: usize, w: &Worker<P>) -> Result<(), String> { - match self.deny { - Some(id) if id == w.id => Err("no".into()), - _ => Ok(()), - } - } - fn on_start(&self, i: usize, _w: &Worker<P>) { - self.started.lock().unwrap().push(i); - } - fn on_finish(&self, _i: usize, r: &String) { - self.finished.lock().unwrap().push(r.clone()); - } -} - -#[test] -fn concurrency_is_clamped_to_batch_size() { - assert_eq!(resolve_concurrency(0, 5), 5); - assert_eq!(resolve_concurrency(8, 3), 3); - assert_eq!(resolve_concurrency(2, 5), 2); - assert_eq!(resolve_concurrency(0, 0), 1); -} - -#[tokio::test(flavor = "multi_thread")] -async fn denied_worker_is_skipped_but_keeps_its_slot_and_progress_fires() { - let hooks = Rec { - deny: Some("b"), - ..Default::default() - }; - let built: Mutex<Vec<usize>> = Mutex::default(); - let out = run_with_hooks( - workers(&["a", "b", "c"]), - 2, - CancellationToken::new(), - &hooks, - |i| { - built.lock().unwrap().push(i); - }, - |w, _| async move { format!("ran:{}", w.id) }, - |w, reason| format!("denied:{}:{reason}", w.id), - ) - .await - .unwrap(); - assert_eq!(out, vec!["ran:a", "denied:b:no", "ran:c"]); - // A denied worker never gets a branch context built for it. - let mut built = built.lock().unwrap().clone(); - built.sort(); - assert_eq!(built, vec![0, 2]); - let mut started = hooks.started.lock().unwrap().clone(); - started.sort(); - assert_eq!(started, vec![0, 2]); - let mut finished = hooks.finished.lock().unwrap().clone(); - finished.sort(); - let mut expected = out.clone(); - expected.sort(); - assert_eq!(finished, expected); -} diff --git a/crates/tinyagents-orchestration/src/parallel/mod.rs b/crates/tinyagents-orchestration/src/parallel/mod.rs deleted file mode 100644 index 6ef63b9fc..000000000 --- a/crates/tinyagents-orchestration/src/parallel/mod.rs +++ /dev/null @@ -1,44 +0,0 @@ -//! Parallel fan-out of independent branches over [`tinyagents_graph::parallel`]. -//! -//! A host builds a [`Plan`] of [`Task`]s, supplies an [`AdmissionPolicy`] that -//! decides (in its own vocabulary) whether each task may run and what it needs -//! from the shared workspace, and runs the staged [`Worker`]s with [`run`]. -//! -//! - [`staging`]: admission, shared-workspace arbitration and the ownership -//! boundary prompt. Conflicts are reported as data ([`Rejection`]) so the -//! rejection sentence stays a host decision. -//! - [`workers`]: serial fallback for shared-workspace writers, bounded -//! `map_reduce` fan-out otherwise. Results are returned in **input order** -//! and every branch receives its own context from a factory (per-branch -//! isolation). -//! - [`collect`]: aggregate results and cross-branch overlap detection. -//! -//! Nothing here performs I/O beyond what the host's branch executor does; -//! policy, progress and approvals remain host concerns. - -pub mod collect; -pub mod hooks; -pub mod staging; -pub mod tool; -pub mod workers; - -pub use collect::{BranchReport, Collected, OverlapWarning, overlap_warnings}; -pub use hooks::{RunHooks, resolve_concurrency, run_with_hooks}; -pub use staging::{ - AdmissionPolicy, AdmissionVerdict, Plan, Preflight, Rejection, RejectionKind, Task, Worker, - WorkerDispatchMode, with_ownership_boundary, -}; -pub use tool::ParallelFanoutTool; -pub use workers::run; - -#[cfg(test)] -#[path = "mod_tests.rs"] -mod tests; - -#[cfg(test)] -#[path = "hooks_tests.rs"] -mod hooks_tests; - -#[cfg(test)] -#[path = "tool_tests.rs"] -mod tool_tests; diff --git a/crates/tinyagents-orchestration/src/parallel/mod_tests.rs b/crates/tinyagents-orchestration/src/parallel/mod_tests.rs deleted file mode 100644 index 3b224ceb3..000000000 --- a/crates/tinyagents-orchestration/src/parallel/mod_tests.rs +++ /dev/null @@ -1,262 +0,0 @@ -use super::*; -use std::sync::Arc; -use std::sync::atomic::{AtomicUsize, Ordering}; -use std::time::Duration; -use tinyagents_harness::CancellationToken; - -#[derive(Clone, Debug)] -struct P { - kind: &'static str, - paths: Vec<&'static str>, -} - -struct Policy; -impl AdmissionPolicy<P> for Policy { - fn admit(&self, _i: usize, t: &Task<P>) -> AdmissionVerdict { - match t.payload.kind { - "read" => AdmissionVerdict::ReadOnly, - "iso" => AdmissionVerdict::Isolated, - "write" => AdmissionVerdict::Writes(t.payload.paths.iter().map(Into::into).collect()), - _ => AdmissionVerdict::Reject("nope".into()), - } - } -} - -fn task(id: &str, kind: &'static str, paths: Vec<&'static str>) -> Task<P> { - Task { - id: id.into(), - prompt: format!("do {id}"), - ownership: None, - payload: P { kind, paths }, - } -} - -fn prepared(p: Vec<Preflight<P>>) -> Vec<Worker<P>> { - p.into_iter() - .map(|x| match x { - Preflight::Prepared(w) => w, - Preflight::Rejected(r) => panic!("rejected {:?}", r.kind), - }) - .collect() -} - -#[test] -fn policy_rejection_keeps_input_position() { - let out = Plan::new(vec![ - task("a", "read", vec![]), - task("b", "bad", vec![]), - task("c", "iso", vec![]), - ]) - .stage(&Policy); - assert!(matches!(out[0], Preflight::Prepared(_))); - assert!( - matches!(&out[1], Preflight::Rejected(r) if r.kind == RejectionKind::Policy("nope".into())) - ); - assert!(matches!(out[2], Preflight::Prepared(_))); -} - -#[test] -fn disjoint_writers_serialize_and_overlapping_writers_conflict() { - let out = Plan::new(vec![ - task("a", "write", vec!["x.rs"]), - task("b", "write", vec!["y.rs"]), - ]) - .stage(&Policy); - assert!(out.iter().all(|p| matches!(p, Preflight::Prepared(w) if w.dispatch_mode == WorkerDispatchMode::SerialSharedWorkspaceWrite))); - - let out = Plan::new(vec![ - task("a", "write", vec!["x.rs"]), - task("b", "write", vec!["x.rs"]), - ]) - .stage(&Policy); - assert!(matches!(out[0], Preflight::Prepared(_))); - assert!( - matches!(&out[1], Preflight::Rejected(r) if matches!(r.kind, RejectionKind::Conflict(_))) - ); -} - -#[test] -fn ownership_boundary_is_folded_into_prompt() { - let mut t = task("a", "read", vec![]); - t.ownership = Some("files: a.rs".into()); - let w = prepared(Plan::new(vec![t]).stage(&Policy)); - assert!(w[0].prompt.starts_with("[Ownership Boundary]\nfiles: a.rs")); - assert_eq!(with_ownership_boundary("x", Some(" ")), "x"); -} - -#[tokio::test(flavor = "multi_thread")] -async fn results_are_in_input_order_with_interleaved_completion() { - let ws = prepared( - Plan::new(vec![ - task("a", "read", vec![]), - task("b", "read", vec![]), - task("c", "read", vec![]), - ]) - .stage(&Policy), - ); - // Earlier workers finish later. - let out = run( - ws, - 0, - CancellationToken::new(), - |i| i, - |w, i| async move { - tokio::time::sleep(Duration::from_millis(60 - 25 * i as u64)).await; - (w.id, i) - }, - ) - .await - .unwrap(); - assert_eq!( - out, - vec![ - ("a".to_string(), 0), - ("b".to_string(), 1), - ("c".to_string(), 2) - ] - ); -} - -#[tokio::test(flavor = "multi_thread")] -async fn each_branch_gets_its_own_context() { - let ws = prepared( - Plan::new(vec![task("a", "iso", vec![]), task("b", "iso", vec![])]).stage(&Policy), - ); - let made = Arc::new(AtomicUsize::new(0)); - let m = made.clone(); - let out = run( - ws, - 2, - CancellationToken::new(), - move |_| Arc::new(AtomicUsize::new(m.fetch_add(1, Ordering::SeqCst))), - |_w, ctx| async move { ctx.fetch_add(100, Ordering::SeqCst) }, - ) - .await - .unwrap(); - assert_eq!(made.load(Ordering::SeqCst), 2); - assert_eq!(out, vec![0, 1]); -} - -#[tokio::test] -async fn serial_path_runs_one_at_a_time_and_honours_cancel() { - let ws = prepared( - Plan::new(vec![ - task("a", "write", vec!["x"]), - task("b", "write", vec!["y"]), - ]) - .stage(&Policy), - ); - let active = Arc::new(AtomicUsize::new(0)); - let peak = Arc::new(AtomicUsize::new(0)); - let (a, p) = (active.clone(), peak.clone()); - run( - ws.clone(), - 4, - CancellationToken::new(), - |_| (), - move |_w, _| { - let (a, p) = (a.clone(), p.clone()); - async move { - let now = a.fetch_add(1, Ordering::SeqCst) + 1; - p.fetch_max(now, Ordering::SeqCst); - tokio::task::yield_now().await; - a.fetch_sub(1, Ordering::SeqCst); - } - }, - ) - .await - .unwrap(); - assert_eq!(peak.load(Ordering::SeqCst), 1); - - let cancel = CancellationToken::new(); - let c2 = cancel.clone(); - let ran = Arc::new(AtomicUsize::new(0)); - let r = ran.clone(); - let err = run( - ws, - 4, - cancel, - |_| (), - move |_w, _| { - let (c, r) = (c2.clone(), r.clone()); - async move { - r.fetch_add(1, Ordering::SeqCst); - c.cancel(); - } - }, - ) - .await - .unwrap_err(); - assert!(matches!( - err, - tinyagents_harness::TinyAgentsError::Cancelled - )); - assert_eq!(ran.load(Ordering::SeqCst), 1); -} - -#[tokio::test] -async fn async_preflight_prepares_each_task_before_admission_staging() { - let out = Plan::new(vec![task("a", "read", vec![]), task("b", "read", vec![])]) - .stage_with_preflight(&Policy, |index, mut task| async move { - tokio::task::yield_now().await; - if index == 1 { - return Err("workspace isolation failed".to_owned()); - } - task.payload.kind = "iso"; - Ok(task) - }) - .await; - - assert!( - matches!(out[0], Preflight::Prepared(ref w) if w.dispatch_mode == WorkerDispatchMode::Parallel) - ); - assert!( - matches!(&out[1], Preflight::Rejected(r) if r.kind == RejectionKind::Policy("workspace isolation failed".into())) - ); -} - -struct R { - id: &'static str, - ok: bool, - files: Vec<String>, -} -impl BranchReport for R { - fn branch_id(&self) -> &str { - self.id - } - fn succeeded(&self) -> bool { - self.ok - } - fn changed_files(&self) -> &[String] { - &self.files - } -} - -#[test] -fn collected_counts_failures_and_reports_overlaps() { - let c = Collected::new(vec![ - R { - id: "a", - ok: true, - files: vec!["x.rs".into(), "y.rs".into()], - }, - R { - id: "b", - ok: false, - files: vec!["x.rs".into()], - }, - R { - id: "c", - ok: true, - files: vec![], - }, - ]); - assert_eq!((c.total(), c.succeeded(), c.failures), (3, 2, 1)); - assert_eq!( - c.overlap_warnings, - vec![OverlapWarning { - file: "x.rs".into(), - workers: vec!["a".into(), "b".into()] - }] - ); -} diff --git a/crates/tinyagents-orchestration/src/parallel/staging.rs b/crates/tinyagents-orchestration/src/parallel/staging.rs deleted file mode 100644 index 19a0a933d..000000000 --- a/crates/tinyagents-orchestration/src/parallel/staging.rs +++ /dev/null @@ -1,223 +0,0 @@ -//! Admission and shared-workspace arbitration for a batch of parallel tasks. - -use std::collections::HashMap; -use std::future::Future; -use std::path::PathBuf; - -use tinyagents_graph::parallel::{ - ClaimConflict, DispatchMode, WorkspaceClaim, plan_shared_workspace_dispatch, -}; - -/// One branch requested by the caller. `P` is the host's own payload. -#[derive(Debug, Clone)] -pub struct Task<P> { - /// Stable branch id (also the workspace-claim worker id). - pub id: String, - /// The instruction the branch runs. - pub prompt: String, - /// Free-form ownership boundary, folded into the prompt when present. - pub ownership: Option<String>, - /// Host payload carried through to the worker untouched. - pub payload: P, -} - -/// The host's verdict on one task, in terms the arbiter understands. -#[derive(Debug, Clone)] -pub enum AdmissionVerdict { - /// The branch never writes the shared workspace. - ReadOnly, - /// The branch has its own isolated workspace root. - Isolated, - /// The branch writes the shared workspace under these relative paths. - Writes(Vec<PathBuf>), - /// The host refuses the branch; the string is the host's own sentence. - Reject(String), -} - -/// Host policy: identity, allowlists, sandbox mode and whether a branch can -/// write. Everything product-specific lives behind this trait. -pub trait AdmissionPolicy<P>: Send + Sync { - fn admit(&self, index: usize, task: &Task<P>) -> AdmissionVerdict; -} - -/// How a prepared worker runs relative to its siblings. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum WorkerDispatchMode { - Parallel, - /// Writes the shared workspace; the whole batch falls back to serial. - SerialSharedWorkspaceWrite, -} - -/// Why a task did not become a [`Worker`]. -#[derive(Debug, Clone, PartialEq, Eq)] -pub enum RejectionKind { - /// The host policy refused it. - Policy(String), - /// Its claim conflicts with a sibling; render the data into a message. - Conflict(ClaimConflict), -} - -/// A rejected task, kept in input position. -#[derive(Debug, Clone)] -pub struct Rejection<P> { - pub task: Task<P>, - pub kind: RejectionKind, -} - -/// A task cleared for execution. -#[derive(Debug, Clone)] -pub struct Worker<P> { - pub id: String, - /// Prompt with the ownership boundary applied. - pub prompt: String, - pub ownership: Option<String>, - pub payload: P, - pub dispatch_mode: WorkerDispatchMode, -} - -/// Staging outcome for one input task. -#[derive(Debug, Clone)] -pub enum Preflight<P> { - Prepared(Worker<P>), - Rejected(Rejection<P>), -} - -/// An ordered batch of tasks. -#[derive(Debug, Clone)] -pub struct Plan<P> { - pub tasks: Vec<Task<P>>, -} - -impl<P> Plan<P> { - pub fn new(tasks: Vec<Task<P>>) -> Self { - Self { tasks } - } - - /// Admit every task, then arbitrate the shared-workspace claims of the - /// admitted ones in a single planner call. The result is a pure function - /// of the input order, one entry per input task. - pub fn stage(self, policy: &impl AdmissionPolicy<P>) -> Vec<Preflight<P>> { - enum Slot<P> { - Admitted(Task<P>, WorkspaceClaim), - Rejected(Task<P>, String), - } - let slots: Vec<Slot<P>> = self - .tasks - .into_iter() - .enumerate() - .map(|(i, task)| match policy.admit(i, &task) { - AdmissionVerdict::Reject(msg) => Slot::Rejected(task, msg), - AdmissionVerdict::ReadOnly => { - let c = WorkspaceClaim::read_only(task.id.clone()); - Slot::Admitted(task, c) - } - AdmissionVerdict::Isolated => { - let c = WorkspaceClaim::isolated(task.id.clone()); - Slot::Admitted(task, c) - } - AdmissionVerdict::Writes(paths) => { - let c = WorkspaceClaim::writing(task.id.clone(), paths); - Slot::Admitted(task, c) - } - }) - .collect(); - - let claims: Vec<WorkspaceClaim> = slots - .iter() - .filter_map(|s| match s { - Slot::Admitted(_, c) => Some(c.clone()), - Slot::Rejected(..) => None, - }) - .collect(); - let plan = plan_shared_workspace_dispatch(&claims); - let conflicts: HashMap<usize, ClaimConflict> = plan.conflicts.iter().cloned().collect(); - - let mut admitted_index = 0usize; - slots - .into_iter() - .map(|slot| match slot { - Slot::Rejected(task, msg) => Preflight::Rejected(Rejection { - task, - kind: RejectionKind::Policy(msg), - }), - Slot::Admitted(task, _) => { - let index = admitted_index; - admitted_index += 1; - if let Some(conflict) = conflicts.get(&index) { - return Preflight::Rejected(Rejection { - task, - kind: RejectionKind::Conflict(conflict.clone()), - }); - } - let dispatch_mode = match plan.modes.get(index).copied().flatten() { - Some(DispatchMode::Serial) => { - WorkerDispatchMode::SerialSharedWorkspaceWrite - } - Some(DispatchMode::Parallel) | None => WorkerDispatchMode::Parallel, - }; - let prompt = with_ownership_boundary(&task.prompt, task.ownership.as_deref()); - Preflight::Prepared(Worker { - id: task.id, - prompt, - ownership: task.ownership, - payload: task.payload, - dispatch_mode, - }) - } - }) - .collect() - } - - /// Asynchronously prepares each task before admission and workspace - /// arbitration. Hosts can use this phase to create isolated workspaces or - /// attach prepared resources to the task payload. A failed preparation is - /// returned in its original position as a policy rejection, and the - /// admission policy still receives the original task index. - pub async fn stage_with_preflight<F, Fut>( - self, - policy: &dyn AdmissionPolicy<P>, - mut preflight: F, - ) -> Vec<Preflight<P>> - where - P: Clone, - F: FnMut(usize, Task<P>) -> Fut, - Fut: Future<Output = Result<Task<P>, String>>, - { - let mut tasks = Vec::with_capacity(self.tasks.len()); - let mut failures = HashMap::new(); - for (index, task) in self.tasks.into_iter().enumerate() { - match preflight(index, task.clone()).await { - Ok(prepared) => tasks.push(prepared), - Err(reason) => { - failures.insert(index, reason); - tasks.push(task); - } - } - } - Plan::new(tasks).stage(&PreflightAdmission { policy, failures }) - } -} - -struct PreflightAdmission<'a, P> { - policy: &'a dyn AdmissionPolicy<P>, - failures: HashMap<usize, String>, -} - -impl<P> AdmissionPolicy<P> for PreflightAdmission<'_, P> { - fn admit(&self, index: usize, task: &Task<P>) -> AdmissionVerdict { - match self.failures.get(&index) { - Some(reason) => AdmissionVerdict::Reject(reason.clone()), - None => self.policy.admit(index, task), - } - } -} - -/// Fold an ownership boundary into the branch prompt. -pub fn with_ownership_boundary(prompt: &str, ownership: Option<&str>) -> String { - match ownership.map(str::trim).filter(|s| !s.is_empty()) { - Some(boundary) => format!( - "[Ownership Boundary]\n{boundary}\n\n[Task]\n{prompt}\n\nDo not work outside the ownership boundary unless the parent explicitly asks you to." - ), - None => prompt.to_string(), - } -} diff --git a/crates/tinyagents-orchestration/src/parallel/tool.rs b/crates/tinyagents-orchestration/src/parallel/tool.rs deleted file mode 100644 index ab993e60d..000000000 --- a/crates/tinyagents-orchestration/src/parallel/tool.rs +++ /dev/null @@ -1,175 +0,0 @@ -//! Host-bound parallel fan-out tool. - -use std::{future::Future, marker::PhantomData, pin::Pin, sync::Arc}; - -use async_trait::async_trait; -use serde::Serialize; -use serde_json::{Value, json}; -use tinyagents_harness::{context::RunContext, ids::CallId, tool::ToolDispatch}; -use tinytools::{Tool, ToolResult}; - -use super::{AdmissionPolicy, Plan, Preflight, RejectionKind, Task, Worker, run}; - -type Prepare<'a, P> = Pin<Box<dyn Future<Output = Result<Task<P>, String>> + Send + 'a>>; -type PreflightFn<P, Ctx> = - dyn for<'a> Fn(usize, Task<P>, &'a RunContext<Ctx>) -> Prepare<'a, P> + Send + Sync; - -/// Host-composed model tool that stages and runs independent branches. -pub struct ParallelFanoutTool<State, Ctx, P, BranchCtx, Output, MakeTasks, MakeBranch, Execute, Fut> -where - Ctx: Send + Sync + 'static, - P: Send + Sync + 'static, - BranchCtx: Send + 'static, - Output: Serialize + Send + 'static, - MakeTasks: Fn(&State, &Value) -> anyhow::Result<Vec<Task<P>>> + Send + Sync + 'static, - MakeBranch: Fn(usize, &RunContext<Ctx>) -> BranchCtx + Send + Sync + 'static, - Execute: Fn(Worker<P>, BranchCtx) -> Fut + Send + Sync + 'static, - Fut: Future<Output = Output> + Send + 'static, -{ - name: String, - description: String, - schema: Value, - tasks: Arc<MakeTasks>, - policy: Arc<dyn AdmissionPolicy<P> + Send + Sync>, - preflight: Arc<PreflightFn<P, Ctx>>, - branch_context: Arc<MakeBranch>, - execute: Arc<Execute>, - max_concurrency: usize, - state: PhantomData<fn(&State)>, -} - -impl<State, Ctx, P, BranchCtx, Output, MakeTasks, MakeBranch, Execute, Fut> - ParallelFanoutTool<State, Ctx, P, BranchCtx, Output, MakeTasks, MakeBranch, Execute, Fut> -where - Ctx: Send + Sync + 'static, - P: Send + Sync + 'static, - BranchCtx: Send + 'static, - Output: Serialize + Send + 'static, - MakeTasks: Fn(&State, &Value) -> anyhow::Result<Vec<Task<P>>> + Send + Sync + 'static, - MakeBranch: Fn(usize, &RunContext<Ctx>) -> BranchCtx + Send + Sync + 'static, - Execute: Fn(Worker<P>, BranchCtx) -> Fut + Send + Sync + 'static, - Fut: Future<Output = Output> + Send + 'static, -{ - /// Creates a tool from host-supplied task parsing, policy, async preflight, - /// per-branch context factory, and executor. - #[allow(clippy::too_many_arguments)] - pub fn new( - name: impl Into<String>, - description: impl Into<String>, - schema: Value, - tasks: MakeTasks, - policy: Arc<dyn AdmissionPolicy<P> + Send + Sync>, - preflight: impl for<'a> Fn(usize, Task<P>, &'a RunContext<Ctx>) -> Prepare<'a, P> - + Send - + Sync - + 'static, - branch_context: MakeBranch, - execute: Execute, - max_concurrency: usize, - ) -> Self { - Self { - name: name.into(), - description: description.into(), - schema, - tasks: Arc::new(tasks), - policy, - preflight: Arc::new(preflight), - branch_context: Arc::new(branch_context), - execute: Arc::new(execute), - max_concurrency, - state: PhantomData, - } - } -} - -#[async_trait] -impl<State, Ctx, P, BranchCtx, Output, MakeTasks, MakeBranch, Execute, Fut> ToolDispatch<State, Ctx> - for ParallelFanoutTool<State, Ctx, P, BranchCtx, Output, MakeTasks, MakeBranch, Execute, Fut> -where - State: Send + Sync + 'static, - Ctx: Send + Sync + 'static, - P: Send + Sync + Clone + 'static, - BranchCtx: Send + 'static, - Output: Serialize + Send + 'static, - MakeTasks: Fn(&State, &Value) -> anyhow::Result<Vec<Task<P>>> + Send + Sync + 'static, - MakeBranch: Fn(usize, &RunContext<Ctx>) -> BranchCtx + Send + Sync + 'static, - Execute: Fn(Worker<P>, BranchCtx) -> Fut + Send + Sync + 'static, - Fut: Future<Output = Output> + Send + 'static, -{ - fn tool(&self) -> Arc<dyn Tool> { - Arc::new(FanoutDeclaration { - name: self.name.clone(), - description: self.description.clone(), - schema: self.schema.clone(), - }) - } - - async fn execute( - &self, - state: &State, - _call_id: CallId, - args: Value, - _options: tinytools::ToolCallOptions, - parent: &RunContext<Ctx>, - ) -> anyhow::Result<ToolResult> { - let tasks = (self.tasks)(state, &args)?; - let prepared = Plan::new(tasks) - .stage_with_preflight(self.policy.as_ref(), |index, task| { - (self.preflight)(index, task, parent) - }) - .await; - let mut workers = Vec::new(); - let mut rejected = Vec::new(); - for (index, item) in prepared.into_iter().enumerate() { - match item { - Preflight::Prepared(worker) => workers.push((index, worker)), - Preflight::Rejected(rejection) => rejected.push(json!({ - "index": index, "task_id": rejection.task.id, - "reason": match rejection.kind { RejectionKind::Policy(reason) => reason, - RejectionKind::Conflict(conflict) => format!("workspace conflict: {conflict:?}") } - })), - } - } - let indices: Vec<_> = workers.iter().map(|(index, _)| *index).collect(); - let workers: Vec<_> = workers.into_iter().map(|(_, worker)| worker).collect(); - let contexts = Arc::clone(&self.branch_context); - let parent_ref = parent; - let executor = Arc::clone(&self.execute); - let results = run( - workers, - self.max_concurrency, - parent.cancellation.clone(), - move |branch| contexts(indices[branch], parent_ref), - move |worker, context| executor(worker, context), - ) - .await?; - let outcomes: Vec<_> = results - .into_iter() - .map(serde_json::to_value) - .collect::<Result<_, _>>()?; - Ok(ToolResult::json( - json!({"results":outcomes,"rejected":rejected}), - )) - } -} - -struct FanoutDeclaration { - name: String, - description: String, - schema: Value, -} -#[async_trait] -impl Tool for FanoutDeclaration { - fn name(&self) -> &str { - &self.name - } - fn description(&self) -> &str { - &self.description - } - fn parameters_schema(&self) -> Value { - self.schema.clone() - } - async fn execute(&self, _args: Value) -> anyhow::Result<ToolResult> { - anyhow::bail!("parallel fan-out requires typed-parent dispatch") - } -} diff --git a/crates/tinyagents-orchestration/src/parallel/tool_tests.rs b/crates/tinyagents-orchestration/src/parallel/tool_tests.rs deleted file mode 100644 index 10adb3159..000000000 --- a/crates/tinyagents-orchestration/src/parallel/tool_tests.rs +++ /dev/null @@ -1,67 +0,0 @@ -use super::*; -use serde_json::json; -use std::sync::Arc; -use tinyagents_harness::{ - context::{RunConfig, RunContext}, - ids::CallId, - tool::ToolDispatch, -}; - -struct ReadOnly; -impl AdmissionPolicy<String> for ReadOnly { - fn admit(&self, _: usize, _: &Task<String>) -> AdmissionVerdict { - AdmissionVerdict::ReadOnly - } -} - -#[tokio::test] -async fn fanout_runs_preflighted_branches_in_input_order_and_reports_rejections() { - let tool = ParallelFanoutTool::<(), (), String, (), String, _, _, _, _>::new( - "spawn_parallel", - "Run independent branches", - json!({"type":"object"}), - |_, _| { - Ok(vec![ - Task { - id: "a".into(), - prompt: "first".into(), - ownership: None, - payload: "one".into(), - }, - Task { - id: "b".into(), - prompt: "second".into(), - ownership: None, - payload: "two".into(), - }, - ]) - }, - Arc::new(ReadOnly), - |_, mut task, _| { - Box::pin(async move { - task.prompt.push_str(" prepared"); - Ok(task) - }) - }, - |_, _| (), - |worker, ()| async move { format!("{}:{}", worker.id, worker.prompt) }, - 2, - ); - let parent = RunContext::new(RunConfig::new("parent"), ()); - let result = ToolDispatch::<(), ()>::execute( - &tool, - &(), - CallId::new("parallel"), - json!({}), - tinytools::ToolCallOptions::default(), - &parent, - ) - .await - .unwrap(); - let value: serde_json::Value = serde_json::from_str(&result.output()).unwrap(); - assert_eq!( - value["results"], - json!(["a:first prepared", "b:second prepared"]) - ); - assert_eq!(value["rejected"], json!([])); -} diff --git a/crates/tinyagents-orchestration/src/parallel/workers.rs b/crates/tinyagents-orchestration/src/parallel/workers.rs deleted file mode 100644 index 2eaadc45d..000000000 --- a/crates/tinyagents-orchestration/src/parallel/workers.rs +++ /dev/null @@ -1,86 +0,0 @@ -//! Running staged workers: serial fallback for shared-workspace writers, -//! bounded `map_reduce` fan-out otherwise. - -use std::future::Future; - -use tinyagents_graph::parallel::{FailurePolicy, ParallelOptions, map_reduce}; -use tinyagents_harness::{CancellationToken, Result, TinyAgentsError}; - -use super::staging::{Worker, WorkerDispatchMode}; - -/// Run `workers` and return one result per worker, in input order. -/// -/// `branch_ctx` is called once per worker to build that branch's own context -/// (per-branch isolation: no state is shared between branches unless the host -/// puts it in the context deliberately). If any worker is -/// [`WorkerDispatchMode::SerialSharedWorkspaceWrite`] the whole batch runs one -/// at a time, checking `cancel` before each; otherwise up to `max_concurrency` -/// (0 = one slot per worker) run at once. Cancellation surfaces as -/// [`TinyAgentsError::Cancelled`]. -pub async fn run<P, C, R, B, F, Fut>( - workers: Vec<Worker<P>>, - max_concurrency: usize, - cancel: CancellationToken, - branch_ctx: B, - exec: F, -) -> Result<Vec<R>> -where - B: Fn(usize) -> C, - F: Fn(Worker<P>, C) -> Fut, - Fut: Future<Output = R>, -{ - let n = workers.len(); - let serial = workers - .iter() - .any(|w| w.dispatch_mode == WorkerDispatchMode::SerialSharedWorkspaceWrite); - if serial { - tracing::debug!( - workers = n, - "[parallel] serial fallback for shared-workspace writers" - ); - let mut results = Vec::with_capacity(n); - for (i, worker) in workers.into_iter().enumerate() { - if cancel.is_cancelled() { - return Err(TinyAgentsError::Cancelled); - } - results.push(exec(worker, branch_ctx(i)).await); - } - return Ok(results); - } - - let concurrency = if max_concurrency == 0 { - n.max(1) - } else { - max_concurrency - }; - tracing::debug!(workers = n, concurrency, "[parallel] bounded fan-out"); - let options = ParallelOptions::default() - .with_max_concurrency(concurrency) - .with_failure_policy(FailurePolicy::CollectAll) - .with_cancellation(cancel); - let outcome = map_reduce(workers, options, |i, worker| { - let fut = exec(worker, branch_ctx(i)); - async move { Ok(fut.await) } - }) - .await?; - - let mut results = Vec::with_capacity(n); - for item in outcome.outcomes { - match item.result { - Ok(v) => results.push(v), - Err(err) => { - return Err(TinyAgentsError::Graph(format!( - "parallel fan-out: worker {} failed: {err}", - item.index - ))); - } - } - } - if results.len() != n { - return Err(TinyAgentsError::Graph(format!( - "parallel fan-out: expected {n} result(s), got {}", - results.len() - ))); - } - Ok(results) -} diff --git a/crates/tinyagents-orchestration/src/status/mod.rs b/crates/tinyagents-orchestration/src/status/mod.rs deleted file mode 100644 index 565df4204..000000000 --- a/crates/tinyagents-orchestration/src/status/mod.rs +++ /dev/null @@ -1,45 +0,0 @@ -//! Conversions between the status vocabularies that describe one unit of -//! background agent work. -//! -//! [`OrchestrationTaskStatus`] is the one durable lifecycle status. The other -//! enums stay because their serialized forms are persisted or consumed -//! elsewhere (ledger rows, completion records, job snapshots, projected -//! transcripts), but each is expressed in terms of the canonical one through -//! `From` (total) or `TryFrom` (fallible, [`NoEquivalentStatus`]) impls. The -//! single mapping table is in the `tinyagents-tasks` README. -//! -//! Dependency direction decides where an impl lives: the ledger and -//! transcript-view conversions sit in `tinyagents-session`, the completion -//! conversions in `tinyagents-tasks`, and the job, detached and outcome ones -//! beside their types in this crate. The free functions below are the former -//! entry points and now delegate to those impls. - -use tinyagents_session::run_ledger::AgentRunStatus; -use tinyagents_tasks::OrchestrationTaskStatus; - -pub use tinyagents_tasks::NoEquivalentStatus; - -/// Maps a managed-task status onto the run-ledger status. -/// -/// Lossy (many-to-one): `CancelRequested` -> `Running`, `TimedOut` -> -/// `Failed`, `Abandoned` -> `Interrupted`, `Awaiting` -> `AwaitingUser`. -#[deprecated(since = "2.1.4", note = "use `AgentRunStatus::from(task_status)`")] -pub fn task_status_to_run_status(status: OrchestrationTaskStatus) -> AgentRunStatus { - AgentRunStatus::from(status) -} - -/// Maps a run-ledger status onto the managed-task status. -/// -/// Lossy (many-to-one): `AwaitingUser` and `Paused` -> `Awaiting`, -/// `Interrupted` -> `Abandoned`. -#[deprecated( - since = "2.1.4", - note = "use `OrchestrationTaskStatus::from(run_status)`" -)] -pub fn run_status_to_task_status(status: AgentRunStatus) -> OrchestrationTaskStatus { - OrchestrationTaskStatus::from(status) -} - -#[cfg(test)] -#[path = "mod_tests.rs"] -mod tests; diff --git a/crates/tinyagents-orchestration/src/status/mod_tests.rs b/crates/tinyagents-orchestration/src/status/mod_tests.rs deleted file mode 100644 index 5303182b9..000000000 --- a/crates/tinyagents-orchestration/src/status/mod_tests.rs +++ /dev/null @@ -1,45 +0,0 @@ -use super::*; -use tinyagents_session::run_ledger::AgentRunStatus as Run; -use tinyagents_tasks::OrchestrationTaskStatus as Task; - -const ALL_TASKS: [Task; 9] = [ - Task::Pending, - Task::Running, - Task::Awaiting, - Task::Completed, - Task::Failed, - Task::CancelRequested, - Task::Cancelled, - Task::TimedOut, - Task::Abandoned, -]; - -const ALL_RUNS: [Run; 8] = [ - Run::Pending, - Run::Running, - Run::AwaitingUser, - Run::Paused, - Run::Completed, - Run::Failed, - Run::Cancelled, - Run::Interrupted, -]; - -// The deprecated free functions are kept for source compatibility; this is the -// one place that still calls them, to prove they equal the `From` impls. -#[test] -#[allow(deprecated)] -fn deprecated_free_functions_delegate_to_the_from_impls() { - for task in ALL_TASKS { - assert_eq!(task_status_to_run_status(task), Run::from(task), "{task:?}"); - } - for run in ALL_RUNS { - assert_eq!(run_status_to_task_status(run), Task::from(run), "{run:?}"); - } -} - -#[test] -fn no_equivalent_status_is_the_tasks_crate_type() { - let err: tinyagents_tasks::NoEquivalentStatus = NoEquivalentStatus::new("awaiting", "X"); - assert_eq!(err.to_string(), "status `awaiting` has no equivalent in X"); -} diff --git a/crates/tinyagents-orchestration/src/steering.rs b/crates/tinyagents-orchestration/src/steering.rs index 287176d53..107d4704d 100644 --- a/crates/tinyagents-orchestration/src/steering.rs +++ b/crates/tinyagents-orchestration/src/steering.rs @@ -32,7 +32,7 @@ pub enum SteerError { UnsupportedLane, /// The `request_id` exceeds [`RecentRequestIds::MAX_REQUEST_ID_BYTES`]. RequestIdTooLong, - /// An [`SteerAccess::Owner`] steer's text exceeds [`MAX_STEER_TEXT_BYTES`]. + /// An owner supplied a message larger than the bounded steer payload. TextTooLong, /// The registry failed internally (its lock was poisoned by a panic); the /// subagent may still be live, so this is not [`Self::Unknown`]. @@ -72,12 +72,9 @@ pub struct SteerReceipt { /// Whose authority a steer is made under. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum SteerAccess<'a> { - /// Only the owning parent session may steer. The text is a parent - /// agent's, so it is length-capped, fenced and labelled as agent-origin - /// (see [`fence_agent_message`]), never presented as the user's. + /// Only the owning parent session may steer. Owner(&'a str), - /// A trusted control surface (a real user): no ownership check, and the - /// text is framed as the user's. + /// A trusted control surface: no ownership check. Trusted, } @@ -114,9 +111,6 @@ pub fn queue_lane_name(lane: QueueLane) -> &'static str { /// The steering command a lane maps to; `None` for a lane a detached subagent /// does not accept. -/// -/// The text is framed as the user's, or as the parent agent's when it is a -/// [`fence_agent_message`] fence ([`lane_prefix`]). pub fn steering_command_for_lane(lane: QueueLane, text: &str) -> Option<SteeringCommand> { match lane { QueueLane::Steer | QueueLane::Collect => Some(SteeringCommand::InjectMessage( @@ -136,11 +130,32 @@ fn send_registered(handle: &SteeringHandle, text: &str, lane: QueueLane) -> bool } } +fn live_handle_or_fallback<M>( + registry: &DetachedTaskRegistry<M, DetachedSubagentStatus>, + task_id: &TaskId, + access: SteerAccess<'_>, +) -> Result<Option<SteeringHandle>, SteerError> +where + M: Clone + Send + Sync + 'static, +{ + let result = match access { + SteerAccess::Owner(owner) => registry.steering_handle(task_id, owner), + SteerAccess::Trusted => registry.steering_handle_trusted(task_id), + }; + match result { + Ok(handle) => Ok(Some(handle)), + Err( + DetachedTaskRegistryError::NoSteeringHandle | DetachedTaskRegistryError::AlreadyDone, + ) => Ok(None), + Err(error) => Err(error.into()), + } +} + /// Inject `text` into the running subagent `task_id`. /// -/// Rejects an unsupported lane or unknown / unowned task, then prefers the -/// live steering handle and falls back to the host when the child is finished -/// or has no live handle. The text is never logged here. +/// Rejects an unsupported lane or unknown or unowned task, then prefers the +/// live steering handle and falls back to the host when the handle is absent +/// or the retained task has finished. The text is never logged here. pub async fn steer_detached<M, F>( registry: &DetachedTaskRegistry<M, DetachedSubagentStatus>, task_id: &str, @@ -153,6 +168,11 @@ where M: Clone + Send + Sync + 'static, F: SteerFallback<M>, { + if matches!(access, SteerAccess::Owner(_)) + && text.len() > tinyagents_harness::steering::MAX_STEER_TEXT_BYTES + { + return Err(SteerError::TextTooLong); + } let receipt = steer_detached_with_request_id(registry, task_id, access, text, lane, None, fallback) .await?; @@ -168,9 +188,7 @@ where /// never consumes its id and can be retried. A terminal task with retained /// metadata is delivered through the host fallback. A `request_id` /// longer than [`RecentRequestIds::MAX_REQUEST_ID_BYTES`] is rejected with -/// [`SteerError::RequestIdTooLong`]. An [`SteerAccess::Owner`] steer longer -/// than [`MAX_STEER_TEXT_BYTES`] is rejected with [`SteerError::TextTooLong`]; -/// a shorter one is delivered fenced, labelled as the parent agent's. +/// [`SteerError::RequestIdTooLong`]. /// /// # Delivery semantics /// @@ -210,6 +228,9 @@ where SteerAccess::Owner(_) => fence_agent_message(&text), SteerAccess::Trusted => text, }; + // A task may turn terminal between the retained snapshot and this live + // handle check. Resolve that race before claiming the request id. + let handle = live_handle_or_fallback(registry, &key, access)?; if let Some(request_id) = request_id && !registry.claim_steer_request(&key, request_id)? { @@ -219,22 +240,14 @@ where duplicate: true, }); } - let handle = match access { - SteerAccess::Owner(owner) => registry.steering_handle(&key, owner), - SteerAccess::Trusted => registry.steering_handle_trusted(&key), - }; match handle { - Ok(handle) if !snapshot.status.is_terminal() && send_registered(&handle, &text, lane) => { + Some(handle) if send_registered(&handle, &text, lane) => { return Ok(SteerReceipt { route: Some(SteerRoute::Registry), duplicate: false, }); } - Err( - DetachedTaskRegistryError::NoSteeringHandle | DetachedTaskRegistryError::AlreadyDone, - ) - | Ok(_) => {} - Err(error) => return Err(error.into()), + Some(_) | None => {} } fallback.enqueue(snapshot.metadata, lane, text).await; Ok(SteerReceipt { diff --git a/crates/tinyagents-orchestration/src/steering_tests.rs b/crates/tinyagents-orchestration/src/steering_tests.rs index fc558c571..fc7e48089 100644 --- a/crates/tinyagents-orchestration/src/steering_tests.rs +++ b/crates/tinyagents-orchestration/src/steering_tests.rs @@ -1,9 +1,8 @@ use super::*; -use std::sync::Arc; use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::{Arc, Mutex}; -use std::future::Ready; use tinyagents_harness::CancellationToken; use tinyagents_harness::ids::TaskId; use tinyagents_harness::run_queue::QueueLane; @@ -29,17 +28,7 @@ impl SubagentIdentity for Meta { } type Registry = DetachedTaskRegistry<Meta, DetachedSubagentStatus>; - -struct RecordingFallback(Arc<std::sync::Mutex<Option<(String, QueueLane, String)>>>); - -impl SteerFallback<Meta> for RecordingFallback { - type Future = Ready<()>; - - fn enqueue(self, metadata: Meta, lane: QueueLane, text: String) -> Self::Future { - *self.0.lock().unwrap() = Some((metadata.agent_id().to_string(), lane, text)); - std::future::ready(()) - } -} +type RecordedSteer = (Option<&'static str>, String, QueueLane, String); fn registry(steering: SteeringRegistry) -> Registry { DetachedTaskRegistry::new(steering, 256, DetachedSubagentStatus::is_terminal) @@ -87,10 +76,10 @@ fn lane_names_and_commands_are_literal() { } #[tokio::test] -async fn steer_rejects_followup_unknown_unowned_and_falls_back_for_terminal() { +async fn steer_rejects_followup_unknown_and_unowned() { let steering = SteeringRegistry::default(); let reg = registry(steering.clone()); - let tx = add(®, "t1", "p1", None); + let _tx = add(®, "t1", "p1", None); let (calls, fb) = counting(); assert_eq!( @@ -132,27 +121,48 @@ async fn steer_rejects_followup_unknown_unowned_and_falls_back_for_terminal() { .await, Err(SteerError::NotOwned) ); - tx.send(DetachedSubagentStatus::Failed { error: "e".into() }) - .unwrap(); +} + +#[tokio::test] +async fn terminal_child_uses_fallback_with_metadata_and_does_not_drain_stale_handle() { + let steering = SteeringRegistry::default(); + let reg = registry(steering.clone()); + let tx = add(®, "t1", "p1", Some("thread-1")); let handle = SteeringHandle::allow_all(); steering.register(TaskId::new("t1"), handle.clone()); - let delivered = Arc::new(std::sync::Mutex::new(None)); + tx.send(DetachedSubagentStatus::Failed { + error: "finished".into(), + }) + .unwrap(); + + let delivered: Arc<Mutex<Option<RecordedSteer>>> = Arc::default(); let captured = delivered.clone(); + let fallback = move |metadata: Meta, lane: QueueLane, text: String| { + *captured.lock().unwrap() = + Some((metadata.thread, metadata.agent_id().to_string(), lane, text)); + std::future::ready(()) + }; + assert_eq!( steer_detached( ®, "t1", - SteerAccess::Trusted, - "x".into(), - QueueLane::Steer, - RecordingFallback(captured) + SteerAccess::Owner("p1"), + "continue with the result".into(), + QueueLane::Collect, + fallback, ) .await, Ok(SteerRoute::Fallback) ); assert_eq!( *delivered.lock().unwrap(), - Some(("researcher".to_string(), QueueLane::Steer, "x".to_string())) + Some(( + Some("thread-1"), + "researcher".into(), + QueueLane::Collect, + "<parent_agent_message>\n\"continue with the result\"\n</parent_agent_message>".into(), + )) ); assert!( handle.drain().is_empty(), @@ -160,6 +170,24 @@ async fn steer_rejects_followup_unknown_unowned_and_falls_back_for_terminal() { ); } +#[tokio::test] +async fn terminal_transition_after_snapshot_selects_fallback() { + let steering = SteeringRegistry::default(); + let reg = registry(steering); + let tx = add(®, "t1", "p1", None); + let task_id = TaskId::new("t1"); + + let snapshot = reg.snapshot(&task_id, "p1").unwrap(); + assert!(!snapshot.status.is_terminal()); + tx.send(DetachedSubagentStatus::Failed { error: "e".into() }) + .unwrap(); + + assert!(matches!( + live_handle_or_fallback(®, &task_id, SteerAccess::Owner("p1")), + Ok(None) + )); +} + #[tokio::test] async fn steer_prefers_the_live_handle_then_falls_back() { let steering = SteeringRegistry::default(); @@ -393,6 +421,7 @@ async fn owner_steer_over_the_cap_is_refused_but_trusted_is_not() { ); assert_eq!(calls.load(Ordering::SeqCst), 0); let delivered = Arc::new(std::sync::Mutex::new(None)); + let delivered_by_fallback = Arc::clone(&delivered); assert_eq!( steer_detached( ®, @@ -400,7 +429,11 @@ async fn owner_steer_over_the_cap_is_refused_but_trusted_is_not() { SteerAccess::Trusted, long.clone(), QueueLane::Steer, - RecordingFallback(delivered.clone()) + move |metadata: Meta, lane, text| { + *delivered_by_fallback.lock().unwrap() = + Some((metadata.agent_id().to_owned(), lane, text)); + std::future::ready(()) + } ) .await, Ok(SteerRoute::Fallback) diff --git a/crates/tinyagents-orchestration/src/subagent/README.md b/crates/tinyagents-orchestration/src/subagent/README.md index 38331b89a..1c5be5ec3 100644 --- a/crates/tinyagents-orchestration/src/subagent/README.md +++ b/crates/tinyagents-orchestration/src/subagent/README.md @@ -2,27 +2,62 @@ `tinyagents_orchestration::subagent` owns child-agent orchestration policy. Generic job and message tool shells live in `tinyagents_harness_tools::subagents`. -It contains two complementary surfaces. - -The invocation surface provides `SubAgent`, `SubAgentTool`, -`SubAgentJobRegistry`, `SubAgentSession`, and `ChildDataPolicy`. It wraps an -`AgentHarness` as a child run, enforces recursion depth, and inherits live -parent cancellation, events, and host authority. Model-facing delegation is -background by default: `SubAgentTool` returns a job id immediately while the -child continues in the background. Callers can request `mode: "inline"` to -wait for the child and receive its final result in the same tool call. A host -constructs `tinyagents_harness_tools::subagents::SubAgentJobsTool` and +It exposes a driver-backed asynchronous job registry and durable lifecycle ports. + +The async surface provides `SubagentDriver` and `SubAgentJobRegistry`. A +`DriverSubAgentTool` call always registers and starts a background job, then +returns its task reference immediately. Hosts wait separately through the +registry or detached settlement APIs. A host constructs +`tinyagents_harness_tools::subagents::SubAgentJobsTool` and `SubAgentMessageTool` over the same registry to query status/results and send messages to queued or running jobs. Message delivery is cooperative at the child loop's next safe steering checkpoint. The registry is process-local and host-owned; durable lifecycle persistence remains the separate surface below. +The registry retains up to 128 unpinned terminal records. Settling or directly +cancelling a job prunes the oldest eligible terminal records immediately; +waiter-pinned results remain available until their wait futures consume or +drop. Total entries may temporarily exceed 128 while jobs are active or +terminal jobs have waiters. Queued/running work and `AwaitingInput` sessions +are never cap-pruned, so the latter remain available for reuse by their scoped +task key. A completed result remains queryable until retention pressure +removes it. + +Hosts that need complete results from one top-level run should acquire +`SubAgentJobRegistry::retain_root_run_descendants` before the root starts +admitting child jobs. After the root turn returns and can no longer admit +direct children, call `RootRunDescendantLease::settle` with one absolute +deadline and the caller's cancellation token. The lease keeps completed +ancestors and descendants available across terminal-retention pressure while +the subtree is active, then releases its pins on success, timeout, +cancellation, or drop. Owner filtering is applied to every descendant. A new +lease captures currently retained records and descendants admitted while it +is active; it cannot recover records already pruned before acquisition or +after an earlier lease ended. + +Driver jobs retain native parent-to-child-attempt edges. +`DriverSubAgentTool` records the actual caller and child context run ids; +retries and durable continuations append new edges while preserving earlier +attempts. Settlement traverses only edges reachable from the requested run +and returns final snapshots, including failures and cancellations. Paused +`AwaitingInput` jobs remain nonterminal and keep settlement open. The older +`settle_root_run_descendants` method remains available for callers that only +need currently retained records; it does not pin results before settlement +starts. A root is not itself a registry job, so hosts must close its direct +admissions before calling either settlement method. + +The legacy `register_driver_job` API can recover a historical parent run from +its durable key but cannot infer the native child attempt id. Use +`register_driver_job_with_parent_run` when the live caller context is +available, then report every actual attempt with `set_attempt_run_id`; the +settlement traversal can only follow descendants whose actual child attempt +ids have been recorded. The durable lifecycle surface coordinates a host-resolved subagent lifecycle without making product-policy decisions. Its dependency direction is: ```text tinyagents-harness-tools::subagents -> orchestration::subagent -orchestration::subagent -> {tinyagents-graph, tinyagents-harness, tinyagents-session} +orchestration::subagent -> {tinyagents-harness, tinyagents-session} ``` Hosts provide three object-safe lifecycle seams. `SubagentPlanner<C, H>` @@ -115,27 +150,26 @@ Enforcement uses *reservation* semantics: `try_reserve` checks every limit and claims the slot under one lock, returning a `SpawnReservation` guard, so concurrent spawns cannot race past a cap. Dropping the guard releases the live slot; dropping it without `commit()` (the spawn failed before the child -launched) also refunds the total budget. `SubAgentTool` (background and inline) +launched) also refunds the total budget. `DriverSubAgentTool` reserves before creating the child and moves the guard into the child's task -for its whole lifetime (a panic, abort, or dropped inline call still releases +for its whole lifetime (a panic or abort still releases it); `SubagentDriver::run` reserves after the cancellation and resume checks and before the planner runs, and commits just before the executor launches. Coalesced followers and cached terminal results never reserve; a resumed lifecycle takes a live slot but no new total budget. -An over-limit `SubAgentTool` call returns a tool error worded as a limit signal +An over-limit `DriverSubAgentTool` call returns a tool error worded as a limit signal ("...treat this as a delegated-agent limit signal, not a completed answer"), like the depth-limit error; the driver returns `SubagentError::SpawnRejected`. Because the neutral `SubagentRequest` carries no agent identity, a driver host that configures `allowed_targets` must name each request with `SubagentRequest::with_target`; an unnamed request is refused (fail closed). -**Bypass.** Only `SubAgentTool` and `SubagentDriver` enforce admission. Calling -`SubAgent::invoke_in_parent` / `invoke_hosted_in_parent` (or `SubAgentSession`) +**Bypass.** Admission is enforced through `SubagentDriver`. directly bypasses it; a host exposing those paths must reserve itself. Every limit defaults to `None` (unlimited), so nothing changes until a host opts -in with `SubAgentTool::with_spawn_admission` / +in with `DriverSubAgentTool::with_spawn_admission` / `SubagentDriver::with_spawn_admission`. Recommended guard-rail starting point: `max_children_per_parent: Some(5)` (OpenClaw's default) and a `max_total_per_root` of a few dozen. @@ -180,18 +214,18 @@ watching across a pause and skips a child whose ledger shows a cancellation. ## Policy, result and role (D5/D6/D9) One `SubAgentPolicy { timeout, retry, budget, retry_after_tool_calls }` governs -every path (it is defined in `tinyagents-graph` and re-exported here, because -orchestration depends on graph). +every path (it is defined in orchestration and temporarily re-exported by +the former graph subagent node). - **Timeout.** `SubagentDriver` runs the executor under its own child token; on elapse the child is cancelled (the caller's token is not) and the outcome is - `Incomplete(kind: Timeout)`. `SubAgentTool` cancels the job token and marks + `Incomplete(kind: Timeout)`. `DriverSubAgentTool` cancels the job token and marks the job `Incomplete` (`incomplete_kind: "timeout"`). A timeout is never retried. - **Retry.** Only a failure the `RetryPolicy` deems retryable, and never after the attempt ran tools unless `retry_after_tool_calls`. The driver retries `SubagentError::Transient { tools_ran, .. }` (adapters classify; anything else is never retried) and needs `PreparedSubagent::retry_context`, a factory for a - fresh run context per attempt; without it one attempt runs. `SubAgentTool` + fresh run context per attempt; without it one attempt runs. `DriverSubAgentTool` pre-mints one child per attempt and watches `ToolStarted` events to learn whether tools ran. A host-visible side effect of retry: the job's recorded `subagent_run_id` is the first attempt's. @@ -225,45 +259,25 @@ orchestration depends on graph). `restrict_tools` drops `subagent_jobs`, `subagent_message` and the host's `delegation_tools` for a leaf and intersects with `tool_ceiling` so a child never widens what it inherited; the driver applies it to - `PreparedSubagent::tools`. `SubAgentTool` cannot filter a shared harness, so a + `PreparedSubagent::tools`. The host driver cannot filter a shared harness, so a leaf tool refuses to spawn when the child harness exposes a delegation tool. `subagent_framing(role, depth, task)` is an optional neutral preamble hosts may prepend; nothing applies it by default. **Retry notes.** Subagent retry compounds with the harness's per-call model retry (each attempt may retry model calls first), so keep one of the two low. -`SubAgentTool` attempts after the first get run ids `{first}-a{n}` (no extra +Driver attempts after the first get run ids `{first}-a{n}` (no extra child ordinals are consumed, so sibling ids do not depend on retry) and the job's `subagent_run_id` follows the attempt that is running. The retry predicate (`RetryPolicy::retry_on`) sees the real error on the tool path; the driver only has the `Transient` message, so it sees a synthetic tool error. Usage of a failed or timed-out attempt is not reported (an error carries no run). -## Agent steps: one engine for teams, workflows and the driver (D11) - -`run_agent_step(&AgentStepConfig, AgentStepIdentity, CancellationToken, work)` -adapts an opaque host worker (a team member closure, a workflow child -executor) to the driver's planner / executor / in-memory persistence seams and -runs it on a real `SubagentDriver`. `teams::run_member_graph_with` and -`WorkflowEngine::with_step_config` call it, so those steps honour `SpawnPolicy` -admission (scope = team id / workflow run id, target = member / agent id), -`SubAgentPolicy::timeout` and `ResultPolicy` and the typed -`SubagentOutcomeKind`. Retry (`StepWorkError::Transient`), token budget (`StepSuccess::usage`) and -`SubagentRole` are supported by `run_agent_step` itself, but the team and -workflow adapters have no source for them (their workers report no transient -failures or usage and expose no tools), so they only take effect for callers -that use `run_agent_step` directly. - -`AgentStepConfig::default()` is inert (unlimited admission, no timeout, one -attempt, no trimming), so `run_member_graph` and a `WorkflowEngine` without -`with_step_config` behave as before. A worker `Err` is returned unchanged as -`AgentStepError::Worker`. - -`tinyagents_graph::SubAgentNode` / `subagent_node` are `#[deprecated]` rather -than adapters: graph cannot depend on orchestration. Their behaviour is -unchanged; `SubAgentPolicy` is shared and not deprecated. - -`StepContext` hands the worker the role and the model/tool call caps (`max_model_calls`, `max_tool_calls`) to enforce, since the crate cannot count an opaque worker's calls. +## Graph compatibility + +the removed graph crate re-exported `SubAgentBudget`, `SubAgentPolicy`, +and `SubAgentOutput` from this crate while its deprecated subagent node is +removed. ## Breaking changes diff --git a/crates/tinyagents-orchestration/src/subagent/admission.rs b/crates/tinyagents-orchestration/src/subagent/admission.rs index 0c72090f2..cfef6b3c7 100644 --- a/crates/tinyagents-orchestration/src/subagent/admission.rs +++ b/crates/tinyagents-orchestration/src/subagent/admission.rs @@ -44,10 +44,8 @@ //! one ledger, and a host shares one instance across every tool/driver whose //! spawns should count against the same limits. //! -//! Only [`SubAgentTool`](super::SubAgentTool) and -//! [`SubagentDriver`](super::SubagentDriver) enforce admission. Calling -//! `SubAgent::invoke_in_parent` / `invoke_hosted_in_parent` directly bypasses -//! it; hosts that expose those paths must reserve themselves. +//! [`SubagentDriver`](super::SubagentDriver) enforces admission for every +//! async spawn. Hosts should not expose a second path around the driver. use std::collections::HashMap; use std::sync::Arc; @@ -250,9 +248,8 @@ impl SpawnAdmission { /// paid its budget. The caller must only use it after confirming the child /// exists. The driver treats a request that carries resume state (loaded /// from persistence, or supplied by the trusted host through - /// `SubagentRequest`) as a continuation; the model can only trigger fresh - /// spawns, through [`SubAgentTool`](super::SubAgentTool), which reserves - /// with [`Self::try_reserve`]. Calling this for a fresh spawn would bypass + /// `SubagentRequest`) as a continuation; fresh model spawns reserve with + /// [`Self::try_reserve`]. Calling this for a fresh spawn would bypass /// `max_total_per_root`. pub fn try_reserve_continuation( &self, diff --git a/crates/tinyagents-orchestration/src/subagent/awaiting_user.rs b/crates/tinyagents-orchestration/src/subagent/awaiting_user.rs index 0d2bd8c95..7211299c7 100644 --- a/crates/tinyagents-orchestration/src/subagent/awaiting_user.rs +++ b/crates/tinyagents-orchestration/src/subagent/awaiting_user.rs @@ -3,8 +3,10 @@ /// A subagent pause that needs a user answer before it can resume. #[derive(Clone, Debug, PartialEq, Eq)] pub struct SubagentAwaitingUser { - /// Durable task id accepted by the continuation tool. + /// Durable task id used for the child operation. pub task_id: String, + /// Stable session selector used to respawn/reuse the paused child. + pub task_key: String, /// Host-resolved agent identity accepted by the continuation tool. pub agent_id: String, /// Worker conversation id, when the host uses a separate worker thread. @@ -22,23 +24,25 @@ impl SubagentAwaitingUser { /// identifier cannot inject a second field or terminate the marker. pub fn render(&self) -> String { let task_id = encode(&self.task_id); + let task_key = encode(&self.task_key); let agent_id = encode(&self.agent_id); let worker_thread_id = encode(self.worker_thread_id.as_deref().unwrap_or("(none)")); let question = encode(&self.question); let resume_caveat = if self.checkpointed { "" } else { - " NOTE: this pause could NOT be saved to disk, so resuming may fail. If continue_subagent reports no checkpoint and no durable session, tell the user the sub-agent's progress was lost instead of retrying." + " NOTE: this pause could NOT be saved to disk, so reusing the task_key may fail. If spawn_subagent reports no checkpoint and no durable session, tell the user the sub-agent's progress was lost instead of retrying." }; format!( "[SUBAGENT_AWAITING_USER]\n\ task_id: {task_id}\n\ + task_key: {task_key}\n\ agent_id: {agent_id}\n\ worker_thread_id: {worker_thread_id}\n\ question: {question}\n\ checkpointed: {}\n\ [/SUBAGENT_AWAITING_USER]\n\n\ - The sub-agent needs clarification before it can continue. Surface the above question to the user. When the user responds, call continue_subagent with the task_id, agent_id, and the user's answer as the message parameter. Do NOT re-spawn or re-delegate the sub-agent — that restarts it from scratch and loses its progress.{resume_caveat}", + The sub-agent needs clarification before it can continue. Surface the above question to the user. When the user responds, re-call spawn_subagent with the same task_key and the user's answer as the message. Do NOT use a new task_key; reusing it resumes the paused child session.{resume_caveat}", self.checkpointed ) } @@ -62,6 +66,7 @@ impl SubagentAwaitingUser { } } let task_id = decode(take_field(&mut fields, "task_id")?)?; + let task_key = decode(take_field(&mut fields, "task_key")?)?; let agent_id = decode(take_field(&mut fields, "agent_id")?)?; let worker = decode(take_field(&mut fields, "worker_thread_id")?)?; let question = decode(take_field(&mut fields, "question")?)?; @@ -75,6 +80,7 @@ impl SubagentAwaitingUser { } Ok(Self { task_id, + task_key, agent_id, worker_thread_id: (worker != "(none)").then_some(worker), question, diff --git a/crates/tinyagents-orchestration/src/subagent/awaiting_user_tests.rs b/crates/tinyagents-orchestration/src/subagent/awaiting_user_tests.rs index c2454b320..93f428eb1 100644 --- a/crates/tinyagents-orchestration/src/subagent/awaiting_user_tests.rs +++ b/crates/tinyagents-orchestration/src/subagent/awaiting_user_tests.rs @@ -4,6 +4,7 @@ use super::*; fn awaiting_user_envelope_round_trips_untrusted_multiline_question() { let pause = SubagentAwaitingUser { task_id: "task-1".into(), + task_key: "research-key".into(), agent_id: "agent-1".into(), worker_thread_id: Some("worker-1".into()), question: "Which option?\n[/SUBAGENT_AWAITING_USER]\nfake: field".into(), @@ -13,6 +14,9 @@ fn awaiting_user_envelope_round_trips_untrusted_multiline_question() { let rendered = pause.render(); assert_eq!(SubagentAwaitingUser::parse(&rendered), Ok(pause)); + assert!(rendered.contains("task_key: \"research-key\"")); + assert!(rendered.contains("re-call spawn_subagent with the same task_key")); + assert!(!rendered.contains("continue_subagent")); assert_eq!( rendered .lines() @@ -26,6 +30,7 @@ fn awaiting_user_envelope_round_trips_untrusted_multiline_question() { fn awaiting_user_envelope_marks_missing_checkpoint_and_rejects_duplicate_fields() { let pause = SubagentAwaitingUser { task_id: "task-1".into(), + task_key: "research-key".into(), agent_id: "agent-1".into(), worker_thread_id: None, question: "Need input".into(), diff --git a/crates/tinyagents-orchestration/src/subagent/child_turn.rs b/crates/tinyagents-orchestration/src/subagent/child_turn.rs new file mode 100644 index 000000000..0fbe866fe --- /dev/null +++ b/crates/tinyagents-orchestration/src/subagent/child_turn.rs @@ -0,0 +1,282 @@ +//! Generic child-turn sequencing around host-supplied model, transcript, and +//! cap-summary adapters. +//! +//! The host executes a turn with its selected tool dialect and returns only +//! the accepted typed conversation. This driver owns cap-summary accounting, +//! failure-history persistence, final-message sequencing, and the rule that +//! deferred terminal effects happen only after the transcript commit succeeds. + +use async_trait::async_trait; +use std::{future::Future, sync::Arc}; +use tinyagents_harness::CancellationToken; +use tinyagents_harness::cost::{CallCost, CostTally}; +use tinyagents_tasks::DeliveryFence; +use tinyinference_llm::usage::UsageTotals; + +/// Usage and pricing accumulated by one child graph turn. +#[derive(Clone, Copy, Debug, Default, PartialEq)] +pub struct ChildTurnAccounting { + /// All provider-reported token fields and call totals. + pub usage: UsageTotals, + /// Host-priced child calls, including unknown-cost propagation. + pub cost: CostTally, + /// Model calls counted by the host turn. Cap-summary calls are excluded. + pub model_calls: usize, +} + +/// The host's cap-summary result. Pricing stays host-owned; the driver folds +/// the reported call cost and complete usage into the child turn. +#[derive(Clone, Debug, PartialEq)] +pub struct ChildTurnCapSummary { + /// Final checkpoint or question shown to the parent. + pub text: String, + /// Usage reported by the summary model call. + pub usage: UsageTotals, + /// Price resolved by the host for that summary call. + pub cost: CallCost, +} + +/// A successfully executed child turn before optional cap summarization. +#[derive(Clone, Debug, PartialEq)] +pub struct ChildTurnRun<Message> { + /// Accepted transcript history in the host's native typed representation. + /// Provider-rejected or malformed suffixes must not be included. + pub accepted_history: Vec<Message>, + /// Visible turn output before an optional cap summary. + pub output: String, + /// Model calls made by the child turn. + pub model_calls: usize, + /// Provider usage and host-priced cost from child calls. + pub accounting: ChildTurnAccounting, + /// Whether the model-call cap stopped the turn with work pending. + pub hit_cap: bool, + /// Whether a tool requested an early exit such as awaiting user input. + pub early_exit: bool, +} + +/// Turn state returned after summary, transcript commit, and deferred effects. +#[derive(Clone, Debug, PartialEq)] +pub struct ChildTurnResult<Message> { + /// Accepted history, including any final checkpoint/question message. + pub accepted_history: Vec<Message>, + /// Final visible output. + pub output: String, + /// Folded child and summary accounting. + pub accounting: ChildTurnAccounting, + /// Whether the child stopped at its model-call cap. + pub hit_cap: bool, + /// Whether the child stopped at an early-exit tool. + pub early_exit: bool, +} + +/// A failed model turn with only the accepted transcript available for recovery. +#[derive(Clone, Debug, PartialEq)] +pub struct ChildTurnFailure<Message, Error> { + /// Original model, tool, or turn error. + pub error: Error, + /// Accepted typed history to persist before returning the error. + pub accepted_history: Vec<Message>, + /// Accounting for calls completed before failure. + pub accounting: ChildTurnAccounting, +} + +/// Cancellation and ownership capability retained for the whole child turn. +#[derive(Clone)] +pub struct ChildTurnControl { + cancellation: CancellationToken, + fence: Arc<dyn DeliveryFence>, +} + +impl ChildTurnControl { + /// Creates a turn control from the cancellation token and durable fence + /// acquired by the caller. A turn must never reacquire after fence loss. + pub fn new(cancellation: CancellationToken, fence: Arc<dyn DeliveryFence>) -> Self { + Self { + cancellation, + fence, + } + } + + /// Returns the cancellation token passed to host provider/tool work. + pub fn cancellation(&self) -> CancellationToken { + self.cancellation.clone() + } + + fn ensure_admitted(&self) -> Result<(), ChildTurnStop> { + if self.cancellation.is_cancelled() { + Err(ChildTurnStop::Cancelled) + } else if !self.fence.admitted() { + self.cancellation.cancel(); + Err(ChildTurnStop::LeaseLost) + } else { + Ok(()) + } + } + + async fn race<T, E>( + &self, + future: impl Future<Output = Result<T, E>>, + ) -> Result<Result<T, E>, ChildTurnStop> { + self.ensure_admitted()?; + tokio::select! { + biased; + _ = self.cancellation.cancelled() => Err(ChildTurnStop::Cancelled), + _ = self.fence.lost() => { + self.cancellation.cancel(); + Err(ChildTurnStop::LeaseLost) + } + result = future => { + self.ensure_admitted()?; + Ok(result) + } + } + } + + async fn race_host<T, E>( + &self, + future: impl Future<Output = Result<T, E>>, + ) -> Result<T, ChildTurnError<E>> { + match self.race(future).await.map_err(ChildTurnError::from)? { + Ok(value) => Ok(value), + Err(error) => Err(ChildTurnError::Host(error)), + } + } +} + +enum ChildTurnStop { + Cancelled, + LeaseLost, +} + +impl<Error> From<ChildTurnStop> for ChildTurnError<Error> { + fn from(stop: ChildTurnStop) -> Self { + match stop { + ChildTurnStop::Cancelled => Self::Cancelled, + ChildTurnStop::LeaseLost => Self::LeaseLost, + } + } +} + +/// Failure from host work or loss of the turn's cancellation/fencing authority. +#[derive(Debug, PartialEq, thiserror::Error)] +pub enum ChildTurnError<Error> { + /// A provider, transcript, or deferred host operation failed. + #[error("child turn host operation failed: {0}")] + Host(Error), + /// The parent explicitly cancelled this child turn. + #[error("child turn cancelled")] + Cancelled, + /// The durable thread/turn lease was lost and cannot be reacquired. + #[error("child turn lease lost")] + LeaseLost, +} + +/// Host-owned ports used by [`drive_child_turn`]. +#[async_trait] +pub trait ChildTurnHost: Send { + /// Typed conversation message representation selected by the host. + type Message: Clone + Send + Sync; + /// Host error type covering model, summary, and persistence failures. + type Error: Send; + + /// Run the provider/model loop and return accepted typed history. On + /// failure, retain and return the accepted-history snapshot for recovery. + async fn run_turn( + &mut self, + cancellation: CancellationToken, + ) -> Result<ChildTurnRun<Self::Message>, ChildTurnFailure<Self::Message, Self::Error>>; + + /// Produce a resumable checkpoint only when the child hit its cap. + async fn summarize_cap( + &mut self, + run: &ChildTurnRun<Self::Message>, + ) -> Result<ChildTurnCapSummary, Self::Error>; + + /// Append a final assistant checkpoint or early-exit question using the + /// selected dialect's typed transcript representation. + fn append_final_message( + &mut self, + history: &mut Vec<Self::Message>, + text: &str, + ) -> Result<(), Self::Error>; + + /// Commit successful transcript and metadata before deferred terminal work. + async fn persist_success( + &mut self, + result: &ChildTurnResult<Self::Message>, + ) -> Result<(), Self::Error>; + + /// Persist accepted partial history before the original turn error returns. + async fn persist_failure( + &mut self, + history: &[Self::Message], + accounting: &ChildTurnAccounting, + ) -> Result<(), Self::Error>; + + /// Perform deferred terminal effects after a successful transcript commit. + async fn after_commit( + &mut self, + result: &ChildTurnResult<Self::Message>, + cancellation: CancellationToken, + ) -> Result<(), Self::Error>; +} + +/// Drives one host-configured child turn while centralizing summary accounting +/// and durable commit ordering. +pub async fn drive_child_turn<H: ChildTurnHost>( + host: &mut H, + control: ChildTurnControl, +) -> Result<ChildTurnResult<H::Message>, ChildTurnError<H::Error>> { + let result = match control + .race(host.run_turn(control.cancellation())) + .await + .map_err(ChildTurnError::from)? + { + Ok(run) => run, + Err(failure) => { + control + .race_host(host.persist_failure(&failure.accepted_history, &failure.accounting)) + .await?; + return Err(ChildTurnError::Host(failure.error)); + } + }; + let mut run = result; + // This is the one canonical child-turn model-call counter. A cap summary + // contributes usage and price, but not to this host-defined count. + run.accounting.model_calls = run.model_calls; + + if run.hit_cap { + let summary = control.race_host(host.summarize_cap(&run)).await?; + run.output = summary.text; + run.accounting.usage.calls = run + .accounting + .usage + .calls + .saturating_add(summary.usage.calls); + run.accounting.usage.usage += summary.usage.usage; + run.accounting.cost.add(summary.cost); + } + + if (run.hit_cap || run.early_exit) && !run.output.trim().is_empty() { + control.ensure_admitted()?; + host.append_final_message(&mut run.accepted_history, &run.output) + .map_err(ChildTurnError::Host)?; + } + + let result = ChildTurnResult { + accepted_history: run.accepted_history, + output: run.output, + accounting: run.accounting, + hit_cap: run.hit_cap, + early_exit: run.early_exit, + }; + control.race_host(host.persist_success(&result)).await?; + control + .race_host(host.after_commit(&result, control.cancellation())) + .await?; + Ok(result) +} + +#[cfg(test)] +#[path = "child_turn_tests.rs"] +mod tests; diff --git a/crates/tinyagents-orchestration/src/subagent/child_turn_tests.rs b/crates/tinyagents-orchestration/src/subagent/child_turn_tests.rs new file mode 100644 index 000000000..a25c2d8c5 --- /dev/null +++ b/crates/tinyagents-orchestration/src/subagent/child_turn_tests.rs @@ -0,0 +1,394 @@ +use super::*; +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; +use tinyagents_harness::CancellationToken; +use tinyagents_harness::cost::{CallCost, CostSource}; +use tinyinference_llm::usage::{ChargedAmount, Usage}; +use tokio::sync::{Notify, oneshot}; + +#[derive(Default)] +struct FakeHost { + events: Vec<&'static str>, + run: Option<Result<ChildTurnRun<String>, ChildTurnFailure<String, &'static str>>>, + summary: Option<ChildTurnCapSummary>, + persisted: Option<ChildTurnResult<String>>, + failure_history: Option<Vec<String>>, + fail_persist: bool, + wait_on_run: bool, + run_started: Option<oneshot::Sender<()>>, + wait_on_summary: bool, + summary_started: Option<oneshot::Sender<()>>, + wait_on_persist: bool, + persist_started: Option<oneshot::Sender<()>>, +} + +#[async_trait::async_trait] +impl ChildTurnHost for FakeHost { + type Message = String; + type Error = &'static str; + + async fn run_turn( + &mut self, + cancellation: CancellationToken, + ) -> Result<ChildTurnRun<String>, ChildTurnFailure<String, &'static str>> { + if let Some(started) = self.run_started.take() { + let _ = started.send(()); + } + if self.wait_on_run { + cancellation.cancelled().await; + } + self.events.push("run"); + self.run.take().unwrap() + } + + async fn summarize_cap( + &mut self, + _: &ChildTurnRun<String>, + ) -> Result<ChildTurnCapSummary, &'static str> { + self.events.push("summary"); + if let Some(started) = self.summary_started.take() { + let _ = started.send(()); + } + if self.wait_on_summary { + std::future::pending().await + } + Ok(self.summary.take().unwrap()) + } + + fn append_final_message( + &mut self, + history: &mut Vec<String>, + text: &str, + ) -> Result<(), &'static str> { + self.events.push("append-final"); + history.push(format!("native:{text}")); + Ok(()) + } + + async fn persist_success( + &mut self, + result: &ChildTurnResult<String>, + ) -> Result<(), &'static str> { + self.events.push("persist-success"); + if let Some(started) = self.persist_started.take() { + let _ = started.send(()); + } + if self.wait_on_persist { + std::future::pending().await + } + if self.fail_persist { + return Err("persist failed"); + } + self.persisted = Some(result.clone()); + Ok(()) + } + + async fn persist_failure( + &mut self, + history: &[String], + _: &ChildTurnAccounting, + ) -> Result<(), &'static str> { + self.events.push("persist-failure"); + self.failure_history = Some(history.to_vec()); + Ok(()) + } + + async fn after_commit( + &mut self, + _: &ChildTurnResult<String>, + _: CancellationToken, + ) -> Result<(), &'static str> { + self.events.push("after-commit"); + Ok(()) + } +} + +fn active_control() -> ChildTurnControl { + ChildTurnControl::new(CancellationToken::new(), Arc::new(|| true)) +} + +struct RevocableFence { + admitted: AtomicBool, + lost: Notify, +} + +impl RevocableFence { + fn new() -> Self { + Self { + admitted: AtomicBool::new(true), + lost: Notify::new(), + } + } + + fn revoke(&self) { + self.admitted.store(false, Ordering::SeqCst); + self.lost.notify_waiters(); + } +} + +impl tinyagents_tasks::DeliveryFence for RevocableFence { + fn admitted(&self) -> bool { + self.admitted.load(Ordering::SeqCst) + } + + fn lost(&self) -> std::pin::Pin<Box<dyn std::future::Future<Output = ()> + Send + '_>> { + Box::pin(async move { self.lost.notified().await }) + } +} + +fn totals( + input: u64, + output: u64, + cache_read: u64, + cache_create: u64, + reasoning: u64, +) -> UsageTotals { + let mut totals = UsageTotals::default(); + totals.record(Usage { + input_tokens: input, + output_tokens: output, + total_tokens: input + output, + cache_read_tokens: cache_read, + cache_creation_tokens: cache_create, + reasoning_tokens: reasoning, + charged_amount: Some(ChargedAmount::usd_micros(250)), + charged_amount_nanos: None, + context_window_tokens: Some(32_000), + }); + totals +} + +#[tokio::test] +async fn cap_summary_folds_complete_usage_and_runs_after_transcript_commit() { + let accounting = ChildTurnAccounting { + usage: totals(10, 4, 3, 2, 1), + cost: { + let mut cost = CostTally::default(); + cost.add(CallCost::Charged(1.25)); + cost + }, + model_calls: 5, + }; + let mut host = FakeHost { + run: Some(Ok(ChildTurnRun { + accepted_history: vec!["accepted-native-assistant-tool-call".into()], + output: "partial".into(), + model_calls: 5, + accounting, + hit_cap: true, + early_exit: false, + })), + summary: Some(ChildTurnCapSummary { + text: "resume from checkpoint".into(), + usage: totals(7, 5, 2, 1, 3), + cost: CallCost::Estimated(2.25), + }), + ..FakeHost::default() + }; + + let result = drive_child_turn(&mut host, active_control()).await.unwrap(); + + assert_eq!( + host.events, + [ + "run", + "summary", + "append-final", + "persist-success", + "after-commit" + ] + ); + assert_eq!( + result.accepted_history, + [ + "accepted-native-assistant-tool-call", + "native:resume from checkpoint" + ] + ); + assert_eq!(result.output, "resume from checkpoint"); + assert_eq!(result.accounting.model_calls, 5); + assert_eq!(result.accounting.usage.calls, 2); + assert_eq!(result.accounting.usage.usage.input_tokens, 17); + assert_eq!(result.accounting.usage.usage.output_tokens, 9); + assert_eq!(result.accounting.usage.usage.total_tokens, 26); + assert_eq!(result.accounting.usage.usage.cache_read_tokens, 5); + assert_eq!(result.accounting.usage.usage.cache_creation_tokens, 3); + assert_eq!(result.accounting.usage.usage.reasoning_tokens, 4); + assert_eq!( + result.accounting.usage.usage.charged_amount, + Some(ChargedAmount::usd_micros(500)) + ); + assert_eq!(result.accounting.cost.known_usd, 3.5); + assert_eq!(result.accounting.cost.source, CostSource::Estimated); + assert_eq!(host.persisted, Some(result)); +} + +#[tokio::test] +async fn early_exit_appends_typed_question_without_cap_summary() { + let mut host = FakeHost { + run: Some(Ok(ChildTurnRun { + accepted_history: vec!["accepted-native-tool-result".into()], + output: "Which account should I use?".into(), + model_calls: 3, + accounting: ChildTurnAccounting::default(), + hit_cap: false, + early_exit: true, + })), + ..FakeHost::default() + }; + + let result = drive_child_turn(&mut host, active_control()).await.unwrap(); + + assert_eq!( + host.events, + ["run", "append-final", "persist-success", "after-commit"] + ); + assert_eq!( + result.accepted_history, + [ + "accepted-native-tool-result", + "native:Which account should I use?" + ] + ); + assert!(result.early_exit); + assert!(!result.hit_cap); + assert_eq!(result.accounting.model_calls, 3); + assert_eq!(host.persisted, Some(result)); +} + +#[tokio::test] +async fn failed_turn_persists_only_accepted_history_before_returning_error() { + let mut host = FakeHost { + run: Some(Err(ChildTurnFailure { + error: "provider failed", + accepted_history: vec!["accepted-user".into(), "accepted-assistant".into()], + accounting: ChildTurnAccounting { + model_calls: 1, + usage: totals(3, 2, 0, 0, 0), + ..ChildTurnAccounting::default() + }, + })), + ..FakeHost::default() + }; + + assert_eq!( + drive_child_turn(&mut host, active_control()).await, + Err(ChildTurnError::Host("provider failed")) + ); + assert_eq!(host.events, ["run", "persist-failure"]); + assert_eq!( + host.failure_history, + Some(vec!["accepted-user".into(), "accepted-assistant".into()]) + ); + assert!(host.persisted.is_none()); +} + +#[tokio::test] +async fn failed_transcript_commit_never_runs_deferred_terminal_effects() { + let mut host = FakeHost { + run: Some(Ok(ChildTurnRun { + accepted_history: vec!["accepted-xml-suffix".into()], + output: "done".into(), + model_calls: 2, + accounting: ChildTurnAccounting::default(), + hit_cap: false, + early_exit: false, + })), + fail_persist: true, + ..FakeHost::default() + }; + + assert_eq!( + drive_child_turn(&mut host, active_control()).await, + Err(ChildTurnError::Host("persist failed")) + ); + assert_eq!(host.events, ["run", "persist-success"]); + assert!(host.persisted.is_none()); +} + +#[tokio::test] +async fn losing_turn_fence_cancels_host_work_and_skips_persistence() { + let fence = Arc::new(RevocableFence::new()); + let cancellation = CancellationToken::new(); + let control = ChildTurnControl::new(cancellation.clone(), fence.clone()); + let (started, started_rx) = oneshot::channel(); + let mut host = FakeHost { + wait_on_run: true, + run_started: Some(started), + ..FakeHost::default() + }; + + let (outcome, ()) = tokio::join!(drive_child_turn(&mut host, control), async { + started_rx.await.unwrap(); + fence.revoke(); + }); + + assert_eq!(outcome, Err(ChildTurnError::LeaseLost)); + assert!(cancellation.is_cancelled()); + assert!(host.persisted.is_none()); + assert!(!host.events.contains(&"persist-success")); + assert!(!host.events.contains(&"after-commit")); +} + +#[tokio::test] +async fn losing_turn_fence_during_cap_summary_skips_transcript_commit() { + let fence = Arc::new(RevocableFence::new()); + let cancellation = CancellationToken::new(); + let control = ChildTurnControl::new(cancellation.clone(), fence.clone()); + let (started, started_rx) = oneshot::channel(); + let mut host = FakeHost { + run: Some(Ok(ChildTurnRun { + accepted_history: vec!["accepted".into()], + output: "partial".into(), + model_calls: 4, + accounting: ChildTurnAccounting::default(), + hit_cap: true, + early_exit: false, + })), + wait_on_summary: true, + summary_started: Some(started), + ..FakeHost::default() + }; + + let (outcome, ()) = tokio::join!(drive_child_turn(&mut host, control), async { + started_rx.await.unwrap(); + fence.revoke(); + }); + + assert_eq!(outcome, Err(ChildTurnError::LeaseLost)); + assert!(cancellation.is_cancelled()); + assert!(host.persisted.is_none()); + assert!(!host.events.contains(&"persist-success")); + assert!(!host.events.contains(&"after-commit")); +} + +#[tokio::test] +async fn losing_turn_fence_during_transcript_commit_skips_deferred_terminal_work() { + let fence = Arc::new(RevocableFence::new()); + let cancellation = CancellationToken::new(); + let control = ChildTurnControl::new(cancellation.clone(), fence.clone()); + let (started, started_rx) = oneshot::channel(); + let mut host = FakeHost { + run: Some(Ok(ChildTurnRun { + accepted_history: vec!["accepted".into()], + output: "done".into(), + model_calls: 1, + accounting: ChildTurnAccounting::default(), + hit_cap: false, + early_exit: false, + })), + wait_on_persist: true, + persist_started: Some(started), + ..FakeHost::default() + }; + + let (outcome, ()) = tokio::join!(drive_child_turn(&mut host, control), async { + started_rx.await.unwrap(); + fence.revoke(); + }); + + assert_eq!(outcome, Err(ChildTurnError::LeaseLost)); + assert!(cancellation.is_cancelled()); + assert!(host.persisted.is_none()); + assert!(!host.events.contains(&"after-commit")); +} diff --git a/crates/tinyagents-orchestration/src/subagent/completion.rs b/crates/tinyagents-orchestration/src/subagent/completion.rs index 0632c8bba..983894e7d 100644 --- a/crates/tinyagents-orchestration/src/subagent/completion.rs +++ b/crates/tinyagents-orchestration/src/subagent/completion.rs @@ -44,6 +44,7 @@ impl From<&AppliedResult> for CompletionResult { pub(crate) struct CompletionOrigin { task_id: String, parent_key: String, + parent_epoch: u64, agent_id: String, notify_mode: NotifyMode, } @@ -53,8 +54,11 @@ impl CompletionOrigin { pub(crate) fn new<C>( task_key: &SubagentTaskKey, prepared: &PreparedSubagent<C>, - ) -> Option<Self> { - let notify_mode = prepared.notify_mode?; + router: &CompletionRouter, + ) -> Result<Option<Self>, super::SubagentError> { + let Some(notify_mode) = prepared.notify_mode else { + return Ok(None); + }; // The thread outlives a restart; the parent run id does not, so it is // only the fallback for a thread-less parent. let parent_key = prepared @@ -62,23 +66,34 @@ impl CompletionOrigin { .clone() .or_else(|| task_key.thread_id.clone()) .unwrap_or_else(|| task_key.parent_run_id.clone()); - Some(Self { + let parent_epoch = router + .admit_child_generation(&parent_key) + .map_err(|error| super::SubagentError::Persistence(error.to_string()))? + .ok_or_else(|| { + super::SubagentError::Persistence(format!( + "parent lifecycle is closed: {parent_key:?}" + )) + })?; + Ok(Some(Self { task_id: task_key.task_id.clone(), parent_key, + parent_epoch, agent_id: prepared.agent_key.clone(), notify_mode, - }) + })) } fn record(&self, status: CompletionStatus, result: CompletionResult) -> CompletionRecord { - CompletionRecord::new( + let mut record = CompletionRecord::new( self.task_id.clone(), self.parent_key.clone(), self.agent_id.clone(), status, result, ) - .with_notify_mode(self.notify_mode) + .with_notify_mode(self.notify_mode); + record.parent_epoch = self.parent_epoch; + record } /// The record for a finished lifecycle, or `None` when the outcome is not a diff --git a/crates/tinyagents-orchestration/src/subagent/detached/README.md b/crates/tinyagents-orchestration/src/subagent/detached/README.md index fc641591f..4ebfcf92f 100644 --- a/crates/tinyagents-orchestration/src/subagent/detached/README.md +++ b/crates/tinyagents-orchestration/src/subagent/detached/README.md @@ -33,7 +33,11 @@ store lives, progress projection and policy. for children it runs itself, via `with_completion_router`. - Roster helpers: `snapshot_for_owner`, `task_id_for_session`, `task_id_for_session_in_records`, `resume_ref_for_task`, - `resume_ref_from_record`. + `resume_ref_from_record`, and `parent_roster_for_owner` / + `format_parent_roster_for_owner` for the actual subagent-job registry plus + durable session summaries. The latter requires summaries read through the + matching `ScopedSubagentSessionStore`; it preserves `awaiting_user` state for + host wording and orders entries by agent then task. ## Operational constraints diff --git a/crates/tinyagents-orchestration/src/subagent/detached/completion.rs b/crates/tinyagents-orchestration/src/subagent/detached/completion.rs index adc0cbb77..c661e9563 100644 --- a/crates/tinyagents-orchestration/src/subagent/detached/completion.rs +++ b/crates/tinyagents-orchestration/src/subagent/detached/completion.rs @@ -10,6 +10,7 @@ use std::sync::Arc; +use tinyagents_harness::error::TinyAgentsError; use tinyagents_harness::ids::TaskId; use tinyagents_tasks::{ CompletionRecord, CompletionResult, CompletionRouter, CompletionStatus, NotifyMode, @@ -21,6 +22,18 @@ use super::ledger::record_status_with_retries; use super::types::DetachedSubagentStatus; const LOG_PREFIX: &str = "[detached-completion]"; +const PAUSE_RECORD_PREFIX: &str = "\u{1}subagent-pause:"; + +/// Stable prefix for independently deliverable pause notices belonging to a +/// task. Hosts may use it to withdraw notices when the task is collected. +pub fn pause_record_prefix(task_id: &str) -> String { + let encoded_task: String = task_id + .as_bytes() + .iter() + .map(|byte| format!("{byte:02x}")) + .collect(); + format!("{PAUSE_RECORD_PREFIX}{encoded_task}:") +} /// Where a detached child's completion goes and how it announces itself. #[derive(Clone)] @@ -35,6 +48,8 @@ pub struct DetachedCompletionTarget { pub label: Option<String>, /// How the parent is told. pub notify_mode: NotifyMode, + /// Lifecycle generation captured when the child was admitted. + pub parent_epoch: u64, } impl DetachedCompletionTarget { @@ -50,9 +65,33 @@ impl DetachedCompletionTarget { agent_id: agent_id.into(), label: None, notify_mode: NotifyMode::default(), + parent_epoch: 0, } } + /// Creates a target only when the parent lifecycle is open and captures + /// the generation that every later pause and terminal record must retain. + pub fn try_new( + router: Arc<CompletionRouter>, + parent_key: impl Into<String>, + agent_id: impl Into<String>, + ) -> Result<Self, TinyAgentsError> { + let parent_key = parent_key.into(); + let parent_epoch = router.admit_child_generation(&parent_key)?.ok_or_else(|| { + TinyAgentsError::Graph(format!( + "detached child parent lifecycle is closed: {parent_key:?}" + )) + })?; + Ok(Self { + router, + parent_key, + agent_id: agent_id.into(), + label: None, + notify_mode: NotifyMode::default(), + parent_epoch, + }) + } + /// Sets the human label. pub fn with_label(mut self, label: impl Into<String>) -> Self { self.label = Some(label.into()); @@ -90,8 +129,47 @@ impl DetachedCompletionTarget { ) .with_notify_mode(self.notify_mode); record.label = self.label.clone(); + record.parent_epoch = self.parent_epoch; Some(record) } + + /// Creates a separate deliverable notice for a child pause. The unique + /// record id leaves `task_id` available for the child's eventual terminal + /// completion and permits the same child to pause more than once. + pub fn pause_record_for(&self, task_id: &str, notice: impl Into<String>) -> CompletionRecord { + let record_id = format!( + "{}{}", + pause_record_prefix(task_id), + uuid::Uuid::new_v4().simple() + ); + let mut record = CompletionRecord::new( + record_id, + self.parent_key.clone(), + self.agent_id.clone(), + CompletionStatus::Incomplete, + CompletionResult::text(notice), + ) + .with_notify_mode(self.notify_mode); + record.label = self.label.clone(); + record.parent_epoch = self.parent_epoch; + record + } + + /// Records an independently deliverable pause notice with the same retry + /// policy as terminal detached completions. + pub async fn record_pause(&self, task_id: &str, notice: impl Into<String>) -> bool { + let record = self.pause_record_for(task_id, notice); + match self.router.record_with_retries(record, 3).await { + Ok(outcome) => { + tracing::debug!(task_id, outcome = ?outcome, "{LOG_PREFIX} recorded pause notice"); + true + } + Err(error) => { + tracing::warn!(task_id, error = %error, "{LOG_PREFIX} could not record pause notice"); + false + } + } + } } /// Records the completion `status` stands for, if it is one. A router failure @@ -141,8 +219,9 @@ pub fn spawn_status_watcher_with_completions( let snapshot = status.borrow_and_update().clone(); match snapshot { DetachedSubagentStatus::Running => {} - DetachedSubagentStatus::AwaitingUser { .. } => { + DetachedSubagentStatus::AwaitingUser { ref question } => { record_status_with_retries(store.as_ref(), &task_id, &snapshot); + target.record_pause(&task_id, question.clone()).await; } _ => break snapshot, } diff --git a/crates/tinyagents-orchestration/src/subagent/detached/completion_tests.rs b/crates/tinyagents-orchestration/src/subagent/detached/completion_tests.rs index 1b696a4a5..1d172ffa6 100644 --- a/crates/tinyagents-orchestration/src/subagent/detached/completion_tests.rs +++ b/crates/tinyagents-orchestration/src/subagent/detached/completion_tests.rs @@ -18,12 +18,18 @@ fn router() -> Arc<CompletionRouter> { } fn target(router: &Arc<CompletionRouter>) -> DetachedCompletionTarget { - DetachedCompletionTarget::new(router.clone(), "p", "worker").with_label("indexer") + DetachedCompletionTarget::try_new(router.clone(), "p", "worker") + .expect("test parent starts open") + .with_label("indexer") } -async fn settle(router: &CompletionRouter, parent: &str) { +async fn settle(router: &CompletionRouter, parent: &str, task_id: &str) { for _ in 0..100 { - if !router.pending_for(parent).is_empty() { + if router + .pending_for(parent) + .iter() + .any(|record| record.task_id == task_id) + { return; } tokio::time::sleep(Duration::from_millis(10)).await; @@ -90,6 +96,42 @@ fn only_final_statuses_are_completions() { assert_eq!(failed.result.text, "bad"); } +#[tokio::test] +async fn pause_notices_use_independent_ids_and_leave_terminal_task_id_free() { + let router = router(); + let target = target(&router); + let first = target.pause_record_for("task:1", "question one"); + let second = target.pause_record_for("task:1", "question two"); + let prefix = pause_record_prefix("task:1"); + assert!(first.task_id.starts_with(&prefix)); + assert!(second.task_id.starts_with(&prefix)); + assert_ne!(first.task_id, second.task_id); + assert_eq!(first.status, CompletionStatus::Incomplete); + assert_eq!(first.parent_key, "p"); + assert_eq!(first.result.text, "question one"); + + assert!(target.record_pause("task:1", "question one").await); + assert!( + record_detached_completion( + &target, + "task:1", + &DetachedSubagentStatus::Completed { + output: "done".into(), + iterations: 1, + }, + ) + .await + ); + let pending = router.pending_for("p"); + assert_eq!(pending.len(), 2); + assert!(pending.iter().any(|record| record.task_id == "task:1")); + assert!( + pending + .iter() + .any(|record| record.task_id.starts_with(&prefix)) + ); +} + #[tokio::test] async fn the_watcher_mirrors_the_status_and_records_the_completion() { let router = router(); @@ -101,7 +143,7 @@ async fn the_watcher_mirrors_the_status_and_records_the_completion() { iterations: 1, }) .unwrap(); - settle(&router, "p").await; + settle(&router, "p", "t1").await; let pending = router.pending_for("p"); assert_eq!(pending[0].result.text, "found it"); let record = store.get(&"t1".into()).unwrap(); @@ -115,7 +157,7 @@ async fn a_dropped_sender_is_recorded_as_a_failure() { let (tx, rx) = watch::channel(DetachedSubagentStatus::Running); spawn_status_watcher_with_completions(store, "t1".into(), rx, target(&router)); drop(tx); - settle(&router, "p").await; + settle(&router, "p", "t1").await; assert_eq!(router.pending_for("p")[0].status, CompletionStatus::Failed); } @@ -143,11 +185,17 @@ async fn the_watcher_pushes_to_a_live_parent_by_notify_mode() { } tokio::time::sleep(Duration::from_millis(10)).await; } - assert_eq!(queue.drain(QueueLane::Collect).await.len(), 1); + let admitted = queue + .take_admitted( + QueueLane::Collect, + tinyagents_harness::run_queue::QueueMode::All, + ) + .await; + assert_eq!(admitted.len(), 1); } #[tokio::test] -async fn a_paused_child_records_nothing() { +async fn a_paused_child_records_a_separate_notice_not_its_terminal_task_id() { let router = router(); let store = spawned_store("t1"); let (tx, rx) = watch::channel(DetachedSubagentStatus::Running); @@ -157,7 +205,10 @@ async fn a_paused_child_records_nothing() { }) .unwrap(); tokio::time::sleep(Duration::from_millis(100)).await; - assert!(router.pending_for("p").is_empty()); + let pauses = router.pending_for("p"); + assert_eq!(pauses.len(), 1); + assert!(pauses[0].task_id.starts_with(&pause_record_prefix("t1"))); + assert_eq!(pauses[0].result.text, "which?"); } #[tokio::test] @@ -183,12 +234,97 @@ async fn a_resumed_child_is_recorded_when_it_finally_finishes() { }) .unwrap(); tokio::time::sleep(Duration::from_millis(50)).await; - assert!(router.pending_for("p").is_empty()); + assert_eq!(router.pending_for("p").len(), 1); tx.send(DetachedSubagentStatus::Completed { output: "answered".into(), iterations: 2, }) .unwrap(); - settle(&router, "p").await; - assert_eq!(router.pending_for("p")[0].result.text, "answered"); + settle(&router, "p", "t1").await; + let pending = router.pending_for("p"); + assert_eq!(pending.len(), 2); + assert!( + pending + .iter() + .any(|record| record.result.text == "answered") + ); +} + +#[tokio::test] +async fn watcher_targets_keep_the_admitted_parent_generation() { + let router = router(); + let old_target = DetachedCompletionTarget::try_new(router.clone(), "p", "worker").unwrap(); + let old_epoch = old_target.parent_epoch; + router.cancel_parent("p").unwrap(); + router.resume_parent("p"); + + assert!(old_target.record_pause("late-pause", "old question").await); + let old_store = spawned_store("old-child"); + let (old_tx, old_rx) = watch::channel(DetachedSubagentStatus::Running); + spawn_status_watcher_with_completions(old_store, "old-child".into(), old_rx, old_target); + old_tx + .send(DetachedSubagentStatus::Completed { + output: "old answer".into(), + iterations: 1, + }) + .unwrap(); + tokio::time::sleep(Duration::from_millis(100)).await; + assert!(router.pending_for("p").is_empty()); + + let new_target = DetachedCompletionTarget::try_new(router.clone(), "p", "worker").unwrap(); + assert!(new_target.parent_epoch > old_epoch); + let new_epoch = new_target.parent_epoch; + let new_store = spawned_store("new-child"); + let (new_tx, new_rx) = watch::channel(DetachedSubagentStatus::Running); + spawn_status_watcher_with_completions(new_store, "new-child".into(), new_rx, new_target); + new_tx + .send(DetachedSubagentStatus::AwaitingUser { + question: "new question".into(), + }) + .unwrap(); + let mut saw_pause = false; + for _ in 0..100 { + if router.pending_for("p").iter().any(|record| { + record + .task_id + .starts_with(&pause_record_prefix("new-child")) + }) { + saw_pause = true; + break; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + assert!( + saw_pause, + "watcher must persist the pause before it resumes" + ); + new_tx + .send(DetachedSubagentStatus::Completed { + output: "new answer".into(), + iterations: 1, + }) + .unwrap(); + for _ in 0..100 { + if router + .pending_for("p") + .iter() + .any(|record| record.task_id == "new-child") + { + break; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + let pending = router.pending_for("p"); + assert_eq!(pending.len(), 2); + assert!( + pending + .iter() + .all(|record| record.parent_epoch == new_epoch) + ); + assert!( + pending + .iter() + .any(|record| record.result.text == "new question") + ); + assert!(pending.iter().any(|record| record.task_id == "new-child")); } diff --git a/crates/tinyagents-orchestration/src/subagent/detached/ledger.rs b/crates/tinyagents-orchestration/src/subagent/detached/ledger.rs index bee16b36a..0deff034e 100644 --- a/crates/tinyagents-orchestration/src/subagent/detached/ledger.rs +++ b/crates/tinyagents-orchestration/src/subagent/detached/ledger.rs @@ -126,7 +126,7 @@ pub fn record_cancelled(store: &dyn TaskStore, task_id: &str) -> tinyagents_harn /// Watch a child's status channel and mirror its first terminal status into /// `store`. A dropped sender without a terminal status is recorded as /// [`DetachedSubagentStatus::ended_without_result`]. A failed write is retried -/// up to [`STATUS_WRITE_ATTEMPTS`] times before the watcher gives up; hosts +/// up to `STATUS_WRITE_ATTEMPTS` times before the watcher gives up; hosts /// that must observe persistence failures call [`record_status`] themselves. pub fn spawn_status_watcher( store: Arc<dyn TaskStore>, diff --git a/crates/tinyagents-orchestration/src/subagent/detached/mod.rs b/crates/tinyagents-orchestration/src/subagent/detached/mod.rs index 728ac022a..9d43a3c9c 100644 --- a/crates/tinyagents-orchestration/src/subagent/detached/mod.rs +++ b/crates/tinyagents-orchestration/src/subagent/detached/mod.rs @@ -32,7 +32,8 @@ mod status_map; mod types; pub use completion::{ - DetachedCompletionTarget, record_detached_completion, spawn_status_watcher_with_completions, + DetachedCompletionTarget, pause_record_prefix, record_detached_completion, + spawn_status_watcher_with_completions, }; pub use ledger::{ DETACHED_LEDGER_TIMEOUT_MS, list_subagent_records, orphaned_subagent_reason, record_agent_id, @@ -47,8 +48,9 @@ pub use ops::{ wait_detached_or_durable, wait_detached_or_durable_with_owners, }; pub use roster::{ - resume_ref_for_task, resume_ref_from_record, snapshot_for_owner, task_id_for_session, - task_id_for_session_in_records, + ParentRosterEntry, format_parent_roster_for_owner, format_roster_for_owner, + parent_roster_for_owner, resume_ref_for_task, resume_ref_from_record, snapshot_for_owner, + task_id_for_session, task_id_for_session_in_records, }; pub use status::{wait_detached, wait_error_from_registry}; pub use types::{ diff --git a/crates/tinyagents-orchestration/src/subagent/detached/mod_tests.rs b/crates/tinyagents-orchestration/src/subagent/detached/mod_tests.rs index 4a48bce15..595abb8b9 100644 --- a/crates/tinyagents-orchestration/src/subagent/detached/mod_tests.rs +++ b/crates/tinyagents-orchestration/src/subagent/detached/mod_tests.rs @@ -635,3 +635,126 @@ fn record_cancelled_is_a_noop_when_terminal_or_missing() { record_cancelled(&store, "c1").unwrap(); record_cancelled(&store, "missing").unwrap(); } + +#[tokio::test] +async fn host_roster_formatter_receives_sorted_owner_scoped_snapshot() { + let reg = registry(); + let _b = add( + ®, + "task-b", + "owner-a", + Meta { + agent: "worker-b", + session: None, + }, + ); + let _a = add( + ®, + "task-a", + "owner-a", + Meta { + agent: "worker-a", + session: None, + }, + ); + let _foreign = add( + ®, + "task-foreign", + "owner-b", + Meta { + agent: "worker-foreign", + session: None, + }, + ); + + let rendered = format_roster_for_owner(®, "owner-a", |snapshot| { + assert_eq!( + snapshot + .iter() + .map(|row| row.task_id.as_str()) + .collect::<Vec<_>>(), + ["task-a", "task-b"] + ); + format!("host roster: {} workers", snapshot.len()) + }) + .unwrap(); + assert_eq!(rendered, "host roster: 2 workers"); +} + +#[test] +fn actual_job_roster_merges_owner_filtered_live_jobs_and_waiting_sessions() { + let jobs = crate::subagent::SubAgentJobRegistry::new(); + let workspace = tinytools::WorkspaceDescriptor::new("/workspace/a"); + let scope = crate::subagent::ParentJobScope::new(workspace.clone(), "thread-a"); + let task_key = |task: &str, thread: &str| crate::subagent::SubagentTaskKey { + root_run_id: "root".into(), + parent_run_id: "parent".into(), + thread_id: Some(thread.into()), + task_id: task.into(), + }; + for (task, owner, thread, job_workspace) in [ + ("task-a", "owner-a", "thread-a", workspace.clone()), + ("task-foreign", "owner-b", "thread-a", workspace.clone()), + ( + "task-other-thread", + "owner-a", + "thread-b", + workspace.clone(), + ), + ] { + jobs.register_driver_job(crate::subagent::DriverJobRegistration { + key: task_key(task, thread), + agent: "worker".into(), + owner: owner.into(), + workspace: Some(job_workspace), + cancellation: CancellationToken::new(), + steering: None, + run_queue: None, + continuing: false, + }) + .unwrap(); + } + + let session = crate::subagent::DurableSubagentSessionSummary { + subagent_session_id: "session-a".into(), + parent_thread_id: Some("thread-a".into()), + worker_thread_id: Some("worker-thread".into()), + agent_id: "worker".into(), + display_name: Some("Worker A".into()), + model: Some("model".into()), + sandbox_mode: "none".into(), + action_root: None, + task_key: "research".into(), + task_title: "research".into(), + current_task_id: Some("task-a".into()), + status: crate::subagent::DurableSubagentStatus::AwaitingUser, + reusable: true, + latest_error: None, + created_at: "created".into(), + updated_at: "updated".into(), + last_used_at: "used".into(), + }; + let wrong_thread = crate::subagent::DurableSubagentSessionSummary { + parent_thread_id: Some("thread-b".into()), + current_task_id: Some("task-hidden".into()), + ..session.clone() + }; + let formatted = format_parent_roster_for_owner( + &jobs, + "owner-a", + &scope, + &[wrong_thread, session], + |entries| { + assert_eq!(entries.len(), 1); + assert_eq!(entries[0].task_id, "task-a"); + assert!(entries[0].live_job.is_some()); + assert_eq!( + entries[0].session.as_ref().unwrap().status, + crate::subagent::DurableSubagentStatus::AwaitingUser + ); + "host wording: worker is waiting for user input".to_owned() + }, + ) + .unwrap(); + assert_eq!(formatted, "host wording: worker is waiting for user input"); +} diff --git a/crates/tinyagents-orchestration/src/subagent/detached/roster.rs b/crates/tinyagents-orchestration/src/subagent/detached/roster.rs index 7a690c1fc..111fc7595 100644 --- a/crates/tinyagents-orchestration/src/subagent/detached/roster.rs +++ b/crates/tinyagents-orchestration/src/subagent/detached/roster.rs @@ -6,6 +6,89 @@ use super::types::{ DetachedSubagentStatus, SubagentIdentity, SubagentResumeRef, SubagentSnapshot, WaitError, }; +/// One owner-scoped live job and its matching durable reusable-session view. +#[derive(Clone, Debug)] +pub struct ParentRosterEntry { + /// Stable child task identity. + pub task_id: String, + /// Child agent identifier. + pub agent_id: String, + /// Current live job, when the child is still registered in this process. + pub live_job: Option<crate::subagent::SubAgentJob>, + /// Durable resumable-session summary for this task, when present. + pub session: Option<crate::subagent::DurableSubagentSessionSummary>, +} + +/// Builds a stable roster from the actual subagent-job registry and sessions +/// already read through a [`ScopedSubagentSessionStore`][crate::subagent::ScopedSubagentSessionStore] +/// bound to the same authenticated parent/workspace. Session summaries do not +/// carry workspace ownership, so callers must not pass summaries from a global +/// store or another workspace. Sessions with another parent thread are still +/// excluded here; the owner and workspace/thread scope live jobs. +pub fn parent_roster_for_owner( + jobs: &crate::subagent::SubAgentJobRegistry, + owner: &str, + scope: &crate::subagent::ParentJobScope, + sessions: &[crate::subagent::DurableSubagentSessionSummary], +) -> Result<Vec<ParentRosterEntry>, crate::subagent::SubAgentJobError> { + let mut entries = std::collections::BTreeMap::<String, ParentRosterEntry>::new(); + for job in jobs.list_for_parent_scope(owner, scope)? { + let task_id = job + .extensions + .get("task_id") + .and_then(serde_json::Value::as_str) + .unwrap_or(job.id.as_str()) + .to_owned(); + entries.insert( + task_id.clone(), + ParentRosterEntry { + task_id, + agent_id: job.agent.clone(), + live_job: Some(job), + session: None, + }, + ); + } + for session in sessions + .iter() + .filter(|session| session.parent_thread_id.as_deref() == Some(scope.thread_id.as_str())) + { + let Some(task_id) = session.current_task_id.as_deref() else { + continue; + }; + let entry = entries + .entry(task_id.to_owned()) + .or_insert_with(|| ParentRosterEntry { + task_id: task_id.to_owned(), + agent_id: session.agent_id.clone(), + live_job: None, + session: None, + }); + entry.session = Some(session.clone()); + } + let mut entries: Vec<_> = entries.into_values().collect(); + entries.sort_by(|left, right| { + left.agent_id + .cmp(&right.agent_id) + .then_with(|| left.task_id.cmp(&right.task_id)) + }); + Ok(entries) +} + +/// Applies host wording to the actual owner-filtered live and durable roster. +/// The callback receives stable data, including `AwaitingUser` session state, +/// without TinyAgents imposing product-specific prose. +pub fn format_parent_roster_for_owner( + jobs: &crate::subagent::SubAgentJobRegistry, + owner: &str, + scope: &crate::subagent::ParentJobScope, + sessions: &[crate::subagent::DurableSubagentSessionSummary], + format: impl FnOnce(&[ParentRosterEntry]) -> String, +) -> Result<String, crate::subagent::SubAgentJobError> { + let entries = parent_roster_for_owner(jobs, owner, scope, sessions)?; + Ok(format(&entries)) +} + /// Snapshot the subagents registered under `owner`, with live status, ordered /// by `agent_id` then `task_id` so a rendered roster is stable across turns. /// A poisoned registry lock surfaces as [`WaitError::RegistryPoisoned`]. @@ -34,6 +117,22 @@ where Ok(out) } +/// Formats an owner's stable roster through host-supplied wording. +/// +/// The callback receives only the owner-filtered, deterministically ordered +/// snapshot. TinyAgents supplies data and ordering; hosts own prompt wording. +pub fn format_roster_for_owner<M>( + registry: &DetachedTaskRegistry<M, DetachedSubagentStatus>, + owner: &str, + format: impl FnOnce(&[SubagentSnapshot]) -> String, +) -> Result<String, WaitError> +where + M: SubagentIdentity + Clone + Send + Sync + 'static, +{ + let snapshot = snapshot_for_owner(registry, owner)?; + Ok(format(&snapshot)) +} + /// Resolve a durable session id to the live task id, enforcing ownership. A /// non-terminal entry wins over a terminal one. pub fn task_id_for_session<M>( diff --git a/crates/tinyagents-orchestration/src/subagent/driver.rs b/crates/tinyagents-orchestration/src/subagent/driver.rs index b2c158380..4f44a65a2 100644 --- a/crates/tinyagents-orchestration/src/subagent/driver.rs +++ b/crates/tinyagents-orchestration/src/subagent/driver.rs @@ -1,6 +1,6 @@ -use std::{collections::HashMap, sync::Arc}; +use std::{collections::HashMap, future::Future, sync::Arc}; -use tinyagents_tasks::CompletionRouter; +use tinyagents_tasks::{CompletionRouter, DeliveryFence}; use tokio::sync::Mutex as AsyncMutex; use super::completion::{CompletionOrigin, deliver}; @@ -18,6 +18,44 @@ use tinyagents_harness::error::TinyAgentsError; const LOG_PREFIX: &str = "[subagent-driver]"; +fn ensure_turn_fence( + fence: Option<&Arc<dyn DeliveryFence>>, + cancellation: Option<&CancellationToken>, +) -> Result<(), SubagentError> { + if fence.is_some_and(|fence| !fence.admitted()) { + if let Some(cancellation) = cancellation { + cancellation.cancel(); + } + Err(SubagentError::Cancelled) + } else { + Ok(()) + } +} + +async fn await_with_turn_fence<T>( + fence: Option<&Arc<dyn DeliveryFence>>, + cancellation: Option<&CancellationToken>, + future: impl Future<Output = Result<T, SubagentError>>, +) -> Result<T, SubagentError> { + ensure_turn_fence(fence, cancellation)?; + let Some(fence) = fence else { + return future.await; + }; + tokio::select! { + biased; + _ = fence.lost() => { + if let Some(cancellation) = cancellation { + cancellation.cancel(); + } + Err(SubagentError::Cancelled) + } + result = future => { + ensure_turn_fence(Some(fence), cancellation)?; + result + } + } +} + /// Cloneable process-scope in-flight registry that can be shared by drivers. pub type SubagentInFlight = SharedFlightMap<SubagentTaskKey, SubagentRunResult, SubagentError>; use tinyagents_harness::{ @@ -174,15 +212,47 @@ impl<C: Send + 'static, H: Send + 'static> SubagentDriver<C, H> { request: SubagentRequest<C, H>, cancellation: CancellationToken, ) -> Result<SubagentRunResult, SubagentError> { + self.run_with_attempt_observer(request, cancellation, |_| {}) + .await + } + + /// Runs a lifecycle and reports each actual native attempt run id before + /// that attempt reaches the executor. Detached-job registries use this to + /// retain the complete parent-to-attempt lineage across retries. + pub async fn run_with_attempt_observer<F>( + &self, + request: SubagentRequest<C, H>, + cancellation: CancellationToken, + on_attempt: F, + ) -> Result<SubagentRunResult, SubagentError> + where + F: Fn(&RunId) + Send + Sync, + { let task_key = request.task_key().clone(); - if let Some(outcome) = self.terminal_outcomes.lock().await.get(&task_key).cloned() { + let turn_fence = request.turn_fence(); + ensure_turn_fence(turn_fence.as_ref(), None)?; + if let Some(outcome) = await_with_turn_fence(turn_fence.as_ref(), None, async { + Ok(self.terminal_outcomes.lock().await.get(&task_key).cloned()) + }) + .await? + { return Ok(SubagentRunResult::new( outcome, SubagentPersistenceDisposition::TerminalExisting, )); } - if let Some(outcome) = self.persistence.load_terminal(&task_key).await? { - self.cache_terminal(task_key, &outcome).await; + if let Some(outcome) = await_with_turn_fence( + turn_fence.as_ref(), + None, + self.persistence.load_terminal(&task_key), + ) + .await? + { + await_with_turn_fence(turn_fence.as_ref(), None, async { + self.cache_terminal(task_key, &outcome).await; + Ok(()) + }) + .await?; return Ok(SubagentRunResult::new( outcome, SubagentPersistenceDisposition::TerminalExisting, @@ -191,7 +261,11 @@ impl<C: Send + 'static, H: Send + 'static> SubagentDriver<C, H> { let (entry, is_leader) = self.in_flight.claim(task_key.clone()); if !is_leader { - return match entry.wait_or_cancel(&cancellation).await { + return match await_with_turn_fence(turn_fence.as_ref(), None, async { + Ok(entry.wait_or_cancel(&cancellation).await) + }) + .await? + { None => Ok(SubagentRunResult::new( SubagentOutcome::cancelled(task_key.task_id.clone()), SubagentPersistenceDisposition::ObserverCancelled, @@ -227,7 +301,11 @@ impl<C: Send + 'static, H: Send + 'static> SubagentDriver<C, H> { // A preceding caller may have committed a terminal result between the // initial cache check and this reservation. Do not reopen that task // after its in-flight entry has been removed. - if let Some(outcome) = self.terminal_outcomes.lock().await.get(&task_key).cloned() { + if let Some(outcome) = await_with_turn_fence(turn_fence.as_ref(), None, async { + Ok(self.terminal_outcomes.lock().await.get(&task_key).cloned()) + }) + .await? + { reservation.finish(Ok(SubagentRunResult::new( outcome.clone(), SubagentPersistenceDisposition::TerminalExisting, @@ -239,7 +317,7 @@ impl<C: Send + 'static, H: Send + 'static> SubagentDriver<C, H> { } let result = self - .run_reserved(request, task_key.clone(), cancellation) + .run_reserved(request, task_key.clone(), cancellation, &on_attempt) .await; reservation.finish(result.clone()); result @@ -250,17 +328,30 @@ impl<C: Send + 'static, H: Send + 'static> SubagentDriver<C, H> { mut request: SubagentRequest<C, H>, task_key: SubagentTaskKey, cancellation: CancellationToken, + on_attempt: &(impl Fn(&RunId) + Send + Sync), ) -> Result<SubagentRunResult, SubagentError> { let task_id = request.task_id().to_owned(); + let turn_fence = request.turn_fence(); + + ensure_turn_fence(turn_fence.as_ref(), None)?; if cancellation.is_cancelled() { - return self - .persist_cancelled(task_key, SubagentOutcome::cancelled(task_id), None) - .await; + let result = await_with_turn_fence( + turn_fence.as_ref(), + None, + self.persist_cancelled(task_key, SubagentOutcome::cancelled(task_id), None), + ) + .await?; + self.record_completion_under_fence(turn_fence.as_ref(), None, &result, 0) + .await?; + return Ok(result); } if request.resume().is_none() { - request.set_resume(self.persistence.load(&task_key).await?); + let resume = + await_with_turn_fence(turn_fence.as_ref(), None, self.persistence.load(&task_key)) + .await?; + request.set_resume(resume); } // Keep the exact pause seen at the durable read boundary. The // persistence seam uses it as a scoped compare-and-swap expectation @@ -268,13 +359,19 @@ impl<C: Send + 'static, H: Send + 'static> SubagentDriver<C, H> { // duplicate continuation returns the newer durable winner. let expected_pause = request.resume().cloned(); if cancellation.is_cancelled() { - return self - .persist_cancelled( + let result = await_with_turn_fence( + turn_fence.as_ref(), + None, + self.persist_cancelled( task_key, SubagentOutcome::cancelled(task_id), expected_pause, - ) - .await; + ), + ) + .await?; + self.record_completion_under_fence(turn_fence.as_ref(), None, &result, 0) + .await?; + return Ok(result); } // Reserve before any planner/executor work. A resumed lifecycle is an @@ -313,7 +410,9 @@ impl<C: Send + 'static, H: Send + 'static> SubagentDriver<C, H> { }; let mut reservation = reservation.map_err(SubagentError::SpawnRejected)?; - let mut prepared = self.planner.prepare(request).await?; + ensure_turn_fence(turn_fence.as_ref(), None)?; + let mut prepared = + await_with_turn_fence(turn_fence.as_ref(), None, self.planner.prepare(request)).await?; if prepared.task_id != task_id { return Err(SubagentError::TaskIdMismatch { expected: task_id, @@ -322,20 +421,30 @@ impl<C: Send + 'static, H: Send + 'static> SubagentDriver<C, H> { } // Known as soon as the planner has named the spawn's notify mode, so a // cancel from here on is recorded too. - let completion_origin = self - .completions - .as_ref() - .and_then(|_| CompletionOrigin::new(&task_key, &prepared)); + let completion_origin = if let Some(router) = &self.completions { + CompletionOrigin::new(&task_key, &prepared, router)? + } else { + None + }; + ensure_turn_fence(turn_fence.as_ref(), None)?; if cancellation.is_cancelled() { - let result = self - .persist_cancelled( + let result = await_with_turn_fence( + turn_fence.as_ref(), + None, + self.persist_cancelled( task_key, SubagentOutcome::cancelled(prepared.task_id), expected_pause, - ) - .await?; - self.record_completion(completion_origin.as_ref(), &result, 0) - .await; + ), + ) + .await?; + self.record_completion_under_fence( + turn_fence.as_ref(), + completion_origin.as_ref(), + &result, + 0, + ) + .await?; return Ok(result); } prepared.tools = restrict_tools_with_scope( @@ -367,32 +476,60 @@ impl<C: Send + 'static, H: Send + 'static> SubagentDriver<C, H> { prepared, task_key: task_key.clone(), cancellation: child_token, + turn_fence: turn_fence.clone(), }; + on_attempt(execution.prepared.run_context.run_id()); let mut timed_out = false; + let mut turn_lease_lost = false; let mut omitted_chars = 0usize; - let mut token; + let mut token = execution.cancellation.clone(); let executed = loop { + if turn_fence.as_ref().is_some_and(|fence| !fence.admitted()) { + turn_lease_lost = true; + break Err(SubagentError::Cancelled); + } token = execution.cancellation.clone(); let run = self.executor.execute(execution); - let result = match policy.timeout { - Some(limit) => match tokio::time::timeout(limit, run).await { - Ok(result) => result, - Err(_) => { - token.cancel(); - timed_out = true; - tracing::debug!( - "{LOG_PREFIX} timeout task_id={task_id} after_ms={}", - limit.as_millis() - ); - break Ok(SubagentOutcome::incomplete( - task_id.clone(), - SubagentIncomplete::new(format!("subagent timed out after {limit:?}")) + let timeout_token = token.clone(); + let run = async { + match policy.timeout { + Some(limit) => match tokio::time::timeout(limit, run).await { + Ok(result) => result, + Err(_) => { + timeout_token.cancel(); + timed_out = true; + tracing::debug!( + "{LOG_PREFIX} timeout task_id={task_id} after_ms={}", + limit.as_millis() + ); + Ok(SubagentOutcome::incomplete( + task_id.clone(), + SubagentIncomplete::new(format!( + "subagent timed out after {limit:?}" + )) .with_kind(IncompleteKind::Timeout), - )); + )) + } + }, + None => run.await, + } + }; + let result = if let Some(fence) = &turn_fence { + tokio::select! { + biased; + _ = fence.lost() => { + token.cancel(); + turn_lease_lost = true; + Err(SubagentError::Cancelled) } - }, - None => run.await, + result = run => result, + } + } else { + run.await }; + if turn_lease_lost { + break result; + } let Err(SubagentError::Transient { message, tools_ran }) = &result else { break result; }; @@ -406,24 +543,42 @@ impl<C: Send + 'static, H: Send + 'static> SubagentDriver<C, H> { } attempt += 1; tracing::debug!("{LOG_PREFIX} retry task_id={task_id} attempt={attempt}"); - tokio::select! { - biased; - _ = cancellation.cancelled() => break Err(SubagentError::Cancelled), - _ = policy.retry.sleep_backoff(attempt) => {} + if let Some(fence) = &turn_fence { + tokio::select! { + biased; + _ = fence.lost() => { + token.cancel(); + turn_lease_lost = true; + break Err(SubagentError::Cancelled); + } + _ = cancellation.cancelled() => break Err(SubagentError::Cancelled), + _ = policy.retry.sleep_backoff(attempt) => {} + } + } else { + tokio::select! { + biased; + _ = cancellation.cancelled() => break Err(SubagentError::Cancelled), + _ = policy.retry.sleep_backoff(attempt) => {} + } } let next_token = cancellation.child_token(); match attempts.next(attempt as u32, next_token.clone()) { Ok(next) => { + on_attempt(next.run_context.run_id()); execution = SubagentExecution { prepared: next, task_key: task_key.clone(), cancellation: next_token, + turn_fence: turn_fence.clone(), }; } Err(error) => break Err(error), } }; + if turn_lease_lost || turn_fence.as_ref().is_some_and(|fence| !fence.admitted()) { + return Err(SubagentError::Cancelled); + } let outcome = match executed { Ok(outcome) => { if outcome.task_id != task_id { @@ -439,7 +594,10 @@ impl<C: Send + 'static, H: Send + 'static> SubagentDriver<C, H> { outcome.cancelled_preserving() } else { let (outcome, omitted) = - apply_outcome_policies(outcome, &policy, &result_policy).await; + await_with_turn_fence(turn_fence.as_ref(), Some(&token), async { + Ok(apply_outcome_policies(outcome, &policy, &result_policy).await) + }) + .await?; omitted_chars = omitted; outcome } @@ -454,19 +612,45 @@ impl<C: Send + 'static, H: Send + 'static> SubagentDriver<C, H> { && let (Some(router), Some(origin)) = (&self.completions, &completion_origin) && let Some(record) = origin.record_for_error(&error) { - deliver(router, record).await; + await_with_turn_fence(turn_fence.as_ref(), Some(&token), async { + deliver(router, record).await; + Ok(()) + }) + .await?; } return Err(error); } }; - let result = self - .persist(task_key, outcome, expected_pause, &cancellation) - .await?; - self.record_completion(completion_origin.as_ref(), &result, omitted_chars) - .await; + let result = await_with_turn_fence( + turn_fence.as_ref(), + Some(&token), + self.persist(task_key, outcome, expected_pause, &cancellation), + ) + .await?; + self.record_completion_under_fence( + turn_fence.as_ref(), + completion_origin.as_ref(), + &result, + omitted_chars, + ) + .await?; Ok(result) } + async fn record_completion_under_fence( + &self, + fence: Option<&Arc<dyn DeliveryFence>>, + origin: Option<&CompletionOrigin>, + result: &SubagentRunResult, + omitted_chars: usize, + ) -> Result<(), SubagentError> { + await_with_turn_fence(fence, None, async { + self.record_completion(origin, result, omitted_chars).await; + Ok(()) + }) + .await + } + /// Records `result` when this invocation owns its host effects and the /// outcome maps to a completion (a pause does not). async fn record_completion( diff --git a/crates/tinyagents-orchestration/src/subagent/driver_completion_tests.rs b/crates/tinyagents-orchestration/src/subagent/driver_completion_tests.rs index 1d3bb32cd..2cc5967fc 100644 --- a/crates/tinyagents-orchestration/src/subagent/driver_completion_tests.rs +++ b/crates/tinyagents-orchestration/src/subagent/driver_completion_tests.rs @@ -619,8 +619,16 @@ async fn a_replayed_terminal_result_is_not_recorded_again() { .run(request("t1"), CancellationToken::new()) .await .unwrap(); - assert_eq!(router.pull("p", 10).unwrap().len(), 1); - router.mark_delivered(&["t1"]).unwrap(); + let claimed = router.pull("p", 10).unwrap(); + assert_eq!(claimed.len(), 1); + router + .mark_delivered( + &claimed + .iter() + .filter_map(|record| record.receipt()) + .collect::<Vec<_>>(), + ) + .unwrap(); let replay = driver .run(request("t1"), CancellationToken::new()) .await @@ -719,7 +727,10 @@ fn the_overflow_artifact_is_the_one_a_completion_names() { RunContext::new(RunConfig::new("c"), String::new()), ) .with_notify_mode(NotifyMode::Off); - let origin = completion::CompletionOrigin::new(&origin_task, &prepared).unwrap(); + let router = router(); + let origin = completion::CompletionOrigin::new(&origin_task, &prepared, &router) + .unwrap() + .unwrap(); let mut outcome = SubagentOutcome::completed("t1", "preview"); for id in ["pre-existing", "overflow"] { outcome.artifacts.push(ArtifactReference { @@ -731,3 +742,268 @@ fn the_overflow_artifact_is_the_one_a_completion_names() { assert_eq!(record.result.artifact.unwrap().id, "overflow"); assert_eq!(record.parent_key, "p"); } + +struct RevocableTurnFence { + state: tokio::sync::watch::Sender<bool>, + receiver: tokio::sync::watch::Receiver<bool>, +} + +impl RevocableTurnFence { + fn new() -> Self { + let (state, receiver) = tokio::sync::watch::channel(true); + Self { state, receiver } + } + + fn revoke(&self) { + self.state.send_replace(false); + } +} + +impl tinyagents_tasks::DeliveryFence for RevocableTurnFence { + fn admitted(&self) -> bool { + *self.receiver.borrow() + } + + fn lost(&self) -> std::pin::Pin<Box<dyn std::future::Future<Output = ()> + Send + '_>> { + let mut receiver = self.receiver.clone(); + Box::pin(async move { + loop { + if !*receiver.borrow() || receiver.changed().await.is_err() { + return; + } + } + }) + } +} + +struct WaitingExecutor { + started: Arc<tokio::sync::Notify>, + observed_cancellation: Arc<Mutex<Option<CancellationToken>>>, +} + +struct BarrierExecutor { + started: Arc<tokio::sync::Notify>, + release: Arc<tokio::sync::Notify>, +} + +#[async_trait] +impl SubagentExecutor<String> for BarrierExecutor { + async fn execute( + &self, + execution: SubagentExecution<String>, + ) -> Result<SubagentOutcome, SubagentError> { + self.started.notify_one(); + self.release.notified().await; + Ok(SubagentOutcome::completed( + execution.prepared.task_id, + "finished after parent reopened", + )) + } +} + +#[async_trait] +impl SubagentExecutor<String> for WaitingExecutor { + async fn execute( + &self, + execution: SubagentExecution<String>, + ) -> Result<SubagentOutcome, SubagentError> { + *self.observed_cancellation.lock().unwrap() = Some(execution.cancellation); + self.started.notify_one(); + std::future::pending().await + } +} + +#[tokio::test] +async fn generic_child_driver_cancels_dispatch_and_skips_persistence_when_turn_fence_is_lost() { + let fence = Arc::new(RevocableTurnFence::new()); + let started = Arc::new(tokio::sync::Notify::new()); + let observed_cancellation = Arc::new(Mutex::new(None)); + let persistence = Arc::new(Memory::default()); + let driver = SubagentDriver::new(SubagentCapabilities { + planner: Some(Arc::new(Planner { + mode: None, + parent: None, + })), + executor: Some(Arc::new(WaitingExecutor { + started: started.clone(), + observed_cancellation: observed_cancellation.clone(), + })), + persistence: Some(persistence.clone()), + }) + .unwrap(); + let request = request_in_scope("parent-a", "root-a", "thread-a").with_turn_fence(fence.clone()); + let turn = tokio::spawn(async move { driver.run(request, CancellationToken::new()).await }); + started.notified().await; + + fence.revoke(); + assert!(matches!(turn.await.unwrap(), Err(SubagentError::Cancelled))); + assert!( + observed_cancellation + .lock() + .unwrap() + .as_ref() + .is_some_and(CancellationToken::is_cancelled), + "loss cancels the executor's in-flight turn token" + ); + assert!( + persistence.terminals.lock().unwrap().is_empty(), + "a stale turn must not persist a terminal outcome after losing its fence" + ); +} + +struct HeldArtifactStore { + started: tokio::sync::Notify, + release: tokio::sync::Notify, +} + +#[async_trait] +impl ArtifactStore for HeldArtifactStore { + async fn store(&self, _: &str, _: &str) -> Result<ArtifactReference, String> { + self.started.notify_one(); + self.release.notified().await; + Ok(ArtifactReference { + id: "stored-artifact".into(), + ..ArtifactReference::default() + }) + } +} + +struct ArtifactPlanner(Arc<HeldArtifactStore>); + +#[async_trait] +impl SubagentPlanner<String> for ArtifactPlanner { + async fn prepare( + &self, + request: SubagentRequest<String>, + ) -> Result<PreparedSubagent<String>, SubagentError> { + let parts = request.into_parts(); + Ok(PreparedSubagent::new( + parts.task_key.task_id, + "worker", + vec![Message::user(parts.input)], + ToolSnapshot::new(vec![]).unwrap(), + parts.run_context, + ) + .with_notify_mode(NotifyMode::Followup) + .with_completion_parent("parent-a") + .with_result_policy( + ResultPolicy::new() + .with_max_bytes(1) + .with_overflow(ResultOverflow::Artifact) + .with_artifact_store(self.0.clone()), + )) + } +} + +#[tokio::test] +async fn losing_turn_fence_during_artifact_storage_skips_terminal_and_completion_writes() { + let fence = Arc::new(RevocableTurnFence::new()); + let artifacts = Arc::new(HeldArtifactStore { + started: tokio::sync::Notify::new(), + release: tokio::sync::Notify::new(), + }); + let persistence = Arc::new(Memory::default()); + let router = router(); + let driver = SubagentDriver::new(SubagentCapabilities { + planner: Some(Arc::new(ArtifactPlanner(artifacts.clone()))), + executor: Some(Arc::new(Executor(Arc::new(|task| { + Ok(SubagentOutcome::completed(task, "oversized child output")) + })))), + persistence: Some(persistence.clone()), + }) + .unwrap() + .with_completion_router(router.clone()); + let request = request_in_scope("parent-a", "root-a", "thread-a").with_turn_fence(fence.clone()); + let turn = tokio::spawn(async move { driver.run(request, CancellationToken::new()).await }); + artifacts.started.notified().await; + + fence.revoke(); + artifacts.release.notify_one(); + + assert!(matches!(turn.await.unwrap(), Err(SubagentError::Cancelled))); + assert!( + persistence.terminals.lock().unwrap().is_empty(), + "a result policy await after executor completion must not reopen terminal persistence" + ); + assert!( + router.pending_for("parent-a").is_empty(), + "lease loss must not dispatch a completion" + ); +} + +#[tokio::test] +async fn a_running_child_completion_keeps_its_spawned_parent_generation() { + let store = Arc::new(InMemoryCompletionStore::new()); + let router_a = Arc::new(CompletionRouter::new(store.clone())); + let router_b = Arc::new(CompletionRouter::new(Arc::new(store.as_ref().clone()))); + let spawned_generation = router_a + .admit_child_generation("parent-a") + .unwrap() + .unwrap(); + let started = Arc::new(tokio::sync::Notify::new()); + let release = Arc::new(tokio::sync::Notify::new()); + let driver = SubagentDriver::new(SubagentCapabilities { + planner: Some(Arc::new(Planner { + mode: Some(NotifyMode::Followup), + parent: Some("parent-a"), + })), + executor: Some(Arc::new(BarrierExecutor { + started: started.clone(), + release: release.clone(), + })), + persistence: Some(Arc::new(Memory::default())), + }) + .unwrap() + .with_completion_router(router_a.clone()); + let run = tokio::spawn(async move { + driver + .run( + request_in_scope("parent-a-run", "root-a", "thread-a"), + CancellationToken::new(), + ) + .await + }); + + started.notified().await; + router_b.cancel_parent("parent-a").unwrap(); + router_b.resume_parent("parent-a"); + assert!( + router_b + .admit_child_generation("parent-a") + .unwrap() + .unwrap() + > spawned_generation + ); + release.notify_one(); + + assert!(run.await.unwrap().is_ok()); + let fresh_driver = SubagentDriver::new(SubagentCapabilities { + planner: Some(Arc::new(Planner { + mode: Some(NotifyMode::Followup), + parent: Some("parent-a"), + })), + executor: Some(Arc::new(Executor(Arc::new(|task| { + Ok(SubagentOutcome::completed(task, "fresh-generation child")) + })))), + persistence: Some(Arc::new(Memory::default())), + }) + .unwrap() + .with_completion_router(router_b.clone()); + assert!( + fresh_driver + .run( + request_in_scope("parent-b-run", "root-b", "thread-a"), + CancellationToken::new(), + ) + .await + .is_ok() + ); + + let pending = router_b.pending_for("parent-a"); + assert!( + pending.len() == 1, + "the old generation remains suppressed while a fresh child is deliverable" + ); + assert_eq!(pending[0].result.text, "fresh-generation child"); + assert!(pending[0].parent_epoch > spawned_generation); +} diff --git a/crates/tinyagents-orchestration/src/subagent/driver_policy_tests.rs b/crates/tinyagents-orchestration/src/subagent/driver_policy_tests.rs index 354ccc912..71347d992 100644 --- a/crates/tinyagents-orchestration/src/subagent/driver_policy_tests.rs +++ b/crates/tinyagents-orchestration/src/subagent/driver_policy_tests.rs @@ -42,6 +42,7 @@ type Behaviour = Arc< struct Executor { attempts: AtomicU32, + observed_attempts: Mutex<Option<Arc<Mutex<Vec<String>>>>>, seen_tools: Mutex<Vec<Vec<String>>>, seen_caps: Mutex<Vec<(Option<usize>, Option<usize>)>>, child_tokens: Mutex<Vec<CancellationToken>>, @@ -53,6 +54,7 @@ impl Executor { fn new(behaviour: Behaviour) -> Arc<Self> { Arc::new(Self { attempts: AtomicU32::new(0), + observed_attempts: Mutex::new(None), seen_tools: Mutex::default(), seen_caps: Mutex::default(), child_tokens: Mutex::default(), @@ -75,6 +77,17 @@ impl SubagentExecutor<String> for Executor { &self, execution: SubagentExecution<String>, ) -> Result<SubagentOutcome, SubagentError> { + if let Some(observed) = self.observed_attempts.lock().unwrap().as_ref() { + let run_id = execution.prepared.run_context.run_id().as_str(); + assert!( + observed + .lock() + .unwrap() + .iter() + .any(|observed| observed == run_id), + "attempt {run_id} reached the executor before the observer recorded it" + ); + } let attempt = self.attempts.fetch_add(1, Ordering::SeqCst); self.seen_tools.lock().unwrap().push( execution @@ -244,6 +257,33 @@ async fn retries_a_transient_failure_that_ran_no_tools() { assert_eq!(executor.attempts.load(Ordering::SeqCst), 3); } +#[tokio::test] +async fn attempt_observer_receives_each_actual_run_id_before_execution() { + let executor = Executor::new(flaky(1, false)); + let observed = Arc::new(Mutex::new(Vec::new())); + *executor.observed_attempts.lock().unwrap() = Some(observed.clone()); + let plan: Planned = Arc::new(|p| with_factory(p.with_policy(retry_policy(2)))); + let observed_by_driver = observed.clone(); + driver(plan, executor) + .run_with_attempt_observer( + request("t-observed-attempts"), + CancellationToken::new(), + move |run_id| { + observed_by_driver + .lock() + .unwrap() + .push(run_id.as_str().to_owned()); + }, + ) + .await + .unwrap(); + + assert_eq!( + *observed.lock().unwrap(), + vec!["run-t-observed-attempts", "retry-1"] + ); +} + #[tokio::test] async fn does_not_retry_after_tools_ran_unless_allowed() { let executor = Executor::new(flaky(1, true)); diff --git a/crates/tinyagents-orchestration/src/subagent/invocation/delegation_tool.rs b/crates/tinyagents-orchestration/src/subagent/invocation/delegation_tool.rs deleted file mode 100644 index ef096ff3b..000000000 --- a/crates/tinyagents-orchestration/src/subagent/invocation/delegation_tool.rs +++ /dev/null @@ -1,441 +0,0 @@ -//! Model-facing multi-target delegation adapter. - -use std::marker::PhantomData; -use std::sync::Arc; - -use async_trait::async_trait; -use serde::{Deserialize, Serialize}; -use serde_json::{Value, json}; -use tinyagents_graph::delegation::{ - DelegationConfig, DelegationOutcome, DelegationStage, DelegationStageOutput, DelegationState, - run_delegation, run_or_resume_delegation, -}; -use tinyagents_harness::{context::RunContext, ids::CallId, tool::ToolDispatch}; -use tinytools::{ - PermissionLevel, Tool, ToolCategory, ToolExposure, ToolResult, ToolScope, ToolTimeout, -}; - -use crate::subagent::SUBAGENT_MODE_FIELD; - -/// A task envelope passed to the selected target agent. -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -pub struct DelegationEnvelope { - /// The requested work. - pub prompt: String, - /// Optional objective statement. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub objective: Option<String>, - /// Relevant evidence. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub evidence: Option<String>, - /// Constraints the target must follow. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub constraints: Option<String>, - /// Claims the target must not assume. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub must_not_assume: Option<String>, - /// Expected response shape. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub expected_output: Option<String>, - /// Citation expectations. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub citation_requirement: Option<String>, - /// Optional model override, interpreted by the host callback. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub model: Option<String>, - /// Whether the caller waits for completion; maps to [`SUBAGENT_MODE_FIELD`]. - #[serde(default)] - pub blocking: bool, -} - -/// Parsed target and task passed to a host request factory. -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct DelegationRequest { - /// The selected agent id. - pub agent: String, - /// Structured task instructions. - pub envelope: DelegationEnvelope, - /// Driver mode selected from `blocking`. - pub mode: String, -} - -/// Driver-backed multi-target tool whose schema and execution remain host owned. -pub struct MultiTargetDelegationTool<State, Ctx, Worker, Fut> -where - Ctx: Send + Sync + 'static, - Worker: Fn(DelegationRequest) -> Fut + Send + Sync + 'static, - Fut: std::future::Future<Output = Result<DelegationStageOutput, String>> + Send + 'static, -{ - name: String, - description: String, - agents: Vec<(String, String)>, - worker: Arc<Worker>, - _context: PhantomData<fn(&State, &Ctx)>, - exposure: ToolExposure, -} - -impl<State, Ctx, Worker, Fut> MultiTargetDelegationTool<State, Ctx, Worker, Fut> -where - Ctx: Send + Sync + 'static, - Worker: Fn(DelegationRequest) -> Fut + Send + Sync + 'static, - Fut: std::future::Future<Output = Result<DelegationStageOutput, String>> + Send + 'static, -{ - /// Creates a multi-target delegation tool. `agents` entries are `(id, description)`. - pub fn new( - name: impl Into<String>, - description: impl Into<String>, - agents: Vec<(String, String)>, - worker: Worker, - ) -> Self { - Self { - name: name.into(), - description: description.into(), - agents, - worker: Arc::new(worker), - _context: PhantomData, - exposure: ToolExposure::Direct, - } - } - - /// Sets whether the tool is directly advertised, deferred, or host-only. - pub fn with_exposure(mut self, exposure: ToolExposure) -> Self { - self.exposure = exposure; - self - } -} - -#[async_trait] -impl<State, Ctx, Worker, Fut> ToolDispatch<State, Ctx> - for MultiTargetDelegationTool<State, Ctx, Worker, Fut> -where - State: Send + Sync + 'static, - Ctx: Send + Sync + 'static, - Worker: Fn(DelegationRequest) -> Fut + Send + Sync + 'static, - Fut: std::future::Future<Output = Result<DelegationStageOutput, String>> + Send + 'static, -{ - fn tool(&self) -> Arc<dyn Tool> { - Arc::new(DelegationDeclaration { - name: self.name.clone(), - description: self.description.clone(), - parameters: json!({"type":"object","required":["agent","prompt"],"properties":{ - "agent":{"type":"string","enum":self.agents.iter().map(|(id,_)|id).collect::<Vec<_>>()}, - "prompt":{"type":"string"},"objective":{"type":"string"},"evidence":{"type":"string"}, - "constraints":{"type":"string"},"must_not_assume":{"type":"string"}, - "expected_output":{"type":"string"},"citation_requirement":{"type":"string"}, - "model":{"type":"string"},"blocking":{"type":"boolean","description":"Wait for the target result when true."} - }}), - metadata: None, - exposure: self.exposure, - }) - } - - async fn execute( - &self, - _state: &State, - _call_id: CallId, - args: Value, - _options: tinytools::ToolCallOptions, - _parent: &RunContext<Ctx>, - ) -> anyhow::Result<ToolResult> { - let agent = args - .get("agent") - .and_then(Value::as_str) - .ok_or_else(|| anyhow::anyhow!("agent must be a string"))?; - if !self.agents.iter().any(|(id, _)| id == agent) { - anyhow::bail!("agent is not an available delegation target"); - } - let envelope: DelegationEnvelope = serde_json::from_value(args.clone())?; - let request = DelegationRequest { - agent: agent.into(), - mode: if envelope.blocking { - "inline" - } else { - "background" - } - .into(), - envelope, - }; - let result = (self.worker)(request.clone()) - .await - .map_err(|error| anyhow::anyhow!("delegation failed: {error}"))?; - Ok(ToolResult::json( - json!({"agent":agent,"mode":request.mode,"output":result.text}), - )) - } -} - -struct DelegationDeclaration { - name: String, - description: String, - parameters: Value, - metadata: Option<DelegationToolMetadata>, - exposure: ToolExposure, -} -#[async_trait] -impl Tool for DelegationDeclaration { - fn name(&self) -> &str { - &self.name - } - fn description(&self) -> &str { - &self.description - } - fn parameters_schema(&self) -> Value { - self.parameters.clone() - } - fn exposure(&self) -> ToolExposure { - self.exposure - } - fn permission_level(&self) -> PermissionLevel { - self.metadata - .map_or(PermissionLevel::ReadOnly, |metadata| metadata.permission) - } - fn scope(&self) -> ToolScope { - self.metadata - .map_or(ToolScope::All, |metadata| metadata.scope) - } - fn category(&self) -> ToolCategory { - self.metadata - .map_or(ToolCategory::System, |metadata| metadata.category) - } - fn timeout_policy(&self, _args: &Value) -> ToolTimeout { - self.metadata - .map_or(ToolTimeout::Inherit, |metadata| metadata.timeout) - } - async fn execute(&self, _args: Value) -> anyhow::Result<ToolResult> { - anyhow::bail!("delegation requires typed-parent dispatch") - } -} - -/// Parses a reviewer approval exactly when the whole normalized response is `APPROVE`. -pub fn parse_approval(text: &str) -> bool { - text.trim().eq_ignore_ascii_case("APPROVE") -} - -/// Tool argument field used to map blocking mode into the subagent invocation mode. -pub const DELEGATION_MODE_FIELD: &str = SUBAGENT_MODE_FIELD; - -/// Generic prompt guidance for a graph delegation planning worker. -pub const DELEGATION_PLAN_PROMPT: &str = "Create a concise, verifiable plan for the requested task. Preserve its constraints and expected output."; -/// Generic prompt guidance for a graph delegation execution worker. -pub const DELEGATION_EXECUTE_PROMPT: &str = "Execute the current plan. Report concrete changes and evidence, and do not claim work that was not performed."; -/// Generic prompt guidance for a graph delegation review worker. -pub const DELEGATION_REVIEW_PROMPT: &str = "Review the latest execution against the request and plan. Reply with exactly APPROVE only when it is complete; otherwise explain the needed revision."; - -/// Creates a review-stage result using the canonical standalone APPROVE parser. -pub fn review_stage_output(text: impl Into<String>) -> DelegationStageOutput { - let text = text.into(); - let approved = parse_approval(&text); - DelegationStageOutput { - text, - approved, - prompt: None, - } -} - -type ConfigFactory<State, Ctx> = - dyn Fn(&State, &RunContext<Ctx>, &Value) -> DelegationConfig + Send + Sync; -type ResultProjector = dyn Fn(Result<DelegationOutcome, String>) -> ToolResult + Send + Sync; - -#[derive(Debug, Clone, Copy)] -struct DelegationToolMetadata { - permission: PermissionLevel, - scope: ToolScope, - category: ToolCategory, - timeout: ToolTimeout, -} - -/// Determines how review-stage approval is interpreted by [`DelegationTool`]. -#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] -pub enum ReviewApprovalPolicy { - /// Require a standalone `APPROVE` response (the generic default). - #[default] - StandaloneApprove, - /// Trust the approval boolean returned by the host worker. - WorkerOutput, -} - -/// Single-target graph delegation tool with an injected worker for each stage. -pub struct DelegationTool<State, Ctx, Worker, Fut> -where - Ctx: Send + Sync + 'static, - Worker: Fn(Value, DelegationStage, DelegationState) -> Fut + Send + Sync + 'static, - Fut: std::future::Future<Output = Result<DelegationStageOutput, String>> + Send + 'static, -{ - config: Arc<ConfigFactory<State, Ctx>>, - worker: Arc<Worker>, - name: String, - description: String, - parameters: Value, - resumable: bool, - review_approval_policy: ReviewApprovalPolicy, - result_projector: Option<Arc<ResultProjector>>, - metadata: Option<DelegationToolMetadata>, -} - -impl<State, Ctx, Worker, Fut> DelegationTool<State, Ctx, Worker, Fut> -where - Ctx: Send + Sync + 'static, - Worker: Fn(Value, DelegationStage, DelegationState) -> Fut + Send + Sync + 'static, - Fut: std::future::Future<Output = Result<DelegationStageOutput, String>> + Send + 'static, -{ - /// Constructs the tool from a host config factory and stage worker. - pub fn new( - config: impl Fn(&State, &RunContext<Ctx>, &Value) -> DelegationConfig + Send + Sync + 'static, - worker: Worker, - ) -> Self { - Self { - config: Arc::new(config), - worker: Arc::new(worker), - name: "delegate_graph".into(), - description: "Plan, execute, and review delegated work.".into(), - parameters: json!({"type":"object","properties":{"task":{"type":"string"}},"required":["task"]}), - resumable: false, - review_approval_policy: ReviewApprovalPolicy::StandaloneApprove, - result_projector: None, - metadata: None, - } - } - - /// Overrides the model-facing name while keeping graph execution generic. - pub fn with_tool_name(mut self, name: impl Into<String>) -> Self { - self.name = name.into(); - self - } - - /// Overrides the model-facing description for a host-owned alias. - pub fn with_description(mut self, description: impl Into<String>) -> Self { - self.description = description.into(); - self - } - - /// Supplies a host-owned schema, including dynamic target enums. - pub fn with_parameters(mut self, parameters: Value) -> Self { - self.parameters = parameters; - self - } - - /// Enables durable classify-and-resume execution for checkpointed threads. - pub fn with_resumable_execution(mut self) -> Self { - self.resumable = true; - self - } - - /// Selects how the tool interprets the review worker's approval result. - pub fn with_review_approval_policy(mut self, policy: ReviewApprovalPolicy) -> Self { - self.review_approval_policy = policy; - self - } - - /// Projects the complete graph outcome or raw graph error into a tool result. - pub fn with_result_projection( - mut self, - project: impl Fn(Result<DelegationOutcome, String>) -> ToolResult + Send + Sync + 'static, - ) -> Self { - self.result_projector = Some(Arc::new(project)); - self - } - - /// Preserves host authorization and execution metadata on the declaration. - pub fn with_metadata( - mut self, - permission: PermissionLevel, - scope: ToolScope, - category: ToolCategory, - timeout: ToolTimeout, - ) -> Self { - self.metadata = Some(DelegationToolMetadata { - permission, - scope, - category, - timeout, - }); - self - } -} - -#[async_trait] -impl<State, Ctx, Worker, Fut> ToolDispatch<State, Ctx> for DelegationTool<State, Ctx, Worker, Fut> -where - State: Send + Sync + 'static, - Ctx: Send + Sync + 'static, - Worker: Fn(Value, DelegationStage, DelegationState) -> Fut + Send + Sync + 'static, - Fut: std::future::Future<Output = Result<DelegationStageOutput, String>> + Send + 'static, -{ - fn tool(&self) -> Arc<dyn Tool> { - Arc::new(DelegationDeclaration { - name: self.name.clone(), - description: self.description.clone(), - parameters: self.parameters.clone(), - metadata: self.metadata, - exposure: ToolExposure::Direct, - }) - } - async fn execute( - &self, - state: &State, - _call_id: CallId, - args: Value, - _options: tinytools::ToolCallOptions, - parent: &RunContext<Ctx>, - ) -> anyhow::Result<ToolResult> { - if args.get("task").and_then(Value::as_str).is_none() { - anyhow::bail!("task must be a string"); - } - let config = (self.config)(state, parent, &args); - let worker = Arc::clone(&self.worker); - let review_approval_policy = self.review_approval_policy; - let stage_runner = move |stage, graph| { - let worker = Arc::clone(&worker); - let mut stage_args = args.clone(); - stage_args["stage_prompt"] = Value::String(stage_prompt(stage).to_owned()); - async move { - let mut output = worker(stage_args, stage, graph).await?; - if stage == DelegationStage::Review - && review_approval_policy == ReviewApprovalPolicy::StandaloneApprove - { - output.approved = parse_approval(&output.text); - } - Ok(output) - } - }; - let outcome = if self.resumable { - run_or_resume_delegation(config, stage_runner).await - } else { - run_delegation(config, stage_runner) - .await - .map(|state| DelegationOutcome { - state, - pending: None, - }) - }; - - if let Some(project) = &self.result_projector { - return Ok(project(outcome)); - } - let outcome = outcome.map_err(|error| anyhow::anyhow!("delegation failed: {error}"))?; - if self.resumable { - return Ok(ToolResult::json(json!({ - "state": outcome.state, - "pending": outcome.pending.map(|pending| json!({ - "interrupt_id": pending.interrupt_id, - "node": pending.node, - "payload": pending.payload, - "thread_id": pending.thread_id, - })), - }))); - } - Ok(ToolResult::json(json!({ - "output": outcome.state.final_output, - "approved": outcome.state.approved, - "cancelled": outcome.state.cancelled, - }))) - } -} - -fn stage_prompt(stage: DelegationStage) -> &'static str { - match stage { - DelegationStage::Plan => DELEGATION_PLAN_PROMPT, - DelegationStage::Execute => DELEGATION_EXECUTE_PROMPT, - DelegationStage::Review => DELEGATION_REVIEW_PROMPT, - } -} diff --git a/crates/tinyagents-orchestration/src/subagent/invocation/delegation_tool_tests.rs b/crates/tinyagents-orchestration/src/subagent/invocation/delegation_tool_tests.rs deleted file mode 100644 index 865296645..000000000 --- a/crates/tinyagents-orchestration/src/subagent/invocation/delegation_tool_tests.rs +++ /dev/null @@ -1,490 +0,0 @@ -use super::*; -use tinyagents_graph::delegation::{DelegationConfig, DelegationStage, DelegationStageOutput}; -use tinytools::{ToolCategory, ToolExposure, ToolResult, ToolScope, ToolTimeout}; - -#[test] -fn approval_parser_requires_a_standalone_approve_token() { - assert!(parse_approval(" APPROVE\n")); - assert!(!parse_approval("APPROVE, because it looks good")); - assert!(!parse_approval("not approved")); -} - -#[test] -fn multi_target_schema_exposes_only_configured_agents_and_blocking_envelope() { - let tool = MultiTargetDelegationTool::<(), (), _, _>::new( - "delegate_to", - "Delegate to a specialist", - vec![ - ("coder".into(), "Writes code".into()), - ("reviewer".into(), "Reviews".into()), - ], - |_| async { Ok(DelegationStageOutput::done("ok")) }, - ); - let declaration = tool.with_exposure(ToolExposure::Hidden).tool(); - let schema = declaration.parameters_schema(); - assert_eq!( - schema["properties"]["agent"]["enum"], - json!(["coder", "reviewer"]) - ); - assert_eq!(schema["properties"]["blocking"]["type"], "boolean"); - assert_eq!(DELEGATION_MODE_FIELD, SUBAGENT_MODE_FIELD); - assert_eq!(declaration.exposure(), ToolExposure::Hidden); -} - -#[test] -fn graph_alias_preserves_host_execution_metadata() { - use tinytools::PermissionLevel; - - let declaration = DelegationTool::<(), (), _, _>::new( - |_, _, _| DelegationConfig::default(), - |_, _, _| async { Ok(DelegationStageOutput::done("unused")) }, - ) - .with_metadata( - PermissionLevel::Execute, - ToolScope::AgentOnly, - ToolCategory::Workflow, - ToolTimeout::Unbounded, - ) - .tool(); - - assert_eq!(declaration.permission_level(), PermissionLevel::Execute); - assert_eq!(declaration.scope(), ToolScope::AgentOnly); - assert_eq!(declaration.category(), ToolCategory::Workflow); - assert_eq!( - declaration.timeout_policy(&Value::Null), - ToolTimeout::Unbounded - ); -} - -#[test] -fn delegation_envelope_keeps_structured_instructions_and_mode() { - let envelope: DelegationEnvelope = serde_json::from_value(json!({ - "prompt": "Inspect the patch", "objective": "Find regressions", - "must_not_assume": "Tests passed", "blocking": true - })) - .unwrap(); - assert!(envelope.blocking); - assert_eq!(envelope.objective.as_deref(), Some("Find regressions")); - assert_eq!(envelope.must_not_assume.as_deref(), Some("Tests passed")); -} - -#[tokio::test] -async fn multi_target_blocking_maps_to_subagent_mode_and_validates_target() { - use tinyagents_harness::{ - context::{RunConfig, RunContext}, - ids::new_call_id, - tool::ToolDispatch, - }; - let tool = MultiTargetDelegationTool::<(), (), _, _>::new( - "delegate_to", - "Delegate", - vec![("coder".into(), "Writes code".into())], - |request| async move { Ok(DelegationStageOutput::done(request.mode)) }, - ); - let parent = RunContext::new(RunConfig::new("parent"), ()); - let inline = ToolDispatch::<(), ()>::execute( - &tool, - &(), - new_call_id(), - json!({"agent":"coder","prompt":"fix","blocking":true}), - tinytools::ToolCallOptions::default(), - &parent, - ) - .await - .unwrap(); - let value: serde_json::Value = serde_json::from_str(&inline.output()).unwrap(); - assert_eq!(value["mode"], "inline"); - let invalid = ToolDispatch::<(), ()>::execute( - &tool, - &(), - new_call_id(), - json!({"agent":"unlisted","prompt":"fix"}), - tinytools::ToolCallOptions::default(), - &parent, - ) - .await; - assert!(invalid.is_err()); -} - -#[tokio::test] -async fn graph_delegation_runs_stage_worker_with_original_task_and_approval_parser() { - use tinyagents_harness::{ - context::{RunConfig, RunContext}, - ids::new_call_id, - tool::ToolDispatch, - }; - let seen = std::sync::Arc::new(std::sync::Mutex::new(Vec::new())); - let worker_seen = seen.clone(); - let tool = DelegationTool::<(), (), _, _>::new( - |_, _, _| DelegationConfig::default(), - move |args, stage, _| { - let seen = worker_seen.clone(); - async move { - seen.lock().unwrap().push(( - args["task"].as_str().unwrap().to_owned(), - stage, - args["stage_prompt"].as_str().unwrap().to_owned(), - )); - Ok(match stage { - DelegationStage::Plan => DelegationStageOutput::done("plan"), - DelegationStage::Execute => DelegationStageOutput::done("result"), - DelegationStage::Review => DelegationStageOutput::done("APPROVE, looks good"), - }) - } - }, - ); - let parent = RunContext::new(RunConfig::new("parent"), ()); - let result = ToolDispatch::<(), ()>::execute( - &tool, - &(), - new_call_id(), - json!({"task":"ship a safe fix"}), - tinytools::ToolCallOptions::default(), - &parent, - ) - .await - .unwrap(); - assert!(!result.is_error); - let seen = seen.lock().unwrap(); - assert!(seen.len() >= 3); - assert!( - seen.len() > 3, - "non-standalone APPROVE must trigger another review cycle" - ); - assert!(seen.iter().all(|(task, _, _)| task == "ship a safe fix")); - assert_eq!(seen[0].1, DelegationStage::Plan); - assert_eq!(seen[0].2, DELEGATION_PLAN_PROMPT); - assert_eq!(seen[1].2, DELEGATION_EXECUTE_PROMPT); - assert_eq!(seen[2].1, DelegationStage::Review); - assert_eq!(seen[2].2, DELEGATION_REVIEW_PROMPT); - assert!(!parse_approval("APPROVE, looks good")); -} - -#[tokio::test] -async fn graph_alias_keeps_host_declaration_and_routes_arguments_to_stage_worker() { - use tinyagents_harness::{ - context::{RunConfig, RunContext}, - ids::new_call_id, - tool::ToolDispatch, - }; - let seen = std::sync::Arc::new(std::sync::Mutex::new(Vec::new())); - let worker_seen = seen.clone(); - let schema = json!({ - "type":"object", - "required":["agent_id","task"], - "properties":{ - "agent_id":{"type":"string","enum":["planner","reviewer"]}, - "task":{"type":"string"}, - "revisions":{"type":"integer","minimum":0,"maximum":5,"default":2} - } - }); - let tool = DelegationTool::<(), (), _, _>::new( - |_, _, _| DelegationConfig::default(), - move |args, stage, _| { - let seen = worker_seen.clone(); - async move { - seen.lock().unwrap().push(( - args["agent_id"].as_str().unwrap().to_owned(), - args["revisions"].as_u64().unwrap(), - stage, - )); - Ok(DelegationStageOutput::done( - if stage == DelegationStage::Review { - "APPROVE" - } else { - "ok" - }, - )) - } - }, - ) - .with_tool_name("delegate_graph") - .with_description("Plan and review delegated work.") - .with_parameters(schema.clone()); - let declaration = tool.tool(); - assert_eq!(declaration.name(), "delegate_graph"); - assert_eq!(declaration.description(), "Plan and review delegated work."); - assert_eq!(declaration.parameters_schema(), schema); - - let parent = RunContext::new(RunConfig::new("parent"), ()); - ToolDispatch::<(), ()>::execute( - &tool, - &(), - new_call_id(), - json!({"agent_id":"planner","task":"ship the fix","revisions":3}), - tinytools::ToolCallOptions::default(), - &parent, - ) - .await - .expect("custom alias executes the typed graph worker"); - let seen = seen.lock().unwrap(); - assert!(seen.len() >= 3); - assert!( - seen.iter() - .all(|(agent, revisions, _)| agent == "planner" && *revisions == 3) - ); -} - -#[tokio::test] -async fn resumable_delegation_restarts_at_the_pending_stage() { - use std::sync::atomic::{AtomicBool, Ordering}; - - use tinyagents_graph::{DelegationState, checkpoint::InMemoryCheckpointer}; - use tinyagents_harness::{ - context::{RunConfig, RunContext}, - ids::new_call_id, - tool::ToolDispatch, - }; - - let checkpointer = std::sync::Arc::new(InMemoryCheckpointer::<DelegationState>::new()); - let fail_execute_once = std::sync::Arc::new(AtomicBool::new(true)); - let calls = std::sync::Arc::new(std::sync::Mutex::new(Vec::new())); - let make_tool = || { - let cp = checkpointer.clone(); - let fail = fail_execute_once.clone(); - let calls = calls.clone(); - DelegationTool::<(), (), _, _>::new( - move |_, _, _| DelegationConfig { - checkpointer: Some(cp.clone()), - thread_id: Some("durable-delegation".into()), - ..DelegationConfig::default() - }, - move |_, stage, _| { - let fail = fail.clone(); - let calls = calls.clone(); - async move { - calls.lock().unwrap().push(stage); - if stage == DelegationStage::Execute && fail.swap(false, Ordering::SeqCst) { - return Err("simulated process interruption".into()); - } - Ok(DelegationStageOutput::done( - if stage == DelegationStage::Review { - "APPROVE" - } else { - "completed" - }, - )) - } - }, - ) - .with_resumable_execution() - }; - let parent = RunContext::new(RunConfig::new("parent"), ()); - let first = make_tool(); - assert!( - ToolDispatch::<(), ()>::execute( - &first, - &(), - new_call_id(), - json!({"task":"resume me"}), - tinytools::ToolCallOptions::default(), - &parent, - ) - .await - .is_err() - ); - - // A new dispatcher models a host process starting from the persisted graph. - let restarted = make_tool(); - assert!( - ToolDispatch::<(), ()>::execute( - &restarted, - &(), - new_call_id(), - json!({"task":"resume me"}), - tinytools::ToolCallOptions::default(), - &parent, - ) - .await - .is_ok() - ); - - let before = { - let calls = calls.lock().unwrap(); - assert_eq!( - calls - .iter() - .filter(|&&stage| stage == DelegationStage::Plan) - .count(), - 1 - ); - assert_eq!( - calls - .iter() - .filter(|&&stage| stage == DelegationStage::Execute) - .count(), - 2 - ); - calls.len() - }; - let terminal = make_tool(); - ToolDispatch::<(), ()>::execute( - &terminal, - &(), - new_call_id(), - json!({"task":"resume me"}), - tinytools::ToolCallOptions::default(), - &parent, - ) - .await - .expect("a terminal checkpoint is returned without rerunning workers"); - assert_eq!(calls.lock().unwrap().len(), before); -} - -#[tokio::test] -async fn resumable_delegation_projects_pending_approval_and_full_state() { - use tinyagents_graph::{ - DelegationState, checkpoint::InMemoryCheckpointer, stream::CollectingSink, - }; - use tinyagents_harness::{ - context::{RunConfig, RunContext}, - ids::new_call_id, - tool::ToolDispatch, - }; - - let checkpointer = std::sync::Arc::new(InMemoryCheckpointer::<DelegationState>::new()); - let event_sink = std::sync::Arc::new(CollectingSink::new()); - let config_checkpointer = checkpointer.clone(); - let config_event_sink = event_sink.clone(); - let tool = DelegationTool::<(), (), _, _>::new( - move |_, _, _| DelegationConfig { - checkpointer: Some(config_checkpointer.clone()), - thread_id: Some("pending-review".into()), - require_review_approval: true, - event_sink: Some(config_event_sink.clone()), - ..DelegationConfig::default() - }, - |_, stage, _| async move { - Ok(if stage == DelegationStage::Review { - DelegationStageOutput { - text: "APPROVE after the final checks".into(), - approved: true, - prompt: None, - } - } else { - DelegationStageOutput::done("completed") - }) - }, - ) - .with_resumable_execution() - .with_review_approval_policy(ReviewApprovalPolicy::WorkerOutput) - .with_result_projection(|result| match result { - Ok(outcome) => ToolResult::success( - json!({ - "final_output": outcome.state.final_output, - "approved": outcome.state.approved, - "cancelled": outcome.state.cancelled, - "revisions": outcome.state.revisions, - "reviews": outcome.state.reviews, - "pending": outcome.pending.as_ref().map(|pending| json!({ - "node": pending.node, - "thread_id": pending.thread_id, - "payload": pending.payload, - })), - }) - .to_string(), - ), - Err(error) => ToolResult::error(error), - }); - - let parent = RunContext::new(RunConfig::new("parent"), ()); - let result = ToolDispatch::<(), ()>::execute( - &tool, - &(), - new_call_id(), - json!({"task":"approve durably"}), - tinytools::ToolCallOptions::default(), - &parent, - ) - .await - .expect("the durable approval interrupt is a result, not an error"); - let value: Value = serde_json::from_str(&result.output()).unwrap(); - assert!(!result.is_error); - assert_eq!(value["final_output"], Value::Null); - assert_eq!(value["approved"], true); - assert_eq!(value["reviews"], json!(["APPROVE after the final checks"])); - assert_eq!(value["pending"]["node"], "approval"); - assert_eq!(value["pending"]["thread_id"], "pending-review"); - assert!(!event_sink.is_empty()); -} - -#[tokio::test] -async fn resumable_delegation_preserves_caller_cancellation() { - use tinyagents_graph::{DelegationState, checkpoint::InMemoryCheckpointer}; - use tinyagents_harness::{ - cancel::CancellationToken, - context::{RunConfig, RunContext}, - ids::new_call_id, - tool::ToolDispatch, - }; - - let checkpointer = std::sync::Arc::new(InMemoryCheckpointer::<DelegationState>::new()); - let cancellation = CancellationToken::new(); - cancellation.cancel(); - let config_checkpointer = checkpointer.clone(); - let config_cancellation = cancellation.clone(); - let tool = DelegationTool::<(), (), _, _>::new( - move |_, _, _| DelegationConfig { - checkpointer: Some(config_checkpointer.clone()), - thread_id: Some("cancelled-run".into()), - cancel: config_cancellation.clone(), - ..DelegationConfig::default() - }, - |_, _, _| async { Ok(DelegationStageOutput::done("must not run")) }, - ) - .with_resumable_execution() - .with_result_projection(|result| match result { - Ok(outcome) if outcome.state.cancelled => ToolResult::error("CANCELLED"), - Ok(_) => ToolResult::success("unexpected completion"), - Err(error) => ToolResult::error(error), - }); - let parent = RunContext::new(RunConfig::new("parent"), ()); - let result = ToolDispatch::<(), ()>::execute( - &tool, - &(), - new_call_id(), - json!({"task":"cancel before start"}), - tinytools::ToolCallOptions::default(), - &parent, - ) - .await - .unwrap(); - assert!(result.is_error); - assert_eq!(result.output(), "CANCELLED"); -} - -#[tokio::test] -async fn result_projection_receives_raw_graph_errors() { - use tinyagents_harness::{ - context::{RunConfig, RunContext}, - ids::new_call_id, - tool::ToolDispatch, - }; - - let tool = DelegationTool::<(), (), _, _>::new( - |_, _, _| DelegationConfig::default(), - |_, _, _| async { Err("worker database unavailable".into()) }, - ) - .with_resumable_execution() - .with_result_projection(|result| match result { - Ok(_) => ToolResult::success("unexpected completion"), - Err(error) => ToolResult::error(format!("HOST ERROR: {error}")), - }); - let parent = RunContext::new(RunConfig::new("parent"), ()); - let result = ToolDispatch::<(), ()>::execute( - &tool, - &(), - new_call_id(), - json!({"task":"surface worker failure"}), - tinytools::ToolCallOptions::default(), - &parent, - ) - .await - .expect("the projection formats the graph error as a tool result"); - - assert!(result.is_error); - assert_eq!( - result.output(), - "HOST ERROR: delegation graph run failed: model error: worker database unavailable" - ); -} diff --git a/crates/tinyagents-orchestration/src/subagent/invocation/jobs.rs b/crates/tinyagents-orchestration/src/subagent/invocation/jobs.rs index aef2f3023..4dc8a6b97 100644 --- a/crates/tinyagents-orchestration/src/subagent/invocation/jobs.rs +++ b/crates/tinyagents-orchestration/src/subagent/invocation/jobs.rs @@ -1,6 +1,7 @@ //! Asynchronous subagent job registry and host-facing control tools. use std::sync::Arc; +use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering}; use tinyagents_harness::cancel::CancellationToken; use tinyagents_harness::error::TinyAgentsError; @@ -13,14 +14,17 @@ use tinyagents_tasks::DetachedTaskRegistryError; use tinyinference_llm::message::Message; use tokio::sync::watch; -use crate::subagent::{AppliedResult, IncompleteKind}; - -use super::types::{DriverJobRegistration, JobControl, JobMeta}; -use super::{ - JobLink, SubAgentJob, SubAgentJobError, SubAgentJobId, SubAgentJobRegistry, SubAgentJobStatus, -}; +use super::types::{DriverJobRegistration, JobControl, JobMeta, ParentJobScope}; +use super::{SubAgentJob, SubAgentJobError, SubAgentJobId, SubAgentJobRegistry, SubAgentJobStatus}; const LOG_PREFIX: &str = "[subagent-jobs]"; +type DriverJobControls = ( + SubAgentJobId, + SteeringHandle, + tinyagents_harness::run_queue::RunQueueHandle, + CancellationToken, + bool, +); fn registry_failure(error: DetachedTaskRegistryError) -> SubAgentJobError { SubAgentJobError::Registry(error.to_string()) @@ -32,16 +36,46 @@ struct JobExecutionControls { run_queue: Option<tinyagents_harness::run_queue::RunQueueHandle>, } -/// Settles an inline job if its tool future is dropped (tool timeout, parent -/// stream drop) or unwinds from a panic before the result is recorded. +struct JobInsert { + id: SubAgentJobId, + agent: String, + owner: String, + parent_run_id: tinyagents_harness::ids::RunId, + metadata: JobMeta, + execution: JobExecutionControls, +} + +pub(super) struct JobWaitGuard { + jobs: SubAgentJobRegistry, + waiters: Arc<AtomicUsize>, +} + +impl JobWaitGuard { + pub(super) fn pin(jobs: SubAgentJobRegistry, waiters: Arc<AtomicUsize>) -> Self { + waiters.fetch_add(1, Ordering::AcqRel); + Self { jobs, waiters } + } +} + +impl Drop for JobWaitGuard { + fn drop(&mut self) { + let previous = self.waiters.fetch_sub(1, Ordering::AcqRel); + if previous == 1 { + let _ = self.jobs.prune_terminal_over_cap(); + } + } +} + +/// Settles a driver-backed job if its future is dropped or unwinds before its +/// result is recorded. /// Call [`Self::disarm`] once the result has been written. -pub struct InlineJobGuard { +pub struct JobSettleGuard { jobs: SubAgentJobRegistry, id: SubAgentJobId, armed: bool, } -impl InlineJobGuard { +impl JobSettleGuard { pub fn new(jobs: SubAgentJobRegistry, id: SubAgentJobId) -> Self { Self { jobs, @@ -55,7 +89,7 @@ impl InlineJobGuard { } } -impl Drop for InlineJobGuard { +impl Drop for JobSettleGuard { fn drop(&mut self) { if self.armed { self.jobs.mark_aborted(&self.id, std::thread::panicking()); @@ -75,78 +109,37 @@ impl SubAgentJobRegistry { &self.tasks } - #[allow(dead_code)] - pub(crate) fn create(&self, agent: &str, owner: u64) -> (SubAgentJobId, SteeringHandle) { - self.create_with_cancellation(agent, owner, CancellationToken::new(), JobLink::default()) - .expect("a fresh job id registers") - } - - /// Registers a job whose child run observes `cancellation`, so - /// [`Self::cancel_owned`] can stop exactly this job. - pub(crate) fn create_with_cancellation( - &self, - agent: &str, - owner: u64, - cancellation: CancellationToken, - link: JobLink, - ) -> Result<(SubAgentJobId, SteeringHandle), SubAgentJobError> { - self.create_with_cancellation_and_queue(agent, owner, cancellation, link, None) - } - - /// Registers a job with an optional loop-consumed queue for content - /// messages routed to a specific lane. A job the registry could not record - /// is an error: its handle would be untracked, so it could be neither - /// listed, steered nor cancelled. - pub(crate) fn create_with_cancellation_and_queue( - &self, - agent: &str, - owner: u64, - cancellation: CancellationToken, - link: JobLink, - run_queue: Option<tinyagents_harness::run_queue::RunQueueHandle>, - ) -> Result<(SubAgentJobId, SteeringHandle), SubAgentJobError> { - let id = SubAgentJobId(format!("subagent-job-{}", next_seq())); - let steering = - SteeringHandle::new(SteeringPolicy::new().allow(SteeringCommandKind::InjectMessage)); - if let Err(error) = self.insert_job( - id.clone(), - agent, - owner.to_string(), - link, - JobExecutionControls { - cancellation, - steering: steering.clone(), - run_queue, - }, - ) { - tracing::error!("{LOG_PREFIX} register.failed job_id={id} error={error}"); - return Err(error); - } - Ok((id, steering)) - } - fn insert_job( &self, - id: SubAgentJobId, - agent: &str, - owner: String, - link: JobLink, - execution: JobExecutionControls, + insert: JobInsert, + driver_jobs: &mut std::collections::HashMap< + crate::subagent::SubagentTaskKey, + SubAgentJobId, + >, ) -> Result<(), SubAgentJobError> { + let JobInsert { + id, + agent, + owner, + parent_run_id, + metadata, + execution, + } = insert; if id.as_str().is_empty() { return Err(SubAgentJobError::Registry("empty job id".into())); } + let lineage = super::types::JobLineage::new(Some(parent_run_id), None); let task_id = TaskId::new(id.as_str()); let job = SubAgentJob { id: id.clone(), - agent: agent.to_owned(), + agent, status: SubAgentJobStatus::Queued, output: None, error: None, - subagent_run_id: link.subagent_run_id, + subagent_run_id: None, worker_thread_id: None, extensions: std::collections::BTreeMap::new(), - parent_tool_call_id: link.parent_tool_call_id, + parent_tool_call_id: None, incomplete_kind: None, artifacts: Vec::new(), schema_error: None, @@ -157,12 +150,13 @@ impl SubAgentJobRegistry { if controls.contains_key(&id) { return Err(SubAgentJobError::AlreadyExists(id.as_str().to_owned())); } + self.prune_terminal_before_insert(driver_jobs, &mut controls)?; self.steering .register(task_id.clone(), execution.steering.clone()); - if let Err(error) = self.tasks.register_cooperative( + if let Err(error) = self.tasks.register_cooperative_managed( task_id.clone(), - owner, - JobMeta, + owner.clone(), + metadata, watcher, execution.cancellation.clone(), ) { @@ -173,12 +167,102 @@ impl SubAgentJobRegistry { id.clone(), JobControl { status, + owner, cancellation_requested: false, run_queue: execution.run_queue, steering: execution.steering, cancellation: execution.cancellation, + waiters: Arc::new(AtomicUsize::new(0)), + terminal_sequence: AtomicU64::new(0), + lineage, }, ); + self.changes + .send_modify(|generation| *generation = generation.wrapping_add(1)); + Ok(()) + } + + /// Prunes unobserved terminal jobs before the generic registry reaches its + /// finite soft cap. The shared runtime entry, control sender and driver key + /// are removed together; live jobs and jobs with current waiters survive. + fn prune_terminal_before_insert( + &self, + driver_jobs: &mut std::collections::HashMap< + crate::subagent::SubagentTaskKey, + SubAgentJobId, + >, + controls: &mut std::sync::MutexGuard< + '_, + std::collections::HashMap<SubAgentJobId, JobControl>, + >, + ) -> Result<(), SubAgentJobError> { + self.prune_terminal_over_cap_locked(driver_jobs, controls) + } + + /// Keeps unpinned terminal history within the soft cap while allowing + /// live work and terminal jobs pinned by waiters to exceed it temporarily. + pub(super) fn prune_terminal_over_cap(&self) -> Result<(), SubAgentJobError> { + let mut driver_jobs = self + .driver_jobs + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let mut controls = self.controls(); + self.prune_terminal_over_cap_locked(&mut driver_jobs, &mut controls) + } + + fn prune_terminal_over_cap_locked( + &self, + driver_jobs: &mut std::collections::HashMap< + crate::subagent::SubagentTaskKey, + SubAgentJobId, + >, + controls: &mut std::sync::MutexGuard< + '_, + std::collections::HashMap<SubAgentJobId, JobControl>, + >, + ) -> Result<(), SubAgentJobError> { + let terminal_count = controls + .values() + .filter(|control| { + control.status.borrow().status.is_terminal() + && control.waiters.load(Ordering::Acquire) == 0 + }) + .count(); + let excess = terminal_count.saturating_sub(super::types::SUBAGENT_JOB_SOFT_CAP); + if excess == 0 { + return Ok(()); + } + + let mut candidates: Vec<_> = controls + .iter() + .filter(|(_, control)| { + control.status.borrow().status.is_terminal() + && control.waiters.load(Ordering::Acquire) == 0 + }) + .map(|(id, control)| { + ( + control.terminal_sequence.load(Ordering::Acquire), + id.clone(), + ) + }) + .collect(); + candidates.sort_by_key(|(sequence, _)| *sequence); + let mut removed = std::collections::HashSet::with_capacity(candidates.len()); + for (_, id) in candidates.into_iter().take(excess) { + let task_id = TaskId::new(id.as_str()); + if self + .tasks + .prune_terminal_trusted(&task_id) + .map_err(registry_failure)? + .is_some() + { + controls.remove(&id); + removed.insert(id); + } + } + if !removed.is_empty() { + driver_jobs.retain(|_, mapped_id| !removed.contains(mapped_id)); + } Ok(()) } @@ -198,10 +282,40 @@ impl SubAgentJobRegistry { ), SubAgentJobError, > { + self.register_driver_job_inner(registration, None) + } + + /// Registers a driver job with the actual live caller run. The driver + /// reports each planned attempt with [`Self::set_attempt_run_id`]. Resuming + /// the same durable task appends another parent edge and retains earlier + /// attempts. + pub fn register_driver_job_with_parent_run( + &self, + registration: DriverJobRegistration, + parent_run_id: tinyagents_harness::ids::RunId, + ) -> Result< + ( + SubAgentJobId, + SteeringHandle, + tinyagents_harness::run_queue::RunQueueHandle, + CancellationToken, + bool, + ), + SubAgentJobError, + > { + self.register_driver_job_inner(registration, Some(parent_run_id)) + } + + fn register_driver_job_inner( + &self, + registration: DriverJobRegistration, + parent_run_id: Option<tinyagents_harness::ids::RunId>, + ) -> Result<DriverJobControls, SubAgentJobError> { let DriverJobRegistration { key, agent, owner, + workspace, cancellation, steering, run_queue, @@ -211,14 +325,16 @@ impl SubAgentJobRegistry { .driver_jobs .lock() .unwrap_or_else(std::sync::PoisonError::into_inner); + // Lock order: driver_jobs -> root_leases -> controls -> detached task registry. + let mut root_leases = self.root_leases(); if let Some(id) = driver_jobs.get(&key) { - if self.get_owned(id.as_str(), &owner).is_none() { - return Err(SubAgentJobError::NotFound(id.as_str().to_owned())); - } - let controls = self.controls(); + let mut controls = self.controls(); let control = controls - .get(id) + .get_mut(id) .ok_or_else(|| SubAgentJobError::NotFound(id.as_str().to_owned()))?; + if control.owner != owner { + return Err(SubAgentJobError::NotFound(id.as_str().to_owned())); + } let status = control.status.borrow().status; if status.is_terminal() { return Ok(( @@ -239,17 +355,24 @@ impl SubAgentJobRegistry { .clone() .ok_or_else(|| SubAgentJobError::LaneUnavailable(id.as_str().to_owned()))?; if status == SubAgentJobStatus::AwaitingInput && continuing { + if let Some(parent_run_id) = parent_run_id { + control.lineage.add(Some(parent_run_id), None); + } control.status.send_modify(|job| { job.status = SubAgentJobStatus::Running; job.error = None; }); - return Ok(( + self.changes + .send_modify(|generation| *generation = generation.wrapping_add(1)); + let result = ( id.clone(), control.steering.clone(), queue, control.cancellation.clone(), true, - )); + ); + self.refresh_root_leases_locked(&mut root_leases, &mut controls); + return Ok(result); } return Ok(( id.clone(), @@ -266,18 +389,34 @@ impl SubAgentJobRegistry { let steering = steering.unwrap_or_else(|| { SteeringHandle::new(SteeringPolicy::new().allow(SteeringCommandKind::InjectMessage)) }); + let parent_run_id = parent_run_id.unwrap_or_else(|| { + // A legacy caller's durable key proves its historical parent, + // but does not prove the current child attempt's native RunId. + tinyagents_harness::ids::RunId::new(key.parent_run_id.clone()) + }); self.insert_job( - id.clone(), - &agent, - owner, - JobLink::default(), - JobExecutionControls { - cancellation, - steering: steering.clone(), - run_queue: Some(queue.clone()), + JobInsert { + id: id.clone(), + agent, + owner, + parent_run_id, + metadata: JobMeta { + parent_thread_id: key.thread_id.clone(), + workspace, + }, + execution: JobExecutionControls { + cancellation, + steering: steering.clone(), + run_queue: Some(queue.clone()), + }, }, + &mut driver_jobs, )?; driver_jobs.insert(key, id.clone()); + { + let mut controls = self.controls(); + self.refresh_root_leases_locked(&mut root_leases, &mut controls); + } let token = self .controls() .get(&id) @@ -289,7 +428,7 @@ impl SubAgentJobRegistry { pub fn mark_running(&self, id: &SubAgentJobId) { if let Some(control) = self.controls().get(id) { - control.status.send_if_modified(|job| { + let changed = control.status.send_if_modified(|job| { if matches!( job.status, SubAgentJobStatus::Queued | SubAgentJobStatus::AwaitingInput @@ -300,6 +439,10 @@ impl SubAgentJobRegistry { false } }); + if changed { + self.changes + .send_modify(|generation| *generation = generation.wrapping_add(1)); + } } } @@ -309,80 +452,43 @@ impl SubAgentJobRegistry { /// terminal state wins: a settled job is left untouched and `false` is /// returned. fn settle(&self, id: &SubAgentJobId, settle: impl FnOnce(&mut SubAgentJob, bool)) -> bool { - let controls = self.controls(); - let Some(control) = controls.get(id) else { - return false; + let settled = { + let controls = self.controls(); + let Some(control) = controls.get(id) else { + return false; + }; + if control.status.borrow().status.is_terminal() { + return false; + } + let task_id = TaskId::new(id.as_str()); + let _ = self.tasks.release_cancellation_trusted(&task_id); + let cancellation_requested = control.cancellation_requested; + control + .status + .send_modify(|job| settle(job, cancellation_requested)); + if control.status.borrow().status.is_terminal() { + self.stamp_terminal_sequence(control); + self.changes + .send_modify(|generation| *generation = generation.wrapping_add(1)); + } + self.steering.deregister(&task_id); + true }; - if control.status.borrow().status.is_terminal() { - return false; + if settled { + let _ = self.prune_terminal_over_cap(); } - let task_id = TaskId::new(id.as_str()); - let _ = self.tasks.release_cancellation_trusted(&task_id); - let cancellation_requested = control.cancellation_requested; - control - .status - .send_modify(|job| settle(job, cancellation_requested)); - self.steering.deregister(&task_id); - true + settled } - pub(crate) fn mark_result( - &self, - id: &SubAgentJobId, - result: Result<tinyagents_harness::middleware::AgentRun, TinyAgentsError>, - ) { - self.mark_result_applied(id, result, None); - } - - /// Settles a job like [`Self::mark_result`], publishing the - /// result-policy-applied output in the same registry write so no reader - /// ever sees a terminal job with the raw, unpolicied output. - pub(crate) fn mark_result_applied( - &self, - id: &SubAgentJobId, - result: Result<tinyagents_harness::middleware::AgentRun, TinyAgentsError>, - applied: Option<AppliedResult>, - ) { - let settled = self.settle(id, |job, cancellation_requested| match result { - // An owner cancel wins over whatever the run ended with, as in - // `mark_budget_overrun`. - _ if cancellation_requested => { - job.status = SubAgentJobStatus::Cancelled; - job.error = Some(TinyAgentsError::Cancelled.to_string()); - } - Ok(run) => { - job.status = SubAgentJobStatus::Completed; - job.output = run.text(); - if let Some(applied) = applied { - job.output = Some(applied.text); - job.artifacts.extend(applied.artifact); - job.schema_error = applied.schema_error; - job.artifact_error = applied.artifact_error; - } - } - Err(TinyAgentsError::Cancelled) => { - job.status = SubAgentJobStatus::Cancelled; - job.error = Some(TinyAgentsError::Cancelled.to_string()); - } - Err(error @ TinyAgentsError::LimitExceeded(_)) => { - job.status = SubAgentJobStatus::Incomplete; - job.incomplete_kind = Some(IncompleteKind::BudgetExceeded); - job.error = Some(error.to_string()); - } - Err(error @ TinyAgentsError::Timeout(_)) => { - job.status = SubAgentJobStatus::Incomplete; - job.incomplete_kind = Some(IncompleteKind::Timeout); - job.error = Some(error.to_string()); - } - Err(error) => { - job.status = SubAgentJobStatus::Failed; - job.error = Some(error.to_string()); - } - }); - if !settled { - // Already settled (e.g. cancelled by the owner): the first - // terminal state wins. - tracing::debug!("{LOG_PREFIX} mark_result.ignored job_id={id}"); + fn stamp_terminal_sequence(&self, control: &JobControl) { + if control.terminal_sequence.load(Ordering::Acquire) == 0 { + let sequence = self.next_terminal_sequence.fetch_add(1, Ordering::Relaxed); + let _ = control.terminal_sequence.compare_exchange( + 0, + sequence, + Ordering::Release, + Ordering::Relaxed, + ); } } @@ -409,7 +515,12 @@ impl SubAgentJobRegistry { job.status = SubAgentJobStatus::Cancelled; job.error = Some(TinyAgentsError::Cancelled.to_string()); }); + self.changes + .send_modify(|generation| *generation = generation.wrapping_add(1)); + self.stamp_terminal_sequence(control); self.steering.deregister(&task_id); + drop(controls); + let _ = self.prune_terminal_over_cap(); } else if let Ok(outcome) = result { control.status.send_modify(|job| { job.output = Some(outcome.output); @@ -505,43 +616,56 @@ impl SubAgentJobRegistry { self.cancel_owned(id.as_str(), owner) } - /// Points the job link at the attempt that is now running. - pub(crate) fn set_attempt_run_id(&self, id: &SubAgentJobId, run_id: &str) { - if let Some(control) = self.controls().get(id) { - control.status.send_if_modified(|job| { - if job.status.is_terminal() { - return false; - } - job.subagent_run_id = Some(run_id.to_owned()); - true - }); - } + /// Cancels a driver-backed job and immediately tombstones its public + /// completion record. The router reservation suppresses both an already + /// pending result and a late result arriving after this call. A foreign + /// owner is rejected before the router is touched. + pub fn cancel_key_owned_and_tombstone( + &self, + key: &crate::subagent::SubagentTaskKey, + owner: impl ToString, + completions: &tinyagents_tasks::CompletionRouter, + ) -> Result<SubAgentJob, SubAgentJobError> { + let owner = owner.to_string(); + let id = self + .driver_job_for_key(key, &owner) + .ok_or_else(|| SubAgentJobError::NotFound(key.task_id.clone()))?; + self.cancel_owned_with_change( + id.as_str(), + owner, + Some((completions, key.task_id.as_str())), + ) + .map(|(job, _)| job) } - /// Settles a job whose run finished but overshot a budget: it keeps the - /// (policy-applied) output and ends `Incomplete(BudgetExceeded)`. - pub(crate) fn mark_budget_overrun( - &self, - id: &SubAgentJobId, - applied: AppliedResult, - reason: String, - ) { - self.settle(id, |job, cancellation_requested| { - if cancellation_requested { - // An owner cancellation that raced the finish wins, as in - // `mark_result`. - job.status = SubAgentJobStatus::Cancelled; - job.error = Some(TinyAgentsError::Cancelled.to_string()); + /// Points the job link at the attempt that is now running. + /// Records the actual native run id of one driver attempt for descendant + /// settlement. Call once for every attempt, including retries. + pub fn set_attempt_run_id(&self, id: &SubAgentJobId, run_id: &str) { + let mut root_leases = self.root_leases(); + let mut controls = self.controls(); + if let Some(control) = controls.get_mut(id) { + if control.status.borrow().status.is_terminal() { return; } - job.status = SubAgentJobStatus::Incomplete; - job.incomplete_kind = Some(IncompleteKind::BudgetExceeded); - job.error = Some(reason); - job.output = Some(applied.text); - job.artifacts.extend(applied.artifact); - job.schema_error = applied.schema_error; - job.artifact_error = applied.artifact_error; - }); + let run_id = tinyagents_harness::ids::RunId::new(run_id); + let new_edge = control.lineage.children.insert(run_id.clone()); + for lease in root_leases.values_mut() { + if lease.owner == control.owner && lease.pinned_jobs.contains(id) { + lease.reachable_runs.insert(run_id.clone()); + } + } + let changed = control.status.send_if_modified(|job| { + let changed = job.subagent_run_id.as_deref() != Some(run_id.as_str()); + job.subagent_run_id = Some(run_id.as_str().to_owned()); + changed + }); + if new_edge || changed { + self.changes + .send_modify(|generation| *generation = generation.wrapping_add(1)); + } + self.refresh_root_leases_locked(&mut root_leases, &mut controls); + } } /// Marks a job `Failed` because its child task panicked or was aborted @@ -567,13 +691,63 @@ impl SubAgentJobRegistry { job_id: &str, owner: impl ToString, ) -> Result<SubAgentJob, SubAgentJobError> { + self.cancel_owned_with_change(job_id, owner, None) + .map(|(job, _)| job) + } + + fn cancel_owned_with_change( + &self, + job_id: &str, + owner: impl ToString, + completion_tombstone: Option<(&tinyagents_tasks::CompletionRouter, &str)>, + ) -> Result<(SubAgentJob, bool), SubAgentJobError> { + let owner = owner.to_string(); let id = SubAgentJobId(job_id.to_owned()); + let task_id = TaskId::new(job_id); + match self.tasks.snapshot(&task_id, &owner) { + Ok(_) => {} + Err(DetachedTaskRegistryError::Unknown | DetachedTaskRegistryError::NotOwned) => { + return Err(SubAgentJobError::NotFound(job_id.to_owned())); + } + Err(error) => return Err(registry_failure(error)), + } let mut controls = self.controls(); let control = controls .get_mut(&id) .ok_or_else(|| SubAgentJobError::NotFound(job_id.to_owned()))?; - let task_id = TaskId::new(job_id); - let mut snapshot = match self.tasks.cancel_cooperative(&task_id, &owner.to_string()) { + if control.status.borrow().status.is_terminal() + && (completion_tombstone.is_none() || !control.cancellation_requested) + { + return Err(SubAgentJobError::Terminal { + job_id: job_id.to_owned(), + status: control.status.borrow().status, + }); + } + if let Some((completions, task_id)) = completion_tombstone { + completions + .tombstone(task_id) + .map_err(|error| SubAgentJobError::Registry(error.to_string()))?; + } + if control.cancellation_requested { + if completion_tombstone.is_some() && !control.status.borrow().status.is_terminal() { + let task_id = TaskId::new(job_id); + let _ = self.tasks.release_cancellation_trusted(&task_id); + control.status.send_modify(|job| { + job.status = SubAgentJobStatus::Cancelled; + job.error = Some(TinyAgentsError::Cancelled.to_string()); + }); + self.changes + .send_modify(|generation| *generation = generation.wrapping_add(1)); + self.stamp_terminal_sequence(control); + self.steering.deregister(&task_id); + let snapshot = control.status.borrow().clone(); + drop(controls); + let _ = self.prune_terminal_over_cap(); + return Ok((snapshot, true)); + } + return Ok((control.status.borrow().clone(), false)); + } + let mut snapshot = match self.tasks.cancel_cooperative(&task_id, &owner) { Ok(snapshot) => snapshot.status, Err(DetachedTaskRegistryError::AlreadyDone) => { return Err(SubAgentJobError::Terminal { @@ -588,18 +762,76 @@ impl SubAgentJobRegistry { }; tracing::debug!("{LOG_PREFIX} cancel_owned job_id={job_id}"); control.cancellation_requested = true; - if control.status.borrow().status == SubAgentJobStatus::AwaitingInput { + if completion_tombstone.is_some() + || control.status.borrow().status == SubAgentJobStatus::AwaitingInput + { let _ = self.tasks.release_cancellation_trusted(&task_id); control.status.send_modify(|job| { job.status = SubAgentJobStatus::Cancelled; job.error = Some(TinyAgentsError::Cancelled.to_string()); }); + self.changes + .send_modify(|generation| *generation = generation.wrapping_add(1)); + self.stamp_terminal_sequence(control); self.steering.deregister(&task_id); - return Ok(control.status.borrow().clone()); + let snapshot = control.status.borrow().clone(); + drop(controls); + let _ = self.prune_terminal_over_cap(); + return Ok((snapshot, true)); } snapshot.error = Some("cancellation requested; job will be cancelled when the child unwinds".to_owned()); - Ok(snapshot) + Ok((snapshot, true)) + } + + /// Requests cancellation for jobs belonging to the exact parent workspace + /// and thread. Jobs with incomplete metadata never match this scoped call. + pub fn cancel_for_thread(&self, scope: &ParentJobScope) -> Result<usize, SubAgentJobError> { + self.cancel_matching(|metadata| { + metadata.parent_thread_id.as_deref() == Some(scope.thread_id.as_str()) + && metadata.workspace.as_ref() == Some(&scope.workspace) + }) + } + + /// Requests cancellation for every live job in one exact workspace. + /// Jobs with incomplete workspace metadata never match this scoped call. + pub fn cancel_all_for_workspace( + &self, + workspace: &tinytools::WorkspaceDescriptor, + ) -> Result<usize, SubAgentJobError> { + self.cancel_matching(|metadata| metadata.workspace.as_ref() == Some(workspace)) + } + + /// Requests cancellation for every live job while retaining each public + /// job record for the normal settlement path. This is trusted process-wide + /// supervision; tenant controls should use [`Self::cancel_all_for_workspace`]. + pub fn cancel_all(&self) -> Result<usize, SubAgentJobError> { + self.cancel_matching(|_| true) + } + + fn cancel_matching( + &self, + matches: impl Fn(&JobMeta) -> bool, + ) -> Result<usize, SubAgentJobError> { + let ids = self.controls().keys().cloned().collect::<Vec<_>>(); + let mut requested = 0; + for id in ids { + let task_id = TaskId::new(id.as_str()); + let snapshot = match self.tasks.snapshot_trusted(&task_id) { + Ok(snapshot) if !snapshot.status.status.is_terminal() => snapshot, + Ok(_) | Err(DetachedTaskRegistryError::Unknown) => continue, + Err(error) => return Err(registry_failure(error)), + }; + if !matches(&snapshot.metadata) { + continue; + } + match self.cancel_owned_with_change(id.as_str(), snapshot.owner_id, None) { + Ok((_, changed)) => requested += usize::from(changed), + Err(SubAgentJobError::Terminal { .. } | SubAgentJobError::NotFound(_)) => {} + Err(error) => return Err(error), + } + } + Ok(requested) } /// Returns a snapshot for `job_id` when it belongs to `owner`. @@ -657,24 +889,42 @@ impl SubAgentJobRegistry { job_id: &str, owner: impl ToString, ) -> Result<SubAgentJob, SubAgentJobError> { - self.try_get_owned(job_id, owner)? - .ok_or_else(|| SubAgentJobError::NotFound(job_id.to_owned()))?; - let mut status = { + let owner = owner.to_string(); + let id = SubAgentJobId(job_id.to_owned()); + let task_id = TaskId::new(job_id); + let (mut status, waiters) = { + // Pin the control before releasing its lock so cap pruning cannot + // remove a terminal snapshot between ownership validation and the + // watch subscription. let controls = self.controls(); - controls - .get(&SubAgentJobId(job_id.to_owned())) - .ok_or_else(|| SubAgentJobError::NotFound(job_id.to_owned()))? - .status - .subscribe() + match self.tasks.snapshot(&task_id, &owner) { + Ok(_) => {} + Err(DetachedTaskRegistryError::Unknown | DetachedTaskRegistryError::NotOwned) => { + return Err(SubAgentJobError::NotFound(job_id.to_owned())); + } + Err(error) => return Err(registry_failure(error)), + } + let control = controls + .get(&id) + .ok_or_else(|| SubAgentJobError::NotFound(job_id.to_owned()))?; + control.waiters.fetch_add(1, Ordering::AcqRel); + (control.status.subscribe(), control.waiters.clone()) + }; + let waiter_guard = JobWaitGuard { + jobs: self.clone(), + waiters, }; loop { let snapshot = status.borrow().clone(); if snapshot.status.is_terminal() || snapshot.status == SubAgentJobStatus::AwaitingInput { + drop(waiter_guard); return Ok(snapshot); } if status.changed().await.is_err() { - return Ok(status.borrow().clone()); + let snapshot = status.borrow().clone(); + drop(waiter_guard); + return Ok(snapshot); } } } @@ -686,6 +936,47 @@ impl SubAgentJobRegistry { Self::sorted(self.tasks.snapshots(None)) } + /// Returns live jobs for one exact owner and parent workspace/thread. + /// Incomplete parent metadata never matches this scoped roster query. + pub fn list_for_parent_scope( + &self, + owner: &str, + scope: &ParentJobScope, + ) -> Result<Vec<SubAgentJob>, SubAgentJobError> { + let snapshots = self + .tasks + .snapshots(Some(owner)) + .map_err(registry_failure)?; + let driver_jobs = self + .driver_jobs + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let mut jobs: Vec<SubAgentJob> = snapshots + .into_iter() + .filter(|snapshot| { + snapshot.metadata.workspace.as_ref() == Some(&scope.workspace) + && snapshot.metadata.parent_thread_id.as_deref() + == Some(scope.thread_id.as_str()) + }) + .map(|snapshot| { + let mut job = snapshot.status; + if !job.extensions.contains_key("task_id") + && let Some(key) = driver_jobs + .iter() + .find_map(|(key, id)| (id == &job.id).then_some(key)) + { + job.extensions.insert( + "task_id".into(), + serde_json::Value::String(key.task_id.clone()), + ); + } + job + }) + .collect(); + jobs.sort_by(|left, right| left.id.0.cmp(&right.id.0)); + Ok(jobs) + } + /// Jobs in stable id order, independent of the detached registry's own /// iteration order. A registry failure is logged, never silent. fn sorted( @@ -814,11 +1105,20 @@ impl SubAgentJobRegistry { Ok(false) } - fn controls( + pub(super) fn controls( &self, ) -> std::sync::MutexGuard<'_, std::collections::HashMap<SubAgentJobId, JobControl>> { self.controls .lock() .unwrap_or_else(std::sync::PoisonError::into_inner) } + + pub(super) fn root_leases( + &self, + ) -> std::sync::MutexGuard<'_, std::collections::HashMap<u64, super::types::RootLeaseRecord>> + { + self.root_leases + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + } } diff --git a/crates/tinyagents-orchestration/src/subagent/invocation/mod.rs b/crates/tinyagents-orchestration/src/subagent/invocation/mod.rs index be39c99c3..7f6aeaf1d 100644 --- a/crates/tinyagents-orchestration/src/subagent/invocation/mod.rs +++ b/crates/tinyagents-orchestration/src/subagent/invocation/mod.rs @@ -1,1162 +1,19 @@ -//! First-class sub-agents with recursion-depth tracking. +//! Driver-backed asynchronous subagent jobs and their lifecycle controls. //! -//! This is the harness's flagship recursion surface: it lets one agent run -//! another agent as a child of itself, so a language model orchestrating tools -//! is, transparently, a language model orchestrating *other models*. It is the -//! concrete "agents calling agents" mechanism behind the crate's -//! recursive-language-model framing — the in-harness analogue of the graph-side -//! `subgraph` recursion (in `tinyagents-graph`). -//! -//! This module provides the agent-calling-agent compositional primitive: -//! -//! - [`SubAgent`] wraps an [`AgentHarness`] and runs it as a *child run* one -//! level deeper in the recursion tree than its caller. -//! - [`SubAgentTool`] adapts a [`SubAgent`] into a typed [`ToolDispatch`] that -//! runs in the background by default, returning a job id immediately; callers -//! can request `mode: "inline"` to await the child and receive its final -//! result in the same call. -//! - [`tinyagents_harness_tools::subagents::SubAgentJobsTool`] and [`tinyagents_harness_tools::subagents::SubAgentMessageTool`] let the host or parent -//! query those jobs and steer live children by job id. -//! - [`SubAgentSession`] keeps a single [`SubAgent`] alive across multiple -//! turns, *reusing* the same harness while accumulating the conversation -//! transcript — the post-completion, human-in-the-loop reuse primitive. -//! -//! # Reuse vs. steering -//! -//! There are two ways an orchestrator keeps a sub-agent "in play" across human -//! input: -//! -//! - **Reuse** ([`SubAgentSession`]): the child run *completes*, the -//! orchestrator obtains human input, then calls the **same** sub-agent again -//! carrying the prior transcript. Nothing is killed or restarted. -//! - **Steering** ([`tinyagents_harness_tools::subagents::SubAgentMessageTool`]): an orchestrator or human sends a -//! message to a **still-running** job. Delivery occurs at a safe checkpoint. -//! -//! `SubAgentSession` implements the first. The flow is: -//! -//! 1. `session.send(state, ctx, vec![Message::user("…")])` — runs the sub-agent -//! over the retained transcript and folds its reply back in. -//! 2. Inspect the returned [`AgentRun`]; obtain human input out-of-band. -//! 3. `session.send(state, ctx, vec![Message::user(human_reply)])` — the same -//! sub-agent answers with the full prior context still in the transcript. -//! -//! Every send after the first emits [`AgentEvent::SubAgentReused`] so the reuse -//! is visible alongside the per-send -//! [`SubAgentStarted`][AgentEvent::SubAgentStarted]/[`SubAgentCompleted`][AgentEvent::SubAgentCompleted] -//! bracket. -//! -//! # Depth tracking -//! -//! Every run carries a `depth` in its [`RunConfig`] (top-level runs are depth -//! `0`). When a sub-agent is invoked at `parent_depth`, its child run is created -//! at `parent_depth + 1`. The depth cap is -//! [`RunLimits::max_depth`][tinyagents_harness::limits::RunLimits::max_depth] -//! (default [`RunLimits::DEFAULT_MAX_DEPTH`][tinyagents_harness::limits::RunLimits::DEFAULT_MAX_DEPTH], -//! i.e. `8`), read from the child harness's [`RunPolicy`][tinyagents_harness::runtime::RunPolicy]. -//! If the child depth would exceed the cap, the invocation fails fast with -//! [`TinyAgentsError::SubAgentDepth`] *before* any model call — a deterministic, -//! cheap guard against unbounded recursion. -//! -//! # Observability -//! -//! Each invocation emits [`AgentEvent::SubAgentStarted`] and -//! [`AgentEvent::SubAgentCompleted`] (carrying the sub-agent name and child -//! depth). When invoked with a shared [`EventSink`] — via -//! [`SubAgent::invoke_with_events`] or [`SubAgent::invoke_in_parent`] — the child -//! run's own events also flow onto the parent sink, so a parent observer sees -//! the full nested run tree. -//! -//! # Layout -//! -//! - `types` holds the public type definitions. -//! - This file holds the impls (constructors, the invoke methods, and the -//! typed-parent dispatcher). -//! - `*_tests.rs` files hold the focused tests. +//! Hosts create subagent requests through `SubagentDriver` and register the +//! resulting jobs here for status, cancellation, messaging, and descendant +//! settlement. This module does not construct or synchronously invoke child +//! harnesses. -mod delegation_tool; mod jobs; -mod policy_run; +mod settlement; mod status_map; mod types; -const LOG_PREFIX: &str = "[subagent-tool]"; - -pub use delegation_tool::{ - DELEGATION_EXECUTE_PROMPT, DELEGATION_MODE_FIELD, DELEGATION_PLAN_PROMPT, - DELEGATION_REVIEW_PROMPT, DelegationEnvelope, DelegationRequest, DelegationTool, - MultiTargetDelegationTool, ReviewApprovalPolicy, parse_approval, review_stage_output, -}; -pub use jobs::InlineJobGuard; +pub use jobs::JobSettleGuard; +pub use settlement::RootRunDescendantLease; pub use types::*; -#[cfg(test)] -#[path = "delegation_tool_tests.rs"] -mod delegation_tool_tests; - -use std::sync::Arc; - -use async_trait::async_trait; -use serde_json::{Value, json}; - -use tinyagents_harness::context::{RunConfig, RunContext}; -use tinyagents_harness::error::{Result, TinyAgentsError}; -use tinyagents_harness::events::{AgentEvent, EventSink}; -use tinyagents_harness::ids::{RunId, ThreadId, next_seq}; -use tinyagents_harness::middleware::AgentRun; -use tinyagents_harness::runtime::AgentHarness; -use tinyagents_harness::tool::ToolDispatch; -use tinyinference_llm::message::Message; - -use super::SpawnAdmission; - -impl<State: Send + Sync, Ctx: Send + Sync + 'static> SubAgent<State, Ctx> { - /// Creates a sub-agent wrapping `harness` with a stable `name` and - /// `description`. - pub fn new( - name: impl Into<String>, - description: impl Into<String>, - harness: Arc<AgentHarness<State, Ctx>>, - ) -> Self { - Self { - harness, - name: name.into(), - description: description.into(), - system_prompt: None, - } - } - - /// Sets a fixed system prompt prepended to every child run as a leading - /// system message. Returns `self` for chaining. - pub fn with_system_prompt(mut self, prompt: impl Into<String>) -> Self { - self.system_prompt = Some(prompt.into()); - self - } - - /// Returns the sub-agent's name. - pub fn name(&self) -> &str { - &self.name - } - - /// Returns the sub-agent's description. - pub fn description(&self) -> &str { - &self.description - } - - /// Returns the wrapped harness. - pub fn harness(&self) -> &Arc<AgentHarness<State, Ctx>> { - &self.harness - } - - /// Builds the child run's seed messages from the optional system prompt and - /// the caller `input`. - fn seed_messages(&self, input: String) -> Vec<Message> { - let mut messages = Vec::with_capacity(2); - if let Some(prompt) = &self.system_prompt { - messages.push(Message::system(prompt.clone())); - } - messages.push(Message::user(input)); - messages - } - - /// Builds the child [`RunConfig`] for an invocation at `parent_depth`, - /// enforcing the depth cap and deriving an isolated child thread from a - /// parent thread when one is available. - /// - /// `parent` is `Some((parent_run_id, ordinal))` for every entry point that - /// has a live parent [`RunContext`] to derive from — `ordinal` comes from - /// [`RunContext::next_child_ordinal`], a counter scoped to that one - /// context instance (not process-global). The child run id is then a pure - /// function of the parent's run id and that ordinal - /// (`{name}-d{depth}-{parent_run_id}-{ordinal}`), so two processes - /// replaying the identical sequence of calls against the identical parent - /// run derive the identical child run ids (M-2) — unlike the historical - /// `ids::next_seq()` suffix, which restarts at a different value every - /// process and made replayed journals of nested runs diverge across - /// processes. - /// - /// `parent` is `None` only for the standalone entry points - /// ([`Self::invoke`]/[`Self::invoke_with_events`]) that are not called - /// with a live parent context at all; those fall back to - /// [`tinyagents_harness::ids::next_seq`] since there is no parent run to derive - /// determinism from. - /// - /// Returns [`TinyAgentsError::SubAgentDepth`] when the child depth - /// (`parent_depth + 1`) would exceed the harness policy's `max_depth`. - fn child_config( - &self, - parent_depth: usize, - thread_id: Option<&ThreadId>, - max_turn_output_tokens: Option<u32>, - parent: Option<(&str, u64)>, - ) -> Result<RunConfig> { - let max_depth = self.harness.policy().limits.max_depth; - let child_depth = RunConfig::checked_child_depth(parent_depth, max_depth)?; - let child_run_id = match parent { - Some((parent_run_id, ordinal)) => { - format!("{}-d{child_depth}-{parent_run_id}-{ordinal}", self.name) - } - // No parent context to derive determinism from: suffix a - // process-unique sequence so each invocation still gets its own - // run id (a bare `{name}-d{depth}` was reused across invocations, - // which interleaved journals and status stores keyed by run id). - None => format!("{}-d{child_depth}-{}", self.name, next_seq()), - }; - let mut config = RunConfig::new(child_run_id.clone()) - .with_depth(child_depth) - .with_max_depth(max_depth); - config.thread_id = thread_id.map(|parent| child_thread_id(parent, &child_run_id)); - config.max_turn_output_tokens = max_turn_output_tokens; - Ok(config) - } - - /// Runs the sub-agent as a child run at `parent_depth`, returning the - /// child's [`AgentRun`]. - /// - /// The child run is created at `parent_depth + 1`. `ctx_data` seeds the - /// child [`RunContext`]. Sub-agent lifecycle events are emitted on the - /// child's own (fresh) event sink; use [`Self::invoke_with_events`] to fan - /// them out to a shared parent sink. - /// - /// # Errors - /// - /// Returns [`TinyAgentsError::SubAgentDepth`] if the child depth would - /// exceed the configured `max_depth`, or any error surfaced by the child - /// agent loop. - pub async fn invoke( - &self, - state: &State, - ctx_data: Ctx, - parent_depth: usize, - input: impl Into<String>, - ) -> Result<AgentRun> { - let config = self.child_config(parent_depth, None, None, None)?; - let ctx = RunContext::new(config, ctx_data); - self.run_child(state, ctx, input.into(), false).await - } - - /// Like [`Self::invoke`] but routes the child run's events (and the - /// sub-agent lifecycle events) onto the shared `events` sink so a parent - /// observer sees the nested run. - pub async fn invoke_with_events( - &self, - state: &State, - ctx_data: Ctx, - parent_depth: usize, - input: impl Into<String>, - events: &EventSink, - ) -> Result<AgentRun> { - let config = self.child_config(parent_depth, None, None, None)?; - let ctx = RunContext::new(config, ctx_data).with_events(events.clone()); - self.run_child(state, ctx, input.into(), false).await - } - - /// Runs the sub-agent as a child of the live `parent` context. - /// - /// This is the fully context-threaded entry point: the child depth is - /// derived from `parent.depth()` and the child inherits the parent's event - /// sink so all nested events share one stream. `ctx_data` seeds the child - /// context's user data. - /// - /// # Errors - /// - /// Identical to [`Self::invoke`]. - pub async fn invoke_in_parent( - &self, - state: &State, - ctx_data: Ctx, - parent: &RunContext<Ctx>, - input: impl Into<String>, - ) -> Result<AgentRun> { - // This is the generic explicit-model entry point, intentionally kept - // available to borrowed `State` callers. A hosted parent must enter - // through `invoke_hosted_in_parent` below, where the `State: 'static` - // bound makes the invocation authority type-safe. Falling through to - // the explicit loop here would discard the parent's definition and - // approval authority, so reject it before constructing a child. - if parent.is_hosted() { - return Err(TinyAgentsError::Validation( - "hosted parent delegation requires invoke_hosted_in_parent".into(), - )); - } - // The child harness may tighten the tree cap, but it may never widen - // the explicit parent lineage cap. `RunContext::child` is the one - // place that copies the live recursive capabilities and creates the - // isolated counters/control slot for this invocation. - let config = self.child_config( - parent.depth(), - parent.thread_id(), - parent.config.max_turn_output_tokens, - Some((parent.run_id().as_str(), parent.next_child_ordinal())), - )?; - let ctx = parent.child(config, ctx_data)?; - self.run_child(state, ctx, input.into(), parent.streaming) - .await - } - - /// Shared driver: emits the sub-agent lifecycle events around the child - /// agent loop. - /// - /// When `streaming` is `true` the child runs through the streaming loop - /// path, so its per-token model/reasoning deltas are emitted onto the - /// shared [`EventSink`] and reach the parent's stream (stamped with the - /// child's own `run_id` and `depth`). When `false` the child runs the - /// unary path, leaving the parent's event stream unchanged. - async fn run_child( - &self, - state: &State, - ctx: RunContext<Ctx>, - input: String, - streaming: bool, - ) -> Result<AgentRun> { - let depth = ctx.depth(); - let child_run_id = ctx.run_id().clone(); - let messages = self.seed_messages(input.clone()); - // Clone the sink (it shares listeners and the offset counter with the - // context) so the completion event can be emitted after `ctx` is moved - // into the child agent loop. - let events = ctx.events.clone(); - - events.emit(AgentEvent::SubAgentStarted { - name: self.name.clone(), - depth, - }); - events.emit(AgentEvent::SubAgentRunStarted { - name: self.name.clone(), - depth, - run_id: child_run_id.clone(), - }); - - let run = if streaming { - self.harness - .invoke_streaming_in_context(state, ctx, messages) - .await? - } else { - self.harness.invoke_in_context(state, ctx, messages).await? - }; - - events.emit(AgentEvent::SubAgentCompleted { - name: self.name.clone(), - depth, - }); - events.emit(AgentEvent::SubAgentRunCompleted { - run_id: child_run_id, - }); - - Ok(run) - } -} - -impl<State: Send + Sync + 'static, Ctx: Send + Sync + 'static> SubAgent<State, Ctx> { - /// Runs this child under the exact host authority installed on `parent`. - /// - /// Unlike [`Self::invoke_in_parent`], this path resolves the parent's - /// delegate allowlist and re-enters the child through the parent's shared - /// invocation bundle. A context that is not hosted fails closed rather - /// than silently acquiring an unrelated harness configuration. - pub async fn invoke_hosted_in_parent( - &self, - state: &State, - ctx_data: Ctx, - parent: &RunContext<Ctx>, - input: impl Into<String>, - ) -> Result<AgentRun> { - if !parent.is_hosted() { - return Err(TinyAgentsError::Validation( - "hosted subagent invocation requires parent host authority".into(), - )); - } - let config = self.child_config( - parent.depth(), - parent.thread_id(), - parent.config.max_turn_output_tokens, - Some((parent.run_id().as_str(), parent.next_child_ordinal())), - )?; - let child = parent.child(config, ctx_data)?; - self.run_hosted_child(state, child, input.into(), parent.streaming) - .await - } - - /// Hosted recursive driver. Kept separate from the generic explicit-model - /// path so a borrowed `State` never has to interact with live host - /// authority stored on a `RunContext`. - async fn run_hosted_child( - &self, - state: &State, - ctx: RunContext<Ctx>, - input: String, - streaming: bool, - ) -> Result<AgentRun> { - if !ctx.is_hosted() { - return self.run_child(state, ctx, input, streaming).await; - } - let depth = ctx.depth(); - let child_run_id = ctx.run_id().clone(); - let messages = self.seed_messages(input.clone()); - let events = ctx.events.clone(); - events.emit(AgentEvent::SubAgentStarted { - name: self.name.clone(), - depth, - }); - events.emit(AgentEvent::SubAgentRunStarted { - name: self.name.clone(), - depth, - run_id: child_run_id.clone(), - }); - let run = match self - .harness - .invoke_authorized_child(state, ctx, self.name.clone(), messages, streaming) - .await? - { - Some(run) => run, - None => { - return Err(TinyAgentsError::Validation( - "hosted child invocation lost its parent authority".into(), - )); - } - }; - events.emit(AgentEvent::SubAgentCompleted { - name: self.name.clone(), - depth, - }); - events.emit(AgentEvent::SubAgentRunCompleted { - run_id: child_run_id, - }); - Ok(run) - } -} - -/// Stamps the explicit parent-call -> child-run link onto a child run's -/// metadata, preserving whatever metadata the child already carries. -fn stamp_link_metadata( - metadata: &mut Value, - subagent_run_id: &str, - job_id: &str, - tool_call_id: Option<&str>, -) { - if !metadata.is_object() { - let original = std::mem::take(metadata); - *metadata = json!({"value": original}); - } - if let Value::Object(map) = metadata { - map.insert("subagent_run_id".into(), json!(subagent_run_id)); - map.insert("subagent_job_id".into(), json!(job_id)); - if let Some(tool_call_id) = tool_call_id { - map.insert("parent_tool_call_id".into(), json!(tool_call_id)); - } else { - map.remove("parent_tool_call_id"); - } - } -} - -/// Derives an isolated child thread id from the parent thread and the child's -/// run id. The run id already carries a process-unique sequence, so the thread -/// id inherits its uniqueness. -fn child_thread_id(parent: &ThreadId, child_run_id: &str) -> ThreadId { - ThreadId::new(format!("{}-subagent-{child_run_id}", parent.as_str())) -} - -impl<State: Send + Sync, Ctx: Send + Sync> SubAgentSession<State, Ctx> { - /// Creates a session that reuses `subagent` across turns. - /// - /// The child runs at depth `1` by default (caller `parent_depth = 0`); use - /// [`Self::with_parent_depth`] to express deeper nesting. - pub fn new(subagent: Arc<SubAgent<State, Ctx>>) -> Self { - Self { - subagent, - transcript: Vec::new(), - turn: 0, - parent_depth: 0, - events: EventSink::new(), - seeded: false, - } - } - - /// Creates a session from an owned [`SubAgent`], wrapping it in an `Arc`. - pub fn from_subagent(subagent: SubAgent<State, Ctx>) -> Self { - Self::new(Arc::new(subagent)) - } - - /// Routes the reuse lifecycle and the child run's own events onto `events` - /// so an external observer (or testkit recorder) sees every send. Returns - /// `self` for chaining. - pub fn with_events(mut self, events: EventSink) -> Self { - self.events = events; - self - } - - /// Sets the caller depth the child runs at; the child run is created at - /// `parent_depth + 1`. Returns `self` for chaining. - pub fn with_parent_depth(mut self, parent_depth: usize) -> Self { - self.parent_depth = parent_depth; - self - } - - /// Returns the reused sub-agent. The same `Arc` is shared across every - /// send, so this is how callers confirm the harness was never rebuilt. - pub fn subagent(&self) -> &Arc<SubAgent<State, Ctx>> { - &self.subagent - } - - /// Returns the accumulated conversation transcript carried across sends. - pub fn transcript(&self) -> &[Message] { - &self.transcript - } - - /// Returns the number of completed sends (turns). - pub fn turns(&self) -> usize { - self.turn - } - - /// Clears the retained transcript and turn counter, so the next [`Self::send`] - /// starts a fresh conversation (re-seeding the fixed system prompt). The - /// underlying [`SubAgent`]/harness is left untouched and still reused. - pub fn reset(&mut self) { - self.transcript.clear(); - self.turn = 0; - self.seeded = false; - } - - /// Builds the child [`RunConfig`] for the current turn, enforcing the depth - /// cap exactly as [`SubAgent::invoke`] does. - fn child_config(&self) -> Result<RunConfig> { - let max_depth = self.subagent.harness.policy().limits.max_depth; - let child_depth = RunConfig::checked_child_depth(self.parent_depth, max_depth)?; - // As in `SubAgent::child_config`, suffix a process-unique sequence so - // two sessions reusing the same sub-agent never share run ids for the - // same turn number. - Ok(RunConfig::new(format!( - "{}-t{}-d{child_depth}-{}", - self.subagent.name, - self.turn, - next_seq() - )) - .with_depth(child_depth) - .with_max_depth(max_depth)) - } - - /// Runs the reused sub-agent for one turn over the FULL accumulated - /// transcript, then folds the produced assistant/tool messages back into - /// the transcript so the next send continues with full context. - /// - /// `input` (typically a single [`Message::user`] carrying human input) is - /// appended to the retained transcript before the run. On the first send - /// the sub-agent's fixed [`SubAgent::with_system_prompt`] is prepended once. - /// The same underlying harness is reused on every call — nothing is - /// reconstructed. - /// - /// # Errors - /// - /// Returns [`TinyAgentsError::SubAgentDepth`] if the child depth would - /// exceed the configured `max_depth`, or any error surfaced by the child - /// agent loop. - pub async fn send( - &mut self, - state: &State, - ctx_data: Ctx, - input: Vec<Message>, - ) -> Result<AgentRun> { - // Seed the fixed system prompt once, on the first send. - if !self.seeded { - if let Some(prompt) = &self.subagent.system_prompt { - self.transcript.push(Message::system(prompt.clone())); - } - self.seeded = true; - } - - // Append the new (e.g. human/user) input to the retained transcript. - self.transcript.extend(input); - - let config = self.child_config()?; - let depth = config.depth(); - let ctx = RunContext::new(config, ctx_data).with_events(self.events.clone()); - let child_run_id = ctx.run_id().clone(); - - // Clone the sink so we can emit the completion event after `ctx` is - // moved into the child agent loop. - let events = self.events.clone(); - events.emit(AgentEvent::SubAgentStarted { - name: self.subagent.name.clone(), - depth, - }); - events.emit(AgentEvent::SubAgentRunStarted { - name: self.subagent.name.clone(), - depth, - run_id: child_run_id.clone(), - }); - if self.turn > 0 { - events.emit(AgentEvent::SubAgentReused { - name: self.subagent.name.clone(), - turn: self.turn, - }); - } - - // REUSE the same underlying harness/SubAgent (no reconstruction), - // running it over the full accumulated transcript. - let run = self - .subagent - .harness - .invoke_in_context(state, ctx, self.transcript.clone()) - .await?; - - events.emit(AgentEvent::SubAgentCompleted { - name: self.subagent.name.clone(), - depth, - }); - events.emit(AgentEvent::SubAgentRunCompleted { - run_id: child_run_id, - }); - - // Carry the produced assistant/tool messages forward. `run.messages` is - // the working transcript the loop ended with (everything we passed plus - // the new assistant/tool messages), so the next send continues with the - // full context. - self.transcript = run.messages.clone(); - self.turn += 1; - - Ok(run) - } -} - -impl<State: Clone + Send + Sync + 'static, Ctx: Send + Sync + 'static> SubAgentTool<State, Ctx> { - /// Default JSON Schema for a sub-agent tool: an object with one required - /// string field named [`SUBAGENT_INPUT_FIELD`]. - fn default_parameters() -> Value { - json!({ - "type": "object", - "properties": { - SUBAGENT_INPUT_FIELD: { - "type": "string", - "description": "The task or question to delegate to the sub-agent." - }, - SUBAGENT_MODE_FIELD: { - "type": "string", - "enum": ["background", "inline"], - "description": "`background` (default) returns a job id immediately while the sub-agent keeps running; `inline` waits for the sub-agent and returns its final result in this call." - } - }, - "required": [SUBAGENT_INPUT_FIELD] - }) - } - - /// Wraps `subagent` as a typed-parent tool. - /// - /// `child_data` is required: it makes application-data inheritance explicit - /// for every recursive invocation. - pub fn new(subagent: Arc<SubAgent<State, Ctx>>, child_data: ChildDataPolicy<Ctx>) -> Self { - let tool_name = subagent.name().to_owned(); - let tool_description = subagent.description().to_owned(); - Self { - subagent, - tool_name, - tool_description, - child_data, - parameters: Self::default_parameters(), - declaration: std::sync::OnceLock::new(), - jobs: SubAgentJobRegistry::new(), - admission: SpawnAdmission::default(), - policy: crate::subagent::SubAgentPolicy::default(), - role: crate::subagent::SubagentRole::default(), - delegation_tools: Vec::new(), - result_policy: crate::subagent::ResultPolicy::default(), - job_registry: None, - job_owner_scope: Arc::new(|owner| owner.to_string()), - } - } - - /// Applies a timeout/retry/budget policy to every spawned child. - /// - /// Call caps tighten the child's run limits; token caps are checked when it - /// finishes; a retry needs `retry.max_attempts > 1` and re-runs a fresh - /// child only after a retryable failure that ran no tools (unless - /// [`SubAgentPolicy::retry_after_tool_calls`](crate::subagent::SubAgentPolicy)). - pub fn with_policy(mut self, policy: crate::subagent::SubAgentPolicy) -> Self { - self.policy = policy; - self - } - - /// Declares the child's role. A [`SubagentRole::Leaf`](crate::subagent::SubagentRole) - /// child's harness must not expose delegation tools: the spawn is refused - /// otherwise (the shared harness cannot be filtered per call). - pub fn with_role(mut self, role: crate::subagent::SubagentRole) -> Self { - self.role = role; - self.warn_if_leaf_misconfigured(); - self - } - - /// Names the host's own delegation tools, which a leaf must not expose. - pub fn with_delegation_tools(mut self, names: Vec<String>) -> Self { - self.delegation_tools = names; - self.warn_if_leaf_misconfigured(); - self - } - - /// Why a leaf cannot spawn: its harness exposes a delegation tool (the job - /// and message tools, a host-named tool, or a tool with this tool's own - /// name). `None` for an orchestrator or a clean leaf. - fn leaf_violation(&self) -> Option<String> { - if self.role.can_delegate() { - return None; - } - let exposed = policy_run::delegation_tools_exposed( - &self.subagent, - &self.delegation_tools, - &self.tool_name, - ); - (!exposed.is_empty()).then(|| { - format!( - "You are a leaf agent: do this work yourself; do not call this tool again. (Its harness exposes delegation tools {exposed:?}, so `{}` was not started.)", - self.tool_name - ) - }) - } - - /// Construction-time visibility for a misconfigured leaf: the spawn is - /// refused per call, and this logs the cause once up front. - fn warn_if_leaf_misconfigured(&self) { - if let Some(message) = self.leaf_violation() { - tracing::warn!( - "{LOG_PREFIX} leaf_misconfigured tool={} {message}", - self.tool_name - ); - } - } - - /// Trims and schema-checks each child's final output. - pub fn with_result_policy(mut self, policy: crate::subagent::ResultPolicy) -> Self { - self.result_policy = policy; - self - } - - /// Enforces spawn limits through `admission` (see [`super::SpawnPolicy`]). - /// - /// Share one [`SpawnAdmission`] across every tool whose spawns should - /// count against the same limits. Without this call spawning is unlimited. - pub fn with_spawn_admission(mut self, admission: SpawnAdmission) -> Self { - self.admission = admission; - self - } - - /// Returns the admission ledger this tool reserves spawn slots from. - pub fn spawn_admission(&self) -> &SpawnAdmission { - &self.admission - } - - /// Uses a host-shared registry for spawned jobs and control tools. - pub fn with_job_registry(mut self, jobs: SubAgentJobRegistry) -> Self { - self.jobs = jobs.clone(); - self.job_registry = Some(jobs); - self - } - - /// Resolves ephemeral turn run ids to a stable host-owned job scope. - /// Configure the same resolver on this spawn tool and its Jobs/Message - /// aliases when the host recreates run contexts between turns. - pub fn with_job_owner_scope( - mut self, - resolve: Arc<dyn Fn(u64) -> String + Send + Sync>, - ) -> Self { - self.job_owner_scope = resolve; - self - } - - /// Returns the registry that owns jobs spawned by this tool. - pub fn job_registry(&self) -> &SubAgentJobRegistry { - &self.jobs - } - - /// Overrides the model-visible tool name. - pub fn with_tool_name(mut self, name: impl Into<String>) -> Self { - self.tool_name = name.into(); - self.declaration = std::sync::OnceLock::new(); - self - } - - /// Overrides the model-facing description for a named driver alias. - pub fn with_description(mut self, description: impl Into<String>) -> Self { - self.tool_description = description.into(); - self.declaration = std::sync::OnceLock::new(); - self - } - - /// Overrides the model-visible JSON Schema for the tool arguments. - pub fn with_parameters(mut self, parameters: Value) -> Self { - self.parameters = parameters; - self.declaration = std::sync::OnceLock::new(); - self - } - - /// Extracts the child input string from model-supplied `arguments`. - /// - /// Accepts either an object carrying a string [`SUBAGENT_INPUT_FIELD`] field - /// or a bare JSON string; anything else yields the empty string. - fn extract_input(arguments: &Value) -> String { - match arguments { - Value::String(s) => s.clone(), - Value::Object(map) => map - .get(SUBAGENT_INPUT_FIELD) - .and_then(Value::as_str) - .unwrap_or_default() - .to_owned(), - _ => String::new(), - } - } - - /// Reads the optional `mode` argument; absent means [`SubAgentMode::Background`]. - fn extract_mode(arguments: &Value) -> std::result::Result<SubAgentMode, String> { - match arguments.get(SUBAGENT_MODE_FIELD) { - None | Some(Value::Null) => Ok(SubAgentMode::Background), - Some(Value::String(mode)) if mode == "background" => Ok(SubAgentMode::Background), - Some(Value::String(mode)) if mode == "inline" => Ok(SubAgentMode::Inline), - Some(_) => Err(format!( - "`{SUBAGENT_MODE_FIELD}` must be \"background\" or \"inline\"" - )), - } - } - - /// Spawns this sub-agent from the actual parent [`RunContext`]. - /// - /// This is the agent-native recursive-tool boundary. It is intentionally - /// separate from `tinytools::Tool`: TinyTools only receives the narrow - /// workspace/thread/output vocabulary it needs for ordinary tools, while - /// a child agent must inherit the parent run's live lineage, cancellation, - /// stores, events, workspace, steering, and streaming state. The harness - /// tool dispatcher registers this typed entry point explicitly; it does - /// not downcast a generic tool or recover a parent from a global map. - pub async fn invoke_in_parent_context( - &self, - state: &State, - args: Value, - options: tinytools::ToolCallOptions, - parent: &RunContext<Ctx>, - ) -> Result<tinytools::ToolResult> { - self.invoke_in_parent_context_for_call(state, args, options, parent, None) - .await - } - - /// Like [`Self::invoke_in_parent_context`], additionally recording the - /// parent's `call_id` as the explicit parent-call -> child-run link. - /// - /// The link is carried by the queued result (`subagent_run_id`, - /// `parent_tool_call_id`, `job_id`), the job snapshot, and the child run's - /// metadata (`subagent_run_id`, `subagent_job_id`, `parent_tool_call_id`). - pub async fn invoke_in_parent_context_for_call( - &self, - state: &State, - args: Value, - _options: tinytools::ToolCallOptions, - parent: &RunContext<Ctx>, - call_id: Option<&tinyagents_harness::ids::CallId>, - ) -> Result<tinytools::ToolResult> { - let input = Self::extract_input(&args); - let mode = match Self::extract_mode(&args) { - Ok(mode) => mode, - Err(message) => { - tracing::debug!("{LOG_PREFIX} invalid_mode tool={}", self.tool_name); - return Ok(tinytools::ToolResult::error(message)); - } - }; - if let Some(message) = self.leaf_violation() { - tracing::debug!("{LOG_PREFIX} leaf_violation tool={}", self.tool_name); - return Ok(tinytools::ToolResult::error(message)); - } - // Reserve the slot atomically before anything is spawned. The guard - // refunds on every early return below; it is committed once the child - // is registered and then lives exactly as long as the child runs. - let mut reservation = match self - .admission - .try_reserve(&parent.config, self.subagent.name()) - { - Ok(reservation) => reservation, - Err(rejection) => { - tracing::debug!( - "{LOG_PREFIX} spawn_rejected tool={} reason={rejection}", - self.tool_name - ); - return Ok(tinytools::ToolResult::error(format!( - "Sub-agent `{}` was not started because a spawn limit was reached: {rejection}. The parent orchestrator should treat this as a delegated-agent limit signal, not a completed answer.", - self.tool_name - ))); - } - }; - let mut config = match self.subagent.child_config( - parent.depth(), - parent.thread_id(), - parent.config.max_turn_output_tokens, - Some((parent.run_id().as_str(), parent.next_child_ordinal())), - ) { - Ok(config) => config, - Err(error) => { - return Ok(tinytools::ToolResult::error(format!( - "Sub-agent `{}` stopped before completing because it hit its recursion depth limit: {error}. The parent orchestrator should treat this as a delegated-agent limit signal, not a completed answer.", - self.tool_name - ))); - } - }; - self.policy.budget.apply_call_caps(&mut config); - let retry_base = (self.policy.retry.max_attempts > 1).then(|| config.clone()); - let child_data = self.child_data.child_data(&parent.data); - let child = match parent.child(config, child_data) { - Ok(child) => child, - Err(TinyAgentsError::SubAgentDepth(_)) => { - return Ok(tinytools::ToolResult::error(format!( - "Sub-agent `{}` stopped before completing because it hit its recursion depth limit. The parent orchestrator should treat this as a delegated-agent limit signal, not a completed answer.", - self.tool_name - ))); - } - Err(error) => return Err(error), - }; - let subagent_run_id = child.run_id().as_str().to_owned(); - let tool_call_id = call_id.map(|id| id.as_str().to_owned()); - let run_queue = Arc::new(tinyagents_harness::run_queue::RunQueue::new()); - let (job_id, steering) = match self.jobs.create_with_cancellation_and_queue( - &self.tool_name, - parent.instance_id(), - child.cancellation.clone(), - JobLink { - subagent_run_id: Some(subagent_run_id.clone()), - parent_tool_call_id: tool_call_id.clone(), - }, - Some(run_queue.clone()), - ) { - Ok(created) => created, - Err(error) => { - return Ok(tinytools::ToolResult::error(format!( - "Sub-agent `{}` was not started: its job could not be registered ({error}).", - self.tool_name - ))); - } - }; - tracing::debug!( - "{LOG_PREFIX} spawn job_id={job_id} subagent_run_id={subagent_run_id} tool_call_id={tool_call_id:?}" - ); - let mut child = child - .with_steering(steering.clone()) - .with_run_queue(run_queue.clone()); - stamp_link_metadata( - &mut child.config.metadata, - &subagent_run_id, - job_id.as_str(), - tool_call_id.as_deref(), - ); - // Retries need fresh child contexts, and a background child outlives the - // borrow of `parent`, so the (cheap) contexts are built up front. Their - // ids derive from the first child's (`{first}-a{n}`), so no extra child - // ordinals are consumed and sibling ids do not depend on retries. - let attempt_count = self.policy.retry.max_attempts.max(1); - let watch = attempt_count > 1; - let job_token = child.cancellation.clone(); - let mut attempts = vec![policy_run::Attempt::new(child, watch)]; - if let Some(base) = retry_base { - for n in 1..attempt_count { - let spare_id = format!("{subagent_run_id}-a{n}"); - let mut config = base.clone(); - config.run_id = RunId::new(spare_id.clone()); - config.thread_id = parent - .thread_id() - .map(|thread| child_thread_id(thread, &spare_id)); - match parent.child(config, self.child_data.child_data(&parent.data)) { - Ok(spare) => { - let mut spare = spare - .with_cancellation(job_token.clone()) - .with_steering(steering.clone()); - stamp_link_metadata( - &mut spare.config.metadata, - &spare_id, - job_id.as_str(), - tool_call_id.as_deref(), - ); - attempts.push(policy_run::Attempt::new(spare, watch)); - } - Err(error) => { - tracing::debug!("{LOG_PREFIX} retry_context_unavailable error={error}"); - break; - } - } - } - } - reservation.commit(); - let streaming = parent.streaming; - if mode == SubAgentMode::Inline { - tracing::debug!("{LOG_PREFIX} inline.start job_id={job_id}"); - self.jobs.mark_running(&job_id); - let mut guard = jobs::InlineJobGuard::new(self.jobs.clone(), job_id.clone()); - let finished = policy_run::run_attempts( - &self.subagent, - &self.policy, - state, - attempts, - input, - streaming, - &job_token, - &self.jobs, - &job_id, - ) - .await; - policy_run::settle(&self.jobs, &job_id, finished, &self.result_policy).await; - guard.disarm(); - drop(reservation); - let Some(job) = self.jobs.get(job_id.as_str()) else { - return Ok(tinytools::ToolResult::error(format!( - "Sub-agent job `{job_id}` was removed before its result could be read." - ))); - }; - tracing::debug!( - "{LOG_PREFIX} inline.done job_id={job_id} status={:?}", - job.status - ); - // The job snapshot, plus the `job_id` key the queued result uses - // so both modes name the job the same way. - let mut payload = serde_json::to_value(&job)?; - payload["job_id"] = json!(job_id); - return Ok(match job.status { - SubAgentJobStatus::Completed => tinytools::ToolResult::json(payload), - _ => tinytools::ToolResult::error(payload.to_string()), - }); - } - let jobs = self.jobs.clone(); - let task_job_id = job_id.clone(); - let subagent = self.subagent.clone(); - let owned_state = state.clone(); - let policy = self.policy.clone(); - let result_policy = self.result_policy.clone(); - // The child runs in its own task and a supervisor awaits its - // `JoinHandle`: a panic inside the child surfaces as a `JoinError` - // there, so the job can never stay `Running` forever. - let child_job_id = task_job_id.clone(); - let child_jobs = jobs.clone(); - let child_task = tokio::spawn(async move { - // The slot is held for the child's whole lifetime and released - // when this task ends, however it ends (result, panic, abort). - let _reservation = reservation; - child_jobs.mark_running(&child_job_id); - let finished = policy_run::run_attempts( - &subagent, - &policy, - &owned_state, - attempts, - input, - streaming, - &job_token, - &child_jobs, - &child_job_id, - ) - .await; - policy_run::settle(&child_jobs, &child_job_id, finished, &result_policy).await; - }); - tokio::spawn(async move { - if let Err(join_error) = child_task.await { - jobs.mark_aborted(&task_job_id, join_error.is_panic()); - } - }); - - let mut queued = json!({ - "job_id": job_id, - "status": "queued", - "subagent_run_id": subagent_run_id, - }); - if let Some(tool_call_id) = tool_call_id { - queued["parent_tool_call_id"] = Value::String(tool_call_id); - } - Ok(tinytools::ToolResult::json(queued)) - } -} - -#[async_trait] -impl<State, Ctx> ToolDispatch<State, Ctx> for SubAgentTool<State, Ctx> -where - State: Clone + Send + Sync + 'static, - Ctx: Send + Sync + 'static, -{ - fn tool(&self) -> Arc<dyn tinytools::Tool> { - // Built once and cached (M-4): `tool()` is called several times per - // admitted call and once per tool per run for `schemas()`, and a - // fresh `Arc<SubAgentToolDeclaration>` with a cloned `parameters` - // `Value` on every call is unnecessary allocation for a declaration - // that never changes after registration. - Arc::clone(self.declaration.get_or_init(|| { - Arc::new(SubAgentToolDeclaration { - name: self.tool_name.clone(), - description: self.tool_description.clone(), - parameters: self.parameters.clone(), - }) - })) - } - - fn output_origin(&self) -> tinyagents_harness::host::ContentOrigin { - tinyagents_harness::host::ContentOrigin::Agent - } - - async fn execute( - &self, - state: &State, - call_id: tinyagents_harness::ids::CallId, - arguments: Value, - options: tinytools::ToolCallOptions, - parent: &RunContext<Ctx>, - ) -> anyhow::Result<tinytools::ToolResult> { - self.invoke_in_parent_context_for_call(state, arguments, options, parent, Some(&call_id)) - .await - .map_err(anyhow::Error::from) - } -} - -/// Pure canonical declaration for a typed-parent sub-agent dispatcher. -/// -/// Calling it through `tinytools::Tool` directly is refused because that trait -/// intentionally lacks the parent `RunContext`; register the enclosing -/// [`SubAgentTool`] with [`tinyagents_harness::tool::ToolRegistry::register_dispatch`]. -struct SubAgentToolDeclaration { - name: String, - description: String, - parameters: Value, -} - -#[async_trait] -impl tinytools::Tool for SubAgentToolDeclaration { - fn name(&self) -> &str { - &self.name - } - - fn description(&self) -> &str { - &self.description - } - - fn parameters_schema(&self) -> Value { - self.parameters.clone() - } - - async fn execute(&self, _args: Value) -> anyhow::Result<tinytools::ToolResult> { - anyhow::bail!( - "sub-agent `{}` requires typed-parent dispatch; register SubAgentTool with ToolRegistry::register_dispatch", - self.name - ) - } -} - -#[cfg(test)] -#[path = "mod_tests.rs"] -mod test; - -#[cfg(test)] -#[path = "test_helpers.rs"] -mod test_helpers; - -#[cfg(test)] -#[path = "mod_link_tests.rs"] -mod link_test; - -#[cfg(test)] -#[path = "mod_inline_tests.rs"] -mod inline_test; - -#[cfg(test)] -#[path = "mod_admission_tests.rs"] -mod admission_test; - -#[cfg(test)] -#[path = "mod_policy_tests.rs"] -mod policy_test; - #[cfg(test)] #[path = "mod_registry_tests.rs"] mod registry_tests; diff --git a/crates/tinyagents-orchestration/src/subagent/invocation/mod_admission_tests.rs b/crates/tinyagents-orchestration/src/subagent/invocation/mod_admission_tests.rs deleted file mode 100644 index a52167e5c..000000000 --- a/crates/tinyagents-orchestration/src/subagent/invocation/mod_admission_tests.rs +++ /dev/null @@ -1,289 +0,0 @@ -//! Spawn-admission contracts for [`SubAgentTool`]: caps are enforced for both -//! modes, concurrent spawns cannot race past a cap, and slots are returned -//! when a child finishes or its spawn fails. - -use super::*; -use std::sync::Arc; -use std::time::Duration; - -use serde_json::{Value, json}; - -use super::test::BlockedModel; -use super::test_helpers::PanickingModel; -use crate::subagent::{SpawnAdmission, SpawnPolicy}; -use tinyagents_harness::context::{RunConfig, RunContext}; -use tinyagents_harness::runtime::AgentHarness; -use tinyinference_llm::providers::MockModel; - -type Gate = Arc<tokio::sync::Semaphore>; - -fn policy(parent: Option<usize>, root: Option<usize>, targets: Option<&[&str]>) -> SpawnPolicy { - SpawnPolicy { - max_children_per_parent: parent, - max_total_per_root: root, - allowed_targets: targets.map(|t| t.iter().map(|s| (*s).to_owned()).collect()), - } -} - -fn blocked_tool(policy: SpawnPolicy) -> (Arc<SubAgentTool<(), ()>>, Gate, Gate) { - let started: Gate = Arc::new(tokio::sync::Semaphore::new(0)); - let release: Gate = Arc::new(tokio::sync::Semaphore::new(0)); - let mut harness = AgentHarness::new(); - harness.register_model( - "worker", - Arc::new(BlockedModel { - started: started.clone(), - release: release.clone(), - }), - ); - let tool = SubAgentTool::new( - Arc::new(SubAgent::new("worker", "works", Arc::new(harness))), - ChildDataPolicy::new(|_: &()| ()), - ) - .with_spawn_admission(SpawnAdmission::new(policy)); - (Arc::new(tool), started, release) -} - -fn constant_tool(policy: SpawnPolicy) -> SubAgentTool<(), ()> { - let mut harness = AgentHarness::new(); - harness.register_model("child", Arc::new(MockModel::constant("done"))); - SubAgentTool::new( - Arc::new(SubAgent::new("worker", "works", Arc::new(harness))), - ChildDataPolicy::new(|_: &()| ()), - ) - .with_spawn_admission(SpawnAdmission::new(policy)) -} - -async fn call( - tool: &SubAgentTool<(), ()>, - parent: &RunContext<()>, - args: Value, -) -> tinytools::ToolResult { - tool.invoke_in_parent_context(&(), args, tinytools::ToolCallOptions::default(), parent) - .await - .expect("tool call returns a result") -} - -fn parent() -> RunContext<()> { - RunContext::new(RunConfig::new("parent"), ()) -} - -fn assert_limit_signal(result: &tinytools::ToolResult) { - assert!(result.is_error, "{}", result.output()); - assert!( - result - .output() - .contains("limit signal, not a completed answer"), - "{}", - result.output() - ); -} - -async fn slots_settle_to(tool: &SubAgentTool<(), ()>, parent: &str, expected: usize) { - tokio::time::timeout(Duration::from_secs(2), async { - while tool.spawn_admission().active_children(parent) != expected { - tokio::time::sleep(Duration::from_millis(5)).await; - } - }) - .await - .expect("slots settle"); -} - -#[tokio::test] -async fn background_spawn_over_the_parent_cap_is_a_limit_signal() { - let (tool, started, release) = blocked_tool(policy(Some(1), None, None)); - let parent = parent(); - - let first = call(&tool, &parent, json!({"input": "a"})).await; - assert!(!first.is_error); - let second = call(&tool, &parent, json!({"input": "b"})).await; - - assert_limit_signal(&second); - assert!(second.output().contains("1/1"), "{}", second.output()); - assert_eq!(tool.job_registry().list().len(), 1, "no job was created"); - - // The slot returns when the child reaches a terminal state. - let _started = started.acquire().await.unwrap(); - release.add_permits(1); - slots_settle_to(&tool, "run:parent", 0).await; - let third = call(&tool, &parent, json!({"input": "c"})).await; - assert!(!third.is_error, "{}", third.output()); - release.add_permits(1); -} - -#[tokio::test] -async fn inline_spawn_over_the_parent_cap_is_a_limit_signal() { - let (tool, _started, release) = blocked_tool(policy(Some(1), None, None)); - let parent = parent(); - let held = call(&tool, &parent, json!({"input": "a"})).await; - assert!(!held.is_error); - - let inline = call(&tool, &parent, json!({"input": "b", "mode": "inline"})).await; - - assert_limit_signal(&inline); - assert_eq!(tool.job_registry().list().len(), 1); - release.add_permits(1); -} - -#[tokio::test] -async fn inline_slot_is_released_when_the_call_returns() { - let tool = constant_tool(policy(Some(1), None, None)); - let parent = parent(); - for _ in 0..3 { - let result = call(&tool, &parent, json!({"input": "x", "mode": "inline"})).await; - assert!(!result.is_error, "{}", result.output()); - } - assert_eq!(tool.spawn_admission().active_children("run:parent"), 0); -} - -#[tokio::test] -async fn total_budget_per_root_is_spent_even_by_finished_children() { - let tool = constant_tool(policy(None, Some(2), None)); - let parent = parent(); - for _ in 0..2 { - let result = call(&tool, &parent, json!({"input": "x", "mode": "inline"})).await; - assert!(!result.is_error); - } - let over = call(&tool, &parent, json!({"input": "x", "mode": "inline"})).await; - assert_limit_signal(&over); - assert!(over.output().contains("2/2"), "{}", over.output()); -} - -#[tokio::test] -async fn disallowed_target_is_refused_before_spawning() { - let tool = constant_tool(policy(None, None, Some(&["researcher"]))); - let result = call(&tool, &parent(), json!({"input": "x"})).await; - assert_limit_signal(&result); - assert!(result.output().contains("`worker`"), "{}", result.output()); - assert!(tool.job_registry().list().is_empty()); -} - -#[tokio::test] -async fn allowed_target_spawns() { - let tool = constant_tool(policy(None, None, Some(&["worker"]))); - let result = call(&tool, &parent(), json!({"input": "x", "mode": "inline"})).await; - assert!(!result.is_error, "{}", result.output()); -} - -#[tokio::test] -async fn failed_spawn_refunds_the_reservation() { - let tool = constant_tool(policy(Some(1), Some(1), None)); - // Depth limit hit after admission: nothing spawns, nothing stays reserved. - let too_deep = RunContext::new(RunConfig::new("parent").with_max_depth(0), ()); - let result = call(&tool, &too_deep, json!({"input": "x"})).await; - assert_limit_signal(&result); - assert_eq!(tool.spawn_admission().active_children("run:parent"), 0); - assert_eq!(tool.spawn_admission().spawned_in_scope("run:parent"), 0); - - let ok = call(&tool, &parent(), json!({"input": "x", "mode": "inline"})).await; - assert!(!ok.is_error, "{}", ok.output()); -} - -#[tokio::test(flavor = "multi_thread", worker_threads = 4)] -async fn concurrent_background_spawns_cannot_exceed_the_cap() { - let (tool, _started, release) = blocked_tool(policy(Some(3), None, None)); - let parent = Arc::new(parent()); - let tasks: Vec<_> = (0..12) - .map(|_| { - let tool = tool.clone(); - let parent = parent.clone(); - tokio::spawn(async move { call(&tool, &parent, json!({"input": "x"})).await }) - }) - .collect(); - let mut admitted = 0; - for task in tasks { - if !task.await.unwrap().is_error { - admitted += 1; - } - } - assert_eq!(admitted, 3); - assert_eq!(tool.job_registry().list().len(), 3); - release.add_permits(3); -} - -#[tokio::test] -async fn shared_admission_counts_across_tools() { - let shared = SpawnAdmission::new(policy(Some(1), None, None)); - let (tool_a, _sa, release_a) = blocked_tool(SpawnPolicy::default()); - let (tool_b, _sb, release_b) = blocked_tool(SpawnPolicy::default()); - let tool_a = Arc::new( - Arc::into_inner(tool_a) - .unwrap() - .with_spawn_admission(shared.clone()), - ); - let tool_b = Arc::new( - Arc::into_inner(tool_b) - .unwrap() - .with_spawn_admission(shared), - ); - let parent = parent(); - - assert!(!call(&tool_a, &parent, json!({"input": "a"})).await.is_error); - assert_limit_signal(&call(&tool_b, &parent, json!({"input": "b"})).await); - release_a.add_permits(1); - release_b.add_permits(1); -} - -fn parent_on_thread(run_id: &str, thread: &str) -> RunContext<()> { - RunContext::new(RunConfig::new(run_id).with_thread(thread), ()) -} - -#[tokio::test] -async fn a_later_turn_shares_the_cap_with_a_background_child_still_alive() { - let (tool, started, release) = blocked_tool(policy(Some(1), None, None)); - - // Turn 1 spawns a background child and ends; turn 2 has a fresh run id on - // the same conversation thread. - let turn_one = parent_on_thread("run-turn-1", "thread-A"); - assert!(!call(&tool, &turn_one, json!({"input": "a"})).await.is_error); - let _started = started.acquire().await.unwrap(); - drop(turn_one); - - let turn_two = parent_on_thread("run-turn-2", "thread-A"); - let blocked = call(&tool, &turn_two, json!({"input": "b"})).await; - assert_limit_signal(&blocked); - assert_eq!(tool.spawn_admission().active_children("thread:thread-A"), 1); - - // A different conversation is unaffected. - let other = parent_on_thread("run-turn-3", "thread-B"); - assert!(!call(&tool, &other, json!({"input": "c"})).await.is_error); - - release.add_permits(2); -} - -#[tokio::test] -async fn dropping_an_inline_call_releases_its_slot() { - let (tool, started, _release) = blocked_tool(policy(Some(1), Some(1), None)); - let parent = parent(); - - let abandoned = tokio::time::timeout( - Duration::from_millis(50), - call(&tool, &parent, json!({"input": "a", "mode": "inline"})), - ) - .await; - assert!(abandoned.is_err(), "the blocked inline call timed out"); - let _started = started.acquire().await.unwrap(); - - // The dropped future released the live slot, and refunded nothing it - // already spent: the child had started, so the budget stays consumed. - assert_eq!(tool.spawn_admission().active_children("run:parent"), 0); - assert_eq!(tool.spawn_admission().spawned_in_scope("run:parent"), 1); -} - -#[tokio::test] -async fn a_panicking_background_child_releases_its_slot() { - let mut harness = AgentHarness::new(); - harness.register_model("worker", Arc::new(PanickingModel)); - let tool = SubAgentTool::new( - Arc::new(SubAgent::new("worker", "works", Arc::new(harness))), - ChildDataPolicy::new(|_: &()| ()), - ) - .with_spawn_admission(SpawnAdmission::new(policy(Some(1), None, None))); - let parent = parent(); - - assert!(!call(&tool, &parent, json!({"input": "a"})).await.is_error); - slots_settle_to(&tool, "run:parent", 0).await; - - let again = call(&tool, &parent, json!({"input": "b"})).await; - assert!(!again.is_error, "{}", again.output()); -} diff --git a/crates/tinyagents-orchestration/src/subagent/invocation/mod_inline_tests.rs b/crates/tinyagents-orchestration/src/subagent/invocation/mod_inline_tests.rs deleted file mode 100644 index 1800c67ee..000000000 --- a/crates/tinyagents-orchestration/src/subagent/invocation/mod_inline_tests.rs +++ /dev/null @@ -1,240 +0,0 @@ -//! `mode: "inline"` contracts for [`SubAgentTool`]: the child is awaited in -//! the same tool call and its final result is returned directly. - -use super::*; -use std::sync::Arc; -use std::time::Duration; - -use serde_json::{Value, json}; - -use super::test::BlockedModel; -use super::test_helpers::PanickingModel; -use tinyagents_harness::cancel::CancellationToken; -use tinyagents_harness::context::{RunConfig, RunContext}; -use tinyagents_harness::ids::CallId; -use tinyagents_harness::runtime::AgentHarness; -use tinyagents_harness::tool::ToolDispatch; -use tinyinference_llm::providers::MockModel; - -fn constant_tool(answer: &str) -> SubAgentTool<(), ()> { - let mut harness = AgentHarness::new(); - harness.register_model("child", Arc::new(MockModel::constant(answer))); - SubAgentTool::new( - Arc::new(SubAgent::new("worker", "works", Arc::new(harness))), - ChildDataPolicy::new(|_: &()| ()), - ) -} - -async fn call( - tool: &SubAgentTool<(), ()>, - parent: &RunContext<()>, - args: Value, -) -> tinytools::ToolResult { - ToolDispatch::<(), ()>::execute( - tool, - &(), - CallId::new("call-inline"), - args, - tinytools::ToolCallOptions::default(), - parent, - ) - .await - .expect("tool call returns a result") -} - -fn json_of(result: &tinytools::ToolResult) -> Value { - serde_json::from_str(&result.output()).expect("tool result is JSON") -} - -#[tokio::test] -async fn inline_mode_returns_the_final_output_in_the_same_call() { - let tool = constant_tool("inline answer"); - let parent = RunContext::new(RunConfig::new("parent"), ()); - - let result = call(&tool, &parent, json!({"input": "work", "mode": "inline"})).await; - - assert!(!result.is_error); - let payload = json_of(&result); - assert_eq!(payload["status"], "completed"); - assert_eq!(payload["output"], "inline answer"); - assert_eq!(payload["parent_tool_call_id"], "call-inline"); - let job_id = payload["job_id"].as_str().expect("job id"); - assert!( - payload["subagent_run_id"] - .as_str() - .is_some_and(|id| id.starts_with("worker-d1-parent-")) - ); - let job = tool - .job_registry() - .get_owned(job_id, parent.instance_id()) - .expect("inline run is registered"); - assert_eq!(job.status, SubAgentJobStatus::Completed); -} - -#[tokio::test] -async fn background_remains_the_default_mode() { - let tool = constant_tool("later"); - let parent = RunContext::new(RunConfig::new("parent"), ()); - - for args in [ - json!({"input": "work"}), - json!({"input": "work", "mode": "background"}), - ] { - let payload = json_of(&call(&tool, &parent, args).await); - assert_eq!(payload["status"], "queued"); - assert!(payload.get("output").is_none()); - } -} - -#[tokio::test] -async fn unknown_mode_is_a_recoverable_tool_error() { - let tool = constant_tool("unused"); - let parent = RunContext::new(RunConfig::new("parent"), ()); - - let result = call(&tool, &parent, json!({"input": "work", "mode": "sideways"})).await; - - assert!(result.is_error); - assert!(result.output().contains("mode")); - assert!(tool.job_registry().list().is_empty(), "nothing was spawned"); -} - -#[tokio::test] -async fn inline_mode_enforces_the_depth_limit_like_background() { - let tool = constant_tool("unused"); - let parent = RunContext::new(RunConfig::new("parent").with_max_depth(0), ()); - - let result = call(&tool, &parent, json!({"input": "work", "mode": "inline"})).await; - - assert!(result.is_error); - assert!(result.output().contains("delegated-agent limit signal")); -} - -#[tokio::test] -async fn inline_mode_reports_a_cancelled_child() { - let tool = constant_tool("unused"); - let cancellation = CancellationToken::new(); - cancellation.cancel(); - let parent = RunContext::new(RunConfig::new("parent"), ()).with_cancellation(cancellation); - - let result = call(&tool, &parent, json!({"input": "work", "mode": "inline"})).await; - - assert!(result.is_error); - assert_eq!(json_of(&result)["status"], "cancelled"); -} - -#[tokio::test] -async fn inline_mode_observes_parent_cancellation_mid_run() { - let started = Arc::new(tokio::sync::Semaphore::new(0)); - let release = Arc::new(tokio::sync::Semaphore::new(0)); - let mut harness = AgentHarness::new(); - harness.register_model( - "worker", - Arc::new(BlockedModel { - started: started.clone(), - release, - }), - ); - let tool = SubAgentTool::new( - Arc::new(SubAgent::new("worker", "works", Arc::new(harness))), - ChildDataPolicy::new(|_: &()| ()), - ); - let parent = RunContext::new(RunConfig::new("parent"), ()); - let cancellation = parent.cancellation.clone(); - - let (result, ()) = tokio::join!( - call(&tool, &parent, json!({"input": "work", "mode": "inline"})), - async { - let _permit = started.acquire().await.unwrap(); - cancellation.cancel(); - } - ); - - assert!(result.is_error); - assert_eq!(json_of(&result)["status"], "cancelled"); -} - -#[test] -fn schema_advertises_the_mode_argument() { - let tool = constant_tool("unused"); - let schema = ToolDispatch::<(), ()>::tool(&tool).parameters_schema(); - assert_eq!( - schema["properties"]["mode"]["enum"], - json!(["background", "inline"]) - ); - assert_eq!(schema["required"], json!(["input"])); -} - -#[tokio::test] -async fn dropping_an_inline_call_settles_its_job() { - let started = Arc::new(tokio::sync::Semaphore::new(0)); - let mut harness = AgentHarness::new(); - harness.register_model( - "worker", - Arc::new(BlockedModel { - started: started.clone(), - release: Arc::new(tokio::sync::Semaphore::new(0)), - }), - ); - let tool = SubAgentTool::new( - Arc::new(SubAgent::new("worker", "works", Arc::new(harness))), - ChildDataPolicy::new(|_: &()| ()), - ); - let parent = RunContext::new(RunConfig::new("parent"), ()); - - // The harness drops the tool future on a tool timeout or stream drop; - // model that by abandoning the call once the child is running. - tokio::select! { - _ = call(&tool, &parent, json!({"input": "work", "mode": "inline"})) => { - panic!("the blocked child cannot finish") - } - _ = started.acquire() => {} - } - - tokio::time::timeout(Duration::from_secs(1), async { - loop { - if tool - .job_registry() - .list() - .first() - .is_some_and(|job| job.status.is_terminal()) - { - break; - } - tokio::task::yield_now().await; - } - }) - .await - .expect("dropped inline job settles"); - let jobs = tool.job_registry().list(); - assert_eq!(jobs.len(), 1); - assert!( - jobs[0].status.is_terminal(), - "dropped inline job must not stay running: {:?}", - jobs[0].status - ); -} - -#[tokio::test] -async fn a_panicking_inline_child_marks_its_job_failed() { - let mut harness = AgentHarness::new(); - harness.register_model("worker", Arc::new(PanickingModel)); - let tool = Arc::new(SubAgentTool::new( - Arc::new(SubAgent::new("worker", "works", Arc::new(harness))), - ChildDataPolicy::new(|_: &()| ()), - )); - let task_tool = tool.clone(); - let joined = tokio::spawn(async move { - let parent = RunContext::new(RunConfig::new("parent"), ()); - call( - &task_tool, - &parent, - json!({"input": "work", "mode": "inline"}), - ) - .await - }) - .await; - assert!(joined.expect_err("the panic propagates").is_panic()); - - let jobs = tool.job_registry().list(); - assert_eq!(jobs[0].status, SubAgentJobStatus::Failed); -} diff --git a/crates/tinyagents-orchestration/src/subagent/invocation/mod_link_tests.rs b/crates/tinyagents-orchestration/src/subagent/invocation/mod_link_tests.rs deleted file mode 100644 index 7776d48da..000000000 --- a/crates/tinyagents-orchestration/src/subagent/invocation/mod_link_tests.rs +++ /dev/null @@ -1,219 +0,0 @@ -//! Explicit parent-call -> child-run link carried by sub-agent results. - -use super::*; -use std::sync::{Arc, Mutex}; - -use serde_json::{Value, json}; - -use super::test::wait_for_terminal; -use tinyagents_harness::context::{RunConfig, RunContext}; -use tinyagents_harness::error::Result; -use tinyagents_harness::events::{AgentEvent, EventSink, RecordingListener}; -use tinyagents_harness::ids::CallId; -use tinyagents_harness::middleware::Middleware; -use tinyagents_harness::runtime::AgentHarness; -use tinyagents_harness::tool::ToolDispatch; -use tinyinference_llm::providers::MockModel; - -fn tool() -> SubAgentTool<(), ()> { - let mut harness = AgentHarness::new(); - harness.register_model("child", Arc::new(MockModel::constant("done"))); - SubAgentTool::new( - Arc::new(SubAgent::new("worker", "works", Arc::new(harness))), - ChildDataPolicy::new(|_: &()| ()), - ) -} - -fn json_of(result: &tinytools::ToolResult) -> Value { - serde_json::from_str(&result.output()).expect("tool result is JSON") -} - -#[tokio::test] -async fn queued_result_and_completed_job_carry_the_explicit_link() { - let tool = tool(); - let jobs = tool.job_registry().clone(); - let events = EventSink::new(); - let recorder = Arc::new(RecordingListener::new()); - events.subscribe(recorder.clone()); - let parent = RunContext::new(RunConfig::new("parent"), ()).with_events(events); - - let result = ToolDispatch::<(), ()>::execute( - &tool, - &(), - CallId::new("call-42"), - json!({"input": "work"}), - tinytools::ToolCallOptions::default(), - &parent, - ) - .await - .unwrap(); - let queued = json_of(&result); - let job_id = queued["job_id"].as_str().unwrap().to_owned(); - let run_id = queued["subagent_run_id"] - .as_str() - .expect("run id in queued result"); - assert!(run_id.starts_with("worker-d1-parent-"), "run id: {run_id}"); - assert_eq!(queued["parent_tool_call_id"], "call-42"); - - let job = wait_for_terminal(&jobs, &job_id, parent.instance_id()).await; - assert_eq!(job.subagent_run_id.as_deref(), Some(run_id)); - assert_eq!(job.parent_tool_call_id.as_deref(), Some("call-42")); - let completed = serde_json::to_value(&job).unwrap(); - assert_eq!(completed["subagent_run_id"], run_id); - assert_eq!(completed["parent_tool_call_id"], "call-42"); - - // The child run itself is stamped with the same ids. - let started = recorder - .events() - .iter() - .any(|record| matches!(&record.event, AgentEvent::RunStarted { run_id: id, .. } if id.as_str() == run_id)); - assert!( - started, - "child run id matches the advertised subagent_run_id" - ); -} - -struct MetadataProbe(Arc<Mutex<Option<Value>>>); - -#[async_trait::async_trait] -impl Middleware<(), ()> for MetadataProbe { - fn name(&self) -> &str { - "metadata-probe" - } - - async fn before_agent(&self, ctx: &mut RunContext<()>, _state: &()) -> Result<()> { - *self.0.lock().unwrap() = Some(ctx.config.metadata.clone()); - Ok(()) - } -} - -#[tokio::test] -async fn child_run_metadata_names_the_parent_call_and_job() { - let seen = Arc::new(Mutex::new(None)); - let mut harness = AgentHarness::new(); - harness.register_model("child", Arc::new(MockModel::constant("done"))); - harness.push_middleware(Arc::new(MetadataProbe(seen.clone()))); - let tool = SubAgentTool::new( - Arc::new(SubAgent::new("worker", "works", Arc::new(harness))), - ChildDataPolicy::new(|_: &()| ()), - ); - let jobs = tool.job_registry().clone(); - let parent = RunContext::new( - RunConfig::new("parent").with_metadata(json!({"keep": 1})), - (), - ); - let result = ToolDispatch::<(), ()>::execute( - &tool, - &(), - CallId::new("call-7"), - json!({"input": "work"}), - tinytools::ToolCallOptions::default(), - &parent, - ) - .await - .unwrap(); - let queued = json_of(&result); - wait_for_terminal( - &jobs, - queued["job_id"].as_str().unwrap(), - parent.instance_id(), - ) - .await; - - let metadata = seen.lock().unwrap().clone().expect("child ran"); - assert_eq!(metadata["keep"], 1, "parent metadata is preserved"); - assert_eq!(metadata["parent_tool_call_id"], "call-7"); - assert_eq!(metadata["subagent_job_id"], queued["job_id"]); - assert_eq!(metadata["subagent_run_id"], queued["subagent_run_id"]); -} - -#[tokio::test] -async fn child_run_metadata_preserves_non_object_parent_metadata() { - let seen = Arc::new(Mutex::new(None)); - let mut harness = AgentHarness::new(); - harness.register_model("child", Arc::new(MockModel::constant("done"))); - harness.push_middleware(Arc::new(MetadataProbe(seen.clone()))); - let tool = SubAgentTool::new( - Arc::new(SubAgent::new("worker", "works", Arc::new(harness))), - ChildDataPolicy::new(|_: &()| ()), - ); - let jobs = tool.job_registry().clone(); - let parent = RunContext::new( - RunConfig::new("parent").with_metadata(json!(["keep", 1])), - (), - ); - let result = ToolDispatch::<(), ()>::execute( - &tool, - &(), - CallId::new("call-array"), - json!({"input": "work"}), - tinytools::ToolCallOptions::default(), - &parent, - ) - .await - .unwrap(); - let queued = json_of(&result); - wait_for_terminal( - &jobs, - queued["job_id"].as_str().unwrap(), - parent.instance_id(), - ) - .await; - - let metadata = seen.lock().unwrap().clone().expect("child ran"); - assert_eq!(metadata["value"], json!(["keep", 1])); - assert_eq!(metadata["parent_tool_call_id"], "call-array"); -} - -#[tokio::test] -async fn tool_call_id_is_omitted_when_the_caller_has_none() { - let tool = tool(); - let parent = RunContext::new(RunConfig::new("parent"), ()); - let result = tool - .invoke_in_parent_context( - &(), - json!({"input": "work"}), - tinytools::ToolCallOptions::default(), - &parent, - ) - .await - .unwrap(); - let queued = json_of(&result); - assert!(queued.get("parent_tool_call_id").is_none()); - assert!(queued.get("subagent_run_id").is_some()); -} - -/// The session view recovers the run id from the child's thread id, so the -/// harness-derived thread id must keep embedding it. -#[tokio::test] -async fn child_thread_id_embeds_the_advertised_run_id() { - let tool = tool(); - let events = EventSink::new(); - let recorder = Arc::new(RecordingListener::new()); - events.subscribe(recorder.clone()); - let jobs = tool.job_registry().clone(); - let parent = - RunContext::new(RunConfig::new("parent").with_thread("thr-1"), ()).with_events(events); - let result = tool - .invoke_in_parent_context( - &(), - json!({"input": "work"}), - tinytools::ToolCallOptions::default(), - &parent, - ) - .await - .unwrap(); - let queued = json_of(&result); - wait_for_terminal( - &jobs, - queued["job_id"].as_str().unwrap(), - parent.instance_id(), - ) - .await; - let run_id = queued["subagent_run_id"].as_str().unwrap(); - let expected = format!("thr-1-subagent-{run_id}"); - assert!(recorder.events().iter().any(|record| matches!( - &record.event, - AgentEvent::RunStarted { thread_id: Some(thread), .. } if thread.as_str() == expected - ))); -} diff --git a/crates/tinyagents-orchestration/src/subagent/invocation/mod_policy_tests.rs b/crates/tinyagents-orchestration/src/subagent/invocation/mod_policy_tests.rs deleted file mode 100644 index 31ae701ef..000000000 --- a/crates/tinyagents-orchestration/src/subagent/invocation/mod_policy_tests.rs +++ /dev/null @@ -1,481 +0,0 @@ -//! Policy contracts for [`SubAgentTool`]: timeout, retry, result policy, role. - -use super::*; -use std::sync::Arc; -use std::sync::atomic::{AtomicUsize, Ordering}; -use std::time::Duration; - -use serde_json::{Value, json}; - -use super::test::BlockedModel; -use crate::subagent::{IncompleteKind, ResultPolicy, SubAgentPolicy, SubagentRole}; -use tinyagents_harness::context::{RunConfig, RunContext}; -use tinyagents_harness::ids::CallId; -use tinyagents_harness::retry::RetryPolicy; -use tinyagents_harness::runtime::{AgentHarness, RunPolicy}; -use tinyagents_harness::tool::ToolDispatch; -use tinyinference_llm::model::{ChatModel, ModelRequest, ModelResponse}; -use tinyinference_llm::providers::MockModel; - -fn tool_over(harness: AgentHarness<(), ()>) -> SubAgentTool<(), ()> { - SubAgentTool::new( - Arc::new(SubAgent::new("worker", "works", Arc::new(harness))), - ChildDataPolicy::new(|_: &()| ()), - ) -} - -fn constant(answer: &str) -> AgentHarness<(), ()> { - let mut harness = AgentHarness::new(); - harness.register_model("child", Arc::new(MockModel::constant(answer))); - harness -} - -async fn call_inline(tool: &SubAgentTool<(), ()>) -> tinytools::ToolResult { - let parent = RunContext::new(RunConfig::new("parent"), ()); - ToolDispatch::<(), ()>::execute( - tool, - &(), - CallId::new("c"), - json!({"input": "work", "mode": "inline"}), - tinytools::ToolCallOptions::default(), - &parent, - ) - .await - .expect("tool call returns a result") -} - -fn payload(result: &tinytools::ToolResult) -> Value { - serde_json::from_str(&result.output()).expect("JSON payload") -} - -#[tokio::test] -async fn timeout_cancels_the_child_and_marks_the_job_incomplete() { - let mut harness = AgentHarness::new(); - harness.register_model( - "child", - Arc::new(BlockedModel { - started: Arc::new(tokio::sync::Semaphore::new(0)), - release: Arc::new(tokio::sync::Semaphore::new(0)), - }), - ); - let tool = tool_over(harness) - .with_policy(SubAgentPolicy::default().with_timeout(Duration::from_millis(50))); - let result = call_inline(&tool).await; - assert!(result.is_error); - let payload = payload(&result); - assert_eq!(payload["status"], "incomplete"); - assert_eq!(payload["incomplete_kind"], "timeout"); - let job = &tool.job_registry().list()[0]; - assert_eq!(job.status, SubAgentJobStatus::Incomplete); - assert!(job.status.is_terminal()); -} - -struct FailingModel(Arc<AtomicUsize>); - -#[async_trait::async_trait] -impl ChatModel<()> for FailingModel { - async fn invoke(&self, _: &(), _: ModelRequest) -> tinyinference_llm::Result<ModelResponse> { - self.0.fetch_add(1, Ordering::SeqCst); - Err(tinyinference_llm::Error::Model( - "connection reset by peer".into(), - )) - } -} - -async fn failing_calls(max_attempts: usize) -> usize { - let calls = Arc::new(AtomicUsize::new(0)); - let mut harness = AgentHarness::new(); - harness.register_model("child", Arc::new(FailingModel(calls.clone()))); - harness.with_policy(RunPolicy { - retry: RetryPolicy::default().with_max_attempts(1), - ..RunPolicy::default() - }); - let tool = tool_over(harness).with_policy( - SubAgentPolicy::default().with_retry( - RetryPolicy::default() - .with_max_attempts(max_attempts) - .with_backoff_sleep(false), - ), - ); - let result = call_inline(&tool).await; - assert!(result.is_error); - calls.load(Ordering::SeqCst) -} - -#[tokio::test] -async fn a_retryable_failure_that_ran_no_tools_is_retried_with_a_fresh_child() { - let once = failing_calls(1).await; - let thrice = failing_calls(3).await; - assert!(once >= 1); - assert_eq!(thrice, once * 3, "three attempts, each a fresh child run"); -} - -#[tokio::test] -async fn result_policy_trims_the_job_output_and_reports_schema_errors() { - let tool = tool_over(constant(&"0123456789".repeat(20))).with_result_policy( - ResultPolicy::new() - .with_max_chars(100) - .with_schema(json!({"type": "object"})), - ); - let result = call_inline(&tool).await; - assert!( - !result.is_error, - "a schema mismatch is reported, not failed" - ); - let payload = payload(&result); - assert_eq!(payload["status"], "completed"); - assert!( - payload["output"] - .as_str() - .unwrap() - .contains("chars omitted") - ); - assert!( - payload["schema_error"] - .as_str() - .unwrap() - .contains("not valid JSON") - ); -} - -#[tokio::test] -async fn default_policies_leave_the_output_untouched() { - let result = call_inline(&tool_over(constant("full text"))).await; - let payload = payload(&result); - assert_eq!(payload["output"], "full text"); - assert!(payload.get("schema_error").is_none() && payload.get("artifacts").is_none()); -} - -#[tokio::test] -async fn a_leaf_refuses_to_spawn_when_its_harness_exposes_delegation_tools() { - let mut harness = constant("x"); - harness.register_tool(Arc::new(ExposedDelegationTool)); - let tool = tool_over(harness).with_role(SubagentRole::Leaf); - let result = call_inline(&tool).await; - assert!(result.is_error); - assert!(result.output().contains("leaf")); - assert!(tool.job_registry().list().is_empty(), "nothing was spawned"); - - let ok = tool_over(constant("x")).with_role(SubagentRole::Leaf); - assert!(!call_inline(&ok).await.is_error); -} - -struct ExposedDelegationTool; - -#[async_trait::async_trait] -impl tinytools::Tool for ExposedDelegationTool { - fn name(&self) -> &str { - crate::subagent::SUBAGENT_JOBS_TOOL - } - - fn description(&self) -> &str { - "test delegation tool" - } - - fn parameters_schema(&self) -> Value { - json!({"type": "object"}) - } - - async fn execute(&self, _args: Value) -> anyhow::Result<tinytools::ToolResult> { - Ok(tinytools::ToolResult::success("unused")) - } -} - -#[test] -fn incomplete_kind_serializes_as_snake_case() { - assert_eq!( - serde_json::to_value(IncompleteKind::BudgetExceeded).unwrap(), - "budget_exceeded" - ); -} - -// ---- attempt ids, tool safety, background retry, token overrun, artifacts ---- - -use tinyagents_harness::cancel::CancellationToken; -use tinyinference_llm::usage::Usage; - -struct ProbeTool { - token: Arc<std::sync::Mutex<Option<CancellationToken>>>, - hang: bool, -} - -struct ProbeDeclaration; - -#[async_trait::async_trait] -impl tinytools::Tool for ProbeDeclaration { - fn name(&self) -> &str { - "probe" - } - fn description(&self) -> &str { - "records the child's cancellation token" - } - fn parameters_schema(&self) -> Value { - json!({"type": "object"}) - } - async fn execute(&self, _: Value) -> anyhow::Result<tinytools::ToolResult> { - unreachable!("dispatch only") - } -} - -#[async_trait::async_trait] -impl ToolDispatch<(), ()> for ProbeTool { - fn tool(&self) -> Arc<dyn tinytools::Tool> { - Arc::new(ProbeDeclaration) - } - async fn execute( - &self, - _: &(), - _: CallId, - _: Value, - _: tinytools::ToolCallOptions, - parent: &RunContext<()>, - ) -> anyhow::Result<tinytools::ToolResult> { - *self.token.lock().unwrap() = Some(parent.cancellation.clone()); - if self.hang { - std::future::pending::<()>().await; - } - Ok(tinytools::ToolResult::json(json!({}))) - } -} - -/// Calls `probe` first, then fails transiently once the tool result is in. -struct ToolThenFail(Arc<AtomicUsize>); - -#[async_trait::async_trait] -impl ChatModel<()> for ToolThenFail { - async fn invoke( - &self, - state: &(), - request: ModelRequest, - ) -> tinyinference_llm::Result<ModelResponse> { - self.0.fetch_add(1, Ordering::SeqCst); - if request.messages.len() <= 1 { - MockModel::with_tool_call("probe", json!({})) - .invoke(state, request) - .await - } else { - Err(tinyinference_llm::Error::Model( - "connection reset by peer".into(), - )) - } - } -} - -fn probe_harness( - model: Arc<dyn ChatModel<()>>, - hang: bool, -) -> ( - AgentHarness<(), ()>, - Arc<std::sync::Mutex<Option<CancellationToken>>>, -) { - let token = Arc::new(std::sync::Mutex::new(None)); - let mut harness = AgentHarness::new(); - harness.register_model("child", model); - harness.register_tool_dispatch(Arc::new(ProbeTool { - token: token.clone(), - hang, - })); - harness.with_policy(RunPolicy { - retry: RetryPolicy::default().with_max_attempts(1), - ..RunPolicy::default() - }); - (harness, token) -} - -fn retrying(attempts: usize) -> SubAgentPolicy { - SubAgentPolicy::default().with_retry( - RetryPolicy::default() - .with_max_attempts(attempts) - .with_backoff_sleep(false), - ) -} - -#[tokio::test] -async fn no_retry_once_a_tool_ran_unless_the_policy_allows_it() { - let calls = Arc::new(AtomicUsize::new(0)); - let (harness, _) = probe_harness(Arc::new(ToolThenFail(calls.clone())), false); - let tool = tool_over(harness).with_policy(retrying(3)); - assert!(call_inline(&tool).await.is_error); - assert_eq!( - calls.load(Ordering::SeqCst), - 2, - "one attempt: tool call, then the failure" - ); - - let calls = Arc::new(AtomicUsize::new(0)); - let (harness, _) = probe_harness(Arc::new(ToolThenFail(calls.clone())), false); - let tool = tool_over(harness).with_policy(retrying(3).with_retry_after_tool_calls(true)); - assert!(call_inline(&tool).await.is_error); - assert_eq!( - calls.load(Ordering::SeqCst), - 6, - "three attempts when allowed" - ); -} - -#[tokio::test] -async fn timeout_cancels_the_childs_own_token() { - let calls = Arc::new(AtomicUsize::new(0)); - let (harness, token) = probe_harness(Arc::new(ToolThenFail(calls)), true); - let tool = tool_over(harness) - .with_policy(SubAgentPolicy::default().with_timeout(Duration::from_millis(100))); - let result = call_inline(&tool).await; - assert_eq!(payload(&result)["incomplete_kind"], "timeout"); - let token = token - .lock() - .unwrap() - .clone() - .expect("the child ran its tool"); - assert!(token.is_cancelled(), "the child observed cancellation"); -} - -async fn wait_terminal(tool: &SubAgentTool<(), ()>) -> SubAgentJob { - for _ in 0..200 { - if let Some(job) = tool - .job_registry() - .list() - .into_iter() - .find(|j| j.status.is_terminal()) - { - return job; - } - tokio::time::sleep(Duration::from_millis(25)).await; - } - panic!("job never settled"); -} - -#[tokio::test] -async fn background_mode_retries_and_the_job_link_follows_the_final_attempt() { - let calls = Arc::new(AtomicUsize::new(0)); - let mut harness = AgentHarness::new(); - harness.register_model("child", Arc::new(FailingModel(calls.clone()))); - harness.with_policy(RunPolicy { - retry: RetryPolicy::default().with_max_attempts(1), - ..RunPolicy::default() - }); - let tool = tool_over(harness).with_policy(retrying(3)); - let parent = RunContext::new(RunConfig::new("parent"), ()); - let queued = ToolDispatch::<(), ()>::execute( - &tool, - &(), - CallId::new("c"), - json!({"input": "work"}), - tinytools::ToolCallOptions::default(), - &parent, - ) - .await - .unwrap(); - let first_run = payload(&queued)["subagent_run_id"] - .as_str() - .unwrap() - .to_owned(); - let job = wait_terminal(&tool).await; - assert_eq!(job.status, SubAgentJobStatus::Failed); - assert_eq!(calls.load(Ordering::SeqCst), 3); - assert_eq!( - job.subagent_run_id.as_deref(), - Some(format!("{first_run}-a2").as_str()) - ); -} - -async fn sibling_run_ids(tool: &SubAgentTool<(), ()>) -> Vec<String> { - let parent = RunContext::new(RunConfig::new("parent"), ()); - let mut ids = Vec::new(); - for _ in 0..2 { - let queued = ToolDispatch::<(), ()>::execute( - tool, - &(), - CallId::new("c"), - json!({"input": "work"}), - tinytools::ToolCallOptions::default(), - &parent, - ) - .await - .unwrap(); - ids.push( - payload(&queued)["subagent_run_id"] - .as_str() - .unwrap() - .to_owned(), - ); - } - ids -} - -#[tokio::test] -async fn enabling_retries_does_not_change_sibling_child_ids() { - let plain = sibling_run_ids(&tool_over(constant("ok"))).await; - let with_retry = sibling_run_ids(&tool_over(constant("ok")).with_policy(retrying(4))).await; - assert_eq!(plain, with_retry); -} - -struct UsageModel; - -#[async_trait::async_trait] -impl ChatModel<()> for UsageModel { - async fn invoke(&self, _: &(), _: ModelRequest) -> tinyinference_llm::Result<ModelResponse> { - let mut response = ModelResponse::assistant("long answer"); - response.usage = Some(Usage { - output_tokens: 500, - ..Usage::default() - }); - Ok(response) - } -} - -#[tokio::test] -async fn a_token_overrun_keeps_the_output_applies_the_result_policy_and_is_incomplete() { - let mut harness = AgentHarness::new(); - harness.register_model("child", Arc::new(UsageModel)); - let tool = - tool_over(harness) - .with_policy(SubAgentPolicy::default().with_budget( - crate::subagent::SubAgentBudget::unlimited().with_max_output_tokens(100), - )) - .with_result_policy(ResultPolicy::new().with_schema(json!({"type": "object"}))); - let result = call_inline(&tool).await; - assert!(result.is_error); - let payload = payload(&result); - assert_eq!(payload["status"], "incomplete"); - assert_eq!(payload["incomplete_kind"], "budget_exceeded"); - assert_eq!(payload["output"], "long answer", "completed work is kept"); - assert!( - payload["schema_error"].is_string(), - "the result policy still ran" - ); -} - -#[tokio::test] -async fn artifact_overflow_without_a_store_surfaces_artifact_error_on_the_job() { - let tool = tool_over(constant(&"x".repeat(500))).with_result_policy( - ResultPolicy::new() - .with_max_chars(100) - .with_overflow(crate::subagent::ResultOverflow::Artifact), - ); - let payload = payload(&call_inline(&tool).await); - assert!( - payload["artifact_error"] - .as_str() - .unwrap() - .contains("no artifact store") - ); -} - -#[tokio::test] -async fn a_leaf_is_refused_with_actionable_wording_and_catches_its_own_tool_name() { - let mut harness = constant("x"); - harness.register_tool_dispatch(Arc::new(ProbeTool { - token: Arc::default(), - hang: false, - })); - let tool = tool_over(harness) - .with_tool_name("probe") - .with_role(SubagentRole::Leaf); - let result = call_inline(&tool).await; - assert!(result.is_error); - assert!( - result - .output() - .contains("You are a leaf agent: do this work yourself") - ); -} diff --git a/crates/tinyagents-orchestration/src/subagent/invocation/mod_registry_tests.rs b/crates/tinyagents-orchestration/src/subagent/invocation/mod_registry_tests.rs index 549ab8994..b798b28ab 100644 --- a/crates/tinyagents-orchestration/src/subagent/invocation/mod_registry_tests.rs +++ b/crates/tinyagents-orchestration/src/subagent/invocation/mod_registry_tests.rs @@ -2,13 +2,135 @@ //! these tests pin the shared-registry view and the lifecycle hand-offs. use super::*; +use crate::subagent::{SubagentOutcome, SubagentOutcomeKind, SubagentPause, SubagentResume}; use tinyagents_harness::error::TinyAgentsError; use tinyagents_harness::ids::TaskId; +/// Test fixture builder: all fixture jobs enter through the same driver +/// registration path used by `DriverSubAgentTool`. +#[derive(Default)] +struct RegistryFixtureMetadata { + parent_thread_id: Option<String>, + workspace: Option<tinytools::WorkspaceDescriptor>, +} + +trait RegistryTestHelpers { + fn create_fixture_job( + &self, + agent: &str, + owner: impl ToString, + ) -> (SubAgentJobId, tinyagents_harness::steering::SteeringHandle); + fn register_fixture_job( + &self, + agent: &str, + owner: impl ToString, + cancellation: tinyagents_harness::cancel::CancellationToken, + link: RegistryFixtureMetadata, + ) -> (SubAgentJobId, tinyagents_harness::steering::SteeringHandle); + fn finish_fixture_job(&self, id: &SubAgentJobId, result: Result<(), TinyAgentsError>); + fn finish_budget_overrun(&self, id: &SubAgentJobId); +} + +impl RegistryTestHelpers for SubAgentJobRegistry { + fn create_fixture_job( + &self, + agent: &str, + owner: impl ToString, + ) -> (SubAgentJobId, tinyagents_harness::steering::SteeringHandle) { + self.register_fixture_job( + agent, + owner, + tinyagents_harness::cancel::CancellationToken::new(), + RegistryFixtureMetadata::default(), + ) + } + + fn register_fixture_job( + &self, + agent: &str, + owner: impl ToString, + cancellation: tinyagents_harness::cancel::CancellationToken, + link: RegistryFixtureMetadata, + ) -> (SubAgentJobId, tinyagents_harness::steering::SteeringHandle) { + use tinyagents_harness::ids::next_seq; + let task = format!("registry-fixture-{}", next_seq()); + let key = crate::subagent::SubagentTaskKey { + root_run_id: "fixture-root".into(), + parent_run_id: "fixture-parent".into(), + thread_id: link.parent_thread_id, + task_id: task, + }; + let (id, steering, _, _, created) = self + .register_driver_job(DriverJobRegistration { + key, + agent: agent.to_owned(), + owner: owner.to_string(), + workspace: link.workspace, + cancellation, + steering: None, + run_queue: None, + continuing: false, + }) + .expect("driver fixture registration"); + assert!(created); + (id, steering) + } + + fn finish_fixture_job(&self, id: &SubAgentJobId, result: Result<(), TinyAgentsError>) { + let owner = self.controls.lock().unwrap().get(id).unwrap().owner.clone(); + if matches!(&result, Err(TinyAgentsError::Cancelled)) { + let _ = self.cancel_owned(id.as_str(), owner); + self.mark_driver_result(id, Err(TinyAgentsError::Cancelled.to_string())); + } else { + let result = result + .map(|()| SubagentOutcome::completed(id.as_str(), "fixture complete")) + .map_err(|error| error.to_string()); + self.mark_driver_result(id, result); + } + } + + fn finish_budget_overrun(&self, id: &SubAgentJobId) { + let mut outcome = SubagentOutcome::completed(id.as_str(), "late"); + outcome.status = SubagentOutcomeKind::Incomplete( + crate::subagent::SubagentIncomplete::new("over budget") + .with_kind(crate::subagent::IncompleteKind::BudgetExceeded), + ); + self.mark_driver_result(id, Ok(outcome)); + } +} + +fn register_driver_job( + jobs: &SubAgentJobRegistry, + task_id: impl Into<String>, +) -> (crate::subagent::SubagentTaskKey, SubAgentJobId) { + use tinyagents_harness::cancel::CancellationToken; + + let key = crate::subagent::SubagentTaskKey { + root_run_id: "root".into(), + parent_run_id: "parent".into(), + thread_id: Some("thread".into()), + task_id: task_id.into(), + }; + let (id, _, _, _, inserted) = jobs + .register_driver_job(DriverJobRegistration { + key: key.clone(), + agent: "worker".into(), + owner: "owner".into(), + workspace: None, + cancellation: CancellationToken::new(), + steering: None, + run_queue: None, + continuing: false, + }) + .unwrap(); + assert!(inserted); + (key, id) +} + #[tokio::test] async fn jobs_are_registered_in_the_shared_detached_registry() { let jobs = SubAgentJobRegistry::new(); - let (job_id, _steering) = jobs.create("worker", 7); + let (job_id, _steering) = jobs.create_fixture_job("worker", 7); let task_id = TaskId::new(job_id.as_str()); let tasks = jobs.tasks(); @@ -20,14 +142,1028 @@ async fn jobs_are_registered_in_the_shared_detached_registry() { assert!(tasks.snapshot(&task_id, "8").is_err()); } +#[tokio::test] +async fn root_settlement_includes_late_grandchildren_and_only_owned_root_jobs() { + use tinyagents_harness::cancel::CancellationToken; + use tinyagents_harness::ids::RunId; + + let jobs = SubAgentJobRegistry::new(); + let registration = |task: &str, root: &str, parent: &str, owner: &str| { + let key = crate::subagent::SubagentTaskKey { + root_run_id: root.into(), + parent_run_id: parent.into(), + thread_id: None, + task_id: task.into(), + }; + DriverJobRegistration { + key, + agent: task.into(), + owner: owner.into(), + workspace: None, + cancellation: CancellationToken::new(), + steering: None, + run_queue: None, + continuing: false, + } + }; + let register_with_runs = |registration, parent: &str, child: &str| { + let (id, _, _, _, _) = jobs + .register_driver_job_with_parent_run(registration, RunId::new(parent)) + .unwrap(); + jobs.set_attempt_run_id(&id, child); + id + }; + + let child = register_with_runs( + registration("child", "root", "key-parent", "owner"), + "root", + "child-run", + ); + let unrelated = jobs + .register_driver_job(registration( + "unrelated", + "other-root", + "other-root", + "owner", + )) + .unwrap() + .0; + let foreign = jobs + .register_driver_job(registration("foreign", "root", "root", "other-owner")) + .unwrap() + .0; + jobs.mark_running(&child); + jobs.mark_driver_result(&unrelated, Err("unrelated failure".into())); + jobs.mark_driver_result(&foreign, Err("foreign failure".into())); + + let settling = { + let jobs = jobs.clone(); + tokio::spawn(async move { + jobs.settle_root_run_descendants( + &RunId::new("root"), + "owner", + tokio::time::Instant::now() + std::time::Duration::from_secs(30), + &CancellationToken::new(), + ) + .await + }) + }; + tokio::task::yield_now().await; + + // A running descendant can still admit a grandchild. Its terminal + // transition happens only after that registration, so settlement must + // observe and wait for both jobs. + let grandchild = register_with_runs( + registration("grandchild", "root", "durable-parent", "owner"), + "child-run", + "grandchild-run", + ); + jobs.mark_running(&grandchild); + jobs.mark_driver_result(&child, Err("child failed".into())); + assert!(!settling.is_finished()); + jobs.mark_driver_result(&grandchild, Err("grandchild failed".into())); + + let settled = settling.await.unwrap().unwrap(); + let mut settled_ids = settled + .iter() + .map(|job| job.id.as_str()) + .collect::<Vec<_>>(); + settled_ids.sort_unstable(); + let mut expected_ids = vec![child.as_str(), grandchild.as_str()]; + expected_ids.sort_unstable(); + assert_eq!(settled_ids, expected_ids); + assert!( + settled + .iter() + .all(|job| job.status == SubAgentJobStatus::Failed) + ); + + let repeated = jobs + .settle_root_run_descendants( + &RunId::new("root"), + "owner", + tokio::time::Instant::now(), + &CancellationToken::new(), + ) + .await + .unwrap(); + assert_eq!(repeated, settled); +} + +#[tokio::test] +async fn root_settlement_uses_one_deadline_and_propagates_cancellation() { + use tinyagents_harness::cancel::CancellationToken; + use tinyagents_harness::ids::RunId; + + let jobs = SubAgentJobRegistry::new(); + let (child, _, _, _, _) = jobs + .register_driver_job_with_parent_run( + DriverJobRegistration { + key: crate::subagent::SubagentTaskKey { + root_run_id: "root".into(), + parent_run_id: "legacy-parent".into(), + thread_id: None, + task_id: "pending".into(), + }, + agent: "worker".into(), + owner: "owner".into(), + workspace: None, + cancellation: CancellationToken::new(), + steering: None, + run_queue: None, + continuing: false, + }, + RunId::new("root"), + ) + .unwrap(); + jobs.set_attempt_run_id(&child, "child"); + jobs.mark_running(&child); + let deadline_result = jobs + .settle_root_run_descendants( + &RunId::new("root"), + "owner", + tokio::time::Instant::now(), + &CancellationToken::new(), + ) + .await; + assert_eq!(deadline_result, Err(SubAgentJobError::SettlementTimedOut)); + + let cancellation = CancellationToken::new(); + cancellation.cancel(); + let cancelled_result = jobs + .settle_root_run_descendants( + &RunId::new("root"), + "owner", + tokio::time::Instant::now() + std::time::Duration::from_secs(30), + &cancellation, + ) + .await; + assert_eq!(cancelled_result, Err(SubAgentJobError::SettlementCancelled)); +} + +#[tokio::test] +async fn root_settlement_returns_empty_for_unlinked_jobs_and_empty_roots() { + use tinyagents_harness::cancel::CancellationToken; + use tinyagents_harness::ids::RunId; + + let jobs = SubAgentJobRegistry::new(); + let (legacy, _, _, _, _) = jobs + .register_driver_job(DriverJobRegistration { + key: crate::subagent::SubagentTaskKey { + root_run_id: "root".into(), + parent_run_id: "legacy-parent".into(), + thread_id: None, + task_id: "legacy-unlinked".into(), + }, + agent: "worker".into(), + owner: "owner".into(), + workspace: None, + cancellation: CancellationToken::new(), + steering: None, + run_queue: None, + continuing: false, + }) + .unwrap(); + jobs.mark_driver_result(&legacy, Err("legacy terminal".into())); + assert!( + jobs.settle_root_run_descendants( + &RunId::new("root"), + "owner", + tokio::time::Instant::now(), + &CancellationToken::new(), + ) + .await + .unwrap() + .is_empty() + ); + assert!( + jobs.settle_root_run_descendants( + &RunId::new("empty-root"), + "owner", + tokio::time::Instant::now(), + &CancellationToken::new(), + ) + .await + .unwrap() + .is_empty() + ); +} + +#[tokio::test(start_paused = true)] +async fn root_settlement_keeps_one_deadline_across_lifecycle_wakes() { + use tinyagents_harness::cancel::CancellationToken; + use tinyagents_harness::ids::RunId; + + let jobs = SubAgentJobRegistry::new(); + let (child, _, _, _, _) = jobs + .register_driver_job_with_parent_run( + DriverJobRegistration { + key: crate::subagent::SubagentTaskKey { + root_run_id: "root".into(), + parent_run_id: "ignored".into(), + thread_id: None, + task_id: "deadline-child".into(), + }, + agent: "worker".into(), + owner: "owner".into(), + workspace: None, + cancellation: CancellationToken::new(), + steering: None, + run_queue: None, + continuing: false, + }, + RunId::new("root"), + ) + .unwrap(); + jobs.set_attempt_run_id(&child, "deadline-child-attempt"); + jobs.mark_running(&child); + let cancellation = CancellationToken::new(); + let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(5); + let root_run_id = RunId::new("root"); + let settling = jobs.settle_root_run_descendants(&root_run_id, "owner", deadline, &cancellation); + tokio::pin!(settling); + tokio::select! { + biased; + result = &mut settling => panic!("live descendant settled unexpectedly: {result:?}"), + _ = tokio::task::yield_now() => {}, + } + tokio::time::advance(std::time::Duration::from_secs(4)).await; + jobs.set_attempt_run_id(&child, "deadline-child-another-observed-attempt"); + tokio::select! { + biased; + result = &mut settling => panic!("live descendant settled unexpectedly after wake: {result:?}"), + _ = tokio::task::yield_now() => {}, + } + tokio::time::advance(std::time::Duration::from_secs(1)).await; + tokio::select! { + biased; + result = &mut settling => assert_eq!(result, Err(SubAgentJobError::SettlementTimedOut)), + _ = tokio::task::yield_now() => panic!("settlement extended the caller's absolute deadline"), + } +} + +#[tokio::test] +async fn root_settlement_cancellation_while_blocked_releases_waiter_pins() { + use tinyagents_harness::cancel::CancellationToken; + use tinyagents_harness::ids::RunId; + + let jobs = SubAgentJobRegistry::new(); + let (child, _, _, _, _) = jobs + .register_driver_job_with_parent_run( + DriverJobRegistration { + key: crate::subagent::SubagentTaskKey { + root_run_id: "root".into(), + parent_run_id: "ignored".into(), + thread_id: None, + task_id: "cancel-child".into(), + }, + agent: "worker".into(), + owner: "owner".into(), + workspace: None, + cancellation: CancellationToken::new(), + steering: None, + run_queue: None, + continuing: false, + }, + RunId::new("root"), + ) + .unwrap(); + jobs.set_attempt_run_id(&child, "cancel-child-attempt"); + jobs.mark_running(&child); + let cancellation = CancellationToken::new(); + let root_run_id = RunId::new("root"); + let settling = jobs.settle_root_run_descendants( + &root_run_id, + "owner", + tokio::time::Instant::now() + std::time::Duration::from_secs(30), + &cancellation, + ); + tokio::pin!(settling); + tokio::select! { + biased; + result = &mut settling => panic!("live descendant settled unexpectedly: {result:?}"), + _ = tokio::task::yield_now() => {}, + } + cancellation.cancel(); + assert_eq!(settling.await, Err(SubAgentJobError::SettlementCancelled)); + assert_eq!( + jobs.controls()[&child] + .waiters + .load(std::sync::atomic::Ordering::Acquire), + 0 + ); + assert_eq!( + jobs.get(child.as_str()).unwrap().status, + SubAgentJobStatus::Running + ); +} + +#[tokio::test] +async fn root_settlement_pins_terminal_ancestors_while_descendants_remain_live() { + use tinyagents_harness::cancel::CancellationToken; + use tinyagents_harness::ids::RunId; + + let jobs = SubAgentJobRegistry::new(); + let registration = |task: &str, parent: &str| DriverJobRegistration { + key: crate::subagent::SubagentTaskKey { + root_run_id: "settlement-retention-root".into(), + parent_run_id: parent.into(), + thread_id: None, + task_id: task.into(), + }, + agent: task.into(), + owner: "owner".into(), + workspace: None, + cancellation: CancellationToken::new(), + steering: None, + run_queue: None, + continuing: false, + }; + let (parent, _, _, _, _) = jobs + .register_driver_job_with_parent_run( + registration("settlement-parent", "unused"), + RunId::new("settlement-root"), + ) + .unwrap(); + jobs.set_attempt_run_id(&parent, "settlement-parent-attempt"); + jobs.mark_running(&parent); + let (child, _, _, _, _) = jobs + .register_driver_job_with_parent_run( + registration("settlement-child", "unused"), + RunId::new("settlement-parent-attempt"), + ) + .unwrap(); + jobs.set_attempt_run_id(&child, "settlement-child-attempt"); + jobs.mark_running(&child); + jobs.mark_driver_result( + &parent, + Ok(SubagentOutcome::completed( + parent.as_str(), + "parent finished", + )), + ); + + let cancellation = CancellationToken::new(); + let root_run_id = RunId::new("settlement-root"); + let settling = jobs.settle_root_run_descendants( + &root_run_id, + "owner", + tokio::time::Instant::now() + std::time::Duration::from_secs(30), + &cancellation, + ); + tokio::pin!(settling); + tokio::select! { + biased; + result = &mut settling => panic!("live child settled unexpectedly: {result:?}"), + _ = tokio::task::yield_now() => {}, + } + + for index in 0..=super::types::SUBAGENT_JOB_SOFT_CAP { + let task = format!("pressure-{index}"); + let (id, _, _, _, _) = jobs + .register_driver_job_with_parent_run( + registration(&task, "unused"), + RunId::new("unrelated-root"), + ) + .unwrap(); + jobs.set_attempt_run_id(&id, &format!("{task}-attempt")); + jobs.mark_driver_result( + &id, + Ok(SubagentOutcome::completed( + id.as_str(), + "unrelated terminal", + )), + ); + } + + assert!( + jobs.get(parent.as_str()).is_some(), + "the pinned ancestor must survive terminal eviction" + ); + jobs.mark_driver_result( + &child, + Ok(SubagentOutcome::completed(child.as_str(), "child finished")), + ); + let settled = settling.await.unwrap(); + assert_eq!(settled.len(), 2); + assert!(settled.iter().any(|job| { + job.id == parent + && job.output.as_deref() == Some("parent finished") + && job.status == SubAgentJobStatus::Completed + })); + assert!(settled.iter().any(|job| job.id == child)); +} + +#[tokio::test] +async fn root_lease_preserves_parent_and_grandchild_before_settlement_starts() { + use tinyagents_harness::cancel::CancellationToken; + use tinyagents_harness::ids::RunId; + + let jobs = SubAgentJobRegistry::new(); + let lease = jobs + .retain_root_run_descendants(&RunId::new("leased-root"), "owner") + .unwrap(); + let registration = |task: &str, parent: &str| DriverJobRegistration { + key: crate::subagent::SubagentTaskKey { + root_run_id: "leased-root".into(), + parent_run_id: parent.into(), + thread_id: None, + task_id: task.into(), + }, + agent: task.into(), + owner: "owner".into(), + workspace: None, + cancellation: CancellationToken::new(), + steering: None, + run_queue: None, + continuing: false, + }; + let (parent, _, _, _, _) = jobs + .register_driver_job_with_parent_run( + registration("leased-parent", "ignored"), + RunId::new("leased-root"), + ) + .unwrap(); + jobs.set_attempt_run_id(&parent, "leased-parent-attempt"); + jobs.mark_running(&parent); + let (grandchild, _, _, _, _) = jobs + .register_driver_job_with_parent_run( + registration("leased-grandchild", "ignored"), + RunId::new("leased-parent-attempt"), + ) + .unwrap(); + jobs.set_attempt_run_id(&grandchild, "leased-grandchild-attempt"); + jobs.mark_running(&grandchild); + jobs.mark_driver_result( + &parent, + Ok(SubagentOutcome::completed( + parent.as_str(), + "retained parent output", + )), + ); + + for index in 0..=super::types::SUBAGENT_JOB_SOFT_CAP { + let task = format!("lease-pressure-{index}"); + let (id, _, _, _, _) = jobs + .register_driver_job_with_parent_run( + registration(&task, "ignored"), + RunId::new("unrelated-root"), + ) + .unwrap(); + jobs.set_attempt_run_id(&id, &format!("{task}-attempt")); + jobs.mark_driver_result(&id, Ok(SubagentOutcome::completed(id.as_str(), "pressure"))); + } + assert!(jobs.get(parent.as_str()).is_some()); + + let cancellation = CancellationToken::new(); + let settling = lease.settle( + tokio::time::Instant::now() + std::time::Duration::from_secs(30), + &cancellation, + ); + tokio::pin!(settling); + tokio::select! { + biased; + result = &mut settling => panic!("live grandchild settled unexpectedly: {result:?}"), + _ = tokio::task::yield_now() => {}, + } + jobs.mark_driver_result( + &grandchild, + Ok(SubagentOutcome::completed( + grandchild.as_str(), + "grandchild output", + )), + ); + let settled = settling.await.unwrap(); + assert_eq!(settled.len(), 2); + assert!(settled.iter().any(|job| { + job.id == parent && job.output.as_deref() == Some("retained parent output") + })); + assert!(settled.iter().any(|job| job.id == grandchild)); +} + +#[tokio::test] +async fn root_lease_keeps_completed_history_and_overlapping_lease_pins_independent() { + use tinyagents_harness::cancel::CancellationToken; + use tinyagents_harness::ids::RunId; + + let jobs = SubAgentJobRegistry::new(); + let root = RunId::new("history-root"); + let first = jobs.retain_root_run_descendants(&root, "owner").unwrap(); + let second = jobs.retain_root_run_descendants(&root, "owner").unwrap(); + let settling_lease = jobs.retain_root_run_descendants(&root, "owner").unwrap(); + let foreign = jobs + .retain_root_run_descendants(&root, "foreign-owner") + .unwrap(); + let (parent, _, _, _, _) = jobs + .register_driver_job_with_parent_run( + DriverJobRegistration { + key: crate::subagent::SubagentTaskKey { + root_run_id: "history-root".into(), + parent_run_id: "ignored".into(), + thread_id: None, + task_id: "history-parent".into(), + }, + agent: "worker".into(), + owner: "owner".into(), + workspace: None, + cancellation: CancellationToken::new(), + steering: None, + run_queue: None, + continuing: false, + }, + root.clone(), + ) + .unwrap(); + jobs.set_attempt_run_id(&parent, "history-parent-attempt"); + let (child, _, _, _, _) = jobs + .register_driver_job_with_parent_run( + DriverJobRegistration { + key: crate::subagent::SubagentTaskKey { + root_run_id: "history-root".into(), + parent_run_id: "ignored".into(), + thread_id: None, + task_id: "history-child".into(), + }, + agent: "worker".into(), + owner: "owner".into(), + workspace: None, + cancellation: CancellationToken::new(), + steering: None, + run_queue: None, + continuing: false, + }, + RunId::new("history-parent-attempt"), + ) + .unwrap(); + jobs.set_attempt_run_id(&child, "history-child-attempt"); + jobs.mark_driver_result( + &parent, + Ok(SubagentOutcome::completed(parent.as_str(), "parent")), + ); + jobs.mark_driver_result( + &child, + Ok(SubagentOutcome::completed(child.as_str(), "child")), + ); + assert_eq!( + jobs.controls()[&parent] + .waiters + .load(std::sync::atomic::Ordering::Acquire), + 3 + ); + drop(foreign); + drop(first); + assert_eq!( + jobs.controls()[&parent] + .waiters + .load(std::sync::atomic::Ordering::Acquire), + 2 + ); + assert!(jobs.get(parent.as_str()).is_some()); + + for index in 0..=super::types::SUBAGENT_JOB_SOFT_CAP { + let task = format!("history-pressure-{index}"); + let (id, _, _, _, _) = jobs + .register_driver_job_with_parent_run( + DriverJobRegistration { + key: crate::subagent::SubagentTaskKey { + root_run_id: "history-root".into(), + parent_run_id: "ignored".into(), + thread_id: None, + task_id: task.clone(), + }, + agent: "worker".into(), + owner: "owner".into(), + workspace: None, + cancellation: CancellationToken::new(), + steering: None, + run_queue: None, + continuing: false, + }, + RunId::new("pressure-root"), + ) + .unwrap(); + jobs.set_attempt_run_id(&id, &format!("{task}-attempt")); + jobs.mark_driver_result(&id, Ok(SubagentOutcome::completed(id.as_str(), "pressure"))); + } + assert!(jobs.get(parent.as_str()).is_some()); + drop(second); + assert!(jobs.get(parent.as_str()).is_some()); + let settled = settling_lease + .settle( + tokio::time::Instant::now() + std::time::Duration::from_secs(30), + &CancellationToken::new(), + ) + .await + .unwrap(); + assert_eq!(settled.len(), 2); + assert!( + settled + .iter() + .any(|job| job.id == parent && job.output.as_deref() == Some("parent")) + ); + assert!( + settled + .iter() + .any(|job| job.id == child && job.output.as_deref() == Some("child")) + ); + assert!( + jobs.get(parent.as_str()).is_none(), + "last lease drop releases retention and allows pruning" + ); +} + +#[tokio::test] +async fn cancelled_root_lease_releases_history_retention() { + use tinyagents_harness::cancel::CancellationToken; + use tinyagents_harness::ids::RunId; + + let jobs = SubAgentJobRegistry::new(); + let root = RunId::new("cancelled-lease-root"); + let lease = jobs.retain_root_run_descendants(&root, "owner").unwrap(); + let (parent, _, _, _, _) = jobs + .register_driver_job_with_parent_run( + DriverJobRegistration { + key: crate::subagent::SubagentTaskKey { + root_run_id: root.as_str().into(), + parent_run_id: "ignored".into(), + thread_id: None, + task_id: "cancelled-lease-parent".into(), + }, + agent: "worker".into(), + owner: "owner".into(), + workspace: None, + cancellation: CancellationToken::new(), + steering: None, + run_queue: None, + continuing: false, + }, + root.clone(), + ) + .unwrap(); + jobs.set_attempt_run_id(&parent, "cancelled-lease-attempt"); + jobs.mark_running(&parent); + for index in 0..=super::types::SUBAGENT_JOB_SOFT_CAP { + let task = format!("cancel-lease-pressure-{index}"); + let (id, _, _, _, _) = jobs + .register_driver_job_with_parent_run( + DriverJobRegistration { + key: crate::subagent::SubagentTaskKey { + root_run_id: root.as_str().into(), + parent_run_id: "ignored".into(), + thread_id: None, + task_id: task.clone(), + }, + agent: "worker".into(), + owner: "owner".into(), + workspace: None, + cancellation: CancellationToken::new(), + steering: None, + run_queue: None, + continuing: false, + }, + RunId::new("unrelated-root"), + ) + .unwrap(); + jobs.set_attempt_run_id(&id, &format!("{task}-attempt")); + jobs.mark_driver_result(&id, Ok(SubagentOutcome::completed(id.as_str(), "pressure"))); + } + assert!(jobs.get(parent.as_str()).is_some()); + let cancellation = CancellationToken::new(); + cancellation.cancel(); + assert_eq!( + lease + .settle( + tokio::time::Instant::now() + std::time::Duration::from_secs(30), + &cancellation, + ) + .await, + Err(SubAgentJobError::SettlementCancelled) + ); + assert_eq!( + jobs.controls()[&parent] + .waiters + .load(std::sync::atomic::Ordering::Acquire), + 0, + "cancellation consumes and releases the lease pin" + ); +} + +#[tokio::test] +async fn resumed_job_keeps_old_attempt_edges_without_importing_new_root_siblings() { + use tinyagents_harness::cancel::CancellationToken; + use tinyagents_harness::ids::RunId; + + let jobs = SubAgentJobRegistry::new(); + let registration = |task: &str, key_parent: &str| DriverJobRegistration { + key: crate::subagent::SubagentTaskKey { + root_run_id: "durable-root".into(), + parent_run_id: key_parent.into(), + thread_id: None, + task_id: task.into(), + }, + agent: task.into(), + owner: "owner".into(), + workspace: None, + cancellation: CancellationToken::new(), + steering: None, + run_queue: None, + continuing: false, + }; + let root_lease = jobs + .retain_root_run_descendants(&RunId::new("old-root"), "owner") + .unwrap(); + + let (job, _, _, _, _) = jobs + .register_driver_job_with_parent_run( + registration("resumed", "durable-parent"), + RunId::new("old-root"), + ) + .unwrap(); + jobs.set_attempt_run_id(&job, "old-attempt"); + jobs.set_attempt_run_id(&job, "old-retry-attempt"); + + let (old_grandchild, _, _, _, _) = jobs + .register_driver_job_with_parent_run( + registration("old-grandchild", "durable-parent"), + RunId::new("old-retry-attempt"), + ) + .unwrap(); + jobs.set_attempt_run_id(&old_grandchild, "old-grandchild-run"); + jobs.mark_running(&old_grandchild); + + let mut paused = SubagentOutcome::completed(job.as_str(), "waiting"); + paused.status = SubagentOutcomeKind::AwaitingInput(SubagentPause { + reason: "resume later".into(), + resume: SubagentResume::default(), + }); + jobs.mark_driver_result(&old_grandchild, Err("old attempt finished".into())); + jobs.mark_driver_result(&job, Ok(paused)); + let paused_settlement = { + let jobs = jobs.clone(); + tokio::spawn(async move { + jobs.settle_root_run_descendants( + &RunId::new("old-root"), + "owner", + tokio::time::Instant::now() + std::time::Duration::from_secs(30), + &CancellationToken::new(), + ) + .await + }) + }; + tokio::task::yield_now().await; + assert!(!paused_settlement.is_finished()); + let mut resumed = registration("resumed", "durable-parent"); + resumed.continuing = true; + let (resumed_id, _, _, _, continued) = jobs + .register_driver_job_with_parent_run(resumed, RunId::new("new-root")) + .unwrap(); + assert_eq!(resumed_id, job); + assert!(continued); + jobs.set_attempt_run_id(&resumed_id, "new-attempt"); + + let (new_grandchild, _, _, _, _) = jobs + .register_driver_job_with_parent_run( + registration("new-grandchild", "durable-parent"), + RunId::new("new-attempt"), + ) + .unwrap(); + jobs.set_attempt_run_id(&new_grandchild, "new-grandchild-run"); + jobs.mark_running(&new_grandchild); + let (new_root_sibling, _, _, _, _) = jobs + .register_driver_job_with_parent_run( + registration("new-root-sibling", "new-root"), + RunId::new("new-root"), + ) + .unwrap(); + jobs.set_attempt_run_id(&new_root_sibling, "sibling-run"); + jobs.mark_running(&new_root_sibling); + + jobs.mark_driver_result(&job, Err("resumed job finished".into())); + jobs.mark_driver_result(&new_grandchild, Err("new attempt finished".into())); + jobs.mark_driver_result(&new_root_sibling, Err("sibling finished".into())); + paused_settlement.await.unwrap().unwrap(); + + let settle = |root: &'static str| { + let jobs = jobs.clone(); + async move { + jobs.settle_root_run_descendants( + &RunId::new(root), + "owner", + tokio::time::Instant::now() + std::time::Duration::from_secs(30), + &CancellationToken::new(), + ) + .await + .unwrap() + } + }; + let old_root_jobs = root_lease + .settle( + tokio::time::Instant::now() + std::time::Duration::from_secs(30), + &CancellationToken::new(), + ) + .await + .unwrap(); + assert_eq!(old_root_jobs.len(), 3); + assert!(old_root_jobs.iter().any(|entry| entry.id == old_grandchild)); + assert!(old_root_jobs.iter().any(|entry| entry.id == new_grandchild)); + assert!( + !old_root_jobs + .iter() + .any(|entry| entry.id == new_root_sibling) + ); + + let new_root_jobs = settle("new-root").await; + assert_eq!(new_root_jobs.len(), 4); + assert!(new_root_jobs.iter().any(|entry| entry.id == job)); + assert!(new_root_jobs.iter().any(|entry| entry.id == old_grandchild)); + assert!(new_root_jobs.iter().any(|entry| entry.id == new_grandchild)); + assert!( + new_root_jobs + .iter() + .any(|entry| entry.id == new_root_sibling) + ); +} + +#[tokio::test] +async fn owner_cancel_tombstones_a_driver_job_against_late_completion() { + use tinyagents_harness::cancel::CancellationToken; + use tinyagents_tasks::{ + CompletionRecord, CompletionResult, CompletionRouter, CompletionStatus, + InMemoryCompletionStore, RecordOutcome, + }; + + let jobs = SubAgentJobRegistry::new(); + let completions = CompletionRouter::new(std::sync::Arc::new(InMemoryCompletionStore::new())); + let key = crate::subagent::SubagentTaskKey { + root_run_id: "root".into(), + parent_run_id: "parent".into(), + thread_id: Some("thread".into()), + task_id: "child".into(), + }; + let token = CancellationToken::new(); + let (id, _, _, registered_token, inserted) = jobs + .register_driver_job(DriverJobRegistration { + key: key.clone(), + agent: "worker".into(), + owner: "owner-a".into(), + workspace: Some(tinytools::WorkspaceDescriptor::new("/workspace/a")), + cancellation: token.clone(), + steering: None, + run_queue: None, + continuing: false, + }) + .unwrap(); + assert!(inserted); + assert!(!registered_token.is_cancelled()); + + assert!( + jobs.cancel_key_owned_and_tombstone(&key, "owner-b", &completions) + .is_err() + ); + assert!(!token.is_cancelled()); + let late_completion = CompletionRecord::new( + "child", + "thread", + "worker", + CompletionStatus::Success, + CompletionResult::text("late output"), + ); + assert!(matches!( + completions.record(late_completion.clone()).await.unwrap(), + RecordOutcome::Recorded { .. } + )); + let canceled = jobs + .cancel_key_owned_and_tombstone(&key, "owner-a", &completions) + .unwrap(); + assert!(token.is_cancelled()); + assert_eq!(canceled.status, SubAgentJobStatus::Cancelled); + + jobs.mark_driver_result( + &id, + Ok(crate::subagent::SubagentOutcome::completed( + "child", + "late output", + )), + ); + let retained = jobs.get_owned(id.as_str(), "owner-a").unwrap(); + assert_eq!(retained.status, SubAgentJobStatus::Cancelled); + assert_eq!(retained.output, None); + jobs.mark_driver_result(&id, Err("late failure".into())); + let retained = jobs.get_owned(id.as_str(), "owner-a").unwrap(); + assert_eq!(retained.status, SubAgentJobStatus::Cancelled); + assert_eq!(retained.error, Some(TinyAgentsError::Cancelled.to_string())); + assert_eq!( + completions.record(late_completion).await.unwrap(), + RecordOutcome::Suppressed + ); + assert!(completions.pending_for("thread").is_empty()); +} + +#[tokio::test] +async fn live_job_metadata_retains_parent_thread_and_workspace() { + let jobs = SubAgentJobRegistry::new(); + let workspace = tinytools::WorkspaceDescriptor::new("/workspace/project"); + let (job_id, _) = jobs.register_fixture_job( + "worker", + 7, + tinyagents_harness::cancel::CancellationToken::new(), + RegistryFixtureMetadata { + parent_thread_id: Some("thread-7".into()), + workspace: Some(workspace.clone()), + }, + ); + let snapshot = jobs + .tasks() + .snapshot(&TaskId::new(job_id.as_str()), "7") + .unwrap(); + assert_eq!( + snapshot.metadata.parent_thread_id.as_deref(), + Some("thread-7") + ); + assert_eq!(snapshot.metadata.workspace, Some(workspace)); +} + +#[tokio::test] +async fn thread_and_global_cancellation_select_jobs_by_retained_metadata() { + use tinyagents_harness::cancel::CancellationToken; + + let jobs = SubAgentJobRegistry::new(); + let workspace = tinytools::WorkspaceDescriptor::new("/workspace/project"); + let thread_token = CancellationToken::new(); + let other_token = CancellationToken::new(); + let (thread_job, _) = jobs.register_fixture_job( + "worker", + 7, + thread_token.clone(), + RegistryFixtureMetadata { + parent_thread_id: Some("thread-a".into()), + workspace: Some(workspace.clone()), + }, + ); + let (other_job, _) = jobs.register_fixture_job( + "worker", + 7, + other_token.clone(), + RegistryFixtureMetadata { + parent_thread_id: Some("thread-b".into()), + workspace: Some(workspace.clone()), + }, + ); + + let scope = ParentJobScope::new(workspace, "thread-a"); + assert_eq!(jobs.cancel_for_thread(&scope).unwrap(), 1); + assert!(thread_token.is_cancelled()); + assert!(!other_token.is_cancelled()); + assert!(jobs.get(thread_job.as_str()).is_some()); + assert!(!jobs.get(thread_job.as_str()).unwrap().status.is_terminal()); + assert_eq!(jobs.cancel_for_thread(&scope).unwrap(), 0); + + assert_eq!(jobs.cancel_all().unwrap(), 1); + assert!(other_token.is_cancelled()); + assert!(jobs.get(other_job.as_str()).is_some()); +} + +#[tokio::test] +async fn thread_cancellation_is_scoped_to_workspace() { + use tinyagents_harness::cancel::CancellationToken; + + let jobs = SubAgentJobRegistry::new(); + let workspace_a = tinytools::WorkspaceDescriptor::new("/workspace/a"); + let workspace_b = tinytools::WorkspaceDescriptor::new("/workspace/b"); + let token_a = CancellationToken::new(); + let token_b = CancellationToken::new(); + for (token, workspace) in [ + (token_a.clone(), workspace_a.clone()), + (token_b.clone(), workspace_b), + ] { + jobs.register_fixture_job( + "worker", + 7, + token, + RegistryFixtureMetadata { + parent_thread_id: Some("same-thread".into()), + workspace: Some(workspace), + }, + ); + } + + assert_eq!( + jobs.cancel_for_thread(&ParentJobScope::new(workspace_a, "same-thread")) + .unwrap(), + 1 + ); + assert!(token_a.is_cancelled()); + assert!(!token_b.is_cancelled()); +} + #[tokio::test] async fn steering_is_reachable_while_live_and_released_when_settled() { let jobs = SubAgentJobRegistry::new(); - let (job_id, _steering) = jobs.create("worker", 1); + let (job_id, _steering) = jobs.create_fixture_job("worker", 1); let task_id = TaskId::new(job_id.as_str()); assert!(jobs.tasks().steering_handle_trusted(&task_id).is_ok()); - jobs.mark_result(&job_id, Err(TinyAgentsError::Cancelled)); + jobs.finish_fixture_job(&job_id, Err(TinyAgentsError::Cancelled)); assert!(jobs.tasks().steering_handle_trusted(&task_id).is_err()); // A settled job stays queryable; the registry never prunes it. @@ -39,26 +1175,297 @@ async fn steering_is_reachable_while_live_and_released_when_settled() { } #[tokio::test] -async fn terminal_jobs_survive_many_registrations() { - // The detached registry sweeps terminal entries at its soft cap; the job - // registry must never lose a settled job that way. - let jobs = SubAgentJobRegistry::new(); - let mut ids = Vec::new(); - for _ in 0..2_000 { - let (id, _s) = jobs.create("worker", 1); - jobs.mark_result(&id, Err(TinyAgentsError::Cancelled)); - ids.push(id); +async fn settled_driver_jobs_have_bounded_retention_across_all_registry_maps() { + let jobs = SubAgentJobRegistry::new(); + let total = 2_000; + for index in 0..total { + let (_, id) = register_driver_job(&jobs, format!("child-{index}")); + jobs.mark_driver_result(&id, Err("finished".into())); + } + + let retained = jobs.tasks().len().unwrap(); + assert!(retained < total, "settled tasks retained: {retained}"); + assert_eq!(jobs.controls.lock().unwrap().len(), retained); + assert_eq!(jobs.driver_jobs.lock().unwrap().len(), retained); + assert_eq!(jobs.list().len(), retained); +} + +#[tokio::test] +async fn recent_terminal_driver_jobs_replay_until_oldest_cache_eviction() { + use tinyagents_harness::cancel::CancellationToken; + + let jobs = SubAgentJobRegistry::new(); + let (old_key, old_id) = register_driver_job(&jobs, "old-cache-entry"); + jobs.mark_running(&old_id); + jobs.mark_driver_result( + &old_id, + Ok(crate::subagent::SubagentOutcome::completed( + "old", + "old result", + )), + ); + + let (recent_key, recent_id) = register_driver_job(&jobs, "recent-cache-entry"); + jobs.mark_running(&recent_id); + jobs.mark_driver_result( + &recent_id, + Ok(crate::subagent::SubagentOutcome::completed( + "recent", + "replayable result", + )), + ); + let waited = jobs + .wait_owned(recent_id.as_str(), "owner") + .await + .expect("terminal result"); + assert_eq!(waited.output.as_deref(), Some("replayable result")); + + let (replayed_id, _, _, _, inserted) = jobs + .register_driver_job(DriverJobRegistration { + key: recent_key.clone(), + agent: "worker".into(), + owner: "owner".into(), + workspace: None, + cancellation: CancellationToken::new(), + steering: None, + run_queue: None, + continuing: false, + }) + .expect("reuse retained terminal driver job"); + assert_eq!(replayed_id, recent_id); + assert!(!inserted); + + for index in 0..127 { + let (_, id) = register_driver_job(&jobs, format!("cache-filler-{index}")); + jobs.mark_driver_result( + &id, + Ok(crate::subagent::SubagentOutcome::completed( + "filler", "done", + )), + ); + } + + assert!(jobs.get(old_id.as_str()).is_none()); + assert_eq!( + jobs.get(recent_id.as_str()).unwrap().output.as_deref(), + Some("replayable result") + ); + assert_eq!( + jobs.driver_job_for_key(&recent_key, "owner"), + Some(recent_id) + ); + assert!(jobs.driver_job_for_key(&old_key, "owner").is_none()); + assert_eq!(jobs.tasks().len().unwrap(), 128); + assert_eq!(jobs.controls.lock().unwrap().len(), 128); + assert_eq!(jobs.driver_jobs.lock().unwrap().len(), 128); +} + +#[tokio::test] +async fn settling_many_prestarted_jobs_bounds_terminal_retention() { + let jobs = SubAgentJobRegistry::new(); + let total = 256; + let ids: Vec<_> = (0..total) + .map(|index| register_driver_job(&jobs, format!("prestarted-{index}")).1) + .collect(); + for id in &ids { + jobs.mark_running(id); + } + for id in &ids { + jobs.mark_driver_result( + id, + Ok(crate::subagent::SubagentOutcome::completed( + "prestarted", + "done", + )), + ); + } + + let retained = jobs.tasks().len().unwrap(); + assert!( + retained <= 128, + "terminal tasks retained after a settlement batch: {retained}" + ); + assert_eq!(jobs.controls.lock().unwrap().len(), retained); + assert_eq!(jobs.driver_jobs.lock().unwrap().len(), retained); + assert_eq!(jobs.list().len(), retained); +} + +#[tokio::test] +async fn dropping_terminal_waiters_releases_retention_pins() { + use std::future::Future; + use std::task::{Context, Poll, Waker}; + + let jobs = SubAgentJobRegistry::new(); + let ids: Vec<_> = (0..132) + .map(|index| register_driver_job(&jobs, format!("pinned-{index}")).1) + .collect(); + for id in &ids { + jobs.mark_running(id); } - assert!(ids.iter().all(|id| jobs.get(id.as_str()).is_some())); - assert_eq!(jobs.list().len(), 2_000); + + let mut waiters: Vec<_> = ids + .iter() + .map(|id| Box::pin(jobs.wait_owned(id.as_str(), "owner"))) + .collect(); + let waker = Waker::noop(); + let mut context = Context::from_waker(waker); + for waiter in &mut waiters { + assert!(matches!(waiter.as_mut().poll(&mut context), Poll::Pending)); + } + + for id in &ids { + jobs.mark_driver_result( + id, + Ok(crate::subagent::SubagentOutcome::completed( + "pinned", "done", + )), + ); + } + assert_eq!(jobs.tasks().len().unwrap(), ids.len()); + + drop(waiters); + let retained = jobs.tasks().len().unwrap(); + assert!( + retained <= 128, + "terminal tasks retained after waiter drops: {retained}" + ); + assert_eq!(jobs.controls.lock().unwrap().len(), retained); + assert_eq!(jobs.driver_jobs.lock().unwrap().len(), retained); +} + +#[tokio::test] +async fn direct_cancellation_keeps_terminal_retention_bounded() { + let jobs = SubAgentJobRegistry::new(); + let ids: Vec<_> = (0..256) + .map(|index| register_driver_job(&jobs, format!("cancel-{index}")).1) + .collect(); + for id in &ids { + let mut awaiting = crate::subagent::SubagentOutcome::completed("cancel", "paused"); + awaiting.status = + crate::subagent::SubagentOutcomeKind::AwaitingInput(crate::subagent::SubagentPause { + reason: "clarification".into(), + resume: crate::subagent::SubagentResume::default(), + }); + jobs.mark_driver_result(id, Ok(awaiting)); + } + for id in &ids { + jobs.cancel_owned(id.as_str(), "owner").expect("cancel"); + } + + let retained = jobs.tasks().len().unwrap(); + assert!(retained <= 128, "cancelled tasks retained: {retained}"); + assert_eq!(jobs.controls.lock().unwrap().len(), retained); + assert_eq!(jobs.driver_jobs.lock().unwrap().len(), retained); +} + +#[tokio::test] +async fn active_and_awaiting_jobs_survive_terminal_retention_pressure() { + let jobs = SubAgentJobRegistry::new(); + let (active_key, active_id) = register_driver_job(&jobs, "active"); + jobs.mark_running(&active_id); + let (awaiting_key, awaiting_id) = register_driver_job(&jobs, "awaiting"); + let mut awaiting = crate::subagent::SubagentOutcome::completed("awaiting", "paused"); + awaiting.status = + crate::subagent::SubagentOutcomeKind::AwaitingInput(crate::subagent::SubagentPause { + reason: "clarification".into(), + resume: crate::subagent::SubagentResume::default(), + }); + jobs.mark_driver_result(&awaiting_id, Ok(awaiting)); + + for index in 0..2_000 { + let (_, id) = register_driver_job(&jobs, format!("settled-{index}")); + jobs.mark_driver_result(&id, Err("finished".into())); + } + + assert!(jobs.tasks().len().unwrap() < 2_002); + assert!(jobs.get(active_id.as_str()).is_some()); + assert_eq!( + jobs.get(awaiting_id.as_str()).unwrap().status, + SubAgentJobStatus::AwaitingInput + ); + assert_eq!( + jobs.driver_job_for_key(&active_key, "owner"), + Some(active_id) + ); + assert_eq!( + jobs.driver_job_for_key(&awaiting_key, "owner"), + Some(awaiting_id.clone()) + ); + + let (resumed_id, _, _, _, should_run) = jobs + .register_driver_job(DriverJobRegistration { + key: awaiting_key, + agent: "worker".into(), + owner: "owner".into(), + workspace: None, + cancellation: tinyagents_harness::cancel::CancellationToken::new(), + steering: None, + run_queue: None, + continuing: true, + }) + .unwrap(); + assert!(should_run); + assert_eq!(resumed_id, awaiting_id); + assert_eq!( + jobs.get(resumed_id.as_str()).unwrap().status, + SubAgentJobStatus::Running + ); +} + +#[tokio::test] +async fn waiter_keeps_terminal_snapshot_when_pruning_races_its_observation() { + use std::future::Future; + + let jobs = SubAgentJobRegistry::new(); + let (key, id) = register_driver_job(&jobs, "awaited"); + jobs.mark_running(&id); + + let waiter_jobs = jobs.clone(); + let waiter_id = id.clone(); + let (started_tx, started_rx) = tokio::sync::oneshot::channel(); + let waiter = tokio::spawn(async move { + let wait = waiter_jobs.wait_owned(waiter_id.as_str(), "owner"); + tokio::pin!(wait); + let mut started_tx = Some(started_tx); + std::future::poll_fn(|cx| { + let result = wait.as_mut().poll(cx); + if result.is_pending() + && let Some(started_tx) = started_tx.take() + { + let _ = started_tx.send(()); + } + result + }) + .await + }); + started_rx.await.unwrap(); + + jobs.mark_driver_result( + &id, + Ok(crate::subagent::SubagentOutcome::completed( + "awaited", + "preserved output", + )), + ); + for index in 0..2_000 { + let (_, pressure_id) = register_driver_job(&jobs, format!("pressure-{index}")); + jobs.mark_driver_result(&pressure_id, Err("finished".into())); + } + + assert_eq!(jobs.driver_job_for_key(&key, "owner"), Some(id.clone())); + let observed = waiter.await.unwrap().unwrap(); + assert_eq!(observed.id, id); + assert_eq!(observed.status, SubAgentJobStatus::Completed); + assert_eq!(observed.output.as_deref(), Some("preserved output")); + assert!(jobs.driver_job_for_key(&key, "owner").is_none()); } #[tokio::test] async fn control_errors_keep_their_documented_order() { let jobs = SubAgentJobRegistry::new(); - let (live, _s) = jobs.create("worker", 1); - let (done, _s2) = jobs.create("worker", 1); - jobs.mark_result(&done, Err(TinyAgentsError::Cancelled)); + let (live, _s) = jobs.create_fixture_job("worker", 1); + let (done, _s2) = jobs.create_fixture_job("worker", 1); + jobs.finish_fixture_job(&done, Err(TinyAgentsError::Cancelled)); let not_found = |id: &str| SubAgentJobError::NotFound(id.to_owned()); // Unknown id and foreign owner are both NotFound, even for a settled job. @@ -112,7 +1519,7 @@ async fn control_errors_keep_their_documented_order() { #[tokio::test] async fn poisoned_detached_registry_errors_remain_typed_across_job_controls() { let jobs = SubAgentJobRegistry::new(); - let (id, _) = jobs.create("worker", 7); + let (id, _) = jobs.create_fixture_job("worker", 7); let panic = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { super::types::PANIC_ON_JOB_META_CLONE.with(|panic| panic.set(true)); @@ -152,7 +1559,7 @@ async fn poisoned_detached_registry_errors_remain_typed_across_job_controls() { async fn listings_are_sorted_by_job_id() { let jobs = SubAgentJobRegistry::new(); for _ in 0..40 { - jobs.create("worker", 1); + jobs.create_fixture_job("worker", 1); } let all: Vec<String> = jobs.list().into_iter().map(|j| j.id.0).collect(); let mut sorted = all.clone(); @@ -164,9 +1571,9 @@ async fn listings_are_sorted_by_job_id() { #[tokio::test] async fn terminal_jobs_never_return_to_running() { let jobs = SubAgentJobRegistry::new(); - let (id, _) = jobs.create("worker", 1); + let (id, _) = jobs.create_fixture_job("worker", 1); jobs.cancel_owned(id.as_str(), 1).expect("cancel"); - jobs.mark_result(&id, Err(TinyAgentsError::Cancelled)); + jobs.finish_fixture_job(&id, Err(TinyAgentsError::Cancelled)); jobs.mark_running(&id); assert_eq!( jobs.get(id.as_str()).unwrap().status, @@ -177,16 +1584,16 @@ async fn terminal_jobs_never_return_to_running() { #[tokio::test] async fn settled_jobs_release_their_cancellation_token() { let jobs = SubAgentJobRegistry::new(); - let (id, _) = jobs.create("worker", 1); + let (id, _) = jobs.create_fixture_job("worker", 1); let task_id = TaskId::new(id.as_str()); assert!(jobs.tasks().holds_cancellation(&task_id).unwrap()); jobs.cancel_owned(id.as_str(), 1).expect("cancel"); assert!(!jobs.tasks().holds_cancellation(&task_id).unwrap()); assert!(!jobs.get(id.as_str()).unwrap().status.is_terminal()); - jobs.mark_result(&id, Err(TinyAgentsError::Cancelled)); + jobs.finish_fixture_job(&id, Err(TinyAgentsError::Cancelled)); assert!(jobs.get(id.as_str()).unwrap().status.is_terminal()); - let (aborted, _) = jobs.create("worker", 1); + let (aborted, _) = jobs.create_fixture_job("worker", 1); jobs.mark_aborted(&aborted, true); assert!( !jobs @@ -199,16 +1606,9 @@ async fn settled_jobs_release_their_cancellation_token() { #[tokio::test] async fn a_budget_overrun_that_raced_an_owner_cancel_settles_cancelled() { let jobs = SubAgentJobRegistry::new(); - let (id, _) = jobs.create("worker", 1); + let (id, _) = jobs.create_fixture_job("worker", 1); jobs.cancel_owned(id.as_str(), 1).expect("cancel"); - jobs.mark_budget_overrun( - &id, - crate::subagent::AppliedResult { - text: "late".into(), - ..Default::default() - }, - "over budget".into(), - ); + jobs.finish_budget_overrun(&id); let job = jobs.get(id.as_str()).unwrap(); assert_eq!(job.status, SubAgentJobStatus::Cancelled); assert!(job.incomplete_kind.is_none() && job.output.is_none()); @@ -217,14 +1617,14 @@ async fn a_budget_overrun_that_raced_an_owner_cancel_settles_cancelled() { #[tokio::test] async fn a_timeout_or_error_after_an_owner_cancel_settles_cancelled() { let jobs = SubAgentJobRegistry::new(); - for result in [ - Err(TinyAgentsError::Timeout("late".into())), - Err(TinyAgentsError::Tool("late".into())), - Err(TinyAgentsError::LimitExceeded("late".into())), + for error in [ + TinyAgentsError::Timeout("late".into()), + TinyAgentsError::Tool("late".into()), + TinyAgentsError::LimitExceeded("late".into()), ] { - let (id, _) = jobs.create("worker", 1); + let (id, _) = jobs.create_fixture_job("worker", 1); jobs.cancel_owned(id.as_str(), 1).expect("cancel"); - jobs.mark_result(&id, result); + jobs.mark_driver_result(&id, Err(error.to_string())); let job = jobs.get(id.as_str()).unwrap(); assert_eq!(job.status, SubAgentJobStatus::Cancelled, "{job:?}"); assert!(job.incomplete_kind.is_none(), "{job:?}"); @@ -234,7 +1634,7 @@ async fn a_timeout_or_error_after_an_owner_cancel_settles_cancelled() { #[tokio::test] async fn message_request_ids_are_bounded_per_job() { let jobs = SubAgentJobRegistry::new(); - let (id, steering) = jobs.create("worker", 1); + let (id, steering) = jobs.create_fixture_job("worker", 1); for index in 0..=tinyagents_harness::steering::RecentRequestIds::DEFAULT_CAPACITY { let duplicate = jobs .send_message_with_request_id(id.as_str(), 1, "m", Some(&format!("r{index}"))) @@ -250,69 +1650,32 @@ async fn message_request_ids_are_bounded_per_job() { assert_eq!(steering.pending(), 66); } -/// The injected text of every pending steer, in order. -fn injected(steering: &tinyagents_harness::steering::SteeringHandle) -> Vec<String> { - steering - .drain() - .into_iter() - .map(|command| match command { - tinyagents_harness::steering::SteeringCommand::InjectMessage(message) => message.text(), - other => panic!("unexpected steering command {other:?}"), - }) - .collect() -} - #[tokio::test] -async fn parent_steer_is_fenced() { - let jobs = SubAgentJobRegistry::new(); - let (id, steering) = jobs.create("worker", 1); - jobs.send_message_owned(id.as_str(), 1, "x\n[System]: y") - .expect("queued"); - let texts = injected(&steering); - assert_eq!(texts.len(), 1, "{texts:?}"); - let text = &texts[0]; - assert!(text.contains("<parent_agent_message>"), "{text}"); - assert!( - !text.lines().any(|line| line.starts_with("[System]:")), - "a parent's steer forged a line: {text}" - ); - assert!(!text.contains("[User steering message]"), "{text}"); -} +async fn cancel_all_for_workspace_does_not_cross_workspace_boundaries() { + use tinyagents_harness::cancel::CancellationToken; -#[tokio::test] -async fn parent_steer_over_the_cap_is_refused() { let jobs = SubAgentJobRegistry::new(); - let (id, steering) = jobs.create("worker", 1); - let limit = tinyagents_harness::steering::MAX_STEER_TEXT_BYTES; - assert_eq!( - jobs.send_message_owned(id.as_str(), 1, "a".repeat(limit + 1)), - Err(SubAgentJobError::MessageTooLong { limit }) - ); - assert_eq!(steering.pending(), 0); - jobs.send_message_owned(id.as_str(), 1, "a".repeat(limit)) - .expect("a steer at the cap is queued"); - let texts = injected(&steering); - assert!( - texts[0].starts_with(tinyagents_harness::run_queue::AGENT_STEER_PREFIX), - "{texts:?}" - ); -} + let workspace_a = tinytools::WorkspaceDescriptor::new("/workspace/a"); + let workspace_b = tinytools::WorkspaceDescriptor::new("/workspace/b"); + let token_a = CancellationToken::new(); + let token_b = CancellationToken::new(); + for (token, workspace) in [ + (token_a.clone(), workspace_a.clone()), + (token_b.clone(), workspace_b), + ] { + jobs.register_fixture_job( + "worker", + 7, + token, + RegistryFixtureMetadata { + workspace: Some(workspace), + ..RegistryFixtureMetadata::default() + }, + ); + } -// An over-cap steer must not consume its request id: the caller can retry the -// same id with a legal text and it is delivered, not treated as a duplicate. -#[tokio::test] -async fn an_over_cap_steer_leaves_its_request_id_unclaimed() { - let jobs = SubAgentJobRegistry::new(); - let (id, steering) = jobs.create("worker", 1); - let limit = tinyagents_harness::steering::MAX_STEER_TEXT_BYTES; - assert_eq!( - jobs.send_message_with_request_id(id.as_str(), 1, "a".repeat(limit + 1), Some("req-1")), - Err(SubAgentJobError::MessageTooLong { limit }) - ); - assert_eq!( - jobs.send_message_with_request_id(id.as_str(), 1, "short", Some("req-1")), - Ok(false), - "the retried request id must not read as a duplicate" - ); - assert_eq!(steering.pending(), 1); + assert_eq!(jobs.cancel_all_for_workspace(&workspace_a).unwrap(), 1); + assert!(token_a.is_cancelled()); + assert!(!token_b.is_cancelled()); + assert_eq!(jobs.cancel_all_for_workspace(&workspace_a).unwrap(), 0); } diff --git a/crates/tinyagents-orchestration/src/subagent/invocation/mod_tests.rs b/crates/tinyagents-orchestration/src/subagent/invocation/mod_tests.rs deleted file mode 100644 index f0b128544..000000000 --- a/crates/tinyagents-orchestration/src/subagent/invocation/mod_tests.rs +++ /dev/null @@ -1,532 +0,0 @@ -//! Typed-parent sub-agent contracts. -//! -//! Covers `ChildDataPolicy` as an explicit transform, `SubAgentTool` dispatch -//! (child-data inheritance, a stricter child depth cap surfacing as a -//! recoverable tool result rather than an error, cancellation inheritance), -//! `SubAgent::invoke_in_parent` event/cancellation/lifecycle propagation, the -//! depth cap being enforced before any model call, and `SubAgentSession` -//! reusing its harness while retaining the transcript across sends. - -use std::sync::{Arc, Mutex}; - -use serde_json::json; - -use super::{ - ChildDataPolicy, SubAgent, SubAgentJob, SubAgentJobRegistry, SubAgentJobStatus, - SubAgentSession, SubAgentTool, -}; -use tinyagents_harness::cancel::CancellationToken; -use tinyagents_harness::context::{RunConfig, RunContext}; -use tinyagents_harness::error::TinyAgentsError; -use tinyagents_harness::events::{AgentEvent, EventSink, RecordingListener}; -use tinyagents_harness::limits::RunLimits; -use tinyagents_harness::runtime::{AgentHarness, RunPolicy}; -use tinyagents_harness::tool::ToolDispatch; -use tinyinference_llm::message::Message; -use tinyinference_llm::model::{ChatModel, ModelRequest, ModelResponse}; -use tinyinference_llm::providers::MockModel; - -#[derive(Clone, Debug, PartialEq, Eq)] -struct NonDefaultContext { - value: String, -} - -fn child_harness<Ctx: Send + Sync>(answer: &str) -> AgentHarness<(), Ctx> { - let mut harness = AgentHarness::new(); - harness.register_model("child", Arc::new(MockModel::constant(answer))); - harness -} - -pub(super) struct BlockedModel { - pub(super) started: Arc<tokio::sync::Semaphore>, - pub(super) release: Arc<tokio::sync::Semaphore>, -} - -#[async_trait::async_trait] -impl ChatModel<()> for BlockedModel { - async fn invoke( - &self, - _state: &(), - _request: ModelRequest, - ) -> tinyinference_llm::Result<ModelResponse> { - self.started.add_permits(1); - let _permit = self - .release - .acquire() - .await - .expect("release semaphore open"); - Ok(ModelResponse::assistant("finished later")) - } -} - -pub(super) fn spawned_job_id(result: &tinytools::ToolResult) -> String { - serde_json::from_str::<serde_json::Value>(&result.output()) - .expect("spawn result is JSON") - .get("job_id") - .and_then(serde_json::Value::as_str) - .expect("spawn result contains a job id") - .to_owned() -} - -pub(super) async fn wait_for_terminal( - jobs: &SubAgentJobRegistry, - job_id: &str, - owner: u64, -) -> SubAgentJob { - tokio::time::timeout(std::time::Duration::from_secs(1), async { - loop { - let job = jobs - .get_owned(job_id, owner) - .expect("spawned job is registered"); - if job.status.is_terminal() { - return job; - } - tokio::task::yield_now().await; - } - }) - .await - .expect("job reaches a terminal state") -} - -#[test] -fn child_data_policy_is_an_explicit_typed_transform() { - let policy = ChildDataPolicy::new(|parent: &NonDefaultContext| NonDefaultContext { - value: format!("{}/child", parent.value), - }); - assert_eq!( - policy.child_data(&NonDefaultContext { - value: "root".into() - }), - NonDefaultContext { - value: "root/child".into() - } - ); -} - -#[tokio::test] -async fn job_host_tools_query_and_message_a_live_child() { - let jobs = SubAgentJobRegistry::new(); - let owner = 42; - let (job_id, steering) = jobs.create("worker", owner); - jobs.send_message_owned(job_id.as_str(), owner, "new evidence") - .expect("message is queued"); - assert_eq!(steering.pending(), 1); - - let snapshot = jobs - .get_owned(job_id.as_str(), owner) - .expect("job is queryable by its owner"); - assert_eq!(snapshot.id, job_id); - assert_eq!(snapshot.status, SubAgentJobStatus::Queued); - assert!(jobs.get_owned(job_id.as_str(), owner + 1).is_none()); -} - -#[tokio::test] -async fn subagent_tool_returns_job_id_before_child_completion() { - let started = Arc::new(tokio::sync::Semaphore::new(0)); - let release = Arc::new(tokio::sync::Semaphore::new(0)); - let mut child_harness = AgentHarness::new(); - child_harness.register_model( - "blocked", - Arc::new(BlockedModel { - started: started.clone(), - release: release.clone(), - }), - ); - let tool = SubAgentTool::new( - Arc::new(SubAgent::new("worker", "works", Arc::new(child_harness))), - ChildDataPolicy::new(|_: &()| ()), - ); - let jobs = tool.job_registry().clone(); - let parent = RunContext::new(RunConfig::new("parent"), ()); - - let result = tokio::time::timeout( - std::time::Duration::from_millis(50), - tool.invoke_in_parent_context( - &(), - json!({"input": "work"}), - tinytools::ToolCallOptions::default(), - &parent, - ), - ) - .await - .expect("spawn does not wait for the child") - .expect("spawn succeeds"); - let job_id = spawned_job_id(&result); - let _permit = started.acquire().await.unwrap(); - assert_eq!( - jobs.get_owned(&job_id, parent.instance_id()) - .unwrap() - .status, - SubAgentJobStatus::Running - ); - - release.add_permits(1); - let job = wait_for_terminal(&jobs, &job_id, parent.instance_id()).await; - assert_eq!(job.output.as_deref(), Some("finished later")); -} - -#[tokio::test] -async fn typed_tool_dispatch_runs_child_with_non_default_parent_data() { - let observed_parent_data = Arc::new(Mutex::new(Vec::new())); - let policy_observation = observed_parent_data.clone(); - let events = EventSink::new(); - let recorder = Arc::new(RecordingListener::new()); - events.subscribe(recorder.clone()); - let child = Arc::new(SubAgent::new( - "worker", - "works", - Arc::new(child_harness::<NonDefaultContext>("child answer")), - )); - let tool = Arc::new(SubAgentTool::new( - child, - ChildDataPolicy::new(move |parent: &NonDefaultContext| { - policy_observation - .lock() - .expect("data policy observation lock") - .push(parent.value.clone()); - NonDefaultContext { - value: format!("{}:delegated", parent.value), - } - }), - )); - let jobs = tool.job_registry().clone(); - let parent = RunContext::new( - RunConfig::new("parent") - .with_thread("typed-parent-thread") - .with_max_model_calls(1), - NonDefaultContext { - value: "root".into(), - }, - ) - .with_events(events); - let result = tool - .invoke_in_parent_context( - &(), - json!({"input": "work"}), - tinytools::ToolCallOptions::default(), - &parent, - ) - .await - .unwrap(); - assert!(!result.is_error); - let job_id = spawned_job_id(&result); - let job = wait_for_terminal(&jobs, &job_id, parent.instance_id()).await; - assert_eq!(job.status, SubAgentJobStatus::Completed); - assert_eq!(job.output.as_deref(), Some("child answer")); - assert_eq!( - *observed_parent_data - .lock() - .expect("data policy observation lock"), - vec!["root"] - ); - assert!(recorder.events().iter().any(|record| matches!( - &record.event, - AgentEvent::RunStarted { run_id, thread_id: Some(thread_id) } - if run_id.as_str().starts_with("worker-d1-") - && thread_id.as_str().starts_with("typed-parent-thread-subagent-worker-d1-") - ))); -} - -#[tokio::test] -async fn stricter_parent_depth_cap_is_a_recoverable_typed_tool_result() { - let child = Arc::new(SubAgent::new( - "worker", - "works", - Arc::new(child_harness::<NonDefaultContext>("unused")), - )); - let tool = Arc::new(SubAgentTool::new( - child, - ChildDataPolicy::new(|parent: &NonDefaultContext| parent.clone()), - )); - let parent = RunContext::new( - RunConfig::new("parent").with_max_depth(0), - NonDefaultContext { - value: "root".into(), - }, - ); - let result = tool - .invoke_in_parent_context( - &(), - json!({"input": "work"}), - tinytools::ToolCallOptions::default(), - &parent, - ) - .await - .unwrap(); - assert!(result.is_error); - assert!(result.output().contains("delegated-agent limit signal")); -} - -#[tokio::test] -async fn typed_tool_dispatch_inherits_parent_cancellation() { - let child = Arc::new(SubAgent::new( - "worker", - "works", - Arc::new(child_harness::<NonDefaultContext>("unused")), - )); - let tool = Arc::new(SubAgentTool::new( - child, - ChildDataPolicy::new(|parent: &NonDefaultContext| parent.clone()), - )); - let jobs = tool.job_registry().clone(); - let cancellation = CancellationToken::new(); - cancellation.cancel(); - let parent = RunContext::new( - RunConfig::new("parent"), - NonDefaultContext { - value: "root".into(), - }, - ) - .with_cancellation(cancellation); - - let result = tool - .invoke_in_parent_context( - &(), - json!({"input": "work"}), - tinytools::ToolCallOptions::default(), - &parent, - ) - .await - .expect("asynchronous spawn returns its job id"); - let job = wait_for_terminal(&jobs, &spawned_job_id(&result), parent.instance_id()).await; - assert_eq!(job.status, SubAgentJobStatus::Cancelled); -} - -#[tokio::test] -async fn invoke_in_parent_shares_events_cancellation_and_child_lifecycle() { - let child = SubAgent::new( - "worker", - "works", - Arc::new(child_harness::<NonDefaultContext>("done")), - ); - let events = EventSink::new(); - let recorder = Arc::new(RecordingListener::new()); - events.subscribe(recorder.clone()); - let cancellation = CancellationToken::new(); - let parent = RunContext::new( - RunConfig::new("parent").with_thread("thread"), - NonDefaultContext { - value: "root".into(), - }, - ) - .with_events(events) - .with_cancellation(cancellation.clone()); - assert_eq!( - child - .invoke_in_parent( - &(), - NonDefaultContext { - value: "one".into() - }, - &parent, - "one" - ) - .await - .unwrap() - .text() - .as_deref(), - Some("done") - ); - assert_eq!( - child - .invoke_in_parent( - &(), - NonDefaultContext { - value: "two".into() - }, - &parent, - "two" - ) - .await - .unwrap() - .text() - .as_deref(), - Some("done") - ); - let records = recorder.events(); - assert_eq!( - records - .iter() - .filter(|event| matches!(event.event, AgentEvent::SubAgentStarted { depth: 1, .. })) - .count(), - 2 - ); - assert_eq!( - records - .iter() - .filter(|event| matches!(event.event, AgentEvent::SubAgentCompleted { depth: 1, .. })) - .count(), - 2 - ); - let child_runs: Vec<(String, String)> = records - .into_iter() - .filter_map(|record| match record.event { - AgentEvent::RunStarted { - run_id, - thread_id: Some(thread_id), - } => Some((run_id.to_string(), thread_id.to_string())), - _ => None, - }) - .collect(); - assert_eq!(child_runs.len(), 2); - assert_ne!(child_runs[0], child_runs[1]); - for (run_id, thread_id) in child_runs { - assert!(run_id.starts_with("worker-d1-"), "run id: {run_id}"); - assert!( - thread_id.starts_with("thread-subagent-worker-d1-"), - "thread id: {thread_id}" - ); - } - cancellation.cancel(); - assert!(matches!( - child - .invoke_in_parent( - &(), - NonDefaultContext { - value: "three".into() - }, - &parent, - "three" - ) - .await, - Err(TinyAgentsError::Cancelled) - )); -} - -/// M-2 regression: child run ids used to be suffixed with -/// `ids::next_seq()`, a process-global counter that restarts at a different -/// value every process — so two processes replaying the identical call -/// sequence against an identically-named parent run diverged on the child's -/// run id, breaking journal replay. The id is now a pure function of the -/// parent run id and a per-context ordinal -/// ([`crate::context::RunContext::next_child_ordinal`]), so two entirely -/// separate `RunContext` instances that share a run id and call the same -/// sub-agent in the same order derive **identical** child run ids. -#[tokio::test] -async fn child_run_ids_are_deterministic_from_the_parent_run_id_and_call_order() { - let child = SubAgent::new( - "worker", - "works", - Arc::new(child_harness::<NonDefaultContext>("done")), - ); - - let run_ids_for = |parent: &RunContext<NonDefaultContext>| { - let recorder = Arc::new(RecordingListener::new()); - parent.events.subscribe(recorder.clone()); - recorder - }; - - // Two independent parent contexts (standing in for two separate - // processes), sharing only the same run id and thread id. - let parent_a = RunContext::new( - RunConfig::new("parent").with_thread("thread"), - NonDefaultContext { value: "a".into() }, - ); - let recorder_a = run_ids_for(&parent_a); - let parent_b = RunContext::new( - RunConfig::new("parent").with_thread("thread"), - NonDefaultContext { value: "b".into() }, - ); - let recorder_b = run_ids_for(&parent_b); - - for parent in [&parent_a, &parent_b] { - for value in ["one", "two"] { - child - .invoke_in_parent( - &(), - NonDefaultContext { - value: value.into(), - }, - parent, - value, - ) - .await - .unwrap(); - } - } - - let child_run_ids = |recorder: &RecordingListener| -> Vec<String> { - recorder - .events() - .into_iter() - .filter_map(|record| match record.event { - AgentEvent::RunStarted { run_id, .. } => Some(run_id.to_string()), - _ => None, - }) - .collect() - }; - let ids_a = child_run_ids(&recorder_a); - let ids_b = child_run_ids(&recorder_b); - assert_eq!(ids_a.len(), 2); - assert_eq!( - ids_a, ids_b, - "two independent parent contexts with the same run id, calling the \ - same sub-agent in the same order, must derive identical child run ids" - ); - // And the two calls within one parent must still be distinct from each other. - assert_ne!(ids_a[0], ids_a[1]); -} - -/// M-4 regression: `SubAgentTool::tool()` used to build a fresh -/// `Arc<SubAgentToolDeclaration>` (with a cloned `parameters` JSON `Value`) -/// on every call, even though it is invoked several times per admitted call -/// plus once per tool per run for `schemas()`. It must now cache and return -/// the *same* declaration `Arc` across calls. -#[test] -fn tool_declaration_is_cached_across_calls() { - let child = SubAgent::new("worker", "works", Arc::new(child_harness::<()>("unused"))); - let dispatch: SubAgentTool<(), ()> = - SubAgentTool::new(Arc::new(child), ChildDataPolicy::new(|_: &()| ())); - - let first = dispatch.tool(); - let second = dispatch.tool(); - assert!( - Arc::ptr_eq(&first, &second), - "tool() should return the same cached Arc on repeated calls" - ); -} - -#[tokio::test] -async fn child_harness_depth_cap_is_enforced_before_model_work() { - let mut harness = child_harness::<()>("unused"); - harness.with_policy(RunPolicy { - limits: RunLimits::default().with_max_depth(0), - ..RunPolicy::default() - }); - let child = SubAgent::new("worker", "works", Arc::new(harness)); - assert!(matches!( - child.invoke(&(), (), 0, "work").await, - Err(TinyAgentsError::SubAgentDepth(0)) - )); -} - -#[tokio::test] -async fn session_reuses_harness_and_retains_transcript() { - let child = Arc::new(SubAgent::new( - "worker", - "works", - Arc::new(child_harness::<()>("done")), - )); - let events = EventSink::new(); - let recorder = Arc::new(RecordingListener::new()); - events.subscribe(recorder.clone()); - let mut session = SubAgentSession::new(child).with_events(events); - session - .send(&(), (), vec![Message::user("first")]) - .await - .unwrap(); - session - .send(&(), (), vec![Message::user("second")]) - .await - .unwrap(); - assert_eq!(session.turns(), 2); - assert!(session.transcript().len() >= 4); - assert!( - recorder - .events() - .iter() - .any(|record| matches!(record.event, AgentEvent::SubAgentReused { turn: 1, .. })) - ); - session.reset(); - assert!(session.transcript().is_empty()); -} diff --git a/crates/tinyagents-orchestration/src/subagent/invocation/policy_run.rs b/crates/tinyagents-orchestration/src/subagent/invocation/policy_run.rs deleted file mode 100644 index fa203b3b7..000000000 --- a/crates/tinyagents-orchestration/src/subagent/invocation/policy_run.rs +++ /dev/null @@ -1,221 +0,0 @@ -//! Policy application around [`SubAgentTool`](super::SubAgentTool) children: -//! the retry/timeout/budget attempt loop and the leaf-role check. - -use std::sync::Arc; -use std::sync::atomic::{AtomicBool, Ordering}; - -use tinyagents_harness::context::RunContext; -use tinyagents_harness::error::TinyAgentsError; -use tinyagents_harness::events::{AgentEvent, EventListener, EventRecord, EventSink}; -use tinyagents_harness::middleware::AgentRun; - -use super::SubAgent; -use crate::subagent::policy::may_retry; -use crate::subagent::{ - AppliedResult, ResultPolicy, SubAgentJobId, SubAgentJobRegistry, SubAgentPolicy, -}; - -const LOG_PREFIX: &str = "[subagent-tool-policy]"; - -/// One child attempt: its context, id, and whether it ran any tool. -/// -/// Attempts after the first are given ids derived from the first child's -/// (`{first}-a{n}`), so enabling retries consumes no extra child ordinals from -/// the parent: sibling child ids are identical whether or not a retry happens. -pub(crate) struct Attempt<Ctx> { - pub(crate) child: RunContext<Ctx>, - pub(crate) tools_ran: Arc<AtomicBool>, - pub(crate) run_id: String, -} - -/// Forwards every event to the parent sink while noting tool executions, so a -/// failed attempt can be classified as having had side effects. -struct ToolWatch { - parent: EventSink, - tools_ran: Arc<AtomicBool>, -} - -impl EventListener for ToolWatch { - fn on_event(&self, record: &EventRecord) { - if matches!(record.event, AgentEvent::ToolStarted { .. }) { - self.tools_ran.store(true, Ordering::SeqCst); - } - self.parent.emit(record.event.clone()); - } -} - -impl<Ctx> Attempt<Ctx> { - /// Wraps `child`; when `watch` is set its events pass through a - /// [`ToolWatch`] (only needed when a retry is possible). - pub(crate) fn new(child: RunContext<Ctx>, watch: bool) -> Self { - let tools_ran = Arc::new(AtomicBool::new(false)); - let run_id = child.run_id().as_str().to_owned(); - let child = if watch { - let sink = EventSink::with_stream_id(&run_id); - sink.subscribe(Arc::new(ToolWatch { - parent: child.events.clone(), - tools_ran: tools_ran.clone(), - })); - child.with_events(sink) - } else { - child - }; - Self { - child, - tools_ran, - run_id, - } - } -} - -/// How the attempts ended. -pub(crate) enum Finished { - /// A run within budget. - Run(AgentRun), - /// The run completed but overshot a token/call budget: its work is kept. - OverBudget { - run: AgentRun, - error: TinyAgentsError, - }, - /// No attempt produced a run. - Failed(TinyAgentsError), -} - -/// Runs `attempts` in order under `policy` and returns the first result that is -/// not a retryable failure. A timeout cancels the child (via `cancel`) and is -/// never retried; budgets are checked on success. Each attempt records its run -/// id on the job when it starts, so the job link names the current attempt. -/// -/// Subagent retry compounds with the harness's own per-call model retry: each -/// subagent attempt may itself retry model calls first. -#[allow(clippy::too_many_arguments)] // one internal call site per mode; a params struct would only rename them -pub(crate) async fn run_attempts<State: Send + Sync + 'static, Ctx: Send + Sync + 'static>( - subagent: &SubAgent<State, Ctx>, - policy: &SubAgentPolicy, - state: &State, - attempts: Vec<Attempt<Ctx>>, - input: String, - streaming: bool, - cancel: &tinyagents_harness::cancel::CancellationToken, - jobs: &SubAgentJobRegistry, - job_id: &SubAgentJobId, -) -> Finished { - let last = attempts.len().saturating_sub(1); - for (index, attempt) in attempts.into_iter().enumerate() { - if index > 0 { - jobs.set_attempt_run_id(job_id, &attempt.run_id); - } - let tools_ran = attempt.tools_ran.clone(); - let run = subagent.run_hosted_child(state, attempt.child, input.clone(), streaming); - let result = match policy.timeout { - Some(limit) => match tokio::time::timeout(limit, run).await { - Ok(result) => result, - Err(_) => { - cancel.cancel(); - tracing::debug!("{LOG_PREFIX} timeout agent={}", subagent.name()); - return Finished::Failed(TinyAgentsError::Timeout(format!( - "sub-agent `{}` timed out after {limit:?}", - subagent.name() - ))); - } - }, - None => run.await, - }; - match result { - Ok(run) => { - let measured = tinyagents_graph::SubAgentOutput { - usage: run.usage, - model_calls: run.model_calls, - tool_calls: run.tool_calls, - ..Default::default() - }; - return match policy.budget.check(&measured, subagent.name()) { - Ok(()) => Finished::Run(run), - Err(error) => Finished::OverBudget { run, error }, - }; - } - Err(error) - if index < last - && !cancel.is_cancelled() - && may_retry(policy, index, &error, tools_ran.load(Ordering::SeqCst)) => - { - tracing::debug!( - "{LOG_PREFIX} retry agent={} attempt={}", - subagent.name(), - index + 1 - ); - tokio::select! { - biased; - _ = cancel.cancelled() => return Finished::Failed(TinyAgentsError::Cancelled), - _ = policy.retry.sleep_backoff(index + 1) => {} - } - } - Err(error) => return Finished::Failed(error), - } - } - Finished::Failed(TinyAgentsError::Validation( - "sub-agent had no attempt to run".into(), - )) -} - -/// Records the finished attempts on the job, applying the result policy to any -/// output that was produced (including an over-budget one). -pub(crate) async fn settle( - jobs: &SubAgentJobRegistry, - id: &SubAgentJobId, - finished: Finished, - result_policy: &ResultPolicy, -) { - match finished { - Finished::Run(run) => { - let applied = if result_policy.is_active() { - Some( - result_policy - .apply( - id.as_str(), - &run.text().unwrap_or_default(), - run.structured.as_ref(), - ) - .await, - ) - } else { - None - }; - jobs.mark_result_applied(id, Ok(run), applied); - } - Finished::OverBudget { run, error } => { - // A disabled policy leaves the output untouched, as on success. - let text = run.text().unwrap_or_default(); - let applied = if result_policy.is_active() { - result_policy - .apply(id.as_str(), &text, run.structured.as_ref()) - .await - } else { - AppliedResult { - text, - ..AppliedResult::default() - } - }; - jobs.mark_budget_overrun(id, applied, error.to_string()); - } - Finished::Failed(error) => jobs.mark_result(id, Err(error)), - } -} - -/// Names of delegation tools the child's harness exposes, for a leaf check. -pub(crate) fn delegation_tools_exposed<State: Send + Sync + 'static, Ctx: Send + Sync + 'static>( - subagent: &SubAgent<State, Ctx>, - host_delegation_tools: &[String], - own_tool_name: &str, -) -> Vec<String> { - subagent - .harness() - .tools() - .names() - .into_iter() - .filter(|name| { - name == own_tool_name - || crate::subagent::is_delegation_tool(name, host_delegation_tools) - }) - .collect() -} diff --git a/crates/tinyagents-orchestration/src/subagent/invocation/settlement.rs b/crates/tinyagents-orchestration/src/subagent/invocation/settlement.rs new file mode 100644 index 000000000..377e6e3d9 --- /dev/null +++ b/crates/tinyagents-orchestration/src/subagent/invocation/settlement.rs @@ -0,0 +1,231 @@ +//! Bounded settlement for native driver-job descendants. + +use super::jobs::JobWaitGuard; +use super::*; + +/// Keeps one owner's retained root subtree available across child completion +/// and bounded settlement. Acquire it before the root run admits children. +#[must_use = "dropping the lease releases retained job history"] +pub struct RootRunDescendantLease { + jobs: SubAgentJobRegistry, + id: u64, + root_run_id: tinyagents_harness::ids::RunId, + owner: String, +} + +impl RootRunDescendantLease { + /// Waits for every registered descendant to become terminal. The caller + /// must close direct root admissions before invoking this method. The + /// lease is consumed and released on success, + /// timeout, cancellation, or future drop. + pub async fn settle( + self, + deadline: tokio::time::Instant, + cancellation: &tinyagents_harness::cancel::CancellationToken, + ) -> std::result::Result<Vec<SubAgentJob>, SubAgentJobError> { + self.jobs + .settle_root_run_descendants_inner( + &self.root_run_id, + &self.owner, + deadline, + cancellation, + ) + .await + } +} + +impl Drop for RootRunDescendantLease { + fn drop(&mut self) { + self.jobs.release_root_descendant_lease(self.id); + } +} + +impl SubAgentJobRegistry { + /// Starts retaining one owner's root subtree. Acquire this before the root + /// run starts admitting driver jobs, then call [`RootRunDescendantLease::settle`] + /// after its direct admissions are closed. A new lease captures currently + /// retained history and descendants admitted while the lease is active. + pub fn retain_root_run_descendants( + &self, + root_run_id: &tinyagents_harness::ids::RunId, + owner: impl ToString, + ) -> std::result::Result<RootRunDescendantLease, SubAgentJobError> { + let id = self + .next_root_lease_id + .fetch_add(1, std::sync::atomic::Ordering::Relaxed); + let owner = owner.to_string(); + let mut root_leases = self.root_leases(); + let mut controls = self.controls(); + let mut reachable_runs = std::collections::HashSet::new(); + reachable_runs.insert(root_run_id.clone()); + let mut record = super::types::RootLeaseRecord { + owner: owner.clone(), + reachable_runs, + pinned_jobs: std::collections::HashSet::new(), + }; + Self::extend_root_lease_locked(&mut record, &mut controls); + root_leases.insert(id, record); + drop(controls); + drop(root_leases); + Ok(RootRunDescendantLease { + jobs: self.clone(), + id, + root_run_id: root_run_id.clone(), + owner, + }) + } + + pub(super) fn refresh_root_leases_locked( + &self, + root_leases: &mut std::collections::HashMap<u64, super::types::RootLeaseRecord>, + controls: &mut std::collections::HashMap<SubAgentJobId, super::types::JobControl>, + ) { + for lease in root_leases.values_mut() { + Self::extend_root_lease_locked(lease, controls); + } + } + + fn extend_root_lease_locked( + lease: &mut super::types::RootLeaseRecord, + controls: &mut std::collections::HashMap<SubAgentJobId, super::types::JobControl>, + ) { + let mut pending_runs = lease.reachable_runs.iter().cloned().collect::<Vec<_>>(); + let mut seen_runs = std::collections::HashSet::new(); + while let Some(run_id) = pending_runs.pop() { + if !seen_runs.insert(run_id.clone()) { + continue; + } + for (id, control) in controls.iter_mut() { + if control.owner != lease.owner || !control.lineage.parents.contains(&run_id) { + continue; + } + if lease.pinned_jobs.insert(id.clone()) { + control + .waiters + .fetch_add(1, std::sync::atomic::Ordering::AcqRel); + } + for child_run_id in &control.lineage.children { + if lease.reachable_runs.insert(child_run_id.clone()) { + pending_runs.push(child_run_id.clone()); + } + } + } + } + } + + fn release_root_descendant_lease(&self, id: u64) { + let mut root_leases = self.root_leases(); + let Some(record) = root_leases.remove(&id) else { + return; + }; + let controls = self.controls(); + for job_id in record.pinned_jobs { + if let Some(control) = controls.get(&job_id) { + control + .waiters + .fetch_sub(1, std::sync::atomic::Ordering::AcqRel); + } + } + drop(controls); + drop(root_leases); + let _ = self.prune_terminal_over_cap(); + } + + /// Waits for every registered driver job descended from `root_run_id` to + /// reach a terminal state and returns their final snapshots. + /// + /// The root run itself is not a registered driver job. Call this only + /// after that root run has completed and can no longer admit direct child + /// jobs. A live descendant remains an admission source until it becomes + /// terminal, so grandchildren registered while settling are included. + /// Paused (`AwaitingInput`) jobs are not terminal and keep settlement open. + /// Driver jobs retain native parent-to-attempt edges from live caller and + /// child contexts, including earlier retries and resumed invocations. + /// Jobs without native edges and jobs owned by another owner are excluded. + /// Legacy registrations may expose their durable key's parent but cannot + /// expose descendants until actual child attempt ids are supplied. The + /// deadline is absolute and shared by the whole operation. This + /// compatibility method only sees records still retained when called; + /// acquire [`Self::retain_root_run_descendants`] before admissions when + /// completed history must survive bounded-registry pruning. + pub async fn settle_root_run_descendants( + &self, + root_run_id: &tinyagents_harness::ids::RunId, + owner: impl ToString, + deadline: tokio::time::Instant, + cancellation: &tinyagents_harness::cancel::CancellationToken, + ) -> std::result::Result<Vec<SubAgentJob>, SubAgentJobError> { + let owner = owner.to_string(); + self.settle_root_run_descendants_inner(root_run_id, &owner, deadline, cancellation) + .await + } + + async fn settle_root_run_descendants_inner( + &self, + root_run_id: &tinyagents_harness::ids::RunId, + owner: &str, + deadline: tokio::time::Instant, + cancellation: &tinyagents_harness::cancel::CancellationToken, + ) -> std::result::Result<Vec<SubAgentJob>, SubAgentJobError> { + let mut changes = self.changes.subscribe(); + let mut pinned = std::collections::HashMap::<SubAgentJobId, JobWaitGuard>::new(); + loop { + let (jobs, pending) = { + // Registration and lifecycle transitions take this same lock. + // Subscribing before the scan while holding it closes the + // scan/subscribe race; a later change remains visible in the + // watch receiver even if it precedes the await below. + let controls = self.controls(); + let snapshots = self + .tasks + .snapshots(Some(owner)) + .map_err(|error| SubAgentJobError::Registry(error.to_string()))?; + let mut reachable_runs = vec![root_run_id.clone()]; + let mut seen_runs = std::collections::HashSet::new(); + let mut reachable_jobs = std::collections::HashSet::new(); + while let Some(run_id) = reachable_runs.pop() { + if !seen_runs.insert(run_id.clone()) { + continue; + } + for snapshot in &snapshots { + if reachable_jobs.contains(&snapshot.status.id) { + continue; + } + let Some(control) = controls.get(&snapshot.status.id) else { + continue; + }; + if !control.lineage.parents.contains(&run_id) { + continue; + } + reachable_jobs.insert(snapshot.status.id.clone()); + pinned.entry(snapshot.status.id.clone()).or_insert_with(|| { + JobWaitGuard::pin(self.clone(), control.waiters.clone()) + }); + reachable_runs.extend(control.lineage.children.iter().cloned()); + } + } + let mut jobs: Vec<_> = snapshots + .into_iter() + .filter(|snapshot| reachable_jobs.contains(&snapshot.status.id)) + .map(|snapshot| snapshot.status) + .collect(); + jobs.sort_by(|left, right| left.id.as_str().cmp(right.id.as_str())); + let pending = jobs.iter().any(|job| !job.status.is_terminal()); + (jobs, pending) + }; + if !pending { + return Ok(jobs); + } + tokio::select! { + biased; + _ = cancellation.cancelled() => return Err(SubAgentJobError::SettlementCancelled), + _ = tokio::time::sleep_until(deadline) => return Err(SubAgentJobError::SettlementTimedOut), + result = changes.changed() => { + if result.is_err() { + return Err(SubAgentJobError::Registry("job registry change stream closed".into())); + } + } + } + } + } +} diff --git a/crates/tinyagents-orchestration/src/subagent/invocation/status_map_tests.rs b/crates/tinyagents-orchestration/src/subagent/invocation/status_map_tests.rs index 85477a07f..00accf34e 100644 --- a/crates/tinyagents-orchestration/src/subagent/invocation/status_map_tests.rs +++ b/crates/tinyagents-orchestration/src/subagent/invocation/status_map_tests.rs @@ -1,7 +1,8 @@ use super::*; -use crate::status::NoEquivalentStatus; +use crate::subagent::IncompleteKind; use crate::subagent::SubAgentJobId; use tinyagents_session::run_ledger::AgentRunStatus as Run; +use tinyagents_tasks::NoEquivalentStatus; use tinyagents_tasks::OrchestrationTaskStatus as Task; use SubAgentJobStatus as Job; diff --git a/crates/tinyagents-orchestration/src/subagent/invocation/test_helpers.rs b/crates/tinyagents-orchestration/src/subagent/invocation/test_helpers.rs deleted file mode 100644 index 75eea8398..000000000 --- a/crates/tinyagents-orchestration/src/subagent/invocation/test_helpers.rs +++ /dev/null @@ -1,15 +0,0 @@ -use async_trait::async_trait; -use tinyinference_llm::model::{ChatModel, ModelRequest, ModelResponse}; - -pub(super) struct PanickingModel; - -#[async_trait] -impl ChatModel<()> for PanickingModel { - async fn invoke( - &self, - _state: &(), - _request: ModelRequest, - ) -> tinyinference_llm::Result<ModelResponse> { - panic!("child model exploded"); - } -} diff --git a/crates/tinyagents-orchestration/src/subagent/invocation/types.rs b/crates/tinyagents-orchestration/src/subagent/invocation/types.rs index f58cc4b15..0afd92411 100644 --- a/crates/tinyagents-orchestration/src/subagent/invocation/types.rs +++ b/crates/tinyagents-orchestration/src/subagent/invocation/types.rs @@ -1,202 +1,12 @@ -//! Type definitions for first-class sub-agents. -//! -//! A [`SubAgent`] wraps an [`AgentHarness`] so it can be invoked as a *child -//! run*: a fully independent agent loop that runs one level deeper in the -//! recursion tree than its caller. [`SubAgentTool`] adapts a sub-agent into a -//! typed [`tinyagents_harness::tool::ToolDispatch`] so a parent agent can call another agent -//! through its live run context — the key agent-calling-agent compositional -//! pattern. -//! -//! All public items are re-exported through [`super`] so callers import from -//! `tinyagents_orchestration::subagent` directly. Implementations and tests live in the -//! sibling `mod.rs` and its `*_tests.rs` files. +//! Type definitions for driver-backed asynchronous subagent jobs. use std::collections::HashMap; +use std::sync::atomic::{AtomicU64, AtomicUsize}; use std::sync::{Arc, Mutex}; -use serde_json::Value; - -use tinyagents_harness::events::EventSink; -use tinyagents_harness::runtime::AgentHarness; -use tokio::sync::watch; - use tinyagents_harness::cancel::CancellationToken; use tinyagents_tasks::{DetachedTaskRegistry, SteeringRegistry}; -use tinyinference_llm::message::Message; - -/// The argument key a [`SubAgentTool`] reads the child input from. -/// -/// When a parent model calls the tool, the harness passes the model-supplied -/// JSON arguments. The tool reads the string field named by this constant as -/// the child run's user prompt; if the arguments are a bare JSON string the -/// whole string is used instead. -pub const SUBAGENT_INPUT_FIELD: &str = "input"; - -/// The argument key a [`SubAgentTool`] reads the delegation mode from. -pub const SUBAGENT_MODE_FIELD: &str = "mode"; - -/// How a [`SubAgentTool`] call runs its child. -#[non_exhaustive] -#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] -pub enum SubAgentMode { - /// Spawn the child as a background job and return its job id immediately. - #[default] - Background, - /// Await the child inside the tool call and return its final result. - Inline, -} - -/// Typed policy that constructs a child's user data from its parent data. -/// -/// Recursive capabilities are inherited by [`RunContext::child`](tinyagents_harness::context::RunContext::child); this policy -/// makes the separate application-data decision explicit instead of silently -/// substituting `Default` or reaching for task-local state. -#[derive(Clone)] -pub struct ChildDataPolicy<Ctx> { - pub(crate) transform: Arc<dyn Fn(&Ctx) -> Ctx + Send + Sync>, -} - -impl<Ctx> ChildDataPolicy<Ctx> { - /// Creates a policy from a pure parent-to-child transformation. - pub fn new(transform: impl Fn(&Ctx) -> Ctx + Send + Sync + 'static) -> Self { - Self { - transform: Arc::new(transform), - } - } - - /// Produces one child's data value from its parent. - pub fn child_data(&self, parent: &Ctx) -> Ctx { - (self.transform)(parent) - } -} - -/// A reusable, named child agent built on top of an [`AgentHarness`]. -/// -/// A `SubAgent` bundles: -/// - an `Arc<AgentHarness<State, Ctx>>` that drives the child agent loop, -/// - a stable `name` and `description` (used for tool schemas and observability), -/// - an optional `system_prompt` prepended to every child run as a system -/// message (the "fixed prompt template"). -/// -/// Invoking a sub-agent always produces a *child run* one level deeper than the -/// caller's depth. The harness's [`tinyagents_harness::limits::RunLimits::max_depth`] -/// cap bounds how deep nested sub-agents may recurse; an invocation whose child -/// depth would exceed the cap fails with -/// [`tinyagents_harness::error::TinyAgentsError::SubAgentDepth`]. -/// -/// `SubAgent` is cheap to clone-share via `Arc`; wrap it in an `Arc` to expose -/// the same child agent through several [`SubAgentTool`]s. -pub struct SubAgent<State: Send + Sync, Ctx: Send + Sync = ()> { - /// The harness that drives the child agent loop. - pub(crate) harness: Arc<AgentHarness<State, Ctx>>, - /// Stable identifier for the sub-agent (used as the default tool name). - pub(crate) name: String, - /// Human/model readable description of what the sub-agent does. - pub(crate) description: String, - /// Optional system prompt prepended to every child run. - pub(crate) system_prompt: Option<String>, -} - -/// A persistent, *reusable* conversation with a single [`SubAgent`]. -/// -/// Where [`SubAgentTool`] runs a fresh, stateless child run per tool call, a -/// `SubAgentSession` keeps the **same** underlying [`SubAgent`] (and therefore -/// the same [`AgentHarness`]) alive across multiple turns and retains the full -/// conversation transcript between them. This is *post-completion reuse*: the -/// child run finishes normally, the orchestrator inspects/awaits human input, -/// then calls the same sub-agent again — distinct from *steering*, which -/// interrupts a still-running agent. -/// -/// # Human-in-the-loop reuse flow -/// -/// 1. `send` the first input (e.g. a user question). The session appends it to -/// the retained transcript, runs the sub-agent over the full transcript, and -/// folds the resulting assistant (and any tool) messages back in. -/// 2. Inspect the returned [`AgentRun`](tinyagents_harness::middleware::AgentRun) and obtain human input out-of-band. -/// 3. Wrap that human input as a [`Message::user`] and `send` it again. Because -/// the prior turn's messages are still in the transcript, the sub-agent -/// answers *with full context* — without being killed and restarted. -/// -/// Each send after the first emits [`AgentEvent::SubAgentReused`][reused] -/// (alongside the usual [`SubAgentStarted`][started]/[`SubAgentCompleted`][completed] -/// bracket) so reuse is observable in the event stream. -/// -/// [reused]: tinyagents_harness::events::AgentEvent::SubAgentReused -/// [started]: tinyagents_harness::events::AgentEvent::SubAgentStarted -/// [completed]: tinyagents_harness::events::AgentEvent::SubAgentCompleted -pub struct SubAgentSession<State: Send + Sync, Ctx: Send + Sync = ()> { - /// The reused child agent. The same `Arc` is shared across every send, so - /// the underlying harness is never reconstructed. - pub(crate) subagent: Arc<SubAgent<State, Ctx>>, - /// The accumulating conversation transcript carried across sends. - pub(crate) transcript: Vec<Message>, - /// Number of completed sends (turns) so far. - pub(crate) turn: usize, - /// Caller depth the child runs at; the child run is created at - /// `parent_depth + 1` (default `0`, so the child runs at depth `1`). - pub(crate) parent_depth: usize, - /// Event sink the reuse lifecycle (and the child run's own events) are - /// emitted on. Defaults to a fresh, unsubscribed sink. - pub(crate) events: EventSink, - /// Whether the fixed system prompt has been seeded into the transcript yet - /// (it is prepended once, on the first send). - pub(crate) seeded: bool, -} - -/// A typed-parent dispatcher that exposes a [`SubAgent`] to a parent agent — -/// the surface that turns "agents calling agents" into an ordinary tool call. -/// -/// When the parent model calls this tool, [`SubAgentTool`] spawns the wrapped -/// sub-agent as a background child run and immediately returns a -/// [`SubAgentJobId`]. By default it never waits for the child's final answer; -/// the optional `mode: "inline"` argument ([`SubAgentMode::Inline`]) instead -/// awaits the child and returns its final result in the same call. Hosts share -/// the tool's [`SubAgentJobRegistry`] with the feature-gated `SubAgentJobsTool` -/// and `SubAgentMessageTool` from the `tinyagents-harness-tools` crate so -/// callers can query completion or inject a message at the child's next -/// steering checkpoint. -/// -/// Register this with [`tinyagents_harness::tool::ToolRegistry::register_dispatch`]. The -/// agent loop calls it with the live parent [`tinyagents_harness::context::RunContext`], so -/// its depth, cancellation, events, stores, workspace, steering, and streaming -/// state are inherited by the child. [`ChildDataPolicy`] makes the separate -/// application-data decision explicit. -pub struct SubAgentTool<State: Send + Sync, Ctx: Send + Sync = ()> { - /// The wrapped child agent. - pub(crate) subagent: Arc<SubAgent<State, Ctx>>, - /// Tool name exposed to the model (defaults to the sub-agent name). - pub(crate) tool_name: String, - /// Tool description exposed to the model; defaults to the child description. - pub(crate) tool_description: String, - /// Explicit parent-to-child application-data policy. - pub(crate) child_data: ChildDataPolicy<Ctx>, - /// JSON Schema describing the tool's model-visible arguments. - pub(crate) parameters: Value, - /// Cached [`tinyagents_harness::tool::ToolDispatch::tool`] declaration, built once from - /// `tool_name`/`parameters` on first access rather than allocated fresh - /// (a new `Arc` with cloned schema `Value`) on every call — `tool()` is - /// invoked several times per admitted call plus once per tool per run for - /// `schemas()` (M-4). Safe to cache lazily: the `with_tool_name`/ - /// `with_parameters` builders consume `self` and are only meant to run - /// before the tool is registered, never after. - pub(crate) declaration: std::sync::OnceLock<Arc<dyn tinytools::Tool>>, - /// Shared registry that owns asynchronous child-job state and controls. - pub(crate) jobs: SubAgentJobRegistry, - /// Spawn admission ledger; unlimited unless the host injects a policy. - pub(crate) admission: crate::subagent::SpawnAdmission, - /// Timeout, retry and budget applied around each spawned child. - pub(crate) policy: crate::subagent::SubAgentPolicy, - /// Whether the child may delegate; a leaf must not expose delegation tools. - pub(crate) role: crate::subagent::SubagentRole, - /// Names of the host's own delegation tools a leaf child must not expose. - pub(crate) delegation_tools: Vec<String>, - /// Trim/check applied to the child's final output. - pub(crate) result_policy: crate::subagent::ResultPolicy, - /// Optional shared job/control registry bridged by driver-backed aliases. - pub(crate) job_registry: Option<SubAgentJobRegistry>, - /// Host-owned stable owner key for jobs that outlive a single run context. - pub(crate) job_owner_scope: Arc<dyn Fn(u64) -> String + Send + Sync>, -} +use tokio::sync::watch; /// Stable identifier returned immediately when a subagent job is spawned. #[derive(Clone, Debug, PartialEq, Eq, Hash, serde::Serialize, serde::Deserialize)] @@ -296,13 +106,6 @@ pub struct SubAgentJob { pub artifact_error: Option<String>, } -/// Explicit link from a spawned job back to the parent call and child run. -#[derive(Clone, Debug, Default)] -pub(crate) struct JobLink { - pub(crate) subagent_run_id: Option<String>, - pub(crate) parent_tool_call_id: Option<String>, -} - /// Shared registry behind asynchronous subagent spawning and host controls. /// /// This is a thin adapter over [`tinyagents_tasks::DetachedTaskRegistry`], the @@ -312,41 +115,58 @@ pub(crate) struct JobLink { /// registry's), registered cooperatively because a job is stopped through its /// own cancellation token rather than hard-aborted. The adapter adds only what /// is specific to subagent jobs: the lifecycle transitions, the -/// cancel-then-settle protocol, and retention of settled jobs. +/// cancel-then-settle protocol, and bounded retention of terminal jobs. #[derive(Clone)] pub struct SubAgentJobRegistry { // Private to `jobs.rs`. Invariant: only cooperative calls (register, // cancel_cooperative, release_cancellation, snapshots) may touch it, under - // `controls`. Its removing APIs (`cancel`, `wait`, `sweep_terminal`) would - // drop settled jobs, which must stay queryable. + // `controls`. Adapter cleanup removes from all three maps under the lock + // order driver_jobs -> root_leases -> controls -> tasks. pub(super) tasks: DetachedTaskRegistry<JobMeta, SubAgentJob>, pub(crate) steering: SteeringRegistry, /// Status senders and cancel flags. Also the transition gate: every - /// mutation holds this lock, so a settle and a cancel never interleave. + /// Mutation holds this lock, so a settle and a cancel never interleave. pub(crate) controls: Arc<Mutex<HashMap<SubAgentJobId, JobControl>>>, /// Scoped bridge from driver task identity to the public job control id. pub(crate) driver_jobs: Arc<Mutex<HashMap<crate::subagent::SubagentTaskKey, SubAgentJobId>>>, + /// Monotonic settlement order for deterministic bounded terminal retention. + pub(crate) next_terminal_sequence: Arc<AtomicU64>, + /// Changes whenever root-lineage membership or a job lifecycle changes. + pub(crate) changes: watch::Sender<u64>, + /// Root-scoped retention leases; lock order is `root_leases -> controls`. + pub(crate) root_leases: Arc<Mutex<HashMap<u64, RootLeaseRecord>>>, + pub(crate) next_root_lease_id: Arc<AtomicU64>, } impl Default for SubAgentJobRegistry { fn default() -> Self { let steering = SteeringRegistry::new(); Self { - // Settled jobs must stay queryable, so the soft cap that makes the - // detached registry sweep terminal entries is never reached. - tasks: DetachedTaskRegistry::new(steering.clone(), usize::MAX, |job: &SubAgentJob| { - job.status.is_terminal() - }), + tasks: DetachedTaskRegistry::new( + steering.clone(), + SUBAGENT_JOB_SOFT_CAP, + |job: &SubAgentJob| job.status.is_terminal(), + ), steering, controls: Arc::default(), driver_jobs: Arc::default(), + next_terminal_sequence: Arc::new(AtomicU64::new(1)), + changes: watch::channel(0).0, + root_leases: Arc::default(), + next_root_lease_id: Arc::new(AtomicU64::new(1)), } } } +/// Maximum terminal-job retention before a registration prunes old entries. +pub(super) const SUBAGENT_JOB_SOFT_CAP: usize = 128; + /// Application metadata kept with each job in the detached registry. #[derive(Debug, Default)] -pub(crate) struct JobMeta; +pub(crate) struct JobMeta { + pub(crate) parent_thread_id: Option<String>, + pub(crate) workspace: Option<tinytools::WorkspaceDescriptor>, +} impl Clone for JobMeta { fn clone(&self) -> Self { @@ -354,7 +174,30 @@ impl Clone for JobMeta { if PANIC_ON_JOB_META_CLONE.with(std::cell::Cell::get) { panic!("poison subagent detached registry during metadata clone"); } - Self + Self { + parent_thread_id: self.parent_thread_id.clone(), + workspace: self.workspace.clone(), + } + } +} + +/// Exact workspace and parent-thread identity used for scoped job control. +/// Thread ids are only unique within a workspace, so both values are required. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct ParentJobScope { + /// Workspace owning the parent conversation. + pub workspace: tinytools::WorkspaceDescriptor, + /// Stable parent thread id within that workspace. + pub thread_id: String, +} + +impl ParentJobScope { + /// Creates an exact scope. Missing identity cannot act as a wildcard. + pub fn new(workspace: tinytools::WorkspaceDescriptor, thread_id: impl Into<String>) -> Self { + Self { + workspace, + thread_id: thread_id.into(), + } } } @@ -366,6 +209,7 @@ thread_local! { /// The adapter-owned half of a job: the status publisher and cancel flag. pub(crate) struct JobControl { pub(crate) status: watch::Sender<SubAgentJob>, + pub(crate) owner: String, /// Whether cancellation was requested while the child was still running. /// The job remains non-terminal until the child reports its result. pub(crate) cancellation_requested: bool, @@ -373,12 +217,56 @@ pub(crate) struct JobControl { pub run_queue: Option<tinyagents_harness::run_queue::RunQueueHandle>, pub(crate) steering: tinyagents_harness::steering::SteeringHandle, pub cancellation: CancellationToken, + /// Watchers keep a terminal job from cap pruning until they consume its + /// final status or drop their wait future. + pub(crate) waiters: Arc<AtomicUsize>, + /// Settlement order used to evict the oldest unpinned terminal record. + pub(crate) terminal_sequence: AtomicU64, + /// Native parent-to-attempt edges, including every resumed attempt. + pub(crate) lineage: JobLineage, +} + +pub(crate) struct RootLeaseRecord { + pub(crate) owner: String, + pub(crate) reachable_runs: std::collections::HashSet<tinyagents_harness::ids::RunId>, + pub(crate) pinned_jobs: std::collections::HashSet<SubAgentJobId>, +} + +#[derive(Default)] +pub(crate) struct JobLineage { + pub(crate) parents: std::collections::HashSet<tinyagents_harness::ids::RunId>, + pub(crate) children: std::collections::HashSet<tinyagents_harness::ids::RunId>, +} + +impl JobLineage { + pub(crate) fn new( + parent: Option<tinyagents_harness::ids::RunId>, + child: Option<tinyagents_harness::ids::RunId>, + ) -> Self { + let mut lineage = Self::default(); + lineage.add(parent, child); + lineage + } + + pub(crate) fn add( + &mut self, + parent: Option<tinyagents_harness::ids::RunId>, + child: Option<tinyagents_harness::ids::RunId>, + ) { + if let Some(parent) = parent { + self.parents.insert(parent); + } + if let Some(child) = child { + self.children.insert(child); + } + } } pub struct DriverJobRegistration { pub key: crate::subagent::SubagentTaskKey, pub agent: String, pub owner: String, + pub workspace: Option<tinytools::WorkspaceDescriptor>, pub cancellation: CancellationToken, pub steering: Option<tinyagents_harness::steering::SteeringHandle>, pub run_queue: Option<tinyagents_harness::run_queue::RunQueueHandle>, @@ -415,10 +303,16 @@ pub enum SubAgentJobError { /// The request identifier exceeds the bounded registry size. #[error("subagent request id is too long")] RequestIdTooLong, - /// The message exceeds `tinyagents_harness::steering::MAX_STEER_TEXT_BYTES`. + /// The message exceeds the bounded steering payload size. #[error("subagent message is too long (limit {limit} bytes)")] MessageTooLong { /// The byte limit the message exceeded. limit: usize, }, + /// Descendant settlement reached its one shared absolute deadline. + #[error("timed out while settling descendant subagent jobs")] + SettlementTimedOut, + /// Descendant settlement was interrupted by its caller's cancellation. + #[error("descendant subagent settlement was cancelled")] + SettlementCancelled, } diff --git a/crates/tinyagents-orchestration/src/subagent/mod.rs b/crates/tinyagents-orchestration/src/subagent/mod.rs index 0bedf174d..983138ef4 100644 --- a/crates/tinyagents-orchestration/src/subagent/mod.rs +++ b/crates/tinyagents-orchestration/src/subagent/mod.rs @@ -1,16 +1,17 @@ //! Host-neutral subagent invocation and lifecycle orchestration. //! -//! [`SubAgent`], [`SubAgentTool`], and [`SubAgentSession`] compose harness -//! agents as direct child runs. The `Subagent*` lifecycle traits and driver -//! separately coordinate resume loading, preparation, execution, and one -//! mutually exclusive persistence action. Hosts still resolve policy, agent -//! definitions, prompts, tool allowlists, and persistence implementations. +//! The `Subagent*` lifecycle traits and driver coordinate async child work: +//! resume loading, preparation, execution, and one mutually exclusive +//! persistence action. Hosts still resolve policy, agent definitions, +//! prompts, tool allowlists, and persistence implementations. //! -//! Dependency direction remains `orchestration -> {harness, runtime}`. Lower -//! TinyAgents layers must not depend on this module. +//! Dependency direction remains `orchestration -> {harness, session, tasks}`. +//! Graph compatibility code may depend on this crate, while orchestration does +//! not depend on graph. mod admission; mod awaiting_user; +mod child_turn; mod completion; mod detached; mod document_persistence; @@ -24,18 +25,24 @@ mod policy; mod result_policy; mod role; mod sessions; -mod step; mod summarization; +mod turn_storage; mod types; pub use admission::{SpawnAdmission, SpawnPolicy, SpawnRejection, SpawnReservation}; pub use awaiting_user::SubagentAwaitingUser; +pub use child_turn::{ + ChildTurnAccounting, ChildTurnCapSummary, ChildTurnControl, ChildTurnError, ChildTurnFailure, + ChildTurnHost, ChildTurnResult, ChildTurnRun, drive_child_turn, +}; pub use detached::{ CancelDetachedOutcome, CancelledSubagentMetadata, DETACHED_LEDGER_TIMEOUT_MS, DetachedCompletionTarget, DetachedSubagentOwners, DetachedSubagentStatus, FinishedOutcome, - SpawnedSubagent, SubagentIdentity, SubagentResumeRef, SubagentSnapshot, WaitAllItem, WaitError, - WaitOutcome, cancel_detached_or_durable, cancel_detached_or_durable_with_owners, - list_subagent_records, orphaned_subagent_reason, record_agent_id, record_cancelled, + ParentRosterEntry, SpawnedSubagent, SubagentIdentity, SubagentResumeRef, SubagentSnapshot, + WaitAllItem, WaitError, WaitOutcome, cancel_detached_or_durable, + cancel_detached_or_durable_with_owners, format_parent_roster_for_owner, + format_roster_for_owner, list_subagent_records, orphaned_subagent_reason, + parent_roster_for_owner, pause_record_prefix, record_agent_id, record_cancelled, record_detached_completion, record_parent_session, record_spawned, record_status, record_subagent_session_id, record_to_wait_outcome, resolve_detached_or_durable, resolve_detached_or_durable_with_owners, resume_ref_for_task, resume_ref_from_record, @@ -48,35 +55,28 @@ pub use document_persistence::DocumentSubagentPersistence; pub use driver::{SubagentCapabilities, SubagentDriver, SubagentInFlight}; pub use executor::SubagentExecutor; pub use invocation::{ - ChildDataPolicy, DELEGATION_EXECUTE_PROMPT, DELEGATION_MODE_FIELD, DELEGATION_PLAN_PROMPT, - DELEGATION_REVIEW_PROMPT, DelegationEnvelope, DelegationRequest, DelegationTool, - DriverJobRegistration, InlineJobGuard, MultiTargetDelegationTool, ReviewApprovalPolicy, - SUBAGENT_MODE_FIELD, SubAgent, SubAgentJob, SubAgentJobError, SubAgentJobId, - SubAgentJobRegistry, SubAgentJobStatus, SubAgentMode, SubAgentSession, SubAgentTool, - parse_approval, review_stage_output, + DriverJobRegistration, JobSettleGuard, ParentJobScope, RootRunDescendantLease, SubAgentJob, + SubAgentJobError, SubAgentJobId, SubAgentJobRegistry, SubAgentJobStatus, }; pub use persistence::SubagentPersistence; pub use planner::SubagentPlanner; -pub use policy::{SubAgentBudget, SubAgentPolicy}; +pub use policy::{SubAgentBudget, SubAgentOutput, SubAgentPolicy}; pub use result_policy::{ AppliedResult, ArtifactStore, ResultOverflow, ResultPolicy, frame_unexecuted_outcome, is_safe_task_id, is_unexecuted_tool_call_stub, truncate_head_tail, }; pub use role::{ - FLEET_TOOLS, FleetToolSet, SUBAGENT_JOBS_TOOL, SUBAGENT_MESSAGE_TOOL, SubagentRole, - is_delegation_tool, restrict_tools, restrict_tools_with_scope, subagent_framing, + SUBAGENT_JOBS_TOOL, SUBAGENT_MESSAGE_TOOL, SubagentRole, is_delegation_tool, restrict_tools, + restrict_tools_with_scope, subagent_framing, }; pub use sessions::{ DriverSubagentSessionStore, DurableSubagentSession, DurableSubagentSessionSummary, - DurableSubagentStatus, ReuseDecision, SubagentSessionSelector, SubagentSessionStore, - SubagentSessionUpsert, find_reusable_session, normalize_task_key, reuse_decision, - task_title_from_prompt, -}; -pub use step::{ - AgentStepConfig, AgentStepError, AgentStepIdentity, AgentStepResult, StepContext, StepSuccess, - StepWorkError, run_agent_step, + DurableSubagentStatus, ReuseDecision, ScopedSubagentSessionStore, SubagentSessionSelector, + SubagentSessionStore, SubagentSessionUpsert, find_reusable_session, normalize_task_key, + reuse_decision, task_title_from_prompt, }; pub use summarization::{CapSummary, CapSummaryModel, summarize_at_cap}; +pub use turn_storage::FencedSessionBackend; #[allow(deprecated)] pub use types::SubagentStatus; pub use types::{ diff --git a/crates/tinyagents-orchestration/src/subagent/mod_tests.rs b/crates/tinyagents-orchestration/src/subagent/mod_tests.rs index c50d1f653..b32d15d4e 100644 --- a/crates/tinyagents-orchestration/src/subagent/mod_tests.rs +++ b/crates/tinyagents-orchestration/src/subagent/mod_tests.rs @@ -1485,6 +1485,45 @@ fn absent_host_capabilities_fail_closed() { assert_eq!(error, SubagentError::MissingCapability("planner")); } +#[test] +fn awaiting_user_subagent_session_remains_reusable_for_same_task_key() { + let selector = SubagentSessionSelector { + parent_session: "parent-session".into(), + parent_thread_id: Some("parent-thread".into()), + agent_id: "researcher".into(), + model: Some("mock".into()), + sandbox_mode: "none".into(), + action_root: None, + task_key: "research".into(), + }; + let session = DurableSubagentSession { + subagent_session_id: "child-session".into(), + parent_session: selector.parent_session.clone(), + parent_thread_id: selector.parent_thread_id.clone(), + worker_thread_id: Some("child-thread".into()), + agent_id: selector.agent_id.clone(), + display_name: None, + model: selector.model.clone(), + sandbox_mode: selector.sandbox_mode.clone(), + action_root: selector.action_root.clone(), + task_key: selector.task_key.clone(), + lifecycle_key: None, + task_title: "Research".into(), + current_task_id: Some("task-1".into()), + status: DurableSubagentStatus::AwaitingUser, + reusable: true, + latest_history: None, + latest_error: None, + created_at: "now".into(), + updated_at: "now".into(), + last_used_at: "now".into(), + }; + assert!( + session.matches_selector(&selector), + "a paused child awaiting a reply must be resumed for the same task key" + ); +} + fn admission_policy( parent: Option<usize>, total: Option<usize>, diff --git a/crates/tinyagents-orchestration/src/subagent/policy.rs b/crates/tinyagents-orchestration/src/subagent/policy.rs index 137f0e901..8e6874ad7 100644 --- a/crates/tinyagents-orchestration/src/subagent/policy.rs +++ b/crates/tinyagents-orchestration/src/subagent/policy.rs @@ -1,10 +1,9 @@ //! One timeout/retry/budget policy for every subagent path. //! -//! [`SubAgentPolicy`] and [`SubAgentBudget`] are defined beside the graph -//! `SubAgentNode` (`tinyagents-graph`, which this crate depends on, so the -//! types cannot live here without a cycle) and re-exported so a host names -//! them from either place. [`SubagentDriver`](super::SubagentDriver) and -//! [`SubAgentTool`](super::SubAgentTool) apply them like the graph node does: +//! [`SubAgentPolicy`], [`SubAgentBudget`], and [`SubAgentOutput`] live here so +//! orchestration owns the policy used by the async driver. The graph crate +//! temporarily re-exports these types while its deprecated subagent node is +//! removed. [`SubagentDriver`](super::SubagentDriver) applies the policy: //! //! - **timeout** cancels the child and ends the run `Incomplete(Timeout)`; //! a timed-out attempt is never retried (it may have run for the whole @@ -16,7 +15,200 @@ //! run); token caps are checked against the reported usage afterwards; //! `max_cost` is carried but not enforced (see [`SubAgentBudget`]). -pub use tinyagents_graph::{SubAgentBudget, SubAgentPolicy}; +use std::time::Duration; + +use serde_json::Value; +use tinyagents_harness::context::RunConfig; +use tinyagents_harness::middleware::BudgetLimits; +use tinyagents_harness::retry::RetryPolicy; +use tinyinference_llm::usage::UsageTotals; + +/// Structured output and accounting returned by an agent invocation. +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub struct SubAgentOutput { + /// The final assistant text. + pub text: String, + /// Parsed structured output, when present. + pub structured: Option<Value>, + /// Cumulative token usage across model calls. + pub usage: UsageTotals, + /// Number of dispatched model calls. + pub model_calls: usize, + /// Number of executed tool invocations. + pub tool_calls: usize, +} + +/// Optional work caps for one subagent invocation. +#[derive(Clone, Copy, Debug, Default, PartialEq)] +pub struct SubAgentBudget { + /// Maximum model calls (`None` means unbounded). + pub max_model_calls: Option<usize>, + /// Maximum tool calls (`None` means unbounded). + pub max_tool_calls: Option<usize>, + /// Maximum cumulative input tokens (`None` means unbounded). + pub max_input_tokens: Option<u64>, + /// Maximum cumulative output tokens (`None` means unbounded). + pub max_output_tokens: Option<u64>, + /// Maximum estimated cost (`None` means unbounded); carried, not enforced + /// here because enforcement requires a pricing table. + pub max_cost: Option<f64>, +} + +impl SubAgentBudget { + /// A budget that imposes no cap. + pub fn unlimited() -> Self { + Self::default() + } + + /// Caps model calls, returning `self` for chaining. + pub fn with_max_model_calls(mut self, max: usize) -> Self { + self.max_model_calls = Some(max); + self + } + + /// Caps tool calls, returning `self` for chaining. + pub fn with_max_tool_calls(mut self, max: usize) -> Self { + self.max_tool_calls = Some(max); + self + } + + /// Caps cumulative input tokens, returning `self` for chaining. + pub fn with_max_input_tokens(mut self, max: u64) -> Self { + self.max_input_tokens = Some(max); + self + } + + /// Caps cumulative output tokens, returning `self` for chaining. + pub fn with_max_output_tokens(mut self, max: u64) -> Self { + self.max_output_tokens = Some(max); + self + } + + /// Caps cumulative estimated cost, returning `self` for chaining. + pub fn with_max_cost(mut self, max: f64) -> Self { + self.max_cost = Some(max); + self + } + + /// Tightens `config`'s call caps without widening existing restrictions. + pub fn apply_call_caps(&self, config: &mut RunConfig) { + if let Some(max) = self.max_model_calls { + config.max_model_calls = Some(config.max_model_calls.map_or(max, |set| set.min(max))); + } + if let Some(max) = self.max_tool_calls { + config.max_tool_calls = Some(config.max_tool_calls.map_or(max, |set| set.min(max))); + } + } + + /// Converts token and cost limits to harness middleware limits. + pub fn to_budget_limits(&self) -> BudgetLimits { + BudgetLimits { + max_input_tokens: self.max_input_tokens, + max_output_tokens: self.max_output_tokens, + max_cost: self.max_cost, + ..BudgetLimits::default() + } + } + + /// Returns an error when the measured output exceeds a configured cap. + pub fn check( + &self, + output: &SubAgentOutput, + agent: &str, + ) -> Result<(), tinyagents_harness::error::TinyAgentsError> { + if let Some(max) = self.max_model_calls + && output.model_calls > max + { + return Err(tinyagents_harness::error::TinyAgentsError::LimitExceeded( + format!( + "sub-agent `{agent}` exceeded model-call budget: {} > {max}", + output.model_calls + ), + )); + } + if let Some(max) = self.max_tool_calls + && output.tool_calls > max + { + return Err(tinyagents_harness::error::TinyAgentsError::LimitExceeded( + format!( + "sub-agent `{agent}` exceeded tool-call budget: {} > {max}", + output.tool_calls + ), + )); + } + if let Some(max) = self.max_input_tokens + && output.usage.usage.input_tokens > max + { + return Err(tinyagents_harness::error::TinyAgentsError::LimitExceeded( + format!( + "sub-agent `{agent}` exceeded input-token budget: {} > {max}", + output.usage.usage.input_tokens + ), + )); + } + if let Some(max) = self.max_output_tokens + && output.usage.usage.output_tokens > max + { + return Err(tinyagents_harness::error::TinyAgentsError::LimitExceeded( + format!( + "sub-agent `{agent}` exceeded output-token budget: {} > {max}", + output.usage.usage.output_tokens + ), + )); + } + Ok(()) + } +} + +/// Timeout, retry, and budget policy for an async subagent invocation. +#[derive(Clone, Debug)] +pub struct SubAgentPolicy { + /// Optional wall-clock timeout for each attempt. + pub timeout: Option<Duration>, + /// Retry and backoff policy. + pub retry: RetryPolicy, + /// Optional work caps. + pub budget: SubAgentBudget, + /// Whether a failed attempt may be retried after tools ran. + pub retry_after_tool_calls: bool, +} + +impl Default for SubAgentPolicy { + fn default() -> Self { + Self { + timeout: None, + retry: RetryPolicy::default().with_max_attempts(1), + budget: SubAgentBudget::unlimited(), + retry_after_tool_calls: false, + } + } +} + +impl SubAgentPolicy { + /// Sets the per-attempt timeout. + pub fn with_timeout(mut self, timeout: Duration) -> Self { + self.timeout = Some(timeout); + self + } + + /// Sets the retry policy. + pub fn with_retry(mut self, retry: RetryPolicy) -> Self { + self.retry = retry; + self + } + + /// Sets the work budget. + pub fn with_budget(mut self, budget: SubAgentBudget) -> Self { + self.budget = budget; + self + } + + /// Allows retries after tools have run, for idempotent tool sets. + pub fn with_retry_after_tool_calls(mut self, allow: bool) -> Self { + self.retry_after_tool_calls = allow; + self + } +} use tinyagents_harness::CancellationToken; use tinyagents_harness::context::RunContext; @@ -143,7 +335,7 @@ pub(crate) async fn apply_outcome_policies( if !matches!(outcome.status, SubagentOutcomeKind::Completed) { return (outcome, omitted_chars); } - let measured = tinyagents_graph::SubAgentOutput { + let measured = SubAgentOutput { usage: outcome.usage, ..Default::default() }; diff --git a/crates/tinyagents-orchestration/src/subagent/policy_tests.rs b/crates/tinyagents-orchestration/src/subagent/policy_tests.rs index 509374df9..b7d137171 100644 --- a/crates/tinyagents-orchestration/src/subagent/policy_tests.rs +++ b/crates/tinyagents-orchestration/src/subagent/policy_tests.rs @@ -1,9 +1,8 @@ use std::sync::Arc; use std::time::Duration; -use tinyagents_harness::retry::RetryPolicy; - use super::*; +use tinyagents_harness::retry::RetryPolicy; fn transient() -> TinyAgentsError { TinyAgentsError::Tool("flaky".into()) @@ -72,8 +71,7 @@ fn does_not_retry_after_tools_ran_unless_the_policy_allows_it() { #[test] fn the_orchestration_path_names_the_same_policy_type_as_the_graph() { - let policy: tinyagents_graph::SubAgentPolicy = - SubAgentPolicy::default().with_timeout(Duration::from_secs(1)); + let policy: SubAgentPolicy = SubAgentPolicy::default().with_timeout(Duration::from_secs(1)); assert_eq!(policy.timeout, Some(Duration::from_secs(1))); } diff --git a/crates/tinyagents-orchestration/src/subagent/role.rs b/crates/tinyagents-orchestration/src/subagent/role.rs index c72514a34..691f82318 100644 --- a/crates/tinyagents-orchestration/src/subagent/role.rs +++ b/crates/tinyagents-orchestration/src/subagent/role.rs @@ -3,88 +3,13 @@ //! //! A [`SubagentRole::Leaf`] child never receives the delegation surface this //! crate registers (the job and message tools, plus whichever -//! [`SubAgentTool`](super::SubAgentTool) names the host lists), and +//! spawn tools the host lists), and //! [`restrict_tools`] guarantees a child's declarations are a subset of the //! set it inherits. [`subagent_framing`] is an optional, neutral system //! preamble a host may prepend; nothing applies it by default. use tinyagents_harness::tool::ToolSnapshot; -/// Fleet-control tools whose visibility shapes orchestration prompts and -/// decisions. The names are stable tool identifiers supplied by the host. -pub const FLEET_TOOLS: &[&str] = &[ - "steer_subagent", - "wait_subagent", - "wait", - "wait_loop", - "close_subagent", - "continue_subagent", - "list_subagents", -]; - -/// The fleet-control subset visible to one parent agent. -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct FleetToolSet { - available: Vec<&'static str>, -} - -impl FleetToolSet { - /// Treats an unavailable parent definition as having the full fleet - /// vocabulary, allowing host prompts to retain their prior fallback. - pub fn all() -> Self { - Self { - available: FLEET_TOOLS.to_vec(), - } - } - - /// Resolves fleet visibility from a definition's tool scope and denylist. - /// A trailing `*` denies every fleet name with that prefix. - pub fn from_scope(scope: &tinyagents_definition::ToolScope, disallowed: &[String]) -> Self { - let denied = |name: &str| { - disallowed - .iter() - .any(|entry| match entry.strip_suffix('*') { - Some(prefix) => name.starts_with(prefix), - None => entry == name, - }) - }; - let available = FLEET_TOOLS - .iter() - .copied() - .filter(|name| match scope { - tinyagents_definition::ToolScope::Wildcard => true, - tinyagents_definition::ToolScope::Named(names) => { - names.iter().any(|candidate| candidate == name) - } - }) - .filter(|name| !denied(name)) - .collect(); - Self { available } - } - - /// Resolves fleet visibility from the effective, already-filtered tool - /// names for the current turn. - pub fn from_visible_tool_names(names: &std::collections::HashSet<String>) -> Self { - Self { - available: FLEET_TOOLS - .iter() - .copied() - .filter(|name| names.contains(*name)) - .collect(), - } - } - - /// Whether the parent can call `tool`. - pub fn has(&self, tool: &str) -> bool { - self.available.contains(&tool) - } - - /// Whether the parent can block on or poll a worker. - pub fn can_wait(&self) -> bool { - self.has("wait_subagent") - } -} - /// Name of the job query tool `register_subagent_job_tools` from the /// feature-gated `tinyagents-harness-tools` crate registers. pub const SUBAGENT_JOBS_TOOL: &str = "subagent_jobs"; @@ -126,7 +51,7 @@ impl From<tinyagents_definition::AgentTier> for SubagentRole { } /// Whether `name` is part of the delegation surface: this crate's fixed tools -/// or one of the host's [`SubAgentTool`](super::SubAgentTool) names. +/// or one of the host's registered spawn-tool names. pub fn is_delegation_tool(name: &str, host_delegation_tools: &[String]) -> bool { name == SUBAGENT_JOBS_TOOL || name == SUBAGENT_MESSAGE_TOOL diff --git a/crates/tinyagents-orchestration/src/subagent/role_tests.rs b/crates/tinyagents-orchestration/src/subagent/role_tests.rs index 49de3c6f8..c959cbcc7 100644 --- a/crates/tinyagents-orchestration/src/subagent/role_tests.rs +++ b/crates/tinyagents-orchestration/src/subagent/role_tests.rs @@ -129,28 +129,3 @@ fn authored_tiers_map_to_runtime_delegation_roles() { ); assert_eq!(SubagentRole::from(AgentTier::Worker), SubagentRole::Leaf); } - -#[test] -fn fleet_tools_follow_named_scope_and_deny_patterns() { - let fleet = FleetToolSet::from_scope( - &tinyagents_definition::ToolScope::Named(vec![ - "wait_subagent".into(), - "list_subagents".into(), - "read".into(), - ]), - &["list_*".into()], - ); - assert!(fleet.has("wait_subagent")); - assert!(!fleet.has("list_subagents")); - assert!(!fleet.has("read")); - assert!(fleet.can_wait()); -} - -#[test] -fn wildcard_fleet_visibility_can_use_effective_visible_names() { - let names = ["wait_subagent".to_owned()].into_iter().collect(); - let fleet = FleetToolSet::from_visible_tool_names(&names); - assert!(fleet.has("wait_subagent")); - assert!(!fleet.has("close_subagent")); - assert!(!fleet.has("continue_subagent")); -} diff --git a/crates/tinyagents-orchestration/src/subagent/sessions.rs b/crates/tinyagents-orchestration/src/subagent/sessions.rs index c68992a36..0d65788f5 100644 --- a/crates/tinyagents-orchestration/src/subagent/sessions.rs +++ b/crates/tinyagents-orchestration/src/subagent/sessions.rs @@ -3,17 +3,20 @@ use std::sync::Arc; use async_trait::async_trait; -use serde::{Deserialize, Serialize}; +use serde::{Deserialize, Deserializer, Serialize}; +use serde_json::Value; use tinyagents_storage_kit::{ Blocking, CollectionSpec, DocumentStore, DocumentStoreExt, ErrorKind, Precondition, Query, }; -use tinyinference_llm::message::Message; +use tinyinference_llm::message::{AssistantMessage, ContentBlock, Message}; +use tinyinference_llm::tool::ToolCall; use tokio::sync::OnceCell; use super::{SubagentError, SubagentTaskKey}; const COLLECTION: &str = "subagent_sessions"; const CAS_ATTEMPTS: usize = 32; +type SessionCloseCallback = dyn Fn(&str, &str, &SubagentTaskKey) + Send + Sync; fn session_id_for_selector(selector: &SubagentSessionSelector) -> String { let encoded = @@ -38,7 +41,7 @@ pub enum DurableSubagentStatus { impl DurableSubagentStatus { /// Whether a session may be selected for reuse. pub fn reusable(self) -> bool { - matches!(self, Self::Running | Self::Idle) + matches!(self, Self::Running | Self::Idle | Self::AwaitingUser) } } @@ -78,7 +81,11 @@ pub struct DurableSubagentSession { pub current_task_id: Option<String>, pub status: DurableSubagentStatus, pub reusable: bool, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde( + default, + skip_serializing_if = "Option::is_none", + deserialize_with = "deserialize_session_history" + )] pub latest_history: Option<Vec<Message>>, pub latest_error: Option<String>, pub created_at: String, @@ -86,6 +93,75 @@ pub struct DurableSubagentSession { pub last_used_at: String, } +#[derive(Deserialize)] +struct LegacyTranscriptMessage { + role: String, + #[serde(default)] + content: String, + #[serde(default)] + tool_calls: Vec<LegacyTranscriptToolCall>, + #[serde(default)] + tool_call_id: Option<String>, +} + +#[derive(Deserialize)] +struct LegacyTranscriptToolCall { + id: String, + name: String, + #[serde(default)] + arguments: String, +} + +fn deserialize_session_history<'de, D>(deserializer: D) -> Result<Option<Vec<Message>>, D::Error> +where + D: Deserializer<'de>, +{ + let Some(rows) = Option::<Vec<Value>>::deserialize(deserializer)? else { + return Ok(None); + }; + rows.into_iter() + .map(|row| { + if let Ok(message) = serde_json::from_value::<Message>(row.clone()) { + return Ok(message); + } + let legacy: LegacyTranscriptMessage = + serde_json::from_value(row).map_err(serde::de::Error::custom)?; + let message = match legacy.role.as_str() { + "system" => Message::system(legacy.content), + "user" => Message::user(legacy.content), + "tool" => Message::tool( + legacy.tool_call_id.ok_or_else(|| { + serde::de::Error::custom("legacy tool row has no call id") + })?, + legacy.content, + ), + "assistant" => Message::Assistant(AssistantMessage { + id: None, + content: vec![ContentBlock::Text(legacy.content)], + tool_calls: legacy + .tool_calls + .into_iter() + .map(|call| match serde_json::from_str(&call.arguments) { + Ok(arguments) => ToolCall::new(call.id, call.name, arguments), + Err(error) => ToolCall::invalid( + call.id, + call.name, + call.arguments, + error.to_string(), + ), + }) + .collect(), + usage: None, + origin: None, + }), + _ => return Err(serde::de::Error::custom("unknown legacy transcript role")), + }; + Ok(message) + }) + .collect::<Result<Vec<_>, D::Error>>() + .map(Some) +} + impl DurableSubagentSession { /// Whether this session is reusable for `selector`. pub fn matches_selector(&self, selector: &SubagentSessionSelector) -> bool { @@ -201,6 +277,21 @@ pub trait SubagentSessionStore: Send + Sync { ) -> Result<(DurableSubagentSession, ReuseDecision), SubagentError>; /// Closes a session so it cannot be reused. async fn close(&self, id: &str) -> Result<bool, SubagentError>; + /// Closes only the current task lifecycle. A stale owner cannot close a + /// replacement that reused the same durable session id. + async fn close_if_current( + &self, + id: &str, + task_id: &str, + lifecycle_key: &SubagentTaskKey, + parent_session: &str, + parent_thread_id: Option<&str>, + ) -> Result<bool, SubagentError> { + let _ = (id, task_id, lifecycle_key, parent_session, parent_thread_id); + Err(SubagentError::Persistence( + "session store does not support atomic lifecycle close".into(), + )) + } /// Records the latest result and reusable state after a completed turn. async fn mark_finished( &self, @@ -209,6 +300,33 @@ pub trait SubagentSessionStore: Send + Sync { status: DurableSubagentStatus, history: Vec<Message>, ) -> Result<bool, SubagentError>; + /// Owner- and lifecycle-checked completion for an owner-bound session. + // Each argument is part of the atomic ownership predicate or the new + // session result; combining them would hide which values are checked. + #[allow(clippy::too_many_arguments)] + async fn mark_finished_if_current( + &self, + id: &str, + task_id: &str, + lifecycle_key: &SubagentTaskKey, + parent_session: &str, + parent_thread_id: Option<&str>, + status: DurableSubagentStatus, + history: Vec<Message>, + ) -> Result<bool, SubagentError> { + let _ = ( + id, + task_id, + lifecycle_key, + parent_session, + parent_thread_id, + status, + history, + ); + Err(SubagentError::Persistence( + "session store does not support atomic owner-checked completion".into(), + )) + } /// Records a failed turn while preserving the durable session for listing. async fn mark_failed( &self, @@ -216,6 +334,234 @@ pub trait SubagentSessionStore: Send + Sync { task_id: &str, error: String, ) -> Result<bool, SubagentError>; + /// Owner- and lifecycle-checked failure for an owner-bound session. + async fn mark_failed_if_current( + &self, + id: &str, + task_id: &str, + lifecycle_key: &SubagentTaskKey, + parent_session: &str, + parent_thread_id: Option<&str>, + error: String, + ) -> Result<bool, SubagentError> { + let _ = ( + id, + task_id, + lifecycle_key, + parent_session, + parent_thread_id, + error, + ); + Err(SubagentError::Persistence( + "session store does not support atomic owner-checked failure".into(), + )) + } +} + +/// Owner-bound view of a session store for one exact parent session/thread. +/// Missing thread identity is exact (threadless), never a wildcard. +#[derive(Clone)] +pub struct ScopedSubagentSessionStore<S> { + inner: Arc<S>, + parent_session: String, + parent_thread_id: Option<String>, + close_callback: Option<Arc<SessionCloseCallback>>, +} + +impl<S> ScopedSubagentSessionStore<S> { + /// Binds the underlying store to one parent owner. + pub fn new( + inner: Arc<S>, + parent_session: impl Into<String>, + parent_thread_id: Option<String>, + ) -> Self { + Self { + inner, + parent_session: parent_session.into(), + parent_thread_id, + close_callback: None, + } + } + + /// Binds a host cancellation callback to successful lifecycle closes. + /// The callback receives the closed session id, task id and lifecycle key + /// only after the store's conditional close wins. + pub fn with_close_callback( + inner: Arc<S>, + parent_session: impl Into<String>, + parent_thread_id: Option<String>, + callback: impl Fn(&str, &str, &SubagentTaskKey) + Send + Sync + 'static, + ) -> Self { + Self { + inner, + parent_session: parent_session.into(), + parent_thread_id, + close_callback: Some(Arc::new(callback)), + } + } + + fn owns(&self, session: &DurableSubagentSession) -> bool { + session.parent_session == self.parent_session + && session.parent_thread_id == self.parent_thread_id + } + + fn owns_selector(&self, selector: &SubagentSessionSelector) -> bool { + selector.parent_session == self.parent_session + && selector.parent_thread_id == self.parent_thread_id + } + + fn owner_error() -> SubagentError { + SubagentError::Persistence("subagent session is outside the bound owner scope".into()) + } +} + +#[async_trait] +impl<S: SubagentSessionStore + 'static> SubagentSessionStore for ScopedSubagentSessionStore<S> { + async fn list_for_parent( + &self, + parent_session: &str, + parent_thread_id: Option<&str>, + ) -> Result<Vec<DurableSubagentSession>, SubagentError> { + if parent_session != self.parent_session + || parent_thread_id != self.parent_thread_id.as_deref() + { + return Err(Self::owner_error()); + } + Ok(self + .inner + .list_for_parent(&self.parent_session, self.parent_thread_id.as_deref()) + .await? + .into_iter() + .filter(|session| self.owns(session)) + .collect()) + } + + async fn get(&self, id: &str) -> Result<Option<DurableSubagentSession>, SubagentError> { + Ok(self + .inner + .get(id) + .await? + .filter(|session| self.owns(session))) + } + + async fn upsert_running( + &self, + upsert: SubagentSessionUpsert, + ) -> Result<(DurableSubagentSession, ReuseDecision), SubagentError> { + if !self.owns_selector(&upsert.selector) { + return Err(Self::owner_error()); + } + let (session, decision) = self.inner.upsert_running(upsert).await?; + if !self.owns(&session) { + return Err(Self::owner_error()); + } + Ok((session, decision)) + } + + async fn close(&self, id: &str) -> Result<bool, SubagentError> { + let Some(session) = self.get(id).await? else { + return Ok(false); + }; + let Some(task_id) = session.current_task_id.as_deref() else { + return Ok(false); + }; + let Some(lifecycle_key) = session.lifecycle_key.as_ref() else { + return Ok(false); + }; + let closed = self + .inner + .close_if_current( + id, + task_id, + lifecycle_key, + &self.parent_session, + self.parent_thread_id.as_deref(), + ) + .await?; + if closed && let Some(callback) = &self.close_callback { + callback(id, task_id, lifecycle_key); + } + Ok(closed) + } + + async fn close_if_current( + &self, + id: &str, + task_id: &str, + lifecycle_key: &SubagentTaskKey, + parent_session: &str, + parent_thread_id: Option<&str>, + ) -> Result<bool, SubagentError> { + if parent_session != self.parent_session + || parent_thread_id != self.parent_thread_id.as_deref() + { + return Ok(false); + } + let closed = self + .inner + .close_if_current( + id, + task_id, + lifecycle_key, + &self.parent_session, + self.parent_thread_id.as_deref(), + ) + .await?; + if closed && let Some(callback) = &self.close_callback { + callback(id, task_id, lifecycle_key); + } + Ok(closed) + } + + async fn mark_finished( + &self, + id: &str, + task_id: &str, + status: DurableSubagentStatus, + history: Vec<Message>, + ) -> Result<bool, SubagentError> { + let Some(session) = self.get(id).await? else { + return Ok(false); + }; + let Some(lifecycle_key) = session.lifecycle_key.as_ref() else { + return Ok(false); + }; + self.inner + .mark_finished_if_current( + id, + task_id, + lifecycle_key, + &self.parent_session, + self.parent_thread_id.as_deref(), + status, + history, + ) + .await + } + + async fn mark_failed( + &self, + id: &str, + task_id: &str, + error: String, + ) -> Result<bool, SubagentError> { + let Some(session) = self.get(id).await? else { + return Ok(false); + }; + let Some(lifecycle_key) = session.lifecycle_key.as_ref() else { + return Ok(false); + }; + self.inner + .mark_failed_if_current( + id, + task_id, + lifecycle_key, + &self.parent_session, + self.parent_thread_id.as_deref(), + error, + ) + .await + } } /// A `SubagentSessionStore` backed by the host's `tinystoragedrivers` document @@ -385,6 +731,12 @@ impl SubagentSessionStore for DriverSubagentSessionStore { } else { find_reusable_session(all, &upsert.selector) }; + let selected_owner = reuse.as_ref().map(|session| { + ( + session.current_task_id.clone(), + session.lifecycle_key.clone(), + ) + }); let decision = reuse_decision(reuse.as_ref(), upsert.force_fresh); let now = chrono::Utc::now().to_rfc3339(); let mut session = reuse.clone().unwrap_or_else(|| DurableSubagentSession { @@ -427,7 +779,14 @@ impl SubagentSessionStore for DriverSubagentSessionStore { if reuse.is_some() { let update = this .mutate(&session_id, |current| { - if !current.status.reusable() || !current.reusable { + if !current.status.reusable() + || !current.reusable + || selected_owner.as_ref() + != Some(&( + current.current_task_id.clone(), + current.lifecycle_key.clone(), + )) + { return false; } current.lifecycle_key = session.lifecycle_key.clone(); @@ -514,6 +873,70 @@ impl SubagentSessionStore for DriverSubagentSessionStore { .await } + async fn mark_finished_if_current( + &self, + id: &str, + task_id: &str, + lifecycle_key: &SubagentTaskKey, + parent_session: &str, + parent_thread_id: Option<&str>, + status: DurableSubagentStatus, + history: Vec<Message>, + ) -> Result<bool, SubagentError> { + let task_id = task_id.to_owned(); + let lifecycle_key = lifecycle_key.clone(); + let parent_session = parent_session.to_owned(); + let parent_thread_id = parent_thread_id.map(str::to_owned); + self.mutate(id, move |session| { + if session.status == DurableSubagentStatus::Closed + || session.current_task_id.as_deref() != Some(task_id.as_str()) + || session.lifecycle_key.as_ref() != Some(&lifecycle_key) + || session.parent_session != parent_session + || session.parent_thread_id != parent_thread_id + { + return false; + } + session.status = status; + session.reusable = status.reusable(); + session.latest_history = Some(history.clone()); + session.latest_error = None; + session.updated_at = chrono::Utc::now().to_rfc3339(); + session.last_used_at = session.updated_at.clone(); + true + }) + .await + } + + async fn close_if_current( + &self, + id: &str, + task_id: &str, + lifecycle_key: &SubagentTaskKey, + parent_session: &str, + parent_thread_id: Option<&str>, + ) -> Result<bool, SubagentError> { + let id = id.to_owned(); + let task_id = task_id.to_owned(); + let lifecycle_key = lifecycle_key.clone(); + let parent_session = parent_session.to_owned(); + let parent_thread_id = parent_thread_id.map(str::to_owned); + self.mutate(&id, move |session| { + if session.current_task_id.as_deref() != Some(task_id.as_str()) + || session.lifecycle_key.as_ref() != Some(&lifecycle_key) + || session.parent_session != parent_session + || session.parent_thread_id != parent_thread_id + { + return false; + } + session.status = DurableSubagentStatus::Closed; + session.reusable = false; + session.updated_at = chrono::Utc::now().to_rfc3339(); + session.last_used_at = session.updated_at.clone(); + true + }) + .await + } + async fn mark_finished( &self, id: &str, @@ -559,6 +982,38 @@ impl SubagentSessionStore for DriverSubagentSessionStore { }) .await } + + async fn mark_failed_if_current( + &self, + id: &str, + task_id: &str, + lifecycle_key: &SubagentTaskKey, + parent_session: &str, + parent_thread_id: Option<&str>, + error: String, + ) -> Result<bool, SubagentError> { + let task_id = task_id.to_owned(); + let lifecycle_key = lifecycle_key.clone(); + let parent_session = parent_session.to_owned(); + let parent_thread_id = parent_thread_id.map(str::to_owned); + self.mutate(id, move |session| { + if session.status == DurableSubagentStatus::Closed + || session.current_task_id.as_deref() != Some(task_id.as_str()) + || session.lifecycle_key.as_ref() != Some(&lifecycle_key) + || session.parent_session != parent_session + || session.parent_thread_id != parent_thread_id + { + return false; + } + session.status = DurableSubagentStatus::Failed; + session.reusable = false; + session.latest_error = Some(error.clone()); + session.updated_at = chrono::Utc::now().to_rfc3339(); + session.last_used_at = session.updated_at.clone(); + true + }) + .await + } } /// Selects the most recently used reusable session matching `selector`. @@ -634,3 +1089,7 @@ pub fn task_title_from_prompt(prompt: &str) -> String { title } } + +#[cfg(test)] +#[path = "sessions_tests.rs"] +mod tests; diff --git a/crates/tinyagents-orchestration/src/subagent/sessions_tests.rs b/crates/tinyagents-orchestration/src/subagent/sessions_tests.rs new file mode 100644 index 000000000..53a369eaf --- /dev/null +++ b/crates/tinyagents-orchestration/src/subagent/sessions_tests.rs @@ -0,0 +1,629 @@ +use super::*; +use serde_json::Value; +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; +use tinyagents_storage_kit::{ + Capabilities, CollectionSpec, DocumentStore, Filter, Page, Precondition, Query, Scope, Sort, + StorageBackend, Version, Versioned, +}; +use tinystoragedrivers_testkit::MemoryBackend; + +fn key(task_id: &str) -> SubagentTaskKey { + SubagentTaskKey { + root_run_id: "root".into(), + parent_run_id: "parent".into(), + thread_id: Some("thread".into()), + task_id: task_id.into(), + } +} + +fn upsert(task_id: &str, lifecycle_key: SubagentTaskKey) -> SubagentSessionUpsert { + SubagentSessionUpsert { + selector: SubagentSessionSelector { + parent_session: "parent-session".into(), + parent_thread_id: Some("thread".into()), + agent_id: "worker".into(), + model: None, + sandbox_mode: "none".into(), + action_root: None, + task_key: "research".into(), + }, + lifecycle_key, + display_name: None, + task_title: "research".into(), + worker_thread_id: None, + task_id: task_id.into(), + force_fresh: false, + } +} + +#[tokio::test] +async fn stale_session_owner_cannot_close_a_reused_lifecycle() { + let backend = MemoryBackend::new(); + let docs = Arc::clone( + backend + .for_scope(&Scope::new("local").unwrap()) + .unwrap() + .documents(), + ); + let store = DriverSubagentSessionStore::new(docs).unwrap(); + let first_key = key("task-a"); + let (first, _) = store + .upsert_running(upsert("task-a", first_key.clone())) + .await + .unwrap(); + let second_key = key("task-b"); + let (second, _) = store + .upsert_running(upsert("task-b", second_key.clone())) + .await + .unwrap(); + assert_eq!(first.subagent_session_id, second.subagent_session_id); + + assert!( + !store + .close_if_current( + &second.subagent_session_id, + "task-a", + &first_key, + "parent-session", + Some("thread"), + ) + .await + .unwrap() + ); + assert_eq!( + store + .get(&second.subagent_session_id) + .await + .unwrap() + .unwrap() + .status, + DurableSubagentStatus::Running + ); + assert!( + store + .close_if_current( + &second.subagent_session_id, + "task-b", + &second_key, + "parent-session", + Some("thread"), + ) + .await + .unwrap() + ); +} + +#[tokio::test] +async fn scoped_session_store_rejects_other_parent_identity() { + let backend = MemoryBackend::new(); + let docs = Arc::clone( + backend + .for_scope(&Scope::new("local").unwrap()) + .unwrap() + .documents(), + ); + let inner = Arc::new(DriverSubagentSessionStore::new(docs).unwrap()); + let scoped = + ScopedSubagentSessionStore::new(inner, "parent-session", Some("thread".to_owned())); + assert!( + scoped + .list_for_parent("other-session", Some("thread")) + .await + .is_err() + ); + assert!( + scoped + .list_for_parent("parent-session", None) + .await + .is_err() + ); + assert!(scoped.get("not-owned").await.unwrap().is_none()); + let mut foreign = upsert("task", key("task")); + foreign.selector.parent_session = "other-session".into(); + assert!(scoped.upsert_running(foreign).await.is_err()); +} + +#[tokio::test] +async fn scoped_session_close_notifies_only_after_current_owner_cas_wins() { + let backend = MemoryBackend::new(); + let docs = Arc::clone( + backend + .for_scope(&Scope::new("local").unwrap()) + .unwrap() + .documents(), + ); + let driver = Arc::new(DriverSubagentSessionStore::new(docs).unwrap()); + let lifecycle_key = key("close-me"); + let (session, _) = driver + .upsert_running(upsert("close-me", lifecycle_key.clone())) + .await + .unwrap(); + let callbacks = Arc::new(std::sync::Mutex::new(Vec::new())); + let callback_log = Arc::clone(&callbacks); + let scoped = ScopedSubagentSessionStore::with_close_callback( + Arc::clone(&driver), + "parent-session", + Some("thread".into()), + move |session_id, task_id, key| { + callback_log.lock().unwrap().push(( + session_id.to_owned(), + task_id.to_owned(), + key.clone(), + )); + }, + ); + + assert!( + !scoped + .close_if_current( + &session.subagent_session_id, + "foreign-task", + &lifecycle_key, + "parent-session", + Some("thread"), + ) + .await + .unwrap() + ); + assert!(callbacks.lock().unwrap().is_empty()); + assert!( + !scoped + .close_if_current( + &session.subagent_session_id, + "close-me", + &lifecycle_key, + "other-parent", + Some("thread"), + ) + .await + .unwrap() + ); + assert!(callbacks.lock().unwrap().is_empty()); + assert!(scoped.close(&session.subagent_session_id).await.unwrap()); + assert_eq!( + callbacks.lock().unwrap().as_slice(), + &[( + session.subagent_session_id.clone(), + "close-me".to_owned(), + lifecycle_key, + )] + ); + + let second_key = key("close-direct"); + let mut second_upsert = upsert("close-direct", second_key.clone()); + second_upsert.force_fresh = true; + let (second, _) = driver.upsert_running(second_upsert).await.unwrap(); + assert!( + scoped + .close_if_current( + &second.subagent_session_id, + "close-direct", + &second_key, + "parent-session", + Some("thread"), + ) + .await + .unwrap() + ); + assert_eq!(callbacks.lock().unwrap().len(), 2); +} + +#[test] +fn durable_session_decodes_legacy_transcript_message_history() { + let value = serde_json::json!({ + "subagentSessionId": "session", + "parentSession": "parent", + "parentThreadId": "thread", + "workerThreadId": null, + "agentId": "worker", + "displayName": null, + "model": null, + "sandboxMode": "none", + "actionRoot": null, + "taskKey": "research", + "taskTitle": "research", + "currentTaskId": "task", + "status": "awaiting_user", + "reusable": true, + "latestHistory": [ + { "role": "user", "content": "need help" }, + { + "role": "assistant", + "content": "checking", + "tool_calls": [{ "id": "call-1", "name": "lookup", "arguments": "{\"q\":\"x\"}" }] + }, + { "role": "tool", "content": "found", "tool_call_id": "call-1" } + ], + "latestError": null, + "createdAt": "now", + "updatedAt": "now", + "lastUsedAt": "now" + }); + + let session: DurableSubagentSession = serde_json::from_value(value).unwrap(); + let history = session.latest_history.unwrap(); + assert_eq!(history[0], Message::user("need help")); + assert_eq!(history[1].text(), "checking"); + assert_eq!(history[2], Message::tool("call-1", "found")); + let Message::Assistant(assistant) = &history[1] else { + panic!("legacy assistant row must remain an assistant message") + }; + assert_eq!(assistant.tool_calls.len(), 1); + assert_eq!(assistant.tool_calls[0].name, "lookup"); +} + +#[derive(Clone)] +struct PauseAfterSessionRead { + inner: Arc<DriverSubagentSessionStore>, + entered: Arc<tokio::sync::Notify>, + release: Arc<tokio::sync::Notify>, + pause_once: Arc<std::sync::atomic::AtomicBool>, +} + +#[async_trait::async_trait] +impl SubagentSessionStore for PauseAfterSessionRead { + async fn list_for_parent( + &self, + parent_session: &str, + parent_thread_id: Option<&str>, + ) -> Result<Vec<DurableSubagentSession>, SubagentError> { + self.inner + .list_for_parent(parent_session, parent_thread_id) + .await + } + + async fn get(&self, id: &str) -> Result<Option<DurableSubagentSession>, SubagentError> { + let session = self.inner.get(id).await?; + if self + .pause_once + .swap(false, std::sync::atomic::Ordering::SeqCst) + { + self.entered.notify_one(); + self.release.notified().await; + } + Ok(session) + } + + async fn upsert_running( + &self, + upsert: SubagentSessionUpsert, + ) -> Result<(DurableSubagentSession, ReuseDecision), SubagentError> { + self.inner.upsert_running(upsert).await + } + + async fn close(&self, id: &str) -> Result<bool, SubagentError> { + self.inner.close(id).await + } + + async fn close_if_current( + &self, + id: &str, + task_id: &str, + lifecycle_key: &SubagentTaskKey, + parent_session: &str, + parent_thread_id: Option<&str>, + ) -> Result<bool, SubagentError> { + self.inner + .close_if_current(id, task_id, lifecycle_key, parent_session, parent_thread_id) + .await + } + + async fn mark_finished( + &self, + id: &str, + task_id: &str, + status: DurableSubagentStatus, + history: Vec<Message>, + ) -> Result<bool, SubagentError> { + self.inner.mark_finished(id, task_id, status, history).await + } + + async fn mark_finished_if_current( + &self, + id: &str, + task_id: &str, + lifecycle_key: &SubagentTaskKey, + parent_session: &str, + parent_thread_id: Option<&str>, + status: DurableSubagentStatus, + history: Vec<Message>, + ) -> Result<bool, SubagentError> { + self.inner + .mark_finished_if_current( + id, + task_id, + lifecycle_key, + parent_session, + parent_thread_id, + status, + history, + ) + .await + } + + async fn mark_failed( + &self, + id: &str, + task_id: &str, + error: String, + ) -> Result<bool, SubagentError> { + self.inner.mark_failed(id, task_id, error).await + } + + async fn mark_failed_if_current( + &self, + id: &str, + task_id: &str, + lifecycle_key: &SubagentTaskKey, + parent_session: &str, + parent_thread_id: Option<&str>, + error: String, + ) -> Result<bool, SubagentError> { + self.inner + .mark_failed_if_current( + id, + task_id, + lifecycle_key, + parent_session, + parent_thread_id, + error, + ) + .await + } +} + +#[tokio::test] +async fn scoped_mutations_recheck_reused_task_owner_after_a_paused_read() { + for mutation in ["close", "finish", "fail"] { + let backend = MemoryBackend::new(); + let docs = Arc::clone( + backend + .for_scope(&Scope::new("local").unwrap()) + .unwrap() + .documents(), + ); + let driver = Arc::new(DriverSubagentSessionStore::new(docs).unwrap()); + let old_key = key("same-task"); + let (old, _) = driver + .upsert_running(upsert("same-task", old_key.clone())) + .await + .unwrap(); + let barrier = PauseAfterSessionRead { + inner: Arc::clone(&driver), + entered: Arc::new(tokio::sync::Notify::new()), + release: Arc::new(tokio::sync::Notify::new()), + pause_once: Arc::new(std::sync::atomic::AtomicBool::new(true)), + }; + let scoped = Arc::new(ScopedSubagentSessionStore::new( + Arc::new(barrier.clone()), + "parent-session", + Some("thread".into()), + )); + let waiting_for_read = barrier.entered.notified(); + let pending = { + let scoped = Arc::clone(&scoped); + let id = old.subagent_session_id.clone(); + tokio::spawn(async move { + match mutation { + "close" => scoped.close(&id).await, + "finish" => { + scoped + .mark_finished( + &id, + "same-task", + DurableSubagentStatus::Idle, + vec![Message::user("stale")], + ) + .await + } + "fail" => scoped.mark_failed(&id, "same-task", "stale".into()).await, + _ => unreachable!(), + } + }) + }; + waiting_for_read.await; + + let new_key = SubagentTaskKey { + root_run_id: "replacement-root".into(), + parent_run_id: "replacement-parent".into(), + thread_id: Some("thread".into()), + task_id: "same-task".into(), + }; + let (replacement, _) = driver + .upsert_running(upsert("same-task", new_key.clone())) + .await + .unwrap(); + assert_eq!(replacement.subagent_session_id, old.subagent_session_id); + barrier.release.notify_one(); + assert!( + !pending.await.unwrap().unwrap(), + "stale {mutation} must be rejected" + ); + + assert!(match mutation { + "close" => driver + .close_if_current( + &replacement.subagent_session_id, + "same-task", + &new_key, + "parent-session", + Some("thread"), + ) + .await + .unwrap(), + "finish" => driver + .mark_finished_if_current( + &replacement.subagent_session_id, + "same-task", + &new_key, + "parent-session", + Some("thread"), + DurableSubagentStatus::Idle, + vec![Message::user("current")], + ) + .await + .unwrap(), + "fail" => driver + .mark_failed_if_current( + &replacement.subagent_session_id, + "same-task", + &new_key, + "parent-session", + Some("thread"), + "current failure".into(), + ) + .await + .unwrap(), + _ => unreachable!(), + }); + } +} + +#[derive(Debug)] +struct PauseAfterSelectorRead { + inner: Arc<dyn DocumentStore>, + pause_once: AtomicBool, + entered: tokio::sync::Notify, + release: tokio::sync::Notify, +} + +#[async_trait::async_trait] +impl DocumentStore for PauseAfterSelectorRead { + fn capabilities(&self) -> Capabilities { + self.inner.capabilities() + } + + async fn ensure_collection( + &self, + spec: &CollectionSpec, + ) -> tinyagents_storage_kit::error::Result<()> { + self.inner.ensure_collection(spec).await + } + + async fn get( + &self, + collection: &str, + id: &str, + ) -> tinyagents_storage_kit::error::Result<Option<Versioned<Value>>> { + self.inner.get(collection, id).await + } + + async fn put( + &self, + collection: &str, + id: &str, + doc: Value, + precondition: Precondition, + ) -> tinyagents_storage_kit::error::Result<Version> { + self.inner.put(collection, id, doc, precondition).await + } + + async fn delete( + &self, + collection: &str, + id: &str, + precondition: Precondition, + ) -> tinyagents_storage_kit::error::Result<bool> { + self.inner.delete(collection, id, precondition).await + } + + async fn query( + &self, + collection: &str, + query: &Query, + ) -> tinyagents_storage_kit::error::Result<Page<Versioned<Value>>> { + let page = self.inner.query(collection, query).await?; + if self.pause_once.swap(false, Ordering::SeqCst) { + self.entered.notify_one(); + self.release.notified().await; + } + Ok(page) + } + + async fn count( + &self, + collection: &str, + filter: &Filter, + ) -> tinyagents_storage_kit::error::Result<u64> { + self.inner.count(collection, filter).await + } + + async fn delete_where( + &self, + collection: &str, + filter: &Filter, + ) -> tinyagents_storage_kit::error::Result<u64> { + self.inner.delete_where(collection, filter).await + } + + async fn claim( + &self, + collection: &str, + filter: &Filter, + sort: &[Sort], + patch: &Value, + ) -> tinyagents_storage_kit::error::Result<Option<Versioned<Value>>> { + self.inner.claim(collection, filter, sort, patch).await + } + + async fn drop_collection(&self, collection: &str) -> tinyagents_storage_kit::error::Result<()> { + self.inner.drop_collection(collection).await + } +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn stale_reuse_selection_cannot_overwrite_a_newer_session_owner() { + let backend = MemoryBackend::new(); + let docs = Arc::clone( + backend + .for_scope(&Scope::new("local").unwrap()) + .unwrap() + .documents(), + ); + let gate = Arc::new(PauseAfterSelectorRead { + inner: docs.clone(), + pause_once: AtomicBool::new(false), + entered: tokio::sync::Notify::new(), + release: tokio::sync::Notify::new(), + }); + let stale_store = Arc::new(DriverSubagentSessionStore::new(gate.clone()).unwrap()); + let current_store = DriverSubagentSessionStore::new(docs).unwrap(); + let old_key = key("task-old"); + let (old, _) = stale_store + .upsert_running(upsert("task-old", old_key.clone())) + .await + .unwrap(); + + gate.pause_once.store(true, Ordering::SeqCst); + let selector_read = gate.entered.notified(); + let stale_attempt = { + let store = Arc::clone(&stale_store); + tokio::spawn(async move { store.upsert_running(upsert("task-old", old_key)).await }) + }; + selector_read.await; + + let new_key = key("task-new"); + let (current, _) = current_store + .upsert_running(upsert("task-new", new_key.clone())) + .await + .unwrap(); + assert_eq!(current.subagent_session_id, old.subagent_session_id); + gate.release.notify_one(); + + assert!(stale_attempt.await.unwrap().is_err()); + let retained = current_store + .get(¤t.subagent_session_id) + .await + .unwrap() + .unwrap(); + assert_eq!(retained.current_task_id.as_deref(), Some("task-new")); + assert_eq!(retained.lifecycle_key, Some(new_key.clone())); + + let (fresh, _) = current_store + .upsert_running(upsert("task-fresh", key("task-fresh"))) + .await + .unwrap(); + assert_eq!(fresh.subagent_session_id, current.subagent_session_id); + assert_eq!(fresh.current_task_id.as_deref(), Some("task-fresh")); +} diff --git a/crates/tinyagents-orchestration/src/subagent/step/mod.rs b/crates/tinyagents-orchestration/src/subagent/step/mod.rs deleted file mode 100644 index 2f23b26e8..000000000 --- a/crates/tinyagents-orchestration/src/subagent/step/mod.rs +++ /dev/null @@ -1,145 +0,0 @@ -//! Run one host-supplied unit of agent work through [`SubagentDriver`]. -//! -//! Team members and workflow phase agents are not driven by a harness agent -//! loop the crate controls: the host hands over an opaque async worker. This -//! module adapts such a worker to the driver's planner / executor / -//! persistence seams so those steps get the same lifecycle features as every -//! other subagent path: spawn admission ([`SpawnPolicy`](super::SpawnPolicy)), -//! timeout / retry / budget ([`SubAgentPolicy`]), the [`ResultPolicy`], the -//! [`SubagentRole`], and the typed [`SubagentOutcomeKind`]. -//! -//! [`AgentStepConfig::default`] is deliberately inert (unlimited admission, no -//! timeout, a single attempt, no result trimming), so a step run with the -//! default config behaves like calling the worker directly. - -mod seams; -mod types; - -use std::future::Future; -use std::sync::{Arc, Mutex}; - -use tinyagents_harness::CancellationToken; -use tinyagents_harness::context::{RunConfig, RunContext}; - -use self::seams::{Slot, StepExecutor, StepPersistence, StepPlanner, Work}; -pub use self::types::{ - AgentStepConfig, AgentStepError, AgentStepIdentity, AgentStepResult, StepContext, StepSuccess, - StepWorkError, -}; -use super::{ - SubagentCapabilities, SubagentDriver, SubagentError, SubagentOutcomeKind, SubagentRequest, -}; - -const LOG_PREFIX: &str = "[agent-step]"; - -/// Runs `work` as one subagent lifecycle on a [`SubagentDriver`] configured -/// from `config`. -/// -/// `work` receives a [`StepContext`] (the child cancellation token, cancelled on -/// lifecycle cancellation or policy timeout, and the role it must enforce) and may be invoked more than once when the -/// retry policy retries a [`StepWorkError::Transient`] failure. -pub async fn run_agent_step<T, F, Fut>( - config: &AgentStepConfig, - identity: AgentStepIdentity, - cancellation: CancellationToken, - work: F, -) -> Result<AgentStepResult<T>, AgentStepError> -where - T: Send + 'static, - F: Fn(StepContext) -> Fut + Send + Sync + 'static, - Fut: Future<Output = Result<StepSuccess<T>, StepWorkError>> + Send + 'static, -{ - let target = identity - .target - .clone() - .unwrap_or_else(|| "agent-step".to_owned()); - tracing::debug!( - "{LOG_PREFIX} start parent={} task={} target={target}", - identity.parent_run_id, - identity.task_id - ); - let slot = Arc::new(Mutex::new(Slot { - value: None, - error: None, - })); - let work: Work<T> = Arc::new(move |token| Box::pin(work(token))); - let driver = SubagentDriver::new(SubagentCapabilities { - planner: Some(Arc::new(StepPlanner { - config: config.clone(), - target: target.clone(), - })), - executor: Some(Arc::new(StepExecutor { - work, - slot: slot.clone(), - })), - persistence: Some(Arc::new(StepPersistence::default())), - }) - .map_err(AgentStepError::Driver)? - .with_spawn_admission(config.admission.clone()); - - let parent = RunContext::new(RunConfig::new(identity.parent_run_id.as_str()), ()); - let child = parent - .child( - RunConfig::new(format!("{}:{}", identity.parent_run_id, identity.task_id)), - (), - ) - .map_err(|e| AgentStepError::Driver(SubagentError::InvalidRequest(e.to_string())))?; - let mut request = SubagentRequest::fresh_from_parent( - &parent, - child, - identity.task_id.as_str(), - (), - identity.task_id.as_str(), - None, - ) - .map_err(AgentStepError::Driver)?; - if identity.target.is_some() { - request = request.with_target(target); - } - - let run = driver.run(request, cancellation).await; - let mut slot = slot.lock().expect("agent-step slot poisoned"); - match run { - Ok(result) => { - tracing::debug!( - "{LOG_PREFIX} done parent={} task={} status={}", - identity.parent_run_id, - identity.task_id, - status_label(&result.outcome.status) - ); - if matches!(result.outcome.status, SubagentOutcomeKind::Cancelled) - && slot.value.is_none() - { - return Err(AgentStepError::Cancelled); - } - Ok(AgentStepResult { - outcome: result.outcome, - value: slot.value.take(), - }) - } - Err(SubagentError::SpawnRejected(rejection)) => { - tracing::debug!( - "{LOG_PREFIX} rejected parent={} task={} reason={rejection}", - identity.parent_run_id, - identity.task_id - ); - Err(AgentStepError::Rejected(rejection)) - } - Err(SubagentError::Cancelled) => Err(AgentStepError::Cancelled), - Err(SubagentError::Execution(_) | SubagentError::Transient { .. }) - if slot.error.is_some() => - { - Err(AgentStepError::Worker(slot.error.take().expect("checked"))) - } - Err(error) => Err(AgentStepError::Driver(error)), - } -} - -fn status_label(status: &SubagentOutcomeKind) -> &'static str { - match status { - SubagentOutcomeKind::Completed => "completed", - SubagentOutcomeKind::AwaitingInput(_) => "awaiting_input", - SubagentOutcomeKind::Incomplete(_) => "incomplete", - SubagentOutcomeKind::Cancelled => "cancelled", - } -} diff --git a/crates/tinyagents-orchestration/src/subagent/step/mod_tests.rs b/crates/tinyagents-orchestration/src/subagent/step/mod_tests.rs deleted file mode 100644 index 893edb7da..000000000 --- a/crates/tinyagents-orchestration/src/subagent/step/mod_tests.rs +++ /dev/null @@ -1,222 +0,0 @@ -//! `run_agent_step`: opaque host work goes through the real `SubagentDriver`. - -use super::*; - -use std::sync::Arc; -use std::sync::atomic::{AtomicU32, Ordering}; -use std::time::Duration; - -use tinyagents_harness::retry::RetryPolicy; - -use crate::subagent::{IncompleteKind, SpawnPolicy, SubAgentPolicy, SubagentRole}; - -fn ident(task: &str) -> AgentStepIdentity { - AgentStepIdentity::new("parent", task) -} - -async fn ok_text(text: &'static str) -> Result<StepSuccess<u32>, StepWorkError> { - Ok(StepSuccess::new(text, 7)) -} - -#[tokio::test] -async fn default_config_is_a_passthrough() { - let result = run_agent_step( - &AgentStepConfig::default(), - ident("a"), - CancellationToken::new(), - |_| ok_text("hello"), - ) - .await - .unwrap(); - assert_eq!(result.outcome.status, SubagentOutcomeKind::Completed); - assert_eq!(result.outcome.output, "hello"); - assert_eq!(result.value, Some(7)); -} - -#[tokio::test] -async fn result_policy_cap_trims_output() { - let config = - AgentStepConfig::default().with_result_policy(ResultPolicy::new().with_max_chars(5)); - let result = run_agent_step(&config, ident("a"), CancellationToken::new(), |_| { - ok_text("0123456789abcdefghij") - }) - .await - .unwrap(); - assert_eq!(result.outcome.output, "01234", "hard cut to the cap"); -} - -#[tokio::test] -async fn spawn_policy_denial_never_runs_the_worker() { - let ran = Arc::new(AtomicU32::new(0)); - let config = AgentStepConfig::default().with_admission(SpawnAdmission::new(SpawnPolicy { - max_total_per_root: Some(1), - ..Default::default() - })); - for task in ["a", "b"] { - let ran = ran.clone(); - let res = run_agent_step(&config, ident(task), CancellationToken::new(), move |_| { - ran.fetch_add(1, Ordering::SeqCst); - ok_text("x") - }) - .await; - if task == "b" { - assert!(matches!(res, Err(AgentStepError::Rejected(_)))); - } else { - assert!(res.is_ok()); - } - } - assert_eq!(ran.load(Ordering::SeqCst), 1); -} - -#[tokio::test] -async fn allowed_targets_check_the_identity_target() { - let config = AgentStepConfig::default().with_admission(SpawnAdmission::new(SpawnPolicy { - allowed_targets: Some(vec!["good".into()]), - ..Default::default() - })); - let bad = run_agent_step( - &config, - ident("a").with_target("evil"), - CancellationToken::new(), - |_| ok_text("x"), - ) - .await; - assert!(matches!(bad, Err(AgentStepError::Rejected(_)))); - let good = run_agent_step( - &config, - ident("b").with_target("good"), - CancellationToken::new(), - |_| ok_text("x"), - ) - .await; - assert!(good.is_ok()); -} - -#[tokio::test] -async fn timeout_ends_incomplete_with_typed_kind() { - let config = AgentStepConfig::default() - .with_policy(SubAgentPolicy::default().with_timeout(Duration::from_millis(20))); - let seen_token = Arc::new(std::sync::Mutex::new(None)); - let capture = seen_token.clone(); - let result = run_agent_step(&config, ident("a"), CancellationToken::new(), move |ctx| { - *capture.lock().unwrap() = Some(ctx.cancellation.clone()); - async { - std::future::pending::<()>().await; - ok_text("never").await - } - }) - .await - .unwrap(); - assert!( - seen_token.lock().unwrap().as_ref().unwrap().is_cancelled(), - "the worker's token was cancelled" - ); - match result.outcome.status { - SubagentOutcomeKind::Incomplete(inc) => assert_eq!(inc.kind, IncompleteKind::Timeout), - other => panic!("expected timeout, got {other:?}"), - } -} - -#[tokio::test] -async fn transient_failures_retry_per_policy() { - let attempts = Arc::new(AtomicU32::new(0)); - let config = AgentStepConfig::default().with_policy( - SubAgentPolicy::default().with_retry( - RetryPolicy::default() - .with_max_attempts(3) - .with_backoff_sleep(false), - ), - ); - let seen = attempts.clone(); - let result = run_agent_step(&config, ident("a"), CancellationToken::new(), move |_| { - let n = seen.fetch_add(1, Ordering::SeqCst); - async move { - if n < 2 { - Err(StepWorkError::Transient { - error: anyhow::anyhow!("flaky"), - tools_ran: false, - }) - } else { - ok_text("recovered").await - } - } - }) - .await - .unwrap(); - assert_eq!(result.outcome.output, "recovered"); - assert_eq!(attempts.load(Ordering::SeqCst), 3); -} - -#[tokio::test] -async fn fatal_worker_error_keeps_its_original_message() { - let err = run_agent_step::<(), _, _>( - &AgentStepConfig::default(), - ident("a"), - CancellationToken::new(), - |_| async { Err(anyhow::anyhow!("boom 42").into()) }, - ) - .await - .err() - .unwrap(); - match err { - AgentStepError::Worker(e) => assert_eq!(e.to_string(), "boom 42"), - other => panic!("expected worker error, got {other}"), - } -} - -#[tokio::test] -async fn pre_cancelled_step_reports_cancelled_without_running() { - let token = CancellationToken::new(); - token.cancel(); - let ran = Arc::new(AtomicU32::new(0)); - let seen = ran.clone(); - let err = run_agent_step(&AgentStepConfig::default(), ident("a"), token, move |_| { - seen.fetch_add(1, Ordering::SeqCst); - ok_text("x") - }) - .await - .err() - .unwrap(); - assert!(matches!(err, AgentStepError::Cancelled)); - assert_eq!(ran.load(Ordering::SeqCst), 0, "the worker never ran"); -} - -#[tokio::test] -async fn worker_receives_the_configured_role() { - let config = AgentStepConfig::default().with_role(SubagentRole::Leaf); - let result = run_agent_step( - &config, - ident("a"), - CancellationToken::new(), - |ctx| async move { Ok(StepSuccess::new(format!("{:?}", ctx.role), ())) }, - ) - .await - .unwrap(); - assert_eq!(result.outcome.output, "Leaf"); -} - -#[tokio::test] -async fn call_caps_reach_the_worker_through_the_context() { - use crate::subagent::SubAgentBudget; - let config = AgentStepConfig::default().with_policy( - SubAgentPolicy::default().with_budget( - SubAgentBudget::default() - .with_max_model_calls(3) - .with_max_tool_calls(5), - ), - ); - let result = run_agent_step( - &config, - ident("a"), - CancellationToken::new(), - |ctx| async move { - Ok(StepSuccess::new( - format!("{:?}/{:?}", ctx.max_model_calls, ctx.max_tool_calls), - (), - )) - }, - ) - .await - .unwrap(); - assert_eq!(result.outcome.output, "Some(3)/Some(5)"); -} diff --git a/crates/tinyagents-orchestration/src/subagent/step/seams.rs b/crates/tinyagents-orchestration/src/subagent/step/seams.rs deleted file mode 100644 index 8d9dca6c4..000000000 --- a/crates/tinyagents-orchestration/src/subagent/step/seams.rs +++ /dev/null @@ -1,165 +0,0 @@ -//! The driver seams ([`SubagentPlanner`], [`SubagentExecutor`], -//! [`SubagentPersistence`]) an agent step plugs an opaque worker into. - -use std::collections::HashMap; -use std::future::Future; -use std::pin::Pin; -use std::sync::{Arc, Mutex}; - -use async_trait::async_trait; -use tinyagents_harness::context::{RunConfig, RunContext}; -use tinyagents_harness::tool::ToolSnapshot; -use tinyinference_llm::message::Message; - -use super::types::{AgentStepConfig, StepContext, StepSuccess, StepWorkError}; -use crate::subagent::{ - PersistedSubagentPause, PreparedSubagent, SubagentError, SubagentExecution, SubagentExecutor, - SubagentIncomplete, SubagentOutcome, SubagentPausePersistenceDisposition, SubagentPersistence, - SubagentPlanner, SubagentRequest, SubagentResume, SubagentTaskKey, - SubagentTerminalPersistenceDisposition, -}; - -pub(super) type Work<T> = Arc< - dyn Fn( - StepContext, - ) -> Pin<Box<dyn Future<Output = Result<StepSuccess<T>, StepWorkError>> + Send>> - + Send - + Sync, ->; - -pub(super) struct Slot<T> { - pub(super) value: Option<T>, - pub(super) error: Option<anyhow::Error>, -} - -pub(super) struct StepPlanner { - pub(super) config: AgentStepConfig, - pub(super) target: String, -} - -#[async_trait] -impl SubagentPlanner<(), ()> for StepPlanner { - async fn prepare( - &self, - request: SubagentRequest<(), ()>, - ) -> Result<PreparedSubagent<()>, SubagentError> { - let parts = request.into_parts(); - let task_id = parts.task_key.task_id.clone(); - let factory_task = task_id.clone(); - Ok(PreparedSubagent::new( - task_id, - self.target.clone(), - vec![Message::user(parts.input)], - ToolSnapshot::new(vec![]).map_err(|e| SubagentError::Planning(e.to_string()))?, - parts.run_context, - ) - .with_role(self.config.role) - .with_policy(self.config.policy.clone()) - .with_result_policy(self.config.result_policy.clone()) - .with_retry_context(Arc::new(move |attempt| { - Ok(RunContext::new( - RunConfig::new(format!("{factory_task}-attempt-{attempt}")), - (), - )) - }))) - } -} - -pub(super) struct StepExecutor<T> { - pub(super) work: Work<T>, - pub(super) slot: Arc<Mutex<Slot<T>>>, -} - -#[async_trait] -impl<T: Send + 'static> SubagentExecutor<()> for StepExecutor<T> { - async fn execute( - &self, - execution: SubagentExecution<()>, - ) -> Result<SubagentOutcome, SubagentError> { - let task_id = execution.prepared.task_id.clone(); - if execution.cancellation.is_cancelled() { - return Err(SubagentError::Cancelled); - } - let result = (self.work)(StepContext { - cancellation: execution.cancellation.clone(), - role: execution.prepared.role, - max_model_calls: execution.prepared.run_context.config.max_model_calls, - max_tool_calls: execution.prepared.run_context.config.max_tool_calls, - }) - .await; - let mut slot = self.slot.lock().expect("agent-step slot poisoned"); - match result { - Ok(success) => { - slot.value = Some(success.value); - slot.error = None; - let mut outcome = match success.incomplete_reason { - Some(reason) => { - SubagentOutcome::incomplete(task_id, SubagentIncomplete::new(reason)) - } - None => SubagentOutcome::completed(task_id, success.output), - }; - outcome.usage = success.usage; - Ok(outcome) - } - Err(StepWorkError::Fatal(error)) => { - let message = error.to_string(); - slot.error = Some(error); - Err(SubagentError::Execution(message)) - } - Err(StepWorkError::Transient { error, tools_ran }) => { - let message = error.to_string(); - slot.error = Some(error); - Err(SubagentError::Transient { message, tools_ran }) - } - } - } -} - -/// In-memory persistence: a step is a single-process lifecycle with no pause. -#[derive(Default)] -pub(super) struct StepPersistence(Mutex<HashMap<SubagentTaskKey, SubagentOutcome>>); - -#[async_trait] -impl SubagentPersistence for StepPersistence { - async fn load_terminal( - &self, - key: &SubagentTaskKey, - ) -> Result<Option<SubagentOutcome>, SubagentError> { - Ok(self - .0 - .lock() - .expect("step persistence poisoned") - .get(key) - .cloned()) - } - async fn load(&self, _: &SubagentTaskKey) -> Result<Option<SubagentResume>, SubagentError> { - Ok(None) - } - async fn load_pause( - &self, - _: &SubagentTaskKey, - ) -> Result<Option<SubagentOutcome>, SubagentError> { - Ok(None) - } - async fn save_pause( - &self, - _: PersistedSubagentPause, - ) -> Result<SubagentPausePersistenceDisposition, SubagentError> { - Err(SubagentError::Persistence( - "agent steps never pause".to_owned(), - )) - } - async fn record_terminal( - &self, - key: &SubagentTaskKey, - outcome: &SubagentOutcome, - _: Option<&SubagentResume>, - ) -> Result<SubagentTerminalPersistenceDisposition, SubagentError> { - let mut map = self.0.lock().expect("step persistence poisoned"); - if map.contains_key(key) { - return Ok(SubagentTerminalPersistenceDisposition::Existing); - } - map.insert(key.clone(), outcome.clone()); - Ok(SubagentTerminalPersistenceDisposition::Inserted) - } -} diff --git a/crates/tinyagents-orchestration/src/subagent/step/types.rs b/crates/tinyagents-orchestration/src/subagent/step/types.rs deleted file mode 100644 index 18214953e..000000000 --- a/crates/tinyagents-orchestration/src/subagent/step/types.rs +++ /dev/null @@ -1,206 +0,0 @@ -//! Public types of the agent-step adapter: config, identity, worker -//! contract, and errors. - -use tinyagents_harness::CancellationToken; -use tinyinference_llm::usage::UsageTotals; - -use crate::subagent::{ - ResultPolicy, SpawnAdmission, SpawnRejection, SubAgentPolicy, SubagentError, SubagentOutcome, - SubagentRole, -}; - -/// Policies applied around every step run with this config. -/// -/// Cloning shares the [`SpawnAdmission`] ledger, so one config reused across a -/// team's members (or a workflow's children) enforces its limits across all -/// of them. The ledger scope is the parent run id passed in -/// [`AgentStepIdentity::parent_run_id`] (the team id / workflow run id). -#[derive(Clone, Default)] -pub struct AgentStepConfig { - /// Spawn limits (D2). Default: unlimited. - pub admission: SpawnAdmission, - /// Timeout, retry and budget (D6). Default: none. - pub policy: SubAgentPolicy, - /// Output trimming / schema check (D5). Default: inactive. - pub result_policy: ResultPolicy, - /// Delegation role (D9). Default: orchestrator. Passed to the worker in - /// [`StepContext::role`]; the worker enforces it. - pub role: SubagentRole, -} - -impl AgentStepConfig { - /// Replaces the spawn admission ledger. - pub fn with_admission(mut self, admission: SpawnAdmission) -> Self { - self.admission = admission; - self - } - - /// Replaces the timeout/retry/budget policy. - pub fn with_policy(mut self, policy: SubAgentPolicy) -> Self { - self.policy = policy; - self - } - - /// Replaces the result policy. - pub fn with_result_policy(mut self, result_policy: ResultPolicy) -> Self { - self.result_policy = result_policy; - self - } - - /// Replaces the delegation role. - pub fn with_role(mut self, role: SubagentRole) -> Self { - self.role = role; - self - } -} - -/// Durable-style identity of one step. -#[derive(Clone, Debug, PartialEq, Eq)] -pub struct AgentStepIdentity { - /// The enclosing run (team id, workflow run id). Spawn admission is scoped - /// to it and it becomes the step's parent/root run id. - pub parent_run_id: String, - /// Id of this step inside the parent run. - pub task_id: String, - /// Name checked against [`SpawnPolicy::allowed_targets`](crate::subagent::SpawnPolicy). - pub target: Option<String>, -} - -impl AgentStepIdentity { - /// An identity with no target. - pub fn new(parent_run_id: impl Into<String>, task_id: impl Into<String>) -> Self { - Self { - parent_run_id: parent_run_id.into(), - task_id: task_id.into(), - target: None, - } - } - - /// Names the step's target for allowlist checks. - pub fn with_target(mut self, target: impl Into<String>) -> Self { - self.target = Some(target.into()); - self - } -} - -/// A worker's successful result: the final text the result policy governs, -/// the usage the budget is checked against, and a host value carried back -/// untouched. -pub struct StepSuccess<T> { - /// Final visible output. - pub output: String, - /// Model usage the worker consumed. - pub usage: UsageTotals, - /// Host payload returned alongside the outcome. - pub value: T, - /// When set, the worker finished without a usable answer (a failure it - /// reports itself, not an error): the outcome is `Incomplete` with this - /// reason instead of `Completed`. - pub incomplete_reason: Option<String>, -} - -impl<T> StepSuccess<T> { - /// A success with no reported usage. - pub fn new(output: impl Into<String>, value: T) -> Self { - Self { - output: output.into(), - usage: UsageTotals::default(), - value, - incomplete_reason: None, - } - } - - /// A worker-reported failure: the outcome is `Incomplete(reason)`. - pub fn incomplete(reason: impl Into<String>, value: T) -> Self { - Self { - output: String::new(), - usage: UsageTotals::default(), - value, - incomplete_reason: Some(reason.into()), - } - } - - /// Sets the reported usage. - pub fn with_usage(mut self, usage: UsageTotals) -> Self { - self.usage = usage; - self - } -} - -/// How a worker failed. -#[derive(Debug)] -pub enum StepWorkError { - /// Not retried. - Fatal(anyhow::Error), - /// Retried per [`SubAgentPolicy::retry`] (after tool calls only when - /// `tools_ran` is false or the policy allows it). - Transient { - /// The underlying failure. - error: anyhow::Error, - /// Whether the attempt had already executed tools. - tools_ran: bool, - }, -} - -impl From<anyhow::Error> for StepWorkError { - fn from(error: anyhow::Error) -> Self { - Self::Fatal(error) - } -} - -/// Why a step did not produce an outcome. -#[derive(Debug)] -pub enum AgentStepError { - /// The spawn policy refused the step; the worker never ran. - Rejected(SpawnRejection), - /// The worker failed; carries its original error. - Worker(anyhow::Error), - /// The step was cancelled before or while running, with no worker result. - Cancelled, - /// A lifecycle failure unrelated to the worker. - Driver(SubagentError), -} - -impl std::fmt::Display for AgentStepError { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - match self { - Self::Rejected(rejection) => write!(f, "agent step was not admitted: {rejection}"), - Self::Worker(error) => error.fmt(f), - Self::Cancelled => f.write_str("agent step was cancelled"), - Self::Driver(error) => error.fmt(f), - } - } -} - -impl std::error::Error for AgentStepError {} - -/// The driver's outcome for a step plus the worker's payload, when it ran to -/// a result. -pub struct AgentStepResult<T> { - /// The post-policy outcome (status, trimmed output, usage). - pub outcome: SubagentOutcome, - /// The worker's payload from its last successful attempt. Present for - /// every status the worker produced a result for, including an outcome the - /// policies later downgraded to incomplete or cancelled. - pub value: Option<T>, -} - -/// What the driver hands the worker for one attempt. -/// -/// The crate cannot see the tools an opaque worker holds, so -/// [`AgentStepConfig::role`] cannot filter them itself: a worker that can -/// delegate must honour [`Self::role`] (drop its delegation tools when it is -/// [`SubagentRole::Leaf`]). -#[derive(Clone, Debug)] -pub struct StepContext { - /// Child token: cancelled on lifecycle cancellation or policy timeout. - pub cancellation: CancellationToken, - /// The configured delegation role the worker must enforce. - pub role: SubagentRole, - /// Model-call cap from [`SubAgentBudget`](crate::subagent::SubAgentBudget) call caps - /// (already tightened onto the child's `RunConfig`); the worker must - /// enforce it, as the crate cannot count an opaque worker's calls. - pub max_model_calls: Option<usize>, - /// Tool-call cap, as for [`Self::max_model_calls`]. - pub max_tool_calls: Option<usize>, -} diff --git a/crates/tinyagents-orchestration/src/subagent/turn_storage.rs b/crates/tinyagents-orchestration/src/subagent/turn_storage.rs new file mode 100644 index 000000000..bb3b46f5a --- /dev/null +++ b/crates/tinyagents-orchestration/src/subagent/turn_storage.rs @@ -0,0 +1,78 @@ +//! Scope-restricted, lease-conditioned storage for a parent/thread turn. + +use std::sync::Arc; + +use tinyagents_session::port::AgentStorageFactory; +use tinyagents_storage_kit::{Scope, ScopedStorage, StorageBackend, StorageError, WriteCondition}; + +/// Host storage restricted to the admitted tenant and conditioned on its +/// durable turn lease. The lease condition is checked atomically by the +/// backend for every mutation performed through the returned scoped handle. +#[derive(Clone)] +pub struct FencedSessionBackend { + backend: Arc<dyn StorageBackend>, + tenant: Scope, + condition: WriteCondition, +} + +impl std::fmt::Debug for FencedSessionBackend { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("FencedSessionBackend") + .field("tenant", &self.tenant.as_str()) + .field("condition_clauses", &self.condition.clauses().len()) + .finish_non_exhaustive() + } +} + +impl FencedSessionBackend { + /// Restricts storage to agent scopes under `tenant`, conditioned on the + /// exact durable lease tenure supplied by the caller. + pub fn new( + backend: Arc<dyn StorageBackend>, + tenant: Scope, + condition: WriteCondition, + ) -> Result<Self, StorageError> { + if condition.clauses().is_empty() { + return Err(StorageError::invalid_input( + "a session write condition must contain a lease guard", + )); + } + if tenant.as_str() == "cluster" || tenant.as_str().contains('/') { + return Err(StorageError::invalid_input( + "a conditioned session backend requires a tenant scope", + )); + } + Ok(Self { + backend, + tenant, + condition, + }) + } + + /// Builds the session provider factory for an admitted actor. Each agent + /// receives a distinct `tenant/agent:<id>` scope, and the backend opens + /// that scope only with the lease condition. Backends without atomic + /// conditional-write support fail closed. + pub fn storage_factory(self: Arc<Self>) -> AgentStorageFactory { + Arc::new(move |agent_id| self.for_agent(agent_id)) + } + + /// Opens one admitted agent's conditioned storage. + pub fn for_agent(&self, agent_id: &str) -> Result<ScopedStorage, StorageError> { + if agent_id.is_empty() { + return Err(StorageError::invalid_input("agent id must not be empty")); + } + let scope = self.tenant.child(&format!("agent:{agent_id}"))?; + self.backend.for_scope_conditioned(&scope, &self.condition) + } + + /// Returns the tenant admitted by this backend. + pub fn tenant(&self) -> &Scope { + &self.tenant + } +} + +#[cfg(test)] +#[path = "turn_storage_tests.rs"] +mod tests; diff --git a/crates/tinyagents-orchestration/src/subagent/turn_storage_tests.rs b/crates/tinyagents-orchestration/src/subagent/turn_storage_tests.rs new file mode 100644 index 000000000..bc9eacb3b --- /dev/null +++ b/crates/tinyagents-orchestration/src/subagent/turn_storage_tests.rs @@ -0,0 +1,136 @@ +use super::*; +use std::sync::Arc; +use std::sync::atomic::{AtomicU64, Ordering}; + +use tinyagents_storage_kit::{ + CollectionSpec, DocumentStore, Leases, Precondition, Scope, Storage, + StorageBackend as ScopedBackend, WriteCondition, +}; +use tinystoragedrivers_connect::sqlite::SqliteConnection; +use tinystoragedrivers_document::sqlite::SqliteDocuments; + +#[derive(Clone)] +struct SqliteBackend { + connection: SqliteConnection, +} + +impl ScopedBackend for SqliteBackend { + fn for_scope(&self, scope: &Scope) -> Result<ScopedStorage, StorageError> { + let docs = Arc::new(SqliteDocuments::new(self.connection.clone(), scope.clone())) + as Arc<dyn DocumentStore>; + Ok(ScopedStorage::new(storage_with(docs), scope.clone())) + } + + fn for_scope_conditioned( + &self, + scope: &Scope, + condition: &WriteCondition, + ) -> Result<ScopedStorage, StorageError> { + let docs = Arc::new(SqliteDocuments::conditioned( + self.connection.clone(), + scope.clone(), + condition.clone(), + )) as Arc<dyn DocumentStore>; + Ok(ScopedStorage::new(storage_with(docs), scope.clone())) + } +} + +fn storage_with(documents: Arc<dyn DocumentStore>) -> Storage { + let unsupported = Storage::unsupported(); + Storage::new( + documents, + unsupported.streams().clone(), + unsupported.kv().clone(), + unsupported.blobs().clone(), + unsupported.queues().clone(), + unsupported.secrets().clone(), + ) +} + +#[tokio::test] +async fn retained_session_documents_reject_expired_lease_after_takeover() { + let directory = tempfile::tempdir().unwrap(); + let now = Arc::new(AtomicU64::new(100)); + let clock = { + let now = Arc::clone(&now); + Arc::new(move || now.load(Ordering::SeqCst)) as Arc<dyn Fn() -> u64 + Send + Sync> + }; + let connection = + SqliteConnection::open_with_clock(directory.path().join("session.db"), clock).unwrap(); + let backend: Arc<dyn ScopedBackend> = Arc::new(SqliteBackend { connection }); + let lease_scope = Scope::new("cluster").unwrap(); + + let lease_docs = backend.for_scope(&lease_scope).unwrap().documents().clone(); + lease_docs + .ensure_collection(&CollectionSpec::new("session_probe")) + .await + .unwrap(); + let leases = Leases::new(&lease_docs).with_clock(Arc::new({ + let now = Arc::clone(&now); + move || now.load(Ordering::SeqCst) + })); + let lease_a = leases + .acquire("turn", "thread", "holder-a", 10) + .await + .unwrap() + .unwrap(); + + let adapter = Arc::new( + FencedSessionBackend::new( + Arc::clone(&backend), + Scope::new("u:alice").unwrap(), + lease_a.write_condition(lease_scope.clone()).unwrap(), + ) + .unwrap(), + ); + let stale_docs = adapter.for_agent("worker").unwrap().documents().clone(); + stale_docs + .ensure_collection(&CollectionSpec::new("session_probe")) + .await + .unwrap(); + stale_docs + .put( + "session_probe", + "stale", + serde_json::json!({"holder":"a"}), + Precondition::None, + ) + .await + .unwrap(); + + now.store(110, Ordering::SeqCst); + let lease_b = leases + .acquire("turn", "thread", "holder-b", 30) + .await + .unwrap() + .unwrap(); + let stale_write = stale_docs + .put( + "session_probe", + "stale", + serde_json::json!({"holder":"stale"}), + Precondition::None, + ) + .await; + assert!( + stale_write.is_err(), + "expired holder's retained document handle must fail closed" + ); + + let current = FencedSessionBackend::new( + backend, + Scope::new("u:alice").unwrap(), + lease_b.write_condition(lease_scope).unwrap(), + ) + .unwrap(); + let current_docs = current.for_agent("worker").unwrap().documents().clone(); + current_docs + .put( + "session_probe", + "current", + serde_json::json!({"holder":"b"}), + Precondition::None, + ) + .await + .unwrap(); +} diff --git a/crates/tinyagents-orchestration/src/subagent/types.rs b/crates/tinyagents-orchestration/src/subagent/types.rs index 5a45f683c..3d8d34797 100644 --- a/crates/tinyagents-orchestration/src/subagent/types.rs +++ b/crates/tinyagents-orchestration/src/subagent/types.rs @@ -3,7 +3,7 @@ use std::sync::Arc; use tinyagents_harness::tool::ToolSnapshot; use tinyagents_harness::{CancellationToken, context::RunContext}; -use tinyagents_tasks::NotifyMode; +use tinyagents_tasks::{DeliveryFence, NotifyMode}; use tinyinference_llm::{message::Message, usage::UsageTotals}; use super::{ResultPolicy, SubAgentPolicy, SubagentRole}; @@ -22,6 +22,7 @@ pub struct SubagentRequest<C = (), H = ()> { host_request: H, input: String, resume: Option<SubagentResume>, + turn_fence: Option<Arc<dyn DeliveryFence>>, } /// One-way planner view of a validated request. @@ -139,6 +140,7 @@ impl<C, H> SubagentRequest<C, H> { host_request, input: input.into(), resume, + turn_fence: None, }) } @@ -188,6 +190,7 @@ impl<C, H> SubagentRequest<C, H> { host_request, input: input.into(), resume, + turn_fence: None, }) } @@ -213,6 +216,7 @@ impl<C, H> SubagentRequest<C, H> { host_request, input: input.into(), resume, + turn_fence: None, }) } @@ -290,6 +294,18 @@ impl<C, H> SubagentRequest<C, H> { self.resume.as_ref() } + /// Supplies a host-owned durable turn fence for this child lifecycle. The + /// generic driver keeps the guard alive and cancels execution on loss; the + /// executor uses the same authority to bind its transcript writer. + pub fn with_turn_fence(mut self, fence: Arc<dyn DeliveryFence>) -> Self { + self.turn_fence = Some(fence); + self + } + + pub(crate) fn turn_fence(&self) -> Option<Arc<dyn DeliveryFence>> { + self.turn_fence.clone() + } + pub(crate) fn set_resume(&mut self, resume: Option<SubagentResume>) { self.resume = resume; } @@ -360,8 +376,8 @@ pub struct PreparedSubagent<C = ()> { /// The tool set the child inherited. When set, the driver intersects /// [`Self::tools`] with it so a child can never widen what it inherited. pub tool_ceiling: Option<ToolSnapshot>, - /// Names of the host's own [`SubAgentTool`](super::SubAgentTool)s, which a - /// leaf must not keep (the crate cannot know custom tool names). + /// Names of the host's own spawn tools, which a leaf must not keep (the + /// crate cannot know custom tool names). pub delegation_tools: Vec<String>, /// Authored capability scope to intersect with the supplied declarations. pub tool_scope: Option<tinyagents_definition::ToolScope>, @@ -490,6 +506,9 @@ pub struct SubagentExecution<C = ()> { pub task_key: SubagentTaskKey, /// Cooperative cancellation shared with the lifecycle owner. pub cancellation: CancellationToken, + /// Host-owned thread-turn authority retained across every executor await. + /// When set, the driver stops dispatch and persistence after ownership loss. + pub turn_fence: Option<Arc<dyn DeliveryFence>>, } /// A neutral reference to a host-owned artifact. diff --git a/crates/tinyagents-orchestration/src/teams/README.md b/crates/tinyagents-orchestration/src/teams/README.md deleted file mode 100644 index 1dc9b5808..000000000 --- a/crates/tinyagents-orchestration/src/teams/README.md +++ /dev/null @@ -1,84 +0,0 @@ -# `orchestration::teams` — durable agent-team coordination - -A **team** is a durable group of worker agents (members) who claim and -complete collaborative tasks. The module enforces team structure invariants -(no duplicate member names, valid task dependencies) and manages the event -log that records every state change: members added, tasks created, claims, -completions, and inter-member messages. - -## Public surface - -- **`TeamService<L: TeamLedger>`** — the main API: creates teams, adds members, - creates tasks, handles claims and completions, and routes messages. Generic - over a durable ledger. -- **`TeamMemberExecutor` / `TeamService::start_member`** — runs a claimed task - through a host-supplied executor, persists its evidence, and releases the - task and idles the member when execution fails or completion is rejected. -- **`TeamLedger`** — the trait for durable team state: CRUD operations for - teams, members, tasks, and events. Implementations handle all persistence. -- **`SessionTeamLedger`** — built-in ledger backed by `tinyagents-session`'s - run ledger at a caller-supplied workspace root. -- **`NewMember`, `TeamView`, `MemberShutdown`** — data types for snapshots and - operations. -- **`TeamError`** — validation errors: duplicate members, unknown members, - cyclic task dependencies, etc. -- **`run_member_graph`** — executes a member's work as a generic DAG (execute → - complete/fail → done), bridging host worker callbacks and the graph layer. - The `execute` step runs the worker through `SubagentDriver` - (`subagent::run_agent_step`) under an inert default policy, so existing - callers behave exactly as before. -- **`run_member_graph_with` / `MemberStep`** — the same graph with explicit - driver policy (`AgentStepConfig`: `SpawnPolicy` admission scoped to the team - id with the member id as target, `SubAgentPolicy::timeout`, `ResultPolicy`). - Retry, token budget and role need worker-reported transient failures / usage / - a tool surface that `MemberOutcome` does not carry, so they do not take - effect through this adapter (use `subagent::run_agent_step` directly). - `MemberStep::with_cancellation` ties a step to the run's token (a pre-cancelled member never runs and goes to `on_failed`). `MemberStep::default()` shares one scope and has no target: give a limited - config a real `MemberStep::new(config, team, member)`. - A completed run goes to `on_complete` with the policy-trimmed output; a - spawn rejection, timeout, exceeded budget or worker-reported failure goes to - `on_failed` with the reason; a worker `Err` still fails the graph run. - -## Design and invariants - -- **Durable state at every step:** Team members, tasks, events, and watermarks - are all durably persisted via the ledger. Hosts can use any storage backend - that implements [`TeamLedger`]. -- **Event-sourced messaging:** All team events (messages, member adds, task - completions) are logged in the durable run event stream. Members read their - undelivered messages and advance a delivery watermark — making repeated - reads idempotent and safe for retries. -- **Dependency-aware tasks:** Tasks form a directed acyclic graph with - dependencies. The service validates this at creation time (no cycles, no - self-dependencies, no unknown member/task references). -- **Atomic task claims:** Task claiming uses atomic compare-and-swap to ensure - exactly one member claims each task. Completion is similarly atomic, - advancing task status and optionally requiring evidence. -- **Watermark-based delivery:** Members have a message delivery watermark - (sequence number) that advances durably as messages are read. Hosts can call - message-delivery functions repeatedly without duplicating messages. - -## File map - -- **`service.rs`** — [`TeamService`] implementation: team/member/task CRUD, - validation, and coordination. [`TeamLedger`] trait definition. - [`SessionTeamLedger`] — session run-ledger adapter. -- **`types.rs`** — data types: [`NewMember`], [`TeamView`], [`MemberShutdown`], - [`TeamError`]. -- **`graph.rs`** — member worker graph: `execute` → `complete`/`fail` → `done` - DAG. Host supplies `run_worker`, `on_complete`, `on_failed` callbacks. - [`MemberOutcome`] routes on success/failure. -- **`runtime.rs`** — message delivery and prompt composition: - `deliver_pending_messages()` reads undelivered messages and advances - watermarks. `build_member_prompt()` composes the stable worker prompt from - task and delivered messages. -- **`tests.rs`** — tests for team creation, member lifecycle, task coordination, - concurrency, and message delivery. - -## Relationship to other modules - -- **Depends on:** `tinyagents-graph` (DAG execution), `tinyagents-session` (run - ledger), `tinyagents-harness` (error types). -- **Used by:** `orchestration::workflow` (to model multi-agent phases), - host code (for team management). -- **Integration:** The workflow engine can model a phase as a team of agents. diff --git a/crates/tinyagents-orchestration/src/teams/graph.rs b/crates/tinyagents-orchestration/src/teams/graph.rs deleted file mode 100644 index 1c0b5ec46..000000000 --- a/crates/tinyagents-orchestration/src/teams/graph.rs +++ /dev/null @@ -1,308 +0,0 @@ -//! Member worker graph execution: a generic three-step DAG that invokes a host -//! worker, routes on success/failure, and calls back to the host. -//! -//! This module owns the structure of member execution as seen by the graph -//! layer (entry → execute → complete/fail → done), while letting the host -//! supply the actual work (`run_worker`) and post-work effects (`on_complete` -//! and `on_failed`). It is the bridge between durable team coordination -//! (membership, tasks, event delivery) and the graph's lifecycle and event -//! streaming. - -use std::future::Future; -use std::sync::Arc; - -use anyhow::Result; -use tinyagents_graph::export::GraphTopology; -use tinyagents_graph::stream::GraphEventSink; -use tinyagents_graph::{ - ClosureStateReducer, Command, CompiledGraph, GraphBuilder, NodeContext, NodeResult, -}; - -use super::types::MemberStep; -use crate::subagent::{AgentStepError, StepSuccess, SubagentOutcomeKind, run_agent_step}; - -const LOG_PREFIX: &str = "[team-member-step]"; - -/// Terminal classification of a host worker run. -/// -/// Returned by a host's `run_worker` callback to signal whether the member -/// completed its work successfully or failed. The member graph routes on this -/// outcome and calls the appropriate host callback (`on_complete` or -/// `on_failed`). -pub enum MemberOutcome { - /// Member completed successfully; carries the output to be recorded. - Completed { output: String }, - /// Member failed; carries the failure reason to be recorded. - Failed { reason: String }, -} - -#[derive(Clone, Default)] -struct MemberState { - payload: Option<String>, -} - -enum MemberUpdate { - Payload(String), - Noop, -} - -fn graph_err(error: anyhow::Error) -> tinyagents_harness::TinyAgentsError { - tinyagents_harness::TinyAgentsError::Graph(error.to_string()) -} - -/// Run the generic complete-or-fail member graph with host supplied effects. -/// -/// `event_sink` is optional because observability belongs to the embedding host; -/// when supplied it receives the graph executor's lifecycle events unchanged. -pub async fn run_member_graph<W, WF, C, CF, F, FF>( - event_sink: Option<Arc<dyn GraphEventSink>>, - run_worker: W, - on_complete: C, - on_failed: F, -) -> Result<()> -where - W: Fn() -> WF + Clone + Send + Sync + 'static, - WF: Future<Output = Result<MemberOutcome>> + Send + 'static, - C: Fn(String) -> CF + Clone + Send + Sync + 'static, - CF: Future<Output = Result<()>> + Send + 'static, - F: Fn(String) -> FF + Clone + Send + Sync + 'static, - FF: Future<Output = Result<()>> + Send + 'static, -{ - run_member_graph_with( - event_sink, - MemberStep::default(), - run_worker, - on_complete, - on_failed, - ) - .await -} - -/// [`run_member_graph`] with explicit driver policy. -/// -/// The `execute` node runs `run_worker` through -/// [`SubagentDriver`](crate::subagent::SubagentDriver) (via -/// [`run_agent_step`]) under `step`. Outcomes map onto the existing routing: -/// a completed run goes to `on_complete` with its (result-policy trimmed) -/// output; a worker-reported failure, spawn rejection, timeout or exceeded -/// budget goes to `on_failed` with the reason; a worker `Err` still fails the -/// graph run, unchanged. -pub async fn run_member_graph_with<W, WF, C, CF, F, FF>( - event_sink: Option<Arc<dyn GraphEventSink>>, - step: MemberStep, - run_worker: W, - on_complete: C, - on_failed: F, -) -> Result<()> -where - W: Fn() -> WF + Clone + Send + Sync + 'static, - WF: Future<Output = Result<MemberOutcome>> + Send + 'static, - C: Fn(String) -> CF + Clone + Send + Sync + 'static, - CF: Future<Output = Result<()>> + Send + 'static, - F: Fn(String) -> FF + Clone + Send + Sync + 'static, - FF: Future<Output = Result<()>> + Send + 'static, -{ - let step = Arc::new(step); - let mut graph = build_member_graph( - move || { - let step = step.clone(); - let run_worker = run_worker.clone(); - async move { drive_member(&step, run_worker).await } - }, - on_complete, - on_failed, - )?; - if let Some(event_sink) = event_sink { - graph = graph.with_event_sink(event_sink); - } - graph - .run(MemberState::default()) - .await - .map_err(|error| anyhow::anyhow!("member graph run failed: {error}"))?; - Ok(()) -} - -/// Runs the host worker as one driver lifecycle and projects the typed -/// outcome back onto [`MemberOutcome`]. -async fn drive_member<W, WF>(step: &MemberStep, run_worker: W) -> Result<MemberOutcome> -where - W: Fn() -> WF + Clone + Send + Sync + 'static, - WF: Future<Output = Result<MemberOutcome>> + Send + 'static, -{ - let result = run_agent_step( - &step.config, - step.identity.clone(), - step.cancellation.clone(), - move |ctx| { - let fut = run_worker(); - async move { - // A stuck worker cannot observe the token, so race it: a - // cancelled lifecycle (or policy timeout) ends the step. - let outcome = tokio::select! { - // A worker that has already produced its result wins over - // a cancellation observed afterwards. - biased; - outcome = fut => outcome?, - _ = ctx.cancellation.cancelled() => { - // A typed incomplete result, never a worker error: - // the member is routed to `on_failed`. - return Ok(StepSuccess::incomplete("member step was cancelled", ())); - } - }; - Ok(match outcome { - MemberOutcome::Completed { output } => StepSuccess::new(output, ()), - MemberOutcome::Failed { reason } => StepSuccess::incomplete(reason, ()), - }) - } - }, - ) - .await; - match result { - Ok(result) => Ok(match result.outcome.status { - SubagentOutcomeKind::Completed => { - // `on_complete` carries only the text; findings are logged so a - // schema miss or failed artifact store is not silent. - if let Some(error) = &result.outcome.schema_error { - tracing::warn!( - "{LOG_PREFIX} schema_error member={} error={error}", - step.identity.task_id - ); - } - if let Some(error) = &result.outcome.artifact_error { - tracing::warn!( - "{LOG_PREFIX} artifact_error member={} error={error}", - step.identity.task_id - ); - } - MemberOutcome::Completed { - output: result.outcome.output, - } - } - SubagentOutcomeKind::Incomplete(incomplete) => MemberOutcome::Failed { - reason: incomplete.reason, - }, - SubagentOutcomeKind::Cancelled | SubagentOutcomeKind::AwaitingInput(_) => { - MemberOutcome::Failed { - reason: "member step was cancelled".to_owned(), - } - } - }), - Err(AgentStepError::Worker(error)) => Err(error), - Err(error @ AgentStepError::Rejected(_)) => { - tracing::debug!( - "{LOG_PREFIX} rejected member={} {error}", - step.identity.task_id - ); - Ok(MemberOutcome::Failed { - reason: error.to_string(), - }) - } - Err(AgentStepError::Cancelled) => Ok(MemberOutcome::Failed { - reason: "member step was cancelled".to_owned(), - }), - Err(error) => Err(anyhow::anyhow!("{error}")), - } -} - -fn build_member_graph<W, WF, C, CF, F, FF>( - run_worker: W, - on_complete: C, - on_failed: F, -) -> Result<CompiledGraph<MemberState, MemberUpdate>> -where - W: Fn() -> WF + Clone + Send + Sync + 'static, - WF: Future<Output = Result<MemberOutcome>> + Send + 'static, - C: Fn(String) -> CF + Clone + Send + Sync + 'static, - CF: Future<Output = Result<()>> + Send + 'static, - F: Fn(String) -> FF + Clone + Send + Sync + 'static, - FF: Future<Output = Result<()>> + Send + 'static, -{ - let mut builder = GraphBuilder::<MemberState, MemberUpdate>::new().set_reducer( - ClosureStateReducer::new(|mut state: MemberState, update: MemberUpdate| { - if let MemberUpdate::Payload(payload) = update { - state.payload = Some(payload); - } - Ok(state) - }), - ); - builder = builder.add_node( - "execute", - move |_state: MemberState, _context: NodeContext| { - let run_worker = run_worker.clone(); - async move { - match run_worker().await.map_err(graph_err)? { - MemberOutcome::Completed { output } => Ok(NodeResult::Command( - Command::default() - .with_update(MemberUpdate::Payload(output)) - .with_goto(["complete"]), - )), - MemberOutcome::Failed { reason } => Ok(NodeResult::Command( - Command::default() - .with_update(MemberUpdate::Payload(reason)) - .with_goto(["fail"]), - )), - } - } - }, - ); - builder = builder.add_node( - "complete", - move |state: MemberState, _context: NodeContext| { - let on_complete = on_complete.clone(); - async move { - on_complete(state.payload.unwrap_or_default()) - .await - .map_err(graph_err)?; - Ok(NodeResult::Update(MemberUpdate::Noop)) - } - }, - ); - builder = builder.add_node("fail", move |state: MemberState, _context: NodeContext| { - let on_failed = on_failed.clone(); - async move { - on_failed(state.payload.unwrap_or_default()) - .await - .map_err(graph_err)?; - Ok(NodeResult::Update(MemberUpdate::Noop)) - } - }); - builder - .add_node( - "done", - |_state: MemberState, _context: NodeContext| async move { - Ok(NodeResult::Update(MemberUpdate::Noop)) - }, - ) - .add_edge("complete", "done") - .add_edge("fail", "done") - .set_entry("execute") - .mark_command_routing("execute") - .set_finish("done") - .compile() - .map_err(|error| anyhow::anyhow!("member graph compile failed: {error}")) -} - -/// Structure-only view of the generic member execution graph. -/// -/// Returns a topology that reflects the nodes and edges (entry, execute, -/// complete, fail, done, routing rules) without running any real worker or -/// calling effects. Used for introspection and documentation. -pub fn member_graph_topology() -> Result<GraphTopology> { - Ok(build_member_graph( - || async { - Ok(MemberOutcome::Completed { - output: String::new(), - }) - }, - |_| async { Ok(()) }, - |_| async { Ok(()) }, - )? - .topology()) -} - -#[cfg(test)] -mod tests; - -#[cfg(test)] -#[path = "graph_driver_tests.rs"] -mod driver_tests; diff --git a/crates/tinyagents-orchestration/src/teams/graph/tests.rs b/crates/tinyagents-orchestration/src/teams/graph/tests.rs deleted file mode 100644 index c0fdd55d8..000000000 --- a/crates/tinyagents-orchestration/src/teams/graph/tests.rs +++ /dev/null @@ -1,100 +0,0 @@ -//! Tests for member worker graph execution. - -use std::sync::Arc; -use std::sync::atomic::{AtomicBool, Ordering}; - -use tinyagents_graph::stream::CollectingSink; - -use super::*; - -#[tokio::test] -async fn member_graph_routes_completed_and_failed_workers() { - let complete = Arc::new(AtomicBool::new(false)); - let failed = Arc::new(AtomicBool::new(false)); - let complete_flag = complete.clone(); - let failed_flag = failed.clone(); - run_member_graph( - None, - || async { - Ok(MemberOutcome::Completed { - output: "done".into(), - }) - }, - move |output| { - let complete = complete_flag.clone(); - async move { - assert_eq!(output, "done"); - complete.store(true, Ordering::SeqCst); - Ok(()) - } - }, - move |_| { - let failed = failed_flag.clone(); - async move { - failed.store(true, Ordering::SeqCst); - Ok(()) - } - }, - ) - .await - .unwrap(); - assert!(complete.load(Ordering::SeqCst)); - assert!(!failed.load(Ordering::SeqCst)); - - let failed = Arc::new(AtomicBool::new(false)); - let failed_flag = failed.clone(); - run_member_graph( - None, - || async { - Ok(MemberOutcome::Failed { - reason: "boom".into(), - }) - }, - |_| async { Ok(()) }, - move |reason| { - let failed = failed_flag.clone(); - async move { - assert_eq!(reason, "boom"); - failed.store(true, Ordering::SeqCst); - Ok(()) - } - }, - ) - .await - .unwrap(); - assert!(failed.load(Ordering::SeqCst)); -} - -#[tokio::test] -async fn worker_engine_errors_propagate() { - let result = run_member_graph( - None, - || async { Err(anyhow::anyhow!("worker unavailable")) }, - |_| async { Ok(()) }, - |_| async { Ok(()) }, - ) - .await; - assert!(result.is_err()); -} - -#[tokio::test] -async fn injected_event_sink_observes_member_graph_lifecycle() { - let sink = Arc::new(CollectingSink::new()); - run_member_graph( - Some(sink.clone()), - || async { - Ok(MemberOutcome::Completed { - output: "done".into(), - }) - }, - |_| async { Ok(()) }, - |_| async { Ok(()) }, - ) - .await - .unwrap(); - - assert!( - !sink.is_empty(), - "an injected sink must receive the member graph lifecycle" - ); -} diff --git a/crates/tinyagents-orchestration/src/teams/graph_driver_tests.rs b/crates/tinyagents-orchestration/src/teams/graph_driver_tests.rs deleted file mode 100644 index 43a1392fc..000000000 --- a/crates/tinyagents-orchestration/src/teams/graph_driver_tests.rs +++ /dev/null @@ -1,225 +0,0 @@ -//! Member steps run through `SubagentDriver`: configured driver policy takes -//! effect, and the default step is behaviour-neutral. - -use super::*; - -use std::sync::{Arc, Mutex}; -use std::time::Duration; - -use tinyagents_harness::CancellationToken; - -use crate::subagent::{AgentStepConfig, ResultPolicy, SpawnAdmission, SpawnPolicy, SubAgentPolicy}; - -#[derive(Default, Clone)] -struct Seen { - completed: Arc<Mutex<Vec<String>>>, - failed: Arc<Mutex<Vec<String>>>, -} - -async fn run(step: MemberStep, seen: &Seen, worker: MemberOutcome) -> Result<()> { - let (c, f) = (seen.completed.clone(), seen.failed.clone()); - let worker = Arc::new(Mutex::new(Some(worker))); - run_member_graph_with( - None, - step, - move || { - let worker = worker.clone(); - async move { Ok(worker.lock().unwrap().take().expect("worker ran once")) } - }, - move |o| { - let c = c.clone(); - async move { - c.lock().unwrap().push(o); - Ok(()) - } - }, - move |r| { - let f = f.clone(); - async move { - f.lock().unwrap().push(r); - Ok(()) - } - }, - ) - .await -} - -fn done(text: &str) -> MemberOutcome { - MemberOutcome::Completed { - output: text.into(), - } -} - -#[tokio::test] -async fn result_policy_cap_applies_to_a_member_output() { - let seen = Seen::default(); - let config = - AgentStepConfig::default().with_result_policy(ResultPolicy::new().with_max_chars(8)); - let long = "x".repeat(500); - run(MemberStep::new(config, "team", "m1"), &seen, done(&long)) - .await - .unwrap(); - let completed = seen.completed.lock().unwrap(); - assert_eq!(completed.len(), 1); - assert!( - completed[0].chars().count() < 500, - "output was trimmed by ResultPolicy" - ); -} - -#[tokio::test] -async fn spawn_policy_denial_fails_the_member_without_running_the_worker() { - let seen = Seen::default(); - let config = AgentStepConfig::default().with_admission(SpawnAdmission::new(SpawnPolicy { - allowed_targets: Some(vec!["allowed".into()]), - ..Default::default() - })); - // The worker closure panics if it is ever taken: run() unwraps `take()`. - run( - MemberStep::new(config, "team", "intruder"), - &seen, - done("never"), - ) - .await - .unwrap(); - assert!(seen.completed.lock().unwrap().is_empty()); - let failed = seen.failed.lock().unwrap(); - assert_eq!(failed.len(), 1); - assert!(failed[0].contains("not admitted"), "got {failed:?}"); -} - -#[tokio::test] -async fn shared_admission_caps_members_across_a_team() { - let seen = Seen::default(); - let config = AgentStepConfig::default().with_admission(SpawnAdmission::new(SpawnPolicy { - max_total_per_root: Some(1), - ..Default::default() - })); - run( - MemberStep::new(config.clone(), "team", "m1"), - &seen, - done("a"), - ) - .await - .unwrap(); - run(MemberStep::new(config, "team", "m2"), &seen, done("b")) - .await - .unwrap(); - assert_eq!(seen.completed.lock().unwrap().as_slice(), ["a"]); - assert_eq!(seen.failed.lock().unwrap().len(), 1); -} - -#[tokio::test] -async fn timeout_policy_fails_a_hung_member() { - let seen = Seen::default(); - let config = AgentStepConfig::default() - .with_policy(SubAgentPolicy::default().with_timeout(Duration::from_millis(20))); - let f = seen.failed.clone(); - run_member_graph_with( - None, - MemberStep::new(config, "team", "m1"), - || async { - std::future::pending::<()>().await; - Ok(done("never")) - }, - |_| async { Ok(()) }, - move |r| { - let f = f.clone(); - async move { - f.lock().unwrap().push(r); - Ok(()) - } - }, - ) - .await - .unwrap(); - assert!(seen.failed.lock().unwrap()[0].contains("timed out")); -} - -#[tokio::test] -async fn default_step_is_neutral_and_worker_errors_still_fail_the_graph() { - let seen = Seen::default(); - run(MemberStep::default(), &seen, done("same")) - .await - .unwrap(); - assert_eq!(seen.completed.lock().unwrap().as_slice(), ["same"]); - let err = run_member_graph( - None, - || async { Err::<MemberOutcome, _>(anyhow::anyhow!("worker exploded")) }, - |_| async { Ok(()) }, - |_| async { Ok(()) }, - ) - .await - .unwrap_err(); - assert!(err.to_string().contains("worker exploded"), "{err}"); -} - -#[tokio::test] -async fn pre_cancelled_member_never_runs_its_worker() { - let seen = Seen::default(); - let token = CancellationToken::new(); - token.cancel(); - let step = MemberStep::default().with_cancellation(token); - let f = seen.failed.clone(); - let ran = Arc::new(Mutex::new(0u32)); - let counter = ran.clone(); - run_member_graph_with( - None, - step, - move || { - *counter.lock().unwrap() += 1; - async { Ok(done("never")) } - }, - |_| async { Ok(()) }, - move |r| { - let f = f.clone(); - async move { - f.lock().unwrap().push(r); - Ok(()) - } - }, - ) - .await - .unwrap(); - assert_eq!(*ran.lock().unwrap(), 0); - assert!(seen.failed.lock().unwrap()[0].contains("cancelled")); -} - -#[tokio::test] -async fn cancelling_while_the_worker_runs_routes_to_on_failed() { - let seen = Seen::default(); - let token = CancellationToken::new(); - let step = MemberStep::default().with_cancellation(token.clone()); - let f = seen.failed.clone(); - let cancel = token.clone(); - let started = Arc::new(tokio::sync::Notify::new()); - let signal = started.clone(); - tokio::spawn(async move { - // Cancel only once the worker has actually started. - started.notified().await; - cancel.cancel(); - }); - run_member_graph_with( - None, - step, - move || { - let signal = signal.clone(); - async move { - signal.notify_one(); - std::future::pending::<()>().await; - Ok(done("never")) - } - }, - |_| async { Ok(()) }, - move |r| { - let f = f.clone(); - async move { - f.lock().unwrap().push(r); - Ok(()) - } - }, - ) - .await - .unwrap(); - assert!(seen.failed.lock().unwrap()[0].contains("cancelled")); -} diff --git a/crates/tinyagents-orchestration/src/teams/mod.rs b/crates/tinyagents-orchestration/src/teams/mod.rs deleted file mode 100644 index 912f05575..000000000 --- a/crates/tinyagents-orchestration/src/teams/mod.rs +++ /dev/null @@ -1,34 +0,0 @@ -//! Durable, dependency-aware agent-team composition. -//! -//! A **team** is a durable group of worker agents (members) who claim and -//! complete collaborative tasks. This module owns team creation, member -//! lifecycle, task assignment and completion, and the event log that records -//! all durable state changes. -//! -//! [`TeamService`] is the primary API: it validates team structure, manages -//! member and task persistence (via the [`TeamLedger`] trait), and enforces -//! coordination invariants (no duplicate names, no cycles in task dependencies, -//! no dangling member or task references). `runtime` handles the per-member -//! details: reading undelivered messages from the event log and composing the -//! prompt a worker should receive. `graph` executes a member's work as a -//! generic execute → complete/fail → done DAG, bridging the graph layer and -//! durable team state. - -mod graph; -mod runtime; -mod service; -mod types; - -pub use graph::{MemberOutcome, member_graph_topology, run_member_graph, run_member_graph_with}; -pub use runtime::{ - DeliveredMessages, EVENT_PAGE_SIZE, MESSAGE_DELIVERED_EVENT, TEAM_MESSAGE_EVENT, - build_member_prompt, deliver_pending_messages, drain_run_events, truncate_chars, -}; -pub use service::{ - SessionTeamLedger, TeamLedger, TeamMemberExecutor, TeamMemberStartError, TeamMemberStarted, - TeamService, claimable_task, -}; -pub use types::{LEAD_SENDER, MemberShutdown, MemberStep, NewMember, TeamError, TeamView}; - -#[cfg(test)] -mod tests; diff --git a/crates/tinyagents-orchestration/src/teams/runtime.rs b/crates/tinyagents-orchestration/src/teams/runtime.rs deleted file mode 100644 index 32ba92e1a..000000000 --- a/crates/tinyagents-orchestration/src/teams/runtime.rs +++ /dev/null @@ -1,150 +0,0 @@ -//! Host-neutral mechanics used when a durable team member starts work. -//! -//! Hosts decide when and how to run a worker. This module only reads the -//! durable event log, advances delivery watermarks, and produces the prompt -//! that a worker should receive at its start boundary. - -use anyhow::Result; -use serde_json::json; -use tinyagents_session::run_ledger::{ - AgentTeamTask, RunEvent, RunEventAppend, RunEventListRequest, -}; - -use super::TeamLedger; - -/// Event type used for durable lead and teammate messages. -pub const TEAM_MESSAGE_EVENT: &str = "team_message"; -/// Event type used to record that a member consumed messages through a sequence. -pub const MESSAGE_DELIVERED_EVENT: &str = "team_message_delivered"; -/// Maximum page accepted by the session run ledger. -pub const EVENT_PAGE_SIZE: u32 = 1_000; - -/// The result of reading one member's undelivered team messages. -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct DeliveredMessages { - /// Message bodies in durable sequence order. - pub messages: Vec<String>, - /// The highest delivered message sequence, when a watermark was written. - pub up_to_sequence: Option<u64>, -} - -/// Compose the stable worker prompt from the claimed task and delivered messages. -pub fn build_member_prompt(task: &AgentTeamTask, messages: &[String]) -> String { - let mut prompt = format!("You are a teammate on an agent team. Task: {}", task.title); - if let Some(objective) = task - .objective - .as_deref() - .filter(|value| !value.trim().is_empty()) - { - prompt.push_str("\n\nObjective:\n"); - prompt.push_str(objective.trim()); - } - if !messages.is_empty() { - prompt.push_str("\n\nMessages from your lead / teammates:\n"); - for message in messages { - prompt.push_str("- "); - prompt.push_str(message); - prompt.push('\n'); - } - } - prompt.push_str("\n\nComplete the task and report what you did."); - prompt -} - -/// Drain every durable event in ascending sequence order. -/// -/// The session ledger caps each request, so callers must use this rather than -/// assuming that one event-list request represents a complete team history. -pub fn drain_run_events<L: TeamLedger>(ledger: &L, team_id: &str) -> Result<Vec<RunEvent>> { - let mut events = Vec::new(); - let mut after_sequence = None; - loop { - let page = ledger.list_events(&RunEventListRequest { - run_id: team_id.to_string(), - after_sequence, - limit: Some(EVENT_PAGE_SIZE), - })?; - let exhausted = page.len() < EVENT_PAGE_SIZE as usize; - after_sequence = page.last().map(|event| event.sequence); - events.extend(page); - if exhausted || after_sequence.is_none() { - return Ok(events); - } - } -} - -/// Read messages for a member, then durably advance that member's watermark. -/// -/// Direct messages and broadcasts are selected; messages for other members -/// stay untouched. The watermark is append-only, which makes repeated calls -/// idempotent while retaining the full event history for audit and replay. -pub fn deliver_pending_messages<L: TeamLedger>( - ledger: &L, - team_id: &str, - member_id: &str, -) -> Result<DeliveredMessages> { - let events = drain_run_events(ledger, team_id)?; - let watermark = events - .iter() - .filter(|event| event.event_type == MESSAGE_DELIVERED_EVENT) - .filter(|event| { - event - .payload - .get("memberId") - .and_then(|value| value.as_str()) - == Some(member_id) - }) - .filter_map(|event| { - event - .payload - .get("upToSeq") - .and_then(|value| value.as_u64()) - }) - .max() - .unwrap_or_default(); - - let mut up_to_sequence = watermark; - let mut messages = Vec::new(); - for event in &events { - if event.event_type != TEAM_MESSAGE_EVENT || event.sequence <= watermark { - continue; - } - let recipient = event.payload.get("to").and_then(|value| value.as_str()); - if recipient.is_none() || recipient == Some(member_id) { - if let Some(content) = event - .payload - .get("content") - .and_then(|value| value.as_str()) - { - messages.push(content.to_string()); - } - up_to_sequence = up_to_sequence.max(event.sequence); - } - } - - let delivered_up_to = (up_to_sequence > watermark).then_some(up_to_sequence); - if let Some(up_to_sequence) = delivered_up_to { - ledger.append_event(RunEventAppend { - run_id: team_id.to_string(), - event_type: MESSAGE_DELIVERED_EVENT.to_string(), - payload: json!({ "memberId": member_id, "upToSeq": up_to_sequence }), - })?; - } - Ok(DeliveredMessages { - messages, - up_to_sequence: delivered_up_to, - }) -} - -/// Truncate text at a Unicode character boundary and append an ellipsis. -pub fn truncate_chars(value: &str, max_chars: usize) -> String { - if value.chars().count() <= max_chars { - return value.to_string(); - } - let mut output: String = value.chars().take(max_chars).collect(); - output.push('…'); - output -} - -#[cfg(test)] -mod tests; diff --git a/crates/tinyagents-orchestration/src/teams/runtime/tests.rs b/crates/tinyagents-orchestration/src/teams/runtime/tests.rs deleted file mode 100644 index f9cd9b4ac..000000000 --- a/crates/tinyagents-orchestration/src/teams/runtime/tests.rs +++ /dev/null @@ -1,89 +0,0 @@ -//! Tests for member prompt composition and message delivery. - -use std::sync::Arc; - -use tempfile::TempDir; -use tinystoragedrivers_testkit::MemoryBackend; - -use super::*; -use crate::teams::{NewMember, SessionTeamLedger, TeamService}; - -fn team(_dir: &TempDir) -> (SessionTeamLedger, String, String) { - let ledger_backend = MemoryBackend::new(); - let ledger = - SessionTeamLedger::new(Arc::new(crate::test_support::ledger_in(&ledger_backend))).unwrap(); - let service = TeamService::new(ledger.clone()); - let team = service - .create_team( - "lead", - None, - None, - &[NewMember { - name: "member".into(), - agent_id: None, - }], - ) - .unwrap(); - (ledger, team.team.id, team.members[0].id.clone()) -} - -#[test] -fn delivery_selects_direct_and_broadcast_messages_once() { - let dir = TempDir::new().unwrap(); - let (ledger, team_id, member_id) = team(&dir); - let service = TeamService::new(ledger.clone()); - service - .message_member(&team_id, None, Some(&member_id), "direct", None) - .unwrap(); - service - .message_member(&team_id, None, None, "broadcast", None) - .unwrap(); - let first = deliver_pending_messages(&ledger, &team_id, &member_id).unwrap(); - assert_eq!(first.messages, ["direct", "broadcast"]); - assert!(first.up_to_sequence.is_some()); - assert!( - deliver_pending_messages(&ledger, &team_id, &member_id) - .unwrap() - .messages - .is_empty() - ); -} - -#[test] -fn delivery_pages_past_the_session_ledger_cap() { - let dir = TempDir::new().unwrap(); - let (ledger, team_id, member_id) = team(&dir); - for index in 0..=EVENT_PAGE_SIZE { - ledger - .append_event(tinyagents_session::run_ledger::RunEventAppend { - run_id: team_id.clone(), - event_type: "noise".into(), - payload: serde_json::json!({ "index": index }), - }) - .unwrap(); - } - let service = TeamService::new(ledger.clone()); - service - .message_member(&team_id, None, Some(&member_id), "late", None) - .unwrap(); - assert_eq!( - deliver_pending_messages(&ledger, &team_id, &member_id) - .unwrap() - .messages, - ["late"] - ); -} - -#[test] -fn prompt_and_truncation_preserve_text_boundaries() { - let dir = TempDir::new().unwrap(); - let (ledger, team_id, _) = team(&dir); - let service = TeamService::new(ledger); - let task = service - .assign_task(&team_id, "Ship", Some(" Build it "), None, &[]) - .unwrap(); - let prompt = build_member_prompt(&task, &["coordinate".into()]); - assert!(prompt.contains("Build it")); - assert!(prompt.contains("coordinate")); - assert_eq!(truncate_chars("aébc", 2), "aé…"); -} diff --git a/crates/tinyagents-orchestration/src/teams/service.rs b/crates/tinyagents-orchestration/src/teams/service.rs deleted file mode 100644 index c1e750bbd..000000000 --- a/crates/tinyagents-orchestration/src/teams/service.rs +++ /dev/null @@ -1,1001 +0,0 @@ -//! Team service and durable ledger: the public API for team management. -//! -//! [`TeamService`] validates team structure, manages member and task -//! persistence (via [`TeamLedger`]), and enforces coordination invariants. -//! [`SessionTeamLedger`] provides a built-in `tinyagents-session` backend; -//! hosts can supply their own ledger implementation for testing or custom -//! storage. - -use std::collections::HashSet; -use std::sync::Arc; - -use anyhow::{Result, anyhow}; -use async_trait::async_trait; -use chrono::Utc; -use serde_json::json; -use tinyagents_graph::dag::{DagNode, has_cycle}; -use tinyagents_session::run_ledger::{ - AgentTeam, AgentTeamListRequest, AgentTeamListResponse, AgentTeamMember, - AgentTeamMemberRunAdmission, AgentTeamMemberStatus, AgentTeamMemberUpsert, AgentTeamStatus, - AgentTeamTask, AgentTeamTaskStatus, AgentTeamTaskUpsert, AgentTeamUpsert, ClaimOutcome, - CompletionOutcome, RunEvent, RunEventAppend, RunEventListRequest, RunLedgerStore, -}; -use tinyagents_storage_kit::Blocking; -use uuid::Uuid; - -use super::runtime::drain_run_events; -use super::{LEAD_SENDER, MemberShutdown, NewMember, TEAM_MESSAGE_EVENT, TeamError, TeamView}; - -/// Host-provided execution for a claimed team member task. -#[async_trait] -pub trait TeamMemberExecutor: Send + Sync { - /// Run `task` as `member` and return the evidence to submit at completion. - async fn execute(&self, member: &AgentTeamMember, task: &AgentTeamTask) -> Result<Vec<String>>; -} - -/// Typed claim failures from [`TeamService::start_member_with_started`]. -#[derive(Debug)] -pub enum TeamMemberStartError { - /// The task was already claimed. - AlreadyClaimed { task_id: String }, - /// One or more task dependencies are incomplete. - Blocked { task_id: String, unmet: Vec<String> }, - /// The task does not exist in the requested team. - UnknownTask { task_id: String }, - /// The member already owns another active task. - AlreadyActive { - /// Member whose active run prevented admission. - member_id: String, - /// Task currently assigned to that member, if recorded. - current_task_id: Option<String>, - }, - /// Team/member validation or a ledger operation failed. - Other(anyhow::Error), -} - -impl std::fmt::Display for TeamMemberStartError { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - match self { - Self::AlreadyClaimed { task_id } => write!(f, "task is already claimed: {task_id}"), - Self::Blocked { unmet, .. } => { - write!(f, "task dependencies are incomplete: {}", unmet.join(", ")) - } - Self::UnknownTask { task_id } => write!(f, "unknown task: {task_id}"), - Self::AlreadyActive { member_id, .. } => { - write!(f, "team member is already active: {member_id}") - } - Self::Other(error) => error.fmt(f), - } - } -} - -impl std::error::Error for TeamMemberStartError { - fn source(&self) -> Option<&(dyn std::error::Error + 'static)> { - match self { - Self::Other(error) => Some(error.as_ref()), - _ => None, - } - } -} - -/// Durable identity delivered after a team member claim becomes active. -#[derive(Clone, Debug)] -pub struct TeamMemberStarted { - /// Run id recorded on the active member row. - pub run_id: String, - /// Active member row, including its worker thread and run ids. - pub member: AgentTeamMember, - /// Task returned by the successful durable claim. - pub task: AgentTeamTask, -} - -/// Durable team state required by [`TeamService`]. -/// -/// The trait intentionally mirrors the session run-ledger operations rather -/// than introducing another task store. Hosts can substitute a fake ledger in -/// tests or select their own durable implementation. -pub trait TeamLedger: Send + Sync { - /// Upserts a team row: creates if absent, updates if present. - fn upsert_team(&self, upsert: AgentTeamUpsert) -> Result<AgentTeam>; - /// Retrieves a team by id. - fn get_team(&self, id: &str) -> Result<Option<AgentTeam>>; - /// Lists teams according to the request filters. - fn list_teams(&self, request: &AgentTeamListRequest) -> Result<AgentTeamListResponse>; - /// Upserts a team member: creates if absent, updates if present. - fn upsert_member(&self, upsert: AgentTeamMemberUpsert) -> Result<AgentTeamMember>; - /// Lists members of a team in creation order. - fn list_members(&self, team_id: &str) -> Result<Vec<AgentTeamMember>>; - /// Lists tasks assigned to a team in creation order. - fn list_tasks(&self, team_id: &str) -> Result<Vec<AgentTeamTask>>; - /// Upserts a task: creates if absent, updates if present. - fn upsert_task(&self, upsert: AgentTeamTaskUpsert) -> Result<AgentTeamTask>; - /// Attempts to claim a task for a member using an atomic CAS operation. - /// Returns the claim outcome (success, already claimed, not found). - fn claim_task( - &self, - team_id: &str, - task_id: &str, - member_id: &str, - claim_token: &str, - ) -> Result<ClaimOutcome>; - /// Records a task completion with optional evidence, atomically advancing - /// the task status and optionally validating evidence before transition. - fn complete_task( - &self, - team_id: &str, - task_id: &str, - member_id: &str, - evidence: &[String], - require_evidence: bool, - ) -> Result<CompletionOutcome>; - /// Completes a task only while `member_id` still owns `claim_token`. - fn complete_task_if_claimed( - &self, - team_id: &str, - task_id: &str, - member_id: &str, - claim_token: &str, - evidence: &[String], - require_evidence: bool, - ) -> Result<CompletionOutcome>; - /// Releases a claimed task back to the queue after a worker failure. - /// Implementations must only release the claim held by `member_id` with - /// `claim_token`, so cleanup from an old worker cannot release a successor. - fn release_task( - &self, - team_id: &str, - task_id: &str, - member_id: &str, - claim_token: &str, - ) -> Result<bool>; - /// Legacy unconditional member-running setter. It does not reserve against - /// another active run; new task admission must use - /// [`TeamLedger::try_mark_member_running`]. - fn mark_member_running( - &self, - team_id: &str, - member_id: &str, - task_id: &str, - worker_thread_id: &str, - run_id: &str, - ) -> Result<Option<AgentTeamMember>>; - /// Atomically reserves the member for a task. The default fails closed so - /// legacy ledgers cannot silently permit overlapping member runs. - fn try_mark_member_running( - &self, - team_id: &str, - member_id: &str, - task_id: &str, - worker_thread_id: &str, - run_id: &str, - ) -> Result<AgentTeamMemberRunAdmission> { - let _ = (team_id, member_id, task_id, worker_thread_id, run_id); - Err(anyhow!( - "atomic team member admission is unsupported by this ledger" - )) - } - /// Marks a running member idle after its worker finishes. - fn mark_member_idle( - &self, - team_id: &str, - member_id: &str, - task_id: &str, - worker_thread_id: &str, - run_id: &str, - ) -> Result<bool>; - /// Stops a member and releases its claimed tasks, returning the stopped - /// member and the released task ids. - fn shutdown_member( - &self, - team_id: &str, - member_id: &str, - ) -> Result<Option<(AgentTeamMember, Vec<String>)>>; - /// Appends an event to the durable run event log. - fn append_event(&self, event: RunEventAppend) -> Result<RunEvent>; - /// Lists events from the run event log according to the request filters. - fn list_events(&self, request: &RunEventListRequest) -> Result<Vec<RunEvent>>; -} - -/// [`TeamLedger`] over a bound `tinyagents-session` run-ledger handle. -/// -/// The handle is whatever backend the host opened (SQLite, MongoDB, memory); -/// this adapter opens no file, picks no workspace and has no fallback. Every -/// call reaches the ledger's async operations through a [`Blocking`] bridge -/// thread, so it works from inside or outside a tokio runtime. The claim, -/// completion and shutdown operations are the ledger's atomic ones: a contested -/// task goes to exactly one member. -#[derive(Clone)] -pub struct SessionTeamLedger { - ledger: Arc<dyn RunLedgerStore>, - bridge: Blocking, -} - -impl std::fmt::Debug for SessionTeamLedger { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("SessionTeamLedger").finish_non_exhaustive() - } -} - -impl SessionTeamLedger { - /// Adapt `ledger`, starting a bridge thread for the synchronous calls. - /// - /// # Errors - /// - /// The bridge thread cannot start. - pub fn new(ledger: Arc<dyn RunLedgerStore>) -> Result<Self> { - Ok(Self::with_bridge(ledger, Blocking::new()?)) - } - - /// Adapt `ledger`, sharing an existing bridge. - pub fn with_bridge(ledger: Arc<dyn RunLedgerStore>, bridge: Blocking) -> Self { - Self { ledger, bridge } - } - - fn run<T: Send + 'static>( - &self, - work: impl std::future::Future< - Output = std::result::Result<T, tinyagents_session::run_ledger::LedgerError>, - > + Send - + 'static, - ) -> Result<T> { - Ok(self.bridge.run(work)??) - } -} - -impl TeamLedger for SessionTeamLedger { - fn upsert_team(&self, upsert: AgentTeamUpsert) -> Result<AgentTeam> { - let ledger = Arc::clone(&self.ledger); - self.run(async move { ledger.upsert_agent_team(upsert).await }) - } - fn get_team(&self, id: &str) -> Result<Option<AgentTeam>> { - let (ledger, id) = (Arc::clone(&self.ledger), id.to_owned()); - self.run(async move { ledger.get_agent_team(&id).await }) - } - fn list_teams(&self, request: &AgentTeamListRequest) -> Result<AgentTeamListResponse> { - let (ledger, request) = (Arc::clone(&self.ledger), request.clone()); - self.run(async move { ledger.list_agent_teams(&request).await }) - } - fn upsert_member(&self, upsert: AgentTeamMemberUpsert) -> Result<AgentTeamMember> { - let ledger = Arc::clone(&self.ledger); - self.run(async move { ledger.upsert_agent_team_member(upsert).await }) - } - fn list_members(&self, team_id: &str) -> Result<Vec<AgentTeamMember>> { - let (ledger, team_id) = (Arc::clone(&self.ledger), team_id.to_owned()); - self.run(async move { ledger.list_agent_team_members(&team_id).await }) - } - fn list_tasks(&self, team_id: &str) -> Result<Vec<AgentTeamTask>> { - let (ledger, team_id) = (Arc::clone(&self.ledger), team_id.to_owned()); - self.run(async move { ledger.list_agent_team_tasks(&team_id).await }) - } - fn upsert_task(&self, upsert: AgentTeamTaskUpsert) -> Result<AgentTeamTask> { - let ledger = Arc::clone(&self.ledger); - self.run(async move { ledger.upsert_agent_team_task(upsert).await }) - } - fn claim_task( - &self, - team_id: &str, - task_id: &str, - member_id: &str, - claim_token: &str, - ) -> Result<ClaimOutcome> { - let ledger = Arc::clone(&self.ledger); - let (team, task, member, token) = ( - team_id.to_owned(), - task_id.to_owned(), - member_id.to_owned(), - claim_token.to_owned(), - ); - self.run(async move { - ledger - .claim_agent_team_task(&team, &task, &member, &token) - .await - }) - } - fn complete_task( - &self, - team_id: &str, - task_id: &str, - member_id: &str, - evidence: &[String], - require_evidence: bool, - ) -> Result<CompletionOutcome> { - let ledger = Arc::clone(&self.ledger); - let (team, task, member, evidence) = ( - team_id.to_owned(), - task_id.to_owned(), - member_id.to_owned(), - evidence.to_vec(), - ); - self.run(async move { - ledger - .complete_agent_team_task(&team, &task, &member, &evidence, require_evidence) - .await - }) - } - fn complete_task_if_claimed( - &self, - team_id: &str, - task_id: &str, - member_id: &str, - claim_token: &str, - evidence: &[String], - require_evidence: bool, - ) -> Result<CompletionOutcome> { - let (ledger, team, task, member, token, evidence) = ( - Arc::clone(&self.ledger), - team_id.to_owned(), - task_id.to_owned(), - member_id.to_owned(), - claim_token.to_owned(), - evidence.to_vec(), - ); - self.run(async move { - ledger - .complete_agent_team_task_if_claimed( - &team, - &task, - &member, - &token, - &evidence, - require_evidence, - ) - .await - }) - } - fn release_task( - &self, - team_id: &str, - task_id: &str, - member_id: &str, - claim_token: &str, - ) -> Result<bool> { - let (ledger, team, task, member, token) = ( - Arc::clone(&self.ledger), - team_id.to_owned(), - task_id.to_owned(), - member_id.to_owned(), - claim_token.to_owned(), - ); - self.run(async move { - ledger - .release_agent_team_task_if_claimed(&team, &task, &member, &token) - .await - }) - } - fn mark_member_running( - &self, - team_id: &str, - member_id: &str, - task_id: &str, - worker_thread_id: &str, - run_id: &str, - ) -> Result<Option<AgentTeamMember>> { - let (ledger, team, member, task, thread, run) = ( - Arc::clone(&self.ledger), - team_id.to_owned(), - member_id.to_owned(), - task_id.to_owned(), - worker_thread_id.to_owned(), - run_id.to_owned(), - ); - self.run(async move { - ledger - .mark_agent_team_member_running(&team, &member, &task, &thread, &run) - .await - }) - } - fn try_mark_member_running( - &self, - team_id: &str, - member_id: &str, - task_id: &str, - worker_thread_id: &str, - run_id: &str, - ) -> Result<AgentTeamMemberRunAdmission> { - let (ledger, team, member, task, thread, run) = ( - Arc::clone(&self.ledger), - team_id.to_owned(), - member_id.to_owned(), - task_id.to_owned(), - worker_thread_id.to_owned(), - run_id.to_owned(), - ); - self.run(async move { - ledger - .try_start_agent_team_member(&team, &member, &task, &thread, &run) - .await - }) - } - fn mark_member_idle( - &self, - team_id: &str, - member_id: &str, - task_id: &str, - worker_thread_id: &str, - run_id: &str, - ) -> Result<bool> { - let (ledger, team, member, task, thread, run) = ( - Arc::clone(&self.ledger), - team_id.to_owned(), - member_id.to_owned(), - task_id.to_owned(), - worker_thread_id.to_owned(), - run_id.to_owned(), - ); - self.run(async move { - Ok(ledger - .mark_agent_team_member_idle_if_current(&team, &member, &task, &thread, &run) - .await? - .is_some()) - }) - } - fn shutdown_member( - &self, - team_id: &str, - member_id: &str, - ) -> Result<Option<(AgentTeamMember, Vec<String>)>> { - let ledger = Arc::clone(&self.ledger); - let (team, member) = (team_id.to_owned(), member_id.to_owned()); - self.run(async move { ledger.shutdown_agent_team_member(&team, &member).await }) - } - fn append_event(&self, event: RunEventAppend) -> Result<RunEvent> { - let ledger = Arc::clone(&self.ledger); - self.run(async move { ledger.append_run_event(event).await }) - } - fn list_events(&self, request: &RunEventListRequest) -> Result<Vec<RunEvent>> { - let (ledger, request) = (Arc::clone(&self.ledger), request.clone()); - self.run(async move { Ok(ledger.list_recent_run_events(&request).await?.events) }) - } -} - -/// Host-neutral service for durable, dependency-aware agent teams. -/// -/// Provides a high-level API for team management: creating teams, adding -/// members, creating and claiming tasks, composing prompts, and shutting down -/// members. All mutations are durably persisted via a caller-supplied -/// [`TeamLedger`]; the service enforces coordination invariants (no duplicate -/// member names, valid task dependencies, no cycles). -/// -/// Generic over the ledger to allow hosts to inject their own storage -/// implementation or a test double. -#[derive(Debug, Clone)] -pub struct TeamService<L> { - ledger: L, -} - -struct MemberExecutionCleanup<'a, L: TeamLedger> { - service: &'a TeamService<L>, - team_id: &'a str, - task_id: &'a str, - member_id: &'a str, - claim_token: &'a str, - execution: Option<(String, String)>, - armed: bool, -} - -impl<L: TeamLedger> MemberExecutionCleanup<'_, L> { - fn cleanup(&mut self) -> Result<()> { - self.service.release_and_idle( - self.team_id, - self.task_id, - self.member_id, - self.claim_token, - self.execution - .as_ref() - .map(|(thread, run)| (thread.as_str(), run.as_str())), - )?; - self.armed = false; - Ok(()) - } - - fn disarm(&mut self) { - self.armed = false; - } -} - -impl<L: TeamLedger> Drop for MemberExecutionCleanup<'_, L> { - fn drop(&mut self) { - if self.armed { - let _ = self.service.release_and_idle( - self.team_id, - self.task_id, - self.member_id, - self.claim_token, - self.execution - .as_ref() - .map(|(thread, run)| (thread.as_str(), run.as_str())), - ); - } - } -} - -impl<L> TeamService<L> { - /// Creates a service wrapping the provided [`TeamLedger`]. - pub fn new(ledger: L) -> Self { - Self { ledger } - } - - /// Returns a reference to the underlying ledger. - pub fn ledger(&self) -> &L { - &self.ledger - } -} - -impl<L: TeamLedger> TeamService<L> { - pub fn create_team( - &self, - lead_agent_id: &str, - parent_thread_id: Option<&str>, - summary: Option<&str>, - members: &[NewMember], - ) -> Result<TeamView> { - let mut seen = HashSet::new(); - for member in members { - if !seen.insert(member.name.as_str()) { - return Err(anyhow!(TeamError::DuplicateMemberName { - name: member.name.clone() - })); - } - } - let team_id = format!("team-{}", Uuid::new_v4().simple()); - self.ledger.upsert_team(AgentTeamUpsert { - id: team_id.clone(), - parent_thread_id: parent_thread_id.map(str::to_string), - lead_agent_id: lead_agent_id.to_string(), - status: AgentTeamStatus::Active, - summary: summary.map(str::to_string), - created_at: None, - closed_at: None, - })?; - for member in members { - self.ledger.upsert_member(AgentTeamMemberUpsert { - id: format!("member-{}", Uuid::new_v4().simple()), - team_id: team_id.clone(), - name: member.name.clone(), - agent_id: member.agent_id.clone(), - member_status: AgentTeamMemberStatus::Pending, - current_task_id: None, - worker_thread_id: None, - run_id: None, - created_at: None, - })?; - } - self.team_view(&team_id) - } - - pub fn list_teams(&self, request: &AgentTeamListRequest) -> Result<AgentTeamListResponse> { - self.ledger.list_teams(request) - } - - pub fn get_team(&self, team_id: &str) -> Result<Option<TeamView>> { - if self.ledger.get_team(team_id)?.is_some() { - self.team_view(team_id).map(Some) - } else { - Ok(None) - } - } - - pub fn assign_task( - &self, - team_id: &str, - title: &str, - objective: Option<&str>, - owner_member_id: Option<&str>, - depends_on: &[String], - ) -> Result<AgentTeamTask> { - let team = self - .ledger - .get_team(team_id)? - .ok_or_else(|| anyhow!("unknown team: {team_id}"))?; - if team.status == AgentTeamStatus::Closed { - return Err(anyhow!("team is closed: {team_id}")); - } - let existing = self.ledger.list_tasks(team_id)?; - if let Some(owner) = owner_member_id - && !self.ledger.list_members(team_id)?.iter().any(|member| { - member.id == owner && member.member_status != AgentTeamMemberStatus::Stopped - }) - { - return Err(anyhow!(TeamError::UnknownMember { - member_id: owner.to_string() - })); - } - let task_id = format!("task-{}", Uuid::new_v4().simple()); - validate_dependencies(&task_id, depends_on, &existing)?; - self.ledger.upsert_task(AgentTeamTaskUpsert { - id: task_id, - team_id: team_id.to_string(), - title: title.to_string(), - objective: objective.map(str::to_string), - status: AgentTeamTaskStatus::Todo, - owner_member_id: owner_member_id.map(str::to_string), - depends_on: depends_on.to_vec(), - gate_status: None, - gate_reason: None, - evidence: vec![], - source_run_id: None, - order_index: existing.len() as i64, - created_at: None, - }) - } - - pub fn claim_task( - &self, - team_id: &str, - task_id: &str, - member_id: &str, - claim_token: &str, - ) -> Result<ClaimOutcome> { - self.ensure_team_active(team_id)?; - self.ensure_member(team_id, member_id)?; - self.ledger - .claim_task(team_id, task_id, member_id, claim_token) - } - - pub fn message_member( - &self, - team_id: &str, - from_member_id: Option<&str>, - to_member_id: Option<&str>, - content: &str, - visibility: Option<&str>, - ) -> Result<RunEvent> { - self.ledger - .get_team(team_id)? - .ok_or_else(|| anyhow!("unknown team: {team_id}"))?; - if let Some(from) = from_member_id { - self.ensure_member(team_id, from)?; - } - if let Some(to) = to_member_id { - self.ensure_member(team_id, to)?; - } - self.ledger.append_event(RunEventAppend { - run_id: team_id.to_string(), - event_type: TEAM_MESSAGE_EVENT.to_string(), - payload: json!({"from": from_member_id.unwrap_or(LEAD_SENDER), "to": to_member_id, - "content": content, "visibility": visibility.unwrap_or("team")}), - }) - } - - pub fn list_messages(&self, team_id: &str, limit: Option<u32>) -> Result<Vec<RunEvent>> { - Ok(drain_run_events(&self.ledger, team_id)? - .into_iter() - .filter(|event| event.event_type == TEAM_MESSAGE_EVENT) - .take(limit.unwrap_or(u32::MAX) as usize) - .collect()) - } - - pub fn complete_task( - &self, - team_id: &str, - task_id: &str, - member_id: &str, - evidence: &[String], - require_evidence: bool, - ) -> Result<CompletionOutcome> { - self.ensure_member(team_id, member_id)?; - self.ledger - .complete_task(team_id, task_id, member_id, evidence, require_evidence) - } - - /// Claim and execute one team task, then durably complete it. - /// - /// Worker failures release the task and return the member to idle before - /// the original error is propagated. Completion gate failures are also - /// released so another attempt can claim the task. - pub async fn start_member<E: TeamMemberExecutor + ?Sized>( - &self, - team_id: &str, - task_id: &str, - member_id: &str, - executor: &E, - require_evidence: bool, - ) -> Result<CompletionOutcome> { - self.start_member_with_started( - team_id, - task_id, - member_id, - executor, - require_evidence, - |_| {}, - ) - .await - } - - /// Claims and activates a member task, then calls `on_started` before the - /// executor is polled. A host can return an immediate response from the - /// callback while this same future continues execution and retains the - /// normal owner-controlled evidence and cleanup path. - pub async fn start_member_with_started<E, Started>( - &self, - team_id: &str, - task_id: &str, - member_id: &str, - executor: &E, - require_evidence: bool, - on_started: Started, - ) -> Result<CompletionOutcome> - where - E: TeamMemberExecutor + ?Sized, - Started: FnOnce(TeamMemberStarted), - { - self.ensure_member(team_id, member_id)?; - let claim_token = Uuid::new_v4().to_string(); - let claimed = match self.claim_task(team_id, task_id, member_id, &claim_token)? { - ClaimOutcome::Claimed(task) => *task, - ClaimOutcome::AlreadyClaimed => { - return Err(anyhow!(TeamMemberStartError::AlreadyClaimed { - task_id: task_id.to_string(), - })); - } - ClaimOutcome::Blocked { unmet } => { - return Err(anyhow!(TeamMemberStartError::Blocked { - task_id: task_id.to_string(), - unmet, - })); - } - ClaimOutcome::UnknownTask => { - return Err(anyhow!(TeamMemberStartError::UnknownTask { - task_id: task_id.to_string(), - })); - } - }; - let mut cleanup = MemberExecutionCleanup { - service: self, - team_id, - task_id, - member_id, - claim_token: &claim_token, - execution: None, - armed: true, - }; - let worker_thread_id = format!("team-worker-{}", Uuid::new_v4().simple()); - let run_id = format!("team-run-{}", Uuid::new_v4().simple()); - let member = match self.ledger.try_mark_member_running( - team_id, - member_id, - task_id, - &worker_thread_id, - &run_id, - ) { - Ok(AgentTeamMemberRunAdmission::Started(member)) => member, - Ok(AgentTeamMemberRunAdmission::AlreadyActive(member)) => { - return Err(anyhow!(TeamMemberStartError::AlreadyActive { - member_id: member_id.to_owned(), - current_task_id: member.current_task_id, - })); - } - Ok( - AgentTeamMemberRunAdmission::NotAdmissible(_) - | AgentTeamMemberRunAdmission::Missing, - ) => { - return Err(anyhow!(TeamError::UnknownMember { - member_id: member_id.to_owned(), - })); - } - Err(error) => { - return Err(error.context("admit team member run")); - } - }; - cleanup.execution = Some(( - member - .worker_thread_id - .clone() - .unwrap_or_else(|| worker_thread_id.clone()), - member.run_id.clone().unwrap_or_else(|| run_id.clone()), - )); - - on_started(TeamMemberStarted { - run_id: member.run_id.clone().unwrap_or_else(|| run_id.clone()), - member: member.clone(), - task: claimed.clone(), - }); - - let execution = executor.execute(&member, &claimed).await; - let completion = match execution { - Ok(evidence) => self.ledger.complete_task_if_claimed( - team_id, - task_id, - member_id, - &claim_token, - &evidence, - require_evidence, - ), - Err(error) => { - return match cleanup.cleanup() { - Ok(()) => Err(error.context("team member execution failed")), - Err(cleanup_error) => Err(anyhow!( - "team member execution failed: {error}; cleanup failed: {cleanup_error}" - )), - }; - } - }; - match completion { - Ok(outcome @ CompletionOutcome::Completed(_)) => { - self.ledger.mark_member_idle( - team_id, - member_id, - task_id, - member - .worker_thread_id - .as_deref() - .unwrap_or(&worker_thread_id), - member.run_id.as_deref().unwrap_or(&run_id), - )?; - cleanup.disarm(); - Ok(outcome) - } - Ok(outcome) => { - cleanup.cleanup()?; - Ok(outcome) - } - Err(error) => match cleanup.cleanup() { - Ok(()) => Err(error.context("complete team task")), - Err(cleanup_error) => Err(anyhow!( - "complete team task failed: {error}; cleanup failed: {cleanup_error}" - )), - }, - } - } - - fn release_and_idle( - &self, - team_id: &str, - task_id: &str, - member_id: &str, - claim_token: &str, - execution: Option<(&str, &str)>, - ) -> Result<()> { - let release = self - .ledger - .release_task(team_id, task_id, member_id, claim_token); - let idle = execution.map_or(Ok(false), |(thread, run)| { - self.ledger - .mark_member_idle(team_id, member_id, task_id, thread, run) - }); - match (release, idle) { - (Ok(_), Ok(_)) => Ok(()), - (Err(error), Ok(_)) | (Ok(_), Err(error)) => Err(error), - (Err(release), Err(idle)) => Err(anyhow!( - "failed to release team task: {release}; failed to idle team member: {idle}" - )), - } - } - - pub fn shutdown_member(&self, team_id: &str, member_id: &str) -> Result<MemberShutdown> { - self.ledger - .shutdown_member(team_id, member_id)? - .map(|(member, released_task_ids)| MemberShutdown { - member, - released_task_ids, - }) - .ok_or_else(|| { - anyhow!(TeamError::UnknownMember { - member_id: member_id.to_string() - }) - }) - } - - pub fn close_team(&self, team_id: &str, summary: Option<&str>) -> Result<AgentTeam> { - let existing = self - .ledger - .get_team(team_id)? - .ok_or_else(|| anyhow!("unknown team: {team_id}"))?; - self.ledger.upsert_team(AgentTeamUpsert { - id: team_id.to_string(), - parent_thread_id: existing.parent_thread_id, - lead_agent_id: existing.lead_agent_id, - status: AgentTeamStatus::Closed, - summary: summary.map(str::to_string), - created_at: Some(existing.created_at), - closed_at: Some(Utc::now()), - }) - } - - fn team_view(&self, team_id: &str) -> Result<TeamView> { - let team = self - .ledger - .get_team(team_id)? - .ok_or_else(|| anyhow!("team missing after creation: {team_id}"))?; - Ok(TeamView { - team, - members: self.ledger.list_members(team_id)?, - tasks: self.ledger.list_tasks(team_id)?, - }) - } - - fn ensure_member(&self, team_id: &str, member_id: &str) -> Result<()> { - if self.ledger.list_members(team_id)?.iter().any(|member| { - member.id == member_id && member.member_status != AgentTeamMemberStatus::Stopped - }) { - Ok(()) - } else { - Err(anyhow!(TeamError::UnknownMember { - member_id: member_id.to_string() - })) - } - } - - fn ensure_team_active(&self, team_id: &str) -> Result<()> { - let team = self - .ledger - .get_team(team_id)? - .ok_or_else(|| anyhow!("unknown team: {team_id}"))?; - if team.status == AgentTeamStatus::Active { - Ok(()) - } else { - Err(anyhow!("team is closed: {team_id}")) - } - } -} - -#[cfg(test)] -#[path = "service_tests.rs"] -mod tests; - -fn validate_dependencies( - new_task_id: &str, - depends_on: &[String], - existing: &[AgentTeamTask], -) -> Result<()> { - let known: HashSet<&str> = existing.iter().map(|task| task.id.as_str()).collect(); - for dependency in depends_on { - if dependency == new_task_id { - return Err(anyhow!(TeamError::SelfDependency { - task_id: new_task_id.to_string() - })); - } - if !known.contains(dependency.as_str()) { - return Err(anyhow!(TeamError::UnknownDependency { - depends_on: dependency.clone() - })); - } - } - if has_task_cycle(new_task_id, depends_on, existing) { - return Err(anyhow!(TeamError::CyclicDependency)); - } - Ok(()) -} - -fn has_task_cycle(new_task_id: &str, depends_on: &[String], existing: &[AgentTeamTask]) -> bool { - let mut nodes: Vec<DagNode<'_>> = existing - .iter() - .map(|task| DagNode::new(task.id.as_str(), task.depends_on.iter().map(String::as_str))) - .collect(); - nodes.push(DagNode::new( - new_task_id, - depends_on.iter().map(String::as_str), - )); - has_cycle(&nodes) -} - -/// Select the next task a member may claim without making any policy decision. -pub fn claimable_task<'a>( - tasks: &'a [AgentTeamTask], - member_id: &str, -) -> Option<&'a AgentTeamTask> { - let done: HashSet<&str> = tasks - .iter() - .filter(|task| task.status == AgentTeamTaskStatus::Done) - .map(|task| task.id.as_str()) - .collect(); - tasks.iter().find(|task| { - matches!( - task.status, - AgentTeamTaskStatus::Todo | AgentTeamTaskStatus::Ready - ) && task.claimed_by_member_id.is_none() - && task - .owner_member_id - .as_deref() - .map(|owner| owner == member_id) - .unwrap_or(true) - && task - .depends_on - .iter() - .all(|dependency| done.contains(dependency.as_str())) - }) -} - -#[cfg(test)] -#[path = "service_dependency_tests.rs"] -mod dependency_tests; diff --git a/crates/tinyagents-orchestration/src/teams/service_dependency_tests.rs b/crates/tinyagents-orchestration/src/teams/service_dependency_tests.rs deleted file mode 100644 index 4b13b5a87..000000000 --- a/crates/tinyagents-orchestration/src/teams/service_dependency_tests.rs +++ /dev/null @@ -1,50 +0,0 @@ -use chrono::Utc; -use tinyagents_session::run_ledger::AgentTeamTaskStatus; - -use super::*; - -fn task(id: &str, depends_on: &[&str]) -> AgentTeamTask { - let now = Utc::now(); - AgentTeamTask { - id: id.to_string(), - team_id: "team".to_string(), - title: id.to_string(), - objective: None, - status: AgentTeamTaskStatus::Todo, - owner_member_id: None, - claimed_by_member_id: None, - claim_token: None, - depends_on: depends_on - .iter() - .map(|dependency| (*dependency).to_string()) - .collect(), - gate_status: "pending".to_string(), - gate_reason: None, - evidence: Vec::new(), - source_run_id: None, - order_index: 0, - created_at: now, - updated_at: now, - } -} - -#[test] -fn rejects_self_dependency() { - let error = validate_dependencies("task-self", &["task-self".to_string()], &[]).unwrap_err(); - assert_eq!( - error.downcast::<TeamError>().unwrap(), - TeamError::SelfDependency { - task_id: "task-self".to_string() - } - ); -} - -#[test] -fn rejects_dependency_cycle() { - let existing = vec![task("task-a", &["task-new"])]; - let error = validate_dependencies("task-new", &["task-a".to_string()], &existing).unwrap_err(); - assert_eq!( - error.downcast::<TeamError>().unwrap(), - TeamError::CyclicDependency - ); -} diff --git a/crates/tinyagents-orchestration/src/teams/service_tests.rs b/crates/tinyagents-orchestration/src/teams/service_tests.rs deleted file mode 100644 index f7b956d21..000000000 --- a/crates/tinyagents-orchestration/src/teams/service_tests.rs +++ /dev/null @@ -1,602 +0,0 @@ -use super::*; - -use crate::teams::NewMember; -use tempfile::TempDir; -use tinyagents_session::run_ledger::{AgentTeamMemberStatus, AgentTeamTaskStatus}; -use tinystoragedrivers_testkit::MemoryBackend; - -fn service(_dir: &TempDir) -> TeamService<SessionTeamLedger> { - let ledger_backend = MemoryBackend::new(); - service_with_backend(&ledger_backend) -} - -fn service_with_backend(ledger_backend: &MemoryBackend) -> TeamService<SessionTeamLedger> { - TeamService::new( - SessionTeamLedger::new(Arc::new(crate::test_support::ledger_in(ledger_backend))).unwrap(), - ) -} - -struct FixedExecutor(Result<Vec<String>, &'static str>); - -#[async_trait] -impl TeamMemberExecutor for FixedExecutor { - async fn execute(&self, _: &AgentTeamMember, _: &AgentTeamTask) -> Result<Vec<String>> { - self.0.clone().map_err(anyhow::Error::msg) - } -} - -struct RecordingExecutor(Arc<std::sync::Mutex<Option<AgentTeamMember>>>); - -#[async_trait] -impl TeamMemberExecutor for RecordingExecutor { - async fn execute(&self, member: &AgentTeamMember, _: &AgentTeamTask) -> Result<Vec<String>> { - *self.0.lock().unwrap() = Some(member.clone()); - Ok(vec!["done".into()]) - } -} - -struct HangingExecutor(Arc<tokio::sync::Notify>); - -#[async_trait] -impl TeamMemberExecutor for HangingExecutor { - async fn execute(&self, _: &AgentTeamMember, _: &AgentTeamTask) -> Result<Vec<String>> { - self.0.notify_one(); - std::future::pending().await - } -} - -struct CountingExecutor(Arc<std::sync::atomic::AtomicUsize>); - -#[async_trait] -impl TeamMemberExecutor for CountingExecutor { - async fn execute(&self, _: &AgentTeamMember, _: &AgentTeamTask) -> Result<Vec<String>> { - self.0.fetch_add(1, std::sync::atomic::Ordering::SeqCst); - Ok(vec!["unexpected second execution".into()]) - } -} - -struct PausedExecutor { - entered: Arc<tokio::sync::Notify>, - resume: Arc<tokio::sync::Notify>, -} - -#[async_trait] -impl TeamMemberExecutor for PausedExecutor { - async fn execute(&self, _: &AgentTeamMember, _: &AgentTeamTask) -> Result<Vec<String>> { - self.entered.notify_one(); - self.resume.notified().await; - Ok(vec!["stale completion".into()]) - } -} - -#[tokio::test] -async fn start_member_runs_and_completes_claimed_task() { - let dir = TempDir::new().unwrap(); - let service = service(&dir); - let team = service - .create_team( - "lead", - None, - None, - &[NewMember { - name: "alice".into(), - agent_id: None, - }], - ) - .unwrap(); - let member = &team.members[0]; - let task = service - .assign_task(&team.team.id, "write", None, Some(&member.id), &[]) - .unwrap(); - - let executed_member = Arc::new(std::sync::Mutex::new(None)); - let outcome = service - .start_member( - &team.team.id, - &task.id, - &member.id, - &RecordingExecutor(executed_member.clone()), - true, - ) - .await - .unwrap(); - - assert!(matches!(outcome, CompletionOutcome::Completed(_))); - let executed_member = executed_member.lock().unwrap().clone().unwrap(); - assert_eq!(executed_member.member_status, AgentTeamMemberStatus::Active); - assert_eq!( - executed_member.current_task_id.as_deref(), - Some(task.id.as_str()) - ); - assert!(executed_member.worker_thread_id.is_some()); - assert!(executed_member.run_id.is_some()); - assert_eq!( - service.ledger().list_tasks(&team.team.id).unwrap()[0].status, - AgentTeamTaskStatus::Done - ); - assert_eq!( - service.ledger().list_members(&team.team.id).unwrap()[0].member_status, - AgentTeamMemberStatus::Idle - ); -} - -#[tokio::test] -async fn a_member_cannot_start_a_second_task_while_its_first_task_is_running() { - // Independent handles prove the admission fence is in the shared backend, - // rather than a mutex local to one TeamService instance. - let ledger_backend = MemoryBackend::new(); - let first_service = service_with_backend(&ledger_backend); - let second_service = service_with_backend(&ledger_backend); - let team = first_service - .create_team( - "lead", - None, - None, - &[NewMember { - name: "alice".into(), - agent_id: None, - }], - ) - .unwrap(); - let member = &team.members[0]; - let first_task = first_service - .assign_task(&team.team.id, "first", None, Some(&member.id), &[]) - .unwrap(); - let second_task = first_service - .assign_task(&team.team.id, "second", None, Some(&member.id), &[]) - .unwrap(); - - let first_entered = Arc::new(tokio::sync::Notify::new()); - let first_entered_for_worker = Arc::clone(&first_entered); - let first_worker_service = first_service.clone(); - let first_team_id = team.team.id.clone(); - let first_task_id = first_task.id.clone(); - let member_id = member.id.clone(); - let first_worker = tokio::spawn(async move { - first_worker_service - .start_member( - &first_team_id, - &first_task_id, - &member_id, - &HangingExecutor(first_entered_for_worker), - false, - ) - .await - }); - first_entered.notified().await; - - let executions = Arc::new(std::sync::atomic::AtomicUsize::new(0)); - let mut loser_callback_called = false; - let second = second_service - .start_member_with_started( - &team.team.id, - &second_task.id, - &member.id, - &CountingExecutor(Arc::clone(&executions)), - false, - |_| loser_callback_called = true, - ) - .await; - - let error = second.expect_err("an active member must reject a second task"); - assert!(matches!( - error.downcast_ref::<TeamMemberStartError>(), - Some(TeamMemberStartError::AlreadyActive { member_id, current_task_id }) - if member_id == &member.id && current_task_id.as_deref() == Some(first_task.id.as_str()) - )); - assert!( - !loser_callback_called, - "denied admission must not report Started" - ); - assert_eq!(executions.load(std::sync::atomic::Ordering::SeqCst), 0); - let after = first_service.get_team(&team.team.id).unwrap().unwrap(); - assert_eq!( - after.members[0].member_status, - AgentTeamMemberStatus::Active - ); - assert_eq!( - after.members[0].current_task_id.as_deref(), - Some(first_task.id.as_str()) - ); - assert!( - after.members[0] - .worker_thread_id - .as_deref() - .is_some_and(|id| id.starts_with("team-worker-")) - ); - assert!( - after.members[0] - .run_id - .as_deref() - .is_some_and(|id| id.starts_with("team-run-")) - ); - let first_after = after - .tasks - .iter() - .find(|task| task.id == first_task.id) - .unwrap(); - let second_after = after - .tasks - .iter() - .find(|task| task.id == second_task.id) - .unwrap(); - assert_eq!(first_after.status, AgentTeamTaskStatus::InProgress); - assert!(first_after.claim_token.is_some()); - assert_eq!(second_after.status, AgentTeamTaskStatus::Todo); - assert!(second_after.claim_token.is_none()); - - first_worker.abort(); - assert!(first_worker.await.unwrap_err().is_cancelled()); -} - -#[tokio::test] -async fn started_callback_fires_after_claim_before_executor_finishes() { - let dir = TempDir::new().unwrap(); - let service = service(&dir); - let team = service - .create_team( - "lead", - None, - None, - &[NewMember { - name: "alice".into(), - agent_id: None, - }], - ) - .unwrap(); - let member = &team.members[0]; - let task = service - .assign_task(&team.team.id, "write", None, Some(&member.id), &[]) - .unwrap(); - let (started_tx, started_rx) = tokio::sync::oneshot::channel(); - let entered = Arc::new(tokio::sync::Notify::new()); - let executor = HangingExecutor(entered.clone()); - let worker_service = service.clone(); - let team_id = team.team.id.clone(); - let task_id = task.id.clone(); - let member_id = member.id.clone(); - let worker = tokio::spawn(async move { - worker_service - .start_member_with_started( - &team_id, - &task_id, - &member_id, - &executor, - false, - |started| { - let _ = started_tx.send(started); - }, - ) - .await - }); - - let started = started_rx.await.unwrap(); - assert!(!started.run_id.is_empty()); - assert_eq!(started.task.id, task.id); - assert_eq!(started.member.member_status, AgentTeamMemberStatus::Active); - assert_eq!( - started.member.run_id.as_deref(), - Some(started.run_id.as_str()) - ); - assert!( - started - .member - .worker_thread_id - .as_deref() - .is_some_and(|id| id.starts_with("team-worker-")) - ); - let durable = service.get_team(&team.team.id).unwrap().unwrap(); - assert_eq!( - durable.members[0].run_id.as_deref(), - Some(started.run_id.as_str()) - ); - assert_eq!(durable.tasks[0].status, AgentTeamTaskStatus::InProgress); - entered.notified().await; - assert!(!worker.is_finished()); - worker.abort(); - let _ = worker.await; - let released = service.get_team(&team.team.id).unwrap().unwrap(); - assert_eq!( - released.members[0].member_status, - AgentTeamMemberStatus::Idle - ); - assert_eq!(released.tasks[0].status, AgentTeamTaskStatus::Todo); -} - -#[tokio::test] -async fn start_member_exposes_claim_conflicts_as_typed_errors() { - let dir = TempDir::new().unwrap(); - let service = service(&dir); - let team = service - .create_team( - "lead", - None, - None, - &[NewMember { - name: "alice".into(), - agent_id: None, - }], - ) - .unwrap(); - let member = &team.members[0]; - let task = service - .assign_task(&team.team.id, "write", None, Some(&member.id), &[]) - .unwrap(); - service - .claim_task(&team.team.id, &task.id, &member.id, "existing-claim") - .unwrap(); - - let error = service - .start_member_with_started( - &team.team.id, - &task.id, - &member.id, - &FixedExecutor(Ok(vec![])), - false, - |_| panic!("claim conflict must not notify the host as started"), - ) - .await - .unwrap_err(); - assert!(matches!( - error.downcast_ref::<TeamMemberStartError>(), - Some(TeamMemberStartError::AlreadyClaimed { task_id }) if task_id == &task.id - )); -} - -#[tokio::test] -async fn started_callback_unwind_releases_its_claim_and_idles_the_member() { - let dir = TempDir::new().unwrap(); - let service = service(&dir); - let team = service - .create_team( - "lead", - None, - None, - &[NewMember { - name: "alice".into(), - agent_id: None, - }], - ) - .unwrap(); - let member = &team.members[0]; - let task = service - .assign_task(&team.team.id, "write", None, Some(&member.id), &[]) - .unwrap(); - let worker_service = service.clone(); - let team_id = team.team.id.clone(); - let task_id = task.id.clone(); - let member_id = member.id.clone(); - let worker = tokio::spawn(async move { - worker_service - .start_member_with_started( - &team_id, - &task_id, - &member_id, - &FixedExecutor(Ok(vec![])), - false, - |_| panic!("observer panic"), - ) - .await - }); - assert!(worker.await.unwrap_err().is_panic()); - - let after = service.get_team(&team.team.id).unwrap().unwrap(); - assert_eq!(after.members[0].member_status, AgentTeamMemberStatus::Idle); - assert_eq!(after.tasks[0].status, AgentTeamTaskStatus::Todo); -} - -#[tokio::test] -async fn start_member_releases_task_and_idles_member_on_executor_failure() { - let dir = TempDir::new().unwrap(); - let service = service(&dir); - let team = service - .create_team( - "lead", - None, - None, - &[NewMember { - name: "alice".into(), - agent_id: None, - }], - ) - .unwrap(); - let member = &team.members[0]; - let task = service - .assign_task(&team.team.id, "write", None, Some(&member.id), &[]) - .unwrap(); - - let error = service - .start_member( - &team.team.id, - &task.id, - &member.id, - &FixedExecutor(Err("worker stopped")), - true, - ) - .await - .unwrap_err(); - - assert!(format!("{error:#}").contains("worker stopped")); - assert_eq!( - service.ledger().list_tasks(&team.team.id).unwrap()[0].status, - AgentTeamTaskStatus::Todo - ); - assert_eq!( - service.ledger().list_members(&team.team.id).unwrap()[0].member_status, - AgentTeamMemberStatus::Idle - ); -} - -#[tokio::test] -async fn cancelling_start_member_releases_task_and_idles_member() { - let dir = TempDir::new().unwrap(); - let service = service(&dir); - let team = service - .create_team( - "lead", - None, - None, - &[NewMember { - name: "alice".into(), - agent_id: None, - }], - ) - .unwrap(); - let member = &team.members[0]; - let task = service - .assign_task(&team.team.id, "write", None, Some(&member.id), &[]) - .unwrap(); - let entered = Arc::new(tokio::sync::Notify::new()); - let executor = Arc::new(HangingExecutor(entered.clone())); - let worker_service = service.clone(); - let team_id = team.team.id.clone(); - let task_id = task.id.clone(); - let member_id = member.id.clone(); - let worker = tokio::spawn(async move { - worker_service - .start_member(&team_id, &task_id, &member_id, executor.as_ref(), true) - .await - }); - entered.notified().await; - worker.abort(); - let _ = worker.await; - - assert_eq!( - service.ledger().list_tasks(&team.team.id).unwrap()[0].status, - AgentTeamTaskStatus::Todo - ); - assert_eq!( - service.ledger().list_members(&team.team.id).unwrap()[0].member_status, - AgentTeamMemberStatus::Idle - ); -} - -#[tokio::test] -async fn stale_executor_cannot_complete_successor_claim_for_same_member() { - let dir = TempDir::new().unwrap(); - let service = service(&dir); - let team = service - .create_team( - "lead", - None, - None, - &[NewMember { - name: "alice".into(), - agent_id: None, - }], - ) - .unwrap(); - let member = &team.members[0]; - let task = service - .assign_task(&team.team.id, "write", None, Some(&member.id), &[]) - .unwrap(); - let entered = Arc::new(tokio::sync::Notify::new()); - let resume = Arc::new(tokio::sync::Notify::new()); - let executor = Arc::new(PausedExecutor { - entered: entered.clone(), - resume: resume.clone(), - }); - let worker_service = service.clone(); - let (team_id, task_id, member_id) = (team.team.id.clone(), task.id.clone(), member.id.clone()); - let worker = tokio::spawn(async move { - worker_service - .start_member(&team_id, &task_id, &member_id, executor.as_ref(), false) - .await - }); - entered.notified().await; - - service.shutdown_member(&team.team.id, &member.id).unwrap(); - let ClaimOutcome::Claimed(_) = service - .ledger() - .claim_task(&team.team.id, &task.id, &member.id, "successor-token") - .unwrap() - else { - panic!("successor should reclaim the released task"); - }; - service - .ledger() - .mark_member_running( - &team.team.id, - &member.id, - &task.id, - "successor-thread", - "successor-run", - ) - .unwrap() - .unwrap(); - - resume.notify_one(); - assert!(matches!( - worker.await.unwrap().unwrap(), - CompletionOutcome::NotClaimed - )); - let current_task = service - .ledger() - .list_tasks(&team.team.id) - .unwrap() - .remove(0); - assert_eq!(current_task.status, AgentTeamTaskStatus::InProgress); - assert_eq!(current_task.claim_token.as_deref(), Some("successor-token")); - assert!(current_task.evidence.is_empty()); - let current_member = service - .ledger() - .list_members(&team.team.id) - .unwrap() - .remove(0); - assert_eq!(current_member.member_status, AgentTeamMemberStatus::Active); - assert_eq!(current_member.run_id.as_deref(), Some("successor-run")); -} - -#[test] -fn stale_claim_cleanup_cannot_release_a_successor_claim() { - let dir = TempDir::new().unwrap(); - let service = service(&dir); - let team = service - .create_team( - "lead", - None, - None, - &[NewMember { - name: "alice".into(), - agent_id: None, - }], - ) - .unwrap(); - let member = &team.members[0]; - let task = service - .assign_task(&team.team.id, "write", None, Some(&member.id), &[]) - .unwrap(); - let ClaimOutcome::Claimed(first_claim) = service - .claim_task(&team.team.id, &task.id, &member.id, "first-generation") - .unwrap() - else { - panic!("first generation should claim the task"); - }; - assert!( - service - .ledger() - .release_task(&team.team.id, &task.id, &member.id, "first-generation") - .unwrap() - ); - let ClaimOutcome::Claimed(second_claim) = service - .claim_task(&team.team.id, &task.id, &member.id, "successor-generation") - .unwrap() - else { - panic!("successor should claim the released task"); - }; - - assert!( - !service - .ledger() - .release_task(&team.team.id, &task.id, &member.id, "first-generation") - .unwrap() - ); - let current = service - .ledger() - .list_tasks(&team.team.id) - .unwrap() - .remove(0); - assert_eq!(current.claim_token.as_deref(), Some("successor-generation")); - assert_eq!(current.claim_token, second_claim.claim_token); - assert_ne!(first_claim.claim_token, second_claim.claim_token); -} diff --git a/crates/tinyagents-orchestration/src/teams/tests.rs b/crates/tinyagents-orchestration/src/teams/tests.rs deleted file mode 100644 index 15338c71f..000000000 --- a/crates/tinyagents-orchestration/src/teams/tests.rs +++ /dev/null @@ -1,428 +0,0 @@ -//! Tests for team service: creation, member lifecycle, task coordination, -//! and messaging flow. - -use std::sync::{Arc, Barrier, Mutex}; -use std::thread; - -use chrono::Utc; -use tempfile::TempDir; -use tinyagents_session::run_ledger::{ - AgentTeam, AgentTeamListRequest, AgentTeamListResponse, AgentTeamMember, AgentTeamMemberStatus, - AgentTeamMemberUpsert, AgentTeamStatus, AgentTeamTask, AgentTeamTaskUpsert, AgentTeamUpsert, - ClaimOutcome, CompletionOutcome, RunEvent, RunEventAppend, RunEventListRequest, -}; -use tinystoragedrivers_testkit::MemoryBackend; - -use super::*; - -fn service(_dir: &TempDir) -> TeamService<SessionTeamLedger> { - let ledger_backend = MemoryBackend::new(); - TeamService::new( - SessionTeamLedger::new(Arc::new(crate::test_support::ledger_in(&ledger_backend))).unwrap(), - ) -} - -fn solo_team(service: &TeamService<SessionTeamLedger>) -> (String, String) { - let view = service - .create_team( - "lead", - None, - None, - &[NewMember { - name: "alice".into(), - agent_id: None, - }], - ) - .unwrap(); - (view.team.id, view.members[0].id.clone()) -} - -fn team_error(error: anyhow::Error) -> TeamError { - error.downcast::<TeamError>().unwrap() -} - -#[test] -fn rejects_duplicate_members_and_unknown_dependencies() { - let dir = TempDir::new().unwrap(); - let service = service(&dir); - let error = service - .create_team( - "lead", - None, - None, - &[ - NewMember { - name: "alice".into(), - agent_id: None, - }, - NewMember { - name: "alice".into(), - agent_id: None, - }, - ], - ) - .unwrap_err(); - assert_eq!( - team_error(error), - TeamError::DuplicateMemberName { - name: "alice".into() - } - ); - - let (team_id, _) = solo_team(&service); - let error = service - .assign_task(&team_id, "task", None, None, &["missing".into()]) - .unwrap_err(); - assert_eq!( - team_error(error), - TeamError::UnknownDependency { - depends_on: "missing".into() - } - ); -} - -#[test] -fn task_claim_completion_and_quality_gate_are_durable() { - let dir = TempDir::new().unwrap(); - let service = service(&dir); - let (team_id, member_id) = solo_team(&service); - let task = service - .assign_task(&team_id, "ship", None, None, &[]) - .unwrap(); - assert!(matches!( - service - .claim_task(&team_id, &task.id, &member_id, "claim-1") - .unwrap(), - ClaimOutcome::Claimed(_) - )); - assert!(matches!( - service - .complete_task(&team_id, &task.id, &member_id, &[], true) - .unwrap(), - CompletionOutcome::GateFailed { .. } - )); - let after_failed_gate = service.get_team(&team_id).unwrap().unwrap(); - let durable_task = after_failed_gate - .tasks - .iter() - .find(|candidate| candidate.id == task.id) - .unwrap(); - assert_eq!( - durable_task.status, - tinyagents_session::run_ledger::AgentTeamTaskStatus::InProgress - ); - assert_eq!(durable_task.gate_status, "failed"); - assert!(matches!( - service - .complete_task(&team_id, &task.id, &member_id, &["proof".into()], true) - .unwrap(), - CompletionOutcome::Completed(_) - )); -} - -#[test] -fn racing_claims_have_one_winner_and_one_already_claimed_loser() { - let dir = TempDir::new().unwrap(); - let service = service(&dir); - let view = service - .create_team( - "lead", - None, - None, - &[ - NewMember { - name: "alice".into(), - agent_id: None, - }, - NewMember { - name: "bob".into(), - agent_id: None, - }, - ], - ) - .unwrap(); - let task = service - .assign_task(&view.team.id, "race", None, None, &[]) - .unwrap(); - let barrier = Arc::new(Barrier::new(2)); - let mut handles = Vec::new(); - for (member, token) in [ - (view.members[0].id.clone(), "alice-token"), - (view.members[1].id.clone(), "bob-token"), - ] { - let service = service.clone(); - let team_id = view.team.id.clone(); - let task_id = task.id.clone(); - let barrier = barrier.clone(); - handles.push(thread::spawn(move || { - barrier.wait(); - service - .claim_task(&team_id, &task_id, &member, token) - .unwrap() - })); - } - let outcomes: Vec<_> = handles - .into_iter() - .map(|handle| handle.join().unwrap()) - .collect(); - assert_eq!( - outcomes - .iter() - .filter(|outcome| matches!(outcome, ClaimOutcome::Claimed(_))) - .count(), - 1 - ); - assert_eq!( - outcomes - .iter() - .filter(|outcome| matches!(outcome, ClaimOutcome::AlreadyClaimed)) - .count(), - 1 - ); -} - -#[test] -fn completion_rejects_non_claimants_and_owner_mismatches() { - let dir = TempDir::new().unwrap(); - let service = service(&dir); - let view = service - .create_team( - "lead", - None, - None, - &[ - NewMember { - name: "alice".into(), - agent_id: None, - }, - NewMember { - name: "bob".into(), - agent_id: None, - }, - ], - ) - .unwrap(); - let alice = &view.members[0].id; - let bob = &view.members[1].id; - - let claimed_by_alice = service - .assign_task(&view.team.id, "alice work", None, None, &[]) - .unwrap(); - service - .claim_task(&view.team.id, &claimed_by_alice.id, alice, "alice-token") - .unwrap(); - assert!(matches!( - service - .complete_task( - &view.team.id, - &claimed_by_alice.id, - bob, - &["proof".into()], - false - ) - .unwrap(), - CompletionOutcome::NotClaimed - )); - - let owned_by_alice = service - .assign_task(&view.team.id, "owned work", None, Some(alice), &[]) - .unwrap(); - service - .claim_task(&view.team.id, &owned_by_alice.id, bob, "bob-token") - .unwrap(); - let outcome = service - .complete_task( - &view.team.id, - &owned_by_alice.id, - bob, - &["proof".into()], - false, - ) - .unwrap(); - assert!(matches!( - outcome, - CompletionOutcome::GateFailed { ref reasons } - if reasons.iter().any(|reason| reason.contains("owned by")) - )); -} - -#[test] -fn messages_remain_ordered_and_member_shutdown_releases_work() { - let dir = TempDir::new().unwrap(); - let service = service(&dir); - let view = service - .create_team( - "lead", - None, - None, - &[ - NewMember { - name: "alice".into(), - agent_id: None, - }, - NewMember { - name: "bob".into(), - agent_id: None, - }, - ], - ) - .unwrap(); - let team_id = view.team.id; - let alice = view.members[0].id.clone(); - let bob = view.members[1].id.clone(); - service - .message_member(&team_id, Some(&alice), Some(&bob), "first", None) - .unwrap(); - service - .message_member(&team_id, None, Some(&alice), "second", None) - .unwrap(); - let messages = service.list_messages(&team_id, None).unwrap(); - assert_eq!( - messages - .iter() - .map(|event| event.payload["content"].as_str()) - .collect::<Vec<_>>(), - vec![Some("first"), Some("second")] - ); - assert_eq!(messages[1].payload["from"], LEAD_SENDER); - - let task = service - .assign_task(&team_id, "ship", None, None, &[]) - .unwrap(); - service - .claim_task(&team_id, &task.id, &alice, "claim-1") - .unwrap(); - let shutdown = service.shutdown_member(&team_id, &alice).unwrap(); - assert_eq!(shutdown.released_task_ids, vec![task.id]); - assert_eq!( - shutdown.member.member_status, - AgentTeamMemberStatus::Stopped - ); -} - -#[test] -fn fake_ledger_exercises_member_validation_without_session_storage() { - let ledger = FakeLedger::default(); - let service = TeamService::new(ledger.clone()); - ledger.teams.lock().unwrap().push(team("team-1")); - let error = service - .claim_task("team-1", "task-1", "unknown", "token") - .unwrap_err(); - assert_eq!( - team_error(error), - TeamError::UnknownMember { - member_id: "unknown".into() - } - ); -} - -fn team(id: &str) -> AgentTeam { - AgentTeam { - id: id.into(), - parent_thread_id: None, - lead_agent_id: "lead".into(), - status: AgentTeamStatus::Active, - summary: None, - created_at: Utc::now(), - updated_at: Utc::now(), - closed_at: None, - } -} - -#[derive(Clone, Default)] -struct FakeLedger { - teams: Arc<Mutex<Vec<AgentTeam>>>, -} - -impl TeamLedger for FakeLedger { - fn upsert_team(&self, upsert: AgentTeamUpsert) -> anyhow::Result<AgentTeam> { - Ok(team(&upsert.id)) - } - fn get_team(&self, id: &str) -> anyhow::Result<Option<AgentTeam>> { - Ok(self - .teams - .lock() - .unwrap() - .iter() - .find(|team| team.id == id) - .cloned()) - } - fn list_teams(&self, _: &AgentTeamListRequest) -> anyhow::Result<AgentTeamListResponse> { - Ok(AgentTeamListResponse { - teams: self.teams.lock().unwrap().clone(), - count: self.teams.lock().unwrap().len(), - }) - } - fn upsert_member(&self, _: AgentTeamMemberUpsert) -> anyhow::Result<AgentTeamMember> { - unreachable!() - } - fn list_members(&self, _: &str) -> anyhow::Result<Vec<AgentTeamMember>> { - Ok(vec![]) - } - fn list_tasks(&self, _: &str) -> anyhow::Result<Vec<AgentTeamTask>> { - Ok(vec![]) - } - fn upsert_task(&self, _: AgentTeamTaskUpsert) -> anyhow::Result<AgentTeamTask> { - unreachable!() - } - fn claim_task(&self, _: &str, _: &str, _: &str, _: &str) -> anyhow::Result<ClaimOutcome> { - unreachable!() - } - fn complete_task( - &self, - _: &str, - _: &str, - _: &str, - _: &[String], - _: bool, - ) -> anyhow::Result<CompletionOutcome> { - unreachable!() - } - fn complete_task_if_claimed( - &self, - _: &str, - _: &str, - _: &str, - _: &str, - _: &[String], - _: bool, - ) -> anyhow::Result<CompletionOutcome> { - unreachable!() - } - fn release_task(&self, _: &str, _: &str, _: &str, _: &str) -> anyhow::Result<bool> { - Ok(false) - } - fn mark_member_running( - &self, - _: &str, - _: &str, - _: &str, - _: &str, - _: &str, - ) -> anyhow::Result<Option<AgentTeamMember>> { - Ok(None) - } - fn mark_member_idle( - &self, - _: &str, - _: &str, - _: &str, - _: &str, - _: &str, - ) -> anyhow::Result<bool> { - Ok(false) - } - fn shutdown_member( - &self, - _: &str, - _: &str, - ) -> anyhow::Result<Option<(AgentTeamMember, Vec<String>)>> { - Ok(None) - } - fn append_event(&self, _: RunEventAppend) -> anyhow::Result<RunEvent> { - unreachable!() - } - fn list_events(&self, _: &RunEventListRequest) -> anyhow::Result<Vec<RunEvent>> { - Ok(vec![]) - } -} diff --git a/crates/tinyagents-orchestration/src/teams/types.rs b/crates/tinyagents-orchestration/src/teams/types.rs deleted file mode 100644 index b436cf125..000000000 --- a/crates/tinyagents-orchestration/src/teams/types.rs +++ /dev/null @@ -1,146 +0,0 @@ -//! Data types for team coordination: members, tasks, shutdown results, and -//! validation errors. - -use serde::Serialize; -use tinyagents_harness::CancellationToken; -use tinyagents_session::run_ledger::{AgentTeam, AgentTeamMember, AgentTeamTask}; - -use crate::subagent::{AgentStepConfig, AgentStepIdentity}; - -/// Sentinel sender for a lead or user message rather than a member row. -/// -/// Used in the event log to distinguish lead/user messages from member-to-member -/// messages. All events with this sender should be treated as external input -/// rather than team member output. -pub const LEAD_SENDER: &str = "lead"; - -/// One member supplied when a team is created. -/// -/// Carries only the identity and (optional) agent-id hint needed for -/// initialization; the full persistent [`AgentTeamMember`] row is created -/// by the ledger. -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct NewMember { - /// Human-readable member name for prompts and display. - pub name: String, - /// Optional reference to an agent definition in the host registry. When - /// supplied, the host may use it to resolve member prompts and tool access. - pub agent_id: Option<String>, -} - -/// A durable team and the member/task rows needed to render it. -/// -/// Projects a complete team state at one point in time: the team metadata, -/// its members, and its tasks. Used for snapshots, UI display, and audit logs. -#[derive(Debug, Clone, PartialEq, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct TeamView { - /// The team metadata and status. - pub team: AgentTeam, - /// All members of the team, in creation order. - pub members: Vec<AgentTeamMember>, - /// All tasks the team has been assigned, in creation order. - pub tasks: Vec<AgentTeamTask>, -} - -/// Result of stopping a member and releasing its active claims. -/// -/// Returned when a running member is shut down. Records the stopped member -/// and the task ids it had claimed (now released and available for -/// reassignment). -#[derive(Debug, Clone, PartialEq, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct MemberShutdown { - /// The member that was shut down. - pub member: AgentTeamMember, - /// Task ids that member had claimed, now released. - pub released_task_ids: Vec<String>, -} - -/// Coordination validation errors that are independent of a host or storage backend. -/// -/// These errors detect structural issues in team definition or task dependency -/// graphs that do not require access to a host's agent registry or persistence -/// layer. They are deterministic and stable across runs. -#[derive(Debug, Clone, PartialEq, Eq, Serialize)] -#[serde(rename_all = "camelCase", tag = "kind", content = "detail")] -pub enum TeamError { - /// A team member name appears more than once. - DuplicateMemberName { name: String }, - /// A task references a member id that is not registered. - UnknownMember { member_id: String }, - /// A task depends on itself. - SelfDependency { task_id: String }, - /// Task dependencies form a cycle. - CyclicDependency, - /// A task depends on another task that is not registered. - UnknownDependency { depends_on: String }, -} - -impl std::fmt::Display for TeamError { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - match self { - Self::DuplicateMemberName { name } => write!(f, "duplicate member name: {name}"), - Self::UnknownMember { member_id } => write!(f, "unknown member: {member_id}"), - Self::SelfDependency { task_id } => write!(f, "task {task_id} cannot depend on itself"), - Self::CyclicDependency => write!(f, "dependency cycle detected"), - Self::UnknownDependency { depends_on } => { - write!(f, "unknown dependency: {depends_on}") - } - } - } -} - -impl std::error::Error for TeamError {} - -/// The driver-side policy and identity of one member step. -/// -/// [`run_member_graph`](super::run_member_graph) uses [`MemberStep::default`], which is inert (unlimited -/// admission, no timeout, one attempt, no result trimming), so the worker runs -/// exactly as it did before the step went through -/// [`SubagentDriver`](crate::subagent::SubagentDriver). Hosts that want the -/// driver's limits call [`run_member_graph_with`](super::run_member_graph_with). -#[derive(Clone)] -pub struct MemberStep { - /// Spawn admission, timeout/retry/budget, result policy and role. - pub config: AgentStepConfig, - /// Team (admission scope) and member (task id, allowlist target). - pub identity: AgentStepIdentity, - /// Lifecycle cancellation of the enclosing run. A token cancelled before - /// the step starts keeps the worker from running (the member is routed to - /// `on_failed`); cancelling it later cancels the worker's child token. - pub cancellation: CancellationToken, -} - -impl MemberStep { - /// A step for `member_id` of `team_id` under `config`; the member id is - /// also the allowlist target. - pub fn new( - config: AgentStepConfig, - team_id: impl Into<String>, - member_id: impl Into<String>, - ) -> Self { - let member_id = member_id.into(); - Self { - config, - identity: AgentStepIdentity::new(team_id, member_id.clone()).with_target(member_id), - cancellation: CancellationToken::new(), - } - } - - /// Ties the step to the enclosing run's cancellation token. - pub fn with_cancellation(mut self, cancellation: CancellationToken) -> Self { - self.cancellation = cancellation; - self - } -} - -impl Default for MemberStep { - fn default() -> Self { - Self { - config: AgentStepConfig::default(), - identity: AgentStepIdentity::new("team", "member"), - cancellation: CancellationToken::new(), - } - } -} diff --git a/crates/tinyagents-orchestration/src/test_support.rs b/crates/tinyagents-orchestration/src/test_support.rs deleted file mode 100644 index f3f3a6346..000000000 --- a/crates/tinyagents-orchestration/src/test_support.rs +++ /dev/null @@ -1,13 +0,0 @@ -//! Shared testkit fixtures for the orchestration drivers. - -use tinyagents_session::run_ledger::DriverRunLedger; -use tinyagents_storage_kit::{Scope, StorageBackend}; -use tinystoragedrivers_testkit::MemoryBackend; - -/// Opens a run ledger over the explicit testkit backend owned by one test. -pub(crate) fn ledger_in(backend: &MemoryBackend) -> DriverRunLedger { - let scoped = backend - .for_scope(&Scope::new("local").unwrap()) - .expect("local scope"); - DriverRunLedger::new(scoped) -} diff --git a/crates/tinyagents-orchestration/src/topology.rs b/crates/tinyagents-orchestration/src/topology.rs deleted file mode 100644 index f8ffc107a..000000000 --- a/crates/tinyagents-orchestration/src/topology.rs +++ /dev/null @@ -1,36 +0,0 @@ -//! Topology export for the orchestration graphs, for debug / inspection. -//! -//! Each orchestration graph exposes a `*_topology()` builder that constructs -//! its structure with no-op stub closures and returns a behaviour-free -//! [`GraphTopology`](tinyagents_graph::GraphTopology). [`all_graph_topologies`] -//! collects them so a UI or debug endpoint can render them as JSON or Mermaid -//! and surface structural defects. - -use tinyagents_graph::GraphTopologyReport; - -/// Structure-only topologies of every fixed-shape orchestration graph. -/// -/// Graphs that fail to build are skipped. Item-count-driven `map_reduce` -/// fan-outs are not exported: their node set is decided per run. -pub fn all_graph_topologies() -> Vec<GraphTopologyReport> { - let mut out = Vec::new(); - - if let Ok(t) = crate::teams::member_graph_topology() { - out.push(GraphTopologyReport::describe("agent_teams:member", &t)); - } - if let Ok(t) = tinyagents_graph::delegation::delegation_graph_topology() { - out.push(GraphTopologyReport::describe("delegation", &t)); - } - if let Ok(t) = crate::workflow::scheduler_topology_preview() { - out.push(GraphTopologyReport::describe( - "workflow_runs:scheduler_preview", - &t, - )); - } - - out -} - -#[cfg(test)] -#[path = "topology_tests.rs"] -mod tests; diff --git a/crates/tinyagents-orchestration/src/topology_tests.rs b/crates/tinyagents-orchestration/src/topology_tests.rs deleted file mode 100644 index 3323c7354..000000000 --- a/crates/tinyagents-orchestration/src/topology_tests.rs +++ /dev/null @@ -1,26 +0,0 @@ -use super::*; - -#[test] -fn all_topologies_includes_member_delegation_and_workflow_scheduler() { - let reports = all_graph_topologies(); - for name in [ - "agent_teams:member", - "delegation", - "workflow_runs:scheduler_preview", - ] { - let report = reports - .iter() - .find(|r| r.name == name) - .unwrap_or_else(|| panic!("the {name} graph should be exported")); - assert!( - report.ok, - "{name} graph should validate structurally: {:?}", - report.errors - ); - assert!( - report.mermaid.contains("flowchart"), - "{name} mermaid should render: {}", - report.mermaid - ); - } -} diff --git a/crates/tinyagents-orchestration/src/workflow/README.md b/crates/tinyagents-orchestration/src/workflow/README.md deleted file mode 100644 index 2ab080469..000000000 --- a/crates/tinyagents-orchestration/src/workflow/README.md +++ /dev/null @@ -1,127 +0,0 @@ -# `orchestration::workflow` — durable phase DAG execution - -A **workflow** is a directed acyclic graph of phases, each with associated -agents, dependencies, and concurrency limits. The [`WorkflowEngine`] schedules -phases topologically, spawns bounded-concurrent child tasks per phase, -collects results, and manages state transitions durably. - -## Public surface - -- **`WorkflowEngine<E: WorkflowExecutor>`** — the execution engine: schedules - runnable phases, spawns child tasks, handles concurrency limits, collects - results, and persists state. Generic over a host-supplied executor. -- **`WorkflowExecutor`** — the trait for host-supplied work: create and - monitor child tasks, cancel them, and query their status. -- **`WorkflowStore`, `SessionWorkflowStore`** — durable run state: loads, - saves, claims (lease), and compare-and-swap updates. [`SessionWorkflowStore`] - wraps the session run ledger. -- **`WorkflowRunSupervisor`** — process-local start/resume/stop coordination. - Hosts supply their task-context spawner and wrap the drive future with any - host parent or origin scopes. Every start installs a fresh cancellation - generation; stale handles cannot cancel a resumed loop. -- **`WorkflowDefinition`, `WorkflowPhase`** — the declarative phase DAG: phase - names, descriptions, agent ids, dependencies, concurrency settings. -- **`PhaseStatus`** — phase state machine: Pending → Running → (Completed | - Failed). Interrupted phases can reset to Pending for retry. -- **State projection functions** — query and mutate the JSON phase-state - document: `phase_status()`, `next_runnable_phase()`, `all_phases_completed()`, - `reset_running_phases()`, `phase_prompt()`, `synthesize_summary()`. - -## Design and invariants - -- **Durable at every boundary:** Workflow runs and phase states are durably - persisted via [`WorkflowStore`]. The engine is safe to interrupt and resume. -- **Topological scheduling:** Phases are scheduled in dependency order. The - engine finds the next runnable phase (all dependencies met), preventing - partial execution and cycles. -- **Bounded parallelism:** The engine respects `default_concurrency` (agents - per phase in parallel) and `max_children` (total spawned children at once). - Phase work is fan-out (parallel agents) + fan-in (collect results). -- **Deterministic retry:** If the engine is interrupted while running phases, - all running phases are reset to Pending (outputs cleared). Completed phases - remain immutable and are never retried. -- **Result aggregation:** Upstream outputs (from dependency phases) are - collected and passed to downstream phases' prompts, allowing workflows to - reason over prior results. -- **JSON phase state:** The phase-state document is JSON (a BTreeMap of phase - names to `{ status, outputs, reason, ... }`). This projection is durable, - queryable, and renderable in UIs. - -## File map - -- **`engine.rs`** — [`WorkflowEngine`] implementation: phase scheduling, child - task spawning and monitoring, state persistence, concurrency enforcement. - [`WorkflowStore`], [`SessionWorkflowStore`], [`WorkflowExecutor`] trait. -- **`supervisor.rs`** — generation-safe cancellation and scheduling handles for - host-launched drive futures. -- **`types.rs`** — data types: [`WorkflowDefinition`], [`WorkflowPhase`], - [`DefinitionError`], [`WorkflowDefinitionListResponse`]. -- **`state.rs`** — phase-state projection: JSON document queries and mutations. - [`PhaseStatus`], `next_runnable_phase()`, `phase_prompt()`, - `synthesize_summary()`, etc. -- **`graph.rs`** — scheduler DAG: models workflow phases as a directed graph - for topological sorting and dependency resolution. -- **`child_step.rs`** — runs each agent child through `SubagentDriver` - (`subagent::run_agent_step`). `WorkflowEngine::with_step_config` applies an - `AgentStepConfig` (spawn admission scoped per workflow run id and targeted - at the agent id, `SubAgentPolicy::timeout`, result policy). The default is - inert; a refused or timed-out child fails its phase with the reason (a - timed-out child is also cancelled by its registered id), and a - result-policy-trimmed output replaces the raw output (as a string). Retry, - token budget and role cannot take effect here: the executor reports no - usage, transient failures or tools. Admission slots are not refunded when an - interrupted phase re-runs on resume, and a per-parent live cap below - `default_concurrency` rejects siblings (failing the phase) instead of queueing. - Each child gets a child cancellation token of the run's token. -- **`validate.rs`** — structural validation: no duplicate phases, valid - dependencies, no cycles, valid concurrency settings, etc. -- **`tests.rs`** — tests for scheduling, phase transitions, concurrency, - result aggregation, and interruption/retry behavior. - -## Relationship to other modules - -- **Depends on:** `tinyagents-graph` (DAG validation), `tinyagents-session` - (run ledger), `tinyagents-harness` (cancellation, error types), and this crate's `subagent` - module (`run_agent_step`). -- **Used by:** host orchestration logic (workflow management and execution). -- **Integration:** Can model multi-agent phases as teams (via - `orchestration::teams`). - -## Typical usage - -1. Define a [`WorkflowDefinition`] with phases and dependencies. -2. Create a [`WorkflowEngine`] with a host-supplied [`WorkflowExecutor`] and - [`WorkflowStore`]. -3. Call `engine.run()`: the engine schedules phases, spawns bounded child - tasks, collects results, and persists state. -4. On interruption, retry `engine.run()`: running phases reset to Pending; - completed phases remain done. - -## Asynchronous authoritative persistence - -AsyncWorkflowStore and LedgerWorkflowStore retain a caller-supplied -RunLedgerStore handle. The adapter awaits load, upsert, claim, compare-and-swap -and lease renewal without choosing a workspace or falling back to SQLite. -WorkflowPersistenceError retains LedgerError, including native driver kinds. - -AsyncWorkflowChildRegistration requires the executor to await registration; -its synchronous entrypoint refuses. LegacyWorkflowRegistration explicitly -adapts existing desktop callbacks. - -WorkflowEngine::new keeps the synchronous desktop contracts compatible. -WorkflowEngine::new_async selects only the supplied AsyncWorkflowStore and -AsyncWorkflowExecutor. Its initialise_async and drive_async entrypoints await -persistence, child registration and heartbeat renewal, preserving typed sources. -Legacy initialise/drive entrypoints refuse on an authoritative instance. - -An async executor must await register_async before waiting for a child. A -refused registration interrupts the phase even if an executor ignores the -error, cancels attempted children, and prevents later success. Results without -a committed registration are rejected. Native fenced/backend failures remain -typed and do not become a successful lifecycle hand-off or a SQLite fallback. - -This implements workflow-row persistence only. Child progress/status/events, -host provider routing, generation-bound import and owner-aware recovery are -separate prerequisites. Hosts must not advertise a complete cloud workflow -migration or opt in before those owning contracts and dependencies are -available upstream. diff --git a/crates/tinyagents-orchestration/src/workflow/child_step.rs b/crates/tinyagents-orchestration/src/workflow/child_step.rs deleted file mode 100644 index e632c311b..000000000 --- a/crates/tinyagents-orchestration/src/workflow/child_step.rs +++ /dev/null @@ -1,238 +0,0 @@ -//! One workflow child invocation, run as a [`SubagentDriver`](crate::subagent::SubagentDriver) -//! lifecycle so workflow agent phases share the spawn-admission, timeout / -//! retry / budget, result-policy and typed-outcome behaviour of every other -//! subagent path. - -use std::sync::Arc; - -use serde_json::Value; -use tinyagents_harness::CancellationToken; - -use super::engine::{ - OrchestrationError, WorkflowChildRegistration, WorkflowChildRequest, WorkflowChildResult, - WorkflowExecutor, render_compat_output, -}; -use super::persistence::{ - AsyncWorkflowChildRegistration, AsyncWorkflowExecutor, WorkflowPersistenceError, - WorkflowPersistenceResult, -}; -use crate::subagent::{ - AgentStepConfig, AgentStepIdentity, IncompleteKind, ResultPolicy, StepSuccess, SubagentOutcome, - SubagentOutcomeKind, run_agent_step, -}; - -/// Remembers the child ids one step registered so a timed-out child can be -/// cancelled by id (the engine only cancels on interruption or lease loss). -struct RecordingRegistration { - inner: Arc<dyn WorkflowChildRegistration>, - async_inner: Option<Arc<dyn AsyncWorkflowChildRegistration>>, - ids: parking_lot::Mutex<Vec<String>>, -} - -impl WorkflowChildRegistration for RecordingRegistration { - fn register(&self, child_id: String) -> Result<(), OrchestrationError> { - self.ids.lock().push(child_id.clone()); - self.inner.register(child_id) - } -} - -#[async_trait::async_trait] -impl AsyncWorkflowChildRegistration for RecordingRegistration { - fn register(&self, child_id: String) -> WorkflowPersistenceResult<()> { - self.ids.lock().push(child_id.clone()); - match &self.async_inner { - Some(inner) => inner.register(child_id), - None => Ok(self.inner.register(child_id)?), - } - } - - async fn register_async(&self, child_id: String) -> WorkflowPersistenceResult<()> { - self.ids.lock().push(child_id.clone()); - match &self.async_inner { - Some(inner) => inner.register_async(child_id).await, - None => Ok(self.inner.register(child_id)?), - } - } -} - -const LOG_PREFIX: &str = "[workflow-child-step]"; - -#[cfg(test)] -pub(super) async fn run_child_step<E: WorkflowExecutor + 'static>( - config: &AgentStepConfig, - executor: Arc<E>, - request: WorkflowChildRequest, - cancel: CancellationToken, - registration: Arc<dyn WorkflowChildRegistration>, -) -> Result<(WorkflowChildResult, Option<Value>), OrchestrationError> { - run_child_step_authoritative(config, executor, None, request, cancel, registration, None) - .await - .map_err(WorkflowPersistenceError::into_legacy) -} - -pub(super) async fn run_child_step_authoritative<E: WorkflowExecutor + 'static>( - config: &AgentStepConfig, - executor: Arc<E>, - async_executor: Option<Arc<dyn AsyncWorkflowExecutor>>, - request: WorkflowChildRequest, - cancel: CancellationToken, - registration: Arc<dyn WorkflowChildRegistration>, - async_registration: Option<Arc<dyn AsyncWorkflowChildRegistration>>, -) -> WorkflowPersistenceResult<(WorkflowChildResult, Option<Value>)> { - let identity = AgentStepIdentity::new( - request.run_id.clone(), - format!("{}:{}", request.phase, request.index_in_phase), - ) - .with_target(request.agent_id.clone()); - tracing::debug!( - "{LOG_PREFIX} start run={} phase={} index={} agent={}", - request.run_id, - request.phase, - request.index_in_phase, - request.agent_id - ); - let work_request = request.clone(); - // The driver validates a schema against the output *text* only, which - // mis-reads a JSON string. Validate the structured value here instead and - // keep the schema out of the driver's policy so truncation and artifact - // storage keep seeing the display text. - let schema = config.result_policy.schema.clone(); - let mut step_config = config.clone(); - step_config.result_policy.schema = None; - let config = &step_config; - let registration = Arc::new(RecordingRegistration { - inner: registration, - async_inner: async_registration, - ids: parking_lot::Mutex::new(Vec::new()), - }); - let recorded = registration.clone(); - let cancel_executor = executor.clone(); - let cancel_async_executor = async_executor.clone(); - // The generic step outcome carries textual reasons. Keep the exact - // executor source across that projection for the authoritative caller. - let typed_failure = Arc::new(parking_lot::Mutex::new(None)); - let worker_failure = typed_failure.clone(); - // `cancel` is the lifecycle token: cancelling the run reaches the child, - // whose own token the driver derives from it. - let result = run_agent_step(config, identity, cancel, move |ctx| { - let executor = executor.clone(); - let async_executor = async_executor.clone(); - let typed_failure = worker_failure.clone(); - let request = work_request.clone(); - let registration = registration.clone(); - async move { - let result = match async_executor { - Some(executor) => { - executor - .execute_async(request, ctx.cancellation, registration) - .await - } - None => executor - .execute(request, ctx.cancellation, registration) - .await - .map_err(WorkflowPersistenceError::from), - } - .map_err(|error| { - let message = error.to_string(); - *typed_failure.lock() = Some(error); - anyhow::anyhow!(message) - })?; - let rendered = render_compat_output(&result.output); - Ok(StepSuccess::new(rendered, result)) - } - }) - .await; - if let Some(error) = typed_failure.lock().take() { - return Err(error); - } - match result { - Ok(step) => match step.outcome.status { - SubagentOutcomeKind::Completed => { - let mut child = step.value.ok_or_else(|| { - OrchestrationError("workflow child completed without a result".to_owned()) - })?; - let original = child.output.clone(); - // Keep the executor's structured output unless the result - // policy actually changed the text. - if step.outcome.output != render_compat_output(&child.output) { - child.output = Value::String(step.outcome.output.clone()); - } - let schema_error = match &schema { - Some(schema) => { - // Validate the executor's structured value, not the - // possibly trimmed display text. - ResultPolicy::new() - .with_schema(schema.clone()) - .apply(&step.outcome.task_id, &step.outcome.output, Some(&original)) - .await - .schema_error - } - None => None, - }; - Ok(( - child, - result_policy_annotations(&step.outcome, schema_error), - )) - } - SubagentOutcomeKind::Incomplete(incomplete) => { - tracing::debug!( - "{LOG_PREFIX} incomplete run={} phase={} agent={} kind={:?}", - request.run_id, - request.phase, - request.agent_id, - incomplete.kind - ); - if incomplete.kind == IncompleteKind::Timeout { - // The driver dropped the executor future; the host's real - // child may still be running, so cancel what it registered. - let ids = recorded.ids.lock().clone(); - if !ids.is_empty() { - match &cancel_async_executor { - Some(executor) => executor.cancel_children(&ids).await, - None => cancel_executor.cancel_children(&ids).await, - } - } - } - Err(OrchestrationError(incomplete.reason).into()) - } - // The child registered its id before it could be cancelled, so the - // engine keeps it (via the registration) without a result. - SubagentOutcomeKind::Cancelled => { - Err(OrchestrationError("workflow child cancelled".to_owned()).into()) - } - SubagentOutcomeKind::AwaitingInput(_) => { - Err(OrchestrationError("workflow child paused for input".to_owned()).into()) - } - }, - Err(error) => { - tracing::debug!( - "{LOG_PREFIX} failed run={} phase={} agent={} error={error}", - request.run_id, - request.phase, - request.agent_id - ); - Err(OrchestrationError(error.to_string()).into()) - } - } -} - -/// Result-policy findings (schema check, artifact overflow) the driver -/// recorded on the outcome, for the phase output metadata. `None` when the -/// policy had nothing to report. -fn result_policy_annotations( - outcome: &SubagentOutcome, - schema_error: Option<String>, -) -> Option<Value> { - if schema_error.is_none() && outcome.artifact_error.is_none() && outcome.artifacts.is_empty() { - return None; - } - Some(serde_json::json!({ - "schemaError": schema_error, - "artifactError": outcome.artifact_error, - "artifacts": outcome.artifacts, - })) -} - -#[cfg(test)] -#[path = "child_step_tests.rs"] -mod tests; diff --git a/crates/tinyagents-orchestration/src/workflow/child_step_tests.rs b/crates/tinyagents-orchestration/src/workflow/child_step_tests.rs deleted file mode 100644 index 64d124605..000000000 --- a/crates/tinyagents-orchestration/src/workflow/child_step_tests.rs +++ /dev/null @@ -1,199 +0,0 @@ -//! Workflow agent children run through `SubagentDriver`: configured driver -//! policy takes effect on a workflow phase. - -use super::*; - -use std::sync::Arc; - -use serde_json::json; -use tinyagents_harness::CancellationToken; -use tinyagents_session::run_ledger::WorkflowRunStatus; - -use crate::subagent::{ResultPolicy, SpawnAdmission, SpawnPolicy}; -use crate::workflow::tests::{FakeExecutor, MemoryStore, definition}; -use crate::workflow::{WorkflowEngine, WorkflowStore}; - -fn run_with( - config: AgentStepConfig, -) -> ( - Arc<MemoryStore>, - Arc<FakeExecutor>, - WorkflowEngine<MemoryStore, FakeExecutor>, -) { - let store = Arc::new(MemoryStore::default()); - let executor = Arc::new(FakeExecutor::default()); - let engine = WorkflowEngine::new(store.clone(), executor.clone()).with_step_config(config); - (store, executor, engine) -} - -async fn drive( - engine: &WorkflowEngine<MemoryStore, FakeExecutor>, - store: &MemoryStore, -) -> serde_json::Value { - let def = definition(); - engine - .initialise("run".into(), &def, json!("q"), None) - .unwrap(); - engine - .drive("run", &def, CancellationToken::new()) - .await - .unwrap(); - serde_json::to_value(store.load("run").unwrap().unwrap().phase_states).unwrap() -} - -#[tokio::test] -async fn result_policy_cap_trims_a_workflow_child_output() { - let (store, _executor, engine) = run_with( - AgentStepConfig::default().with_result_policy(ResultPolicy::new().with_max_chars(4)), - ); - let states = drive(&engine, &store).await; - let out = states["plan"]["outputs"][0]["output"] - .as_str() - .unwrap() - .to_owned(); - assert_ne!(out, "plan output", "ResultPolicy trimmed the child output"); - assert!( - out.chars().count() <= 4 + 64, - "capped near max_chars: {out}" - ); - // Trimmed output is carried verbatim as the raw output too. - assert_eq!( - states["plan"]["outputs"][0]["metadata"]["rawOutput"], - json!(out) - ); -} - -#[tokio::test] -async fn spawn_policy_denial_fails_the_phase_and_never_runs_the_child() { - let (store, executor, engine) = run_with(AgentStepConfig::default().with_admission( - SpawnAdmission::new(SpawnPolicy { - allowed_targets: Some(vec!["only-this-agent".into()]), - ..Default::default() - }), - )); - drive(&engine, &store).await; - let run = store.load("run").unwrap().unwrap(); - assert_eq!(run.status, WorkflowRunStatus::Failed); - assert!(executor.calls.lock().is_empty(), "no child was executed"); - assert!( - run.phase_states.to_string().contains("not admitted"), - "{}", - run.phase_states - ); -} - -#[tokio::test] -async fn default_step_config_leaves_the_output_untouched() { - let (store, _executor, engine) = run_with(AgentStepConfig::default()); - let states = drive(&engine, &store).await; - assert_eq!(states["plan"]["outputs"][0]["output"], json!("plan output")); - assert_eq!( - store.load("run").unwrap().unwrap().status, - WorkflowRunStatus::Completed - ); -} - -#[tokio::test] -async fn timed_out_child_is_cancelled_by_its_registered_id() { - use std::time::Duration; - - use crate::subagent::SubAgentPolicy; - use crate::workflow::tests::BlockingExecutor; - - let store = Arc::new(MemoryStore::default()); - let executor = Arc::new(BlockingExecutor::default()); - let engine = WorkflowEngine::new(store.clone(), executor.clone()).with_step_config( - AgentStepConfig::default() - .with_policy(SubAgentPolicy::default().with_timeout(Duration::from_millis(30))), - ); - let def = definition(); - engine - .initialise("run".into(), &def, json!("q"), None) - .unwrap(); - engine - .drive("run", &def, CancellationToken::new()) - .await - .unwrap(); - let run = store.load("run").unwrap().unwrap(); - assert_eq!(run.status, WorkflowRunStatus::Failed); - assert!( - run.phase_states.to_string().contains("timed out"), - "{}", - run.phase_states - ); - assert!( - executor - .cancelled - .lock() - .iter() - .any(|id| id == "live-plan-0"), - "the timed-out child was cancelled: {:?}", - executor.cancelled.lock() - ); -} - -#[tokio::test] -async fn schema_failure_is_recorded_in_phase_metadata() { - let (store, _executor, engine) = run_with(AgentStepConfig::default().with_result_policy( - ResultPolicy::new().with_schema(json!({"type": "object", "required": ["x"]})), - )); - let states = drive(&engine, &store).await; - let meta = &states["plan"]["outputs"][0]["metadata"]; - assert!(meta["resultPolicy"]["schemaError"].is_string(), "{meta}"); -} - -#[tokio::test] -async fn string_output_validates_against_a_string_schema() { - let (store, _executor, engine) = run_with( - AgentStepConfig::default() - .with_result_policy(ResultPolicy::new().with_schema(json!({"type": "string"}))), - ); - let states = drive(&engine, &store).await; - let out = &states["plan"]["outputs"][0]; - assert!(out["metadata"].get("resultPolicy").is_none(), "{out}"); - assert_eq!(out["output"], json!("plan output")); -} - -#[tokio::test] -async fn schema_and_cap_together_keep_the_display_text_for_truncation() { - // "plan output" is 11 chars; a cap of 11 must leave it untouched even - // though its JSON form (13 chars) would not fit. - let (store, _executor, engine) = run_with( - AgentStepConfig::default().with_result_policy( - ResultPolicy::new() - .with_schema(json!({"type": "string"})) - .with_max_chars(11), - ), - ); - let states = drive(&engine, &store).await; - assert_eq!(states["plan"]["outputs"][0]["output"], json!("plan output")); -} - -#[tokio::test] -async fn pre_cancelled_child_never_reaches_the_executor() { - struct NoRegistration; - impl WorkflowChildRegistration for NoRegistration { - fn register(&self, _: String) -> Result<(), OrchestrationError> { - Ok(()) - } - } - let executor = Arc::new(FakeExecutor::default()); - let token = CancellationToken::new(); - token.cancel(); - let result = run_child_step( - &AgentStepConfig::default(), - executor.clone(), - WorkflowChildRequest { - run_id: "run".into(), - phase: "plan".into(), - agent_id: "planner".into(), - index_in_phase: 0, - prompt: "p".into(), - }, - token, - Arc::new(NoRegistration), - ) - .await; - assert!(result.is_err()); - assert!(executor.calls.lock().is_empty(), "no child was executed"); -} diff --git a/crates/tinyagents-orchestration/src/workflow/engine.rs b/crates/tinyagents-orchestration/src/workflow/engine.rs deleted file mode 100644 index c0734d43f..000000000 --- a/crates/tinyagents-orchestration/src/workflow/engine.rs +++ /dev/null @@ -1,1284 +0,0 @@ -//! Workflow execution engine: orchestrates phase scheduling, task spawning, -//! and result collection. -//! -//! [`WorkflowEngine`] drives the scheduler, manages child task creation and -//! claim-based assignment, handles phase state persistence and transitions, -//! and enforces bounded concurrency. Legacy hosts supply [`WorkflowExecutor`]; -//! authoritative hosts select [`AsyncWorkflowExecutor`] and await the explicit -//! async persistence entrypoints. - -use std::fmt; -use std::sync::Arc; -use std::sync::atomic::{AtomicU64, Ordering}; -use std::time::Duration; - -use async_trait::async_trait; -use chrono::Utc; -use serde_json::{Value, json}; -use tinyagents_graph::GraphEventSink; -use tinyagents_graph::parallel::{FailurePolicy, ParallelOptions, map_reduce}; -use tinyagents_harness::CancellationToken; -use tinyagents_session::run_ledger::{ - RunLedgerStore, WorkflowLeaseClaim, WorkflowRun, WorkflowRunStatus, WorkflowRunUpsert, -}; -use tinyagents_storage_kit::Blocking; - -use super::child_step::run_child_step_authoritative; -use super::persistence::{ - AsyncWorkflowChildRegistration, AsyncWorkflowExecutor, AsyncWorkflowStore, - WorkflowPersistenceError, WorkflowPersistenceResult, -}; -use super::state::{ - PhaseStatus, all_phases_completed, init_phase_states, next_runnable_phase, phase_prompt, - reset_running_phases, set_phase_reason, set_phase_status, synthesize_summary, upstream_outputs, -}; -use super::{WorkflowDefinition, WorkflowPhase}; -use crate::subagent::AgentStepConfig; - -/// Error returned by a host child executor. -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct OrchestrationError(pub String); - -impl fmt::Display for OrchestrationError { - fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { - self.0.fmt(formatter) - } -} - -impl std::error::Error for OrchestrationError {} - -impl From<anyhow::Error> for OrchestrationError { - fn from(error: anyhow::Error) -> Self { - Self(error.to_string()) - } -} - -impl From<tinyagents_harness::TinyAgentsError> for OrchestrationError { - fn from(error: tinyagents_harness::TinyAgentsError) -> Self { - Self(error.to_string()) - } -} - -/// Durable workflow rows supplied by the session owner. -pub trait WorkflowStore: Send + Sync { - fn load(&self, id: &str) -> Result<Option<WorkflowRun>, OrchestrationError>; - fn upsert(&self, row: WorkflowRunUpsert) -> Result<WorkflowRun, OrchestrationError>; - fn claim( - &self, - id: &str, - owner: &str, - lease_for: Duration, - ) -> Result<WorkflowLeaseClaim, OrchestrationError>; - fn compare_and_swap( - &self, - row: WorkflowRunUpsert, - expected_revision: u64, - owner: &str, - lease_for: Duration, - ) -> Result<Option<WorkflowRun>, OrchestrationError>; - fn renew(&self, id: &str, owner: &str, lease_for: Duration) - -> Result<bool, OrchestrationError>; -} - -/// Synchronous [`WorkflowStore`] over a bound run-ledger handle. -/// -/// Every call reaches the ledger's async operations through a [`Blocking`] -/// bridge thread, so it works from inside or outside a tokio runtime. The -/// handle is whatever backend the host opened (SQLite, MongoDB, memory); this -/// adapter opens no file and has no fallback. Prefer -/// [`LedgerWorkflowStore`](super::LedgerWorkflowStore) with -/// `WorkflowEngine::new_async`, which keeps typed errors and needs no bridge. -#[derive(Clone)] -pub struct SessionWorkflowStore { - ledger: Arc<dyn RunLedgerStore>, - bridge: Blocking, -} - -impl std::fmt::Debug for SessionWorkflowStore { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.debug_struct("SessionWorkflowStore") - .finish_non_exhaustive() - } -} - -impl SessionWorkflowStore { - /// Adapt `ledger`, starting a bridge thread for the synchronous calls. - /// - /// # Errors - /// - /// The bridge thread cannot start. - pub fn new(ledger: Arc<dyn RunLedgerStore>) -> Result<Self, OrchestrationError> { - Ok(Self::with_bridge( - ledger, - Blocking::new().map_err(|error| OrchestrationError(error.to_string()))?, - )) - } - - /// Adapt `ledger`, sharing an existing bridge. - pub fn with_bridge(ledger: Arc<dyn RunLedgerStore>, bridge: Blocking) -> Self { - Self { ledger, bridge } - } - - fn run<T: Send + 'static>( - &self, - work: impl std::future::Future<Output = Result<T, tinyagents_session::run_ledger::LedgerError>> - + Send - + 'static, - ) -> Result<T, OrchestrationError> { - self.bridge - .run(work) - .map_err(|error| OrchestrationError(error.to_string()))? - .map_err(|error| OrchestrationError(error.to_string())) - } -} - -impl WorkflowStore for SessionWorkflowStore { - fn load(&self, id: &str) -> Result<Option<WorkflowRun>, OrchestrationError> { - let (ledger, id) = (Arc::clone(&self.ledger), id.to_owned()); - self.run(async move { ledger.get_workflow_run(&id).await }) - } - - fn upsert(&self, row: WorkflowRunUpsert) -> Result<WorkflowRun, OrchestrationError> { - let ledger = Arc::clone(&self.ledger); - self.run(async move { ledger.upsert_workflow_run(row).await }) - } - - fn claim( - &self, - id: &str, - owner: &str, - lease_for: Duration, - ) -> Result<WorkflowLeaseClaim, OrchestrationError> { - let lease = chrono::Duration::from_std(lease_for) - .map_err(|error| OrchestrationError(error.to_string()))?; - let (ledger, id, owner) = (Arc::clone(&self.ledger), id.to_owned(), owner.to_owned()); - self.run(async move { ledger.try_claim_workflow_run(&id, &owner, lease).await }) - } - - fn compare_and_swap( - &self, - row: WorkflowRunUpsert, - expected_revision: u64, - owner: &str, - lease_for: Duration, - ) -> Result<Option<WorkflowRun>, OrchestrationError> { - let lease = chrono::Duration::from_std(lease_for) - .map_err(|error| OrchestrationError(error.to_string()))?; - let (ledger, owner) = (Arc::clone(&self.ledger), owner.to_owned()); - self.run(async move { - ledger - .compare_and_swap_workflow_run(row, expected_revision, &owner, lease) - .await - }) - } - - fn renew( - &self, - id: &str, - owner: &str, - lease_for: Duration, - ) -> Result<bool, OrchestrationError> { - let lease = chrono::Duration::from_std(lease_for) - .map_err(|error| OrchestrationError(error.to_string()))?; - let (ledger, id, owner) = (Arc::clone(&self.ledger), id.to_owned(), owner.to_owned()); - self.run(async move { ledger.renew_workflow_run_lease(&id, &owner, lease).await }) - } -} - -/// One host-authorized child invocation. -#[derive(Debug, Clone, PartialEq)] -pub struct WorkflowChildRequest { - pub run_id: String, - pub phase: String, - pub agent_id: String, - pub index_in_phase: usize, - pub prompt: String, -} - -/// One child's terminal result. `output` is retained verbatim in phase state. -#[derive(Debug, Clone, PartialEq)] -pub struct WorkflowChildResult { - pub child_id: String, - pub output: Value, -} - -/// Called by a host immediately after it has created a real child. Registering -/// before waiting makes the child visible to a concurrent cancellation request -/// even when the worker is still in flight. -pub trait WorkflowChildRegistration: Send + Sync { - fn register(&self, child_id: String) -> Result<(), OrchestrationError>; -} - -/// Host-owned execution and cancellation mechanism. -#[async_trait] -pub trait WorkflowExecutor: Send + Sync { - async fn execute( - &self, - request: WorkflowChildRequest, - cancel: CancellationToken, - registration: Arc<dyn WorkflowChildRegistration>, - ) -> Result<WorkflowChildResult, OrchestrationError>; - - async fn cancel_children(&self, child_ids: &[String]); -} - -/// Generic durable workflow engine. It never creates tasks: hosts decide where -/// work runs and which authorization context is in force via [`WorkflowExecutor`]. -pub struct WorkflowEngine<S, E> { - store: Arc<S>, - executor: Arc<E>, - async_store: Option<Arc<dyn AsyncWorkflowStore>>, - async_executor: Option<Arc<dyn AsyncWorkflowExecutor>>, - event_sink: Option<Arc<dyn GraphEventSink>>, - event_seq: AtomicU64, - lease_for: Duration, - step_config: AgentStepConfig, -} - -const WORKFLOW_LEASE: Duration = Duration::from_secs(10 * 60); - -struct PersistRequest { - phase_states: Value, - child_run_ids: Vec<String>, - status: WorkflowRunStatus, - summary: Option<String>, - terminal: bool, -} - -struct PhaseRegistration<S: WorkflowStore> { - store: Arc<S>, - owner: String, - run: parking_lot::Mutex<WorkflowRun>, - phase_states: Value, - lease_for: Duration, - async_store: Option<Arc<dyn AsyncWorkflowStore>>, - operation: tokio::sync::Mutex<()>, - failure: parking_lot::Mutex<Option<Arc<WorkflowPersistenceError>>>, - failure_notify: tokio::sync::Notify, - attempted: parking_lot::Mutex<Vec<String>>, -} - -impl<S: WorkflowStore> PhaseRegistration<S> { - fn current(&self) -> WorkflowRun { - self.run.lock().clone() - } -} - -impl<S: WorkflowStore + 'static> WorkflowChildRegistration for PhaseRegistration<S> { - fn register(&self, child_id: String) -> Result<(), OrchestrationError> { - if self.async_store.is_some() { - return Err(OrchestrationError( - "authoritative registration requires register_async".into(), - )); - } - let mut run = self.run.lock(); - if run.child_run_ids.iter().any(|known| known == &child_id) { - return Ok(()); - } - let mut children = run.child_run_ids.clone(); - children.push(child_id); - let Some(updated) = self.store.compare_and_swap( - WorkflowRunUpsert { - id: run.id.clone(), - definition_id: run.definition_id.clone(), - parent_thread_id: run.parent_thread_id.clone(), - input: run.input.clone(), - phase_states: self.phase_states.clone(), - child_run_ids: children, - status: WorkflowRunStatus::Running, - summary: None, - started_at: Some(run.started_at), - completed_at: None, - }, - run.revision, - &self.owner, - self.lease_for, - )? - else { - return Err(OrchestrationError( - "workflow lease lost while registering child".into(), - )); - }; - *run = updated; - Ok(()) - } -} - -#[async_trait] -impl<S: WorkflowStore + 'static> AsyncWorkflowChildRegistration for PhaseRegistration<S> { - fn register(&self, child_id: String) -> WorkflowPersistenceResult<()> { - if self.async_store.is_none() { - return Ok(WorkflowChildRegistration::register(self, child_id)?); - } - self.attempted.lock().push(child_id); - let error = Arc::new(WorkflowPersistenceError::AsyncRegistrationRequired); - *self.failure.lock() = Some(error.clone()); - self.failure_notify.notify_one(); - Err(WorkflowPersistenceError::Registration(error)) - } - - async fn register_async(&self, child_id: String) -> WorkflowPersistenceResult<()> { - let Some(store) = &self.async_store else { - return Ok(WorkflowChildRegistration::register(self, child_id)?); - }; - let _operation = self.operation.lock().await; - let run = self.current(); - if run.child_run_ids.contains(&child_id) { - return Ok(()); - } - self.attempted.lock().push(child_id.clone()); - let mut children = run.child_run_ids.clone(); - children.push(child_id); - let result = store - .compare_and_swap( - WorkflowRunUpsert { - id: run.id.clone(), - definition_id: run.definition_id.clone(), - parent_thread_id: run.parent_thread_id.clone(), - input: run.input.clone(), - phase_states: self.phase_states.clone(), - child_run_ids: children, - status: WorkflowRunStatus::Running, - summary: None, - started_at: Some(run.started_at), - completed_at: None, - }, - run.revision, - &self.owner, - self.lease_for, - ) - .await; - match result { - Ok(Some(updated)) => { - *self.run.lock() = updated; - Ok(()) - } - other => { - let error = Arc::new(match other { - Err(error) => error, - Ok(None) => WorkflowPersistenceError::LeaseLost, - Ok(Some(_)) => unreachable!(), - }); - *self.failure.lock() = Some(error.clone()); - self.failure_notify.notify_one(); - Err(WorkflowPersistenceError::Registration(error)) - } - } - } -} - -impl<S, E> WorkflowEngine<S, E> -where - S: WorkflowStore + 'static, - E: WorkflowExecutor + 'static, -{ - pub fn new(store: Arc<S>, executor: Arc<E>) -> Self { - Self { - store, - executor, - async_store: None, - async_executor: None, - event_sink: None, - event_seq: AtomicU64::new(0), - lease_for: WORKFLOW_LEASE, - step_config: AgentStepConfig::default(), - } - } - - /// Attach an optional host sink for ordinary graph lifecycle tracing. - pub fn with_event_sink(mut self, sink: Arc<dyn GraphEventSink>) -> Self { - self.event_sink = Some(sink); - self - } - - /// Apply driver policy to every agent child: spawn admission (scoped per - /// workflow run, targeted at the agent id), timeout and result policy. - /// Retry, token budget and role have no effect here: the executor reports - /// no transient failures, usage or tools. - /// - /// Without this call the config is inert and children run exactly as - /// before. A child refused by the spawn policy or timed out fails its - /// phase with that reason. - pub fn with_step_config(mut self, config: AgentStepConfig) -> Self { - self.step_config = config; - self - } - - /// Override the driver lease for deterministic tests or hosts with a - /// deliberately shorter failure-detection window. - pub fn with_lease_duration(mut self, lease_for: Duration) -> Self { - self.lease_for = lease_for.max(Duration::from_millis(3)); - self - } - - /// Initialise a durable run before the host schedules [`Self::drive`]. - pub fn initialise( - &self, - id: String, - definition: &WorkflowDefinition, - input: Value, - parent_thread_id: Option<String>, - ) -> Result<WorkflowRun, OrchestrationError> { - self.store.upsert(WorkflowRunUpsert { - id, - definition_id: definition.id.clone(), - parent_thread_id, - input, - phase_states: init_phase_states(definition), - child_run_ids: Vec::new(), - status: WorkflowRunStatus::Running, - summary: None, - started_at: None, - completed_at: None, - }) - } - - /// Initialise using the authoritative asynchronous store when selected. - pub async fn initialise_async( - &self, - id: String, - definition: &WorkflowDefinition, - input: Value, - parent_thread_id: Option<String>, - ) -> WorkflowPersistenceResult<WorkflowRun> { - let row = WorkflowRunUpsert { - id, - definition_id: definition.id.clone(), - parent_thread_id, - input, - phase_states: init_phase_states(definition), - child_run_ids: Vec::new(), - status: WorkflowRunStatus::Running, - summary: None, - started_at: None, - completed_at: None, - }; - match &self.async_store { - Some(store) => store.upsert(row).await, - None => Ok(self.store.upsert(row)?), - } - } - - /// Drive using the typed authoritative contract. - pub async fn drive_async( - &self, - run_id: &str, - definition: &WorkflowDefinition, - cancel: CancellationToken, - ) -> WorkflowPersistenceResult<()> { - self.drive_inner(run_id, definition, cancel).await - } - - /// Drive a run to a terminal state. Completed phases are never executed - /// again, so a host may safely call this after process restart or resume. - pub async fn drive( - &self, - run_id: &str, - definition: &WorkflowDefinition, - cancel: CancellationToken, - ) -> Result<(), OrchestrationError> { - if self.async_store.is_some() { - return Err(OrchestrationError( - "authoritative workflow requires drive_async".into(), - )); - } - self.drive_inner(run_id, definition, cancel) - .await - .map_err(WorkflowPersistenceError::into_legacy) - } - - async fn load(&self, id: &str) -> WorkflowPersistenceResult<Option<WorkflowRun>> { - match &self.async_store { - Some(store) => store.load(id).await, - None => Ok(self.store.load(id)?), - } - } - async fn claim( - &self, - id: &str, - owner: &str, - lease: Duration, - ) -> WorkflowPersistenceResult<WorkflowLeaseClaim> { - match &self.async_store { - Some(store) => store.claim(id, owner, lease).await, - None => Ok(self.store.claim(id, owner, lease)?), - } - } - async fn renew( - &self, - id: &str, - owner: &str, - lease: Duration, - ) -> WorkflowPersistenceResult<bool> { - match &self.async_store { - Some(store) => store.renew(id, owner, lease).await, - None => Ok(self.store.renew(id, owner, lease)?), - } - } - async fn compare_and_swap( - &self, - row: WorkflowRunUpsert, - revision: u64, - owner: &str, - lease: Duration, - ) -> WorkflowPersistenceResult<Option<WorkflowRun>> { - match &self.async_store { - Some(store) => store.compare_and_swap(row, revision, owner, lease).await, - None => Ok(self.store.compare_and_swap(row, revision, owner, lease)?), - } - } - async fn cancel_children(&self, ids: &[String]) { - match &self.async_executor { - Some(executor) => { - let mut seen = std::collections::HashSet::new(); - let unique = ids - .iter() - .filter(|id| seen.insert((*id).clone())) - .cloned() - .collect::<Vec<_>>(); - executor.cancel_children(&unique).await; - } - None => self.executor.cancel_children(ids).await, - } - } - - async fn drive_inner( - &self, - run_id: &str, - definition: &WorkflowDefinition, - cancel: CancellationToken, - ) -> WorkflowPersistenceResult<()> { - // A driver lease is acquired before looking for runnable work. This - // is deliberately separate from the in-process cancellation token: - // resume can race in another process, and only the durable lease - // prevents both drivers from spawning the same phase. - let owner = uuid::Uuid::new_v4().to_string(); - let mut run = match self.claim(run_id, &owner, self.lease_for).await? { - WorkflowLeaseClaim::Acquired(run) => run, - WorkflowLeaseClaim::Busy(_) => return Ok(()), - WorkflowLeaseClaim::Missing => { - return Err(OrchestrationError(format!( - "workflow run {run_id} vanished before start" - )) - .into()); - } - }; - // A crashed owner can leave the durable phase marked `running`. That - // marker is intentionally not runnable, so reclaiming the lease must - // turn it back into retryable work before scheduling. The claim above - // fences every prior owner; this CAS is the new owner's durable - // recovery transition rather than a read/modify/write race. - if run.phase_states.as_object().is_some_and(|phases| { - phases - .values() - .any(|phase| phase.get("status").and_then(Value::as_str) == Some("running")) - }) { - // The previous owner may have crashed after registering remote - // children. Fence side effects before making its phase runnable. - self.cancel_children(&run.child_run_ids).await; - let mut phase_states = run.phase_states.clone(); - reset_running_phases( - &mut phase_states, - "workflow owner expired; phase will retry after lease takeover", - ); - run = self - .persist( - &run, - PersistRequest { - phase_states, - child_run_ids: Vec::new(), - status: WorkflowRunStatus::Running, - summary: None, - terminal: false, - }, - &owner, - ) - .await?; - } - self.emit( - run_id, - tinyagents_graph::GraphEvent::RunStarted { - run_id: tinyagents_harness::ids::RunId::new(run_id), - }, - ); - // Registered children from an interrupted attempt are historical, not - // part of the retry's spawn budget. - let mut total_spawned = 0; - - loop { - if cancel.is_cancelled() { - self.cancel_children(&run.child_run_ids).await; - let mut phase_states = run.phase_states.clone(); - reset_running_phases( - &mut phase_states, - "workflow interrupted; phase will retry on resume", - ); - if let Err(error) = self - .persist( - &run, - PersistRequest { - phase_states, - child_run_ids: run.child_run_ids.clone(), - status: WorkflowRunStatus::Interrupted, - summary: None, - terminal: false, - }, - &owner, - ) - .await - { - if self.async_store.is_some() { - return Err(error); - } - if self.owner_lost(run_id, &owner).await.unwrap_or(false) - || self - .emit_recorded_terminal(run_id, total_spawned as usize) - .await - .unwrap_or(false) - { - return Ok(()); - } - self.finish_failed(run_id, error.to_string()); - return Err(error); - } - self.finish_cancelled(run_id); - return Ok(()); - } - let Some(phase) = next_runnable_phase(definition, &run.phase_states).cloned() else { - if all_phases_completed(definition, &run.phase_states) { - if let Err(error) = self - .persist( - &run, - PersistRequest { - phase_states: run.phase_states.clone(), - child_run_ids: run.child_run_ids.clone(), - status: WorkflowRunStatus::Completed, - summary: synthesize_summary(definition, &run.phase_states), - terminal: true, - }, - &owner, - ) - .await - { - if self.async_store.is_some() { - return Err(error); - } - if self.owner_lost(run_id, &owner).await.unwrap_or(false) { - return Ok(()); - } - self.finish_failed(run_id, error.to_string()); - return Err(error); - } - self.finish_completed(run_id, total_spawned as usize); - } else { - let reason = "no runnable phase (dependency deadlock)".to_owned(); - if let Err(error) = self - .persist( - &run, - PersistRequest { - phase_states: run.phase_states.clone(), - child_run_ids: run.child_run_ids.clone(), - status: WorkflowRunStatus::Failed, - summary: Some(reason.clone()), - terminal: true, - }, - &owner, - ) - .await - { - if self.async_store.is_some() { - return Err(error); - } - if self.owner_lost(run_id, &owner).await.unwrap_or(false) { - return Ok(()); - } - self.finish_failed(run_id, error.to_string()); - return Err(error); - } - self.finish_failed(run_id, reason); - } - return Ok(()); - }; - self.emit( - run_id, - tinyagents_graph::GraphEvent::NodeStarted { - node: tinyagents_harness::ids::NodeId::new("run_phase"), - step: total_spawned as usize + 1, - }, - ); - let phase_result = self - .run_phase( - &run, - definition, - &phase, - total_spawned, - cancel.clone(), - &owner, - ) - .await; - let (updated, spawned) = match phase_result { - Ok(result) => result, - Err(error) => { - // A host stop/resume fences this owner with a revision CAS. - // Do not turn that intentional hand-off into a stale - // failure event or overwrite the newer durable state. - if self.async_store.is_some() { - return Err(error); - } - if self.owner_lost(run_id, &owner).await.unwrap_or(false) - || self - .emit_recorded_terminal(run_id, total_spawned as usize) - .await - .unwrap_or(false) - { - return Ok(()); - } - self.finish_failed(run_id, error.to_string()); - return Err(error); - } - }; - run = updated; - self.emit( - run_id, - tinyagents_graph::GraphEvent::NodeCompleted { - node: tinyagents_harness::ids::NodeId::new("run_phase"), - step: total_spawned as usize + 1, - }, - ); - total_spawned += spawned; - if run.status != WorkflowRunStatus::Running { - match run.status { - WorkflowRunStatus::Completed => { - self.finish_completed(run_id, total_spawned as usize) - } - WorkflowRunStatus::Interrupted | WorkflowRunStatus::Cancelled => { - self.finish_cancelled(run_id) - } - WorkflowRunStatus::Failed => self.finish_failed( - run_id, - run.summary - .clone() - .unwrap_or_else(|| "workflow phase failed".to_owned()), - ), - WorkflowRunStatus::Pending | WorkflowRunStatus::Running => {} - } - return Ok(()); - } - } - } - - async fn run_phase( - &self, - run: &WorkflowRun, - definition: &WorkflowDefinition, - phase: &WorkflowPhase, - total_spawned: u32, - cancel: CancellationToken, - owner: &str, - ) -> WorkflowPersistenceResult<(WorkflowRun, u32)> { - let mut phase_states = run.phase_states.clone(); - let mut child_ids = run.child_run_ids.clone(); - set_phase_status(&mut phase_states, &phase.name, PhaseStatus::Running, None); - let running = self - .persist( - run, - PersistRequest { - phase_states: phase_states.clone(), - child_run_ids: child_ids.clone(), - status: WorkflowRunStatus::Running, - summary: None, - terminal: false, - }, - owner, - ) - .await?; - - let budget = definition.max_children.saturating_sub(total_spawned) as usize; - if budget == 0 { - return self - .fail_phase( - &running, - &mut phase_states, - child_ids, - phase, - format!( - "max_children cap ({}) reached before phase '{}' completed", - definition.max_children, phase.name - ), - owner, - ) - .await; - } - let capacity = phase.agent_ids.len().min(budget); - let capped = capacity != phase.agent_ids.len(); - let upstream = upstream_outputs(phase, &phase_states); - let requests = phase.agent_ids[..capacity] - .iter() - .enumerate() - .map(|(index_in_phase, agent_id)| WorkflowChildRequest { - run_id: run.id.clone(), - phase: phase.name.clone(), - agent_id: agent_id.clone(), - index_in_phase, - prompt: phase_prompt(&run.input, phase, index_in_phase, &upstream), - }) - .collect::<Vec<_>>(); - let registration = Arc::new(PhaseRegistration { - store: self.store.clone(), - owner: owner.to_owned(), - run: parking_lot::Mutex::new(running.clone()), - phase_states: phase_states.clone(), - lease_for: self.lease_for, - async_store: self.async_store.clone(), - operation: tokio::sync::Mutex::new(()), - failure: parking_lot::Mutex::new(None), - failure_notify: tokio::sync::Notify::new(), - attempted: parking_lot::Mutex::new(Vec::new()), - }); - let executor = self.executor.clone(); - let async_executor = self.async_executor.clone(); - // Policy/parallel helpers carry textual ordinary child reasons. - // Retain authoritative sources and wake the phase even while an - // unrelated worker never finishes. - let typed_failure = Arc::new(parking_lot::Mutex::new(None)); - let worker_failure = typed_failure.clone(); - let typed_notify = Arc::new(tokio::sync::Notify::new()); - let worker_notify = typed_notify.clone(); - let step_config = self.step_config.clone(); - let worker_cancel = cancel.clone(); - let worker_registration = registration.clone(); - let outcomes = map_reduce( - requests, - ParallelOptions::default() - .with_max_concurrency(definition.default_concurrency as usize) - .with_failure_policy(FailurePolicy::CollectAll) - .with_cancellation(cancel.clone()), - move |_index, request| { - let executor = executor.clone(); - let async_executor = async_executor.clone(); - let typed_failure = worker_failure.clone(); - let typed_notify = worker_notify.clone(); - let cancel = worker_cancel.clone(); - let registration = worker_registration.clone(); - let step_config = step_config.clone(); - async move { - let async_registration = async_executor - .as_ref() - .map(|_| registration.clone() as Arc<dyn AsyncWorkflowChildRegistration>); - run_child_step_authoritative( - &step_config, - executor, - async_executor, - request, - cancel, - registration, - async_registration, - ) - .await - .map_err(|error| { - let message = error.to_string(); - if !matches!(error, WorkflowPersistenceError::Legacy(_)) { - *typed_failure.lock() = Some(error); - typed_notify.notify_one(); - } - tinyagents_harness::TinyAgentsError::Graph(message) - }) - } - }, - ); - tokio::pin!(outcomes); - let renew_every = self.lease_for.div_f32(3.0).max(Duration::from_millis(1)); - let mut heartbeat = tokio::time::interval(renew_every); - // Ignore interval's eager first tick: `claim` has just installed this - // lease, so renewals begin only while a child can be in flight. - heartbeat.tick().await; - let outcomes = loop { - tokio::select! { - outcomes = &mut outcomes => break outcomes, - _ = registration.failure_notify.notified() => { - let error = registration.failure.lock().take().expect("registration notified failure"); - cancel.cancel(); - let mut children = registration.current().child_run_ids; - children.extend(registration.attempted.lock().iter().cloned()); - self.cancel_children(&children).await; - return Err(WorkflowPersistenceError::Registration(error)); - } - _ = typed_notify.notified() => { - let error = typed_failure.lock().take().expect("child notified failure"); - cancel.cancel(); - let mut children = registration.current().child_run_ids; - children.extend(registration.attempted.lock().iter().cloned()); - self.cancel_children(&children).await; - return Err(error); - } - - _ = heartbeat.tick() => { - let renewed = self.renew(&run.id, owner, self.lease_for).await; - if !matches!(renewed, Ok(true)) { - cancel.cancel(); - let mut children = registration.current().child_run_ids; - children.extend(registration.attempted.lock().iter().cloned()); - self.cancel_children(&children).await; - renewed?; - return Err(WorkflowPersistenceError::LeaseRenewalLost); - } - } - } - }; - let registration_error = registration.failure.lock().take(); - let child_error = typed_failure.lock().take(); - if let Some(error) = registration_error - .map(WorkflowPersistenceError::Registration) - .or(child_error) - { - let mut children = registration.current().child_run_ids; - children.extend(registration.attempted.lock().iter().cloned()); - self.cancel_children(&children).await; - return Err(error); - } - let outcomes = match outcomes { - Ok(outcomes) => outcomes, - Err(tinyagents_harness::TinyAgentsError::Cancelled) => { - let children = registration.current().child_run_ids; - let mut to_cancel = children.clone(); - to_cancel.extend(registration.attempted.lock().iter().cloned()); - self.cancel_children(&to_cancel).await; - reset_running_phases( - &mut phase_states, - "workflow interrupted; phase will retry on resume", - ); - let updated = self - .persist( - ®istration.current(), - PersistRequest { - phase_states, - child_run_ids: children, - status: WorkflowRunStatus::Interrupted, - summary: None, - terminal: false, - }, - owner, - ) - .await?; - return Ok((updated, 0)); - } - Err(error) => return Err(OrchestrationError(error.to_string()).into()), - }; - child_ids = registration.current().child_run_ids; - let mut outputs = Vec::new(); - let mut failure = None; - let mut spawned = 0_u32; - for outcome in outcomes.outcomes { - match outcome.result { - Ok((result, annotations)) => { - spawned += 1; - // The executor registered the real id before it could - // await completion. Keep older executors harmlessly - // compatible by accepting an already-present id only. - if !child_ids.iter().any(|id| id == &result.child_id) { - if self.async_store.is_some() { - self.cancel_children(std::slice::from_ref(&result.child_id)) - .await; - return Err(WorkflowPersistenceError::UnregisteredChild( - result.child_id, - )); - } - child_ids.push(result.child_id.clone()); - } - let mut entry = json!({ - // Preserve the persisted/RPC v1 projection while the - // v2 metadata remains lossless for engine consumers. - "agentId": phase.agent_ids[outcome.index], - "output": render_compat_output(&result.output), - "metadata": { "version": 2, "rawOutput": result.output }, - }); - if let Some(annotations) = annotations { - entry["metadata"]["resultPolicy"] = annotations; - } - outputs.push(entry); - } - Err(error) if failure.is_none() => failure = Some(error), - Err(_) => {} - } - } - if cancel.is_cancelled() { - let children = registration.current().child_run_ids; - let mut to_cancel = children.clone(); - to_cancel.extend(registration.attempted.lock().iter().cloned()); - self.cancel_children(&to_cancel).await; - reset_running_phases( - &mut phase_states, - "workflow interrupted; phase will retry on resume", - ); - let updated = self - .persist( - ®istration.current(), - PersistRequest { - phase_states, - child_run_ids: children, - status: WorkflowRunStatus::Interrupted, - summary: None, - terminal: false, - }, - owner, - ) - .await?; - return Ok((updated, 0)); - } - if let Some(reason) = failure.or_else(|| { - capped.then(|| { - format!( - "max_children cap ({}) reached before phase '{}' completed", - definition.max_children, phase.name - ) - }) - }) { - return self - .fail_phase( - ®istration.current(), - &mut phase_states, - child_ids, - phase, - reason, - owner, - ) - .await; - } - set_phase_status( - &mut phase_states, - &phase.name, - PhaseStatus::Completed, - Some(Value::Array(outputs)), - ); - let updated = self - .persist( - ®istration.current(), - PersistRequest { - phase_states, - child_run_ids: child_ids, - status: WorkflowRunStatus::Running, - summary: None, - terminal: false, - }, - owner, - ) - .await?; - Ok((updated, spawned)) - } - - async fn fail_phase( - &self, - run: &WorkflowRun, - phase_states: &mut Value, - child_ids: Vec<String>, - phase: &WorkflowPhase, - reason: String, - owner: &str, - ) -> WorkflowPersistenceResult<(WorkflowRun, u32)> { - set_phase_status( - phase_states, - &phase.name, - PhaseStatus::Failed, - Some(json!([])), - ); - set_phase_reason(phase_states, &phase.name, &reason); - let updated = self - .persist( - run, - PersistRequest { - phase_states: phase_states.clone(), - child_run_ids: child_ids, - status: WorkflowRunStatus::Failed, - summary: Some(reason), - terminal: true, - }, - owner, - ) - .await?; - Ok((updated, 0)) - } - - async fn persist( - &self, - run: &WorkflowRun, - request: PersistRequest, - owner: &str, - ) -> WorkflowPersistenceResult<WorkflowRun> { - self.compare_and_swap( - WorkflowRunUpsert { - id: run.id.clone(), - definition_id: run.definition_id.clone(), - parent_thread_id: run.parent_thread_id.clone(), - input: run.input.clone(), - phase_states: request.phase_states, - child_run_ids: request.child_run_ids, - status: request.status, - summary: request.summary, - started_at: Some(run.started_at), - completed_at: request.terminal.then(Utc::now), - }, - run.revision, - owner, - self.lease_for, - ) - .await? - .ok_or_else(|| WorkflowPersistenceError::LeaseLost) - } - - fn emit(&self, run_id: &str, event: tinyagents_graph::GraphEvent) { - if let Some(sink) = &self.event_sink { - sink.emit(tinyagents_graph::GraphEventEnvelope { - run_id: tinyagents_harness::ids::RunId::new(run_id), - task_id: None, - ns: Vec::new(), - seq: self.event_seq.fetch_add(1, Ordering::Relaxed), - event, - }); - } - } - - fn finish_completed(&self, run_id: &str, steps: usize) { - self.emit( - run_id, - tinyagents_graph::GraphEvent::RunCompleted { - run_id: tinyagents_harness::ids::RunId::new(run_id), - steps, - }, - ); - self.flush_terminal_events(); - } - - fn finish_failed(&self, run_id: &str, error: String) { - self.emit( - run_id, - tinyagents_graph::GraphEvent::RunFailed { - run_id: tinyagents_harness::ids::RunId::new(run_id), - error, - }, - ); - self.flush_terminal_events(); - } - - fn finish_cancelled(&self, run_id: &str) { - self.emit( - run_id, - tinyagents_graph::GraphEvent::RunCancelled { - run_id: tinyagents_harness::ids::RunId::new(run_id), - }, - ); - self.flush_terminal_events(); - } - - fn flush_terminal_events(&self) { - if let Some(sink) = &self.event_sink { - sink.flush(); - } - } - - /// A lifecycle hand-off or lease takeover has fenced this driver. It must - /// not manufacture a terminal graph event for the replacement owner. - async fn owner_lost(&self, run_id: &str, owner: &str) -> WorkflowPersistenceResult<bool> { - Ok(self.load(run_id).await?.is_some_and(|current| { - current.lease_owner.as_deref() != Some(owner) - || current - .lease_expires_at - .is_none_or(|expires| expires <= Utc::now()) - })) - } - - async fn emit_recorded_terminal( - &self, - run_id: &str, - steps: usize, - ) -> WorkflowPersistenceResult<bool> { - let Some(current) = self.load(run_id).await? else { - return Ok(false); - }; - if !current.status.is_terminal() { - return Ok(false); - } - match current.status { - WorkflowRunStatus::Completed => self.finish_completed(run_id, steps), - WorkflowRunStatus::Interrupted | WorkflowRunStatus::Cancelled => { - self.finish_cancelled(run_id) - } - WorkflowRunStatus::Failed => self.finish_failed( - run_id, - current - .summary - .unwrap_or_else(|| "workflow phase failed".to_owned()), - ), - WorkflowRunStatus::Pending | WorkflowRunStatus::Running => return Ok(false), - } - Ok(true) - } -} - -pub(super) fn render_compat_output(output: &Value) -> String { - match output { - Value::String(text) => text.clone(), - _ => serde_json::to_string(output) - .unwrap_or_else(|_| "<unserializable workflow output>".to_owned()), - } -} - -/// Placeholder legacy types for an explicitly authoritative engine. -/// These callbacks always refuse; they never open a database. -pub struct AuthoritativeWorkflow; - -impl WorkflowStore for AuthoritativeWorkflow { - fn load(&self, _: &str) -> Result<Option<WorkflowRun>, OrchestrationError> { - Err(OrchestrationError("async store required".into())) - } - fn upsert(&self, _: WorkflowRunUpsert) -> Result<WorkflowRun, OrchestrationError> { - Err(OrchestrationError("async store required".into())) - } - fn claim( - &self, - _: &str, - _: &str, - _: Duration, - ) -> Result<WorkflowLeaseClaim, OrchestrationError> { - Err(OrchestrationError("async store required".into())) - } - fn compare_and_swap( - &self, - _: WorkflowRunUpsert, - _: u64, - _: &str, - _: Duration, - ) -> Result<Option<WorkflowRun>, OrchestrationError> { - Err(OrchestrationError("async store required".into())) - } - fn renew(&self, _: &str, _: &str, _: Duration) -> Result<bool, OrchestrationError> { - Err(OrchestrationError("async store required".into())) - } -} -#[async_trait] -impl WorkflowExecutor for AuthoritativeWorkflow { - async fn execute( - &self, - _: WorkflowChildRequest, - _: CancellationToken, - _: Arc<dyn WorkflowChildRegistration>, - ) -> Result<WorkflowChildResult, OrchestrationError> { - Err(OrchestrationError("async executor required".into())) - } - async fn cancel_children(&self, _: &[String]) {} -} -impl WorkflowEngine<AuthoritativeWorkflow, AuthoritativeWorkflow> { - /// Select only the supplied asynchronous authority and child executor. - /// No legacy persistence callback or workspace fallback is installed. - pub fn new_async( - store: Arc<dyn AsyncWorkflowStore>, - executor: Arc<dyn AsyncWorkflowExecutor>, - ) -> Self { - let mut engine = Self::new( - Arc::new(AuthoritativeWorkflow), - Arc::new(AuthoritativeWorkflow), - ); - engine.async_store = Some(store); - engine.async_executor = Some(executor); - engine - } -} diff --git a/crates/tinyagents-orchestration/src/workflow/engine_async_tests.rs b/crates/tinyagents-orchestration/src/workflow/engine_async_tests.rs deleted file mode 100644 index 4191f19bf..000000000 --- a/crates/tinyagents-orchestration/src/workflow/engine_async_tests.rs +++ /dev/null @@ -1,520 +0,0 @@ -use super::*; -use async_trait::async_trait; -use serde_json::json; -use std::sync::Arc; -use tinyagents_harness::CancellationToken; -use tinyagents_session::run_ledger::{RunLedgerStore, WorkflowRunStatus}; -use tinystoragedrivers_testkit::MemoryBackend; - -struct RegisteredExecutor; -#[async_trait] -impl AsyncWorkflowExecutor for RegisteredExecutor { - async fn execute_async( - &self, - request: WorkflowChildRequest, - _: CancellationToken, - registration: Arc<dyn AsyncWorkflowChildRegistration>, - ) -> WorkflowPersistenceResult<WorkflowChildResult> { - let id = format!("{}-{}", request.phase, request.index_in_phase); - registration.register_async(id.clone()).await?; - Ok(WorkflowChildResult { - child_id: id, - output: json!({"ok":true}), - }) - } - async fn cancel_children(&self, _: &[String]) {} -} - -#[tokio::test(flavor = "current_thread")] -async fn authoritative_engine_awaits_registration_and_completes_on_current_thread() { - let ledger_backend = MemoryBackend::new(); - let _workspace = tempfile::tempdir().unwrap(); - let ledger = Arc::new(crate::test_support::ledger_in(&ledger_backend)); - let engine = WorkflowEngine::new_async( - Arc::new(LedgerWorkflowStore::new(ledger.clone())), - Arc::new(RegisteredExecutor), - ); - let definition = tests::definition(); - engine - .initialise_async("run".into(), &definition, json!({}), None) - .await - .unwrap(); - engine - .drive_async("run", &definition, CancellationToken::new()) - .await - .unwrap(); - let run = ledger.get_workflow_run("run").await.unwrap().unwrap(); - assert_eq!(run.status, WorkflowRunStatus::Completed); - assert_eq!(run.child_run_ids.len(), 4); -} - -use std::sync::atomic::{AtomicUsize, Ordering}; -use std::time::Duration; -use tinyagents_session::run_ledger::{WorkflowLeaseClaim, WorkflowRun, WorkflowRunUpsert}; - -fn refusal() -> WorkflowPersistenceError { - WorkflowPersistenceError::Ledger( - tinyagents_storage_kit::StorageError::new( - tinyagents_storage_kit::ErrorKind::Unavailable, - "registration backend unavailable", - ) - .into(), - ) -} - -struct RegistrationRefusal { - inner: LedgerWorkflowStore, - cas: AtomicUsize, -} -#[async_trait] -impl AsyncWorkflowStore for RegistrationRefusal { - async fn load(&self, id: &str) -> WorkflowPersistenceResult<Option<WorkflowRun>> { - self.inner.load(id).await - } - async fn upsert(&self, row: WorkflowRunUpsert) -> WorkflowPersistenceResult<WorkflowRun> { - self.inner.upsert(row).await - } - async fn claim( - &self, - id: &str, - owner: &str, - lease: Duration, - ) -> WorkflowPersistenceResult<WorkflowLeaseClaim> { - self.inner.claim(id, owner, lease).await - } - async fn renew( - &self, - id: &str, - owner: &str, - lease: Duration, - ) -> WorkflowPersistenceResult<bool> { - self.inner.renew(id, owner, lease).await - } - async fn compare_and_swap( - &self, - row: WorkflowRunUpsert, - rev: u64, - owner: &str, - lease: Duration, - ) -> WorkflowPersistenceResult<Option<WorkflowRun>> { - if self.cas.fetch_add(1, Ordering::SeqCst) >= 1 { - return Err(refusal()); - } - self.inner.compare_and_swap(row, rev, owner, lease).await - } -} - -#[derive(Default)] -struct IgnoreFailedRegistration { - cancelled: parking_lot::Mutex<Vec<String>>, -} -#[async_trait] -impl AsyncWorkflowExecutor for IgnoreFailedRegistration { - async fn execute_async( - &self, - _: WorkflowChildRequest, - _: CancellationToken, - registration: Arc<dyn AsyncWorkflowChildRegistration>, - ) -> WorkflowPersistenceResult<WorkflowChildResult> { - assert!( - registration - .register_async("real-child".into()) - .await - .is_err() - ); - std::future::pending().await - } - async fn cancel_children(&self, ids: &[String]) { - self.cancelled.lock().extend_from_slice(ids); - } -} - -fn innermost(error: &WorkflowPersistenceError) -> &WorkflowPersistenceError { - match error { - WorkflowPersistenceError::Registration(error) => innermost(error), - other => other, - } -} - -#[tokio::test(flavor = "current_thread", start_paused = true)] -async fn ignored_registration_refusal_cancels_child_and_preserves_source_without_waiting() { - let ledger_backend = MemoryBackend::new(); - let _workspace = tempfile::tempdir().unwrap(); - let ledger = Arc::new(crate::test_support::ledger_in(&ledger_backend)); - let store = Arc::new(RegistrationRefusal { - inner: LedgerWorkflowStore::new(ledger.clone()), - cas: AtomicUsize::new(0), - }); - let executor = Arc::new(IgnoreFailedRegistration::default()); - let engine = WorkflowEngine::new_async(store, executor.clone()); - let mut definition = tests::definition(); - definition.phases.truncate(1); - engine - .initialise_async("run".into(), &definition, json!({}), None) - .await - .unwrap(); - let result = tokio::time::timeout( - Duration::from_secs(1), - engine.drive_async("run", &definition, CancellationToken::new()), - ) - .await; - let error = result - .expect("registration refusal must interrupt an executor that ignores it") - .unwrap_err(); - assert!(matches!( - innermost(&error), - WorkflowPersistenceError::Ledger(_) - )); - assert!( - error - .to_string() - .contains("registration backend unavailable") - ); - assert_eq!(*executor.cancelled.lock(), vec!["real-child"]); - assert_ne!( - ledger - .get_workflow_run("run") - .await - .unwrap() - .unwrap() - .status, - WorkflowRunStatus::Completed - ); -} - -struct UnregisteredExecutor; -#[async_trait] -impl AsyncWorkflowExecutor for UnregisteredExecutor { - async fn execute_async( - &self, - _: WorkflowChildRequest, - _: CancellationToken, - _: Arc<dyn AsyncWorkflowChildRegistration>, - ) -> WorkflowPersistenceResult<WorkflowChildResult> { - Ok(WorkflowChildResult { - child_id: "unregistered".into(), - output: json!("uncommitted"), - }) - } - async fn cancel_children(&self, _: &[String]) {} -} - -#[tokio::test(flavor = "current_thread")] -async fn authoritative_success_cannot_substitute_for_committed_child_registration() { - let ledger_backend = MemoryBackend::new(); - let _workspace = tempfile::tempdir().unwrap(); - let ledger = Arc::new(crate::test_support::ledger_in(&ledger_backend)); - let engine = WorkflowEngine::new_async( - Arc::new(LedgerWorkflowStore::new(ledger.clone())), - Arc::new(UnregisteredExecutor), - ); - let definition = tests::definition(); - engine - .initialise_async("run".into(), &definition, json!({}), None) - .await - .unwrap(); - assert!( - matches!(engine.drive_async("run",&definition,CancellationToken::new()).await, - Err(WorkflowPersistenceError::UnregisteredChild(id)) if id=="unregistered") - ); - let run = ledger.get_workflow_run("run").await.unwrap().unwrap(); - assert!(run.child_run_ids.is_empty()); - assert_ne!(run.status, WorkflowRunStatus::Completed); - assert!( - engine - .initialise("legacy".into(), &definition, json!({}), None) - .is_err() - ); - assert!( - engine - .drive("run", &definition, CancellationToken::new()) - .await - .is_err() - ); -} - -struct RefusingChild; -#[async_trait] -impl AsyncWorkflowExecutor for RefusingChild { - async fn execute_async( - &self, - request: WorkflowChildRequest, - _: CancellationToken, - registration: Arc<dyn AsyncWorkflowChildRegistration>, - ) -> WorkflowPersistenceResult<WorkflowChildResult> { - if request.index_in_phase == 0 { - registration.register_async("refused-child".into()).await?; - return Err(refusal()); - } - std::future::pending().await - } - async fn cancel_children(&self, _: &[String]) {} -} - -#[tokio::test(flavor = "current_thread", start_paused = true)] -async fn typed_child_failure_interrupts_parallel_waiters_without_string_flattening() { - let ledger_backend = MemoryBackend::new(); - let _workspace = tempfile::tempdir().unwrap(); - let ledger = Arc::new(crate::test_support::ledger_in(&ledger_backend)); - let engine = WorkflowEngine::new_async( - Arc::new(LedgerWorkflowStore::new(ledger.clone())), - Arc::new(RefusingChild), - ); - let mut definition = tests::definition(); - definition.phases.truncate(1); - definition.phases[0].agent_ids = vec!["refused".into(), "waiting".into()]; - engine - .initialise_async("run".into(), &definition, json!({}), None) - .await - .unwrap(); - let error = tokio::time::timeout( - Duration::from_secs(1), - engine.drive_async("run", &definition, CancellationToken::new()), - ) - .await - .expect("typed child error cannot wait for unrelated child") - .unwrap_err(); - assert!(matches!( - innermost(&error), - WorkflowPersistenceError::Ledger(_) - )); - assert!( - error - .to_string() - .contains("registration backend unavailable") - ); - assert_ne!( - ledger - .get_workflow_run("run") - .await - .unwrap() - .unwrap() - .status, - WorkflowRunStatus::Completed - ); -} - -struct SyncRegistrationAttempt; -#[async_trait] -impl AsyncWorkflowExecutor for SyncRegistrationAttempt { - async fn execute_async( - &self, - _: WorkflowChildRequest, - _: CancellationToken, - registration: Arc<dyn AsyncWorkflowChildRegistration>, - ) -> WorkflowPersistenceResult<WorkflowChildResult> { - assert!(registration.register("sync-child".into()).is_err()); - std::future::pending().await - } - async fn cancel_children(&self, _: &[String]) {} -} - -#[tokio::test(flavor = "current_thread", start_paused = true)] -async fn ignored_synchronous_registration_refusal_cannot_leave_authoritative_driver_waiting() { - let ledger_backend = MemoryBackend::new(); - let _workspace = tempfile::tempdir().unwrap(); - let store = Arc::new(LedgerWorkflowStore::new(Arc::new( - crate::test_support::ledger_in(&ledger_backend), - ))); - let engine = WorkflowEngine::new_async(store, Arc::new(SyncRegistrationAttempt)); - let definition = tests::definition(); - engine - .initialise_async("run".into(), &definition, json!({}), None) - .await - .unwrap(); - let error = tokio::time::timeout( - Duration::from_secs(1), - engine.drive_async("run", &definition, CancellationToken::new()), - ) - .await - .expect("synchronous registration refused without blocking") - .unwrap_err(); - assert!(matches!( - innermost(&error), - WorkflowPersistenceError::AsyncRegistrationRequired - )); -} - -#[path = "engine_driver_tests.rs"] -mod driver_tests; - -struct RefusedLegacyRegistration; -impl WorkflowChildRegistration for RefusedLegacyRegistration { - fn register(&self, _: String) -> Result<(), OrchestrationError> { - Err(OrchestrationError("registration refused".into())) - } -} - -#[tokio::test(flavor = "current_thread")] -async fn authoritative_placeholder_callbacks_cannot_supply_legacy_fallback() { - let marker = AuthoritativeWorkflow; - let row = WorkflowRunUpsert { - id: "run".into(), - definition_id: "definition".into(), - parent_thread_id: None, - input: json!({}), - phase_states: json!({}), - child_run_ids: vec![], - status: WorkflowRunStatus::Running, - summary: None, - started_at: None, - completed_at: None, - }; - assert!(marker.load("run").is_err()); - assert!(marker.upsert(row.clone()).is_err()); - assert!( - marker - .claim("run", "owner", Duration::from_secs(30)) - .is_err() - ); - assert!( - marker - .compare_and_swap(row, 0, "owner", Duration::from_secs(30)) - .is_err() - ); - assert!( - marker - .renew("run", "owner", Duration::from_secs(30)) - .is_err() - ); - let request = WorkflowChildRequest { - run_id: "run".into(), - phase: "phase".into(), - agent_id: "agent".into(), - index_in_phase: 0, - prompt: String::new(), - }; - assert!( - marker - .execute( - request, - CancellationToken::new(), - Arc::new(RefusedLegacyRegistration) - ) - .await - .is_err() - ); - WorkflowExecutor::cancel_children(&marker, &[]).await; -} - -struct PendingRegistration { - inner: tests::MemoryStore, - started: tokio::sync::Notify, - pending: std::sync::atomic::AtomicBool, - refuse_renewal: bool, -} -#[async_trait] -impl AsyncWorkflowStore for PendingRegistration { - async fn load(&self, id: &str) -> WorkflowPersistenceResult<Option<WorkflowRun>> { - Ok(self.inner.load(id)?) - } - async fn upsert(&self, row: WorkflowRunUpsert) -> WorkflowPersistenceResult<WorkflowRun> { - Ok(self.inner.upsert(row)?) - } - async fn claim( - &self, - id: &str, - owner: &str, - lease: Duration, - ) -> WorkflowPersistenceResult<WorkflowLeaseClaim> { - Ok(self.inner.claim(id, owner, lease)?) - } - async fn renew( - &self, - id: &str, - owner: &str, - lease: Duration, - ) -> WorkflowPersistenceResult<bool> { - if self.refuse_renewal && self.pending.load(Ordering::SeqCst) { - return Err(refusal()); - } - Ok(self.inner.renew(id, owner, lease)?) - } - async fn compare_and_swap( - &self, - row: WorkflowRunUpsert, - rev: u64, - owner: &str, - lease: Duration, - ) -> WorkflowPersistenceResult<Option<WorkflowRun>> { - if !row.child_run_ids.is_empty() { - self.pending.store(true, Ordering::SeqCst); - self.started.notify_one(); - return std::future::pending().await; - } - Ok(self.inner.compare_and_swap(row, rev, owner, lease)?) - } -} -#[derive(Default)] -struct PendingChild { - cancelled: parking_lot::Mutex<Vec<String>>, -} -#[async_trait] -impl AsyncWorkflowExecutor for PendingChild { - async fn execute_async( - &self, - _: WorkflowChildRequest, - _: CancellationToken, - registration: Arc<dyn AsyncWorkflowChildRegistration>, - ) -> WorkflowPersistenceResult<WorkflowChildResult> { - registration.register_async("pending-child".into()).await?; - unreachable!("registration remains pending") - } - async fn cancel_children(&self, ids: &[String]) { - self.cancelled.lock().extend_from_slice(ids); - } -} -async fn pending_child_case(renewal: bool) { - let store = Arc::new(PendingRegistration { - inner: tests::MemoryStore::default(), - started: tokio::sync::Notify::new(), - pending: std::sync::atomic::AtomicBool::new(false), - refuse_renewal: renewal, - }); - let executor = Arc::new(PendingChild::default()); - let engine = WorkflowEngine::new_async(store.clone(), executor.clone()) - .with_lease_duration(Duration::from_millis(30)); - let mut definition = tests::definition(); - definition.phases.truncate(1); - definition.phases[0].agent_ids.truncate(1); - engine - .initialise_async("run".into(), &definition, json!({}), None) - .await - .unwrap(); - let cancel = CancellationToken::new(); - let control = async { - store.started.notified().await; - if !renewal { - cancel.cancel(); - } - }; - let drive = engine.drive_async("run", &definition, cancel.clone()); - let (result, _) = tokio::time::timeout(Duration::from_secs(1), async { - tokio::join!(drive, control) - }) - .await - .unwrap(); - if renewal { - assert!(matches!( - innermost(&result.unwrap_err()), - WorkflowPersistenceError::Ledger(_) - )); - } else { - result.unwrap(); - let run = store.inner.load("run").unwrap().unwrap(); - assert_eq!(run.status, WorkflowRunStatus::Interrupted); - assert!( - run.child_run_ids.is_empty(), - "uncommitted ids are cancellation targets, not persisted children" - ); - } - assert_eq!(*executor.cancelled.lock(), vec!["pending-child"]); -} -#[tokio::test(flavor = "current_thread", start_paused = true)] -async fn heartbeat_loss_cancels_a_child_with_pending_registration() { - pending_child_case(true).await; -} -#[tokio::test(flavor = "current_thread", start_paused = true)] -async fn host_cancellation_cancels_a_child_with_pending_registration() { - pending_child_case(false).await; -} diff --git a/crates/tinyagents-orchestration/src/workflow/engine_driver_tests.rs b/crates/tinyagents-orchestration/src/workflow/engine_driver_tests.rs deleted file mode 100644 index 7ec9952e8..000000000 --- a/crates/tinyagents-orchestration/src/workflow/engine_driver_tests.rs +++ /dev/null @@ -1,314 +0,0 @@ -use super::*; -use tinyagents_session::run_ledger::{DriverRunLedger, LedgerError}; -use tinyagents_storage_kit::{ErrorKind, Scope, StorageBackend}; -use tinystoragedrivers_testkit::MemoryBackend; - -fn native_kind(error: &WorkflowPersistenceError) -> ErrorKind { - match innermost(error) { - WorkflowPersistenceError::Ledger(LedgerError::Driver(error)) => error.kind(), - other => panic!("native source flattened: {other:?}"), - } -} - -#[tokio::test(flavor = "current_thread")] -async fn independent_scoped_drivers_keep_complete_workflow_runs_separate() { - let backend = MemoryBackend::new(); - let a = Arc::new(DriverRunLedger::new( - backend - .for_scope(&Scope::new("profile-a").unwrap()) - .unwrap(), - )); - let b = Arc::new(DriverRunLedger::new( - backend - .for_scope(&Scope::new("profile-b").unwrap()) - .unwrap(), - )); - let engine = WorkflowEngine::new_async( - Arc::new(LedgerWorkflowStore::new(a.clone())), - Arc::new(RegisteredExecutor), - ); - let definition = tests::definition(); - engine - .initialise_async("run".into(), &definition, json!({}), None) - .await - .unwrap(); - engine - .drive_async("run", &definition, CancellationToken::new()) - .await - .unwrap(); - assert_eq!( - a.get_workflow_run("run").await.unwrap().unwrap().status, - WorkflowRunStatus::Completed - ); - assert!(b.get_workflow_run("run").await.unwrap().is_none()); -} - -#[tokio::test(flavor = "current_thread", start_paused = true)] -async fn registration_and_child_failures_retain_native_unavailable_kind() { - let ledger_backend = MemoryBackend::new(); - // Invoke both owning regressions under the driver's native error variant. - let error = refusal(); - assert_eq!(native_kind(&error), ErrorKind::Unavailable); - let _workspace = tempfile::tempdir().unwrap(); - let ledger = Arc::new(crate::test_support::ledger_in(&ledger_backend)); - let store = Arc::new(RegistrationRefusal { - inner: LedgerWorkflowStore::new(ledger), - cas: AtomicUsize::new(0), - }); - let engine = WorkflowEngine::new_async(store, Arc::new(IgnoreFailedRegistration::default())); - let definition = tests::definition(); - engine - .initialise_async("run".into(), &definition, json!({}), None) - .await - .unwrap(); - let error = engine - .drive_async("run", &definition, CancellationToken::new()) - .await - .unwrap_err(); - assert_eq!(native_kind(&error), ErrorKind::Unavailable); -} - -#[derive(Default)] -struct HeldChild { - started: tokio::sync::Notify, - release: tokio::sync::Notify, - executions: AtomicUsize, -} -#[async_trait] -impl AsyncWorkflowExecutor for HeldChild { - async fn execute_async( - &self, - _: WorkflowChildRequest, - _: CancellationToken, - registration: Arc<dyn AsyncWorkflowChildRegistration>, - ) -> WorkflowPersistenceResult<WorkflowChildResult> { - self.executions.fetch_add(1, Ordering::SeqCst); - registration.register_async("held-child".into()).await?; - self.started.notify_one(); - self.release.notified().await; - Ok(WorkflowChildResult { - child_id: "held-child".into(), - output: json!("done"), - }) - } - async fn cancel_children(&self, _: &[String]) {} -} - -#[tokio::test(flavor = "current_thread")] -async fn independent_authoritative_engines_acquire_only_one_live_owner() { - let backend = MemoryBackend::new(); - let scope = Scope::new("profile-a").unwrap(); - let first = Arc::new(DriverRunLedger::new(backend.for_scope(&scope).unwrap())); - let second = Arc::new(DriverRunLedger::new(backend.for_scope(&scope).unwrap())); - let executor = Arc::new(HeldChild::default()); - let a = Arc::new(WorkflowEngine::new_async( - Arc::new(LedgerWorkflowStore::new(first.clone())), - executor.clone(), - )); - let b = WorkflowEngine::new_async(Arc::new(LedgerWorkflowStore::new(second)), executor.clone()); - let mut definition = tests::definition(); - definition.phases.truncate(1); - a.initialise_async("run".into(), &definition, json!({}), None) - .await - .unwrap(); - let task = tokio::spawn({ - let a = a.clone(); - let definition = definition.clone(); - async move { - a.drive_async("run", &definition, CancellationToken::new()) - .await - } - }); - executor.started.notified().await; - b.drive_async("run", &definition, CancellationToken::new()) - .await - .unwrap(); - assert_eq!(executor.executions.load(Ordering::SeqCst), 1); - executor.release.notify_one(); - task.await.unwrap().unwrap(); - assert_eq!( - first.get_workflow_run("run").await.unwrap().unwrap().status, - WorkflowRunStatus::Completed - ); -} - -struct RefusedAuthority(ErrorKind); -impl RefusedAuthority { - fn error<T>(&self) -> WorkflowPersistenceResult<T> { - Err(WorkflowPersistenceError::Ledger( - tinyagents_storage_kit::StorageError::new(self.0, "configured authority refused") - .into(), - )) - } -} -#[async_trait] -impl AsyncWorkflowStore for RefusedAuthority { - async fn load(&self, _: &str) -> WorkflowPersistenceResult<Option<WorkflowRun>> { - self.error() - } - async fn upsert(&self, _: WorkflowRunUpsert) -> WorkflowPersistenceResult<WorkflowRun> { - self.error() - } - async fn claim( - &self, - _: &str, - _: &str, - _: Duration, - ) -> WorkflowPersistenceResult<WorkflowLeaseClaim> { - self.error() - } - async fn renew(&self, _: &str, _: &str, _: Duration) -> WorkflowPersistenceResult<bool> { - self.error() - } - async fn compare_and_swap( - &self, - _: WorkflowRunUpsert, - _: u64, - _: &str, - _: Duration, - ) -> WorkflowPersistenceResult<Option<WorkflowRun>> { - self.error() - } -} - -#[tokio::test(flavor = "current_thread")] -async fn configured_refusal_never_uses_legacy_initialise_or_drive() { - let ambient_db = std::env::current_dir() - .unwrap() - .join("session_db/sessions.db"); - let existed_before = ambient_db.exists(); - for kind in [ - ErrorKind::Unavailable, - ErrorKind::Conflict, - ErrorKind::Unsupported(tinyagents_storage_kit::Capability::Transactions), - ] { - let engine = WorkflowEngine::new_async( - Arc::new(RefusedAuthority(kind)), - Arc::new(RegisteredExecutor), - ); - let definition = tests::definition(); - assert_eq!( - native_kind( - &engine - .initialise_async("run".into(), &definition, json!({}), None) - .await - .unwrap_err() - ), - kind - ); - assert_eq!( - native_kind( - &engine - .drive_async("run", &definition, CancellationToken::new()) - .await - .unwrap_err() - ), - kind - ); - assert!( - engine - .initialise("legacy".into(), &definition, json!({}), None) - .is_err() - ); - assert!( - engine - .drive("legacy", &definition, CancellationToken::new()) - .await - .is_err() - ); - } - assert_eq!( - ambient_db.exists(), - existed_before, - "configured refusal created an ambient SQLite fallback" - ); -} - -struct WorkflowTakeoverChild { - ledger: Arc<DriverRunLedger>, - replacement: parking_lot::Mutex<Option<WorkflowRun>>, - cancelled: parking_lot::Mutex<Vec<String>>, -} -#[async_trait] -impl AsyncWorkflowExecutor for WorkflowTakeoverChild { - async fn execute_async( - &self, - _: WorkflowChildRequest, - _: CancellationToken, - registration: Arc<dyn AsyncWorkflowChildRegistration>, - ) -> WorkflowPersistenceResult<WorkflowChildResult> { - registration.register_async("takeover-child".into()).await?; - let current = self.ledger.get_workflow_run("run").await.unwrap().unwrap(); - let revision = current.revision; - let next = WorkflowRunUpsert { - id: current.id, - definition_id: current.definition_id, - parent_thread_id: current.parent_thread_id, - input: current.input, - phase_states: current.phase_states, - child_run_ids: current.child_run_ids, - status: WorkflowRunStatus::Running, - summary: None, - started_at: Some(current.started_at), - completed_at: None, - }; - self.ledger - .compare_and_swap_workflow_run_lifecycle(next, revision) - .await - .unwrap() - .unwrap(); - let WorkflowLeaseClaim::Acquired(replacement) = self - .ledger - .try_claim_workflow_run("run", "replacement-owner", chrono::Duration::minutes(1)) - .await - .unwrap() - else { - panic!("replacement claim refused") - }; - *self.replacement.lock() = Some(replacement); - std::future::pending().await - } - async fn cancel_children(&self, ids: &[String]) { - self.cancelled.lock().extend_from_slice(ids); - } -} - -#[tokio::test(flavor = "current_thread", start_paused = true)] -async fn workflow_owner_takeover_returns_typed_loss_without_mutating_replacement() { - let backend = MemoryBackend::new(); - let scope = Scope::new("profile-a").unwrap(); - let ledger = Arc::new(DriverRunLedger::new(backend.for_scope(&scope).unwrap())); - let executor = Arc::new(WorkflowTakeoverChild { - ledger: ledger.clone(), - replacement: parking_lot::Mutex::new(None), - cancelled: parking_lot::Mutex::new(vec![]), - }); - let engine = WorkflowEngine::new_async( - Arc::new(LedgerWorkflowStore::new(ledger.clone())), - executor.clone(), - ) - .with_lease_duration(Duration::from_secs(60)); - let mut definition = tests::definition(); - definition.phases.truncate(1); - engine - .initialise_async("run".into(), &definition, json!({}), None) - .await - .unwrap(); - let error = tokio::time::timeout( - Duration::from_secs(30), - engine.drive_async("run", &definition, CancellationToken::new()), - ) - .await - .expect("stale workflow owner cannot keep renewing") - .unwrap_err(); - assert!( - !matches!(error, WorkflowPersistenceError::Legacy(_)), - "authoritative ownership loss was flattened into the legacy string error" - ); - assert!(error.to_string().contains("lease renewal failed")); - assert_eq!( - ledger.get_workflow_run("run").await.unwrap(), - executor.replacement.lock().clone() - ); - assert_eq!(*executor.cancelled.lock(), vec!["takeover-child"]); -} diff --git a/crates/tinyagents-orchestration/src/workflow/graph.rs b/crates/tinyagents-orchestration/src/workflow/graph.rs deleted file mode 100644 index ab28b7481..000000000 --- a/crates/tinyagents-orchestration/src/workflow/graph.rs +++ /dev/null @@ -1,78 +0,0 @@ -//! Workflow scheduler DAG: phase scheduling, dependency resolution, and -//! topological ordering. -//! -//! Builds a directed acyclic graph representing the workflow's phase -//! dependencies, computes runnable phases, and projects the schedule into -//! the graph layer for topology introspection. - -use anyhow::{Result, anyhow}; -use tinyagents_graph::export::GraphTopology; -use tinyagents_graph::recursion::RecursionPolicy; -use tinyagents_graph::{ - ClosureStateReducer, Command, CompiledGraph, GraphBuilder, NodeContext, NodeResult, -}; - -#[derive(Clone, Default)] -pub(crate) struct SchedulerState; - -pub(crate) enum SchedulerUpdate { - Noop, -} - -/// Structure-only *preview* of the scheduler topology. -/// -/// `WorkflowEngine` is the effectful scheduler. This helper deliberately does -/// not execute that engine; it only supplies a stable topology to diagnostic -/// UIs, so callers must never present it as the graph that ran a workflow. -pub fn scheduler_topology_preview() -> Result<GraphTopology> { - Ok(build_scheduler_graph(1)?.topology()) -} - -pub(crate) fn build_scheduler_graph( - phase_count: usize, -) -> Result<CompiledGraph<SchedulerState, SchedulerUpdate>> { - let mut builder = GraphBuilder::<SchedulerState, SchedulerUpdate>::new().set_reducer( - ClosureStateReducer::new(|state: SchedulerState, update| { - match update { - SchedulerUpdate::Noop => {} - } - Ok(state) - }), - ); - // Effects are supplied at invocation time through the execution context. - // These nodes only make the validated dispatch/run/done topology inspectable; - // `WorkflowEngine` installs its effectful variants below. - builder = builder.add_node( - "dispatch", - |_state: SchedulerState, _context: NodeContext| async move { - Ok(NodeResult::Command(Command::default().with_goto(["done"]))) - }, - ); - let graph = builder - .add_node( - "run_phase", - |_state: SchedulerState, _context: NodeContext| async move { - Ok(NodeResult::Command( - Command::default().with_goto(["dispatch"]), - )) - }, - ) - .add_node( - "done", - |_state: SchedulerState, _context: NodeContext| async move { - Ok(NodeResult::Update(SchedulerUpdate::Noop)) - }, - ) - .set_entry("dispatch") - .mark_command_routing("dispatch") - .mark_command_routing("run_phase") - .set_finish("done") - .compile() - .map_err(|error| anyhow!("workflow scheduler graph compile failed: {error}"))? - .with_recursion_policy(RecursionPolicy { - max_visits_per_node: Some(phase_count + 2), - max_total_steps: (phase_count + 1) * 3 + 16, - ..RecursionPolicy::default() - }); - Ok(graph) -} diff --git a/crates/tinyagents-orchestration/src/workflow/mod.rs b/crates/tinyagents-orchestration/src/workflow/mod.rs deleted file mode 100644 index 2f7c78c94..000000000 --- a/crates/tinyagents-orchestration/src/workflow/mod.rs +++ /dev/null @@ -1,42 +0,0 @@ -//! Durable, host-neutral workflow definitions and execution. -//! -//! The workflow engine owns phase scheduling, bounded fan-out, cancellation, -//! resume semantics, and the JSON phase-state projection. A host supplies a -//! [`WorkflowStore`] and [`WorkflowExecutor`] for the classic path, or selects -//! [`AsyncWorkflowStore`] and [`AsyncWorkflowExecutor`] explicitly. This module -//! has no knowledge of credentials, model selection, host progress, or RPC. - -mod child_step; -mod engine; -mod graph; -mod persistence; -mod state; -mod supervisor; -mod types; -mod validate; - -pub use engine::{ - AuthoritativeWorkflow, OrchestrationError, SessionWorkflowStore, WorkflowChildRegistration, - WorkflowChildRequest, WorkflowChildResult, WorkflowEngine, WorkflowExecutor, WorkflowStore, -}; -pub use graph::scheduler_topology_preview; -pub use persistence::{ - AsyncWorkflowChildRegistration, AsyncWorkflowExecutor, AsyncWorkflowStore, LedgerWorkflowStore, - LegacyWorkflowRegistration, WorkflowPersistenceError, WorkflowPersistenceResult, -}; -pub use state::{ - PhaseStatus, all_phases_completed, init_phase_states, next_runnable_phase, phase_prompt, - phase_status, reset_running_phases, synthesize_summary, upstream_outputs, -}; -pub use supervisor::{WorkflowDriveFuture, WorkflowRunHandle, WorkflowRunSupervisor}; -pub use types::{ - DefinitionError, WorkflowDefinition, WorkflowDefinitionListResponse, WorkflowPhase, -}; -pub use validate::{validate_agents, validate_structure}; - -#[cfg(test)] -mod tests; - -#[cfg(test)] -#[path = "engine_async_tests.rs"] -mod async_tests; diff --git a/crates/tinyagents-orchestration/src/workflow/persistence.rs b/crates/tinyagents-orchestration/src/workflow/persistence.rs deleted file mode 100644 index a7fb6b4f2..000000000 --- a/crates/tinyagents-orchestration/src/workflow/persistence.rs +++ /dev/null @@ -1,214 +0,0 @@ -//! Explicit asynchronous workflow persistence over an already-bound ledger. -//! -//! This contract does not choose a workspace, scope, backend, or fallback. -//! The legacy engine is not switched by constructing this adapter. - -use super::{ - OrchestrationError, WorkflowChildRegistration, WorkflowChildRequest, WorkflowChildResult, -}; -use async_trait::async_trait; -use std::{sync::Arc, time::Duration}; -use tinyagents_harness::CancellationToken; -use tinyagents_session::run_ledger::{ - LedgerError, RunLedgerStore, WorkflowLeaseClaim, WorkflowRun, WorkflowRunUpsert, -}; - -/// Typed persistence or registration failure, retained until the host boundary. -#[derive(Debug, thiserror::Error)] -#[non_exhaustive] -pub enum WorkflowPersistenceError { - /// The authoritative ledger refused or failed an operation. - #[error(transparent)] - Ledger(#[from] LedgerError), - /// An orchestration failure uses the compatible legacy error type. - #[error(transparent)] - Legacy(#[from] OrchestrationError), - /// A std duration cannot be represented by the ledger. - #[error(transparent)] - LeaseDuration(#[from] chrono::OutOfRangeError), - /// Synchronous registration cannot persist to an asynchronous authority. - #[error("authoritative workflow registration requires register_async")] - AsyncRegistrationRequired, - /// A child identity was not committed before its result was accepted. - #[error("workflow child {0} completed without committed registration")] - UnregisteredChild(String), - /// The revision or owner no longer authorizes this transition. - #[error("workflow lease lost before durable state transition")] - LeaseLost, - /// Heartbeat discovered a lease now owned by another driver or expired. - #[error("workflow lease renewal failed; cancelled registered children")] - LeaseRenewalLost, - /// A registration failed; shared ownership retains its exact error source. - #[error(transparent)] - Registration(Arc<WorkflowPersistenceError>), -} - -/// Result returned by the authoritative asynchronous contract. -pub type WorkflowPersistenceResult<T> = Result<T, WorkflowPersistenceError>; - -/// Minimal object-safe asynchronous workflow persistence required by an engine. -#[async_trait] -pub trait AsyncWorkflowStore: Send + Sync { - /// Read one workflow from the admitted authority. - async fn load(&self, id: &str) -> WorkflowPersistenceResult<Option<WorkflowRun>>; - /// Create or merge a workflow, preserving ledger revision semantics. - async fn upsert(&self, row: WorkflowRunUpsert) -> WorkflowPersistenceResult<WorkflowRun>; - /// Acquire the durable lease or return its current claim outcome. - async fn claim( - &self, - id: &str, - owner: &str, - lease_for: Duration, - ) -> WorkflowPersistenceResult<WorkflowLeaseClaim>; - /// Persist only while the expected revision and live owner still match. - async fn compare_and_swap( - &self, - row: WorkflowRunUpsert, - expected_revision: u64, - owner: &str, - lease_for: Duration, - ) -> WorkflowPersistenceResult<Option<WorkflowRun>>; - /// Renew only the matching unexpired owner lease. - async fn renew( - &self, - id: &str, - owner: &str, - lease_for: Duration, - ) -> WorkflowPersistenceResult<bool>; -} - -/// Workflow adapter over a caller-admitted ledger handle. -/// -/// Every operation uses that handle; there is no SQLite fallback, runtime -/// blocking bridge or ambient scope lookup. Select it explicitly through -/// WorkflowEngine::new_async; constructing it does not switch a legacy engine. -pub struct LedgerWorkflowStore { - ledger: Arc<dyn RunLedgerStore>, -} - -impl LedgerWorkflowStore { - /// Retain the explicit authoritative handle selected by the session owner. - pub fn new(ledger: Arc<dyn RunLedgerStore>) -> Self { - Self { ledger } - } -} - -#[async_trait] -impl AsyncWorkflowStore for LedgerWorkflowStore { - async fn load(&self, id: &str) -> WorkflowPersistenceResult<Option<WorkflowRun>> { - Ok(self.ledger.get_workflow_run(id).await?) - } - - async fn upsert(&self, row: WorkflowRunUpsert) -> WorkflowPersistenceResult<WorkflowRun> { - Ok(self.ledger.upsert_workflow_run(row).await?) - } - - async fn claim( - &self, - id: &str, - owner: &str, - lease_for: Duration, - ) -> WorkflowPersistenceResult<WorkflowLeaseClaim> { - Ok(self - .ledger - .try_claim_workflow_run(id, owner, chrono::Duration::from_std(lease_for)?) - .await?) - } - - async fn compare_and_swap( - &self, - row: WorkflowRunUpsert, - expected_revision: u64, - owner: &str, - lease_for: Duration, - ) -> WorkflowPersistenceResult<Option<WorkflowRun>> { - Ok(self - .ledger - .compare_and_swap_workflow_run( - row, - expected_revision, - owner, - chrono::Duration::from_std(lease_for)?, - ) - .await?) - } - - async fn renew( - &self, - id: &str, - owner: &str, - lease_for: Duration, - ) -> WorkflowPersistenceResult<bool> { - Ok(self - .ledger - .renew_workflow_run_lease(id, owner, chrono::Duration::from_std(lease_for)?) - .await?) - } -} - -/// Authoritative registration that an async child executor must await. -/// -/// The synchronous entrypoint refuses by default; a caller cannot mistake a -/// queued or omitted registration for a committed child identity. -#[async_trait] -pub trait AsyncWorkflowChildRegistration: Send + Sync { - /// Commit the child identity before the executor proceeds to waiting. - async fn register_async(&self, child_id: String) -> WorkflowPersistenceResult<()>; - - /// Refuse synchronous use of an asynchronous authority. - fn register(&self, _child_id: String) -> WorkflowPersistenceResult<()> { - Err(WorkflowPersistenceError::AsyncRegistrationRequired) - } -} - -/// Explicit adapter for an existing synchronous desktop registration callback. -pub struct LegacyWorkflowRegistration { - inner: Arc<dyn WorkflowChildRegistration>, -} - -impl LegacyWorkflowRegistration { - /// Preserve a supplied legacy callback without choosing storage. - pub fn new(inner: Arc<dyn WorkflowChildRegistration>) -> Self { - Self { inner } - } -} - -#[async_trait] -impl AsyncWorkflowChildRegistration for LegacyWorkflowRegistration { - async fn register_async(&self, child_id: String) -> WorkflowPersistenceResult<()> { - Ok(self.inner.register(child_id)?) - } - - fn register(&self, child_id: String) -> WorkflowPersistenceResult<()> { - Ok(self.inner.register(child_id)?) - } -} - -#[cfg(test)] -#[path = "persistence_tests.rs"] -mod tests; - -/// Child execution that can await authoritative registration and retain errors. -/// Cloud hosts must separately support child progress/status/event durability. -#[async_trait] -pub trait AsyncWorkflowExecutor: Send + Sync { - /// Register the real child before waiting, then return its terminal result. - async fn execute_async( - &self, - request: WorkflowChildRequest, - cancel: CancellationToken, - registration: Arc<dyn AsyncWorkflowChildRegistration>, - ) -> WorkflowPersistenceResult<WorkflowChildResult>; - /// Cancel real children admitted by this executor. - async fn cancel_children(&self, child_ids: &[String]); -} - -impl WorkflowPersistenceError { - /// Project only at the explicitly legacy API boundary. - pub(super) fn into_legacy(self) -> OrchestrationError { - match self { - Self::Legacy(error) => error, - other => OrchestrationError(other.to_string()), - } - } -} diff --git a/crates/tinyagents-orchestration/src/workflow/persistence_refusal_tests.rs b/crates/tinyagents-orchestration/src/workflow/persistence_refusal_tests.rs deleted file mode 100644 index 76134ff69..000000000 --- a/crates/tinyagents-orchestration/src/workflow/persistence_refusal_tests.rs +++ /dev/null @@ -1,134 +0,0 @@ -use super::*; -use serde_json::Value; -use tinyagents_session::run_ledger::DriverRunLedger; -use tinyagents_storage_kit::{ - Capabilities, Capability, CollectionSpec, DocumentStore, ErrorKind, Filter, Page, Precondition, - Query, Scope, ScopedStorage, Sort, Storage, StorageBackend, StorageError, Version, Versioned, -}; -use tinystoragedrivers_testkit::MemoryBackend; - -#[derive(Debug)] -struct RefusedDocuments(ErrorKind); - -impl RefusedDocuments { - fn refused<T>(&self) -> tinyagents_storage_kit::Result<T> { - Err(StorageError::new(self.0, "admitted backend refused")) - } -} - -#[async_trait] -impl DocumentStore for RefusedDocuments { - async fn drop_collection(&self, _: &str) -> tinyagents_storage_kit::Result<()> { - self.refused() - } - fn capabilities(&self) -> Capabilities { - Capabilities::default() - } - async fn ensure_collection(&self, _: &CollectionSpec) -> tinyagents_storage_kit::Result<()> { - self.refused() - } - async fn get( - &self, - _: &str, - _: &str, - ) -> tinyagents_storage_kit::Result<Option<Versioned<Value>>> { - self.refused() - } - async fn put( - &self, - _: &str, - _: &str, - _: Value, - _: Precondition, - ) -> tinyagents_storage_kit::Result<Version> { - self.refused() - } - async fn delete( - &self, - _: &str, - _: &str, - _: Precondition, - ) -> tinyagents_storage_kit::Result<bool> { - self.refused() - } - async fn query( - &self, - _: &str, - _: &Query, - ) -> tinyagents_storage_kit::Result<Page<Versioned<Value>>> { - self.refused() - } - async fn count(&self, _: &str, _: &Filter) -> tinyagents_storage_kit::Result<u64> { - self.refused() - } - async fn delete_where(&self, _: &str, _: &Filter) -> tinyagents_storage_kit::Result<u64> { - self.refused() - } - async fn claim( - &self, - _: &str, - _: &Filter, - _: &[Sort], - _: &Value, - ) -> tinyagents_storage_kit::Result<Option<Versioned<Value>>> { - self.refused() - } -} - -#[tokio::test(flavor = "current_thread")] -async fn authoritative_refusals_preserve_native_kind_and_never_create_sqlite() { - let ambient_db = std::env::current_dir() - .unwrap() - .join("session_db/sessions.db"); - let existed_before = ambient_db.exists(); - for kind in [ - ErrorKind::Unavailable, - ErrorKind::Conflict, - ErrorKind::Unsupported(Capability::Transactions), - ] { - let memory = MemoryBackend::new(); - let scope = Scope::new("explicit-profile").unwrap(); - let handles = memory.for_scope(&scope).unwrap(); - let admitted = ScopedStorage::new( - Storage::new( - Arc::new(RefusedDocuments(kind)), - handles.streams().clone(), - handles.kv().clone(), - handles.blobs().clone(), - handles.queues().clone(), - handles.secrets().clone(), - ), - scope, - ); - let store = LedgerWorkflowStore::new(Arc::new(DriverRunLedger::new(admitted))); - let outcomes = [ - store.load("workflow").await.map(|_| ()), - store.upsert(row()).await.map(|_| ()), - store - .claim("workflow", "owner", Duration::from_secs(30)) - .await - .map(|_| ()), - store - .compare_and_swap(row(), 0, "owner", Duration::from_secs(30)) - .await - .map(|_| ()), - store - .renew("workflow", "owner", Duration::from_secs(30)) - .await - .map(|_| ()), - ]; - for outcome in outcomes { - let WorkflowPersistenceError::Ledger(LedgerError::Driver(error)) = outcome.unwrap_err() - else { - panic!("native error flattened") - }; - assert_eq!(error.kind(), kind); - assert_eq!(error.message(), "admitted backend refused"); - } - } - assert_eq!( - ambient_db.exists(), - existed_before, - "refusal created an ambient SQLite database" - ); -} diff --git a/crates/tinyagents-orchestration/src/workflow/persistence_tests.rs b/crates/tinyagents-orchestration/src/workflow/persistence_tests.rs deleted file mode 100644 index 8d424b999..000000000 --- a/crates/tinyagents-orchestration/src/workflow/persistence_tests.rs +++ /dev/null @@ -1,226 +0,0 @@ -use super::*; -use serde_json::json; -use tinyagents_session::run_ledger::WorkflowRunStatus; -use tinystoragedrivers_testkit::MemoryBackend; - -fn row() -> WorkflowRunUpsert { - WorkflowRunUpsert { - id: "workflow".into(), - definition_id: "definition".into(), - parent_thread_id: None, - input: json!({"input":1}), - phase_states: json!({}), - child_run_ids: vec![], - status: WorkflowRunStatus::Running, - summary: None, - started_at: None, - completed_at: None, - } -} - -#[tokio::test(flavor = "current_thread")] -async fn authoritative_adapter_loads_persisted_ledger_without_runtime_bridge() { - let ledger_backend = MemoryBackend::new(); - let _workspace = tempfile::tempdir().unwrap(); - let ledger = Arc::new(crate::test_support::ledger_in(&ledger_backend)); - let expected = ledger.upsert_workflow_run(row()).await.unwrap(); - let store: Arc<dyn AsyncWorkflowStore> = Arc::new(LedgerWorkflowStore::new(ledger)); - assert_eq!(store.load("workflow").await.unwrap(), Some(expected)); -} - -#[tokio::test(flavor = "current_thread")] -async fn independent_adapters_preserve_claim_revision_and_stale_owner_rules() { - let ledger_backend = MemoryBackend::new(); - let _workspace = tempfile::tempdir().unwrap(); - let first = LedgerWorkflowStore::new(Arc::new(crate::test_support::ledger_in(&ledger_backend))); - let second = - LedgerWorkflowStore::new(Arc::new(crate::test_support::ledger_in(&ledger_backend))); - let initial = first.upsert(row()).await.unwrap(); - assert_eq!(initial.revision, 0); - let WorkflowLeaseClaim::Acquired(claimed) = first - .claim("workflow", "owner-a", Duration::from_secs(60)) - .await - .unwrap() - else { - panic!("claim refused") - }; - assert!(matches!( - second - .claim("workflow", "owner-b", Duration::from_secs(60)) - .await - .unwrap(), - WorkflowLeaseClaim::Busy(_) - )); - assert!( - !second - .renew("workflow", "owner-b", Duration::from_secs(60)) - .await - .unwrap() - ); - assert!( - second - .compare_and_swap(row(), claimed.revision, "owner-b", Duration::from_secs(60)) - .await - .unwrap() - .is_none() - ); - assert!( - first - .renew("workflow", "owner-a", Duration::from_secs(60)) - .await - .unwrap() - ); - let updated = first - .compare_and_swap(row(), claimed.revision, "owner-a", Duration::from_secs(60)) - .await - .unwrap() - .unwrap(); - assert_eq!(updated.revision, claimed.revision + 1); - assert!( - first - .compare_and_swap(row(), claimed.revision, "owner-a", Duration::from_secs(60)) - .await - .unwrap() - .is_none() - ); -} - -struct AuthoritativeRegistration; -#[async_trait] -impl AsyncWorkflowChildRegistration for AuthoritativeRegistration { - async fn register_async(&self, _: String) -> WorkflowPersistenceResult<()> { - Ok(()) - } -} - -#[tokio::test(flavor = "current_thread")] -async fn authoritative_registration_refuses_synchronous_success() { - let registration: Arc<dyn AsyncWorkflowChildRegistration> = Arc::new(AuthoritativeRegistration); - assert!(matches!( - registration.register("child".into()), - Err(WorkflowPersistenceError::AsyncRegistrationRequired) - )); - registration.register_async("child".into()).await.unwrap(); -} - -#[tokio::test(flavor = "current_thread")] -async fn explicit_driver_scopes_isolate_workflows_without_local_files() { - use tinyagents_session::run_ledger::DriverRunLedger; - use tinyagents_storage_kit::{Scope, StorageBackend}; - let backend = MemoryBackend::new(); - let first = LedgerWorkflowStore::new(Arc::new(DriverRunLedger::new( - backend - .for_scope(&Scope::new("profile-a").unwrap()) - .unwrap(), - ))); - let second = LedgerWorkflowStore::new(Arc::new(DriverRunLedger::new( - backend - .for_scope(&Scope::new("profile-b").unwrap()) - .unwrap(), - ))); - first.upsert(row()).await.unwrap(); - assert!(second.load("workflow").await.unwrap().is_none()); - assert!(first.load("workflow").await.unwrap().is_some()); -} - -#[path = "persistence_refusal_tests.rs"] -mod refusal_tests; - -#[tokio::test(flavor = "current_thread")] -async fn invalid_lease_duration_refuses_before_touching_the_ledger() { - let ledger_backend = MemoryBackend::new(); - let _workspace = tempfile::tempdir().unwrap(); - let store = LedgerWorkflowStore::new(Arc::new(crate::test_support::ledger_in(&ledger_backend))); - let lease = Duration::MAX; - assert!(matches!( - store.claim("workflow", "owner", lease).await, - Err(WorkflowPersistenceError::LeaseDuration(_)) - )); - assert!(matches!( - store.compare_and_swap(row(), 0, "owner", lease).await, - Err(WorkflowPersistenceError::LeaseDuration(_)) - )); - assert!(matches!( - store.renew("workflow", "owner", lease).await, - Err(WorkflowPersistenceError::LeaseDuration(_)) - )); - // Nothing reached the ledger: the refused claim created no workflow. - assert!(store.load("workflow").await.unwrap().is_none()); -} - -struct LegacyCallback(std::sync::atomic::AtomicUsize); -impl WorkflowChildRegistration for LegacyCallback { - fn register(&self, child_id: String) -> Result<(), OrchestrationError> { - if child_id == "refused" { - return Err(OrchestrationError("legacy callback refused".into())); - } - self.0.fetch_add(1, std::sync::atomic::Ordering::SeqCst); - Ok(()) - } -} - -#[tokio::test(flavor = "current_thread")] -async fn explicit_legacy_registration_preserves_callback_results() { - let callback = Arc::new(LegacyCallback(std::sync::atomic::AtomicUsize::new(0))); - let adapter = LegacyWorkflowRegistration::new(callback.clone()); - adapter.register("sync-child".into()).unwrap(); - adapter.register_async("async-child".into()).await.unwrap(); - assert_eq!(callback.0.load(std::sync::atomic::Ordering::SeqCst), 2); - assert!(matches!( - adapter.register_async("refused".into()).await, - Err(WorkflowPersistenceError::Legacy(_)) - )); - assert!(matches!( - adapter.register("refused".into()), - Err(WorkflowPersistenceError::Legacy(_)) - )); - assert_eq!(callback.0.load(std::sync::atomic::Ordering::SeqCst), 2); -} - -#[tokio::test(flavor = "current_thread")] -async fn expired_owner_takeover_fences_prior_owner_without_sleeping() { - let ledger_backend = MemoryBackend::new(); - let _workspace = tempfile::tempdir().unwrap(); - let ledger = Arc::new(crate::test_support::ledger_in(&ledger_backend)); - ledger.upsert_workflow_run(row()).await.unwrap(); - let WorkflowLeaseClaim::Acquired(expired) = ledger - .try_claim_workflow_run("workflow", "old-owner", chrono::Duration::seconds(-1)) - .await - .unwrap() - else { - panic!("initial claim refused") - }; - let store = LedgerWorkflowStore::new(ledger); - let WorkflowLeaseClaim::Acquired(replacement) = store - .claim("workflow", "new-owner", Duration::from_secs(60)) - .await - .unwrap() - else { - panic!("expired takeover refused") - }; - assert!(replacement.revision > expired.revision); - assert!( - !store - .renew("workflow", "old-owner", Duration::from_secs(60)) - .await - .unwrap() - ); - assert!( - store - .compare_and_swap( - row(), - replacement.revision, - "old-owner", - Duration::from_secs(60) - ) - .await - .unwrap() - .is_none() - ); - assert!( - store - .renew("workflow", "new-owner", Duration::from_secs(60)) - .await - .unwrap() - ); -} diff --git a/crates/tinyagents-orchestration/src/workflow/state.rs b/crates/tinyagents-orchestration/src/workflow/state.rs deleted file mode 100644 index 70a37d7df..000000000 --- a/crates/tinyagents-orchestration/src/workflow/state.rs +++ /dev/null @@ -1,259 +0,0 @@ -//! Phase state projection and advancement for workflow runs. -//! -//! This module owns the JSON phase-state document (a BTreeMap of phase names -//! to status + outputs) that the workflow engine persists. It provides queries -//! (what phases are runnable?) and mutations (mark complete, reset after -//! interruption) over this state without touching the underlying persistence layer. - -use serde_json::{Value, json}; - -use super::{WorkflowDefinition, WorkflowPhase}; - -/// Durable status of one phase in a workflow run. -/// -/// Phases transition: Pending → Running → (Completed | Failed). A phase -/// interrupted while running can be reset to Pending for retry; completed -/// phases are immutable. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum PhaseStatus { - /// Phase has not yet started. - Pending, - /// Phase is currently executing (likely in a child). - Running, - /// Phase completed successfully; immutable. - Completed, - /// Phase failed; may be retried by resetting to Pending. - Failed, -} - -impl PhaseStatus { - /// Returns the string representation of this status (used in JSON). - pub const fn as_str(self) -> &'static str { - match self { - Self::Pending => "pending", - Self::Running => "running", - Self::Completed => "completed", - Self::Failed => "failed", - } - } -} - -/// Initializes a phase-state document from a workflow definition. -/// -/// All phases start in the Pending status with empty outputs. The returned -/// JSON structure maps phase names to `{ status, outputs, ... }` objects. -pub fn init_phase_states(definition: &WorkflowDefinition) -> Value { - Value::Object( - definition - .phases - .iter() - .map(|phase| { - ( - phase.name.clone(), - json!({ "status": "pending", "outputs": [] }), - ) - }) - .collect(), - ) -} - -/// Queries the current status string of a phase, or None if the phase is unknown. -pub fn phase_status<'a>(phase_states: &'a Value, name: &str) -> Option<&'a str> { - phase_states.get(name)?.get("status")?.as_str() -} - -pub(crate) fn set_phase_status( - phase_states: &mut Value, - name: &str, - status: PhaseStatus, - outputs: Option<Value>, -) { - let Some(entries) = phase_states.as_object_mut() else { - return; - }; - let entry = entries - .entry(name.to_owned()) - .or_insert_with(|| json!({ "status": "pending", "outputs": [] })); - if let Some(object) = entry.as_object_mut() { - object.insert("status".to_owned(), json!(status.as_str())); - if let Some(outputs) = outputs { - object.insert("outputs".to_owned(), outputs); - } - } -} - -pub(crate) fn set_phase_reason(phase_states: &mut Value, name: &str, reason: &str) { - if let Some(object) = phase_states.get_mut(name).and_then(Value::as_object_mut) { - object.insert("reason".to_owned(), json!(reason)); - } -} - -/// Make an interrupted phase runnable again. A stopped phase may have spawned -/// children whose results were never durably collected; retrying the whole -/// phase is the only safe, deterministic recovery. Completed phases remain -/// immutable and are never retried. -pub fn reset_running_phases(phase_states: &mut Value, reason: &str) { - let Some(phases) = phase_states.as_object_mut() else { - return; - }; - for entry in phases.values_mut() { - let Some(state) = entry.as_object_mut() else { - continue; - }; - if state.get("status").and_then(Value::as_str) == Some("running") { - state.insert("status".to_owned(), json!(PhaseStatus::Pending.as_str())); - state.insert("outputs".to_owned(), json!([])); - state.insert("reason".to_owned(), json!(reason)); - } - } -} - -/// Finds the next phase that should run: a phase that is not already -/// completed, running, or failed, and all of whose dependencies are completed. -/// -/// Returns the first such phase in definition order, or `None` if no phase is -/// ready (either all are done, or some have unmet dependencies). -pub fn next_runnable_phase<'a>( - definition: &'a WorkflowDefinition, - phase_states: &Value, -) -> Option<&'a WorkflowPhase> { - definition.phases.iter().find(|phase| { - !matches!( - phase_status(phase_states, &phase.name), - Some("completed" | "running" | "failed") - ) && phase - .depends_on - .iter() - .all(|dependency| phase_status(phase_states, dependency) == Some("completed")) - }) -} - -/// Checks whether all phases in the workflow have completed. -pub fn all_phases_completed(definition: &WorkflowDefinition, phase_states: &Value) -> bool { - definition - .phases - .iter() - .all(|phase| phase_status(phase_states, &phase.name) == Some("completed")) -} - -/// Collects outputs from all upstream (dependency) phases for a given phase. -/// Filters out empty and null outputs to yield only meaningful results. -pub fn upstream_outputs(phase: &WorkflowPhase, phase_states: &Value) -> Vec<Value> { - phase - .depends_on - .iter() - .flat_map(|dependency| { - phase_states - .get(dependency) - .and_then(|entry| entry.get("outputs")) - .and_then(Value::as_array) - .into_iter() - .flatten() - .filter_map(move |item| { - durable_output(item) - .filter(|output| match output { - Value::Null => false, - Value::String(text) => !text.trim().is_empty(), - _ => true, - }) - .map(|output| json!({ "phase": dependency, "output": output })) - }) - }) - .collect() -} - -/// Composes the prompt for a worker in a phase. -/// -/// Includes the phase name and description, the input question, the worker's -/// index (if multiple workers), and relevant upstream outputs from dependencies. -pub fn phase_prompt( - input: &Value, - phase: &WorkflowPhase, - index: usize, - upstream: &[Value], -) -> String { - let question = input - .get("question") - .or_else(|| input.get("input")) - .and_then(Value::as_str) - .map(str::to_owned) - .unwrap_or_else(|| input.to_string()); - let mut prompt = format!( - "Workflow phase: {}\n{}\n\nInput:\n{}\n", - phase.name, phase.description, question - ); - if phase.agent_ids.len() > 1 { - prompt.push_str(&format!( - "\n(You are worker #{} in this phase.)\n", - index + 1 - )); - } - if !upstream.is_empty() { - prompt.push_str("\nContext from prior phases:\n"); - for item in upstream { - if let (Some(source), Some(output)) = ( - item.get("phase").and_then(Value::as_str), - item.get("output"), - ) { - prompt.push_str(&format!("- [{source}] {}\n", render_output(output))); - } - } - } - prompt -} - -/// Composes a workflow summary from all final phase outputs. -/// -/// Returns `None` if all phases are complete and there are no outputs; -/// otherwise returns a formatted summary of all non-empty outputs in phase order. -pub fn synthesize_summary(definition: &WorkflowDefinition, phase_states: &Value) -> Option<String> { - let outputs_for = |name: &str| { - phase_states - .get(name) - .and_then(|entry| entry.get("outputs")) - .and_then(Value::as_array) - .map(|outputs| { - outputs - .iter() - .filter_map(|output| durable_output(output).map(|value| render_output(&value))) - .filter(|output| !output.trim().is_empty() && output != "null") - .collect::<Vec<_>>() - .join("\n") - }) - .filter(|summary| !summary.trim().is_empty()) - }; - outputs_for("synthesize").or_else(|| { - definition - .phases - .iter() - .rev() - .find_map(|phase| outputs_for(&phase.name)) - }) -} - -/// The public/RPC projection remains `{ agentId, output: String }` for -/// compatibility. New rows carry the exact result in `metadata.rawOutput` so -/// future phases retain arbitrary JSON without changing the old wire shape. -fn durable_output(item: &Value) -> Option<Value> { - item.get("metadata") - .and_then(|metadata| metadata.get("version")) - .and_then(Value::as_u64) - .filter(|version| *version >= 2) - .and_then(|_| { - item.get("metadata") - .and_then(|metadata| metadata.get("rawOutput")) - }) - .cloned() - .or_else(|| item.get("output").cloned()) -} - -/// Preserve every JSON output in prompt context and summaries. JSON object's -/// map ordering is canonical under serde_json's default map implementation, -/// so repeated resume/synthesis renders the same bytes rather than silently -/// discarding structured child results. -fn render_output(value: &Value) -> String { - match value { - Value::String(text) => text.clone(), - _ => serde_json::to_string(value).unwrap_or_else(|_| "<unserializable output>".to_owned()), - } -} diff --git a/crates/tinyagents-orchestration/src/workflow/supervisor.rs b/crates/tinyagents-orchestration/src/workflow/supervisor.rs deleted file mode 100644 index 687e76215..000000000 --- a/crates/tinyagents-orchestration/src/workflow/supervisor.rs +++ /dev/null @@ -1,468 +0,0 @@ -//! Process-local ownership of workflow drive tasks. -//! -//! A host supplies its task spawner (for example, a runtime context's `spawn`) -//! and the closure that drives a run. Every start, including a resume, gets a -//! fresh cancellation generation. A handle from an older generation cannot -//! stop a replacement loop. - -use std::collections::HashMap; -use std::convert::Infallible; -use std::future::Future; -use std::pin::Pin; -use std::sync::{Arc, Mutex}; -use std::task::{Context, Poll}; - -use tinyagents_harness::CancellationToken; - -#[derive(Debug, Clone)] -struct Generation { - id: u64, - token: CancellationToken, -} - -/// Handle for one scheduled drive generation. -#[derive(Debug, Clone)] -pub struct WorkflowRunHandle { - run_id: Arc<str>, - generation: u64, - token: CancellationToken, -} - -/// A scheduled workflow future that releases its cancellation generation when -/// it completes or is dropped. -pub struct WorkflowDriveFuture<F: Future<Output = ()>> { - future: Pin<Box<F>>, - supervisor: WorkflowRunSupervisor, - handle: WorkflowRunHandle, -} - -impl<F: Future<Output = ()>> Future for WorkflowDriveFuture<F> { - type Output = (); - - fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { - let this = self.get_mut(); - match this.future.as_mut().poll(cx) { - Poll::Ready(()) => { - this.supervisor.finish(&this.handle); - Poll::Ready(()) - } - Poll::Pending => Poll::Pending, - } - } -} - -impl<F: Future<Output = ()>> Drop for WorkflowDriveFuture<F> { - fn drop(&mut self) { - self.supervisor.finish(&self.handle); - } -} - -impl WorkflowRunHandle { - /// Cancellation token the workflow engine should observe while driving. - pub fn cancellation_token(&self) -> CancellationToken { - self.token.clone() - } - - /// Durable workflow id associated with this generation. - pub fn run_id(&self) -> &str { - &self.run_id - } -} - -#[derive(Debug, Default)] -struct Registry { - next_generation: u64, - current: HashMap<String, Generation>, -} - -/// Starts and cancels host-scheduled workflow drive loops. -/// -/// The supervisor owns only process-local task identity and cancellation. -/// Durable lifecycle transitions remain the responsibility of the caller. -#[derive(Debug, Clone, Default)] -pub struct WorkflowRunSupervisor { - registry: Arc<Mutex<Registry>>, -} - -impl WorkflowRunSupervisor { - /// Install a fresh generation and schedule its drive future. - /// - /// `spawn` is a host seam: pass a closure backed by the host's task context. - /// The future factory receives the generation's cancellation token so the - /// caller can pass it directly to `WorkflowEngine::drive` or - /// `drive_async`. The scheduled future clears its generation when it - /// completes or is dropped. - pub fn start<F, MakeFuture, Spawn>( - &self, - run_id: impl Into<String>, - make_future: MakeFuture, - spawn: Spawn, - ) -> WorkflowRunHandle - where - F: Future<Output = ()> + Send + 'static, - MakeFuture: FnOnce(CancellationToken) -> F, - Spawn: FnOnce(WorkflowDriveFuture<F>), - { - self.start_scoped_with_handle( - run_id, - make_future, - |future| future, - |future| future, - |_, future| spawn(future), - ) - } - - /// Start a drive with explicit host scope wrappers. - /// - /// `root_parent` and `origin` wrap the drive future in that order before - /// it reaches the host-provided spawner. Hosts can bind their root parent - /// context and turn origin without adding either policy to TinyAgents. - pub fn start_scoped< - Drive, - MakeFuture, - RootScope, - RootFuture, - OriginScope, - OriginFuture, - Spawn, - >( - &self, - run_id: impl Into<String>, - make_future: MakeFuture, - root_parent: RootScope, - origin: OriginScope, - spawn: Spawn, - ) -> WorkflowRunHandle - where - Drive: Future<Output = ()> + Send + 'static, - MakeFuture: FnOnce(CancellationToken) -> Drive, - RootScope: FnOnce(Drive) -> RootFuture, - RootFuture: Future<Output = ()> + Send + 'static, - OriginScope: FnOnce(RootFuture) -> OriginFuture, - OriginFuture: Future<Output = ()> + Send + 'static, - Spawn: FnOnce(WorkflowDriveFuture<OriginFuture>), - { - self.start_scoped_with_handle(run_id, make_future, root_parent, origin, |_, future| { - spawn(future) - }) - } - - /// Start a drive and pass its installed handle to the host before it can - /// schedule or poll the future. This lets a host register the generation - /// with its durable supervisor before a fast drive can finish. - pub fn start_scoped_with_handle< - Drive, - MakeFuture, - RootScope, - RootFuture, - OriginScope, - OriginFuture, - Spawn, - >( - &self, - run_id: impl Into<String>, - make_future: MakeFuture, - root_parent: RootScope, - origin: OriginScope, - spawn: Spawn, - ) -> WorkflowRunHandle - where - Drive: Future<Output = ()> + Send + 'static, - MakeFuture: FnOnce(CancellationToken) -> Drive, - RootScope: FnOnce(Drive) -> RootFuture, - RootFuture: Future<Output = ()> + Send + 'static, - OriginScope: FnOnce(RootFuture) -> OriginFuture, - OriginFuture: Future<Output = ()> + Send + 'static, - Spawn: FnOnce(WorkflowRunHandle, WorkflowDriveFuture<OriginFuture>), - { - self.start_scoped_with_commit( - run_id, - || Ok::<_, Infallible>(()), - make_future, - root_parent, - origin, - spawn, - ) - .expect("infallible workflow commit") - .0 - } - - /// Commit a host-owned lifecycle transition and install a drive generation - /// as one serialized operation. - /// - /// `commit` runs first while the supervisor registry is locked. If it - /// returns an error, the current generation is left unchanged and `spawn` - /// is not called. If it succeeds, the previous generation is cancelled and - /// the new one is installed before the lock is released. Future construction - /// and spawning happen after the lock is released. Hosts can use this to - /// serialize a durable compare-and-set with generation replacement. The - /// callback must not call back into this supervisor because the registry - /// mutex is not reentrant. - pub fn start_scoped_with_commit< - CommitValue, - CommitError, - Drive, - MakeFuture, - RootScope, - RootFuture, - OriginScope, - OriginFuture, - Spawn, - >( - &self, - run_id: impl Into<String>, - commit: impl FnOnce() -> Result<CommitValue, CommitError>, - make_future: MakeFuture, - root_parent: RootScope, - origin: OriginScope, - spawn: Spawn, - ) -> Result<(WorkflowRunHandle, CommitValue), CommitError> - where - Drive: Future<Output = ()> + Send + 'static, - MakeFuture: FnOnce(CancellationToken) -> Drive, - RootScope: FnOnce(Drive) -> RootFuture, - RootFuture: Future<Output = ()> + Send + 'static, - OriginScope: FnOnce(RootFuture) -> OriginFuture, - OriginFuture: Future<Output = ()> + Send + 'static, - Spawn: FnOnce(WorkflowRunHandle, WorkflowDriveFuture<OriginFuture>), - { - let run_id = run_id.into(); - let (generation, handle, commit_value) = { - let mut registry = self.registry.lock().expect("workflow registry poisoned"); - let commit_value = commit()?; - registry.next_generation = registry.next_generation.wrapping_add(1); - if let Some(previous) = registry.current.get(&run_id) { - previous.token.cancel(); - } - let generation = Generation { - id: registry.next_generation, - token: CancellationToken::new(), - }; - registry.current.insert(run_id.clone(), generation.clone()); - let handle = WorkflowRunHandle { - run_id: Arc::from(run_id), - generation: generation.id, - token: generation.token.clone(), - }; - (generation, handle, commit_value) - }; - let future = WorkflowDriveFuture { - future: Box::pin(origin(root_parent(make_future(generation.token)))), - supervisor: self.clone(), - handle: handle.clone(), - }; - spawn(handle.clone(), future); - Ok((handle, commit_value)) - } - - /// Install a new generation for a resumed run and schedule its drive. - pub fn resume<F, MakeFuture, Spawn>( - &self, - run_id: impl Into<String>, - make_future: MakeFuture, - spawn: Spawn, - ) -> WorkflowRunHandle - where - F: Future<Output = ()> + Send + 'static, - MakeFuture: FnOnce(CancellationToken) -> F, - Spawn: FnOnce(WorkflowDriveFuture<F>), - { - self.start(run_id, make_future, spawn) - } - - /// Resume a drive with explicit host root-parent and origin scopes. - pub fn resume_scoped< - Drive, - MakeFuture, - RootScope, - RootFuture, - OriginScope, - OriginFuture, - Spawn, - >( - &self, - run_id: impl Into<String>, - make_future: MakeFuture, - root_parent: RootScope, - origin: OriginScope, - spawn: Spawn, - ) -> WorkflowRunHandle - where - Drive: Future<Output = ()> + Send + 'static, - MakeFuture: FnOnce(CancellationToken) -> Drive, - RootScope: FnOnce(Drive) -> RootFuture, - RootFuture: Future<Output = ()> + Send + 'static, - OriginScope: FnOnce(RootFuture) -> OriginFuture, - OriginFuture: Future<Output = ()> + Send + 'static, - Spawn: FnOnce(WorkflowDriveFuture<OriginFuture>), - { - self.resume_scoped_with_handle(run_id, make_future, root_parent, origin, |_, future| { - spawn(future) - }) - } - - /// Resume a drive and provide its installed handle before scheduling it. - pub fn resume_scoped_with_handle< - Drive, - MakeFuture, - RootScope, - RootFuture, - OriginScope, - OriginFuture, - Spawn, - >( - &self, - run_id: impl Into<String>, - make_future: MakeFuture, - root_parent: RootScope, - origin: OriginScope, - spawn: Spawn, - ) -> WorkflowRunHandle - where - Drive: Future<Output = ()> + Send + 'static, - MakeFuture: FnOnce(CancellationToken) -> Drive, - RootScope: FnOnce(Drive) -> RootFuture, - RootFuture: Future<Output = ()> + Send + 'static, - OriginScope: FnOnce(RootFuture) -> OriginFuture, - OriginFuture: Future<Output = ()> + Send + 'static, - Spawn: FnOnce(WorkflowRunHandle, WorkflowDriveFuture<OriginFuture>), - { - self.start_scoped_with_handle(run_id, make_future, root_parent, origin, spawn) - } - - /// Resume a drive only after the host's lifecycle commit succeeds. - /// - /// The commit, replacement of any previous generation, and installation - /// of this generation are serialized by the registry lock. Future - /// construction and spawning happen after releasing that lock. The commit - /// callback must not call back into this supervisor because the registry - /// mutex is not reentrant. - pub fn resume_scoped_with_commit< - CommitValue, - CommitError, - Drive, - MakeFuture, - RootScope, - RootFuture, - OriginScope, - OriginFuture, - Spawn, - >( - &self, - run_id: impl Into<String>, - commit: impl FnOnce() -> Result<CommitValue, CommitError>, - make_future: MakeFuture, - root_parent: RootScope, - origin: OriginScope, - spawn: Spawn, - ) -> Result<(WorkflowRunHandle, CommitValue), CommitError> - where - Drive: Future<Output = ()> + Send + 'static, - MakeFuture: FnOnce(CancellationToken) -> Drive, - RootScope: FnOnce(Drive) -> RootFuture, - RootFuture: Future<Output = ()> + Send + 'static, - OriginScope: FnOnce(RootFuture) -> OriginFuture, - OriginFuture: Future<Output = ()> + Send + 'static, - Spawn: FnOnce(WorkflowRunHandle, WorkflowDriveFuture<OriginFuture>), - { - self.start_scoped_with_commit(run_id, commit, make_future, root_parent, origin, spawn) - } - - /// Cancel the current generation for `run_id`, if one is scheduled. - /// - /// Returns `false` if no generation is registered. - pub fn stop(&self, run_id: &str) -> bool { - let registry = self.registry.lock().expect("workflow registry poisoned"); - let Some(generation) = registry.current.get(run_id) else { - return false; - }; - generation.token.cancel(); - true - } - - /// Cancel only when `handle` still owns the run's current generation. - /// This is safe for late callbacks from a stopped loop after resume. - pub fn stop_generation(&self, handle: &WorkflowRunHandle) -> bool { - let registry = self.registry.lock().expect("workflow registry poisoned"); - let Some(current) = registry.current.get(handle.run_id()) else { - return false; - }; - if current.id != handle.generation { - return false; - } - current.token.cancel(); - true - } - - /// Snapshot the current generation handle for `run_id`, if one is - /// registered. - /// - /// The returned handle may become stale immediately after this call. It is - /// safe to use with [`Self::stop_generation`]: that method checks the - /// generation identity and will not cancel a replacement. This snapshot - /// does not make a durable lifecycle transition atomic with the registry. - pub fn current_handle(&self, run_id: &str) -> Option<WorkflowRunHandle> { - let registry = self.registry.lock().expect("workflow registry poisoned"); - registry - .current - .get(run_id) - .map(|generation| WorkflowRunHandle { - run_id: Arc::from(run_id), - generation: generation.id, - token: generation.token.clone(), - }) - } - - /// Returns whether `handle` still identifies the current generation. - /// Stopping a generation cancels it but leaves it current until it is - /// finished or dropped; replacement and completion make it non-current. - pub fn is_current(&self, handle: &WorkflowRunHandle) -> bool { - self.registry - .lock() - .expect("workflow registry poisoned") - .current - .get(handle.run_id()) - .is_some_and(|current| current.id == handle.generation) - } - - /// Run a callback while `handle` still identifies the installed generation. - /// - /// The callback runs while the supervisor registry is locked, so a - /// replacement cannot be installed between the identity check and the - /// callback. The callback must not call back into this supervisor, since - /// that would attempt to lock the same non-reentrant mutex. - pub fn with_current<T>( - &self, - handle: &WorkflowRunHandle, - callback: impl FnOnce() -> T, - ) -> Option<T> { - let registry = self.registry.lock().expect("workflow registry poisoned"); - if registry - .current - .get(handle.run_id()) - .is_none_or(|current| current.id != handle.generation) - { - return None; - } - Some(callback()) - } - - /// Remove a finished generation without disturbing a newer replacement. - pub fn finish(&self, handle: &WorkflowRunHandle) -> bool { - let mut registry = self.registry.lock().expect("workflow registry poisoned"); - if registry - .current - .get(handle.run_id()) - .is_some_and(|current| current.id == handle.generation) - { - registry.current.remove(handle.run_id()); - true - } else { - false - } - } -} - -#[cfg(test)] -#[path = "supervisor_tests.rs"] -mod tests; diff --git a/crates/tinyagents-orchestration/src/workflow/supervisor_tests.rs b/crates/tinyagents-orchestration/src/workflow/supervisor_tests.rs deleted file mode 100644 index 7a4cab17d..000000000 --- a/crates/tinyagents-orchestration/src/workflow/supervisor_tests.rs +++ /dev/null @@ -1,463 +0,0 @@ -use super::*; - -#[test] -fn stale_generation_cannot_cancel_its_resumed_successor() { - let supervisor = WorkflowRunSupervisor::default(); - let pending = Arc::new(std::sync::Mutex::new(Vec::< - Pin<Box<dyn Future<Output = ()> + Send>>, - >::new())); - let first_pending = pending.clone(); - let second_pending = pending.clone(); - let first = supervisor.start( - "run", - |_| async {}, - move |future| { - first_pending.lock().unwrap().push(Box::pin(future)); - }, - ); - let second = supervisor.resume( - "run", - |_| async {}, - move |future| { - second_pending.lock().unwrap().push(Box::pin(future)); - }, - ); - - assert!(first.cancellation_token().is_cancelled()); - assert!(!supervisor.stop_generation(&first)); - assert!(!second.cancellation_token().is_cancelled()); - drop(pending.lock().unwrap().remove(0)); - assert!(supervisor.stop("run")); - assert!(second.cancellation_token().is_cancelled()); -} - -#[test] -fn finish_only_removes_the_generation_that_finished() { - let supervisor = WorkflowRunSupervisor::default(); - let pending = Arc::new(std::sync::Mutex::new(Vec::< - Pin<Box<dyn Future<Output = ()> + Send>>, - >::new())); - let old_pending = pending.clone(); - let current_pending = pending.clone(); - let old = supervisor.start( - "run", - |_| async {}, - move |future| { - old_pending.lock().unwrap().push(Box::pin(future)); - }, - ); - let current = supervisor.start( - "run", - |_| async {}, - move |future| { - current_pending.lock().unwrap().push(Box::pin(future)); - }, - ); - - assert!(!supervisor.finish(&old)); - assert!(supervisor.stop("run")); - assert!(current.cancellation_token().is_cancelled()); - pending.lock().unwrap().clear(); - assert!(!supervisor.stop("run")); -} - -#[test] -fn dropping_unpolled_drive_clears_its_generation() { - let supervisor = WorkflowRunSupervisor::default(); - let handle = supervisor.start("run", |_| async {}, drop); - - assert!(!supervisor.stop_generation(&handle)); - assert!(!supervisor.stop("run")); -} - -#[test] -fn handle_aware_start_installs_current_handle_before_spawning_fast_future() { - let supervisor = WorkflowRunSupervisor::default(); - let mut observed = None; - let returned = supervisor.start_scoped_with_handle( - "run", - |_| async {}, - |future| future, - |future| future, - |handle, future| { - assert!(supervisor.is_current(&handle)); - drop(future); - observed = Some(handle); - }, - ); - - assert_eq!(observed.unwrap().generation, returned.generation); - assert!(!supervisor.is_current(&returned)); -} - -#[test] -fn old_handle_is_not_current_after_replacement_and_stop_keeps_current_identity() { - let supervisor = WorkflowRunSupervisor::default(); - let old = supervisor.start("run", |_| async {}, drop); - let pending = Arc::new(std::sync::Mutex::new( - None::<Pin<Box<dyn Future<Output = ()> + Send>>>, - )); - let pending_slot = pending.clone(); - let current = supervisor.resume_scoped_with_handle( - "run", - |_| async {}, - |future| future, - |future| future, - |handle, future| { - assert!(supervisor.is_current(&handle)); - *pending_slot.lock().unwrap() = Some(Box::pin(future)); - }, - ); - - assert!(!supervisor.is_current(&old)); - assert!(supervisor.is_current(¤t)); - assert!(supervisor.stop_generation(¤t)); - assert!(supervisor.is_current(¤t)); - pending.lock().unwrap().take(); - assert!(!supervisor.is_current(¤t)); -} - -#[test] -fn current_handle_snapshot_cannot_cancel_a_replacement() { - let supervisor = WorkflowRunSupervisor::default(); - let pending = Arc::new(std::sync::Mutex::new(Vec::< - Pin<Box<dyn Future<Output = ()> + Send>>, - >::new())); - let first_pending = pending.clone(); - let second_pending = pending.clone(); - - assert!(supervisor.current_handle("run").is_none()); - supervisor.start( - "run", - |_| async {}, - move |future| { - first_pending.lock().unwrap().push(Box::pin(future)); - }, - ); - let snapshot = supervisor.current_handle("run").unwrap(); - assert!(supervisor.is_current(&snapshot)); - - let replacement = supervisor.resume( - "run", - |_| async {}, - move |future| { - second_pending.lock().unwrap().push(Box::pin(future)); - }, - ); - assert!(!supervisor.stop_generation(&snapshot)); - assert!(!replacement.cancellation_token().is_cancelled()); - assert!(supervisor.is_current(&replacement)); - - pending.lock().unwrap().clear(); - assert!(supervisor.current_handle("run").is_none()); -} - -#[test] -fn current_handle_tracks_stopped_and_finished_generations() { - let supervisor = WorkflowRunSupervisor::default(); - let pending = Arc::new(std::sync::Mutex::new( - None::<Pin<Box<dyn Future<Output = ()> + Send>>>, - )); - let slot = pending.clone(); - let handle = supervisor.start( - "run", - |_| async {}, - move |future| { - *slot.lock().unwrap() = Some(Box::pin(future)); - }, - ); - - let snapshot = supervisor.current_handle("run").unwrap(); - assert_eq!(snapshot.generation, handle.generation); - assert!(supervisor.current_handle("other-run").is_none()); - assert!(supervisor.stop_generation(&snapshot)); - let stopped = supervisor.current_handle("run").unwrap(); - assert!(stopped.cancellation_token().is_cancelled()); - assert!(supervisor.is_current(&stopped)); - assert!(supervisor.finish(&snapshot)); - assert!(supervisor.current_handle("run").is_none()); - pending.lock().unwrap().take(); -} - -#[test] -fn ready_drive_clears_its_generation_while_future_is_retained() { - let supervisor = WorkflowRunSupervisor::default(); - let pending = Arc::new(std::sync::Mutex::new( - None::<Pin<Box<dyn Future<Output = ()> + Send>>>, - )); - let slot = pending.clone(); - let handle = supervisor.start( - "run", - |_| async {}, - move |future| { - *slot.lock().unwrap() = Some(Box::pin(future)); - }, - ); - let mut future = pending.lock().unwrap().take().unwrap(); - let waker = futures::task::noop_waker(); - let mut context = Context::from_waker(&waker); - - assert!(matches!( - future.as_mut().poll(&mut context), - Poll::Ready(()) - )); - assert!(!supervisor.stop_generation(&handle)); - drop(future); -} - -#[test] -fn completed_drive_clears_its_generation_and_stale_drop_preserves_successor() { - let supervisor = WorkflowRunSupervisor::default(); - let old_pending = Arc::new(std::sync::Mutex::new( - None::<Pin<Box<dyn Future<Output = ()> + Send>>>, - )); - let slot = old_pending.clone(); - let old = supervisor.start( - "run", - |_| async {}, - move |future| { - *slot.lock().unwrap() = Some(Box::pin(future)); - }, - ); - let successor_pending = Arc::new(std::sync::Mutex::new( - None::<Pin<Box<dyn Future<Output = ()> + Send>>>, - )); - let successor_slot = successor_pending.clone(); - let successor = supervisor.resume( - "run", - |_| futures::future::pending(), - move |future| { - *successor_slot.lock().unwrap() = Some(Box::pin(future)); - }, - ); - assert!(old.cancellation_token().is_cancelled()); - futures::executor::block_on(old_pending.lock().unwrap().take().unwrap()); - assert!(supervisor.stop_generation(&successor)); - - let completed = supervisor.start("completed", |_| async {}, futures::executor::block_on); - assert!(!supervisor.stop_generation(&completed)); -} - -#[test] -fn scoped_start_wraps_drive_with_root_parent_then_origin_before_spawning() { - let supervisor = WorkflowRunSupervisor::default(); - let scopes = Arc::new(std::sync::Mutex::new(Vec::new())); - let root_scopes = scopes.clone(); - let origin_scopes = scopes.clone(); - - supervisor.start_scoped( - "run", - |_| async {}, - move |future| { - root_scopes.lock().unwrap().push("root-parent"); - future - }, - move |future| { - origin_scopes.lock().unwrap().push("origin"); - future - }, - futures::executor::block_on, - ); - - assert_eq!(*scopes.lock().unwrap(), ["root-parent", "origin"]); -} - -#[test] -fn failed_transactional_start_keeps_previous_generation_and_does_not_spawn() { - let supervisor = WorkflowRunSupervisor::default(); - let held = Arc::new(std::sync::Mutex::new(Vec::new())); - let held_old = held.clone(); - let old = supervisor.start( - "run", - |_| futures::future::pending(), - move |future| held_old.lock().unwrap().push(future), - ); - let spawned = Arc::new(std::sync::atomic::AtomicBool::new(false)); - let spawned_flag = spawned.clone(); - - let result = supervisor.start_scoped_with_commit( - "run", - || Err::<(), _>("durable claim rejected"), - |_| async {}, - |future| future, - |future| future, - move |_, future| { - spawned_flag.store(true, std::sync::atomic::Ordering::SeqCst); - drop(future); - }, - ); - - assert_eq!(result.unwrap_err(), "durable claim rejected"); - assert!(!spawned.load(std::sync::atomic::Ordering::SeqCst)); - assert!(supervisor.is_current(&old)); - assert!(!old.cancellation_token().is_cancelled()); -} - -#[test] -fn transactional_commit_serializes_generation_install_until_commit_finishes() { - let supervisor = WorkflowRunSupervisor::default(); - let held = Arc::new(std::sync::Mutex::new(Vec::new())); - let held_old = held.clone(); - let old = supervisor.start( - "run", - |_| futures::future::pending(), - move |future| held_old.lock().unwrap().push(future), - ); - let commit_entered = Arc::new(std::sync::Barrier::new(2)); - let release_commit = Arc::new(std::sync::Barrier::new(2)); - let supervisor_for_start = supervisor.clone(); - let entered_for_commit = commit_entered.clone(); - let release_for_commit = release_commit.clone(); - let held_new = held.clone(); - let start_thread = std::thread::spawn(move || { - supervisor_for_start.start_scoped_with_commit( - "run", - move || { - entered_for_commit.wait(); - release_for_commit.wait(); - Ok::<_, ()>(()) - }, - |_| futures::future::pending(), - |future| future, - |future| future, - move |_, future| held_new.lock().unwrap().push(future), - ) - }); - - commit_entered.wait(); - let (snapshot_attempted_tx, snapshot_attempted_rx) = std::sync::mpsc::channel(); - let supervisor_for_snapshot = supervisor.clone(); - let snapshot_thread = std::thread::spawn(move || { - snapshot_attempted_tx.send(()).unwrap(); - supervisor_for_snapshot.current_handle("run") - }); - snapshot_attempted_rx.recv().unwrap(); - - // The commit owns the registry lock, so a current-generation snapshot - // cannot observe a partially committed durable transition. - assert!(matches!( - supervisor.registry.try_lock(), - Err(std::sync::TryLockError::WouldBlock) - )); - release_commit.wait(); - let (new, ()) = start_thread.join().unwrap().unwrap(); - let snapshot = snapshot_thread.join().unwrap(); - - assert!(old.cancellation_token().is_cancelled()); - assert_eq!(snapshot.unwrap().generation, new.generation); - assert!(supervisor.is_current(&new)); -} - -#[test] -fn losing_transactional_commit_does_not_cancel_a_concurrent_winner() { - let supervisor = WorkflowRunSupervisor::default(); - let held = Arc::new(std::sync::Mutex::new(Vec::new())); - let held_old = held.clone(); - let old = supervisor.start( - "run", - |_| futures::future::pending(), - move |future| held_old.lock().unwrap().push(future), - ); - let held_new = held.clone(); - let gate = Arc::new(std::sync::Barrier::new(2)); - let thread_gate = gate.clone(); - let supervisor_for_loser = supervisor.clone(); - let loser = std::thread::spawn(move || { - supervisor_for_loser.start_scoped_with_commit( - "run", - move || { - thread_gate.wait(); - Err::<(), _>("compare-and-swap lost") - }, - |_| futures::future::pending(), - |future| future, - |future| future, - |_, future| drop(future), - ) - }); - gate.wait(); - let winner = supervisor.resume_scoped_with_commit( - "run", - || Ok::<_, ()>(()), - |_| futures::future::pending(), - |future| future, - |future| future, - move |_, future| held_new.lock().unwrap().push(future), - ); - - assert_eq!(loser.join().unwrap().unwrap_err(), "compare-and-swap lost"); - let (winner_handle, ()) = winner.unwrap(); - assert!(old.cancellation_token().is_cancelled()); - assert!(!winner_handle.cancellation_token().is_cancelled()); - assert!(supervisor.is_current(&winner_handle)); -} - -#[test] -fn with_current_skips_stale_callback_and_serializes_current_callback() { - let supervisor = WorkflowRunSupervisor::default(); - let held = Arc::new(std::sync::Mutex::new(Vec::new())); - let held_old = held.clone(); - let old = supervisor.start( - "run", - |_| futures::future::pending(), - move |future| held_old.lock().unwrap().push(future), - ); - let called = Arc::new(std::sync::atomic::AtomicBool::new(false)); - let called_flag = called.clone(); - let held_current = held.clone(); - let current = supervisor.resume( - "run", - |_| futures::future::pending(), - move |future| held_current.lock().unwrap().push(future), - ); - - assert_eq!( - supervisor.with_current(&old, move || { - called_flag.store(true, std::sync::atomic::Ordering::SeqCst); - }), - None - ); - assert!(!called.load(std::sync::atomic::Ordering::SeqCst)); - - let callback_entered = Arc::new(std::sync::Barrier::new(2)); - let release_callback = Arc::new(std::sync::Barrier::new(2)); - let entered = callback_entered.clone(); - let release = release_callback.clone(); - let supervisor_for_callback = supervisor.clone(); - let handle = current.clone(); - let callback_thread = std::thread::spawn(move || { - supervisor_for_callback.with_current(&handle, move || { - entered.wait(); - release.wait(); - }) - }); - callback_entered.wait(); - - let replacement_done = Arc::new(std::sync::atomic::AtomicBool::new(false)); - let done = replacement_done.clone(); - let held_replacement = held.clone(); - let (replacement_attempted_tx, replacement_attempted_rx) = std::sync::mpsc::channel(); - let supervisor_for_replacement = supervisor.clone(); - let replacement_thread = std::thread::spawn(move || { - replacement_attempted_tx.send(()).unwrap(); - let replacement = supervisor_for_replacement.start( - "run", - |_| futures::future::pending(), - move |future| held_replacement.lock().unwrap().push(future), - ); - done.store(true, std::sync::atomic::Ordering::SeqCst); - replacement - }); - - replacement_attempted_rx.recv().unwrap(); - assert!(matches!( - supervisor.registry.try_lock(), - Err(std::sync::TryLockError::WouldBlock) - )); - assert!(!replacement_done.load(std::sync::atomic::Ordering::SeqCst)); - release_callback.wait(); - assert_eq!(callback_thread.join().unwrap(), Some(())); - let replacement = replacement_thread.join().unwrap(); - assert!(current.cancellation_token().is_cancelled()); - assert!(supervisor.is_current(&replacement)); -} diff --git a/crates/tinyagents-orchestration/src/workflow/tests.rs b/crates/tinyagents-orchestration/src/workflow/tests.rs deleted file mode 100644 index 22a7e6e1c..000000000 --- a/crates/tinyagents-orchestration/src/workflow/tests.rs +++ /dev/null @@ -1,853 +0,0 @@ -//! Tests for workflow execution: scheduling, phase transitions, concurrency, -//! and error handling. - -use std::collections::BTreeMap; -use std::collections::HashMap; -use std::sync::Arc; -use std::sync::atomic::{AtomicUsize, Ordering}; -use std::time::Duration; - -use async_trait::async_trait; -use chrono::Utc; -use parking_lot::Mutex; -use serde_json::json; -use tinyagents_graph::CollectingSink; -use tinyagents_harness::CancellationToken; -use tinyagents_session::run_ledger::{ - WorkflowLeaseClaim, WorkflowRun, WorkflowRunStatus, WorkflowRunUpsert, -}; - -use super::state::set_phase_status; -use super::*; - -pub(super) fn definition() -> WorkflowDefinition { - WorkflowDefinition { - id: "test".into(), - name: "Test".into(), - description: "test workflow".into(), - phases: vec![ - WorkflowPhase { - name: "plan".into(), - description: "plan".into(), - agent_ids: vec!["planner".into()], - depends_on: vec![], - }, - WorkflowPhase { - name: "research".into(), - description: "research".into(), - agent_ids: vec!["researcher".into(), "researcher".into()], - depends_on: vec!["plan".into()], - }, - WorkflowPhase { - name: "synthesize".into(), - description: "synthesize".into(), - agent_ids: vec!["writer".into()], - depends_on: vec!["research".into()], - }, - ], - default_concurrency: 2, - max_children: 8, - extensions: BTreeMap::new(), - } -} - -#[derive(Default)] -pub(super) struct MemoryStore(Mutex<HashMap<String, WorkflowRun>>); - -impl WorkflowStore for MemoryStore { - fn load(&self, id: &str) -> Result<Option<WorkflowRun>, OrchestrationError> { - Ok(self.0.lock().get(id).cloned()) - } - - fn upsert(&self, update: WorkflowRunUpsert) -> Result<WorkflowRun, OrchestrationError> { - let now = Utc::now(); - let prior = self.0.lock().get(&update.id).cloned(); - let row = WorkflowRun { - id: update.id.clone(), - definition_id: update.definition_id, - parent_thread_id: update.parent_thread_id, - input: update.input, - phase_states: update.phase_states, - child_run_ids: update.child_run_ids, - status: update.status, - summary: update - .summary - .or_else(|| prior.as_ref().and_then(|row| row.summary.clone())), - started_at: update - .started_at - .unwrap_or_else(|| prior.as_ref().map(|row| row.started_at).unwrap_or(now)), - updated_at: now, - completed_at: update - .completed_at - .or_else(|| prior.as_ref().and_then(|row| row.completed_at)), - revision: prior.as_ref().map_or(0, |row| row.revision + 1), - lease_owner: prior.as_ref().and_then(|row| row.lease_owner.clone()), - lease_expires_at: prior.as_ref().and_then(|row| row.lease_expires_at), - }; - self.0.lock().insert(row.id.clone(), row.clone()); - Ok(row) - } - - fn claim( - &self, - id: &str, - owner: &str, - lease_for: Duration, - ) -> Result<WorkflowLeaseClaim, OrchestrationError> { - let mut rows = self.0.lock(); - let Some(row) = rows.get_mut(id) else { - return Ok(WorkflowLeaseClaim::Missing); - }; - let now = Utc::now(); - if row - .lease_owner - .as_deref() - .is_some_and(|current| current != owner) - && row.lease_expires_at.is_some_and(|until| until > now) - { - return Ok(WorkflowLeaseClaim::Busy(row.clone())); - } - row.lease_owner = Some(owner.to_owned()); - row.lease_expires_at = chrono::Duration::from_std(lease_for) - .ok() - .map(|duration| now + duration); - row.revision += 1; - Ok(WorkflowLeaseClaim::Acquired(row.clone())) - } - - fn compare_and_swap( - &self, - update: WorkflowRunUpsert, - expected_revision: u64, - owner: &str, - lease_for: Duration, - ) -> Result<Option<WorkflowRun>, OrchestrationError> { - let mut rows = self.0.lock(); - let Some(prior) = rows.get(&update.id).cloned() else { - return Ok(None); - }; - if prior.revision != expected_revision || prior.lease_owner.as_deref() != Some(owner) { - return Ok(None); - } - let now = Utc::now(); - let row = WorkflowRun { - id: update.id, - definition_id: update.definition_id, - parent_thread_id: update.parent_thread_id, - input: update.input, - phase_states: update.phase_states, - child_run_ids: update.child_run_ids, - status: update.status, - summary: update.summary.or(prior.summary), - started_at: update.started_at.unwrap_or(prior.started_at), - updated_at: now, - completed_at: update.completed_at.or(prior.completed_at), - revision: prior.revision + 1, - lease_owner: (!update.status.is_terminal()).then(|| owner.to_owned()), - lease_expires_at: (!update.status.is_terminal()) - .then(|| { - chrono::Duration::from_std(lease_for) - .ok() - .map(|duration| now + duration) - }) - .flatten(), - }; - rows.insert(row.id.clone(), row.clone()); - Ok(Some(row)) - } - - fn renew( - &self, - id: &str, - owner: &str, - lease_for: Duration, - ) -> Result<bool, OrchestrationError> { - let mut rows = self.0.lock(); - let Some(row) = rows.get_mut(id) else { - return Ok(false); - }; - let now = Utc::now(); - if row.lease_owner.as_deref() != Some(owner) - || row.lease_expires_at.is_none_or(|expires| expires <= now) - { - return Ok(false); - } - row.lease_expires_at = chrono::Duration::from_std(lease_for) - .ok() - .map(|duration| now + duration); - Ok(true) - } -} - -#[derive(Default)] -struct RenewFailStore(MemoryStore); - -impl WorkflowStore for RenewFailStore { - fn load(&self, id: &str) -> Result<Option<WorkflowRun>, OrchestrationError> { - self.0.load(id) - } - - fn upsert(&self, update: WorkflowRunUpsert) -> Result<WorkflowRun, OrchestrationError> { - self.0.upsert(update) - } - - fn claim( - &self, - id: &str, - owner: &str, - lease_for: Duration, - ) -> Result<WorkflowLeaseClaim, OrchestrationError> { - self.0.claim(id, owner, lease_for) - } - - fn compare_and_swap( - &self, - update: WorkflowRunUpsert, - expected_revision: u64, - owner: &str, - lease_for: Duration, - ) -> Result<Option<WorkflowRun>, OrchestrationError> { - self.0 - .compare_and_swap(update, expected_revision, owner, lease_for) - } - - fn renew( - &self, - _id: &str, - _owner: &str, - _lease_for: Duration, - ) -> Result<bool, OrchestrationError> { - Ok(false) - } -} - -#[derive(Default)] -pub(super) struct FakeExecutor { - pub(super) calls: Mutex<Vec<WorkflowChildRequest>>, - active: AtomicUsize, - peak: AtomicUsize, - fail_agent: Mutex<Option<String>>, - cancelled: AtomicUsize, -} - -#[derive(Default)] -pub(super) struct BlockingExecutor { - started: tokio::sync::Notify, - pub(super) cancelled: Mutex<Vec<String>>, - calls: AtomicUsize, -} - -#[async_trait] -impl WorkflowExecutor for BlockingExecutor { - async fn execute( - &self, - request: WorkflowChildRequest, - cancel: CancellationToken, - registration: Arc<dyn WorkflowChildRegistration>, - ) -> Result<WorkflowChildResult, OrchestrationError> { - let id = format!("live-{}-{}", request.phase, request.index_in_phase); - registration.register(id.clone())?; - self.calls.fetch_add(1, Ordering::SeqCst); - self.started.notify_one(); - cancel.cancelled().await; - Err(OrchestrationError("cancelled while child was live".into())) - } - - async fn cancel_children(&self, child_ids: &[String]) { - self.cancelled.lock().extend(child_ids.iter().cloned()); - } -} - -#[async_trait] -impl WorkflowExecutor for FakeExecutor { - async fn execute( - &self, - request: WorkflowChildRequest, - cancel: CancellationToken, - registration: Arc<dyn WorkflowChildRegistration>, - ) -> Result<WorkflowChildResult, OrchestrationError> { - if cancel.is_cancelled() { - return Err(OrchestrationError("cancelled".into())); - } - let active = self.active.fetch_add(1, Ordering::SeqCst) + 1; - self.peak.fetch_max(active, Ordering::SeqCst); - self.calls.lock().push(request.clone()); - registration.register(format!("{}-{}", request.phase, request.index_in_phase))?; - tokio::task::yield_now().await; - self.active.fetch_sub(1, Ordering::SeqCst); - if self.fail_agent.lock().as_deref() == Some(request.agent_id.as_str()) { - return Err(OrchestrationError("child failure".into())); - } - Ok(WorkflowChildResult { - child_id: format!("{}-{}", request.phase, request.index_in_phase), - output: json!(format!("{} output", request.phase)), - }) - } - - async fn cancel_children(&self, _child_ids: &[String]) { - self.cancelled.fetch_add(1, Ordering::SeqCst); - } -} - -fn engine() -> ( - Arc<MemoryStore>, - Arc<FakeExecutor>, - WorkflowEngine<MemoryStore, FakeExecutor>, -) { - let store = Arc::new(MemoryStore::default()); - let executor = Arc::new(FakeExecutor::default()); - let engine = WorkflowEngine::new(store.clone(), executor.clone()); - (store, executor, engine) -} - -#[test] -fn structural_validation_covers_invalid_definitions() { - let mut empty = definition(); - empty.phases.clear(); - assert_eq!(validate_structure(&empty), vec![DefinitionError::NoPhases]); - let mut bad = definition(); - bad.default_concurrency = 0; - bad.max_children = 0; - bad.phases[1].name = "plan".into(); - bad.phases[2].depends_on = vec!["missing".into()]; - let errors = validate_structure(&bad); - assert!( - errors - .iter() - .any(|error| matches!(error, DefinitionError::DuplicatePhase { .. })) - ); - assert!( - errors - .iter() - .any(|error| matches!(error, DefinitionError::UnknownDependency { .. })) - ); - assert!( - errors - .iter() - .any(|error| matches!(error, DefinitionError::InvalidConcurrency { .. })) - ); - let mut cyclic = definition(); - cyclic.phases[0].depends_on = vec!["synthesize".into()]; - assert!(validate_structure(&cyclic).contains(&DefinitionError::CyclicDependency)); -} - -#[test] -fn scheduler_topology_preview_exposes_dispatch_run_and_done() { - let topology = scheduler_topology_preview().expect("topology"); - let nodes = topology - .nodes - .iter() - .map(|node| node.id.as_str()) - .collect::<Vec<_>>(); - assert!(nodes.contains(&"dispatch") && nodes.contains(&"run_phase") && nodes.contains(&"done")); -} - -#[tokio::test] -async fn engine_runs_in_deterministic_dependency_order_and_threads_context() { - let (store, executor, engine) = engine(); - let def = definition(); - engine - .initialise("run".into(), &def, json!({"question":"q"}), None) - .unwrap(); - engine - .drive("run", &def, CancellationToken::new()) - .await - .unwrap(); - let run = store.load("run").unwrap().unwrap(); - assert_eq!(run.status, WorkflowRunStatus::Completed); - let calls = executor.calls.lock(); - assert_eq!( - calls - .iter() - .map(|call| call.phase.as_str()) - .collect::<Vec<_>>(), - vec!["plan", "research", "research", "synthesize"] - ); - assert!( - calls - .last() - .unwrap() - .prompt - .contains("Context from prior phases") - ); - assert!(run.summary.unwrap().contains("synthesize output")); -} - -#[tokio::test] -async fn engine_respects_concurrency_global_cap_and_partial_failure() { - let (store, executor, engine) = engine(); - let mut def = definition(); - def.default_concurrency = 1; - engine - .initialise("run".into(), &def, json!("q"), None) - .unwrap(); - engine - .drive("run", &def, CancellationToken::new()) - .await - .unwrap(); - assert!(executor.peak.load(Ordering::SeqCst) <= 1); - let mut cap = definition(); - cap.max_children = 2; - engine - .initialise("cap".into(), &cap, json!("q"), None) - .unwrap(); - engine - .drive("cap", &cap, CancellationToken::new()) - .await - .unwrap(); - assert_eq!( - store.load("cap").unwrap().unwrap().status, - WorkflowRunStatus::Failed - ); - *executor.fail_agent.lock() = Some("planner".into()); - let failing = definition(); - engine - .initialise("failed".into(), &failing, json!("q"), None) - .unwrap(); - engine - .drive("failed", &failing, CancellationToken::new()) - .await - .unwrap(); - assert_eq!( - store.load("failed").unwrap().unwrap().status, - WorkflowRunStatus::Failed - ); -} - -#[tokio::test] -async fn cancellation_and_resume_do_not_repeat_completed_phases() { - let (store, executor, engine) = engine(); - let def = definition(); - engine - .initialise("run".into(), &def, json!("q"), None) - .unwrap(); - let cancel = CancellationToken::new(); - cancel.cancel(); - engine.drive("run", &def, cancel).await.unwrap(); - assert_eq!( - store.load("run").unwrap().unwrap().status, - WorkflowRunStatus::Interrupted - ); - let mut states = init_phase_states(&def); - set_phase_status( - &mut states, - "plan", - PhaseStatus::Completed, - Some(json!([{ "output": "already" }])), - ); - let run = store.load("run").unwrap().unwrap(); - store - .upsert(WorkflowRunUpsert { - id: run.id, - definition_id: run.definition_id, - parent_thread_id: run.parent_thread_id, - input: run.input, - phase_states: states, - child_run_ids: vec!["old".into()], - status: WorkflowRunStatus::Running, - summary: None, - started_at: Some(run.started_at), - completed_at: None, - }) - .unwrap(); - engine - .drive("run", &def, CancellationToken::new()) - .await - .unwrap(); - assert!( - !executor - .calls - .lock() - .iter() - .any(|call| call.phase == "plan") - ); - assert_eq!( - store.load("run").unwrap().unwrap().status, - WorkflowRunStatus::Completed - ); -} - -#[tokio::test] -async fn cancellation_after_workers_start_cancels_durably_registered_children() { - let store = Arc::new(MemoryStore::default()); - let executor = Arc::new(BlockingExecutor::default()); - let engine = Arc::new(WorkflowEngine::new(store.clone(), executor.clone())); - let mut def = definition(); - def.phases = vec![WorkflowPhase { - name: "live".into(), - description: "live".into(), - agent_ids: vec!["one".into(), "two".into()], - depends_on: vec![], - }]; - def.default_concurrency = 2; - engine - .initialise("run".into(), &def, json!("q"), None) - .unwrap(); - let cancel = CancellationToken::new(); - let drive = { - let engine = engine.clone(); - let def = def.clone(); - let cancel = cancel.clone(); - tokio::spawn(async move { engine.drive("run", &def, cancel).await }) - }; - executor.started.notified().await; - // The notification occurs only after register(), so this races precisely - // the old orphan window between real child spawn and ledger persistence. - cancel.cancel(); - drive.await.unwrap().unwrap(); - let run = store.load("run").unwrap().unwrap(); - assert_eq!(run.status, WorkflowRunStatus::Interrupted); - assert_eq!( - run.phase_states["live"]["status"], - json!("pending"), - "an interrupted in-flight phase must be retryable on resume" - ); - assert!(!run.child_run_ids.is_empty()); - let cancelled = executor.cancelled.lock().clone(); - for id in &run.child_run_ids { - assert!(cancelled.contains(id), "missing cancellation for {id}"); - } -} - -#[tokio::test] -async fn concurrent_drives_acquire_one_lease_and_do_not_duplicate_children() { - let (store, executor, engine) = engine(); - let def = definition(); - engine - .initialise("run".into(), &def, json!("q"), None) - .unwrap(); - let engine = Arc::new(engine); - let first = { - let engine = engine.clone(); - let def = def.clone(); - tokio::spawn(async move { engine.drive("run", &def, CancellationToken::new()).await }) - }; - let second = { - let engine = engine.clone(); - let def = def.clone(); - tokio::spawn(async move { engine.drive("run", &def, CancellationToken::new()).await }) - }; - first.await.unwrap().unwrap(); - second.await.unwrap().unwrap(); - assert_eq!( - store.load("run").unwrap().unwrap().status, - WorkflowRunStatus::Completed - ); - assert_eq!(executor.calls.lock().len(), 4, "one driver owns all phases"); -} - -#[tokio::test] -async fn expired_owner_takeover_resets_running_phase_and_retries_once() { - let (store, executor, engine) = engine(); - let mut def = definition(); - def.phases = vec![WorkflowPhase { - name: "recover".into(), - description: "recover".into(), - agent_ids: vec!["worker".into()], - depends_on: vec![], - }]; - engine - .initialise("takeover".into(), &def, json!("q"), None) - .unwrap(); - - // Simulate a process death after the owner has persisted `running`, but - // before it can complete or reset the phase. - let old = match store - .claim("takeover", "crashed-owner", Duration::from_millis(3)) - .unwrap() - { - WorkflowLeaseClaim::Acquired(run) => run, - other => panic!("expected lease, got {other:?}"), - }; - store - .compare_and_swap( - WorkflowRunUpsert { - id: old.id.clone(), - definition_id: old.definition_id.clone(), - parent_thread_id: old.parent_thread_id.clone(), - input: old.input.clone(), - phase_states: json!({ - "recover": {"status": "running", "outputs": []} - }), - child_run_ids: vec!["recover-0".into()], - status: WorkflowRunStatus::Running, - summary: None, - started_at: Some(old.started_at), - completed_at: None, - }, - old.revision, - "crashed-owner", - Duration::from_millis(3), - ) - .unwrap() - .expect("crashed owner persists phase start"); - tokio::time::sleep(Duration::from_millis(10)).await; - - engine - .drive("takeover", &def, CancellationToken::new()) - .await - .expect("new owner retries reclaimed phase"); - - let run = store.load("takeover").unwrap().expect("run remains"); - assert_eq!(run.status, WorkflowRunStatus::Completed); - assert_eq!(run.phase_states["recover"]["status"], json!("completed")); - assert_eq!( - executor - .calls - .lock() - .iter() - .filter(|call| call.phase == "recover") - .count(), - 1, - "lease takeover must schedule the reclaimed phase once" - ); - assert_eq!( - run.child_run_ids - .iter() - .filter(|id| id.as_str() == "recover-0") - .count(), - 1, - "the retry must not duplicate a durably registered child id" - ); -} - -#[tokio::test] -async fn heartbeat_renews_a_short_lease_while_a_child_is_running() { - let store = Arc::new(MemoryStore::default()); - let executor = Arc::new(BlockingExecutor::default()); - let engine = Arc::new( - WorkflowEngine::new(store.clone(), executor.clone()) - .with_lease_duration(Duration::from_millis(30)), - ); - let mut def = definition(); - def.phases = vec![WorkflowPhase { - name: "live".into(), - description: "live".into(), - agent_ids: vec!["one".into()], - depends_on: vec![], - }]; - engine - .initialise("run".into(), &def, json!("q"), None) - .unwrap(); - let cancel = CancellationToken::new(); - let first = { - let engine = engine.clone(); - let def = def.clone(); - let cancel = cancel.clone(); - tokio::spawn(async move { engine.drive("run", &def, cancel).await }) - }; - executor.started.notified().await; - tokio::time::sleep(Duration::from_millis(75)).await; - // This exceeds the original lease, so it only remains busy if the - // in-flight driver's heartbeat kept renewing it. - engine - .drive("run", &def, CancellationToken::new()) - .await - .unwrap(); - assert_eq!(executor.calls.load(Ordering::SeqCst), 1); - cancel.cancel(); - first.await.unwrap().unwrap(); -} - -#[tokio::test] -async fn lost_heartbeat_cancels_registered_children_and_fails_closed() { - let store = Arc::new(RenewFailStore::default()); - let executor = Arc::new(BlockingExecutor::default()); - let engine = Arc::new( - WorkflowEngine::new(store.clone(), executor.clone()) - .with_lease_duration(Duration::from_millis(30)), - ); - let mut def = definition(); - def.phases = vec![WorkflowPhase { - name: "live".into(), - description: "live".into(), - agent_ids: vec!["one".into()], - depends_on: vec![], - }]; - engine - .initialise("run".into(), &def, json!("q"), None) - .unwrap(); - let drive = { - let engine = engine.clone(); - let def = def.clone(); - tokio::spawn(async move { engine.drive("run", &def, CancellationToken::new()).await }) - }; - executor.started.notified().await; - let error = drive - .await - .unwrap() - .expect_err("renewal loss must fail closed"); - assert!(error.0.contains("lease renewal failed")); - assert!( - executor - .cancelled - .lock() - .iter() - .any(|id| id == "live-live-0"), - "the child registered before waiting must be cancelled" - ); -} - -#[tokio::test] -async fn terminal_events_are_truthful_and_flushed() { - let (store, executor, engine) = engine(); - let sink = Arc::new(CollectingSink::new()); - let engine = engine.with_event_sink(sink.clone()); - let def = definition(); - engine - .initialise("ok".into(), &def, json!("q"), None) - .unwrap(); - engine - .drive("ok", &def, CancellationToken::new()) - .await - .unwrap(); - assert!( - sink.events() - .iter() - .any(|event| matches!(event, tinyagents_graph::GraphEvent::RunCompleted { .. })) - ); - - *executor.fail_agent.lock() = Some("planner".into()); - engine - .initialise("failed".into(), &def, json!("q"), None) - .unwrap(); - engine - .drive("failed", &def, CancellationToken::new()) - .await - .unwrap(); - assert!(sink.events().iter().any(|event| matches!( - event, - tinyagents_graph::GraphEvent::RunFailed { error, .. } if error.contains("child failure") - ))); - assert_eq!( - store.load("failed").unwrap().unwrap().status, - WorkflowRunStatus::Failed - ); -} - -#[tokio::test] -async fn fenced_driver_exits_silently_when_a_replacement_is_running() { - let store = Arc::new(RenewFailStore::default()); - let executor = Arc::new(BlockingExecutor::default()); - let sink = Arc::new(CollectingSink::new()); - let engine = Arc::new( - WorkflowEngine::new(store.clone(), executor.clone()) - .with_lease_duration(Duration::from_millis(30)) - .with_event_sink(sink.clone()), - ); - let mut def = definition(); - def.phases = vec![WorkflowPhase { - name: "live".into(), - description: "live".into(), - agent_ids: vec!["one".into()], - depends_on: vec![], - }]; - engine - .initialise("handoff".into(), &def, json!("q"), None) - .unwrap(); - let first = { - let engine = engine.clone(); - let def = def.clone(); - tokio::spawn(async move { - engine - .drive("handoff", &def, CancellationToken::new()) - .await - }) - }; - executor.started.notified().await; - - // A replacement owner acquires after expiry before the old driver's - // heartbeat sees its loss. The old loop must not flush a false failure. - let current = store.load("handoff").unwrap().unwrap(); - store - .compare_and_swap( - WorkflowRunUpsert { - id: current.id.clone(), - definition_id: current.definition_id.clone(), - parent_thread_id: current.parent_thread_id.clone(), - input: current.input.clone(), - phase_states: current.phase_states.clone(), - child_run_ids: current.child_run_ids.clone(), - status: WorkflowRunStatus::Running, - summary: None, - started_at: Some(current.started_at), - completed_at: None, - }, - current.revision, - current.lease_owner.as_deref().unwrap(), - Duration::from_millis(3), - ) - .unwrap() - .unwrap(); - tokio::time::sleep(Duration::from_millis(5)).await; - match store - .claim("handoff", "replacement", Duration::from_secs(1)) - .unwrap() - { - WorkflowLeaseClaim::Acquired(run) => { - assert_eq!(run.lease_owner.as_deref(), Some("replacement")); - } - other => panic!("expected replacement lease, got {other:?}"), - } - first.await.unwrap().expect("fenced driver exits cleanly"); - assert!( - !sink - .events() - .iter() - .any(|event| matches!(event, tinyagents_graph::GraphEvent::RunFailed { .. })), - "a fenced driver must not report the active replacement as failed" - ); -} - -#[test] -fn structured_outputs_are_preserved_in_context_and_summary() { - let def = definition(); - let mut states = init_phase_states(&def); - set_phase_status( - &mut states, - "plan", - PhaseStatus::Completed, - Some(json!([{ - "output": { "claims": ["a", "b"], "score": 7 } - }])), - ); - let upstream = upstream_outputs(&def.phases[1], &states); - let prompt = phase_prompt(&json!("q"), &def.phases[1], 0, &upstream); - assert!(prompt.contains(r#"{"claims":["a","b"],"score":7}"#)); - set_phase_status( - &mut states, - "synthesize", - PhaseStatus::Completed, - Some(json!([{ - "output": { "answer": "kept" } - }])), - ); - assert_eq!( - synthesize_summary(&def, &states).as_deref(), - Some(r#"{"answer":"kept"}"#) - ); -} - -#[tokio::test] -async fn output_wire_shape_remains_compatible_while_json_stays_lossless() { - let (store, _executor, engine) = engine(); - let mut def = definition(); - def.phases = vec![WorkflowPhase { - name: "only".into(), - description: "only".into(), - agent_ids: vec!["planner".into()], - depends_on: vec![], - }]; - engine - .initialise("run".into(), &def, json!("q"), None) - .unwrap(); - engine - .drive("run", &def, CancellationToken::new()) - .await - .unwrap(); - let output = store.load("run").unwrap().unwrap().phase_states["only"]["outputs"][0].clone(); - assert_eq!(output["agentId"], json!("planner")); - assert!(output["output"].is_string()); - assert_eq!(output["metadata"]["version"], json!(2)); - assert_eq!(output["metadata"]["rawOutput"], json!("only output")); -} diff --git a/crates/tinyagents-orchestration/src/workflow/types.rs b/crates/tinyagents-orchestration/src/workflow/types.rs deleted file mode 100644 index bda9d2879..000000000 --- a/crates/tinyagents-orchestration/src/workflow/types.rs +++ /dev/null @@ -1,87 +0,0 @@ -//! Workflow definition types: phases, definitions, and validation errors. - -use std::collections::BTreeMap; - -use serde::{Deserialize, Serialize}; -use serde_json::Value; - -/// One phase of a declarative workflow. -/// -/// A phase specifies a set of agents to work on it concurrently and its -/// dependencies (other phases that must complete first). Agents run in -/// parallel within a phase; the phase itself is the unit of scheduling. -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct WorkflowPhase { - /// The human-readable name of this phase. - pub name: String, - /// A description of what this phase does. - pub description: String, - /// Agent ids that should work on this phase in parallel. - pub agent_ids: Vec<String>, - /// Phase names this phase depends on (must complete first). - pub depends_on: Vec<String>, -} - -/// A host-neutral declarative phase DAG. -/// -/// Defines a workflow as a directed acyclic graph of phases, each with -/// associated agents, concurrency limits, and dependencies. The orchestration -/// engine schedules phases topologically and manages bounded parallelism via -/// `default_concurrency` and `max_children`. -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct WorkflowDefinition { - /// Unique identifier for this workflow. - pub id: String, - /// Human-readable workflow name. - pub name: String, - /// Description of what the workflow does. - pub description: String, - /// All phases in the workflow. - pub phases: Vec<WorkflowPhase>, - /// Default concurrency limit for agents in phases that don't specify one. - pub default_concurrency: u32, - /// Maximum number of child tasks that can be spawned concurrently. - pub max_children: u32, - /// Host-defined wire metadata which the orchestration engine never reads. - #[serde(flatten, default)] - pub extensions: BTreeMap<String, Value>, -} - -/// List response used by hosts that expose a workflow catalog. -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct WorkflowDefinitionListResponse { - /// The workflow definitions in this response. - pub definitions: Vec<WorkflowDefinition>, - /// Total count of workflows in the host's catalog. - pub count: usize, -} - -/// A structural or host-supplied lookup problem in a workflow definition. -/// -/// These errors detect issues that prevent a workflow from running: -/// unknown agent references, missing dependencies, cycles, or invalid -/// concurrency settings. They are host-independent validation issues. -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -#[serde(rename_all = "camelCase")] -pub enum DefinitionError { - /// A phase references an agent the host does not know about. - UnknownAgent { phase: String, agent_id: String }, - /// A phase depends on another phase that is not defined. - UnknownDependency { phase: String, depends_on: String }, - /// Multiple phases have the same name. - DuplicatePhase { name: String }, - /// A phase has no agents assigned to it. - EmptyPhase { phase: String }, - /// The phase dependencies form a cycle. - CyclicDependency, - /// The workflow has no phases. - NoPhases, - /// The concurrency settings are invalid (e.g. default > max). - InvalidConcurrency { - default_concurrency: u32, - max_children: u32, - }, -} diff --git a/crates/tinyagents-orchestration/src/workflow/validate.rs b/crates/tinyagents-orchestration/src/workflow/validate.rs deleted file mode 100644 index fb5254410..000000000 --- a/crates/tinyagents-orchestration/src/workflow/validate.rs +++ /dev/null @@ -1,71 +0,0 @@ -//! Workflow definition validation: structural checks and error detection. -//! -//! Validates a workflow definition for structural issues (missing phases, -//! duplicate names, cycles, invalid concurrency) and host-specific issues -//! (unknown agents). Errors are deterministic and host-independent. - -use tinyagents_graph::dag::{DagIssue, DagNode, validate_dag}; - -use super::{DefinitionError, WorkflowDefinition}; - -/// Validate properties that do not require a host agent registry. -pub fn validate_structure(definition: &WorkflowDefinition) -> Vec<DefinitionError> { - if definition.phases.is_empty() { - return vec![DefinitionError::NoPhases]; - } - - let mut errors = definition - .phases - .iter() - .filter(|phase| phase.agent_ids.is_empty()) - .map(|phase| DefinitionError::EmptyPhase { - phase: phase.name.clone(), - }) - .collect::<Vec<_>>(); - let nodes = definition - .phases - .iter() - .map(|phase| { - DagNode::new( - phase.name.as_str(), - phase.depends_on.iter().map(String::as_str), - ) - }) - .collect::<Vec<_>>(); - errors.extend(validate_dag(&nodes).into_iter().map(|issue| match issue { - DagIssue::DuplicateNode { id } => DefinitionError::DuplicatePhase { name: id }, - DagIssue::UnknownDependency { node, depends_on } => DefinitionError::UnknownDependency { - phase: node, - depends_on, - }, - DagIssue::Cycle => DefinitionError::CyclicDependency, - })); - if definition.default_concurrency == 0 || definition.max_children == 0 { - errors.push(DefinitionError::InvalidConcurrency { - default_concurrency: definition.default_concurrency, - max_children: definition.max_children, - }); - } - errors -} - -/// Validate agent identifiers using a host-owned registry lookup. -pub fn validate_agents<F>(definition: &WorkflowDefinition, is_known: F) -> Vec<DefinitionError> -where - F: Fn(&str) -> bool, -{ - definition - .phases - .iter() - .flat_map(|phase| { - phase - .agent_ids - .iter() - .filter(|agent_id| !is_known(agent_id)) - .map(|agent_id| DefinitionError::UnknownAgent { - phase: phase.name.clone(), - agent_id: agent_id.clone(), - }) - }) - .collect() -} diff --git a/crates/tinyagents-orchestration/tests/e2e_fuzz_graph_agents.rs b/crates/tinyagents-orchestration/tests/e2e_fuzz_graph_agents.rs deleted file mode 100644 index c382dc1ff..000000000 --- a/crates/tinyagents-orchestration/tests/e2e_fuzz_graph_agents.rs +++ /dev/null @@ -1,782 +0,0 @@ -//! Deterministic end-to-end fuzz matrix for graph/agent composition. -//! -//! These tests generate a compact matrix of graph shapes and agent behaviors: -//! whole-state (overwrite) graphs, durable partial-update linear/conditional -//! graphs, command fan-out graphs, and adapter-subgraph graphs. Graph nodes -//! drive real -//! `AgentHarness` runs whose scripted models may make ordinary tool calls, -//! sub-agent tool calls, or both. The assertions are structural: graph visit -//! paths, reducer state, model/tool event counts, and tool-result transcript -//! evidence rather than exact prose. - -use std::sync::Arc; - -use serde_json::json; - -use tinyagents_graph::ClosureStateReducer; -use tinyagents_graph::*; -use tinyagents_harness::context::{RunConfig, RunContext}; -use tinyagents_harness::events::EventSink; -use tinyagents_harness::runtime::AgentHarness; -use tinyagents_harness::testkit::{EventRecorder, FakeTool, ScriptedModel, Trajectory}; -use tinyagents_orchestration::subagent::{ - ChildDataPolicy, SubAgent, SubAgentJobRegistry, SubAgentTool, -}; -use tinyinference_llm::message::{AssistantMessage, ContentBlock, Message}; -use tinyinference_llm::model::ModelResponse; -use tinyinference_llm::tool::ToolCall; -use tinyinference_llm::usage::Usage; - -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -enum GraphShape { - WholeStateLinear, - DurableLinear, - DurableConditional, - DurableCommandFanout, - DurableAdapterSubgraph, -} - -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -struct Scenario { - id: usize, - shape: GraphShape, - parallel: bool, - branch_count: usize, - use_regular_tool: bool, - use_subagent: bool, - route_right: bool, -} - -impl Scenario { - fn label(self) -> String { - format!( - "scenario-{}::{:?}::parallel={}::branches={}::regular={}::subagent={}::right={}", - self.id, - self.shape, - self.parallel, - self.branch_count, - self.use_regular_tool, - self.use_subagent, - self.route_right - ) - } - - fn expected_tools_per_agent(self) -> usize { - usize::from(self.use_regular_tool) + usize::from(self.use_subagent) - } - - fn expected_parent_model_calls_per_agent(self) -> usize { - if self.expected_tools_per_agent() == 0 { - 1 - } else { - 2 - } - } - - fn expected_observed_model_calls_per_agent(self) -> usize { - self.expected_parent_model_calls_per_agent() + usize::from(self.use_subagent) - } -} - -#[derive(Clone, Debug, Default, PartialEq, Eq)] -struct FuzzState { - scenario_id: usize, - path: Vec<String>, - agent_answers: Vec<String>, - regular_tool_results: usize, - subagent_tool_results: usize, - model_calls: usize, - tool_calls: usize, - branch_values: Vec<i32>, - total: Option<i32>, -} - -#[derive(Clone, Debug, PartialEq, Eq)] -enum FuzzUpdate { - Mark(String), - Agent { - node: String, - answer: String, - model_calls: usize, - tool_calls: usize, - saw_regular_tool: bool, - saw_subagent_tool: bool, - }, - Branch { - node: String, - value: i32, - }, - Total(i32), -} - -fn scenarios() -> Vec<Scenario> { - vec![ - Scenario { - id: 0, - shape: GraphShape::WholeStateLinear, - parallel: false, - branch_count: 0, - use_regular_tool: false, - use_subagent: false, - route_right: false, - }, - Scenario { - id: 1, - shape: GraphShape::WholeStateLinear, - parallel: false, - branch_count: 0, - use_regular_tool: true, - use_subagent: true, - route_right: false, - }, - Scenario { - id: 2, - shape: GraphShape::DurableLinear, - parallel: false, - branch_count: 0, - use_regular_tool: true, - use_subagent: false, - route_right: false, - }, - Scenario { - id: 3, - shape: GraphShape::DurableLinear, - parallel: false, - branch_count: 0, - use_regular_tool: false, - use_subagent: true, - route_right: false, - }, - Scenario { - id: 4, - shape: GraphShape::DurableConditional, - parallel: false, - branch_count: 0, - use_regular_tool: true, - use_subagent: true, - route_right: true, - }, - Scenario { - id: 5, - shape: GraphShape::DurableConditional, - parallel: false, - branch_count: 0, - use_regular_tool: false, - use_subagent: true, - route_right: false, - }, - Scenario { - id: 6, - shape: GraphShape::DurableCommandFanout, - parallel: false, - branch_count: 2, - use_regular_tool: true, - use_subagent: false, - route_right: false, - }, - Scenario { - id: 7, - shape: GraphShape::DurableCommandFanout, - parallel: true, - branch_count: 3, - use_regular_tool: true, - use_subagent: true, - route_right: false, - }, - Scenario { - id: 8, - shape: GraphShape::DurableCommandFanout, - parallel: true, - branch_count: 4, - use_regular_tool: false, - use_subagent: true, - route_right: false, - }, - Scenario { - id: 9, - shape: GraphShape::DurableAdapterSubgraph, - parallel: false, - branch_count: 0, - use_regular_tool: true, - use_subagent: false, - route_right: false, - }, - Scenario { - id: 10, - shape: GraphShape::DurableAdapterSubgraph, - parallel: false, - branch_count: 0, - use_regular_tool: true, - use_subagent: true, - route_right: false, - }, - ] -} - -fn reducer() --> ClosureStateReducer<FuzzState, FuzzUpdate, impl Fn(FuzzState, FuzzUpdate) -> Result<FuzzState>> { - ClosureStateReducer::new(|mut state: FuzzState, update: FuzzUpdate| { - match update { - FuzzUpdate::Mark(node) => state.path.push(node), - FuzzUpdate::Agent { - node, - answer, - model_calls, - tool_calls, - saw_regular_tool, - saw_subagent_tool, - } => { - state.path.push(node); - state.agent_answers.push(answer); - state.model_calls += model_calls; - state.tool_calls += tool_calls; - if saw_regular_tool { - state.regular_tool_results += 1; - } - if saw_subagent_tool { - state.subagent_tool_results += 1; - } - } - FuzzUpdate::Branch { node, value } => { - state.path.push(node); - state.branch_values.push(value); - } - FuzzUpdate::Total(total) => { - state.path.push("aggregate".to_string()); - state.total = Some(total); - } - } - Ok(state) - }) -} - -fn tool_call_response(calls: Vec<ToolCall>) -> ModelResponse { - ModelResponse { - message: AssistantMessage { - id: Some("tool-turn".to_string()), - content: Vec::new(), - tool_calls: calls, - usage: Some(Usage::new(7, 3)), - origin: None, - }, - usage: Some(Usage::new(7, 3)), - finish_reason: Some("tool_calls".to_string()), - raw: None, - resolved_model: None, - continue_turn: None, - served_from_cache: false, - correlation: None, - resolved_route: None, - } -} - -fn text_response(text: impl Into<String>) -> ModelResponse { - ModelResponse { - message: AssistantMessage { - id: None, - content: vec![ContentBlock::Text(text.into())], - tool_calls: Vec::new(), - usage: Some(Usage::new(5, 2)), - origin: None, - }, - usage: Some(Usage::new(5, 2)), - finish_reason: Some("stop".to_string()), - raw: None, - resolved_model: None, - continue_turn: None, - served_from_cache: false, - correlation: None, - resolved_route: None, - } -} - -fn child_harness(answer: String) -> AgentHarness<()> { - let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("child", Arc::new(ScriptedModel::replies(vec![answer]))); - harness -} - -fn parent_harness(scenario: Scenario, node: &str) -> (AgentHarness<()>, SubAgentJobRegistry) { - let mut harness: AgentHarness<()> = AgentHarness::new(); - let jobs = SubAgentJobRegistry::new(); - - let mut calls = Vec::new(); - if scenario.use_regular_tool { - harness.register_tool(Arc::new(FakeTool::returning( - "lookup", - format!("lookup:{}:{node}", scenario.id), - ))); - calls.push(ToolCall::new( - format!("{}-{node}-lookup", scenario.id), - "lookup", - json!({ "node": node }), - )); - } - if scenario.use_subagent { - let child = Arc::new(SubAgent::new( - "delegate", - "delegate deterministic work to a child agent", - Arc::new(child_harness(format!("child:{}:{node}", scenario.id))), - )); - harness.register_tool_dispatch(Arc::new( - SubAgentTool::new(child, ChildDataPolicy::new(|parent: &()| *parent)) - .with_job_registry(jobs.clone()), - )); - calls.push(ToolCall::new( - format!("{}-{node}-delegate", scenario.id), - "delegate", - json!({ "input": format!("child task for {node}") }), - )); - } - - let final_answer = format!("final:{}:{node}", scenario.id); - let responses = if calls.is_empty() { - vec![text_response(final_answer)] - } else { - vec![tool_call_response(calls), text_response(final_answer)] - }; - harness.register_model("parent", Arc::new(ScriptedModel::new(responses))); - (harness, jobs) -} - -async fn run_agent_node( - scenario: Scenario, - node: &'static str, - events: EventSink, -) -> Result<FuzzUpdate> { - let (harness, jobs) = parent_harness(scenario, node); - let ctx = RunContext::new(RunConfig::new(format!("agent-{}-{node}", scenario.id)), ()) - .with_events(events); - - let run = harness - .invoke_in_context( - &(), - ctx, - vec![Message::user(format!("run {} in {node}", scenario.label()))], - ) - .await?; - - let saw_regular_tool = run - .messages - .iter() - .any(|message| matches!(message, Message::Tool(_)) && message.text().contains("lookup:")); - let saw_subagent_tool = if scenario.use_subagent { - let job = tokio::time::timeout(std::time::Duration::from_secs(1), async { - loop { - if let Some(job) = jobs.list().into_iter().next() - && job.status.is_terminal() - { - return job; - } - tokio::task::yield_now().await; - } - }) - .await - .expect("spawned fuzz subagent job reaches a terminal state"); - assert_eq!( - job.output.as_deref(), - Some(format!("child:{}:{node}", scenario.id).as_str()) - ); - true - } else { - false - }; - - Ok(FuzzUpdate::Agent { - node: node.to_string(), - answer: run.text().unwrap_or_default(), - model_calls: run.model_calls, - tool_calls: run.tool_calls, - saw_regular_tool, - saw_subagent_tool, - }) -} - -async fn run_whole_state_linear(scenario: Scenario, recorder: &EventRecorder) -> Result<FuzzState> { - let events = recorder.sink(); - // A whole-state (overwrite) durable graph: each node returns the full next - // state. The agent node folds its update through the shared reducer itself. - let graph = GraphBuilder::<FuzzState, FuzzState>::overwrite() - .add_node( - "prep", - |mut state: FuzzState, _ctx: NodeContext| async move { - state.path.push("prep".to_string()); - Ok(NodeResult::Update(state)) - }, - ) - .add_node("agent", move |state: FuzzState, _ctx: NodeContext| { - let events = events.clone(); - async move { - let update = run_agent_node(scenario, "agent", events).await?; - let state = reducer().apply(state, update)?; - Ok(NodeResult::Update(state)) - } - }) - .set_entry("prep") - .add_edge("prep", "agent") - .set_finish("agent") - .compile()?; - - let run = graph - .run(FuzzState { - scenario_id: scenario.id, - ..FuzzState::default() - }) - .await?; - - assert_visited(&run.visited, &["prep", "agent"], &scenario.label()); - Ok(run.state) -} - -async fn run_durable_scenario(scenario: Scenario, recorder: &EventRecorder) -> Result<FuzzState> { - match scenario.shape { - GraphShape::DurableLinear => run_durable_linear(scenario, recorder).await, - GraphShape::DurableConditional => run_durable_conditional(scenario, recorder).await, - GraphShape::DurableCommandFanout => run_durable_command_fanout(scenario, recorder).await, - GraphShape::DurableAdapterSubgraph => { - run_durable_adapter_subgraph(scenario, recorder).await - } - GraphShape::WholeStateLinear => { - unreachable!("whole-state scenarios run via run_whole_state_linear") - } - } -} - -async fn run_durable_linear(scenario: Scenario, recorder: &EventRecorder) -> Result<FuzzState> { - let events = recorder.sink(); - let graph = GraphBuilder::<FuzzState, FuzzUpdate>::new() - .set_reducer(reducer()) - .add_node("prep", |_state: FuzzState, _ctx: NodeContext| async move { - Ok(NodeResult::Update(FuzzUpdate::Mark("prep".to_string()))) - }) - .add_node("agent", move |_state: FuzzState, _ctx: NodeContext| { - let events = events.clone(); - async move { - Ok(NodeResult::Update( - run_agent_node(scenario, "agent", events).await?, - )) - } - }) - .set_entry("prep") - .add_edge("prep", "agent") - .set_finish("agent") - .compile()?; - - let run = graph - .run(FuzzState { - scenario_id: scenario.id, - ..FuzzState::default() - }) - .await?; - assert_visited(&run.visited, &["prep", "agent"], &scenario.label()); - Ok(run.state) -} - -async fn run_durable_conditional( - scenario: Scenario, - recorder: &EventRecorder, -) -> Result<FuzzState> { - let left_events = recorder.sink(); - let right_events = recorder.sink(); - let expected_agent = if scenario.route_right { - "agent_right" - } else { - "agent_left" - }; - - let graph = GraphBuilder::<FuzzState, FuzzUpdate>::new() - .set_reducer(reducer()) - .add_node("prep", |_state: FuzzState, _ctx: NodeContext| async move { - Ok(NodeResult::Update(FuzzUpdate::Mark("prep".to_string()))) - }) - .add_node("route", |_state: FuzzState, _ctx: NodeContext| async move { - Ok(NodeResult::Update(FuzzUpdate::Mark("route".to_string()))) - }) - .add_node("agent_left", move |_state: FuzzState, _ctx: NodeContext| { - let events = left_events.clone(); - async move { - Ok(NodeResult::Update( - run_agent_node(scenario, "agent_left", events).await?, - )) - } - }) - .add_node( - "agent_right", - move |_state: FuzzState, _ctx: NodeContext| { - let events = right_events.clone(); - async move { - Ok(NodeResult::Update( - run_agent_node(scenario, "agent_right", events).await?, - )) - } - }, - ) - .set_entry("prep") - .add_edge("prep", "route") - .add_conditional_edges( - "route", - move |_state: &FuzzState| { - if scenario.route_right { - "right".to_string() - } else { - "left".to_string() - } - }, - [("left", "agent_left"), ("right", "agent_right")], - ) - .set_finish("agent_left") - .set_finish("agent_right") - .compile()?; - - let run = graph - .run(FuzzState { - scenario_id: scenario.id, - ..FuzzState::default() - }) - .await?; - assert_visited( - &run.visited, - &["prep", "route", expected_agent], - &scenario.label(), - ); - Ok(run.state) -} - -async fn run_durable_command_fanout( - scenario: Scenario, - recorder: &EventRecorder, -) -> Result<FuzzState> { - let mut builder = GraphBuilder::<FuzzState, FuzzUpdate>::new() - .with_parallel(scenario.parallel) - .set_reducer(reducer()) - .add_node( - "dispatch", - move |_state: FuzzState, _ctx: NodeContext| async move { - let targets = (0..scenario.branch_count) - .map(|idx| format!("branch_{idx}")) - .collect::<Vec<_>>(); - Ok(NodeResult::Command(Command::default().with_goto(targets))) - }, - ) - .add_node( - "aggregate", - |_state: FuzzState, _ctx: NodeContext| async move { - Ok(NodeResult::Update(FuzzUpdate::Total(100))) - }, - ) - .set_entry("dispatch") - .mark_command_routing("dispatch") - .set_finish("aggregate"); - - for idx in 0..scenario.branch_count { - let node = format!("branch_{idx}"); - let node_for_handler = node.clone(); - let events = recorder.sink(); - builder = builder - .add_node(node.clone(), move |_state: FuzzState, _ctx: NodeContext| { - let events = events.clone(); - let node_for_update = node_for_handler.clone(); - async move { - if idx == 0 { - Ok(NodeResult::Update( - run_agent_node(scenario, "branch_0", events).await?, - )) - } else { - Ok(NodeResult::Update(FuzzUpdate::Branch { - node: node_for_update, - value: idx as i32, - })) - } - } - }) - .add_edge(node, "aggregate"); - } - - let graph = builder.compile()?; - let run = graph - .run(FuzzState { - scenario_id: scenario.id, - ..FuzzState::default() - }) - .await?; - - assert!( - run.visited.iter().any(|node| node.as_str() == "dispatch"), - "{}", - scenario.label() - ); - assert!( - run.visited.iter().any(|node| node.as_str() == "aggregate"), - "{}", - scenario.label() - ); - assert_eq!(run.state.total, Some(100), "{}", scenario.label()); - Ok(run.state) -} - -async fn run_durable_adapter_subgraph( - scenario: Scenario, - recorder: &EventRecorder, -) -> Result<FuzzState> { - let child = GraphBuilder::<FuzzState, FuzzUpdate>::new() - .set_reducer(reducer()) - .add_node( - "child_prep", - |_state: FuzzState, _ctx: NodeContext| async move { - Ok(NodeResult::Update(FuzzUpdate::Mark( - "child_prep".to_string(), - ))) - }, - ) - .add_node( - "child_done", - |_state: FuzzState, _ctx: NodeContext| async move { - Ok(NodeResult::Update(FuzzUpdate::Branch { - node: "child_done".to_string(), - value: 7, - })) - }, - ) - .set_entry("child_prep") - .add_edge("child_prep", "child_done") - .set_finish("child_done") - .compile()?; - - let events = recorder.sink(); - let graph = GraphBuilder::<FuzzState, FuzzUpdate>::new() - .set_reducer(reducer()) - .add_node( - "subgraph", - tinyagents_graph::subgraph::adapter_subgraph_node( - child, - |parent: &FuzzState| parent.clone(), - |_parent: &FuzzState, child_state: FuzzState| FuzzUpdate::Branch { - node: "subgraph".to_string(), - value: child_state.branch_values.iter().sum(), - }, - ), - ) - .add_node("agent", move |_state: FuzzState, _ctx: NodeContext| { - let events = events.clone(); - async move { - Ok(NodeResult::Update( - run_agent_node(scenario, "agent", events).await?, - )) - } - }) - .set_entry("subgraph") - .add_edge("subgraph", "agent") - .set_finish("agent") - .compile()?; - - let run = graph - .run(FuzzState { - scenario_id: scenario.id, - ..FuzzState::default() - }) - .await?; - assert_visited(&run.visited, &["subgraph", "agent"], &scenario.label()); - Ok(run.state) -} - -fn assert_visited(actual: &[tinyagents_harness::ids::NodeId], expected: &[&str], label: &str) { - let actual = actual.iter().map(ToString::to_string).collect::<Vec<_>>(); - assert_eq!(actual, expected, "{label}"); -} - -fn assert_state_invariants(scenario: Scenario, state: &FuzzState) { - assert_eq!(state.scenario_id, scenario.id, "{}", scenario.label()); - assert!( - !state.path.is_empty(), - "graph should commit at least one update for {}", - scenario.label() - ); - - let expected_agent_runs = match scenario.shape { - GraphShape::DurableCommandFanout => 1, - _ => 1, - }; - assert_eq!( - state.agent_answers.len(), - expected_agent_runs, - "{}", - scenario.label() - ); - assert_eq!( - state.model_calls, - expected_agent_runs * scenario.expected_parent_model_calls_per_agent(), - "{}", - scenario.label() - ); - assert_eq!( - state.tool_calls, - expected_agent_runs * scenario.expected_tools_per_agent(), - "{}", - scenario.label() - ); - assert_eq!( - state.regular_tool_results, - expected_agent_runs * usize::from(scenario.use_regular_tool), - "{}", - scenario.label() - ); - assert_eq!( - state.subagent_tool_results, - expected_agent_runs * usize::from(scenario.use_subagent), - "{}", - scenario.label() - ); - - if matches!(scenario.shape, GraphShape::DurableCommandFanout) { - assert_eq!(state.total, Some(100), "{}", scenario.label()); - assert_eq!( - state.branch_values.len(), - scenario.branch_count.saturating_sub(1), - "{}", - scenario.label() - ); - } - - if matches!(scenario.shape, GraphShape::DurableAdapterSubgraph) { - assert_eq!(state.branch_values, vec![7], "{}", scenario.label()); - } -} - -fn assert_event_invariants(scenario: Scenario, recorder: &EventRecorder) { - let trajectory = Trajectory::from_events(recorder.events()); - trajectory.assert_completed(); - trajectory.assert_model_called_times(scenario.expected_observed_model_calls_per_agent()); - if scenario.use_regular_tool { - trajectory.assert_tool_called("lookup"); - assert_eq!( - trajectory.tool_call_count("lookup"), - 1, - "{}", - scenario.label() - ); - } - if scenario.use_subagent { - trajectory.assert_tool_called("delegate"); - assert_eq!( - trajectory.tool_call_count("delegate"), - 1, - "{}", - scenario.label() - ); - } -} - -#[tokio::test] -async fn generated_graph_agent_scenarios_keep_e2e_invariants() { - for scenario in scenarios() { - let recorder = EventRecorder::new(); - let state = match scenario.shape { - GraphShape::WholeStateLinear => run_whole_state_linear(scenario, &recorder).await, - _ => run_durable_scenario(scenario, &recorder).await, - } - .unwrap_or_else(|error| panic!("{} failed: {error}", scenario.label())); - - assert_state_invariants(scenario, &state); - assert_event_invariants(scenario, &recorder); - } -} diff --git a/crates/tinyagents-orchestration/tests/e2e_graph_subagent_node.rs b/crates/tinyagents-orchestration/tests/e2e_graph_subagent_node.rs deleted file mode 100644 index 9b7e44b77..000000000 --- a/crates/tinyagents-orchestration/tests/e2e_graph_subagent_node.rs +++ /dev/null @@ -1,180 +0,0 @@ -//! End-to-end coverage for graph sub-agent nodes through the public registry, -//! harness, and graph execution surfaces. -// The deprecated `SubAgentNode` keeps its behaviour and its coverage. -#![allow(deprecated)] - -use std::collections::HashMap; -use std::sync::Arc; - -use async_trait::async_trait; -use serde_json::json; - -use tinyagents_graph::*; -use tinyinference_llm::usage::Usage; - -/// Host-owned test adapter for the graph's explicit agent-invocation boundary. -/// The graph receives it per run, so it cannot retain a fabricated harness -/// context or accidentally share run lineage across executions. -#[derive(Clone, Default)] -struct TestAgentInvoker { - outputs: Arc<HashMap<String, SubAgentOutput>>, -} - -impl TestAgentInvoker { - fn with_output(name: &str, output: SubAgentOutput) -> Arc<Self> { - Arc::new(Self { - outputs: Arc::new(HashMap::from([(name.to_owned(), output)])), - }) - } -} - -#[async_trait] -impl AgentInvoker for TestAgentInvoker { - async fn invoke(&self, request: AgentInvocation) -> tinyagents_graph::Result<SubAgentOutput> { - self.outputs.get(&request.agent_id).cloned().ok_or_else(|| { - TinyAgentsError::Capability(format!("unknown test agent `{}`", request.agent_id)) - }) - } -} - -fn binding(invoker: Arc<dyn AgentInvoker>) -> AgentInvocationBinding { - AgentInvocationBinding::new( - invoker, - tinyagents_harness::events::EventSink::new(), - tinyagents_harness::cancel::CancellationToken::new(), - ) -} - -fn graph_delegating_to( - node: SubAgentNode<String, String>, -) -> tinyagents_graph::CompiledGraph<String, String> { - GraphBuilder::<String, String>::overwrite() - .add_node("delegate", subagent_node(node)) - .set_entry("delegate") - .set_finish("delegate") - .compile() - .expect("graph compiles") -} - -#[tokio::test] -async fn subagent_node_delegates_records_child_run_and_forwards_events() { - let mut usage = tinyinference_llm::usage::UsageTotals::new(); - usage.record(Usage::new(7, 3)); - let invoker = TestAgentInvoker::with_output( - "researcher", - SubAgentOutput { - text: "answer: 42".to_owned(), - usage, - model_calls: 1, - ..SubAgentOutput::default() - }, - ); - let node = SubAgentNode::<String, String>::from_fns( - "researcher", - |state: &String| SubAgentInput::prompt(format!("question: {state}")), - |out: SubAgentOutput| { - assert_eq!(out.text, "answer: 42"); - assert!(out.model_calls >= 1); - assert_eq!(out.tool_calls, 0); - out.text - }, - ); - let graph = graph_delegating_to(node); - - let run = graph - .run_with_agent_binding("life?".to_string(), binding(invoker)) - .await - .expect("graph run"); - assert_eq!(run.state, "answer: 42"); - assert_eq!(run.child_runs.len(), 1); - let child = &run.child_runs[0]; - assert_eq!(child.node.as_str(), "delegate"); - assert_eq!(child.graph_id.as_str(), "agent:researcher"); - assert_eq!(child.root_run_id, run.root_run_id); - assert_ne!(child.run_id, run.run_id); - assert!(child.usage.usage.effective_total() > 0); - assert_eq!(run.run_tree().children.len(), 1); -} - -#[tokio::test] -async fn subagent_node_errors_for_missing_agent_and_budget_excess() { - let missing_node = SubAgentNode::<String, String>::from_fns( - "missing", - |state: &String| SubAgentInput::prompt(state.clone()).with_data(json!({ "source": "e2e" })), - |out: SubAgentOutput| out.text, - ); - let missing_graph = graph_delegating_to(missing_node); - let err = missing_graph.run("go".to_string()).await.unwrap_err(); - assert!(matches!(err, TinyAgentsError::Capability(_)), "got {err:?}"); - - let invoker = TestAgentInvoker::with_output( - "twostep", - SubAgentOutput { - text: "final".to_owned(), - model_calls: 2, - tool_calls: 1, - ..SubAgentOutput::default() - }, - ); - - let policy = SubAgentPolicy::default().with_budget(SubAgentBudget { - max_model_calls: Some(1), - ..SubAgentBudget::unlimited() - }); - let node = SubAgentNode::<String, String>::from_fns( - "twostep", - |state: &String| SubAgentInput { - prompt: state.clone(), - data: Some(json!({ "kind": "budget-check" })), - }, - |out: SubAgentOutput| out.text, - ) - .with_policy(policy); - let graph = graph_delegating_to(node); - - let err = graph - .run_with_agent_binding("go".to_string(), binding(invoker)) - .await - .unwrap_err(); - assert!( - matches!(err, TinyAgentsError::LimitExceeded(_)), - "got {err:?}" - ); - - assert!( - SubAgentBudget::unlimited() - .check( - &SubAgentOutput { - text: "ok".to_string(), - structured: None, - usage: Default::default(), - model_calls: 0, - tool_calls: 0, - }, - "agent" - ) - .is_ok() - ); -} - -#[tokio::test] -async fn graph_uses_the_execution_scoped_agent_invoker() { - let invoker = TestAgentInvoker::with_output( - "adapter", - SubAgentOutput { - text: "adapter answer".to_owned(), - model_calls: 1, - ..SubAgentOutput::default() - }, - ); - let node = SubAgentNode::<String, String>::from_fns( - "adapter", - |state: &String| SubAgentInput::prompt(state.clone()), - |out: SubAgentOutput| out.text, - ); - let run = graph_delegating_to(node) - .run_with_agent_binding("delegated prompt".to_owned(), binding(invoker)) - .await - .expect("adapter run"); - assert_eq!(run.state, "adapter answer"); -} diff --git a/crates/tinyagents-orchestration/tests/e2e_orchestrator_subagents.rs b/crates/tinyagents-orchestration/tests/e2e_orchestrator_subagents.rs deleted file mode 100644 index 9a250e92c..000000000 --- a/crates/tinyagents-orchestration/tests/e2e_orchestrator_subagents.rs +++ /dev/null @@ -1,302 +0,0 @@ -//! TRUE end-to-end, fully OFFLINE: the flagship registry showcase. -//! -//! An *orchestrator* agent designs which specialized sub-agents to run by -//! emitting a structured selection of **names**, those names are resolved out of -//! a [`CapabilityRegistry`] at runtime, the resolved sub-agents are run in -//! parallel, and their outputs are composed into a final answer. -//! -//! Everything is deterministic: the orchestrator and every sub-agent are backed -//! by [`ScriptedModel`]s, so the test asserts on *structure* — which capabilities -//! the registry resolved and ran (testkit [`Trajectory`]) and that the unchosen -//! capability was never resolved or invoked — never on model prose. -//! -//! Topology: -//! -//! ```text -//! registry: { researcher, coder, summarizer } (3 named SubAgentTools) -//! │ -//! orchestrator (structured output) ── selects ──> ["researcher", "summarizer"] -//! │ -//! resolve each name from the registry, run in parallel, compose -//! ``` - -use std::collections::HashMap; -use std::sync::Arc; - -use futures::future::join_all; -use serde_json::{Value, json}; - -use tinyagents_graph::*; -use tinyagents_harness::context::{RunConfig, RunContext}; -use tinyagents_harness::events::{AgentEvent, EventSink}; -use tinyagents_harness::ids::{CallId, RunId}; -use tinyagents_harness::middleware::AgentRun; -use tinyagents_harness::runtime::{AgentHarness, RunPolicy}; -use tinyagents_harness::testkit::{EventRecorder, ScriptedModel, Trajectory}; -use tinyagents_harness::tool::ToolDispatch; -use tinyagents_orchestration::subagent::{ - ChildDataPolicy, SubAgent, SubAgentJobRegistry, SubAgentTool, -}; -use tinyagents_registry::*; -use tinyinference_llm::message::Message; -use tinyinference_llm::model::ResponseFormat; - -// ── Helpers ──────────────────────────────────────────────────────────────── - -/// Wraps a [`ScriptedModel`] as a named [`SubAgentTool`] whose child run always -/// answers with `answer`. -fn specialist( - name: &str, - description: &str, - model: Arc<ScriptedModel>, - jobs: SubAgentJobRegistry, -) -> SubAgentTool<()> { - let mut harness: AgentHarness<()> = AgentHarness::new(); - harness - .register_model("model", model) - .set_default_model("model"); - let subagent = SubAgent::new(name, description, Arc::new(harness)) - .with_system_prompt(format!("You are the {name}.")); - SubAgentTool::new( - Arc::new(subagent), - ChildDataPolicy::new(|parent: &()| *parent), - ) - .with_job_registry(jobs) -} - -/// Reads the `{ "agents": [..] }` selection out of an [`AgentRun`], preferring -/// extracted structured output and falling back to parsing the raw text. -fn parse_selection(run: &AgentRun) -> Vec<String> { - let value: Value = run - .structured - .clone() - .or_else(|| run.text().and_then(|t| serde_json::from_str(&t).ok())) - .unwrap_or(Value::Null); - value - .get("agents") - .and_then(Value::as_array) - .map(|arr| { - arr.iter() - .filter_map(|v| v.as_str().map(str::to_owned)) - .collect() - }) - .unwrap_or_default() -} - -#[tokio::test] -async fn orchestrator_resolves_and_runs_only_the_chosen_subagents() -> Result<()> { - let task = "research and summarize Rust trait objects"; - - // Keep concrete handles so we can assert *which* sub-agent models were - // actually invoked (and which were not). - let researcher_model = Arc::new(ScriptedModel::replies(vec!["RESEARCH_NOTES_ON_TRAITS"])); - let coder_model = Arc::new(ScriptedModel::replies(vec!["CODE_SNIPPET_DYN_TRAIT"])); - let summarizer_model = Arc::new(ScriptedModel::replies(vec!["PLAIN_SUMMARY_DONE"])); - - // 1. Register three named specialist sub-agents in the capability registry. - let mut registry: CapabilityRegistry<()> = CapabilityRegistry::new(); - let mut dispatches: HashMap<String, Arc<SubAgentTool<()>>> = HashMap::new(); - let jobs = SubAgentJobRegistry::new(); - for (name, description, model) in [ - ( - "researcher", - "Gathers factual background.", - researcher_model.clone(), - ), - ("coder", "Writes code snippets.", coder_model.clone()), - ( - "summarizer", - "Condenses material.", - summarizer_model.clone(), - ), - ] { - let dispatch = Arc::new(specialist(name, description, model, jobs.clone())); - registry.register_tool(dispatch.tool())?; - dispatches.insert(name.to_owned(), dispatch); - } - - let available = registry.names(ComponentKind::Tool); - assert_eq!(available, vec!["coder", "researcher", "summarizer"]); - - // 2. Orchestrator designs the plan: a scripted structured selection of TWO - // of the three registered names (deliberately skipping `coder`). - let mut orchestrator: AgentHarness<()> = AgentHarness::new(); - orchestrator - .register_model( - "model", - Arc::new(ScriptedModel::replies(vec![ - r#"{"agents":["researcher","summarizer"]}"#, - ])), - ) - .set_default_model("model") - .with_policy(RunPolicy { - default_response_format: Some(ResponseFormat::json_schema( - "agent_selection", - json!({ - "type": "object", - "properties": { - "agents": { "type": "array", "items": { "type": "string" } } - }, - "required": ["agents"] - }), - )), - ..RunPolicy::default() - }); - - let plan = orchestrator - .invoke_default(&(), vec![Message::user(task)]) - .await?; - - // The orchestrator's selection is extracted as structured output. - assert_eq!( - plan.structured, - Some(json!({ "agents": ["researcher", "summarizer"] })) - ); - - let chosen = parse_selection(&plan); - assert_eq!(chosen, vec!["researcher", "summarizer"]); - assert!( - !chosen.iter().any(|n| n == "coder"), - "the orchestrator did not choose `coder`" - ); - - // 3. Bind at runtime: resolve each chosen name from the registry and run the - // resolved sub-agents in parallel, bracketing each dispatch with tool - // events on a shared sink so the trajectory is observable. - let recorder = EventRecorder::new(); - let sink: EventSink = recorder.sink(); - sink.emit(AgentEvent::RunStarted { - run_id: RunId::new("orchestrator"), - thread_id: None, - }); - - let dispatches = chosen.iter().enumerate().map(|(i, name)| { - let name = name.clone(); - let dispatch = dispatches - .get(&name) - .cloned() - .expect("a chosen name must resolve in the registry"); - let sink = sink.clone(); - let call_id = format!("dispatch-{i}"); - async move { - sink.emit(AgentEvent::ToolStarted { - parent_call_id: None, - call_id: CallId::new(call_id.clone()), - tool_name: name.clone(), - input: None, - }); - let parent = RunContext::new(RunConfig::new(format!("dispatch-{i}")), ()); - let result = dispatch - .invoke_in_parent_context( - &(), - json!({ "input": task }), - tinytools::ToolCallOptions::default(), - &parent, - ) - .await?; - sink.emit(AgentEvent::ToolCompleted { - parent_call_id: None, - call_id: CallId::new(call_id), - tool_name: name.clone(), - started_at_ms: None, - input: None, - output: None, - duration_ms: None, - output_bytes: None, - error: None, - metadata: None, - }); - let job_id = serde_json::from_str::<Value>(&result.output())? - .get("job_id") - .and_then(Value::as_str) - .expect("spawn result contains a job id") - .to_owned(); - Ok::<(String, String), TinyAgentsError>((name, job_id)) - } - }); - let spawned: Vec<(String, String)> = join_all(dispatches) - .await - .into_iter() - .collect::<Result<_>>()?; - - let outputs = tokio::time::timeout(std::time::Duration::from_secs(1), async { - loop { - let completed = spawned - .iter() - .filter_map(|(name, job_id)| { - let job = jobs.get(job_id)?; - job.status - .is_terminal() - .then(|| (name.clone(), job.output.unwrap_or_default())) - }) - .collect::<Vec<_>>(); - if completed.len() == spawned.len() { - return completed; - } - tokio::task::yield_now().await; - } - }) - .await - .expect("chosen subagent jobs reach terminal states"); - - sink.emit(AgentEvent::RunCompleted { - run_id: RunId::new("orchestrator"), - outcome: None, - }); - - // 4. Compose the resolved sub-agents' outputs into one final answer. - let composed = outputs - .iter() - .map(|(name, text)| format!("[{name}] {text}")) - .collect::<Vec<_>>() - .join("\n"); - - // ── Assertions ─────────────────────────────────────────────────────────── - - // Registry resolution: all three are registered, but only the chosen two - // were ever resolved into a runnable handle. - assert!(registry.has(ComponentKind::Tool, "coder")); - assert!(registry.tool("researcher").is_some()); - assert!(registry.tool("summarizer").is_some()); - - // Trajectory: ONLY the chosen sub-agents ran; the unchosen one never did. - let traj = Trajectory::from_events(recorder.events()); - traj.assert_tool_called("researcher"); - traj.assert_tool_called("summarizer"); - assert_eq!(traj.tool_call_count("researcher"), 1); - assert_eq!(traj.tool_call_count("summarizer"), 1); - assert!( - !traj.tool_was_called("coder"), - "the unchosen `coder` capability must never run" - ); - traj.assert_completed(); - traj.assert_order(&["run.started", "run.completed"]) - .expect("dispatch is bracketed by run start/completion"); - - // Model-level proof: only the chosen sub-agents' models were invoked. - assert_eq!( - researcher_model.requests().len(), - 1, - "researcher sub-agent ran exactly once" - ); - assert_eq!( - summarizer_model.requests().len(), - 1, - "summarizer sub-agent ran exactly once" - ); - assert_eq!( - coder_model.requests().len(), - 0, - "the unchosen coder sub-agent was never invoked" - ); - - // The composition weaves in BOTH chosen sub-agents' outputs. - assert!(composed.contains("RESEARCH_NOTES_ON_TRAITS")); - assert!(composed.contains("PLAIN_SUMMARY_DONE")); - assert!( - !composed.contains("CODE_SNIPPET_DYN_TRAIT"), - "the unchosen sub-agent's output must not appear in the composition" - ); - - Ok(()) -} diff --git a/crates/tinyagents-orchestration/tests/e2e_subagent_error.rs b/crates/tinyagents-orchestration/tests/e2e_subagent_error.rs deleted file mode 100644 index ad17c53ef..000000000 --- a/crates/tinyagents-orchestration/tests/e2e_subagent_error.rs +++ /dev/null @@ -1,149 +0,0 @@ -//! TRUE end-to-end (offline): sub-agent ERROR propagation. -//! -//! A [`SubAgent`] is built over a child [`AgentHarness`] whose model always -//! requests a tool that fails ([`FakeTool::failing`]). The tool failure -//! propagates out of the child agent loop, so: -//! -//! - [`SubAgent::invoke`] returns `Err(TinyAgentsError::Tool(..))`, -//! - [`SubAgentTool`] returns a job id immediately, and the job registry later -//! records the sanitized failure, and -//! - an orchestrator remains live after spawning a failing child and can query -//! that failure through host tools. -//! -//! All assertions are structural / on the error variant — never on model prose. - -use std::sync::Arc; - -use serde_json::json; - -use tinyagents_harness::context::{RunConfig, RunContext}; -use tinyagents_harness::error::TinyAgentsError; -use tinyagents_harness::runtime::AgentHarness; -use tinyagents_harness::testkit::{EventRecorder, FakeTool, ScriptedModel, Trajectory}; -use tinyagents_orchestration::subagent::{ - ChildDataPolicy, SubAgent, SubAgentJobRegistry, SubAgentJobStatus, SubAgentTool, -}; -use tinyinference_llm::message::Message; -use tinyinference_llm::model::ModelResponse; -use tinyinference_llm::providers::MockModel; -use tinyinference_llm::tool::ToolCall; - -async fn wait_for_failed_job(jobs: &SubAgentJobRegistry) { - tokio::time::timeout(std::time::Duration::from_secs(1), async { - loop { - if let Some(job) = jobs.list().into_iter().next() - && job.status.is_terminal() - { - assert_eq!(job.status, SubAgentJobStatus::Failed); - assert_eq!( - job.error.as_deref(), - Some("tool error: tool dispatch failed") - ); - return; - } - tokio::task::yield_now().await; - } - }) - .await - .expect("failing child job reaches a terminal state"); -} - -/// Builds a child harness whose model always asks for the `broken` tool, which -/// fails with a foreign `anyhow` error. The harness maps that to its stable, -/// non-sensitive tool-failure surface before a sub-agent can propagate it. -fn failing_child_harness() -> AgentHarness<()> { - let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( - "child-model", - Arc::new(MockModel::with_tool_call("broken", json!({}))), - ); - harness.register_tool(Arc::new(FakeTool::failing("broken", "boom"))); - harness -} - -#[tokio::test] -async fn subagent_invoke_propagates_tool_failure() { - let subagent = SubAgent::new( - "broken_worker", - "a worker whose tool always fails", - Arc::new(failing_child_harness()), - ); - - let err = subagent - .invoke(&(), (), 0, "do the thing") - .await - .expect_err("the child tool failure must propagate out of SubAgent::invoke"); - - match err { - TinyAgentsError::Tool(msg) => assert_eq!(msg, "tool dispatch failed"), - other => panic!("expected a sanitized TinyAgentsError::Tool, got {other:?}"), - } -} - -#[tokio::test] -async fn subagent_tool_call_returns_job_id_and_records_failure() { - let subagent = Arc::new(SubAgent::new( - "broken_worker", - "a worker whose tool always fails", - Arc::new(failing_child_harness()), - )); - let tool = SubAgentTool::new(subagent, ChildDataPolicy::new(|parent: &()| *parent)); - let jobs = tool.job_registry().clone(); - let parent = RunContext::new(RunConfig::new("parent"), ()); - - let result = tool - .invoke_in_parent_context( - &(), - json!({ "input": "x" }), - tinytools::ToolCallOptions::default(), - &parent, - ) - .await - .expect("spawning a child is independent from its eventual result"); - assert!(result.output().contains("job_id")); - wait_for_failed_job(&jobs).await; -} - -#[tokio::test] -async fn orchestrator_survives_a_failing_subagent_job() { - let subagent = Arc::new(SubAgent::new( - "broken_worker", - "a worker whose tool always fails", - Arc::new(failing_child_harness()), - )); - let tool = Arc::new(SubAgentTool::new( - subagent, - ChildDataPolicy::new(|parent: &()| *parent), - )); - let jobs = tool.job_registry().clone(); - - let mut orchestrator: AgentHarness<()> = AgentHarness::new(); - orchestrator.register_tool_dispatch(tool); - let mut delegation = ModelResponse::assistant(""); - delegation.message.tool_calls.push(ToolCall::new( - "delegate", - "broken_worker", - json!({ "input": "delegate" }), - )); - orchestrator.register_model( - "parent-model", - Arc::new(ScriptedModel::new(vec![ - delegation, - ModelResponse::assistant("job spawned"), - ])), - ); - - let recorder = EventRecorder::new(); - let ctx = RunContext::new(RunConfig::new("orchestrator-run"), ()).with_events(recorder.sink()); - - let run = orchestrator - .invoke_in_context(&(), ctx, vec![Message::user("delegate this")]) - .await - .expect("child failure does not abort the spawning orchestrator"); - assert_eq!(run.text().as_deref(), Some("job spawned")); - wait_for_failed_job(&jobs).await; - - // The orchestrator run emitted a RunFailed event (on_error fan-out path). - let traj = Trajectory::from_events(recorder.events()); - traj.assert_completed(); -} diff --git a/crates/tinyagents-orchestration/tests/e2e_subagent_reuse.rs b/crates/tinyagents-orchestration/tests/e2e_subagent_reuse.rs deleted file mode 100644 index 7c83b0f52..000000000 --- a/crates/tinyagents-orchestration/tests/e2e_subagent_reuse.rs +++ /dev/null @@ -1,110 +0,0 @@ -//! TRUE end-to-end (offline): post-completion sub-agent REUSE with human input. -//! -//! An orchestrator calls a sub-agent, then takes (simulated) human input, then -//! calls the **same** sub-agent again — reusing it and carrying the prior -//! conversation context forward instead of killing/restarting it. This is the -//! reuse flow exposed by [`SubAgentSession`], distinct from steering (which -//! interrupts a still-running agent). -//! -//! The sub-agent is driven by a deterministic [`ScriptedModel`] so the test is -//! fully offline. We assert *structurally* — the composed result, the carried -//! context (via the model's recorded second request), and the reuse trajectory -//! (via a testkit [`Trajectory`]) — never on LLM prose. - -use std::sync::Arc; - -use tinyagents_harness::runtime::AgentHarness; -use tinyagents_harness::testkit::{EventRecorder, ScriptedModel, Trajectory}; -use tinyagents_orchestration::subagent::{SubAgent, SubAgentSession}; -use tinyinference_llm::message::Message; - -#[tokio::test] -async fn orchestrator_reuses_subagent_with_human_input_between() { - // A persistent "travel concierge" sub-agent. The scripted model answers the - // first question, then a context-dependent follow-up. - let model = Arc::new(ScriptedModel::replies(vec![ - "Tokyo is a great choice for a spring trip.", - "For Tokyo in spring, pack light layers and a rain jacket.", - ])); - - let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("scripted", model.clone()); - - let subagent = Arc::new( - SubAgent::new( - "concierge", - "a persistent travel concierge", - Arc::new(harness), - ) - .with_system_prompt("You are a helpful travel concierge."), - ); - let subagent_handle = Arc::clone(&subagent); - - // The orchestrator wires the session's events into a recorder so the reuse - // is observable as a trajectory. - let recorder = EventRecorder::new(); - let mut session = SubAgentSession::new(subagent).with_events(recorder.sink()); - - // ── Turn 1: orchestrator asks on behalf of the user ────────────────────── - let first = session - .send( - &(), - (), - vec![Message::user("Where should I travel this spring?")], - ) - .await - .expect("first sub-agent run"); - assert_eq!( - first.text(), - Some("Tokyo is a great choice for a spring trip.".to_string()) - ); - - // ── Simulated human-in-the-loop input arrives out-of-band ──────────────── - let human_followup = "Great — what should I pack for there?"; - - // ── Turn 2: the SAME sub-agent is reused, carrying prior context ────────── - let second = session - .send(&(), (), vec![Message::user(human_followup)]) - .await - .expect("second sub-agent run (reused)"); - assert_eq!( - second.text(), - Some("For Tokyo in spring, pack light layers and a rain jacket.".to_string()) - ); - - // The same SubAgent instance was reused — never reconstructed. - assert!( - Arc::ptr_eq(session.subagent(), &subagent_handle), - "the orchestrator reused the same sub-agent across turns" - ); - assert_eq!(session.turns(), 2); - - // Carried context: the second model request contained the first turn's - // question, the first answer, AND the human follow-up. - let requests = model.requests(); - assert_eq!(requests.len(), 2, "one model request per reused send"); - let second_texts: Vec<String> = requests[1].messages.iter().map(Message::text).collect(); - assert!(second_texts.contains(&"Where should I travel this spring?".to_string())); - assert!(second_texts.contains(&"Tokyo is a great choice for a spring trip.".to_string())); - assert!(second_texts.contains(&human_followup.to_string())); - - // The accumulating session transcript reflects the full conversation. - let transcript_texts: Vec<String> = session.transcript().iter().map(Message::text).collect(); - assert!(transcript_texts.contains(&human_followup.to_string())); - assert!( - transcript_texts - .contains(&"For Tokyo in spring, pack light layers and a rain jacket.".to_string()) - ); - - // Trajectory: both child runs completed, and the second send is recorded as - // a reuse (`subagent.reused`) bracketed by start/complete events. - let traj = Trajectory::from_events(recorder.events()); - traj.assert_completed(); - traj.assert_order(&[ - "subagent.started", - "subagent.completed", - "subagent.reused", - "subagent.completed", - ]) - .expect("first run completes, then the sub-agent is reused for the second"); -} diff --git a/crates/tinyagents-orchestration/tests/e2e_subagent_timeout.rs b/crates/tinyagents-orchestration/tests/e2e_subagent_timeout.rs deleted file mode 100644 index bf1b1f424..000000000 --- a/crates/tinyagents-orchestration/tests/e2e_subagent_timeout.rs +++ /dev/null @@ -1,72 +0,0 @@ -//! TRUE end-to-end (offline, deterministic): sub-agent TIMEOUT enforcement. -//! -//! A [`SubAgent`] is built over a child harness driven by a [`SlowModel`] that -//! sleeps 200ms before replying. The child harness policy sets a 20ms -//! wall-clock cap, so the per-model-call budget interrupts the slow call and -//! [`SubAgent::invoke`] returns `Err(TinyAgentsError::Timeout(..))` — even -//! though the child [`RunConfig`] built internally by `invoke` carries no -//! per-run timeout (the loop also honors the harness policy's wall-clock cap). -//! -//! A control case shows a *fast* model under the same small cap completes -//! successfully — the timeout fires only on genuinely slow calls. - -use std::sync::Arc; -use std::time::Duration; - -use tinyagents_harness::error::TinyAgentsError; -use tinyagents_harness::limits::RunLimits; -use tinyagents_harness::runtime::{AgentHarness, RunPolicy}; -use tinyagents_harness::testkit::SlowModel; -use tinyagents_orchestration::subagent::SubAgent; -use tinyinference_llm::providers::MockModel; - -/// Builds a child harness whose policy caps wall-clock time at `timeout_ms` and -/// whose registered model is `model`. -fn capped_child_harness<M>(model: M, timeout_ms: u64) -> AgentHarness<()> -where - M: tinyinference_llm::model::ChatModel<()> + 'static, -{ - let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("model", Arc::new(model)); - harness.with_policy(RunPolicy { - limits: RunLimits::default().with_max_wall_clock_ms(Some(timeout_ms)), - ..RunPolicy::default() - }); - harness -} - -#[tokio::test] -async fn subagent_invoke_times_out_on_slow_model() { - // 20ms cap against a 200ms model call: the per-call budget must fire. - let child = capped_child_harness(SlowModel::new(Duration::from_millis(200), "eventually"), 20); - let subagent = SubAgent::new( - "slow_worker", - "a worker backed by a slow model", - Arc::new(child), - ); - - let err = subagent - .invoke(&(), (), 0, "go") - .await - .expect_err("a 200ms model call under a 20ms cap must time out"); - - assert!(matches!(err, TinyAgentsError::Timeout(_)), "got {err:?}"); -} - -#[tokio::test] -async fn fast_subagent_invoke_succeeds_under_same_cap() { - // Control: a fast model under the same 20ms cap completes cleanly. - let child = capped_child_harness(MockModel::constant("quick"), 20); - let subagent = SubAgent::new( - "fast_worker", - "a worker backed by a fast model", - Arc::new(child), - ); - - let run = subagent - .invoke(&(), (), 0, "go") - .await - .expect("a fast model call completes within the cap"); - - assert_eq!(run.text(), Some("quick".to_string())); -} diff --git a/crates/tinyagents-orchestration/tests/e2e_subagents.rs b/crates/tinyagents-orchestration/tests/e2e_subagents.rs deleted file mode 100644 index 07013fc8b..000000000 --- a/crates/tinyagents-orchestration/tests/e2e_subagents.rs +++ /dev/null @@ -1,276 +0,0 @@ -//! TRUE end-to-end: agent-calling-agent composition (sub-agents). -//! -//! A parent [`AgentHarness`] whose scripted [`MockModel`] calls a -//! [`SubAgentTool`] drives a child [`AgentHarness`] (also a `MockModel`) and -//! composes the child's answer into a final assistant reply. The parent run's -//! [`EventSink`] is wired to a testkit [`EventRecorder`] so we can reconstruct -//! a [`Trajectory`] and assert *structurally* that the sub-agent really ran — -//! never on model prose. -//! -//! A second test exercises the deterministic recursion-depth guard: nesting a -//! sub-agent past the harness's `max_depth` fails fast with -//! [`TinyAgentsError::SubAgentDepth`] *before* any model call, both through the -//! direct invoke path and through the tool path. - -use std::sync::Arc; - -use serde_json::json; - -use tinyagents_harness::context::{RunConfig, RunContext}; -use tinyagents_harness::error::TinyAgentsError; -use tinyagents_harness::events::AgentEvent; -use tinyagents_harness::limits::RunLimits; -use tinyagents_harness::runtime::{AgentHarness, RunPolicy}; -use tinyagents_harness::testkit::{EventRecorder, Trajectory}; -use tinyagents_orchestration::subagent::{ - ChildDataPolicy, SubAgent, SubAgentJobStatus, SubAgentTool, -}; -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; - -// ── Helpers ────────────────────────────────────────────────────────────────── - -/// A tool-call assistant turn: no text, a single tool call. -fn tool_call_response(id: &str, name: &str, arguments: serde_json::Value) -> ModelResponse { - ModelResponse { - message: AssistantMessage { - id: Some(format!("msg-{id}")), - content: Vec::new(), - tool_calls: vec![ToolCall::new(id, name, arguments)], - usage: Some(Usage::new(9, 4)), - origin: None, - }, - usage: Some(Usage::new(9, 4)), - finish_reason: Some("tool_calls".to_string()), - raw: None, - resolved_model: None, - continue_turn: None, - served_from_cache: false, - correlation: None, - resolved_route: None, - } -} - -/// A plain-text assistant turn. -fn text_response(text: &str) -> ModelResponse { - ModelResponse { - message: AssistantMessage { - id: None, - content: vec![ContentBlock::Text(text.to_string())], - tool_calls: Vec::new(), - usage: Some(Usage::new(5, 3)), - origin: None, - }, - usage: Some(Usage::new(5, 3)), - finish_reason: Some("stop".to_string()), - raw: None, - resolved_model: None, - continue_turn: None, - served_from_cache: false, - correlation: None, - resolved_route: None, - } -} - -/// A child harness whose model always answers with `answer`. -fn child_harness(answer: &str) -> AgentHarness<()> { - let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("child-model", Arc::new(MockModel::constant(answer))); - harness -} - -/// A child harness capped at `max_depth`. -fn child_harness_with_max_depth(answer: &str, max_depth: usize) -> AgentHarness<()> { - let mut harness = child_harness(answer); - harness.with_policy(RunPolicy { - limits: RunLimits::default().with_max_depth(max_depth), - ..RunPolicy::default() - }); - harness -} - -// ── Tests ──────────────────────────────────────────────────────────────────── - -#[tokio::test] -async fn parent_drives_subagent_and_composes_answer() { - // Child agent that "researches" and returns a fixed answer. - let child = Arc::new(SubAgent::new( - "researcher", - "answers research questions", - Arc::new(child_harness("RUST_IS_A_SYSTEMS_LANGUAGE")), - )); - let tool = Arc::new(SubAgentTool::new( - child, - ChildDataPolicy::new(|parent: &()| *parent), - )); - let jobs = tool.job_registry().clone(); - - // Parent: first turn delegates to the sub-agent tool, second turn composes - // the final answer. - let mut parent: AgentHarness<()> = AgentHarness::new(); - parent.register_tool_dispatch(tool); - parent.register_model( - "parent-model", - Arc::new(MockModel::with_responses(vec![ - tool_call_response("c1", "researcher", json!({ "input": "what is rust?" })), - text_response("Based on the researcher: a systems language."), - ])), - ); - - // Wire the parent run's events into a shared recorder so we can assert on - // the trajectory (the sub-agent tool really ran) afterward. - let recorder = EventRecorder::new(); - let ctx = RunContext::new(RunConfig::new("parent-run"), ()).with_events(recorder.sink()); - - let run = parent - .invoke_in_context(&(), ctx, vec![Message::user("delegate this")]) - .await - .expect("parent run succeeds"); - - // Behavior / structure assertions — never on the exact prose. - assert_eq!(run.tool_calls, 1, "parent invoked the sub-agent tool once"); - assert_eq!(run.model_calls, 2, "parent made two model calls"); - assert_eq!( - run.text(), - Some("Based on the researcher: a systems language.".to_string()) - ); - - // The tool result contains a stable job id rather than blocking for output. - let job_id = jobs - .list() - .into_iter() - .next() - .expect("the subagent tool registers one job") - .id - .to_string(); - let job = tokio::time::timeout(std::time::Duration::from_secs(1), async { - loop { - let job = jobs.get(&job_id).expect("job is registered"); - if job.status.is_terminal() { - break job; - } - tokio::task::yield_now().await; - } - }) - .await - .expect("job completes"); - assert_eq!(job.status, SubAgentJobStatus::Completed); - assert_eq!(job.output.as_deref(), Some("RUST_IS_A_SYSTEMS_LANGUAGE")); - - // Trajectory assertion: the sub-agent (exposed as the `researcher` tool) - // really ran, and the run completed cleanly. - let traj = Trajectory::from_events(recorder.events()); - traj.assert_tool_called("researcher"); - assert_eq!(traj.tool_call_count("researcher"), 1); - traj.assert_model_called_times(3); - traj.assert_completed(); - traj.assert_order(&["run.started", "researcher", "run.completed"]) - .expect("sub-agent tool runs between run start and completion"); -} - -#[tokio::test] -async fn child_subagents_derive_unique_thread_ids_from_parent_thread() { - let child = SubAgent::new( - "researcher", - "answers research questions", - Arc::new(child_harness("RUST_IS_A_SYSTEMS_LANGUAGE")), - ); - let recorder = EventRecorder::new(); - let parent = RunContext::new( - RunConfig::new("parent-run").with_thread("parent-thread"), - (), - ) - .with_events(recorder.sink()); - - child - .invoke_in_parent(&(), (), &parent, "first question") - .await - .expect("first child run succeeds"); - child - .invoke_in_parent(&(), (), &parent, "second question") - .await - .expect("second child run succeeds"); - - let child_threads: Vec<String> = recorder - .events() - .into_iter() - .filter_map(|event| match event { - AgentEvent::RunStarted { - thread_id: Some(thread_id), - .. - } => Some(thread_id.to_string()), - _ => None, - }) - .collect(); - - assert_eq!(child_threads.len(), 2); - assert_ne!( - child_threads[0], child_threads[1], - "each child run should get an isolated thread" - ); - for thread in child_threads { - assert!( - thread.starts_with("parent-thread-subagent-researcher-d1-"), - "child thread should inherit the parent thread as a hyphenated prefix: {thread}" - ); - assert!( - !thread.contains('/'), - "child thread id should not use slash separators: {thread}" - ); - } -} - -#[tokio::test] -async fn nesting_past_max_depth_is_a_deterministic_error() { - // Cap the child harness at depth 1: a child run is allowed at depth 1 - // (parent_depth 0) but not at depth 2 (parent_depth 1). - let subagent = Arc::new(SubAgent::new( - "deep", - "a deep agent", - Arc::new(child_harness_with_max_depth("ok", 1)), - )); - - // Within the cap: parent_depth 0 -> child depth 1. - let ok_run = subagent - .invoke(&(), (), 0, "ok") - .await - .expect("child run at depth 1 is within the cap"); - assert_eq!(ok_run.text(), Some("ok".to_string())); - - // Direct invoke past the cap: parent_depth 1 -> child depth 2 > cap of 1. - let err = subagent - .invoke(&(), (), 1, "too deep") - .await - .expect_err("child depth 2 exceeds the cap"); - assert!( - matches!(err, TinyAgentsError::SubAgentDepth(1)), - "expected SubAgentDepth(1), got {err:?}" - ); - - // Tool path past the cap: the typed parent context carries depth 1, so the - // child run would exceed its maximum depth of 1. - let tool = SubAgentTool::new(subagent, ChildDataPolicy::new(|parent: &()| *parent)); - let parent = RunContext::new( - RunConfig::new("deep-parent") - .with_depth(1) - .with_max_depth(1), - (), - ); - let tool_result = tool - .invoke_in_parent_context( - &(), - json!({ "input": "x" }), - tinytools::ToolCallOptions::default(), - &parent, - ) - .await - .expect("the tool returns a failed tool result"); - let tool_error = tool_result.output(); - assert!( - tool_error.contains("recursion depth limit") && tool_error.contains("maximum depth of 1"), - "expected SubAgentDepth(1) from the tool path, got {tool_error:?}" - ); -} diff --git a/crates/tinyagents-orchestration/tests/hosted_subagents.rs b/crates/tinyagents-orchestration/tests/hosted_subagents.rs deleted file mode 100644 index 5499d88fd..000000000 --- a/crates/tinyagents-orchestration/tests/hosted_subagents.rs +++ /dev/null @@ -1,255 +0,0 @@ -//! Hosted subagent authorization and capability-bundle propagation. - -use std::sync::Arc; - -use async_trait::async_trait; -use serde_json::json; -use tinyagents_definition::{AgentDefinition, InMemoryDefinitionRegistry}; -use tinyagents_harness::context::{RunConfig, RunContext}; -use tinyagents_harness::host::{ - AllowAllSecurityGate, HostCapabilities, ModelResolveRequest, ModelResolver, - StaticContextComposer, -}; -use tinyagents_harness::runtime::{ - AgentHarness, AgentInvocation, AgentTurnRequest, InvocationRuntime, -}; -use tinyagents_harness::testkit::{FakeTool, ScriptedModel}; -use tinyagents_orchestration::subagent::{ - ChildDataPolicy, SubAgent, SubAgentJobRegistry, SubAgentJobStatus, SubAgentTool, -}; -use tinyinference_llm::model::{ChatModel, ModelResponse}; - -fn delegation_call() -> ModelResponse { - let mut response = ModelResponse::assistant(""); - response - .message - .tool_calls - .push(tinyinference_llm::tool::ToolCall::new( - "delegate", - "worker", - json!({"input": "child task"}), - )); - response -} - -struct AgentResolver { - parent: Arc<dyn ChatModel<()>>, - worker: Arc<dyn ChatModel<()>>, -} - -#[async_trait] -impl ModelResolver<()> for AgentResolver { - async fn resolve( - &self, - request: &ModelResolveRequest, - ) -> tinyagents_harness::error::Result<Arc<dyn ChatModel<()>>> { - Ok(if request.agent_id == "worker" { - self.worker.clone() - } else { - self.parent.clone() - }) - } -} - -fn host( - parent: AgentDefinition, - parent_model: Arc<dyn ChatModel<()>>, - worker_model: Arc<dyn ChatModel<()>>, -) -> HostCapabilities<()> { - HostCapabilities::new( - Arc::new(StaticContextComposer::empty()), - Arc::new(InMemoryDefinitionRegistry::new(vec![ - parent, - AgentDefinition::new("worker", "Worker", "child"), - ])), - Arc::new(AllowAllSecurityGate), - Arc::new(AgentResolver { - parent: parent_model, - worker: worker_model, - }), - ) -} - -fn runtime_with_worker( - child_harness: AgentHarness<()>, -) -> (InvocationRuntime<(), ()>, SubAgentJobRegistry) { - let child = Arc::new(SubAgent::new("worker", "child", Arc::new(child_harness))); - let jobs = SubAgentJobRegistry::new(); - let mut overlay = AgentHarness::new(); - overlay.register_tool_dispatch(Arc::new( - SubAgentTool::new(child, ChildDataPolicy::new(|_: &()| ())).with_job_registry(jobs.clone()), - )); - (InvocationRuntime::new(overlay), jobs) -} - -#[tokio::test] -async fn authorized_child_reuses_the_parent_host_bundle() { - let parent_model = Arc::new(ScriptedModel::new(vec![ - delegation_call(), - ModelResponse::assistant("parent answer"), - ])); - let worker_model = Arc::new(ScriptedModel::replies(vec!["child answer"])); - let parent = AgentDefinition::new("parent", "Parent", "delegates") - .with_subagents(["worker"]) - .with_tools(["worker"]); - let entry = AgentHarness::new(); - let (runtime, jobs) = runtime_with_worker(AgentHarness::new()); - let run = entry - .invoke_agent( - AgentInvocation::new( - host(parent, parent_model.clone(), worker_model), - AgentTurnRequest::new( - "parent", - vec![tinyinference_llm::message::Message::user("delegate")], - ), - RunContext::new(RunConfig::new("authorized-child"), ()), - ) - .with_runtime(runtime), - &(), - ) - .await - .expect("authorized child runs through the parent's host bundle"); - - assert_eq!(run.text().as_deref(), Some("parent answer")); - let job = tokio::time::timeout(std::time::Duration::from_secs(1), async { - loop { - if let Some(job) = jobs.list().into_iter().next() - && job.status.is_terminal() - { - return job; - } - tokio::task::yield_now().await; - } - }) - .await - .expect("hosted child job reaches a terminal state"); - assert_eq!(job.status, SubAgentJobStatus::Completed); - assert_eq!(job.output.as_deref(), Some("child answer")); -} - -#[tokio::test] -async fn parent_denial_cannot_fall_back_to_the_child_harness() { - let parent_model = Arc::new(ScriptedModel::new(vec![delegation_call()])); - let local_child_model = Arc::new(ScriptedModel::replies(vec!["local bypass"])); - let host_worker_model = Arc::new(ScriptedModel::replies(vec!["host worker"])); - let mut child_harness = AgentHarness::new(); - child_harness.register_model("local", local_child_model.clone()); - - let (runtime, _jobs) = runtime_with_worker(child_harness); - let error = AgentHarness::new() - .invoke_agent( - AgentInvocation::new( - host( - AgentDefinition::new("parent", "Parent", "does not delegate") - .with_tools(["worker"]), - parent_model, - host_worker_model.clone(), - ), - AgentTurnRequest::new( - "parent", - vec![tinyinference_llm::message::Message::user("delegate")], - ), - RunContext::new(RunConfig::new("denied-child"), ()), - ) - .with_runtime(runtime), - &(), - ) - .await - .expect_err("the parent's delegate allowlist denies the child"); - - assert_eq!( - error.to_string(), - "hosted agent invocation failed at the model provider" - ); - assert!( - local_child_model.requests().is_empty(), - "the child cannot substitute its own model authority" - ); - assert!( - host_worker_model.requests().is_empty(), - "a denied child must not run on the host's worker model either" - ); -} - -#[tokio::test] -async fn hosted_child_cannot_call_a_tool_the_parent_was_denied() { - // The parent may only call `worker`; the worker's own definition is an - // explicit wildcard. A hosted child never widens what its parent holds. - let parent_model = Arc::new(ScriptedModel::new(vec![ - delegation_call(), - ModelResponse::assistant("parent answer"), - ])); - let mut shell_call = ModelResponse::assistant(""); - shell_call - .message - .tool_calls - .push(tinyinference_llm::tool::ToolCall::new( - "shell-1", - "shell", - json!({"command": "rm -rf /"}), - )); - let worker_model = Arc::new(ScriptedModel::new(vec![ - shell_call, - ModelResponse::assistant("child answer"), - ])); - let parent = AgentDefinition::new("parent", "Parent", "delegates") - .with_subagents(["worker"]) - .with_tools(["worker"]); - let mut worker = AgentDefinition::new("worker", "Worker", "child"); - worker.disallowed_tools = vec!["unused".into()]; - let host = HostCapabilities::new( - Arc::new(StaticContextComposer::empty()), - Arc::new(InMemoryDefinitionRegistry::new(vec![parent, worker])), - Arc::new(AllowAllSecurityGate), - Arc::new(AgentResolver { - parent: parent_model, - worker: worker_model, - }), - ); - let shell = Arc::new(FakeTool::returning("shell", "ran")); - let mut child_harness = AgentHarness::new(); - child_harness.register_tool(shell.clone()); - let child = Arc::new(SubAgent::new("worker", "child", Arc::new(child_harness))); - let jobs = SubAgentJobRegistry::new(); - let mut overlay = AgentHarness::new(); - overlay.register_tool(shell.clone()); - overlay.register_tool_dispatch(Arc::new( - SubAgentTool::new(child, ChildDataPolicy::new(|_: &()| ())).with_job_registry(jobs.clone()), - )); - let mut entry = AgentHarness::new(); - entry.register_tool(shell.clone()); - - entry - .invoke_agent( - AgentInvocation::new( - host, - AgentTurnRequest::new( - "parent", - vec![tinyinference_llm::message::Message::user("delegate")], - ), - RunContext::new(RunConfig::new("child-ceiling"), ()), - ) - .with_runtime(InvocationRuntime::new(overlay)), - &(), - ) - .await - .expect("parent run"); - tokio::time::timeout(std::time::Duration::from_secs(1), async { - loop { - if let Some(job) = jobs.list().into_iter().next() - && job.status.is_terminal() - { - return job; - } - tokio::task::yield_now().await; - } - }) - .await - .expect("hosted child job reaches a terminal state"); - - assert!( - shell.calls().is_empty(), - "the child ran `shell`, which its parent was never allowed: {:?}", - shell.calls() - ); -} diff --git a/crates/tinyagents-orchestration/tests/live_orchestrator_subagents.rs b/crates/tinyagents-orchestration/tests/live_orchestrator_subagents.rs deleted file mode 100644 index 0a75480a9..000000000 --- a/crates/tinyagents-orchestration/tests/live_orchestrator_subagents.rs +++ /dev/null @@ -1,256 +0,0 @@ -//! LIVE end-to-end: a real OpenAI orchestrator designs which sub-agents to run -//! by resolving them **by name** from a [`CapabilityRegistry`], runs them, and -//! composes a final answer. -//! -//! This is the network-backed sibling of `e2e_orchestrator_subagents.rs`. The -//! orchestrator, every specialist sub-agent, and the composer are driven by a -//! real [`OpenAiModel`]. We assert *structurally* — the orchestrator selected at -//! least one registered capability, those capabilities were resolved out of the -//! registry and run, and a non-empty composed answer came back — never on the -//! exact prose. -//! -//! # Skips gracefully -//! -//! This test is `#[ignore]`d and only runs opted in via -//! `tests/common/live.rs::require_live`, so `cargo test` passes with no key -//! configured and never dials a real provider by accident. - -mod common; - -#[tokio::test] -#[ignore = "network: set TINYAGENTS_LIVE=1 and run with --ignored"] -async fn live_openai_orchestrator_designs_subagents_via_registry() { - use std::collections::HashMap; - use std::sync::Arc; - - use futures::future::join_all; - use serde_json::{Value, json}; - - use tinyagents_harness::middleware::AgentRun; - use tinyagents_harness::runtime::{AgentHarness, RunPolicy}; - use tinyagents_harness::tool::ToolDispatch; - use tinyagents_orchestration::subagent::{ - ChildDataPolicy, SubAgent, SubAgentJobRegistry, SubAgentTool, - }; - use tinyagents_registry::*; - use tinyinference_llm::message::Message; - use tinyinference_llm::model::{ChatModel, ResponseFormat}; - use tinyinference_llm::providers::openai::OpenAiModel; - - if !common::live::require_live(&["OPENAI_API_KEY"]) { - return; - } - - fn parse_selection(run: &AgentRun) -> Vec<String> { - let value: Value = run - .structured - .clone() - .or_else(|| run.text().and_then(|t| serde_json::from_str(&t).ok())) - .unwrap_or(Value::Null); - value - .get("agents") - .and_then(Value::as_array) - .map(|arr| { - arr.iter() - .filter_map(|v| v.as_str().map(str::to_owned)) - .collect() - }) - .unwrap_or_default() - } - - let model: Arc<dyn ChatModel<()>> = - Arc::new(OpenAiModel::from_env().expect("OPENAI_API_KEY present")); - - // Three named specialists, each a SubAgent over the real model. - let specs = [ - ( - "researcher", - "Gathers and explains factual background on a topic. No code.", - "You are a meticulous researcher. Reply with a couple of factual bullet points.", - ), - ( - "coder", - "Writes small, focused code snippets.", - "You are a senior Rust engineer. Reply with a short, correct code snippet.", - ), - ( - "summarizer", - "Condenses material into a short plain-language summary.", - "You are an editor. Reply with a crisp 1-2 sentence summary.", - ), - ]; - - let mut registry: CapabilityRegistry<()> = CapabilityRegistry::new(); - let mut dispatches: HashMap<String, Arc<SubAgentTool<()>>> = HashMap::new(); - let jobs = SubAgentJobRegistry::new(); - for (name, description, system_prompt) in specs { - let mut harness: AgentHarness<()> = AgentHarness::new(); - harness - .register_model("model", model.clone()) - .set_default_model("model"); - let subagent = - SubAgent::new(name, description, Arc::new(harness)).with_system_prompt(system_prompt); - let dispatch = Arc::new( - SubAgentTool::new( - Arc::new(subagent), - ChildDataPolicy::new(|parent: &()| *parent), - ) - .with_job_registry(jobs.clone()), - ); - registry - .register_tool(dispatch.tool()) - .expect("unique specialist name"); - dispatches.insert(name.to_owned(), dispatch); - } - - // Discover the menu from the registry. - let available = registry.names(ComponentKind::Tool); - let menu_text = available - .iter() - .map(|name| { - let desc = registry - .tool(name) - .map(|t| t.description().to_owned()) - .unwrap_or_default(); - format!("- {name}: {desc}") - }) - .collect::<Vec<_>>() - .join("\n"); - - let task = "Summarize what a Rust trait object is for a beginner."; - - // The orchestrator designs the plan via structured output constrained to the - // registered names. - let mut orchestrator: AgentHarness<()> = AgentHarness::new(); - orchestrator - .register_model("model", model.clone()) - .set_default_model("model") - .with_policy(RunPolicy { - default_response_format: Some(ResponseFormat::json_schema( - "agent_selection", - json!({ - "type": "object", - "properties": { - "agents": { - "type": "array", - "items": { "type": "string", "enum": available } - } - }, - "required": ["agents"], - "additionalProperties": false - }), - )), - ..RunPolicy::default() - }); - - let plan = orchestrator - .invoke_default( - &(), - vec![ - Message::system(format!( - "You are an orchestrator with these named sub-agents available:\n{menu_text}\n\n\ - Choose the minimal subset whose skills solve the task. Respond ONLY with the \ - requested JSON object." - )), - Message::user(task), - ], - ) - .await - .expect("orchestrator run succeeds"); - - let mut chosen = parse_selection(&plan); - chosen.retain(|name| registry.has(ComponentKind::Tool, name)); - assert!( - !chosen.is_empty(), - "the orchestrator selected at least one registered sub-agent (got {chosen:?})" - ); - - // Resolve each chosen name from the registry and run them in parallel. - let dispatches = chosen.iter().enumerate().map(|(i, name)| { - let name = name.clone(); - let dispatch = dispatches - .get(&name) - .cloned() - .expect("a chosen name resolves in the registry"); - async move { - let parent = tinyagents_harness::context::RunContext::new( - tinyagents_harness::context::RunConfig::new(format!("dispatch-{i}")), - (), - ); - let result = dispatch - .invoke_in_parent_context( - &(), - json!({ "input": task }), - tinytools::ToolCallOptions::default(), - &parent, - ) - .await - .expect("sub-agent run succeeds"); - let job_id = serde_json::from_str::<Value>(&result.output()) - .expect("spawn result is JSON") - .get("job_id") - .and_then(Value::as_str) - .expect("spawn result contains a job id") - .to_owned(); - (name, job_id) - } - }); - let spawned: Vec<(String, String)> = join_all(dispatches).await; - let outputs = tokio::time::timeout(std::time::Duration::from_secs(120), async { - loop { - let outputs = spawned - .iter() - .filter_map(|(name, job_id)| { - let job = jobs.get(job_id)?; - job.status - .is_terminal() - .then(|| (name.clone(), job.output.unwrap_or_default())) - }) - .collect::<Vec<_>>(); - if outputs.len() == spawned.len() { - return outputs; - } - tokio::time::sleep(std::time::Duration::from_millis(25)).await; - } - }) - .await - .expect("live subagent jobs finish within two minutes"); - - assert!( - outputs.iter().any(|(_, text)| !text.trim().is_empty()), - "at least one resolved sub-agent returned non-empty text" - ); - - // Compose the resolved sub-agents' outputs into one final answer. - let context = outputs - .iter() - .map(|(name, text)| format!("[{name}]\n{text}")) - .collect::<Vec<_>>() - .join("\n\n"); - - let mut composer: AgentHarness<()> = AgentHarness::new(); - composer - .register_model("model", model) - .set_default_model("model"); - let composed = composer - .invoke_default( - &(), - vec![ - Message::system( - "Combine the labeled sub-agent outputs into one coherent answer. \ - Do not mention the sub-agents.", - ), - Message::user(format!("Task: {task}\n\nOutputs:\n{context}")), - ], - ) - .await - .expect("composer run succeeds"); - - let final_text = composed.text().unwrap_or_default(); - assert!( - !final_text.trim().is_empty(), - "the composed answer is non-empty" - ); - - eprintln!("orchestrator chose {chosen:?}; composed answer: {final_text}"); -} diff --git a/crates/tinyagents-orchestration/tests/live_subagent_error.rs b/crates/tinyagents-orchestration/tests/live_subagent_error.rs deleted file mode 100644 index 9be5b9358..000000000 --- a/crates/tinyagents-orchestration/tests/live_subagent_error.rs +++ /dev/null @@ -1,65 +0,0 @@ -//! LIVE: a real OpenAI sub-agent whose only tool fails surfaces the failure. -//! -//! The child [`AgentHarness`] is driven by a real [`OpenAiModel`] and given a -//! single tool [`FakeTool::failing`] plus a directive system prompt that makes -//! the model call it. The tool returns `Err(TinyAgentsError::Tool(..))`, which -//! propagates out of the child agent loop, so [`SubAgent::invoke`] returns the -//! failure. -//! -//! # Skips gracefully -//! -//! This test is `#[ignore]`d and only runs opted in via -//! `tests/common/live.rs::require_live`, so the default `cargo test` passes -//! with no key configured and never dials a real provider by accident. - -mod common; - -#[tokio::test] -#[ignore = "network: set TINYAGENTS_LIVE=1 and run with --ignored"] -async fn live_openai_subagent_surfaces_tool_failure() { - use std::sync::Arc; - - use tinyagents_harness::error::TinyAgentsError; - use tinyagents_harness::runtime::AgentHarness; - use tinyagents_harness::testkit::FakeTool; - use tinyagents_orchestration::subagent::SubAgent; - use tinyinference_llm::providers::openai::OpenAiModel; - - if !common::live::require_live(&["OPENAI_API_KEY"]) { - return; - } - - // Child agent backed by a real model, with a single tool that always fails. - let mut child: AgentHarness<()> = AgentHarness::new(); - child - .register_model( - "openai", - Arc::new(OpenAiModel::from_env().expect("OPENAI_API_KEY present")), - ) - .set_default_model("openai"); - child.register_tool(Arc::new(FakeTool::failing("lookup", "upstream 500"))); - - let subagent = SubAgent::new( - "lookup_agent", - "Looks things up. Always uses the `lookup` tool to answer.", - Arc::new(child), - ) - .with_system_prompt( - "You must answer every request by calling the `lookup` tool exactly once with any \ - arguments. Never answer directly; always call `lookup` first.", - ); - - let err = subagent - .invoke(&(), (), 0, "Look up the capital of France.") - .await - .expect_err("the failing tool must surface as an error from the sub-agent run"); - - // The failure surfaces as the tool error (its message is propagated). - match err { - TinyAgentsError::Tool(msg) => assert!( - msg.contains("upstream 500"), - "expected the tool error message to propagate, got {msg:?}" - ), - other => panic!("expected TinyAgentsError::Tool, got {other:?}"), - } -} diff --git a/crates/tinyagents-orchestration/tests/live_subagent_reuse.rs b/crates/tinyagents-orchestration/tests/live_subagent_reuse.rs deleted file mode 100644 index fa3809470..000000000 --- a/crates/tinyagents-orchestration/tests/live_subagent_reuse.rs +++ /dev/null @@ -1,100 +0,0 @@ -//! LIVE end-to-end: post-completion sub-agent REUSE against a real OpenAI model. -//! -//! A real OpenAI sub-agent answers a first question, then a follow-up that -//! *depends on the first answer* — proving the sub-agent was reused with its -//! prior conversation context intact (rather than killed and restarted). This -//! is the network-backed sibling of `e2e_subagent_reuse.rs`. -//! -//! # Skips gracefully -//! -//! This test is `#[ignore]`d and only runs opted in via -//! `tests/common/live.rs::require_live`, so the default `cargo test` passes -//! with no key configured and never dials a real provider by accident. - -mod common; - -#[tokio::test] -#[ignore = "network: set TINYAGENTS_LIVE=1 and run with --ignored"] -async fn live_openai_subagent_reused_with_carried_context() { - use std::sync::Arc; - - use tinyagents_harness::runtime::AgentHarness; - use tinyagents_harness::testkit::{EventRecorder, Trajectory}; - use tinyagents_orchestration::subagent::{SubAgent, SubAgentSession}; - use tinyinference_llm::message::Message; - use tinyinference_llm::providers::openai::OpenAiModel; - - if !common::live::require_live(&["OPENAI_API_KEY"]) { - return; - } - - // A real OpenAI sub-agent we will reuse across two turns. - let mut harness: AgentHarness<()> = AgentHarness::new(); - harness - .register_model( - "openai", - Arc::new(OpenAiModel::from_env().expect("OPENAI_API_KEY present")), - ) - .set_default_model("openai"); - - let subagent = Arc::new( - SubAgent::new( - "assistant", - "a helpful assistant that remembers the conversation", - Arc::new(harness), - ) - .with_system_prompt( - "You are a concise assistant. Answer using the full prior conversation context.", - ), - ); - - let recorder = EventRecorder::new(); - let mut session = SubAgentSession::new(subagent).with_events(recorder.sink()); - - // ── Turn 1: establish a fact the follow-up will depend on. ─────────────── - let first = session - .send( - &(), - (), - vec![Message::user( - "My favorite color is teal. Please acknowledge it in one short sentence.", - )], - ) - .await - .expect("first live sub-agent run"); - let first_text = first.text().unwrap_or_default(); - assert!( - !first_text.trim().is_empty(), - "first turn should produce a non-empty answer" - ); - - // ── Turn 2: a follow-up that ONLY succeeds if prior context was reused. ─── - let second = session - .send( - &(), - (), - vec![Message::user( - "What is my favorite color? Answer with just the color word.", - )], - ) - .await - .expect("second live sub-agent run (reused)"); - let second_text = second.text().unwrap_or_default().to_lowercase(); - - assert_eq!(session.turns(), 2); - assert!( - second_text.contains("teal"), - "the reused sub-agent should recall the earlier answer (favorite color), got: {second_text:?}" - ); - - // The reuse is observable in the trajectory. - let traj = Trajectory::from_events(recorder.events()); - traj.assert_completed(); - assert!( - recorder - .events() - .iter() - .any(|e| matches!(e, tinyagents_harness::events::AgentEvent::SubAgentReused { turn, .. } if *turn == 1)), - "the second send should record a SubAgentReused event" - ); -} diff --git a/crates/tinyagents-orchestration/tests/live_subagent_timeout.rs b/crates/tinyagents-orchestration/tests/live_subagent_timeout.rs deleted file mode 100644 index be55684cd..000000000 --- a/crates/tinyagents-orchestration/tests/live_subagent_timeout.rs +++ /dev/null @@ -1,57 +0,0 @@ -//! LIVE: a real OpenAI sub-agent with a tiny wall-clock cap times out. -//! -//! The child [`AgentHarness`] is driven by a real [`OpenAiModel`] but its policy -//! caps wall-clock time at 1ms. Any real HTTP round-trip vastly exceeds that, so -//! the per-model-call budget interrupts the in-flight request and -//! [`SubAgent::invoke`] returns `Err(TinyAgentsError::Timeout(..))`. -//! -//! # Skips gracefully -//! -//! This test is `#[ignore]`d and only runs opted in via -//! `tests/common/live.rs::require_live`, so the default `cargo test` passes -//! with no key configured and never dials a real provider by accident. - -mod common; - -#[tokio::test] -#[ignore = "network: set TINYAGENTS_LIVE=1 and run with --ignored"] -async fn live_openai_subagent_times_out_on_tiny_budget() { - use std::sync::Arc; - - use tinyagents_harness::error::TinyAgentsError; - use tinyagents_harness::limits::RunLimits; - use tinyagents_harness::runtime::{AgentHarness, RunPolicy}; - use tinyagents_orchestration::subagent::SubAgent; - use tinyinference_llm::providers::openai::OpenAiModel; - - if !common::live::require_live(&["OPENAI_API_KEY"]) { - return; - } - - // Child agent backed by a real model, with a 1ms wall-clock cap. A real - // network call cannot complete in 1ms, so the per-call budget trips. - let mut child: AgentHarness<()> = AgentHarness::new(); - child - .register_model( - "openai", - Arc::new(OpenAiModel::from_env().expect("OPENAI_API_KEY present")), - ) - .set_default_model("openai"); - child.with_policy(RunPolicy { - limits: RunLimits::default().with_max_wall_clock_ms(Some(1)), - ..RunPolicy::default() - }); - - let subagent = SubAgent::new( - "slow_lookup", - "Answers questions via a real model.", - Arc::new(child), - ); - - let err = subagent - .invoke(&(), (), 0, "What is 2 + 2?") - .await - .expect_err("a 1ms budget must trip on any real network call"); - - assert!(matches!(err, TinyAgentsError::Timeout(_)), "got {err:?}"); -} diff --git a/crates/tinyagents-orchestration/tests/public_parent_job_scope.rs b/crates/tinyagents-orchestration/tests/public_parent_job_scope.rs new file mode 100644 index 000000000..6fab3ac49 --- /dev/null +++ b/crates/tinyagents-orchestration/tests/public_parent_job_scope.rs @@ -0,0 +1,62 @@ +use tinyagents_harness::cancel::CancellationToken; +use tinyagents_orchestration::subagent::{ + DriverJobRegistration, ParentJobScope, SubAgentJobRegistry, format_parent_roster_for_owner, +}; +use tinytools::WorkspaceDescriptor; + +#[test] +fn parent_job_scope_is_available_from_the_public_subagent_facade() { + let constructor: fn(WorkspaceDescriptor, String) -> ParentJobScope = ParentJobScope::new; + let _ = constructor; + let registry = SubAgentJobRegistry::new(); + let scope = ParentJobScope::new(WorkspaceDescriptor::new("/workspace"), "thread"); + let rendered = + format_parent_roster_for_owner(®istry, "owner", &scope, &[], |_| "empty".to_owned()) + .unwrap(); + assert_eq!(rendered, "empty"); +} + +#[test] +fn parent_scope_can_cancel_only_its_jobs_through_the_public_facade() { + let registry = SubAgentJobRegistry::new(); + let workspace_a = WorkspaceDescriptor::new("/workspace/a"); + let workspace_b = WorkspaceDescriptor::new("/workspace/b"); + let token_a = CancellationToken::new(); + let token_b = CancellationToken::new(); + for (task, thread, workspace, cancellation) in [ + ( + "task-a", + "same-thread", + workspace_a.clone(), + token_a.clone(), + ), + ("task-b", "same-thread", workspace_b, token_b.clone()), + ] { + registry + .register_driver_job(DriverJobRegistration { + key: tinyagents_orchestration::subagent::SubagentTaskKey { + root_run_id: "root".into(), + parent_run_id: "parent".into(), + thread_id: Some(thread.into()), + task_id: task.into(), + }, + agent: "worker".into(), + owner: "owner".into(), + workspace: Some(workspace), + cancellation, + steering: None, + run_queue: None, + continuing: false, + }) + .unwrap(); + } + + assert_eq!( + registry + .cancel_for_thread(&ParentJobScope::new(workspace_a, "same-thread")) + .unwrap(), + 1 + ); + assert!(token_a.is_cancelled()); + assert!(!token_b.is_cancelled()); +} diff --git a/crates/tinyagents-registry/Cargo.toml b/crates/tinyagents-registry/Cargo.toml index 43e32757e..a4622b1d8 100644 --- a/crates/tinyagents-registry/Cargo.toml +++ b/crates/tinyagents-registry/Cargo.toml @@ -13,8 +13,8 @@ anyhow = { workspace = true } chrono = { workspace = true } serde = { workspace = true } serde_json = { workspace = true } -tinyagents-definition = { path = "../tinyagents-definition", version = "2.11.0" } -tinyagents-harness = { path = "../tinyagents-harness", version = "2.11.0", default-features = false } +tinyagents-definition = { path = "../tinyagents-definition", version = "3.0.0" } +tinyagents-harness = { path = "../tinyagents-harness", version = "3.0.0", default-features = false } tinyinference-llm = { path = "../../vendor/tinyinference/crates/tinyinference-llm", version = "0.5.0" } tinyinference-providers = { path = "../../vendor/tinyinference/crates/tinyinference-providers", version = "0.5.0" } tinytools = { path = "../../vendor/tinytools/crates/tinytools", version = "0.6.0" } diff --git a/crates/tinyagents-registry/src/README.md b/crates/tinyagents-registry/src/README.md index b6d489eda..909b25980 100644 --- a/crates/tinyagents-registry/src/README.md +++ b/crates/tinyagents-registry/src/README.md @@ -47,8 +47,8 @@ only capabilities a human explicitly registered can be invoked. ## File and module map - **`capability/`** — the [`CapabilityRegistry`] implementation: registration, - lookup, aliasing, duplicate detection, and bridges to the harness - ([`ModelRegistry`], [`ToolRegistry`]). + lookup, aliasing, duplicate detection, and tool projection to the harness + ([`ToolRegistry`]); model execution uses the harness resolver port. - `types.rs` — the [`CapabilityRegistry`] struct and storage maps. - `mod.rs` — registration and accessor methods. - `mod_tests.rs` — tests for registration, alias, and lookup. diff --git a/crates/tinyagents-registry/src/capability/mod.rs b/crates/tinyagents-registry/src/capability/mod.rs index d94ac0e86..a94f64399 100644 --- a/crates/tinyagents-registry/src/capability/mod.rs +++ b/crates/tinyagents-registry/src/capability/mod.rs @@ -7,31 +7,29 @@ //! //! See `types` for the data definitions. This module provides registration, //! lookup, aliasing, duplicate validation, and conveniences for handing the -//! catalog's models and tools to a harness ([`to_model_registry`] / -//! [`to_tool_registry`]). +//! catalog's tools to a harness ([`to_tool_registry`]). //! -//! [`to_model_registry`]: CapabilityRegistry::to_model_registry //! [`to_tool_registry`]: CapabilityRegistry::to_tool_registry mod types; +mod value; use std::sync::Arc; use crate::component::{ComponentKind, ComponentMetadata}; use tinyagents_harness::error::{Result, TinyAgentsError}; -use tinyagents_harness::model_registry::ModelRegistry; use tinyagents_harness::tool::ToolRegistry; use tinyinference_llm::model::ChatModel; use tinytools::Tool; pub use types::*; +pub use value::Capability; impl<State: Send + Sync> CapabilityRegistry<State> { /// Creates an empty registry. pub fn new() -> Self { Self { models: std::collections::HashMap::new(), - model_order: Vec::new(), tools: std::collections::HashMap::new(), agents: std::collections::HashMap::new(), meta: std::collections::HashMap::new(), @@ -108,7 +106,7 @@ impl<State: Send + Sync> CapabilityRegistry<State> { /// Replaces the [`ComponentMetadata`] recorded for `(kind, name)`. /// - /// Unlike [`record_meta`](Self::record_meta) (which only fills in a + /// Unlike `record_meta` (which only fills in a /// default the first time a name is registered), this overwrites whatever /// metadata is already there — so `with_description`/`with_tag` builders /// (`crate::component::ComponentMetadata`) actually reach a registered @@ -156,7 +154,6 @@ impl<State: Send + Sync> CapabilityRegistry<State> { match kind { ComponentKind::Model => { self.models.remove(name); - self.model_order.retain(|n| n != name); } ComponentKind::Tool => { self.tools.remove(name); @@ -195,7 +192,6 @@ impl<State: Send + Sync> CapabilityRegistry<State> { let name = name.into(); self.ensure_absent(ComponentKind::Model, &name)?; self.record_meta(ComponentKind::Model, &name); - self.remember_model_order(&name); self.models.insert(name, model); Ok(self) } @@ -209,22 +205,12 @@ impl<State: Send + Sync> CapabilityRegistry<State> { ) -> &mut Self { let name = name.into(); self.record_meta(ComponentKind::Model, &name); - self.remember_model_order(&name); self.models.insert(name, model); self } - /// Appends `name` to [`Self::model_order`] the first time it is - /// registered. Re-registering an existing name (via - /// [`replace_model`](Self::replace_model)) keeps its original position. - fn remember_model_order(&mut self, name: &str) { - if !self.models.contains_key(name) { - self.model_order.push(name.to_owned()); - } - } - /// Registers a model under `name` with explicit [`ComponentMetadata`] - /// instead of the bare default [`record_meta`](Self::record_meta) would + /// instead of the bare default `record_meta` would /// attach, so a description/tags are attached atomically with /// registration rather than needing a follow-up /// [`set_metadata`](Self::set_metadata) call. @@ -243,7 +229,6 @@ impl<State: Send + Sync> CapabilityRegistry<State> { self.ensure_absent(ComponentKind::Model, &name)?; self.meta .insert((ComponentKind::Model, name.clone()), metadata); - self.remember_model_order(&name); self.models.insert(name, model); Ok(self) } @@ -338,8 +323,7 @@ impl<State: Send + Sync> CapabilityRegistry<State> { // Registration: capability bundles (gap G3) // ----------------------------------------------------------------------- - /// Registers a [`tinyagents_harness::capability::Capability`] bundle - /// under its own [`Capability::name`](tinyagents_harness::capability::Capability::name). + /// Registers a [`Capability`] bundle under its own name. /// /// # Why `Ctx` is a method type parameter, not stored on `Self` /// @@ -348,12 +332,9 @@ impl<State: Send + Sync> CapabilityRegistry<State> { /// definitions — is `Ctx`-free: `Arc<dyn ChatModel<State>>`, `Arc<dyn /// tinytools::Tool>`, and `AgentDefinition` none of them name a `Ctx` /// type. A `Capability`, - /// however, composes `tinyagents_harness::tool::toolset::ToolSet<State, - /// Ctx>` and `tinyagents_harness::middleware::Middleware<State, Ctx>` - /// trait objects (see `tinyagents-harness`'s `capability` module doc - /// comment for why `Capability` itself cannot live in - /// `tinyagents-definition`, the lower crate that would otherwise be the - /// natural home for a registry-stored declarative bundle). + /// however, composes harness `ToolSet<State, Ctx>` and `Middleware<State, + /// Ctx>` trait objects, so its owning type lives here rather than in the + /// lower `tinyagents-definition` crate. /// /// Adding a `Ctx` type parameter to `CapabilityRegistry<State>` itself — /// so every one of its dozen kinds carries a `Ctx` dimension only this @@ -372,7 +353,7 @@ impl<State: Send + Sync> CapabilityRegistry<State> { /// already registered under this name. pub fn register_capability<Ctx>( &mut self, - capability: tinyagents_harness::capability::Capability<State, Ctx>, + capability: Capability<State, Ctx>, ) -> Result<&mut Self> where State: 'static, @@ -388,10 +369,7 @@ impl<State: Send + Sync> CapabilityRegistry<State> { /// Registers or overwrites a capability bundle under its own name, /// preserving any existing metadata. See [`Self::register_capability`] /// for why `Ctx` is a method type parameter. - pub fn replace_capability<Ctx>( - &mut self, - capability: tinyagents_harness::capability::Capability<State, Ctx>, - ) -> &mut Self + pub fn replace_capability<Ctx>(&mut self, capability: Capability<State, Ctx>) -> &mut Self where State: 'static, Ctx: Send + Sync + 'static, @@ -409,10 +387,7 @@ impl<State: Send + Sync> CapabilityRegistry<State> { /// when one is registered but under a different `Ctx` than the one /// requested here (a caller-programming-error case — a registry is /// normally used with one consistent `Ctx` per application). - pub fn capability<Ctx>( - &self, - name: &str, - ) -> Option<&tinyagents_harness::capability::Capability<State, Ctx>> + pub fn capability<Ctx>(&self, name: &str) -> Option<&Capability<State, Ctx>> where State: 'static, Ctx: Send + Sync + 'static, @@ -593,61 +568,6 @@ impl<State: Send + Sync> CapabilityRegistry<State> { self.meta.get(&(kind, canonical)) } - // ----------------------------------------------------------------------- - // Handoff to harness registries - // ----------------------------------------------------------------------- - - /// Builds a harness [`ModelRegistry`] from the registered models, including - /// alias names bound to the same model handle. - /// - /// Models are registered onto the result in [`Self::model_order`] — the - /// order `register_model`/`replace_model` first saw each name in — so the - /// harness registry's "first-registered model becomes the default" rule - /// ([`ModelRegistry::register`]) is deterministic and reproducible across - /// runs, rather than following `HashMap` iteration order. That default is - /// still whichever model happened to be registered first; callers who - /// need a specific default regardless of registration order should use - /// [`Self::to_model_registry_with_default`] or call `set_default` - /// explicitly on the result. - pub fn to_model_registry(&self) -> ModelRegistry<State> { - let mut registry = ModelRegistry::new(); - for name in &self.model_order { - if let Some(model) = self.models.get(name) { - registry.register(name.clone(), model.clone()); - } - } - let mut aliases: Vec<(&String, &String)> = self - .aliases - .iter() - .filter(|((kind, _), _)| *kind == ComponentKind::Model) - .map(|((_, alias), target)| (alias, target)) - .collect(); - aliases.sort(); - for (alias, target) in aliases { - if let Some(model) = self.models.get(target) { - registry.register(alias.clone(), model.clone()); - } - } - registry - } - - /// Builds a harness [`ModelRegistry`] exactly like [`Self::to_model_registry`], - /// but with the default model explicitly set to `name` instead of - /// whichever model was registered first. - /// - /// # Errors - /// - /// Returns [`TinyAgentsError::ModelNotFound`] if `name` (or an alias of - /// it) is not a registered model. - pub fn to_model_registry_with_default(&self, name: &str) -> Result<ModelRegistry<State>> { - if self.model(name).is_none() { - return Err(TinyAgentsError::ModelNotFound(name.to_string())); - } - let mut registry = self.to_model_registry(); - registry.set_default(name); - Ok(registry) - } - /// Builds a harness [`ToolRegistry`] from the registered tools. /// /// The harness [`ToolRegistry`] keys tools by their own [`Tool::name`], so diff --git a/crates/tinyagents-registry/src/capability/mod_tests.rs b/crates/tinyagents-registry/src/capability/mod_tests.rs index 2647fad51..5a642d991 100644 --- a/crates/tinyagents-registry/src/capability/mod_tests.rs +++ b/crates/tinyagents-registry/src/capability/mod_tests.rs @@ -230,7 +230,7 @@ fn alias_validation_rejects_unknown_target_and_duplicates() { } #[tokio::test] -async fn builds_harness_registries_with_model_aliases() { +async fn model_aliases_remain_in_the_capability_catalog() { let mut reg = CapabilityRegistry::<()>::new(); reg.register_model("gpt-4o", Arc::new(FakeModel("hello"))) .unwrap(); @@ -239,77 +239,17 @@ async fn builds_harness_registries_with_model_aliases() { reg.alias(ComponentKind::Model, "default", "gpt-4o") .unwrap(); - let models = reg.to_model_registry(); - assert!(models.get("gpt-4o").is_some()); - assert!(models.get("default").is_some()); + assert!(reg.model("gpt-4o").is_some()); + assert!(reg.model("default").is_some()); + assert!( + reg.names(ComponentKind::Model) + .contains(&"gpt-4o".to_string()) + ); let tools = reg.to_tool_registry::<()>(); assert_eq!(tools.names(), vec!["lookup_user"]); } -/// Builds a registry with `charlie`, `alpha`, `bravo` registered in that -/// exact order (deliberately not alphabetical, so a name-sorted iteration -/// would pick a different "first" model than registration order does). -fn registry_with_three_models_in_order() -> CapabilityRegistry<()> { - let mut reg = CapabilityRegistry::<()>::new(); - reg.register_model("charlie", Arc::new(FakeModel("c"))) - .unwrap(); - reg.register_model("alpha", Arc::new(FakeModel("a"))) - .unwrap(); - reg.register_model("bravo", Arc::new(FakeModel("b"))) - .unwrap(); - reg -} - -#[test] -fn to_model_registry_default_is_the_first_registered_model_every_time() { - // Build the same registry several times over; a `HashMap`-order default - // would vary run to run (or construction to construction within a - // process, depending on hash-seed timing), while first-registration - // order should not. - for _ in 0..5 { - let reg = registry_with_three_models_in_order(); - let models = reg.to_model_registry(); - assert_eq!( - models.default_name(), - Some("charlie"), - "default model should always be the first one registered" - ); - assert!(models.get("charlie").is_some()); - assert!(models.get("alpha").is_some()); - assert!(models.get("bravo").is_some()); - } -} - -#[test] -fn replace_model_does_not_move_an_existing_name_in_registration_order() { - let mut reg = registry_with_three_models_in_order(); - // Re-registering "bravo" (already registered second) must not make it - // the new first-registered name. - reg.replace_model("bravo", Arc::new(FakeModel("b2"))); - reg.register_model("delta", Arc::new(FakeModel("d"))) - .unwrap(); - - let models = reg.to_model_registry(); - assert_eq!(models.default_name(), Some("charlie")); -} - -#[test] -fn to_model_registry_with_default_overrides_first_registered() { - let reg = registry_with_three_models_in_order(); - - let models = reg - .to_model_registry_with_default("bravo") - .expect("bravo is registered"); - assert_eq!(models.default_name(), Some("bravo")); - assert!(models.get("charlie").is_some()); - - let err = reg - .to_model_registry_with_default("not-registered") - .unwrap_err(); - assert!(matches!(err, TinyAgentsError::ModelNotFound(name) if name == "not-registered")); -} - #[test] fn snapshot_lists_components_sorted_by_kind_and_name() { let mut reg = CapabilityRegistry::<()>::new(); @@ -601,3 +541,17 @@ fn with_router_builder_installs_the_routing_policy() { assert!(reg.router().route("chat-v1").is_some()); assert!(reg.route_workload("chat-v1").is_some()); } + +#[test] +fn stores_and_looks_up_registry_owned_capability_data() { + let capability = Capability::<(), ()>::new("research").with_instructions("Use research tools."); + let mut registry = CapabilityRegistry::new(); + registry.register_capability(capability).unwrap(); + + assert_eq!( + registry + .capability::<()>("research") + .and_then(|capability| capability.instructions.as_deref()), + Some("Use research tools.") + ); +} diff --git a/crates/tinyagents-registry/src/capability/types.rs b/crates/tinyagents-registry/src/capability/types.rs index 63197ea3b..169c0ca15 100644 --- a/crates/tinyagents-registry/src/capability/types.rs +++ b/crates/tinyagents-registry/src/capability/types.rs @@ -2,9 +2,9 @@ //! addressable by name — the data structure that lets a workflow reference //! sub-capabilities it didn't hardcode. //! -//! This layer is deliberately distinct from the harness' -//! [`tinyagents_harness::model_registry::ModelRegistry`] and -//! [`tinyagents_harness::tool::ToolRegistry`], which are per-run executable stores. +//! This layer is the named catalog. Hosts adapt selected models through the +//! harness model resolver port and can project tools into a per-run +//! [`tinyagents_harness::tool::ToolRegistry`]. //! The [`CapabilityRegistry`] is a *capability catalog*: it owns named models, //! tools, agents, graphs, routers, and reducers so host code can resolve //! capabilities by name. @@ -39,16 +39,6 @@ where State: Send + Sync, { pub(crate) models: HashMap<String, Arc<dyn ChatModel<State>>>, - /// Canonical model names in first-registration order. - /// - /// `models` is a `HashMap`, whose iteration order is randomized per - /// process; without this, [`CapabilityRegistry::to_model_registry`]'s - /// "first-registered model becomes the default" choice - /// ([`tinyagents_harness::model_registry::ModelRegistry::register`]) - /// would silently vary per run. `register_model`/`replace_model` append a - /// name here the first time it is registered; re-registering an existing - /// name (via `replace_model`) does not move it. - pub(crate) model_order: Vec<String>, pub(crate) tools: HashMap<String, Arc<dyn Tool>>, /// Declarative agent definitions keyed by their stable id. Execution is /// host-owned through graph's explicit `AgentInvoker` boundary. @@ -58,7 +48,7 @@ where /// Alias map, keyed by `(kind, alias)` -> canonical name. pub(crate) aliases: HashMap<(ComponentKind, String), String>, /// Capability bundles (gap G3), type-erased since - /// `tinyagents_harness::capability::Capability<State, Ctx>` is generic + /// `crate::Capability<State, Ctx>` is generic /// over `Ctx`, a dimension this registry does not otherwise carry — see /// [`crate::CapabilityRegistry::register_capability`]'s doc comment for /// why. Keyed by canonical name, same as every other kind here. diff --git a/crates/tinyagents-registry/src/capability/value.rs b/crates/tinyagents-registry/src/capability/value.rs new file mode 100644 index 000000000..701eaa5cd --- /dev/null +++ b/crates/tinyagents-registry/src/capability/value.rs @@ -0,0 +1,79 @@ +//! Registry-owned executable capability data. + +use std::sync::Arc; + +use tinyagents_harness::middleware::Middleware; +use tinyagents_harness::tool::toolset::ToolSet; + +/// A named bundle of instructions, tools, and middleware stored by the +/// [`CapabilityRegistry`](crate::CapabilityRegistry). +/// +/// The registry type-erases this value because `Ctx` belongs to each bundle, +/// not to the registry as a whole. This type describes stored capability data; +/// installing it into a running harness is a host responsibility. +pub struct Capability<State: Send + Sync, Ctx: Send + Sync> { + /// Stable name used to register and look up this capability. + pub name: String, + /// Instructions associated with the capability. + pub instructions: Option<String>, + /// Tools associated with the capability. + pub toolset: Option<Arc<dyn ToolSet<State, Ctx>>>, + /// Middleware associated with the capability. + pub middleware: Vec<Arc<dyn Middleware<State, Ctx>>>, +} + +impl<State: Send + Sync, Ctx: Send + Sync> Capability<State, Ctx> { + /// Creates a capability with the given stable name and no attachments. + pub fn new(name: impl Into<String>) -> Self { + Self { + name: name.into(), + instructions: None, + toolset: None, + middleware: Vec::new(), + } + } + + /// Sets the instructions associated with this capability. + #[must_use] + pub fn with_instructions(mut self, instructions: impl Into<String>) -> Self { + self.instructions = Some(instructions.into()); + self + } + + /// Sets the toolset associated with this capability. + #[must_use] + pub fn with_toolset(mut self, toolset: Arc<dyn ToolSet<State, Ctx>>) -> Self { + self.toolset = Some(toolset); + self + } + + /// Appends middleware to this capability. + #[must_use] + pub fn with_middleware(mut self, middleware: Arc<dyn Middleware<State, Ctx>>) -> Self { + self.middleware.push(middleware); + self + } +} + +impl<State: Send + Sync, Ctx: Send + Sync> Clone for Capability<State, Ctx> { + fn clone(&self) -> Self { + Self { + name: self.name.clone(), + instructions: self.instructions.clone(), + toolset: self.toolset.clone(), + middleware: self.middleware.clone(), + } + } +} + +impl<State: Send + Sync, Ctx: Send + Sync> std::fmt::Debug for Capability<State, Ctx> { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("Capability") + .field("name", &self.name) + .field("instructions", &self.instructions) + .field("has_toolset", &self.toolset.is_some()) + .field("middleware_count", &self.middleware.len()) + .finish() + } +} diff --git a/crates/tinyagents-registry/src/component/types.rs b/crates/tinyagents-registry/src/component/types.rs index 07d7dbab5..53073eaf6 100644 --- a/crates/tinyagents-registry/src/component/types.rs +++ b/crates/tinyagents-registry/src/component/types.rs @@ -49,11 +49,11 @@ pub enum ComponentKind { TaskStore, /// An event listener descriptor (name-only for now). Listener, - /// An executable capability bundle (gap G3): instructions, toolset, - /// middleware, model defaults, exposure, and `defer_loading`. Unlike + /// An executable capability bundle: instructions, toolset, and + /// middleware. Unlike /// every other executable kind above (`Model`, `Tool`, `Graph`, /// `Agent`), the stored value is `Ctx`-generic - /// (`tinyagents_harness::capability::Capability<State, Ctx>`), so + /// (`crate::Capability<State, Ctx>`), so /// [`crate::CapabilityRegistry::register_capability`] stores it /// type-erased rather than in a `Ctx`-free typed field — see that /// method's doc comment. diff --git a/crates/tinyagents-registry/src/lib.rs b/crates/tinyagents-registry/src/lib.rs index 893624783..06151fa2b 100644 --- a/crates/tinyagents-registry/src/lib.rs +++ b/crates/tinyagents-registry/src/lib.rs @@ -34,7 +34,7 @@ pub mod router; pub use tinyagents_harness::error::{Result, TinyAgentsError}; -pub use capability::CapabilityRegistry; +pub use capability::{Capability, CapabilityRegistry}; pub use catalog::{ ModelCapabilities, ModelCatalog, ModelCatalogEntry, ModelCatalogSnapshot, ModelCatalogSource, }; diff --git a/crates/tinyagents-registry/src/router/mod.rs b/crates/tinyagents-registry/src/router/mod.rs index 89391cc87..23eba72c0 100644 --- a/crates/tinyagents-registry/src/router/mod.rs +++ b/crates/tinyagents-registry/src/router/mod.rs @@ -4,9 +4,9 @@ //! //! # Why this exists //! -//! [`CapabilityRegistry`](crate::CapabilityRegistry) and the harness -//! [`ModelRegistry`](tinyagents_harness::ModelRegistry) resolve a model *by -//! name* and the agent loop can fail over across a +//! [`CapabilityRegistry`](crate::CapabilityRegistry) stores models by name and +//! the harness model resolver adapts execution through its host port; the agent +//! loop can fail over across a //! [`FallbackPolicy`] — but neither owns //! the *policy* that maps a host's **workload tiers** (`chat-v1`, `reasoning-v1`, //! `vision-v1`, …) onto concrete models, nor the per-tier capability gates and diff --git a/crates/tinyagents-registry/src/router/types.rs b/crates/tinyagents-registry/src/router/types.rs index 5949f70a3..eac101322 100644 --- a/crates/tinyagents-registry/src/router/types.rs +++ b/crates/tinyagents-registry/src/router/types.rs @@ -78,8 +78,7 @@ pub struct FallbackDiagnostics { /// /// A route is pure metadata — it names a model rather than owning one, so the /// same route table can be declared once and projected onto any -/// [`CapabilityRegistry`](crate::CapabilityRegistry) / -/// [`ModelRegistry`](tinyagents_harness::ModelRegistry) that has registered +/// [`CapabilityRegistry`](crate::CapabilityRegistry) that has registered /// those model names. This is what lets a host describe its tiered routing /// (workload aliases → concrete BYOK/managed/local models) declaratively and hand /// it to the crate, instead of re-implementing alias resolution + fallback diff --git a/crates/tinyagents-runtime/Cargo.toml b/crates/tinyagents-runtime/Cargo.toml index a30ba1bf1..4b714e7ff 100644 --- a/crates/tinyagents-runtime/Cargo.toml +++ b/crates/tinyagents-runtime/Cargo.toml @@ -13,11 +13,12 @@ async-trait = { workspace = true } anyhow = { workspace = true } chrono = { workspace = true } serde_json = { workspace = true } +sha2 = "0.11" thiserror = "2" -tinyagents-harness = { path = "../tinyagents-harness", version = "2.11.0" } -tinyagents-hooks = { path = "../tinyagents-hooks", version = "2.11.0" } -tinyagents-registry = { path = "../tinyagents-registry", version = "2.11.0" } -tinyagents-session = { path = "../tinyagents-session", version = "2.11.0" } +tinyagents-harness = { path = "../tinyagents-harness", version = "3.0.0" } +tinyagents-hooks = { path = "../tinyagents-hooks", version = "3.0.0" } +tinyagents-registry = { path = "../tinyagents-registry", version = "3.0.0" } +tinyagents-session = { path = "../tinyagents-session", version = "3.0.0" } tinyinference-llm = { path = "../../vendor/tinyinference/crates/tinyinference-llm", version = "0.5.0" } tinytools = { path = "../../vendor/tinytools/crates/tinytools", version = "0.6.0" } tokio = { workspace = true, features = ["io-util", "macros", "process", "rt", "sync", "time"] } @@ -26,7 +27,11 @@ tracing = { workspace = true } [dev-dependencies] anyhow = { workspace = true } tinyagents-storage-kit = { workspace = true } +tinyagents-session = { path = "../tinyagents-session", version = "3.0.0", features = ["testkit"] } tinystoragedrivers-api = { workspace = true } +tinystoragedrivers = { workspace = true } +tinystoragedrivers-connect = { workspace = true } +tinystoragedrivers-document = { workspace = true } tinystoragedrivers-testkit = { workspace = true } tinystoragedrivers-types = { workspace = true } serde_json = { workspace = true } diff --git a/crates/tinyagents-runtime/README.md b/crates/tinyagents-runtime/README.md index d9ece737b..2bf906e33 100644 --- a/crates/tinyagents-runtime/README.md +++ b/crates/tinyagents-runtime/README.md @@ -6,6 +6,14 @@ the mutable model history, a stable prefix, and append-only transcript state. Tool declarations are prepared afresh for every driver call. It has no host configuration, credentials, prompt construction, tool authorization, model selection, or event system. +Sub-agent turns are leaf turns by default: registration refuses spawn tools +even when the child request or tool allowlist includes them. A host may opt a +turn into nested spawning by setting `ToolAdmission::nested_spawn_depth` to a +validated `NestedSpawnDepth`. The runtime admits a spawn tool only when it is +also present in the ordinary tool allowlist and another child fits within the +supplied depth ceiling. The host remains responsible for enforcing that same +ceiling when it executes the spawn. + ## Host responsibilities The host supplies three narrow seams: diff --git a/crates/tinyagents-runtime/src/builder.rs b/crates/tinyagents-runtime/src/builder.rs index 3764b256f..1289f9c49 100644 --- a/crates/tinyagents-runtime/src/builder.rs +++ b/crates/tinyagents-runtime/src/builder.rs @@ -3,6 +3,8 @@ use tinyagents_harness::tool::ToolSnapshot; use tinyagents_session::transcript::{SessionRef, TranscriptLocator, TranscriptMeta, session_stem}; +use crate::SessionIdentityResolver; +use crate::identity::IdentityConfig; use crate::{ NoopSessionHooks, PrefixSnapshot, RuntimeError, Session, SessionDriver, SessionHooks, TranscriptCodec, @@ -13,6 +15,8 @@ pub struct SessionBuilder<C: Clone + Send + Sync + 'static = ()> { driver: Arc<dyn SessionDriver<C>>, codec: Option<Arc<dyn TranscriptCodec<C>>>, hooks: Arc<dyn SessionHooks<C>>, + post_commit_sinks: Vec<Arc<dyn crate::hooks::PostCommitSink<C>>>, + identity: Option<IdentityConfig<C>>, prefix: PrefixSnapshot, tools: ToolSnapshot, retain_recorded_tools: bool, @@ -34,6 +38,8 @@ impl<C: Clone + Send + Sync + 'static> SessionBuilder<C> { driver, codec: None, hooks: Arc::new(NoopSessionHooks), + post_commit_sinks: Vec::new(), + identity: None, prefix: PrefixSnapshot::default(), tools: ToolSnapshot::default(), retain_recorded_tools: false, @@ -53,6 +59,48 @@ impl<C: Clone + Send + Sync + 'static> SessionBuilder<C> { self } + /// Adds a host observer for successful durable commits. + pub fn post_commit_sink(mut self, sink: Arc<dyn crate::hooks::PostCommitSink<C>>) -> Self { + self.post_commit_sinks.push(sink); + self + } + + /// Resolves a root session's logical key against host-authorized storage. + pub fn identity_resolver( + mut self, + session_key: impl Into<String>, + meta: TranscriptMeta, + resolver: Arc<dyn SessionIdentityResolver<C>>, + ) -> Self { + self.identity = Some(IdentityConfig { + resolver, + session_key: session_key.into(), + parent: None, + child_key: None, + meta, + }); + self + } + + /// Resolves a child session while preserving its logical parent lineage. + pub fn child_identity_resolver( + mut self, + parent: SessionRef, + child_key: impl Into<String>, + meta: TranscriptMeta, + resolver: Arc<dyn SessionIdentityResolver<C>>, + ) -> Self { + let child_key = child_key.into(); + self.identity = Some(IdentityConfig { + resolver, + session_key: child_key.clone(), + parent: Some(parent), + child_key: Some(child_key), + meta, + }); + self + } + /// Freezes the prefix used to initialize this session's history. pub fn prefix(mut self, prefix: PrefixSnapshot) -> Self { self.prefix = prefix; @@ -150,16 +198,19 @@ impl<C: Clone + Send + Sync + 'static> SessionBuilder<C> { session: config.session, meta: config.meta, }); - if target.is_some() && self.codec.is_none() { + let identity = self.identity; + if (target.is_some() || identity.is_some()) && self.codec.is_none() { return Err(RuntimeError::MissingDependency("TranscriptCodec")); } let mut session = Session::<C>::new( self.driver, self.codec, self.hooks, + self.post_commit_sinks, self.prefix, self.tools, target, + identity, ); session.set_retain_recorded_tools(self.retain_recorded_tools); Ok(session) diff --git a/crates/tinyagents-runtime/src/driver.rs b/crates/tinyagents-runtime/src/driver.rs index 0d0031f57..6512e67d0 100644 --- a/crates/tinyagents-runtime/src/driver.rs +++ b/crates/tinyagents-runtime/src/driver.rs @@ -39,6 +39,8 @@ pub struct DriverOutcome { /// outcome it would derive from `interrupted`, so a run that stopped on a /// cap or a deferral is not reported as `Completed`. pub outcome: Option<tinyagents_harness::terminal::TerminalOutcome>, + /// Tool calls executed during the turn, captured before their results are folded into history. + pub tool_outcomes: Vec<tinyagents_harness::turn_outcome::ToolCallOutcome>, } /// A driver error which may retain model history and display-only partial text. @@ -150,6 +152,7 @@ pub(crate) fn outcome_from_partial( .map(TranscriptPartial::new), interrupted: partial.run.paused.is_some(), outcome: terminal.clone(), + tool_outcomes: Vec::new(), }; match partial.error { Some(error) => Err(DriverFailure { diff --git a/crates/tinyagents-runtime/src/hooks.rs b/crates/tinyagents-runtime/src/hooks.rs index 44f8a1e62..d3f46afd0 100644 --- a/crates/tinyagents-runtime/src/hooks.rs +++ b/crates/tinyagents-runtime/src/hooks.rs @@ -1,11 +1,22 @@ use async_trait::async_trait; use tinyagents_harness::terminal::TerminalOutcome; +use crate::types::PostCommitObservation; use crate::{ CommitReceipt, ResumePreparation, RuntimeError, SessionStateView, SessionTerminal, SessionTurnOutcome, SessionTurnRequest, TranscriptTurnOptions, TurnOptions, TurnPreparation, }; +/// Receives captured execution facts after a successful durable commit. +#[async_trait] +pub trait PostCommitSink<C: Clone + Send + Sync + 'static = ()>: Send + Sync { + /// Observes a committed turn. Errors cannot roll back the durable append. + async fn after_commit( + &self, + observation: &PostCommitObservation<C>, + ) -> Result<(), RuntimeError>; +} + /// Host observation/preparation around a session turn. /// /// Hooks do not grant tools, select models, compose product prompts, or own diff --git a/crates/tinyagents-runtime/src/hooks_tests.rs b/crates/tinyagents-runtime/src/hooks_tests.rs index 59610acfa..95cd4e6d5 100644 --- a/crates/tinyagents-runtime/src/hooks_tests.rs +++ b/crates/tinyagents-runtime/src/hooks_tests.rs @@ -70,6 +70,7 @@ async fn callback_adapter_implements_runtime_session_hooks() { transcript_target: None, committed_turns: 4, resumed: true, + identity: None, }; let mut request = SessionTurnRequest::new(Message::user("original")); let mut options = TurnOptions::default(); @@ -138,6 +139,7 @@ async fn callback_adapter_propagates_hook_errors_through_runtime_trait() { transcript_target: None, committed_turns: 0, resumed: false, + identity: None, }; let mut request = SessionTurnRequest::new(Message::user("x")); let mut options = TurnOptions::default(); diff --git a/crates/tinyagents-runtime/src/identity.rs b/crates/tinyagents-runtime/src/identity.rs new file mode 100644 index 000000000..8c86526e6 --- /dev/null +++ b/crates/tinyagents-runtime/src/identity.rs @@ -0,0 +1,48 @@ +use std::sync::Arc; + +use tinyagents_session::transcript::TranscriptMeta; + +use crate::RuntimeError; +use crate::types::{ResolvedSessionStores, SessionIdentity, SessionIdentityRequest}; + +/// Resolves logical session naming inputs against host-authorized storage. +/// +/// Implementations must reuse the actor-bound provider (and its pinned store +/// lease) for the lifetime of one admitted runtime session. A new admission +/// epoch must create a fresh fenced binding and runtime session; returning a +/// different tenant/backend locator during an existing admission is rejected. +/// Identity-changing preparation belongs in `before_resume`, which runs before +/// admission. `before_turn` is a trusted execution hook and must preserve the +/// admitted actor/backend authority for that turn. +pub trait SessionIdentityResolver<C: Send + Sync + 'static>: Send + Sync { + /// Admits the current host context and returns a pinned transcript store lease. + fn resolve( + &self, + request: &SessionIdentityRequest, + context: &C, + ) -> Result<ResolvedSessionStores, RuntimeError>; +} + +pub(crate) struct IdentityConfig<C: Send + Sync + 'static> { + pub(crate) resolver: Arc<dyn SessionIdentityResolver<C>>, + pub(crate) session_key: String, + pub(crate) parent: Option<tinyagents_session::transcript::SessionRef>, + pub(crate) child_key: Option<String>, + pub(crate) meta: TranscriptMeta, +} + +pub(crate) fn same_admission(a: &SessionIdentity, b: &SessionIdentity) -> bool { + if a.session.first_generation() != b.session.first_generation() || a.agent_id != b.agent_id { + return false; + } + let left = a.stores.stores().transcripts.destination_key(); + let right = b.stores.stores().transcripts.destination_key(); + match (left, right) { + (Some(left), Some(right)) => left == right, + (None, None) => Arc::ptr_eq( + &a.stores.stores().transcripts, + &b.stores.stores().transcripts, + ), + _ => false, + } +} diff --git a/crates/tinyagents-runtime/src/lib.rs b/crates/tinyagents-runtime/src/lib.rs index 85282247e..7bb9ca4c6 100644 --- a/crates/tinyagents-runtime/src/lib.rs +++ b/crates/tinyagents-runtime/src/lib.rs @@ -11,10 +11,13 @@ pub use tinyagents_hooks::command as command_hooks; mod driver; mod error; mod hooks; +mod identity; mod managed_tools; mod permanent_tool; mod prefix; mod session; +mod summary; +pub use summary::turn_summary; mod tool_registration; pub mod transcript; pub mod turn_driver; @@ -24,18 +27,22 @@ pub use builder::SessionBuilder; pub use capability_projection::{project_descriptors, register_admitted_projected}; pub use driver::{DriverFailure, DriverOutcome, DriverRequest, HarnessDriver, SessionDriver}; pub use error::RuntimeError; +pub use hooks::PostCommitSink; pub use hooks::{FnSessionHooks, NoopSessionHooks, SessionHookFuture, SessionHooks}; +pub use identity::SessionIdentityResolver; pub use managed_tools::{ManagedToolNames, refresh_prefix, reject_synthesized_collisions}; pub use permanent_tool::PermanentTool; pub use prefix::PrefixSnapshot; pub use session::Session; pub use tinyagents_session::transcript::TranscriptPartial; pub use tool_registration::{ - RegistrationPlan, ToolAdmission, ToolRegistration, ToolSurfaceDiff, plan_registration, + HostOnlyToolPolicy, NestedSpawnDepth, RegistrationPlan, ToolAdmission, ToolRegistration, + ToolSurfaceDiff, dedup_visible_tool_specs, drop_synthesized_name_collisions, plan_registration, register_admitted, tool_surface_diff, }; pub use types::{ - CommitReceipt, ResumeMode, ResumePreparation, SessionResume, SessionStateView, SessionTerminal, + CommitReceipt, PostCommitObservation, ResolvedSessionStores, ResumeMode, ResumePreparation, + SessionIdentity, SessionIdentityRequest, SessionResume, SessionStateView, SessionTerminal, SessionTurnOutcome, SessionTurnRequest, TranscriptCommitReceipt, TranscriptDelta, TranscriptTarget, TranscriptTurnOptions, TurnOptions, TurnPreparation, }; @@ -46,6 +53,9 @@ pub use types::{ /// It is object-safe so a host can keep its conversion and metadata ownership /// outside this reusable runtime. The runtime never narrows transcript data /// itself; a codec must explicitly decide how every field is represented. +/// The default reconciliation assumes decoded messages retain positional +/// correspondence with durable rows; codecs that reorder, combine, or split +/// rows must override [`Self::reconcile`] to preserve their associations. pub trait TranscriptCodec<C: Clone + Send + Sync + 'static = ()>: Send + Sync { /// Decodes a durable transcript for model execution. fn decode_history( @@ -54,12 +64,16 @@ pub trait TranscriptCodec<C: Clone + Send + Sync + 'static = ()>: Send + Sync { ) -> Result<Vec<tinyinference_llm::message::Message>, RuntimeError>; /// Reconciles a model-history transition with the prior lossless durable - /// rows. `prior` is authoritative for fields absent from inference - /// messages (provider metadata, raw arguments, reasoning, ids, and host - /// extensions): unchanged model positions must retain their corresponding - /// raw row, including when `next` is a compaction replacement. The returned - /// rows are the complete next logical durable set; the history layer writes - /// its delta atomically. `options` is captured after `before_turn` mutates + /// rows. `prior` is the lossless durable row set aligned with `previous`, + /// authoritative for fields absent from inference messages (provider + /// metadata, raw arguments, reasoning, ids, and host extensions): unchanged + /// model positions must retain their corresponding raw row, including + /// when `next` is a compaction replacement. Replay-only repairs can omit + /// rows from this correspondence; those source rows remain in the session's + /// unmodified write baseline and are preserved in a sealed generation if + /// the repaired history commits as a replacement. The returned rows are + /// the complete next logical durable set; the history layer writes its + /// delta atomically. `options` is captured after `before_turn` mutates /// the explicit turn options and before the driver consumes `RunContext`. fn reconcile( &self, diff --git a/crates/tinyagents-runtime/src/lib_tests.rs b/crates/tinyagents-runtime/src/lib_tests.rs index 02a90a328..595773769 100644 --- a/crates/tinyagents-runtime/src/lib_tests.rs +++ b/crates/tinyagents-runtime/src/lib_tests.rs @@ -24,11 +24,70 @@ use tinytools::{Tool, ToolResult, ToolSpec}; use crate::{ CommitReceipt, DriverFailure, DriverOutcome, DriverRequest, HarnessDriver, PrefixSnapshot, - ResumeMode, ResumePreparation, RuntimeError, SessionBuilder, SessionDriver, SessionHooks, - SessionStateView, SessionTerminal, SessionTurnOutcome, SessionTurnRequest, TranscriptCodec, - TranscriptTarget, TranscriptTurnOptions, TurnOptions, TurnPreparation, + ResolvedSessionStores, ResumeMode, ResumePreparation, RuntimeError, SessionBuilder, + SessionDriver, SessionHooks, SessionIdentityRequest, SessionIdentityResolver, SessionStateView, + SessionTerminal, SessionTurnOutcome, SessionTurnRequest, TranscriptCodec, TranscriptTarget, + TranscriptTurnOptions, TurnOptions, TurnPreparation, }; +struct FixedIdentityResolver { + provider: Arc<dyn tinyagents_session::port::SessionStoreProvider>, + calls: Arc<Mutex<usize>>, + agent_id: Arc<Mutex<String>>, +} + +struct FailingIdentityResolver; + +impl SessionIdentityResolver<()> for FailingIdentityResolver { + fn resolve( + &self, + _: &SessionIdentityRequest, + _: &(), + ) -> Result<ResolvedSessionStores, RuntimeError> { + Err(RuntimeError::Driver("actor admission denied".into())) + } +} + +struct SwitchableActorResolver { + provider: Arc<Mutex<Arc<dyn tinyagents_session::port::SessionStoreProvider>>>, +} + +impl SessionIdentityResolver<()> for SwitchableActorResolver { + fn resolve( + &self, + _: &SessionIdentityRequest, + _: &(), + ) -> Result<ResolvedSessionStores, RuntimeError> { + let provider = self.provider.lock().unwrap().clone(); + let stores = tinyagents_session::port::SessionStoreProvider::lease_for_agent( + provider.as_ref(), + "shared-agent", + ) + .map_err(|error| RuntimeError::InvalidSessionState(error.to_string()))?; + Ok(ResolvedSessionStores { + agent_id: "shared-agent".into(), + stores, + }) + } +} + +impl SessionIdentityResolver<()> for FixedIdentityResolver { + fn resolve( + &self, + _: &SessionIdentityRequest, + _: &(), + ) -> Result<ResolvedSessionStores, RuntimeError> { + *self.calls.lock().unwrap() += 1; + let agent_id = self.agent_id.lock().unwrap().clone(); + let stores = tinyagents_session::port::SessionStoreProvider::lease_for_agent( + self.provider.as_ref(), + &agent_id, + ) + .map_err(|error| RuntimeError::InvalidSessionState(error.to_string()))?; + Ok(ResolvedSessionStores { agent_id, stores }) + } +} + struct Driver { results: Mutex<VecDeque<Result<DriverOutcome, DriverFailure>>>, requests: Mutex<Vec<DriverRequest>>, @@ -85,6 +144,7 @@ impl SessionDriver for CancelAwareDriver { output: None, partial: None, interrupted: true, + tool_outcomes: Vec::new(), }), }) } @@ -115,6 +175,7 @@ fn outcome(history: Vec<Message>) -> DriverOutcome { output: Some("ok".into()), partial: None, interrupted: false, + tool_outcomes: Vec::new(), } } @@ -636,6 +697,7 @@ async fn persistence_failure_rolls_back_and_a_partial_never_falls_back_to_two_wr output: None, partial: Some(crate::TranscriptPartial::new("display only")), interrupted: true, + tool_outcomes: Vec::new(), }; let mut session = SessionBuilder::new(Arc::new(Driver::new(vec![Err(DriverFailure { outcome: None, @@ -753,7 +815,12 @@ async fn post_commit_errors_cannot_relabel_a_successful_turn() { #[tokio::test] async fn harness_driver_keeps_the_explicit_model_and_tool_snapshot_boundary() { let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", Arc::new(MockModel::constant("from harness"))); + harness.with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named( + "mock", + Arc::new(MockModel::constant("from harness")), + ), + )); harness.register_tool(Arc::new(RegisteredTool)); let snapshot = ToolSnapshot::new(harness.tools().declared_specs()).unwrap(); let driver = Arc::new(HarnessDriver::new(Arc::new(harness), Arc::new(()))); @@ -774,7 +841,12 @@ async fn harness_driver_keeps_the_explicit_model_and_tool_snapshot_boundary() { #[tokio::test] async fn harness_driver_rejects_a_snapshot_not_registered_by_the_harness() { let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", Arc::new(MockModel::constant("unreachable"))); + harness.with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named( + "mock", + Arc::new(MockModel::constant("unreachable")), + ), + )); let driver = Arc::new(HarnessDriver::new(Arc::new(harness), Arc::new(()))); let mut session = SessionBuilder::new(driver) .tool_snapshot(tools("not-registered")) @@ -804,6 +876,7 @@ async fn history_commits_partial_model_history_and_display_only_partial_together output: None, partial: Some(crate::TranscriptPartial::new("display partial")), interrupted: true, + tool_outcomes: Vec::new(), }; let mut session = SessionBuilder::new(Arc::new(Driver::new(vec![Err(DriverFailure { outcome: None, @@ -902,6 +975,7 @@ async fn partial_usage_error_leaves_the_session_and_target_entirely_uncommitted( output: None, partial: Some(crate::TranscriptPartial::new("display partial")), interrupted: true, + tool_outcomes: Vec::new(), }; let mut session = SessionBuilder::new(Arc::new(Driver::new(vec![Err(DriverFailure { outcome: None, @@ -2527,6 +2601,7 @@ fn session_outcome(history: Vec<Message>, output: &str) -> DriverOutcome { output: Some(output.into()), partial: None, interrupted: false, + tool_outcomes: Vec::new(), } } @@ -2701,7 +2776,7 @@ async fn a_compaction_opens_the_next_generation_and_leaves_the_sealed_one_intact #[tokio::test] async fn a_restart_after_a_compaction_resumes_the_head_generation() { let session_ref = SessionRef::scoped("thread-1", "agent-id"); - let locator = memory(); + let locator = crate::test_store::in_memory(); locator .open_session(&session_ref, meta()) .unwrap() @@ -3262,7 +3337,6 @@ async fn thread_resume_on_a_session_bound_target_does_not_corrupt_its_own_destin ], &legacy, ); - // The driver extends whatever it was handed by exactly one message. let mut session = SessionBuilder::new(Arc::new(Driver::new(vec![Ok(session_outcome( vec![ @@ -3471,7 +3545,6 @@ async fn a_first_session_resume_adopts_a_pre_identity_conversation() { &[TranscriptMessage::new("user", "plan a trip to kashmir")], &legacy, ); - let mut session = SessionBuilder::new(Arc::new(Driver::new(Vec::new()))) .codec(Arc::new(Codec::default())) .session( @@ -4504,3 +4577,536 @@ async fn dropping_the_turn_while_the_outcome_hook_is_pending_still_delivers_on_t .expect("on_terminal must still be delivered"); assert_eq!(hook.terminals.lock().unwrap().len(), 1); } + +#[tokio::test] +async fn direct_resume_resolves_and_uses_the_admitted_lease() { + let provider = crate::test_store::driver_session_stores( + tinyagents_storage_kit::Scope::new("runtime-tests").unwrap(), + ); + let calls = Arc::new(Mutex::new(0)); + let resolver = Arc::new(FixedIdentityResolver { + provider, + calls: calls.clone(), + agent_id: Arc::new(Mutex::new("agent-resolved".into())), + }); + let mut session = SessionBuilder::new(Arc::new(Driver::new(vec![]))) + .codec(Arc::new(Codec::default())) + .identity_resolver("thread-identity", meta(), resolver) + .build() + .unwrap(); + let options = TurnOptions { + resume: ResumeMode::Session, + thread_id: Some("thread-identity".into()), + ..TurnOptions::default() + }; + assert!(!session.resume(&options).await.unwrap().loaded); + assert_eq!(*calls.lock().unwrap(), 1); +} + +#[tokio::test] +async fn child_identity_resolution_preserves_parent_lineage() { + let provider = crate::test_store::driver_session_stores( + tinyagents_storage_kit::Scope::new("runtime-tests").unwrap(), + ); + let parent = SessionRef::scoped("root-thread", "parent-agent"); + let child = SessionRef::child_of(&parent, "delegate:worker"); + let resolver = Arc::new(FixedIdentityResolver { + provider: provider.clone(), + calls: Arc::new(Mutex::new(0)), + agent_id: Arc::new(Mutex::new("worker-agent".into())), + }); + let driver = Arc::new(Driver::new(vec![Ok(outcome(vec![ + Message::user("child task"), + Message::assistant("child result"), + ]))])); + let mut session = SessionBuilder::new(driver) + .codec(Arc::new(Codec::default())) + .child_identity_resolver(parent.clone(), "delegate:worker", meta(), resolver) + .build() + .unwrap(); + session + .turn( + SessionTurnRequest::new(Message::user("child task")), + TurnOptions { + resume: ResumeMode::Session, + thread_id: Some("root-thread".into()), + ..TurnOptions::default() + }, + ) + .await + .unwrap(); + let lease = tinyagents_session::port::SessionStoreProvider::lease_for_agent( + provider.as_ref(), + "worker-agent", + ) + .unwrap(); + let transcripts = &lease.stores().transcripts; + assert!(transcripts.read_session_transcript(&child).is_some()); + assert!(transcripts.read_session_transcript(&parent).is_none()); + assert!(child.is_subagent()); + assert!(session_stem(&child).starts_with(&format!("{}__", session_stem(&parent)))); +} + +struct ReplaceTargetHook(Arc<dyn TranscriptLocator>); + +#[async_trait] +impl SessionHooks for ReplaceTargetHook { + async fn before_resume( + &self, + _: &mut SessionTurnRequest, + _: &mut TurnOptions, + _: SessionStateView<'_>, + ) -> Result<ResumePreparation, RuntimeError> { + Ok(ResumePreparation { + transcript: Some(TranscriptTarget::for_session( + self.0.clone(), + SessionRef::scoped("thread-identity", "agent-resolved"), + meta(), + )), + }) + } + async fn before_turn( + &self, + _: &mut SessionTurnRequest, + _: &mut TurnOptions, + _: SessionStateView<'_>, + ) -> Result<TurnPreparation, RuntimeError> { + Ok(TurnPreparation::default()) + } + async fn before_commit( + &self, + _: &SessionTurnOutcome, + _: &TranscriptTurnOptions, + ) -> Result<(), RuntimeError> { + Ok(()) + } + async fn on_terminal(&self, _: SessionTerminal) -> Result<(), RuntimeError> { + Ok(()) + } +} + +#[tokio::test] +async fn hook_cannot_replace_resolved_identity_locator_before_any_read() { + let provider = crate::test_store::driver_session_stores( + tinyagents_storage_kit::Scope::new("runtime-tests").unwrap(), + ); + let resolver = Arc::new(FixedIdentityResolver { + provider, + calls: Arc::new(Mutex::new(0)), + agent_id: Arc::new(Mutex::new("agent-resolved".into())), + }); + let (wrong, wrong_history) = locator(None); + let driver = Arc::new(Driver::new(vec![])); + let mut session = SessionBuilder::new(driver.clone()) + .codec(Arc::new(Codec::default())) + .hooks(Arc::new(ReplaceTargetHook(wrong.clone()))) + .identity_resolver("thread-identity", meta(), resolver) + .build() + .unwrap(); + let result = session + .turn( + SessionTurnRequest::new(Message::user("next")), + TurnOptions { + resume: ResumeMode::Session, + ..TurnOptions::default() + }, + ) + .await; + assert!(matches!(result, Err(RuntimeError::InvalidSessionState(_)))); + assert!(driver.requests.lock().unwrap().is_empty()); + assert!(wrong.latest_agents.lock().unwrap().is_empty()); + assert_eq!(*wrong_history.opens.lock().unwrap(), 0); +} + +#[tokio::test] +async fn resolver_rejects_a_later_agent_change_before_another_driver_call() { + let provider = crate::test_store::driver_session_stores( + tinyagents_storage_kit::Scope::new("runtime-tests").unwrap(), + ); + let agent_id = Arc::new(Mutex::new("agent-one".to_owned())); + let resolver = Arc::new(FixedIdentityResolver { + provider, + calls: Arc::new(Mutex::new(0)), + agent_id: agent_id.clone(), + }); + let driver = Arc::new(Driver::new(vec![Ok(outcome(vec![ + Message::user("one"), + Message::assistant("done"), + ]))])); + let mut session = SessionBuilder::new(driver.clone()) + .codec(Arc::new(Codec::default())) + .identity_resolver("thread-identity", meta(), resolver) + .build() + .unwrap(); + let options = || TurnOptions { + resume: ResumeMode::Session, + thread_id: Some("thread-identity".into()), + ..TurnOptions::default() + }; + session + .turn(SessionTurnRequest::new(Message::user("one")), options()) + .await + .unwrap(); + *agent_id.lock().unwrap() = "agent-two".into(); + assert!(matches!( + session + .turn(SessionTurnRequest::new(Message::user("two")), options()) + .await, + Err(RuntimeError::InvalidSessionState(_)), + )); + assert_eq!(driver.requests.lock().unwrap().len(), 1); +} + +#[tokio::test] +async fn same_admission_can_resume_the_current_generation_twice() { + let provider = crate::test_store::driver_session_stores( + tinyagents_storage_kit::Scope::new("runtime-tests").unwrap(), + ); + let resolver = Arc::new(FixedIdentityResolver { + provider, + calls: Arc::new(Mutex::new(0)), + agent_id: Arc::new(Mutex::new("agent-resolved".into())), + }); + let driver = Arc::new(Driver::new(vec![ + Ok(outcome(vec![ + Message::user("one"), + Message::assistant("done"), + ])), + Ok(outcome(vec![ + Message::user("two"), + Message::assistant("done"), + ])), + ])); + let mut session = SessionBuilder::new(driver.clone()) + .codec(Arc::new(Codec::default())) + .identity_resolver("thread-identity", meta(), resolver) + .build() + .unwrap(); + let options = || TurnOptions { + resume: ResumeMode::Session, + thread_id: Some("thread-identity".into()), + ..TurnOptions::default() + }; + session + .turn(SessionTurnRequest::new(Message::user("one")), options()) + .await + .unwrap(); + session + .turn(SessionTurnRequest::new(Message::user("two")), options()) + .await + .unwrap(); + assert_eq!(driver.requests.lock().unwrap().len(), 2); +} + +#[tokio::test] +async fn resolver_error_does_not_touch_builder_fallback_locator() { + let (fallback, history) = locator(None); + let mut session = SessionBuilder::new(Arc::new(Driver::new(vec![]))) + .codec(Arc::new(Codec::default())) + .transcript(fallback.clone(), "fallback", meta()) + .identity_resolver("thread-identity", meta(), Arc::new(FailingIdentityResolver)) + .build() + .unwrap(); + let options = TurnOptions { + resume: ResumeMode::LatestForAgent, + ..TurnOptions::default() + }; + assert!(matches!( + session.resume(&options).await, + Err(RuntimeError::Driver(_)) + )); + assert!(fallback.latest_agents.lock().unwrap().is_empty()); + assert_eq!(*history.opens.lock().unwrap(), 0); +} + +#[tokio::test] +async fn first_admission_ignores_same_identity_fallback_at_generation_nine() { + let provider = crate::test_store::driver_session_stores( + tinyagents_storage_kit::Scope::new("runtime-tests").unwrap(), + ); + let mut fallback_session = SessionRef::scoped("thread-first-admission", "agent-resolved"); + for _ in 0..9 { + fallback_session = fallback_session.next_generation(); + } + let (fallback, fallback_history) = locator(None); + let resolver = Arc::new(FixedIdentityResolver { + provider: provider.clone(), + calls: Arc::new(Mutex::new(0)), + agent_id: Arc::new(Mutex::new("agent-resolved".into())), + }); + let driver = Arc::new(Driver::new(vec![Ok(outcome(vec![ + Message::user("first"), + Message::assistant("admitted root"), + ]))])); + let mut session = SessionBuilder::new(driver) + .codec(Arc::new(Codec::default())) + .session(fallback.clone(), fallback_session.clone(), meta()) + .identity_resolver("thread-first-admission", meta(), resolver) + .build() + .unwrap(); + session + .turn( + SessionTurnRequest::new(Message::user("first")), + TurnOptions { + resume: ResumeMode::Never, + thread_id: Some("thread-first-admission".into()), + ..TurnOptions::default() + }, + ) + .await + .unwrap(); + let lease = tinyagents_session::port::SessionStoreProvider::lease_for_agent( + provider.as_ref(), + "agent-resolved", + ) + .unwrap(); + let transcripts = &lease.stores().transcripts; + let admitted_root = SessionRef::scoped("thread-first-admission", "agent-resolved"); + assert!( + transcripts + .read_session_transcript(&admitted_root) + .is_some() + ); + assert!( + transcripts + .read_session_transcript(&fallback_session) + .is_none() + ); + assert_eq!(*fallback_history.opens.lock().unwrap(), 0); + assert!(fallback.latest_agents.lock().unwrap().is_empty()); +} + +#[tokio::test] +async fn explicit_session_override_cannot_change_resolved_identity() { + let provider = crate::test_store::driver_session_stores( + tinyagents_storage_kit::Scope::new("runtime-tests").unwrap(), + ); + let calls = Arc::new(Mutex::new(0)); + let resolver = Arc::new(FixedIdentityResolver { + provider, + calls: calls.clone(), + agent_id: Arc::new(Mutex::new("agent-resolved".into())), + }); + let mut session = SessionBuilder::new(Arc::new(Driver::new(vec![]))) + .codec(Arc::new(Codec::default())) + .identity_resolver("thread-identity", meta(), resolver) + .build() + .unwrap(); + let options = TurnOptions { + resume: ResumeMode::Session, + session: Some(SessionRef::scoped("other-thread", "agent-resolved")), + ..TurnOptions::default() + }; + assert!(matches!( + session.resume(&options).await, + Err(RuntimeError::InvalidSessionState(_)) + )); + assert_eq!(*calls.lock().unwrap(), 1); +} + +#[tokio::test] +async fn admitted_identity_keeps_compacted_head_on_following_never_resume_turn() { + let provider = Arc::new(tinyagents_session::port::InMemorySessionStores::new()); + let session_ref = SessionRef::scoped("thread-identity-compact", "agent-resolved"); + let resolver = Arc::new(FixedIdentityResolver { + provider: provider.clone(), + calls: Arc::new(Mutex::new(0)), + agent_id: Arc::new(Mutex::new("agent-resolved".into())), + }); + let driver = Arc::new(Driver::new(vec![ + Ok(outcome(vec![ + Message::user("one"), + Message::assistant("first"), + Message::user("two"), + Message::assistant("second"), + ])), + Ok(outcome(vec![ + Message::user("three"), + Message::assistant("third"), + ])), + Ok(outcome(vec![ + Message::user("three"), + Message::assistant("third"), + Message::user("four"), + Message::assistant("fourth"), + ])), + ])); + let mut session = SessionBuilder::new(driver) + .codec(Arc::new(Codec::default())) + .identity_resolver("thread-identity-compact", meta(), resolver) + .build() + .unwrap(); + session + .turn( + SessionTurnRequest::new(Message::user("one")), + session_turn_options(ResumeMode::Session, "thread-identity-compact"), + ) + .await + .unwrap(); + let lease = tinyagents_session::port::SessionStoreProvider::lease_for_agent( + provider.as_ref(), + "agent-resolved", + ) + .unwrap(); + let locator = &lease.stores().transcripts; + session + .turn( + SessionTurnRequest::new(Message::user("three")), + session_turn_options(ResumeMode::Never, "thread-identity-compact"), + ) + .await + .unwrap(); + let sealed = locator.read_session_transcript(&session_ref).unwrap(); + let sealed_bytes = fingerprint(sealed.as_ref()); + assert_eq!(locator.head_generation(&session_ref).generation, 1); + session + .turn( + SessionTurnRequest::new(Message::user("four")), + session_turn_options(ResumeMode::Never, "thread-identity-compact"), + ) + .await + .unwrap(); + assert_eq!( + fingerprint(sealed.as_ref()), + sealed_bytes, + "re-admission must not reset or rewrite the sealed generation" + ); + assert_eq!(locator.head_generation(&session_ref).generation, 1); + let active = locator + .read_session_transcript(&session_ref.next_generation()) + .unwrap() + .read_session() + .unwrap() + .unwrap(); + assert_eq!( + active + .messages + .iter() + .map(|row| row.content.as_str()) + .collect::<Vec<_>>(), + ["three", "third", "four", "fourth"] + ); +} + +#[tokio::test] +async fn same_agent_backend_change_is_rejected_until_a_fresh_session_admits_it() { + use tinyagents_session::port::SessionStoreProvider; + + let backend: Arc<dyn tinystoragedrivers_api::StorageBackend> = + Arc::new(tinystoragedrivers_testkit::MemoryBackend::new()); + let alice_root = crate::test_store::driver_session_stores_over( + backend.clone(), + tinyagents_storage_kit::Scope::new("tenant:alice").unwrap(), + ); + let bob_root = crate::test_store::driver_session_stores_over( + backend, + tinyagents_storage_kit::Scope::new("tenant:bob").unwrap(), + ); + let alice: Arc<dyn SessionStoreProvider> = alice_root; + let bob: Arc<dyn SessionStoreProvider> = bob_root; + let provider = Arc::new(Mutex::new(alice)); + let resolver = Arc::new(SwitchableActorResolver { + provider: provider.clone(), + }); + let first_driver = Arc::new(Driver::new(vec![Ok(outcome(vec![ + Message::user("first"), + Message::assistant("alice"), + ]))])); + let options = || TurnOptions { + resume: ResumeMode::Session, + thread_id: Some("tenant-thread".into()), + ..TurnOptions::default() + }; + let mut established = SessionBuilder::new(first_driver.clone()) + .codec(Arc::new(Codec::default())) + .identity_resolver("tenant-thread", meta(), resolver.clone()) + .build() + .unwrap(); + established + .turn(SessionTurnRequest::new(Message::user("first")), options()) + .await + .unwrap(); + *provider.lock().unwrap() = bob; + assert!(matches!( + established + .turn(SessionTurnRequest::new(Message::user("second")), options()) + .await, + Err(RuntimeError::InvalidSessionState(_)), + )); + assert_eq!(first_driver.requests.lock().unwrap().len(), 1); + + let fresh_driver = Arc::new(Driver::new(vec![Ok(outcome(vec![ + Message::user("fresh"), + Message::assistant("bob"), + ]))])); + let mut fresh = SessionBuilder::new(fresh_driver.clone()) + .codec(Arc::new(Codec::default())) + .identity_resolver("tenant-thread", meta(), resolver) + .build() + .unwrap(); + fresh + .turn(SessionTurnRequest::new(Message::user("fresh")), options()) + .await + .unwrap(); + assert_eq!(fresh_driver.requests.lock().unwrap().len(), 1); +} + +#[tokio::test] +async fn direct_resume_reads_admitted_transcript_and_rejects_later_actor_change() { + use tinyagents_session::port::SessionStoreProvider; + + let backend: Arc<dyn tinystoragedrivers_api::StorageBackend> = + Arc::new(tinystoragedrivers_testkit::MemoryBackend::new()); + let alice_root = crate::test_store::driver_session_stores_over( + backend.clone(), + tinyagents_storage_kit::Scope::new("tenant:alice").unwrap(), + ); + let bob_root = crate::test_store::driver_session_stores_over( + backend, + tinyagents_storage_kit::Scope::new("tenant:bob").unwrap(), + ); + let alice: Arc<dyn SessionStoreProvider> = alice_root; + let bob: Arc<dyn SessionStoreProvider> = bob_root; + let provider = Arc::new(Mutex::new(alice)); + let resolver = Arc::new(SwitchableActorResolver { + provider: provider.clone(), + }); + let driver = Arc::new(Driver::new(vec![Ok(outcome(vec![ + Message::user("persisted"), + Message::assistant("admitted transcript"), + ]))])); + let mut session = SessionBuilder::new(driver.clone()) + .codec(Arc::new(Codec::default())) + .identity_resolver("direct-resume-actor", meta(), resolver) + .build() + .unwrap(); + let options = TurnOptions { + resume: ResumeMode::Session, + thread_id: Some("direct-resume-actor".into()), + ..TurnOptions::default() + }; + session + .turn( + SessionTurnRequest::new(Message::user("persisted")), + TurnOptions { + resume: ResumeMode::Never, + thread_id: Some("direct-resume-actor".into()), + ..TurnOptions::default() + }, + ) + .await + .unwrap(); + assert!( + session.resume(&options).await.unwrap().loaded, + "direct resume should read the transcript from the admitted actor lease" + ); + *provider.lock().unwrap() = bob; + assert!( + matches!( + session.resume(&options).await, + Err(RuntimeError::InvalidSessionState(_)) + ), + "direct resume must reject a later actor/backend binding before it reads another history" + ); + assert_eq!(driver.requests.lock().unwrap().len(), 1); +} diff --git a/crates/tinyagents-runtime/src/session.rs b/crates/tinyagents-runtime/src/session.rs index 38f47a782..0c820bacd 100644 --- a/crates/tinyagents-runtime/src/session.rs +++ b/crates/tinyagents-runtime/src/session.rs @@ -9,6 +9,7 @@ use tinyagents_session::transcript::{ }; use tinyinference_llm::message::Message; +use crate::identity::{IdentityConfig, same_admission}; use crate::{ CommitReceipt, DriverRequest, PrefixSnapshot, ResumeMode, ResumePreparation, RuntimeError, SessionDriver, SessionHooks, SessionResume, SessionStateView, SessionTerminal, @@ -25,10 +26,20 @@ pub struct Session<C: Clone + Send + Sync + 'static = ()> { driver: Arc<dyn SessionDriver<C>>, codec: Option<Arc<dyn TranscriptCodec<C>>>, hooks: Arc<dyn SessionHooks<C>>, + post_commit_sinks: Vec<Arc<dyn crate::hooks::PostCommitSink<C>>>, + identity_config: Option<IdentityConfig<C>>, + established_identity: Option<crate::SessionIdentity>, + identity: Option<crate::SessionIdentity>, prefix: PrefixSnapshot, default_tools: ToolSnapshot, history: Vec<Message>, persisted: Vec<TranscriptMessage>, + /// Lossless rows aligned with the replay history passed to the codec. + /// + /// This can be shorter than persisted after orphan-head repair. The + /// original rows remain the write baseline so a successful turn rolls + /// forward into a successor generation without erasing its source. + reconcile_prior: Vec<TranscriptMessage>, target: Option<TranscriptTarget>, transcript: Option<Arc<dyn TranscriptHistory>>, committed_turns: usize, @@ -50,18 +61,25 @@ impl<C: Clone + Send + Sync + 'static> Session<C> { driver: Arc<dyn SessionDriver<C>>, codec: Option<Arc<dyn TranscriptCodec<C>>>, hooks: Arc<dyn SessionHooks<C>>, + post_commit_sinks: Vec<Arc<dyn crate::hooks::PostCommitSink<C>>>, prefix: PrefixSnapshot, default_tools: ToolSnapshot, target: Option<TranscriptTarget>, + identity_config: Option<IdentityConfig<C>>, ) -> Self { Self { driver, codec, hooks, + post_commit_sinks, + identity_config, + established_identity: None, + identity: None, history: prefix.messages().to_vec(), prefix: prefix.frozen(), default_tools, persisted: Vec::new(), + reconcile_prior: Vec::new(), target, transcript: None, committed_turns: 0, @@ -113,7 +131,88 @@ impl<C: Clone + Send + Sync + 'static> Session<C> { )); } self.history = self.with_prefix(history); - self.persisted = raw; + self.persisted = raw.clone(); + self.reconcile_prior = raw; + Ok(()) + } + + fn resolve_and_bind( + &mut self, + options: &TurnOptions<C>, + hook_target: Option<&TranscriptTarget>, + ) -> Result<(), RuntimeError> { + let Some(config) = self.identity_config.as_ref() else { + return Ok(()); + }; + self.identity = None; + let request = crate::SessionIdentityRequest { + thread_id: options.thread_id.clone(), + session_key: config.session_key.clone(), + parent: config.parent.clone(), + child_key: config.child_key.clone(), + }; + let resolved = config + .resolver + .resolve(&request, &options.run_context.data)?; + if resolved.agent_id.is_empty() { + return Err(RuntimeError::InvalidSessionState( + "resolver returned empty agent id".into(), + )); + } + let first_session = match (&request.parent, &request.child_key) { + (Some(parent), Some(child_key)) => SessionRef::child_of(parent, child_key), + (None, None) => SessionRef::scoped(&request.session_key, &resolved.agent_id), + _ => { + return Err(RuntimeError::InvalidSessionState( + "identity request has an incomplete child lineage".into(), + )); + } + }; + let identity = crate::SessionIdentity { + session: first_session.clone(), + agent_id: resolved.agent_id, + stores: resolved.stores, + }; + if options + .session + .as_ref() + .is_some_and(|selected| selected.first_generation() != first_session.first_generation()) + { + return Err(RuntimeError::InvalidSessionState( + "TurnOptions.session conflicts with resolved identity".into(), + )); + } + let session = if let Some(previous) = self.established_identity.as_ref() { + if !same_admission(previous, &identity) { + return Err(RuntimeError::InvalidSessionState( + "session identity or pinned store changed".into(), + )); + } + self.target + .as_ref() + .and_then(|target| target.session.clone()) + .filter(|current| current.first_generation() == first_session.first_generation()) + .unwrap_or(first_session) + } else { + first_session + }; + let mut meta = config.meta.clone(); + meta.agent_id = Some(identity.agent_id.clone()); + let target = TranscriptTarget::for_session( + identity.stores.stores().transcripts.clone(), + session.clone(), + meta, + ); + if hook_target.is_some_and(|selected| !selected.same_binding(&target)) { + return Err(RuntimeError::InvalidSessionState( + "before_resume target conflicts with resolved identity".into(), + )); + } + self.target = Some(target); + let mut identity = identity; + identity.session = session; + self.established_identity = Some(identity.clone()); + self.identity = Some(identity); Ok(()) } @@ -126,6 +225,14 @@ impl<C: Clone + Send + Sync + 'static> Session<C> { if options.cancellation.is_cancelled() { return Err(RuntimeError::Cancelled); } + self.resolve_and_bind(options, None)?; + self.resume_admitted(options).await + } + + async fn resume_admitted( + &mut self, + options: &TurnOptions<C>, + ) -> Result<SessionResume, RuntimeError> { let Some(target) = self.target.as_ref() else { return Ok(SessionResume { loaded: false, @@ -204,6 +311,7 @@ impl<C: Clone + Send + Sync + 'static> Session<C> { history: self.history.clone(), }); }; + let scanned_session_id = transcript.meta.session_id.clone(); if let Some(selected) = session_binding.as_ref() { let expected = selected.session_id(); if transcript @@ -221,7 +329,36 @@ impl<C: Clone + Send + Sync + 'static> Session<C> { .codec .as_ref() .ok_or(RuntimeError::MissingDependency("TranscriptCodec"))?; - let mut decoded = codec.decode_history(&transcript)?; + let replay_rows = + tinyagents_session::transcript::repair_orphaned_tool_head(&transcript.messages); + let mut reconcile_prior = replay_rows.as_ref().to_vec(); + let mut decoded = match replay_rows { + std::borrow::Cow::Borrowed(_) => codec.decode_history(&transcript)?, + std::borrow::Cow::Owned(messages) => { + let mut replay_transcript = transcript.clone(); + replay_transcript.messages = messages; + codec.decode_history(&replay_transcript)? + } + }; + // Keep the lossless prior rows aligned with the decoded replay view. + // Some legacy codecs recognize tool rows whose raw role is not + // `tool`; when decoded repair removes such a leading tool message, its + // corresponding raw row must not remain at that positional index for + // the default reconciler to reuse for the following user message. + let decoded_system_prefix = decoded + .iter() + .take_while(|message| matches!(message, Message::System(_))) + .count(); + let decoded_orphan_count = decoded[decoded_system_prefix..] + .iter() + .take_while(|message| matches!(message, Message::Tool(_))) + .count(); + if decoded_orphan_count != 0 && reconcile_prior.len() == decoded.len() { + reconcile_prior + .drain(decoded_system_prefix..decoded_system_prefix + decoded_orphan_count); + } + decoded = tinyagents_session::transcript::repair_orphaned_tool_head_messages(&decoded) + .into_owned(); // A compacted head can start with a System summary immediately after // the original frozen prompt. Its role does not make it prefix // material. Determine how many *stored* rows to strip independently @@ -379,6 +516,7 @@ impl<C: Clone + Send + Sync + 'static> Session<C> { .map_or(0, |tools| tools.specs().len()) ); self.persisted = transcript.messages; + self.reconcile_prior = reconcile_prior; // The discovered metadata, not the builder seed, is authoritative for // the subsequent append. This keeps resume-only host fields intact. if let Some(target) = self.target.as_mut() { @@ -456,7 +594,25 @@ impl<C: Clone + Send + Sync + 'static> Session<C> { self.recorded_tools = Self::decode_recorded_tools(destination_transcript.tools.as_ref()); self.persisted_prefix_len = destination_transcript.meta.prefix_message_count; - self.persisted = destination_transcript.messages; + self.persisted = destination_transcript.messages.clone(); + let destination = target + .session + .as_ref() + .map(|session| target.locator.head_generation(session)) + .and_then(|head| target.locator.read_session_transcript(&head)); + let same_source = destination_transcript + .meta + .session_id + .as_deref() + .is_some_and(|session_id| { + scanned_session_id.as_deref() == Some(session_id) + }) + || destination + .as_ref() + .is_some_and(|destination| destination.path() == read.path()); + if !same_source { + self.reconcile_prior = destination_transcript.messages.clone(); + } if let Some(target) = self.target.as_mut() { target.meta = destination_transcript.meta; } @@ -472,6 +628,7 @@ impl<C: Clone + Send + Sync + 'static> Session<C> { self.recorded_tools = None; self.persisted_prefix_len = None; self.persisted = Vec::new(); + self.reconcile_prior = Vec::new(); if let Some(target) = self.target.as_mut() { target.meta = pre_scan_meta; if let Some(session) = target.session.clone() { @@ -533,6 +690,7 @@ impl<C: Clone + Send + Sync + 'static> Session<C> { options: &mut TurnOptions<C>, terminal_guard: &mut TerminalGuard<C>, ) -> Result<SessionTurnOutcome, RuntimeError> { + let original_input = request.input.clone(); if options.cancellation.is_cancelled() { return Err(RuntimeError::Cancelled); } @@ -543,7 +701,9 @@ impl<C: Clone + Send + Sync + 'static> Session<C> { .before_resume(request, options, self.state_view(false)), ) .await?; + let hook_target = resume_preparation.transcript.clone(); self.apply_resume_preparation(resume_preparation)?; + self.resolve_and_bind(options, hook_target.as_ref())?; // Resume preparation may install a lazy target or change the selected // session, so choose the shared locks only after applying it. The lock // still covers the complete resume read and subsequent write. @@ -552,7 +712,7 @@ impl<C: Clone + Send + Sync + 'static> Session<C> { let resumed = if options.resume == ResumeMode::Never { false } else { - self.resume(options).await?.loaded + self.resume_admitted(options).await?.loaded }; // `resume` is synchronous after its read, so this explicit boundary // makes cancellation between loading and before-turn preparation @@ -662,7 +822,8 @@ impl<C: Clone + Send + Sync + 'static> Session<C> { self.remember_sent_tools(record_tools.as_ref()); self.prefix = prefix; self.history = partial_history; - self.persisted = raw; + self.persisted = raw.clone(); + self.reconcile_prior = raw; if receipt.is_some() { self.committed_turns += 1; } @@ -704,7 +865,8 @@ impl<C: Clone + Send + Sync + 'static> Session<C> { self.remember_sent_tools(record_tools.as_ref()); self.prefix = prefix; self.history = committed.history.clone(); - self.persisted = raw; + self.persisted = raw.clone(); + self.reconcile_prior = raw; self.committed_turns += 1; // Post-commit hooks may deliver another background message into this // session. Do not hold the non-reentrant turn lock while they run. @@ -718,7 +880,19 @@ impl<C: Clone + Send + Sync + 'static> Session<C> { options: codec_options, transcript, }; - let finalization = terminal_guard.finalize_commit(receipt); + let observation = crate::PostCommitObservation { + receipt: receipt.clone(), + input: original_input, + usage: turn_usage.clone(), + tool_outcomes: outcome.tool_outcomes.clone(), + thread_id, + agent_id: self + .identity + .as_ref() + .map(|identity| identity.agent_id().to_owned()), + }; + let finalization = + terminal_guard.finalize_commit(receipt, observation, self.post_commit_sinks.clone()); // This await deliberately does not observe cancellation. If this turn // future is dropped, dropping `JoinHandle` detaches rather than aborts // the owned finalization task. @@ -871,6 +1045,7 @@ impl<C: Clone + Send + Sync + 'static> Session<C> { transcript_target: self.target.as_ref(), committed_turns: self.committed_turns, resumed, + identity: self.established_identity.as_ref(), } } @@ -881,7 +1056,7 @@ impl<C: Clone + Send + Sync + 'static> Session<C> { options: &TranscriptTurnOptions<C>, ) -> Result<Vec<TranscriptMessage>, RuntimeError> { match &self.codec { - Some(codec) => codec.reconcile(&self.persisted, previous, next, options), + Some(codec) => codec.reconcile(&self.reconcile_prior, previous, next, options), None => Ok(Vec::new()), } } @@ -946,6 +1121,7 @@ impl<C: Clone + Send + Sync + 'static> Session<C> { .messages() .map_err(|error| RuntimeError::Persistence(error.to_string()))?; self.transcript = Some(handle); + self.reconcile_prior = persisted.clone(); self.persisted = persisted; } @@ -1147,7 +1323,12 @@ impl<C: Clone + Send + Sync + 'static> TerminalGuard<C> { Some((terminal, outcome)) } - fn finalize_commit(&mut self, receipt: CommitReceipt<C>) -> tokio::task::JoinHandle<()> { + fn finalize_commit( + &mut self, + receipt: CommitReceipt<C>, + observation: crate::PostCommitObservation<C>, + sinks: Vec<Arc<dyn crate::hooks::PostCommitSink<C>>>, + ) -> tokio::task::JoinHandle<()> { let terminal = SessionTerminal::Completed(receipt.outcome.clone()); let outcome = self .driver_success_outcome @@ -1159,7 +1340,14 @@ impl<C: Clone + Send + Sync + 'static> TerminalGuard<C> { self.committed = true; let hooks = self.hooks.clone(); tokio::spawn(async move { - let _ = hooks.after_commit(receipt).await; + if let Err(error) = hooks.after_commit(receipt).await { + tracing::warn!("[session] after_commit hook failed: {error}"); + } + for sink in sinks { + if let Err(error) = sink.after_commit(&observation).await { + tracing::warn!("[session] post-commit sink failed: {error}"); + } + } let _ = hooks.on_terminal_outcome(outcome).await; let _ = hooks.on_terminal(terminal).await; }) diff --git a/crates/tinyagents-runtime/src/session_tests.rs b/crates/tinyagents-runtime/src/session_tests.rs index e4afaebe4..ffada9309 100644 --- a/crates/tinyagents-runtime/src/session_tests.rs +++ b/crates/tinyagents-runtime/src/session_tests.rs @@ -1,16 +1,805 @@ use super::*; -use crate::test_store::{memory, read, seed}; +use crate::test_store::{fingerprint, memory, read, seed}; use crate::{DriverFailure, DriverOutcome, SessionBuilder, TranscriptCodec}; use async_trait::async_trait; -use std::sync::Arc; +use std::sync::{ + Arc, Mutex, + atomic::{AtomicUsize, Ordering}, +}; use tinyagents_session::transcript::{ - BackgroundAppend, BackgroundAppendOutcome, SessionRef, SessionTranscript, TranscriptMessage, - TranscriptMeta, session_stem, + BackgroundAppend, BackgroundAppendOutcome, SessionRef, SessionTranscript, TranscriptLocator, + TranscriptMessage, TranscriptMeta, TranscriptTurn, session_stem, }; use tinyinference_llm::message::Message; use tokio::sync::Notify; use tokio::time::{Duration, timeout}; +struct ReplayCaptureDriver(Arc<Mutex<Vec<Message>>>); + +#[async_trait] +impl SessionDriver for ReplayCaptureDriver { + async fn execute(&self, request: DriverRequest) -> Result<DriverOutcome, DriverFailure> { + *self.0.lock().unwrap() = request.history; + Err(DriverFailure { + error: RuntimeError::Driver("stop after replay capture".into()), + partial: None, + outcome: None, + }) + } +} + +struct AppendReplyDriver; + +#[async_trait] +impl SessionDriver for AppendReplyDriver { + async fn execute(&self, request: DriverRequest) -> Result<DriverOutcome, DriverFailure> { + let mut history = request.history; + history.push(Message::assistant("reply after replay")); + Ok(DriverOutcome { + history, + output: Some("reply after replay".into()), + partial: None, + interrupted: false, + outcome: None, + tool_outcomes: Vec::new(), + }) + } +} + +struct NativeReplayCodec; + +impl TranscriptCodec for NativeReplayCodec { + fn decode_history(&self, transcript: &SessionTranscript) -> Result<Vec<Message>, RuntimeError> { + crate::transcript::native_transcript_to_messages(&transcript.messages) + } +} + +struct ReplayProbeCodec { + seen_rows: Arc<Mutex<Vec<TranscriptMessage>>>, + synthesize_leading_tool: bool, +} + +impl TranscriptCodec for ReplayProbeCodec { + fn decode_history(&self, transcript: &SessionTranscript) -> Result<Vec<Message>, RuntimeError> { + *self.seen_rows.lock().unwrap() = transcript.messages.clone(); + let mut decoded = transcript + .messages + .iter() + .map(|row| match row.role.as_str() { + "system" => Message::system(&row.content), + "user" => Message::user(&row.content), + "tool" => Message::tool(row.id.as_deref().unwrap_or("missing-id"), &row.content), + _ => Message::assistant(&row.content), + }) + .collect::<Vec<_>>(); + if self.synthesize_leading_tool { + let prefix = decoded + .iter() + .take_while(|message| matches!(message, Message::System(_))) + .count(); + decoded.insert(prefix, Message::tool("codec-orphan", "decoded orphan")); + } + Ok(decoded) + } + + fn reconcile( + &self, + prior: &[TranscriptMessage], + _: &[Message], + _: &[Message], + _: &crate::TranscriptTurnOptions, + ) -> Result<Vec<TranscriptMessage>, RuntimeError> { + Ok(prior.to_vec()) + } +} + +fn seed_replay_transcript( + locator: &dyn TranscriptLocator, + session: &SessionRef, + messages: &[TranscriptMessage], +) { + let mut transcript_meta = meta(); + transcript_meta.session_id = Some(session.session_id()); + transcript_meta.agent_id = session.agent_id.clone(); + transcript_meta.thread_id = Some(session.session_key.clone()); + let handle = locator + .open_session(session, transcript_meta.clone()) + .unwrap(); + handle + .append_turn(TranscriptTurn { + prev: &[], + next: messages, + meta: &transcript_meta, + turn_usage: None, + request_id: None, + tools: None, + }) + .unwrap(); +} + +#[tokio::test] +async fn resume_repairs_only_the_leading_orphan_tool_head_for_dispatch() { + let locator = memory(); + let session_ref = SessionRef::scoped("thread-orphan-runtime", "agent"); + let rows = vec![ + TranscriptMessage::system("frozen prompt"), + TranscriptMessage::tool("orphan result").with_id("orphan-call"), + TranscriptMessage::user("keep this user turn"), + TranscriptMessage::tool("later malformed tool").with_id("later-call"), + TranscriptMessage::assistant("trailing assistant remains"), + ]; + seed_replay_transcript(locator.as_ref(), &session_ref, &rows); + let before = fingerprint( + locator + .read_session_transcript(&session_ref) + .unwrap() + .as_ref(), + ); + let seen_rows = Arc::new(Mutex::new(Vec::new())); + let dispatched = Arc::new(Mutex::new(Vec::new())); + let mut session = SessionBuilder::new(Arc::new(ReplayCaptureDriver(dispatched.clone()))) + .codec(Arc::new(ReplayProbeCodec { + seen_rows: seen_rows.clone(), + synthesize_leading_tool: false, + })) + .session(locator.clone(), session_ref.clone(), meta()) + .build() + .unwrap(); + + assert!( + session + .turn( + SessionTurnRequest::new(Message::user("next")), + TurnOptions { + resume: ResumeMode::Session, + thread_id: Some(session_ref.session_key.clone()), + ..TurnOptions::default() + }, + ) + .await + .is_err() + ); + + let replay_rows = seen_rows.lock().unwrap().clone(); + assert_eq!( + replay_rows + .iter() + .map(|row| row.content.as_str()) + .collect::<Vec<_>>(), + [ + "frozen prompt", + "keep this user turn", + "later malformed tool", + "trailing assistant remains" + ], + "the codec should receive the replay-only raw-row view" + ); + let replay = dispatched.lock().unwrap().clone(); + assert_eq!( + replay.iter().map(Message::text).collect::<Vec<_>>(), + [ + "frozen prompt", + "keep this user turn", + "later malformed tool", + "trailing assistant remains", + "next" + ], + "later malformed tool rows and trailing assistant rows remain visible to dispatch" + ); + let after = fingerprint( + locator + .read_session_transcript(&session_ref) + .unwrap() + .as_ref(), + ); + assert_eq!( + after, before, + "replay repair must never rewrite durable transcript rows" + ); +} + +#[tokio::test] +async fn resume_repairs_a_leading_tool_synthesized_by_the_decoder() { + let locator = memory(); + let session_ref = SessionRef::scoped("thread-decoded-orphan-runtime", "agent"); + let rows = vec![ + TranscriptMessage::system("frozen prompt"), + TranscriptMessage::user("durable user row"), + ]; + seed_replay_transcript(locator.as_ref(), &session_ref, &rows); + let before = fingerprint( + locator + .read_session_transcript(&session_ref) + .unwrap() + .as_ref(), + ); + let seen_rows = Arc::new(Mutex::new(Vec::new())); + let dispatched = Arc::new(Mutex::new(Vec::new())); + let mut session = SessionBuilder::new(Arc::new(ReplayCaptureDriver(dispatched.clone()))) + .codec(Arc::new(ReplayProbeCodec { + seen_rows, + synthesize_leading_tool: true, + })) + .session(locator.clone(), session_ref.clone(), meta()) + .build() + .unwrap(); + + assert!( + session + .turn( + SessionTurnRequest::new(Message::user("next")), + TurnOptions { + resume: ResumeMode::Session, + thread_id: Some(session_ref.session_key.clone()), + ..TurnOptions::default() + }, + ) + .await + .is_err() + ); + assert_eq!( + dispatched + .lock() + .unwrap() + .iter() + .map(Message::text) + .collect::<Vec<_>>(), + ["frozen prompt", "durable user row", "next"], + "decoded history must be repaired before dispatch" + ); + let after = fingerprint( + locator + .read_session_transcript(&session_ref) + .unwrap() + .as_ref(), + ); + assert_eq!( + after, before, + "decoded repair must leave durable rows unchanged" + ); +} + +#[tokio::test] +async fn successful_turn_after_orphan_replay_preserves_user_row_and_original_generation() { + let locator = memory(); + let session_ref = SessionRef::scoped("thread-orphan-append", "agent"); + let mut system = TranscriptMessage::system("frozen prompt"); + system.extra_metadata = Some(serde_json::json!({"host": "system metadata"})); + let mut orphan = TranscriptMessage::tool("orphan result").with_id("orphan-call"); + orphan.extra_metadata = Some(serde_json::json!({"host": "retained only in sealed source"})); + let mut user = TranscriptMessage::user("persisted user"); + user.extra_metadata = Some(serde_json::json!({"source": "original user row"})); + let rows = vec![system, orphan, user]; + seed_replay_transcript(locator.as_ref(), &session_ref, &rows); + let original = locator + .read_session_transcript(&session_ref) + .unwrap() + .read_session() + .unwrap() + .unwrap(); + let original_fingerprint = fingerprint( + locator + .read_session_transcript(&session_ref) + .unwrap() + .as_ref(), + ); + let mut session = SessionBuilder::new(Arc::new(AppendReplyDriver)) + .codec(Arc::new(NativeReplayCodec)) + .session(locator.clone(), session_ref.clone(), meta()) + .build() + .unwrap(); + + session + .turn( + SessionTurnRequest::new(Message::user("new user turn")), + TurnOptions { + resume: ResumeMode::Session, + thread_id: Some(session_ref.session_key.clone()), + ..TurnOptions::default() + }, + ) + .await + .unwrap(); + + let old_generation = locator + .read_session_transcript(&session_ref) + .unwrap() + .read_session() + .unwrap() + .unwrap(); + assert_eq!(old_generation.messages, original.messages); + assert_eq!( + fingerprint( + locator + .read_session_transcript(&session_ref) + .unwrap() + .as_ref() + ), + original_fingerprint, + "the sealed source generation must remain unchanged" + ); + let active_ref = locator.head_generation(&session_ref); + assert_eq!( + active_ref.generation, 1, + "repairing history commits a successor generation" + ); + let active = locator + .read_session_transcript(&active_ref) + .unwrap() + .read_session() + .unwrap() + .unwrap(); + assert_eq!( + active + .messages + .iter() + .map(|row| (row.role.as_str(), row.content.as_str())) + .collect::<Vec<_>>(), + [ + ("system", "frozen prompt"), + ("user", "persisted user"), + ("user", "new user turn"), + ("assistant", "reply after replay"), + ], + "the active generation retains prior conversation and commits the new turn" + ); + assert_eq!( + active.messages[1].extra_metadata, + Some(serde_json::json!({"source": "original user row"})), + "the matched user row keeps its opaque durable metadata" + ); + assert_eq!( + active.meta.session_id.as_deref(), + Some(active_ref.session_id().as_str()) + ); + assert_eq!( + active.meta.thread_id.as_deref(), + Some("thread-orphan-append") + ); +} + +#[tokio::test] +async fn successful_replay_repairs_same_source_for_thread_and_latest_modes() { + for resume in [ResumeMode::Thread, ResumeMode::LatestForAgent] { + let locator = memory(); + let session_ref = SessionRef::scoped("thread-orphan-scan", "agent"); + let mut user = TranscriptMessage::user("persisted user"); + user.extra_metadata = Some(serde_json::json!({"source": "kept"})); + let rows = vec![ + TranscriptMessage::system("frozen prompt"), + TranscriptMessage::tool("orphan result").with_id("orphan-call"), + user, + ]; + seed_replay_transcript(locator.as_ref(), &session_ref, &rows); + let mut session = SessionBuilder::new(Arc::new(AppendReplyDriver)) + .codec(Arc::new(NativeReplayCodec)) + .session(locator.clone(), session_ref.clone(), meta()) + .resume_agent("agent") + .build() + .unwrap(); + + session + .turn( + SessionTurnRequest::new(Message::user("new user turn")), + TurnOptions { + resume, + thread_id: Some(session_ref.session_key.clone()), + ..TurnOptions::default() + }, + ) + .await + .unwrap(); + + let active_ref = locator.head_generation(&session_ref); + let active = locator + .read_session_transcript(&active_ref) + .unwrap() + .read_session() + .unwrap() + .unwrap(); + assert_eq!(active_ref.generation, 1, "{resume:?}"); + assert_eq!( + active + .messages + .iter() + .map(|row| (row.role.as_str(), row.content.as_str())) + .collect::<Vec<_>>(), + [ + ("system", "frozen prompt"), + ("user", "persisted user"), + ("user", "new user turn"), + ("assistant", "reply after replay") + ], + "same-source {resume:?} replay keeps repaired raw/message correspondence" + ); + assert_eq!( + active.messages[1].extra_metadata, + Some(serde_json::json!({"source": "kept"})) + ); + } +} + +struct LegacyToolReplayCodec; + +impl TranscriptCodec for LegacyToolReplayCodec { + fn decode_history(&self, transcript: &SessionTranscript) -> Result<Vec<Message>, RuntimeError> { + Ok(transcript + .messages + .iter() + .map(|row| match row.role.as_str() { + "system" => Message::system(&row.content), + "user" => Message::user(&row.content), + "legacy_tool" => Message::tool(row.id.as_deref().unwrap_or("legacy"), &row.content), + _ => Message::assistant(&row.content), + }) + .collect()) + } +} + +#[tokio::test] +async fn successful_replay_keeps_raw_rows_aligned_after_decoded_tool_repair() { + let locator = memory(); + let session_ref = SessionRef::scoped("thread-legacy-tool", "agent"); + let mut legacy_tool = TranscriptMessage::tool("orphan result").with_id("legacy-call"); + legacy_tool.role = "legacy_tool".into(); + legacy_tool.extra_metadata = Some(serde_json::json!({"opaque": "preserve"})); + let rows = vec![ + TranscriptMessage::system("frozen prompt"), + legacy_tool, + TranscriptMessage::user("persisted user"), + ]; + seed_replay_transcript(locator.as_ref(), &session_ref, &rows); + let original_fingerprint = fingerprint( + locator + .read_session_transcript(&session_ref) + .unwrap() + .as_ref(), + ); + let mut session = SessionBuilder::new(Arc::new(AppendReplyDriver)) + .codec(Arc::new(LegacyToolReplayCodec)) + .session(locator.clone(), session_ref.clone(), meta()) + .build() + .unwrap(); + + session + .turn( + SessionTurnRequest::new(Message::user("new user turn")), + TurnOptions { + resume: ResumeMode::Session, + thread_id: Some(session_ref.session_key.clone()), + ..TurnOptions::default() + }, + ) + .await + .unwrap(); + + assert_eq!( + fingerprint( + locator + .read_session_transcript(&session_ref) + .unwrap() + .as_ref() + ), + original_fingerprint, + "the original legacy raw row remains in the sealed source generation" + ); + let active_ref = locator.head_generation(&session_ref); + let active = locator + .read_session_transcript(&active_ref) + .unwrap() + .read_session() + .unwrap() + .unwrap(); + assert_eq!(active_ref.generation, 1); + assert_eq!( + active + .messages + .iter() + .map(|row| (row.role.as_str(), row.content.as_str())) + .collect::<Vec<_>>(), + [ + ("system", "frozen prompt"), + ("user", "persisted user"), + ("user", "new user turn"), + ("assistant", "reply after replay") + ] + ); +} + +struct BlockingReadSink { + locator: Arc<dyn TranscriptLocator>, + session: SessionRef, + started: Mutex<Option<tokio::sync::oneshot::Sender<()>>>, + release: Arc<Notify>, + finished: Arc<Notify>, +} + +#[async_trait] +impl crate::PostCommitSink for BlockingReadSink { + async fn after_commit( + &self, + observation: &crate::PostCommitObservation, + ) -> Result<(), RuntimeError> { + assert!(observation.receipt.transcript.is_some()); + assert_eq!(observation.input, Message::user("commit this")); + let head = self.locator.head_generation(&self.session); + let read = self.locator.read_session_transcript(&head).unwrap(); + let transcript = read.read_session().unwrap().unwrap(); + assert!( + transcript + .messages + .iter() + .any(|row| row.content == "committed reply") + ); + if let Some(started) = self.started.lock().unwrap().take() { + let _ = started.send(()); + } + self.release.notified().await; + self.finished.notify_one(); + Ok(()) + } +} + +struct ReplyDriver; + +#[async_trait] +impl SessionDriver for ReplyDriver { + async fn execute(&self, request: DriverRequest) -> Result<DriverOutcome, DriverFailure> { + let mut history = request.history; + history.push(Message::assistant("committed reply")); + Ok(DriverOutcome { + history, + output: Some("committed reply".into()), + partial: None, + interrupted: false, + outcome: None, + tool_outcomes: Vec::new(), + }) + } +} + +#[tokio::test] +async fn owned_post_commit_sink_finishes_after_caller_abort() { + let locator = memory(); + let session_ref = SessionRef::scoped("thread-post-commit", "agent"); + let release = Arc::new(Notify::new()); + let finished = Arc::new(Notify::new()); + let (started_tx, started_rx) = tokio::sync::oneshot::channel(); + let sink = Arc::new(BlockingReadSink { + locator: locator.clone(), + session: session_ref.clone(), + started: Mutex::new(Some(started_tx)), + release: release.clone(), + finished: finished.clone(), + }); + let session = SessionBuilder::new(Arc::new(ReplyDriver)) + .codec(Arc::new(RoleCodec)) + .session(locator, session_ref, meta()) + .post_commit_sink(sink) + .build() + .unwrap(); + let turn = tokio::spawn(async move { + let mut session = session; + session + .turn( + SessionTurnRequest::new(Message::user("commit this")), + TurnOptions::default(), + ) + .await + }); + started_rx.await.unwrap(); + turn.abort(); + release.notify_one(); + timeout(Duration::from_secs(1), finished.notified()) + .await + .unwrap(); +} + +struct CountSink(Arc<AtomicUsize>); + +#[async_trait] +impl crate::PostCommitSink for CountSink { + async fn after_commit(&self, _: &crate::PostCommitObservation) -> Result<(), RuntimeError> { + self.0.fetch_add(1, Ordering::SeqCst); + Ok(()) + } +} + +struct FailureTerminalRecorder(Arc<Mutex<Vec<SessionTerminal>>>); + +#[async_trait] +impl SessionHooks for FailureTerminalRecorder { + async fn before_turn( + &self, + _: &mut SessionTurnRequest, + _: &mut TurnOptions, + _: SessionStateView<'_>, + ) -> Result<TurnPreparation, RuntimeError> { + Ok(TurnPreparation::default()) + } + async fn before_commit( + &self, + _: &SessionTurnOutcome, + _: &TranscriptTurnOptions, + ) -> Result<(), RuntimeError> { + Ok(()) + } + async fn on_terminal(&self, terminal: SessionTerminal) -> Result<(), RuntimeError> { + self.0.lock().unwrap().push(terminal); + Ok(()) + } +} + +struct FailingPartialDriver; + +#[async_trait] +impl SessionDriver for FailingPartialDriver { + async fn execute(&self, request: DriverRequest) -> Result<DriverOutcome, DriverFailure> { + let mut history = request.history; + history.push(Message::assistant("durable partial")); + Err(DriverFailure { + error: RuntimeError::Driver("provider failed".into()), + partial: Some(DriverOutcome { + history, + output: Some("durable partial".into()), + partial: None, + interrupted: false, + outcome: None, + tool_outcomes: Vec::new(), + }), + outcome: None, + }) + } +} + +#[tokio::test] +async fn persisted_failure_partial_skips_post_commit_sinks_and_fails_once() { + let locator = memory(); + let session_ref = SessionRef::scoped("thread-failure-partial", "agent"); + let sink_calls = Arc::new(AtomicUsize::new(0)); + let terminals = Arc::new(Mutex::new(Vec::new())); + let mut session = SessionBuilder::new(Arc::new(FailingPartialDriver)) + .codec(Arc::new(RoleCodec)) + .hooks(Arc::new(FailureTerminalRecorder(terminals.clone()))) + .post_commit_sink(Arc::new(CountSink(sink_calls.clone()))) + .session(locator.clone(), session_ref.clone(), meta()) + .build() + .unwrap(); + assert!( + session + .turn( + SessionTurnRequest::new(Message::user("fail")), + TurnOptions::default() + ) + .await + .is_err() + ); + assert_eq!(sink_calls.load(Ordering::SeqCst), 0); + let terminals = terminals.lock().unwrap(); + assert_eq!(terminals.len(), 1); + assert!(matches!(terminals[0], SessionTerminal::Failed(_))); + let head = locator.head_generation(&session_ref); + let transcript = locator + .read_session_transcript(&head) + .unwrap() + .read_session() + .unwrap() + .unwrap(); + assert!( + transcript + .messages + .iter() + .any(|row| row.content == "durable partial") + ); +} + +struct RejectBeforeCommit; + +#[async_trait] +impl SessionHooks for RejectBeforeCommit { + async fn before_turn( + &self, + _: &mut SessionTurnRequest, + _: &mut TurnOptions, + _: SessionStateView<'_>, + ) -> Result<TurnPreparation, RuntimeError> { + Ok(TurnPreparation::default()) + } + async fn before_commit( + &self, + _: &SessionTurnOutcome, + _: &TranscriptTurnOptions, + ) -> Result<(), RuntimeError> { + Err(RuntimeError::Driver("reject before durable append".into())) + } + async fn on_terminal(&self, _: SessionTerminal) -> Result<(), RuntimeError> { + Ok(()) + } +} + +#[tokio::test] +async fn precommit_error_does_not_call_post_commit_sink() { + let calls = Arc::new(AtomicUsize::new(0)); + let mut session = SessionBuilder::new(Arc::new(ReplyDriver)) + .hooks(Arc::new(RejectBeforeCommit)) + .post_commit_sink(Arc::new(CountSink(calls.clone()))) + .build() + .unwrap(); + assert!( + session + .turn( + SessionTurnRequest::new(Message::user("reject")), + TurnOptions::default() + ) + .await + .is_err() + ); + assert_eq!(calls.load(Ordering::SeqCst), 0); +} + +struct InterruptedDriver; + +#[async_trait] +impl SessionDriver for InterruptedDriver { + async fn execute(&self, request: DriverRequest) -> Result<DriverOutcome, DriverFailure> { + Ok(DriverOutcome { + history: request.history, + output: None, + partial: None, + interrupted: true, + outcome: None, + tool_outcomes: Vec::new(), + }) + } +} + +#[tokio::test] +async fn normal_interrupted_commit_still_invokes_post_commit_sink() { + let calls = Arc::new(AtomicUsize::new(0)); + let mut session = SessionBuilder::new(Arc::new(InterruptedDriver)) + .post_commit_sink(Arc::new(CountSink(calls.clone()))) + .build() + .unwrap(); + let outcome = session + .turn( + SessionTurnRequest::new(Message::user("stop")), + TurnOptions::default(), + ) + .await + .unwrap(); + assert!(outcome.interrupted); + assert_eq!(calls.load(Ordering::SeqCst), 1); +} + +struct ErrorSink; + +#[async_trait] +impl crate::PostCommitSink for ErrorSink { + async fn after_commit(&self, _: &crate::PostCommitObservation) -> Result<(), RuntimeError> { + Err(RuntimeError::Driver("sink failed after commit".into())) + } +} + +#[tokio::test] +async fn post_commit_sink_error_does_not_change_successful_turn_result() { + let locator = memory(); + let session_ref = SessionRef::scoped("thread-sink-error", "agent"); + let mut session = SessionBuilder::new(Arc::new(ReplyDriver)) + .codec(Arc::new(RoleCodec)) + .post_commit_sink(Arc::new(ErrorSink)) + .session(locator, session_ref, meta()) + .build() + .unwrap(); + let outcome = session + .turn( + SessionTurnRequest::new(Message::user("succeed")), + TurnOptions::default(), + ) + .await + .unwrap(); + assert_eq!(outcome.output.as_deref(), Some("committed reply")); +} + fn meta() -> TranscriptMeta { TranscriptMeta { session_id: None, @@ -88,6 +877,7 @@ impl SessionDriver for GatedDriver { output: Some("Sunny.".into()), partial: None, interrupted: false, + tool_outcomes: Vec::new(), }) } } diff --git a/crates/tinyagents-runtime/src/summary.rs b/crates/tinyagents-runtime/src/summary.rs new file mode 100644 index 000000000..073424194 --- /dev/null +++ b/crates/tinyagents-runtime/src/summary.rs @@ -0,0 +1,56 @@ +use crate::PostCommitObservation; + +/// Builds the content-minimized learning summary from captured execution facts. +pub fn turn_summary<C: Clone + Send + Sync + 'static>( + observation: &PostCommitObservation<C>, +) -> Option<tinyagents_harness::host::TurnSummary> { + use sha2::{Digest, Sha256}; + let thread = observation.thread_id.as_deref()?.trim(); + if thread.is_empty() { + return None; + } + let usage = observation + .usage + .as_ref() + .map_or_else(Default::default, |usage| tinyinference_llm::usage::Usage { + input_tokens: usage.usage.input, + output_tokens: usage.usage.output, + total_tokens: usage.usage.input.saturating_add(usage.usage.output), + cache_read_tokens: usage.usage.cached_input, + context_window_tokens: (usage.usage.context_window > 0) + .then_some(usage.usage.context_window), + ..Default::default() + }); + let mut summary = tinyagents_harness::host::TurnSummary::new( + tinyagents_harness::ids::ThreadId::from(thread), + observation.agent_id.as_deref()?, + ) + .with_text( + observation.input.text(), + observation + .receipt + .outcome + .output + .clone() + .unwrap_or_default(), + ) + .with_usage(usage); + for call in &observation.tool_outcomes { + let bytes = serde_json::to_vec(&call.arguments).ok()?; + let digest = Sha256::digest(bytes); + let args_digest = digest.iter().map(|byte| format!("{byte:02x}")).collect(); + summary.record_tool_outcome(tinyagents_harness::host::ToolOutcomeRecord { + name: call.name.clone(), + args_digest, + success: call.success, + failure_class: (!call.success) + .then(|| tinyagents_harness::host::classify_tool_failure(&call.content).to_owned()), + duration_ms: call.duration_ms, + }); + } + Some(summary) +} + +#[cfg(test)] +#[path = "summary_tests.rs"] +mod tests; diff --git a/crates/tinyagents-runtime/src/summary_tests.rs b/crates/tinyagents-runtime/src/summary_tests.rs new file mode 100644 index 000000000..cfde62d2f --- /dev/null +++ b/crates/tinyagents-runtime/src/summary_tests.rs @@ -0,0 +1,105 @@ +use super::*; +use crate::types::{CommitReceipt, SessionTurnOutcome, TranscriptTurnOptions}; +use tinyagents_session::transcript::{MessageUsage, TurnUsage}; +use tinyinference_llm::message::Message; + +fn observation(thread_id: Option<&str>, agent_id: Option<&str>) -> PostCommitObservation { + PostCommitObservation { + receipt: CommitReceipt { + outcome: SessionTurnOutcome { + history: vec![Message::user("private input")], + output: Some("private output".into()), + interrupted: false, + }, + options: TranscriptTurnOptions { + request_id: None, + thread_id: thread_id.map(str::to_owned), + stream: false, + resume: crate::ResumeMode::Never, + context: (), + }, + transcript: None, + }, + input: Message::user("private input"), + usage: Some(TurnUsage { + provider: "test-provider".into(), + model: "test-model".into(), + usage: MessageUsage { + input: 12, + output: 5, + cached_input: 3, + context_window: 128, + ..MessageUsage::default() + }, + ts: String::new(), + reasoning_content: None, + tool_calls: Vec::new(), + iteration: 1, + }), + tool_outcomes: vec![tinyagents_harness::turn_outcome::ToolCallOutcome { + call_id: "call-private".into(), + name: "private_tool".into(), + arguments: serde_json::json!({"secret": "do-not-persist"}), + success: false, + content: "timeout: private details".into(), + duration_ms: 23, + }], + thread_id: thread_id.map(str::to_owned), + agent_id: agent_id.map(str::to_owned), + } +} + +#[test] +fn summary_requires_resolved_thread_and_agent_identity() { + assert!(turn_summary(&observation(None, Some("agent"))).is_none()); + assert!(turn_summary(&observation(Some(" "), Some("agent"))).is_none()); + assert!(turn_summary(&observation(Some("thread"), None)).is_none()); +} + +#[test] +fn summary_keeps_only_tool_digest_classification_and_captured_usage() { + use sha2::{Digest, Sha256}; + let summary = turn_summary(&observation(Some("thread-1"), Some("agent-1"))).unwrap(); + let expected = Sha256::digest( + serde_json::to_vec(&serde_json::json!({ + "secret": "do-not-persist" + })) + .unwrap(), + ); + let expected: String = expected.iter().map(|byte| format!("{byte:02x}")).collect(); + assert_eq!(summary.input, "private input"); + assert_eq!(summary.output, "private output"); + assert_eq!(summary.tool_outcomes.len(), 1); + assert_eq!(summary.tool_outcomes[0].args_digest, expected); + assert_eq!( + summary.tool_outcomes[0].failure_class.as_deref(), + Some("timeout") + ); + assert_eq!(summary.tool_outcomes[0].duration_ms, 23); + assert_eq!(summary.usage.input_tokens, 12); + assert_eq!(summary.usage.output_tokens, 5); + assert_eq!(summary.usage.total_tokens, 17); + assert_eq!(summary.usage.cache_read_tokens, 3); + assert_eq!(summary.usage.context_window_tokens, Some(128)); + let encoded = serde_json::to_string(&summary).unwrap(); + assert!(!encoded.contains("do-not-persist")); + assert!(!encoded.contains("timeout: private details")); +} + +#[test] +fn summary_usage_total_saturates_on_overflow() { + let mut observation = observation(Some("thread-1"), Some("agent-1")); + let usage = observation.usage.as_mut().unwrap(); + usage.usage.input = u64::MAX; + usage.usage.output = 1; + let summary = turn_summary(&observation).unwrap(); + assert_eq!(summary.usage.total_tokens, u64::MAX); +} + +#[test] +fn successful_tool_summary_has_no_failure_class() { + let mut observation = observation(Some("thread-1"), Some("agent-1")); + observation.tool_outcomes[0].success = true; + let summary = turn_summary(&observation).unwrap(); + assert!(summary.tool_outcomes[0].failure_class.is_none()); +} diff --git a/crates/tinyagents-runtime/src/test_store.rs b/crates/tinyagents-runtime/src/test_store.rs index 9a3909698..3578ad746 100644 --- a/crates/tinyagents-runtime/src/test_store.rs +++ b/crates/tinyagents-runtime/src/test_store.rs @@ -6,15 +6,77 @@ use std::sync::Arc; -use tinyagents_session::port::{AgentStorageFactory, DriverSessionStores, SessionStoreProvider}; +use tinyagents_session::port::{ + AgentStorageFactory, ConditionedAgentStorageFactory, DriverSessionStores, SessionStoreProvider, +}; use tinyagents_session::transcript::{ - DisplaySessionTranscript, SessionRef, SessionTranscript, ThreadTranscripts, TranscriptHistory, - TranscriptLocator, TranscriptMessage, TranscriptMeta, TranscriptRead, TranscriptTurn, + DisplaySessionTranscript, SessionAdoption, SessionRef, SessionTranscript, ThreadTranscripts, + TranscriptHistory, TranscriptLocator, TranscriptMessage, TranscriptMeta, TranscriptRead, + TranscriptTurn, }; -use tinyagents_storage_kit::{Scope, ScopedStorage}; -use tinystoragedrivers_api::StorageBackend; -use tinystoragedrivers_testkit::MemoryBackend; -use tinystoragedrivers_types::StorageError; +use tinyagents_storage_kit::{Scope, ScopedStorage, WriteCondition}; +use tinystoragedrivers_api::{DocumentStore, Storage, StorageBackend, StreamStore}; +use tinystoragedrivers_types::{Capabilities, Result}; + +/// A driver-backed session provider bound to `scope`, for identity-admission tests. +pub(crate) fn driver_session_stores(scope: Scope) -> Arc<DriverSessionStores> { + driver_session_stores_over(sqlite_backend(), scope) +} + +/// A SQLite backend whose temporary directory stays alive as long as the +/// returned backend. Driver session tests exercise lease-conditioned generation +/// writes, which the testkit's simpler in-memory backend does not implement. +fn sqlite_backend() -> Arc<dyn StorageBackend> { + let directory = tempfile::tempdir().expect("SQLite session test directory"); + let config = tinystoragedrivers::StorageConfig::parse(&format!( + "sqlite:{}", + directory.path().join("storage").display() + )) + .expect("SQLite session test URL"); + let backend = std::thread::spawn(move || { + tokio::runtime::Builder::new_multi_thread() + .enable_all() + .build() + .expect("SQLite session test runtime") + .block_on(tinystoragedrivers::open(&config)) + .expect("SQLite session test backend") + }) + .join() + .expect("SQLite session test thread"); + Arc::new(TemporarySqliteBackend { + _directory: directory, + backend, + }) +} + +#[derive(Debug)] +struct TemporarySqliteBackend { + _directory: tempfile::TempDir, + backend: Arc<dyn StorageBackend>, +} + +impl StorageBackend for TemporarySqliteBackend { + fn capabilities(&self) -> Capabilities { + self.backend.capabilities() + } + + fn for_scope(&self, scope: &Scope) -> Result<Storage> { + self.backend.for_scope(scope) + } +} + +/// A driver-backed session provider bound to `scope` over a shared backend. +pub(crate) fn driver_session_stores_over( + backend: Arc<dyn StorageBackend>, + scope: Scope, +) -> Arc<DriverSessionStores> { + let scope_for_factory = scope.clone(); + let factory: AgentStorageFactory = Arc::new(move |_| { + let storage = backend.for_scope(&scope_for_factory)?; + Ok(ScopedStorage::new(storage, scope_for_factory.clone())) + }); + Arc::new(DriverSessionStores::new(factory).expect("driver session stores over memory")) +} /// An empty transcript store: the production driver-backed locator, over memory. pub(crate) fn memory() -> Arc<dyn TranscriptLocator> { @@ -177,19 +239,121 @@ impl TranscriptLocator for Destination { } } -/// An empty driver-backed transcript store (the production store, over memory). +/// An empty driver-backed transcript store (the production store, over SQLite). pub(crate) fn driver() -> Arc<dyn TranscriptLocator> { - let backend = Arc::new(MemoryBackend::new()); - let factory: AgentStorageFactory = Arc::new(move |agent_id| { - if agent_id != "runtime-tests" { - return Err(StorageError::invalid_input("unknown test agent")); - } - let scope = Scope::new("runtime-tests")?; - let storage = backend.for_scope(&scope)?; - Ok(ScopedStorage::new(storage, scope)) + let scope = Scope::new("runtime-tests").expect("runtime test scope"); + let stores = conditioned_driver_session_stores(scope.clone()); + Arc::new(OwnedTestLocator { + locator: stores.for_agent("runtime-tests").transcripts, + _stores: stores, + }) +} + +fn conditioned_driver_session_stores(scope: Scope) -> Arc<DriverSessionStores> { + use tinystoragedrivers_connect::sqlite::SqliteConnection; + use tinystoragedrivers_document::sqlite::{SqliteDocuments, SqliteStreams}; + + fn storage( + path: &std::path::Path, + scope: &Scope, + memory: &tinystoragedrivers_testkit::MemoryBackend, + condition: Option<&WriteCondition>, + ) -> Result<ScopedStorage> { + let connection = SqliteConnection::open(path) + .map_err(|error| tinystoragedrivers_types::StorageError::backend(error.to_string()))?; + let documents = match condition { + Some(condition) => Arc::new(SqliteDocuments::conditioned( + connection.clone(), + scope.clone(), + condition.clone(), + )) as Arc<dyn DocumentStore>, + None => Arc::new(SqliteDocuments::new(connection.clone(), scope.clone())) + as Arc<dyn DocumentStore>, + }; + let streams = + Arc::new(SqliteStreams::new(connection, scope.clone())) as Arc<dyn StreamStore>; + let ports = memory + .for_scope(scope) + .map_err(|error| tinystoragedrivers_types::StorageError::backend(error.to_string()))?; + Ok(ScopedStorage::new( + Storage::new( + documents, + streams, + ports.kv().clone(), + ports.blobs().clone(), + ports.queues().clone(), + ports.secrets().clone(), + ), + scope.clone(), + )) + } + + let directory = Arc::new(tempfile::tempdir().expect("SQLite session test directory")); + let path = directory.path().join("session.sqlite"); + let plain_path = path.clone(); + let memory = Arc::new(tinystoragedrivers_testkit::MemoryBackend::new()); + let plain_scope = scope.clone(); + let plain_dir = Arc::clone(&directory); + let plain_memory = Arc::clone(&memory); + let factory: AgentStorageFactory = Arc::new(move |_| { + let _keep_directory = &plain_dir; + storage(&plain_path, &plain_scope, &plain_memory, None) + }); + let condition_scope = scope; + let condition_dir = Arc::clone(&directory); + let condition_memory = memory; + let conditioned: ConditionedAgentStorageFactory = Arc::new(move |_, condition| { + let _keep_directory = &condition_dir; + storage(&path, &condition_scope, &condition_memory, Some(condition)) }); - DriverSessionStores::new(factory) - .expect("driver session stores over memory") + Arc::new( + DriverSessionStores::new(factory) + .expect("conditioned SQLite test stores") + .with_conditioned_storage_factory(conditioned), + ) +} + +/// An empty in-memory session locator for behavior tests that do not certify +/// the driver's storage adapter. +pub(crate) fn in_memory() -> Arc<dyn TranscriptLocator> { + tinyagents_session::port::InMemorySessionStores::new() .for_agent("runtime-tests") .transcripts } + +/// A transcript handle plus its backend, which retains the fixture directory. +struct OwnedTestLocator { + locator: Arc<dyn TranscriptLocator>, + _stores: Arc<DriverSessionStores>, +} + +impl TranscriptLocator for OwnedTestLocator { + fn destination_key(&self) -> Option<String> { + self.locator.destination_key() + } + + fn latest_for_agent(&self, agent_name: &str) -> Option<Arc<dyn TranscriptRead>> { + self.locator.latest_for_agent(agent_name) + } + + fn root_for_thread(&self, thread_id: &str) -> Option<Arc<dyn TranscriptRead>> { + self.locator.root_for_thread(thread_id) + } + + fn open_stem( + &self, + stem: &str, + seed: TranscriptMeta, + ) -> anyhow::Result<Arc<dyn TranscriptHistory>> { + self.locator.open_stem(stem, seed) + } + + fn adopt_legacy( + &self, + session: &SessionRef, + thread_id: &str, + seed: &TranscriptMeta, + ) -> anyhow::Result<Option<SessionAdoption>> { + self.locator.adopt_legacy(session, thread_id, seed) + } +} diff --git a/crates/tinyagents-runtime/src/tool_registration.rs b/crates/tinyagents-runtime/src/tool_registration.rs index 538b98433..68781def5 100644 --- a/crates/tinyagents-runtime/src/tool_registration.rs +++ b/crates/tinyagents-runtime/src/tool_registration.rs @@ -3,15 +3,81 @@ //! The host decides *which* tools exist and how each is wrapped; this module //! owns the generic rules for which candidates register: de-duplication by //! name, an allowlist whose empty set denies everything, thread-scoped tools -//! dropped on thread-less turns, and spawn/delegate tools never handed to a -//! sub-agent run whatever the allowlist says. +//! dropped on thread-less turns, and spawn/delegate tools withheld from +//! sub-agent runs by default. A host may grant bounded nested spawning for a +//! turn; those tools still need to pass the ordinary allowlist, and the host +//! remains responsible for enforcing the depth ceiling when a child executes. use std::collections::HashSet; use std::sync::Arc; use tinyagents_harness::runtime::AgentHarness; use tinyagents_harness::tool::ToolDispatch; -use tinytools::Tool; +use tinyagents_harness::{ + context::RunContext, + middleware::{PolicyDecision, ToolCallPolicy}, +}; +use tinyinference_llm::tool::ToolCall; +use tinytools::{Tool, ToolSpec}; + +/// A deny-by-default host belt that may delegate allowed calls to host policy. +pub struct HostOnlyToolPolicy<C: Send + Sync + 'static> { + allowed: HashSet<String>, + host: Option<Arc<dyn ToolCallPolicy<C>>>, +} + +impl<C: Send + Sync + 'static> HostOnlyToolPolicy<C> { + pub fn new( + allowed: impl IntoIterator<Item = String>, + host: Option<Arc<dyn ToolCallPolicy<C>>>, + ) -> Self { + Self { + allowed: allowed.into_iter().collect(), + host, + } + } +} + +#[async_trait::async_trait] +impl<C: Send + Sync + 'static> ToolCallPolicy<C> for HostOnlyToolPolicy<C> { + fn name(&self) -> &str { + "host_only" + } + + async fn check(&self, ctx: &RunContext<C>, call: &ToolCall) -> PolicyDecision { + if !self.allowed.contains(&call.name) { + return PolicyDecision::deny(format!( + "tool `{}` is not in this host-only belt", + call.name + )); + } + match &self.host { + Some(host) => host.check(ctx, call).await, + None => PolicyDecision::Allow, + } + } +} + +/// Keeps the first declaration for each name, preserving order and allocation. +pub fn dedup_visible_tool_specs(specs: Vec<Arc<ToolSpec>>) -> Vec<Arc<ToolSpec>> { + let mut names = HashSet::new(); + specs + .into_iter() + .filter(|spec| names.insert(spec.name.clone())) + .collect() +} + +/// Removes synthesized declarations shadowed by durable declarations. +pub fn drop_synthesized_name_collisions( + durable: &[Arc<ToolSpec>], + synthesized: Vec<Arc<ToolSpec>>, +) -> Vec<Arc<ToolSpec>> { + let durable_names: HashSet<&str> = durable.iter().map(|spec| spec.name.as_str()).collect(); + synthesized + .into_iter() + .filter(|spec| !durable_names.contains(spec.name.as_str())) + .collect() +} /// How one admitted tool is registered on the harness. pub enum ToolRegistration<State, C> { @@ -21,6 +87,28 @@ pub enum ToolRegistration<State, C> { Dispatch(Arc<dyn ToolDispatch<State, C>>), } +/// Depth and ceiling supplied by a host that explicitly permits nested +/// sub-agent creation for this turn. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct NestedSpawnDepth { + depth: usize, + max_depth: usize, +} + +impl NestedSpawnDepth { + /// Creates a depth grant only when the current depth is within its ceiling. + pub fn try_new(depth: usize, max_depth: usize) -> Option<Self> { + (depth <= max_depth).then_some(Self { depth, max_depth }) + } + + /// Returns whether one more child fits within the host's ceiling. + pub fn allows_child(self) -> bool { + self.depth + .checked_add(1) + .is_some_and(|child_depth| child_depth <= self.max_depth) + } +} + /// Admission rules for one turn's tool registration. #[derive(Clone, Default)] pub struct ToolAdmission { @@ -33,8 +121,14 @@ pub struct ToolAdmission { pub has_thread: bool, /// Names of tools that cannot run without a thread. pub thread_scoped: HashSet<String>, - /// Names of spawn/delegate tools a sub-agent must never receive. + /// Names of spawn/delegate tools denied to sub-agents unless the host + /// supplies a [`NestedSpawnDepth`] grant that permits another child. + /// These tools must also pass `allowed`; the grant does not expand that + /// allowlist. pub spawn_tools: HashSet<String>, + /// Explicit host grant for bounded nested spawning. `None` keeps the + /// default leaf-subagent behavior. + pub nested_spawn_depth: Option<NestedSpawnDepth>, } /// The outcome of [`plan_registration`]. @@ -100,7 +194,12 @@ where if !allowed { continue; } - if admission.is_subagent_run && admission.spawn_tools.contains(name) { + if admission.is_subagent_run + && admission.spawn_tools.contains(name) + && !admission + .nested_spawn_depth + .is_some_and(NestedSpawnDepth::allows_child) + { tracing::warn!( tool = name, "[tool-registration] refused spawn tool on sub-agent" diff --git a/crates/tinyagents-runtime/src/tool_registration_tests.rs b/crates/tinyagents-runtime/src/tool_registration_tests.rs index 7dcfbb1be..baf5f2d4b 100644 --- a/crates/tinyagents-runtime/src/tool_registration_tests.rs +++ b/crates/tinyagents-runtime/src/tool_registration_tests.rs @@ -1,4 +1,5 @@ use super::*; +use std::sync::atomic::{AtomicUsize, Ordering}; fn set(names: &[&str]) -> HashSet<String> { names.iter().map(|n| n.to_string()).collect() @@ -11,6 +12,83 @@ fn open() -> ToolAdmission { } } +struct CountingPolicy(Arc<AtomicUsize>); + +#[async_trait::async_trait] +impl ToolCallPolicy<()> for CountingPolicy { + fn name(&self) -> &str { + "counting-host" + } + async fn check(&self, _: &RunContext<()>, _: &ToolCall) -> PolicyDecision { + self.0.fetch_add(1, Ordering::SeqCst); + PolicyDecision::Allow + } +} + +fn policy_call(name: &str) -> ToolCall { + ToolCall { + id: "call-1".into(), + name: name.into(), + arguments: serde_json::json!({}), + invalid: None, + } +} + +#[tokio::test] +async fn host_only_policy_denies_before_delegating_and_allows_named_calls() { + let calls = Arc::new(AtomicUsize::new(0)); + let policy = HostOnlyToolPolicy::new( + ["read_file".to_owned()], + Some(Arc::new(CountingPolicy(calls.clone()))), + ); + let context = RunContext::new( + tinyagents_harness::context::RunConfig::new("host-only-test"), + (), + ); + let denied = policy.check(&context, &policy_call("write_file")).await; + assert!(matches!(denied, PolicyDecision::Deny { .. })); + assert_eq!(calls.load(Ordering::SeqCst), 0); + assert_eq!( + policy.check(&context, &policy_call("read_file")).await, + PolicyDecision::Allow + ); + assert_eq!(calls.load(Ordering::SeqCst), 1); + let empty = HostOnlyToolPolicy::new(Vec::<String>::new(), None); + assert!(matches!( + empty.check(&context, &policy_call("read_file")).await, + PolicyDecision::Deny { .. } + )); +} + +fn tool_spec(name: &str, description: &str) -> Arc<ToolSpec> { + Arc::new(ToolSpec { + name: name.into(), + description: description.into(), + parameters: serde_json::json!({"type":"object"}), + }) +} + +#[test] +fn visible_spec_dedup_keeps_first_order_and_allocation() { + let first = tool_spec("read", "first"); + let duplicate = tool_spec("read", "duplicate"); + let second = tool_spec("write", "write"); + let result = dedup_visible_tool_specs(vec![first.clone(), duplicate, second.clone()]); + assert_eq!(result.len(), 2); + assert!(Arc::ptr_eq(&result[0], &first)); + assert!(Arc::ptr_eq(&result[1], &second)); +} + +#[test] +fn durable_declarations_remove_synthesized_name_collisions() { + let durable = tool_spec("search", "durable"); + let search = tool_spec("search", "synthesized"); + let render = tool_spec("render", "render"); + let result = drop_synthesized_name_collisions(&[durable], vec![search, render.clone()]); + assert_eq!(result.len(), 1); + assert!(Arc::ptr_eq(&result[0], &render)); +} + #[test] fn no_allowlist_admits_all_deduplicated_in_order() { let plan = plan_registration(["b", "a", "b"], &open()); @@ -63,6 +141,47 @@ fn subagent_never_gets_spawn_tools_even_if_allowlisted() { assert_eq!(plan.spawn_refused, ["spawn"]); } +#[test] +fn bounded_subagent_spawn_requires_depth_room_and_exact_allowlisting() { + let below_cap = ToolAdmission { + allowed: Some(set(&["spawn", "read"])), + is_subagent_run: true, + spawn_tools: set(&["spawn"]), + nested_spawn_depth: Some(NestedSpawnDepth::try_new(1, 3).unwrap()), + ..open() + }; + let plan = plan_registration(["spawn", "read", "other"], &below_cap); + assert_eq!(plan.admitted, ["spawn", "read"]); + assert!(plan.spawn_refused.is_empty()); + + let at_cap = ToolAdmission { + nested_spawn_depth: Some(NestedSpawnDepth::try_new(3, 3).unwrap()), + ..below_cap.clone() + }; + let plan = plan_registration(["spawn", "read"], &at_cap); + assert_eq!(plan.admitted, ["read"]); + assert_eq!(plan.spawn_refused, ["spawn"]); + + let not_allowlisted = ToolAdmission { + allowed: Some(set(&["read"])), + nested_spawn_depth: Some(NestedSpawnDepth::try_new(1, 3).unwrap()), + ..below_cap + }; + let plan = plan_registration(["spawn", "read"], ¬_allowlisted); + assert_eq!(plan.admitted, ["read"]); + assert!(plan.spawn_refused.is_empty()); +} + +#[test] +fn nested_spawn_depth_rejects_invalid_or_overflowing_depths() { + assert!(NestedSpawnDepth::try_new(4, 3).is_none()); + assert!( + !NestedSpawnDepth::try_new(usize::MAX, usize::MAX) + .unwrap() + .allows_child() + ); +} + #[test] fn root_run_keeps_spawn_tools() { let admission = ToolAdmission { diff --git a/crates/tinyagents-runtime/src/transcript/mod.rs b/crates/tinyagents-runtime/src/transcript/mod.rs index 47e9f7a17..9104bd6e2 100644 --- a/crates/tinyagents-runtime/src/transcript/mod.rs +++ b/crates/tinyagents-runtime/src/transcript/mod.rs @@ -629,7 +629,7 @@ pub fn message_to_display_transcript(msg: &Message) -> Option<TranscriptMessage> /// with pictures keeps its text and image blocks as ordered `parts`. Without /// this the provider sees an assistant with no `tool_calls` followed by an /// orphan tool message and drops the round — breaking multi-turn native tool -/// calling (e.g. the orchestrator's `spawn_parallel_agents` → synthesis hop) — +/// calling (e.g. a parallel-agent fanout followed by synthesis) — /// and a native-tool provider silently loses every pasted image. /// /// Returns `None` for [`Message::Custom`]: that variant is a host-side diff --git a/crates/tinyagents-runtime/src/turn_driver/driver.rs b/crates/tinyagents-runtime/src/turn_driver/driver.rs index 03ac9a3d4..898942106 100644 --- a/crates/tinyagents-runtime/src/turn_driver/driver.rs +++ b/crates/tinyagents-runtime/src/turn_driver/driver.rs @@ -23,6 +23,7 @@ use async_trait::async_trait; use tinyagents_harness::agent_loop::PartialRunOutcome; use tinyagents_harness::context::RunContext; use tinyagents_harness::events::EventListener; +use tinyagents_harness::host::PrebuiltModelResolver; use tinyagents_harness::middleware::AgentRun; use tinyagents_harness::middleware::{ ApprovalResolver, Middleware, ToolCallPolicy, ToolPolicyGateMiddleware, @@ -171,6 +172,7 @@ pub struct TurnTail { max_iterations: usize, tools: ToolSetReport, completion: Option<Arc<dyn TurnCompletionSink>>, + outcomes: tinyagents_harness::turn_outcome::ToolOutcomeSink, } impl TurnTail { @@ -189,6 +191,7 @@ impl TurnTail { /// `None` when the turn failed before producing one (no cap is detected /// then). pub fn finish(self, run: Option<&AgentRun>) -> TurnReport { + let tool_outcomes = tinyagents_harness::turn_outcome::drain_tool_outcomes(&self.outcomes); // Stop forwarding and deregister now, not at end of scope. drop(self.guard); let early_exit = self.hook.and_then(|hook| hook.take()); @@ -212,6 +215,7 @@ impl TurnTail { model_calls: run.map_or(0, |run| run.model_calls), max_iterations: self.max_iterations, tools: self.tools, + tool_outcomes, completion_deferred: self.controls.defer_turn_completed, ..TurnReport::default() } @@ -383,12 +387,18 @@ impl<State: Send + Sync + 'static, C: Send + Sync + 'static> TurnDriver<State, C let mut harness: AgentHarness<State, C> = AgentHarness::new(); harness.with_policy(policy); - for (index, (name, model)) in models.into_iter().enumerate() { - harness.register_model(name.clone(), model); - if index == 0 { - harness.set_default_model(name); - } + if let Some((primary_name, primary_model)) = models.first().cloned() { + let routes = models.iter().skip(1).cloned().collect::<Vec<_>>(); + let resolver = + PrebuiltModelResolver::new(primary_model, routes).with_primary_name(primary_name); + harness.with_model_resolver(Arc::new(resolver)); } + let outcome_sink = tinyagents_harness::turn_outcome::new_tool_outcome_sink(); + harness.push_middleware(Arc::new( + tinyagents_harness::middleware::library::ToolOutcomeCaptureMiddleware::new(Arc::clone( + &outcome_sink, + )), + )); for middleware in context { harness.push_middleware(middleware); } @@ -479,6 +489,7 @@ impl<State: Send + Sync + 'static, C: Send + Sync + 'static> TurnDriver<State, C max_iterations, tools: tool_report, completion, + outcomes: outcome_sink, }, }) } @@ -499,6 +510,15 @@ impl<State: Send + Sync + 'static, C: Send + Sync + 'static> SessionDriver<C> } = prepared.run().await; let mut outcome = outcome_from_partial(partial, input_len); + let tool_outcomes = report.tool_outcomes.clone(); + match &mut outcome { + Ok(done) => done.tool_outcomes = tool_outcomes, + Err(failure) => { + if let Some(partial) = &mut failure.partial { + partial.tool_outcomes = tool_outcomes; + } + } + } // An early-exit tool's output stands in for the paused assistant turn, // which carries the call and no answer. if let (Some(exit), Ok(done)) = (&report.early_exit, &mut outcome) { diff --git a/crates/tinyagents-runtime/src/turn_driver/driver_tests.rs b/crates/tinyagents-runtime/src/turn_driver/driver_tests.rs index ef88f8553..8f0628c3a 100644 --- a/crates/tinyagents-runtime/src/turn_driver/driver_tests.rs +++ b/crates/tinyagents-runtime/src/turn_driver/driver_tests.rs @@ -422,6 +422,27 @@ async fn tool_policy_rejection_is_returned_to_the_model_and_never_executes() { ); } +#[tokio::test] +async fn host_only_tool_policy_blocks_unknown_calls_in_the_turn_plan() { + let executions = Arc::new(AtomicUsize::new(0)); + let model = Arc::new(ScriptedModel::new(vec![ + tool_call_response(call("c1", "danger", 1)), + text_response("understood"), + ])); + let tools: Vec<Box<dyn Tool>> = vec![Box::new(CountingTool(executions.clone(), "danger"))]; + let configure: Configure = Arc::new(|plan| { + plan.tool_policy = Some(Arc::new(crate::HostOnlyToolPolicy::new( + Vec::<String>::new(), + None, + ))); + }); + let mut session = session(&model, tools, &["danger"], configure); + let outcome = turn(&mut session, TurnControls::default()).await.unwrap(); + assert_eq!(outcome.output.as_deref(), Some("understood")); + assert_eq!(executions.load(Ordering::SeqCst), 0); + assert!(all_text(&model.requests()[1]).contains("host-only belt")); +} + struct Trace { name: &'static str, log: Arc<Mutex<Vec<String>>>, @@ -812,6 +833,11 @@ async fn a_split_turn_is_driven_by_the_host_and_settled_by_its_tail() { let chat: Arc<dyn ChatModel<()>> = model.clone(); let mut plan = TurnPlan::new(vec![("scripted".to_string(), chat)]); plan.tools = ToolSetBuilder::new().set(set.clone()); + plan.context.push(Arc::new(Trace { + name: "rewrite", + log: Arc::new(Mutex::new(Vec::new())), + rewrite_result: true, + })); Ok(plan) } }; @@ -840,6 +866,9 @@ async fn a_split_turn_is_driven_by_the_host_and_settled_by_its_tail() { let report = tail.finish(Some(&run)); assert_eq!(report.early_exit.unwrap().question, "Which file?"); + assert_eq!(report.tool_outcomes.len(), 1); + assert_eq!(report.tool_outcomes[0].name, "ask_user"); + assert_eq!(report.tool_outcomes[0].content, "[capped]"); assert_eq!( slot.get().unwrap().model_calls, run.model_calls, @@ -848,6 +877,46 @@ async fn a_split_turn_is_driven_by_the_host_and_settled_by_its_tail() { assert_eq!(model.requests().len(), 1); } +#[tokio::test] +async fn prepared_run_exposes_the_same_rewritten_tool_outcome() { + let model = Arc::new(ScriptedModel::new(vec![ + tool_call_response(call("c1", "ask_user", 1)), + text_response("never requested"), + ])); + let tools: Vec<Box<dyn Tool>> = vec![Box::new(FakeTool::returning("ask_user", "question"))]; + let snapshot = snapshot_of(&tools, &["ask_user"]); + let set: SharedToolSet = Arc::new(tools); + let planner = { + let model = model.clone(); + move |_: &DriverRequest<()>| { + let chat: Arc<dyn ChatModel<()>> = model.clone(); + let mut plan = TurnPlan::new(vec![("scripted".to_string(), chat)]); + plan.tools = ToolSetBuilder::new().set(set.clone()); + plan.context.push(Arc::new(Trace { + name: "rewrite", + log: Arc::new(Mutex::new(Vec::new())), + rewrite_result: true, + })); + Ok(plan) + } + }; + let driver = TurnDriver::new(Arc::new(planner), Arc::new(())); + let request = DriverRequest { + history: vec![Message::user("go")], + tools: snapshot, + run_context: RunContext::new( + tinyagents_harness::context::RunConfig::new("prepared-run"), + (), + ), + stream: false, + controls: TurnControls::default(), + }; + let run = driver.prepare(request).await.unwrap().run().await; + assert_eq!(run.report.tool_outcomes.len(), 1); + assert_eq!(run.report.tool_outcomes[0].name, "ask_user"); + assert_eq!(run.report.tool_outcomes[0].content, "[capped]"); +} + /// The cap pause reports `(completed, cap)` to the host's recorder. #[tokio::test] async fn the_cap_pause_is_reported_to_the_recorder() { diff --git a/crates/tinyagents-runtime/src/turn_driver/types.rs b/crates/tinyagents-runtime/src/turn_driver/types.rs index 2957fb9ac..980d548a7 100644 --- a/crates/tinyagents-runtime/src/turn_driver/types.rs +++ b/crates/tinyagents-runtime/src/turn_driver/types.rs @@ -85,6 +85,8 @@ pub struct TurnReport { pub max_iterations: usize, /// How tool registration resolved (candidates, registered, refused). pub tools: ToolSetReport, + /// Tool-call outcomes captured at the harness boundary. + pub tool_outcomes: Vec<tinyagents_harness::turn_outcome::ToolCallOutcome>, /// The completion signal was deferred to the caller. Call /// [`Self::emit_completed`] once its post-run work is done. pub completion_deferred: bool, @@ -126,6 +128,7 @@ impl fmt::Debug for TurnReport { .field("model_calls", &self.model_calls) .field("max_iterations", &self.max_iterations) .field("tools", &self.tools) + .field("tool_outcomes", &self.tool_outcomes) .field("completion_deferred", &self.completion_deferred) .field("completed_emitted", &self.completed_emitted()) .finish() diff --git a/crates/tinyagents-runtime/src/types.rs b/crates/tinyagents-runtime/src/types.rs index 85f5be59f..62421184e 100644 --- a/crates/tinyagents-runtime/src/types.rs +++ b/crates/tinyagents-runtime/src/types.rs @@ -241,6 +241,43 @@ pub struct SessionStateView<'a> { /// `true` only when this call loaded and decoded a durable transcript /// before `before_turn` ran. pub resumed: bool, + /// Last established host identity binding, if this session uses a resolver. + pub identity: Option<&'a SessionIdentity>, +} + +/// Logical naming inputs passed to the host identity resolver. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct SessionIdentityRequest { + pub thread_id: Option<String>, + pub session_key: String, + /// Logical parent lineage, never authorization. + pub parent: Option<SessionRef>, + /// Child key paired with `parent`, never an actor or tenant selector. + pub child_key: Option<String>, +} + +/// Host-authorized agent identity and pinned session stores. +#[derive(Clone)] +pub struct ResolvedSessionStores { + pub agent_id: String, + pub stores: tinyagents_session::port::SessionStoreLease, +} + +/// Runtime-created identity binding established after host admission. +#[derive(Clone)] +pub struct SessionIdentity { + pub(crate) session: SessionRef, + pub(crate) agent_id: String, + pub(crate) stores: tinyagents_session::port::SessionStoreLease, +} + +impl SessionIdentity { + pub fn session(&self) -> &SessionRef { + &self.session + } + pub fn agent_id(&self) -> &str { + &self.agent_id + } } /// The shape of a successful logical transcript transition. @@ -279,6 +316,23 @@ pub struct CommitReceipt<C = ()> { pub transcript: Option<TranscriptCommitReceipt>, } +/// Separate facts captured at execution time for post-commit host adapters. +#[derive(Clone, Debug)] +pub struct PostCommitObservation<C: Clone + Send + Sync + 'static = ()> { + /// The durable commit and its session result. + pub receipt: CommitReceipt<C>, + /// The original request input before host hooks can replace it. + pub input: Message, + /// Provider usage captured for the committed turn. + pub usage: Option<tinyagents_session::transcript::TurnUsage>, + /// Tool results captured at the harness execution boundary. + pub tool_outcomes: Vec<tinyagents_harness::turn_outcome::ToolCallOutcome>, + /// Host-selected conversation identity. + pub thread_id: Option<String>, + /// Resolved agent identity, when one was supplied by the host. + pub agent_id: Option<String>, +} + impl<C: Clone> TurnOptions<C> { pub(crate) fn transcript_options(&self) -> TranscriptTurnOptions<C> { TranscriptTurnOptions { diff --git a/crates/tinyagents-session/Cargo.toml b/crates/tinyagents-session/Cargo.toml index 4271cc9fd..8e16dd90a 100644 --- a/crates/tinyagents-session/Cargo.toml +++ b/crates/tinyagents-session/Cargo.toml @@ -15,10 +15,10 @@ anyhow = { workspace = true } thiserror = "2" serde = { workspace = true } serde_json = { workspace = true } -tinyagents-harness = { path = "../tinyagents-harness", version = "2.11.0", default-features = false, features = ["sqlite"] } +tinyagents-harness = { path = "../tinyagents-harness", version = "3.0.0", default-features = false, features = ["sqlite"] } # `OrchestrationTaskStatus` is the canonical lifecycle status the run-ledger and # transcript-view statuses convert to and from. -tinyagents-tasks = { path = "../tinyagents-tasks", version = "2.11.0", default-features = false } +tinyagents-tasks = { path = "../tinyagents-tasks", version = "3.0.0", default-features = false } tinytools-agent = { path = "../../vendor/tinytools/crates/tinytools-agent", version = "0.6.0", default-features = false } hex = "0.4" sha2 = "0.11" @@ -47,6 +47,8 @@ tracing = ["tinyagents-harness/tracing"] storage-drivers = [] [dev-dependencies] +tinystoragedrivers-connect = { workspace = true } +tinystoragedrivers-document = { version = "2.1.0", git = "https://github.com/tinyhumansai/tinystoragedrivers", rev = "05eaada962decbab33bafa4763d2bd1a67f6b456" } tinystoragedrivers-testkit = { workspace = true } tempfile = { workspace = true } tokio = { workspace = true, default-features = true, features = ["macros", "rt-multi-thread"] } diff --git a/crates/tinyagents-session/src/lib.rs b/crates/tinyagents-session/src/lib.rs index cc986502f..0dd579f42 100644 --- a/crates/tinyagents-session/src/lib.rs +++ b/crates/tinyagents-session/src/lib.rs @@ -22,8 +22,6 @@ //! //! - `tinyagents_harness::store` is namespaced key-value storage for live //! runtime data. It is a substrate runs read and write during execution. -//! - `tinyagents_graph::checkpoint` is durability for *resuming* an interrupted -//! graph run. //! - This crate is queryable history. Nothing resumes from it; it answers //! "what happened", supports cross-session search, and lets a host recover //! orchestration state after a restart. diff --git a/crates/tinyagents-session/src/port/drivers/mod.rs b/crates/tinyagents-session/src/port/drivers/mod.rs index d00b603aa..2af10d99d 100644 --- a/crates/tinyagents-session/src/port/drivers/mod.rs +++ b/crates/tinyagents-session/src/port/drivers/mod.rs @@ -39,7 +39,10 @@ use tinyagents_storage_kit::{ use tinystoragedrivers_api::DocumentStore; use tinystoragedrivers_types::{CollectionSpec, StorageError}; -use super::{AgentStorageFactory, AgentStores, SessionStoreProvider}; +use super::{ + AgentStorageFactory, AgentStores, ConditionedAgentStorageFactory, SessionStoreLease, + SessionStoreLeaseError, SessionStoreProvider, +}; pub use rebound::{RebindableSessionStores, SessionStoreRegistry}; pub use schema::storage_schema; @@ -53,6 +56,7 @@ const KV_OWNER: &str = "session"; /// A [`SessionStoreProvider`] that obtains already-bound storage from its host. pub struct DriverSessionStores { factory: AgentStorageFactory, + conditioned_factory: Option<ConditionedAgentStorageFactory>, bridge: Blocking, agents: Mutex<HashMap<String, AgentStores>>, recover_on_open: bool, @@ -114,6 +118,7 @@ impl DriverSessionStores { ) -> Self { Self { factory, + conditioned_factory: None, bridge, agents: Mutex::new(HashMap::new()), recover_on_open: false, @@ -124,6 +129,18 @@ impl DriverSessionStores { } } + /// Enables native lease-conditioned writes for the first entry in a + /// reserved transcript generation. The factory must open the same agent + /// scope as `factory`, with the supplied condition enforced atomically. + /// Without this adapter, generation writes fail closed. + pub fn with_conditioned_storage_factory( + mut self, + factory: ConditionedAgentStorageFactory, + ) -> Self { + self.conditioned_factory = Some(factory); + self + } + /// Marks an agent's in-flight turns interrupted the first time this /// provider opens its stores. /// @@ -235,7 +252,7 @@ impl DriverSessionStores { Ok::<_, StorageError>(()) })??; } - let (stores, turn_states) = self.build(&scoped); + let (stores, turn_states) = self.build(agent_id, &scoped); if self.recover_on_open { // Fail closed: stores whose recovery did not run must not start // turns a retried sweep would then mistake for crash residue. @@ -297,7 +314,11 @@ impl DriverSessionStores { } /// Builds public stores and retains their exact typed recovery handle. - fn build(&self, scoped: &ScopedStorage) -> (AgentStores, Arc<DriverTurnStates>) { + fn build( + &self, + agent_id: &str, + scoped: &ScopedStorage, + ) -> (AgentStores, Arc<DriverTurnStates>) { let docs = Arc::clone(scoped.documents()); // The provider's own id keeps two providers with the same driver and // scope from claiming one destination; unlike an address, it is @@ -308,13 +329,38 @@ impl DriverSessionStores { self.bridge.clone(), self.schema_admission, )); - let stores = AgentStores { - transcripts: Arc::new(DriverTranscriptLocator::new_with_schema_admission( + let transcripts = match self.conditioned_factory.as_ref() { + Some(factory) => { + let factory = Arc::clone(factory); + let agent_id = agent_id.to_owned(); + let expected_scope = scoped.scope().clone(); + let opener: transcripts::ConditionedEntryOpener = Arc::new(move |condition| { + let conditioned = factory(&agent_id, condition)?; + if conditioned.scope().as_str() != expected_scope.as_str() { + return Err(StorageError::invalid_input( + "conditioned transcript storage changed the agent scope", + )); + } + Ok(Arc::clone(conditioned.documents())) + }); + Arc::new(DriverTranscriptLocator::new_with_conditioned_entry_opener( + Arc::clone(&docs), + self.bridge.clone(), + label, + self.schema_admission, + scoped.scope().clone(), + opener, + )) + } + None => Arc::new(DriverTranscriptLocator::new_with_schema_admission( Arc::clone(&docs), self.bridge.clone(), label, self.schema_admission, )), + }; + let stores = AgentStores { + transcripts, turn_states: turn_states.clone(), kv: Arc::new(DriverStore::with_owner_and_schema_admission( docs, @@ -358,6 +404,12 @@ impl SessionStoreProvider for DriverSessionStores { DriverSessionStores::try_for_agent(self, agent_id) } + fn lease_for_agent(&self, agent_id: &str) -> Result<SessionStoreLease, SessionStoreLeaseError> { + self.try_for_agent(agent_id) + .map(SessionStoreLease::pinned) + .map_err(SessionStoreLeaseError::Open) + } + fn bind_actor( &self, factory: AgentStorageFactory, diff --git a/crates/tinyagents-session/src/port/drivers/mod_tests.rs b/crates/tinyagents-session/src/port/drivers/mod_tests.rs index 57c7753d9..61d1fb2d9 100644 --- a/crates/tinyagents-session/src/port/drivers/mod_tests.rs +++ b/crates/tinyagents-session/src/port/drivers/mod_tests.rs @@ -2,7 +2,8 @@ use super::*; use crate::testkit::MemoryStorage; use crate::testkit::conformance::{session_store_conformance, session_store_isolation_conformance}; use crate::turn_state::{TurnLifecycle, TurnState}; -use tinystoragedrivers_api::StorageBackend; +use tinyagents_storage_kit::WriteCondition; +use tinystoragedrivers_api::{DocumentStore, StorageBackend}; fn storage_factory(backend: Arc<dyn StorageBackend>) -> AgentStorageFactory { let bindings = [ @@ -25,8 +26,74 @@ fn storage_factory(backend: Arc<dyn StorageBackend>) -> AgentStorageFactory { crate::testkit::mapped_storage_factory(backend, bindings) } -fn memory() -> DriverSessionStores { - DriverSessionStores::new(storage_factory(Arc::new(MemoryStorage::new()))).unwrap() +fn conditioned_factories() -> (AgentStorageFactory, ConditionedAgentStorageFactory) { + use tinystoragedrivers_connect::sqlite::SqliteConnection; + use tinystoragedrivers_document::sqlite::SqliteDocuments; + + let directory = Arc::new(tempfile::tempdir().unwrap()); + let connection = SqliteConnection::open(directory.path().join("session-tests.db")).unwrap(); + let memory: Arc<dyn StorageBackend> = Arc::new(MemoryStorage::new()); + let scopes = [ + ("plain", "test-bound-plain"), + ("a", "test-bound-a"), + ("b", "test-bound-b"), + ("agent", "test-bound-agent"), + ("agent-1", "test-bound-agent-1"), + ("other", "test-bound-other"), + ("alice", "test-bound-alice"), + ("bob", "test-bound-bob"), + ("prepared-actor", "test-bound-prepared-actor"), + ("contract-store-agent", "test-bound-contract-store"), + ("contract-iso-a", "test-bound-contract-iso-a"), + ("contract-iso-b", "test-bound-contract-iso-b"), + ] + .into_iter() + .map(|(agent, label)| (agent.to_owned(), Scope::new(label).unwrap())) + .collect::<std::collections::HashMap<_, _>>(); + let plain = { + let connection = connection.clone(); + let memory = Arc::clone(&memory); + let directory = Arc::clone(&directory); + let scopes = scopes.clone(); + Arc::new(move |agent_id: &str| { + let _keep_directory = &directory; + let scope = scopes + .get(agent_id) + .cloned() + .ok_or_else(|| StorageError::invalid_input("test agent scope is not mapped"))?; + let native = ScopedStorage::new(memory.for_scope(&scope)?, scope.clone()); + let docs = Arc::new(SqliteDocuments::new(connection.clone(), scope.clone())) + as Arc<dyn DocumentStore>; + Ok(crate::testkit::with_documents(&native, docs)) + }) as AgentStorageFactory + }; + let conditioned = { + let connection = connection.clone(); + let memory = Arc::clone(&memory); + let directory = Arc::clone(&directory); + Arc::new(move |agent_id: &str, condition: &WriteCondition| { + let _keep_directory = &directory; + let scope = scopes + .get(agent_id) + .cloned() + .ok_or_else(|| StorageError::invalid_input("test agent scope is not mapped"))?; + let native = ScopedStorage::new(memory.for_scope(&scope)?, scope.clone()); + let docs = Arc::new(SqliteDocuments::conditioned( + connection.clone(), + scope.clone(), + condition.clone(), + )) as Arc<dyn DocumentStore>; + Ok(crate::testkit::with_documents(&native, docs)) + }) as ConditionedAgentStorageFactory + }; + (plain, conditioned) +} + +fn sqlite_stores() -> DriverSessionStores { + let (factory, conditioned) = conditioned_factories(); + DriverSessionStores::new(factory) + .unwrap() + .with_conditioned_storage_factory(conditioned) } #[test] @@ -99,7 +166,10 @@ async fn bootstrap_preparation_is_idempotent_and_preserves_native_session_data() .doc["value"], "retained" ); - let provider = DriverSessionStores::new(storage_factory(Arc::clone(&backend))).unwrap(); + let (factory, conditioned) = conditioned_factories(); + let provider = DriverSessionStores::new(factory) + .unwrap() + .with_conditioned_storage_factory(conditioned); session_store_conformance(&provider).await; session_store_isolation_conformance(&provider).await; } @@ -132,23 +202,23 @@ async fn opening_declares_every_session_collection_before_migration() { } #[tokio::test] -async fn memory_backed_stores_meet_the_session_store_contract() { - session_store_conformance(&memory()).await; +async fn driver_stores_meet_the_session_store_contract() { + session_store_conformance(&sqlite_stores()).await; } #[tokio::test] async fn memory_backed_stores_keep_agents_apart() { - session_store_isolation_conformance(&memory()).await; + session_store_isolation_conformance(&sqlite_stores()).await; } #[tokio::test(flavor = "multi_thread")] async fn stores_work_from_a_multi_threaded_runtime() { - session_store_conformance(&memory()).await; + session_store_conformance(&sqlite_stores()).await; } #[test] fn stores_work_outside_any_runtime() { - let stores = memory().for_agent("plain"); + let stores = sqlite_stores().for_agent("plain"); let turn = TurnState::started("t", "r", 8, "2026-01-01T00:00:00Z"); stores.turn_states.put(&turn).unwrap(); assert_eq!(stores.turn_states.get("t").unwrap(), Some(turn)); @@ -186,8 +256,8 @@ fn host_bound_factory_keeps_agents_in_their_supplied_scopes() { #[test] fn two_backends_never_share_a_destination() { - let one = memory().for_agent("a"); - let two = memory().for_agent("a"); + let one = sqlite_stores().for_agent("a"); + let two = sqlite_stores().for_agent("a"); assert_ne!( one.transcripts.destination_key(), two.transcripts.destination_key() @@ -196,7 +266,7 @@ fn two_backends_never_share_a_destination() { #[test] fn an_agent_gets_the_same_stores_every_time() { - let provider = memory(); + let provider = sqlite_stores(); let first = provider.for_agent("a"); let again = provider.for_agent("a"); assert!(Arc::ptr_eq(&first.turn_states, &again.turn_states)); @@ -238,7 +308,7 @@ fn data_survives_a_new_provider_on_the_same_backend() { #[test] fn recover_marks_open_agents_turns_interrupted() { - let provider = memory(); + let provider = sqlite_stores(); let stores = provider.for_agent("a"); stores .turn_states @@ -771,8 +841,8 @@ async fn native_provider_returns_stable_bound_ledger_and_rebinds_actor_independe ); } -async fn prepared_restricted_actor_contract(backend: Arc<dyn StorageBackend>) { - let factory = storage_factory(Arc::clone(&backend)); +async fn prepared_restricted_actor_contract() { + let (factory, conditioned) = conditioned_factories(); let bootstrap = factory("prepared-actor").unwrap(); DriverSessionStores::prepare_schema(bootstrap.documents().as_ref()) .await @@ -780,6 +850,7 @@ async fn prepared_restricted_actor_contract(backend: Arc<dyn StorageBackend>) { let actor = DriverSessionStores::new_with_schema_admission(factory, SchemaAdmission::ValidateOnly) .unwrap() + .with_conditioned_storage_factory(conditioned) .recover_on_open(true); let opened = actor.try_for_agent("prepared-actor"); assert!( @@ -792,8 +863,8 @@ async fn prepared_restricted_actor_contract(backend: Arc<dyn StorageBackend>) { } #[tokio::test] -async fn prepared_restricted_actor_recovers_and_uses_memory_stores() { - prepared_restricted_actor_contract(Arc::new(MemoryStorage::new())).await; +async fn prepared_restricted_actor_recovers_and_uses_conditioned_stores() { + prepared_restricted_actor_contract().await; } #[tokio::test] diff --git a/crates/tinyagents-session/src/port/drivers/rebound.rs b/crates/tinyagents-session/src/port/drivers/rebound.rs index 0391d9218..ce5f74379 100644 --- a/crates/tinyagents-session/src/port/drivers/rebound.rs +++ b/crates/tinyagents-session/src/port/drivers/rebound.rs @@ -13,7 +13,7 @@ use tinyagents_harness::Result as HarnessResult; use tinyagents_harness::store::{AppendStore, Store}; use tinystoragedrivers_types::StorageError; -use super::{AgentStores, SessionStoreProvider}; +use super::{AgentStores, SessionStoreLease, SessionStoreLeaseError, SessionStoreProvider}; use crate::DriverSessionStores; use crate::port::AgentStorageFactory; use crate::port::{TurnState, TurnStates}; @@ -387,6 +387,10 @@ impl TranscriptLocator for LiveTranscriptLocator { /// Each resolved trait object is stable, while its backing stores are replaced /// atomically when the owning provider rebinds. impl SessionStoreProvider for RebindableSessionStores { + fn lease_for_agent(&self, agent_id: &str) -> Result<SessionStoreLease, SessionStoreLeaseError> { + self.current().lease_for_agent(agent_id) + } + fn for_agent(&self, agent_id: &str) -> AgentStores { let mut live = self.live.lock().unwrap_or_else(PoisonError::into_inner); live.retain(|_, weak| weak.strong_count() > 0); diff --git a/crates/tinyagents-session/src/port/drivers/transcripts.rs b/crates/tinyagents-session/src/port/drivers/transcripts.rs index e1dabb086..4ecca743e 100644 --- a/crates/tinyagents-session/src/port/drivers/transcripts.rs +++ b/crates/tinyagents-session/src/port/drivers/transcripts.rs @@ -36,11 +36,11 @@ use std::sync::Arc; use serde::{Deserialize, Serialize}; use serde_json::{Value, json}; use tinyagents_harness::store::SchemaAdmission; -use tinyagents_storage_kit::primitives::{Leases, Repo}; +use tinyagents_storage_kit::primitives::{Lease, Leases, Repo}; use tinyagents_storage_kit::{Blocking, DocumentStoreExt}; use tinyagents_storage_kit::{ - CollectionSpec, ErrorKind, Filter, ID_FIELD, IndexSpec, Key, Query, Sort, StorageError, - WriteOptions, + CollectionSpec, ErrorKind, Filter, ID_FIELD, IndexSpec, Key, Query, Scope, Sort, StorageError, + WriteCondition, WriteOptions, }; use tinystoragedrivers_api::DocumentStore; use tokio::sync::{Mutex, OnceCell}; @@ -302,6 +302,8 @@ pub struct DriverTranscriptHistory { struct HistoryInner { docs: Arc<dyn DocumentStore>, + scope: Option<Scope>, + conditioned_entries: Option<ConditionedEntryOpener>, tables: Arc<Tables>, declared: Arc<Declared>, stem: String, @@ -309,6 +311,9 @@ struct HistoryInner { replay: Mutex<Replay>, } +pub(crate) type ConditionedEntryOpener = + Arc<dyn Fn(&WriteCondition) -> Result<Arc<dyn DocumentStore>, StorageError> + Send + Sync>; + impl std::fmt::Debug for DriverTranscriptHistory { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("DriverTranscriptHistory") @@ -436,20 +441,47 @@ impl HistoryInner { let Some(entry) = build(&replay, &self.seed)? else { return Ok(false); }; - self.claim_reservation(&mut replay).await?; + let claimed = self.claim_reservation(&mut replay).await?; let seq = replay.next_seq; let mut doc = serde_json::to_value(&entry).map_err(|error| serialization(&error))?; if let Value::Object(fields) = &mut doc { fields.insert("stem".into(), json!(self.stem)); fields.insert("seq".into(), json!(seq)); } - match self - .tables - .entries - .insert(&entry_id(&self.stem, seq)?, &doc) - .await - { + let reserved = claimed.is_some(); + let insert = if let Some(lease) = claimed { + let opener = self.conditioned_entries.as_ref().ok_or_else(|| { + anyhow::anyhow!( + "session generation {} requires atomic lease-conditioned entry writes", + self.stem + ) + })?; + let scope = self.scope.as_ref().ok_or_else(|| { + anyhow::anyhow!( + "session generation {} has no bound storage scope", + self.stem + ) + })?; + let condition = lease.write_condition(scope.clone())?; + let docs = opener(&condition)?; + Repo::new(&docs, ENTRIES)? + .with_options(WriteOptions::new().schema(SCHEMA_VERSION)) + .insert(&entry_id(&self.stem, seq)?, &doc) + .await + } else { + self.tables + .entries + .insert(&entry_id(&self.stem, seq)?, &doc) + .await + }; + match insert { Ok(_) => { + // Keep the reservation on a failed insert so retries must + // renew the same epoch and can never fall back to an + // unconditioned write after a takeover. + if reserved { + replay.reservation = None; + } replay.apply(seq, entry); // The entry is durable: report the write as done. A // failed index refresh is retried a few times here, then @@ -491,13 +523,13 @@ impl HistoryInner { /// process took the reservation over in the meantime (this one stalled /// past the reservation's lifetime), the generation is theirs and this /// handle's writes are refused instead of overwriting it. - async fn claim_reservation(&self, replay: &mut Replay) -> anyhow::Result<()> { + async fn claim_reservation(&self, replay: &mut Replay) -> anyhow::Result<Option<Lease>> { let Some(held) = replay.reservation.as_ref() else { - return Ok(()); + return Ok(None); }; - reservation::claim(&self.tables.leases, &self.stem, held).await?; - replay.reservation = None; - Ok(()) + Ok(Some( + reservation::claim(&self.tables.leases, &self.stem, held).await?, + )) } /// Releases the reservation of a generation this handle never wrote. @@ -764,6 +796,8 @@ impl TranscriptHistory for DriverTranscriptHistory { #[derive(Clone)] pub struct DriverTranscriptLocator { docs: Arc<dyn DocumentStore>, + scope: Option<Scope>, + conditioned_entries: Option<ConditionedEntryOpener>, tables: Arc<Tables>, bridge: Blocking, label: String, @@ -779,6 +813,11 @@ impl std::fmt::Debug for DriverTranscriptLocator { } impl DriverTranscriptLocator { + #[cfg(test)] + pub(crate) fn documents(&self) -> &Arc<dyn DocumentStore> { + &self.docs + } + /// Transcripts in `docs`, with synchronous calls run on `bridge`. /// /// `label` names where `docs` points (driver, database, scope). It is @@ -795,10 +834,43 @@ impl DriverTranscriptLocator { bridge: Blocking, label: impl Into<String>, schema_admission: SchemaAdmission, + ) -> Self { + Self::with_conditioned_entry_opener(docs, bridge, label, schema_admission, None, None) + } + + /// Transcripts whose first entry in a reserved generation is opened with + /// the lease's native atomic write condition. + pub(crate) fn new_with_conditioned_entry_opener( + docs: Arc<dyn DocumentStore>, + bridge: Blocking, + label: impl Into<String>, + schema_admission: SchemaAdmission, + scope: Scope, + opener: ConditionedEntryOpener, + ) -> Self { + Self::with_conditioned_entry_opener( + docs, + bridge, + label, + schema_admission, + Some(scope), + Some(opener), + ) + } + + fn with_conditioned_entry_opener( + docs: Arc<dyn DocumentStore>, + bridge: Blocking, + label: impl Into<String>, + schema_admission: SchemaAdmission, + scope: Option<Scope>, + conditioned_entries: Option<ConditionedEntryOpener>, ) -> Self { Self { tables: Arc::new(Tables::new(&docs)), docs, + scope, + conditioned_entries, bridge, label: label.into(), declared: Arc::new(Declared { @@ -823,6 +895,8 @@ impl DriverTranscriptLocator { DriverTranscriptHistory { inner: Arc::new(HistoryInner { docs: Arc::clone(&self.docs), + scope: self.scope.clone(), + conditioned_entries: self.conditioned_entries.clone(), tables: Arc::clone(&self.tables), declared: Arc::clone(&self.declared), stem: stem.to_string(), diff --git a/crates/tinyagents-session/src/port/drivers/transcripts_adoption.rs b/crates/tinyagents-session/src/port/drivers/transcripts_adoption.rs index 011623a42..7b94696cb 100644 --- a/crates/tinyagents-session/src/port/drivers/transcripts_adoption.rs +++ b/crates/tinyagents-session/src/port/drivers/transcripts_adoption.rs @@ -136,6 +136,12 @@ impl DriverTranscriptLocator { if self.session_exists(session) { return Ok(None); } + // Legacy roots predate delegated session lineage. A child has no + // independent root identity to adopt, and treating its absent + // `agent_id` as a wildcard would copy the parent's history into it. + if session.is_subagent() { + return Ok(None); + } let stem = session_stem(session); let candidates = self.adoption_candidates(session, thread_id, &stem)?; diff --git a/crates/tinyagents-session/src/port/drivers/transcripts_adoption_tests.rs b/crates/tinyagents-session/src/port/drivers/transcripts_adoption_tests.rs index dc2d07118..7d2ea6211 100644 --- a/crates/tinyagents-session/src/port/drivers/transcripts_adoption_tests.rs +++ b/crates/tinyagents-session/src/port/drivers/transcripts_adoption_tests.rs @@ -1,22 +1,17 @@ use super::*; -use crate::testkit::MemoryStorage; use std::sync::Arc; +use super::super::{leases_spec, reservation::RESERVATION_TTL_MS}; use crate::port::drivers::transcripts::tests::meta; use crate::transcript::TranscriptMessage; -use tinyagents_storage_kit::{Blocking, DocumentStore, Scope}; +use crate::transcript::test_locator::conditioned_driver_locator; +use tinyagents_storage_kit::Blocking; +use tinyagents_storage_kit::primitives::Leases; const THREAD: &str = "t"; fn locator() -> DriverTranscriptLocator { - let storage = MemoryStorage::new(); - let docs: Arc<dyn DocumentStore> = Arc::clone( - storage - .for_scope(&Scope::new("local").unwrap()) - .unwrap() - .documents(), - ); - DriverTranscriptLocator::new(docs, Blocking::new().unwrap(), "memory://adoption") + conditioned_driver_locator() } /// A root of `THREAD` under `stem`, written with `meta` (as an import does). @@ -218,15 +213,18 @@ fn a_reserved_stem_reports_nothing_adopted() { let session = SessionRef::scoped(THREAD, "planner"); // A concurrent writer holds the session's stem. let stem = session_stem(&session); - let docs = Arc::clone(&locator.docs); - let tables = Arc::clone(&locator.tables); - let declared = Arc::clone(&locator.declared); - run_on(&locator.bridge, async move { - declared.ensure(&docs).await?; - Ok(reservation::reserve(&docs, &tables.leases, &stem).await) - }) - .unwrap() - .unwrap(); + let docs = Arc::clone(locator.documents()); + Blocking::new() + .unwrap() + .run(async move { + docs.ensure_collection(&leases_spec()).await?; + Leases::new(&docs) + .acquire("session", &stem, "concurrent-adopter", RESERVATION_TTL_MS) + .await + }) + .unwrap() + .unwrap() + .unwrap(); assert!( locator .adopt_legacy(&session, THREAD, &meta(THREAD)) @@ -289,7 +287,7 @@ fn a_root_missing_from_the_index_is_still_folded() { &["two"], ); // The second write was durable but every index refresh after it failed. - let docs = Arc::clone(&locator.docs); + let docs = Arc::clone(locator.documents()); Blocking::new() .unwrap() .run(async move { diff --git a/crates/tinyagents-session/src/port/drivers/transcripts_reservation.rs b/crates/tinyagents-session/src/port/drivers/transcripts_reservation.rs index 631812c7d..dc43cee51 100644 --- a/crates/tinyagents-session/src/port/drivers/transcripts_reservation.rs +++ b/crates/tinyagents-session/src/port/drivers/transcripts_reservation.rs @@ -17,7 +17,7 @@ use std::sync::Arc; -use tinyagents_storage_kit::primitives::{self, Leases}; +use tinyagents_storage_kit::primitives::{self, Lease, Leases}; use tinyagents_storage_kit::{CollectionSpec, Filter, StorageError}; use tinystoragedrivers_api::DocumentStore; @@ -97,7 +97,7 @@ pub(super) async fn claim( leases: &Leases, stem: &str, reservation: &Reservation, -) -> anyhow::Result<()> { +) -> anyhow::Result<Lease> { let renewed = leases .renew( LEASE_OWNER, @@ -107,11 +107,11 @@ pub(super) async fn claim( RESERVATION_TTL_MS, ) .await?; - anyhow::ensure!( - renewed.is_some(), - "session generation {stem} was taken over by another writer; reload the session" - ); - Ok(()) + renewed.ok_or_else(|| { + anyhow::anyhow!( + "session generation {stem} was taken over by another writer; reload the session" + ) + }) } /// Gives `reservation` back. A reservation another process has since taken diff --git a/crates/tinyagents-session/src/port/drivers/transcripts_tests.rs b/crates/tinyagents-session/src/port/drivers/transcripts_tests.rs index 87a3b1c6a..b6b162b54 100644 --- a/crates/tinyagents-session/src/port/drivers/transcripts_tests.rs +++ b/crates/tinyagents-session/src/port/drivers/transcripts_tests.rs @@ -1,7 +1,8 @@ use super::*; use crate::testkit::MemoryStorage; use crate::testkit::conformance::transcript_history_conformance; -use tinyagents_storage_kit::{Precondition, Scope, WriteOptions}; +use tinyagents_storage_kit::{Precondition, Scope, WriteCondition, WriteOptions}; +use tinystoragedrivers_types::{Filter, Query, Sort}; pub(in crate::port::drivers) fn meta(thread_id: &str) -> TranscriptMeta { let mut meta = discovered_seed("planner"); @@ -17,18 +18,58 @@ fn message(role: &str, content: &str) -> TranscriptMessage { serde_json::from_value(json!({ "role": role, "content": content })).unwrap() } -fn docs() -> Arc<dyn DocumentStore> { - let storage = MemoryStorage::new(); - Arc::clone( - storage - .for_scope(&Scope::new("local").unwrap()) - .unwrap() - .documents(), - ) +struct TestDocs { + docs: Arc<dyn DocumentStore>, + connection: tinystoragedrivers_connect::sqlite::SqliteConnection, + scope: Scope, + _directory: tempfile::TempDir, +} + +impl std::ops::Deref for TestDocs { + type Target = Arc<dyn DocumentStore>; + + fn deref(&self) -> &Self::Target { + &self.docs + } +} + +fn docs() -> TestDocs { + use tinystoragedrivers_connect::sqlite::SqliteConnection; + use tinystoragedrivers_document::sqlite::SqliteDocuments; + + let directory = tempfile::tempdir().unwrap(); + let connection = SqliteConnection::open(directory.path().join("transcripts.db")).unwrap(); + let scope = Scope::new("local").unwrap(); + let docs = + Arc::new(SqliteDocuments::new(connection.clone(), scope.clone())) as Arc<dyn DocumentStore>; + TestDocs { + docs, + connection, + scope, + _directory: directory, + } } -fn locator(docs: &Arc<dyn DocumentStore>) -> DriverTranscriptLocator { - DriverTranscriptLocator::new(Arc::clone(docs), Blocking::new().unwrap(), "memory://test") +fn locator(docs: &TestDocs) -> DriverTranscriptLocator { + use tinystoragedrivers_document::sqlite::SqliteDocuments; + + let connection = docs.connection.clone(); + let scope = docs.scope.clone(); + let opener: ConditionedEntryOpener = Arc::new(move |condition| { + Ok(Arc::new(SqliteDocuments::conditioned( + connection.clone(), + scope.clone(), + condition.clone(), + )) as Arc<dyn DocumentStore>) + }); + DriverTranscriptLocator::new_with_conditioned_entry_opener( + Arc::clone(&docs.docs), + Blocking::new().unwrap(), + "memory://test", + tinyagents_harness::store::SchemaAdmission::Declare, + docs.scope.clone(), + opener, + ) } /// Lets the reservation of `stem` lapse, as if its owner stopped long ago. @@ -95,6 +136,33 @@ fn on_bridge<T: Send + 'static>( Blocking::new().unwrap().run(future).unwrap().unwrap() } +/// Expires and takes over a reserved successor while the old handle is +/// opening its conditioned entry store, after claim and immediately before +/// the first insert. +fn take_over_reservation(docs: Arc<dyn DocumentStore>, stem: String) { + on_bridge(async move { + let id = Key::of(["session", stem.as_str()])?; + let lease = docs + .get("tsd.leases", id.as_str()) + .await? + .expect("successor reservation exists"); + let mut body = lease.doc; + body["expires_ms"] = json!(0); + docs.put( + "tsd.leases", + id.as_str(), + body, + Precondition::Version(lease.version), + ) + .await?; + Leases::new(&docs) + .acquire("session", &stem, "replacement-holder", 30_000) + .await? + .expect("expired reservation can be taken over"); + Ok(()) + }); +} + #[test] fn a_handle_meets_the_transcript_history_contract() { let docs = docs(); @@ -102,6 +170,74 @@ fn a_handle_meets_the_transcript_history_contract() { transcript_history_conformance(history.as_ref()); } +#[test] +fn a_reserved_generation_rejects_entry_insert_after_lease_takeover() { + use tinystoragedrivers_connect::sqlite::SqliteConnection; + use tinystoragedrivers_document::sqlite::SqliteDocuments; + + let directory = tempfile::tempdir().unwrap(); + let connection = SqliteConnection::open(directory.path().join("transcripts.db")).unwrap(); + let scope = Scope::new("local").unwrap(); + let base = + Arc::new(SqliteDocuments::new(connection.clone(), scope.clone())) as Arc<dyn DocumentStore>; + let armed = Arc::new(std::sync::atomic::AtomicBool::new(false)); + let (conditioned_open_tx, conditioned_open_rx) = std::sync::mpsc::channel(); + let (takeover_done_tx, takeover_done_rx) = std::sync::mpsc::channel(); + let takeover_done_rx = std::sync::Mutex::new(takeover_done_rx); + let opener_scope = scope.clone(); + let opener: ConditionedEntryOpener = { + let armed = Arc::clone(&armed); + Arc::new(move |condition: &WriteCondition| { + if armed.swap(false, std::sync::atomic::Ordering::SeqCst) { + conditioned_open_tx.send(()).unwrap(); + takeover_done_rx.lock().unwrap().recv().unwrap(); + } + Ok(Arc::new(SqliteDocuments::conditioned( + connection.clone(), + opener_scope.clone(), + condition.clone(), + )) as Arc<dyn DocumentStore>) + }) + }; + let locator = DriverTranscriptLocator::new_with_conditioned_entry_opener( + Arc::clone(&base), + Blocking::new().unwrap(), + "sqlite://transcript-reservation-race", + tinyagents_harness::store::SchemaAdmission::Declare, + scope, + opener, + ); + let parent = SessionRef::scoped("reservation-race", "planner"); + let (successor, history) = locator + .begin_generation(&parent, meta("reservation-race")) + .unwrap(); + let stem = session_stem(&successor); + armed.store(true, std::sync::atomic::Ordering::SeqCst); + + let writer = std::thread::spawn(move || { + let rows = vec![message("user", "first successor write")]; + history.append_turn(turn(&[], &rows, &meta("reservation-race"))) + }); + conditioned_open_rx.recv().unwrap(); + take_over_reservation(Arc::clone(&base), stem.clone()); + takeover_done_tx.send(()).unwrap(); + let result = writer.join().unwrap(); + + assert!( + result.is_err(), + "the old reservation holder must not insert after another holder takes over" + ); + assert_eq!( + entries(&base, &stem).len(), + 0, + "the stale insert must leave no entry" + ); + assert!( + reservation_epoch(&base, &stem) > 1, + "the test took over the lease" + ); +} + #[test] fn an_ordinary_turn_stores_only_its_new_rows() { let docs = docs(); diff --git a/crates/tinyagents-session/src/port/memory/mod.rs b/crates/tinyagents-session/src/port/memory/mod.rs index a8d9b2d93..97063030b 100644 --- a/crates/tinyagents-session/src/port/memory/mod.rs +++ b/crates/tinyagents-session/src/port/memory/mod.rs @@ -16,7 +16,7 @@ use tinyagents_harness::store::{InMemoryAppendStore, InMemoryStore}; use super::helpers::{ COMPLETED_RETENTION, MAX_GENERATIONS, completed_newest_first, is_terminal, newest_first, }; -use super::{AgentStores, SessionStoreProvider, TurnStates}; +use super::{AgentStores, SessionStoreLease, SessionStoreProvider, TurnStates}; use crate::testkit::InMemoryTranscriptHistory; use crate::transcript::{ SessionRef, ThreadTranscript, ThreadTranscripts, TranscriptHistory, TranscriptLocator, @@ -65,6 +65,13 @@ impl SessionStoreProvider for InMemorySessionStores { .clone() } + fn lease_for_agent( + &self, + agent_id: &str, + ) -> Result<SessionStoreLease, super::SessionStoreLeaseError> { + Ok(SessionStoreLease::pinned(self.for_agent(agent_id))) + } + fn recover(&self) -> anyhow::Result<()> { let now = chrono::Utc::now().to_rfc3339(); let agents: Vec<AgentStores> = self diff --git a/crates/tinyagents-session/src/port/memory/mod_tests.rs b/crates/tinyagents-session/src/port/memory/mod_tests.rs index a786abbaa..a8733ea08 100644 --- a/crates/tinyagents-session/src/port/memory/mod_tests.rs +++ b/crates/tinyagents-session/src/port/memory/mod_tests.rs @@ -54,6 +54,24 @@ fn an_agent_gets_the_same_stores_every_time() { ); } +#[test] +fn in_memory_stores_can_lease_their_stable_agent_handles() { + let provider = InMemorySessionStores::new(); + let opened = provider.for_agent("a"); + let lease = provider + .lease_for_agent("a") + .expect("the process-lifetime memory provider pins its handles"); + + assert!(Arc::ptr_eq( + &opened.turn_states, + &lease.stores().turn_states + )); + assert_eq!( + opened.transcripts.destination_key(), + lease.stores().transcripts.destination_key() + ); +} + #[test] fn clearing_in_memory_history_clears_display_partials() { let history = InMemoryTranscriptHistory::new("clear-partial", test_meta("clear-partial")); diff --git a/crates/tinyagents-session/src/port/mod.rs b/crates/tinyagents-session/src/port/mod.rs index d4c96fc14..4713cb674 100644 --- a/crates/tinyagents-session/src/port/mod.rs +++ b/crates/tinyagents-session/src/port/mod.rs @@ -59,7 +59,7 @@ pub use drivers::{ }; #[cfg(any(test, feature = "testkit"))] pub use memory::{InMemorySessionStores, InMemoryTranscriptLocator, InMemoryTurnStates}; -pub use types::AgentStores; +pub use types::{AgentStores, SessionStoreLease, SessionStoreLeaseError}; pub use tinyagents_harness::store::{AppendStore, Store}; @@ -75,6 +75,23 @@ pub type AgentStorageFactory = Arc< + Sync, >; +/// Host-owned opener for storage conditioned on an exact durable lease. +/// +/// The session driver uses this only for the first entry written to a +/// reserved transcript generation. The returned scope must be the same +/// agent scope supplied by [`AgentStorageFactory`], and the backend must +/// enforce the condition atomically on every mutation. Backends that cannot +/// provide that guarantee must return an error. +pub type ConditionedAgentStorageFactory = Arc< + dyn Fn( + &str, + &tinyagents_storage_kit::WriteCondition, + ) + -> Result<tinyagents_storage_kit::ScopedStorage, tinystoragedrivers_types::StorageError> + + Send + + Sync, +>; + use crate::turn_state::{TurnLifecycle, TurnState}; /// Hands out each agent's stores. Installed once per runtime by the host. @@ -101,6 +118,15 @@ pub trait SessionStoreProvider: Send + Sync { Ok(self.for_agent(agent_id)) } + /// Opens stores whose scoped storage remains pinned to this provider's + /// current binding for the lifetime of the returned handle. + fn lease_for_agent( + &self, + _agent_id: &str, + ) -> Result<SessionStoreLease, SessionStoreLeaseError> { + Err(SessionStoreLeaseError::Unsupported) + } + /// Optionally create an independent actor provider over a host-admitted backend. /// /// The backend must already embody the actor authority supplied by the @@ -306,6 +332,10 @@ impl<P: SessionStoreProvider + ?Sized> SessionStoreProvider for Arc<P> { (**self).try_for_agent(agent_id) } + fn lease_for_agent(&self, agent_id: &str) -> Result<SessionStoreLease, SessionStoreLeaseError> { + (**self).lease_for_agent(agent_id) + } + fn bind_actor( &self, factory: AgentStorageFactory, diff --git a/crates/tinyagents-session/src/port/types.rs b/crates/tinyagents-session/src/port/types.rs index f02bf6ea3..3227ec759 100644 --- a/crates/tinyagents-session/src/port/types.rs +++ b/crates/tinyagents-session/src/port/types.rs @@ -25,6 +25,34 @@ pub struct AgentStores { pub journal: Arc<dyn AppendStore>, } +/// A set of session stores whose scoped storage was captured at open time. +#[derive(Clone)] +pub struct SessionStoreLease { + stores: AgentStores, +} + +impl SessionStoreLease { + /// The stores pinned by this lease. + pub fn stores(&self) -> &AgentStores { + &self.stores + } + + pub(crate) fn pinned(stores: AgentStores) -> Self { + Self { stores } + } +} + +/// Why a provider could not return a pinned set of stores. +#[derive(Debug, thiserror::Error)] +pub enum SessionStoreLeaseError { + /// The provider cannot guarantee that its handles retain one backend. + #[error("provider cannot guarantee pinned session stores")] + Unsupported, + /// Opening the pinned stores failed. + #[error("opening pinned session stores failed: {0}")] + Open(#[source] tinystoragedrivers_types::StorageError), +} + impl std::fmt::Debug for AgentStores { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("AgentStores") diff --git a/crates/tinyagents-session/src/testkit/conformance.rs b/crates/tinyagents-session/src/testkit/conformance.rs index 5cc49a1b1..cd096b856 100644 --- a/crates/tinyagents-session/src/testkit/conformance.rs +++ b/crates/tinyagents-session/src/testkit/conformance.rs @@ -740,6 +740,32 @@ pub async fn session_store_conformance(provider: &dyn SessionStoreProvider) { ); } +/// Runs the explicit pinned-lease contract against a provider that supports it. +/// +/// This checks that a lease opens usable stores without depending on any +/// provider-specific storage implementation. Providers that cannot guarantee +/// pinned handles are covered by the fail-closed `Unsupported` contract test. +pub async fn session_store_lease_conformance(provider: &dyn SessionStoreProvider) { + let lease = provider + .lease_for_agent("contract-lease-agent") + .expect("provider supports pinned session-store leases"); + lease + .stores() + .kv + .put("contract", "lease", serde_json::json!("pinned")) + .await + .expect("leased key-value store writes"); + assert_eq!( + lease + .stores() + .kv + .get("contract", "lease") + .await + .expect("leased key-value store reads"), + Some(serde_json::json!("pinned")) + ); +} + /// Checks that `provider` keeps agents apart: whatever one agent writes, /// another agent's stores cannot see. /// diff --git a/crates/tinyagents-session/src/testkit/mod.rs b/crates/tinyagents-session/src/testkit/mod.rs index 95d11663a..b27da6117 100644 --- a/crates/tinyagents-session/src/testkit/mod.rs +++ b/crates/tinyagents-session/src/testkit/mod.rs @@ -1,8 +1,7 @@ //! Session-persistence test building blocks. //! //! This is the session-layer counterpart to the harness -//! [`testkit`](tinyagents_harness::testkit) and the graph -//! [`testkit`](../../tinyagents_graph/testkit/index.html): reusable doubles and +//! [`testkit`](tinyagents_harness::testkit): reusable doubles and //! **conformance** (contract) suites so a downstream backend author can prove //! their implementation behaves like the bundled ones, and the bundled ones //! can be pinned against each other. diff --git a/crates/tinyagents-session/src/transcript.rs b/crates/tinyagents-session/src/transcript.rs index 3a1aca15b..417f69657 100644 --- a/crates/tinyagents-session/src/transcript.rs +++ b/crates/tinyagents-session/src/transcript.rs @@ -93,7 +93,9 @@ pub mod view; #[cfg(test)] mod writer; -pub use adoption::SessionAdoption; +pub use adoption::{ + SessionAdoption, repair_orphaned_tool_head, repair_orphaned_tool_head_messages, +}; pub use background::append_background_message_in; pub use display::BACKGROUND_METADATA_KEY; pub(crate) use display::{display_from_session, display_message_of, partial_display_message}; diff --git a/crates/tinyagents-session/src/transcript/adoption.rs b/crates/tinyagents-session/src/transcript/adoption.rs index 07e8bca88..6033f9e30 100644 --- a/crates/tinyagents-session/src/transcript/adoption.rs +++ b/crates/tinyagents-session/src/transcript/adoption.rs @@ -1,8 +1,62 @@ //! The result of folding a conversation's pre-identity transcripts into its //! session ([`TranscriptLocator::adopt_legacy`](super::TranscriptLocator::adopt_legacy)). +use std::borrow::Cow; use std::path::PathBuf; +use super::TranscriptMessage; +use tinyagents_harness::tinyinference_llm::message::Message; + +/// Removes only the contiguous orphan tool-result head following the system prefix. +/// +/// This is a replay view: durable rows are left untouched, and later tool rows +/// remain visible to strict transcript validation. +pub fn repair_orphaned_tool_head(rows: &[TranscriptMessage]) -> Cow<'_, [TranscriptMessage]> { + let prefix = rows.iter().take_while(|row| row.role == "system").count(); + let orphans = rows[prefix..] + .iter() + .take_while(|row| row.role == "tool") + .count(); + if orphans == 0 { + Cow::Borrowed(rows) + } else { + Cow::Owned( + rows[..prefix] + .iter() + .chain(&rows[prefix + orphans..]) + .cloned() + .collect(), + ) + } +} + +/// Removes only the contiguous orphan tool-result head from decoded replay history. +pub fn repair_orphaned_tool_head_messages(history: &[Message]) -> Cow<'_, [Message]> { + let prefix = history + .iter() + .take_while(|message| matches!(message, Message::System(_))) + .count(); + let orphans = history[prefix..] + .iter() + .take_while(|message| matches!(message, Message::Tool(_))) + .count(); + if orphans == 0 { + Cow::Borrowed(history) + } else { + Cow::Owned( + history[..prefix] + .iter() + .chain(&history[prefix + orphans..]) + .cloned() + .collect(), + ) + } +} + +#[cfg(test)] +#[path = "adoption_tests.rs"] +mod tests; + /// What adoption did for one session. #[derive(Debug, Clone, PartialEq, Eq)] pub struct SessionAdoption { diff --git a/crates/tinyagents-session/src/transcript/adoption_tests.rs b/crates/tinyagents-session/src/transcript/adoption_tests.rs new file mode 100644 index 000000000..7148f8fbc --- /dev/null +++ b/crates/tinyagents-session/src/transcript/adoption_tests.rs @@ -0,0 +1,60 @@ +use super::*; +use crate::transcript::{TranscriptMessage, repair_orphaned_tool_head}; +use std::borrow::Cow; +use tinyagents_harness::tinyinference_llm::message::Message; + +#[test] +fn repair_removes_only_the_leading_tool_run_after_system_prefix() { + let rows = vec![ + TranscriptMessage::system("system"), + TranscriptMessage::tool("orphan one").with_id("one"), + TranscriptMessage::tool("orphan two").with_id("two"), + TranscriptMessage::user("keep"), + TranscriptMessage::tool("later orphan").with_id("later"), + ]; + let repaired = repair_orphaned_tool_head(&rows); + assert!(matches!(repaired, Cow::Owned(_))); + assert_eq!(repaired.len(), 3); + assert_eq!(repaired[0].role, "system"); + assert_eq!(repaired[1].content, "keep"); + assert_eq!(repaired[2].id.as_deref(), Some("later")); + assert_eq!(rows.len(), 5); + assert_eq!(rows[1].id.as_deref(), Some("one")); +} + +#[test] +fn repair_preserves_nonleading_tool_rows_and_borrows_without_orphans() { + let rows = vec![ + TranscriptMessage::system("system"), + TranscriptMessage::user("user"), + TranscriptMessage::tool("later").with_id("later"), + ]; + assert!(matches!(repair_orphaned_tool_head(&rows), Cow::Borrowed(_))); +} + +#[test] +fn message_repair_matches_transcript_repair_and_preserves_later_orphans() { + let history = vec![ + Message::system("system"), + Message::tool("call-one", "orphan one"), + Message::tool("call-two", "orphan two"), + Message::user("keep"), + Message::tool("later", "later orphan"), + ]; + let repaired = repair_orphaned_tool_head_messages(&history); + assert!(matches!(repaired, Cow::Owned(_))); + assert_eq!(repaired.len(), 3); + assert!(matches!(repaired[0], Message::System(_))); + assert!(matches!(repaired[1], Message::User(_))); + assert!(matches!(repaired[2], Message::Tool(_))); + assert_eq!(history.len(), 5); +} + +#[test] +fn message_repair_borrows_history_without_a_leading_orphan() { + let history = vec![Message::system("system"), Message::user("user")]; + assert!(matches!( + repair_orphaned_tool_head_messages(&history), + Cow::Borrowed(_) + )); +} diff --git a/crates/tinyagents-session/src/transcript/session_tests.rs b/crates/tinyagents-session/src/transcript/session_tests.rs index 28beb5d1f..47e8ccd63 100644 --- a/crates/tinyagents-session/src/transcript/session_tests.rs +++ b/crates/tinyagents-session/src/transcript/session_tests.rs @@ -19,6 +19,15 @@ fn a_stem_is_deterministic_and_carries_no_timestamp() { ); } +#[test] +fn canonical_identity_is_stable_and_distinct_from_legacy_stems() { + let first = SessionRef::scoped("thread-legacy", "agent-legacy"); + let second = SessionRef::scoped("thread-legacy", "agent-legacy"); + assert_eq!(session_stem(&first), session_stem(&second)); + assert_ne!(session_stem(&first), "1700000100_agent-legacy"); + assert_eq!(first.first_generation(), second.first_generation()); +} + #[test] fn two_agents_on_one_key_get_distinct_stems() { let left = session_stem(&SessionRef::scoped("thread-1", "orchestrator")); diff --git a/crates/tinyagents-session/src/transcript/test_locator.rs b/crates/tinyagents-session/src/transcript/test_locator.rs index c4774b095..1de18272e 100644 --- a/crates/tinyagents-session/src/transcript/test_locator.rs +++ b/crates/tinyagents-session/src/transcript/test_locator.rs @@ -14,9 +14,42 @@ use super::{ }; use crate::port::SessionStoreProvider; -/// An empty driver-backed locator (the production store, over memory). +/// An empty driver-backed locator with native conditional writes available. pub(crate) fn driver_locator() -> Arc<dyn TranscriptLocator> { - crate::testkit::in_memory_transcript() + Arc::new(conditioned_driver_locator()) +} + +pub(crate) fn conditioned_driver_locator() -> crate::port::DriverTranscriptLocator { + use crate::port::DriverTranscriptLocator; + use tinyagents_harness::store::SchemaAdmission; + use tinyagents_storage_kit::{Blocking, Scope, WriteCondition}; + use tinystoragedrivers_api::DocumentStore; + use tinystoragedrivers_connect::sqlite::SqliteConnection; + use tinystoragedrivers_document::sqlite::SqliteDocuments; + + let directory = tempfile::tempdir().expect("create transcript fixture database"); + let connection = SqliteConnection::open(directory.path().join("transcripts.db")) + .expect("open transcript fixture database"); + let scope = Scope::new("transcript-tests").expect("valid test scope"); + let opener_scope = scope.clone(); + let docs = + Arc::new(SqliteDocuments::new(connection.clone(), scope.clone())) as Arc<dyn DocumentStore>; + let opener = Arc::new(move |condition: &WriteCondition| { + let _keep_directory = &directory; + Ok(Arc::new(SqliteDocuments::conditioned( + connection.clone(), + opener_scope.clone(), + condition.clone(), + )) as Arc<dyn DocumentStore>) + }); + DriverTranscriptLocator::new_with_conditioned_entry_opener( + docs, + Blocking::new().expect("start transcript test bridge"), + "sqlite://transcript-tests", + SchemaAdmission::Declare, + scope, + opener, + ) } /// A fixture locator holding parsed transcript rows from `<workspace>`. diff --git a/crates/tinyagents-storage-kit/src/lib.rs b/crates/tinyagents-storage-kit/src/lib.rs index 452e83bb8..37e143240 100644 --- a/crates/tinyagents-storage-kit/src/lib.rs +++ b/crates/tinyagents-storage-kit/src/lib.rs @@ -48,6 +48,20 @@ impl ScopedStorage { /// Adapter for opening a scope-bound storage bundle from a port backend. pub trait StorageBackend: Send + Sync { fn for_scope(&self, scope: &tinystoragedrivers_types::Scope) -> error::Result<ScopedStorage>; + + /// Opens a scope whose mutations atomically require `condition` to hold. + /// Implementations that cannot provide this guarantee fail closed. + fn for_scope_conditioned( + &self, + scope: &tinystoragedrivers_types::Scope, + condition: &tinystoragedrivers_types::WriteCondition, + ) -> error::Result<ScopedStorage> { + let _ = (scope, condition); + Err(tinystoragedrivers_types::StorageError::unsupported( + tinystoragedrivers_types::Capability::ConditionalWrites, + "atomic condition-bound writes are unsupported by this backend", + )) + } } impl<T: PortStorageBackend> StorageBackend for T { @@ -57,6 +71,17 @@ impl<T: PortStorageBackend> StorageBackend for T { scope.clone(), )) } + + fn for_scope_conditioned( + &self, + scope: &tinystoragedrivers_types::Scope, + condition: &tinystoragedrivers_types::WriteCondition, + ) -> error::Result<ScopedStorage> { + Ok(ScopedStorage::new( + PortStorageBackend::for_scope_conditioned(self, scope, condition)?, + scope.clone(), + )) + } } pub mod blocking; diff --git a/crates/tinyagents-storage-kit/src/primitives/leases.rs b/crates/tinyagents-storage-kit/src/primitives/leases.rs index 8040e9c4b..84dfb68f0 100644 --- a/crates/tinyagents-storage-kit/src/primitives/leases.rs +++ b/crates/tinyagents-storage-kit/src/primitives/leases.rs @@ -18,7 +18,7 @@ use serde::{Deserialize, Serialize}; use super::LEASES; use super::repo::{Repo, RetryPolicy}; use tinystoragedrivers_api::DocumentStore; -use tinystoragedrivers_types::{Result, StorageError}; +use tinystoragedrivers_types::{Filter, Result, Scope, StorageError, WriteCondition}; /// Injectable Unix epoch-millisecond clock. pub type Clock = Arc<dyn Fn() -> u64 + Send + Sync>; @@ -149,6 +149,23 @@ impl Lease { endpoint: body.endpoint, } } + + /// Builds an atomic write condition for this exact lease tenure. + /// + /// The condition is tied to the lease's private document id, holder, + /// epoch, unreleased state, and expiry. A later takeover or expiry makes + /// every conditioned mutation fail in the backend's atomic write boundary. + pub fn write_condition(&self, scope: Scope) -> Result<WriteCondition> { + WriteCondition::new( + scope, + LEASES, + self.id.clone(), + Filter::eq("holder", self.holder.clone()) + .and(Filter::eq("epoch", self.epoch)) + .and(Filter::ne("released", true)), + )? + .with_valid_until_field("expires_ms") + } } /// Leases over [`LEASES`], one record per `(owner, key)`. @@ -504,6 +521,10 @@ impl Leases { } } +#[cfg(test)] +#[path = "leases_tests.rs"] +mod tests; + fn id(owner: &str, key: &str) -> Result<String> { super::super::composite_key([owner, key]) } diff --git a/crates/tinyagents-storage-kit/src/primitives/leases_tests.rs b/crates/tinyagents-storage-kit/src/primitives/leases_tests.rs new file mode 100644 index 000000000..17acbf038 --- /dev/null +++ b/crates/tinyagents-storage-kit/src/primitives/leases_tests.rs @@ -0,0 +1,62 @@ +use super::*; +use tinystoragedrivers_types::Meta; + +#[test] +fn lease_write_condition_pins_holder_epoch_release_state_and_expiry() { + let lease = Lease { + id: "lease-id".to_owned(), + holder: "holder-a".to_owned(), + epoch: 7, + expires_ms: 500, + endpoint: None, + }; + let condition = lease + .write_condition(tinystoragedrivers_types::Scope::new("cluster").unwrap()) + .unwrap(); + let [guard] = condition.clauses() else { + panic!("one lease must produce one guard clause"); + }; + + assert_eq!(guard.collection(), LEASES); + assert_eq!(guard.id(), "lease-id"); + assert_eq!(guard.valid_until_field(), Some("expires_ms")); + + let active = serde_json::json!({ + "holder": "holder-a", + "epoch": 7, + "expires_ms": 500 + }); + assert!(guard.matches_at("lease-id", &active, &Meta::unknown(), 499)); + assert!(!guard.matches_at("lease-id", &active, &Meta::unknown(), 500)); + assert!(!guard.matches_at( + "lease-id", + &serde_json::json!({ + "holder": "holder-b", + "epoch": 7, + "expires_ms": 500 + }), + &Meta::unknown(), + 499 + )); + assert!(!guard.matches_at( + "lease-id", + &serde_json::json!({ + "holder": "holder-a", + "epoch": 8, + "expires_ms": 500 + }), + &Meta::unknown(), + 499 + )); + assert!(!guard.matches_at( + "lease-id", + &serde_json::json!({ + "holder": "holder-a", + "epoch": 7, + "expires_ms": 500, + "released": true + }), + &Meta::unknown(), + 499 + )); +} diff --git a/crates/tinyagents-tasks/Cargo.toml b/crates/tinyagents-tasks/Cargo.toml index b8dca0519..e22d9b22a 100644 --- a/crates/tinyagents-tasks/Cargo.toml +++ b/crates/tinyagents-tasks/Cargo.toml @@ -11,10 +11,11 @@ description = "Graph-independent detached-task registry, task stores and orchest [dependencies] serde = { workspace = true } serde_json = { workspace = true } -tinyagents-harness = { path = "../tinyagents-harness", version = "2.11.0", default-features = false } +tinyagents-harness = { path = "../tinyagents-harness", version = "3.0.0", default-features = false } tinyinference-llm = { path = "../../vendor/tinyinference/crates/tinyinference-llm", version = "0.5.0" } tokio = { workspace = true, features = ["sync", "time", "rt", "rt-multi-thread"] } tracing = { workspace = true } +uuid = { workspace = true } # Storage ports (no database client). `blocking` bridges the synchronous # `TaskStore` / `CompletionStore` traits onto the async ports. tinyagents-storage-kit = { workspace = true } diff --git a/crates/tinyagents-tasks/src/README.md b/crates/tinyagents-tasks/src/README.md index 09da9d45d..980919f39 100644 --- a/crates/tinyagents-tasks/src/README.md +++ b/crates/tinyagents-tasks/src/README.md @@ -2,7 +2,7 @@ Graph-independent detached-task machinery: task model, `TaskStore` implementations, the detached task registry and restart reconciliation. -`tinyagents-graph` and `tinyagents-orchestration` depend on this crate; it does +`tinyagents-orchestration` depends on this crate; it does not depend on the graph engine. The task store and runtime registry provide hosts with stable task ids without exposing executor handles such as `tokio::JoinHandle`. @@ -57,6 +57,9 @@ hosts with stable task ids without exposing executor handles such as `cancel_all`, `snapshot(s)`, or fetch a `steering_handle` for a task it still owns in-process. `sweep_terminal` / the `soft_cap` passed to `new` bound unbounded growth from tasks nobody ever waited on. + Adapters with sidecar maps can use `register_cooperative_managed` and + `prune_terminal_trusted` to coordinate removal with their own waiter pins; + that registration path requires the adapter to enforce its finite cap. ### Orphan reconciliation (`reconcile.rs`) diff --git a/crates/tinyagents-tasks/src/completions/admission.rs b/crates/tinyagents-tasks/src/completions/admission.rs new file mode 100644 index 000000000..d84b3ac7a --- /dev/null +++ b/crates/tinyagents-tasks/src/completions/admission.rs @@ -0,0 +1,141 @@ +//! Turn-queue admission tied to a durable parent lifecycle generation. + +use std::sync::{Arc, Mutex}; + +use tinyagents_harness::run_queue::{QueueAdmission, QueueAdmissionError, QueueAdmissionPermit}; + +use super::CompletionStore; +use super::lifecycle::ParentLifecycle; +use super::{CompletionRecord, CompletionRouter}; + +/// Retained admission capability for one queued completion message. +/// +/// The harness checks this capability when the queue consumer applies the +/// message to a model turn. Parent stop, deletion, or reopen changes the +/// durable generation and rejects messages queued under an older generation. +#[derive(Clone)] +pub struct CompletionAdmission { + store: Arc<Mutex<super::router::AdmissionStoreSlot>>, + lifecycle: Arc<ParentLifecycle>, + parent_key: String, + parent_epoch: u64, +} + +impl std::fmt::Debug for CompletionAdmission { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("CompletionAdmission") + .field("parent_key", &self.parent_key) + .field("parent_epoch", &self.parent_epoch) + .finish_non_exhaustive() + } +} + +impl CompletionAdmission { + /// Binds a queued completion to the lifecycle generation that recorded it. + pub fn for_record(router: Arc<CompletionRouter>, record: &CompletionRecord) -> Self { + Self::for_router(&router, record) + } + + /// Binds a queued completion directly to its durable store and lifecycle. + pub fn for_store(store: Arc<dyn CompletionStore>, record: &CompletionRecord) -> Self { + Self { + store: Arc::new(Mutex::new(super::router::AdmissionStoreSlot { + store: Arc::clone(&store), + migrating: false, + })), + lifecycle: if store.is_process_local() { + store + .parent_lifecycle() + .unwrap_or_else(|| ParentLifecycle::for_store(store.as_ref())) + } else { + ParentLifecycle::new() + }, + parent_key: record.parent_key.clone(), + parent_epoch: record.parent_epoch, + } + } + + /// Binds the admission to the router's live store slot so a queued item + /// follows a successful durable-store migration. + pub(crate) fn for_router(router: &CompletionRouter, record: &CompletionRecord) -> Self { + Self { + store: router.admission_store_handle(), + lifecycle: router.parent_lifecycle_handle(), + parent_key: record.parent_key.clone(), + parent_epoch: record.parent_epoch, + } + } + + fn lifecycle_for(&self, store: &Arc<dyn CompletionStore>) -> Arc<ParentLifecycle> { + if store.is_process_local() { + store + .parent_lifecycle() + .unwrap_or_else(|| Arc::clone(&self.lifecycle)) + } else { + Arc::clone(&self.lifecycle) + } + } + + fn key_for( + &self, + store: &Arc<dyn CompletionStore>, + lifecycle: &Arc<ParentLifecycle>, + ) -> String { + let identity = if store.is_process_local() { + format!("lifecycle:{}", lifecycle.identity()) + } else { + store + .coordination_identity() + .unwrap_or_else(|| format!("store:{:p}", Arc::as_ptr(store))) + }; + serde_json::to_string(&(identity, self.parent_key.as_str(), self.parent_epoch)) + .expect("serializing a completion admission batch key cannot fail") + } +} + +impl QueueAdmission for CompletionAdmission { + fn acquire(&self) -> Result<Arc<dyn QueueAdmissionPermit>, QueueAdmissionError> { + self.acquire_for(self.batch_key().as_deref()) + } + + fn acquire_for( + &self, + expected_key: Option<&str>, + ) -> Result<Arc<dyn QueueAdmissionPermit>, QueueAdmissionError> { + let store = CompletionRouter::admission_store_snapshot(&self.store)?; + let lifecycle = self.lifecycle_for(&store); + if expected_key.is_some_and(|expected| self.key_for(&store, &lifecycle) != expected) { + // The queue's key came from the immediately preceding snapshot. + // A mismatch therefore means migration won between those steps; + // carry the transient classification in this same result rather + // than querying the mutable slot again later. + return Err(QueueAdmissionError::Retry); + } + let permit = super::router::ParentAdmissionPermit::acquire( + Arc::clone(&store), + Arc::clone(&lifecycle), + &self.parent_key, + self.parent_epoch, + ) + .map_err(|_| QueueAdmissionError::Rejected)? + .ok_or(QueueAdmissionError::Rejected)?; + CompletionRouter::admission_store_is_current(&self.store, &store)?; + if store.is_process_local() + && store + .parent_lifecycle() + .is_some_and(|current| !Arc::ptr_eq(¤t, &lifecycle)) + { + return Err(QueueAdmissionError::Retry); + } + Ok(Arc::new(permit)) + } + + fn batch_key(&self) -> Option<String> { + let store = CompletionRouter::admission_store_snapshot(&self.store).ok()?; + // Keep each component framed: scope and parent keys are caller-owned + // strings and may contain any punctuation, including colons. + let lifecycle = self.lifecycle_for(&store); + Some(self.key_for(&store, &lifecycle)) + } +} diff --git a/crates/tinyagents-tasks/src/completions/delivery.rs b/crates/tinyagents-tasks/src/completions/delivery.rs index bd6b178c6..f7c825720 100644 --- a/crates/tinyagents-tasks/src/completions/delivery.rs +++ b/crates/tinyagents-tasks/src/completions/delivery.rs @@ -1,12 +1,12 @@ //! Host-driven delivery timing around durable completion claims. -use std::collections::HashSet; +use std::collections::{HashMap, HashSet}; use std::future::Future; use std::pin::Pin; use std::sync::{Arc, Mutex}; use std::time::Duration; -use super::{CompletionRecord, CompletionRouter}; +use super::{CompletionReceipt, CompletionRecord, CompletionRouter}; /// Timing and batch bounds for one delivery drain. #[derive(Clone, Copy, Debug)] @@ -47,8 +47,10 @@ pub enum DeliveryOutcome { Delivered { count: usize }, /// Delivery failed and the records remain retryable or gave up. Failed { count: usize, gave_up: usize }, + /// Another drain already owns the tenant-scoped thread delivery lease. + LeaseHeld, /// The parent remained busy beyond the configured polling window. - Busy, + ParentBusy, /// The host's ownership fence stopped this drain. Fenced, } @@ -60,6 +62,8 @@ pub struct DeliveryTarget<'a> { pub router: &'a CompletionRouter, /// Stable key recorded on child completions. pub parent_key: &'a str, + /// Tenant/profile namespace for thread-level coordination. + pub tenant_key: &'a str, /// Thread key used for one-delivery-at-a-time coordination. pub thread_key: &'a str, } @@ -105,6 +109,33 @@ impl CompletionDelivery { .await } + /// Delivers one batch with a call-specific initial delay. The override + /// applies only to this drain; recovery and later calls keep their own + /// configured delays. + pub async fn deliver_once_with_initial_delay<B, F, D, Fut, E, U>( + &self, + target: DeliveryTarget<'_>, + initial_delay: Duration, + is_busy: B, + admitted: F, + deliver: D, + write_undelivered: U, + ) -> DeliveryOutcome + where + B: FnMut() -> bool, + F: DeliveryFence, + D: FnMut(&[CompletionRecord]) -> Fut, + Fut: Future<Output = std::result::Result<(), E>>, + E: std::fmt::Display, + U: FnMut(&[CompletionRecord]), + { + let mut config = self.config; + config.debounce = initial_delay; + Self::new(config, self.thread_leases.clone()) + .deliver_once(target, is_busy, admitted, deliver, write_undelivered) + .await + } + /// Error-aware form of [`Self::deliver_once`]. The callback receives the /// final delivery error for records that give up, synchronously under the /// thread lease after the durable give-up transition. @@ -124,8 +155,11 @@ impl CompletionDelivery { E: std::fmt::Display, U: FnMut(&[CompletionRecord], &str), { - let Some(_lease) = self.thread_leases.try_acquire(target.thread_key) else { - return DeliveryOutcome::Busy; + let Some(_lease) = self + .thread_leases + .try_acquire(TenantThreadKey::new(target.tenant_key, target.thread_key)) + else { + return DeliveryOutcome::LeaseHeld; }; self.drain( target, @@ -181,8 +215,11 @@ impl CompletionDelivery { E: std::fmt::Display, U: FnMut(&[CompletionRecord], &str), { - let Some(_lease) = self.thread_leases.try_acquire(target.thread_key) else { - return DeliveryOutcome::Busy; + let Some(_lease) = self + .thread_leases + .try_acquire(TenantThreadKey::new(target.tenant_key, target.thread_key)) + else { + return DeliveryOutcome::LeaseHeld; }; self.drain( target, @@ -230,7 +267,7 @@ impl CompletionDelivery { return DeliveryOutcome::Fenced; } if is_busy() { - return DeliveryOutcome::Busy; + return DeliveryOutcome::ParentBusy; } let mut batch = match target @@ -263,8 +300,9 @@ impl CompletionDelivery { } match delivery_result { Ok(()) => { - let ids: Vec<_> = batch.iter().map(|record| record.task_id.as_str()).collect(); - match target.router.mark_delivered(&ids) { + let receipts: Vec<_> = + batch.iter().filter_map(CompletionRecord::receipt).collect(); + match target.router.mark_delivered(&receipts) { Ok(_) => { claim.clear(); delivered_total += batch.len(); @@ -308,8 +346,9 @@ impl CompletionDelivery { } Err(error) => { tracing::warn!(parent_key = target.parent_key, error = %error, "completion delivery turn failed"); - let ids: Vec<_> = batch.iter().map(|record| record.task_id.as_str()).collect(); - let gave_up = match target.router.mark_failed(&ids) { + let receipts: Vec<_> = + batch.iter().filter_map(CompletionRecord::receipt).collect(); + let gave_up = match target.router.mark_failed(&receipts) { Ok(records) => records, Err(error) => { tracing::warn!(parent_key = target.parent_key, error = %error, "completion failure update failed"); @@ -377,54 +416,100 @@ async fn sleep_if_admitted<F: DeliveryFence>(duration: Duration, admitted: &F) - struct BatchClaim<'a> { router: &'a CompletionRouter, - ids: Vec<String>, + receipts: Vec<CompletionReceipt>, } impl<'a> BatchClaim<'a> { fn new(router: &'a CompletionRouter, records: &[CompletionRecord]) -> Self { Self { router, - ids: records + receipts: records .iter() - .map(|record| record.task_id.clone()) + .filter_map(CompletionRecord::receipt) .collect(), } } fn replace(&mut self, records: &[CompletionRecord]) { - self.ids = records + self.receipts = records .iter() - .map(|record| record.task_id.clone()) + .filter_map(CompletionRecord::receipt) .collect(); } fn clear(&mut self) { - self.ids.clear(); + self.receipts.clear(); } } impl Drop for BatchClaim<'_> { fn drop(&mut self) { - self.router.release(&self.ids); + self.router.release(&self.receipts); + } +} + +/// Tenant-scoped identity used by parent and thread delivery gates. +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +pub struct TenantThreadKey { + tenant: String, + key: String, +} + +impl TenantThreadKey { + /// Names a thread or session within a tenant/profile. + pub fn new(tenant: impl Into<String>, key: impl Into<String>) -> Self { + Self { + tenant: tenant.into(), + key: key.into(), + } + } + + /// Tenant/profile component. + pub fn tenant(&self) -> &str { + &self.tenant + } + + /// Thread or session component. + pub fn key(&self) -> &str { + &self.key + } +} + +impl From<String> for TenantThreadKey { + fn from(key: String) -> Self { + Self::new("", key) + } +} + +impl From<&str> for TenantThreadKey { + fn from(key: &str) -> Self { + Self::new("", key) } } +#[derive(Default)] +struct ParentBusyState { + leases: HashSet<(TenantThreadKey, u64)>, + /// Maps an active session id to the corresponding stable thread id. + session_threads: HashMap<TenantThreadKey, TenantThreadKey>, +} + /// Parent turn activity gate. Dropping a lease clears only the guard it owns. #[derive(Clone, Default)] -pub struct ParentBusyGate(Arc<Mutex<HashSet<(String, u64)>>>); +pub struct ParentBusyGate(Arc<Mutex<ParentBusyState>>); /// RAII guard marking one parent turn busy. pub struct ParentBusyLease { gate: ParentBusyGate, - key: String, + key: TenantThreadKey, id: u64, } impl ParentBusyGate { /// Marks one user turn busy. Overlapping leases keep the parent busy until /// every guard is dropped; `None` indicates a poisoned lock. - pub fn enter(&self, parent_key: impl Into<String>) -> Option<ParentBusyLease> { + pub fn enter(&self, parent_key: impl Into<TenantThreadKey>) -> Option<ParentBusyLease> { static NEXT: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(1); let key = parent_key.into(); let mut entries = self.0.lock().ok()?; let id = NEXT.fetch_add(1, std::sync::atomic::Ordering::Relaxed); - entries.insert((key.clone(), id)); + entries.leases.insert((key.clone(), id)); Some(ParentBusyLease { gate: self.clone(), key, @@ -434,48 +519,112 @@ impl ParentBusyGate { /// Whether a parent has an active busy lease. pub fn is_busy(&self, parent_key: &str) -> bool { + self.is_busy_scoped(&TenantThreadKey::from(parent_key)) + } + + /// Maps a session alias to its stable thread key for busy lookups. + pub fn alias_session_to_thread( + &self, + session: TenantThreadKey, + thread: TenantThreadKey, + ) -> bool { + if session.tenant != thread.tenant { + return false; + } self.0 .lock() - .is_ok_and(|entries| entries.iter().any(|(key, _)| key == parent_key)) + .is_ok_and(|mut state| state.session_threads.insert(session, thread).is_none()) + } + + /// Whether a tenant-scoped thread or one of its session aliases is busy. + pub fn is_busy_scoped(&self, key: &TenantThreadKey) -> bool { + self.0.lock().is_ok_and(|state| { + state.leases.iter().any(|(busy, _)| { + busy == key + || state.session_threads.get(busy) == Some(key) + || state.session_threads.get(key) == Some(busy) + }) + }) + } + + /// Clears active leases for a parent and removes its session aliases. + pub fn clear(&self, parent_key: &TenantThreadKey) -> usize { + let Ok(mut state) = self.0.lock() else { + return 0; + }; + let mut related = HashSet::from([parent_key.clone()]); + for (session, thread) in &state.session_threads { + if session == parent_key || thread == parent_key { + related.insert(session.clone()); + related.insert(thread.clone()); + } + } + state + .session_threads + .retain(|session, thread| !related.contains(session) && !related.contains(thread)); + let before = state.leases.len(); + state.leases.retain(|(key, _)| !related.contains(key)); + before - state.leases.len() } } impl Drop for ParentBusyLease { fn drop(&mut self) { if let Ok(mut entries) = self.gate.0.lock() { - entries.remove(&(self.key.clone(), self.id)); + entries.leases.remove(&(self.key.clone(), self.id)); } } } /// At-most-one live delivery turn per thread, released automatically on drop. #[derive(Clone, Default)] -pub struct ThreadDeliveryLeases(Arc<Mutex<HashSet<String>>>); +pub struct ThreadDeliveryLeases(Arc<Mutex<HashMap<TenantThreadKey, u64>>>); /// RAII exclusive delivery lease for one thread. pub struct ThreadDeliveryLease { leases: ThreadDeliveryLeases, - key: String, + key: TenantThreadKey, + id: u64, } impl ThreadDeliveryLeases { /// Acquires the per-thread delivery slot, if currently free. - pub fn try_acquire(&self, thread_key: impl Into<String>) -> Option<ThreadDeliveryLease> { + pub fn try_acquire( + &self, + thread_key: impl Into<TenantThreadKey>, + ) -> Option<ThreadDeliveryLease> { let key = thread_key.into(); let mut leases = self.0.lock().ok()?; - if !leases.insert(key.clone()) { + if leases.contains_key(&key) { return None; } + static NEXT: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(1); + let id = NEXT.fetch_add(1, std::sync::atomic::Ordering::Relaxed); + leases.insert(key.clone(), id); Some(ThreadDeliveryLease { leases: self.clone(), key, + id, + }) + } + + /// Releases every process-local delivery lease and returns the count. + /// Existing guards become inert and cannot remove a subsequently acquired + /// lease for the same thread when they are dropped. + pub fn release_all(&self) -> usize { + self.0.lock().map_or(0, |mut leases| { + let count = leases.len(); + leases.clear(); + count }) } } impl Drop for ThreadDeliveryLease { fn drop(&mut self) { - if let Ok(mut leases) = self.leases.0.lock() { + if let Ok(mut leases) = self.leases.0.lock() + && leases.get(&self.key) == Some(&self.id) + { leases.remove(&self.key); } } diff --git a/crates/tinyagents-tasks/src/completions/delivery_tests.rs b/crates/tinyagents-tasks/src/completions/delivery_tests.rs index 1dba0364a..07151c82d 100644 --- a/crates/tinyagents-tasks/src/completions/delivery_tests.rs +++ b/crates/tinyagents-tasks/src/completions/delivery_tests.rs @@ -13,6 +13,7 @@ fn target(router: &CompletionRouter) -> DeliveryTarget<'_> { DeliveryTarget { router, parent_key: "parent", + tenant_key: "", thread_key: "thread", } } @@ -456,7 +457,7 @@ async fn shared_thread_lease_prevents_overlapping_delivery_turns() { |_| {}, ) .await; - assert_eq!(second, DeliveryOutcome::Busy); + assert_eq!(second, DeliveryOutcome::LeaseHeld); release_tx.send(()).unwrap(); assert_eq!( first.await.unwrap(), @@ -464,6 +465,67 @@ async fn shared_thread_lease_prevents_overlapping_delivery_turns() { ); } +#[test] +fn same_thread_id_is_isolated_between_tenants_for_delivery_and_busy_gates() { + let delivery = ThreadDeliveryLeases::default(); + let thread_a = TenantThreadKey::new("tenant-a", "thread"); + let thread_b = TenantThreadKey::new("tenant-b", "thread"); + let lease_a = delivery.try_acquire(thread_a.clone()).unwrap(); + let lease_b = delivery.try_acquire(thread_b.clone()).unwrap(); + assert!(delivery.try_acquire(thread_a.clone()).is_none()); + + let busy = ParentBusyGate::default(); + let session_a = TenantThreadKey::new("tenant-a", "session"); + assert!(busy.alias_session_to_thread(session_a.clone(), thread_a.clone())); + let _busy_a = busy.enter(session_a).unwrap(); + assert!(busy.is_busy_scoped(&thread_a)); + assert!(!busy.is_busy_scoped(&thread_b)); + assert_eq!(busy.clear(&thread_a), 1); + assert!(!busy.is_busy_scoped(&thread_a)); + + drop((lease_a, lease_b)); +} + +#[tokio::test] +async fn delivery_accepts_a_call_specific_initial_delay() { + let router = CompletionRouter::new(Arc::new(InMemoryCompletionStore::new())); + let delivery = CompletionDelivery::new( + CompletionDeliveryConfig::default(), + ThreadDeliveryLeases::default(), + ); + let outcome = delivery + .deliver_once_with_initial_delay( + target(&router), + Duration::ZERO, + || false, + || true, + |_| async { Ok::<(), &'static str>(()) }, + |_| {}, + ) + .await; + assert_eq!(outcome, DeliveryOutcome::Empty); +} + +#[test] +fn release_all_clears_delivery_leases_without_touching_busy_state() { + let delivery = ThreadDeliveryLeases::default(); + let thread_a = TenantThreadKey::new("tenant-a", "thread-a"); + let thread_b = TenantThreadKey::new("tenant-b", "thread-b"); + let busy = ParentBusyGate::default(); + let lease_a = delivery.try_acquire(thread_a.clone()).unwrap(); + let lease_b = delivery.try_acquire(thread_b.clone()).unwrap(); + let _busy = busy.enter(thread_a.clone()).unwrap(); + + assert_eq!(delivery.release_all(), 2); + let renewed_a = delivery.try_acquire(thread_a.clone()).unwrap(); + let renewed_b = delivery.try_acquire(thread_b.clone()).unwrap(); + drop((lease_a, lease_b)); + assert!(delivery.try_acquire(thread_a.clone()).is_none()); + assert!(delivery.try_acquire(thread_b).is_none()); + assert!(busy.is_busy_scoped(&thread_a)); + drop((renewed_a, renewed_b)); +} + #[tokio::test] async fn aborted_delivery_releases_its_claimed_batch() { let router = Arc::new(CompletionRouter::new(Arc::new( diff --git a/crates/tinyagents-tasks/src/completions/driver.rs b/crates/tinyagents-tasks/src/completions/driver.rs index 24f387c75..08c3000fa 100644 --- a/crates/tinyagents-tasks/src/completions/driver.rs +++ b/crates/tinyagents-tasks/src/completions/driver.rs @@ -12,9 +12,9 @@ //! body (`task_id`, `parent_key`, `agent_id`, `state`, `result`, `finished_ms`, //! ...). `parent_key` and `state` are indexed for [`CompletionStore::list`] and //! [`CompletionStore::compact`]. The record's last-update time is the document's -//! `updated_ms` metadata, which the driver stamps on every write. `put` is an -//! upsert, which is all the contract asks of it: the router serialises its -//! read-modify-write sequences. +//! `updated_ms` metadata, which the driver stamps on every write. Ordinary +//! `put` is an upsert; claim, reclaim, and settlement operations use a row +//! version precondition so concurrent routers cannot overwrite one another. //! //! # Sync trait //! @@ -26,19 +26,25 @@ use std::sync::Arc; use std::time::{Duration, SystemTime}; use tinyagents_harness::error::{Result, TinyAgentsError}; +use tinyagents_storage_kit::primitives::LEASES; use tinyagents_storage_kit::{ - Blocking, CollectionSpec, DocumentStore, DocumentStoreExt, Filter, IndexSpec, Precondition, - Query, Repo, StorageError, UPDATED_FIELD, Versioned, WriteOptions, + Blocking, CollectionSpec, DocumentStore, DocumentStoreExt, ErrorKind, Filter, IndexSpec, + Leases, Precondition, Query, Scope, StorageBackend, StorageError, UPDATED_FIELD, Versioned, + WriteCondition, WriteOptions, }; use tokio::sync::OnceCell; -use super::store::CompletionStore; +use super::store::{ + ClaimResolution, CompletionStore, ParentGateClaim, PendingClaimRequest, ResolveClaimRequest, + TaskClaimRequest, +}; use super::types::{CompletionRecord, CompletionState}; use crate::storage::{ COMPLETIONS, CompletionDoc, SCHEMA_VERSION, completions_spec, task_id_key, to_ms, }; const LOG_PREFIX: &str = "[completion-store]"; +const PARENT_GATE_OWNER: &str = "tinyagents.tasks.completions"; fn map_error(error: StorageError) -> TinyAgentsError { TinyAgentsError::Storage(format!("completion store: {error}")) @@ -68,40 +74,108 @@ fn decode_strict(row: Versioned<serde_json::Value>) -> Result<CompletionRecord> } /// A durable [`CompletionStore`] stored in a driver [`DocumentStore`]. -#[derive(Debug, Clone)] +#[derive(Clone)] pub struct DriverCompletionStore { + backend: Arc<dyn StorageBackend>, + scope: Scope, + coordination_domain: String, docs: Arc<dyn DocumentStore>, bridge: Blocking, - repo: Repo<CompletionDoc>, declared: Arc<OnceCell<()>>, } +impl std::fmt::Debug for DriverCompletionStore { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("DriverCompletionStore") + .field("scope", &self.scope.as_str()) + .field("coordination_domain", &self.coordination_domain) + .finish_non_exhaustive() + } +} + impl DriverCompletionStore { /// Store completions in `tasks.completions`, with its own bridge. /// + /// The default coordination identity is tied to this exact `Arc` + /// handle. Use `new_with_coordination_domain` when distinct handles may + /// point at the same shared database; all processes using that database + /// and scope must be configured with the same trusted domain. + /// /// # Errors /// /// When the bridge's runtime thread cannot start. - pub fn new(docs: Arc<dyn DocumentStore>) -> std::result::Result<Self, StorageError> { - Ok(Self::with_bridge(docs, Blocking::new()?)) + pub fn new( + backend: Arc<dyn StorageBackend>, + scope: Scope, + ) -> std::result::Result<Self, StorageError> { + let coordination_domain = format!("backend-handle:{:p}", Arc::as_ptr(&backend)); + Self::new_with_coordination_domain(backend, scope, coordination_domain) + } + + /// Creates a driver-backed store with an explicit identity for a shared + /// backend domain. Hosts opening distinct handles to the same database + /// must pass the same opaque domain id; unrelated databases must use + /// different ids. Scope remains part of the resulting identity. + pub fn new_with_coordination_domain( + backend: Arc<dyn StorageBackend>, + scope: Scope, + coordination_domain: impl Into<String>, + ) -> std::result::Result<Self, StorageError> { + Self::with_bridge_and_coordination_domain( + backend, + scope, + Blocking::new()?, + coordination_domain, + ) } /// Store completions in `tasks.completions`, running the synchronous trait /// on `bridge` so several stores share one thread. /// + /// The default coordination identity is tied to this exact `Arc` + /// handle. Use `with_bridge_and_coordination_domain` when distinct + /// handles may point at the same shared database. + /// /// # Panics /// /// Never: the collection name is a valid constant. #[allow(clippy::expect_used, reason = "the collection name is a constant")] - pub fn with_bridge(docs: Arc<dyn DocumentStore>, bridge: Blocking) -> Self { - Self { - repo: Repo::new(&docs, COMPLETIONS) - .expect("tasks.completions is a valid collection") - .with_options(WriteOptions::new().schema(SCHEMA_VERSION)), + pub fn with_bridge( + backend: Arc<dyn StorageBackend>, + scope: Scope, + bridge: Blocking, + ) -> std::result::Result<Self, StorageError> { + let coordination_domain = format!("backend-handle:{:p}", Arc::as_ptr(&backend)); + Self::with_bridge_and_coordination_domain(backend, scope, bridge, coordination_domain) + } + + /// Like [`Self::with_bridge`], with the caller's stable shared-backend + /// coordination identity. + pub fn with_bridge_and_coordination_domain( + backend: Arc<dyn StorageBackend>, + scope: Scope, + bridge: Blocking, + coordination_domain: impl Into<String>, + ) -> std::result::Result<Self, StorageError> { + // Fail during setup when the backend cannot issue durable handles + // whose writes are atomically guarded by the lease document. + let probe = WriteCondition::new( + scope.clone(), + "tinyagents.condition_probe", + "completion-store", + Filter::eq("holder", "probe"), + )?; + backend.for_scope_conditioned(&scope, &probe)?; + let scoped = backend.for_scope(&scope)?; + let docs = scoped.documents().clone(); + Ok(Self { docs, + backend, + scope, + coordination_domain: coordination_domain.into(), bridge, declared: Arc::new(OnceCell::new()), - } + }) } /// The declaration a store using `collection` needs. @@ -111,8 +185,8 @@ impl DriverCompletionStore { .index(IndexSpec::new("by_state", ["state"])) } - /// Declares the backing collection eagerly so an unavailable backend is - /// reported before the host accepts work. + /// Declares the backing collection eagerly so a host can decide whether to + /// keep this durable store or fall back to memory before accepting work. pub fn declare_collections(&self) -> Result<()> { let this = self.clone(); self.run(async move { this.declare().await }) @@ -133,14 +207,113 @@ impl DriverCompletionStore { self.docs .ensure_collection(&completions_spec()) .await + .map_err(map_error)?; + self.docs + .ensure_collection(&CollectionSpec::new(LEASES)) + .await .map_err(map_error) }) .await .map(|_| ()) } + + async fn fenced_docs( + &self, + parent_key: &str, + claim: &ParentGateClaim, + ) -> Result<Arc<dyn DocumentStore>> { + let lease = Leases::new(&self.docs) + .current(PARENT_GATE_OWNER, parent_key) + .await + .map_err(map_error)? + .filter(|lease| lease.holder == claim.holder && lease.epoch == claim.epoch) + .ok_or_else(|| { + TinyAgentsError::Storage("completion store: parent admission lease was lost".into()) + })?; + let condition = lease + .write_condition(self.scope.clone()) + .map_err(map_error)?; + let fenced = self + .backend + .for_scope_conditioned(&self.scope, &condition) + .map_err(map_error)?; + Ok(Arc::clone(fenced.documents())) + } } impl CompletionStore for DriverCompletionStore { + fn coordination_identity(&self) -> Option<String> { + Some(format!( + "driver:{}", + serde_json::to_string(&(self.coordination_domain.as_str(), self.scope.as_str())) + .expect("driver coordination identity is serializable") + )) + } + + fn acquire_parent_gate( + &self, + parent_key: &str, + holder: &str, + ttl: Duration, + ) -> Result<Option<ParentGateClaim>> { + let this = self.clone(); + let parent_key = parent_key.to_owned(); + let holder = holder.to_owned(); + let ttl_ms = duration_ms(ttl)?; + self.run(async move { + this.declare().await?; + Leases::new(&this.docs) + .acquire(PARENT_GATE_OWNER, &parent_key, &holder, ttl_ms) + .await + .map(|lease| { + lease.map(|lease| ParentGateClaim { + holder: lease.holder, + epoch: lease.epoch, + }) + }) + .map_err(map_error) + }) + } + + fn renew_parent_gate( + &self, + parent_key: &str, + claim: &ParentGateClaim, + ttl: Duration, + ) -> Result<bool> { + let this = self.clone(); + let parent_key = parent_key.to_owned(); + let claim = claim.clone(); + let ttl_ms = duration_ms(ttl)?; + self.run(async move { + this.declare().await?; + Leases::new(&this.docs) + .renew( + PARENT_GATE_OWNER, + &parent_key, + &claim.holder, + claim.epoch, + ttl_ms, + ) + .await + .map(|lease| lease.is_some()) + .map_err(map_error) + }) + } + + fn release_parent_gate(&self, parent_key: &str, claim: &ParentGateClaim) -> Result<bool> { + let this = self.clone(); + let parent_key = parent_key.to_owned(); + let claim = claim.clone(); + self.run(async move { + this.declare().await?; + Leases::new(&this.docs) + .release(PARENT_GATE_OWNER, &parent_key, &claim.holder, claim.epoch) + .await + .map_err(map_error) + }) + } + fn get(&self, task_id: &str) -> Option<CompletionRecord> { let this = self.clone(); let task_id = task_id.to_owned(); @@ -155,17 +328,257 @@ impl CompletionStore for DriverCompletionStore { .flatten() } + fn try_get(&self, task_id: &str) -> Result<Option<CompletionRecord>> { + let this = self.clone(); + let task_id = task_id.to_owned(); + self.run(async move { + this.declare().await?; + let id = task_id_key(&task_id).map_err(map_error)?; + let found = this.docs.get(COMPLETIONS, &id).await.map_err(map_error)?; + found.map(decode_strict).transpose() + }) + } + fn put(&self, record: &CompletionRecord) -> Result<()> { + let _ = record; + Err(TinyAgentsError::Storage( + "completion store: durable writes require a parent gate claim".into(), + )) + } + + fn put_fenced( + &self, + record: &CompletionRecord, + parent_key: &str, + gate: &ParentGateClaim, + ) -> Result<()> { + let this = self.clone(); + let record = record.clone(); + let parent_key = parent_key.to_owned(); + let gate = gate.clone(); + self.run(async move { + this.declare().await?; + let docs = this.fenced_docs(&parent_key, &gate).await?; + let id = task_id_key(&record.task_id).map_err(map_error)?; + let current = docs.get(COMPLETIONS, &id).await.map_err(map_error)?; + let precondition = current.as_ref().map_or(Precondition::Absent, |row| { + Precondition::Version(row.version) + }); + docs.put_with( + COMPLETIONS, + &id, + serde_json::to_value(CompletionDoc::from_record(&record))?, + precondition, + WriteOptions::new().schema(SCHEMA_VERSION), + ) + .await + .map(|_| ()) + .map_err(map_error) + }) + } + + fn insert_if_absent(&self, record: &CompletionRecord, gate: &ParentGateClaim) -> Result<bool> { let this = self.clone(); let record = record.clone(); + let gate = gate.clone(); self.run(async move { this.declare().await?; + let docs = this.fenced_docs(&record.parent_key, &gate).await?; let id = task_id_key(&record.task_id).map_err(map_error)?; - this.repo - .put(&id, &CompletionDoc::from_record(&record)) + let body = serde_json::to_value(CompletionDoc::from_record(&record))?; + match docs + .put_with( + COMPLETIONS, + &id, + body, + Precondition::Absent, + WriteOptions::new().schema(SCHEMA_VERSION), + ) .await - .map(|_| ()) - .map_err(map_error) + { + Ok(_) => Ok(true), + Err(error) if error.kind() == ErrorKind::Conflict => Ok(false), + Err(error) => Err(map_error(error)), + } + }) + } + + fn claim_pending(&self, request: PendingClaimRequest) -> Result<Vec<CompletionRecord>> { + let this = self.clone(); + let PendingClaimRequest { + parent_key, + parent_epoch, + gate, + mode, + max, + token, + visible_after, + } = request; + self.run(async move { + this.declare().await?; + let docs = this.fenced_docs(&parent_key, &gate).await?; + let rows = docs + .query_all( + COMPLETIONS, + &Query::filter(Filter::eq("parent_key", parent_key.as_str())), + ) + .await + .map_err(map_error)?; + let now = SystemTime::now(); + let mut rows = rows + .into_iter() + .filter_map(|row| decode_strict(row.clone()).ok().map(|record| (row, record))) + .filter(|(_, record)| { + record.parent_epoch == parent_epoch + && record.state == CompletionState::Pending + && mode.accepts(record.notify_mode) + && record + .claim + .as_ref() + .is_none_or(|claim| claim.visible_after <= now) + }) + .collect::<Vec<_>>(); + rows.sort_by(|(_, left), (_, right)| { + left.finished_at + .cmp(&right.finished_at) + .then_with(|| left.task_id.cmp(&right.task_id)) + }); + let mut claimed = Vec::new(); + for (row, mut record) in rows.into_iter().take(max) { + record.attempts = record.attempts.saturating_add(1); + record.updated_at = now; + record.claim = Some(crate::completions::CompletionClaim { + token: format!("{token}:{}", record.task_id), + visible_after, + }); + let body = serde_json::to_value(CompletionDoc::from_record(&record))?; + match docs + .put_with( + COMPLETIONS, + &row.id, + body, + Precondition::Version(row.version), + WriteOptions::new().schema(SCHEMA_VERSION), + ) + .await + { + Ok(_) => claimed.push(record), + Err(error) if error.kind() == ErrorKind::Conflict => {} + Err(error) => return Err(map_error(error)), + } + } + Ok(claimed) + }) + } + + fn claim_task(&self, request: TaskClaimRequest) -> Result<Option<CompletionRecord>> { + let this = self.clone(); + let TaskClaimRequest { + task_id, + parent_key, + parent_epoch, + gate, + mode, + token, + visible_after, + } = request; + self.run(async move { + this.declare().await?; + let docs = this.fenced_docs(&parent_key, &gate).await?; + let id = task_id_key(&task_id).map_err(map_error)?; + let Some(row) = docs.get(COMPLETIONS, &id).await.map_err(map_error)? else { + return Ok(None); + }; + let mut record = decode_strict(row.clone())?; + let now = SystemTime::now(); + if record.parent_key != parent_key + || record.parent_epoch != parent_epoch + || record.state != CompletionState::Pending + || !mode.accepts(record.notify_mode) + || record + .claim + .as_ref() + .is_some_and(|claim| claim.visible_after > now) + { + return Ok(None); + } + record.attempts = record.attempts.saturating_add(1); + record.updated_at = now; + record.claim = Some(crate::completions::CompletionClaim { + token: format!("{token}:{task_id}"), + visible_after, + }); + let body = serde_json::to_value(CompletionDoc::from_record(&record))?; + match docs + .put_with( + COMPLETIONS, + &id, + body, + Precondition::Version(row.version), + WriteOptions::new().schema(SCHEMA_VERSION), + ) + .await + { + Ok(_) => Ok(Some(record)), + Err(error) if error.kind() == ErrorKind::Conflict => Ok(None), + Err(error) => Err(map_error(error)), + } + }) + } + + fn resolve_claim(&self, request: ResolveClaimRequest) -> Result<Option<CompletionRecord>> { + let this = self.clone(); + let ResolveClaimRequest { + task_id, + parent_key, + parent_epoch, + gate, + token, + resolution, + } = request; + self.run(async move { + this.declare().await?; + let docs = this.fenced_docs(&parent_key, &gate).await?; + let id = task_id_key(&task_id).map_err(map_error)?; + let Some(row) = docs.get(COMPLETIONS, &id).await.map_err(map_error)? else { + return Ok(None); + }; + let mut record = decode_strict(row.clone())?; + let now = SystemTime::now(); + if record.parent_key != parent_key + || record.parent_epoch != parent_epoch + || record.state != CompletionState::Pending + || record + .claim + .as_ref() + .is_none_or(|claim| claim.token != token || claim.visible_after <= now) + { + return Ok(None); + } + record.claim = None; + record.updated_at = SystemTime::now(); + match resolution { + ClaimResolution::Delivered => record.state = CompletionState::Delivered, + ClaimResolution::Failed { max_attempts } if record.attempts >= max_attempts => { + record.state = CompletionState::GaveUp; + } + ClaimResolution::Failed { .. } | ClaimResolution::Release => {} + } + let body = serde_json::to_value(CompletionDoc::from_record(&record))?; + match docs + .put_with( + COMPLETIONS, + &id, + body, + Precondition::Version(row.version), + WriteOptions::new().schema(SCHEMA_VERSION), + ) + .await + { + Ok(_) => Ok(Some(record)), + Err(error) if error.kind() == ErrorKind::Conflict => Ok(None), + Err(error) => Err(map_error(error)), + } }) } @@ -224,6 +637,11 @@ impl CompletionStore for DriverCompletionStore { .map_err(map_error)?; let mut dropped = 0; for row in rows { + let is_parent_marker = decode_strict(row.clone()) + .is_ok_and(|record| record.task_id.starts_with("\u{1}parent-cancelled:")); + if is_parent_marker { + continue; + } // Only the version we saw: a record refreshed since is live again. match this .docs @@ -231,7 +649,7 @@ impl CompletionStore for DriverCompletionStore { .await { Ok(removed) => dropped += usize::from(removed), - Err(error) if error.kind() == tinyagents_storage_kit::ErrorKind::Conflict => {} + Err(error) if error.kind() == ErrorKind::Conflict => {} Err(error) => return Err(map_error(error)), } } @@ -241,6 +659,12 @@ impl CompletionStore for DriverCompletionStore { } } +fn duration_ms(duration: Duration) -> Result<i64> { + i64::try_from(duration.as_millis()).map_err(|_| { + TinyAgentsError::Storage("completion store: parent gate ttl out of range".into()) + }) +} + #[cfg(test)] #[path = "driver_tests.rs"] mod tests; diff --git a/crates/tinyagents-tasks/src/completions/driver_tests.rs b/crates/tinyagents-tasks/src/completions/driver_tests.rs index d5ebb1a6e..7551490fc 100644 --- a/crates/tinyagents-tasks/src/completions/driver_tests.rs +++ b/crates/tinyagents-tasks/src/completions/driver_tests.rs @@ -1,15 +1,13 @@ -use std::path::Path; -use std::sync::Arc; -use std::time::{Duration, SystemTime, UNIX_EPOCH}; - use super::*; +use crate::completions::CompletionReceipt; use crate::completions::{CompletionResult, RecordOutcome}; use crate::completions::{CompletionRouter, CompletionStatus, InMemoryCompletionStore, NotifyMode}; -use crate::storage::{COMPLETIONS, task_id_key}; -use tinyagents_storage_kit::{Scope, StorageBackend}; -use tinystoragedrivers_document::memory::MemoryConnection; -use tinystoragedrivers_document::sqlite::SqliteDocuments; -use tinystoragedrivers_testkit::MemoryBackend; +use crate::conditional_sqlite::MemoryStorage; +use crate::open_driver_completion_store_with_coordination_domain; +use std::sync::Arc; +use std::time::{Duration, SystemTime, UNIX_EPOCH}; +use tinyagents_harness::run_queue::QueueAdmission; +use tinyagents_storage_kit::{ErrorKind, Precondition, Scope, StorageBackend}; fn record(task: &str, parent: &str) -> CompletionRecord { CompletionRecord::new( @@ -21,34 +19,653 @@ fn record(task: &str, parent: &str) -> CompletionRecord { ) } -fn docs_of(storage: &dyn StorageBackend, scope: &str) -> Arc<dyn DocumentStore> { - Arc::clone( - storage - .for_scope(&Scope::new(scope).unwrap()) +fn receipts(records: &[CompletionRecord]) -> Vec<CompletionReceipt> { + records + .iter() + .filter_map(CompletionRecord::receipt) + .collect() +} + +fn parent_gate( + store: &dyn CompletionStore, + parent_key: &str, +) -> crate::completions::ParentGateClaim { + store + .acquire_parent_gate(parent_key, "driver-test-owner", Duration::from_secs(30)) + .unwrap() + .unwrap() +} + +fn put_record(store: &dyn CompletionStore, record: &CompletionRecord) { + let parent_key = if record.parent_key.is_empty() { + format!("completion-task:{}", record.task_id) + } else { + record.parent_key.clone() + }; + let gate = parent_gate(store, &parent_key); + store.put_fenced(record, &parent_key, &gate).unwrap(); +} + +#[test] +fn stale_parent_admission_cannot_create_a_completion_after_takeover() { + let backend: Arc<dyn StorageBackend> = Arc::new(MemoryStorage::new()); + let store = + DriverCompletionStore::new(Arc::clone(&backend), Scope::new("shared").unwrap()).unwrap(); + let stale = store + .acquire_parent_gate("parent", "router-a", Duration::from_secs(30)) + .unwrap() + .unwrap(); + assert!(store.release_parent_gate("parent", &stale).unwrap()); + let current = store + .acquire_parent_gate("parent", "router-b", Duration::from_secs(30)) + .unwrap() + .unwrap(); + assert!(current.epoch > stale.epoch); + + assert!( + store + .insert_if_absent(&record("late", "parent"), ¤t) .unwrap() - .documents(), + ); + assert!( + store + .insert_if_absent(&record("late", "other-parent"), &stale) + .is_err() + ); + let preserved = store.get("late").unwrap(); + assert_eq!(preserved.parent_key, "parent"); + + assert!( + store + .insert_if_absent(&record("stale-create", "parent"), &stale) + .is_err() + ); + assert!(store.get("stale-create").is_none()); +} + +#[test] +fn cloned_memory_stores_share_a_batch_admission_identity() { + let store = InMemoryCompletionStore::new(); + let first: Arc<dyn CompletionStore> = Arc::new(store.clone()); + let second: Arc<dyn CompletionStore> = Arc::new(store.clone()); + let item = record("same-parent-a", "parent"); + let first = crate::completions::CompletionAdmission::for_store(first, &item); + let second = crate::completions::CompletionAdmission::for_store(second, &item); + + assert_eq!(first.batch_key(), second.batch_key()); +} + +#[test] +fn driver_stores_can_share_an_explicit_backend_domain_without_coalescing_other_domains() { + let first_backend: Arc<dyn StorageBackend> = Arc::new(MemoryStorage::new()); + let other_backend: Arc<dyn StorageBackend> = Arc::new(MemoryStorage::new()); + let scope = Scope::new("tenant").unwrap(); + let first = DriverCompletionStore::new_with_coordination_domain( + Arc::clone(&first_backend), + scope.clone(), + "shared-backend", + ) + .unwrap(); + let second = DriverCompletionStore::new_with_coordination_domain( + Arc::clone(&first_backend), + scope.clone(), + "shared-backend", ) + .unwrap(); + let unrelated = + DriverCompletionStore::new_with_coordination_domain(other_backend, scope, "other-backend") + .unwrap(); + + assert_eq!( + first.coordination_identity(), + second.coordination_identity() + ); + assert_ne!( + first.coordination_identity(), + unrelated.coordination_identity() + ); } -/// A router over the explicitly owned testkit backend. -fn open_router(storage: &MemoryBackend) -> CompletionRouter { - let store = DriverCompletionStore::new(docs_of(storage, "local")).unwrap(); - CompletionRouter::new(Arc::new(store)) +#[tokio::test] +async fn separate_database_domains_do_not_share_a_same_parent_queue_permit() { + use tinyagents_harness::run_queue::{QueueLane, QueueMode, RunQueue}; + use tinyinference_llm::message::Message; + + let backend_a: Arc<dyn StorageBackend> = Arc::new(MemoryStorage::new()); + let backend_b: Arc<dyn StorageBackend> = Arc::new(MemoryStorage::new()); + let scope = Scope::new("shared-tenant").unwrap(); + let router_a = Arc::new(CompletionRouter::new(Arc::new( + DriverCompletionStore::new_with_coordination_domain(backend_a, scope.clone(), "database-a") + .unwrap(), + ))); + let router_b = Arc::new(CompletionRouter::new(Arc::new( + DriverCompletionStore::new_with_coordination_domain(backend_b, scope, "database-b") + .unwrap(), + ))); + router_a + .record(record("database-a-child", "same-parent")) + .await + .unwrap(); + router_b + .record(record("database-b-child", "same-parent")) + .await + .unwrap(); + let completion_a = router_a.pending_for("same-parent").remove(0); + let completion_b = router_b.pending_for("same-parent").remove(0); + router_b.cancel_parent("same-parent").unwrap(); + + let queue = RunQueue::new(); + for (router, task_id, completion) in [ + (&router_a, "database-a-child", completion_a), + (&router_b, "database-b-child", completion_b), + ] { + assert_eq!(completion.task_id, task_id); + queue + .push_guarded( + QueueLane::Followup, + Message::user(task_id), + Arc::new(crate::completions::CompletionAdmission::for_record( + Arc::clone(router), + &completion, + )), + ) + .await; + } + + let admitted = queue + .take_admitted(QueueLane::Followup, QueueMode::All) + .await; + assert_eq!( + admitted + .iter() + .map(|item| item.value.clone()) + .collect::<Vec<_>>(), + [Message::user("database-a-child")], + "a permit from one database must not authorize an unrelated store with the same scope and parent text" + ); +} + +#[test] +fn in_memory_store_rejects_mutations_from_a_released_parent_gate() { + let store = InMemoryCompletionStore::new(); + let stale = store + .acquire_parent_gate("parent", "router-a", Duration::from_secs(30)) + .unwrap() + .unwrap(); + assert!(store.release_parent_gate("parent", &stale).unwrap()); + let current = store + .acquire_parent_gate("parent", "router-b", Duration::from_secs(30)) + .unwrap() + .unwrap(); + + let mut accepted = record("same-task", "parent"); + accepted.result = CompletionResult::text("current owner"); + store.put_fenced(&accepted, "parent", ¤t).unwrap(); + + let mut stale_update = accepted.clone(); + stale_update.result = CompletionResult::text("stale owner"); + assert!(store.put_fenced(&stale_update, "parent", &stale).is_err()); + assert_eq!( + store.get("same-task").unwrap().result, + CompletionResult::text("current owner") + ); + assert!( + store + .insert_if_absent(&record("stale-create", "parent"), &stale) + .is_err() + ); + assert!(store.get("stale-create").is_none()); +} + +#[tokio::test] +async fn a_claim_prepared_before_takeover_cannot_land_after_the_new_epoch() { + let store = DriverCompletionStore::new( + Arc::new(MemoryStorage::new()), + Scope::new("shared").unwrap(), + ) + .unwrap(); + let stale = store + .acquire_parent_gate("parent", "router-a", Duration::from_secs(30)) + .unwrap() + .unwrap(); + let stale_docs = store.fenced_docs("parent", &stale).await.unwrap(); + + assert!(store.release_parent_gate("parent", &stale).unwrap()); + let current = store + .acquire_parent_gate("parent", "router-b", Duration::from_secs(30)) + .unwrap() + .unwrap(); + let pending = record("late-claim", "parent"); + assert!(store.insert_if_absent(&pending, ¤t).unwrap()); + + let id = task_id_key(&pending.task_id).unwrap(); + let row = stale_docs.get(COMPLETIONS, &id).await.unwrap().unwrap(); + let mut claimed = decode_strict(row.clone()).unwrap(); + claimed.attempts = 1; + claimed.claim = Some(crate::completions::CompletionClaim { + token: "router-a:late-claim".into(), + visible_after: SystemTime::now() + Duration::from_secs(30), + }); + let error = stale_docs + .put_with( + COMPLETIONS, + &id, + serde_json::to_value(CompletionDoc::from_record(&claimed)).unwrap(), + Precondition::Version(row.version), + WriteOptions::new().schema(SCHEMA_VERSION), + ) + .await + .unwrap_err(); + assert_eq!(error.kind(), ErrorKind::Conflict); + let persisted = store.get(&pending.task_id).unwrap(); + assert_eq!(persisted.attempts, 0); + assert!(persisted.claim.is_none()); +} + +#[tokio::test] +async fn closed_parent_markers_survive_compaction_and_a_router_restart() { + let backend: Arc<dyn StorageBackend> = Arc::new(MemoryStorage::new()); + let open = || { + Arc::new( + DriverCompletionStore::new(Arc::clone(&backend), Scope::new("shared").unwrap()) + .unwrap(), + ) as Arc<dyn CompletionStore> + }; + let first = CompletionRouter::new(open()); + first.cancel_parent("cancelled-thread").unwrap(); + first.delete_parent("deleted-thread").unwrap(); + std::thread::sleep(Duration::from_millis(5)); + assert_eq!(first.compact(Duration::ZERO).unwrap(), 0); + + let restarted = CompletionRouter::new(open()); + assert_eq!( + restarted + .record(record("after-delete", "deleted-thread")) + .await + .unwrap(), + RecordOutcome::Suppressed + ); + assert_eq!( + restarted + .record(record("after-stop", "cancelled-thread")) + .await + .unwrap(), + RecordOutcome::Suppressed + ); + assert!(restarted.pending_for("deleted-thread").is_empty()); +} + +#[tokio::test] +async fn resuming_an_open_parent_preserves_its_generation_and_pending_work() { + let store = Arc::new(InMemoryCompletionStore::new()); + let router = CompletionRouter::new(store.clone()); + router.record(record("pending", "parent")).await.unwrap(); + let before = store.get("pending").unwrap().parent_epoch; + + router.resume_parent("parent"); + + assert!(router.parent_generation_admitted("parent", before).unwrap()); + assert_eq!(router.claim_pending("parent", 1).unwrap().len(), 1); +} + +#[tokio::test] +async fn deleted_parent_cannot_be_reopened_by_resume() { + let router = CompletionRouter::new(Arc::new(InMemoryCompletionStore::new())); + router.delete_parent("deleted").unwrap(); + + router.resume_parent("deleted"); + + assert!(!router.admit_child("deleted").unwrap()); + assert_eq!( + router + .record(record("after-delete", "deleted")) + .await + .unwrap(), + RecordOutcome::Suppressed + ); +} + +#[tokio::test] +async fn queued_completion_generation_is_rejected_after_cross_router_close_and_reopen() { + let backend: Arc<dyn StorageBackend> = Arc::new(MemoryStorage::new()); + let open = || { + Arc::new( + DriverCompletionStore::new(Arc::clone(&backend), Scope::new("shared").unwrap()) + .unwrap(), + ) as Arc<dyn CompletionStore> + }; + let delivery_router = CompletionRouter::new(open()); + let lifecycle_router = CompletionRouter::new(open()); + delivery_router + .record(record("queued", "parent")) + .await + .unwrap(); + let claimed = delivery_router.claim_pending("parent", 1).unwrap(); + let stale_receipt = claimed[0].receipt().unwrap(); + + lifecycle_router.cancel_parent("parent").unwrap(); + lifecycle_router.resume_parent("parent"); + assert!( + !delivery_router + .parent_generation_admitted("parent", stale_receipt.parent_epoch) + .unwrap() + ); + assert_eq!( + delivery_router + .mark_delivered(std::slice::from_ref(&stale_receipt)) + .unwrap(), + 0 + ); + let current_generation = delivery_router + .admit_child_generation("parent") + .unwrap() + .unwrap(); + let mut new_generation = record("new-generation", "parent"); + new_generation.parent_epoch = current_generation; + assert_eq!( + delivery_router.record(new_generation).await.unwrap(), + RecordOutcome::Recorded { + lane: None, + receipt: None + } + ); + assert!(delivery_router.pending_for("parent").iter().any(|item| { + item.task_id == "new-generation" && item.parent_epoch == current_generation + })); +} + +#[tokio::test] +async fn cross_router_cancellation_rejects_a_queued_completion_at_turn_consumption() { + use tinyagents_harness::run_queue::{QueueLane, QueueMode, RunQueue}; + use tinyinference_llm::message::Message; + + let backend: Arc<dyn StorageBackend> = Arc::new(MemoryStorage::new()); + let open = || { + Arc::new( + DriverCompletionStore::new(Arc::clone(&backend), Scope::new("shared").unwrap()) + .unwrap(), + ) as Arc<dyn CompletionStore> + }; + let producer = Arc::new(CompletionRouter::new(open())); + let canceller = CompletionRouter::new(open()); + producer.record(record("queued", "parent")).await.unwrap(); + let completion = producer.pending_for("parent").remove(0); + let queue = RunQueue::new(); + queue + .push_guarded( + QueueLane::Followup, + Message::user("child completed"), + Arc::new(crate::completions::CompletionAdmission::for_record( + Arc::clone(&producer), + &completion, + )), + ) + .await; + + canceller.cancel_parent("parent").unwrap(); + canceller.resume_parent("parent"); + + assert!( + queue + .take_admitted(QueueLane::Followup, QueueMode::All) + .await + .is_empty(), + "a completion queued under the old generation must not reach the model turn" + ); } -fn open_sqlite_router(path: &Path) -> CompletionRouter { - let scope = Scope::new("local").unwrap(); - let connection = tinystoragedrivers_connect::sqlite::SqliteConnection::open(path).unwrap(); - let docs = Arc::new(SqliteDocuments::new(connection, scope)) - as Arc<dyn tinyagents_storage_kit::DocumentStore>; - CompletionRouter::new(Arc::new(DriverCompletionStore::new(docs).unwrap())) +#[tokio::test] +async fn router_live_push_closed_before_agent_turn_never_reaches_provider() { + use tinyagents_harness::context::{RunConfig, RunContext}; + use tinyagents_harness::run_queue::RunQueue; + use tinyagents_harness::runtime::AgentHarness; + use tinyagents_harness::testkit::ScriptedModel; + use tinyinference_llm::message::Message; + use tinyinference_llm::model::ModelResponse; + + let backend: Arc<dyn StorageBackend> = Arc::new(MemoryStorage::new()); + let open = || { + Arc::new( + DriverCompletionStore::new(Arc::clone(&backend), Scope::new("shared").unwrap()) + .unwrap(), + ) as Arc<dyn CompletionStore> + }; + let producer = CompletionRouter::new(open()); + let canceller = CompletionRouter::new(open()); + let queue = Arc::new(RunQueue::new()); + producer.attach_parent("parent", queue.clone()); + let queued = producer + .record(record("queued-child", "parent")) + .await + .unwrap(); + assert!(matches!( + queued, + RecordOutcome::Recorded { lane: Some(_), .. } + )); + assert_eq!(queue.status().await.followups, 1); + canceller.cancel_parent("parent").unwrap(); + + let model = Arc::new(ScriptedModel::new(vec![ModelResponse::assistant("done")])); + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness.with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named("mock", model.clone()), + )); + let context = RunContext::new(RunConfig::new("parent-turn"), ()).with_run_queue(queue); + harness + .invoke_in_context(&(), context, vec![Message::user("start")]) + .await + .expect("the canceled queued completion is discarded before another provider turn"); + + assert_eq!( + model.requests().len(), + 1, + "a completion rejected by the durable lifecycle gate must never start a provider call" + ); +} + +#[tokio::test] +async fn router_live_push_reaches_provider_when_parent_remains_open() { + use tinyagents_harness::context::{RunConfig, RunContext}; + use tinyagents_harness::run_queue::RunQueue; + use tinyagents_harness::runtime::AgentHarness; + use tinyagents_harness::testkit::ScriptedModel; + use tinyinference_llm::message::Message; + use tinyinference_llm::model::ModelResponse; + + let backend: Arc<dyn StorageBackend> = Arc::new(MemoryStorage::new()); + let store = + Arc::new(DriverCompletionStore::new(backend, Scope::new("shared").unwrap()).unwrap()); + let producer = CompletionRouter::new(store); + let queue = Arc::new(RunQueue::new()); + producer.attach_parent("parent", queue.clone()); + let outcome = producer + .record(record("live-child", "parent")) + .await + .unwrap(); + assert!(matches!( + outcome, + RecordOutcome::Recorded { lane: Some(_), .. } + )); + assert_eq!(queue.status().await.followups, 1); + + let model = Arc::new(ScriptedModel::new(vec![ + ModelResponse::assistant("first response"), + ModelResponse::assistant("completion handled"), + ])); + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness.with_model_resolver(Arc::new( + tinyagents_harness::host::FixedModelResolver::named("mock", model.clone()), + )); + let context = RunContext::new(RunConfig::new("parent-turn"), ()).with_run_queue(queue); + harness + .invoke_in_context(&(), context, vec![Message::user("start")]) + .await + .expect("an open parent receives the live completion on its next turn"); + + let requests = model.requests(); + assert_eq!(requests.len(), 2); + assert!( + requests[1] + .messages + .iter() + .any(|message| message.text().contains("live-child")) + ); +} + +#[tokio::test] +async fn child_completion_keeps_its_spawn_generation_across_cancel_and_resume() { + let backend: Arc<dyn StorageBackend> = Arc::new(MemoryStorage::new()); + let open = || { + CompletionRouter::new(Arc::new( + DriverCompletionStore::new(Arc::clone(&backend), Scope::new("shared").unwrap()) + .unwrap(), + )) + }; + let spawned = open(); + let lifecycle = open(); + let stale_generation = spawned.admit_child_generation("parent").unwrap().unwrap(); + lifecycle.cancel_parent("parent").unwrap(); + lifecycle.resume_parent("parent"); + + let mut stale = record("late-child", "parent"); + stale.parent_epoch = stale_generation; + assert_eq!( + spawned.record(stale).await.unwrap(), + RecordOutcome::Suppressed, + "a child finishing under its original lifecycle cannot be relabeled into a later generation" + ); + assert!(spawned.pending_for("parent").is_empty()); + + let current_generation = lifecycle.admit_child_generation("parent").unwrap().unwrap(); + let mut current = record("current-child", "parent"); + current.parent_epoch = current_generation; + assert!(matches!( + lifecycle.record(current).await.unwrap(), + RecordOutcome::Recorded { .. } + )); + assert!( + lifecycle.pending_for("parent").iter().any(|row| { + row.task_id == "current-child" && row.parent_epoch == current_generation + }) + ); +} + +#[tokio::test] +async fn same_parent_batch_across_routers_shares_one_lifecycle_admission_permit() { + use tinyagents_harness::run_queue::{QueueLane, QueueMode, RunQueue}; + use tinyinference_llm::message::Message; + + let shared_memory = MemoryStorage::new(); + let backend_a: Arc<dyn StorageBackend> = Arc::new(shared_memory.clone()); + let backend_b: Arc<dyn StorageBackend> = Arc::new(shared_memory); + let scope = Scope::new("shared").unwrap(); + let producer_store = open_driver_completion_store_with_coordination_domain( + Arc::clone(&backend_a), + scope.clone(), + "shared-backend-aliases", + ) + .unwrap(); + let observer_store = open_driver_completion_store_with_coordination_domain( + backend_b, + scope, + "shared-backend-aliases", + ) + .unwrap(); + let producer = Arc::new(CompletionRouter::new(producer_store)); + let observer = Arc::new(CompletionRouter::new(observer_store)); + producer.record(record("first", "parent")).await.unwrap(); + producer.record(record("second", "parent")).await.unwrap(); + let records = producer.pending_for("parent"); + assert_eq!(records.len(), 2); + + let queue = RunQueue::new(); + for (index, completion) in records.into_iter().enumerate() { + let admission_router = if index == 0 { + Arc::clone(&producer) + } else { + Arc::clone(&observer) + }; + queue + .push_guarded( + QueueLane::Followup, + Message::user(completion.task_id.clone()), + Arc::new(crate::completions::CompletionAdmission::for_record( + admission_router, + &completion, + )), + ) + .await; + } + + let admitted = queue + .take_admitted(QueueLane::Followup, QueueMode::All) + .await; + assert_eq!(admitted.len(), 2); + assert_eq!(admitted[0].value, Message::user("first")); + assert_eq!(admitted[1].value, Message::user("second")); +} + +#[tokio::test] +async fn batch_admission_does_not_alias_distinct_scopes_and_parent_keys() { + use tinyagents_harness::run_queue::{QueueLane, QueueMode, RunQueue}; + use tinyinference_llm::message::Message; + + let backend: Arc<dyn StorageBackend> = Arc::new(MemoryStorage::new()); + let scope_a = Arc::new( + DriverCompletionStore::new(Arc::clone(&backend), Scope::new("user:a").unwrap()).unwrap(), + ) as Arc<dyn CompletionStore>; + let scope_b = Arc::new( + DriverCompletionStore::new(Arc::clone(&backend), Scope::new("user:a:b").unwrap()).unwrap(), + ) as Arc<dyn CompletionStore>; + let router_a = Arc::new(CompletionRouter::new(scope_a)); + let router_b = Arc::new(CompletionRouter::new(scope_b)); + router_a.record(record("a", "b:c")).await.unwrap(); + router_b.record(record("b", "c")).await.unwrap(); + let a = router_a.pending_for("b:c").remove(0); + let b = router_b.pending_for("c").remove(0); + router_b.cancel_parent("c").unwrap(); + + let queue = RunQueue::new(); + for (router, completion) in [(&router_a, &a), (&router_b, &b)] { + queue + .push_guarded( + QueueLane::Followup, + Message::user(completion.task_id.clone()), + Arc::new(crate::completions::CompletionAdmission::for_record( + Arc::clone(router), + completion, + )), + ) + .await; + } + + let admitted = queue + .take_admitted(QueueLane::Followup, QueueMode::All) + .await; + assert_eq!( + admitted + .iter() + .map(|item| item.value.clone()) + .collect::<Vec<_>>(), + [Message::user("a")], + "a closed tenant-parent must not reuse another tenant's batch permit" + ); +} + +/// A router over the SQLite file at `path`, opened afresh (a simulated restart). +fn open_router(path: &std::path::Path) -> CompletionRouter { + let storage = MemoryStorage::open(path); + let store = + DriverCompletionStore::new(Arc::new(storage), Scope::new("local").unwrap()).unwrap(); + CompletionRouter::new(Arc::new(store)) } #[tokio::test] async fn pending_completions_survive_a_reopen() { - let storage = MemoryBackend::new(); + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("completions.sqlite"); { - let router = open_router(&storage); + let router = open_router(&path); router.record(record("t1", "p")).await.unwrap(); router .record(record("t2", "p").with_notify_mode(NotifyMode::Off)) @@ -56,7 +673,7 @@ async fn pending_completions_survive_a_reopen() { .unwrap(); router.record(record("t3", "other")).await.unwrap(); } - let router = open_router(&storage); + let router = open_router(&path); let mut pending = router.pending_for("p"); pending.sort_by(|a, b| a.task_id.cmp(&b.task_id)); assert_eq!( @@ -71,38 +688,278 @@ async fn pending_completions_survive_a_reopen() { assert_eq!(batch.len(), 1, "the redelivery path finds t1 again"); } +#[tokio::test] +async fn two_routers_on_one_backend_do_not_claim_the_same_completion() { + let storage = MemoryStorage::new(); + let backend: Arc<dyn StorageBackend> = Arc::new(storage); + let first = CompletionRouter::new(Arc::new( + DriverCompletionStore::new(Arc::clone(&backend), Scope::new("shared").unwrap()).unwrap(), + )); + let second = CompletionRouter::new(Arc::new( + DriverCompletionStore::new(backend, Scope::new("shared").unwrap()).unwrap(), + )); + first.record(record("shared-task", "parent")).await.unwrap(); + + let first_claim = first.claim_pending("parent", 1).unwrap(); + let second_claim = second.claim_pending("parent", 1).unwrap(); + + assert_eq!(first_claim.len(), 1); + assert!( + second_claim.is_empty(), + "another router already owns this visible claim" + ); +} + +#[tokio::test] +async fn a_live_push_claim_blocks_a_second_driver_router() { + let storage = MemoryStorage::new(); + let backend: Arc<dyn StorageBackend> = Arc::new(storage); + let first = CompletionRouter::new(Arc::new( + DriverCompletionStore::new(Arc::clone(&backend), Scope::new("shared").unwrap()).unwrap(), + )); + let second = CompletionRouter::new(Arc::new( + DriverCompletionStore::new(backend, Scope::new("shared").unwrap()).unwrap(), + )); + first.attach_parent( + "parent", + Arc::new(tinyagents_harness::run_queue::RunQueue::< + tinyinference_llm::message::Message, + >::new()), + ); + first.record(record("live-push", "parent")).await.unwrap(); + + assert!(second.claim_pending("parent", 1).unwrap().is_empty()); +} + +#[test] +fn driver_claims_choose_oldest_before_task_id_tiebreaking() { + use crate::completions::CompletionClaimMode; + + let storage = MemoryStorage::new(); + let store = + DriverCompletionStore::new(Arc::new(storage), Scope::new("shared").unwrap()).unwrap(); + let now = SystemTime::now(); + put_record(&store, &record("z-older", "p").with_finished_at(now)); + put_record( + &store, + &record("a-newer", "p").with_finished_at(now + Duration::from_secs(1)), + ); + + let gate = parent_gate(&store, "p"); + let claimed = store + .claim_pending(crate::completions::PendingClaimRequest { + parent_key: "p".into(), + parent_epoch: 0, + gate: gate.clone(), + mode: CompletionClaimMode::Followup, + max: 1, + token: "owner".into(), + visible_after: now + Duration::from_secs(30), + }) + .unwrap(); + assert_eq!( + claimed + .iter() + .map(|record| record.task_id.as_str()) + .collect::<Vec<_>>(), + ["z-older"] + ); +} + +#[test] +fn conditional_driver_create_preserves_an_existing_claim() { + use crate::completions::{ClaimResolution, CompletionClaimMode}; + + let storage = MemoryStorage::new(); + let store = + DriverCompletionStore::new(Arc::new(storage), Scope::new("shared").unwrap()).unwrap(); + let original = record("claim-owner", "parent"); + let gate = parent_gate(&store, "parent"); + assert!(store.insert_if_absent(&original, &gate).unwrap()); + let claimed = store + .claim_task(crate::completions::TaskClaimRequest { + task_id: "claim-owner".into(), + parent_key: "parent".into(), + parent_epoch: 0, + gate: gate.clone(), + mode: CompletionClaimMode::Followup, + token: "owner-token".into(), + visible_after: SystemTime::now() + Duration::from_secs(30), + }) + .unwrap() + .unwrap(); + + let other_parent_gate = parent_gate(&store, "other-parent"); + assert!( + !store + .insert_if_absent(&record("claim-owner", "other-parent"), &other_parent_gate,) + .unwrap() + ); + let persisted = store.get("claim-owner").unwrap(); + assert_eq!(persisted.parent_key, claimed.parent_key); + assert_eq!(persisted.attempts, claimed.attempts); + assert_eq!(persisted.receipt(), claimed.receipt()); + assert_eq!(persisted.result, claimed.result); + let receipt = claimed.receipt().unwrap(); + assert_eq!( + store + .resolve_claim(crate::completions::ResolveClaimRequest { + task_id: receipt.task_id, + parent_key: "parent".into(), + parent_epoch: receipt.parent_epoch, + gate, + token: receipt.token, + resolution: ClaimResolution::Delivered, + }) + .unwrap() + .unwrap() + .state, + CompletionState::Delivered + ); +} + +#[test] +fn conditional_driver_create_does_not_replace_a_tombstone() { + let storage = MemoryStorage::new(); + let store = + DriverCompletionStore::new(Arc::new(storage), Scope::new("shared").unwrap()).unwrap(); + let tombstone = CompletionRecord::tombstone_stub("deleted"); + put_record(&store, &tombstone); + let gate = parent_gate(&store, "parent"); + + assert!( + !store + .insert_if_absent(&record("deleted", "parent"), &gate) + .unwrap() + ); + let persisted = store.get("deleted").unwrap(); + assert_eq!(persisted.task_id, tombstone.task_id); + assert_eq!(persisted.state, CompletionState::Tombstoned); + assert!(persisted.parent_key.is_empty()); +} + +#[tokio::test] +async fn expired_claim_can_be_reclaimed_and_stale_owner_cannot_settle_it() { + use crate::completions::{ClaimResolution, CompletionClaimMode}; + + let storage = MemoryStorage::new(); + let backend: Arc<dyn StorageBackend> = Arc::new(storage); + let first = + DriverCompletionStore::new(Arc::clone(&backend), Scope::new("shared").unwrap()).unwrap(); + let second = DriverCompletionStore::new(backend, Scope::new("shared").unwrap()).unwrap(); + let record = record("lease-task", "parent"); + put_record(&first, &record); + let now = SystemTime::now(); + let first_gate = parent_gate(&first, "parent"); + let first_claim = first + .claim_pending(crate::completions::PendingClaimRequest { + parent_key: "parent".into(), + parent_epoch: 0, + gate: first_gate.clone(), + mode: CompletionClaimMode::Followup, + max: 1, + token: "first".into(), + visible_after: now - Duration::from_secs(1), + }) + .unwrap() + .pop() + .unwrap(); + first.release_parent_gate("parent", &first_gate).unwrap(); + let first_token = first_claim.claim.unwrap().token; + let second_gate = parent_gate(&second, "parent"); + assert!( + first + .resolve_claim(crate::completions::ResolveClaimRequest { + task_id: "lease-task".into(), + parent_key: "parent".into(), + parent_epoch: first_claim.parent_epoch, + gate: second_gate.clone(), + token: first_token.clone(), + resolution: ClaimResolution::Delivered, + },) + .unwrap() + .is_none(), + "an expired owner cannot settle before another router reclaims" + ); + let second_claim = second + .claim_pending(crate::completions::PendingClaimRequest { + parent_key: "parent".into(), + parent_epoch: 0, + gate: second_gate.clone(), + mode: CompletionClaimMode::Followup, + max: 1, + token: "second".into(), + visible_after: now + Duration::from_secs(10), + }) + .unwrap(); + assert_eq!( + second_claim.len(), + 1, + "the expired visibility lease is reclaimable" + ); + let second_token = second_claim[0].claim.as_ref().unwrap().token.clone(); + + assert!( + first + .resolve_claim(crate::completions::ResolveClaimRequest { + task_id: "lease-task".into(), + parent_key: "parent".into(), + parent_epoch: first_claim.parent_epoch, + gate: second_gate.clone(), + token: first_token, + resolution: ClaimResolution::Delivered, + },) + .unwrap() + .is_none() + ); + let settled = second + .resolve_claim(crate::completions::ResolveClaimRequest { + task_id: "lease-task".into(), + parent_key: "parent".into(), + parent_epoch: second_claim[0].parent_epoch, + gate: second_gate, + token: second_token, + resolution: ClaimResolution::Delivered, + }) + .unwrap() + .unwrap(); + assert_eq!(settled.state, CompletionState::Delivered); +} + #[tokio::test] async fn attempts_and_give_up_survive_a_reopen() { - let storage = MemoryBackend::new(); + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("completions.sqlite"); { - let router = open_router(&storage).with_max_attempts(2); + let router = open_router(&path).with_max_attempts(2); router.record(record("t1", "p")).await.unwrap(); - router.claim_pending("p", 1).unwrap(); - router.mark_failed(&["t1"]).unwrap(); + let batch = router.claim_pending("p", 1).unwrap(); + router.mark_failed(&receipts(&batch)).unwrap(); } - let router = open_router(&storage).with_max_attempts(2); + let router = open_router(&path).with_max_attempts(2); let batch = router.claim_pending("p", 1).unwrap(); assert_eq!(batch[0].attempts, 2, "the earlier attempt was kept"); - assert_eq!(router.mark_failed(&["t1"]).unwrap().len(), 1); + assert_eq!(router.mark_failed(&receipts(&batch)).unwrap().len(), 1); drop(router); - let router = open_router(&storage); + let router = open_router(&path); assert!(router.pending_for("p").is_empty()); assert!(router.claim_pending("p", 1).unwrap().is_empty()); } #[tokio::test] async fn delivered_and_tombstoned_stay_deduped_after_a_reopen() { - let storage = MemoryBackend::new(); + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("completions.sqlite"); { - let router = open_router(&storage); + let router = open_router(&path); router.record(record("done", "p")).await.unwrap(); - router.claim_pending("p", 1).unwrap(); - router.mark_delivered(&["done"]).unwrap(); + let batch = router.claim_pending("p", 1).unwrap(); + router.mark_delivered(&receipts(&batch)).unwrap(); router.record(record("waited", "p")).await.unwrap(); router.tombstone("waited").unwrap(); router.tombstone("early").unwrap(); } - let router = open_router(&storage); + let router = open_router(&path); for task in ["done", "waited", "early"] { let outcome = router.record(record(task, "p")).await.unwrap(); assert!( @@ -116,134 +973,22 @@ async fn delivered_and_tombstoned_stay_deduped_after_a_reopen() { assert!(router.pending_for("p").is_empty()); } -#[tokio::test] -async fn a_settled_record_is_persisted_without_its_result() { - let storage = MemoryBackend::new(); - let docs = docs_of(&storage, "local"); - let store = Arc::new(DriverCompletionStore::new(docs.clone()).unwrap()); - let router = CompletionRouter::new(store.clone()).with_max_attempts(1); - let doc_of = |task: &str| { - let docs = docs.clone(); - let id = task_id_key(task).unwrap(); - async move { docs.get(COMPLETIONS, &id).await.unwrap().unwrap().doc } - }; - for (index, task) in ["delivered", "tombstoned", "gave-up", "cancelled"] - .into_iter() - .enumerate() - { - let mut record = record(task, "p"); - record.finished_at = UNIX_EPOCH + Duration::from_secs(index as u64 + 1); - record.result.omitted_chars = 7; - record.result.artifact = Some(crate::completions::CompletionArtifact { - id: format!("art-{task}"), - ..Default::default() - }); - router.record(record).await.unwrap(); - let persisted = doc_of(task).await; - assert_eq!(persisted["result"]["text"], format!("result of {task}")); - assert_eq!(persisted["result"]["artifact"]["id"], format!("art-{task}")); - } - // Oldest first: the claim leases the first three; `cancelled` stays - // unclaimed for `cancel_parent` to withdraw. - let claimed = router.claim_pending("p", 3).unwrap(); - assert_eq!(claimed.len(), 3); - router.mark_delivered(&["delivered"]).unwrap(); - assert_eq!(router.mark_failed(&["gave-up"]).unwrap().len(), 1); - router.tombstone("tombstoned").unwrap(); - assert_eq!(router.cancel_parent("p").unwrap(), 1); - for (task, state) in [ - ("delivered", "delivered"), - ("tombstoned", "tombstoned"), - ("gave-up", "gave_up"), - ("cancelled", "tombstoned"), - ] { - let persisted = doc_of(task).await; - assert_eq!(persisted["state"], state, "{task}"); - assert_eq!(persisted["task_id"], task, "{task}"); - assert_eq!(persisted["parent_key"], "p", "{task}"); - assert_eq!(persisted["result"]["text"], "", "{task}"); - assert_eq!(persisted["result"]["omitted_chars"], 0, "{task}"); - assert!( - persisted["result"].get("artifact").is_none(), - "{task}: {}", - persisted["result"] - ); - let read = store.get(task).unwrap(); - assert_eq!(read.result, CompletionResult::default(), "{task}"); - assert_eq!(read.task_id, task, "{task}"); - } -} - -#[tokio::test] -async fn settled_records_reopen_without_their_result_and_stay_deduped() { - let dir = tempfile::tempdir().unwrap(); - let path = dir.path().join("completions.sqlite"); - { - let router = open_sqlite_router(&path).with_max_attempts(1); - router.record(record("delivered", "p")).await.unwrap(); - router.record(record("gave-up", "p")).await.unwrap(); - router.record(record("waited", "p")).await.unwrap(); - router.record(record("live", "p")).await.unwrap(); - router.claim_pending("p", 2).unwrap(); - router.mark_delivered(&["delivered"]).unwrap(); - let gave_up = router.mark_failed(&["gave-up"]).unwrap(); - assert_eq!(gave_up[0].result.text, "result of gave-up"); - router.tombstone("waited").unwrap(); - } - let router = open_sqlite_router(&path); - let store = router.store(); - let pending = router.pending_for("p"); - assert_eq!(pending.len(), 1); - assert_eq!(pending[0].task_id, "live"); - assert_eq!(pending[0].result.text, "result of live"); - for (task, state) in [ - ("delivered", CompletionState::Delivered), - ("gave-up", CompletionState::GaveUp), - ("waited", CompletionState::Tombstoned), - ] { - let settled = store.get(task).unwrap(); - assert_eq!(settled.state, state, "{task}"); - assert_eq!(settled.result, CompletionResult::default(), "{task}"); - assert_eq!( - settled.attempts, - u32::from(state != CompletionState::Tombstoned), - "{task}" - ); - let outcome = router.record(record(task, "p")).await.unwrap(); - let expected = if state == CompletionState::Tombstoned { - RecordOutcome::Suppressed - } else { - RecordOutcome::Duplicate - }; - assert_eq!(outcome, expected, "{task}"); - } - assert_eq!(router.claim_pending("p", 10).unwrap().len(), 1); - assert_eq!(router.cancel_parent("p").unwrap(), 1); - drop(router); - let router = open_sqlite_router(&path); - assert!(router.pending_for("p").is_empty()); - assert_eq!( - router.record(record("live", "p")).await.unwrap(), - RecordOutcome::Suppressed, - "a withdrawn record stays tombstoned across a reopen" - ); -} - #[tokio::test] async fn the_store_keeps_the_latest_record_per_task() { - let storage = MemoryBackend::new(); - let store = DriverCompletionStore::new(docs_of(&storage, "local")).unwrap(); + let storage = MemoryStorage::new(); + let store = + DriverCompletionStore::new(Arc::new(storage), Scope::new("local").unwrap()).unwrap(); let mut first = record("t1", "p"); - store.put(&first).unwrap(); + put_record(&store, &first); first.state = CompletionState::Delivered; first.attempts = 3; - store.put(&first).unwrap(); + put_record(&store, &first); assert_eq!(store.list(None).len(), 1); let read = store.get("t1").unwrap(); assert_eq!(read.state, CompletionState::Delivered); assert_eq!(read.attempts, 3); assert!(store.get("missing").is_none()); - store.put(&record("t2", "q")).unwrap(); + put_record(&store, &record("t2", "q")); assert_eq!(store.list(Some("p")).len(), 1); assert_eq!(store.list(Some("nobody")).len(), 0); assert_eq!(store.list(None).len(), 2); @@ -254,9 +999,10 @@ async fn migration_listing_fails_on_corrupt_records_while_legacy_listing_skips_t use serde_json::json; use tinyagents_storage_kit::Precondition; - let storage = MemoryBackend::new(); - let store = DriverCompletionStore::new(docs_of(&storage, "local")).unwrap(); - store.put(&record("corrupt", "p")).unwrap(); + let storage = MemoryStorage::new(); + let store = DriverCompletionStore::new(Arc::new(storage.clone()), Scope::new("local").unwrap()) + .unwrap(); + put_record(&store, &record("corrupt", "p")); store .docs .put( @@ -280,17 +1026,18 @@ async fn migration_listing_fails_on_corrupt_records_while_legacy_listing_skips_t #[tokio::test] async fn compaction_drops_settled_records_only_and_keeps_pending_work() { - let storage = MemoryBackend::new(); - let router = open_router(&storage); + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("completions.sqlite"); + let router = open_router(&path); router.record(record("old", "p")).await.unwrap(); - router.claim_pending("p", 1).unwrap(); - router.mark_delivered(&["old"]).unwrap(); + let batch = router.claim_pending("p", 1).unwrap(); + router.mark_delivered(&receipts(&batch)).unwrap(); router.record(record("live", "p")).await.unwrap(); tokio::time::sleep(Duration::from_millis(20)).await; assert_eq!(router.compact(Duration::from_millis(1)).unwrap(), 1); router.record(record("after", "p")).await.unwrap(); drop(router); - let router = open_router(&storage); + let router = open_router(&path); let mut ids: Vec<_> = router .pending_for("p") .into_iter() @@ -312,11 +1059,9 @@ fn clocked_store() -> (DriverCompletionStore, Arc<std::sync::atomic::AtomicU64>) let now = Arc::clone(&now); Arc::new(move || now.load(std::sync::atomic::Ordering::SeqCst)) }; - let connection = MemoryConnection::with_clock(clock); - let docs = connection - .for_scope(Scope::new("local").unwrap()) - .documents(); - let store = DriverCompletionStore::new(docs).unwrap(); + let storage = MemoryStorage::with_clock(clock); + let store = + DriverCompletionStore::new(Arc::new(storage), Scope::new("local").unwrap()).unwrap(); (store, now) } @@ -336,7 +1081,7 @@ fn compaction_drops_settled_records_older_than_the_window_like_the_in_memory_sto let mut r = record(task, "p"); r.state = state; r.updated_at = at; - store.put(&r).unwrap(); + put_record(store, &r); } }; now.store(ms(old), std::sync::atomic::Ordering::SeqCst); @@ -345,8 +1090,7 @@ fn compaction_drops_settled_records_older_than_the_window_like_the_in_memory_sto now.store(ms(real_now), std::sync::atomic::Ordering::SeqCst); let mut recent = record("recent", "p"); recent.state = CompletionState::Delivered; - driver.put(&recent).unwrap(); - recent.updated_at = real_now; + put_record(&driver, &recent); memory.put(&recent).unwrap(); assert_eq!(memory.compact(Duration::from_secs(1)).unwrap(), 2); assert_eq!(driver.compact(Duration::from_secs(1)).unwrap(), 2); @@ -360,21 +1104,25 @@ fn compaction_drops_settled_records_older_than_the_window_like_the_in_memory_sto #[tokio::test] async fn a_cancelled_parent_stays_cancelled_after_a_reopen() { - let storage = MemoryBackend::new(); + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("completions.sqlite"); { - let router = open_router(&storage); + let router = open_router(&path); router.cancel_parent("p").unwrap(); } - let router = open_router(&storage); + let router = open_router(&path); assert_eq!( router.record(record("late", "p")).await.unwrap(), RecordOutcome::Suppressed ); router.resume_parent("p"); + let generation = router.admit_child_generation("p").unwrap().unwrap(); drop(router); - let router = open_router(&storage); + let router = open_router(&path); + let mut later = record("later", "p"); + later.parent_epoch = generation; assert!(matches!( - router.record(record("later", "p")).await.unwrap(), + router.record(later).await.unwrap(), RecordOutcome::Recorded { .. } )); } @@ -382,10 +1130,16 @@ async fn a_cancelled_parent_stays_cancelled_after_a_reopen() { #[tokio::test] async fn an_unreadable_record_is_skipped_without_hiding_the_rest() { use tinyagents_storage_kit::Precondition; - let storage = MemoryBackend::new(); - let docs = docs_of(&storage, "local"); - let store = DriverCompletionStore::new(docs.clone()).unwrap(); - store.put(&record("good", "p")).unwrap(); + let storage = MemoryStorage::new(); + let docs = Arc::clone( + storage + .for_scope(&Scope::new("local").unwrap()) + .unwrap() + .documents(), + ); + let store = DriverCompletionStore::new(Arc::new(storage.clone()), Scope::new("local").unwrap()) + .unwrap(); + put_record(&store, &record("good", "p")); docs.put( "tasks.completions", "corrupt", @@ -401,27 +1155,37 @@ async fn an_unreadable_record_is_skipped_without_hiding_the_rest() { #[tokio::test] async fn scopes_keep_completions_apart() { - let storage = MemoryBackend::new(); - let alice = DriverCompletionStore::new(docs_of(&storage, "alice")).unwrap(); - let bob = DriverCompletionStore::new(docs_of(&storage, "bob")).unwrap(); - alice.put(&record("t", "p")).unwrap(); + let storage = MemoryStorage::new(); + let backend: Arc<dyn StorageBackend> = Arc::new(storage); + let alice = + DriverCompletionStore::new(Arc::clone(&backend), Scope::new("alice").unwrap()).unwrap(); + let bob = DriverCompletionStore::new(backend, Scope::new("bob").unwrap()).unwrap(); + put_record(&alice, &record("t", "p")); assert!(bob.get("t").is_none()); assert!(bob.list(None).is_empty()); } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn the_sync_trait_works_from_a_multi_thread_runtime() { - let store = DriverCompletionStore::new(docs_of(&MemoryBackend::new(), "local")).unwrap(); - store.put(&record("mt", "p")).unwrap(); + let store = + DriverCompletionStore::new(Arc::new(MemoryStorage::new()), Scope::new("local").unwrap()) + .unwrap(); + put_record(&store, &record("mt", "p")); assert!(store.get("mt").is_some()); } #[tokio::test] async fn completions_are_v2_documents_with_key_ids_and_metadata_update_times() { - let storage = MemoryBackend::new(); - let docs = docs_of(&storage, "local"); - let store = DriverCompletionStore::new(docs.clone()).unwrap(); - store.put(&record("c/1", "p")).unwrap(); + let storage = MemoryStorage::new(); + let docs = Arc::clone( + storage + .for_scope(&Scope::new("local").unwrap()) + .unwrap() + .documents(), + ); + let store = DriverCompletionStore::new(Arc::new(storage.clone()), Scope::new("local").unwrap()) + .unwrap(); + put_record(&store, &record("c/1", "p")); let row = docs .get("tasks.completions", "c%2F1") .await diff --git a/crates/tinyagents-tasks/src/completions/durable_lease.rs b/crates/tinyagents-tasks/src/completions/durable_lease.rs new file mode 100644 index 000000000..46f13c862 --- /dev/null +++ b/crates/tinyagents-tasks/src/completions/durable_lease.rs @@ -0,0 +1,301 @@ +//! TSD-backed ownership for one tenant-scoped parent-thread turn. + +use std::future::Future; +use std::pin::Pin; +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::time::Duration; + +use tinyagents_storage_kit::primitives::Lease; +use tinyagents_storage_kit::{Leases, Scope, StorageBackend, StorageError}; +use tokio::sync::watch; +use tokio::task::JoinHandle; + +use super::DeliveryFence; + +/// Factory for durable thread-turn leases stored in `tsd.leases`. +/// +/// Every process serving the same database must use the same explicit +/// `coordination_domain`, lease scope, tenant, and thread key. The domain is +/// host configuration, never model input. Different databases must use +/// different domains even when their tenant/thread text matches. +#[derive(Clone)] +pub struct DurableThreadDeliveryLeases { + leases: Leases, + lease_scope: Scope, + coordination_domain: String, + ttl: Duration, + renewal_interval: Duration, +} + +impl std::fmt::Debug for DurableThreadDeliveryLeases { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("DurableThreadDeliveryLeases") + .field("lease_scope", &self.lease_scope) + .field("coordination_domain", &self.coordination_domain) + .field("ttl", &self.ttl) + .field("renewal_interval", &self.renewal_interval) + .finish_non_exhaustive() + } +} + +impl DurableThreadDeliveryLeases { + /// Opens the durable lease collection over `lease_scope`. + /// + /// `lease_scope` is normally the trusted cluster scope shared by every + /// process serving the database. `coordination_domain` names that shared + /// deployment namespace and must not be derived from a thread or tenant + /// payload. + pub fn new( + backend: &dyn StorageBackend, + lease_scope: Scope, + coordination_domain: impl Into<String>, + ttl: Duration, + ) -> Result<Self, StorageError> { + let coordination_domain = coordination_domain.into(); + if coordination_domain.trim().is_empty() || ttl < Duration::from_millis(3) { + return Err(StorageError::invalid_input( + "durable thread lease requires a domain and a TTL of at least 3 ms", + )); + } + let documents = Arc::clone(backend.for_scope(&lease_scope)?.documents()); + let renewal_interval = (ttl / 3).max(Duration::from_millis(1)); + Ok(Self { + leases: Leases::new(&documents), + lease_scope, + coordination_domain, + ttl, + renewal_interval, + }) + } + + /// Attempts the tenant/thread lease. `None` means another live holder owns + /// it; backend errors remain errors and never fall back to local locking. + pub async fn try_acquire( + &self, + tenant: &Scope, + thread_key: &str, + ) -> Result<Option<DurableThreadDeliveryLease>, StorageError> { + if thread_key.is_empty() { + return Err(StorageError::invalid_input( + "durable thread lease key must not be empty", + )); + } + // TSD v2 scopes are opaque identifiers. Tenant scopes are single path + // segments; agent/module child scopes contain `/`, and cluster is the + // shared coordination namespace rather than a tenant. + if tenant.as_str() == "cluster" || tenant.as_str().contains('/') { + return Err(StorageError::invalid_input( + "durable thread lease requires a tenant scope without a principal", + )); + } + let owner = "tinyagents.thread-turn"; + let key = serde_json::to_string(&( + self.coordination_domain.as_str(), + tenant.as_str(), + thread_key, + )) + .map_err(|error| StorageError::invalid_input(error.to_string()))?; + let holder = uuid::Uuid::new_v4().to_string(); + let ttl_ms = duration_ms(self.ttl)?; + let Some(lease) = self.leases.acquire(owner, &key, &holder, ttl_ms).await? else { + return Ok(None); + }; + let epoch = lease.epoch; + let (admission, receiver) = watch::channel(true); + let renewer = spawn_renewer(LeaseRenewalContext { + leases: self.leases.clone(), + ttl: self.ttl, + interval: self.renewal_interval, + owner: owner.to_owned(), + key: key.clone(), + holder: holder.clone(), + epoch: lease.epoch, + admission: admission.clone(), + }); + Ok(Some(DurableThreadDeliveryLease { + leases: self.leases.clone(), + lease_scope: self.lease_scope.clone(), + owner: owner.to_owned(), + key, + holder, + lease, + epoch, + admission, + receiver, + renewer: Some(renewer), + released: AtomicBool::new(false), + })) + } +} + +/// One renewable TSD lease tenure. Renewal never reacquires a lost lease; a +/// loss permanently closes admission for this handle. +pub struct DurableThreadDeliveryLease { + leases: Leases, + lease_scope: Scope, + owner: String, + key: String, + holder: String, + lease: Lease, + epoch: u64, + admission: watch::Sender<bool>, + receiver: watch::Receiver<bool>, + renewer: Option<JoinHandle<()>>, + released: AtomicBool, +} + +impl std::fmt::Debug for DurableThreadDeliveryLease { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("DurableThreadDeliveryLease") + .field("epoch", &self.epoch) + .field("admitted", &self.admitted()) + .finish_non_exhaustive() + } +} + +impl DurableThreadDeliveryLease { + /// The stable lease epoch used to fence session writes. + pub fn epoch(&self) -> u64 { + self.epoch + } + + /// Whether renewal still confirms this exact tenure. + pub fn admitted(&self) -> bool { + *self.receiver.borrow() && !self.released.load(Ordering::Acquire) + } + + /// Creates the native epoch fence for storage in `lease_scope`. + pub fn write_condition(&self) -> Result<tinyagents_storage_kit::WriteCondition, StorageError> { + self.lease.write_condition(self.lease_scope.clone()) + } + + /// Permanently cancels the supplied token if this exact lease tenure is lost. + /// + /// Hosts pass the token through the full turn/runtime operation. Renewal + /// failure closes the token; it never reacquires a lease or restores + /// admission on this handle. + pub fn cancel_on_loss( + &self, + cancellation: tinyagents_harness::CancellationToken, + ) -> JoinHandle<()> { + let mut receiver = self.receiver.clone(); + tokio::spawn(async move { + loop { + if !*receiver.borrow() { + cancellation.cancel(); + return; + } + if receiver.changed().await.is_err() { + cancellation.cancel(); + return; + } + } + }) + } + + /// Releases this exact holder/epoch and stops renewal. A later holder's + /// acquisition advances the epoch and fences every retained old writer. + pub async fn release(mut self) -> Result<bool, StorageError> { + self.released.store(true, Ordering::Release); + let _ = self.admission.send(false); + if let Some(renewer) = self.renewer.take() { + renewer.abort(); + let _ = renewer.await; + } + self.leases + .release(&self.owner, &self.key, &self.holder, self.epoch) + .await + } +} + +impl DeliveryFence for DurableThreadDeliveryLease { + fn admitted(&self) -> bool { + self.admitted() + } + + fn lost(&self) -> Pin<Box<dyn Future<Output = ()> + Send + '_>> { + let mut receiver = self.receiver.clone(); + Box::pin(async move { + loop { + if !*receiver.borrow() { + return; + } + if receiver.changed().await.is_err() { + return; + } + } + }) + } +} + +impl Drop for DurableThreadDeliveryLease { + fn drop(&mut self) { + if self.released.swap(true, Ordering::AcqRel) { + return; + } + let _ = self.admission.send(false); + if let Some(renewer) = self.renewer.take() { + renewer.abort(); + } + if let Ok(runtime) = tokio::runtime::Handle::try_current() { + let leases = self.leases.clone(); + let owner = self.owner.clone(); + let key = self.key.clone(); + let holder = self.holder.clone(); + let epoch = self.epoch; + runtime.spawn(async move { + if let Err(error) = leases.release(&owner, &key, &holder, epoch).await { + tracing::warn!(error = %error, "durable thread lease release failed"); + } + }); + } + } +} + +struct LeaseRenewalContext { + leases: Leases, + ttl: Duration, + interval: Duration, + owner: String, + key: String, + holder: String, + epoch: u64, + admission: watch::Sender<bool>, +} + +fn spawn_renewer(context: LeaseRenewalContext) -> JoinHandle<()> { + let LeaseRenewalContext { + leases, + ttl, + interval, + owner, + key, + holder, + epoch, + admission, + } = context; + tokio::spawn(async move { + let Ok(ttl_ms) = duration_ms(ttl) else { + let _ = admission.send(false); + return; + }; + loop { + tokio::time::sleep(interval).await; + match leases.renew(&owner, &key, &holder, epoch, ttl_ms).await { + Ok(Some(_)) => {} + Ok(None) | Err(_) => { + let _ = admission.send(false); + return; + } + } + } + }) +} + +fn duration_ms(duration: Duration) -> Result<i64, StorageError> { + i64::try_from(duration.as_millis()) + .map_err(|_| StorageError::invalid_input("durable thread lease TTL is out of range")) +} diff --git a/crates/tinyagents-tasks/src/completions/durable_lease_tests.rs b/crates/tinyagents-tasks/src/completions/durable_lease_tests.rs new file mode 100644 index 000000000..d95f8ee22 --- /dev/null +++ b/crates/tinyagents-tasks/src/completions/durable_lease_tests.rs @@ -0,0 +1,136 @@ +use super::*; +use crate::conditional_sqlite::MemoryStorage; +use std::sync::Arc; +use std::time::Duration; +use tinyagents_harness::CancellationToken; +use tinyagents_storage_kit::{Scope, StorageBackend}; + +#[tokio::test] +async fn releasing_a_turn_lease_cancels_its_admission_and_successor_gets_new_epoch() { + let backend: Arc<dyn StorageBackend> = Arc::new(MemoryStorage::new()); + let lease_scope = Scope::new("cluster").unwrap(); + let tenant = Scope::new("u:alice").unwrap(); + let first_factory = DurableThreadDeliveryLeases::new( + backend.as_ref(), + lease_scope.clone(), + "test-cluster", + Duration::from_secs(30), + ) + .unwrap(); + let second_factory = DurableThreadDeliveryLeases::new( + backend.as_ref(), + lease_scope, + "test-cluster", + Duration::from_secs(30), + ) + .unwrap(); + + let first = first_factory + .try_acquire(&tenant, "thread-1") + .await + .unwrap() + .unwrap(); + let first_epoch = first.epoch(); + let turn_cancellation = CancellationToken::new(); + let loss_listener = first.cancel_on_loss(turn_cancellation.clone()); + + assert!( + second_factory + .try_acquire(&tenant, "thread-1") + .await + .unwrap() + .is_none(), + "a live thread turn has one durable holder" + ); + assert!(first.release().await.unwrap()); + tokio::time::timeout(Duration::from_secs(1), turn_cancellation.cancelled()) + .await + .expect("loss of the held tenure cancels the executing turn"); + loss_listener.await.unwrap(); + + let successor = second_factory + .try_acquire(&tenant, "thread-1") + .await + .unwrap() + .unwrap(); + assert!(successor.epoch() > first_epoch); + assert_eq!( + successor.write_condition().unwrap().clauses()[0].collection(), + "tsd.leases" + ); + successor.release().await.unwrap(); +} + +#[tokio::test(start_paused = true)] +async fn a_takeover_lost_during_renewal_permanently_cancels_the_turn() { + use serde_json::json; + use tinyagents_storage_kit::{Key, Precondition}; + + let backend: Arc<dyn StorageBackend> = Arc::new(MemoryStorage::new()); + let lease_scope = Scope::new("cluster").unwrap(); + let tenant = Scope::new("u:alice").unwrap(); + let factory = DurableThreadDeliveryLeases::new( + backend.as_ref(), + lease_scope.clone(), + "test-cluster", + Duration::from_millis(300), + ) + .unwrap(); + let first = factory + .try_acquire(&tenant, "thread-1") + .await + .unwrap() + .unwrap(); + let old_epoch = first.epoch(); + let turn_cancellation = CancellationToken::new(); + let loss_listener = first.cancel_on_loss(turn_cancellation.clone()); + + // Simulate another process taking over through the released document CAS + // contract. The lease renewer observes the new holder/epoch and permanently + // closes this tenure; no private clock or test-only lease hook is involved. + let encoded_key = + serde_json::to_string(&("test-cluster", tenant.as_str(), "thread-1")).unwrap(); + let lease_id = Key::of(["tinyagents.thread-turn", encoded_key.as_str()]) + .unwrap() + .as_str() + .to_owned(); + let documents = backend.for_scope(&lease_scope).unwrap().documents().clone(); + let current = documents + .get("tsd.leases", &lease_id) + .await + .unwrap() + .unwrap(); + let successor_expiry = current.doc["expires_ms"].as_i64().unwrap() + 60_000; + documents + .put( + "tsd.leases", + &lease_id, + json!({ + "holder": "successor-process", + "epoch": old_epoch + 1, + "expires_ms": successor_expiry, + }), + Precondition::Version(current.version), + ) + .await + .unwrap(); + assert!( + first.admitted(), + "the takeover is written before renewal runs" + ); + tokio::time::advance(Duration::from_millis(100)).await; + + tokio::time::timeout(Duration::from_secs(2), turn_cancellation.cancelled()) + .await + .expect("renewal detects the takeover and cancels the active turn"); + loss_listener.await.unwrap(); + assert!(!first.admitted(), "lost tenure cannot regain admission"); + assert!( + factory + .try_acquire(&tenant, "thread-1") + .await + .unwrap() + .is_none(), + "the old turn must not reacquire while a successor owns the thread" + ); +} diff --git a/crates/tinyagents-tasks/src/completions/lifecycle.rs b/crates/tinyagents-tasks/src/completions/lifecycle.rs new file mode 100644 index 000000000..ec208f9f3 --- /dev/null +++ b/crates/tinyagents-tasks/src/completions/lifecycle.rs @@ -0,0 +1,147 @@ +//! Process-local parent lifecycle authority shared across router cache entries. + +use std::collections::HashMap; +use std::sync::{Arc, Condvar, Mutex}; +use std::time::{Duration, Instant}; + +use tinyagents_harness::error::{Result, TinyAgentsError}; + +const LOCAL_GATE_WAIT: Duration = Duration::from_secs(65); + +/// A lifecycle snapshot retained for one parent in one trusted pool scope. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(super) enum ParentLifecycleStatus { + Open(u64), + Stopped(u64), + Deleted(u64), +} + +#[derive(Debug, Default)] +struct ParentLifecycleEntry { + held: bool, + status: Option<ParentLifecycleStatus>, +} + +/// Scope-owned state that outlives cached routers and their backing stores. +#[doc(hidden)] +#[derive(Debug)] +pub struct ParentLifecycle { + identity: String, + entries: Mutex<HashMap<String, ParentLifecycleEntry>>, + changed: Condvar, +} + +impl ParentLifecycle { + pub(super) fn new() -> Arc<Self> { + Self::with_identity(uuid::Uuid::new_v4().to_string()) + } + + pub(super) fn for_store(store: &dyn super::CompletionStore) -> Arc<Self> { + Self::with_identity( + store.coordination_identity().unwrap_or_else(|| { + format!("store:{:p}", store as *const dyn super::CompletionStore) + }), + ) + } + + fn with_identity(identity: String) -> Arc<Self> { + Arc::new(Self { + identity, + entries: Mutex::new(HashMap::new()), + changed: Condvar::new(), + }) + } + + pub(super) fn identity(&self) -> &str { + &self.identity + } + + pub(super) fn acquire(self: &Arc<Self>, parent_key: &str) -> Result<ParentLifecycleGuard> { + let started = Instant::now(); + let mut entries = self + .entries + .lock() + .map_err(|_| TinyAgentsError::Graph("parent lifecycle lock poisoned".into()))?; + loop { + let entry = entries.entry(parent_key.to_owned()).or_default(); + if !entry.held { + entry.held = true; + return Ok(ParentLifecycleGuard { + lifecycle: Arc::clone(self), + parent_key: parent_key.to_owned(), + }); + } + let remaining = LOCAL_GATE_WAIT.saturating_sub(started.elapsed()); + if remaining.is_zero() { + return Err(TinyAgentsError::Graph(format!( + "completion router: timed out waiting for local parent admission gate `{parent_key}`" + ))); + } + let (next, timeout) = self + .changed + .wait_timeout(entries, remaining) + .map_err(|_| TinyAgentsError::Graph("parent lifecycle lock poisoned".into()))?; + entries = next; + if timeout.timed_out() && entries.get(parent_key).is_some_and(|entry| entry.held) { + return Err(TinyAgentsError::Graph(format!( + "completion router: timed out waiting for local parent admission gate `{parent_key}`" + ))); + } + } + } + + pub(super) fn status(&self, parent_key: &str) -> Result<Option<ParentLifecycleStatus>> { + let entries = self + .entries + .lock() + .map_err(|_| TinyAgentsError::Graph("parent lifecycle lock poisoned".into()))?; + Ok(entries.get(parent_key).and_then(|entry| entry.status)) + } + + pub(super) fn set_status( + &self, + guard: &ParentLifecycleGuard, + status: ParentLifecycleStatus, + ) -> Result<()> { + if !std::ptr::eq(self, Arc::as_ptr(&guard.lifecycle)) { + return Err(TinyAgentsError::Graph( + "completion router: parent lifecycle guard scope mismatch".into(), + )); + } + let mut entries = self + .entries + .lock() + .map_err(|_| TinyAgentsError::Graph("parent lifecycle lock poisoned".into()))?; + let entry = entries.entry(guard.parent_key.clone()).or_default(); + if !entry.held { + return Err(TinyAgentsError::Graph( + "completion router: parent lifecycle guard is no longer held".into(), + )); + } + entry.status = Some(status); + Ok(()) + } +} + +pub(super) struct ParentLifecycleGuard { + lifecycle: Arc<ParentLifecycle>, + parent_key: String, +} + +impl ParentLifecycleGuard { + pub(super) fn parent_key(&self) -> &str { + &self.parent_key + } +} + +impl Drop for ParentLifecycleGuard { + fn drop(&mut self) { + let Ok(mut entries) = self.lifecycle.entries.lock() else { + return; + }; + if let Some(entry) = entries.get_mut(&self.parent_key) { + entry.held = false; + } + self.lifecycle.changed.notify_all(); + } +} diff --git a/crates/tinyagents-tasks/src/completions/mod.rs b/crates/tinyagents-tasks/src/completions/mod.rs index d79cda601..b817d7bf9 100644 --- a/crates/tinyagents-tasks/src/completions/mod.rs +++ b/crates/tinyagents-tasks/src/completions/mod.rs @@ -20,34 +20,50 @@ //! Live pushes are at-least-once: the host acknowledges with //! `mark_delivered`. Nothing here relaunches a child. +mod admission; mod delivery; mod driver; +mod durable_lease; mod format; +mod lifecycle; mod registry; mod router; mod store; mod types; +pub use admission::CompletionAdmission; pub use delivery::{ CompletionDelivery, CompletionDeliveryConfig, DeliveryFence, DeliveryOutcome, DeliveryTarget, - ParentBusyGate, ParentBusyLease, ThreadDeliveryLeases, + ParentBusyGate, ParentBusyLease, TenantThreadKey, ThreadDeliveryLeases, }; pub use driver::DriverCompletionStore; +pub use durable_lease::{DurableThreadDeliveryLease, DurableThreadDeliveryLeases}; pub use format::{CompletionFormatter, NeutralCompletionFormatter}; +#[doc(hidden)] +pub use lifecycle::ParentLifecycle; pub use registry::{ CompletionRouterPool, CompletionRouterPoolError, CompletionRouterReplaceOutcome, - open_driver_completion_store, + CompletionStoreFallback, open_driver_completion_store, open_driver_completion_store_or_memory, + open_driver_completion_store_or_memory_with_coordination_domain, + open_driver_completion_store_with_coordination_domain, }; pub use router::{ CompletionRouter, CompletionStoreReplaceOutcome, DEFAULT_MAX_ATTEMPTS, RecordOutcome, TombstoneOutcome, }; -pub use store::{CompletionStore, InMemoryCompletionStore}; +pub use store::{ + ClaimResolution, CompletionClaimMode, CompletionStore, InMemoryCompletionStore, + ParentGateClaim, PendingClaimRequest, ResolveClaimRequest, TaskClaimRequest, +}; pub use types::{ - CompletionArtifact, CompletionRecord, CompletionResult, CompletionState, CompletionStatus, - NotifyMode, + CompletionArtifact, CompletionClaim, CompletionReceipt, CompletionRecord, CompletionResult, + CompletionState, CompletionStatus, NotifyMode, }; +#[cfg(test)] +#[path = "durable_lease_tests.rs"] +mod durable_lease_tests; + #[cfg(test)] #[path = "router_tests.rs"] mod router_tests; diff --git a/crates/tinyagents-tasks/src/completions/registry.rs b/crates/tinyagents-tasks/src/completions/registry.rs index 4c1663f5a..2adcbbd4f 100644 --- a/crates/tinyagents-tasks/src/completions/registry.rs +++ b/crates/tinyagents-tasks/src/completions/registry.rs @@ -4,6 +4,10 @@ use std::collections::HashMap; use std::hash::Hash; use std::sync::{Arc, Mutex}; +use tinyagents_harness::error::Result as HarnessResult; +use tinyagents_storage_kit::{Scope, StorageBackend}; + +use super::lifecycle::ParentLifecycle; use super::{CompletionRouter, CompletionStore, DriverCompletionStore}; /// Outcome of replacing one cached router only if it is still current. @@ -21,12 +25,19 @@ pub enum CompletionRouterReplaceOutcome { pub enum CompletionRouterPoolError { /// The pool mutex was poisoned by a panic in another task. Lock(String), + /// A router could not release its active durable claims. Earlier routers + /// in the same call may already have released theirs; retry is safe. + Release(String), } impl std::fmt::Display for CompletionRouterPoolError { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { match self { Self::Lock(detail) => write!(f, "completion router pool lock poisoned: {detail}"), + Self::Release(detail) => write!( + f, + "completion router pool could not release claims: {detail}" + ), } } } @@ -35,9 +46,23 @@ impl std::error::Error for CompletionRouterPoolError {} type OpenFn<K> = dyn Fn(&K) -> Arc<CompletionRouter> + Send + Sync; +/// Whether a failed durable completion-store open may use process-local +/// storage. Shared-backend callers must use +/// [`CompletionStoreFallback::FailClosed`]. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum CompletionStoreFallback { + /// Return the backend or declaration error; never make cross-process work + /// appear durable when it is only held in process memory. + FailClosed, + /// Explicitly permit ephemeral local behavior if durable storage cannot be + /// opened. The caller owns deciding that the scope is local-only. + LocalMemory, +} + /// Process-local cache of one durable completion router per scope key. pub struct CompletionRouterPool<K: Eq + Hash + Clone + Send + Sync> { routers: Mutex<HashMap<K, Arc<CompletionRouter>>>, + lifecycles: Mutex<HashMap<(K, String), Arc<ParentLifecycle>>>, open: Arc<OpenFn<K>>, } @@ -46,6 +71,7 @@ impl<K: Eq + Hash + Clone + Send + Sync> CompletionRouterPool<K> { pub fn new(open: impl Fn(&K) -> Arc<CompletionRouter> + Send + Sync + 'static) -> Self { Self { routers: Mutex::new(HashMap::new()), + lifecycles: Mutex::new(HashMap::new()), open: Arc::new(open), } } @@ -57,6 +83,28 @@ impl<K: Eq + Hash + Clone + Send + Sync> CompletionRouterPool<K> { return Ok(Arc::clone(router)); } let router = (self.open)(key); + let store = router.store(); + let domain = if store.is_process_local() { + "process-local-memory".to_owned() + } else if let Some(identity) = store.coordination_identity() { + format!("coordination:{identity}") + } else { + format!("store:{:p}", Arc::as_ptr(&store)) + }; + let lifecycle = { + let mut lifecycles = self + .lifecycles + .lock() + .map_err(|error| CompletionRouterPoolError::Lock(error.to_string()))?; + Arc::clone( + lifecycles + .entry((key.clone(), domain)) + .or_insert_with(ParentLifecycle::new), + ) + }; + router + .attach_parent_lifecycle(lifecycle) + .map_err(|error| CompletionRouterPoolError::Lock(error.to_string()))?; routers.insert(key.clone(), Arc::clone(&router)); Ok(router) } @@ -80,12 +128,28 @@ impl<K: Eq + Hash + Clone + Send + Sync> CompletionRouterPool<K> { } /// Drops cached routers so their stores can be reopened after a lifecycle - /// boundary. Durable records remain in the backing store. + /// boundary. Durable records remain in the backing store, and the + /// pool-scoped parent lifecycle authority remains available to the new + /// router and to queue admissions retained from the old one. pub fn clear(&self) -> Result<(), CompletionRouterPoolError> { self.lock()?.clear(); Ok(()) } + /// Releases all active claims held by cached routers without clearing the + /// pool. Each router is snapshotted before store I/O, so the pool mutex is + /// not held while a backend is accessed. On failure, earlier routers may + /// already have released claims; calling this again is safe. + pub fn release_all(&self) -> Result<usize, CompletionRouterPoolError> { + let routers: Vec<_> = self.lock()?.values().map(Arc::clone).collect(); + routers.into_iter().try_fold(0usize, |released, router| { + router + .release_all() + .map(|count| released + count) + .map_err(|error| CompletionRouterPoolError::Release(error.to_string())) + }) + } + /// Atomically replaces `expected` only while it remains the cached router. /// /// A concurrent replacement wins and is returned as `Current`. If the @@ -139,18 +203,82 @@ impl<K: Eq + Hash + Clone + Send + Sync> std::fmt::Debug for CompletionRouterPoo } } -/// Opens the durable completion store after applying the tasks owner's schema -/// migrations to this already scope-bound storage handle. -pub async fn open_driver_completion_store( - storage: tinyagents_storage_kit::ScopedStorage, -) -> tinyagents_harness::error::Result<Arc<dyn CompletionStore>> { - crate::open_storage(storage.clone()) - .await - .map_err(|error| tinyagents_harness::error::TinyAgentsError::Storage(error.to_string()))?; - let store = DriverCompletionStore::new(Arc::clone(storage.documents())) - .map_err(|error| tinyagents_harness::error::TinyAgentsError::Storage(error.to_string()))?; - store - .declare_collections() +/// Opens a durable completion store over the tenant-scoped backend. +/// +/// This convenience constructor uses the identity of the supplied backend +/// `Arc`. If independent handles may point to the same database, use +/// [`open_driver_completion_store_with_coordination_domain`] with one +/// trusted domain shared by all processes opening that database and scope. +/// Shared-store setup fails closed. +pub fn open_driver_completion_store( + backend: Arc<dyn StorageBackend>, + scope: Scope, +) -> HarnessResult<Arc<dyn CompletionStore>> { + let store = DriverCompletionStore::new(backend, scope) .map_err(|error| tinyagents_harness::error::TinyAgentsError::Storage(error.to_string()))?; + store.declare_collections()?; + Ok(Arc::new(store)) +} + +/// Opens a durable completion store with a host-configured shared-backend +/// coordination domain. +/// +/// The opaque domain is trusted host configuration, never caller or model +/// input. Processes sharing a database and scope must use the same domain; +/// different databases must use different domains. +pub fn open_driver_completion_store_with_coordination_domain( + backend: Arc<dyn StorageBackend>, + scope: Scope, + coordination_domain: impl Into<String>, +) -> HarnessResult<Arc<dyn CompletionStore>> { + let store = + DriverCompletionStore::new_with_coordination_domain(backend, scope, coordination_domain) + .map_err(|error| { + tinyagents_harness::error::TinyAgentsError::Storage(error.to_string()) + })?; + store.declare_collections()?; Ok(Arc::new(store)) } + +/// Opens a driver-backed store, optionally falling back to explicit local +/// memory behavior when the caller has established that the scope is local. +/// +/// Shared/durable scopes must pass [`CompletionStoreFallback::FailClosed`]. +/// The local-memory option is an explicit policy choice, never an automatic +/// response to a shared-backend failure. +pub fn open_driver_completion_store_or_memory( + backend: Arc<dyn StorageBackend>, + scope: Scope, + fallback: CompletionStoreFallback, +) -> HarnessResult<Arc<dyn CompletionStore>> { + match open_driver_completion_store(backend, scope) { + Ok(store) => Ok(store), + Err(error) if fallback == CompletionStoreFallback::LocalMemory => { + tracing::warn!(error = %error, "durable completion store unavailable; using explicitly permitted local memory"); + Ok(Arc::new(super::InMemoryCompletionStore::new())) + } + Err(error) => Err(error), + } +} + +/// Opens a durable store in an explicit coordination domain, with optional +/// local-only memory fallback. +/// +/// Use `FailClosed` for shared/durable scopes. `LocalMemory` is an explicit +/// host policy for a scope that does not require cross-process delivery. +pub fn open_driver_completion_store_or_memory_with_coordination_domain( + backend: Arc<dyn StorageBackend>, + scope: Scope, + coordination_domain: impl Into<String>, + fallback: CompletionStoreFallback, +) -> HarnessResult<Arc<dyn CompletionStore>> { + match open_driver_completion_store_with_coordination_domain(backend, scope, coordination_domain) + { + Ok(store) => Ok(store), + Err(error) if fallback == CompletionStoreFallback::LocalMemory => { + tracing::warn!(error = %error, "durable completion store unavailable; using explicitly permitted local memory"); + Ok(Arc::new(super::InMemoryCompletionStore::new())) + } + Err(error) => Err(error), + } +} diff --git a/crates/tinyagents-tasks/src/completions/registry_tests.rs b/crates/tinyagents-tasks/src/completions/registry_tests.rs index 424351737..72d5f9a6b 100644 --- a/crates/tinyagents-tasks/src/completions/registry_tests.rs +++ b/crates/tinyagents-tasks/src/completions/registry_tests.rs @@ -2,6 +2,52 @@ use super::*; use std::sync::Arc; use crate::completions::{CompletionRouter, InMemoryCompletionStore}; +use crate::conditional_sqlite::MemoryStorage; +use tinyagents_storage_kit::{Scope, ScopedStorage, StorageBackend, StorageError}; + +struct UnconditionedBackend(MemoryStorage); + +impl StorageBackend for UnconditionedBackend { + fn for_scope(&self, scope: &Scope) -> Result<ScopedStorage, StorageError> { + self.0.for_scope(scope) + } +} + +#[tokio::test] +async fn durable_completion_open_fails_closed_while_explicit_local_fallback_remains_usable() { + let backend: Arc<dyn StorageBackend> = Arc::new(UnconditionedBackend(MemoryStorage::new())); + let scope = Scope::new("shared").unwrap(); + assert!( + open_driver_completion_store_or_memory( + Arc::clone(&backend), + scope.clone(), + CompletionStoreFallback::FailClosed, + ) + .is_err() + ); + + let local = open_driver_completion_store_or_memory( + backend, + scope, + CompletionStoreFallback::LocalMemory, + ) + .unwrap(); + let router = CompletionRouter::new(local); + let outcome = router + .record(crate::completions::CompletionRecord::new( + "local-task", + "local-parent", + "worker", + crate::completions::CompletionStatus::Success, + crate::completions::CompletionResult::default(), + )) + .await + .unwrap(); + assert!(matches!( + outcome, + crate::completions::RecordOutcome::Recorded { .. } + )); +} #[test] fn router_pool_opens_once_per_key_and_reuses_the_router() { @@ -86,6 +132,126 @@ fn router_pool_does_not_resurrect_a_router_after_clear() { assert!(!Arc::ptr_eq(&reopened, &stale)); } +#[test] +fn stopped_parent_remains_closed_after_pool_clear_and_reopen() { + let pool = CompletionRouterPool::new(|_| { + Arc::new(CompletionRouter::new(Arc::new( + InMemoryCompletionStore::new(), + ))) + }); + let first = pool.get_or_open(&"workspace").unwrap(); + assert_eq!(first.admit_child_generation("parent").unwrap(), Some(1)); + let mut queued = CompletionRecord::new( + "queued-child", + "parent", + "worker", + crate::completions::CompletionStatus::Success, + crate::completions::CompletionResult::default(), + ); + queued.parent_epoch = 1; + let old_admission = CompletionAdmission::for_router(&first, &queued); + let old_store_alias = CompletionAdmission::for_store(first.store(), &queued); + assert_eq!( + tinyagents_harness::run_queue::QueueAdmission::batch_key(&old_admission), + tinyagents_harness::run_queue::QueueAdmission::batch_key(&old_store_alias), + "for_store aliases must use the pool-bound lifecycle coordination key" + ); + let permit = tinyagents_harness::run_queue::QueueAdmission::acquire(&old_admission) + .expect("the current lifecycle admits its queued child"); + drop(permit); + first.cancel_parent("parent").unwrap(); + assert!(tinyagents_harness::run_queue::QueueAdmission::acquire(&old_admission).is_err()); + + pool.clear().unwrap(); + let reopened = pool.get_or_open(&"workspace").unwrap(); + let reopened_store_alias = CompletionAdmission::for_store(reopened.store(), &queued); + assert_eq!( + tinyagents_harness::run_queue::QueueAdmission::batch_key(&old_admission), + tinyagents_harness::run_queue::QueueAdmission::batch_key(&reopened_store_alias), + "queued handles across clear must preserve the same scoped lifecycle key" + ); + + assert!( + !reopened.admit_child("parent").unwrap(), + "clearing the router cache must not implicitly reopen a stopped parent" + ); + + reopened.resume_parent("parent"); + assert_eq!( + reopened.admit_child_generation("parent").unwrap(), + Some(3), + "only explicit resume opens a new generation" + ); + assert!(tinyagents_harness::run_queue::QueueAdmission::acquire(&old_admission).is_err()); + assert_eq!( + first.admit_child_generation("parent").unwrap(), + Some(3), + "an old router must reconcile its stale marker with the retained newer generation" + ); + assert!(tinyagents_harness::run_queue::QueueAdmission::acquire(&old_admission).is_err()); + let mut current = queued; + current.parent_epoch = 3; + let current_admission = CompletionAdmission::for_router(&reopened, ¤t); + assert!(tinyagents_harness::run_queue::QueueAdmission::acquire(¤t_admission).is_ok()); +} + +#[test] +fn deleted_parent_stays_deleted_across_pool_clear() { + let pool = CompletionRouterPool::new(|_| { + Arc::new(CompletionRouter::new(Arc::new( + InMemoryCompletionStore::new(), + ))) + }); + let first = pool.get_or_open(&"workspace").unwrap(); + assert_eq!(first.admit_child_generation("parent").unwrap(), Some(1)); + first.delete_parent("parent").unwrap(); + + pool.clear().unwrap(); + let reopened = pool.get_or_open(&"workspace").unwrap(); + reopened.resume_parent("parent"); + + assert_eq!(first.admit_child_generation("parent").unwrap(), None); + let mut old_record = CompletionRecord::new( + "queued-child", + "parent", + "worker", + crate::completions::CompletionStatus::Success, + crate::completions::CompletionResult::default(), + ); + old_record.parent_epoch = 1; + let old_admission = CompletionAdmission::for_router(&first, &old_record); + assert!(tinyagents_harness::run_queue::QueueAdmission::acquire(&old_admission).is_err()); + + assert!( + !reopened.admit_child("parent").unwrap(), + "a deleted parent cannot be reopened after its router is evicted" + ); +} + +#[test] +fn parent_lifecycle_isolated_between_pool_scope_keys() { + let pool = CompletionRouterPool::new(|_| { + Arc::new(CompletionRouter::new(Arc::new( + InMemoryCompletionStore::new(), + ))) + }); + let first_scope = pool.get_or_open(&"workspace-a").unwrap(); + let second_scope = pool.get_or_open(&"workspace-b").unwrap(); + assert_eq!( + first_scope.admit_child_generation("same-parent").unwrap(), + Some(1) + ); + assert_eq!( + second_scope.admit_child_generation("same-parent").unwrap(), + Some(1) + ); + + first_scope.cancel_parent("same-parent").unwrap(); + + assert!(!first_scope.admit_child("same-parent").unwrap()); + assert!(second_scope.admit_child("same-parent").unwrap()); +} + #[tokio::test] async fn evicting_one_key_preserves_other_cached_routers_and_held_claims() { let pool = CompletionRouterPool::new(|_| { @@ -108,20 +274,102 @@ async fn evicting_one_key_preserves_other_cached_routers_and_held_claims() { ) .await .unwrap(); - assert_eq!(removed.pull("parent", 1).unwrap().len(), 1); + let claimed = removed.pull("parent", 1).unwrap(); + assert_eq!(claimed.len(), 1); let evicted = pool.remove(&"workspace-a").unwrap().unwrap(); assert!(Arc::ptr_eq(&evicted, &removed)); assert!(pool.get(&"workspace-a").unwrap().is_none()); let still_cached = pool.get(&"workspace-b").unwrap().unwrap(); assert!(Arc::ptr_eq(&still_cached, &unrelated)); - assert_eq!(removed.mark_delivered(&["held-claim"]).unwrap(), 1); + assert_eq!( + removed + .mark_delivered(&claimed[0].receipt().into_iter().collect::<Vec<_>>()) + .unwrap(), + 1 + ); assert_eq!( removed.store().get("held-claim").unwrap().state, crate::completions::CompletionState::Delivered ); } +#[test] +fn one_memory_store_rejects_a_second_pool_scope_without_breaking_the_first() { + let store: Arc<dyn crate::completions::CompletionStore> = + Arc::new(InMemoryCompletionStore::new()); + let shared = Arc::clone(&store); + let pool = CompletionRouterPool::new(move |_| Arc::new(CompletionRouter::new(shared.clone()))); + let first = pool.get_or_open(&"workspace-a").unwrap(); + assert!(pool.get_or_open(&"workspace-b").is_err()); + + let epoch = first.admit_child_generation("parent").unwrap().unwrap(); + let mut record = CompletionRecord::new( + "child", + "parent", + "worker", + crate::completions::CompletionStatus::Success, + crate::completions::CompletionResult::default(), + ); + record.parent_epoch = epoch; + let router_alias = CompletionAdmission::for_router(&first, &record); + let store_alias = CompletionAdmission::for_store(first.store(), &record); + assert_eq!( + tinyagents_harness::run_queue::QueueAdmission::batch_key(&router_alias), + tinyagents_harness::run_queue::QueueAdmission::batch_key(&store_alias) + ); + drop(tinyagents_harness::run_queue::QueueAdmission::acquire(&router_alias).unwrap()); + assert!(tinyagents_harness::run_queue::QueueAdmission::acquire(&store_alias).is_ok()); +} + +#[tokio::test] +async fn pool_release_all_releases_claims_from_every_cached_router() { + let pool = CompletionRouterPool::new(|_| { + Arc::new(CompletionRouter::new(Arc::new( + InMemoryCompletionStore::new(), + ))) + }); + let first = pool.get_or_open(&"workspace-a").unwrap(); + let second = pool.get_or_open(&"workspace-b").unwrap(); + let first_store = first.store(); + let second_store = second.store(); + first + .record(CompletionRecord::new( + "first", + "parent-a", + "worker", + crate::completions::CompletionStatus::Success, + crate::completions::CompletionResult::text("one"), + )) + .await + .unwrap(); + second + .record(CompletionRecord::new( + "second", + "parent-b", + "worker", + crate::completions::CompletionStatus::Success, + crate::completions::CompletionResult::text("two"), + )) + .await + .unwrap(); + assert_eq!(first.claim_pending("parent-a", 1).unwrap().len(), 1); + assert_eq!(second.claim_pending("parent-b", 1).unwrap().len(), 1); + + assert_eq!(pool.release_all().unwrap(), 2); + + let restarted_first = CompletionRouter::new(first_store); + let restarted_second = CompletionRouter::new(second_store); + assert_eq!( + restarted_first.claim_pending("parent-a", 1).unwrap().len(), + 1 + ); + assert_eq!( + restarted_second.claim_pending("parent-b", 1).unwrap().len(), + 1 + ); +} + #[test] fn concurrent_fallback_replacements_have_one_winner_and_keep_other_keys_isolated() { let pool = Arc::new(CompletionRouterPool::new(|_| { diff --git a/crates/tinyagents-tasks/src/completions/router.rs b/crates/tinyagents-tasks/src/completions/router.rs index ee88085b7..4fc061542 100644 --- a/crates/tinyagents-tasks/src/completions/router.rs +++ b/crates/tinyagents-tasks/src/completions/router.rs @@ -8,16 +8,23 @@ use tinyagents_harness::error::{Result, TinyAgentsError}; use tinyagents_harness::run_queue::{QueueLane, RunQueueHandle}; use super::format::{CompletionFormatter, NeutralCompletionFormatter}; -use super::store::CompletionStore; -use super::types::{CompletionRecord, CompletionResult, CompletionState, NotifyMode}; +use super::lifecycle::{ParentLifecycle, ParentLifecycleGuard, ParentLifecycleStatus}; +use super::store::{ + ClaimResolution, CompletionClaimMode, CompletionStore, ParentGateClaim, PendingClaimRequest, + ResolveClaimRequest, TaskClaimRequest, +}; +use super::types::{CompletionReceipt, CompletionRecord, CompletionState, NotifyMode}; const LOG_PREFIX: &str = "[completion-router]"; +const CLAIM_VISIBILITY: Duration = Duration::from_secs(30); +const PARENT_GATE_TTL: Duration = Duration::from_secs(60); +const PARENT_GATE_WAIT: Duration = Duration::from_secs(65); /// Delivery attempts before a record gives up, unless overridden. pub const DEFAULT_MAX_ATTEMPTS: u32 = 5; /// What [`CompletionRouter::record`] did with a completion. -#[derive(Clone, Copy, Debug, PartialEq, Eq)] +#[derive(Clone, Debug, PartialEq, Eq)] pub enum RecordOutcome { /// Stored. `lane` is set when it was pushed onto a live parent's queue; the /// record is then pending and leased until the host calls @@ -25,6 +32,8 @@ pub enum RecordOutcome { Recorded { /// The lane the completion was pushed onto, if any. lane: Option<QueueLane>, + /// The receipt authorizing settlement of a live push. + receipt: Option<CompletionReceipt>, }, /// A record for this task id already exists; nothing changed. Duplicate, @@ -61,37 +70,296 @@ struct RouterState { /// Store operations and in-process leases share this lock so a migration /// cannot detach a claim from the state that owns its acknowledgement. store: Arc<dyn CompletionStore>, + /// Process-local lifecycle authority retained by the owning scope pool. + lifecycle: Arc<ParentLifecycle>, /// Task ids handed out by a claim and not yet resolved. Process-local on /// purpose: after a restart nobody holds a lease, so everything still /// pending is claimable again. leased: HashSet<String>, + /// Durable claim tokens owned by this router, used to reject stale settles. + active_receipts: HashMap<String, CompletionReceipt>, /// Live parents and their steering queues. parents: HashMap<String, RunQueueHandle>, } /// Durable, deduplicating hand-off of finished children to their parents. /// -/// See the `completions` module docs for the division of labour with the host. One -/// router serves one process: its read-modify-write sequences are serialised -/// by an in-process lock, so two processes must not share a store. +/// See the `completions` module docs for the division of labour with the host. +/// Router locks protect local bookkeeping; shared stores coordinate delivery +/// ownership through atomic claims, so independent routers may safely share a +/// backend that implements those claims. pub struct CompletionRouter { formatter: Arc<dyn CompletionFormatter>, max_attempts: u32, state: Mutex<RouterState>, + admission_store: Arc<Mutex<AdmissionStoreSlot>>, + lifecycle: Mutex<Arc<ParentLifecycle>>, +} + +pub(super) struct AdmissionStoreSlot { + pub(super) store: Arc<dyn CompletionStore>, + pub(super) migrating: bool, +} + +struct StoreMigration<'a> { + slot: &'a Mutex<AdmissionStoreSlot>, + replacement: Option<Arc<dyn CompletionStore>>, +} + +impl StoreMigration<'_> { + fn publish(mut self, replacement: Arc<dyn CompletionStore>) { + self.replacement = Some(replacement); + } +} + +impl Drop for StoreMigration<'_> { + fn drop(&mut self) { + let mut slot = self + .slot + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if let Some(replacement) = self.replacement.take() { + slot.store = replacement; + } + slot.migrating = false; + } +} + +struct ParentGateGuard { + store: Arc<dyn CompletionStore>, + _local: Option<ParentLifecycleGuard>, + parent_key: String, + claim: ParentGateClaim, +} + +/// Short-lived durable authority held while a queued completion crosses into +/// its parent turn. Dropping it releases the parent lifecycle gate. +pub struct ParentAdmissionPermit { + _guard: ParentGateGuard, +} + +impl tinyagents_harness::run_queue::QueueAdmissionPermit for ParentAdmissionPermit {} + +impl std::fmt::Debug for ParentAdmissionPermit { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("ParentAdmissionPermit") + .finish_non_exhaustive() + } +} + +impl ParentAdmissionPermit { + pub(crate) fn acquire( + store: Arc<dyn CompletionStore>, + lifecycle: Arc<ParentLifecycle>, + parent_key: &str, + parent_epoch: u64, + ) -> Result<Option<Self>> { + let gate = ParentGateGuard::acquire(Arc::clone(&store), lifecycle.clone(), parent_key)?; + let current = + CompletionRouter::parent_lifecycle_status(&store, &lifecycle, gate.local_guard())?; + let admitted = current == Some(ParentLifecycleStatus::Open(parent_epoch)); + Ok(admitted.then_some(Self { _guard: gate })) + } +} + +impl ParentGateGuard { + pub(crate) fn acquire( + store: Arc<dyn CompletionStore>, + lifecycle: Arc<ParentLifecycle>, + parent_key: &str, + ) -> Result<Self> { + let mut gate = Self::acquire_native(store, parent_key)?; + gate._local = Some(lifecycle.acquire(parent_key)?); + if !gate.renew()? { + return Err(TinyAgentsError::Graph(format!( + "completion router: parent gate expired while waiting for local lifecycle `{parent_key}`" + ))); + } + Ok(gate) + } + + fn acquire_native(store: Arc<dyn CompletionStore>, parent_key: &str) -> Result<Self> { + let holder = uuid::Uuid::new_v4().to_string(); + let started = std::time::Instant::now(); + loop { + if let Some(claim) = store.acquire_parent_gate(parent_key, &holder, PARENT_GATE_TTL)? { + return Ok(Self { + store, + _local: None, + parent_key: parent_key.to_owned(), + claim, + }); + } + if started.elapsed() >= PARENT_GATE_WAIT { + return Err(TinyAgentsError::Graph(format!( + "completion store: timed out waiting for parent gate {parent_key}" + ))); + } + std::thread::sleep(Duration::from_millis(5)); + } + } + + fn renew(&self) -> Result<bool> { + self.store + .renew_parent_gate(&self.parent_key, &self.claim, PARENT_GATE_TTL) + } + + fn claim(&self) -> &ParentGateClaim { + &self.claim + } + + pub(super) fn local_guard(&self) -> &ParentLifecycleGuard { + self._local + .as_ref() + .expect("router gates retain local lifecycle") + } +} + +impl Drop for ParentGateGuard { + fn drop(&mut self) { + if let Err(error) = self + .store + .release_parent_gate(&self.parent_key, &self.claim) + { + tracing::warn!(parent_key = %self.parent_key, error = %error, "{LOG_PREFIX} parent gate release failed"); + } + } } impl CompletionRouter { /// A router over `store` with the neutral formatter and /// [`DEFAULT_MAX_ATTEMPTS`]. pub fn new(store: Arc<dyn CompletionStore>) -> Self { + let lifecycle = if store.is_process_local() { + store + .parent_lifecycle() + .unwrap_or_else(|| ParentLifecycle::for_store(store.as_ref())) + } else { + ParentLifecycle::new() + }; + let admission_store = Arc::new(Mutex::new(AdmissionStoreSlot { + store: Arc::clone(&store), + migrating: false, + })); Self { formatter: Arc::new(NeutralCompletionFormatter), max_attempts: DEFAULT_MAX_ATTEMPTS, state: Mutex::new(RouterState { store, + lifecycle: Arc::clone(&lifecycle), leased: HashSet::new(), + active_receipts: HashMap::new(), parents: HashMap::new(), }), + admission_store, + lifecycle: Mutex::new(lifecycle), + } + } + + pub(super) fn attach_parent_lifecycle(&self, lifecycle: Arc<ParentLifecycle>) -> Result<()> { + let mut state = self.lock()?; + if state.store.is_process_local() { + state.store.bind_parent_lifecycle(Arc::clone(&lifecycle))?; + } + state.lifecycle = Arc::clone(&lifecycle); + *self + .lifecycle + .lock() + .map_err(|_| TinyAgentsError::Graph("router lifecycle lock poisoned".into()))? = + lifecycle; + Ok(()) + } + + pub(super) fn parent_lifecycle_handle(&self) -> Arc<ParentLifecycle> { + Arc::clone( + &self + .lifecycle + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner), + ) + } + + pub(super) fn parent_lifecycle_status( + store: &Arc<dyn CompletionStore>, + lifecycle: &Arc<ParentLifecycle>, + gate: &ParentLifecycleGuard, + ) -> Result<Option<ParentLifecycleStatus>> { + let parent_key = gate.parent_key(); + let marker = store.try_get(&Self::parent_marker_id(parent_key))?; + let marker_status = marker.map(|marker| match marker.state { + CompletionState::Delivered => ParentLifecycleStatus::Open(marker.parent_epoch), + CompletionState::ParentDeleted => ParentLifecycleStatus::Deleted(marker.parent_epoch), + _ => ParentLifecycleStatus::Stopped(marker.parent_epoch), + }); + + if store.is_process_local() { + let cached = lifecycle.status(parent_key)?; + let status = match (marker_status, cached) { + (Some(ParentLifecycleStatus::Deleted(marker_epoch)), snapshot) => { + Some(ParentLifecycleStatus::Deleted( + marker_epoch.max(snapshot.map_or(0, status_epoch)), + )) + } + (Some(marker), Some(ParentLifecycleStatus::Deleted(snapshot_epoch))) => Some( + ParentLifecycleStatus::Deleted(snapshot_epoch.max(status_epoch(marker))), + ), + (None, Some(ParentLifecycleStatus::Deleted(epoch))) => { + Some(ParentLifecycleStatus::Deleted(epoch)) + } + (Some(marker), Some(snapshot)) => { + Some(if status_epoch(marker) >= status_epoch(snapshot) { + marker + } else { + snapshot + }) + } + (Some(marker), None) => Some(marker), + (None, snapshot) => snapshot, + }; + if let Some(status) = status { + lifecycle.set_status(gate, status)?; + } + return Ok(status); + } + if let Some(status) = marker_status { + lifecycle.set_status(gate, status)?; + return Ok(Some(status)); + } + lifecycle.status(parent_key) + } + + pub(super) fn admission_store_handle(&self) -> Arc<Mutex<AdmissionStoreSlot>> { + Arc::clone(&self.admission_store) + } + + pub(super) fn admission_store_snapshot( + slot: &Mutex<AdmissionStoreSlot>, + ) -> std::result::Result< + Arc<dyn CompletionStore>, + tinyagents_harness::run_queue::QueueAdmissionError, + > { + let slot = slot + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if slot.migrating { + Err(tinyagents_harness::run_queue::QueueAdmissionError::Retry) + } else { + Ok(Arc::clone(&slot.store)) + } + } + + pub(super) fn admission_store_is_current( + slot: &Mutex<AdmissionStoreSlot>, + store: &Arc<dyn CompletionStore>, + ) -> std::result::Result<(), tinyagents_harness::run_queue::QueueAdmissionError> { + let slot = slot + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if slot.migrating || !Arc::ptr_eq(&slot.store, store) { + Err(tinyagents_harness::run_queue::QueueAdmissionError::Retry) + } else { + Ok(()) } } @@ -125,11 +393,82 @@ impl CompletionRouter { &state.store, ))); } + // Stop new queue admissions without holding a slot lock while waiting + // on durable parent gates. Existing permits finish first; admissions + // racing the migration fail closed and remain recoverable in the store. + { + let mut slot = self + .admission_store + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + slot.migrating = true; + } + let migration = StoreMigration { + slot: &self.admission_store, + replacement: None, + }; + if replacement.is_process_local() { + replacement.bind_parent_lifecycle(Arc::clone(&state.lifecycle))?; + } + // The first snapshot discovers which parent gates must be acquired. + // It is deliberately not copied: another router may close or update a + // parent while this migration is waiting for those gates. + let discovered = state.store.try_list(None)?; + let mut parents = discovered + .iter() + .map(|record| { + if record.parent_key.is_empty() { + Self::parent_from_marker_id(&record.task_id) + .unwrap_or_else(|| record.task_id.clone()) + } else { + record.parent_key.clone() + } + }) + .collect::<Vec<_>>(); + parents.sort(); + parents.dedup(); + let _source_gates = parents + .iter() + .map(|parent| ParentGateGuard::acquire_native(Arc::clone(&state.store), parent)) + .collect::<Result<Vec<_>>>()?; + let _local_gates = parents + .iter() + .map(|parent| state.lifecycle.acquire(parent)) + .collect::<Result<Vec<_>>>()?; + for gate in &_source_gates { + if !gate.renew()? { + return Err(TinyAgentsError::Graph( + "completion store: source parent gate expired while waiting for local lifecycle".into(), + )); + } + } let records = state.store.try_list(None)?; + let gated_parents = parents.iter().cloned().collect::<HashSet<_>>(); + if records.iter().any(|record| { + let parent = if record.parent_key.is_empty() { + Self::parent_from_marker_id(&record.task_id) + .unwrap_or_else(|| record.task_id.clone()) + } else { + record.parent_key.clone() + }; + !gated_parents.contains(&parent) + }) { + return Err(TinyAgentsError::Graph( + "completion store: migration snapshot gained a parent while acquiring gates; retry migration".into(), + )); + } for record in records { - replacement.put(&record)?; + let parent_key = if record.parent_key.is_empty() { + Self::parent_from_marker_id(&record.task_id) + .unwrap_or_else(|| record.task_id.clone()) + } else { + record.parent_key.clone() + }; + let gate = ParentGateGuard::acquire_native(Arc::clone(&replacement), &parent_key)?; + replacement.put_fenced(&record, &parent_key, gate.claim())?; } state.store = Arc::clone(&replacement); + migration.publish(Arc::clone(&replacement)); Ok(CompletionStoreReplaceOutcome::Replaced(replacement)) } @@ -185,22 +524,89 @@ impl CompletionRouter { } fn release_parent_locked(&self, state: &mut RouterState, parent_key: &str) { + let Ok(gate) = ParentGateGuard::acquire( + Arc::clone(&state.store), + Arc::clone(&state.lifecycle), + parent_key, + ) else { + return; + }; for record in state.store.list(Some(parent_key)) { state.leased.remove(&record.task_id); + if let Some(receipt) = state.active_receipts.remove(&record.task_id) { + let _ = state.store.resolve_claim(ResolveClaimRequest { + task_id: record.task_id, + parent_key: parent_key.to_owned(), + parent_epoch: receipt.parent_epoch, + gate: gate.claim().clone(), + token: receipt.token, + resolution: ClaimResolution::Release, + }); + } } } /// Releases the lease on `task_ids` without counting a failure, so they can /// be claimed again. Use it when a claim or push was abandoned (a dropped /// delivery future, a cleared queue). Attempts already counted stay counted. - pub fn release<S: AsRef<str>>(&self, task_ids: &[S]) { + pub fn release(&self, receipts: &[CompletionReceipt]) { if let Ok(mut state) = self.lock() { - for id in task_ids { - state.leased.remove(id.as_ref()); + for receipt in receipts { + if state.active_receipts.get(&receipt.task_id) == Some(receipt) { + state.leased.remove(&receipt.task_id); + state.active_receipts.remove(&receipt.task_id); + } + let Ok(gate) = ParentGateGuard::acquire( + Arc::clone(&state.store), + Arc::clone(&state.lifecycle), + &receipt.parent_key, + ) else { + continue; + }; + let _ = state.store.resolve_claim(ResolveClaimRequest { + task_id: receipt.task_id.clone(), + parent_key: receipt.parent_key.clone(), + parent_epoch: receipt.parent_epoch, + gate: gate.claim().clone(), + token: receipt.token.clone(), + resolution: ClaimResolution::Release, + }); } } } + /// Releases every receipt currently owned by this router so another + /// router can retry the pending records immediately. Store failures leave + /// the affected receipt locally owned and are returned to the caller. + pub fn release_all(&self) -> Result<usize> { + let mut state = self.lock()?; + let receipts: Vec<_> = state.active_receipts.values().cloned().collect(); + let mut released = 0; + for receipt in receipts { + let gate = ParentGateGuard::acquire( + Arc::clone(&state.store), + Arc::clone(&state.lifecycle), + &receipt.parent_key, + )?; + let resolved = state.store.resolve_claim(ResolveClaimRequest { + task_id: receipt.task_id.clone(), + parent_key: receipt.parent_key.clone(), + parent_epoch: receipt.parent_epoch, + gate: gate.claim().clone(), + token: receipt.token.clone(), + resolution: ClaimResolution::Release, + })?; + if state.active_receipts.get(&receipt.task_id) == Some(&receipt) { + state.leased.remove(&receipt.task_id); + state.active_receipts.remove(&receipt.task_id); + } + if resolved.is_some() { + released += 1; + } + } + Ok(released) + } + /// Records pushed onto `parent_key`'s live queue (or claimed) and not yet /// acknowledged, oldest first. pub fn in_flight_for(&self, parent_key: &str) -> Vec<CompletionRecord> { @@ -217,6 +623,22 @@ impl CompletionRouter { records } + fn classify_existing( + existing: &CompletionRecord, + incoming: &CompletionRecord, + ) -> RecordOutcome { + if existing.state == CompletionState::Tombstoned { + tracing::debug!(task_id = %incoming.task_id, "{LOG_PREFIX} dropped: tombstoned"); + RecordOutcome::Suppressed + } else if !existing.parent_key.is_empty() && existing.parent_key != incoming.parent_key { + tracing::warn!(task_id = %incoming.task_id, "{LOG_PREFIX} task id already recorded for another parent"); + RecordOutcome::IdCollision + } else { + tracing::debug!(task_id = %incoming.task_id, state = existing.state.as_str(), "{LOG_PREFIX} duplicate completion ignored"); + RecordOutcome::Duplicate + } + } + /// Records a finished child, idempotent per task id. /// /// A new record is stored `Pending`. If its [`NotifyMode`] is `Followup` or @@ -227,48 +649,53 @@ impl CompletionRouter { /// is lost first ([`Self::detach_parent`], [`Self::release`], or a restart) /// the record is claimable again. Delivery is at-least-once. pub async fn record(&self, mut record: CompletionRecord) -> Result<RecordOutcome> { - let push = { + let (push, gate) = { let mut state = self.lock()?; - if Self::parent_cancelled(&state, &record.parent_key) { + let gate = ParentGateGuard::acquire( + Arc::clone(&state.store), + Arc::clone(&state.lifecycle), + &record.parent_key, + )?; + let marker_id = Self::parent_marker_id(&record.parent_key); + let marker_before = state.store.try_get(&marker_id)?; + let Some(parent_epoch) = Self::admit_child_locked(&state, &record.parent_key, &gate)? + else { tracing::debug!( task_id = %record.task_id, parent_key = %record.parent_key, - "{LOG_PREFIX} dropped: parent cancelled" + "{LOG_PREFIX} dropped: parent is closed" ); return Ok(RecordOutcome::Suppressed); + }; + if let Some(existing) = state.store.get(&record.task_id) { + return Ok(Self::classify_existing(&existing, &record)); } - match state.store.get(&record.task_id) { - Some(existing) if existing.state == CompletionState::Tombstoned => { - tracing::debug!( - task_id = %record.task_id, - "{LOG_PREFIX} dropped: tombstoned" - ); - return Ok(RecordOutcome::Suppressed); - } - Some(existing) - if !existing.parent_key.is_empty() - && existing.parent_key != record.parent_key => - { - tracing::warn!( - task_id = %record.task_id, - "{LOG_PREFIX} task id already recorded for another parent" - ); - return Ok(RecordOutcome::IdCollision); - } - Some(existing) => { - tracing::debug!( - task_id = %record.task_id, - state = existing.state.as_str(), - "{LOG_PREFIX} duplicate completion ignored" - ); - return Ok(RecordOutcome::Duplicate); - } - None => {} + if (record.parent_epoch != 0 && record.parent_epoch != parent_epoch) + || (record.parent_epoch == 0 + && marker_before.is_some_and(|marker| marker.parent_epoch > 1)) + { + tracing::debug!( + task_id = %record.task_id, + parent_key = %record.parent_key, + supplied_epoch = record.parent_epoch, + current_epoch = parent_epoch, + "{LOG_PREFIX} suppressed completion from a stale or unversioned child" + ); + return Ok(RecordOutcome::Suppressed); } record.state = CompletionState::Pending; + record.parent_epoch = parent_epoch; record.attempts = 0; record.updated_at = SystemTime::now(); - state.store.put(&record)?; + if !state.store.insert_if_absent(&record, gate.claim())? { + let existing = state.store.try_get(&record.task_id)?.ok_or_else(|| { + TinyAgentsError::Graph(format!( + "completion store: conditional create conflicted but `{}` was not readable", + record.task_id + )) + })?; + return Ok(Self::classify_existing(&existing, &record)); + } tracing::debug!( task_id = %record.task_id, parent_key = %record.parent_key, @@ -285,34 +712,68 @@ impl CompletionRouter { // stays pending and leased until the host acknowledges it with // `mark_delivered`, so a crash, a cleared queue or a detached parent // never loses it (it is redelivered through `claim_pending`). - match lane.and_then(|lane| { + let push = match lane.and_then(|lane| { state .parents .get(&record.parent_key) .map(|queue| (lane, queue.clone())) }) { Some((lane, queue)) => { - record.attempts = 1; - record.updated_at = SystemTime::now(); - state.store.put(&record)?; + let token = uuid::Uuid::new_v4().to_string(); + let visible_after = SystemTime::now() + CLAIM_VISIBILITY; + let Some(claimed) = state.store.claim_task(TaskClaimRequest { + task_id: record.task_id.clone(), + parent_key: record.parent_key.clone(), + parent_epoch: record.parent_epoch, + gate: gate.claim().clone(), + mode: CompletionClaimMode::Followup, + token, + visible_after, + })? + else { + return Ok(RecordOutcome::Recorded { + lane: None, + receipt: None, + }); + }; + record = claimed; + let receipt = record.receipt().expect("claimed record has a receipt"); state.leased.insert(record.task_id.clone()); - Some((lane, queue)) + state + .active_receipts + .insert(record.task_id.clone(), receipt.clone()); + Some((lane, queue, receipt)) } None => None, - } + }; + (push, gate) }; - let Some((lane, queue)) = push else { - return Ok(RecordOutcome::Recorded { lane: None }); + let Some((lane, queue, push_receipt)) = push else { + return Ok(RecordOutcome::Recorded { + lane: None, + receipt: None, + }); }; + if !gate.renew()? { + self.release(std::slice::from_ref(&push_receipt)); + return Err(TinyAgentsError::Graph(format!( + "completion router: parent admission gate lost for `{}`", + record.parent_key + ))); + } let message = self.formatter.to_message(std::slice::from_ref(&record)); - queue.push(lane, message).await; + let admission = Arc::new(super::CompletionAdmission::for_router(self, &record)); + queue.push_guarded(lane, message, admission).await; tracing::debug!( task_id = %record.task_id, parent_key = %record.parent_key, lane = lane.as_str(), "{LOG_PREFIX} pushed to live parent" ); - Ok(RecordOutcome::Recorded { lane: Some(lane) }) + Ok(RecordOutcome::Recorded { + lane: Some(lane), + receipt: Some(push_receipt), + }) } /// [`Self::record`], retried up to `attempts` times with a short backoff @@ -343,25 +804,65 @@ impl CompletionRouter { } } + /// Records with retries and, only when the caller explicitly permits a + /// process-local scope, migrates this router to memory after durable + /// writes fail. Existing records must be readable and copyable before the + /// replacement is published; otherwise the durable store remains active + /// and the migration error is returned. + pub async fn record_resilient( + &self, + record: CompletionRecord, + attempts: u32, + fallback: super::registry::CompletionStoreFallback, + ) -> Result<RecordOutcome> { + match self.record_with_retries(record.clone(), attempts).await { + Ok(outcome) => Ok(outcome), + Err(error) if fallback == super::registry::CompletionStoreFallback::FailClosed => { + Err(error) + } + Err(_) => { + let current = self.store(); + let replacement: Arc<dyn CompletionStore> = + Arc::new(super::InMemoryCompletionStore::new()); + self.replace_store_if_current(¤t, replacement)?; + self.record_with_retries(record, attempts).await + } + } + } + /// Withdraws a child's completion because the parent collected it itself /// (waited on it, or read its result). Suppresses routing from now on; a /// message already pushed onto a live queue is not retracted. A completion that has not arrived - /// yet is dropped when it does. A withdrawn record keeps its id and state - /// but not its result. + /// yet is dropped when it does. pub fn tombstone(&self, task_id: &str) -> Result<TombstoneOutcome> { let mut state = self.lock()?; let outcome = match state.store.get(task_id) { Some(mut existing) if existing.state == CompletionState::Pending => { - existing.settle(CompletionState::Tombstoned); - state.store.put(&existing)?; + let gate_key = existing.parent_key.clone(); + let gate = ParentGateGuard::acquire( + Arc::clone(&state.store), + Arc::clone(&state.lifecycle), + &gate_key, + )?; + existing.state = CompletionState::Tombstoned; + existing.updated_at = SystemTime::now(); + state.store.put_fenced(&existing, &gate_key, gate.claim())?; state.leased.remove(task_id); TombstoneOutcome::Suppressed } Some(_) => TombstoneOutcome::AlreadyFinal, None => { - state - .store - .put(&CompletionRecord::tombstone_stub(task_id))?; + let gate_key = format!("completion-task:{task_id}"); + let gate = ParentGateGuard::acquire( + Arc::clone(&state.store), + Arc::clone(&state.lifecycle), + &gate_key, + )?; + state.store.put_fenced( + &CompletionRecord::tombstone_stub(task_id), + &gate_key, + gate.claim(), + )?; TombstoneOutcome::Reserved } }; @@ -372,19 +873,49 @@ impl CompletionRouter { /// Drops every pending completion for `parent_key` and everything that /// finishes for it later, until [`Self::resume_parent`]. For a deleted or /// stopped parent. A message already pushed onto the parent's queue is not - /// retracted (the host clears that queue when it stops the parent). The flag is stored durably, so it holds across a restart - /// (until [`Self::compact`] drops it after its retention window). Returns how many pending records were withdrawn. - /// Withdrawn records lose their result at once, so a deleted parent's - /// content does not wait for compaction. + /// retracted; queued generations must be rejected when the parent turn is + /// admitted. The flag is durable and survives completion compaction and + /// process restarts. Returns how many pending records were withdrawn. pub fn cancel_parent(&self, parent_key: &str) -> Result<usize> { + self.close_parent(parent_key, CompletionState::Tombstoned) + } + + /// Permanently closes a deleted parent thread. The durable marker is + /// retained through completion compaction so late children remain dropped + /// after a restart. + pub fn delete_parent(&self, parent_key: &str) -> Result<usize> { + self.close_parent(parent_key, CompletionState::ParentDeleted) + } + + fn close_parent(&self, parent_key: &str, marker_state: CompletionState) -> Result<usize> { let mut state = self.lock()?; - Self::set_parent_marker(&state, parent_key, CompletionState::Tombstoned)?; + let gate = ParentGateGuard::acquire( + Arc::clone(&state.store), + Arc::clone(&state.lifecycle), + parent_key, + )?; + let current = + Self::parent_lifecycle_status(&state.store, &state.lifecycle, gate.local_guard())?; + if marker_state != CompletionState::ParentDeleted + && matches!(current, Some(ParentLifecycleStatus::Deleted(_))) + { + state.parents.remove(parent_key); + return Ok(0); + } + if marker_state == CompletionState::ParentDeleted + && matches!(current, Some(ParentLifecycleStatus::Deleted(_))) + { + state.parents.remove(parent_key); + return Ok(0); + } + Self::set_parent_marker(&state, parent_key, marker_state, &gate)?; state.parents.remove(parent_key); let mut withdrawn = 0; - for mut record in state.store.list(Some(parent_key)) { + for mut record in state.store.try_list(Some(parent_key))? { if record.state == CompletionState::Pending { - record.settle(CompletionState::Tombstoned); - state.store.put(&record)?; + record.state = CompletionState::Tombstoned; + record.updated_at = SystemTime::now(); + state.store.put_fenced(&record, parent_key, gate.claim())?; state.leased.remove(&record.task_id); withdrawn += 1; } @@ -393,11 +924,26 @@ impl CompletionRouter { Ok(withdrawn) } - /// Lets completions for `parent_key` through again after - /// [`Self::cancel_parent`] (a stopped thread that the user reopened). + /// Reopens a parent stopped with [`Self::cancel_parent`]. Calling this on + /// an already open parent is idempotent. A parent closed with + /// [`Self::delete_parent`] cannot be reopened. pub fn resume_parent(&self, parent_key: &str) { let result = self.lock().and_then(|state| { - Self::set_parent_marker(&state, parent_key, CompletionState::Delivered) + let gate = ParentGateGuard::acquire( + Arc::clone(&state.store), + Arc::clone(&state.lifecycle), + parent_key, + )?; + match Self::parent_lifecycle_status(&state.store, &state.lifecycle, gate.local_guard())? + { + Some(ParentLifecycleStatus::Open(_)) => Ok(()), + Some(ParentLifecycleStatus::Deleted(_)) => Err(TinyAgentsError::Graph(format!( + "completion router: deleted parent {parent_key:?} cannot be resumed" + ))), + Some(ParentLifecycleStatus::Stopped(_)) | None => { + Self::set_parent_marker(&state, parent_key, CompletionState::Delivered, &gate) + } + } }); if let Err(error) = result { tracing::warn!(parent_key = %parent_key, error = %error, "{LOG_PREFIX} could not resume parent"); @@ -411,21 +957,138 @@ impl CompletionRouter { format!("\u{1}parent-cancelled:{parent_key}") } - fn parent_cancelled(state: &RouterState, parent_key: &str) -> bool { - state - .store - .get(&Self::parent_marker_id(parent_key)) - .is_some_and(|m| m.state == CompletionState::Tombstoned) + fn parent_from_marker_id(task_id: &str) -> Option<String> { + task_id + .strip_prefix("\u{1}parent-cancelled:") + .map(str::to_owned) + } + + /// Opens a parent on first sight, or reports whether its durable lifecycle + /// marker currently admits child completions. Closed parents are never + /// reopened implicitly. + pub fn admit_child(&self, parent_key: &str) -> Result<bool> { + Ok(self.admit_child_generation(parent_key)?.is_some()) + } + + /// Admits a child into the currently open durable lifecycle and returns + /// the generation it must carry when it later records its completion. + /// The generation changes only at an explicit parent stop/delete/resume, + /// not at ordinary router or lease handoff. + pub fn admit_child_generation(&self, parent_key: &str) -> Result<Option<u64>> { + let state = self.lock()?; + let gate = ParentGateGuard::acquire( + Arc::clone(&state.store), + Arc::clone(&state.lifecycle), + parent_key, + )?; + Self::admit_child_locked(&state, parent_key, &gate) + } + + /// Whether a generation carried by a queued completion still belongs to + /// the currently open parent lifecycle. Hosts use this at turn admission + /// to discard queue entries that raced with stop, delete, or reopen. + pub fn parent_generation_admitted(&self, parent_key: &str, parent_epoch: u64) -> Result<bool> { + let state = self.lock()?; + let gate = ParentGateGuard::acquire( + Arc::clone(&state.store), + Arc::clone(&state.lifecycle), + parent_key, + )?; + Ok( + Self::parent_lifecycle_status(&state.store, &state.lifecycle, gate.local_guard())? + == Some(ParentLifecycleStatus::Open(parent_epoch)), + ) + } + + fn admit_child_locked( + router: &RouterState, + parent_key: &str, + gate: &ParentGateGuard, + ) -> Result<Option<u64>> { + let marker_id = Self::parent_marker_id(parent_key); + match Self::parent_lifecycle_status(&router.store, &router.lifecycle, gate.local_guard())? { + Some(ParentLifecycleStatus::Open(epoch)) => { + if !router.store.is_process_local() || router.store.try_get(&marker_id)?.is_none() { + Self::write_parent_marker( + router, + parent_key, + ParentLifecycleStatus::Open(epoch), + gate, + )?; + } + Ok(Some(epoch)) + } + Some(ParentLifecycleStatus::Stopped(_)) | Some(ParentLifecycleStatus::Deleted(_)) => { + if !router.store.is_process_local() && router.store.try_get(&marker_id)?.is_none() { + let status = router.lifecycle.status(parent_key)?.ok_or_else(|| { + TinyAgentsError::Graph( + "completion router: lifecycle marker disappeared".into(), + ) + })?; + Self::write_parent_marker(router, parent_key, status, gate)?; + } + Ok(None) + } + None => { + // Records written before lifecycle markers existed carry + // generation zero. First sight admits those durable records + // into the new open generation under the same parent gate. + for mut legacy in router.store.try_list(Some(parent_key))? { + if legacy.parent_epoch == 0 { + legacy.parent_epoch = 1; + router.store.put_fenced(&legacy, parent_key, gate.claim())?; + } + } + Self::set_parent_marker(router, parent_key, CompletionState::Delivered, gate)?; + let marker = router.store.try_get(&marker_id)?.ok_or_else(|| { + TinyAgentsError::Graph( + "completion router: parent marker disappeared after admission".into(), + ) + })?; + Ok(Some(marker.parent_epoch)) + } + } } fn set_parent_marker( router: &RouterState, parent_key: &str, state: CompletionState, + gate: &ParentGateGuard, ) -> Result<()> { let mut marker = CompletionRecord::tombstone_stub(&Self::parent_marker_id(parent_key)); + marker.parent_key = parent_key.to_owned(); marker.state = state; - router.store.put(&marker) + let current = router.store.try_get(&marker.task_id)?; + let local = router.lifecycle.status(parent_key)?; + let epoch = if router.store.is_process_local() { + local + .map(status_epoch) + .or_else(|| current.map(|record| record.parent_epoch)) + } else { + current + .map(|record| record.parent_epoch) + .or_else(|| local.map(status_epoch)) + }; + marker.parent_epoch = epoch.map_or(1, |epoch| epoch.saturating_add(1)); + router.store.put_fenced(&marker, parent_key, gate.claim())?; + router.lifecycle.set_status( + gate.local_guard(), + status_from_marker(state, marker.parent_epoch), + ) + } + + fn write_parent_marker( + router: &RouterState, + parent_key: &str, + status: ParentLifecycleStatus, + gate: &ParentGateGuard, + ) -> Result<()> { + let mut marker = CompletionRecord::tombstone_stub(&Self::parent_marker_id(parent_key)); + marker.parent_key = parent_key.to_owned(); + marker.parent_epoch = status_epoch(status); + marker.state = marker_state(status); + router.store.put_fenced(&marker, parent_key, gate.claim()) } /// Every undelivered completion for `parent_key`, oldest first, in every @@ -445,84 +1108,145 @@ impl CompletionRouter { pending } + /// Parent keys with at least one pending completion in this store's scope. + /// Tenant isolation is provided by the scoped completion store. + pub fn parents_with_pending(&self) -> Result<Vec<String>> { + let state = self.lock()?; + Ok(state + .store + .try_list(None)? + .into_iter() + .filter(|record| record.state == CompletionState::Pending) + .map(|record| record.parent_key) + .collect::<std::collections::BTreeSet<_>>() + .into_iter() + .collect()) + } + + /// Closes every parent with pending completions and returns the number of + /// records withdrawn. Each marker uses the same durable fenced path as + /// [`Self::cancel_parent`]. + pub fn cancel_all_parents(&self) -> Result<usize> { + let (mut parents, records) = { + let state = self.lock()?; + let records = state.store.try_list(None)?; + let mut parents = state + .parents + .keys() + .cloned() + .collect::<std::collections::BTreeSet<_>>(); + for record in &records { + if let Some(parent) = Self::parent_from_marker_id(&record.task_id) { + if record.state == CompletionState::Delivered { + parents.insert(parent); + } + } else if record.state == CompletionState::Pending && !record.parent_key.is_empty() + { + parents.insert(record.parent_key.clone()); + } + } + (parents, records) + }; + // Pending rows may predate lifecycle markers, so retain their parent + // keys even when they have no open marker yet. + parents.extend( + records + .into_iter() + .filter(|record| record.state == CompletionState::Pending) + .map(|record| record.parent_key), + ); + parents.into_iter().try_fold(0usize, |total, parent| { + self.cancel_parent(&parent) + .map(|withdrawn| total + withdrawn) + }) + } + /// Claims up to `max` `Followup`/`Collect` completions for the idle-time /// batch, oldest first, counting an attempt on each. A claimed record is /// leased until [`Self::mark_delivered`] or [`Self::mark_failed`]. pub fn claim_pending(&self, parent_key: &str, max: usize) -> Result<Vec<CompletionRecord>> { - self.claim(parent_key, max, |mode| { - matches!(mode, NotifyMode::Followup | NotifyMode::Collect) - }) + self.claim(parent_key, max, CompletionClaimMode::Followup) } /// Claims the `HoldForNextTurn` completions for the turn that is starting. pub fn begin_turn(&self, parent_key: &str) -> Result<Vec<CompletionRecord>> { - self.claim(parent_key, usize::MAX, |mode| { - mode == NotifyMode::HoldForNextTurn - }) + self.claim(parent_key, usize::MAX, CompletionClaimMode::HoldForNextTurn) } /// Claims up to `max` `Off` completions the parent asked to pull. pub fn pull(&self, parent_key: &str, max: usize) -> Result<Vec<CompletionRecord>> { - self.claim(parent_key, max, |mode| mode == NotifyMode::Off) + self.claim(parent_key, max, CompletionClaimMode::Pull) } fn claim( &self, parent_key: &str, max: usize, - eligible: impl Fn(NotifyMode) -> bool, + mode: CompletionClaimMode, ) -> Result<Vec<CompletionRecord>> { let mut state = self.lock()?; - let mut candidates: Vec<_> = state - .store - .list(Some(parent_key)) - .into_iter() - .filter(|r| { - r.state == CompletionState::Pending - && eligible(r.notify_mode) - && !state.leased.contains(&r.task_id) - }) - .collect(); + let _gate = ParentGateGuard::acquire( + Arc::clone(&state.store), + Arc::clone(&state.lifecycle), + parent_key, + )?; + let Some(parent_epoch) = Self::admit_child_locked(&state, parent_key, &_gate)? else { + return Ok(Vec::new()); + }; + let token = uuid::Uuid::new_v4().to_string(); + let visible_after = SystemTime::now() + CLAIM_VISIBILITY; + let mut candidates = state.store.claim_pending(PendingClaimRequest { + parent_key: parent_key.to_owned(), + parent_epoch, + gate: _gate.claim().clone(), + mode, + max, + token, + visible_after, + })?; sort_oldest_first(&mut candidates); - candidates.truncate(max); - let mut claimed: Vec<CompletionRecord> = Vec::with_capacity(candidates.len()); - for mut record in candidates { - record.attempts = record.attempts.saturating_add(1); - record.updated_at = SystemTime::now(); - if let Err(error) = state.store.put(&record) { - // The caller gets no ids on error, so none may stay leased. - for done in &claimed { - state.leased.remove(&done.task_id); - } - return Err(error); - } + for record in &candidates { state.leased.insert(record.task_id.clone()); - claimed.push(record); + if let Some(receipt) = record.receipt() { + state + .active_receipts + .insert(record.task_id.clone(), receipt); + } } - if !claimed.is_empty() { + if !candidates.is_empty() { tracing::debug!( parent_key = %parent_key, - claimed = claimed.len(), + claimed = candidates.len(), "{LOG_PREFIX} claimed" ); } - Ok(claimed) + Ok(candidates) } /// Settles claimed (or pulled) records as delivered. Returns how many - /// changed; records that are not pending are left alone. A delivered - /// record keeps its id, state and attempts for dedupe but drops its result. - pub fn mark_delivered<S: AsRef<str>>(&self, task_ids: &[S]) -> Result<usize> { + /// changed; records that are not pending are left alone. + pub fn mark_delivered(&self, receipts: &[CompletionReceipt]) -> Result<usize> { let mut state = self.lock()?; let mut changed = 0; - for id in task_ids { - let id = id.as_ref(); - state.leased.remove(id); - if let Some(mut record) = state.store.get(id) - && record.state == CompletionState::Pending - { - record.settle(CompletionState::Delivered); - state.store.put(&record)?; + for receipt in receipts { + let gate = ParentGateGuard::acquire( + Arc::clone(&state.store), + Arc::clone(&state.lifecycle), + &receipt.parent_key, + )?; + let resolved = state.store.resolve_claim(ResolveClaimRequest { + task_id: receipt.task_id.clone(), + parent_key: receipt.parent_key.clone(), + parent_epoch: receipt.parent_epoch, + gate: gate.claim().clone(), + token: receipt.token.clone(), + resolution: ClaimResolution::Delivered, + })?; + if state.active_receipts.get(&receipt.task_id) == Some(receipt) { + state.active_receipts.remove(&receipt.task_id); + state.leased.remove(&receipt.task_id); + } + if resolved.is_some() { changed += 1; } } @@ -534,57 +1258,50 @@ impl CompletionRouter { /// claimed again, unless its attempts have reached [`Self::max_attempts`]: /// then it becomes [`CompletionState::GaveUp`] and is returned so the host /// can apply its own give-up policy (the router never delivers it again). - /// The returned record carries the full result; the stored one does not, - /// so the host's give-up writer is the result's last reader. - /// - /// A record whose give-up cannot be stored stays `Pending` (its lease is - /// released) and is retried; the ones already given up are still returned, - /// so a partly written batch never strands a give-up the host has not seen. - pub fn mark_failed<S: AsRef<str>>(&self, task_ids: &[S]) -> Result<Vec<CompletionRecord>> { + /// A failure resolving one receipt does not hide give-ups already resolved + /// in this batch; the failed receipt can be retried after its store claim + /// expires. + pub fn mark_failed(&self, receipts: &[CompletionReceipt]) -> Result<Vec<CompletionRecord>> { let mut state = self.lock()?; let mut gave_up = Vec::new(); - for id in task_ids { - let id = id.as_ref(); - state.leased.remove(id); - let Some(mut record) = state.store.get(id) else { - continue; - }; - if record.state != CompletionState::Pending || record.attempts < self.max_attempts { - continue; + for receipt in receipts { + let resolved = ParentGateGuard::acquire( + Arc::clone(&state.store), + Arc::clone(&state.lifecycle), + &receipt.parent_key, + ) + .and_then(|gate| { + state.store.resolve_claim(ResolveClaimRequest { + task_id: receipt.task_id.clone(), + parent_key: receipt.parent_key.clone(), + parent_epoch: receipt.parent_epoch, + gate: gate.claim().clone(), + token: receipt.token.clone(), + resolution: ClaimResolution::Failed { + max_attempts: self.max_attempts, + }, + }) + }); + if state.active_receipts.get(&receipt.task_id) == Some(receipt) { + state.active_receipts.remove(&receipt.task_id); + state.leased.remove(&receipt.task_id); } - // The stored copy keeps its result until the host has the record: - // the cleanup below drops it only for records returned from here. - let result = std::mem::take(&mut record.result); - record.settle(CompletionState::GaveUp); - record.result = result; - if let Err(error) = state.store.put(&record) { - tracing::warn!(task_id = %record.task_id, %error, "{LOG_PREFIX} could not store a give-up; it stays pending"); - continue; - } - tracing::warn!( - task_id = %record.task_id, - parent_key = %record.parent_key, - attempts = record.attempts, - "{LOG_PREFIX} gave up" - ); - gave_up.push(record); - } - for record in &gave_up { - let settled = CompletionRecord { - result: CompletionResult::default(), - ..record.clone() - }; - if let Err(error) = state.store.put(&settled) { - tracing::warn!(task_id = %record.task_id, %error, "{LOG_PREFIX} could not drop a given-up result"); + match resolved { + Ok(Some(record)) if record.state == CompletionState::GaveUp => { + tracing::warn!(task_id = %record.task_id, parent_key = %record.parent_key, attempts = record.attempts, "{LOG_PREFIX} gave up"); + gave_up.push(record); + } + Ok(_) => {} + Err(error) => { + tracing::warn!(task_id = %receipt.task_id, %error, "{LOG_PREFIX} could not resolve failed delivery"); + } } } Ok(gave_up) } - /// Drops settled records older than `retain` from the store. Their results - /// went when they settled; this frees the dedupe entries too, so pick a - /// window longer than any child can outlive its parent's interest. See - /// [`CompletionStore::compact`]. + /// Drops settled records older than `retain` from the store (and, for the + /// JSONL store, the log). See [`CompletionStore::compact`]. pub fn compact(&self, retain: Duration) -> Result<usize> { let state = self.lock()?; let dropped = state.store.compact(retain)?; @@ -600,3 +1317,27 @@ fn sort_oldest_first(records: &mut [CompletionRecord]) { .then_with(|| a.task_id.cmp(&b.task_id)) }); } + +fn status_epoch(status: ParentLifecycleStatus) -> u64 { + match status { + ParentLifecycleStatus::Open(epoch) + | ParentLifecycleStatus::Stopped(epoch) + | ParentLifecycleStatus::Deleted(epoch) => epoch, + } +} + +fn marker_state(status: ParentLifecycleStatus) -> CompletionState { + match status { + ParentLifecycleStatus::Open(_) => CompletionState::Delivered, + ParentLifecycleStatus::Stopped(_) => CompletionState::Tombstoned, + ParentLifecycleStatus::Deleted(_) => CompletionState::ParentDeleted, + } +} + +fn status_from_marker(state: CompletionState, epoch: u64) -> ParentLifecycleStatus { + match state { + CompletionState::Delivered => ParentLifecycleStatus::Open(epoch), + CompletionState::ParentDeleted => ParentLifecycleStatus::Deleted(epoch), + _ => ParentLifecycleStatus::Stopped(epoch), + } +} diff --git a/crates/tinyagents-tasks/src/completions/router_tests.rs b/crates/tinyagents-tasks/src/completions/router_tests.rs index 45dc55e17..94c897fbe 100644 --- a/crates/tinyagents-tasks/src/completions/router_tests.rs +++ b/crates/tinyagents-tasks/src/completions/router_tests.rs @@ -1,6 +1,9 @@ use std::sync::Arc; +use std::time::{Duration, SystemTime}; -use tinyagents_harness::run_queue::{QueueLane, RunQueue, RunQueueHandle}; +use tinyagents_harness::run_queue::{ + QueueAdmission, QueueLane, QueueMode, RunQueue, RunQueueHandle, +}; use tinyinference_llm::message::Message; use super::*; @@ -27,6 +30,13 @@ fn ids(records: &[CompletionRecord]) -> Vec<&str> { records.iter().map(|r| r.task_id.as_str()).collect() } +fn receipts(records: &[CompletionRecord]) -> Vec<CompletionReceipt> { + records + .iter() + .filter_map(CompletionRecord::receipt) + .collect() +} + #[tokio::test] async fn recording_twice_stores_once() { let router = router(); @@ -47,7 +57,7 @@ async fn a_delivered_task_is_not_recorded_again() { let router = router(); router.record(record("t1", "p")).await.unwrap(); let claimed = router.claim_pending("p", 10).unwrap(); - router.mark_delivered(&["t1"]).unwrap(); + router.mark_delivered(&receipts(&claimed)).unwrap(); assert_eq!(claimed.len(), 1); assert_eq!( router.record(record("t1", "p")).await.unwrap(), @@ -79,6 +89,149 @@ async fn tombstone_before_the_completion_drops_it() { assert!(router.claim_pending("p", 10).unwrap().is_empty()); } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn cancellation_fences_a_completion_that_observed_the_old_parent_state() { + let inner = InMemoryCompletionStore::new(); + let observed = Arc::new(std::sync::Barrier::new(2)); + let resume = Arc::new(std::sync::Barrier::new(2)); + let cancel_waiting = Arc::new(std::sync::Barrier::new(2)); + let store = Arc::new(BlockingParentMarkerRead { + inner: inner.clone(), + marker_id: "\u{1}parent-cancelled:p".to_owned(), + observed: Arc::clone(&observed), + resume: Arc::clone(&resume), + gate_attempts: std::sync::atomic::AtomicUsize::new(0), + marker_reads: std::sync::atomic::AtomicUsize::new(0), + cancel_waiting: Arc::clone(&cancel_waiting), + }); + let router = Arc::new(CompletionRouter::new(store.clone())); + let cancel_router = CompletionRouter::new(store); + let q = queue(); + router.attach_parent("p", q.clone()); + + let recording = { + let router = Arc::clone(&router); + tokio::spawn(async move { router.record(record("late", "p")).await }) + }; + observed.wait(); + let (started_tx, started_rx) = std::sync::mpsc::channel(); + let (done_tx, done_rx) = std::sync::mpsc::channel(); + let cancellation = std::thread::spawn(move || { + started_tx.send(()).unwrap(); + let result = cancel_router.cancel_parent("p"); + done_tx.send(result).unwrap(); + }); + started_rx.recv().unwrap(); + cancel_waiting.wait(); + resume.wait(); + let outcome = recording.await.unwrap().unwrap(); + assert!(matches!(outcome, RecordOutcome::Recorded { .. })); + assert_eq!(done_rx.recv().unwrap().unwrap(), 1); + cancellation.join().unwrap(); + assert!(router.pending_for("p").is_empty()); + let delivered_before_rejected_admission = q.status().await.total; + assert_eq!( + router.record(record("after-cancel", "p")).await.unwrap(), + RecordOutcome::Suppressed + ); + assert_eq!(q.status().await.total, delivered_before_rejected_admission); +} + +struct BlockingParentMarkerRead { + inner: InMemoryCompletionStore, + marker_id: String, + observed: Arc<std::sync::Barrier>, + resume: Arc<std::sync::Barrier>, + gate_attempts: std::sync::atomic::AtomicUsize, + marker_reads: std::sync::atomic::AtomicUsize, + cancel_waiting: Arc<std::sync::Barrier>, +} + +impl CompletionStore for BlockingParentMarkerRead { + fn acquire_parent_gate( + &self, + parent_key: &str, + holder: &str, + ttl: std::time::Duration, + ) -> tinyagents_harness::error::Result<Option<ParentGateClaim>> { + use std::sync::atomic::Ordering::SeqCst; + if self.gate_attempts.fetch_add(1, SeqCst) == 1 { + self.cancel_waiting.wait(); + } + self.inner.acquire_parent_gate(parent_key, holder, ttl) + } + + fn renew_parent_gate( + &self, + parent_key: &str, + claim: &ParentGateClaim, + ttl: std::time::Duration, + ) -> tinyagents_harness::error::Result<bool> { + self.inner.renew_parent_gate(parent_key, claim, ttl) + } + + fn release_parent_gate( + &self, + parent_key: &str, + claim: &ParentGateClaim, + ) -> tinyagents_harness::error::Result<bool> { + self.inner.release_parent_gate(parent_key, claim) + } + + fn get(&self, task_id: &str) -> Option<CompletionRecord> { + let record = self.inner.get(task_id); + if task_id == self.marker_id + && self + .marker_reads + .fetch_add(1, std::sync::atomic::Ordering::SeqCst) + == 0 + { + self.observed.wait(); + self.resume.wait(); + } + record + } + + fn try_get( + &self, + task_id: &str, + ) -> tinyagents_harness::error::Result<Option<CompletionRecord>> { + Ok(self.get(task_id)) + } + + fn put(&self, record: &CompletionRecord) -> tinyagents_harness::error::Result<()> { + self.inner.put(record) + } + + fn put_fenced( + &self, + record: &CompletionRecord, + parent_key: &str, + gate: &ParentGateClaim, + ) -> tinyagents_harness::error::Result<()> { + self.inner.put_fenced(record, parent_key, gate) + } + + fn insert_if_absent( + &self, + record: &CompletionRecord, + gate: &ParentGateClaim, + ) -> tinyagents_harness::error::Result<bool> { + self.inner.insert_if_absent(record, gate) + } + + fn claim_task( + &self, + request: TaskClaimRequest, + ) -> tinyagents_harness::error::Result<Option<CompletionRecord>> { + self.inner.claim_task(request) + } + + fn list(&self, parent_key: Option<&str>) -> Vec<CompletionRecord> { + self.inner.list(parent_key) + } +} + #[tokio::test] async fn tombstone_never_reaches_a_live_parent_queue() { let router = router(); @@ -93,8 +246,8 @@ async fn tombstone_never_reaches_a_live_parent_queue() { async fn tombstone_after_delivery_changes_nothing() { let router = router(); router.record(record("t1", "p")).await.unwrap(); - router.claim_pending("p", 1).unwrap(); - router.mark_delivered(&["t1"]).unwrap(); + let claimed = router.claim_pending("p", 1).unwrap(); + router.mark_delivered(&receipts(&claimed)).unwrap(); assert_eq!( router.tombstone("t1").unwrap(), TombstoneOutcome::AlreadyFinal @@ -113,6 +266,337 @@ async fn claim_counts_attempts_and_does_not_hand_out_a_leased_record_twice() { ); } +#[tokio::test] +async fn another_router_cannot_ack_or_fail_an_active_claim_by_task_id() { + let store = Arc::new(InMemoryCompletionStore::new()); + let owner = CompletionRouter::new(store.clone()).with_max_attempts(1); + let stranger = CompletionRouter::new(store); + owner.record(record("owned", "p")).await.unwrap(); + let claimed = owner.claim_pending("p", 1).unwrap(); + let forged = CompletionReceipt { + task_id: "owned".into(), + parent_key: "p".into(), + parent_epoch: 1, + token: "not-the-owner".into(), + }; + assert_eq!( + stranger + .mark_delivered(std::slice::from_ref(&forged)) + .unwrap(), + 0 + ); + assert_eq!(stranger.mark_failed(&[forged]).unwrap().len(), 0); + assert_eq!(owner.pending_for("p").len(), 1); + assert_eq!(owner.mark_delivered(&receipts(&claimed)).unwrap(), 1); + assert_eq!(claimed.len(), 1); +} + +#[tokio::test] +async fn an_old_claim_receipt_cannot_settle_a_reclaimed_attempt() { + let store = Arc::new(InMemoryCompletionStore::new()); + let router = CompletionRouter::new(store.clone()); + router.record(record("reclaimed", "p")).await.unwrap(); + let old = router.claim_pending("p", 1).unwrap().remove(0); + + // Simulate another process reclaiming the durable row after visibility. + let mut expired = store.get("reclaimed").unwrap(); + expired.claim.as_mut().unwrap().visible_after = + SystemTime::now() - std::time::Duration::from_secs(1); + store.put(&expired).unwrap(); + let new = router.claim_pending("p", 1).unwrap(); + assert_eq!(new.len(), 1); + let old_receipt = old.receipt().unwrap(); + assert_eq!( + router + .mark_delivered(std::slice::from_ref(&old_receipt)) + .unwrap(), + 0 + ); + assert_eq!( + router + .mark_failed(std::slice::from_ref(&old_receipt)) + .unwrap() + .len(), + 0 + ); + router.release(std::slice::from_ref(&old_receipt)); + assert_eq!(router.pending_for("p").len(), 1); + assert_eq!(router.mark_delivered(&receipts(&new)).unwrap(), 1); +} + +#[tokio::test] +async fn live_push_uses_the_same_durable_claim_as_batch_delivery() { + let store = Arc::new(InMemoryCompletionStore::new()); + let live = CompletionRouter::new(store.clone()); + let other = CompletionRouter::new(store); + live.attach_parent("p", queue()); + live.record(record("live", "p")).await.unwrap(); + + assert!(other.claim_pending("p", 1).unwrap().is_empty()); +} + +#[test] +fn in_memory_claims_choose_oldest_before_task_id_tiebreaking() { + let store = InMemoryCompletionStore::new(); + let now = SystemTime::now(); + store + .put(&record("z-older", "p").with_finished_at(now)) + .unwrap(); + store + .put(&record("a-newer", "p").with_finished_at(now + std::time::Duration::from_secs(1))) + .unwrap(); + + let gate = store + .acquire_parent_gate("p", "test", std::time::Duration::from_secs(30)) + .unwrap() + .unwrap(); + let claimed = store + .claim_pending(PendingClaimRequest { + parent_key: "p".into(), + parent_epoch: 0, + gate, + mode: CompletionClaimMode::Followup, + max: 1, + token: "owner".into(), + visible_after: now + std::time::Duration::from_secs(30), + }) + .unwrap(); + assert_eq!(ids(&claimed), ["z-older"]); +} + +struct BlockingFirstReadStore { + inner: InMemoryCompletionStore, + blocked_task: &'static str, + reads: std::sync::atomic::AtomicUsize, + entered: std::sync::Barrier, + resume: std::sync::Barrier, +} + +impl CompletionStore for BlockingFirstReadStore { + fn acquire_parent_gate( + &self, + parent_key: &str, + holder: &str, + ttl: std::time::Duration, + ) -> tinyagents_harness::error::Result<Option<ParentGateClaim>> { + self.inner.acquire_parent_gate(parent_key, holder, ttl) + } + fn renew_parent_gate( + &self, + parent_key: &str, + claim: &ParentGateClaim, + ttl: std::time::Duration, + ) -> tinyagents_harness::error::Result<bool> { + self.inner.renew_parent_gate(parent_key, claim, ttl) + } + fn release_parent_gate( + &self, + parent_key: &str, + claim: &ParentGateClaim, + ) -> tinyagents_harness::error::Result<bool> { + self.inner.release_parent_gate(parent_key, claim) + } + + fn get(&self, task_id: &str) -> Option<CompletionRecord> { + let observed = self.inner.get(task_id); + if task_id == self.blocked_task + && self.reads.fetch_add(1, std::sync::atomic::Ordering::SeqCst) == 0 + { + self.entered.wait(); + self.resume.wait(); + } + observed + } + + fn put(&self, record: &CompletionRecord) -> tinyagents_harness::error::Result<()> { + self.inner.put(record) + } + + fn put_fenced( + &self, + record: &CompletionRecord, + parent_key: &str, + gate: &ParentGateClaim, + ) -> tinyagents_harness::error::Result<()> { + self.inner.put_fenced(record, parent_key, gate) + } + + fn insert_if_absent( + &self, + record: &CompletionRecord, + gate: &ParentGateClaim, + ) -> tinyagents_harness::error::Result<bool> { + self.inner.insert_if_absent(record, gate) + } + + fn claim_pending( + &self, + request: PendingClaimRequest, + ) -> tinyagents_harness::error::Result<Vec<CompletionRecord>> { + self.inner.claim_pending(request) + } + + fn claim_task( + &self, + request: TaskClaimRequest, + ) -> tinyagents_harness::error::Result<Option<CompletionRecord>> { + self.inner.claim_task(request) + } + + fn resolve_claim( + &self, + request: ResolveClaimRequest, + ) -> tinyagents_harness::error::Result<Option<CompletionRecord>> { + self.inner.resolve_claim(request) + } + + fn list(&self, parent_key: Option<&str>) -> Vec<CompletionRecord> { + self.inner.list(parent_key) + } +} + +fn blocking_first_read_store(blocked_task: &'static str) -> Arc<BlockingFirstReadStore> { + Arc::new(BlockingFirstReadStore { + inner: InMemoryCompletionStore::new(), + blocked_task, + reads: std::sync::atomic::AtomicUsize::new(0), + entered: std::sync::Barrier::new(2), + resume: std::sync::Barrier::new(2), + }) +} + +#[tokio::test] +async fn concurrent_first_record_cannot_replace_another_parent_live_claim() { + let store = blocking_first_read_store("racing"); + let first = Arc::new(CompletionRouter::new(store.clone())); + let second = CompletionRouter::new(store.clone()); + let first_queue = queue(); + let second_queue = queue(); + first.attach_parent("parent-a", first_queue.clone()); + second.attach_parent("parent-b", second_queue.clone()); + + let first_router = Arc::clone(&first); + let first_writer = std::thread::spawn(move || { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + runtime.block_on(first_router.record(record("racing", "parent-a"))) + }); + store.entered.wait(); + + let second_outcome = second.record(record("racing", "parent-b")).await.unwrap(); + let second_receipt = match second_outcome { + RecordOutcome::Recorded { + receipt: Some(receipt), + .. + } => receipt, + other => panic!("second writer should win the create: {other:?}"), + }; + store.resume.wait(); + let first_outcome = first_writer.join().unwrap().unwrap(); + + assert_eq!(first_outcome, RecordOutcome::IdCollision); + assert_eq!(store.inner.get("racing").unwrap().parent_key, "parent-b"); + let winner = store.inner.get("racing").unwrap(); + assert_eq!(winner.receipt(), Some(second_receipt.clone())); + assert_eq!( + winner.attempts, 1, + "the losing create must not reset attempt state" + ); + assert_eq!(first_queue.status().await.total, 0); + assert_eq!(second_queue.status().await.total, 1); +} + +#[tokio::test] +async fn concurrent_duplicate_record_cannot_erase_the_first_claim_or_push_twice() { + let store = Arc::new(InMemoryCompletionStore::new()); + let first = Arc::new(CompletionRouter::new(store.clone())); + let second = Arc::new(CompletionRouter::new(store.clone())); + let first_queue = queue(); + let second_queue = queue(); + first.attach_parent("parent", first_queue.clone()); + second.attach_parent("parent", second_queue.clone()); + + let start = Arc::new(tokio::sync::Barrier::new(3)); + let first_writer = { + let router = Arc::clone(&first); + let start = Arc::clone(&start); + tokio::spawn(async move { + start.wait().await; + router + .record(record("same-parent", "parent")) + .await + .unwrap() + }) + }; + let second_writer = { + let router = Arc::clone(&second); + let start = Arc::clone(&start); + tokio::spawn(async move { + start.wait().await; + router + .record(record("same-parent", "parent")) + .await + .unwrap() + }) + }; + start.wait().await; + let outcomes = [first_writer.await.unwrap(), second_writer.await.unwrap()]; + let recorded = outcomes + .iter() + .find_map(|outcome| match outcome { + RecordOutcome::Recorded { + receipt: Some(receipt), + .. + } => Some(receipt), + _ => None, + }) + .expect("exactly one writer should own the live claim"); + assert_eq!( + outcomes + .iter() + .filter(|outcome| matches!(outcome, RecordOutcome::Duplicate)) + .count(), + 1 + ); + + let winner = store.get("same-parent").unwrap(); + assert_eq!(winner.receipt().as_ref(), Some(recorded)); + assert_eq!(winner.attempts, 1); + assert_eq!( + first_queue.status().await.total + second_queue.status().await.total, + 1 + ); +} + +#[tokio::test] +async fn concurrent_tombstone_wins_over_a_first_completion_write() { + let store = blocking_first_read_store("deleted"); + let router = Arc::new(CompletionRouter::new(store.clone())); + let tombstoner = CompletionRouter::new(store.clone()); + let writer_router = Arc::clone(&router); + let writer = std::thread::spawn(move || { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + runtime.block_on(writer_router.record(record("deleted", "parent"))) + }); + store.entered.wait(); + + assert_eq!( + tombstoner.tombstone("deleted").unwrap(), + TombstoneOutcome::Reserved + ); + store.resume.wait(); + assert_eq!(writer.join().unwrap().unwrap(), RecordOutcome::Suppressed); + assert_eq!( + store.inner.get("deleted").unwrap().state, + CompletionState::Tombstoned + ); +} + #[tokio::test] async fn claim_respects_max_and_orders_by_finish_time() { let router = router(); @@ -140,11 +624,11 @@ async fn failed_attempts_move_to_gave_up_at_max_attempts() { for attempt in 1..=2 { let batch = router.claim_pending("p", 1).unwrap(); assert_eq!(batch[0].attempts, attempt); - assert!(router.mark_failed(&["t1"]).unwrap().is_empty()); + assert!(router.mark_failed(&receipts(&batch)).unwrap().is_empty()); } let batch = router.claim_pending("p", 1).unwrap(); assert_eq!(batch[0].attempts, 3); - let gave_up = router.mark_failed(&["t1"]).unwrap(); + let gave_up = router.mark_failed(&receipts(&batch)).unwrap(); assert_eq!(ids(&gave_up), ["t1"]); assert_eq!(gave_up[0].state, CompletionState::GaveUp); assert_eq!(gave_up[0].result.text, "result of t1"); @@ -158,154 +642,322 @@ async fn default_max_attempts_is_five() { let router = router(); router.record(record("t1", "p")).await.unwrap(); for _ in 0..4 { - router.claim_pending("p", 1).unwrap(); - assert!(router.mark_failed(&["t1"]).unwrap().is_empty()); + let batch = router.claim_pending("p", 1).unwrap(); + assert!(router.mark_failed(&receipts(&batch)).unwrap().is_empty()); } - router.claim_pending("p", 1).unwrap(); - assert_eq!(router.mark_failed(&["t1"]).unwrap().len(), 1); + let batch = router.claim_pending("p", 1).unwrap(); + assert_eq!(router.mark_failed(&receipts(&batch)).unwrap().len(), 1); } #[tokio::test] async fn mark_failed_on_a_settled_record_is_ignored() { let router = router(); router.record(record("t1", "p")).await.unwrap(); - router.claim_pending("p", 1).unwrap(); - router.mark_delivered(&["t1"]).unwrap(); - assert!(router.mark_failed(&["t1", "unknown"]).unwrap().is_empty()); -} - -fn record_with_artifact(task: &str, parent: &str) -> CompletionRecord { - let mut record = record(task, parent).with_label("indexer"); - record.result.omitted_chars = 42; - record.result.artifact = Some(CompletionArtifact { - id: "art-1".into(), - media_type: Some("text/plain".into()), - metadata: [("bytes".to_string(), "9000".to_string())].into(), + let batch = router.claim_pending("p", 1).unwrap(); + router.mark_delivered(&receipts(&batch)).unwrap(); + let mut stale_and_unknown = receipts(&batch); + stale_and_unknown.push(CompletionReceipt { + task_id: "unknown".into(), + parent_key: "p".into(), + parent_epoch: 1, + token: "no-claim".into(), }); - record + assert!(router.mark_failed(&stale_and_unknown).unwrap().is_empty()); } #[tokio::test] -async fn a_delivered_record_keeps_its_bookkeeping_but_not_its_result() { - let store = Arc::new(InMemoryCompletionStore::new()); - let router = CompletionRouter::new(store.clone()); - let recorded = record_with_artifact("t1", "p"); - router.record(recorded.clone()).await.unwrap(); - router.claim_pending("p", 1).unwrap(); - assert_eq!(store.get("t1").unwrap().result.text, "result of t1"); - assert_eq!(router.mark_delivered(&["t1"]).unwrap(), 1); - let settled = store.get("t1").unwrap(); - assert_eq!(settled.state, CompletionState::Delivered); - assert_eq!(settled.result, CompletionResult::default()); - assert_eq!(settled.attempts, 1); - assert_eq!(settled.task_id, "t1"); - assert_eq!(settled.parent_key, "p"); - assert_eq!(settled.agent_id, "worker"); - assert_eq!(settled.status, CompletionStatus::Success); - assert_eq!(settled.label.as_deref(), Some("indexer")); - assert_eq!(settled.finished_at, recorded.finished_at); - assert!(settled.updated_at >= recorded.updated_at); - // Dedupe and the tombstone guard still work on the emptied record. - assert_eq!( - router.record(record("t1", "p")).await.unwrap(), - RecordOutcome::Duplicate - ); +async fn followup_mode_pushes_onto_the_followup_lane_of_a_live_parent() { + let router = router(); + let q = queue(); + router.attach_parent("p", q.clone()); + let outcome = router + .record(record("t1", "p").with_notify_mode(NotifyMode::Followup)) + .await + .unwrap(); + assert!(matches!( + &outcome, + RecordOutcome::Recorded { + lane: Some(QueueLane::Followup), + receipt: Some(_) + } + )); + let receipt = match outcome { + RecordOutcome::Recorded { + receipt: Some(receipt), + .. + } => receipt, + _ => unreachable!(), + }; + let status = q.status().await; assert_eq!( - router.tombstone("t1").unwrap(), - TombstoneOutcome::AlreadyFinal + (status.followups, status.collects, status.steers), + (1, 0, 0) ); + let pushed = q + .take_admitted( + QueueLane::Followup, + tinyagents_harness::run_queue::QueueMode::All, + ) + .await; + assert!(format!("{:?}", pushed[0].value).contains("t1")); + // Taking the item models the queue-to-turn handoff. Once the turn has + // accepted the message, the short admission gate must be released before + // the router can settle or reclaim its delivery claim. + drop(pushed); + // Pushed but not acknowledged: still pending, leased, not claimable. + assert_eq!(router.in_flight_for("p").len(), 1); + assert!(router.claim_pending("p", 10).unwrap().is_empty()); + router.mark_delivered(&[receipt]).unwrap(); assert!(router.pending_for("p").is_empty()); + assert!(router.in_flight_for("p").is_empty()); } #[tokio::test] -async fn tombstone_and_cancel_parent_drop_the_withdrawn_result() { - let store = Arc::new(InMemoryCompletionStore::new()); - let router = CompletionRouter::new(store.clone()); +async fn queued_admission_after_store_replacement_observes_new_cancellation_marker() { + let source = Arc::new(InMemoryCompletionStore::new()); + let source_port: Arc<dyn CompletionStore> = source; + let router = CompletionRouter::new(source_port.clone()); + let q = queue(); + router.attach_parent("p", q.clone()); + router.record(record("queued", "p")).await.unwrap(); + + let replacement = Arc::new(InMemoryCompletionStore::new()); + let replacement_port: Arc<dyn CompletionStore> = replacement; router - .record(record_with_artifact("waited", "p")) - .await + .replace_store_if_current(&source_port, replacement_port) .unwrap(); + router.cancel_parent("p").unwrap(); + + let admitted = q + .take_admitted( + QueueLane::Followup, + tinyagents_harness::run_queue::QueueMode::All, + ) + .await; + assert!( + admitted.is_empty(), + "queued pre-migration capability must not bypass cancellation on replacement store" + ); +} + +#[tokio::test] +async fn queued_admission_survives_store_replacement_without_cancellation() { + let source = Arc::new(InMemoryCompletionStore::new()); + let source_port: Arc<dyn CompletionStore> = source; + let router = CompletionRouter::new(source_port.clone()); + let q = queue(); + router.attach_parent("p", q.clone()); + router.record(record("queued", "p")).await.unwrap(); + + let replacement: Arc<dyn CompletionStore> = Arc::new(InMemoryCompletionStore::new()); router - .record(record_with_artifact("orphaned", "p")) - .await + .replace_store_if_current(&source_port, replacement) .unwrap(); - router.record(record("other", "q")).await.unwrap(); - assert_eq!( - router.tombstone("waited").unwrap(), - TombstoneOutcome::Suppressed - ); - assert_eq!(router.cancel_parent("p").unwrap(), 1); - for task in ["waited", "orphaned"] { - let withdrawn = store.get(task).unwrap(); - assert_eq!(withdrawn.state, CompletionState::Tombstoned, "{task}"); - assert_eq!(withdrawn.result, CompletionResult::default(), "{task}"); - assert_eq!(withdrawn.parent_key, "p", "{task}"); - assert_eq!( - router.record(record(task, "p")).await.unwrap(), - RecordOutcome::Suppressed, - "{task}" - ); + + let admitted = q + .take_admitted( + QueueLane::Followup, + tinyagents_harness::run_queue::QueueMode::All, + ) + .await; + assert_eq!(admitted.len(), 1); + assert!(format!("{:?}", admitted[0].value).contains("queued")); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn queue_admission_does_not_deadlock_with_store_migration_waiting_for_parent_gate() { + let source = Arc::new(BlockingTryList { + inner: Arc::new(InMemoryCompletionStore::new()), + entered: Arc::new(std::sync::Barrier::new(2)), + resume: Arc::new(std::sync::Barrier::new(2)), + pause_list: std::sync::atomic::AtomicBool::new(false), + gate_wait: Arc::new(std::sync::Mutex::new(None)), + }); + let source_port: Arc<dyn CompletionStore> = source.clone(); + let router = Arc::new(CompletionRouter::new(source_port.clone())); + let q = queue(); + router.attach_parent("p", q.clone()); + router.record(record("queued", "p")).await.unwrap(); + + let (entered_tx, entered_rx) = std::sync::mpsc::channel(); + let (resume_tx, resume_rx) = std::sync::mpsc::channel(); + *source + .gate_wait + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) = Some((entered_tx, resume_rx)); + let migration_router = Arc::clone(&router); + let migration_source = Arc::clone(&source_port); + let destination: Arc<dyn CompletionStore> = Arc::new(InMemoryCompletionStore::new()); + let migration = std::thread::spawn(move || { + migration_router + .replace_store_if_current(&migration_source, destination) + .unwrap() + }); + entered_rx.recv().unwrap(); + + // The migrator is paused inside durable gate acquisition. Queue admission + // must see the slot as migrating and return the item to its lane without + // holding the slot lock while waiting for that same parent gate. + let queue_for_take = Arc::clone(&q); + let mut take = tokio::task::spawn_blocking(move || { + tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap() + .block_on(queue_for_take.take_admitted(QueueLane::Followup, QueueMode::All)) + }); + let take_before_release = tokio::time::timeout(Duration::from_secs(2), &mut take).await; + let take_timed_out = take_before_release.is_err(); + let (item_was_deferred, queued_after_take) = match take_before_release { + Ok(result) => { + let items = result.unwrap(); + let deferred = items.is_empty(); + let queued = q.status().await.followups; + drop(items); + (deferred, queued) + } + Err(_) => (false, q.status().await.followups), + }; + resume_tx.send(()).unwrap(); + migration.join().unwrap(); + if take_timed_out { + drop(take.await.unwrap()); + } + assert!(item_was_deferred); + assert_eq!(queued_after_take, 1); + let delivered = q.take_admitted(QueueLane::Followup, QueueMode::All).await; + assert_eq!(delivered.len(), 1); + assert!(format!("{:?}", delivered[0].value).contains("queued")); +} + +#[test] +fn admission_rejects_a_stale_batch_key_when_store_changes_during_acquire() { + let source = Arc::new(BlockingTryList { + inner: Arc::new(InMemoryCompletionStore::new()), + entered: Arc::new(std::sync::Barrier::new(2)), + resume: Arc::new(std::sync::Barrier::new(2)), + pause_list: std::sync::atomic::AtomicBool::new(false), + gate_wait: Arc::new(std::sync::Mutex::new(None)), + }); + let source_port: Arc<dyn CompletionStore> = source.clone(); + let router = Arc::new(CompletionRouter::new(source_port.clone())); + let queued = record("queued", "p"); + tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap() + .block_on(router.record(queued.clone())) + .unwrap(); + let queued = source.get("queued").unwrap(); + let admission = Arc::new(crate::completions::CompletionAdmission::for_router( + &router, &queued, + )); + let old_key = admission.batch_key().unwrap(); + + // Pause after acquire_for has snapshotted the old binding but before it + // can obtain the old store's parent gate. + let (entered_tx, entered_rx) = std::sync::mpsc::channel(); + let (resume_tx, resume_rx) = std::sync::mpsc::channel(); + *source + .gate_wait + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) = Some((entered_tx, resume_rx)); + let acquiring = Arc::clone(&admission); + let expected_key = old_key.clone(); + let acquire = std::thread::spawn(move || acquiring.acquire_for(Some(&expected_key))); + entered_rx.recv().unwrap(); + + let migration_router = Arc::clone(&router); + let migration_source = Arc::clone(&source_port); + let destination: Arc<dyn CompletionStore> = Arc::new(InMemoryCompletionStore::new()); + let migration = std::thread::spawn(move || { + migration_router + .replace_store_if_current(&migration_source, destination) + .unwrap() + }); + migration.join().unwrap(); + resume_tx.send(()).unwrap(); + + assert!(matches!( + acquire.join().unwrap(), + Err(tinyagents_harness::run_queue::QueueAdmissionError::Retry) + )); + assert_ne!(admission.batch_key().as_deref(), Some(old_key.as_str())); +} + +struct LostRenewalGateStore { + inner: InMemoryCompletionStore, + acquired: std::sync::Mutex<Option<std::sync::mpsc::Sender<()>>>, +} + +impl CompletionStore for LostRenewalGateStore { + fn acquire_parent_gate( + &self, + parent_key: &str, + holder: &str, + ttl: Duration, + ) -> tinyagents_harness::error::Result<Option<ParentGateClaim>> { + let claim = self.inner.acquire_parent_gate(parent_key, holder, ttl)?; + if let Some(sender) = claim.as_ref().and_then(|_| { + self.acquired + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .take() + }) { + let _ = sender.send(()); + } + Ok(claim) + } + + fn renew_parent_gate( + &self, + _parent_key: &str, + _claim: &ParentGateClaim, + _ttl: Duration, + ) -> tinyagents_harness::error::Result<bool> { + Ok(false) + } + + fn release_parent_gate( + &self, + parent_key: &str, + claim: &ParentGateClaim, + ) -> tinyagents_harness::error::Result<bool> { + self.inner.release_parent_gate(parent_key, claim) + } + + fn get(&self, task_id: &str) -> Option<CompletionRecord> { + self.inner.get(task_id) + } + + fn put(&self, record: &CompletionRecord) -> tinyagents_harness::error::Result<()> { + self.inner.put(record) } - assert_eq!( - store.get("other").unwrap().result.text, - "result of other", - "another parent's pending record is untouched" - ); -} -#[tokio::test] -async fn the_give_up_writer_gets_the_full_result_and_the_store_does_not() { - let store = Arc::new(InMemoryCompletionStore::new()); - let router = CompletionRouter::new(store.clone()).with_max_attempts(1); - let recorded = record_with_artifact("t1", "p"); - router.record(recorded.clone()).await.unwrap(); - router.claim_pending("p", 1).unwrap(); - let gave_up = router.mark_failed(&["t1"]).unwrap(); - assert_eq!(gave_up.len(), 1); - assert_eq!(gave_up[0].state, CompletionState::GaveUp); - assert_eq!(gave_up[0].attempts, 1); - assert_eq!(gave_up[0].result, recorded.result); - let stored = store.get("t1").unwrap(); - assert_eq!(stored.state, CompletionState::GaveUp); - assert_eq!(stored.attempts, 1); - assert_eq!(stored.updated_at, gave_up[0].updated_at); - assert_eq!(stored.result, CompletionResult::default()); - assert!(router.claim_pending("p", 1).unwrap().is_empty()); - assert_eq!( - router.record(record("t1", "p")).await.unwrap(), - RecordOutcome::Duplicate - ); - assert!(router.mark_failed(&["t1"]).unwrap().is_empty()); + fn list(&self, parent_key: Option<&str>) -> Vec<CompletionRecord> { + self.inner.list(parent_key) + } } -#[tokio::test] -async fn followup_mode_pushes_onto_the_followup_lane_of_a_live_parent() { - let router = router(); - let q = queue(); - router.attach_parent("p", q.clone()); - let outcome = router - .record(record("t1", "p").with_notify_mode(NotifyMode::Followup)) - .await - .unwrap(); - assert_eq!( - outcome, - RecordOutcome::Recorded { - lane: Some(QueueLane::Followup) - } - ); - let status = q.status().await; - assert_eq!( - (status.followups, status.collects, status.steers), - (1, 0, 0) - ); - let pushed = q.drain(QueueLane::Followup).await; - assert!(format!("{:?}", pushed[0]).contains("t1")); - // Pushed but not acknowledged: still pending, leased, not claimable. - assert_eq!(router.in_flight_for("p").len(), 1); - assert!(router.claim_pending("p", 10).unwrap().is_empty()); - router.mark_delivered(&["t1"]).unwrap(); - assert!(router.pending_for("p").is_empty()); - assert!(router.in_flight_for("p").is_empty()); +#[test] +fn admission_fails_closed_if_native_gate_expires_while_waiting_for_local_lifecycle() { + let lifecycle = super::lifecycle::ParentLifecycle::new(); + let local = lifecycle.acquire("p").unwrap(); + let (acquired_tx, acquired_rx) = std::sync::mpsc::channel(); + let store: Arc<dyn CompletionStore> = Arc::new(LostRenewalGateStore { + inner: InMemoryCompletionStore::new(), + acquired: std::sync::Mutex::new(Some(acquired_tx)), + }); + let waiting_store = Arc::clone(&store); + let waiting_lifecycle = Arc::clone(&lifecycle); + let admission = std::thread::spawn(move || { + super::router::ParentAdmissionPermit::acquire(waiting_store, waiting_lifecycle, "p", 1) + }); + + acquired_rx.recv().unwrap(); + drop(local); + + assert!(admission.join().unwrap().is_err()); } #[tokio::test] @@ -317,12 +969,13 @@ async fn collect_mode_pushes_onto_the_collect_lane() { .record(record("t1", "p").with_notify_mode(NotifyMode::Collect)) .await .unwrap(); - assert_eq!( + assert!(matches!( outcome, RecordOutcome::Recorded { - lane: Some(QueueLane::Collect) + lane: Some(QueueLane::Collect), + receipt: Some(_) } - ); + )); let status = q.status().await; assert_eq!((status.followups, status.collects), (0, 1)); } @@ -331,7 +984,13 @@ async fn collect_mode_pushes_onto_the_collect_lane() { async fn followup_without_a_live_parent_stays_pending() { let router = router(); let outcome = router.record(record("t1", "p")).await.unwrap(); - assert_eq!(outcome, RecordOutcome::Recorded { lane: None }); + assert_eq!( + outcome, + RecordOutcome::Recorded { + lane: None, + receipt: None + } + ); assert_eq!(ids(&router.claim_pending("p", 10).unwrap()), ["t1"]); } @@ -354,7 +1013,13 @@ async fn hold_mode_waits_for_the_next_turn() { .record(record("t1", "p").with_notify_mode(NotifyMode::HoldForNextTurn)) .await .unwrap(); - assert_eq!(outcome, RecordOutcome::Recorded { lane: None }); + assert_eq!( + outcome, + RecordOutcome::Recorded { + lane: None, + receipt: None + } + ); assert_eq!(q.status().await.total, 0); assert!( router.claim_pending("p", 10).unwrap().is_empty(), @@ -364,7 +1029,7 @@ async fn hold_mode_waits_for_the_next_turn() { let released = router.begin_turn("p").unwrap(); assert_eq!(ids(&released), ["t1"]); assert_eq!(released[0].attempts, 1); - router.mark_delivered(&["t1"]).unwrap(); + router.mark_delivered(&receipts(&released)).unwrap(); assert!(router.begin_turn("p").unwrap().is_empty()); } @@ -377,7 +1042,13 @@ async fn off_mode_is_recorded_for_the_parent_to_pull() { .record(record("t1", "p").with_notify_mode(NotifyMode::Off)) .await .unwrap(); - assert_eq!(outcome, RecordOutcome::Recorded { lane: None }); + assert_eq!( + outcome, + RecordOutcome::Recorded { + lane: None, + receipt: None + } + ); assert_eq!(q.status().await.total, 0); assert!(router.claim_pending("p", 10).unwrap().is_empty()); assert!(router.begin_turn("p").unwrap().is_empty()); @@ -407,8 +1078,11 @@ async fn cancel_parent_drops_pending_and_late_completions() { ); assert!(router.pending_for("p").is_empty()); router.resume_parent("p"); + let generation = router.admit_child_generation("p").unwrap().unwrap(); + let mut next = record("t3", "p"); + next.parent_epoch = generation; assert!(matches!( - router.record(record("t3", "p")).await.unwrap(), + router.record(next).await.unwrap(), RecordOutcome::Recorded { .. } )); } @@ -419,11 +1093,74 @@ async fn pending_for_lists_only_undelivered_records_of_the_parent() { router.record(record("t1", "p")).await.unwrap(); router.record(record("t2", "p")).await.unwrap(); router.record(record("t3", "other")).await.unwrap(); - router.claim_pending("p", 1).unwrap(); - router.mark_delivered(&["t1"]).unwrap(); + let first = router.claim_pending("p", 1).unwrap(); + router.mark_delivered(&receipts(&first)).unwrap(); assert_eq!(ids(&router.pending_for("p")), ["t2"]); } +#[tokio::test] +async fn parents_with_pending_and_cancel_all_use_durable_parent_markers() { + let router = router(); + router + .record(record("pending-a", "parent-a").with_notify_mode(NotifyMode::Off)) + .await + .unwrap(); + router + .record(record("pending-b", "parent-b").with_notify_mode(NotifyMode::Off)) + .await + .unwrap(); + router + .record(record("settled", "parent-c").with_notify_mode(NotifyMode::Followup)) + .await + .unwrap(); + let settled = router.claim_pending("parent-c", 1).unwrap(); + router.mark_delivered(&receipts(&settled)).unwrap(); + + assert_eq!( + router.parents_with_pending().unwrap(), + vec!["parent-a", "parent-b"] + ); + assert_eq!(router.cancel_all_parents().unwrap(), 2); + assert!(router.parents_with_pending().unwrap().is_empty()); + assert!(!router.admit_child("parent-a").unwrap()); + assert!(!router.admit_child("parent-b").unwrap()); + assert!( + !router.admit_child("parent-c").unwrap(), + "the bulk close also closes a known open parent with only settled records" + ); +} + +#[test] +fn cancelling_a_deleted_parent_does_not_reopen_its_lifecycle() { + let router = router(); + router.delete_parent("deleted").unwrap(); + router.cancel_parent("deleted").unwrap(); + router.resume_parent("deleted"); + + assert!( + !router.admit_child("deleted").unwrap(), + "a later stop marker must not downgrade permanent deletion" + ); +} + +#[tokio::test] +async fn cancel_all_closes_an_attached_parent_before_its_first_completion() { + let router = router(); + let queue = queue(); + router.attach_parent("idle-parent", queue.clone()); + + assert_eq!(router.cancel_all_parents().unwrap(), 0); + assert!(matches!( + router + .record(record("late-child", "idle-parent")) + .await + .unwrap(), + RecordOutcome::Suppressed + )); + assert_eq!(queue.status().await.total, 0); + assert!(!router.admit_child("idle-parent").unwrap()); +} + #[test] fn neutral_formatter_escapes_child_text() { let hostile = CompletionRecord::new( @@ -452,8 +1189,13 @@ async fn a_custom_formatter_frames_the_pushed_message() { let q = queue(); router.attach_parent("p", q.clone()); router.record(record("t1", "p")).await.unwrap(); - let pushed = q.drain(QueueLane::Followup).await; - assert!(format!("{:?}", pushed[0]).contains("background_agent_result count=1")); + let pushed = q + .take_admitted( + QueueLane::Followup, + tinyagents_harness::run_queue::QueueMode::All, + ) + .await; + assert!(format!("{:?}", pushed[0].value).contains("background_agent_result count=1")); } #[tokio::test] @@ -473,9 +1215,9 @@ async fn a_push_lost_with_the_queue_is_redelivered() { async fn release_makes_an_abandoned_claim_claimable_again() { let router = router(); router.record(record("t1", "p")).await.unwrap(); - router.claim_pending("p", 1).unwrap(); + let first = router.claim_pending("p", 1).unwrap(); assert!(router.claim_pending("p", 1).unwrap().is_empty()); - router.release(&["t1"]); + router.release(&receipts(&first)); assert_eq!(router.claim_pending("p", 1).unwrap().len(), 1); } @@ -502,7 +1244,7 @@ async fn replacing_a_router_preserves_its_leases_and_claim_acknowledgements() { )); let duplicate_claim = source_router.pull("p", 1).unwrap(); - let old_ack = source_router.mark_delivered(&["leased"]).unwrap(); + let old_ack = source_router.mark_delivered(&receipts(&claimed)).unwrap(); let replacement_record = replacement_store.get("leased").unwrap(); assert!( duplicate_claim.is_empty(), @@ -526,13 +1268,13 @@ fn store_migration_copies_all_states_and_keeps_router_policy() { source .put(&record("pending", "p").with_notify_mode(NotifyMode::Off)) .unwrap(); - router.pull("p", 1).unwrap(); + let _pending_claim = router.pull("p", 1).unwrap(); source .put(&record("delivered", "p").with_notify_mode(NotifyMode::Off)) .unwrap(); - router.pull("p", 1).unwrap(); - router.mark_delivered(&["delivered"]).unwrap(); + let delivered = router.pull("p", 1).unwrap(); + router.mark_delivered(&receipts(&delivered)).unwrap(); source .put(&record("tombstoned", "p").with_notify_mode(NotifyMode::Off)) @@ -542,10 +1284,10 @@ fn store_migration_copies_all_states_and_keeps_router_policy() { source .put(&record("gave-up", "p").with_notify_mode(NotifyMode::Off)) .unwrap(); - router.pull("p", 1).unwrap(); - router.release(&["gave-up"]); - router.pull("p", 1).unwrap(); - router.mark_failed(&["gave-up"]).unwrap(); + let attempt = router.pull("p", 1).unwrap(); + router.release(&receipts(&attempt)); + let attempt = router.pull("p", 1).unwrap(); + router.mark_failed(&receipts(&attempt)).unwrap(); router.cancel_parent("q").unwrap(); let destination = Arc::new(InMemoryCompletionStore::new()); @@ -602,9 +1344,39 @@ struct FailingSecondPut { } impl CompletionStore for FailingSecondPut { + fn acquire_parent_gate( + &self, + parent_key: &str, + holder: &str, + ttl: std::time::Duration, + ) -> tinyagents_harness::error::Result<Option<ParentGateClaim>> { + self.inner.acquire_parent_gate(parent_key, holder, ttl) + } + fn renew_parent_gate( + &self, + parent_key: &str, + claim: &ParentGateClaim, + ttl: std::time::Duration, + ) -> tinyagents_harness::error::Result<bool> { + self.inner.renew_parent_gate(parent_key, claim, ttl) + } + fn release_parent_gate( + &self, + parent_key: &str, + claim: &ParentGateClaim, + ) -> tinyagents_harness::error::Result<bool> { + self.inner.release_parent_gate(parent_key, claim) + } + fn get(&self, task_id: &str) -> Option<CompletionRecord> { self.inner.get(task_id) } + fn try_get( + &self, + task_id: &str, + ) -> tinyagents_harness::error::Result<Option<CompletionRecord>> { + self.inner.try_get(task_id) + } fn put(&self, record: &CompletionRecord) -> tinyagents_harness::error::Result<()> { use std::sync::atomic::Ordering::SeqCst; if self.puts.fetch_add(1, SeqCst) >= self.fail_from.load(SeqCst) { @@ -614,6 +1386,83 @@ impl CompletionStore for FailingSecondPut { } self.inner.put(record) } + fn put_fenced( + &self, + record: &CompletionRecord, + parent_key: &str, + gate: &ParentGateClaim, + ) -> tinyagents_harness::error::Result<()> { + use std::sync::atomic::Ordering::SeqCst; + if self.puts.fetch_add(1, SeqCst) >= self.fail_from.load(SeqCst) { + return Err(tinyagents_harness::error::TinyAgentsError::Graph( + "disk".into(), + )); + } + self.inner.put_fenced(record, parent_key, gate) + } + fn insert_if_absent( + &self, + record: &CompletionRecord, + gate: &ParentGateClaim, + ) -> tinyagents_harness::error::Result<bool> { + use std::sync::atomic::Ordering::SeqCst; + if self.puts.fetch_add(1, SeqCst) >= self.fail_from.load(SeqCst) { + return Err(tinyagents_harness::error::TinyAgentsError::Graph( + "disk".into(), + )); + } + self.inner.insert_if_absent(record, gate) + } + fn claim_pending( + &self, + request: PendingClaimRequest, + ) -> tinyagents_harness::error::Result<Vec<CompletionRecord>> { + use std::sync::atomic::Ordering::SeqCst; + let claimed = self.inner.claim_pending(request.clone())?; + if (0..claimed.len()).any(|_| self.puts.fetch_add(1, SeqCst) >= self.fail_from.load(SeqCst)) + { + for record in &claimed { + if let Some(receipt) = record.receipt() { + self.inner.resolve_claim(ResolveClaimRequest { + task_id: receipt.task_id, + parent_key: receipt.parent_key, + parent_epoch: receipt.parent_epoch, + gate: request.gate.clone(), + token: receipt.token, + resolution: ClaimResolution::Release, + })?; + } + } + return Err(tinyagents_harness::error::TinyAgentsError::Graph( + "disk".into(), + )); + } + Ok(claimed) + } + fn claim_task( + &self, + request: TaskClaimRequest, + ) -> tinyagents_harness::error::Result<Option<CompletionRecord>> { + use std::sync::atomic::Ordering::SeqCst; + if self.puts.fetch_add(1, SeqCst) >= self.fail_from.load(SeqCst) { + return Err(tinyagents_harness::error::TinyAgentsError::Graph( + "disk".into(), + )); + } + self.inner.claim_task(request) + } + fn resolve_claim( + &self, + request: ResolveClaimRequest, + ) -> tinyagents_harness::error::Result<Option<CompletionRecord>> { + use std::sync::atomic::Ordering::SeqCst; + if self.puts.fetch_add(1, SeqCst) >= self.fail_from.load(SeqCst) { + return Err(tinyagents_harness::error::TinyAgentsError::Graph( + "disk".into(), + )); + } + self.inner.resolve_claim(request) + } fn list(&self, parent_key: Option<&str>) -> Vec<CompletionRecord> { self.inner.list(parent_key) } @@ -623,19 +1472,111 @@ struct FailingList { inner: InMemoryCompletionStore, } +#[test] +fn parent_cancellation_propagates_pending_enumeration_failure() { + let store = Arc::new(FailingList { + inner: InMemoryCompletionStore::new(), + }); + let router = CompletionRouter::new(store.clone()); + + assert!(router.cancel_parent("parent").is_err()); + assert_eq!( + store + .try_get("\u{1}parent-cancelled:parent") + .unwrap() + .unwrap() + .state, + CompletionState::Tombstoned, + "the durable closed marker is written before enumeration and remains fail-closed" + ); +} + +type GateWait = + Arc<std::sync::Mutex<Option<(std::sync::mpsc::Sender<()>, std::sync::mpsc::Receiver<()>)>>>; + struct BlockingTryList { - inner: InMemoryCompletionStore, + inner: Arc<InMemoryCompletionStore>, entered: Arc<std::sync::Barrier>, resume: Arc<std::sync::Barrier>, + pause_list: std::sync::atomic::AtomicBool, + gate_wait: GateWait, } impl CompletionStore for BlockingTryList { + fn acquire_parent_gate( + &self, + parent_key: &str, + holder: &str, + ttl: std::time::Duration, + ) -> tinyagents_harness::error::Result<Option<ParentGateClaim>> { + let wait = self + .gate_wait + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .take(); + if let Some((entered, resume)) = wait { + entered.send(()).expect("test observer remains connected"); + resume + .recv() + .expect("test releases blocked gate acquisition"); + } + self.inner.acquire_parent_gate(parent_key, holder, ttl) + } + fn renew_parent_gate( + &self, + parent_key: &str, + claim: &ParentGateClaim, + ttl: std::time::Duration, + ) -> tinyagents_harness::error::Result<bool> { + self.inner.renew_parent_gate(parent_key, claim, ttl) + } + fn release_parent_gate( + &self, + parent_key: &str, + claim: &ParentGateClaim, + ) -> tinyagents_harness::error::Result<bool> { + self.inner.release_parent_gate(parent_key, claim) + } + fn get(&self, task_id: &str) -> Option<CompletionRecord> { self.inner.get(task_id) } fn put(&self, record: &CompletionRecord) -> tinyagents_harness::error::Result<()> { self.inner.put(record) } + fn put_fenced( + &self, + record: &CompletionRecord, + parent_key: &str, + gate: &ParentGateClaim, + ) -> tinyagents_harness::error::Result<()> { + self.inner.put_fenced(record, parent_key, gate) + } + fn insert_if_absent( + &self, + record: &CompletionRecord, + gate: &ParentGateClaim, + ) -> tinyagents_harness::error::Result<bool> { + self.inner.insert_if_absent(record, gate) + } + fn claim_task( + &self, + request: TaskClaimRequest, + ) -> tinyagents_harness::error::Result<Option<CompletionRecord>> { + self.inner.claim_task(request) + } + fn claim_pending( + &self, + request: PendingClaimRequest, + ) -> tinyagents_harness::error::Result<Vec<CompletionRecord>> { + self.inner.claim_pending(request) + } + fn resolve_claim( + &self, + request: ResolveClaimRequest, + ) -> tinyagents_harness::error::Result<Option<CompletionRecord>> { + self.inner.resolve_claim(request) + } fn list(&self, parent_key: Option<&str>) -> Vec<CompletionRecord> { self.inner.list(parent_key) } @@ -643,19 +1584,57 @@ impl CompletionStore for BlockingTryList { &self, parent_key: Option<&str>, ) -> tinyagents_harness::error::Result<Vec<CompletionRecord>> { - self.entered.wait(); - self.resume.wait(); - Ok(self.inner.list(parent_key)) + let records = self.inner.list(parent_key); + if self + .pause_list + .swap(false, std::sync::atomic::Ordering::SeqCst) + { + self.entered.wait(); + self.resume.wait(); + } + Ok(records) } } impl CompletionStore for FailingList { + fn acquire_parent_gate( + &self, + parent_key: &str, + holder: &str, + ttl: std::time::Duration, + ) -> tinyagents_harness::error::Result<Option<ParentGateClaim>> { + self.inner.acquire_parent_gate(parent_key, holder, ttl) + } + fn renew_parent_gate( + &self, + parent_key: &str, + claim: &ParentGateClaim, + ttl: std::time::Duration, + ) -> tinyagents_harness::error::Result<bool> { + self.inner.renew_parent_gate(parent_key, claim, ttl) + } + fn release_parent_gate( + &self, + parent_key: &str, + claim: &ParentGateClaim, + ) -> tinyagents_harness::error::Result<bool> { + self.inner.release_parent_gate(parent_key, claim) + } + fn get(&self, task_id: &str) -> Option<CompletionRecord> { self.inner.get(task_id) } fn put(&self, record: &CompletionRecord) -> tinyagents_harness::error::Result<()> { self.inner.put(record) } + fn put_fenced( + &self, + record: &CompletionRecord, + parent_key: &str, + gate: &ParentGateClaim, + ) -> tinyagents_harness::error::Result<()> { + self.inner.put_fenced(record, parent_key, gate) + } fn list(&self, parent_key: Option<&str>) -> Vec<CompletionRecord> { self.inner.list(parent_key) } @@ -713,9 +1692,11 @@ fn a_source_enumeration_error_does_not_replace_the_store() { #[test] fn a_record_racing_store_migration_is_not_lost() { let source = Arc::new(BlockingTryList { - inner: InMemoryCompletionStore::new(), + inner: Arc::new(InMemoryCompletionStore::new()), entered: Arc::new(std::sync::Barrier::new(2)), resume: Arc::new(std::sync::Barrier::new(2)), + pause_list: std::sync::atomic::AtomicBool::new(true), + gate_wait: Arc::new(std::sync::Mutex::new(None)), }); let source_port: Arc<dyn CompletionStore> = source.clone(); let router = Arc::new(CompletionRouter::new(source_port.clone())); @@ -755,6 +1736,49 @@ fn a_record_racing_store_migration_is_not_lost() { assert!(destination.get("during").is_some()); } +#[tokio::test] +async fn store_migration_copies_parent_state_after_acquiring_source_gates() { + let source = Arc::new(BlockingTryList { + inner: Arc::new(InMemoryCompletionStore::new()), + entered: Arc::new(std::sync::Barrier::new(2)), + resume: Arc::new(std::sync::Barrier::new(2)), + pause_list: std::sync::atomic::AtomicBool::new(false), + gate_wait: Arc::new(std::sync::Mutex::new(None)), + }); + let source_port: Arc<dyn CompletionStore> = source.clone(); + let router = Arc::new(CompletionRouter::new(source_port.clone())); + source.put(&record("pending", "p")).unwrap(); + source + .pause_list + .store(true, std::sync::atomic::Ordering::SeqCst); + + let (gate_entered_tx, gate_entered_rx) = std::sync::mpsc::channel(); + let (gate_resume_tx, gate_resume_rx) = std::sync::mpsc::channel(); + *source + .gate_wait + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) = + Some((gate_entered_tx, gate_resume_rx)); + let destination: Arc<dyn CompletionStore> = Arc::new(InMemoryCompletionStore::new()); + let migration_router = Arc::clone(&router); + let migration_source = Arc::clone(&source_port); + let migration = std::thread::spawn(move || { + migration_router + .replace_store_if_current(&migration_source, destination) + .unwrap() + }); + + source.entered.wait(); + let cancellation_router = CompletionRouter::new(source.inner.clone()); + cancellation_router.cancel_parent("p").unwrap(); + source.resume.wait(); + gate_entered_rx.recv().unwrap(); + gate_resume_tx.send(()).unwrap(); + migration.join().unwrap(); + + assert_eq!(router.admit_child_generation("p").unwrap(), None); +} + #[tokio::test] async fn a_failed_batch_claim_leaves_nothing_leased() { use std::sync::atomic::Ordering::SeqCst; @@ -779,7 +1803,7 @@ async fn record_with_retries_survives_a_transient_store_failure() { let store = Arc::new(FailingSecondPut { inner: InMemoryCompletionStore::new(), puts: Default::default(), - fail_from: std::sync::atomic::AtomicUsize::new(0), + fail_from: std::sync::atomic::AtomicUsize::new(1), }); let router = CompletionRouter::new(store.clone()); assert!( @@ -790,9 +1814,9 @@ async fn record_with_retries_survives_a_transient_store_failure() { ); // Fails once, then the store recovers. store.puts.store(0, SeqCst); - store.fail_from.store(1, SeqCst); + store.fail_from.store(usize::MAX, SeqCst); let attempt = router.record_with_retries(record("t2", "p"), 1).await; - assert!(attempt.is_ok()); + assert!(attempt.is_ok(), "retry result: {attempt:?}"); store.puts.store(0, SeqCst); store.fail_from.store(0, SeqCst); let handle = { @@ -811,30 +1835,94 @@ async fn record_with_retries_survives_a_transient_store_failure() { handle.await.unwrap(); } -// A failed put mid-batch must not lose an already-given-up result. +#[tokio::test] +async fn record_resilient_uses_local_memory_only_when_explicitly_allowed() { + use crate::completions::CompletionStoreFallback; + let store = Arc::new(FailingSecondPut { + inner: InMemoryCompletionStore::new(), + puts: Default::default(), + fail_from: std::sync::atomic::AtomicUsize::new(0), + }); + let router = CompletionRouter::new(store.clone()); + assert!( + router + .record_resilient( + record("shared", "p"), + 1, + CompletionStoreFallback::FailClosed + ) + .await + .is_err() + ); + assert!(Arc::ptr_eq( + &router.store(), + &(store.clone() as Arc<dyn CompletionStore>) + )); + + let outcome = router + .record_resilient( + record("local", "p"), + 1, + CompletionStoreFallback::LocalMemory, + ) + .await + .unwrap(); + assert!(matches!(outcome, RecordOutcome::Recorded { .. })); + assert!(router.store().get("local").is_some()); +} + #[tokio::test] async fn partial_mark_failed_keeps_a_gave_up_result() { use std::sync::atomic::Ordering::SeqCst; let store = Arc::new(FailingSecondPut { inner: InMemoryCompletionStore::new(), - puts: std::sync::atomic::AtomicUsize::new(0), + puts: Default::default(), fail_from: std::sync::atomic::AtomicUsize::new(usize::MAX), }); let router = CompletionRouter::new(store.clone()).with_max_attempts(1); router.record(record("t1", "p")).await.unwrap(); router.record(record("t2", "p")).await.unwrap(); - router.claim_pending("p", 10).unwrap(); + let batch = router.claim_pending("p", 10).unwrap(); store.fail_from.store(store.puts.load(SeqCst) + 1, SeqCst); - // t2's give-up cannot be stored: t1's is still handed to the host with its - // result, and t2 stays pending so a later failure gives it up again. - let gave_up = router.mark_failed(&["t1", "t2"]).unwrap(); + + let gave_up = router.mark_failed(&receipts(&batch)).unwrap(); assert_eq!( gave_up .iter() - .map(|r| (r.task_id.as_str(), r.result.text.as_str())) + .map(|record| (record.task_id.as_str(), record.result.text.as_str())) .collect::<Vec<_>>(), [("t1", "result of t1")] ); - let t2 = store.inner.get("t2").unwrap(); - assert_eq!(t2.state, CompletionState::Pending); + assert_eq!( + store.inner.get("t1").unwrap().state, + CompletionState::GaveUp + ); + assert_eq!( + store.inner.get("t2").unwrap().state, + CompletionState::Pending + ); +} + +#[tokio::test] +async fn release_all_counts_only_the_claim_receipts_it_still_owns() { + let store = Arc::new(InMemoryCompletionStore::new()); + let first = CompletionRouter::new(store.clone()); + let second = CompletionRouter::new(store.clone()); + first.record(record("task", "p")).await.unwrap(); + let stale = first.claim_pending("p", 1).unwrap()[0].receipt().unwrap(); + second.mark_failed(std::slice::from_ref(&stale)).unwrap(); + let current = second.claim_pending("p", 1).unwrap()[0].receipt().unwrap(); + + assert_eq!(first.release_all().unwrap(), 0); + assert_eq!( + second + .mark_delivered(std::slice::from_ref(¤t)) + .unwrap(), + 1 + ); + + second.record(record("next", "p")).await.unwrap(); + assert_eq!(second.claim_pending("p", 1).unwrap().len(), 1); + assert_eq!(second.release_all().unwrap(), 1); + assert_eq!(first.claim_pending("p", 1).unwrap().len(), 1); } diff --git a/crates/tinyagents-tasks/src/completions/store.rs b/crates/tinyagents-tasks/src/completions/store.rs index a0ccbceaa..ffd3af991 100644 --- a/crates/tinyagents-tasks/src/completions/store.rs +++ b/crates/tinyagents-tasks/src/completions/store.rs @@ -6,7 +6,8 @@ use std::time::{Duration, SystemTime}; use tinyagents_harness::error::{Result, TinyAgentsError}; -use super::types::{CompletionRecord, CompletionState}; +use super::lifecycle::ParentLifecycle; +use super::types::{CompletionClaim, CompletionRecord, CompletionState, NotifyMode}; fn store_err(what: &str, err: impl std::fmt::Display) -> TinyAgentsError { TinyAgentsError::Graph(format!("completion store: {what}: {err}")) @@ -14,17 +15,148 @@ fn store_err(what: &str, err: impl std::fmt::Display) -> TinyAgentsError { /// Durable bookkeeping behind [`CompletionRouter`](super::CompletionRouter). /// -/// A store holds the latest [`CompletionRecord`] per task id. It is a plain -/// persistence seam: the router serialises its read-modify-write sequences, so -/// an implementation only has to make `put` durable before it returns. -/// Synchronous on purpose, like [`TaskStore`](crate::TaskStore). +/// A store holds the latest [`CompletionRecord`] per task id. `put` persists +/// unconditional record changes; delivery ownership uses `claim_pending`, +/// `claim_task`, and `resolve_claim`, which implementations shared by routers +/// must perform atomically. Synchronous on purpose, like +/// [`TaskStore`](crate::TaskStore). pub trait CompletionStore: Send + Sync { + /// Whether this store is intentionally process-local and may use the + /// router pool's retained lifecycle state when a new memory store opens. + fn is_process_local(&self) -> bool { + false + } + + /// Identity shared by separately opened stores coordinating through the + /// same durable backend and tenant scope. `None` falls back to this store + /// object's identity, suitable for deliberately shared in-memory stores. + fn coordination_identity(&self) -> Option<String> { + None + } + + /// Returns the lifecycle authority bound by a router pool, when this + /// process-local store belongs to exactly one pool scope. + fn parent_lifecycle(&self) -> Option<Arc<ParentLifecycle>> { + None + } + + /// Binds this store to its pool-owned lifecycle authority. Implementations + /// that opt into process-local routing must preserve the binding for all + /// handles sharing the same backing store and reject conflicting scopes. + fn bind_parent_lifecycle(&self, lifecycle: Arc<ParentLifecycle>) -> Result<()> { + let _ = lifecycle; + Err(store_err( + "parent lifecycle", + "store cannot retain a process-local lifecycle binding", + )) + } + + /// Acquires the cross-router parent admission gate. Shared durable stores + /// must implement this with a backend lease; the default fails closed. + fn acquire_parent_gate( + &self, + parent_key: &str, + holder: &str, + ttl: Duration, + ) -> Result<Option<ParentGateClaim>> { + let _ = (parent_key, holder, ttl); + Err(store_err( + "parent gate", + "store does not support durable leases", + )) + } + + /// Renews the exact holder and epoch returned by `acquire_parent_gate`. + fn renew_parent_gate( + &self, + parent_key: &str, + claim: &ParentGateClaim, + ttl: Duration, + ) -> Result<bool> { + let _ = (parent_key, claim, ttl); + Err(store_err( + "parent gate", + "store does not support durable leases", + )) + } + + /// Releases only the holder and epoch returned by `acquire_parent_gate`. + fn release_parent_gate(&self, parent_key: &str, claim: &ParentGateClaim) -> Result<bool> { + let _ = (parent_key, claim); + Err(store_err( + "parent gate", + "store does not support durable leases", + )) + } + /// The latest record for `task_id`. fn get(&self, task_id: &str) -> Option<CompletionRecord>; + /// Fallible read used to classify a conflicting conditional create. + fn try_get(&self, task_id: &str) -> Result<Option<CompletionRecord>> { + Ok(self.get(task_id)) + } + /// Upserts a record. Durable stores persist it before returning. fn put(&self, record: &CompletionRecord) -> Result<()>; + /// Upserts a record only through the supplied current parent tenure. + /// Shared durable stores must fence the write atomically; the default + /// fails closed rather than accepting an unguarded write. + fn put_fenced( + &self, + record: &CompletionRecord, + parent_key: &str, + gate: &ParentGateClaim, + ) -> Result<()> { + let _ = (record, parent_key, gate); + Err(store_err( + "fenced write", + "store does not support epoch-fenced writes", + )) + } + + /// Creates a record only if its key is absent. Shared stores must make the + /// absence check and write atomic; unsupported stores fail closed. + fn insert_if_absent(&self, record: &CompletionRecord, gate: &ParentGateClaim) -> Result<bool> { + let _ = (record, gate); + Err(store_err( + "conditional create", + "store does not support atomic create", + )) + } + + /// Atomically leases matching pending completions until `visible_after`. + /// Implementations shared across processes must use backend CAS/claim + /// semantics; the default fails closed rather than emulating a local lock. + fn claim_pending(&self, request: PendingClaimRequest) -> Result<Vec<CompletionRecord>> { + let _ = request; + Err(store_err( + "claim", + "store does not support atomic delivery claims", + )) + } + + /// Atomically claims one pending record for a live push. The default fails + /// closed; a process-local lock is not a substitute for backend ownership. + fn claim_task(&self, request: TaskClaimRequest) -> Result<Option<CompletionRecord>> { + let _ = request; + Err(store_err( + "claim task", + "store does not support atomic delivery claims", + )) + } + + /// Resolves a claim only while its token still owns the record. `None` + /// means the claim expired, was reclaimed, or the record was withdrawn. + fn resolve_claim(&self, request: ResolveClaimRequest) -> Result<Option<CompletionRecord>> { + let _ = request; + Err(store_err( + "resolve claim", + "store does not support atomic delivery claims", + )) + } + /// Latest records, for one parent or (`None`) for all. Unordered. fn list(&self, parent_key: Option<&str>) -> Vec<CompletionRecord>; @@ -45,6 +177,82 @@ pub trait CompletionStore: Send + Sync { } } +/// Fencing identity for one held parent admission lease. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct ParentGateClaim { + /// Unique holder identity for this lease tenure. + pub holder: String, + /// Backend epoch that fences any earlier tenure. + pub epoch: u64, +} + +/// Which pending completion lane a delivery worker may lease. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum CompletionClaimMode { + /// Follow-up or explicit collect delivery. + Followup, + /// Completions held for the parent's next turn. + HoldForNextTurn, + /// Pull-only completions. + Pull, +} + +/// Inputs for an atomic batch claim from a parent queue. +#[derive(Debug, Clone)] +pub struct PendingClaimRequest { + pub parent_key: String, + pub parent_epoch: u64, + pub gate: ParentGateClaim, + pub mode: CompletionClaimMode, + pub max: usize, + pub token: String, + pub visible_after: SystemTime, +} + +/// Inputs for an atomic claim of one live-push record. +#[derive(Debug, Clone)] +pub struct TaskClaimRequest { + pub task_id: String, + pub parent_key: String, + pub parent_epoch: u64, + pub gate: ParentGateClaim, + pub mode: CompletionClaimMode, + pub token: String, + pub visible_after: SystemTime, +} + +/// Inputs for settling one owned delivery claim. +#[derive(Debug, Clone)] +pub struct ResolveClaimRequest { + pub task_id: String, + pub parent_key: String, + pub parent_epoch: u64, + pub gate: ParentGateClaim, + pub token: String, + pub resolution: ClaimResolution, +} + +impl CompletionClaimMode { + pub(crate) fn accepts(self, mode: NotifyMode) -> bool { + match self { + Self::Followup => matches!(mode, NotifyMode::Followup | NotifyMode::Collect), + Self::HoldForNextTurn => mode == NotifyMode::HoldForNextTurn, + Self::Pull => mode == NotifyMode::Off, + } + } +} + +/// How a currently owned delivery claim is settled. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum ClaimResolution { + /// The parent accepted the completion. + Delivered, + /// The delivery failed; the claim is released or the record gives up. + Failed { max_attempts: u32 }, + /// The delivery was abandoned without counting another attempt. + Release, +} + fn expired(record: &CompletionRecord, now: SystemTime, retain: Duration) -> bool { record.state != CompletionState::Pending && now @@ -53,32 +261,357 @@ fn expired(record: &CompletionRecord, now: SystemTime, retain: Duration) -> bool } /// Thread-safe in-memory [`CompletionStore`]. -#[derive(Clone, Debug, Default)] +#[derive(Clone, Debug)] pub struct InMemoryCompletionStore { inner: Arc<Mutex<HashMap<String, CompletionRecord>>>, + parent_gates: Arc<Mutex<HashMap<String, MemoryParentGate>>>, + parent_gate_epochs: Arc<Mutex<HashMap<String, u64>>>, + identity: Arc<str>, + lifecycle: Arc<Mutex<Option<Arc<ParentLifecycle>>>>, +} + +impl Default for InMemoryCompletionStore { + fn default() -> Self { + Self::new() + } +} + +#[derive(Clone, Debug)] +struct MemoryParentGate { + claim: ParentGateClaim, + expires_at: SystemTime, } impl InMemoryCompletionStore { /// An empty store. pub fn new() -> Self { - Self::default() + Self { + inner: Arc::default(), + parent_gates: Arc::default(), + parent_gate_epochs: Arc::default(), + identity: Arc::from(format!("memory:{}", uuid::Uuid::new_v4().simple())), + lifecycle: Arc::default(), + } } fn lock(&self) -> Result<std::sync::MutexGuard<'_, HashMap<String, CompletionRecord>>> { self.inner.lock().map_err(|_| store_err("lock", "poisoned")) } + + fn with_current_gate<T>( + &self, + parent_key: &str, + claim: &ParentGateClaim, + operation: impl FnOnce(&mut HashMap<String, CompletionRecord>) -> Result<T>, + ) -> Result<T> { + let gates = self + .parent_gates + .lock() + .map_err(|_| store_err("parent gate", "lock poisoned"))?; + let now = SystemTime::now(); + if !gates + .get(parent_key) + .is_some_and(|gate| gate.claim == *claim && gate.expires_at > now) + { + return Err(store_err("parent gate", "claim is stale or expired")); + } + let mut records = self.lock()?; + if !gates + .get(parent_key) + .is_some_and(|gate| gate.claim == *claim && gate.expires_at > SystemTime::now()) + { + return Err(store_err("parent gate", "claim expired before mutation")); + } + operation(&mut records) + } } impl CompletionStore for InMemoryCompletionStore { + fn is_process_local(&self) -> bool { + true + } + + fn coordination_identity(&self) -> Option<String> { + Some(self.identity.to_string()) + } + + fn parent_lifecycle(&self) -> Option<Arc<ParentLifecycle>> { + self.lifecycle + .lock() + .ok() + .and_then(|lifecycle| lifecycle.clone()) + } + + fn bind_parent_lifecycle(&self, lifecycle: Arc<ParentLifecycle>) -> Result<()> { + let mut current = self + .lifecycle + .lock() + .map_err(|_| store_err("parent lifecycle", "lock poisoned"))?; + if current + .as_ref() + .is_some_and(|existing| !Arc::ptr_eq(existing, &lifecycle)) + { + return Err(store_err( + "parent lifecycle", + "store is already bound to a different pool scope", + )); + } + *current = Some(lifecycle); + Ok(()) + } + + fn acquire_parent_gate( + &self, + parent_key: &str, + holder: &str, + ttl: Duration, + ) -> Result<Option<ParentGateClaim>> { + let now = SystemTime::now(); + let expires_at = now + .checked_add(ttl) + .ok_or_else(|| store_err("parent gate", "ttl out of range"))?; + let mut gates = self + .parent_gates + .lock() + .map_err(|_| store_err("parent gate", "lock poisoned"))?; + if let Some(existing) = gates.get(parent_key) + && existing.expires_at > now + { + return if existing.claim.holder == holder { + Ok(Some(existing.claim.clone())) + } else { + Ok(None) + }; + } + let mut epochs = self + .parent_gate_epochs + .lock() + .map_err(|_| store_err("parent gate epoch", "lock poisoned"))?; + let epoch = epochs + .get(parent_key) + .copied() + .unwrap_or(0) + .checked_add(1) + .ok_or_else(|| store_err("parent gate epoch", "epoch overflow"))?; + epochs.insert(parent_key.to_owned(), epoch); + let claim = ParentGateClaim { + holder: holder.to_owned(), + epoch, + }; + gates.insert( + parent_key.to_owned(), + MemoryParentGate { + claim: claim.clone(), + expires_at, + }, + ); + Ok(Some(claim)) + } + + fn renew_parent_gate( + &self, + parent_key: &str, + claim: &ParentGateClaim, + ttl: Duration, + ) -> Result<bool> { + let now = SystemTime::now(); + let expires_at = now + .checked_add(ttl) + .ok_or_else(|| store_err("parent gate", "ttl out of range"))?; + let mut gates = self + .parent_gates + .lock() + .map_err(|_| store_err("parent gate", "lock poisoned"))?; + let Some(gate) = gates.get_mut(parent_key) else { + return Ok(false); + }; + if gate.claim != *claim || gate.expires_at <= now { + return Ok(false); + } + gate.expires_at = expires_at; + Ok(true) + } + + fn release_parent_gate(&self, parent_key: &str, claim: &ParentGateClaim) -> Result<bool> { + let mut gates = self + .parent_gates + .lock() + .map_err(|_| store_err("parent gate", "lock poisoned"))?; + if gates + .get(parent_key) + .is_some_and(|gate| gate.claim == *claim && gate.expires_at > SystemTime::now()) + { + gates.remove(parent_key); + Ok(true) + } else { + Ok(false) + } + } + fn get(&self, task_id: &str) -> Option<CompletionRecord> { self.inner.lock().ok()?.get(task_id).cloned() } fn put(&self, record: &CompletionRecord) -> Result<()> { - self.lock()?.insert(record.task_id.clone(), record.clone()); + let mut records = self.lock()?; + let mut record = record.clone(); + if record.parent_epoch == 0 && !record.parent_key.is_empty() { + let marker_id = format!("\u{1}parent-cancelled:{}", record.parent_key); + if let Some(marker) = records + .get(&marker_id) + .filter(|m| m.state == CompletionState::Delivered) + { + record.parent_epoch = marker.parent_epoch; + } + } + records.insert(record.task_id.clone(), record); Ok(()) } + fn put_fenced( + &self, + record: &CompletionRecord, + parent_key: &str, + gate: &ParentGateClaim, + ) -> Result<()> { + self.with_current_gate(parent_key, gate, |records| { + let mut record = record.clone(); + if record.parent_epoch == 0 && !record.parent_key.is_empty() { + let marker_id = format!("\u{1}parent-cancelled:{}", record.parent_key); + if let Some(marker) = records + .get(&marker_id) + .filter(|m| m.state == CompletionState::Delivered) + { + record.parent_epoch = marker.parent_epoch; + } + } + records.insert(record.task_id.clone(), record); + Ok(()) + }) + } + + fn insert_if_absent(&self, record: &CompletionRecord, gate: &ParentGateClaim) -> Result<bool> { + self.with_current_gate(&record.parent_key, gate, |records| { + if records.contains_key(&record.task_id) { + return Ok(false); + } + records.insert(record.task_id.clone(), record.clone()); + Ok(true) + }) + } + + fn claim_pending(&self, request: PendingClaimRequest) -> Result<Vec<CompletionRecord>> { + let PendingClaimRequest { + parent_key, + parent_epoch, + gate, + mode, + max, + token, + visible_after, + } = request; + self.with_current_gate(&parent_key, &gate, |records| { + let now = SystemTime::now(); + let mut candidates: Vec<_> = records + .values() + .filter(|record| { + record.parent_key == parent_key + && record.parent_epoch == parent_epoch + && record.state == CompletionState::Pending + && mode.accepts(record.notify_mode) + && record + .claim + .as_ref() + .is_none_or(|claim| claim.visible_after <= now) + }) + .cloned() + .collect(); + candidates.sort_by(|a, b| { + a.finished_at + .cmp(&b.finished_at) + .then_with(|| a.task_id.cmp(&b.task_id)) + }); + candidates.truncate(max); + let mut claimed = Vec::with_capacity(candidates.len()); + for candidate in candidates { + let Some(record) = records.get_mut(&candidate.task_id) else { + continue; + }; + record.attempts = record.attempts.saturating_add(1); + record.updated_at = now; + record.claim = Some(CompletionClaim { + token: format!("{token}:{}", record.task_id), + visible_after, + }); + claimed.push(record.clone()); + } + Ok(claimed) + }) + } + + fn claim_task(&self, request: TaskClaimRequest) -> Result<Option<CompletionRecord>> { + let TaskClaimRequest { + task_id, + parent_key, + parent_epoch, + gate, + mode, + token, + visible_after, + } = request; + self.with_current_gate(&parent_key, &gate, |records| { + let now = SystemTime::now(); + let Some(record) = records.get_mut(&task_id) else { + return Ok(None); + }; + if record.parent_key != parent_key + || record.parent_epoch != parent_epoch + || record.state != CompletionState::Pending + || !mode.accepts(record.notify_mode) + || record + .claim + .as_ref() + .is_some_and(|claim| claim.visible_after > now) + { + return Ok(None); + } + record.attempts = record.attempts.saturating_add(1); + record.updated_at = now; + record.claim = Some(CompletionClaim { + token: format!("{token}:{task_id}"), + visible_after, + }); + Ok(Some(record.clone())) + }) + } + + fn resolve_claim(&self, request: ResolveClaimRequest) -> Result<Option<CompletionRecord>> { + self.with_current_gate(&request.parent_key, &request.gate, |records| { + let Some(record) = records.get_mut(&request.task_id) else { + return Ok(None); + }; + if record.parent_key != request.parent_key + || record.parent_epoch != request.parent_epoch + || record.state != CompletionState::Pending + || record.claim.as_ref().is_none_or(|claim| { + claim.token != request.token || claim.visible_after <= SystemTime::now() + }) + { + return Ok(None); + } + record.updated_at = SystemTime::now(); + record.claim = None; + match request.resolution { + ClaimResolution::Delivered => record.state = CompletionState::Delivered, + ClaimResolution::Failed { max_attempts } if record.attempts >= max_attempts => { + record.state = CompletionState::GaveUp; + } + ClaimResolution::Failed { .. } | ClaimResolution::Release => {} + } + Ok(Some(record.clone())) + }) + } + fn list(&self, parent_key: Option<&str>) -> Vec<CompletionRecord> { let Ok(guard) = self.inner.lock() else { return Vec::new(); @@ -94,7 +627,9 @@ impl CompletionStore for InMemoryCompletionStore { let now = SystemTime::now(); let mut guard = self.lock()?; let before = guard.len(); - guard.retain(|_, r| !expired(r, now, retain)); + guard.retain(|task_id, r| { + task_id.starts_with("\u{1}parent-cancelled:") || !expired(r, now, retain) + }); Ok(before - guard.len()) } } diff --git a/crates/tinyagents-tasks/src/completions/types.rs b/crates/tinyagents-tasks/src/completions/types.rs index d0b348ddc..a1a1ec952 100644 --- a/crates/tinyagents-tasks/src/completions/types.rs +++ b/crates/tinyagents-tasks/src/completions/types.rs @@ -77,6 +77,9 @@ pub enum CompletionState { Delivered, /// The parent collected the child explicitly, so it is not pushed from then on. Tombstoned, + /// The parent thread was deleted; this lifecycle marker is retained so + /// late child completions stay suppressed across compaction and restart. + ParentDeleted, /// Too many failed attempts. The host applies its own give-up policy. GaveUp, } @@ -88,11 +91,35 @@ impl CompletionState { Self::Pending => "pending", Self::Delivered => "delivered", Self::Tombstoned => "tombstoned", + Self::ParentDeleted => "parent_deleted", Self::GaveUp => "gave_up", } } } +/// A visibility claim held by one delivery attempt. +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct CompletionClaim { + /// Unique token for the claim; reclaiming a record replaces it. + pub token: String, + /// A later router may reclaim the record after this instant. + pub visible_after: SystemTime, +} + +/// Capability carried by the consumer that owns one specific delivery attempt. +/// A task id alone never authorizes settling a completion. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct CompletionReceipt { + /// The completion this attempt owns. + pub task_id: String, + /// Parent thread/session whose admission tenure accepted this delivery. + pub parent_key: String, + /// Durable parent admission epoch carried through host turn admission. + pub parent_epoch: u64, + /// The opaque token for this attempt. + pub token: String, +} + /// A neutral reference to a host-owned artifact holding the full output. #[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)] pub struct CompletionArtifact { @@ -141,6 +168,10 @@ pub struct CompletionRecord { /// restart. A tombstone written before the completion exists carries an /// empty key. pub parent_key: String, + /// Fencing epoch of the parent admission lease that accepted this record. + /// Older records remain readable as epoch zero. + #[serde(default)] + pub parent_epoch: u64, /// The child's agent id. pub agent_id: String, /// Optional human label. @@ -148,10 +179,7 @@ pub struct CompletionRecord { pub label: Option<String>, /// How the child ended. pub status: CompletionStatus, - /// The child's output. Emptied once the record settles (`Delivered`, - /// `Tombstoned` or `GaveUp`): the store then keeps only the bookkeeping - /// that dedupe and tombstones need, so a withdrawn or delivered result does - /// not outlive its parent's interest in it. + /// The child's output. pub result: CompletionResult, /// When the child finished. pub finished_at: SystemTime, @@ -165,6 +193,9 @@ pub struct CompletionRecord { /// How the parent wants to hear about it. #[serde(default)] pub notify_mode: NotifyMode, + /// Current visibility lease, if this record is claimed for delivery. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub claim: Option<CompletionClaim>, } impl CompletionRecord { @@ -180,6 +211,7 @@ impl CompletionRecord { Self { task_id: task_id.into(), parent_key: parent_key.into(), + parent_epoch: 0, agent_id: agent_id.into(), label: None, status, @@ -189,6 +221,7 @@ impl CompletionRecord { attempts: 0, state: CompletionState::Pending, notify_mode: NotifyMode::default(), + claim: None, } } @@ -210,19 +243,14 @@ impl CompletionRecord { self } - /// Moves the record to a terminal `state` now, dropping its result. The - /// task id, parent key, status, attempts and timestamps stay, which is all - /// dedupe, tombstones and compaction read; the payload has no reader once - /// the record is settled. - pub(crate) fn settle(&mut self, state: CompletionState) { - debug_assert_ne!( - state, - CompletionState::Pending, - "settle needs a terminal state" - ); - self.state = state; - self.updated_at = SystemTime::now(); - self.result = CompletionResult::default(); + /// Returns the receipt for this delivery attempt, if this record is claimed. + pub fn receipt(&self) -> Option<CompletionReceipt> { + self.claim.as_ref().map(|claim| CompletionReceipt { + task_id: self.task_id.clone(), + parent_key: self.parent_key.clone(), + parent_epoch: self.parent_epoch, + token: claim.token.clone(), + }) } /// A tombstone for a task whose completion has not been recorded yet. diff --git a/crates/tinyagents-tasks/src/conditional_sqlite.rs b/crates/tinyagents-tasks/src/conditional_sqlite.rs new file mode 100644 index 000000000..e874c3584 --- /dev/null +++ b/crates/tinyagents-tasks/src/conditional_sqlite.rs @@ -0,0 +1,78 @@ +//! SQLite-backed conditional storage for lease-fencing regressions. + +use std::sync::Arc; + +use tempfile::TempDir; +use tinyagents_storage_kit::{ + DocumentStore, Scope, ScopedStorage, Storage, StorageBackend, StorageError, WriteCondition, +}; +use tinystoragedrivers_connect::sqlite::SqliteConnection; +use tinystoragedrivers_document::sqlite::SqliteDocuments; + +#[derive(Clone)] +pub(crate) struct MemoryStorage { + _directory: Arc<TempDir>, + connection: SqliteConnection, +} + +impl MemoryStorage { + pub(crate) fn new() -> Self { + Self::with_clock(Arc::new(|| { + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap_or_default() + .as_millis() as u64 + })) + } + + pub(crate) fn with_clock(clock: Arc<dyn Fn() -> u64 + Send + Sync>) -> Self { + let directory = Arc::new(tempfile::tempdir().expect("temporary storage directory")); + let connection = + SqliteConnection::open_with_clock(directory.path().join("storage.db"), clock) + .expect("temporary sqlite storage opens"); + Self { + _directory: directory, + connection, + } + } + + pub(crate) fn open(path: impl AsRef<std::path::Path>) -> Self { + Self { + _directory: Arc::new(tempfile::tempdir().expect("temporary storage directory")), + connection: SqliteConnection::open(path).expect("sqlite storage opens"), + } + } +} + +impl StorageBackend for MemoryStorage { + fn for_scope(&self, scope: &Scope) -> Result<ScopedStorage, StorageError> { + let docs = Arc::new(SqliteDocuments::new(self.connection.clone(), scope.clone())) + as Arc<dyn DocumentStore>; + Ok(ScopedStorage::new(storage_with(docs), scope.clone())) + } + + fn for_scope_conditioned( + &self, + scope: &Scope, + condition: &WriteCondition, + ) -> Result<ScopedStorage, StorageError> { + let docs = Arc::new(SqliteDocuments::conditioned( + self.connection.clone(), + scope.clone(), + condition.clone(), + )) as Arc<dyn DocumentStore>; + Ok(ScopedStorage::new(storage_with(docs), scope.clone())) + } +} + +fn storage_with(documents: Arc<dyn DocumentStore>) -> Storage { + let unsupported = Storage::unsupported(); + Storage::new( + documents, + unsupported.streams().clone(), + unsupported.kv().clone(), + unsupported.blobs().clone(), + unsupported.queues().clone(), + unsupported.secrets().clone(), + ) +} diff --git a/crates/tinyagents-tasks/src/driver_store_tests.rs b/crates/tinyagents-tasks/src/driver_store_tests.rs index 61975fd40..776b46915 100644 --- a/crates/tinyagents-tasks/src/driver_store_tests.rs +++ b/crates/tinyagents-tasks/src/driver_store_tests.rs @@ -1,10 +1,10 @@ use std::sync::Arc; use super::*; +use crate::conditional_sqlite::MemoryStorage; use crate::store::TaskStore; use crate::types::{OrchestrationTaskKind, OrchestrationTaskStatus}; -use tinyagents_storage_kit::{Scope, StorageBackend}; -use tinystoragedrivers_testkit::MemoryBackend; +use tinyagents_storage_kit::{DocumentStore, Scope, StorageBackend}; fn docs_of(storage: &dyn StorageBackend, scope: &str) -> Arc<dyn DocumentStore> { Arc::clone( @@ -25,7 +25,7 @@ fn spec(id: &str) -> OrchestrationTaskSpec { } fn store() -> DriverTaskStore { - DriverTaskStore::new(docs_of(&MemoryBackend::new(), "local")).unwrap() + DriverTaskStore::new(docs_of(&MemoryStorage::new(), "local")).unwrap() } #[test] @@ -98,10 +98,12 @@ fn an_already_requested_cancellation_is_idempotent() { } #[test] -fn state_and_history_survive_a_driver_reopen_on_the_same_backend() { - let storage = MemoryBackend::new(); +fn state_and_history_survive_a_reopened_sqlite_file() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("tasks.sqlite"); let id = TaskId::new("task-a"); { + let storage = MemoryStorage::open(&path); let store = DriverTaskStore::new(docs_of(&storage, "local")).unwrap(); store.insert(spec("task-a")).unwrap(); store.mark_running(&id).unwrap(); @@ -110,10 +112,9 @@ fn state_and_history_survive_a_driver_reopen_on_the_same_backend() { .unwrap(); assert_eq!(store.history(&id).len(), 3); } + let storage = MemoryStorage::open(&path); let reopened = DriverTaskStore::new(docs_of(&storage, "local")).unwrap(); - let record = reopened - .get(&id) - .expect("task survives reopening a driver on the same backend"); + let record = reopened.get(&id).expect("task survives a restart"); assert_eq!(record.status, OrchestrationTaskStatus::Completed); assert_eq!(record.result.unwrap().text.as_deref(), Some("done")); assert_eq!(reopened.history(&id).len(), 3); @@ -204,7 +205,7 @@ fn list_filters_by_lineage_status_kind_and_time() { /// sees a terminal task rather than overwriting the winner. #[test] fn racing_settlements_have_one_winner() { - let storage = MemoryBackend::new(); + let storage = MemoryStorage::new(); let docs = docs_of(&storage, "local"); let first = DriverTaskStore::new(docs.clone()).unwrap(); first.insert(spec("race")).unwrap(); @@ -232,7 +233,7 @@ fn racing_settlements_have_one_winner() { #[test] fn tenants_never_share_tasks() { - let storage = MemoryBackend::new(); + let storage = MemoryStorage::new(); let alice = DriverTaskStore::new(docs_of(&storage, "alice")).unwrap(); let bob = DriverTaskStore::new(docs_of(&storage, "bob")).unwrap(); alice.insert(spec("shared")).unwrap(); @@ -269,23 +270,23 @@ async fn works_inside_a_current_thread_runtime() { assert_eq!(store.history(&id).len(), 2); } -#[tokio::test] -async fn the_opener_migrates_and_returns_a_durable_store() { - let backend = MemoryBackend::new(); - let scope = Scope::new("local").unwrap(); - let durable = crate::open_driver_task_store(backend.for_scope(&scope).unwrap()) - .await - .unwrap(); - durable.insert(spec("t1")).unwrap(); - let reopened = crate::open_driver_task_store(backend.for_scope(&scope).unwrap()) - .await - .unwrap(); +#[test] +fn the_opener_returns_a_durable_store_and_degrades_to_memory_when_unusable() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("opener.sqlite"); + { + let storage = MemoryStorage::open(&path); + let durable = crate::open_driver_task_store_or_memory(docs_of(&storage, "local")); + durable.insert(spec("t1")).unwrap(); + } + let storage = MemoryStorage::open(&path); + let reopened = crate::open_driver_task_store_or_memory(docs_of(&storage, "local")); assert_eq!(reopened.list(OrchestrationTaskFilter::default()).len(), 1); } #[tokio::test] async fn records_and_history_are_v2_documents_with_key_ids_and_ms_times() { - let storage = MemoryBackend::new(); + let storage = MemoryStorage::new(); let docs = docs_of(&storage, "local"); let store = DriverTaskStore::new(docs.clone()).unwrap(); let id = TaskId::new("a/b"); diff --git a/crates/tinyagents-tasks/src/lib.rs b/crates/tinyagents-tasks/src/lib.rs index b5a871a48..0fff4900d 100644 --- a/crates/tinyagents-tasks/src/lib.rs +++ b/crates/tinyagents-tasks/src/lib.rs @@ -3,7 +3,7 @@ //! This crate is the managed child-work surface: task model, `TaskStore` //! implementations, the detached task registry and restart reconciliation. //! It does not depend on the graph engine, so -//! `tinyagents-graph` and `tinyagents-orchestration` both build on it. +//! `tinyagents-orchestration` builds on it. //! //! Hosts can track child work by stable task id without exposing raw executor //! handles such as `tokio::JoinHandle`. @@ -20,15 +20,21 @@ mod store_registry; mod types; pub use completions::{ - CompletionArtifact, CompletionFormatter, CompletionRecord, CompletionResult, CompletionRouter, + ClaimResolution, CompletionAdmission, CompletionArtifact, CompletionClaim, CompletionClaimMode, + CompletionFormatter, CompletionReceipt, CompletionRecord, CompletionResult, CompletionRouter, CompletionState, CompletionStatus, CompletionStore, CompletionStoreReplaceOutcome, DEFAULT_MAX_ATTEMPTS, DriverCompletionStore, InMemoryCompletionStore, - NeutralCompletionFormatter, NotifyMode, RecordOutcome, TombstoneOutcome, + NeutralCompletionFormatter, NotifyMode, ParentGateClaim, ParentLifecycle, PendingClaimRequest, + RecordOutcome, ResolveClaimRequest, TaskClaimRequest, TombstoneOutcome, }; pub use completions::{ CompletionDelivery, CompletionDeliveryConfig, CompletionRouterPool, CompletionRouterPoolError, - CompletionRouterReplaceOutcome, DeliveryFence, DeliveryOutcome, DeliveryTarget, ParentBusyGate, - ParentBusyLease, ThreadDeliveryLeases, open_driver_completion_store, + CompletionRouterReplaceOutcome, CompletionStoreFallback, DeliveryFence, DeliveryOutcome, + DeliveryTarget, DurableThreadDeliveryLease, DurableThreadDeliveryLeases, ParentBusyGate, + ParentBusyLease, TenantThreadKey, ThreadDeliveryLeases, open_driver_completion_store, + open_driver_completion_store_or_memory, + open_driver_completion_store_or_memory_with_coordination_domain, + open_driver_completion_store_with_coordination_domain, }; pub use driver_store::DriverTaskStore; pub use reconcile::{ @@ -36,7 +42,7 @@ pub use reconcile::{ }; pub use runtime::{DetachedAbortEvent, DetachedTaskRegistry}; pub use steering_registry::SteeringRegistry; -pub use storage::{collections, open_storage, storage_schema}; +pub use storage::{open_storage, storage_schema}; pub use store::{InMemoryTaskStore, TaskStore}; pub use store_registry::{ TaskStoreRegistry, TaskStoreRegistryError, open_driver_task_store, @@ -44,6 +50,10 @@ pub use store_registry::{ }; pub use types::*; +#[cfg(test)] +#[path = "conditional_sqlite.rs"] +mod conditional_sqlite; + #[cfg(test)] #[path = "lib_tests.rs"] mod test; diff --git a/crates/tinyagents-tasks/src/lib_tests.rs b/crates/tinyagents-tasks/src/lib_tests.rs index 7b0e84224..8699ce730 100644 --- a/crates/tinyagents-tasks/src/lib_tests.rs +++ b/crates/tinyagents-tasks/src/lib_tests.rs @@ -104,6 +104,39 @@ async fn detached_registry_enforces_owner_and_waits_for_terminal_status() { join.abort(); } +#[tokio::test] +async fn trusted_terminal_pruning_keeps_live_tasks_and_releases_steering() { + let steering = SteeringRegistry::new(); + let registry = DetachedTaskRegistry::new(steering.clone(), 8, |status: &RuntimeStatus| { + matches!(status, RuntimeStatus::Completed(_)) + }); + let task_id = TaskId::new("trusted-terminal-prune"); + let (tx, rx, cancellation, join) = detached_handles(); + registry + .register_cooperative( + task_id.clone(), + "parent", + "metadata".to_owned(), + rx, + cancellation, + ) + .unwrap(); + steering.register(task_id.clone(), SteeringHandle::allow_all()); + + assert_eq!(registry.prune_terminal_trusted(&task_id).unwrap(), None); + assert!(registry.snapshot_trusted(&task_id).is_ok()); + assert!(steering.get(&task_id).is_some()); + + tx.send(RuntimeStatus::Completed("done".into())).unwrap(); + assert_eq!( + registry.prune_terminal_trusted(&task_id).unwrap(), + Some(RuntimeStatus::Completed("done".into())) + ); + assert!(registry.snapshot_trusted(&task_id).is_err()); + assert!(steering.get(&task_id).is_none()); + join.abort(); +} + #[tokio::test] async fn batch_wait_preserves_order_and_metadata_when_pruning_terminal_entries() { let registry = runtime_registry(); diff --git a/crates/tinyagents-tasks/src/runtime.rs b/crates/tinyagents-tasks/src/runtime.rs index e99c00183..8079b3483 100644 --- a/crates/tinyagents-tasks/src/runtime.rs +++ b/crates/tinyagents-tasks/src/runtime.rs @@ -53,6 +53,12 @@ pub struct DetachedAbortEvent<Metadata, Status> { type RegistryGuard<'a, Metadata, Status> = MutexGuard<'a, HashMap<TaskId, DetachedTaskEntry<Metadata, Status>>>; +enum RegistrationKind { + Spawned(AbortHandle), + Cooperative, + ManagedCooperative, +} + /// Ownership-aware registry for live detached task executors. /// /// The registry is generic over application metadata and status payloads. A @@ -125,7 +131,7 @@ where metadata, status, cancellation, - Some(abort), + RegistrationKind::Spawned(abort), ) } @@ -146,7 +152,30 @@ where metadata, status, cancellation, - None, + RegistrationKind::Cooperative, + ) + } + + /// Registers a cooperative task without the automatic soft-cap sweep. + /// + /// Use this only when an adapter coordinates terminal pruning with + /// sidecar state and protects entries that still have active waiters. The + /// adapter must enforce its own finite retention policy before registering. + pub fn register_cooperative_managed( + &self, + task_id: TaskId, + owner_id: impl Into<String>, + metadata: Metadata, + status: watch::Receiver<Status>, + cancellation: CancellationToken, + ) -> Result<()> { + self.register_entry( + task_id, + owner_id.into(), + metadata, + status, + cancellation, + RegistrationKind::ManagedCooperative, ) } @@ -157,9 +186,14 @@ where metadata: Metadata, status: watch::Receiver<Status>, cancellation: CancellationToken, - abort: Option<AbortHandle>, + registration: RegistrationKind, ) -> Result<()> { - if self.len().map_err(Self::tinyagents_error)? >= self.soft_cap { + let (abort, sweep_at_cap) = match registration { + RegistrationKind::Spawned(abort) => (Some(abort), true), + RegistrationKind::Cooperative => (None, true), + RegistrationKind::ManagedCooperative => (None, false), + }; + if sweep_at_cap && self.len().map_err(Self::tinyagents_error)? >= self.soft_cap { self.sweep_terminal().map_err(Self::tinyagents_error)?; } let mut guard = self.lock().map_err(Self::tinyagents_error)?; @@ -565,6 +599,33 @@ where status.has_changed().is_err() || (self.is_terminal)(&status.borrow()) } + /// Removes one terminal entry for trusted adapter-side lifecycle cleanup. + /// + /// Returns the terminal snapshot when this call removed the entry, and + /// `None` when it is still live or has already been removed. Ownership is + /// intentionally not checked; callers must establish trust before using + /// this process-local supervision API. Steering is deregistered with the + /// runtime entry, as it is for every other removal path. + pub fn prune_terminal_trusted( + &self, + task_id: &TaskId, + ) -> std::result::Result<Option<Status>, DetachedTaskRegistryError> { + let status = { + let mut guard = self.lock()?; + let Some(entry) = guard.get(task_id) else { + return Ok(None); + }; + let status = entry.status.borrow().clone(); + if !self.is_settled(&entry.status) { + return Ok(None); + } + guard.remove(task_id); + status + }; + self.steering.deregister(task_id); + Ok(Some(status)) + } + fn ensure_live_owned( &self, task_id: &TaskId, @@ -618,7 +679,7 @@ where } if let Some(abort) = &entry.abort { let status = entry.status.borrow().clone(); - if !(self.is_terminal)(&status) { + if !self.is_settled(&entry.status) { (self.on_abort)(DetachedAbortEvent { task_id: task_id.clone(), owner_id: entry.owner_id.clone(), diff --git a/crates/tinyagents-tasks/src/storage/dto.rs b/crates/tinyagents-tasks/src/storage/dto.rs index b524fb0af..cb6d476f9 100644 --- a/crates/tinyagents-tasks/src/storage/dto.rs +++ b/crates/tinyagents-tasks/src/storage/dto.rs @@ -104,6 +104,8 @@ impl TaskDoc { pub(crate) struct CompletionDoc { pub(crate) task_id: String, pub(crate) parent_key: String, + #[serde(default)] + pub(crate) parent_epoch: u64, pub(crate) agent_id: String, #[serde(default, skip_serializing_if = "Option::is_none")] pub(crate) label: Option<String>, @@ -115,6 +117,8 @@ pub(crate) struct CompletionDoc { pub(crate) state: CompletionState, #[serde(default)] pub(crate) notify_mode: NotifyMode, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub(crate) claim: Option<crate::completions::CompletionClaim>, } impl CompletionDoc { @@ -122,6 +126,7 @@ impl CompletionDoc { Self { task_id: record.task_id.clone(), parent_key: record.parent_key.clone(), + parent_epoch: record.parent_epoch, agent_id: record.agent_id.clone(), label: record.label.clone(), status: record.status, @@ -130,6 +135,7 @@ impl CompletionDoc { attempts: record.attempts, state: record.state, notify_mode: record.notify_mode, + claim: record.claim.clone(), } } @@ -139,6 +145,7 @@ impl CompletionDoc { CompletionRecord { task_id: self.task_id, parent_key: self.parent_key, + parent_epoch: self.parent_epoch, agent_id: self.agent_id, label: self.label, status: self.status, @@ -148,6 +155,7 @@ impl CompletionDoc { attempts: self.attempts, state: self.state, notify_mode: self.notify_mode, + claim: self.claim, } } } diff --git a/docs/audit.md b/docs/audit.md index d51f3c80e..fcca1e97c 100644 --- a/docs/audit.md +++ b/docs/audit.md @@ -81,51 +81,6 @@ candidate filter: select a capable hinted model over an incapable default after middleware adds required capabilities. -### Resolved: runtime `Command::goto` targets are validated before persistence - -Runtime command routing now fails before a bad target can become the next active -set or be written into a checkpoint: - -- `crates/tinyagents-graph/src/compiled/mod.rs`: command targets are validated in `route()` before - scheduler activation. `END` is allowed, `START` is rejected, and every other - target must name a compiled node. -- `crates/tinyagents-graph/src/compiled/test.rs`: unit coverage rejects unknown targets, rejects - `START`, and proves invalid command routes fail before checkpoint persistence. -- `tests/graph_durable.rs`: integration coverage proves a threaded durable run - with an invalid command target writes no poisoned checkpoint. -- `docs/modules/graph/routing.md`: documents the runtime validation contract. - -### Resolved: embedded subgraph runs inherit the parent thread id - -Subgraph node adapters now propagate the parent thread id into embedded child -runs when the parent run is threaded: - -- `crates/tinyagents-graph/src/subgraph/mod.rs`: shared-state and adapter subgraph nodes call the - child with `run_with_thread` when `NodeContext::thread_id` is present, - preserving the existing unthreaded `run` behavior otherwise. -- `crates/tinyagents-graph/src/subgraph/test.rs`: unit coverage proves an embedded child with a - checkpointer persists under the parent thread and embedding-node namespace. -- `tests/graph_durable.rs`: integration coverage proves the public durable graph - surface writes both parent and child checkpoints for a threaded subgraph run. -- `docs/modules/graph/subgraphs.md`: documents the inherited thread and - namespace persistence contract. - -### Resolved: interrupts fail closed without resumable durability - -The executor now returns a resume error instead of an interrupted execution when -a node emits an interrupt without the durability required to resume it: - -- `crates/tinyagents-graph/src/compiled/mod.rs`: interrupt handling requires both a configured - checkpointer and a thread id before persisting the interrupt checkpoint and - returning `ExecutionStatus::Interrupted`. -- `crates/tinyagents-graph/src/compiled/test.rs`: unit coverage proves interrupts without a - checkpointer and interrupts without a thread id return errors rather than - unresumable paused executions. -- `tests/graph_durable.rs`: integration coverage proves the public run surface - rejects an interrupt emitted without a thread id. -- `docs/modules/graph/interrupts.md`: documents that interrupted results are - only returned after the checkpoint needed for resume is persisted. - ### Resolved: event sink listeners are notified outside the sink lock The harness event sink no longer holds its mutex while invoking listener diff --git a/docs/feature-e2e-test-matrix.md b/docs/feature-e2e-test-matrix.md index d2d461589..ad28b86cd 100644 --- a/docs/feature-e2e-test-matrix.md +++ b/docs/feature-e2e-test-matrix.md @@ -11,12 +11,10 @@ tests remain valuable but do not close an E2E gap. | Surface | Existing integration coverage | Status | | --- | --- | --- | -| Graph construction, routing, fan-out, checkpoints, interrupts, export, streaming, subgraphs, parallelism | `graph_durable.rs`, `e2e_complex_graph.rs`, `e2e_durable_interrupt.rs`, `e2e_graph_export.rs`, `feature_graph_{routing,fanout,parallel,streaming}.rs` | Covered | -| Graph goals, todo board, task dispatch, subagents, observability | integration tests plus `crates/tinyagents-orchestration/tests/e2e_graph_subagent_node.rs` | Covered; delegation durability remains unit-only | | Harness loop, tools, structured output, middleware, retry, cache, streams, hosted execution | `harness_agent_loop.rs`, `feature_harness_*`, `e2e_{middleware,streaming_cancel,tool_policy}.rs`, `wave2_*` | Covered; optional features remain gaps | | Registry catalog, capability binding, diagnostics, observability | `e2e_registry_binding.rs`, `e2e_registry_observability_contracts.rs`, `feature_registry_{catalog,diagnostics}.rs` | Covered; router-to-runtime fallback is a gap | | Session records, retention, search, lifecycle, transcripts | `e2e_session_lifecycle.rs`, `feature_session_retention.rs`, `feature_session_transcript.rs`, persistence contracts | Covered; run-ledger workflows remain gaps | -| Subagent orchestration | `crates/tinyagents-orchestration/tests/e2e_*subagent*.rs` and `hosted_subagents.rs` | Direct invocation, tools, reuse, limits, graph nodes, and hosted orchestration covered | +| Subagent orchestration | `crates/tinyagents-orchestration/tests/e2e_*subagent*.rs` and `hosted_subagents.rs` | Direct invocation, tools, reuse, limits, and hosted orchestration covered | ## Prioritized execution backlog @@ -33,8 +31,6 @@ tests remain valuable but do not close an E2E gap. | Feature | Test to add | | --- | --- | -| Durable reviewed delegation | `e2e_graph_delegation.rs` with `DriverCheckpointer`, approval interrupt/resume, revision, and denial. | -| Typed graph channels/barriers | `e2e_graph_channels.rs` for aggregate, barrier, ephemeral, conflict, events, and checkpoints. | | Workspace dispatch claims | A mixed tool-side-effect batch proving permitted parallelism, writer serialization, and unsafe-path rejection. | | Goal tools in a model loop | A scripted model invoking registered goal tools and an over-budget preflight block. | | SQLite cache across harness restart | Reopen a disk cache with a new harness and prove a cache hit avoids model call and billing. | @@ -51,11 +47,10 @@ tests remain valuable but do not close an E2E gap. | Public stream iterator | Consume `AgentHarness::invoke_stream` through terminal and early-drop paths. | | Artifact, handoff, and workspace utilities | Temp-root artifact-offload, handoff, and workspace lifecycle API contract. | | Namespaced stores and control primitives | Namespace/TTL/batch/search, `RunQueue`, and no-progress/repeat tracking as direct feature contracts. | -| Extended language grammar | Execute materialized command/send/join/retry/checkpoint constructs only where graph runtime supports them. | ## Feature flags -`sqlite` is integration-tested through graph/session persistence. `tools` and +`sqlite` is integration-tested through session persistence. `tools` and `multimodal` are compiled and unit-tested in CI but need the P2 behavioral integration tests above. `tracing` is compiled under all-features; adding behavioral tracing tests is lower priority because instrumentation is opt-in. diff --git a/docs/host-adoption/481-host-adoption-map.md b/docs/host-adoption/481-host-adoption-map.md index 957a45128..1907f40cf 100644 --- a/docs/host-adoption/481-host-adoption-map.md +++ b/docs/host-adoption/481-host-adoption-map.md @@ -7,7 +7,7 @@ redaction rules. | OpenHuman copy | TinyAgents API | OpenHuman work after the gitlink is merged | | --- | --- | --- | | `crates/openhuman-core/src/agent/tinyagents/use_skill_dispatch.rs` | `tinyagents_harness::tool::packs::{UseSkillDispatch, PackDispatchSelector, PackSelection}`. The dispatch pins one `PackSelection` at `bind_call`; the host selector builds its existing typed dispatch from `selection.registry()` and `selection.tool_name()`. That dispatch receives the original State, CallId, options, and live `RunContext`. | Delete `use_skill_dispatch.rs` and its `UseSkillDispatch`/`BoundUseSkillDispatch` implementation. Keep `tool_registration::typed_dispatch_for`, `shared_tool_adapter`, the catalog, registry binding, and the call site that registers `UseSkillDispatch`. The selector closure is host wiring, not generic dispatch logic. | -| `crates/openhuman-core/src/agent/tools/plan_exit.rs` | `tinyagents_harness::tool::plan_exit::{PlanExitTool, PlanExitModeLookup, PLAN_EXIT_MARKER}`. `PlanExitTool::new(Arc<PlanExitModeLookup>)` validates `plan`, emits the stable marker, and sets a returned `RunModeHandle` to `Build`. | Delete the schema, validation, result formatting, and marker implementation. Keep the callback that resolves the thread's live handle and checks tenant context plus a parked plan review before returning `Some(handle)`. Keep registration and mode middleware wiring. | +| `crates/openhuman-core/src/agent/tools/plan_exit.rs` | `tinyagents_harness_tools::PlanExitTool` validates `plan`, emits the stable marker, refuses a parked review, and switches to Build through the injected `RunModeStore` (which publishes the mode change). Its host seam resolves the opaque scope/thread identity. | Delete the schema, validation, result formatting, and marker implementation. Keep the host identity resolver, store wiring, and mode middleware wiring. `PlanExitTool::new` receives the host resolver, run-mode store, and optional plan-review store. | | `crates/openhuman-core/src/agent/harness/artifact_offload/mod.rs` | `ArtifactOffload::{validate_materialization_target, materialize_to_fs}` and `artifacts::offload_oversized_result_to_fs`. The filesystem entry point checks the actual target through the installed `ArtifactPathPolicy` before the blob write, then materializes the stored redacted body with a same-directory atomic rename. Unix uses directory file descriptors, `mkdirat`/`openat`, and no-follow traversal; other platforms use `cap_std` directory capabilities for relative directory creation and rename. | Delete the host's pre-store target check and private filesystem `materialize` helper. Route oversized results through `offload_oversized_result_to_fs`. Keep `WorkspaceGuard`, `SanitizingRedactor`, scoped blob-store selection, root selection, prompt contract/tool names, and host logging/soft-failure integration. | `UseSkillDispatch` currently requires `State: 'static` and `Ctx: 'static`, the diff --git a/docs/modules/graph/README.md b/docs/modules/graph/README.md deleted file mode 100644 index 3667b9f19..000000000 --- a/docs/modules/graph/README.md +++ /dev/null @@ -1,147 +0,0 @@ -# Graph Module Specification - -The graph module is TinyAgents' workflow runtime. It owns topology, state -transition semantics, routing, execution history, checkpointing, interrupts, -streaming, parallel execution, subgraph invocation, sub-agent nodes, recursive -calls, and graph-level observability. - -The graph runtime does not know or care where a graph came from. - -The implementation lives in `crates/tinyagents-graph/src/` and includes the -typed builder, compiled executor, channels, checkpointing, interrupts, -streaming, parallel execution, subgraphs, sub-agent nodes, and observability. -The specifications below describe the contract those modules implement. - -## Source Inspiration - -Primary references: - -- LangGraph repository: <https://github.com/langchain-ai/langgraph> -- LangGraph graph API: - <https://docs.langchain.com/oss/python/langgraph/graph-api> -- LangGraph persistence: - <https://docs.langchain.com/oss/python/langgraph/persistence> -- LangGraph durable execution: - <https://docs.langchain.com/oss/python/langgraph/durable-execution> -- LangGraph checkpointers: - <https://docs.langchain.com/oss/python/langgraph/checkpointers> -- LangGraph interrupts: - <https://docs.langchain.com/oss/python/langgraph/interrupts> -- LangGraph streaming: - <https://docs.langchain.com/oss/python/langgraph/streaming> -- LangGraph event streaming: - <https://docs.langchain.com/oss/python/langgraph/event-streaming> -- LangGraph subgraphs: - <https://docs.langchain.com/oss/python/langgraph/use-subgraphs> -- LangGraph fault tolerance: - <https://docs.langchain.com/oss/python/langgraph/fault-tolerance> -- `rust-langgraph` docs: - <https://docs.rs/rust-langgraph/latest/rust_langgraph/> -- OpenHuman PR #4261: - <https://github.com/tinyhumansai/openhuman/pull/4261> - -Useful upstream code references: - -- `libs/langgraph/langgraph/graph/state.py`: `StateGraph`, channels, - conditional edges, compile-time validation, node defaults, and subgraph - attachment. -- `libs/langgraph/langgraph/pregel/main.py`: executable graph runtime, - superstep loop, streaming, state update APIs, durability, recursion errors, - and subgraph stream propagation. -- `libs/langgraph/langgraph/types.py`: `Command`, `Send`, `Interrupt`, - `StateSnapshot`, stream part types, retry/cache/timeout policy, and durability - modes. -- `libs/langgraph/langgraph/channels/`: reducer/channel implementations such as - last-value, binary operator aggregate, topic, ephemeral, named barrier, and - delta channels. -- `libs/checkpoint/langgraph/checkpoint/base/__init__.py`: checkpoint tuple, - pending writes, thread operations, copy/prune semantics, and delta-channel - history. - -Rust-specific precedent: - -- `rust-langgraph` is useful as a small Rust-native reference for stateful LLM - workflows with nodes, conditional edges, checkpoints, streaming, optional - model adapters, and ReAct/tool helpers. -- OpenHuman PR #4261 is useful as a product reference for keeping a generic - graph engine independent from the agent harness while exposing checkpoints, - HITL, observability events, and RPC run control. -- TinyAgents should preserve the ergonomic Rust builder surface from these - precedents, but the target graph runtime must also be rich enough to be - inspected by UIs, driven from a host orchestrator, and tested with - deterministic state/channel snapshots. - -## Design Stance - -The graph module is the stable execution contract below all graph authoring -surfaces. A graph may be built by Rust code, assembled by a host orchestration -session, or restored from durable state. Once compiled, those origins must converge on the same immutable -`CompiledGraph` behavior: - -- topology is validated before execution -- node capabilities are resolved through registries -- state writes pass through channel reducers -- commands and `Send` packets are explicit runtime values -- checkpoints and pending writes are owned by the graph runtime -- nested graph and sub-agent calls preserve run hierarchy -- events contain enough source/origin metadata to explain generated graphs - -The graph runtime must not accept unchecked model-generated topology or -executable code. - -## Responsibilities - -- Build named node graphs with direct, conditional, barrier, and command-based - routing. -- Validate topology before execution and freeze it into an immutable compiled - graph. -- Execute async nodes under graph and node policies. -- Apply partial state updates through typed channel/reducer policies. -- Execute multiple active nodes in a superstep. -- Support dynamic fanout through `Send`-style packets. -- Enforce recursion, step, timeout, concurrency, retry, and cache policy. -- Persist checkpoints, pending writes, task outcomes, and interrupt state at - execution boundaries. -- Support human-in-the-loop interrupts and resumable commands. -- Support manual state inspection, state history, state update, forks, and time - travel when checkpointing is enabled. -- Support targeted steering of parent orchestrator runs, child sub-agent runs, - graph tasks, and pending interrupts through policy-checked commands. -- Emit typed graph, task, checkpoint, interrupt, and streamed-output events. -- Represent subgraphs as executable nodes with namespaced checkpoints and nested - streams. -- Represent harness sub-agents as graph nodes with child run identity. -- Export graph structure for visualization, tests, and generated docs. - -## Non-Responsibilities - -- It does not own chat model provider logic. -- It does not own tool schema validation or tool dispatch. -- It does not implement prompt templating. -- It does not manage long-term application memory. -- It does not own model/tool usage accounting, though it must forward and roll - up child events from harness nodes. - -## Feature Specifications - -- [Package and core types](package.md) -- [Builder and compile contract](builder.md) -- [Node model](nodes.md) -- [State, channels, and updates](state-channels.md) -- [Edges, routing, commands, and sends](routing.md) -- [Execution model and parallelization](execution.md) -- [Parallel agents and context forking](parallel-agents-forking.md) -- [Checkpointing, durability, state inspection, and time travel](checkpointing.md) -- [Interrupts and resume](interrupts.md) -- [Streaming and events](streaming.md) -- [Observability and tracing](observability.md) -- [Runtime context, node defaults, and policies](runtime-policy.md) -- [Error handling and fault tolerance](fault-tolerance.md) -- [Subgraphs](subgraphs.md) -- [Sub-agents, recursion, and depth tracking](subagents-recursion.md) -- [Memory and stores boundary](memory-boundary.md) -- [Visualization, introspection, and testkit](visualization-testkit.md) -- [Per-thread goal and graph-native continuation](goals.md) -- [Per-thread todo list](todos.md) -- [Per-thread scratchpad: the model's notes](scratchpad.md) -- [Implementation milestones](milestones.md) diff --git a/docs/modules/graph/builder.md b/docs/modules/graph/builder.md deleted file mode 100644 index a8897d44d..000000000 --- a/docs/modules/graph/builder.md +++ /dev/null @@ -1,169 +0,0 @@ -# Graph Builder And Compile Contract - -The builder supports: - -- `add_node` -- `add_node_shared` -- `add_sequence` -- `add_edge` -- `add_waiting_edge` -- `add_conditional_edges` -- `set_entry_point` -- `set_conditional_entry_point` -- `set_finish_point` -- `set_node_defaults` -- `set_defaults` -- `with_max_concurrency` -- `with_node_timeout` -- `interrupt_before` / `interrupt_after` / `mark_interrupt` -- `compile` - -`add_conditional_edges` accepts route labels and a router return value that are -any `impl ToString`, so a user-defined route enum that implements `Display` (or -the `Route` newtype) can label edges directly; plain `&str`/`String` labels keep -working unchanged. - -`add_sequence([a, b, c])` is convenience sugar for a chain of direct edges -(`add_edge(a, b).add_edge(b, c)`). `add_waiting_edge(from, to)` is a barrier edge: -`to` activates only once *all* of its registered predecessors have completed, -even across supersteps. - -`add_edge`/`add_waiting_edge` accumulate: calling `add_edge("a", "b")` then -`add_edge("a", "c")` schedules **both** `b` and `c` as static fan-out targets -of `a` (deduplicated — adding the same edge twice does not schedule the -target twice), matching the "one or more node names" routing contract. This -is a change from earlier versions, where a second `add_edge` call from the -same source silently overwrote the first. - -`add_conditional_edges_checked(from, router, all_labels)` is -`add_conditional_edges` plus an exhaustive `all_labels` list tied to the -router's own return type. `compile()` cross-checks every declared label -against the node's route table and rejects a mismatch with -`TinyAgentsError::MissingRoute` **at build time**, instead of only at run -time when the router happens to return the mistyped label. Plain -`add_conditional_edges` (no exhaustive label list) still only fails at run -time, since an opaque closure's possible outputs can't be enumerated ahead of -time. - -Graph defaults are settable in one call: - -```rust -pub struct GraphDefaults { - pub recursion_limit: Option<usize>, - pub parallel: Option<bool>, - pub max_concurrency: Option<usize>, - pub node_timeout: Option<Duration>, -} -``` - -`set_defaults(GraphDefaults { .. })` applies only the `Some` fields. -`with_max_concurrency(n)` bounds the number of node handlers in flight per -parallel superstep (the active set runs in chunks of at most `n`). -`with_node_timeout(d)` fails the run with `TinyAgentsError::Timeout` if any node -handler does not resolve within `d`. - -`CompiledGraph::with_run_deadline(d)` bounds the *whole run* by a wall-clock -`d`, checked at every super-step boundary: when the elapsed run time first -reaches `d` the run stops *between* super-steps with `TinyAgentsError::Timeout`, -leaving the last committed boundary checkpoint intact and resumable. Prefer this -over wrapping `run` in an external `tokio::time::timeout`, which aborts -mid-super-step and cannot leave a clean checkpoint. It bounds scheduling, not a -single in-flight node — pair it with `with_node_timeout` to also bound -individual handlers. - -## `add_node` vs `add_node_shared` (state cloning) - -```rust -// By value — clones out of the step's Arc<State> once per invocation. -.add_node("summarize", |state: State, ctx: NodeContext| async move { .. }) - -// Zero-clone — receives the step's Arc<State> directly. -.add_node_shared("summarize", |state: Arc<State>, ctx: NodeContext| async move { .. }) -``` - -Internally every node handler receives the step's committed state as an -`Arc<State>` (`docs/runtime-comparison/code-review-graph.md` M2): a -superstep clones `State` at most once, and every branch of a parallel -`Send` fan-out and every retry attempt of a task shares that `Arc` via a -cheap `Arc::clone`. `add_node` is the pre-existing, source-compatible entry -point — it wraps the closure so it still receives an owned `State`, cloning -out of the `Arc` on each call. `add_node_shared` skips that clone entirely -by handing the closure the `Arc<State>` directly; prefer it for a node with -a large state or one that only reads state. See -[nodes.md](nodes.md#actual-handler-signature-and-state-cloning-m2) for the -full rationale and the matching `NodeContext::send_arg` change. - -## Interrupt selectors - -```rust -GraphBuilder::<State, Update>::new() - .interrupt_before(["approve"]) - .interrupt_after(["plan"]) // requires Update: Serialize + DeserializeOwned - .mark_interrupt("review") // == interrupt_before(["review"]) -``` - -`interrupt_before(nodes)` pauses the run instead of invoking a listed node; -`interrupt_after(nodes)` lets the node run, then pauses *before* its result -is applied, persisting the result as a deferred write that resume replays -without re-running the handler. Both inject an `Interrupt` with payload -`{"phase": "before" | "after"}`, need a checkpointer and thread like any -interrupt, are validated at `compile()` (`MissingNode`), and set the -export's `NodeInfo::interrupt` marker. `mark_interrupt` used to set only that -marker; it is now an alias for `interrupt_before` — a node that already -pauses itself should not be listed (it would pause twice; annotate it with -`with_node_metadata` for the export instead). Full semantics, including how -acknowledgements survive repeated pauses, are in -[interrupts.md](interrupts.md#interrupt_before--interrupt_after-selectors). - -## `NodeContext::durable_task` - -```rust -.add_node("charge", |state, ctx: NodeContext| async move { - let receipt: Receipt = ctx - .durable_task("charge-card", async { payments.charge(&state).await }) - .await?; - if ctx.resume.is_none() { - return Ok(NodeResult::Interrupt(Interrupt::new("charge", json!(receipt)))); - } - Ok(NodeResult::Update(state.with_receipt(receipt))) -}) -``` - -A node is re-run from its start after an interrupt/resume, a failure/retry, -or an in-process node retry. `durable_task(key, fut)` runs `fut` at most -once per `(task_id, key)`: the first execution awaits it and records its -`Ok` output (`T: Serialize + DeserializeOwned`) as a -`PendingWrite::durable_task` memo in the task's checkpoint write ledger; a -re-run of the same task returns the stored value without polling `fut`. An -`Err` is returned unmemoised. Keys are independent — a handler that memoised -`"a"` then failed before `"b"` replays `"a"` and runs `"b"` fresh on retry. -Memos are scoped to one task in one thread (not a cross-run cache; see -`CompiledGraph::with_cached_node` for that) and are durable across process -restarts on a checkpointed thread; a graph without a checkpointer still -dedupes within one run. The executor pre-seeds each re-run's context from -the checkpoint and persists new memos at whichever boundary the task -stalls at (interrupt or failure); a task that completes drops its memos. - -## Not implemented: compile-time options - -An earlier draft of this page described a `CompileOptions` struct -(`checkpointer: CheckpointerChoice`, `cache`, `store`, `debug`, -`stream_transformers`) mirroring LangGraph's compile-time bag. It does not -exist: checkpointers, caches, sinks, and policies are attached to the -`CompiledGraph` with `with_*` methods after `compile()`, and the interrupt -selectors live on the builder as above. - -Validation rules: - -- graph must have at least one `START` path -- `START` cannot be an edge target -- `END` cannot be an edge source -- every edge source exists, except `START` -- every edge target exists, except `END` -- duplicate node ids are rejected -- duplicate branch names from a source are rejected -- conditional route targets are validated at compile time when known -- interrupt targets must exist -- waiting-edge sources and targets must exist -- command destinations used only for rendering are marked as such -- node additions after compile do not mutate an existing compiled graph diff --git a/docs/modules/graph/checkpointing.md b/docs/modules/graph/checkpointing.md deleted file mode 100644 index d807083eb..000000000 --- a/docs/modules/graph/checkpointing.md +++ /dev/null @@ -1,438 +0,0 @@ -# Graph Checkpointing, Durability, State Inspection, And Time Travel - -Checkpointing is graph runtime persistence. It is separate from harness memory -and long-term stores. - -LangGraph stores graph state through checkpointers. The checkpoint package -defines this as the graph persistence layer: it saves graph state at every -superstep and enables human-in-the-loop, memory between interactions, durable -execution, and replay. Its core persistence unit is a checkpoint tuple: the -checkpoint itself plus config, metadata, parent config, and pending writes. - -Checkpoint coordinates are thread-based: - -- `thread_id` identifies the checkpoint lineage for a conversation, workflow, or - tenant-isolated run series. -- `checkpoint_id` optionally selects a specific point in that thread, including - time-travel or replay from the middle of a thread. -- `checkpoint_ns` scopes nested graph/subgraph state so subgraphs can share a - checkpointer without colliding with parent checkpoints. - -```rust -#[async_trait] -pub trait Checkpointer: Send + Sync { - async fn get_tuple( - &self, - query: CheckpointQuery, - ) -> Result<Option<CheckpointTuple>>; - - async fn list( - &self, - query: CheckpointListQuery, - ) -> Result<Vec<CheckpointTuple>>; - - async fn put( - &self, - checkpoint: Checkpoint, - metadata: CheckpointMetadata, - new_versions: ChannelVersions, - ) -> Result<CheckpointConfig>; - - async fn put_writes( - &self, - config: CheckpointConfig, - writes: Vec<PendingWrite>, - task_id: TaskId, - task_path: TaskPath, - ) -> Result<()>; -} -``` - -Checkpoint tuple: - -```rust -pub struct CheckpointTuple { - pub config: CheckpointConfig, - pub checkpoint: Checkpoint, - pub metadata: CheckpointMetadata, - pub parent_config: Option<CheckpointConfig>, - pub pending_writes: Vec<PendingWrite>, -} -``` - -Checkpoint fields (verified against `crates/tinyagents-graph/src/checkpoint/types.rs`, -`Checkpoint<State>` and `CheckpointMetadata`): - -Implemented today: - -- checkpoint id -- thread id -- checkpoint namespace -- run id -- `version: u32` — the on-disk record shape. `2` (`CHECKPOINT_FORMAT_VERSION`) - for every record this crate writes; see "Checkpoint format v2" below. -- `created_at: u64` — wall-clock write time in Unix millis - (`tinyagents_harness::ids::now_ms()`); `0` on a v1 record, which never - carried a timestamp. -- committed state (`state: State`, not a per-channel value map) -- `tasks: Vec<PendingActivation>` — the single source of truth for what runs - on resume, preserving each pending node's `Send` argument and task id. -- `completed: Vec<CompletedTask>` — the single source of truth for what - completed in the step that produced this checkpoint, each entry carrying - its own task id, node, and explicit `Command::goto` routing (empty when - the task returned none — route via static/conditional edges instead). - Replaces the v1 `completed_tasks`/`completed_routes` pair: persisting a - carried-forward branch's `Command::goto` is what lets it survive an - interrupt/failure + resume instead of re-resolving via static/conditional - edges only. -- barrier (waiting-edge) arrivals (`barrier_arrivals`) -- pending writes -- interrupts -- parent checkpoint id -- free-form `metadata: serde_json::Value`, including well-known keys the - executor itself reads back on resume: `source`, `step` (kept monotonic - across a resume — see `resume::resume_from_inner`'s `initial_steps` - seeding), `node_visits` (per-node visit counts, seeded from here on resume - so `RecursionPolicy::max_visits_per_node` bounds a thread's whole - lifetime, not just one run), `interrupted_nodes`, and (on a failure - boundary) `failed_node`/`error` -- metadata source: `input`, `loop`, `update`, or `fork` (`CheckpointMetadata::source`) - -**Target (not implemented):** the following LangGraph-derived fields do not -exist on `Checkpoint`/`CheckpointMetadata` today — no per-channel -`channel_values`/`channel_versions`, no `versions_seen` map, no -`updated_channels` list, no `graph_id` field, and no structured per-task -outcome record beyond `completed`/`pending_writes`. Introducing these would -require a channel-based state model this crate does not have (state here is a -single typed `State`, not a set of named channels). - -## Checkpoint format v2 - -Checkpoint format v2 (`Checkpoint::version == CHECKPOINT_FORMAT_VERSION == 2`) -replaces four overlapping v1 projections of pending/completed work — -`next_nodes: Vec<NodeId>`, `completed_tasks: Vec<NodeId>` + -`completed_routes: Vec<Vec<RouteTarget>>` (a parallel pair that had to stay -positionally aligned by convention, not by type), and -`pending_activations: Option<Vec<PendingActivation>>` (an `Option`-wrapped -superset of `next_nodes` carrying the same information) — with exactly two -fields: `tasks: Vec<PendingActivation>` and `completed: Vec<CompletedTask>`. -`pending_writes` and `barrier_arrivals` are unchanged; they already carried -information (the resume/error/interrupt control-plane slots, and barrier -arrival sets) that v1's schedule fields never could. - -### Decode path: `Checkpoint::normalize` - -The four v1 fields (`next_nodes`, `completed_tasks`, `completed_routes`, -`pending_activations`) still exist on `Checkpoint<State>` as -`#[serde(default, skip_serializing_if = "…")]` **decode-only** inputs: a -record written by a build that predates checkpoint format v2 still -deserializes, with `version` defaulting to `1` via -`#[serde(default = "checkpoint_version_v1")]`. Every writer in this crate -leaves those four fields at their empty default, so a freshly-written record -serializes with none of them present — only `tasks`/`completed` describe -pending/completed work going forward. - -`Checkpoint::normalize(&mut self)` is the single place that reads the legacy -fields and folds them into the v2 shape: when `version < 2`, `tasks` is -derived from `pending_activations` (falling back to `next_nodes`, projected -to plain node activations with no `Send` arg and an empty task id) and -`completed` is derived by zipping `completed_tasks` with `completed_routes` -(padding a shorter/missing `completed_routes` with empty routing — the -pre-`completed_routes` behavior); the legacy fields are then cleared and -`version` is stamped to `2`. It is a no-op on an already-v2 record. - -Every bundled `Checkpointer` backend calls `normalize()` on every decode -path — `get`, `get_scoped`, `list` (via a normalize-equivalent header -projection that never needs a full-record decode; see "Listing" below), -`state_history`, and `get_thread` — for both backends -(`InMemoryCheckpointer`, `DriverCheckpointer`). Callers -outside this module therefore never observe a v1 record: `checkpoint.tasks`/ -`checkpoint.completed` are always populated and correct regardless of which -format version the stored record was originally written in. Readers -elsewhere in this crate (`compiled::resume`, `compiled::boundary`, -`compiled::state_api`, `compiled::mod`) read `tasks`/`completed` only — the -v1-derivation logic that used to live in those modules (a `next_nodes` -fallback in `resume_from_inner`, a zip of `completed_tasks`/ -`completed_routes` in `boundary::advance`'s carried-completion handling) was -deleted once `normalize()` centralized it. - -`Checkpoint::to_metadata()` (which backs `Checkpointer::list` and the -inspection API) resolves the same v2-or-derived-from-v1 pending set without -mutating `self`, so a listing path that only ever decodes a header (the file -backend's `CheckpointHeader`, the SQLite backend's projected `next_nodes` -column) can report the correct schedule for a v1 record without a full -record decode. - -### Building a v2 checkpoint - -`Checkpoint::new(state, tasks)` builds a fresh v2 record with sensible -defaults: a freshly-minted `checkpoint_id`, the current `created_at`, -`version == CHECKPOINT_FORMAT_VERSION`, and empty `thread_id`/`completed`/ -`pending_writes`/`interrupts`/`barrier_arrivals`/`namespace` (`metadata` is -`null`). `Checkpoint::builder(state)` is an alias with no pending tasks yet. -Every field is also directly `pub`, but the fluent `with_*` setters -(`with_thread_id`, `with_checkpoint_id`, `with_run_id`, -`with_parent_checkpoint_id`, `with_namespace`, `with_tasks`, -`with_completed`, `with_pending_writes`, `with_interrupts`, -`with_barrier_arrivals`, `with_metadata`, `with_channel_versions`, -`with_versions_seen`, `with_channel_deltas`) are what every checkpoint -construction site in this crate uses (`boundary::{advance, -handle_failure_boundary, handle_interrupt_boundary, persist_cancel_checkpoint, -persist_failure_checkpoint, build_loop_checkpoint}`, -`state_api::{update_state, fork_state}`, the conformance suite) instead of a -~15-field struct literal repeating the same four-projection duplication at -every call site. - -### Channel bookkeeping fields (I5/R3) - -Three fields round out the v2 record, all `#[serde(default)]` so an older -checkpoint decodes with them empty: - -- `channel_versions: BTreeMap<String, u64>` — cumulative per-channel write - counters as of this boundary. For a `channel::ChannelState` graph this is - `ChannelState::channel_versions()` verbatim; for a plain whole-`State` - graph it is a single `{"state": n}` entry, bumped once per checkpoint. -- `versions_seen: BTreeMap<String, BTreeMap<String, u64>>` — per-node - snapshot of `channel_versions` as of the last time each node ran (keyed by - node id string — `NodeId` has no `Ord` impl to key a `BTreeMap` directly). - Backs `NodeContext::changed_since_last_run`. -- `channel_deltas: BTreeMap<String, Vec<serde_json::Value>>` — this - checkpoint's own step's raw writes to every - `channel::ChannelSet::with_delta`-tracked channel (not cumulative — see - `docs/modules/graph/state-channels.md`'s "Delta-channel history" section). - Replayed across a lineage by `Checkpointer::delta_history`. - -Every checkpoint-construction call site fills these through one shared -function, `channel::channel_bookkeeping(state, fallback_version)`, so a normal -superstep boundary and a manual `update_state`/`fork_state` write can never -disagree about what they persist here. - -Durability modes: - -- `sync`: persist before the next step starts. -- `async`: persist while the next step executes. -- `exit`: persist only when the graph exits. - -Backends: - -- in-memory -- file-backed JSON/JSONL -- SQLite -- Postgres later - -Thread operations: - -- list checkpoints for a thread -- delete all checkpoints for a thread -- delete checkpoints by run id -- copy a thread to a new thread id -- prune checkpoints with a documented strategy - -Delta channels require careful copy and prune semantics. A checkpoint backend -must not keep only the latest checkpoint if the latest checkpoint depends on -ancestor pending writes or a previous delta snapshot. - -## Long-Term Stores - -LangGraph also has a `BaseStore`, but that is not the same thing as graph state -checkpointing. Stores provide long-term memory that can persist across threads -and conversations. They support hierarchical namespaces, key-value items, -metadata, and optional vector search. - -TinyAgents should mirror this separation: - -- checkpointers store execution state needed to resume a graph exactly -- stores hold application memory, records, artifacts, and searchable data that - graph or harness nodes may read and write - -Compiled graphs may receive a store registry at compile/run time, and -`GraphContext` may expose stores to nodes, but the executor must not use stores -as a substitute for checkpoints. - -## State Inspection And Time Travel - -Compiled graphs with checkpointing should expose: - -- `get_state(thread_id, checkpoint_id)` -- `get_state_history(thread_id, before, limit, filter)` -- `update_state(thread_id, values, as_node, task_id)` -- `bulk_update_state(thread_id, supersteps)` -- `fork_state(source_checkpoint, target_thread_id)` - -State snapshots contain: - -- current values -- next node names -- config used to fetch the snapshot -- checkpoint metadata -- creation timestamp -- parent config -- tasks for the next step -- task errors and results from attempted work -- pending interrupts - -Manual state updates are graph writes. They must pass through the same channel -reducers, produce checkpoint metadata with source `update`, and validate -`as_node` when a caller attributes the write to a node. - -Time travel is implemented by invoking or streaming from an older checkpoint -config or by forking a thread. It must not mutate old checkpoint records. - -## Implemented (TinyAgents) - -The checkpoint core lives in `crates/tinyagents-graph/src/checkpoint/`: - -- `Checkpoint<State>` — the persisted superstep snapshot (thread/checkpoint ids, - parent lineage, namespace, committed state, next/completed nodes, pending - writes, interrupts, and free-form metadata). -- `CheckpointMetadata` — the lightweight listing record. Its `source` field is a - typed `CheckpointSource` (no longer a bare string). -- `CheckpointSource` — `Input | Loop | Update | Fork` with serde (lowercase wire - form) and `Display`. `CheckpointSource::parse` recovers it from a string. -- `DurabilityMode` — `Sync | Async | Exit`, default `Sync`. `Sync` persists a - checkpoint before the next step starts. `Async` hands non-terminal boundary - writes to spawned background tasks so checkpoint I/O stays off the superstep - critical path; the checkpoint id is minted up front so lineage stays chained. - A failed background write is never silently lost: the run fails at the next - durability boundary that observes it, and every in-flight write is awaited - (drained) at the terminal, interrupt, and failure boundaries so the run - result reflects persistence failures. The terminal and interrupt checkpoints - themselves are always written synchronously; `CheckpointSaved` events for - background writes arrive when the write completes, so their ordering - relative to later step events is not deterministic. Outside a tokio runtime - `Async` degrades to `Sync`. `Exit` persists only the terminal checkpoint and - any interrupt boundary (interrupts must persist so the run can resume). Set - it with `CompiledGraph::with_durability(mode)`. - -Two guarantees make `Async` durability safe to reason about. **Background writes -are ordered**: each spawned write awaits the previously spawned one before its -own `put`, because every bundled backend defines a thread's "latest" checkpoint -by insertion order — an unserialized writer could let boundary N land after -boundary N+1 and make a concurrent reader (or a `retry` racing a straggler -append) see a stale record as latest. A write whose predecessor failed is -skipped entirely rather than appending a record whose `parent_checkpoint_id` -points at something that never persisted; the predecessor's error is what the -run reports. **Every exit drains**: not only the terminal, interrupt, and -failure boundaries but also the mid-run aborts — recursion limit, run deadline, -node-visit limit, reducer and routing errors — settle in-flight writes before -returning `Err`. Dropping the tracker would detach those tasks, discarding their -outcome and racing a caller that immediately calls `retry(thread_id)`. - -`prune` is namespace-aware: the recency window applies per namespace, so the -lineages of embedded subgraphs (which share the parent's thread id under their -own namespace, and are never referenced by a parent record's -`parent_checkpoint_id`) are retained rather than deleted along with old parent -records. - -- `CheckpointConfig { thread_id, checkpoint_id, namespace }` — checkpoint - coordinates. `CheckpointConfig::latest(thread_id)` addresses the newest - checkpoint at the root namespace. -- `CheckpointTuple<State> { config, checkpoint, parent_config, pending_writes }` - — the documented core persistence unit. -- `Checkpointer::get_tuple(config)` — a default trait method composed from `get` - so every backend gets it for free; it resolves the concrete config and the - parent's config from the loaded record. - -The `Checkpointer` trait retains its existing `put` / `get` / `list` surface. -Two backends are bundled: - -- `InMemoryCheckpointer` — an `Arc<Mutex<..>>` map, cheap to clone (clones share - storage), for tests and ephemeral runs. -- `DriverCheckpointer` — a durable backend over a `tinystoragedrivers` - `DocumentStore`, so the host's storage backend (SQLite, MongoDB, memory) - holds the checkpoints. Each checkpoint is one document keyed by thread, - namespace and sequence; the full record is stored as JSON, and the parent - id, namespace, source, step and run id are projected into indexed fields so - listing and parent-chain walks do not decode whole states. `list` decodes - headers only (`Checkpoint<IgnoredAny>`), and `get`/`get_scoped`/`get_thread`/ - `state_history` call `Checkpoint::normalize` on the decoded record so a v1 - document is indistinguishable from a v2 one. Pending writes live in their - own collection keyed by checkpoint id. `put_with_writes` commits the - checkpoint and its writes in one atomic batch. The impl is bound by - `State: Serialize + DeserializeOwned`. - -The earlier `FileCheckpointer` and `SqliteCheckpointer` were removed; see the -CHANGELOG for the migration. - -### `put_with_writes` - -`Checkpointer::put_with_writes(checkpoint, writes)` is a default trait method -(so every out-of-tree backend keeps compiling unchanged) composed from `put` -followed by `put_writes`. `DriverCheckpointer` overrides it to run both -writes inside one atomic batch, so a boundary that needs to persist -both a checkpoint and its pending writes commits them atomically instead of -as two independent writes. - -### Thread operations - -`Checkpoint` and `CheckpointMetadata` now carry an optional `run_id` -(back-compatible — pre-existing/manual records leave it `None`). The executor -stamps every boundary checkpoint with the producing run id. - -The `Checkpointer` trait exposes the documented thread operations. Three are -storage-specific primitives (no default body): `list_threads`, `delete_thread`, -and the low-level `delete_checkpoints(thread_id, ids)`. The higher-level -operations are default trait methods composed from those plus `list`/`get`/`put`, -so every backend inherits them: - -- `delete_by_run(thread_id, run_id)` — deletes the checkpoints stamped with a - run id (composed from `list` + `delete_checkpoints`), returning the count. -- `get_thread(thread_id)` — bulk-reads every full checkpoint record for a - thread in listing order. The default is composed from `list` + `get` (one - lookup per id); all three bundled backends override it with a single-pass - read (map clone / one file parse / one indexed range query). -- `copy_thread(source, target)` — deep-copies every checkpoint into a new thread - id, preserving each record's `checkpoint_id` and `parent_checkpoint_id` so the - lineage spine stays walkable for time-travel/resume (composed from - `get_thread` + `put`, so the source thread is read once). -- `prune(thread_id, keep_last)` — retains the most recent `keep_last` - checkpoints **plus the full `parent_checkpoint_id` ancestor chain of every - retained checkpoint**, then deletes the rest. Protecting the entire ancestor - chain is what honors the delta-channel warning: a kept checkpoint that stores - only a delta (or depends on an ancestor's pending writes/snapshot) can never - be orphaned from the state it needs. `keep_last == 0` is clamped to `1` so the - latest checkpoint always survives. - -`InMemoryCheckpointer` implements the three storage primitives (plus a -single-pass `get_thread`); it inherits `delete_by_run`, `copy_thread`, and -`prune` from the trait defaults. - -### State inspection & time travel - -`CompiledGraph` exposes the documented inspection/time-travel surface when a -checkpointer is configured (every method returns `TinyAgentsError::Checkpoint` -if it is not). A `StateSnapshot<State>` bundles the committed `values`, the -`next_nodes`/`tasks` that would run on resume, the `config` addressing the -snapshot, its `parent_config`, the listing `metadata`, and any -`pending_interrupts`. - -- `get_state(thread_id, checkpoint_id)` — loads a snapshot (latest when - `checkpoint_id` is `None`); `Ok(None)` for an unknown thread/checkpoint. -- `get_state_history(thread_id, limit)` — snapshots newest-first, walking the - `parent_checkpoint_id` lineage back from the latest checkpoint; `limit` caps - the count. -- `update_state(thread_id, update, as_node)` — a manual graph write. The - `update` is folded through the same `StateReducer` the executor uses, on top - of the thread's latest committed state, and persisted as a new checkpoint with - source `update`. `as_node` must name a real node (`MissingNode` otherwise); the - write is attributed to it: the node is treated as just-completed (it leaves - the pending set) and its routing successors are merged into the base - checkpoint's remaining pending work, so branches the write never touched keep - running — with their `Send` args intact. A successor behind a waiting edge is - barrier-gated exactly as during a run, and the retained predecessors are what - later clear the join. With `as_node == None` the latest pending set is - preserved. -- `bulk_update_state(thread_id, updates)` — applies a sequence of - `(update, as_node)` pairs as successive `update` checkpoints, each layered on - the previous one's committed state; returns the last config (errors on an - empty sequence). -- `fork_state(source_thread, source_checkpoint_id, target_thread)` — copies a - checkpoint into a new thread as a fresh root (no parent) with source `fork`. - The source record is read with `get` and never mutated, so forks are - non-destructive time travel. - -Time-travel resume is `resume_from(thread_id, target, command)` where `target` -is a `ResumeTarget` (`Latest` or `Checkpoint(id)`). `resume` is shorthand for -`ResumeTarget::Latest`. Resuming from an older checkpoint replays its pending -nodes forward (applying `command`'s resume value to any interrupted node) -without rewriting history — new boundary checkpoints are appended to the thread. diff --git a/docs/modules/graph/execution.md b/docs/modules/graph/execution.md deleted file mode 100644 index d3408e820..000000000 --- a/docs/modules/graph/execution.md +++ /dev/null @@ -1,167 +0,0 @@ -# Graph Execution Model And Parallelization - -Milestone 1 executor: - -- sequential -- one active node at a time -- whole-state updates -- direct or string conditional routes -- recursion limit - -Target executor: - -- Pregel-like bulk synchronous parallel supersteps -- multiple active tasks per step -- immutable state snapshot during a step -- writes visible in the next step only -- deterministic write ordering before reducer application -- reducer/channel updates at step boundaries -- checkpoint after step completion -- pending writes from completed tasks -- cached writes replay without rerunning nodes -- resume from checkpoint -- recursive call tracking -- child graph and child agent run tracking - -Superstep lifecycle: - -1. Load checkpoint, active tasks, and pending writes. On resume, a checkpoint - left mid-step by an interrupt/failure (see step 6 below and - `code-review-graph.md` C1/C2) reloads its `completed_tasks` too, not just - its pending set — those already-completed siblings are not re-run. -2. Emit step started event. -3. **Target (not implemented):** match cached task writes when cache policy - allows it — no `cache_policy`/cached-writes-replay mechanism (an explicit, - opt-in policy for replaying a *stored write's payload* instead of - re-running a node) exists in `crates/tinyagents-graph/src` today; every - *active* task still re-runs. What does exist, unconditionally, is a - narrower but load-bearing guarantee: a task that already **completed** - within the same still-in-progress superstep is never re-run just because - an interrupt/failure boundary was hit — see step 6. -4. Run active tasks under concurrency, timeout, retry, and cancellation policy. - `State` is cloned at most once per superstep here (into an `Arc<State>`, - `code-review-graph.md` M2): every active task and every retry attempt of - one task shares that `Arc` via a cheap `Arc::clone` instead of a fresh - `State` clone per attempt/branch — see - [nodes.md](nodes.md#actual-handler-signature-and-state-cloning-m2). -5. Collect writes, commands, sends, interrupts, and errors — every result, - not only the ones before the first stalled (errored/interrupted) branch - (`crates/tinyagents-graph/src/compiled/step.rs::fold_step`). -6. Persist task writes as pending writes when checkpointing supports it. At an - interrupt/failure boundary this includes every branch that completed this - step (`Checkpoint::completed_tasks`), regardless of its index relative to - the stalled branch; only the stalled branch(es) become the resumed run's - pending set. A completed branch's *routing* (step 8) is deferred rather - than resolved here — resolving it before the stalled branches are known - would let a downstream node observe a state missing their eventual - updates — and is carried forward until the whole step finishes - (`crates/tinyagents-graph/src/compiled/boundary.rs::advance`'s - `carried_completed` handling). -7. Apply channel reducers at the step boundary. An additive channel-per-field - state model does exist (`crates/tinyagents-graph/src/channel/`: - `Channel`, `ChannelSet`, `ChannelState`, with `LastValue`/`Topic`/ - `Delta`/`Messages`/`Barrier`/`BinaryAggregate` merge strategies and - concurrent-write conflict detection), but there is no per-channel - **version** tracking (see checkpointing.md) — conflict detection is - step-stamped, not versioned. -8. Select next active tasks from routing commands and `Send` fan-out. - (**Target:** selection specifically by "channel version changes" does not - apply, since channel versions are not tracked — see above.) -9. Persist the checkpoint according to durability mode. -10. Emit checkpoint, update, task, and step completion events. - -Checkpointing mid-node should be avoided. Async Rust stack suspension is not a -stable persistence primitive; rerunning a node from the beginning is easier to -reason about and matches interrupt semantics. A side effect a node performs -*before* it pauses or fails can be made re-run-safe with -`NodeContext::durable_task(key, fut)`, which memoises the sub-step's output -per `(task_id, key)` in the checkpoint write ledger (see -[builder.md](builder.md#nodecontextdurable_task)). - -## Stopping a run: cancellation and graceful drain - -Two run-level stop signals travel in `RunOptions` (accepted by -`run_with_options`, `run_with_thread_options`, and `resume_with_options`); -they differ only in what happens to the superstep in flight. - -```rust -pub struct RunOptions { - pub cancellation: Option<tinyagents_harness::CancellationToken>, - pub drain: Option<tinyagents_graph::DrainSignal>, -} - -let (handle, signal) = DrainSignal::new(); -let run = graph - .run_with_thread_options("thread-1", state, RunOptions::with_drain(signal)) - .await?; -// elsewhere, e.g. a shutdown hook or a node handler: handle.drain(); -``` - -| | cancellation | drain | -|---|---|---| -| checked | before every superstep **and** raced against the step's handlers | before every superstep only | -| step in flight | abandoned (nothing from it is trusted to have applied) | runs to completion; reducer and boundary checkpoint commit as normal | -| checkpoint | pending = the whole active set | pending = the *next* step's activations, metadata `"drained": true` | -| status / event | `ExecutionStatus::Cancelled`, `GraphEvent::RunCancelled` | `ExecutionStatus::Drained`, `GraphEvent::RunDrained { run_id, steps }` | -| result | `Ok(GraphExecution)`, `drained: false` | `Ok(GraphExecution)`, `drained: true` | - -A drain is the graceful counterpart to cancellation: raise it from a shutdown -hook (or from inside a node via a cloned `DrainHandle`) and the run finishes -the work it has started, persists a checkpoint naming exactly what it did not -start, and returns. `GraphExecution::drained` is `true` only on that path and -`false` on every other outcome (completed, interrupted, failed, cancelled); -`GraphRunStatus::is_terminal()` treats `Drained` like `Cancelled` — the run -id will not advance on its own. `resume`/`retry` continue the thread from the -drained checkpoint like any other pending-tasks checkpoint, reaching the same -final state an undrained run would. A drain raised before the first step -persists the entry activation and runs nothing; without a thread id the run -still stops with `Drained` but has no checkpoint to persist against. The -`Drop` guard and cancellation details are in -[fault-tolerance.md](fault-tolerance.md#cooperative-cancellation). - -## Parallelization - -Parallel execution is represented as multiple active tasks in one superstep. A -node can route to more than one next node through conditional routing, a -command, or `Send` packets. - -```rust -Command::new() - .update(update) - .goto([ - Send::new("retrieve_docs", json!({ "query": "billing" })), - Send::new("retrieve_docs", json!({ "query": "refund" })), - Send::new("score_risk", json!({ "account_id": 42 })), - ]) -``` - -Parallel execution rules: - -- all active nodes in a superstep read the same committed state snapshot -- node-local reads may optionally include that node's own pending writes for - branch decisions -- each node returns partial updates, commands, interrupts, sends, or errors -- channels merge successful writes at the step boundary -- conflicting writes produce reducer/channel errors, not arbitrary last-writer - behavior -- a failed required node fails the step unless an error handler or policy routes - it elsewhere -- completed writes can be preserved as pending writes when other nodes fail -- concurrency is bounded by graph defaults and run config - -Parallelism must be visible in events. The `GraphEvent` enum actually shipped -(`crates/tinyagents-graph/src/stream/types.rs`) uses `Node*` naming, not -`Task*`, and has no cached-task variant: - -- `StepStarted { active: [...] }` -- `TaskScheduled` (not `TaskStarted`) -- `NodeStarted` / `NodeCompleted` / `NodeFailed` (not `TaskCompleted`/`TaskFailed`) -- `NodeRetryScheduled` -- **Target (not implemented):** `TaskCached` — no cached-write replay exists -- `StateUpdated` -- `RouteSelected` -- `CheckpointSaved` / `CheckpointRestored` -- `StepCompleted` - -For agent-specific fanout, forked runtime context, and shared-cache semantics, -see [Parallel agents and context forking](parallel-agents-forking.md). diff --git a/docs/modules/graph/fault-tolerance.md b/docs/modules/graph/fault-tolerance.md deleted file mode 100644 index bda4ea341..000000000 --- a/docs/modules/graph/fault-tolerance.md +++ /dev/null @@ -1,139 +0,0 @@ -# Graph Error Handling And Fault Tolerance - -The graph runtime survives transient network failures and makes a hard failure -restartable, using two opt-in mechanisms plus the existing checkpoint/resume -machinery. - -## Node retry (surviving network blips) - -`CompiledGraph::with_node_retry(RetryPolicy)` wraps every node handler. When a -handler fails with a [retryable] error — `TinyAgentsError::Model` or -`TinyAgentsError::Tool`, the transient class classified by -`harness::retry::is_retryable` — the executor re-runs the node **from its start** -up to the policy's `max_attempts`, emitting `GraphEvent::NodeRetryScheduled` -before each retry. Because a node is never suspended mid-flight, a retry rebuilds -the handler future and context from scratch, matching the durable-execution -model. - -Backoff between attempts is computed from the `RetryPolicy` (exponential, with -optional jitter) but only **slept on** when the policy opts in via -`RetryPolicy::with_backoff_sleep(true)`. The default is sleep-free so unit tests -stay deterministic; production callers enable it so retries wait a real, growing -delay. The same opt-in switch governs the harness model-call retry loop and -`RetryMiddleware`. - -Both the sequential and parallel runners drive handlers through one shared -`run_node_with_retry` helper, which also applies the per-node timeout — so a -retried attempt is still bounded by `node_timeout`. - -## Resumable failures (restarting / continuing after a hard failure) - -When a handler fails **beyond** the retry budget, or with a non-retryable error, -the run is not lost. On a checkpointed thread the executor: - -1. folds the updates of the branches that completed *before* the failing one - into committed state (partial parallel progress is preserved, not discarded); -2. persists a **failure-boundary checkpoint** whose `next_nodes` schedule the - failed node (and the not-yet-run tail of the step) for re-run, records the - successful nodes as `completed_tasks`, and stamps the rendered `error` and - `failed_node` into the checkpoint metadata; -3. records a `Failed` run status carrying that checkpoint id; -4. returns the error to the caller. - -Without a checkpointer/thread the run aborts immediately with the error, exactly -as before — the failure checkpoint is a no-op. - -### Restarting and continuing - -Because `resume`/`resume_from` replay from a checkpoint's `next_nodes`, a failed -run resumes exactly like an interrupted one: - -- `CompiledGraph::retry(thread)` — shorthand for `resume` with an empty command; - re-runs the failed node and the not-yet-run tail from the failure boundary. -- **Continue on user feedback** — inspect committed state with `get_state`, edit - it with `update_state` (attributed through the reducers), then `retry`/`resume`. - The re-run sees the edited state. - -See the `resilient_graph` example (`cargo run --example resilient_graph`) for -both mechanisms end to end. - -### Per-thread execution lease - -`CompiledGraph::execute` holds a per-`(thread_id, namespace)` lock for a run's -whole lifetime: an in-process `ThreadLockMap` guard first (always active), and, -when a checkpointer is configured, a durable lease claimed via -`Checkpointer::try_claim`/`renew`/`release` (owner = run id, default TTL 5 -minutes). A second concurrent `run_with_thread`/`resume`/`retry` call for the -same thread either waits on the in-process lock (same process) or is refused -with `TinyAgentsError::Validation` if a live lease is held by another owner -(cross-process). A lease whose owner crashed without releasing it is -reclaimable once its TTL elapses. `DriverCheckpointer` and `InMemoryCheckpointer` -both implement the lease; the trait's default is a no-op that always succeeds, -so out-of-tree backends keep compiling unprotected. - -## Panic safety - -Each node handler future is polled through -`futures::FutureExt::catch_unwind(AssertUnwindSafe(..))`. A panic inside a -handler no longer unwinds through the executor: it is converted into -`TinyAgentsError::Graph("node `{id}` panicked: {msg}")` and flows through the -same resumable failure boundary as any other node error (checkpoint, -`RunFailed` event, `Failed` status). `retry(thread)` re-runs the panicking -node exactly as it would a node that returned `Err`. - -## Cooperative cancellation - -`RunOptions { cancellation: Option<tinyagents_harness::CancellationToken> }` -carries an optional cancellation token into `CompiledGraph::run_with_options`, -`run_with_thread_options`, and `resume_with_options`. The token is checked -before every superstep and raced against that step's in-flight handler -futures. On cancellation the executor persists a checkpoint whose -`next_nodes` is the still-pending active set, records a `Cancelled` run -status (and emits `GraphEvent::RunCancelled`), and returns `Ok` rather than -an error — a cancelled run is resumable exactly like an interrupted one. - -`RunOptions::drain` (a `DrainSignal`, paired with the `DrainHandle` that -raises it) is the graceful variant: it is checked only *between* supersteps, -so the step in flight always completes and commits before the run stops with -`ExecutionStatus::Drained` / `GraphEvent::RunDrained` and -`GraphExecution::drained == true`. See -[execution.md](execution.md#stopping-a-run-cancellation-and-graceful-drain) -for the side-by-side contract. - -A `Drop` guard armed for the run's lifetime protects against the run future -itself being dropped mid-flight (a host timeout racing `tokio::select!`, -`JoinHandle::abort`, and similar). It disarms on every normal terminal exit -(success, failure, interrupt, explicit cancel); if it is still armed when -dropped, it best-effort persists a `Cancelled` status from a detached task so -the run is never left stuck at `Running`. This is a status guarantee, not a -full flush guarantee: an in-flight `AsyncCheckpointWrites` write already -spawned before the drop still completes in the background (tokio does not -abort a `JoinHandle`'s underlying task on drop), but the guard's own tracking -of that write's outcome is lost. - -## Error taxonomy - -`TinyAgentsError` distinguishes structural/config errors (non-resumable) from -node failures (resumable on a checkpointed thread): - -- structural / not retried, not resumable: missing start, missing node, missing - edge target, missing route, invalid command/parent/send target, recursion - limit, node-visit limit, sub-agent depth, checkpoint required/missing, resume - mismatch, reducer conflict, invalid concurrent update, serialization failure, - checkpoint backend failure; -- transient / retryable and resumable: `Model`, `Tool` (and `Timeout` aborts the - node but leaves a resumable boundary like any other node failure). - -## Future work - -- Route node errors to a node-specific or default error-handler node. -- Cooperative drain/shutdown with a drain reason. -- Populate the checkpoint `pending_writes` list explicitly (today partial - progress is folded into committed state instead). -- Renew the durable execution lease mid-run for long-running steps that could - outlive its TTL (today it is claimed once, at `execute` entry, and released - at exit — no heartbeat loop). -- Force-drain (not just best-effort track) in-flight async checkpoint writes - from the run-future drop guard. - -[retryable]: ../../../crates/tinyagents-harness/src/retry/mod.rs diff --git a/docs/modules/graph/goals.md b/docs/modules/graph/goals.md deleted file mode 100644 index 170b553a5..000000000 --- a/docs/modules/graph/goals.md +++ /dev/null @@ -1,6 +0,0 @@ -# Goals moved to harness-tools - -Per-thread goals are generic harness functionality, not graph executor policy. -The `GoalStore`, observer, atomic accounting, budget guard, and host-facing turn -preparation API live in `crates/tinyagents-harness-tools/src/goals/`. See -[`Goals-and-Todos`](../../../wiki/Goals-and-Todos.md). diff --git a/docs/modules/graph/interrupts.md b/docs/modules/graph/interrupts.md deleted file mode 100644 index faeb5654d..000000000 --- a/docs/modules/graph/interrupts.md +++ /dev/null @@ -1,191 +0,0 @@ -# Graph Interrupts And Resume - -Interrupts pause execution and return control to the caller. Everything in -this section — the `Interrupt` type, `Command::resume`/`resume_tasks`, -`CompiledGraph::resume`/`resume_from`/`retry`, the `interrupt_before` / -`interrupt_after` selectors, and `response_schema` validation — is -implemented today (`crates/tinyagents-graph/src/command/`, -`crates/tinyagents-graph/src/builder/mod.rs`, -`crates/tinyagents-graph/src/compiled/{step,boundary,resume}.rs`). Only the -"Targeted Human Steering" section at the end remains a target. - -The struct as shipped — `id` is a bare `String` (not `InterruptId`) and -there is no `order` field; `task_id` is stamped by the interrupt boundary -with the pausing branch's task id, and `response_schema` is optional: - -```rust -pub struct Interrupt { - pub id: String, - pub node: NodeId, - pub payload: serde_json::Value, - pub task_id: Option<TaskId>, - pub response_schema: Option<serde_json::Value>, -} -``` - -Resume API: - -```rust -compiled_graph - .resume("support-123", Command::resume(json!({ "approved": true }))) - .await?; -``` - -Rules: - -- interrupts require both a checkpointer and a `thread_id` -- if a node emits an interrupt without resumable durability, the run returns a - resume error instead of an interrupted execution -- interrupted executions are returned only after the checkpoint needed for - resume has been persisted -- the interrupted node restarts from the beginning (a node-emitted interrupt - and an `interrupt_before` pause); an `interrupt_after` pause is the one - exception — the node already ran and its stored result is replayed instead -- every branch of a step that interrupts is surfaced (`GraphExecution::interrupts` - carries all of them, not just the lowest-index one) — a `Send` fan-out of - one node interrupting on several concurrent activations is matched by task - id, each stamped onto its own `Interrupt::task_id`, rather than by an - `order` field -- resume values as a map from task id to value: `Command::resume_tasks(..)` / - `Command::resume_by_task` deliver a distinct value per interrupted task in - one resume call, keyed by `TaskId` — `Command::resume(value)` (one value, - fanned to every task named by the checkpoint's stamped `interrupted_nodes` - or, absent that, to every pending task) still works and is consulted as the - fallback for any task the map does not name -- **Target (not implemented):** resume values as a map keyed by interrupt id - specifically (rather than task id) -- node code before an interrupt must be deterministic or idempotent, or - wrapped in `NodeContext::durable_task` (see - [builder.md](builder.md#nodecontextdurable_task)) so a side effect ahead of - the pause is memoised rather than repeated on the re-run - -## `interrupt_before` / `interrupt_after` selectors - -```rust -GraphBuilder::<State, Update>::new() - // ... - .interrupt_before(["approve"]) - .interrupt_after(["plan", "act"]) // needs Update: Serialize + DeserializeOwned - .mark_interrupt("review") // alias for interrupt_before(["review"]) -``` - -Executor-level pauses at named nodes, for debugging, approvals, and human -review at arbitrary graph boundaries without editing node code. They are -distinct from a node returning `NodeResult::Interrupt` itself: the executor -injects the `Interrupt`, with payload `{"phase": "before"}` or -`{"phase": "after"}`, stamped with the activation's task id, and persists the -usual interrupt-boundary checkpoint (`has_interrupts: true` in -`get_state_history`). Both selectors are validated at `compile()` -(`TinyAgentsError::MissingNode` for an unknown node) and mark the node as an -interrupt point in the topology export. - -`interrupt_before(nodes)` pauses *instead of* invoking a listed node's -handler. The checkpoint's pending set is that activation; `resume` / -`retry` re-schedules it and runs the handler normally, delivering any -`Command::resume` value on `NodeContext::resume`. The handler runs exactly -once overall. - -`interrupt_after(nodes)` lets the handler run to completion, then holds its -`Update`/`Command` back from committed state: the result is serialized as a -deferred-result write (`PendingWrite::interrupt_after`, control-plane idx -`WRITES_IDX_INTERRUPT_AFTER`) in the checkpoint's `pending_writes`, the -checkpoint's `state` and the returned `GraphExecution::state` are the -*pre-update* state, and the node stays in the pending set. On resume the -executor **replays** the stored result — decoding the update through the -`Update` codec and honouring the node's `goto` — without invoking the -handler again, so it still ran exactly once and nothing is lost. This is -why `interrupt_after` requires `Update: Serialize + DeserializeOwned` (the -same bound `CompiledGraph::with_cached_node` takes); `interrupt_before` has -no such bound. A node that returns its own `NodeResult::Interrupt` is not -paused a second time, and an `Err` passes through to the failure boundary. - -Acknowledgement: resuming a checkpoint acknowledges the executor-injected -pauses it recorded (per `(phase, task_id)`), so the re-run skips that phase. -The acks of a still-pending task are carried forward in checkpoint metadata -(`acknowledged_interrupts`) across any later pause of the same task — an -`interrupt_before` node that then pauses at `interrupt_after`, or emits its -own interrupt, is not paused *before* again. `update_state` carries both the -acks and the deferred-result write forward for tasks it leaves pending. - -`mark_interrupt(node)` is `interrupt_before([node])`: earlier versions set -only the export marker, the runtime pause is now real. - -## `response_schema` - -```rust -Interrupt::new("approve", json!({ "ask": "approve?" })) - .with_response_schema(json!({ - "type": "object", - "required": ["approved"], - "properties": { "approved": { "type": "boolean" } } - })) -``` - -An interrupt can declare the JSON-schema subset (`type`, `required`, -`properties`, `additionalProperties`, `items`, `enum`) its resume value must -satisfy. `resume`/`resume_from` validate **fail-closed, before any -checkpoint write**: the value each schema-bearing pending interrupt would -receive — from `Command::resume_tasks` by task id, or `Command::resume` -fanned out exactly as `NodeContext::resume` would be filled — is checked -with `tinyagents_harness::tool::validate_against_schema` before the resumed -run claims the thread. A mismatch returns `TinyAgentsError::Validation` -naming the interrupt and the offending path, and the thread's latest -checkpoint is untouched (same checkpoint id, same pending interrupt, no new -history entry). A bare `retry()` delivers no value and validates nothing. -`response_schema` is `None` by default and omitted from checkpoint JSON when -unset, so legacy records decode unchanged. - -## Targeted Human Steering - -**Target (not implemented; see `docs/runtime-comparison/plan.md`).** Nothing -below this point — `ResumeTarget` as a targeted-steering struct, -`resume_targeted`, or per-run/per-task/per-namespace resume routing — exists -in `crates/tinyagents-graph/src` today. The `ResumeTarget` type that does -exist (`crates/tinyagents-graph/src/compiled/types.rs`) is unrelated: it is a -`Latest`/`Checkpoint(CheckpointId)` enum selecting which checkpoint a resume -replays from, not a target-selection struct with `run_id`/`task_id`/ -`interrupt_id`/`namespace` fields. - -Human input during an interrupt is one form of steering. A control surface -should be able to target: - -- the parent orchestrator run -- a specific child sub-agent run -- a graph task id -- a node namespace inside a subgraph -- a specific interrupt id - -Targeted resume shape: - -```rust -pub struct ResumeTarget { - pub run_id: RunId, - pub task_id: Option<TaskId>, - pub interrupt_id: Option<InterruptId>, - pub namespace: Vec<String>, -} - -compiled_graph - .resume_targeted( - ResumeTarget { - run_id, - task_id: Some(child_task), - interrupt_id: Some(approval_interrupt), - namespace: vec!["supervisor".into(), "research_agent".into()], - }, - Command::resume(json!({ "approved": true })), - ) - .await?; -``` - -Rules: - -- resuming a child interrupt resumes that child task, not all paused siblings -- resuming the parent orchestrator may leave child interrupts pending unless - policy cancels or resolves them -- a human can add steering instructions while resuming, but those instructions - must be recorded separately from the interrupt answer -- stale resume targets are rejected with the latest run/checkpoint metadata -- UI clients should present pending interrupts with run tree path, node id, - task id, sub-agent id, and checkpoint id so humans can steer the intended - target diff --git a/docs/modules/graph/memory-boundary.md b/docs/modules/graph/memory-boundary.md deleted file mode 100644 index ae839c15c..000000000 --- a/docs/modules/graph/memory-boundary.md +++ /dev/null @@ -1,32 +0,0 @@ -# Graph Memory And Stores Boundary - -Graph checkpointing is short-term execution persistence. Harness memory and -stores are separate concerns. - -LangGraph makes this distinction directly: - -- checkpointers store graph state by thread and checkpoint id -- stores provide long-term memory across threads and conversations, with - hierarchical namespaces, metadata, and optional vector search - -Graph owns: - -- checkpoints -- channel values -- channel versions -- pending writes -- active tasks -- interrupts -- state history and forks - -Harness/store owns: - -- chat history storage -- semantic memory -- user/application records -- model/tool artifacts -- usage and cost records -- provider traces - -The graph can pass a `StoreRegistry` through `GraphContext`, but graph execution -must not assume any specific store backend. diff --git a/docs/modules/graph/milestones.md b/docs/modules/graph/milestones.md deleted file mode 100644 index f50a3ff5e..000000000 --- a/docs/modules/graph/milestones.md +++ /dev/null @@ -1,87 +0,0 @@ -# Graph Implementation Milestones - -## G1: Current Sequential Runtime - -- `Node` -- `NodeOutput` -- `StateGraph` -- direct edges -- conditional edges -- recursion limit - -## G2: Builder And Compile Step - -- introduce `GraphBuilder` -- introduce `CompiledGraph` -- move validation to `compile` -- add `START` and `END` -- add input/output schema concepts -- preserve current closure-node ergonomics - -## G3: Commands, Sends, And Typed Events - -- add `Command` -- add `Send` -- add `GraphEvent` -- add event recorder -- add destination hints for command-rendered edges - -## G4: Channels And Partial Updates - -- add `Channel` trait -- add state update type -- add last-value, append, aggregate, topic, and message reducers -- add invalid concurrent update errors -- add partial update examples - -## G5: Supersteps And Parallel Execution - -- add multi-active-task executor -- apply writes at step boundaries -- add waiting/barrier edges -- add concurrency limits -- add task stream events - -## G6: Checkpointing, Interrupts, And Time Travel - -- add `Checkpointer` -- add in-memory backend -- add checkpoint tuples and pending writes -- add interrupt/resume API -- add `get_state`, `get_state_history`, and `update_state` -- add durability modes - -## G7: Policies, Cache, And Error Handlers - -- add graph and node defaults -- add retry, timeout, cache, and error-handler policies -- replay cached writes -- add cooperative drain - -## G8: Subgraphs And Namespaces - -- add subgraph node -- add checkpoint namespaces -- add parent command routing -- stream nested subgraph events -- expose child state in task metadata - -## G9: Sub-Agents And Recursion - -- add `SubAgentNode` -- add child run hierarchy -- add context forking for parallel sub-agents -- add shared-cache fork policy -- add recursion stack -- add depth events -- add max-depth policy -- roll up child usage and cost events - -## G10: Full Graph Streaming And Introspection - -- expose async graph event streams -- add stream projections for values, updates, messages, custom data, - checkpoints, tasks, and debug -- forward harness streams with node context -- add stream transformers -- add JSON and Mermaid graph exports diff --git a/docs/modules/graph/nodes.md b/docs/modules/graph/nodes.md deleted file mode 100644 index af0c4c0ef..000000000 --- a/docs/modules/graph/nodes.md +++ /dev/null @@ -1,94 +0,0 @@ -# Graph Node Model - -Nodes are async units of work. They receive a state view plus runtime context -and return an update, command, interrupt, or no-op. - -```rust -#[async_trait] -pub trait GraphNode<State, Ctx = ()>: Send + Sync { - async fn run( - &self, - state: StateView<'_, State>, - ctx: &mut GraphContext<Ctx>, - ) -> Result<NodeResult>; -} - -pub enum NodeResult { - Update(StateUpdate), - Command(Command), - Interrupt(Interrupt), - None, -} -``` - -Closure-backed nodes stay as an ergonomic adapter: - -```rust -Node::new("agent", |state| async move { - Ok(NodeOutput::continue_with(state)) -}) -``` - -Target node spec: - -```rust -pub struct NodeSpec<State, Ctx = ()> { - pub id: NodeId, - pub node: Arc<dyn GraphNode<State, Ctx>>, - pub input: ChannelSelection, - pub destinations: Option<DestinationHints>, - pub metadata: serde_json::Value, - pub defer: bool, - pub retry: Option<RetryPolicy>, - pub cache: Option<CachePolicy>, - pub timeout: Option<TimeoutPolicy>, - pub error_handler: Option<NodeId>, - pub is_error_handler: bool, -} -``` - -Important node kinds: - -- closure node -- trait-backed node -- harness model/tool/agent-loop node -- sub-agent node -- subgraph node -- router node -- error-handler node -- test node - -Deferred nodes run near graph termination, after normal active work is drained. -Use them for cleanup, final scoring, final summarization, or output shaping. - -## Actual handler signature and state cloning (M2) - -The shipped `NodeHandler<State, Update>` (`crates/tinyagents-graph/src/builder/types.rs`) -is `Fn(Arc<State>, NodeContext) -> NodeFuture<Update>` — every handler -receives the step's committed state as an `Arc<State>`, not an owned -`State`. This is `docs/runtime-comparison/code-review-graph.md` finding M2: -before it, the executor cloned `State` once per handler invocation (every -retry attempt, every parallel `Send`/fan-out branch), so an N-way fan-out -with retries could clone the whole state many times in a single superstep. -Now `StepRunner::run_step` (`crates/tinyagents-graph/src/compiled/step.rs`) -clones `State` **at most once per superstep**, into that `Arc`; every -branch/attempt within the step shares it via a cheap `Arc::clone`. - -Two builder entry points wire a closure into a node: - -- `GraphBuilder::add_node(id, |state: State, ctx| async { .. })` — the - by-value convenience form kept for source compatibility. It is a thin - adapter that clones out of the `Arc` once per invocation, so every - existing handler written against the pre-M2 signature still compiles - unchanged. -- `GraphBuilder::add_node_shared(id, |state: Arc<State>, ctx| async { .. })` - — the zero-clone form. Prefer it for a node with a large `State` (e.g. one - carrying message history) or one that only reads its state; it forwards - the step's `Arc<State>` directly with no clone at all. - -`NodeContext::send_arg` (the per-invocation `Send` argument) is likewise -`Option<Arc<serde_json::Value>>` rather than `Option<serde_json::Value>`, so -a repeated `Send` fan-out of the same node and every retry of one activation -share the argument's allocation too. It serializes exactly like a bare -`serde_json::Value` in checkpoint records (serde's `Arc<T>` impl is -transparent), so on-disk format is unaffected. diff --git a/docs/modules/graph/observability.md b/docs/modules/graph/observability.md deleted file mode 100644 index 44775af3b..000000000 --- a/docs/modules/graph/observability.md +++ /dev/null @@ -1,419 +0,0 @@ -# Graph Observability And Tracing - -LangGraph has several observability surfaces: - -- stream modes for values, updates, messages, custom data, checkpoints, tasks, - debug payloads, and execution events -- `debug` streams that wrap checkpoint, task, and task-result payloads with step - and timestamp metadata -- callback/tracing integration through the LangChain callback manager -- `ls_integration = "langgraph"` metadata so traces can identify LangGraph runs -- graph lifecycle callbacks for interrupt and resume events -- LangSmith tracing configuration in the SDK -- thread-level stream modes for run lifecycle and state updates -- nested namespace propagation for subgraph streams - -TinyAgents should expose observability as a graph feature, not only as a side -effect of streaming. Streaming is the transport; observability is the semantic -contract for what every run, step, task, checkpoint, interrupt, subgraph, and -child agent must report. - -## Execution Status Store - -Graph runs need a readable status surface in addition to transient event -streams. A caller, web UI, supervisor, or test should be able to ask what a graph -is doing without holding the original stream handle. - -The graph runtime should write a compact status record at every execution -boundary: - -```rust -pub struct GraphRunStatus { - pub run_id: RunId, - pub root_run_id: RunId, - pub parent_run_id: Option<RunId>, - pub thread_id: Option<ThreadId>, - pub graph_id: GraphId, - pub checkpoint_id: Option<CheckpointId>, - pub checkpoint_namespace: Vec<String>, - pub status: ExecutionStatus, - pub current_step: usize, - pub active_nodes: Vec<NodeId>, - pub pending_interrupts: Vec<InterruptId>, - pub last_event_id: Option<EventId>, - pub started_at: SystemTime, - pub updated_at: SystemTime, - pub ended_at: Option<SystemTime>, - pub error: Option<GraphErrorSummary>, - pub metadata: serde_json::Value, -} - -pub enum ExecutionStatus { - Pending, - Running, - Waiting, - Interrupted, - Draining, - Completed, - Failed, - Cancelled, -} -``` - -Required status transitions: - -- `Pending` before the first step is scheduled. -- `Running` while one or more graph tasks are executing. -- `Waiting` when the executor is idle but expects more work, such as async - checkpoint persistence or external resume input. -- `Interrupted` when a human-in-the-loop interrupt has been emitted and the run - cannot continue without a resume command. -- `Draining` after cancellation or failure while child tasks, streams, and - checkpoint writes are closing. -- `Completed`, `Failed`, or `Cancelled` as terminal states. - -The status store is not a replacement for checkpoints. Checkpoints preserve -resumable state; status records summarize live and recent execution for -observers. A checkpointer may persist status records in the same backend, but -the API should keep `Checkpointer` and `GraphStatusStore` separate. - -```rust -#[async_trait] -pub trait GraphStatusStore: Send + Sync { - async fn put_status(&self, status: GraphRunStatus) -> Result<()>; - async fn get_status(&self, run_id: RunId) -> Result<Option<GraphRunStatus>>; - async fn list_by_thread(&self, thread_id: ThreadId) -> Result<Vec<GraphRunStatus>>; -} -``` - -Status writes should be idempotent by `run_id` and monotonic by `updated_at` or -step. Readers must never need to deserialize full graph state to answer basic -questions such as "is this run still active?", "which node is executing?", or -"which interrupt is waiting?". - -Run-level latency can be derived from `started_at` and either `ended_at` or -`updated_at`. This status-derived view is compact and cheap, but detailed step -and node latency comes from durable observations. - -## Observability Event Model - -Every graph event should carry stable envelope data: - -```rust -pub struct GraphObservation { - pub event_id: EventId, - pub run_id: RunId, - pub root_run_id: RunId, - pub parent_run_id: Option<RunId>, - pub thread_id: Option<ThreadId>, - pub checkpoint_id: Option<CheckpointId>, - pub checkpoint_namespace: Vec<String>, - pub graph_id: GraphId, - pub node_id: Option<NodeId>, - pub task_id: Option<TaskId>, - pub step: Option<usize>, - pub time: SystemTime, - pub tags: Vec<String>, - pub metadata: serde_json::Value, - pub kind: GraphObservationKind, -} -``` - -Required observation kinds: - -- `graph.started` -- `graph.completed` -- `graph.failed` -- `graph.draining` -- `step.started` -- `step.completed` -- `task.started` -- `task.completed` -- `task.failed` -- `task.cached` -- `state.updated` -- `route.selected` -- `checkpoint.saved` -- `checkpoint.loaded` -- `interrupt.emitted` -- `interrupt.resumed` -- `subgraph.started` -- `subgraph.completed` -- `subagent.started` -- `subagent.completed` -- `stream.opened` -- `stream.closed` -- `debug.payload` - -## Latency Metrics - -The observability module exposes `GraphLatencyMetrics` for durable graph -latency rollups. Metrics are derived from `GraphObservation` timestamps, so a -late UI or supervisor can compute them from the event journal without holding -the original in-process stream. - -Graph latency metrics include: - -- `run_elapsed_ms`: first `run.started` to first terminal `run.completed` or - `run.failed` -- step latency: `step.started` to `step.completed`, correlated by superstep -- node latency: `node.started` to `node.completed` or `node.failed`, - correlated by `(node_id, step)` -- total, max, and average latency helpers for completed steps -- total, max, and average latency helpers for completed node handlers - -```rust -let observations = journal.read_from(run_id, 0).await?; -let metrics = GraphLatencyMetrics::from_observations(&observations); - -let run_ms = metrics.run_elapsed_ms; -let slowest_node_ms = metrics.max_node_ms; -let average_step_ms = metrics.average_step_ms(); -``` - -Incomplete work is ignored. A `step.started` without `step.completed`, or a -`node.started` without `node.completed` / `node.failed`, does not produce a -latency record because there is no terminal timestamp. Failed nodes are retained -with `failed = true` so dashboards can separate slow successes from slow -failures. - -### Node / Tool Health - -A graph node is the graph's unit of work and is frequently a delegated agent or -tool call (`SubAgentNode`), so per-node success/failure counts double as **tool -health** telemetry. `GraphHealthSummary::from_observations` rolls the same -durable stream into: - -- `total_started` / `total_completed` / `total_failed` node counts -- `run_failed` — whether the run emitted `run.failed` -- a per-node `GraphNodeHealth` entry (`started`, `completed`, `failed`), sorted - by node id, with `failure_rate()` and `is_healthy()` helpers -- `is_healthy()`, `failure_rate()`, and `unhealthy_nodes()` at the run level - -```rust -let observations = journal.read_from(run_id, 0).await?; -let health = GraphHealthSummary::from_observations(&observations); - -if !health.is_healthy() { - for node in health.unhealthy_nodes() { - eprintln!("{} failed {}x", node.node.as_str(), node.failed); - } -} -``` - -The summary is attached to the Langfuse trace metadata (below) so it is visible -alongside the exported spans. - -## Event Journal And Listener Replay - -Outside listeners need two paths: - -- live subscription for events emitted after the listener attaches -- replay from a durable offset for events that were emitted earlier - -The graph should therefore support an append-only event journal when configured. -Each journal record should include the observation envelope, stream namespace, -monotonic offset, and redaction status. Consumers can store the last seen offset -and resume later without inspecting checkpoints. - -```rust -#[async_trait] -pub trait GraphEventJournal: Send + Sync { - async fn append(&self, observation: GraphObservation) -> Result<EventOffset>; - async fn read_from( - &self, - run_id: RunId, - offset: EventOffset, - ) -> Result<Vec<GraphObservationRecord>>; -} -``` - -Live listeners should be registered against filters rather than ad hoc callback -closures: - -- run id, root run id, thread id, graph id, node id, or namespace -- event kind families such as graph, step, task, state, checkpoint, interrupt, - subgraph, sub-agent, harness child, or debug -- redaction profile -- replay offset - -Listener delivery is best-effort by default. If a deployment requires durable -delivery, the listener should read from the event journal and acknowledge -offsets outside the graph executor. - -### Implementation status (runtime-comparison Phase 3, C3/G-M7) - -`GraphEventJournal`/`GraphObservation`/`JournalGraphSink` already exist as -described above (offset-addressable, best-effort delivery). Two additions -this phase: - -- **`JournalGraphSink::dropped()`** (G-M7) exposes the count of observations - dropped because the background drain's bounded queue was full when they - were submitted — the lossy-under-load behavior was already real (`emit` - never blocks the executor on I/O) but previously had no way to detect it - happened. A non-zero value means the journal for that run is an - incomplete record; size the drain capacity for the workload, or poll this - counter and alert, rather than assume completeness. -- **`GraphEventSink::emit` now takes a `GraphEventEnvelope`**, not a bare - `GraphEvent` — see `docs/modules/graph/streaming.md`'s "Implementation - status" for `run_id`/`task_id`/`ns`/`seq`. `JournalGraphSink` unwraps the - envelope's `event` for its own `GraphObservation` (which keeps its own - independent `run_id`/`namespace`/`offset`, configured at sink construction) - and forwards the whole envelope to any configured `inner` live sink. - -## Observability Cache - -Some observability projections are expensive or repetitive to compute. The graph -module should permit cache-backed projections without making cached data the -source of truth. - -Cacheable graph observability records include: - -- latest status by run id -- latest status by thread id -- node/task timing rollups -- graph latency rollups from durable observations -- latest state-update summary per step -- checkpoint metadata summaries -- visualization/introspection snapshots -- listener replay cursors for test and UI adapters - -Cache records must include source coordinates: run id, checkpoint id, -checkpoint namespace, step, event offset, and projection version. If any source -coordinate changes, the projection is stale and must be recomputed. - -## Trace Integration - -Graph tracing should support: - -- run names -- run ids supplied by callers -- inherited root and parent run ids -- tags and metadata -- per-node/task timing -- retry and cache metadata -- checkpoint ids and checkpoint namespaces -- graph recursion depth -- subgraph namespace paths -- child harness model/tool events linked to the graph node/task that emitted - them - -The graph should not require a specific tracing backend. It should expose typed -events and adapters: - -- no-op observer -- in-memory observer for tests -- JSONL observer for local debugging -- tracing-span adapter -- a Langfuse ingestion exporter (`GraphLangfuseExporter`, see below) -- OpenTelemetry adapter later -- LangSmith-compatible exporter later - -## Langfuse Export - -`GraphLangfuseExporter` is the concrete, implemented exporter for the Langfuse -ingestion API. It is pull-based and best-effort: read a run's observations back -from a journal, build a batch, and send it — it never sits on the executor hot -path. - -The exporter reuses the harness `LangfuseClient` for transport (Basic or Bearer -auth, endpoint normalization, `207 Multi-Status` handling), so a single set of -credentials serves both the agent and graph sides: - -```rust -let exporter = GraphLangfuseExporter::from_env()?; // or ::new(client) -let observations = journal.read_from(run_id, 0).await?; -exporter - .send_observations(LangfuseTraceConfig::default(), &observations) - .await?; -``` - -Hosts that provide their own transport can call -`tinyagents_graph::build_graph_langfuse_batch(trace, &observations, metadata_fn)`. -This pure builder requires no `LangfuseClient`, credentials, environment -variables, or network access. It returns the same ingestion payload used by -`GraphLangfuseExporter`, including graph hierarchy and health metadata. The -optional span metadata hook follows the exporter's override rules. An empty -journal returns a validation error. The host then hands the payload to its -own export sink, retaining ownership of endpoint and authentication policy. - -The batch maps the graph run onto Langfuse's trace/observation model: - -- one `trace-create` (id defaults to the run's `root_run_id`; name defaults to - the `graph_id`; session defaults to the run's `thread_id`) with the - `GraphHealthSummary` folded into trace metadata -- one structural `span-create` for the graph run (`{trace}:run:{run_id}`, named - for the `graph_id`), parented to the trace, bracketing the whole run -- one timed `span-create` per superstep (`{trace}:step:{n}`), parented to the - graph-run span -- one timed `span-create` per node handler (`{trace}:node:{name}:{step}`), - parented to its superstep span; `node.failed` promotes the span to `ERROR` - level with the rendered error as `statusMessage` -- one timed `span-create` per embedded subgraph, parented to the graph-run span -- an `event-create` for every remaining observation (routes, checkpoints, - interrupts, custom writes, run lifecycle), parented to the graph-run span, - with `run.failed` mapped to `ERROR` - -Still-running work (a `node.started` with no terminal event) is exported as an -open span with a start time but no end time. - -### Unified Traces With Agents And Tools - -Because a graph run and the agent runs its nodes spawn share the same -`root_run_id`, and both the graph and harness exporters default their `traceId` -to that root run id, exporting a graph run **and** its child agent runs lands -every graph step, node, model generation, and tool call under one Langfuse -trace. This is what makes full end-to-end telemetry — including tool health and -tool timing — visible in a single trace tree. - -The nesting is exact, not just co-located: the graph-run span -(`{trace}:run:{run_id}`) uses the same id scheme the harness exporter parents -its agent run spans to. A sub-agent a node spawns carries `parent_run_id` equal -to the graph run id, so its harness-exported run span resolves to -`{trace}:run:{graph_run_id}` and nests directly under the graph-run span rather -than floating at the trace root — the graph, its nodes, and their agents form -one contiguous tree. - -## Debug Payloads - -Debug streams are for executor introspection. They should be structured enough -for tests and UIs, but they are allowed to expose implementation details that -normal event consumers do not need. - -Debug payloads should include: - -- checkpoint payloads with state values, next nodes, tasks, metadata, and parent - config -- task start payloads with task id, node name, input, triggers, and metadata -- task result payloads with task id, node name, error, interrupts, and writes -- step number and timestamp for every debug item - -## Thread-Level Observability - -Server and SDK use cases need thread-level observation in addition to run-level -streams. The graph layer should make it possible to subscribe to: - -- run modes plus run completion events -- lifecycle only -- state updates only - -This is separate from a single graph run stream. Thread streams watch a durable -thread as runs are created, resumed, interrupted, updated, or completed. - -## Redaction And Export - -Observation sinks must support redaction at export boundaries. Internal graph -events may include full state updates and task inputs; persistent or remote -observers may need to redact: - -- state values -- model messages -- tool arguments/results -- store keys -- checkpoint metadata -- provider payloads -- user identifiers -- secrets and credentials - -Redaction must preserve structural fields such as ids, node names, step numbers, -event kinds, timings, and checkpoint coordinates so traces remain useful. diff --git a/docs/modules/graph/package.md b/docs/modules/graph/package.md deleted file mode 100644 index 7b8847a82..000000000 --- a/docs/modules/graph/package.md +++ /dev/null @@ -1,66 +0,0 @@ -# Graph Package And Core Types - -## Package Shape - -Package layout: - -```text -crates/tinyagents-graph/src/ - lib.rs - builder/ - channel/ - checkpoint/ - command/ - compiled/ - parallel/ - recursion/ - reducer/ - stream/ - subagent_node/ - subgraph/ - testkit/ -``` - -## Core Types - -```rust -pub struct GraphBuilder<State, Ctx = (), Input = State, Output = State> { - graph_id: GraphId, - nodes: IndexMap<NodeId, NodeSpec<State, Ctx>>, - edges: EdgeSet, - branches: BranchSet<State, Ctx>, - channels: ChannelSet<State>, - input_schema: SchemaRef<Input>, - output_schema: SchemaRef<Output>, - defaults: GraphDefaults, -} - -pub struct CompiledGraph<State, Ctx = (), Input = State, Output = State> { - graph_id: GraphId, - nodes: Arc<IndexMap<NodeId, CompiledNode<State, Ctx>>>, - edges: Arc<EdgeSet>, - branches: Arc<BranchSet<State, Ctx>>, - channels: Arc<ChannelSet<State>>, - input_channels: ChannelSelection, - output_channels: ChannelSelection, - defaults: GraphDefaults, -} - -pub struct GraphRun<Output> { - pub run_id: RunId, - pub thread_id: Option<ThreadId>, - pub checkpoint_id: Option<CheckpointId>, - pub output: Output, - pub interrupts: Vec<Interrupt>, - pub visited: Vec<NodeId>, - pub steps: usize, - pub max_depth: usize, -} -``` - -The builder is mutable and ergonomic. The compiled graph is immutable, -validated, cheap to clone, and safe to run concurrently. - -`State`, `Input`, and `Output` should be separate generic concepts. Many real -graphs accept a narrow input shape, maintain richer internal state, and expose a -filtered output shape. diff --git a/docs/modules/graph/parallel-agents-forking.md b/docs/modules/graph/parallel-agents-forking.md deleted file mode 100644 index e39b2598b..000000000 --- a/docs/modules/graph/parallel-agents-forking.md +++ /dev/null @@ -1,347 +0,0 @@ -# Graph Parallel Agents And Context Forking - -Parallel agents are a graph pattern where one supervisor node starts multiple -child agents in the same superstep, lets them work from the same committed -state, and merges their outputs through normal graph reducers. - -Context forking is the runtime mechanism that makes this safe. A fork creates a -child task context from the current graph context. The child receives shared -read-only runtime dependencies, including cache and stores, while getting its -own task id, node id, event namespace, cancellation scope, and mutable scratch -space. - -This pattern is intended for workflows such as: - -- run two specialist sub-agents against the same user request -- compare two strategies and merge the better result -- dispatch research and verification agents in parallel -- fan out over multiple tools or retrievers while reusing cached model/tool - calls -- run a critique agent and a builder agent from the same initial context - -## Target Shape - -```rust -pub struct ContextForkOptions { - pub node: NodeId, - pub task_name: Option<String>, - pub namespace_label: Option<String>, - pub input: serde_json::Value, - pub cache: CacheForkPolicy, - pub stores: StoreForkPolicy, - pub timeout: Option<TimeoutPolicy>, - pub metadata: serde_json::Value, -} - -pub enum CacheForkPolicy { - InheritReadWrite, - InheritReadOnly, - Isolated, -} - -pub enum StoreForkPolicy { - Inherit, - IsolatedNamespace(StoreNamespace), -} -``` - -Default behavior should be `InheritReadWrite` for cache and `Inherit` for -stores. That lets parallel child agents reuse expensive cached model/tool -results while still writing their own child-run events and outputs. - -## Parallel Sub-Agent Fanout - -A supervisor node can spawn two sub-agents by returning `Send` packets or a -command with multiple destinations: - -```rust -Command::new() - .goto([ - Send::new("research_agent", json!({ "question": state.question })), - Send::new("critic_agent", json!({ "question": state.question })), - ]) -``` - -Each child agent runs as a separate graph task: - -- unique child `task_id` -- unique child `run_id` -- shared `root_run_id` -- parent run id set to the supervisor task run -- namespace such as `supervisor:<task_id>/research_agent:<task_id>` -- inherited `thread_id` -- inherited checkpointer according to subgraph/sub-agent policy -- inherited cache handle according to `CacheForkPolicy` -- inherited stores according to `StoreForkPolicy` - -The children must not mutate the parent state directly. They return writes that -are merged at the superstep boundary by channels/reducers. - -## Forked Context Contract - -Forking a context copies identity and dependency handles deliberately: - -Shared across forks: - -- root run id -- thread id -- immutable user context -- store registry unless isolated by policy -- cache handle unless isolated by policy -- event sink -- stream sink -- cancellation parent -- budget counters when policy says children share budget - -Unique per fork: - -- task id -- node id -- child run id -- checkpoint namespace suffix -- event namespace -- recursion frame -- mutable scratchpad -- task-local resume values -- task-local pending writes - -The fork must be cheap. It should clone `Arc` handles for shared services rather -than copying stores, cache contents, or large state snapshots. - -## Shared Cache Semantics - -Parallel agents should be able to reuse cache entries produced before the fork -and entries produced by sibling forks when durability policy allows it. - -Rules: - -- cache keys must include the called component, normalized input, relevant - config, provider/model id, tool version, and cache namespace -- cache keys must not include child task ids unless the cached result is - intentionally task-local -- child cache writes are visible to sibling forks only after the cache backend - commits them -- cache hits are emitted as task events with `TaskCached` -- cached task writes are replayed through graph writes, not injected as final - parent state -- cache failures do not fail the child task unless cache is marked required - -This allows two sub-agents to share expensive retrieval, embedding, prompt, -model, or tool results without sharing mutable state. - -## Merge Semantics - -Parallel child outputs merge through state channels: - -- independent fields use `LastValue` only when a single writer is guaranteed -- lists use append or topic reducers -- ranked candidates use custom merge reducers -- chat messages use message merge by id -- conflicting writes fail with an invalid concurrent update unless a reducer - defines deterministic resolution - -Example: - -```rust -pub struct AgentFanoutUpdate { - pub candidates: Vec<CandidateAnswer>, - pub critiques: Vec<Critique>, - pub usage: UsageDelta, -} -``` - -The parent graph should have reducers that append `candidates`, append -`critiques`, and aggregate `usage`. - -## Checkpointing - -Forked child tasks participate in normal checkpointing: - -- task start appears in checkpoint task metadata -- completed child writes can be persisted as pending writes -- child checkpoints include namespace and parent checkpoint config - -"Failed sibling tasks do not force successful child agents to rerun once -pending writes are saved" and "resuming from interrupt restarts the -interrupted child task, not unrelated completed siblings" hold today (fixed -per `code-review-graph.md` Critical C1/C2; see -`crates/tinyagents-graph/src/compiled/step.rs::fold_step` and -`crates/tinyagents-graph/src/compiled/boundary.rs::advance`). When a parallel -step interrupts or fails at branch index `i`, every branch that completed — -regardless of whether its index is above or below `i` — is folded into -committed state and is **not** re-run on resume/retry; only the -interrupted/failed branch(es) become the resumed run's pending set. A -completed branch's own routing is deferred (not resolved at the interrupt/ -failure boundary itself) rather than dropped: it is carried forward in the -checkpoint's `completed_tasks` and routed together with the rest of that -step's results once the pending branches finish, so a downstream node sees -the same merged state an uninterrupted run would have produced. One caveat: a -carried-forward branch's routing is re-resolved via static/conditional edges -only — an explicit `Command::goto` it returned is not itself persisted across -the boundary. Regression coverage: -`higher_index_completed_sibling_not_rerun_after_interrupt_then_resume`, -`higher_index_completed_sibling_not_rerun_after_failure_then_retry`, and -`interrupted_and_uninterrupted_runs_reach_the_same_state` -(`crates/tinyagents-graph/src/compiled/mod_tests.rs`). - -If a forked sub-agent interrupts, the parent run should surface the interrupt -with enough namespace information to resume the correct child. - -## Observability - -Parallel agents must be visible as distinct child runs, not hidden futures. - -Required events: - -- `ContextForked` -- `SubAgentStarted` -- `TaskStarted` -- `TaskCached` -- `TaskCompleted` -- `TaskFailed` -- `SubAgentCompleted` -- `ContextForkJoined` -- `StateUpdated` - -Every event must include: - -- root run id -- parent run id -- child run id -- task id -- node id -- namespace -- recursion depth -- checkpoint id when available - -UIs should be able to render the supervisor task with two child sub-agent lanes, -their streamed messages, cache hits, writes, and final reducer merge. - -## Failure And Cancellation - -Default policy: - -- required child failure fails the superstep -- optional child failure records an error update and lets siblings complete -- parent cancellation cancels all children -- child timeout fails only that child unless policy escalates -- successful sibling writes can be preserved as pending writes - -Future policies: - -- race: first successful child wins and cancels siblings -- quorum: continue when N of M children succeed -- compare: require all children, then run a judge node -- fallback: run secondary agent only if primary fails - -These should be explicit policies, not implicit behavior hidden in node code. - -## Reusable map/reduce helper {#map-reduce} - -`graph::parallel::map_reduce` is a standalone helper — independent of the graph -executor — for the common "run N items concurrently and reduce the results" -pattern. Where `Send` fanout is the low-level superstep primitive, `map_reduce` -gives a reusable async function with deterministic input-order results, a -concurrency cap, per-item success/failure isolation, and a configurable failure -policy: - -```rust -pub async fn map_reduce<I, T, F, Fut>( - items: Vec<I>, - options: ParallelOptions, - f: F, -) -> Result<ParallelOutcome<T>> -where - F: Fn(usize, I) -> Fut, - Fut: Future<Output = Result<T>>; -``` - -`f` receives each item's 0-based input index and the item. Results are collected -into a `ParallelOutcome<T>` of per-item `ItemOutcome<T>` values **in input order** -regardless of completion order (concurrency is driven by an ordered buffered -stream). `ParallelOutcome` exposes `success_count`, `failure_count`, `successes`, -and `into_successes`. - -### Failure policy - -`FailurePolicy` decides what happens when items fail: - -- `FailFast` — stop at the first failing item (in input order) and return its - error; remaining in-flight work is dropped. -- `CollectAll` (default) — run every item to completion, always return `Ok`, - recording per-item success/failure for the caller to inspect. -- `Quorum(n)` — run every item; return `TinyAgentsError::Graph` unless at least - `n` items succeeded. -- `BestEffort` — always `Ok`, silently keeping only the successful outputs. - -### Execution controls - -`ParallelOptions` is a builder. It is `Clone` (no longer `Copy`) because it can -hold a cancellation token. - -| Builder | Effect | -| --- | --- | -| `with_max_concurrency(n)` | Cap simultaneous items (`0` = unbounded). | -| `with_failure_policy(p)` | Set the `FailurePolicy`. | -| `with_item_timeout(dur)` | Per-item wall-clock bound; a slow item becomes a recoverable per-item failure (then handled by the policy) rather than stalling the batch. | -| `with_total_timeout(dur)` | Overall wall-clock bound; on elapse the call returns `TinyAgentsError::Timeout` and drops remaining in-flight work. | -| `with_cancellation(token)` | Attach a cooperative `CancellationToken`; when cancelled the call returns `TinyAgentsError::Cancelled` and drops in-flight work. | - -Per-item timeout — a slow item fails in isolation while fast items still -succeed: - -```rust -use tinyagents_graph::parallel::{map_reduce, FailurePolicy, ParallelOptions}; -use tinyagents_harness::TinyAgentsError; - -// Item 1 hangs; a 50ms per-item timeout turns it into a recoverable failure. -let out = map_reduce( - vec![0u64, 3_600_000, 0], - ParallelOptions::default() - .with_failure_policy(FailurePolicy::CollectAll) - .with_item_timeout(std::time::Duration::from_millis(50)), - |_i, ms| async move { - tokio::time::sleep(std::time::Duration::from_millis(ms)).await; - Ok::<_, TinyAgentsError>(ms) - }, -) -.await?; -assert_eq!(out.success_count(), 2); -assert_eq!(out.failure_count(), 1); // the hanging item timed out -``` - -Total timeout — the whole batch aborts and drops in-flight work: - -```rust -let err = map_reduce( - vec![3_600_000u64], - ParallelOptions::default().with_total_timeout(std::time::Duration::from_millis(20)), - |_i, ms| async move { - tokio::time::sleep(std::time::Duration::from_millis(ms)).await; - Ok::<_, TinyAgentsError>(ms) - }, -) -.await -.expect_err("total timeout aborts the batch"); -assert!(matches!(err, TinyAgentsError::Timeout(_))); -``` - -Cancellation — a cancelled token stops the batch: - -```rust -use tinyagents_harness::CancellationToken; - -let token = CancellationToken::new(); -token.cancel(); -let err = map_reduce( - vec![0u64], - ParallelOptions::default().with_cancellation(token), - |_i, _n| async move { - tokio::time::sleep(std::time::Duration::from_secs(3600)).await; - Ok::<_, TinyAgentsError>(0u64) - }, -) -.await -.expect_err("a cancelled token aborts"); -assert!(matches!(err, TinyAgentsError::Cancelled)); -``` diff --git a/docs/modules/graph/routing.md b/docs/modules/graph/routing.md deleted file mode 100644 index a23e7980e..000000000 --- a/docs/modules/graph/routing.md +++ /dev/null @@ -1,158 +0,0 @@ -# Graph Edges, Routing, Commands, And Sends - -Reserved virtual nodes: - -- `START` -- `END` - -Direct edge: - -```text -START -> agent -agent -> summarize -summarize -> END -``` - -Conditional edge: - -```text -agent --tool--> tools -agent --final--> END -tools ---------> agent -``` - -Waiting/barrier edge: - -```text -retrieve_docs --\ -lookup_user ----> synthesize -score_risk --/ -``` - -Command routing: - -```rust -Command::new() - .update(update) - .goto(["tools"]) -``` - -Router closures may return any `impl ToString` (plain `&str`/`String`, or a -user-defined enum implementing `Display`), for example: - -```rust -enum AgentRoute { - Tool, - Final, -} -``` - -Internally, route labels are represented as a `Route` newtype (with -`From<Route> for String` / `From<String> for Route` / `From<&str> for -Route`), so `RouterFn`'s public signature stays `Fn(&State) -> R where R: -ToString` — existing closures returning `String` keep compiling unchanged. -`GraphBuilder::add_conditional_edges_checked` additionally takes the -exhaustive set of labels a router can return, so `compile()` can validate -every branch label against the node's route table before the graph ever -runs (see `builder.md`). - -Routing outputs: - -- node name -- `END` -- one or more node names -- `Send` packet -- one or more `Send` packets -- `Command` - -Branch metadata should include an optional path map and typed route labels so -visualization does not have to assume every conditional edge can reach every -node. - -## Commands And Send Packets - -Commands combine state update, routing, parent/subgraph targeting, and interrupt -resume values. - -```rust -pub struct Command { - pub graph: CommandGraphTarget, - pub update: Option<StateUpdate>, - pub resume: Option<ResumeValue>, - pub goto: Vec<RouteTarget>, -} - -pub enum CommandGraphTarget { - Current, - Parent, - Graph(GraphId), -} - -pub enum RouteTarget { - Node(NodeId), - Send(Send), -} - -pub struct Send { - pub node: NodeId, - pub arg: serde_json::Value, -} -``` - -`Command::goto([..])` (plain node activations) and `Command::send([Send::new(node, -arg), ..])` (per-invocation fanout) both populate `Command::goto: -Vec<RouteTarget>`; `with_goto`/`with_sends` append to it. A `Send`-scheduled node -receives its `arg` on `NodeContext::send_arg` (`None` for normal activations); -many `Send`s may target the *same* node, producing one parallel activation each -(map-reduce). Plain activations are deduplicated by node; `Send` activations are -not. - -Runtime command targets are validated before they become the next active set or -are persisted in a checkpoint: - -- `END` is allowed and terminates that branch. -- `START` is rejected because it is a virtual entry marker, not an executable - runtime target. -- every other target must name a compiled node. -- invalid runtime targets fail the run before the boundary checkpoint is - written, so durable state cannot be poisoned with missing next nodes. - -Use `Command` for: - -- dynamic routing -- node-local state update plus routing -- human approval resume values -- parent graph handoff from subgraphs -- supervisor/worker handoff - -Use `Send` for dynamic fanout where each target invocation receives custom -input that can differ from the graph's main state. This is the primitive for -map-reduce, search fanout, parallel tool calls, and per-item scoring. - -## External Run Inputs - -Most runs enter through the compiled `START -> entry` edge: - -```rust -graph.run(state).await?; -``` - -When the caller needs to seed multiple graph loops at once, use -`GraphInput`: - -```rust -graph - .run_with_inputs( - state, - [ - GraphInput::start(json!({ "message": "hello" })), - GraphInput::new("tool_loop", json!({ "tool": "search" })), - ], - ) - .await?; -``` - -`GraphInput::start(..)` resolves to the compiled entry node and delivers its -payload through `NodeContext::send_arg`. `GraphInput::new(node, ..)` targets a -real node directly in the first superstep. Inputs are not deduplicated, so two -inputs aimed at the same node run two activations with distinct payloads. diff --git a/docs/modules/graph/runtime-policy.md b/docs/modules/graph/runtime-policy.md deleted file mode 100644 index a8f2d888d..000000000 --- a/docs/modules/graph/runtime-policy.md +++ /dev/null @@ -1,154 +0,0 @@ -# Graph Runtime Context, Node Defaults, And Policies - -This page describes the per-task runtime context and the per-node execution -policy surface as implemented in `crates/tinyagents-graph/src/builder/` -(`types.rs`, `policy.rs`) and `crates/tinyagents-graph/src/compiled/` -(`run_ctx.rs`, `step.rs`, `boundary.rs`). It replaces an older sketch of a -`GraphContext<Ctx>` / `TimeoutPolicy` / `error_handler: NodeId` design that -was never built; nothing below is a target. - -## `NodeContext` - -Every node handler receives `(State, NodeContext)`. The context is created -per activation by the executor (`RunCtx::node_context`), scoped to that one -task, and never a global. It exposes: - -- identity: `graph_id`, `node_id`, `run_id`, `thread_id` (when - checkpointing), `root_run_id`, `recursion_frames` (the live recursion - stack, root-first), `task_id` (stable within the superstep; distinguishes - `Send` fan-out activations of one node), `siblings` (how many activations - of this node share the step), and `fork` (the branch in a parallel step) -- `step`: the 1-based superstep number -- `resume`: the value a `Command::resume`/`resume_tasks` delivered to this - task, `None` otherwise -- `send_arg`: the per-invocation argument of a `Send` packet or `GraphInput` -- `channel_versions` / `versions_seen` and `changed_since_last_run(channel)` - for channel-model graphs -- `child_runs`: the sink a subgraph node reports spawned child runs to -- `agent_binding`: the host-supplied recursive-agent binding for this run -- `heartbeat()`: restarts the node's idle-timeout window -- `durable_task(key, fut)`: memoises a side-effecting sub-step per - `(task_id, key)` across re-runs of the task — see - [builder.md](builder.md#nodecontextdurable_task) - -Deadline and cancellation are *not* on the context. The whole-run deadline -(`CompiledGraph::with_run_deadline`) and the run-level stop signals -(`RunOptions::cancellation`, `RunOptions::drain`) are enforced by the -executor between and around supersteps (see -[execution.md](execution.md) and [fault-tolerance.md](fault-tolerance.md)). -There is no store registry or custom stream writer on the context; events -flow through the graph's `GraphEventSink`. - -## Node policies - -```rust -pub struct NodePolicy<State, Update> { - pub retry: Option<RetryPolicy>, - pub timeout: Option<Duration>, - pub idle_timeout: Option<Duration>, - pub cache: Option<NodeCachePolicy<State>>, - pub on_error: Option<Arc<OnErrorFn<State, Update>>>, - pub defer: bool, -} - -pub struct NodeCachePolicy<State> { - pub key: Arc<CacheKeyFn<State>>, // Fn(&State, Option<&send_arg>) -> String - pub ttl: Option<Duration>, -} -``` - -Attach a policy per node with `GraphBuilder::with_node_policy(node, policy)` -or graph-wide with `GraphBuilder::set_node_defaults(policy)`. Every field is -optional and additive; `NodePolicy::default()` changes nothing. Builder -sugar: `NodePolicy::default().with_retry(..).with_timeout(..) -.with_idle_timeout(..).with_cache(..).with_on_error(..).deferred()`. - -### Resolution - -The executor resolves an *effective* policy per activation -(`CompiledGraph::effective_policy` → `NodePolicy::resolve`), field by field: - -1. the node's own policy, when that field is `Some`; -2. else the `set_node_defaults` policy's field; -3. else the legacy graph-wide setting — `CompiledGraph::with_node_retry` - for `retry`, `GraphBuilder::with_node_timeout` / `GraphDefaults - { node_timeout }` for `timeout`; -4. else nothing. - -`defer` is the logical OR of the per-node and default flags. Per-node values -therefore always win; graph-level values only fill gaps. - -### Retry - -A handler that fails with a retryable error (`tinyagents_harness::retry::is_retryable` -— the transient model/tool class) is re-run from its start, with a fresh -handler future and a cloned context, up to the policy's attempt cap, -emitting `GraphEvent::NodeRetryScheduled` and sleeping the backoff only when -the policy opts in (`RetryPolicy::with_backoff_sleep`). A cloned context -shares the task's `durable_task` memos, so a retried attempt does not repeat -a side effect the first attempt already memoised. Non-retryable errors and -exhausted budgets escalate to the failure boundary (a resumable checkpoint -on a checkpointed thread; see fault-tolerance.md). - -### Timeouts - -Two independent ceilings race every attempt: `timeout` (flat wall-clock time -per attempt) and `idle_timeout` (the maximum gap between two -`NodeContext::heartbeat()` calls, or between start and the first one). Either -firing fails the attempt with `TinyAgentsError::Timeout`, which is not -retryable. Timeout cancellation is cooperative in the sense that the handler -future is dropped at the race, not pre-empted mid-instruction. There is no -"refresh on graph progress" mode; only an explicit heartbeat re-arms the -idle window. - -### `on_error` - -`on_error: Fn(&State, &TinyAgentsError) -> Option<Command<Update>>` is -consulted once retries are exhausted (or the error is not retryable). -Returning `Some(command)` makes the node complete with that command — its -update applied, its `goto` honoured — instead of failing the run. There is no -separate "error-handler node" concept; a policy that wants one routes to it -via the returned `Command::goto`. - -### Caching - -Caching is a two-part opt-in: - -- `NodeCachePolicy` on the node (via `NodePolicy::cache`, or directly with - `CompiledGraph::with_cached_node(node, policy)`, which also installs the - `Update: Serialize + DeserializeOwned` codec) derives a key from the state - snapshot and the activation's `send_arg`, with an optional TTL; -- a backend on the graph: `CompiledGraph::with_task_cache(Arc<dyn TaskCache>)` - (`InMemoryTaskCache`, or `DriverTaskCache`). - -Without a backend every cache policy is inert. With one, a hit replays the -stored `Update` without invoking the handler and emits -`GraphEvent::TaskCompleted { cached: true }`; a miss runs the handler and, -on an `Update`/`Command` result, stores it (`TaskCompleted { cached: false }`). -Errors and interrupts are never cached. The key function is called once per -activation. The cache key is entirely the policy's to compute — include -whatever inputs, config version, or namespace make two activations -interchangeable — and keys are scoped by `(graph_id, node_id, hash)`. Cache -errors are treated as misses; caching is an optimisation, never a -correctness requirement. This is distinct from `NodeContext::durable_task`, -which is per-task memoisation inside one thread's checkpoint ledger, not a -cross-run cache. - -### Defer - -`defer: true` (or `GraphBuilder::mark_deferred(node)`, which also sets the -export marker) holds an activation back from the frontier while any -non-deferred activation is also ready, accumulating held activations across -supersteps and releasing them all at once the first time a boundary would -otherwise route to an empty frontier — a "run once everything else is done" -synthesis join without an explicit barrier (`boundary::apply_defer`). Held -deferred activations are not persisted: a run resumed mid-hold starts with -an empty hold. - -## Graph defaults - -`GraphBuilder::set_defaults(GraphDefaults { .. })` applies only the `Some` -fields of `recursion_limit`, `parallel`, `max_concurrency`, and -`node_timeout`; `with_max_concurrency(n)` bounds in-flight handlers per -parallel superstep (`0` = unbounded); `with_node_timeout(d)` is the legacy -graph-wide flat timeout that `NodePolicy::resolve` falls back to. diff --git a/docs/modules/graph/state-channels.md b/docs/modules/graph/state-channels.md deleted file mode 100644 index a8c788cce..000000000 --- a/docs/modules/graph/state-channels.md +++ /dev/null @@ -1,211 +0,0 @@ -# Graph State, Channels, And Updates - -The current scaffold returns whole state from each node. Durable graphs should -move to partial state updates applied through channels. A channel owns the -current value, the accepted update type, checkpoint representation, reducer, and -step-boundary merge behavior. - -```rust -pub trait Channel: Send + Sync { - type Value; - type Update; - type Checkpoint; - - fn get(&self) -> Result<Self::Value>; - fn update(&mut self, values: Vec<Self::Update>) -> Result<ChannelChange>; - fn checkpoint(&self) -> Result<Self::Checkpoint>; - fn restore(checkpoint: Self::Checkpoint) -> Result<Self> - where - Self: Sized; - fn consume(&mut self) -> Result<ChannelChange> { - Ok(ChannelChange::Unchanged) - } - fn finish(&mut self) -> Result<ChannelChange> { - Ok(ChannelChange::Unchanged) - } -} -``` - -Required channel policies: - -- `LastValue`: accepts one update per step and overwrites the value. -- `Overwrite`: explicit overwrite marker for aggregate channels. -- `BinaryAggregate`: applies a binary reducer such as append, add, min, or max. -- `Topic`: pub/sub collection, optionally accumulating across steps. -- `Ephemeral`: value exists only for one step or one trigger. -- `Barrier`: waits until named sources have all arrived. -- `NamedBarrier`: tracks named arrivals for join semantics. -- `Messages`: message merge by id for chat histories. -- `Delta`: stores compact deltas plus periodic snapshots for large append-heavy - channels. -- `Untracked`: excluded from checkpointing when safe. - -Why channel-level reducers matter: - -- parallel branches can write different fields safely -- map-reduce fanout can aggregate many outputs in one step -- checkpoints can store pending writes instead of only final whole-state values -- failed parallel nodes can rerun without discarding completed writes -- tests can assert exact writes and reducer behavior -- generated/language-defined nodes can use simple partial update contracts - -The graph should support both root state and object state. A single root channel -is useful for scalar workflows; multi-key state is the default for agent graphs. - -## Implemented additive model (`graph::channel`) - -The channel model is shipped **additively**: the monolithic `State` + -`StateReducer` path is unchanged, and channels are an opt-in alternative that -runs on the *existing* executor. The implementation lives in -`crates/tinyagents-graph/src/channel/` and is `serde_json::Value`-backed for generality. - -- `Channel` (object-safe trait): per-key `merge(current, incoming) -> value` - plus `allows_concurrent`, `is_ephemeral`, `is_tracked`, and `is_ready` - (barrier) hooks. Concrete channels: `LastValue` (overwrite), `Topic` - (append into an array), `Delta` (numeric accumulate), `Messages` (merge by - `id`), `Ephemeral` (overwrite, cleared next step), `Untracked` (overwrite, - excluded from snapshots), `Barrier`/`NamedBarrier` (count/name fan-in with - readiness), and `BinaryAggregate` (fold via a closure or any - `Reducer<Value>`). -- `ChannelSet`: a named map of `Box<dyn Channel>` plus their current values, - with `add_channel`/`with_channel`, `apply_update(name, value)`, `get`, - `is_ready`, and `snapshot()` (the tracked, durable view). -- `ChannelState`: a graph `State` wrapping a `ChannelSet`. It implements - `StateReducer<ChannelState, ChannelUpdate>` for itself (the `&self` reducer - receiver is unused; merge rules travel inside the running state), so a - channel graph is built directly with - `GraphBuilder::<ChannelState, ChannelUpdate>::new().set_reducer(ChannelState::new())`. -- `ChannelUpdate`: a batch of `(name, value)` writes a node returns - (`ChannelUpdate::new().set(..).set(..)`). Stamp it with the producing node's - superstep via `.at_step(ctx.step)`. - -### Concurrent-write conflict detection - -The executor folds a step's branch updates one at a time. Stamping each -`ChannelUpdate` with `ctx.step` lets the reducer group a step's writes: when the -step number advances it resets its per-step bookkeeping and clears `Ephemeral` -channels. A second same-step write to a non-aggregate channel (`LastValue`, -`Ephemeral`, `Untracked` — `allows_concurrent == false`) raises -`TinyAgentsError::InvalidConcurrentUpdate`; aggregate channels -(`allows_concurrent == true`) merge both writes in deterministic active-set -index order. Cross-step overwrites and repeated writes inside one update are -last-wins, not conflicts. Unstamped updates are treated as independent steps -(no conflict detection, no ephemeral clearing), preserving the simplest path. - -## Serializable channels - -`Channel` carries a `config(&self) -> serde_json::Value` hook (default -`Value::Null`) alongside `kind()`. A `ChannelSet` serializes as an ordered map -of `{ kind, config, value }` entries — no `dyn Channel` trait objects on the -wire — and its `Deserialize` impl reconstructs each channel from `kind` + -`config` via an internal `channel_from_config` dispatcher, with no external -context. `ChannelState` derives `Serialize`/`DeserializeOwned` on top of that, -so `Checkpoint<ChannelState>` satisfies every bundled `Checkpointer`'s durable -bounds and round-trips through `DriverCheckpointer` like any -other graph state. - -Built-in channels need no config beyond `kind` (`LastValue`, `Topic`, `Delta`, -`Messages`, `Ephemeral`, `Untracked`), or a small literal payload -(`Barrier`/`NamedBarrier` persist their `expected` set). `BinaryAggregate` is -the one exception: its merge rule is a closure, which cannot serialize. Instead -of persisting the closure, a `BinaryAggregate` channel persists a **reducer -name**, resolved through a process-wide `ReducerRegistry`: - -```rust -// Register once, anywhere before the checkpoint is decoded (typically at -// startup, alongside the graph builder that will use it): -let graph = GraphBuilder::<ChannelState, ChannelUpdate>::new() - .register_reducer("double", |current: Value, incoming: Value| { - Ok(json!(current.as_i64().unwrap() * incoming.as_i64().unwrap())) - }) - // ... - .set_reducer(ChannelState::new()); - -// Build the channel from the registered name: -let initial = ChannelState::new() - .with_channel("product", BinaryAggregate::named("double")?); -``` - -`BinaryAggregate::named(name)` looks `name` up in the registry immediately -(returning `TinyAgentsError::Checkpoint` if it is not registered yet) and -records `name` in the channel's `config()`. Decoding a `binary_aggregate` -channel repeats that lookup against whatever the *decoding* process has -registered; a name nobody registered decodes to -`TinyAgentsError::Checkpoint("unknown reducer ...")` rather than silently -losing the merge rule. `GraphBuilder::register_reducer` is the intended entry -point — it just delegates to `ReducerRegistry::register`, and the registry is -global (not scoped to one builder) because a checkpoint decode has no builder -in scope at all. - -The built-ins `"append"`, `"last"`, `"sum"`, `"max"`, `"min"`, and -`"set_union"` are always registered, needing no call to `register_reducer`. -`BinaryAggregate::new`/`from_reducer` (a bare closure, no name) still work for -a channel that is never checkpointed — they just cannot round-trip. - -## Channel versions - -Every `ChannelState` tracks a cumulative per-channel version counter, -`channel_versions: BTreeMap<String, u64>`, bumped once per distinct channel -name touched by each folded `ChannelUpdate` (a channel written twice in one -update still bumps once). It is exposed via `ChannelState::channel_versions()` -and persisted on `Checkpoint::channel_versions` at every boundary — the -executor's normal/failure/cancel boundaries and the manual `update_state` -write all go through the same `channel::channel_bookkeeping` extraction -function, so they cannot disagree about what they persist. A plain whole-state -graph (any `State` that is not `ChannelState`) reports a single `"state"` -channel, bumped once per checkpoint, so `channel_versions` is always populated -regardless of which state model a graph uses. - -`NodeContext` carries two views built from this: `channel_versions` (what the -node's invocation currently observes) and `versions_seen` (this node's own -snapshot from the last time it ran — empty on its first-ever activation). -`NodeContext::changed_since_last_run(channel)` compares the two, so a node can -skip work when nothing it depends on has changed: - -```rust -.add_node("summarize", |state: ChannelState, ctx: NodeContext| async move { - if !ctx.changed_since_last_run("messages") { - return Ok(NodeResult::Update(ChannelUpdate::new())); - } - // ... recompute the summary ... -}) -``` - -The per-node snapshot is itself persisted (`Checkpoint::versions_seen`, keyed -by node id string — `NodeId` has no `Ord` impl to key a `BTreeMap` on -directly) and restored on resume, so this bookkeeping survives a restart -instead of resetting. - -## Delta-channel history and `Overwrite` - -`ChannelSet::with_delta(name, snapshot_every)` marks an already-registered -append-style channel (typically `Topic`, or an `"append"`/`"set_union"` -`BinaryAggregate`) for **delta tracking**: every write to `name` also records -its raw incoming value into `ChannelState::step_deltas()` for the current -step, which the checkpoint-construction call sites copy into -`Checkpoint::channel_deltas`. Unlike `channel_versions`/the channel's own -`value`, this field is **not cumulative** — each checkpoint carries only its -own step's writes to a delta-tracked channel, which is what keeps a single -checkpoint's size bounded no matter how long the channel's full value grows -(a 200-step append thread's checkpoint bytes grow roughly linearly with step -count, not quadratically). `Checkpointer::delta_history(config, channel)` -replays the full per-step write sequence for a delta-tracked channel; the -default implementation walks `state_history` oldest-first and concatenates -each checkpoint's own `channel_deltas` entry. - -`ChannelUpdate::overwrite(name, value)` (backed by `ChannelWrite::Overwrite`) -bypasses the channel's merge rule entirely and replaces its value outright. -This is what lets a delta/append channel reset its baseline: the overwritten -value becomes what subsequent `.set(name, ..)` merges build on, and it also -rebases the channel's `step_deltas` entry (prior accumulated deltas for that -step are cleared before the overwrite's own value is recorded), so a consumer -walking `delta_history` sees the overwrite as a fresh starting point rather -than an append onto stale history. - -`update_state` and `fork_state` never bypass this: `update_state`'s manual -write folds through the exact same `ChannelState::merge` → `ChannelSet::apply_ -channel_write` dispatch the executor boundary uses (`apply_channel_write` is -the single write-path function every channel-graph write funnels through), and -`fork_state` copies a source checkpoint's `channel_versions`/`channel_deltas`/ -`versions_seen` verbatim rather than re-deriving them — so replay and a manual -write can never diverge on what a checkpoint records. diff --git a/docs/modules/graph/storage-schema.md b/docs/modules/graph/storage-schema.md deleted file mode 100644 index 818375280..000000000 --- a/docs/modules/graph/storage-schema.md +++ /dev/null @@ -1,48 +0,0 @@ -# Graph and tasks storage schema v2 - -`tinyagents_graph::storage_schema()` (owner `graph`) and -`tinyagents_tasks::storage_schema()` (owner `tasks`), both scope class Agent. - -## Graph - -One checkpointer. The old collection prefix becomes the `ns` body field, part of -every id and index. - -| v1 | v2 | -|---|---| -| `graph_*`, `delegation_graph_*`, `flows_graph_*` `{checkpoints,threads,writes}` | `graph.{checkpoints,threads,writes}`, `ns` = `graph`, `delegation`, `flows` | -| checkpoint doc `{thread, namespace, seq, checkpoint_id, record}` | `{ns, thread, subgraph, seq, checkpoint_id, checkpoint}` under `Key::of([ns, thread, seq])` | -| `<prefix>_leases` `{owner, expires_at_ms}` | `tsd.leases`, owner `graph`, key `Key::of([ns, thread])` (epoch fence) | -| `graph_task_cache` `{graph, node, hash, value, expiresAtMs}` | `tsd.kv` owner `graph`, ns `task_cache:<graph>`, `expires_ms` metadata | -| goals, todos and the scratchpad (harness `Store`) | `tsd.kv` through the host's store, unchanged in code | - -A collection gets one step per legacy prefix and a step chain cannot fork, so -the steps chain `1 -> 2 -> 3 -> 4` (declared version 4). A migrated row carries -the version its step ends at (`graph` 2, `delegation` 3, `flows` 4); the body -shape is the same. - -`tsd.leases` and `tsd.kv` belong to `tsd`, so `open_storage(scoped)` calls -`import_legacy_leases` and `import_legacy_task_cache` after `Registry::migrate`. -Both record completion in `tsd.migrations`. `DriverCheckpointer::open(scoped, ns)` -runs `open_storage` first; the sync `new`/`with_namespace` do not. - -Breaking: `DriverCheckpointer::with_prefix` is now `with_namespace`, and -`DriverTaskCache::with_collection` is `with_namespace`. A task cache needs a -driver with record metadata when a TTL is used. - -## Tasks - -| v1 | v2 | -|---|---| -| `orchestration_tasks` | `tasks.records`, id `Key::of([task_id])` | -| `orchestration_task_history` | `tasks.history`, id `Key::of([task_id, rev])`, unique index `task_rev` | -| `orchestration_completions` | `tasks.completions`, id `Key::of([task_id])` | - -Bodies are bare snake_case DTOs with `*_ms` times; `OrchestrationTaskRecord` and -`CompletionRecord` (tool and RPC output types) are unchanged. The last-update -time is the document's `updated_ms` metadata; a migrated record's is the -migration time. Transitions use `Repo::update`. The history snapshot is written -after the record commit, not in the same batch. - -`DriverTaskStore::with_prefix`, `collection_specs(prefix)` and -`DriverCompletionStore::with_collection` are removed: the names are fixed. diff --git a/docs/modules/graph/streaming.md b/docs/modules/graph/streaming.md deleted file mode 100644 index ac8b91892..000000000 --- a/docs/modules/graph/streaming.md +++ /dev/null @@ -1,147 +0,0 @@ -# Graph Streaming And Events - -The graph stream should support low-level event streams and high-level -projection streams. - -Core stream modes: - -- `values`: full state values after each step -- `updates`: per-node/per-task state updates -- `messages`: harness message or token deltas emitted by model nodes -- `custom`: arbitrary user stream writes from inside nodes -- `checkpoints`: checkpoint payloads -- `tasks`: task start and task result payloads -- `debug`: checkpoints plus task internals -- `events`: all graph lifecycle events - -Typed stream part: - -```rust -pub enum StreamPart<State, Output> { - Values { - namespace: Vec<String>, - data: Output, - interrupts: Vec<Interrupt>, - }, - Updates { - namespace: Vec<String>, - data: IndexMap<NodeId, StateUpdate>, - }, - Messages { - namespace: Vec<String>, - message: Message, - metadata: StreamMetadata, - }, - Custom { - namespace: Vec<String>, - data: serde_json::Value, - }, - Checkpoint { - namespace: Vec<String>, - data: CheckpointPayload<State>, - }, - Tasks { - namespace: Vec<String>, - data: TaskStreamPayload, - }, - Debug { - namespace: Vec<String>, - data: DebugPayload<State>, - }, -} -``` - -Event stream: - -```rust -pub enum GraphEvent { - RunStarted { run_id: RunId, graph_id: GraphId }, - RunStreamingStarted { run_id: RunId }, - StepStarted { step: usize, active: Vec<NodeId> }, - TaskStarted { task_id: TaskId, node: NodeId, triggers: Vec<String> }, - TaskCompleted { task_id: TaskId, node: NodeId }, - TaskCached { task_id: TaskId, node: NodeId }, - TaskFailed { task_id: TaskId, node: NodeId, error: String }, - StateUpdated { node: NodeId, update: serde_json::Value }, - RouteSelected { node: NodeId, routes: Vec<RouteTarget> }, - ContextForked { parent_task_id: TaskId, child_task_id: TaskId }, - ContextForkJoined { parent_task_id: TaskId, child_task_id: TaskId }, - SubgraphStarted { node: NodeId, child_run_id: RunId, namespace: Vec<String> }, - SubgraphCompleted { node: NodeId, child_run_id: RunId }, - SubAgentStarted { node: NodeId, agent: ComponentId, child_run_id: RunId }, - SubAgentCompleted { node: NodeId, agent: ComponentId, child_run_id: RunId }, - RecursionDepthChanged { depth: usize }, - CheckpointSaved { checkpoint_id: CheckpointId }, - InterruptEmitted { interrupt: Interrupt }, - RunDraining { run_id: RunId, reason: String }, - RunCompleted { run_id: RunId }, - RunFailed { run_id: RunId, error: String }, - Custom { name: String, payload: serde_json::Value }, -} -``` - -Streaming requirements: - -- graph runs can be consumed as an async stream -- streaming does not require waiting for final state -- every streamed event carries run id, thread id, namespace, step, and node/task - metadata when available -- subgraph streams preserve nested namespaces -- harness streams from model/tool/sub-agent nodes are forwarded with graph node - context -- subscribers can filter graph events, harness events, sub-agent events, state - updates, task payloads, messages, and checkpoints -- a typed run stream should expose final output, interrupted status, and pending - interrupts even when the caller only subscribed to a subset of projections - -## Implementation status (runtime-comparison Phase 3, C3) - -The design above is the target shape; the current implementation is a -scoped-down but real subset built around `crate::stream`: - -- **`GraphEventEnvelope { run_id, task_id, ns, seq, event }`** - (`stream/types.rs`) wraps every `GraphEvent` the executor emits. - `GraphEventSink::emit` takes the envelope, not the bare event — `ns` is the - emitting graph instance's checkpoint namespace (empty at the top level, one - segment deeper per level of subgraph embedding) and `seq` is a per-instance - monotonic counter (`CompiledGraph::sequence`, an `Arc<AtomicU64>`). The - counter is shared across a clone that only swaps `event_sink` (journal - wrapping keeps counting where the plain run left off) but is reset to a - fresh one for a subgraph embedded as a node - (`subgraph::namespaced`/`CompiledGraph::with_fresh_sequence`) — the deeper - `ns` already disambiguates that stream, so nothing is gained by chaining - the parent's sequence into it. `task_id` is `None` until per-task - correlation ids exist end to end (`docs/runtime-comparison/feature-gaps.md` - D4); adding it later is additive. -- **`StreamMode::{Tasks, Checkpoints}`** are implemented. `GraphEvent::mode()` - maps `TaskScheduled`/`TaskStarted`/`TaskCompleted`/`NodeStarted`/ - `NodeCompleted`/`NodeFailed`/`NodeRetryScheduled` onto `Tasks`, - `CheckpointSaved`/`CheckpointRestored` onto `Checkpoints`; every other kind - (including the plain run/step lifecycle events) is debug-only detail. - `stream::project::project_graph_event(event, modes)` applies that mapping. - `GraphEvent::TaskStarted`/`TaskCompleted { cached }` are new variants - emitted alongside `NodeStarted`/`NodeCompleted`/`NodeFailed` at the same - boundary; `cached` is always `false` today (no per-node task cache exists - yet — D2). -- **`StreamProjection`** (`stream/project.rs`) is the cross-source fold this - doc's "harness streams … are forwarded with graph node context" bullet - calls for: `fold_graph_event`/`fold_agent_event` take a `GraphEventEnvelope` - and a harness `AgentEvent` respectively and append into `messages`, - `tool_calls`, and `subagents` — each a `Vec<Cursored<T>>` sharing one - monotonic `cursor` across all three views. `StreamProjection::since(cursor)` - is the late-attach replay primitive: a consumer that connects after a run - is already underway asks for everything past the cursor it last saw - instead of re-reading history. Only `GraphEvent::SubgraphStarted`/ - `SubgraphCompleted` project from the graph side today (as `subagents` - entries); the harness side covers model deltas - (`AgentEvent::ModelDelta` → `messages`), tool lifecycle - (`ToolStarted`/`ToolCompleted`/`ToolFailed` → `tool_calls`), and sub-agent - lifecycle (`SubAgentStarted`/`SubAgentCompleted` → `subagents`). -- **Not yet implemented**: the typed `StreamPart<State, Output>` enum, the - `values`/`updates`/`custom` projections from live graph state (those still - require the graph-state side channel this doc's header note already calls - out), `RunStreamingStarted`, `ContextForkJoined`, and `RunDraining`. A - subgraph node does not automatically inherit its parent's `event_sink` — - nested observability today requires configuring the same sink on both - explicitly (see `subgraph::test::nested_subgraph_run_yields_envelopes_with_correct_namespace_depth_and_seq` - for the pattern); full automatic propagation is future work alongside D4. diff --git a/docs/modules/graph/subagents-recursion.md b/docs/modules/graph/subagents-recursion.md deleted file mode 100644 index 344d44ca0..000000000 --- a/docs/modules/graph/subagents-recursion.md +++ /dev/null @@ -1,205 +0,0 @@ -# Graph Sub-Agents, Recursion, And Depth Tracking - -Sub-agents are harness agents registered in the registry and invoked from graph -nodes. They are not opaque side effects; they are child runs with their own -events, usage, cost, failures, and optional streams. - -> **Deprecated.** `SubAgentNode` / `subagent_node` are `#[deprecated(note = "use tinyagents_orchestration SubagentDriver")]`: graph cannot depend on orchestration, so they cannot adapt the driver. Behaviour is unchanged. - -```rust -pub struct SubAgentNode { - pub agent: ComponentId, - pub input_mapper: InputMapper, - pub output_mapper: OutputMapper, - pub policy: SubAgentPolicy, -} -``` - -Sub-agent requirements: - -- create a child `run_id` -- preserve `root_run_id` -- set `parent_run_id` to the graph node/task run -- forward harness events through graph/registry streaming -- apply timeout, retry, cache, and budget policy -- map child output into parent graph update -- make child usage and cost visible in parent rollups -- include child run references in checkpoints and task events - -Sub-agents allow a graph to coordinate specialized workers while keeping each -worker independently observable and testable. - -Parallel sub-agent fanout should use the graph's context-forking contract so -child agents inherit run identity, stores, stream sinks, and cache handles while -receiving isolated mutable task scope. See -[Parallel agents and context forking](parallel-agents-forking.md). - -## Steering Sub-Agents - -Graph-backed sub-agents can be steered by a parent orchestrator, a graph -supervisor node, middleware, a human operator, or a test. Steering should use -the harness steering command contract, then lower into graph commands, -checkpoint updates, or child-run delivery depending on the target state. See -[Sub-agent and orchestrator steering](../harness/subagent-steering.md). - -Supported graph-level steering cases: - -- parent orchestrator sends additional instructions to a running child - sub-agent -- parent orchestrator narrows a child sub-agent's tool or model policy -- parent orchestrator asks a child sub-agent for status before join -- human pauses, resumes, redirects, approves, rejects, or cancels a specific - child sub-agent run -- human steers the parent orchestrator while children continue or pause under - policy -- graph supervisor node redirects a child task to a review, retry, or finalize - node - -Graph steering targets must be addressable by run-tree metadata: - -```rust -pub enum GraphSteeringTarget { - ParentRun(RunId), - ChildRun { parent_run_id: RunId, child_run_id: RunId }, - Task { run_id: RunId, task_id: TaskId }, - Node { run_id: RunId, node_id: NodeId }, - Namespace { run_id: RunId, namespace: Vec<String> }, -} -``` - -Rules: - -- parent runs can steer descendants by default, not unrelated runs -- humans can steer any run or child task allowed by the control surface policy -- steering a child must not mutate parent state directly -- child steering is delivered as structured input with actor/provenance, not as - anonymous user text -- accepted child steering produces child events and parent rollup events -- rejected child steering includes a policy reason -- graph reducers own any state change caused by steering -- checkpoint metadata records pending, applied, and rejected steering commands - -When a sub-agent is paused on an interrupt, steering can target that exact -interrupt or child namespace. Resuming a child interrupt should restart or -continue only that child task unless the steering policy explicitly escalates to -the parent graph. - -## Orchestrator Steering - -The parent orchestrator is itself steerable. Human or supervisor steering can: - -- add or replace orchestration instructions at the next safe boundary -- narrow delegation policy before more sub-agents are spawned -- cancel pending child fanout -- redirect the graph to a review, retry, join, or finalize node -- apply reducer-mediated state corrections - -The orchestrator must treat steering as external control with provenance. It -should not silently merge human instructions into its own assistant messages or -pretend parent-agent steering came from the original user. - -## Steering Events - -Graph event streams should include steering alongside sub-agent lifecycle: - -- `steering.requested` -- `steering.accepted` -- `steering.rejected` -- `steering.delivered` -- `steering.applied` -- `subagent.steering_received` -- `subagent.steering_rejected` - -Every steering event must include root run id, parent run id, child run id when -available, task id when graph-backed, namespace, actor, command kind, policy id, -correlation id, and checkpoint id when available. - -## Detached Task Runtime Registry - -`DetachedTaskRegistry<Metadata, Status>` complements the durable `TaskStore` -with process-local executor handles. Applications register a task's owner, -metadata, status watch receiver, cancellation token, and abort handle, while -the task's live `SteeringHandle` remains addressable through the shared -`SteeringRegistry`. - -The registry enforces owner checks for wait and steering, keeps timed-out waits -registered, prunes terminal waits, cancels cooperatively before hard-aborting, -and sweeps terminal entries when its soft cap is reached. It never evicts live -work to satisfy the cap and does not duplicate durable lifecycle records. -After a restart, applications reconcile live `TaskStore` records that have no -matching runtime entry according to their executor policy. - -## Listing Detached Tasks - -The former model-facing `orchestrate_*` tools have been removed. Hosts that -need to list managed tasks can query the `TaskStore` directly with an -`OrchestrationTaskFilter`: - -```rust -use tinyagents_tasks::{OrchestrationTaskFilter, TaskStore}; - -let subagents = task_store.list( - OrchestrationTaskFilter::default().with_kind("sub_agent"), -); -``` - -The filter supports run-tree fields (`parent_run_id`, `root_run_id`, -`thread_id`, `node_id`, `status`), `kind` labels such as `"graph"` and -`"sub_agent"`, and inclusive `created_between` bounds using `SystemTime`. -`TaskStore` remains the durable lifecycle API; hosts that expose model-facing -task controls must provide their own tools and authorization policy. - -## Recursion And Depth Tracking - -The graph should allow recursive execution, but only with explicit limits and -tracking. - -Recursive cases: - -- a graph invokes itself as a subgraph -- a graph invokes a subgraph that eventually invokes the parent graph -- an agent node calls another agent that calls back into the same graph -- a router intentionally loops through nodes until state converges -- a `Send` fanout schedules many recursive child tasks - -Required tracking: - -```rust -pub struct RecursionFrame { - pub graph_id: GraphId, - pub node_id: Option<NodeId>, - pub run_id: RunId, - pub task_id: Option<TaskId>, - pub namespace: Vec<String>, - pub depth: usize, - pub parent: Option<RunId>, -} - -pub struct RecursionPolicy { - pub max_depth: usize, - pub max_visits_per_node: Option<usize>, - pub max_total_steps: usize, -} -``` - -Rules: - -- every graph/subgraph/sub-agent call pushes a recursion frame -- every return pops the frame -- depth is emitted on graph and registry events -- exceeding `max_depth` fails with a clear graph recursion error -- node-loop recursion and graph-call recursion are tracked separately -- checkpoint metadata includes the current recursion stack -- UIs should be able to render nested runs without reconstructing them from logs - -## See also - -- [Harness state graph runtime](../harness/state-graph.md) documents - `tinyagents_graph::agent_loop` (A5): the harness's own `plan -> model -> - tools -> settle` loop, compiled into an ordinary `CompiledGraph`. It is a - leaf graph, not a sub-agent host, but it composes with everything on this - page the same way any other compiled graph does — a `plan`/`model`/`tools` - node it drives can itself call a sub-agent tool, and the compiled loop - graph can itself be embedded as a subgraph node, so "graphs run graphs" - still holds all the way down into the loop that is usually the - recursion's innermost, non-graph leaf. diff --git a/docs/modules/graph/subgraphs.md b/docs/modules/graph/subgraphs.md deleted file mode 100644 index c9b78b57b..000000000 --- a/docs/modules/graph/subgraphs.md +++ /dev/null @@ -1,82 +0,0 @@ -# Graph Subgraphs - -Subgraphs are compiled graphs used as nodes. - -Two state modes: - -```text -shared-state subgraph -parent State == child State - -adapter subgraph -parent State -> child Input -> child Output -> parent Update -``` - -Subgraph requirements (implemented unless marked target; verified against -`crates/tinyagents-graph/src/subgraph/`): - -- namespace checkpoint ids -- preserve `root_run_id` -- set child `parent_run_id` -- propagate thread id by default -- **Target (not implemented):** allow isolated child thread ids by explicit - configuration — today an embedded child always runs on the parent's thread - when one is set, or unthreaded (no checkpoints) when it is not; there is no - opt-in for a child to have its own independent thread id -- **Target (not implemented):** inherit, override, or disable the parent - checkpointer per subgraph — today the child always inherits the parent's - checkpointer -- emit nested events with parent node id and namespace -- stream child values, updates, messages, tasks, and checkpoints when requested -- **Target (not implemented):** allow `Command::Parent` handoff from child - graph to parent graph — no `Parent` variant exists on `Command` -- each `ChildRun` entry (the `child_runs` array embedded in the parent's - boundary-checkpoint metadata) carries the child's latest checkpoint id - alongside its run id and node, so the association between a parent - activation and the exact child checkpoint it drove is explicit -- a `Send` fan-out of the same subgraph node — several concurrent - activations of one node within a step (map-reduce over a subgraph) — - namespaces each activation's child under `[node_id, task_id]` instead of - every activation sharing one `[node_id]` namespace; a node activated only - once keeps the plain `[node_id]` namespace, so existing checkpoints stay - readable -- a subgraph child that failed (or is interrupted with a resume value - already in hand) is continued through the parent rather than restarted: - `retry()`/`resume()` on the parent detects the child's own resumable - checkpoint and retries/resumes it in place, instead of re-running the - child's already-completed nodes from scratch - -Subgraph persistence must be explicit. Inherited checkpointing is convenient for -shared-state subgraphs; isolated checkpointing is safer for reusable child -graphs that may also run independently. - -## Run hierarchy - -When a subgraph node runs its embedded `CompiledGraph`, the child run is wired -into the parent's recursion tree: - -- the child gets its own `run_id`, preserves the enclosing run's `root_run_id`, - and sets `parent_run_id` to the parent run; -- the child run extends the parent's recursion frames (seeded from - `NodeContext::recursion_frames`) and its root frame names the embedding node, - so depth tracking is correct without mutating the parent's live stack; -- the spawned child is reported back through `NodeContext::child_runs` and - surfaces on `GraphExecution::child_runs` (a `ChildRun` list keyed by node) and - in every boundary checkpoint's metadata under a `child_runs` array; -- callers read the parent/child lineage after a run via - `GraphExecution::run_tree()` (a `RunTree`: this run's id, the shared root, the - parent run, and the spawned children). - -When the parent `NodeContext` carries a thread id, embedded child graphs run on -that same thread by default. The child checkpoint namespace is still extended -with the embedding node id, so parent and child checkpoints share a thread but -remain separately inspectable: - -```text -thread: t -parent checkpoint namespace: [] -child checkpoint namespace: ["child_node"] -``` - -If the parent run has no thread id, the child keeps the unthreaded execution -behavior and does not persist checkpoints. diff --git a/docs/modules/graph/todos.md b/docs/modules/graph/todos.md deleted file mode 100644 index a7c675483..000000000 --- a/docs/modules/graph/todos.md +++ /dev/null @@ -1,3 +0,0 @@ -# Todos moved to harness-tools - -Per-thread todo tools, persistence, and sequence middleware are generic harness functionality. They live in `crates/tinyagents-harness-tools/src/todos/`. See [`Goals-and-Todos`](../../../wiki/Goals-and-Todos.md). diff --git a/docs/modules/graph/visualization-testkit.md b/docs/modules/graph/visualization-testkit.md deleted file mode 100644 index 6efefa0c2..000000000 --- a/docs/modules/graph/visualization-testkit.md +++ /dev/null @@ -1,135 +0,0 @@ -# Graph Visualization, Introspection, And Testkit - -The graph exports (implemented in `crates/tinyagents-graph/src/export`): - -- graph id and name — `GraphTopology::graph_id` / `GraphTopology::name` - (set via `GraphBuilder::with_name`) -- node ids — `NodeInfo::id` -- node metadata — `NodeInfo::kind`, `NodeInfo::metadata` - (`GraphBuilder::with_node_kind` / `with_node_metadata`) -- node policy summaries — `NodeInfo::policy` (`NodePolicySummary`: derived - routing discipline plus barrier/interrupt/deferred/subgraph roles) and the - graph-level `GraphTopology::policy` (`GraphPolicySummary`: recursion limit, - parallel, max concurrency, node timeout) -- direct edges — `GraphTopology::edges` -- waiting/barrier edges — `GraphTopology::waiting_edges` (`WaitingEdgeInfo`, - lifted out of the direct-edge set) -- conditional route labels — `GraphTopology::conditional_edges` -- command destination hints — `NodeInfo::command_destinations` - (`GraphBuilder::with_command_destinations`) -- start and end paths — `GraphTopology::entry` / `GraphTopology::finish_nodes` -- interrupt markers — `NodeInfo::interrupt` (`GraphBuilder::mark_interrupt`) -- deferred-node markers — `NodeInfo::deferred` (`GraphBuilder::mark_deferred`) -- subgraph nodes — `NodeInfo::subgraph` (`GraphBuilder::mark_subgraph`) -- validation report — `GraphTopology::validation` (`ValidationReport`: - structural errors for dangling references, warnings for unreachable/dead-end - nodes) - -Export formats: - -- structured JSON — `to_json` / `from_json` (the structured `GraphTopology` is - the single source of truth; tests snapshot it) -- Mermaid — `to_mermaid` (subgraph nodes use the subroutine shape; interrupt, - deferred, and subgraph nodes get `classDef` markers; barrier and command - `goto` edges render dotted) -- DOT — later - -Both extraction sources — `CompiledGraph::topology` and -`GraphBuilder::topology` — produce the same -`GraphTopology`, so visualization and test snapshots share one truth. - -## Testkit - -`graph::testkit` (implemented in `crates/tinyagents-graph/src/testkit`) provides graph-test -building blocks distinct from the harness testkit. Each node helper returns a -closure ready for `GraphBuilder::add_node`: - -- no-op node — `noop_node` -- scripted update node — `scripted_update_node` -- scripted route node — `scripted_route_node` -- send/fanout node — `fanout_node` -- failing node — `failing_node` -- retry-counting node — `RetryCountingNode` (`.handler(..)` / `.attempts()`) -- interrupting node — `interrupting_node` -- subgraph test node — `subgraph_test_node` (wraps `shared_subgraph_node`) -- sub-agent fake node — `subagent_fake_node` (records a `ChildRun`) -- in-memory checkpointer — reuse `graph::InMemoryCheckpointer` - -Observation and assertion surfaces: - -- event recorder — `GraphEventRecorder` (`.sink()`, `.events()`, `.kinds()`) -- stream collector — `StreamCollector` (`node_order`, `updates`, `routes`, - `interrupts`, `checkpoint_count`, `custom`) -- run bundling — `run_recorded` produces a `GraphRun` (execution + recorded - events + checkpoint history), the single test truth -- fluent assertions — `assert_graph(run)` → `GraphAssertions`: `.visited`, - `.routed` (route assertion), `.checkpoint_count`, `.state_history` (state - history assertion), `.checkpoint` (checkpoint assertion), `.completed`, - `.interrupted` - -`run_recorded` wires the recorder, runs the (cloned) graph, and — when a thread -id is supplied — collects the durable checkpoint history, so the export/event/ -checkpoint truth all read from one `GraphRun`. - -## Storage conformance {#conformance} - -`graph::testkit::conformance` encodes durable-store **contracts** once so any -backend — the built-in ones or a caller-supplied adapter — behaves -interchangeably. Each function panics with a descriptive message on the first -violation, so call them from a `#[test]` / `#[tokio::test]`. Five contracts -ship: - -| Contract | Asserts | -| --- | --- | -| `checkpointer_contract(cp)` | Single-threaded put/get (latest + specific), insertion-order listing, `list_threads`, `delete_thread`, `prune`. | -| `taskstore_contract(store)` | Full task lifecycle state machine, cooperative cancel, kill, deadline updates, status filtering, terminal-transition rejection. | -| `checkpointer_concurrent_contract(Arc<C>)` | Many concurrent tasks put distinct checkpoints on one shared instance; every write lands and is retrievable (no lost writes). | -| `taskstore_concurrent_contract(Arc<S>)` | Many threads insert and advance distinct tasks against one shared store; every write lands exactly once. | -| `taskstore_replay_contract(reopen)` | Durable state written through one handle survives re-opening the backing store (terminal status + transition history replay). | - -The concurrent contracts take an `Arc<_>` (the store is shared across -threads/tasks); the replay contract takes a `Fn() -> S` that re-opens the same -durable backing so the same store can be closed and reconstructed. Wire them -across every backend in one test file: - -```rust -use tinyagents_graph::checkpoint::{DriverCheckpointer, InMemoryCheckpointer}; -use tinyagents_tasks::InMemoryTaskStore; -use tinyagents_graph::testkit::conformance::{ - checkpointer_concurrent_contract, taskstore_concurrent_contract, taskstore_replay_contract, -}; - -#[tokio::test] -async fn checkpointer_handles_concurrent_writes() { - checkpointer_concurrent_contract(std::sync::Arc::new(InMemoryCheckpointer::<i32>::new())).await; -} - -#[test] -fn task_store_handles_concurrent_writes() { - taskstore_concurrent_contract(std::sync::Arc::new(InMemoryTaskStore::new())); -} - -#[test] -fn driver_task_store_replays_after_restart() { - // `open` returns a `DriverTaskStore` over the same SQLite file each call. - let dir = tempfile::tempdir().unwrap(); - taskstore_replay_contract(|| open_driver_task_store(dir.path())); -} -``` - -Only durable backends (`DriverCheckpointer`, `DriverTaskStore` on a durable driver) can pass -the replay contract — an in-memory store drops state on reopen, so run it only -against the durable ones. - -Example: - -```rust -use tinyagents_graph::testkit::{assert_graph, run_recorded}; - -let run = run_recorded(&graph, Some("t1"), 0).await?; -assert_graph(&run) - .visited(["agent", "tools", "agent"]) - .routed("agent", "tools") - .checkpoint_count(3) - .completed(); -``` diff --git a/docs/modules/harness/README.md b/docs/modules/harness/README.md index 5ca23a2c3..97a6a5e45 100644 --- a/docs/modules/harness/README.md +++ b/docs/modules/harness/README.md @@ -16,15 +16,14 @@ registration, tool registration, prompt assembly, middleware, event streaming, tracing, retries, limits, summarization, caching, usage accounting, pricing, sub-agent/orchestrator steering, and test support. -The harness should be usable in three modes: +The harness is currently used in two modes: 1. direct model invocation 2. model plus tools agent loop -3. graph node runtime dependency -It should not require the graph module. The graph module can depend on harness -traits, but a user should be able to call a model or run a tool loop without -constructing a graph. +The graph runtime formerly provided a third mode, but was removed in TinyAgents +#61. The harness now owns the direct model-tool-model loop and has no graph +runtime dependency. ## Source Inspiration @@ -100,9 +99,6 @@ tests for every adapter. - Dispatch tool calls through a registry with schema validation. - Expose retrievers and vector stores for retrieval-augmented context assembly. - Run the standard model-tool-model agent loop. -- Represent the standard agent loop as an explicit state machine when callers - need inspection, checkpointing, HITL, branch/fork semantics, or graph-node - execution. - Apply middleware before and after model calls, tool calls, retries, and errors. - Enforce model-call limits, tool-call limits, timeouts, and retry policy. - Emit typed events for tracing, streaming, and tests. @@ -121,29 +117,23 @@ tests for every adapter. - Provide deterministic test utilities. - Describe provider capability profiles so middleware can choose safe defaults. - Translate between provider-native message formats and TinyAgents messages. -- Persist resolved model identity in response metadata, run status, events, - usage/cost rows, and durable agent or graph state for reuse. +- Persist resolved model identity in response metadata, run status, events, and + usage/cost rows for reuse. - Support dynamic runtime context injection into tools and middleware without exposing private state to model-visible schemas. - Support model fallback, tool retry, rate limiting, and human interruption as explicit policies rather than ad hoc callbacks. - Support parent-orchestrator and human steering of sub-agents, orchestrator - agents, graph tasks, and harness loops through typed commands delivered at + agents, and harness loops through typed commands delivered at safe boundaries. - Queue host-owned payloads in steer, follow-up, and collected-context lanes without importing host transport or UI metadata into the harness. -- Support durable graph runs with pause/resume, checkpoint listing, and - inspectable node transitions. -- Support per-agent execution blueprints that describe how an agent runs - separately from what its prompt says. -- Provide a standard conformance test suite for model, tool, store, stream, and +- Provide a standard conformance test suite for model, tool, store, and middleware implementations. ## Non-Responsibilities -- It does not own graph topology. -- It does not decide graph routing except inside an agent-loop node. -- It does not persist graph checkpoints; that belongs to the graph module. +- It does not provide a workflow or state-graph runtime. - It does not hide provider-specific metadata when users need it. - It does not execute arbitrary workflow language source directly. - It does not require every provider to support every modality or output @@ -168,7 +158,6 @@ crates/tinyagents-harness/src/ cost.rs embeddings.rs events.rs - graph_runtime.rs limits.rs message.rs middleware.rs @@ -182,7 +171,6 @@ crates/tinyagents-harness/src/ types.rs runtime.rs steering.rs - stream.rs summarization.rs structured.rs store.rs @@ -207,8 +195,6 @@ Feature ownership: - `embeddings`: embedding providers, vector stores, retrievers, indexing, and retrieval-context records. - `events`: typed harness events, sinks, streams, redaction adapters. -- `graph_runtime`: explicit state graphs, node commands, reducers, - checkpointing, HITL, run records, and graph execution blueprints. - `limits`: model-call, tool-call, concurrency, timeout, and recursion policy. - `message`: structured messages, content blocks, tool call correlation. - `middleware`: before/after/wrap hooks and middleware stack ordering. @@ -220,8 +206,9 @@ Feature ownership: collected-context payloads; hosts retain ownership of payload metadata. - `runtime`: high-level `AgentHarness` builder/facade. - `steering`: policy-checked parent/human steering of orchestrators, - sub-agents, graph tasks, and harness loops. -- `stream`: token streams, tool progress streams, event streams, adapters. + sub-agents, and harness loops. +- Streaming is provided through the agent loop and event APIs; the former + graph stream projection was removed in TinyAgents #61. - `summarization`: context summaries, message compaction, summary provenance. - `structured`: typed response formats and validation. - `store`: persistence contracts supplied by hosts. @@ -247,7 +234,6 @@ Feature details: - [Local models and embeddings (Ollama, LM Studio)](local-models.md) - [Embeddings and retrieval feature](embeddings.md) - [Media generation tools](media.md) -- [State graph runtime feature](state-graph.md) - [Prompt feature](prompt.md) - [Tool feature](tool.md) - [Tool execution context and rich returns (B1/B2)](tool-context.md) diff --git a/docs/modules/harness/cache.md b/docs/modules/harness/cache.md index e0fcee12b..567b95d9b 100644 --- a/docs/modules/harness/cache.md +++ b/docs/modules/harness/cache.md @@ -92,8 +92,7 @@ Request builders and middleware should follow these rules: - fingerprint prompt segments separately from the full request - include middleware policy fingerprints when a middleware can affect model-visible bytes -- preserve `thread_id` across parent agents, sub-agents, graph subgraphs, and - harness calls so provider-side prompt/KV-cache attribution remains stable for +- preserve `thread_id` across parent agents, sub-agents, and harness calls so provider-side prompt/KV-cache attribution remains stable for the logical conversation - map stable `thread_id` or tenant/user identity into provider-specific cache headers when required by a provider policy; for example, Fireworks-style diff --git a/docs/modules/harness/context.md b/docs/modules/harness/context.md index bf3bc0689..739ac5af6 100644 --- a/docs/modules/harness/context.md +++ b/docs/modules/harness/context.md @@ -12,9 +12,9 @@ and model context-window pressure. - Carry cancellation. - Carry store, event, cache, usage, and cost handles. - Track context-window budget for the selected model. -- Expose context to models, tools, middleware, and graph nodes. -- Provide inherited context to nested model calls, tools, sub-agents, and graph - nodes without relying on global variables. +- Expose context to models, tools, and middleware. +- Provide inherited context to nested model calls, tools, and sub-agents + without relying on global variables. - Hide runtime-only values from model-visible tool schemas. ## Source Inspiration @@ -94,8 +94,8 @@ mid-tool or mid-chunk, only at well-defined checkpoints. `RunContext` carries a one-shot control request that a middleware or step can set to steer the agent loop from outside its `Result<()>` return channel. This is the -harness-native complement to the graph `Command`/`Interrupt` vocabulary; the loop -drains any request at its safe checkpoints (after each model response). +harness-native control surface; the loop drains requests at safe checkpoints +(after each model response). ```rust pub enum MiddlewareControl { diff --git a/docs/modules/harness/cost.md b/docs/modules/harness/cost.md index 6d842eeb3..7ce12496c 100644 --- a/docs/modules/harness/cost.md +++ b/docs/modules/harness/cost.md @@ -9,7 +9,7 @@ render spend in web UIs. - Price prompt, completion, cached, and reasoning tokens. - Track per-request fixed fees when needed. - Enforce per-run and per-thread budgets. -- Roll costs up across sub-agents and graph nodes. +- Roll costs up across sub-agents. - Emit cost events. - Keep price data updateable outside provider adapter code. - Mark estimated versus provider-confirmed cost. diff --git a/docs/modules/harness/design-notes.md b/docs/modules/harness/design-notes.md index 6ff03dc3f..3bbf14b53 100644 --- a/docs/modules/harness/design-notes.md +++ b/docs/modules/harness/design-notes.md @@ -50,13 +50,13 @@ pub struct RunContext<Ctx = ()> { runtime dependency container. This split keeps tests deterministic and prevents global singletons. -Nested model calls, tools, sub-agents, and graph nodes must inherit the root run +Nested model calls, tools, and sub-agents must inherit the root run id, selected tags, inherited metadata, event sink, cancellation token, stores, usage tracker, cost tracker, and configured budget policy. They may add local tags and metadata, but they must not mutate parent config in place. Nested runs may also receive steering commands. Steering is explicit runtime -control from a parent orchestrator, human, graph supervisor, middleware, or +control from a parent orchestrator, human, middleware, or test. A steered run must record actor, target, policy, payload summary, and the safe boundary where the command was applied. See [Sub-agent and orchestrator steering](subagent-steering.md). diff --git a/docs/modules/harness/embeddings.md b/docs/modules/harness/embeddings.md index 08bf19196..9af4f48a3 100644 --- a/docs/modules/harness/embeddings.md +++ b/docs/modules/harness/embeddings.md @@ -59,7 +59,6 @@ enters a model request. ## Non-Responsibilities -- It does not decide graph routing. - It does not automatically inject retrieved documents into every prompt. - It does not replace durable application stores. - It does not require a vector store for every embedding call. diff --git a/docs/modules/harness/hosting.md b/docs/modules/harness/hosting.md index 45966722b..819bb85ea 100644 --- a/docs/modules/harness/hosting.md +++ b/docs/modules/harness/hosting.md @@ -45,13 +45,10 @@ assemble a run without host capabilities. Hosted invocations require `State: 'static` because their live capability authority must be retained in the recursive context. The explicit-model `invoke*`, streaming, and direct `SubAgent` paths do not install or inspect -that authority and continue to support borrowed state. +that authority and continue to support borrowed state. Child execution is +planned by the host through `SubagentDriver`; the host supplies the child +context and enforces its own delegation policy. -For a child of a hosted parent, call -`SubAgent::invoke_hosted_in_parent`; it rechecks the parent's delegate -allowlist and inherits the exact bundle. The borrowed-state-compatible -`SubAgent::invoke_in_parent` is explicit-only and rejects a hosted parent -context before it can start a child. ## Hosted failures and timeout bounds diff --git a/docs/modules/harness/langchain-parity.md b/docs/modules/harness/langchain-parity.md index 6225512b1..912ef76b7 100644 --- a/docs/modules/harness/langchain-parity.md +++ b/docs/modules/harness/langchain-parity.md @@ -8,7 +8,7 @@ surface area that TinyAgents should intentionally support, adapt, or reject. | LangChain area | Source | TinyAgents harness implication | | ---------------------------- | --------------------------------------------------------------------------------------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -| `create_agent` factory | `libs/langchain_v1/langchain/agents/factory.py` | `AgentHarness` should compose model selection, tool execution, middleware, structured output, runtime context, and graph-node compatibility behind one builder while keeping traits reusable outside the facade. | +| `create_agent` factory | `libs/langchain_v1/langchain/agents/factory.py` | `AgentHarness` should compose model selection, tool execution, middleware, structured output, and runtime context behind one builder while keeping traits reusable outside the facade. | | Agent middleware | `libs/langchain_v1/langchain/agents/middleware/types.py` | Middleware needs before/after hooks, streaming delta hooks, and wrap hooks that can replace the model/tool call, inject commands, short-circuit, or jump to `model`, `tools`, or `end`. | | Built-in middleware | `libs/langchain_v1/langchain/agents/middleware/*.py` | Ship focused middleware for summarization, context compression, transcript compression, retrieval compression, output compression, prompt cache layout guards, context editing, PII redaction, model/tool limits, retries, fallback, tool selection, human-in-the-loop, shell/file-search style privileged tools, and todo/task state. | | Structured output | `libs/langchain_v1/langchain/agents/structured_output.py` | Support provider-native schemas and artificial tool-call schemas, with typed validation, retryable validation errors, union/oneOf variants, and configurable error handling. | @@ -21,7 +21,7 @@ surface area that TinyAgents should intentionally support, adapt, or reject. | OpenHuman checkpointer | `openhuman#4261`, `src/openhuman/agent_graph/checkpoint/*` | Persist graph runs and checkpoints through a pluggable `Checkpointer`, with in-memory tests and durable SQLite-style production storage. | | OpenHuman graph blueprints | `openhuman#4261`, `src/openhuman/agent_graph/blueprint/*` | Keep per-agent execution topology in `graph.rs`-style blueprints next to prompts, so "what the agent says" and "how the agent runs" are inspectable separately. | | OpenHuman live turn graph | `openhuman#4261`, `src/openhuman/agent_graph/live/*` and `agent/harness/engine/core.rs` | Preserve the hot-path turn contract while making phases explicit: dispatch, parse, stop check, tools, compact, loop, finalize, max-iteration checkpoint. | -| OpenHuman sub-agent steering | `spawn_subagent`, `spawn_async_subagent`, `steer_subagent`, `wait_subagent` product pattern | Generalize steering into typed commands so parent orchestrators, humans, middleware, UIs, and tests can guide sub-agents or orchestrators without prompt-injection side channels. | +| OpenHuman sub-agent steering | `spawn_subagent` with driver-backed job control | Generalize steering into typed commands so parent orchestrators, humans, middleware, UIs, and tests can guide sub-agents or orchestrators without prompt-injection side channels. | | Vector stores | `libs/core/langchain_core/vectorstores/base.py`, `in_memory.py` | Support add/update/delete/get-by-id, similarity search, score-threshold search, MMR search, metadata filters, async variants, and in-memory test stores. | | Retrievers and indexing | `libs/core/langchain_core/retrievers.py`, `indexing/*.py` | Treat retrievers as query-to-document components with events, tags, metadata, and record-manager-backed incremental indexing for dedupe and cleanup. | | Tool runtime injection | `langgraph.prebuilt.ToolRuntime` as re-exported by `libs/langchain_v1/langchain/tools/tool_node.py` | Tools should receive typed runtime context, state, store handles, stream writers, and cancellation handles through Rust parameters, not model-visible JSON schema fields. | @@ -29,6 +29,6 @@ surface area that TinyAgents should intentionally support, adapt, or reject. | Runnables config | `libs/core/langchain_core/runnables/config.py` | `RunConfig` should carry tags, metadata, configurable values, concurrency, recursion, callbacks/events, and stable run identity through nested calls. | | Retry/fallback/rate limit | `libs/core/langchain_core/runnables/retry.py`, `fallbacks.py`, `rate_limiters.py` | Policies should distinguish retryable transport errors, provider errors, validation errors, tool errors, budget failures, and rate-limit waits. | | Cache | `libs/core/langchain_core/caches.py` | Separate local response cache from provider prompt/KV-cache reuse, preserve stable prefix layout, and include all behavior-affecting request fields in keys. | -| Stores and chat history | `libs/core/langchain_core/stores.py`, `chat_history.py` | Keep generic stores separate from conversation memory and graph checkpoints. | +| Stores and chat history | `libs/core/langchain_core/stores.py`, `chat_history.py` | Keep generic stores separate from conversation memory. | | Standard tests | `libs/standard-tests` | Add reusable conformance tests so provider adapters prove tool calling, structured output, streaming, usage, callbacks/events, multimodal input, Unicode, and error behavior. | diff --git a/docs/modules/harness/limits-retry.md b/docs/modules/harness/limits-retry.md index 71eae780f..e017140a5 100644 --- a/docs/modules/harness/limits-retry.md +++ b/docs/modules/harness/limits-retry.md @@ -17,7 +17,7 @@ LangChain and LangChain v1 provide related behavior in several places: <https://github.com/langchain-ai/langchain/tree/master/libs/langchain_v1/langchain/agents/middleware> TinyAgents should centralize these policies so every model call, tool call, -agent loop, and graph node can share predictable behavior. +agent loop can share predictable behavior. ## Responsibilities @@ -28,7 +28,6 @@ agent loop, and graph node can share predictable behavior. - Enforce per-call timeouts. - Enforce maximum concurrency. - Enforce optional recursion depth. -- Enforce graph super-step caps. - Support cancellation. - Classify retryable and non-retryable errors. - Support exponential backoff with jitter. @@ -46,7 +45,6 @@ pub struct RunLimits { pub max_retries: usize, pub max_concurrency: usize, pub max_recursion_depth: usize, - pub max_graph_steps: usize, pub wall_clock_timeout: Option<Duration>, pub model_call_timeout: Option<Duration>, pub tool_call_timeout: Option<Duration>, @@ -56,10 +54,6 @@ pub struct RunLimits { Limits should fail closed. When a limit is reached, the harness should stop the agent loop and return a classified error that includes the current counters. -For graph-backed runs, max-step failures should include graph name, run id, -current super-step, frontier nodes, and the configured cap. The run should be -persisted as failed when a checkpointer is configured. - ## Retry Policy ```rust diff --git a/docs/modules/harness/middleware.md b/docs/modules/harness/middleware.md index 0e0b5f950..1cd21cc98 100644 --- a/docs/modules/harness/middleware.md +++ b/docs/modules/harness/middleware.md @@ -1,7 +1,7 @@ # Harness Middleware Feature Middleware is the main extension point for behavior that cuts across models, -tools, memory, stores, streaming, and graph nodes. +tools, memory, stores, and streaming. ## Source Inspiration @@ -35,8 +35,6 @@ understand graph internals for normal harness usage. - Allow middleware to emit events. - Allow middleware to add local state updates without mutating unrelated state. - Allow middleware to jump to model, tools, or end when used inside an agent loop. -- Translate middleware control outcomes into state-graph commands when a run is - graph-backed. - Expose errors to middleware for logging, redaction, fallback, or recovery. - Keep middleware testable with fake models, tools, and event sinks. @@ -144,27 +142,6 @@ Middleware should be able to return: - resume from checkpoint - fail with classified error -Graph-specific commands should be translated at the graph boundary. The harness -should expose harness-native control outcomes so it remains usable without a -graph. - -## Graph Boundary - -Middleware must not need to know whether the caller is using the simple loop or -the state-graph runtime. The runtime adapter maps harness-native outcomes onto -graph commands: - -- continue -> `Command::Continue` -- jump to model/tools/end -> `Command::Goto(...)` or `Command::End` -- human interrupt -> `Command::Interrupt` -- accepted steering -> `Command::Update`, `Command::Goto`, - `Command::Interrupt`, or queued child-run delivery depending on target -- branch/fan-out middleware -> `Command::Fork` -- retry/fallback -> handled inside the node or wrap hook before command return - -When middleware mutates graph-visible state, it must emit an explicit state -update event so checkpoint replay can explain the change. - ## Built-In Middleware Initial built-ins should include: diff --git a/docs/modules/harness/model.md b/docs/modules/harness/model.md index 628a05a88..cd3f70eed 100644 --- a/docs/modules/harness/model.md +++ b/docs/modules/harness/model.md @@ -43,7 +43,6 @@ feature-gated provider crates/modules and a shared conformance suite. ## Non-Responsibilities - It does not execute tools. -- It does not choose graph routes. - It does not own provider API keys globally. - It does not hardcode time-sensitive model prices. - It does not assume all models support tool calling, structured output, @@ -103,8 +102,7 @@ pub struct ModelResponse { Agents and model calls should be able to define their own model preferences without hardcoding one provider everywhere. The harness resolves a model at the start of a model call, records the selected model, and makes that resolved model -available to later calls, child agents, graph nodes, event streams, and persisted -state. +available to later calls, child agents, event streams, and persisted state. The design mirrors OpenHuman-style smart model resolution from hints: callers can provide candidate models and capability requirements, while the registry and @@ -205,8 +203,7 @@ Every model call should attach `ResolvedModel` to: - model lifecycle events - usage and cost records - agent run status -- graph task metadata when graph-backed -- checkpoint state when the graph or harness state is durable +- session state when the run is persisted Agent state should have a standard place for reusable model resolution: diff --git a/docs/modules/harness/observability.md b/docs/modules/harness/observability.md index 85f18072e..85351d62b 100644 --- a/docs/modules/harness/observability.md +++ b/docs/modules/harness/observability.md @@ -31,14 +31,12 @@ TinyAgents should use typed Rust events rather than stringly callback names. - Preserve parent/root run relationships. - Include usage, cost, cache, retry, fallback, and latency data. - Capture model-call and tool-call latency as first-class metrics. -- Include graph run ids, graph names, super-step indexes, node ids, routing - commands, checkpoint ids, and interrupt metadata for state-graph runs. + ## Execution Status Store -Harness runs need a readable status surface just like graph runs. Direct model -calls, model-tool loops, and harness calls made from graph nodes should all -publish their current execution status so external listeners can inspect runs +Harness runs need a readable status surface. Direct model calls and model-tool +loops publish their current execution status so external listeners can inspect runs without holding an in-process stream. ```rust @@ -81,8 +79,8 @@ pub enum HarnessPhase { Status should move through `Pending`, `Running`, `Waiting`, `Completed`, `Failed`, or `Cancelled` states while `current_phase` describes the active -harness operation. A graph node that invokes a harness should be able to read -the child harness status through `parent_run_id` and `root_run_id`. +harness operation. Nested harness calls can read child status through `parent_run_id` and +`root_run_id`. Run-level latency can be derived from `started_at` and either `ended_at` or `updated_at`. This status-derived view is compact and cheap, but it only reports @@ -149,13 +147,6 @@ Event kinds should include: - `rate_limit.waited` - `limit.reached` - `stream.closed` -- `graph.run.started` -- `graph.node.entered` -- `graph.node.completed` -- `graph.run.paused` -- `graph.run.completed` -- `graph.run.failed` -- `graph.checkpoint.saved` ## Physical Provider Dispatches @@ -240,22 +231,6 @@ Live event delivery should not block model or tool execution unless the run policy explicitly requires a mandatory observer. Durable listeners should replay from the journal rather than relying on in-memory broadcast delivery. -### Implementation status (runtime-comparison Phase 3, C2) - -`HarnessEventJournal`/`StoreEventJournal` exist as designed above (offset -replay, no block-level detail). What is new this phase is a durable shape for -*in-progress* assistant messages specifically: -`tinyagents_harness::stream::{AssistantFrame, FrameEncoder, reduce_frames}` -(`crates/tinyagents-harness/src/stream/frame.rs`) — see -`docs/modules/harness/streaming.md`'s "Implementation status" for the codec -itself. A journal that persists `AssistantFrame`s alongside its ordinary -`HarnessEvent` stream lets a reconnecting consumer rebuild a partial -in-progress message with `reduce_frames` instead of only replaying whichever -already-flattened `AgentEvent::ModelDelta`s the journal captured — the -journal wiring (an `AssistantFrame`-aware `HarnessEventJournal` -implementation, or a `Model*` frame projection alongside `ModelDelta`) is not -itself implemented yet; the codec is the missing primitive that unblocks it. - ## Observability Cache The harness cache feature can store derived observability projections. These @@ -380,8 +355,7 @@ Input/Output panels of a generation (`generation-create`) and a tool observation (`span-create`). With capture disabled the fields are absent and the observation renders with ids, timing, and usage only. The harness exporter also defaults the trace-level metadata from the run lineage (root/first run ids and -parent), mirroring the graph exporter, so a trace is correlatable even when the -caller passes no metadata. +parent), so a trace is correlatable even when the caller passes no metadata. ## Langfuse Observation Tree @@ -397,9 +371,7 @@ Because run-span ids are deterministic and trace-namespaced, the tree spans **separately exported batches**: a sub-agent run whose observations are shipped in their own batch carries `parent_run_id`, so its run span resolves to `{trace_id}:run:{parent_run_id}` and nests under the parent agent's span in the -same trace. The graph exporter emits a matching `{trace_id}:run:{run_id}` span -for the graph run, so a graph run and the agent runs its nodes spawn -reconstruct one unified nested tree. +same trace, preserving the nested agent-run lineage. ## Scores @@ -412,22 +384,3 @@ deterministic value derived from the target and metric name, so re-scoring the same target upserts rather than duplicating. Point the score's `trace_id` at an exported trace (recall both exporters default the trace id to the run's `root_run_id`) to attach the evaluation to that run. - -## Graph Events - -State-graph event payloads should be rich enough for a UI to render run -progress without loading checkpoints: - -- run id -- graph name -- step or super-step index -- node id -- rendered command such as `continue`, `goto:node`, `fork:[a,b]`, - `interrupt:approval`, or `end` -- elapsed milliseconds -- lifecycle status -- interrupt kind and resume node when paused -- error class and message when failed - -Graph event tests should filter by both event kind and run id. A shared global -event bus can otherwise pick up unrelated graph events from parallel tests. diff --git a/docs/modules/harness/runtime.md b/docs/modules/harness/runtime.md index 21fe1dbd6..92e1d622c 100644 --- a/docs/modules/harness/runtime.md +++ b/docs/modules/harness/runtime.md @@ -274,8 +274,8 @@ installed (calling it again replaces the previous one). ## Middleware -Middleware is the main extension point for behavior that cuts across providers, -tools, and graph nodes. +Middleware is the main extension point for behavior that cuts across providers +and tools. ```rust #[async_trait] diff --git a/docs/modules/harness/state-graph.md b/docs/modules/harness/state-graph.md deleted file mode 100644 index 03d3c9f7a..000000000 --- a/docs/modules/harness/state-graph.md +++ /dev/null @@ -1,171 +0,0 @@ -# Harness State Graph Runtime - -This document describes what actually exists in the workspace today. An -earlier draft of this file sketched an aspirational `StateGraph<S>` design -inspired by an OpenHuman PR; none of those types (`GraphState`, `Node<S>`, -`Command::Fork`, a `StateGraph` builder) were ever implemented under -`tinyagents-harness`. The real state-graph runtime lives in the -`tinyagents-graph` crate (`crates/tinyagents-graph/src/`), documented in -`docs/modules/graph/`, and this page now only covers the one harness-facing -surface that is genuinely new: the compiled-graph rendition of the agent loop -(A5, `docs/runtime-comparison/langgraph.md` §4 and -`docs/runtime-comparison/pydantic-ai.md` §4's `iter` API). - -## Where the graph runtime actually lives - -`tinyagents-graph`'s `GraphBuilder`/`CompiledGraph` (`crates/tinyagents-graph/src/builder/`, -`crates/tinyagents-graph/src/compiled/`) is the durable, typed, Pregel-style -state graph: partial updates and reducers (`reducer`), commands and -interrupts (`command`), checkpointing (`checkpoint`, with `InMemoryCheckpointer`, -and `DriverCheckpointer`), streaming/events -(`stream`), subgraph embedding (`subgraph`), and recursion tracking -(`recursion`). See `docs/modules/graph/` for that design in full; nothing on -this page duplicates it. - -`tinyagents-harness` does not implement its own competing graph engine. What -it exposes, in `tinyagents_harness::agent_loop::phases`, is a **seam**: typed -phase contracts (`TurnPlan`, `ModelOutcome`, `ToolBatchOutcome`, `Settlement`) -and a `LoopDriver` trait that lets `AgentHarness` delegate its loop execution -to an alternate engine. `tinyagents-graph` is the only implementor of that -seam, because the dependency direction in this workspace is graph → harness -(never the reverse) — a compiled-graph rendition of the harness's own loop -therefore has to live in the graph crate, wired back in through this trait. - -## The agent loop as a `CompiledGraph` (A5) - -`tinyagents_graph::agent_loop` (`crates/tinyagents-graph/src/agent_loop/`) -compiles the harness's default `plan -> model -> tools -> settle` loop into -an ordinary `CompiledGraph<LoopState, LoopUpdate>`. `LoopState` is a -whole-state graph value (`LoopUpdate = LoopState`, built with -`GraphBuilder::overwrite()`): every node returns the complete next state -rather than a partial patch. It is `Serialize`/`Deserialize`, unlike the -harness's own `RunContext` (deliberately not serializable), so a -graph-driven run can be checkpointed mid-run — including at an interrupt — -and resumed. - -Three entry points, sharing one set of node bodies -(`runtime::plan_node`/`model_node`/`tools_node`/`settle_node`): - -- **`compile_loop(rt: Arc<LoopRuntime<State, Ctx>>) -> Result<CompiledGraph<LoopState, LoopUpdate>>`** - builds the real compiled graph over an owned `LoopRuntime` (harness/state - `Arc`'d, `RunContext`/`AgentRun`/`HarnessRunStatus` owned). This is the - graph that gets checkpointing, `resume`, and step-by-step control. -- **`AgentLoopGraphExt::iter(harness: Arc<AgentHarness<..>>, app_state, ctx, input) -> LoopIter`** - mirrors pydantic-ai's `Agent.iter`: `LoopIter::next()` runs exactly one - node activation and reports a `LoopStep` (which node ran, what runs next, - whether it interrupted); `override_next(node)` redirects the next - activation; `state()` reads the committed `LoopState`; `run_to_end()` steps - until completion or the first unresolved interrupt. -- **`GraphLoopDriver`** implements `phases::LoopDriver` and is installed with - `AgentHarness::with_loop_driver` plus - `RunPolicy::execution = LoopExecution::Graph` (the default, - `LoopExecution::Direct`, is the original `run_loop` body, byte-for-byte - unchanged). This makes `AgentHarness::invoke` itself run the graph-rendition - node bodies. - -### Why `GraphLoopDriver` does not build a `CompiledGraph` - -`phases::LoopDriver::drive` is handed a **borrowed** `&AgentHarness`/`&State` -and `&mut RunContext`/`&mut AgentRun`/`&mut HarnessRunStatus` — the same -shape `run_loop` itself uses. `GraphBuilder::add_node`'s closures must be -`Send + Sync + 'static`, so a `CompiledGraph`'s nodes can only close over -genuinely owned (`Arc`'d) data. Building an `Arc<AgentHarness>` from a bare -`&AgentHarness` is not possible in safe Rust, and this workspace denies -`unsafe_code`. Forcing every `with_loop_driver` caller to already hold an -`Arc<AgentHarness>` before installing a driver would also be circular (the -harness does not exist as an `Arc` until after it is fully built, including -its driver). - -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`. This is sound with zero unsafe code -because a borrowed async call needs no `'static` bound. The trade-off: an -interrupt raised through `AgentHarness::invoke` (`RunPolicy::execution == -Graph`) still ends the run cleanly (mirroring a steering pause, or -`TinyAgentsError::Interrupted` for a `MiddlewareControl::Interrupt` — see -below) but is **not** a resumable `CompiledGraph` checkpoint. A caller that -wants graph-level checkpoint/resume across an interrupt should drive the loop -through `compile_loop`/`LoopIter` directly instead of through -`AgentHarness::invoke`. - -### Routing and interrupts - -Every node routes explicitly (`mark_command_routing` on all four nodes) via -`Command::goto`, driven by `tinyagents_harness::context::MiddlewareControl`: - -| `MiddlewareControl` | Graph routing | -| --- | --- | -| `Continue` / `UpdateState` | falls through to whatever the calling node had already determined the turn's natural next step to be (tool routing from `model`, `plan` from `tools`) | -| `JumpTo(Model)` | routes to `plan` (closing any unanswered tool calls first) | -| `JumpTo(Tools)` | routes to `tools` if there are pending calls, else `settle` | -| `JumpTo(End)` / `StopWithFinal` | routes to `settle` with `finished = true` | -| `Interrupt { node, message }` | `NodeResult::Interrupt` — a real, checkpointable `crate::Interrupt` returned directly by the `plan`/`model`/`tools` node bodies through `compile_loop`/`LoopIter` (these nodes are marked as interrupt points for the export only — via the `NodeMeta` flag directly, not `GraphBuilder::mark_interrupt`, which now aliases the real `interrupt_before` pause and would double-pause every activation); through `GraphLoopDriver` this instead surfaces as `TinyAgentsError::Interrupted`, matching `run_loop`'s own behavior for this control | - -A steering pause (`SteeringOutcome::Pause`, checked in `plan_node`) also -produces a `NodeResult::Interrupt`, distinguished from a middleware interrupt -by its id suffix (`-steering-pause`); `GraphLoopDriver` treats that one as a -clean pause (`run.paused`, `Ok(())`), exactly like `run_loop`'s -`LoopExit::Paused`. - -`RunLimits` (model/tool call caps) are enforced through the same -`RunContext::limits`/`LimitTracker` the direct loop uses, reconciled against -`RunPolicy::limits` with the identical (`runtime::reconcile_call_limits`) -logic `run_loop_body` applies at the top of a run, so -`RunConfig::with_max_model_calls` and `RunPolicy::limits` interact -identically under both engines. A cap hit surfaces as -`TinyAgentsError::LimitExceeded` (or a clean `LimitStop`-style finish under -`LimitBehavior::StopWithPartial`) either way. - -### Tool batch execution - -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 call. The -direct loop's serial-admission / serial-or-concurrent-dispatch decision (see -`tinyagents_harness::agent_loop::tools`) is call-count- and -middleware-dependent; re-deriving it as graph topology would either duplicate -that decision as a router (two sources of truth) or lose the exact -ordering/budget guarantees the direct loop promises. Reusing the harness -function as-is keeps ordering, concurrency-eligibility, and budget/limit -semantics identical to the direct loop by construction — at the cost of a -coarser graph: a `tools` activation is atomic from the executor's point of -view, so resuming after an interrupt mid-batch re-runs the whole batch, not -just its unfinished calls. - -### Documented scope (not a byte-for-byte reimplementation) - -`tinyagents_graph::agent_loop` covers the common path: tool calling, -structured output (`ResponseFormat::Auto`/`JsonSchema`, provider-schema and -tool-call-fallback strategies), the output-validation retry loop (A3), run -limits, `MiddlewareControl` routing, and steering. It intentionally does -**not** cover: host-model routing (`HostCapabilities`), cross-provider -handoff transforms, the deferred-tool discovery bridge, truncated-empty- -response recovery/retry, `RunPolicy::retry`/`fallback`'s built-in retry loop -(a registered `ModelMiddleware` still runs), response caching, and -`StructuredStrategy::Prompted`/`ToolCallUnion` or -`EndStrategy::Early`/`Exhaustive` (A6). `resolve_structured_plan` also -resolves a profile-driven `ResponseFormat::Auto` choice against the -*default* model binding rather than the turn's actually-resolved model. -`RunPolicy::execution` defaults to `LoopExecution::Direct`, so no existing -caller is affected unless it opts in. - -Equivalence between the two engines is covered by -`crates/tinyagents-integration-tests/tests/loop_as_graph.rs`: scripted tool -call, structured output, output-validation retry, limit stop, approval -interrupt, and steering-inject scenarios run through both `Direct` and -`Graph`, asserting matching transcript/structured output/usage and that the -direct run's `AgentEvent` kind sequence is a subsequence of the graph run's -(the graph engine may emit extra events — for example `ToolsAdvertised` once -per turn instead of once per run — but never fewer or reordered ones); plus -checkpoint+resume across an interrupt (`InMemoryCheckpointer` and -`DriverCheckpointer`) and `LoopIter` stepping with `override_next`. - -## See also - -- `docs/modules/graph/` — the underlying `GraphBuilder`/`CompiledGraph` - runtime this rendition is built from. -- `docs/modules/graph/subagents-recursion.md` — how a compiled graph (this - one included) nests inside the recursive sub-agent/subgraph architecture. -- `docs/runtime-comparison/langgraph.md` §4 and - `docs/runtime-comparison/pydantic-ai.md` §4 — the comparative design notes - A5 responds to. diff --git a/docs/modules/harness/store.md b/docs/modules/harness/store.md index 1a518ad38..550fe9e47 100644 --- a/docs/modules/harness/store.md +++ b/docs/modules/harness/store.md @@ -3,16 +3,10 @@ The store feature provides durable and ephemeral storage for the harness. It is used by memory, events, model calls, tool calls, artifacts, web UIs, and tests. -The store is not graph checkpointing. Graph checkpoints belong to the graph -module. Harness stores record application/runtime data around LLM orchestration. -When the harness state-graph runtime is enabled, graph checkpointers own the -serialized graph state snapshots while harness stores can still index the run's -events, artifacts, model calls, tool calls, and UI-facing summaries. - -The store may also back status and event-journal features for the harness and -for graph integrations. That does not make those records generic memory: status -records are operational snapshots, event records are append-only execution -history, and checkpoints remain graph-owned resumable state. +Harness stores record application/runtime data around LLM orchestration: +events, artifacts, model calls, tool calls, and UI-facing summaries. Status +records are operational snapshots and event records are append-only execution +history. ## Source Inspiration @@ -38,8 +32,6 @@ runtime injection. TinyAgents should expose stores through `RunContext` and - Provide append-only stream storage. - Store messages, run records, events, tool records, model records, and artifacts. - Store embedding records and vector/retrieval metadata when configured. -- Store graph run indexes, graph event journals, and graph UI summaries when - configured, without taking ownership of checkpoint snapshots. - Support in-memory tests. - Support JSONL local development. - Support MongoDB server deployments. @@ -56,8 +48,6 @@ runtime injection. TinyAgents should expose stores through `RunContext` and - It does not decide what enters a prompt. - It does not summarize memory. -- It does not checkpoint graph execution. -- It does not replace the graph checkpointer used for pause/resume. - It does not replace the registry. - It does not require one global database. @@ -300,9 +290,7 @@ Store-backed observability uses two different record families: - event journal records are append-only and answer "what happened?" Harness status records should store the compact `HarnessRunStatus` shape from -[observability](observability.md). Graph integrations may store graph status -records in the same backend, but graph-specific fields such as checkpoint id, -checkpoint namespace, step, active nodes, and interrupts remain graph-owned. +[observability](observability.md), separate from append-only event history. Event journal records should be addressable by stream and offset: diff --git a/docs/modules/harness/streaming.md b/docs/modules/harness/streaming.md index ea6c2fa60..2a126b964 100644 --- a/docs/modules/harness/streaming.md +++ b/docs/modules/harness/streaming.md @@ -1,8 +1,6 @@ # Harness Streaming Feature -The harness must support streaming independently from the graph. Graph streaming -can forward harness events, but direct harness users should also be able to -consume streams. +The harness supports caller-consumable run events and provider model streams. ## Responsibilities @@ -41,8 +39,8 @@ tracer streams: - `events`: all harness events - `final`: final result only -Every stream item should carry run ids and component ids so web UIs can merge -harness streams with graph streams. +Every stream item should carry run ids and component ids so web UIs can identify +the run that emitted it. ## Stream Items @@ -77,8 +75,8 @@ Streaming adapters must merge chunks deterministically: ## Implementation status (runtime-comparison Phase 3, C1/C2) The design above names the target shape; here is what exists in code today, -one layer down at `tinyinference_llm::model::ModelStreamItem` and -`tinyagents_harness::stream`. +one layer down at `tinyinference_llm::model::ModelStreamItem` and the +`agent_loop::stream` caller API. **Block-indexed streaming (C1, vendor `tinyinference-llm`).** Content block indexes are respected via new `ModelStreamItem` variants: `BlockStart { @@ -106,56 +104,3 @@ specifically `ProviderError` — now carries `partial_message: Option<AssistantMessage>` and `stop_reason: Option<String>` on both adapters, so a mid-stream failure does not discard whatever content had already arrived. - -**Frame codec (C2, `crates/tinyagents-harness/src/stream/frame.rs`).** -`AssistantFrame` is the durable, journal-friendly encoding of a block-aware -stream: `FrameEncoder` turns a sequence of `ModelStreamItem`s into frames -(only the block-indexed and terminal items are framed — `Started` and the -flat compatibility deltas carry no information a reducer needs beyond what -the block items already have), emitting a periodic -`AssistantFrame::ToolArgsCheckpoint { index, json_so_far }` full snapshot for -long tool-argument streams. `reduce_frames(&[AssistantFrame]) -> -PartialAssistantMessage` folds a — possibly truncated — frame sequence back -into closed `content` blocks plus whatever `open_blocks` were still -in-progress, without needing the original provider stream. A checkpoint is a -full snapshot, not a delta, so a reader that only has frames from a -checkpoint onward (earlier per-fragment frames pruned from the journal) still -reduces to a consistent result. - -**Mid-tool progress (C2).** A running tool reports through -`tinytools::ToolRunContext::report_progress(ToolProgress { message, fraction, -partial })` (a default no-op, so existing tools are unchanged). The loop gives -every executing call a progress gate (`crates/tinyagents-harness/src/tool/progress/mod.rs`) -that emits `AgentEvent::ToolProgressDetail { call_id, message, fraction, partial }` -live and queues a `ToolDelta` for `Middleware::on_tool_delta`. Guarantees: - -- Every `ToolProgressDetail` for a call falls between its `ToolStarted` and its - terminal `ToolCompleted`/`ToolFailed`, in the order the tool reported it. -- The gate closes when the tool's future settles (return, error, timeout), - before the terminal event. An update reported afterwards (for example from a - task the tool spawned) is dropped, never emitted late; the open check and the - emit share one lock, so there is no race window. -- A concurrent batch interleaves progress across calls; each call's progress - still precedes its own terminal event (terminals are emitted by the fold, - in call order, after the batch). -- `on_tool_delta` needs `&mut RunContext`, which the executing tool holds, so - middleware sees a call's deltas replayed in order immediately after the call - settles and before `after_tool` and the terminal event. It observes; the live - event is already out. A failing hook is logged and does not fail the call. -- The gate also closes if the dispatch future is dropped (run cancelled or - timed out mid-call), so a task the tool spawned cannot emit afterwards. -- Flooding is coalesced: at most 32 events per second per call (fixed, - crate-private `ToolProgressLimits`); beyond that the newest value of each - field replaces the held update. There is no timer: the held update is emitted - on the first update of the next window or when the call settles (that final - flush is exempt from the window limit), so the final reported state is never - lost. -- The middleware replay queue is bounded (newest 64 deltas, each `content` - capped at 4 KiB, partials serialized only up to the cap) and is not filled at - all when the run has no middleware. -- An event listener must not call `report_progress` re-entrantly (events are - emitted under the call's gate lock). - -The typed -`HarnessStreamItem` enum, `StreamMode::{tools, usage, cost, events, final}`, -and stream replay from event stores are still design-only. diff --git a/docs/modules/harness/subagent-steering.md b/docs/modules/harness/subagent-steering.md index a42f69341..e3c2ea1cb 100644 --- a/docs/modules/harness/subagent-steering.md +++ b/docs/modules/harness/subagent-steering.md @@ -2,9 +2,8 @@ Steering is runtime control sent to an already-created agent run. It is broader than a prompt edit and narrower than mutating arbitrary state. TinyAgents should -support steering at both the harness level and graph level so parent -orchestrators, humans, tests, and UIs can guide long-running work without -breaking run hierarchy, checkpoints, or observability. +support harness-level steering so parent orchestrators, humans, tests, and UIs +can guide long-running work without breaking run hierarchy or observability. ## Steering Relationships @@ -13,19 +12,18 @@ Supported steering relationships: - parent orchestrator agent steers a sub-agent - human steers a sub-agent - human steers a parent orchestrator agent -- graph supervisor node steers child sub-agent tasks - middleware steers an agent loop under explicit policy - test harness steers fake agents deterministically Every steering operation must preserve `root_run_id`, `parent_run_id`, -`thread_id`, `agent_id`, `task_id` when graph-backed, source actor, source -scope, and checkpoint or event offsets when available. +`thread_id`, `agent_id`, `task_id` when present, source actor, source scope, +and event offsets when available. ## Non-Goals - Steering is not hidden prompt injection. - Steering is not direct mutation of another run's private memory. -- Steering is not a way to bypass model, tool, graph, store, or budget +- Steering is not a way to bypass model, tool, store, or budget allowlists. - Steering is not an untracked side channel; every accepted steering instruction emits events and appears in run inspection. @@ -47,15 +45,12 @@ pub struct SteeringCommand { pub enum SteeringTarget { AgentRun(RunId), SubAgentRun { parent_run_id: RunId, child_run_id: RunId }, - GraphTask { run_id: RunId, task_id: TaskId }, - GraphNode { run_id: RunId, node_id: NodeId }, Thread(ThreadId), } pub enum SteeringActor { Human { user_id: String }, ParentAgent { run_id: RunId, agent_id: ComponentId }, - GraphSupervisor { run_id: RunId, node_id: NodeId }, Middleware { name: String }, Test, } @@ -85,8 +80,7 @@ commands. ## Harness-Level Steering -Harness-level steering applies to direct model calls and model-tool loops that -are not graph-backed. +Harness-level steering applies to direct model calls and model-tool loops. Required behavior: @@ -171,25 +165,6 @@ opts in with `SteeringPolicy::allow_all_with_model_switch()` or - The orchestration `steer_agent` tool does not expose `switch_model`; it is a host-level control for now. -## Graph-Level Steering - -Graph-backed steering is a command against graph state and task scheduling. It -should use the graph runtime whenever the target run has checkpoints, active -tasks, interrupts, child sub-agents, or subgraphs. - -Graph-level steering can lower to: - -- `Command::resume` for an interrupt -- `Command::goto` for an explicit redirect -- `Command::update` for a reducer-mediated state update -- `Command::parent` for child-to-parent handoff -- task cancellation through graph policy -- a queued steering event delivered to a child sub-agent node - -Graph steering must be checkpoint-aware. If a child task is paused, a steering -command targets the child namespace and resumes only that child unless policy -explicitly escalates to the parent run. - ## Parent Orchestrator Steering A parent orchestrator may steer sub-agents it created when policy grants that @@ -202,21 +177,20 @@ orchestrator -> spawn research_subagent orchestrator -> steer research_subagent: "focus only on billing evidence" orchestrator -> steer critic_subagent: "compare answer against policy version 3" orchestrator -> wait research_subagent, critic_subagent -orchestrator -> merge child outputs through graph reducers +orchestrator -> merge child outputs under host policy ``` Rules: - the parent can steer only descendants in its run tree by default -- sibling agents cannot steer each other unless a graph supervisor grants it +- sibling agents cannot steer each other unless host policy grants it - steering cannot expand a child's tool/model allowlist beyond the child policy - parent steering is visible to the child as a structured instruction with provenance, not as anonymous user text - the child may reject steering that violates its policy and must emit a rejection event -- parent steering and child output both merge through normal graph or harness - state contracts +- parent steering and child output follow the harness run contracts ### Child Limit Coordination @@ -225,18 +199,11 @@ hallucinate indefinitely or loop forever. Harness-level caps such as `RunLimits::max_model_calls`, `max_tool_calls`, wall-clock deadlines, and `max_depth` are enforced by the child run. -When a parent calls a sub-agent through `SubAgentTool`, child limit exhaustion is -reported back as an error `ToolResult` instead of crashing the parent run. The -tool result content must tell the parent orchestrator that the delegated -sub-agent stopped because it hit a configured limit, not because it completed -the task. The parent can then narrow the task, split it across more workers, -retry with a different budget, or report the partial failure. - -Direct sub-agent APIs (`SubAgent::invoke`, `invoke_with_events`, -`invoke_in_parent`, and `SubAgentSession::send`) still return the original -classified error. That lets application code decide whether a direct child -limit should fail the caller, be retried, or be converted into a graph/state -transition. +`DriverSubAgentTool` runs child work asynchronously and immediately returns a +job reference. A child limit is recorded as an incomplete job with its typed +reason; the host can query or await that job separately and decide whether to +retry, narrow the task, or report the partial result. The driver does not run +child work inline in the parent model tool call. ### Agent-origin text is fenced, capped and labelled @@ -261,23 +228,21 @@ Human steering modes: boundary - policy edit: narrow model/tool/budget limits for a run - pause/cancel: stop a run or child task cooperatively -- redirect: send a graph-backed run to a review, retry, or finalize node -- state update: apply a reducer-mediated correction to graph state +- redirect: change the host-defined next action for a run Human steering must be auditable: - actor id -- target run/task/node +- target run/task - timestamp - accepted/rejected/deferred status - policy that allowed the action -- model/tool/graph operation affected +- model/tool operation affected - redacted payload for UI and event export Human steering of sub-agents should not require steering the parent first. A UI -must be able to target a specific child run or graph task using run-tree -metadata. The parent is notified through events and rollups, not by opaque -shared state. +must be able to target a specific child run using run-tree metadata. The parent +is notified through events and rollups, not by opaque shared state. ## Delivery And Conflict Rules @@ -295,7 +260,7 @@ Conflict rules: - human `ConstrainTools` can narrow but not widen a parent-granted tool set. - parent `AddInstruction` appends to the child steering context unless a policy allows replacement. -- simultaneous redirects require a graph policy: first accepted, highest +- simultaneous redirects require an explicit policy: first accepted, highest priority, or explicit conflict error. - stale steering commands that refer to a completed checkpoint are rejected unless they target a new fork. @@ -313,9 +278,9 @@ Required events: - `steering.superseded` - `steering.failed` -Events must include root run id, parent run id, target run id, target task/node -when graph-backed, actor, kind, correlation id, policy id, checkpoint id when -available, and redacted payload summary. +Events must include root run id, parent run id, target run id, target task when +available, actor, kind, correlation id, policy id, and redacted payload +summary. ## Testkit @@ -327,5 +292,4 @@ The testkit should include: - rejected steering because of missing authority - rejected tool allowlist expansion - deterministic ordering of simultaneous steering commands -- graph-backed targeted child resume - harness-loop steering delivered at the next safe boundary diff --git a/docs/modules/harness/terminal-outcome.md b/docs/modules/harness/terminal-outcome.md index 07816a5da..ad656b992 100644 --- a/docs/modules/harness/terminal-outcome.md +++ b/docs/modules/harness/terminal-outcome.md @@ -11,10 +11,10 @@ record whether the deadline landed before a provider call, during one, or after a call. Provider-start state is tracked separately so a host can tell a first-call failure from a preflight failure. -The direct and graph loop drivers publish the outcome on `AgentRun` before -running `after_agent` middleware. This lets middleware observe the same -classification as the terminal event. Session hooks receive the typed outcome -before the legacy terminal notification. +The direct loop publishes the outcome on `AgentRun` before running +`after_agent` middleware. This lets middleware observe the same classification +as the terminal event. Session hooks receive the typed outcome before the +legacy terminal notification. Turn lifecycle events announce appended messages and retract only messages that were previously announced. Initial input messages are treated as a seed @@ -35,28 +35,11 @@ crate ends; it is `None` only when deserializing journals written before the field existed. Scope: the typed outcome and the turn and message lifecycle events -(`TurnStarted`, `TurnCompleted`, `MessageAppended`) are published by both the -direct loop and the graph loop driver, with the same semantics: input messages -are never announced, each appended message is announced exactly once, and a -nested tool call never produces a `MessageAppended`. The graph driver announces -appends at the same points as the direct loop (before a turn's `ModelStarted`, -after the assistant reply, after a tool batch, and at the end of the run, -including after `after_agent`). `MessageRetracted` and `TranscriptRewritten` -come from direct-loop recovery paths the graph rendition does not implement. The -step-by-step `LoopIter` and `compile_loop` graph announce the same events and -treat the transcript present at their first node activation (any node) as the seed. A node that -interrupts discards its state and re-runs on resume, so the appends it announced -are retracted with `MessageRetracted` first. A model node's turn is closed on an -interrupt; a tools node leaves its turn open, and the re-run closes it with the -real results. A fresh runtime resuming from a checkpoint continues the turn -numbering from the checkpoint and re-opens the in-flight tool turn (without a -second `TurnStarted`). - -Known difference from the direct loop: the graph rendition closes a turn before -an output-retry prompt is announced (the direct loop announces the prompt first). -A model or tools node that fails closes the turn it opened (and announces the -results of calls that ran before the failure), so `LoopIter` and `compile_loop` -need no driver epilogue. +(`TurnStarted`, `TurnCompleted`, `MessageAppended`) are published by the direct +loop. Input messages are never announced, each appended message is announced +exactly once, and a nested tool call never produces a `MessageAppended`. +`MessageRetracted` and `TranscriptRewritten` describe direct-loop recovery +paths. ## `provider_started` and summarizers diff --git a/docs/modules/harness/testkit.md b/docs/modules/harness/testkit.md index 2b3975f23..47022328b 100644 --- a/docs/modules/harness/testkit.md +++ b/docs/modules/harness/testkit.md @@ -30,7 +30,6 @@ adapters can assert behavior without live providers. - Provide in-memory stores and event collectors. - Provide deterministic clocks and id generators. - Provide trajectory assertions for agent loops. -- Provide graph trajectory assertions for state-graph runs. - Provide conformance tests for provider adapters. - Provide replay fixtures from recorded events. - Provide helpers for structured-output validation tests. @@ -105,8 +104,6 @@ Tests should be able to assert: - no unexpected provider calls - retrieved document ids and scores - indexed document hashes -- graph run status -- graph node transition sequence - checkpoint labels and state snapshots - HITL interrupt kind/options/resume node @@ -131,21 +128,5 @@ Provider adapters should run a shared suite that covers: - timeout - rate-limit error classification -State-graph conformance should cover: - -- linear route -- conditional route -- fork and merge reducer -- cycle guarded by max steps -- cancellation -- node error propagation -- HITL pause/resume/reject -- in-memory and durable checkpointer round trip -- run list pagination and newest-first order -- per-agent blueprint validation -- live turn graph tool advertisement, allowlist enforcement, malformed - argument handling, and final output depending on actual tool execution -- provider options passthrough - Tests requiring live providers should be opt-in and documented with environment variables. Unit-level conformance should run with fakes in normal CI. diff --git a/docs/modules/harness/usage.md b/docs/modules/harness/usage.md index 70930304e..d540966d2 100644 --- a/docs/modules/harness/usage.md +++ b/docs/modules/harness/usage.md @@ -11,7 +11,7 @@ Usage accounting normalizes token and cache usage across providers. - Track embedding input tokens, characters, vector counts, dimensions, and batch counts when providers expose them. - Track tool/model call counts. -- Roll usage up by call, node, agent, graph, thread, and run. +- Roll usage up by call, agent, thread, and run. - Emit usage events. - Preserve provider-specific token details. - Mark provider-reported versus locally estimated usage. diff --git a/docs/modules/registry/README.md b/docs/modules/registry/README.md index c0eba4c39..b3fb38f4b 100644 --- a/docs/modules/registry/README.md +++ b/docs/modules/registry/README.md @@ -1,7 +1,7 @@ # Registry Module Specification The registry module is the coordination and discovery layer for TinyAgents. It -registers agents, graphs, tools, models, stores, middleware, listeners, and +registers agents, tools, models, stores, middleware, listeners, and runtime data catalogs, then provides an event-friendly surface for web UIs, CLIs, tests, and distributed supervisors. @@ -21,7 +21,6 @@ server. ## Responsibilities - Register named agents. -- Register compiled graphs. - Register tools. - Register model providers. - Register middleware and stores. @@ -32,17 +31,13 @@ server. - Validate duplicate or incompatible registrations. - Emit registration, lifecycle, and execution events. - Route events to outside listeners. -- Support parallel agent and graph execution with correlation ids. - Provide a test recorder for event assertions. - Provide local lookup for model prices, capabilities, and context windows. ## Non-Responsibilities -- It does not execute graph nodes itself. - It does not call LLM providers itself. - It does not validate tool arguments beyond registry-level schema checks. -- It does not persist checkpoints unless a checkpointer is registered as a - component. - It does not force a singleton global registry. - It does not guarantee live provider pricing without an explicit snapshot refresh. @@ -59,7 +54,6 @@ crates/tinyagents-registry/src/ component.rs discovery.rs events.rs - graph.rs listener.rs model.rs names.rs @@ -90,7 +84,6 @@ The registry owns named component handles and local metadata catalogs: ```rust pub struct Registry<State, Ctx = ()> { agents: AgentRegistry<State, Ctx>, - graphs: GraphRegistry<State, Ctx>, models: ModelRegistry<State, Ctx>, tools: ToolRegistry<State, Ctx>, stores: StoreRegistry, @@ -101,12 +94,11 @@ pub struct Registry<State, Ctx = ()> { } ``` -Runtime execution still belongs to the harness and graph modules. The registry +Agent execution belongs to the harness. The registry answers questions such as: - Which model aliases are available? - Which tool schemas are available? -- Which graph ids can be invoked? - What are the known context-window and price details for this model? - Which events should be emitted and where should they go? @@ -119,7 +111,7 @@ live handles. ### Component kinds `ComponentKind` partitions the registry namespace and now has **12** variants. -Alongside `Model`, `Tool`, `Graph`, `Router`, `Reducer`, `Store`, `Agent`, and +Alongside `Model`, `Tool`, `Router`, `Reducer`, `Store`, `Agent`, and four kinds cover the runtime's durable roles: diff --git a/docs/modules/registry/design.md b/docs/modules/registry/design.md index 25382397b..db6d68dd6 100644 --- a/docs/modules/registry/design.md +++ b/docs/modules/registry/design.md @@ -10,7 +10,7 @@ Parent module: [Registry](README.md). The registry module is the coordination layer for TinyAgents. It registers -agents, graphs, tools, models, stores, middleware, and listeners, then provides +agents, tools, models, stores, middleware, and listeners, then provides an event-friendly runtime surface for external systems such as web UIs, CLIs, logs, tests, and distributed supervisors. @@ -42,7 +42,6 @@ Primary code references: ## Responsibilities - Register named agents. -- Register compiled graphs. - Register tools. - Register model providers. - Register model aliases, resolver policies, and agent model defaults. @@ -53,16 +52,12 @@ Primary code references: - Validate duplicate or incompatible registrations. - Emit registration, lifecycle, and execution events. - Route events to outside listeners. -- Support parallel agent and graph execution with correlation ids. - Provide a test recorder for event assertions. ## Non-Responsibilities -- It does not execute graph nodes itself. - It does not call LLM providers itself. - It does not validate tool arguments beyond registry-level schema checks. -- It does not persist checkpoints unless a checkpointer is registered as a - component. - It does not force a singleton global registry. ## Package Shape @@ -77,7 +72,6 @@ crates/tinyagents-registry/src/ component.rs discovery.rs events.rs - graph.rs listener.rs model.rs names.rs @@ -96,7 +90,6 @@ The registry owns named component handles: ```rust pub struct Registry<State, Ctx = ()> { agents: AgentRegistry<State, Ctx>, - graphs: GraphRegistry<State, Ctx>, models: ModelRegistry<State, Ctx>, tools: ToolRegistry<State, Ctx>, stores: StoreRegistry, @@ -115,7 +108,6 @@ pub type SharedRegistry<State, Ctx = ()> = Arc<Registry<State, Ctx>>; The registry should be passable into: - harness runs -- graph runs - tool calls - web server request handlers - background workers @@ -132,7 +124,7 @@ ComponentRef -> component-owned runtime events ``` -Registry operations emit registry lifecycle events. Model, tool, graph, and +Registry operations emit registry lifecycle events. Model, tool, and agent execution spans are emitted by the component runtimes. ## Component Identity @@ -149,7 +141,6 @@ pub struct ComponentId { pub enum ComponentKind { Agent, - Graph, Model, Tool, Store, @@ -206,7 +197,6 @@ pub struct ComponentMetadata { Examples: - a tool exposes input/output schema -- a graph exposes node and edge summary - an agent exposes model name and tool names - a listener exposes supported event filters @@ -248,8 +238,7 @@ Use direct `Arc<dyn Trait>` entries for simple tools and stores. ## Agent Registration -Agents are executable harness configurations. They may be simple model-tool loops -or wrappers around compiled graphs. +Agents are executable harness configurations, typically model-tool loops. ```rust #[async_trait] @@ -308,49 +297,8 @@ The registry stores these declarations; the harness applies them during a run. This keeps model selection inspectable without making the registry call model providers. -## Graph Registration - -Graphs are compiled graph runtimes. - -```rust -pub struct RegisteredGraph<State, Ctx = ()> { - pub metadata: ComponentMetadata, - pub graph: Arc<CompiledGraph<State, Ctx>>, -} -``` - -Registration: - -```rust -registry - .graphs() - .register("support_flow", compiled_graph) - .await?; -``` - -Graph metadata should include: - -- node ids -- edge count -- route names -- start node -- whether checkpoints are required -- whether interrupts are possible - declared input/output schemas when available -Graph events: - -- `graph.registered` -- `graph.started` -- `graph.step_started` -- `graph.node_started` -- `graph.node_completed` -- `graph.route_selected` -- `graph.interrupted` -- `graph.checkpoint_saved` -- `graph.completed` -- `graph.failed` - ## Tool Registration Tools are named callable capabilities. The registry should own normalized name @@ -471,6 +419,6 @@ state, events, checkpoints, usage, and cost rows. --- -Continues in [`events.md`](events.md) (store/checkpointer registration, the +Continues in [`events.md`](events.md) (store registration, the event model, event bus, and filters) and [`operations.md`](operations.md) (lifecycle, discovery, error model, testkit, and milestones). diff --git a/docs/modules/registry/events.md b/docs/modules/registry/events.md index 7ad4c5dc7..527d6f6e7 100644 --- a/docs/modules/registry/events.md +++ b/docs/modules/registry/events.md @@ -6,21 +6,19 @@ > [`implementation-status.md`](implementation-status.md) for what is actually > shipped. -Continues from [`design.md`](design.md): store and checkpointer -registration, listener registration, the event model, event bus, and +Continues from [`design.md`](design.md): store registration, listener +registration, the event model, event bus, and event filters. ## Store And Checkpointer Registration -Stores and checkpointers should be registry components so agents and graphs can -share them without globals. +Stores should be registry components so agents can share them without globals. Registered stores: - short-term memory - long-term key/value store - vector store -- graph checkpointer - event sink Store events: @@ -136,7 +134,6 @@ Event payloads: ```rust pub enum EventPayload { Agent(AgentEvent), - Graph(GraphEvent), Model(ModelEvent), Tool(ToolEvent), Store(StoreEvent), @@ -162,10 +159,6 @@ on_agent_start on_agent_stream on_agent_end on_agent_error -on_graph_start -on_graph_stream -on_graph_end -on_graph_error on_tool_start on_tool_stream on_tool_end @@ -230,7 +223,6 @@ pub struct EventFilter { Example filters: - all events for one `run_id` -- only graph node events - only tool failures - only events tagged `tenant:acme` - only model token deltas for streaming text diff --git a/docs/modules/registry/implementation-status.md b/docs/modules/registry/implementation-status.md index e6551415e..20de2f4d8 100644 --- a/docs/modules/registry/implementation-status.md +++ b/docs/modules/registry/implementation-status.md @@ -114,43 +114,13 @@ in Phase 1c — see below. an explicit default instead, returning `TinyAgentsError::ModelNotFound` if `name` is not registered. -## Capability bundle (gap G3) +## Capability data -`CapabilityRegistry::register_capability(name, Capability<State, Ctx>)` -(`capability/mod.rs`) registers a -`tinyagents_harness::capability::Capability<State, Ctx>` — instructions + -toolset + middleware + model defaults + exposure + `defer_loading`, the -runtime-level unit Pydantic AI v2's `AbstractCapability` occupies -(`docs/runtime-comparison/pydantic-ai.md` §4). The bundle type itself lives -in `tinyagents-harness::capability` rather than in `tinyagents-registry` or -the dependency-free `tinyagents-definition` crate: it composes -`tool::toolset::ToolSet<State, Ctx>` and `middleware::Middleware<State, Ctx>` -trait objects that are native to `tinyagents-harness`, and -`tinyagents-definition` has zero dependency on `tinyagents-harness` by -design (that asymmetry is what keeps the crate graph acyclic). Since -`CapabilityRegistry<State>` carries no `Ctx` type parameter (none of its -other stored kinds need one), registry storage for this specific -`Ctx`-generic bundle goes through type-erased `Box<dyn Any>` storage instead -of adding a `Ctx` parameter to the whole registry for one feature; see -`capability/mod.rs`'s `register_capability`/`capability` doc comments for the -erasure mechanics. - -`Capability::from_spec(serde_json::Value)` builds a bundle from data (no -toolset/middleware — those are still wired programmatically). Hosts choose -which JSON-declared capabilities to install and retain control of the -available tool, model, and subagent handles. - -On the harness side, `AgentHarness::with_capability(capability)` installs a -bundle: its middleware is appended in installation order, its model defaults -are applied onto the harness's `RunPolicy`, and its toolset is -folded into a `CapabilityToolSet` combined with whatever toolset was already -installed via `with_toolset` before the first `with_capability` call. When -any installed capability has `defer_loading: true`, a synthetic -`load_capability` tool (`LOAD_CAPABILITY_TOOL_NAME`) is auto-registered so a -model can bring a deferred capability's tools into scope on demand; the -resulting mid-run tool-set change is recorded through the B6 transcript -patch mechanism (`agent_loop/tool_changes.rs`) rather than silently -reshaping the next request. +`tinyagents_registry::Capability<State, Ctx>` is the registry-owned value for +named instructions, toolsets, and middleware. The registry keeps it type-erased +because each value has its own `Ctx` type. It is stored data only; the retired +harness capability loader, model defaults, exposure policy, and deferred +`load_capability` tool are no longer part of TinyAgents. ## Why the gap diff --git a/docs/modules/registry/operations.md b/docs/modules/registry/operations.md index 924db08a8..be1547226 100644 --- a/docs/modules/registry/operations.md +++ b/docs/modules/registry/operations.md @@ -21,7 +21,7 @@ Static components: - registered before a run - listed in discovery APIs - have durable ids -- validated when graphs compile +- validated before an agent run Runtime components: @@ -38,7 +38,7 @@ name and produce a registered-like descriptor. Parallel execution requires explicit hierarchy. -When an agent spawns child agents or graph branches: +When an agent spawns child agents: - every child run gets a new `run_id` - every child run inherits `root_run_id` @@ -51,9 +51,9 @@ Example: ```text root run: support_agent - child run: retrieval_graph - node: search - node: rerank + child run: retrieval_agent + model: default + tool: search child run: draft_agent model: default tool: lookup_user @@ -72,7 +72,6 @@ GET /runs/{run_id}/events GET /threads/{thread_id}/runs GET /events/stream?run_id=... POST /agents/{name}/invoke -POST /graphs/{name}/invoke ``` The Rust library should not ship a web server initially, but the event and @@ -82,7 +81,6 @@ UI event needs: - stable component ids - display names -- graph topology metadata - tool schemas - run tree correlation - streaming model deltas @@ -109,11 +107,10 @@ Examples: - tool-call timeline transformer - subagent tree transformer - model-token text transformer -- graph-state diff transformer - cost/usage accumulator Transformers should be registry extensions. They should not be hardcoded into -every tool, graph, or model implementation. +tool or model implementations. ## Redaction And Safety @@ -149,7 +146,6 @@ Registration should be explicit and validated. ```rust registry.register_tool(tool).await?; registry.register_model("default", model).await?; -registry.register_graph("support_flow", graph).await?; registry.register_agent("support_agent", agent).await?; registry.register_listener(websocket_listener).await?; ``` @@ -225,7 +221,6 @@ listener is marked required. - event recorder listener - event snapshot assertions - fake tool registration -- fake graph registration - fake agent registration - deterministic event ids - deterministic timestamps @@ -257,7 +252,6 @@ recorder.assert() - `ComponentMetadata` - `ToolRegistry` - `ModelRegistry` -- `GraphRegistry` - duplicate validation ### R2: Event Envelope @@ -274,10 +268,9 @@ recorder.assert() - event recorder - stdout listener -### R4: Agent And Graph Registration +### R4: Agent Registration - registered agent trait -- registered graph wrapper - lookup and invoke helpers ### R5: Parallel Run Correlation diff --git a/docs/plans/2026-10-19-close-observability-host-seams.md b/docs/plans/2026-10-19-close-observability-host-seams.md new file mode 100644 index 000000000..931e2d21e --- /dev/null +++ b/docs/plans/2026-10-19-close-observability-host-seams.md @@ -0,0 +1,377 @@ +# Plan: move OpenHuman close and observability seams into TinyAgents + +Issue: [senamakel/tinyagents#64](https://github.com/senamakel/tinyagents/issues/64), part of [tinyhumansai/openhuman#8088](https://github.com/tinyhumansai/openhuman/issues/8088). + +This plan targets the current `epic/8088-tinyagents-slim` fork branch. It is based on the TinyAgents fork at `aca9fd54666ce8fdcc02ae4d6499baa2a3354cd5` and the OpenHuman super-worktree's current source. It does not assume future interfaces from issue #60 (session reuse), #63 (session phases), or OpenHuman #212. Recheck this plan against the merged #60 integration before starting implementation; merge/review intervening fork changes and adapt the plan without discarding either side's behavior. + +## Goal and boundaries + +Move generic close, turn outcome, journal, trace, no-progress ladder, and model-budget behavior from OpenHuman's TinyAgents adapter into `tinyagents-harness`. Keep OpenHuman-specific model sourcing, stop taxonomy/copy, recovery headroom, terminal-inference fast-halt policy, security and sandbox decisions, provider usage accounting, and persistence adapters in OpenHuman. No `openhuman_core` types or dependencies may enter TinyAgents. + +The current harness already owns `close.rs`, `turn_stop.rs`, `turn_outcome.rs`, `observability` (including `StoreEventJournal`), `no_progress`, and `cost::ModelBudget`. Treat existing harness implementations as the baseline and extend or consolidate them; do not duplicate equivalent algorithms under new names. Issue #64's “rows are gone” means the generic implementations and duplicate shims disappear from the specified OpenHuman source directories once the host adapts to existing/new upstream APIs. + +Required user-visible invariant: using OpenHuman's copy injection produces byte-for-byte identical checkpoint, deterministic final-close, and stop-summary output to current OpenHuman. Port the existing host golden assertions before deleting their source. Keep every behavior assertion unless the corresponding capability is expressly removed by this issue; none of the listed capabilities is being removed. + +## Repository-grounded ownership and dependencies + +| Existing OpenHuman source | Destination / action | Host behavior retained | +|---|---|---| +| `crates/openhuman-core/src/agent/session_host/driver/grounded_close.rs`: `RepairUsage` (30-98), `SummarizerSourceModel` (116-225), `classified_halt_close` (235-249), `repair_required_output` (253-334), and `close_if_needed` (365-464) | Move generic close accounting, source-model decorator, halt-close selector, required-output repair and orchestration to `crates/tinyagents-harness/src/close.rs`; use host source/cost ports below | `TurnModelSource` implementation, provider/model selection, thread context and output token cap, native `BilledUsage` conversion/catalog `CallCost`, provider pricing lookup, progress event delivery, host `RequiredOutput` construction and `Prepend` fallback selection | +| `agent/session_host/turn_checkpoint.rs`: result truncation/rendering, harness-instruction matching, duplicate `CloseViolation`, deterministic checkpoint/final summary, repair and verifier prompt assembly | Reuse `harness::close`; delete duplicate helpers after call sites migrate | OpenHuman's `CloseCopy` implementation containing product strings and exact existing rendering behavior | +| `agent/session_host/stop_summary.rs`: wrappers (36-82), `StopReasonKind` (85-220), scrubbed grouping/renderer (223-319) | Move generic bounded grouping/rendering to `harness::turn_stop` | Product `StopReasonKind` parser and failure-class mapping, stable `turn_stop.*` keys, scrubber, host `StopCopy` implementation | +| `agent/tinyagents/turn_outcome.rs`: `DeferredTurnTerminal` (12-94) | Move deferred `TurnStopped`, `RunFailed`, and `RunFinalized` event ownership/emission to `harness::turn_outcome` | `TinyagentsTurnOutcome`, session sidecar plumbing, and durable status calls currently reached through its `TurnJournal` field | +| `agent/tinyagents/journal.rs`: `TurnJournal` (123-225) | No new harness facade: generic event persistence already uses `StoreEventJournal`, `JournalSink`, and `HarnessEventJournal` | Keep status snapshots, `DispatchCtx`, request/run map, storage selection, journal attachment and host replay APIs | +| `agent/tinyagents/supplemental_trace.rs` (13-169) | Move call-id lifecycle and event emission to a generic harness decorator accepting an injected enrichment callback | Keep provider rate lookup, billed usage conversion, cost estimation, pricing-model selection and host usage projection | +| `agent/tinyagents/middleware/tool_outcome_capture.rs` (22-188) | Delete redundant wrapper and use existing harness library `ToolOutcomeCaptureMiddleware` directly | Keep policy-block enrichment, TinyTools failure classification, metadata filtering, and `ToolFailureMap` UI compatibility projection | +| `agent/tinyagents/middleware/repeated_failure.rs` (23-540) | Current source already delegates generic ladder mechanics to harness. Move only a source-audited generic driver with no OpenHuman-specific type dependency; do not add a duplicate ladder. | Keep `HostFailurePolicy`, classification/scope/assessment, sanitization, security and recovery copy, flow body failures, `StopLatch`, steering, additional recoverable-failure headroom, and fast halt for terminal inference failures. | +| `agent/tinyagents/budget.rs` (17-37) `ModelBudget` wrapper | Replace with existing `harness::cost::ModelBudget` (`new`, `wait_for_capacity`, `wrap_model`) | Keep host depth constants, `depth`, `install_depth`, and `TurnModels::with_budget` route wiring | + +Harness already exposes `close`, `cost`, `no_progress`, `observability`, `turn_outcome`, and `turn_stop`. `close.rs` owns tool-result rendering, deterministic summaries, quote detection, repair, and grounded-close state machine. `turn_stop.rs` owns classification and typed stops. `turn_outcome.rs` owns tool outcomes/sinks. Observability owns `StoreEventJournal` over `AppendStore`. Cost owns `ModelBudget::new`, `wait_for_capacity`, and `wrap_model`. `middleware::library::ToolOutcomeCaptureMiddleware` is already available. Add only the copy and source/cost ports needed for issue #64; introduce no OpenHuman/RPC/TinyHumans/storage dependency. + +## Copy seams: complete API contract + +The harness already accepts copy as template arguments in `deterministic_checkpoint`, `deterministic_final_summary`, `close_repair_instruction`, and `GroundedCloseConfig`. Bundle these existing seams rather than creating string fields unrelated to the current renderers. Add the following values in `close.rs` and `turn_stop.rs`: + +```rust +// close.rs +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct CloseCopy { + pub checkpoint_template: &'static str, + pub final_summary_template: &'static str, + pub empty_checkpoint_results: &'static str, + pub empty_final_summary: &'static str, + pub repair_template: &'static str, + pub max_iteration_instruction: &'static str, + pub final_write_instruction: &'static str, + pub final_answer_instruction: &'static str, + pub truncated_answer_instruction: &'static str, + pub stop_note_preamble: &'static str, + pub no_tool_calls: &'static str, + pub repair_no_reply: &'static str, + pub repair_quoted_instruction: &'static str, + pub repair_unverified: &'static str, + pub instruction_needles: &'static [&'static str], + pub minimum_quoted_words: usize, + pub verification_prompt_template: &'static str, +} + +pub const ENGLISH_CLOSE_COPY: CloseCopy = CloseCopy { + checkpoint_template: "I reached the tool-call limit for this turn ({steps} steps), so I paused here before I could write up what I found.\n\n**Raw results from this turn**\n{results}\n**Next steps:** I'll continue from here — just reply (e.g. \"continue\") and I'll pick up where I left off.", + final_summary_template: "I finished this turn without writing up a result. Here is what each tool call returned:\n{results}\nTell me how you'd like to proceed.", + empty_checkpoint_results: "\n- (no tools completed yet)\n", + empty_final_summary: "I finished this turn but produced no result to report.", + repair_template: "{repair_reason} Write the message again.\n\n{instruction}", + max_iteration_instruction: MAX_ITER_CHECKPOINT_INSTRUCTION, + final_write_instruction: FINAL_WRITE_INSTRUCTION, + final_answer_instruction: FINAL_ANSWER_INSTRUCTION, + truncated_answer_instruction: TRUNCATED_ANSWER_INSTRUCTION, + stop_note_preamble: STOP_NOTE_PREAMBLE, + no_tool_calls: "(no tool calls completed)", + repair_no_reply: "Your previous reply was empty or tried to call a tool. No tools will run: write the message itself.", + repair_quoted_instruction: "Your previous reply repeated these directions, or the stop note, back to the user. That text is addressed to you alone and must never appear in the message. Do not reproduce any part of it, do not refer to it, and do not explain that you were given it.", + repair_unverified: "Your previous reply did not pass the check that it reports what actually happened, grounded in the records below.", + instruction_needles: &[FINAL_ANSWER_INSTRUCTION, TRUNCATED_ANSWER_INSTRUCTION, + MAX_ITER_CHECKPOINT_INSTRUCTION, FINAL_WRITE_INSTRUCTION, STOP_NOTE_PREAMBLE], + minimum_quoted_words: 10, + verification_prompt_template: "You are checking a reply before it is shown to a user. Below are the user's request, the records of the tool calls made while handling it, and the reply.\n\nAnswer REJECT if any of these is true:\n1. The reply only says what the assistant will do or is about to do, instead of reporting what happened.\n2. The reply states something the tool records contradict, for example that something does not exist or did not work when a record shows it succeeded, or that something succeeded when its record shows it failed.\n3. The request was not completed, a failed record gives the reason, and the reply does not pass that reason on.\n4. The reply narrates the assistant's own deliberation — thinking aloud, correcting itself part way through, or weighing what to do — rather than stating the outcome. Reporting what a tool call returned is not deliberation.\n5. The reply recites or enumerates the tools the assistant has, instead of reporting what this turn's tool calls produced. Naming a tool the reply actually used, or answering a request that asked what the assistant can do, is not a violation.\n\nOtherwise answer ACCEPT. Reply with the single word ACCEPT or REJECT.\n\n<user_request>\n{user_request}\n</user_request>\n\n<tool_records>\n{records}\n</tool_records>\n\n<reply>\n{reply}\n</reply>", +}; + +pub fn checkpoint_with_copy(copy: &CloseCopy, steps: usize, rendered_results: &str, has_results: bool) -> String { + copy.checkpoint_template + .replace("{steps}", &steps.to_string()) + .replace("{results}", if has_results { rendered_results } else { copy.empty_checkpoint_results }) +} + +pub fn final_summary_with_copy(copy: &CloseCopy, rendered_results: &str, has_results: bool) -> String { + if !has_results { + return copy.empty_final_summary.to_owned(); + } + copy.final_summary_template.replace("{results}", rendered_results) +} + +pub fn final_answer_instruction(copy: &CloseCopy, stop_reason: Option<&str>, truncated: bool, records: &str) -> String { + use crate::middleware::wrap_harness_instruction; + let mut directive = String::new(); + if let Some(reason) = stop_reason { + directive.push_str(copy.stop_note_preamble); + directive.push_str("\n<stop_note>\n"); + directive.push_str(reason.trim()); + directive.push_str("\n</stop_note>\n\n"); + } + directive.push_str(if truncated && stop_reason.is_none() { + copy.truncated_answer_instruction + } else { + copy.final_answer_instruction + }); + let mut out = wrap_harness_instruction(&directive); + out.push_str("\n\n<tool_records>\n"); + out.push_str(if records.trim().is_empty() { copy.no_tool_calls } else { records.trim() }); + out.push_str("\n</tool_records>"); + out +} + +pub fn verification_prompt(copy: &CloseCopy, request: &str, records: &str, reply: &str, request_char_cap: usize) -> String { + let records = if records.trim().is_empty() { copy.no_tool_calls } else { records.trim() }; + interpolate_once(copy.verification_prompt_template, &[ + ("{user_request}", &truncate_chars(request, request_char_cap)), + ("{records}", records), + ("{reply}", reply.trim()), + ]) +} + +// Scan only the template. Inserted values are copied verbatim and never rescanned. +pub(crate) fn interpolate_once(template: &str, fields: &[(&str, &str)]) -> String { + let mut output = String::with_capacity(template.len()); + let mut remaining = template; + while let Some((index, key, value)) = fields.iter() + .filter_map(|(key, value)| remaining.find(key).map(|index| (index, *key, *value))) + .min_by_key(|(index, _, _)| *index) + { + output.push_str(&remaining[..index]); + output.push_str(value); + remaining = &remaining[index + key.len()..]; + } + output.push_str(remaining); + output +} + +pub fn instruction_needles(copy: &CloseCopy, stop_reason: Option<&str>) -> Vec<String> { + let mut sources = copy.instruction_needles.to_vec(); + if let Some(reason) = stop_reason { sources.push(reason); } + harness_instruction_needles(&sources, copy.minimum_quoted_words) +} + +pub fn repair_with_copy(copy: &CloseCopy, instruction: &str, violation: CloseViolation) -> String { + let reason = match violation { + CloseViolation::NoReply => copy.repair_no_reply, + CloseViolation::QuotedInstruction => copy.repair_quoted_instruction, + CloseViolation::Unverified => copy.repair_unverified, + }; + interpolate_once(copy.repair_template, &[("{repair_reason}", reason), ("{instruction}", instruction)]) +} + +// turn_stop.rs +use crate::close::interpolate_once; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct StopCopy { + pub opening_template: &'static str, + pub what_happened_heading: &'static str, + pub success_template: &'static str, + pub failure_template: &'static str, + pub hidden_calls_template: &'static str, + pub proceed_prompt: &'static str, +} + +pub const ENGLISH_STOP_COPY: StopCopy = StopCopy { + opening_template: "I stopped this turn early because {reason}.\n", + what_happened_heading: "\n**What happened**\n", + success_template: "- `{name}` succeeded{count}", + failure_template: "- `{name}` failed{count}{error}", + hidden_calls_template: "- …and {count} more tool call(s)", + proceed_prompt: "\nTell me how you'd like to proceed, or ask me to try again.", +}; + +pub fn english_stop_reason(key: &str) -> &'static str { + match key { + "turn_stop.transient" => "a service I needed kept timing out or was temporarily unavailable", + "turn_stop.uncertain_side_effect" => "an action timed out and I could not confirm whether it went through, so I did not repeat it", + "turn_stop.authentication" => "a service rejected my credentials", + "turn_stop.permission" => "I don't have permission for a step this needs", + "turn_stop.policy" => "a step this needs is blocked by your security settings", + "turn_stop.missing_program" => "a program or feature this needs is not available on this device", + "turn_stop.not_found" => "something this needs could not be found", + "turn_stop.invalid_calls" => "a tool kept rejecting my requests as invalid", + "turn_stop.unavailable" => "a tool this needs is not available right now", + "turn_stop.site_refused" => "a website kept refusing my requests", + "turn_stop.repeated_failure" => "the same step kept failing", + "turn_stop.varied_failures" => "several different tool calls in a row failed", + "turn_stop.repeating_step" => "I kept repeating the same step without getting anything new", + "turn_stop.out_of_credits" => "the account is out of inference credits; add credits, then ask me to try again", + "turn_stop.provider_config" => "the configured model or provider rejected the request; check the model and API key in Connections", + "turn_stop.missing_connection" => "a service this needs isn't connected; connect it in Connections, then ask me to try again", + _ => "my tool calls were not getting anywhere", + } +} + +pub struct StopSummaryGroup { pub name: String, pub success: bool, pub error: Option<String>, pub count: usize } + +pub fn group_stop_results<F>(results: &[CloseToolResult], mut scrubbed_error_line: F) -> Vec<StopSummaryGroup> +where F: FnMut(&str) -> String { + let mut groups: Vec<StopSummaryGroup> = Vec::new(); + for result in results { + let error = if result.success { None } else { Some(scrubbed_error_line(&result.content)).filter(|line| !line.is_empty()) }; + if let Some(group) = groups.iter_mut().find(|g| g.name == result.name && g.success == result.success && g.error == error) { + group.count += 1; + } else { + groups.push(StopSummaryGroup { name: result.name.clone(), success: result.success, error, count: 1 }); + } + } + groups +} + +pub fn render_stop_summary_with_copy(reason: &str, groups: &[StopSummaryGroup], copy: &StopCopy) -> String { + use std::fmt::Write as _; + let mut out = interpolate_once(copy.opening_template, &[("{reason}", reason)]); + if !groups.is_empty() { + out.push_str(copy.what_happened_heading); + let shown_from = groups.len().saturating_sub(8); + for group in groups.iter().skip(shown_from) { + let name = truncate_with_ellipsis(&group.name, 60); + let count = if group.count == 1 { String::new() } else { format!(" {} times", group.count) }; + let line = if group.success { + interpolate_once(copy.success_template, &[("{name}", &name), ("{count}", &count)]) + } else { + let error = group.error.as_ref().map(|e| format!(": {e}")).unwrap_or_default(); + interpolate_once(copy.failure_template, &[("{name}", &name), ("{count}", &count), ("{error}", &error)]) + }; + let _ = writeln!(out, "{line}"); + } + let hidden: usize = groups.iter().take(shown_from).map(|group| group.count).sum(); + if hidden > 0 { + let hidden = hidden.to_string(); + let _ = writeln!(out, "{}", interpolate_once(copy.hidden_calls_template, &[("{count}", &hidden)])); + } + } + out.push_str(copy.proceed_prompt); + out +} + +fn truncate_with_ellipsis(text: &str, max_chars: usize) -> String { + match text.char_indices().nth(max_chars) { + Some((idx, _)) => format!("{}...", text[..idx].trim_end()), + None => text.to_owned(), + } +} +``` + +`StopSummaryGroup.error` contains only the already-scrubbed `short_error_line` output. Move the host grouping and scrub logic into harness and pass the resulting groups to the renderer above. Host `StopCopy` provides the reason sentence through `StopReasonKind.copy()`; it does not move the product's failure-note parser or security scrubber. Keep the existing total-budget result rendering APIs and let OpenHuman supply its existing `CloseResultRenderOptions` block/omission closures. + +Complete `CloseCopy` with these additional literal fields; each initializer value is mechanically moved from the named source constant without rewording: + +| Field | Source value | Golden test | +|---|---|---| +| `max_iteration_instruction` | `turn_checkpoint.rs::MAX_ITER_CHECKPOINT_INSTRUCTION` | checkpoint body and no-completed-tools assertions | +| `final_write_instruction` | `turn_checkpoint.rs::FINAL_WRITE_INSTRUCTION` | `the_final_write_instruction_steers_code_tasks_to_real_edits` and requested-file counterpart | +| `final_answer_instruction` | `turn_checkpoint.rs::FINAL_ANSWER_INSTRUCTION` plus `final_answer_body!()` | final-answer framing/no-deliberation test | +| `truncated_answer_instruction` | `turn_checkpoint.rs::TRUNCATED_ANSWER_INSTRUCTION` plus `final_answer_body!()` | truncated-answer and breaker-truncated tests | +| `stop_note_preamble` | `turn_checkpoint.rs::STOP_NOTE_PREAMBLE` | stop-note carry/quote tests | +| `no_tool_calls` | exact literal `(no tool calls completed)` | close-verification and final-answer record tests | +| `repair_no_reply`, `repair_quoted_instruction`, `repair_unverified` | three arms of `turn_checkpoint.rs::close_repair_instruction` | `the_repair_re_ask_names_the_violation_and_repeats_the_instruction` | +| `verification_prompt_template` | full format string in `turn_checkpoint.rs::close_verification_prompt` | verification-prompt and deliberation/toolset rejection tests | + +Place `interpolate_once` in `close.rs` as `pub(crate) fn` so both copy modules share one implementation; `turn_stop.rs` imports it with `use crate::close::interpolate_once;`. Move the named source constants and `final_answer_body!()` into `close.rs` unchanged before initializing `ENGLISH_CLOSE_COPY`; this makes every referenced initializer expression defined and preserves exact bytes. `CloseCopy` also carries the host instruction-needle sources and `MIN_QUOTED_WORDS = 10`; the harness already implements normalization, span selection, and matching. Its tests use whole-output equality for checkpoint and close prompt paths. These are source inputs required by #64; they are not left for the implementer to rediscover. + +Keep existing APIs source-compatible. Existing `deterministic_checkpoint` and `deterministic_final_summary` remain; add the exact functions above. Existing `CloseViolation` has `NoReply`, `QuotedInstruction`, and `Unverified`; map OpenHuman's `QuotedHarnessText` variant to `QuotedInstruction` at its host boundary. The host passes the copy-aware close renderers explicitly. + +## TDD implementation sequence + +Do each step as a small red/green/refactor slice. Rust tests must be sibling `*_tests.rs` files declared with `#[cfg(test)]` and `#[path = "<module>_tests.rs"] mod tests;`, start with `use super::*;`, and contain no inline tests. For behavior that already exists upstream, first add/port a test that exercises the required public behavior and fails for the actual gap; do not claim that a missing API alone is a behavioral regression. + +### 1. Establish copy compatibility before moving call sites + +Files: `crates/tinyagents-harness/src/close.rs`, `close_tests.rs`, `turn_stop.rs`, `turn_stop_tests.rs`, and corresponding `mod.rs` exports if needed. + +1. Port the exact behavior assertions from OpenHuman `turn_checkpoint_tests.rs`: `the_deterministic_checkpoint_reproduces_each_result_body`, `checkpoint_records_keep_the_captured_tool_failure_status_and_body`, `tool_result_renderer_preserves_the_openhuman_block_format`, `results_over_the_budget_are_omitted_with_a_disclosure`, `one_oversized_result_is_still_shown`, `no_completed_tools_says_so`, `newline_heavy_results_are_charged_what_they_render`, `the_final_summary_quotes_each_failure_message`, `the_final_summary_of_a_halted_turn_gives_a_reason_and_names_the_tools`, `the_final_answer_instruction_carries_the_records_and_the_stop_note`, `the_close_verification_prompt_holds_request_records_and_reply`, `the_final_answer_instruction_is_framed_and_forbids_quoting_tools_and_deliberation`, both quote guards, the repair re-ask, truncated/breaker copy, and both final-write checks. Keep full equality assertions unchanged, including `"\n- `read_file` — ok\n > first line\n > second line\n"`. +2. Port `stop_summary_tests.rs`: exact expected full output from `a_halted_turn_gives_one_reason_and_collapses_identical_failures`, redaction/truncation, secret suppression, successful-repeat semantics, every key/copy, all stop-note categories, policy-refusal copy, and newest-eight retention. The golden has one reason line, `**What happened**`, grouped lines, and the exact final prompt; retain its assertions verbatim. + Add a whole-string assertion for the production three-timeout case with the transient English copy: `I stopped this turn early because a service I needed kept timing out or was temporarily unavailable.\n\n**What happened**\n- \`get_profile\` succeeded\n- \`use_skill\` failed 3 times: tool \`use_skill\` timed out after 30000 ms\n\nTell me how you'd like to proceed, or ask me to try again.` The existing component assertions remain too. +3. First add/run the output behavior tests against existing functions. The OpenHuman block-format golden is a real behavioral gap against the harness's current default renderer; it fails because the harness defaults to inline `- \`name\` (ok):` instead of OpenHuman's quoted multiline block. Do not count adding a new API's compile failure as a behavioral RED. After the algorithmic output test is red, add the sentinel copy test with every field, then implement the new copy API and confirm both pass. +4. Make the red test fail on a real behavior mismatch, then implement. No host migration occurs in #64. Confirm existing APIs and tests still pass unchanged. +5. Add a full-string verifier golden for empty records and repeat it with whitespace-only records. Both calls must equal this exact output (including the no-tools fallback): + ```rust + let expected = "You are checking a reply before it is shown to a user. Below are the user's request, the records of the tool calls made while handling it, and the reply.\n\nAnswer REJECT if any of these is true:\n1. The reply only says what the assistant will do or is about to do, instead of reporting what happened.\n2. The reply states something the tool records contradict, for example that something does not exist or did not work when a record shows it succeeded, or that something succeeded when its record shows it failed.\n3. The request was not completed, a failed record gives the reason, and the reply does not pass that reason on.\n4. The reply narrates the assistant's own deliberation — thinking aloud, correcting itself part way through, or weighing what to do — rather than stating the outcome. Reporting what a tool call returned is not deliberation.\n5. The reply recites or enumerates the tools the assistant has, instead of reporting what this turn's tool calls produced. Naming a tool the reply actually used, or answering a request that asked what the assistant can do, is not a violation.\n\nOtherwise answer ACCEPT. Reply with the single word ACCEPT or REJECT.\n\n<user_request>\nstatus?\n</user_request>\n\n<tool_records>\n(no tool calls completed)\n</tool_records>\n\n<reply>\nNothing ran.\n</reply>"; + assert_eq!(verification_prompt(&ENGLISH_CLOSE_COPY, "status?", "", "Nothing ran.", 200), expected); + assert_eq!(verification_prompt(&ENGLISH_CLOSE_COPY, "status?", " \n\t ", "Nothing ran.", 200), expected); + ``` +6. Add `verification_prompt_does_not_rescan_placeholder_shaped_user_content`: pass request `Explain {records}`, records `record contains {reply} and {user_request}`, and reply `reply contains {records}`; assert the full result retains each of those braces and token strings literally in its corresponding XML section. Add `stop_summary_preserves_placeholder_shaped_tool_names_and_errors`: use one failed group with name `tool_{count}_{error}` and scrubbed error `failed {name}`; assert the exact line is `- `tool_{count}_{error}` failed: failed {name}`. These tests pin single-pass interpolation: scan template tokens only and append inserted values verbatim. + +### 2. Consolidate checkpoint and grounded close + +Files: `close.rs`, `close_tests.rs`; public paths remain `tinyagents_harness::close`. + +1. Inventory existing `close.rs` entry points before editing: `CloseToolResult`, `results_from_tool_outcomes`, truncation/rendering, deterministic builders, `CloseViolation`, repair config/options and usage callbacks, required-output repair, grounded close and fallback. Extend those APIs rather than creating parallel algorithms. +2. Add focused tests for both checkpoint and final close with custom copy and byte equality against OpenHuman goldens. Preserve the cases in host `grounded_close_tests.rs`: classified halt with no provider usage; selected model and verifier excluded from final-context gauge; failed close request counted without usage; correlated native start/completion events; EOF and timeout traces without timeout usage; all required-output placement variants; truncated and capped turns; timeout repair retains completed-call usage; cache event clears provider charge; free-route catalog rates are zero. Assert output bytes, request sequence, model/temperature/output cap, tool-less calls, callback kind/response presence, token gauges, and event correlation. +3. Move `SummarizerSourceModel` to harness as `CloseSourceModel<'a> { source: &'a dyn CloseModelSource, model: &'a str, temperature: f64, thread_id: Option<&'a str>, trace: Option<&'a Arc<SupplementalCallTrace>> }`. Define `pub trait CloseModelSource: Send + Sync { fn build_summarizer(&self, model: &str, temperature: f64, thread_id: Option<&str>) -> Result<Arc<dyn ChatModel<()>>, String>; }`. Its `ChatModel<()>` implementation obtains the per-call model and forwards request, invoke, and stream. In `invoke`, call `SupplementalCallGuard::start(trace, model, &mut request)` before provider invocation; on success call `complete(response)`, while error propagation drops the guard and emits one failure. In `stream`, start before opening the provider stream, then retain the guard in the mapped stream: terminal `Completed` calls `complete` once, `Failed`/`ProviderFailed` drops it once, and EOF or consumer drop before a terminal item drops it as an incomplete failure. This preserves the host's current `Arc` ownership so the guard keeps the trace alive if the stream outlives the wrapper. Map the source's `String` build error to `tinyinference_llm::Error::Model`; OpenHuman implements the source port by delegating to `TurnModelSource`. +4. Move `RepairUsage`'s generic fields and accumulation into harness `CloseUsage`: `model_calls`, `input_tokens`, `output_tokens`, `cached_input_tokens`, `last_call_input_tokens`, `last_call_output_tokens`, plus `cost: CostTally` and `pricing_model: String` (both types already belong to harness/inference contracts). Implement `for_model(model)`, `record(kind, Option<&Usage>, measures_context)`, and `record_close_call(kind, Option<&ModelResponse>)`. Preserve exact rules from source lines 45-98: count each attempted request; add usage totals only when present; record final context tokens only for non-verification requests; a successful close/repair with context tokens supersedes the loop's final call. Keep the callback seam so OpenHuman's `BilledUsage` conversion and catalog pricing call `CostTally::add(CallCost)` without importing OpenHuman types upstream. Existing `grounded_close_tests.rs` covers callback complete/fail/timeout paths. +5. Move `classified_halt_close` into harness as a generic selector with exact contract `pub fn classified_halt_close<T>(stop_reason: Option<&str>, is_classified_halt: impl FnOnce(&str) -> bool, build_close: impl FnOnce() -> T) -> Option<T>`. Its complete implementation is: `let reason = stop_reason?; is_classified_halt(reason).then(build_close)`. Host predicate is `|reason| reason.starts_with("Stopping after ")`; its fallback closure renders the host `StopCopy` summary from tool outcomes. This keeps the product marker and classification host-side while moving the generic short-circuit. Port `classified_halt_returns_partial_work_without_provider_usage`; assert no model request/usage, partial successes listed, success payload omitted, and failed line retained. +6. `repair_required_output` already exists in harness at `close.rs:559-698`; consolidate the host call to it rather than porting a duplicate. Preserve host construction of `RequiredOutput`, selected repair instruction, `Prepend` placement and post-repair progress delta. Port its exact behavior assertions from `required_output_repair_preserves_all_host_output_placements` and retain model-repaired vs deterministic fallback results. + +### 3. Adopt generic turn-stop output and typed terminal outcome + +Files: `turn_stop.rs`, `turn_stop_tests.rs`, `turn_outcome.rs`, `turn_outcome_tests.rs`, `lib.rs` exports as needed. + +1. Add tests for custom stop copy using exact host golden output and retain every existing `turn_stop_tests.rs` assertion: typed content-free causes, legacy fallback, serde sanitization, stop precedence, incomplete-reason mapping, and status text. +2. Move `DeferredTurnTerminal`'s event-owned fields (`RunId`, `EventSink`) and methods to `turn_outcome.rs`; remove its host `TurnJournal` field. Preserve event order and exact error/outcome literals: success emits optional `TurnStopped` then `RunFinalized`; cancellation/failure emits `RunFailed` then `RunFinalized`. Host session sidecar owns both the harness terminal token and journal; it emits terminal events then calls the appropriate journal status method. Add tests asserting exact event variants/order, not just call counts. Keep `TinyagentsTurnOutcome` and unobserved usage tests on the host. +3. Host-specific `StopReasonKind::from_stop_note` / stable product keys remain in OpenHuman. `TurnStop` construction can call the existing generic `classify`, `classify_with_cause`, and `from_incomplete_reason`; product fallback closure stays supplied by OpenHuman. Remove only forwarding wrappers once all uses call canonical API. + +### 4. Move observability bridge behavior without moving host persistence + +Files: `observability/mod.rs` and its existing `mod_journal_tests.rs` only for changes required by supplemental tracing. No new journal facade/test is planned. Host files are evidence inputs; #64 does not edit OpenHuman. + +1. Do not add/move `TurnJournal` into harness. The inspected type at host lines 123-225 includes `DriverStatusStore`, `HarnessRunStatus`, sealing, and host logging. The generic event path already uses `StoreEventJournal`, `JournalSink`, and `HarnessEventJournal`. +2. Keep attachment, status lifecycle and `merge_run_observations` on the host. `journal_tests.rs` only tests that the request/run map stays scoped to the registering `DispatchCtx`; keep it there. Existing harness observability tests cover generic append/read/offset/lineage behavior; add no duplicate tests. +3. Move `SupplementalTrace` call lifecycle into `harness::observability::SupplementalCallTrace`, with concrete request correlation and RAII guard API. Define `pub trait SupplementalEnricher: Send + Sync { fn enrich(&self, model: &str, response: &ModelResponse) -> (Option<String>, GenerationEnrichment); }`; the trace owns `EventSink`, `RunId`, `AtomicUsize`, `Mutex<Option<String>>` actual model, `AtomicBool` cache flag, and `Arc<dyn SupplementalEnricher>`. `SupplementalCallTrace::new(events: EventSink, run_id: RunId, enricher: Arc<dyn SupplementalEnricher>) -> Arc<Self>` initializes the next-call counter to 1 and cache/model state empty. Its accessors are `run_id(&self) -> &RunId`, `actual_model(&self) -> Option<String>`, and `is_cache_hit(&self) -> bool`. `SupplementalCallGuard::start(trace: &Arc<SupplementalCallTrace>, model: &str, request: &mut ModelRequest) -> Self` allocates `CallId::new(format!("{}-host-supplemental-{n}", trace.run_id().as_str()))`, captures `now_ms()`, sets `request.correlation = Some(ModelCallCorrelation::new(trace.run_id().as_str(), call_id.as_str()))`, serializes request messages with `serde_json::to_value(&request.messages).unwrap_or(Value::Null)`, emits `ModelStarted`, and stores an `Arc::clone(trace)` plus call id, model, start time, input and `completed: false` in the guard. `complete(&mut self, response: &ModelResponse)` builds mutable local `native_usage = response.usage.or(response.message.usage)`; on cache hit emit `ModelResponseCached` and clear `native_usage.charged_amount`; then update cache state, emit `UsageRecorded` if present, emit `ModelCompleted` with captured input and serialized message output, ask the enricher for `(actual_model, enrichment)`, store actual model, emit `GenerationEnriched` only when nonempty, and set `completed = true`. `Drop` emits exactly one generic `ModelFailed` if incomplete, including invoke error, stream error, EOF, or consumer drop. `CloseSourceModel` owns `Option<&Arc<SupplementalCallTrace>>` and uses this guard for invoke and stream as specified in Step 2.3, preserving correlation and lifetime through the whole call. The host enricher reproduces `supplemental_trace.rs:101-148`: actual-route model fallback, provider id validation, provider pricing, `usage_info_from_response`, and estimate suppression for cached/reported charges. Move `SupplementalCall` RAII lifecycle from `grounded_close.rs:124-225`. Test success correlation and event order, invoke failure, stream error, EOF, and consumer drop; each failure path must emit one `ModelFailed`, and a completed stream must emit one completion only. Retain cached-response charge clearing and zero-price route tests; no external calls. +4. Keep all persistence through `AppendStore` and the existing journal ports. No database/storage dependencies may be added to `tinyagents-harness`. + +### 5. Move capture middleware and repeated-failure ladder core + +Files: `turn_outcome.rs`, `turn_outcome_tests.rs`, `close.rs`, `close_tests.rs`, and `observability/mod.rs` plus its existing sibling tests. The existing `classified_ladder.rs` is already the generic owner and does not need a duplicate implementation. Host files are evidence inputs; #64 does not edit OpenHuman. + +1. The host capture wrapper delegates to `tinyagents_harness::middleware::library::ToolOutcomeCaptureMiddleware` already. In #212 delete only this adapter after direct construction is wired. Keep host policy enrichment, TinyTools failure classification, `mcp_result` metadata filtering and `ToolFailureMap`. Preserve the host tests for same-name calls by call id, forwarding structured MCP presentation, and dropping undecorated MCP envelopes. Harness `turn_outcome_tests.rs` already tests poison-safe drain and one-time take. +2. The inspected 541-line `RepeatedToolFailureMiddleware` already delegates to harness `FailureScopeMiddleware`, `FailurePolicy`, `ClassifiedFailureTracker`, and `classified_step`; its remaining `HostFailurePolicy` owns host `RunContext`, call-resource scope, sanitization, flow-body failure detection, TinyTools classifier, recoverable headroom, terminal-inference fast halt, StopLatch and OpenHuman copy. Make no source move for this file in #64; preserving host tests is the acceptance evidence. +3. Existing harness `classified_ladder_tests.rs`, `recoverable_tests.rs`, and failure-scope tests own generic ladder behavior. Do not mirror them in OpenHuman or duplicate them in harness. +4. Do not relocate `HostFailurePolicy`, inference `failure_copy`, `OpenHumanRunContext`, fetched-site metadata, or `StopLatch` into TinyAgents. + +### 6. Replace the host model-budget wrapper + +Files: existing `crates/tinyagents-harness/src/cost/model_budget.rs` and `model_budget_tests.rs`; host files are evidence inputs only. + +1. Keep the existing `wrapped_model_settles_the_reported_charge_once_into_shared_ledger` test in `model_budget_tests.rs`; its assertion requires normalized charge `7` micros and ledger spend `{ tokens: 5, cost_micros: 7 }`. Add a behavior test only if changing upstream wrapper behavior; issue #64 requires no accounting algorithm rewrite. +2. Make the upstream wrapper public from the harness facade if it is not already exported, then replace the host `ModelBudget` struct with a type alias to `tinyagents_harness::cost::ModelBudget` if host call-site visibility allows it. Keep `TurnModels::with_budget` route wrapping and depth policy in OpenHuman. +3. Ensure issue #60's runtime/job changes do not create a harness dependency on orchestration. No speculative session/runtime API is part of this step. + +### 7. Future host adoption in OpenHuman #212 (not a #64 implementation task) + +After the TinyAgents #64 PR merges into the fork epic branch, OpenHuman #212 consumes that published fork API and moves the gitlink. This section gives #212 acceptance constraints only; #64 changes no OpenHuman source. + +1. Implement `CloseModelSource` for `TurnModelSource`; inject OpenHuman `CloseCopy` templates and the host result-render options. Keep model selection, output caps, billing/pricing callback, progress deltas, required-output contract construction, and fallback placement in host adapters. +2. Implement the supplemental enricher with OpenHuman provider/catalog lookups. Keep `StopReasonKind`, failure-note parsing, identifier/secret scrub, failure classification, `HostFailurePolicy`, `StopLatch`, recoverable headroom, terminal-inference fast halt, flow-specific body failures, request correlation, status storage and replay in OpenHuman. +3. Replace the host budget wrapper with the upstream `ModelBudget` at `TurnModels::with_budget`; retain depth and policy setup. +4. Remove only generic duplicates actually consumed from TinyAgents; retain host policy/adapters. Ported goldens remain in the owning tests. Keep `journal_tests.rs`, UI metadata tests, policy and budget tests because their behavior is still host-owned. + +## Deferred dependencies and integration order + +- TinyAgents #60 is the delivery/full-source change; it has no phase pin. Integrate its actual public source where needed and review/merge against the latest fork epic tip. Do not assume a phase-only API from #60. +- TinyAgents #63's session-only approved phase is `120410141b9e59896072a657985073b11cb53768`. Its later source work waits on #60; no API beyond that approved phase is assumed. +- TinyAgents #65 plan-v3 is `bc4910f49b98839d16e9dd91f192f02a589e3a8b`; it is not #63 and this work does not depend on its proposed APIs. +- OpenHuman #212 is the future host-adoption leaf. This #64 plan covers TinyAgents only; #212 adapts and deletes the host rows after #64 lands on the senamakel epic branch and the fork gitlink is advanced there. +- Preserve all unrelated #8088 tracker requirements. The issue #64 scope is the rows and behaviors in its body; it does not narrow the epic's other leaves. + +## Acceptance and verification + +Issue #64 acceptance (TinyAgents repository): + +1. Ported host golden outputs are byte-identical when host `CloseCopy`/`StopCopy` is injected: deterministic checkpoint, final summary, breaker halt, tool/no-tool output, no-reply/quoted/unverified close repairs, output-budget exhaustion and iteration cap. +2. Preserve all generic close state-machine assertions in `close_tests.rs` (tool-less requests, verifier behavior, one repair, deterministic fallback, timeout/cancellation, usage callbacks, request properties, required-output placement) and stop serialization/classification assertions in `turn_stop_tests.rs`. +3. Preserve supplemental trace event order, cache charge clearing, usage/enrichment, run correlation and cancellation/failure terminal event behavior. Retain the existing cost budget test unchanged. No external service calls in tests. +4. Source audit confirms TinyAgents contains no OpenHuman type/dependency, new ports use only harness and inference contracts, no duplicated existing journal/close/cost algorithm was added, and all requested generic source rows have an upstream implementation or documented existing implementation. + +OpenHuman row removal, adapter compile, storage integration and all retained host policy tests are acceptance for future OpenHuman #212, not this PR. + +Run required repository gates from the super-worktree root through its long-command cancellation wrapper, pointing Cargo at the TinyAgents manifest (no temp build directory and no `CARGO_TARGET_DIR`): + +```sh +scripts/ci-cancel-aware.sh cargo fmt --manifest-path vendor/tinyagents/Cargo.toml --all -- --check +scripts/ci-cancel-aware.sh cargo check --manifest-path vendor/tinyagents/Cargo.toml --workspace --all-targets +scripts/ci-cancel-aware.sh cargo clippy --manifest-path vendor/tinyagents/Cargo.toml --workspace --all-targets -- -D warnings +scripts/ci-cancel-aware.sh cargo test --manifest-path vendor/tinyagents/Cargo.toml --workspace +``` + +Also run focused tests for changed harness modules during implementation. The OpenHuman adapter phase separately runs its required focused tests and layout checks. This plan phase runs no builds or tests and makes no gate claim. + +## Resolved source audits + +1. CloseCopy inventory is complete: prompt/copy literals and render rules are in `turn_checkpoint.rs:49-126,152-169,198-258,260-372,388-514`. Existing generic algorithms are in harness `close.rs`; move all user/model-facing literals into `CloseCopy`, with default values copied verbatim from those exact source ranges, and reuse existing result-render options. Keep the OpenHuman golden test strings unchanged. +2. Do not add `TurnJournal` to harness. `journal.rs:123-225` couples terminal status writes to host `DriverStatusStore`; the generic event side already uses `StoreEventJournal`, `JournalSink`, and `HarnessEventJournal`. The listed generic row is already represented upstream. Keep attachment, status lifecycle, correlation, and replay host-side. Move only event terminal emission from `DeferredTurnTerminal` and leave durable status finalization with its host owner. +3. `SupplementalTrace` (13-169) uses harness `EventSink`, `AgentEvent`, ids, `ModelResponse`, `GenerationPricing`, and `GenerationEnrichment` for call lifecycle. It uses host `agent::cost` and provider conversion for catalog pricing, charge estimates, and `RepairUsage` fields. Move event lifecycle to harness behind the typed host enrichment callback; retain catalog/billing and final cost projection on the host callback. +4. Upstream `ModelBudget::new`, `wait_for_capacity`, and `wrap_model` are verified in `tinyagents-harness/src/cost/model_budget.rs:22-42`; retain host depth/policy setup and replace only the wrapper type. diff --git a/docs/plans/2026-10-19-harness-tools-exec-gates.md b/docs/plans/2026-10-19-harness-tools-exec-gates.md new file mode 100644 index 000000000..22499bae1 --- /dev/null +++ b/docs/plans/2026-10-19-harness-tools-exec-gates.md @@ -0,0 +1,848 @@ +# Plan: move standard tools into TinyAgents and add gated `exec` + +This plan implements `senamakel/tinyagents#65`, which is part of +`tinyhumansai/openhuman#8088`. It covers the whole leaf: moving the +`tinytools-std` implementation families and their tests into +`tinyagents-harness-tools`, adding shell and install tools behind generic host +ports, moving the plan-review and media-artifact tool wrappers found by the +OpenHuman core-shrink audit, and removing the harness's last dependency on +`tinytools-std`. + +The implementation belongs in TinyAgents. `tinytools` remains the protocol, +types, and policy vocabulary owner; `tinyagents-harness` owns the agent loop, +portable runtime utilities, and multimodal resolution; `tinyagents-harness-tools` +owns reusable tool implementations; OpenHuman keeps the concrete policy, +approval, sandbox, audit, plan-review, and artifact-storage adapters for its +later #213 integration. No OpenHuman type or dependency is introduced here. + +## Starting state and source audit + +The plan was prepared against TinyAgents `aca9fd54666ce8fdcc02ae4d6499baa2a3354cd5` +on branch `epic-8088-65-harness-tools`. Issue #65 is open and has no comments. +OpenHuman audit comment 6254297338 adds three rows to the issue's original +scope: `sandbox/shell_progress.rs`, `plan_review/tool.rs`, and +`media/generation/artifact_tool.rs`. The OpenHuman task #213 will consume these +new generic APIs later; it does not block this crate-local change. + +The current `tinytools-std` public modules are `detect_tools`, `file_state`, +`filesystem`, `network`, `sanitize`, `timeout`, and `url_guard`. The prior +`tinytools#4` change already moved `sanitize` into `tinytools`; do not copy it +back. The harness currently references `tinytools-std` only from +`crates/tinyagents-harness/src/multimodal/resolve.rs` and its tests, for +`url_guard`. That is the only dependency edge that would otherwise prevent +removing `tinytools-std` from the harness manifest. + +The important dependency shape is: + +```text +tinyagents-harness-tools ──depends on──> tinyagents-harness ──does not depend on──> harness-tools +``` + +Do not reverse this edge. Move the shared URL-validation implementation into +an unconditional `tinyagents-harness::url_guard` module. It uses only `std`, +`anyhow`, Tokio's already-enabled `rt` support, and `tracing`; convert the +source `log::{debug,warn}` calls to equivalent `tracing` macros. It must not +add reqwest, flate2, tar, zip, a resolver, or any multimodal feature as a +dependency. Multimodal resolution uses it directly and `harness-tools/net` +consumes it through its existing dependency on `tinyagents-harness`. This removes the optional +`tinytools-std` dependency without creating a crate cycle or adding a shim +crate. The module must retain the current `ValidatedUrl`, validation, +DNS-resolution, public-address checks, and exact address-pinning contract; +preserve all URL-guard tests with the move. + +Before coding, re-read the current source and tests from the exact parent +checkout that the OpenHuman #213 adapter will use. In particular inspect: + +| Current OpenHuman source | Behavior to preserve / destination | +| --- | --- | +| `crates/openhuman-core/src/tools/impl/system/shell.rs` and sibling `shell_tests*.rs` | `shell` tool metadata, schema, output limits, command analysis and policy call order, sandbox/backend selection, audit, progress, cancellation, and errors; rename the generic tool to `exec` | +| `tools/impl/system/shell_support.rs`, `shell_description.rs` | Shell descriptions, argv/environment helpers, timeout text, exit-124 attribution, status mapping, and output formatting | +| `tools/impl/system/install_tool.rs` and `install_tool_tests.rs` | Package validation, allowed managers and exact argv, interactive/feature/autonomy/rate checks, timeout, output cap, and error messages | +| `tools/schema_cache.rs` | `static_schema!` and parse-once embedded JSON behavior | +| `sandbox/shell_progress.rs` and `shell_progress_tests.rs` | Byte counting, bounded last-line tracking, heartbeat interval/text, and context progress events | +| `plan_review/tool.rs` and tests | Tool contract and current behavior; the host's gate/event state machine is behind `PlanReviewPort` | +| `media/generation/artifact_tool.rs` and tests | Inner-tool metadata/results, multiple files, blob/path handling, artifact errors, markdown replacement, and result preservation; host filing/storage is behind `ArtifactFiling` | + +The tests described in issue #65 are the source of truth for behavior. Copy +them with their assertions and fixtures, then adapt only test imports and +host-owned fakes to the new ports. Do not copy OpenHuman security policy, +`SecurityPolicy`, `AuditLogger`, TinyBox/TinyBus contracts, artifact stores, +`Production`, SaaS switches, Docker/TinyBox selection, or OpenHuman context +types into TinyAgents. The generic tool asks its injected ports; OpenHuman's later adapter +implements those ports. Keep the existing host's approval gate and turn +routing in the host. The tool runs `ExecGate::authorize` inside `execute` +immediately before invoking the selected backend. A gate denial or selected +backend preflight refusal never attempts process creation, even if the outer +approval gate let the call reach the tool. + +## Contract to establish before porting + +Put public seams in `crates/tinyagents-harness-tools/src/exec/ports.rs`, +`plan_review/port.rs`, and `media/artifact_filing.rs`. Keep the port types free +of host-specific dependencies. These definitions are complete contracts for +the implementation; adapt names only if the inspected current +`ToolRunContext` or media-output types require a mechanical signature change. + +### `exec` host boundary + +`ExecGate` decides whether an operation may proceed and owns host audit +recording. Its authorization call is awaited before backend execution. It +returns a permit only after host policy approval; a denial is returned to the +model as a stable `ToolResult::error` and does not call `ExecBackend`. The +selected backend may still refuse before launch after configuration/preflight. +The host can use `ToolRunContext::host_extension()` to read its own turn state +without TinyAgents importing or downcasting to a host type. + +```rust +use std::{any::Any, collections::BTreeMap, ffi::OsString, path::PathBuf, sync::Arc, time::Duration}; +use async_trait::async_trait; +use serde_json::Value; +use tinytools::ToolRunContext; +use crate::exec::ExecProgress; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum ExecOperation { + Shell { command: String, category: Option<String> }, + Install { package: String, manager: String }, +} + +#[derive(Debug, Clone)] +pub struct ExecAuthorizationRequest { + pub operation: ExecOperation, + pub timeout: Option<Duration>, +} + +/// Opaque host-owned data selected during authorization. TinyAgents can carry +/// it unchanged but cannot interpret or construct a host policy type. +#[derive(Clone)] +pub struct ExecBackendCapability(Arc<dyn Any + Send + Sync>); +impl ExecBackendCapability { + pub fn new<T: Any + Send + Sync>(value: T) -> Self { Self(Arc::new(value)) } + pub fn downcast_ref<T: Any>(&self) -> Option<&T> { self.0.as_ref().downcast_ref() } +} + +#[derive(Clone)] +pub struct ExecPermit { + pub action_dir: PathBuf, + pub environment: BTreeMap<OsString, OsString>, + pub backend_selection: String, + pub backend_capability: ExecBackendCapability, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ExecDenied(pub String); + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum ExecRefusalKind { BackendUnavailable, Configuration, Preflight } + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum ExecAuditOutcome { + GateDenied, + BackendRefused { kind: ExecRefusalKind }, + LaunchFailed, + Completed { success: bool, timed_out: bool }, +} + +#[derive(Debug, Clone)] +pub struct ExecAuditEvent { + pub operation: ExecOperation, + /// The source audit's approval-required proxy from the exact tool args. + pub approval_required: bool, + /// True only after host authorization produced a permit. + pub authorized: bool, + /// Mirrors the source `(allowed, ToolResult)` gate result: false for a + /// policy/rate denial or a backend refusal before launch. + pub allowed: bool, + pub backend_selection: Option<String>, + /// True only when process creation was attempted. + pub process_attempted: bool, + pub outcome: ExecAuditOutcome, + pub duration: Duration, +} + +#[derive(Clone)] +pub struct ExecRequest { + pub program: OsString, + pub args: Vec<OsString>, + pub action_dir: PathBuf, + pub environment: BTreeMap<OsString, OsString>, + pub timeout: Option<Duration>, + pub shell_command: Option<String>, + pub backend_selection: String, + pub backend_capability: ExecBackendCapability, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ExecOutput { + pub stdout: Vec<u8>, pub stderr: Vec<u8>, + pub exit_code: Option<i32>, pub timed_out: bool, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum ExecBackendOutcome { + /// A selected backend declined before process creation (for example a + /// SaaS backend is disabled, or Docker configuration/preflight failed). + Refused { kind: ExecRefusalKind, reason: String }, + /// Process creation was attempted and failed. + LaunchFailed { error: String }, + /// A process ran to completion or timed out and was cleaned up. + Completed(ExecOutput), +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ShellFlavor { Posix, Cmd } + +#[async_trait] +pub trait ExecGate: Send + Sync { + /// Exact TinyTools `Tool` authorization-preparation hook. The host may + /// return serialized command analysis for its outer approval flow. + async fn prepare_authorization(&self, args: &Value) -> anyhow::Result<Option<Value>>; + /// Exact argument-aware external-effect hook; retain the source shell's + /// conservative classification for every nonempty command. + fn external_effect_with_args(&self, args: &Value) -> bool; + /// Revalidate and authorize immediately before backend dispatch. + async fn authorize(&self, request: &ExecAuthorizationRequest, + context: Option<&dyn ToolRunContext>) -> Result<ExecPermit, ExecDenied>; + fn record(&self, event: ExecAuditEvent); +} + +#[async_trait] +pub trait ExecBackend: Send + Sync { + fn shell_flavor(&self) -> ShellFlavor; + async fn execute(&self, request: ExecRequest, + context: Option<&dyn ToolRunContext>, progress: Option<ExecProgress>) -> ExecBackendOutcome; +} +``` + +These signatures are the actual `tinytools::Tool` hooks to implement and +forward: `prepare_authorization(&Value) -> anyhow::Result<Option<Value>>` and +`external_effect_with_args(&Value) -> bool`. Do not invent a second authorization +hook or narrow classification to a command prefix. `ExecTool` and +`InstallTool` delegate those hooks to the injected gate; the host adapter +reuses TinyBox analysis for `prepare_authorization` and returns its JSON +analysis exactly as the source shell does. `authorize` is a separate final +revalidation and must occur immediately before dispatch. + +The host gate returns a permit only after its approval, autonomy, rate, +workspace, and other policy checks succeed. The permit carries the authorized +backend selection and opaque capability through `ExecRequest`; the backend +must use that selection and capability rather than independently choosing a +more permissive route. OpenHuman retains authority over native/sandbox/Docker/ +TinyBox/SaaS selection. A SaaS-unavailable refusal and Docker configuration or +preflight refusal are typed `Refused` outcomes before process creation; they +are not policy denials and are not launch failures. `LaunchFailed` means a +process-creation attempt failed. Audit exactly once for each well-formed +operation that reaches the gate, recording the source approval-required proxy +(`external_effect_with_args(args)`) separately from authorization and allowed +status. Gate denial is `authorized=false, allowed=false, +process_attempted=false`; backend refusal (including SaaS refusal or Docker +configuration/preflight failure) is `authorized=true, allowed=false, +process_attempted=false`; launch failure is `authorized=true, allowed=true, +process_attempted=true`; completed execution is allowed and attempted, with +success derived from the source result mapping. Include the selected backend +when a permit exists. Never represent a backend refusal as an allowed successful +execution. + +The shell result conversion preserves current behavior: decode output lossily +as UTF-8; cap each stream at 1 MiB at a valid UTF-8 boundary with the existing +truncation marker; retain `max_result_size_chars = 30_000`; append successful +stderr as `[stderr]`; use the current `tinytools::command_failure` mapping for +nonzero/signal exits; append the source “command's own timeout” note for exit +124; and retain explicit timeout text including captured stdout/stderr. +`ExecBackendOutcome::Refused` maps to the matching existing source refusal +text; launch failure maps to the stable sandbox-execution failure text. The +Python UTF-8 defaults and caller overrides remain unchanged. + +The fake gate captures audit events. Test each branch and assert exactly one +event per well-formed call: gate denial records `GateDenied` with +`authorized=false, allowed=false, process_attempted=false` and zero backend +calls; SaaS refusal and Docker config/preflight refusal each record +`BackendRefused` with `authorized=true, allowed=false, +process_attempted=false`; launch failure records `LaunchFailed` with +`authorized=true, allowed=true, process_attempted=true, success=false`; and +completed success records `Completed` with the selected backend, +`authorized=true, allowed=true, process_attempted=true, success=true`. +Completed command failure/timeout also records exactly one `Completed` event +with success matching the final `ToolResult`. Preserve shell-analysis +preparation and `external_effect_with_args` tests using source arguments, +including empty and nonempty commands. + +### Plan-review port + +`RequestPlanReviewTool` owns only parsing, public tool metadata, and mapping the +resolution to the existing stable result text. The port owns the host review +surface, user interaction, event publication, expiry, plan-mode check, and +non-interactive policy. A caller without an installed port fails closed with +an actionable error; it must not silently auto-approve. The host adapter can +preserve OpenHuman's non-interactive auto-approve rule explicitly. + +```rust +use async_trait::async_trait; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct PlanReviewRequest { + pub summary: String, + pub steps: Vec<String>, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum PlanReviewResolution { + /// Interactive review surface returned approval. + Approve, + /// Host policy approved a non-interactive turn without a review surface. + NonInteractiveApprove, + /// The host established that this thread is outside plan mode. + NotApplicable, + Reject, + Revise { feedback: String }, +} + +#[async_trait] +pub trait PlanReviewPort: Send + Sync { + async fn request_review( + &self, + request: PlanReviewRequest, + context: Option<&dyn tinytools::ToolRunContext>, + ) -> Result<PlanReviewResolution, String>; +} + +pub struct RequestPlanReviewTool { + port: Option<std::sync::Arc<dyn PlanReviewPort>>, +} + +impl RequestPlanReviewTool { + pub fn new(port: Option<std::sync::Arc<dyn PlanReviewPort>>) -> Self { + Self { port } + } +} +``` + +`RequestPlanReviewTool` stores `Option<Arc<dyn PlanReviewPort>>`; its +constructor accepts that option. The adapter returns `NonInteractiveApprove` +only for the existing non-interactive policy, and `NotApplicable` only for the +existing outside-plan-mode success path. Interactive approval, rejection, and +revision map to separate variants. A missing port returns a stable error +before any resolution, making the fail-closed test meaningful. +The tool keeps `request_plan_review`, `PermissionLevel::None`, no external +effect, and an unbounded timeout because the port's review deadline is +authoritative. It trims summary and nonempty string steps exactly as the source +does. The port reads the provider-assigned call id from the generic +`ToolExecutionContext` exposed through `ToolRunContext::host_extension()`; +the request itself contains no OpenHuman identifier type. Preserve exact source result strings as constants/tests: + +```text +interactive approve: approved: the user approved the plan — proceed and execute it now. +non-interactive approve: approved: non-interactive turn (no review surface) — proceed with the plan. +not applicable: not applicable: this thread is not in plan mode, so there is no plan to review. Do not call `request_plan_review` again this turn. +reject: rejected: the user rejected the plan — do NOT execute it. Briefly ask what they would like to do instead. +revise: revise: the user requested changes before executing. Their feedback: +{feedback} +Revise the plan accordingly, then call `request_plan_review` again before executing. +``` +Tests assert all five outcomes and exact strings. A timeout/dropped review +continues to resolve as the host's fail-closed `Reject` (with its expiry +metadata); a missing port or port error remains an error and can never map to +an approval variant. + +### Media artifact filing port + +The generic wrapper is feature-gated as `media` and introduces no +`tinyinference-video`, `tinyinference-image`, multimodal, or provider transport +dependency. The media progress callback is a transport-free generic seam: + +```rust +#[derive(Debug, Clone, PartialEq)] +pub struct MediaJobProgress { + pub id: String, + /// Exact provider-neutral display state (`JobState::as_str()` in the + /// OpenHuman adapter, including provider `Other` values). + pub state: String, + pub outputs: usize, + pub output_indices: Vec<usize>, + pub cost_usd: Option<f64>, + pub error: Option<String>, +} +pub type MediaProgressFn = std::sync::Arc<dyn Fn(&MediaJobProgress) + Send + Sync>; +``` + +Keep `MediaArtifactTool<T>` generic over `T: tinytools::Tool` and pass the +`ToolRunContext` to a per-call inner-tool factory so host code can read its +own approval/chat extension. The factory accepts +`(MediaProgressFn, Option<&dyn ToolRunContext>)`; it does not mention +`ApprovalChatContext` or provider types. The OpenHuman adapter converts the +provider callback/status into this value (`state.to_string()` preserves +`JobState::as_str()` including arbitrary provider `Other` values), preserving +all fields and progress heartbeat behavior. Compare it against +`VideoJobStatus` and the old wrapper's progress tests; do not discard an +observable status. The wrapper keeps the inner tool's name, description, +schema, policy, permissions, scope, category, result content, and progress +behavior. It parses +each JSON result's `artifacts` array and invokes the injected `ArtifactFiling` +port once per recognizable output. It does not store files, create product +artifact records, or depend on TinyArtifacts/OpenHuman. + +```rust +use async_trait::async_trait; +use std::path::PathBuf; +use std::sync::Arc; +use tinytools::{Tool, ToolRunContext}; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum MediaArtifactKind { Image, Video } + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct GeneratedArtifact { + pub path: PathBuf, + /// Wrapper configuration retained from the source host adapter and used + /// by its filing port to select the image/video artifact category. + pub artifact_kind: MediaArtifactKind, + /// The generated output's own reported kind (for example image/video). + pub kind: String, + pub media_type: Option<String>, + pub bytes: Option<u64>, + pub title: String, + pub tool_call_id: Option<String>, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct FiledArtifact { + pub id: String, + pub path: String, +} + +#[async_trait] +pub trait ArtifactFiling: Send + Sync { + async fn file(&self, generated: GeneratedArtifact, + context: Option<&dyn ToolRunContext>) -> Result<FiledArtifact, String>; +} + +pub struct MediaArtifactTool<T: Tool> { + /// Retained instance is the stable metadata source (`name`, `description`, + /// schema and other Tool declarations) and fallback executor. + inner: T, + kind: MediaArtifactKind, + /// Per-invocation tool, built only for execution/progress; its borrowed + /// name/description are never returned from this temporary value. + per_call: Option<Arc<dyn for<'a> Fn(MediaProgressFn, Option<&'a dyn ToolRunContext>) -> T + Send + Sync>>, + filing: Arc<dyn ArtifactFiling>, +} +impl<T: Tool> MediaArtifactTool<T> { + pub fn new(inner: T, kind: MediaArtifactKind, filing: Arc<dyn ArtifactFiling>) -> Self { + Self { inner, kind, per_call: None, filing } + } + pub fn with_per_call_inner( + mut self, + factory: Arc<dyn for<'a> Fn(MediaProgressFn, Option<&'a dyn ToolRunContext>) -> T + Send + Sync>, + ) -> Self { + self.per_call = Some(factory); + self + } +} +``` + +The wrapper passes each invocation's original `ToolRunContext` through every +`ArtifactFiling::file` call. This per-call argument is required because the +host adapter derives approval/chat routing from that call's extension; do not +capture one ambient chat context in the shared wrapper or factory. Retain the +stable `inner: T` instance alongside the per-call factory: delegate `Tool::name` +and `Tool::description` (and the other static Tool metadata) to it. Do not +return references borrowed from the temporary tool built by the factory. Store +`MediaArtifactKind::{Image, Video}` on the wrapper and pass it with every +`GeneratedArtifact`; use it to preserve the existing heartbeat labels +`Generating image` / `Generating video`. The wrapper +replaces only the entry's `path`, adds `artifact_id` on success, +or adds `artifact_error` on a per-entry failure; it keeps processing later +entries and leaves unrecognized shapes untouched. It rewrites matching source +paths in `markdown_formatted`. Keep the `n > 1` title suffix and 60-character +prompt title cap. OpenHuman's later adapter implements `ArtifactFiling` and +retains the existing versioning, trusted-file-root checks, blob cleanup, and +reference allowlist policy. Metadata tests use a fake inner tool with distinct +name, description, schema, category, policy, permission, scope, timeout, +external-effect behavior, and result-size limit; assert every static/argument- +based `Tool` declaration delegates to the retained instance with and without a +per-call factory. In particular, assert `name()` and `description()` continue +to return the retained instance's references after invoking a factory that +returns a different temporary tool. Image/video configuration tests assert the exact configured +kind passed to each filer and exact heartbeat label (`Generating image` / +`Generating video`), including when the factory returns a new temporary inner +tool. This catches both metadata borrowing from a temporary and accidental +use of the fallback image kind for video calls. + +## Bite-sized implementation tasks + +Complete these in order. Each task can be a separate commit; do not combine or +squash checkpoints. Every changed Rust test stays in the repository's sibling +`*_tests.rs` convention and begins with `use super::*;`. No inline test modules, +`tests.rs`, `test.rs`, or `<module>_test.rs`. Ported tests retain the original +assertions. Tests must use fakes or loopback listeners and never call a real +third-party service. + +### 1. Move the shared URL guard and detach the harness from `tinytools-std` + +Files: `crates/tinyagents-harness/src/url_guard.rs` (new), +`crates/tinyagents-harness/src/lib.rs`, +`crates/tinyagents-harness/src/multimodal/resolve.rs`, +`crates/tinyagents-harness/src/multimodal/resolve_tests.rs`, +`crates/tinyagents-harness/Cargo.toml`, +`crates/tinyagents-harness-tools/Cargo.toml`. + +Move `vendor/tinytools/crates/tinytools-std/src/url_guard/mod.rs` to +`crates/tinyagents-harness/src/url_guard.rs` and +`url_guard/mod_tests.rs` to `crates/tinyagents-harness/src/url_guard_tests.rs`. +Preserve its full implementation and tests, changing only the logging macros +and crate paths. Export `pub mod url_guard;` from the harness root. Change +multimodal resolution and tests to use +`crate::url_guard`. Update the public docs so the SSRF rules and DNS pinning +contract name the new owner. Remove the optional `tinytools-std` dependency and +the `dep:tinytools-std` edge from the harness `multimodal` feature. Do not add +the URL module to `harness-tools`; its `net` module imports the utility from +its existing harness dependency. + +Red/green: run the focused `tinyagents-harness` URL guard and multimodal +resolve tests before the move and record their baseline; after the move run the +same tests plus `cargo tree -p tinyagents-harness -e normal` and confirm no +`tinytools-std` edge remains. This is a relocation, so compare every test +assertion and fixture, not just the test count. + +### 2. Add feature gates and port all remaining standard tools + +Files: `crates/tinyagents-harness-tools/Cargo.toml`, +`crates/tinyagents-harness-tools/src/lib.rs`, +new `src/fs/`, `src/net/`, and `src/exec/` modules; remove the harness's +`tinytools-std` manifest entry. Add opt-in `fs`, `net`, and `exec` features; +preserve the existing `default = ["todos", "goals", "skills"]` behavior. +Add non-default `media` and `plan-review` features too, so migrating those +OpenHuman-only tools does not widen the defaults. Gate their modules in +`src/lib.rs` by those features. +Use this dependency-to-feature mapping, retaining existing mandatory +dependencies (`anyhow`, `async-trait`, `chrono`, `serde`, `serde_json`, +`tinyagents-harness`, `tinytools`, `tokio` base, and `tracing`) for old tool +families: + +| Dependency / Tokio capability | Feature | +| --- | --- | +| `base64`, `fs2`, `glob`, `parking_lot`, `regex`, `walkdir`, `libc`, `log` | `fs` | +| `futures-util`, `reqwest`, `sha2`, `log` | `net` | +| `rustix` (`fs`, `process` as target permits) | `exec` | +| Tokio `fs`, `io-util` capabilities | `fs`, `net`, and `exec` where used | +| Tokio `process` capability | `exec` | +| no provider/media dependency | `media` (generic progress callback only) | +| no added dependency | `plan-review` | +| `tempfile`, Tokio `test-util` | dev-dependencies only | + +The URL guard itself uses no optional dependency; replace its `log` calls with +`tracing`, because `tinyagents-harness` already depends on tracing. Do not add +a provider dependency or forward `media` to the harness's `multimodal` feature. +Make dependencies optional at the narrowest family feature that uses them. Preserve target-specific Unix/Windows dependencies. +Do not make a new feature activate an unrelated old default. + +Move all `tinytools-std` files except `sanitize` with their fixtures and tests, +preserving modules and public APIs under `tinyagents_harness_tools::{fs,net,exec}`: + +| New module | Source modules and contents | +| --- | --- | +| `fs` | `filesystem/` including `FsGate`, every listed filesystem tool (`file_read`, `file_write`, `edit_file`, `apply_patch`, `grep`, `glob_search`, `list_files`, `csv_export`, `read_diff`, `git_operations`, `run_linter`, `run_tests`, `update_memory_md`, `image_info`, `workspace_state`), `file_sink`, `read_only`, text helpers, fixtures, contract tests, and test support; `file_state/` including agent context, ops, types, tests; preserve `WorkspaceStateTool` and `FsGate` APIs | +| `net` | `network/` (`http_request`, `web_fetch`, `curl`, `pushover`, `NetGate`, payment/extractor traits, test support, fixtures, contract tests); `url_guard` call sites are changed to harness utility imports | +| `exec` | `detect_tools/` and `timeout/` (settings, command environment, output collection, process-group kill/reap, cleanup tracking, tests); `ExecGate`, `ExecBackend`, shell/install, and progress are added in later tasks | + +The exact source-root map is: + +| Source root | Destination root | +| --- | --- | +| `vendor/tinytools/crates/tinytools-std/src/filesystem/**` | `crates/tinyagents-harness-tools/src/fs/filesystem/**` | +| `vendor/tinytools/crates/tinytools-std/src/file_state/{agent_context.rs,mod.rs,ops.rs,types.rs}` | `crates/tinyagents-harness-tools/src/fs/file_state/{agent_context.rs,mod.rs,ops.rs,types.rs}` | +| `vendor/tinytools/crates/tinytools-std/src/file_state/test/{agent_context.rs,ops.rs}` | Combine all assertions into `crates/tinyagents-harness-tools/src/fs/file_state/file_state_tests.rs`; do not copy the old `test/mod.rs` aggregator | +| `vendor/tinytools/crates/tinytools-std/src/network/**` | `crates/tinyagents-harness-tools/src/net/network/**` | +| `vendor/tinytools/crates/tinytools-std/src/detect_tools/**` | `crates/tinyagents-harness-tools/src/exec/detect_tools/**` | +| `vendor/tinytools/crates/tinytools-std/src/timeout/**` | `crates/tinyagents-harness-tools/src/exec/timeout/**` | +| `vendor/tinytools/crates/tinytools-std/src/url_guard/mod.rs` | `crates/tinyagents-harness/src/url_guard.rs` | +| `vendor/tinytools/crates/tinytools-std/src/url_guard/mod_tests.rs` | `crates/tinyagents-harness/src/url_guard_tests.rs` | + +`**` includes all source files, nested modules, sibling test files, and +fixtures below that directory, keeping relative paths. For `file_state`, the +table lists production files explicitly because the two test files must move +to the mandated sibling test module. In `file_state/mod.rs`, replace +`mod test;` with `#[cfg(test)] #[path = "file_state_tests.rs"] mod tests;`. +The consolidated test file starts with `use super::*;` and retains every +assertion from both source test files. Do not copy +`tinytools-std/src/lib.rs`, its manifest, or `sanitize/**`; explicitly declare +the new root modules and keep sanitization with TinyTools. The remaining +source map is directory-to-directory, with file contents preserved +mechanically. Before deleting the source dependency, search all workspace +members and examples for references to the retired `tinytools-std` crate and +package; replace each actual implementation import with the new owning +crate/module. Do not retain a +compatibility re-export in TinyTools. Check every former `tinytools-std` public +item and test/fixture is accounted for; `sanitize` is the only intentional +retirement from the source crate because tinytools#4 owns it now. + +At the end of this task, these must compile independently: + +```sh +cargo check -p tinyagents-harness-tools --no-default-features --features fs +cargo check -p tinyagents-harness-tools --no-default-features --features net +cargo check -p tinyagents-harness-tools --no-default-features --features exec +cargo check -p tinyagents-harness-tools --no-default-features --features media +cargo check -p tinyagents-harness-tools --no-default-features --features plan-review +cargo check -p tinyagents-harness-tools --no-default-features +cargo check -p tinyagents-harness-tools --no-default-features --features fs,net,exec,media,plan-review +``` + +The final command proves a disabled-tools build remains lean and does not +accidentally rely on a new module's imports. Add or update crate docs showing all opt-in features (`fs`, `net`, `exec`, +`media`, `plan-review`), their public module paths, and the unchanged default +features. Validate actual dependency +resolution per feature with `cargo tree -e features -p tinyagents-harness-tools +--no-default-features --features media` and the equivalent `--features net`. +The `media` graph must not introduce `tinyinference-video`, +`tinyinference-image`, or any media-only dependency edge; do not assert that +`reqwest` is absent from the full graph if another mandatory dependency already +brings it in. Compare resolved package trees against `--no-default-features` +to attribute any graph difference correctly. The `net` graph must show its +intended `reqwest` edge under the `tinyagents-harness-tools feature "net"` +path. Also assert the `media`-only graph has no new `flate2`, `tar`, or `zip` +edge relative to the no-default baseline. Add these checks to a CI assertion +or a checked-in script; the required shell checks are: + +```sh +diff -u \ + <(cargo tree -p tinyagents-harness-tools -e normal --no-default-features) \ + <(cargo tree -p tinyagents-harness-tools -e normal --no-default-features --features media) +! cargo tree -p tinyagents-harness-tools -e normal --no-default-features --features media \ + | rg 'tinyinference-(video|image)|(^| )((flate2|tar|zip))($| )' +cargo tree -p tinyagents-harness-tools -e features -i reqwest --no-default-features --features net \ + | rg 'tinyagents-harness-tools feature "net"' +``` + + +### 3. Establish `ExecGate`, backend, result, and progress contract + +Files: `crates/tinyagents-harness-tools/src/exec/ports.rs`, +`exec/progress.rs`, `exec/types.rs`, `exec/mod.rs`, plus sibling +`ports_tests.rs` and `progress_tests.rs`. + +Implement the contract above with imports and derives exactly as shown. `Exec` +types are public through `exec`; the progress meter itself is cloneable and +thread-safe and follows the current `ShellOutputMeter` behavior: count bytes +from both streams, process `\n` and `\r` as line endings, retain a maximum +1,024 characters of pending line and only the last complete nonempty line of +up to 120 characters, and recover poisoned mutexes with `into_inner()`. Keep +the 15-second interval and exact progress messages (`running Ns · no output +yet`, byte/KB/MB figures, optional last line). The progress wrapper reports +through `ToolRunContext::report_progress` when context is present and performs +no ticker task otherwise. + +Tests use a fake `ExecGate` and `ExecBackend` to prove the public seam is +object-safe, thread-safe, and that no-context calls work. For this API-addition +task, the expected first red is compile failure because the public types do +not exist yet; it is not a behavioral bug regression. In task 4, first add +the required denial-before-backend test and verify its assertion fails against +an implementation that calls the backend before authorization, then implement +the ordered call sequence and confirm it passes. + +### 4. Port `shell` as `ExecTool` and map the host contract + +Files: `crates/tinyagents-harness-tools/src/exec/tool.rs`, +`exec/tool_tests.rs`, `exec/shell_support.rs`, +`exec/shell_support_tests.rs`, `exec/shell_description.rs`, and static schema +JSON beside the tool. Copy `shell.rs`, `shell_support.rs`, +`shell_description.rs`, and the source behavior tests from the OpenHuman +checkout into this module. Rename the tool's public struct to `ExecTool` and +wire it through required injected `Arc<dyn ExecGate>` and +`Arc<dyn ExecBackend>` fields. It is named `exec` (wire identity), preserving +the shell schema (`command`, optional `category`, `timeout_secs`), exact shell +description for POSIX/cmd flavors supplied as a generic `ShellFlavor`, +`PermissionLevel::Execute`, unbounded default timeout, and 30,000-character +result cap. + +Execution order is a contract and has focused tests: parse/validate the +command; form `ExecAuthorizationRequest`; await `ExecGate::authorize`. A gate +deny maps to the source refusal result and records exactly one +`GateDenied` event (`authorized=false`, `allowed=false`, +`process_attempted=false`), without calling the backend. On permit, build +`ExecRequest` with the permit's authorized backend selection/capability and +call the backend with context and output meter. Map each typed result and +record exactly one matching event: `Refused` maps to the source refusal result +and records `BackendRefused` (`authorized=true`, `allowed=false`, +`process_attempted=false`); `LaunchFailed` maps to the stable execution-failure +result and records `LaunchFailed` (`authorized=true`, `allowed=true`, +`process_attempted=true`, `success=false`); `Completed` maps its output to the +source `ToolResult` and records `Completed` (`authorized=true`, `allowed=true`, +`process_attempted=true`, with success derived from that final result). Do not +return early from a permitted path before its one audit record. Do not call +policy a second time after approval; the host's gate must perform its final +revalidation immediately before dispatch, and the selected backend must refuse +before process creation when configuration/preflight fails. + +Port/retain behavior assertions for schema/name/permission/timeout; output +Unicode and per-stream truncation boundaries; success with/without stderr; +nonzero exit and signal mapping; exit 124 note; explicit timeout text and +captured output; progress heartbeats and last-line accounting; environment +defaults/override; native, sandbox, Docker and TinyBox adapter fixture results +where they are representable through `ExecOutput`. Keep adapter-specific +scenario tests in OpenHuman #213; the TinyAgents suite tests only fake ports +and the in-crate native backend. Verify installer argv is not routed through +shell string construction. + +### 5. Port process supervision and argv-only install + +Files: `crates/tinyagents-harness-tools/src/exec/install_tool.rs`, +`exec/install_tool_tests.rs`, and `exec/native_backend.rs` with its sibling +`native_backend_tests.rs` if the existing timeout module does not fully cover +the behavior after relocation. + +Copy the install tool into the new `exec` module as `InstallTool`, requiring +the same injected gate/backend as `ExecTool`. Preserve the manager allowlist +and each exact binary/argv prefix; validate package names against the source +ASCII set and 200-byte cap; reject missing PATH executables; preserve the +interactive approval, feature-enabled, autonomy, and action-budget checks by +representing these as `ExecGate` denials in the host adapter, not as +TinyAgents policy. The tool creates `ExecOperation::Install`, awaits the gate, +then gives the backend an `ExecRequest` with `shell_command: None`, exact +separate `program` and `args`, workspace action dir, 300-second timeout, and +authorized environment. There is no `sh -c`, `cmd /C`, or concatenated argv +path for installs. Keep source error/output/truncation wording. + +The native backend configures child process groups, streams stdout and stderr +into the bounded progress meter and captured output, and invokes the existing +`output_or_kill` / `output_unbounded` cleanup primitives. Strengthen +`output_or_kill` so it establishes a `ProcessCleanup` scope for the timed +operation, and on `Elapsed` awaits that scope's `wait()` before returning the +same `Elapsed` result. `output_unbounded` continues to signal its owned waiter +when its future is dropped; it does not wait in `Drop`. For cancellation of a +whole tool call, require the owner to run the call inside a retained +`ProcessCleanup` scope; after dropping the call future the owner awaits +`wait()` before acknowledging cancellation. Wire and document this in the +native backend and the OpenHuman #213 turn adapter. Keep the scope handle alive +until its wait has completed. On Windows use the existing platform-supported +cleanup path; do not claim descendant termination where the platform does not +provide it. Tests use a temporary executable +script/process and assert: only the expected argv reaches it; malformed +package inputs never launch; denied calls have a zero backend count; timeout +does not return until the owned waiter's `Reaped` acknowledgment after +process-group kill, output-pipe drain, and child reap. On return, assert that +neither the child nor its descendant remains alive and that captured pipes are +closed. Separate stdout/stderr truncation remains Unicode safe; exit status +survives. Cancellation tests run +the call under `ProcessCleanup::scope`, drop the call future, then await +`ProcessCleanup::wait()` and assert the child and descendant are gone. The +cleanup acknowledgement is the `wait()` completion after the owned waiter has +killed, drained, and reaped; do not add a public test-only hook to observe it. +Do not assert or imply that `Drop` itself waits. No real package manager or network install is +run. + +Port the OpenHuman `install_tool` behavioral tests with fake ports. Write +behavioral RED first for injection attempts in the package parameter and +denial-before-spawn, then make the tests green. The timeout/cancel tests must fail if only the direct child is killed, if +`output_or_kill` returns before the timeout reaper acknowledgment, or if the +owner reports cancellation complete before `ProcessCleanup::wait()` resolves. + +### 6. Move parse-once schema utility without the host module tree + +Files: `crates/tinyagents-harness-tools/src/schema.rs` (or +`src/schema_cache.rs`), `src/lib.rs`, and sibling schema tests. Copy the +`static_schema!` macro and `parse_embedded` semantics from OpenHuman's +`tools/schema_cache.rs`, including call-site panic text and cloned +`LazyLock<Value>`. Export the macro and parse helper only as required by +downstream tool crates; do not create an OpenHuman-shaped `tools` tree. Move +the shell and install schema JSON into their respective exec tool directories, +use the macro in both tools, and retain malformed-schema and parsed-once/value +tests. No dynamic schema cache or host type is introduced. + +### 7. Move plan review and media artifact wrappers behind ports + +Files: `crates/tinyagents-harness-tools/src/plan_review/{mod.rs,port.rs,tool.rs,tool_tests.rs}`, +`src/media/{mod.rs,artifact_filing.rs,artifact_tool.rs,artifact_tool_tests.rs}`, +and `src/lib.rs`. + +Port `RequestPlanReviewTool` to `plan_review::RequestPlanReviewTool` with an +`Option<Arc<dyn PlanReviewPort>>`. Preserve tool name, schema, permission, +external-effect flag, unbounded timeout, argument trimming, and result text. +The port owns interactive/noninteractive policy, review event delivery, +plan-mode gating, and its own expiry. If absent or failing, return a clear +error and never approve. Tests use a deterministic fake port to assert request +fields, preserved resolution-to-result mapping, missing-port fail-closed +behavior, and that feedback is not silently discarded. + +Port `MediaArtifactTool<T>` to `media::MediaArtifactTool<T>` with an injected +`Arc<dyn ArtifactFiling>`. Retain the stable inner tool instance for all +borrowed metadata and static Tool declarations alongside the per-call factory. +Add the configured `MediaArtifactKind::{Image, Video}` and preserve its filer +kind and heartbeat label. Add the transport-free `MediaJobProgress` and +`MediaProgressFn` contract above; do not add tinyinference-video/image or +reqwest solely to name a callback type. Update the per-call inner factory to +accept `Option<&dyn ToolRunContext>` instead of `ApprovalChatContext`, letting +the host-specific factory downcast its own context in the adapter. Pass that +same context to each `ArtifactFiling::file` call. Test wrapper metadata +(name/description/schema/category/policy/permissions) against the retained +inner both with and without a factory, and test image/video config mapping to +filer kind and heartbeat labels. Add a concurrent routing test using one +shared wrapper and two simultaneous invocations with distinct +fake `ToolRunContext` host extensions; the fake filer must observe each +invocation's own extension for every generated artifact, with the corresponding +IDs and paths rewritten only in that invocation's result. Remove OpenHuman `Production`, artifact store, +TinyArtifacts, TinyStorageDrivers, file-root, blob-store, approval-chat, and +product reference-allowlist dependencies. Keep the generic output walk, +per-entry continuation, ID/error insertion, path rewrite, title cap/suffix, +and all `Tool` metadata/progress delegation. Where a generated `path` is a +blob key, pass that key as `PathBuf` to the port unchanged; the host filing +implementation decides how to fetch it from its own blob store. Tests use +fake filers and retain the source tests for single/multiple outputs, success, +failure isolation, unknown result shape, markdown path replacement, and +per-call progress reporting. OpenHuman's file-root/trusted-reference and +blob-cleanup checks remain host-adapter responsibilities in #213. + +### 8. Public API, manifest cleanup, docs, and migration audit + +Files: all previous changed files plus crate README/docs if applicable. + +Remove `tinytools-std` from every workspace package manifest and from +`Cargo.lock` through Cargo resolution. Keep TinyTools' definitions-only +boundary intact and do not touch the TinyTools gitlink. Search the full +TinyAgents workspace and test fixtures for references to the retired +`tinytools-std` crate or package, its former sanitization module, OpenHuman +module names, and OpenHuman crate paths. The only remaining `sanitize` imports must resolve to +the already-merged TinyTools API; there must be no `tinytools-std` dependency +or compatibility path. Update harness-tools README with feature selection, +ports, constructors, and the distinction between reusable tool mechanisms +and host authorization/storage. + +## Final verification + +Run these in the TinyAgents checkout after all tasks, through +`scripts/ci-cancel-aware.sh` for long commands. Do not export +`CARGO_TARGET_DIR` or place build output under `/tmp`. + +```sh +cargo fmt --all -- --check +cargo check --workspace --all-targets +cargo clippy --workspace --all-targets -- -D warnings +cargo test --workspace +cargo check -p tinyagents-harness-tools --no-default-features +cargo check -p tinyagents-harness-tools --no-default-features --features fs +cargo check -p tinyagents-harness-tools --no-default-features --features net +cargo check -p tinyagents-harness-tools --no-default-features --features exec +cargo check -p tinyagents-harness-tools --no-default-features --features media +cargo check -p tinyagents-harness-tools --no-default-features --features plan-review +cargo check -p tinyagents-harness-tools --no-default-features --features fs,net,exec,media,plan-review +cargo test -p tinyagents-harness-tools --no-default-features --features fs +cargo test -p tinyagents-harness-tools --no-default-features --features net +cargo test -p tinyagents-harness-tools --no-default-features --features exec +cargo test -p tinyagents-harness-tools --no-default-features --features media +cargo test -p tinyagents-harness-tools --no-default-features --features plan-review +cargo test -p tinyagents-harness-tools --no-default-features --features fs,net,exec,media,plan-review +cargo tree -p tinyagents-harness -e normal +cargo tree -p tinyagents-harness-tools -e normal +cargo tree -p tinyagents-harness-tools -e normal --no-default-features +cargo machete +``` + +Also re-run focused `tinyagents-harness` multimodal/URL guard tests and the +`tinyagents-harness-tools` tests for each feature, both no-default and +combined-feature builds. Confirm the resolved dependency tree has no +`tinytools-std` package. Review copied tests against source and report any +assertion that cannot be retained with the reason; do not silently drop +coverage. Formatting, check, clippy, workspace tests, feature-disabled builds, +and dependency audit are required before the TinyAgents PR. The OpenHuman +adapter and its integration tests are separate work in #213 and are not +claimed by this leaf. diff --git a/docs/plans/2026-10-19-runtime-session-host-seams.md b/docs/plans/2026-10-19-runtime-session-host-seams.md new file mode 100644 index 000000000..846e79ec1 --- /dev/null +++ b/docs/plans/2026-10-19-runtime-session-host-seams.md @@ -0,0 +1,2777 @@ +# Runtime and session host seams + +Issue: `senamakel/tinyagents#63`, under `tinyhumansai/openhuman#8088`. + +This plan moves generic session behavior into TinyAgents and leaves OpenHuman +with adapters for tenant authority, SaaS storage leases, approvals, progress, +memory, and learning. Implementation starts only after the integrated +`senamakel/tinyagents#60` changes are merged into this fork branch and the +affected baseline has been re-audited. The starting worktree currently contains +the `#57`, `#59`, and `#67` changes; do not recreate their types or lifecycle. + +The source state already provides `SessionRef`, `TranscriptMeta`, +`RebindableSessionStores`, `WorkspaceStores`, `CommitReceipt`, +`SessionHooks::after_commit`, harness `TurnSummary` and `ToolOutcomeRecord`, +runtime tool admission, and harness `ToolSnapshot`. In particular, +`ToolSnapshot::new` rejects different specs sharing a name; it is not the +first-wins deduplicator this migration needs. + +## Boundaries and invariants + +* Keep session identity and transcript repair in `tinyagents-session`; runtime + lifecycle and generic tool policy in `tinyagents-runtime`; context + configuration and compaction strategy in `tinyagents-harness` (the + `#7390`/`#8088` audit found no upstream equivalent). +* No `tinyagents-harness` or `tinyagents-session` dependency on runtime, no + runtime re-export of `ToolSnapshot`, and no compatibility shims in the host. + Runtime may use the already existing harness snapshot type. +* Tenant authority is resolved by the host. A session key, model argument, + transcript metadata field, or tool argument never selects a tenant or store + scope. +* Preserve frozen prompt and recorded-tool resume behavior. Persisted tools + remain available through deferred executors when the host can rebuild them; + explicit per-turn overrides and host-only belts remain authoritative. +* Preserve two distinct rebind behaviors: generic retained `AgentStores` + handles keep following a provider rebind; an explicit host lease pins the + backend and tenant admitted when that lease was opened. New handles opened + after rebind use the new backend. Rebinding never changes tenant scope. +* Preserve every existing test assertion unless behavior is deliberately + retired by this issue. Add regression tests that fail against the old + behavior before implementing each fix. Keep Rust tests in sibling + `*_tests.rs` files, declared using `#[cfg(test)]` and + `#[path = "<module>_tests.rs"]`; each test file begins with `use super::*;`. + +## Proposed public seams + +The signatures below are concrete API contracts. Implement them as written, +using the named module's existing error types and without adding dependency +edges beyond the existing crate chain. + +### Transcript identity and repair (`tinyagents-session`) + +Do not add a raw-stem override or `SessionRef::from_legacy_stem`. The +canonical stem builder in `transcript/session.rs` sanitizes and digests +components, so it cannot reproduce arbitrary historical filenames. Build the +canonical identity from the stable logical thread key and resolved agent id +using the existing `SessionRef::scoped` or `SessionRef::child_of`, then use +`TranscriptLocator::adopt_legacy` to read old timestamped stems and adopt +their rows into that canonical identity. Adoption is one-time and read-only +with respect to the old files. `TranscriptTarget::for_session` and +`SessionBuilder::session` already derive `session_id` and `parent_session_id`; +do not add a raw legacy-stem identity constructor or a duplicate identity +builder. + +Add both replay helpers to the exact existing module +`crates/tinyagents-session/src/transcript/adoption.rs` and export them from +`transcript.rs`. The transcript representation is used by codec replay; the +inference `Message` representation is used by the runtime driver. Both must +remove exactly the same leading tool head after the contiguous system prefix. + +```rust +pub fn repair_orphaned_tool_head( + rows: &[TranscriptMessage], +) -> Cow<'_, [TranscriptMessage]> { + let prefix = rows.iter().take_while(|row| row.role == "system").count(); + let orphans = rows[prefix..] + .iter() + .take_while(|row| row.role == "tool") + .count(); + if orphans == 0 { + Cow::Borrowed(rows) + } else { + Cow::Owned(rows[..prefix].iter() + .chain(&rows[prefix + orphans..]) + .cloned() + .collect()) + } +} + +pub fn repair_orphaned_tool_head_messages( + history: &[tinyinference_llm::message::Message], +) -> Cow<'_, [tinyinference_llm::message::Message]> { + use tinyinference_llm::message::Message; + let prefix = history.iter() + .take_while(|message| matches!(message, Message::System(_))) + .count(); + let orphans = history[prefix..].iter() + .take_while(|message| matches!(message, Message::Tool(_))) + .count(); + if orphans == 0 { + Cow::Borrowed(history) + } else { + Cow::Owned(history[..prefix].iter() + .chain(&history[prefix + orphans..]) + .cloned() + .collect()) + } +} +``` + +Add `use std::borrow::Cow;` in `adoption.rs`. Do not parse tool-call ids or +delete assistant rows: the historical repair is the leading run of orphan +tool-result rows only. + +This deliberately removes only leading `tool` rows immediately after the +contiguous system prefix, matching OpenHuman +`driver.rs::repair_orphaned_tool_head` and `codec.rs::replayable_host_rows`. +It does not remove a trailing assistant +tool-call round, does not remove later orphan rows, and does not rewrite the +durable source during resume. Later orphans remain visible to normal strict +transcript validation and fail there. The return value is a replay view; when +there are no leading orphans, the helper returns the original slice through +`Cow::Borrowed` rather than cloning it. + +### Host identity resolution and store leases (`tinyagents-session`, +`tinyagents-runtime`) + +Keep authority resolution at the host boundary and keep generic rebind behavior +intact. Add an explicit lease wrapper in `port/types.rs`: +Its constructor remains private; only a provider implementation that can +attest to pinned lifetime may create it. + +Add this error and lease contract in `port/types.rs`: + +```rust +#[derive(Debug, thiserror::Error)] +pub enum SessionStoreLeaseError { + #[error("provider cannot guarantee pinned session stores")] + Unsupported, + #[error("opening pinned session stores failed: {0}")] + Open(#[source] tinystoragedrivers_core::StorageError), +} + +#[derive(Clone)] +pub struct SessionStoreLease { + stores: AgentStores, +} + +impl SessionStoreLease { + pub fn stores(&self) -> &AgentStores { &self.stores } + pub(crate) fn pinned(stores: AgentStores) -> Self { + Self { stores } + } +} +``` + +Add `SessionStoreProvider::lease_for_agent(&self, agent_id: &str)` returning +`Result<SessionStoreLease, SessionStoreLeaseError>`. Its default returns +`Err(SessionStoreLeaseError::Unsupported)`; it must never wrap +`try_for_agent`, because that may return forwarding handles and does not prove +pinned lifetime. Add the method body to the trait: + +```rust +fn lease_for_agent( + &self, + _agent_id: &str, +) -> Result<SessionStoreLease, SessionStoreLeaseError> { + Err(SessionStoreLeaseError::Unsupported) +} +``` + +`DriverSessionStores` and `RebindableSessionStores` support leases. The former +is immutable after construction, including actor providers returned by +`bind_actor`, so it may open direct stores and return them in the lease. The +latter takes one `current()` snapshot and delegates lease opening to that +immutable provider. In both cases, the lease owns the returned non-forwarding +`AgentStores`; each direct store adapter owns its scoped backend handle. If +agent open fails, return `SessionStoreLeaseError::Open(error)` without +creating a lease. Keep `for_agent`, `try_for_agent`, and existing retained handles +live-following `rebind` exactly as they are today. `bind_actor` remains the +host-admission seam: the host resolves the actor and supplies the already +admitted backend; TinyAgents neither resolves tenant scope nor infers it from +agent/session keys. A provider that cannot attest that its handles pin their +backend returns `Unsupported`. + +The immutable driver implementation is: + +```rust +fn lease_for_agent( + &self, + agent_id: &str, +) -> Result<SessionStoreLease, SessionStoreLeaseError> { + let stores = DriverSessionStores::try_for_agent(self, agent_id) + .map_err(SessionStoreLeaseError::Open)?; + Ok(SessionStoreLease::pinned(stores)) +} +``` + +The rebindable implementation uses one immutable snapshot: + +```rust +fn lease_for_agent( + &self, + agent_id: &str, +) -> Result<SessionStoreLease, SessionStoreLeaseError> { + self.current().lease_for_agent(agent_id) +} +``` + +`current()` takes and releases the `RwLock` read guard while cloning the +current provider `Arc`; `lease_for_agent` then runs on that immutable +snapshot. `DriverSessionStores` carries the tenant `Scope`, so no second scope +lookup occurs after opening. Because `bind_actor` in `drivers/mod.rs` returns +`Arc<DriverSessionStores>` as `Arc<dyn SessionStoreProvider>`, its result +supports `lease_for_agent` directly; do not rely on the trait's unsupported +default for actor-bound storage. + +Rebind may update the provider's current backend for future opens, but must +not mutate an already returned lease. `WorkspaceStores::open` +continues to cache its ordinary workspace-bound handles; it must not become a +global mutable routing point. + +### Post-commit bridge (`tinyagents-runtime`) + +`CommitReceipt<C>` contains only `SessionTurnOutcome`, transcript options and +the optional transcript commit record. `SessionTurnOutcome` has history, +output, and interrupted status; it has no tool call durations, argument +digests, or failure classifications. Do not derive those values from +transcript prose. The runtime must capture actual tool outcomes at the +`tinyagents_harness::turn_outcome::ToolCallOutcome` boundary and pass a +separate observation alongside the receipt. Add this exact value in +`runtime/types.rs`: + +```rust +#[derive(Clone, Debug)] +pub struct PostCommitObservation<C: Clone + Send + Sync + 'static = ()> { + pub receipt: CommitReceipt<C>, + pub input: tinyinference_llm::message::Message, + pub usage: Option<tinyagents_session::transcript::TurnUsage>, + pub tool_outcomes: Vec<tinyagents_harness::turn_outcome::ToolCallOutcome>, + pub thread_id: Option<String>, + pub agent_id: Option<String>, +} +``` + +Thread +the generic observation sink from `DriverOutcome` through the runtime-owned +finalization task. Add a `tool_outcomes` field to `DriverOutcome`, populated +by the TinyAgents harness adapter by draining its existing +`ToolOutcomeSink`; custom drivers default to an empty vector. Keep the +argument and result text only for the duration of the callback. The summary +adapter hashes serialized arguments, classifies failure using +`classify_tool_failure`, and emits only the digest, name, success, class and +duration in `ToolOutcomeRecord`. + +Add `PostCommitSink<C>` next to `SessionHooks` in `runtime/hooks.rs`: + +```rust +#[async_trait] +pub trait PostCommitSink<C: Clone + Send + Sync + 'static = ()>: Send + Sync { + async fn after_commit(&self, observation: &PostCommitObservation<C>) + -> Result<(), RuntimeError>; +} +``` + +Add imports for `PostCommitObservation` and the trait-object `Arc`. Add +`post_commit_sinks: Vec<Arc<dyn PostCommitSink<C>>>` to `SessionBuilder` and +`Session`, initializing it with `Vec::new()` in each constructor. Add this +builder method and move the vector from `build` into `Session::new`: + +```rust +pub fn post_commit_sink(mut self, sink: Arc<dyn PostCommitSink<C>>) -> Self { + self.post_commit_sinks.push(sink); + self +} +``` + +The finalization owner already spawns a task after durable append and awaits +its `JoinHandle` without observing cancellation (`runtime/session.rs`, commit +path). Invoke sinks inside that same owned task so dropping the caller future +cannot cancel a durable transition's finalization. Do not put additional +awaits before the existing commit/finalization handoff. Sink errors remain +observational. Registration is an explicit ordered `Vec<Arc<dyn +PostCommitSink<C>>>` on `SessionBuilder`; no-op is an empty vector. Preserve +`SessionHooks::after_commit` as the existing lifecycle hook and call it once +from the same task; do not replace it or introduce an independently +cancellable second tail. Add the `TurnSummary` adapter as a helper that +consumes `PostCommitObservation`, while hosts continue to supply product +progress/memory/learning sinks. + +The existing `TerminalGuard::finalize_commit` owns terminal delivery. Change +its signature to +`finalize_commit(receipt, observation, sinks) -> JoinHandle<()>`, where +`observation: PostCommitObservation<C>` and `sinks` is the already-cloned +ordered sink vector. Inside the spawned task, preserve this order exactly: +`hooks.after_commit(receipt.clone())`, each `sink.after_commit(&observation)` +in vector order, `hooks.on_terminal_outcome(outcome)`, then +`hooks.on_terminal(terminal)`. Ignore/log each postcommit error without +short-circuiting. The existing call site after durable append constructs the +observation and immediately stores/awaits the returned join handle; add no +await between the successful append and spawning the owner task. + +Construct summary metadata using only captured facts. Convert input/output to +their message text; hash `serde_json::to_vec(&call.arguments)` with SHA-256; +classify failure text only when `success == false`; copy measured duration. +Never invent a duration, success flag, thread id, agent id, or usage value. +Because `TurnSummary` requires a thread, its adapter returns `None` when +`thread_id` is absent or empty. For usage, map captured input/output/cache +token counts and context window; leave charge absent unless an explicit +provider charge was captured. + +Implement that adapter as this pure function in `runtime/summary.rs`: + +```rust +pub fn turn_summary<C: Clone + Send + Sync + 'static>( + observation: &PostCommitObservation<C>, +) -> Option<tinyagents_harness::host::TurnSummary> { + use sha2::{Digest, Sha256}; + let thread = observation.thread_id.as_deref()?.trim(); + if thread.is_empty() { return None; } + let usage = observation.usage.as_ref().map_or_else(Default::default, |usage| { + tinyinference_llm::usage::Usage { + input_tokens: usage.usage.input, + output_tokens: usage.usage.output, + total_tokens: usage.usage.input.saturating_add(usage.usage.output), + cache_read_tokens: usage.usage.cached_input, + context_window_tokens: (usage.usage.context_window > 0) + .then_some(usage.usage.context_window), + ..Default::default() + } + }); + let mut summary = tinyagents_harness::host::TurnSummary::new( + tinyagents_harness::ids::ThreadId::from(thread), + observation.agent_id.as_deref()?, + ).with_text( + observation.input.text(), + observation.receipt.outcome.output.clone().unwrap_or_default(), + ).with_usage(usage); + for call in &observation.tool_outcomes { + let bytes = serde_json::to_vec(&call.arguments).ok()?; + let digest = Sha256::digest(bytes); + let args_digest = digest.iter().map(|byte| format!("{byte:02x}")).collect(); + summary.record_tool_outcome(tinyagents_harness::host::ToolOutcomeRecord { + name: call.name.clone(), + args_digest, + success: call.success, + failure_class: (!call.success).then(|| { + tinyagents_harness::host::classify_tool_failure(&call.content).to_owned() + }), + duration_ms: call.duration_ms, + }); + } + Some(summary) +} +``` + +Add `sha2` as a normal runtime dependency only if it is not already present in +the runtime manifest. Test missing/blank identity returns `None`, a failed +tool gets a digest and failure class without storing arguments/content, a +successful tool has no failure class, and the captured usage fields map +exactly. + +Put these executable tests in `runtime/summary_tests.rs`, declared from +`summary.rs` with `#[cfg(test)] #[path = "summary_tests.rs"] mod tests;`: + +```rust +use super::*; +use crate::types::{CommitReceipt, SessionTurnOutcome, TranscriptTurnOptions}; +use tinyagents_session::transcript::{MessageUsage, TurnUsage}; +use tinyinference_llm::message::Message; + +fn observation(thread_id: Option<&str>, agent_id: Option<&str>) -> PostCommitObservation { + PostCommitObservation { + receipt: CommitReceipt { + outcome: SessionTurnOutcome { + history: vec![Message::user("private input")], + output: Some("private output".into()), + interrupted: false, + }, + options: TranscriptTurnOptions { + request_id: None, + thread_id: thread_id.map(str::to_owned), + stream: false, + resume: crate::ResumeMode::Never, + context: (), + }, + transcript: None, + }, + input: Message::user("private input"), + usage: Some(TurnUsage { + provider: "test-provider".into(), + model: "test-model".into(), + usage: MessageUsage { + input: 12, + output: 5, + cached_input: 3, + context_window: 128, + ..MessageUsage::default() + }, + ts: String::new(), + reasoning_content: None, + tool_calls: Vec::new(), + iteration: 1, + }), + tool_outcomes: vec![tinyagents_harness::turn_outcome::ToolCallOutcome { + call_id: "call-private".into(), + name: "private_tool".into(), + arguments: serde_json::json!({"secret": "do-not-persist"}), + success: false, + content: "timeout: private details".into(), + duration_ms: 23, + }], + thread_id: thread_id.map(str::to_owned), + agent_id: agent_id.map(str::to_owned), + } +} + +#[test] +fn summary_requires_resolved_thread_and_agent_identity() { + assert!(turn_summary(&observation(None, Some("agent"))).is_none()); + assert!(turn_summary(&observation(Some(" "), Some("agent"))).is_none()); + assert!(turn_summary(&observation(Some("thread"), None)).is_none()); +} + +#[test] +fn summary_keeps_only_tool_digest_classification_and_captured_usage() { + use sha2::{Digest, Sha256}; + let summary = turn_summary(&observation(Some("thread-1"), Some("agent-1"))).unwrap(); + let expected = Sha256::digest(serde_json::to_vec(&serde_json::json!({ + "secret": "do-not-persist" + })).unwrap()); + let expected: String = expected.iter().map(|byte| format!("{byte:02x}")).collect(); + assert_eq!(summary.input, "private input"); + assert_eq!(summary.output, "private output"); + assert_eq!(summary.tool_outcomes.len(), 1); + assert_eq!(summary.tool_outcomes[0].args_digest, expected); + assert_eq!(summary.tool_outcomes[0].failure_class.as_deref(), Some("timeout")); + assert_eq!(summary.tool_outcomes[0].duration_ms, 23); + assert_eq!(summary.usage.input_tokens, 12); + assert_eq!(summary.usage.output_tokens, 5); + assert_eq!(summary.usage.total_tokens, 17); + assert_eq!(summary.usage.cache_read_tokens, 3); + assert_eq!(summary.usage.context_window_tokens, Some(128)); + let encoded = serde_json::to_string(&summary).unwrap(); + assert!(!encoded.contains("do-not-persist")); + assert!(!encoded.contains("timeout: private details")); +} + +#[test] +fn summary_usage_total_saturates_on_overflow() { + let mut observation = observation(Some("thread-1"), Some("agent-1")); + let usage = observation.usage.as_mut().unwrap(); + usage.usage.input = u64::MAX; + usage.usage.output = 1; + let summary = turn_summary(&observation).unwrap(); + assert_eq!(summary.usage.total_tokens, u64::MAX); +} + +#[test] +fn successful_tool_summary_has_no_failure_class() { + let mut observation = observation(Some("thread-1"), Some("agent-1")); + observation.tool_outcomes[0].success = true; + let summary = turn_summary(&observation).unwrap(); + assert!(summary.tool_outcomes[0].failure_class.is_none()); +} +``` + +### Tool surface policy and durable extensions (`tinyagents-runtime`, +`tinyagents-harness`) + +Do not implement host-only behavior as a registration filter alone. The +OpenHuman behavior is a deny-by-default `ToolCallPolicy` that wraps the +host's own gate and is installed in the actual call-policy chain. Add the +generic equivalent in `runtime/tool_registration.rs`: + +```rust +pub struct HostOnlyToolPolicy<C: Send + Sync + 'static> { + allowed: HashSet<String>, + host: Option<Arc<dyn tinyagents_harness::middleware::ToolCallPolicy<C>>>, +} + +impl<C: Send + Sync + 'static> HostOnlyToolPolicy<C> { + pub fn new( + allowed: impl IntoIterator<Item = String>, + host: Option<Arc<dyn tinyagents_harness::middleware::ToolCallPolicy<C>>>, + ) -> Self { + Self { allowed: allowed.into_iter().collect(), host } + } +} + +#[async_trait] +impl<C: Send + Sync + 'static> ToolCallPolicy<C> for HostOnlyToolPolicy<C> { + fn name(&self) -> &str { "host_only" } + + async fn check( + &self, + ctx: &RunContext<C>, + call: &ToolCall, + ) -> PolicyDecision { + if !self.allowed.contains(&call.name) { + return PolicyDecision::deny(format!( + "tool `{}` is not in this host-only belt", call.name + )); + } + match &self.host { + Some(host) => host.check(ctx, call).await, + None => PolicyDecision::Allow, + } + } +} +``` + +The generic declaration and impl both carry `C: Send + Sync + 'static` (the +same bound as `ToolCallPolicy<C>`); put the bound on the struct, impl, and +constructor so the public type is well-formed. Install it on the actual +`TurnPlan::tool_policy`, which `TurnDriver::prepare` wraps with +`ToolPolicyGateMiddleware`; registration filtering is only the model-visible +half of this policy. + +Use the harness's actual `ToolCallPolicy<C>` trait and `RunContext<C>` call +signature rather than a new policy protocol. The `check` implementation +denies names outside `allowed` before delegating admitted names to `host`; +empty `allowed` denies all. Runtime admission still controls model-visible +registration, while this policy independently guards actual invocations. +OpenHuman approvals remain the wrapped host policy. Extract first-wins +visible-spec deduplication as a separate helper over `Vec<Arc<ToolSpec>>` that +preserves input order and the winning allocation. Remove synthesized name +collisions before deduplication so the durable declaration wins. Keep +`ToolSnapshot::new` conflict rejection unchanged. Implement the two helpers +with these signatures: `dedup_visible_tool_specs(Vec<Arc<ToolSpec>>) -> +Vec<Arc<ToolSpec>>`, which retains the first name and original allocation in +input order, and `drop_synthesized_name_collisions(&[Arc<ToolSpec>], +Vec<Arc<ToolSpec>>) -> Vec<Arc<ToolSpec>>`, which removes synthesized specs +whose names occur in the durable set before the first-wins pass. + +Implement those helpers with this code (imports are `HashSet`, `Arc`, and +`tinytools::ToolSpec`): + +```rust +pub fn dedup_visible_tool_specs(specs: Vec<Arc<ToolSpec>>) -> Vec<Arc<ToolSpec>> { + let mut names = HashSet::new(); + specs.into_iter() + .filter(|spec| names.insert(spec.name.clone())) + .collect() +} + +pub fn drop_synthesized_name_collisions( + durable: &[Arc<ToolSpec>], + synthesized: Vec<Arc<ToolSpec>>, +) -> Vec<Arc<ToolSpec>> { + let durable_names: HashSet<&str> = durable.iter() + .map(|spec| spec.name.as_str()) + .collect(); + synthesized.into_iter() + .filter(|spec| !durable_names.contains(spec.name.as_str())) + .collect() +} +``` + +Test order and allocation directly with two `Arc<ToolSpec>` values named +`"read"`: the returned first element is `Arc::ptr_eq` to the first allocation, +and the second distinct name remains second. For durable collision, pass a +durable `"search"` and synthesized `["search", "render"]`; assert only +`"render"` survives before deduplication. + +The host-policy test uses a counter-backed host policy and calls +`ToolCallPolicy::check` directly with a real `RunContext::new(RunConfig::new("test"), ())`: +assert unknown name yields `PolicyDecision::Deny` with counter `0`, admitted +name forwards the same call and counter becomes `1`, and empty allowlist +denies. Also run it through `TurnPlan::tool_policy` and assert the underlying +tool executor is never invoked for an unknown call. + +Extend harness `ToolSnapshot` in `harness/src/tool/snapshot.rs` with a +Rust-type keyed extension carrier. TinyAgents defines no Composio/REPL names +or payload schemas; the host owns each payload type and its lifecycle: + +```rust +pub trait ToolSnapshotExtension: std::any::Any + std::fmt::Debug + Send + Sync { + fn as_any(&self) -> &dyn std::any::Any; +} +impl<T: std::any::Any + std::fmt::Debug + Send + Sync> ToolSnapshotExtension for T { + fn as_any(&self) -> &dyn std::any::Any { self } +} + +#[derive(Clone, Default)] +pub struct ToolSnapshotExtensions { + values: HashMap<TypeId, Arc<dyn ToolSnapshotExtension>>, +} + +impl ToolSnapshotExtensions { + pub fn insert<T: ToolSnapshotExtension>(&mut self, value: T) { + self.values.insert(TypeId::of::<T>(), Arc::new(value)); + } + + pub fn get<T: ToolSnapshotExtension>(&self) -> Option<&T> { + self.values.get(&TypeId::of::<T>())?.as_ref().as_any().downcast_ref::<T>() + } +} + +pub struct ToolSnapshot { + specs: Vec<ToolSpec>, + exact: bool, + extensions: ToolSnapshotExtensions, +} + +impl ToolSnapshot { + pub fn with_extension<T: ToolSnapshotExtension>(mut self, extension: T) -> Self { + self.extensions.insert(extension); + self + } + pub fn extension<T: ToolSnapshotExtension>(&self) -> Option<&T> { + self.extensions.get::<T>() + } +} +``` + +Implement `Debug` for the collection as: + +```rust +impl std::fmt::Debug for ToolSnapshotExtensions { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("ToolSnapshotExtensions") + .field("len", &self.values.len()) + .finish() + } +} +``` + +This lets `ToolSnapshot` keep its current `Debug` derive. Replacing an existing value of the same Rust +type replaces that value; different host types coexist. Clones share immutable +`Arc` payloads. `retaining` and `with_retained` copy the extension collection; +`exact` semantics remain unchanged. + +In the existing `ToolSnapshot::new`, initialize `extensions` to +`ToolSnapshotExtensions::default()`. Replace `retaining` with this complete +body: + +```rust +pub fn retaining(&self, keep: impl Fn(&ToolSpec) -> bool) -> Self { + Self { + specs: self.specs.iter().filter(|spec| keep(spec)).cloned().collect(), + exact: self.exact, + extensions: self.extensions.clone(), + } +} +``` + +Keep the existing early return in `with_retained`, and replace its success +tail with this exact code. The `Self::new` call continues to validate name +collisions: + +```rust +let count = retained.len(); +let mut merged = Self::new(self.specs.iter().cloned().chain(retained).collect())?; +merged.extensions = self.extensions.clone(); +Ok((merged, count)) +``` + +Keep `from_json` going through `Self::new`; decoded snapshots therefore have +an empty extension collection, and the existing array wire format stays +unchanged. + +Keep `ToolSnapshot::to_json` and `from_json` as the existing array-of-specs +format with no extension serialization. The harness transcript record already +stores arbitrary JSON, but TinyAgents cannot serialize a host-defined Rust +type without a host-owned codec. The host rebuilds typed extension values when +it reconstructs the current snapshot after resume. If a concrete host later +proves it needs durable extension payloads, that host adds a separate codec +and transcript record while the generic runtime carrier remains typed. + +### Context configuration (`tinyagents-harness`) + +Move `ContextConfig`, `CompactionSettings`, and `CompactionStrategy` from +OpenHuman `crates/openhuman-core/src/config/schema/context.rs` into the exact +TinyAgents files `crates/tinyagents-harness/src/config/context.rs`, +`config/mod.rs`, and `src/lib.rs`. Preserve the complete field set and serde +shape: `enabled`, `microcompact_enabled`, `autocompact_enabled`, +`microcompact_keep_recent`, `tool_result_budget_bytes`, +`summarizer_payload_threshold_tokens` (alias +`summarizer_payload_threshold_bytes`), `summarizer_max_payload_tokens` (alias +`summarizer_max_payload_bytes`), `summarizer_model`, +`prefer_markdown_tool_output`, `compaction_enabled`, +`compaction_trigger_tokens`, and `compaction_strategy`. Keep every default +function's current value, `#[serde(default)]` behavior, strategy spelling, +and `CompactionSettings` normalization of a zero trigger to `None`. +`CompactionStrategy::parse` must continue accepting +`task_state`/`typed`/`s4` and `summary`/`free_form`/`s0` after trimming, +case-folding and converting hyphens to underscores. + +The moved type may retain the fields whose comments refer to TinyJuice or +OpenHuman's compaction setting only if TinyAgents already consumes the +matching generic knob. The host retains environment variable loading and +product policy; `ContextConfig` itself owns only serializable context knobs. +The default for `microcompact_keep_recent` becomes a harness constant with the +same value as `DEFAULT_KEEP_RECENT_TOOL_RESULTS`; it must not import OpenHuman. +The existing host `CompactionSettings` call site and full TOML compatibility +are part of this task, even though host deletion/adoption of the rest of +`session_host` is later in #212. + +In `config/context_tests.rs`, use a TOML string with exactly these keys and +values: `enabled=false`, `microcompact_enabled=false`, +`autocompact_enabled=true`, `microcompact_keep_recent=9`, +`tool_result_budget_bytes=1234`, +`summarizer_payload_threshold_tokens=4000`, +`summarizer_max_payload_tokens=64000`, `summarizer_model="test-model"`, +`prefer_markdown_tool_output=false`, `compaction_enabled=true`, +`compaction_trigger_tokens=123`, `compaction_strategy="summary"`; assert +every field. A second fixture substitutes both aliases for their canonical +threshold keys and asserts the same parsed values. Check default round-trip by +serializing both values to `serde_json::Value` and comparing; do not add +`PartialEq` solely for this test. + +### Session identity resolution (`tinyagents-runtime`) + +Issue #63 implements and wires this seam; OpenHuman's concrete resolver and +adapter migration remain in #212. Put request/result in `runtime/types.rs`, +the resolver trait in `runtime/identity.rs`, and the resolver on +`SessionBuilder`/`Session`. One authority operation, `resolve_and_bind`, is +used by every transcript-reading path. `Session::turn` calls it after +`before_resume` and `apply_resume_preparation`, so hook-provided target +replacements are visible, but before locks or locator methods. Public +`Session::resume` has no hook in its existing contract; it calls the same +helper directly before any locator method and never reuses a prior operation's +identity. Both then call private `resume_admitted`. A +resolver-enabled session rejects absent identity, conflicting +`TurnOptions::session`, or changed hook-selected target before +`head_generation`, `read_session_transcript`, legacy adoption, and lock lookup. + +```rust +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct SessionIdentityRequest { + pub thread_id: Option<String>, + pub session_key: String, + /// Logical parent lineage for a delegated session; never authorization. + pub parent: Option<SessionRef>, + /// Child key paired with `parent`; never an actor or tenant selector. + pub child_key: Option<String>, +} + +#[derive(Clone)] +pub struct ResolvedSessionStores { + pub agent_id: String, + pub stores: tinyagents_session::port::SessionStoreLease, +} + +#[derive(Clone)] +pub struct SessionIdentity { + pub(crate) session: SessionRef, + pub(crate) agent_id: String, + pub(crate) stores: tinyagents_session::port::SessionStoreLease, +} + +impl SessionIdentity { + pub fn session(&self) -> &SessionRef { &self.session } + pub fn agent_id(&self) -> &str { &self.agent_id } +} + +pub trait SessionIdentityResolver<C: Send + Sync + 'static>: Send + Sync { + fn resolve( + &self, + request: &SessionIdentityRequest, + context: &C, + ) -> Result<ResolvedSessionStores, RuntimeError>; +} +``` + +Add the builder/storage wiring with these exact ownership fields. Put the +resolver and identity seed in a private builder config so existing +`.transcript(...)` and `.session(...)` remain source-compatible: + +```rust +pub(crate) struct IdentityConfig<C: Send + Sync + 'static> { + pub(crate) resolver: Arc<dyn SessionIdentityResolver<C>>, + pub(crate) session_key: String, + pub(crate) parent: Option<SessionRef>, + pub(crate) child_key: Option<String>, + pub(crate) meta: TranscriptMeta, +} +``` + +Define `IdentityConfig<C>` in `runtime/identity.rs` and import it from both +`builder.rs` and `session.rs`; its crate visibility is required because both +modules store it. In `builder.rs`, add `identity: Option<IdentityConfig<C>>` to +`SessionBuilder<C>`, initialize it to `None` in `new`, and add the method +implementation below. Keep `IdentityConfig<C>` in a private field so the +existing `.transcript(...)` and `.session(...)` methods remain source +compatible: + +```rust +pub fn identity_resolver( + mut self, + session_key: impl Into<String>, + meta: TranscriptMeta, + resolver: Arc<dyn SessionIdentityResolver<C>>, +) -> Self { + self.identity = Some(IdentityConfig { + resolver, + session_key: session_key.into(), + parent: None, + child_key: None, + meta, + }); + self +} + +pub fn child_identity_resolver( + mut self, + parent: SessionRef, + child_key: impl Into<String>, + meta: TranscriptMeta, + resolver: Arc<dyn SessionIdentityResolver<C>>, +) -> Self { + let child_key = child_key.into(); + self.identity = Some(IdentityConfig { + resolver, + session_key: child_key.clone(), + parent: Some(parent), + child_key: Some(child_key), + meta, + }); + self +} +``` + +In `build`, destructure the identity configuration before moving other +builder fields. Require `TranscriptCodec` if either transcript or identity +configuration exists, preserve the existing target construction, and pass +`identity` as the final constructor argument: + +```rust +let identity = self.identity; +let target = self.transcript.map(|config| crate::TranscriptTarget { + locator: config.locator, + stem: config.stem, + resume_agent: config.resume_agent, + session: config.session, + meta: config.meta, +}); +if (target.is_some() || identity.is_some()) && self.codec.is_none() { + return Err(RuntimeError::MissingDependency("TranscriptCodec")); +} +let mut session = Session::<C>::new( + self.driver, self.codec, self.hooks, self.prefix, self.tools, target, identity, +); +session.set_retain_recorded_tools(self.retain_recorded_tools); +Ok(session) +``` + +Change `Session::new` to accept `identity_config: Option<IdentityConfig<C>>` +after `target` and initialize `identity_config`, +`established_identity: None`, and `identity: None`. Add +`identity: Option<&'a SessionIdentity>` to `SessionStateView<'a>` and populate +it from `established_identity` (never the in-flight result). Export +`SessionIdentity`, `SessionIdentityRequest`, and `ResolvedSessionStores` from +`types`; export `SessionIdentityResolver` from `identity` in `lib.rs`. Define +`IdentityConfig`, `SessionIdentityResolver`, and `same_admission` in +`identity.rs`; import `SessionIdentityRequest`, `ResolvedSessionStores`, and +`SessionIdentity` from `crate::types`, `RuntimeError` from `crate::error`, and +`Arc` from `std::sync`. Because `session.rs` calls the comparison helper, +declare `same_admission` `pub(crate)` and import it as +`crate::identity::same_admission`. Keep `identity` private in `lib.rs` and +add `pub use identity::SessionIdentityResolver;` plus +`pub use types::{ResolvedSessionStores, SessionIdentity, SessionIdentityRequest};`. + +In `SessionBuilder::build`, require a codec when either transcript or +identity config is present, then pass the identity config into `Session::new`. +Do not require a fallback `TranscriptTarget`; the resolver's lease is the +target. Retain an optional legacy transcript target only as an unselected +fallback: resolver failure must not access it. Add this exact helper body to +`Session<C>`; pass the hook-selected target separately so the builder's +fallback is never mistaken for hook authority: + +```rust +fn resolve_and_bind( + &mut self, + options: &TurnOptions<C>, + hook_target: Option<&TranscriptTarget>, +) -> Result<(), RuntimeError> { + let Some(config) = self.identity_config.as_ref() else { return Ok(()); }; + self.identity = None; + let request = SessionIdentityRequest { + thread_id: options.thread_id.clone(), + session_key: config.session_key.clone(), + parent: config.parent.clone(), + child_key: config.child_key.clone(), + }; + let resolved = config.resolver.resolve(&request, &options.run_context.data)?; + if resolved.agent_id.is_empty() { + return Err(RuntimeError::InvalidSessionState("resolver returned empty agent id".into())); + } + let first_session = match (&request.parent, &request.child_key) { + (Some(parent), Some(child_key)) => SessionRef::child_of(parent, child_key), + (None, None) => SessionRef::scoped(&request.session_key, &resolved.agent_id), + _ => return Err(RuntimeError::InvalidSessionState( + "identity request has an incomplete child lineage".into(), + )), + }; + let identity = SessionIdentity { + session: first_session.clone(), + agent_id: resolved.agent_id, + stores: resolved.stores, + }; + if options.session.as_ref().is_some_and(|selected| + selected.first_generation() != first_session.first_generation()) + { + return Err(RuntimeError::InvalidSessionState("TurnOptions.session conflicts with resolved identity".into())); + } + // Only a previously admitted identity authorizes carrying a nonzero + // generation forward. A builder target is an unselected fallback, even + // when it names the same logical session at generation 9. + let session = if let Some(previous) = self.established_identity.as_ref() { + if !same_admission(previous, &identity) { + return Err(RuntimeError::InvalidSessionState( + "session identity or pinned store changed".into(), + )); + } + self.target.as_ref() + .and_then(|target| target.session.clone()) + .filter(|current| current.first_generation() == first_session.first_generation()) + .unwrap_or(first_session) + } else { + first_session + }; + let mut meta = config.meta.clone(); + meta.agent_id = Some(identity.agent_id.clone()); + let target = TranscriptTarget::for_session( + identity.stores.stores().transcripts.clone(), + session.clone(), + meta, + ); + if hook_target.is_some_and(|selected| !selected.same_binding(&target)) { + return Err(RuntimeError::InvalidSessionState("before_resume target conflicts with resolved identity".into())); + } + self.target = Some(target); + let mut identity = identity; + identity.session = session; + self.established_identity = Some(identity.clone()); + self.identity = Some(identity); + Ok(()) +} +``` + +`SessionIdentity` is a runtime-created private binding; the host result is +only `ResolvedSessionStores { agent_id, stores }`. Runtime constructs a root +or child `SessionRef` from configured logical naming inputs and the resolved +agent id; it accepts no host-provided stem. `same_admission` compares the +identity's first generation, agent id, and locator key/pointer as specified +below. + +Implement the binding predicate in `runtime/identity.rs` with the following +comparison semantics (the current target locator is read from the target that +will actually be resumed, after hook preparation): + +```rust +pub(crate) fn same_admission(a: &SessionIdentity, b: &SessionIdentity) -> bool { + if a.session.first_generation() != b.session.first_generation() + || a.agent_id != b.agent_id + { + return false; + } + let left = a.stores.stores().transcripts.destination_key(); + let right = b.stores.stores().transcripts.destination_key(); + match (left, right) { + (Some(left), Some(right)) => left == right, + (None, None) => Arc::ptr_eq( + &a.stores.stores().transcripts, + &b.stores.stores().transcripts, + ), + _ => false, + } +} +``` + +`resolve_and_bind` returns `InvalidSessionState` if this predicate fails for +the established identity, if `options.session` has another +`first_generation()`, or if the hook-selected target locator/session does not +match the newly resolved admission. It updates the target only after all +checks pass. `Session::turn` must call it after `apply_resume_preparation`; +public `Session::resume` calls it before the first locator access. This order +is intentional: hooks may prepare a lazy target but cannot grant authority to +change its destination. + +Extract the storage-reading body of `Session::resume` into private +`resume_admitted`. Public `resume` invokes `resolve_and_bind(options, None)` +and then `resume_admitted`; it must not add hooks or a synthetic request. +`turn_inner` retains its existing `before_resume` and preparation calls, then +snapshots `resume_preparation.transcript`, applies the preparation, then calls +`resolve_and_bind(options, hook_target.as_ref())` and acquires +`lock_turn(options)`, and calls `resume_admitted` directly. Both paths share +one authority predicate without changing public resume hook semantics. The +private reader performs no hook, resolver, lock, or target-selection work. +Validate `TurnOptions::session` and any hook-selected target before the lock +or any locator operation. + +Use this complete public wrapper (the old locator-reading match moves intact +to `resume_admitted`): + +```rust +pub async fn resume( + &mut self, + options: &TurnOptions<C>, +) -> Result<SessionResume, RuntimeError> { + if options.cancellation.is_cancelled() { + return Err(RuntimeError::Cancelled); + } + self.resolve_and_bind(options, None)?; + self.resume_admitted(options).await +} +``` + +The relevant `turn_inner` sequence is: + +```rust +let resume_preparation = cancelable( + &cancellation, + self.hooks.before_resume(request, options, self.state_view(false)), +).await?; +let hook_target = resume_preparation.transcript.clone(); +self.apply_resume_preparation(resume_preparation)?; +self.resolve_and_bind(options, hook_target.as_ref())?; +let _turn_lock = cancelable(&cancellation, async { Ok(self.lock_turn(options).await) }).await?; +let resumed = if options.resume == ResumeMode::Never { + false +} else { + self.resume_admitted(options).await?.loaded +}; +``` + +Keep the existing cancellation check immediately following the resume call. +`resume_admitted` is the old locator-reading body verbatim, renamed and made +private; the only structural change inside it is to remove the old early +return for absent `self.target` when identity config is active, because +`resolve_and_bind` has already installed the admitted target. + +The request carries logical naming inputs only: root key or parent plus child +key, and the optional conversation `thread_id`. It never carries an agent or +tenant selector. The host configures the parent/child seed when it builds the +session; that seed determines `SessionRef` lineage and is not authorization. +The resolver uses `options.run_context.data` to authorize the actor, choose +the resolved agent id, and open a tenant-scoped lease. Map +`SessionStoreLeaseError::Unsupported` or `Open(_)` into +`RuntimeError::InvalidSessionState(error.to_string())`; host resolver errors +remain the host's `RuntimeError`. Runtime constructs a root with +`SessionRef::scoped(request.session_key, resolved.agent_id)` or a child with +`SessionRef::child_of(request.parent, request.child_key)`. It binds the target +locator from `resolved.stores.stores().transcripts`; `for_session` stamps +`agent_id`, `session_id`, `parent_session_id`, and the canonical stem. A child +must never be replaced with a root merely because the resolver returns an +agent id: child lineage is preserved in the child `SessionRef` and its stem. +Resolver failure returns immediately without touching the current or fallback +locator. A session-bound target may keep its own fixed locator only when no +resolver is installed. + +Only when a prior `established_identity` exists and passes the first-generation, +agent, and pinned-lease comparison may re-admission retain the current +`self.target.session` generation. Verify its `first_generation()` against the +new logical identity before carrying it forward. With no established identity, +ignore `self.target` completely: a builder-provided target is an unselected +fallback and cannot choose the first admitted generation, even if it uses the +same canonical session at generation 9. Use the validated current generation +to construct the new target and identity. In +`ResumeMode::Never`, do not call `head_generation` just to discover a known +head: preserve the already-bound generation, then let normal commit/compaction +advance it. This keeps the sealed generation immutable and ensures the next +turn still appends to the current head. The first admission starts from +generation zero. Re-admission does not reset a compacted root or child to zero. + +Put the last established binding in `SessionStateView` as +`Option<&SessionIdentity>` for diagnostics and hook decisions; it is not +authority for the current turn. The current resolver result is produced only +after `before_resume` and `apply_resume_preparation`, then checked before any +lock or locator access. Store it on `Session<C>` for the duration of each operation, clear it before +each resolution attempt, and reject a later identity whose `first_generation` +or admitted lease destination differs from the established identity. This +prevents one live `Session` from crossing tenant boundaries. The resolver is +called on every turn before resume so a newly constructed runtime/process +repeats host identity admission; host adapters do not cache authority in the +serialized transcript. Compare each resolved `SessionRef::first_generation()`, +agent id, and locator `destination_key()` with the established binding. Reject +a change with `RuntimeError::InvalidSessionState` before opening a history +handle. When a locator has no destination key, require `Arc::ptr_eq`; a host +resolver may cache its lease for repeated calls in one actor context. + +## TDD work sequence + +Each step is independently reviewable. For each behavior, first add the +described failing test, run only that focused test to observe the regression, +then implement the smallest API/code change and rerun it. Keep existing +assertions and run the named crate's full tests after each API family. No host +source edits are part of this plan. + +Run every focused Rust test with this helper from the plan. It explicitly +enters the nested TinyAgents workspace and uses the super-worktree's +cancel-aware wrapper: + +```sh +run_tinyagents_test() { + cd -- "/home/enamakel/work/openhuman/worktrees/epic-8088-63-session-runtime/vendor/tinyagents" && \ + CARGO_BUILD_JOBS=1 CARGO_PROFILE_DEV_DEBUG=0 \ + "/home/enamakel/work/openhuman/worktrees/epic-8088-63-session-runtime/scripts/ci-cancel-aware.sh" cargo test "$@" +} +``` + +### 1. Repair only the historical leading orphan-tool head + +**Files:** `crates/tinyagents-session/src/transcript/adoption.rs`, +`transcript.rs`, and new `transcript/adoption_tests.rs` declared from +`transcript.rs` as `#[cfg(test)] #[path = "transcript/adoption_tests.rs"] mod adoption_tests;`. + +First add this failing runtime behavior test in +`crates/tinyagents-runtime/src/session_tests.rs`. The codec deliberately +retains prior raw rows while encoding only new inference messages, so the test +proves replay repair does not rewrite durable history: + +```rust +struct CaptureHistoryDriver(Arc<Mutex<Option<Vec<Message>>>>); + +#[async_trait] +impl SessionDriver for CaptureHistoryDriver { + async fn execute(&self, request: DriverRequest) -> Result<DriverOutcome, DriverFailure> { + *self.0.lock().unwrap() = Some(request.history.clone()); + let mut history = request.history; + history.push(Message::assistant("done")); + Ok(DriverOutcome { + output: Some("done".into()), + history, + partial: None, + interrupted: false, + outcome: None, + tool_outcomes: Vec::new(), + }) + } +} + +struct ToolAwareCodec; + +impl TranscriptCodec for ToolAwareCodec { + fn decode_history(&self, transcript: &SessionTranscript) -> Result<Vec<Message>, RuntimeError> { + Ok(transcript.messages.iter().map(|row| match row.role.as_str() { + "system" => Message::system(&row.content), + "user" => Message::user(&row.content), + "tool" => Message::tool(row.id.as_deref().unwrap_or("missing-call"), &row.content), + _ => Message::assistant(&row.content), + }).collect()) + } + fn reconcile( + &self, + prior: &[TranscriptMessage], + previous: &[Message], + next: &[Message], + _: &TranscriptTurnOptions, + ) -> Result<Vec<TranscriptMessage>, RuntimeError> { + let mut rows = prior.to_vec(); + rows.extend(next[previous.len().min(next.len())..].iter().map(|message| match message { + Message::System(_) => TranscriptMessage::system(message.text()), + Message::User(_) => TranscriptMessage::user(message.text()), + Message::Tool(_) => TranscriptMessage::tool(message.text()), + Message::Assistant(_) => TranscriptMessage::assistant(message.text()), + })); + Ok(rows) + } +} + +#[tokio::test] +async fn resume_repairs_leading_orphan_for_driver_without_rewriting_source_rows() { + let locator = memory(); + let session_ref = SessionRef::scoped("thread-orphan", "agent"); + let stem = session_stem(&session_ref); + let durable = vec![ + TranscriptMessage::system("frozen"), + TranscriptMessage::tool("orphan result").with_id("orphan-call"), + TranscriptMessage::user("next turn"), + ]; + seed(locator.as_ref(), &stem, &durable, &meta()); + let seen = Arc::new(Mutex::new(None)); + let mut session = SessionBuilder::new(Arc::new(CaptureHistoryDriver(seen.clone()))) + .codec(Arc::new(ToolAwareCodec)) + .session(locator.clone(), session_ref, meta()) + .build().unwrap(); + session.turn( + Message::user("next turn"), + TurnOptions { resume: ResumeMode::Session, ..TurnOptions::default() }, + ).await.unwrap(); + let history = seen.lock().unwrap().clone().unwrap(); + assert_eq!(history.iter().map(Message::text).collect::<Vec<_>>(), ["frozen", "next turn"]); + assert_eq!(&read(locator.as_ref(), &stem).messages[..durable.len()], durable.as_slice()); +} +``` + +Then add direct parity tests in the new session `adoption_tests.rs` for both +`TranscriptMessage` and inference `Message` representations. The direct rows +fixture is: + +```rust +#[test] +fn repair_removes_only_tool_rows_after_the_leading_system_prefix() { + let rows = vec![ + TranscriptMessage::system("frozen"), + TranscriptMessage::tool("legacy orphan").with_id("orphan"), + TranscriptMessage::user("next turn"), + TranscriptMessage::tool("later invalid").with_id("later"), + ]; + let repaired = repair_orphaned_tool_head(&rows); + assert_eq!(repaired.len(), 3); + assert_eq!(repaired[0], rows[0]); + assert_eq!(repaired[1], rows[2]); + assert_eq!(repaired[2], rows[3]); + assert_eq!(rows[1].id.as_deref(), Some("orphan")); // source untouched +} + +#[test] +fn later_orphans_are_left_for_strict_validation() { + let rows = vec![ + TranscriptMessage::system("frozen"), + TranscriptMessage::user("turn"), + TranscriptMessage::tool("later orphan").with_id("call-later"), + ]; + let repaired = repair_orphaned_tool_head(&rows); + assert_eq!(repaired.as_ref(), rows.as_slice()); +} + +#[test] +fn clean_head_borrows_the_original_rows() { + let rows = vec![TranscriptMessage::system("frozen"), TranscriptMessage::user("turn")]; + assert!(matches!(repair_orphaned_tool_head(&rows), Cow::Borrowed(_))); +} +``` + +Add the corresponding inference-view test in the same test file: + +```rust +#[test] +fn message_repair_removes_the_same_leading_tool_head() { + use tinyinference_llm::message::Message; + let history = vec![ + Message::system("frozen"), + Message::tool("orphan", "old result"), + Message::user("next"), + Message::tool("later", "invalid"), + ]; + let repaired = repair_orphaned_tool_head_messages(&history); + assert_eq!(repaired.len(), 3); + assert!(matches!(repaired[0], Message::System(_))); + assert!(matches!(repaired[1], Message::User(_))); + assert!(matches!(repaired[2], Message::Tool(_))); + assert_eq!(history.len(), 4); +} +``` + +`Message::tool` takes the call id and result content in that order. + +After the runtime red test fails, add the helper implementation to `adoption.rs`, call it on raw rows before codec decoding and on decoded `Message` history before dispatch, and keep `persisted` raw rows unchanged. Add a second test with a later orphan and assert it remains in the returned replay view, plus a no-orphan test asserting `matches!(..., Cow::Borrowed(_))`. Run `run_tinyagents_test -p tinyagents-runtime orphaned_tool_head` for the red/green behavioral regression, then `run_tinyagents_test -p tinyagents-session adoption_tests`, and `run_tinyagents_test -p tinyagents-runtime`. + +### 2. Adopt legacy transcripts into canonical identity without changing old files + +**Files:** `crates/tinyagents-session/src/transcript/session_tests.rs` and +`crates/tinyagents-session/src/port/drivers/transcripts_adoption_tests.rs`. + +Do not parse a legacy filename into a new canonical stem. Preserve the +existing `a_stem_is_deterministic_and_carries_no_timestamp` regression, and +add this explicit logical-key assertion to `transcript/session_tests.rs`: + +```rust +#[test] +fn canonical_identity_is_stable_and_distinct_from_legacy_stems() { + let first = SessionRef::scoped("thread-legacy", "agent-legacy"); + let second = SessionRef::scoped("thread-legacy", "agent-legacy"); + assert_eq!(session_stem(&first), session_stem(&second)); + assert_ne!(session_stem(&first), "1700000100_agent-legacy"); + assert_eq!(first.first_generation(), second.first_generation()); +} +``` + +Keep the existing complete `legacy_roots_are_folded_oldest_first_once` test +in `port/drivers/transcripts_adoption_tests.rs`. It creates timestamped stems, +calls `adopt_legacy`, checks oldest-first imported messages and canonical +metadata, verifies a repeated call is a no-op, and leaves legacy roots intact. +Do not duplicate that integration fixture in `transcript_tests.rs`; this plan +adds no new adoption algorithm. Run +`run_tinyagents_test -p tinyagents-session legacy_roots_are_folded_oldest_first_once` +and `run_tinyagents_test -p tinyagents-session canonical_identity_is_stable_and_distinct_from_legacy_stems`. +The runtime calls legacy adoption only through the existing +`ResumeMode::Session` path. Do not add a constructor that maps arbitrary old +raw filenames to `SessionRef`. + +Add this separate child case in `port/drivers/transcripts_adoption_tests.rs`. +Legacy root adoption must remain attached to the root `SessionRef`; asking a +delegated child to resume must not copy the parent's legacy root into the +child transcript: + +```rust +#[test] +fn legacy_root_adoption_does_not_replace_a_child_session() { + let locator = locator(); + write_root( + &locator, + "1700000100_planner", + &legacy("planner", "2026-01-01", 1, 5), + &["root history"], + ); + let root = SessionRef::scoped(THREAD, "planner"); + let child = SessionRef::child_of(&root, "delegate:worker"); + + assert!(locator.adopt_legacy(&child, THREAD, &meta(THREAD)).unwrap().is_none()); + assert!(locator.read_session_transcript(&child).is_none()); + assert!(locator.adopt_legacy(&root, THREAD, &meta(THREAD)).unwrap().is_some()); + assert_eq!(contents(&locator, &root), vec!["root history"]); + assert!(locator.read_session_transcript(&child).is_none()); +} +``` + +Implement this as a child guard in the driver-backed `adopt` before candidate +discovery: if `session.is_subagent()`, return `Ok(None)`. Keep root adoption's +existing integration test and this child no-cross-adoption test separate. + +### 3. Resolve identity before resume and bind the admitted lease + +**Files:** new `crates/tinyagents-runtime/src/identity.rs`, +`runtime/types.rs`, `runtime/builder.rs`, `runtime/session.rs`, `runtime/lib.rs`, +and `runtime/lib_tests.rs` (the existing sibling test module). + +Add these red tests first: (a) direct `Session::resume` uses the resolved +locator; (b) a `before_resume` hook replacing `ResumePreparation.transcript` +with another destination is rejected before either locator runs; (c) +`TurnOptions::session` naming a different first generation is rejected before +locator access; (d) resolver failure leaves the fallback locator untouched; +(e) changing agent, first generation, or destination on a later turn returns +`InvalidSessionState` before locator access; (f) a new generation with the +same first generation and destination is accepted. Preserve the existing +`lazy_target_is_opened_only_after_before_resume_selects_it` test to prove +ordinary hook-selected lazy targets still work when no identity resolver is +installed. Each recording locator +asserts zero calls to `head_generation`, `read_session_transcript`, +`adopt_legacy`, and lock lookup after rejection. Implement +`resolve_and_bind(options, hook_target)` in `session.rs`; turn calls it after +hooks and public `resume` calls it directly with `None`, then both invoke +private `resume_admitted`. +Bind the target locator from the lease, canonical `SessionRef`, and metadata +seed; expose the identity through `SessionStateView`. Require metadata seed at +builder build time if a resolver is installed. Check any explicit option or +hook-selected target against this identity; they cannot replace authority. +Add these test helpers to the top of `lib_tests.rs` and use the exact direct +resume regression below: + +```rust +struct FixedIdentityResolver { + provider: Arc<tinyagents_session::port::DriverSessionStores>, + calls: Arc<Mutex<usize>>, + agent_id: Arc<Mutex<String>>, +} + +impl crate::SessionIdentityResolver<()> for FixedIdentityResolver { + fn resolve( + &self, + request: &crate::SessionIdentityRequest, + _: &(), + ) -> Result<crate::ResolvedSessionStores, RuntimeError> { + if request.parent.is_none() { + assert!(request.session_key.starts_with("thread-")); + assert!(request.child_key.is_none()); + } else { + assert_eq!(request.session_key, "delegate:worker"); + assert_eq!(request.child_key.as_deref(), Some("delegate:worker")); + } + *self.calls.lock().unwrap() += 1; + let agent_id = self.agent_id.lock().unwrap().clone(); + let stores = tinyagents_session::port::SessionStoreProvider::lease_for_agent( + self.provider.as_ref(), &agent_id, + ).map_err(|error| RuntimeError::InvalidSessionState(error.to_string()))?; + Ok(crate::ResolvedSessionStores { agent_id, stores }) + } +} + +#[tokio::test] +async fn direct_resume_resolves_and_uses_the_admitted_lease() { + use tinystoragedrivers_core::{MemoryStorage, Scope}; + let provider = Arc::new( + tinyagents_session::port::DriverSessionStores::new( + Arc::new(MemoryStorage::new()), Scope::local(), + ).unwrap(), + ); + let calls = Arc::new(Mutex::new(0)); + let resolver = Arc::new(FixedIdentityResolver { + provider, calls: calls.clone(), agent_id: Arc::new(Mutex::new("agent-resolved".into())), + }); + let mut session = SessionBuilder::new(Arc::new(Driver::new(vec![]))) + .codec(Arc::new(Codec::default())) + .identity_resolver("thread-identity", meta(), resolver) + .build().unwrap(); + let mut options = TurnOptions::default(); + options.resume = ResumeMode::Session; + assert!(!session.resume(&options).await.unwrap().loaded); + assert_eq!(*calls.lock().unwrap(), 1); +} + +#[tokio::test] +async fn child_identity_resolution_keeps_parent_lineage_and_does_not_adopt_as_root() { + use tinystoragedrivers_core::{MemoryStorage, Scope}; + let provider = Arc::new( + tinyagents_session::port::DriverSessionStores::new( + Arc::new(MemoryStorage::new()), Scope::local(), + ).unwrap(), + ); + let parent = SessionRef::scoped("root-thread", "parent-agent"); + let child = SessionRef::child_of(&parent, "delegate:worker"); + let resolver = Arc::new(FixedIdentityResolver { + provider: provider.clone(), + calls: Arc::new(Mutex::new(0)), + agent_id: Arc::new(Mutex::new("worker-agent".into())), + }); + let driver = Arc::new(Driver::new(vec![Ok(outcome(vec![ + Message::user("child task"), Message::assistant("child result"), + ]))])); + let mut session = SessionBuilder::new(driver) + .codec(Arc::new(Codec::default())) + .child_identity_resolver( + parent.clone(), "delegate:worker", meta(), resolver, + ) + .build().unwrap(); + let options = TurnOptions { + resume: ResumeMode::Session, + thread_id: Some("root-thread".into()), + ..TurnOptions::default() + }; + assert!(!session.resume(&options).await.unwrap().loaded); + session.turn( + Message::user("child task"), + options, + ).await.unwrap(); + + let stores = tinyagents_session::port::SessionStoreProvider::lease_for_agent( + provider.as_ref(), "worker-agent", + ).unwrap(); + let transcripts = &stores.stores().transcripts; + assert_eq!(session_stem(&child), session_stem(&child.first_generation())); + assert!(session_stem(&child).starts_with(&format!("{}__", session_stem(&parent)))); + assert!(transcripts.read_session_transcript(&child).is_some()); + assert!(transcripts.read_session_transcript(&parent).is_none()); +} + +#[tokio::test] +async fn admitted_identity_keeps_compacted_head_on_following_never_resume_turn() { + use tinystoragedrivers_core::{MemoryStorage, Scope}; + let provider = Arc::new(tinyagents_session::port::DriverSessionStores::new( + Arc::new(MemoryStorage::new()), Scope::local(), + ).unwrap()); + let session_ref = SessionRef::scoped("thread-identity-compact", "agent-resolved"); + let resolver = Arc::new(FixedIdentityResolver { + provider: provider.clone(), + calls: Arc::new(Mutex::new(0)), + agent_id: Arc::new(Mutex::new("agent-resolved".into())), + }); + let driver = Arc::new(Driver::new(vec![ + Ok(outcome(vec![ + Message::user("one"), Message::assistant("first"), + Message::user("two"), Message::assistant("second"), + ])), + Ok(outcome(vec![Message::user("three"), Message::assistant("third")])), + Ok(outcome(vec![ + Message::user("three"), Message::assistant("third"), + Message::user("four"), Message::assistant("fourth"), + ])), + ])); + let mut session = SessionBuilder::new(driver) + .codec(Arc::new(Codec::default())) + .identity_resolver("thread-identity-compact", meta(), resolver) + .build().unwrap(); + session.turn( + Message::user("one"), + session_turn_options(ResumeMode::Session, "thread-identity-compact"), + ).await.unwrap(); + let stores = tinyagents_session::port::SessionStoreProvider::lease_for_agent( + provider.as_ref(), "agent-resolved", + ).unwrap(); + let locator = &stores.stores().transcripts; + session.turn( + Message::user("three"), + session_turn_options(ResumeMode::Never, "thread-identity-compact"), + ).await.unwrap(); + + // Capture the generation zero bytes after the compaction has committed. + let sealed = locator.read_session_transcript(&session_ref).unwrap(); + let sealed_bytes = fingerprint(sealed.as_ref()); + assert_eq!(locator.head_generation(&session_ref).generation, 1); + + // This third turn re-admits the identity after compaction. The old bug + // rebuilt generation zero here and attempted to write through the sealed + // handle instead of appending to the known generation-one head. + session.turn( + Message::user("four"), + session_turn_options(ResumeMode::Never, "thread-identity-compact"), + ).await.unwrap(); + + assert_eq!(fingerprint(sealed.as_ref()), sealed_bytes, + "re-admission must not reset or rewrite the sealed generation"); + assert_eq!(locator.head_generation(&session_ref).generation, 1); + let active = locator.read_session_transcript(&session_ref.next_generation()) + .unwrap().read_session().unwrap().unwrap(); + assert_eq!(active.messages.iter().map(|row| row.content.as_str()).collect::<Vec<_>>(), + ["three", "third", "four", "fourth"]); +} + +#[tokio::test] +async fn first_admission_ignores_same_identity_fallback_at_generation_nine() { + use tinystoragedrivers_core::{MemoryStorage, Scope}; + let provider = Arc::new(tinyagents_session::port::DriverSessionStores::new( + Arc::new(MemoryStorage::new()), Scope::local(), + ).unwrap()); + let mut fallback_session = SessionRef::scoped("thread-first-admission", "agent-resolved"); + for _ in 0..9 { + fallback_session = fallback_session.next_generation(); + } + let (fallback, fallback_history) = locator(None); + let resolver = Arc::new(FixedIdentityResolver { + provider: provider.clone(), + calls: Arc::new(Mutex::new(0)), + agent_id: Arc::new(Mutex::new("agent-resolved".into())), + }); + let driver = Arc::new(Driver::new(vec![Ok(outcome(vec![ + Message::user("first"), Message::assistant("admitted root"), + ]))])); + let mut session = SessionBuilder::new(driver) + .codec(Arc::new(Codec::default())) + .session(fallback.clone(), fallback_session.clone(), meta()) + .identity_resolver("thread-first-admission", meta(), resolver) + .build().unwrap(); + session.turn( + Message::user("first"), + session_turn_options(ResumeMode::Never, "thread-first-admission"), + ).await.unwrap(); + + let stores = tinyagents_session::port::SessionStoreProvider::lease_for_agent( + provider.as_ref(), "agent-resolved", + ).unwrap(); + let transcripts = &stores.stores().transcripts; + let admitted_root = SessionRef::scoped("thread-first-admission", "agent-resolved"); + assert!(transcripts.read_session_transcript(&admitted_root).is_some()); + assert!(transcripts.read_session_transcript(&fallback_session).is_none()); + assert_eq!(*fallback_history.opens.lock().unwrap(), 0); + assert!(fallback.latest_agents.lock().unwrap().is_empty()); +} +``` + +Use these exact adversarial test bodies. They reuse `locator(None)`, `meta()`, +and `Driver` from `lib_tests.rs`: + +```rust +struct ReplaceTargetHook { + target: Arc<dyn TranscriptLocator>, + session_override: Option<SessionRef>, +} + +#[async_trait] +impl SessionHooks for ReplaceTargetHook { + async fn before_resume( + &self, _: &mut SessionTurnRequest, options: &mut TurnOptions, + _: SessionStateView<'_>, + ) -> Result<ResumePreparation, RuntimeError> { + options.session = self.session_override.clone(); + Ok(ResumePreparation { transcript: Some(TranscriptTarget::for_session( + self.target.clone(), SessionRef::scoped("thread-identity", "agent-resolved"), meta(), + )) }) + } + async fn before_turn( + &self, _: &mut SessionTurnRequest, _: &mut TurnOptions, + _: SessionStateView<'_>, + ) -> Result<TurnPreparation, RuntimeError> { Ok(TurnPreparation::default()) } + async fn before_commit( + &self, _: &SessionTurnOutcome, _: &TranscriptTurnOptions, + ) -> Result<(), RuntimeError> { Ok(()) } + async fn on_terminal(&self, _: SessionTerminal) -> Result<(), RuntimeError> { Ok(()) } +} + +#[tokio::test] +async fn hook_cannot_replace_the_identity_locator() { + use tinystoragedrivers_core::{MemoryStorage, Scope}; + use tinyagents_session::port::DriverSessionStores; + let provider = Arc::new(DriverSessionStores::new(Arc::new(MemoryStorage::new()), Scope::local()).unwrap()); + let calls = Arc::new(Mutex::new(0)); + let resolver = Arc::new(FixedIdentityResolver { + provider, calls, agent_id: Arc::new(Mutex::new("agent-resolved".into())), + }); + let (wrong, wrong_history) = locator(None); + let driver = Arc::new(Driver::new(vec![])); + let mut session = SessionBuilder::new(driver.clone()) + .codec(Arc::new(Codec::default())) + .hooks(Arc::new(ReplaceTargetHook { + target: wrong.clone(), session_override: None, + })) + .identity_resolver("thread-identity", meta(), resolver) + .build().unwrap(); + let options = TurnOptions { resume: ResumeMode::Session, ..TurnOptions::default() }; + assert!(matches!(session.turn(Message::user("next"), options).await, + Err(RuntimeError::InvalidSessionState(_)))); + assert!(driver.requests.lock().unwrap().is_empty()); + assert_eq!(*wrong_history.opens.lock().unwrap(), 0); + assert!(wrong.latest_agents.lock().unwrap().is_empty()); +} + +#[tokio::test] +async fn hook_cannot_override_the_resolved_session_option() { + use tinystoragedrivers_core::{MemoryStorage, Scope}; + use tinyagents_session::port::DriverSessionStores; + let provider = Arc::new(DriverSessionStores::new( + Arc::new(MemoryStorage::new()), Scope::local(), + ).unwrap()); + let resolver = Arc::new(FixedIdentityResolver { + provider, calls: Arc::new(Mutex::new(0)), + agent_id: Arc::new(Mutex::new("agent-resolved".into())), + }); + let (target_locator, target_history) = locator(None); + let (foreign_locator, foreign_history) = locator(None); + let driver = Arc::new(Driver::new(vec![])); + let mut session = SessionBuilder::new(driver.clone()) + .codec(Arc::new(Codec::default())) + .hooks(Arc::new(OverrideSessionHook( + SessionRef::scoped("foreign-thread", "agent-resolved"), + ))) + .transcript(foreign_locator.clone(), "fallback", meta()) + .identity_resolver("thread-identity", meta(), resolver) + .build().unwrap(); + let options = TurnOptions { resume: ResumeMode::Session, ..TurnOptions::default() }; + assert!(matches!(session.turn(Message::user("next"), options).await, + Err(RuntimeError::InvalidSessionState(_)))); + assert!(driver.requests.lock().unwrap().is_empty()); + assert!(target_locator.latest_agents.lock().unwrap().is_empty()); + assert_eq!(*target_history.opens.lock().unwrap(), 0); + assert!(foreign_locator.latest_agents.lock().unwrap().is_empty()); + assert_eq!(*foreign_history.opens.lock().unwrap(), 0); + assert!(target_locator.latest_agents.lock().unwrap().is_empty()); + assert_eq!(*target_history.opens.lock().unwrap(), 0); +} + +struct OverrideSessionHook(SessionRef); + +#[async_trait] +impl SessionHooks for OverrideSessionHook { + async fn before_resume( + &self, _: &mut SessionTurnRequest, options: &mut TurnOptions, + _: SessionStateView<'_>, + ) -> Result<ResumePreparation, RuntimeError> { + options.session = Some(self.0.clone()); + Ok(ResumePreparation::default()) + } + async fn before_turn( + &self, _: &mut SessionTurnRequest, _: &mut TurnOptions, + _: SessionStateView<'_>, + ) -> Result<TurnPreparation, RuntimeError> { Ok(TurnPreparation::default()) } + async fn before_commit( + &self, _: &SessionTurnOutcome, _: &TranscriptTurnOptions, + ) -> Result<(), RuntimeError> { Ok(()) } + async fn on_terminal(&self, _: SessionTerminal) -> Result<(), RuntimeError> { Ok(()) } +} + +#[tokio::test] +async fn public_resume_rejects_foreign_session_before_locator_access() { + use tinyagents_session::port::DriverSessionStores; + use tinystoragedrivers_core::{MemoryStorage, Scope}; + let provider = Arc::new(DriverSessionStores::new( + Arc::new(MemoryStorage::new()), Scope::local(), + ).unwrap()); + let resolver = Arc::new(FixedIdentityResolver { + provider, calls: Arc::new(Mutex::new(0)), + agent_id: Arc::new(Mutex::new("agent-resolved".into())), + }); + let (foreign, foreign_history) = locator(None); + let mut session = SessionBuilder::new(Arc::new(Driver::new(vec![]))) + .codec(Arc::new(Codec::default())) + .transcript(foreign.clone(), "foreign-fallback", meta()) + .identity_resolver("thread-identity", meta(), resolver) + .build().unwrap(); + let mut options = TurnOptions { resume: ResumeMode::Session, ..TurnOptions::default() }; + options.session = Some(SessionRef::scoped("foreign-thread", "agent-resolved")); + assert!(matches!(session.resume(&options).await, + Err(RuntimeError::InvalidSessionState(_)))); + assert!(foreign.latest_agents.lock().unwrap().is_empty()); + assert_eq!(*foreign_history.opens.lock().unwrap(), 0); +} + +#[tokio::test] +async fn explicit_session_override_cannot_change_resolved_identity() { + use tinyagents_session::port::DriverSessionStores; + use tinystoragedrivers_core::{MemoryStorage, Scope}; + let provider = Arc::new(DriverSessionStores::new(Arc::new(MemoryStorage::new()), Scope::local()).unwrap()); + let calls = Arc::new(Mutex::new(0)); + let resolver = Arc::new(FixedIdentityResolver { + provider, calls: calls.clone(), agent_id: Arc::new(Mutex::new("agent-resolved".into())), + }); + let mut session = SessionBuilder::new(Arc::new(Driver::new(vec![]))) + .codec(Arc::new(Codec::default())) + .identity_resolver("thread-identity", meta(), resolver) + .build().unwrap(); + let options = TurnOptions { + resume: ResumeMode::Session, + session: Some(SessionRef::scoped("other-thread", "agent-resolved")), + ..TurnOptions::default() + }; + assert!(matches!(session.resume(&options).await, + Err(RuntimeError::InvalidSessionState(_)))); + assert_eq!(*calls.lock().unwrap(), 1); +} + +struct FailingIdentityResolver; + +impl SessionIdentityResolver<()> for FailingIdentityResolver { + fn resolve( + &self, + _: &SessionIdentityRequest, + _: &(), + ) -> Result<ResolvedSessionStores, RuntimeError> { + Err(RuntimeError::Driver("actor admission denied".into())) + } +} + +#[tokio::test] +async fn resolver_error_does_not_touch_builder_fallback_locator() { + let (fallback, history) = locator(None); + let mut session = SessionBuilder::new(Arc::new(Driver::new(vec![]))) + .codec(Arc::new(Codec::default())) + .transcript(fallback.clone(), "fallback", meta()) + .identity_resolver("thread-identity", meta(), Arc::new(FailingIdentityResolver)) + .build().unwrap(); + let options = TurnOptions { + resume: ResumeMode::LatestForAgent, + ..TurnOptions::default() + }; + assert!(matches!(session.resume(&options).await, Err(RuntimeError::Driver(_)))); + assert!(fallback.latest_agents.lock().unwrap().is_empty()); + assert_eq!(*history.opens.lock().unwrap(), 0); +} + +#[tokio::test] +async fn later_agent_change_is_rejected_before_second_driver_call() { + use tinystoragedrivers_core::{MemoryStorage, Scope}; + use tinyagents_session::port::DriverSessionStores; + let provider = Arc::new(DriverSessionStores::new(Arc::new(MemoryStorage::new()), Scope::local()).unwrap()); + let calls = Arc::new(Mutex::new(0)); + let agent_id = Arc::new(Mutex::new("agent-one".into())); + let resolver = Arc::new(FixedIdentityResolver { provider, calls, agent_id: agent_id.clone() }); + let driver = Arc::new(Driver::new(vec![ + Ok(outcome(vec![Message::user("one"), Message::assistant("done")])), + Ok(outcome(vec![Message::user("two"), Message::assistant("done")])), + ])); + let mut session = SessionBuilder::new(driver.clone()) + .codec(Arc::new(Codec::default())) + .identity_resolver("thread-identity", meta(), resolver) + .build().unwrap(); + let options = || TurnOptions { resume: ResumeMode::Session, ..TurnOptions::default() }; + session.turn(Message::user("one"), options()).await.unwrap(); + *agent_id.lock().unwrap() = "agent-two".into(); + assert!(matches!(session.turn(Message::user("two"), options()).await, + Err(RuntimeError::InvalidSessionState(_)))); + assert_eq!(driver.requests.lock().unwrap().len(), 1); +} +``` + +The option mismatch must precede the `head_generation` call; the hook mismatch +must precede both locators; the later identity mismatch must precede the +second driver invocation. Add this accepted case using the same immutable +`DriverSessionStores` destination and resolved agent for both admissions. The +second public turn must pass the established-binding guard and reach the +driver: + +```rust +#[tokio::test] +async fn same_admission_can_resume_the_current_generation_twice() { + use tinystoragedrivers_core::{MemoryStorage, Scope}; + use tinyagents_session::port::DriverSessionStores; + let provider = Arc::new(DriverSessionStores::new( + Arc::new(MemoryStorage::new()), Scope::local(), + ).unwrap()); + let resolver = Arc::new(FixedIdentityResolver { + provider, calls: Arc::new(Mutex::new(0)), + agent_id: Arc::new(Mutex::new("agent-resolved".into())), + }); + let driver = Arc::new(Driver::new(vec![ + Ok(outcome(vec![Message::user("one"), Message::assistant("done one")])), + Ok(outcome(vec![Message::user("two"), Message::assistant("done two")])), + ])); + let mut session = SessionBuilder::new(driver.clone()) + .codec(Arc::new(Codec::default())) + .identity_resolver("thread-identity", meta(), resolver) + .build().unwrap(); + let options = || TurnOptions { + resume: ResumeMode::Session, + thread_id: Some("thread-identity".into()), + ..TurnOptions::default() + }; + session.turn(Message::user("one"), options()).await.unwrap(); + session.turn(Message::user("two"), options()).await.unwrap(); + assert_eq!(driver.requests.lock().unwrap().len(), 2); +} +``` + +Run the direct test filter and the runtime suite. +Run `run_tinyagents_test -p tinyagents-runtime direct_resume_resolves_and_uses_the_admitted_lease` +and `run_tinyagents_test -p tinyagents-runtime identity`, then the runtime crate suite. + +### 4. Add explicit pinned store leases without changing live rebind + +**Files:** `crates/tinyagents-session/src/port/{mod.rs,types.rs,drivers/rebound.rs}`, +`port/drivers/rebound_tests.rs`, `port/workspace_tests.rs`, and +`testkit/conformance.rs`. + +In `rebound_tests.rs`, add `a_lease_keeps_its_backend_while_new_opens_follow_rebind`: open a lease over `MemoryStorage`, write a transcript row, rebind to fresh storage, and assert the lease still uses the old backend while a new ordinary open follows rebind. Add `actor_bound_driver_provider_can_open_a_pinned_lease`: call `bind_actor(admitted_backend)`, then open a lease from the returned provider and assert its tenant-scoped stores access the admitted backend. Add `unsupported_provider_cannot_issue_a_pinned_lease` for a provider implementing only `for_agent`; assert `Unsupported` and no attempted handle open. Add two-tenant isolation coverage with distinct `Scope`s. Implement the preceding `SessionStoreLease { stores }` exactly, with private construction and a read-only `stores()` accessor. The `SessionStoreProvider` default returns `Err(Unsupported)`; `DriverSessionStores::lease_for_agent` opens immutable, direct handles, and `RebindableSessionStores::lease_for_agent` delegates to `self.current()` so a live forwarding handle never escapes. Map open errors to `Open(error)`. Do not change the host-admitted `bind_actor` contract; test unsupported providers return `Ok(None)` and do not silently borrow another actor. Retain the existing live-handle rebind test unchanged. Run focused lease tests, `run_tinyagents_test -p tinyagents-session port::drivers::rebound`, then the crate suite. + +Use the existing `rebound_tests.rs` fixture imports and these exact bodies: + +```rust +#[tokio::test] +async fn a_lease_keeps_its_backend_while_new_opens_follow_rebind() { + let registry = SessionStoreRegistry::default(); + let old_backend = Arc::new(MemoryStorage::new()); + let provider = RebindableSessionStores::new_in(®istry, old_backend.clone(), false).unwrap(); + let lease = provider.lease_for_agent("agent-one").unwrap(); + lease.stores().kv.put("prefs", "key", serde_json::json!("leased")) + .await.unwrap(); + let new_backend = Arc::new(MemoryStorage::new()); + assert_eq!(registry.rebind(new_backend.clone()).unwrap(), 1); + assert_eq!(lease.stores().kv.get("prefs", "key").await.unwrap(), + Some(serde_json::json!("leased"))); + let new_lease = provider.lease_for_agent("agent-one").unwrap(); + assert_eq!(new_lease.stores().kv.get("prefs", "key").await.unwrap(), None); + new_lease.stores().kv.put("prefs", "key", serde_json::json!("new")) + .await.unwrap(); + let old = DriverSessionStores::new(old_backend, Scope::local()).unwrap() + .try_for_agent("agent-one").unwrap(); + assert_eq!(old.kv.get("prefs", "key").await.unwrap(), Some(serde_json::json!("leased"))); +} + +#[tokio::test] +async fn actor_bound_driver_provider_can_open_a_pinned_lease() { + let backend: Arc<dyn StorageBackend> = Arc::new(MemoryStorage::new()); + let root = DriverSessionStores::new(Arc::new(MemoryStorage::new()), Scope::tenant("alice").unwrap()).unwrap(); + let actor = root.bind_actor(Arc::clone(&backend)).unwrap().unwrap(); + let lease = actor.lease_for_agent("agent").unwrap(); + lease.stores().kv.put("prefs", "key", serde_json::json!("actor")) + .await.unwrap(); + let direct = DriverSessionStores::new(backend, Scope::tenant("alice").unwrap()).unwrap() + .try_for_agent("agent").unwrap(); + assert_eq!(direct.kv.get("prefs", "key").await.unwrap(), Some(serde_json::json!("actor"))); +} + +#[tokio::test] +async fn actor_bound_leases_keep_the_admitted_tenant_separate() { + let backend: Arc<dyn StorageBackend> = Arc::new(MemoryStorage::new()); + let alice_root = DriverSessionStores::new( + Arc::new(MemoryStorage::new()), Scope::tenant("alice").unwrap(), + ).unwrap(); + let bob_root = DriverSessionStores::new( + Arc::new(MemoryStorage::new()), Scope::tenant("bob").unwrap(), + ).unwrap(); + let alice_provider = alice_root.bind_actor(Arc::clone(&backend)).unwrap().unwrap(); + let bob_provider = bob_root.bind_actor(Arc::clone(&backend)).unwrap().unwrap(); + let alice = alice_provider.lease_for_agent("same-agent").unwrap(); + let bob = bob_provider.lease_for_agent("same-agent").unwrap(); + alice.stores().kv.put("prefs", "key", serde_json::json!("alice")) + .await.unwrap(); + assert_eq!(bob.stores().kv.get("prefs", "key").await.unwrap(), None); + assert_eq!(alice.stores().kv.get("prefs", "key").await.unwrap(), + Some(serde_json::json!("alice"))); +} + +#[tokio::test] +async fn leases_keep_tenant_scope_and_agent_isolation() { + let backend = Arc::new(MemoryStorage::new()); + let alice = DriverSessionStores::new(backend.clone(), Scope::tenant("alice").unwrap()).unwrap(); + let bob = DriverSessionStores::new(backend, Scope::tenant("bob").unwrap()).unwrap(); + let alice_lease = alice.lease_for_agent("agent").unwrap(); + let bob_lease = bob.lease_for_agent("agent").unwrap(); + alice_lease.stores().kv.put("prefs", "key", serde_json::json!("alice")) + .await.unwrap(); + assert_eq!(bob_lease.stores().kv.get("prefs", "key").await.unwrap(), None); +} +``` + +Add this default-provider negative test in `port/mod_tests.rs`: + +```rust +struct NonPinningProvider { + stores: InMemorySessionStores, + opens: Arc<std::sync::atomic::AtomicUsize>, +} + +impl SessionStoreProvider for NonPinningProvider { + fn for_agent(&self, agent_id: &str) -> AgentStores { + self.opens.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + self.stores.for_agent(agent_id) + } +} + +#[test] +fn providers_without_a_pin_contract_fail_closed() { + let opens = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let provider = NonPinningProvider { + stores: InMemorySessionStores::new(), + opens: opens.clone(), + }; + assert!(matches!(provider.lease_for_agent("agent"), + Err(SessionStoreLeaseError::Unsupported))); + assert_eq!(opens.load(std::sync::atomic::Ordering::SeqCst), 0); +} +``` + +Keep `retained_and_cloned_handles_follow_rebind_without_crossing_agent_scopes` +unchanged; the new test proves leases pin while that test proves ordinary +handles continue to follow rebind. + +### 5. Capture tool observations and route post-commit sinks in the owned finalizer + +**Files:** `crates/tinyagents-runtime/src/{driver.rs,hooks.rs,session.rs,types.rs}`, +`turn_driver/driver.rs`, and their sibling `driver_tests.rs`, +`hooks_tests.rs`, `session_tests.rs`, `turn_driver/driver_tests.rs`. + +First add a recording `SessionDriver` returning one `ToolCallOutcome`; assert +the returned `DriverOutcome.tool_outcomes` equals it. Then extend the existing +`a_split_turn_is_driven_by_the_host_and_settled_by_its_tail` test: after +`tail.finish(Some(&run))`, assert `report.tool_outcomes.len() == 1`, its name +is `ask_user`, and its content is the result after `RewriteResult` middleware. +Add a paired `prepared_run_exposes_the_same_rewritten_tool_outcome` test that +uses the same scripted tool/model setup but calls `prepared.run().await`; assert +the same three fields on `TurnRun.report.tool_outcomes`. Put a +`RewriteResult` middleware after capture in the planner's `TurnPlan.context` +and assert both public drive paths observe the rewritten text `[capped]`. The +capture is registered first/outermost; the +harness runs `after_tool` in reverse order, so capture observes the result +after each planner rewriter. Then add a recording +sink test that reads the appended transcript row inside the callback, a +cancellation test that drops the caller future after durable commit and +awaits a oneshot sent by the sink, a pre-commit failure test asserting no +callback, and a sink-error test asserting the turn result remains successful. +Each test must fail before the corresponding API/code exists. `DriverOutcome` derives +`PartialEq`, so add `PartialEq` to `ToolCallOutcome`'s existing `Debug, Clone` +derive and preserve the `DriverOutcome` derive and its existing assertions. +Add a value test comparing call id, name, structured arguments, success, +content, and duration. Extend `DriverOutcome` +with `Vec<ToolCallOutcome>` defaulted empty; have `TurnDriver` build a shared +`ToolOutcomeSink` during `prepare` and carry it through `PreparedTurn` into +`TurnTail`, which drains it for both `run` and `split`. Keep custom drivers' +empty default. Add the sink list to `SessionBuilder`, pass it to `Session`, and +invoke each sink in registration order from the existing owned +`terminal_guard.finalize_commit(receipt, observation, sinks)` task. Keep `SessionHooks::after_commit` +in that task and in its current position. The `TurnSummary` adapter consumes +outcomes and must not persist the raw argument/result. Verify with +`run_tinyagents_test -p tinyagents-runtime post_commit` and then +`run_tinyagents_test -p tinyagents-runtime`. + +The value test body in `turn_driver/driver_tests.rs` is: + +```rust +#[test] +fn tool_call_outcome_equality_covers_every_observed_field() { + use tinyagents_harness::turn_outcome::ToolCallOutcome; + + let original = ToolCallOutcome { + call_id: "call-1".into(), + name: "read_file".into(), + arguments: serde_json::json!({"path": "notes.txt"}), + success: false, + content: "denied".into(), + duration_ms: 17, + }; + assert_eq!(original, original.clone()); + + let mut changed = original.clone(); + changed.call_id.push_str("-changed"); + assert_ne!(original, changed); + let mut changed = original.clone(); + changed.name.push_str("_changed"); + assert_ne!(original, changed); + let mut changed = original.clone(); + changed.arguments["path"] = serde_json::json!("other.txt"); + assert_ne!(original, changed); + let mut changed = original.clone(); + changed.success = true; + assert_ne!(original, changed); + let mut changed = original.clone(); + changed.content.push_str(" changed"); + assert_ne!(original, changed); + let mut changed = original.clone(); + changed.duration_ms += 1; + assert_ne!(original, changed); +} +``` + +In `a_split_turn_is_driven_by_the_host_and_settled_by_its_tail`, add this +planner context middleware after its tool builder setup: + +```rust +plan.context.push(Arc::new(Trace { + name: "rewrite", + log: Arc::new(Mutex::new(Vec::new())), + rewrite_result: true, +})); +``` + +After `tail.finish(Some(&run))`, assert: + +```rust +assert_eq!(report.tool_outcomes.len(), 1); +assert_eq!(report.tool_outcomes[0].name, "ask_user"); +assert_eq!(report.tool_outcomes[0].content, "[capped]"); +``` + +Add this paired run-path test body in `turn_driver/driver_tests.rs`: + +```rust +#[tokio::test] +async fn prepared_run_exposes_the_same_rewritten_tool_outcome() { + let model = Arc::new(ScriptedModel::new(vec![ + tool_call_response(call("c1", "ask_user", 1)), + text_response("never requested"), + ])); + let tools: Vec<Box<dyn Tool>> = vec![Box::new(FakeTool::returning("ask_user", "question"))]; + let snapshot = snapshot_of(&tools, &["ask_user"]); + let set: SharedToolSet = Arc::new(tools); + let planner = { + let model = model.clone(); + move |_: &DriverRequest<()>| { + let chat: Arc<dyn ChatModel<()>> = model.clone(); + let mut plan = TurnPlan::new(vec![("scripted".to_string(), chat)]); + plan.tools = ToolSetBuilder::new().set(set.clone()); + plan.context.push(Arc::new(Trace { + name: "rewrite", + log: Arc::new(Mutex::new(Vec::new())), + rewrite_result: true, + })); + Ok(plan) + } + }; + let driver = TurnDriver::new(Arc::new(planner), Arc::new(())); + let request = DriverRequest { + history: vec![Message::user("go")], + tools: snapshot, + run_context: RunContext::new( + tinyagents_harness::context::RunConfig::new("prepared-run"), (), + ), + stream: false, + controls: TurnControls::default(), + }; + let run = driver.prepare(request).await.unwrap().run().await; + assert_eq!(run.report.tool_outcomes.len(), 1); + assert_eq!(run.report.tool_outcomes[0].name, "ask_user"); + assert_eq!(run.report.tool_outcomes[0].content, "[capped]"); +} +``` + +Use this sink and test in `session_tests.rs` to prove the post-commit task owns +completion after the caller task is aborted. The callback reads the just- +appended transcript before it signals the test: + +```rust +struct BlockingReadSink { + locator: Arc<dyn TranscriptLocator>, + session: SessionRef, + started: Mutex<Option<tokio::sync::oneshot::Sender<()>>>, + release: Arc<tokio::sync::Notify>, + finished: Arc<tokio::sync::Notify>, +} + +#[async_trait] +impl PostCommitSink for BlockingReadSink { + async fn after_commit( + &self, + observation: &PostCommitObservation, + ) -> Result<(), RuntimeError> { + assert!(observation.receipt.transcript.is_some()); + let head = self.locator.head_generation(&self.session); + let read = self.locator.read_session_transcript(&head).unwrap(); + let transcript = read.read_session().unwrap().unwrap(); + assert!(transcript.messages.iter().any(|row| row.content == "committed reply")); + if let Some(started) = self.started.lock().unwrap().take() { + let _ = started.send(()); + } + self.release.notified().await; + self.finished.notify_one(); + Ok(()) + } +} + +struct ReplyDriver; + +#[async_trait] +impl SessionDriver for ReplyDriver { + async fn execute(&self, request: DriverRequest) -> Result<DriverOutcome, DriverFailure> { + let mut history = request.history; + history.push(Message::assistant("committed reply")); + Ok(DriverOutcome { + history, + output: Some("committed reply".into()), + partial: None, + interrupted: false, + outcome: None, + tool_outcomes: Vec::new(), + }) + } +} + +#[tokio::test] +async fn owned_post_commit_sink_finishes_after_caller_abort() { + let locator = memory(); + let session_ref = SessionRef::scoped("thread-post-commit", "agent"); + let release = Arc::new(tokio::sync::Notify::new()); + let finished = Arc::new(tokio::sync::Notify::new()); + let (started_tx, started_rx) = tokio::sync::oneshot::channel(); + let sink = Arc::new(BlockingReadSink { + locator: locator.clone(), + session: session_ref.clone(), + started: Mutex::new(Some(started_tx)), + release: release.clone(), + finished: finished.clone(), + }); + let session = SessionBuilder::new(Arc::new(ReplyDriver)) + .codec(Arc::new(RoleCodec)) + .session(locator, session_ref, meta()) + .post_commit_sink(sink) + .build().unwrap(); + let turn = tokio::spawn(async move { + let mut session = session; + session.turn(Message::user("commit this"), TurnOptions::default()).await + }); + started_rx.await.unwrap(); + turn.abort(); + release.notify_one(); + timeout(Duration::from_secs(1), finished.notified()).await.unwrap(); +} +``` + +Add this failure-partial regression in the same test module. It documents that +the host sink is successful-turn-only while preserving existing durable +partial-history and terminal behavior: + +```rust +struct CountSink(Arc<std::sync::atomic::AtomicUsize>); + +#[async_trait] +impl PostCommitSink for CountSink { + async fn after_commit(&self, _: &PostCommitObservation) -> Result<(), RuntimeError> { + self.0.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + Ok(()) + } +} + +struct FailureTerminalRecorder(Arc<Mutex<Vec<SessionTerminal>>>); + +#[async_trait] +impl SessionHooks for FailureTerminalRecorder { + async fn before_turn( + &self, _: &mut SessionTurnRequest, _: &mut TurnOptions, + _: SessionStateView<'_>, + ) -> Result<TurnPreparation, RuntimeError> { + Ok(TurnPreparation::default()) + } + async fn before_commit( + &self, _: &SessionTurnOutcome, _: &TranscriptTurnOptions, + ) -> Result<(), RuntimeError> { Ok(()) } + async fn on_terminal(&self, terminal: SessionTerminal) -> Result<(), RuntimeError> { + self.0.lock().unwrap().push(terminal); + Ok(()) + } +} + +struct FailingPartialDriver; + +#[async_trait] +impl SessionDriver for FailingPartialDriver { + async fn execute(&self, request: DriverRequest) -> Result<DriverOutcome, DriverFailure> { + let mut history = request.history; + history.push(Message::assistant("durable partial")); + Err(DriverFailure { + error: RuntimeError::Driver("provider failed".into()), + partial: Some(DriverOutcome { + history, + output: Some("durable partial".into()), + partial: None, + interrupted: false, + outcome: None, + tool_outcomes: Vec::new(), + }), + outcome: None, + }) + } +} + +#[tokio::test] +async fn persisted_failure_partial_skips_post_commit_sinks_and_fails_once() { + let locator = memory(); + let session_ref = SessionRef::scoped("thread-failure-partial", "agent"); + let sink_calls = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let terminals = Arc::new(Mutex::new(Vec::new())); + let mut session = SessionBuilder::new(Arc::new(FailingPartialDriver)) + .codec(Arc::new(RoleCodec)) + .hooks(Arc::new(FailureTerminalRecorder(terminals.clone()))) + .post_commit_sink(Arc::new(CountSink(sink_calls.clone()))) + .session(locator.clone(), session_ref.clone(), meta()) + .build().unwrap(); + assert!(session.turn(Message::user("fail"), TurnOptions::default()).await.is_err()); + assert_eq!(sink_calls.load(std::sync::atomic::Ordering::SeqCst), 0); + let terminals = terminals.lock().unwrap(); + assert_eq!(terminals.len(), 1); + assert!(matches!(terminals[0], SessionTerminal::Failed(_))); + let head = locator.head_generation(&session_ref); + let transcript = locator.read_session_transcript(&head).unwrap() + .read_session().unwrap().unwrap(); + assert!(transcript.messages.iter().any(|row| row.content == "durable partial")); +} + +struct RejectBeforeCommit; + +#[async_trait] +impl SessionHooks for RejectBeforeCommit { + async fn before_turn( + &self, _: &mut SessionTurnRequest, _: &mut TurnOptions, + _: SessionStateView<'_>, + ) -> Result<TurnPreparation, RuntimeError> { Ok(TurnPreparation::default()) } + async fn before_commit( + &self, _: &SessionTurnOutcome, _: &TranscriptTurnOptions, + ) -> Result<(), RuntimeError> { + Err(RuntimeError::Driver("reject before durable append".into())) + } + async fn on_terminal(&self, _: SessionTerminal) -> Result<(), RuntimeError> { Ok(()) } +} + +#[tokio::test] +async fn precommit_error_does_not_call_post_commit_sink() { + let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let mut session = SessionBuilder::new(Arc::new(ReplyDriver)) + .hooks(Arc::new(RejectBeforeCommit)) + .post_commit_sink(Arc::new(CountSink(calls.clone()))) + .build().unwrap(); + assert!(session.turn(Message::user("reject"), TurnOptions::default()).await.is_err()); + assert_eq!(calls.load(std::sync::atomic::Ordering::SeqCst), 0); +} + +struct InterruptedDriver; + +#[async_trait] +impl SessionDriver for InterruptedDriver { + async fn execute(&self, request: DriverRequest) -> Result<DriverOutcome, DriverFailure> { + Ok(DriverOutcome { + history: request.history, + output: None, + partial: None, + interrupted: true, + outcome: None, + tool_outcomes: Vec::new(), + }) + } +} + +#[tokio::test] +async fn normal_interrupted_commit_still_invokes_post_commit_sink() { + let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let mut session = SessionBuilder::new(Arc::new(InterruptedDriver)) + .post_commit_sink(Arc::new(CountSink(calls.clone()))) + .build().unwrap(); + let outcome = session.turn(Message::user("stop"), TurnOptions::default()).await.unwrap(); + assert!(outcome.interrupted); + assert_eq!(calls.load(std::sync::atomic::Ordering::SeqCst), 1); +} + +struct ErrorSink; + +#[async_trait] +impl PostCommitSink for ErrorSink { + async fn after_commit(&self, _: &PostCommitObservation) -> Result<(), RuntimeError> { + Err(RuntimeError::Driver("sink failed after commit".into())) + } +} + +#[tokio::test] +async fn post_commit_sink_error_does_not_change_successful_turn_result() { + let locator = memory(); + let session_ref = SessionRef::scoped("thread-sink-error", "agent"); + let mut session = SessionBuilder::new(Arc::new(ReplyDriver)) + .codec(Arc::new(RoleCodec)) + .post_commit_sink(Arc::new(ErrorSink)) + .session(locator, session_ref, meta()) + .build().unwrap(); + let outcome = session.turn(Message::user("succeed"), TurnOptions::default()).await.unwrap(); + assert_eq!(outcome.output.as_deref(), Some("committed reply")); +} +``` + +Capture must live in `TurnDriver::prepare`, not only `execute`: callers may +drive `PreparedTurn::split` themselves and never call `execute`. Create the +sink while preparing, push capture middleware before planner middleware so it +is outermost and observes rewritten final results, and carry the sink in +`PreparedTurn` into `TurnTail`. `TurnTail::finish` drains it into a new +`TurnReport.tool_outcomes` field; `PreparedTurn::run` already settles through +`TurnTail::finish`, and split callers do the same. `TurnDriver::execute` copies +the report's outcomes into `DriverOutcome`; `outcome_from_partial` initializes +custom-driver outcomes to an empty vector. This gives direct `prepare/run`, +`prepare/split`, and `execute` the same captured result. + +Make these exact structural edits in `turn_driver/types.rs` and +`turn_driver/driver.rs`: + +```rust +// TurnTail +pub outcomes: tinyagents_harness::turn_outcome::ToolOutcomeSink, + +// TurnReport +pub tool_outcomes: Vec<tinyagents_harness::turn_outcome::ToolCallOutcome>, + +// In TurnDriver::prepare, before pushing planner middleware: +let outcome_sink = tinyagents_harness::turn_outcome::new_tool_outcome_sink(); +harness.push_middleware(Arc::new( + tinyagents_harness::middleware::library::ToolOutcomeCaptureMiddleware::new( + Arc::clone(&outcome_sink), + ), +)); + +// In the existing TurnTail initializer: +outcomes: Arc::clone(&outcome_sink), + +// In TurnTail::finish, before constructing TurnReport: +let tool_outcomes = + tinyagents_harness::turn_outcome::drain_tool_outcomes(&self.outcomes); +// In the TurnReport literal: +tool_outcomes, +``` + +The exact declarations and report construction are: + +```rust +pub struct TurnTail { + guard: Option<FeedGuard>, + hook: Option<EarlyExitHook>, + controls: TurnControls, + max_iterations: usize, + tools: ToolSetReport, + completion: Option<Arc<dyn TurnCompletionSink>>, + outcomes: tinyagents_harness::turn_outcome::ToolOutcomeSink, +} + +// At the beginning of TurnTail::finish, before moving other fields: +let tool_outcomes = + tinyagents_harness::turn_outcome::drain_tool_outcomes(&self.outcomes); + +// Include this field in the existing TurnReport literal: +tool_outcomes, +``` + +`TurnReport` receives +`pub tool_outcomes: Vec<tinyagents_harness::turn_outcome::ToolCallOutcome>`; +its manual `Debug` implementation receives +`.field("tool_outcomes", &self.tool_outcomes)`. `TurnReport::default()` keeps +failed-before-run reports empty. Keep `PreparedTurn`'s existing `tail` field: +that owns the sink for both `run` and `split`. + +In `TurnDriver::prepare`, create `outcome_sink`, push +`Arc::new(tinyagents_harness::middleware::library::ToolOutcomeCaptureMiddleware::new(Arc::clone(&outcome_sink)))` into +the harness before the `for middleware in context` loop, and place the sink in +`TurnTail`. In `TurnTail::finish`, call `drain_tool_outcomes(&self.outcomes)` +and put the result into `TurnReport`. Add `tool_outcomes: Vec<ToolCallOutcome>` +to `TurnReport`; add `tool_outcomes: Vec<ToolCallOutcome>` to +`DriverOutcome` and initialize it to `Vec::new()` in `outcome_from_partial` +and each custom-driver literal. In `TurnDriver::execute`, preserve observations +on both result branches: + +```rust +let TurnRun { partial, report, input_len } = prepared.run().await; +let mut outcome = outcome_from_partial(partial, input_len); +let tool_outcomes = report.tool_outcomes.clone(); +match &mut outcome { + Ok(done) => done.tool_outcomes = tool_outcomes, + Err(failure) => { + if let Some(partial) = &mut failure.partial { + partial.tool_outcomes = tool_outcomes; + } + } +} +// Keep the existing early-exit output rewrite and completion emission here. +if let (Some(exit), Ok(done)) = (&report.early_exit, &mut outcome) { + done.output = Some(exit.question.clone()); +} +if outcome.is_ok() && !defer && !scope.is_subagent() { + report.emit_completed(); +} +outcome +``` + +Initialize the new `DriverOutcome.tool_outcomes` field to `Vec::new()` in +`outcome_from_partial` and in every custom-driver literal. Keep +`ToolCallOutcome`'s new `PartialEq` derive so the preserved `DriverOutcome` +`PartialEq` derive and existing equality assertions remain intact. + +Change the owned finalizer to accept the separate observation and sinks, not +to reconstruct observations from `CommitReceipt`: + +```rust +fn finalize_commit( + &mut self, + receipt: CommitReceipt<C>, + observation: PostCommitObservation<C>, + sinks: Vec<Arc<dyn PostCommitSink<C>>>, +) -> tokio::task::JoinHandle<()> { + let terminal = SessionTerminal::Completed(receipt.outcome.clone()); + let outcome = self.driver_success_outcome.take() + .unwrap_or_else(|| terminal.outcome()); + self.terminal = None; + self.committed = true; + let hooks = self.hooks.clone(); + tokio::spawn(async move { + let _ = hooks.after_commit(receipt.clone()).await; + for sink in sinks { + let _ = sink.after_commit(&observation).await; + } + let _ = hooks.on_terminal_outcome(outcome).await; + let _ = hooks.on_terminal(terminal).await; + }) +} +``` + +Keep the existing `after_commit` call at the same place in the task. Build +`PostCommitObservation` from the committed receipt, the original request +input, usage captured from the completed turn, `DriverOutcome.tool_outcomes`, +and the resolved thread/agent identity before transferring ownership to the +spawned task. At the start of `turn_inner`, before `before_resume` can mutate +the request, capture `let original_input = request.input.clone();`. At the +existing successful commit site, construct and pass the +values in this order: + +```rust +let observation = PostCommitObservation { + receipt: receipt.clone(), + input: original_input, + usage: turn_usage.clone(), + tool_outcomes: outcome.tool_outcomes.clone(), + thread_id: thread_id.clone(), + agent_id: self.identity.as_ref().map(|identity| identity.agent_id().to_owned()), +}; +let finalization = terminal_guard.finalize_commit( + receipt, observation, self.post_commit_sinks.clone(), +); +let _ = finalization.await; +``` + +Perform no await between successful append and this call. `PostCommitSink` is +for successful durable turn commits, including normal interrupted outcomes, +matching the existing `SessionHooks::after_commit` and host committed-turn +finalizer. A `DriverFailure.partial` may be persisted for history recovery, +but it does not create a successful `CommitReceipt` and must not invoke these +sinks. Keep the existing failed-terminal finalizer behavior. Add tests proving +a persisted failure partial invokes no sink and emits one `Failed` terminal, +pre-commit failure invokes no sink, and a normal interrupted durable commit +does invoke sinks. + +### 6. Move tool policy and declaration preservation + +**Files:** `crates/tinyagents-runtime/src/{tool_registration.rs,managed_tools.rs}`, +`crates/tinyagents-harness/src/tool/{snapshot.rs,mod.rs}`, plus sibling +`tool_registration_tests.rs`, `managed_tools_tests.rs`, and +`snapshot_tests.rs`. + +In `tool_registration_tests.rs`, add tests that run the wrapper's real +`ToolCallPolicy::check`: unknown name denies without calling the host gate; +allowed name delegates exactly once; empty set denies all. Implement +`HostOnlyToolPolicy<C>` as `ToolCallPolicy<C>` with +`RunContext<C>`, `ToolCall`, and `PolicyDecision`, then expose a helper to +compose the wrapper around an optional host policy. Add first-wins spec test +that checks names/order and `Arc::ptr_eq` with the winning allocation; add +durable-wins-synthesized collision test. In `snapshot_tests.rs`, assert old +array decode and byte-compatible `to_json` without extensions, typed +retrieval and replacement using two host-defined fixture types, clone +preservation, extension preservation through `retaining` and `with_retained`, +and exact snapshots never retaining recorded tools. Do not add object-envelope +decoding or extension serialization. Run each new filter and verify red before +code, then run +`run_tinyagents_test -p tinyagents-runtime -p tinyagents-harness`. + +The typed-carrier regression test is concrete and host-neutral: + +```rust +#[derive(Clone, Debug, PartialEq, Eq)] +struct HostAlpha(u8); +#[derive(Clone, Debug, PartialEq, Eq)] +struct HostBeta(&'static str); + +#[test] +fn typed_extensions_survive_snapshot_copies_without_changing_json() { + let original_json = serde_json::json!([]); + let snapshot = ToolSnapshot::from_json(&original_json).unwrap() + .with_extension(HostAlpha(7)) + .with_extension(HostBeta("host-owned")); + assert_eq!(snapshot.extension::<HostAlpha>(), Some(&HostAlpha(7))); + assert_eq!(snapshot.extension::<HostBeta>(), Some(&HostBeta("host-owned"))); + assert_eq!(snapshot.to_json(), original_json); + + let clone = snapshot.clone(); + let narrowed = snapshot.retaining(|_| true); + let (retained, count) = snapshot.with_retained(&ToolSnapshot::default()).unwrap(); + assert_eq!(count, 0); + for copy in [&clone, &narrowed, &retained] { + assert_eq!(copy.extension::<HostAlpha>(), Some(&HostAlpha(7))); + assert_eq!(copy.extension::<HostBeta>(), Some(&HostBeta("host-owned"))); + assert_eq!(copy.to_json(), original_json); + } + let replaced = clone.with_extension(HostAlpha(8)); + assert_eq!(replaced.extension::<HostAlpha>(), Some(&HostAlpha(8))); + assert_eq!(snapshot.extension::<HostAlpha>(), Some(&HostAlpha(7))); +} +``` + +Add these policy tests to `tool_registration_tests.rs`. They call the exact +runtime policy trait directly and prove both deny-before-host and delegation: + +```rust +struct CountingPolicy(Arc<std::sync::atomic::AtomicUsize>); + +#[async_trait::async_trait] +impl tinyagents_harness::middleware::ToolCallPolicy<()> for CountingPolicy { + fn name(&self) -> &str { "counting-host" } + + async fn check( + &self, + _: &tinyagents_harness::context::RunContext<()>, + _: &tinyinference_llm::tool::ToolCall, + ) -> tinyagents_harness::middleware::PolicyDecision { + self.0.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + tinyagents_harness::middleware::PolicyDecision::Allow + } +} + +fn policy_call(name: &str) -> tinyinference_llm::tool::ToolCall { + tinyinference_llm::tool::ToolCall { + id: "call-1".into(), + name: name.into(), + arguments: serde_json::json!({}), + invalid: None, + } +} + +#[tokio::test] +async fn host_only_policy_denies_before_delegating_and_allows_named_calls() { + use tinyagents_harness::middleware::{PolicyDecision, ToolCallPolicy}; + let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let policy = HostOnlyToolPolicy::new( + ["read_file".to_owned()], + Some(Arc::new(CountingPolicy(calls.clone()))), + ); + let context = tinyagents_harness::context::RunContext::new( + tinyagents_harness::context::RunConfig::new("host-only-test"), + (), + ); + let denied = policy.check(&context, &policy_call("write_file")).await; + assert!(matches!(denied, PolicyDecision::Deny { .. })); + assert_eq!(calls.load(std::sync::atomic::Ordering::SeqCst), 0); + + assert_eq!(policy.check(&context, &policy_call("read_file")).await, + PolicyDecision::Allow); + assert_eq!(calls.load(std::sync::atomic::Ordering::SeqCst), 1); + + let empty = HostOnlyToolPolicy::new(Vec::<String>::new(), None); + assert!(matches!(empty.check(&context, &policy_call("read_file")).await, + PolicyDecision::Deny { .. })); +} +``` + +### 7. Move context configuration into harness + +**Files:** `crates/tinyagents-harness/src/config/context.rs`, +`config/mod.rs`, `src/lib.rs`, and new `config/context_tests.rs` declared +through `#[cfg(test)] #[path = "context_tests.rs"] mod tests;`. Update +TinyAgents compaction consumers and config exports; do not change OpenHuman +sources in this TinyAgents plan. + +Add the sibling `config/context_tests.rs` first. The test body is: + +```rust +use super::*; + +#[test] +fn context_config_preserves_defaults_aliases_and_strategy_spellings() { + let default = ContextConfig::default(); + let parsed: ContextConfig = toml::from_str("").unwrap(); + assert_eq!(serde_json::to_value(parsed).unwrap(), serde_json::to_value(default).unwrap()); + + let canonical = r#" +enabled = false +microcompact_enabled = false +autocompact_enabled = true +microcompact_keep_recent = 9 +tool_result_budget_bytes = 1234 +summarizer_payload_threshold_tokens = 4000 +summarizer_max_payload_tokens = 64000 +summarizer_model = "test-model" +prefer_markdown_tool_output = false +compaction_enabled = true +compaction_trigger_tokens = 123 +compaction_strategy = "summary" +"#; + let value: ContextConfig = toml::from_str(canonical).unwrap(); + assert!(!value.enabled); + assert!(!value.microcompact_enabled); + assert!(value.autocompact_enabled); + assert_eq!(value.microcompact_keep_recent, 9); + assert_eq!(value.tool_result_budget_bytes, 1234); + assert_eq!(value.summarizer_payload_threshold_tokens, 4000); + assert_eq!(value.summarizer_max_payload_tokens, 64000); + assert_eq!(value.summarizer_model.as_deref(), Some("test-model")); + assert!(!value.prefer_markdown_tool_output); + assert!(value.compaction_enabled); + assert_eq!(value.compaction_trigger_tokens, Some(123)); + assert_eq!(value.compaction_settings().trigger_tokens, Some(123)); + assert_eq!(value.compaction_settings().strategy, CompactionStrategy::Summary); + + let aliases: ContextConfig = toml::from_str( + "summarizer_payload_threshold_bytes=41\nsummarizer_max_payload_bytes=42", + ).unwrap(); + assert_eq!(aliases.summarizer_payload_threshold_tokens, 41); + assert_eq!(aliases.summarizer_max_payload_tokens, 42); + + for spelling in ["task_state", "typed", "s4"] { + assert_eq!(CompactionStrategy::parse(spelling), Some(CompactionStrategy::TaskState)); + } + for spelling in ["summary", "free_form", "s0"] { + assert_eq!(CompactionStrategy::parse(spelling), Some(CompactionStrategy::Summary)); + } + assert_eq!(CompactionStrategy::parse(" TASK-STATE "), Some(CompactionStrategy::TaskState)); + assert_eq!(CompactionStrategy::parse("unknown"), None); + let mut zero = value.clone(); + zero.compaction_trigger_tokens = Some(0); + assert_eq!(zero.compaction_settings().trigger_tokens, None); + let encoded = toml::to_string(&default).unwrap(); + let round_trip: ContextConfig = toml::from_str(&encoded).unwrap(); + assert_eq!(serde_json::to_value(round_trip).unwrap(), serde_json::to_value(default).unwrap()); +} +``` + +Then move the full type definitions below into `config/context.rs`, preserving +the 12 fields and serde attributes exactly. No host `JsonSchema` derive or +OpenHuman imports belong in the harness. The tests fail to compile before the +move because the harness has no `ContextConfig`; run the focused test after +each migration substep. + +```rust +use serde::{Deserialize, Serialize}; + +pub const DEFAULT_KEEP_RECENT_TOOL_RESULTS: usize = 5; + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(default)] +pub struct ContextConfig { + #[serde(default = "default_enabled")] + pub enabled: bool, + #[serde(default = "default_true")] + pub microcompact_enabled: bool, + #[serde(default = "default_true")] + pub autocompact_enabled: bool, + #[serde(default = "default_microcompact_keep_recent")] + pub microcompact_keep_recent: usize, + #[serde(default = "default_tool_result_budget_bytes")] + pub tool_result_budget_bytes: usize, + #[serde(default = "default_summarizer_payload_threshold_tokens", alias = "summarizer_payload_threshold_bytes")] + pub summarizer_payload_threshold_tokens: usize, + #[serde(default = "default_summarizer_max_payload_tokens", alias = "summarizer_max_payload_bytes")] + pub summarizer_max_payload_tokens: usize, + #[serde(default)] + pub summarizer_model: Option<String>, + #[serde(default = "default_true")] + pub prefer_markdown_tool_output: bool, + #[serde(default = "default_true")] + pub compaction_enabled: bool, + #[serde(default)] + pub compaction_trigger_tokens: Option<u64>, + #[serde(default)] + pub compaction_strategy: CompactionStrategy, +} + +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub struct CompactionSettings { + pub trigger_tokens: Option<u64>, + pub strategy: CompactionStrategy, +} + +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum CompactionStrategy { + #[default] + TaskState, + Summary, +} + +impl ContextConfig { + pub fn compaction_settings(&self) -> CompactionSettings { + CompactionSettings { + trigger_tokens: self.compaction_trigger_tokens.filter(|value| *value > 0), + strategy: self.compaction_strategy, + } + } +} + +impl CompactionStrategy { + pub fn parse(raw: &str) -> Option<Self> { + match raw.trim().to_ascii_lowercase().replace('-', "_").as_str() { + "task_state" | "typed" | "s4" => Some(Self::TaskState), + "summary" | "free_form" | "s0" => Some(Self::Summary), + _ => None, + } + } +} + +fn default_enabled() -> bool { true } +fn default_true() -> bool { true } +fn default_microcompact_keep_recent() -> usize { DEFAULT_KEEP_RECENT_TOOL_RESULTS } +fn default_tool_result_budget_bytes() -> usize { + crate::config::TurnConfig::default().tool_result_budget_bytes +} +fn default_summarizer_payload_threshold_tokens() -> usize { 4_000 } +fn default_summarizer_max_payload_tokens() -> usize { 64_000 } + +impl Default for ContextConfig { + fn default() -> Self { + Self { + enabled: default_enabled(), + microcompact_enabled: default_true(), + autocompact_enabled: default_true(), + microcompact_keep_recent: default_microcompact_keep_recent(), + tool_result_budget_bytes: default_tool_result_budget_bytes(), + summarizer_payload_threshold_tokens: default_summarizer_payload_threshold_tokens(), + summarizer_max_payload_tokens: default_summarizer_max_payload_tokens(), + summarizer_model: None, + prefer_markdown_tool_output: default_true(), + compaction_enabled: default_true(), + compaction_trigger_tokens: None, + compaction_strategy: CompactionStrategy::default(), + } + } +} +``` + +Export all three types and the default constant from `config/mod.rs` and the +harness crate root. Keep the `ContextConfig` field docs describing only the +generic behavior; host-only product policy and environment loading stay in +OpenHuman. + +### 8. Reconcile full crate and workspace integration + +Update remaining TinyAgents call sites to the new upstream-owned APIs without +adding a runtime-to-host or harness-to-runtime dependency. Re-run all +previously affected tests from the audit, including tenant routing, recorded +tool resume, host-only tool gating, transcript adoption, workspace isolation, +and post-commit behavior. Do not delete a test merely because an assertion +currently disagrees; map it to the explicit live-handle or lease contract +above. + +Run the final gates from the TinyAgents workspace root through the +super-worktree cancel-aware wrapper, with one job and debug info disabled: + +```sh +cd -- "/home/enamakel/work/openhuman/worktrees/epic-8088-63-session-runtime/vendor/tinyagents" && CARGO_BUILD_JOBS=1 CARGO_PROFILE_DEV_DEBUG=0 "/home/enamakel/work/openhuman/worktrees/epic-8088-63-session-runtime/scripts/ci-cancel-aware.sh" cargo check --workspace --all-targets +cd -- "/home/enamakel/work/openhuman/worktrees/epic-8088-63-session-runtime/vendor/tinyagents" && CARGO_BUILD_JOBS=1 CARGO_PROFILE_DEV_DEBUG=0 "/home/enamakel/work/openhuman/worktrees/epic-8088-63-session-runtime/scripts/ci-cancel-aware.sh" cargo clippy --workspace --all-targets -- -Dwarnings +cd -- "/home/enamakel/work/openhuman/worktrees/epic-8088-63-session-runtime/vendor/tinyagents" && CARGO_BUILD_JOBS=1 CARGO_PROFILE_DEV_DEBUG=0 "/home/enamakel/work/openhuman/worktrees/epic-8088-63-session-runtime/scripts/ci-cancel-aware.sh" cargo test --workspace +``` + +Do not set `CARGO_TARGET_DIR`. + +## Risks to resolve during implementation + +* Lease wrappers must snapshot both backend and tenant scope atomically. + Capturing a backend while deriving scope later would permit cross-tenant + reads after actor context changes. +* Snapshot encoding must preserve existing transcript records and exact + override behavior. Additive extension metadata must not change model-facing + tool specs or silently revive revoked tools. +* `PostCommitSink` errors occur after the durable boundary and therefore can + only be reported; they cannot veto or roll back the committed turn. +* `ContextConfig` has product-only fields mixed with generic context knobs. + Before moving each field, document its owner and preserve existing config + aliases/defaults. Do not leave a generic implementation or dependency + backedge in OpenHuman as a compatibility shim. + +This document is the plan phase for issue #63. It does not claim the issue is +implemented or complete. Open the upstream PR against `senamakel/tinyagents` +with base `epic/8088-tinyagents-slim` only after implementation, review, and +the required workspace gates pass; do not squash the individual commits. diff --git a/docs/runtime-comparison/README.md b/docs/runtime-comparison/README.md index 80c46a08c..30e5f9410 100644 --- a/docs/runtime-comparison/README.md +++ b/docs/runtime-comparison/README.md @@ -2,6 +2,10 @@ Date: 2026-09-19. Baseline: TinyAgents v2.1.2 (`fc33c43`). +This comparison records the pre-#61 baseline. Claims about TinyAgents graph +execution, graph checkpointing, and graph-specific gaps are historical; the +graph runtime was removed in TinyAgents #61. + This directory compares TinyAgents with three agent *runtime libraries* and pairs that with a deep code review of our own crates. It answers three questions: @@ -22,17 +26,14 @@ as calibration for the runtime/harness split. |---|---| | [`feature-gaps.md`](feature-gaps.md) | Consolidated gap matrix across all three runtimes with a runtime-vs-host layer assignment for every item. Start here. | | [`plan.md`](plan.md) | Phased execution plan: PR-sized work items, dependencies, and what moves between OpenHuman and TinyAgents. | -| [`langgraph.md`](langgraph.md) | LangGraph 1.2 / LangChain 1.4 / Deep Agents 0.7: feature inventory, ranked gaps, design lessons. | | [`pydantic-ai.md`](pydantic-ai.md) | Pydantic AI v2 / pydantic-graph / pydantic-ai-harness: same structure. | | [`pi.md`](pi.md) | pi (`earendil-works/pi`): `pi-ai`, `pi-agent-core` and its `AgentHarness`: same structure. | | [`code-review-harness.md`](code-review-harness.md) | `tinyagents-harness` findings (3 critical, 13 important, 14 minor), refactors, test gaps. | -| [`code-review-graph.md`](code-review-graph.md) | `tinyagents-graph` / `-orchestration` / `-session` findings (4 critical, 12 important, 12 minor). | ## Executive summary **Where TinyAgents is ahead.** None of the three runtimes has our combination -of a durable typed graph (channels, `Send`, subgraphs, checkpoints with time -travel), a policy-checked steering channel, detached sub-agent registry with +of a policy-checked steering channel, detached sub-agent registry with parallel failure policies, fail-closed `ToolPolicy`, prompt-cache segment layout and capability-set model resolution. LangGraph OSS puts double-texting, cron and background runs in its paid @@ -53,11 +54,10 @@ sub-agents, no structured output, no budgets. boundaries or indices, so interleaved thinking/text and durable partials cannot be reconstructed (pi's `*_start/_delta/_end` + frame codec; LangGraph `StreamPart{ns, seq}`). -3. **Durability semantics.** The graph executor discards completed - higher-index parallel siblings on interrupt/failure and re-runs them, and - routes their successors into the wrong superstep (critical findings C1/C2). - Checkpoints are unversioned full-state snapshots; per-node - retry/cache/timeout policies and delta channel history are missing. +3. **Durability semantics.** Graph-executor findings in the original review + are historical and no longer describe shipped code. Session state remains + linear; per-node workflow retry/cache/timeout policies are outside the + current TinyAgents runtime. 4. **Tools.** `ToolExecutionContext` lacks call id, store and state access; tool results are text/JSON only (no `ToolReturn`-style multimodal follow-up + metadata); no composable toolsets; no generic MCP client; no diff --git a/docs/runtime-comparison/code-review-graph.md b/docs/runtime-comparison/code-review-graph.md deleted file mode 100644 index 4847b3e3c..000000000 --- a/docs/runtime-comparison/code-review-graph.md +++ /dev/null @@ -1,387 +0,0 @@ -# tinyagents-graph / -orchestration / -session: durable-runtime code review - -Reviewed at v2.1.2 (`fc33c43`), read-only. Paths are relative to the repository root. - -`cargo clippy -p tinyagents-graph -p tinyagents-orchestration -p tinyagents-session --all-targets -- -W clippy::pedantic` is clean under default lints; 564 pedantic warnings for graph (mostly `must_use`, `missing_errors_doc`, casts), notably: `execute_run` 499 lines, `run_active_parallel` 121 lines, `state_api::update_state` 121 lines. - -## 1. Architecture as-built - -Entry points `CompiledGraph::{run, run_with_thread, run_with_inputs, resume, resume_from, retry}` -(`compiled/executor.rs:18-244`) all funnel into `execute` → `execute_run` (`executor.rs:505`). -One run = a loop over supersteps on an `active: Vec<Activation>` (`compiled/mod.rs:176`, -`Activation {node, send_arg, task_id}`): - -1. Guards: `steps >= min(recursion_limit, policy.max_total_steps)` (`executor.rs:601-616`), - wall-clock `run_deadline` (`:623`), `max_visits_per_node` (`:643`). -2. Assign task ids `"{steps}:{index}:{node}"` (`:661-665`), emit `StepStarted`. -3. Run handlers: sequential (`run_active_sequential`, `:1377`) or, with `with_parallel(true)` - and >1 active, `run_active_parallel` (`:1483`; `join_all` or a `select_all` rolling window - under `max_concurrency`). Every handler gets `state.clone()` and a fresh `NodeContext` - (`:1297`, `:1247`). Per-node timeout + retry policy wrap each attempt (`:1265-1310`). -4. Fold results in index order (`fold_result`, `:1320`): updates → `Vec<Update>`, `Command.goto` - → `goto_map[index]`, first `Interrupt`/error → stop folding (`break` at `:1444/:1458/:1617/:1631`). -5. Boundary: `state = reducer.apply(state, update)` for each update (`:726`). -6. Failure boundary (`:764-830`): route the completed prefix `active[..failed_index]`, then - `pending = successors ++ active[failed_index..]`, persist a failure checkpoint, status Failed. -7. Interrupt boundary (`:838-926`): same shape with `active[..index]`, persist checkpoint with - `interrupts` + `metadata.interrupted_nodes`, return `GraphExecution{interrupts}`. -8. Otherwise `route_completed(active)` (`compiled/routing.rs:15`): `goto` > static edge > - conditional branch; waiting-edge barriers accumulate in `barrier_arrivals`; `Send` targets - are not deduped, plain targets are. Persist per `DurabilityMode` (`:947-1038`), chain - `parent_checkpoint`, `active = next`. - -Checkpoint record: `Checkpoint<State>` (`checkpoint/types.rs:151`) = full `state` clone + -`next_nodes` + `completed_tasks` + inline `pending_writes` (completion markers only) + -`pending_activations` + `barrier_arrivals` + `interrupts` + free-form `metadata` JSON. The -executor calls `put` then `put_writes` (`executor.rs:1696-1697`), so the same markers are -stored twice. There is no format version, no timestamp, no channel versions, no delta scheme: -every superstep snapshots the whole state. - -Resume (`resume_from_inner`, `executor.rs:246`): load the scoped checkpoint, take -`pending_activations` (fallback `next_nodes`), filter out tasks with a completion marker, -key the resume value on `metadata.interrupted_nodes` (fallback: fan across pending), then call -`execute` with `steps = 0` and `initial_parent = loaded id`. - -State API (`compiled/state_api.rs`): `get_state`, `get_state_history`, `update_state` -(latest only), `bulk_update_state`, `fork_state` (new root, no parent). - -Subgraphs (`subgraph/mod.rs`): a node handler that clones the child graph with namespace -`parent_ns ++ [node_id]` and calls `run_with_thread` / `resume` depending on `ctx.resume` -(`drive_child`, `:135-164`). A child interrupt is re-emitted as the parent node's result. - -Other crates: `tinyagents-orchestration/workflow/engine.rs` is a *second* scheduler: phases -are a JSON `phase_states` blob persisted through the session run ledger (CAS + lease), with -`map_reduce` for fan-out; it does not use `CompiledGraph` at all (the graph in -`workflow/graph.rs` is a topology preview only). `teams/graph.rs` wraps one worker future in -an unthreaded, uncheckpointed 2-node graph. `tinyagents-session` is a SQLite ledger -(one `Connection::open` per operation, `store.rs:65`, WAL + busy_timeout). - -Divergence from docs (see also findings): `docs/modules/graph/execution.md` lists a target -executor with channel versions, cached-write replay, task-level `TaskStarted/TaskCached` -events; none exist. `checkpointing.md:56-77` lists checkpoint fields (version, timestamp, -channel versions, versions seen, task outcomes) that the record does not have. -`interrupts.md` shows `interrupt_before/after`, resume-by-interrupt-id maps, and -`resume_targeted`; none are implemented (`mark_interrupt` is export-only, -`builder/mod.rs:365`). `parallel-agents-forking.md:156-166` promises "resuming from -interrupt restarts the interrupted child task, not unrelated completed siblings", which is -false for higher-index siblings (Critical 1). - -## 2. Findings - -### Critical - -**C1. Completed parallel siblings after the interrupting/failing index are discarded and -re-executed on resume.** -`executor.rs:1600-1633` (parallel fold): -```rust -for (index, (activation, result)) in active.iter().zip(results).enumerate() { - ... - Err(error) => { ... failure = Some(StepFailure{failed_index: index, error}); break; } - ... - if let Some(found) = self.fold_result(...) { interrupt = Some(found); break; } -``` -and `executor.rs:869` / `:790`: `pending.extend(active[index..].iter().cloned());`. -`join_all` has already driven every branch to completion, so branches `index+1..` finished -(LLM calls, tool side effects, sub-agent runs), yet their `NodeResult`s are dropped and they -are scheduled to run again. The docs' own contract ("failed sibling tasks do not force -successful child agents to rerun once pending writes are saved") is not met; the pending-write -ledger only records the *lower-index* prefix. Fix: fold every `Ok` result regardless of -position; record all completed tasks in `completed_tasks`/`pending_writes`; make `pending` -= successors of *all* completed + only the interrupted/failed activations. The test -`parallel_interrupt_pauses_at_lowest_index_branch` (`compiled/test.rs:1002`) pins the current -behaviour and would need updating. Effort: M. - -**C2. Successors of completed siblings run in the *same* superstep as the re-run of the -interrupted/failed node, observing a different state than an uninterrupted run would.** -`executor.rs:856-869`: `pending = route_completed(active[..index]) ++ active[index..]`. -Uninterrupted: step N = [A, B]; step N+1 = [S_A, S_B] and S_A sees B's update. -Interrupted at B: resume step 1 = [S_A, B]; S_A now reads a state *without* B's update, and -B's update lands one step later. Any node that aggregates over sibling output through a plain -edge (not a waiting edge) silently produces different results depending on whether an -interrupt happened. LangGraph avoids this by re-running only the incomplete tasks of the -*same* step and applying all writes together. Fix: persist the completed tasks' updates as -real pending writes (requires `Update: Serialize` for durable backends, or store the -already-reduced partial state plus "tasks still owed for step N"), and on resume finish step N -before routing anyone. Effort: L (ties to R2 below). - -**C3. No per-thread execution lock: two concurrent `run_with_thread`/`resume` on the same -thread interleave one lineage.** -`executor.rs:246-390` takes no lock; `delegation/run.rs:139-160` and `:200-212` each add their -own `ThreadLockMap` and say why: "neither this wrapper nor the `CompiledGraph` run/resume -entry points hold a per-thread lock across that gap on their own". Backends allow duplicate -`(thread_id, checkpoint_id)` rows (`sqlite.rs:135-147` has no UNIQUE constraint) and `get(None)` -returns the last inserted row, so two writers produce a thread whose "latest" flips between -two histories and whose `parent_checkpoint_id` chains cross. Fix: an in-process -`ThreadLockMap` guard in `execute` keyed by `(thread_id, namespace)` plus an optional durable -lease in `Checkpointer` (`try_claim(thread, owner, ttl)`) like `run_ledger::try_claim_workflow_run` -(`session/run_ledger/ops.rs:169`). Effort: S (in-process) / M (durable). - -**C4. Subgraph failures are not resumable through the parent; `retry` restarts the child from -scratch and re-runs its completed nodes.** -`subgraph/mod.rs:161`: `(Some(thread_id), None, None) => child.run_with_thread(thread_id, state)`. -After a child node fails, the child writes a failure checkpoint in its namespace and the parent -writes one scheduling the subgraph node. `parent.retry()` re-runs the node with `resume = None` -→ fresh `run_with_thread` with `initial_parent = None`: the child's partial progress is ignored, -its completed nodes (with side effects) re-execute, and a second root-less lineage is appended -to the same namespace. Fix: in `drive_child`, before running fresh, check -`get_scoped(thread, None, child_ns)` for a checkpoint with non-empty pending activations -whose `metadata.failed_node` is set (or any pending), and `retry` it; also record the child's -checkpoint id in the parent's activation so the association is explicit. Effort: M. - -### Important - -**I1. `Send` fan-out of the same subgraph node shares one checkpoint namespace.** -`subgraph/mod.rs:109-113`: namespace = parent ++ `[ctx.node_id]` only. N concurrent -activations of node `worker` (map-reduce over a subgraph) write N interleaved lineages under -`["worker"]`; on parent resume each activation calls `child.resume(thread)` → `get_scoped(..., -None)` → all N resume from whichever child wrote last. Fix: namespace by -`[node_id, task_id]` (task id is already on the `Activation`, `compiled/mod.rs:179`) and expose -`NodeContext.task_id`. Effort: S. (Same identity gap: `resume_map` is keyed by `NodeId`, -`executor.rs:361-374`, `:1247`, so only the first same-node activation receives the value.) - -**I2. `update_state` erases interrupt provenance, so a later `resume(value)` fans the value to -every pending node.** `state_api.rs:228-231` writes `interrupts: Vec::new()` and metadata -`{source, step}` without `interrupted_nodes`; `compiled/mod.rs:229-251` then finds nothing -stamped and nothing in `interrupts`, and `executor.rs:363-366` falls back to fanning across -`active`. The documented flow "inspect → `update_state` → `resume`" therefore hands -`ctx.resume` to nodes that never interrupted. Fix: carry `interrupts` and `interrupted_nodes` -through an `update` checkpoint unless `as_node` names the interrupted node. Effort: S. - -**I3. Step counter, node-visit counts and recursion caps reset on every resume.** -`executor.rs:518` `let mut steps = 0usize;` and `:513-520` fresh `node_visits`. A loop that -interrupts (or fails and is auto-retried by a host) every step is unbounded; checkpoint -`metadata.step` restarts at 1 after each resume, so `get_state_history` shows non-monotonic -steps and `update_state`'s `parent_step + 1` (`state_api.rs:231`) is meaningless across -resumes. Fix: seed `steps` from the loaded checkpoint's `to_metadata().step` and persist -`node_visits` in metadata. Effort: S. - -**I4. Node panics and mid-superstep cancellation leave the run `Running` forever with no -terminal event.** No `catch_unwind` anywhere in `compiled/` (grep). A panic in a handler -unwinds through `join_all`/`fut.await` (`:1297`); dropping the run future (host timeout, -`tokio::select!`) skips `fail_run`, `save_status`, `RunFailed`, and drops -`AsyncCheckpointWrites` (detaching in-flight writes, contrary to its own contract at -`types.rs:141-149`). The executor also has no `CancellationToken` input at all (`map_reduce` -and `SubAgentPolicy` do). Fix: wrap handler futures in `AssertUnwindSafe(..).catch_unwind()` -mapped to `TinyAgentsError::Graph`; add `run_with_cancel(token)` checked at step boundaries and -raced against the handler set; add a `Drop` guard on `execute_run` that writes a `Cancelled` -status. Effort: M. - -**I5. The channel state model is not durable.** `channel/types.rs:240` `ChannelState` (and -`ChannelSet`, holding `Box<dyn Channel>`) has no `Serialize`/`Deserialize`, so the only state -model with per-key reducers, conflict detection and barrier channels cannot be used with -`FileCheckpointer` or `SqliteCheckpointer` (both bound `State: Serialize + DeserializeOwned`, -`file.rs:372`, `sqlite.rs:203`). `state-channels.md` presents channels as the way to get -"checkpoints can store pending writes". Fix: serialize `ChannelSet` as `{kind, config}` per -channel + values (each `Channel` already has `kind()`), with a registry for -`BinaryAggregate` closures. Effort: M. - -**I6. Checkpoint format has no version and four overlapping projections of "what runs next".** -`checkpoint/types.rs:151-201`: `next_nodes`, `completed_tasks`, `pending_writes`, -`pending_activations` all derive from the same activation list and can disagree when written -by hand/`update_state` (`state_api.rs` has a long comment defending the merge for that reason). -No `format_version` field means a future incompatible change can only be detected by serde -`Category::Data` heuristics (`checkpoint/mod.rs:70-99` documents this as an accepted risk). -`delegation` works around it with its own `schema_version` in the state (`delegation/run.rs`). -Fix: v2 record with `version: u32`, `tasks: Vec<PendingActivation>` as the single source, and -`completed: Vec<TaskId>`; keep a v1 → v2 decoder. Effort: M. - -**I7. `Interrupt::new` ids are process-local counters, unlike checkpoint/run ids.** -`command/mod.rs:110-112`: `INTERRUPT_SEQ.fetch_add` + `format!("interrupt-{node}-{seq}")`. -After a restart the same `(node, seq)` is re-minted; `GraphRunStatus.pending_interrupts` -(`status/types.rs`) and any UI keyed on interrupt id conflate two pauses. `ids::new_checkpoint_id` -already solved this with a process nonce (`harness/src/ids/mod.rs:135`). Fix: use it. Effort: S. - -**I8. SQLite backend: synchronous I/O on the async executor path, no WAL, no busy timeout, -and full-state loads for `state_history(limit)`.** Only `put` uses `spawn_blocking` -(`sqlite.rs:207`); `get`, `get_scoped`, `list`, `state_history`, `get_thread`, `put_writes`, -`delete_*` lock a `std::sync::Mutex<Connection>` and run queries inline in `async fn` -(`:268`, `:314`, `:359`, `:436`, `:472`, `:491`, `:619`). `from_connection` never sets -`journal_mode=WAL`/`busy_timeout`/`synchronous` (the session crate does, `session/store.rs:49-55`). -`state_history` reads and deserializes *every* record of the namespace even for `limit: Some(1)` -(`:353-374`). `put` and `put_writes` are two transactions. Fix: `spawn_blocking` everywhere, -pragmas on open, `LIMIT`-driven lineage query (walk parents with a recursive CTE), one -transaction per boundary. Effort: M. - -**I9. File backend: `list`/`get_scoped`/resume deserialize every full record; `put_writes` -rewrites and fsyncs the whole sidecar every superstep.** `file.rs:451` -`list` → `read_records` (`:247`, full `Checkpoint<State>` decode); `FileCheckpointer` does not -override `get_scoped` (no match in `file.rs`), so `resume`/`update_state`/`fork_state` cost -O(H) full-state decodes plus a second streaming pass in `get`. The trait doc at -`checkpoint/mod.rs:129` claims "both durable backends do [override]" — false. `put_writes` -(`:621-653`) is read-modify-rewrite with `write_atomic` (fsync), synchronously on the executor -task. Fix: header-only decode for `list` (as `get` already does), override `get_scoped`, -append-only writes sidecar keyed by checkpoint id, `spawn_blocking`. Effort: M. - -**I10. `GraphBuilder::add_edge` silently overwrites a previous edge from the same node.** -`builder/mod.rs:195-197`: `self.edges.insert(from.into(), to.into())` (also `add_waiting_edge`, -`:226`). `edges: HashMap<NodeId, NodeId>` means static fan-out is impossible and -`add_edge("a","b").add_edge("a","c")` compiles to `a → c` with no error; `routing.md:49-52` -lists "one or more node names" as a routing output. Fix: `HashMap<NodeId, Vec<NodeId>>`, or at -least return `Validation` on a second insert. Effort: S/M. - -**I11. `#![cfg_attr(not(feature = "tracing"), allow(dead_code, unused_imports, -unused_variables))]` at crate root hides dead code in the default build.** -`tinyagents-graph/src/lib.rs:25-28`, `tinyagents-session/src/lib.rs:1-4`. The default feature -set is `tracing` off, so every unused item in both crates is invisible to `cargo clippy -D -warnings`. Evidence of accumulation: `StreamMode` (`stream/types.rs:220`) is exported and -consumed nowhere; `WRITES_IDX_RESUME/ERROR/INTERRUPT` (`checkpoint/types.rs:272-278`) are only -used by the conformance suite — the executor never persists a resume value, so a resumed node -that crashes loses it and `retry` re-interrupts; `CompiledGraph.command_nodes` and -`BuilderNode.id` carry explicit `#[allow(dead_code)]`. Fix: make the tracing macros expand to -`{ let _ = (...); }` in the off configuration and delete the crate-level allow. Effort: S. - -**I12. Two durable schedulers with different semantics.** `orchestration/workflow/engine.rs` -implements phase DAG scheduling, lease/CAS persistence, cancellation and resume in ~830 lines -of JSON manipulation (`state.rs` string-typed statuses `"running"`, `"completed"`), emitting -`GraphEvent::NodeStarted{node: "run_phase", step: total_spawned+1}` (`engine.rs:428`) to look -like a graph. `workflow/graph.rs:37` says "`WorkflowEngine` installs its effectful variants -below" but the engine never runs that graph. The graph runtime already has waiting edges, -`Send` fan-out, `max_concurrency`, checkpoints and resume; the workflow engine re-implements -them without any of the executor's tests. Fix: lower a `WorkflowDefinition` to a -`CompiledGraph` (phase = node, `depends_on` = waiting edges, agents = `Send` fan-out) and keep -the run ledger as a status projection. Effort: L. - -### Minor - -**M1. `execute_run` is 499 lines with 12 `#[allow(clippy::too_many_arguments)]`** across -`compiled/` (`executor.rs`, `mod.rs`). Thirteen positional args of `Option<&..>`/`&HashMap` -(`persist_checkpoint`, `:1652-1666`) are error-prone. Introduce a `RunCtx` struct (run id, -thread, started_at, recursion meta, barriers, async writes) and a `Boundary` struct. Effort: M. - -**M2. Hot-path cloning.** `executor.rs:1297` clones `State` per handler *per attempt*; -`:1188`/`:1855` clone it again per checkpoint; `serde_json::to_string(&checkpoint)` per -boundary; `Send` args cloned into `Activation`, `PendingActivation`, `NodeContext`. For -`serde_json::Value` states with message histories this is O(history) per node per step. Fix: -`NodeHandler = Fn(Arc<State>, NodeContext)` (one clone per step), `Arc<Value>` send args. -Effort: M (API break). - -**M3. `run_deadline` uses `SystemTime::elapsed`** (`executor.rs:624`); wall-clock jumps make -runs fail or never time out. Use `Instant`. Effort: S. - -**M4. Status-store errors are swallowed** (`executor.rs:500` `let _ = store.put_status`), so a -dead status backend is invisible. Log at warn at minimum. Effort: S. - -**M5. Async durability skips `put_writes`.** `executor.rs:1764` background task only calls -`put`; sync path calls both. Ledger tooling sees no markers for async-mode runs. Effort: S. - -**M6. `GraphEvent` carries no run id/task id/namespace on step/node events** -(`stream/types.rs:46-104`); the module doc (`stream/mod.rs:8-10`) claims nested streams can be -"attributed back up the run tree", but only the journal wrapper adds that. Add a `RunRef` -header or emit `(run_id, event)`. Effort: S/M. - -**M7. `GraphEventJournal` is lossy under load.** `observability/mod.rs:464` uses the harness -`AppendWorker`, which `try_send`s and increments `dropped` when the bounded queue is full -(`harness/src/observability/worker.rs:250-253`). A "durable ... replay" journal that drops -events should at least expose `dropped()` on `JournalGraphSink` and document it. Effort: S. - -**M8. Router closures return `String`, `Route` is unused by the executor** (`builder/types.rs` -`RouterFn = dyn Fn(&State) -> String`; `Route` newtype defined but branches key on `String`). -A typo in a route label is a runtime `MissingRoute`. Effort: S. - -**M9. `language::build_graph` drops declared conditional routes**: `language.rs:41` -`Routing::Conditional(_) => builder.mark_command_routing(...)` ignores the route table, so a -blueprint's routing is not validated against what the node's `Command` may return. Effort: S. - -**M10. Session ledger opens a new SQLite connection per operation** (`session/store.rs:65`), -including the lease heartbeat every `lease/3` (`engine.rs:585`), and all ledger calls are sync -inside `async fn drive`. A small connection cache or `spawn_blocking` would help. Effort: S. - -**M11. `PhaseRegistration::register` performs a blocking DB CAS under a `parking_lot::Mutex` -inside an async executor callback** (`engine.rs:217-224`). Effort: S. - -**M12. Subgraph doc lists unimplemented requirements as if present** (`subgraphs.md:15-25`: -isolated child thread ids, checkpointer override, `Command::Parent`, child state in parent -task metadata). Mark them as target or remove. Effort: S. - -## 3. Structural refactors worth doing - -**R1. Split `compiled/executor.rs` along boundaries, not along "public vs private".** -`RunCtx` (identity, clocks, recursion meta, async writes), `StepRunner` (sequential/parallel, -returning a `StepOutcome` with *all* results), `Boundary` (apply reducer, route, persist), and -`Resume` (load + filter + resume map). This is the precondition for C1/C2 because the "which -results do we keep" decision is currently spread over four `break`s and two `extend`s. Risk: -low, pure code motion with the existing 100 unit tests as a net. - -**R2. Real pending-writes: make the executor persist task outputs, not just markers.** -Add `Update: Serialize + DeserializeOwned` bounds on the durable checkpointers only (the -in-memory path stays bound-free), store `PendingWrite.payload = serde_json::to_value(update)`, -and on resume replay stored writes for completed tasks of the interrupted step instead of -re-running them (this is what `execution.md`'s "cached writes replay without rerunning nodes" -describes). Fixes C1 and C2 and gives `WRITES_IDX_RESUME` a purpose (persist the resume value -so `retry` after a crash still has it). Migration risk: medium; the `Update` bound is a public -API change for durable users, hidden behind a `DurableUpdate` marker trait with a blanket impl. - -**R3. Checkpoint format v2 + versioned channels.** Add `version`, `created_at`, `task_id` -on `Interrupt`, one `tasks` list, and an optional `channel_versions: BTreeMap<String, u64>` -written by `ChannelState` (I5). Versioned channels are what make subgraph output merging -correct (`shared_subgraph_node` returns the child's *whole* final state as the parent update, -so an append reducer double-applies inherited messages) and what enables delta snapshots -(`prune` already documents "delta-channel" semantics that nothing implements, -`checkpoint/mod.rs:462-476`). Risk: medium; needs a decode shim for v1 rows and a -`FileCheckpointer`/`SqliteCheckpointer` migration test. - -**R4. Per-thread run lease in the executor (C3)** with an in-process `ThreadLockMap` (already -exists at `thread_locks/mod.rs`) and an optional `Checkpointer::try_claim/renew/release` -defaulting to no-op; `delegation/run.rs` can then drop its private lock map. Risk: low. - -**R5. Typed task identity.** `TaskId` exists in `harness::ids` (used by `RecursionFrame`), -but the executor uses `String` task ids, `Interrupt` has none, and `NodeContext` does not -expose it. Threading `TaskId` through `Activation`, `PendingActivation`, `PendingWrite`, -`Interrupt`, `NodeContext`, and the subgraph namespace (I1) closes several identity bugs at -once. Risk: low-medium (public struct fields). - -**R6. Retire the workflow engine's private scheduler (I12)** in favour of a -`WorkflowDefinition → CompiledGraph` lowering; keep `WorkflowStore` as the status projection -and the lease as R4's durable lock. Risk: high (host-facing `phase_states` JSON contract); -do it behind a feature flag with the existing `workflow/tests.rs` as the acceptance suite. - -## 4. Test-quality assessment - -Strengths: `compiled/test.rs` (100 tests) covers routing, `Send` args across resume, barrier -arrivals across resume, async-durability ordering/failure/drain, attributed `update_state` -merges, retry/failure checkpoints, and legacy-checkpoint fallbacks. `testkit/conformance.rs` -has checkpointer contract suites (writes, lineage, concurrency) run against all three -backends. `delegation/test.rs` (1.3k lines) exercises schema-version guards and lock -serialization. Tests are descriptive and mostly assert observable durable state, not -internals. - -Gaps (each maps to a finding): -- No test that a *higher-index* completed parallel branch is not re-run after an interrupt - or failure, nor that its update survives (C1). `parallel_interrupt_pauses_at_lowest_index_branch` - asserts the lossy behaviour. -- No test comparing final state of an interrupted-then-resumed run against the same graph - run uninterrupted (C2) — the single most valuable durability property test. -- No concurrent-same-thread test for the executor (C3); only delegation tests it indirectly. -- No subgraph failure → parent `retry` test (C4); `subgraph/test.rs` covers interrupt resume - only. No `Send` fan-out of a checkpointed subgraph node (I1). -- No `update_state` (as_node None) followed by `resume(value)` test asserting the value - reaches only the interrupted node (I2). -- No test for step monotonicity across resume or cumulative recursion limits (I3). -- No node-panic test and no cancellation (dropped future) test asserting status/checkpoint - consistency (I4). `orchestration/test.rs:254` is the only `catch_unwind` in the crate. -- No channel-graph durability test (I5) — impossible today, which is the point. -- No test exercising a `FileCheckpointer`/`SqliteCheckpointer` behind the *executor* for - interrupt → resume across a fresh `Connection`/process nonce (the "restart" path); the - durable-backend tests are conformance-only and executor tests use `InMemoryCheckpointer`. -- No performance guard on `state_history(limit)` or on `list` decode cost (I8/I9). -- `workflow/tests.rs` (850 lines) is thorough on lease/CAS races but relies on string statuses; - a typo in `"completed"` would pass type-check and fail only at runtime. - -## 5. Things that are genuinely good - -- Deterministic fold order and explicit `goto_map` keyed by index (repeated `Send` of one node - keeps its own routing) — a real correctness win over node-keyed maps. -- `AsyncCheckpointWrites` chaining with first-error propagation and drain-on-every-exit is - carefully reasoned and tested (`compiled/types.rs:153-249`). -- Barrier-relief with `reaches_deterministically` (`routing.rs:68-158`) solves a real - deadlock without weakening the barrier, and the comment explains the multi-hop pitfall. -- Failure boundaries are resumable and `retry`/`update_state` compose with them; the - `interrupted_nodes` metadata keeps resume values off nodes that never paused. -- Checkpoint id minting uses a process nonce; `state_history` has a cycle guard; - `copy_thread` refuses to interleave lineages; torn trailing JSONL lines are tolerated. -- `ThreadLockMap` weak-value map is a neat leak-free per-key mutex. -- Schema-version fencing and per-thread locking in `delegation/run.rs` show the team knows - the executor's gaps; the fixes above mostly move that discipline into the runtime. -- `map_reduce` gets ordering, fail-fast-in-input-order, cancellation and timeouts right. -- The session ledger's lease + revision CAS (`run_ledger/ops.rs:169-320`) is a solid pattern - the graph checkpointer should borrow. diff --git a/docs/runtime-comparison/code-review-harness.md b/docs/runtime-comparison/code-review-harness.md index 24a0f8ec3..28071a546 100644 --- a/docs/runtime-comparison/code-review-harness.md +++ b/docs/runtime-comparison/code-review-harness.md @@ -31,8 +31,7 @@ tool-wrap onion, or `join_all` when eligible) → ordered fold (`after_tool`, ho append) → loop. Termination: no tool calls (after truncated-empty retry and `continue_turn` nudge) → structured extraction → `LoopExit::Finished`; or `LimitStop`, `Paused`, or `Err`. -**Sub-agents** (`subagent/mod.rs`): `SubAgentTool` is a `ToolDispatch` whose `execute` builds `parent.child(...)` and runs the child -harness inside the parent's tool call, sharing `EventSink`, `CancellationToken`, `SteeringHandle`, stores and workspace. +**Sub-agents** (`tinyagents-orchestration`): `DriverSubAgentTool` starts a child through the host-planned `SubagentDriver` and returns a job reference immediately. The host supplies authorization, persistence, and completion routing; the registry exposes status, cancellation, steering, and messaging. **Providers.** Only two adapters live in this crate (`providers/claude_code`, `providers/claude_agent_sdk` — subprocess/JSON-RPC drivers). OpenAI/Anthropic/local adapters live in `vendor/tinyinference/crates/tinyinference-llm`; the harness never sees wire JSON. @@ -282,7 +281,7 @@ and, on success, clear `invalid` and proceed to normal validation (emit `Invalid `From`. Add `#[non_exhaustive]` to it and `AgentEvent` now (cheap, prevents the next break). Risk: medium for downstream matches. 5. **R-5. Feature-gate heavy optional surfaces**: `claude-code` (subprocess driver + `uuid`, `tempfile`, `wait-timeout`, `dirs`), `langfuse` (`reqwest`), keep `tools` owning `chrono`. Risk: low; integration tests already gate `sqlite`/`tools`. -6. **R-6. Move `handoff`, `run_queue`, `no_progress`, `artifacts`, `workspace/git` into a `tinyagents-host-utils` crate** or under a +6. **R-6. Move `run_queue`, `no_progress`, and `artifacts` into a `tinyagents-host-utils` crate** or under a `host-utils` feature. They are not on any loop path and carry product-specific heuristics (M-9). --- diff --git a/docs/runtime-comparison/feature-gaps.md b/docs/runtime-comparison/feature-gaps.md index 5fe362cd8..a6551a22c 100644 --- a/docs/runtime-comparison/feature-gaps.md +++ b/docs/runtime-comparison/feature-gaps.md @@ -3,8 +3,8 @@ Cross-runtime view of what LangGraph/LangChain (LG), Pydantic AI (PA) and pi have that TinyAgents lacks or only partly has, with a layer decision for each. Per-runtime detail, API shapes and code excerpts live in -[`langgraph.md`](langgraph.md), [`pydantic-ai.md`](pydantic-ai.md) and -[`pi.md`](pi.md); the section numbers in the "Source" column point there. +[`pydantic-ai.md`](pydantic-ai.md) and [`pi.md`](pi.md). LangGraph references +in this historical gap inventory describe the external project. Layer legend: **RT** = TinyAgents runtime (crate named), **OH** = OpenHuman, **RT+OH** = primitive in the runtime, policy/UI in OpenHuman. @@ -20,9 +20,9 @@ that the feature would extend. | A2 | Deferred tool calls as a typed, resumable output: `DeferredToolRequests{calls, approvals, metadata}` / `DeferredToolResults` with per-call ids, approve / edit-args / reject / respond decisions, external execution (`CallDeferred`), durable across process restart | PA, LG (HITL middleware) | PA §3.2, LG §3.2 | `HumanApprovalMiddleware` (`Fn(&ToolCall) -> bool`), `Err(Interrupted)`, graph `Interrupt`, delegation `PendingApproval`, `ToolPolicy.access.approval_required` | RT+OH (OpenHuman `security/approval::ApprovalGate` becomes a handler) | 1 | shipped | | A3 | Output-validation retry loop: validator raises `ModelRetry`, loop re-asks with the error, bounded by `retries.output`; unified retry/`ToolFailed` vocabulary for tool errors, arg errors and output errors | PA | PA §3.1, §4 | `StructuredOutcome.error`, `structured/repair.rs`, `InvalidArgsPolicy`; docs already promise `StructuredOutputErrorPolicy` | RT harness | 1 | shipped | | A4 | Steering and follow-up as two message lanes polled at turn boundaries with `QueueMode` (one-at-a-time / all); follow-ups run "one more turn" after the agent would stop | pi, PA (`ctx.enqueue`) | pi §4.4 | `RunQueue{Steer, Followup, Collect}` (no consumer), `SteeringCommand` | RT harness (+ OpenHuman `agent/harness/run_queue` moves down) | 1 | shipped | -| A5 | Agent loop compiled to the graph runtime so checkpoints, interrupts, streaming and time travel apply to the loop itself; step iteration (`agent.iter()`) | LG (`create_agent` → `StateGraph`), PA (`iter`, `UserPromptNode`…) | LG §4, PA §4 | `subagent_node`, `docs/modules/harness/state-graph.md` (promised) | RT harness+graph | 2 | partial (`compile_loop`/`LoopIter` are a real `CompiledGraph` with checkpoint+resume tests; the drop-in `GraphLoopDriver` drives the same node bodies without a checkpointer) | +| A5 | Agent loop compiled to a graph runtime so checkpoints, interrupts, streaming and time travel apply to the loop itself; step iteration (`agent.iter()`) | LG (`create_agent` → `StateGraph`), PA (`iter`, `UserPromptNode`…) | LG §4, PA §4 | Removed with the graph crate in TinyAgents #61; no current implementation | Runtime | 2 | descoped: the graph crate and `GraphLoopDriver` were removed; harness runs use the direct loop | | A6 | Prompted / union / output-function structured modes; `EndStrategy` when output tools and function tools co-occur | PA, LG (`AutoStrategy`) | PA §3.8 | `StructuredStrategy::{ProviderSchema, ToolCall}` | RT harness | 3 | shipped | -| A7 | `interrupt_before/after` selectors, `interrupt(response_schema)`, graceful drain (`RunControl`) | LG | LG §3.8 | documented in `interrupts.md`, `mark_interrupt` export-only, `CancellationToken` | RT graph | 3 | shipped | +| A7 | `interrupt_before/after` selectors, `interrupt(response_schema)`, graceful drain (`RunControl`) | LG | LG §3.8 | harness steering and cancellation primitives | Runtime | 3 | partial: direct-loop steering is supported; graph selectors and checkpointed interrupt routing are not | ## B. Tools @@ -42,25 +42,20 @@ that the feature would extend. |---|---|---|---|---|---|---|---| | C1 | Block-indexed stream events (`text/thinking/toolcall _start/_delta/_end` with `content_index`), error terminals that carry the partial assistant message and stop reason | pi | pi §4.2, `sdk-gaps/streaming.md` §3 | `ModelStreamItem::{MessageDelta, ToolCallDelta, UsageDelta, Completed}` | RT tinyinference+harness | 1 | shipped | | C2 | Compact durable frame codec + reducer for partial messages (crash recovery, reconnecting clients) | pi | pi §4.2 | `AgentEvent::ModelDelta` journaled, `HarnessEventJournal` | RT harness | 2 | partial (frame codec/reducer shipped and journaled; `ToolProgress`/`on_tool_delta` still has no real mid-execution caller pending a tinytools progress seam) | -| C3 | Unified stream envelope with `ns`/`seq`, `tasks` and `checkpoints` stream modes, projections (`stream.messages`, `stream.tool_calls`, `stream.subagents`) | LG | LG §3.5 | `GraphEvent` (no run/task id on step events), `GraphObservation` journal, `StreamMode` (unused) | RT graph+harness; OH keeps format adapters | 2 | shipped (`GraphEventEnvelope{run_id, task_id, ns, seq}`, `StreamMode::{Tasks, Checkpoints}` and `StreamProjection` shipped; `task_id` is `Some(..)` on every per-task/task-scoped event — `TaskScheduled`/`TaskStarted`/`NodeStarted`/`NodeCompleted`/`NodeFailed`/`TaskCompleted`/`NodeRetryScheduled`/`StateUpdated`/`InterruptEmitted`/`ContextForked`/`RouteSelected` — and `None` only for run/step-level events with no single owning task) | +| C3 | Unified stream envelope with `ns`/`seq`, `tasks` and `checkpoints` stream modes, projections (`stream.messages`, `stream.tool_calls`, `stream.subagents`) | LG | LG §3.5 | harness event streams and session records | RT harness+session; OH keeps format adapters | 2 | partial: direct harness events and session streams are available; graph projections were removed in TinyAgents #61 | | C4 | OpenTelemetry GenAI semantic-convention sink | PA | PA §3.12 | Langfuse exporter, `JournalSink` | RT harness optional feature | 3 | OpenHuman | -## D. Graph durability +## D. Graph runtime scope -| # | Gap | Who has it | Source | Existing seam | Layer | Value | Status | -|---|---|---|---|---|---|---|---| -| D1 | Real pending writes: completed parallel siblings are never re-run after an interrupt/failure and their successors run in the right superstep (today's C1/C2 bugs) | LG | [`code-review-graph.md`](code-review-graph.md) C1, C2, R2 | `PendingWrite` (markers only), `completed_tasks` | RT graph | 1 | shipped | -| D2 | Per-node `RetryPolicy`, `CachePolicy` (task cache with backends), `TimeoutPolicy{run, idle}`, `error_handler`, real `defer` scheduling | LG | LG §3.3 | graph-wide `with_node_retry` / `with_node_timeout`, harness `ResponseCache`, `mark_deferred` (export-only) | RT graph | 2 | shipped | -| D3 | Checkpoint v2: format version, single task list, channel versions, serialisable `ChannelState`, delta-channel history with snapshot frequency, `Overwrite` | LG (`DeltaChannel`, `Overwrite`) | LG §3.4, code-review-graph I5/I6/R3 | `Checkpoint<State>` full snapshot, `ChannelSet` (not serialisable), `prune` doc mentions deltas | RT graph | 2 | shipped | -| D4 | Executor-level per-thread lease (in-process + optional durable claim), typed `TaskId` end to end, `Send` fan-out of subgraphs namespaced by task, subgraph failure resumable via parent, restart-safe interrupt ids, panic/cancel safety | LG (task ids, ns) | code-review-graph C3, C4, I1, I4, I7, R4, R5 | `ThreadLockMap`, `ids::new_checkpoint_id`, `run_ledger::try_claim` | RT graph | 1 | shipped | -| D5 | Durable task memoisation inside a node (`durable_task(ctx, key, fut)`) so side effects before an interrupt are not repeated | LG (functional API) | LG §3.9 | `PendingWrite`, `interrupts.md` "must use idempotency keys" | RT graph | 3 | shipped | -| D6 | Semantic search on the namespaced store (`IndexConfig{embed, dims, fields}`) | LG | LG §3.7 | `NamespacedStore::search` (substring), `retriever/`, `tinyinference-embeddings` | RT harness; backend choice in OH `memory/` | 3 | OpenHuman | +The graph-specific durability gaps in this comparison were descoped when the +TinyAgents graph crate was removed in #61. The current workspace does not ship +graph execution, graph streams, graph checkpointing, or graph task policies. ## E. Sessions and context | # | Gap | Who has it | Source | Existing seam | Layer | Value | Status | |---|---|---|---|---|---|---|---| -| E1 | Conversation entry tree: `id/parent_id`, in-place branching, labels, fork/clone with `parent_session`, `BranchSummaryEntry`, context projection that never reads past the newest compaction | pi | pi §4.3 | `tinyagents-session` keeps linear transcripts; graph exposes `fork_state` | RT session; `/fork`, `/tree`, labels UI in OH | 2 | removed | +| E1 | Conversation entry tree: `id/parent_id`, in-place branching, labels, fork/clone with `parent_session`, `BranchSummaryEntry`, context projection that never reads past the newest compaction | pi | pi §4.3 | `tinyagents-session` keeps linear transcripts; the former graph `fork_state` API was removed in #61 | RT session; `/fork`, `/tree`, labels UI in OH | 2 | removed | | E2 | Compaction rules: cut points at user/assistant boundaries, `keep_recent_tokens`, split turns, iterative summaries, durable `CompactionRecord{first_kept, tokens_before, usage}`, overflow classifier → compact → retry same turn, hook may decline/replace | pi, PA (`Compaction` capability), LG (`SummarizationMiddleware`) | pi §4.5 | `summarization/` (policy, trim, pairing, `Summarizer`), `ContextCompressionMiddleware`, `MicrocompactMiddleware` | RT harness+session; summary prompt wording in OH | 2 | shipped | | E3 | Cross-provider handoff transform: record `origin{provider, api, model}` on assistant messages; drop/convert thinking, normalise tool-call ids, downgrade images per target profile | pi | pi §4.6 | `AssistantMessage.id`, `ContentBlock::Thinking` doc note | RT tinyinference+harness | 2 | shipped | | E4 | `sanitize_history()` for untrusted client-supplied history (strip system prompts, non-HTTP file URLs, dangling tool calls) | PA | PA §5 | `summarization/trim.rs` orphan handling | RT harness (cheap helper); transport in OH | 3 | OpenHuman | @@ -71,7 +66,7 @@ that the feature would extend. | # | Gap | Who has it | Source | Existing seam | Layer | Value | Status | |---|---|---|---|---|---|---|---| | F1 | `ModelProfile` as behaviour: `json_schema_transformer`, `default_structured_output_mode`, `prompted_output_template`, `thinking_tags` parsing, per-API compat matrix (~30 flags), `thinking_level_map` | PA, pi | PA §3.6, pi §4.8 | `ModelProfile` + `CapabilitySet` (capability data), `SchemaPreparation`, `ReasoningConfig` | RT tinyinference | 2 | shipped | -| F2 | Generated model catalog (models.dev) with tiered pricing, `refresh_models()`, availability filtered by resolved auth | pi, PA (`genai-prices`) | pi §4.8 | `ModelCatalogEntry`, 5-model stale seed (`code-review-workspace.md` M7), `ModelPricing` | RT registry+tinyinference | 2 | shipped | +| F2 | Generated model catalog (models.dev) with tiered pricing, `refresh_models()`, availability filtered by resolved auth | pi, PA (`genai-prices`) | pi §4.8 | `ModelCatalogEntry`, `ModelPricing` | RT registry+tinyinference | 2 | shipped | | F3 | Audio / video / document content blocks in the message model, SSRF-guarded URL download | PA | PA §3.9 | `ContentBlock::Image(ImageRef)`, `multimodal/` | RT tinyinference (OpenHuman `agent/multimodal.rs` markers go away) | 2 | partial (`ContentBlock::{Audio,Video,Document}` shipped; SSRF-guarded URL download is OpenHuman -- see below) | | F4 | Provider breadth (Google, Vertex, Bedrock, Mistral native protocols), generalised OAuth flows, `CredentialStore` trait | pi (10 APIs, 41 presets, 7 OAuth flows) | pi §2, §4.8 | OpenAI + Anthropic + local in tinyinference, Codex `OAuthFlow` | RT tinyinference (trait); keychain and login UI in OH `security/credentials` | 3 | OpenHuman | | F5 | Deferred / background model responses (`DeferredHandle`, `stop_reason: deferred`) | pi | pi §4.9 | none | RT tinyinference | 3 | shipped (vendor `ModelStreamItem::Deferred` + `ChatModel::fetch_deferred`; no adapter overrides it yet) | @@ -84,7 +79,7 @@ that the feature would extend. | G1 | Evals: `Dataset` / `Case` / `Evaluator` trait, `LLMJudge`, span-based evaluators, report | PA (`pydantic_evals`), LG (`agentevals`) | PA §3.7 | `testkit::Trajectory` | RT new crate `tinyagents-evals`; datasets in OH | 2 | OpenHuman | | G2 | Schema-driven `TestModel` (auto-calls every tool with generated args), process-wide `deny_network_models()` kill-switch | PA | PA §3.10 | `ScriptedModel`, `StreamingMock`, `FakeTool` | RT harness testkit + tinyinference | 3 | shipped (`SchemaDrivenModel` in harness testkit; `deny_network_models()` in vendor tinyinference) | | G3 | Capability bundle: instructions + toolset + middleware + model defaults + exposure, loadable on demand, `from_spec`; what a "skill"/"plugin" is at runtime level | PA v2 (`Capability`) | PA §4 | `Middleware`, `AgentDefinition`, `ToolExposure::Deferred` | RT registry; discovery from disk in OH `skills/` | 2 | shipped | -| G4 | Backend conformance suites for sessions / stores run against every backend | pi (`session/testing/conformance`), LG | pi §3, `sdk-gaps/durability.md` §17 | graph `testkit/conformance.rs` (checkpointers, task stores) | RT | 3 | shipped | +| G4 | Backend conformance suites for sessions / stores run against every backend | pi (`session/testing/conformance`), LG | pi §3, `sdk-gaps/durability.md` §17 | session and store tests | RT | 3 | partial: session/store tests exist; the graph testkit and graph checkpointer conformance suite were removed in TinyAgents #61 | ## Stays in OpenHuman diff --git a/docs/runtime-comparison/langgraph.md b/docs/runtime-comparison/langgraph.md deleted file mode 100644 index a66bf6ede..000000000 --- a/docs/runtime-comparison/langgraph.md +++ /dev/null @@ -1,319 +0,0 @@ -# LangGraph / LangChain 1.x vs TinyAgents — runtime comparison - -Research date: 2026-09-19. TinyAgents baseline: v2.1.2 (`fc33c43`). -Paths below are relative to this repository unless prefixed `openhuman:`. - -## 1. What LangGraph/LangChain is today - -**Runtime layer — LangGraph.** `langgraph` 1.0.0 shipped 2025-10-17; 1.1.0 on 2026-03-10; 1.2.0 on -2026-05-12; latest core is 1.2.11 (2026-08-11), with `langgraph-checkpoint` 4.2.0 (2026-08-07). The -runtime is a Pregel/BSP engine: a superstep plans tasks by comparing per-channel version counters -(`channel_versions` vs each node's `versions_seen`), runs all tasks in parallel, then applies writes -through channel reducers and writes a checkpoint (`{v:2, id: uuid6, ts, channel_values, -channel_versions, versions_seen, updated_channels}` + metadata `{source: input|loop|update|fork, -step, parents}`). Nodes are triggered by channel-version changes, not by edges as such (edges compile -to `branch:to:<node>` channels). Control values are first-class: `Command(update, goto, resume, -graph=PARENT)`, `Send(node, arg, timeout)`, `Overwrite(value)` (bypass reducer), `interrupt(value, -response_schema)` which raises `GraphInterrupt` and re-runs the node from the top on resume, with -resume values matched by call index (scalar `Command(resume=)`) or by interrupt id (dict form). -1.x adds durability modes `sync|async|exit` (default `async`), per-node `RetryPolicy`, -`CachePolicy(key_func, ttl)` + `BaseCache`, `defer=True` join nodes, per-node `TimeoutPolicy` -and `error_handler` (1.2), `DeltaChannel` writes-only checkpoint history (1.2, beta), `RunControl` -graceful drain, `stream(version="v2")` unified `StreamPart{type, ns, data}` and -`stream_events(version="v3")` protocol events, `BaseStore` with namespaces/TTL/semantic `IndexConfig`, -and the functional API (`@entrypoint`/`@task`). Double-texting strategies, thread TTL, cron and -background runs are **Agent Server / LangSmith Deployment features, not OSS library features**. - -**Agent layer — LangChain.** `langchain` 1.0.0 (2025-10-17), 1.3.0 (2026-05-12), 1.4.2 (2026-09-18); -`langchain-core` 1.6.3 (2026-09-11). The whole agent layer is `create_agent(model, tools, -system_prompt, middleware, response_format, context_schema, checkpointer, store, ...)`, which -compiles to a small LangGraph `StateGraph` (`model` node, `tools` node, one node per middleware -hook). `AgentMiddleware` is the extension point: node hooks `before_agent/before_model/after_model/ -after_agent` (return state dicts, may `jump_to`), and wrap hooks `wrap_model_call(request, handler)` -/ `wrap_tool_call(request, handler)` (nest, first middleware outermost). A middleware can extend the -state schema, contribute tools, own stream transformers and trace policy. Built-ins cover -summarization, HITL (durable `interrupt()` with approve/edit/reject/respond), call limits, model -retry/fallback, PII, tool retry/error, LLM tool selection, provider tool search, context editing, -todo list, shell, file search. `ToolRuntime` injects state/context/store/stream_writer/tool_call_id; -`response_format` chooses `ToolStrategy|ProviderStrategy|AutoStrategy`. 1.4.0 added `langchain.mcp` -`MCPAdapter`. Deep Agents (`deepagents` 0.7.15, 2026-09-16) is a harness on top of `create_agent`, -explicitly "does not introduce a new runtime": filesystem tools over pluggable backends, subagents via -a `task` tool, summarization-with-offload, AGENTS.md memory, SKILL.md skills, path permissions, -sandbox backends. `langgraph-supervisor`/`langgraph-swarm` are in maintenance mode; the docs now -recommend `create_agent` + tool-wrapped subagents or middleware-driven handoffs. - -## 2. Feature inventory - -Legend: Yes / Partial / No. File paths are what I checked. - -| Feature | LangGraph/LangChain | TinyAgents | Notes | -|---|---|---|---| -| Pregel supersteps, parallel tasks, reducer at boundary | Yes | Yes — `crates/tinyagents-graph/src/compiled/executor.rs` | Same BSP shape. | -| Channel-version triggering (`versions_seen`) | Yes | No — `channel/mod.rs` is a reducer bridge; scheduling is edge/active-set based | See §4. | -| Typed channels (LastValue/Topic/BinaryOp/Ephemeral/NamedBarrier/Untracked) | Yes | Yes — `graph/src/channel/types.rs` | `Delta` in TinyAgents is a numeric accumulator, unrelated to LangGraph's `DeltaChannel`. | -| `DeltaChannel` writes-only checkpoint history, `snapshot_frequency` | Yes (1.2, beta) | No — `checkpoint/types.rs` stores full `state: State` per checkpoint | Deepagents uses it for `messages` (O(N) vs O(N²)). | -| `Overwrite` (bypass reducer) | Yes | No | Reducer runs on every write. | -| `Command(update, goto, resume)` | Yes | Yes — `graph/src/command/types.rs` | | -| `Command(graph=PARENT)` from subgraph/tool | Yes | Partial — `subgraph/`, `resume_targeted` exists; no parent-directed goto from a child node found | | -| `Send` fanout, persisted send args | Yes | Yes — `checkpoint/types.rs::PendingActivation` | `Send.timeout` (1.2) absent. | -| `interrupt()` inside a node, index/id-matched resume | Yes | Yes — node returns `NodeResult::Interrupt`; `docs/modules/graph/interrupts.md`, `compiled/executor.rs::resume` | TinyAgents interrupt is a return value, not a call; node still re-runs from start. `response_schema` absent. | -| `interrupt_before/after` compile options | Yes | No — documented in `interrupts.md`, no code hit for `interrupt_before` | Docs promise it; not implemented. | -| Durability `sync/async/exit` | Yes (default async) | Yes — `checkpoint/types.rs::DurabilityMode` (default Sync) | TinyAgents async mode has stricter failure semantics (fails run at next boundary). | -| Pending writes / skip completed tasks on resume | Yes | Yes — `checkpoint/types.rs::PendingWrite`, `merge_writes` | | -| Checkpoint namespaces for subgraphs | Yes (`ns|ns`, `node:task_id`) | Yes — `CheckpointConfig.namespace: Vec<String>` | | -| `get_state`, `get_state_history`, `update_state(as_node)`, fork | Yes | Yes — `compiled/state_api.rs` (`bulk_update_state`, `fork_state` extra) | | -| Checkpointer `prune`, `copy_thread`, `delete_for_runs` | Yes (checkpoint 4.1) | Yes — `checkpoint/mod.rs` (`prune`, `copy_thread`, `delete_by_run`) | | -| Checkpoint backends | Memory, SQLite, Postgres, Redis, Mongo… | Memory, File, SQLite — `checkpoint/{file,sqlite}.rs` | `docs/sdk-gaps/durability.md` §5: rusqlite version coupling. | -| Per-node `RetryPolicy` (list, `retry_on`) | Yes | Partial — graph-wide `with_node_retry` in `compiled/mod.rs`; no per-node | | -| Per-node `CachePolicy` + `BaseCache` (task-level cache) | Yes | No in graph; harness has `ResponseCache`/`CachePolicy` for model calls (`harness/src/cache/`) | | -| `defer=True` | Yes | Partial — `builder/mod.rs::mark_deferred` is export-only; real join is `add_waiting_edge` + `add_barrier_relief` | Waiting edges ≈ `NamedBarrierValue`; "run when nothing else is left" semantics absent. | -| Per-node `TimeoutPolicy(run, idle)`, `error_handler`/`NodeError` | Yes (1.2) | Partial — graph-wide `with_node_timeout`; no idle timeout, no node error handler | | -| `RunControl` graceful drain (`GraphDrained`) | Yes (1.2) | Partial — `CancellationToken` (hard cancel) in `harness/src/cancel/` | | -| `Runtime` (context, store, stream_writer, previous, execution_info) | Yes | Yes — `builder/types.rs::NodeContext`, harness `RunContext` | | -| Stream modes | values/updates/messages/custom/checkpoints/tasks/debug, `StreamPart{type,ns,data}` | Partial — `stream/types.rs::StreamMode` {Values, Updates, Messages, Debug, Interrupts, Custom}; typed `GraphEvent` enum instead of `StreamPart` | No `tasks`/`checkpoints` modes; comment says "StreamPart projection is future work". | -| `stream_events(version="v3")` protocol events | Yes (1.2) | No | Observability journal exists instead (`observability/`). | -| `BaseStore` namespaces, batch, TTL, `list_namespaces` | Yes | Yes — `harness/src/store/namespaced/` | Deliberately modelled on LangGraph's batch design. | -| Store semantic search (`IndexConfig{embed,dims,fields}`) | Yes | Partial — `SearchQuery.query` is a substring seam; embeddings exist in `retriever/` but are not wired to the store | | -| Functional API (`@entrypoint`/`@task`, `previous`) | Yes | No | | -| Double texting / multitask strategies | Platform only | Partial — `run_queue/` lanes (steer/follow-up), `thread_locks.rs` | Neither library has it in OSS; TinyAgents' queue is closer than LangGraph OSS. | -| Thread TTL, cron, background runs | Platform only | Partial — store TTL yes; `orchestration/` detached tasks + `JsonlTaskStore`; no cron | OpenHuman has `cron/`. | -| `create_agent` factory | Yes | Yes — `harness/src/runtime/mod.rs::AgentHarness` | | -| Middleware node hooks returning state updates / `jump_to` | Yes | No — `middleware/types.rs` hooks return `Result<()>`; `MiddlewareModelOutcome` has one variant | `sdk-gaps/orchestration.md` §13 already lists it. | -| Middleware `wrap_model_call`/`wrap_tool_call` nesting | Yes | Yes — `ModelMiddleware::wrap_model`, `ToolMiddleware::wrap_tool` | Same onion order. | -| Middleware extends state schema / contributes tools / stream transformers | Yes | No | Tools are registered on the harness, not by middleware. | -| Dynamic prompt | Yes | Yes — `DynamicPromptMiddleware` | | -| HITL middleware (durable `interrupt`, approve/edit/reject/respond, `when`) | Yes | Partial — `HumanApprovalMiddleware` is `Fn(&ToolCall) -> bool` (`library/types.rs:375`); durable pause only via graph `Interrupt` or `SteeringCommand::Pause` | | -| Summarization middleware (fraction/tokens/messages triggers) | Yes | Yes — `summarization/`, `ContextCompressionMiddleware`, `MessageTrimMiddleware` | | -| Context editing (`ClearToolUsesEdit`) | Yes | Partial — `MicrocompactMiddleware` (`library/context.rs`) | No tool-input clearing config. | -| PII middleware | Yes | Partial — `RedactionMiddleware`, `RedactingSink` | No detector/strategy matrix. | -| Model/tool call limits, model retry/fallback, tool retry, rate limit | Yes | Yes — `BudgetMiddleware`, `RetryMiddleware`, `ModelFallbackMiddleware`, `RateLimitMiddleware`, `limits/` | | -| LLM tool selector / provider tool search | Yes | Yes/No — `DynamicToolSelectionMiddleware`, `ContextualToolSelectionMiddleware` (`tool/select/`); no provider-side tool search | | -| Todo list middleware | Yes (opt-in) | Yes — `graph/src/todos/` (`TodoList`, same shape) | | -| Shell / file-search middleware | Yes | No in harness (`tools/` has `time.rs` only); `workspace/` gives roots | OpenHuman owns tools. | -| `ToolRuntime` injection (state, store, stream_writer, tool_call_id) | Yes | Partial — `ToolExecutionContext` (run/thread/depth/events/cancel/workspace); no state/store/tool_call_id | `tool/injected.rs` exists for hidden args. | -| Tool returns `Command` (state update + routing) | Yes | No — `ToolResult{content,is_error}` (vendor `tinytools`) | | -| `return_direct`, `ToolMessage.artifact` | Yes | No / Partial — no early-exit flag; `artifacts/` + `handoff.rs` offload large results instead | `sdk-gaps/orchestration.md` §13 "early-exit tools". | -| `ToolNode.handle_tool_errors` matrix | Yes | Partial — tool errors are recoverable results; unknown tool aborts (`sdk-gaps/tools.md` §2) | | -| Structured output strategies | Tool/Provider/Auto, unions, `handle_errors` | Yes — `structured/types.rs::StructuredStrategy`, `StructuredOutcome`, `repair.rs` | No union-to-many-tools; no auto-select from profile. | -| Standard content blocks | Yes | Yes — vendor `tinyinference/src/message/types.rs::ContentBlock` (Text/Json/Image/Thinking/Redacted…) | No citations/server_tool_call blocks. | -| `init_chat_model`, model profiles | Yes | Yes — `model_registry/`, `ModelProfile` | | -| Prompt caching middleware | Yes (Anthropic/Bedrock) | Yes — `PromptCacheGuardMiddleware`, `cache/layout.rs` | | -| MCP adapter | Yes (1.4) | Partial — only inside `providers/claude_code/` | | -| Supervisor / swarm / handoffs | Legacy libs; docs pattern | Yes — `subagent_node/`, `delegation/`, `orchestration/`, `parallel::map_reduce`, `SteeringRegistry` | TinyAgents is richer here. | -| Deep Agents: FS backends, subagent `task` tool, memory/skills files, permissions, sandboxes | Yes (harness) | No in TinyAgents; OpenHuman has `skills/`, `memory/`, `sandbox/`, `security/approval` | Correct layering already. | -| Dangling tool-call patch | Yes (deepagents) | Partial — `summarization/trim.rs` orphan handling | | -| Test fakes / trajectory eval | `GenericFakeChatModel`, `agentevals` | Yes — `testkit/` (trajectory assertions), `graph/testkit/` | | -| Tracing hooks (`trace_policy`, LangSmith) | Yes | Yes — Langfuse exporter, `TracingMiddleware`, `RedactingSink` | No per-node trace policy. | - -## 3. Features TinyAgents lacks, ranked by value - -### 3.1 Middleware control outcomes (`jump_to`, state updates, `Command` from tools) -What: LangChain node hooks return a state dict and may set `jump_to: "model"|"tools"|"end"` (declared -via `@hook_config(can_jump_to=[...])`); `wrap_model_call` may return `ExtendedModelResponse(model_response, -command)`; `wrap_tool_call` and tools may return `Command(update=..., goto=...)`. -```python -class ModelCallLimitMiddleware(AgentMiddleware): - @hook_config(can_jump_to=["end"]) - def before_model(self, state, runtime): - if state["model_call_count"] >= self.limit: - return {"jump_to": "end", "messages": [AIMessage("limit reached")]} -``` -Why: this is what makes limits, HITL, budget stops, early-exit tools and fallback re-routing -composable without side channels. `docs/sdk-gaps/orchestration.md` §13 already asks for it. -Mapping: add a `MiddlewareControl` enum returned from `before_model`/`after_model`/`before_tool`/ -`after_tool` (`Continue | JumpTo(LoopTarget) | StopWith(AgentRun) | Interrupt(Interrupt)`), let -`MiddlewareModelOutcome`/`MiddlewareToolOutcome` (already `#[non_exhaustive]`) gain `Command` -variants carrying `graph::Command<Update>`-style updates, and let `ToolResult` carry an optional -`ToolCommand { state_update: Value, goto: Option<LoopTarget>, return_direct: bool }` so -`agent_loop/tools.rs` can honour it. Precedence rule: first control outcome wins in hook order -(LangChain accumulates commands inner-first; keep it simple and documented). - -### 3.2 Durable human-in-the-loop in the harness (`HumanInTheLoopMiddleware`) -What: `after_model` issues one `interrupt(HITLRequest{action_requests, review_configs})` for the -whole tool-call batch; resume with `Command(resume={"decisions":[{"type":"approve"} | -{"type":"edit","edited_action":{name,args}} | {"type":"reject","message"} | -{"type":"respond","message"}]})`; `InterruptOnConfig{allowed_decisions, description, args_schema, -when}`; requires checkpointer + thread_id. -Why: TinyAgents' `HumanApprovalMiddleware` is an in-process `Fn(&ToolCall) -> bool`; a desktop -assistant needs approval to survive process restart and to edit args. Requires 3.1. -Mapping: `HumanApprovalMiddleware` returns `MiddlewareControl::Interrupt(Interrupt{payload: -HitlRequest})`; the agent loop persists a `harness` checkpoint (via `tinyagents-session` run ledger -or a graph `Checkpointer` when the loop runs as a graph node) and `AgentHarness::resume(thread, -HitlResponse)` applies decisions in `before_tool`. OpenHuman's `security/approval` becomes the UI. - -### 3.3 Task-level cache and per-node retry/timeout/error handler -What: `add_node(name, fn, retry_policy=RetryPolicy(max_attempts, backoff_factor, retry_on), -cache_policy=CachePolicy(key_func, ttl), timeout=TimeoutPolicy(run_timeout, idle_timeout), -error_handler=fn(state, NodeError) -> Command|None, defer=True)`; `set_node_defaults(...)`. -Cache key = `(namespace, xxh3(key_func(input)))`, backends `InMemoryCache`, `SqliteCache`; cached -tasks stream with `cached=True`. `error_handler` runs after retries are exhausted and is durable. -Why: deterministic replays, cheap re-runs of expensive nodes during development, per-node -resilience without wrapping every handler. -Mapping: `GraphBuilder::with_node_policy(node, NodePolicy{retry: Option<RetryPolicy>, timeout, -idle_timeout, cache: Option<NodeCachePolicy>, on_error: Option<ErrorHandler<State,Update>>})`; -reuse `tinyagents_harness::retry::RetryPolicy` and `harness/src/cache` traits (`ResponseCache` -already keys by hash; add a `TaskCache` trait keyed by `(graph_id, node_id, hash(send_arg|state -projection))`). Emit `TaskCompleted{cached: true}` in `GraphEvent`. - -### 3.4 `DeltaChannel` / writes-only checkpoint history and `Overwrite` -What: `Annotated[list, DeltaChannel(reducer, snapshot_frequency=50)]` stores only per-step writes; -`BaseCheckpointSaver.get_delta_channel_history` replays them; `Overwrite(value)` resets the -baseline. LangGraph shipped it in 1.2 and spent 1.2.5–1.2.11 fixing `update_state`/round-trip bugs. -Why: TinyAgents writes the full `state: State` every superstep; for message-heavy long threads -checkpoint volume is O(N²). Deepagents adopted it precisely for `messages`. -Mapping: add `Checkpoint.channel_deltas: Option<BTreeMap<String, Vec<ChannelUpdate>>>` with -`Checkpointer::delta_history(config, channel)`; make `Messages`/`Topic` channels opt into delta -persistence via `ChannelSet::with_delta(name, snapshot_every)`; add `ChannelUpdate::overwrite(v)`. -Learn from LangGraph's bugs: keep `update` and `replay` on one code path. - -### 3.5 Unified stream parts + `tasks`/`checkpoints` modes, `stream_events v3` -What: `stream(version="v2")` yields `StreamPart{type, ns: tuple, data}` for -`values|updates|messages|custom|checkpoints|tasks|debug`; `subgraphs=True` fills `ns` with -`("node:task_id", ...)`; `stream_events(version="v3")` gives `ProtocolEvent{seq, method, params}` -with projections (`stream.messages`, `stream.tool_calls`, `stream.subagents`). -Why: UIs need one cursor over nested runs. `sdk-gaps/streaming.md` §3/§6 ask for the same (reasoning/tool-arg -deltas, late-attach replay). -Mapping: TinyAgents already has `GraphEvent` + `GraphObservation` journal; add `ns: Vec<String>` -and `seq` to the envelope, add `TaskStarted/TaskResult` and `CheckpointSaved` projections as -`StreamMode::Tasks|Checkpoints`, and a `StreamProjection` adapter that folds `GraphEvent` + -harness `AgentEvent` into `stream.messages`/`stream.tool_calls` views. - -### 3.6 `ToolRuntime` parity and `return_direct` -What: `ToolRuntime{state, context, config, stream_writer, tool_call_id, store, tools, -execution_info}` injected into tools by parameter type; `@tool(return_direct=True)` exits the loop -when all executed tools are return_direct; `response_format="content_and_artifact"` puts a -non-model payload in `ToolMessage.artifact`. -Mapping: extend `ToolExecutionContext` with `call_id`, `store: Option<Arc<dyn NamespacedStore>>`, -`state_view: Option<Arc<dyn Any>>` (or a typed `StateHandle<State>`), and `stream: EventSink` -custom-write helper; add `ToolSchema.return_direct: bool` and `ToolResult.artifact: Option<Value>` -(vendor `tinytools` change). Pair with 3.1 for the loop exit. - -### 3.7 Semantic search on the namespaced store -What: `InMemoryStore(index={"embed": embeddings, "dims": 1536, "fields": ["text", "$"]})`; -`store.search(ns, query="...")` ranks by cosine; `put(..., index=False)` opts out. -Mapping: `NamespacedStore::search` already has `query`; add `IndexConfig{embedder: Arc<dyn -Embeddings>, dims, fields}` to `InMemoryNamespacedStore` using `tinyinference-embeddings`, and a -`VectorNamespacedStore` adapter over the existing `retriever/` vector store trait. - -### 3.8 `interrupt_before/after`, `response_schema`, graceful drain -Small items: compile-time `interrupt_before/after` selectors are documented in -`docs/modules/graph/interrupts.md` but absent from code; `interrupt(value, response_schema=)` -validates the resume payload; `RunControl.request_drain(reason)` lets a node finish the current -superstep and stop cleanly (`GraphDrained`). Mapping: `GraphBuilder::interrupt_before(nodes)`, -`Interrupt.response_schema: Option<Value>` validated in `resume`, and a `Drain` variant on -`SteeringCommand` that the executor honours at the boundary. - -### 3.9 Functional API -`@entrypoint(checkpointer)` + `@task` turn ordinary code into a one-node Pregel where each task -result is a `RETURN` pending write, so resume replays completed tasks without re-running them. -Value for TinyAgents is moderate (Rust closures are already ergonomic), but a `durable_task(ctx, -key, async fn)` helper inside a node handler — memoising by `PendingWrite` — would give the same -"side effects before an interrupt are not repeated" guarantee the interrupts doc currently pushes -onto users ("must be guarded by idempotency keys"). - -### 3.10 MCP adapter, provider tool search, node `trace_policy` -`langchain.mcp.MCPAdapter` (1.4) lists/executes MCP tools as `BaseTool`s; `ProviderToolSearchMiddleware` -defers large tool catalogs to provider-side search; `add_node(trace_policy=TracePolicy(process_inputs, -process_outputs))` redacts per node. TinyAgents has MCP only inside the Claude Code provider and -redaction only at the sink; both are modest additions on existing traits. - -## 4. Design lessons - -**Channel semantics.** LangGraph schedules nodes from channel versions; edges are sugar over -`branch:to:*` channels. This gives uniform semantics for `Send`, deferred joins -(`NamedBarrierValueAfterFinish`) and subscriptions, but it makes "why did this node run?" hard to -explain and is the source of a long tail of `update_state`/`versions_seen` bugs. TinyAgents' -explicit active-set + waiting-edge model is easier to export (`export/`) and debug, at the cost of -special mechanisms (`add_barrier_relief`) where LangGraph gets them for free. Keep the explicit -model; borrow `defer` as a scheduling flag rather than a channel type. - -**Checkpoint format.** LangGraph's checkpoint is channel-level (`channel_values` + versions), so a -saver can be schema-agnostic and delta history is natural. TinyAgents persists a typed `State` -blob plus `pending_activations`/`barrier_arrivals`, which is simpler and type-safe but -all-or-nothing. TinyAgents is better on failure semantics: async durability surfaces write errors -at the next boundary and always syncs terminal/interrupt checkpoints; LangGraph's default `async` -mode can lose the last step on crash and the docs say so. - -**Interrupt resume.** Both re-run the node from the top. LangGraph's `interrupt()` is a call that -can appear anywhere (index-matched), which is ergonomic but fragile (docs list "don't wrap in -try/except", "don't reorder interrupts"). TinyAgents' `NodeResult::Interrupt` forces one interrupt -per node return, which is explicit and serialisable but makes multi-step approvals inside one node -awkward. A `ctx.interrupt(payload)` helper that reads `ctx.resume` by index would get both. - -**Streaming.** LangGraph went through three formats (tuples, `StreamPart` v2, protocol events -v3) because the first had no namespace or sequence. TinyAgents should add `ns`/`seq` to -`GraphEvent` now rather than later. Typed enums beat `dict` payloads; keep them. - -**Middleware composition.** LangChain's split of node hooks (return state, may jump) vs wrap hooks -(onion) is the same as TinyAgents', but LangChain lets middleware own state keys, tools and stream -transformers, which is what made summarization/todo/HITL/PII shippable as single classes. TinyAgents' -hooks that only return `Result<()>` push that logic into the loop or host adapters (`sdk-gaps/orchestration.md` -§13). LangChain's weakness: `state_schema` merging across middleware is untyped `TypedDict` -unioning; TinyAgents can do better with a typed `MiddlewareState` extension slot on `RunContext`. - -**Agent loop as a graph.** `create_agent` compiles the loop into a `StateGraph`, so every graph -feature (checkpoints, interrupts, streaming, time travel) applies to the agent loop with zero extra -code. TinyAgents runs the loop in `agent_loop/run_loop.rs` and separately offers `subagent_node`; -the harness README even describes the loop "as an explicit state machine when callers need -inspection, checkpointing, HITL". Finishing that (loop-as-`CompiledGraph`) would close 3.1/3.2 in -one move. - -**Where TinyAgents is ahead.** Steering commands with policy checks, detached task registry with -durable stores, parallel failure policies (`quorum`/`race`/`compare`), goal/task-board primitives, -budget middleware, prompt-cache layout guards, artifact offload, tool timeouts with grace, and a -first-class registry. LangGraph OSS has none of these; LangGraph *Platform* has -some (background runs, cron, multitask) as hosted services. - -## 5. Runtime-level vs harness-level split - -Calibration: `openhuman:crates/openhuman-core/src` (the OpenHuman desktop host) already owns `skills/`, `memory/`, `sandbox/`, -`cron/`, `hooks/`, `security/{approval,audit,bubblewrap}`, `agent/tools`, `agent/tool_policy.rs`, -`agent/orchestration/{worktree,spawn_parallel_graph,running_subagents}` and `agent/harness/{memory_context, -artifact_offload,tool_result_artifacts}`. Those are the OpenHuman analogues of Deep Agents. - -| Gap | Layer | Reason | -|---|---|---| -| 3.1 middleware control outcomes / tool `Command` | **Runtime (harness crate)** | Loop control vocabulary; hosts cannot add it from outside. | -| 3.2 durable HITL request/decision protocol | **Runtime** for the interrupt/resume/decision types; **OpenHuman** for the approval UI and `security/approval` policy | Mirrors LangChain: middleware in library, policy in app. | -| 3.3 per-node retry/cache/timeout/error handler | **Runtime (graph crate)** | Pure scheduler policy. | -| 3.4 delta checkpoint history / `Overwrite` | **Runtime (graph crate)** | Checkpoint format. | -| 3.5 unified stream parts, tasks/checkpoints modes | **Runtime** | Contract UIs depend on; OpenHuman keeps only format adapters (`sdk-gaps/streaming.md` §6). | -| 3.6 `ToolRuntime` parity, `return_direct`, artifact | **Runtime** (`ToolExecutionContext`, vendor `tinytools`) | Tool contract. | -| 3.7 semantic store search | **Runtime** (trait + in-memory impl); backend choice in OpenHuman `memory/` | | -| 3.8 `interrupt_before/after`, `response_schema`, drain | **Runtime** | | -| 3.9 durable task memoisation | **Runtime** | Idempotency guard belongs beside `PendingWrite`. | -| 3.10 MCP adapter | **Runtime** (generic `McpToolSource` in harness) — OpenHuman already has `mcp/` | Move generic bits down only if OpenHuman's is reusable. | -| Provider tool search, node trace policy | Runtime, low priority | | -| Deep Agents FS tools, backends, sandboxes, permissions | **OpenHuman** (`sandbox/`, `agent/tools`, `tool_policy.rs`) | LangChain also keeps these out of the runtime; only `WorkspaceIsolation` roots stay in TinyAgents. | -| AGENTS.md memory, SKILL.md skills, harness profiles | **OpenHuman** (`memory/`, `skills/`) | Prompt conventions over a file backend. | -| Summarization-with-offload thresholds, large-result eviction | Mechanism in runtime (`artifacts/`, `handoff.rs` exist); thresholds/prompts in OpenHuman | Already split this way. | -| Subagent `task` tool, fork mode | Runtime has `SubAgentTool`/`subagent_node`; prompts and delegation policy in OpenHuman | Already split. | -| Double texting, cron, thread TTL, background runs | Runtime has `run_queue`, `DetachedTaskRegistry`, store TTL; scheduler/cron and UI in OpenHuman (`cron/`) | LangGraph puts these in the paid server. | -| Rubric grader, LLM tool emulator | OpenHuman / test tooling | | - -## 6. Sources - -- LangGraph releases: https://github.com/langchain-ai/langgraph/releases ; changelog https://docs.langchain.com/oss/python/langgraph/changelog-py -- Pregel internals: https://raw.githubusercontent.com/langchain-ai/langgraph/main/libs/langgraph/langgraph/pregel/_algo.py , `_loop.py`, `_runner.py`, `_retry.py`; https://docs.langchain.com/oss/python/langgraph/pregel -- Types (`Command`, `Send`, `Overwrite`, `interrupt`, `RetryPolicy`, `CachePolicy`, `TimeoutPolicy`): https://raw.githubusercontent.com/langchain-ai/langgraph/main/libs/langgraph/langgraph/types.py ; `graph/state.py`; `runtime.py`; `func/__init__.py` -- Checkpoint base: https://raw.githubusercontent.com/langchain-ai/langgraph/main/libs/checkpoint/langgraph/checkpoint/base/__init__.py ; https://docs.langchain.com/oss/python/langgraph/checkpointers ; https://docs.langchain.com/oss/python/langgraph/use-time-travel -- Interrupts: https://docs.langchain.com/oss/python/langgraph/interrupts ; fault tolerance: https://docs.langchain.com/oss/python/langgraph/fault-tolerance -- Streaming: https://docs.langchain.com/oss/python/langgraph/streaming ; https://docs.langchain.com/oss/python/langgraph/event-streaming ; https://docs.langchain.com/oss/python/langchain/event-streaming -- Stores: https://docs.langchain.com/oss/python/langgraph/stores ; functional API: https://docs.langchain.com/oss/python/langgraph/functional-api -- DeltaChannel bugs: https://github.com/langchain-ai/langgraph/issues/8821 ; https://github.com/langchain-ai/langgraph/pull/8961 -- Platform-only: https://raw.githubusercontent.com/langchain-ai/langgraph/main/libs/sdk-py/langgraph_sdk/schema.py ; https://docs.langchain.com/langgraph-platform/double-texting ; https://docs.langchain.com/langsmith/cron-jobs ; https://docs.langchain.com/langsmith/configure-ttl -- LangChain releases: https://github.com/langchain-ai/langchain/releases ; migration: https://docs.langchain.com/oss/python/migrate/langgraph-v1 -- `create_agent`: https://raw.githubusercontent.com/langchain-ai/langchain/master/libs/langchain_v1/langchain/agents/factory.py -- Middleware types and built-ins: https://raw.githubusercontent.com/langchain-ai/langchain/master/libs/langchain_v1/langchain/agents/middleware/types.py ; https://github.com/langchain-ai/langchain/tree/master/libs/langchain_v1/langchain/agents/middleware ; https://docs.langchain.com/oss/python/langchain/middleware/built-in ; https://docs.langchain.com/oss/python/langchain/middleware/custom -- Structured output: https://raw.githubusercontent.com/langchain-ai/langchain/master/libs/langchain_v1/langchain/agents/structured_output.py ; https://docs.langchain.com/oss/python/langchain/structured-output -- ToolNode / ToolRuntime: https://raw.githubusercontent.com/langchain-ai/langgraph/main/libs/prebuilt/langgraph/prebuilt/tool_node.py ; https://docs.langchain.com/oss/python/langchain/tools -- Messages/content blocks: https://docs.langchain.com/oss/python/langchain/messages ; prompt caching middleware: https://raw.githubusercontent.com/langchain-ai/langchain/master/libs/partners/anthropic/langchain_anthropic/middleware/prompt_caching.py -- Testing/eval: https://docs.langchain.com/oss/python/langchain/test/unit-testing ; https://docs.langchain.com/oss/python/langchain/test/evals -- Multi-agent: https://docs.langchain.com/oss/python/langchain/multi-agent (+ `subagents`, `handoffs`, `router`) ; https://docs.langchain.com/oss/python/migrate/langgraph-supervisor ; https://github.com/langchain-ai/langgraph-supervisor-py ; https://github.com/langchain-ai/langgraph-swarm-py ; https://docs.langchain.com/oss/python/langgraph/use-subgraphs ; https://docs.langchain.com/oss/python/langgraph/graph-api -- Deep Agents: https://github.com/langchain-ai/deepagents (libs/deepagents/deepagents/graph.py, middleware/, backends/, ARCHITECTURE.md) ; https://docs.langchain.com/oss/python/deepagents/human-in-the-loop ; PyPI `deepagents`, `deepagents-code`, `deepagents-acp`, `deepagents-talon` -- TinyAgents files checked: `docs/spec/README.md`, `docs/sdk-gaps/README.md`, `ROADMAP.md`, `docs/modules/{harness,graph}/README.md`, `docs/modules/graph/interrupts.md`, `crates/tinyagents-{harness,graph}/src/lib.rs`, `graph/src/{command,checkpoint,stream,builder,channel,compiled}/`, `harness/src/{middleware,tool,structured,store/namespaced,cache,artifacts,handoff.rs,steering}`, `vendor/tinytools/crates/tinytools/src/result/types.rs`, `vendor/tinyinference/crates/tinyinference/src/message/types.rs` - -Unverified (flagged by the research agents): exact version that introduced `Overwrite` and -`interrupt(response_schema=)`; whether "enqueue" is the server's default multitask strategy; -HITL `interrupt_mode` (PR title only, absent from source); `Runnable.as_tool` as a recommended -agent-as-tool path; AgentCore sandbox package location; formal deprecation of `langgraph-swarm`. diff --git a/docs/runtime-comparison/pi.md b/docs/runtime-comparison/pi.md index e0e853f91..772a8d9e3 100644 --- a/docs/runtime-comparison/pi.md +++ b/docs/runtime-comparison/pi.md @@ -90,14 +90,14 @@ Harness/product layer: | Tool hooks / terminate | `beforeToolCall`/`afterToolCall`, `terminate` hint, `shouldStopAfterTurn`, `prepareNextTurn` | Partial: middleware `before_tool/after_tool/wrap` (`ta:.../middleware/types.rs:155-244`); no terminate/stop-after-turn hook (sdk-gaps §13) | | | `transformContext` / `convertToLlm` split | Yes | Partial: `before_model` middleware mutates request; no app-message vs LLM-message layer | | | Abort | `AbortSignal`, `stopReason:"aborted"`, `continue()` | Yes: `CancellationToken`, `AbortOnDrop` (`ti:.../model/types.rs:656`) | | -| Conversation tree / branching | JSONL v3 tree with `id/parentId`, labels, `/fork`, `/tree` (`pi:packages/coding-agent/src/core/session-manager.ts:53`); `AgentHarness` entry tree + Branches + `ForkOptions` | No for conversations: `tinyagents-session` is linear JSONL + SQLite history; only graph checkpoints fork (`ta:crates/tinyagents-graph/src/checkpoint/types.rs:24-35`) | See gap 3. | +| Conversation tree / branching | JSONL v3 tree with `id/parentId`, labels, `/fork`, `/tree` (`pi:packages/coding-agent/src/core/session-manager.ts:53`); `AgentHarness` entry tree + Branches + `ForkOptions` | No: `tinyagents-session` is linear JSONL + SQLite history | See gap 3. | | Compaction | cut point / `keepRecentTokens` / split turns / iterative summary / `CompactionEntry{firstKeptEntryId,tokensBefore}`; branch summaries | Partial: `SummarizationPolicy`, `Summarizer` trait, trim strategies, `ContextCompressionMiddleware`, `MicrocompactMiddleware` (`ta:.../summarization/`, `middleware/library/context.rs`) | TA has no durable compaction record or turn-boundary rules; no overflow detector (grep `overflow` empty). | -| Durable per-step run state | `AgentHarness` intent/settlement transactions, `replay: "never"\|"safe"` (`types.ts:422`) | Partial: graph checkpoints + pending writes; harness loop itself not durable per tool call | See gap 7. | +| Durable per-step run state | `AgentHarness` intent/settlement transactions, `replay: "never"\|"safe"` (`types.ts:422`) | No: the harness loop is not durable per tool call | See gap 7. | | Session backends + conformance | memory, JSONL, SQLite, `session/testing/conformance` | Partial: SQLite session DB + JSONL transcript; no backend conformance suite (sdk-gaps §17) | | | Extension API / event bus | 40+ product events, `registerTool/Command/Provider`, UI | Partial: middleware + registry; no plugin loader (host concern) | | | Skills / prompt templates | Yes (`Skill`, `PromptTemplate` types in agent-core `harness/types.ts`; loaders in coding-agent) | No skills (grep empty); `prompt/` module exists | | -| Sub-agents | Example extension only (`examples/extensions/subagent/`) | Yes: `SubAgent`, `SubAgentTool`, parallel policies | TA ahead. | -| Graph/workflow runtime | None | Yes (`tinyagents-graph`) | TA ahead. | +| Sub-agents | Example extension only (`examples/extensions/subagent/`) | Yes: async `DriverSubAgentTool`, job registry, steering, and detached settlement | TA ahead. | +| Graph/workflow runtime | None | None in the current TinyAgents workspace; removed in #61 | | | Per-tool timeouts / limits / budgets | None (signal only) | Yes: `ToolTimeout`, `RunLimits`, budget middleware | TA ahead. | | Telemetry | typed schema spans (`pi-telemetry`) | Yes: `AgentEvent` + Langfuse exporter | | | Remote session protocol | CBOR `pi-protocol` | No (host concern) | | @@ -233,14 +233,14 @@ transition; `AgentTool.replay: "never" | "safe"` (`types.ts:422`) decides whethe settle in completion order but materialize in source order; queued input is staged as `pi.pending.entry` before placement. Hooks are classified pass-local / request-local / transition-consumed. -Why: TinyAgents' graph has checkpoints and pending writes, but the harness loop (the thing most hosts run) -is not resumable mid-batch; `append_interrupted_partial` (`ta:.../transcript/writer.rs:175`) is the only -crash artefact. sdk-gaps §1 asks for `idempotency`/`retry` tool metadata; pi's `replay` is the minimal +Why: the harness loop is not resumable mid-batch; `append_interrupted_partial` +(`ta:.../transcript/writer.rs:175`) is the only crash artefact. sdk-gaps §1 +asks for `idempotency`/`retry` tool metadata; pi's `replay` is the minimal version of that. -Mapping: add `ToolReplay::{Never, Safe}` to `ToolSchema`/policy; model the loop as a `tinyagents-graph` -graph with per-tool-call task checkpoints (the graph already has task outcomes and pending writes), rather -than porting pi's bespoke state machine. +Mapping: add `ToolReplay::{Never, Safe}` to `ToolSchema`/policy and design +per-tool-call durable state before attempting to port pi's bespoke state +machine. The graph runtime was removed in TinyAgents #61. ### 4.8 Model catalog breadth, compat matrix, thinking-level map, auth resolution @@ -282,8 +282,8 @@ nodes, forks as path copies) is a clean, replay-only design. Its `AgentHarness` AgentLane, "context never reads past a compaction", append-only-context invariant for KV cache) is sharper than TinyAgents' linear session history. Worse: pi has two coexisting designs (coding-agent v3 JSONL and the harness entry tree, plus pico v5 in `pi-durable`) and an admitted O(history) copy on first -divergence of an uncompacted SQLite branch (§2.6). TinyAgents' graph checkpointer already has -fork/time-travel; the gap is applying the same idea to conversation history. +divergence of an uncompacted SQLite branch (§2.6). TinyAgents has no graph +checkpointer; durable conversation branching remains an open gap. **Compaction.** pi's is more complete operationally (cut-point rules, split turns, iterative summaries, overflow-triggered compaction with retry, hook can decline/replace, usage attributed). TinyAgents is @@ -310,9 +310,10 @@ and `ProviderRequestOptions.fetch` injection are cheap, useful escape hatches Ti **General.** pi optimizes for one product and moves fast (100 commits/week, version-locked monorepo); its runtime layers are extracted from that product and carry two generations of design at once. -TinyAgents is spec-first and layered (graph/harness/registry/session/orchestration with enforced -dependency direction). pi has no graph runtime, no sub-agent primitives, no budgets, no structured -output, and no permission model; TinyAgents has all of those and should not import pi's product coupling. +TinyAgents is spec-first and layered (harness/registry/session/orchestration with enforced +dependency direction). Neither current workspace has a graph runtime. pi has no +sub-agent primitives, budgets, structured output, or permission model; TinyAgents +has these and should not import pi's product coupling. ## 6. Runtime-level vs harness-level split @@ -324,7 +325,7 @@ output, and no permission model; TinyAgents has all of those and should not impo | 4.4 Steer/follow-up lanes | Yes: wire `RunQueue` into loop, queue modes, terminate hint | Which channel/message becomes a steer vs follow-up (`agent/harness/run_queue/` already exists in OpenHuman, so this moves down) | | 4.5 Compaction rules + overflow | Yes: cut points, overflow classifier, durable record | Summary prompt wording, file-tracking `details`, user-facing `/compact` | | 4.6 Handoff transform | Yes (pure function over messages + profile) | — | -| 4.7 Durable loop state | Yes, via graph-backed loop and `ToolReplay` metadata | Approval gate and sandbox decisions stay in `security/approval`, `security/bubblewrap.rs` | +| 4.7 Durable loop state | No: `ToolReplay` metadata and per-tool checkpoints remain open work | Approval gate and sandbox decisions stay in `security/approval`, `security/bubblewrap.rs` | | 4.8 Catalog/compat/auth | Catalog generator, compat data, `CredentialStore` trait, generic OAuth flow: runtime. Actual key storage, keychain, login UI: host | `security/credentials/`, model picker | | 4.9 Deferred, custom roles, skill types | Runtime types; deferred stream item | Skill/template discovery from disk, prompt assembly (`agent/prompts/`) | | Extension loader, commands, UI, packages | No | Yes (OpenHuman RPC/bus, `agent/registry`) | @@ -342,11 +343,10 @@ output, and no permission model; TinyAgents has all of those and should not impo - `pi:packages/coding-agent/docs/{session-format,compaction,extensions,skills,prompt-templates,packages,sdk,models}.md`, `src/core/{session-manager,agent-session}.ts`, `src/experimental/` - GitHub API: repo metadata, releases, commits, contributors (fetched 2026-09-19) -- TinyAgents: `ta:docs/spec/README.md`, `ta:docs/modules/harness/README.md`, `ta:docs/modules/graph/README.md`, +- TinyAgents: `ta:docs/spec/README.md`, `ta:docs/modules/harness/README.md`, `ta:docs/sdk-gaps/README.md`, `ta:ROADMAP.md`, `ta:crates/*/src/lib.rs`, - `ta:crates/tinyagents-harness/src/{agent_loop,steering,run_queue,summarization,middleware,events,stream,cache,cost}/`, + `ta:crates/tinyagents-harness/src/{agent_loop,steering,run_queue,summarization,middleware,events,cache,cost}/`, `ta:crates/tinyagents-session/src/transcript/`, `ta:crates/tinyagents-registry/src/catalog.rs`, - `ta:crates/tinyagents-graph/src/checkpoint/types.rs` - tinyinference (main checkout): `ti:crates/tinyinference-llm/src/{message,model,catalog,usage}/types.rs`, `ti:crates/tinyinference-llm/src/providers/`, `ti:crates/tinyinference-providers/src/oauth.rs` - OpenHuman calibration: `openhuman/crates/openhuman-core/src/agent/README.md`, diff --git a/docs/runtime-comparison/plan.md b/docs/runtime-comparison/plan.md index 7c623c851..1a77d0731 100644 --- a/docs/runtime-comparison/plan.md +++ b/docs/runtime-comparison/plan.md @@ -5,11 +5,9 @@ format and contract changes (stream, checkpoint, error, tool result) once, early, and non-exhaustively so later work does not break downstream matches; prefer the smallest change that closes a documented promise over a new abstraction. Every item is one PR unless marked (multi-PR). Finding ids -(`H-C1`, `G-I3`, `W-I5`) refer to -[`code-review-harness.md`](code-review-harness.md), -[`code-review-graph.md`](code-review-graph.md) and -[`code-review-workspace.md`](code-review-workspace.md); gap ids (`A1`, -`D3`) refer to [`feature-gaps.md`](feature-gaps.md). +(`H-C1`, `G-I3`, `W-I5`) refer to [`code-review-harness.md`](code-review-harness.md); +graph and workspace review documents were removed in TinyAgents #61. Gap ids +(`A1`, `D3`) refer to [`feature-gaps.md`](feature-gaps.md). ## Execution log @@ -49,9 +47,9 @@ were all merged into `runtime-comparison`. The loop-as-graph piece (A5) landed only partially — see (b). - **Phase 6** (tool ecosystem) — landed. The `ToolSet` trait plus `Combined`/`Filtered`/`Prefixed`/`Renamed`/`Prepared`/`ApprovalRequired`/ - `ExternalToolSet` (B3), transcript-carried tool add/remove (B6), and the - `Capability` bundle (`instructions`/`toolset`/`middleware`/`model_defaults`/ - `exposure`/`defer_loading`) (G3) all shipped. The MCP client crate (B4) and + `ExternalToolSet` (B3) and transcript-carried tool add/remove (B6) shipped. + The former G3 capability bundle was removed in TinyAgents #62; the registry + retains only its stored data type. The MCP client crate (B4) and provider-executed tool content parts (B7) are OpenHuman, out of scope. - **Phase 7** (models, providers, evals) — landed partially. `ModelProfile` behaviour fields including `thinking_tags` (F1), the models.dev catalog @@ -74,18 +72,15 @@ were all merged into `runtime-comparison`. - **`Turn` object** (H-R1, "built once per turn, cache the tools fingerprint"): skipped. No `struct Turn` exists anywhere in `crates/`; each phase still re-derives what it needs from `RunContext` directly. -- **`GraphLoopDriver` is not itself checkpointable**: by design, per its own - doc comment in `crates/tinyagents-graph/src/agent_loop/mod.rs` — it drives - the same phase node bodies without a checkpointer so it can be a drop-in - `LoopDriver`. `compile_loop`/`LoopIter` are the real `CompiledGraph` path - and have interrupt → checkpoint → resume tests (`loop_as_graph.rs`). +- **Compiled graph loop (A5) descoped in TinyAgents #61**: the graph crate and + `GraphLoopDriver` were removed. The harness now drives runs through its + direct loop; graph-backed checkpointing and iteration are not available. - **Lease renewal loop**: exists at the orchestration layer only. `WorkflowEngine` has a real heartbeat that renews a short lease while a child run is in flight (`crates/tinyagents-orchestration/src/workflow/ - tests.rs::heartbeat_renews_a_short_lease_while_a_child_is_running`), but - the compiled-graph executor's per-thread lease (`ThreadLockMap` + - `Checkpointer::try_claim`) is claim-once with a fixed TTL and release-at- - end — no periodic renewal during a long-running graph execution. + tests.rs::heartbeat_renews_a_short_lease_while_a_child_is_running`). The + comparison with a compiled-graph executor described the pre-#61 runtime and + is no longer applicable. - **`ToolProgress` still unwired**: `docs/sdk-gaps/streaming.md` states verbatim that "`run_on_tool_delta`/`ToolProgress` still have no real caller; that needs a tinytools progress seam" — confirmed by grep: no @@ -210,17 +205,17 @@ Builds the vocabulary that A1–A4 and B1–B2 share; everything later | Per-node `NodePolicy{retry, timeout, idle_timeout, cache, on_error}` reusing `harness::retry::RetryPolicy`; `TaskCache` trait keyed by `(graph_id, node_id, hash)`; `defer` as a scheduling flag; `TaskCompleted{cached: true}` | M | D2 | | `interrupt_before/after`, `Interrupt.response_schema`, `SteeringCommand::Drain` honoured at the boundary | S | A7 | | `durable_task(ctx, key, fut)` memoised through `PendingWrite` | S | D5 | -| Lower `WorkflowDefinition` to a `CompiledGraph` behind a feature flag; keep `WorkflowStore` as the status projection and the lease as the durable lock; `workflow/tests.rs` as acceptance | L | G-I12, G-R6 | +| WorkflowDefinition-to-graph execution | descoped in TinyAgents #61 with graph runtime removal; requires a replacement runtime design | L | G-I12, G-R6 | | `NodeHandler` takes `Arc<State>`; `Arc<Value>` send args | M | G-M2 | -## Phase 5: sessions, context, and the loop as a graph (3–4 weeks) +## Phase 5: sessions, context, and loop hardening (3–4 weeks) | Item | Effort | Refs | OpenHuman hand-off | |---|---|---|---| | Compaction rules: `find_cut_point`, split turns, iterative summaries, `CompactionRecord`, `OverflowClassifier`, `overflow → compact → retry` in `ContextCompressionMiddleware`; `before_compaction` hook may decline/replace | M | E2 | summary prompt text and `/compact` stay | | `AssistantMessage.origin{provider, api, model}`; `prepare_for_model(&[Message], &ModelProfile)` handoff pass | S | E3 | mid-session model switch just works | | Tool-effect ledger keyed by `CallId` with `ToolReplay::{Never, Safe}` on `ToolSchema`/policy; `list_unresolved_tool_effects(run_id)` | M | B5 | OpenHuman marks its tools' replay class | -| Agent loop as a `CompiledGraph` (`plan → model → tools → settle` nodes) so checkpoints, interrupts and time travel apply to the loop; `AgentHarness::iter()` | L | A5 | `agent_graph.rs` collapses onto the SDK graph | +| Checkpointed agent loop and `AgentHarness::iter()` | descoped in TinyAgents #61: no graph runtime is shipped; revisit only with a replacement design | A5 | no current OpenHuman hand-off | | `sanitize_history()` helper; `Message::Custom` | S | E4, E5 | web_chat sanitises via the SDK | ## Phase 6: tool ecosystem (2–3 weeks) @@ -230,7 +225,6 @@ Builds the vocabulary that A1–A4 and B1–B2 share; everything later | `ToolSet<State,Ctx>` trait with `Combined`, `Filtered`, `Prefixed`, `Renamed`, `Prepared`, `ApprovalRequired`, `External`; `ToolRegistry` becomes one `ToolSet`; per-tool `prepare` | M | B3 | `agent/tool_policy.rs` and `tinyagents/middleware` tool filtering shrink to predicates | | `tinytools-mcp` crate (stdio / HTTP / SSE, prefixes, `process_tool_call`, sampling, elicitation, config loading) implementing `ToolSet` | L | B4 | `mcp/` keeps server config, auth and UI; evaluate reusing its transport | | Transcript-carried `SystemMessage{sections, tools_added, tools_removed}`, `replay_system_state`, `declare_tool_changes` before `ModelStarted`, `ModelProfile.mid_conversation_system_messages` | M | B6 | tool loadout changes stop busting the cache | -| `Capability{instructions, tools, middleware, model_defaults, exposure}` bundle in `tinyagents-registry`, referenced by `AgentDefinition`, `defer_loading` | M | G3 | `skills/` discovery produces `Capability` values | | Provider-executed tool parts in `ContentBlock` | S | B7 | — | ## Phase 7: models, providers, evals (2–3 weeks) diff --git a/docs/runtime-comparison/pydantic-ai.md b/docs/runtime-comparison/pydantic-ai.md index 0db7550a4..1788a91c1 100644 --- a/docs/runtime-comparison/pydantic-ai.md +++ b/docs/runtime-comparison/pydantic-ai.md @@ -55,8 +55,7 @@ and `tinytools` at `a14e24d`. ## 2. Feature inventory Paths are relative to this repository unless prefixed `vendor:` -(`vendor/tinyinference` / `vendor/tinytools`). H = `crates/tinyagents-harness/src`, -G = `crates/tinyagents-graph/src`. +(`vendor/tinyinference` / `vendor/tinytools`). H = `crates/tinyagents-harness/src`. The former graph crate was removed in TinyAgents #61; the current workspace has no state-graph runtime. | Feature | Pydantic AI | TinyAgents | Notes | |---|---|---|---| @@ -76,11 +75,11 @@ G = `crates/tinyagents-graph/src`. | Sequential vs parallel tool execution | default concurrent; `sequential=True` per tool; `parallel_tool_call_execution_mode('sequential')` | **Yes** — `is_concurrency_safe()` default `false` (`vendor:tinytools/src/tool/types.rs`), serial admission/concurrent fold (`H/agent_loop/tools.rs`) | Opposite default (TinyAgents fails safe). | | Tool timeouts | `timeout=` per tool, `tool_timeout=` agent-wide; timeout -> retry prompt | **Yes** — `ToolTimeout::{Inherit,Millis,Unbounded}`, `with_tool_timeout_settings` (`H/tool/timeout.rs`) | | | Native/provider tools (web search, code exec, image gen) | `WebSearch`, `WebFetch`, `XSearch`, `ImageGeneration`, `MCP(native=True)`; `BuiltinToolCallPart` | **No** — no provider-executed tool parts in `ContentBlock` (`vendor:tinyinference-llm/src/message/types.rs:21`) | | -| On-demand capability loading / tool search | `defer_loading=True`, `load_capability` tool, `ToolSearch`, `ctx.loaded_capability_ids` reconstructed from history | **Partial** — `ToolExposure::Deferred` + `H/tool/select` (keyword ranking) | No bundle (instructions+tools+settings) loading. | +| On-demand capability loading / tool search | `defer_loading=True`, `load_capability` tool, `ToolSearch`, `ctx.loaded_capability_ids` reconstructed from history | **No** — TinyAgents removed the harness capability loader; generic tool selection remains available. | | Model fallback | `FallbackModel(fallback_on=exceptions | response predicates)` | **Yes** — `FallbackPolicy` (`H/retry/types.rs:183`), `ModelFallbackMiddleware`, `CapabilitySet`-filtered resolution (`H/model_registry`) | TinyAgents resolution is richer (hints, capabilities). | | Wrapper/decorator models | `WrapperModel`, `InstrumentedModel`, `ConcurrencyLimitedModel` | **Yes** — `ProfileOverrideModel`, `MaxTokensModel`, `RouteRecordingModel`, `ObservingModel` (`vendor:tinyinference-llm/src/model/decorators.rs`), `RateLimiter` | | | Model profiles | `ModelProfile` TypedDict: `json_schema_transformer`, `default_structured_output_mode`, `prompted_output_template`, `thinking_tags`, `supported_native_tools`, `tool_deferral_mode`, `context_window`... | **Partial** — `ModelProfile` (`vendor:tinyinference-llm/src/model/types.rs:193`): modalities, tool calling, streaming chunks, native structured output, reasoning, token windows | No schema transformer, no thinking-tag parsing, no prompted-output template. | -| Model settings precedence (model < agent < run), callable per step | Yes | **Yes** — `ModelRequestDefaults`, request overrides (`H/model_registry`) | | +| Model settings precedence (model < agent < run), callable per step | Yes | **Partial** — per-request overrides remain; named harness model selection was removed. | | Usage accounting + cost | `RunUsage{requests,tool_calls,input/output,cache_read/write,audio, cost: Decimal|None, details}`; cost via `genai-prices` | **Yes** — `Usage`, `UsageTotals`, `ModelPricing`, `CostTotals` (`H/cost/types.rs`), catalog prices (`vendor:tinyinference-llm/src/catalog/types.rs:20`) | Pydantic's price DB is an external, maintained package. | | Usage limits | `UsageLimits{request_limit=50, tool_calls_limit, input/output/total_tokens_limit, per_request_input_tokens_limit, cost_limit, count_tokens_before_request}` -> `UsageLimitExceeded` | **Yes** — `RunLimits`, `LimitBehavior` (`H/limits/types.rs`), `BudgetMiddleware`/`BudgetLimits` (`H/middleware/library/budget.rs`) | | | Shared usage across delegated agents | `delegate.run(..., usage=ctx.usage)` | **Yes** — parent/child usage roll-up in `subagent`, `RunTree` (`G/recursion`) | | @@ -90,18 +89,18 @@ G = `crates/tinyagents-graph/src`. | Pending message queue / mid-run steering | `ctx.enqueue(msg, priority='asap'|'when_idle')`, `agent_run.enqueue` | **Yes** — `RunQueue` lanes `Steer/Followup/Collect` (`H/run_queue/types.rs`), `steering/` | TinyAgents is broader (policy-checked commands). | | Cancellation | `ctx.cancel()`, `CancellationToken`, `RunCancelled` | **Yes** — `CancellationToken` (`H/cancel`) | | | Lifecycle hooks | `Hooks`: before/after/wrap/error for run, node, model_request, tool_validate, tool_execute, output_validate, output_process; `SkipModelRequest`, `SkipToolExecution`; per-tool filter; timeouts | **Yes** — `Middleware` before/after model+tool, `ModelMiddleware`/`ToolMiddleware` wrap, `MiddlewareControl::{StopWithFinal, Interrupt}` (`H/middleware/types.rs`) | Pydantic has no "wrap node"/"output validate" equivalents in TinyAgents; TinyAgents has no output-validate hook. | -| Node-level iteration (`agent.iter()` / `next(node)`) | Yes; override next node | **Partial** — graph `run/resume/retry/resume_from` (`G/compiled/executor.rs`), `GraphEvent` stream, `get_state_history`/`update_state`/`fork_state` | No harness-level step iterator; agent loop is opaque. | +| Node-level iteration (`agent.iter()` / `next(node)`) | Yes; override next node | **No** — graph runtime removed in TinyAgents #61; the direct agent loop has no step iterator | The agent loop is opaque. | | Streaming | `run_stream` (`stream_text/stream_output`, partial-validated snapshots), `run_stream_events` (`PartStart/Delta/End`, tool events, `AgentRunResultEvent`) | **Yes** — `ModelDelta{content, reasoning, tool_call}`, `invoke_agent_stream`, `AgentEvent` (`H/stream`, `H/events`) | No partial-structured-output snapshots. | | Thinking/reasoning parts | `ThinkingPart`, `Thinking(effort=)` capability, `thinking_tags` parsing | **Yes** — `ContentBlock::Thinking{signature}`, `RedactedThinking`, `ReasoningConfig` | | | Multimodal input | `ImageUrl`, `AudioUrl`, `VideoUrl`, `DocumentUrl`, `BinaryContent`, `UploadedFile`, `force_download`, SSRF guard | **Partial** — `ContentBlock::Image(ImageRef)` only; `multimodal/` handles data URIs and `FilePayload` (feature-gated) | No audio/video/document blocks in the message model. | -| Testing | `TestModel` (schema-driven auto tool calls), `FunctionModel`, `Agent.override(model/deps/toolsets/capabilities)`, `ALLOW_MODEL_REQUESTS`, `capture_run_messages` | **Yes** — `ScriptedModel`, `StreamingMock`, `SlowModel`, `FakeTool`, `EventRecorder`, `Trajectory` (`H/testkit`), graph testkit | No schema-driven auto-responder; no global network kill-switch. | +| Testing | `TestModel` (schema-driven auto tool calls), `FunctionModel`, `Agent.override(model/deps/toolsets/capabilities)`, `ALLOW_MODEL_REQUESTS`, `capture_run_messages` | **Yes** — `ScriptedModel`, `StreamingMock`, `SlowModel`, `FakeTool`, `EventRecorder`, `Trajectory` (`H/testkit`) | No schema-driven auto-responder; no global network kill-switch. | | Evals | `pydantic_evals`: `Dataset`/`Case`/`Evaluator`, `LLMJudge`, span-based evaluators, YAML datasets, Logfire | **No** — no eval crate | | | Observability | `Instrumentation` capability, OTel GenAI semconv v5, Logfire | **Partial** — Langfuse client/exporter, event journals, `RedactingSink`; no OTel `gen_ai.*` (grep: none) | | | MCP | `MCP` capability / `MCPToolset` (stdio, streamable HTTP, SSE), prefixes, `process_tool_call`, sampling, elicitation, prompts/resources, `load_mcp_toolsets(config.json)` | **No** — MCP appears only inside the Claude Code provider bridge (`H/providers/claude_code`) | `ToolResult` is "MCP-shaped" but no client. | -| Durable execution | Temporal/DBOS/Prefect/Restate/Kitaru/Airflow/Lambda capabilities; deterministic replay; activities for model/tool/MCP | **Yes (different model)** — built-in `Checkpointer` (`InMemory`, `File`, `Sqlite`), `DurabilityMode`, pending writes, `resume`, time travel (`G/checkpoint`) | See §4. | -| Run-state persistence (message-history checkpoints) | Harness `StepPersistence`: `InMemory/File/Sqlite/Mongo` step stores, `continue_run`, `fork_run`, tool-effect ledger `(run_id, tool_call_id)` with `started/completed/failed`, `MediaStore` externalization | **Partial** — `tinyagents-session` (history/run ledger), harness `store/`, graph checkpoints | No tool-effect ledger for crash-time "did this side effect happen". | -| Graph library | `pydantic_graph`: `BaseNode`/`End`, builder `@g.step`, `g.join(reducer)`, `edge_from().map().to()`, broadcast, `transform`, `graph.render()` mermaid; persistence removed in v2 | **Yes, richer** — `GraphBuilder`, channels/reducers (`LastValue`, `Topic`, `BinaryAggregate`, `Barrier`, `NamedBarrier`), `Send`, `Command`, subgraphs, `map_reduce`, `to_mermaid` (`G/`) | | -| Sub-agents / delegation | tool-call delegation, `Subagents` (`delegate_task`), `Advisor`, `DynamicWorkflow` | **Yes** — `SubAgent`, `SubAgentTool`, `subagent_node`, `DetachedTaskRegistry`, orchestration tools, steering | TinyAgents is deeper (detached, wait, steer). | +| Durable execution | Temporal/DBOS/Prefect/Restate/Kitaru/Airflow/Lambda capabilities; deterministic replay; activities for model/tool/MCP | **No** — no graph checkpointer or durable graph execution runtime is shipped | See §4. | +| Run-state persistence (message-history checkpoints) | Harness `StepPersistence`: `InMemory/File/Sqlite/Mongo` step stores, `continue_run`, `fork_run`, tool-effect ledger `(run_id, tool_call_id)` with `started/completed/failed`, `MediaStore` externalization | **Partial** — `tinyagents-session` (history/run ledger), harness `store/`; no graph checkpoints | No tool-effect ledger for crash-time "did this side effect happen". | +| Graph library | `pydantic_graph`: `BaseNode`/`End`, builder `@g.step`, `g.join(reducer)`, `edge_from().map().to()`, broadcast, `transform`, `graph.render()` mermaid; persistence removed in v2 | **No** — the graph crate and its builder/runtime were removed in TinyAgents #61 | | +| Sub-agents / delegation | tool-call delegation, `Subagents` (`delegate_task`), `Advisor`, `DynamicWorkflow` | **Yes** — async `DriverSubAgentTool`, `SubAgentJobRegistry`, `DetachedTaskRegistry`, and steering | TinyAgents supports detached work, waiting, and steering. | | Declarative agent spec | `Agent.from_file('agent.yaml')`, `from_spec`, capability `from_spec` | **Yes** — `AgentDefinition` (`tinyagents-definition`) | | | Guardrails | Harness `InputGuardrail/OutputGuardrail/ToolGuardrail` with allow/block/replace/retry/approve; secret/PII detectors | **Partial** — `RedactionMiddleware`, `ToolPolicyMiddleware`, `RedactingSink` | No output-retry guardrail outcome. | | Coding-agent batteries | Harness `Coder`, `FileSystem`, `Shell`, `Planning`, `Memory`, `Skills`, `CodeMode` | **Partial** — `goals`, `todos` (task board), `workspace` isolation, `tools/time` | TinyAgents deliberately leaves file/shell tools to the host. | @@ -257,8 +256,8 @@ Snapshots are taken at settled `CallToolsNode` boundaries; each tool call has a `idempotency_key` / `effect_summary` so the orchestrator decides whether to `continue_run(include_interrupted=True)` or `fork_run`. Large/binary content is externalised to a `MediaStore` (`media+sha256://` URIs). **Mapping.** -TinyAgents' checkpointer already has pending writes; add a per-tool-call -effect row to `tinyagents-session`'s run ledger keyed by `CallId`, with +TinyAgents' session ledger is distinct from graph checkpointing; add a +per-tool-call effect row to `tinyagents-session`'s run ledger keyed by `CallId`, with `ToolPolicy.runtime.idempotent` feeding the decision. ### 3.12 OTel GenAI semantic conventions + UI adapters @@ -273,8 +272,8 @@ sanitisation belong to the host (§5). - **Typed deps.** Both inject a typed context; Pydantic passes `RunContext[Deps]` into tools, instructions, validators, `prepare`, and toolset factories with - one signature. TinyAgents splits `&State` (graph state) from `Ctx` - (runtime deps), which is the right split for a graph runtime but means + one signature. TinyAgents separates application `State` from `Ctx` + (runtime deps), which is useful for nested calls but means tools receive `ToolRunContext` (tinytools) rather than the harness `RunContext`; the injected-argument mechanism (`tool/injected.rs`) is a workaround Pydantic does not need. Worse in Pydantic: no separate durable @@ -305,12 +304,10 @@ sanitisation belong to the host (§5). - **Evals.** Pydantic ships them in-repo and wires them to spans; TinyAgents has nothing. Low cost to add a minimal crate. - **iter API.** `agent.iter()` exposing `UserPromptNode/ModelRequestNode/ - CallToolsNode` is the same idea as TinyAgents' explicit graph, but Pydantic - applies it to the *agent loop itself*, so any host can step, inspect, or - override the next node. TinyAgents' loop is a closed `run_loop.rs`; - exposing it as a `CompiledGraph` (the docs promise this in - `docs/modules/harness/state-graph.md`) would close the gap and reuse the - checkpointer for free. + CallToolsNode` lets a host step, inspect, or override the next node. + TinyAgents' direct loop remains closed in `run_loop.rs`; its former graph + runtime was removed in #61. A future replacement design could address this + gap, but no such runtime is currently shipped. - **Durable execution.** Pydantic explicitly refuses to own a checkpointer: "durability is not storage", and it delegates replay to Temporal/DBOS/etc. with the cost that deps, settings, and metadata must be Pydantic-serializable, @@ -352,12 +349,12 @@ sanitisation belong to the host (§5). | 3.8 prompted/union output, `EndStrategy` | **Runtime** | | | 3.9 audio/video/document blocks | **Runtime** (tinyinference) | OpenHuman's `agent/multimodal.rs` currently works around it with markers. | | 3.10 schema-driven test model, network kill-switch | **Runtime** (testkit) | | -| 3.11 tool-effect ledger, media externalisation | **Runtime** (session + graph checkpoint) with **Host** deciding idempotency | Matches sdk-gaps §4/§6. | +| 3.11 tool-effect ledger, media externalisation | **Runtime** (session state) with **Host** deciding idempotency | Matches sdk-gaps §4/§6; no graph checkpoint runtime is currently shipped. | | 3.12 OTel semconv sink | **Runtime** (optional sink) | Langfuse already lives there. | | UI adapters (AG-UI/Vercel), `sanitize_messages` | **Host** | Transport; but a `sanitize_history()` helper is cheap to put in harness. | | Harness batteries (`Coder`, `Shell`, `FileSystem`, `Memory`, `Guardrails`, `Planning`) | **Host** (OpenHuman `tools`, `security`, `agent/goals`, `agent/todos`) | TinyAgents' stance (tools are host vocabulary) is the same as Pydantic's core/Harness split; note Pydantic still ships them as an *official* library, which TinyAgents could mirror as an optional `tinyagents-batteries` crate. | | Capability bundle type, agent spec loading | **Runtime** (registry) | `AgentDefinition` already points this way. | -| Durable-execution engines (Temporal etc.) | Neither now | Desktop host has no workflow engine; keep the built-in checkpointer. | +| Durable-execution engines (Temporal etc.) | Neither now | Desktop host has no workflow engine; TinyAgents currently has no graph checkpointer. | | Realtime voice, CLI, web chat | **Host** | | ## 6. Sources diff --git a/docs/sdk-gaps.md b/docs/sdk-gaps.md index 9838204c5..38205e03c 100644 --- a/docs/sdk-gaps.md +++ b/docs/sdk-gaps.md @@ -305,7 +305,7 @@ Status: partially present. TinyAgents graph has `Send`, `Command`, reducers, interrupts, parallel execution, and max concurrency. OpenHuman still added `run_parallel_fanout` to provide an -ordered, bounded map/reduce helper for council runs and `spawn_parallel_agents`. +ordered, bounded map/reduce helper for council runs. Implement: @@ -327,7 +327,7 @@ Implement: Acceptance criteria: - OpenHuman can delete most of `run_parallel_fanout` and use the SDK helper. -- `spawn_parallel_agents` can be expressed as graph configuration plus +- Parallel-agent dispatch can be expressed as graph configuration plus OpenHuman policy adapters. - Results remain deterministic in input order even when workers complete out of order. diff --git a/docs/sdk-gaps/README.md b/docs/sdk-gaps/README.md index 492828c8b..80181ef9a 100644 --- a/docs/sdk-gaps/README.md +++ b/docs/sdk-gaps/README.md @@ -17,11 +17,11 @@ implementation candidates, with tests last once API and storage surfaces settle. ## Executive Summary -TinyAgents already has strong primitives for harness runs, graph execution, -middleware, event streams, model profiles, usage/cost accounting, checkpointers, -and sub-agent orchestration. The biggest remaining gaps are durable -orchestration stores, richer streaming events, graph fanout ergonomics, and -SDK-owned adapters for the lifecycle controls OpenHuman implements around it. +TinyAgents has strong primitives for harness runs, middleware, event streams, +model profiles, usage/cost accounting, and sub-agent orchestration. The graph +runtime, graph checkpointers, and graph fanout APIs were removed in TinyAgents +#61. Remaining gaps include durable orchestration stores, richer streaming +events, and SDK-owned adapters for lifecycle controls OpenHuman implements. OpenHuman can migrate more of `src/openhuman/agent/` if TinyAgents grows: @@ -32,12 +32,10 @@ OpenHuman can migrate more of `src/openhuman/agent/` if TinyAgents grows: redaction, and cancellation semantics. - Storage compatibility options for SQLite users that already depend on a different `rusqlite` / `libsqlite3-sys` version. -- Higher-level map/reduce and parallel-agent orchestration helpers on top of - graph `Send`. +- Higher-level map/reduce and parallel-agent orchestration helpers. - Budget enforcement and provider/model catalog metadata that can drive preflight, fallback, and reconciliation. -- Conformance suites for providers, tools, middleware, graph stores, and - checkpointers. +- Conformance suites for providers, tools, middleware, sessions, and stores. ## Backlog By Topic diff --git a/docs/sdk-gaps/cost-and-model-catalog.md b/docs/sdk-gaps/cost-and-model-catalog.md index acf0cf491..4f16713d2 100644 --- a/docs/sdk-gaps/cost-and-model-catalog.md +++ b/docs/sdk-gaps/cost-and-model-catalog.md @@ -85,12 +85,12 @@ TinyAgents has registry primitives. OpenHuman still needs richer diagnostics for duplicate components, alias resolution, component health, model/provider/tool capabilities, and event listener wiring. Implement: -- Registry snapshot export with models, tools, middleware, graph nodes, - checkpointers, task stores, event listeners, and aliases. +- Registry snapshot export with models, tools, middleware, task stores, event + listeners, and aliases. - Duplicate and shadowing diagnostics. - Health/status probes for registered providers and stores. - Machine-readable component dependency graph. -- Optional DOT/JSON graph export for runtime components, not only graph nodes. +- Optional DOT/JSON export of runtime component dependencies. Acceptance criteria: diff --git a/docs/sdk-gaps/durability.md b/docs/sdk-gaps/durability.md index fddc05c83..3f8ba52e2 100644 --- a/docs/sdk-gaps/durability.md +++ b/docs/sdk-gaps/durability.md @@ -59,22 +59,19 @@ Acceptance criteria: or a thin schema integration. - Storage features remain opt-in and keep the default crate dependency-light. -### 17. Storage And Graph Conformance +### 17. Storage Conformance Status: shipped. Every persistence layer has a shared contract-test suite, so a backend swap or a caller-supplied implementation can be certified rather than trusted: -`tinyagents_graph::testkit::conformance` (checkpointer + task-store, run -against memory/file/SQLite and JSONL in `tests/conformance.rs` and -`tests/persistence_conformance.rs`); `tinyagents_harness::store::conformance` +`tinyagents_harness::store::conformance` (`run_store_conformance`/`run_namespaced_store_conformance`, run against `InMemoryStore`/`FileStore`/`InMemoryNamespacedStore` in `tests/store_conformance.rs` — no in-tree SQLite `Store` exists yet); and `tinyagents_session::testkit::conformance` (run ledger + transcript history, run against SQLite and an in-memory double in `tests/session_conformance.rs`). See `docs/modules/harness/store.md` and -`crates/tinyagents-session/src/README.md` for details. Parallel-agent -order/failure/timeout/cancellation regression tests remain open (see the -fuzz/e2e graph-agent orchestration tests instead). +`crates/tinyagents-session/src/README.md` for details. Detached-task +order/failure/timeout/cancellation regression tests remain open. diff --git a/docs/sdk-gaps/orchestration.md b/docs/sdk-gaps/orchestration.md index d68c68b2d..53db1d895 100644 --- a/docs/sdk-gaps/orchestration.md +++ b/docs/sdk-gaps/orchestration.md @@ -12,7 +12,7 @@ Status: partially present. TinyAgents graph has `Send`, `Command`, reducers, interrupts, parallel execution, and max concurrency. OpenHuman still added `run_parallel_fanout` to provide an -ordered, bounded map/reduce helper for council runs and `spawn_parallel_agents`. +ordered, bounded map/reduce helper for council runs. Implement: @@ -34,7 +34,7 @@ Implement: Acceptance criteria: - OpenHuman can delete most of `run_parallel_fanout` and use the SDK helper. -- `spawn_parallel_agents` can be expressed as graph configuration plus +- Parallel-agent dispatch can be expressed as graph configuration plus OpenHuman policy adapters. - Results remain deterministic in input order even when workers complete out of order. @@ -55,7 +55,7 @@ and at a natural finish, `Followup` at a natural finish (one more turn), `Collect` onto `AgentRun::collected`; `RunPolicy::queue_mode` (`All` | `OneAtATime`); `AgentEvent::QueuedMessageApplied`. `SteeringHandle` is unchanged; OpenHuman's `agent/harness/run_queue/` can go. See -[`runtime.md`](modules/harness/runtime.md#queued-steering-and-follow-ups-a4). +[`runtime.md`](../modules/harness/runtime.md#queued-steering-and-follow-ups-a4). Implement: @@ -98,7 +98,7 @@ Acceptance criteria: ### 13. Middleware Control Outcomes -Status: shipped (harness loop control), partial (graph/sub-agent defer). +Status: shipped (harness loop control), partial (sub-agent defer); graph control was descoped in TinyAgents #61. Landed as part of `docs/runtime-comparison/plan.md` Phase 2 item A1. Every lifecycle `Middleware` hook has a `_control`-suffixed counterpart @@ -119,12 +119,12 @@ non-observer hooks are skipped (`Middleware::is_observer`). `BudgetMiddleware`/`HumanApprovalMiddleware` are rebased on `JumpTo(End)` and `Interrupt` respectively. See `crates/tinyagents-harness/src/{context,middleware,agent_loop}/*` and -[`docs/modules/harness/middleware.md`](modules/harness/middleware.md#middleware-control-a1). +[`docs/modules/harness/middleware.md`](../modules/harness/middleware.md#middleware-control-a1). -Still open: routing a harness-level control outcome onto a graph `Command` -node when the loop runs as a graph node, and "defer to task/sub-agent" as a -control outcome — both remain graph/A5 (loop-as-`CompiledGraph`) territory, -not yet implemented. +Still open: a future workflow runtime could route harness-level control +outcomes through graph commands and add a "defer to task/sub-agent" outcome. +TinyAgents #61 removed the graph runtime; this remains future design work and +is not a current harness capability. Acceptance criteria (harness scope): diff --git a/docs/sdk-gaps/streaming.md b/docs/sdk-gaps/streaming.md index 27a23c9f3..81b797228 100644 --- a/docs/sdk-gaps/streaming.md +++ b/docs/sdk-gaps/streaming.md @@ -54,28 +54,14 @@ TinyAgents has `HarnessEventJournal`, `StoreEventJournal`, `HarnessStatusStore`, and `HarnessRunStatus`. OpenHuman still bridges TinyAgents events into its own progress system, cost tracker, run ledger, and UI status stream. -Late attach is now partially solved. `tinyagents_harness::stream::{AssistantFrame, -FrameEncoder, reduce_frames}` give a durable per-block frame codec: -`FrameEncoder` turns a `ModelStreamItem` sequence into frames (periodic -`ToolArgsCheckpoint` snapshots bound replay depth), and `reduce_frames` folds -a — possibly truncated — sequence into a `PartialAssistantMessage`. Every -`GraphEvent` is now wrapped in a `GraphEventEnvelope { run_id, task_id, ns, -seq, event }` (`seq` monotonic per emitting graph instance, fresh for an -embedded subgraph), and `tinyagents_graph::stream::StreamProjection` folds -graph envelopes plus harness `AgentEvent`s into cursor-ordered -`messages`/`tool_calls`/`subagents` views; `StreamProjection::since(cursor)` -is the late-attach replay primitive. `seq` does not chain across a subgraph -boundary into one run-tree-wide sequence yet (D4's `TaskId` is the natural -place for that). `JournalGraphSink::dropped()` exposes its best-effort drop -counter so lossy-under-load is observable; the harness-side -`HarnessEventJournal` has no equivalent yet. Filters/compaction/redaction -hooks are still missing on both sides. +The harness exposes typed `AgentEvent`s and durable event/status journals. Hosts can +replay those records, but the journals do not yet provide cursor-addressable +replay, replay windows, or stream projections. Remaining work: - Replay windows, filters, compaction, and redaction hooks on the durable - journals (cursors/`since` now exist for `StreamProjection`; the journals - themselves still lack cursor-addressable replay). + journals (the journals still lack cursor-addressable replay). - Status stores with parent/root lineage, thread-scoped listing, phase details, active tool/model call ids, usage totals, cost totals, and terminal summaries. - Redaction policies for prompts, tool args, tool results, PII, secrets, and diff --git a/docs/spec/README.md b/docs/spec/README.md index 42895f530..796118fd5 100644 --- a/docs/spec/README.md +++ b/docs/spec/README.md @@ -1,38 +1,29 @@ # TinyAgents System Specification -TinyAgents is a small, provider-neutral agent harness for Rust, plus a durable -typed state-graph runtime. It takes its shape from LangChain (models, tools, -middleware, structured output, streaming, usage/cost) and LangGraph -(`START`/`END`, nodes, conditional edges, channels/reducers, checkpoints, -interrupts, subgraphs, time travel) — rebuilt as ordinary, typed Rust. The -system is organized as focused workspace packages: +TinyAgents is a small, provider-neutral agent harness for Rust. It takes its +shape from LangChain (models, tools, middleware, structured output, streaming, +usage/cost) — rebuilt as ordinary, typed Rust. The system is organized as +focused workspace packages: 1. the harness 2. reusable harness tools -3. the graph -4. the registry -5. durable sessions -6. host-neutral session runtime -7. host-neutral orchestration +3. the registry +4. durable sessions +5. host-neutral session runtime +6. host-neutral orchestration The goal is to make agent systems easy to define, inspect, run, test, and serialize without hiding the Rust types that make production systems reliable. ## Reference Positioning -TinyAgents synthesizes the reference systems rather than cloning either one: - -- LangGraph contributes the durable execution model: explicit state graphs, - virtual `START` and `END`, Pregel-style supersteps, reducers/channels, - commands, `Send` fanout, checkpointing, interrupts, subgraphs, streaming, and - time travel. -- LangChain contributes the harness model: provider-neutral models, tools, - middleware, runtime context, memory, retrieval, structured output, tracing, - usage, cost, and conformance tests for integrations. +TinyAgents builds on the harness model: provider-neutral models, tools, +middleware, runtime context, memory, retrieval, structured output, tracing, +usage, cost, and conformance tests for integrations. The target architecture is layered: the harness owns model/tool execution and -policies, the graph owns deterministic state transition and durability, the -registry owns named capabilities. No layer should bypass another layer's safety, policy, +policies, the session crate owns durable conversation state, and the registry +owns named capabilities. No layer should bypass another layer's safety, policy, observability, or test contracts. ## Detailed Module Docs @@ -61,25 +52,6 @@ observability, or test contracts. - [Observability and events](../modules/harness/observability.md) - [Performance and capacity testing](../modules/harness/performance.md) - [Testkit](../modules/harness/testkit.md) -- [Graph module](../modules/graph/README.md) - - [Package and core types](../modules/graph/package.md) - - [Builder and compile contract](../modules/graph/builder.md) - - [Node model](../modules/graph/nodes.md) - - [State, channels, and updates](../modules/graph/state-channels.md) - - [Edges, routing, commands, and sends](../modules/graph/routing.md) - - [Execution model and parallelization](../modules/graph/execution.md) - - [Parallel agents and context forking](../modules/graph/parallel-agents-forking.md) - - [Checkpointing, durability, state inspection, and time travel](../modules/graph/checkpointing.md) - - [Interrupts and resume](../modules/graph/interrupts.md) - - [Streaming and events](../modules/graph/streaming.md) - - [Observability and tracing](../modules/graph/observability.md) - - [Runtime context and policies](../modules/graph/runtime-policy.md) - - [Fault tolerance](../modules/graph/fault-tolerance.md) - - [Subgraphs](../modules/graph/subgraphs.md) - - [Sub-agents and recursion](../modules/graph/subagents-recursion.md) - - [Memory and stores boundary](../modules/graph/memory-boundary.md) - - [Visualization, introspection, and testkit](../modules/graph/visualization-testkit.md) - - [Implementation milestones](../modules/graph/milestones.md) - [Registry module](../modules/registry/README.md) - [Design](../modules/registry/design.md) - [Model catalog and local snapshots](../modules/registry/model-catalog.md) @@ -94,16 +66,12 @@ it. ## Design Goals - Make simple agent workflows concise. -- Make complex workflows explicit, inspectable, and testable. -- Treat graph execution as a first-class runtime, not an incidental callback - chain. +- Make agent workflows explicit, inspectable, and testable. - Keep model providers, tools, memory, and tracing behind stable traits. - Allow parent orchestrators and humans to steer orchestrator agents and sub-agents through typed, policy-checked, observable commands. - Prefer deterministic state transitions around inherently nondeterministic LLM calls. -- Keep every generated or hand-authored graph explainable as topology, - capabilities, policies, state channels, checkpoints, and events. ## Module 1: Harness @@ -133,15 +101,6 @@ resume. The `tools` array stays stable between discoveries when no per-turn exposure middleware changes the advertised set. `RunPolicy::tool_schemas` optionally projects and byte-budgets every schema. -## Module 2: Graph - -The graph is the durable, typed state-graph runtime: `START`/`END`, nodes, -reducers/channels, routing, supersteps, checkpointing, interrupts, streaming, -subgraphs, and execution guarantees. See [`graph-spec.md`](graph-spec.md) for -the full specification, and -[`docs/modules/graph/README.md`](../modules/graph/README.md) for the -per-topic implementation docs. - ## Package Layout The repository root is a virtual Cargo workspace. There is no `tinyagents` @@ -152,13 +111,12 @@ they use. Shared runtime errors live in `tinyagents-harness`. crates/ tinyagents-harness/ # models, tools, middleware, providers, runtime tinyagents-harness-tools/ # reusable tool implementations and stateful support - tinyagents-graph/ # durable typed state graphs - tinyagents-tasks/ # graph-independent detached-task registry, task stores, orchestration tools + tinyagents-tasks/ # detached-task registry, task stores, orchestration tools tinyagents-registry/ # named capabilities and model catalog tinyagents-session/ # durable session history and run ledger tinyagents-definition/ # host-owned agent definition vocabulary tinyagents-runtime/ # host-neutral stateful harness sessions - tinyagents-orchestration/ # host-neutral team/workflow composition over graph+harness+session + tinyagents-orchestration/ # host-neutral multi-agent composition over harness+session tinyagents-integration-tests/ # cross-crate tests and runnable examples ``` @@ -175,8 +133,9 @@ All four milestones below have shipped as of v1.5.0. ### Milestone 1: Core Runtime (shipped) -Chat message primitives, the model and tool traits, the state graph with -direct and conditional edges, and the initial test/example suite. +Chat message primitives, the model and tool traits, and the initial +test/example suite. The original state-graph runtime was removed in TinyAgents +#61. ### Milestone 2: Harness (shipped) @@ -192,9 +151,9 @@ mock provider. ### Milestone 4: Production Runtime Features (shipped) -Streaming events, checkpointing and resume support, the graph run status -store, event journal with listener replay, graph export, and an embedded -Langfuse tracing integration (`LangfuseClient`, `GraphLangfuseExporter`). +Streaming events, session checkpointing and resume support, an event journal +with listener replay, and an embedded Langfuse tracing integration. The graph +run status store and graph exporter were removed in TinyAgents #61. ## Open Questions @@ -204,10 +163,6 @@ Historical decisions that have since been settled, kept for context: than becoming one crate per provider. - Memory and embeddings are async, matching the rest of the harness surface. -Remaining open question: - -- Should graph nodes support typed route enums as a stronger alternative to - string-keyed conditional routing before further serialization work lands? Transcript accounting preserves cost provenance in durable usage records. Hosts should use `TranscriptSpend::cost_split` to separate charged or estimated diff --git a/docs/spec/graph-spec.md b/docs/spec/graph-spec.md deleted file mode 100644 index ce671b32e..000000000 --- a/docs/spec/graph-spec.md +++ /dev/null @@ -1,338 +0,0 @@ -# Graph Module Specification - -The graph is the workflow runtime. It executes stateful nodes, applies state -updates, follows direct or conditional edges, records execution history, handles -interrupts, and returns a final state. - -The first implementation can stay sequential, but the module should be designed -toward LangGraph's durable execution model: compiled graphs, virtual `START` and -`END` nodes, supersteps, reducer-driven state updates, checkpoints, interrupts, -commands, streaming, and subgraphs. - -### Source Inspiration - -The graph design is informed by LangGraph's docs on the graph API, reducers, -commands, persistence, checkpointers, interrupts, streaming, subgraphs, and fault -tolerance: - -- <https://docs.langchain.com/oss/python/langgraph/graph-api> -- <https://docs.langchain.com/oss/python/langgraph/persistence> -- <https://docs.langchain.com/oss/python/langgraph/checkpointers> -- <https://docs.langchain.com/oss/python/langgraph/interrupts> -- <https://docs.langchain.com/oss/python/langgraph/streaming> -- <https://docs.langchain.com/oss/python/langgraph/event-streaming> -- <https://docs.langchain.com/oss/python/langgraph/use-subgraphs> -- <https://docs.langchain.com/oss/python/langgraph/fault-tolerance> - -### Responsibilities - -- Store named nodes. -- Store direct and conditional edges. -- Validate graph structure at compile time. -- Produce an immutable executable graph. -- Run async node handlers. -- Route based on node output or command output. -- Apply partial state updates through reducers. -- Enforce recursion limits. -- Persist checkpoints at safe boundaries. -- Support interrupts and resume. -- Stream typed execution events. -- Write readable execution status records for graph runs. -- Maintain append-only graph event journals for external listeners. -- Cache derived graph observability projections without making them the source - of truth. -- Return final state and execution history. -- Support graph visualization and serialization later. - -### Core Concepts - -`State` is user-owned application state. TinyAgents should never require a -specific state shape for hand-written Rust graphs. - -`Node<State>` is an async unit of work. - -`NodeOutput<State>` controls execution in the current scaffold: - -- `Continue(State)` follows a direct edge. -- `Route { state, route }` follows a conditional edge. -- `End(State)` stops execution. - -The target design should evolve this into partial updates and commands: - -```rust -pub enum NodeResult<Update> { - Update(Update), - Command(Command<Update>), - Interrupt(Interrupt), -} - -pub struct Command<Update> { - pub update: Option<Update>, - pub goto: Vec<NodeId>, - pub resume: Option<serde_json::Value>, -} -``` - -`GraphBuilder<State, Update>` should own graph construction. `CompiledGraph` -should own execution. This separates user-friendly mutation during setup from a -validated immutable runtime. - -```rust -let graph = GraphBuilder::new() - .add_node("agent", agent_node) - .add_node("tools", tools_node) - .add_edge(START, "agent") - .add_conditional_edges("agent", route_agent) - .add_edge("tools", "agent") - .compile()?; -``` - -### State Updates And Reducers - -LangGraph nodes return partial state updates. TinyAgents should adopt the same -direction because it enables parallel execution, replay, checkpointing, and -clearer node contracts. - -The default reducer should be overwrite. Users should be able to opt into -reducers for fields that accumulate values: - -- append list -- merge messages by id -- set union -- numeric min/max -- custom reducer - -Possible Rust shape: - -```rust -pub trait Reducer<T>: Send + Sync { - fn reduce(&self, current: T, update: T) -> Result<T>; -} - -pub trait StateReducer<State, Update>: Send + Sync { - fn apply(&self, state: State, update: Update) -> Result<State>; -} -``` - -For milestone 1, whole-state updates are acceptable. For durable parallel graph -execution, partial updates and reducers should be introduced before -checkpoint/resume semantics harden. - -### Graph Lifecycle - -1. Define state. -2. Define update type if partial updates are enabled. -3. Create graph builder. -4. Add nodes. -5. Add direct or conditional edges. -6. Add `START` edge. -7. Compile and validate the graph. -8. Run graph with initial state and runtime config. -9. Inspect final state, checkpoints, events, and visited nodes. - -### Routing Semantics - -Direct routing: - -```text -START -> agent -> summarize -> END -``` - -Conditional routing: - -```text -START -> agent -agent --tool--> tools -agent --final--> END -tools ---------> agent -``` - -Conditional routes may start as explicit strings. Later versions should support -typed route enums or route newtypes so Rust users can avoid typo-prone strings. - -Nodes should not mix static outgoing edges and dynamic command-based routing in -the same execution mode unless the behavior is deliberately specified. A strict -compile-time validation rule is preferable: a node has either normal outgoing -edges or command routing, not both. - -### Supersteps - -The target executor should be superstep-based: - -1. Take the current active node set. -2. Run all active nodes for the step, respecting concurrency policy. -3. Collect partial state updates, commands, interrupts, and errors. -4. Apply reducers at the step boundary. -5. Persist a checkpoint. -6. Select the next active nodes. -7. Stop when the active set is empty or reaches `END`. - -The first implementation can run one node at a time, but checkpointing and -parallel execution should use superstep boundaries as the durable unit. Do not -checkpoint mid-node. - -### Checkpointing And Persistence - -Graph checkpointing is not the same as harness memory. Checkpoints are -thread-scoped graph execution snapshots used for resume, interrupts, and fault -tolerance. - -```rust -#[async_trait] -pub trait Checkpointer<State>: Send + Sync { - async fn put(&self, checkpoint: Checkpoint<State>) -> Result<CheckpointId>; - async fn get(&self, thread_id: &str, checkpoint_id: Option<&str>) -> Result<Option<Checkpoint<State>>>; - async fn list(&self, thread_id: &str) -> Result<Vec<CheckpointMetadata>>; -} -``` - -A checkpoint should contain: - -- thread id -- checkpoint id -- parent checkpoint id -- namespace -- state snapshot -- next active nodes -- completed tasks for the superstep -- pending writes -- interrupts -- metadata - -Interrupted or failed nodes may rerun from the beginning. Node authors must make -side effects idempotent or isolate side effects behind tools/middleware that can -record exactly-once intent. - -### Interrupts And Resume - -Interrupts support human-in-the-loop and external approval flows. - -```rust -pub struct Interrupt { - pub id: String, - pub node: NodeId, - pub payload: serde_json::Value, -} -``` - -Resume should use a command-style API: - -```rust -graph.resume( - RunConfig::thread("support-123"), - Command::resume(json!({ "approved": true })), -).await?; -``` - -The default semantic should match LangGraph: resuming restarts the interrupted -node and replays until the interrupt point using stored resume values. That is -more durable than trying to suspend an async Rust stack. - -### Streaming - -The graph should expose low-level runtime events, higher-level projections, a -status store, and optional durable replay for outside listeners. The canonical -feature references are: - -- [Graph streaming and events](../modules/graph/streaming.md) -- [Graph observability and tracing](../modules/graph/observability.md) -- [Graph checkpointing and state inspection](../modules/graph/checkpointing.md) -- [Graph memory and stores boundary](../modules/graph/memory-boundary.md) - -Low-level events: - -- node started -- node completed -- node failed -- state update -- checkpoint saved -- task scheduled -- interrupt emitted -- route selected - -High-level stream modes: - -- values: full state snapshots -- updates: partial state updates -- messages: model/message deltas emitted by harness nodes -- debug: verbose executor events -- interrupts: interrupt payloads -- custom: user events - -The graph should also expose a compact run-status record: - -```rust -pub struct GraphRunStatus { - pub run_id: RunId, - pub root_run_id: RunId, - pub parent_run_id: Option<RunId>, - pub thread_id: Option<ThreadId>, - pub graph_id: GraphId, - pub checkpoint_id: Option<CheckpointId>, - pub checkpoint_namespace: Vec<String>, - pub status: ExecutionStatus, - pub current_step: usize, - pub active_nodes: Vec<NodeId>, - pub pending_interrupts: Vec<InterruptId>, - pub last_event_id: Option<EventId>, - pub started_at: SystemTime, - pub updated_at: SystemTime, - pub ended_at: Option<SystemTime>, - pub error: Option<GraphErrorSummary>, -} -``` - -Graph status records are not checkpoints. Checkpoints preserve resumable graph -state; status records summarize live and recent execution for observers. A -graph event journal should let listeners subscribe live or replay from a stored -offset by run id, root run id, thread id, graph id, node id, event kind, or -namespace. Derived projections such as latest status by thread, task timing -rollups, checkpoint summaries, and introspection snapshots may be cached when -they include source coordinates: run id, checkpoint id, namespace, step, event -offset, and projection version. - -### Subgraphs - -Subgraphs should be executable graphs that can be used as nodes. - -Two modes are needed: - -- shared-state subgraph: parent and child graph use the same state channels -- adapter subgraph: wrapper node maps parent state into child state and maps the - child result back into parent state - -Checkpoint namespaces are required so parent and child checkpoint ids do not -collide. - -### Execution Guarantees - -The graph runtime should guarantee: - -- every visited node existed at validation time -- every configured edge points to an existing node -- conditional routes fail clearly when missing -- recursion limit failures are deterministic -- checkpoint writes happen at configured execution boundaries -- interrupted runs can be resumed only when checkpointing is configured -- final state is returned exactly once - -The graph runtime should not guarantee: - -- deterministic LLM output -- tool idempotency -- provider-specific retry behavior -- persistence across process restarts unless a checkpointer is configured -- exactly-once side effects inside node code - -### Future Graph Features - -- graph serialization to JSON -- Mermaid export -- parallel branches -- joins -- typed route enums -- static graph analysis -- graph diffing -- graph snapshots for tests -- durable task queue integration - diff --git a/docs/spec/harness-spec.md b/docs/spec/harness-spec.md index c698a0568..4e4b78c6b 100644 --- a/docs/spec/harness-spec.md +++ b/docs/spec/harness-spec.md @@ -8,7 +8,7 @@ middleware and context seams. The harness must stay composable. It should not be a single monolithic `Agent` type that hides every behavior. A direct model call, a model-plus-tools loop, and -a graph node that invokes a model should all share the same harness primitives. +nested harness calls should share the same primitives. ### Source Inspiration @@ -48,15 +48,15 @@ testing: without replacing provider adapters. - Emit typed events for observability and streaming. - Write readable run status records for direct model calls, agent loops, and - graph-node child harness calls. + nested harness calls. - Maintain append-only event journals when durable listener replay is configured. - Enforce retry, timeout, model-call, tool-call, and recursion policies. - Accept sub-agent and orchestrator steering commands from humans, parent - agents, graph supervisors, middleware, and tests at safe loop boundaries. + agents, middleware, and tests at safe loop boundaries. - Normalize model and tool errors into framework errors. - Persist resolved model identity in responses, events, usage/cost records, run - status, and durable agent or graph state so later calls can reuse it. + status, and session state so later calls can reuse it. - Provide test doubles for models, tools, stores, clocks, and ids. ### Core Types @@ -93,8 +93,8 @@ unit tests straightforward. ### Model Abstraction -Models should be provider-agnostic. The graph layer should never know whether a -node uses OpenAI, Anthropic, Ollama, a local model, or a test fake. +Models should be provider-agnostic. Harness callers can use OpenAI, Anthropic, +Ollama, a local model, or a test fake. ```rust #[async_trait] @@ -261,7 +261,7 @@ Limits are not optional. The harness should enforce: ### Middleware Middleware is the primary extension point for behavior that should not be baked -into the model or graph APIs. +into the model API. ```rust #[async_trait] @@ -304,9 +304,8 @@ Expected middleware: ### Memory -Memory is host policy applied through harness middleware. The graph runtime -handles checkpointed graph execution; hosts decide how conversation and -application memory are loaded and persisted. +Memory is host policy applied through harness middleware; hosts decide how +conversation and application memory are loaded and persisted. Memory is split into two concepts: diff --git a/integration-proofs/run-ledger-mongodb/Cargo.lock b/integration-proofs/run-ledger-mongodb/Cargo.lock index 0cb6a3eb2..7d7eaa5d9 100644 --- a/integration-proofs/run-ledger-mongodb/Cargo.lock +++ b/integration-proofs/run-ledger-mongodb/Cargo.lock @@ -2324,7 +2324,7 @@ dependencies = [ [[package]] name = "tinyagents-definition" -version = "2.11.0" +version = "3.0.0" dependencies = [ "anyhow", "async-trait", @@ -2337,7 +2337,7 @@ dependencies = [ [[package]] name = "tinyagents-harness" -version = "2.11.0" +version = "3.0.0" dependencies = [ "anyhow", "async-trait", @@ -2387,7 +2387,7 @@ dependencies = [ [[package]] name = "tinyagents-session" -version = "2.11.0" +version = "3.0.0" dependencies = [ "anyhow", "async-trait", @@ -2411,7 +2411,7 @@ dependencies = [ [[package]] name = "tinyagents-storage-kit" -version = "2.11.0" +version = "3.0.0" dependencies = [ "data-encoding", "serde", @@ -2425,7 +2425,7 @@ dependencies = [ [[package]] name = "tinyagents-tasks" -version = "2.11.0" +version = "3.0.0" dependencies = [ "serde", "serde_json", @@ -2434,6 +2434,7 @@ dependencies = [ "tinyinference-llm", "tokio", "tracing", + "uuid", ] [[package]] diff --git a/integration-proofs/run-ledger-mongodb/tests/session_driver_conformance.rs b/integration-proofs/run-ledger-mongodb/tests/session_driver_conformance.rs index 41bb85900..cfe8b6eba 100644 --- a/integration-proofs/run-ledger-mongodb/tests/session_driver_conformance.rs +++ b/integration-proofs/run-ledger-mongodb/tests/session_driver_conformance.rs @@ -6,13 +6,13 @@ use std::sync::Arc; use tinyagents_session::{ - port::{AgentStorageFactory, DriverSessionStores}, + port::{AgentStorageFactory, ConditionedAgentStorageFactory, DriverSessionStores}, run_ledger::{AgentRunListRequest, DriverRunLedger, RunLedgerStore}, testkit::conformance::{ run_ledger_conformance, session_store_conformance, session_store_isolation_conformance, }, }; -use tinyagents_storage_kit::{Scope, ScopedStorage}; +use tinyagents_storage_kit::{Scope, ScopedStorage, WriteCondition}; use tinystoragedrivers_api::{DocumentStore, Storage, StorageBackend, StreamStore}; use tinystoragedrivers_connect::MongoConnection; use tinystoragedrivers_document::mongodb::{MongoDocuments, MongoStreams}; @@ -43,13 +43,24 @@ async fn local_connection(label: &str) -> MongoConnection { .expect("local MongoDB replica set is reachable") } -fn scoped_storage(connection: &MongoConnection, scope: &Scope) -> ScopedStorage { - let documents = - Arc::new(MongoDocuments::new(connection.clone(), scope.clone())) as Arc<dyn DocumentStore>; +fn scoped_storage( + connection: &MongoConnection, + scope: &Scope, + condition: Option<&WriteCondition>, +) -> Result<ScopedStorage, tinystoragedrivers_types::StorageError> { + let documents = match condition { + Some(condition) => Arc::new(MongoDocuments::conditioned( + connection.clone(), + scope.clone(), + condition.clone(), + )?) as Arc<dyn DocumentStore>, + None => Arc::new(MongoDocuments::new(connection.clone(), scope.clone())) + as Arc<dyn DocumentStore>, + }; let streams = Arc::new(MongoStreams::new(connection.clone(), scope.clone())) as Arc<dyn StreamStore>; let memory = MemoryBackend::new().for_scope(scope).unwrap(); - ScopedStorage::new( + Ok(ScopedStorage::new( Storage::new( documents, streams, @@ -59,7 +70,7 @@ fn scoped_storage(connection: &MongoConnection, scope: &Scope) -> ScopedStorage memory.secrets().clone(), ), scope.clone(), - ) + )) } #[tokio::test(flavor = "multi_thread")] @@ -71,9 +82,17 @@ async fn mongodb_session_store_round_trip_list_delete_and_agent_isolation() { let factory_connection = connection.clone(); let factory: AgentStorageFactory = Arc::new(move |agent_id| { let agent_scope = factory_scope.child(agent_id)?; - Ok(scoped_storage(&factory_connection, &agent_scope)) + scoped_storage(&factory_connection, &agent_scope, None) }); - let stores = DriverSessionStores::new(factory).unwrap(); + let conditioned_scope = scope.clone(); + let conditioned_connection = connection.clone(); + let conditioned: ConditionedAgentStorageFactory = Arc::new(move |agent_id, condition| { + let agent_scope = conditioned_scope.child(agent_id)?; + scoped_storage(&conditioned_connection, &agent_scope, Some(condition)) + }); + let stores = DriverSessionStores::new(factory) + .unwrap() + .with_conditioned_storage_factory(conditioned); session_store_conformance(&stores).await; session_store_isolation_conformance(&stores).await; } @@ -83,11 +102,11 @@ async fn mongodb_session_store_round_trip_list_delete_and_agent_isolation() { async fn mongodb_run_ledger_contract_and_scope_isolation() { let scope = Scope::new("run_ledger_conformance").unwrap(); let connection = local_connection("run_ledger").await; - let storage = scoped_storage(&connection, &scope); + let storage = scoped_storage(&connection, &scope, None).unwrap(); let ledger = Arc::new(DriverRunLedger::new(storage)) as Arc<dyn RunLedgerStore>; run_ledger_conformance(ledger.clone()).await; let peer_scope = Scope::new("run_ledger_peer").unwrap(); - let peer_storage = scoped_storage(&connection, &peer_scope); + let peer_storage = scoped_storage(&connection, &peer_scope, None).unwrap(); let peer = Arc::new(DriverRunLedger::new(peer_storage)) as Arc<dyn RunLedgerStore>; assert!( peer.list_agent_runs(&AgentRunListRequest::default()) diff --git a/scripts/check-harness-tools-features.sh b/scripts/check-harness-tools-features.sh new file mode 100755 index 000000000..de3c2ef22 --- /dev/null +++ b/scripts/check-harness-tools-features.sh @@ -0,0 +1,13 @@ +#!/usr/bin/env bash +set -euo pipefail + +baseline=$(cargo tree -p tinyagents-harness-tools -e normal --no-default-features) +media=$(cargo tree -p tinyagents-harness-tools -e normal --no-default-features --features media) +diff -u <(printf '%s\n' "$baseline") <(printf '%s\n' "$media") +if cargo tree -p tinyagents-harness-tools -e normal --no-default-features --features media \ + | rg 'tinyinference-(video|image)|(^| )((flate2|tar|zip))($| )'; then + echo "media feature introduced a provider/media dependency" >&2 + exit 1 +fi +cargo tree -p tinyagents-harness-tools -e features -i reqwest --no-default-features --features net \ + | rg 'tinyagents-harness-tools feature "net"'