diff --git a/.gitmodules b/.gitmodules index 693dded..a01a23c 100644 --- a/.gitmodules +++ b/.gitmodules @@ -1,3 +1,3 @@ [submodule "wiki"] path = wiki - url = git@github.com:tinyhumansai/tinyagents.wiki.git + url = https://github.com/tinyhumansai/tinyagents.wiki.git diff --git a/AGENTS.md b/AGENTS.md index fa3eec6..091c50a 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -3,10 +3,12 @@ ## Project Structure & Module Organization TinyAgents is a Rust 2024 library crate rooted at `Cargo.toml`. Public API -exports live in `src/lib.rs`, with core modules split across `src/chat.rs`, -`src/model.rs`, `src/tool.rs`, `src/graph.rs`, and `src/error.rs`. Additional -architecture work is staged under module directories such as `src/harness/`, -`src/language/`, and `src/registry/`. +exports live in `src/lib.rs`, with the crate-wide error type in `src/error.rs`. +The five surfaces each live in their own module directory: `src/graph/` +(durable typed state graphs), `src/harness/` (provider-neutral model calls, +tools, middleware, streaming), `src/language/` (the declarative `.rag` +blueprint format), `src/registry/` (the named capability catalog), and +`src/repl/` (the imperative `.ragsh` session runtime). Prefer small, focused modules that do one thing extremely well. New feature areas should live in module directories instead of accumulating broad, @@ -15,10 +17,22 @@ dedicated `types.rs` file and keep module-local unit tests in a dedicated `test.rs` file. The module root should wire the pieces together and expose the smallest useful API. -Integration tests are in `tests/`, currently focused on serialization behavior. -Runnable usage examples are in `examples/`, especially -`examples/basic_graph.rs`. Design notes and module-level specifications live in -`docs/`, with `docs/spec/README.md` as the top-level architecture reference. +Two Cargo features gate optional dependencies: `sqlite` (embedded SQLite +checkpointer, `graph::checkpoint::SqliteCheckpointer`) and `repl` (embedded +Rhai engine backing `repl::session`); every other provider and surface is +compiled in by default. + +Integration tests are in `tests/`, covering serialization, graph routing, +registry binding, the expressive and REPL languages, streaming, subagents, +and provider contracts (including live, network-gated tests such as +`tests/live_*.rs`). Runnable usage examples are in `examples/`, especially +`examples/basic_graph.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/`, `expressive-language/`, `repl-language/`). 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. ## Build, Test, and Development Commands diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 3ae99dd..5a9da7b 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -20,6 +20,31 @@ The bundled example should also run: cargo run --example basic_graph ``` +To build with the optional embedded SQLite checkpointer or the `.ragsh` Rhai +session runtime, enable the relevant feature: + +```sh +cargo test --features sqlite +cargo test --features repl +``` + +### Wiki Submodule + +The published GitHub wiki lives in `wiki/`, checked out as a git submodule +pointing at the `tinyhumansai/tinyagents.wiki` repository. It is not part of +the crate build and is not covered by this project's Markdown line-length or +review rules. Clone with submodules to pull it down: + +```sh +git clone --recurse-submodules https://github.com/tinyhumansai/tinyagents.git +# or, in an existing checkout: +git submodule update --init wiki +``` + +Do not edit `wiki/` content as part of an unrelated code or docs change; wiki +edits should be their own commit (or made directly on the wiki) so the +submodule pointer update stays easy to review in isolation. + ## Project Philosophy TinyAgents should make agent systems explicit and inspectable. Prefer: diff --git a/README.md b/README.md index 198fada..2a644b2 100644 --- a/README.md +++ b/README.md @@ -157,21 +157,23 @@ Add TinyAgents to your project: ```toml [dependencies] -tinyagents = "0.1" +tinyagents = "1.5" ``` -The OpenAI (and OpenAI-compatible) provider is compiled in by default and pulls -no extra dependencies; the build stays offline unless you actually make a call. +The OpenAI (and OpenAI-compatible) provider is compiled in by default; the +build stays offline unless you actually make a call. Two optional Cargo +features gate heavier dependencies: `sqlite` (embedded SQLite checkpointer) +and `repl` (embedded Rhai engine for the `.ragsh` session runtime). To explore locally: ```sh -git clone git@github.com:tinyhumansai/rustagents.git -cd rustagents +git clone git@github.com:tinyhumansai/tinyagents.git +cd tinyagents cargo run --example basic_graph ``` -OpenAI-backed examples need the feature flag and an API key: +OpenAI-backed examples need an API key: ```sh export OPENAI_API_KEY=... @@ -236,6 +238,10 @@ All live in [`examples/`](examples/): - **`openai_self_blueprint`** — **the deepest recursion:** a model authors a `.rag` blueprint that is compiled and run on the same runtime. - **`rag_blueprint`** — load and run a declarative `.rag` workflow. +- **`goals_and_todos`** — a durable `ThreadGoal` driving a `TaskBoard` kanban + on one thread. +- **`subconscious_loop`** — an offline, testable autonomous closed-loop + harness (see [`examples/subconscious_loop/README.md`](examples/subconscious_loop/README.md)). - **`openai_chat`** — a single provider-backed chat turn. - **`openai_tools`** — tool calling against a hosted model. - **`openai_structured`** — typed structured output. diff --git a/ROADMAP.md b/ROADMAP.md index b7187c9..d48fcdb 100644 --- a/ROADMAP.md +++ b/ROADMAP.md @@ -1,66 +1,50 @@ # Roadmap -TinyAgents is pre-1.0. The roadmap favors small, well-tested modules that build -toward a production-grade Rust agent runtime. - -## Current Foundation - -- chat message primitives -- model request and response types -- async chat model trait -- async tool trait -- executable state graph -- direct and conditional routing -- graph validation and recursion limits -- basic examples and serialization tests +TinyAgents is at v1.5.0. The roadmap favors small, well-tested modules that +build toward a production-grade Rust agent runtime. + +## Shipped Foundation + +- 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 `TaskBoard` productivity primitives, exposed as + harness tools +- named capability registry (models, tools, agents, graphs, stores, + middleware, policy) bound by name (`registry/`) +- the declarative `.rag` blueprint language: lexer, parser, compiler, and + registry-backed binding (`language/`) +- the imperative `.ragsh` REPL language for capability-bound interactive + orchestration (`repl/`) +- an optional SQLite-backed checkpointer (`sqlite` feature) and an optional + Rhai-backed `.ragsh` session runtime (`repl` feature) +- an embedded Langfuse client and graph exporter for observability ## Near-Term Work -- split broad modules into focused module directories with `types.rs` and - `test.rs` -- expand graph tests for routing, recursion, validation, and error behavior -- strengthen harness model, tool, prompt, context, middleware, and usage APIs -- add more examples for model calls, tools, and graph composition -- define reducer and state-channel APIs for parallel writes -- document stable public API boundaries as modules mature - -## Declarative Workflow Language - -The `.rag` language should let humans and LLMs describe agent workflows without -embedding arbitrary host code. - -Planned capabilities: - -- graph topology declarations -- allowed models, tools, agents, stores, middleware, and subgraphs -- state channels and reducers -- direct routes, conditional routes, commands, sends, joins, and barriers -- parallel sub-agent fanout -- blocking and optional child-agent policies -- checkpoint, interrupt, timeout, retry, budget, and concurrency policies -- source spans and diagnostics -- blueprint review before execution +- broaden `.rag`/`.ragsh` example coverage for less-common routing and + parallel-fanout shapes +- continue splitting any module or doc that grows past the 500-line limit + into focused files +- expand live (network-gated) provider contract tests as new + OpenAI-compatible endpoints are added +- track and close the internal SDK feature-parity backlog in + [`docs/sdk-gaps.md`](docs/sdk-gaps.md) ## Parallel Agents And Sub-Agents -TinyAgents should support workflow-native parallelism: - -- forked child contexts -- shared caches with explicit isolation policy -- child event namespaces -- parent and child run ids -- deterministic reducer-based merges -- optional, blocking, race, quorum, fallback, and compare policies -- resumable checkpoints across parallel branches - -## REPL And Agent-Authored Workflows - -The `.ragsh` REPL layer should let agents and humans inspect, script, and -control graph runs through capability-bound functions. It should be able to -propose `.rag` workflows, but those workflows must pass through parser, -registry, policy, and compiler checks before execution. +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. -## Pre-1.0 Stability +## Stability -APIs may change before 1.0. Changes should be documented, tested, and shaped by -real examples rather than speculative abstraction. +The public API is versioned via semver starting at 1.0. Breaking changes are +documented in release notes, tested, and shaped by real examples rather than +speculative abstraction. diff --git a/docs/audit.md b/docs/audit.md index 66cbd67..a264917 100644 --- a/docs/audit.md +++ b/docs/audit.md @@ -2,7 +2,7 @@ Date: 2026-06-29 -This audit reviewed the current RustAgents codebase for correctness gaps, +This audit reviewed the current TinyAgents codebase for correctness gaps, security issues, implementation weaknesses, and spec-contract mismatches. No code changes were made as part of the audit. diff --git a/docs/modules/expressive-language/README.md b/docs/modules/expressive-language/README.md index 49c76ab..e86e327 100644 --- a/docs/modules/expressive-language/README.md +++ b/docs/modules/expressive-language/README.md @@ -386,429 +386,8 @@ error[E-rag-unknown-node]: route target `toolz` does not exist help: did you mean `tools`? ``` -## Node Kinds +--- -Initial built-in node kinds: - -### `agent` - -Uses the harness agent loop or one model call depending on config. - -Supported fields: - -- `model` -- `system` -- `prompt` -- `tools` -- `routes` -- `retry` -- `timeout` - -### `model` - -Single model invocation. Does not automatically execute tools. - -Supported fields: - -- `model` -- `system` -- `prompt` -- `routes` -- `retry` -- `timeout` - -### `tool_executor` - -Executes tool calls already present in state. - -Supported fields: - -- `tools` -- `next` -- `retry` -- `timeout` - -### `router` - -Routes based on a named route function provided from Rust. - -Supported fields: - -- `routes` -- `metadata` - -### `subgraph` - -Calls another compiled graph. - -Supported fields: - -- `graph` -- `next` -- `routes` - -### `subagent` - -Calls a registered harness agent as a graph node. - -Supported fields: - -- `agent` -- `input` -- `next` -- `routes` -- `retry` -- `timeout` -- `steering` - -Example: - -```tinyagents -node research { - kind subagent - agent "researcher" - steering { - parent allow ["add_instruction", "request_status", "cancel"] - human allow ["add_instruction", "pause", "resume", "cancel"] - delivery "safe_boundary" - } - next synthesize -} -``` - -Steering policies lower into harness steering policy and graph task policy. They -can narrow a child agent's model/tool/runtime limits but cannot grant -capabilities absent from the registry or parent run policy. - -### `repl_agent` - -Runs a registered REPL script or model-driven CodeAct loop through the harness -REPL runtime. - -Supported fields: - -- `model` -- `script` -- `tools` -- `routes` -- `retry` -- `timeout` - -### `interrupt` - -Emits a resumable human-in-the-loop interrupt. - -Supported fields: - -- `prompt` -- `options` -- `routes` -- `metadata` - -### `join` - -Waits for named upstream nodes or barrier channels before continuing. - -Supported fields: - -- `sources` -- `next` -- `timeout` - -## Binding To Rust - -The language should not define arbitrary Rust closures. Instead, it should bind -to Rust-provided registries: - -```rust -let workflow = LanguageCompiler::new() - .with_models(models) - .with_tools(tools) - .with_node_templates(templates) - .compile_source("support.rag", source)?; -``` - -Registries: - -- model registry -- tool registry -- agent registry -- node template registry -- route function registry -- reducer registry -- graph registry for subgraphs -- store registry -- middleware registry -- REPL script registry - -When a graph is generated by a REPL session, the session may call the compiler -with source text or an AST, but the compiler must use the same registries and -policy checks. Generated source can request capabilities only from the allowed -set attached to the parent run or registry namespace. - -This keeps source files declarative and prevents unsafe dynamic execution. - -## State Model - -Version 1 should keep state Rust-owned. The language can refer to standard -channels by convention and bind them to registered reducers: - -- `messages` -- `tool_calls` -- `structured_response` -- `metadata` -- `artifacts` -- `candidates` -- `usage` -- `interrupts` - -Example: - -```tinyagents -channel messages messages -channel candidates append -channel usage aggregate "usage_delta" -channel review overwrite -``` - -Future versions may add state schema declarations: - -```tinyagents -state SupportState { - messages: messages append - customer_id: string overwrite - ticket_id: string? overwrite -} -``` - -State schemas should be delayed until reducer-based graph execution exists. - -## Routes - -Routes are named outcomes. - -```tinyagents -routes { - tool_call -> tools - final -> END - escalate -> human_review -} -``` - -Rules: - -- route names are unique per node -- route targets must exist or be `END` -- route names are ASCII identifiers -- a node may use `routes` or `next`, not both -- `END` is reserved - -Future typed route support can generate Rust enums from route declarations. - -## Policies - -Node-level policies: - -```tinyagents -node agent { - timeout 30s - retry { - max_attempts: 3 - backoff: "exponential" - } -} -``` - -Graph-level defaults: - -```tinyagents -graph support_agent { - defaults { - timeout 60s - recursion_limit 50 - } -} -``` - -Policies lower into graph node policies and harness request policies. - -## Comments And Strings - -Comments: - -```tinyagents -// line comment -``` - -Strings: - -```tinyagents -system "single line" - -prompt """ -multi-line prompt -""" -``` - -Multi-line strings should preserve content exactly except for one predictable -dedent rule. - -## Safety - -The language must be safe to parse and validate from untrusted text. - -Safety rules: - -- no arbitrary code execution -- no filesystem access from language source -- no network access from language source -- no dynamic provider lookup without registry binding -- no environment variable interpolation in v1 -- bounded parser recursion -- bounded source size - -## Examples - -### Minimal Model Graph - -```tinyagents -graph summarize { - start model - - node model { - kind model - model "default" - system "Summarize the user request." - next END - } -} -``` - -### Agent With Tools - -```tinyagents -graph support_agent { - start agent - - node agent { - kind agent - model "default" - system "Resolve support requests using tools when useful." - tools ["lookup_user", "create_ticket"] - routes { - tool_call -> tools - final -> END - } - } - - node tools { - kind tool_executor - next agent - } -} -``` - -### Human Review - -```tinyagents -graph approval_flow { - start draft - - node draft { - kind model - model "default" - system "Draft a response." - next review - } - - node review { - kind interrupt - prompt "Approve this response?" - routes { - approved -> send - rejected -> draft - } - } - - node send { - kind tool_executor - tools ["send_email"] - next END - } -} -``` - -## Formatting - -Formatter goals: - -- stable ordering within declarations -- preserve comments -- normalize indentation to two spaces -- one item per line for lists longer than one entry -- avoid rewriting prompt body content - -The formatter can come after parser and diagnostics. - -## Testkit - -`language::testkit` should include: - -- parse snapshot helper -- diagnostic snapshot helper -- compile helper with fake registries -- golden source fixtures -- round-trip formatter tests once formatter exists - -## Implementation Milestones - -### L1: Parser Skeleton - -- token model -- spans -- lexer -- parser for graph, start, node, next, routes - -### L2: Diagnostics - -- structured diagnostics -- duplicate node validation -- unknown route target validation - -### L3: Compiler Preview - -- compile topology into `GraphBuilder` -- support `kind model` -- support `kind agent` -- bind model names - -### L4: Tool Binding - -- validate tool names -- compile `tool_executor` -- add agent/tool graph example - -### L5: Policies And Subgraphs - -- parse timeout and retry -- bind subgraphs -- compile node policies - -### L6: Channels, Commands, And Fanout - -- parse channel declarations -- bind reducer registry entries -- lower `command` and `sends` -- compile join/barrier nodes - -### L7: Agent-Authored Graphs - -- compile generated source under parent run policy -- require review gates when policy marks generated graphs as untrusted -- store generated blueprint provenance -- expose graph diff and preview diagnostics - -### L8: Formatter - -- stable formatting -- golden tests +Continues in [`reference.md`](reference.md) (node kinds, binding to +Rust, state model, routes, policies, comments/strings, safety, +examples, formatting, testkit, milestones). diff --git a/docs/modules/expressive-language/implementation-status.md b/docs/modules/expressive-language/implementation-status.md index 3de1481..8edf335 100644 --- a/docs/modules/expressive-language/implementation-status.md +++ b/docs/modules/expressive-language/implementation-status.md @@ -74,7 +74,12 @@ The registry-backed binding path (`DEFAULT_NODE_KINDS`) accepts `agent`, `compile` rejects: duplicate nodes, missing/undefined `start`, unknown `next`/`route`/`edge`/`command goto`/`send`/`join` targets, duplicate route labels, and mixing static routing with `routes`. Registry binding additionally -checks model/tool/subgraph/router/reducer references and node kinds. +checks model/tool/subgraph/router/agent/script/reducer references and node +kinds. A single shared policy (`CapabilityResolver::classify_reference`) maps +each node kind to the reference it must resolve, so the compiler blueprint gate +and both `Resolver` paths cannot drift: `subagent` binds its `agent` reference +against the registered agents and `repl_agent` binds its `script` reference +against the registered scripts. ## Not yet implemented diff --git a/docs/modules/expressive-language/reference.md b/docs/modules/expressive-language/reference.md new file mode 100644 index 0000000..a3cdc71 --- /dev/null +++ b/docs/modules/expressive-language/reference.md @@ -0,0 +1,432 @@ +# Expressive Language: Node Kinds, Binding, State, Routes, Examples + +Continues from [`README.md`](README.md): node kinds, binding to Rust, +the state model, routes, policies, comments/strings, safety, examples, +formatting, testkit, and implementation milestones. + +## Node Kinds + +Initial built-in node kinds: + +### `agent` + +Uses the harness agent loop or one model call depending on config. + +Supported fields: + +- `model` +- `system` +- `prompt` +- `tools` +- `routes` +- `retry` +- `timeout` + +### `model` + +Single model invocation. Does not automatically execute tools. + +Supported fields: + +- `model` +- `system` +- `prompt` +- `routes` +- `retry` +- `timeout` + +### `tool_executor` + +Executes tool calls already present in state. + +Supported fields: + +- `tools` +- `next` +- `retry` +- `timeout` + +### `router` + +Routes based on a named route function provided from Rust. + +Supported fields: + +- `routes` +- `metadata` + +### `subgraph` + +Calls another compiled graph. + +Supported fields: + +- `graph` +- `next` +- `routes` + +### `subagent` + +Calls a registered harness agent as a graph node. + +Supported fields: + +- `agent` +- `input` +- `next` +- `routes` +- `retry` +- `timeout` +- `steering` + +Example: + +```tinyagents +node research { + kind subagent + agent "researcher" + steering { + parent allow ["add_instruction", "request_status", "cancel"] + human allow ["add_instruction", "pause", "resume", "cancel"] + delivery "safe_boundary" + } + next synthesize +} +``` + +Steering policies lower into harness steering policy and graph task policy. They +can narrow a child agent's model/tool/runtime limits but cannot grant +capabilities absent from the registry or parent run policy. + +### `repl_agent` + +Runs a registered REPL script or model-driven CodeAct loop through the harness +REPL runtime. + +Supported fields: + +- `model` +- `script` +- `tools` +- `routes` +- `retry` +- `timeout` + +### `interrupt` + +Emits a resumable human-in-the-loop interrupt. + +Supported fields: + +- `prompt` +- `options` +- `routes` +- `metadata` + +### `join` + +Waits for named upstream nodes or barrier channels before continuing. + +Supported fields: + +- `sources` +- `next` +- `timeout` + +## Binding To Rust + +The language should not define arbitrary Rust closures. Instead, it should bind +to Rust-provided registries: + +```rust +let workflow = LanguageCompiler::new() + .with_models(models) + .with_tools(tools) + .with_node_templates(templates) + .compile_source("support.rag", source)?; +``` + +Registries: + +- model registry +- tool registry +- agent registry +- node template registry +- route function registry +- reducer registry +- graph registry for subgraphs +- store registry +- middleware registry +- REPL script registry + +When a graph is generated by a REPL session, the session may call the compiler +with source text or an AST, but the compiler must use the same registries and +policy checks. Generated source can request capabilities only from the allowed +set attached to the parent run or registry namespace. + +This keeps source files declarative and prevents unsafe dynamic execution. + +## State Model + +Version 1 should keep state Rust-owned. The language can refer to standard +channels by convention and bind them to registered reducers: + +- `messages` +- `tool_calls` +- `structured_response` +- `metadata` +- `artifacts` +- `candidates` +- `usage` +- `interrupts` + +Example: + +```tinyagents +channel messages messages +channel candidates append +channel usage aggregate "usage_delta" +channel review overwrite +``` + +Future versions may add state schema declarations: + +```tinyagents +state SupportState { + messages: messages append + customer_id: string overwrite + ticket_id: string? overwrite +} +``` + +State schemas should be delayed until reducer-based graph execution exists. + +## Routes + +Routes are named outcomes. + +```tinyagents +routes { + tool_call -> tools + final -> END + escalate -> human_review +} +``` + +Rules: + +- route names are unique per node +- route targets must exist or be `END` +- route names are ASCII identifiers +- a node may use `routes` or `next`, not both +- `END` is reserved + +Future typed route support can generate Rust enums from route declarations. + +## Policies + +Node-level policies: + +```tinyagents +node agent { + timeout 30s + retry { + max_attempts: 3 + backoff: "exponential" + } +} +``` + +Graph-level defaults: + +```tinyagents +graph support_agent { + defaults { + timeout 60s + recursion_limit 50 + } +} +``` + +Policies lower into graph node policies and harness request policies. + +## Comments And Strings + +Comments: + +```tinyagents +// line comment +``` + +Strings: + +```tinyagents +system "single line" + +prompt """ +multi-line prompt +""" +``` + +Multi-line strings should preserve content exactly except for one predictable +dedent rule. + +## Safety + +The language must be safe to parse and validate from untrusted text. + +Safety rules: + +- no arbitrary code execution +- no filesystem access from language source +- no network access from language source +- no dynamic provider lookup without registry binding +- no environment variable interpolation in v1 +- bounded parser recursion +- bounded source size + +## Examples + +### Minimal Model Graph + +```tinyagents +graph summarize { + start model + + node model { + kind model + model "default" + system "Summarize the user request." + next END + } +} +``` + +### Agent With Tools + +```tinyagents +graph support_agent { + start agent + + node agent { + kind agent + model "default" + system "Resolve support requests using tools when useful." + tools ["lookup_user", "create_ticket"] + routes { + tool_call -> tools + final -> END + } + } + + node tools { + kind tool_executor + next agent + } +} +``` + +### Human Review + +```tinyagents +graph approval_flow { + start draft + + node draft { + kind model + model "default" + system "Draft a response." + next review + } + + node review { + kind interrupt + prompt "Approve this response?" + routes { + approved -> send + rejected -> draft + } + } + + node send { + kind tool_executor + tools ["send_email"] + next END + } +} +``` + +## Formatting + +Formatter goals: + +- stable ordering within declarations +- preserve comments +- normalize indentation to two spaces +- one item per line for lists longer than one entry +- avoid rewriting prompt body content + +The formatter can come after parser and diagnostics. + +## Testkit + +`language::testkit` should include: + +- parse snapshot helper +- diagnostic snapshot helper +- compile helper with fake registries +- golden source fixtures +- round-trip formatter tests once formatter exists + +## Implementation Milestones + +### L1: Parser Skeleton + +- token model +- spans +- lexer +- parser for graph, start, node, next, routes + +### L2: Diagnostics + +- structured diagnostics +- duplicate node validation +- unknown route target validation + +### L3: Compiler Preview + +- compile topology into `GraphBuilder` +- support `kind model` +- support `kind agent` +- bind model names + +### L4: Tool Binding + +- validate tool names +- compile `tool_executor` +- add agent/tool graph example + +### L5: Policies And Subgraphs + +- parse timeout and retry +- bind subgraphs +- compile node policies + +### L6: Channels, Commands, And Fanout + +- parse channel declarations +- bind reducer registry entries +- lower `command` and `sends` +- compile join/barrier nodes + +### L7: Agent-Authored Graphs + +- compile generated source under parent run policy +- require review gates when policy marks generated graphs as untrusted +- store generated blueprint provenance +- expose graph diff and preview diagnostics + +### L8: Formatter + +- stable formatting +- golden tests diff --git a/docs/modules/harness/README.md b/docs/modules/harness/README.md index 3a82be7..eaf9bed 100644 --- a/docs/modules/harness/README.md +++ b/docs/modules/harness/README.md @@ -214,6 +214,11 @@ Feature ownership: - `workspace`: per-agent filesystem/sandbox isolation, allowed-root descriptors, and fail-closed path enforcement for tools that touch real files. +Continued specification: [runtime.md](runtime.md) (tool registry, agent loop, +middleware, memory/stores) and +[observability-overview.md](observability-overview.md) (structured output, +events/streaming, errors, testkit, milestones). + Feature details: - [Context feature](context.md) @@ -465,477 +470,9 @@ registry default, and fallback policy. The selected model is recorded as a `ResolvedModel` in the response, event stream, run status, usage/cost records, and durable state when configured. -## Tool Registry - -The tool registry owns available tools and their schemas. - -```rust -pub struct ToolRegistry { - tools: HashMap>>, -} - -#[async_trait] -pub trait Tool: Send + Sync { - fn name(&self) -> &str; - fn description(&self) -> &str; - fn schema(&self) -> ToolSchema; - - async fn call( - &self, - state: &State, - ctx: &mut RunContext, - call: ToolCall, - ) -> Result; -} -``` - -Tool schema requirements: - -- name -- description -- JSON schema compatible input shape -- optional output schema -- safety metadata -- timeout override -- retry override -- model-visible flag for each argument -- injected-runtime argument declarations that are hidden from model schemas -- side-effect and idempotency metadata -- confirmation policy for destructive operations -- artifact output policy - -Tool call requirements: - -- `id` -- `name` -- `arguments` -- provider metadata -- originating model call id -- validation status -- retry attempt - -Tool result requirements: - -- `tool_call_id` -- `name` -- content -- raw structured value -- elapsed time -- error flag -- artifact references -- user-visible summary -- redacted event payload - -Tool names should default to ASCII `snake_case`. The registry should reject -duplicate names and invalid names. - -## Agent Loop - -The default loop is the LangChain-style model-tool loop: - -```text -input messages - -> build request - -> call model - -> if assistant has tool calls: - validate tool calls - execute tools - append tool messages - repeat - -> final assistant message -``` - -Detailed lifecycle: - -1. Create `RunConfig` and `RunContext`. -2. Load short-term memory for `thread_id` if configured. -3. Normalize input into messages. -4. Apply prompt templates and dynamic context. -5. Select model. -6. Select exposed tools. -7. Run `before_model` middleware, including prompt/cache-layout guards and - pre-call compression. -8. Invoke or stream the model through `wrap_model` middleware. -9. Run `on_model_delta` middleware for streamed chunks. -10. Run `after_model` middleware, including post-call compression and summary - persistence. -11. Emit model events and append assistant message. -12. If tool calls exist, validate name, schema, and limits. -13. Run `before_tool` middleware per call. -14. Execute tools serially or concurrently according to policy. -15. Run `on_tool_delta` middleware for tool progress streams. -16. Run `after_tool` middleware per result. -17. Append tool messages. -18. Repeat until no tool calls remain. -19. Validate structured output if configured. -20. Persist short-term memory. -21. Emit final event and return `AgentRun`. - -Hard limits: - -- `max_model_calls` -- `max_tool_calls` -- `max_concurrency` -- wall-clock timeout -- per-call timeout -- retry budget - -The loop must fail closed when a limit is reached. - -## Middleware - -Middleware is the main extension point for behavior that cuts across providers, -tools, and graph nodes. - -```rust -#[async_trait] -pub trait Middleware: Send + Sync { - async fn before_model( - &self, - state: &State, - ctx: &mut RunContext, - request: &mut ModelRequest, - ) -> Result<()>; - - async fn on_model_delta( - &self, - state: &State, - ctx: &mut RunContext, - delta: &mut ModelDelta, - ) -> Result<()>; - - async fn after_model( - &self, - state: &State, - ctx: &mut RunContext, - response: &mut ModelResponse, - ) -> Result<()>; - - async fn before_tool( - &self, - state: &State, - ctx: &mut RunContext, - call: &mut ToolCall, - ) -> Result<()>; - - async fn on_tool_delta( - &self, - state: &State, - ctx: &mut RunContext, - delta: &mut ToolDelta, - ) -> Result<()>; - - async fn after_tool( - &self, - state: &State, - ctx: &mut RunContext, - result: &mut ToolResult, - ) -> Result<()>; - - async fn on_error( - &self, - state: &State, - ctx: &mut RunContext, - error: &TinyAgentsError, - ) -> Result<()>; -} -``` - -Middleware ordering is stable and explicit. Middleware runs in registration -order for `before_*` hooks, registration order for streaming delta hooks, and -reverse order for `after_*` hooks. Wrap hooks should surround the full model or -tool operation when middleware needs setup, streaming inspection, and teardown -as one unit. - -Built-in middleware candidates: - -- tracing middleware -- retry middleware -- timeout middleware -- model fallback middleware -- token-bucket rate limiter middleware -- prompt cache layout guard middleware -- message trimming middleware -- summarization middleware -- context compression middleware -- transcript compression middleware -- retrieval compression middleware -- streaming delta compression middleware -- output compression middleware -- context editing middleware -- tool allowlist middleware -- dynamic tool selection middleware -- guardrail middleware -- PII detection/redaction middleware -- human-in-the-loop middleware -- shell/filesystem privilege boundary middleware -- structured output validator -- rate limiter - -Wrap hooks should exist in addition to before/after hooks. A wrap hook receives a -request plus a handler and can call the handler, replace the request, retry, -fallback to another model/tool, short-circuit with a response, or return a -control command. Before/after hooks are simpler and should remain available for -common mutation and observation cases. - -## Memory And Stores - -Memory and storage are related but not the same feature. `memory` owns -conversation semantics. `store` owns persistence backends. - -Memory has two layers conceptually: - -```text -short-term memory: thread-scoped conversation state -long-term store: cross-thread application data -``` - -Short-term memory: - -- keyed by `thread_id` -- loaded before an agent loop -- updated after successful loop completion -- optionally trimmed or summarized -- useful for conversation continuity - -Stores: - -- available through `RunContext` -- namespaced -- typed where possible -- usable by tools and middleware -- not automatically injected into prompts unless middleware does it -- reusable by memory, event recording, tool artifacts, and web UIs - -Suggested traits: - -```rust -#[async_trait] -pub trait ShortTermMemory: Send + Sync { - async fn load(&self, thread_id: &ThreadId) -> Result>; - async fn save(&self, thread_id: &ThreadId, state: &State) -> Result<()>; -} -``` - -The storage layer should be a separate harness feature: - -```rust -#[async_trait] -pub trait Store: Send + Sync { - async fn get(&self, key: StoreKey) -> Result>; - async fn put(&self, key: StoreKey, value: StoreValue) -> Result<()>; - async fn delete(&self, key: StoreKey) -> Result<()>; - async fn scan(&self, prefix: StoreKeyPrefix) -> Result>; -} - -#[async_trait] -pub trait AppendStore: Send + Sync { - async fn append(&self, stream: StoreStream, value: StoreValue) -> Result; - async fn read_from(&self, stream: StoreStream, offset: StoreOffset) -> Result>; -} - -pub enum StoreValue { - Json(serde_json::Value), - Bytes(Vec), - Text(String), -} -``` - -Initial store backends: - -- `InMemoryStore`: deterministic tests and examples. -- `JsonlStore`: append-only local development, replayable event logs, and cheap - debugging. -- `FileStore`: local artifacts such as tool outputs, provider payload snapshots, - and prompt fixtures. -- `MongoStore`: durable application/runtime records for server deployments. - -Later store backends: - -- SQLite for single-node durable local apps. -- Postgres for multi-tenant production apps. -- S3-compatible blob store for large artifacts. -- Redis for short-lived cache/session data. - -Store data classes: - -- run records -- thread records -- normalized messages -- event envelopes -- tool call records -- model call records -- structured outputs -- user/application memory -- tool artifacts and blobs - -Backend selection should be per store namespace: - -```rust -let stores = StoreRegistry::new() - .register("events", JsonlStore::new("./data/events.jsonl")) - .register("threads", MongoStore::new(mongo, "threads")) - .register("artifacts", FileStore::new("./data/artifacts")); -``` - -Store events should flow through `harness::events` or the registry event bus: - -- `store.read` -- `store.write` -- `store.append` -- `store.delete` -- `store.error` - -Sensitive store fields must support redaction before event emission. - -## Structured Output - -Structured output should support: - -- provider-native schema mode -- tool-call fallback mode -- JSON parsing mode for simple local models - -```rust -pub enum ResponseFormat { - Text, - JsonSchema(JsonSchema), - ProviderNative(JsonSchema), - ToolStrategy { tool_name: String, schema: JsonSchema }, -} -``` - -The final run result should keep messages and structured output separate: - -```rust -pub struct AgentRun { - pub state: State, - pub messages: Vec, - pub structured_response: Option, - pub events: Vec, -} -``` - -## Events And Streaming - -The harness event stream should be typed, not a string callback. - -```rust -pub enum AgentEvent { - RunStarted { run_id: RunId, thread_id: Option }, - ModelStarted { call_id: CallId, model: ModelName }, - ModelDelta { call_id: CallId, delta: MessageDelta }, - ModelCompleted { call_id: CallId, usage: Option }, - ToolStarted { call_id: CallId, tool_name: ToolName }, - ToolDelta { call_id: CallId, delta: ToolDelta }, - ToolCompleted { call_id: CallId, tool_name: ToolName }, - MiddlewareStarted { name: String }, - MiddlewareCompleted { name: String }, - RetryScheduled { call_id: CallId, attempt: usize }, - Custom { name: String, payload: serde_json::Value }, - RunCompleted { run_id: RunId }, - RunFailed { run_id: RunId, error: String }, -} -``` - -Streaming modes: - -- `messages`: model deltas and final messages -- `tools`: tool start, progress, result -- `updates`: state or memory updates -- `events`: all low-level events -- `final`: final output only - -## Errors - -Harness errors should distinguish: - -- invalid request -- missing model -- missing tool -- invalid tool schema -- invalid tool arguments -- provider authentication failure -- provider rate limit -- provider server error -- timeout -- retry exhausted -- structured output validation failure -- middleware failure -- memory failure - -Retry policy should only retry explicitly retryable classes by default: - -- network interruption -- timeout -- rate limit -- provider 5xx - -Do not retry authentication, schema, malformed request, or missing tool errors -unless a user explicitly overrides policy. - -## Testkit - -`harness::testkit` should be part of the early API. - -Utilities: - -- `FakeChatModel` -- `ScriptedChatModel` -- `FakeStreamingModel` -- `FakeTool` -- `InMemoryShortTermMemory` -- `InMemoryStore` -- `EventRecorder` -- deterministic ids -- deterministic clock -- trajectory assertions - -Example trajectory assertion: - -```rust -assert_trajectory(run.events()) - .model_called("default") - .tool_called("lookup_user") - .model_called("default") - .completed(); -``` - -## Implementation Milestones - -### H1: Current Minimal Traits - -- Keep `ChatMessage`. -- Keep `ChatModel`. -- Keep `Tool`. -- Add better tool call ids. - -### H2: Registries And Context - -- Add `ModelRegistry`. -- Add `ToolRegistry`. -- Add `RunConfig`. -- Add `RunContext`. - -### H3: Agent Loop - -- Implement model-tool loop. -- Enforce limits. -- Add fake model and fake tool tests. - -### H4: Middleware And Events - -- Add middleware stack. -- Add typed events. -- Add event recorder. - -### H5: Memory And Structured Output - -- Add short-term memory trait. -- Add store trait. -- Add structured response format. -### H6: Providers +--- -- Add feature-gated provider crates or modules. -- Start with mock and one hosted provider. +Continues in [`runtime.md`](runtime.md) (tool registry, agent loop, +middleware, memory/stores) and [`observability-overview.md`](observability-overview.md) +(structured output, events/streaming, errors, testkit, milestones). diff --git a/docs/modules/harness/model.md b/docs/modules/harness/model.md index e092e44..1adc6de 100644 --- a/docs/modules/harness/model.md +++ b/docs/modules/harness/model.md @@ -412,6 +412,12 @@ pub enum ModelStreamItem { } ``` +`MessageDelta` carries separate visible text and reasoning/thinking fields. +Provider adapters must map exposed thinking fragments, such as +OpenAI-compatible `reasoning_content` or `reasoning` deltas, onto the reasoning +field so UIs can render them without appending them to final assistant text. +Middleware receives the same side channel through `ModelDelta.reasoning`. + The final merged stream must be equivalent to `invoke` for the same provider where the provider supports both paths. If a provider emits cumulative usage rather than delta usage, the adapter must normalize it before sending usage diff --git a/docs/modules/harness/observability-overview.md b/docs/modules/harness/observability-overview.md new file mode 100644 index 0000000..65a12f7 --- /dev/null +++ b/docs/modules/harness/observability-overview.md @@ -0,0 +1,157 @@ +# Harness Observability: Structured Output, Events, Testkit + +Continues from [`README.md`](README.md) and [`runtime.md`](runtime.md): +structured output, events and streaming, errors, testkit, and +implementation milestones. + +## Structured Output + +Structured output should support: + +- provider-native schema mode +- tool-call fallback mode +- JSON parsing mode for simple local models + +```rust +pub enum ResponseFormat { + Text, + JsonSchema(JsonSchema), + ProviderNative(JsonSchema), + ToolStrategy { tool_name: String, schema: JsonSchema }, +} +``` + +The final run result should keep messages and structured output separate: + +```rust +pub struct AgentRun { + pub state: State, + pub messages: Vec, + pub structured_response: Option, + pub events: Vec, +} +``` + +## Events And Streaming + +The harness event stream should be typed, not a string callback. + +```rust +pub enum AgentEvent { + RunStarted { run_id: RunId, thread_id: Option }, + ModelStarted { call_id: CallId, model: ModelName }, + ModelDelta { call_id: CallId, delta: MessageDelta }, + ModelCompleted { call_id: CallId, usage: Option }, + ToolStarted { call_id: CallId, tool_name: ToolName }, + ToolDelta { call_id: CallId, delta: ToolDelta }, + ToolCompleted { call_id: CallId, tool_name: ToolName }, + MiddlewareStarted { name: String }, + MiddlewareCompleted { name: String }, + RetryScheduled { call_id: CallId, attempt: usize }, + Custom { name: String, payload: serde_json::Value }, + RunCompleted { run_id: RunId }, + RunFailed { run_id: RunId, error: String }, +} +``` + +Streaming modes: + +- `messages`: model deltas and final messages +- `tools`: tool start, progress, result +- `updates`: state or memory updates +- `events`: all low-level events +- `final`: final output only + +## Errors + +Harness errors should distinguish: + +- invalid request +- missing model +- missing tool +- invalid tool schema +- invalid tool arguments +- provider authentication failure +- provider rate limit +- provider server error +- timeout +- retry exhausted +- structured output validation failure +- middleware failure +- memory failure + +Retry policy should only retry explicitly retryable classes by default: + +- network interruption +- timeout +- rate limit +- provider 5xx + +Do not retry authentication, schema, malformed request, or missing tool errors +unless a user explicitly overrides policy. + +## Testkit + +`harness::testkit` should be part of the early API. + +Utilities: + +- `FakeChatModel` +- `ScriptedChatModel` +- `FakeStreamingModel` +- `FakeTool` +- `InMemoryShortTermMemory` +- `InMemoryStore` +- `EventRecorder` +- deterministic ids +- deterministic clock +- trajectory assertions + +Example trajectory assertion: + +```rust +assert_trajectory(run.events()) + .model_called("default") + .tool_called("lookup_user") + .model_called("default") + .completed(); +``` + +## Implementation Milestones + +### H1: Current Minimal Traits + +- Keep `ChatMessage`. +- Keep `ChatModel`. +- Keep `Tool`. +- Add better tool call ids. + +### H2: Registries And Context + +- Add `ModelRegistry`. +- Add `ToolRegistry`. +- Add `RunConfig`. +- Add `RunContext`. + +### H3: Agent Loop + +- Implement model-tool loop. +- Enforce limits. +- Add fake model and fake tool tests. + +### H4: Middleware And Events + +- Add middleware stack. +- Add typed events. +- Add event recorder. + +### H5: Memory And Structured Output + +- Add short-term memory trait. +- Add store trait. +- Add structured response format. + +### H6: Providers + +- Add feature-gated provider crates or modules. +- Start with mock and one hosted provider. diff --git a/docs/modules/harness/runtime.md b/docs/modules/harness/runtime.md new file mode 100644 index 0000000..cfde138 --- /dev/null +++ b/docs/modules/harness/runtime.md @@ -0,0 +1,328 @@ +# Harness Runtime: Tools, Agent Loop, Middleware, Memory + +Continues from [`README.md`](README.md): tool registry, agent loop, +middleware, and memory/stores. + +## Tool Registry + +The tool registry owns available tools and their schemas. + +```rust +pub struct ToolRegistry { + tools: HashMap>>, +} + +#[async_trait] +pub trait Tool: Send + Sync { + fn name(&self) -> &str; + fn description(&self) -> &str; + fn schema(&self) -> ToolSchema; + + async fn call( + &self, + state: &State, + ctx: &mut RunContext, + call: ToolCall, + ) -> Result; +} +``` + +Tool schema requirements: + +- name +- description +- JSON schema compatible input shape +- optional output schema +- safety metadata +- timeout override +- retry override +- model-visible flag for each argument +- injected-runtime argument declarations that are hidden from model schemas +- side-effect and idempotency metadata +- confirmation policy for destructive operations +- artifact output policy + +Tool call requirements: + +- `id` +- `name` +- `arguments` +- provider metadata +- originating model call id +- validation status +- retry attempt + +Tool result requirements: + +- `tool_call_id` +- `name` +- content +- raw structured value +- elapsed time +- error flag +- artifact references +- user-visible summary +- redacted event payload + +Tool names should default to ASCII `snake_case`. The registry should reject +duplicate names and invalid names. + +## Agent Loop + +The default loop is the LangChain-style model-tool loop: + +```text +input messages + -> build request + -> call model + -> if assistant has tool calls: + validate tool calls + execute tools + append tool messages + repeat + -> final assistant message +``` + +Detailed lifecycle: + +1. Create `RunConfig` and `RunContext`. +2. Load short-term memory for `thread_id` if configured. +3. Normalize input into messages. +4. Apply prompt templates and dynamic context. +5. Select model. +6. Select exposed tools. +7. Run `before_model` middleware, including prompt/cache-layout guards and + pre-call compression. +8. Invoke or stream the model through `wrap_model` middleware. +9. Run `on_model_delta` middleware for streamed chunks. +10. Run `after_model` middleware, including post-call compression and summary + persistence. +11. Emit model events and append assistant message. +12. If tool calls exist, validate name, schema, and limits. +13. Run `before_tool` middleware per call. +14. Execute tools serially or concurrently according to policy. +15. Run `on_tool_delta` middleware for tool progress streams. +16. Run `after_tool` middleware per result. +17. Append tool messages. +18. Repeat until no tool calls remain. +19. Validate structured output if configured. +20. Persist short-term memory. +21. Emit final event and return `AgentRun`. + +Hard limits: + +- `max_model_calls` +- `max_tool_calls` +- `max_concurrency` +- wall-clock timeout +- per-call timeout +- retry budget + +The loop must fail closed when a limit is reached. + +## Middleware + +Middleware is the main extension point for behavior that cuts across providers, +tools, and graph nodes. + +```rust +#[async_trait] +pub trait Middleware: Send + Sync { + async fn before_model( + &self, + state: &State, + ctx: &mut RunContext, + request: &mut ModelRequest, + ) -> Result<()>; + + async fn on_model_delta( + &self, + state: &State, + ctx: &mut RunContext, + delta: &mut ModelDelta, + ) -> Result<()>; + + async fn after_model( + &self, + state: &State, + ctx: &mut RunContext, + response: &mut ModelResponse, + ) -> Result<()>; + + async fn before_tool( + &self, + state: &State, + ctx: &mut RunContext, + call: &mut ToolCall, + ) -> Result<()>; + + async fn on_tool_delta( + &self, + state: &State, + ctx: &mut RunContext, + delta: &mut ToolDelta, + ) -> Result<()>; + + async fn after_tool( + &self, + state: &State, + ctx: &mut RunContext, + result: &mut ToolResult, + ) -> Result<()>; + + async fn on_error( + &self, + state: &State, + ctx: &mut RunContext, + error: &TinyAgentsError, + ) -> Result<()>; +} +``` + +Middleware ordering is stable and explicit. Middleware runs in registration +order for `before_*` hooks, registration order for streaming delta hooks, and +reverse order for `after_*` hooks. Wrap hooks should surround the full model or +tool operation when middleware needs setup, streaming inspection, and teardown +as one unit. + +Built-in middleware candidates: + +- tracing middleware +- retry middleware +- timeout middleware +- model fallback middleware +- token-bucket rate limiter middleware +- prompt cache layout guard middleware +- message trimming middleware +- summarization middleware +- context compression middleware +- transcript compression middleware +- retrieval compression middleware +- streaming delta compression middleware +- output compression middleware +- context editing middleware +- tool allowlist middleware +- dynamic tool selection middleware +- guardrail middleware +- PII detection/redaction middleware +- human-in-the-loop middleware +- shell/filesystem privilege boundary middleware +- structured output validator +- rate limiter + +Wrap hooks should exist in addition to before/after hooks. A wrap hook receives a +request plus a handler and can call the handler, replace the request, retry, +fallback to another model/tool, short-circuit with a response, or return a +control command. Before/after hooks are simpler and should remain available for +common mutation and observation cases. + +## Memory And Stores + +Memory and storage are related but not the same feature. `memory` owns +conversation semantics. `store` owns persistence backends. + +Memory has two layers conceptually: + +```text +short-term memory: thread-scoped conversation state +long-term store: cross-thread application data +``` + +Short-term memory: + +- keyed by `thread_id` +- loaded before an agent loop +- updated after successful loop completion +- optionally trimmed or summarized +- useful for conversation continuity + +Stores: + +- available through `RunContext` +- namespaced +- typed where possible +- usable by tools and middleware +- not automatically injected into prompts unless middleware does it +- reusable by memory, event recording, tool artifacts, and web UIs + +Suggested traits: + +```rust +#[async_trait] +pub trait ShortTermMemory: Send + Sync { + async fn load(&self, thread_id: &ThreadId) -> Result>; + async fn save(&self, thread_id: &ThreadId, state: &State) -> Result<()>; +} +``` + +The storage layer should be a separate harness feature: + +```rust +#[async_trait] +pub trait Store: Send + Sync { + async fn get(&self, key: StoreKey) -> Result>; + async fn put(&self, key: StoreKey, value: StoreValue) -> Result<()>; + async fn delete(&self, key: StoreKey) -> Result<()>; + async fn scan(&self, prefix: StoreKeyPrefix) -> Result>; +} + +#[async_trait] +pub trait AppendStore: Send + Sync { + async fn append(&self, stream: StoreStream, value: StoreValue) -> Result; + async fn read_from(&self, stream: StoreStream, offset: StoreOffset) -> Result>; +} + +pub enum StoreValue { + Json(serde_json::Value), + Bytes(Vec), + Text(String), +} +``` + +Initial store backends: + +- `InMemoryStore`: deterministic tests and examples. +- `JsonlStore`: append-only local development, replayable event logs, and cheap + debugging. +- `FileStore`: local artifacts such as tool outputs, provider payload snapshots, + and prompt fixtures. +- `MongoStore`: durable application/runtime records for server deployments. + +Later store backends: + +- SQLite for single-node durable local apps. +- Postgres for multi-tenant production apps. +- S3-compatible blob store for large artifacts. +- Redis for short-lived cache/session data. + +Store data classes: + +- run records +- thread records +- normalized messages +- event envelopes +- tool call records +- model call records +- structured outputs +- user/application memory +- tool artifacts and blobs + +Backend selection should be per store namespace: + +```rust +let stores = StoreRegistry::new() + .register("events", JsonlStore::new("./data/events.jsonl")) + .register("threads", MongoStore::new(mongo, "threads")) + .register("artifacts", FileStore::new("./data/artifacts")); +``` + +Store events should flow through `harness::events` or the registry event bus: + +- `store.read` +- `store.write` +- `store.append` +- `store.delete` +- `store.error` + +Sensitive store fields must support redaction before event emission. + diff --git a/docs/modules/harness/streaming.md b/docs/modules/harness/streaming.md index f5850a0..1a8df84 100644 --- a/docs/modules/harness/streaming.md +++ b/docs/modules/harness/streaming.md @@ -67,6 +67,7 @@ consumer owns the run cancellation token. Streaming adapters must merge chunks deterministically: - text chunks preserve order +- reasoning/thinking chunks preserve order on a side channel - content block indexes are respected - tool-call chunks are correlated by id or index - cumulative usage is converted into deltas or clearly marked cumulative diff --git a/docs/modules/registry/README.md b/docs/modules/registry/README.md index 090d496..f19c57a 100644 --- a/docs/modules/registry/README.md +++ b/docs/modules/registry/README.md @@ -12,6 +12,8 @@ server. ## Detailed Module Docs - [Design](design.md) + - [Events and persistence](events.md) + - [Operations and lifecycle](operations.md) - [Model catalog and local snapshots](model-catalog.md) ## Responsibilities @@ -73,6 +75,8 @@ Documentation layout: docs/modules/registry/ README.md design.md + events.md + operations.md model-catalog.md model-catalog.snapshot.json ``` @@ -112,9 +116,10 @@ live handles. ### Component kinds -`ComponentKind` partitions the registry namespace and now has **11** variants. -Alongside `Model`, `Tool`, `Graph`, `Router`, `Reducer`, `Store`, and `Agent`, -four kinds cover the runtime's durable roles: +`ComponentKind` partitions the registry namespace and now has **12** variants. +Alongside `Model`, `Tool`, `Graph`, `Router`, `Reducer`, `Store`, `Agent`, and +`Script` (a REPL script a `repl_agent` node may reference), four kinds cover the +runtime's durable roles: | Kind | `as_str` | | --- | --- | diff --git a/docs/modules/registry/design.md b/docs/modules/registry/design.md index fc210f3..6219bf5 100644 --- a/docs/modules/registry/design.md +++ b/docs/modules/registry/design.md @@ -461,510 +461,9 @@ Resolution should return both an executable handle and a durable `ResolvedModel` record. The handle is process-local; the record is persisted in state, events, checkpoints, usage, and cost rows. -## Store And Checkpointer Registration -Stores and checkpointers should be registry components so agents and graphs can -share them without globals. +--- -Registered stores: - -- short-term memory -- long-term key/value store -- vector store -- graph checkpointer -- event sink - -Store events: - -- `store.registered` -- `store.read` -- `store.write` -- `store.failed` - -Store read/write events may need redaction controls. - -## Listener Registration - -Listeners observe events. They should be trait objects so web UIs, tests, -OpenTelemetry exporters, logs, and custom dashboards can all subscribe. - -```rust -#[async_trait] -pub trait EventListener: Send + Sync { - fn id(&self) -> ListenerId; - fn filter(&self) -> EventFilter; - async fn on_event(&self, event: RegistryEvent) -> Result<()>; -} -``` - -Listener examples: - -- in-memory event recorder -- stdout logger -- tracing subscriber bridge -- websocket broadcaster -- Server-Sent Events broadcaster -- OpenTelemetry exporter -- metrics collector -- test assertion recorder - -Listener lifecycle events: - -- `listener.registered` -- `listener.failed` -- `listener.removed` - -## Event Model - -All events should share an envelope. - -```rust -pub struct RegistryEvent { - pub id: EventId, - pub time: SystemTime, - pub run: Option, - pub parent: Option, - pub component: ComponentId, - pub kind: EventKind, - pub level: EventLevel, - pub tags: Vec, - pub metadata: serde_json::Value, - pub payload: EventPayload, -} -``` - -Correlation fields: - -```rust -pub struct RunRef { - pub run_id: RunId, - pub thread_id: Option, - pub parent_run_id: Option, - pub root_run_id: RunId, - pub span_id: SpanId, -} -``` - -`RunRef` should be created from one propagated `RunConfig`. Do not pass tracing -ids, registry ids, and tool ids through separate ad hoc channels. - -Metadata inheritance: - -```rust -pub struct RunConfig { - pub run_id: RunId, - pub parent_run_id: Option, - pub root_run_id: RunId, - pub thread_id: Option, - pub tags: MetadataScope>, - pub metadata: MetadataScope, - pub max_concurrency: Option, - pub recursion_limit: Option, - pub configurable: serde_json::Value, -} - -pub struct MetadataScope { - pub local: T, - pub inherited: T, -} -``` - -Inherited tags and metadata propagate to child components. Local metadata stays -on the current component span. This distinction lets UIs filter by durable -context like tenant, thread, agent, or tool package without polluting every child -span with local retry and provider details. - -Event levels: - -- trace -- debug -- info -- warn -- error - -Event payloads: - -```rust -pub enum EventPayload { - Agent(AgentEvent), - Graph(GraphEvent), - Model(ModelEvent), - Tool(ToolEvent), - Store(StoreEvent), - Listener(ListenerEvent), - Registry(RegistryLifecycleEvent), - Custom(serde_json::Value), -} -``` - -The envelope gives external listeners enough context to render parallel runs -without guessing parent/child relationships. - -Recommended serialized event names: - -```text -on_registry_lookup_start -on_registry_lookup_end -on_registry_resolve_start -on_registry_resolve_end -on_component_instantiate_start -on_component_instantiate_end -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 -on_tool_error -on_model_start -on_model_stream -on_model_end -on_model_error -``` - -The Rust enum names can be idiomatic, but serialized event names should remain -stable for web UIs and persisted traces. - -## Event Bus - -The event bus is responsible for fanout. - -```rust -#[async_trait] -pub trait EventBus: Send + Sync { - async fn emit(&self, event: RegistryEvent) -> Result<()>; - async fn subscribe(&self, filter: EventFilter) -> Result; -} -``` - -Implementation requirements: - -- nonblocking emit path where possible -- bounded queues -- backpressure policy -- listener error isolation -- event redaction hook -- deterministic in-memory implementation for tests -- async stream subscription for web UIs - -Backpressure policies: - -- block -- drop oldest -- drop newest -- fail run - -Default policy should be bounded and fail noisy in tests, but avoid crashing a -production run because a UI listener disconnected. - -## Event Filters - -Listeners need filters. - -```rust -pub struct EventFilter { - pub kinds: Option>, - pub component_kinds: Option>, - pub component_names: Option>, - pub run_id: Option, - pub thread_id: Option, - pub tags: Vec, - pub min_level: EventLevel, -} -``` - -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 - -## Static And Dynamic Components - -The registry should distinguish static registered components from runtime -components. - -Static components: - -- registered before a run -- listed in discovery APIs -- have durable ids -- validated when graphs compile - -Runtime components: - -- supplied by middleware or run config -- may not be discoverable before a run -- must be executable through a dynamic resolver -- must still emit events with a component id or temporary runtime id - -Dynamic tools are useful for middleware and subagents, but unknown tool names -should not silently execute. A dynamic resolver must explicitly accept the tool -name and produce a registered-like descriptor. - -## Parallel Agents And Hierarchical Runs - -Parallel execution requires explicit hierarchy. - -When an agent spawns child agents or graph branches: - -- every child run gets a new `run_id` -- every child run inherits `root_run_id` -- every child run records `parent_run_id` -- tags and metadata inherit by default -- child components may append tags -- event ordering is per listener stream, not global causality - -Example: - -```text -root run: support_agent - child run: retrieval_graph - node: search - node: rerank - child run: draft_agent - model: default - tool: lookup_user -``` - -The web UI should be able to reconstruct this tree from event envelopes alone. - -## Web UI Integration - -The registry should support UI-facing APIs: - -```rust -GET /registry/components -GET /registry/components/{kind}/{name} -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 -discovery APIs should make one straightforward. - -UI event needs: - -- stable component ids -- display names -- graph topology metadata -- tool schemas -- run tree correlation -- streaming model deltas -- interrupt payloads -- checkpoint ids -- final outputs -- redacted error details - -## Stream Transformers - -Stream transformers derive UI-friendly views from raw events. - -```rust -#[async_trait] -pub trait StreamTransformer: Send + Sync { - fn id(&self) -> &str; - fn input_filter(&self) -> EventFilter; - async fn transform(&self, event: RegistryEvent) -> Result>; -} -``` - -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. - -## Redaction And Safety - -Events may contain sensitive data. The registry must support redaction. - -Redaction hooks: - -- before event leaves component -- before event reaches bus -- per listener - -Fields that may require redaction: - -- API keys -- provider request bodies -- tool arguments -- tool outputs -- memory values -- user messages -- environment details - -Default behavior: - -- metadata values are JSON only -- known secret keys are redacted -- raw provider payloads are opt-in -- listeners can request safe summaries - -## Registration Lifecycle - -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?; -``` - -Lifecycle events: - -- `registry.component_registered` -- `registry.component_replaced` -- `registry.component_removed` -- `registry.lookup_started` -- `registry.lookup_completed` -- `registry.resolve_started` -- `registry.resolve_completed` -- `registry.instantiate_started` -- `registry.instantiate_completed` -- `registry.alias_resolved` -- `registry.validation_failed` - -Validation: - -- duplicate names -- invalid names -- missing dependencies -- incompatible state type where statically knowable -- tool schema invalid -- graph references missing component - -## Discovery API - -Discovery lets UIs and orchestrators inspect capabilities. - -```rust -pub trait Discoverable { - fn metadata(&self) -> ComponentMetadata; - fn dependencies(&self) -> Vec; -} -``` - -Discovery output should include: - -- component id -- component kind -- description -- tags -- schema -- dependencies -- event kinds emitted -- run modes supported - -## Error Model - -Registry errors should distinguish: - -- duplicate component -- component not found -- invalid component name -- invalid schema -- missing dependency -- listener failure -- event queue full -- component type mismatch -- redaction failure -- registration locked - -Listener failures should emit events but should not fail the run unless the -listener is marked required. - -## Testkit - -`registry::testkit` should include: - -- in-memory registry builder -- event recorder listener -- event snapshot assertions -- fake tool registration -- fake graph registration -- fake agent registration -- deterministic event ids -- deterministic timestamps - -Example: - -```rust -let recorder = EventRecorder::new(); -let registry = TestRegistry::new() - .with_listener(recorder.clone()) - .with_tool(fake_tool("lookup_user")) - .with_agent(fake_agent("support_agent")) - .build(); - -registry.agent("support_agent")?.invoke(input).await?; - -recorder.assert() - .saw("agent.started") - .saw("tool.started") - .saw("tool.completed") - .saw("agent.completed"); -``` - -## Implementation Milestones - -### R1: Component Registries - -- `ComponentId` -- `ComponentMetadata` -- `ToolRegistry` -- `ModelRegistry` -- `GraphRegistry` -- duplicate validation - -### R2: Event Envelope - -- `RegistryEvent` -- `RunRef` -- `EventPayload` -- in-memory event bus - -### R3: Listeners - -- `EventListener` -- filters -- event recorder -- stdout listener - -### R4: Agent And Graph Registration - -- registered agent trait -- registered graph wrapper -- lookup and invoke helpers - -### R5: Parallel Run Correlation - -- parent/root run ids -- child run creation -- inherited tags and metadata -- event tree assertions - -### R6: UI Streaming Surface - -- async subscriptions -- websocket/SSE example -- redaction policy -- component discovery JSON +Continues in [`events.md`](events.md) (store/checkpointer 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 new file mode 100644 index 0000000..449366c --- /dev/null +++ b/docs/modules/registry/events.md @@ -0,0 +1,231 @@ +# Registry Events And Persistence + +Continues from [`design.md`](design.md): store and checkpointer +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. + +Registered stores: + +- short-term memory +- long-term key/value store +- vector store +- graph checkpointer +- event sink + +Store events: + +- `store.registered` +- `store.read` +- `store.write` +- `store.failed` + +Store read/write events may need redaction controls. + +## Listener Registration + +Listeners observe events. They should be trait objects so web UIs, tests, +OpenTelemetry exporters, logs, and custom dashboards can all subscribe. + +```rust +#[async_trait] +pub trait EventListener: Send + Sync { + fn id(&self) -> ListenerId; + fn filter(&self) -> EventFilter; + async fn on_event(&self, event: RegistryEvent) -> Result<()>; +} +``` + +Listener examples: + +- in-memory event recorder +- stdout logger +- tracing subscriber bridge +- websocket broadcaster +- Server-Sent Events broadcaster +- OpenTelemetry exporter +- metrics collector +- test assertion recorder + +Listener lifecycle events: + +- `listener.registered` +- `listener.failed` +- `listener.removed` + +## Event Model + +All events should share an envelope. + +```rust +pub struct RegistryEvent { + pub id: EventId, + pub time: SystemTime, + pub run: Option, + pub parent: Option, + pub component: ComponentId, + pub kind: EventKind, + pub level: EventLevel, + pub tags: Vec, + pub metadata: serde_json::Value, + pub payload: EventPayload, +} +``` + +Correlation fields: + +```rust +pub struct RunRef { + pub run_id: RunId, + pub thread_id: Option, + pub parent_run_id: Option, + pub root_run_id: RunId, + pub span_id: SpanId, +} +``` + +`RunRef` should be created from one propagated `RunConfig`. Do not pass tracing +ids, registry ids, and tool ids through separate ad hoc channels. + +Metadata inheritance: + +```rust +pub struct RunConfig { + pub run_id: RunId, + pub parent_run_id: Option, + pub root_run_id: RunId, + pub thread_id: Option, + pub tags: MetadataScope>, + pub metadata: MetadataScope, + pub max_concurrency: Option, + pub recursion_limit: Option, + pub configurable: serde_json::Value, +} + +pub struct MetadataScope { + pub local: T, + pub inherited: T, +} +``` + +Inherited tags and metadata propagate to child components. Local metadata stays +on the current component span. This distinction lets UIs filter by durable +context like tenant, thread, agent, or tool package without polluting every child +span with local retry and provider details. + +Event levels: + +- trace +- debug +- info +- warn +- error + +Event payloads: + +```rust +pub enum EventPayload { + Agent(AgentEvent), + Graph(GraphEvent), + Model(ModelEvent), + Tool(ToolEvent), + Store(StoreEvent), + Listener(ListenerEvent), + Registry(RegistryLifecycleEvent), + Custom(serde_json::Value), +} +``` + +The envelope gives external listeners enough context to render parallel runs +without guessing parent/child relationships. + +Recommended serialized event names: + +```text +on_registry_lookup_start +on_registry_lookup_end +on_registry_resolve_start +on_registry_resolve_end +on_component_instantiate_start +on_component_instantiate_end +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 +on_tool_error +on_model_start +on_model_stream +on_model_end +on_model_error +``` + +The Rust enum names can be idiomatic, but serialized event names should remain +stable for web UIs and persisted traces. + +## Event Bus + +The event bus is responsible for fanout. + +```rust +#[async_trait] +pub trait EventBus: Send + Sync { + async fn emit(&self, event: RegistryEvent) -> Result<()>; + async fn subscribe(&self, filter: EventFilter) -> Result; +} +``` + +Implementation requirements: + +- nonblocking emit path where possible +- bounded queues +- backpressure policy +- listener error isolation +- event redaction hook +- deterministic in-memory implementation for tests +- async stream subscription for web UIs + +Backpressure policies: + +- block +- drop oldest +- drop newest +- fail run + +Default policy should be bounded and fail noisy in tests, but avoid crashing a +production run because a UI listener disconnected. + +## Event Filters + +Listeners need filters. + +```rust +pub struct EventFilter { + pub kinds: Option>, + pub component_kinds: Option>, + pub component_names: Option>, + pub run_id: Option, + pub thread_id: Option, + pub tags: Vec, + pub min_level: EventLevel, +} +``` + +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/operations.md b/docs/modules/registry/operations.md new file mode 100644 index 0000000..2e7191d --- /dev/null +++ b/docs/modules/registry/operations.md @@ -0,0 +1,289 @@ +# Registry Operations And Lifecycle + +Continues from [`design.md`](design.md) and [`events.md`](events.md): +static/dynamic components, parallel agents, web UI integration, stream +transformers, redaction, registration lifecycle, discovery, error model, +testkit, and implementation milestones. + +## Static And Dynamic Components + +The registry should distinguish static registered components from runtime +components. + +Static components: + +- registered before a run +- listed in discovery APIs +- have durable ids +- validated when graphs compile + +Runtime components: + +- supplied by middleware or run config +- may not be discoverable before a run +- must be executable through a dynamic resolver +- must still emit events with a component id or temporary runtime id + +Dynamic tools are useful for middleware and subagents, but unknown tool names +should not silently execute. A dynamic resolver must explicitly accept the tool +name and produce a registered-like descriptor. + +## Parallel Agents And Hierarchical Runs + +Parallel execution requires explicit hierarchy. + +When an agent spawns child agents or graph branches: + +- every child run gets a new `run_id` +- every child run inherits `root_run_id` +- every child run records `parent_run_id` +- tags and metadata inherit by default +- child components may append tags +- event ordering is per listener stream, not global causality + +Example: + +```text +root run: support_agent + child run: retrieval_graph + node: search + node: rerank + child run: draft_agent + model: default + tool: lookup_user +``` + +The web UI should be able to reconstruct this tree from event envelopes alone. + +## Web UI Integration + +The registry should support UI-facing APIs: + +```rust +GET /registry/components +GET /registry/components/{kind}/{name} +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 +discovery APIs should make one straightforward. + +UI event needs: + +- stable component ids +- display names +- graph topology metadata +- tool schemas +- run tree correlation +- streaming model deltas +- interrupt payloads +- checkpoint ids +- final outputs +- redacted error details + +## Stream Transformers + +Stream transformers derive UI-friendly views from raw events. + +```rust +#[async_trait] +pub trait StreamTransformer: Send + Sync { + fn id(&self) -> &str; + fn input_filter(&self) -> EventFilter; + async fn transform(&self, event: RegistryEvent) -> Result>; +} +``` + +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. + +## Redaction And Safety + +Events may contain sensitive data. The registry must support redaction. + +Redaction hooks: + +- before event leaves component +- before event reaches bus +- per listener + +Fields that may require redaction: + +- API keys +- provider request bodies +- tool arguments +- tool outputs +- memory values +- user messages +- environment details + +Default behavior: + +- metadata values are JSON only +- known secret keys are redacted +- raw provider payloads are opt-in +- listeners can request safe summaries + +## Registration Lifecycle + +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?; +``` + +Lifecycle events: + +- `registry.component_registered` +- `registry.component_replaced` +- `registry.component_removed` +- `registry.lookup_started` +- `registry.lookup_completed` +- `registry.resolve_started` +- `registry.resolve_completed` +- `registry.instantiate_started` +- `registry.instantiate_completed` +- `registry.alias_resolved` +- `registry.validation_failed` + +Validation: + +- duplicate names +- invalid names +- missing dependencies +- incompatible state type where statically knowable +- tool schema invalid +- graph references missing component + +## Discovery API + +Discovery lets UIs and orchestrators inspect capabilities. + +```rust +pub trait Discoverable { + fn metadata(&self) -> ComponentMetadata; + fn dependencies(&self) -> Vec; +} +``` + +Discovery output should include: + +- component id +- component kind +- description +- tags +- schema +- dependencies +- event kinds emitted +- run modes supported + +## Error Model + +Registry errors should distinguish: + +- duplicate component +- component not found +- invalid component name +- invalid schema +- missing dependency +- listener failure +- event queue full +- component type mismatch +- redaction failure +- registration locked + +Listener failures should emit events but should not fail the run unless the +listener is marked required. + +## Testkit + +`registry::testkit` should include: + +- in-memory registry builder +- event recorder listener +- event snapshot assertions +- fake tool registration +- fake graph registration +- fake agent registration +- deterministic event ids +- deterministic timestamps + +Example: + +```rust +let recorder = EventRecorder::new(); +let registry = TestRegistry::new() + .with_listener(recorder.clone()) + .with_tool(fake_tool("lookup_user")) + .with_agent(fake_agent("support_agent")) + .build(); + +registry.agent("support_agent")?.invoke(input).await?; + +recorder.assert() + .saw("agent.started") + .saw("tool.started") + .saw("tool.completed") + .saw("agent.completed"); +``` + +## Implementation Milestones + +### R1: Component Registries + +- `ComponentId` +- `ComponentMetadata` +- `ToolRegistry` +- `ModelRegistry` +- `GraphRegistry` +- duplicate validation + +### R2: Event Envelope + +- `RegistryEvent` +- `RunRef` +- `EventPayload` +- in-memory event bus + +### R3: Listeners + +- `EventListener` +- filters +- event recorder +- stdout listener + +### R4: Agent And Graph Registration + +- registered agent trait +- registered graph wrapper +- lookup and invoke helpers + +### R5: Parallel Run Correlation + +- parent/root run ids +- child run creation +- inherited tags and metadata +- event tree assertions + +### R6: UI Streaming Surface + +- async subscriptions +- websocket/SSE example +- redaction policy +- component discovery JSON diff --git a/docs/modules/repl-language/README.md b/docs/modules/repl-language/README.md index 1a2faa3..92d0c8c 100644 --- a/docs/modules/repl-language/README.md +++ b/docs/modules/repl-language/README.md @@ -18,6 +18,7 @@ system, recursion policy, or run limits. ## Detailed Module Docs - [Design](design.md) + - [RLM feature map, embedding, safety, events, testkit](operations.md) ## Responsibilities diff --git a/docs/modules/repl-language/design.md b/docs/modules/repl-language/design.md index 79ee1f3..168fc42 100644 --- a/docs/modules/repl-language/design.md +++ b/docs/modules/repl-language/design.md @@ -461,347 +461,9 @@ answer.ready = true; The function style should be the default because it is harder to accidentally partially mutate. -## RLM Feature Map - -The goal is to port the useful `rlm` behavior into TinyAgents without porting -Python's unsafe local execution model. - -| `rlm` feature | TinyAgents REPL equivalent | -| --------------------------- | ----------------------------------------------------- | -| Python `context` variable | Rhai `context` variable | -| Python persistent locals | `ReplSession::variables` | -| fenced `repl` blocks | fenced `ragsh` blocks | -| `llm_query` | `model_query` | -| `llm_query_batched` | `model_query_batched` | -| `rlm_query` | `agent_query` or `repl_query` | -| `rlm_query_batched` | `agent_query_batched` or `repl_query_batched` | -| custom Python tools | registered Rust tool capabilities | -| generated Python programs | `.ragsh` cells plus generated `.rag` graph blueprints | -| `SHOW_VARS()` | `show_vars()` | -| `answer["ready"] = True` | `answer(...)` | -| max iterations | `ReplPolicy::max_iterations` for CodeAct loops | -| max depth | graph/harness recursion policy | -| max budget | harness cost policy | -| token compaction | harness summarization feature | -| JSONL trajectory logger | typed event stream plus store backend | -| Docker/cloud REPL isolation | future `PythonSandboxRepl` backend | - -## CodeAct Loop - -A model-driven REPL agent has this lifecycle: - -1. Create `ReplSession`. -2. Load `context`, `state`, `messages`, `history`, and `run` variables. -3. Build a model request explaining the available REPL functions. -4. Invoke the model through the harness. -5. Extract fenced `ragsh` blocks from the assistant message. -6. Execute each block in the REPL session. -7. Capture stdout, changed variables, call records, events, and errors. -8. Append a compact execution result as the next user message. -9. Repeat until `answer(...)` is called or limits are reached. -10. Persist events, usage, cost, and final answer. - -This loop is a harness feature. When used inside a graph node, the graph still -owns node routing, checkpointing, interrupts, recursion depth, and failure -policy. - -If the model writes `.rag` source, the loop should treat it as a graph proposal. -The REPL may validate, diff, compile, and run that proposal only through the -expressive-language compiler and the graph registry policy. This is how an -agent can define its own graph without acquiring arbitrary topology mutation or -host-code execution privileges. - -## Example Session -```rhai -let lines = context.split("\n"); -let candidates = []; - -for line in lines { - if line.contains("SECRET_NUMBER=") { - candidates.push(line); - } -} - -emit("candidates_found", #{ count: candidates.len() }); - -let result = model_query(#{ - model: "default", - prompt: "Return only the digits from this candidate line:\n" + candidates[0] -}); - -answer(result); -``` - -## Example Graph Node - -```tinyagents -graph support_repl { - start investigate - - node investigate { - kind repl_agent - model "default" - script "support-investigation.ragsh" - tools ["lookup_user", "create_ticket"] - routes { - final -> END - needs_review -> review - } - } - - node review { - kind interrupt - prompt "Approve escalation?" - routes { - approved -> END - rejected -> investigate - } - } -} -``` - -The `repl_agent` node is a harness-backed node template. It may execute a fixed -script, a model-driven CodeAct loop, or a combination where a fixed prologue -sets up variables before the model starts writing cells. - -## Rhai Embedding Plan - -The Rhai runtime should be isolated behind an interface so future Python or WASM -backends can reuse the same TinyAgents semantics. - -```rust -#[async_trait] -pub trait ReplBackend: Send { - async fn execute_cell( - &mut self, - session: &mut ReplSession, - source: SourceCell, - ) -> Result; -} - -pub struct RhaiReplBackend { - engine: rhai::Engine, - ast_cache: AstCache, -} -``` - -Rhai-specific requirements: - -- configure `Engine::set_max_operations` -- disable or avoid unneeded packages -- register only TinyAgents capability functions -- expose data through `Dynamic`, maps, and arrays with explicit conversion -- compile and cache ASTs for repeated scripts -- keep each session's `Scope` separate -- restore reserved names after each cell -- truncate stdout and returned values according to policy -- convert Rhai errors into structured diagnostics with spans - -Async adapter requirement: - -Rhai host functions are easiest to expose as synchronous functions. TinyAgents -model, tool, and graph calls are async. The backend should not hide blocking in -unbounded threads. Use one of these designs: - -1. command recording: Rhai functions create `ReplCommand` values, then the - Rust async runtime executes those commands after the cell -2. blocking bridge: host functions call into a bounded runtime handle with - strict timeouts -3. staged syntax: `let x = model_query(...)` is transformed before evaluation - into host-executed calls - -Recommendation for v1: use a blocking bridge only in examples and tests, but -design the public API around command recording. Command recording is easier to -make deterministic and safer under async graph execution. - -## Python Compatibility Backend - -Python should be a compatibility backend, not the default embedded runtime. - -```rust -pub struct PythonSandboxReplBackend { - sandbox: SandboxClient, -} -``` - -Potential use cases: - -- training model behavior that already expects Python -- local research workflows -- compatibility with RLM-style prompts -- data-heavy scripts where Python libraries are explicitly useful - -Requirements: - -- must run out of process -- must have no direct host filesystem access by default -- must communicate through a framed JSON protocol -- must expose the same TinyAgents capability functions -- must enforce the same `ReplPolicy` -- must emit the same `ReplEvent` stream - -This lets TinyAgents support Python-like RLM ergonomics without making Python a -trusted in-process extension language. - -## Safety - -Safety rules: - -- no arbitrary filesystem access in the default Rhai backend -- no environment variable interpolation from scripts -- no direct network access -- no process spawning -- no unregistered native functions -- bounded script size -- bounded operation count -- bounded output size -- bounded model/tool/graph calls -- bounded recursion depth -- bounded concurrency -- typed conversion at every capability boundary -- redaction before event and store writes - -The REPL is an orchestration surface, not a privilege escalation surface. - -## Events - -The REPL event stream should compose with graph and harness events. - -```rust -pub enum ReplEvent { - SessionStarted { session_id: SessionId, run_id: RunId }, - CellStarted { cell_id: CellId, source_name: String }, - CellStdout { cell_id: CellId, chunk: String }, - CellCompleted { cell_id: CellId, elapsed: Duration }, - CellFailed { cell_id: CellId, diagnostic: Diagnostic }, - VariableChanged { cell_id: CellId, name: String }, - CapabilityCallStarted { cell_id: CellId, call_id: CallId, name: String }, - CapabilityCallCompleted { cell_id: CellId, call_id: CallId }, - GraphBlueprintDefined { cell_id: CellId, graph_name: String }, - GraphBlueprintValidated { cell_id: CellId, graph_name: String }, - GraphBlueprintCompiled { cell_id: CellId, graph_name: String }, - GraphBlueprintRegistered { cell_id: CellId, graph_name: String }, - FinalAnswer { cell_id: CellId, content: String }, - SessionCompleted { session_id: SessionId }, - SessionFailed { session_id: SessionId, error: String }, -} -``` - -When the REPL calls a model, tool, agent, or graph, the child harness/graph -events should preserve: - -- root run id -- parent run id -- cell id -- node id when used inside a graph -- recursion depth -- capability name - -## Diagnostics - -Diagnostics should preserve source spans from scripts and model-generated cells. - -Required errors: - -- invalid script syntax -- unknown capability -- unknown model -- unknown tool -- unknown graph -- invalid graph source -- graph compilation failed -- generated graph review required -- graph registration denied -- invalid arguments -- unsupported value type -- operation limit exceeded -- timeout exceeded -- output limit exceeded -- call limit exceeded -- recursion limit exceeded -- unsafe backend requested -- reserved name overwrite - -Example: - -```text -error[E-ragsh-unknown-tool]: tool `lookup_usr` is not registered - --> support.ragsh:8:18 - | -8 | let user = tool_call(#{ tool: "lookup_usr", arguments: #{ id: id } }); - | ^^^^^^^^^^^^^^^^^^^^^^^^^ unknown tool - | -help: did you mean `lookup_user`? -``` - -## Testkit - -`repl::testkit` should include: - -- fake model capability -- fake tool capability -- fake graph capability -- deterministic event recorder -- script execution helper -- CodeAct loop helper with scripted model responses -- operation-limit assertion -- output-limit assertion -- recursive-call assertion -- batched-call ordering assertion -- golden trajectory fixtures - -## Implementation Milestones - -### R1: Documentation And Types - -- add this module doc -- add `repl` package shape to the spec -- define `ReplSession`, `ReplPolicy`, `ReplResult`, and `ReplEvent` -- no Rhai dependency yet - -### R2: Rhai Prototype - -- add optional `repl-rhai` feature -- embed Rhai behind `ReplBackend` -- support persistent variables -- support `show_vars`, `emit`, and `answer` -- enforce operation and output limits - -### R3: Harness Capabilities - -- add `model_query` -- add `model_query_batched` -- add fake-model tests -- forward harness events through REPL events - -### R4: Tool And Agent Capabilities - -- add `tool_call` -- add `agent_query` -- validate schemas and limits -- record usage and cost rollups - -### R5: Graph Capability - -- add `graph_run` -- add `graph_define`, `graph_validate`, `graph_compile`, `graph_diff`, and - `graph_register` -- support graph-node `kind repl_agent` -- preserve node id, parent run id, and depth in child events -- require generated-graph review gates when policy enables them - -### R6: CodeAct Loop - -- parse fenced `ragsh` blocks from assistant messages -- execute cells iteratively -- append compact execution feedback to model history -- stop on `answer(...)` -- add trajectory logging and tests - -### R7: Python Sandbox Backend +--- -- add optional out-of-process backend -- expose the same capability protocol -- run RLM-compatible Python scripts under explicit policy -- keep it disabled by default +Continues in [`operations.md`](operations.md) (RLM feature map, +CodeAct loop, examples, Rhai embedding plan, Python compatibility +backend, safety, events, diagnostics, testkit, milestones). diff --git a/docs/modules/repl-language/operations.md b/docs/modules/repl-language/operations.md new file mode 100644 index 0000000..3e4fb85 --- /dev/null +++ b/docs/modules/repl-language/operations.md @@ -0,0 +1,351 @@ +# REPL Language: RLM Feature Map, Embedding, Safety, Events, Testkit + +Continues from [`design.md`](design.md): RLM feature map, CodeAct +loop, example session/graph node, Rhai embedding plan, Python +compatibility backend, safety, events, diagnostics, testkit, and +implementation milestones. + +## RLM Feature Map + +The goal is to port the useful `rlm` behavior into TinyAgents without porting +Python's unsafe local execution model. + +| `rlm` feature | TinyAgents REPL equivalent | +| --------------------------- | ----------------------------------------------------- | +| Python `context` variable | Rhai `context` variable | +| Python persistent locals | `ReplSession::variables` | +| fenced `repl` blocks | fenced `ragsh` blocks | +| `llm_query` | `model_query` | +| `llm_query_batched` | `model_query_batched` | +| `rlm_query` | `agent_query` or `repl_query` | +| `rlm_query_batched` | `agent_query_batched` or `repl_query_batched` | +| custom Python tools | registered Rust tool capabilities | +| generated Python programs | `.ragsh` cells plus generated `.rag` graph blueprints | +| `SHOW_VARS()` | `show_vars()` | +| `answer["ready"] = True` | `answer(...)` | +| max iterations | `ReplPolicy::max_iterations` for CodeAct loops | +| max depth | graph/harness recursion policy | +| max budget | harness cost policy | +| token compaction | harness summarization feature | +| JSONL trajectory logger | typed event stream plus store backend | +| Docker/cloud REPL isolation | future `PythonSandboxRepl` backend | + +## CodeAct Loop + +A model-driven REPL agent has this lifecycle: + +1. Create `ReplSession`. +2. Load `context`, `state`, `messages`, `history`, and `run` variables. +3. Build a model request explaining the available REPL functions. +4. Invoke the model through the harness. +5. Extract fenced `ragsh` blocks from the assistant message. +6. Execute each block in the REPL session. +7. Capture stdout, changed variables, call records, events, and errors. +8. Append a compact execution result as the next user message. +9. Repeat until `answer(...)` is called or limits are reached. +10. Persist events, usage, cost, and final answer. + +This loop is a harness feature. When used inside a graph node, the graph still +owns node routing, checkpointing, interrupts, recursion depth, and failure +policy. + +If the model writes `.rag` source, the loop should treat it as a graph proposal. +The REPL may validate, diff, compile, and run that proposal only through the +expressive-language compiler and the graph registry policy. This is how an +agent can define its own graph without acquiring arbitrary topology mutation or +host-code execution privileges. + +## Example Session + +```rhai +let lines = context.split("\n"); +let candidates = []; + +for line in lines { + if line.contains("SECRET_NUMBER=") { + candidates.push(line); + } +} + +emit("candidates_found", #{ count: candidates.len() }); + +let result = model_query(#{ + model: "default", + prompt: "Return only the digits from this candidate line:\n" + candidates[0] +}); + +answer(result); +``` + +## Example Graph Node + +```tinyagents +graph support_repl { + start investigate + + node investigate { + kind repl_agent + model "default" + script "support-investigation.ragsh" + tools ["lookup_user", "create_ticket"] + routes { + final -> END + needs_review -> review + } + } + + node review { + kind interrupt + prompt "Approve escalation?" + routes { + approved -> END + rejected -> investigate + } + } +} +``` + +The `repl_agent` node is a harness-backed node template. It may execute a fixed +script, a model-driven CodeAct loop, or a combination where a fixed prologue +sets up variables before the model starts writing cells. + +## Rhai Embedding Plan + +The Rhai runtime should be isolated behind an interface so future Python or WASM +backends can reuse the same TinyAgents semantics. + +```rust +#[async_trait] +pub trait ReplBackend: Send { + async fn execute_cell( + &mut self, + session: &mut ReplSession, + source: SourceCell, + ) -> Result; +} + +pub struct RhaiReplBackend { + engine: rhai::Engine, + ast_cache: AstCache, +} +``` + +Rhai-specific requirements: + +- configure `Engine::set_max_operations` +- disable or avoid unneeded packages +- register only TinyAgents capability functions +- expose data through `Dynamic`, maps, and arrays with explicit conversion +- compile and cache ASTs for repeated scripts +- keep each session's `Scope` separate +- restore reserved names after each cell +- truncate stdout and returned values according to policy +- convert Rhai errors into structured diagnostics with spans + +Async adapter requirement: + +Rhai host functions are easiest to expose as synchronous functions. TinyAgents +model, tool, and graph calls are async. The backend should not hide blocking in +unbounded threads. Use one of these designs: + +1. command recording: Rhai functions create `ReplCommand` values, then the + Rust async runtime executes those commands after the cell +2. blocking bridge: host functions call into a bounded runtime handle with + strict timeouts +3. staged syntax: `let x = model_query(...)` is transformed before evaluation + into host-executed calls + +Recommendation for v1: use a blocking bridge only in examples and tests, but +design the public API around command recording. Command recording is easier to +make deterministic and safer under async graph execution. + +## Python Compatibility Backend + +Python should be a compatibility backend, not the default embedded runtime. + +```rust +pub struct PythonSandboxReplBackend { + sandbox: SandboxClient, +} +``` + +Potential use cases: + +- training model behavior that already expects Python +- local research workflows +- compatibility with RLM-style prompts +- data-heavy scripts where Python libraries are explicitly useful + +Requirements: + +- must run out of process +- must have no direct host filesystem access by default +- must communicate through a framed JSON protocol +- must expose the same TinyAgents capability functions +- must enforce the same `ReplPolicy` +- must emit the same `ReplEvent` stream + +This lets TinyAgents support Python-like RLM ergonomics without making Python a +trusted in-process extension language. + +## Safety + +Safety rules: + +- no arbitrary filesystem access in the default Rhai backend +- no environment variable interpolation from scripts +- no direct network access +- no process spawning +- no unregistered native functions +- bounded script size +- bounded operation count +- bounded output size +- bounded model/tool/graph calls +- bounded recursion depth +- bounded concurrency +- typed conversion at every capability boundary +- redaction before event and store writes + +The REPL is an orchestration surface, not a privilege escalation surface. + +## Events + +The REPL event stream should compose with graph and harness events. + +```rust +pub enum ReplEvent { + SessionStarted { session_id: SessionId, run_id: RunId }, + CellStarted { cell_id: CellId, source_name: String }, + CellStdout { cell_id: CellId, chunk: String }, + CellCompleted { cell_id: CellId, elapsed: Duration }, + CellFailed { cell_id: CellId, diagnostic: Diagnostic }, + VariableChanged { cell_id: CellId, name: String }, + CapabilityCallStarted { cell_id: CellId, call_id: CallId, name: String }, + CapabilityCallCompleted { cell_id: CellId, call_id: CallId }, + GraphBlueprintDefined { cell_id: CellId, graph_name: String }, + GraphBlueprintValidated { cell_id: CellId, graph_name: String }, + GraphBlueprintCompiled { cell_id: CellId, graph_name: String }, + GraphBlueprintRegistered { cell_id: CellId, graph_name: String }, + FinalAnswer { cell_id: CellId, content: String }, + SessionCompleted { session_id: SessionId }, + SessionFailed { session_id: SessionId, error: String }, +} +``` + +When the REPL calls a model, tool, agent, or graph, the child harness/graph +events should preserve: + +- root run id +- parent run id +- cell id +- node id when used inside a graph +- recursion depth +- capability name + +## Diagnostics + +Diagnostics should preserve source spans from scripts and model-generated cells. + +Required errors: + +- invalid script syntax +- unknown capability +- unknown model +- unknown tool +- unknown graph +- invalid graph source +- graph compilation failed +- generated graph review required +- graph registration denied +- invalid arguments +- unsupported value type +- operation limit exceeded +- timeout exceeded +- output limit exceeded +- call limit exceeded +- recursion limit exceeded +- unsafe backend requested +- reserved name overwrite + +Example: + +```text +error[E-ragsh-unknown-tool]: tool `lookup_usr` is not registered + --> support.ragsh:8:18 + | +8 | let user = tool_call(#{ tool: "lookup_usr", arguments: #{ id: id } }); + | ^^^^^^^^^^^^^^^^^^^^^^^^^ unknown tool + | +help: did you mean `lookup_user`? +``` + +## Testkit + +`repl::testkit` should include: + +- fake model capability +- fake tool capability +- fake graph capability +- deterministic event recorder +- script execution helper +- CodeAct loop helper with scripted model responses +- operation-limit assertion +- output-limit assertion +- recursive-call assertion +- batched-call ordering assertion +- golden trajectory fixtures + +## Implementation Milestones + +### R1: Documentation And Types + +- add this module doc +- add `repl` package shape to the spec +- define `ReplSession`, `ReplPolicy`, `ReplResult`, and `ReplEvent` +- no Rhai dependency yet + +### R2: Rhai Prototype + +- add optional `repl-rhai` feature +- embed Rhai behind `ReplBackend` +- support persistent variables +- support `show_vars`, `emit`, and `answer` +- enforce operation and output limits + +### R3: Harness Capabilities + +- add `model_query` +- add `model_query_batched` +- add fake-model tests +- forward harness events through REPL events + +### R4: Tool And Agent Capabilities + +- add `tool_call` +- add `agent_query` +- validate schemas and limits +- record usage and cost rollups + +### R5: Graph Capability + +- add `graph_run` +- add `graph_define`, `graph_validate`, `graph_compile`, `graph_diff`, and + `graph_register` +- support graph-node `kind repl_agent` +- preserve node id, parent run id, and depth in child events +- require generated-graph review gates when policy enables them + +### R6: CodeAct Loop + +- parse fenced `ragsh` blocks from assistant messages +- execute cells iteratively +- append compact execution feedback to model history +- stop on `answer(...)` +- add trajectory logging and tests + +### R7: Python Sandbox Backend + +- add optional out-of-process backend +- expose the same capability protocol +- run RLM-compatible Python scripts under explicit policy +- keep it disabled by default diff --git a/goal.md b/docs/sdk-gaps.md similarity index 98% rename from goal.md rename to docs/sdk-gaps.md index 94fea22..85287be 100644 --- a/goal.md +++ b/docs/sdk-gaps.md @@ -1,5 +1,10 @@ # TinyAgents SDK Gaps +> **Internal migration backlog.** This is a working document tracking an +> internal OpenHuman-to-TinyAgents migration effort, not a general public +> roadmap or API reference. See [`ROADMAP.md`](../ROADMAP.md) for the +> project's public-facing roadmap. + This document lists TinyAgents SDK features that are missing or only partially available from the perspective of migrating OpenHuman's Rust agent core onto TinyAgents. diff --git a/docs/spec/README.md b/docs/spec/README.md index 3f3d9b7..bc85d07 100644 --- a/docs/spec/README.md +++ b/docs/spec/README.md @@ -118,1004 +118,104 @@ it. ## Module 1: Harness -The harness is the outer runtime for LLM applications. In LangChain terms, this -is the layer around a model call that owns the agent loop, prompt/context -assembly, tool execution, middleware, memory, streaming, tracing, retries, and -testability. - -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. - -### Source Inspiration - -The harness design is informed by LangChain's docs on agents, chat models, tools, -runtime context, memory, structured output, middleware, streaming, tracing, and -testing: - -- -- -- -- -- -- -- -- -- -- - -### Responsibilities - -- Register chat model providers. -- Resolve model calls from request overrides, reusable state, model hints, - agent defaults, registry defaults, and fallback policy. -- Register tools and validate tool calls against schemas. -- Build model requests from state, prompts, memory, and runtime context. -- Apply prompt and message templates. -- Preserve provider prompt/KV-cache stability by keeping cacheable prompt - prefixes deterministic and isolating volatile context near the tail of model - requests. -- Manage per-run config such as run ids, thread ids, metadata, tags, deadlines, - max concurrency, model limits, tool limits, and - cancellation. -- Provide middleware hooks before and after model calls, tool calls, and errors. -- Provide middleware hooks during streaming model calls so compression, - redaction, observability, and adaptive context algorithms can inspect deltas - 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. -- 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. -- 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. -- Provide test doubles for models, tools, stores, clocks, and ids. - -### Core Types - -```rust -pub struct AgentHarness { - models: ModelRegistry, - tools: ToolRegistry, - middleware: MiddlewareStack, - policy: RunPolicy, -} - -pub struct RunConfig { - pub run_id: String, - pub thread_id: Option, - pub tags: Vec, - pub metadata: serde_json::Value, - pub timeout_ms: Option, - pub max_model_calls: usize, - pub max_tool_calls: usize, -} - -pub struct RunContext { - pub config: RunConfig, - pub data: Ctx, - pub stores: StoreRegistry, - pub events: EventSink, -} -``` - -`RunConfig` is stable invocation identity and policy. `RunContext` is the -per-run dependency bag. Keeping those separate prevents global state and makes -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. - -```rust -#[async_trait] -pub trait ChatModel: Send + Sync { - async fn invoke( - &self, - state: &State, - request: ModelRequest, - ) -> Result; - - async fn stream( - &self, - state: &State, - request: ModelRequest, - ) -> Result { - default_stream_from_invoke(self, state, request).await - } -} -``` - -`ModelRequest` should grow beyond the current minimal version: - -- model hints and reusable resolved-model policy -- messages -- tools available for this call -- tool choice policy -- response format -- model id/provider override -- temperature -- max tokens -- timeout -- retry policy -- local response cache policy -- provider prompt-cache policy -- cacheable prompt prefix boundaries -- ephemeral/non-cacheable context boundaries -- prompt layout fingerprint -- tags and metadata - -`ModelResponse` and agent state should record a `ResolvedModel` with registry -name, provider, provider model id, catalog snapshot/entry when known, resolver -source, and fallback history. This record is the durable answer to which model -actually ran, and it may be reused by later calls when policy allows. - -Provider prompt caching is different from local response caching. The harness -must support extreme prompt caching for providers with KV-cache or -prompt-prefix-cache behavior. That means request construction must be able to -mark stable message and tool-schema prefixes, preserve their byte/token order -across turns, and append volatile state, retrieved context, scratchpads, and -per-run metadata after those stable prefixes. Middleware that compresses, -trims, summarizes, or injects context must declare whether it changes the -cacheable prefix, the volatile tail, or only non-model-visible metadata. - -The cache contract should prevent accidental KV-cache busting: - -- stable system prompts, policy text, tool declarations, schema text, and - reusable instruction blocks should have explicit prefix segment ids -- volatile values such as timestamps, run ids, retrieved documents, current - tool results, and user-specific ephemeral context should stay out of the - cacheable prefix unless a policy explicitly opts in -- request builders should preserve segment order and canonical serialization -- middleware must emit a cache-layout event when it mutates prompt segments -- tests should be able to assert whether a change preserves or invalidates the - provider prompt-cache prefix - -Initial provider implementations should be optional feature flags: - -- `openai` -- `anthropic` -- `ollama` -- `mock` - -### Message Model - -Messages are the internal currency of the harness. The framework should not pass -raw strings after initial user input normalization. - -```rust -pub enum Message { - System(SystemMessage), - User(UserMessage), - Assistant(AssistantMessage), - Tool(ToolMessage), -} - -pub enum ContentBlock { - Text(String), - Json(serde_json::Value), - Image(ImageRef), - ProviderExtension(serde_json::Value), -} -``` - -The message model should preserve: - -- role -- content blocks -- assistant tool calls -- tool call ids -- tool result ids -- usage metadata -- provider extensions - -Tool call ids are mandatory once tool execution is implemented because they are -the correlation key between assistant requests and tool messages. - -### Tool Abstraction - -Tools are typed capabilities exposed to agents. The initial executor can accept -JSON arguments, but the registry should store schema metadata from the start. - -```rust -#[async_trait] -pub trait Tool: Send + Sync { - fn name(&self) -> &str; - fn description(&self) -> &str; - fn schema(&self) -> ToolSchema; - async fn call(&self, state: &State, call: ToolCall) -> Result; -} -``` - -Tool calls must be observable and replayable. Each call should record: - -- tool name -- arguments -- result content -- raw provider result when available -- elapsed time -- error details - -Tool names should be ASCII and `snake_case` by default. This keeps names -portable across providers that are strict about tool naming. - -### Agent Loop - -The default harness loop should be: - -1. Build `RunContext`. -2. Load short-term memory for `thread_id` when configured. -3. Build a `ModelRequest`. -4. Run pre-request middleware that can edit prompts, context, cache layout, - compression state, and provider options. -5. Run wrap middleware around the invoke or stream call for retry, fallback, - rate limiting, tracing, and replacement. -6. Run streaming middleware while model deltas arrive, including compression, - redaction, tool-call reconstruction, usage accounting, and adaptive - cancellation. -7. Run post-response middleware that can validate, compress, summarize, - persist, or transform the model response. -8. If the assistant produced tool calls, validate and execute them. -9. Append tool result messages. -10. Repeat until no tool calls remain or limits are reached. -11. Persist updated short-term memory and return the final output. - -Limits are not optional. The harness should enforce: - -- maximum model calls per run -- maximum tool calls per run -- maximum wall-clock duration -- maximum retries per call -- optional maximum concurrency for parallel tool calls - -### Middleware - -Middleware is the primary extension point for behavior that should not be baked -into the model or graph APIs. - -```rust -#[async_trait] -pub trait Middleware: Send + Sync { - async fn before_agent(&self, ctx: &mut RunContext, state: &State) -> Result<()>; - async fn after_agent(&self, ctx: &mut RunContext, state: &State, run: &mut AgentRun) -> Result<()>; - async fn before_model(&self, ctx: &mut RunContext, state: &State, request: &mut ModelRequest) -> Result<()>; - async fn on_model_delta(&self, ctx: &mut RunContext, state: &State, delta: &mut ModelDelta) -> Result<()>; - async fn after_model(&self, ctx: &mut RunContext, state: &State, response: &mut ModelResponse) -> Result<()>; - async fn before_tool(&self, ctx: &mut RunContext, state: &State, call: &mut ToolCall) -> Result<()>; - async fn on_tool_delta(&self, ctx: &mut RunContext, state: &State, delta: &mut ToolDelta) -> Result<()>; - async fn after_tool(&self, ctx: &mut RunContext, state: &State, result: &mut ToolResult) -> Result<()>; - async fn on_error(&self, ctx: &mut RunContext, error: &TinyAgentsError) -> Result<()>; -} -``` - -Wrap middleware should also exist around model calls and tool calls. A -compression algorithm often needs to wrap the entire model operation so it can -prepare context before the call, inspect streaming deltas during the call, and -commit summaries or cache metadata after the final response. - -Expected middleware: - -- retry and timeout policy -- prompt injection -- prompt cache layout protection -- provider prompt-cache/KV-cache hints -- dynamic tool filtering -- guardrails -- context compression -- transcript compression -- retrieved-context compression -- output compression -- streaming delta compression -- message trimming -- summarization -- structured output validation -- tracing -- rate limiting - -### Memory - -Memory should be a harness capability. The graph runtime should handle -checkpointed graph execution; the harness should handle conversation and -application memory. - -Memory is split into two concepts: - -- short-term memory: thread-scoped conversation state, usually backed by graph - checkpoints or a conversation checkpoint store -- long-term memory: cross-thread application data exposed through a store trait - -Memory backends should start with: - -- in-memory store for tests -- file-backed store for local development -- trait boundary for external stores - -Trimming and summarization should be explicit policies, not hidden behavior. -Compression is a broader middleware family than summarization. The harness -should support pre-call compression of old messages and retrieved context, -during-call compression or redaction of streaming deltas, and post-call -compression of transcripts, tool artifacts, reasoning traces, and memory -records. Compression middleware must preserve provenance: the original source -ids, token estimates, cache segment ids, and enough metadata to explain why a -message was removed, replaced, or summarized. - -### Structured Output - -The harness should support typed output using two strategies: - -- provider-native schema enforcement when the model supports it -- tool-call-based structured output fallback - -The user-facing API should allow: - -```rust -let output: MyType = harness - .with_response_format(ResponseFormat::json_schema::()) - .invoke(state) - .await? - .structured_response()?; -``` - -The final structured value should be separate from final chat messages so users -can inspect both. - -### Observability - -Every run should be traceable through typed events and readable through a -compact execution status store. The status store is the answer to "what is this -run doing now?"; the event stream and journal are the answer to "what happened?" - -The canonical feature references are: - -- [Harness observability and events](../modules/harness/observability.md) -- [Harness store](../modules/harness/store.md) -- [Harness streaming](../modules/harness/streaming.md) -- [Harness cache](../modules/harness/cache.md) - -At minimum, the harness should emit: - -- run started -- model requested -- model token delta -- model responded -- tool requested -- tool token or progress delta -- tool responded -- state update -- middleware started -- middleware completed -- retry scheduled -- route selected -- run completed -- run failed - -The event stream should be structured data so it can feed logs, -OpenTelemetry, test recorders, durable JSONL/MongoDB journals, or a custom UI. - -```rust -pub enum AgentEvent { - RunStarted { run_id: String, thread_id: Option }, - ModelStarted { call_id: String, model: String }, - ModelDelta { call_id: String, delta: MessageDelta }, - ModelCompleted { call_id: String, usage: Option }, - ToolStarted { call_id: String, tool_name: String }, - ToolCompleted { call_id: String, tool_name: String }, - RetryScheduled { call_id: String, attempt: usize }, - RunCompleted { run_id: String }, - RunFailed { run_id: String, error: String }, -} -``` - -The harness should also expose a compact run-status record: - -```rust -pub struct HarnessRunStatus { - pub run_id: RunId, - pub parent_run_id: Option, - pub root_run_id: RunId, - pub thread_id: Option, - pub component: ComponentId, - pub status: ExecutionStatus, - pub current_phase: HarnessPhase, - pub model_calls: usize, - pub tool_calls: usize, - pub active_model_call: Option, - pub active_tool_calls: Vec, - pub last_event_id: Option, - pub usage: UsageTotals, - pub cost: CostTotals, - pub started_at: SystemTime, - pub updated_at: SystemTime, - pub ended_at: Option, - pub error: Option, -} -``` - -Status records are operational snapshots. They should not include full prompts, -tool outputs, or raw provider payloads. Event journals are append-only and -should support listener replay by stream offset. Derived observability -projections such as latest status, usage rollups, cost rollups, and timing -summaries may be cached, but every cached projection must include a source event -offset and projection version. - -### Testability - -The harness should ship a `testkit` module early. It should include: - -- fake chat model with scripted responses -- fake streaming model -- fake tool -- in-memory stores -- deterministic run id generator -- deterministic clock -- event recorder -- trajectory assertions that check tool calls and state changes without relying - on exact LLM prose +The harness is the provider-neutral runtime for model calls, tools, +middleware, structured output, streaming, usage/cost, retry/limits, cache, +memory/embeddings, sub-agents, and steering. See +[`harness-spec.md`](harness-spec.md) for the full specification (core types, +model/tool/message abstractions, agent loop, middleware, memory, structured +output, observability, and testability), and +[`docs/modules/harness/README.md`](../modules/harness/README.md) for the +per-topic implementation docs. ## Module 2: Graph -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: - -- -- -- -- -- -- -- -- - -### 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` is an async unit of work. - -`NodeOutput` 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), - Command(Command), - Interrupt(Interrupt), -} - -pub struct Command { - pub update: Option, - pub goto: Vec, - pub resume: Option, -} -``` - -`GraphBuilder` 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: Send + Sync { - fn reduce(&self, current: T, update: T) -> Result; -} - -pub trait StateReducer: Send + Sync { - fn apply(&self, state: State, update: Update) -> Result; -} -``` - -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: Send + Sync { - async fn put(&self, checkpoint: Checkpoint) -> Result; - async fn get(&self, thread_id: &str, checkpoint_id: Option<&str>) -> Result>>; - async fn list(&self, thread_id: &str) -> Result>; -} -``` - -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, - pub thread_id: Option, - pub graph_id: GraphId, - pub checkpoint_id: Option, - pub checkpoint_namespace: Vec, - pub status: ExecutionStatus, - pub current_step: usize, - pub active_nodes: Vec, - pub pending_interrupts: Vec, - pub last_event_id: Option, - pub started_at: SystemTime, - pub updated_at: SystemTime, - pub ended_at: Option, - pub error: Option, -} -``` - -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 +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. ## Module 3: Expressive Language -The expressive language is a compact way to define agent workflows without -writing all builder calls manually. It should compile into the same graph and -harness types as Rust code. - -This language is not meant to replace Rust. It is a workflow definition layer for -fast iteration, examples, documentation, and eventually user-authored agent -plans. - -It is also the safe boundary for agent-authored graph plans. A REPL or model may -propose `.rag` source, but that source must pass through the same parser, -diagnostics, registry binding, allowlist checks, review gates, and graph -compiler as human-authored source before it can run. - -### Goals - -- Make common agent graphs readable at a glance. -- Keep syntax close to graph intent. -- Compile into explicit TinyAgents structures. -- Preserve source locations for helpful errors. -- Avoid embedding arbitrary code in the first version. -- Describe state channels, reducers, policies, subgraphs, sub-agents, - interrupts, joins, and fanout as declarative graph primitives. -- Produce inspectable blueprints that can be reviewed, diffed, registered, and - tested. - -### Non-Goals - -- It is not a general-purpose programming language. -- It is not a prompt templating language by itself. -- It should not execute untrusted code. -- It should not bypass Rust type checks for stateful logic. -- It should not install model-generated topology directly into the graph - runtime. - -### Initial Syntax Sketch - -```tinyagents -graph support_agent { - defaults { - recursion_limit 50 - checkpoint inherit - } - - start agent - - channel messages messages - channel tool_calls append - - node agent { - kind agent - model "default" - prompt "You are a concise support agent." - tools ["lookup_user", "create_ticket"] - routes { - tool_call -> tools - final -> END - } - } - - node tools { - kind tool_executor - next agent - } -} -``` - -### Minimal Grammar - -```text -program = graph_decl* -graph_decl = "graph" ident "{" graph_item* "}" -graph_item = start_decl | defaults_decl | channel_decl | node_decl | edge_decl -start_decl = "start" ident -defaults_decl = "defaults" object -channel_decl = "channel" ident reducer_ref -node_decl = "node" ident "{" node_item* "}" -node_item = kind_decl | model_decl | prompt_decl | tools_decl | next_decl | routes_decl -kind_decl = "kind" ident -model_decl = "model" string -prompt_decl = "prompt" string -tools_decl = "tools" "[" string_list? "]" -next_decl = "next" ident -routes_decl = "routes" "{" route_decl* "}" -route_decl = ident "->" (ident | "END") -edge_decl = ident "->" ident -``` - -The full language target is broader than this minimal grammar. It should grow -toward commands, `Send` fanout, joins/barriers, subgraphs, sub-agents, -`repl_agent` nodes, interrupts, registered route functions, graph defaults, -capability allowlists, blueprint provenance, and deterministic graph diffs. See -[the expressive language module](../modules/expressive-language/README.md) for -the canonical target. - -### Compilation Pipeline - -1. Parse source into an AST. -2. Validate identifiers and route targets. -3. Lower AST into graph builder calls. -4. Bind model and tool references through the harness. -5. Return a compiled workflow object. - -### Error Requirements - -Errors should include: - -- file name when available -- line and column -- invalid token or missing token -- unknown node name -- duplicate node name -- missing start node -- route target that does not exist -- model or tool reference that is not registered in the harness - -### Runtime Relationship - -The expressive language should produce the same runtime structures as hand-written -Rust: - -```text -source -> parser -> AST -> compiler -> StateGraph + Harness bindings -``` - -The graph runtime should not know whether a graph came from Rust builders or the -expressive language. - -For generated source, the runtime relationship is: - -```text -REPL/model proposal -> .rag source or AST -> parser -> diagnostics -> resolver - -> policy/review gate -> compiler -> GraphBuilder + Harness bindings - -> CompiledGraph -> optional registry registration -``` +The `.rag` expressive language is a declarative, side-effect-free blueprint +format that compiles through lexer -> parser -> compiler into the same +graph/harness runtime types as hand-written Rust. See +[`expressive-language-spec.md`](expressive-language-spec.md) for the goals, +grammar sketch, and compilation pipeline, and +[`docs/modules/expressive-language/README.md`](../modules/expressive-language/README.md) +for implementation status. ## Package Layout -Target module layout: +The crate is a single library at the repository root (`Cargo.toml`), with +`src/lib.rs` re-exporting the public surface and `src/error.rs` holding the +crate-wide error type. Each of the five surfaces lives in its own module +directory: ```text src/ - chat.rs error.rs - graph.rs - harness.rs - language/ - ast.rs - lexer.rs - parser.rs - compiler.rs - mod.rs - model.rs - tool.rs + lib.rs + graph/ # durable typed state graphs (checkpoint, interrupt, streaming, ...) + harness/ # provider-neutral model calls, tools, middleware, streaming, ... + language/ # the declarative `.rag` blueprint format (lexer/parser/compiler) + registry/ # the named capability catalog (models, tools, agents, stores, ...) + repl/ # the imperative `.ragsh` session runtime ``` -Provider implementations should live behind feature flags: - -```text -src/providers/ - openai.rs - anthropic.rs - ollama.rs - mock.rs -``` +Provider implementations (OpenAI and the OpenAI-compatible endpoints for +Anthropic, Ollama, DeepSeek, Groq, xAI, OpenRouter, Together, and Mistral) +live inside `src/harness/providers/` and are compiled in unconditionally. +Two Cargo features gate optional dependencies: `sqlite` (embedded SQLite +checkpointer) and `repl` (embedded Rhai engine for `.ragsh` sessions). ## Milestones -### Milestone 1: Core Runtime +All five milestones below have shipped as of v1.5.0. + +### Milestone 1: Core Runtime (shipped) -- Chat message primitives. -- Model trait. -- Tool trait. -- State graph with direct and conditional edges. -- Basic tests and examples. +Chat message primitives, the model and tool traits, the state graph with +direct and conditional edges, and the initial test/example suite. -### Milestone 2: Harness +### Milestone 2: Harness (shipped) -- Harness type. -- Model registry. -- Tool registry. -- Run context. -- Callback events. -- Run status store. -- Durable event journal. -- Cache-backed observability projections. -- Mock model and mock tool utilities. +The `AgentHarness` type, model and tool registries, run context, callback +events, run status store, durable event journal, cache-backed observability +projections, and mock model/tool testkit utilities. -### Milestone 3: Expressive Language Preview +### Milestone 3: Expressive Language (shipped) -- AST. -- Parser for a small graph definition language. -- Compiler into `StateGraph`. -- Helpful parse and validation errors. -- Example `.rag` or `.tinyagents` workflow file. +The `.rag` AST, lexer, parser, compiler into the graph runtime, parse/ +validation diagnostics with source spans, and example `.rag` workflow files +(see `examples/rag_blueprint.rs`, `examples/openai_self_blueprint.rs`). -### Milestone 4: Provider Integrations +### Milestone 4: Provider Integrations (shipped) -- OpenAI chat model provider. -- Anthropic chat model provider. -- Local/mock provider. -- Provider feature flags. +OpenAI and OpenAI-compatible provider adapters (Anthropic, Ollama, DeepSeek, +Groq, xAI, OpenRouter, Together, Mistral), plus the offline deterministic +mock provider. -### Milestone 5: Production Runtime Features +### Milestone 5: Production Runtime Features (shipped) -- Streaming events. -- Checkpointing. -- Resume support. -- Graph run status store. -- Graph event journal and listener replay. -- Graph export. -- Tracing integration. +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`). ## Open Questions -- Should the expressive language file extension be `.rag`, `.tinyagents`, or - something shorter? -- Should state schemas be declared in the language, or should state remain purely - Rust-owned? -- Should graph nodes support typed route enums before serialization support? -- Should provider crates live in this crate behind feature flags or in separate - crates? -- Should memory be synchronous, async, or both? +Historical decisions that have since been settled, kept for context: + +- The expressive language file extension is `.rag` (interactive/imperative + orchestration uses the separate `.ragsh` extension). +- State schemas remain Rust-owned; `.rag` binds to them by name through the + registry rather than declaring schemas itself. +- Provider crates live in this crate as always-compiled modules behind + `src/harness/providers/`, not separate crates or feature flags. +- 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? diff --git a/docs/spec/expressive-language-spec.md b/docs/spec/expressive-language-spec.md new file mode 100644 index 0000000..3993c41 --- /dev/null +++ b/docs/spec/expressive-language-spec.md @@ -0,0 +1,137 @@ +# Expressive Language Module Specification + +The expressive language is a compact way to define agent workflows without +writing all builder calls manually. It should compile into the same graph and +harness types as Rust code. + +This language is not meant to replace Rust. It is a workflow definition layer for +fast iteration, examples, documentation, and eventually user-authored agent +plans. + +It is also the safe boundary for agent-authored graph plans. A REPL or model may +propose `.rag` source, but that source must pass through the same parser, +diagnostics, registry binding, allowlist checks, review gates, and graph +compiler as human-authored source before it can run. + +### Goals + +- Make common agent graphs readable at a glance. +- Keep syntax close to graph intent. +- Compile into explicit TinyAgents structures. +- Preserve source locations for helpful errors. +- Avoid embedding arbitrary code in the first version. +- Describe state channels, reducers, policies, subgraphs, sub-agents, + interrupts, joins, and fanout as declarative graph primitives. +- Produce inspectable blueprints that can be reviewed, diffed, registered, and + tested. + +### Non-Goals + +- It is not a general-purpose programming language. +- It is not a prompt templating language by itself. +- It should not execute untrusted code. +- It should not bypass Rust type checks for stateful logic. +- It should not install model-generated topology directly into the graph + runtime. + +### Initial Syntax Sketch + +```tinyagents +graph support_agent { + defaults { + recursion_limit 50 + checkpoint inherit + } + + start agent + + channel messages messages + channel tool_calls append + + node agent { + kind agent + model "default" + prompt "You are a concise support agent." + tools ["lookup_user", "create_ticket"] + routes { + tool_call -> tools + final -> END + } + } + + node tools { + kind tool_executor + next agent + } +} +``` + +### Minimal Grammar + +```text +program = graph_decl* +graph_decl = "graph" ident "{" graph_item* "}" +graph_item = start_decl | defaults_decl | channel_decl | node_decl | edge_decl +start_decl = "start" ident +defaults_decl = "defaults" object +channel_decl = "channel" ident reducer_ref +node_decl = "node" ident "{" node_item* "}" +node_item = kind_decl | model_decl | prompt_decl | tools_decl | next_decl | routes_decl +kind_decl = "kind" ident +model_decl = "model" string +prompt_decl = "prompt" string +tools_decl = "tools" "[" string_list? "]" +next_decl = "next" ident +routes_decl = "routes" "{" route_decl* "}" +route_decl = ident "->" (ident | "END") +edge_decl = ident "->" ident +``` + +The full language target is broader than this minimal grammar. It should grow +toward commands, `Send` fanout, joins/barriers, subgraphs, sub-agents, +`repl_agent` nodes, interrupts, registered route functions, graph defaults, +capability allowlists, blueprint provenance, and deterministic graph diffs. See +[the expressive language module](../modules/expressive-language/README.md) for +the canonical target. + +### Compilation Pipeline + +1. Parse source into an AST. +2. Validate identifiers and route targets. +3. Lower AST into graph builder calls. +4. Bind model and tool references through the harness. +5. Return a compiled workflow object. + +### Error Requirements + +Errors should include: + +- file name when available +- line and column +- invalid token or missing token +- unknown node name +- duplicate node name +- missing start node +- route target that does not exist +- model or tool reference that is not registered in the harness + +### Runtime Relationship + +The expressive language should produce the same runtime structures as hand-written +Rust: + +```text +source -> parser -> AST -> compiler -> StateGraph + Harness bindings +``` + +The graph runtime should not know whether a graph came from Rust builders or the +expressive language. + +For generated source, the runtime relationship is: + +```text +REPL/model proposal -> .rag source or AST -> parser -> diagnostics -> resolver + -> policy/review gate -> compiler -> GraphBuilder + Harness bindings + -> CompiledGraph -> optional registry registration +``` + diff --git a/docs/spec/graph-spec.md b/docs/spec/graph-spec.md new file mode 100644 index 0000000..ce671b3 --- /dev/null +++ b/docs/spec/graph-spec.md @@ -0,0 +1,338 @@ +# 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: + +- +- +- +- +- +- +- +- + +### 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` is an async unit of work. + +`NodeOutput` 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), + Command(Command), + Interrupt(Interrupt), +} + +pub struct Command { + pub update: Option, + pub goto: Vec, + pub resume: Option, +} +``` + +`GraphBuilder` 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: Send + Sync { + fn reduce(&self, current: T, update: T) -> Result; +} + +pub trait StateReducer: Send + Sync { + fn apply(&self, state: State, update: Update) -> Result; +} +``` + +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: Send + Sync { + async fn put(&self, checkpoint: Checkpoint) -> Result; + async fn get(&self, thread_id: &str, checkpoint_id: Option<&str>) -> Result>>; + async fn list(&self, thread_id: &str) -> Result>; +} +``` + +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, + pub thread_id: Option, + pub graph_id: GraphId, + pub checkpoint_id: Option, + pub checkpoint_namespace: Vec, + pub status: ExecutionStatus, + pub current_step: usize, + pub active_nodes: Vec, + pub pending_interrupts: Vec, + pub last_event_id: Option, + pub started_at: SystemTime, + pub updated_at: SystemTime, + pub ended_at: Option, + pub error: Option, +} +``` + +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 new file mode 100644 index 0000000..c7e3cb8 --- /dev/null +++ b/docs/spec/harness-spec.md @@ -0,0 +1,441 @@ +# Harness Module Specification + +The harness is the outer runtime for LLM applications. In LangChain terms, this +is the layer around a model call that owns the agent loop, prompt/context +assembly, tool execution, middleware, memory, streaming, tracing, retries, and +testability. + +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. + +### Source Inspiration + +The harness design is informed by LangChain's docs on agents, chat models, tools, +runtime context, memory, structured output, middleware, streaming, tracing, and +testing: + +- +- +- +- +- +- +- +- +- +- + +### Responsibilities + +- Register chat model providers. +- Resolve model calls from request overrides, reusable state, model hints, + agent defaults, registry defaults, and fallback policy. +- Register tools and validate tool calls against schemas. +- Build model requests from state, prompts, memory, and runtime context. +- Apply prompt and message templates. +- Preserve provider prompt/KV-cache stability by keeping cacheable prompt + prefixes deterministic and isolating volatile context near the tail of model + requests. +- Manage per-run config such as run ids, thread ids, metadata, tags, deadlines, + max concurrency, model limits, tool limits, and + cancellation. +- Provide middleware hooks before and after model calls, tool calls, and errors. +- Provide middleware hooks during streaming model calls so compression, + redaction, observability, and adaptive context algorithms can inspect deltas + 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. +- 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. +- 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. +- Provide test doubles for models, tools, stores, clocks, and ids. + +### Core Types + +```rust +pub struct AgentHarness { + models: ModelRegistry, + tools: ToolRegistry, + middleware: MiddlewareStack, + policy: RunPolicy, +} + +pub struct RunConfig { + pub run_id: String, + pub thread_id: Option, + pub tags: Vec, + pub metadata: serde_json::Value, + pub timeout_ms: Option, + pub max_model_calls: usize, + pub max_tool_calls: usize, +} + +pub struct RunContext { + pub config: RunConfig, + pub data: Ctx, + pub stores: StoreRegistry, + pub events: EventSink, +} +``` + +`RunConfig` is stable invocation identity and policy. `RunContext` is the +per-run dependency bag. Keeping those separate prevents global state and makes +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. + +```rust +#[async_trait] +pub trait ChatModel: Send + Sync { + async fn invoke( + &self, + state: &State, + request: ModelRequest, + ) -> Result; + + async fn stream( + &self, + state: &State, + request: ModelRequest, + ) -> Result { + default_stream_from_invoke(self, state, request).await + } +} +``` + +`ModelRequest` should grow beyond the current minimal version: + +- model hints and reusable resolved-model policy +- messages +- tools available for this call +- tool choice policy +- response format +- model id/provider override +- temperature +- max tokens +- timeout +- retry policy +- local response cache policy +- provider prompt-cache policy +- cacheable prompt prefix boundaries +- ephemeral/non-cacheable context boundaries +- prompt layout fingerprint +- tags and metadata + +`ModelResponse` and agent state should record a `ResolvedModel` with registry +name, provider, provider model id, catalog snapshot/entry when known, resolver +source, and fallback history. This record is the durable answer to which model +actually ran, and it may be reused by later calls when policy allows. + +Provider prompt caching is different from local response caching. The harness +must support extreme prompt caching for providers with KV-cache or +prompt-prefix-cache behavior. That means request construction must be able to +mark stable message and tool-schema prefixes, preserve their byte/token order +across turns, and append volatile state, retrieved context, scratchpads, and +per-run metadata after those stable prefixes. Middleware that compresses, +trims, summarizes, or injects context must declare whether it changes the +cacheable prefix, the volatile tail, or only non-model-visible metadata. + +The cache contract should prevent accidental KV-cache busting: + +- stable system prompts, policy text, tool declarations, schema text, and + reusable instruction blocks should have explicit prefix segment ids +- volatile values such as timestamps, run ids, retrieved documents, current + tool results, and user-specific ephemeral context should stay out of the + cacheable prefix unless a policy explicitly opts in +- request builders should preserve segment order and canonical serialization +- middleware must emit a cache-layout event when it mutates prompt segments +- tests should be able to assert whether a change preserves or invalidates the + provider prompt-cache prefix + +Initial provider implementations should be optional feature flags: + +- `openai` +- `anthropic` +- `ollama` +- `mock` + +### Message Model + +Messages are the internal currency of the harness. The framework should not pass +raw strings after initial user input normalization. + +```rust +pub enum Message { + System(SystemMessage), + User(UserMessage), + Assistant(AssistantMessage), + Tool(ToolMessage), +} + +pub enum ContentBlock { + Text(String), + Json(serde_json::Value), + Image(ImageRef), + ProviderExtension(serde_json::Value), +} +``` + +The message model should preserve: + +- role +- content blocks +- assistant tool calls +- tool call ids +- tool result ids +- usage metadata +- provider extensions + +Tool call ids are mandatory once tool execution is implemented because they are +the correlation key between assistant requests and tool messages. + +### Tool Abstraction + +Tools are typed capabilities exposed to agents. The initial executor can accept +JSON arguments, but the registry should store schema metadata from the start. + +```rust +#[async_trait] +pub trait Tool: Send + Sync { + fn name(&self) -> &str; + fn description(&self) -> &str; + fn schema(&self) -> ToolSchema; + async fn call(&self, state: &State, call: ToolCall) -> Result; +} +``` + +Tool calls must be observable and replayable. Each call should record: + +- tool name +- arguments +- result content +- raw provider result when available +- elapsed time +- error details + +Tool names should be ASCII and `snake_case` by default. This keeps names +portable across providers that are strict about tool naming. + +### Agent Loop + +The default harness loop should be: + +1. Build `RunContext`. +2. Load short-term memory for `thread_id` when configured. +3. Build a `ModelRequest`. +4. Run pre-request middleware that can edit prompts, context, cache layout, + compression state, and provider options. +5. Run wrap middleware around the invoke or stream call for retry, fallback, + rate limiting, tracing, and replacement. +6. Run streaming middleware while model deltas arrive, including compression, + redaction, tool-call reconstruction, usage accounting, and adaptive + cancellation. +7. Run post-response middleware that can validate, compress, summarize, + persist, or transform the model response. +8. If the assistant produced tool calls, validate and execute them. +9. Append tool result messages. +10. Repeat until no tool calls remain or limits are reached. +11. Persist updated short-term memory and return the final output. + +Limits are not optional. The harness should enforce: + +- maximum model calls per run +- maximum tool calls per run +- maximum wall-clock duration +- maximum retries per call +- optional maximum concurrency for parallel tool calls + +### Middleware + +Middleware is the primary extension point for behavior that should not be baked +into the model or graph APIs. + +```rust +#[async_trait] +pub trait Middleware: Send + Sync { + async fn before_agent(&self, ctx: &mut RunContext, state: &State) -> Result<()>; + async fn after_agent(&self, ctx: &mut RunContext, state: &State, run: &mut AgentRun) -> Result<()>; + async fn before_model(&self, ctx: &mut RunContext, state: &State, request: &mut ModelRequest) -> Result<()>; + async fn on_model_delta(&self, ctx: &mut RunContext, state: &State, delta: &mut ModelDelta) -> Result<()>; + async fn after_model(&self, ctx: &mut RunContext, state: &State, response: &mut ModelResponse) -> Result<()>; + async fn before_tool(&self, ctx: &mut RunContext, state: &State, call: &mut ToolCall) -> Result<()>; + async fn on_tool_delta(&self, ctx: &mut RunContext, state: &State, delta: &mut ToolDelta) -> Result<()>; + async fn after_tool(&self, ctx: &mut RunContext, state: &State, result: &mut ToolResult) -> Result<()>; + async fn on_error(&self, ctx: &mut RunContext, error: &TinyAgentsError) -> Result<()>; +} +``` + +Wrap middleware should also exist around model calls and tool calls. A +compression algorithm often needs to wrap the entire model operation so it can +prepare context before the call, inspect streaming deltas during the call, and +commit summaries or cache metadata after the final response. + +Expected middleware: + +- retry and timeout policy +- prompt injection +- prompt cache layout protection +- provider prompt-cache/KV-cache hints +- dynamic tool filtering +- guardrails +- context compression +- transcript compression +- retrieved-context compression +- output compression +- streaming delta compression +- message trimming +- summarization +- structured output validation +- tracing +- rate limiting + +### Memory + +Memory should be a harness capability. The graph runtime should handle +checkpointed graph execution; the harness should handle conversation and +application memory. + +Memory is split into two concepts: + +- short-term memory: thread-scoped conversation state, usually backed by graph + checkpoints or a conversation checkpoint store +- long-term memory: cross-thread application data exposed through a store trait + +Memory backends should start with: + +- in-memory store for tests +- file-backed store for local development +- trait boundary for external stores + +Trimming and summarization should be explicit policies, not hidden behavior. +Compression is a broader middleware family than summarization. The harness +should support pre-call compression of old messages and retrieved context, +during-call compression or redaction of streaming deltas, and post-call +compression of transcripts, tool artifacts, reasoning traces, and memory +records. Compression middleware must preserve provenance: the original source +ids, token estimates, cache segment ids, and enough metadata to explain why a +message was removed, replaced, or summarized. + +### Structured Output + +The harness should support typed output using two strategies: + +- provider-native schema enforcement when the model supports it +- tool-call-based structured output fallback + +The user-facing API should allow: + +```rust +let output: MyType = harness + .with_response_format(ResponseFormat::json_schema::()) + .invoke(state) + .await? + .structured_response()?; +``` + +The final structured value should be separate from final chat messages so users +can inspect both. + +### Observability + +Every run should be traceable through typed events and readable through a +compact execution status store. The status store is the answer to "what is this +run doing now?"; the event stream and journal are the answer to "what happened?" + +The canonical feature references are: + +- [Harness observability and events](../modules/harness/observability.md) +- [Harness store](../modules/harness/store.md) +- [Harness streaming](../modules/harness/streaming.md) +- [Harness cache](../modules/harness/cache.md) + +At minimum, the harness should emit: + +- run started +- model requested +- model token delta +- model responded +- tool requested +- tool token or progress delta +- tool responded +- state update +- middleware started +- middleware completed +- retry scheduled +- route selected +- run completed +- run failed + +The event stream should be structured data so it can feed logs, +OpenTelemetry, test recorders, durable JSONL/MongoDB journals, or a custom UI. + +```rust +pub enum AgentEvent { + RunStarted { run_id: String, thread_id: Option }, + ModelStarted { call_id: String, model: String }, + ModelDelta { call_id: String, delta: MessageDelta }, + ModelCompleted { call_id: String, usage: Option }, + ToolStarted { call_id: String, tool_name: String }, + ToolCompleted { call_id: String, tool_name: String }, + RetryScheduled { call_id: String, attempt: usize }, + RunCompleted { run_id: String }, + RunFailed { run_id: String, error: String }, +} +``` + +The harness should also expose a compact run-status record: + +```rust +pub struct HarnessRunStatus { + pub run_id: RunId, + pub parent_run_id: Option, + pub root_run_id: RunId, + pub thread_id: Option, + pub component: ComponentId, + pub status: ExecutionStatus, + pub current_phase: HarnessPhase, + pub model_calls: usize, + pub tool_calls: usize, + pub active_model_call: Option, + pub active_tool_calls: Vec, + pub last_event_id: Option, + pub usage: UsageTotals, + pub cost: CostTotals, + pub started_at: SystemTime, + pub updated_at: SystemTime, + pub ended_at: Option, + pub error: Option, +} +``` + +Status records are operational snapshots. They should not include full prompts, +tool outputs, or raw provider payloads. Event journals are append-only and +should support listener replay by stream offset. Derived observability +projections such as latest status, usage rollups, cost rollups, and timing +summaries may be cached, but every cached projection must include a source event +offset and projection version. + +### Testability + +The harness should ship a `testkit` module early. It should include: + +- fake chat model with scripted responses +- fake streaming model +- fake tool +- in-memory stores +- deterministic run id generator +- deterministic clock +- event recorder +- trajectory assertions that check tool calls and state changes without relying + on exact LLM prose + diff --git a/examples/openai_self_blueprint.rs b/examples/openai_self_blueprint.rs index 001b75e..3393d9c 100644 --- a/examples/openai_self_blueprint.rs +++ b/examples/openai_self_blueprint.rs @@ -29,9 +29,8 @@ use tinyagents::graph::{CompiledGraph, END, NodeFuture}; use tinyagents::harness::message::Message; use tinyagents::harness::providers::openai::OpenAiModel; use tinyagents::harness::runtime::AgentHarness; -use tinyagents::language::compiler::{ - BoxedNode, CapabilityResolver, NodeFactory, bind_capabilities, build_graph, compile, -}; +use tinyagents::language::capability_resolver::{CapabilityResolver, bind_capabilities}; +use tinyagents::language::compiler::{BoxedNode, NodeFactory, build_graph, compile}; use tinyagents::language::parser::parse_str; use tinyagents::language::types::{NodeSpec, Routing}; use tinyagents::{Command, NodeContext, NodeResult}; diff --git a/examples/rag_blueprint.rs b/examples/rag_blueprint.rs index 707104f..bbf0ad9 100644 --- a/examples/rag_blueprint.rs +++ b/examples/rag_blueprint.rs @@ -12,7 +12,8 @@ //! ``` use tinyagents::Result; -use tinyagents::language::compiler::{CapabilityResolver, bind_capabilities, compile}; +use tinyagents::language::capability_resolver::{CapabilityResolver, bind_capabilities}; +use tinyagents::language::compiler::compile; use tinyagents::language::parser::parse_str; use tinyagents::language::types::Routing; diff --git a/src/error.rs b/src/error.rs index a3f5741..0c3e2ed 100644 --- a/src/error.rs +++ b/src/error.rs @@ -72,9 +72,30 @@ pub enum TinyAgentsError { /// A model provider call failed (transport error, non-2xx status, or a /// malformed response). The payload is a human-readable, provider-normalized /// description. + /// + /// Prefer [`TinyAgentsError::Provider`] when the structured failure detail + /// (HTTP status, provider error code, retryability) is available — this + /// variant remains for transport-level and parsing failures that have no + /// such structure to preserve. #[error("model error: {0}")] Model(String), + /// A model provider call failed with the full structured detail preserved + /// — HTTP status, provider error code, and whether retrying the same + /// request may succeed — instead of flattened into a display string. + /// + /// Real provider adapters (for example the OpenAI unary and streaming + /// paths) raise this instead of [`TinyAgentsError::Model`] whenever they + /// have a [`crate::harness::model::ProviderError`] in hand, so + /// [`crate::harness::retry::is_retryable`] can classify retryability from + /// [`crate::harness::model::ProviderError::retryable`] (a 429 is + /// retryable; a 401 is not) rather than retrying every provider failure + /// indiscriminately. Boxed so this one variant's larger payload does not + /// inflate every `Result` in the crate + /// (`clippy::result_large_err`). + #[error("model error: {0}")] + Provider(Box), + /// A tool invocation returned an error. The payload describes the failure. #[error("tool error: {0}")] Tool(String), diff --git a/src/graph/README.md b/src/graph/README.md new file mode 100644 index 0000000..68a3e4a --- /dev/null +++ b/src/graph/README.md @@ -0,0 +1,83 @@ +# graph + +TinyAgents' durable workflow runtime (LangGraph-style), and one of the +load-bearing surfaces of the crate's recursive language-model (RLM) +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 authored from a `.rag` +blueprint or driven from the `.ragsh` REPL lowers into exactly these same +types, so a model can describe, compile, and re-enter the very runtime it is +executing inside. + +Each submodule keeps type definitions in `types.rs`, behavior in `mod.rs`, and +unit tests in `test.rs` (per repo convention); 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`. | +| `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. | +| `compiled` | The superstep executor, `CompiledGraph` — sequential/parallel steps, node retry, resumable failure, run/resume/state APIs. See [`compiled/README.md`](compiled/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` | Managed child-work controls (`spawn`/`await`/`cancel`/... ) exposed as harness tools, backed by a `TaskStore`. See [`orchestration/README.md`](orchestration/README.md). | +| `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. | +| `reducer` | `StateReducer`/`Reducer` implementations (overwrite, append, min/max, set-union, closures) that fold branch updates into committed state at a superstep boundary. | +| `status` | `GraphRunStatus` — a compact run-status snapshot. | +| `stream` | `GraphEvent`, `GraphEventSink`, and streaming modes — the live, in-process event surface `observability` makes durable. | +| `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 kanban task board and harness tools. 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/src/graph/builder/mod.rs b/src/graph/builder/mod.rs index 031c5eb..d148fae 100644 --- a/src/graph/builder/mod.rs +++ b/src/graph/builder/mod.rs @@ -48,7 +48,7 @@ where /// recursion limit of 50. A reducer must be set before [`Self::compile`]. pub fn new() -> Self { Self { - graph_id: GraphId::new(format!("graph-{}", crate::graph::compiled::next_seq())), + graph_id: GraphId::new(format!("graph-{}", crate::harness::ids::next_seq())), name: None, nodes: HashMap::new(), edges: HashMap::new(), diff --git a/src/graph/channel/README.md b/src/graph/channel/README.md new file mode 100644 index 0000000..8260090 --- /dev/null +++ b/src/graph/channel/README.md @@ -0,0 +1,81 @@ +# 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. | +| `test.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/src/graph/channel/mod.rs b/src/graph/channel/mod.rs index d79635b..f92847a 100644 --- a/src/graph/channel/mod.rs +++ b/src/graph/channel/mod.rs @@ -43,7 +43,7 @@ pub use types::{ LastValue, Messages, NamedBarrier, Topic, Untracked, }; -use std::collections::{BTreeMap, HashSet}; +use std::collections::{BTreeMap, HashMap, HashSet}; use std::sync::Arc; use serde_json::Value; @@ -58,7 +58,7 @@ impl Channel for LastValue { "last_value" } - fn merge(&self, _current: Option<&Value>, incoming: Value) -> Result { + fn merge(&self, _current: Option, incoming: Value) -> Result { Ok(incoming) } @@ -72,10 +72,11 @@ impl Channel for Topic { "topic" } - fn merge(&self, current: Option<&Value>, incoming: Value) -> Result { + 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.clone(), - Some(other) => vec![other.clone()], + Some(Value::Array(items)) => items, + Some(other) => vec![other], None => Vec::new(), }; match incoming { @@ -99,7 +100,7 @@ impl Channel for Delta { "delta" } - fn merge(&self, current: Option<&Value>, incoming: Value) -> Result { + 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)?; @@ -130,9 +131,10 @@ impl Channel for Messages { "messages" } - fn merge(&self, current: Option<&Value>, incoming: Value) -> Result { + 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.clone(), + Some(Value::Array(items)) => items, Some(_) => { return Err(TinyAgentsError::Graph( "Messages channel value must be a JSON array".to_string(), @@ -144,13 +146,32 @@ impl Channel for Messages { 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 { - let id = msg.get("id").and_then(Value::as_str).map(str::to_string); - match id.and_then(|id| { - list.iter_mut() - .find(|existing| existing.get("id").and_then(Value::as_str) == Some(&id)) - }) { - Some(existing) => *existing = msg, + 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), } } @@ -171,7 +192,7 @@ impl Channel for Ephemeral { "ephemeral" } - fn merge(&self, _current: Option<&Value>, incoming: Value) -> Result { + fn merge(&self, _current: Option, incoming: Value) -> Result { Ok(incoming) } @@ -189,7 +210,7 @@ impl Channel for Untracked { "untracked" } - fn merge(&self, _current: Option<&Value>, incoming: Value) -> Result { + fn merge(&self, _current: Option, incoming: Value) -> Result { Ok(incoming) } @@ -214,10 +235,11 @@ impl Channel for Barrier { "barrier" } - fn merge(&self, current: Option<&Value>, incoming: Value) -> Result { + 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.clone(), - Some(other) => vec![other.clone()], + Some(Value::Array(items)) => items, + Some(other) => vec![other], None => Vec::new(), }; match incoming { @@ -257,9 +279,10 @@ impl Channel for NamedBarrier { "named_barrier" } - fn merge(&self, current: Option<&Value>, incoming: Value) -> Result { + 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.clone(), + Some(Value::Object(map)) => map, Some(_) => { return Err(TinyAgentsError::Graph( "NamedBarrier channel value must be a JSON object".to_string(), @@ -320,9 +343,9 @@ impl Channel for BinaryAggregate { "binary_aggregate" } - fn merge(&self, current: Option<&Value>, incoming: Value) -> Result { + fn merge(&self, current: Option, incoming: Value) -> Result { match current { - Some(current) => (self.fold)(current.clone(), incoming), + Some(current) => (self.fold)(current, incoming), None => Ok(incoming), } } @@ -385,9 +408,24 @@ impl ChannelSet { /// 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<()> { - let channel = self.channel(name)?; - let merged = channel.merge(self.values.get(name), value)?; + // 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(()) } diff --git a/src/graph/channel/test.rs b/src/graph/channel/test.rs index 37f480e..7ef3ae9 100644 --- a/src/graph/channel/test.rs +++ b/src/graph/channel/test.rs @@ -15,7 +15,7 @@ use serde_json::{Value, json}; 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_eq!(c.merge(Some(json!(1)), json!(2)).unwrap(), json!(2)); assert!(!c.allows_concurrent()); } @@ -23,8 +23,8 @@ fn last_value_overwrites() { 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(); + 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()); } @@ -33,12 +33,12 @@ fn topic_appends_scalars_and_arrays() { fn delta_accumulates_numbers() { let c = Delta; let v = c.merge(None, json!(2)).unwrap(); - let v = c.merge(Some(&v), json!(3)).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(); + 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()); + assert!(c.merge(Some(json!(1)), json!("x")).is_err()); } #[test] @@ -46,11 +46,11 @@ 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"}])) + .merge(Some(v), json!([{"id": "2", "text": "yo"}])) .unwrap(); // upsert existing id 1 let v = c - .merge(Some(&v), json!({"id": "1", "text": "hello"})) + .merge(Some(v), json!({"id": "1", "text": "hello"})) .unwrap(); assert_eq!( v, @@ -59,10 +59,48 @@ fn messages_merge_by_id() { 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_eq!(c.merge(Some(json!(1)), json!(2)).unwrap(), json!(2)); assert!(c.is_ephemeral()); assert!(!c.allows_concurrent()); } @@ -80,8 +118,8 @@ fn binary_aggregate_folds_with_closure() { 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(); + 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()); } @@ -95,8 +133,8 @@ fn binary_aggregate_from_reducer() { 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(); + let v = c.merge(Some(v), json!(7)).unwrap(); + let v = c.merge(Some(v), json!(1)).unwrap(); assert_eq!(v, json!(7)); } @@ -105,7 +143,7 @@ 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(); + let v = c.merge(Some(v), json!("b")).unwrap(); assert!(c.is_ready(Some(&v))); } @@ -114,7 +152,7 @@ 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(); + let v = c.merge(Some(v), json!({"right": 2})).unwrap(); assert!(c.is_ready(Some(&v))); assert!(c.allows_concurrent()); } @@ -148,6 +186,38 @@ fn channel_set_unknown_channel_errors() { )); } +#[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 { diff --git a/src/graph/channel/types.rs b/src/graph/channel/types.rs index 19ad3a1..37b2c1c 100644 --- a/src/graph/channel/types.rs +++ b/src/graph/channel/types.rs @@ -48,7 +48,12 @@ pub trait Channel: Send + Sync { /// Folds `incoming` into the channel's `current` value, returning the new /// value. `current` is `None` the first time the channel is written. - fn merge(&self, current: Option<&Value>, incoming: Value) -> Result; + /// + /// 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; /// Whether more than one concurrent branch may write this channel within a /// single superstep. Aggregates (append/fold/accumulate/barrier) return diff --git a/src/graph/checkpoint/README.md b/src/graph/checkpoint/README.md new file mode 100644 index 0000000..c7f39a4 --- /dev/null +++ b/src/graph/checkpoint/README.md @@ -0,0 +1,95 @@ +# 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 + +- `FileCheckpointer` (`file.rs`) — durable, file-backed implementation; no + external dependencies. +- `SqliteCheckpointer` (`sqlite.rs`, feature `sqlite`) — Sqlite-backed + implementation for concurrent/multi-process access. +- An in-memory backend is also available for tests (see `mod.rs`/`test.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. | +| `file.rs` | `FileCheckpointer`. | +| `sqlite.rs` | `SqliteCheckpointer` (feature `sqlite`). | +| `test.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. +- `FileCheckpointer` and `SqliteCheckpointer` are 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. diff --git a/src/graph/checkpoint/file.rs b/src/graph/checkpoint/file.rs index 8715334..feba4f0 100644 --- a/src/graph/checkpoint/file.rs +++ b/src/graph/checkpoint/file.rs @@ -22,7 +22,14 @@ use async_trait::async_trait; use serde::Serialize; use serde::de::DeserializeOwned; -use super::{Checkpoint, CheckpointMetadata, Checkpointer}; +/// Minimal projection used to read a checkpoint's id without deserializing its +/// `State` payload, so `get` can pick the target line and decode only that one. +#[derive(serde::Deserialize)] +struct CheckpointIdHeader { + checkpoint_id: String, +} + +use super::{Checkpoint, CheckpointConfig, CheckpointMetadata, CheckpointTuple, Checkpointer}; use crate::harness::ids::CheckpointId; use crate::{Result, TinyAgentsError}; @@ -93,6 +100,32 @@ fn io_err(context: &str, err: impl std::fmt::Display) -> TinyAgentsError { TinyAgentsError::Checkpoint(format!("file checkpointer: {context}: {err}")) } +/// Builds a [`CheckpointTuple`] from an owned checkpoint, mirroring the +/// addressing/parent/pending-writes wiring of the default +/// [`Checkpointer::get_tuple`]. +fn tuple_from_checkpoint(checkpoint: Checkpoint) -> CheckpointTuple { + let config = CheckpointConfig { + thread_id: checkpoint.thread_id.clone(), + checkpoint_id: Some(checkpoint.checkpoint_id.clone()), + namespace: checkpoint.namespace.clone(), + }; + let parent_config = checkpoint + .parent_checkpoint_id + .as_ref() + .map(|parent| CheckpointConfig { + thread_id: checkpoint.thread_id.clone(), + checkpoint_id: Some(parent.clone()), + namespace: checkpoint.namespace.clone(), + }); + let pending_writes = checkpoint.pending_writes.clone(); + CheckpointTuple { + config, + checkpoint, + parent_config, + pending_writes, + } +} + impl FileCheckpointer where State: DeserializeOwned, @@ -158,18 +191,26 @@ where { async fn put(&self, checkpoint: Checkpoint) -> Result { let id = CheckpointId::new(checkpoint.checkpoint_id.clone()); - fs::create_dir_all(&self.base_dir).map_err(|e| io_err("create base dir", e))?; + // The serialize + filesystem append is synchronous, blocking work; run + // it on the blocking pool so it never stalls a tokio worker on the + // step-critical path. + let base_dir = self.base_dir.clone(); let path = self.thread_path(&checkpoint.thread_id); - let mut line = - serde_json::to_string(&checkpoint).map_err(|e| io_err("encode record", e))?; - line.push('\n'); - let mut file = OpenOptions::new() - .create(true) - .append(true) - .open(&path) - .map_err(|e| io_err("open thread file for append", e))?; - file.write_all(line.as_bytes()) - .map_err(|e| io_err("append record", e))?; + tokio::task::spawn_blocking(move || -> Result<()> { + fs::create_dir_all(&base_dir).map_err(|e| io_err("create base dir", e))?; + let mut line = + serde_json::to_string(&checkpoint).map_err(|e| io_err("encode record", e))?; + line.push('\n'); + let mut file = OpenOptions::new() + .create(true) + .append(true) + .open(&path) + .map_err(|e| io_err("open thread file for append", e))?; + file.write_all(line.as_bytes()) + .map_err(|e| io_err("append record", e)) + }) + .await + .map_err(|e| io_err("join blocking put task", e))??; Ok(id) } @@ -178,12 +219,41 @@ where thread_id: &str, checkpoint_id: Option<&str>, ) -> Result>> { - let records = self.read_records(thread_id)?; - let found = match checkpoint_id { - Some(id) => records.into_iter().rev().find(|c| c.checkpoint_id == id), - None => records.into_iter().next_back(), + // Stream lines and fully decode only the single target line, instead of + // deserializing every record's `State` just to pick one. Selection + // matches the previous `rev().find` / `next_back` semantics: the last + // matching line (or the last line, for `None`) wins. + let path = self.thread_path(thread_id); + let file = match File::open(&path) { + Ok(f) => f, + Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(None), + Err(e) => return Err(io_err("open thread file", e)), }; - Ok(found) + let reader = BufReader::new(file); + let mut target: Option = None; + for line in reader.lines() { + let line = line.map_err(|e| io_err("read line", e))?; + if line.trim().is_empty() { + continue; + } + match checkpoint_id { + Some(id) => { + // Decode only the id header to test the match, not `State`. + let header: CheckpointIdHeader = + serde_json::from_str(&line).map_err(|e| io_err("decode header", e))?; + if header.checkpoint_id == id { + target = Some(line); + } + } + None => target = Some(line), + } + } + match target { + Some(line) => Ok(Some( + serde_json::from_str(&line).map_err(|e| io_err("decode record", e))?, + )), + None => Ok(None), + } } async fn list(&self, thread_id: &str) -> Result> { @@ -194,6 +264,55 @@ where .collect()) } + async fn state_history( + &self, + thread_id: &str, + namespace: &[String], + limit: Option, + ) -> Result>> { + // Read the whole thread once, then walk the parent lineage in memory + // (O(H)), instead of re-reading and re-parsing the file per hop (O(H²)). + let records = self.read_records(thread_id)?; + if records.is_empty() { + return Ok(Vec::new()); + } + + // id -> checkpoint, last write wins for duplicate ids (matching `get`, + // which takes the last matching record). Track the latest checkpoint in + // the target namespace as the walk's starting point. + let mut by_id: std::collections::HashMap> = + std::collections::HashMap::with_capacity(records.len()); + let mut cursor: Option = None; + for record in records { + if record.namespace.as_slice() == namespace { + cursor = Some(record.checkpoint_id.clone()); + } + by_id.insert(record.checkpoint_id.clone(), record); + } + + let mut out = Vec::new(); + while let Some(id) = cursor { + if let Some(limit) = limit + && out.len() >= limit + { + break; + } + // `remove` doubles as a cycle guard: each id is visited at most once. + let Some(checkpoint) = by_id.remove(&id) else { + break; + }; + // A checkpoint outside the target namespace is not visible under + // namespace-scoped lookup, so the lineage walk stops (matching the + // `get_scoped`-based default). + if checkpoint.namespace.as_slice() != namespace { + break; + } + cursor = checkpoint.parent_checkpoint_id.clone(); + out.push(tuple_from_checkpoint(checkpoint)); + } + Ok(out) + } + async fn list_threads(&self) -> Result> { let entries = match fs::read_dir(&self.base_dir) { Ok(e) => e, diff --git a/src/graph/checkpoint/mod.rs b/src/graph/checkpoint/mod.rs index 6841ea5..a07651f 100644 --- a/src/graph/checkpoint/mod.rs +++ b/src/graph/checkpoint/mod.rs @@ -22,8 +22,8 @@ pub use file::FileCheckpointer; #[cfg(feature = "sqlite")] pub use sqlite::SqliteCheckpointer; pub use types::{ - Checkpoint, CheckpointConfig, CheckpointMetadata, CheckpointSource, CheckpointTuple, - DurabilityMode, PendingWrite, + BarrierArrivals, Checkpoint, CheckpointConfig, CheckpointMetadata, CheckpointSource, + CheckpointTuple, DurabilityMode, PendingActivation, PendingWrite, }; use std::collections::{HashMap, HashSet}; @@ -51,6 +51,42 @@ where 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 only for a cheaper scoped query. + 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>; @@ -62,7 +98,11 @@ where /// `config.checkpoint_id` is `None` the latest checkpoint is returned. async fn get_tuple(&self, config: CheckpointConfig) -> Result>> { let Some(checkpoint) = self - .get(&config.thread_id, config.checkpoint_id.as_deref()) + .get_scoped( + &config.thread_id, + config.checkpoint_id.as_deref(), + &config.namespace, + ) .await? else { return Ok(None); @@ -90,6 +130,48 @@ where })) } + /// 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 + /// that re-reads the whole thread per lookup (the file/JSONL backend) is + /// O(H²) over the lineage. Such backends override this to read the thread + /// once and walk the lineage in memory (O(H)). The observable result is + /// identical to iterating `get_tuple` by parent pointer. + async fn state_history( + &self, + thread_id: &str, + namespace: &[String], + limit: Option, + ) -> Result>> { + let mut out = Vec::new(); + let mut cursor: Option = None; + 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; + }; + let parent = tuple.checkpoint.parent_checkpoint_id.clone(); + out.push(tuple); + match parent { + Some(parent) => cursor = Some(parent), + None => break, + } + } + Ok(out) + } + // ---- Thread operations ------------------------------------------------- // // Three storage-specific primitives (`list_threads`, `delete_thread`, @@ -284,8 +366,12 @@ where 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().find(|c| c.checkpoint_id == id), + Some(id) => list.iter().rfind(|c| c.checkpoint_id == id), None => list.last(), }; Ok(found.cloned()) diff --git a/src/graph/checkpoint/sqlite.rs b/src/graph/checkpoint/sqlite.rs index 6264885..fd02d4a 100644 --- a/src/graph/checkpoint/sqlite.rs +++ b/src/graph/checkpoint/sqlite.rs @@ -177,34 +177,47 @@ where { async fn put(&self, checkpoint: Checkpoint) -> Result { let id = CheckpointId::new(checkpoint.checkpoint_id.clone()); - let meta = checkpoint.to_metadata(); - let namespace = serde_json::to_string(&checkpoint.namespace) - .map_err(|e| sqlite_err("encode namespace", e))?; - let next_nodes = serde_json::to_string(&checkpoint.next_nodes) - .map_err(|e| sqlite_err("encode next_nodes", e))?; - let record = - serde_json::to_string(&checkpoint).map_err(|e| sqlite_err("encode record", e))?; + // Serialize + the synchronous rusqlite insert (which also blocks on the + // connection mutex) is blocking work; run it on the blocking pool so it + // never stalls a tokio worker on the step-critical path. + let conn = self.conn.clone(); + tokio::task::spawn_blocking(move || -> Result<()> { + let meta = checkpoint.to_metadata(); + let namespace = serde_json::to_string(&checkpoint.namespace) + .map_err(|e| sqlite_err("encode namespace", e))?; + let next_nodes = serde_json::to_string(&checkpoint.next_nodes) + .map_err(|e| sqlite_err("encode next_nodes", e))?; + let record = + serde_json::to_string(&checkpoint).map_err(|e| sqlite_err("encode record", e))?; - let conn = self.lock()?; - conn.execute( - "INSERT INTO checkpoints ( + let conn = conn.lock().map_err(|_| { + TinyAgentsError::Checkpoint( + "sqlite checkpointer: connection lock poisoned".to_string(), + ) + })?; + conn.execute( + "INSERT INTO checkpoints ( thread_id, checkpoint_id, parent_checkpoint_id, run_id, namespace, next_nodes, source, step, has_interrupts, record ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10)", - params![ - checkpoint.thread_id, - checkpoint.checkpoint_id, - checkpoint.parent_checkpoint_id, - checkpoint.run_id, - namespace, - next_nodes, - meta.source.as_str(), - meta.step as i64, - i64::from(meta.has_interrupts), - record, - ], - ) - .map_err(|e| sqlite_err("insert checkpoint", e))?; + params![ + checkpoint.thread_id, + checkpoint.checkpoint_id, + checkpoint.parent_checkpoint_id, + checkpoint.run_id, + namespace, + next_nodes, + meta.source.as_str(), + meta.step as i64, + i64::from(meta.has_interrupts), + record, + ], + ) + .map_err(|e| sqlite_err("insert checkpoint", e))?; + Ok(()) + }) + .await + .map_err(|e| sqlite_err("join blocking put task", e))??; Ok(id) } diff --git a/src/graph/checkpoint/test.rs b/src/graph/checkpoint/test.rs index 5693ec5..f694278 100644 --- a/src/graph/checkpoint/test.rs +++ b/src/graph/checkpoint/test.rs @@ -18,6 +18,8 @@ fn checkpoint(thread: &str, id: &str, parent: Option<&str>, step: usize) -> Chec completed_tasks: vec![], pending_writes: vec![], interrupts: vec![], + pending_activations: None, + barrier_arrivals: vec![], metadata: json!({ "source": "loop", "step": step }), } } @@ -49,6 +51,61 @@ async fn put_get_list_roundtrip() { 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 cp: Checkpoint = serde_json::from_value(legacy).unwrap(); + assert_eq!(cp.state, 7); + assert_eq!(cp.next_nodes.len(), 2); + assert!(cp.pending_activations.is_none()); + assert!(cp.barrier_arrivals.is_empty()); +} + +#[test] +fn pending_activation_send_arg_roundtrips() { + let cp = Checkpoint { + thread_id: "t".into(), + checkpoint_id: "c1".into(), + run_id: None, + parent_checkpoint_id: None, + namespace: vec![], + state: 1i32, + next_nodes: vec![NodeId::from("w")], + completed_tasks: vec![], + pending_writes: vec![], + interrupts: vec![], + pending_activations: Some(vec![super::PendingActivation { + node: NodeId::from("w"), + send_arg: Some(json!({ "item": 42 })), + }]), + barrier_arrivals: vec![super::BarrierArrivals { + node: NodeId::from("join"), + arrived: vec![NodeId::from("p1")], + }], + metadata: json!({ "source": "loop", "step": 1 }), + }; + let round: Checkpoint = + serde_json::from_str(&serde_json::to_string(&cp).unwrap()).unwrap(); + let pa = round.pending_activations.unwrap(); + assert_eq!(pa[0].send_arg, Some(json!({ "item": 42 }))); + assert_eq!(round.barrier_arrivals[0].arrived, vec![NodeId::from("p1")]); +} + #[tokio::test] async fn clones_share_storage() { let cp = InMemoryCheckpointer::::new(); diff --git a/src/graph/checkpoint/types.rs b/src/graph/checkpoint/types.rs index 001d7cb..8901298 100644 --- a/src/graph/checkpoint/types.rs +++ b/src/graph/checkpoint/types.rs @@ -163,10 +163,55 @@ pub struct Checkpoint { pub pending_writes: Vec, /// Interrupts that paused the run at this boundary. pub interrupts: Vec, + /// Pending activations to schedule on resume, preserving each pending + /// node's per-invocation [`Send`](crate::graph::Send) argument. + /// + /// A richer superset of [`next_nodes`](Self::next_nodes) (which stays the + /// node-id projection used for listing and status). `#[serde(default)]` + /// keeps checkpoints written before this field loadable: they deserialize + /// to `None`, and resume falls back to `next_nodes` (node-only, no send + /// arg) — exactly the pre-field behavior. + #[serde(default)] + pub pending_activations: Option>, + /// Barrier (waiting-edge) arrivals accumulated across supersteps, persisted + /// so a join node's precondition survives an interrupt/failure + resume. + /// + /// `#[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, /// Free-form metadata (source, step, etc.). pub metadata: serde_json::Value, } +/// One pending node activation persisted in a checkpoint: the node to run on +/// resume plus the optional per-invocation [`Send`](crate::graph::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`). + #[serde(default, skip_serializing_if = "Option::is_none")] + pub send_arg: Option, +} + +/// 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 { /// Builds the lightweight [`CheckpointMetadata`] summary for this checkpoint. /// diff --git a/src/graph/compiled/README.md b/src/graph/compiled/README.md new file mode 100644 index 0000000..0b878d6 --- /dev/null +++ b/src/graph/compiled/README.md @@ -0,0 +1,143 @@ +# 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). +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. + +### 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). | +| `test.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/src/graph/compiled/executor.rs b/src/graph/compiled/executor.rs new file mode 100644 index 0000000..1eca370 --- /dev/null +++ b/src/graph/compiled/executor.rs @@ -0,0 +1,1319 @@ +//! 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). + +use super::*; + +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( + state, + vec![Activation::node(self.entry.clone())], + None, + HashMap::new(), + HashMap::new(), + None, + ) + .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::graph::NodeContext::send_arg). + pub async fn run_with_inputs( + &self, + state: State, + inputs: impl IntoIterator, + ) -> Result> { + let active = self.initial_inputs(inputs)?; + self.execute(state, active, None, HashMap::new(), HashMap::new(), 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( + state, + vec![Activation::node(self.entry.clone())], + Some(thread_id.into()), + HashMap::new(), + HashMap::new(), + None, + ) + .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( + state, + active, + Some(thread_id.into()), + HashMap::new(), + HashMap::new(), + None, + ) + .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 + } + + /// 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 + } + + /// 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> { + let checkpointer = self + .checkpointer + .as_ref() + .ok_or_else(|| TinyAgentsError::Resume("no checkpointer configured".to_string()))?; + let thread_id = thread_id.into(); + + let checkpoint_id = match &target { + ResumeTarget::Latest => None, + ResumeTarget::Checkpoint(id) => Some(id.as_str()), + }; + let checkpoint = checkpointer + .get_scoped(thread_id.as_str(), checkpoint_id, &self.namespace) + .await? + .ok_or_else(|| match &target { + ResumeTarget::Latest => { + TinyAgentsError::Resume(format!("no checkpoint found for thread `{thread_id}`")) + } + ResumeTarget::Checkpoint(id) => TinyAgentsError::Resume(format!( + "no checkpoint `{id}` found for thread `{thread_id}`" + )), + })?; + // Resume *loads* this checkpoint — it is a read, not a write — so emit a + // restore event, not `CheckpointSaved` (which would falsely inflate + // persisted-checkpoint counts and mislead durability observers). + self.emit(GraphEvent::CheckpointRestored { + checkpoint_id: CheckpointId::new(checkpoint.checkpoint_id.clone()), + }); + + // Prefer the persisted pending activations (which preserve each pending + // node's `Send` arg); fall back to the node-id projection for + // checkpoints written before that field existed. + let active: Vec = match &checkpoint.pending_activations { + Some(pending) if !pending.is_empty() => pending.iter().map(Activation::from).collect(), + _ => checkpoint + .next_nodes + .iter() + .cloned() + .map(Activation::node) + .collect(), + }; + if active.is_empty() { + return Err(TinyAgentsError::Resume( + "checkpoint has no pending nodes to resume".to_string(), + )); + } + + let mut resume_map = HashMap::new(); + if let Some(value) = command.resume { + for activation in &active { + resume_map.insert(activation.node.clone(), value.clone()); + } + } + + // Restore accumulated barrier arrivals so a join's precondition survives + // the interrupt/failure boundary this checkpoint recorded. + let initial_barriers = barriers_from_persisted(&checkpoint.barrier_arrivals); + // Chain the first post-resume boundary onto the checkpoint we loaded so + // the lineage spine stays connected across the resume. + let initial_parent = Some(checkpoint.checkpoint_id.clone()); + + self.execute( + checkpoint.state, + active, + Some(thread_id), + resume_map, + initial_barriers, + initial_parent, + ) + .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, + }); + } + 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. + async fn execute( + &self, + state: State, + initial_active: Vec, + thread_id: Option, + resume_map: HashMap, + initial_barriers: HashMap>, + initial_parent: Option, + ) -> Result> { + let run_id = crate::harness::ids::new_run_id(); + // When a durable journal is configured, run against a clone whose event + // sink wraps every emitted event into a `GraphObservation` and appends + // it (while still forwarding to any pre-existing live sink). The journal + // sink carries this graph's checkpoint namespace so subgraph runs record + // their nested path. Default (no journal) leaves `self` untouched. + if self.journal.is_some() { + let this = self.clone_with_journal_sink(&run_id, &thread_id); + this.execute_run( + run_id, + state, + initial_active, + thread_id, + resume_map, + initial_barriers, + initial_parent, + ) + .await + } else { + self.execute_run( + run_id, + state, + initial_active, + thread_id, + resume_map, + initial_barriers, + initial_parent, + ) + .await + } + } + + /// 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::graph::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 + } + + /// Best-effort status write; never aborts the run on a status-store error. + async fn save_status(&self, status: GraphRunStatus) { + if let Some(store) = &self.status_store { + let _ = store.put_status(status).await; + } + } + + #[allow(clippy::too_many_arguments)] + async fn execute_run( + &self, + run_id: RunId, + mut state: State, + initial_active: Vec, + thread_id: Option, + mut resume_map: HashMap, + initial_barriers: HashMap>, + initial_parent: Option, + ) -> Result> { + let started_at = SystemTime::now(); + let mut visited: Vec = Vec::new(); + let mut steps = 0usize; + let mut last_checkpoint: Option = None; + // On resume this is the loaded checkpoint's id, so the first boundary + // checkpoint after a resume chains onto pre-interrupt history rather + // than orphaning the lineage (which would stop `get_state_history` at + // the resume point and let `prune` delete the ancestors). + let mut parent_checkpoint: Option = initial_parent; + + // Build this run's recursion stack from the inherited parent frames and + // push the frame for this graph call. A push that would exceed + // `max_depth` fails the run with a clear recursion error before any + // node executes. Graph-call depth (the stack) is tracked separately + // from node-loop visits (`node_visits`, below). + let mut recursion = + RecursionStack::with_frames(self.recursion_frames.clone(), self.recursion_policy); + // Run lineage: the root is the first inherited frame's run (the top of + // the recursion tree) or this run when top-level; the parent is the + // enclosing run, if any. + let root_run_id = self + .recursion_frames + .first() + .map(|f| f.run_id.clone()) + .unwrap_or_else(|| run_id.clone()); + let parent_run_id = self.recursion_frames.last().map(|f| f.run_id.clone()); + let this_frame = RecursionFrame { + graph_id: self.graph_id.clone(), + node_id: self.recursion_node.clone(), + run_id: run_id.clone(), + task_id: None, + namespace: self.namespace.clone(), + depth: recursion.depth(), + parent: parent_run_id.clone(), + }; + if let Err(err) = recursion.push(this_frame) { + self.emit(GraphEvent::RunStarted { + run_id: run_id.clone(), + }); + self.fail_run(&run_id, &thread_id, started_at, steps, &err, None) + .await; + return Err(err); + } + // Serialized once per run for embedding in every checkpoint's metadata. + let recursion_meta = + serde_json::to_value(recursion.frames()).unwrap_or(serde_json::Value::Null); + // The live frame stack handed to node contexts so a subgraph node can + // seed an embedded child with this run's recursion path, plus the + // per-run sink the node reports its spawned child run into. + let live_frames = recursion.frames().to_vec(); + let child_sink = ChildRunSink::new(); + // Accumulates every child run spawned across all supersteps for the + // final `GraphExecution::child_runs`. + let mut all_child_runs: Vec = Vec::new(); + // Per-node activation counts for `max_visits_per_node` enforcement. + let mut node_visits: HashMap = HashMap::new(); + let mut active = initial_active; + // Barrier/waiting-edge arrivals accumulate across supersteps: a waiting + // node only activates once every required predecessor has arrived. + // Seeded from the resumed checkpoint so a join's precondition survives + // an interrupt/failure boundary. + let mut barrier_arrivals: HashMap> = initial_barriers; + + self.emit(GraphEvent::RunStarted { + run_id: run_id.clone(), + }); + // Surface this run's recursion depth so observers can attribute nested + // runs without reconstructing the tree from logs. + self.emit(GraphEvent::RecursionDepthChanged { + depth: recursion.depth(), + }); + // Record the run as live before the first superstep is scheduled. + let mut running = self.base_status(&run_id, &thread_id, started_at); + running.active_nodes = activation_nodes(&active); + self.save_status(running).await; + + while !active.is_empty() { + // The effective step cap is the smaller of the builder's recursion + // limit and the policy's `max_total_steps`, so a policy never + // loosens an existing limit. Both surface a `RecursionLimit`. + let step_limit = self + .recursion_limit + .min(self.recursion_policy.max_total_steps); + if steps >= step_limit { + let err = TinyAgentsError::RecursionLimit(step_limit); + self.fail_run(&run_id, &thread_id, started_at, steps, &err, None) + .await; + return Err(err); + } + // Node-loop recursion: enforce `max_visits_per_node` per activation. + for activation in &active { + if let Err(err) = recursion.record_node_visit(&mut node_visits, &activation.node) { + self.fail_run(&run_id, &thread_id, started_at, steps, &err, None) + .await; + return Err(err); + } + } + steps += 1; + self.emit(GraphEvent::StepStarted { + step: steps, + active: activation_nodes(&active), + }); + + let run_result = if self.parallel && active.len() > 1 { + self.run_active_parallel( + &active, + &state, + &run_id, + &thread_id, + steps, + &mut resume_map, + &mut visited, + &root_run_id, + &live_frames, + &child_sink, + ) + .await + } else { + self.run_active_sequential( + &active, + &state, + &run_id, + &thread_id, + steps, + &mut resume_map, + &mut visited, + &root_run_id, + &live_frames, + &child_sink, + ) + .await + }; + let StepRun { + updates, + goto_map, + interrupt, + failure, + } = match run_result { + Ok(step_run) => step_run, + Err(err) => { + self.fail_run(&run_id, &thread_id, started_at, steps, &err, None) + .await; + return Err(err); + } + }; + + // Apply collected updates through the reducer at the boundary. A + // reducer error here must still fail the run (not just unwind + // leaving it `Running`). + for update in updates { + state = match self.reducer.apply(state, update) { + Ok(state) => state, + Err(err) => { + return self + .fail_and_return(&run_id, &thread_id, started_at, steps, err) + .await; + } + }; + } + + // Collect any child runs spawned by subgraph nodes this step. They + // are embedded into this boundary's checkpoint metadata (keyed by + // node) and accumulated onto the final `GraphExecution`. + let step_child_runs = child_sink.drain(); + all_child_runs.extend(step_child_runs.iter().cloned()); + let child_runs_meta = + serde_json::to_value(&step_child_runs).unwrap_or(serde_json::Value::Null); + + // Node-handler failure (survived any node-retry policy): the updates + // of the branches that completed before it are already folded into + // `state` above, so persist a resumable failure-boundary checkpoint + // scheduling the failed node (and the not-yet-run tail) for a later + // `resume`/`retry`, record a `Failed` status carrying the error and + // that checkpoint, and abort. Without a checkpointer/thread the + // checkpoint is a no-op and the run aborts exactly as before. + if let Some(fail) = failure { + let StepFailure { + failed_index, + error, + } = fail; + let failed_node = active[failed_index].node.clone(); + // Schedule the successors of the branches that completed before + // the failure (they succeeded; their routing must not be lost) + // followed by the failed branch and the not-yet-run tail, which + // re-run on resume with their `Send` args preserved. + let successors = match self.route_completed( + &active[..failed_index], + &goto_map, + &state, + &mut barrier_arrivals, + ) { + Ok(successors) => successors, + Err(route_err) => { + return self + .fail_and_return(&run_id, &thread_id, started_at, steps, route_err) + .await; + } + }; + let mut pending = successors; + pending.extend(active[failed_index..].iter().cloned()); + let completed_nodes = activation_nodes(&active[..failed_index]); + // A failure-boundary persist error must not replace the original + // node error: keep reporting the node error and just drop the + // resumable checkpoint reference. + let checkpoint_id = self + .persist_failure_checkpoint( + &thread_id, + &run_id, + &state, + &pending, + &completed_nodes, + &barrier_arrivals, + parent_checkpoint.clone(), + steps, + &failed_node, + &error, + &recursion_meta, + &child_runs_meta, + ) + .await + .unwrap_or(None); + self.fail_run( + &run_id, + &thread_id, + started_at, + steps, + &error, + checkpoint_id, + ) + .await; + return Err(error); + } + + // Interrupt: persist a checkpoint whose pending activations are the + // successors of the branches that completed before the interrupt + // (their routing must survive) followed by the not-yet-completed + // members of this step (interrupted node first). Each pending branch + // keeps its `Send` arg; accumulated barrier arrivals are persisted + // too. Then return control to the caller. + if let Some((index, emitted)) = interrupt { + if let Err(err) = self.require_interrupt_durability(&thread_id) { + self.fail_run(&run_id, &thread_id, started_at, steps, &err, None) + .await; + return Err(err); + } + let successors = match self.route_completed( + &active[..index], + &goto_map, + &state, + &mut barrier_arrivals, + ) { + Ok(successors) => successors, + Err(route_err) => { + return self + .fail_and_return(&run_id, &thread_id, started_at, steps, route_err) + .await; + } + }; + let mut pending = successors; + pending.extend(active[index..].iter().cloned()); + let pending_nodes = activation_nodes(&pending); + let interrupt_id = InterruptId::new(emitted.id.clone()); + let checkpoint_id = match self + .persist_checkpoint( + &thread_id, + &run_id, + &state, + &pending, + &activation_nodes(&active[..index]), + vec![emitted.clone()], + &barrier_arrivals, + parent_checkpoint.clone(), + steps, + "loop", + &recursion_meta, + &child_runs_meta, + ) + .await + { + Ok(id) => id, + Err(persist_err) => { + return self + .fail_and_return(&run_id, &thread_id, started_at, steps, persist_err) + .await; + } + }; + + let mut status = self.base_status(&run_id, &thread_id, started_at); + status.status = ExecutionStatus::Interrupted; + status.current_step = steps; + status.active_nodes = pending_nodes; + status.pending_interrupts = vec![interrupt_id]; + status.checkpoint_id = checkpoint_id.clone(); + self.save_status(status.clone()).await; + + return Ok(GraphExecution { + state, + run_id: run_id.clone(), + graph_id: self.graph_id.clone(), + root_run_id: root_run_id.clone(), + parent_run_id: parent_run_id.clone(), + child_runs: all_child_runs, + visited, + steps, + interrupts: vec![emitted], + status, + checkpoint_id, + }); + } + + // Select the next active set from commands or static/conditional + // edges, evaluated against the freshly-committed state. Barrier + // arrivals accumulate into `barrier_arrivals` (persisted below). + let completed_nodes = activation_nodes(&active); + let next = match self.route_completed(&active, &goto_map, &state, &mut barrier_arrivals) + { + Ok(next) => next, + Err(route_err) => { + return self + .fail_and_return(&run_id, &thread_id, started_at, steps, route_err) + .await; + } + }; + + // Persist a boundary checkpoint. Under `Exit` durability only the + // terminal boundary (the step that empties the active set) is + // written; `Sync`/`Async` persist every boundary. + let persist_now = match self.durability { + DurabilityMode::Exit => next.is_empty(), + DurabilityMode::Sync | DurabilityMode::Async => true, + }; + let checkpoint_id = if persist_now { + match self + .persist_checkpoint( + &thread_id, + &run_id, + &state, + &next, + &completed_nodes, + Vec::new(), + &barrier_arrivals, + parent_checkpoint.clone(), + steps, + "loop", + &recursion_meta, + &child_runs_meta, + ) + .await + { + Ok(id) => id, + Err(persist_err) => { + return self + .fail_and_return(&run_id, &thread_id, started_at, steps, persist_err) + .await; + } + } + } else { + None + }; + if let Some(id) = &checkpoint_id { + last_checkpoint = Some(id.clone()); + parent_checkpoint = Some(id.to_string()); + } + + self.emit(GraphEvent::StepCompleted { step: steps }); + active = next; + } + + let mut status = self.base_status(&run_id, &thread_id, started_at); + status.status = ExecutionStatus::Completed; + status.current_step = steps; + status.checkpoint_id = last_checkpoint.clone(); + status.ended_at = Some(SystemTime::now()); + self.save_status(status.clone()).await; + self.emit(GraphEvent::RunCompleted { + run_id: run_id.clone(), + steps, + }); + + Ok(GraphExecution { + state, + run_id: run_id.clone(), + graph_id: self.graph_id.clone(), + root_run_id, + parent_run_id, + child_runs: all_child_runs, + visited, + steps, + interrupts: Vec::new(), + status, + checkpoint_id: last_checkpoint, + }) + } + + /// Emits a [`GraphEvent::RunFailed`] and records a terminal `Failed` status + /// for a run that aborted with `err`. + /// + /// `checkpoint_id` is the resumable failure-boundary checkpoint when the run + /// left one (a node-handler failure on a checkpointed thread), or `None` for + /// a structural/non-resumable abort. When present it is recorded on the + /// status so an observer can locate the checkpoint to `resume`/`retry` from. + async fn fail_run( + &self, + run_id: &RunId, + thread_id: &Option, + started_at: SystemTime, + steps: usize, + err: &TinyAgentsError, + checkpoint_id: Option, + ) { + self.emit(GraphEvent::RunFailed { + run_id: run_id.clone(), + error: err.to_string(), + }); + let mut status = self.base_status(run_id, thread_id, started_at); + status.status = ExecutionStatus::Failed; + status.current_step = steps; + status.ended_at = Some(SystemTime::now()); + status.error = Some(err.to_string()); + status.checkpoint_id = checkpoint_id; + self.save_status(status).await; + } + + /// Records a terminal `Failed` status for `err` (via [`Self::fail_run`]) and + /// returns it as `Err`. + /// + /// Used at the step boundary so an error raised *after* the node runners — + /// a reducer merge, a routing resolution, or a checkpoint persist — still + /// transitions the run to `Failed` (rather than leaving observers to see it + /// stuck in `Running` forever) before the error unwinds out of the run. + async fn fail_and_return( + &self, + run_id: &RunId, + thread_id: &Option, + started_at: SystemTime, + steps: usize, + err: TinyAgentsError, + ) -> Result { + self.fail_run(run_id, thread_id, started_at, 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. + #[allow(clippy::too_many_arguments)] + async fn persist_failure_checkpoint( + &self, + thread_id: &Option, + run_id: &RunId, + state: &State, + pending: &[Activation], + completed_tasks: &[NodeId], + barrier_arrivals: &HashMap>, + parent: Option, + step: usize, + failed_node: &NodeId, + error: &TinyAgentsError, + recursion: &serde_json::Value, + child_runs: &serde_json::Value, + ) -> Result> { + let (Some(checkpointer), Some(thread)) = (&self.checkpointer, thread_id) else { + return Ok(None); + }; + let checkpoint = Checkpoint { + thread_id: thread.to_string(), + checkpoint_id: next_checkpoint_id(), + run_id: Some(run_id.to_string()), + parent_checkpoint_id: parent, + namespace: self.namespace.clone(), + state: state.clone(), + next_nodes: activation_nodes(pending), + completed_tasks: completed_tasks.to_vec(), + pending_writes: Vec::new(), + interrupts: Vec::new(), + pending_activations: Some(pending.iter().map(PendingActivation::from).collect()), + barrier_arrivals: barriers_to_persisted(barrier_arrivals), + metadata: serde_json::json!({ + "source": "loop", + "step": step, + "recursion": recursion, + "child_runs": child_runs, + "failed_node": failed_node.as_str(), + "error": error.to_string(), + }), + }; + let id = checkpointer.put(checkpoint).await?; + self.emit(GraphEvent::CheckpointSaved { + checkpoint_id: id.clone(), + }); + Ok(Some(id)) + } + + /// Builds the per-task [`NodeContext`] for `node_id` at the given branch. + /// + /// `fork` carries the branch identity in a concurrent step (`None` in + /// sequential mode or single-node steps). The resume value for the node is + /// consumed from `resume_map`. + #[allow(clippy::too_many_arguments)] + fn node_context( + &self, + node_id: &NodeId, + run_id: &RunId, + thread_id: &Option, + step: usize, + resume_map: &mut HashMap, + fork: Option, + send_arg: Option, + root_run_id: &RunId, + frames: &[RecursionFrame], + child_runs: &ChildRunSink, + ) -> NodeContext { + NodeContext { + node_id: node_id.clone(), + run_id: run_id.clone(), + thread_id: thread_id.clone(), + step, + resume: resume_map.remove(node_id), + fork, + send_arg, + root_run_id: Some(root_run_id.clone()), + recursion_frames: frames.to_vec(), + child_runs: Some(child_runs.clone()), + } + } + + /// Wraps a node future in the configured per-node timeout (if any), mapping + /// an elapsed deadline onto [`TinyAgentsError::Timeout`]. + async fn run_node_future( + &self, + node_id: &NodeId, + fut: NodeFuture, + ) -> Result> { + match self.node_timeout { + Some(timeout) => match tokio::time::timeout(timeout, fut).await { + Ok(result) => result, + Err(_) => Err(TinyAgentsError::Timeout(format!( + "node `{node_id}` exceeded its {timeout:?} timeout" + ))), + }, + None => fut.await, + } + } + + /// Runs one node handler under the graph's node-retry policy. + /// + /// Builds a fresh handler future (and re-clones the context) for each + /// attempt, so a retried node re-runs from its start — matching the durable + /// execution model, where a node is never suspended mid-flight. On a + /// [retryable][crate::harness::retry::is_retryable] error, when a + /// [`RetryPolicy`](crate::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: &State, + ctx: NodeContext, + step: usize, + ) -> Result> { + let mut attempt = 0usize; + loop { + let fut = handler(state.clone(), ctx.clone()); + match self.run_node_future(node_id, fut).await { + Ok(result) => return Ok(result), + Err(error) => { + let retry = self + .node_retry + .as_ref() + .filter(|policy| policy.should_retry(attempt) && is_retryable(&error)); + let Some(policy) = retry else { + return Err(error); + }; + attempt += 1; + self.emit(GraphEvent::NodeRetryScheduled { + node: node_id.clone(), + step, + attempt, + }); + policy.sleep_backoff(attempt).await; + } + } + } + } + + /// Folds a single successful branch result into the step accumulators. + /// + /// Pushes the node to `visited`, records updates/goto, emits the matching + /// events, and returns the interrupt (with its branch index) when the branch + /// paused. Shared by the sequential and parallel run paths so both fold + /// results identically; only the *running* of handlers differs. + #[allow(clippy::too_many_arguments)] + fn fold_result( + &self, + index: usize, + node_id: &NodeId, + step: usize, + result: NodeResult, + updates: &mut Vec, + goto_map: &mut HashMap>, + visited: &mut Vec, + ) -> Option<(usize, Interrupt)> { + visited.push(node_id.clone()); + match result { + NodeResult::Update(update) => { + updates.push(update); + self.emit(GraphEvent::StateUpdated { + node: node_id.clone(), + step, + }); + } + NodeResult::Command(command) => { + if let Some(update) = command.update { + updates.push(update); + self.emit(GraphEvent::StateUpdated { + node: node_id.clone(), + step, + }); + } + if !command.goto.is_empty() { + goto_map.insert(index, command.goto); + } + } + NodeResult::Interrupt(emitted) => { + self.emit(GraphEvent::InterruptEmitted { + interrupt: emitted.clone(), + }); + return Some((index, emitted)); + } + } + self.emit(GraphEvent::NodeCompleted { + node: node_id.clone(), + step, + }); + None + } + + /// Runs the active node set one node at a time (default behavior). + /// + /// Short-circuits on the first error (run aborts) or interrupt (later nodes + /// in the step are not started), exactly preserving milestone-1 semantics. + #[allow(clippy::too_many_arguments)] + async fn run_active_sequential( + &self, + active: &[Activation], + state: &State, + run_id: &RunId, + thread_id: &Option, + step: usize, + resume_map: &mut HashMap, + visited: &mut Vec, + root_run_id: &RunId, + frames: &[RecursionFrame], + child_runs: &ChildRunSink, + ) -> Result> { + let mut updates: Vec = Vec::new(); + let mut goto_map: HashMap> = HashMap::new(); + let mut interrupt: Option<(usize, Interrupt)> = None; + let mut failure: Option = None; + + for (index, activation) in active.iter().enumerate() { + let node_id = &activation.node; + let node = self + .nodes + .get(node_id) + .ok_or_else(|| TinyAgentsError::MissingNode(node_id.to_string()))?; + + self.emit(GraphEvent::TaskScheduled { + node: node_id.clone(), + step, + }); + self.emit(GraphEvent::NodeStarted { + node: node_id.clone(), + step, + }); + + let ctx = self.node_context( + node_id, + run_id, + thread_id, + step, + resume_map, + None, + activation.send_arg.clone(), + root_run_id, + frames, + child_runs, + ); + let result = match self + .run_node_with_retry(node_id, &node.handler, state, ctx, step) + .await + { + Ok(result) => result, + Err(error) => { + self.emit(GraphEvent::NodeFailed { + node: node_id.clone(), + step, + error: error.to_string(), + }); + // Preserve the progress of the branches that already ran: + // the executor records them as completed and schedules their + // successors plus this node and the not-yet-run tail for a + // resumable retry. + failure = Some(StepFailure { + failed_index: index, + error, + }); + break; + } + }; + + if let Some(found) = self.fold_result( + index, + node_id, + step, + result, + &mut updates, + &mut goto_map, + visited, + ) { + interrupt = Some(found); + break; + } + } + + Ok(StepRun { + updates, + goto_map, + interrupt, + failure, + }) + } + + /// Runs the active node set concurrently (opt-in via `with_parallel`). + /// + /// Each branch executes on its own cloned `State` snapshot and a distinct + /// [`ForkId`], optionally with the [`Send`] argument that scheduled it. With + /// no `max_concurrency` bound every branch starts before any is awaited and + /// all are driven via [`futures::future::join_all`]; with a bound the active + /// set is run in chunks of at most that many futures, so at most that many + /// node handlers are in flight at once. Results are folded in active-set + /// index order — the reducer is the join/fan-in — so the merged state is + /// reproducible regardless of completion order. The lowest-index branch that + /// errors or interrupts is the step's terminal outcome; lower-index + /// successful branches still contribute their updates. + #[allow(clippy::too_many_arguments)] + async fn run_active_parallel( + &self, + active: &[Activation], + state: &State, + run_id: &RunId, + thread_id: &Option, + step: usize, + resume_map: &mut HashMap, + visited: &mut Vec, + root_run_id: &RunId, + frames: &[RecursionFrame], + child_runs: &ChildRunSink, + ) -> Result> { + // Build one forked context + future per branch. Node lookup and resume + // consumption happen up front so the futures borrow nothing mutable; each + // branch drives its handler through the node-retry policy (which also + // applies the per-node timeout), so a transient failure in one branch is + // retried without disturbing its siblings. + let mut futures = Vec::with_capacity(active.len()); + for (index, activation) in active.iter().enumerate() { + let node_id = &activation.node; + let node = self + .nodes + .get(node_id) + .ok_or_else(|| TinyAgentsError::MissingNode(node_id.to_string()))?; + + self.emit(GraphEvent::TaskScheduled { + node: node_id.clone(), + step, + }); + self.emit(GraphEvent::NodeStarted { + node: node_id.clone(), + step, + }); + + self.emit(GraphEvent::ContextForked { + node: node_id.clone(), + fork: index, + step, + }); + let fork = Some(ForkId::new(index, node_id.clone())); + let ctx = self.node_context( + node_id, + run_id, + thread_id, + step, + resume_map, + fork, + activation.send_arg.clone(), + root_run_id, + frames, + child_runs, + ); + let handler = node.handler.clone(); + let owned_node = node_id.clone(); + // Box each branch future behind a concrete `Send` bound. This keeps + // the `buffer_unordered` rolling window below (used for a + // `max_concurrency` bound) from requiring a higher-ranked `Send` + // proof over the borrowed recursion frames, which the compiler + // cannot discharge for the bare `async` blocks. + let fut: std::pin::Pin< + Box>> + Send + '_>, + > = Box::pin(async move { + self.run_node_with_retry(&owned_node, &handler, state, ctx, step) + .await + }); + futures.push(fut); + } + + // Drive branches to completion, bounding in-flight count when configured. + // With a bound, keep a rolling window of `limit` branches in flight + // instead of fixed `join_all` chunks. A chunked join runs each chunk to + // completion before starting the next, so a single slow branch + // head-of-line blocks the whole chunk; the rolling window starts a new + // branch as soon as *any* in-flight one finishes. `select_all` reports + // which pending future completed; a parallel index Vec maps it back to + // the branch's active-set position, so results are re-ordered into + // deterministic order for the fold below. + let results = match self.max_concurrency { + Some(limit) if limit < futures.len() => { + let total = futures.len(); + let mut slots: Vec>>> = + (0..total).map(|_| None).collect(); + let mut source = futures.into_iter().enumerate(); + let mut running = Vec::with_capacity(limit); + let mut running_index = Vec::with_capacity(limit); + for (index, fut) in source.by_ref().take(limit) { + running.push(fut); + running_index.push(index); + } + while !running.is_empty() { + let (result, completed, rest) = futures::future::select_all(running).await; + let index = running_index.remove(completed); + slots[index] = Some(result); + running = rest; + if let Some((index, fut)) = source.next() { + running.push(fut); + running_index.push(index); + } + } + slots + .into_iter() + .map(|slot| slot.expect("every branch produced a result")) + .collect::>() + } + _ => futures::future::join_all(futures).await, + }; + + // Fold in deterministic active-set index order. + let mut updates: Vec = Vec::new(); + let mut goto_map: HashMap> = HashMap::new(); + let mut interrupt: Option<(usize, Interrupt)> = None; + let mut failure: Option = None; + + for (index, (activation, result)) in active.iter().zip(results).enumerate() { + let node_id = &activation.node; + let result = match result { + Ok(result) => result, + Err(error) => { + self.emit(GraphEvent::NodeFailed { + node: node_id.clone(), + step, + error: error.to_string(), + }); + // The lowest-index failing branch is terminal: fold the + // lower-index successes (already applied above) and schedule + // their successors plus this branch and the rest for a + // resumable retry. + failure = Some(StepFailure { + failed_index: index, + error, + }); + break; + } + }; + + if let Some(found) = self.fold_result( + index, + node_id, + step, + result, + &mut updates, + &mut goto_map, + visited, + ) { + interrupt = Some(found); + break; + } + } + + Ok(StepRun { + updates, + goto_map, + interrupt, + failure, + }) + } + + /// Routes a set of completed activations into their successor activations. + /// + /// Honors per-activation command `goto` (keyed by active-set index), static + /// and conditional edges, barrier gating (a waiting node is held until every + /// required predecessor has arrived, accumulating into `barrier_arrivals` + /// across supersteps), and per-node dedup — while preserving each `Send` + /// packet's per-invocation argument. Emits a + /// [`GraphEvent::RouteSelected`] per selected edge. + /// + /// Shared by the normal step boundary (routes the whole active set) and the + /// interrupt/failure boundaries (route just the branches that completed + /// before the pause, so their successors are still scheduled on resume). + #[allow(clippy::too_many_arguments)] + async fn persist_checkpoint( + &self, + thread_id: &Option, + run_id: &RunId, + state: &State, + pending: &[Activation], + completed_tasks: &[NodeId], + interrupts: Vec, + barrier_arrivals: &HashMap>, + parent: Option, + step: usize, + source: &str, + recursion: &serde_json::Value, + child_runs: &serde_json::Value, + ) -> Result> { + let (Some(checkpointer), Some(thread)) = (&self.checkpointer, thread_id) else { + return Ok(None); + }; + let checkpoint_id = next_checkpoint_id(); + let checkpoint = Checkpoint { + thread_id: thread.to_string(), + checkpoint_id, + run_id: Some(run_id.to_string()), + parent_checkpoint_id: parent, + namespace: self.namespace.clone(), + state: state.clone(), + next_nodes: activation_nodes(pending), + completed_tasks: completed_tasks.to_vec(), + pending_writes: Vec::new(), + pending_activations: Some(pending.iter().map(PendingActivation::from).collect()), + barrier_arrivals: barriers_to_persisted(barrier_arrivals), + interrupts, + metadata: serde_json::json!({ + "source": source, + "step": step, + "recursion": recursion, + "child_runs": child_runs, + }), + }; + let id = checkpointer.put(checkpoint).await?; + self.emit(GraphEvent::CheckpointSaved { + checkpoint_id: id.clone(), + }); + Ok(Some(id)) + } + + fn base_status( + &self, + run_id: &RunId, + thread_id: &Option, + started_at: SystemTime, + ) -> GraphRunStatus { + let mut status = GraphRunStatus::new( + run_id.clone(), + self.graph_id.clone(), + ExecutionStatus::Running, + ); + status.thread_id = thread_id.clone(); + status.checkpoint_namespace = self.namespace.clone(); + status.started_at = started_at; + status.updated_at = SystemTime::now(); + status + } +} diff --git a/src/graph/compiled/mod.rs b/src/graph/compiled/mod.rs index 01ecb92..32aef11 100644 --- a/src/graph/compiled/mod.rs +++ b/src/graph/compiled/mod.rs @@ -64,20 +64,23 @@ //! [`CompiledGraph::update_state`] before resuming. Without a checkpointer the //! run aborts immediately, exactly as before. +mod executor; +mod routing; +mod state_api; mod types; pub use types::{CompiledGraph, GraphExecution, GraphInput, ResumeTarget, StateSnapshot}; use std::collections::{HashMap, HashSet}; use std::sync::Arc; -use std::sync::atomic::{AtomicU64, Ordering}; use std::time::{Duration, SystemTime}; use crate::graph::builder::{ Branch, BuilderNode, END, ForkId, NodeContext, NodeFuture, NodeHandler, NodeMeta, START, }; use crate::graph::checkpoint::{ - Checkpoint, CheckpointConfig, CheckpointTuple, Checkpointer, DurabilityMode, + BarrierArrivals, Checkpoint, CheckpointConfig, CheckpointTuple, Checkpointer, DurabilityMode, + PendingActivation, }; use crate::graph::command::{Command, Interrupt, NodeResult, RouteTarget}; use crate::graph::recursion::{ @@ -92,15 +95,18 @@ use crate::harness::ids::{ use crate::harness::retry::is_retryable; use crate::{Result, TinyAgentsError}; -static SEQ: AtomicU64 = AtomicU64::new(0); - -/// Returns a process-unique monotonic sequence number for id generation. -pub(crate) fn next_seq() -> u64 { - SEQ.fetch_add(1, Ordering::Relaxed) -} - -fn next_id(prefix: &str) -> String { - format!("{prefix}-{}", next_seq()) +/// Allocates a fresh checkpoint id (string form) that is collision-free across +/// process restarts. +/// +/// Delegates to [`crate::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 { + crate::harness::ids::new_checkpoint_id() + .as_str() + .to_string() } /// Projects a loaded [`CheckpointTuple`] onto a [`StateSnapshot`] for the @@ -133,8 +139,13 @@ struct StepRun { /// Branch updates in deterministic active-set index order. updates: Vec, /// Explicit routing (plain `goto` nodes and/or [`Send`] packets) keyed by the - /// node that produced it. - goto_map: HashMap>, + /// producing branch's active-set index. + /// + /// Keyed by index rather than node id so repeated [`Send`] activations of + /// the *same* node within a step (map-reduce fanout) each keep their own + /// [`Command::goto`] — a node-keyed map would let a later activation's + /// command clobber an earlier one's routing. + goto_map: HashMap>, /// The lowest-index branch interrupt, if any (its active-set index + value). interrupt: Option<(usize, Interrupt)>, /// A node-handler failure that survived the node-retry policy, if any. When @@ -147,14 +158,12 @@ struct StepRun { /// 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 { - /// The node whose handler ultimately failed (after any retries). - failed_node: NodeId, - /// Nodes that completed successfully before the failure, in active-set order - /// (recorded as the checkpoint's completed tasks). - completed: Vec, - /// Nodes to re-run on resume: the failed node first, then the not-yet-folded - /// members of the step (recorded as the checkpoint's next nodes). - pending: Vec, + /// 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, } @@ -181,6 +190,42 @@ impl Activation { } } +impl From<&Activation> for PendingActivation { + fn from(a: &Activation) -> Self { + PendingActivation { + node: a.node.clone(), + send_arg: a.send_arg.clone(), + } + } +} + +impl From<&PendingActivation> for Activation { + fn from(p: &PendingActivation) -> Self { + Activation { + node: p.node.clone(), + send_arg: p.send_arg.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() +} + /// Maps an [`Activation`] slice to its node ids (for events, status, and /// checkpoint records, which are node-keyed). fn activation_nodes(active: &[Activation]) -> Vec { @@ -366,1417 +411,18 @@ impl CompiledGraph { fn emit(&self, event: GraphEvent) { if let Some(sink) = &self.event_sink { - sink.emit(event); - } - } -} - -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( - state, - vec![Activation::node(self.entry.clone())], - None, - HashMap::new(), - ) - .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::graph::NodeContext::send_arg). - pub async fn run_with_inputs( - &self, - state: State, - inputs: impl IntoIterator, - ) -> Result> { - let active = self.initial_inputs(inputs)?; - self.execute(state, active, None, HashMap::new()).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( - state, - vec![Activation::node(self.entry.clone())], - Some(thread_id.into()), - HashMap::new(), - ) - .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(state, active, Some(thread_id.into()), HashMap::new()) - .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 - } - - /// 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 - } - - /// 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> { - let checkpointer = self - .checkpointer - .as_ref() - .ok_or_else(|| TinyAgentsError::Resume("no checkpointer configured".to_string()))?; - let thread_id = thread_id.into(); - - let checkpoint_id = match &target { - ResumeTarget::Latest => None, - ResumeTarget::Checkpoint(id) => Some(id.as_str()), - }; - let checkpoint = checkpointer - .get(thread_id.as_str(), checkpoint_id) - .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}`" - )), - })?; - self.emit(GraphEvent::CheckpointSaved { - checkpoint_id: CheckpointId::new(checkpoint.checkpoint_id.clone()), - }); - - let active = checkpoint.next_nodes.clone(); - if active.is_empty() { - return Err(TinyAgentsError::Resume( - "checkpoint has no pending nodes to resume".to_string(), - )); - } - - let mut resume_map = HashMap::new(); - if let Some(value) = command.resume { - for node in &active { - resume_map.insert(node.clone(), value.clone()); - } - } - - self.execute( - checkpoint.state, - active.into_iter().map(Activation::node).collect(), - Some(thread_id), - resume_map, - ) - .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, - }); - } - 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. - 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(), - } - } - - /// Loads a [`StateSnapshot`] for a thread. - /// - /// With `checkpoint_id == None` the thread's latest checkpoint is returned; - /// otherwise the specific checkpoint is addressed. Returns `Ok(None)` when no - /// matching checkpoint exists. Requires a configured checkpointer. - 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()?; - let mut out: Vec> = Vec::new(); - let mut cursor: Option = None; - loop { - if let Some(limit) = limit - && out.len() >= limit - { - break; - } - let config = self.config_for(thread_id, cursor.as_deref()); - let Some(tuple) = checkpointer.get_tuple(config).await? else { - break; - }; - let parent = tuple.checkpoint.parent_checkpoint_id.clone(); - out.push(snapshot_from_tuple(tuple)); - match parent { - Some(parent) => cursor = Some(parent), - None => break, - } - } - Ok(out) - } - - /// 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::graph::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 and the new checkpoint's pending nodes become that - /// node's routing successors (so a subsequent resume continues from after the - /// attributed node). 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 - && !self.nodes.contains_key(node) - { - return Err(TinyAgentsError::MissingNode(node.to_string())); - } - - let base = checkpointer.get(thread_id, None).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(); - let new_state = self.reducer.apply(base.state, update)?; - - // Pending nodes: the attributed node's successors, or the inherited set. - let next_nodes: Vec = match &as_node { - Some(node) => self - .route(node, &HashMap::new(), &new_state)? - .into_iter() - .map(|t| t.node().clone()) - .filter(|n| n.as_str() != END) - .collect(), - None => base.next_nodes.clone(), - }; - let completed_tasks: Vec = as_node.iter().cloned().collect(); - - let checkpoint_id = next_id("ckpt"); - let config = self.config_for(thread_id, Some(&checkpoint_id)); - let checkpoint = Checkpoint { - thread_id: thread_id.to_string(), - checkpoint_id, - run_id: None, - parent_checkpoint_id: Some(parent_id), - namespace: self.namespace.clone(), - state: new_state, - next_nodes, - completed_tasks, - pending_writes: Vec::new(), - interrupts: Vec::new(), - metadata: serde_json::json!({ "source": "update", "step": parent_step + 1 }), - }; - let id = checkpointer.put(checkpoint).await?; - self.emit(GraphEvent::CheckpointSaved { checkpoint_id: id }); - 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(source_thread, source_checkpoint_id) - .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_id("ckpt"); - let config = self.config_for(target_thread, Some(&checkpoint_id)); - let forked = Checkpoint { - thread_id: target_thread.to_string(), - checkpoint_id, - run_id: None, - parent_checkpoint_id: None, - namespace: source.namespace.clone(), - state: source.state.clone(), - next_nodes: source.next_nodes.clone(), - completed_tasks: source.completed_tasks.clone(), - pending_writes: source.pending_writes.clone(), - interrupts: source.interrupts.clone(), - metadata: serde_json::json!({ "source": "fork", "step": step }), - }; - let id = checkpointer.put(forked).await?; - self.emit(GraphEvent::CheckpointSaved { checkpoint_id: id }); - Ok(config) - } - - async fn execute( - &self, - state: State, - initial_active: Vec, - thread_id: Option, - resume_map: HashMap, - ) -> Result> { - let run_id = RunId::new(next_id("run")); - // When a durable journal is configured, run against a clone whose event - // sink wraps every emitted event into a `GraphObservation` and appends - // it (while still forwarding to any pre-existing live sink). The journal - // sink carries this graph's checkpoint namespace so subgraph runs record - // their nested path. Default (no journal) leaves `self` untouched. - if self.journal.is_some() { - let this = self.clone_with_journal_sink(&run_id, &thread_id); - this.execute_run(run_id, state, initial_active, thread_id, resume_map) - .await - } else { - self.execute_run(run_id, state, initial_active, thread_id, resume_map) - .await - } - } - - /// 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::graph::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 - } - - /// Best-effort status write; never aborts the run on a status-store error. - async fn save_status(&self, status: GraphRunStatus) { - if let Some(store) = &self.status_store { - let _ = store.put_status(status).await; - } - } - - async fn execute_run( - &self, - run_id: RunId, - mut state: State, - initial_active: Vec, - thread_id: Option, - mut resume_map: HashMap, - ) -> Result> { - let started_at = SystemTime::now(); - let mut visited: Vec = Vec::new(); - let mut steps = 0usize; - let mut last_checkpoint: Option = None; - let mut parent_checkpoint: Option = None; - - // Build this run's recursion stack from the inherited parent frames and - // push the frame for this graph call. A push that would exceed - // `max_depth` fails the run with a clear recursion error before any - // node executes. Graph-call depth (the stack) is tracked separately - // from node-loop visits (`node_visits`, below). - let mut recursion = - RecursionStack::with_frames(self.recursion_frames.clone(), self.recursion_policy); - // Run lineage: the root is the first inherited frame's run (the top of - // the recursion tree) or this run when top-level; the parent is the - // enclosing run, if any. - let root_run_id = self - .recursion_frames - .first() - .map(|f| f.run_id.clone()) - .unwrap_or_else(|| run_id.clone()); - let parent_run_id = self.recursion_frames.last().map(|f| f.run_id.clone()); - let this_frame = RecursionFrame { - graph_id: self.graph_id.clone(), - node_id: self.recursion_node.clone(), - run_id: run_id.clone(), - task_id: None, - namespace: self.namespace.clone(), - depth: recursion.depth(), - parent: parent_run_id.clone(), - }; - if let Err(err) = recursion.push(this_frame) { - self.emit(GraphEvent::RunStarted { - run_id: run_id.clone(), - }); - self.fail_run(&run_id, &thread_id, started_at, steps, &err, None) - .await; - return Err(err); - } - // Serialized once per run for embedding in every checkpoint's metadata. - let recursion_meta = - serde_json::to_value(recursion.frames()).unwrap_or(serde_json::Value::Null); - // The live frame stack handed to node contexts so a subgraph node can - // seed an embedded child with this run's recursion path, plus the - // per-run sink the node reports its spawned child run into. - let live_frames = recursion.frames().to_vec(); - let child_sink = ChildRunSink::new(); - // Accumulates every child run spawned across all supersteps for the - // final `GraphExecution::child_runs`. - let mut all_child_runs: Vec = Vec::new(); - // Per-node activation counts for `max_visits_per_node` enforcement. - let mut node_visits: HashMap = HashMap::new(); - let mut active = initial_active; - // Barrier/waiting-edge arrivals accumulate across supersteps: a waiting - // node only activates once every required predecessor has arrived. - let mut barrier_arrivals: HashMap> = HashMap::new(); - - self.emit(GraphEvent::RunStarted { - run_id: run_id.clone(), - }); - // Surface this run's recursion depth so observers can attribute nested - // runs without reconstructing the tree from logs. - self.emit(GraphEvent::RecursionDepthChanged { - depth: recursion.depth(), - }); - // Record the run as live before the first superstep is scheduled. - let mut running = self.base_status(&run_id, &thread_id, started_at); - running.active_nodes = activation_nodes(&active); - self.save_status(running).await; - - while !active.is_empty() { - // The effective step cap is the smaller of the builder's recursion - // limit and the policy's `max_total_steps`, so a policy never - // loosens an existing limit. Both surface a `RecursionLimit`. - let step_limit = self - .recursion_limit - .min(self.recursion_policy.max_total_steps); - if steps >= step_limit { - let err = TinyAgentsError::RecursionLimit(step_limit); - self.fail_run(&run_id, &thread_id, started_at, steps, &err, None) - .await; - return Err(err); - } - // Node-loop recursion: enforce `max_visits_per_node` per activation. - for activation in &active { - if let Err(err) = recursion.record_node_visit(&mut node_visits, &activation.node) { - self.fail_run(&run_id, &thread_id, started_at, steps, &err, None) - .await; - return Err(err); - } - } - steps += 1; - self.emit(GraphEvent::StepStarted { - step: steps, - active: activation_nodes(&active), - }); - - let run_result = if self.parallel && active.len() > 1 { - self.run_active_parallel( - &active, - &state, - &run_id, - &thread_id, - steps, - &mut resume_map, - &mut visited, - &root_run_id, - &live_frames, - &child_sink, - ) - .await - } else { - self.run_active_sequential( - &active, - &state, - &run_id, - &thread_id, - steps, - &mut resume_map, - &mut visited, - &root_run_id, - &live_frames, - &child_sink, - ) - .await - }; - let StepRun { - updates, - goto_map, - interrupt, - failure, - } = match run_result { - Ok(step_run) => step_run, - Err(err) => { - self.fail_run(&run_id, &thread_id, started_at, steps, &err, None) - .await; - return Err(err); - } - }; - - // Apply collected updates through the reducer at the boundary. - for update in updates { - state = self.reducer.apply(state, update)?; - } - - // Collect any child runs spawned by subgraph nodes this step. They - // are embedded into this boundary's checkpoint metadata (keyed by - // node) and accumulated onto the final `GraphExecution`. - let step_child_runs = child_sink.drain(); - all_child_runs.extend(step_child_runs.iter().cloned()); - let child_runs_meta = - serde_json::to_value(&step_child_runs).unwrap_or(serde_json::Value::Null); - - // Node-handler failure (survived any node-retry policy): the updates - // of the branches that completed before it are already folded into - // `state` above, so persist a resumable failure-boundary checkpoint - // scheduling the failed node (and the not-yet-run tail) for a later - // `resume`/`retry`, record a `Failed` status carrying the error and - // that checkpoint, and abort. Without a checkpointer/thread the - // checkpoint is a no-op and the run aborts exactly as before. - if let Some(fail) = failure { - let StepFailure { - failed_node, - completed, - pending, - error, - } = fail; - let checkpoint_id = self - .persist_failure_checkpoint( - &thread_id, - &run_id, - &state, - &pending, - &completed, - parent_checkpoint.clone(), - steps, - &failed_node, - &error, - &recursion_meta, - &child_runs_meta, - ) - .await?; - self.fail_run( - &run_id, - &thread_id, - started_at, - steps, - &error, - checkpoint_id, - ) - .await; - return Err(error); - } - - // Interrupt: persist a checkpoint whose next nodes are the - // not-yet-completed members of this step (interrupted node first), - // then return control to the caller. - if let Some((index, emitted)) = interrupt { - if let Err(err) = self.require_interrupt_durability(&thread_id) { - self.fail_run(&run_id, &thread_id, started_at, steps, &err, None) - .await; - return Err(err); - } - let pending: Vec = activation_nodes(&active[index..]); - let interrupt_id = InterruptId::new(emitted.id.clone()); - let checkpoint_id = self - .persist_checkpoint( - &thread_id, - &run_id, - &state, - &pending, - &activation_nodes(&active[..index]), - vec![emitted.clone()], - parent_checkpoint.clone(), - steps, - "loop", - &recursion_meta, - &child_runs_meta, - ) - .await?; - - let mut status = self.base_status(&run_id, &thread_id, started_at); - status.status = ExecutionStatus::Interrupted; - status.current_step = steps; - status.active_nodes = pending; - status.pending_interrupts = vec![interrupt_id]; - status.checkpoint_id = checkpoint_id.clone(); - self.save_status(status.clone()).await; - - return Ok(GraphExecution { - state, - run_id: run_id.clone(), - graph_id: self.graph_id.clone(), - root_run_id: root_run_id.clone(), - parent_run_id: parent_run_id.clone(), - child_runs: all_child_runs, - visited, - steps, - interrupts: vec![emitted], - status, - checkpoint_id, - }); - } - - // Select the next active set from commands or static/conditional - // edges, evaluated against the freshly-committed state. - let completed = active.clone(); - let mut next: Vec = Vec::new(); - let mut next_seen: HashSet = HashSet::new(); - for activation in &completed { - let node_id = &activation.node; - let targets = self.route(node_id, &goto_map, &state)?; - for target in targets { - let tnode = target.node().clone(); - if tnode.as_str() == END { - continue; - } - self.emit(GraphEvent::RouteSelected { - node: node_id.clone(), - target: tnode.clone(), - }); - // 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(); - if send_arg.is_some() { - next.push(Activation { - node: tnode, - send_arg, - }); - } else if next_seen.insert(tnode.clone()) { - next.push(Activation { - node: tnode, - send_arg: None, - }); - } - } - } - - // Persist a boundary checkpoint (node-keyed records). Under - // `Exit` durability only the terminal boundary (the step that - // empties the active set) is written; `Sync`/`Async` persist every - // boundary. - let completed_nodes = activation_nodes(&completed); - let next_nodes = activation_nodes(&next); - let persist_now = match self.durability { - DurabilityMode::Exit => next.is_empty(), - DurabilityMode::Sync | DurabilityMode::Async => true, - }; - let checkpoint_id = if persist_now { - self.persist_checkpoint( - &thread_id, - &run_id, - &state, - &next_nodes, - &completed_nodes, - Vec::new(), - parent_checkpoint.clone(), - steps, - "loop", - &recursion_meta, - &child_runs_meta, - ) - .await? - } else { - None - }; - if let Some(id) = &checkpoint_id { - last_checkpoint = Some(id.clone()); - parent_checkpoint = Some(id.to_string()); - } - - self.emit(GraphEvent::StepCompleted { step: steps }); - active = next; - } - - let mut status = self.base_status(&run_id, &thread_id, started_at); - status.status = ExecutionStatus::Completed; - status.current_step = steps; - status.checkpoint_id = last_checkpoint.clone(); - status.ended_at = Some(SystemTime::now()); - self.save_status(status.clone()).await; - self.emit(GraphEvent::RunCompleted { - run_id: run_id.clone(), - steps, - }); - - Ok(GraphExecution { - state, - run_id: run_id.clone(), - graph_id: self.graph_id.clone(), - root_run_id, - parent_run_id, - child_runs: all_child_runs, - visited, - steps, - interrupts: Vec::new(), - status, - checkpoint_id: last_checkpoint, - }) - } - - /// Emits a [`GraphEvent::RunFailed`] and records a terminal `Failed` status - /// for a run that aborted with `err`. - /// - /// `checkpoint_id` is the resumable failure-boundary checkpoint when the run - /// left one (a node-handler failure on a checkpointed thread), or `None` for - /// a structural/non-resumable abort. When present it is recorded on the - /// status so an observer can locate the checkpoint to `resume`/`retry` from. - async fn fail_run( - &self, - run_id: &RunId, - thread_id: &Option, - started_at: SystemTime, - steps: usize, - err: &TinyAgentsError, - checkpoint_id: Option, - ) { - self.emit(GraphEvent::RunFailed { - run_id: run_id.clone(), - error: err.to_string(), - }); - let mut status = self.base_status(run_id, thread_id, started_at); - status.status = ExecutionStatus::Failed; - status.current_step = steps; - status.ended_at = Some(SystemTime::now()); - status.error = Some(err.to_string()); - status.checkpoint_id = checkpoint_id; - self.save_status(status).await; - } - - /// 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. - #[allow(clippy::too_many_arguments)] - async fn persist_failure_checkpoint( - &self, - thread_id: &Option, - run_id: &RunId, - state: &State, - next_nodes: &[NodeId], - completed_tasks: &[NodeId], - parent: Option, - step: usize, - failed_node: &NodeId, - error: &TinyAgentsError, - recursion: &serde_json::Value, - child_runs: &serde_json::Value, - ) -> Result> { - let (Some(checkpointer), Some(thread)) = (&self.checkpointer, thread_id) else { - return Ok(None); - }; - let checkpoint = Checkpoint { - thread_id: thread.to_string(), - checkpoint_id: next_id("ckpt"), - run_id: Some(run_id.to_string()), - parent_checkpoint_id: parent, - namespace: self.namespace.clone(), - state: state.clone(), - next_nodes: next_nodes.to_vec(), - completed_tasks: completed_tasks.to_vec(), - pending_writes: Vec::new(), - interrupts: Vec::new(), - metadata: serde_json::json!({ - "source": "loop", - "step": step, - "recursion": recursion, - "child_runs": child_runs, - "failed_node": failed_node.as_str(), - "error": error.to_string(), - }), - }; - let id = checkpointer.put(checkpoint).await?; - self.emit(GraphEvent::CheckpointSaved { - checkpoint_id: id.clone(), - }); - Ok(Some(id)) - } - - /// Builds the per-task [`NodeContext`] for `node_id` at the given branch. - /// - /// `fork` carries the branch identity in a concurrent step (`None` in - /// sequential mode or single-node steps). The resume value for the node is - /// consumed from `resume_map`. - #[allow(clippy::too_many_arguments)] - fn node_context( - &self, - node_id: &NodeId, - run_id: &RunId, - thread_id: &Option, - step: usize, - resume_map: &mut HashMap, - fork: Option, - send_arg: Option, - root_run_id: &RunId, - frames: &[RecursionFrame], - child_runs: &ChildRunSink, - ) -> NodeContext { - NodeContext { - node_id: node_id.clone(), - run_id: run_id.clone(), - thread_id: thread_id.clone(), - step, - resume: resume_map.remove(node_id), - fork, - send_arg, - root_run_id: Some(root_run_id.clone()), - recursion_frames: frames.to_vec(), - child_runs: Some(child_runs.clone()), - } - } - - /// Wraps a node future in the configured per-node timeout (if any), mapping - /// an elapsed deadline onto [`TinyAgentsError::Timeout`]. - async fn run_node_future( - &self, - node_id: &NodeId, - fut: NodeFuture, - ) -> Result> { - match self.node_timeout { - Some(timeout) => match tokio::time::timeout(timeout, fut).await { - Ok(result) => result, - Err(_) => Err(TinyAgentsError::Timeout(format!( - "node `{node_id}` exceeded its {timeout:?} timeout" - ))), - }, - None => fut.await, - } - } - - /// Runs one node handler under the graph's node-retry policy. - /// - /// Builds a fresh handler future (and re-clones the context) for each - /// attempt, so a retried node re-runs from its start — matching the durable - /// execution model, where a node is never suspended mid-flight. On a - /// [retryable][crate::harness::retry::is_retryable] error, when a - /// [`RetryPolicy`](crate::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: &State, - ctx: NodeContext, - step: usize, - ) -> Result> { - let mut attempt = 0usize; - loop { - let fut = handler(state.clone(), ctx.clone()); - match self.run_node_future(node_id, fut).await { - Ok(result) => return Ok(result), - Err(error) => { - let retry = self - .node_retry - .as_ref() - .filter(|policy| policy.should_retry(attempt) && is_retryable(&error)); - let Some(policy) = retry else { - return Err(error); - }; - attempt += 1; - self.emit(GraphEvent::NodeRetryScheduled { - node: node_id.clone(), - step, - attempt, - }); - policy.sleep_backoff(attempt).await; - } - } - } - } - - /// Folds a single successful branch result into the step accumulators. - /// - /// Pushes the node to `visited`, records updates/goto, emits the matching - /// events, and returns the interrupt (with its branch index) when the branch - /// paused. Shared by the sequential and parallel run paths so both fold - /// results identically; only the *running* of handlers differs. - #[allow(clippy::too_many_arguments)] - fn fold_result( - &self, - index: usize, - node_id: &NodeId, - step: usize, - result: NodeResult, - updates: &mut Vec, - goto_map: &mut HashMap>, - visited: &mut Vec, - ) -> Option<(usize, Interrupt)> { - visited.push(node_id.clone()); - match result { - NodeResult::Update(update) => { - updates.push(update); - self.emit(GraphEvent::StateUpdated { - node: node_id.clone(), - step, - }); - } - NodeResult::Command(command) => { - if let Some(update) = command.update { - updates.push(update); - self.emit(GraphEvent::StateUpdated { - node: node_id.clone(), - step, - }); - } - if !command.goto.is_empty() { - goto_map.insert(node_id.clone(), command.goto); - } - } - NodeResult::Interrupt(emitted) => { - self.emit(GraphEvent::InterruptEmitted { - interrupt: emitted.clone(), - }); - return Some((index, emitted)); - } - } - self.emit(GraphEvent::NodeCompleted { - node: node_id.clone(), - step, - }); - None - } - - /// Runs the active node set one node at a time (default behavior). - /// - /// Short-circuits on the first error (run aborts) or interrupt (later nodes - /// in the step are not started), exactly preserving milestone-1 semantics. - #[allow(clippy::too_many_arguments)] - async fn run_active_sequential( - &self, - active: &[Activation], - state: &State, - run_id: &RunId, - thread_id: &Option, - step: usize, - resume_map: &mut HashMap, - visited: &mut Vec, - root_run_id: &RunId, - frames: &[RecursionFrame], - child_runs: &ChildRunSink, - ) -> Result> { - let mut updates: Vec = Vec::new(); - let mut goto_map: HashMap> = HashMap::new(); - let mut interrupt: Option<(usize, Interrupt)> = None; - let mut failure: Option = None; - - for (index, activation) in active.iter().enumerate() { - let node_id = &activation.node; - let node = self - .nodes - .get(node_id) - .ok_or_else(|| TinyAgentsError::MissingNode(node_id.to_string()))?; - - self.emit(GraphEvent::TaskScheduled { - node: node_id.clone(), - step, - }); - self.emit(GraphEvent::NodeStarted { - node: node_id.clone(), - step, - }); - - let ctx = self.node_context( - node_id, - run_id, - thread_id, - step, - resume_map, - None, - activation.send_arg.clone(), - root_run_id, - frames, - child_runs, + // 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 { .. } ); - let result = match self - .run_node_with_retry(node_id, &node.handler, state, ctx, step) - .await - { - Ok(result) => result, - Err(error) => { - self.emit(GraphEvent::NodeFailed { - node: node_id.clone(), - step, - error: error.to_string(), - }); - // Preserve the progress of the branches that already ran: - // record them as completed and schedule this node plus the - // not-yet-run tail for a resumable retry. - failure = Some(StepFailure { - failed_node: node_id.clone(), - completed: activation_nodes(&active[..index]), - pending: activation_nodes(&active[index..]), - error, - }); - break; - } - }; - - if let Some(found) = self.fold_result( - index, - node_id, - step, - result, - &mut updates, - &mut goto_map, - visited, - ) { - interrupt = Some(found); - break; - } - } - - Ok(StepRun { - updates, - goto_map, - interrupt, - failure, - }) - } - - /// Runs the active node set concurrently (opt-in via `with_parallel`). - /// - /// Each branch executes on its own cloned `State` snapshot and a distinct - /// [`ForkId`], optionally with the [`Send`] argument that scheduled it. With - /// no `max_concurrency` bound every branch starts before any is awaited and - /// all are driven via [`futures::future::join_all`]; with a bound the active - /// set is run in chunks of at most that many futures, so at most that many - /// node handlers are in flight at once. Results are folded in active-set - /// index order — the reducer is the join/fan-in — so the merged state is - /// reproducible regardless of completion order. The lowest-index branch that - /// errors or interrupts is the step's terminal outcome; lower-index - /// successful branches still contribute their updates. - #[allow(clippy::too_many_arguments)] - async fn run_active_parallel( - &self, - active: &[Activation], - state: &State, - run_id: &RunId, - thread_id: &Option, - step: usize, - resume_map: &mut HashMap, - visited: &mut Vec, - root_run_id: &RunId, - frames: &[RecursionFrame], - child_runs: &ChildRunSink, - ) -> Result> { - // Build one forked context + future per branch. Node lookup and resume - // consumption happen up front so the futures borrow nothing mutable; each - // branch drives its handler through the node-retry policy (which also - // applies the per-node timeout), so a transient failure in one branch is - // retried without disturbing its siblings. - let mut futures = Vec::with_capacity(active.len()); - for (index, activation) in active.iter().enumerate() { - let node_id = &activation.node; - let node = self - .nodes - .get(node_id) - .ok_or_else(|| TinyAgentsError::MissingNode(node_id.to_string()))?; - - self.emit(GraphEvent::TaskScheduled { - node: node_id.clone(), - step, - }); - self.emit(GraphEvent::NodeStarted { - node: node_id.clone(), - step, - }); - - self.emit(GraphEvent::ContextForked { - node: node_id.clone(), - fork: index, - step, - }); - let fork = Some(ForkId::new(index, node_id.clone())); - let ctx = self.node_context( - node_id, - run_id, - thread_id, - step, - resume_map, - fork, - activation.send_arg.clone(), - root_run_id, - frames, - child_runs, - ); - let handler = node.handler.clone(); - let owned_node = node_id.clone(); - futures.push(async move { - self.run_node_with_retry(&owned_node, &handler, state, ctx, step) - .await - }); - } - - // Drive branches to completion, bounding in-flight count when configured. - let results = match self.max_concurrency { - Some(limit) if limit < futures.len() => { - let mut out = Vec::with_capacity(futures.len()); - let mut iter = futures.into_iter(); - loop { - let chunk: Vec<_> = iter.by_ref().take(limit).collect(); - if chunk.is_empty() { - break; - } - out.extend(futures::future::join_all(chunk).await); - } - out - } - _ => futures::future::join_all(futures).await, - }; - - // Fold in deterministic active-set index order. - let mut updates: Vec = Vec::new(); - let mut goto_map: HashMap> = HashMap::new(); - let mut interrupt: Option<(usize, Interrupt)> = None; - let mut failure: Option = None; - - for (index, (activation, result)) in active.iter().zip(results).enumerate() { - let node_id = &activation.node; - let result = match result { - Ok(result) => result, - Err(error) => { - self.emit(GraphEvent::NodeFailed { - node: node_id.clone(), - step, - error: error.to_string(), - }); - // The lowest-index failing branch is terminal: fold the - // lower-index successes (already applied above) and schedule - // this branch plus the rest for a resumable retry. - failure = Some(StepFailure { - failed_node: node_id.clone(), - completed: activation_nodes(&active[..index]), - pending: activation_nodes(&active[index..]), - error, - }); - break; - } - }; - - if let Some(found) = self.fold_result( - index, - node_id, - step, - result, - &mut updates, - &mut goto_map, - visited, - ) { - interrupt = Some(found); - break; - } - } - - Ok(StepRun { - updates, - goto_map, - interrupt, - failure, - }) - } - - /// 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. - fn route( - &self, - node_id: &NodeId, - goto_map: &HashMap>, - state: &State, - ) -> Result> { - if let Some(targets) = goto_map.get(node_id) { - self.validate_route_targets(node_id, targets)?; - return Ok(targets.clone()); - } - if let Some(target) = self.edges.get(node_id) { - return Ok(vec![RouteTarget::Node(target.clone())]); - } - if let Some(branch) = self.branches.get(node_id) { - let route = (branch.router)(state); - 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())); + sink.emit(event); + if terminal { + sink.flush(); } } - Ok(()) - } - - 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(()) - } - - #[allow(clippy::too_many_arguments)] - async fn persist_checkpoint( - &self, - thread_id: &Option, - run_id: &RunId, - state: &State, - next_nodes: &[NodeId], - completed_tasks: &[NodeId], - interrupts: Vec, - parent: Option, - step: usize, - source: &str, - recursion: &serde_json::Value, - child_runs: &serde_json::Value, - ) -> Result> { - let (Some(checkpointer), Some(thread)) = (&self.checkpointer, thread_id) else { - return Ok(None); - }; - let checkpoint_id = next_id("ckpt"); - let checkpoint = Checkpoint { - thread_id: thread.to_string(), - checkpoint_id, - run_id: Some(run_id.to_string()), - parent_checkpoint_id: parent, - namespace: self.namespace.clone(), - state: state.clone(), - next_nodes: next_nodes.to_vec(), - completed_tasks: completed_tasks.to_vec(), - pending_writes: Vec::new(), - interrupts, - metadata: serde_json::json!({ - "source": source, - "step": step, - "recursion": recursion, - "child_runs": child_runs, - }), - }; - let id = checkpointer.put(checkpoint).await?; - self.emit(GraphEvent::CheckpointSaved { - checkpoint_id: id.clone(), - }); - Ok(Some(id)) - } - - fn base_status( - &self, - run_id: &RunId, - thread_id: &Option, - started_at: SystemTime, - ) -> GraphRunStatus { - let mut status = GraphRunStatus::new( - run_id.clone(), - self.graph_id.clone(), - ExecutionStatus::Running, - ); - status.thread_id = thread_id.clone(); - status.checkpoint_namespace = self.namespace.clone(); - status.started_at = started_at; - status.updated_at = SystemTime::now(); - status } } diff --git a/src/graph/compiled/routing.rs b/src/graph/compiled/routing.rs new file mode 100644 index 0000000..bf5062b --- /dev/null +++ b/src/graph/compiled/routing.rs @@ -0,0 +1,131 @@ +//! 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, +{ + pub(super) fn route_completed( + &self, + completed: &[Activation], + goto_map: &HashMap>, + state: &State, + barrier_arrivals: &mut HashMap>, + ) -> Result> { + let mut next: Vec = Vec::new(); + let mut next_seen: HashSet = HashSet::new(); + for (index, activation) in completed.iter().enumerate() { + let node_id = &activation.node; + let targets = self.route(node_id, goto_map.get(&index).map(Vec::as_slice), state)?; + for target in targets { + let tnode = target.node().clone(); + if tnode.as_str() == END { + continue; + } + self.emit(GraphEvent::RouteSelected { + node: node_id.clone(), + target: tnode.clone(), + }); + // 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(); + if send_arg.is_some() { + next.push(Activation { + node: tnode, + send_arg, + }); + } else if next_seen.insert(tnode.clone()) { + next.push(Activation { + node: tnode, + send_arg: None, + }); + } + } + } + Ok(next) + } + + /// 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(target) = self.edges.get(node_id) { + return Ok(vec![RouteTarget::Node(target.clone())]); + } + if let Some(branch) = self.branches.get(node_id) { + let route = (branch.router)(state); + 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(()) + } + + 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/src/graph/compiled/state_api.rs b/src/graph/compiled/state_api.rs new file mode 100644 index 0000000..de40080 --- /dev/null +++ b/src/graph/compiled/state_api.rs @@ -0,0 +1,224 @@ +//! 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::graph::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 and the new checkpoint's pending nodes become that + /// node's routing successors (so a subsequent resume continues from after the + /// attributed node). A command node cannot be used as `as_node` (it routes + /// dynamically and has no static successors); doing so returns + /// [`TinyAgentsError::Graph`] rather than silently producing a non-resumable + /// checkpoint. With `as_node == None` the latest pending node set is + /// 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(); + let new_state = self.reducer.apply(base.state, update)?; + + // Pending nodes: the attributed node's successors, or the inherited set. + let next_nodes: Vec = match &as_node { + Some(node) => self + .route(node, None, &new_state)? + .into_iter() + .map(|t| t.node().clone()) + .filter(|n| n.as_str() != END) + .collect(), + None => base.next_nodes.clone(), + }; + let completed_tasks: Vec = as_node.iter().cloned().collect(); + // With `as_node`, pending becomes that node's (plain) successors, so no + // send args carry over; without it, inherit the base checkpoint's + // pending activations verbatim so any pending `Send` args survive. + let pending_activations = match &as_node { + Some(_) => None, + None => base.pending_activations.clone(), + }; + // Manual writes preserve any accumulated barrier arrivals. + let barrier_arrivals = base.barrier_arrivals.clone(); + + let checkpoint_id = next_checkpoint_id(); + let config = self.config_for(thread_id, Some(&checkpoint_id)); + let checkpoint = Checkpoint { + thread_id: thread_id.to_string(), + checkpoint_id, + run_id: None, + parent_checkpoint_id: Some(parent_id), + namespace: self.namespace.clone(), + state: new_state, + next_nodes, + completed_tasks, + pending_writes: Vec::new(), + interrupts: Vec::new(), + pending_activations, + barrier_arrivals, + metadata: serde_json::json!({ "source": "update", "step": parent_step + 1 }), + }; + let id = checkpointer.put(checkpoint).await?; + self.emit(GraphEvent::CheckpointSaved { checkpoint_id: id }); + 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)); + let forked = Checkpoint { + thread_id: target_thread.to_string(), + checkpoint_id, + run_id: None, + parent_checkpoint_id: None, + namespace: source.namespace.clone(), + state: source.state.clone(), + next_nodes: source.next_nodes.clone(), + completed_tasks: source.completed_tasks.clone(), + pending_writes: source.pending_writes.clone(), + interrupts: source.interrupts.clone(), + pending_activations: source.pending_activations.clone(), + barrier_arrivals: source.barrier_arrivals.clone(), + metadata: serde_json::json!({ "source": "fork", "step": step }), + }; + let id = checkpointer.put(forked).await?; + self.emit(GraphEvent::CheckpointSaved { checkpoint_id: id }); + Ok(config) + } +} diff --git a/src/graph/compiled/test.rs b/src/graph/compiled/test.rs index 27a10c8..6aeff47 100644 --- a/src/graph/compiled/test.rs +++ b/src/graph/compiled/test.rs @@ -12,7 +12,7 @@ use crate::harness::ids::ExecutionStatus; use crate::harness::retry::RetryPolicy; use serde_json::json; use std::sync::Arc; -use std::sync::atomic::{AtomicUsize, Ordering as AtomicOrdering}; +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering as AtomicOrdering}; use std::time::Duration; #[derive(Clone, Debug, PartialEq)] @@ -329,6 +329,133 @@ async fn interrupt_then_resume_reruns_node() { 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() @@ -537,6 +664,43 @@ async fn update_state_as_node_sets_successor_pending_nodes() { ); } +#[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(); +} + #[tokio::test] async fn bulk_update_state_applies_successive_updates() { use crate::graph::CheckpointSource; @@ -879,6 +1043,239 @@ async fn parallel_interrupt_pauses_at_lowest_index_branch() { ); } +#[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); +} + +#[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 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(); @@ -952,6 +1349,69 @@ async fn send_fanout_delivers_distinct_args_to_parallel_branches() { 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() @@ -1299,6 +1759,89 @@ async fn max_concurrency_bounds_in_flight_branches() { 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)] diff --git a/src/graph/goals/README.md b/src/graph/goals/README.md index 3504c94..6d2700e 100644 --- a/src/graph/goals/README.md +++ b/src/graph/goals/README.md @@ -116,7 +116,8 @@ let exec = graph.run_with_thread("thread-1", St::default()).await?; | File | Role | | --- | --- | -| `types.rs` | Data model + `GoalProgress` + `active_goal_context_block`. | +| `types.rs` | Data model: `ThreadGoal`, `ThreadGoalStatus`, `GoalProgress`. | +| `prompt.rs` | `active_goal_context_block` — renders the per-iteration prompt context block. | | `store.rs` | `Store`-backed CRUD, per-thread RMW lock, budget + CAS guards. | | `tool.rs` | `GoalTool` / `GoalToolKind` harness tools. | | `continuation.rs` | `goal_gate_node`, `run_continuation_tick`, `note_user_turn`. | diff --git a/src/graph/goals/continuation.rs b/src/graph/goals/continuation.rs index 417a08b..1f9b6d2 100644 --- a/src/graph/goals/continuation.rs +++ b/src/graph/goals/continuation.rs @@ -132,7 +132,7 @@ where F: Fn(ThreadGoal) -> Fut, Fut: Future>, { - let now = store::now_ms(); + let now = crate::harness::ids::now_ms(); let idle_ms = idle.as_millis() as u64; let mut candidates: Vec = store::list_all(store) .await? diff --git a/src/graph/goals/mod.rs b/src/graph/goals/mod.rs index b614879..c438f41 100644 --- a/src/graph/goals/mod.rs +++ b/src/graph/goals/mod.rs @@ -13,15 +13,15 @@ //! primitive is provider-neutral and drives off the graph runtime. mod continuation; +mod prompt; pub mod store; mod tool; mod types; pub use continuation::{goal_gate_node, note_user_turn, run_continuation_tick}; +pub use prompt::active_goal_context_block; pub use tool::{GoalTool, GoalToolKind, goal_tools, register_goal_tools}; -pub use types::{ - GoalProgress, ThreadGoal, ThreadGoalStatus, TurnOutcome, active_goal_context_block, -}; +pub use types::{GoalProgress, ThreadGoal, ThreadGoalStatus, TurnOutcome}; #[cfg(test)] mod test; diff --git a/src/graph/goals/prompt.rs b/src/graph/goals/prompt.rs new file mode 100644 index 0000000..69605ed --- /dev/null +++ b/src/graph/goals/prompt.rs @@ -0,0 +1,35 @@ +//! Prompt-rendering for the durable [`ThreadGoal`]: the per-iteration +//! context block a caller prepends to a work node's prompt. +//! +//! Split out of `goals/types.rs`; kept separate from the plain type +//! definitions since this is presentation, not state. + +use crate::graph::goals::types::{ThreadGoal, ThreadGoalStatus}; + +/// Renders the per-iteration context block a caller can prepend to a work +/// node's prompt so the model knows the active objective and how to close it +/// out. Returns `None` for statuses that should not drive further work +/// ([`Paused`](ThreadGoalStatus::Paused) / [`Complete`](ThreadGoalStatus::Complete)). +pub fn active_goal_context_block(goal: &ThreadGoal) -> Option { + match goal.status { + ThreadGoalStatus::Active => { + let budget = match goal.budget_remaining() { + Some(remaining) => format!(" (~{remaining} tokens of budget remain)"), + None => String::new(), + }; + Some(format!( + "[thread goal] You are working toward this thread's durable goal{budget}.\n\n\ + Goal: {objective}\n\n\ + Assess progress against concrete evidence, then take the next useful step. \ + If the goal is already satisfied, call `goal_complete` now.", + objective = goal.objective, + )) + } + ThreadGoalStatus::BudgetLimited => Some(format!( + "[thread goal] The token budget for this goal is exhausted. Summarise progress \ + and stop; do not start new substantive work.\n\nGoal: {}", + goal.objective, + )), + ThreadGoalStatus::Paused | ThreadGoalStatus::Complete => None, + } +} diff --git a/src/graph/goals/store.rs b/src/graph/goals/store.rs index 39d833d..28da269 100644 --- a/src/graph/goals/store.rs +++ b/src/graph/goals/store.rs @@ -24,13 +24,12 @@ use std::collections::HashMap; use std::sync::{Arc, Mutex as StdMutex, OnceLock}; -use std::time::{SystemTime, UNIX_EPOCH}; use tokio::sync::Mutex; use super::types::{ThreadGoal, ThreadGoalStatus}; use crate::error::{Result, TinyAgentsError}; -use crate::harness::ids::next_seq; +use crate::harness::ids::{next_seq, now_ms}; use crate::harness::store::Store; /// The [`Store`] namespace holding one [`ThreadGoal`] per thread. @@ -46,14 +45,6 @@ fn thread_lock(thread_id: &str) -> Arc> { guard.entry(thread_id.to_string()).or_default().clone() } -/// Current unix time in milliseconds. Dependency-free (no `chrono`). -pub(crate) fn now_ms() -> u64 { - SystemTime::now() - .duration_since(UNIX_EPOCH) - .map(|d| d.as_millis() as u64) - .unwrap_or(0) -} - /// Hex-encodes the thread id into a [`Store`]-safe key. Required because /// [`FileStore`](crate::harness::store::FileStore) rejects key bytes outside /// `[A-Za-z0-9._-]`; hex is uniform across every backend. diff --git a/src/graph/goals/test.rs b/src/graph/goals/test.rs index e40dcd2..365c2c8 100644 --- a/src/graph/goals/test.rs +++ b/src/graph/goals/test.rs @@ -1,5 +1,6 @@ //! Unit tests for the thread-goal domain types. +use super::prompt::*; use super::types::*; fn goal(status: ThreadGoalStatus, token_budget: Option, tokens_used: u64) -> ThreadGoal { diff --git a/src/graph/goals/types.rs b/src/graph/goals/types.rs index 25daaef..2eb1c95 100644 --- a/src/graph/goals/types.rs +++ b/src/graph/goals/types.rs @@ -131,31 +131,3 @@ pub struct GoalProgress { /// driver ([`super::run_continuation_tick`]). Identical in shape to /// [`GoalProgress`]; named distinctly at the driver boundary for clarity. pub type TurnOutcome = GoalProgress; - -/// Renders the per-iteration context block a caller can prepend to a work -/// node's prompt so the model knows the active objective and how to close it -/// out. Returns `None` for statuses that should not drive further work -/// ([`Paused`](ThreadGoalStatus::Paused) / [`Complete`](ThreadGoalStatus::Complete)). -pub fn active_goal_context_block(goal: &ThreadGoal) -> Option { - match goal.status { - ThreadGoalStatus::Active => { - let budget = match goal.budget_remaining() { - Some(remaining) => format!(" (~{remaining} tokens of budget remain)"), - None => String::new(), - }; - Some(format!( - "[thread goal] You are working toward this thread's durable goal{budget}.\n\n\ - Goal: {objective}\n\n\ - Assess progress against concrete evidence, then take the next useful step. \ - If the goal is already satisfied, call `goal_complete` now.", - objective = goal.objective, - )) - } - ThreadGoalStatus::BudgetLimited => Some(format!( - "[thread goal] The token budget for this goal is exhausted. Summarise progress \ - and stop; do not start new substantive work.\n\nGoal: {}", - goal.objective, - )), - ThreadGoalStatus::Paused | ThreadGoalStatus::Complete => None, - } -} diff --git a/src/graph/mod.rs b/src/graph/mod.rs index f0dd7f1..67b961b 100644 --- a/src/graph/mod.rs +++ b/src/graph/mod.rs @@ -51,8 +51,9 @@ pub use channel::{ #[cfg(feature = "sqlite")] pub use checkpoint::SqliteCheckpointer; pub use checkpoint::{ - Checkpoint, CheckpointConfig, CheckpointMetadata, CheckpointSource, CheckpointTuple, - Checkpointer, DurabilityMode, FileCheckpointer, InMemoryCheckpointer, PendingWrite, + BarrierArrivals, Checkpoint, CheckpointConfig, CheckpointMetadata, CheckpointSource, + CheckpointTuple, Checkpointer, DurabilityMode, FileCheckpointer, InMemoryCheckpointer, + PendingActivation, PendingWrite, }; pub use command::{Command, Interrupt, NodeResult, RouteTarget, Send}; pub use compiled::{CompiledGraph, GraphExecution, GraphInput, ResumeTarget, StateSnapshot}; diff --git a/src/graph/observability/README.md b/src/graph/observability/README.md new file mode 100644 index 0000000..0a218c8 --- /dev/null +++ b/src/graph/observability/README.md @@ -0,0 +1,79 @@ +# 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. +- `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 shares the + harness `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. + +## 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 reported to stderr 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. | +| `test.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/src/graph/observability/langfuse.rs b/src/graph/observability/langfuse.rs index 5e96a26..8421229 100644 --- a/src/graph/observability/langfuse.rs +++ b/src/graph/observability/langfuse.rs @@ -567,4 +567,4 @@ fn pop( } #[cfg(test)] -mod tests; +mod test; diff --git a/src/graph/observability/langfuse/tests.rs b/src/graph/observability/langfuse/test.rs similarity index 100% rename from src/graph/observability/langfuse/tests.rs rename to src/graph/observability/langfuse/test.rs diff --git a/src/graph/observability/mod.rs b/src/graph/observability/mod.rs index 0939683..66eb0e5 100644 --- a/src/graph/observability/mod.rs +++ b/src/graph/observability/mod.rs @@ -34,8 +34,12 @@ //! 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 with `futures::executor::block_on` and treats -//! persistence as best-effort: a backend error never aborts the run. +//! the async journal API through a background drain (an +//! [`AppendWorker`](crate::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 types; @@ -53,7 +57,8 @@ use async_trait::async_trait; use crate::error::Result; use crate::graph::status::GraphRunStatus; use crate::graph::stream::{GraphEvent, GraphEventSink}; -use crate::harness::ids::{CheckpointId, EventId, GraphId, NodeId, RunId, ThreadId}; +use crate::harness::ids::{CheckpointId, EventId, GraphId, NodeId, RunId, ThreadId, now_ms}; +use crate::harness::observability::{AppendWorker, DEFAULT_DRAIN_CAPACITY}; use crate::harness::store::AppendStore; // --------------------------------------------------------------------------- @@ -378,8 +383,16 @@ impl JournalGraphSink { /// 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 { - journal, + worker, inner: None, root_run_id: run_id.clone(), run_id, @@ -454,12 +467,19 @@ impl JournalGraphSink { impl GraphEventSink for JournalGraphSink { fn emit(&self, event: GraphEvent) { let obs = self.observe(&event); - // Best-effort durable append; never abort the run on a journal error. - let _ = futures::executor::block_on(self.journal.append(obs)); + // Hand off to the background drain; never block the executor on I/O. + self.worker.submit(obs); if let Some(inner) = &self.inner { inner.emit(event); } } + + fn flush(&self) { + self.worker.flush(); + if let Some(inner) = &self.inner { + inner.flush(); + } + } } /// Extracts the checkpoint id a [`GraphEvent::CheckpointSaved`] carries, so the @@ -499,14 +519,6 @@ fn duration_ms(start: SystemTime, end: SystemTime) -> Option { .map(|duration| duration.as_millis() as u64) } -/// Returns the current time in Unix-epoch milliseconds, saturating at `0`. -pub(crate) fn now_ms() -> u64 { - std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .map(|d| d.as_millis() as u64) - .unwrap_or(0) -} - /// Builds a uniform poisoned-lock validation error for the in-memory backends. fn poisoned(what: &str, err: E) -> crate::error::TinyAgentsError { crate::error::TinyAgentsError::Validation(format!("{what} lock poisoned: {err}")) diff --git a/src/graph/observability/test.rs b/src/graph/observability/test.rs index 53f84d6..94c2073 100644 --- a/src/graph/observability/test.rs +++ b/src/graph/observability/test.rs @@ -420,6 +420,9 @@ async fn journal_sink_used_directly_forwards_to_inner() { // 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(); diff --git a/src/graph/observability/types.rs b/src/graph/observability/types.rs index f03954d..aea9e23 100644 --- a/src/graph/observability/types.rs +++ b/src/graph/observability/types.rs @@ -22,6 +22,7 @@ use crate::error::Result; use crate::graph::status::GraphRunStatus; use crate::graph::stream::{GraphEvent, GraphEventSink}; use crate::harness::ids::{CheckpointId, EventId, GraphId, NodeId, RunId, ThreadId}; +use crate::harness::observability::AppendWorker; use crate::harness::store::AppendStore; // --------------------------------------------------------------------------- @@ -337,16 +338,19 @@ pub struct InMemoryGraphStatusStore { /// 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 -/// async append is bridged synchronously with `futures::executor::block_on`, -/// and append errors are swallowed so a failing journal never aborts the run. +/// 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::graph::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::graph::stream::CollectingSink`]) so a /// single configured sink can both persist and broadcast. #[derive(Clone)] pub struct JournalGraphSink { - /// The journal observations are appended to. - pub(crate) journal: Arc, + /// 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. diff --git a/src/graph/orchestration/README.md b/src/graph/orchestration/README.md new file mode 100644 index 0000000..ac4a883 --- /dev/null +++ b/src/graph/orchestration/README.md @@ -0,0 +1,92 @@ +# graph::orchestration + +Graph-level orchestration controls. + +This module is the graph runtime's managed child-work surface. It gives +language-model orchestrators stable task ids and typed controls — `spawn`, +`await`, `cancel`, `kill`, `status`, `list`, `timeout`, `race`, `yield`, and +`steer` — without exposing raw executor handles such as `tokio::JoinHandle`. +A model asks for work by task id and observes lifecycle status; it never +touches an in-process future/handle directly, so the same tool surface works +whether the "child" is an in-process subgraph, a sub-agent, or (in principle) +out-of-process work fronted by a `TaskStore` implementation. + +The controls are ordinary harness tools. Use `OrchestrationTool` directly, +call `orchestration_tools` to build the full set, or call +`register_orchestration_tools` to insert them into a +`harness::tool::ToolRegistry` alongside any other tools. + +## Public surface + +### Tools (`tool.rs`) + +- `OrchestrationTool` — a single control, constructed with an + `OrchestrationToolKind` and a `TaskStore`; `.with_steering(..)` wires it to a + `SteeringRegistry` for controls that need to reach a running task (e.g. + `steer`, `cancel`). +- `orchestration_tools(store)` — builds the full default set of controls + against one store. +- `orchestration_tools_with_steering(store, steering)` — same, with steering + wired. +- `register_orchestration_tools(registry, store, ..)` — inserts the full set + into a `ToolRegistry`. +- `orchestration_tool_schema(kind)` / `orchestration_tool_schemas()` — the + JSON tool schemas, independent of a concrete store (useful for prompting or + schema inspection without constructing tools). +- `SteeringRegistry` — a concurrent map from `TaskId` to `SteeringHandle` + (`register` / `deregister` / `get`), letting a `steer`/`cancel` control reach + a task that is currently running. + +### Task model (`types.rs`) + +- `OrchestrationTaskKind` — what a task *is* (subgraph run, sub-agent run, + ...); `as_str()`. +- `OrchestrationTaskStatus` — lifecycle state; `is_terminal()` / `is_live()` + predicates. +- `OrchestrationTaskSpec` — the request to spawn a task: kind, lineage + (thread/node), timeout, input payload, metadata. Built with `new` + + `with_lineage` / `with_thread` / `with_node` / `with_timeout_ms` / + `with_input` / `with_metadata`. +- `OrchestrationTaskResult` — a completed task's output (`text` / `output` + constructors for the common shapes). +- `OrchestrationTaskRecord` — the durable record a `TaskStore` holds: spec + + current status + result once terminal. `pending(spec)` constructs the + initial record; `task_id()` / `is_terminal()` accessors. +- `OrchestrationTaskFilter` — filters for `list` (`with_kind`, + `created_between`, `matches(&record)`). +- `OrchestrationToolKind` — the control identifiers (`spawn`, `await`, + `cancel`, `kill`, `status`, `list`, `timeout`, `race`, `yield`, `steer`); + `name()` / `description()` give the tool-facing strings. +- `OrchestrationControlOutcome` — the typed result of invoking a control. + +### Storage (`store.rs`) + +- `TaskStore` (trait) — durable task bookkeeping: create/update/list/get + records by id, apply an `OrchestrationTaskFilter`. + - `InMemoryTaskStore` — in-process implementation; `from_records(..)` seeds + it for tests. + - `JsonlTaskStore` — append-only JSONL-backed implementation; + `JsonlTaskStore::open(path)`. + +## Files + +| File | Role | +| --- | --- | +| `types.rs` | Task kind/status/spec/result/record/filter types, `OrchestrationToolKind`, `OrchestrationControlOutcome`. | +| `tool.rs` | `OrchestrationTool`, `SteeringRegistry`, tool constructors and schemas. | +| `store.rs` | `TaskStore` trait, `InMemoryTaskStore`, `JsonlTaskStore`. | +| `test.rs` | Unit tests (spawn/await/cancel/timeout/race semantics, store round-trips, filters). | + +## Operational constraints + +- Controls that reach a *running* task (`steer`, `cancel`) require a + `SteeringRegistry` populated with that task's `SteeringHandle`; without one, + those controls only affect terminal-state bookkeeping in the `TaskStore`, + not the in-flight task itself. +- `JsonlTaskStore` is append-only: task-record updates are appended, not + rewritten in place, so a long-lived store grows monotonically. Compact or + rotate externally if that matters for your deployment. +- `OrchestrationTaskFilter::matches` is evaluated in-process against loaded + records; a `TaskStore` backend is not required to push the filter down, so + `list` cost scales with total record count unless a given implementation + optimizes it. diff --git a/src/graph/parallel/mod.rs b/src/graph/parallel/mod.rs index d711778..d168c8e 100644 --- a/src/graph/parallel/mod.rs +++ b/src/graph/parallel/mod.rs @@ -9,12 +9,14 @@ //! is what [`map_reduce`] provides, independent of the graph executor. //! //! Results are always returned in **input order** even though items complete out -//! of order, because the concurrency is driven by -//! [`buffered`](futures::stream::StreamExt::buffered), which preserves order. +//! 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 types; -pub use types::*; +pub use types::{FailurePolicy, ItemOutcome, ParallelOptions, ParallelOutcome}; use futures::stream::StreamExt; use std::future::Future; @@ -78,7 +80,12 @@ where .buffer_unordered(concurrency); let mut slots: Vec>> = (0..total).map(|_| None).collect(); - let mut fail_fast_error: Option = None; + // 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. @@ -88,8 +95,20 @@ where Ok(value) => slots[index] = Some(Ok(value)), Err(err) => { if options.failure_policy == FailurePolicy::FailFast { - fail_fast_error = Some(err); - break; // dropping the stream cancels remaining work + // 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())); } @@ -125,7 +144,7 @@ where None => raced.await?, } - if let Some(err) = fail_fast_error { + if let Some((_, err)) = fail_fast_error { return Err(err); } diff --git a/src/graph/parallel/test.rs b/src/graph/parallel/test.rs index c9d9fad..f3ca1bd 100644 --- a/src/graph/parallel/test.rs +++ b/src/graph/parallel/test.rs @@ -61,6 +61,33 @@ async fn fail_fast_returns_first_error() { 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)); diff --git a/src/graph/stream/mod.rs b/src/graph/stream/mod.rs index 7bc0ca9..aaa92ee 100644 --- a/src/graph/stream/mod.rs +++ b/src/graph/stream/mod.rs @@ -23,6 +23,14 @@ use std::sync::{Arc, Mutex}; pub trait GraphEventSink: Send + Sync { /// Receives one graph event. Implementations must not block the executor. fn emit(&self, event: GraphEvent); + + /// 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. diff --git a/src/graph/stream/types.rs b/src/graph/stream/types.rs index 8836aee..1f7f3d6 100644 --- a/src/graph/stream/types.rs +++ b/src/graph/stream/types.rs @@ -114,6 +114,11 @@ pub enum GraphEvent { /// Persisted checkpoint id. checkpoint_id: CheckpointId, }, + /// 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. @@ -177,6 +182,7 @@ impl GraphEvent { 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", diff --git a/src/graph/subagent_node/mod.rs b/src/graph/subagent_node/mod.rs index 4f875f3..0eb99ed 100644 --- a/src/graph/subagent_node/mod.rs +++ b/src/graph/subagent_node/mod.rs @@ -38,9 +38,9 @@ use async_trait::async_trait; use crate::graph::builder::NodeContext; use crate::graph::command::NodeResult; -use crate::graph::compiled::next_seq; use crate::graph::recursion::ChildRun; use crate::harness::events::EventSink; +use crate::harness::ids::next_seq; use crate::harness::ids::{GraphId, RunId}; use crate::harness::subagent::SubAgent; use crate::registry::{CapabilityRegistry, ComponentId}; diff --git a/src/graph/subgraph/README.md b/src/graph/subgraph/README.md new file mode 100644 index 0000000..8e9aa79 --- /dev/null +++ b/src/graph/subgraph/README.md @@ -0,0 +1,77 @@ +# 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). + +## 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. | +| `test.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/src/graph/subgraph/mod.rs b/src/graph/subgraph/mod.rs index c84d294..21fe3f6 100644 --- a/src/graph/subgraph/mod.rs +++ b/src/graph/subgraph/mod.rs @@ -24,7 +24,7 @@ use std::pin::Pin; use crate::Result; use crate::graph::builder::NodeContext; -use crate::graph::command::NodeResult; +use crate::graph::command::{Command, NodeResult}; use crate::graph::compiled::{CompiledGraph, GraphExecution}; use crate::graph::recursion::ChildRun; @@ -47,10 +47,16 @@ where 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 recorder = ChildRunRecorder::new(&ctx); Box::pin(async move { - let execution = run_child(child, thread_id, state).await?; + let execution = drive_child(child, thread_id, state, resume).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)) }) }) @@ -77,13 +83,19 @@ where 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 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 = run_child(child, thread_id, child_input).await?; + let execution = drive_child(child, thread_id, child_input, resume).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)) }) @@ -109,21 +121,41 @@ fn child_for(child: &CompiledGraph, ctx: &NodeContext) -> CompiledGr .with_recursion_node(ctx.node_id.clone()) } -async fn run_child( +/// 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. +async fn drive_child( child: CompiledGraph, thread_id: Option, state: S, + resume: Option, ) -> Result> where S: Clone + Send + Sync + 'static, U: Send + 'static, { - match thread_id { - Some(thread_id) => child.run_with_thread(thread_id, state).await, - None => child.run(state).await, + match (thread_id, resume) { + (Some(thread_id), Some(value)) => child.resume(thread_id, Command::resume(value)).await, + (Some(thread_id), None) => child.run_with_thread(thread_id, state).await, + (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::graph::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. diff --git a/src/graph/subgraph/test.rs b/src/graph/subgraph/test.rs index 1a0ff1f..af2fdd2 100644 --- a/src/graph/subgraph/test.rs +++ b/src/graph/subgraph/test.rs @@ -229,6 +229,86 @@ async fn embedded_child_persists_under_parent_thread_and_child_namespace() { 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::graph::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::graph::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 subgraph_child_run_distinct_and_shares_root() { // A parent embedding one child: the parent run records exactly one child run diff --git a/src/graph/testkit/README.md b/src/graph/testkit/README.md new file mode 100644 index 0000000..0dbc28b --- /dev/null +++ b/src/graph/testkit/README.md @@ -0,0 +1,89 @@ +# 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: + +- `taskstore_contract(store)` — basic CRUD/lifecycle contract for a + `graph::orchestration::TaskStore` implementation. +- `taskstore_concurrent_contract(store: Arc)` — concurrent-access contract + (no lost updates, no corrupted records under concurrent writers). +- `taskstore_replay_contract(reopen)` — durability contract: a store reopened + via `reopen` must replay to the same state. + +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` | `taskstore_contract`, `taskstore_concurrent_contract`, `taskstore_replay_contract`. | +| `test.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/src/graph/testkit/conformance.rs b/src/graph/testkit/conformance.rs index bb22a7b..d1f4ad3 100644 --- a/src/graph/testkit/conformance.rs +++ b/src/graph/testkit/conformance.rs @@ -32,6 +32,8 @@ fn contract_checkpoint( completed_tasks: vec![], pending_writes: vec![], interrupts: vec![], + pending_activations: None, + barrier_arrivals: vec![], metadata: serde_json::json!({ "source": "loop", "step": step }), } } @@ -65,6 +67,25 @@ where .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!( @@ -74,6 +95,36 @@ where .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"); diff --git a/src/graph/testkit/mod.rs b/src/graph/testkit/mod.rs index bade19b..61babbc 100644 --- a/src/graph/testkit/mod.rs +++ b/src/graph/testkit/mod.rs @@ -266,7 +266,7 @@ where graph_id: GraphId::new(format!("agent:{agent}")), run_id: RunId::new(format!( "subagent-fake-{}", - crate::graph::compiled::next_seq() + crate::harness::ids::next_seq() )), root_run_id, usage, diff --git a/src/harness/agent_loop/README.md b/src/harness/agent_loop/README.md new file mode 100644 index 0000000..b4da19d --- /dev/null +++ b/src/harness/agent_loop/README.md @@ -0,0 +1,95 @@ +# harness::agent_loop + +The default model-tool-model agent loop: the innermost turn of the recursive +(RLM-style) harness. + +This loop is where one model call is driven to completion. Because a whole +harness can be exposed as a tool (`harness::subagent::SubAgentTool`), 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 +usage/cost up to a parent run, and observes cooperative cancellation and +steering at safe checkpoints. + +The module is implemented as inherent methods on +`harness::runtime::AgentHarness` rather than a free function or +separate type — there is no standalone "AgentLoop" struct to construct. + +## Lifecycle + +1. Build a `RunContext` from the `RunConfig` and emit `AgentEvent::RunStarted`. +2. Run `before_agent` middleware. +3. Repeatedly: + - enforce the model-call cap and wall-clock deadline (fail-closed), + - build the `ModelRequest` from the working messages, registered tool + schemas, and the policy's default response format, + - run `before_model` middleware, emit `AgentEvent::ModelStarted`, + - resolve and invoke the model with retry + fallback, + - run `after_model` middleware, emit `AgentEvent::ModelCompleted`, fold + usage into the `AgentRun`, append the assistant message, + - if the assistant requested tools, execute each (enforcing the tool-call + cap, running `before_tool`/`after_tool`, emitting tool events) and append + the tool results, then continue, + - otherwise extract structured output when configured and break. +4. Run `after_agent` middleware and emit `AgentEvent::RunCompleted`. + +On any error the loop emits `AgentEvent::RunFailed`, fans the error out +through `on_error` middleware, and returns the error. + +## Limits + +Model- and tool-call caps come from `runtime::RunPolicy::limits` and are +enforced *before* each call, returning `TinyAgentsError::LimitExceeded`. The +wall-clock deadline (from the run config) is checked each iteration and +surfaces as `TinyAgentsError::Timeout`. The run context's own +`limits::LimitTracker` is also advanced so its counters stay consistent with +the enforced caps. + +## Backoff + +Retry backoff durations are *computed* via +`retry::RetryPolicy::backoff_for_attempt`, but whether the loop actually +sleeps for that duration is opt-in — off by default (keeping tests fast and +deterministic) and enabled per policy via +`retry::RetryPolicy::with_backoff_sleep`. A real provider integration retries +after a genuine, growing delay while unit tests stay sleep-free. + +## Public surface + +- `AgentHarness::invoke(state, ctx_data, config, input) -> Result` — + runs the loop, returns only the accumulated `AgentRun`. +- `AgentHarness::invoke_with_status(..) -> Result` — same run, + also returns a compact `HarnessRunStatus` snapshot (phase, counters, timing, + error summary) alongside the `AgentRun`. +- `AgentLoopResult { run: AgentRun, status: HarnessRunStatus }` — the richer + return type; the only public type this module owns beyond the + `AgentHarness` methods themselves. + +## Errors + +`TinyAgentsError::LimitExceeded` (model/tool cap reached), +`TinyAgentsError::Timeout` (wall-clock deadline elapsed), +`TinyAgentsError::ModelNotFound` (no model resolvable), +`TinyAgentsError::ToolNotFound` (model called an unregistered tool), or any +error surfaced by a model, tool, middleware, or structured-output extraction. + +## Files + +| File | Role | +| --- | --- | +| `mod.rs` | Module wiring: shared imports and the module-level doc comment. | +| `entry.rs` | Public entry points (`invoke`/`invoke_with_status`/`invoke_streaming*`) and the shared `drive` lifecycle wrapper. | +| `run_loop.rs` | The core loop body (`run_loop`) and response-cache decision logic. | +| `model_call.rs` | Cache-aware retry/fallback model dispatch, the streaming variant, and the innermost `ModelBaseCall`/`ToolBaseCall` impls the middleware wrap-onion terminates into. | +| `types.rs` | `AgentLoopResult`. | +| `test.rs` | Unit tests (limits, retry/fallback, tool execution, structured extraction). | + +## Operational constraints + +- The loop assumes `state: &State` is safe to read concurrently with any + nested sub-agent call — it never mutates it directly. +- Unregistered-tool calls fail closed per `runtime::UnknownToolPolicy`; there + is no silent skip. +- Identifiers (`CallId`, `ComponentId`) are derived deterministically from the + `RunConfig`, not randomly or from wall-clock time, so repeated calls with the + same input and config produce the same ids. diff --git a/src/harness/agent_loop/entry.rs b/src/harness/agent_loop/entry.rs new file mode 100644 index 0000000..7e4562f --- /dev/null +++ b/src/harness/agent_loop/entry.rs @@ -0,0 +1,210 @@ +//! Public entry points for the default agent loop: `invoke*`/ +//! `invoke_streaming*` and the shared `drive` lifecycle wrapper. +//! +//! Split out of `agent_loop/mod.rs`; see that module's doc comment for +//! the full loop lifecycle, limits, and backoff design. + +use super::*; + +impl AgentHarness { + /// Runs the default agent loop and returns the accumulated [`AgentRun`]. + /// + /// `state` is shared, read-only application data passed to every model and + /// tool call. `ctx_data` is moved into the [`RunContext`] for the run. + /// `config` supplies the run identity and limits, and `input` seeds the + /// working message transcript. + /// + /// # Errors + /// + /// Returns [`TinyAgentsError::LimitExceeded`] when the model- or tool-call + /// cap is reached, [`TinyAgentsError::Timeout`] when the wall-clock deadline + /// elapses, [`TinyAgentsError::ModelNotFound`] when no model can be + /// resolved, [`TinyAgentsError::ToolNotFound`] when the model calls an + /// unregistered tool, or any error surfaced by a model, tool, middleware, + /// or structured-output extraction. + pub async fn invoke( + &self, + state: &State, + ctx_data: Ctx, + config: RunConfig, + input: Vec, + ) -> Result { + self.invoke_with_status(state, ctx_data, config, input) + .await + .map(|result| result.run) + } + + /// Runs the default agent loop with a generated default [`RunConfig`]. + /// + /// Builds `RunConfig::new("run")` and a default `Ctx`. Identifiers are + /// derived deterministically from the config (no random or time-based ids), + /// so repeated calls with the same input behave identically. + pub async fn invoke_default(&self, state: &State, input: Vec) -> Result + where + Ctx: Default, + { + self.invoke(state, Ctx::default(), RunConfig::new("run"), input) + .await + } + + /// Runs the default agent loop and returns both the [`AgentRun`] and a + /// compact [`HarnessRunStatus`] snapshot describing how the run ended. + /// + /// This is the underlying entry point used by [`AgentHarness::invoke`]; use + /// it directly when you also need lifecycle/status information (phase, + /// counters, timing, error summary). On error the returned status would have + /// been marked failed, but the error is propagated instead so callers see + /// the failure; use the event stream for failed-run status. + pub async fn invoke_with_status( + &self, + state: &State, + ctx_data: Ctx, + config: RunConfig, + input: Vec, + ) -> Result { + let ctx = RunContext::new(config, ctx_data); + self.drive(state, ctx, input, false).await + } + + /// Runs the default agent loop inside a caller-supplied [`RunContext`], + /// returning the accumulated [`AgentRun`]. + /// + /// Use this when you need to control the run's dependencies — for example + /// to attach your own [`crate::harness::events::EventSink`] (so an external + /// listener or [`crate::harness::testkit::EventRecorder`] receives every + /// event), inject a custom [`crate::harness::store::StoreRegistry`], or carry + /// pre-populated `Ctx` data. The context's [`RunConfig`] supplies the run + /// identity and limits, exactly as for [`AgentHarness::invoke`]. + /// + /// # Errors + /// + /// Identical to [`AgentHarness::invoke`]. + pub async fn invoke_in_context( + &self, + state: &State, + ctx: RunContext, + input: Vec, + ) -> Result { + self.drive(state, ctx, input, false) + .await + .map(|result| result.run) + } + + /// Like [`AgentHarness::invoke_in_context`] but also returns the compact + /// [`HarnessRunStatus`] snapshot. + pub async fn invoke_in_context_with_status( + &self, + state: &State, + ctx: RunContext, + input: Vec, + ) -> Result { + self.drive(state, ctx, input, false).await + } + + /// Streaming counterpart of [`AgentHarness::invoke`]. + /// + /// Behaves exactly like [`AgentHarness::invoke`] except each model call is + /// driven through [`crate::harness::model::ChatModel::stream`] rather than + /// [`crate::harness::model::ChatModel::invoke`]: incremental message deltas + /// are emitted as [`AgentEvent::ModelDelta`] events and threaded through + /// every middleware's + /// [`on_model_delta`][crate::harness::middleware::Middleware::on_model_delta] + /// hook before the chunks are merged back into the final + /// [`crate::harness::model::ModelResponse`]. Tool execution, limits, retry, + /// fallback, structured output, and all other lifecycle behavior are + /// identical to the non-streaming path. + pub async fn invoke_streaming( + &self, + state: &State, + ctx_data: Ctx, + config: RunConfig, + input: Vec, + ) -> Result { + let ctx = RunContext::new(config, ctx_data); + self.drive(state, ctx, input, true) + .await + .map(|result| result.run) + } + + /// Streaming counterpart of [`AgentHarness::invoke_default`]. + pub async fn invoke_streaming_default( + &self, + state: &State, + input: Vec, + ) -> Result + where + Ctx: Default, + { + self.invoke_streaming(state, Ctx::default(), RunConfig::new("run"), input) + .await + } + + /// Streaming counterpart of [`AgentHarness::invoke_in_context`]. + pub async fn invoke_streaming_in_context( + &self, + state: &State, + ctx: RunContext, + input: Vec, + ) -> Result { + self.drive(state, ctx, input, true) + .await + .map(|result| result.run) + } + + /// Streaming counterpart of [`AgentHarness::invoke_in_context_with_status`]. + pub async fn invoke_streaming_in_context_with_status( + &self, + state: &State, + ctx: RunContext, + input: Vec, + ) -> Result { + self.drive(state, ctx, input, true).await + } + + /// Shared driver: runs the loop inside `ctx` and owns lifecycle + /// bookkeeping (status transitions plus `RunFailed`/`on_error` on error). + /// + /// `streaming` selects whether each model call is driven through + /// [`crate::harness::model::ChatModel::stream`] (firing `on_model_delta` + /// middleware per delta) or the unary + /// [`crate::harness::model::ChatModel::invoke`] path. + async fn drive( + &self, + state: &State, + mut ctx: RunContext, + input: Vec, + streaming: bool, + ) -> Result { + let run_id = ctx.config.run_id.clone(); + let thread_id = ctx.config.thread_id.clone(); + + let mut status = HarnessRunStatus::new(run_id.clone(), ComponentId::new("agent_loop")); + if let Some(thread) = thread_id { + status = status.with_thread(thread); + } + + let mut run = AgentRun::new(); + + match self + .run_loop(state, &mut ctx, &mut run, &mut status, input, streaming) + .await + { + Ok(()) => { + status.mark_completed(); + Ok(AgentLoopResult { run, status }) + } + Err(error) => { + let record = ctx.emit(AgentEvent::RunFailed { + run_id, + error: error.to_string(), + }); + status.set_last_event(record.id); + status.mark_failed(error.to_string()); + // Surface the failure to every middleware. Inner errors are + // ignored so the originating error is never masked. + let _ = self.middleware.run_on_error(&mut ctx, &error).await; + Err(error) + } + } + } +} diff --git a/src/harness/agent_loop/mod.rs b/src/harness/agent_loop/mod.rs index 3d2efcf..45a63ef 100644 --- a/src/harness/agent_loop/mod.rs +++ b/src/harness/agent_loop/mod.rs @@ -40,12 +40,16 @@ //! //! # Limits //! -//! Model and tool caps come from [`RunPolicy::limits`][crate::harness::runtime::RunPolicy] -//! and are enforced *before* each call, returning -//! [`TinyAgentsError::LimitExceeded`]. The wall-clock deadline (from the run -//! config) is checked each iteration and surfaces as -//! [`TinyAgentsError::Timeout`]. The run context's own [`crate::harness::limits::LimitTracker`] -//! is also advanced so its counters stay consistent. +//! Model and tool caps are enforced by the run context's own +//! [`crate::harness::limits::LimitTracker`], which is synced with +//! [`RunPolicy::limits`][crate::harness::runtime::RunPolicy] once at the start +//! of each run (see [`crate::harness::limits::LimitTracker::sync_call_limits`]) +//! so the harness policy and the per-run [`RunConfig`] agree on a single +//! enforced cap instead of silently disagreeing. Each call is checked +//! *before* it is made, returning [`TinyAgentsError::LimitExceeded`] whose +//! message always names the limit that actually tripped. The wall-clock +//! deadline (from the run config) is checked each iteration and surfaces as +//! [`TinyAgentsError::Timeout`]. //! //! # Backoff //! @@ -78,1018 +82,15 @@ use crate::harness::model::{ ChatModel, ModelDelta, ModelRequest, ModelResolutionSource, ModelResponse, ModelStreamItem, ResolvedModel, ResolvedModelBinding, ResponseFormat, StreamAccumulator, ToolChoice, }; -use crate::harness::retry::is_retryable; use crate::harness::runtime::{AgentHarness, UnknownToolPolicy}; use crate::harness::structured::{StructuredExtractor, StructuredStrategy}; use crate::harness::tool::{Tool, ToolCall, ToolSchema}; use futures::StreamExt; use serde_json::Value; -impl AgentHarness { - /// Runs the default agent loop and returns the accumulated [`AgentRun`]. - /// - /// `state` is shared, read-only application data passed to every model and - /// tool call. `ctx_data` is moved into the [`RunContext`] for the run. - /// `config` supplies the run identity and limits, and `input` seeds the - /// working message transcript. - /// - /// # Errors - /// - /// Returns [`TinyAgentsError::LimitExceeded`] when the model- or tool-call - /// cap is reached, [`TinyAgentsError::Timeout`] when the wall-clock deadline - /// elapses, [`TinyAgentsError::ModelNotFound`] when no model can be - /// resolved, [`TinyAgentsError::ToolNotFound`] when the model calls an - /// unregistered tool, or any error surfaced by a model, tool, middleware, - /// or structured-output extraction. - pub async fn invoke( - &self, - state: &State, - ctx_data: Ctx, - config: RunConfig, - input: Vec, - ) -> Result { - self.invoke_with_status(state, ctx_data, config, input) - .await - .map(|result| result.run) - } - - /// Runs the default agent loop with a generated default [`RunConfig`]. - /// - /// Builds `RunConfig::new("run")` and a default `Ctx`. Identifiers are - /// derived deterministically from the config (no random or time-based ids), - /// so repeated calls with the same input behave identically. - pub async fn invoke_default(&self, state: &State, input: Vec) -> Result - where - Ctx: Default, - { - self.invoke(state, Ctx::default(), RunConfig::new("run"), input) - .await - } - - /// Runs the default agent loop and returns both the [`AgentRun`] and a - /// compact [`HarnessRunStatus`] snapshot describing how the run ended. - /// - /// This is the underlying entry point used by [`AgentHarness::invoke`]; use - /// it directly when you also need lifecycle/status information (phase, - /// counters, timing, error summary). On error the returned status would have - /// been marked failed, but the error is propagated instead so callers see - /// the failure; use the event stream for failed-run status. - pub async fn invoke_with_status( - &self, - state: &State, - ctx_data: Ctx, - config: RunConfig, - input: Vec, - ) -> Result { - let ctx = RunContext::new(config, ctx_data); - self.drive(state, ctx, input, false).await - } - - /// Runs the default agent loop inside a caller-supplied [`RunContext`], - /// returning the accumulated [`AgentRun`]. - /// - /// Use this when you need to control the run's dependencies — for example - /// to attach your own [`crate::harness::events::EventSink`] (so an external - /// listener or [`crate::harness::testkit::EventRecorder`] receives every - /// event), inject a custom [`crate::harness::store::StoreRegistry`], or carry - /// pre-populated `Ctx` data. The context's [`RunConfig`] supplies the run - /// identity and limits, exactly as for [`AgentHarness::invoke`]. - /// - /// # Errors - /// - /// Identical to [`AgentHarness::invoke`]. - pub async fn invoke_in_context( - &self, - state: &State, - ctx: RunContext, - input: Vec, - ) -> Result { - self.drive(state, ctx, input, false) - .await - .map(|result| result.run) - } - - /// Like [`AgentHarness::invoke_in_context`] but also returns the compact - /// [`HarnessRunStatus`] snapshot. - pub async fn invoke_in_context_with_status( - &self, - state: &State, - ctx: RunContext, - input: Vec, - ) -> Result { - self.drive(state, ctx, input, false).await - } - - /// Streaming counterpart of [`AgentHarness::invoke`]. - /// - /// Behaves exactly like [`AgentHarness::invoke`] except each model call is - /// driven through [`crate::harness::model::ChatModel::stream`] rather than - /// [`crate::harness::model::ChatModel::invoke`]: incremental message deltas - /// are emitted as [`AgentEvent::ModelDelta`] events and threaded through - /// every middleware's - /// [`on_model_delta`][crate::harness::middleware::Middleware::on_model_delta] - /// hook before the chunks are merged back into the final - /// [`crate::harness::model::ModelResponse`]. Tool execution, limits, retry, - /// fallback, structured output, and all other lifecycle behavior are - /// identical to the non-streaming path. - pub async fn invoke_streaming( - &self, - state: &State, - ctx_data: Ctx, - config: RunConfig, - input: Vec, - ) -> Result { - let ctx = RunContext::new(config, ctx_data); - self.drive(state, ctx, input, true) - .await - .map(|result| result.run) - } - - /// Streaming counterpart of [`AgentHarness::invoke_default`]. - pub async fn invoke_streaming_default( - &self, - state: &State, - input: Vec, - ) -> Result - where - Ctx: Default, - { - self.invoke_streaming(state, Ctx::default(), RunConfig::new("run"), input) - .await - } - - /// Streaming counterpart of [`AgentHarness::invoke_in_context`]. - pub async fn invoke_streaming_in_context( - &self, - state: &State, - ctx: RunContext, - input: Vec, - ) -> Result { - self.drive(state, ctx, input, true) - .await - .map(|result| result.run) - } - - /// Streaming counterpart of [`AgentHarness::invoke_in_context_with_status`]. - pub async fn invoke_streaming_in_context_with_status( - &self, - state: &State, - ctx: RunContext, - input: Vec, - ) -> Result { - self.drive(state, ctx, input, true).await - } - - /// Shared driver: runs the loop inside `ctx` and owns lifecycle - /// bookkeeping (status transitions plus `RunFailed`/`on_error` on error). - /// - /// `streaming` selects whether each model call is driven through - /// [`crate::harness::model::ChatModel::stream`] (firing `on_model_delta` - /// middleware per delta) or the unary - /// [`crate::harness::model::ChatModel::invoke`] path. - async fn drive( - &self, - state: &State, - mut ctx: RunContext, - input: Vec, - streaming: bool, - ) -> Result { - let run_id = ctx.config.run_id.clone(); - let thread_id = ctx.config.thread_id.clone(); - - let mut status = HarnessRunStatus::new(run_id.clone(), ComponentId::new("agent_loop")); - if let Some(thread) = thread_id { - status = status.with_thread(thread); - } - - let mut run = AgentRun::new(); - - match self - .run_loop(state, &mut ctx, &mut run, &mut status, input, streaming) - .await - { - Ok(()) => { - status.mark_completed(); - Ok(AgentLoopResult { run, status }) - } - Err(error) => { - let record = ctx.emit(AgentEvent::RunFailed { - run_id, - error: error.to_string(), - }); - status.set_last_event(record.id); - status.mark_failed(error.to_string()); - // Surface the failure to every middleware. Inner errors are - // ignored so the originating error is never masked. - let _ = self.middleware.run_on_error(&mut ctx, &error).await; - Err(error) - } - } - } - - /// Drives the loop body, returning `Ok(())` on a clean finish or the first - /// error encountered. The caller owns lifecycle bookkeeping (final status - /// transition, `RunFailed`/`on_error` on error). - async fn run_loop( - &self, - state: &State, - ctx: &mut RunContext, - run: &mut AgentRun, - status: &mut HarnessRunStatus, - input: Vec, - streaming: bool, - ) -> Result<()> { - 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); - - let mut messages = input; - - status.mark_running(HarnessPhase::Middleware); - self.middleware.run_before_agent(ctx, state).await?; - - loop { - // Safe cancellation checkpoint: if an orchestrator requested - // cooperative cancellation, stop before doing any further work - // (steering, request build, or model call) for this turn. - if ctx.cancellation.is_cancelled() { - return Err(TinyAgentsError::Cancelled); - } - - // Safe steering checkpoint: drain any orchestrator/human steering - // commands and apply the policy-permitted ones before the next - // model call. Cancel terminates the run; Pause short-circuits it. - match crate::harness::steering::apply_pending_steering(ctx, &mut messages)? { - crate::harness::steering::SteeringOutcome::Cancel => { - return Err(TinyAgentsError::Cancelled); - } - crate::harness::steering::SteeringOutcome::Pause => break, - crate::harness::steering::SteeringOutcome::Continue => {} - } - - // Fail-closed limit and deadline checks before each model call. - if ctx.check_deadline().is_err() { - ctx.emit(AgentEvent::LimitReached { - kind: LimitKind::WallClock, - }); - return Err(TinyAgentsError::Timeout(format!( - "run `{}` exceeded its wall-clock deadline", - ctx.run_id() - ))); - } - if run.model_calls >= self.policy.limits.max_model_calls { - ctx.emit(AgentEvent::LimitReached { - kind: LimitKind::ModelCalls, - }); - return Err(TinyAgentsError::LimitExceeded(format!( - "max model calls ({}) reached", - self.policy.limits.max_model_calls - ))); - } - - // Build the request from the working transcript, tool schemas, and - // policy response format. - status.mark_running(HarnessPhase::BuildingRequest); - let mut request = ModelRequest::new(messages.clone()).with_tools(self.tools.schemas()); - if let Some(format) = &self.policy.default_response_format { - request = request.with_response_format(format.clone()); - } - if let Some(cap) = ctx.config.max_turn_output_tokens { - request.max_tokens = - Some(request.max_tokens.map_or(cap, |current| current.min(cap))); - } - - status.mark_running(HarnessPhase::Middleware); - self.middleware - .run_before_model(ctx, state, &mut request) - .await?; - - // Record against the context tracker too (keeps its counters - // consistent); map its error onto the deterministic limit error. - ctx.record_model_call().map_err(|_| { - TinyAgentsError::LimitExceeded(format!( - "max model calls ({}) reached", - self.policy.limits.max_model_calls - )) - })?; - - // Resolve the model for the event/log name before invoking. - let binding = self - .models - .resolve_request(&request, None, None) - .ok_or_else(|| { - TinyAgentsError::ModelNotFound( - request - .model - .clone() - .unwrap_or_else(|| "".to_string()), - ) - })?; - let model_name = binding.resolved.name.clone(); - - // Resolve the structured-output plan against the resolved model. - // `Auto` consults the model profile to choose provider-native schema - // mode versus a tool-call fallback; an explicit `JsonSchema` always - // uses provider-native mode. The chosen strategy drives extraction of - // the final response below. - let structured_plan: Option<(StructuredStrategy, String, Value)> = - match request.response_format.clone() { - Some(ResponseFormat::Auto { name, schema }) => { - let strategy = StructuredStrategy::for_profile(binding.model.profile()); - match strategy { - StructuredStrategy::ProviderSchema => { - request.response_format = - Some(ResponseFormat::json_schema(name.clone(), schema.clone())); - } - StructuredStrategy::ToolCall => { - request.response_format = Some(ResponseFormat::Text); - request.tools.push(ToolSchema { - name: name.clone(), - description: format!("Return the result as `{name}`."), - parameters: schema.clone(), - format: crate::harness::tool::ToolFormat::Json, - }); - request.tool_choice = ToolChoice::Tool(name.clone()); - } - } - Some((strategy, name, schema)) - } - Some(ResponseFormat::JsonSchema { name, schema }) => { - Some((StructuredStrategy::ProviderSchema, name, schema)) - } - _ => None, - }; - - let call_id = CallId::new(format!("{}-model-{}", ctx.run_id(), run.model_calls + 1)); - status.mark_running(HarnessPhase::Model); - status.active_model_call = Some(call_id.clone()); - let record = ctx.emit(AgentEvent::ModelStarted { - call_id: call_id.clone(), - model: model_name, - }); - status.set_last_event(record.id); - - // The real model call (cache + retry + fallback core) is the - // innermost base of the model-wrap onion. Lifecycle `before_model` - // already ran above; the wrap onion runs here; lifecycle - // `after_model` runs below — so ordering is: - // before_model -> wrap onion (outer..inner..base) -> after_model. - let base = ModelCallBase { - harness: self, - call_id: call_id.clone(), - resolved: binding.resolved, - model: binding.model, - streaming, - }; - // Snapshot the request messages for observability before `request` - // is moved into the model-wrap onion, gated by the capture policy so - // payload-free runs never serialize prompt text. - let captured_input = self - .policy - .capture - .model_io - .then(|| serde_json::to_value(&request.messages).unwrap_or(Value::Null)); - let mut response = self - .middleware - .run_wrapped_model(ctx, state, request, &base) - .await? - .into_response(); - - status.mark_running(HarnessPhase::Middleware); - self.middleware - .run_after_model(ctx, state, &mut response) - .await?; - - // Accounting. - run.model_calls += 1; - run.steps += 1; - status.model_calls = run.model_calls; - status.active_model_call = None; - if let Some(usage) = response.usage { - run.usage.record(usage); - status.usage = run.usage; - let record = ctx.emit(AgentEvent::UsageRecorded { usage }); - status.set_last_event(record.id); - } - let captured_output = self - .policy - .capture - .model_io - .then(|| serde_json::to_value(&response.message).unwrap_or(Value::Null)); - let record = ctx.emit(AgentEvent::ModelCompleted { - call_id, - usage: response.usage, - input: captured_input, - output: captured_output, - }); - status.set_last_event(record.id); - - messages.push(Message::Assistant(response.message.clone())); - - // Safe checkpoint: honor any control outcome a middleware requested - // during this turn (for example an early-exit tool or a budget stop - // hook), before executing further tools. - if let Some(control) = ctx.take_control() { - let record = ctx.emit(AgentEvent::ControlApplied { - control: control.kind().to_string(), - detail: match &control { - MiddlewareControl::StopWithFinal(text) => text.clone(), - MiddlewareControl::Interrupt { node, message } => { - format!("{node}: {message}") - } - }, - }); - status.set_last_event(record.id); - match control { - MiddlewareControl::StopWithFinal(text) => { - run.final_response = Some(ModelResponse::assistant(text)); - break; - } - MiddlewareControl::Interrupt { node, message } => { - return Err(TinyAgentsError::Interrupted { node, message }); - } - } - } - - let tool_calls = response.tool_calls().to_vec(); - - // A tool-call structured-output strategy produces an artificial tool - // call that is not a registered tool; treat it as the final response - // rather than attempting to execute it. - let structured_tool_hit = matches!( - &structured_plan, - Some((StructuredStrategy::ToolCall, name, _)) - if tool_calls.iter().any(|c| &c.name == name) - ); - - if tool_calls.is_empty() || structured_tool_hit { - // Final response: optionally extract structured output using the - // resolved plan (provider-native schema or tool-call arguments). - if let Some((strategy, name, schema)) = &structured_plan { - let extractor = - StructuredExtractor::new(*strategy, name.clone(), schema.clone()); - let output = extractor.extract(&response)?; - run.structured = Some(output.value); - } - run.final_response = Some(response); - break; - } - - // Execute requested tools serially. - status.mark_running(HarnessPhase::Tools); - for mut call in tool_calls { - // Safe cancellation checkpoint: stop before invoking the next - // (side-effecting) tool if cancellation was requested. - if ctx.cancellation.is_cancelled() { - return Err(TinyAgentsError::Cancelled); - } - if ctx.check_deadline().is_err() { - ctx.emit(AgentEvent::LimitReached { - kind: LimitKind::WallClock, - }); - return Err(TinyAgentsError::Timeout(format!( - "run `{}` exceeded its wall-clock deadline", - ctx.run_id() - ))); - } - if run.tool_calls >= self.policy.limits.max_tool_calls { - ctx.emit(AgentEvent::LimitReached { - kind: LimitKind::ToolCalls, - }); - return Err(TinyAgentsError::LimitExceeded(format!( - "max tool calls ({}) reached", - self.policy.limits.max_tool_calls - ))); - } - ctx.record_tool_call().map_err(|_| { - TinyAgentsError::LimitExceeded(format!( - "max tool calls ({}) reached", - self.policy.limits.max_tool_calls - )) - })?; - - self.middleware - .run_before_tool(ctx, state, &mut call) - .await?; - - let tool = match self.tools.get(&call.name) { - Some(tool) => tool, - None => { - // The model called an unregistered tool. Apply the run's - // `UnknownToolPolicy` instead of unconditionally aborting. - let requested = call.name.clone(); - let arguments = call.arguments.clone(); - let call_id = CallId::new(call.id.clone()); - - // Rewrite mode: retarget to a fixed compatibility tool if - // that tool exists, otherwise fall through to recovery. - let rewrite_target = match &self.policy.unknown_tool { - UnknownToolPolicy::Rewrite { tool_name } => { - self.tools.get(tool_name).map(|t| (tool_name.clone(), t)) - } - _ => None, - }; - - if let Some((tool_name, tool)) = rewrite_target { - call.name = tool_name.clone(); - let record = ctx.emit(AgentEvent::UnknownToolCall { - call_id, - requested_name: requested, - arguments, - recovery: format!("rewrite:{tool_name}"), - }); - status.set_last_event(record.id); - tool - } else if matches!(self.policy.unknown_tool, UnknownToolPolicy::Fail) { - return Err(TinyAgentsError::ToolNotFound(requested)); - } else { - // `ReturnToolError` (or a Rewrite whose target is also - // missing): inject a tool-error result naming the - // requested tool and the valid tools, then continue so - // the model can correct itself. This consumed one - // tool-call budget slot above, bounding the loop. - let valid = self.tools.names().join(", "); - let args_repr = serde_json::to_string(&arguments) - .unwrap_or_else(|_| "".to_string()); - let message = format!( - "unknown tool `{requested}` (arguments: {args_repr}); \ - valid tools: [{valid}]" - ); - let record = ctx.emit(AgentEvent::UnknownToolCall { - call_id, - requested_name: requested.clone(), - arguments, - recovery: "tool_error".to_string(), - }); - status.set_last_event(record.id); - run.tool_calls += 1; - status.tool_calls = run.tool_calls; - messages.push(Message::tool(call.id.clone(), message)); - continue; - } - } - }; - tool.schema().validate_call(&call)?; - - let tool_call_id = CallId::new(call.id.clone()); - let tool_name = call.name.clone(); - status.active_tool_calls.push(tool_call_id.clone()); - let record = ctx.emit(AgentEvent::ToolStarted { - call_id: tool_call_id.clone(), - tool_name: tool_name.clone(), - }); - status.set_last_event(record.id); - - // Snapshot the arguments for observability before `call` is - // moved into the tool-wrap onion, gated by the capture policy. - let captured_input = self.policy.capture.tool_io.then(|| call.arguments.clone()); - - // The real tool call is the innermost base of the tool-wrap - // onion (same before -> wrap -> after ordering as the model - // path): lifecycle `before_tool` ran above, the wrap onion runs - // here, and lifecycle `after_tool` runs below. - let base = ToolCallBase { tool }; - let mut result = self - .middleware - .run_wrapped_tool(ctx, state, call, &base) - .await? - .into_result(); - - self.middleware - .run_after_tool(ctx, state, &mut result) - .await?; - - run.tool_calls += 1; - status.tool_calls = run.tool_calls; - status.active_tool_calls.retain(|c| c != &tool_call_id); - let captured_output = self - .policy - .capture - .tool_io - .then(|| Value::String(result.content.clone())); - let record = ctx.emit(AgentEvent::ToolCompleted { - call_id: tool_call_id, - tool_name, - input: captured_input, - output: captured_output, - }); - status.set_last_event(record.id); - - messages.push(Message::tool( - result.call_id.clone(), - result.content.clone(), - )); - } - } - - run.messages = messages; - - status.mark_running(HarnessPhase::Middleware); - self.middleware.run_after_agent(ctx, state, run).await?; - - let record = ctx.emit(AgentEvent::RunCompleted { - run_id: ctx.run_id().clone(), - }); - status.set_last_event(record.id); - - Ok(()) - } - - /// Resolves the effective response-cache decision for `request`. - /// - /// Returns `Some((cache, key))` when a [`ResponseCache`] is attached to the - /// harness *and* caching is enabled for this call. The per-request - /// [`ModelRequest::cache_policy`] takes precedence over the harness-level - /// [`RunPolicy::cache`][crate::harness::runtime::RunPolicy]; when the request - /// carries no policy the run policy's - /// [`response_cache_enabled`][crate::harness::cache::CachePolicy] decides. - /// Returns `None` (caching disabled) when no cache is attached or the - /// effective policy disables it. - fn response_cache_decision( - &self, - request: &ModelRequest, - ) -> Option<(Arc, String)> { - let cache = self.response_cache.as_ref()?; - let enabled = match &request.cache_policy { - Some(policy) => policy.response_cache_enabled, - None => self.policy.cache.response_cache_enabled, - }; - if !enabled { - return None; - } - Some((Arc::clone(cache), cache_key(request))) - } - - /// Invokes a model, consulting the local response cache around the - /// retry/fallback path. - /// - /// When caching is enabled for this call (see - /// [`Self::response_cache_decision`]) the cache is checked **before** any - /// provider call: on a hit an [`AgentEvent::CacheHit`] is emitted and the - /// cached [`ModelResponse`] is returned *without* invoking the underlying - /// [`ChatModel`] (the retry/fallback path is skipped entirely); on a miss an - /// [`AgentEvent::CacheMiss`] is emitted, the provider is invoked normally, - /// and the successful response is written back to the cache. - /// - /// # Accounting - /// - /// A cache hit is still counted as a model "step"/call by the caller - /// ([`Self::run_loop`] increments `model_calls`/`steps` and emits - /// [`AgentEvent::ModelCompleted`] after this returns) so usage and limit - /// bookkeeping stay consistent whether or not a call was served from cache. - /// The behavioral guarantee is only that the underlying provider is not - /// contacted on a hit. - async fn invoke_model_with_retry( - &self, - state: &State, - ctx: &mut RunContext, - request: &ModelRequest, - call_id: &CallId, - binding: ResolvedModelBinding, - streaming: bool, - ) -> Result { - let decision = self.response_cache_decision(request); - - if let Some((cache, key)) = decision.as_ref() { - if let Some(mut cached) = cache.get(key).await? { - ctx.emit(AgentEvent::CacheHit { - call_id: call_id.clone(), - key: key.clone(), - }); - if cached.resolved_model.is_none() { - cached.resolved_model = Some(binding.resolved.clone()); - } - return Ok(cached); - } - ctx.emit(AgentEvent::CacheMiss { - call_id: call_id.clone(), - key: key.clone(), - }); - } - - let response = self - .invoke_model_resolving(state, ctx, request, call_id, binding, streaming) - .await?; - - if let Some((cache, key)) = decision.as_ref() { - cache.put(key, response.clone()).await?; - } - - Ok(response) - } - - /// Invokes a model with retry and fallback (no caching). - /// - /// Retries are governed by [`RunPolicy::retry`][crate::harness::runtime::RunPolicy] - /// and apply only to retryable errors (see [`is_retryable`]); each scheduled - /// retry emits [`AgentEvent::RetryScheduled`]. When retries are exhausted - /// (or the error is non-retryable) and a [`crate::harness::retry::FallbackPolicy`] - /// is configured, the next model in the chain is tried. The computed backoff - /// duration is intentionally not slept on (see the module docs). - async fn invoke_model_resolving( - &self, - state: &State, - ctx: &mut RunContext, - request: &ModelRequest, - call_id: &CallId, - binding: ResolvedModelBinding, - streaming: bool, - ) -> Result { - let mut current_name = binding.resolved.name.clone(); - let mut model = binding.model; - let mut resolved = binding.resolved; - let run_id = ctx.run_id().clone(); - - loop { - // Retry loop for the current model. - let mut attempt = 0usize; - let outcome = loop { - // Observe cancellation before (re)issuing a model attempt so a - // cancel requested during a retry/rate-limit wait stops the run - // promptly instead of firing another provider call or falling - // through to the fallback chain. - if ctx.cancellation.is_cancelled() { - return Err(TinyAgentsError::Cancelled); - } - // Bound this individual provider call by the run's *remaining* - // wall-clock budget so a hung or slow model call is interrupted - // mid-flight, not merely detected by the between-call deadline - // check. reqwest/futures are cancel-safe, so dropping the future - // on elapse cancels the underlying request. When neither the run - // config nor the harness policy configures a timeout the call is - // awaited unbounded. - let remaining = self.call_budget(ctx); - let attempt_result = if streaming { - let fut = - self.invoke_model_streaming_once(state, ctx, &model, request, call_id); - Self::with_call_budget(remaining, run_id.as_str(), fut).await - } else { - let fut = model.invoke(state, request.clone()); - Self::with_call_budget(remaining, run_id.as_str(), fut).await - }; - match attempt_result { - Ok(response) => break Ok(response), - Err(error) => { - if is_retryable(&error) && self.policy.retry.should_retry(attempt) { - attempt += 1; - ctx.emit(AgentEvent::RetryScheduled { - call_id: call_id.clone(), - attempt, - }); - // Sleep for the backoff only when the policy opts in - // (`with_backoff_sleep`); otherwise this is a no-op so - // the loop stays fast and deterministic in tests. - self.policy.retry.sleep_backoff(attempt).await; - continue; - } - break Err(error); - } - } - }; - - match outcome { - Ok(mut response) => { - if response.resolved_model.is_none() { - response.resolved_model = Some(resolved); - } - return Ok(response); - } - Err(error) => { - // Retries exhausted (or non-retryable): try the next model - // in the fallback chain, if any. - let next = self - .policy - .fallback - .as_ref() - .and_then(|fallback| fallback.next_after(¤t_name)) - .map(str::to_owned); - match next.and_then(|name| self.models.get(&name).map(|m| (name, m))) { - Some((name, next_model)) => { - resolved = ResolvedModel { - name: name.clone(), - requested: Some(name.clone()), - source: ModelResolutionSource::Hint, - }; - current_name = name; - model = next_model; - continue; - } - None => return Err(error), - } - } - } - } - } - - /// Computes the wall-clock budget for the next individual model call. - /// - /// The budget is the *tighter* of two remaining-time sources: - /// - /// - the run config's `timeout_ms` (the same deadline the between-call - /// [`RunContext::check_deadline`] enforces), tracked by the run's - /// [`crate::harness::limits::LimitTracker`], and - /// - the harness policy's - /// [`RunLimits::max_wall_clock_ms`][crate::harness::limits::RunLimits::max_wall_clock_ms], - /// measured against the same tracker start. - /// - /// Either source may be absent; when both are absent the call is unbounded - /// (`None`). Honoring the policy source lets a sub-agent whose child - /// [`RunConfig`] carries no per-run timeout still be bounded by its - /// harness's policy-level wall-clock cap. - fn call_budget(&self, ctx: &RunContext) -> Option { - let config_budget = ctx.remaining_wall_clock(); - let policy_budget = self.policy.limits.max_wall_clock_ms.map(|ms| { - Duration::from_millis(ms) - .checked_sub(ctx.limits.elapsed()) - .unwrap_or(Duration::ZERO) - }); - match (config_budget, policy_budget) { - (Some(a), Some(b)) => Some(a.min(b)), - (Some(a), None) => Some(a), - (None, Some(b)) => Some(b), - (None, None) => None, - } - } - - /// Awaits a single model-call future, optionally bounded by `budget`. - /// - /// When `budget` is `Some`, the future is wrapped in - /// [`tokio::time::timeout`]; if it elapses the future is dropped (cancelling - /// the in-flight provider request) and a - /// [`TinyAgentsError::Timeout`] is returned. When `budget` is `None` (no - /// run timeout configured) the future is awaited without a bound. - /// - /// `budget` is the run's *remaining* wall-clock budget at the time the call - /// is issued, so each successive call gets a tighter bound as the deadline - /// approaches. - async fn with_call_budget(budget: Option, run_id: &str, fut: F) -> F::Output - where - F: Future>, - { - match budget { - Some(budget) => match tokio::time::timeout(budget, fut).await { - Ok(result) => result, - Err(_) => Err(TinyAgentsError::Timeout(format!( - "model call for run `{run_id}` exceeded its remaining wall-clock budget \ - ({} ms)", - budget.as_millis() - ))), - }, - None => fut.await, - } - } - - /// Drives one streaming model call to completion. - /// - /// Consumes [`crate::harness::model::ChatModel::stream`], emitting an - /// [`AgentEvent::ModelDelta`] and running every middleware's - /// [`on_model_delta`][crate::harness::middleware::Middleware::on_model_delta] - /// hook for each [`ModelStreamItem::MessageDelta`] (and standalone - /// [`ModelStreamItem::ToolCallDelta`]), then folds the items into the final - /// [`ModelResponse`] via [`StreamAccumulator`]. The merged response is - /// equivalent to what the unary [`crate::harness::model::ChatModel::invoke`] - /// path would have produced, so the rest of the loop is unaffected. - async fn invoke_model_streaming_once( - &self, - state: &State, - ctx: &mut RunContext, - model: &Arc>, - request: &ModelRequest, - call_id: &CallId, - ) -> Result { - let mut stream = model.stream(state, request.clone()).await?; - let mut accumulator = StreamAccumulator::new(); - - // Clone the cheap token so the cancellation future does not borrow - // `ctx` for the duration of the stream loop (the body still needs - // `&mut ctx` for events and middleware). - let cancellation = ctx.cancellation.clone(); - - loop { - // Race the next provider chunk against cooperative cancellation. If - // cancellation wins we drop the partially consumed stream and unwind - // with `Cancelled`; the `cancelled()` future is cancel-safe. - let item = tokio::select! { - biased; - _ = cancellation.cancelled() => { - return Err(TinyAgentsError::Cancelled); - } - next = stream.next() => match next { - Some(item) => item, - None => break, - }, - }; - - // Surface incremental message/tool-call fragments through events and - // the `on_model_delta` middleware hook before merging them. - let message_delta = match &item { - ModelStreamItem::MessageDelta(delta) => Some(delta.clone()), - ModelStreamItem::ToolCallDelta(tool_delta) => Some(MessageDelta { - text: String::new(), - reasoning: String::new(), - tool_call: Some(tool_delta.clone()), - }), - _ => None, - }; - - if let Some(message_delta) = message_delta { - let record = ctx.emit(AgentEvent::ModelDelta { - run_id: ctx.config.run_id.clone(), - call_id: call_id.clone(), - delta: message_delta.clone(), - }); - let _ = record; - - let mut model_delta = ModelDelta { - call_id: call_id.as_str().to_string(), - content: message_delta.text.clone(), - tool_call: message_delta.tool_call.clone(), - }; - self.middleware - .run_on_model_delta(ctx, state, &mut model_delta) - .await?; - } - - accumulator.push(&item); - } - - accumulator.finish() - } -} - -/// The innermost model call wrapped by the model-wrap onion. -/// -/// Implements [`ModelBaseCall`] over the harness's cache + retry + fallback core -/// ([`AgentHarness::invoke_model_with_retry`]) so a [`crate::harness::middleware::ModelMiddleware`] -/// can proceed, short-circuit, retry, or fall back around the *whole* real model -/// call. The resolved binding is rebuilt per invocation so a wrap middleware -/// that retries `next` issues a fresh provider call each time. -struct ModelCallBase<'h, State: Send + Sync, Ctx: Send + Sync> { - harness: &'h AgentHarness, - call_id: CallId, - resolved: ResolvedModel, - model: Arc>, - streaming: bool, -} - -impl ModelBaseCall - for ModelCallBase<'_, State, Ctx> -{ - fn call<'a>( - &'a self, - ctx: &'a mut RunContext, - state: &'a State, - request: ModelRequest, - ) -> BoxModelFuture<'a> { - Box::pin(async move { - let binding = ResolvedModelBinding { - resolved: self.resolved.clone(), - model: Arc::clone(&self.model), - }; - self.harness - .invoke_model_with_retry( - state, - ctx, - &request, - &self.call_id, - binding, - self.streaming, - ) - .await - }) - } -} - -/// The innermost tool call wrapped by the tool-wrap onion. -/// -/// Implements [`ToolBaseCall`] over a single resolved [`Tool`] so a -/// [`crate::harness::middleware::ToolMiddleware`] can wrap the real tool -/// invocation. -struct ToolCallBase { - tool: Arc>, -} - -impl ToolBaseCall for ToolCallBase { - fn call<'a>( - &'a self, - ctx: &'a mut RunContext, - state: &'a State, - call: ToolCall, - ) -> BoxToolFuture<'a> { - Box::pin(async move { - self.tool - .call_with_context( - state, - call, - crate::harness::tool::ToolExecutionContext::from_run_context(ctx), - ) - .await - }) - } -} +mod entry; +mod model_call; +mod run_loop; #[cfg(test)] mod test; diff --git a/src/harness/agent_loop/model_call.rs b/src/harness/agent_loop/model_call.rs new file mode 100644 index 0000000..b423c8f --- /dev/null +++ b/src/harness/agent_loop/model_call.rs @@ -0,0 +1,426 @@ +//! Model invocation: cache-aware retry/fallback dispatch +//! (`invoke_model_with_retry`, `invoke_model_resolving`), the streaming +//! variant, and the innermost `ModelBaseCall`/`ToolBaseCall` impls that +//! the middleware wrap-onion terminates into. +//! +//! Split out of `agent_loop/mod.rs`; see that module's doc comment for +//! the full loop lifecycle, limits, and backoff design. + +use super::*; + +impl AgentHarness { + /// Invokes a model, consulting the local response cache around the + /// retry/fallback path. + /// + /// When caching is enabled for this call (see + /// [`Self::response_cache_decision`]) the cache is checked **before** any + /// provider call: on a hit an [`AgentEvent::CacheHit`] is emitted and the + /// cached [`ModelResponse`] is returned *without* invoking the underlying + /// [`ChatModel`] (the retry/fallback path is skipped entirely); on a miss an + /// [`AgentEvent::CacheMiss`] is emitted, the provider is invoked normally, + /// and the successful response is written back to the cache. + /// + /// # Accounting + /// + /// A cache hit is still counted as a model "step"/call by the caller + /// ([`Self::run_loop`] increments `model_calls`/`steps` and emits + /// [`AgentEvent::ModelCompleted`] after this returns) so usage and limit + /// bookkeeping stay consistent whether or not a call was served from cache. + /// The behavioral guarantee is only that the underlying provider is not + /// contacted on a hit. + async fn invoke_model_with_retry( + &self, + state: &State, + ctx: &mut RunContext, + request: &ModelRequest, + call_id: &CallId, + binding: ResolvedModelBinding, + streaming: bool, + ) -> Result { + let decision = self.response_cache_decision(request); + + if let Some((cache, key)) = decision.as_ref() { + if let Some(mut cached) = cache.get(key).await? { + ctx.emit(AgentEvent::CacheHit { + call_id: call_id.clone(), + key: key.clone(), + }); + if cached.resolved_model.is_none() { + cached.resolved_model = Some(binding.resolved.clone()); + } + return Ok(cached); + } + ctx.emit(AgentEvent::CacheMiss { + call_id: call_id.clone(), + key: key.clone(), + }); + } + + let response = self + .invoke_model_resolving(state, ctx, request, call_id, binding, streaming) + .await?; + + if let Some((cache, key)) = decision.as_ref() { + cache.put(key, response.clone()).await?; + } + + Ok(response) + } + + /// Invokes a model with retry and fallback (no caching). + /// + /// Retries are governed by [`RunPolicy::retry`][crate::harness::runtime::RunPolicy] + /// and apply only to retryable errors (see + /// [`is_retryable`][crate::harness::retry::is_retryable]); each scheduled + /// retry emits [`AgentEvent::RetryScheduled`]. When retries are exhausted + /// (or the error is non-retryable) and a [`crate::harness::retry::FallbackPolicy`] + /// is configured, the next model in the chain is tried. The computed backoff + /// duration is intentionally not slept on (see the module docs). + async fn invoke_model_resolving( + &self, + state: &State, + ctx: &mut RunContext, + request: &ModelRequest, + call_id: &CallId, + binding: ResolvedModelBinding, + streaming: bool, + ) -> Result { + let mut current_name = binding.resolved.name.clone(); + let mut model = binding.model; + let mut resolved = binding.resolved; + let run_id = ctx.run_id().clone(); + // Tracks every model name already attempted in this fallback chain so + // a chain containing a repeated name (e.g. `[primary, backup, + // primary]`) cannot alternate between the same two models forever; + // once a name has been tried it is never tried again. + let mut visited: std::collections::HashSet = std::collections::HashSet::new(); + visited.insert(current_name.clone()); + + loop { + // Retry loop for the current model. + let mut attempt = 0usize; + let outcome = loop { + // Observe cancellation before (re)issuing a model attempt so a + // cancel requested during a retry/rate-limit wait stops the run + // promptly instead of firing another provider call or falling + // through to the fallback chain. + if ctx.cancellation.is_cancelled() { + return Err(TinyAgentsError::Cancelled); + } + // Bound this individual provider call by the run's *remaining* + // wall-clock budget so a hung or slow model call is interrupted + // mid-flight, not merely detected by the between-call deadline + // check. reqwest/futures are cancel-safe, so dropping the future + // on elapse cancels the underlying request. When neither the run + // config nor the harness policy configures a timeout the call is + // awaited unbounded. + let remaining = self.call_budget(ctx); + let attempt_result = if streaming { + let fut = + self.invoke_model_streaming_once(state, ctx, &model, request, call_id); + Self::with_call_budget(remaining, run_id.as_str(), "model call", fut).await + } else { + let fut = model.invoke(state, request.clone()); + Self::with_call_budget(remaining, run_id.as_str(), "model call", fut).await + }; + match attempt_result { + Ok(response) => break Ok(response), + Err(error) => { + // `RunLimits::max_retries_per_call` is a hard ceiling + // that a looser `RetryPolicy::max_attempts` cannot + // exceed; whichever is stricter wins. + let max_attempts = self + .policy + .retry + .max_attempts_capped_at(self.policy.limits.max_retries_per_call); + // Route the retry decision through the shared + // `RetryPolicy::should_retry_error` engine (same + // classification + attempt-cap logic RetryMiddleware + // uses), applying the harness ceiling by capping a + // cloned policy first so the two sites cannot drift. + let capped = self.policy.retry.clone().with_max_attempts(max_attempts); + if capped.should_retry_error(attempt, &error) { + // Compute the backoff from the *pre-increment* + // attempt number: `attempt == 0` is the first + // retry and must sleep `initial_backoff_ms` + // (`RetryPolicy::backoff_for_attempt(0)`). Sleeping + // on the post-increment value skipped + // `initial_backoff_ms` entirely and shifted the + // whole exponential schedule one step too high. + let backoff_attempt = attempt; + attempt += 1; + ctx.emit(AgentEvent::RetryScheduled { + call_id: call_id.clone(), + attempt, + }); + // Sleep for the backoff only when the policy opts in + // (`with_backoff_sleep`); otherwise this is a no-op so + // the loop stays fast and deterministic in tests. + self.policy.retry.sleep_backoff(backoff_attempt).await; + continue; + } + break Err(error); + } + } + }; + + match outcome { + Ok(mut response) => { + if response.resolved_model.is_none() { + response.resolved_model = Some(resolved); + } + return Ok(response); + } + Err(error) => { + // A non-retryable, deadline-driven timeout must not feed + // back into the fallback chain: the run itself is out of + // wall-clock budget, so trying another model would just + // spin until the *next* deadline check fails identically. + if matches!(error, TinyAgentsError::Timeout(_)) { + return Err(error); + } + // Retries exhausted (or non-retryable): try the next model + // in the fallback chain, if any, skipping any name already + // visited in this chain so a chain with a repeated name + // cannot alternate between the same models forever. + let next = self + .policy + .fallback + .as_ref() + .and_then(|fallback| fallback.next_after(¤t_name)) + .map(str::to_owned) + .filter(|name| !visited.contains(name)); + match next.and_then(|name| self.models.get(&name).map(|m| (name, m))) { + Some((name, next_model)) => { + visited.insert(name.clone()); + resolved = ResolvedModel { + name: name.clone(), + requested: Some(name.clone()), + source: ModelResolutionSource::Hint, + }; + current_name = name; + model = next_model; + continue; + } + None => return Err(error), + } + } + } + } + } + + /// Computes the wall-clock budget for the next individual model call. + /// + /// The budget is the *tighter* of two remaining-time sources: + /// + /// - the run config's `timeout_ms` (the same deadline the between-call + /// [`RunContext::check_deadline`] enforces), tracked by the run's + /// [`crate::harness::limits::LimitTracker`], and + /// - the harness policy's + /// [`RunLimits::max_wall_clock_ms`][crate::harness::limits::RunLimits::max_wall_clock_ms], + /// measured against the same tracker start. + /// + /// Either source may be absent; when both are absent the call is unbounded + /// (`None`). Honoring the policy source lets a sub-agent whose child + /// [`RunConfig`] carries no per-run timeout still be bounded by its + /// harness's policy-level wall-clock cap. + pub(super) fn call_budget(&self, ctx: &RunContext) -> Option { + let config_budget = ctx.remaining_wall_clock(); + let policy_budget = self.policy.limits.max_wall_clock_ms.map(|ms| { + Duration::from_millis(ms) + .checked_sub(ctx.limits.elapsed()) + .unwrap_or(Duration::ZERO) + }); + match (config_budget, policy_budget) { + (Some(a), Some(b)) => Some(a.min(b)), + (Some(a), None) => Some(a), + (None, Some(b)) => Some(b), + (None, None) => None, + } + } + + /// Awaits a single call future (model or tool), optionally bounded by + /// `budget`. + /// + /// When `budget` is `Some`, the future is wrapped in + /// [`tokio::time::timeout`]; if it elapses the future is dropped (cancelling + /// the in-flight provider/tool request) and a + /// [`TinyAgentsError::Timeout`] is returned. When `budget` is `None` (no + /// run timeout configured) the future is awaited without a bound. + /// + /// `budget` is the run's *remaining* wall-clock budget at the time the call + /// is issued, so each successive call gets a tighter bound as the deadline + /// approaches. `what` names the kind of call in the timeout message (e.g. + /// `"model call"`, `"tool call"`). + pub(super) async fn with_call_budget( + budget: Option, + run_id: &str, + what: &str, + fut: F, + ) -> Result + where + F: Future>, + { + match budget { + Some(budget) => match tokio::time::timeout(budget, fut).await { + Ok(result) => result, + Err(_) => Err(TinyAgentsError::Timeout(format!( + "{what} for run `{run_id}` exceeded its remaining wall-clock budget \ + ({} ms)", + budget.as_millis() + ))), + }, + None => fut.await, + } + } + + /// Drives one streaming model call to completion. + /// + /// Consumes [`crate::harness::model::ChatModel::stream`], emitting an + /// [`AgentEvent::ModelDelta`] and running every middleware's + /// [`on_model_delta`][crate::harness::middleware::Middleware::on_model_delta] + /// hook for each [`ModelStreamItem::MessageDelta`] (and standalone + /// [`ModelStreamItem::ToolCallDelta`]), then folds the items into the final + /// [`ModelResponse`] via [`StreamAccumulator`]. The merged response is + /// equivalent to what the unary [`crate::harness::model::ChatModel::invoke`] + /// path would have produced, so the rest of the loop is unaffected. + async fn invoke_model_streaming_once( + &self, + state: &State, + ctx: &mut RunContext, + model: &Arc>, + request: &ModelRequest, + call_id: &CallId, + ) -> Result { + let mut stream = model.stream(state, request.clone()).await?; + let mut accumulator = StreamAccumulator::new(); + + // Clone the cheap token so the cancellation future does not borrow + // `ctx` for the duration of the stream loop (the body still needs + // `&mut ctx` for events and middleware). + let cancellation = ctx.cancellation.clone(); + + loop { + // Race the next provider chunk against cooperative cancellation. If + // cancellation wins we drop the partially consumed stream and unwind + // with `Cancelled`; the `cancelled()` future is cancel-safe. + let item = tokio::select! { + biased; + _ = cancellation.cancelled() => { + return Err(TinyAgentsError::Cancelled); + } + next = stream.next() => match next { + Some(item) => item, + None => break, + }, + }; + + // Surface incremental message/tool-call fragments through events and + // the `on_model_delta` middleware hook before merging them. + let message_delta = match &item { + ModelStreamItem::MessageDelta(delta) => Some(delta.clone()), + ModelStreamItem::ToolCallDelta(tool_delta) => Some(MessageDelta { + text: String::new(), + reasoning: String::new(), + tool_call: Some(tool_delta.clone()), + }), + _ => None, + }; + + if let Some(message_delta) = message_delta { + // Build the middleware-facing delta first (it needs owned + // copies of the fields), then move `message_delta` into the + // event so the hot path clones the payload once instead of + // twice per streamed token. + let mut model_delta = ModelDelta { + call_id: call_id.as_str().to_string(), + content: message_delta.text.clone(), + reasoning: message_delta.reasoning.clone(), + tool_call: message_delta.tool_call.clone(), + }; + ctx.emit(AgentEvent::ModelDelta { + run_id: ctx.config.run_id.clone(), + call_id: call_id.clone(), + delta: message_delta, + }); + self.middleware + .run_on_model_delta(ctx, state, &mut model_delta) + .await?; + } + + accumulator.push(&item); + } + + accumulator.finish() + } +} +/// The innermost model call wrapped by the model-wrap onion. +/// +/// Implements [`ModelBaseCall`] over the harness's cache + retry + fallback core +/// ([`AgentHarness::invoke_model_with_retry`]) so a [`crate::harness::middleware::ModelMiddleware`] +/// can proceed, short-circuit, retry, or fall back around the *whole* real model +/// call. The resolved binding is rebuilt per invocation so a wrap middleware +/// that retries `next` issues a fresh provider call each time. +pub(super) struct ModelCallBase<'h, State: Send + Sync, Ctx: Send + Sync> { + pub(super) harness: &'h AgentHarness, + pub(super) call_id: CallId, + pub(super) resolved: ResolvedModel, + pub(super) model: Arc>, + pub(super) streaming: bool, +} + +impl ModelBaseCall + for ModelCallBase<'_, State, Ctx> +{ + fn call<'a>( + &'a self, + ctx: &'a mut RunContext, + state: &'a State, + request: ModelRequest, + ) -> BoxModelFuture<'a> { + Box::pin(async move { + let binding = ResolvedModelBinding { + resolved: self.resolved.clone(), + model: Arc::clone(&self.model), + }; + self.harness + .invoke_model_with_retry( + state, + ctx, + &request, + &self.call_id, + binding, + self.streaming, + ) + .await + }) + } +} + +/// The innermost tool call wrapped by the tool-wrap onion. +/// +/// Implements [`ToolBaseCall`] over a single resolved [`Tool`] so a +/// [`crate::harness::middleware::ToolMiddleware`] can wrap the real tool +/// invocation. +pub(super) struct ToolCallBase { + pub(super) tool: Arc>, +} + +impl ToolBaseCall for ToolCallBase { + fn call<'a>( + &'a self, + ctx: &'a mut RunContext, + state: &'a State, + call: ToolCall, + ) -> BoxToolFuture<'a> { + Box::pin(async move { + self.tool + .call_with_context( + state, + call, + crate::harness::tool::ToolExecutionContext::from_run_context(ctx), + ) + .await + }) + } +} diff --git a/src/harness/agent_loop/run_loop.rs b/src/harness/agent_loop/run_loop.rs new file mode 100644 index 0000000..e037142 --- /dev/null +++ b/src/harness/agent_loop/run_loop.rs @@ -0,0 +1,462 @@ +//! The core superstep loop body: `run_loop` drives one model call, +//! any requested tool calls, and repeats until the model finishes or a +//! configured limit is reached. +//! +//! Split out of `agent_loop/mod.rs`; see that module's doc comment for +//! the full loop lifecycle, limits, and backoff design. + +use super::model_call::{ModelCallBase, ToolCallBase}; +use super::*; + +impl AgentHarness { + /// Drives the loop body, returning `Ok(())` on a clean finish or the first + /// error encountered. The caller owns lifecycle bookkeeping (final status + /// transition, `RunFailed`/`on_error` on error). + pub(super) async fn run_loop( + &self, + state: &State, + ctx: &mut RunContext, + run: &mut AgentRun, + status: &mut HarnessRunStatus, + input: Vec, + streaming: bool, + ) -> Result<()> { + 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); + + // Reconcile the `RunConfig`-derived limit tracker with the harness's + // `RunPolicy::limits` so model/tool call caps have one enforced + // source of truth instead of the two silently disagreeing (see + // `LimitTracker::sync_call_limits`). + ctx.limits.sync_call_limits( + self.policy.limits.max_model_calls, + self.policy.limits.max_tool_calls, + ); + + let mut messages = input; + + // The tool set is fixed for the duration of a run, so build the sorted + // schema vec once here instead of re-collecting, re-calling every tool's + // `schema()`, and re-sorting on every turn (per model call). + let tool_schemas = self.tools.schemas(); + + status.mark_running(HarnessPhase::Middleware); + self.middleware.run_before_agent(ctx, state).await?; + + loop { + // Safe cancellation checkpoint: if an orchestrator requested + // cooperative cancellation, stop before doing any further work + // (steering, request build, or model call) for this turn. + if ctx.cancellation.is_cancelled() { + return Err(TinyAgentsError::Cancelled); + } + + // Safe steering checkpoint: drain any orchestrator/human steering + // commands and apply the policy-permitted ones before the next + // model call. Cancel terminates the run; Pause short-circuits it. + match crate::harness::steering::apply_pending_steering(ctx, &mut messages)? { + crate::harness::steering::SteeringOutcome::Cancel => { + return Err(TinyAgentsError::Cancelled); + } + crate::harness::steering::SteeringOutcome::Pause => break, + crate::harness::steering::SteeringOutcome::Continue => {} + } + + // Fail-closed limit and deadline checks before each model call. + if ctx.check_deadline().is_err() { + ctx.emit(AgentEvent::LimitReached { + kind: LimitKind::WallClock, + }); + return Err(TinyAgentsError::Timeout(format!( + "run `{}` exceeded its wall-clock deadline", + ctx.run_id() + ))); + } + // The context's `LimitTracker` (synced with `RunPolicy::limits` + // above) is the single enforced source of truth for the model-call + // cap, so the reported limit always matches the one that trips. + if let Err(err) = ctx.record_model_call() { + ctx.emit(AgentEvent::LimitReached { + kind: LimitKind::ModelCalls, + }); + return Err(TinyAgentsError::LimitExceeded(err.to_string())); + } + + // Build the request from the working transcript, tool schemas, and + // policy response format. + status.mark_running(HarnessPhase::BuildingRequest); + let mut request = ModelRequest::new(messages.clone()).with_tools(tool_schemas.clone()); + if let Some(format) = &self.policy.default_response_format { + request = request.with_response_format(format.clone()); + } + if let Some(cap) = ctx.config.max_turn_output_tokens { + request.max_tokens = + Some(request.max_tokens.map_or(cap, |current| current.min(cap))); + } + + status.mark_running(HarnessPhase::Middleware); + self.middleware + .run_before_model(ctx, state, &mut request) + .await?; + + // Resolve the model for the event/log name before invoking. + let binding = self + .models + .resolve_request(&request, None, None) + .ok_or_else(|| { + TinyAgentsError::ModelNotFound( + request + .model + .clone() + .unwrap_or_else(|| "".to_string()), + ) + })?; + let model_name = binding.resolved.name.clone(); + + // Resolve the structured-output plan against the resolved model. + // `Auto` consults the model profile to choose provider-native schema + // mode versus a tool-call fallback; an explicit `JsonSchema` always + // uses provider-native mode. The chosen strategy drives extraction of + // the final response below. + let structured_plan: Option<(StructuredStrategy, String, Value)> = + match request.response_format.clone() { + Some(ResponseFormat::Auto { name, schema }) => { + let strategy = StructuredStrategy::for_profile(binding.model.profile()); + match strategy { + StructuredStrategy::ProviderSchema => { + request.response_format = + Some(ResponseFormat::json_schema(name.clone(), schema.clone())); + } + StructuredStrategy::ToolCall => { + request.response_format = Some(ResponseFormat::Text); + request.tools.push(ToolSchema { + name: name.clone(), + description: format!("Return the result as `{name}`."), + parameters: schema.clone(), + format: crate::harness::tool::ToolFormat::Json, + }); + request.tool_choice = ToolChoice::Tool(name.clone()); + } + } + Some((strategy, name, schema)) + } + Some(ResponseFormat::JsonSchema { name, schema }) => { + Some((StructuredStrategy::ProviderSchema, name, schema)) + } + _ => None, + }; + + let call_id = CallId::new(format!("{}-model-{}", ctx.run_id(), run.model_calls + 1)); + status.mark_running(HarnessPhase::Model); + status.active_model_call = Some(call_id.clone()); + let record = ctx.emit(AgentEvent::ModelStarted { + call_id: call_id.clone(), + model: model_name, + }); + status.set_last_event(record.id); + + // The real model call (cache + retry + fallback core) is the + // innermost base of the model-wrap onion. Lifecycle `before_model` + // already ran above; the wrap onion runs here; lifecycle + // `after_model` runs below — so ordering is: + // before_model -> wrap onion (outer..inner..base) -> after_model. + let base = ModelCallBase { + harness: self, + call_id: call_id.clone(), + resolved: binding.resolved, + model: binding.model, + streaming, + }; + // Snapshot the request messages for observability before `request` + // is moved into the model-wrap onion, gated by the capture policy so + // payload-free runs never serialize prompt text. + let captured_input = self + .policy + .capture + .model_io + .then(|| serde_json::to_value(&request.messages).unwrap_or(Value::Null)); + let mut response = self + .middleware + .run_wrapped_model(ctx, state, request, &base) + .await? + .into_response(); + + status.mark_running(HarnessPhase::Middleware); + self.middleware + .run_after_model(ctx, state, &mut response) + .await?; + + // Accounting. + run.model_calls += 1; + run.steps += 1; + status.model_calls = run.model_calls; + status.active_model_call = None; + if let Some(usage) = response.usage { + run.usage.record(usage); + status.usage = run.usage; + let record = ctx.emit(AgentEvent::UsageRecorded { usage }); + status.set_last_event(record.id); + } + let captured_output = self + .policy + .capture + .model_io + .then(|| serde_json::to_value(&response.message).unwrap_or(Value::Null)); + let record = ctx.emit(AgentEvent::ModelCompleted { + call_id, + usage: response.usage, + input: captured_input, + output: captured_output, + }); + status.set_last_event(record.id); + + messages.push(Message::Assistant(response.message.clone())); + + // Safe checkpoint: honor any control outcome a middleware requested + // during this turn (for example an early-exit tool or a budget stop + // hook), before executing further tools. + if let Some(control) = ctx.take_control() { + let record = ctx.emit(AgentEvent::ControlApplied { + control: control.kind().to_string(), + detail: match &control { + MiddlewareControl::StopWithFinal(text) => text.clone(), + MiddlewareControl::Interrupt { node, message } => { + format!("{node}: {message}") + } + }, + }); + status.set_last_event(record.id); + match control { + MiddlewareControl::StopWithFinal(text) => { + run.final_response = Some(ModelResponse::assistant(text)); + break; + } + MiddlewareControl::Interrupt { node, message } => { + return Err(TinyAgentsError::Interrupted { node, message }); + } + } + } + + let tool_calls = response.tool_calls().to_vec(); + + // A tool-call structured-output strategy produces an artificial tool + // call that is not a registered tool; treat it as the final response + // rather than attempting to execute it. + let structured_tool_hit = matches!( + &structured_plan, + Some((StructuredStrategy::ToolCall, name, _)) + if tool_calls.iter().any(|c| &c.name == name) + ); + + if tool_calls.is_empty() || structured_tool_hit { + // Final response: optionally extract structured output using the + // resolved plan (provider-native schema or tool-call arguments). + if let Some((strategy, name, schema)) = &structured_plan { + let extractor = + StructuredExtractor::new(*strategy, name.clone(), schema.clone()); + let output = extractor.extract(&response)?; + run.structured = Some(output.value); + } + run.final_response = Some(response); + break; + } + + // Execute requested tools serially. + status.mark_running(HarnessPhase::Tools); + for mut call in tool_calls { + // Safe cancellation checkpoint: stop before invoking the next + // (side-effecting) tool if cancellation was requested. + if ctx.cancellation.is_cancelled() { + return Err(TinyAgentsError::Cancelled); + } + if ctx.check_deadline().is_err() { + ctx.emit(AgentEvent::LimitReached { + kind: LimitKind::WallClock, + }); + return Err(TinyAgentsError::Timeout(format!( + "run `{}` exceeded its wall-clock deadline", + ctx.run_id() + ))); + } + // The context's `LimitTracker` (synced with `RunPolicy::limits` + // above) is the single enforced source of truth for the + // tool-call cap, so the reported limit always matches the one + // that trips. + if let Err(err) = ctx.record_tool_call() { + ctx.emit(AgentEvent::LimitReached { + kind: LimitKind::ToolCalls, + }); + return Err(TinyAgentsError::LimitExceeded(err.to_string())); + } + + self.middleware + .run_before_tool(ctx, state, &mut call) + .await?; + + let tool = match self.tools.get(&call.name) { + Some(tool) => tool, + None => { + // The model called an unregistered tool. Apply the run's + // `UnknownToolPolicy` instead of unconditionally aborting. + let requested = call.name.clone(); + let arguments = call.arguments.clone(); + let call_id = CallId::new(call.id.clone()); + + // Rewrite mode: retarget to a fixed compatibility tool if + // that tool exists, otherwise fall through to recovery. + let rewrite_target = match &self.policy.unknown_tool { + UnknownToolPolicy::Rewrite { tool_name } => { + self.tools.get(tool_name).map(|t| (tool_name.clone(), t)) + } + _ => None, + }; + + if let Some((tool_name, tool)) = rewrite_target { + call.name = tool_name.clone(); + let record = ctx.emit(AgentEvent::UnknownToolCall { + call_id, + requested_name: requested, + arguments, + recovery: format!("rewrite:{tool_name}"), + }); + status.set_last_event(record.id); + tool + } else if matches!(self.policy.unknown_tool, UnknownToolPolicy::Fail) { + return Err(TinyAgentsError::ToolNotFound(requested)); + } else { + // `ReturnToolError` (or a Rewrite whose target is also + // missing): inject a tool-error result naming the + // requested tool and the valid tools, then continue so + // the model can correct itself. This consumed one + // tool-call budget slot above, bounding the loop. + let valid = self.tools.names().join(", "); + let args_repr = serde_json::to_string(&arguments) + .unwrap_or_else(|_| "".to_string()); + let message = format!( + "unknown tool `{requested}` (arguments: {args_repr}); \ + valid tools: [{valid}]" + ); + let record = ctx.emit(AgentEvent::UnknownToolCall { + call_id, + requested_name: requested.clone(), + arguments, + recovery: "tool_error".to_string(), + }); + status.set_last_event(record.id); + run.tool_calls += 1; + status.tool_calls = run.tool_calls; + messages.push(Message::tool(call.id.clone(), message)); + continue; + } + } + }; + tool.schema().validate_call(&call)?; + + let tool_call_id = CallId::new(call.id.clone()); + let tool_name = call.name.clone(); + status.active_tool_calls.push(tool_call_id.clone()); + let record = ctx.emit(AgentEvent::ToolStarted { + call_id: tool_call_id.clone(), + tool_name: tool_name.clone(), + }); + status.set_last_event(record.id); + + // Snapshot the arguments for observability before `call` is + // moved into the tool-wrap onion, gated by the capture policy. + let captured_input = self.policy.capture.tool_io.then(|| call.arguments.clone()); + + // The real tool call is the innermost base of the tool-wrap + // onion (same before -> wrap -> after ordering as the model + // path): lifecycle `before_tool` ran above, the wrap onion runs + // here, and lifecycle `after_tool` runs below. Bounded by the + // same remaining wall-clock budget as a model call, so a + // hanging tool cannot block the run past its deadline either. + let base = ToolCallBase { tool }; + let remaining = self.call_budget(ctx); + let run_id = ctx.run_id().as_str().to_string(); + let fut = self.middleware.run_wrapped_tool(ctx, state, call, &base); + let mut result = Self::with_call_budget(remaining, &run_id, "tool call", fut) + .await? + .into_result(); + + self.middleware + .run_after_tool(ctx, state, &mut result) + .await?; + + run.tool_calls += 1; + status.tool_calls = run.tool_calls; + status.active_tool_calls.retain(|c| c != &tool_call_id); + let captured_output = self + .policy + .capture + .tool_io + .then(|| Value::String(result.content.clone())); + let record = ctx.emit(AgentEvent::ToolCompleted { + call_id: tool_call_id, + tool_name, + input: captured_input, + output: captured_output, + }); + status.set_last_event(record.id); + + messages.push(Message::tool( + result.call_id.clone(), + result.content.clone(), + )); + } + } + + run.messages = messages; + + status.mark_running(HarnessPhase::Middleware); + self.middleware.run_after_agent(ctx, state, run).await?; + + let record = ctx.emit(AgentEvent::RunCompleted { + run_id: ctx.run_id().clone(), + }); + status.set_last_event(record.id); + + Ok(()) + } + + /// Resolves the effective response-cache decision for `request`. + /// + /// Returns `Some((cache, key))` when a [`ResponseCache`] is attached to the + /// harness *and* caching is enabled for this call. The per-request + /// [`ModelRequest::cache_policy`] takes precedence over the harness-level + /// [`RunPolicy::cache`][crate::harness::runtime::RunPolicy]; when the request + /// carries no policy the run policy's + /// [`response_cache_enabled`][crate::harness::cache::CachePolicy] decides. + /// Returns `None` (caching disabled) when no cache is attached or the + /// effective policy disables it. + pub(super) fn response_cache_decision( + &self, + request: &ModelRequest, + ) -> Option<(Arc, String)> { + let cache = self.response_cache.as_ref()?; + let enabled = match &request.cache_policy { + Some(policy) => policy.response_cache_enabled, + None => self.policy.cache.response_cache_enabled, + }; + if !enabled { + return None; + } + // Skip caching multi-turn requests. Once the transcript contains a prior + // assistant turn (or tool result), every subsequent call carries a + // unique history and can never be re-served, so caching it only pays the + // hashing/serialization cost and grows the cache with dead entries. The + // first, history-free call is the only reusable one. + if request + .messages + .iter() + .any(|m| matches!(m, Message::Assistant(_) | Message::Tool(_))) + { + return None; + } + Some((Arc::clone(cache), cache_key(request))) + } +} diff --git a/src/harness/agent_loop/test.rs b/src/harness/agent_loop/test.rs index 30db8bf..7fd452e 100644 --- a/src/harness/agent_loop/test.rs +++ b/src/harness/agent_loop/test.rs @@ -14,13 +14,14 @@ use serde_json::json; use crate::error::{Result, TinyAgentsError}; use crate::harness::context::{RunConfig, RunContext}; use crate::harness::limits::RunLimits; -use crate::harness::message::{AssistantMessage, ContentBlock, Message}; +use crate::harness::message::{AssistantMessage, ContentBlock, Message, MessageDelta}; use crate::harness::middleware::{ AgentRun, Middleware, MiddlewareModelOutcome, MiddlewareToolOutcome, ModelHandler, ModelMiddleware, ToolHandler, ToolMiddleware, }; use crate::harness::model::{ - ChatModel, ModelProfile, ModelRequest, ModelResponse, ResponseFormat, ToolChoice, + ChatModel, ModelProfile, ModelRequest, ModelResponse, ModelStreamItem, ResponseFormat, + ToolChoice, }; use crate::harness::providers::MockModel; use crate::harness::retry::{FallbackPolicy, RetryPolicy}; @@ -64,6 +65,30 @@ impl Tool<()> for FakeTool { } } +/// A tool that sleeps `delay` before returning a fixed reply, used to prove a +/// hanging tool call is bounded by the run's remaining wall-clock budget the +/// same way a hanging model call is. +struct SlowTool { + delay: std::time::Duration, +} + +#[async_trait] +impl Tool<()> for SlowTool { + fn name(&self) -> &str { + "slow" + } + fn description(&self) -> &str { + "slow tool" + } + fn schema(&self) -> ToolSchema { + ToolSchema::new("slow", "slow tool", json!({"type": "object"})) + } + async fn call(&self, _state: &(), call: ToolCall) -> Result { + tokio::time::sleep(self.delay).await; + Ok(ToolResult::text(call.id, "slow", "too late")) + } +} + /// A strict tool used to prove harness-level schema validation runs before the /// tool implementation is invoked. struct StrictLookupTool { @@ -216,6 +241,51 @@ impl ChatModel<()> for FailingModel { } } +/// A model that always fails with a retryable error and records the +/// (virtual) `tokio::time::Instant` of each `invoke` call, so a test can +/// assert on the actual elapsed time between retries rather than just the +/// attempt count. +struct TimestampingFailingModel { + timestamps: Mutex>, +} + +#[async_trait] +impl ChatModel<()> for TimestampingFailingModel { + async fn invoke(&self, _state: &(), _request: ModelRequest) -> Result { + self.timestamps + .lock() + .unwrap() + .push(tokio::time::Instant::now()); + Err(TinyAgentsError::Model("transient boom".to_string())) + } +} + +/// A model that always fails with a structured `TinyAgentsError::Provider` +/// error whose `retryable` flag is fixed at construction, and counts +/// attempts. Used to prove the agent loop's retry decision consults the +/// structured flag rather than retrying every provider failure. +struct ProviderFailingModel { + retryable: bool, + status: u16, + attempts: Mutex, +} + +#[async_trait] +impl ChatModel<()> for ProviderFailingModel { + async fn invoke(&self, _state: &(), _request: ModelRequest) -> Result { + *self.attempts.lock().unwrap() += 1; + Err(TinyAgentsError::Provider(Box::new( + crate::harness::model::ProviderError { + provider: "test-provider".to_string(), + status: Some(self.status), + retryable: self.retryable, + message: "boom".to_string(), + ..crate::harness::model::ProviderError::default() + }, + ))) + } +} + /// Around-model wrap middleware that calls the inner pipeline then stamps the /// finish reason on the resulting response. struct StampModelWrap; @@ -399,6 +469,42 @@ async fn max_model_calls_limit_triggers_limit_exceeded() { ); } +#[tokio::test] +async fn policy_model_call_limit_above_run_config_default_is_honored() { + // Regression test: `RunConfig::new` defaults `max_model_calls` to 25, but + // a harness-wide `RunPolicy` can configure a higher cap. Before the two + // limit sources were unified, the context's tracker (seeded from the + // `RunConfig` default) tripped at call 26 while the error message + // incorrectly reported the policy's higher limit. This asserts the run + // survives past 25 calls and, once it does trip, reports the limit that + // actually applies. + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness.register_model( + "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(30) + .with_max_tool_calls(1000), + ..RunPolicy::default() + }); + + let err = harness + .invoke_default(&(), vec![Message::user("go")]) + .await + .expect_err("limit should be exceeded"); + assert!( + matches!(err, TinyAgentsError::LimitExceeded(_)), + "got {err:?}" + ); + assert!( + err.to_string().contains("30"), + "expected error to report the policy's limit (30), got: {err}" + ); +} + #[tokio::test] async fn max_tool_calls_limit_triggers_limit_exceeded() { let mut harness: AgentHarness<()> = AgentHarness::new(); @@ -655,6 +761,56 @@ async fn no_model_registered_errors() { ); } +#[tokio::test(start_paused = true)] +async fn retry_backoff_sleeps_the_documented_schedule() { + // Regression test: the loop used to compute the backoff from the + // *post-increment* attempt number, so the first retry's sleep skipped + // `initial_backoff_ms` entirely and the whole exponential schedule was + // shifted one step higher than `RetryPolicy::backoff_for_attempt` + // documents. With `initial_backoff_ms = 100`, `multiplier = 2.0`, no + // jitter: attempt 0 -> 100ms, attempt 1 -> 200ms, attempt 2 -> 400ms. + use std::time::Duration; + + let policy = RetryPolicy::default() + .with_max_attempts(4) + .with_initial_backoff_ms(100) + .with_multiplier(2.0) + .with_jitter(false) + .with_backoff_sleep(true); + + let model = Arc::new(TimestampingFailingModel { + timestamps: Mutex::new(Vec::new()), + }); + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness.register_model("flaky", model.clone()); + harness.with_policy(RunPolicy { + retry: policy, + ..RunPolicy::default() + }); + + harness + .invoke_default(&(), vec![Message::user("hi")]) + .await + .expect_err("all 4 attempts fail"); + + let timestamps = model.timestamps.lock().unwrap().clone(); + assert_eq!(timestamps.len(), 4, "expected exactly max_attempts calls"); + + let gaps: Vec = timestamps + .windows(2) + .map(|w| w[1].duration_since(w[0])) + .collect(); + assert_eq!( + gaps, + vec![ + Duration::from_millis(100), // before retry 1 (attempt 0's backoff) + Duration::from_millis(200), // before retry 2 (attempt 1's backoff) + Duration::from_millis(400), // before retry 3 (attempt 2's backoff) + ], + "backoff schedule does not match RetryPolicy::backoff_for_attempt" + ); +} + #[tokio::test] async fn retry_then_fallback_succeeds() { let mut harness: AgentHarness<()> = AgentHarness::new(); @@ -680,6 +836,80 @@ async fn retry_then_fallback_succeeds() { assert_eq!(*failing.attempts.lock().unwrap(), 2); } +#[tokio::test] +async fn run_limits_max_retries_per_call_caps_a_looser_retry_policy() { + // Regression test: `RunLimits::max_retries_per_call` was parsed but never + // enforced, so a `RetryPolicy` with a higher `max_attempts` silently + // ignored the harness's "hard" limit. `max_retries_per_call: 1` (one + // retry, so 2 attempts total) must win over `max_attempts: 5`. + let mut harness: AgentHarness<()> = AgentHarness::new(); + let failing = Arc::new(FailingModel { + attempts: Mutex::new(0), + }); + harness.register_model("primary", failing.clone()); + harness.with_policy(RunPolicy { + retry: RetryPolicy::default().with_max_attempts(5), + limits: RunLimits::default().with_max_retries_per_call(1), + ..RunPolicy::default() + }); + + let err = harness + .invoke_default(&(), vec![Message::user("hi")]) + .await + .expect_err("no fallback, retries capped by RunLimits"); + assert!(matches!(err, TinyAgentsError::Model(_)), "got {err:?}"); + assert_eq!(*failing.attempts.lock().unwrap(), 2); +} + +#[tokio::test] +async fn provider_error_401_is_not_retried() { + // Regression test: before `ProviderError` was preserved structurally, a + // 401 flattened into `Model(String)` was retried like any other model + // error. A non-retryable `Provider` error must fail on the first attempt. + let mut harness: AgentHarness<()> = AgentHarness::new(); + let model = Arc::new(ProviderFailingModel { + retryable: false, + status: 401, + attempts: Mutex::new(0), + }); + harness.register_model("primary", model.clone()); + harness.with_policy(RunPolicy { + retry: RetryPolicy::default().with_max_attempts(5), + ..RunPolicy::default() + }); + + let err = harness + .invoke_default(&(), vec![Message::user("hi")]) + .await + .expect_err("401 is not retryable"); + assert!(matches!(err, TinyAgentsError::Provider(_)), "got {err:?}"); + assert_eq!(*model.attempts.lock().unwrap(), 1); +} + +#[tokio::test] +async fn provider_error_429_is_retried_up_to_max_attempts() { + // Contrast with the 401 case: a retryable `Provider` error (e.g. a 429) + // must still be retried up to `max_attempts`. + let mut harness: AgentHarness<()> = AgentHarness::new(); + let model = Arc::new(ProviderFailingModel { + retryable: true, + status: 429, + attempts: Mutex::new(0), + }); + harness.register_model("primary", model.clone()); + harness.with_policy(RunPolicy { + retry: RetryPolicy::default().with_max_attempts(3), + ..RunPolicy::default() + }); + + let err = harness + .invoke_default(&(), vec![Message::user("hi")]) + .await + .expect_err("retries exhausted"); + assert!(matches!(err, TinyAgentsError::Provider(_)), "got {err:?}"); + assert_eq!(*model.attempts.lock().unwrap(), 3); +} + #[tokio::test] async fn non_retryable_or_exhausted_without_fallback_errors() { let mut harness: AgentHarness<()> = AgentHarness::new(); @@ -701,6 +931,43 @@ async fn non_retryable_or_exhausted_without_fallback_errors() { assert!(matches!(err, TinyAgentsError::Model(_)), "got {err:?}"); } +#[tokio::test] +async fn fallback_chain_with_repeated_model_name_terminates() { + // Regression test: a fallback chain that repeats a model name + // (`[primary, backup, primary]`) used to alternate primary <-> backup + // forever because `FallbackPolicy::next_after` always resolves from the + // *first* occurrence of the current name. Both models fail every call, so + // without a visited-set/hop-cap this run would never terminate. + let mut harness: AgentHarness<()> = AgentHarness::new(); + let primary = Arc::new(FailingModel { + attempts: Mutex::new(0), + }); + let backup = Arc::new(FailingModel { + attempts: Mutex::new(0), + }); + harness.register_model("primary", primary.clone()); + harness.register_model("backup", backup.clone()); + harness.with_policy(RunPolicy { + retry: RetryPolicy::default().with_max_attempts(1), + fallback: Some(FallbackPolicy::new(["primary", "backup", "primary"])), + ..RunPolicy::default() + }); + + let err = tokio::time::timeout( + std::time::Duration::from_secs(5), + harness.invoke_default(&(), vec![Message::user("hi")]), + ) + .await + .expect("fallback chain must terminate, not hang") + .expect_err("both models fail, so the run must error out"); + + assert!(matches!(err, TinyAgentsError::Model(_)), "got {err:?}"); + // Each model is visited at most once: primary, then backup, then the + // chain's repeated `primary` entry is skipped as already-visited. + assert_eq!(*primary.attempts.lock().unwrap(), 1); + assert_eq!(*backup.attempts.lock().unwrap(), 1); +} + #[tokio::test] async fn invoke_with_status_reports_completed() { use crate::harness::ids::{ExecutionStatus, HarnessPhase}; @@ -734,6 +1001,7 @@ fn agent_run_default_is_empty() { struct DeltaRecorder { count: Arc>, texts: Arc>>, + reasonings: Arc>>, } #[async_trait] @@ -749,6 +1017,10 @@ impl Middleware<(), ()> for DeltaRecorder { ) -> Result<()> { *self.count.lock().unwrap() += 1; self.texts.lock().unwrap().push(delta.content.clone()); + self.reasonings + .lock() + .unwrap() + .push(delta.reasoning.clone()); Ok(()) } } @@ -759,6 +1031,7 @@ async fn invoke_streaming_fires_on_model_delta_per_delta_and_accumulates() { let count = Arc::new(Mutex::new(0usize)); let texts = Arc::new(Mutex::new(Vec::new())); + let reasonings = Arc::new(Mutex::new(Vec::new())); let mut harness: AgentHarness<()> = AgentHarness::new(); harness.register_model( @@ -768,6 +1041,7 @@ async fn invoke_streaming_fires_on_model_delta_per_delta_and_accumulates() { harness.push_middleware(Arc::new(DeltaRecorder { count: count.clone(), texts: texts.clone(), + reasonings: reasonings.clone(), })); let run = harness @@ -790,6 +1064,64 @@ async fn invoke_streaming_fires_on_model_delta_per_delta_and_accumulates() { *texts.lock().unwrap(), vec!["Hel".to_string(), "lo, ".to_string(), "world".to_string()] ); + assert_eq!( + *reasonings.lock().unwrap(), + vec![String::new(), String::new(), String::new()] + ); +} + +#[tokio::test] +async fn invoke_streaming_forwards_reasoning_deltas_to_middleware_and_events() { + use crate::harness::testkit::{EventRecorder, StreamingMock}; + + let count = Arc::new(Mutex::new(0usize)); + let texts = Arc::new(Mutex::new(Vec::new())); + let reasonings = Arc::new(Mutex::new(Vec::new())); + + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness.register_model( + "stream", + Arc::new(StreamingMock::new(vec![ + ModelStreamItem::Started, + ModelStreamItem::MessageDelta(MessageDelta::reasoning("think ")), + ModelStreamItem::MessageDelta(MessageDelta::text("answer")), + ModelStreamItem::Completed(ModelResponse::assistant("answer")), + ])), + ); + harness.push_middleware(Arc::new(DeltaRecorder { + count: count.clone(), + texts: texts.clone(), + reasonings: reasonings.clone(), + })); + + let recorder = EventRecorder::new(); + let ctx = RunContext::new(RunConfig::new("stream-run"), ()).with_events(recorder.sink()); + + let run = harness + .invoke_streaming_in_context(&(), ctx, vec![Message::user("hi")]) + .await + .expect("streaming run succeeds"); + + assert_eq!(run.text(), Some("answer".to_string())); + assert_eq!(*count.lock().unwrap(), 2); + assert_eq!( + *texts.lock().unwrap(), + vec![String::new(), "answer".to_string()] + ); + assert_eq!( + *reasonings.lock().unwrap(), + vec!["think ".to_string(), String::new()] + ); + + let event_reasoning: String = recorder + .events() + .into_iter() + .filter_map(|event| match event { + crate::harness::events::AgentEvent::ModelDelta { delta, .. } => Some(delta.reasoning), + _ => None, + }) + .collect(); + assert_eq!(event_reasoning, "think "); } #[tokio::test] @@ -997,6 +1329,33 @@ async fn slow_streaming_model_call_is_timed_out() { assert!(matches!(err, TinyAgentsError::Timeout(_)), "got {err:?}"); } +#[tokio::test] +async fn slow_tool_call_is_timed_out_by_remaining_budget() { + use std::time::Duration; + + // Regression test: the remaining wall-clock budget was previously only + // enforced around model calls, so a hanging tool call could block the run + // past its deadline. The tool sleeps far longer (200ms) than the run's + // 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( + "mock", + Arc::new(MockModel::with_tool_call("slow", json!({}))), + ); + harness.register_tool(Arc::new(SlowTool { + delay: Duration::from_millis(200), + })); + + let config = RunConfig::new("tool-timeout-run").with_timeout_ms(20); + let err = harness + .invoke(&(), (), config, vec![Message::user("go")]) + .await + .expect_err("a tool call slower than the budget must time out"); + + assert!(matches!(err, TinyAgentsError::Timeout(_)), "got {err:?}"); +} + // ── Response caching ────────────────────────────────────────────────────────── #[tokio::test] @@ -1062,6 +1421,44 @@ async fn response_cache_serves_repeated_request_without_calling_model() { assert_eq!(run2.model_calls, 1); } +#[tokio::test] +async fn multi_turn_request_with_prior_assistant_turn_is_not_cached() { + use crate::harness::cache::InMemoryResponseCache; + + // A request whose transcript already contains an assistant turn can never be + // re-served identically, so it must bypass the cache entirely: the model is + // invoked on every run even for identical multi-turn input. + let model = Arc::new(MockModel::with_responses(vec![ + text_response("a1", 4, 2), + text_response("a2", 4, 2), + ])); + let cache = Arc::new(InMemoryResponseCache::new()); + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness.register_model("mock", model.clone()); + harness.with_response_cache(cache.clone()); + + let convo = vec![ + Message::user("q1"), + Message::assistant("prior answer"), + Message::user("q2"), + ]; + + harness + .invoke_default(&(), convo.clone()) + .await + .expect("first run succeeds"); + harness + .invoke_default(&(), convo) + .await + .expect("second run succeeds"); + + assert_eq!( + model.call_count(), + 2, + "multi-turn requests bypass the cache, so the model runs each time" + ); +} + #[tokio::test] async fn no_cache_attached_invokes_model_each_run() { // Control: without a cache the model is invoked on every run. diff --git a/src/harness/cache/mod.rs b/src/harness/cache/mod.rs index 1b461e0..df57561 100644 --- a/src/harness/cache/mod.rs +++ b/src/harness/cache/mod.rs @@ -36,13 +36,6 @@ use crate::harness::model::{ModelRequest, ModelResponse}; // ── Deterministic hash ──────────────────────────────────────────────────────── -/// Computes a deterministic SHA-256 digest over `data` and returns it as a -/// 64-character lowercase hex string. -fn sha256_hex(data: &[u8]) -> String { - let digest = Sha256::digest(data); - digest.iter().map(|byte| format!("{byte:02x}")).collect() -} - /// Computes a deterministic FNV-1a 64-bit hash over `data` and returns it as /// a 16-character lowercase hex string. /// @@ -101,35 +94,92 @@ fn canonical_value(v: Value) -> Value { pub fn cache_key(request: &ModelRequest) -> String { let value = serde_json::to_value(request).unwrap_or(Value::Null); let canonical = canonical_value(value); - let bytes = serde_json::to_vec(&canonical).unwrap_or_default(); - sha256_hex(&bytes) + // Stream the canonical JSON directly into the digest instead of allocating + // an intermediate `Vec`: `Sha256` implements `std::io::Write`. + let mut hasher = Sha256::new(); + if serde_json::to_writer(&mut hasher, &canonical).is_err() { + // Serialization of an already-materialized `Value` does not fail in + // practice; fall back to hashing empty data for a stable result. + hasher = Sha256::new(); + } + let digest = hasher.finalize(); + digest.iter().map(|byte| format!("{byte:02x}")).collect() } // ── InMemoryResponseCache ───────────────────────────────────────────────────── impl InMemoryResponseCache { - /// Creates a new, empty in-memory response cache. + /// Default LRU capacity when constructed via [`new`](Self::new) or + /// [`Default`]. + pub const DEFAULT_CAPACITY: usize = 1024; + + /// Creates a new, empty in-memory response cache bounded by + /// [`DEFAULT_CAPACITY`](Self::DEFAULT_CAPACITY) entries. pub fn new() -> Self { - Self::default() + Self::with_capacity(Self::DEFAULT_CAPACITY) + } + + /// Creates a new, empty in-memory response cache retaining at most + /// `capacity` entries (least-recently-used evicted first). A `capacity` of + /// zero is treated as `1` so the cache always retains the last write. + pub fn with_capacity(capacity: usize) -> Self { + Self { + inner: std::sync::Arc::new(std::sync::Mutex::new(LruResponseMap { + data: std::collections::HashMap::new(), + order: std::collections::VecDeque::new(), + capacity: capacity.max(1), + })), + } + } +} + +impl Default for InMemoryResponseCache { + fn default() -> Self { + Self::new() + } +} + +impl LruResponseMap { + /// Moves `key` to the most-recently-used end of the order queue. + fn touch(&mut self, key: &str) { + if let Some(pos) = self.order.iter().position(|k| k == key) { + let k = self.order.remove(pos).expect("position is valid"); + self.order.push_back(k); + } } } #[async_trait] impl ResponseCache for InMemoryResponseCache { async fn get(&self, key: &str) -> Result> { - let data = self - .data + let mut inner = self + .inner .lock() .map_err(|e| TinyAgentsError::Validation(format!("cache lock poisoned: {e}")))?; - Ok(data.get(key).cloned()) + let hit = inner.data.get(key).cloned(); + if hit.is_some() { + inner.touch(key); + } + Ok(hit) } async fn put(&self, key: &str, value: ModelResponse) -> Result<()> { - let mut data = self - .data + let mut inner = self + .inner .lock() .map_err(|e| TinyAgentsError::Validation(format!("cache lock poisoned: {e}")))?; - data.insert(key.to_string(), value); + if inner.data.insert(key.to_string(), value).is_some() { + // Existing key: refresh its recency without changing the length. + inner.touch(key); + } else { + inner.order.push_back(key.to_string()); + // Evict least-recently-used entries until within capacity. + while inner.order.len() > inner.capacity { + if let Some(evicted) = inner.order.pop_front() { + inner.data.remove(&evicted); + } + } + } Ok(()) } } diff --git a/src/harness/cache/test.rs b/src/harness/cache/test.rs index 80807c0..eb2ff0d 100644 --- a/src/harness/cache/test.rs +++ b/src/harness/cache/test.rs @@ -18,6 +18,35 @@ async fn response_cache_put_get() { assert_eq!(fetched.text(), "hello"); } +#[tokio::test] +async fn response_cache_evicts_least_recently_used() { + let cache = InMemoryResponseCache::with_capacity(2); + cache.put("a", ModelResponse::assistant("a")).await.unwrap(); + cache.put("b", ModelResponse::assistant("b")).await.unwrap(); + + // Touch "a" so "b" becomes the least-recently-used entry. + assert!(cache.get("a").await.unwrap().is_some()); + + // Inserting a third key evicts "b", not the recently-read "a". + cache.put("c", ModelResponse::assistant("c")).await.unwrap(); + assert!( + cache.get("a").await.unwrap().is_some(), + "a was recently used" + ); + assert!(cache.get("c").await.unwrap().is_some(), "c is newest"); + assert!(cache.get("b").await.unwrap().is_none(), "b evicted as LRU"); +} + +#[tokio::test] +async fn response_cache_capacity_zero_retains_last() { + // A zero capacity is clamped to 1: the cache still retains the last write. + let cache = InMemoryResponseCache::with_capacity(0); + cache.put("x", ModelResponse::assistant("x")).await.unwrap(); + cache.put("y", ModelResponse::assistant("y")).await.unwrap(); + assert!(cache.get("x").await.unwrap().is_none()); + assert!(cache.get("y").await.unwrap().is_some()); +} + #[test] fn cache_key_is_deterministic() { let req = ModelRequest::new(vec![]).with_model("gpt-4"); diff --git a/src/harness/cache/types.rs b/src/harness/cache/types.rs index 5724156..11be36f 100644 --- a/src/harness/cache/types.rs +++ b/src/harness/cache/types.rs @@ -18,7 +18,7 @@ //! //! All public types in this module are re-exported through [`super`]. -use std::collections::HashMap; +use std::collections::{HashMap, VecDeque}; use std::sync::{Arc, Mutex}; use async_trait::async_trait; @@ -47,9 +47,26 @@ pub trait ResponseCache: Send + Sync { /// /// Intended for unit tests and short-lived local runs. Contains no durable /// storage: all entries are lost when the value is dropped. -#[derive(Clone, Debug, Default)] +/// +/// Entries are bounded by an LRU eviction policy (default +/// [`InMemoryResponseCache::DEFAULT_CAPACITY`]) so a long-lived cache attached +/// to a busy harness cannot grow without limit. Reads and writes move a key to +/// the most-recently-used end; once the map is full the least-recently-used key +/// is evicted on insert. +#[derive(Clone, Debug)] pub struct InMemoryResponseCache { - pub(crate) data: Arc>>, + pub(crate) inner: Arc>, +} + +/// LRU-ordered map backing [`InMemoryResponseCache`]. +#[derive(Debug)] +pub(crate) struct LruResponseMap { + /// Cached responses keyed by cache key. + pub(crate) data: HashMap, + /// Keys in least- to most-recently-used order. + pub(crate) order: VecDeque, + /// Maximum number of entries retained before LRU eviction. + pub(crate) capacity: usize, } // ── PromptCacheLayout ───────────────────────────────────────────────────────── diff --git a/src/harness/context/mod.rs b/src/harness/context/mod.rs index 7287653..904f4a6 100644 --- a/src/harness/context/mod.rs +++ b/src/harness/context/mod.rs @@ -124,6 +124,22 @@ impl RunConfig { self } + /// Derives a child run's depth from its parent, enforcing the recursion cap. + /// + /// The single source of truth for the sub-agent depth guard: returns + /// `parent_depth + 1`, or [`crate::error::TinyAgentsError::SubAgentDepth`] + /// carrying `max_depth` when the child would exceed the cap. Every recursion + /// surface — [`crate::harness::subagent::SubAgent`], its reuse-session tool, + /// and the REPL sub-run builtin — funnels its `depth + 1` check through here + /// so the fail-closed guard cannot drift out of sync between them. + pub fn checked_child_depth(parent_depth: usize, max_depth: usize) -> Result { + let child_depth = parent_depth + 1; + if child_depth > max_depth { + return Err(crate::error::TinyAgentsError::SubAgentDepth(max_depth)); + } + Ok(child_depth) + } + /// Builds the [`RunConfig`] for a child run one level deeper than this one. /// /// The returned config keeps this config's `max_depth` and thread, sets diff --git a/src/harness/context/test.rs b/src/harness/context/test.rs index 15a48f1..ffe30b7 100644 --- a/src/harness/context/test.rs +++ b/src/harness/context/test.rs @@ -141,3 +141,18 @@ fn request_control_keeps_highest_precedence() { Some(MiddlewareControl::Interrupt { .. }) )); } + +#[test] +fn checked_child_depth_is_the_shared_depth_guard() { + use crate::error::TinyAgentsError; + + // Below the cap: returns parent_depth + 1. + assert_eq!(RunConfig::checked_child_depth(0, 8).unwrap(), 1); + assert_eq!(RunConfig::checked_child_depth(7, 8).unwrap(), 8); + + // At the boundary (child would be max_depth + 1): fail closed with the cap. + match RunConfig::checked_child_depth(8, 8) { + Err(TinyAgentsError::SubAgentDepth(cap)) => assert_eq!(cap, 8), + other => panic!("expected SubAgentDepth(8), got {other:?}"), + } +} diff --git a/src/harness/cost/mod.rs b/src/harness/cost/mod.rs index b85d2f7..8184ae5 100644 --- a/src/harness/cost/mod.rs +++ b/src/harness/cost/mod.rs @@ -55,12 +55,22 @@ impl AddAssign for CostTotals { /// Missing prices contribute zero. Cache read and cache creation tokens are /// priced independently when the catalog provides those rates and folded into /// `cache_cost`. +/// +/// Providers report `cache_read_tokens` as a *subset* of `input_tokens` (and +/// `reasoning_tokens` as a subset of `output_tokens`), not an addition to +/// them. Pricing the full `input_tokens`/`output_tokens` count at the +/// standard rate *and* separately pricing the cached/reasoning subset would +/// double-charge those tokens, so the standard-rate cost is computed on the +/// non-cached/non-reasoning remainder only. pub fn estimate_cost(pricing: &ModelPricing, usage: &Usage) -> CostTotals { let price = |rate: Option, tokens: u64| rate.unwrap_or(0.0) * tokens as f64; + let billable_input_tokens = usage.input_tokens.saturating_sub(usage.cache_read_tokens); + let billable_output_tokens = usage.output_tokens.saturating_sub(usage.reasoning_tokens); + let mut totals = CostTotals { - input_cost: price(pricing.input_per_token, usage.input_tokens), - output_cost: price(pricing.output_per_token, usage.output_tokens), + input_cost: price(pricing.input_per_token, billable_input_tokens), + output_cost: price(pricing.output_per_token, billable_output_tokens), cache_cost: price(pricing.cache_read_input_per_token, usage.cache_read_tokens) + price( pricing.cache_creation_input_per_token, diff --git a/src/harness/cost/test.rs b/src/harness/cost/test.rs index 2e6dd25..c75c037 100644 --- a/src/harness/cost/test.rs +++ b/src/harness/cost/test.rs @@ -31,12 +31,35 @@ fn estimates_each_component() { reasoning_tokens: 10, }; let cost = estimate_cost(&pricing(), &usage); - assert!((cost.input_cost - 1.0).abs() < 1e-9); - assert!((cost.output_cost - 1.0).abs() < 1e-9); + // cache_read_tokens (100) and reasoning_tokens (10) are subsets of + // input_tokens/output_tokens, so the standard-rate cost only applies to + // the non-cached/non-reasoning remainder: (1000-100)*0.001 = 0.9, + // (500-10)*0.002 = 0.98. + assert!((cost.input_cost - 0.9).abs() < 1e-9); + assert!((cost.output_cost - 0.98).abs() < 1e-9); // 100*0.0001 + 200*0.0005 = 0.01 + 0.1 = 0.11 assert!((cost.cache_cost - 0.11).abs() < 1e-9); assert!((cost.reasoning_cost - 0.03).abs() < 1e-9); - assert!((cost.total_cost - 2.14).abs() < 1e-9); + assert!((cost.total_cost - 2.02).abs() < 1e-9); +} + +#[test] +fn cache_and_reasoning_tokens_are_not_double_counted() { + // A fully-cached prompt: all 90k input tokens were served from cache. + // Charging the full input_tokens at the standard rate *and* the + // cache_read_tokens at the cache rate would double-bill the cached + // tokens; only the cache rate should apply to them. + let usage = Usage { + input_tokens: 90_000, + output_tokens: 100, + total_tokens: 90_100, + cache_read_tokens: 90_000, + cache_creation_tokens: 0, + reasoning_tokens: 0, + }; + let cost = estimate_cost(&pricing(), &usage); + assert_eq!(cost.input_cost, 0.0); + assert!((cost.cache_cost - 9.0).abs() < 1e-9); } #[test] diff --git a/src/harness/embeddings/mod.rs b/src/harness/embeddings/mod.rs index d882795..a113628 100644 --- a/src/harness/embeddings/mod.rs +++ b/src/harness/embeddings/mod.rs @@ -212,185 +212,7 @@ impl Retriever { } } -// ── OpenAiEmbeddingModel ────────────────────────────────────────────────────── - -mod openai { - use async_trait::async_trait; - use serde_json::{Value, json}; - - use super::EmbeddingModel; - use crate::error::{Result, TinyAgentsError}; - - /// Default OpenAI embedding model id. - const DEFAULT_MODEL: &str = "text-embedding-3-small"; - /// Default OpenAI API base URL (no trailing slash). - const DEFAULT_BASE_URL: &str = "https://api.openai.com/v1"; - /// Default dimensionality of `text-embedding-3-small`. - const DEFAULT_DIMENSIONS: usize = 1536; - - /// An [`EmbeddingModel`] backed by the hosted OpenAI embeddings endpoint - /// (`POST {base_url}/embeddings`). - /// - /// Construct one with [`OpenAiEmbeddingModel::new`] (plus the `with_*` - /// builders) or [`OpenAiEmbeddingModel::from_env`]. The model holds a - /// reusable [`reqwest::Client`] so repeated calls share a connection pool. - /// - /// This adapter intentionally mirrors the transport pattern of - /// [`OpenAiModel`](crate::harness::providers::openai::OpenAiModel). - /// - /// # Example - /// ```no_run - /// use tinyagents::harness::embeddings::OpenAiEmbeddingModel; - /// - /// # fn main() -> tinyagents::Result<()> { - /// let model = OpenAiEmbeddingModel::from_env()?; - /// # let _ = model; - /// # Ok(()) - /// # } - /// ``` - pub struct OpenAiEmbeddingModel { - client: reqwest::Client, - api_key: String, - model: String, - base_url: String, - dimensions: usize, - } - - impl OpenAiEmbeddingModel { - /// Creates a model with the given API key, the default model - /// (`text-embedding-3-small`), and the default base URL. - pub fn new(api_key: impl Into) -> Self { - Self { - client: reqwest::Client::new(), - api_key: api_key.into(), - model: DEFAULT_MODEL.to_string(), - base_url: DEFAULT_BASE_URL.to_string(), - dimensions: DEFAULT_DIMENSIONS, - } - } - - /// Overrides the embedding model id. - pub fn with_model(mut self, model: impl Into) -> Self { - self.model = model.into(); - self - } - - /// Overrides the API base URL; a trailing slash is trimmed so the joined - /// endpoint is always `{base_url}/embeddings`. - pub fn with_base_url(mut self, base_url: impl Into) -> Self { - self.base_url = base_url.into().trim_end_matches('/').to_string(); - self - } - - /// Overrides the reported dimensionality (and requests it from the API - /// via the `dimensions` parameter, which `text-embedding-3-*` supports). - pub fn with_dimensions(mut self, dimensions: usize) -> Self { - self.dimensions = dimensions; - self - } - - /// Builds a model from environment variables. - /// - /// Reads `OPENAI_API_KEY` (required), `OPENAI_EMBEDDING_MODEL` - /// (optional), and `OPENAI_BASE_URL` (optional). - /// - /// # Errors - /// Returns [`TinyAgentsError::Validation`] when `OPENAI_API_KEY` is - /// missing or empty. - pub fn from_env() -> Result { - let api_key = std::env::var("OPENAI_API_KEY") - .ok() - .filter(|k| !k.trim().is_empty()) - .ok_or_else(|| { - TinyAgentsError::Validation( - "OPENAI_API_KEY is not set; export it or add it to a .env file".to_string(), - ) - })?; - let mut model = Self::new(api_key); - if let Ok(name) = std::env::var("OPENAI_EMBEDDING_MODEL") - && !name.trim().is_empty() - { - model = model.with_model(name); - } - if let Ok(url) = std::env::var("OPENAI_BASE_URL") - && !url.trim().is_empty() - { - model = model.with_base_url(url); - } - Ok(model) - } - } - - #[async_trait] - impl EmbeddingModel for OpenAiEmbeddingModel { - async fn embed(&self, texts: &[String]) -> Result>> { - if texts.is_empty() { - return Ok(Vec::new()); - } - let url = format!("{}/embeddings", self.base_url); - let body = json!({ - "model": self.model, - "input": texts, - "dimensions": self.dimensions, - }); - - let response = self - .client - .post(&url) - .header("Authorization", format!("Bearer {}", self.api_key)) - .json(&body) - .send() - .await - .map_err(|e| { - TinyAgentsError::Embedding(format!( - "openai embeddings request to {url} failed: {e}" - )) - })?; - - let status = response.status(); - let text = response.text().await.map_err(|e| { - TinyAgentsError::Embedding(format!("openai embeddings body read failed: {e}")) - })?; - if !status.is_success() { - return Err(TinyAgentsError::Embedding(format!( - "openai embeddings returned HTTP {status}: {text}" - ))); - } - - let value: Value = serde_json::from_str(&text)?; - let data = value - .get("data") - .and_then(|d| d.as_array()) - .ok_or_else(|| { - TinyAgentsError::Embedding( - "openai embeddings response missing `data` array".into(), - ) - })?; - let mut vectors = Vec::with_capacity(data.len()); - for item in data { - let embedding = item - .get("embedding") - .and_then(|e| e.as_array()) - .ok_or_else(|| { - TinyAgentsError::Embedding( - "openai embeddings response missing `embedding` array".into(), - ) - })?; - vectors.push( - embedding - .iter() - .map(|n| n.as_f64().unwrap_or(0.0) as f32) - .collect(), - ); - } - Ok(vectors) - } - - fn dimensions(&self) -> usize { - self.dimensions - } - } -} +mod openai; pub use openai::OpenAiEmbeddingModel; diff --git a/src/harness/embeddings/openai.rs b/src/harness/embeddings/openai.rs new file mode 100644 index 0000000..8316965 --- /dev/null +++ b/src/harness/embeddings/openai.rs @@ -0,0 +1,179 @@ +//! [`OpenAiEmbeddingModel`]: an [`EmbeddingModel`] backed by the hosted +//! OpenAI embeddings endpoint (`POST {base_url}/embeddings`). +//! +//! Split out of `embeddings/mod.rs`; mirrors the transport pattern of +//! [`crate::harness::providers::openai::OpenAiModel`]. + +use async_trait::async_trait; +use serde_json::{Value, json}; + +use super::EmbeddingModel; +use crate::error::{Result, TinyAgentsError}; + +/// Default OpenAI embedding model id. +const DEFAULT_MODEL: &str = "text-embedding-3-small"; +/// Default OpenAI API base URL (no trailing slash). +const DEFAULT_BASE_URL: &str = "https://api.openai.com/v1"; +/// Default dimensionality of `text-embedding-3-small`. +const DEFAULT_DIMENSIONS: usize = 1536; + +/// An [`EmbeddingModel`] backed by the hosted OpenAI embeddings endpoint +/// (`POST {base_url}/embeddings`). +/// +/// Construct one with [`OpenAiEmbeddingModel::new`] (plus the `with_*` +/// builders) or [`OpenAiEmbeddingModel::from_env`]. The model holds a +/// reusable [`reqwest::Client`] so repeated calls share a connection pool. +/// +/// This adapter intentionally mirrors the transport pattern of +/// [`OpenAiModel`](crate::harness::providers::openai::OpenAiModel). +/// +/// # Example +/// ```no_run +/// use tinyagents::harness::embeddings::OpenAiEmbeddingModel; +/// +/// # fn main() -> tinyagents::Result<()> { +/// let model = OpenAiEmbeddingModel::from_env()?; +/// # let _ = model; +/// # Ok(()) +/// # } +/// ``` +pub struct OpenAiEmbeddingModel { + client: reqwest::Client, + api_key: String, + model: String, + base_url: String, + dimensions: usize, +} + +impl OpenAiEmbeddingModel { + /// Creates a model with the given API key, the default model + /// (`text-embedding-3-small`), and the default base URL. + pub fn new(api_key: impl Into) -> Self { + Self { + client: reqwest::Client::new(), + api_key: api_key.into(), + model: DEFAULT_MODEL.to_string(), + base_url: DEFAULT_BASE_URL.to_string(), + dimensions: DEFAULT_DIMENSIONS, + } + } + + /// Overrides the embedding model id. + pub fn with_model(mut self, model: impl Into) -> Self { + self.model = model.into(); + self + } + + /// Overrides the API base URL; a trailing slash is trimmed so the joined + /// endpoint is always `{base_url}/embeddings`. + pub fn with_base_url(mut self, base_url: impl Into) -> Self { + self.base_url = base_url.into().trim_end_matches('/').to_string(); + self + } + + /// Overrides the reported dimensionality (and requests it from the API + /// via the `dimensions` parameter, which `text-embedding-3-*` supports). + pub fn with_dimensions(mut self, dimensions: usize) -> Self { + self.dimensions = dimensions; + self + } + + /// Builds a model from environment variables. + /// + /// Reads `OPENAI_API_KEY` (required), `OPENAI_EMBEDDING_MODEL` + /// (optional), and `OPENAI_BASE_URL` (optional). + /// + /// # Errors + /// Returns [`TinyAgentsError::Validation`] when `OPENAI_API_KEY` is + /// missing or empty. + pub fn from_env() -> Result { + let api_key = std::env::var("OPENAI_API_KEY") + .ok() + .filter(|k| !k.trim().is_empty()) + .ok_or_else(|| { + TinyAgentsError::Validation( + "OPENAI_API_KEY is not set; export it or add it to a .env file".to_string(), + ) + })?; + let mut model = Self::new(api_key); + if let Ok(name) = std::env::var("OPENAI_EMBEDDING_MODEL") + && !name.trim().is_empty() + { + model = model.with_model(name); + } + if let Ok(url) = std::env::var("OPENAI_BASE_URL") + && !url.trim().is_empty() + { + model = model.with_base_url(url); + } + Ok(model) + } +} + +#[async_trait] +impl EmbeddingModel for OpenAiEmbeddingModel { + async fn embed(&self, texts: &[String]) -> Result>> { + if texts.is_empty() { + return Ok(Vec::new()); + } + let url = format!("{}/embeddings", self.base_url); + let body = json!({ + "model": self.model, + "input": texts, + "dimensions": self.dimensions, + }); + + let response = self + .client + .post(&url) + .header("Authorization", format!("Bearer {}", self.api_key)) + .json(&body) + .send() + .await + .map_err(|e| { + TinyAgentsError::Embedding(format!( + "openai embeddings request to {url} failed: {e}" + )) + })?; + + let status = response.status(); + let text = response.text().await.map_err(|e| { + TinyAgentsError::Embedding(format!("openai embeddings body read failed: {e}")) + })?; + if !status.is_success() { + return Err(TinyAgentsError::Embedding(format!( + "openai embeddings returned HTTP {status}: {text}" + ))); + } + + let value: Value = serde_json::from_str(&text)?; + let data = value + .get("data") + .and_then(|d| d.as_array()) + .ok_or_else(|| { + TinyAgentsError::Embedding("openai embeddings response missing `data` array".into()) + })?; + let mut vectors = Vec::with_capacity(data.len()); + for item in data { + let embedding = item + .get("embedding") + .and_then(|e| e.as_array()) + .ok_or_else(|| { + TinyAgentsError::Embedding( + "openai embeddings response missing `embedding` array".into(), + ) + })?; + vectors.push( + embedding + .iter() + .map(|n| n.as_f64().unwrap_or(0.0) as f32) + .collect(), + ); + } + Ok(vectors) + } + + fn dimensions(&self) -> usize { + self.dimensions + } +} diff --git a/src/harness/ids/mod.rs b/src/harness/ids/mod.rs index fa2d13d..6384797 100644 --- a/src/harness/ids/mod.rs +++ b/src/harness/ids/mod.rs @@ -64,7 +64,9 @@ impl_string_id!(CellId); impl_string_id!(CheckpointId); impl_string_id!(InterruptId); +use std::sync::OnceLock; use std::sync::atomic::{AtomicU64, Ordering}; +use std::time::{SystemTime, UNIX_EPOCH}; /// Process-unique monotonic sequence source for deterministic, dependency-free /// id generation. @@ -82,6 +84,58 @@ pub fn next_seq() -> u64 { ID_SEQ.fetch_add(1, Ordering::Relaxed) } +/// Returns a per-process nonce, stable for the lifetime of the process and +/// (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. +pub fn process_nonce() -> u64 { + static NONCE: OnceLock = OnceLock::new(); + *NONCE.get_or_init(|| { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|d| d.as_nanos() as u64) + .unwrap_or(0) + }) +} + +/// Returns the current wall-clock time in milliseconds since the Unix epoch, +/// or `0` if the clock is set before the epoch. +/// +/// The single `now_ms` used across the crate for timestamping records, +/// checkpoints, goals, and observability events, so the epoch/`unwrap_or(0)` +/// convention lives in exactly one place instead of being re-hand-rolled in +/// every module that needs a millisecond timestamp. +pub fn now_ms() -> u64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|d| d.as_millis() as u64) + .unwrap_or(0) +} + +/// Allocates a fresh [`RunId`] of the form `run--`, collision-free +/// across process restarts (see [`process_nonce`]). +pub fn new_run_id() -> RunId { + RunId(format!("run-{}-{}", process_nonce(), next_seq())) +} + +/// Allocates a fresh [`CheckpointId`] of the form `ckpt--`, +/// 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. +pub fn new_checkpoint_id() -> CheckpointId { + CheckpointId(format!("ckpt-{}-{}", process_nonce(), next_seq())) +} + /// Allocates a fresh, process-unique [`SessionId`] of the form `session-`. pub fn new_session_id() -> SessionId { SessionId(format!("session-{}", next_seq())) diff --git a/src/harness/ids/test.rs b/src/harness/ids/test.rs index 2965b64..2ecd603 100644 --- a/src/harness/ids/test.rs +++ b/src/harness/ids/test.rs @@ -30,6 +30,31 @@ fn ids_are_distinct_types_but_serialize_as_strings() { assert_eq!(back, call); } +#[test] +fn generated_run_and_checkpoint_ids_carry_restart_nonce() { + // Regression for the cross-restart collision: ids must not be the bare + // `run-0`/`ckpt-0` monotonic form (which repeats in every fresh process), + // but `--` so a restarted, resumed thread never re-mints + // an id it already used. + let ckpt = new_checkpoint_id(); + let parts: Vec<&str> = ckpt.as_str().split('-').collect(); + assert_eq!(parts.len(), 3, "checkpoint id has prefix-nonce-seq shape"); + assert_eq!(parts[0], "ckpt"); + assert_eq!(parts[1], process_nonce().to_string(), "middle is the nonce"); + + let run = new_run_id(); + assert!(run.as_str().starts_with("run-")); + assert!( + run.as_str().contains(&format!("-{}-", process_nonce())), + "run id embeds the process nonce" + ); + + // Ids are unique within a process. + let a = new_checkpoint_id(); + let b = new_checkpoint_id(); + assert_ne!(a, b, "consecutive checkpoint ids differ"); +} + #[test] fn status_and_phase_use_snake_case() { assert_eq!( @@ -41,3 +66,15 @@ fn status_and_phase_use_snake_case() { "\"building_request\"" ); } + +#[test] +fn now_ms_returns_a_recent_unix_millis_timestamp() { + // The shared clock helper must return a plausible wall-clock timestamp: at + // or after a fixed 2020-01-01 epoch anchor and monotonic non-decreasing + // across two reads. (2020-01-01T00:00:00Z in ms.) + const JAN_2020_MS: u64 = 1_577_836_800_000; + let first = now_ms(); + assert!(first >= JAN_2020_MS, "timestamp {first} predates 2020"); + let second = now_ms(); + assert!(second >= first, "now_ms went backwards: {second} < {first}"); +} diff --git a/src/harness/limits/mod.rs b/src/harness/limits/mod.rs index 11bd695..991c994 100644 --- a/src/harness/limits/mod.rs +++ b/src/harness/limits/mod.rs @@ -49,18 +49,13 @@ impl RunLimits { self } - /// Sets the per-call retry cap. + /// Sets the per-call retry cap (a retry *count*, not counting the first + /// attempt). See [`RunLimits::max_retries_per_call`]. pub fn with_max_retries_per_call(mut self, n: usize) -> Self { self.max_retries_per_call = n; self } - /// Sets the concurrency cap. `None` removes the limit. - pub fn with_max_concurrency(mut self, n: Option) -> Self { - self.max_concurrency = n; - self - } - /// Sets the maximum sub-agent / recursion depth for the run tree. pub fn with_max_depth(mut self, n: usize) -> Self { self.max_depth = n; @@ -190,6 +185,21 @@ impl LimitTracker { pub fn limits(&self) -> &RunLimits { &self.limits } + + /// Overrides the model-call and tool-call caps in place, preserving + /// already-recorded counts and the wall-clock start time. + /// + /// A `RunContext` derives its tracker's initial limits from its + /// `RunConfig`, which always carries a concrete default. That can + /// silently disagree with a harness-wide `RunPolicy` configured with a + /// different cap, so the *reported* limit (the policy's) and the limit + /// that actually trips (the tracker's) diverge. The agent loop calls this + /// once per run to reconcile the two into a single enforced source of + /// truth before the loop begins. + pub fn sync_call_limits(&mut self, max_model_calls: usize, max_tool_calls: usize) { + self.limits.max_model_calls = max_model_calls; + self.limits.max_tool_calls = max_tool_calls; + } } #[cfg(test)] diff --git a/src/harness/limits/types.rs b/src/harness/limits/types.rs index a47f6e9..0ce6564 100644 --- a/src/harness/limits/types.rs +++ b/src/harness/limits/types.rs @@ -27,10 +27,13 @@ pub struct RunLimits { pub max_tool_calls: usize, /// Maximum elapsed wall-clock time in milliseconds. `None` means no limit. pub max_wall_clock_ms: Option, - /// Maximum number of retry attempts per individual call. + /// Maximum number of retry *attempts* (not counting the first try) + /// permitted for an individual model call. Reconciled with + /// [`crate::harness::retry::RetryPolicy::max_attempts`] by the agent loop + /// (see [`crate::harness::retry::RetryPolicy::max_attempts_capped_at`]): + /// whichever of the two is stricter wins, so this is a hard ceiling a + /// looser `RetryPolicy` cannot exceed. pub max_retries_per_call: usize, - /// Maximum number of concurrent in-flight calls. `None` means no limit. - pub max_concurrency: Option, /// Maximum sub-agent / recursion depth allowed for the run tree rooted at /// this run. A top-level run is depth `0`; each nested child run increments /// the depth. A sub-agent invocation whose child depth would exceed this cap @@ -51,7 +54,6 @@ impl Default for RunLimits { max_tool_calls: 50, max_wall_clock_ms: None, max_retries_per_call: 3, - max_concurrency: None, max_depth: Self::DEFAULT_MAX_DEPTH, } } diff --git a/src/harness/memory/mod.rs b/src/harness/memory/mod.rs index e5c7f09..47f72d7 100644 --- a/src/harness/memory/mod.rs +++ b/src/harness/memory/mod.rs @@ -74,6 +74,19 @@ impl ChatHistory for InMemoryChatHistory { Ok(()) } + async fn replace(&self, thread_id: &str, messages: Vec) -> Result<()> { + let mut threads = self + .threads + .lock() + .map_err(|e| TinyAgentsError::Memory(format!("chat history lock poisoned: {e}")))?; + if messages.is_empty() { + threads.remove(thread_id); + } else { + threads.insert(thread_id.to_string(), messages); + } + Ok(()) + } + async fn clear(&self, thread_id: &str) -> Result<()> { let mut threads = self .threads @@ -120,6 +133,17 @@ impl ChatHistory for StoreChatHistory { self.store.put(Self::NAMESPACE, thread_id, value).await } + async fn replace(&self, thread_id: &str, messages: Vec) -> Result<()> { + // A single put (or delete for an empty list) instead of a + // clear-then-N-appends: the whole thread is rewritten atomically, so a + // failed write cannot leave the thread with its history destroyed. + if messages.is_empty() { + return self.store.delete(Self::NAMESPACE, thread_id).await; + } + let value = serde_json::to_value(&messages)?; + self.store.put(Self::NAMESPACE, thread_id, value).await + } + async fn clear(&self, thread_id: &str) -> Result<()> { self.store.delete(Self::NAMESPACE, thread_id).await } @@ -168,15 +192,13 @@ impl ShortTermMemory { /// Replaces the thread's history with `messages` (trimmed first). /// - /// Clears the existing history then re-appends the trimmed list, so the - /// stored state matches what [`load`](Self::load) would return. + /// Delegates to [`ChatHistory::replace`], a single bulk write on the durable + /// and in-memory backends, so the stored state matches what + /// [`load`](Self::load) would return without clearing history first (an + /// intermediate failure can no longer destroy the thread). pub async fn save(&self, messages: Vec) -> Result<()> { let trimmed = self.apply_trim(messages); - self.history.clear(&self.thread_id).await?; - for message in trimmed { - self.history.append(&self.thread_id, message).await?; - } - Ok(()) + self.history.replace(&self.thread_id, trimmed).await } /// Clears the thread's history. diff --git a/src/harness/memory/test.rs b/src/harness/memory/test.rs index e4b4202..ed10fe6 100644 --- a/src/harness/memory/test.rs +++ b/src/harness/memory/test.rs @@ -111,6 +111,76 @@ async fn short_term_memory_trim_hook_applies() { assert_eq!(loaded[0].text(), "3"); } +/// A store whose `put` always fails, wrapping a working `get`/`delete`. Used to +/// prove that a mid-save write failure does not destroy pre-existing history. +#[derive(Clone)] +struct PutFailsStore { + inner: InMemoryStore, +} + +#[async_trait::async_trait] +impl crate::harness::store::Store for PutFailsStore { + async fn get(&self, namespace: &str, key: &str) -> Result> { + self.inner.get(namespace, key).await + } + async fn put(&self, _namespace: &str, _key: &str, _value: serde_json::Value) -> Result<()> { + Err(TinyAgentsError::Memory("write refused".into())) + } + async fn delete(&self, namespace: &str, key: &str) -> Result<()> { + self.inner.delete(namespace, key).await + } + async fn list(&self, namespace: &str) -> Result> { + self.inner.list(namespace).await + } +} + +#[tokio::test] +async fn save_does_not_destroy_history_when_write_fails() { + // Seed a thread's history through the working inner store. + let inner = InMemoryStore::new(); + let seed = StoreChatHistory::new(inner.clone()); + seed.append("t", Message::user("kept-one")).await.unwrap(); + seed.append("t", Message::user("kept-two")).await.unwrap(); + + // A memory whose backing store rejects writes: the atomic replace must fail + // WITHOUT having cleared the existing history first (the old clear-then- + // append path would have wiped the thread before the failing append). + let mem = ShortTermMemory::new(StoreChatHistory::new(PutFailsStore { inner }), "t"); + let err = mem.save(vec![Message::system("new")]).await; + assert!(err.is_err(), "save should surface the write failure"); + + // The original two messages survive because nothing was cleared. + let survivor = StoreChatHistory::new(seed.store().clone()); + let msgs = survivor.messages("t").await.unwrap(); + assert_eq!(msgs.len(), 2); + assert_eq!(msgs[0].text(), "kept-one"); + assert_eq!(msgs[1].text(), "kept-two"); +} + +#[tokio::test] +async fn replace_is_single_write_bulk() { + let store = InMemoryStore::new(); + let history = StoreChatHistory::new(store.clone()); + history.append("t", Message::user("old")).await.unwrap(); + + history + .replace( + "t", + vec![Message::user("a"), Message::user("b"), Message::user("c")], + ) + .await + .unwrap(); + + let msgs = history.messages("t").await.unwrap(); + assert_eq!(msgs.len(), 3); + assert_eq!(msgs[0].text(), "a"); + assert_eq!(msgs[2].text(), "c"); + + // Replacing with an empty list removes the thread entirely. + history.replace("t", vec![]).await.unwrap(); + assert!(history.messages("t").await.unwrap().is_empty()); +} + #[test] fn memory_scope_serializes_snake_case() { assert_eq!( diff --git a/src/harness/memory/types.rs b/src/harness/memory/types.rs index a2ad81f..063bda0 100644 --- a/src/harness/memory/types.rs +++ b/src/harness/memory/types.rs @@ -54,6 +54,21 @@ pub trait ChatHistory: Send + Sync { /// Appends `message` to the end of `thread_id`'s history. async fn append(&self, thread_id: &str, message: Message) -> Result<()>; + /// Replaces `thread_id`'s entire history with `messages` in one operation. + /// + /// The default implementation clears then re-appends one message at a time, + /// which is O(n) writes and — critically — non-atomic: a failure partway + /// through leaves the thread with the already-cleared prefix lost. Durable + /// and in-memory backends override this with a single bulk write so a + /// mid-write failure leaves the prior history intact. + async fn replace(&self, thread_id: &str, messages: Vec) -> Result<()> { + self.clear(thread_id).await?; + for message in messages { + self.append(thread_id, message).await?; + } + Ok(()) + } + /// Removes all history for `thread_id`. /// /// This is a no-op if the thread has no history; it does not error. diff --git a/src/harness/message/mod.rs b/src/harness/message/mod.rs index cab7cfc..26772f7 100644 --- a/src/harness/message/mod.rs +++ b/src/harness/message/mod.rs @@ -75,6 +75,26 @@ impl Message { Message::Tool(m) => concat_text(&m.content), } } + + /// Returns the total number of Unicode scalar values across all text content + /// blocks, without allocating the concatenated string. + /// + /// Equivalent to `self.text().chars().count()` but avoids the intermediate + /// `String` allocation, which matters on hot paths such as token estimation + /// over a whole transcript. + pub fn char_len(&self) -> usize { + let content = match self { + Message::System(m) => &m.content, + Message::User(m) => &m.content, + Message::Assistant(m) => &m.content, + Message::Tool(m) => &m.content, + }; + content + .iter() + .filter_map(ContentBlock::as_text) + .map(|t| t.chars().count()) + .sum() + } } #[cfg(test)] diff --git a/src/harness/message/test.rs b/src/harness/message/test.rs index 5c44c72..bfd0f8a 100644 --- a/src/harness/message/test.rs +++ b/src/harness/message/test.rs @@ -30,6 +30,25 @@ fn text_concatenates_and_ignores_non_text() { assert_eq!(msg.text(), "ab"); } +#[test] +fn char_len_matches_text_char_count_including_multibyte() { + // Mixed text blocks with multi-byte scalar values; non-text blocks ignored. + let msg = Message::User(UserMessage { + content: vec![ + ContentBlock::Text("héllo".into()), + ContentBlock::Json(json!({"k": "v"})), + ContentBlock::Text("🌍!".into()), + ], + }); + // char_len counts Unicode scalar values, matching text().chars().count() + // without allocating the joined string. + assert_eq!(msg.char_len(), msg.text().chars().count()); + assert_eq!( + msg.char_len(), + "héllo".chars().count() + "🌍!".chars().count() + ); +} + #[test] fn assistant_holds_tool_calls_and_usage() { let msg = Message::Assistant(AssistantMessage { diff --git a/src/harness/middleware/README.md b/src/harness/middleware/README.md new file mode 100644 index 0000000..72f0d31 --- /dev/null +++ b/src/harness/middleware/README.md @@ -0,0 +1,115 @@ +# harness::middleware + +Cross-cutting extension points that wrap agent, model, and tool execution. + +In the recursive (RLM-style) harness a sub-agent or sub-graph is just another +agent loop, so the same before/after hooks bracket the parent run *and* every +nested model/tool/agent call beneath it. That uniform wrapping is what lets +concerns like tracing, usage/cost roll-up, and guardrails compose consistently +as models call models and graphs run graphs. + +## Two extension shapes + +- **Lifecycle middleware** (`Middleware` trait) — observes and optionally + mutates values flowing past fixed points: `before_agent` / `after_agent`, + `before_model` / `on_model_delta` / `after_model`, `before_tool` / + `on_tool_delta` / `after_tool`, and `on_error`. Every hook has a no-op + default, is async, and returns `Result<()>`; an `Err` short-circuits the + stack. +- **Wrap ("around-call") middleware** (`ModelMiddleware`, `ToolMiddleware`) — + surrounds the inner call with a `next` handler (`ModelHandler` / + `ToolHandler`) instead of only observing before/after values. A wrap + middleware can proceed (call `next.run(..)` once), short-circuit (never call + it, returning a replacement `MiddlewareModelOutcome` / `MiddlewareToolOutcome` + directly), retry (call `next.run(..)` in a loop), or fall back (call it, then + substitute a response on error). This is the only extension point expressive + enough for retry/fallback/caching semantics. + +Both shapes are composed by `MiddlewareStack`, which holds three ordered +lists: `Middleware`, `ModelMiddleware`, `ToolMiddleware`. + +## Onion ordering + +`before_*` lifecycle hooks run in registration order; `after_*` hooks run in +**reverse** registration order, so the first-registered middleware is the +outermost layer — it sets up first and tears down last. Wrap middleware +compose the same way as a nested onion around the real model/tool call +(`ModelBaseCall` / `ToolBaseCall`), with the first-registered middleware +outermost and the base call innermost. + +Every per-middleware hook invocation is bracketed by +`AgentEvent::MiddlewareStarted` / `MiddlewareCompleted` events emitted through +the `RunContext`, so hook activity is independently observable via the event +sink regardless of what a middleware itself records. The pair is always +balanced: a hook that returns `Err` still emits its `MiddlewareCompleted` +before the error short-circuits the stack, so an observer never sees a dangling +`Started`. (The per-delta streaming hook `on_model_delta` is the one deliberate +exception — it emits no bracketing events, to stay cheap on the token hot path.) + +## Error handling + +The first hook that returns `Err` short-circuits the stack: every +middleware's `Middleware::on_error` is invoked (so all middleware get a chance +to log/redact/react), then the *original* error is returned to the caller. +Errors raised from `on_error` itself are ignored — they cannot mask the root +cause or replace it with a different error. + +## Public surface + +- `Middleware` — the lifecycle trait described above. +- `ModelMiddleware` / `ToolMiddleware` — the wrap + traits, each with a single `wrap_model` / `wrap_tool` method. +- `MiddlewareStack` — the composer; `push` / `push_model` / + `push_tool` register middleware, `run_before_agent` / `run_after_agent` / + `run_wrapped_model` / `run_wrapped_tool` / etc. run them. +- `AgentRun` — the accumulated result of a run (messages, final response, + structured output, usage, call/step counters) threaded through + `after_agent`. +- `ModelHandler<'a, State, Ctx>` / `ToolHandler<'a, State, Ctx>` — the `next` + handle passed to wrap middleware; `.run(ctx, state, request/call)` proceeds + to the next layer. +- `ModelBaseCall` / `ToolBaseCall` — the innermost real + call, supplied by the agent loop as the base of the onion. +- `MiddlewareModelOutcome` / `MiddlewareToolOutcome` — `#[non_exhaustive]` + result enums wrap middleware resolve to; construct via `.into()` from a + `ModelResponse` / `ToolResult`. + +### Built-in middleware (`library/`) + +- `LoggingMiddleware` — observation-only; counts how often each lifecycle hook + fires (`HookCounts`, readable via `.counts()`). Emits no events of its own — + the stack already emits start/completed events. +- `MessageTrimMiddleware` — replaces `request.messages` with the result of + `summarization::trim_messages` under a configured `TrimStrategy` in + `before_model`. +- `ContextCompressionMiddleware` — consults a `SummarizationPolicy` in + `before_model` and is a complete no-op below the context-window threshold; + above it, condenses older messages via a `Summarizer` (default + `ConcatSummarizer`) into a summary message, keeps the recent window and + system messages verbatim, records a `SummaryRecord`, and emits + `AgentEvent::Compressed`. +- `PromptCacheGuardMiddleware` — computes the request's + `cache::PromptCacheLayout` in `before_model` and records a + `CacheLayoutEvent` whenever the cacheable prefix changes from the previous + call, making KV-cache/prompt-cache regressions observable. +- `UsageAccountingMiddleware` — folds each `response.usage` into a running + `UsageTotals` in `after_model`, readable via `.totals()`. + +## Files + +| File | Role | +| --- | --- | +| `types.rs` | Every public type: traits, `MiddlewareStack`, built-in middleware structs. | +| `mod.rs` | Behavioral code: `AgentRun` helpers, the stack runner. | +| `library/` | Constructors and impls for every built-in middleware, split by concern: `resilience.rs` (retry/timeout/fallback/rate-limit), `budget.rs` (token/cost tracking and enforcement), `tool_policy.rs` (allowlisting, policy, dynamic/contextual selection, human approval), `context.rs` (message trim, summarization-based compression, prompt-cache guard), `observe.rs` (structured-output validation, dynamic prompt, redaction, tracing, logging, usage accounting). | +| `test.rs` | Unit tests (ordering, short-circuiting, each built-in middleware). | + +## Operational constraints + +- Hooks are invoked with `&mut RunContext` and a shared `&State`; a + middleware must not assume exclusive ownership of `State` — it is read-only + from the middleware's perspective. +- Because `on_error` failures are swallowed, a middleware must not rely on + `on_error` for anything beyond best-effort side effects (logging, metrics). +- Wrap middleware that retry `next.run(..)` are responsible for their own + budget/backoff; the stack does not cap retry attempts. diff --git a/src/harness/middleware/library/budget.rs b/src/harness/middleware/library/budget.rs new file mode 100644 index 0000000..d7b1dc5 --- /dev/null +++ b/src/harness/middleware/library/budget.rs @@ -0,0 +1,346 @@ +//! Budget middleware: token/cost preflight reservation and enforcement +//! (`BudgetTracker`, `BudgetLimits`, `BudgetMiddleware`). +//! +//! Split out of `library/mod.rs`; see that module's doc comment for the +//! full built-in middleware library overview. + +use super::*; + +// ── BudgetMiddleware ────────────────────────────────────────────────────────── + +impl BudgetTracker { + /// Creates an empty tracker. + pub fn new() -> Self { + Self::default() + } + + /// Locks the inner spend, recovering the last-known state if the mutex + /// was poisoned by a panicking holder. + /// + /// A poisoned mutex still holds a valid (if possibly stale) last-written + /// value; treating poisoning as "spend unknown, default to zero" would + /// make the budget enforcer fail *open* (every subsequent call sees an + /// empty budget and is admitted). Recovering the poisoned guard keeps + /// enforcement fail-closed: accumulated spend is never lost. + fn lock_recovering(&self) -> std::sync::MutexGuard<'_, BudgetSpend> { + self.inner + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + } + + /// Returns a snapshot of the accumulated spend. + pub fn snapshot(&self) -> BudgetSpend { + *self.lock_recovering() + } + + /// Folds a model call's usage and estimated cost into the tracker. + pub fn record( + &self, + usage: crate::harness::usage::Usage, + cost: crate::harness::cost::CostTotals, + ) { + let mut guard = self.lock_recovering(); + guard.usage += usage; + guard.cost += cost; + } +} + +impl BudgetLimits { + /// Returns a human-readable reason when `spend` meets or exceeds any limit. + pub(super) fn exceeded_reason(&self, spend: &BudgetSpend) -> Option { + let u = &spend.usage.usage; + if let Some(max) = self.max_input_tokens + && u.input_tokens >= max + { + return Some(format!("input tokens {} >= budget {max}", u.input_tokens)); + } + if let Some(max) = self.max_cached_input_tokens + && u.cache_read_tokens >= max + { + return Some(format!( + "cached input tokens {} >= budget {max}", + u.cache_read_tokens + )); + } + if let Some(max) = self.max_output_tokens + && u.output_tokens >= max + { + return Some(format!("output tokens {} >= budget {max}", u.output_tokens)); + } + if let Some(max) = self.max_total_tokens + && u.effective_total() >= max + { + return Some(format!( + "total tokens {} >= budget {max}", + u.effective_total() + )); + } + if let Some(max) = self.max_reasoning_tokens + && u.reasoning_tokens >= max + { + return Some(format!( + "reasoning tokens {} >= budget {max}", + u.reasoning_tokens + )); + } + if let Some(max) = self.max_cost + && spend.cost.total_cost >= max + { + return Some(format!( + "cost {:.6} >= budget {max:.6}", + spend.cost.total_cost + )); + } + None + } + + /// Returns a reason when `spend` crosses `warn_fraction` of any set limit. + fn warn_reason(&self, spend: &BudgetSpend) -> Option { + let frac = self.warn_fraction?; + let u = &spend.usage.usage; + let over = |value: f64, max: Option| -> bool { + max.is_some_and(|m| m > 0 && value >= frac * m as f64) + }; + if over(u.input_tokens as f64, self.max_input_tokens) { + return Some("input token budget".to_string()); + } + if over(u.cache_read_tokens as f64, self.max_cached_input_tokens) { + return Some("cached input token budget".to_string()); + } + if over(u.output_tokens as f64, self.max_output_tokens) { + return Some("output token budget".to_string()); + } + if over(u.effective_total() as f64, self.max_total_tokens) { + return Some("total token budget".to_string()); + } + if over(u.reasoning_tokens as f64, self.max_reasoning_tokens) { + return Some("reasoning token budget".to_string()); + } + if let Some(max) = self.max_cost + && max > 0.0 + && spend.cost.total_cost >= frac * max + { + return Some("cost budget".to_string()); + } + None + } +} + +/// Estimates the input tokens a request will consume by summing a +/// heuristic token estimate over every message's text. Used for budget +/// preflight reservation, which only needs an order-of-magnitude bound. +fn estimated_input_tokens(request: &ModelRequest) -> u64 { + request + .messages + .iter() + .map(|m| crate::harness::summarization::estimate_tokens(&m.text())) + .sum() +} + +impl BudgetMiddleware { + /// Creates a budget middleware with its own fresh tracker and no pricing + /// table (token budgets only; cost stays zero until pricing is supplied). + pub fn new(limits: BudgetLimits) -> Self { + Self { + label: "budget", + limits, + tracker: BudgetTracker::new(), + pricing: std::collections::HashMap::new(), + pending_reservation: std::sync::Mutex::new(0), + } + } + + /// Shares an existing [`BudgetTracker`] so this middleware's spend rolls up + /// into a run-tree-wide budget (hand the same tracker to sub-agents). + pub fn with_tracker(mut self, tracker: BudgetTracker) -> Self { + self.tracker = tracker; + self + } + + /// Supplies a per-model-name [`ModelPricing`] table so `after_model` can + /// price usage and enforce the money budget. + pub fn with_pricing( + mut self, + pricing: std::collections::HashMap, + ) -> Self { + self.pricing = pricing; + self + } + + /// Returns the shared tracker (for reading accumulated spend). + pub fn tracker(&self) -> BudgetTracker { + self.tracker.clone() + } + + fn price(&self, response: &ModelResponse) -> crate::harness::cost::CostTotals { + let Some(usage) = response.usage else { + return crate::harness::cost::CostTotals::default(); + }; + let Some(name) = response.resolved_model.as_ref().map(|r| r.name.as_str()) else { + return crate::harness::cost::CostTotals::default(); + }; + match self.pricing.get(name) { + Some(pricing) => crate::harness::cost::estimate_cost(pricing, &usage), + None => crate::harness::cost::CostTotals::default(), + } + } +} + +#[async_trait] +impl Middleware for BudgetMiddleware { + fn name(&self) -> &str { + self.label + } + + async fn before_model( + &self, + ctx: &mut RunContext, + _state: &State, + request: &mut ModelRequest, + ) -> Result<()> { + // Preflight is check-and-reserve under a single lock acquisition so + // concurrent runs sharing this tracker cannot all observe capacity + // and reserve past it (a separate check-then-lock-then-write window + // lets N concurrent callers each pass the check before any of them + // records a reservation, collectively overshooting the budget). + let estimated = estimated_input_tokens(request); + { + let mut guard = self.tracker.lock_recovering(); + // (1) Already exhausted before this call. + if let Some(reason) = self.limits.exceeded_reason(&guard) { + drop(guard); + ctx.emit(AgentEvent::BudgetExceeded { + reason: reason.clone(), + blocked: true, + }); + return Err(TinyAgentsError::LimitExceeded(format!( + "budget exhausted: {reason}" + ))); + } + + // (2) Preflight reservation: estimate this call's input tokens + // (plus every other in-flight reservation on this shared + // tracker) and block *before* dispatching if it would breach the + // input budget, so a single large call — or several concurrent + // ones — cannot collectively overshoot it. + if let Some(max) = self.limits.max_input_tokens + && guard.usage.usage.input_tokens + guard.reserved_input_total + estimated > max + { + let reason = format!( + "reserved input tokens {} + {estimated} > budget {max}", + guard.usage.usage.input_tokens + guard.reserved_input_total + ); + drop(guard); + ctx.emit(AgentEvent::BudgetExceeded { + reason: reason.clone(), + blocked: true, + }); + return Err(TinyAgentsError::LimitExceeded(format!( + "budget reservation exceeded: {reason}" + ))); + } + guard.reserved_input_total += estimated; + } + + // Remember this run's own outstanding reservation for reconciliation + // in `after_model` (local to this middleware instance, so concurrent + // runs sharing the tracker never clobber each other's amount). + *self + .pending_reservation + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) = estimated; + + ctx.emit(AgentEvent::BudgetReserved { + estimated_input_tokens: estimated, + }); + Ok(()) + } + + async fn after_model( + &self, + ctx: &mut RunContext, + _state: &State, + response: &mut ModelResponse, + ) -> Result<()> { + // Release this run's outstanding reservation regardless of whether + // usage came back, so a call that fails to report usage (or errors + // out before this hook) never leaks a permanent reservation that + // starves later calls on a shared tracker. + let reserved = std::mem::take( + &mut *self + .pending_reservation + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner), + ); + { + let mut guard = self.tracker.lock_recovering(); + guard.reserved_input_total = guard.reserved_input_total.saturating_sub(reserved); + } + + let Some(usage) = response.usage else { + return Ok(()); + }; + let cost = self.price(response); + ctx.emit(AgentEvent::BudgetReconciled { + estimated_input_tokens: reserved, + actual_input_tokens: usage.input_tokens, + }); + self.tracker.record(usage, cost); + + ctx.emit(AgentEvent::UsageRecorded { usage }); + if cost.total_cost > 0.0 { + ctx.emit(AgentEvent::CostRecorded { cost }); + } + + // Warn-once on threshold crossing: check-and-set the `warned` flag + // under a single lock so two concurrent calls crossing the + // threshold at once can't both observe `warned == false` and both + // emit the warning. + let (spend, warning) = { + let mut guard = self.tracker.lock_recovering(); + let warning = if !guard.warned { + let reason = self.limits.warn_reason(&guard); + if reason.is_some() { + guard.warned = true; + } + reason + } else { + None + }; + (*guard, warning) + }; + if let Some(reason) = warning { + ctx.emit(AgentEvent::BudgetWarning { + reason: format!("approaching {reason}"), + }); + } + if let Some(reason) = self.limits.exceeded_reason(&spend) { + ctx.emit(AgentEvent::BudgetExceeded { + reason, + blocked: false, + }); + } + Ok(()) + } + + async fn on_error(&self, _ctx: &mut RunContext, _error: &TinyAgentsError) -> Result<()> { + // A model call that fails (retries/fallback exhausted, hard provider + // error, middleware timeout, ...) short-circuits with `?` before + // `after_model` ever runs, so the reservation `before_model` added to + // the shared tracker would otherwise never be released. Release it + // here so a run of failures cannot permanently inflate + // `reserved_input_total` and starve every future call on a + // process-lifetime-shared `BudgetTracker`. + let reserved = std::mem::take( + &mut *self + .pending_reservation + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner), + ); + if reserved > 0 { + let mut guard = self.tracker.lock_recovering(); + guard.reserved_input_total = guard.reserved_input_total.saturating_sub(reserved); + } + Ok(()) + } +} diff --git a/src/harness/middleware/library/context.rs b/src/harness/middleware/library/context.rs new file mode 100644 index 0000000..ff8c3d0 --- /dev/null +++ b/src/harness/middleware/library/context.rs @@ -0,0 +1,331 @@ +//! Context-management middleware: message trimming, summarization-based +//! compression, and prompt-cache-layout guarding. +//! +//! Split out of `middleware/mod.rs`; see that module's doc comment for the +//! full middleware pipeline overview. + +use super::*; +use crate::harness::cache::{CacheLayoutEvent, PromptCacheLayout}; +use crate::harness::middleware::{ + ContextCompressionMiddleware, DEFAULT_CACHE_GUARD_EVENT_CAP, DEFAULT_COMPRESSION_RECORD_CAP, + MessageTrimMiddleware, MicrocompactMiddleware, PromptCacheGuardMiddleware, +}; +use crate::harness::summarization::{ + ConcatSummarizer, SummarizationPolicy, Summarizer, SummaryRecord, TrimStrategy, + estimate_tokens, trim_messages, +}; + +// ── MessageTrimMiddleware ───────────────────────────────────────────────────── + +impl MessageTrimMiddleware { + /// Creates a trim middleware using the given [`TrimStrategy`]. + pub fn new(strategy: TrimStrategy) -> Self { + Self { strategy } + } +} + +#[async_trait] +impl Middleware for MessageTrimMiddleware { + fn name(&self) -> &str { + "message_trim" + } + + async fn before_model( + &self, + _ctx: &mut RunContext, + _state: &State, + request: &mut ModelRequest, + ) -> Result<()> { + request.messages = trim_messages(&request.messages, &self.strategy); + Ok(()) + } +} + +// ── ContextCompressionMiddleware ────────────────────────────────────────────── + +/// Estimate the total tokens of a message slice using the same per-message +/// heuristic the [`SummarizationPolicy`] uses internally. +fn total_message_tokens(messages: &[crate::harness::message::Message]) -> u64 { + messages.iter().map(|m| estimate_tokens(&m.text())).sum() +} + +impl ContextCompressionMiddleware { + /// Creates a compression middleware backed by the default + /// [`ConcatSummarizer`]. + pub fn new(policy: SummarizationPolicy) -> Self { + Self::with_summarizer(policy, Box::new(ConcatSummarizer)) + } + + /// Creates a compression middleware with a custom [`Summarizer`]. + pub fn with_summarizer(policy: SummarizationPolicy, summarizer: Box) -> Self { + Self { + label: "context_compression", + policy, + summarizer, + records: std::sync::Mutex::new(std::collections::VecDeque::new()), + max_records: DEFAULT_COMPRESSION_RECORD_CAP, + } + } + + /// Sets the maximum number of [`SummaryRecord`]s retained before the + /// oldest is evicted. `0` disables recording entirely. + pub fn with_max_records(mut self, max_records: usize) -> Self { + self.max_records = max_records; + let mut records = self.records.lock().expect("records mutex poisoned"); + while records.len() > max_records { + records.pop_front(); + } + drop(records); + self + } + + /// Returns the configured [`SummarizationPolicy`]. + pub fn policy(&self) -> &SummarizationPolicy { + &self.policy + } + + /// Returns the [`SummaryRecord`]s produced so far, in order. Bounded to + /// at most [`ContextCompressionMiddleware::with_max_records`] entries + /// (default [`DEFAULT_COMPRESSION_RECORD_CAP`]); older records are + /// evicted first. + pub fn records(&self) -> Vec { + self.records + .lock() + .expect("records mutex poisoned") + .iter() + .cloned() + .collect() + } +} + +#[async_trait] +impl Middleware for ContextCompressionMiddleware { + fn name(&self) -> &str { + self.label + } + + async fn before_model( + &self, + ctx: &mut RunContext, + _state: &State, + request: &mut ModelRequest, + ) -> Result<()> { + // Below the window threshold: pass through untouched (no-op, no event). + if !self.policy.should_summarize(&request.messages) { + return Ok(()); + } + + let (to_summarize, mut to_keep) = self.policy.plan(&request.messages); + // Nothing old enough to compress (e.g. keep_last covers everything): + // leave the transcript untouched rather than summarizing an empty set. + if to_summarize.is_empty() { + return Ok(()); + } + + let from_tokens = total_message_tokens(&request.messages); + let record = self.summarizer.summarize(&to_summarize).await?; + + // `plan` returns `to_keep` as `[system prompts..., recent turns...]`. + // Insert the summary *after* the leading system prompts, not at index 0: + // a system prompt must stay first so its persistent instructions keep + // priority and the cacheable prefix is not churned. The summary of the + // elided older turns then sits between the system prompt and the kept + // recent turns, in chronological position. + let system_prefix = to_keep + .iter() + .take_while(|m| matches!(m, crate::harness::message::Message::System(_))) + .count(); + let recent = to_keep.split_off(system_prefix); + let mut new_messages = Vec::with_capacity(to_keep.len() + recent.len() + 1); + new_messages.append(&mut to_keep); + new_messages.push(record.summary.clone()); + new_messages.extend(recent); + let to_tokens = total_message_tokens(&new_messages); + + { + let mut records = self.records.lock().expect("records mutex poisoned"); + if self.max_records > 0 { + if records.len() >= self.max_records { + records.pop_front(); + } + records.push_back(record); + } + } + request.messages = new_messages; + + ctx.emit(AgentEvent::Compressed { + from_tokens, + to_tokens, + }); + Ok(()) + } +} + +// ── MicrocompactMiddleware ──────────────────────────────────────────────────── + +impl MicrocompactMiddleware { + /// Creates a micro-compaction middleware that keeps the newest `keep_recent` + /// tool-result bodies verbatim and blanks older ones with `placeholder`. + /// Event emission is off by default; enable it with + /// [`MicrocompactMiddleware::with_events`]. + pub fn new(keep_recent: usize, placeholder: impl Into) -> Self { + Self { + label: "microcompact", + keep_recent, + placeholder: placeholder.into(), + emit_events: false, + } + } + + /// Enable or disable emitting an + /// [`AgentEvent::Compressed`][crate::harness::events::AgentEvent::Compressed] + /// event whenever at least one tool body is cleared. Off by default so the + /// middleware can be a silent transcript rewrite. + pub fn with_events(mut self, emit_events: bool) -> Self { + self.emit_events = emit_events; + self + } + + /// The number of most-recent tool-result bodies kept verbatim. + pub fn keep_recent(&self) -> usize { + self.keep_recent + } + + /// The placeholder text swapped in for cleared tool-result bodies. + pub fn placeholder(&self) -> &str { + &self.placeholder + } +} + +#[async_trait] +impl Middleware for MicrocompactMiddleware { + fn name(&self) -> &str { + self.label + } + + async fn before_model( + &self, + ctx: &mut RunContext, + _state: &State, + request: &mut ModelRequest, + ) -> Result<()> { + // Indices of every tool-result message, oldest → newest. + let tool_idxs: Vec = request + .messages + .iter() + .enumerate() + .filter(|(_, m)| matches!(m, Message::Tool(_))) + .map(|(i, _)| i) + .collect(); + if tool_idxs.len() <= self.keep_recent { + return Ok(()); + } + + let from_tokens = if self.emit_events { + total_message_tokens(&request.messages) + } else { + 0 + }; + + let cut = tool_idxs.len() - self.keep_recent; + let mut cleared = 0usize; + for &i in &tool_idxs[..cut] { + // Skip messages already reduced to the placeholder; otherwise swap the + // body for it (idempotent, preserves the tool_call_id). + if request.messages[i].text() == self.placeholder { + continue; + } + if let Message::Tool(t) = &request.messages[i] { + let id = t.tool_call_id.clone(); + request.messages[i] = Message::tool(id, self.placeholder.clone()); + cleared += 1; + } + } + + if self.emit_events && cleared > 0 { + let to_tokens = total_message_tokens(&request.messages); + ctx.emit(AgentEvent::Compressed { + from_tokens, + to_tokens, + }); + } + Ok(()) + } +} + +// ── PromptCacheGuardMiddleware ──────────────────────────────────────────────── + +impl PromptCacheGuardMiddleware { + /// Creates a cache-guard middleware with the default label + /// `"prompt_cache_guard"`. + pub fn new() -> Self { + Self { + label: "prompt_cache_guard", + previous: std::sync::Mutex::new(None), + events: std::sync::Mutex::new(std::collections::VecDeque::new()), + max_events: DEFAULT_CACHE_GUARD_EVENT_CAP, + } + } + + /// Sets the maximum number of [`CacheLayoutEvent`]s retained before the + /// oldest is evicted. `0` disables recording entirely. + pub fn with_max_events(mut self, max_events: usize) -> Self { + self.max_events = max_events; + let mut events = self.events.lock().expect("events mutex poisoned"); + while events.len() > max_events { + events.pop_front(); + } + drop(events); + self + } + + /// Returns the cache-layout change events recorded so far, in order. + /// Bounded to at most [`PromptCacheGuardMiddleware::with_max_events`] + /// entries (default [`DEFAULT_CACHE_GUARD_EVENT_CAP`]); older events are + /// evicted first. + pub fn layout_events(&self) -> Vec { + self.events + .lock() + .expect("events mutex poisoned") + .iter() + .cloned() + .collect() + } +} + +impl Default for PromptCacheGuardMiddleware { + fn default() -> Self { + Self::new() + } +} + +#[async_trait] +impl Middleware for PromptCacheGuardMiddleware { + fn name(&self) -> &str { + self.label + } + + async fn before_model( + &self, + _ctx: &mut RunContext, + _state: &State, + request: &mut ModelRequest, + ) -> Result<()> { + let layout = PromptCacheLayout::from_request(request); + let mut previous = self.previous.lock().expect("previous mutex poisoned"); + if let Some(prev) = previous.as_ref() + && !prev.is_prefix_stable_against(&layout) + { + let event = CacheLayoutEvent::new(prev, &layout); + let mut events = self.events.lock().expect("events mutex poisoned"); + if self.max_events > 0 { + if events.len() >= self.max_events { + events.pop_front(); + } + events.push_back(event); + } + } + *previous = Some(layout); + Ok(()) + } +} diff --git a/src/harness/middleware/library/mod.rs b/src/harness/middleware/library/mod.rs index 14265a9..339bfd0 100644 --- a/src/harness/middleware/library/mod.rs +++ b/src/harness/middleware/library/mod.rs @@ -29,7 +29,7 @@ mod types; pub use types::*; -use std::collections::HashSet; +use std::collections::{HashSet, VecDeque}; use std::marker::PhantomData; use std::sync::{Arc, Mutex}; use std::time::{Duration, Instant}; @@ -49,1288 +49,11 @@ use crate::harness::retry::{RateLimiter, RetryPolicy, is_retryable}; use crate::harness::structured::{StructuredExtractor, StructuredStrategy}; use crate::harness::tool::{ToolCall, ToolDelta, ToolResult, ToolSchema}; -// ── RetryMiddleware ─────────────────────────────────────────────────────────── - -impl RetryMiddleware { - /// Creates a retry middleware using the given [`RetryPolicy`]. - pub fn new(policy: RetryPolicy) -> Self { - Self { - label: "retry", - policy, - } - } - - /// Creates a retry middleware with the default [`RetryPolicy`]. - pub fn with_default_policy() -> Self { - Self::new(RetryPolicy::default()) - } - - /// Returns the policy-derived backoff for the given retry `attempt`. - /// - /// Exposed for callers that want to inspect the backoff. The middleware - /// itself sleeps between retries only when the policy opts in via - /// [`RetryPolicy::with_backoff_sleep`]; otherwise it retries back-to-back. - pub fn backoff_for_attempt(&self, attempt: usize) -> Duration { - self.policy.backoff_for_attempt(attempt) - } -} - -#[async_trait] -impl ModelMiddleware for RetryMiddleware { - fn name(&self) -> &str { - self.label - } - - async fn wrap_model( - &self, - ctx: &mut RunContext, - state: &State, - request: ModelRequest, - next: ModelHandler<'_, State, Ctx>, - ) -> Result { - let mut attempt = 0usize; - loop { - match next.run(ctx, state, request.clone()).await { - Ok(outcome) => return Ok(outcome), - Err(error) => { - if is_retryable(&error) && self.policy.should_retry(attempt) { - attempt += 1; - let call_id = CallId::new(format!("{}-model", ctx.run_id())); - ctx.emit(AgentEvent::RetryScheduled { call_id, attempt }); - // Sleep for the backoff only when the policy opts in - // (`with_backoff_sleep`); a no-op otherwise. - self.policy.sleep_backoff(attempt).await; - continue; - } - return Err(error); - } - } - } - } -} - -// ── TimeoutMiddleware ───────────────────────────────────────────────────────── - -impl TimeoutMiddleware { - /// Creates a timeout middleware bounding each model call to `timeout`. - pub fn new(timeout: Duration) -> Self { - Self { - label: "timeout", - timeout, - } - } - - /// Creates a timeout middleware from a millisecond duration. - pub fn from_millis(ms: u64) -> Self { - Self::new(Duration::from_millis(ms)) - } -} - -#[async_trait] -impl ModelMiddleware for TimeoutMiddleware { - fn name(&self) -> &str { - self.label - } - - async fn wrap_model( - &self, - ctx: &mut RunContext, - state: &State, - request: ModelRequest, - next: ModelHandler<'_, State, Ctx>, - ) -> Result { - let run_id = ctx.run_id().as_str().to_string(); - let fut = next.run(ctx, state, request); - match tokio::time::timeout(self.timeout, fut).await { - Ok(result) => result, - Err(_) => Err(TinyAgentsError::Timeout(format!( - "model call for run `{run_id}` exceeded the {} ms middleware timeout", - self.timeout.as_millis() - ))), - } - } -} - -// ── ModelFallbackMiddleware ─────────────────────────────────────────────────── - -impl ModelFallbackMiddleware { - /// Creates a fallback middleware that tries each model name in order after - /// the primary call fails. - pub fn new(fallbacks: impl IntoIterator>) -> Self { - Self { - label: "model_fallback", - fallbacks: fallbacks.into_iter().map(Into::into).collect(), - } - } -} - -#[async_trait] -impl ModelMiddleware for ModelFallbackMiddleware { - fn name(&self) -> &str { - self.label - } - - async fn wrap_model( - &self, - ctx: &mut RunContext, - state: &State, - request: ModelRequest, - next: ModelHandler<'_, State, Ctx>, - ) -> Result { - match next.run(ctx, state, request.clone()).await { - Ok(outcome) => Ok(outcome), - Err(mut last_error) => { - let mut current = request.model.clone().unwrap_or_default(); - for fallback in &self.fallbacks { - ctx.emit(AgentEvent::FallbackSelected { - from: current.clone(), - to: fallback.clone(), - }); - let mut req = request.clone(); - req.model = Some(fallback.clone()); - match next.run(ctx, state, req).await { - Ok(outcome) => return Ok(outcome), - Err(error) => { - last_error = error; - current = fallback.clone(); - } - } - } - Err(last_error) - } - } - } -} - -// ── RateLimitMiddleware ─────────────────────────────────────────────────────── - -impl RateLimitMiddleware { - /// Creates a rate-limit middleware gating one token per call through - /// `limiter`, failing immediately when the bucket is empty - /// ([`RateLimitBehavior::Error`]). - pub fn new(limiter: Arc) -> Self { - Self { - label: "rate_limit", - limiter, - tokens: 1, - behavior: RateLimitBehavior::Error, - poll_interval: Duration::from_millis(50), - now: Arc::new(Instant::now), - } - } - - /// Sets the number of tokens each call consumes. - pub fn with_tokens(mut self, tokens: u64) -> Self { - self.tokens = tokens; - self - } - - /// Sets the behavior when the bucket lacks capacity. - pub fn with_behavior(mut self, behavior: RateLimitBehavior) -> Self { - self.behavior = behavior; - self - } - - /// Switches to [`RateLimitBehavior::Wait`] with the given poll interval. - pub fn waiting(mut self, poll_interval: Duration) -> Self { - self.behavior = RateLimitBehavior::Wait; - self.poll_interval = poll_interval; - self - } - - /// Replaces the clock used to read the current instant (for deterministic - /// tests). - pub fn with_clock(mut self, now: NowFn) -> Self { - self.now = now; - self - } -} - -#[async_trait] -impl ModelMiddleware for RateLimitMiddleware { - fn name(&self) -> &str { - self.label - } - - async fn wrap_model( - &self, - ctx: &mut RunContext, - state: &State, - request: ModelRequest, - next: ModelHandler<'_, State, Ctx>, - ) -> Result { - loop { - let now = (self.now)(); - if self.limiter.try_acquire(self.tokens, now) { - break; - } - match self.behavior { - RateLimitBehavior::Error => { - return Err(TinyAgentsError::LimitExceeded(format!( - "rate limit: could not acquire {} token(s)", - self.tokens - ))); - } - RateLimitBehavior::Wait => { - ctx.emit(AgentEvent::RateLimitWaited { - waited_ms: self.poll_interval.as_millis() as u64, - }); - tokio::time::sleep(self.poll_interval).await; - } - } - } - next.run(ctx, state, request).await - } -} - -// ── ToolAllowlistMiddleware ─────────────────────────────────────────────────── - -impl ToolAllowlistMiddleware { - /// Creates an allowlist middleware permitting only the named tools. - pub fn new(allowed: impl IntoIterator>) -> Self { - Self { - label: "tool_allowlist", - allowed: allowed.into_iter().map(Into::into).collect(), - } - } - - /// Returns `true` if `name` is on the allowlist. - pub fn allows(&self, name: &str) -> bool { - self.allowed.contains(name) - } -} - -#[async_trait] -impl Middleware for ToolAllowlistMiddleware { - fn name(&self) -> &str { - self.label - } - - async fn before_tool( - &self, - _ctx: &mut RunContext, - _state: &State, - call: &mut ToolCall, - ) -> Result<()> { - if !self.allowed.contains(&call.name) { - return Err(TinyAgentsError::Validation(format!( - "tool `{}` is not on the allowlist", - call.name - ))); - } - Ok(()) - } -} - -// ── BudgetMiddleware ────────────────────────────────────────────────────────── - -impl BudgetTracker { - /// Creates an empty tracker. - pub fn new() -> Self { - Self::default() - } - - /// Returns a snapshot of the accumulated spend. - pub fn snapshot(&self) -> BudgetSpend { - self.inner.lock().map(|g| *g).unwrap_or_default() - } - - /// Folds a model call's usage and estimated cost into the tracker. - pub fn record( - &self, - usage: crate::harness::usage::Usage, - cost: crate::harness::cost::CostTotals, - ) { - if let Ok(mut guard) = self.inner.lock() { - guard.usage += usage; - guard.cost += cost; - } - } -} - -impl BudgetLimits { - /// Returns a human-readable reason when `spend` meets or exceeds any limit. - fn exceeded_reason(&self, spend: &BudgetSpend) -> Option { - let u = &spend.usage.usage; - if let Some(max) = self.max_input_tokens - && u.input_tokens >= max - { - return Some(format!("input tokens {} >= budget {max}", u.input_tokens)); - } - if let Some(max) = self.max_cached_input_tokens - && u.cache_read_tokens >= max - { - return Some(format!( - "cached input tokens {} >= budget {max}", - u.cache_read_tokens - )); - } - if let Some(max) = self.max_output_tokens - && u.output_tokens >= max - { - return Some(format!("output tokens {} >= budget {max}", u.output_tokens)); - } - if let Some(max) = self.max_total_tokens - && u.effective_total() >= max - { - return Some(format!( - "total tokens {} >= budget {max}", - u.effective_total() - )); - } - if let Some(max) = self.max_reasoning_tokens - && u.reasoning_tokens >= max - { - return Some(format!( - "reasoning tokens {} >= budget {max}", - u.reasoning_tokens - )); - } - if let Some(max) = self.max_cost - && spend.cost.total_cost >= max - { - return Some(format!( - "cost {:.6} >= budget {max:.6}", - spend.cost.total_cost - )); - } - None - } - - /// Returns a reason when `spend` crosses `warn_fraction` of any set limit. - fn warn_reason(&self, spend: &BudgetSpend) -> Option { - let frac = self.warn_fraction?; - let u = &spend.usage.usage; - let over = |value: f64, max: Option| -> bool { - max.is_some_and(|m| m > 0 && value >= frac * m as f64) - }; - if over(u.input_tokens as f64, self.max_input_tokens) { - return Some("input token budget".to_string()); - } - if over(u.cache_read_tokens as f64, self.max_cached_input_tokens) { - return Some("cached input token budget".to_string()); - } - if over(u.output_tokens as f64, self.max_output_tokens) { - return Some("output token budget".to_string()); - } - if over(u.effective_total() as f64, self.max_total_tokens) { - return Some("total token budget".to_string()); - } - if over(u.reasoning_tokens as f64, self.max_reasoning_tokens) { - return Some("reasoning token budget".to_string()); - } - if let Some(max) = self.max_cost - && max > 0.0 - && spend.cost.total_cost >= frac * max - { - return Some("cost budget".to_string()); - } - None - } -} - -/// Estimates the input tokens a request will consume by summing a -/// heuristic token estimate over every message's text. Used for budget -/// preflight reservation, which only needs an order-of-magnitude bound. -fn estimated_input_tokens(request: &ModelRequest) -> u64 { - request - .messages - .iter() - .map(|m| crate::harness::summarization::estimate_tokens(&m.text())) - .sum() -} - -impl BudgetMiddleware { - /// Creates a budget middleware with its own fresh tracker and no pricing - /// table (token budgets only; cost stays zero until pricing is supplied). - pub fn new(limits: BudgetLimits) -> Self { - Self { - label: "budget", - limits, - tracker: BudgetTracker::new(), - pricing: std::collections::HashMap::new(), - } - } - - /// Shares an existing [`BudgetTracker`] so this middleware's spend rolls up - /// into a run-tree-wide budget (hand the same tracker to sub-agents). - pub fn with_tracker(mut self, tracker: BudgetTracker) -> Self { - self.tracker = tracker; - self - } - - /// Supplies a per-model-name [`ModelPricing`] table so `after_model` can - /// price usage and enforce the money budget. - pub fn with_pricing( - mut self, - pricing: std::collections::HashMap, - ) -> Self { - self.pricing = pricing; - self - } - - /// Returns the shared tracker (for reading accumulated spend). - pub fn tracker(&self) -> BudgetTracker { - self.tracker.clone() - } - - fn price(&self, response: &ModelResponse) -> crate::harness::cost::CostTotals { - let Some(usage) = response.usage else { - return crate::harness::cost::CostTotals::default(); - }; - let Some(name) = response.resolved_model.as_ref().map(|r| r.name.as_str()) else { - return crate::harness::cost::CostTotals::default(); - }; - match self.pricing.get(name) { - Some(pricing) => crate::harness::cost::estimate_cost(pricing, &usage), - None => crate::harness::cost::CostTotals::default(), - } - } -} - -#[async_trait] -impl Middleware for BudgetMiddleware { - fn name(&self) -> &str { - self.label - } - - async fn before_model( - &self, - ctx: &mut RunContext, - _state: &State, - request: &mut ModelRequest, - ) -> Result<()> { - let spend = self.tracker.snapshot(); - // (1) Already exhausted before this call. - if let Some(reason) = self.limits.exceeded_reason(&spend) { - ctx.emit(AgentEvent::BudgetExceeded { - reason: reason.clone(), - blocked: true, - }); - return Err(TinyAgentsError::LimitExceeded(format!( - "budget exhausted: {reason}" - ))); - } - - // (2) Preflight reservation: estimate this call's input tokens and block - // *before* dispatching if the reservation would breach the input budget, - // so a single large call cannot overshoot arbitrarily. - let estimated = estimated_input_tokens(request); - if let Some(max) = self.limits.max_input_tokens - && spend.usage.usage.input_tokens + estimated > max - { - let reason = format!( - "reserved input tokens {} + {estimated} > budget {max}", - spend.usage.usage.input_tokens - ); - ctx.emit(AgentEvent::BudgetExceeded { - reason: reason.clone(), - blocked: true, - }); - return Err(TinyAgentsError::LimitExceeded(format!( - "budget reservation exceeded: {reason}" - ))); - } - if let Ok(mut guard) = self.tracker.inner.lock() { - guard.last_reserved_input = estimated; - } - ctx.emit(AgentEvent::BudgetReserved { - estimated_input_tokens: estimated, - }); - Ok(()) - } - - async fn after_model( - &self, - ctx: &mut RunContext, - _state: &State, - response: &mut ModelResponse, - ) -> Result<()> { - let Some(usage) = response.usage else { - return Ok(()); - }; - let cost = self.price(response); - // Reconcile the preflight reservation against the actual usage. - let reserved = self - .tracker - .inner - .lock() - .map(|mut guard| std::mem::take(&mut guard.last_reserved_input)) - .unwrap_or(0); - ctx.emit(AgentEvent::BudgetReconciled { - estimated_input_tokens: reserved, - actual_input_tokens: usage.input_tokens, - }); - self.tracker.record(usage, cost); - - ctx.emit(AgentEvent::UsageRecorded { usage }); - if cost.total_cost > 0.0 { - ctx.emit(AgentEvent::CostRecorded { cost }); - } - - let mut spend = self.tracker.snapshot(); - // Warn-once on threshold crossing. - if !spend.warned - && let Some(reason) = self.limits.warn_reason(&spend) - { - ctx.emit(AgentEvent::BudgetWarning { - reason: format!("approaching {reason}"), - }); - if let Ok(mut guard) = self.tracker.inner.lock() { - guard.warned = true; - } - spend.warned = true; - } - if let Some(reason) = self.limits.exceeded_reason(&spend) { - ctx.emit(AgentEvent::BudgetExceeded { - reason, - blocked: false, - }); - } - Ok(()) - } -} - -// ── ToolPolicyMiddleware ────────────────────────────────────────────────────── - -impl ToolPolicyMiddleware { - /// Creates a policy middleware from a name→policy snapshot (typically - /// [`ToolRegistry::policies`][crate::harness::tool::ToolRegistry::policies]). - /// - /// Defaults are permissive: nothing is required or denied until configured. - /// Use [`strict`](Self::strict) for a fail-closed baseline. - pub fn new( - policies: std::collections::HashMap, - ) -> Self { - Self { - label: "tool_policy", - policies, - require_classification: false, - require_background_safe: false, - deny: crate::harness::tool::ToolSideEffects::default(), - require_sandbox: false, - require_approval: false, - approved: std::collections::HashSet::new(), - enforce_result_bytes: false, - } - } - - /// Creates a fail-closed policy middleware: unclassified tools are rejected, - /// and tools declaring `destructive` or `payment` side effects are denied. - pub fn strict( - policies: std::collections::HashMap, - ) -> Self { - Self { - label: "tool_policy", - policies, - require_classification: true, - require_background_safe: false, - deny: crate::harness::tool::ToolSideEffects { - destructive: true, - payment: true, - ..crate::harness::tool::ToolSideEffects::default() - }, - require_sandbox: false, - require_approval: false, - approved: std::collections::HashSet::new(), - enforce_result_bytes: false, - } - } - - /// Requires every tool to carry a classified policy (fail closed on - /// unclassified or unknown tools). - pub fn require_classification(mut self, require: bool) -> Self { - self.require_classification = require; - self - } - - /// Requires every exposed/executed tool to be `background_safe`. - pub fn require_background_safe(mut self, require: bool) -> Self { - self.require_background_safe = require; - self - } - - /// Denies tools declaring any side effect present in `mask`. - pub fn deny_side_effects(mut self, mask: crate::harness::tool::ToolSideEffects) -> Self { - self.deny = mask; - self - } - - /// Enforces that a tool declaring - /// [`SandboxMode::Required`][crate::harness::tool::SandboxMode::Required] - /// only runs when the run carries a sandboxed workspace (fail closed - /// otherwise). See [`RunContext::with_workspace`][crate::harness::context::RunContext::with_workspace]. - pub fn require_sandbox(mut self, require: bool) -> Self { - self.require_sandbox = require; - self - } - - /// Blocks any tool declaring `approval_required` unless its name is in - /// `approved`, turning the declarative approval flag into a fail-closed gate. - pub fn require_approval( - mut self, - approved: impl IntoIterator>, - ) -> Self { - self.require_approval = true; - self.approved = approved.into_iter().map(Into::into).collect(); - self - } - - /// Enforces each tool's declared `max_result_bytes` cap by truncating and - /// flagging oversized results in `after_tool`. - pub fn enforce_result_bytes(mut self, enforce: bool) -> Self { - self.enforce_result_bytes = enforce; - self - } - - /// Returns `Ok(())` if the named tool is permitted, otherwise an explanation - /// of why it is blocked. Used by both the exposure and execution hooks so a - /// hidden tool cannot be executed by a divergent decision. - fn evaluate(&self, name: &str) -> std::result::Result<(), String> { - let Some(policy) = self.policies.get(name) else { - if self.require_classification { - return Err(format!("tool `{name}` has no declared policy")); - } - return Ok(()); - }; - if self.require_classification && !policy.classified { - return Err(format!("tool `{name}` is unclassified")); - } - let s = &policy.side_effects; - let d = &self.deny; - let denied = (d.writes_files && s.writes_files) - || (d.network && s.network) - || (d.installs_dependencies && s.installs_dependencies) - || (d.destructive && s.destructive) - || (d.external_service && s.external_service) - || (d.payment && s.payment); - if denied { - return Err(format!("tool `{name}` declares a denied side effect")); - } - if self.require_background_safe && !policy.access.background_safe { - return Err(format!("tool `{name}` is not background-safe")); - } - if self.require_approval && policy.access.approval_required && !self.approved.contains(name) - { - return Err(format!( - "tool `{name}` requires approval that was not granted" - )); - } - Ok(()) - } - - /// The context-aware slice of policy enforcement: the sandbox requirement - /// depends on the run's workspace, which `evaluate` (name-only) cannot see. - fn evaluate_sandbox( - &self, - name: &str, - ctx: &RunContext, - ) -> std::result::Result<(), String> { - if !self.require_sandbox { - return Ok(()); - } - let Some(policy) = self.policies.get(name) else { - return Ok(()); - }; - if policy.runtime.sandbox != crate::harness::tool::SandboxMode::Required { - return Ok(()); - } - let sandboxed = ctx - .workspace - .as_ref() - .is_some_and(|ws| ws.sandbox == crate::harness::tool::SandboxMode::Required); - if sandboxed { - Ok(()) - } else { - Err(format!( - "tool `{name}` requires a sandbox but the run has none" - )) - } - } -} - -#[async_trait] -impl Middleware for ToolPolicyMiddleware { - fn name(&self) -> &str { - self.label - } - - async fn before_model( - &self, - ctx: &mut RunContext, - _state: &State, - request: &mut ModelRequest, - ) -> Result<()> { - request.tools.retain(|schema| { - self.evaluate(&schema.name).is_ok() && self.evaluate_sandbox(&schema.name, ctx).is_ok() - }); - Ok(()) - } - - async fn before_tool( - &self, - ctx: &mut RunContext, - _state: &State, - call: &mut ToolCall, - ) -> Result<()> { - self.evaluate(&call.name) - .and_then(|_| self.evaluate_sandbox(&call.name, ctx)) - .map_err(TinyAgentsError::Validation) - } - - async fn after_tool( - &self, - _ctx: &mut RunContext, - _state: &State, - result: &mut ToolResult, - ) -> Result<()> { - if !self.enforce_result_bytes { - return Ok(()); - } - if let Some(policy) = self.policies.get(&result.name) - && let Some(limit) = policy.runtime.max_result_bytes - && result.content.len() > limit - { - // Truncate on a char boundary at or below the byte limit so the - // enforced payload is still valid UTF-8. - let mut end = limit; - while end > 0 && !result.content.is_char_boundary(end) { - end -= 1; - } - result.content.truncate(end); - let note = format!("tool result exceeded max_result_bytes ({limit}); truncated"); - result.error = Some(match result.error.take() { - Some(existing) => format!("{existing}; {note}"), - None => note, - }); - } - Ok(()) - } -} - -// ── DynamicToolSelectionMiddleware ──────────────────────────────────────────── - -impl DynamicToolSelectionMiddleware { - /// Creates a selection middleware exposing only tools for which `predicate` - /// returns `true`. - pub fn new(predicate: ToolPredicate) -> Self { - Self { - label: "dynamic_tool_selection", - predicate, - } - } - - /// Creates a selection middleware exposing only the named tools. - pub fn allowing(names: impl IntoIterator>) -> Self { - let allowed: HashSet = names.into_iter().map(Into::into).collect(); - Self::new(Arc::new(move |schema: &ToolSchema| { - allowed.contains(&schema.name) - })) - } -} - -#[async_trait] -impl Middleware - for DynamicToolSelectionMiddleware -{ - fn name(&self) -> &str { - self.label - } - - async fn before_model( - &self, - _ctx: &mut RunContext, - _state: &State, - request: &mut ModelRequest, - ) -> Result<()> { - request.tools.retain(|schema| (self.predicate)(schema)); - Ok(()) - } -} - -// ── ContextualToolSelectionMiddleware ───────────────────────────────────────── - -impl ContextualToolSelectionMiddleware { - /// Creates a selection middleware from a context-aware predicate. - pub fn new(predicate: ContextualToolPredicate) -> Self { - Self { - label: "contextual_tool_selection", - predicate, - } - } - - /// Builds a selection middleware from explicit allow/deny lists. - /// - /// Composition rules (fail-closed): - /// - a tool named in `deny` is always hidden; - /// - when `allow` is `Some`, a tool must be named in it to be exposed - /// (unknown tools are hidden); - /// - when `allow` is `None`, everything not denied is exposed. - pub fn from_lists( - allow: Option>>, - deny: impl IntoIterator>, - ) -> Self { - let allow: Option> = - allow.map(|names| names.into_iter().map(Into::into).collect()); - let deny: HashSet = deny.into_iter().map(Into::into).collect(); - Self::from_resolved_lists(allow, deny) - } - - /// Builds a selection middleware whose effective policy is a child allow/deny - /// pair *composed with* an inherited parent policy, so a delegated sub-agent - /// can only ever narrow — never widen — the tools its parent allowed. - /// - /// Inheritance rules: - /// - **deny is additive**: the effective denylist is `parent_deny ∪ child_deny` - /// (a child cannot un-deny what the parent denied); - /// - **allow is intersective**: if both parent and child restrict to an - /// allowlist, the effective allowlist is their intersection; if only one - /// restricts, that allowlist applies; if neither does, all-not-denied is - /// exposed. - /// - /// The result is fail-closed for the same reasons as [`Self::from_lists`]. - pub fn inheriting( - parent_allow: Option>>, - parent_deny: impl IntoIterator>, - child_allow: Option>>, - child_deny: impl IntoIterator>, - ) -> Self { - let parent_allow: Option> = - parent_allow.map(|n| n.into_iter().map(Into::into).collect()); - let child_allow: Option> = - child_allow.map(|n| n.into_iter().map(Into::into).collect()); - let allow = match (parent_allow, child_allow) { - (Some(p), Some(c)) => Some(p.intersection(&c).cloned().collect()), - (Some(p), None) => Some(p), - (None, Some(c)) => Some(c), - (None, None) => None, - }; - let mut deny: HashSet = parent_deny.into_iter().map(Into::into).collect(); - deny.extend(child_deny.into_iter().map(Into::into)); - Self::from_resolved_lists(allow, deny) - } - - fn from_resolved_lists(allow: Option>, deny: HashSet) -> Self { - Self::new(Arc::new(move |schema: &ToolSchema, _ctx| { - if deny.contains(&schema.name) { - return false; - } - match &allow { - Some(set) => set.contains(&schema.name), - None => true, - } - })) - } -} - -#[async_trait] -impl Middleware - for ContextualToolSelectionMiddleware -{ - fn name(&self) -> &str { - self.label - } - - async fn before_model( - &self, - ctx: &mut RunContext, - _state: &State, - request: &mut ModelRequest, - ) -> Result<()> { - let selection = ToolSelectionContext { - run_id: ctx.config.run_id.as_str().to_string(), - depth: ctx.config.depth, - tags: ctx.config.tags.clone(), - requested_model: request.model.clone(), - }; - let mut excluded = Vec::new(); - request.tools.retain(|schema| { - let keep = (self.predicate)(schema, &selection); - if !keep { - excluded.push(schema.name.clone()); - } - keep - }); - // Make the exposure decision auditable when it actually withheld tools. - if !excluded.is_empty() { - ctx.emit(AgentEvent::ToolsFiltered { - by: self.label.to_string(), - excluded, - remaining: request.tools.len(), - }); - } - Ok(()) - } -} - -// ── HumanApprovalMiddleware ─────────────────────────────────────────────────── - -impl HumanApprovalMiddleware { - /// Creates an approval middleware that interrupts when any flagged tool is - /// called and no approval callback is configured. - pub fn new(flagged: impl IntoIterator>) -> Self { - Self { - label: "human_approval", - flagged: flagged.into_iter().map(Into::into).collect(), - approve: None, - } - } - - /// Attaches an approval callback consulted for flagged tools. Returning - /// `true` admits the call; `false` (or no callback) raises an interrupt. - pub fn with_approval(mut self, approve: ApprovalFn) -> Self { - self.approve = Some(approve); - self - } -} - -#[async_trait] -impl Middleware for HumanApprovalMiddleware { - fn name(&self) -> &str { - self.label - } - - async fn before_tool( - &self, - _ctx: &mut RunContext, - _state: &State, - call: &mut ToolCall, - ) -> Result<()> { - if self.flagged.contains(&call.name) { - let approved = self - .approve - .as_ref() - .map(|approve| approve(call)) - .unwrap_or(false); - if !approved { - return Err(TinyAgentsError::Interrupted { - node: "tool".to_string(), - message: format!("tool `{}` requires human approval", call.name), - }); - } - } - Ok(()) - } -} - -// ── StructuredOutputValidatorMiddleware ─────────────────────────────────────── - -impl StructuredOutputValidatorMiddleware { - /// Creates a validator middleware checking responses against `format`. - pub fn new(format: ResponseFormat) -> Self { - Self { - label: "structured_output_validator", - format, - } - } -} - -#[async_trait] -impl Middleware - for StructuredOutputValidatorMiddleware -{ - fn name(&self) -> &str { - self.label - } - - async fn after_model( - &self, - _ctx: &mut RunContext, - _state: &State, - response: &mut ModelResponse, - ) -> Result<()> { - match &self.format { - ResponseFormat::Text => Ok(()), - ResponseFormat::JsonObject => { - let text = response.text(); - serde_json::from_str::(&text).map_err(|e| { - TinyAgentsError::StructuredOutput(format!( - "response text is not valid JSON: {e}" - )) - })?; - Ok(()) - } - ResponseFormat::JsonSchema { name, schema } | ResponseFormat::Auto { name, schema } => { - let extractor = StructuredExtractor::new( - StructuredStrategy::ProviderSchema, - name.clone(), - schema.clone(), - ); - extractor.extract(response)?; - Ok(()) - } - } - } -} - -// ── DynamicPromptMiddleware ─────────────────────────────────────────────────── - -impl DynamicPromptMiddleware { - /// Creates a dynamic-prompt middleware deriving a system message from - /// `prompt`. - pub fn new(prompt: PromptFn) -> Self { - Self { - label: "dynamic_prompt", - prompt, - _marker: PhantomData, - } - } - - /// Creates a dynamic-prompt middleware from a closure over the shared state - /// and the run's [`RunConfig`]. - pub fn from_fn(f: F) -> Self - where - F: Fn(&State, &RunConfig) -> Option + Send + Sync + 'static, - { - Self::new(Arc::new(f)) - } -} - -#[async_trait] -impl Middleware - for DynamicPromptMiddleware -{ - fn name(&self) -> &str { - self.label - } - - async fn before_model( - &self, - ctx: &mut RunContext, - state: &State, - request: &mut ModelRequest, - ) -> Result<()> { - if let Some(text) = (self.prompt)(state, &ctx.config) { - request.messages.insert(0, Message::system(text)); - } - Ok(()) - } -} - -// ── RedactionMiddleware ─────────────────────────────────────────────────────── - -impl RedactionMiddleware { - /// Creates a redaction middleware replacing each pattern with `"[REDACTED]"`. - pub fn new(patterns: impl IntoIterator>) -> Self { - Self::with_mask(patterns, "[REDACTED]") - } - - /// Creates a redaction middleware replacing each pattern with `mask`. - pub fn with_mask( - patterns: impl IntoIterator>, - mask: impl Into, - ) -> Self { - Self { - label: "redaction", - patterns: patterns - .into_iter() - .map(Into::into) - .filter(|p| !p.is_empty()) - .collect(), - mask: mask.into(), - redactions: Mutex::new(0), - } - } - - /// Returns the total number of pattern occurrences redacted so far. - pub fn redactions(&self) -> usize { - *self.redactions.lock().expect("redactions mutex poisoned") - } - - /// Replaces every configured pattern in `text`, returning the redacted - /// string and the number of occurrences replaced. - fn redact(&self, text: &str) -> (String, usize) { - let mut out = text.to_string(); - let mut hits = 0usize; - for pattern in &self.patterns { - let occurrences = out.matches(pattern.as_str()).count(); - if occurrences > 0 { - hits += occurrences; - out = out.replace(pattern.as_str(), &self.mask); - } - } - (out, hits) - } - - /// Records `hits` redactions against the running total. - fn record(&self, hits: usize) { - if hits > 0 { - *self.redactions.lock().expect("redactions mutex poisoned") += hits; - } - } -} - -#[async_trait] -impl Middleware for RedactionMiddleware { - fn name(&self) -> &str { - self.label - } - - async fn after_model( - &self, - _ctx: &mut RunContext, - _state: &State, - response: &mut ModelResponse, - ) -> Result<()> { - let mut hits = 0usize; - for block in &mut response.message.content { - if let ContentBlock::Text(text) = block { - let (redacted, n) = self.redact(text); - if n > 0 { - *text = redacted; - hits += n; - } - } - } - self.record(hits); - Ok(()) - } - - async fn after_tool( - &self, - _ctx: &mut RunContext, - _state: &State, - result: &mut ToolResult, - ) -> Result<()> { - let (redacted, hits) = self.redact(&result.content); - if hits > 0 { - result.content = redacted; - } - self.record(hits); - Ok(()) - } -} - -// ── TracingMiddleware ───────────────────────────────────────────────────────── - -impl TracingMiddleware { - /// Creates a tracing middleware with the default label `"tracing"`. - pub fn new() -> Self { - Self::with_label("tracing") - } - - /// Creates a tracing middleware with a custom static label. - pub fn with_label(label: &'static str) -> Self { - Self { - label, - records: Mutex::new(Vec::new()), - counts: Mutex::new(TraceCounts::default()), - } - } - - /// Returns the structured begin/end traces recorded so far, in order. - pub fn records(&self) -> Vec { - self.records.lock().expect("records mutex poisoned").clone() - } - - /// Returns a snapshot of the per-phase begin counts. - pub fn counts(&self) -> TraceCounts { - self.counts.lock().expect("counts mutex poisoned").clone() - } - - fn push(&self, phase: &'static str, boundary: TraceBoundary) { - self.records - .lock() - .expect("records mutex poisoned") - .push(PhaseTrace { phase, boundary }); - } -} - -impl Default for TracingMiddleware { - fn default() -> Self { - Self::new() - } -} - -#[async_trait] -impl Middleware for TracingMiddleware { - fn name(&self) -> &str { - self.label - } - - async fn before_agent(&self, _ctx: &mut RunContext, _state: &State) -> Result<()> { - self.counts.lock().expect("counts mutex poisoned").agent += 1; - self.push("agent", TraceBoundary::Begin); - Ok(()) - } - - async fn after_agent( - &self, - _ctx: &mut RunContext, - _state: &State, - _run: &mut crate::harness::middleware::AgentRun, - ) -> Result<()> { - self.push("agent", TraceBoundary::End); - Ok(()) - } - - async fn before_model( - &self, - _ctx: &mut RunContext, - _state: &State, - _request: &mut ModelRequest, - ) -> Result<()> { - self.counts.lock().expect("counts mutex poisoned").model += 1; - self.push("model", TraceBoundary::Begin); - Ok(()) - } - - async fn on_model_delta( - &self, - _ctx: &mut RunContext, - _state: &State, - _delta: &mut ModelDelta, - ) -> Result<()> { - self.counts.lock().expect("counts mutex poisoned").delta += 1; - Ok(()) - } - - async fn after_model( - &self, - _ctx: &mut RunContext, - _state: &State, - _response: &mut ModelResponse, - ) -> Result<()> { - self.push("model", TraceBoundary::End); - Ok(()) - } - - async fn before_tool( - &self, - _ctx: &mut RunContext, - _state: &State, - _call: &mut ToolCall, - ) -> Result<()> { - self.counts.lock().expect("counts mutex poisoned").tool += 1; - self.push("tool", TraceBoundary::Begin); - Ok(()) - } - - async fn on_tool_delta( - &self, - _ctx: &mut RunContext, - _state: &State, - _delta: &mut ToolDelta, - ) -> Result<()> { - self.counts.lock().expect("counts mutex poisoned").delta += 1; - Ok(()) - } - - async fn after_tool( - &self, - _ctx: &mut RunContext, - _state: &State, - _result: &mut ToolResult, - ) -> Result<()> { - self.push("tool", TraceBoundary::End); - Ok(()) - } - - async fn on_error(&self, _ctx: &mut RunContext, _error: &TinyAgentsError) -> Result<()> { - self.counts.lock().expect("counts mutex poisoned").error += 1; - Ok(()) - } -} +mod budget; +mod context; +mod observe; +mod resilience; +mod tool_policy; #[cfg(test)] mod test; diff --git a/src/harness/middleware/library/observe.rs b/src/harness/middleware/library/observe.rs new file mode 100644 index 0000000..a2d410c --- /dev/null +++ b/src/harness/middleware/library/observe.rs @@ -0,0 +1,544 @@ +//! Observation middleware: structured-output validation, dynamic prompt +//! injection, redaction, and tracing. +//! +//! Split out of `library/mod.rs`; see that module's doc comment for the +//! full built-in middleware library overview. + +use super::*; +use crate::harness::middleware::{ + AgentRun, HookCounts, LoggingMiddleware, UsageAccountingMiddleware, +}; +use crate::harness::usage::UsageTotals; + +// ── StructuredOutputValidatorMiddleware ─────────────────────────────────────── + +impl StructuredOutputValidatorMiddleware { + /// Creates a validator middleware checking responses against `format`. + pub fn new(format: ResponseFormat) -> Self { + Self { + label: "structured_output_validator", + format, + } + } +} + +#[async_trait] +impl Middleware + for StructuredOutputValidatorMiddleware +{ + fn name(&self) -> &str { + self.label + } + + async fn after_model( + &self, + _ctx: &mut RunContext, + _state: &State, + response: &mut ModelResponse, + ) -> Result<()> { + match &self.format { + ResponseFormat::Text => Ok(()), + ResponseFormat::JsonObject => { + let text = response.text(); + serde_json::from_str::(&text).map_err(|e| { + TinyAgentsError::StructuredOutput(format!( + "response text is not valid JSON: {e}" + )) + })?; + Ok(()) + } + ResponseFormat::JsonSchema { name, schema } | ResponseFormat::Auto { name, schema } => { + let extractor = StructuredExtractor::new( + StructuredStrategy::ProviderSchema, + name.clone(), + schema.clone(), + ); + extractor.extract(response)?; + Ok(()) + } + } + } +} + +// ── DynamicPromptMiddleware ─────────────────────────────────────────────────── + +impl DynamicPromptMiddleware { + /// Creates a dynamic-prompt middleware deriving a system message from + /// `prompt`. + pub fn new(prompt: PromptFn) -> Self { + Self { + label: "dynamic_prompt", + prompt, + _marker: PhantomData, + } + } + + /// Creates a dynamic-prompt middleware from a closure over the shared state + /// and the run's [`RunConfig`]. + pub fn from_fn(f: F) -> Self + where + F: Fn(&State, &RunConfig) -> Option + Send + Sync + 'static, + { + Self::new(Arc::new(f)) + } +} + +#[async_trait] +impl Middleware + for DynamicPromptMiddleware +{ + fn name(&self) -> &str { + self.label + } + + async fn before_model( + &self, + ctx: &mut RunContext, + state: &State, + request: &mut ModelRequest, + ) -> Result<()> { + if let Some(text) = (self.prompt)(state, &ctx.config) { + request.messages.insert(0, Message::system(text)); + } + Ok(()) + } +} + +// ── RedactionMiddleware ─────────────────────────────────────────────────────── + +impl RedactionMiddleware { + /// Creates a redaction middleware replacing each pattern with `"[REDACTED]"`. + pub fn new(patterns: impl IntoIterator>) -> Self { + Self::with_mask(patterns, "[REDACTED]") + } + + /// Creates a redaction middleware replacing each pattern with `mask`. + pub fn with_mask( + patterns: impl IntoIterator>, + mask: impl Into, + ) -> Self { + Self { + label: "redaction", + patterns: patterns + .into_iter() + .map(Into::into) + .filter(|p| !p.is_empty()) + .collect(), + mask: mask.into(), + redactions: Mutex::new(0), + } + } + + /// Returns the total number of pattern occurrences redacted so far. + pub fn redactions(&self) -> usize { + *self.redactions.lock().expect("redactions mutex poisoned") + } + + /// Replaces every configured pattern in `text`, returning the redacted + /// string and the number of occurrences replaced. + fn redact(&self, text: &str) -> (String, usize) { + let mut out = text.to_string(); + let mut hits = 0usize; + for pattern in &self.patterns { + let occurrences = out.matches(pattern.as_str()).count(); + if occurrences > 0 { + hits += occurrences; + out = out.replace(pattern.as_str(), &self.mask); + } + } + (out, hits) + } + + /// Records `hits` redactions against the running total. + fn record(&self, hits: usize) { + if hits > 0 { + *self.redactions.lock().expect("redactions mutex poisoned") += hits; + } + } +} + +#[async_trait] +impl Middleware for RedactionMiddleware { + fn name(&self) -> &str { + self.label + } + + async fn after_model( + &self, + _ctx: &mut RunContext, + _state: &State, + response: &mut ModelResponse, + ) -> Result<()> { + let mut hits = 0usize; + for block in &mut response.message.content { + if let ContentBlock::Text(text) = block { + let (redacted, n) = self.redact(text); + if n > 0 { + *text = redacted; + hits += n; + } + } + } + self.record(hits); + Ok(()) + } + + async fn after_tool( + &self, + _ctx: &mut RunContext, + _state: &State, + result: &mut ToolResult, + ) -> Result<()> { + let (redacted, hits) = self.redact(&result.content); + if hits > 0 { + result.content = redacted; + } + self.record(hits); + Ok(()) + } +} + +// ── TracingMiddleware ───────────────────────────────────────────────────────── + +impl TracingMiddleware { + /// Creates a tracing middleware with the default label `"tracing"`. + pub fn new() -> Self { + Self::with_label("tracing") + } + + /// Creates a tracing middleware with a custom static label. + pub fn with_label(label: &'static str) -> Self { + Self { + label, + records: Mutex::new(VecDeque::new()), + counts: Mutex::new(TraceCounts::default()), + max_records: DEFAULT_TRACE_RECORD_CAP, + } + } + + /// Sets the maximum number of [`PhaseTrace`] entries retained before the + /// oldest is evicted. `0` disables recording entirely (counts are still + /// tracked). + pub fn with_max_records(mut self, max_records: usize) -> Self { + self.max_records = max_records; + let mut records = self.records.lock().expect("records mutex poisoned"); + while records.len() > max_records { + records.pop_front(); + } + drop(records); + self + } + + /// Returns the structured begin/end traces recorded so far, in order. + /// Bounded to at most [`TracingMiddleware::with_max_records`] entries + /// (default [`DEFAULT_TRACE_RECORD_CAP`]); older traces are evicted first. + pub fn records(&self) -> Vec { + self.records + .lock() + .expect("records mutex poisoned") + .iter() + .cloned() + .collect() + } + + /// Returns a snapshot of the per-phase begin counts. + pub fn counts(&self) -> TraceCounts { + self.counts.lock().expect("counts mutex poisoned").clone() + } + + fn push(&self, phase: &'static str, boundary: TraceBoundary) { + let mut records = self.records.lock().expect("records mutex poisoned"); + if self.max_records == 0 { + return; + } + if records.len() >= self.max_records { + records.pop_front(); + } + records.push_back(PhaseTrace { phase, boundary }); + } +} + +impl Default for TracingMiddleware { + fn default() -> Self { + Self::new() + } +} + +#[async_trait] +impl Middleware for TracingMiddleware { + fn name(&self) -> &str { + self.label + } + + async fn before_agent(&self, _ctx: &mut RunContext, _state: &State) -> Result<()> { + self.counts.lock().expect("counts mutex poisoned").agent += 1; + self.push("agent", TraceBoundary::Begin); + Ok(()) + } + + async fn after_agent( + &self, + _ctx: &mut RunContext, + _state: &State, + _run: &mut crate::harness::middleware::AgentRun, + ) -> Result<()> { + self.push("agent", TraceBoundary::End); + Ok(()) + } + + async fn before_model( + &self, + _ctx: &mut RunContext, + _state: &State, + _request: &mut ModelRequest, + ) -> Result<()> { + self.counts.lock().expect("counts mutex poisoned").model += 1; + self.push("model", TraceBoundary::Begin); + Ok(()) + } + + async fn on_model_delta( + &self, + _ctx: &mut RunContext, + _state: &State, + _delta: &mut ModelDelta, + ) -> Result<()> { + self.counts.lock().expect("counts mutex poisoned").delta += 1; + Ok(()) + } + + async fn after_model( + &self, + _ctx: &mut RunContext, + _state: &State, + _response: &mut ModelResponse, + ) -> Result<()> { + self.push("model", TraceBoundary::End); + Ok(()) + } + + async fn before_tool( + &self, + _ctx: &mut RunContext, + _state: &State, + _call: &mut ToolCall, + ) -> Result<()> { + self.counts.lock().expect("counts mutex poisoned").tool += 1; + self.push("tool", TraceBoundary::Begin); + Ok(()) + } + + async fn on_tool_delta( + &self, + _ctx: &mut RunContext, + _state: &State, + _delta: &mut ToolDelta, + ) -> Result<()> { + self.counts.lock().expect("counts mutex poisoned").delta += 1; + Ok(()) + } + + async fn after_tool( + &self, + _ctx: &mut RunContext, + _state: &State, + _result: &mut ToolResult, + ) -> Result<()> { + self.push("tool", TraceBoundary::End); + Ok(()) + } + + async fn on_error(&self, _ctx: &mut RunContext, _error: &TinyAgentsError) -> Result<()> { + self.counts.lock().expect("counts mutex poisoned").error += 1; + Ok(()) + } +} + +// ── LoggingMiddleware ───────────────────────────────────────────────────────── + +// ── LoggingMiddleware ───────────────────────────────────────────────────────── + +impl LoggingMiddleware { + /// Creates a logging middleware with the default label `"logging"`. + pub fn new() -> Self { + Self::with_label("logging") + } + + /// Creates a logging middleware with a custom static label. + pub fn with_label(label: &'static str) -> Self { + Self { + label, + counts: std::sync::Mutex::new(HookCounts::default()), + } + } + + /// Returns a snapshot of the per-hook invocation counts recorded so far. + pub fn counts(&self) -> HookCounts { + self.counts.lock().expect("counts mutex poisoned").clone() + } +} + +impl Default for LoggingMiddleware { + fn default() -> Self { + Self::new() + } +} + +#[async_trait] +impl Middleware for LoggingMiddleware { + fn name(&self) -> &str { + self.label + } + + async fn before_agent(&self, _ctx: &mut RunContext, _state: &State) -> Result<()> { + self.counts + .lock() + .expect("counts mutex poisoned") + .before_agent += 1; + Ok(()) + } + + async fn after_agent( + &self, + _ctx: &mut RunContext, + _state: &State, + _run: &mut AgentRun, + ) -> Result<()> { + self.counts + .lock() + .expect("counts mutex poisoned") + .after_agent += 1; + Ok(()) + } + + async fn before_model( + &self, + _ctx: &mut RunContext, + _state: &State, + _request: &mut ModelRequest, + ) -> Result<()> { + self.counts + .lock() + .expect("counts mutex poisoned") + .before_model += 1; + Ok(()) + } + + async fn on_model_delta( + &self, + _ctx: &mut RunContext, + _state: &State, + _delta: &mut ModelDelta, + ) -> Result<()> { + self.counts + .lock() + .expect("counts mutex poisoned") + .on_model_delta += 1; + Ok(()) + } + + async fn after_model( + &self, + _ctx: &mut RunContext, + _state: &State, + _response: &mut ModelResponse, + ) -> Result<()> { + self.counts + .lock() + .expect("counts mutex poisoned") + .after_model += 1; + Ok(()) + } + + async fn before_tool( + &self, + _ctx: &mut RunContext, + _state: &State, + _call: &mut ToolCall, + ) -> Result<()> { + self.counts + .lock() + .expect("counts mutex poisoned") + .before_tool += 1; + Ok(()) + } + + async fn on_tool_delta( + &self, + _ctx: &mut RunContext, + _state: &State, + _delta: &mut ToolDelta, + ) -> Result<()> { + self.counts + .lock() + .expect("counts mutex poisoned") + .on_tool_delta += 1; + Ok(()) + } + + async fn after_tool( + &self, + _ctx: &mut RunContext, + _state: &State, + _result: &mut ToolResult, + ) -> Result<()> { + self.counts + .lock() + .expect("counts mutex poisoned") + .after_tool += 1; + Ok(()) + } + + async fn on_error(&self, _ctx: &mut RunContext, _error: &TinyAgentsError) -> Result<()> { + self.counts.lock().expect("counts mutex poisoned").on_error += 1; + Ok(()) + } +} + +// ── UsageAccountingMiddleware ───────────────────────────────────────────────── + +// ── UsageAccountingMiddleware ───────────────────────────────────────────────── + +impl UsageAccountingMiddleware { + /// Creates a usage-accounting middleware with the default label + /// `"usage_accounting"`. + pub fn new() -> Self { + Self { + label: "usage_accounting", + totals: std::sync::Mutex::new(UsageTotals::new()), + } + } + + /// Returns a snapshot of the accumulated usage totals. + pub fn totals(&self) -> UsageTotals { + *self.totals.lock().expect("totals mutex poisoned") + } +} + +impl Default for UsageAccountingMiddleware { + fn default() -> Self { + Self::new() + } +} + +#[async_trait] +impl Middleware for UsageAccountingMiddleware { + fn name(&self) -> &str { + self.label + } + + async fn after_model( + &self, + _ctx: &mut RunContext, + _state: &State, + response: &mut ModelResponse, + ) -> Result<()> { + if let Some(usage) = response.usage { + self.totals + .lock() + .expect("totals mutex poisoned") + .record(usage); + } + Ok(()) + } +} diff --git a/src/harness/middleware/library/resilience.rs b/src/harness/middleware/library/resilience.rs new file mode 100644 index 0000000..3c81c1b --- /dev/null +++ b/src/harness/middleware/library/resilience.rs @@ -0,0 +1,256 @@ +//! Resilience middleware: retry, timeout, model fallback, rate limiting. +//! +//! Split out of `library/mod.rs`; see that module's doc comment for the +//! full built-in middleware library overview. + +use super::*; + +// ── RetryMiddleware ─────────────────────────────────────────────────────────── + +impl RetryMiddleware { + /// Creates a retry middleware using the given [`RetryPolicy`]. + pub fn new(policy: RetryPolicy) -> Self { + Self { + label: "retry", + policy, + } + } + + /// Creates a retry middleware with the default [`RetryPolicy`]. + pub fn with_default_policy() -> Self { + Self::new(RetryPolicy::default()) + } + + /// Returns the policy-derived backoff for the given retry `attempt`. + /// + /// Exposed for callers that want to inspect the backoff. The middleware + /// itself sleeps between retries only when the policy opts in via + /// [`RetryPolicy::with_backoff_sleep`]; otherwise it retries back-to-back. + pub fn backoff_for_attempt(&self, attempt: usize) -> Duration { + self.policy.backoff_for_attempt(attempt) + } +} + +#[async_trait] +impl ModelMiddleware for RetryMiddleware { + fn name(&self) -> &str { + self.label + } + + async fn wrap_model( + &self, + ctx: &mut RunContext, + state: &State, + request: ModelRequest, + next: ModelHandler<'_, State, Ctx>, + ) -> Result { + let mut attempt = 0usize; + loop { + match next.run(ctx, state, request.clone()).await { + Ok(outcome) => return Ok(outcome), + Err(error) => { + if self.policy.should_retry_error(attempt, &error) { + // Compute the backoff from the *pre-increment* attempt + // number: `attempt == 0` is the first retry and must + // sleep `initial_backoff_ms` + // (`RetryPolicy::backoff_for_attempt(0)`). Sleeping on + // the post-increment value skipped `initial_backoff_ms` + // entirely and shifted the whole exponential schedule + // one step too high. + let backoff_attempt = attempt; + attempt += 1; + let call_id = CallId::new(format!("{}-model", ctx.run_id())); + ctx.emit(AgentEvent::RetryScheduled { call_id, attempt }); + // Sleep for the backoff only when the policy opts in + // (`with_backoff_sleep`); a no-op otherwise. + self.policy.sleep_backoff(backoff_attempt).await; + continue; + } + return Err(error); + } + } + } + } +} + +// ── TimeoutMiddleware ───────────────────────────────────────────────────────── + +impl TimeoutMiddleware { + /// Creates a timeout middleware bounding each model call to `timeout`. + pub fn new(timeout: Duration) -> Self { + Self { + label: "timeout", + timeout, + } + } + + /// Creates a timeout middleware from a millisecond duration. + pub fn from_millis(ms: u64) -> Self { + Self::new(Duration::from_millis(ms)) + } +} + +#[async_trait] +impl ModelMiddleware for TimeoutMiddleware { + fn name(&self) -> &str { + self.label + } + + async fn wrap_model( + &self, + ctx: &mut RunContext, + state: &State, + request: ModelRequest, + next: ModelHandler<'_, State, Ctx>, + ) -> Result { + let run_id = ctx.run_id().as_str().to_string(); + let fut = next.run(ctx, state, request); + match tokio::time::timeout(self.timeout, fut).await { + Ok(result) => result, + Err(_) => Err(TinyAgentsError::Timeout(format!( + "model call for run `{run_id}` exceeded the {} ms middleware timeout", + self.timeout.as_millis() + ))), + } + } +} + +// ── ModelFallbackMiddleware ─────────────────────────────────────────────────── + +impl ModelFallbackMiddleware { + /// Creates a fallback middleware that tries each model name in order after + /// the primary call fails. + pub fn new(fallbacks: impl IntoIterator>) -> Self { + Self { + label: "model_fallback", + fallbacks: fallbacks.into_iter().map(Into::into).collect(), + } + } +} + +#[async_trait] +impl ModelMiddleware for ModelFallbackMiddleware { + fn name(&self) -> &str { + self.label + } + + async fn wrap_model( + &self, + ctx: &mut RunContext, + state: &State, + request: ModelRequest, + next: ModelHandler<'_, State, Ctx>, + ) -> Result { + match next.run(ctx, state, request.clone()).await { + Ok(outcome) => Ok(outcome), + Err(mut last_error) => { + let mut current = request.model.clone().unwrap_or_default(); + for fallback in &self.fallbacks { + // Only a *transient* failure justifies trying another model: + // a non-retryable error (auth/validation/schema) will fail + // the same way on every backend, so switching burns quota + // and latency for nothing. Classification is shared with the + // rest of the harness via `is_retryable`. + if !is_retryable(&last_error) { + break; + } + ctx.emit(AgentEvent::FallbackSelected { + from: current.clone(), + to: fallback.clone(), + }); + let mut req = request.clone(); + req.model = Some(fallback.clone()); + match next.run(ctx, state, req).await { + Ok(outcome) => return Ok(outcome), + Err(error) => { + last_error = error; + current = fallback.clone(); + } + } + } + Err(last_error) + } + } + } +} + +// ── RateLimitMiddleware ─────────────────────────────────────────────────────── + +impl RateLimitMiddleware { + /// Creates a rate-limit middleware gating one token per call through + /// `limiter`, failing immediately when the bucket is empty + /// ([`RateLimitBehavior::Error`]). + pub fn new(limiter: Arc) -> Self { + Self { + label: "rate_limit", + limiter, + tokens: 1, + behavior: RateLimitBehavior::Error, + poll_interval: Duration::from_millis(50), + now: Arc::new(Instant::now), + } + } + + /// Sets the number of tokens each call consumes. + pub fn with_tokens(mut self, tokens: u64) -> Self { + self.tokens = tokens; + self + } + + /// Sets the behavior when the bucket lacks capacity. + pub fn with_behavior(mut self, behavior: RateLimitBehavior) -> Self { + self.behavior = behavior; + self + } + + /// Switches to [`RateLimitBehavior::Wait`] with the given poll interval. + pub fn waiting(mut self, poll_interval: Duration) -> Self { + self.behavior = RateLimitBehavior::Wait; + self.poll_interval = poll_interval; + self + } + + /// Replaces the clock used to read the current instant (for deterministic + /// tests). + pub fn with_clock(mut self, now: NowFn) -> Self { + self.now = now; + self + } +} + +#[async_trait] +impl ModelMiddleware for RateLimitMiddleware { + fn name(&self) -> &str { + self.label + } + + async fn wrap_model( + &self, + ctx: &mut RunContext, + state: &State, + request: ModelRequest, + next: ModelHandler<'_, State, Ctx>, + ) -> Result { + loop { + let now = (self.now)(); + if self.limiter.try_acquire(self.tokens, now) { + break; + } + match self.behavior { + RateLimitBehavior::Error => { + return Err(TinyAgentsError::LimitExceeded(format!( + "rate limit: could not acquire {} token(s)", + self.tokens + ))); + } + RateLimitBehavior::Wait => { + ctx.emit(AgentEvent::RateLimitWaited { + waited_ms: self.poll_interval.as_millis() as u64, + }); + tokio::time::sleep(self.poll_interval).await; + } + } + } + next.run(ctx, state, request).await + } +} diff --git a/src/harness/middleware/library/test.rs b/src/harness/middleware/library/test.rs index f859067..c8e9696 100644 --- a/src/harness/middleware/library/test.rs +++ b/src/harness/middleware/library/test.rs @@ -114,6 +114,50 @@ async fn retry_middleware_retries_then_succeeds() { assert_eq!(scheduled, 2); } +#[tokio::test(start_paused = true)] +async fn retry_middleware_sleeps_the_documented_backoff_schedule() { + // Regression test: the middleware used to compute the backoff from the + // *post-increment* attempt number, so the first retry's sleep skipped + // `initial_backoff_ms` entirely and the whole exponential schedule was + // shifted one step higher than `RetryPolicy::backoff_for_attempt` + // documents. With `initial_backoff_ms = 100`, `multiplier = 2.0`, no + // jitter: attempt 0 -> 100ms, attempt 1 -> 200ms. + let (mut ctx, _recorder) = ctx_with_recorder(); + let mut stack: MiddlewareStack<()> = MiddlewareStack::new(); + stack.push_model_middleware(Arc::new(RetryMiddleware::new( + RetryPolicy::default() + .with_max_attempts(3) + .with_initial_backoff_ms(100) + .with_multiplier(2.0) + .with_jitter(false) + .with_backoff_sleep(true), + ))); + + let timestamps = Arc::new(std::sync::Mutex::new(Vec::new())); + let recorded = timestamps.clone(); + let base = FakeModelBase::new(move |_n, _req| { + recorded.lock().unwrap().push(tokio::time::Instant::now()); + Err(TinyAgentsError::Model("transient".to_string())) + }); + + stack + .run_wrapped_model(&mut ctx, &(), ModelRequest::default(), &base) + .await + .expect_err("all 3 attempts fail"); + + let timestamps = timestamps.lock().unwrap().clone(); + assert_eq!(timestamps.len(), 3, "expected exactly max_attempts calls"); + let gaps: Vec = timestamps + .windows(2) + .map(|w| w[1].duration_since(w[0])) + .collect(); + assert_eq!( + gaps, + vec![Duration::from_millis(100), Duration::from_millis(200)], + "backoff schedule does not match RetryPolicy::backoff_for_attempt" + ); +} + #[tokio::test] async fn retry_middleware_does_not_retry_non_retryable() { let (mut ctx, _recorder) = ctx_with_recorder(); @@ -216,6 +260,32 @@ async fn model_fallback_switches_model_on_error() { ); } +#[tokio::test] +async fn model_fallback_does_not_switch_on_non_retryable() { + let (mut ctx, recorder) = ctx_with_recorder(); + let mut stack: MiddlewareStack<()> = MiddlewareStack::new(); + stack.push_model_middleware(Arc::new(ModelFallbackMiddleware::new([ + "backup-a", "backup-b", + ]))); + + // A non-retryable error (validation) will fail identically on every model, + // so the middleware must not burn quota switching backends. + let base = + FakeModelBase::new(|_n, _req| Err(TinyAgentsError::Validation("bad input".to_string()))); + let err = stack + .run_wrapped_model(&mut ctx, &(), ModelRequest::default(), &base) + .await + .expect_err("validation errors are not retryable"); + assert!(matches!(err, TinyAgentsError::Validation(_))); + // Only the primary call happens; no fallback attempts. + assert_eq!(base.calls(), 1); + let selections = events(&recorder) + .into_iter() + .filter(|e| matches!(e, AgentEvent::FallbackSelected { .. })) + .count(); + assert_eq!(selections, 0); +} + #[tokio::test] async fn model_fallback_returns_last_error_when_all_fail() { let (mut ctx, _recorder) = ctx_with_recorder(); @@ -599,6 +669,161 @@ async fn budget_preflight_reserves_and_reconciles() { )); } +#[tokio::test] +async fn reservation_is_released_even_when_response_carries_no_usage() { + // A call whose response never reports `usage` (provider error, dropped + // stream, etc.) must still release its preflight reservation in + // `after_model` — otherwise the reservation leaks forever and + // permanently starves later calls on this (or a shared) tracker. + let (mut ctx, _recorder) = ctx_with_recorder(); + let mw = BudgetMiddleware::new(BudgetLimits { + max_input_tokens: Some(5), + ..BudgetLimits::default() + }); + let tracker = mw.tracker(); + let mut stack: MiddlewareStack<()> = MiddlewareStack::new(); + stack.push(Arc::new(mw)); + + let mut req = ModelRequest::new(vec![Message::user("hi")]); + stack + .run_before_model(&mut ctx, &(), &mut req) + .await + .expect("small call fits the reservation"); + assert!(tracker.snapshot().reserved_input_total > 0); + + let mut resp = ModelResponse::assistant("ok"); // usage: None + stack + .run_after_model(&mut ctx, &(), &mut resp) + .await + .unwrap(); + assert_eq!(tracker.snapshot().reserved_input_total, 0); +} + +#[tokio::test] +async fn reservation_is_released_when_model_call_errors() { + // A model call that ultimately errors (retries/fallback exhausted, hard + // provider error, timeout, ...) never reaches `after_model` — the wrap + // onion returns `Err` and the harness propagates it with `?` before the + // lifecycle `after_model` hook runs. The reservation `before_model` added + // must still be released (via `on_error`), or it leaks forever and + // permanently starves every later call on a shared tracker. + let (mut ctx, _recorder) = ctx_with_recorder(); + let mw = BudgetMiddleware::new(BudgetLimits { + max_input_tokens: Some(5), + ..BudgetLimits::default() + }); + let tracker = mw.tracker(); + let mut stack: MiddlewareStack<()> = MiddlewareStack::new(); + stack.push(Arc::new(mw)); + + let mut req = ModelRequest::new(vec![Message::user("hi")]); + stack + .run_before_model(&mut ctx, &(), &mut req) + .await + .expect("small call fits the reservation"); + assert!(tracker.snapshot().reserved_input_total > 0); + + // Simulate the model call failing before `after_model` ever runs: the + // harness surfaces the failure to every middleware via `on_error` + // instead. + let error = TinyAgentsError::Model("provider down".to_string()); + stack.run_on_error(&mut ctx, &error).await.unwrap(); + + assert_eq!(tracker.snapshot().reserved_input_total, 0); + + // The tracker must have recovered: a fresh call reserving the same + // amount should not be rejected as if capacity were still consumed. + let mut req2 = ModelRequest::new(vec![Message::user("hi")]); + stack + .run_before_model(&mut ctx, &(), &mut req2) + .await + .expect("reservation was released so a new call fits the budget again"); +} + +#[test] +fn poisoned_tracker_stays_fail_closed() { + // A poisoned mutex still holds a valid last-written spend value; a + // `snapshot()` that defaulted to zero on poison would make every + // subsequent budget check see an empty budget and admit calls forever + // (fail-open). Recovering the poisoned guard must keep the previously + // accumulated spend intact. + use crate::harness::cost::CostTotals; + use crate::harness::usage::Usage; + + let tracker = BudgetTracker::new(); + tracker.record(Usage::new(100, 100), CostTotals::default()); + assert!(!tracker.inner.is_poisoned()); + + let poison_tracker = tracker.clone(); + let _ = std::thread::spawn(move || { + let _guard = poison_tracker.inner.lock().unwrap(); + panic!("simulated panic while holding the tracker lock"); + }) + .join(); + assert!(tracker.inner.is_poisoned()); + + // The accumulated spend from before the panic must still be visible, + // not silently reset to zero. + let spend = tracker.snapshot(); + assert_eq!(spend.usage.usage.input_tokens, 100); + assert_eq!(spend.usage.usage.output_tokens, 100); + + let limits = BudgetLimits { + max_total_tokens: Some(200), + ..BudgetLimits::default() + }; + assert!(limits.exceeded_reason(&spend).is_some()); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 8)] +async fn shared_tracker_reservation_is_atomic_under_concurrency() { + // N concurrent "runs" (their own `BudgetMiddleware`, sharing one + // `BudgetTracker` — the documented sub-agent pattern) all reserve the + // same input-token estimate at once. A racy check-then-act preflight + // would let every one of them observe spare capacity and reserve past + // it; an atomic check-and-reserve admits only as many as the budget + // actually allows. + use crate::harness::middleware::Middleware; + + let tracker = BudgetTracker::new(); + let per_call_tokens = 10u64; // "x" * 40 chars / 4 == 10 estimated tokens. + let concurrent_capacity = 4u64; + let attempts = 10usize; + let limits = BudgetLimits { + max_input_tokens: Some(per_call_tokens * concurrent_capacity), + ..BudgetLimits::default() + }; + + let barrier = Arc::new(tokio::sync::Barrier::new(attempts)); + let mut handles = Vec::new(); + for _ in 0..attempts { + let mw = BudgetMiddleware::new(limits).with_tracker(tracker.clone()); + let barrier = barrier.clone(); + handles.push(tokio::spawn(async move { + let mut ctx: RunContext = RunContext::new(RunConfig::new("test-run"), ()); + let mut request = ModelRequest::new(vec![Message::user("x".repeat(40))]); + barrier.wait().await; + mw.before_model(&mut ctx, &(), &mut request).await + })); + } + + let mut ok_count = 0usize; + for handle in handles { + if handle.await.unwrap().is_ok() { + ok_count += 1; + } + } + + // Exactly the number of reservations the budget can hold succeed; the + // rest are rejected before dispatch. A racy implementation would let + // more than `concurrent_capacity` through. + assert_eq!(ok_count, concurrent_capacity as usize); + assert_eq!( + tracker.snapshot().reserved_input_total, + per_call_tokens * concurrent_capacity + ); +} + // ── ToolAllowlistMiddleware ───────────────────────────────────────────────── #[tokio::test] @@ -1020,3 +1245,35 @@ async fn tracing_records_phase_boundaries_and_counts() { assert_eq!(records.last().unwrap().phase, "agent"); assert_eq!(records.last().unwrap().boundary, TraceBoundary::End); } + +#[tokio::test] +async fn tracing_records_are_bounded_by_max_records() { + let (mut ctx, _recorder) = ctx_with_recorder(); + let tracing = Arc::new(TracingMiddleware::new().with_max_records(3)); + let mut stack: MiddlewareStack<()> = MiddlewareStack::new(); + stack.push(tracing.clone()); + + // Each `before_agent`/`after_agent` pair pushes 2 records; run enough + // iterations to far exceed the cap of 3 and confirm memory stays bounded + // rather than growing without limit. + for _ in 0..50 { + stack.run_before_agent(&mut ctx, &()).await.unwrap(); + let mut run = crate::harness::middleware::AgentRun::new(); + stack + .run_after_agent(&mut ctx, &(), &mut run) + .await + .unwrap(); + } + + let records = tracing.records(); + assert_eq!(records.len(), 3); + // The oldest entries were evicted; only the most recent boundary pair + // (plus one) survives. + assert_eq!(records.last().unwrap().phase, "agent"); + assert_eq!(records.last().unwrap().boundary, TraceBoundary::End); + + // Counts are unaffected by the record cap — they track unboundedly by + // design (a single `usize` per phase, not a growing collection). + let counts = tracing.counts(); + assert_eq!(counts.agent, 50); +} diff --git a/src/harness/middleware/library/tool_policy.rs b/src/harness/middleware/library/tool_policy.rs new file mode 100644 index 0000000..11eede2 --- /dev/null +++ b/src/harness/middleware/library/tool_policy.rs @@ -0,0 +1,469 @@ +//! Tool policy and selection middleware: allowlisting, classification-based +//! policy enforcement, dynamic/contextual tool selection, and human +//! approval. +//! +//! Split out of `library/mod.rs`; see that module's doc comment for the +//! full built-in middleware library overview. + +use super::*; + +// ── ToolAllowlistMiddleware ─────────────────────────────────────────────────── + +impl ToolAllowlistMiddleware { + /// Creates an allowlist middleware permitting only the named tools. + pub fn new(allowed: impl IntoIterator>) -> Self { + Self { + label: "tool_allowlist", + allowed: allowed.into_iter().map(Into::into).collect(), + } + } + + /// Returns `true` if `name` is on the allowlist. + pub fn allows(&self, name: &str) -> bool { + self.allowed.contains(name) + } +} + +#[async_trait] +impl Middleware for ToolAllowlistMiddleware { + fn name(&self) -> &str { + self.label + } + + async fn before_tool( + &self, + _ctx: &mut RunContext, + _state: &State, + call: &mut ToolCall, + ) -> Result<()> { + if !self.allowed.contains(&call.name) { + return Err(TinyAgentsError::Validation(format!( + "tool `{}` is not on the allowlist", + call.name + ))); + } + Ok(()) + } +} + +// ── ToolPolicyMiddleware ────────────────────────────────────────────────────── + +impl ToolPolicyMiddleware { + /// Creates a policy middleware from a name→policy snapshot (typically + /// [`ToolRegistry::policies`][crate::harness::tool::ToolRegistry::policies]). + /// + /// Defaults are permissive: nothing is required or denied until configured. + /// Use [`strict`](Self::strict) for a fail-closed baseline. + pub fn new( + policies: std::collections::HashMap, + ) -> Self { + Self { + label: "tool_policy", + policies, + require_classification: false, + require_background_safe: false, + deny: crate::harness::tool::ToolSideEffects::default(), + require_sandbox: false, + require_approval: false, + approved: std::collections::HashSet::new(), + enforce_result_bytes: false, + } + } + + /// Creates a fail-closed policy middleware: unclassified tools are rejected, + /// and tools declaring `destructive` or `payment` side effects are denied. + pub fn strict( + policies: std::collections::HashMap, + ) -> Self { + Self { + label: "tool_policy", + policies, + require_classification: true, + require_background_safe: false, + deny: crate::harness::tool::ToolSideEffects { + destructive: true, + payment: true, + ..crate::harness::tool::ToolSideEffects::default() + }, + require_sandbox: false, + require_approval: false, + approved: std::collections::HashSet::new(), + enforce_result_bytes: false, + } + } + + /// Requires every tool to carry a classified policy (fail closed on + /// unclassified or unknown tools). + pub fn require_classification(mut self, require: bool) -> Self { + self.require_classification = require; + self + } + + /// Requires every exposed/executed tool to be `background_safe`. + pub fn require_background_safe(mut self, require: bool) -> Self { + self.require_background_safe = require; + self + } + + /// Denies tools declaring any side effect present in `mask`. + pub fn deny_side_effects(mut self, mask: crate::harness::tool::ToolSideEffects) -> Self { + self.deny = mask; + self + } + + /// Enforces that a tool declaring + /// [`SandboxMode::Required`][crate::harness::tool::SandboxMode::Required] + /// only runs when the run carries a sandboxed workspace (fail closed + /// otherwise). See [`RunContext::with_workspace`][crate::harness::context::RunContext::with_workspace]. + pub fn require_sandbox(mut self, require: bool) -> Self { + self.require_sandbox = require; + self + } + + /// Blocks any tool declaring `approval_required` unless its name is in + /// `approved`, turning the declarative approval flag into a fail-closed gate. + pub fn require_approval( + mut self, + approved: impl IntoIterator>, + ) -> Self { + self.require_approval = true; + self.approved = approved.into_iter().map(Into::into).collect(); + self + } + + /// Enforces each tool's declared `max_result_bytes` cap by truncating and + /// flagging oversized results in `after_tool`. + pub fn enforce_result_bytes(mut self, enforce: bool) -> Self { + self.enforce_result_bytes = enforce; + self + } + + /// Returns `Ok(())` if the named tool is permitted, otherwise an explanation + /// of why it is blocked. Used by both the exposure and execution hooks so a + /// hidden tool cannot be executed by a divergent decision. + fn evaluate(&self, name: &str) -> std::result::Result<(), String> { + let Some(policy) = self.policies.get(name) else { + if self.require_classification { + return Err(format!("tool `{name}` has no declared policy")); + } + return Ok(()); + }; + if self.require_classification && !policy.classified { + return Err(format!("tool `{name}` is unclassified")); + } + let s = &policy.side_effects; + let d = &self.deny; + let denied = (d.writes_files && s.writes_files) + || (d.network && s.network) + || (d.installs_dependencies && s.installs_dependencies) + || (d.destructive && s.destructive) + || (d.external_service && s.external_service) + || (d.payment && s.payment); + if denied { + return Err(format!("tool `{name}` declares a denied side effect")); + } + if self.require_background_safe && !policy.access.background_safe { + return Err(format!("tool `{name}` is not background-safe")); + } + if self.require_approval && policy.access.approval_required && !self.approved.contains(name) + { + return Err(format!( + "tool `{name}` requires approval that was not granted" + )); + } + Ok(()) + } + + /// The context-aware slice of policy enforcement: the sandbox requirement + /// depends on the run's workspace, which `evaluate` (name-only) cannot see. + fn evaluate_sandbox( + &self, + name: &str, + ctx: &RunContext, + ) -> std::result::Result<(), String> { + if !self.require_sandbox { + return Ok(()); + } + let Some(policy) = self.policies.get(name) else { + return Ok(()); + }; + if policy.runtime.sandbox != crate::harness::tool::SandboxMode::Required { + return Ok(()); + } + let sandboxed = ctx + .workspace + .as_ref() + .is_some_and(|ws| ws.sandbox == crate::harness::tool::SandboxMode::Required); + if sandboxed { + Ok(()) + } else { + Err(format!( + "tool `{name}` requires a sandbox but the run has none" + )) + } + } +} + +#[async_trait] +impl Middleware for ToolPolicyMiddleware { + fn name(&self) -> &str { + self.label + } + + async fn before_model( + &self, + ctx: &mut RunContext, + _state: &State, + request: &mut ModelRequest, + ) -> Result<()> { + request.tools.retain(|schema| { + self.evaluate(&schema.name).is_ok() && self.evaluate_sandbox(&schema.name, ctx).is_ok() + }); + Ok(()) + } + + async fn before_tool( + &self, + ctx: &mut RunContext, + _state: &State, + call: &mut ToolCall, + ) -> Result<()> { + self.evaluate(&call.name) + .and_then(|_| self.evaluate_sandbox(&call.name, ctx)) + .map_err(TinyAgentsError::Validation) + } + + async fn after_tool( + &self, + _ctx: &mut RunContext, + _state: &State, + result: &mut ToolResult, + ) -> Result<()> { + if !self.enforce_result_bytes { + return Ok(()); + } + if let Some(policy) = self.policies.get(&result.name) + && let Some(limit) = policy.runtime.max_result_bytes + && result.content.len() > limit + { + // Truncate on a char boundary at or below the byte limit so the + // enforced payload is still valid UTF-8. + let mut end = limit; + while end > 0 && !result.content.is_char_boundary(end) { + end -= 1; + } + result.content.truncate(end); + let note = format!("tool result exceeded max_result_bytes ({limit}); truncated"); + result.error = Some(match result.error.take() { + Some(existing) => format!("{existing}; {note}"), + None => note, + }); + } + Ok(()) + } +} + +// ── DynamicToolSelectionMiddleware ──────────────────────────────────────────── + +impl DynamicToolSelectionMiddleware { + /// Creates a selection middleware exposing only tools for which `predicate` + /// returns `true`. + pub fn new(predicate: ToolPredicate) -> Self { + Self { + label: "dynamic_tool_selection", + predicate, + } + } + + /// Creates a selection middleware exposing only the named tools. + pub fn allowing(names: impl IntoIterator>) -> Self { + let allowed: HashSet = names.into_iter().map(Into::into).collect(); + Self::new(Arc::new(move |schema: &ToolSchema| { + allowed.contains(&schema.name) + })) + } +} + +#[async_trait] +impl Middleware + for DynamicToolSelectionMiddleware +{ + fn name(&self) -> &str { + self.label + } + + async fn before_model( + &self, + _ctx: &mut RunContext, + _state: &State, + request: &mut ModelRequest, + ) -> Result<()> { + request.tools.retain(|schema| (self.predicate)(schema)); + Ok(()) + } +} + +// ── ContextualToolSelectionMiddleware ───────────────────────────────────────── + +impl ContextualToolSelectionMiddleware { + /// Creates a selection middleware from a context-aware predicate. + pub fn new(predicate: ContextualToolPredicate) -> Self { + Self { + label: "contextual_tool_selection", + predicate, + } + } + + /// Builds a selection middleware from explicit allow/deny lists. + /// + /// Composition rules (fail-closed): + /// - a tool named in `deny` is always hidden; + /// - when `allow` is `Some`, a tool must be named in it to be exposed + /// (unknown tools are hidden); + /// - when `allow` is `None`, everything not denied is exposed. + pub fn from_lists( + allow: Option>>, + deny: impl IntoIterator>, + ) -> Self { + let allow: Option> = + allow.map(|names| names.into_iter().map(Into::into).collect()); + let deny: HashSet = deny.into_iter().map(Into::into).collect(); + Self::from_resolved_lists(allow, deny) + } + + /// Builds a selection middleware whose effective policy is a child allow/deny + /// pair *composed with* an inherited parent policy, so a delegated sub-agent + /// can only ever narrow — never widen — the tools its parent allowed. + /// + /// Inheritance rules: + /// - **deny is additive**: the effective denylist is `parent_deny ∪ child_deny` + /// (a child cannot un-deny what the parent denied); + /// - **allow is intersective**: if both parent and child restrict to an + /// allowlist, the effective allowlist is their intersection; if only one + /// restricts, that allowlist applies; if neither does, all-not-denied is + /// exposed. + /// + /// The result is fail-closed for the same reasons as [`Self::from_lists`]. + pub fn inheriting( + parent_allow: Option>>, + parent_deny: impl IntoIterator>, + child_allow: Option>>, + child_deny: impl IntoIterator>, + ) -> Self { + let parent_allow: Option> = + parent_allow.map(|n| n.into_iter().map(Into::into).collect()); + let child_allow: Option> = + child_allow.map(|n| n.into_iter().map(Into::into).collect()); + let allow = match (parent_allow, child_allow) { + (Some(p), Some(c)) => Some(p.intersection(&c).cloned().collect()), + (Some(p), None) => Some(p), + (None, Some(c)) => Some(c), + (None, None) => None, + }; + let mut deny: HashSet = parent_deny.into_iter().map(Into::into).collect(); + deny.extend(child_deny.into_iter().map(Into::into)); + Self::from_resolved_lists(allow, deny) + } + + fn from_resolved_lists(allow: Option>, deny: HashSet) -> Self { + Self::new(Arc::new(move |schema: &ToolSchema, _ctx| { + if deny.contains(&schema.name) { + return false; + } + match &allow { + Some(set) => set.contains(&schema.name), + None => true, + } + })) + } +} + +#[async_trait] +impl Middleware + for ContextualToolSelectionMiddleware +{ + fn name(&self) -> &str { + self.label + } + + async fn before_model( + &self, + ctx: &mut RunContext, + _state: &State, + request: &mut ModelRequest, + ) -> Result<()> { + let selection = ToolSelectionContext { + run_id: ctx.config.run_id.as_str().to_string(), + depth: ctx.config.depth, + tags: ctx.config.tags.clone(), + requested_model: request.model.clone(), + }; + let mut excluded = Vec::new(); + request.tools.retain(|schema| { + let keep = (self.predicate)(schema, &selection); + if !keep { + excluded.push(schema.name.clone()); + } + keep + }); + // Make the exposure decision auditable when it actually withheld tools. + if !excluded.is_empty() { + ctx.emit(AgentEvent::ToolsFiltered { + by: self.label.to_string(), + excluded, + remaining: request.tools.len(), + }); + } + Ok(()) + } +} + +// ── HumanApprovalMiddleware ─────────────────────────────────────────────────── + +impl HumanApprovalMiddleware { + /// Creates an approval middleware that interrupts when any flagged tool is + /// called and no approval callback is configured. + pub fn new(flagged: impl IntoIterator>) -> Self { + Self { + label: "human_approval", + flagged: flagged.into_iter().map(Into::into).collect(), + approve: None, + } + } + + /// Attaches an approval callback consulted for flagged tools. Returning + /// `true` admits the call; `false` (or no callback) raises an interrupt. + pub fn with_approval(mut self, approve: ApprovalFn) -> Self { + self.approve = Some(approve); + self + } +} + +#[async_trait] +impl Middleware for HumanApprovalMiddleware { + fn name(&self) -> &str { + self.label + } + + async fn before_tool( + &self, + _ctx: &mut RunContext, + _state: &State, + call: &mut ToolCall, + ) -> Result<()> { + if self.flagged.contains(&call.name) { + let approved = self + .approve + .as_ref() + .map(|approve| approve(call)) + .unwrap_or(false); + if !approved { + return Err(TinyAgentsError::Interrupted { + node: "tool".to_string(), + message: format!("tool `{}` requires human approval", call.name), + }); + } + } + Ok(()) + } +} diff --git a/src/harness/middleware/library/types.rs b/src/harness/middleware/library/types.rs index 62df59b..8185022 100644 --- a/src/harness/middleware/library/types.rs +++ b/src/harness/middleware/library/types.rs @@ -19,6 +19,7 @@ //! `mod.rs`; tests live in `test.rs`. Every public item is re-exported through //! `crate::harness::middleware` so callers import from one place. +use std::collections::VecDeque; use std::marker::PhantomData; use std::sync::{Arc, Mutex}; use std::time::{Duration, Instant}; @@ -204,9 +205,14 @@ pub struct BudgetSpend { pub cost: crate::harness::cost::CostTotals, /// Whether a warning has already been emitted (warn-once). pub warned: bool, - /// Input tokens reserved by the most recent preflight, awaiting - /// reconciliation against the provider-reported usage in `after_model`. - pub last_reserved_input: u64, + /// Sum of input tokens preflight-reserved by calls that have not yet + /// reconciled in `after_model`. Shared trackers (handed to concurrent + /// sub-agent runs) can have more than one outstanding reservation at + /// once, so this is a running total, not a single call's estimate; + /// each in-flight call's own reservation is tracked separately by its + /// [`BudgetMiddleware`] instance and released from this total when it + /// reconciles (or is abandoned). + pub reserved_input_total: u64, } /// Around-nothing lifecycle middleware that enforces a token/money @@ -233,6 +239,11 @@ pub struct BudgetMiddleware { pub(crate) limits: BudgetLimits, pub(crate) tracker: BudgetTracker, pub(crate) pricing: std::collections::HashMap, + /// This run's own outstanding preflight reservation (input tokens), + /// awaiting reconciliation in `after_model`. Local to this middleware + /// instance (one per run) so concurrent runs sharing the same + /// [`BudgetTracker`] never clobber each other's reservation. + pub(crate) pending_reservation: std::sync::Mutex, } // ── ToolPolicyMiddleware ────────────────────────────────────────────────────── @@ -454,6 +465,11 @@ pub enum TraceBoundary { End, } +/// Default cap on the number of [`PhaseTrace`] entries a [`TracingMiddleware`] +/// retains before evicting the oldest. Long-running agent loops otherwise grow +/// this recorder without bound. +pub const DEFAULT_TRACE_RECORD_CAP: usize = 1024; + /// Lifecycle middleware that records structured begin/end traces and per-phase /// counts for an entire run. /// @@ -464,10 +480,17 @@ pub enum TraceBoundary { /// giving tests and dashboards a structured timeline without parsing the event /// stream. (The surrounding [`MiddlewareStack`][crate::harness::middleware::MiddlewareStack] /// also emits `MiddlewareStarted`/`MiddlewareCompleted` events around each hook.) +/// +/// The recorder is a bounded ring buffer (default cap +/// [`DEFAULT_TRACE_RECORD_CAP`], configurable via +/// [`TracingMiddleware::with_max_records`]): once full, the oldest trace is +/// dropped to make room for the newest, so an unbounded run cannot grow this +/// middleware's memory footprint forever. pub struct TracingMiddleware { pub(crate) label: &'static str, - pub(crate) records: Mutex>, + pub(crate) records: Mutex>, pub(crate) counts: Mutex, + pub(crate) max_records: usize, } /// Per-phase begin counts captured by [`TracingMiddleware`]. diff --git a/src/harness/middleware/mod.rs b/src/harness/middleware/mod.rs index 1c12eee..294c704 100644 --- a/src/harness/middleware/mod.rs +++ b/src/harness/middleware/mod.rs @@ -37,19 +37,40 @@ pub use library::*; use std::sync::Arc; use crate::error::{Result, TinyAgentsError}; -use crate::harness::cache::{CacheLayoutEvent, PromptCacheLayout}; use crate::harness::context::RunContext; use crate::harness::events::AgentEvent; -use crate::harness::message::Message; use crate::harness::model::{ModelDelta, ModelRequest, ModelResponse}; -use crate::harness::summarization::{ - ConcatSummarizer, SummarizationPolicy, Summarizer, SummaryRecord, TrimStrategy, - estimate_tokens, trim_messages, -}; use crate::harness::tool::{ToolCall, ToolDelta, ToolResult}; -use crate::harness::usage::UsageTotals; -use async_trait::async_trait; +/// Runs one per-middleware lifecycle hook across the whole stack, bracketing +/// each call with `MiddlewareStarted`/`MiddlewareCompleted` events and fanning +/// `on_error` out to every middleware on the first failure (so the originating +/// error is never masked). +/// +/// This is factored as a macro rather than an async helper because each hook +/// takes different arguments and borrows `ctx` mutably across its `await`, which +/// a closure-based helper cannot express without heap-boxing every call. +/// +/// Crucially, `MiddlewareCompleted` is emitted on *both* the success and error +/// paths: a hook that returns `Err` can no longer leave a dangling +/// `MiddlewareStarted` with no matching `Completed` in the event stream. `$iter` +/// selects registration order (`.iter()`) or reverse order (`.iter().rev()`); +/// `$call` is the (un-awaited) hook invocation on `$mw`. +macro_rules! run_stack_hook { + ($self:ident, $ctx:ident, $iter:expr, |$mw:ident| $call:expr) => {{ + for $mw in $iter { + let name = $mw.name().to_string(); + $ctx.emit(AgentEvent::MiddlewareStarted { name: name.clone() }); + let result = $call.await; + $ctx.emit(AgentEvent::MiddlewareCompleted { name }); + if let Err(e) = result { + $self.fan_out_on_error($ctx, &e).await; + return Err(e); + } + } + Ok(()) + }}; +} // ── AgentRun ──────────────────────────────────────────────────────────────── @@ -135,23 +156,8 @@ impl MiddlewareStack { /// Runs every middleware's [`Middleware::before_agent`] in registration /// order. pub async fn run_before_agent(&self, ctx: &mut RunContext, state: &State) -> Result<()> { - for mw in self.middlewares.iter() { - ctx.emit(AgentEvent::MiddlewareStarted { - name: mw.name().to_string(), - }); - match mw.before_agent(ctx, state).await { - Ok(()) => { - ctx.emit(AgentEvent::MiddlewareCompleted { - name: mw.name().to_string(), - }); - } - Err(e) => { - self.fan_out_on_error(ctx, &e).await; - return Err(e); - } - } - } - Ok(()) + run_stack_hook!(self, ctx, self.middlewares.iter(), |mw| mw + .before_agent(ctx, state)) } /// Runs every middleware's [`Middleware::after_agent`] in reverse @@ -162,23 +168,8 @@ impl MiddlewareStack { state: &State, run: &mut AgentRun, ) -> Result<()> { - for mw in self.middlewares.iter().rev() { - ctx.emit(AgentEvent::MiddlewareStarted { - name: mw.name().to_string(), - }); - match mw.after_agent(ctx, state, run).await { - Ok(()) => { - ctx.emit(AgentEvent::MiddlewareCompleted { - name: mw.name().to_string(), - }); - } - Err(e) => { - self.fan_out_on_error(ctx, &e).await; - return Err(e); - } - } - } - Ok(()) + run_stack_hook!(self, ctx, self.middlewares.iter().rev(), |mw| mw + .after_agent(ctx, state, run)) } /// Runs every middleware's [`Middleware::before_model`] in registration @@ -189,27 +180,21 @@ impl MiddlewareStack { state: &State, request: &mut ModelRequest, ) -> Result<()> { - for mw in self.middlewares.iter() { - ctx.emit(AgentEvent::MiddlewareStarted { - name: mw.name().to_string(), - }); - match mw.before_model(ctx, state, request).await { - Ok(()) => { - ctx.emit(AgentEvent::MiddlewareCompleted { - name: mw.name().to_string(), - }); - } - Err(e) => { - self.fan_out_on_error(ctx, &e).await; - return Err(e); - } - } - } - Ok(()) + run_stack_hook!(self, ctx, self.middlewares.iter(), |mw| mw + .before_model(ctx, state, request)) } /// Runs every middleware's [`Middleware::on_model_delta`] in registration /// order for one streamed delta. + /// + /// Unlike the other stack runners, the per-delta hook is deliberately *not* + /// bracketed by `MiddlewareStarted`/`MiddlewareCompleted` events. This runs + /// on the streaming hot path — potentially hundreds of times per second per + /// middleware — and emitting two events (each cloning `mw.name()` and + /// acquiring the recorder mutex) per middleware per token dominated the + /// stream loop's cost for zero observability value. Callers that need to + /// observe delta-level middleware activity should instrument the hook + /// itself. pub async fn run_on_model_delta( &self, ctx: &mut RunContext, @@ -217,19 +202,9 @@ impl MiddlewareStack { delta: &mut ModelDelta, ) -> Result<()> { for mw in self.middlewares.iter() { - ctx.emit(AgentEvent::MiddlewareStarted { - name: mw.name().to_string(), - }); - match mw.on_model_delta(ctx, state, delta).await { - Ok(()) => { - ctx.emit(AgentEvent::MiddlewareCompleted { - name: mw.name().to_string(), - }); - } - Err(e) => { - self.fan_out_on_error(ctx, &e).await; - return Err(e); - } + if let Err(e) = mw.on_model_delta(ctx, state, delta).await { + self.fan_out_on_error(ctx, &e).await; + return Err(e); } } Ok(()) @@ -243,23 +218,8 @@ impl MiddlewareStack { state: &State, response: &mut ModelResponse, ) -> Result<()> { - for mw in self.middlewares.iter().rev() { - ctx.emit(AgentEvent::MiddlewareStarted { - name: mw.name().to_string(), - }); - match mw.after_model(ctx, state, response).await { - Ok(()) => { - ctx.emit(AgentEvent::MiddlewareCompleted { - name: mw.name().to_string(), - }); - } - Err(e) => { - self.fan_out_on_error(ctx, &e).await; - return Err(e); - } - } - } - Ok(()) + run_stack_hook!(self, ctx, self.middlewares.iter().rev(), |mw| mw + .after_model(ctx, state, response)) } /// Runs every middleware's [`Middleware::before_tool`] in registration @@ -270,23 +230,8 @@ impl MiddlewareStack { state: &State, call: &mut ToolCall, ) -> Result<()> { - for mw in self.middlewares.iter() { - ctx.emit(AgentEvent::MiddlewareStarted { - name: mw.name().to_string(), - }); - match mw.before_tool(ctx, state, call).await { - Ok(()) => { - ctx.emit(AgentEvent::MiddlewareCompleted { - name: mw.name().to_string(), - }); - } - Err(e) => { - self.fan_out_on_error(ctx, &e).await; - return Err(e); - } - } - } - Ok(()) + run_stack_hook!(self, ctx, self.middlewares.iter(), |mw| mw + .before_tool(ctx, state, call)) } /// Runs every middleware's [`Middleware::on_tool_delta`] in registration @@ -297,23 +242,8 @@ impl MiddlewareStack { state: &State, delta: &mut ToolDelta, ) -> Result<()> { - for mw in self.middlewares.iter() { - ctx.emit(AgentEvent::MiddlewareStarted { - name: mw.name().to_string(), - }); - match mw.on_tool_delta(ctx, state, delta).await { - Ok(()) => { - ctx.emit(AgentEvent::MiddlewareCompleted { - name: mw.name().to_string(), - }); - } - Err(e) => { - self.fan_out_on_error(ctx, &e).await; - return Err(e); - } - } - } - Ok(()) + run_stack_hook!(self, ctx, self.middlewares.iter(), |mw| mw + .on_tool_delta(ctx, state, delta)) } /// Runs every middleware's [`Middleware::after_tool`] in reverse @@ -324,23 +254,8 @@ impl MiddlewareStack { state: &State, result: &mut ToolResult, ) -> Result<()> { - for mw in self.middlewares.iter().rev() { - ctx.emit(AgentEvent::MiddlewareStarted { - name: mw.name().to_string(), - }); - match mw.after_tool(ctx, state, result).await { - Ok(()) => { - ctx.emit(AgentEvent::MiddlewareCompleted { - name: mw.name().to_string(), - }); - } - Err(e) => { - self.fan_out_on_error(ctx, &e).await; - return Err(e); - } - } - } - Ok(()) + run_stack_hook!(self, ctx, self.middlewares.iter().rev(), |mw| mw + .after_tool(ctx, state, result)) } /// Runs every middleware's [`Middleware::on_error`] in registration order, @@ -426,14 +341,14 @@ impl ModelHandler<'_, State, Ctx> { remaining: tail, base: self.base, }; - ctx.emit(AgentEvent::MiddlewareStarted { - name: head.name().to_string(), - }); - let outcome = head.wrap_model(ctx, state, request, next).await?; - ctx.emit(AgentEvent::MiddlewareCompleted { - name: head.name().to_string(), - }); - Ok(outcome) + let name = head.name().to_string(); + ctx.emit(AgentEvent::MiddlewareStarted { name: name.clone() }); + // Emit `Completed` whether the wrap layer succeeds or errors, so + // a failing layer never leaves a dangling `Started` in the event + // stream (the onion's balance invariant). + let outcome = head.wrap_model(ctx, state, request, next).await; + ctx.emit(AgentEvent::MiddlewareCompleted { name }); + outcome } None => Ok(MiddlewareModelOutcome::Response( self.base.call(ctx, state, request).await?, @@ -457,14 +372,13 @@ impl ToolHandler<'_, State, Ctx> { remaining: tail, base: self.base, }; - ctx.emit(AgentEvent::MiddlewareStarted { - name: head.name().to_string(), - }); - let outcome = head.wrap_tool(ctx, state, call, next).await?; - ctx.emit(AgentEvent::MiddlewareCompleted { - name: head.name().to_string(), - }); - Ok(outcome) + let name = head.name().to_string(); + ctx.emit(AgentEvent::MiddlewareStarted { name: name.clone() }); + // Balance `Started` with `Completed` even when the wrap layer + // errors (see `ModelHandler::run`). + let outcome = head.wrap_tool(ctx, state, call, next).await; + ctx.emit(AgentEvent::MiddlewareCompleted { name }); + outcome } None => Ok(MiddlewareToolOutcome::Result( self.base.call(ctx, state, call).await?, @@ -473,444 +387,5 @@ impl ToolHandler<'_, State, Ctx> { } } -// ── LoggingMiddleware ───────────────────────────────────────────────────────── - -impl LoggingMiddleware { - /// Creates a logging middleware with the default label `"logging"`. - pub fn new() -> Self { - Self::with_label("logging") - } - - /// Creates a logging middleware with a custom static label. - pub fn with_label(label: &'static str) -> Self { - Self { - label, - counts: std::sync::Mutex::new(HookCounts::default()), - } - } - - /// Returns a snapshot of the per-hook invocation counts recorded so far. - pub fn counts(&self) -> HookCounts { - self.counts.lock().expect("counts mutex poisoned").clone() - } -} - -impl Default for LoggingMiddleware { - fn default() -> Self { - Self::new() - } -} - -#[async_trait] -impl Middleware for LoggingMiddleware { - fn name(&self) -> &str { - self.label - } - - async fn before_agent(&self, _ctx: &mut RunContext, _state: &State) -> Result<()> { - self.counts - .lock() - .expect("counts mutex poisoned") - .before_agent += 1; - Ok(()) - } - - async fn after_agent( - &self, - _ctx: &mut RunContext, - _state: &State, - _run: &mut AgentRun, - ) -> Result<()> { - self.counts - .lock() - .expect("counts mutex poisoned") - .after_agent += 1; - Ok(()) - } - - async fn before_model( - &self, - _ctx: &mut RunContext, - _state: &State, - _request: &mut ModelRequest, - ) -> Result<()> { - self.counts - .lock() - .expect("counts mutex poisoned") - .before_model += 1; - Ok(()) - } - - async fn on_model_delta( - &self, - _ctx: &mut RunContext, - _state: &State, - _delta: &mut ModelDelta, - ) -> Result<()> { - self.counts - .lock() - .expect("counts mutex poisoned") - .on_model_delta += 1; - Ok(()) - } - - async fn after_model( - &self, - _ctx: &mut RunContext, - _state: &State, - _response: &mut ModelResponse, - ) -> Result<()> { - self.counts - .lock() - .expect("counts mutex poisoned") - .after_model += 1; - Ok(()) - } - - async fn before_tool( - &self, - _ctx: &mut RunContext, - _state: &State, - _call: &mut ToolCall, - ) -> Result<()> { - self.counts - .lock() - .expect("counts mutex poisoned") - .before_tool += 1; - Ok(()) - } - - async fn on_tool_delta( - &self, - _ctx: &mut RunContext, - _state: &State, - _delta: &mut ToolDelta, - ) -> Result<()> { - self.counts - .lock() - .expect("counts mutex poisoned") - .on_tool_delta += 1; - Ok(()) - } - - async fn after_tool( - &self, - _ctx: &mut RunContext, - _state: &State, - _result: &mut ToolResult, - ) -> Result<()> { - self.counts - .lock() - .expect("counts mutex poisoned") - .after_tool += 1; - Ok(()) - } - - async fn on_error(&self, _ctx: &mut RunContext, _error: &TinyAgentsError) -> Result<()> { - self.counts.lock().expect("counts mutex poisoned").on_error += 1; - Ok(()) - } -} - -// ── MessageTrimMiddleware ───────────────────────────────────────────────────── - -impl MessageTrimMiddleware { - /// Creates a trim middleware using the given [`TrimStrategy`]. - pub fn new(strategy: TrimStrategy) -> Self { - Self { strategy } - } -} - -#[async_trait] -impl Middleware for MessageTrimMiddleware { - fn name(&self) -> &str { - "message_trim" - } - - async fn before_model( - &self, - _ctx: &mut RunContext, - _state: &State, - request: &mut ModelRequest, - ) -> Result<()> { - request.messages = trim_messages(&request.messages, &self.strategy); - Ok(()) - } -} - -// ── ContextCompressionMiddleware ────────────────────────────────────────────── - -/// Estimate the total tokens of a message slice using the same per-message -/// heuristic the [`SummarizationPolicy`] uses internally. -fn total_message_tokens(messages: &[crate::harness::message::Message]) -> u64 { - messages.iter().map(|m| estimate_tokens(&m.text())).sum() -} - -impl ContextCompressionMiddleware { - /// Creates a compression middleware backed by the default - /// [`ConcatSummarizer`]. - pub fn new(policy: SummarizationPolicy) -> Self { - Self::with_summarizer(policy, Box::new(ConcatSummarizer)) - } - - /// Creates a compression middleware with a custom [`Summarizer`]. - pub fn with_summarizer(policy: SummarizationPolicy, summarizer: Box) -> Self { - Self { - label: "context_compression", - policy, - summarizer, - records: std::sync::Mutex::new(Vec::new()), - } - } - - /// Returns the configured [`SummarizationPolicy`]. - pub fn policy(&self) -> &SummarizationPolicy { - &self.policy - } - - /// Returns the [`SummaryRecord`]s produced so far, in order. Each record - /// carries the compression provenance for one compaction. - pub fn records(&self) -> Vec { - self.records.lock().expect("records mutex poisoned").clone() - } -} - -#[async_trait] -impl Middleware for ContextCompressionMiddleware { - fn name(&self) -> &str { - self.label - } - - async fn before_model( - &self, - ctx: &mut RunContext, - _state: &State, - request: &mut ModelRequest, - ) -> Result<()> { - // Below the window threshold: pass through untouched (no-op, no event). - if !self.policy.should_summarize(&request.messages) { - return Ok(()); - } - - let (to_summarize, to_keep) = self.policy.plan(&request.messages); - // Nothing old enough to compress (e.g. keep_last covers everything): - // leave the transcript untouched rather than summarizing an empty set. - if to_summarize.is_empty() { - return Ok(()); - } - - let from_tokens = total_message_tokens(&request.messages); - let record = self.summarizer.summarize(&to_summarize).await?; - - let mut new_messages = Vec::with_capacity(to_keep.len() + 1); - new_messages.push(record.summary.clone()); - new_messages.extend(to_keep); - let to_tokens = total_message_tokens(&new_messages); - - self.records - .lock() - .expect("records mutex poisoned") - .push(record); - request.messages = new_messages; - - ctx.emit(AgentEvent::Compressed { - from_tokens, - to_tokens, - }); - Ok(()) - } -} - -// ── MicrocompactMiddleware ──────────────────────────────────────────────────── - -impl MicrocompactMiddleware { - /// Creates a micro-compaction middleware that keeps the newest `keep_recent` - /// tool-result bodies verbatim and blanks older ones with `placeholder`. - /// Event emission is off by default; enable it with - /// [`MicrocompactMiddleware::with_events`]. - pub fn new(keep_recent: usize, placeholder: impl Into) -> Self { - Self { - label: "microcompact", - keep_recent, - placeholder: placeholder.into(), - emit_events: false, - } - } - - /// Enable or disable emitting an - /// [`AgentEvent::Compressed`][crate::harness::events::AgentEvent::Compressed] - /// event whenever at least one tool body is cleared. Off by default so the - /// middleware can be a silent transcript rewrite. - pub fn with_events(mut self, emit_events: bool) -> Self { - self.emit_events = emit_events; - self - } - - /// The number of most-recent tool-result bodies kept verbatim. - pub fn keep_recent(&self) -> usize { - self.keep_recent - } - - /// The placeholder text swapped in for cleared tool-result bodies. - pub fn placeholder(&self) -> &str { - &self.placeholder - } -} - -#[async_trait] -impl Middleware for MicrocompactMiddleware { - fn name(&self) -> &str { - self.label - } - - async fn before_model( - &self, - ctx: &mut RunContext, - _state: &State, - request: &mut ModelRequest, - ) -> Result<()> { - // Indices of every tool-result message, oldest → newest. - let tool_idxs: Vec = request - .messages - .iter() - .enumerate() - .filter(|(_, m)| matches!(m, Message::Tool(_))) - .map(|(i, _)| i) - .collect(); - if tool_idxs.len() <= self.keep_recent { - return Ok(()); - } - - let from_tokens = if self.emit_events { - total_message_tokens(&request.messages) - } else { - 0 - }; - - let cut = tool_idxs.len() - self.keep_recent; - let mut cleared = 0usize; - for &i in &tool_idxs[..cut] { - // Skip messages already reduced to the placeholder; otherwise swap the - // body for it (idempotent, preserves the tool_call_id). - if request.messages[i].text() == self.placeholder { - continue; - } - if let Message::Tool(t) = &request.messages[i] { - let id = t.tool_call_id.clone(); - request.messages[i] = Message::tool(id, self.placeholder.clone()); - cleared += 1; - } - } - - if self.emit_events && cleared > 0 { - let to_tokens = total_message_tokens(&request.messages); - ctx.emit(AgentEvent::Compressed { - from_tokens, - to_tokens, - }); - } - Ok(()) - } -} - -// ── PromptCacheGuardMiddleware ──────────────────────────────────────────────── - -impl PromptCacheGuardMiddleware { - /// Creates a cache-guard middleware with the default label - /// `"prompt_cache_guard"`. - pub fn new() -> Self { - Self { - label: "prompt_cache_guard", - previous: std::sync::Mutex::new(None), - events: std::sync::Mutex::new(Vec::new()), - } - } - - /// Returns the cache-layout change events recorded so far, in order. - pub fn layout_events(&self) -> Vec { - self.events.lock().expect("events mutex poisoned").clone() - } -} - -impl Default for PromptCacheGuardMiddleware { - fn default() -> Self { - Self::new() - } -} - -#[async_trait] -impl Middleware for PromptCacheGuardMiddleware { - fn name(&self) -> &str { - self.label - } - - async fn before_model( - &self, - _ctx: &mut RunContext, - _state: &State, - request: &mut ModelRequest, - ) -> Result<()> { - let layout = PromptCacheLayout::from_request(request); - let mut previous = self.previous.lock().expect("previous mutex poisoned"); - if let Some(prev) = previous.as_ref() - && !prev.is_prefix_stable_against(&layout) - { - let event = CacheLayoutEvent::new(prev, &layout); - self.events - .lock() - .expect("events mutex poisoned") - .push(event); - } - *previous = Some(layout); - Ok(()) - } -} - -// ── UsageAccountingMiddleware ───────────────────────────────────────────────── - -impl UsageAccountingMiddleware { - /// Creates a usage-accounting middleware with the default label - /// `"usage_accounting"`. - pub fn new() -> Self { - Self { - label: "usage_accounting", - totals: std::sync::Mutex::new(UsageTotals::new()), - } - } - - /// Returns a snapshot of the accumulated usage totals. - pub fn totals(&self) -> UsageTotals { - *self.totals.lock().expect("totals mutex poisoned") - } -} - -impl Default for UsageAccountingMiddleware { - fn default() -> Self { - Self::new() - } -} - -#[async_trait] -impl Middleware for UsageAccountingMiddleware { - fn name(&self) -> &str { - self.label - } - - async fn after_model( - &self, - _ctx: &mut RunContext, - _state: &State, - response: &mut ModelResponse, - ) -> Result<()> { - if let Some(usage) = response.usage { - self.totals - .lock() - .expect("totals mutex poisoned") - .record(usage); - } - Ok(()) - } -} - #[cfg(test)] mod test; diff --git a/src/harness/middleware/test.rs b/src/harness/middleware/test.rs index e58fbe8..65f61e9 100644 --- a/src/harness/middleware/test.rs +++ b/src/harness/middleware/test.rs @@ -194,6 +194,87 @@ async fn emits_started_and_completed_events() { ); } +#[tokio::test] +async fn failing_hook_still_emits_balanced_completed_event() { + // A hook that returns `Err` must still close its `MiddlewareStarted` with a + // matching `MiddlewareCompleted`, so downstream observers never see a + // dangling, unbalanced `Started`. + let mut stack: MiddlewareStack<()> = MiddlewareStack::new(); + stack.push(Arc::new(FailingMiddleware)); + + let recorder = Arc::new(RecordingListener::new()); + let mut c = ctx(); + c.events.subscribe(recorder.clone()); + + let mut request = ModelRequest::default(); + let result = stack.run_before_model(&mut c, &(), &mut request).await; + assert!(matches!(result, Err(TinyAgentsError::Middleware(_)))); + + let brackets: Vec = recorder + .events() + .into_iter() + .map(|r| r.event) + .filter(|e| { + matches!( + e, + AgentEvent::MiddlewareStarted { .. } | AgentEvent::MiddlewareCompleted { .. } + ) + }) + .collect(); + assert_eq!( + brackets, + vec![ + AgentEvent::MiddlewareStarted { + name: "failing".to_string() + }, + AgentEvent::MiddlewareCompleted { + name: "failing".to_string() + }, + ], + "a failing hook must emit a balanced Started/Completed pair" + ); +} + +#[tokio::test] +async fn on_model_delta_hook_emits_no_bracketing_events() { + // The per-delta hook runs on the streaming hot path, so it must NOT emit + // `MiddlewareStarted`/`MiddlewareCompleted` events the way the other stack + // runners do — those two events per middleware per token dominated the + // stream loop for no observability value. + let mut stack: MiddlewareStack<()> = MiddlewareStack::new(); + stack.push(Arc::new(LoggingMiddleware::new())); + + let recorder = Arc::new(RecordingListener::new()); + let mut c = ctx(); + c.events.subscribe(recorder.clone()); + + let mut delta = ModelDelta { + call_id: "call-1".to_string(), + content: "tok".to_string(), + reasoning: String::new(), + tool_call: None, + }; + stack + .run_on_model_delta(&mut c, &(), &mut delta) + .await + .unwrap(); + + let bracketing = recorder + .events() + .into_iter() + .filter(|r| { + matches!( + r.event, + AgentEvent::MiddlewareStarted { .. } | AgentEvent::MiddlewareCompleted { .. } + ) + }) + .count(); + assert_eq!( + bracketing, 0, + "the delta hook must not bracket middleware with events" + ); +} + #[tokio::test] async fn message_trim_middleware_shrinks_request() { let mw = MessageTrimMiddleware::new(TrimStrategy::KeepLast(1)); @@ -316,6 +397,92 @@ async fn context_compression_compresses_at_or_above_threshold() { assert!(compressed[0].1 > 0); } +#[tokio::test] +async fn context_compression_records_are_bounded_by_max_records() { + // A long-running loop that compresses repeatedly must not grow the + // recorder without bound; cap it and confirm eviction happens. + let policy = SummarizationPolicy { + keep_last: 1, + ..SummarizationPolicy::default() + } + .with_context_window(100) + .with_threshold_fraction(0.5); + let mw = Arc::new(ContextCompressionMiddleware::new(policy).with_max_records(2)); + let mut stack: MiddlewareStack<()> = MiddlewareStack::new(); + stack.push(mw.clone()); + + let big = "a".repeat(200); + let mut c = ctx(); + for _ in 0..10 { + let mut request = ModelRequest { + messages: vec![ + user(&format!("{big}-1")), + user(&format!("{big}-2")), + user(&format!("{big}-3")), + ], + ..Default::default() + }; + stack + .run_before_model(&mut c, &(), &mut request) + .await + .unwrap(); + } + + assert_eq!(mw.records().len(), 2); +} + +#[tokio::test] +async fn context_compression_keeps_system_prompt_before_summary() { + // A leading system prompt carries persistent instructions and anchors the + // cacheable prefix; the summary of elided older turns must be inserted + // *after* it, never at position 0. + let policy = SummarizationPolicy { + keep_last: 1, + ..SummarizationPolicy::default() + } + .with_context_window(100) + .with_threshold_fraction(0.5); + let mw = Arc::new(ContextCompressionMiddleware::new(policy)); + let mut stack: MiddlewareStack<()> = MiddlewareStack::new(); + stack.push(mw.clone()); + + let mut c = ctx(); + let big = "a".repeat(200); + let system_prompt = "You are a helpful assistant. Always follow these rules."; + let mut request = ModelRequest { + messages: vec![ + Message::system(system_prompt), + user(&format!("{big}-1")), + user(&format!("{big}-2")), + user(&format!("{big}-3")), + ], + ..Default::default() + }; + stack + .run_before_model(&mut c, &(), &mut request) + .await + .unwrap(); + + // Result is [real system prompt, summary, kept recent turn]. + assert_eq!(request.messages.len(), 3); + assert!(matches!(request.messages[0], Message::System(_))); + assert_eq!( + request.messages[0].text(), + system_prompt, + "the real system prompt must stay at position 0" + ); + assert!( + matches!(request.messages[1], Message::System(_)), + "the summary follows the system prompt" + ); + assert_ne!( + request.messages[1].text(), + system_prompt, + "position 1 is the summary, not a duplicated system prompt" + ); + assert_eq!(request.messages[2].text(), format!("{big}-3")); +} + #[tokio::test] async fn context_compression_none_window_falls_back_to_trigger_tokens() { // No context window → raw trigger_tokens gate (strict `>`). Trigger at 2 @@ -554,6 +721,26 @@ async fn prompt_cache_guard_detects_prefix_change() { assert_eq!(events[0].segment_ids_after, vec!["sys2".to_string()]); } +#[tokio::test] +async fn prompt_cache_guard_events_are_bounded_by_max_events() { + let mw = Arc::new(PromptCacheGuardMiddleware::new().with_max_events(2)); + let mut stack: MiddlewareStack<()> = MiddlewareStack::new(); + stack.push(mw.clone()); + + let mut c = ctx(); + // Each call after the first changes the stable prefix, producing an + // event; run far more iterations than the cap and confirm eviction. + for i in 0..10 { + let mut req = ModelRequest { + cache_segments: vec![segment(&format!("sys{i}"), SegmentRole::System, true)], + ..Default::default() + }; + stack.run_before_model(&mut c, &(), &mut req).await.unwrap(); + } + + assert_eq!(mw.layout_events().len(), 2); +} + // ── wrap middleware: model ──────────────────────────────────────────────────── /// Builds a `ModelResponse` whose single text block is `text`. diff --git a/src/harness/middleware/types.rs b/src/harness/middleware/types.rs index f73778c..2b1ddb4 100644 --- a/src/harness/middleware/types.rs +++ b/src/harness/middleware/types.rs @@ -15,6 +15,7 @@ //! All public items are re-exported through [`super`] so callers import from //! `crate::harness::middleware` directly. +use std::collections::VecDeque; use std::future::Future; use std::pin::Pin; use std::sync::{Arc, Mutex}; @@ -484,11 +485,16 @@ pub struct MessageTrimMiddleware { /// [`ConcatSummarizer`][crate::harness::summarization::ConcatSummarizer] is used /// by default; supply any [`Summarizer`] via /// [`ContextCompressionMiddleware::with_summarizer`]. +/// Default cap on the number of [`SummaryRecord`]s a +/// [`ContextCompressionMiddleware`] retains before evicting the oldest. +pub const DEFAULT_COMPRESSION_RECORD_CAP: usize = 1024; + pub struct ContextCompressionMiddleware { pub(crate) label: &'static str, pub(crate) policy: SummarizationPolicy, pub(crate) summarizer: Box, - pub(crate) records: Mutex>, + pub(crate) records: Mutex>, + pub(crate) max_records: usize, } // ── MicrocompactMiddleware ──────────────────────────────────────────────────── @@ -540,10 +546,15 @@ pub struct MicrocompactMiddleware { /// [`CacheLayoutEvent`] (retrievable via /// [`PromptCacheGuardMiddleware::layout_events`]) so KV-cache regressions are /// observable. This demonstrates provider prompt/KV-cache prefix protection. +/// Default cap on the number of [`CacheLayoutEvent`]s a +/// [`PromptCacheGuardMiddleware`] retains before evicting the oldest. +pub const DEFAULT_CACHE_GUARD_EVENT_CAP: usize = 1024; + pub struct PromptCacheGuardMiddleware { pub(crate) label: &'static str, pub(crate) previous: Mutex>, - pub(crate) events: Mutex>, + pub(crate) events: Mutex>, + pub(crate) max_events: usize, } // ── UsageAccountingMiddleware ───────────────────────────────────────────────── diff --git a/src/harness/model/mod.rs b/src/harness/model/mod.rs index f82cca6..fac2a58 100644 --- a/src/harness/model/mod.rs +++ b/src/harness/model/mod.rs @@ -27,6 +27,32 @@ use crate::harness::usage::Usage; pub use types::*; +impl std::fmt::Display for ProviderError { + /// Renders the same human-readable shape real provider adapters used to + /// build by hand before flattening it into a plain + /// [`crate::error::TinyAgentsError::Model`] string. Preserving this as a + /// `Display` impl lets [`crate::error::TinyAgentsError::Provider`] keep + /// the identical wording while also keeping the structured fields + /// (`status`, `code`, `retryable`) intact for callers — like + /// [`crate::harness::retry::is_retryable`] — that need to reason about + /// the failure rather than just print it. + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!( + f, + "{} returned{}{}: {}", + self.provider, + self.status + .map(|status| format!(" HTTP {status}")) + .unwrap_or_default(), + self.code + .as_deref() + .map(|code| format!(" ({code})")) + .unwrap_or_default(), + self.message + ) + } +} + impl ResponseFormat { /// Constructs a [`ResponseFormat::JsonSchema`] format. pub fn json_schema(name: impl Into, schema: Value) -> Self { @@ -673,10 +699,13 @@ impl StreamAccumulator { } if let Some(mut response) = self.completed { - if response.usage.is_none() { - response.usage = self.usage; - response.message.usage = self.usage; - } + // Reconcile the response and message usage with any streamed + // `UsageDelta`, preferring an already-present value and never + // overwriting a known usage with `None` (which previously clobbered a + // message-level usage the completed response carried). + let merged = response.usage.or(response.message.usage).or(self.usage); + response.usage = merged; + response.message.usage = merged; return Ok(response); } diff --git a/src/harness/model/test.rs b/src/harness/model/test.rs index 8116938..7169882 100644 --- a/src/harness/model/test.rs +++ b/src/harness/model/test.rs @@ -610,3 +610,81 @@ fn stream_accumulator_collects_reasoning_side_channel() { let response = acc.finish().unwrap(); assert_eq!(response.text(), "visible answer"); } + +#[test] +fn finish_preserves_message_usage_from_completed_response() { + // A completed response that carries usage only on the message (not the + // top-level field) and no streamed UsageDelta. finish must not clobber the + // message usage with None; it should promote it to the response too. + let mut response = ModelResponse::assistant("hi"); + response.usage = None; + response.message.usage = Some(Usage::new(10, 20)); + + let mut acc = StreamAccumulator::new(); + acc.push(&ModelStreamItem::Started); + acc.push(&ModelStreamItem::Completed(response)); + + let finished = acc.finish().unwrap(); + assert_eq!(finished.message.usage, Some(Usage::new(10, 20))); + assert_eq!(finished.usage, Some(Usage::new(10, 20))); +} + +#[test] +fn finish_backfills_usage_from_stream_delta_when_completed_lacks_it() { + // No usage anywhere on the completed response, but a UsageDelta arrived. Both + // the response and its message pick up the streamed usage. + let mut response = ModelResponse::assistant("hi"); + response.usage = None; + response.message.usage = None; + + let mut acc = StreamAccumulator::new(); + acc.push(&ModelStreamItem::Started); + acc.push(&ModelStreamItem::UsageDelta(Usage::new(4, 6))); + acc.push(&ModelStreamItem::Completed(response)); + + let finished = acc.finish().unwrap(); + assert_eq!(finished.usage, Some(Usage::new(4, 6))); + assert_eq!(finished.message.usage, Some(Usage::new(4, 6))); +} + +/// Round-trips a [`ModelStreamItem`] through JSON and asserts the re-serialized +/// form is byte-for-byte stable, proving every variant survives serde. +fn roundtrip_stream_item(item: ModelStreamItem) { + let value = serde_json::to_value(&item).expect("serialize ModelStreamItem"); + let back: ModelStreamItem = + serde_json::from_value(value.clone()).expect("deserialize ModelStreamItem"); + let reserialized = serde_json::to_value(&back).expect("re-serialize ModelStreamItem"); + assert_eq!(value, reserialized, "round-trip differs for {value}"); +} + +#[test] +fn model_stream_item_roundtrips_every_variant() { + roundtrip_stream_item(ModelStreamItem::Started); + roundtrip_stream_item(ModelStreamItem::MessageDelta( + crate::harness::message::MessageDelta::text("hi"), + )); + roundtrip_stream_item(ModelStreamItem::ToolCallDelta( + crate::harness::tool::ToolDelta { + call_id: "call-1".into(), + content: "{\"q\":1}".into(), + }, + )); + roundtrip_stream_item(ModelStreamItem::UsageDelta(Usage::new(3, 5))); + roundtrip_stream_item(ModelStreamItem::Completed(ModelResponse::assistant("done"))); + // The scalar-carrying variant an internally tagged enum could not encode. + roundtrip_stream_item(ModelStreamItem::Failed("boom".to_string())); + roundtrip_stream_item(ModelStreamItem::ProviderFailed(ProviderError { + provider: "openai".into(), + message: "nope".into(), + ..ProviderError::default() + })); +} + +#[test] +fn model_stream_item_failed_serializes_without_panicking() { + // Under internal tagging this call errored; adjacent tagging encodes the + // string payload under `content`. + let value = serde_json::to_value(ModelStreamItem::Failed("boom".into())).unwrap(); + assert_eq!(value["type"], json!("failed")); + assert_eq!(value["content"], json!("boom")); +} diff --git a/src/harness/model/types.rs b/src/harness/model/types.rs index eef2387..2994da4 100644 --- a/src/harness/model/types.rs +++ b/src/harness/model/types.rs @@ -432,6 +432,9 @@ pub struct ModelDelta { /// Incremental text content. #[serde(default)] pub content: String, + /// Incremental reasoning/thinking content, kept separate from visible text. + #[serde(default, skip_serializing_if = "String::is_empty")] + pub reasoning: String, /// Incremental tool-call fragment, when present. #[serde(default, skip_serializing_if = "Option::is_none")] pub tool_call: Option, @@ -480,8 +483,15 @@ pub struct ProviderError { /// /// Use [`StreamAccumulator`] (or the [`collect_model_stream`] helper) to fold a /// stream of these items back into a [`ModelResponse`]. +/// # Serialization +/// +/// This enum is **adjacently tagged** (`{"type": …, "content": …}`) rather than +/// internally tagged. Several variants wrap a non-struct payload +/// ([`ModelStreamItem::Failed`] wraps a `String`); an internally tagged enum +/// cannot serialize those (serde errors, or silently corrupts scalar JSON into +/// `{}`), so adjacent tagging is required for every variant to round-trip. #[derive(Clone, Debug, Serialize, Deserialize)] -#[serde(rename_all = "snake_case", tag = "type")] +#[serde(rename_all = "snake_case", tag = "type", content = "content")] pub enum ModelStreamItem { /// The stream has opened; no content has arrived yet. Started, diff --git a/src/harness/no_progress.rs b/src/harness/no_progress/mod.rs similarity index 54% rename from src/harness/no_progress.rs rename to src/harness/no_progress/mod.rs index 1017255..981d2c8 100644 --- a/src/harness/no_progress.rs +++ b/src/harness/no_progress/mod.rs @@ -19,6 +19,11 @@ //! and turns the returned [`NoProgress`] verdict into a steering nudge //! (`Nudge`) or a halt (`Halt`). +mod types; + +use types::LadderState; +pub use types::{NoProgress, NoProgressTracker, ToolAttempt}; + use std::sync::Mutex; /// Consecutive **identical** (tool + args + error) failures tolerated before the @@ -39,67 +44,6 @@ const NO_PROGRESS_NUDGE_THRESHOLD: usize = 4; /// call re-issued unchanged can never succeed. const HARD_REJECT_HALT_THRESHOLD: usize = 2; -/// One tool call's outcome, reduced to the deterministic parts the ladder -/// compares. Built by the driver from a tool result plus the argument -/// fingerprint captured before execution. -pub struct ToolAttempt<'a> { - /// Tool name. - pub tool: &'a str, - /// Stable fingerprint of the call arguments (computed by the driver). Folded - /// into the identical-repeat signature so the "identical arguments" ladder - /// only trips when the args truly repeat. - pub arg_fingerprint: &'a str, - /// `None` on success; otherwise the tool's error text. - pub error: Option<&'a str>, - /// `true` when the result is a hard security/approval rejection that can - /// never succeed re-issued unchanged. - pub hard_reject: bool, - /// `true` for the unknown-tool recovery sentinel — a correctable miss that - /// must not feed the generic any-failure backstop (it still feeds the - /// identical-repeat counter, so re-issuing the *same* unavailable tool - /// halts). - pub recoverable_miss: bool, -} - -/// The ladder's verdict for one recorded attempt. -#[derive(Debug, Clone, PartialEq, Eq)] -pub enum NoProgress { - /// Progress was made, or not enough repetition yet — carry on. - Continue, - /// Same-strategy repetition detected below the retry cap: feed this - /// structured "no progress since step X" corrective back into the loop so - /// the model picks a *different* next action. - Nudge(String), - /// Same-strategy retries exhausted (or the any-failure backstop tripped): - /// halt with this root-cause summary. - Halt(String), -} - -#[derive(Default)] -struct LadderState { - /// Signature of the previous failing call (tool + args + first error line). - last_sig: Option, - /// Consecutive repeats of `last_sig`. - same_count: usize, - /// Consecutive failures of any kind (reset by any success). - consecutive: usize, - /// Signature we have already nudged on, so a nudge fires at most once per - /// distinct failing `(tool, args, error)` before escalating to a halt. - nudged_sig: Option, - /// `true` once the varied-failure nudge fired for the current streak. - nudged_streak: bool, -} - -/// Tracks recent tool outcomes and drives the no-progress escalation ladder. -/// -/// Cheap to construct and interior-mutable, so a middleware can hold one behind -/// a shared reference for the whole turn. `identical_halt_threshold` is the -/// same-strategy retry cap; it is clamped so a nudge always precedes a halt. -pub struct NoProgressTracker { - identical_halt_threshold: usize, - state: Mutex, -} - impl NoProgressTracker { /// Build a tracker whose identical-repeat halt threshold is /// `identical_halt_threshold`, clamped up so it always sits above the nudge @@ -261,144 +205,4 @@ fn truncate_for_halt(text: &str) -> String { } #[cfg(test)] -mod tests { - use super::*; - - fn fail<'a>(tool: &'a str, fp: &'a str, err: &'a str) -> ToolAttempt<'a> { - ToolAttempt { - tool, - arg_fingerprint: fp, - error: Some(err), - hard_reject: false, - recoverable_miss: false, - } - } - - fn ok<'a>(tool: &'a str, fp: &'a str) -> ToolAttempt<'a> { - ToolAttempt { - tool, - arg_fingerprint: fp, - error: None, - hard_reject: false, - recoverable_miss: false, - } - } - - #[test] - fn identical_failure_nudges_then_halts() { - let t = NoProgressTracker::new(DEFAULT_IDENTICAL_HALT_THRESHOLD); - // First identical failure: not enough repetition yet. - assert_eq!( - t.record(1, &fail("read_file", "a", "not found")), - NoProgress::Continue - ); - // Second: nudge the model to change strategy before the retry cap. - match t.record(2, &fail("read_file", "a", "not found")) { - NoProgress::Nudge(msg) => { - assert!(msg.contains("no progress since step 2")); - assert!(msg.contains("read_file")); - assert!(msg.contains("Change strategy")); - } - other => panic!("expected a nudge on the second identical failure, got {other:?}"), - } - // Third: same-strategy retries exhausted → halt. - match t.record(3, &fail("read_file", "a", "not found")) { - NoProgress::Halt(msg) => assert!(msg.contains("retried 3 times")), - other => panic!("expected a halt on the third identical failure, got {other:?}"), - } - } - - #[test] - fn a_success_resets_the_ladder() { - let t = NoProgressTracker::new(DEFAULT_IDENTICAL_HALT_THRESHOLD); - let _ = t.record(1, &fail("t", "a", "boom")); - let _ = t.record(2, &fail("t", "a", "boom")); // nudge - assert_eq!(t.record(3, &ok("t", "a")), NoProgress::Continue); - // After the reset, two more identical failures only re-nudge — no halt. - assert_eq!(t.record(4, &fail("t", "a", "boom")), NoProgress::Continue); - assert!(matches!( - t.record(5, &fail("t", "a", "boom")), - NoProgress::Nudge(_) - )); - } - - #[test] - fn changed_arguments_clear_the_identical_streak() { - let t = NoProgressTracker::new(DEFAULT_IDENTICAL_HALT_THRESHOLD); - let _ = t.record(1, &fail("read_file", "a", "not found")); - // The model heeded the nudge and changed args: a new signature, so the - // identical-repeat counter starts over and never reaches the halt. - assert_eq!( - t.record(2, &fail("read_file", "b", "not found")), - NoProgress::Continue - ); - assert_eq!( - t.record(3, &fail("list_dir", "c", "denied")), - NoProgress::Continue - ); - } - - #[test] - fn a_hard_rejection_halts_on_the_second_identical_repeat() { - let t = NoProgressTracker::new(DEFAULT_IDENTICAL_HALT_THRESHOLD); - let mut a = fail("send_email", "a", "[policy-blocked] denied"); - a.hard_reject = true; - assert_eq!(t.record(1, &a), NoProgress::Continue); - let mut b = fail("send_email", "a", "[policy-blocked] denied"); - b.hard_reject = true; - assert!( - matches!(t.record(2, &b), NoProgress::Halt(msg) if msg.contains("security policy")) - ); - } - - #[test] - fn varied_failures_nudge_then_hit_the_backstop() { - let t = NoProgressTracker::new(DEFAULT_IDENTICAL_HALT_THRESHOLD); - // Distinct signatures each time: the identical counter never climbs, but - // the any-failure streak does. - for i in 1..=3 { - let (fp, err) = (format!("fp{i}"), format!("err{i}")); - assert_eq!(t.record(i, &fail("t", &fp, &err)), NoProgress::Continue); - } - // Fourth consecutive varied failure → nudge to step back. - assert!(matches!( - t.record(4, &fail("t", "fp4", "err4")), - NoProgress::Nudge(_) - )); - assert_eq!(t.record(5, &fail("t", "fp5", "err5")), NoProgress::Continue); - // Sixth → the no-progress backstop halts. - assert!(matches!( - t.record(6, &fail("t", "fp6", "err6")), - NoProgress::Halt(msg) if msg.contains("in a row failed") - )); - } - - #[test] - fn unknown_tool_recovery_does_not_feed_the_backstop() { - let t = NoProgressTracker::new(DEFAULT_IDENTICAL_HALT_THRESHOLD); - // Six correctable misses with *distinct* signatures must not trip the - // any-failure backstop (they don't count toward `consecutive`). - for i in 0..6 { - let (fp, err) = (format!("fp{i}"), format!("unknown tool foo{i}")); - let mut a = fail("__unknown_tool__", &fp, &err); - a.recoverable_miss = true; - assert_eq!(t.record(i, &a), NoProgress::Continue); - } - } - - #[test] - fn identical_halt_threshold_is_clamped_above_the_nudge() { - // A caller asking for a halt threshold of 1 or 2 still gets a nudge - // before the halt (the nudge lands at 2, so the halt must be >= 3). - let t = NoProgressTracker::new(1); - assert_eq!(t.record(1, &fail("t", "a", "boom")), NoProgress::Continue); - assert!(matches!( - t.record(2, &fail("t", "a", "boom")), - NoProgress::Nudge(_) - )); - assert!(matches!( - t.record(3, &fail("t", "a", "boom")), - NoProgress::Halt(_) - )); - } -} +mod test; diff --git a/src/harness/no_progress/test.rs b/src/harness/no_progress/test.rs new file mode 100644 index 0000000..9c5fbde --- /dev/null +++ b/src/harness/no_progress/test.rs @@ -0,0 +1,139 @@ +//! Unit tests for the no-progress escalation ladder. + +use super::*; + +fn fail<'a>(tool: &'a str, fp: &'a str, err: &'a str) -> ToolAttempt<'a> { + ToolAttempt { + tool, + arg_fingerprint: fp, + error: Some(err), + hard_reject: false, + recoverable_miss: false, + } +} + +fn ok<'a>(tool: &'a str, fp: &'a str) -> ToolAttempt<'a> { + ToolAttempt { + tool, + arg_fingerprint: fp, + error: None, + hard_reject: false, + recoverable_miss: false, + } +} + +#[test] +fn identical_failure_nudges_then_halts() { + let t = NoProgressTracker::new(DEFAULT_IDENTICAL_HALT_THRESHOLD); + // First identical failure: not enough repetition yet. + assert_eq!( + t.record(1, &fail("read_file", "a", "not found")), + NoProgress::Continue + ); + // Second: nudge the model to change strategy before the retry cap. + match t.record(2, &fail("read_file", "a", "not found")) { + NoProgress::Nudge(msg) => { + assert!(msg.contains("no progress since step 2")); + assert!(msg.contains("read_file")); + assert!(msg.contains("Change strategy")); + } + other => panic!("expected a nudge on the second identical failure, got {other:?}"), + } + // Third: same-strategy retries exhausted → halt. + match t.record(3, &fail("read_file", "a", "not found")) { + NoProgress::Halt(msg) => assert!(msg.contains("retried 3 times")), + other => panic!("expected a halt on the third identical failure, got {other:?}"), + } +} + +#[test] +fn a_success_resets_the_ladder() { + let t = NoProgressTracker::new(DEFAULT_IDENTICAL_HALT_THRESHOLD); + let _ = t.record(1, &fail("t", "a", "boom")); + let _ = t.record(2, &fail("t", "a", "boom")); // nudge + assert_eq!(t.record(3, &ok("t", "a")), NoProgress::Continue); + // After the reset, two more identical failures only re-nudge — no halt. + assert_eq!(t.record(4, &fail("t", "a", "boom")), NoProgress::Continue); + assert!(matches!( + t.record(5, &fail("t", "a", "boom")), + NoProgress::Nudge(_) + )); +} + +#[test] +fn changed_arguments_clear_the_identical_streak() { + let t = NoProgressTracker::new(DEFAULT_IDENTICAL_HALT_THRESHOLD); + let _ = t.record(1, &fail("read_file", "a", "not found")); + // The model heeded the nudge and changed args: a new signature, so the + // identical-repeat counter starts over and never reaches the halt. + assert_eq!( + t.record(2, &fail("read_file", "b", "not found")), + NoProgress::Continue + ); + assert_eq!( + t.record(3, &fail("list_dir", "c", "denied")), + NoProgress::Continue + ); +} + +#[test] +fn a_hard_rejection_halts_on_the_second_identical_repeat() { + let t = NoProgressTracker::new(DEFAULT_IDENTICAL_HALT_THRESHOLD); + let mut a = fail("send_email", "a", "[policy-blocked] denied"); + a.hard_reject = true; + assert_eq!(t.record(1, &a), NoProgress::Continue); + let mut b = fail("send_email", "a", "[policy-blocked] denied"); + b.hard_reject = true; + assert!(matches!(t.record(2, &b), NoProgress::Halt(msg) if msg.contains("security policy"))); +} + +#[test] +fn varied_failures_nudge_then_hit_the_backstop() { + let t = NoProgressTracker::new(DEFAULT_IDENTICAL_HALT_THRESHOLD); + // Distinct signatures each time: the identical counter never climbs, but + // the any-failure streak does. + for i in 1..=3 { + let (fp, err) = (format!("fp{i}"), format!("err{i}")); + assert_eq!(t.record(i, &fail("t", &fp, &err)), NoProgress::Continue); + } + // Fourth consecutive varied failure → nudge to step back. + assert!(matches!( + t.record(4, &fail("t", "fp4", "err4")), + NoProgress::Nudge(_) + )); + assert_eq!(t.record(5, &fail("t", "fp5", "err5")), NoProgress::Continue); + // Sixth → the no-progress backstop halts. + assert!(matches!( + t.record(6, &fail("t", "fp6", "err6")), + NoProgress::Halt(msg) if msg.contains("in a row failed") + )); +} + +#[test] +fn unknown_tool_recovery_does_not_feed_the_backstop() { + let t = NoProgressTracker::new(DEFAULT_IDENTICAL_HALT_THRESHOLD); + // Six correctable misses with *distinct* signatures must not trip the + // any-failure backstop (they don't count toward `consecutive`). + for i in 0..6 { + let (fp, err) = (format!("fp{i}"), format!("unknown tool foo{i}")); + let mut a = fail("__unknown_tool__", &fp, &err); + a.recoverable_miss = true; + assert_eq!(t.record(i, &a), NoProgress::Continue); + } +} + +#[test] +fn identical_halt_threshold_is_clamped_above_the_nudge() { + // A caller asking for a halt threshold of 1 or 2 still gets a nudge + // before the halt (the nudge lands at 2, so the halt must be >= 3). + let t = NoProgressTracker::new(1); + assert_eq!(t.record(1, &fail("t", "a", "boom")), NoProgress::Continue); + assert!(matches!( + t.record(2, &fail("t", "a", "boom")), + NoProgress::Nudge(_) + )); + assert!(matches!( + t.record(3, &fail("t", "a", "boom")), + NoProgress::Halt(_) + )); +} diff --git a/src/harness/no_progress/types.rs b/src/harness/no_progress/types.rs new file mode 100644 index 0000000..94c8e73 --- /dev/null +++ b/src/harness/no_progress/types.rs @@ -0,0 +1,65 @@ +//! Type definitions for the no-progress escalation ladder +//! (`ToolAttempt`, `NoProgress`, `LadderState`, `NoProgressTracker`). +//! +//! Split out of `no_progress/mod.rs`; see that module's doc comment for +//! the full escalation-ladder design. + +use std::sync::Mutex; + +pub struct ToolAttempt<'a> { + /// Tool name. + pub tool: &'a str, + /// Stable fingerprint of the call arguments (computed by the driver). Folded + /// into the identical-repeat signature so the "identical arguments" ladder + /// only trips when the args truly repeat. + pub arg_fingerprint: &'a str, + /// `None` on success; otherwise the tool's error text. + pub error: Option<&'a str>, + /// `true` when the result is a hard security/approval rejection that can + /// never succeed re-issued unchanged. + pub hard_reject: bool, + /// `true` for the unknown-tool recovery sentinel — a correctable miss that + /// must not feed the generic any-failure backstop (it still feeds the + /// identical-repeat counter, so re-issuing the *same* unavailable tool + /// halts). + pub recoverable_miss: bool, +} + +/// The ladder's verdict for one recorded attempt. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum NoProgress { + /// Progress was made, or not enough repetition yet — carry on. + Continue, + /// Same-strategy repetition detected below the retry cap: feed this + /// structured "no progress since step X" corrective back into the loop so + /// the model picks a *different* next action. + Nudge(String), + /// Same-strategy retries exhausted (or the any-failure backstop tripped): + /// halt with this root-cause summary. + Halt(String), +} + +#[derive(Default)] +pub(super) struct LadderState { + /// Signature of the previous failing call (tool + args + first error line). + pub(super) last_sig: Option, + /// Consecutive repeats of `last_sig`. + pub(super) same_count: usize, + /// Consecutive failures of any kind (reset by any success). + pub(super) consecutive: usize, + /// Signature we have already nudged on, so a nudge fires at most once per + /// distinct failing `(tool, args, error)` before escalating to a halt. + pub(super) nudged_sig: Option, + /// `true` once the varied-failure nudge fired for the current streak. + pub(super) nudged_streak: bool, +} + +/// Tracks recent tool outcomes and drives the no-progress escalation ladder. +/// +/// Cheap to construct and interior-mutable, so a middleware can hold one behind +/// a shared reference for the whole turn. `identical_halt_threshold` is the +/// same-strategy retry cap; it is clamped so a nudge always precedes a halt. +pub struct NoProgressTracker { + pub(super) identical_halt_threshold: usize, + pub(super) state: Mutex, +} diff --git a/src/harness/observability/README.md b/src/harness/observability/README.md new file mode 100644 index 0000000..b993ec3 --- /dev/null +++ b/src/harness/observability/README.md @@ -0,0 +1,84 @@ +# harness::observability + +Durable observability for the harness — journals, status stores, and sinks. + +The live `harness::events` layer fans typed `AgentEvent`s out to in-process +listeners for the duration of a run. This module makes that history +**durable and correlatable** so a UI, supervisor, or test can reconstruct a +recursive run tree after the fact, including across process restarts. + +See `docs/modules/harness/observability.md` for the design rationale +(inspiration, responsibilities, cross-cutting requirements shared with +`graph::observability`). This README documents the module's current public +surface and operational constraints. + +## Public surface + +- `AgentObservation` — a durable envelope pairing an `AgentEvent` with its run + lineage (`run_id` / `parent_run_id` / `root_run_id`), a stream `offset`, and + a `ts_ms` timestamp. This is the unit everything else in the module is built + from. +- `HarnessEventJournal` (trait) — an append-only, offset-addressable journal of + observations. + - `InMemoryEventJournal` — in-process implementation for tests and + short-lived processes. + - `StoreEventJournal` — store-backed implementation; the + stream key is the run id, so a `harness::store::AppendStore` (e.g. a + JSONL file or Sqlite-backed store) durably persists observations per run. +- `HarnessStatusStore` (trait) — a compact "what is running now?" surface + (phase, counters, last-updated) distinct from the full journal. + - `InMemoryStatusStore` — in-process implementation. +- **Sinks**, each implementing `harness::events::EventListener`: + - `FanOutSink` — broadcasts one event to multiple inner listeners. + - `RedactingSink` — masks secrets in event payloads before forwarding. + - `JournalSink` — persists observations into a `HarnessEventJournal`. + - `JsonlSink` — appends records to a JSONL stream via + `harness::store::JsonlAppendStore`. +- `AgentObservation`-derived metrics: + - `AgentCallLatency` — start/end/elapsed for a single model or tool call. + - `AgentLatencyMetrics` — latency rollups for one run, built with + `AgentLatencyMetrics::from_observations(&[AgentObservation])`. +- `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. + +## Persistence bridge + +Persisting sinks (`JournalSink`, `JsonlSink`) bridge the synchronous +`EventListener::on_event` hook to the async journal/store APIs through a +background `AppendWorker`: `on_event` hands the observation to a **bounded** +channel drained on a dedicated thread, so it never blocks the run on I/O. +Persistence is **best-effort**: if the queue is full the observation is dropped +(and counted), backend errors are reported to stderr rather than propagated, +and neither ever aborts the run. Call `JournalSink::flush` / `JsonlSink::flush` +to block until the durable log has caught up (for example before reading it back +or shutting down). Do not rely on a sink for delivery guarantees stronger than +"usually persisted, never run-blocking." + +## Latency metrics semantics + +`AgentLatencyMetrics::from_observations` tolerates redacted payload strings +(structural ids must be preserved by any redaction upstream) but **ignores +incomplete calls** — a `ModelStarted`/`ToolStarted` with no matching +`*Completed`/`*Failed` has no terminal timestamp to measure against and is +silently excluded from the rollup rather than reported with a bogus duration. + +## Files + +| File | Role | +| --- | --- | +| `types.rs` | Every public type: `AgentObservation`, journal/status-store traits and in-memory impls, sink structs, latency types. | +| `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 `test.rs`. | +| `test.rs` | Unit tests (journal round-trips, redaction, latency rollups, sink fan-out). | + +## Operational constraints + +- `StoreEventJournal` keys the stream by run id; using the same `AppendStore` + for unrelated runs is safe (streams are namespaced) but reusing a run id + across logically distinct runs will interleave their observations. +- `RedactingSink` must be composed *outside* any sink that persists or + exports off-process (e.g. wrap before `JournalSink`/`JsonlSink`/Langfuse) — + it only redacts what passes through it, not what already landed elsewhere. +- The Langfuse client performs network I/O; failures there follow the same + best-effort, non-aborting policy as the other persisting sinks. diff --git a/src/harness/observability/langfuse.rs b/src/harness/observability/langfuse/mod.rs similarity index 63% rename from src/harness/observability/langfuse.rs rename to src/harness/observability/langfuse/mod.rs index ba47698..a260e4c 100644 --- a/src/harness/observability/langfuse.rs +++ b/src/harness/observability/langfuse/mod.rs @@ -4,66 +4,17 @@ //! send directly to a self-hosted Langfuse instance with public/secret keys, or //! to the TinyHumans backend proxy with a bearer token. -use serde::Serialize; use serde_json::{Value, json}; use crate::error::{Result, TinyAgentsError}; use crate::harness::events::AgentEvent; +use crate::harness::ids::now_ms; use crate::harness::observability::AgentObservation; use crate::harness::usage::Usage; -/// Authentication mode for [`LangfuseClient`]. -#[derive(Clone, Debug, PartialEq, Eq)] -pub enum LangfuseAuth { - /// Send `Authorization: Basic base64(public_key:secret_key)`. - Basic { - /// Langfuse project public key. - public_key: String, - /// Langfuse project secret key. - secret_key: String, - }, - /// Send `Authorization: Bearer `. - /// - /// Use this when targeting the TinyHumans backend proxy at - /// `/telemetry/langfuse/ingestion`; the backend injects Langfuse Basic Auth. - Bearer { - /// Backend access token. - token: String, - }, -} +mod types; -/// Configuration for a Langfuse trace export. -#[derive(Clone, Debug, Default, PartialEq, Serialize)] -pub struct LangfuseTraceConfig { - /// Stable Langfuse trace id. Defaults to the first observation's root run id. - pub trace_id: Option, - /// Human-readable trace name. - pub name: Option, - /// End-user id to filter by in Langfuse. - pub user_id: Option, - /// Session/thread id to group related traces. - pub session_id: Option, - /// Langfuse environment name. - pub environment: Option, - /// Release identifier. - pub release: Option, - /// Version identifier. - pub version: Option, - /// Tags attached to the trace. - #[serde(default)] - pub tags: Vec, - /// Extra trace metadata. - #[serde(default)] - pub metadata: Value, -} - -/// Async Langfuse ingestion client. -#[derive(Clone, Debug)] -pub struct LangfuseClient { - endpoint: String, - auth: LangfuseAuth, - client: reqwest::Client, -} +pub use types::{LangfuseAuth, LangfuseClient, LangfuseTraceConfig}; impl LangfuseClient { /// Creates a client from a base URL and auth mode. @@ -227,6 +178,19 @@ impl LangfuseClient { "Langfuse ingestion returned {status}: {parsed}" ))); } + // A `207 Multi-Status` reports per-item outcomes: some events may have + // been rejected while the request itself "succeeded". Surface those + // partial failures instead of swallowing them and reporting success. + if status.as_u16() == 207 + && let Some(errors) = parsed.get("errors").and_then(Value::as_array) + && !errors.is_empty() + { + return Err(TinyAgentsError::Model(format!( + "Langfuse ingestion partially failed ({} rejected): {}", + errors.len(), + json!(errors) + ))); + } Ok(parsed) } } @@ -307,12 +271,17 @@ fn trace_metadata(trace: &LangfuseTraceConfig, observations: &[AgentObservation] fn observation_event(trace_id: &str, obs: &AgentObservation) -> Value { let timestamp = iso_ms(obs.ts_ms); + // Attach only run lineage, offset, and the event *kind* — not the full event + // payload. The event's meaningful fields (input/output/usage/name) are + // already lifted into the observation `body`; embedding the whole event here + // duplicated every payload, roughly doubling batch bytes and growing + // O(turns^2) as a run accumulates, which can trip the ~3.5MB batch cap. let 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()), "offset": obs.offset, - "event": obs.event, + "event_kind": obs.event.kind(), }); match &obs.event { AgentEvent::ModelCompleted { @@ -344,7 +313,10 @@ fn observation_event(trace_id: &str, obs: &AgentObservation) -> Value { } => json!({ "id": obs.event_id.as_str(), "timestamp": timestamp, - "type": "tool-create", + // A tool call is modelled as a span. `tool-create` is not a valid + // Langfuse ingestion observation type — older/self-hosted Langfuse + // rejects it, silently dropping every tool observation. + "type": "span-create", "body": clean_nulls(json!({ "id": call_id.as_str(), "traceId": trace_id, @@ -412,14 +384,6 @@ pub(crate) fn iso_ms(ms: u64) -> String { format_unix_iso(secs, millis) } -fn now_ms() -> u64 { - use std::time::{SystemTime, UNIX_EPOCH}; - SystemTime::now() - .duration_since(UNIX_EPOCH) - .map(|d| d.as_millis() as u64) - .unwrap_or(0) -} - fn format_unix_iso(secs: u64, millis: u32) -> String { // Howard Hinnant civil-date conversion for Unix days, dependency-free. let days = (secs / 86_400) as i64; @@ -446,159 +410,4 @@ fn civil_from_days(days: i64) -> (i32, u32, u32) { } #[cfg(test)] -mod tests { - use super::*; - use crate::harness::events::AgentEvent; - use crate::harness::ids::{CallId, EventId, RunId}; - - fn obs(offset: u64, event: AgentEvent) -> AgentObservation { - AgentObservation { - event_id: EventId::new(format!("evt-{offset}")), - run_id: RunId::new("run-1"), - parent_run_id: None, - root_run_id: RunId::new("root-1"), - offset, - ts_ms: 1_704_067_200_000 + offset, - event, - } - } - - #[test] - fn normalizes_langfuse_endpoints() { - let client = LangfuseClient::proxy("https://api.example.test", "token").unwrap(); - assert_eq!( - client.endpoint(), - "https://api.example.test/telemetry/langfuse/ingestion" - ); - let client = LangfuseClient::proxy( - "https://api.example.test/telemetry/langfuse/ingestion", - "token", - ) - .unwrap(); - assert_eq!( - client.endpoint(), - "https://api.example.test/telemetry/langfuse/ingestion" - ); - } - - #[test] - fn builds_trace_and_generation_batch() { - let client = - LangfuseClient::proxy("https://backend.test/telemetry/langfuse/ingestion", "t") - .unwrap(); - let batch = client - .build_ingestion_batch( - LangfuseTraceConfig { - user_id: Some("user-1".to_string()), - session_id: Some("thread-1".to_string()), - ..Default::default() - }, - &[ - obs( - 0, - AgentEvent::RunStarted { - run_id: RunId::new("run-1"), - thread_id: None, - }, - ), - obs( - 1, - AgentEvent::ModelCompleted { - call_id: CallId::new("model-call"), - usage: Some(Usage { - input_tokens: 3, - output_tokens: 4, - total_tokens: 7, - ..Default::default() - }), - input: None, - output: None, - }, - ), - ], - ) - .unwrap(); - - let events = batch["batch"].as_array().unwrap(); - assert_eq!(events[0]["type"], "trace-create"); - assert_eq!(events[0]["body"]["id"], "root-1"); - assert_eq!(events[0]["body"]["userId"], "user-1"); - // Trace metadata is defaulted from run lineage even without a caller value. - assert_eq!(events[0]["body"]["metadata"]["root_run_id"], "root-1"); - assert_eq!(events[0]["body"]["metadata"]["run_id"], "run-1"); - assert_eq!(events[2]["type"], "generation-create"); - assert_eq!(events[2]["body"]["id"], "model-call"); - assert_eq!(events[2]["body"]["usage"]["input"], 3); - // Payload-free generation: no input/output body fields. - assert!(events[2]["body"].get("input").is_none()); - assert!(events[2]["body"].get("output").is_none()); - } - - #[test] - fn populates_generation_and_tool_io_when_captured() { - let client = - LangfuseClient::proxy("https://backend.test/telemetry/langfuse/ingestion", "t") - .unwrap(); - let batch = client - .build_ingestion_batch( - LangfuseTraceConfig::default(), - &[ - obs( - 0, - AgentEvent::ModelCompleted { - call_id: CallId::new("model-call"), - usage: None, - input: Some(json!([{ "role": "user", "content": "hi" }])), - output: Some(json!({ "content": "hello there" })), - }, - ), - obs( - 1, - AgentEvent::ToolCompleted { - call_id: CallId::new("tool-call"), - tool_name: "lookup".to_string(), - input: Some(json!({ "query": "weather" })), - output: Some(json!("sunny")), - }, - ), - ], - ) - .unwrap(); - - let events = batch["batch"].as_array().unwrap(); - assert_eq!(events[1]["type"], "generation-create"); - assert_eq!(events[1]["body"]["input"][0]["content"], "hi"); - assert_eq!(events[1]["body"]["output"]["content"], "hello there"); - assert_eq!(events[2]["type"], "tool-create"); - assert_eq!(events[2]["body"]["input"]["query"], "weather"); - assert_eq!(events[2]["body"]["output"], "sunny"); - } - - #[test] - fn merges_caller_trace_metadata_over_defaults() { - let client = - LangfuseClient::proxy("https://backend.test/telemetry/langfuse/ingestion", "t") - .unwrap(); - let batch = client - .build_ingestion_batch( - LangfuseTraceConfig { - metadata: json!({ "deployment": "prod", "root_run_id": "override" }), - ..Default::default() - }, - &[obs( - 0, - AgentEvent::RunStarted { - run_id: RunId::new("run-1"), - thread_id: None, - }, - )], - ) - .unwrap(); - - let meta = &batch["batch"].as_array().unwrap()[0]["body"]["metadata"]; - assert_eq!(meta["deployment"], "prod"); - // Caller keys win on collision with the defaulted lineage. - assert_eq!(meta["root_run_id"], "override"); - assert_eq!(meta["run_id"], "run-1"); - } -} +mod test; diff --git a/src/harness/observability/langfuse/test.rs b/src/harness/observability/langfuse/test.rs new file mode 100644 index 0000000..6b2365f --- /dev/null +++ b/src/harness/observability/langfuse/test.rs @@ -0,0 +1,163 @@ +//! Unit tests for the Langfuse ingestion exporter. + +use super::*; +use crate::harness::events::AgentEvent; +use crate::harness::ids::{CallId, EventId, RunId}; + +fn obs(offset: u64, event: AgentEvent) -> AgentObservation { + AgentObservation { + event_id: EventId::new(format!("evt-{offset}")), + run_id: RunId::new("run-1"), + parent_run_id: None, + root_run_id: RunId::new("root-1"), + offset, + ts_ms: 1_704_067_200_000 + offset, + event, + } +} + +#[test] +fn normalizes_langfuse_endpoints() { + let client = LangfuseClient::proxy("https://api.example.test", "token").unwrap(); + assert_eq!( + client.endpoint(), + "https://api.example.test/telemetry/langfuse/ingestion" + ); + let client = LangfuseClient::proxy( + "https://api.example.test/telemetry/langfuse/ingestion", + "token", + ) + .unwrap(); + assert_eq!( + client.endpoint(), + "https://api.example.test/telemetry/langfuse/ingestion" + ); +} + +#[test] +fn builds_trace_and_generation_batch() { + let client = + LangfuseClient::proxy("https://backend.test/telemetry/langfuse/ingestion", "t").unwrap(); + let batch = client + .build_ingestion_batch( + LangfuseTraceConfig { + user_id: Some("user-1".to_string()), + session_id: Some("thread-1".to_string()), + ..Default::default() + }, + &[ + obs( + 0, + AgentEvent::RunStarted { + run_id: RunId::new("run-1"), + thread_id: None, + }, + ), + obs( + 1, + AgentEvent::ModelCompleted { + call_id: CallId::new("model-call"), + usage: Some(Usage { + input_tokens: 3, + output_tokens: 4, + total_tokens: 7, + ..Default::default() + }), + input: None, + output: None, + }, + ), + ], + ) + .unwrap(); + + let events = batch["batch"].as_array().unwrap(); + assert_eq!(events[0]["type"], "trace-create"); + assert_eq!(events[0]["body"]["id"], "root-1"); + assert_eq!(events[0]["body"]["userId"], "user-1"); + // Trace metadata is defaulted from run lineage even without a caller value. + assert_eq!(events[0]["body"]["metadata"]["root_run_id"], "root-1"); + assert_eq!(events[0]["body"]["metadata"]["run_id"], "run-1"); + assert_eq!(events[2]["type"], "generation-create"); + assert_eq!(events[2]["body"]["id"], "model-call"); + assert_eq!(events[2]["body"]["usage"]["input"], 3); + // Payload-free generation: no input/output body fields. + assert!(events[2]["body"].get("input").is_none()); + assert!(events[2]["body"].get("output").is_none()); +} + +#[test] +fn populates_generation_and_tool_io_when_captured() { + let client = + LangfuseClient::proxy("https://backend.test/telemetry/langfuse/ingestion", "t").unwrap(); + let batch = client + .build_ingestion_batch( + LangfuseTraceConfig::default(), + &[ + obs( + 0, + AgentEvent::ModelCompleted { + call_id: CallId::new("model-call"), + usage: None, + input: Some(json!([{ "role": "user", "content": "hi" }])), + output: Some(json!({ "content": "hello there" })), + }, + ), + obs( + 1, + AgentEvent::ToolCompleted { + call_id: CallId::new("tool-call"), + tool_name: "lookup".to_string(), + input: Some(json!({ "query": "weather" })), + output: Some(json!("sunny")), + }, + ), + ], + ) + .unwrap(); + + let events = batch["batch"].as_array().unwrap(); + assert_eq!(events[1]["type"], "generation-create"); + assert_eq!(events[1]["body"]["input"][0]["content"], "hi"); + assert_eq!(events[1]["body"]["output"]["content"], "hello there"); + assert_eq!(events[2]["type"], "span-create"); + assert_eq!(events[2]["body"]["input"]["query"], "weather"); + assert_eq!(events[2]["body"]["output"], "sunny"); + + // Observation metadata carries only lineage + event kind, not the whole + // event payload (which would duplicate input/output already in `body`). + let gen_meta = &events[1]["body"]["metadata"]; + assert_eq!(gen_meta["event_kind"], "model.completed"); + assert!( + gen_meta.get("event").is_none(), + "full event payload must not be duplicated into metadata" + ); + assert_eq!(gen_meta["run_id"], "run-1"); +} + +#[test] +fn merges_caller_trace_metadata_over_defaults() { + let client = + LangfuseClient::proxy("https://backend.test/telemetry/langfuse/ingestion", "t").unwrap(); + let batch = client + .build_ingestion_batch( + LangfuseTraceConfig { + metadata: json!({ "deployment": "prod", "root_run_id": "override" }), + ..Default::default() + }, + &[obs( + 0, + AgentEvent::RunStarted { + run_id: RunId::new("run-1"), + thread_id: None, + }, + )], + ) + .unwrap(); + + let meta = &batch["batch"].as_array().unwrap()[0]["body"]["metadata"]; + assert_eq!(meta["deployment"], "prod"); + // Caller keys win on collision with the defaulted lineage. + assert_eq!(meta["root_run_id"], "override"); + assert_eq!(meta["run_id"], "run-1"); +} diff --git a/src/harness/observability/langfuse/types.rs b/src/harness/observability/langfuse/types.rs new file mode 100644 index 0000000..61cdef1 --- /dev/null +++ b/src/harness/observability/langfuse/types.rs @@ -0,0 +1,60 @@ +//! Type definitions for the Langfuse ingestion exporter. +//! +//! Split out of `langfuse/mod.rs`; see that module's doc comment for the +//! exporter overview. + +use serde::Serialize; +use serde_json::Value; + +/// Authentication mode for [`LangfuseClient`]. +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum LangfuseAuth { + /// Send `Authorization: Basic base64(public_key:secret_key)`. + Basic { + /// Langfuse project public key. + public_key: String, + /// Langfuse project secret key. + secret_key: String, + }, + /// Send `Authorization: Bearer `. + /// + /// Use this when targeting the TinyHumans backend proxy at + /// `/telemetry/langfuse/ingestion`; the backend injects Langfuse Basic Auth. + Bearer { + /// Backend access token. + token: String, + }, +} + +/// Configuration for a Langfuse trace export. +#[derive(Clone, Debug, Default, PartialEq, Serialize)] +pub struct LangfuseTraceConfig { + /// Stable Langfuse trace id. Defaults to the first observation's root run id. + pub trace_id: Option, + /// Human-readable trace name. + pub name: Option, + /// End-user id to filter by in Langfuse. + pub user_id: Option, + /// Session/thread id to group related traces. + pub session_id: Option, + /// Langfuse environment name. + pub environment: Option, + /// Release identifier. + pub release: Option, + /// Version identifier. + pub version: Option, + /// Tags attached to the trace. + #[serde(default)] + pub tags: Vec, + /// Extra trace metadata. + #[serde(default)] + pub metadata: Value, +} + +/// Async Langfuse ingestion client. +#[derive(Clone, Debug)] +pub struct LangfuseClient { + pub(super) endpoint: String, + pub(super) auth: LangfuseAuth, + pub(super) client: reqwest::Client, +} diff --git a/src/harness/observability/mod.rs b/src/harness/observability/mod.rs index d789558..8acbbd1 100644 --- a/src/harness/observability/mod.rs +++ b/src/harness/observability/mod.rs @@ -19,11 +19,16 @@ //! to a JSONL stream). //! //! Persisting sinks bridge the synchronous [`EventListener::on_event`] hook to -//! the async journal/store APIs with `futures::executor::block_on` and treat -//! persistence as best-effort: a backend error never aborts the run. +//! the async journal/store APIs through a background [`worker::AppendWorker`]: +//! `on_event` never blocks the run on I/O, persistence is best-effort (a full +//! bounded queue drops rather than stalls; backend errors are reported, not +//! propagated), and `flush` blocks until the durable log has caught up. mod langfuse; mod types; +mod worker; + +pub(crate) use worker::{AppendWorker, DEFAULT_DRAIN_CAPACITY}; pub use langfuse::{LangfuseAuth, LangfuseClient, LangfuseTraceConfig}; // Shared Langfuse payload helpers reused by the graph observability exporter so @@ -32,7 +37,7 @@ pub(crate) use langfuse::{clean_nulls, iso_ms}; pub use types::*; use std::collections::HashMap; -use std::sync::Arc; +use std::sync::{Arc, Mutex}; use std::time::SystemTime; use async_trait::async_trait; @@ -144,18 +149,30 @@ impl AgentLatencyMetrics { // --------------------------------------------------------------------------- impl InMemoryEventJournal { - /// Creates a new, empty in-memory journal. + /// Creates a new, empty in-memory journal with the default run-retention + /// cap ([`DEFAULT_JOURNAL_MAX_RUNS`]). pub fn new() -> Self { Self::default() } + /// Creates a new, empty in-memory journal that retains at most + /// `max_runs` distinct `run_id` streams, evicting the oldest (by first + /// append) once exceeded. `0` means unbounded. + pub fn with_max_runs(max_runs: usize) -> Self { + Self { + state: Arc::new(Mutex::new(EventJournalState::default())), + max_runs, + } + } + /// Returns the number of observations stored for `run_id`. pub fn len(&self, run_id: &str) -> usize { - self.runs + self.state .lock() .expect("InMemoryEventJournal lock poisoned") + .streams .get(run_id) - .map(|v| v.len()) + .map(|s| s.entries.len()) .unwrap_or(0) } @@ -163,30 +180,93 @@ impl InMemoryEventJournal { pub fn is_empty(&self, run_id: &str) -> bool { self.len(run_id) == 0 } + + /// Returns the number of distinct `run_id` streams currently retained. + pub fn run_count(&self) -> usize { + self.state + .lock() + .expect("InMemoryEventJournal lock poisoned") + .streams + .len() + } } #[async_trait] impl HarnessEventJournal for InMemoryEventJournal { async fn append(&self, obs: AgentObservation) -> Result { - let mut runs = self - .runs + let mut state = self + .state .lock() .map_err(|e| poisoned("InMemoryEventJournal", e))?; - let entries = runs.entry(obs.run_id.as_str().to_string()).or_default(); - let offset = entries.len() as u64; - entries.push(obs); + let run_id = obs.run_id.as_str().to_string(); + if !state.streams.contains_key(&run_id) { + state.order.push_back(run_id.clone()); + // Evict the oldest run(s) once the cap is exceeded so a + // long-lived process journaling many runs doesn't grow this + // map without bound. `max_runs == 0` disables the cap. + if self.max_runs > 0 { + while state.order.len() > self.max_runs { + let Some(oldest) = state.order.pop_front() else { + break; + }; + if let Some(stream) = state.streams.remove(&oldest) { + // Remember where this run's numbering left off so a + // later append for the same run_id resumes from + // there instead of restarting at offset 0 — which + // would make `read_from` silently skip entries for a + // consumer resuming from a previously saved offset. + let next_offset = stream.base_offset + stream.entries.len() as u64; + state.evicted.insert(oldest.clone(), next_offset); + state.evicted_order.push_back(oldest); + while state.evicted_order.len() > self.max_runs { + let Some(stale) = state.evicted_order.pop_front() else { + break; + }; + state.evicted.remove(&stale); + } + } + } + } + } + let base_offset = state.evicted.remove(&run_id).unwrap_or(0); + let stream = state.streams.entry(run_id).or_insert_with(|| EventStream { + base_offset, + entries: Vec::new(), + }); + let offset = stream.base_offset + stream.entries.len() as u64; + stream.entries.push(obs); Ok(offset) } async fn read_from(&self, run_id: &str, offset: u64) -> Result> { - let runs = self - .runs + let state = self + .state .lock() .map_err(|e| poisoned("InMemoryEventJournal", e))?; - let Some(entries) = runs.get(run_id) else { - return Ok(Vec::new()); - }; - Ok(entries.iter().skip(offset as usize).cloned().collect()) + match state.streams.get(run_id) { + Some(stream) => { + if offset < stream.base_offset { + return Err(crate::error::TinyAgentsError::Validation(format!( + "run `{run_id}` requested offset {offset} but entries before \ + {} were evicted to bound memory; the caller's saved offset is stale", + stream.base_offset + ))); + } + let skip = (offset - stream.base_offset) as usize; + Ok(stream.entries.iter().skip(skip).cloned().collect()) + } + None => { + if let Some(&next_offset) = state.evicted.get(run_id) + && offset < next_offset + { + return Err(crate::error::TinyAgentsError::Validation(format!( + "run `{run_id}` was evicted to bound memory; the caller's \ + saved offset {offset} is stale" + ))); + } + Ok(Vec::new()) + } + } } } @@ -228,17 +308,39 @@ impl HarnessEventJournal for StoreEventJournal { // InMemoryStatusStore // --------------------------------------------------------------------------- +/// Returns `true` for statuses that are still in flight and must never be +/// evicted to make room for new runs. +fn is_active_status(status: &HarnessRunStatus) -> bool { + use crate::harness::ids::ExecutionStatus; + matches!( + status.status, + ExecutionStatus::Pending | ExecutionStatus::Running | ExecutionStatus::Interrupted + ) +} + impl InMemoryStatusStore { - /// Creates a new, empty in-memory status store. + /// Creates a new, empty in-memory status store with the default + /// run-retention cap ([`DEFAULT_STATUS_STORE_MAX_RUNS`]). pub fn new() -> Self { Self::default() } + /// Creates a new, empty in-memory status store that retains at most + /// `max_runs` distinct runs, evicting the oldest terminal run once + /// exceeded. `0` means unbounded. + pub fn with_max_runs(max_runs: usize) -> Self { + Self { + state: Arc::new(Mutex::new(StatusStoreState::default())), + max_runs, + } + } + /// Returns the number of distinct runs with a recorded status. pub fn len(&self) -> usize { - self.statuses + self.state .lock() .expect("InMemoryStatusStore lock poisoned") + .statuses .len() } @@ -251,28 +353,57 @@ impl InMemoryStatusStore { #[async_trait] impl HarnessStatusStore for InMemoryStatusStore { async fn put_status(&self, status: HarnessRunStatus) -> Result<()> { - let mut statuses = self - .statuses + let mut state = self + .state .lock() .map_err(|e| poisoned("InMemoryStatusStore", e))?; - statuses.insert(status.run_id.as_str().to_string(), status); + let run_id = status.run_id.as_str().to_string(); + if !state.statuses.contains_key(&run_id) { + state.order.push_back(run_id.clone()); + } + state.statuses.insert(run_id, status); + + // Evict the oldest terminal runs once the cap is exceeded so a + // supervisor tracking many short-lived runs doesn't grow this map + // without bound. Active runs are never evicted; if every retained + // run is still active we simply exceed the cap rather than drop + // in-flight state. `max_runs == 0` disables the cap. + if self.max_runs > 0 && state.statuses.len() > self.max_runs { + let mut requeue = Vec::new(); + while state.statuses.len() > self.max_runs { + let Some(candidate) = state.order.pop_front() else { + break; + }; + match state.statuses.get(&candidate) { + Some(s) if is_active_status(s) => requeue.push(candidate), + Some(_) => { + state.statuses.remove(&candidate); + } + None => {} + } + } + for id in requeue.into_iter().rev() { + state.order.push_front(id); + } + } Ok(()) } async fn get_status(&self, run_id: &str) -> Result> { - let statuses = self - .statuses + let state = self + .state .lock() .map_err(|e| poisoned("InMemoryStatusStore", e))?; - Ok(statuses.get(run_id).cloned()) + Ok(state.statuses.get(run_id).cloned()) } async fn list_by_thread(&self, thread_id: &str) -> Result> { - let statuses = self - .statuses + let state = self + .state .lock() .map_err(|e| poisoned("InMemoryStatusStore", e))?; - Ok(statuses + Ok(state + .statuses .values() .filter(|s| { s.thread_id @@ -284,11 +415,12 @@ impl HarnessStatusStore for InMemoryStatusStore { } async fn list_by_root(&self, root_run_id: &str) -> Result> { - let statuses = self - .statuses + let state = self + .state .lock() .map_err(|e| poisoned("InMemoryStatusStore", e))?; - Ok(statuses + Ok(state + .statuses .values() .filter(|s| s.root_run_id.as_str() == root_run_id) .cloned() @@ -296,21 +428,14 @@ impl HarnessStatusStore for InMemoryStatusStore { } async fn list_active(&self) -> Result> { - use crate::harness::ids::ExecutionStatus; - let statuses = self - .statuses + let state = self + .state .lock() .map_err(|e| poisoned("InMemoryStatusStore", e))?; - Ok(statuses + Ok(state + .statuses .values() - .filter(|s| { - matches!( - s.status, - ExecutionStatus::Pending - | ExecutionStatus::Running - | ExecutionStatus::Interrupted - ) - }) + .filter(|s| is_active_status(s)) .cloned() .collect()) } @@ -384,16 +509,21 @@ impl RedactingSink { impl EventListener for RedactingSink { fn on_event(&self, record: &EventRecord) { + // Fast path: with no secrets configured there is nothing to redact, so + // forward the original record unchanged. + if self.secrets.is_empty() { + self.inner.on_event(record); + return; + } // Serialize the event, mask secrets in every string field, and rebuild - // it. On any (de)serialization failure forward the original unchanged - // so observability is never silently dropped. + // it. This is a security boundary, so it must fail closed: if we cannot + // serialize, redact, and rebuild the event, we drop it rather than + // forward a record that may still contain unredacted secrets. let Ok(mut value) = serde_json::to_value(&record.event) else { - self.inner.on_event(record); return; }; redact_value(&mut value, &self.secrets, &self.mask); let Ok(event) = serde_json::from_value::(value) else { - self.inner.on_event(record); return; }; let redacted = EventRecord { @@ -439,11 +569,19 @@ impl JournalSink { /// lineage. `root_run_id` defaults to `run_id` for a top-level run; use /// [`Self::with_lineage`] to set a parent and a different root. pub fn new(journal: Arc, run_id: RunId) -> Self { + let worker = Arc::new(AppendWorker::spawn( + "journal-sink", + DEFAULT_DRAIN_CAPACITY, + move |obs: AgentObservation| { + let journal = Arc::clone(&journal); + async move { journal.append(obs).await.map(|_| ()) } + }, + )); Self { root_run_id: run_id.clone(), run_id, parent_run_id: None, - journal, + worker, } } @@ -453,6 +591,15 @@ impl JournalSink { self.root_run_id = root_run_id; self } + + /// Blocks until every observation submitted so far has been persisted. + /// + /// Persistence is otherwise asynchronous and best-effort; call this before + /// reading the journal back or shutting down to guarantee the durable log + /// has caught up with the events emitted so far. + pub fn flush(&self) { + self.worker.flush(); + } } impl EventListener for JournalSink { @@ -463,8 +610,8 @@ impl EventListener for JournalSink { self.parent_run_id.clone(), self.root_run_id.clone(), ); - // Best-effort durable append; never abort the run on a journal error. - let _ = futures::executor::block_on(self.journal.append(obs)); + // Hand off to the background drain; never block the run on I/O. + self.worker.submit(obs); } } @@ -478,10 +625,26 @@ impl JsonlSink { /// /// [`EventRecord`]: crate::harness::events::EventRecord pub fn new(store: JsonlAppendStore, stream: impl Into) -> Self { - Self { - store, - stream: stream.into(), - } + let stream = stream.into(); + let worker = Arc::new(AppendWorker::spawn( + "jsonl-sink", + DEFAULT_DRAIN_CAPACITY, + move |value: serde_json::Value| { + let store = store.clone(); + let stream = stream.clone(); + async move { store.append(&stream, value).await.map(|_| ()) } + }, + )); + Self { worker } + } + + /// Blocks until every record submitted so far has been appended. + /// + /// Persistence is otherwise asynchronous and best-effort; call this before + /// reading the stream back or shutting down to guarantee the durable log has + /// caught up with the events emitted so far. + pub fn flush(&self) { + self.worker.flush(); } } @@ -490,8 +653,8 @@ impl EventListener for JsonlSink { let Ok(value) = serde_json::to_value(record) else { return; }; - // Best-effort durable append; never abort the run on a store error. - let _ = futures::executor::block_on(self.store.append(&self.stream, value)); + // Hand off to the background drain; never block the run on I/O. + self.worker.submit(value); } } diff --git a/src/harness/observability/test.rs b/src/harness/observability/test.rs index ae378cf..dfc2b80 100644 --- a/src/harness/observability/test.rs +++ b/src/harness/observability/test.rs @@ -5,7 +5,8 @@ use std::sync::Arc; use crate::harness::events::{AgentEvent, EventListener, EventRecord, HarnessRunStatus, LimitKind}; -use crate::harness::ids::{CallId, ComponentId, EventId, RunId, ThreadId}; +use crate::harness::ids::{CallId, ComponentId, EventId, ExecutionStatus, RunId, ThreadId}; +use crate::harness::observability::AppendWorker; use crate::harness::observability::{ AgentLatencyMetrics, AgentObservation, FanOutSink, HarnessEventJournal, HarnessStatusStore, InMemoryEventJournal, InMemoryStatusStore, JournalSink, RedactingSink, StoreEventJournal, @@ -300,6 +301,9 @@ async fn journal_sink_persists_observations() { event: AgentEvent::StreamClosed, }); + // Persistence is asynchronous; block until the durable log catches up. + sink.flush(); + let stored = journal.read_from("run-sink", 0).await.unwrap(); assert_eq!(stored.len(), 2); assert_eq!(stored[0].event.kind(), "run.started"); @@ -309,6 +313,50 @@ async fn journal_sink_persists_observations() { assert_eq!(stored[1].root_run_id, RunId::new("run-sink")); } +#[tokio::test] +async fn append_worker_flush_persists_all_submissions_in_order() { + use std::sync::Mutex; + + let seen: Arc>> = Arc::new(Mutex::new(Vec::new())); + let sink = Arc::clone(&seen); + let worker = AppendWorker::spawn("test", 8, move |n: u64| { + let sink = Arc::clone(&sink); + async move { + sink.lock().unwrap().push(n); + Ok(()) + } + }); + + for n in 0..8 { + worker.submit(n); + } + // Before flush, persistence may still be in flight; after flush every + // submission is durably recorded, in submit order. + worker.flush(); + assert_eq!(*seen.lock().unwrap(), (0..8).collect::>()); +} + +#[tokio::test] +async fn append_worker_drops_and_counts_when_queue_is_full() { + // A slow backend cannot keep up with a burst; the bounded queue drops the + // overflow rather than blocking the caller, and the drops are counted (not + // silently discarded). + let worker = AppendWorker::spawn("test-slow", 1, move |_n: u64| async move { + std::thread::sleep(std::time::Duration::from_millis(20)); + Ok(()) + }); + + for n in 0..50 { + worker.submit(n); + } + worker.flush(); + assert!( + worker.dropped() > 0, + "a saturated bounded queue must drop and count overflow, got {}", + worker.dropped() + ); +} + /// Collects forwarded records for assertions. struct Collector { records: std::sync::Mutex>, @@ -389,6 +437,34 @@ fn redacting_sink_custom_mask_and_passthrough() { } } +#[test] +fn redacting_sink_empty_secrets_forwards_unchanged() { + // With no secrets configured the sink takes the fast path and forwards the + // original record untouched. + let collector = Arc::new(Collector::new()); + let sink = RedactingSink::new(collector.clone(), Vec::new()); + + sink.on_event(&EventRecord { + id: EventId::new("evt-0"), + offset: 7, + event: AgentEvent::RunFailed { + run_id: RunId::new("run-r"), + error: "nothing to redact here".to_string(), + }, + }); + + let events = collector.events(); + assert_eq!(events.len(), 1); + assert_eq!(events[0].id, EventId::new("evt-0")); + assert_eq!(events[0].offset, 7); + match &events[0].event { + AgentEvent::RunFailed { error, .. } => { + assert_eq!(error, "nothing to redact here"); + } + other => panic!("unexpected event: {other:?}"), + } +} + #[test] fn fan_out_sink_reaches_all_listeners() { let a = Arc::new(Collector::new()); @@ -412,3 +488,118 @@ fn fan_out_sink_reaches_all_listeners() { assert_eq!(c.events().len(), 1); assert_eq!(a.events()[0].event.kind(), "state.update"); } + +#[tokio::test] +async fn in_memory_journal_evicts_oldest_run_once_max_runs_exceeded() { + // Regression test: without a cap, a long-lived process journaling many + // runs grows the `run_id -> Vec` map without bound. With a cap of 2, the + // third distinct run must evict the oldest ("run-0"). + let journal = InMemoryEventJournal::with_max_runs(2); + + journal + .append(obs("run-0", 0, AgentEvent::StateUpdate)) + .await + .unwrap(); + journal + .append(obs("run-1", 0, AgentEvent::StateUpdate)) + .await + .unwrap(); + assert_eq!(journal.run_count(), 2); + + journal + .append(obs("run-2", 0, AgentEvent::StateUpdate)) + .await + .unwrap(); + + assert_eq!(journal.run_count(), 2, "cap must not be exceeded"); + assert!(journal.is_empty("run-0"), "oldest run must be evicted"); + assert!(!journal.is_empty("run-1")); + assert!(!journal.is_empty("run-2")); +} + +#[tokio::test] +async fn evicted_run_resuming_append_does_not_reset_offset_or_lose_reads() { + // Regression test: a run's stream can be evicted (FIFO, once max_runs is + // exceeded) while it is still actively appending. If the evicted run + // later appends again, the new stream must NOT restart numbering at 0 — + // otherwise a consumer that saved a durable offset from before eviction + // would have `read_from` silently skip entries instead of erroring. + let journal = InMemoryEventJournal::with_max_runs(2); + + // run-0 appends twice before being evicted by run-1 and run-2. + journal + .append(obs("run-0", 0, AgentEvent::StateUpdate)) + .await + .unwrap(); + journal + .append(obs("run-0", 1, AgentEvent::StateUpdate)) + .await + .unwrap(); + journal + .append(obs("run-1", 0, AgentEvent::StateUpdate)) + .await + .unwrap(); + journal + .append(obs("run-2", 0, AgentEvent::StateUpdate)) + .await + .unwrap(); + assert!(journal.is_empty("run-0"), "run-0 must have been evicted"); + + // run-0 resumes appending after eviction. The offset must continue from + // where it left off (2), not restart at 0. + let offset = journal + .append(obs("run-0", 2, AgentEvent::StateUpdate)) + .await + .unwrap(); + assert_eq!( + offset, 2, + "offset must continue from the evicted stream's length, not reset to 0" + ); + + // A consumer that saved offset 2 (from before eviction) must not have its + // new post-eviction entry silently skipped: it must see exactly the new + // entry. + let resumed = journal.read_from("run-0", 2).await.unwrap(); + assert_eq!(resumed.len(), 1, "post-eviction entry must not be dropped"); + + // A consumer whose saved offset falls within the evicted range must get + // a clear error rather than silently-truncated (or wrongly-offset) data. + let stale = journal.read_from("run-0", 0).await; + assert!( + stale.is_err(), + "reading a stale, evicted offset must fail closed, not silently skip" + ); +} + +#[tokio::test] +async fn status_store_evicts_oldest_terminal_run_but_never_active_ones() { + // Regression test: without a cap, a supervisor tracking many short-lived + // runs grows this map without bound. With a cap of 2: inserting a + // terminal run then an active one, then a third run, must evict the + // oldest *terminal* run rather than the active one. + let store = InMemoryStatusStore::with_max_runs(2); + + let mut terminal = HarnessRunStatus::new(RunId::new("run-a"), ComponentId::new("agent")); + terminal.status = ExecutionStatus::Completed; + store.put_status(terminal).await.unwrap(); + + // Default status is Pending (active). + let active = HarnessRunStatus::new(RunId::new("run-b"), ComponentId::new("agent")); + store.put_status(active).await.unwrap(); + assert_eq!(store.len(), 2); + + let mut run_c = HarnessRunStatus::new(RunId::new("run-c"), ComponentId::new("agent")); + run_c.status = ExecutionStatus::Completed; + store.put_status(run_c).await.unwrap(); + + assert_eq!(store.len(), 2, "cap must not be exceeded"); + assert!( + store.get_status("run-a").await.unwrap().is_none(), + "oldest terminal run must be evicted" + ); + assert!( + store.get_status("run-b").await.unwrap().is_some(), + "active run must never be evicted" + ); + assert!(store.get_status("run-c").await.unwrap().is_some()); +} diff --git a/src/harness/observability/types.rs b/src/harness/observability/types.rs index 007e41d..410bc57 100644 --- a/src/harness/observability/types.rs +++ b/src/harness/observability/types.rs @@ -11,7 +11,7 @@ //! implementations, sink logic, and tests live in the sibling `mod.rs` and //! `test.rs` files. -use std::collections::HashMap; +use std::collections::{HashMap, VecDeque}; use std::sync::{Arc, Mutex}; use async_trait::async_trait; @@ -19,8 +19,10 @@ use serde::{Deserialize, Serialize}; use crate::error::Result; use crate::harness::events::{AgentEvent, EventListener, HarnessRunStatus}; -use crate::harness::ids::{CallId, EventId, RunId}; -use crate::harness::store::{AppendStore, JsonlAppendStore}; +use crate::harness::ids::{CallId, EventId, RunId, now_ms}; +use crate::harness::store::AppendStore; + +use super::worker::AppendWorker; // --------------------------------------------------------------------------- // AgentObservation @@ -193,14 +195,63 @@ pub trait HarnessEventJournal: Send + Sync { } } +/// Default number of distinct `run_id` streams an [`InMemoryEventJournal`] +/// retains before evicting the oldest (by insertion order) to bound memory. +pub const DEFAULT_JOURNAL_MAX_RUNS: usize = 10_000; + +/// A single run's retained observations, plus the offset its first retained +/// entry corresponds to. +/// +/// `base_offset` is normally `0`, but once a run's stream has previously been +/// evicted and then receives another append, the new stream must continue +/// numbering offsets from where the evicted one left off rather than +/// restarting at `0` — otherwise a consumer resuming from a durable offset +/// would have entries silently skipped (see [`EventJournalState::evicted`]). +#[derive(Debug, Default)] +pub(crate) struct EventStream { + pub(crate) base_offset: u64, + pub(crate) entries: Vec, +} + +/// Inner state for [`InMemoryEventJournal`]: the per-run observation streams +/// plus insertion order so the oldest run can be evicted once `max_runs` is +/// exceeded. +#[derive(Debug, Default)] +pub(crate) struct EventJournalState { + pub(crate) streams: HashMap, + /// Oldest-first insertion order of run ids, used for FIFO eviction. + pub(crate) order: VecDeque, + /// Next offset to resume from for runs whose stream was evicted, so a + /// later append for the same `run_id` continues numbering instead of + /// restarting at `0` (which would corrupt any durable offset a consumer + /// has already saved). Bounded the same way `streams` is: entries are + /// dropped oldest-first once `evicted_order` exceeds `max_runs`. + pub(crate) evicted: HashMap, + pub(crate) evicted_order: VecDeque, +} + /// In-memory [`HarnessEventJournal`] 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)] +/// +/// Retains at most [`InMemoryEventJournal::max_runs`] distinct `run_id` +/// streams (default [`DEFAULT_JOURNAL_MAX_RUNS`]); once exceeded, the oldest +/// run (by first-append order) is evicted wholesale to keep memory bounded +/// across long-lived processes that journal many runs. +#[derive(Clone, Debug)] pub struct InMemoryEventJournal { - /// `run_id → ordered observations`. - pub(crate) runs: Arc>>>, + pub(crate) state: Arc>, + pub(crate) max_runs: usize, +} + +impl Default for InMemoryEventJournal { + fn default() -> Self { + Self { + state: Arc::new(Mutex::new(EventJournalState::default())), + max_runs: DEFAULT_JOURNAL_MAX_RUNS, + } + } } /// [`HarnessEventJournal`] backed by any [`AppendStore`]. @@ -255,13 +306,42 @@ pub trait HarnessStatusStore: Send + Sync { } } +/// Default number of distinct runs an [`InMemoryStatusStore`] retains before +/// evicting terminal runs (oldest first) to bound memory. +pub const DEFAULT_STATUS_STORE_MAX_RUNS: usize = 10_000; + +/// Inner state for [`InMemoryStatusStore`]: the `run_id → status` map plus +/// insertion order so terminal runs can be evicted oldest-first once +/// `max_runs` is exceeded. +#[derive(Debug, Default)] +pub(crate) struct StatusStoreState { + pub(crate) statuses: HashMap, + /// Oldest-first insertion order of run ids, used for eviction. + pub(crate) order: VecDeque, +} + /// In-memory [`HarnessStatusStore`] backed by a `run_id → status` map. /// /// Cheaply clonable through an inner [`Arc`]; clones share the same map. -#[derive(Clone, Debug, Default)] +/// +/// Retains at most [`InMemoryStatusStore::max_runs`] distinct runs (default +/// [`DEFAULT_STATUS_STORE_MAX_RUNS`]). Once exceeded, the oldest **terminal** +/// runs (anything not `Pending`/`Running`/`Interrupted`) are evicted first so +/// an in-flight run's status is never dropped out from under it; active runs +/// are only evicted once no terminal run remains to make room. +#[derive(Clone, Debug)] pub struct InMemoryStatusStore { - /// `run_id → latest status`. - pub(crate) statuses: Arc>>, + pub(crate) state: Arc>, + pub(crate) max_runs: usize, +} + +impl Default for InMemoryStatusStore { + fn default() -> Self { + Self { + state: Arc::new(Mutex::new(StatusStoreState::default())), + max_runs: DEFAULT_STATUS_STORE_MAX_RUNS, + } + } } // --------------------------------------------------------------------------- @@ -300,43 +380,37 @@ pub struct RedactingSink { /// a [`HarnessEventJournal`]. /// /// The sink is configured with the emitting run's lineage; each received -/// [`EventRecord`] is wrapped into an [`AgentObservation`] and appended. The -/// async append is bridged synchronously with `futures::executor::block_on`, -/// and append errors are swallowed so a failing journal never aborts the run. +/// [`EventRecord`] is wrapped into an [`AgentObservation`] and handed to a +/// background [`AppendWorker`] that persists it off the emitting thread. The +/// append is best-effort: see [`AppendWorker`] for the backpressure/drop and +/// error policy, and use [`JournalSink::flush`] to block until the durable log +/// has caught up. /// /// [`EventRecord`]: crate::harness::events::EventRecord #[derive(Clone)] pub struct JournalSink { - /// The journal observations are appended to. - pub(crate) journal: Arc, /// The run that owns events delivered to this sink. pub(crate) run_id: RunId, /// Parent run id stamped onto every observation. pub(crate) parent_run_id: Option, /// Root run id stamped onto every observation. pub(crate) root_run_id: RunId, + /// Background drain that persists observations without blocking the run. + pub(crate) worker: Arc>, } /// An [`EventListener`] that appends each [`EventRecord`] as a JSON line into a -/// [`JsonlAppendStore`] stream. +/// [`JsonlAppendStore`](crate::harness::store::JsonlAppendStore) stream. /// /// This is the lightweight durable sink: it persists the live record (id, -/// offset, event) under a fixed stream name. The async append is bridged -/// synchronously and errors are swallowed (best-effort). +/// offset, event) under a fixed stream name. Each record is handed to a +/// background [`AppendWorker`] that appends it off the emitting thread +/// (best-effort — see [`AppendWorker`] for the drop/error policy). Use +/// [`JsonlSink::flush`] to block until the durable log has caught up. /// /// [`EventRecord`]: crate::harness::events::EventRecord #[derive(Clone, Debug)] pub struct JsonlSink { - /// The JSONL append store records are written to. - pub(crate) store: JsonlAppendStore, - /// The stream name appended records land in. - pub(crate) stream: String, -} - -/// Returns the current time in Unix-epoch milliseconds, saturating at `0`. -pub(crate) fn now_ms() -> u64 { - std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .map(|d| d.as_millis() as u64) - .unwrap_or(0) + /// Background drain that appends records without blocking the run. + pub(crate) worker: Arc>, } diff --git a/src/harness/observability/worker.rs b/src/harness/observability/worker.rs new file mode 100644 index 0000000..2a27094 --- /dev/null +++ b/src/harness/observability/worker.rs @@ -0,0 +1,177 @@ +//! Background drain worker for durable event sinks. +//! +//! Durable sinks implement the *synchronous* `EventListener`/`GraphEventSink` +//! hooks but persist through *async* journal/store APIs. Bridging that boundary +//! inline with `futures::executor::block_on` blocked a tokio worker thread for +//! the whole append (a file write, in the worst case) on the run's critical +//! path, and risked a deadlock on a current-thread runtime. +//! +//! [`AppendWorker`] moves persistence off the emitting thread entirely: each +//! `submit` pushes the payload onto a **bounded** channel that a dedicated +//! background thread drains, awaiting the append on its own single-thread +//! runtime. +//! +//! # Backpressure & drop policy +//! The channel is bounded (see [`DEFAULT_DRAIN_CAPACITY`]). `submit` never +//! blocks the emitting thread: if the queue is full the observation is +//! **dropped** and counted in [`AppendWorker::dropped`], trading completeness +//! of the durable log for run latency (a slow or stuck backend must never stall +//! the run). Callers that need a lossless log can inspect the dropped count. +//! +//! # Error policy +//! Append errors are **not** silently discarded: the drain loop reports each +//! failure to stderr with the sink name. Persistence remains best-effort — an +//! error never propagates back into the run — but it is observable. +//! +//! # Ordering & durability boundary +//! A single drain thread preserves submit order. [`AppendWorker::flush`] blocks +//! until every payload queued so far has been persisted, and `Drop` flushes +//! before tearing the thread down, so no buffered observation is lost on a +//! graceful shutdown. + +use std::fmt; +use std::future::Future; +use std::sync::Arc; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::mpsc::{Sender, SyncSender, TrySendError, sync_channel}; +use std::thread::JoinHandle; + +use crate::error::Result; + +/// Default bounded-queue capacity for a durable drain worker. +/// +/// Sized so a transient backend stall buffers a healthy burst of events before +/// the drop policy engages, without letting the queue grow without bound. +pub(crate) const DEFAULT_DRAIN_CAPACITY: usize = 1024; + +/// Messages carried over the drain channel. +enum Msg { + /// A payload to persist. + Item(T), + /// A flush barrier: the drain loop acks once it reaches this marker, which + /// (by FIFO ordering) means every earlier `Item` has been persisted. + Flush(Sender<()>), +} + +/// A background worker that drains submitted payloads into an async append sink. +/// +/// See the [module docs](self) for the backpressure, error, and durability +/// semantics. +pub(crate) struct AppendWorker { + /// Bounded submit channel. Wrapped in `Option` only so [`Drop`] can drop the + /// sender before joining the drain thread; it is always `Some` otherwise. + tx: Option>>, + /// Count of payloads dropped because the queue was full (or disconnected). + dropped: Arc, + /// Handle to the drain thread, joined on drop. + handle: Option>, + /// Human-readable sink name used in error reports. + name: &'static str, +} + +impl AppendWorker { + /// Spawns a drain worker. + /// + /// `append` is invoked once per submitted payload on the drain thread's + /// runtime; it returns the async append future. `name` labels the drain + /// thread and error reports. + pub(crate) fn spawn(name: &'static str, capacity: usize, append: F) -> Self + where + F: Fn(T) -> Fut + Send + 'static, + Fut: Future>, + { + let (tx, rx) = sync_channel::>(capacity.max(1)); + let handle = std::thread::Builder::new() + .name(format!("tinyagents-{name}-drain")) + .spawn(move || { + let rt = match tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + { + Ok(rt) => rt, + // Without a runtime we cannot await appends; drain and drop. + Err(_) => { + while rx.recv().is_ok() {} + return; + } + }; + rt.block_on(async move { + while let Ok(msg) = rx.recv() { + match msg { + Msg::Item(item) => { + if let Err(e) = append(item).await { + eprintln!("tinyagents: {name} durable append failed: {e}"); + } + } + Msg::Flush(ack) => { + let _ = ack.send(()); + } + } + } + }); + }) + .expect("spawn durable-drain thread"); + Self { + tx: Some(tx), + dropped: Arc::new(AtomicU64::new(0)), + handle: Some(handle), + name, + } + } + + /// Queues `item` for durable persistence without blocking. + /// + /// If the bounded queue is full the item is dropped and counted (see + /// [`Self::dropped`]). + pub(crate) fn submit(&self, item: T) { + let Some(tx) = self.tx.as_ref() else { + return; + }; + match tx.try_send(Msg::Item(item)) { + Ok(()) => {} + Err(TrySendError::Full(_)) | Err(TrySendError::Disconnected(_)) => { + self.dropped.fetch_add(1, Ordering::Relaxed); + } + } + } + + /// Returns the number of payloads dropped because the queue was full. + #[cfg_attr(not(test), allow(dead_code))] + pub(crate) fn dropped(&self) -> u64 { + self.dropped.load(Ordering::Relaxed) + } + + /// Blocks until every payload submitted before this call has been persisted. + pub(crate) fn flush(&self) { + let Some(tx) = self.tx.as_ref() else { + return; + }; + let (ack_tx, ack_rx) = std::sync::mpsc::channel(); + // Blocking send (not `try_send`) so a momentarily full queue does not + // drop the flush barrier; the drain thread is actively consuming. + if tx.send(Msg::Flush(ack_tx)).is_ok() { + let _ = ack_rx.recv(); + } + } +} + +impl fmt::Debug for AppendWorker { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("AppendWorker") + .field("name", &self.name) + .field("dropped", &self.dropped.load(Ordering::Relaxed)) + .finish_non_exhaustive() + } +} + +impl Drop for AppendWorker { + fn drop(&mut self) { + // Persist anything still queued, then drop the sender so the drain loop + // observes a closed channel and exits, then join the thread. + self.flush(); + drop(self.tx.take()); + if let Some(handle) = self.handle.take() { + let _ = handle.join(); + } + } +} diff --git a/src/harness/providers/mock.rs b/src/harness/providers/mock.rs new file mode 100644 index 0000000..ef7e5fd --- /dev/null +++ b/src/harness/providers/mock.rs @@ -0,0 +1,282 @@ +//! [`MockModel`]: the deterministic, no-network provider used as the +//! default offline model — constructors, the `ChatModel` impl, and +//! token-estimation helpers. +//! +//! Split out of `providers/mod.rs`; see that module's doc comment for the +//! full provider overview. + +use super::*; + +// --------------------------------------------------------------------------- +// Token-estimation helpers +// --------------------------------------------------------------------------- + +/// Estimates the number of input tokens from a model request. +/// +/// Uses the heuristic of 1 token ≈ 4 characters of UTF-8 text. +fn estimate_input_tokens(request: &ModelRequest) -> u64 { + let total_chars: u64 = request.messages.iter().map(|m| m.text().len() as u64).sum(); + total_chars.div_ceil(4) +} + +/// Estimates output tokens from the response text. +/// +/// Uses the heuristic of 1 token ≈ 4 characters. Returns at least 1. +fn estimate_output_tokens(text: &str) -> u64 { + let chars = text.len() as u64; + std::cmp::max(1, chars.div_ceil(4)) +} + +// --------------------------------------------------------------------------- +// MockModel constructors +// --------------------------------------------------------------------------- + +impl MockModel { + /// Creates a `MockModel` that echoes the last user message back as the + /// assistant reply. + /// + /// If the request contains no user message, the reply is an empty string. + pub fn echo() -> Self { + Self { + behavior: MockBehavior::Echo, + inner: std::sync::Mutex::new(MockInner::default()), + } + } + + /// Creates a `MockModel` that always returns the same fixed assistant text. + pub fn constant(text: impl Into) -> Self { + Self { + behavior: MockBehavior::Constant(text.into()), + inner: std::sync::Mutex::new(MockInner::default()), + } + } + + /// Creates a `MockModel` that returns scripted responses in sequence. + /// + /// Responses are yielded one at a time in the order provided. When all + /// responses have been consumed the sequence **cycles back to the first + /// response**, so the model never errors simply due to exhaustion. + /// + /// # Panics + /// + /// Panics at *construction time* if `responses` is empty, because an empty + /// scripted model cannot produce any response. + pub fn with_responses(responses: Vec) -> Self { + assert!( + !responses.is_empty(), + "MockModel::with_responses: responses must not be empty" + ); + Self { + behavior: MockBehavior::Scripted(responses), + inner: std::sync::Mutex::new(MockInner::default()), + } + } + + /// Creates a `MockModel` that always issues one tool-call request. + /// + /// The returned [`ModelResponse`] has: + /// - An empty `content` block list (no text). + /// - One [`ToolCall`] in `message.tool_calls`. + /// - `finish_reason` set to `"tool_calls"`. + /// + /// `arguments` accepts anything that converts to a `serde_json::Value` + /// (e.g. `serde_json::json!({...})`, a pre-built `Value`, or `Value::Null`). + pub fn with_tool_call(name: impl Into, arguments: impl Into) -> Self { + Self { + behavior: MockBehavior::ToolCall { + name: name.into(), + arguments: arguments.into(), + }, + inner: std::sync::Mutex::new(MockInner::default()), + } + } + + /// Returns the total number of [`ChatModel::invoke`] calls made so far. + /// + /// `stream` calls that delegate to `invoke` also increment this counter. + pub fn call_count(&self) -> u64 { + self.inner + .lock() + .expect("MockModel inner state poisoned") + .call_count + } +} + +// --------------------------------------------------------------------------- +// ChatModel impl +// --------------------------------------------------------------------------- + +/// Returns the shared permissive [`ModelProfile`] advertised by [`MockModel`]. +/// +/// `MockModel` can satisfy any reasonable [`CapabilitySet`][crate::harness::model::CapabilitySet] +/// and supports every structured-output strategy, so its profile enables all +/// capabilities. +fn mock_profile() -> &'static ModelProfile { + static PROFILE: std::sync::OnceLock = std::sync::OnceLock::new(); + PROFILE.get_or_init(ModelProfile::permissive) +} + +#[async_trait] +impl ChatModel for MockModel { + /// Returns a permissive profile advertising every capability. + fn profile(&self) -> Option<&ModelProfile> { + Some(mock_profile()) + } + + /// Invokes the mock model and returns a deterministic response. + /// + /// Increments the internal call counter on every invocation. + async fn invoke(&self, _state: &State, request: ModelRequest) -> Result { + let call_id = { + let mut inner = self + .inner + .lock() + .map_err(|e| TinyAgentsError::Model(format!("MockModel lock poisoned: {e}")))?; + inner.call_count += 1; + inner.call_count + }; + + let msg_id = format!("mock-msg-{call_id}"); + let input_tokens = estimate_input_tokens(&request); + + let response = match &self.behavior { + MockBehavior::Echo => { + let text = request + .messages + .iter() + .rev() + .find_map(|m| { + if let Message::User(_) = m { + Some(m.text()) + } else { + None + } + }) + .unwrap_or_default(); + + let output_tokens = estimate_output_tokens(&text); + ModelResponse::assistant(text) + .with_usage(Usage::new(input_tokens, output_tokens)) + .with_finish_reason("stop") + } + + MockBehavior::Constant(text) => { + let output_tokens = estimate_output_tokens(text); + ModelResponse::assistant(text.clone()) + .with_usage(Usage::new(input_tokens, output_tokens)) + .with_finish_reason("stop") + } + + MockBehavior::Scripted(responses) => { + let index = { + let mut inner = self.inner.lock().map_err(|e| { + TinyAgentsError::Model(format!("MockModel lock poisoned: {e}")) + })?; + // We already incremented call_count above; derive index from + // call_count - 1 (0-based) cycling over the response list. + let idx = ((inner.call_count - 1) as usize) % responses.len(); + inner.scripted_index = idx; + idx + }; + responses[index].clone() + } + + MockBehavior::ToolCall { name, arguments } => { + let tool_call = ToolCall { + id: format!("mock-tool-{call_id}"), + name: name.clone(), + arguments: arguments.clone(), + }; + let usage = Usage::new(input_tokens, 5); + let message = AssistantMessage { + id: Some(msg_id.clone()), + content: Vec::new(), + tool_calls: vec![tool_call], + usage: Some(usage), + }; + ModelResponse { + message, + usage: Some(usage), + finish_reason: Some("tool_calls".to_string()), + raw: None, + resolved_model: None, + } + } + }; + + // Stamp the message id on text-based responses for traceability. + let mut response = response; + if response.message.id.is_none() { + response.message.id = Some(msg_id); + } + + Ok(response) + } + + /// Streams the model response as a real [`ModelStream`]. + /// + /// Internally calls [`invoke`][MockModel::invoke], then replays the response + /// as a [`ModelStreamItem::Started`] item, one or two + /// [`ModelStreamItem::MessageDelta`] items, and a terminal + /// [`ModelStreamItem::Completed`] carrying the full response. Text responses + /// are split into two roughly equal halves (by Unicode scalar value) so + /// streaming consumers observe multiple deltas without real network + /// infrastructure. Tool-call (or otherwise text-less) responses emit a + /// single empty text delta before completing. + async fn stream(&self, state: &State, request: ModelRequest) -> Result { + let response = self.invoke(state, request).await?; + let text = response.text(); + + let mut items = vec![ModelStreamItem::Started]; + + if text.is_empty() { + items.push(ModelStreamItem::MessageDelta(MessageDelta::default())); + } else { + // Split by Unicode scalar values so we never bisect a multi-byte + // char. + let chars: Vec = text.chars().collect(); + let mid = chars.len() / 2; + let first: String = chars[..mid].iter().collect(); + let second: String = chars[mid..].iter().collect(); + items.push(ModelStreamItem::MessageDelta(MessageDelta { + text: first, + reasoning: String::new(), + tool_call: None, + })); + items.push(ModelStreamItem::MessageDelta(MessageDelta { + text: second, + reasoning: String::new(), + tool_call: None, + })); + } + + items.push(ModelStreamItem::Completed(response)); + Ok(Box::pin(futures::stream::iter(items))) + } +} + +// --------------------------------------------------------------------------- +// ContentBlock helper used in tests +// --------------------------------------------------------------------------- + +impl MockModel { + /// Convenience: builds a plain-text [`ModelResponse`] — useful for + /// constructing scripted sequences in tests without importing the full + /// harness message path. + pub fn text_response(text: impl Into) -> ModelResponse { + let s = text.into(); + let output_tokens = estimate_output_tokens(&s); + ModelResponse { + message: AssistantMessage { + id: None, + content: vec![ContentBlock::Text(s)], + tool_calls: Vec::new(), + usage: Some(Usage::new(10, output_tokens)), + }, + usage: Some(Usage::new(10, output_tokens)), + finish_reason: Some("stop".to_string()), + raw: None, + resolved_model: None, + } + } +} diff --git a/src/harness/providers/mod.rs b/src/harness/providers/mod.rs index 7f30bbd..5280751 100644 --- a/src/harness/providers/mod.rs +++ b/src/harness/providers/mod.rs @@ -57,279 +57,7 @@ use crate::harness::model::{ use crate::harness::tool::ToolCall; use crate::harness::usage::Usage; -// --------------------------------------------------------------------------- -// Token-estimation helpers -// --------------------------------------------------------------------------- - -/// Estimates the number of input tokens from a model request. -/// -/// Uses the heuristic of 1 token ≈ 4 characters of UTF-8 text. -fn estimate_input_tokens(request: &ModelRequest) -> u64 { - let total_chars: u64 = request.messages.iter().map(|m| m.text().len() as u64).sum(); - total_chars.div_ceil(4) -} - -/// Estimates output tokens from the response text. -/// -/// Uses the heuristic of 1 token ≈ 4 characters. Returns at least 1. -fn estimate_output_tokens(text: &str) -> u64 { - let chars = text.len() as u64; - std::cmp::max(1, chars.div_ceil(4)) -} - -// --------------------------------------------------------------------------- -// MockModel constructors -// --------------------------------------------------------------------------- - -impl MockModel { - /// Creates a `MockModel` that echoes the last user message back as the - /// assistant reply. - /// - /// If the request contains no user message, the reply is an empty string. - pub fn echo() -> Self { - Self { - behavior: MockBehavior::Echo, - inner: std::sync::Mutex::new(MockInner::default()), - } - } - - /// Creates a `MockModel` that always returns the same fixed assistant text. - pub fn constant(text: impl Into) -> Self { - Self { - behavior: MockBehavior::Constant(text.into()), - inner: std::sync::Mutex::new(MockInner::default()), - } - } - - /// Creates a `MockModel` that returns scripted responses in sequence. - /// - /// Responses are yielded one at a time in the order provided. When all - /// responses have been consumed the sequence **cycles back to the first - /// response**, so the model never errors simply due to exhaustion. - /// - /// # Panics - /// - /// Panics at *construction time* if `responses` is empty, because an empty - /// scripted model cannot produce any response. - pub fn with_responses(responses: Vec) -> Self { - assert!( - !responses.is_empty(), - "MockModel::with_responses: responses must not be empty" - ); - Self { - behavior: MockBehavior::Scripted(responses), - inner: std::sync::Mutex::new(MockInner::default()), - } - } - - /// Creates a `MockModel` that always issues one tool-call request. - /// - /// The returned [`ModelResponse`] has: - /// - An empty `content` block list (no text). - /// - One [`ToolCall`] in `message.tool_calls`. - /// - `finish_reason` set to `"tool_calls"`. - /// - /// `arguments` accepts anything that converts to a `serde_json::Value` - /// (e.g. `serde_json::json!({...})`, a pre-built `Value`, or `Value::Null`). - pub fn with_tool_call(name: impl Into, arguments: impl Into) -> Self { - Self { - behavior: MockBehavior::ToolCall { - name: name.into(), - arguments: arguments.into(), - }, - inner: std::sync::Mutex::new(MockInner::default()), - } - } - - /// Returns the total number of [`ChatModel::invoke`] calls made so far. - /// - /// `stream` calls that delegate to `invoke` also increment this counter. - pub fn call_count(&self) -> u64 { - self.inner - .lock() - .expect("MockModel inner state poisoned") - .call_count - } -} - -// --------------------------------------------------------------------------- -// ChatModel impl -// --------------------------------------------------------------------------- - -/// Returns the shared permissive [`ModelProfile`] advertised by [`MockModel`]. -/// -/// `MockModel` can satisfy any reasonable [`CapabilitySet`][crate::harness::model::CapabilitySet] -/// and supports every structured-output strategy, so its profile enables all -/// capabilities. -fn mock_profile() -> &'static ModelProfile { - static PROFILE: std::sync::OnceLock = std::sync::OnceLock::new(); - PROFILE.get_or_init(ModelProfile::permissive) -} - -#[async_trait] -impl ChatModel for MockModel { - /// Returns a permissive profile advertising every capability. - fn profile(&self) -> Option<&ModelProfile> { - Some(mock_profile()) - } - - /// Invokes the mock model and returns a deterministic response. - /// - /// Increments the internal call counter on every invocation. - async fn invoke(&self, _state: &State, request: ModelRequest) -> Result { - let call_id = { - let mut inner = self - .inner - .lock() - .map_err(|e| TinyAgentsError::Model(format!("MockModel lock poisoned: {e}")))?; - inner.call_count += 1; - inner.call_count - }; - - let msg_id = format!("mock-msg-{call_id}"); - let input_tokens = estimate_input_tokens(&request); - - let response = match &self.behavior { - MockBehavior::Echo => { - let text = request - .messages - .iter() - .rev() - .find_map(|m| { - if let Message::User(_) = m { - Some(m.text()) - } else { - None - } - }) - .unwrap_or_default(); - - let output_tokens = estimate_output_tokens(&text); - ModelResponse::assistant(text) - .with_usage(Usage::new(input_tokens, output_tokens)) - .with_finish_reason("stop") - } - - MockBehavior::Constant(text) => { - let output_tokens = estimate_output_tokens(text); - ModelResponse::assistant(text.clone()) - .with_usage(Usage::new(input_tokens, output_tokens)) - .with_finish_reason("stop") - } - - MockBehavior::Scripted(responses) => { - let index = { - let mut inner = self.inner.lock().map_err(|e| { - TinyAgentsError::Model(format!("MockModel lock poisoned: {e}")) - })?; - // We already incremented call_count above; derive index from - // call_count - 1 (0-based) cycling over the response list. - let idx = ((inner.call_count - 1) as usize) % responses.len(); - inner.scripted_index = idx; - idx - }; - responses[index].clone() - } - - MockBehavior::ToolCall { name, arguments } => { - let tool_call = ToolCall { - id: format!("mock-tool-{call_id}"), - name: name.clone(), - arguments: arguments.clone(), - }; - let usage = Usage::new(input_tokens, 5); - let message = AssistantMessage { - id: Some(msg_id.clone()), - content: Vec::new(), - tool_calls: vec![tool_call], - usage: Some(usage), - }; - ModelResponse { - message, - usage: Some(usage), - finish_reason: Some("tool_calls".to_string()), - raw: None, - resolved_model: None, - } - } - }; - - // Stamp the message id on text-based responses for traceability. - let mut response = response; - if response.message.id.is_none() { - response.message.id = Some(msg_id); - } - - Ok(response) - } - - /// Streams the model response as a real [`ModelStream`]. - /// - /// Internally calls [`invoke`][MockModel::invoke], then replays the response - /// as a [`ModelStreamItem::Started`] item, one or two - /// [`ModelStreamItem::MessageDelta`] items, and a terminal - /// [`ModelStreamItem::Completed`] carrying the full response. Text responses - /// are split into two roughly equal halves (by Unicode scalar value) so - /// streaming consumers observe multiple deltas without real network - /// infrastructure. Tool-call (or otherwise text-less) responses emit a - /// single empty text delta before completing. - async fn stream(&self, state: &State, request: ModelRequest) -> Result { - let response = self.invoke(state, request).await?; - let text = response.text(); - - let mut items = vec![ModelStreamItem::Started]; - - if text.is_empty() { - items.push(ModelStreamItem::MessageDelta(MessageDelta::default())); - } else { - // Split by Unicode scalar values so we never bisect a multi-byte - // char. - let chars: Vec = text.chars().collect(); - let mid = chars.len() / 2; - let first: String = chars[..mid].iter().collect(); - let second: String = chars[mid..].iter().collect(); - items.push(ModelStreamItem::MessageDelta(MessageDelta { - text: first, - reasoning: String::new(), - tool_call: None, - })); - items.push(ModelStreamItem::MessageDelta(MessageDelta { - text: second, - reasoning: String::new(), - tool_call: None, - })); - } - - items.push(ModelStreamItem::Completed(response)); - Ok(Box::pin(futures::stream::iter(items))) - } -} - -// --------------------------------------------------------------------------- -// ContentBlock helper used in tests -// --------------------------------------------------------------------------- - -impl MockModel { - /// Convenience: builds a plain-text [`ModelResponse`] — useful for - /// constructing scripted sequences in tests without importing the full - /// harness message path. - pub fn text_response(text: impl Into) -> ModelResponse { - let s = text.into(); - let output_tokens = estimate_output_tokens(&s); - ModelResponse { - message: AssistantMessage { - id: None, - content: vec![ContentBlock::Text(s)], - tool_calls: Vec::new(), - usage: Some(Usage::new(10, output_tokens)), - }, - usage: Some(Usage::new(10, output_tokens)), - finish_reason: Some("stop".to_string()), - raw: None, - resolved_model: None, - } - } -} +mod mock; #[cfg(test)] mod test; diff --git a/src/harness/providers/openai/README.md b/src/harness/providers/openai/README.md new file mode 100644 index 0000000..a0e711b --- /dev/null +++ b/src/harness/providers/openai/README.md @@ -0,0 +1,107 @@ +# harness::providers::openai + +Real OpenAI Chat Completions provider (feature `openai`). This is one of the +concrete leaves the recursive runtime bottoms out in: a single `OpenAiModel` +backs hosted OpenAI *and* every OpenAI-compatible endpoint (Anthropic +compatibility, Ollama, DeepSeek, Groq, xAI, OpenRouter, Together, Mistral) via +the preset constructors, so the sub-agent / sub-graph layers above never need +to know which provider actually answered. + +## What it does + +`OpenAiModel` implements `ChatModel` against `POST {base_url}/chat/completions`. +It: + +1. Translates the provider-neutral `ModelRequest` into OpenAI's JSON wire + format (message content blocks, tool schemas, tool choice, response format, + image content parts). +2. Performs the HTTP call with a shared, reusable `reqwest::Client` (unary or + streamed). +3. Maps the response back into a `ModelResponse` with a fully populated + `AssistantMessage`, `ToolCall`s, `Usage`, and finish reason — or, for + streaming, a `ModelStream` of `ModelStreamItem`s. + +The wire (de)serialization shapes live in `types.rs`; `transport.rs`, +`convert.rs`, and `sse.rs` own the HTTP transport, request/response +translation, and SSE decoding respectively, keeping OpenAI-specific JSON out +of the rest of the harness. + +## Construction + +- `OpenAiModel::new(api_key)` — bare constructor, hosted OpenAI base URL and + `DEFAULT_MODEL` (`gpt-4.1-mini`). +- `.with_model(..)` / `.with_provider(..)` / `.with_base_url(..)` — builder + overrides. +- `OpenAiModel::from_env()` — reads `OPENAI_API_KEY` (required) and optional + `OPENAI_MODEL` / `OPENAI_BASE_URL`. +- `OpenAiModel::from_spec(spec, api_key)` / `from_spec_env(spec)` — build from a + `providers::ProviderSpec` (base URL, default model, provider id already + resolved). +- **Compatibility presets** — thin wrappers over `new` + `with_base_url` + + `with_model` for endpoints that speak the same Chat Completions wire format: + `compatible(base_url, model)` / `compatible_provider(..)` (arbitrary + endpoint), `deepseek`, `anthropic` (compat endpoint, not the native + Anthropic API), `groq`, `xai`, `openrouter`, `together`, `mistral`, `ollama`. + Override the preset's default model with `.with_model(..)`. + +Accessors: `.model()`, `.provider()`, `.base_url()`. + +## Model discovery + +`list_models()` calls `GET {base_url}/models` with the same credentials as +chat calls. Every OpenAI-compatible endpoint serves the same shape, so this +doubles as runtime model discovery for local/self-hosted providers (Ollama, +Together, Groq, OpenRouter, ...); returned ids can be fed straight into +`.with_model(..)`. + +## Streaming (SSE) + +Streaming responses are decoded by a small state machine (`SseState` / +`OpenAiStreamAcc` in `sse.rs`) built on `futures::stream::unfold`: + +- Bytes are accumulated across chunk boundaries and only lossily decoded once + a complete line is available, so a multi-byte UTF-8 character split across + two HTTP chunks is never corrupted into replacement characters. +- Index-less streamed tool-call fragments (some providers omit the tool-call + index on continuation chunks) are correlated by id rather than position. +- A trailing partial line without a final newline (providers that terminate + the last SSE event without a trailing newline) is still flushed. +- Mid-stream error payloads (`data: {"error": ...}`) surface as a stream error + instead of being silently swallowed. + +## Error handling + +Non-2xx responses are normalized through `parse_error_body` / `provider_error` +into a structured `ProviderError` (HTTP status, provider error code, and a +`retryable` flag derived from the status: 408/409/429/5xx are retryable, +everything else — including 401/400 — is not) and surfaced as +`TinyAgentsError::Provider`, so `harness::retry::is_retryable` can classify +retryability instead of retrying every provider failure indiscriminately. +Transport-level failures (connection errors, body-read failures) have no such +structure to preserve and surface as a plain `TinyAgentsError::Model` string +via `provider_failure_message`. Malformed JSON bodies surface as +`TinyAgentsError::Serialization`. + +## Operational constraints + +- `DEFAULT_CONNECT_TIMEOUT_SECS` (30s) bounds TCP connection establishment on + every call, including streaming, without capping response body time. +- `DEFAULT_REQUEST_TIMEOUT_SECS` (600s) is an overall timeout applied to + **unary** calls only when `ModelRequest::timeout_ms` is unset; streaming + calls get no overall cap by default — callers relying on a hard streaming + deadline must set `timeout_ms` explicitly or enforce one externally. +- o-series models (`o1`, `o3`, ...) require `max_completion_tokens` instead of + `max_tokens`; this is handled internally based on the model id. +- The client is reused across calls for connection pooling — construct one + `OpenAiModel` per logical provider/account rather than per request. + +## Files + +| File | Role | +| --- | --- | +| `mod.rs` | Module wiring: shared imports/constants and re-exports (`OpenAiModel`). | +| `transport.rs` | `OpenAiModel` construction, provider presets, request building, and the `ChatModel` impl (`invoke`/`stream`). | +| `convert.rs` | Request/response translation between harness types and the OpenAI wire format. | +| `sse.rs` | SSE stream parsing and incremental accumulation (`SseState`, `OpenAiStreamAcc`, `sse_next`). | +| `types.rs` | Wire (de)serialization shapes (`ModelListWire`, `ModelListing`, request/response bodies). | +| `test.rs` | Unit tests (SSE boundary decoding, tool-call correlation, error mapping, presets). | diff --git a/src/harness/providers/openai/convert.rs b/src/harness/providers/openai/convert.rs new file mode 100644 index 0000000..1dea259 --- /dev/null +++ b/src/harness/providers/openai/convert.rs @@ -0,0 +1,322 @@ +//! Request/response conversion between the provider-neutral harness +//! types and the OpenAI wire format (`translate_message`, +//! `parse_response`, usage conversion, reasoning-text extraction). +//! +//! Split out of `openai/mod.rs`; see that module's doc comment for the +//! full provider overview. + +use super::*; + +/// Translates one harness [`Message`] into an OpenAI wire message. +/// +/// User messages are rendered as OpenAI content-parts when they carry non-text +/// blocks (for example images), so image inputs are actually sent rather than +/// silently dropped. Blocks that have no faithful OpenAI representation return a +/// [`TinyAgentsError::Validation`] instead of being discarded. +pub(super) fn translate_message(message: &Message) -> Result { + let wire = match message { + Message::System(_) => ChatMessageWire { + role: "system".to_string(), + content: Some(MessageContentWire::Text(message.text())), + tool_calls: Vec::new(), + tool_call_id: None, + }, + Message::User(user) => ChatMessageWire { + role: "user".to_string(), + content: Some(translate_user_content(&user.content)?), + tool_calls: Vec::new(), + tool_call_id: None, + }, + Message::Assistant(assistant) => { + let text = message.text(); + // OpenAI accepts a null content for tool-call-only assistant turns. + let content = if text.is_empty() && !assistant.tool_calls.is_empty() { + None + } else { + Some(MessageContentWire::Text(text)) + }; + let tool_calls = assistant + .tool_calls + .iter() + .map(|call| { + Ok(ToolCallWire { + id: call.id.clone(), + kind: "function".to_string(), + function: FunctionCallWire { + name: call.name.clone(), + // OpenAI expects arguments as a JSON string. + arguments: serde_json::to_string(&call.arguments)?, + }, + }) + }) + .collect::>>()?; + ChatMessageWire { + role: "assistant".to_string(), + content, + tool_calls, + tool_call_id: None, + } + } + Message::Tool(tool) => ChatMessageWire { + role: "tool".to_string(), + content: Some(MessageContentWire::Text(message.text())), + tool_calls: Vec::new(), + tool_call_id: Some(tool.tool_call_id.clone()), + }, + }; + Ok(wire) +} + +/// Renders user-message content blocks into OpenAI message content. +/// +/// Text-only content collapses to a plain string (preserving the historical wire +/// shape). When an image block is present, content is emitted as OpenAI +/// content-parts so the image is actually sent. JSON blocks are serialized into +/// text parts. A [`ContentBlock::ProviderExtension`] has no faithful OpenAI +/// representation, so it fails closed with a validation error rather than being +/// silently dropped. +pub(super) fn translate_user_content(blocks: &[ContentBlock]) -> Result { + let has_image = blocks + .iter() + .any(|block| matches!(block, ContentBlock::Image(_))); + + if !has_image { + // No image: render as a single string, but still fail closed on blocks + // that cannot be represented. + let mut text = String::new(); + for block in blocks { + match block { + ContentBlock::Text(t) => text.push_str(t), + ContentBlock::Json(value) => text.push_str(&value.to_string()), + ContentBlock::Image(_) => unreachable!("guarded by has_image"), + ContentBlock::ProviderExtension(_) => { + return Err(unrepresentable_block_error()); + } + } + } + return Ok(MessageContentWire::Text(text)); + } + + let mut parts = Vec::with_capacity(blocks.len()); + for block in blocks { + match block { + ContentBlock::Text(t) => parts.push(ContentPartWire::Text { text: t.clone() }), + ContentBlock::Json(value) => parts.push(ContentPartWire::Text { + text: value.to_string(), + }), + ContentBlock::Image(image) => parts.push(ContentPartWire::ImageUrl { + image_url: ImageUrlWire { + url: image.url.clone(), + }, + }), + ContentBlock::ProviderExtension(_) => { + return Err(unrepresentable_block_error()); + } + } + } + Ok(MessageContentWire::Parts(parts)) +} + +/// Error returned when a content block cannot be represented in an OpenAI +/// request. Failing closed keeps the block from being silently dropped. +pub(super) fn unrepresentable_block_error() -> TinyAgentsError { + TinyAgentsError::Validation( + "OpenAI request cannot represent a provider-extension content block; \ + remove it or target the originating provider" + .to_string(), + ) +} + +/// Translates a [`ToolChoice`] into the OpenAI `tool_choice` JSON value. +pub(super) fn translate_tool_choice(choice: &ToolChoice) -> Value { + match choice { + ToolChoice::Auto => json!("auto"), + ToolChoice::None => json!("none"), + ToolChoice::Required => json!("required"), + ToolChoice::Tool(name) => json!({ + "type": "function", + "function": { "name": name } + }), + } +} + +/// Translates a [`ResponseFormat`] into the OpenAI `response_format` JSON value. +/// +/// Returns `None` for [`ResponseFormat::Text`] so the field is omitted entirely. +pub(super) fn translate_response_format(format: &ResponseFormat) -> Option { + match format { + ResponseFormat::Text => None, + ResponseFormat::JsonObject => Some(json!({ "type": "json_object" })), + // OpenAI supports native structured output, so `Auto` maps to a JSON + // schema request directly. (The agent loop normally resolves `Auto` + // before reaching the provider; this keeps direct calls correct too.) + ResponseFormat::JsonSchema { name, schema } | ResponseFormat::Auto { name, schema } => { + Some(json!({ + "type": "json_schema", + "json_schema": { + "name": name, + "schema": schema, + "strict": true, + } + })) + } + } +} + +/// Parses an OpenAI response body (already decoded into a [`Value`]) into a +/// provider-neutral [`ModelResponse`]. +/// +/// The first choice is used. The raw JSON is preserved in +/// [`ModelResponse::raw`]. +/// +/// # Errors +/// +/// Returns [`TinyAgentsError::Serialization`] if the value does not match the +/// expected response shape, or [`TinyAgentsError::Model`] when no choices are +/// present. +pub(super) fn parse_response(value: Value) -> Result { + let parsed: ChatCompletionResponse = serde_json::from_value(value.clone())?; + + let choice = parsed.choices.into_iter().next().ok_or_else(|| { + TinyAgentsError::Model("openai response contained no choices".to_string()) + })?; + + let mut content = Vec::new(); + if let Some(text) = choice.message.content.filter(|t| !t.is_empty()) { + content.push(ContentBlock::Text(text)); + } + + let tool_calls = choice + .message + .tool_calls + .into_iter() + .map(|call| { + Ok(ToolCall { + id: call.id.clone(), + name: call.function.name.clone(), + // Tool arguments arrive as a JSON string. Invalid JSON is a + // provider/model error, not an empty/default argument payload. + arguments: parse_tool_arguments( + "openai response", + &call.id, + &call.function.name, + &call.function.arguments, + )?, + }) + }) + .collect::>>()?; + + let usage = parsed.usage.map(convert_usage); + + let message = AssistantMessage { + id: parsed.id, + content, + tool_calls, + usage, + }; + + Ok(ModelResponse { + message, + usage, + finish_reason: choice.finish_reason, + raw: Some(value), + resolved_model: None, + }) +} + +/// Returns the effective call id for a streamed tool-call slot: the +/// provider-assigned id when present, or a stable `tool-{slot}` fallback keyed to +/// the slot's position so delta ids and the final call id always agree. +pub(super) fn tool_call_id(slot: usize, id: &str) -> String { + if id.is_empty() { + format!("tool-{slot}") + } else { + id.to_string() + } +} + +pub(super) fn parse_tool_arguments( + context: &str, + call_id: &str, + name: &str, + raw: &str, +) -> Result { + // Some OpenAI-compatible backends emit an empty arguments string for a + // zero-argument tool call. That is a well-formed "no arguments" payload, not + // malformed JSON, so map it to an empty object instead of failing the call. + if raw.trim().is_empty() { + return Ok(Value::Object(Map::new())); + } + serde_json::from_str(raw).map_err(|err| { + TinyAgentsError::Model(format!( + "{context} contained invalid JSON arguments for tool call `{call_id}` (`{name}`): {err}; raw arguments: {raw:?}" + )) + }) +} + +/// Converts an OpenAI [`UsageWire`] into the harness-neutral [`Usage`]. +pub(super) fn convert_usage(wire: UsageWire) -> Usage { + // OpenAI-compatible endpoints sometimes omit `total_tokens` entirely + // (deserializes to `0` via `#[serde(default)]`); fall back to + // `prompt + completion` so `total_tokens` is never a misleading zero for + // a call that clearly consumed tokens. + let total_tokens = if wire.total_tokens > 0 { + wire.total_tokens + } else { + wire.prompt_tokens + wire.completion_tokens + }; + Usage { + input_tokens: wire.prompt_tokens, + output_tokens: wire.completion_tokens, + total_tokens, + cache_read_tokens: wire + .prompt_tokens_details + .map(|d| d.cached_tokens) + .unwrap_or(0), + reasoning_tokens: wire + .completion_tokens_details + .map(|d| d.reasoning_tokens) + .unwrap_or(0), + ..Usage::default() + } +} + +/// Normalizes provider-specific reasoning/thinking payloads into text. +/// +/// OpenAI-compatible gateways do not agree on this field: some stream a plain +/// `reasoning_content` string, others use `reasoning`, and a few wrap text in +/// an object/array. Preserve renderable text when obvious and ignore opaque +/// shapes rather than failing an otherwise valid completion. +pub(super) fn reasoning_value_text(value: Value) -> Option { + match value { + Value::String(text) => (!text.is_empty()).then_some(text), + Value::Object(map) => ["text", "content", "summary"] + .into_iter() + .find_map(|key| map.get(key).and_then(Value::as_str)) + .filter(|text| !text.is_empty()) + .map(str::to_string), + Value::Array(values) => { + let text = values + .into_iter() + .filter_map(reasoning_value_text) + .collect::(); + (!text.is_empty()).then_some(text) + } + _ => None, + } +} + +/// Extracts the reasoning/thinking text from a streamed delta, accepting the +/// common OpenAI-compatible aliases. +pub(super) fn delta_reasoning_text(delta: &mut ChunkDeltaWire) -> String { + let mut text = String::new(); + for value in [delta.reasoning_content.take(), delta.reasoning.take()] + .into_iter() + .flatten() + { + if let Some(fragment) = reasoning_value_text(value) { + text.push_str(&fragment); + } + } + text +} diff --git a/src/harness/providers/openai/mod.rs b/src/harness/providers/openai/mod.rs index c101f18..68078a4 100644 --- a/src/harness/providers/openai/mod.rs +++ b/src/harness/providers/openai/mod.rs @@ -36,6 +36,7 @@ pub use types::*; use std::collections::VecDeque; use std::pin::Pin; +use std::time::Duration; use async_trait::async_trait; use futures::{Stream, StreamExt}; @@ -56,1022 +57,24 @@ use super::ProviderSpec; const DEFAULT_MODEL: &str = "gpt-4.1-mini"; /// Default OpenAI API base URL. const DEFAULT_BASE_URL: &str = "https://api.openai.com/v1"; - -/// A [`ChatModel`] backed by the hosted OpenAI Chat Completions API. -/// -/// Construct one with [`OpenAiModel::new`] (plus the `with_*` builders) or -/// [`OpenAiModel::from_env`]. The model holds a reusable [`reqwest::Client`] so -/// repeated calls share a connection pool. -pub struct OpenAiModel { - /// Shared HTTP client. - client: reqwest::Client, - /// API key sent as a `Bearer` token. - api_key: String, - /// Default model id used when a request does not override it. - model: String, - /// Provider family identifier used in profiles and normalized errors. - provider: String, - /// API base URL (no trailing slash); `/chat/completions` is appended. - base_url: String, - /// Capability profile derived from the default model id. - profile: ModelProfile, -} - -/// Derives a static [`ModelProfile`] for an OpenAI(-compatible) model id. -/// -/// All targets support tool calling, streaming (including tool-call chunks), -/// and JSON Schema response formats. Modern OpenAI-family models additionally -/// advertise native structured output and (for the o-series) reasoning output. -fn derive_profile(provider: &str, model: &str) -> ModelProfile { - let lower = model.to_ascii_lowercase(); - let native_structured = lower.contains("gpt-4o") - || lower.contains("gpt-4.1") - || lower.starts_with("o1") - || lower.starts_with("o3") - || lower.starts_with("o4"); - let reasoning = lower.starts_with("o1") || lower.starts_with("o3") || lower.starts_with("o4"); - ModelProfile { - provider: Some(provider.to_string()), - model: Some(model.to_string()), - status: ModelStatus::Stable, - modalities: Modalities { - image_in: true, - ..Modalities::default() - }, - tool_calling: true, - parallel_tool_calls: true, - streaming: true, - streaming_tool_chunks: true, - native_structured_output: native_structured, - json_schema: true, - reasoning, - ..ModelProfile::default() - } -} - -impl OpenAiModel { - /// Creates a model with the given API key, the default model - /// (`gpt-4.1-mini`), and the default base URL (`https://api.openai.com/v1`). - pub fn new(api_key: impl Into) -> Self { - Self { - client: reqwest::Client::new(), - api_key: api_key.into(), - model: DEFAULT_MODEL.to_string(), - provider: "openai".to_string(), - base_url: DEFAULT_BASE_URL.to_string(), - profile: derive_profile("openai", DEFAULT_MODEL), - } - } - - /// Overrides the default model id. - pub fn with_model(mut self, model: impl Into) -> Self { - self.model = model.into(); - self.profile = derive_profile(&self.provider, &self.model); - self - } - - /// Overrides the provider family id used in profiles and normalized errors. - pub fn with_provider(mut self, provider: impl Into) -> Self { - self.provider = provider.into(); - self.profile = derive_profile(&self.provider, &self.model); - self - } - - /// Overrides the API base URL. A trailing slash is trimmed so the joined - /// endpoint is always `{base_url}/chat/completions`. - pub fn with_base_url(mut self, base_url: impl Into) -> Self { - self.base_url = base_url.into().trim_end_matches('/').to_string(); - self - } - - /// Builds a model from environment variables. - /// - /// Reads `OPENAI_API_KEY` (required), `OPENAI_MODEL` (optional, defaults to - /// `gpt-4.1-mini`), and `OPENAI_BASE_URL` (optional, defaults to - /// `https://api.openai.com/v1`). - /// - /// # Errors - /// - /// Returns [`TinyAgentsError::Validation`] when `OPENAI_API_KEY` is missing - /// or empty. - pub fn from_env() -> Result { - let api_key = std::env::var("OPENAI_API_KEY") - .ok() - .filter(|k| !k.trim().is_empty()) - .ok_or_else(|| { - TinyAgentsError::Validation( - "OPENAI_API_KEY is not set; export it or add it to a .env file".to_string(), - ) - })?; - - let mut model = Self::new(api_key); - if let Ok(name) = std::env::var("OPENAI_MODEL") - && !name.trim().is_empty() - { - model = model.with_model(name); - } - if let Ok(url) = std::env::var("OPENAI_BASE_URL") - && !url.trim().is_empty() - { - model = model.with_base_url(url); - } - Ok(model) - } - - /// Builds an OpenAI-compatible model from a provider spec and explicit API - /// key. - pub fn from_spec(spec: ProviderSpec, api_key: impl Into) -> Result { - if spec.model.trim().is_empty() { - return Err(TinyAgentsError::Validation( - "provider spec model must not be empty".to_string(), - )); - } - if spec.base_url.trim().is_empty() { - return Err(TinyAgentsError::Validation( - "provider spec base_url must not be empty".to_string(), - )); - } - Ok(Self::compatible_provider( - spec.provider, - api_key, - spec.base_url, - spec.model, - )) - } - - /// Builds an OpenAI-compatible model from a provider spec, reading the API - /// key from the spec's environment variable when required. - pub fn from_spec_env(spec: ProviderSpec) -> Result { - let api_key = if spec.requires_api_key { - let env = spec.api_key_env.as_deref().ok_or_else(|| { - TinyAgentsError::Validation(format!( - "{} requires an api_key_env in ProviderSpec", - spec.provider - )) - })?; - std::env::var(env) - .ok() - .filter(|k| !k.trim().is_empty()) - .ok_or_else(|| { - TinyAgentsError::Validation(format!( - "{env} is not set; export it or provide an explicit API key" - )) - })? - } else { - "local".to_string() - }; - Self::from_spec(spec, api_key) - } - - /// Lists the models the provider advertises via `GET {base_url}/models`. - /// - /// This is a provider/account-level capability (independent of which model - /// this handle is bound to), so it uses the same credentials and base URL as - /// chat calls. Every OpenAI-compatible endpoint (Ollama, Together, Groq, - /// OpenRouter, …) serves the same shape, so this doubles as runtime model - /// discovery for local/self-hosted providers. Returned ids can be fed to - /// [`with_model`](Self::with_model). - /// - /// # Errors - /// - /// Returns [`TinyAgentsError::Model`] on transport failure or a non-2xx - /// status, and [`TinyAgentsError::Serialization`] when the body cannot be - /// decoded. - pub async fn list_models(&self) -> Result> { - let url = format!("{}/models", self.base_url); - - let response = self - .client - .get(&url) - .header("Authorization", format!("Bearer {}", self.api_key)) - .send() - .await - .map_err(|e| { - let error = - self.provider_error(format!("request to {url} failed: {e}"), None, None, None); - TinyAgentsError::Model(self.provider_failure_message(&error)) - })?; - - let status = response.status(); - let text = response.text().await.map_err(|e| { - TinyAgentsError::Model(format!("openai response body read failed: {e}")) - })?; - - if !status.is_success() { - let error = self.parse_error_body(status.as_u16(), &text); - return Err(TinyAgentsError::Model( - self.provider_failure_message(&error), - )); - } - - let listing: ModelListWire = serde_json::from_str(&text)?; - Ok(listing.data) - } - - // ----------------------------------------------------------------------- - // OpenAI-compatible provider presets - // - // DeepSeek, Groq, xAI, OpenRouter, Together, Mistral, Ollama, and - // Anthropic's compatibility endpoint all accept the same Chat Completions - // wire format, so the same [`OpenAiModel`] talks to all of them — only the - // base URL and default model differ. Each preset is a thin wrapper over - // [`OpenAiModel::new`] + [`with_base_url`][Self::with_base_url] + - // [`with_model`][Self::with_model]; override the model with `with_model`. - // ----------------------------------------------------------------------- - - /// Points at an arbitrary OpenAI-compatible endpoint with an explicit base - /// URL and default model. - /// - /// Use this for any provider that implements the Chat Completions API but is - /// not covered by a named preset below. - pub fn compatible( - api_key: impl Into, - base_url: impl Into, - model: impl Into, - ) -> Self { - Self::new(api_key).with_base_url(base_url).with_model(model) - } - - /// Points at an arbitrary OpenAI-compatible endpoint with an explicit - /// provider id, base URL, and default model. - pub fn compatible_provider( - provider: impl Into, - api_key: impl Into, - base_url: impl Into, - model: impl Into, - ) -> Self { - Self::new(api_key) - .with_provider(provider) - .with_base_url(base_url) - .with_model(model) - } - - /// DeepSeek (`https://api.deepseek.com/v1`), default model `deepseek-chat`. - pub fn deepseek(api_key: impl Into) -> Self { - Self::compatible_provider( - "deepseek", - api_key, - "https://api.deepseek.com/v1", - "deepseek-chat", - ) - } - - /// Anthropic's OpenAI-compatible endpoint (`https://api.anthropic.com/v1`), - /// default model `claude-3-5-sonnet-latest`. - pub fn anthropic(api_key: impl Into) -> Self { - Self::compatible_provider( - "anthropic", - api_key, - "https://api.anthropic.com/v1", - "claude-3-5-sonnet-latest", - ) - } - - /// Groq (`https://api.groq.com/openai/v1`), default model - /// `llama-3.3-70b-versatile`. - pub fn groq(api_key: impl Into) -> Self { - Self::compatible_provider( - "groq", - api_key, - "https://api.groq.com/openai/v1", - "llama-3.3-70b-versatile", - ) - } - - /// xAI (`https://api.x.ai/v1`), default model `grok-2-latest`. - pub fn xai(api_key: impl Into) -> Self { - Self::compatible_provider("xai", api_key, "https://api.x.ai/v1", "grok-2-latest") - } - - /// OpenRouter (`https://openrouter.ai/api/v1`), default model - /// `openai/gpt-4o-mini`. - pub fn openrouter(api_key: impl Into) -> Self { - Self::compatible_provider( - "openrouter", - api_key, - "https://openrouter.ai/api/v1", - "openai/gpt-4o-mini", - ) - } - - /// Together AI (`https://api.together.xyz/v1`), default model - /// `meta-llama/Llama-3.3-70B-Instruct-Turbo`. - pub fn together(api_key: impl Into) -> Self { - Self::compatible_provider( - "together", - api_key, - "https://api.together.xyz/v1", - "meta-llama/Llama-3.3-70B-Instruct-Turbo", - ) - } - - /// Mistral (`https://api.mistral.ai/v1`), default model - /// `mistral-small-latest`. - pub fn mistral(api_key: impl Into) -> Self { - Self::compatible_provider( - "mistral", - api_key, - "https://api.mistral.ai/v1", - "mistral-small-latest", - ) - } - - /// A local Ollama server (`http://localhost:11434/v1`), default model - /// `llama3.2`. Ollama ignores the API key, so a placeholder is used. - pub fn ollama() -> Self { - Self::compatible_provider("ollama", "ollama", "http://localhost:11434/v1", "llama3.2") - } - - /// Returns the default model id this instance will request. - pub fn model(&self) -> &str { - &self.model - } - - /// Returns the configured provider family id. - pub fn provider(&self) -> &str { - &self.provider - } - - /// Returns the configured API base URL. - pub fn base_url(&self) -> &str { - &self.base_url - } - - /// Translates a provider-neutral [`ModelRequest`] into the OpenAI wire - /// request body. The per-request `model` override wins over the instance - /// default. - fn translate_request(&self, request: &ModelRequest) -> Result { - let messages = request - .messages - .iter() - .map(translate_message) - .collect::>>()?; - - let tools: Vec = request - .tools - .iter() - .map(|schema| ToolWire { - kind: "function".to_string(), - function: FunctionSchemaWire { - name: schema.name.clone(), - description: schema.description.clone(), - parameters: schema.parameters.clone(), - }, - }) - .collect(); - - // tool_choice is only meaningful when tools are declared. - let tool_choice = if tools.is_empty() { - None - } else { - Some(translate_tool_choice(&request.tool_choice)) - }; - - let response_format = request - .response_format - .as_ref() - .and_then(translate_response_format); - - Ok(ChatCompletionRequest { - model: request.model.clone().unwrap_or_else(|| self.model.clone()), - messages, - tools, - tool_choice, - response_format, - temperature: request.temperature, - top_p: request.top_p, - max_tokens: request.max_tokens, - stop: request.stop_sequences.clone(), - seed: request.seed, - stream: false, - stream_options: None, - extra: provider_extra_options(&request.provider_options)?, - }) - } - - fn provider_error( - &self, - message: impl Into, - status: Option, - code: Option, - raw: Option, - ) -> ProviderError { - let retryable = status.is_some_and(|s| s == 408 || s == 409 || s == 429 || s >= 500); - ProviderError { - provider: self.provider.clone(), - model: Some(self.model.clone()), - status, - code, - message: message.into(), - retryable, - raw, - } - } - - fn provider_failure_message(&self, error: &ProviderError) -> String { - format!( - "{} returned{}{}: {}", - error.provider, - error - .status - .map(|status| format!(" HTTP {status}")) - .unwrap_or_default(), - error - .code - .as_deref() - .map(|code| format!(" ({code})")) - .unwrap_or_default(), - error.message - ) - } - - fn parse_error_body(&self, status: u16, text: &str) -> ProviderError { - let raw = serde_json::from_str::(text).ok(); - let error_obj = raw.as_ref().and_then(|value| value.get("error")); - let message = error_obj - .and_then(|error| error.get("message")) - .and_then(Value::as_str) - .or_else(|| { - raw.as_ref() - .and_then(|value| value.get("message")) - .and_then(Value::as_str) - }) - .filter(|message| !message.trim().is_empty()) - .unwrap_or(text) - .to_string(); - let code = error_obj - .and_then(|error| error.get("code").or_else(|| error.get("type"))) - .and_then(Value::as_str) - .map(str::to_string); - self.provider_error(message, Some(status), code, raw) - } -} - -/// Returns provider-specific top-level fields to flatten into the request body. -/// -/// Core OpenAI-compatible fields are intentionally reserved so normalized -/// TinyAgents fields remain the source of truth. Callers that need local-model -/// controls should use distinct provider fields such as Ollama's `options`. -fn provider_extra_options(options: &Value) -> Result> { - if options.is_null() { - return Ok(Map::new()); - } - let Some(object) = options.as_object() else { - return Err(TinyAgentsError::Validation( - "provider_options for OpenAI-compatible providers must be a JSON object".to_string(), - )); - }; - - const RESERVED: &[&str] = &[ - "model", - "messages", - "tools", - "tool_choice", - "response_format", - "temperature", - "top_p", - "max_tokens", - "stop", - "seed", - "stream", - "stream_options", - ]; - - Ok(object - .iter() - .filter(|(key, _)| !RESERVED.contains(&key.as_str())) - .map(|(key, value)| (key.clone(), value.clone())) - .collect()) -} - -/// Translates one harness [`Message`] into an OpenAI wire message. -fn translate_message(message: &Message) -> Result { - let wire = match message { - Message::System(_) => ChatMessageWire { - role: "system".to_string(), - content: Some(message.text()), - tool_calls: Vec::new(), - tool_call_id: None, - }, - Message::User(_) => ChatMessageWire { - role: "user".to_string(), - content: Some(message.text()), - tool_calls: Vec::new(), - tool_call_id: None, - }, - Message::Assistant(assistant) => { - let text = message.text(); - // OpenAI accepts a null content for tool-call-only assistant turns. - let content = if text.is_empty() && !assistant.tool_calls.is_empty() { - None - } else { - Some(text) - }; - let tool_calls = assistant - .tool_calls - .iter() - .map(|call| { - Ok(ToolCallWire { - id: call.id.clone(), - kind: "function".to_string(), - function: FunctionCallWire { - name: call.name.clone(), - // OpenAI expects arguments as a JSON string. - arguments: serde_json::to_string(&call.arguments)?, - }, - }) - }) - .collect::>>()?; - ChatMessageWire { - role: "assistant".to_string(), - content, - tool_calls, - tool_call_id: None, - } - } - Message::Tool(tool) => ChatMessageWire { - role: "tool".to_string(), - content: Some(message.text()), - tool_calls: Vec::new(), - tool_call_id: Some(tool.tool_call_id.clone()), - }, - }; - Ok(wire) -} - -/// Translates a [`ToolChoice`] into the OpenAI `tool_choice` JSON value. -fn translate_tool_choice(choice: &ToolChoice) -> Value { - match choice { - ToolChoice::Auto => json!("auto"), - ToolChoice::None => json!("none"), - ToolChoice::Required => json!("required"), - ToolChoice::Tool(name) => json!({ - "type": "function", - "function": { "name": name } - }), - } -} - -/// Translates a [`ResponseFormat`] into the OpenAI `response_format` JSON value. -/// -/// Returns `None` for [`ResponseFormat::Text`] so the field is omitted entirely. -fn translate_response_format(format: &ResponseFormat) -> Option { - match format { - ResponseFormat::Text => None, - ResponseFormat::JsonObject => Some(json!({ "type": "json_object" })), - // OpenAI supports native structured output, so `Auto` maps to a JSON - // schema request directly. (The agent loop normally resolves `Auto` - // before reaching the provider; this keeps direct calls correct too.) - ResponseFormat::JsonSchema { name, schema } | ResponseFormat::Auto { name, schema } => { - Some(json!({ - "type": "json_schema", - "json_schema": { - "name": name, - "schema": schema, - "strict": true, - } - })) - } - } -} - -/// Parses an OpenAI response body (already decoded into a [`Value`]) into a -/// provider-neutral [`ModelResponse`]. -/// -/// The first choice is used. The raw JSON is preserved in -/// [`ModelResponse::raw`]. -/// -/// # Errors -/// -/// Returns [`TinyAgentsError::Serialization`] if the value does not match the -/// expected response shape, or [`TinyAgentsError::Model`] when no choices are -/// present. -fn parse_response(value: Value) -> Result { - let parsed: ChatCompletionResponse = serde_json::from_value(value.clone())?; - - let choice = parsed.choices.into_iter().next().ok_or_else(|| { - TinyAgentsError::Model("openai response contained no choices".to_string()) - })?; - - let mut content = Vec::new(); - if let Some(text) = choice.message.content.filter(|t| !t.is_empty()) { - content.push(ContentBlock::Text(text)); - } - - let tool_calls = choice - .message - .tool_calls - .into_iter() - .map(|call| { - Ok(ToolCall { - id: call.id.clone(), - name: call.function.name.clone(), - // Tool arguments arrive as a JSON string. Invalid JSON is a - // provider/model error, not an empty/default argument payload. - arguments: parse_tool_arguments( - "openai response", - &call.id, - &call.function.name, - &call.function.arguments, - )?, - }) - }) - .collect::>>()?; - - let usage = parsed.usage.map(convert_usage); - - let message = AssistantMessage { - id: parsed.id, - content, - tool_calls, - usage, - }; - - Ok(ModelResponse { - message, - usage, - finish_reason: choice.finish_reason, - raw: Some(value), - resolved_model: None, - }) -} - -fn parse_tool_arguments(context: &str, call_id: &str, name: &str, raw: &str) -> Result { - serde_json::from_str(raw).map_err(|err| { - TinyAgentsError::Model(format!( - "{context} contained invalid JSON arguments for tool call `{call_id}` (`{name}`): {err}; raw arguments: {raw:?}" - )) - }) -} - -/// Converts an OpenAI [`UsageWire`] into the harness-neutral [`Usage`]. -fn convert_usage(wire: UsageWire) -> Usage { - Usage { - input_tokens: wire.prompt_tokens, - output_tokens: wire.completion_tokens, - total_tokens: wire.total_tokens, - cache_read_tokens: wire - .prompt_tokens_details - .map(|d| d.cached_tokens) - .unwrap_or(0), - ..Usage::default() - } -} - -// --------------------------------------------------------------------------- -// Streaming (SSE) machinery -// --------------------------------------------------------------------------- - -/// In-progress reconstruction of a single tool call across streamed fragments. -#[derive(Clone, Debug, Default)] -struct ToolCallBuild { - /// Provider-assigned call id (filled from the first fragment carrying it). - id: String, - /// Function name (filled from the first fragment carrying it). - name: String, - /// Concatenated stringified-JSON argument fragments. - args: String, -} - -/// Provider-side accumulator that rebuilds the authoritative [`ModelResponse`] -/// from streamed chunks. Distinct from the generic -/// [`StreamAccumulator`][crate::harness::model::StreamAccumulator]: it tracks -/// tool-call names and ids (which the neutral deltas omit) so the terminal -/// [`ModelStreamItem::Completed`] carries a faithful response. -#[derive(Clone, Debug, Default)] -struct OpenAiStreamAcc { - id: Option, - text: String, - tool_calls: Vec, - usage: Option, - finish_reason: Option, -} - -impl OpenAiStreamAcc { - /// Folds one parsed chunk into the accumulator and pushes the corresponding - /// neutral [`ModelStreamItem`]s onto `pending`. - fn ingest(&mut self, chunk: ChatCompletionChunk, pending: &mut VecDeque) { - if let Some(id) = chunk.id - && self.id.is_none() - { - self.id = Some(id); - } - if let Some(usage_wire) = chunk.usage { - let usage = convert_usage(usage_wire); - self.usage = Some(usage); - pending.push_back(ModelStreamItem::UsageDelta(usage)); - } - for choice in chunk.choices { - if let Some(reason) = choice.finish_reason { - self.finish_reason = Some(reason); - } - if let Some(content) = choice.delta.content.filter(|c| !c.is_empty()) { - self.text.push_str(&content); - pending.push_back(ModelStreamItem::MessageDelta(MessageDelta { - text: content, - reasoning: String::new(), - tool_call: None, - })); - } - for fragment in choice.delta.tool_calls { - let idx = fragment.index as usize; - while self.tool_calls.len() <= idx { - self.tool_calls.push(ToolCallBuild::default()); - } - let slot = &mut self.tool_calls[idx]; - if let Some(id) = fragment.id.filter(|id| !id.is_empty()) { - slot.id = id; - } - if let Some(function) = fragment.function { - if let Some(name) = function.name.filter(|n| !n.is_empty()) { - slot.name = name; - } - if let Some(args) = function.arguments.filter(|a| !a.is_empty()) { - slot.args.push_str(&args); - let call_id = if slot.id.is_empty() { - format!("tool-{idx}") - } else { - slot.id.clone() - }; - pending.push_back(ModelStreamItem::ToolCallDelta(ToolDelta { - call_id, - content: args, - })); - } - } - } - } - } - - /// Consumes the accumulator into the final, merged [`ModelResponse`]. - fn into_response(self) -> Result { - let mut content = Vec::new(); - if !self.text.is_empty() { - content.push(ContentBlock::Text(self.text)); - } - let tool_calls = self - .tool_calls - .into_iter() - .filter(|b| !b.name.is_empty() || !b.args.is_empty()) - .enumerate() - .map(|(idx, b)| { - let id = if b.id.is_empty() { - format!("tool-{idx}") - } else { - b.id.clone() - }; - Ok(ToolCall { - id: id.clone(), - name: b.name.clone(), - arguments: parse_tool_arguments("openai stream", &id, &b.name, &b.args)?, - }) - }) - .collect::>>()?; - let message = AssistantMessage { - id: self.id, - content, - tool_calls, - usage: self.usage, - }; - Ok(ModelResponse { - message, - usage: self.usage, - finish_reason: self.finish_reason, - raw: None, - resolved_model: None, - }) - } -} - -/// Mutable driver state threaded through [`futures::stream::unfold`] while -/// parsing the SSE byte stream into [`ModelStreamItem`]s. -struct SseState { - /// Raw response byte chunks (errors already mapped onto the crate error). - bytes: Pin>> + Send>>, - /// Bytes received but not yet split into complete lines. - buf: String, - /// Parsed items waiting to be yielded, in order. - pending: VecDeque, - /// Provider-side response reconstruction. - acc: OpenAiStreamAcc, - /// Provider family id used in normalized stream failures. - provider: String, - /// Provider model id used in normalized stream failures. - model: String, - /// Whether the leading [`ModelStreamItem::Started`] has been emitted. - started: bool, - /// Whether the byte stream ended or `[DONE]` was seen. - finished: bool, - /// Whether the terminal [`ModelStreamItem::Completed`]/[`ModelStreamItem::Failed`] - /// has been emitted. - terminal_emitted: bool, -} - -impl SseState { - /// Splits buffered bytes into complete lines and folds each SSE `data:` - /// payload into the accumulator. The trailing partial line (if any) is kept - /// for the next chunk. - fn drain_lines(&mut self) { - while let Some(pos) = self.buf.find('\n') { - let line: String = self.buf.drain(..=pos).collect(); - let line = line.trim(); - if line.is_empty() { - continue; - } - let Some(rest) = line.strip_prefix("data:") else { - continue; - }; - let payload = rest.trim(); - if payload == "[DONE]" { - self.finished = true; - continue; - } - // Ignore keepalives / unparseable lines rather than failing the run. - if let Ok(chunk) = serde_json::from_str::(payload) { - let mut pending = std::mem::take(&mut self.pending); - self.acc.ingest(chunk, &mut pending); - self.pending = pending; - } - } - } -} - -/// Advances the SSE [`SseState`] by one item for [`futures::stream::unfold`]. -async fn sse_next(mut state: SseState) -> Option<(ModelStreamItem, SseState)> { - loop { - if let Some(item) = state.pending.pop_front() { - return Some((item, state)); - } - if !state.started { - state.started = true; - return Some((ModelStreamItem::Started, state)); - } - if state.finished { - if state.terminal_emitted { - return None; - } - state.terminal_emitted = true; - return match std::mem::take(&mut state.acc).into_response() { - Ok(response) => Some((ModelStreamItem::Completed(response), state)), - Err(error) => { - let provider_error = ProviderError { - provider: state.provider.clone(), - model: Some(state.model.clone()), - code: Some("invalid_tool_arguments".to_string()), - message: error.to_string(), - retryable: false, - ..ProviderError::default() - }; - Some((ModelStreamItem::ProviderFailed(provider_error), state)) - } - }; - } - match state.bytes.next().await { - Some(Ok(chunk)) => { - state.buf.push_str(&String::from_utf8_lossy(&chunk)); - state.drain_lines(); - } - Some(Err(error)) => { - state.finished = true; - state.terminal_emitted = true; - let provider_error = ProviderError { - provider: state.provider.clone(), - model: Some(state.model.clone()), - message: error.to_string(), - retryable: true, - ..ProviderError::default() - }; - return Some((ModelStreamItem::ProviderFailed(provider_error), state)); - } - None => { - state.finished = true; - } - } - } -} - -#[async_trait] -impl ChatModel for OpenAiModel { - /// Returns the capability profile derived from the configured model id. - fn profile(&self) -> Option<&ModelProfile> { - Some(&self.profile) - } - - /// Invokes the OpenAI Chat Completions endpoint and maps the response into a - /// [`ModelResponse`]. - /// - /// # Errors - /// - /// Returns [`TinyAgentsError::Model`] on transport failure or a non-2xx - /// status (the message includes the status code and response body), and - /// [`TinyAgentsError::Serialization`] when the response cannot be decoded. - async fn invoke(&self, _state: &State, request: ModelRequest) -> Result { - let body = self.translate_request(&request)?; - let url = format!("{}/chat/completions", self.base_url); - - let response = self - .client - .post(&url) - .header("Authorization", format!("Bearer {}", self.api_key)) - .json(&body) - .send() - .await - .map_err(|e| { - let error = - self.provider_error(format!("request to {url} failed: {e}"), None, None, None); - TinyAgentsError::Model(self.provider_failure_message(&error)) - })?; - - let status = response.status(); - let text = response.text().await.map_err(|e| { - TinyAgentsError::Model(format!("openai response body read failed: {e}")) - })?; - - if !status.is_success() { - let error = self.parse_error_body(status.as_u16(), &text); - return Err(TinyAgentsError::Model( - self.provider_failure_message(&error), - )); - } - - let value: Value = serde_json::from_str(&text)?; - parse_response(value) - } - - /// Streams the OpenAI Chat Completions response as a real [`ModelStream`]. - /// - /// Sends the request with `stream: true` (and `stream_options.include_usage` - /// so a usage chunk is delivered), reads the Server-Sent-Events body - /// incrementally with [`reqwest::Response::bytes_stream`], and parses each - /// `data:` line into [`ModelStreamItem`]s: a leading - /// [`ModelStreamItem::Started`], a [`ModelStreamItem::MessageDelta`] per text - /// fragment, a [`ModelStreamItem::ToolCallDelta`] per tool-call argument - /// fragment, a [`ModelStreamItem::UsageDelta`] when usage arrives, and a - /// terminal [`ModelStreamItem::Completed`] carrying the fully merged - /// response (with reassembled tool-call names and ids). Transport errors - /// surface as a terminal [`ModelStreamItem::ProviderFailed`]. - /// - /// # Errors - /// - /// Returns [`TinyAgentsError::Model`] when the initial request fails or the - /// endpoint returns a non-2xx status (per-chunk transport errors are - /// surfaced as [`ModelStreamItem::ProviderFailed`] inside the stream - /// instead). - async fn stream(&self, _state: &State, request: ModelRequest) -> Result { - let mut body = self.translate_request(&request)?; - body.stream = true; - body.stream_options = Some(json!({ "include_usage": true })); - let url = format!("{}/chat/completions", self.base_url); - - let response = self - .client - .post(&url) - .header("Authorization", format!("Bearer {}", self.api_key)) - .json(&body) - .send() - .await - .map_err(|e| { - let error = self.provider_error( - format!("stream request to {url} failed: {e}"), - None, - None, - None, - ); - TinyAgentsError::Model(self.provider_failure_message(&error)) - })?; - - let status = response.status(); - if !status.is_success() { - let text = response.text().await.unwrap_or_default(); - let error = self.parse_error_body(status.as_u16(), &text); - return Err(TinyAgentsError::Model( - self.provider_failure_message(&error), - )); - } - - // Map each raw byte chunk onto an owned `Vec` so the boxed stream's - // item type is nameable without depending on the `bytes` crate. - let bytes = response.bytes_stream().map(|chunk| { - chunk - .map(|b| b.to_vec()) - .map_err(|e| TinyAgentsError::Model(format!("stream chunk failed: {e}"))) - }); - - let state = SseState { - bytes: Box::pin(bytes), - buf: String::new(), - pending: VecDeque::new(), - acc: OpenAiStreamAcc::default(), - provider: self.provider.clone(), - model: self.model.clone(), - started: false, - finished: false, - terminal_emitted: false, - }; - - Ok(Box::pin(futures::stream::unfold(state, sse_next))) - } -} +/// Sane default TCP connect timeout applied to every call. Bounds connection +/// establishment without capping the (potentially long) response body, so it is +/// safe for streaming too. +const DEFAULT_CONNECT_TIMEOUT_SECS: u64 = 30; +/// Default overall timeout applied to unary calls when the request does not set +/// [`ModelRequest::timeout_ms`]. Streaming calls get no overall cap by default. +const DEFAULT_REQUEST_TIMEOUT_SECS: u64 = 600; + +mod convert; +mod sse; +mod transport; + +pub use transport::OpenAiModel; + +use convert::*; +use sse::*; +#[cfg(test)] +use transport::request_timeout; #[cfg(test)] mod test; diff --git a/src/harness/providers/openai/sse.rs b/src/harness/providers/openai/sse.rs new file mode 100644 index 0000000..26a0241 --- /dev/null +++ b/src/harness/providers/openai/sse.rs @@ -0,0 +1,343 @@ +//! Server-sent-events stream parsing and incremental accumulation +//! (`SseState`, `OpenAiStreamAcc`, `sse_next`). +//! +//! Split out of `openai/mod.rs`; see that module's doc comment for the +//! full provider overview. + +use super::*; + +/// In-progress reconstruction of a single tool call across streamed fragments. +#[derive(Clone, Debug, Default)] +pub(super) struct ToolCallBuild { + /// Provider-assigned call id (filled from the first fragment carrying it). + id: String, + /// Function name (filled from the first fragment carrying it). + name: String, + /// Concatenated stringified-JSON argument fragments. + args: String, +} + +/// Provider-side accumulator that rebuilds the authoritative [`ModelResponse`] +/// from streamed chunks. Distinct from the generic +/// [`StreamAccumulator`][crate::harness::model::StreamAccumulator]: it tracks +/// tool-call names and ids (which the neutral deltas omit) so the terminal +/// [`ModelStreamItem::Completed`] carries a faithful response. +#[derive(Clone, Debug, Default)] +pub(super) struct OpenAiStreamAcc { + id: Option, + text: String, + tool_calls: Vec, + usage: Option, + finish_reason: Option, +} + +impl OpenAiStreamAcc { + /// Folds one parsed chunk into the accumulator and pushes the corresponding + /// neutral [`ModelStreamItem`]s onto `pending`. + fn ingest(&mut self, chunk: ChatCompletionChunk, pending: &mut VecDeque) { + if let Some(id) = chunk.id + && self.id.is_none() + { + self.id = Some(id); + } + if let Some(usage_wire) = chunk.usage { + let usage = convert_usage(usage_wire); + self.usage = Some(usage); + pending.push_back(ModelStreamItem::UsageDelta(usage)); + } + for mut choice in chunk.choices { + if let Some(reason) = choice.finish_reason { + self.finish_reason = Some(reason); + } + let reasoning = delta_reasoning_text(&mut choice.delta); + if !reasoning.is_empty() { + pending.push_back(ModelStreamItem::MessageDelta(MessageDelta { + text: String::new(), + reasoning, + tool_call: None, + })); + } + if let Some(content) = choice.delta.content.filter(|c| !c.is_empty()) { + self.text.push_str(&content); + pending.push_back(ModelStreamItem::MessageDelta(MessageDelta { + text: content, + reasoning: String::new(), + tool_call: None, + })); + } + for fragment in choice.delta.tool_calls { + let idx = self.resolve_slot(&fragment); + let slot = &mut self.tool_calls[idx]; + if let Some(id) = fragment.id.filter(|id| !id.is_empty()) { + slot.id = id; + } + if let Some(function) = fragment.function { + if let Some(name) = function.name.filter(|n| !n.is_empty()) { + slot.name = name; + } + if let Some(args) = function.arguments.filter(|a| !a.is_empty()) { + slot.args.push_str(&args); + let call_id = tool_call_id(idx, &slot.id); + pending.push_back(ModelStreamItem::ToolCallDelta(ToolDelta { + call_id, + content: args, + })); + } + } + } + } + } + + /// Resolves the accumulator slot a streamed tool-call fragment belongs to. + /// + /// OpenAI itself always sends a stable `index`; some OpenAI-compatible + /// backends omit it. When `index` is present it selects the slot directly + /// (growing the vector as needed). When it is absent, fragments are + /// correlated by `id`: a fragment carrying a new id opens a new slot, one + /// carrying a known id reuses that slot, and an id-less continuation fragment + /// (arguments only) appends to the most recent slot — so parallel calls no + /// longer all collapse onto slot 0. + fn resolve_slot(&mut self, fragment: &ToolCallChunkWire) -> usize { + if let Some(index) = fragment.index { + let idx = index as usize; + while self.tool_calls.len() <= idx { + self.tool_calls.push(ToolCallBuild::default()); + } + return idx; + } + if let Some(id) = fragment.id.as_deref().filter(|id| !id.is_empty()) { + if let Some(pos) = self.tool_calls.iter().position(|slot| slot.id == id) { + return pos; + } + self.tool_calls.push(ToolCallBuild::default()); + return self.tool_calls.len() - 1; + } + if self.tool_calls.is_empty() { + self.tool_calls.push(ToolCallBuild::default()); + } + self.tool_calls.len() - 1 + } + + /// Consumes the accumulator into the final, merged [`ModelResponse`]. + fn into_response(self) -> Result { + let mut content = Vec::new(); + if !self.text.is_empty() { + content.push(ContentBlock::Text(self.text)); + } + // Enumerate over the full slot vector *before* filtering so the synthetic + // fallback id (`tool-{idx}`) matches the one streamed in `ToolCallDelta` + // items — filtering first would renumber the slots and desynchronize the + // delta ids from the final call ids. + let tool_calls = self + .tool_calls + .into_iter() + .enumerate() + .filter(|(_, b)| !b.name.is_empty() || !b.args.is_empty()) + .map(|(idx, b)| { + let id = tool_call_id(idx, &b.id); + Ok(ToolCall { + id: id.clone(), + name: b.name.clone(), + arguments: parse_tool_arguments("openai stream", &id, &b.name, &b.args)?, + }) + }) + .collect::>>()?; + let message = AssistantMessage { + id: self.id, + content, + tool_calls, + usage: self.usage, + }; + Ok(ModelResponse { + message, + usage: self.usage, + finish_reason: self.finish_reason, + raw: None, + resolved_model: None, + }) + } +} + +/// Mutable driver state threaded through [`futures::stream::unfold`] while +/// parsing the SSE byte stream into [`ModelStreamItem`]s. +pub(super) struct SseState { + /// Raw response byte chunks (errors already mapped onto the crate error). + pub(super) bytes: Pin>> + Send>>, + /// Raw bytes received but not yet split into complete lines. Kept as bytes + /// (not a `String`) so a multi-byte UTF-8 character split across two network + /// chunks is reassembled before decoding, instead of being corrupted into + /// replacement characters by a premature lossy decode. + pub(super) buf: Vec, + /// Parsed items waiting to be yielded, in order. + pub(super) pending: VecDeque, + /// Provider-side response reconstruction. + pub(super) acc: OpenAiStreamAcc, + /// Provider family id used in normalized stream failures. + pub(super) provider: String, + /// Provider model id used in normalized stream failures. + pub(super) model: String, + /// Whether the leading [`ModelStreamItem::Started`] has been emitted. + pub(super) started: bool, + /// Whether the byte stream ended or `[DONE]` was seen. + pub(super) finished: bool, + /// Whether the terminal [`ModelStreamItem::Completed`]/[`ModelStreamItem::Failed`] + /// has been emitted. + pub(super) terminal_emitted: bool, +} + +impl SseState { + /// Splits buffered bytes into complete newline-terminated lines and folds + /// each SSE `data:` payload into the accumulator. The trailing partial line + /// (if any) is kept in `buf` for the next chunk, so a `data:` line split + /// across chunk boundaries — including one that splits a multi-byte UTF-8 + /// character — is only decoded once it is complete. + fn drain_lines(&mut self) { + // Scan with a moving start offset and drain the whole consumed prefix + // once at the end, rather than `drain(..=pos)` per line: the old form + // allocated a `Vec` per line and shifted the remaining buffer down + // on every line (O(n^2) over a chunk carrying many lines). + let mut start = 0; + while let Some(rel) = self.buf[start..].iter().position(|&b| b == b'\n') { + let end = start + rel; + // A complete line (bounded by the ASCII `\n`) is whole UTF-8, so a + // lossy decode here can no longer straddle a chunk boundary. The + // `into_owned` detaches the line from `buf` so `process_line` can + // borrow `self` mutably. + let line = String::from_utf8_lossy(&self.buf[start..end]).into_owned(); + start = end + 1; + self.process_line(&line); + } + if start > 0 { + self.buf.drain(..start); + } + } + + /// Folds any bytes still buffered after the byte stream ends into a final + /// line. Providers that terminate the last SSE event without a trailing + /// newline would otherwise leave the final `data:` payload unprocessed. + fn drain_remaining(&mut self) { + if self.buf.is_empty() { + return; + } + let line = String::from_utf8_lossy(&self.buf).into_owned(); + self.buf.clear(); + self.process_line(&line); + } + + /// Parses one SSE line and folds any resulting chunk into the accumulator. + fn process_line(&mut self, line: &str) { + let line = line.trim(); + if line.is_empty() { + return; + } + let Some(rest) = line.strip_prefix("data:") else { + return; + }; + let payload = rest.trim(); + if payload == "[DONE]" { + self.finished = true; + return; + } + // Ignore keepalives / unparseable lines rather than failing the run. + let Ok(value) = serde_json::from_str::(payload) else { + return; + }; + // Some providers stream a mid-stream `{"error": ...}` payload instead of + // a chunk. This also deserializes cleanly as an all-defaults + // `ChatCompletionChunk`, so it must be detected first and surfaced as a + // terminal failure rather than folded in as an empty chunk and swallowed. + if let Some(error) = value.get("error") { + self.pending + .push_back(ModelStreamItem::ProviderFailed(self.stream_error(error))); + self.finished = true; + self.terminal_emitted = true; + return; + } + if let Ok(chunk) = serde_json::from_value::(value) { + let mut pending = std::mem::take(&mut self.pending); + self.acc.ingest(chunk, &mut pending); + self.pending = pending; + } + } + + /// Builds a normalized [`ProviderError`] from a streamed `error` payload. + fn stream_error(&self, error: &Value) -> ProviderError { + let message = error + .get("message") + .and_then(Value::as_str) + .filter(|message| !message.trim().is_empty()) + .unwrap_or("provider reported a stream error") + .to_string(); + let code = error + .get("code") + .or_else(|| error.get("type")) + .and_then(Value::as_str) + .map(str::to_string); + ProviderError { + provider: self.provider.clone(), + model: Some(self.model.clone()), + code, + message, + retryable: false, + raw: Some(error.clone()), + ..ProviderError::default() + } + } +} + +/// Advances the SSE [`SseState`] by one item for [`futures::stream::unfold`]. +pub(super) async fn sse_next(mut state: SseState) -> Option<(ModelStreamItem, SseState)> { + loop { + if let Some(item) = state.pending.pop_front() { + return Some((item, state)); + } + if !state.started { + state.started = true; + return Some((ModelStreamItem::Started, state)); + } + if state.finished { + if state.terminal_emitted { + return None; + } + state.terminal_emitted = true; + return match std::mem::take(&mut state.acc).into_response() { + Ok(response) => Some((ModelStreamItem::Completed(response), state)), + Err(error) => { + let provider_error = ProviderError { + provider: state.provider.clone(), + model: Some(state.model.clone()), + code: Some("invalid_tool_arguments".to_string()), + message: error.to_string(), + retryable: false, + ..ProviderError::default() + }; + Some((ModelStreamItem::ProviderFailed(provider_error), state)) + } + }; + } + match state.bytes.next().await { + Some(Ok(chunk)) => { + state.buf.extend_from_slice(&chunk); + state.drain_lines(); + } + Some(Err(error)) => { + state.finished = true; + state.terminal_emitted = true; + let provider_error = ProviderError { + provider: state.provider.clone(), + model: Some(state.model.clone()), + message: error.to_string(), + retryable: true, + ..ProviderError::default() + }; + return Some((ModelStreamItem::ProviderFailed(provider_error), state)); + } + None => { + // Drain any final `data:` line the provider sent without a + // trailing newline before terminating. + state.drain_remaining(); + state.finished = true; + } + } + } +} diff --git a/src/harness/providers/openai/test.rs b/src/harness/providers/openai/test.rs index fe6a495..759815c 100644 --- a/src/harness/providers/openai/test.rs +++ b/src/harness/providers/openai/test.rs @@ -255,7 +255,8 @@ fn parses_openai_response_with_content_tool_call_and_usage() { "prompt_tokens": 42, "completion_tokens": 8, "total_tokens": 50, - "prompt_tokens_details": { "cached_tokens": 30 } + "prompt_tokens_details": { "cached_tokens": 30 }, + "completion_tokens_details": { "reasoning_tokens": 6 } } }); @@ -281,6 +282,7 @@ fn parses_openai_response_with_content_tool_call_and_usage() { assert_eq!(usage.output_tokens, 8); assert_eq!(usage.total_tokens, 50); assert_eq!(usage.cache_read_tokens, 30); + assert_eq!(usage.reasoning_tokens, 6); // Raw JSON preserved verbatim. assert_eq!(response.raw, Some(body)); @@ -344,6 +346,83 @@ fn parses_text_only_response_without_usage_details() { assert_eq!(usage.cache_read_tokens, 0); } +#[test] +fn total_tokens_falls_back_to_prompt_plus_completion_when_omitted() { + // Some OpenAI-compatible backends omit `total_tokens` entirely; it must + // not silently deserialize to a misleading `0` when prompt/completion + // tokens were clearly reported. + let body = json!({ + "id": "chatcmpl-omit", + "choices": [ + { + "message": { "role": "assistant", "content": "Hi!" }, + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 5, + "completion_tokens": 2 + } + }); + + let response = parse_response(body).unwrap(); + let usage = response.usage.unwrap(); + assert_eq!(usage.input_tokens, 5); + assert_eq!(usage.output_tokens, 2); + assert_eq!(usage.total_tokens, 7); +} + +#[test] +fn parses_empty_tool_arguments_as_empty_object() { + // Some compat backends send an empty arguments string for a zero-argument + // tool call; it must map to `{}`, not fail as malformed JSON. + let body = json!({ + "id": "chatcmpl-noargs", + "choices": [ + { + "message": { + "role": "assistant", + "tool_calls": [ + { + "id": "call-empty", + "type": "function", + "function": { "name": "ping", "arguments": "" } + } + ] + }, + "finish_reason": "tool_calls" + } + ] + }); + + let response = parse_response(body).unwrap(); + let calls = response.tool_calls(); + assert_eq!(calls.len(), 1); + assert_eq!(calls[0].name, "ping"); + assert_eq!(calls[0].arguments, json!({})); +} + +#[tokio::test] +async fn sse_stream_empty_tool_arguments_reconstruct_as_empty_object() { + // A streamed tool call whose only arguments fragment is empty. The merged + // call must carry `{}` rather than fail terminally. + let raw: Vec> = vec![ + b"data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"id\":\"call-x\",\"function\":{\"name\":\"ping\",\"arguments\":\"\"}}]},\"finish_reason\":\"tool_calls\"}]}\n\n".to_vec(), + b"data: [DONE]\n\n".to_vec(), + ]; + + let items = collect_sse(raw).await; + let mut merged = StreamAccumulator::new(); + for item in &items { + merged.push(item); + } + let response = merged.finish().unwrap(); + let calls = response.tool_calls(); + assert_eq!(calls.len(), 1); + assert_eq!(calls[0].name, "ping"); + assert_eq!(calls[0].arguments, json!({})); +} + #[test] fn parse_response_errors_on_empty_choices() { let body = json!({ "id": "x", "choices": [] }); @@ -351,6 +430,32 @@ fn parse_response_errors_on_empty_choices() { assert!(matches!(err, TinyAgentsError::Model(_))); } +#[test] +fn parse_error_body_classifies_retryability_by_http_status() { + // Regression test: retry used to see every provider failure flattened + // into `TinyAgentsError::Model(String)`, so it could not distinguish a + // retryable 429 from a non-retryable 401 and retried both. The status + // code alone must fully determine `ProviderError::retryable`. + let m = model(); + + let unauthorized = m.parse_error_body( + 401, + r#"{"error":{"message":"Incorrect API key provided","type":"invalid_request_error","code":"invalid_api_key"}}"#, + ); + assert_eq!(unauthorized.status, Some(401)); + assert!(!unauthorized.retryable, "401 must not be retryable"); + + let rate_limited = m.parse_error_body( + 429, + r#"{"error":{"message":"Rate limit reached","type":"requests","code":"rate_limit_exceeded"}}"#, + ); + assert_eq!(rate_limited.status, Some(429)); + assert!(rate_limited.retryable, "429 must be retryable"); + + let server_error = m.parse_error_body(500, r#"{"error":{"message":"internal error"}}"#); + assert!(server_error.retryable, "5xx must be retryable"); +} + #[test] fn compatible_presets_set_base_url_and_default_model() { let deepseek = OpenAiModel::deepseek("k"); @@ -443,7 +548,7 @@ async fn sse_stream_parses_text_tool_calls_and_usage() { let state = SseState { bytes: Box::pin(bytes), - buf: String::new(), + buf: Vec::new(), pending: std::collections::VecDeque::new(), acc: OpenAiStreamAcc::default(), provider: "openai".to_string(), @@ -499,6 +604,47 @@ async fn sse_stream_parses_text_tool_calls_and_usage() { assert_eq!(response.usage.unwrap().total_tokens, 8); } +#[tokio::test] +async fn sse_stream_preserves_reasoning_content_as_side_channel() { + let raw: Vec> = vec![ + b"data: {\"choices\":[{\"delta\":{\"reasoning_content\":\"think \"}}]}\n\n".to_vec(), + b"data: {\"choices\":[{\"delta\":{\"reasoning\":\"carefully\"}}]}\n\n".to_vec(), + b"data: {\"choices\":[{\"delta\":{\"content\":\"answer\"},\"finish_reason\":\"stop\"}]}\n\n" + .to_vec(), + b"data: {\"choices\":[],\"usage\":{\"prompt_tokens\":5,\"completion_tokens\":7,\"total_tokens\":12,\"completion_tokens_details\":{\"reasoning_tokens\":4}}}\n\n".to_vec(), + b"data: [DONE]\n\n".to_vec(), + ]; + + let items = collect_sse(raw).await; + + let reasoning: String = items + .iter() + .filter_map(|item| match item { + ModelStreamItem::MessageDelta(delta) => Some(delta.reasoning.clone()), + _ => None, + }) + .collect(); + let text: String = items + .iter() + .filter_map(|item| match item { + ModelStreamItem::MessageDelta(delta) => Some(delta.text.clone()), + _ => None, + }) + .collect(); + assert_eq!(reasoning, "think carefully"); + assert_eq!(text, "answer"); + + let mut merged = StreamAccumulator::new(); + for item in &items { + merged.push(item); + } + assert_eq!(merged.reasoning(), "think carefully"); + let response = merged.finish().unwrap(); + assert_eq!(response.text(), "answer"); + let usage = response.usage.unwrap(); + assert_eq!(usage.reasoning_tokens, 4); +} + #[tokio::test] async fn sse_stream_invalid_tool_argument_json_fails_terminally() { use futures::StreamExt; @@ -512,7 +658,7 @@ async fn sse_stream_invalid_tool_argument_json_fails_terminally() { let state = SseState { bytes: Box::pin(bytes), - buf: String::new(), + buf: Vec::new(), pending: std::collections::VecDeque::new(), acc: OpenAiStreamAcc::default(), provider: "openai".to_string(), @@ -545,6 +691,313 @@ async fn sse_stream_invalid_tool_argument_json_fails_terminally() { assert!(err.to_string().contains("invalid_tool_arguments")); } +/// Drives an SSE byte stream through the parser and returns every item. +async fn collect_sse(raw: Vec>) -> Vec { + use futures::StreamExt; + + let bytes = futures::stream::iter(raw.into_iter().map(Ok::, TinyAgentsError>)); + let state = SseState { + bytes: Box::pin(bytes), + buf: Vec::new(), + pending: std::collections::VecDeque::new(), + acc: OpenAiStreamAcc::default(), + provider: "openai".to_string(), + model: "gpt-4.1-mini".to_string(), + started: false, + finished: false, + terminal_emitted: false, + }; + futures::stream::unfold(state, sse_next).collect().await +} + +#[tokio::test] +async fn sse_stream_reassembles_multibyte_char_split_across_chunks() { + // A 4-byte emoji in the content payload, split down the middle across two + // network chunks. A lossy per-chunk decode would corrupt it into U+FFFD + // replacement characters; the byte buffer must reassemble it first. + let line = "data: {\"choices\":[{\"delta\":{\"content\":\"hi😀\"}}]}\n\n"; + let bytes = line.as_bytes(); + let split = line.find('😀').unwrap() + 2; // inside the 4-byte sequence + let raw: Vec> = vec![ + bytes[..split].to_vec(), + bytes[split..].to_vec(), + b"data: [DONE]\n\n".to_vec(), + ]; + + let items = collect_sse(raw).await; + + let mut merged = StreamAccumulator::new(); + for item in &items { + merged.push(item); + } + let response = merged.finish().unwrap(); + assert_eq!(response.text(), "hi😀"); +} + +#[tokio::test] +async fn sse_stream_processes_many_lines_in_one_chunk() { + // A single network chunk carrying several complete `data:` lines exercises + // the batched multi-line drain: every line is parsed in order and the buffer + // is consumed exactly once. Each line contributes one content fragment. + let mut chunk = Vec::new(); + for frag in ["a", "b", "c", "d", "e"] { + chunk.extend_from_slice( + format!("data: {{\"choices\":[{{\"delta\":{{\"content\":\"{frag}\"}}}}]}}\n\n") + .as_bytes(), + ); + } + let raw: Vec> = vec![chunk, b"data: [DONE]\n\n".to_vec()]; + + let items = collect_sse(raw).await; + + let text: String = items + .iter() + .filter_map(|item| match item { + ModelStreamItem::MessageDelta(delta) => Some(delta.text.clone()), + _ => None, + }) + .collect(); + assert_eq!(text, "abcde", "all five lines in one chunk parsed in order"); + + let mut merged = StreamAccumulator::new(); + for item in &items { + merged.push(item); + } + assert_eq!(merged.finish().unwrap().text(), "abcde"); +} + +#[tokio::test] +async fn sse_stream_drains_final_line_without_trailing_newline() { + // The provider ends the stream with a final `data:` event that has no + // trailing newline and no `[DONE]` sentinel. The leftover buffer must be + // drained at EOF so the last fragment is not dropped. + let raw: Vec> = + vec![b"data: {\"choices\":[{\"delta\":{\"content\":\"tail\"}}]}".to_vec()]; + + let items = collect_sse(raw).await; + + assert!(matches!(items.last(), Some(ModelStreamItem::Completed(_)))); + let mut merged = StreamAccumulator::new(); + for item in &items { + merged.push(item); + } + let response = merged.finish().unwrap(); + assert_eq!(response.text(), "tail"); +} + +#[tokio::test] +async fn sse_stream_surfaces_mid_stream_error_payload() { + // The provider streams a text delta, then an `{"error": ...}` payload + // instead of a chunk. It must surface as a terminal ProviderFailed rather + // than be swallowed as an empty chunk. + let raw: Vec> = vec![ + b"data: {\"choices\":[{\"delta\":{\"content\":\"hi\"}}]}\n\n".to_vec(), + b"data: {\"error\":{\"message\":\"upstream exploded\",\"code\":\"server_error\"}}\n\n" + .to_vec(), + b"data: [DONE]\n\n".to_vec(), + ]; + + let items = collect_sse(raw).await; + + let failed = items + .iter() + .find_map(|item| match item { + ModelStreamItem::ProviderFailed(error) => Some(error), + _ => None, + }) + .expect("mid-stream error should emit ProviderFailed"); + assert_eq!(failed.code.as_deref(), Some("server_error")); + assert!(failed.message.contains("upstream exploded")); + // No terminal Completed is emitted after the failure. + assert!(matches!( + items.last(), + Some(ModelStreamItem::ProviderFailed(_)) + )); + + let mut merged = StreamAccumulator::new(); + for item in &items { + merged.push(item); + } + let err = merged + .finish() + .expect_err("stream error must reach accumulator"); + assert!(err.to_string().contains("upstream exploded")); +} + +#[tokio::test] +async fn sse_stream_correlates_indexless_parallel_tool_calls_by_id() { + // A compat backend that omits `index` entirely and interleaves two parallel + // tool calls, correlating fragments only by `id`. Without id-based slotting + // both would collapse onto slot 0; the streamed delta ids must also match the + // final reconstructed call ids. + let raw: Vec> = vec![ + b"data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"id\":\"call-a\",\"function\":{\"name\":\"alpha\",\"arguments\":\"{\\\"x\\\":\"}}]}}]}\n\n".to_vec(), + b"data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"id\":\"call-b\",\"function\":{\"name\":\"beta\",\"arguments\":\"{\\\"y\\\":\"}}]}}]}\n\n".to_vec(), + b"data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"id\":\"call-a\",\"function\":{\"arguments\":\"1}\"}}]}}]}\n\n".to_vec(), + b"data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"id\":\"call-b\",\"function\":{\"arguments\":\"2}\"}}]},\"finish_reason\":\"tool_calls\"}]}\n\n".to_vec(), + b"data: [DONE]\n\n".to_vec(), + ]; + + let items = collect_sse(raw).await; + + // Every streamed tool-call delta id must be a real call id, never a slot-0 + // collapse. + let delta_ids: Vec = items + .iter() + .filter_map(|item| match item { + ModelStreamItem::ToolCallDelta(delta) => Some(delta.call_id.clone()), + _ => None, + }) + .collect(); + assert!(delta_ids.contains(&"call-a".to_string())); + assert!(delta_ids.contains(&"call-b".to_string())); + + let mut merged = StreamAccumulator::new(); + for item in &items { + merged.push(item); + } + let response = merged.finish().unwrap(); + let calls = response.tool_calls(); + assert_eq!( + calls.len(), + 2, + "parallel calls must not collapse: {calls:?}" + ); + assert_eq!(calls[0].id, "call-a"); + assert_eq!(calls[0].name, "alpha"); + assert_eq!(calls[0].arguments, json!({ "x": 1 })); + assert_eq!(calls[1].id, "call-b"); + assert_eq!(calls[1].name, "beta"); + assert_eq!(calls[1].arguments, json!({ "y": 2 })); +} + +#[tokio::test] +async fn sse_stream_indexless_fallback_ids_match_between_delta_and_final() { + // No `index` and no `id` at all (arguments-only continuation on the same + // slot). The synthetic fallback id streamed in the delta must equal the id on + // the final reconstructed call. + let raw: Vec> = vec![ + b"data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"function\":{\"name\":\"solo\",\"arguments\":\"{\\\"n\\\":\"}}]}}]}\n\n".to_vec(), + b"data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"function\":{\"arguments\":\"7}\"}}]},\"finish_reason\":\"tool_calls\"}]}\n\n".to_vec(), + b"data: [DONE]\n\n".to_vec(), + ]; + + let items = collect_sse(raw).await; + + let delta_id = items + .iter() + .find_map(|item| match item { + ModelStreamItem::ToolCallDelta(delta) => Some(delta.call_id.clone()), + _ => None, + }) + .expect("a tool-call delta is emitted"); + + let mut merged = StreamAccumulator::new(); + for item in &items { + merged.push(item); + } + let response = merged.finish().unwrap(); + let calls = response.tool_calls(); + assert_eq!(calls.len(), 1); + assert_eq!(calls[0].id, delta_id, "delta id must match final call id"); + assert_eq!(calls[0].name, "solo"); + assert_eq!(calls[0].arguments, json!({ "n": 7 })); +} + +#[test] +fn user_image_blocks_render_as_content_parts() { + use crate::harness::message::{ContentBlock, ImageRef, UserMessage}; + + let request = ModelRequest::new(vec![Message::User(UserMessage { + content: vec![ + ContentBlock::Text("What is in this image?".to_string()), + ContentBlock::Image(ImageRef { + url: "https://example.test/cat.png".to_string(), + mime_type: Some("image/png".to_string()), + }), + ], + })]); + + let value = serde_json::to_value(model().translate_request(&request).unwrap()).unwrap(); + let content = &value["messages"][0]["content"]; + + // Content is an array of parts, not a dropped/plain string. + assert!(content.is_array(), "expected content parts, got {content}"); + assert_eq!(content[0]["type"], json!("text")); + assert_eq!(content[0]["text"], json!("What is in this image?")); + assert_eq!(content[1]["type"], json!("image_url")); + assert_eq!( + content[1]["image_url"]["url"], + json!("https://example.test/cat.png") + ); +} + +#[test] +fn text_only_user_message_stays_a_plain_string() { + // The common text-only case keeps its historical plain-string wire shape. + let request = ModelRequest::new(vec![Message::user("hi")]); + let value = serde_json::to_value(model().translate_request(&request).unwrap()).unwrap(); + assert_eq!(value["messages"][0]["content"], json!("hi")); +} + +#[test] +fn provider_extension_block_fails_closed_instead_of_dropping() { + use crate::harness::message::{ContentBlock, UserMessage}; + + let request = ModelRequest::new(vec![Message::User(UserMessage { + content: vec![ContentBlock::ProviderExtension(json!({ "opaque": true }))], + })]); + + let error = model().translate_request(&request).unwrap_err(); + assert!(matches!(error, TinyAgentsError::Validation(_))); + assert!(error.to_string().contains("provider-extension")); +} + +#[test] +fn routes_max_tokens_to_max_completion_tokens_for_o_series() { + // o-series reasoning models reject `max_tokens` and require + // `max_completion_tokens`. + let request = ModelRequest::new(vec![Message::user("hi")]).with_max_tokens(128); + let model = OpenAiModel::new("k").with_model("o3-mini"); + let value = serde_json::to_value(model.translate_request(&request).unwrap()).unwrap(); + + assert!(value.get("max_tokens").is_none()); + assert_eq!(value["max_completion_tokens"], json!(128)); +} + +#[test] +fn keeps_max_tokens_for_classic_models() { + let request = ModelRequest::new(vec![Message::user("hi")]).with_max_tokens(128); + let value = serde_json::to_value(model().translate_request(&request).unwrap()).unwrap(); + + assert_eq!(value["max_tokens"], json!(128)); + assert!(value.get("max_completion_tokens").is_none()); +} + +#[test] +fn request_timeout_prefers_explicit_override() { + // An explicit per-request timeout wins for both unary and streaming calls. + assert_eq!( + request_timeout(Some(1_500), false), + Some(Duration::from_millis(1_500)) + ); + assert_eq!( + request_timeout(Some(1_500), true), + Some(Duration::from_millis(1_500)) + ); +} + +#[test] +fn request_timeout_defaults_by_call_kind() { + // Unary calls fall back to a sane overall default; streaming calls get no + // overall cap so a long stream is not truncated. + assert_eq!( + request_timeout(None, false), + Some(Duration::from_secs(DEFAULT_REQUEST_TIMEOUT_SECS)) + ); + assert_eq!(request_timeout(None, true), None); +} + #[test] fn from_env_errors_when_api_key_missing() { // Snapshot and clear the key so the missing-key path is exercised diff --git a/src/harness/providers/openai/transport.rs b/src/harness/providers/openai/transport.rs new file mode 100644 index 0000000..f085fff --- /dev/null +++ b/src/harness/providers/openai/transport.rs @@ -0,0 +1,655 @@ +//! HTTP transport: `OpenAiModel` construction, provider presets, request +//! building, and the `ChatModel` impl (`invoke`/`stream`). +//! +//! Split out of `openai/mod.rs`; see that module's doc comment for the +//! full provider overview. + +use super::*; + +/// A [`ChatModel`] backed by the hosted OpenAI Chat Completions API. +/// +/// Construct one with [`OpenAiModel::new`] (plus the `with_*` builders) or +/// [`OpenAiModel::from_env`]. The model holds a reusable [`reqwest::Client`] so +/// repeated calls share a connection pool. +pub struct OpenAiModel { + /// Shared HTTP client. + client: reqwest::Client, + /// API key sent as a `Bearer` token. + api_key: String, + /// Default model id used when a request does not override it. + model: String, + /// Provider family identifier used in profiles and normalized errors. + provider: String, + /// API base URL (no trailing slash); `/chat/completions` is appended. + base_url: String, + /// Capability profile derived from the default model id. + profile: ModelProfile, +} + +/// Returns `true` for OpenAI o-series reasoning models (`o1`/`o3`/`o4`), which +/// reject `max_tokens` and require `max_completion_tokens` instead. +pub(super) fn is_reasoning_model(model: &str) -> bool { + let lower = model.to_ascii_lowercase(); + lower.starts_with("o1") || lower.starts_with("o3") || lower.starts_with("o4") +} + +/// Derives a static [`ModelProfile`] for an OpenAI(-compatible) model id. +/// +/// All targets support tool calling, streaming (including tool-call chunks), +/// and JSON Schema response formats. Modern OpenAI-family models additionally +/// advertise native structured output and (for the o-series) reasoning output. +pub(super) fn derive_profile(provider: &str, model: &str) -> ModelProfile { + let lower = model.to_ascii_lowercase(); + let reasoning = is_reasoning_model(model); + let native_structured = lower.contains("gpt-4o") || lower.contains("gpt-4.1") || reasoning; + ModelProfile { + provider: Some(provider.to_string()), + model: Some(model.to_string()), + status: ModelStatus::Stable, + modalities: Modalities { + image_in: true, + ..Modalities::default() + }, + tool_calling: true, + parallel_tool_calls: true, + streaming: true, + streaming_tool_chunks: true, + native_structured_output: native_structured, + json_schema: true, + reasoning, + ..ModelProfile::default() + } +} + +impl OpenAiModel { + /// Creates a model with the given API key, the default model + /// (`gpt-4.1-mini`), and the default base URL (`https://api.openai.com/v1`). + pub fn new(api_key: impl Into) -> Self { + Self { + client: reqwest::Client::builder() + .connect_timeout(Duration::from_secs(DEFAULT_CONNECT_TIMEOUT_SECS)) + .build() + .expect("default reqwest client builds"), + api_key: api_key.into(), + model: DEFAULT_MODEL.to_string(), + provider: "openai".to_string(), + base_url: DEFAULT_BASE_URL.to_string(), + profile: derive_profile("openai", DEFAULT_MODEL), + } + } + + /// Overrides the default model id. + pub fn with_model(mut self, model: impl Into) -> Self { + self.model = model.into(); + self.profile = derive_profile(&self.provider, &self.model); + self + } + + /// Overrides the provider family id used in profiles and normalized errors. + pub fn with_provider(mut self, provider: impl Into) -> Self { + self.provider = provider.into(); + self.profile = derive_profile(&self.provider, &self.model); + self + } + + /// Overrides the API base URL. A trailing slash is trimmed so the joined + /// endpoint is always `{base_url}/chat/completions`. + pub fn with_base_url(mut self, base_url: impl Into) -> Self { + self.base_url = base_url.into().trim_end_matches('/').to_string(); + self + } + + /// Builds a model from environment variables. + /// + /// Reads `OPENAI_API_KEY` (required), `OPENAI_MODEL` (optional, defaults to + /// `gpt-4.1-mini`), and `OPENAI_BASE_URL` (optional, defaults to + /// `https://api.openai.com/v1`). + /// + /// # Errors + /// + /// Returns [`TinyAgentsError::Validation`] when `OPENAI_API_KEY` is missing + /// or empty. + pub fn from_env() -> Result { + let api_key = std::env::var("OPENAI_API_KEY") + .ok() + .filter(|k| !k.trim().is_empty()) + .ok_or_else(|| { + TinyAgentsError::Validation( + "OPENAI_API_KEY is not set; export it or add it to a .env file".to_string(), + ) + })?; + + let mut model = Self::new(api_key); + if let Ok(name) = std::env::var("OPENAI_MODEL") + && !name.trim().is_empty() + { + model = model.with_model(name); + } + if let Ok(url) = std::env::var("OPENAI_BASE_URL") + && !url.trim().is_empty() + { + model = model.with_base_url(url); + } + Ok(model) + } + + /// Builds an OpenAI-compatible model from a provider spec and explicit API + /// key. + pub fn from_spec(spec: ProviderSpec, api_key: impl Into) -> Result { + if spec.model.trim().is_empty() { + return Err(TinyAgentsError::Validation( + "provider spec model must not be empty".to_string(), + )); + } + if spec.base_url.trim().is_empty() { + return Err(TinyAgentsError::Validation( + "provider spec base_url must not be empty".to_string(), + )); + } + Ok(Self::compatible_provider( + spec.provider, + api_key, + spec.base_url, + spec.model, + )) + } + + /// Builds an OpenAI-compatible model from a provider spec, reading the API + /// key from the spec's environment variable when required. + pub fn from_spec_env(spec: ProviderSpec) -> Result { + let api_key = if spec.requires_api_key { + let env = spec.api_key_env.as_deref().ok_or_else(|| { + TinyAgentsError::Validation(format!( + "{} requires an api_key_env in ProviderSpec", + spec.provider + )) + })?; + std::env::var(env) + .ok() + .filter(|k| !k.trim().is_empty()) + .ok_or_else(|| { + TinyAgentsError::Validation(format!( + "{env} is not set; export it or provide an explicit API key" + )) + })? + } else { + "local".to_string() + }; + Self::from_spec(spec, api_key) + } + + /// Lists the models the provider advertises via `GET {base_url}/models`. + /// + /// This is a provider/account-level capability (independent of which model + /// this handle is bound to), so it uses the same credentials and base URL as + /// chat calls. Every OpenAI-compatible endpoint (Ollama, Together, Groq, + /// OpenRouter, …) serves the same shape, so this doubles as runtime model + /// discovery for local/self-hosted providers. Returned ids can be fed to + /// [`with_model`](Self::with_model). + /// + /// # Errors + /// + /// Returns [`TinyAgentsError::Model`] on transport failure or a non-2xx + /// status, and [`TinyAgentsError::Serialization`] when the body cannot be + /// decoded. + pub async fn list_models(&self) -> Result> { + let url = format!("{}/models", self.base_url); + + let response = self + .send_checked(self.authorized(self.client.get(&url)), "request", &url) + .await?; + + let text = response.text().await.map_err(|e| { + TinyAgentsError::Model(format!("openai response body read failed: {e}")) + })?; + + let listing: ModelListWire = serde_json::from_str(&text)?; + Ok(listing.data) + } + + // ----------------------------------------------------------------------- + // OpenAI-compatible provider presets + // + // DeepSeek, Groq, xAI, OpenRouter, Together, Mistral, Ollama, and + // Anthropic's compatibility endpoint all accept the same Chat Completions + // wire format, so the same [`OpenAiModel`] talks to all of them — only the + // base URL and default model differ. Each preset is a thin wrapper over + // [`OpenAiModel::new`] + [`with_base_url`][Self::with_base_url] + + // [`with_model`][Self::with_model]; override the model with `with_model`. + // ----------------------------------------------------------------------- + + /// Points at an arbitrary OpenAI-compatible endpoint with an explicit base + /// URL and default model. + /// + /// Use this for any provider that implements the Chat Completions API but is + /// not covered by a named preset below. + pub fn compatible( + api_key: impl Into, + base_url: impl Into, + model: impl Into, + ) -> Self { + Self::new(api_key).with_base_url(base_url).with_model(model) + } + + /// Points at an arbitrary OpenAI-compatible endpoint with an explicit + /// provider id, base URL, and default model. + pub fn compatible_provider( + provider: impl Into, + api_key: impl Into, + base_url: impl Into, + model: impl Into, + ) -> Self { + Self::new(api_key) + .with_provider(provider) + .with_base_url(base_url) + .with_model(model) + } + + /// DeepSeek (`https://api.deepseek.com/v1`), default model `deepseek-chat`. + pub fn deepseek(api_key: impl Into) -> Self { + Self::compatible_provider( + "deepseek", + api_key, + "https://api.deepseek.com/v1", + "deepseek-chat", + ) + } + + /// Anthropic's OpenAI-compatible endpoint (`https://api.anthropic.com/v1`), + /// default model `claude-3-5-sonnet-latest`. + pub fn anthropic(api_key: impl Into) -> Self { + Self::compatible_provider( + "anthropic", + api_key, + "https://api.anthropic.com/v1", + "claude-3-5-sonnet-latest", + ) + } + + /// Groq (`https://api.groq.com/openai/v1`), default model + /// `llama-3.3-70b-versatile`. + pub fn groq(api_key: impl Into) -> Self { + Self::compatible_provider( + "groq", + api_key, + "https://api.groq.com/openai/v1", + "llama-3.3-70b-versatile", + ) + } + + /// xAI (`https://api.x.ai/v1`), default model `grok-2-latest`. + pub fn xai(api_key: impl Into) -> Self { + Self::compatible_provider("xai", api_key, "https://api.x.ai/v1", "grok-2-latest") + } + + /// OpenRouter (`https://openrouter.ai/api/v1`), default model + /// `openai/gpt-4o-mini`. + pub fn openrouter(api_key: impl Into) -> Self { + Self::compatible_provider( + "openrouter", + api_key, + "https://openrouter.ai/api/v1", + "openai/gpt-4o-mini", + ) + } + + /// Together AI (`https://api.together.xyz/v1`), default model + /// `meta-llama/Llama-3.3-70B-Instruct-Turbo`. + pub fn together(api_key: impl Into) -> Self { + Self::compatible_provider( + "together", + api_key, + "https://api.together.xyz/v1", + "meta-llama/Llama-3.3-70B-Instruct-Turbo", + ) + } + + /// Mistral (`https://api.mistral.ai/v1`), default model + /// `mistral-small-latest`. + pub fn mistral(api_key: impl Into) -> Self { + Self::compatible_provider( + "mistral", + api_key, + "https://api.mistral.ai/v1", + "mistral-small-latest", + ) + } + + /// A local Ollama server (`http://localhost:11434/v1`), default model + /// `llama3.2`. Ollama ignores the API key, so a placeholder is used. + pub fn ollama() -> Self { + Self::compatible_provider("ollama", "ollama", "http://localhost:11434/v1", "llama3.2") + } + + /// Returns the default model id this instance will request. + pub fn model(&self) -> &str { + &self.model + } + + /// Returns the configured provider family id. + pub fn provider(&self) -> &str { + &self.provider + } + + /// Returns the configured API base URL. + pub fn base_url(&self) -> &str { + &self.base_url + } + + /// Translates a provider-neutral [`ModelRequest`] into the OpenAI wire + /// request body. The per-request `model` override wins over the instance + /// default. + pub(super) fn translate_request( + &self, + request: &ModelRequest, + ) -> Result { + let messages = request + .messages + .iter() + .map(translate_message) + .collect::>>()?; + + let tools: Vec = request + .tools + .iter() + .map(|schema| ToolWire { + kind: "function".to_string(), + function: FunctionSchemaWire { + name: schema.name.clone(), + description: schema.description.clone(), + parameters: schema.parameters.clone(), + }, + }) + .collect(); + + // tool_choice is only meaningful when tools are declared. + let tool_choice = if tools.is_empty() { + None + } else { + Some(translate_tool_choice(&request.tool_choice)) + }; + + let response_format = request + .response_format + .as_ref() + .and_then(translate_response_format); + + let model = request.model.clone().unwrap_or_else(|| self.model.clone()); + // The o-series reasoning models reject `max_tokens` and require + // `max_completion_tokens`; classic Chat Completions models use + // `max_tokens`. Route the request's cap to whichever field the target + // model accepts. + let (max_tokens, max_completion_tokens) = if is_reasoning_model(&model) { + (None, request.max_tokens) + } else { + (request.max_tokens, None) + }; + + Ok(ChatCompletionRequest { + model, + messages, + tools, + tool_choice, + response_format, + temperature: request.temperature, + top_p: request.top_p, + max_tokens, + max_completion_tokens, + stop: request.stop_sequences.clone(), + seed: request.seed, + stream: false, + stream_options: None, + extra: provider_extra_options(&request.provider_options)?, + }) + } + + /// Attaches the provider's bearer credential to an outbound request. + /// + /// The single place the `Authorization` header is set, shared by the chat + /// and model-listing calls. + fn authorized(&self, builder: reqwest::RequestBuilder) -> reqwest::RequestBuilder { + builder.header("Authorization", format!("Bearer {}", self.api_key)) + } + + /// Sends `builder` and returns the checked (2xx) [`reqwest::Response`]. + /// + /// The shared transport tail for every OpenAI call: a send/transport failure + /// is mapped to a [`TinyAgentsError::Model`] describing `what` (e.g. + /// `"request"`, `"stream request"`) against `url`, and any non-2xx status is + /// decoded through [`Self::parse_error_body`] into a structured + /// [`TinyAgentsError::Provider`]. On success the raw response is handed back + /// so the caller can read it as text (unary/list) or stream its body. + async fn send_checked( + &self, + builder: reqwest::RequestBuilder, + what: &str, + url: &str, + ) -> Result { + let response = builder.send().await.map_err(|e| { + let error = + self.provider_error(format!("{what} to {url} failed: {e}"), None, None, None); + TinyAgentsError::Model(self.provider_failure_message(&error)) + })?; + + let status = response.status(); + if !status.is_success() { + let text = response.text().await.unwrap_or_default(); + let error = self.parse_error_body(status.as_u16(), &text); + return Err(TinyAgentsError::Provider(Box::new(error))); + } + Ok(response) + } + + /// Issues an authenticated `POST {base_url}/chat/completions` with `body`, + /// applying the resolved per-request timeout, and returns the checked + /// response. + /// + /// Shared by the unary ([`Self::invoke`]) and streaming ([`Self::stream`]) + /// paths so URL construction, auth, timeout selection, and transport/status + /// handling live in exactly one place. + async fn post_json( + &self, + body: &ChatCompletionRequest, + timeout_ms: Option, + streaming: bool, + what: &str, + ) -> Result { + let url = format!("{}/chat/completions", self.base_url); + let mut builder = self.authorized(self.client.post(&url)).json(body); + if let Some(timeout) = request_timeout(timeout_ms, streaming) { + builder = builder.timeout(timeout); + } + self.send_checked(builder, what, &url).await + } + + fn provider_error( + &self, + message: impl Into, + status: Option, + code: Option, + raw: Option, + ) -> ProviderError { + let retryable = status.is_some_and(|s| s == 408 || s == 409 || s == 429 || s >= 500); + ProviderError { + provider: self.provider.clone(), + model: Some(self.model.clone()), + status, + code, + message: message.into(), + retryable, + raw, + } + } + + fn provider_failure_message(&self, error: &ProviderError) -> String { + format!( + "{} returned{}{}: {}", + error.provider, + error + .status + .map(|status| format!(" HTTP {status}")) + .unwrap_or_default(), + error + .code + .as_deref() + .map(|code| format!(" ({code})")) + .unwrap_or_default(), + error.message + ) + } + + pub(super) fn parse_error_body(&self, status: u16, text: &str) -> ProviderError { + let raw = serde_json::from_str::(text).ok(); + let error_obj = raw.as_ref().and_then(|value| value.get("error")); + let message = error_obj + .and_then(|error| error.get("message")) + .and_then(Value::as_str) + .or_else(|| { + raw.as_ref() + .and_then(|value| value.get("message")) + .and_then(Value::as_str) + }) + .filter(|message| !message.trim().is_empty()) + .unwrap_or(text) + .to_string(); + let code = error_obj + .and_then(|error| error.get("code").or_else(|| error.get("type"))) + .and_then(Value::as_str) + .map(str::to_string); + self.provider_error(message, Some(status), code, raw) + } +} + +/// Resolves the per-request timeout to apply to an outbound HTTP call. +/// +/// An explicit [`ModelRequest::timeout_ms`] always wins. Otherwise a unary call +/// falls back to [`DEFAULT_REQUEST_TIMEOUT_SECS`], while a streaming call gets no +/// overall cap (a total-request timeout would truncate a legitimately +/// long-running stream mid-flight). +pub(super) fn request_timeout(timeout_ms: Option, streaming: bool) -> Option { + match timeout_ms { + Some(ms) => Some(Duration::from_millis(ms)), + None if streaming => None, + None => Some(Duration::from_secs(DEFAULT_REQUEST_TIMEOUT_SECS)), + } +} + +/// Returns provider-specific top-level fields to flatten into the request body. +/// +/// Core OpenAI-compatible fields are intentionally reserved so normalized +/// TinyAgents fields remain the source of truth. Callers that need local-model +/// controls should use distinct provider fields such as Ollama's `options`. +pub(super) fn provider_extra_options(options: &Value) -> Result> { + if options.is_null() { + return Ok(Map::new()); + } + let Some(object) = options.as_object() else { + return Err(TinyAgentsError::Validation( + "provider_options for OpenAI-compatible providers must be a JSON object".to_string(), + )); + }; + + const RESERVED: &[&str] = &[ + "model", + "messages", + "tools", + "tool_choice", + "response_format", + "temperature", + "top_p", + "max_tokens", + "max_completion_tokens", + "stop", + "seed", + "stream", + "stream_options", + ]; + + Ok(object + .iter() + .filter(|(key, _)| !RESERVED.contains(&key.as_str())) + .map(|(key, value)| (key.clone(), value.clone())) + .collect()) +} + +#[async_trait] +impl ChatModel for OpenAiModel { + /// Returns the capability profile derived from the configured model id. + fn profile(&self) -> Option<&ModelProfile> { + Some(&self.profile) + } + + /// Invokes the OpenAI Chat Completions endpoint and maps the response into a + /// [`ModelResponse`]. + /// + /// # Errors + /// + /// Returns [`TinyAgentsError::Model`] on transport failure or a non-2xx + /// status (the message includes the status code and response body), and + /// [`TinyAgentsError::Serialization`] when the response cannot be decoded. + async fn invoke(&self, _state: &State, request: ModelRequest) -> Result { + let body = self.translate_request(&request)?; + + let response = self + .post_json(&body, request.timeout_ms, false, "request") + .await?; + + let text = response.text().await.map_err(|e| { + TinyAgentsError::Model(format!("openai response body read failed: {e}")) + })?; + + let value: Value = serde_json::from_str(&text)?; + parse_response(value) + } + + /// Streams the OpenAI Chat Completions response as a real [`ModelStream`]. + /// + /// Sends the request with `stream: true` (and `stream_options.include_usage` + /// so a usage chunk is delivered), reads the Server-Sent-Events body + /// incrementally with [`reqwest::Response::bytes_stream`], and parses each + /// `data:` line into [`ModelStreamItem`]s: a leading + /// [`ModelStreamItem::Started`], a [`ModelStreamItem::MessageDelta`] per text + /// fragment, a [`ModelStreamItem::ToolCallDelta`] per tool-call argument + /// fragment, a [`ModelStreamItem::UsageDelta`] when usage arrives, and a + /// terminal [`ModelStreamItem::Completed`] carrying the fully merged + /// response (with reassembled tool-call names and ids). Transport errors + /// surface as a terminal [`ModelStreamItem::ProviderFailed`]. + /// + /// # Errors + /// + /// Returns [`TinyAgentsError::Model`] when the initial request fails or the + /// endpoint returns a non-2xx status (per-chunk transport errors are + /// surfaced as [`ModelStreamItem::ProviderFailed`] inside the stream + /// instead). + async fn stream(&self, _state: &State, request: ModelRequest) -> Result { + let mut body = self.translate_request(&request)?; + body.stream = true; + body.stream_options = Some(json!({ "include_usage": true })); + + let response = self + .post_json(&body, request.timeout_ms, true, "stream request") + .await?; + + // Map each raw byte chunk onto an owned `Vec` so the boxed stream's + // item type is nameable without depending on the `bytes` crate. + let bytes = response.bytes_stream().map(|chunk| { + chunk + .map(|b| b.to_vec()) + .map_err(|e| TinyAgentsError::Model(format!("stream chunk failed: {e}"))) + }); + + let state = SseState { + bytes: Box::pin(bytes), + buf: Vec::new(), + pending: VecDeque::new(), + acc: OpenAiStreamAcc::default(), + provider: self.provider.clone(), + model: self.model.clone(), + started: false, + finished: false, + terminal_emitted: false, + }; + + Ok(Box::pin(futures::stream::unfold(state, sse_next))) + } +} diff --git a/src/harness/providers/openai/types.rs b/src/harness/providers/openai/types.rs index 5aa5a50..66fdd3c 100644 --- a/src/harness/providers/openai/types.rs +++ b/src/harness/providers/openai/types.rs @@ -35,9 +35,15 @@ pub struct ChatCompletionRequest { /// Nucleus sampling probability mass. Omitted when unset. #[serde(skip_serializing_if = "Option::is_none")] pub top_p: Option, - /// Maximum number of output tokens. Omitted when unset. + /// Maximum number of output tokens. Omitted when unset. Used for classic + /// Chat Completions models; the o-series uses `max_completion_tokens` + /// instead. #[serde(skip_serializing_if = "Option::is_none")] pub max_tokens: Option, + /// Maximum number of output tokens for reasoning (o-series) models, which + /// reject `max_tokens`. Omitted when unset. + #[serde(skip_serializing_if = "Option::is_none")] + pub max_completion_tokens: Option, /// Stop sequences that terminate generation. Omitted when empty. #[serde(skip_serializing_if = "Vec::is_empty")] pub stop: Vec, @@ -97,6 +103,13 @@ pub struct ChunkDeltaWire { /// Incremental text fragment, when present. #[serde(default)] pub content: Option, + /// Incremental reasoning/thinking fragment used by some OpenAI-compatible + /// providers (for example DeepSeek-compatible streams). + #[serde(default)] + pub reasoning_content: Option, + /// Alternate reasoning/thinking fragment key used by some gateways. + #[serde(default)] + pub reasoning: Option, /// Incremental tool-call fragments, correlated by `index`. #[serde(default)] pub tool_calls: Vec, @@ -106,8 +119,12 @@ pub struct ChunkDeltaWire { #[derive(Clone, Debug, Default, Deserialize)] pub struct ToolCallChunkWire { /// Stable slot index used to correlate fragments across chunks. + /// + /// `None` when the provider omits `index` entirely (some OpenAI-compatible + /// backends do); fragments are then correlated by id/name instead of all + /// collapsing onto slot 0. #[serde(default)] - pub index: u32, + pub index: Option, /// Provider-assigned call id, sent on the first fragment for the slot. #[serde(default)] pub id: Option, @@ -127,15 +144,52 @@ pub struct FunctionChunkWire { pub arguments: Option, } +/// The `content` of a [`ChatMessageWire`]. +/// +/// OpenAI accepts either a plain string (the common case) or an array of typed +/// content parts (needed to attach images alongside text). Serialized untagged +/// so a text-only message keeps its historical plain-string wire shape. +#[derive(Clone, Debug, Serialize)] +#[serde(untagged)] +pub enum MessageContentWire { + /// Plain-string content. + Text(String), + /// Multi-part content (text and/or image parts). + Parts(Vec), +} + +/// One typed content part inside a multi-part [`MessageContentWire::Parts`]. +#[derive(Clone, Debug, Serialize)] +#[serde(tag = "type", rename_all = "snake_case")] +pub enum ContentPartWire { + /// A text fragment. + Text { + /// The text content. + text: String, + }, + /// An image reference, by URL or data URI. + ImageUrl { + /// The `image_url` object. + image_url: ImageUrlWire, + }, +} + +/// The `image_url` payload of a [`ContentPartWire::ImageUrl`]. +#[derive(Clone, Debug, Serialize)] +pub struct ImageUrlWire { + /// URL or data URI of the image. + pub url: String, +} + /// A single message in the request `messages` array. #[derive(Clone, Debug, Serialize)] pub struct ChatMessageWire { /// Role: `system`, `user`, `assistant`, or `tool`. pub role: String, - /// Textual content. `None` (serialized as absent) for assistant messages - /// that only carry tool calls. + /// Message content: a plain string or typed content parts. `None` + /// (serialized as absent) for assistant messages that only carry tool calls. #[serde(skip_serializing_if = "Option::is_none")] - pub content: Option, + pub content: Option, /// Tool calls requested by an assistant message. Omitted when empty. #[serde(skip_serializing_if = "Vec::is_empty")] pub tool_calls: Vec, @@ -225,6 +279,13 @@ pub struct ResponseMessageWire { /// Textual content, when present. #[serde(default)] pub content: Option, + /// Reasoning/thinking content on non-streaming responses, when a compatible + /// provider exposes it. Kept off visible assistant text by the translator. + #[serde(default)] + pub reasoning_content: Option, + /// Alternate reasoning/thinking field used by some gateways. + #[serde(default)] + pub reasoning: Option, /// Tool calls requested by the model. #[serde(default)] pub tool_calls: Vec, @@ -245,6 +306,9 @@ pub struct UsageWire { /// Optional input-token breakdown (carries cached tokens). #[serde(default)] pub prompt_tokens_details: Option, + /// Optional completion-token breakdown (carries reasoning tokens). + #[serde(default)] + pub completion_tokens_details: Option, } /// The `prompt_tokens_details` breakdown of a [`UsageWire`]. @@ -255,6 +319,14 @@ pub struct PromptTokensDetailsWire { pub cached_tokens: u64, } +/// The `completion_tokens_details` breakdown of a [`UsageWire`]. +#[derive(Clone, Debug, Default, Deserialize)] +pub struct CompletionTokensDetailsWire { + /// Output tokens spent on hidden reasoning/thinking. + #[serde(default)] + pub reasoning_tokens: u64, +} + // --------------------------------------------------------------------------- // Model-listing shapes (`GET {base_url}/models`) // --------------------------------------------------------------------------- diff --git a/src/harness/retry/mod.rs b/src/harness/retry/mod.rs index e8906d8..aea92c7 100644 --- a/src/harness/retry/mod.rs +++ b/src/harness/retry/mod.rs @@ -102,6 +102,36 @@ impl RetryPolicy { attempt + 1 < self.max_attempts } + /// The single retry decision shared by every retry loop in the harness. + /// + /// Returns `true` only when `error` is transient ([`is_retryable`]) *and* + /// the policy still permits another attempt ([`RetryPolicy::should_retry`]). + /// Both the agent loop and [`crate::harness::middleware::library::RetryMiddleware`] + /// route their per-attempt decision through here so the retry classification + /// and attempt-cap logic live in exactly one place and cannot drift apart. + /// + /// Callers that need to fold in a harness-level ceiling + /// ([`RetryPolicy::max_attempts_capped_at`]) should apply + /// [`RetryPolicy::with_max_attempts`] first and call this on the capped + /// policy. + pub fn should_retry_error(&self, attempt: usize, error: &TinyAgentsError) -> bool { + is_retryable(error) && self.should_retry(attempt) + } + + /// Reconciles this policy's own `max_attempts` with a harness-level + /// ceiling expressed as a *retry* count (not counting the first attempt) + /// — [`crate::harness::limits::RunLimits::max_retries_per_call`] — and + /// returns whichever total-attempt cap is stricter. + /// + /// Without this, a `RunPolicy` could configure a looser `RetryPolicy` + /// than its own `RunLimits`, silently making the "hard" limit + /// unenforceable; the harness's agent loop always calls this instead of + /// consulting `max_attempts` directly. + pub fn max_attempts_capped_at(&self, max_retries_per_call: usize) -> usize { + self.max_attempts + .min(max_retries_per_call.saturating_add(1)) + } + /// Computes the deterministic (no-jitter) backoff for the given retry /// `attempt`. /// @@ -143,7 +173,8 @@ impl RetryPolicy { /// /// | Variant | Retryable | Rationale | /// |---|---|---| -/// | `Model` | yes | Transient provider 5xx / rate-limit / network glitch. | +/// | `Provider` | depends | Classified from [`crate::harness::model::ProviderError::retryable`] — a 429/408/409/5xx is retryable, a 4xx like 401/400 is not. | +/// | `Model` | yes | No structured detail to classify from (transport/parse failure); transient provider 5xx / rate-limit / network glitch is the common case. | /// | `Tool` | yes | Tool execution may have hit a transient dependency. | /// | `Validation` | **no** | Caller-side schema or policy error; retrying will not help. | /// | `Serialization` | **no** | Malformed data; retrying will not help. | @@ -152,7 +183,11 @@ impl RetryPolicy { /// | `ToolNotFound` / `ModelNotFound` | **no** | Registry errors; not transient. | /// | `StructuredOutput` | **no** | Schema mismatch; retrying the same call will likely fail again. | pub fn is_retryable(err: &TinyAgentsError) -> bool { - matches!(err, TinyAgentsError::Model(_) | TinyAgentsError::Tool(_)) + match err { + TinyAgentsError::Provider(provider_error) => provider_error.retryable, + TinyAgentsError::Model(_) | TinyAgentsError::Tool(_) => true, + _ => false, + } } // ── FallbackPolicy ─────────────────────────────────────────────────────────── diff --git a/src/harness/retry/test.rs b/src/harness/retry/test.rs index c0fafcc..adfcd7d 100644 --- a/src/harness/retry/test.rs +++ b/src/harness/retry/test.rs @@ -99,6 +99,24 @@ fn should_retry_boundary_at_max_attempts() { assert!(!no_retry.should_retry(0)); } +// ── RetryPolicy::max_attempts_capped_at ─────────────────────────────────────── + +#[test] +fn max_attempts_capped_at_takes_the_stricter_of_the_two_caps() { + // A looser `RunLimits::max_retries_per_call` never widens the policy's + // own cap. + let policy = RetryPolicy::default().with_max_attempts(3); + assert_eq!(policy.max_attempts_capped_at(10), 3); + + // A stricter `max_retries_per_call` (a *retry* count, so +1 for the first + // attempt) overrides a looser policy. + assert_eq!(policy.max_attempts_capped_at(1), 2); + + // Zero retries permitted means exactly one attempt, same as + // `max_attempts == 1`. + assert_eq!(policy.max_attempts_capped_at(0), 1); +} + // ── is_retryable per error class ────────────────────────────────────────────── #[test] @@ -113,6 +131,41 @@ fn is_retryable_classification() { assert!(!is_retryable(&TinyAgentsError::Serialization(serde_err))); } +// ── TinyAgentsError::Provider classification ────────────────────────────────── + +#[test] +fn provider_error_retryability_is_read_from_the_structured_flag_not_assumed() { + use crate::harness::model::ProviderError; + + // Regression test: the unary/streaming provider path used to flatten a + // structured `ProviderError` into a plain `Model(String)`, so retry could + // not distinguish a retryable 429 from a non-retryable 401 and retried + // both. `TinyAgentsError::Provider` preserves the `retryable` flag a real + // provider adapter computes from the HTTP status, and `is_retryable` must + // consult it instead of assuming every provider failure is transient. + let rate_limited = ProviderError { + provider: "openai".to_string(), + status: Some(429), + retryable: true, + message: "rate limited".to_string(), + ..ProviderError::default() + }; + assert!(is_retryable(&TinyAgentsError::Provider(Box::new( + rate_limited + )))); + + let unauthorized = ProviderError { + provider: "openai".to_string(), + status: Some(401), + retryable: false, + message: "invalid api key".to_string(), + ..ProviderError::default() + }; + assert!(!is_retryable(&TinyAgentsError::Provider(Box::new( + unauthorized + )))); +} + // ── FallbackPolicy::next_after ──────────────────────────────────────────────── #[test] @@ -211,3 +264,20 @@ async fn sleep_backoff_waits_only_when_enabled() { assert_eq!(t1.elapsed(), expected); assert!(expected > Duration::ZERO); } + +#[test] +fn should_retry_error_combines_classification_and_attempt_cap() { + // 1 try + 2 retries: attempts 0 and 1 may retry, attempt 2 may not. + let policy = RetryPolicy::default().with_max_attempts(3); + + // Retryable error, attempts left → retry. + let retryable = TinyAgentsError::Model("5xx".into()); + assert!(policy.should_retry_error(0, &retryable)); + assert!(policy.should_retry_error(1, &retryable)); + // Retryable error, attempts exhausted → stop. + assert!(!policy.should_retry_error(2, &retryable)); + + // Non-retryable error is never retried regardless of remaining attempts. + let non_retryable = TinyAgentsError::Validation("bad".into()); + assert!(!policy.should_retry_error(0, &non_retryable)); +} diff --git a/src/harness/store/mod.rs b/src/harness/store/mod.rs index 4a5a50a..9c24610 100644 --- a/src/harness/store/mod.rs +++ b/src/harness/store/mod.rs @@ -36,6 +36,11 @@ use serde_json::Value; pub use types::*; use crate::error::{Result, TinyAgentsError}; +use crate::harness::ids::now_ms; + +/// Process-wide counter making [`FileStore`] temp-file names unique so +/// concurrent atomic writes to the same key never collide on their scratch file. +static TMP_COUNTER: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0); // ── InMemoryStore ───────────────────────────────────────────────────────────── @@ -164,8 +169,25 @@ impl Store for FileStore { .map_err(|e| TinyAgentsError::Validation(format!("store mkdir error: {e}")))?; let path = dir.join(format!("{key}.json")); let bytes = serde_json::to_vec_pretty(&value)?; - fs::write(&path, &bytes) + // Write to a uniquely named temp file in the same directory, then rename + // over the destination. Rename is atomic on POSIX/Windows for same-dir + // paths, so a reader never observes a partially written file and a crash + // mid-write leaves the previous value intact (as the type docs promise). + let tmp = dir.join(format!( + "{key}.json.tmp.{}.{}", + std::process::id(), + TMP_COUNTER.fetch_add(1, std::sync::atomic::Ordering::Relaxed) + )); + fs::write(&tmp, &bytes) .map_err(|e| TinyAgentsError::Validation(format!("store write error: {e}")))?; + if let Err(e) = fs::rename(&tmp, &path) { + // Best-effort cleanup of the temp file so a failed rename does not + // leak partial files into the namespace directory. + let _ = fs::remove_file(&tmp); + return Err(TinyAgentsError::Validation(format!( + "store rename error: {e}" + ))); + } Ok(()) } @@ -259,6 +281,7 @@ impl JsonlAppendStore { pub fn new(root_dir: impl Into) -> Self { Self { root_dir: root_dir.into(), + offsets: Default::default(), } } @@ -291,26 +314,58 @@ impl JsonlAppendStore { impl AppendStore for JsonlAppendStore { async fn append(&self, stream: &str, value: Value) -> Result { let path = self.stream_path(stream)?; - fs::create_dir_all(&self.root_dir) - .map_err(|e| TinyAgentsError::Validation(format!("append store mkdir error: {e}")))?; - // The offset is the count of existing lines; re-read to stay correct - // across separate store instances on the same directory. - let offset = Self::read_records(&path)?.len() as u64; - let record = StoreRecord { - offset, - value, - created_at_ms: now_ms(), + let root_dir = self.root_dir.clone(); + let offsets = Arc::clone(&self.offsets); + let stream = stream.to_string(); + + // The append is pure blocking file I/O. Run it off the async runtime so + // it never stalls a tokio worker (`spawn_blocking` when a runtime is + // present, inline otherwise — e.g. a synchronous sink draining outside a + // runtime). The offset cache means we only read the file once per stream + // instead of re-parsing the whole file on every append (previously + // O(n²) per stream). + let work = move || -> Result { + fs::create_dir_all(&root_dir).map_err(|e| { + TinyAgentsError::Validation(format!("append store mkdir error: {e}")) + })?; + // Hold the offset guard across the write so concurrent appends to the + // same store instance get distinct, ordered offsets. + let mut cache = offsets.lock().map_err(|e| { + TinyAgentsError::Validation(format!("append store lock poisoned: {e}")) + })?; + let offset = match cache.get(&stream) { + Some(&next) => next, + // First append for this stream in this instance: learn the length + // from disk once, then track it in memory. + None => Self::read_records(&path)?.len() as u64, + }; + let record = StoreRecord { + offset, + value, + created_at_ms: now_ms(), + }; + let mut line = serde_json::to_string(&record)?; + line.push('\n'); + let mut file = fs::OpenOptions::new() + .create(true) + .append(true) + .open(&path) + .map_err(|e| { + TinyAgentsError::Validation(format!("append store open error: {e}")) + })?; + std::io::Write::write_all(&mut file, line.as_bytes()).map_err(|e| { + TinyAgentsError::Validation(format!("append store write error: {e}")) + })?; + cache.insert(stream, offset + 1); + Ok(offset) }; - let mut line = serde_json::to_string(&record)?; - line.push('\n'); - let mut file = fs::OpenOptions::new() - .create(true) - .append(true) - .open(&path) - .map_err(|e| TinyAgentsError::Validation(format!("append store open error: {e}")))?; - std::io::Write::write_all(&mut file, line.as_bytes()) - .map_err(|e| TinyAgentsError::Validation(format!("append store write error: {e}")))?; - Ok(offset) + + match tokio::runtime::Handle::try_current() { + Ok(handle) => handle.spawn_blocking(work).await.map_err(|e| { + TinyAgentsError::Validation(format!("append store task error: {e}")) + })?, + Err(_) => work(), + } } async fn read_from(&self, stream: &str, offset: u64) -> Result> { @@ -328,15 +383,6 @@ impl AppendStore for JsonlAppendStore { } } -/// Returns the current time in Unix-epoch milliseconds, saturating at `0` for -/// clocks set before the epoch. -fn now_ms() -> u64 { - std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .map(|d| d.as_millis() as u64) - .unwrap_or(0) -} - // ── StoreRegistry ───────────────────────────────────────────────────────────── impl StoreRegistry { diff --git a/src/harness/store/test.rs b/src/harness/store/test.rs index 2d8fbb4..3ef3d08 100644 --- a/src/harness/store/test.rs +++ b/src/harness/store/test.rs @@ -301,6 +301,61 @@ async fn jsonl_round_trips_across_two_store_instances() { ); } +#[tokio::test] +async fn jsonl_append_offsets_stay_correct_under_many_appends() { + // Exercise the cached-offset path: after the first append learns the length + // from disk, later appends must keep incrementing without re-reading, and + // every record must be readable back in order. + let dir = TempDir::new("jsonl-many"); + let store = JsonlAppendStore::new(&dir.0); + + for n in 0..50u64 { + let offset = store.append("evts", json!({ "n": n })).await.unwrap(); + assert_eq!(offset, n, "append offsets must be dense and monotonic"); + } + assert_eq!(store.len("evts").await.unwrap(), 50); + + let all = store.read_from("evts", 0).await.unwrap(); + assert_eq!(all.len(), 50); + for (i, (offset, value)) in all.iter().enumerate() { + assert_eq!(*offset, i as u64); + assert_eq!(value, &json!({ "n": i as u64 })); + } +} + +#[tokio::test] +async fn file_store_put_is_atomic_and_leaves_no_temp_files() { + let dir = TempDir::new("atomic"); + let store = FileStore::new(&dir.0); + + // Overwrite the same key repeatedly; the final read must be a complete, + // well-formed value and no scratch (`.tmp`) files may be left behind. + for n in 0..5 { + store + .put( + "ns", + "key", + json!({ "value": n, "payload": "x".repeat(1024) }), + ) + .await + .unwrap(); + } + + let got = store.get("ns", "key").await.unwrap().unwrap(); + assert_eq!(got["value"], json!(4)); + assert_eq!(store.list("ns").await.unwrap(), vec!["key".to_string()]); + + let leftover: Vec<_> = std::fs::read_dir(dir.0.join("ns")) + .unwrap() + .map(|e| e.unwrap().file_name().to_string_lossy().into_owned()) + .filter(|name| name.contains(".tmp")) + .collect(); + assert!( + leftover.is_empty(), + "atomic put must not leave temp files: {leftover:?}" + ); +} + #[tokio::test] async fn jsonl_rejects_unsafe_stream_names() { let dir = TempDir::new("jsonl-sanitize"); diff --git a/src/harness/store/types.rs b/src/harness/store/types.rs index 6f43893..a201ee6 100644 --- a/src/harness/store/types.rs +++ b/src/harness/store/types.rs @@ -190,16 +190,25 @@ pub(crate) type AppendEntry = (u64, Value); /// /// # Concurrency /// Operations use blocking [`std::fs`] (no async-fs dependency is pulled in for -/// this local backend). Each call opens the file, performs its read or append, -/// and closes it; appends are done with `OpenOptions::append`, which is atomic -/// per line on POSIX for small writes, but no advisory lock is held. For -/// high-concurrency writers, funnel appends through a single writer task or -/// prefer a future server backend. The blocking calls are short-lived; callers -/// inside an async runtime accept that they briefly block the worker thread. -#[derive(Clone, Debug)] +/// this local backend), but `append` runs that I/O on a blocking thread +/// (`spawn_blocking`) when a tokio runtime is present so it never stalls an +/// async worker. Appends use `OpenOptions::append`, which is atomic per line on +/// POSIX for small writes, but no advisory lock is held. To avoid re-parsing the +/// whole file on every append, each store instance caches the next offset per +/// stream (see [`Self::offsets`]); this assumes a single writing process per +/// directory. For multiple concurrent writers, funnel appends through one store +/// instance (its offset guard serialises them) or prefer a server backend. +#[derive(Clone, Debug, Default)] pub struct JsonlAppendStore { /// The root directory under which `.jsonl` files live. pub(crate) root_dir: PathBuf, + /// Per-stream cache of the *next* offset to write, so an append does not + /// have to re-read and re-parse the whole file to learn its length (which + /// made appends O(n²) per stream). Initialised lazily from disk on the + /// first append for a stream and incremented in memory thereafter; the + /// guard is held across the write so concurrent appends to the same store + /// instance stay correctly ordered. Clones share the same cache. + pub(crate) offsets: Arc>>, } // ── StoreRegistry ──────────────────────────────────────────────────────────── diff --git a/src/harness/stream/test.rs b/src/harness/stream/test.rs index 6951735..53381b4 100644 --- a/src/harness/stream/test.rs +++ b/src/harness/stream/test.rs @@ -137,3 +137,32 @@ fn sink_peek_does_not_consume() { 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)); +} diff --git a/src/harness/stream/types.rs b/src/harness/stream/types.rs index 7aa6fb6..88c9c38 100644 --- a/src/harness/stream/types.rs +++ b/src/harness/stream/types.rs @@ -51,8 +51,16 @@ pub enum StreamMode { /// /// 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")] +#[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), diff --git a/src/harness/subagent/mod.rs b/src/harness/subagent/mod.rs index 47c82d1..a787b6c 100644 --- a/src/harness/subagent/mod.rs +++ b/src/harness/subagent/mod.rs @@ -149,10 +149,7 @@ impl SubAgent { max_turn_output_tokens: Option, ) -> Result { let max_depth = self.harness.policy().limits.max_depth; - let child_depth = parent_depth + 1; - if child_depth > max_depth { - return Err(TinyAgentsError::SubAgentDepth(max_depth)); - } + let child_depth = RunConfig::checked_child_depth(parent_depth, max_depth)?; let child_run_id = format!("{}-d{child_depth}", self.name); let mut config = RunConfig::new(child_run_id.clone()) .with_depth(child_depth) @@ -333,10 +330,7 @@ impl SubAgentSession { /// cap exactly as [`SubAgent::invoke`] does. fn child_config(&self) -> Result { let max_depth = self.subagent.harness.policy().limits.max_depth; - let child_depth = self.parent_depth + 1; - if child_depth > max_depth { - return Err(TinyAgentsError::SubAgentDepth(max_depth)); - } + let child_depth = RunConfig::checked_child_depth(self.parent_depth, max_depth)?; Ok(RunConfig::new(format!( "{}-t{}-d{child_depth}", self.subagent.name, self.turn @@ -557,11 +551,26 @@ where ) -> Result { let input = Self::extract_input(&call.arguments); let call_id = call.id; - let config = self.subagent.child_config( + // Route a depth-limit rejection through the same limit-to-tool-error + // conversion as a failure from `run_child` below, rather than letting + // it propagate raw and abort the whole parent run: a sub-agent that + // is simply too deep is a delegated-agent limit signal the parent + // orchestrator should see as a tool result, not a fatal error. + let config = match self.subagent.child_config( context.depth, context.thread_id.as_ref(), context.max_turn_output_tokens, - )?; + ) { + Ok(config) => config, + Err(error) => { + if let Some(result) = + Self::limit_result_for_parent(call_id, &self.tool_name, &error) + { + return Ok(result); + } + return Err(error); + } + }; let ctx = RunContext::new(config, Ctx::default()).with_events(context.events); let run = match self.subagent.run_child(state, ctx, input).await { Ok(run) => run, diff --git a/src/harness/subagent/test.rs b/src/harness/subagent/test.rs index c2b3ddc..db5e891 100644 --- a/src/harness/subagent/test.rs +++ b/src/harness/subagent/test.rs @@ -22,7 +22,7 @@ use crate::harness::providers::MockModel; use crate::harness::runtime::{AgentHarness, RunPolicy}; use crate::harness::subagent::{SubAgent, SubAgentSession, SubAgentTool}; use crate::harness::testkit::ScriptedModel; -use crate::harness::tool::{Tool, ToolCall, ToolResult, ToolSchema}; +use crate::harness::tool::{Tool, ToolCall, ToolExecutionContext, ToolResult, ToolSchema}; use crate::harness::usage::Usage; // ── Helpers ───────────────────────────────────────────────────────────────── @@ -220,6 +220,44 @@ async fn tool_path_enforces_depth_limit() { ); } +#[tokio::test] +async fn call_with_context_enforces_depth_limit_as_a_tool_result() { + // Regression test: `SubAgentTool::call_with_context` used to propagate a + // `child_config` depth-limit rejection raw via `?`, bypassing the + // limit-to-tool-error conversion `call` already applies and aborting the + // whole parent run instead of surfacing a tool-error the model can react + // to. This exercises the `call_with_context` path directly (the one the + // agent loop actually uses), not `call`. + let subagent = Arc::new(SubAgent::new( + "deep", + "a deep agent", + Arc::new(child_harness_with_max_depth("ok", 1)), + )); + let tool = SubAgentTool::new(subagent); + + // Caller (parent) depth 1 -> child depth 2 -> exceeds the cap of 1. + let parent_ctx: RunContext<()> = RunContext::new(RunConfig::new("parent").with_depth(1), ()); + let context = ToolExecutionContext::from_run_context(&parent_ctx); + + let result = tool + .call_with_context( + &(), + ToolCall::new("c1", "deep", json!({ "input": "too deep" })), + context, + ) + .await + .expect("depth limit is converted into a parent-visible tool result, not a fatal error"); + + assert!(result.is_error()); + assert_eq!(result.call_id, "c1"); + assert!( + result.content.contains("recursion depth limit") + && result.content.contains("delegated-agent limit signal"), + "unexpected depth limit message: {}", + result.content + ); +} + #[tokio::test] async fn subagent_tool_reports_child_limit_to_parent_as_tool_error() { let subagent = Arc::new(SubAgent::new( diff --git a/src/harness/summarization/mod.rs b/src/harness/summarization/mod.rs index a3001f8..3938d8e 100644 --- a/src/harness/summarization/mod.rs +++ b/src/harness/summarization/mod.rs @@ -47,10 +47,12 @@ pub fn estimate_tokens(text: &str) -> u64 { if chars == 0 { 0 } else { (chars / 4).max(1) } } -/// Estimate the total tokens for a [`Message`] by running [`estimate_tokens`] -/// over its concatenated text. +/// Estimate the total tokens for a [`Message`] using the same `chars / 4` +/// heuristic as [`estimate_tokens`], but counting characters directly over the +/// message's text blocks rather than allocating the concatenated string first. fn message_token_estimate(msg: &Message) -> u64 { - estimate_tokens(&msg.text()) + let chars = msg.char_len() as u64; + if chars == 0 { 0 } else { (chars / 4).max(1) } } /// Estimate the total tokens for a slice of messages. @@ -122,29 +124,37 @@ pub fn trim_messages(messages: &[Message], strategy: &TrimStrategy) -> Vec = non_system; - while !candidate.is_empty() { - let total = slice_token_estimate(&system) + slice_token_estimate(&candidate); - if total <= limit { - break; - } - candidate.remove(0); + // Precompute each message's token estimate once. The previous + // implementation re-summed the entire slice on every dropped + // message and used `remove(0)` (itself O(n)), making trimming + // O(n^2). Here we sum once and drop from the front by advancing an + // index while subtracting the running total. + let sys_tokens: Vec = system.iter().map(message_token_estimate).collect(); + let non_sys_tokens: Vec = non_system.iter().map(message_token_estimate).collect(); + let sys_total: u64 = sys_tokens.iter().sum(); + + // Drop non-system messages from the front until within budget or + // exhausted. + let mut non_sys_start = 0; + let mut non_sys_total: u64 = non_sys_tokens.iter().sum(); + while non_sys_start < non_system.len() && sys_total + non_sys_total > limit { + non_sys_total -= non_sys_tokens[non_sys_start]; + non_sys_start += 1; } - // If we're still over budget, start dropping system messages from - // the front as a last resort. - let mut sys_candidate = system; - while !sys_candidate.is_empty() { - let total = slice_token_estimate(&sys_candidate) + slice_token_estimate(&candidate); - if total <= limit { - break; - } - sys_candidate.remove(0); + // Still over budget: drop system messages from the front as a last + // resort. + let mut sys_start = 0; + let mut sys_running = sys_total; + while sys_start < system.len() && sys_running + non_sys_total > limit { + sys_running -= sys_tokens[sys_start]; + sys_start += 1; } - let mut result = sys_candidate; - result.extend(candidate); + let mut result = + Vec::with_capacity((system.len() - sys_start) + (non_system.len() - non_sys_start)); + result.extend_from_slice(&system[sys_start..]); + result.extend_from_slice(&non_system[non_sys_start..]); result } } diff --git a/src/harness/tool/mod.rs b/src/harness/tool/mod.rs index fba20d1..3e47577 100644 --- a/src/harness/tool/mod.rs +++ b/src/harness/tool/mod.rs @@ -260,20 +260,27 @@ fn validate_schema_value(schema: &Value, value: &Value, path: &str) -> Result<() validate_type_spec(type_spec, value, path)?; } - if let Some(properties) = schema.get("properties").and_then(Value::as_object) { + // Enforce `required` independently of `properties`. A schema may declare + // required fields without listing per-field property schemas; nesting this + // check under `properties` would let such schemas fail open, silently + // accepting calls that omit required arguments. + if let Some(required) = schema.get("required").and_then(Value::as_array) { let object = value.as_object().ok_or_else(|| { TinyAgentsError::Validation(format!("{path} must be an object with declared fields")) })?; - - if let Some(required) = schema.get("required").and_then(Value::as_array) { - for field in required.iter().filter_map(Value::as_str) { - if !object.contains_key(field) { - return Err(TinyAgentsError::Validation(format!( - "{path}.{field} is required" - ))); - } + for field in required.iter().filter_map(Value::as_str) { + if !object.contains_key(field) { + return Err(TinyAgentsError::Validation(format!( + "{path}.{field} is required" + ))); } } + } + + if let Some(properties) = schema.get("properties").and_then(Value::as_object) { + let object = value.as_object().ok_or_else(|| { + TinyAgentsError::Validation(format!("{path} must be an object with declared fields")) + })?; if schema.get("additionalProperties").and_then(Value::as_bool) == Some(false) { for field in object.keys() { diff --git a/src/harness/tool/test.rs b/src/harness/tool/test.rs index e53153c..9bdae17 100644 --- a/src/harness/tool/test.rs +++ b/src/harness/tool/test.rs @@ -204,6 +204,34 @@ fn schema_validation_rejects_missing_required_fields() { assert!(err.to_string().contains("user_id")); } +#[test] +fn schema_validation_rejects_missing_required_without_properties() { + // A schema may declare `required` without listing `properties`. The + // required check must still fail closed rather than silently accept a call + // that omits the required field. + let schema = ToolSchema::new( + "lookup", + "looks up a user", + json!({ + "type": "object", + "required": ["user_id"] + }), + ); + + let missing = ToolCall::new("c-1", "lookup", json!({})); + let err = schema.validate_call(&missing).expect_err("missing field"); + assert!(err.to_string().contains("user_id")); + + // A non-object argument for a required-bearing schema also fails closed. + let not_object = ToolCall::new("c-2", "lookup", json!("nope")); + schema + .validate_call(¬_object) + .expect_err("non-object arguments"); + + let present = ToolCall::new("c-3", "lookup", json!({ "user_id": "u-1" })); + schema.validate_call(&present).expect("valid call"); +} + #[test] fn schema_validation_rejects_wrong_types_and_extra_fields() { let schema = ToolSchema::new( diff --git a/src/harness/usage/mod.rs b/src/harness/usage/mod.rs index ae8f4d3..699eb89 100644 --- a/src/harness/usage/mod.rs +++ b/src/harness/usage/mod.rs @@ -49,9 +49,15 @@ impl Add for Usage { impl AddAssign for Usage { fn add_assign(&mut self, rhs: Usage) { + // `total_tokens` is only meaningful per-record via `effective_total()` + // (it falls back to `input + output` when the provider omitted it, in + // which case the stored `total_tokens` is a real `0`, not "unknown"). + // Summing the raw fields would silently drop that record's + // contribution, so combine the effective totals instead. + let combined_total = self.effective_total() + rhs.effective_total(); self.input_tokens += rhs.input_tokens; self.output_tokens += rhs.output_tokens; - self.total_tokens += rhs.total_tokens; + self.total_tokens = combined_total; self.cache_read_tokens += rhs.cache_read_tokens; self.cache_creation_tokens += rhs.cache_creation_tokens; self.reasoning_tokens += rhs.reasoning_tokens; diff --git a/src/harness/usage/test.rs b/src/harness/usage/test.rs index d0688d2..32669e3 100644 --- a/src/harness/usage/test.rs +++ b/src/harness/usage/test.rs @@ -49,6 +49,23 @@ fn effective_total_falls_back() { assert_eq!(usage.effective_total(), 10); } +#[test] +fn add_assign_does_not_lose_totals_when_one_side_omits_total_tokens() { + // A record whose provider omitted `total_tokens` (so it's a real `0`) + // combined with a normal record must still sum both records' effective + // totals, not just the raw `total_tokens` fields. + let mut a = Usage { + input_tokens: 100, + output_tokens: 50, + total_tokens: 0, + ..Usage::default() + }; + let b = Usage::new(10, 5); + a += b; + assert_eq!(a.effective_total(), 165); + assert_eq!(a.total_tokens, 165); +} + #[test] fn usage_totals_count_calls() { let mut totals = UsageTotals::new(); diff --git a/src/harness/workspace/mod.rs b/src/harness/workspace/mod.rs index 4e28738..2897557 100644 --- a/src/harness/workspace/mod.rs +++ b/src/harness/workspace/mod.rs @@ -9,6 +9,7 @@ //! Application-specific worktree/sandbox providers implement //! [`WorkspaceIsolation`] themselves. +mod policy; mod types; pub use types::*; diff --git a/src/harness/workspace/policy.rs b/src/harness/workspace/policy.rs new file mode 100644 index 0000000..1822732 --- /dev/null +++ b/src/harness/workspace/policy.rs @@ -0,0 +1,121 @@ +//! Path-gating policy for [`WorkspaceDescriptor`]: the `allows`/`enforce` +//! checks and the lexical path-normalization helpers they rely on. +//! +//! Split out of `workspace/types.rs`; kept separate from the plain type +//! definitions because this is where the fail-closed security guarantee +//! actually lives. + +use std::path::{Path, PathBuf}; + +use crate::Result; +use crate::harness::tool::SandboxMode; +use crate::harness::workspace::types::WorkspaceDescriptor; + +impl WorkspaceDescriptor { + /// Creates a descriptor rooted at `root` with no extra trusted roots. + pub fn new(root: impl Into) -> Self { + Self { + root: root.into(), + trusted_roots: Vec::new(), + policy_id: String::new(), + sandbox: SandboxMode::Inherit, + } + } + + /// Adds a trusted root the tool may also touch. + pub fn with_trusted_root(mut self, root: impl Into) -> Self { + self.trusted_roots.push(root.into()); + self + } + + /// Sets the audit policy identity. + pub fn with_policy_id(mut self, id: impl Into) -> Self { + self.policy_id = id.into(); + self + } + + /// Sets the sandbox mode. + pub fn with_sandbox(mut self, sandbox: SandboxMode) -> Self { + self.sandbox = sandbox; + self + } + + /// Returns `true` when `path` is contained within the root or any trusted + /// root. + /// + /// Comparison is lexical (after normalizing `.`/`..` components) so it does + /// not require the path to exist; it is a policy gate, not a canonicalizing + /// filesystem call. Relative candidates and roots are first anchored to the + /// current working directory so a relative path cannot use leading `..` + /// components to spoof re-entry into a same-named sibling of the root. If + /// the current directory cannot be read, the gate fails closed (`false`). + pub fn allows(&self, path: &Path) -> bool { + let Some(candidate) = anchored_normalize(path) else { + return false; + }; + std::iter::once(&self.root) + .chain(self.trusted_roots.iter()) + .filter_map(|root| anchored_normalize(root)) + .any(|root| candidate.starts_with(&root)) + } + + /// Fail-closed path gate to call *before* a tool touches `path`: when the + /// path is outside every allowed root, emits an + /// [`AgentEvent::WorkspaceViolation`][crate::harness::events::AgentEvent::WorkspaceViolation] + /// on `events` and returns a [`TinyAgentsError::Validation`] so the caller + /// blocks the operation. Returns `Ok(())` when the path is allowed. + pub fn enforce(&self, path: &Path, events: &crate::harness::events::EventSink) -> Result<()> { + if self.allows(path) { + return Ok(()); + } + let rendered = path.display().to_string(); + events.emit(crate::harness::events::AgentEvent::WorkspaceViolation { + path: rendered.clone(), + }); + Err(crate::error::TinyAgentsError::Validation(format!( + "path `{rendered}` is outside the allowed workspace roots" + ))) + } +} + +/// Anchors `path` to an absolute base (the current working directory when +/// relative) and lexically normalizes it. Returns `None` when a relative path +/// cannot be anchored because the current directory is unavailable, so callers +/// fail closed. +fn anchored_normalize(path: &Path) -> Option { + let absolute = if path.is_absolute() { + path.to_path_buf() + } else { + std::env::current_dir().ok()?.join(path) + }; + Some(normalize(&absolute)) +} + +/// Lexically normalizes a path by resolving `.` and `..` components without +/// touching the filesystem. +/// +/// A `..` only pops a preceding *named* segment; a `..` that would escape the +/// accumulated prefix (leading or after another `..`) is preserved rather than +/// discarded. Dropping such components would let a relative path like +/// `ws/../../ws/secret` collapse back onto `ws` and spoof re-entry into a +/// same-named sibling directory outside the workspace. +fn normalize(path: &Path) -> PathBuf { + use std::path::Component; + let mut out = PathBuf::new(); + for component in path.components() { + match component { + Component::ParentDir => match out.components().next_back() { + Some(Component::Normal(_)) => { + out.pop(); + } + Some(Component::RootDir | Component::Prefix(_)) => { + // At a filesystem root; `..` cannot go higher. + } + _ => out.push(Component::ParentDir), + }, + Component::CurDir => {} + other => out.push(other.as_os_str()), + } + } + out +} diff --git a/src/harness/workspace/test.rs b/src/harness/workspace/test.rs index 3d920d4..95709b8 100644 --- a/src/harness/workspace/test.rs +++ b/src/harness/workspace/test.rs @@ -21,6 +21,21 @@ fn descriptor_allows_paths_under_root_and_trusted_roots() { assert_eq!(ws.policy_id, "run-1"); } +#[test] +fn relative_root_rejects_leading_parent_sibling_escape() { + // Relative root: a candidate that walks up past the root and re-descends + // into a same-named sibling must be rejected, not admitted by a leading + // `..` collapsing back onto the root name. + let ws = WorkspaceDescriptor::new("ws"); + + assert!(ws.allows(Path::new("ws/src/main.rs"))); + // `ws/../../ws/secret` resolves outside the anchored `ws` root. + assert!(!ws.allows(Path::new("ws/../../ws/secret"))); + // A bare leading `..` escape is rejected. + assert!(!ws.allows(Path::new("../ws/secret"))); + assert!(!ws.allows(Path::new("../evil"))); +} + #[tokio::test] async fn shared_root_workspace_prepares_and_cleans_up() { let provider = SharedRootWorkspace::new("/work").with_sandbox(SandboxMode::Disabled); diff --git a/src/harness/workspace/types.rs b/src/harness/workspace/types.rs index 929ad5e..e6a1f1d 100644 --- a/src/harness/workspace/types.rs +++ b/src/harness/workspace/types.rs @@ -7,7 +7,7 @@ //! worktrees/sandboxes. TinyAgents does not own any concrete policy; it owns the //! interface so parallel agents can be isolated consistently. -use std::path::{Path, PathBuf}; +use std::path::PathBuf; use async_trait::async_trait; use serde::{Deserialize, Serialize}; @@ -36,84 +36,6 @@ pub struct WorkspaceDescriptor { pub sandbox: SandboxMode, } -impl WorkspaceDescriptor { - /// Creates a descriptor rooted at `root` with no extra trusted roots. - pub fn new(root: impl Into) -> Self { - Self { - root: root.into(), - trusted_roots: Vec::new(), - policy_id: String::new(), - sandbox: SandboxMode::Inherit, - } - } - - /// Adds a trusted root the tool may also touch. - pub fn with_trusted_root(mut self, root: impl Into) -> Self { - self.trusted_roots.push(root.into()); - self - } - - /// Sets the audit policy identity. - pub fn with_policy_id(mut self, id: impl Into) -> Self { - self.policy_id = id.into(); - self - } - - /// Sets the sandbox mode. - pub fn with_sandbox(mut self, sandbox: SandboxMode) -> Self { - self.sandbox = sandbox; - self - } - - /// Returns `true` when `path` is contained within the root or any trusted - /// root. - /// - /// Comparison is lexical (after normalizing `.`/`..` components) so it does - /// not require the path to exist; it is a policy gate, not a canonicalizing - /// filesystem call. - pub fn allows(&self, path: &Path) -> bool { - let candidate = normalize(path); - std::iter::once(&self.root) - .chain(self.trusted_roots.iter()) - .any(|root| candidate.starts_with(normalize(root))) - } - - /// Fail-closed path gate to call *before* a tool touches `path`: when the - /// path is outside every allowed root, emits an - /// [`AgentEvent::WorkspaceViolation`][crate::harness::events::AgentEvent::WorkspaceViolation] - /// on `events` and returns a [`TinyAgentsError::Validation`] so the caller - /// blocks the operation. Returns `Ok(())` when the path is allowed. - pub fn enforce(&self, path: &Path, events: &crate::harness::events::EventSink) -> Result<()> { - if self.allows(path) { - return Ok(()); - } - let rendered = path.display().to_string(); - events.emit(crate::harness::events::AgentEvent::WorkspaceViolation { - path: rendered.clone(), - }); - Err(crate::error::TinyAgentsError::Validation(format!( - "path `{rendered}` is outside the allowed workspace roots" - ))) - } -} - -/// Lexically normalizes a path by resolving `.` and `..` components without -/// touching the filesystem. -fn normalize(path: &Path) -> PathBuf { - use std::path::Component; - let mut out = PathBuf::new(); - for component in path.components() { - match component { - Component::ParentDir => { - out.pop(); - } - Component::CurDir => {} - other => out.push(other.as_os_str()), - } - } - out -} - /// Prepares and tears down per-agent execution environments. /// /// Implementations create a worktree/sandbox for one agent run and clean it up diff --git a/src/language/ast.rs b/src/language/ast.rs index c13d079..36eecb2 100644 --- a/src/language/ast.rs +++ b/src/language/ast.rs @@ -37,7 +37,11 @@ impl Literal { match self { Literal::Str(s) | Literal::Ident(s) => s.clone(), Literal::Num(n) => { - if n.fract() == 0.0 { + if n.fract() == 0.0 + && n.is_finite() + && *n >= i64::MIN as f64 + && *n <= i64::MAX as f64 + { format!("{}", *n as i64) } else { format!("{n}") diff --git a/src/language/capability_resolver.rs b/src/language/capability_resolver.rs new file mode 100644 index 0000000..e28c22d --- /dev/null +++ b/src/language/capability_resolver.rs @@ -0,0 +1,437 @@ +//! Capability binding: resolves and validates every capability reference +//! (model/tool/subgraph/router/reducer/agent/script) a [`Blueprint`] makes +//! against an allowlist, so declarative `.rag` source can only reach +//! capabilities Rust has already registered and allowed. +//! +//! Lives beside [`crate::language::resolver`] (the spanned-diagnostic +//! resolution path) since both implement the same binding policy; +//! [`crate::language::resolver::Resolver`] wraps a [`CapabilityResolver`] +//! internally. Split out of `compiler.rs`; see that module's doc comment +//! for how binding fits into the overall compile pipeline. + +use std::collections::HashSet; + +use crate::error::{Result, TinyAgentsError}; +use crate::language::types::Blueprint; +use crate::registry::CapabilityRegistry; +// =========================================================================== +// Capability binding +// =========================================================================== + +/// The node `kind` values the registry-backed binding path recognises. +/// +/// A `.rag` node may only declare one of these kinds when validated through +/// [`bind_capabilities_with_registry`] (or any resolver built with +/// [`CapabilityResolver::from_registry`]); an unknown kind is a +/// [`TinyAgentsError::Compile`] error. The set deliberately includes `model`, +/// because [`compile`] defaults an unspecified kind to `model`. +/// +/// The kinds carry the following capability-reference conventions, applied by +/// the strict binding path: +/// +/// - `subgraph` / `graph`: the node's `model` field (when present) names a +/// registered graph [`Blueprint`] — a *subgraph reference*. +/// - `router`: the node's `model` field names a registered router function. +/// - everything else (`agent`, `model`, `tool_executor`, `human`): the +/// `model` field names a registered chat model. +pub const DEFAULT_NODE_KINDS: &[&str] = &[ + "agent", + "model", + "tool_executor", + "subgraph", + "graph", + "subagent", + "repl_agent", + "router", + "interrupt", + "join", + "human", +]; + +/// An allowlist of capability names referenced by the expressive language. +/// +/// The expressive language can only *reference* capabilities by name; it can +/// never define them. [`bind_capabilities`] uses a resolver to ensure that +/// every referenced model and tool was already registered and allowed by Rust, +/// which is what makes agent-authored source safe to compile. +/// +/// A resolver holds five name allowlists — models, tools, subgraphs (graph +/// blueprints), routers, and reducers — plus an optional set of allowed node +/// `kind` values. The minimal [`new`](Self::new) / [`from_lists`](Self::from_lists) +/// constructors populate only models and tools and leave `node_kinds` empty, so +/// the legacy [`bind_capabilities`] gate keeps its original behaviour (model and +/// tool checks only). The richer checks — subgraph, router, and reducer +/// references plus node-kind validation — are opt-in through the +/// registry-backed path: [`from_registry`](Self::from_registry) and +/// [`bind_capabilities_with_registry`]. +#[derive(Clone, Debug, Default)] +pub struct CapabilityResolver { + models: HashSet, + tools: HashSet, + subgraphs: HashSet, + routers: HashSet, + reducers: HashSet, + /// Registered agent names (and aliases) a `subagent` node may reference. + agents: HashSet, + /// Registered REPL script names (and aliases) a `repl_agent` node may + /// reference. + scripts: HashSet, + /// Allowed node kinds. When empty, node-kind validation is skipped (the + /// legacy, manual behaviour); when non-empty, the strict binding path + /// rejects any node whose kind is not listed. + node_kinds: HashSet, +} + +/// The class of the primary, kind-specific reference a node carries. +/// +/// This is the shared vocabulary of [`CapabilityResolver::classify_reference`], +/// the one policy that maps a node `kind` to the reference that must resolve and +/// the allowlist it resolves against. Every binding gate — the compiler's +/// [`CapabilityResolver::bind_blueprint`] and both +/// [`crate::language::resolver::Resolver`] paths — routes through it so they +/// cannot drift. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum ReferenceClass { + /// A chat model reference (the `_` default and `router`'s pre-classification). + Model, + /// A subgraph (graph blueprint) reference. + Subgraph, + /// A router-function reference. + Router, + /// A sub-agent reference. + Agent, + /// A REPL script reference. + Script, +} + +impl ReferenceClass { + /// The lowercase noun used in "unknown {word}" diagnostics. + pub fn word(self) -> &'static str { + match self { + ReferenceClass::Model => "model", + ReferenceClass::Subgraph => "subgraph", + ReferenceClass::Router => "router", + ReferenceClass::Agent => "agent", + ReferenceClass::Script => "script", + } + } +} + +/// The primary reference a node carries, resolved by the shared policy. +#[derive(Clone, Copy, Debug)] +pub struct PrimaryReference<'a> { + /// Which allowlist the reference must resolve against. + pub class: ReferenceClass, + /// The referenced name. + pub target: &'a str, +} + +impl CapabilityResolver { + /// Creates an empty resolver that allows nothing. + pub fn new() -> Self { + Self::default() + } + + /// Builds a resolver from iterators of allowed model and tool names. + /// + /// Subgraph, router, and reducer allowlists are left empty and node-kind + /// validation is disabled; use [`from_registry`](Self::from_registry) for a + /// fully populated, registry-backed resolver. + pub fn from_lists(models: M, tools: T) -> Self + where + M: IntoIterator, + T: IntoIterator, + { + Self { + models: models.into_iter().collect(), + tools: tools.into_iter().collect(), + ..Self::default() + } + } + + /// Builds a fully populated resolver from a live [`CapabilityRegistry`]. + /// + /// Every registered model, tool, graph blueprint, router, and reducer name + /// — including their aliases — is added to the corresponding allowlist, and + /// the node-kind allowlist is seeded with [`DEFAULT_NODE_KINDS`]. The + /// resulting resolver therefore validates `.rag` source against exactly what + /// Rust has registered, including subgraph/router/reducer references and + /// node kinds, when used with [`CapabilityResolver::bind_blueprint`] or + /// [`bind_capabilities_with_registry`]. + pub fn from_registry(registry: &CapabilityRegistry) -> Self { + use crate::registry::ComponentKind; + + let collect = |kind| registry.names_including_aliases(kind).into_iter().collect(); + Self { + models: collect(ComponentKind::Model), + tools: collect(ComponentKind::Tool), + subgraphs: collect(ComponentKind::Graph), + routers: collect(ComponentKind::Router), + reducers: collect(ComponentKind::Reducer), + agents: collect(ComponentKind::Agent), + scripts: collect(ComponentKind::Script), + node_kinds: DEFAULT_NODE_KINDS.iter().map(|k| (*k).to_owned()).collect(), + } + } + + /// Allows an additional model name. Returns `self` for chaining. + pub fn allow_model(mut self, name: impl Into) -> Self { + self.models.insert(name.into()); + self + } + + /// Allows an additional tool name. Returns `self` for chaining. + pub fn allow_tool(mut self, name: impl Into) -> Self { + self.tools.insert(name.into()); + self + } + + /// Allows an additional subgraph (graph blueprint) name. Returns `self`. + pub fn allow_subgraph(mut self, name: impl Into) -> Self { + self.subgraphs.insert(name.into()); + self + } + + /// Allows an additional router name. Returns `self` for chaining. + pub fn allow_router(mut self, name: impl Into) -> Self { + self.routers.insert(name.into()); + self + } + + /// Allows an additional reducer name. Returns `self` for chaining. + pub fn allow_reducer(mut self, name: impl Into) -> Self { + self.reducers.insert(name.into()); + self + } + + /// Allows an additional agent name (for `subagent` nodes). Returns `self`. + pub fn allow_agent(mut self, name: impl Into) -> Self { + self.agents.insert(name.into()); + self + } + + /// Allows an additional REPL script name (for `repl_agent` nodes). Returns + /// `self`. + pub fn allow_script(mut self, name: impl Into) -> Self { + self.scripts.insert(name.into()); + self + } + + /// Replaces the set of allowed node kinds. Passing a non-empty set enables + /// node-kind validation in the strict binding path. Returns `self`. + pub fn with_node_kinds(mut self, kinds: I) -> Self + where + I: IntoIterator, + S: Into, + { + self.node_kinds = kinds.into_iter().map(Into::into).collect(); + self + } + + /// Returns true if `name` is an allowed model. + pub fn model_allowed(&self, name: &str) -> bool { + self.models.contains(name) + } + + /// Returns true if `name` is an allowed tool. + pub fn tool_allowed(&self, name: &str) -> bool { + self.tools.contains(name) + } + + /// Returns true if `name` is an allowed subgraph (graph blueprint). + pub fn subgraph_allowed(&self, name: &str) -> bool { + self.subgraphs.contains(name) + } + + /// Returns true if `name` is an allowed router. + pub fn router_allowed(&self, name: &str) -> bool { + self.routers.contains(name) + } + + /// Returns true if `name` is an allowed reducer. + pub fn reducer_allowed(&self, name: &str) -> bool { + self.reducers.contains(name) + } + + /// Returns true if `name` is an allowed agent (for `subagent` nodes). + pub fn agent_allowed(&self, name: &str) -> bool { + self.agents.contains(name) + } + + /// Returns true if `name` is an allowed REPL script (for `repl_agent` + /// nodes). + pub fn script_allowed(&self, name: &str) -> bool { + self.scripts.contains(name) + } + + /// The single kind-to-reference policy every binding gate shares. + /// + /// Given a node `kind` and the reference fields it carries, returns the + /// primary reference that must resolve and the allowlist class it resolves + /// against — or `None` when the node declares no primary reference. The + /// `subgraph` argument is the caller's already-resolved subgraph target + /// (the dedicated graph field falling back to the legacy `model` field). + /// + /// Centralising this mapping is what keeps + /// [`bind_blueprint`](Self::bind_blueprint) and both + /// [`crate::language::resolver::Resolver`] paths from drifting: a new node + /// kind or a changed reference convention is edited here once. + pub fn classify_reference<'a>( + kind: &str, + model: Option<&'a str>, + subgraph: Option<&'a str>, + agent: Option<&'a str>, + script: Option<&'a str>, + ) -> Option> { + let (class, target) = match kind { + "subgraph" | "graph" => (ReferenceClass::Subgraph, subgraph?), + "router" => (ReferenceClass::Router, model?), + "subagent" => (ReferenceClass::Agent, agent?), + "repl_agent" => (ReferenceClass::Script, script?), + // Unknown kinds fall through to a model check, mirroring the + // compiler default of an unspecified kind being `model`. + _ => (ReferenceClass::Model, model?), + }; + Some(PrimaryReference { class, target }) + } + + /// Returns true when `target` is allowed for the given reference `class`. + pub fn reference_allowed(&self, class: ReferenceClass, target: &str) -> bool { + match class { + ReferenceClass::Model => self.model_allowed(target), + ReferenceClass::Subgraph => self.subgraph_allowed(target), + ReferenceClass::Router => self.router_allowed(target), + ReferenceClass::Agent => self.agent_allowed(target), + ReferenceClass::Script => self.script_allowed(target), + } + } + + /// Returns true if `kind` is an allowed node kind, or if node-kind + /// validation is disabled (the allowlist is empty). + pub fn node_kind_allowed(&self, kind: &str) -> bool { + self.node_kinds.is_empty() || self.node_kinds.contains(kind) + } + + /// Runs the full, strict capability binding for `blueprint`. + /// + /// In addition to the model/tool checks performed by [`bind_capabilities`], + /// this validates, per the conventions documented on [`DEFAULT_NODE_KINDS`]: + /// + /// - each node `kind` is in the resolver's node-kind allowlist (a + /// [`TinyAgentsError::Compile`] error otherwise); + /// - `subgraph`/`graph` node references resolve to a registered subgraph, + /// `router` node references to a registered router, `subagent` node + /// references to a registered agent, `repl_agent` node references to a + /// registered script, and all other nodes' `model` references to a + /// registered model (via the shared [`classify_reference`](Self::classify_reference) policy); + /// - every `channel` reducer reference is registered. + /// + /// # Errors + /// + /// Returns [`TinyAgentsError::Compile`] for an unknown node kind, and + /// [`TinyAgentsError::Capability`] for the first unregistered model, tool, + /// subgraph, router, agent, script, or reducer reference. + pub fn bind_blueprint(&self, blueprint: &Blueprint) -> Result<()> { + for node in &blueprint.nodes { + if !self.node_kind_allowed(&node.kind) { + return Err(TinyAgentsError::Compile(format!( + "node `{}` has unknown kind `{}`", + node.name, node.kind + ))); + } + + // Prefer the dedicated `graph "name"` reference, falling back to the + // legacy `model` field for back-compatibility. + let subgraph_target = node.subgraph.as_deref().or(node.model.as_deref()); + if let Some(reference) = Self::classify_reference( + &node.kind, + node.model.as_deref(), + subgraph_target, + node.agent.as_deref(), + node.script.as_deref(), + ) && !self.reference_allowed(reference.class, reference.target) + { + return Err(TinyAgentsError::Capability(format!( + "node `{}` references unknown {} `{}`", + node.name, + reference.class.word(), + reference.target + ))); + } + + for tool in &node.tools { + if !self.tool_allowed(tool) { + return Err(TinyAgentsError::Capability(format!( + "node `{}` references unknown tool `{tool}`", + node.name + ))); + } + } + } + + for channel in &blueprint.channels { + if !self.reducer_allowed(&channel.reducer) { + return Err(TinyAgentsError::Capability(format!( + "channel `{}` references unknown reducer `{}`", + channel.name, channel.reducer + ))); + } + } + + Ok(()) + } +} + +/// Verifies that every model and tool referenced by `blueprint` is allowed by +/// `allow`. +/// +/// This is the minimal, manual gate: it checks only `model` and `tool` +/// references on each node and never inspects node kinds, subgraph/router +/// references, or channel reducers. For full registry-backed validation use +/// [`bind_capabilities_with_registry`]. +/// +/// # Errors +/// +/// Returns [`TinyAgentsError::Capability`] for the first model or tool +/// reference that is not present in the resolver's allowlist. +pub fn bind_capabilities(blueprint: &Blueprint, allow: &CapabilityResolver) -> Result<()> { + for node in &blueprint.nodes { + if let Some(model) = &node.model + && !allow.model_allowed(model) + { + return Err(TinyAgentsError::Capability(format!( + "node `{}` references unknown model `{model}`", + node.name + ))); + } + for tool in &node.tools { + if !allow.tool_allowed(tool) { + return Err(TinyAgentsError::Capability(format!( + "node `{}` references unknown tool `{tool}`", + node.name + ))); + } + } + } + Ok(()) +} + +/// Validates `blueprint` against a live [`CapabilityRegistry`]. +/// +/// This is the registry → language binding gate. It builds a fully populated +/// [`CapabilityResolver`] from `registry` (models, tools, subgraphs, routers, +/// reducers, and the default node kinds) and runs +/// [`CapabilityResolver::bind_blueprint`], so declarative source can only +/// reference capabilities that Rust has actually registered. +/// +/// # Errors +/// +/// Returns [`TinyAgentsError::Compile`] for an unknown node kind, and +/// [`TinyAgentsError::Capability`] for any unregistered model, tool, subgraph, +/// router, or reducer reference. +pub fn bind_capabilities_with_registry( + blueprint: &Blueprint, + registry: &CapabilityRegistry, +) -> Result<()> { + CapabilityResolver::from_registry(registry).bind_blueprint(blueprint) +} diff --git a/src/language/compiler.rs b/src/language/compiler.rs index f415100..a6ca752 100644 --- a/src/language/compiler.rs +++ b/src/language/compiler.rs @@ -29,6 +29,7 @@ use std::sync::Arc; use crate::error::{Result, TinyAgentsError}; use crate::graph::{CompiledGraph, GraphBuilder, NodeHandler}; +use crate::language::capability_resolver::bind_capabilities_with_registry; use crate::language::parser::parse_str; use crate::language::types::{ Blueprint, BlueprintProvenance, ChannelSpec, CommandSpec, END, EdgeSpan, EdgeSpec, IoFieldSpec, @@ -54,6 +55,15 @@ use crate::registry::CapabilityRegistry; /// /// All failures are reported as [`TinyAgentsError::Compile`]. pub fn compile(program: &Program) -> Result> { + let mut graph_ids: HashSet<&str> = HashSet::new(); + for graph in &program.graphs { + if !graph_ids.insert(graph.name.as_str()) { + return Err(TinyAgentsError::Compile(format!( + "duplicate graph `{}`", + graph.name + ))); + } + } program.graphs.iter().map(compile_graph).collect() } @@ -74,6 +84,17 @@ fn compile_graph(graph: &crate::language::types::GraphDecl) -> Result // A target is valid if it is a known node or the virtual `END`. let target_ok = |target: &str| target == END || node_names.contains(target); + // Reject duplicate channel declarations up front. + let mut channel_names: HashSet<&str> = HashSet::new(); + for channel in &graph.channels { + if !channel_names.insert(channel.name.as_str()) { + return Err(compile_err(format!( + "duplicate channel `{}` in graph `{}`", + channel.name, graph.name + ))); + } + } + // 2. Start node must be declared and defined. let start = graph .start @@ -111,12 +132,32 @@ fn compile_graph(graph: &crate::language::types::GraphDecl) -> Result let nodes_with_static_edge: HashSet<&str> = graph.edges.iter().map(|e| e.from.as_str()).collect(); + // Reject contradictory multiple top-level edges from the same source: only + // one static successor can ever be routed to, so a second edge from the + // same node is silent data loss (the first edge wins, the rest are dead + // weight in `blueprint.edges`) rather than a legitimate multi-successor. + for name in &nodes_with_static_edge { + let targets: Vec<&str> = graph + .edges + .iter() + .filter(|e| e.from == *name) + .map(|e| e.to.as_str()) + .collect(); + if targets.len() > 1 { + return Err(compile_err(format!( + "node `{name}` has multiple top-level edges ({}); a node may declare at most one outgoing edge", + targets.join(", ") + ))); + } + } + // 4. Validate and lower each node. let mut nodes = Vec::new(); for node in &graph.nodes { let has_routes = !node.routes.is_empty(); let has_next = node.next.is_some(); let has_static_edge = nodes_with_static_edge.contains(node.name.as_str()); + let has_command_goto = node.command.as_ref().is_some_and(|c| c.goto.is_some()); if has_routes && (has_next || has_static_edge) { return Err(compile_err(format!( @@ -125,6 +166,30 @@ fn compile_graph(graph: &crate::language::types::GraphDecl) -> Result ))); } + // A node may declare at most one of `routes`, `next`, `command { goto + // … }`, or a top-level edge as its routing source. Silently resolving + // by precedence hides a real authoring mistake (e.g. a model-authored + // revision that adds a `command.goto` without removing the old + // `next`), so any additional combination is a compile error. + let routing_sources = [ + (has_routes, "routes"), + (has_next, "`next`"), + (has_command_goto, "`command { goto … }`"), + (has_static_edge, "a top-level edge"), + ]; + let active: Vec<&str> = routing_sources + .iter() + .filter(|(present, _)| *present) + .map(|(_, label)| *label) + .collect(); + if active.len() > 1 { + return Err(compile_err(format!( + "node `{}` declares conflicting routing sources ({}); use exactly one", + node.name, + active.join(", ") + ))); + } + // Validate routes. let mut seen_labels: HashSet<&str> = HashSet::new(); for route in &node.routes { @@ -380,325 +445,6 @@ fn provenance_of( } } -// =========================================================================== -// Capability binding -// =========================================================================== - -/// The node `kind` values the registry-backed binding path recognises. -/// -/// A `.rag` node may only declare one of these kinds when validated through -/// [`bind_capabilities_with_registry`] (or any resolver built with -/// [`CapabilityResolver::from_registry`]); an unknown kind is a -/// [`TinyAgentsError::Compile`] error. The set deliberately includes `model`, -/// because [`compile`] defaults an unspecified kind to `model`. -/// -/// The kinds carry the following capability-reference conventions, applied by -/// the strict binding path: -/// -/// - `subgraph` / `graph`: the node's `model` field (when present) names a -/// registered graph [`Blueprint`] — a *subgraph reference*. -/// - `router`: the node's `model` field names a registered router function. -/// - everything else (`agent`, `model`, `tool_executor`, `human`): the -/// `model` field names a registered chat model. -pub const DEFAULT_NODE_KINDS: &[&str] = &[ - "agent", - "model", - "tool_executor", - "subgraph", - "graph", - "subagent", - "repl_agent", - "router", - "interrupt", - "join", - "human", -]; - -/// An allowlist of capability names referenced by the expressive language. -/// -/// The expressive language can only *reference* capabilities by name; it can -/// never define them. [`bind_capabilities`] uses a resolver to ensure that -/// every referenced model and tool was already registered and allowed by Rust, -/// which is what makes agent-authored source safe to compile. -/// -/// A resolver holds five name allowlists — models, tools, subgraphs (graph -/// blueprints), routers, and reducers — plus an optional set of allowed node -/// `kind` values. The minimal [`new`](Self::new) / [`from_lists`](Self::from_lists) -/// constructors populate only models and tools and leave `node_kinds` empty, so -/// the legacy [`bind_capabilities`] gate keeps its original behaviour (model and -/// tool checks only). The richer checks — subgraph, router, and reducer -/// references plus node-kind validation — are opt-in through the -/// registry-backed path: [`from_registry`](Self::from_registry) and -/// [`bind_capabilities_with_registry`]. -#[derive(Clone, Debug, Default)] -pub struct CapabilityResolver { - models: HashSet, - tools: HashSet, - subgraphs: HashSet, - routers: HashSet, - reducers: HashSet, - /// Allowed node kinds. When empty, node-kind validation is skipped (the - /// legacy, manual behaviour); when non-empty, the strict binding path - /// rejects any node whose kind is not listed. - node_kinds: HashSet, -} - -impl CapabilityResolver { - /// Creates an empty resolver that allows nothing. - pub fn new() -> Self { - Self::default() - } - - /// Builds a resolver from iterators of allowed model and tool names. - /// - /// Subgraph, router, and reducer allowlists are left empty and node-kind - /// validation is disabled; use [`from_registry`](Self::from_registry) for a - /// fully populated, registry-backed resolver. - pub fn from_lists(models: M, tools: T) -> Self - where - M: IntoIterator, - T: IntoIterator, - { - Self { - models: models.into_iter().collect(), - tools: tools.into_iter().collect(), - ..Self::default() - } - } - - /// Builds a fully populated resolver from a live [`CapabilityRegistry`]. - /// - /// Every registered model, tool, graph blueprint, router, and reducer name - /// — including their aliases — is added to the corresponding allowlist, and - /// the node-kind allowlist is seeded with [`DEFAULT_NODE_KINDS`]. The - /// resulting resolver therefore validates `.rag` source against exactly what - /// Rust has registered, including subgraph/router/reducer references and - /// node kinds, when used with [`CapabilityResolver::bind_blueprint`] or - /// [`bind_capabilities_with_registry`]. - pub fn from_registry(registry: &CapabilityRegistry) -> Self { - use crate::registry::ComponentKind; - - let collect = |kind| registry.names_including_aliases(kind).into_iter().collect(); - Self { - models: collect(ComponentKind::Model), - tools: collect(ComponentKind::Tool), - subgraphs: collect(ComponentKind::Graph), - routers: collect(ComponentKind::Router), - reducers: collect(ComponentKind::Reducer), - node_kinds: DEFAULT_NODE_KINDS.iter().map(|k| (*k).to_owned()).collect(), - } - } - - /// Allows an additional model name. Returns `self` for chaining. - pub fn allow_model(mut self, name: impl Into) -> Self { - self.models.insert(name.into()); - self - } - - /// Allows an additional tool name. Returns `self` for chaining. - pub fn allow_tool(mut self, name: impl Into) -> Self { - self.tools.insert(name.into()); - self - } - - /// Allows an additional subgraph (graph blueprint) name. Returns `self`. - pub fn allow_subgraph(mut self, name: impl Into) -> Self { - self.subgraphs.insert(name.into()); - self - } - - /// Allows an additional router name. Returns `self` for chaining. - pub fn allow_router(mut self, name: impl Into) -> Self { - self.routers.insert(name.into()); - self - } - - /// Allows an additional reducer name. Returns `self` for chaining. - pub fn allow_reducer(mut self, name: impl Into) -> Self { - self.reducers.insert(name.into()); - self - } - - /// Replaces the set of allowed node kinds. Passing a non-empty set enables - /// node-kind validation in the strict binding path. Returns `self`. - pub fn with_node_kinds(mut self, kinds: I) -> Self - where - I: IntoIterator, - S: Into, - { - self.node_kinds = kinds.into_iter().map(Into::into).collect(); - self - } - - /// Returns true if `name` is an allowed model. - pub fn model_allowed(&self, name: &str) -> bool { - self.models.contains(name) - } - - /// Returns true if `name` is an allowed tool. - pub fn tool_allowed(&self, name: &str) -> bool { - self.tools.contains(name) - } - - /// Returns true if `name` is an allowed subgraph (graph blueprint). - pub fn subgraph_allowed(&self, name: &str) -> bool { - self.subgraphs.contains(name) - } - - /// Returns true if `name` is an allowed router. - pub fn router_allowed(&self, name: &str) -> bool { - self.routers.contains(name) - } - - /// Returns true if `name` is an allowed reducer. - pub fn reducer_allowed(&self, name: &str) -> bool { - self.reducers.contains(name) - } - - /// Returns true if `kind` is an allowed node kind, or if node-kind - /// validation is disabled (the allowlist is empty). - pub fn node_kind_allowed(&self, kind: &str) -> bool { - self.node_kinds.is_empty() || self.node_kinds.contains(kind) - } - - /// Runs the full, strict capability binding for `blueprint`. - /// - /// In addition to the model/tool checks performed by [`bind_capabilities`], - /// this validates, per the conventions documented on [`DEFAULT_NODE_KINDS`]: - /// - /// - each node `kind` is in the resolver's node-kind allowlist (a - /// [`TinyAgentsError::Compile`] error otherwise); - /// - `subgraph`/`graph` node references resolve to a registered subgraph, - /// `router` node references to a registered router, and all other nodes' - /// `model` references to a registered model; - /// - every `channel` reducer reference is registered. - /// - /// # Errors - /// - /// Returns [`TinyAgentsError::Compile`] for an unknown node kind, and - /// [`TinyAgentsError::Capability`] for the first unregistered model, tool, - /// subgraph, router, or reducer reference. - pub fn bind_blueprint(&self, blueprint: &Blueprint) -> Result<()> { - for node in &blueprint.nodes { - if !self.node_kind_allowed(&node.kind) { - return Err(TinyAgentsError::Compile(format!( - "node `{}` has unknown kind `{}`", - node.name, node.kind - ))); - } - - match node.kind.as_str() { - "subgraph" | "graph" => { - // Prefer the dedicated `graph "name"` reference, falling - // back to the legacy `model` field for back-compatibility. - if let Some(target) = node.subgraph.as_ref().or(node.model.as_ref()) - && !self.subgraph_allowed(target) - { - return Err(TinyAgentsError::Capability(format!( - "node `{}` references unknown subgraph `{target}`", - node.name - ))); - } - } - "router" => { - if let Some(target) = &node.model - && !self.router_allowed(target) - { - return Err(TinyAgentsError::Capability(format!( - "node `{}` references unknown router `{target}`", - node.name - ))); - } - } - _ => { - if let Some(model) = &node.model - && !self.model_allowed(model) - { - return Err(TinyAgentsError::Capability(format!( - "node `{}` references unknown model `{model}`", - node.name - ))); - } - } - } - - for tool in &node.tools { - if !self.tool_allowed(tool) { - return Err(TinyAgentsError::Capability(format!( - "node `{}` references unknown tool `{tool}`", - node.name - ))); - } - } - } - - for channel in &blueprint.channels { - if !self.reducer_allowed(&channel.reducer) { - return Err(TinyAgentsError::Capability(format!( - "channel `{}` references unknown reducer `{}`", - channel.name, channel.reducer - ))); - } - } - - Ok(()) - } -} - -/// Verifies that every model and tool referenced by `blueprint` is allowed by -/// `allow`. -/// -/// This is the minimal, manual gate: it checks only `model` and `tool` -/// references on each node and never inspects node kinds, subgraph/router -/// references, or channel reducers. For full registry-backed validation use -/// [`bind_capabilities_with_registry`]. -/// -/// # Errors -/// -/// Returns [`TinyAgentsError::Capability`] for the first model or tool -/// reference that is not present in the resolver's allowlist. -pub fn bind_capabilities(blueprint: &Blueprint, allow: &CapabilityResolver) -> Result<()> { - for node in &blueprint.nodes { - if let Some(model) = &node.model - && !allow.model_allowed(model) - { - return Err(TinyAgentsError::Capability(format!( - "node `{}` references unknown model `{model}`", - node.name - ))); - } - for tool in &node.tools { - if !allow.tool_allowed(tool) { - return Err(TinyAgentsError::Capability(format!( - "node `{}` references unknown tool `{tool}`", - node.name - ))); - } - } - } - Ok(()) -} - -/// Validates `blueprint` against a live [`CapabilityRegistry`]. -/// -/// This is the registry → language binding gate. It builds a fully populated -/// [`CapabilityResolver`] from `registry` (models, tools, subgraphs, routers, -/// reducers, and the default node kinds) and runs -/// [`CapabilityResolver::bind_blueprint`], so declarative source can only -/// reference capabilities that Rust has actually registered. -/// -/// # Errors -/// -/// Returns [`TinyAgentsError::Compile`] for an unknown node kind, and -/// [`TinyAgentsError::Capability`] for any unregistered model, tool, subgraph, -/// router, or reducer reference. -pub fn bind_capabilities_with_registry( - blueprint: &Blueprint, - registry: &CapabilityRegistry, -) -> Result<()> { - CapabilityResolver::from_registry(registry).bind_blueprint(blueprint) -} - /// Parses, compiles, and registry-binds `.rag`/`.ragsh` `source` in one call. /// /// This is the convenience façade for the common path: it runs diff --git a/src/language/diagnostic.rs b/src/language/diagnostic.rs index 8271309..22e8a44 100644 --- a/src/language/diagnostic.rs +++ b/src/language/diagnostic.rs @@ -183,18 +183,23 @@ impl Diagnostic { /// Folds the diagnostic into a [`TinyAgentsError::Parse`]. /// - /// When `source` is provided the message is the full caret-underline - /// rendering and the `line`/`column` are resolved from the primary span's - /// byte offset; otherwise the message is the source-free rendering and the - /// span's own `line`/`column` anchor is used. Either way the existing - /// `Parse` variant's `line`/`column` fields are preserved. + /// When `source` is provided *and* the primary span carries real byte + /// offsets, the message is the full caret-underline rendering and the + /// `line`/`column` are resolved from that byte offset. Otherwise (no + /// `source`, or a back-compat span built with [`Span::new`] — which + /// anchors only a `line`/`column` and leaves `start`/`end` at `0`) the + /// message is the source-free rendering and the span's own stored + /// `line`/`column` is used directly: resolving byte offset `0` against a + /// real file would otherwise always render `1:1`, silently discarding + /// whatever real position the caller anchored the span at. pub fn into_parse_error(self, source: Option<&SourceFile>) -> TinyAgentsError { + let has_offsets = self.primary.start != 0 || self.primary.end != 0; let (line, column, message) = match source { - Some(file) => { + Some(file) if has_offsets => { let (line, column) = file.location(self.primary.start); (line, column, self.render(file)) } - None => (self.primary.line, self.primary.column, self.render_plain()), + _ => (self.primary.line, self.primary.column, self.render_plain()), }; TinyAgentsError::Parse { message, @@ -241,8 +246,13 @@ fn render_span_block( let _ = writeln!(out, "{line:>gutter$} | {line_text}"); // Caret width: the span's character length on this line, at least one. - let (line_start, line_end) = source.line_range(line).unwrap_or((span.start, span.start)); - let caret_start = span.start.max(line_start); + // Clamp the caret range into the line's byte range so a span that points + // past the end of the source (or past this line) neither trips `clamp`'s + // `min <= max` precondition nor slices out of bounds — the same defensive + // clamping `SourceFile::snippet` applies. + let text_len = source.text().len(); + let (line_start, line_end) = source.line_range(line).unwrap_or((text_len, text_len)); + let caret_start = span.start.clamp(line_start, line_end); let caret_end = span.end.clamp(caret_start, line_end); let caret_width = source.text()[caret_start..caret_end].chars().count().max(1); let indent = column.saturating_sub(1); diff --git a/src/language/diff.rs b/src/language/diff.rs index b70953c..ccba428 100644 --- a/src/language/diff.rs +++ b/src/language/diff.rs @@ -16,7 +16,10 @@ use std::fmt; use serde::{Deserialize, Serialize}; -use crate::language::types::{Blueprint, ChannelSpec, EdgeSpec, NodeSpec, Routing}; +use crate::language::types::{ + Blueprint, ChannelSpec, CommandSpec, EdgeSpec, IoFieldSpec, JoinSpec, Literal, NodeSpec, + Routing, SendSpec, +}; /// A field of a node whose value changed between two blueprints. #[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] @@ -86,6 +89,24 @@ pub struct BlueprintDiff { /// Static edges present only in the old blueprint. #[serde(default, skip_serializing_if = "Vec::is_empty")] pub edges_removed: Vec, + /// `(old, new)` when the graph-level `defaults` entries changed. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub defaults_changed: Option<(String, String)>, + /// `(old, new)` when the declared graph `input` shape changed. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub input_changed: Option<(String, String)>, + /// `(old, new)` when the declared graph `output` shape changed. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub output_changed: Option<(String, String)>, + /// `(old, new)` when the graph-level checkpoint policy changed. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub checkpoint_changed: Option<(String, String)>, + /// `(old, new)` when the graph-level interrupt policy changed. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub interrupt_changed: Option<(String, String)>, + /// `(old, new)` when the compiled join/barrier declarations changed. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub joins_changed: Option<(String, String)>, } impl BlueprintDiff { @@ -102,6 +123,12 @@ impl BlueprintDiff { && self.channels_changed.is_empty() && self.edges_added.is_empty() && self.edges_removed.is_empty() + && self.defaults_changed.is_none() + && self.input_changed.is_none() + && self.output_changed.is_none() + && self.checkpoint_changed.is_none() + && self.interrupt_changed.is_none() + && self.joins_changed.is_none() } } @@ -116,6 +143,24 @@ impl fmt::Display for BlueprintDiff { if let Some((old, new)) = &self.start_changed { writeln!(f, "~ start: {old} -> {new}")?; } + if let Some((old, new)) = &self.defaults_changed { + writeln!(f, "~ defaults: {old} -> {new}")?; + } + if let Some((old, new)) = &self.input_changed { + writeln!(f, "~ input: {old} -> {new}")?; + } + if let Some((old, new)) = &self.output_changed { + writeln!(f, "~ output: {old} -> {new}")?; + } + if let Some((old, new)) = &self.checkpoint_changed { + writeln!(f, "~ checkpoint: {old} -> {new}")?; + } + if let Some((old, new)) = &self.interrupt_changed { + writeln!(f, "~ interrupt: {old} -> {new}")?; + } + if let Some((old, new)) = &self.joins_changed { + writeln!(f, "~ joins: {old} -> {new}")?; + } for name in &self.nodes_added { writeln!(f, "+ node {name}")?; } @@ -167,6 +212,42 @@ pub fn blueprint_diff(old: &Blueprint, new: &Blueprint) -> BlueprintDiff { diff.start_changed = Some((old.start.clone(), new.start.clone())); } + let old_defaults = render_kv_list(&old.defaults); + let new_defaults = render_kv_list(&new.defaults); + if old_defaults != new_defaults { + diff.defaults_changed = Some((old_defaults, new_defaults)); + } + + let old_input = render_io_fields(&old.input); + let new_input = render_io_fields(&new.input); + if old_input != new_input { + diff.input_changed = Some((old_input, new_input)); + } + + let old_output = render_io_fields(&old.output); + let new_output = render_io_fields(&new.output); + if old_output != new_output { + diff.output_changed = Some((old_output, new_output)); + } + + let old_checkpoint = render_opt(&old.checkpoint); + let new_checkpoint = render_opt(&new.checkpoint); + if old_checkpoint != new_checkpoint { + diff.checkpoint_changed = Some((old_checkpoint, new_checkpoint)); + } + + let old_interrupt = render_opt(&old.interrupt); + let new_interrupt = render_opt(&new.interrupt); + if old_interrupt != new_interrupt { + diff.interrupt_changed = Some((old_interrupt, new_interrupt)); + } + + let old_joins = render_joins(&old.joins); + let new_joins = render_joins(&new.joins); + if old_joins != new_joins { + diff.joins_changed = Some((old_joins, new_joins)); + } + // Nodes. for node in &new.nodes { match old.nodes.iter().find(|n| n.name == node.name) { @@ -284,6 +365,22 @@ fn node_field_changes(old: &NodeSpec, new: &NodeSpec) -> Vec { render_opt(&old.timeout), render_opt(&new.timeout), ); + push( + "command", + render_command(&old.command), + render_command(&new.command), + ); + push("sends", render_sends(&old.sends), render_sends(&new.sends)); + push( + "retry", + render_kv_list(&old.retry), + render_kv_list(&new.retry), + ); + push( + "metadata", + render_kv_list(&old.metadata), + render_kv_list(&new.metadata), + ); changes } @@ -328,3 +425,58 @@ fn render_opt(value: &Option) -> String { fn render_list(values: &[String]) -> String { format!("[{}]", values.join(", ")) } + +/// Renders a `command { goto … update { … } }` declaration (`"(none)"` when absent). +fn render_command(command: &Option) -> String { + match command { + None => "(none)".to_string(), + Some(cmd) => { + let goto = cmd.goto.clone().unwrap_or_else(|| "(none)".to_string()); + let update = render_kv_list(&cmd.update); + format!("{{ goto {goto}, update {update} }}") + } + } +} + +/// Renders fanout `send` declarations. +fn render_sends(sends: &[SendSpec]) -> String { + let body = sends + .iter() + .map(|s| match &s.input { + Some(input) => format!("send {} {input}", s.target), + None => format!("send {}", s.target), + }) + .collect::>() + .join(", "); + format!("[{body}]") +} + +/// Renders a `(key, Literal)` list (used for `defaults`, `retry`, `metadata`). +fn render_kv_list(entries: &[(String, Literal)]) -> String { + let body = entries + .iter() + .map(|(k, v)| format!("{k} {}", v.as_display())) + .collect::>() + .join(", "); + format!("{{ {body} }}") +} + +/// Renders a graph `input`/`output` shape. +fn render_io_fields(fields: &[IoFieldSpec]) -> String { + let body = fields + .iter() + .map(|f| format!("{} {}", f.name, f.ty)) + .collect::>() + .join(", "); + format!("{{ {body} }}") +} + +/// Renders compiled join/barrier declarations. +fn render_joins(joins: &[JoinSpec]) -> String { + let body = joins + .iter() + .map(|j| format!("[{}] -> {}", j.sources.join(", "), j.target)) + .collect::>() + .join(", "); + format!("[{body}]") +} diff --git a/src/language/mod.rs b/src/language/mod.rs index 61c00ea..f17a642 100644 --- a/src/language/mod.rs +++ b/src/language/mod.rs @@ -48,6 +48,7 @@ pub mod source; pub mod span; pub mod types; +pub mod capability_resolver; pub mod compiler; pub mod diff; pub mod lexer; diff --git a/src/language/parser.rs b/src/language/parser.rs index 8611f75..4a0edbd 100644 --- a/src/language/parser.rs +++ b/src/language/parser.rs @@ -47,8 +47,18 @@ pub fn parse_str(source: &str) -> Result { /// # Errors /// /// Returns [`TinyAgentsError::Parse`] when the token stream does not match the -/// grammar, with the span of the offending token. +/// grammar, with the span of the offending token. An empty token slice (which +/// violates the lexer contract that every stream ends with an `Eof` sentinel) +/// is rejected as a parse error rather than panicking. pub fn parse(tokens: &[SpannedToken]) -> Result { + if tokens.is_empty() { + return Err(Diagnostic::error( + "empty token stream: expected at least an end-of-input token", + Span::new(1, 1), + ) + .with_primary_label("here") + .into_parse_error(None)); + } Parser { tokens, pos: 0, diff --git a/src/language/resolver.rs b/src/language/resolver.rs index ba6ec19..9e57b9b 100644 --- a/src/language/resolver.rs +++ b/src/language/resolver.rs @@ -31,11 +31,12 @@ //! call, so model-generated source is validated on exactly the same path as a //! checked-in `.rag` file. -use std::collections::HashSet; - use crate::error::{Result, TinyAgentsError}; use crate::language::ast::{ChannelDecl, GraphDecl, NodeDecl, Program}; -use crate::language::compiler::{CapabilityResolver, DEFAULT_NODE_KINDS, compile}; +use crate::language::capability_resolver::{ + CapabilityResolver, DEFAULT_NODE_KINDS, ReferenceClass, +}; +use crate::language::compiler::compile; use crate::language::diagnostic::Diagnostic; use crate::language::parser::parse_str; use crate::language::source::SourceFile; @@ -49,6 +50,7 @@ const CODE_UNKNOWN_TOOL: &str = "E-rag-unknown-tool"; const CODE_UNKNOWN_SUBGRAPH: &str = "E-rag-unknown-subgraph"; const CODE_UNKNOWN_ROUTER: &str = "E-rag-unknown-router"; const CODE_UNKNOWN_AGENT: &str = "E-rag-unknown-agent"; +const CODE_UNKNOWN_SCRIPT: &str = "E-rag-unknown-script"; const CODE_UNKNOWN_REDUCER: &str = "E-rag-unknown-reducer"; const CODE_INVALID_NODE_KIND: &str = "E-rag-invalid-node-kind"; @@ -62,42 +64,30 @@ const CODE_INVALID_NODE_KIND: &str = "E-rag-invalid-node-kind"; /// plan; it only reports references that fall outside the allowlists. #[derive(Clone, Debug)] pub struct Resolver { - /// The overlapping model/tool/subgraph/router/reducer/node-kind allowlists, - /// reused from the compiler's [`CapabilityResolver`] so the two share one - /// policy. + /// The overlapping model/tool/subgraph/router/reducer/agent/script/node-kind + /// allowlists, reused from the compiler's [`CapabilityResolver`] so every + /// binding gate shares one policy. caps: CapabilityResolver, - /// Registered agent names (and aliases) a `subagent` node may reference. - agents: HashSet, } impl Resolver { /// Builds a resolver from a live [`CapabilityRegistry`]. /// - /// Every registered model, tool, graph blueprint, router, reducer, and agent - /// name — including aliases — populates the corresponding allowlist, and the - /// node-kind allowlist is seeded with [`DEFAULT_NODE_KINDS`]. The resolver - /// therefore validates `.rag` source against exactly what Rust has - /// registered. + /// Every registered model, tool, graph blueprint, router, reducer, agent, + /// and script name — including aliases — populates the corresponding + /// allowlist, and the node-kind allowlist is seeded with + /// [`DEFAULT_NODE_KINDS`]. The resolver therefore validates `.rag` source + /// against exactly what Rust has registered. pub fn from_registry(registry: &CapabilityRegistry) -> Self { - use crate::registry::ComponentKind; - let agents = registry - .names_including_aliases(ComponentKind::Agent) - .into_iter() - .collect(); Self { caps: CapabilityResolver::from_registry(registry), - agents, } } - /// Builds a resolver from an existing [`CapabilityResolver`] allowlist, with - /// no extra agent names. Node-kind validation follows the supplied - /// resolver's configuration. + /// Builds a resolver from an existing [`CapabilityResolver`] allowlist. + /// Node-kind validation follows the supplied resolver's configuration. pub fn from_capabilities(caps: CapabilityResolver) -> Self { - Self { - caps, - agents: HashSet::new(), - } + Self { caps } } /// Returns the underlying capability allowlist. @@ -107,13 +97,13 @@ impl Resolver { /// Allows an additional agent name. Returns `self` for chaining. pub fn allow_agent(mut self, name: impl Into) -> Self { - self.agents.insert(name.into()); + self.caps = self.caps.allow_agent(name); self } /// Returns true if `name` is a registered/allowed agent. pub fn agent_allowed(&self, name: &str) -> bool { - self.agents.contains(name) + self.caps.agent_allowed(name) } /// Resolves every reference in `program` against the allowlists, returning a @@ -158,60 +148,27 @@ impl Resolver { // falls through to a model check, which is the compiler default. } - // 2. The kind-specific primary reference. - match kind { - "subgraph" | "graph" => { - if let Some(target) = node.graph.as_deref().or(node.model.as_deref()) { - self.check_ref( - self.caps.subgraph_allowed(target), - &node.name, - "subgraph", - target, - node.span, - CODE_UNKNOWN_SUBGRAPH, - out, - ); - } - } - "router" => { - if let Some(target) = node.model.as_deref() { - self.check_ref( - self.caps.router_allowed(target), - &node.name, - "router", - target, - node.span, - CODE_UNKNOWN_ROUTER, - out, - ); - } - } - "subagent" => { - if let Some(target) = node.agent.as_deref() { - self.check_ref( - self.agent_allowed(target), - &node.name, - "agent", - target, - node.span, - CODE_UNKNOWN_AGENT, - out, - ); - } - } - _ => { - if let Some(target) = node.model.as_deref() { - self.check_ref( - self.caps.model_allowed(target), - &node.name, - "model", - target, - node.span, - CODE_UNKNOWN_MODEL, - out, - ); - } - } + // 2. The kind-specific primary reference, routed through the one shared + // classification policy so this path cannot drift from the blueprint + // gates. + let subgraph_target = node.graph.as_deref().or(node.model.as_deref()); + if let Some(reference) = CapabilityResolver::classify_reference( + kind, + node.model.as_deref(), + subgraph_target, + node.agent.as_deref(), + node.script.as_deref(), + ) { + self.check_ref( + self.caps + .reference_allowed(reference.class, reference.target), + &node.name, + reference.class.word(), + reference.target, + node.span, + code_for(reference.class), + out, + ); } // 3. Every referenced tool must be registered. @@ -314,35 +271,22 @@ impl Resolver { node.name, node.kind ))); } - match node.kind.as_str() { - "subgraph" | "graph" => { - if let Some(target) = node.subgraph.as_deref().or(node.model.as_deref()) - && !self.caps.subgraph_allowed(target) - { - return Err(unregistered("subgraph", &node.name, target)); - } - } - "router" => { - if let Some(target) = node.model.as_deref() - && !self.caps.router_allowed(target) - { - return Err(unregistered("router", &node.name, target)); - } - } - "subagent" => { - if let Some(target) = node.agent.as_deref() - && !self.agent_allowed(target) - { - return Err(unregistered("agent", &node.name, target)); - } - } - _ => { - if let Some(target) = node.model.as_deref() - && !self.caps.model_allowed(target) - { - return Err(unregistered("model", &node.name, target)); - } - } + let subgraph_target = node.subgraph.as_deref().or(node.model.as_deref()); + if let Some(reference) = CapabilityResolver::classify_reference( + &node.kind, + node.model.as_deref(), + subgraph_target, + node.agent.as_deref(), + node.script.as_deref(), + ) && !self + .caps + .reference_allowed(reference.class, reference.target) + { + return Err(unregistered( + reference.class.word(), + &node.name, + reference.target, + )); } for tool in &node.tools { if !self.caps.tool_allowed(tool) { @@ -383,6 +327,17 @@ fn fold_diagnostic(diagnostic: Diagnostic, source: Option<&SourceFile>) -> TinyA /// Builds the span-less "unknown {what}" [`TinyAgentsError::Capability`] used by /// [`Resolver::resolve_blueprint`]. +/// Maps a shared [`ReferenceClass`] to its stable spanned-diagnostic code. +fn code_for(class: ReferenceClass) -> &'static str { + match class { + ReferenceClass::Model => CODE_UNKNOWN_MODEL, + ReferenceClass::Subgraph => CODE_UNKNOWN_SUBGRAPH, + ReferenceClass::Router => CODE_UNKNOWN_ROUTER, + ReferenceClass::Agent => CODE_UNKNOWN_AGENT, + ReferenceClass::Script => CODE_UNKNOWN_SCRIPT, + } +} + fn unregistered(what: &str, node: &str, target: &str) -> TinyAgentsError { TinyAgentsError::Capability(format!( "node `{node}` references unknown {what} `{target}`" diff --git a/src/language/test.rs b/src/language/test.rs deleted file mode 100644 index 403e385..0000000 --- a/src/language/test.rs +++ /dev/null @@ -1,1321 +0,0 @@ -//! Tests for the expressive language pipeline: lexer, parser, compiler, -//! capability binding, and graph materialisation. - -use std::sync::Arc; - -use crate::graph::{Command, NodeContext, NodeFuture, NodeResult}; -use crate::language::compiler::{ - BoxedNode, CapabilityResolver, NodeFactory, bind_capabilities, build_graph, compile, -}; -use crate::language::lexer::tokenize; -use crate::language::parser::{parse, parse_str}; -use crate::language::types::{Literal, NodeSpec, Routing, Token}; - -/// The `support_agent` fixture from the module spec: an agent node with a tool -/// loop plus conditional routing to `END`. -const SUPPORT_AGENT: &str = r#" -// A support workflow with a tool loop. -graph support_agent { - start agent - - defaults { - recursion_limit 50 - backoff "exponential" - checkpoint inherit - } - - channel messages messages - channel tool_calls append - - node agent { - kind agent - model "default" - system "Resolve support requests using tools when useful." - tools ["lookup_user", "create_ticket"] - routes { - tool_call -> tools - final -> END - } - } - - node tools { - kind tool_executor - next agent - } -} -"#; - -// --------------------------------------------------------------------------- -// Lexer -// --------------------------------------------------------------------------- - -#[test] -fn tokenizes_punctuation_and_arrow() { - let tokens = tokenize("a -> b { } [ ] ,").unwrap(); - let kinds: Vec<_> = tokens.into_iter().map(|t| t.token).collect(); - assert_eq!( - kinds, - vec![ - Token::Ident("a".into()), - Token::Arrow, - Token::Ident("b".into()), - Token::LBrace, - Token::RBrace, - Token::LBracket, - Token::RBracket, - Token::Comma, - Token::Eof, - ] - ); -} - -#[test] -fn tokenizes_strings_numbers_and_comments() { - let tokens = tokenize("// comment\n\"hi\\n\" 50 1.5 -3").unwrap(); - let kinds: Vec<_> = tokens.into_iter().map(|t| t.token).collect(); - assert_eq!( - kinds, - vec![ - Token::Str("hi\n".into()), - Token::Num(50.0), - Token::Num(1.5), - Token::Num(-3.0), - Token::Eof, - ] - ); -} - -#[test] -fn tracks_line_and_column_spans() { - let tokens = tokenize("graph\n foo").unwrap(); - assert_eq!(tokens[0].span.line, 1); - assert_eq!(tokens[0].span.column, 1); - assert_eq!(tokens[1].span.line, 2); - assert_eq!(tokens[1].span.column, 3); -} - -#[test] -fn unterminated_string_is_a_parse_error() { - let err = tokenize("\"oops").unwrap_err(); - assert!(matches!(err, crate::error::TinyAgentsError::Parse { .. })); -} - -#[test] -fn invalid_escape_is_a_parse_error() { - let err = tokenize("\"bad\\x\"").unwrap_err(); - assert!(matches!(err, crate::error::TinyAgentsError::Parse { .. })); -} - -// --------------------------------------------------------------------------- -// Spans, source map, and diagnostics -// --------------------------------------------------------------------------- - -use crate::language::diagnostic::{Diagnostic, Severity}; -use crate::language::source::{SourceFile, SourceMap}; -use crate::language::span::Span; - -#[test] -fn span_merge_covers_both_inputs() { - let a = Span::at(2, 5, 1, 3); - let b = Span::at(10, 14, 2, 1); - let merged = a.merge(b); - assert_eq!(merged.start, 2); - assert_eq!(merged.end, 14); - // Anchor comes from the earlier-starting span. - assert_eq!((merged.line, merged.column), (1, 3)); - // Merge is commutative over the covered range. - assert_eq!(b.merge(a).start, 2); - assert_eq!(b.merge(a).end, 14); -} - -#[test] -fn span_len_and_is_empty() { - assert!(Span::new(1, 1).is_empty()); - let s = Span::at(4, 9, 1, 5); - assert_eq!(s.len(), 5); - assert!(!s.is_empty()); -} - -#[test] -fn source_file_maps_offsets_to_line_and_column() { - let file = SourceFile::new("demo.rag", "graph g\n node a\n"); - // `g` is on line 1. - assert_eq!(file.location(6), (1, 7)); - // The `node` keyword starts at byte 10 on line 2, column 3. - let node_byte = file.text().find("node").unwrap(); - assert_eq!(file.location(node_byte), (2, 3)); - assert_eq!(file.line_text(2), Some(" node a")); - assert_eq!( - file.snippet(Span::at(node_byte, node_byte + 4, 2, 3)), - "node" - ); -} - -#[test] -fn source_map_assigns_ids_and_resolves_files() { - let mut map = SourceMap::new(); - assert!(map.is_empty()); - let a = map.add("a.rag", "graph a {}"); - let b = map.add("b.rag", "graph b {}"); - assert_eq!(map.len(), 2); - assert_ne!(a, b); - assert_eq!(map.get(a).unwrap().name(), "a.rag"); - assert_eq!(map.get(b).unwrap().text(), "graph b {}"); -} - -#[test] -fn diagnostic_renders_caret_under_primary_span() { - let source = "graph g {\n tool_call -> toolz\n}\n"; - let file = SourceFile::new("support.rag", source); - let target = source.find("toolz").unwrap(); - let span = Span::at(target, target + "toolz".len(), 2, 16); - let rendered = Diagnostic::error("route target `toolz` does not exist", span) - .with_code("E-rag-unknown-node") - .with_primary_label("unknown node") - .with_help("did you mean `tools`?") - .render(&file); - - assert!( - rendered.contains("error[E-rag-unknown-node]: route target `toolz` does not exist"), - "{rendered}" - ); - assert!(rendered.contains("--> support.rag:2:16"), "{rendered}"); - assert!(rendered.contains("tool_call -> toolz"), "{rendered}"); - // Five carets under the five characters of `toolz`, plus the label. - assert!(rendered.contains("^^^^^ unknown node"), "{rendered}"); - assert!( - rendered.contains("help: did you mean `tools`?"), - "{rendered}" - ); -} - -#[test] -fn severity_labels_are_lowercase() { - assert_eq!(Severity::Error.label(), "error"); - assert_eq!(Severity::Warning.label(), "warning"); - assert_eq!(Severity::Note.label(), "note"); -} - -#[test] -fn parse_error_carries_rendered_caret_for_source() { - // `bogus` is not a valid node item; `parse_str` has the source so the error - // message should render a caret beneath the offending token. - let err = parse_str("graph g {\n node a { bogus x }\n}\n").unwrap_err(); - match err { - crate::error::TinyAgentsError::Parse { - message, - line, - column, - } => { - assert!(message.contains("unknown node item `bogus`"), "{message}"); - assert!(message.contains('^'), "{message}"); - assert!(message.contains("--> :2:12"), "{message}"); - assert_eq!((line, column), (2, 12)); - } - other => panic!("expected parse error, got {other:?}"), - } -} - -#[test] -fn parse_error_without_source_renders_plain() { - // The token-only `parse` entry point has no source text, so the rendered - // message falls back to the source-free presentation (no caret). - let tokens = tokenize("graph { }").unwrap(); - let err = parse(&tokens).unwrap_err(); - match err { - crate::error::TinyAgentsError::Parse { message, .. } => { - assert!(message.contains("expected identifier"), "{message}"); - assert!(!message.contains('^'), "{message}"); - } - other => panic!("expected parse error, got {other:?}"), - } -} - -// --------------------------------------------------------------------------- -// Parser -// --------------------------------------------------------------------------- - -#[test] -fn parses_support_agent_into_ast() { - let program = parse_str(SUPPORT_AGENT).unwrap(); - assert_eq!(program.graphs.len(), 1); - let graph = &program.graphs[0]; - - assert_eq!(graph.name, "support_agent"); - assert_eq!(graph.start.as_deref(), Some("agent")); - assert_eq!(graph.channels.len(), 2); - assert_eq!(graph.channels[0].name, "messages"); - assert_eq!(graph.channels[0].reducer, "messages"); - assert_eq!(graph.channels[1].reducer, "append"); - - // Defaults preserve declared order and literal kinds. - assert_eq!(graph.defaults.len(), 3); - assert_eq!(graph.defaults[0].0, "recursion_limit"); - assert_eq!(graph.defaults[0].1, Literal::Num(50.0)); - assert_eq!(graph.defaults[1].1, Literal::Str("exponential".into())); - assert_eq!(graph.defaults[2].1, Literal::Ident("inherit".into())); - - assert_eq!(graph.nodes.len(), 2); - let agent = &graph.nodes[0]; - assert_eq!(agent.kind.as_deref(), Some("agent")); - assert_eq!(agent.model.as_deref(), Some("default")); - assert_eq!(agent.tools, vec!["lookup_user", "create_ticket"]); - assert_eq!(agent.routes.len(), 2); - assert_eq!(agent.routes[0].label, "tool_call"); - assert_eq!(agent.routes[0].target, "tools"); - assert_eq!(agent.routes[1].target, "END"); - - let tools = &graph.nodes[1]; - assert_eq!(tools.kind.as_deref(), Some("tool_executor")); - assert_eq!(tools.next.as_deref(), Some("agent")); -} - -#[test] -fn parses_top_level_edge() { - let src = "graph g { start a node a { } node b { } a -> b b -> END }"; - let program = parse_str(src).unwrap(); - let graph = &program.graphs[0]; - assert_eq!(graph.edges.len(), 2); - assert_eq!(graph.edges[0].from, "a"); - assert_eq!(graph.edges[0].to, "b"); - assert_eq!(graph.edges[1].to, "END"); -} - -#[test] -fn parse_reports_unexpected_token() { - // Missing graph name. - let tokens = tokenize("graph { }").unwrap(); - let err = parse(&tokens).unwrap_err(); - match err { - crate::error::TinyAgentsError::Parse { message, .. } => { - assert!(message.contains("expected identifier"), "{message}"); - } - other => panic!("expected parse error, got {other:?}"), - } -} - -#[test] -fn parse_rejects_unknown_node_item() { - let src = "graph g { start a node a { bogus x } }"; - let err = parse_str(src).unwrap_err(); - match err { - crate::error::TinyAgentsError::Parse { message, .. } => { - assert!(message.contains("unknown node item"), "{message}"); - } - other => panic!("expected parse error, got {other:?}"), - } -} - -// --------------------------------------------------------------------------- -// Compiler: AST -> Blueprint -// --------------------------------------------------------------------------- - -#[test] -fn compiles_support_agent_blueprint() { - let program = parse_str(SUPPORT_AGENT).unwrap(); - let blueprints = compile(&program).unwrap(); - assert_eq!(blueprints.len(), 1); - let bp = &blueprints[0]; - - assert_eq!(bp.graph_id, "support_agent"); - assert_eq!(bp.start, "agent"); - assert_eq!(bp.channels.len(), 2); - assert_eq!(bp.defaults.len(), 3); - assert_eq!(bp.nodes.len(), 2); - - let agent = &bp.nodes[0]; - assert_eq!(agent.kind, "agent"); - assert_eq!(agent.tools, vec!["lookup_user", "create_ticket"]); - match &agent.routing { - Routing::Conditional(routes) => { - assert_eq!(routes.len(), 2); - assert_eq!(routes[0], ("tool_call".into(), "tools".into())); - assert_eq!(routes[1], ("final".into(), "END".into())); - } - other => panic!("expected conditional routing, got {other:?}"), - } - - let tools = &bp.nodes[1]; - assert_eq!(tools.routing, Routing::Next("agent".into())); -} - -#[test] -fn next_end_lowers_to_terminal() { - let src = "graph g { start a node a { kind model next END } }"; - let bp = &compile(&parse_str(src).unwrap()).unwrap()[0]; - assert_eq!(bp.nodes[0].routing, Routing::Terminal); -} - -#[test] -fn blueprint_round_trips_through_serde() { - let bp = compile(&parse_str(SUPPORT_AGENT).unwrap()) - .unwrap() - .remove(0); - let json = serde_json::to_string(&bp).unwrap(); - let back: crate::language::types::Blueprint = serde_json::from_str(&json).unwrap(); - assert_eq!(bp, back); -} - -#[test] -fn missing_start_is_a_compile_error() { - let src = "graph g { node a { kind model } }"; - let err = compile(&parse_str(src).unwrap()).unwrap_err(); - assert!(matches!(err, crate::error::TinyAgentsError::Compile(_))); - assert!(err.to_string().contains("no `start`"), "{err}"); -} - -#[test] -fn start_not_defined_is_a_compile_error() { - let src = "graph g { start missing node a { kind model } }"; - let err = compile(&parse_str(src).unwrap()).unwrap_err(); - assert!(err.to_string().contains("is not defined"), "{err}"); -} - -#[test] -fn duplicate_node_is_a_compile_error() { - let src = "graph g { start a node a { kind model } node a { kind model } }"; - let err = compile(&parse_str(src).unwrap()).unwrap_err(); - assert!(err.to_string().contains("duplicate node"), "{err}"); -} - -#[test] -fn unknown_route_target_is_a_compile_error() { - let src = "graph g { start a node a { routes { go -> ghost } } }"; - let err = compile(&parse_str(src).unwrap()).unwrap_err(); - assert!(err.to_string().contains("route target"), "{err}"); -} - -#[test] -fn unknown_next_target_is_a_compile_error() { - let src = "graph g { start a node a { next ghost } }"; - let err = compile(&parse_str(src).unwrap()).unwrap_err(); - assert!(err.to_string().contains("next target"), "{err}"); -} - -#[test] -fn mixing_next_and_routes_is_a_compile_error() { - let src = "graph g { start a node a { next b routes { x -> b } } node b { } }"; - let err = compile(&parse_str(src).unwrap()).unwrap_err(); - assert!(err.to_string().contains("mixes static routing"), "{err}"); -} - -#[test] -fn mixing_edge_and_routes_is_a_compile_error() { - let src = "graph g { start a node a { routes { x -> b } } node b { } a -> b }"; - let err = compile(&parse_str(src).unwrap()).unwrap_err(); - assert!(err.to_string().contains("mixes static routing"), "{err}"); -} - -#[test] -fn duplicate_route_label_is_a_compile_error() { - let src = "graph g { start a node a { routes { x -> b\n x -> b } } node b { } }"; - let err = compile(&parse_str(src).unwrap()).unwrap_err(); - assert!(err.to_string().contains("duplicate route label"), "{err}"); -} - -// --------------------------------------------------------------------------- -// Extended grammar (H2): channels+policy, command, send/join, subgraph, -// subagent, repl_agent, interrupt, io shape, checkpoint/interrupt policy. -// --------------------------------------------------------------------------- - -/// A graph exercising every H2 primitive in one declarative blueprint. -const EXTENDED: &str = r#" -graph orchestrator { - start planner - - input { - request string - customer_id string - } - output { - answer string - } - - checkpoint inherit - interrupt manual - - channel messages messages - channel usage aggregate "usage_delta" - channel arrivals barrier 2 - - node planner { - kind agent - model "default" - command { - goto fanout - update { - status "planned" - } - } - } - - node fanout { - kind model - sends [ - send worker_a "split_a" - send worker_b "split_b" - ] - next worker_a - } - - node worker_a { - kind model - next gather - } - - node worker_b { - kind model - next gather - } - - node gather { - kind join - sources [worker_a, worker_b] - next research - } - - node research { - kind subagent - agent "researcher" - input "topic" - timeout 30 - retry { - max_attempts 3 - backoff "exponential" - } - next sub - } - - node sub { - kind subgraph - graph "summarize" - next triage - } - - node triage { - kind repl_agent - model "default" - script "triage_script" - next review - } - - node review { - kind interrupt - prompt "Approve?" - options ["approve", "reject"] - routes { - approve -> END - reject -> planner - } - } - - join [worker_a, worker_b] -> gather -} -"#; - -#[test] -fn extended_grammar_parses_and_compiles_blueprint_shape() { - let program = parse_str(EXTENDED).unwrap(); - let bp = compile(&program).unwrap().remove(0); - - assert_eq!(bp.graph_id, "orchestrator"); - assert_eq!(bp.start, "planner"); - assert_eq!(bp.checkpoint.as_deref(), Some("inherit")); - assert_eq!(bp.interrupt.as_deref(), Some("manual")); - - // Input/output shape. - assert_eq!(bp.input.len(), 2); - assert_eq!(bp.input[0].name, "request"); - assert_eq!(bp.input[0].ty, "string"); - assert_eq!(bp.output.len(), 1); - assert_eq!(bp.output[0].name, "answer"); - - // Channels carry reducer + policy args. - assert_eq!(bp.channels.len(), 3); - let usage = bp.channels.iter().find(|c| c.name == "usage").unwrap(); - assert_eq!(usage.reducer, "aggregate"); - assert_eq!(usage.args, vec![Literal::Str("usage_delta".into())]); - let arrivals = bp.channels.iter().find(|c| c.name == "arrivals").unwrap(); - assert_eq!(arrivals.reducer, "barrier"); - assert_eq!(arrivals.args, vec![Literal::Num(2.0)]); - - // Command lowering + routing precedence (goto becomes a static next). - let planner = bp.nodes.iter().find(|n| n.name == "planner").unwrap(); - let cmd = planner.command.as_ref().unwrap(); - assert_eq!(cmd.goto.as_deref(), Some("fanout")); - assert_eq!( - cmd.update, - vec![("status".into(), Literal::Str("planned".into()))] - ); - assert_eq!(planner.routing, Routing::Next("fanout".into())); - - // Fanout sends. - let fanout = bp.nodes.iter().find(|n| n.name == "fanout").unwrap(); - assert_eq!(fanout.sends.len(), 2); - assert_eq!(fanout.sends[0].target, "worker_a"); - assert_eq!(fanout.sends[0].input.as_deref(), Some("split_a")); - - // Join node. - let gather = bp.nodes.iter().find(|n| n.name == "gather").unwrap(); - assert_eq!(gather.kind, "join"); - assert_eq!(gather.join_sources, vec!["worker_a", "worker_b"]); - - // Sub-agent node with input mapping + policies. - let research = bp.nodes.iter().find(|n| n.name == "research").unwrap(); - assert_eq!(research.kind, "subagent"); - assert_eq!(research.agent.as_deref(), Some("researcher")); - assert_eq!(research.input.as_deref(), Some("topic")); - assert_eq!(research.timeout.as_deref(), Some("30")); - assert_eq!( - research.retry, - vec![ - ("max_attempts".into(), Literal::Num(3.0)), - ("backoff".into(), Literal::Str("exponential".into())), - ] - ); - - // Subgraph node references a registered graph by name. - let sub = bp.nodes.iter().find(|n| n.name == "sub").unwrap(); - assert_eq!(sub.kind, "subgraph"); - assert_eq!(sub.subgraph.as_deref(), Some("summarize")); - - // REPL-backed node names a script capability (declaration only). - let triage = bp.nodes.iter().find(|n| n.name == "triage").unwrap(); - assert_eq!(triage.kind, "repl_agent"); - assert_eq!(triage.script.as_deref(), Some("triage_script")); - - // Interrupt node with options + conditional routing. - let review = bp.nodes.iter().find(|n| n.name == "review").unwrap(); - assert_eq!(review.kind, "interrupt"); - assert_eq!(review.options, vec!["approve", "reject"]); - assert!(matches!(review.routing, Routing::Conditional(_))); - - // Top-level join declaration. - assert_eq!(bp.joins.len(), 1); - assert_eq!(bp.joins[0].target, "gather"); - assert_eq!(bp.joins[0].sources, vec!["worker_a", "worker_b"]); -} - -#[test] -fn extended_blueprint_round_trips_through_serde() { - let bp = compile(&parse_str(EXTENDED).unwrap()).unwrap().remove(0); - let json = serde_json::to_string(&bp).unwrap(); - let back: crate::language::types::Blueprint = serde_json::from_str(&json).unwrap(); - assert_eq!(bp, back); -} - -#[test] -fn command_goto_unknown_target_is_a_compile_error() { - let src = "graph g { start a node a { command { goto ghost } } }"; - let err = compile(&parse_str(src).unwrap()).unwrap_err(); - assert!(err.to_string().contains("command goto target"), "{err}"); -} - -#[test] -fn send_unknown_target_is_a_compile_error() { - let src = "graph g { start a node a { sends [ send ghost ] } }"; - let err = compile(&parse_str(src).unwrap()).unwrap_err(); - assert!(err.to_string().contains("send target"), "{err}"); -} - -#[test] -fn join_unknown_source_is_a_compile_error() { - let src = "graph g { start a node a { } join [ghost] -> a }"; - let err = compile(&parse_str(src).unwrap()).unwrap_err(); - assert!(err.to_string().contains("join source"), "{err}"); -} - -#[test] -fn extended_kinds_bind_against_a_resolver() { - let bp = compile(&parse_str(EXTENDED).unwrap()).unwrap().remove(0); - let resolver = CapabilityResolver::from_lists(["default".to_string()], std::iter::empty()) - .allow_subgraph("summarize") - .allow_reducer("messages") - .allow_reducer("aggregate") - .allow_reducer("barrier") - .with_node_kinds( - crate::language::compiler::DEFAULT_NODE_KINDS - .iter() - .map(|k| k.to_string()), - ); - resolver.bind_blueprint(&bp).unwrap(); -} - -// --------------------------------------------------------------------------- -// Capability binding -// --------------------------------------------------------------------------- - -#[test] -fn bind_capabilities_accepts_allowed_references() { - let bp = compile(&parse_str(SUPPORT_AGENT).unwrap()) - .unwrap() - .remove(0); - let resolver = CapabilityResolver::from_lists( - ["default".to_string()], - ["lookup_user".to_string(), "create_ticket".to_string()], - ); - bind_capabilities(&bp, &resolver).unwrap(); -} - -#[test] -fn bind_capabilities_rejects_unknown_model() { - let bp = compile(&parse_str(SUPPORT_AGENT).unwrap()) - .unwrap() - .remove(0); - let resolver = CapabilityResolver::new() - .allow_tool("lookup_user") - .allow_tool("create_ticket"); - let err = bind_capabilities(&bp, &resolver).unwrap_err(); - assert!(matches!(err, crate::error::TinyAgentsError::Capability(_))); - assert!(err.to_string().contains("unknown model"), "{err}"); -} - -#[test] -fn bind_capabilities_rejects_unknown_tool() { - let bp = compile(&parse_str(SUPPORT_AGENT).unwrap()) - .unwrap() - .remove(0); - let resolver = CapabilityResolver::new().allow_model("default"); - let err = bind_capabilities(&bp, &resolver).unwrap_err(); - assert!(err.to_string().contains("unknown tool"), "{err}"); -} - -// --------------------------------------------------------------------------- -// Graph materialisation + execution -// --------------------------------------------------------------------------- - -#[derive(Clone, Debug, PartialEq, Eq)] -struct TestState { - trail: Vec, - agent_visits: u32, -} - -/// Resolves a conditional route `label` to a *durable* target node id from the -/// blueprint's `(label, target)` table, translating the language `END` sentinel -/// (`"END"`) to the durable graph terminal ([`crate::graph::END`]). Unknown -/// labels fall back to the durable `END`. -fn resolve_durable_target(routes: &[(String, String)], label: &str) -> String { - let target = routes - .iter() - .find(|(l, _)| l == label) - .map(|(_, t)| t.as_str()); - match target { - Some(t) if t != crate::language::types::END => t.to_string(), - _ => crate::graph::END.to_string(), - } -} - -/// A trivial factory that materialises echo/route/end nodes purely from the -/// declarative [`NodeSpec`]. It demonstrates that runnable behaviour comes from -/// Rust, not the source: each node records its name; terminal/`next` nodes -/// commit a whole-state update (static edges route them), and conditional nodes -/// loop once before terminating by emitting an explicit `goto` command. -struct TestFactory; - -impl NodeFactory for TestFactory { - fn make(&self, spec: &NodeSpec) -> crate::error::Result> { - let name = spec.name.clone(); - let routing = spec.routing.clone(); - Ok(Arc::new( - move |mut state: TestState, _ctx: NodeContext| -> NodeFuture { - let name = name.clone(); - let routing = routing.clone(); - Box::pin(async move { - state.trail.push(name.clone()); - let result = match &routing { - // Static edges (Next/Terminal) handle routing; just - // commit the whole-state update. - Routing::Terminal | Routing::Next(_) => NodeResult::Update(state), - Routing::Conditional(routes) => { - state.agent_visits += 1; - // Take the `tool_call -> tools` route until the - // second visit, then take `final -> END`. - let label = if state.agent_visits >= 2 { - "final" - } else { - "tool_call" - }; - let target = resolve_durable_target(routes, label); - NodeResult::Command(Command::goto([target]).with_update(state)) - } - }; - Ok(result) - }) - }, - )) - } -} - -/// Collects the visited node ids into owned strings for comparison. -fn visited_names(run: &crate::graph::GraphExecution) -> Vec { - run.visited.iter().map(ToString::to_string).collect() -} - -#[tokio::test] -async fn build_graph_runs_to_end() { - let bp = compile(&parse_str(SUPPORT_AGENT).unwrap()) - .unwrap() - .remove(0); - let graph = build_graph(&bp, &TestFactory).unwrap(); - - let run = graph - .run(TestState { - trail: Vec::new(), - agent_visits: 0, - }) - .await - .unwrap(); - - // agent -> tools -> agent (ends on second visit). - assert_eq!(visited_names(&run), vec!["agent", "tools", "agent"]); - assert_eq!(run.state.trail, vec!["agent", "tools", "agent"]); - assert_eq!(run.state.agent_visits, 2); -} - -#[tokio::test] -async fn build_graph_handles_linear_terminal() { - let src = "graph g { start a node a { kind model next b } node b { kind model next END } }"; - let bp = compile(&parse_str(src).unwrap()).unwrap().remove(0); - let graph = build_graph(&bp, &TestFactory).unwrap(); - let run = graph - .run(TestState { - trail: Vec::new(), - agent_visits: 0, - }) - .await - .unwrap(); - assert_eq!(visited_names(&run), vec!["a", "b"]); -} - -// --------------------------------------------------------------------------- -// Registry-backed capability binding (registry → language binding) -// --------------------------------------------------------------------------- - -use crate::language::compiler::{ - DEFAULT_NODE_KINDS, bind_capabilities_with_registry, compile_source, -}; -use crate::registry::CapabilityRegistry; - -/// A `.rag` graph that exercises every registry-backed reference kind: a model, -/// a tool, a subgraph reference (`kind subgraph` whose `model` names a -/// registered blueprint), a router reference (`kind router`), and a channel -/// reducer. -const FULL_SOURCE: &str = r#" -graph main { - start agent - - channel messages append - - node agent { - kind agent - model "default" - tools ["lookup_user"] - routes { - retrieve -> sub - classify -> route - done -> END - } - } - - node sub { - kind subgraph - model "retrieval" - next END - } - - node route { - kind router - model "classify" - next END - } -} -"#; - -/// Builds a registry that satisfies every reference in [`FULL_SOURCE`]. -fn full_registry() -> CapabilityRegistry { - let mut reg = CapabilityRegistry::::new(); - reg.register_model( - "default", - std::sync::Arc::new(crate::language::test::testkit::EchoModel), - ) - .unwrap(); - reg.register_tool(std::sync::Arc::new( - crate::language::test::testkit::NoopTool, - )) - .unwrap(); - reg.register_graph_blueprint( - "retrieval", - compile(&parse_str("graph retrieval { start r node r { kind model next END } }").unwrap()) - .unwrap() - .remove(0), - ) - .unwrap(); - reg.register_router("classify").unwrap(); - reg.register_reducer("append").unwrap(); - reg -} - -#[test] -fn compile_source_binds_against_registry() { - let reg = full_registry(); - let blueprints = compile_source(FULL_SOURCE, ®).unwrap(); - assert_eq!(blueprints.len(), 1); - assert_eq!(blueprints[0].graph_id, "main"); -} - -#[test] -fn registry_resolver_allows_all_kinds() { - let reg = full_registry(); - let resolver = reg.capability_resolver(); - assert!(resolver.model_allowed("default")); - assert!(resolver.tool_allowed("lookup_user")); - assert!(resolver.subgraph_allowed("retrieval")); - assert!(resolver.router_allowed("classify")); - assert!(resolver.reducer_allowed("append")); - for kind in DEFAULT_NODE_KINDS { - assert!(resolver.node_kind_allowed(kind)); - } -} - -#[test] -fn registry_bind_rejects_unregistered_model() { - let mut reg = full_registry(); - reg.replace_model( - "other", - std::sync::Arc::new(crate::language::test::testkit::EchoModel), - ); - // Source references `default`, which we did not register here. - let mut bare = CapabilityRegistry::::new(); - bare.register_tool(std::sync::Arc::new( - crate::language::test::testkit::NoopTool, - )) - .unwrap(); - bare.register_graph_blueprint( - "retrieval", - reg.graph_blueprint("retrieval").unwrap().clone(), - ) - .unwrap(); - bare.register_router("classify").unwrap(); - bare.register_reducer("append").unwrap(); - let err = compile_source(FULL_SOURCE, &bare).unwrap_err(); - assert!(matches!(err, crate::error::TinyAgentsError::Capability(_))); - assert!(err.to_string().contains("unknown model"), "{err}"); -} - -#[test] -fn registry_bind_rejects_unregistered_tool() { - let src = r#"graph g { start a channel m append node a { kind agent model "default" tools ["missing"] next END } }"#; - let reg = full_registry(); - let err = compile_source(src, ®).unwrap_err(); - assert!(err.to_string().contains("unknown tool"), "{err}"); -} - -#[test] -fn registry_bind_rejects_unregistered_subgraph() { - let src = r#"graph g { start s node s { kind subgraph model "ghost" next END } }"#; - let reg = full_registry(); - let err = compile_source(src, ®).unwrap_err(); - assert!(matches!(err, crate::error::TinyAgentsError::Capability(_))); - assert!(err.to_string().contains("unknown subgraph"), "{err}"); -} - -#[test] -fn registry_bind_rejects_unregistered_router() { - let src = r#"graph g { start r node r { kind router model "ghost" next END } }"#; - let reg = full_registry(); - let err = compile_source(src, ®).unwrap_err(); - assert!(err.to_string().contains("unknown router"), "{err}"); -} - -#[test] -fn registry_bind_rejects_unregistered_reducer() { - let src = r#"graph g { start a channel messages ghost node a { kind model next END } }"#; - let reg = full_registry(); - let err = compile_source(src, ®).unwrap_err(); - assert!(matches!(err, crate::error::TinyAgentsError::Capability(_))); - assert!(err.to_string().contains("unknown reducer"), "{err}"); -} - -#[test] -fn registry_bind_rejects_unknown_node_kind() { - let src = r#"graph g { start a node a { kind wizard next END } }"#; - let reg = full_registry(); - let err = compile_source(src, ®).unwrap_err(); - assert!(matches!(err, crate::error::TinyAgentsError::Compile(_))); - assert!(err.to_string().contains("unknown kind"), "{err}"); -} - -#[test] -fn manual_bind_path_ignores_kinds_and_reducers() { - // The legacy manual resolver must keep working: a non-empty node kind set is - // never consulted, and reducers/subgraphs are not checked. - let src = r#"graph g { start a channel messages ghost node a { kind wizard model "default" next END } }"#; - let bp = compile(&parse_str(src).unwrap()).unwrap().remove(0); - let resolver = CapabilityResolver::new().allow_model("default"); - // Manual gate only checks model + tool; passes despite the unknown kind, - // unknown reducer, and exotic node kind. - bind_capabilities(&bp, &resolver).unwrap(); -} - -#[test] -fn bind_capabilities_with_registry_matches_compile_source() { - let reg = full_registry(); - let bp = compile(&parse_str(FULL_SOURCE).unwrap()).unwrap().remove(0); - bind_capabilities_with_registry(&bp, ®).unwrap(); -} - -// --------------------------------------------------------------------------- -// Registry-backed Resolver (H3): spanned diagnostics, single binding gate -// --------------------------------------------------------------------------- - -use crate::language::resolver::{Resolver, resolve_source}; - -#[test] -fn resolver_accepts_fully_registered_blueprint() { - let reg = full_registry(); - let program = parse_str(FULL_SOURCE).unwrap(); - let resolver = Resolver::from_registry(®); - // No diagnostics: every model/tool/subgraph/router/reducer is registered. - assert!(resolver.resolve_program(&program).is_empty()); - // And the convenience façade lowers it to a blueprint. - let blueprints = resolve_source(FULL_SOURCE, ®).unwrap(); - assert_eq!(blueprints[0].graph_id, "main"); -} - -#[test] -fn resolver_reports_unregistered_tool_with_spanned_diagnostic() { - // `missing` is not a registered tool. - let src = r#" -graph g { - start a - channel m append - node a { - kind agent - model "default" - tools ["missing"] - next END - } -} -"#; - let reg = full_registry(); - let program = parse_str(src).unwrap(); - let file = SourceFile::new("plan.rag", src); - let resolver = Resolver::from_registry(®); - - let diagnostics = resolver.resolve_program(&program); - assert_eq!(diagnostics.len(), 1, "{diagnostics:?}"); - let diag = &diagnostics[0]; - assert_eq!(diag.code.as_deref(), Some("E-rag-unknown-tool")); - let rendered = diag.render(&file); - assert!( - rendered.contains("node `a` references unknown tool `missing`"), - "{rendered}" - ); - // The diagnostic carries a caret pointing at the offending node span. - assert!(rendered.contains('^'), "{rendered}"); - assert!(rendered.contains("--> plan.rag:"), "{rendered}"); - - // `check_program` folds it into a Capability error with the rendered caret. - let err = resolver.check_program(&program, Some(&file)).unwrap_err(); - assert!(matches!(err, crate::error::TinyAgentsError::Capability(_))); - assert!(err.to_string().contains("unknown tool"), "{err}"); - assert!(err.to_string().contains('^'), "{err}"); -} - -#[test] -fn resolve_source_rejects_unregistered_tool() { - let src = r#"graph g { start a channel m append node a { kind agent model "default" tools ["missing"] next END } }"#; - let reg = full_registry(); - let err = resolve_source(src, ®).unwrap_err(); - assert!(matches!(err, crate::error::TinyAgentsError::Capability(_))); - assert!(err.to_string().contains("unknown tool"), "{err}"); -} - -#[test] -fn resolver_reports_unknown_node_kind_as_compile_error() { - let src = r#"graph g { start a node a { kind wizard next END } }"#; - let reg = full_registry(); - let err = resolve_source(src, ®).unwrap_err(); - assert!(matches!(err, crate::error::TinyAgentsError::Compile(_))); - assert!(err.to_string().contains("unknown kind"), "{err}"); -} - -#[test] -fn resolver_reports_unregistered_agent() { - // A `subagent` node binds its `agent "…"` reference through the registry's - // Agent allowlist. - let src = r#"graph g { start a node a { kind subagent agent "ghost" next END } }"#; - let reg = full_registry(); - let program = parse_str(src).unwrap(); - let diagnostics = Resolver::from_registry(®).resolve_program(&program); - assert_eq!(diagnostics.len(), 1, "{diagnostics:?}"); - assert_eq!(diagnostics[0].code.as_deref(), Some("E-rag-unknown-agent")); - assert!( - diagnostics[0].message.contains("unknown agent `ghost`"), - "{:?}", - diagnostics[0] - ); -} - -#[test] -fn resolver_collects_multiple_diagnostics() { - // Two independent problems: an unregistered model and an unregistered - // reducer. `resolve_program` reports both. - let src = r#"graph g { start a channel m ghost node a { kind model model "nope" next END } }"#; - let reg = full_registry(); - let program = parse_str(src).unwrap(); - let diagnostics = Resolver::from_registry(®).resolve_program(&program); - let codes: Vec<_> = diagnostics - .iter() - .filter_map(|d| d.code.as_deref()) - .collect(); - assert!(codes.contains(&"E-rag-unknown-model"), "{codes:?}"); - assert!(codes.contains(&"E-rag-unknown-reducer"), "{codes:?}"); -} - -#[test] -fn resolver_blueprint_path_matches_registry_binding() { - // The span-less blueprint path mirrors the legacy gate's variants/messages. - let reg = full_registry(); - let bp = compile(&parse_str(FULL_SOURCE).unwrap()).unwrap().remove(0); - Resolver::from_registry(®) - .resolve_blueprint(&bp) - .unwrap(); - - let bad = compile( - &parse_str(r#"graph g { start a channel m append node a { kind subgraph model "ghost" next END } }"#) - .unwrap(), - ) - .unwrap() - .remove(0); - let err = Resolver::from_registry(®) - .resolve_blueprint(&bad) - .unwrap_err(); - assert!(matches!(err, crate::error::TinyAgentsError::Capability(_))); - assert!(err.to_string().contains("unknown subgraph"), "{err}"); -} - -// --------------------------------------------------------------------------- -// Provenance, diff, and language testkit (H4) -// --------------------------------------------------------------------------- - -use crate::language::compiler::compile_with_provenance; -use crate::language::diff::{FieldChange, blueprint_diff}; -use crate::language::testkit as lang_testkit; -use crate::language::types::Origin; - -const DIFF_BASE: &str = r#" -graph flow { - start plan - - channel messages append - - node plan { - kind model - model "default" - routes { - research -> work - done -> END - } - } - - node work { - kind tool_executor - tools ["lookup_user"] - next END - } -} -"#; - -/// Adds a node (`review`) and changes a route target on `plan` -/// (`research -> review` instead of `research -> work`). -const DIFF_NEW: &str = r#" -graph flow { - start plan - - channel messages append - - node plan { - kind model - model "default" - routes { - research -> review - done -> END - } - } - - node review { - kind interrupt - prompt "ok?" - next work - } - - node work { - kind tool_executor - tools ["lookup_user"] - next END - } -} -"#; - -#[test] -fn blueprint_diff_reports_added_node_and_changed_route() { - let old = lang_testkit::blueprint(DIFF_BASE); - let new = lang_testkit::blueprint(DIFF_NEW); - - let diff = blueprint_diff(&old, &new); - assert!(!diff.is_empty()); - - // One node added: `review`. - assert_eq!(diff.nodes_added, vec!["review".to_string()]); - assert!(diff.nodes_removed.is_empty()); - - // `plan`'s routing changed (research target work -> review). - assert_eq!(diff.nodes_changed.len(), 1, "{diff:?}"); - let plan_change = &diff.nodes_changed[0]; - assert_eq!(plan_change.name, "plan"); - assert_eq!( - plan_change.fields, - vec![FieldChange { - field: "routing".to_string(), - old: "{ research -> work, done -> END }".to_string(), - new: "{ research -> review, done -> END }".to_string(), - }] - ); - - // The rendered summary names both the added node and the changed route. - let rendered = diff.to_string(); - assert!(rendered.contains("+ node review"), "{rendered}"); - assert!(rendered.contains("~ node plan"), "{rendered}"); - assert!(rendered.contains("routing:"), "{rendered}"); -} - -#[test] -fn blueprint_diff_of_identical_blueprints_is_empty() { - let a = lang_testkit::blueprint(DIFF_BASE); - let b = lang_testkit::blueprint(DIFF_BASE); - let diff = blueprint_diff(&a, &b); - assert!(diff.is_empty(), "{diff:?}"); - assert_eq!(diff.to_string(), "no changes"); -} - -#[test] -fn blueprint_diff_serializes_round_trip() { - let old = lang_testkit::blueprint(DIFF_BASE); - let new = lang_testkit::blueprint(DIFF_NEW); - let diff = blueprint_diff(&old, &new); - let json = serde_json::to_string(&diff).unwrap(); - let back: crate::language::diff::BlueprintDiff = serde_json::from_str(&json).unwrap(); - assert_eq!(diff, back); -} - -#[test] -fn provenance_points_each_node_at_its_span() { - let program = parse_str(DIFF_NEW).unwrap(); - let bp = compile_with_provenance(&program, Origin::file("flow.rag")) - .unwrap() - .remove(0); - - let prov = bp.provenance().expect("provenance attached"); - assert_eq!(prov.origin, Origin::File("flow.rag".to_string())); - - // Every node in the blueprint has a recorded span anchored at the line of - // its `node ` declaration, and the byte range slices back to the - // `node` keyword that opens it. - let file = SourceFile::new("flow.rag", DIFF_NEW); - for node in &bp.nodes { - let span = prov - .node_span(&node.name) - .unwrap_or_else(|| panic!("no span for node `{}`", node.name)); - // The span's byte range covers the opening `node` keyword. - assert_eq!(&DIFF_NEW[span.start..span.end], "node"); - // The line the span anchors at is the node's declaration line. - let line = file - .line_text(span.line) - .unwrap_or_else(|| panic!("no source line {}", span.line)); - assert!( - line.contains(&format!("node {}", node.name)), - "node `{}` span anchors at line `{line}`, not its declaration", - node.name - ); - } - - // The channel span is recorded too. - assert!(prov.channel_span("messages").is_some()); -} - -#[test] -fn plain_compile_leaves_provenance_none() { - let bp = lang_testkit::blueprint(DIFF_BASE); - assert!(bp.provenance().is_none()); -} - -#[test] -fn generated_origin_renders_label() { - let program = parse_str(DIFF_BASE).unwrap(); - let bp = compile_with_provenance(&program, Origin::generated_by("repl-7")) - .unwrap() - .remove(0); - let prov = bp.provenance().unwrap(); - assert_eq!(prov.origin.as_display(), "generated by repl-7"); -} - -#[test] -fn testkit_assertions_inspect_lowered_topology() { - let bp = lang_testkit::blueprint(DIFF_NEW); - lang_testkit::assert_kind(&bp, "review", "interrupt"); - lang_testkit::assert_next(&bp, "review", "work"); - lang_testkit::assert_terminal(&bp, "work"); - lang_testkit::assert_route(&bp, "plan", "research", "review"); -} - -#[test] -fn testkit_try_compile_surfaces_errors() { - // `start` references an undefined node, so compilation fails. - let err = lang_testkit::try_compile("graph g { start missing node a { kind model next END } }") - .unwrap_err(); - assert!(matches!(err, crate::error::TinyAgentsError::Compile(_))); -} - -/// Minimal fake model/tool used to populate a [`CapabilityRegistry`] in tests. -mod testkit { - use async_trait::async_trait; - use serde_json::json; - - use crate::error::Result; - use crate::harness::model::{ChatModel, ModelRequest, ModelResponse}; - use crate::harness::tool::{Tool, ToolCall, ToolResult, ToolSchema}; - - use super::TestState; - - pub(super) struct EchoModel; - - #[async_trait] - impl ChatModel for EchoModel { - async fn invoke( - &self, - _state: &TestState, - _request: ModelRequest, - ) -> Result { - Ok(ModelResponse::assistant("echo")) - } - } - - pub(super) struct NoopTool; - - #[async_trait] - impl Tool for NoopTool { - fn name(&self) -> &str { - "lookup_user" - } - fn description(&self) -> &str { - "noop" - } - fn schema(&self) -> ToolSchema { - ToolSchema::new("lookup_user", "noop", json!({"type": "object"})) - } - async fn call(&self, _state: &TestState, call: ToolCall) -> Result { - Ok(ToolResult::text(call.id, call.name, "ok")) - } - } -} diff --git a/src/language/test/capability_binding.rs b/src/language/test/capability_binding.rs new file mode 100644 index 0000000..6eed52e --- /dev/null +++ b/src/language/test/capability_binding.rs @@ -0,0 +1,45 @@ +//! Capability-binding tests (the legacy `bind_capabilities` allowlist +//! gate). +//! +//! Split out of `language/test/mod.rs` by pipeline phase. + +use super::*; + +// --------------------------------------------------------------------------- +// Capability binding +// --------------------------------------------------------------------------- + +#[test] +fn bind_capabilities_accepts_allowed_references() { + let bp = compile(&parse_str(SUPPORT_AGENT).unwrap()) + .unwrap() + .remove(0); + let resolver = CapabilityResolver::from_lists( + ["default".to_string()], + ["lookup_user".to_string(), "create_ticket".to_string()], + ); + bind_capabilities(&bp, &resolver).unwrap(); +} + +#[test] +fn bind_capabilities_rejects_unknown_model() { + let bp = compile(&parse_str(SUPPORT_AGENT).unwrap()) + .unwrap() + .remove(0); + let resolver = CapabilityResolver::new() + .allow_tool("lookup_user") + .allow_tool("create_ticket"); + let err = bind_capabilities(&bp, &resolver).unwrap_err(); + assert!(matches!(err, crate::error::TinyAgentsError::Capability(_))); + assert!(err.to_string().contains("unknown model"), "{err}"); +} + +#[test] +fn bind_capabilities_rejects_unknown_tool() { + let bp = compile(&parse_str(SUPPORT_AGENT).unwrap()) + .unwrap() + .remove(0); + let resolver = CapabilityResolver::new().allow_model("default"); + let err = bind_capabilities(&bp, &resolver).unwrap_err(); + assert!(err.to_string().contains("unknown tool"), "{err}"); +} diff --git a/src/language/test/compiler.rs b/src/language/test/compiler.rs new file mode 100644 index 0000000..f0d2f27 --- /dev/null +++ b/src/language/test/compiler.rs @@ -0,0 +1,156 @@ +//! Compiler tests: AST -> Blueprint lowering. +//! +//! Split out of `language/test/mod.rs` by pipeline phase. + +use super::*; + +// --------------------------------------------------------------------------- +// Compiler: AST -> Blueprint +// --------------------------------------------------------------------------- + +#[test] +fn compiles_support_agent_blueprint() { + let program = parse_str(SUPPORT_AGENT).unwrap(); + let blueprints = compile(&program).unwrap(); + assert_eq!(blueprints.len(), 1); + let bp = &blueprints[0]; + + assert_eq!(bp.graph_id, "support_agent"); + assert_eq!(bp.start, "agent"); + assert_eq!(bp.channels.len(), 2); + assert_eq!(bp.defaults.len(), 3); + assert_eq!(bp.nodes.len(), 2); + + let agent = &bp.nodes[0]; + assert_eq!(agent.kind, "agent"); + assert_eq!(agent.tools, vec!["lookup_user", "create_ticket"]); + match &agent.routing { + Routing::Conditional(routes) => { + assert_eq!(routes.len(), 2); + assert_eq!(routes[0], ("tool_call".into(), "tools".into())); + assert_eq!(routes[1], ("final".into(), "END".into())); + } + other => panic!("expected conditional routing, got {other:?}"), + } + + let tools = &bp.nodes[1]; + assert_eq!(tools.routing, Routing::Next("agent".into())); +} + +#[test] +fn next_end_lowers_to_terminal() { + let src = "graph g { start a node a { kind model next END } }"; + let bp = &compile(&parse_str(src).unwrap()).unwrap()[0]; + assert_eq!(bp.nodes[0].routing, Routing::Terminal); +} + +#[test] +fn blueprint_round_trips_through_serde() { + let bp = compile(&parse_str(SUPPORT_AGENT).unwrap()) + .unwrap() + .remove(0); + let json = serde_json::to_string(&bp).unwrap(); + let back: crate::language::types::Blueprint = serde_json::from_str(&json).unwrap(); + assert_eq!(bp, back); +} + +#[test] +fn missing_start_is_a_compile_error() { + let src = "graph g { node a { kind model } }"; + let err = compile(&parse_str(src).unwrap()).unwrap_err(); + assert!(matches!(err, crate::error::TinyAgentsError::Compile(_))); + assert!(err.to_string().contains("no `start`"), "{err}"); +} + +#[test] +fn start_not_defined_is_a_compile_error() { + let src = "graph g { start missing node a { kind model } }"; + let err = compile(&parse_str(src).unwrap()).unwrap_err(); + assert!(err.to_string().contains("is not defined"), "{err}"); +} + +#[test] +fn duplicate_node_is_a_compile_error() { + let src = "graph g { start a node a { kind model } node a { kind model } }"; + let err = compile(&parse_str(src).unwrap()).unwrap_err(); + assert!(err.to_string().contains("duplicate node"), "{err}"); +} + +#[test] +fn unknown_route_target_is_a_compile_error() { + let src = "graph g { start a node a { routes { go -> ghost } } }"; + let err = compile(&parse_str(src).unwrap()).unwrap_err(); + assert!(err.to_string().contains("route target"), "{err}"); +} + +#[test] +fn unknown_next_target_is_a_compile_error() { + let src = "graph g { start a node a { next ghost } }"; + let err = compile(&parse_str(src).unwrap()).unwrap_err(); + assert!(err.to_string().contains("next target"), "{err}"); +} + +#[test] +fn mixing_next_and_routes_is_a_compile_error() { + let src = "graph g { start a node a { next b routes { x -> b } } node b { } }"; + let err = compile(&parse_str(src).unwrap()).unwrap_err(); + assert!(err.to_string().contains("mixes static routing"), "{err}"); +} + +#[test] +fn mixing_edge_and_routes_is_a_compile_error() { + let src = "graph g { start a node a { routes { x -> b } } node b { } a -> b }"; + let err = compile(&parse_str(src).unwrap()).unwrap_err(); + assert!(err.to_string().contains("mixes static routing"), "{err}"); +} + +#[test] +fn duplicate_route_label_is_a_compile_error() { + let src = "graph g { start a node a { routes { x -> b\n x -> b } } node b { } }"; + let err = compile(&parse_str(src).unwrap()).unwrap_err(); + assert!(err.to_string().contains("duplicate route label"), "{err}"); +} + +#[test] +fn duplicate_channel_is_a_compile_error() { + let src = "graph g { start a channel messages append channel messages messages node a { } }"; + let err = compile(&parse_str(src).unwrap()).unwrap_err(); + assert!(err.to_string().contains("duplicate channel"), "{err}"); +} + +#[test] +fn duplicate_graph_id_is_a_compile_error() { + let src = "graph g { start a node a { } } graph g { start b node b { } }"; + let err = compile(&parse_str(src).unwrap()).unwrap_err(); + assert!(err.to_string().contains("duplicate graph"), "{err}"); +} + +#[test] +fn next_and_command_goto_conflict_is_a_compile_error() { + let src = "graph g { start a node a { next b command { goto c } } node b { } node c { } }"; + let err = compile(&parse_str(src).unwrap()).unwrap_err(); + assert!( + err.to_string().contains("conflicting routing sources"), + "{err}" + ); +} + +#[test] +fn command_goto_and_edge_conflict_is_a_compile_error() { + let src = "graph g { start a node a { command { goto b } } node b { } a -> b }"; + let err = compile(&parse_str(src).unwrap()).unwrap_err(); + assert!( + err.to_string().contains("conflicting routing sources"), + "{err}" + ); +} + +#[test] +fn multiple_top_level_edges_from_same_source_is_a_compile_error() { + let src = "graph g { start a node a { } node b { } node c { } a -> b a -> c }"; + let err = compile(&parse_str(src).unwrap()).unwrap_err(); + assert!( + err.to_string().contains("multiple top-level edges"), + "{err}" + ); +} diff --git a/src/language/test/diagnostics.rs b/src/language/test/diagnostics.rs new file mode 100644 index 0000000..e21aaea --- /dev/null +++ b/src/language/test/diagnostics.rs @@ -0,0 +1,210 @@ +//! Span, source-map, and diagnostic-rendering tests. +//! +//! Split out of `language/test/mod.rs` by pipeline phase. + +use super::*; + +// --------------------------------------------------------------------------- +// Spans, source map, and diagnostics +// --------------------------------------------------------------------------- + +use crate::language::diagnostic::{Diagnostic, Severity}; +use crate::language::source::{SourceFile, SourceMap}; +use crate::language::span::Span; + +#[test] +fn span_merge_covers_both_inputs() { + let a = Span::at(2, 5, 1, 3); + let b = Span::at(10, 14, 2, 1); + let merged = a.merge(b); + assert_eq!(merged.start, 2); + assert_eq!(merged.end, 14); + // Anchor comes from the earlier-starting span. + assert_eq!((merged.line, merged.column), (1, 3)); + // Merge is commutative over the covered range. + assert_eq!(b.merge(a).start, 2); + assert_eq!(b.merge(a).end, 14); +} + +#[test] +fn span_len_and_is_empty() { + assert!(Span::new(1, 1).is_empty()); + let s = Span::at(4, 9, 1, 5); + assert_eq!(s.len(), 5); + assert!(!s.is_empty()); +} + +#[test] +fn source_file_maps_offsets_to_line_and_column() { + let file = SourceFile::new("demo.rag", "graph g\n node a\n"); + // `g` is on line 1. + assert_eq!(file.location(6), (1, 7)); + // The `node` keyword starts at byte 10 on line 2, column 3. + let node_byte = file.text().find("node").unwrap(); + assert_eq!(file.location(node_byte), (2, 3)); + assert_eq!(file.line_text(2), Some(" node a")); + assert_eq!( + file.snippet(Span::at(node_byte, node_byte + 4, 2, 3)), + "node" + ); +} + +#[test] +fn source_map_assigns_ids_and_resolves_files() { + let mut map = SourceMap::new(); + assert!(map.is_empty()); + let a = map.add("a.rag", "graph a {}"); + let b = map.add("b.rag", "graph b {}"); + assert_eq!(map.len(), 2); + assert_ne!(a, b); + assert_eq!(map.get(a).unwrap().name(), "a.rag"); + assert_eq!(map.get(b).unwrap().text(), "graph b {}"); +} + +#[test] +fn diagnostic_renders_caret_under_primary_span() { + let source = "graph g {\n tool_call -> toolz\n}\n"; + let file = SourceFile::new("support.rag", source); + let target = source.find("toolz").unwrap(); + let span = Span::at(target, target + "toolz".len(), 2, 16); + let rendered = Diagnostic::error("route target `toolz` does not exist", span) + .with_code("E-rag-unknown-node") + .with_primary_label("unknown node") + .with_help("did you mean `tools`?") + .render(&file); + + assert!( + rendered.contains("error[E-rag-unknown-node]: route target `toolz` does not exist"), + "{rendered}" + ); + assert!(rendered.contains("--> support.rag:2:16"), "{rendered}"); + assert!(rendered.contains("tool_call -> toolz"), "{rendered}"); + // Five carets under the five characters of `toolz`, plus the label. + assert!(rendered.contains("^^^^^ unknown node"), "{rendered}"); + assert!( + rendered.contains("help: did you mean `tools`?"), + "{rendered}" + ); +} + +#[test] +fn diagnostic_renders_span_past_end_of_source_without_panic() { + // A span whose bytes extend past (or start past) the end of the source must + // not panic when rendered — the caret range is clamped into the line. + let source = "graph g {}\n"; + let file = SourceFile::new("plan.rag", source); + let past = source.len() + 50; + let span = Span::at(past, past + 10, 99, 1); + let rendered = Diagnostic::error("dangling span", span) + .with_primary_label("here") + .render(&file); + assert!(rendered.contains("error: dangling span"), "{rendered}"); + // At least one caret is emitted even for an empty clamped range. + assert!(rendered.contains('^'), "{rendered}"); +} + +#[test] +fn into_parse_error_honors_stored_line_col_for_back_compat_spans_even_with_source() { + // `Span::new(line, column)` is the back-compat constructor for callers + // that only have a line/column, not a byte offset — `start`/`end` are + // both left at 0. Even when a `SourceFile` is supplied, resolving byte + // offset 0 against it would always yield 1:1, silently discarding the + // real position the caller anchored the span at. + let source = "graph g {\n start missing\n}\n"; + let file = SourceFile::new("flow.rag", source); + let span = Span::new(2, 9); + let diagnostic = Diagnostic::error("unknown start node", span).with_primary_label("here"); + + let err = diagnostic.into_parse_error(Some(&file)); + match err { + crate::error::TinyAgentsError::Parse { + line, + column, + message, + } => { + assert_eq!( + (line, column), + (2, 9), + "must honor the span's stored anchor" + ); + assert!(message.contains("unknown start node"), "{message}"); + } + other => panic!("expected Parse error, got {other:?}"), + } +} + +#[test] +fn into_parse_error_uses_real_offsets_when_present() { + // A span with real byte offsets (built via `Span::at`) must still resolve + // its line/column from the source, not just echo the stored anchor — + // this pins the happy path the previous test's fix must not regress. + let source = "graph g {\n start missing\n}\n"; + let file = SourceFile::new("flow.rag", source); + let offset = source.find("missing").unwrap(); + let span = Span::at(offset, offset + "missing".len(), 2, 9); + let diagnostic = Diagnostic::error("unknown start node", span).with_primary_label("here"); + + let err = diagnostic.into_parse_error(Some(&file)); + match err { + crate::error::TinyAgentsError::Parse { line, column, .. } => { + assert_eq!((line, column), (2, 9)); + } + other => panic!("expected Parse error, got {other:?}"), + } +} + +#[test] +fn severity_labels_are_lowercase() { + assert_eq!(Severity::Error.label(), "error"); + assert_eq!(Severity::Warning.label(), "warning"); + assert_eq!(Severity::Note.label(), "note"); +} + +#[test] +fn parse_error_carries_rendered_caret_for_source() { + // `bogus` is not a valid node item; `parse_str` has the source so the error + // message should render a caret beneath the offending token. + let err = parse_str("graph g {\n node a { bogus x }\n}\n").unwrap_err(); + match err { + crate::error::TinyAgentsError::Parse { + message, + line, + column, + } => { + assert!(message.contains("unknown node item `bogus`"), "{message}"); + assert!(message.contains('^'), "{message}"); + assert!(message.contains("--> :2:12"), "{message}"); + assert_eq!((line, column), (2, 12)); + } + other => panic!("expected parse error, got {other:?}"), + } +} + +#[test] +fn parse_error_without_source_renders_plain() { + // The token-only `parse` entry point has no source text, so the rendered + // message falls back to the source-free presentation (no caret). + let tokens = tokenize("graph { }").unwrap(); + let err = parse(&tokens).unwrap_err(); + match err { + crate::error::TinyAgentsError::Parse { message, .. } => { + assert!(message.contains("expected identifier"), "{message}"); + assert!(!message.contains('^'), "{message}"); + } + other => panic!("expected parse error, got {other:?}"), + } +} + +#[test] +fn parse_empty_token_slice_returns_error_not_panic() { + // A well-formed token stream always ends with an `Eof` sentinel; an empty + // slice violates that contract and previously underflowed `len() - 1`. + // `parse` must return a parse error instead of panicking. + let err = parse(&[]).unwrap_err(); + match err { + crate::error::TinyAgentsError::Parse { message, .. } => { + assert!(message.contains("empty token stream"), "{message}"); + } + other => panic!("expected parse error, got {other:?}"), + } +} diff --git a/src/language/test/extended_grammar.rs b/src/language/test/extended_grammar.rs new file mode 100644 index 0000000..ad0c0dd --- /dev/null +++ b/src/language/test/extended_grammar.rs @@ -0,0 +1,277 @@ +//! Extended-grammar tests (H2): channels+policy, command, send/join, +//! subgraph, subagent, repl_agent, interrupt, io shape, +//! checkpoint/interrupt policy. +//! +//! Split out of `language/test/mod.rs` by pipeline phase. + +use super::*; + +// --------------------------------------------------------------------------- +// Extended grammar (H2): channels+policy, command, send/join, subgraph, +// subagent, repl_agent, interrupt, io shape, checkpoint/interrupt policy. +// --------------------------------------------------------------------------- + +/// A graph exercising every H2 primitive in one declarative blueprint. +const EXTENDED: &str = r#" +graph orchestrator { + start planner + + input { + request string + customer_id string + } + output { + answer string + } + + checkpoint inherit + interrupt manual + + channel messages messages + channel usage aggregate "usage_delta" + channel arrivals barrier 2 + + node planner { + kind agent + model "default" + command { + goto fanout + update { + status "planned" + } + } + } + + node fanout { + kind model + sends [ + send worker_a "split_a" + send worker_b "split_b" + ] + next worker_a + } + + node worker_a { + kind model + next gather + } + + node worker_b { + kind model + next gather + } + + node gather { + kind join + sources [worker_a, worker_b] + next research + } + + node research { + kind subagent + agent "researcher" + input "topic" + timeout 30 + retry { + max_attempts 3 + backoff "exponential" + } + next sub + } + + node sub { + kind subgraph + graph "summarize" + next triage + } + + node triage { + kind repl_agent + model "default" + script "triage_script" + next review + } + + node review { + kind interrupt + prompt "Approve?" + options ["approve", "reject"] + routes { + approve -> END + reject -> planner + } + } + + join [worker_a, worker_b] -> gather +} +"#; + +#[test] +fn extended_grammar_parses_and_compiles_blueprint_shape() { + let program = parse_str(EXTENDED).unwrap(); + let bp = compile(&program).unwrap().remove(0); + + assert_eq!(bp.graph_id, "orchestrator"); + assert_eq!(bp.start, "planner"); + assert_eq!(bp.checkpoint.as_deref(), Some("inherit")); + assert_eq!(bp.interrupt.as_deref(), Some("manual")); + + // Input/output shape. + assert_eq!(bp.input.len(), 2); + assert_eq!(bp.input[0].name, "request"); + assert_eq!(bp.input[0].ty, "string"); + assert_eq!(bp.output.len(), 1); + assert_eq!(bp.output[0].name, "answer"); + + // Channels carry reducer + policy args. + assert_eq!(bp.channels.len(), 3); + let usage = bp.channels.iter().find(|c| c.name == "usage").unwrap(); + assert_eq!(usage.reducer, "aggregate"); + assert_eq!(usage.args, vec![Literal::Str("usage_delta".into())]); + let arrivals = bp.channels.iter().find(|c| c.name == "arrivals").unwrap(); + assert_eq!(arrivals.reducer, "barrier"); + assert_eq!(arrivals.args, vec![Literal::Num(2.0)]); + + // Command lowering + routing precedence (goto becomes a static next). + let planner = bp.nodes.iter().find(|n| n.name == "planner").unwrap(); + let cmd = planner.command.as_ref().unwrap(); + assert_eq!(cmd.goto.as_deref(), Some("fanout")); + assert_eq!( + cmd.update, + vec![("status".into(), Literal::Str("planned".into()))] + ); + assert_eq!(planner.routing, Routing::Next("fanout".into())); + + // Fanout sends. + let fanout = bp.nodes.iter().find(|n| n.name == "fanout").unwrap(); + assert_eq!(fanout.sends.len(), 2); + assert_eq!(fanout.sends[0].target, "worker_a"); + assert_eq!(fanout.sends[0].input.as_deref(), Some("split_a")); + + // Join node. + let gather = bp.nodes.iter().find(|n| n.name == "gather").unwrap(); + assert_eq!(gather.kind, "join"); + assert_eq!(gather.join_sources, vec!["worker_a", "worker_b"]); + + // Sub-agent node with input mapping + policies. + let research = bp.nodes.iter().find(|n| n.name == "research").unwrap(); + assert_eq!(research.kind, "subagent"); + assert_eq!(research.agent.as_deref(), Some("researcher")); + assert_eq!(research.input.as_deref(), Some("topic")); + assert_eq!(research.timeout.as_deref(), Some("30")); + assert_eq!( + research.retry, + vec![ + ("max_attempts".into(), Literal::Num(3.0)), + ("backoff".into(), Literal::Str("exponential".into())), + ] + ); + + // Subgraph node references a registered graph by name. + let sub = bp.nodes.iter().find(|n| n.name == "sub").unwrap(); + assert_eq!(sub.kind, "subgraph"); + assert_eq!(sub.subgraph.as_deref(), Some("summarize")); + + // REPL-backed node names a script capability (declaration only). + let triage = bp.nodes.iter().find(|n| n.name == "triage").unwrap(); + assert_eq!(triage.kind, "repl_agent"); + assert_eq!(triage.script.as_deref(), Some("triage_script")); + + // Interrupt node with options + conditional routing. + let review = bp.nodes.iter().find(|n| n.name == "review").unwrap(); + assert_eq!(review.kind, "interrupt"); + assert_eq!(review.options, vec!["approve", "reject"]); + assert!(matches!(review.routing, Routing::Conditional(_))); + + // Top-level join declaration. + assert_eq!(bp.joins.len(), 1); + assert_eq!(bp.joins[0].target, "gather"); + assert_eq!(bp.joins[0].sources, vec!["worker_a", "worker_b"]); +} + +#[test] +fn extended_blueprint_round_trips_through_serde() { + let bp = compile(&parse_str(EXTENDED).unwrap()).unwrap().remove(0); + let json = serde_json::to_string(&bp).unwrap(); + let back: crate::language::types::Blueprint = serde_json::from_str(&json).unwrap(); + assert_eq!(bp, back); +} + +#[test] +fn command_goto_unknown_target_is_a_compile_error() { + let src = "graph g { start a node a { command { goto ghost } } }"; + let err = compile(&parse_str(src).unwrap()).unwrap_err(); + assert!(err.to_string().contains("command goto target"), "{err}"); +} + +#[test] +fn send_unknown_target_is_a_compile_error() { + let src = "graph g { start a node a { sends [ send ghost ] } }"; + let err = compile(&parse_str(src).unwrap()).unwrap_err(); + assert!(err.to_string().contains("send target"), "{err}"); +} + +#[test] +fn join_unknown_source_is_a_compile_error() { + let src = "graph g { start a node a { } join [ghost] -> a }"; + let err = compile(&parse_str(src).unwrap()).unwrap_err(); + assert!(err.to_string().contains("join source"), "{err}"); +} + +#[test] +fn extended_kinds_bind_against_a_resolver() { + let bp = compile(&parse_str(EXTENDED).unwrap()).unwrap().remove(0); + let resolver = CapabilityResolver::from_lists(["default".to_string()], std::iter::empty()) + .allow_subgraph("summarize") + .allow_agent("researcher") + .allow_script("triage_script") + .allow_reducer("messages") + .allow_reducer("aggregate") + .allow_reducer("barrier") + .with_node_kinds( + crate::language::capability_resolver::DEFAULT_NODE_KINDS + .iter() + .map(|k| k.to_string()), + ); + resolver.bind_blueprint(&bp).unwrap(); +} + +#[test] +fn bind_blueprint_rejects_unregistered_subagent_and_script() { + // The strict blueprint gate must validate `subagent` agent references and + // `repl_agent` script references, not silently pass them through a model + // check. Both were previously admitted, so this exercises the fail-closed + // path the `Resolver` already covered but `bind_blueprint` did not. + let node_kinds = || { + crate::language::capability_resolver::DEFAULT_NODE_KINDS + .iter() + .map(|k| k.to_string()) + }; + let bp = compile(&parse_str(EXTENDED).unwrap()).unwrap().remove(0); + + // Missing agent `researcher`: rejected with an unknown-agent capability error. + let missing_agent = CapabilityResolver::from_lists(["default".to_string()], std::iter::empty()) + .allow_subgraph("summarize") + .allow_script("triage_script") + .allow_reducer("messages") + .allow_reducer("aggregate") + .allow_reducer("barrier") + .with_node_kinds(node_kinds()); + let err = missing_agent.bind_blueprint(&bp).unwrap_err(); + assert!(matches!(err, crate::error::TinyAgentsError::Capability(_))); + assert!(err.to_string().contains("unknown agent"), "{err}"); + assert!(err.to_string().contains("researcher"), "{err}"); + + // Missing script `triage_script`: rejected with an unknown-script error. + let missing_script = + CapabilityResolver::from_lists(["default".to_string()], std::iter::empty()) + .allow_subgraph("summarize") + .allow_agent("researcher") + .allow_reducer("messages") + .allow_reducer("aggregate") + .allow_reducer("barrier") + .with_node_kinds(node_kinds()); + let err = missing_script.bind_blueprint(&bp).unwrap_err(); + assert!(err.to_string().contains("unknown script"), "{err}"); + assert!(err.to_string().contains("triage_script"), "{err}"); +} diff --git a/src/language/test/graph_materialisation.rs b/src/language/test/graph_materialisation.rs new file mode 100644 index 0000000..91b31b8 --- /dev/null +++ b/src/language/test/graph_materialisation.rs @@ -0,0 +1,113 @@ +//! Graph-materialisation tests: lowering a `Blueprint` into a runnable +//! `CompiledGraph` via `build_graph` and executing it. +//! +//! Split out of `language/test/mod.rs` by pipeline phase. + +use super::*; + +// --------------------------------------------------------------------------- +// Graph materialisation + execution +// --------------------------------------------------------------------------- + +#[derive(Clone, Debug, PartialEq, Eq)] +pub(super) struct TestState { + trail: Vec, + agent_visits: u32, +} + +/// Resolves a conditional route `label` to a *durable* target node id from the +/// blueprint's `(label, target)` table, translating the language `END` sentinel +/// (`"END"`) to the durable graph terminal ([`crate::graph::END`]). Unknown +/// labels fall back to the durable `END`. +fn resolve_durable_target(routes: &[(String, String)], label: &str) -> String { + let target = routes + .iter() + .find(|(l, _)| l == label) + .map(|(_, t)| t.as_str()); + match target { + Some(t) if t != crate::language::types::END => t.to_string(), + _ => crate::graph::END.to_string(), + } +} + +/// A trivial factory that materialises echo/route/end nodes purely from the +/// declarative [`NodeSpec`]. It demonstrates that runnable behaviour comes from +/// Rust, not the source: each node records its name; terminal/`next` nodes +/// commit a whole-state update (static edges route them), and conditional nodes +/// loop once before terminating by emitting an explicit `goto` command. +struct TestFactory; + +impl NodeFactory for TestFactory { + fn make(&self, spec: &NodeSpec) -> crate::error::Result> { + let name = spec.name.clone(); + let routing = spec.routing.clone(); + Ok(Arc::new( + move |mut state: TestState, _ctx: NodeContext| -> NodeFuture { + let name = name.clone(); + let routing = routing.clone(); + Box::pin(async move { + state.trail.push(name.clone()); + let result = match &routing { + // Static edges (Next/Terminal) handle routing; just + // commit the whole-state update. + Routing::Terminal | Routing::Next(_) => NodeResult::Update(state), + Routing::Conditional(routes) => { + state.agent_visits += 1; + // Take the `tool_call -> tools` route until the + // second visit, then take `final -> END`. + let label = if state.agent_visits >= 2 { + "final" + } else { + "tool_call" + }; + let target = resolve_durable_target(routes, label); + NodeResult::Command(Command::goto([target]).with_update(state)) + } + }; + Ok(result) + }) + }, + )) + } +} + +/// Collects the visited node ids into owned strings for comparison. +fn visited_names(run: &crate::graph::GraphExecution) -> Vec { + run.visited.iter().map(ToString::to_string).collect() +} + +#[tokio::test] +async fn build_graph_runs_to_end() { + let bp = compile(&parse_str(SUPPORT_AGENT).unwrap()) + .unwrap() + .remove(0); + let graph = build_graph(&bp, &TestFactory).unwrap(); + + let run = graph + .run(TestState { + trail: Vec::new(), + agent_visits: 0, + }) + .await + .unwrap(); + + // agent -> tools -> agent (ends on second visit). + assert_eq!(visited_names(&run), vec!["agent", "tools", "agent"]); + assert_eq!(run.state.trail, vec!["agent", "tools", "agent"]); + assert_eq!(run.state.agent_visits, 2); +} + +#[tokio::test] +async fn build_graph_handles_linear_terminal() { + let src = "graph g { start a node a { kind model next b } node b { kind model next END } }"; + let bp = compile(&parse_str(src).unwrap()).unwrap().remove(0); + let graph = build_graph(&bp, &TestFactory).unwrap(); + let run = graph + .run(TestState { + trail: Vec::new(), + agent_visits: 0, + }) + .await + .unwrap(); + assert_eq!(visited_names(&run), vec!["a", "b"]); +} diff --git a/src/language/test/lexer.rs b/src/language/test/lexer.rs new file mode 100644 index 0000000..8061547 --- /dev/null +++ b/src/language/test/lexer.rs @@ -0,0 +1,81 @@ +//! Lexer tests (tokenization, escapes, spans, literal formatting). +//! +//! Split out of `language/test/mod.rs` by pipeline phase. + +use super::*; + +// --------------------------------------------------------------------------- +// Lexer +// --------------------------------------------------------------------------- + +#[test] +fn tokenizes_punctuation_and_arrow() { + let tokens = tokenize("a -> b { } [ ] ,").unwrap(); + let kinds: Vec<_> = tokens.into_iter().map(|t| t.token).collect(); + assert_eq!( + kinds, + vec![ + Token::Ident("a".into()), + Token::Arrow, + Token::Ident("b".into()), + Token::LBrace, + Token::RBrace, + Token::LBracket, + Token::RBracket, + Token::Comma, + Token::Eof, + ] + ); +} + +#[test] +fn tokenizes_strings_numbers_and_comments() { + let tokens = tokenize("// comment\n\"hi\\n\" 50 1.5 -3").unwrap(); + let kinds: Vec<_> = tokens.into_iter().map(|t| t.token).collect(); + assert_eq!( + kinds, + vec![ + Token::Str("hi\n".into()), + Token::Num(50.0), + Token::Num(1.5), + Token::Num(-3.0), + Token::Eof, + ] + ); +} + +#[test] +fn tracks_line_and_column_spans() { + let tokens = tokenize("graph\n foo").unwrap(); + assert_eq!(tokens[0].span.line, 1); + assert_eq!(tokens[0].span.column, 1); + assert_eq!(tokens[1].span.line, 2); + assert_eq!(tokens[1].span.column, 3); +} + +#[test] +fn unterminated_string_is_a_parse_error() { + let err = tokenize("\"oops").unwrap_err(); + assert!(matches!(err, crate::error::TinyAgentsError::Parse { .. })); +} + +#[test] +fn invalid_escape_is_a_parse_error() { + let err = tokenize("\"bad\\x\"").unwrap_err(); + assert!(matches!(err, crate::error::TinyAgentsError::Parse { .. })); +} + +#[test] +fn literal_as_display_does_not_saturate_huge_floats() { + // A huge finite float must not be truncated to i64::MAX; it should render + // using the float's own formatting instead. + let huge = Literal::Num(1e30); + assert_eq!(huge.as_display(), format!("{}", 1e30_f64)); + assert_ne!(huge.as_display(), format!("{}", i64::MAX)); + + let nan = Literal::Num(f64::NAN); + assert_eq!(nan.as_display(), "NaN"); + + let inf = Literal::Num(f64::INFINITY); + assert_eq!(inf.as_display(), "inf"); +} diff --git a/src/language/test/mod.rs b/src/language/test/mod.rs new file mode 100644 index 0000000..9b78ff3 --- /dev/null +++ b/src/language/test/mod.rs @@ -0,0 +1,59 @@ +//! Tests for the expressive language pipeline: lexer, parser, compiler, +//! capability binding, and graph materialisation. + +use std::sync::Arc; + +use crate::graph::{Command, NodeContext, NodeFuture, NodeResult}; +use crate::language::capability_resolver::{CapabilityResolver, bind_capabilities}; +use crate::language::compiler::{BoxedNode, NodeFactory, build_graph, compile}; +use crate::language::lexer::tokenize; +use crate::language::parser::{parse, parse_str}; +use crate::language::types::{Literal, NodeSpec, Routing, Token}; + +/// The `support_agent` fixture from the module spec: an agent node with a tool +/// loop plus conditional routing to `END`. +const SUPPORT_AGENT: &str = r#" +// A support workflow with a tool loop. +graph support_agent { + start agent + + defaults { + recursion_limit 50 + backoff "exponential" + checkpoint inherit + } + + channel messages messages + channel tool_calls append + + node agent { + kind agent + model "default" + system "Resolve support requests using tools when useful." + tools ["lookup_user", "create_ticket"] + routes { + tool_call -> tools + final -> END + } + } + + node tools { + kind tool_executor + next agent + } +} +"#; + +mod capability_binding; +mod compiler; +mod diagnostics; +mod extended_grammar; +mod graph_materialisation; +mod lexer; +mod parser; +mod provenance_diff_testkit; +mod registry_binding; +mod resolver; + +use graph_materialisation::TestState; +use registry_binding::{FULL_SOURCE, full_registry}; diff --git a/src/language/test/parser.rs b/src/language/test/parser.rs new file mode 100644 index 0000000..c3cd084 --- /dev/null +++ b/src/language/test/parser.rs @@ -0,0 +1,80 @@ +//! Parser tests. +//! +//! Split out of `language/test/mod.rs` by pipeline phase. + +use super::*; + +// --------------------------------------------------------------------------- +// Parser +// --------------------------------------------------------------------------- + +#[test] +fn parses_support_agent_into_ast() { + let program = parse_str(SUPPORT_AGENT).unwrap(); + assert_eq!(program.graphs.len(), 1); + let graph = &program.graphs[0]; + + assert_eq!(graph.name, "support_agent"); + assert_eq!(graph.start.as_deref(), Some("agent")); + assert_eq!(graph.channels.len(), 2); + assert_eq!(graph.channels[0].name, "messages"); + assert_eq!(graph.channels[0].reducer, "messages"); + assert_eq!(graph.channels[1].reducer, "append"); + + // Defaults preserve declared order and literal kinds. + assert_eq!(graph.defaults.len(), 3); + assert_eq!(graph.defaults[0].0, "recursion_limit"); + assert_eq!(graph.defaults[0].1, Literal::Num(50.0)); + assert_eq!(graph.defaults[1].1, Literal::Str("exponential".into())); + assert_eq!(graph.defaults[2].1, Literal::Ident("inherit".into())); + + assert_eq!(graph.nodes.len(), 2); + let agent = &graph.nodes[0]; + assert_eq!(agent.kind.as_deref(), Some("agent")); + assert_eq!(agent.model.as_deref(), Some("default")); + assert_eq!(agent.tools, vec!["lookup_user", "create_ticket"]); + assert_eq!(agent.routes.len(), 2); + assert_eq!(agent.routes[0].label, "tool_call"); + assert_eq!(agent.routes[0].target, "tools"); + assert_eq!(agent.routes[1].target, "END"); + + let tools = &graph.nodes[1]; + assert_eq!(tools.kind.as_deref(), Some("tool_executor")); + assert_eq!(tools.next.as_deref(), Some("agent")); +} + +#[test] +fn parses_top_level_edge() { + let src = "graph g { start a node a { } node b { } a -> b b -> END }"; + let program = parse_str(src).unwrap(); + let graph = &program.graphs[0]; + assert_eq!(graph.edges.len(), 2); + assert_eq!(graph.edges[0].from, "a"); + assert_eq!(graph.edges[0].to, "b"); + assert_eq!(graph.edges[1].to, "END"); +} + +#[test] +fn parse_reports_unexpected_token() { + // Missing graph name. + let tokens = tokenize("graph { }").unwrap(); + let err = parse(&tokens).unwrap_err(); + match err { + crate::error::TinyAgentsError::Parse { message, .. } => { + assert!(message.contains("expected identifier"), "{message}"); + } + other => panic!("expected parse error, got {other:?}"), + } +} + +#[test] +fn parse_rejects_unknown_node_item() { + let src = "graph g { start a node a { bogus x } }"; + let err = parse_str(src).unwrap_err(); + match err { + crate::error::TinyAgentsError::Parse { message, .. } => { + assert!(message.contains("unknown node item"), "{message}"); + } + other => panic!("expected parse error, got {other:?}"), + } +} diff --git a/src/language/test/provenance_diff_testkit.rs b/src/language/test/provenance_diff_testkit.rs new file mode 100644 index 0000000..3e0fe09 --- /dev/null +++ b/src/language/test/provenance_diff_testkit.rs @@ -0,0 +1,299 @@ +//! Provenance, diff, and language-testkit tests (H4). +//! +//! Split out of `language/test/mod.rs` by pipeline phase. + +use super::*; +use crate::language::source::SourceFile; + +// --------------------------------------------------------------------------- +// Provenance, diff, and language testkit (H4) +// --------------------------------------------------------------------------- + +use crate::language::compiler::compile_with_provenance; +use crate::language::diff::{FieldChange, blueprint_diff}; +use crate::language::testkit as lang_testkit; +use crate::language::types::Origin; + +const DIFF_BASE: &str = r#" +graph flow { + start plan + + channel messages append + + node plan { + kind model + model "default" + routes { + research -> work + done -> END + } + } + + node work { + kind tool_executor + tools ["lookup_user"] + next END + } +} +"#; + +/// Adds a node (`review`) and changes a route target on `plan` +/// (`research -> review` instead of `research -> work`). +const DIFF_NEW: &str = r#" +graph flow { + start plan + + channel messages append + + node plan { + kind model + model "default" + routes { + research -> review + done -> END + } + } + + node review { + kind interrupt + prompt "ok?" + next work + } + + node work { + kind tool_executor + tools ["lookup_user"] + next END + } +} +"#; + +#[test] +fn blueprint_diff_reports_added_node_and_changed_route() { + let old = lang_testkit::blueprint(DIFF_BASE); + let new = lang_testkit::blueprint(DIFF_NEW); + + let diff = blueprint_diff(&old, &new); + assert!(!diff.is_empty()); + + // One node added: `review`. + assert_eq!(diff.nodes_added, vec!["review".to_string()]); + assert!(diff.nodes_removed.is_empty()); + + // `plan`'s routing changed (research target work -> review). + assert_eq!(diff.nodes_changed.len(), 1, "{diff:?}"); + let plan_change = &diff.nodes_changed[0]; + assert_eq!(plan_change.name, "plan"); + assert_eq!( + plan_change.fields, + vec![FieldChange { + field: "routing".to_string(), + old: "{ research -> work, done -> END }".to_string(), + new: "{ research -> review, done -> END }".to_string(), + }] + ); + + // The rendered summary names both the added node and the changed route. + let rendered = diff.to_string(); + assert!(rendered.contains("+ node review"), "{rendered}"); + assert!(rendered.contains("~ node plan"), "{rendered}"); + assert!(rendered.contains("routing:"), "{rendered}"); +} + +#[test] +fn blueprint_diff_of_identical_blueprints_is_empty() { + let a = lang_testkit::blueprint(DIFF_BASE); + let b = lang_testkit::blueprint(DIFF_BASE); + let diff = blueprint_diff(&a, &b); + assert!(diff.is_empty(), "{diff:?}"); + assert_eq!(diff.to_string(), "no changes"); +} + +const DIFF_COMMAND_BASE: &str = r#" +graph flow2 { + start plan + + node plan { + kind agent + model "default" + command { + goto work + update { + status "planned" + } + } + } + + node work { + kind tool_executor + next END + } +} +"#; + +/// Only `plan`'s command `update` value changes (`"planned"` -> `"revised"`); +/// topology, routing target, and everything else stays identical. +const DIFF_COMMAND_NEW: &str = r#" +graph flow2 { + start plan + + node plan { + kind agent + model "default" + command { + goto work + update { + status "revised" + } + } + } + + node work { + kind tool_executor + next END + } +} +"#; + +#[test] +fn blueprint_diff_reports_command_only_change() { + let old = lang_testkit::blueprint(DIFF_COMMAND_BASE); + let new = lang_testkit::blueprint(DIFF_COMMAND_NEW); + + let diff = blueprint_diff(&old, &new); + assert!(!diff.is_empty(), "command-only change must not be silent"); + assert!(diff.nodes_added.is_empty()); + assert!(diff.nodes_removed.is_empty()); + + assert_eq!(diff.nodes_changed.len(), 1, "{diff:?}"); + let plan_change = &diff.nodes_changed[0]; + assert_eq!(plan_change.name, "plan"); + assert!( + plan_change.fields.iter().any(|f| f.field == "command"), + "{plan_change:?}" + ); + + let rendered = diff.to_string(); + assert!(rendered.contains("command:"), "{rendered}"); +} + +#[test] +fn blueprint_diff_serializes_round_trip() { + let old = lang_testkit::blueprint(DIFF_BASE); + let new = lang_testkit::blueprint(DIFF_NEW); + let diff = blueprint_diff(&old, &new); + let json = serde_json::to_string(&diff).unwrap(); + let back: crate::language::diff::BlueprintDiff = serde_json::from_str(&json).unwrap(); + assert_eq!(diff, back); +} + +#[test] +fn provenance_points_each_node_at_its_span() { + let program = parse_str(DIFF_NEW).unwrap(); + let bp = compile_with_provenance(&program, Origin::file("flow.rag")) + .unwrap() + .remove(0); + + let prov = bp.provenance().expect("provenance attached"); + assert_eq!(prov.origin, Origin::File("flow.rag".to_string())); + + // Every node in the blueprint has a recorded span anchored at the line of + // its `node ` declaration, and the byte range slices back to the + // `node` keyword that opens it. + let file = SourceFile::new("flow.rag", DIFF_NEW); + for node in &bp.nodes { + let span = prov + .node_span(&node.name) + .unwrap_or_else(|| panic!("no span for node `{}`", node.name)); + // The span's byte range covers the opening `node` keyword. + assert_eq!(&DIFF_NEW[span.start..span.end], "node"); + // The line the span anchors at is the node's declaration line. + let line = file + .line_text(span.line) + .unwrap_or_else(|| panic!("no source line {}", span.line)); + assert!( + line.contains(&format!("node {}", node.name)), + "node `{}` span anchors at line `{line}`, not its declaration", + node.name + ); + } + + // The channel span is recorded too. + assert!(prov.channel_span("messages").is_some()); +} + +#[test] +fn plain_compile_leaves_provenance_none() { + let bp = lang_testkit::blueprint(DIFF_BASE); + assert!(bp.provenance().is_none()); +} + +#[test] +fn generated_origin_renders_label() { + let program = parse_str(DIFF_BASE).unwrap(); + let bp = compile_with_provenance(&program, Origin::generated_by("repl-7")) + .unwrap() + .remove(0); + let prov = bp.provenance().unwrap(); + assert_eq!(prov.origin.as_display(), "generated by repl-7"); +} + +#[test] +fn testkit_assertions_inspect_lowered_topology() { + let bp = lang_testkit::blueprint(DIFF_NEW); + lang_testkit::assert_kind(&bp, "review", "interrupt"); + lang_testkit::assert_next(&bp, "review", "work"); + lang_testkit::assert_terminal(&bp, "work"); + lang_testkit::assert_route(&bp, "plan", "research", "review"); +} + +#[test] +fn testkit_try_compile_surfaces_errors() { + // `start` references an undefined node, so compilation fails. + let err = lang_testkit::try_compile("graph g { start missing node a { kind model next END } }") + .unwrap_err(); + assert!(matches!(err, crate::error::TinyAgentsError::Compile(_))); +} + +/// Minimal fake model/tool used to populate a [`CapabilityRegistry`] in tests. +pub(super) mod testkit { + use async_trait::async_trait; + use serde_json::json; + + use crate::error::Result; + use crate::harness::model::{ChatModel, ModelRequest, ModelResponse}; + use crate::harness::tool::{Tool, ToolCall, ToolResult, ToolSchema}; + + use super::TestState; + + pub(crate) struct EchoModel; + + #[async_trait] + impl ChatModel for EchoModel { + async fn invoke( + &self, + _state: &TestState, + _request: ModelRequest, + ) -> Result { + Ok(ModelResponse::assistant("echo")) + } + } + + pub(crate) struct NoopTool; + + #[async_trait] + impl Tool for NoopTool { + fn name(&self) -> &str { + "lookup_user" + } + fn description(&self) -> &str { + "noop" + } + fn schema(&self) -> ToolSchema { + ToolSchema::new("lookup_user", "noop", json!({"type": "object"})) + } + async fn call(&self, _state: &TestState, call: ToolCall) -> Result { + Ok(ToolResult::text(call.id, call.name, "ok")) + } + } +} diff --git a/src/language/test/registry_binding.rs b/src/language/test/registry_binding.rs new file mode 100644 index 0000000..813946c --- /dev/null +++ b/src/language/test/registry_binding.rs @@ -0,0 +1,182 @@ +//! Registry-backed capability-binding tests (registry -> language +//! binding via `bind_capabilities_with_registry`/`compile_source`). +//! +//! Split out of `language/test/mod.rs` by pipeline phase. + +use super::*; + +// --------------------------------------------------------------------------- +// Registry-backed capability binding (registry → language binding) +// --------------------------------------------------------------------------- + +use crate::language::capability_resolver::{DEFAULT_NODE_KINDS, bind_capabilities_with_registry}; +use crate::language::compiler::compile_source; +use crate::registry::CapabilityRegistry; + +/// A `.rag` graph that exercises every registry-backed reference kind: a model, +/// a tool, a subgraph reference (`kind subgraph` whose `model` names a +/// registered blueprint), a router reference (`kind router`), and a channel +/// reducer. +pub(super) const FULL_SOURCE: &str = r#" +graph main { + start agent + + channel messages append + + node agent { + kind agent + model "default" + tools ["lookup_user"] + routes { + retrieve -> sub + classify -> route + done -> END + } + } + + node sub { + kind subgraph + model "retrieval" + next END + } + + node route { + kind router + model "classify" + next END + } +} +"#; + +/// Builds a registry that satisfies every reference in [`FULL_SOURCE`]. +pub(super) fn full_registry() -> CapabilityRegistry { + let mut reg = CapabilityRegistry::::new(); + reg.register_model( + "default", + std::sync::Arc::new(crate::language::test::provenance_diff_testkit::testkit::EchoModel), + ) + .unwrap(); + reg.register_tool(std::sync::Arc::new( + crate::language::test::provenance_diff_testkit::testkit::NoopTool, + )) + .unwrap(); + reg.register_graph_blueprint( + "retrieval", + compile(&parse_str("graph retrieval { start r node r { kind model next END } }").unwrap()) + .unwrap() + .remove(0), + ) + .unwrap(); + reg.register_router("classify").unwrap(); + reg.register_reducer("append").unwrap(); + reg +} + +#[test] +fn compile_source_binds_against_registry() { + let reg = full_registry(); + let blueprints = compile_source(FULL_SOURCE, ®).unwrap(); + assert_eq!(blueprints.len(), 1); + assert_eq!(blueprints[0].graph_id, "main"); +} + +#[test] +fn registry_resolver_allows_all_kinds() { + let reg = full_registry(); + let resolver = reg.capability_resolver(); + assert!(resolver.model_allowed("default")); + assert!(resolver.tool_allowed("lookup_user")); + assert!(resolver.subgraph_allowed("retrieval")); + assert!(resolver.router_allowed("classify")); + assert!(resolver.reducer_allowed("append")); + for kind in DEFAULT_NODE_KINDS { + assert!(resolver.node_kind_allowed(kind)); + } +} + +#[test] +fn registry_bind_rejects_unregistered_model() { + let mut reg = full_registry(); + reg.replace_model( + "other", + std::sync::Arc::new(crate::language::test::provenance_diff_testkit::testkit::EchoModel), + ); + // Source references `default`, which we did not register here. + let mut bare = CapabilityRegistry::::new(); + bare.register_tool(std::sync::Arc::new( + crate::language::test::provenance_diff_testkit::testkit::NoopTool, + )) + .unwrap(); + bare.register_graph_blueprint( + "retrieval", + reg.graph_blueprint("retrieval").unwrap().clone(), + ) + .unwrap(); + bare.register_router("classify").unwrap(); + bare.register_reducer("append").unwrap(); + let err = compile_source(FULL_SOURCE, &bare).unwrap_err(); + assert!(matches!(err, crate::error::TinyAgentsError::Capability(_))); + assert!(err.to_string().contains("unknown model"), "{err}"); +} + +#[test] +fn registry_bind_rejects_unregistered_tool() { + let src = r#"graph g { start a channel m append node a { kind agent model "default" tools ["missing"] next END } }"#; + let reg = full_registry(); + let err = compile_source(src, ®).unwrap_err(); + assert!(err.to_string().contains("unknown tool"), "{err}"); +} + +#[test] +fn registry_bind_rejects_unregistered_subgraph() { + let src = r#"graph g { start s node s { kind subgraph model "ghost" next END } }"#; + let reg = full_registry(); + let err = compile_source(src, ®).unwrap_err(); + assert!(matches!(err, crate::error::TinyAgentsError::Capability(_))); + assert!(err.to_string().contains("unknown subgraph"), "{err}"); +} + +#[test] +fn registry_bind_rejects_unregistered_router() { + let src = r#"graph g { start r node r { kind router model "ghost" next END } }"#; + let reg = full_registry(); + let err = compile_source(src, ®).unwrap_err(); + assert!(err.to_string().contains("unknown router"), "{err}"); +} + +#[test] +fn registry_bind_rejects_unregistered_reducer() { + let src = r#"graph g { start a channel messages ghost node a { kind model next END } }"#; + let reg = full_registry(); + let err = compile_source(src, ®).unwrap_err(); + assert!(matches!(err, crate::error::TinyAgentsError::Capability(_))); + assert!(err.to_string().contains("unknown reducer"), "{err}"); +} + +#[test] +fn registry_bind_rejects_unknown_node_kind() { + let src = r#"graph g { start a node a { kind wizard next END } }"#; + let reg = full_registry(); + let err = compile_source(src, ®).unwrap_err(); + assert!(matches!(err, crate::error::TinyAgentsError::Compile(_))); + assert!(err.to_string().contains("unknown kind"), "{err}"); +} + +#[test] +fn manual_bind_path_ignores_kinds_and_reducers() { + // The legacy manual resolver must keep working: a non-empty node kind set is + // never consulted, and reducers/subgraphs are not checked. + let src = r#"graph g { start a channel messages ghost node a { kind wizard model "default" next END } }"#; + let bp = compile(&parse_str(src).unwrap()).unwrap().remove(0); + let resolver = CapabilityResolver::new().allow_model("default"); + // Manual gate only checks model + tool; passes despite the unknown kind, + // unknown reducer, and exotic node kind. + bind_capabilities(&bp, &resolver).unwrap(); +} + +#[test] +fn bind_capabilities_with_registry_matches_compile_source() { + let reg = full_registry(); + let bp = compile(&parse_str(FULL_SOURCE).unwrap()).unwrap().remove(0); + bind_capabilities_with_registry(&bp, ®).unwrap(); +} diff --git a/src/language/test/resolver.rs b/src/language/test/resolver.rs new file mode 100644 index 0000000..a2aaf7a --- /dev/null +++ b/src/language/test/resolver.rs @@ -0,0 +1,138 @@ +//! Registry-backed `Resolver` tests (H3): spanned diagnostics, the +//! single binding gate. +//! +//! Split out of `language/test/mod.rs` by pipeline phase. + +use super::*; +use crate::language::source::SourceFile; + +// --------------------------------------------------------------------------- +// Registry-backed Resolver (H3): spanned diagnostics, single binding gate +// --------------------------------------------------------------------------- + +use crate::language::resolver::{Resolver, resolve_source}; + +#[test] +fn resolver_accepts_fully_registered_blueprint() { + let reg = full_registry(); + let program = parse_str(FULL_SOURCE).unwrap(); + let resolver = Resolver::from_registry(®); + // No diagnostics: every model/tool/subgraph/router/reducer is registered. + assert!(resolver.resolve_program(&program).is_empty()); + // And the convenience façade lowers it to a blueprint. + let blueprints = resolve_source(FULL_SOURCE, ®).unwrap(); + assert_eq!(blueprints[0].graph_id, "main"); +} + +#[test] +fn resolver_reports_unregistered_tool_with_spanned_diagnostic() { + // `missing` is not a registered tool. + let src = r#" +graph g { + start a + channel m append + node a { + kind agent + model "default" + tools ["missing"] + next END + } +} +"#; + let reg = full_registry(); + let program = parse_str(src).unwrap(); + let file = SourceFile::new("plan.rag", src); + let resolver = Resolver::from_registry(®); + + let diagnostics = resolver.resolve_program(&program); + assert_eq!(diagnostics.len(), 1, "{diagnostics:?}"); + let diag = &diagnostics[0]; + assert_eq!(diag.code.as_deref(), Some("E-rag-unknown-tool")); + let rendered = diag.render(&file); + assert!( + rendered.contains("node `a` references unknown tool `missing`"), + "{rendered}" + ); + // The diagnostic carries a caret pointing at the offending node span. + assert!(rendered.contains('^'), "{rendered}"); + assert!(rendered.contains("--> plan.rag:"), "{rendered}"); + + // `check_program` folds it into a Capability error with the rendered caret. + let err = resolver.check_program(&program, Some(&file)).unwrap_err(); + assert!(matches!(err, crate::error::TinyAgentsError::Capability(_))); + assert!(err.to_string().contains("unknown tool"), "{err}"); + assert!(err.to_string().contains('^'), "{err}"); +} + +#[test] +fn resolve_source_rejects_unregistered_tool() { + let src = r#"graph g { start a channel m append node a { kind agent model "default" tools ["missing"] next END } }"#; + let reg = full_registry(); + let err = resolve_source(src, ®).unwrap_err(); + assert!(matches!(err, crate::error::TinyAgentsError::Capability(_))); + assert!(err.to_string().contains("unknown tool"), "{err}"); +} + +#[test] +fn resolver_reports_unknown_node_kind_as_compile_error() { + let src = r#"graph g { start a node a { kind wizard next END } }"#; + let reg = full_registry(); + let err = resolve_source(src, ®).unwrap_err(); + assert!(matches!(err, crate::error::TinyAgentsError::Compile(_))); + assert!(err.to_string().contains("unknown kind"), "{err}"); +} + +#[test] +fn resolver_reports_unregistered_agent() { + // A `subagent` node binds its `agent "…"` reference through the registry's + // Agent allowlist. + let src = r#"graph g { start a node a { kind subagent agent "ghost" next END } }"#; + let reg = full_registry(); + let program = parse_str(src).unwrap(); + let diagnostics = Resolver::from_registry(®).resolve_program(&program); + assert_eq!(diagnostics.len(), 1, "{diagnostics:?}"); + assert_eq!(diagnostics[0].code.as_deref(), Some("E-rag-unknown-agent")); + assert!( + diagnostics[0].message.contains("unknown agent `ghost`"), + "{:?}", + diagnostics[0] + ); +} + +#[test] +fn resolver_collects_multiple_diagnostics() { + // Two independent problems: an unregistered model and an unregistered + // reducer. `resolve_program` reports both. + let src = r#"graph g { start a channel m ghost node a { kind model model "nope" next END } }"#; + let reg = full_registry(); + let program = parse_str(src).unwrap(); + let diagnostics = Resolver::from_registry(®).resolve_program(&program); + let codes: Vec<_> = diagnostics + .iter() + .filter_map(|d| d.code.as_deref()) + .collect(); + assert!(codes.contains(&"E-rag-unknown-model"), "{codes:?}"); + assert!(codes.contains(&"E-rag-unknown-reducer"), "{codes:?}"); +} + +#[test] +fn resolver_blueprint_path_matches_registry_binding() { + // The span-less blueprint path mirrors the legacy gate's variants/messages. + let reg = full_registry(); + let bp = compile(&parse_str(FULL_SOURCE).unwrap()).unwrap().remove(0); + Resolver::from_registry(®) + .resolve_blueprint(&bp) + .unwrap(); + + let bad = compile( + &parse_str(r#"graph g { start a channel m append node a { kind subgraph model "ghost" next END } }"#) + .unwrap(), + ) + .unwrap() + .remove(0); + let err = Resolver::from_registry(®) + .resolve_blueprint(&bad) + .unwrap_err(); + assert!(matches!(err, crate::error::TinyAgentsError::Capability(_))); + assert!(err.to_string().contains("unknown subgraph"), "{err}"); +} diff --git a/src/lib.rs b/src/lib.rs index 9060a35..a0456b1 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -53,10 +53,13 @@ //! //! ## Provider features //! -//! The default build is offline and deterministic ([`harness::providers::MockModel`]). //! Hosted and local providers (OpenAI plus the OpenAI-compatible endpoints for //! Anthropic, Ollama, DeepSeek, Groq, xAI, OpenRouter, Together, and Mistral) -//! live behind the `openai` Cargo feature. +//! are compiled in unconditionally alongside the offline, deterministic +//! [`harness::providers::MockModel`]. Two Cargo features gate optional, +//! heavier dependencies instead: `sqlite` (embedded SQLite checkpointer, +//! [`graph::checkpoint::SqliteCheckpointer`]) and `repl` (embedded Rhai engine +//! powering the `.ragsh` session runtime, [`repl::session`]). //! //! ## Crate-root re-exports //! @@ -76,18 +79,19 @@ pub use error::{Result, TinyAgentsError}; // --- Registry: named capability catalog (.rag/.ragsh binding by name) --- pub use registry::{ - CapabilityRegistry, ComponentId, ComponentKind, ComponentMetadata, DiagnosticSeverity, - RegistryDiagnostic, RegistrySnapshot, + AliasBinding, CapabilityRegistry, ComponentId, ComponentKind, ComponentMetadata, + DiagnosticSeverity, ModelCapabilities, ModelCatalog, ModelCatalogEntry, ModelCatalogSnapshot, + ModelCatalogSource, ModelPricing, RegistryDiagnostic, RegistrySnapshot, }; // --- Language: registry → blueprint binding façade --- // The strict, registry-backed entry points the REPL and orchestrators use to // turn `.rag`/`.ragsh` source into validated blueprints. `compile_source` runs // parse -> compile -> registry-bind in one call. -pub use language::compiler::{ - CapabilityResolver, bind_capabilities, bind_capabilities_with_registry, compile, - compile_source, compile_with_provenance, +pub use language::capability_resolver::{ + CapabilityResolver, bind_capabilities, bind_capabilities_with_registry, }; +pub use language::compiler::{compile, compile_source, compile_with_provenance}; // `Resolver` is the registry-backed binding gate: it resolves every reference in // a `.rag` plan (file-backed or model-generated) against the registry, producing // spanned diagnostics for unknown/disallowed capabilities. `resolve_source` is @@ -146,12 +150,13 @@ pub use harness::observability::{LangfuseAuth, LangfuseClient, LangfuseTraceConf #[cfg(feature = "sqlite")] pub use graph::SqliteCheckpointer; pub use graph::{ - Checkpoint, CheckpointConfig, CheckpointMetadata, CheckpointSource, CheckpointTuple, - Checkpointer, ChildRun, ChildRunSink, ClosureReducer, ClosureStateReducer, Command, - CompiledGraph, DurabilityMode, END, FileCheckpointer, ForkId, GraphBuilder, GraphDefaults, - GraphEvent, GraphExecution, GraphInput, GraphRunStatus, InMemoryCheckpointer, Interrupt, - NodeContext, NodeResult, RecursionFrame, RecursionPolicy, RecursionStack, Reducer, - ResumeTarget, Route, RouteTarget, RunTree, START, StateReducer, StateSnapshot, + BarrierArrivals, Checkpoint, CheckpointConfig, CheckpointMetadata, CheckpointSource, + CheckpointTuple, Checkpointer, ChildRun, ChildRunSink, ClosureReducer, ClosureStateReducer, + Command, CompiledGraph, DurabilityMode, END, FileCheckpointer, ForkId, GraphBuilder, + GraphDefaults, GraphEvent, GraphExecution, GraphInput, GraphRunStatus, InMemoryCheckpointer, + Interrupt, NodeContext, NodeResult, PendingActivation, RecursionFrame, RecursionPolicy, + RecursionStack, Reducer, ResumeTarget, Route, RouteTarget, RunTree, START, StateReducer, + StateSnapshot, }; // --- Graph: sub-agent nodes (delegate a graph step to a registered agent) --- diff --git a/src/registry/capability/mod.rs b/src/registry/capability/mod.rs index 92844bd..b79ce6c 100644 --- a/src/registry/capability/mod.rs +++ b/src/registry/capability/mod.rs @@ -26,7 +26,7 @@ use crate::error::{Result, TinyAgentsError}; use crate::harness::model::{ChatModel, ModelRegistry}; use crate::harness::tool::{Tool, ToolRegistry}; use crate::language::Blueprint; -use crate::language::compiler::CapabilityResolver; +use crate::language::capability_resolver::CapabilityResolver; use crate::registry::component::{ComponentKind, ComponentMetadata}; pub use types::*; @@ -240,8 +240,15 @@ impl CapabilityRegistry { /// Registers a name-only descriptor of an arbitrary [`ComponentKind`]. This /// backs [`register_router`](Self::register_router) and - /// [`register_reducer`](Self::register_reducer) and is exposed for the - /// reserved [`ComponentKind::Store`] / [`ComponentKind::Agent`] kinds. + /// [`register_reducer`](Self::register_reducer), and is the general + /// public fallback for every other kind that has no dedicated typed + /// registration method — [`ComponentKind::Store`], + /// [`ComponentKind::Script`], [`ComponentKind::Middleware`], + /// [`ComponentKind::Checkpointer`], [`ComponentKind::TaskStore`], and + /// [`ComponentKind::Listener`]. [`ComponentKind::Model`], + /// [`ComponentKind::Tool`], [`ComponentKind::Graph`], and + /// [`ComponentKind::Agent`] have their own dedicated `register_*` methods + /// instead. /// /// # Errors /// diff --git a/src/registry/component/mod.rs b/src/registry/component/mod.rs index 0e76104..d89fc83 100644 --- a/src/registry/component/mod.rs +++ b/src/registry/component/mod.rs @@ -49,7 +49,7 @@ impl From for ComponentId { impl ComponentKind { /// All component kinds, in a stable order, for discovery iteration. - pub const ALL: [ComponentKind; 11] = [ + pub const ALL: [ComponentKind; 12] = [ ComponentKind::Model, ComponentKind::Tool, ComponentKind::Graph, @@ -57,6 +57,7 @@ impl ComponentKind { ComponentKind::Reducer, ComponentKind::Store, ComponentKind::Agent, + ComponentKind::Script, ComponentKind::Middleware, ComponentKind::Checkpointer, ComponentKind::TaskStore, @@ -74,6 +75,7 @@ impl ComponentKind { ComponentKind::Reducer => "reducer", ComponentKind::Store => "store", ComponentKind::Agent => "agent", + ComponentKind::Script => "script", ComponentKind::Middleware => "middleware", ComponentKind::Checkpointer => "checkpointer", ComponentKind::TaskStore => "task_store", diff --git a/src/registry/component/types.rs b/src/registry/component/types.rs index 0516b76..0f625bf 100644 --- a/src/registry/component/types.rs +++ b/src/registry/component/types.rs @@ -41,6 +41,8 @@ pub enum ComponentKind { Store, /// An executable agent configuration descriptor (name-only for now). Agent, + /// A REPL script descriptor a `repl_agent` node may reference (name-only). + Script, /// A middleware descriptor (name-only for now). Middleware, /// A graph/harness checkpointer descriptor (name-only for now). diff --git a/src/registry/diagnostics.rs b/src/registry/diagnostics.rs index 3294e94..78d1c7d 100644 --- a/src/registry/diagnostics.rs +++ b/src/registry/diagnostics.rs @@ -78,11 +78,12 @@ impl RegistrySnapshot { kind_label(kind) )); for meta in members { + let escaped_id = escape_dot_string(&meta.id.0); out.push_str(&format!( " \"{}:{}\" [label=\"{}\"];\n", kind_label(kind), - meta.id.0, - meta.id.0 + escaped_id, + escaped_id )); } out.push_str(" }\n"); @@ -119,6 +120,12 @@ fn kind_label(kind: ComponentKind) -> &'static str { kind.as_str() } +/// Escapes backslashes and double quotes so a component id can be embedded in +/// a Graphviz DOT quoted string/id without breaking the surrounding syntax. +fn escape_dot_string(value: &str) -> String { + value.replace('\\', "\\\\").replace('"', "\\\"") +} + pub(crate) fn alias_shadows_component(kind: ComponentKind, alias: &str) -> RegistryDiagnostic { RegistryDiagnostic { severity: DiagnosticSeverity::Warning, @@ -165,3 +172,31 @@ pub(crate) fn dangling_alias( ), } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::registry::component::ComponentId; + + #[test] + fn to_dot_escapes_quotes_and_backslashes_in_component_ids() { + let snapshot = RegistrySnapshot { + components: vec![ComponentMetadata { + id: ComponentId(r#"weird"name\with\backslashes"#.to_string()), + kind: ComponentKind::Tool, + description: None, + tags: Vec::new(), + aliases: Vec::new(), + }], + aliases: Vec::new(), + }; + + let dot = snapshot.to_dot(); + + // The raw id must never appear unescaped inside the DOT output, or a + // malicious/unlucky component name could break out of its quoted + // string and inject arbitrary DOT syntax. + assert!(!dot.contains("\"weird\"name")); + assert!(dot.contains(r#"weird\"name\\with\\backslashes"#)); + } +} diff --git a/src/registry/mod.rs b/src/registry/mod.rs index 4aa6ac1..631108e 100644 --- a/src/registry/mod.rs +++ b/src/registry/mod.rs @@ -32,4 +32,4 @@ pub use catalog::{ ModelPricing, }; pub use component::{ComponentId, ComponentKind, ComponentMetadata}; -pub use diagnostics::{DiagnosticSeverity, RegistryDiagnostic, RegistrySnapshot}; +pub use diagnostics::{AliasBinding, DiagnosticSeverity, RegistryDiagnostic, RegistrySnapshot}; diff --git a/src/repl/mod.rs b/src/repl/mod.rs index bef206b..34a2027 100644 --- a/src/repl/mod.rs +++ b/src/repl/mod.rs @@ -106,10 +106,10 @@ mod test; pub fn parse_command(line: &str) -> crate::error::Result { let trimmed = line.trim(); if trimmed.is_empty() { - return Err(parse_err("empty input", 0, 0)); + return Err(parse_err_at(trimmed, trimmed, "empty input")); } - let (verb, rest) = split_token(trimmed)?; + let (verb, rest) = split_token(trimmed, trimmed)?; match verb.to_lowercase().as_str() { "help" | "?" => Ok(ReplCommand::Help), @@ -117,53 +117,62 @@ pub fn parse_command(line: &str) -> crate::error::Result { "quit" | "exit" | "q" => Ok(ReplCommand::Quit), "load" => { - let (path, _) = require_token(rest, "load ")?; + let (path, _) = require_token(trimmed, rest, "load ")?; Ok(ReplCommand::Load { path }) } "compile" => { - let (name, _) = require_token(rest, "compile ")?; + let (name, _) = require_token(trimmed, rest, "compile ")?; Ok(ReplCommand::Compile { name }) } "run" => { - let (graph, rest) = require_token(rest, "run ")?; - let (input, _) = require_token(rest, "run ")?; + let (graph, rest) = require_token(trimmed, rest, "run ")?; + let (input, _) = require_token(trimmed, rest, "run ")?; Ok(ReplCommand::Run { graph, input }) } "set" => { - let (key, rest) = require_token(rest, "set ")?; - let (value, _) = require_token(rest, "set ")?; + let (key, rest) = require_token(trimmed, rest, "set ")?; + let (value, _) = require_token(trimmed, rest, "set ")?; Ok(ReplCommand::Set { key, value }) } "get" => { - let (key, _) = require_token(rest, "get ")?; + let (key, _) = require_token(trimmed, rest, "get ")?; Ok(ReplCommand::Get { key }) } "show" => { - let (what, _) = require_token(rest, "show ")?; + let (what, _) = require_token(trimmed, rest, "show ")?; Ok(ReplCommand::Show { what }) } "call" => { - let (capability, json_rest) = require_token(rest, "call ")?; + let (capability, json_rest) = require_token(trimmed, rest, "call ")?; let json_str = json_rest.trim(); if json_str.is_empty() { - return Err(parse_err( + return Err(parse_err_at( + trimmed, + json_rest, "call requires a JSON argument: call ", - 0, - 0, )); } - let args: serde_json::Value = serde_json::from_str(json_str) - .map_err(|e| parse_err(&format!("invalid JSON argument for `call`: {e}"), 0, 0))?; + let args: serde_json::Value = serde_json::from_str(json_str).map_err(|e| { + parse_err_at( + trimmed, + json_str, + &format!("invalid JSON argument for `call`: {e}"), + ) + })?; Ok(ReplCommand::Call { capability, args }) } - other => Err(parse_err(&format!("unknown command `{other}`"), 0, 0)), + other => Err(parse_err_at( + trimmed, + trimmed, + &format!("unknown command `{other}`"), + )), } } @@ -179,16 +188,48 @@ fn parse_err(message: &str, line: usize, column: usize) -> crate::error::TinyAge } } +/// Builds a [`crate::error::TinyAgentsError::Parse`] pointing at `at` — a +/// substring slice of `origin` — reporting a real 1-based line/column instead +/// of the placeholder `(0, 0)`. +/// +/// `parse_command` always parses a single command line, so the line is always +/// `1`; the column is the 1-based character offset of `at` within `origin`. +/// Falls back to the end of `origin` if `at` is not actually a subslice of it +/// (defensive; should not happen given how callers use this). +fn parse_err_at(origin: &str, at: &str, message: &str) -> crate::error::TinyAgentsError { + let column = char_column(origin, at); + parse_err(message, 1, column) +} + +/// Computes the 1-based character column of `at` within `origin`, where `at` +/// is a substring slice of `origin` obtained by slicing (not copying). +fn char_column(origin: &str, at: &str) -> usize { + let origin_start = origin.as_ptr() as usize; + let origin_end = origin_start + origin.len(); + let at_start = at.as_ptr() as usize; + let offset = if at_start >= origin_start && at_start <= origin_end { + at_start - origin_start + } else { + // `at` isn't a subslice of `origin` (shouldn't happen); point at the end. + origin.len() + }; + origin[..offset].chars().count() + 1 +} + /// Split the next token from `s`, returning `(token, remainder)`. /// /// Handles quoted strings (`"..."`) with `\\`, `\"`, `\n`, `\t` escapes. /// Bare tokens end at the first whitespace character. /// /// Returns `None` if `s` (after trimming leading whitespace) is empty. -fn split_token(s: &str) -> crate::error::Result<(String, &str)> { +/// +/// `origin` is the full trimmed command line `s` was sliced from; it is used +/// only to compute a real 1-based column for any parse error, via +/// [`parse_err_at`]. +fn split_token<'a>(origin: &str, s: &'a str) -> crate::error::Result<(String, &'a str)> { let s = s.trim_start(); if s.is_empty() { - return Err(parse_err("unexpected end of input", 0, 0)); + return Err(parse_err_at(origin, s, "unexpected end of input")); } if s.starts_with('"') { @@ -206,7 +247,7 @@ fn split_token(s: &str) -> crate::error::Result<(String, &str)> { loop { match chars.next() { None => { - return Err(parse_err("unterminated quoted string", 0, 0)); + return Err(parse_err_at(origin, s, "unterminated quoted string")); } Some((i, '"')) => break 'scan i, Some((_, '\\')) => match chars.next() { @@ -219,7 +260,7 @@ fn split_token(s: &str) -> crate::error::Result<(String, &str)> { token.push(c); } None => { - return Err(parse_err("unterminated escape sequence", 0, 0)); + return Err(parse_err_at(origin, s, "unterminated escape sequence")); } }, Some((_, c)) => token.push(c), @@ -245,14 +286,18 @@ fn split_token(s: &str) -> crate::error::Result<(String, &str)> { /// Like [`split_token`] but returns a [`crate::error::TinyAgentsError::Parse`] /// mentioning the expected usage if the remaining input is empty. -fn require_token<'a>(s: &'a str, usage: &str) -> crate::error::Result<(String, &'a str)> { +fn require_token<'a>( + origin: &str, + s: &'a str, + usage: &str, +) -> crate::error::Result<(String, &'a str)> { let s = s.trim_start(); if s.is_empty() { - return Err(parse_err( + return Err(parse_err_at( + origin, + s, &format!("missing argument — usage: {usage}"), - 0, - 0, )); } - split_token(s) + split_token(origin, s) } diff --git a/src/repl/session/builtins.rs b/src/repl/session/builtins.rs deleted file mode 100644 index 1d9a834..0000000 --- a/src/repl/session/builtins.rs +++ /dev/null @@ -1,1009 +0,0 @@ -//! Capability-bound built-in functions for the Rhai-backed `.ragsh` session -//! (design milestones R3–R5). -//! -//! This module registers the reserved built-in functions on a session's -//! [`rhai::Engine`] as **host capabilities**: each one resolves a name through -//! the session's [`CapabilityRegistry`], enforces the [`ReplPolicy`] call and -//! recursion limits, records a [`ReplCallRecord`], and lowers to the real -//! harness/graph runtime. -//! -//! The surface registered here is: -//! -//! - model calls — `model_query`, `model_query_batched` -//! - agent calls — `agent_query`, `agent_query_batched` -//! - graph runs — `graph_run`, `graph_run_batched` -//! - tool calls — `tool_call`, `tool_call_batched` -//! - session built-ins — `emit`, `show_vars`, `answer`, plus `print`/`debug` -//! capture -//! - graph authoring — `graph_define`, `graph_validate`, `graph_compile`, -//! `graph_diff`, `graph_register`, which lower through the Cluster H `.rag` -//! compiler and capability resolver. Generated topology is never installed -//! directly: a draft only becomes `compiled` after the resolver binds it, and -//! `graph_register` honors [`ReplPolicy::generated_graphs_require_review`]. -//! -//! ## The async adapter (blocking bridge) -//! -//! The Rhai engine is synchronous, but model, tool, agent, and graph calls are -//! async. This slice uses the design's **blocking bridge** for v1: a host -//! function builds the async future and drives it to completion in place with -//! [`futures::executor::block_on`] (see [`bridge_block_on`]). This keeps the -//! capability boundary deterministic for scripted tests (a [`ScriptedModel`] or -//! [`FakeTool`] resolves without yielding to a reactor) and works for real -//! providers when the session is driven from a multi-threaded runtime, where -//! blocking one worker does not starve the call's own I/O. -//! -//! The design's longer-term direction is *command recording* (host functions -//! emit `ReplCommand` values the async runtime executes after the cell); the -//! public `ReplResult`/`ReplCallRecord` types are already shaped for it. The -//! bridge is intentionally the only blocking surface and is confined to this -//! module. -//! -//! [`ScriptedModel`]: crate::harness::testkit::ScriptedModel -//! [`FakeTool`]: crate::harness::testkit::FakeTool - -use std::collections::BTreeMap; -use std::future::Future; -use std::sync::{Arc, Mutex}; -use std::time::{Duration, Instant}; - -use rhai::{Array, Dynamic, Engine, EvalAltResult, Map, Position}; -use serde_json::{Value, json}; - -use super::types::{GraphBlueprintHandle, LanguageCompiler, ReplPolicy}; -use super::{ - ReplCallKind, ReplCallRecord, dynamic_to_repl_value, json_to_repl_value, repl_value_to_dynamic, -}; -use crate::error::TinyAgentsError; -use crate::harness::events::EventSink; -use crate::harness::ids::new_call_id; -use crate::harness::message::Message; -use crate::harness::model::ModelRequest; -use crate::harness::tool::ToolCall; -use crate::language::compiler::compile_with_provenance; -use crate::language::parser::parse_str; -use crate::language::resolver::Resolver; -use crate::language::types::Origin; -use crate::language::{Blueprint, blueprint_diff}; -use crate::registry::CapabilityRegistry; - -/// Session-cumulative counters for capability calls, enforced against the -/// `ReplPolicy` `max_*_calls` limits. Counts accumulate across cells (the -/// limits are documented per session) and are shared with every capability -/// closure on the engine. -#[derive(Debug, Default, Clone, Copy)] -pub(super) struct CallCounters { - /// `model_query` (and per-item `model_query_batched`) calls made. - pub model: usize, - /// `tool_call` (and per-item `tool_call_batched`) calls made. - pub tool: usize, - /// `graph_run` (and per-item `graph_run_batched`) calls made. - pub graph: usize, - /// `agent_query` (and per-item `agent_query_batched`) calls made. - pub agent: usize, - /// `graph_define` blueprints drafted. - pub graph_def: usize, -} - -/// The host-side context shared (via `Arc`) with every capability closure on a -/// session's engine. Holds the live registries, application state, policy, and -/// the shared per-cell buffers / session counters / graph drafts. -pub(super) struct HostContext { - /// The unified capability catalog (models, tools, graphs, agents). - pub registry: Arc>, - /// The application state capability calls are invoked against. - pub state: Arc, - /// The session policy (call/recursion/concurrency limits). - pub policy: ReplPolicy, - /// Optional expressive-language compiler handle (provenance label). - pub language: Option, - /// The session id, used as the generated-graph provenance label. - pub session_label: String, - /// The session's run depth, the parent depth for recursive sub-runs. - pub run_depth: usize, - /// The event sink shared with the run context. - pub events: EventSink, - /// Per-cell shared buffers (stdout, calls, answer, host error, vars). - pub buffers: super::CellBuffers, - /// Session-cumulative call counters. - pub counters: Arc>, - /// Graph blueprints drafted this session, keyed by name. - pub drafts: Arc>>, -} - -/// Drives an async capability future to completion synchronously — the v1 -/// "blocking bridge" adapter (see the [module docs](self)). -fn bridge_block_on(future: F) -> F::Output { - futures::executor::block_on(future) -} - -/// One completed `model_query_batched` item: `(model, text, finish_reason, -/// structured, elapsed)`. -type ModelBatchItem = (String, String, Option, bool, Duration); - -/// One completed `agent_query_batched` item: `(agent, text, elapsed)`. -type AgentBatchItem = (String, String, Duration); - -// ── Error / recording helpers ─────────────────────────────────────────────── - -/// Stashes the precise crate error so `eval_cell` can surface it verbatim, and -/// returns the stringly-typed Rhai runtime error the engine propagates. -fn raise(ctx: &HostContext, err: TinyAgentsError) -> Box { - let message = err.to_string(); - ctx.buffers.set_host_error(err); - Box::new(EvalAltResult::ErrorRuntime( - Dynamic::from(message), - Position::NONE, - )) -} - -/// Raises a [`TinyAgentsError::Validation`] for an invalid script argument. -fn invalid( - ctx: &HostContext, - message: impl Into, -) -> Box { - raise(ctx, TinyAgentsError::Validation(message.into())) -} - -/// Records a capability call (or emitted event) into the per-cell buffer. -fn record( - ctx: &HostContext, - kind: ReplCallKind, - name: &str, - detail: Value, - elapsed: Duration, -) { - ctx.buffers.push_call(ReplCallRecord { - call_id: new_call_id(), - kind, - name: name.to_string(), - detail, - elapsed, - }); -} - -// ── Map argument helpers ──────────────────────────────────────────────────── - -/// Reads a string field from a Rhai object map argument. -fn map_str(map: &Map, key: &str) -> Option { - map.get(key).and_then(|d| d.clone().into_string().ok()) -} - -/// Reads a boolean field from a Rhai object map argument. -fn map_bool(map: &Map, key: &str) -> Option { - map.get(key).and_then(|d| d.as_bool().ok()) -} - -/// Converts a Rhai object map argument into a JSON value (for tool arguments -/// and structured payloads). -fn map_json(map: &Map, key: &str) -> Option { - map.get(key) - .map(|d| dynamic_to_repl_value(d).to_json()) - .filter(|v| !v.is_null()) -} - -// ── Counter limit helpers ─────────────────────────────────────────────────── - -/// Increments and bounds the model-call counter. -fn bump_model(ctx: &HostContext) -> Result<(), Box> { - let mut counters = ctx.counters.lock().expect("counters poisoned"); - if counters.model >= ctx.policy.max_model_calls { - return Err(raise( - ctx, - TinyAgentsError::LimitExceeded(format!( - "model call limit ({}) exceeded", - ctx.policy.max_model_calls - )), - )); - } - counters.model += 1; - Ok(()) -} - -/// Increments and bounds the tool-call counter. -fn bump_tool(ctx: &HostContext) -> Result<(), Box> { - let mut counters = ctx.counters.lock().expect("counters poisoned"); - if counters.tool >= ctx.policy.max_tool_calls { - return Err(raise( - ctx, - TinyAgentsError::LimitExceeded(format!( - "tool call limit ({}) exceeded", - ctx.policy.max_tool_calls - )), - )); - } - counters.tool += 1; - Ok(()) -} - -/// Increments and bounds the graph-run counter. -fn bump_graph(ctx: &HostContext) -> Result<(), Box> { - let mut counters = ctx.counters.lock().expect("counters poisoned"); - if counters.graph >= ctx.policy.max_graph_calls { - return Err(raise( - ctx, - TinyAgentsError::LimitExceeded(format!( - "graph call limit ({}) exceeded", - ctx.policy.max_graph_calls - )), - )); - } - counters.graph += 1; - Ok(()) -} - -/// Increments and bounds the agent-call counter. -fn bump_agent(ctx: &HostContext) -> Result<(), Box> { - let mut counters = ctx.counters.lock().expect("counters poisoned"); - if counters.agent >= ctx.policy.max_model_calls { - return Err(raise( - ctx, - TinyAgentsError::LimitExceeded(format!( - "agent call limit ({}) exceeded", - ctx.policy.max_model_calls - )), - )); - } - counters.agent += 1; - Ok(()) -} - -/// Enforces the recursion-depth bound for a sub-run (agent or graph). -/// -/// Reuses the harness recursion bookkeeping (Cluster G): a sub-run executes one -/// level below the session's run depth, and a child depth past -/// [`ReplPolicy::max_depth`] fails closed with -/// [`TinyAgentsError::SubAgentDepth`]. -fn check_depth(ctx: &HostContext) -> Result<(), Box> { - let child_depth = ctx.run_depth + 1; - if child_depth > ctx.policy.max_depth { - return Err(raise( - ctx, - TinyAgentsError::SubAgentDepth(ctx.policy.max_depth), - )); - } - Ok(()) -} - -// ── Request builders ──────────────────────────────────────────────────────── - -/// Builds a [`ModelRequest`] from a `model_query` argument map. -fn build_model_request(model: &str, params: &Map) -> ModelRequest { - let mut messages = Vec::new(); - if let Some(system) = map_str(params, "system") { - messages.push(Message::system(system)); - } - if let Some(prompt) = map_str(params, "prompt") { - messages.push(Message::user(prompt)); - } - ModelRequest { - messages, - model: Some(model.to_string()), - ..Default::default() - } -} - -/// Wraps a model response text as the script-visible value (a string by -/// default, or a structured map when `structured: true`). -fn model_value(text: String, finish_reason: Option, structured: bool) -> Dynamic { - if structured { - let mut map = Map::new(); - map.insert("content".into(), Dynamic::from(text)); - if let Some(reason) = finish_reason { - map.insert("finish_reason".into(), Dynamic::from(reason)); - } - Dynamic::from_map(map) - } else { - Dynamic::from(text) - } -} - -// ── Single capability implementations ─────────────────────────────────────── - -fn model_query_impl( - ctx: &HostContext, - params: &Map, -) -> Result> { - let model_name = - map_str(params, "model").ok_or_else(|| invalid(ctx, "model_query: missing `model`"))?; - bump_model(ctx)?; - let model = ctx - .registry - .model(&model_name) - .ok_or_else(|| raise(ctx, TinyAgentsError::ModelNotFound(model_name.clone())))?; - let request = build_model_request(&model_name, params); - let start = Instant::now(); - let response = - bridge_block_on(model.invoke(&ctx.state, request)).map_err(|err| raise(ctx, err))?; - let elapsed = start.elapsed(); - let finish_reason = response.finish_reason.clone(); - let text = Message::Assistant(response.message).text(); - record( - ctx, - ReplCallKind::Model, - &model_name, - json!({ "chars": text.len() }), - elapsed, - ); - Ok(model_value( - text, - finish_reason, - map_bool(params, "structured").unwrap_or(false), - )) -} - -fn tool_call_impl( - ctx: &HostContext, - params: &Map, -) -> Result> { - let tool_name = - map_str(params, "tool").ok_or_else(|| invalid(ctx, "tool_call: missing `tool`"))?; - bump_tool(ctx)?; - let tool = ctx - .registry - .tool(&tool_name) - .ok_or_else(|| raise(ctx, TinyAgentsError::ToolNotFound(tool_name.clone())))?; - let arguments = map_json(params, "arguments").unwrap_or(Value::Null); - let call = ToolCall { - id: new_call_id().as_str().to_string(), - name: tool_name.clone(), - arguments: arguments.clone(), - }; - let start = Instant::now(); - let result = bridge_block_on(tool.call(&ctx.state, call)).map_err(|err| raise(ctx, err))?; - let elapsed = start.elapsed(); - record( - ctx, - ReplCallKind::Tool, - &tool_name, - json!({ "arguments": arguments }), - elapsed, - ); - if let Some(error) = result.error { - return Err(raise(ctx, TinyAgentsError::Tool(error))); - } - let structured = map_bool(params, "structured").unwrap_or(false); - if structured && result.raw.is_some() { - let mut map = Map::new(); - map.insert("content".into(), Dynamic::from(result.content)); - map.insert( - "raw".into(), - repl_value_to_dynamic(&json_to_repl_value(&result.raw.unwrap_or(Value::Null))), - ); - Ok(Dynamic::from_map(map)) - } else { - Ok(Dynamic::from(result.content)) - } -} - -fn agent_query_impl( - ctx: &HostContext, - params: &Map, -) -> Result> { - use crate::graph::subagent_node::SubAgentInput; - let agent_name = - map_str(params, "agent").ok_or_else(|| invalid(ctx, "agent_query: missing `agent`"))?; - bump_agent(ctx)?; - check_depth(ctx)?; - let agent = ctx.registry.agent(&agent_name).ok_or_else(|| { - raise( - ctx, - TinyAgentsError::Capability(format!("agent `{agent_name}` is not registered")), - ) - })?; - let prompt = map_str(params, "prompt") - .or_else(|| map_str(params, "input")) - .unwrap_or_default(); - let mut input = SubAgentInput::prompt(prompt); - if let Some(data) = map_json(params, "input") { - input = input.with_data(data); - } - let start = Instant::now(); - let output = - bridge_block_on(agent.run(input, ctx.events.clone())).map_err(|err| raise(ctx, err))?; - record( - ctx, - ReplCallKind::Agent, - &agent_name, - json!({ "model_calls": output.model_calls, "tool_calls": output.tool_calls }), - start.elapsed(), - ); - Ok(Dynamic::from(output.text)) -} - -/// Resolves a registered graph blueprint and records the run, returning a -/// reference to the resolved topology. -/// -/// Resolving a registered graph routes through the capability registry; the -/// REPL hands back the resolved blueprint reference (graph id, start node, node -/// count) rather than installing or stepping topology here. Materializing a -/// `CompiledGraph` and driving its super-steps is owned by the graph runtime -/// and wired in a later slice; this keeps the REPL an orchestration surface, -/// not a topology-mutation surface. -fn graph_run_impl( - ctx: &HostContext, - params: &Map, -) -> Result> { - let graph_name = - map_str(params, "graph").ok_or_else(|| invalid(ctx, "graph_run: missing `graph`"))?; - bump_graph(ctx)?; - check_depth(ctx)?; - let blueprint = ctx - .registry - .graph_blueprint(&graph_name) - .ok_or_else(|| { - raise( - ctx, - TinyAgentsError::Capability(format!("graph `{graph_name}` is not registered")), - ) - })? - .clone(); - record( - ctx, - ReplCallKind::Graph, - &graph_name, - json!({ "nodes": blueprint.nodes.len() }), - Duration::default(), - ); - Ok(blueprint_reference(&blueprint)) -} - -/// Builds the script-visible reference map for a resolved graph blueprint. -fn blueprint_reference(blueprint: &Blueprint) -> Dynamic { - let mut map = Map::new(); - map.insert("graph".into(), Dynamic::from(blueprint.graph_id.clone())); - map.insert("start".into(), Dynamic::from(blueprint.start.clone())); - map.insert("nodes".into(), Dynamic::from(blueprint.nodes.len() as i64)); - map.insert("resolved".into(), Dynamic::from(true)); - Dynamic::from_map(map) -} - -// ── Graph-authoring implementations ───────────────────────────────────────── - -fn graph_define_impl( - ctx: &HostContext, - params: &Map, -) -> Result> { - let name = - map_str(params, "name").ok_or_else(|| invalid(ctx, "graph_define: missing `name`"))?; - let source = - map_str(params, "source").ok_or_else(|| invalid(ctx, "graph_define: missing `source`"))?; - - { - let mut counters = ctx.counters.lock().expect("counters poisoned"); - if counters.graph_def >= ctx.policy.max_graph_definitions { - return Err(raise( - ctx, - TinyAgentsError::LimitExceeded(format!( - "graph definition limit ({}) exceeded", - ctx.policy.max_graph_definitions - )), - )); - } - counters.graph_def += 1; - } - if source.len() > ctx.policy.max_script_bytes { - return Err(raise( - ctx, - TinyAgentsError::LimitExceeded(format!( - "graph source is {} bytes, exceeding max_script_bytes ({})", - source.len(), - ctx.policy.max_script_bytes - )), - )); - } - - let label = ctx - .language - .as_ref() - .map(|l| l.provenance_label.clone()) - .unwrap_or_else(|| ctx.session_label.clone()); - let origin = Origin::generated_by(label); - let program = parse_str(&source).map_err(|err| raise(ctx, err))?; - let blueprints = - compile_with_provenance(&program, origin.clone()).map_err(|err| raise(ctx, err))?; - let blueprint = blueprints - .into_iter() - .find(|b| b.graph_id == name) - .ok_or_else(|| { - invalid( - ctx, - format!("graph_define: source has no graph named `{name}`"), - ) - })?; - - let handle = GraphBlueprintHandle { - name: blueprint.graph_id.clone(), - source, - blueprint: blueprint.clone(), - origin, - compiled: false, - requires_review: ctx.policy.generated_graphs_require_review, - }; - ctx.drafts - .lock() - .expect("drafts poisoned") - .insert(handle.name.clone(), handle.clone()); - record( - ctx, - ReplCallKind::Graph, - "graph_define", - json!({ "name": handle.name }), - Duration::default(), - ); - Ok(draft_descriptor(&handle)) -} - -/// Builds the script-visible descriptor map for a graph draft (carrying its -/// name, node count, and compile/review status). The opaque -/// [`GraphBlueprintHandle`] itself lives host-side in `ctx.drafts`. -fn draft_descriptor(handle: &GraphBlueprintHandle) -> Dynamic { - let mut map = Map::new(); - map.insert("name".into(), Dynamic::from(handle.name.clone())); - map.insert( - "nodes".into(), - Dynamic::from(handle.blueprint.nodes.len() as i64), - ); - map.insert("compiled".into(), Dynamic::from(handle.compiled)); - map.insert( - "requires_review".into(), - Dynamic::from(handle.requires_review), - ); - Dynamic::from_map(map) -} - -/// Looks up a graph draft by the `name` field of a descriptor map. -fn lookup_draft( - ctx: &HostContext, - descriptor: &Map, - func: &str, -) -> Result> { - let name = map_str(descriptor, "name") - .ok_or_else(|| invalid(ctx, format!("{func}: descriptor is missing `name`")))?; - ctx.drafts - .lock() - .expect("drafts poisoned") - .get(&name) - .cloned() - .ok_or_else(|| invalid(ctx, format!("{func}: no graph draft named `{name}`"))) -} - -fn graph_validate_impl( - ctx: &HostContext, - descriptor: &Map, -) -> Result> { - let handle = lookup_draft(ctx, descriptor, "graph_validate")?; - let program = parse_str(&handle.source).map_err(|err| raise(ctx, err))?; - let diagnostics = Resolver::from_registry(&*ctx.registry).resolve_program(&program); - let array: Array = diagnostics - .iter() - .map(|d| Dynamic::from(d.message.clone())) - .collect(); - Ok(Dynamic::from_array(array)) -} - -fn graph_compile_impl( - ctx: &HostContext, - descriptor: &Map, -) -> Result> { - let mut handle = lookup_draft(ctx, descriptor, "graph_compile")?; - // Bind the blueprint through the same resolver gate file-backed `.rag` - // source passes — generated topology is never trusted blindly. - Resolver::from_registry(&*ctx.registry) - .resolve_blueprint(&handle.blueprint) - .map_err(|err| raise(ctx, err))?; - handle.compiled = true; - handle.requires_review = ctx.policy.generated_graphs_require_review; - ctx.drafts - .lock() - .expect("drafts poisoned") - .insert(handle.name.clone(), handle.clone()); - record( - ctx, - ReplCallKind::Graph, - "graph_compile", - json!({ "name": handle.name, "requires_review": handle.requires_review }), - Duration::default(), - ); - Ok(draft_descriptor(&handle)) -} - -fn graph_diff_handles( - ctx: &HostContext, - old: &Blueprint, - new: &Blueprint, -) -> Result> { - let diff = blueprint_diff(old, new); - let value = serde_json::to_value(&diff) - .map_err(|err| raise(ctx, TinyAgentsError::Validation(err.to_string())))?; - Ok(repl_value_to_dynamic(&json_to_repl_value(&value))) -} - -fn graph_register_impl( - ctx: &HostContext, - params: &Map, -) -> Result> { - let graph = params - .get("graph") - .and_then(|d| d.read_lock::().map(|m| m.clone())) - .ok_or_else(|| { - invalid( - ctx, - "graph_register: `graph` must be a compiled graph descriptor", - ) - })?; - let handle = lookup_draft(ctx, &graph, "graph_register")?; - if !handle.compiled { - return Err(raise( - ctx, - TinyAgentsError::Validation( - "graph_register: graph must be compiled via graph_compile first".to_string(), - ), - )); - } - let review_id = map_str(params, "review_id").filter(|s| !s.is_empty()); - if handle.requires_review && review_id.is_none() { - return Err(raise( - ctx, - TinyAgentsError::Validation(format!( - "graph_register: generated graph `{}` requires review (no review_id)", - handle.name - )), - )); - } - // Enforce the review gate and emit a registry intent. The compiled topology - // is handed to the host for installation through the registry resolver — - // the REPL never installs generated topology directly. - record( - ctx, - ReplCallKind::Graph, - "graph_register", - json!({ "name": handle.name, "review_id": review_id }), - Duration::default(), - ); - Ok(Dynamic::from(handle.name)) -} - -// ── Batched implementations ───────────────────────────────────────────────── - -/// Extracts the object-map items of a batched argument array. -fn batch_items( - ctx: &HostContext, - items: &Array, - func: &str, -) -> Result, Box> { - items - .iter() - .map(|item| { - item.read_lock::() - .map(|m| m.clone()) - .ok_or_else(|| invalid(ctx, format!("{func}: each item must be an object map"))) - }) - .collect() -} - -fn model_query_batched_impl( - ctx: &HostContext, - items: &Array, -) -> Result> { - use futures::stream::{self, StreamExt}; - - let items = batch_items(ctx, items, "model_query_batched")?; - let mut prepared = Vec::with_capacity(items.len()); - for params in &items { - let model_name = map_str(params, "model") - .ok_or_else(|| invalid(ctx, "model_query_batched: missing `model`"))?; - bump_model(ctx)?; - let model = ctx - .registry - .model(&model_name) - .ok_or_else(|| raise(ctx, TinyAgentsError::ModelNotFound(model_name.clone())))?; - let request = build_model_request(&model_name, params); - let structured = map_bool(params, "structured").unwrap_or(false); - prepared.push((model_name, model, request, structured)); - } - - let concurrency = ctx.policy.max_concurrency.max(1); - let results: Vec> = bridge_block_on(async { - stream::iter(prepared.iter().map(|(name, model, request, structured)| { - let name = name.clone(); - let structured = *structured; - async move { - let start = Instant::now(); - let response = model.invoke(&ctx.state, request.clone()).await?; - let finish_reason = response.finish_reason.clone(); - let text = Message::Assistant(response.message).text(); - Ok((name, text, finish_reason, structured, start.elapsed())) - } - })) - .buffered(concurrency) - .collect() - .await - }); - - let mut out = Array::with_capacity(results.len()); - for result in results { - let (name, text, finish_reason, structured, elapsed) = - result.map_err(|err| raise(ctx, err))?; - record( - ctx, - ReplCallKind::Model, - &name, - json!({ "chars": text.len() }), - elapsed, - ); - out.push(model_value(text, finish_reason, structured)); - } - Ok(Dynamic::from_array(out)) -} - -fn tool_call_batched_impl( - ctx: &HostContext, - items: &Array, -) -> Result> { - use futures::stream::{self, StreamExt}; - - let items = batch_items(ctx, items, "tool_call_batched")?; - let mut prepared = Vec::with_capacity(items.len()); - for params in &items { - let tool_name = map_str(params, "tool") - .ok_or_else(|| invalid(ctx, "tool_call_batched: missing `tool`"))?; - bump_tool(ctx)?; - let tool = ctx - .registry - .tool(&tool_name) - .ok_or_else(|| raise(ctx, TinyAgentsError::ToolNotFound(tool_name.clone())))?; - let arguments = map_json(params, "arguments").unwrap_or(Value::Null); - prepared.push((tool_name, tool, arguments)); - } - - let concurrency = ctx.policy.max_concurrency.max(1); - let results: Vec< - Result<(String, crate::harness::tool::ToolResult, Duration), TinyAgentsError>, - > = bridge_block_on(async { - stream::iter(prepared.iter().map(|(name, tool, arguments)| { - let name = name.clone(); - let call = ToolCall { - id: new_call_id().as_str().to_string(), - name: name.clone(), - arguments: arguments.clone(), - }; - async move { - let start = Instant::now(); - let result = tool.call(&ctx.state, call).await?; - Ok((name, result, start.elapsed())) - } - })) - .buffered(concurrency) - .collect() - .await - }); - - let mut out = Array::with_capacity(results.len()); - for result in results { - let (name, tool_result, elapsed) = result.map_err(|err| raise(ctx, err))?; - record( - ctx, - ReplCallKind::Tool, - &name, - json!({ "chars": tool_result.content.len() }), - elapsed, - ); - if let Some(error) = tool_result.error { - return Err(raise(ctx, TinyAgentsError::Tool(error))); - } - out.push(Dynamic::from(tool_result.content)); - } - Ok(Dynamic::from_array(out)) -} - -fn agent_query_batched_impl( - ctx: &HostContext, - items: &Array, -) -> Result> { - use crate::graph::subagent_node::SubAgentInput; - use futures::stream::{self, StreamExt}; - - let items = batch_items(ctx, items, "agent_query_batched")?; - let mut prepared = Vec::with_capacity(items.len()); - for params in &items { - let agent_name = map_str(params, "agent") - .ok_or_else(|| invalid(ctx, "agent_query_batched: missing `agent`"))?; - bump_agent(ctx)?; - check_depth(ctx)?; - let agent = ctx.registry.agent(&agent_name).ok_or_else(|| { - raise( - ctx, - TinyAgentsError::Capability(format!("agent `{agent_name}` is not registered")), - ) - })?; - let prompt = map_str(params, "prompt") - .or_else(|| map_str(params, "input")) - .unwrap_or_default(); - let mut input = SubAgentInput::prompt(prompt); - if let Some(data) = map_json(params, "input") { - input = input.with_data(data); - } - prepared.push((agent_name, agent, input)); - } - - let concurrency = ctx.policy.max_concurrency.max(1); - let results: Vec> = bridge_block_on(async { - stream::iter(prepared.iter().map(|(name, agent, input)| { - let name = name.clone(); - async move { - let start = Instant::now(); - let output = agent.run(input.clone(), ctx.events.clone()).await?; - Ok((name, output.text, start.elapsed())) - } - })) - .buffered(concurrency) - .collect() - .await - }); - - let mut out = Array::with_capacity(results.len()); - for result in results { - let (name, text, elapsed) = result.map_err(|err| raise(ctx, err))?; - record(ctx, ReplCallKind::Agent, &name, json!({}), elapsed); - out.push(Dynamic::from(text)); - } - Ok(Dynamic::from_array(out)) -} - -fn graph_run_batched_impl( - ctx: &HostContext, - items: &Array, -) -> Result> { - let items = batch_items(ctx, items, "graph_run_batched")?; - let mut out = Array::with_capacity(items.len()); - for params in &items { - out.push(graph_run_impl(ctx, params)?); - } - Ok(Dynamic::from_array(out)) -} - -// ── Engine construction ───────────────────────────────────────────────────── - -/// Builds a sandboxed Rhai engine for a session, registering every host-backed -/// built-in against the session's live registries and policy. -/// -/// The engine is configured with the policy operation limit (fail-closed on -/// runaway scripts) and is granted no filesystem, network, or process access — -/// the only host surface is the capability functions registered here. -pub(super) fn build_engine(ctx: Arc>) -> Engine { - let mut engine = Engine::new(); - engine.set_max_operations(ctx.policy.max_operations); - - // ── stdout capture ── - let stdout_ctx = ctx.clone(); - engine.on_print(move |text| stdout_ctx.buffers.push_stdout_line(text)); - let debug_ctx = ctx.clone(); - engine.on_debug(move |text, _source, _pos| debug_ctx.buffers.push_stdout_line(text)); - - // ── emit(name) / emit(name, #{ ... }) ── - let emit_ctx = ctx.clone(); - engine.register_fn("emit", move |name: &str| { - record( - &emit_ctx, - ReplCallKind::Emit, - name, - Value::Null, - Duration::default(), - ); - }); - let emit_payload_ctx = ctx.clone(); - engine.register_fn("emit", move |name: &str, data: Map| { - let detail = dynamic_to_repl_value(&Dynamic::from_map(data)).to_json(); - record( - &emit_payload_ctx, - ReplCallKind::Emit, - name, - detail, - Duration::default(), - ); - }); - - // ── answer(content) ── - let answer_ctx = ctx.clone(); - engine.register_fn("answer", move |content: &str| { - answer_ctx.buffers.set_answer(content.to_string()); - }); - - // ── show_vars() ── - let show_ctx = ctx.clone(); - engine.register_fn("show_vars", move || { - show_ctx.buffers.push_stdout_line("# vars"); - for (name, value) in show_ctx.buffers.vars_snapshot() { - show_ctx - .buffers - .push_stdout_line(&format!("{name} = {value}")); - } - }); - - // ── model capabilities ── - let model_ctx = ctx.clone(); - engine.register_fn("model_query", move |params: Map| { - model_query_impl(&model_ctx, ¶ms) - }); - let model_batch_ctx = ctx.clone(); - engine.register_fn("model_query_batched", move |items: Array| { - model_query_batched_impl(&model_batch_ctx, &items) - }); - - // ── tool capabilities ── - let tool_ctx = ctx.clone(); - engine.register_fn("tool_call", move |params: Map| { - tool_call_impl(&tool_ctx, ¶ms) - }); - let tool_batch_ctx = ctx.clone(); - engine.register_fn("tool_call_batched", move |items: Array| { - tool_call_batched_impl(&tool_batch_ctx, &items) - }); - - // ── agent capabilities ── - let agent_ctx = ctx.clone(); - engine.register_fn("agent_query", move |params: Map| { - agent_query_impl(&agent_ctx, ¶ms) - }); - let agent_batch_ctx = ctx.clone(); - engine.register_fn("agent_query_batched", move |items: Array| { - agent_query_batched_impl(&agent_batch_ctx, &items) - }); - - // ── graph run capabilities ── - let graph_ctx = ctx.clone(); - engine.register_fn("graph_run", move |params: Map| { - graph_run_impl(&graph_ctx, ¶ms) - }); - let graph_batch_ctx = ctx.clone(); - engine.register_fn("graph_run_batched", move |items: Array| { - graph_run_batched_impl(&graph_batch_ctx, &items) - }); - - // ── graph authoring (lowering through the `.rag` compiler) ── - let define_ctx = ctx.clone(); - engine.register_fn("graph_define", move |params: Map| { - graph_define_impl(&define_ctx, ¶ms) - }); - let validate_ctx = ctx.clone(); - engine.register_fn("graph_validate", move |descriptor: Map| { - graph_validate_impl(&validate_ctx, &descriptor) - }); - let compile_ctx = ctx.clone(); - engine.register_fn("graph_compile", move |descriptor: Map| { - graph_compile_impl(&compile_ctx, &descriptor) - }); - let diff_name_ctx = ctx.clone(); - engine.register_fn( - "graph_diff", - move |name: &str, draft: Map| -> Result> { - let old = diff_name_ctx - .registry - .graph_blueprint(name) - .ok_or_else(|| { - invalid( - &diff_name_ctx, - format!("graph_diff: graph `{name}` is not registered"), - ) - })? - .clone(); - let new = lookup_draft(&diff_name_ctx, &draft, "graph_diff")?; - graph_diff_handles(&diff_name_ctx, &old, &new.blueprint) - }, - ); - let diff_draft_ctx = ctx.clone(); - engine.register_fn( - "graph_diff", - move |old: Map, new: Map| -> Result> { - let old = lookup_draft(&diff_draft_ctx, &old, "graph_diff")?; - let new = lookup_draft(&diff_draft_ctx, &new, "graph_diff")?; - graph_diff_handles(&diff_draft_ctx, &old.blueprint, &new.blueprint) - }, - ); - let register_ctx = ctx.clone(); - engine.register_fn("graph_register", move |params: Map| { - graph_register_impl(®ister_ctx, ¶ms) - }); - - engine -} diff --git a/src/repl/session/builtins/authoring.rs b/src/repl/session/builtins/authoring.rs new file mode 100644 index 0000000..f9ef98f --- /dev/null +++ b/src/repl/session/builtins/authoring.rs @@ -0,0 +1,231 @@ +//! Graph-authoring implementations (`graph_define`, `graph_validate`, +//! `graph_compile`, `graph_diff`, `graph_register`) lowering through the +//! expressive-language compiler and capability resolver. +//! +//! Split out of `session/builtins/mod.rs`; see that module's doc comment +//! for the full built-in surface and the blocking-bridge design. + +use super::*; + +// ── Graph-authoring implementations ───────────────────────────────────────── + +pub(super) fn graph_define_impl( + ctx: &HostContext, + params: &Map, +) -> Result> { + let name = + map_str(params, "name").ok_or_else(|| invalid(ctx, "graph_define: missing `name`"))?; + let source = + map_str(params, "source").ok_or_else(|| invalid(ctx, "graph_define: missing `source`"))?; + + // Check the limit up front (without consuming a slot) so a session that + // has already hit the cap fails fast instead of paying for a parse and + // compile it can't keep the result of anyway. + if ctx.counters.lock().expect("counters poisoned").graph_def >= ctx.policy.max_graph_definitions + { + return Err(raise( + ctx, + TinyAgentsError::LimitExceeded(format!( + "graph definition limit ({}) exceeded", + ctx.policy.max_graph_definitions + )), + )); + } + if source.len() > ctx.policy.max_script_bytes { + return Err(raise( + ctx, + TinyAgentsError::LimitExceeded(format!( + "graph source is {} bytes, exceeding max_script_bytes ({})", + source.len(), + ctx.policy.max_script_bytes + )), + )); + } + + let label = ctx + .language + .as_ref() + .map(|l| l.provenance_label.clone()) + .unwrap_or_else(|| ctx.session_label.clone()); + let origin = Origin::generated_by(label); + let program = parse_str(&source).map_err(|err| raise(ctx, err))?; + let blueprints = + compile_with_provenance(&program, origin.clone()).map_err(|err| raise(ctx, err))?; + let blueprint = blueprints + .into_iter() + .find(|b| b.graph_id == name) + .ok_or_else(|| { + invalid( + ctx, + format!("graph_define: source has no graph named `{name}`"), + ) + })?; + + // The draft is about to be recorded successfully; consume a slot now + // (re-checking the limit under the same lock to guard against a + // concurrent `graph_define` racing between the check above and here). + { + let mut counters = ctx.counters.lock().expect("counters poisoned"); + if counters.graph_def >= ctx.policy.max_graph_definitions { + return Err(raise( + ctx, + TinyAgentsError::LimitExceeded(format!( + "graph definition limit ({}) exceeded", + ctx.policy.max_graph_definitions + )), + )); + } + counters.graph_def += 1; + } + + let handle = GraphBlueprintHandle { + name: blueprint.graph_id.clone(), + source, + blueprint: blueprint.clone(), + origin, + compiled: false, + requires_review: ctx.policy.generated_graphs_require_review, + }; + ctx.drafts + .lock() + .expect("drafts poisoned") + .insert(handle.name.clone(), handle.clone()); + record( + ctx, + ReplCallKind::Graph, + "graph_define", + json!({ "name": handle.name }), + Duration::default(), + ); + Ok(draft_descriptor(&handle)) +} + +/// Builds the script-visible descriptor map for a graph draft (carrying its +/// name, node count, and compile/review status). The opaque +/// [`GraphBlueprintHandle`] itself lives host-side in `ctx.drafts`. +fn draft_descriptor(handle: &GraphBlueprintHandle) -> Dynamic { + let mut map = Map::new(); + map.insert("name".into(), Dynamic::from(handle.name.clone())); + map.insert( + "nodes".into(), + Dynamic::from(handle.blueprint.nodes.len() as i64), + ); + map.insert("compiled".into(), Dynamic::from(handle.compiled)); + map.insert( + "requires_review".into(), + Dynamic::from(handle.requires_review), + ); + Dynamic::from_map(map) +} + +/// Looks up a graph draft by the `name` field of a descriptor map. +pub(super) fn lookup_draft( + ctx: &HostContext, + descriptor: &Map, + func: &str, +) -> Result> { + let name = map_str(descriptor, "name") + .ok_or_else(|| invalid(ctx, format!("{func}: descriptor is missing `name`")))?; + ctx.drafts + .lock() + .expect("drafts poisoned") + .get(&name) + .cloned() + .ok_or_else(|| invalid(ctx, format!("{func}: no graph draft named `{name}`"))) +} + +pub(super) fn graph_validate_impl( + ctx: &HostContext, + descriptor: &Map, +) -> Result> { + let handle = lookup_draft(ctx, descriptor, "graph_validate")?; + let program = parse_str(&handle.source).map_err(|err| raise(ctx, err))?; + let diagnostics = Resolver::from_registry(&*ctx.registry).resolve_program(&program); + let array: Array = diagnostics + .iter() + .map(|d| Dynamic::from(d.message.clone())) + .collect(); + Ok(Dynamic::from_array(array)) +} + +pub(super) fn graph_compile_impl( + ctx: &HostContext, + descriptor: &Map, +) -> Result> { + let mut handle = lookup_draft(ctx, descriptor, "graph_compile")?; + // Bind the blueprint through the same resolver gate file-backed `.rag` + // source passes — generated topology is never trusted blindly. + Resolver::from_registry(&*ctx.registry) + .resolve_blueprint(&handle.blueprint) + .map_err(|err| raise(ctx, err))?; + handle.compiled = true; + handle.requires_review = ctx.policy.generated_graphs_require_review; + ctx.drafts + .lock() + .expect("drafts poisoned") + .insert(handle.name.clone(), handle.clone()); + record( + ctx, + ReplCallKind::Graph, + "graph_compile", + json!({ "name": handle.name, "requires_review": handle.requires_review }), + Duration::default(), + ); + Ok(draft_descriptor(&handle)) +} + +pub(super) fn graph_diff_handles( + ctx: &HostContext, + old: &Blueprint, + new: &Blueprint, +) -> Result> { + let diff = blueprint_diff(old, new); + let value = serde_json::to_value(&diff) + .map_err(|err| raise(ctx, TinyAgentsError::Validation(err.to_string())))?; + Ok(repl_value_to_dynamic(&json_to_repl_value(&value))) +} + +pub(super) fn graph_register_impl( + ctx: &HostContext, + params: &Map, +) -> Result> { + let graph = params + .get("graph") + .and_then(|d| d.read_lock::().map(|m| m.clone())) + .ok_or_else(|| { + invalid( + ctx, + "graph_register: `graph` must be a compiled graph descriptor", + ) + })?; + let handle = lookup_draft(ctx, &graph, "graph_register")?; + if !handle.compiled { + return Err(raise( + ctx, + TinyAgentsError::Validation( + "graph_register: graph must be compiled via graph_compile first".to_string(), + ), + )); + } + let review_id = map_str(params, "review_id").filter(|s| !s.is_empty()); + if handle.requires_review && review_id.is_none() { + return Err(raise( + ctx, + TinyAgentsError::Validation(format!( + "graph_register: generated graph `{}` requires review (no review_id)", + handle.name + )), + )); + } + // Enforce the review gate and emit a registry intent. The compiled topology + // is handed to the host for installation through the registry resolver — + // the REPL never installs generated topology directly. + record( + ctx, + ReplCallKind::Graph, + "graph_register", + json!({ "name": handle.name, "review_id": review_id }), + Duration::default(), + ); + Ok(Dynamic::from(handle.name)) +} diff --git a/src/repl/session/builtins/batched.rs b/src/repl/session/builtins/batched.rs new file mode 100644 index 0000000..f47c176 --- /dev/null +++ b/src/repl/session/builtins/batched.rs @@ -0,0 +1,228 @@ +//! Batched capability implementations (`model_query_batched`, +//! `tool_call_batched`, `agent_query_batched`, `graph_run_batched`). +//! +//! Split out of `session/builtins/mod.rs`; see that module's doc comment +//! for the full built-in surface and the blocking-bridge design. + +use super::*; + +// ── Batched implementations ───────────────────────────────────────────────── + +/// Extracts the object-map items of a batched argument array. +fn batch_items( + ctx: &HostContext, + items: &Array, + func: &str, +) -> Result, Box> { + items + .iter() + .map(|item| { + item.read_lock::() + .map(|m| m.clone()) + .ok_or_else(|| invalid(ctx, format!("{func}: each item must be an object map"))) + }) + .collect() +} + +pub(super) fn model_query_batched_impl( + ctx: &HostContext, + items: &Array, +) -> Result> { + use futures::stream::{self, StreamExt}; + + let items = batch_items(ctx, items, "model_query_batched")?; + let mut prepared = Vec::with_capacity(items.len()); + for params in &items { + let model_name = map_str(params, "model") + .ok_or_else(|| invalid(ctx, "model_query_batched: missing `model`"))?; + bump_model(ctx)?; + let model = ctx + .registry + .model(&model_name) + .ok_or_else(|| raise(ctx, TinyAgentsError::ModelNotFound(model_name.clone())))?; + let request = build_model_request(&model_name, params); + let structured = map_bool(params, "structured").unwrap_or(false); + prepared.push((model_name, model, request, structured)); + } + + let concurrency = ctx.policy.max_concurrency.max(1); + let results: Vec> = + bridge_block_on_raw(ctx.buffers.deadline(), async { + stream::iter(prepared.iter().map(|(name, model, request, structured)| { + let name = name.clone(); + let structured = *structured; + async move { + let start = Instant::now(); + let response = model.invoke(&ctx.state, request.clone()).await?; + let finish_reason = response.finish_reason.clone(); + let text = Message::Assistant(response.message).text(); + Ok((name, text, finish_reason, structured, start.elapsed())) + } + })) + .buffered(concurrency) + .collect() + .await + }) + .map_err(|err| raise(ctx, err))?; + + let mut out = Array::with_capacity(results.len()); + for result in results { + let (name, text, finish_reason, structured, elapsed) = + result.map_err(|err| raise(ctx, err))?; + record( + ctx, + ReplCallKind::Model, + &name, + json!({ "chars": text.len() }), + elapsed, + ); + out.push(model_value(text, finish_reason, structured)); + } + Ok(Dynamic::from_array(out)) +} + +pub(super) fn tool_call_batched_impl( + ctx: &HostContext, + items: &Array, +) -> Result> { + use futures::stream::{self, StreamExt}; + + let items = batch_items(ctx, items, "tool_call_batched")?; + let mut prepared = Vec::with_capacity(items.len()); + for params in &items { + let tool_name = map_str(params, "tool") + .ok_or_else(|| invalid(ctx, "tool_call_batched: missing `tool`"))?; + bump_tool(ctx)?; + let tool = ctx + .registry + .tool(&tool_name) + .ok_or_else(|| raise(ctx, TinyAgentsError::ToolNotFound(tool_name.clone())))?; + let arguments = map_json(params, "arguments").unwrap_or(Value::Null); + prepared.push((tool_name, tool, arguments)); + } + + let concurrency = ctx.policy.max_concurrency.max(1); + let results: Vec< + Result<(String, crate::harness::tool::ToolResult, Duration), TinyAgentsError>, + > = bridge_block_on_raw(ctx.buffers.deadline(), async { + stream::iter(prepared.iter().map(|(name, tool, arguments)| { + let name = name.clone(); + let call = ToolCall { + id: new_call_id().as_str().to_string(), + name: name.clone(), + arguments: arguments.clone(), + }; + async move { + let start = Instant::now(); + let result = tool.call(&ctx.state, call).await?; + Ok((name, result, start.elapsed())) + } + })) + .buffered(concurrency) + .collect() + .await + }) + .map_err(|err| raise(ctx, err))?; + + // Each item's own tool-reported error is surfaced per item, matching the + // single-call path's behavior for that one call, rather than aborting the + // whole batch and discarding every other item's already-computed result — + // a batch of N independent tool calls should not lose N-1 successes + // because item N/2 failed. A `bridge_block_on_raw`/transport failure + // above (a harness-level failure, not a tool-reported one) still aborts + // the whole batch, since no results exist to preserve in that case. + let mut out = Array::with_capacity(results.len()); + for result in results { + let (name, tool_result, elapsed) = result.map_err(|err| raise(ctx, err))?; + record( + ctx, + ReplCallKind::Tool, + &name, + json!({ "chars": tool_result.content.len() }), + elapsed, + ); + match tool_result.error { + Some(error) => { + let mut map = Map::new(); + map.insert("ok".into(), Dynamic::from(false)); + map.insert("error".into(), Dynamic::from(error)); + out.push(Dynamic::from_map(map)); + } + None => { + let mut map = Map::new(); + map.insert("ok".into(), Dynamic::from(true)); + map.insert("content".into(), Dynamic::from(tool_result.content)); + out.push(Dynamic::from_map(map)); + } + } + } + Ok(Dynamic::from_array(out)) +} + +pub(super) fn agent_query_batched_impl( + ctx: &HostContext, + items: &Array, +) -> Result> { + use crate::graph::subagent_node::SubAgentInput; + use futures::stream::{self, StreamExt}; + + let items = batch_items(ctx, items, "agent_query_batched")?; + let mut prepared = Vec::with_capacity(items.len()); + for params in &items { + let agent_name = map_str(params, "agent") + .ok_or_else(|| invalid(ctx, "agent_query_batched: missing `agent`"))?; + bump_agent(ctx)?; + check_depth(ctx)?; + let agent = ctx.registry.agent(&agent_name).ok_or_else(|| { + raise( + ctx, + TinyAgentsError::Capability(format!("agent `{agent_name}` is not registered")), + ) + })?; + let prompt = map_str(params, "prompt") + .or_else(|| map_str(params, "input")) + .unwrap_or_default(); + let mut input = SubAgentInput::prompt(prompt); + if let Some(data) = map_json(params, "input") { + input = input.with_data(data); + } + prepared.push((agent_name, agent, input)); + } + + let concurrency = ctx.policy.max_concurrency.max(1); + let results: Vec> = + bridge_block_on_raw(ctx.buffers.deadline(), async { + stream::iter(prepared.iter().map(|(name, agent, input)| { + let name = name.clone(); + async move { + let start = Instant::now(); + let output = agent.run(input.clone(), ctx.events.clone()).await?; + Ok((name, output.text, start.elapsed())) + } + })) + .buffered(concurrency) + .collect() + .await + }) + .map_err(|err| raise(ctx, err))?; + + let mut out = Array::with_capacity(results.len()); + for result in results { + let (name, text, elapsed) = result.map_err(|err| raise(ctx, err))?; + record(ctx, ReplCallKind::Agent, &name, json!({}), elapsed); + out.push(Dynamic::from(text)); + } + Ok(Dynamic::from_array(out)) +} + +pub(super) fn graph_run_batched_impl( + ctx: &HostContext, + items: &Array, +) -> Result> { + let items = batch_items(ctx, items, "graph_run_batched")?; + let mut out = Array::with_capacity(items.len()); + for params in &items { + out.push(graph_run_impl(ctx, params)?); + } + Ok(Dynamic::from_array(out)) +} diff --git a/src/repl/session/builtins/capabilities.rs b/src/repl/session/builtins/capabilities.rs new file mode 100644 index 0000000..64a9557 --- /dev/null +++ b/src/repl/session/builtins/capabilities.rs @@ -0,0 +1,168 @@ +//! Single-shot capability implementations (`model_query`, `tool_call`, +//! `agent_query`, `graph_run`). +//! +//! Split out of `session/builtins/mod.rs`; see that module's doc comment +//! for the full built-in surface and the blocking-bridge design. + +use super::*; + +// ── Single capability implementations ─────────────────────────────────────── + +pub(super) fn model_query_impl( + ctx: &HostContext, + params: &Map, +) -> Result> { + let model_name = + map_str(params, "model").ok_or_else(|| invalid(ctx, "model_query: missing `model`"))?; + bump_model(ctx)?; + let model = ctx + .registry + .model(&model_name) + .ok_or_else(|| raise(ctx, TinyAgentsError::ModelNotFound(model_name.clone())))?; + let request = build_model_request(&model_name, params); + let start = Instant::now(); + let response = bridge_block_on(ctx.buffers.deadline(), model.invoke(&ctx.state, request)) + .map_err(|err| raise(ctx, err))?; + let elapsed = start.elapsed(); + let finish_reason = response.finish_reason.clone(); + let text = Message::Assistant(response.message).text(); + record( + ctx, + ReplCallKind::Model, + &model_name, + json!({ "chars": text.len() }), + elapsed, + ); + Ok(model_value( + text, + finish_reason, + map_bool(params, "structured").unwrap_or(false), + )) +} + +pub(super) fn tool_call_impl( + ctx: &HostContext, + params: &Map, +) -> Result> { + let tool_name = + map_str(params, "tool").ok_or_else(|| invalid(ctx, "tool_call: missing `tool`"))?; + bump_tool(ctx)?; + let tool = ctx + .registry + .tool(&tool_name) + .ok_or_else(|| raise(ctx, TinyAgentsError::ToolNotFound(tool_name.clone())))?; + let arguments = map_json(params, "arguments").unwrap_or(Value::Null); + let call = ToolCall { + id: new_call_id().as_str().to_string(), + name: tool_name.clone(), + arguments: arguments.clone(), + }; + let start = Instant::now(); + let result = bridge_block_on(ctx.buffers.deadline(), tool.call(&ctx.state, call)) + .map_err(|err| raise(ctx, err))?; + let elapsed = start.elapsed(); + record( + ctx, + ReplCallKind::Tool, + &tool_name, + json!({ "arguments": arguments }), + elapsed, + ); + if let Some(error) = result.error { + return Err(raise(ctx, TinyAgentsError::Tool(error))); + } + let structured = map_bool(params, "structured").unwrap_or(false); + if structured && result.raw.is_some() { + let mut map = Map::new(); + map.insert("content".into(), Dynamic::from(result.content)); + map.insert( + "raw".into(), + repl_value_to_dynamic(&json_to_repl_value(&result.raw.unwrap_or(Value::Null))), + ); + Ok(Dynamic::from_map(map)) + } else { + Ok(Dynamic::from(result.content)) + } +} + +pub(super) fn agent_query_impl( + ctx: &HostContext, + params: &Map, +) -> Result> { + use crate::graph::subagent_node::SubAgentInput; + let agent_name = + map_str(params, "agent").ok_or_else(|| invalid(ctx, "agent_query: missing `agent`"))?; + bump_agent(ctx)?; + check_depth(ctx)?; + let agent = ctx.registry.agent(&agent_name).ok_or_else(|| { + raise( + ctx, + TinyAgentsError::Capability(format!("agent `{agent_name}` is not registered")), + ) + })?; + let prompt = map_str(params, "prompt") + .or_else(|| map_str(params, "input")) + .unwrap_or_default(); + let mut input = SubAgentInput::prompt(prompt); + if let Some(data) = map_json(params, "input") { + input = input.with_data(data); + } + let start = Instant::now(); + let output = bridge_block_on(ctx.buffers.deadline(), agent.run(input, ctx.events.clone())) + .map_err(|err| raise(ctx, err))?; + record( + ctx, + ReplCallKind::Agent, + &agent_name, + json!({ "model_calls": output.model_calls, "tool_calls": output.tool_calls }), + start.elapsed(), + ); + Ok(Dynamic::from(output.text)) +} + +/// Resolves a registered graph blueprint and records the run, returning a +/// reference to the resolved topology. +/// +/// Resolving a registered graph routes through the capability registry; the +/// REPL hands back the resolved blueprint reference (graph id, start node, node +/// count) rather than installing or stepping topology here. Materializing a +/// `CompiledGraph` and driving its super-steps is owned by the graph runtime +/// and wired in a later slice; this keeps the REPL an orchestration surface, +/// not a topology-mutation surface. +pub(super) fn graph_run_impl( + ctx: &HostContext, + params: &Map, +) -> Result> { + let graph_name = + map_str(params, "graph").ok_or_else(|| invalid(ctx, "graph_run: missing `graph`"))?; + bump_graph(ctx)?; + check_depth(ctx)?; + let blueprint = ctx + .registry + .graph_blueprint(&graph_name) + .ok_or_else(|| { + raise( + ctx, + TinyAgentsError::Capability(format!("graph `{graph_name}` is not registered")), + ) + })? + .clone(); + record( + ctx, + ReplCallKind::Graph, + &graph_name, + json!({ "nodes": blueprint.nodes.len() }), + Duration::default(), + ); + Ok(blueprint_reference(&blueprint)) +} + +/// Builds the script-visible reference map for a resolved graph blueprint. +fn blueprint_reference(blueprint: &Blueprint) -> Dynamic { + let mut map = Map::new(); + map.insert("graph".into(), Dynamic::from(blueprint.graph_id.clone())); + map.insert("start".into(), Dynamic::from(blueprint.start.clone())); + map.insert("nodes".into(), Dynamic::from(blueprint.nodes.len() as i64)); + map.insert("resolved".into(), Dynamic::from(true)); + Dynamic::from_map(map) +} diff --git a/src/repl/session/builtins/mod.rs b/src/repl/session/builtins/mod.rs new file mode 100644 index 0000000..5f624cd --- /dev/null +++ b/src/repl/session/builtins/mod.rs @@ -0,0 +1,578 @@ +//! Capability-bound built-in functions for the Rhai-backed `.ragsh` session +//! (design milestones R3–R5). +//! +//! This module registers the reserved built-in functions on a session's +//! [`rhai::Engine`] as **host capabilities**: each one resolves a name through +//! the session's [`CapabilityRegistry`], enforces the [`ReplPolicy`] call and +//! recursion limits, records a [`ReplCallRecord`], and lowers to the real +//! harness/graph runtime. +//! +//! The surface registered here is: +//! +//! - model calls — `model_query`, `model_query_batched` +//! - agent calls — `agent_query`, `agent_query_batched` +//! - graph runs — `graph_run`, `graph_run_batched` +//! - tool calls — `tool_call`, `tool_call_batched` +//! - session built-ins — `emit`, `show_vars`, `answer`, plus `print`/`debug` +//! capture +//! - graph authoring — `graph_define`, `graph_validate`, `graph_compile`, +//! `graph_diff`, `graph_register`, which lower through the Cluster H `.rag` +//! compiler and capability resolver. Generated topology is never installed +//! directly: a draft only becomes `compiled` after the resolver binds it, and +//! `graph_register` honors [`ReplPolicy::generated_graphs_require_review`]. +//! +//! ## The async adapter (blocking bridge) +//! +//! The Rhai engine is synchronous, but model, tool, agent, and graph calls are +//! async. This slice uses the design's **blocking bridge** for v1: a host +//! function builds the async future and drives it to completion in place with +//! [`futures::executor::block_on`] (see [`bridge_block_on`]). This keeps the +//! capability boundary deterministic for scripted tests (a [`ScriptedModel`] or +//! [`FakeTool`] resolves without yielding to a reactor) and works for real +//! providers when the session is driven from a multi-threaded runtime, where +//! blocking one worker does not starve the call's own I/O. +//! +//! The design's longer-term direction is *command recording* (host functions +//! emit `ReplCommand` values the async runtime executes after the cell); the +//! public `ReplResult`/`ReplCallRecord` types are already shaped for it. The +//! bridge is intentionally the only blocking surface and is confined to this +//! module. +//! +//! [`ScriptedModel`]: crate::harness::testkit::ScriptedModel +//! [`FakeTool`]: crate::harness::testkit::FakeTool + +use std::collections::BTreeMap; +use std::future::Future; +use std::sync::{Arc, Mutex}; +use std::time::{Duration, Instant}; + +use rhai::{Array, Dynamic, Engine, EvalAltResult, Map, Position}; +use serde_json::{Value, json}; + +use super::types::{GraphBlueprintHandle, LanguageCompiler, ReplPolicy}; +use super::{ + ReplCallKind, ReplCallRecord, dynamic_to_repl_value, json_to_repl_value, repl_value_to_dynamic, +}; +use crate::error::TinyAgentsError; +use crate::harness::events::EventSink; +use crate::harness::ids::new_call_id; +use crate::harness::message::Message; +use crate::harness::model::ModelRequest; +use crate::harness::tool::ToolCall; +use crate::language::compiler::compile_with_provenance; +use crate::language::parser::parse_str; +use crate::language::resolver::Resolver; +use crate::language::types::Origin; +use crate::language::{Blueprint, blueprint_diff}; +use crate::registry::CapabilityRegistry; + +/// Session-cumulative counters for capability calls, enforced against the +/// `ReplPolicy` `max_*_calls` limits. Counts accumulate across cells (the +/// limits are documented per session) and are shared with every capability +/// closure on the engine. +#[derive(Debug, Default, Clone, Copy)] +pub(super) struct CallCounters { + /// `model_query` (and per-item `model_query_batched`) calls made. + pub model: usize, + /// `tool_call` (and per-item `tool_call_batched`) calls made. + pub tool: usize, + /// `graph_run` (and per-item `graph_run_batched`) calls made. + pub graph: usize, + /// `agent_query` (and per-item `agent_query_batched`) calls made. + pub agent: usize, + /// `graph_define` blueprints drafted. + pub graph_def: usize, +} + +/// The host-side context shared (via `Arc`) with every capability closure on a +/// session's engine. Holds the live registries, application state, policy, and +/// the shared per-cell buffers / session counters / graph drafts. +pub(super) struct HostContext { + /// The unified capability catalog (models, tools, graphs, agents). + pub registry: Arc>, + /// The application state capability calls are invoked against. + pub state: Arc, + /// The session policy (call/recursion/concurrency limits). + pub policy: ReplPolicy, + /// Optional expressive-language compiler handle (provenance label). + pub language: Option, + /// The session id, used as the generated-graph provenance label. + pub session_label: String, + /// The session's run depth, the parent depth for recursive sub-runs. + pub run_depth: usize, + /// The event sink shared with the run context. + pub events: EventSink, + /// Per-cell shared buffers (stdout, calls, answer, host error, vars). + pub buffers: super::CellBuffers, + /// Session-cumulative call counters. + pub counters: Arc>, + /// Graph blueprints drafted this session, keyed by name. + pub drafts: Arc>>, +} + +/// Drives an async capability future to completion synchronously, bounded by +/// an optional wall-clock `deadline` — the v1 "blocking bridge" adapter (see +/// the [module docs](self)), now with fail-closed enforcement of +/// [`ReplPolicy::timeout`]. +/// +/// `on_progress` (see [`build_engine`]) only fires between Rhai +/// statements/operations, so it can never interrupt a blocked native call: +/// this is the enforcement point for that case. When `deadline` elapses +/// first, the future is dropped — canceling the underlying request, since +/// providers are built on cancel-safe `reqwest`/`futures` — and a `Timeout` +/// error is returned instead of blocking the session forever. +fn bridge_block_on_raw( + deadline: Option, + future: F, +) -> std::result::Result { + let Some(deadline) = deadline else { + return Ok(futures::executor::block_on(future)); + }; + let now = Instant::now(); + if now >= deadline { + return Err(TinyAgentsError::Timeout(format!( + "{DEADLINE_EXCEEDED_TOKEN} before a host capability call could start" + ))); + } + let remaining = deadline - now; + let (tx, rx) = futures::channel::oneshot::channel::<()>(); + // A detached timer thread wakes the race below at the deadline. If the + // capability future finishes first this thread simply sleeps out its + // remainder and exits; nothing observes it after that. + std::thread::spawn(move || { + std::thread::sleep(remaining); + let _ = tx.send(()); + }); + match futures::executor::block_on(futures::future::select(Box::pin(future), rx)) { + futures::future::Either::Left((output, _timer)) => Ok(output), + futures::future::Either::Right(_) => Err(TinyAgentsError::Timeout(format!( + "{DEADLINE_EXCEEDED_TOKEN} during a host capability call" + ))), + } +} + +/// Convenience wrapper over [`bridge_block_on_raw`] for the common case where +/// the capability future itself resolves to a `Result`, flattening the +/// deadline error into the same error channel as the call's own failures. +fn bridge_block_on( + deadline: Option, + future: F, +) -> std::result::Result +where + F: Future>, +{ + bridge_block_on_raw(deadline, future)? +} + +/// One completed `model_query_batched` item: `(model, text, finish_reason, +/// structured, elapsed)`. +type ModelBatchItem = (String, String, Option, bool, Duration); + +/// One completed `agent_query_batched` item: `(agent, text, elapsed)`. +type AgentBatchItem = (String, String, Duration); + +// ── Error / recording helpers ─────────────────────────────────────────────── + +/// Stashes the precise crate error so `eval_cell` can surface it verbatim, and +/// returns the stringly-typed Rhai runtime error the engine propagates. +fn raise(ctx: &HostContext, err: TinyAgentsError) -> Box { + let message = err.to_string(); + ctx.buffers.set_host_error(err); + Box::new(EvalAltResult::ErrorRuntime( + Dynamic::from(message), + Position::NONE, + )) +} + +/// Raises a [`TinyAgentsError::Validation`] for an invalid script argument. +fn invalid( + ctx: &HostContext, + message: impl Into, +) -> Box { + raise(ctx, TinyAgentsError::Validation(message.into())) +} + +/// Records a capability call (or emitted event) into the per-cell buffer. +fn record( + ctx: &HostContext, + kind: ReplCallKind, + name: &str, + detail: Value, + elapsed: Duration, +) { + ctx.buffers.push_call(ReplCallRecord { + call_id: new_call_id(), + kind, + name: name.to_string(), + detail, + elapsed, + }); +} + +// ── Map argument helpers ──────────────────────────────────────────────────── + +/// Reads a string field from a Rhai object map argument. +fn map_str(map: &Map, key: &str) -> Option { + map.get(key).and_then(|d| d.clone().into_string().ok()) +} + +/// Reads a boolean field from a Rhai object map argument. +fn map_bool(map: &Map, key: &str) -> Option { + map.get(key).and_then(|d| d.as_bool().ok()) +} + +/// Converts a Rhai object map argument into a JSON value (for tool arguments +/// and structured payloads). +fn map_json(map: &Map, key: &str) -> Option { + map.get(key) + .map(|d| dynamic_to_repl_value(d).to_json()) + .filter(|v| !v.is_null()) +} + +// ── Counter limit helpers ─────────────────────────────────────────────────── + +/// Increments and bounds the model-call counter. +fn bump_model(ctx: &HostContext) -> Result<(), Box> { + let mut counters = ctx.counters.lock().expect("counters poisoned"); + if counters.model >= ctx.policy.max_model_calls { + return Err(raise( + ctx, + TinyAgentsError::LimitExceeded(format!( + "model call limit ({}) exceeded", + ctx.policy.max_model_calls + )), + )); + } + counters.model += 1; + Ok(()) +} + +/// Increments and bounds the tool-call counter. +fn bump_tool(ctx: &HostContext) -> Result<(), Box> { + let mut counters = ctx.counters.lock().expect("counters poisoned"); + if counters.tool >= ctx.policy.max_tool_calls { + return Err(raise( + ctx, + TinyAgentsError::LimitExceeded(format!( + "tool call limit ({}) exceeded", + ctx.policy.max_tool_calls + )), + )); + } + counters.tool += 1; + Ok(()) +} + +/// Increments and bounds the graph-run counter. +fn bump_graph(ctx: &HostContext) -> Result<(), Box> { + let mut counters = ctx.counters.lock().expect("counters poisoned"); + if counters.graph >= ctx.policy.max_graph_calls { + return Err(raise( + ctx, + TinyAgentsError::LimitExceeded(format!( + "graph call limit ({}) exceeded", + ctx.policy.max_graph_calls + )), + )); + } + counters.graph += 1; + Ok(()) +} + +/// Increments and bounds the agent-call counter. +fn bump_agent(ctx: &HostContext) -> Result<(), Box> { + let mut counters = ctx.counters.lock().expect("counters poisoned"); + if counters.agent >= ctx.policy.max_agent_calls { + return Err(raise( + ctx, + TinyAgentsError::LimitExceeded(format!( + "agent call limit ({}) exceeded", + ctx.policy.max_agent_calls + )), + )); + } + counters.agent += 1; + Ok(()) +} + +/// Enforces the recursion-depth bound for a sub-run (agent or graph). +/// +/// Reuses the harness recursion bookkeeping (Cluster G): a sub-run executes one +/// level below the session's run depth, and a child depth past +/// [`ReplPolicy::max_depth`] fails closed with +/// [`TinyAgentsError::SubAgentDepth`]. +fn check_depth(ctx: &HostContext) -> Result<(), Box> { + // Funnel the depth-cap check through the shared harness guard so the REPL + // sub-run bound stays in lock-step with SubAgent/SubAgentTool. + crate::harness::context::RunConfig::checked_child_depth(ctx.run_depth, ctx.policy.max_depth) + .map(|_| ()) + .map_err(|err| raise(ctx, err)) +} + +// ── Request builders ──────────────────────────────────────────────────────── + +/// Builds a [`ModelRequest`] from a `model_query` argument map. +fn build_model_request(model: &str, params: &Map) -> ModelRequest { + let mut messages = Vec::new(); + if let Some(system) = map_str(params, "system") { + messages.push(Message::system(system)); + } + if let Some(prompt) = map_str(params, "prompt") { + messages.push(Message::user(prompt)); + } + ModelRequest { + messages, + model: Some(model.to_string()), + ..Default::default() + } +} + +/// Wraps a model response text as the script-visible value (a string by +/// default, or a structured map when `structured: true`). +fn model_value(text: String, finish_reason: Option, structured: bool) -> Dynamic { + if structured { + let mut map = Map::new(); + map.insert("content".into(), Dynamic::from(text)); + if let Some(reason) = finish_reason { + map.insert("finish_reason".into(), Dynamic::from(reason)); + } + Dynamic::from_map(map) + } else { + Dynamic::from(text) + } +} + +mod authoring; +mod batched; +mod capabilities; + +use authoring::*; +use batched::*; +use capabilities::*; + +// ── Engine construction ───────────────────────────────────────────────────── + +/// Sentinel exception value `on_progress` terminates a script with when the +/// per-cell [`ReplPolicy::timeout`] deadline elapses. `eval_cell` recognizes +/// this exact string and maps it to `TinyAgentsError::Timeout` instead of the +/// generic runtime-error path. +pub(super) const DEADLINE_EXCEEDED_TOKEN: &str = "ragsh cell exceeded its wall-clock timeout"; + +/// Builds a sandboxed Rhai engine for a session, registering every host-backed +/// built-in against the session's live registries and policy. +/// +/// The engine is configured with the policy operation limit (fail-closed on +/// runaway scripts) and is granted no filesystem, network, or process access — +/// the only host surface is the capability functions registered here. +pub(super) fn build_engine(ctx: Arc>) -> Engine { + let mut engine = Engine::new(); + engine.set_max_operations(ctx.policy.max_operations); + + // Fail-closed wall-clock deadline: `eval_cell` arms `ctx.buffers`'s + // per-cell deadline before running the script. `on_progress` is polled + // between Rhai statements/operations, so this catches runaway *script* + // loops (a busy `while true {}` with no host calls) that `max_operations` + // alone might not bound tightly enough in wall-clock terms. Host + // capability calls (`model_query`, `tool_call`, …) are bounded separately + // by `bridge_block_on`, since a blocked native call never yields back to + // `on_progress`. + let deadline_ctx = ctx.clone(); + engine.on_progress(move |_ops| { + // A fail-closed host check (currently: push_stdout_line's + // max_output_bytes enforcement) may have stashed an error without + // Rhai itself failing; abort promptly instead of letting the script + // keep running until it happens to yield control back naturally. + // `eval_cell` prefers the stashed error over this sentinel's text. + if deadline_ctx.buffers.host_error_pending() { + return Some(Dynamic::from(DEADLINE_EXCEEDED_TOKEN.to_string())); + } + match deadline_ctx.buffers.deadline() { + Some(deadline) if Instant::now() >= deadline => { + Some(Dynamic::from(DEADLINE_EXCEEDED_TOKEN.to_string())) + } + _ => None, + } + }); + + // ── stdout capture ── + let stdout_ctx = ctx.clone(); + engine.on_print(move |text| stdout_ctx.buffers.push_stdout_line(text)); + let debug_ctx = ctx.clone(); + engine.on_debug(move |text, _source, _pos| debug_ctx.buffers.push_stdout_line(text)); + + // ── emit(name) / emit(name, #{ ... }) ── + let emit_ctx = ctx.clone(); + engine.register_fn("emit", move |name: &str| { + record( + &emit_ctx, + ReplCallKind::Emit, + name, + Value::Null, + Duration::default(), + ); + }); + let emit_payload_ctx = ctx.clone(); + engine.register_fn("emit", move |name: &str, data: Map| { + let detail = dynamic_to_repl_value(&Dynamic::from_map(data)).to_json(); + record( + &emit_payload_ctx, + ReplCallKind::Emit, + name, + detail, + Duration::default(), + ); + }); + + // ── answer(content) ── + let answer_ctx = ctx.clone(); + engine.register_fn("answer", move |content: &str| { + answer_ctx.buffers.set_answer(content.to_string()); + }); + + // ── show_vars() ── + let show_ctx = ctx.clone(); + engine.register_fn("show_vars", move || { + show_ctx.buffers.push_stdout_line("# vars"); + for (name, value) in show_ctx.buffers.vars_snapshot() { + show_ctx + .buffers + .push_stdout_line(&format!("{name} = {value}")); + } + }); + + // ── model capabilities ── + let model_ctx = ctx.clone(); + engine.register_fn("model_query", move |params: Map| { + model_query_impl(&model_ctx, ¶ms) + }); + let model_batch_ctx = ctx.clone(); + engine.register_fn("model_query_batched", move |items: Array| { + model_query_batched_impl(&model_batch_ctx, &items) + }); + + // ── tool capabilities ── + let tool_ctx = ctx.clone(); + engine.register_fn("tool_call", move |params: Map| { + tool_call_impl(&tool_ctx, ¶ms) + }); + let tool_batch_ctx = ctx.clone(); + engine.register_fn("tool_call_batched", move |items: Array| { + tool_call_batched_impl(&tool_batch_ctx, &items) + }); + + // ── agent capabilities ── + let agent_ctx = ctx.clone(); + engine.register_fn("agent_query", move |params: Map| { + agent_query_impl(&agent_ctx, ¶ms) + }); + let agent_batch_ctx = ctx.clone(); + engine.register_fn("agent_query_batched", move |items: Array| { + agent_query_batched_impl(&agent_batch_ctx, &items) + }); + + // ── graph run capabilities ── + let graph_ctx = ctx.clone(); + engine.register_fn("graph_run", move |params: Map| { + graph_run_impl(&graph_ctx, ¶ms) + }); + let graph_batch_ctx = ctx.clone(); + engine.register_fn("graph_run_batched", move |items: Array| { + graph_run_batched_impl(&graph_batch_ctx, &items) + }); + + // ── graph authoring (lowering through the `.rag` compiler) ── + let define_ctx = ctx.clone(); + engine.register_fn("graph_define", move |params: Map| { + graph_define_impl(&define_ctx, ¶ms) + }); + let validate_ctx = ctx.clone(); + engine.register_fn("graph_validate", move |descriptor: Map| { + graph_validate_impl(&validate_ctx, &descriptor) + }); + let compile_ctx = ctx.clone(); + engine.register_fn("graph_compile", move |descriptor: Map| { + graph_compile_impl(&compile_ctx, &descriptor) + }); + let diff_name_ctx = ctx.clone(); + engine.register_fn( + "graph_diff", + move |name: &str, draft: Map| -> Result> { + let old = diff_name_ctx + .registry + .graph_blueprint(name) + .ok_or_else(|| { + invalid( + &diff_name_ctx, + format!("graph_diff: graph `{name}` is not registered"), + ) + })? + .clone(); + let new = lookup_draft(&diff_name_ctx, &draft, "graph_diff")?; + graph_diff_handles(&diff_name_ctx, &old, &new.blueprint) + }, + ); + let diff_draft_ctx = ctx.clone(); + engine.register_fn( + "graph_diff", + move |old: Map, new: Map| -> Result> { + let old = lookup_draft(&diff_draft_ctx, &old, "graph_diff")?; + let new = lookup_draft(&diff_draft_ctx, &new, "graph_diff")?; + graph_diff_handles(&diff_draft_ctx, &old.blueprint, &new.blueprint) + }, + ); + let register_ctx = ctx.clone(); + engine.register_fn("graph_register", move |params: Map| { + graph_register_impl(®ister_ctx, ¶ms) + }); + + engine +} + +#[cfg(test)] +mod bridge_deadline_test { + use super::*; + + #[test] + fn no_deadline_awaits_to_completion() { + let out = bridge_block_on::(None, async { Ok(7) }).expect("no deadline"); + assert_eq!(out, 7); + } + + #[test] + fn future_finishing_before_the_deadline_succeeds() { + let deadline = Instant::now() + Duration::from_secs(5); + let out = + bridge_block_on::(Some(deadline), async { Ok(9) }).expect("within deadline"); + assert_eq!(out, 9); + } + + #[test] + fn deadline_already_elapsed_fails_closed_without_starting_the_call() { + // Regression test: `ReplPolicy::timeout` used to be parsed but never + // enforced anywhere a host capability call could hang forever. + let deadline = Instant::now() - Duration::from_millis(1); + let err = bridge_block_on::(Some(deadline), async { Ok(1) }) + .expect_err("deadline already passed"); + assert!(matches!(err, TinyAgentsError::Timeout(_)), "got {err:?}"); + } + + #[test] + fn a_hanging_call_is_cut_off_at_the_deadline_instead_of_blocking_forever() { + // A future that never resolves models a hung provider/tool call. The + // deadline must still return control promptly rather than hanging the + // whole `eval_cell` (and therefore the session) forever. + let start = Instant::now(); + let deadline = start + Duration::from_millis(30); + let err = bridge_block_on::( + Some(deadline), + futures::future::pending::>(), + ) + .expect_err("hanging call must be cut off at the deadline"); + assert!(matches!(err, TinyAgentsError::Timeout(_)), "got {err:?}"); + assert!( + start.elapsed() < Duration::from_secs(5), + "took {:?}, should return promptly at the 30ms deadline", + start.elapsed() + ); + } +} diff --git a/src/repl/session/mod.rs b/src/repl/session/mod.rs index 7543585..6730143 100644 --- a/src/repl/session/mod.rs +++ b/src/repl/session/mod.rs @@ -72,6 +72,19 @@ pub(super) struct CellBuffers { answer: Arc>>, host_error: Arc>>, vars_snapshot: Arc>>, + /// The wall-clock instant the current cell's [`ReplPolicy::timeout`] + /// expires at, if the policy configures one. Set at the start of + /// [`ReplSession::eval_cell`] and read by every host capability call (via + /// [`builtins::bridge_block_on`]) and the engine's `on_progress` hook, so + /// the deadline is enforced fail-closed both for pure script loops and for + /// in-flight model/tool/agent/graph calls. + deadline: Arc>>, + /// The current cell's [`ReplPolicy::max_output_bytes`] budget, armed at + /// the start of every [`ReplSession::eval_cell`] call and enforced + /// fail-closed inside [`CellBuffers::push_stdout_line`] itself, so a + /// print-heavy runaway script cannot buffer unbounded output before the + /// end-of-cell check in `eval_cell` ever runs. + max_output_bytes: Arc>>, } /// The persistent variable namespace of a session. @@ -159,6 +172,24 @@ impl Default for ReplVariables { /// stateless session with [`ReplSession::new`]; supply registries, a custom /// policy, or a run context with [`ReplSession::builder`]-style `with_*` /// methods. +/// +/// # Not to be confused with `repl::ReplSession` +/// +/// This crate has two distinct types named `ReplSession`: +/// +/// - **This type** (`repl::session::ReplSession`, feature `repl` only) — the +/// Rhai-backed scripting session described above; also reachable as +/// `crate::ReplSession` (the crate-root re-export) when the `repl` feature +/// is enabled. +/// - [`crate::repl::ReplSession`] (always available, no feature required) — +/// the line-oriented command skeleton (verbs like `set`/`get`/`run`/`call` +/// parsed from a single line). It is *not* re-exported at the crate root +/// under this feature, so `crate::ReplSession` only ever means this +/// scripting session once `repl` is enabled. +/// +/// The two are unrelated types serving different layers of the `.ragsh` +/// design; always check which module path (`crate::repl::session::ReplSession` +/// vs. `crate::repl::ReplSession`) you imported from. pub struct ReplSession where State: Send + Sync, @@ -195,6 +226,10 @@ where engine: Engine, /// Shared buffers the engine's built-ins write into. buffers: CellBuffers, + /// Number of cells evaluated so far this session, enforced fail-closed + /// against [`ReplPolicy::max_iterations`]. Each `eval_cell` call is one + /// CodeAct-style iteration of a model-driven session. + iterations: usize, } impl ReplSession { @@ -248,6 +283,7 @@ impl ReplSession { drafts: Arc::new(Mutex::new(BTreeMap::new())), engine: Engine::new(), buffers, + iterations: 0, }; session.rebuild_engine(); session @@ -302,9 +338,9 @@ impl ReplSession { /// Returns a shared handle to the application state capability calls are /// invoked against. /// - /// The CodeAct driver ([`crate::repl::codeact`]) needs this to invoke the - /// session's driver model through the same state the in-cell capability - /// functions use, without exposing the private field. + /// A CodeAct-style driver loop needs this to invoke the session's driver + /// model through the same state the in-cell capability functions use, + /// without exposing the private field. pub fn app_state(&self) -> Arc { self.state.clone() } @@ -332,13 +368,27 @@ impl ReplSession { /// /// * [`TinyAgentsError::LimitExceeded`] — the script exceeds /// [`ReplPolicy::max_script_bytes`], the output exceeds - /// [`ReplPolicy::max_output_bytes`], or the script exceeds the engine - /// operation limit (fail-closed runaway protection). + /// [`ReplPolicy::max_output_bytes`], the engine operation limit + /// (fail-closed runaway protection), or the session has already + /// evaluated [`ReplPolicy::max_iterations`] cells. + /// * [`TinyAgentsError::Timeout`] — the cell's wall-clock deadline + /// ([`ReplPolicy::timeout`]) elapsed, either mid-script or during a + /// model/tool/agent/graph call. /// * [`TinyAgentsError::Validation`] — the script failed to compile or /// raised a runtime error. pub fn eval_cell(&mut self, script: &str) -> Result { let start = Instant::now(); + // Each call is one CodeAct-style iteration of a model-driven session; + // enforce the cap fail-closed before doing any other work. + if self.iterations >= self.policy.max_iterations { + return Err(TinyAgentsError::LimitExceeded(format!( + "ragsh session has evaluated {} cells, reaching the max_iterations limit of {}", + self.iterations, self.policy.max_iterations + ))); + } + self.iterations += 1; + if script.len() > self.policy.max_script_bytes { return Err(TinyAgentsError::LimitExceeded(format!( "ragsh cell is {} bytes, exceeding the max_script_bytes limit of {}", @@ -347,16 +397,27 @@ impl ReplSession { ))); } - // Reset per-cell shared buffers. + // Reset per-cell shared buffers and arm the wall-clock deadline (if + // the policy configures one) before any script or host-capability + // work begins. `on_progress` (see `builtins::build_engine`) enforces + // it for pure script execution; `bridge_block_on` enforces it around + // every model/tool/agent/graph call so a hanging host call cannot + // block the session forever either. self.buffers.reset(); - - let before = self.variables.snapshot(); - // Expose the pre-cell namespace to `show_vars()`. + self.buffers + .arm_deadline(self.policy.timeout.map(|d| start + d)); + self.buffers.arm_output_limit(self.policy.max_output_bytes); + + // Snapshot the pre-cell namespace once and move it into the shared + // `vars_snapshot` (read by `show_vars()` during the cell). The diff below + // reads this same baseline back rather than keeping a second full copy, + // so a cell pays one baseline snapshot instead of a snapshot *plus* a + // full O(namespace-bytes) clone. *self .buffers .vars_snapshot .lock() - .expect("vars_snapshot poisoned") = before.clone(); + .expect("vars_snapshot poisoned") = self.variables.snapshot(); // Disjoint field borrows: the engine is read-only while the scope is // mutated in place, so top-level `let` bindings persist into the scope. @@ -369,7 +430,17 @@ impl ReplSession { self.variables.restore_reserved(); let value_dynamic = match eval { - Ok(value) => value, + Ok(value) => { + // The script may have completed "successfully" from Rhai's + // point of view even though a host-side fail-closed check + // (e.g. push_stdout_line's max_output_bytes enforcement) + // stashed an error — `on_print`/`on_debug` cannot themselves + // fail a script, so this is the only place that catches it. + if let Some(host_err) = self.buffers.take_host_error() { + return Err(host_err); + } + value + } Err(err) => { // A fallible capability function stashes its precise crate error // here; prefer it over the generic Rhai runtime wrapper so the @@ -402,7 +473,16 @@ impl ReplSession { } let after = self.variables.snapshot(); - let variables_changed = diff_changed(&before, &after); + // Diff against the baseline stored in `vars_snapshot` instead of a + // separately retained `before` map, avoiding a redundant full copy. + let variables_changed = { + let before = self + .buffers + .vars_snapshot + .lock() + .expect("vars_snapshot poisoned"); + diff_changed(&before, &after) + }; Ok(ReplResult { stdout, @@ -421,6 +501,32 @@ impl CellBuffers { self.calls.lock().expect("calls poisoned").clear(); *self.answer.lock().expect("answer poisoned") = None; *self.host_error.lock().expect("host_error poisoned") = None; + *self.deadline.lock().expect("deadline poisoned") = None; + *self + .max_output_bytes + .lock() + .expect("max_output_bytes poisoned") = None; + } + + /// Arms the per-cell wall-clock deadline, replacing any previous one. + fn arm_deadline(&self, deadline: Option) { + *self.deadline.lock().expect("deadline poisoned") = deadline; + } + + /// Arms the per-cell output-byte budget, replacing any previous one. + /// Read by [`CellBuffers::push_stdout_line`] on every captured line. + fn arm_output_limit(&self, max_bytes: usize) { + *self + .max_output_bytes + .lock() + .expect("max_output_bytes poisoned") = Some(max_bytes); + } + + /// Returns the current cell's wall-clock deadline, if the policy + /// configured a timeout. Read by every host capability call and the + /// engine's `on_progress` hook (see [`builtins::bridge_block_on`]). + pub(super) fn deadline(&self) -> Option { + *self.deadline.lock().expect("deadline poisoned") } fn stdout(&self) -> String { @@ -446,13 +552,44 @@ impl CellBuffers { self.calls.lock().expect("calls poisoned").push(record); } - /// Appends a line to the captured stdout buffer. + /// Appends a line to the captured stdout buffer, enforcing the armed + /// [`ReplPolicy::max_output_bytes`] budget fail-closed: once appending + /// would exceed the budget, the line is dropped (not buffered) and a + /// [`TinyAgentsError::LimitExceeded`] is stashed for `eval_cell` to + /// surface, instead of growing the buffer without bound for the rest of + /// the cell. pub(super) fn push_stdout_line(&self, line: &str) { + let limit = *self + .max_output_bytes + .lock() + .expect("max_output_bytes poisoned"); let mut out = self.stdout.lock().expect("stdout poisoned"); + if let Some(limit) = limit { + let projected = out.len() + line.len() + 1; + if projected > limit { + drop(out); + self.set_host_error(TinyAgentsError::LimitExceeded(format!( + "ragsh cell produced more than {limit} bytes of output, exceeding the max_output_bytes limit" + ))); + return; + } + } out.push_str(line); out.push('\n'); } + /// Returns whether a host error is currently stashed, without consuming + /// it. Used by the engine's `on_progress` hook to abort a script promptly + /// once [`push_stdout_line`](Self::push_stdout_line) has flagged the + /// output budget as exceeded, rather than letting the script keep running + /// until it happens to yield control back naturally. + pub(super) fn host_error_pending(&self) -> bool { + self.host_error + .lock() + .expect("host_error poisoned") + .is_some() + } + /// Sets the session's final answer. pub(super) fn set_answer(&self, content: String) { *self.answer.lock().expect("answer poisoned") = Some(content); @@ -480,6 +617,15 @@ fn map_rhai_error(err: EvalAltResult) -> TinyAgentsError { EvalAltResult::ErrorTooManyOperations(pos) => TinyAgentsError::LimitExceeded(format!( "ragsh cell exceeded the operation limit (max_operations) at {pos}" )), + // The engine's `on_progress` hook (see `builtins::build_engine`) + // terminates the script with this exact sentinel value once the + // per-cell `ReplPolicy::timeout` deadline elapses. + EvalAltResult::ErrorTerminated(token, pos) + if token.clone().into_string().ok().as_deref() + == Some(builtins::DEADLINE_EXCEEDED_TOKEN) => + { + TinyAgentsError::Timeout(format!("{} at {pos}", builtins::DEADLINE_EXCEEDED_TOKEN)) + } other => TinyAgentsError::Validation(format!("ragsh evaluation error: {other}")), } } diff --git a/src/repl/session/test.rs b/src/repl/session/test.rs index 6b7e562..2d59cc3 100644 --- a/src/repl/session/test.rs +++ b/src/repl/session/test.rs @@ -1,5 +1,7 @@ //! Unit tests for the Rhai-backed `.ragsh` session runtime. +use std::time::{Duration, Instant}; + use super::*; use crate::error::TinyAgentsError; @@ -41,6 +43,28 @@ fn variables_persist_across_cells() { assert_eq!(fourth.value, Some(ReplValue::Int(10))); } +#[test] +fn variables_changed_diffs_against_shared_baseline() { + // The pre-cell baseline is snapshotted once into the shared `vars_snapshot` + // and the change diff reads it back from there (no second retained copy). + // A cell that introduces several new bindings and mutates an existing one + // must still report every changed name, and leave an unchanged binding out. + let mut s = session(); + s.eval_cell("let kept = 1; let touched = 2;").expect("seed"); + + let result = s + .eval_cell("let a = 10; let b = 20; touched = 99; kept") + .expect("cell"); + + assert!(result.variables_changed.contains(&"a".to_string())); + assert!(result.variables_changed.contains(&"b".to_string())); + assert!(result.variables_changed.contains(&"touched".to_string())); + assert!( + !result.variables_changed.contains(&"kept".to_string()), + "an unmodified binding is not reported as changed" + ); +} + #[test] fn over_limit_script_fails_closed() { // A tiny operation budget makes an otherwise-bounded loop trip the limit. @@ -62,6 +86,54 @@ fn over_limit_script_fails_closed() { } } +#[test] +fn timeout_fails_closed_on_a_runaway_script() { + // Regression test: `ReplPolicy::timeout` used to be parsed but never + // enforced — a runaway or hanging cell could block the session forever. + // `max_operations` is left effectively unbounded here so only the + // wall-clock deadline (enforced via the engine's `on_progress` hook) can + // stop the loop. + let policy = ReplPolicy { + timeout: Some(Duration::from_millis(30)), + max_operations: 0, + ..ReplPolicy::default() + }; + let mut s = ReplSession::<()>::new().with_policy(policy); + + let start = Instant::now(); + let err = s + .eval_cell("let total = 0; loop { total += 1; }") + .expect_err("should exceed the wall-clock deadline"); + assert!(matches!(err, TinyAgentsError::Timeout(_)), "got {err:?}"); + // The property under test: `eval_cell` returns promptly once the + // deadline elapses rather than running the loop forever. + assert!( + start.elapsed() < Duration::from_secs(5), + "eval_cell took {:?}, should have returned near the 30ms deadline", + start.elapsed() + ); +} + +#[test] +fn max_iterations_limit_fails_closed() { + // Regression test: `ReplPolicy::max_iterations` was parsed and defaulted + // but no code path ever checked it. + let policy = ReplPolicy { + max_iterations: 2, + ..ReplPolicy::default() + }; + let mut s = ReplSession::<()>::new().with_policy(policy); + + s.eval_cell("1").expect("cell 1 within the limit"); + s.eval_cell("2").expect("cell 2 within the limit"); + + let err = s.eval_cell("3").expect_err("cell 3 exceeds max_iterations"); + assert!( + matches!(err, TinyAgentsError::LimitExceeded(_)), + "got {err:?}" + ); +} + #[test] fn script_byte_limit_fails_closed() { let policy = ReplPolicy { @@ -98,6 +170,32 @@ fn reserved_names_are_restored_after_each_cell() { assert!(!result.variables_changed.contains(&"context".to_string())); } +#[test] +fn reserved_names_contains_no_duplicates() { + // `answer` is a capability function (see RESERVED_FUNCTIONS), not a + // readable session variable; it must not also appear in + // RESERVED_VARIABLES, or `ReplVariables::seeded` double-pushes the same + // scope entry. + let names: Vec<&str> = reserved_names().collect(); + let mut seen = std::collections::HashSet::new(); + for name in &names { + assert!(seen.insert(*name), "duplicate reserved name: {name}"); + } + assert!(names.contains(&"answer")); +} + +#[test] +fn answer_variable_is_seeded_exactly_once_in_scope() { + let s = session(); + let count = s + .variables + .scope + .iter() + .filter(|(name, _, _)| *name == "answer") + .count(); + assert_eq!(count, 1, "`answer` must be seeded into scope exactly once"); +} + #[test] fn reserved_capability_name_cannot_be_set_as_a_variable() { let mut s = session(); @@ -153,6 +251,237 @@ fn output_byte_limit_fails_closed() { assert!(matches!(err, TinyAgentsError::LimitExceeded(_))); } +#[test] +fn output_byte_limit_bounds_intra_cell_buffering_in_a_print_loop() { + // A script that prints in a tight loop must not be allowed to buffer + // unbounded output before the limit is noticed: push_stdout_line itself + // must stop growing the buffer (and eval_cell must fail closed) well + // before the loop's total output would otherwise reach many times the + // configured budget. + let policy = ReplPolicy { + max_output_bytes: 100, + max_operations: 1_000_000, + ..ReplPolicy::default() + }; + let mut s = ReplSession::<()>::new().with_policy(policy); + + let err = s + .eval_cell( + r#"for i in 0..100000 { print("0123456789012345678901234567890123456789012345"); }"#, + ) + .expect_err("should exceed the output byte limit"); + assert!(matches!(err, TinyAgentsError::LimitExceeded(_)), "{err:?}"); +} + +#[test] +fn graph_define_does_not_consume_the_limit_on_a_failed_draft() { + // A `graph_define` call whose source parses but names a graph that isn't + // in the source must not consume a definition slot: only a successfully + // recorded draft should count against `max_graph_definitions`. + let policy = ReplPolicy { + max_graph_definitions: 1, + ..ReplPolicy::default() + }; + let mut s = ReplSession::<()>::new().with_policy(policy); + + let source = r#"graph g { start a node a { kind model next END } }"#; + + // First call: wrong graph name, so the draft is never recorded — this + // must fail without spending the one available slot. + let bad = s.eval_cell(&format!( + r#"graph_define(#{{ name: "missing", source: `{source}` }})"# + )); + assert!( + bad.is_err(), + "expected a failure for the unknown graph name" + ); + + // Second call: the correct graph name must still succeed, proving the + // failed attempt above did not consume the definition budget. + let good = s + .eval_cell(&format!( + r#"graph_define(#{{ name: "g", source: `{source}` }})"# + )) + .expect("a valid graph_define should still have a slot available"); + assert!(good.value.is_some()); + + // A third attempt now must fail: the one slot has genuinely been spent. + let over_limit = s.eval_cell(&format!( + r#"graph_define(#{{ name: "g", source: `{source}` }})"# + )); + assert!( + over_limit.is_err(), + "the definition limit must be enforced once a slot is actually consumed" + ); +} + +/// A tool that succeeds for every call except one whose `arguments.id` +/// matches `fail_id`, for which it returns a *tool-reported* error (a +/// `ToolResult` with `error: Some(..)`), not a `Result::Err` — exercising the +/// per-item error path distinct from a harness/transport-level failure. +struct SometimesFailingTool { + fail_id: String, +} + +#[async_trait::async_trait] +impl crate::harness::tool::Tool<()> for SometimesFailingTool { + fn name(&self) -> &str { + "sometimes_fails" + } + + fn description(&self) -> &str { + "Succeeds unless called with the configured failing id." + } + + fn schema(&self) -> crate::harness::tool::ToolSchema { + crate::harness::tool::ToolSchema { + name: self.name().to_string(), + description: self.description().to_string(), + parameters: serde_json::json!({ "type": "object" }), + format: Default::default(), + } + } + + async fn call( + &self, + _state: &(), + call: crate::harness::tool::ToolCall, + ) -> crate::Result { + let id = call + .arguments + .get("id") + .and_then(|v| v.as_str()) + .unwrap_or_default(); + if id == self.fail_id { + Ok(crate::harness::tool::ToolResult::error( + call.id, + call.name, + format!("tool reported an error for id {id}"), + )) + } else { + Ok(crate::harness::tool::ToolResult::text( + call.id, + call.name, + format!("ok:{id}"), + )) + } + } +} + +fn session_with_sometimes_failing_tool(fail_id: &str) -> ReplSession { + let mut registry = crate::registry::CapabilityRegistry::<()>::new(); + registry + .register_tool(std::sync::Arc::new(SometimesFailingTool { + fail_id: fail_id.to_string(), + })) + .expect("register tool"); + let capabilities = ReplCapabilities::new(std::sync::Arc::new(registry)); + ReplSession::<()>::new().with_capabilities(capabilities) +} + +#[test] +fn tool_call_batched_keeps_successes_when_one_item_tool_errors() { + // Regression test: a per-item *tool-reported* error (ToolResult::error, + // as opposed to a harness/transport-level Err) used to abort the whole + // batch, discarding every other item's already-computed successful + // result. Each item's outcome must be reported independently. + let mut s = session_with_sometimes_failing_tool("2"); + + let script = r#" + tool_call_batched([ + #{ tool: "sometimes_fails", arguments: #{ id: "1" } }, + #{ tool: "sometimes_fails", arguments: #{ id: "2" } }, + #{ tool: "sometimes_fails", arguments: #{ id: "3" } }, + ]) + "#; + let result = s.eval_cell(script).expect("batch call should not abort"); + let value = result.value.expect("value").to_json(); + let items = value.as_array().expect("array result"); + assert_eq!(items.len(), 3, "{items:?}"); + + assert_eq!(items[0]["ok"], serde_json::json!(true)); + assert_eq!(items[0]["content"], serde_json::json!("ok:1")); + + assert_eq!(items[1]["ok"], serde_json::json!(false)); + assert!( + items[1]["error"] + .as_str() + .unwrap() + .contains("tool reported an error"), + "{items:?}" + ); + + assert_eq!(items[2]["ok"], serde_json::json!(true)); + assert_eq!(items[2]["content"], serde_json::json!("ok:3")); +} + +/// A trivial [`HarnessAgent`] that returns a fixed response, for exercising +/// `agent_query` without a real model/harness run. +struct StubAgent; + +#[async_trait::async_trait] +impl crate::graph::subagent_node::HarnessAgent for StubAgent { + fn name(&self) -> &str { + "stub" + } + + async fn run( + &self, + input: crate::graph::subagent_node::SubAgentInput, + _events: crate::harness::events::EventSink, + ) -> crate::Result { + Ok(crate::graph::subagent_node::SubAgentOutput { + text: format!("stub replied to: {}", input.prompt), + ..Default::default() + }) + } +} + +fn session_with_stub_agent(policy: ReplPolicy) -> ReplSession { + let mut registry = crate::registry::CapabilityRegistry::<()>::new(); + registry + .register_agent(std::sync::Arc::new(StubAgent)) + .expect("register stub agent"); + let capabilities = ReplCapabilities::new(std::sync::Arc::new(registry)); + ReplSession::<()>::new() + .with_policy(policy) + .with_capabilities(capabilities) +} + +#[test] +fn agent_call_limit_is_independent_of_the_model_call_limit() { + // Regression test: `bump_agent` used to compare the agent-call counter + // against `max_model_calls` (with an "agent call limit" message quoting + // that same number), so a session's *combined* model spend — direct + // `model_query` calls plus every model call a delegated `agent_query` + // itself drives — could reach roughly twice the configured + // `max_model_calls` before anything failed closed. `max_agent_calls` is + // now tracked and enforced independently. + let policy = ReplPolicy { + max_model_calls: 64, + max_agent_calls: 2, + ..ReplPolicy::default() + }; + let mut s = session_with_stub_agent(policy); + + let script = r#"agent_query(#{ agent: "stub", prompt: "hi" })"#; + s.eval_cell(script).expect("call 1 within the limit"); + s.eval_cell(script).expect("call 2 within the limit"); + + let err = s + .eval_cell(script) + .expect_err("call 3 exceeds max_agent_calls"); + match err { + TinyAgentsError::LimitExceeded(msg) => { + assert!( + msg.contains("agent call limit (2)"), + "expected the message to cite max_agent_calls (2), got: {msg}" + ); + } + other => panic!("expected LimitExceeded, got {other:?}"), + } +} + #[test] fn map_and_array_values_round_trip_to_json() { let mut s = session(); diff --git a/src/repl/session/types.rs b/src/repl/session/types.rs index e622790..2058bb1 100644 --- a/src/repl/session/types.rs +++ b/src/repl/session/types.rs @@ -14,7 +14,6 @@ use std::time::Duration; use serde::{Deserialize, Serialize}; -use crate::harness::store::StoreRegistry; use crate::language::types::{Blueprint, Origin}; use crate::registry::CapabilityRegistry; @@ -24,9 +23,11 @@ use crate::registry::CapabilityRegistry; /// /// These are restored to their session baseline after each cell so a script /// can read or temporarily shadow them but cannot permanently replace the -/// session's context, state, or final-answer slots. -pub const RESERVED_VARIABLES: &[&str] = - &["context", "state", "messages", "history", "run", "answer"]; +/// session's context, state, or run slots. `answer` is *not* included here: +/// it is a capability function only (see [`RESERVED_FUNCTIONS`]), never a +/// readable session variable, so listing it in both would seed the scope +/// with a duplicate entry for the same name. +pub const RESERVED_VARIABLES: &[&str] = &["context", "state", "messages", "history", "run"]; /// Reserved built-in *capability function* names. /// @@ -56,12 +57,17 @@ pub const RESERVED_FUNCTIONS: &[&str] = &[ ]; /// Returns every reserved name (variables and capability functions) the -/// runtime must protect across cells. +/// runtime must protect across cells, each name yielded at most once even if +/// it were (accidentally) listed in both [`RESERVED_VARIABLES`] and +/// [`RESERVED_FUNCTIONS`] — callers seed one scope entry per yielded name, so +/// a duplicate here would silently double-push the same variable. pub fn reserved_names() -> impl Iterator { + let mut seen = std::collections::HashSet::new(); RESERVED_VARIABLES .iter() .copied() .chain(RESERVED_FUNCTIONS.iter().copied()) + .filter(move |name| seen.insert(*name)) } // ── Policy ────────────────────────────────────────────────────────────────── @@ -85,6 +91,14 @@ pub struct ReplPolicy { pub max_output_bytes: usize, /// Maximum `model_query` calls per session. pub max_model_calls: usize, + /// Maximum `agent_run`/sub-agent calls per session. + /// + /// Enforced independently of [`max_model_calls`][Self::max_model_calls]: + /// each sub-agent call itself drives one or more model calls, so + /// capping agent calls at the model-call budget (as an earlier + /// implementation did) let a session's *combined* model spend reach + /// roughly twice the configured `max_model_calls`. + pub max_agent_calls: usize, /// Maximum `tool_call` calls per session. pub max_tool_calls: usize, /// Maximum `graph_run` calls per session. @@ -110,6 +124,7 @@ impl Default for ReplPolicy { max_script_bytes: 64 * 1024, max_output_bytes: 256 * 1024, max_model_calls: 64, + max_agent_calls: 32, max_tool_calls: 128, max_graph_calls: 32, max_graph_definitions: 8, @@ -293,19 +308,24 @@ impl Default for LanguageCompiler { /// `GraphRegistry`, and `AgentRegistry` fields. In this crate those four kinds /// are already unified under the single name-addressable /// [`CapabilityRegistry`], so `ReplCapabilities` wraps that registry (shared via -/// `Arc` so a session can be cheaply cloned into a graph node) plus the -/// long-term [`StoreRegistry`] and an optional [`LanguageCompiler`]. The -/// per-kind accessors ([`models`](Self::models), [`tools`](Self::tools), -/// [`graphs`](Self::graphs), [`agents`](Self::agents)) preserve the documented -/// surface. +/// `Arc` so a session can be cheaply cloned into a graph node) plus an optional +/// [`LanguageCompiler`]. The per-kind accessors ([`models`](Self::models), +/// [`tools`](Self::tools), [`graphs`](Self::graphs), [`agents`](Self::agents)) +/// preserve the documented surface. +/// +/// A prior revision also carried a [`crate::harness::store::StoreRegistry`] +/// field, but no built-in +/// (`model_query`, `tool_call`, …) ever read or wrote through it — it was dead +/// weight advertising a capability the engine did not actually expose. It was +/// removed rather than left half-wired; long-term store access can be added +/// back as real `store_get`/`store_set` built-ins (see [`super::builtins`]) +/// once that surface is designed. pub struct ReplCapabilities where State: Send + Sync, { /// The unified capability catalog (models, tools, graphs, agents). pub registry: Arc>, - /// Named long-term stores available to the session. - pub stores: StoreRegistry, /// Optional expressive-language compiler handle for graph drafting. pub language: Option, } @@ -315,7 +335,6 @@ impl ReplCapabilities { pub fn new(registry: Arc>) -> Self { Self { registry, - stores: StoreRegistry::new(), language: None, } } @@ -326,12 +345,6 @@ impl ReplCapabilities { self } - /// Replaces the store registry with a (possibly shared) one. - pub fn with_stores(mut self, stores: StoreRegistry) -> Self { - self.stores = stores; - self - } - /// Returns the registered model names. pub fn models(&self) -> Vec { self.registry.names(crate::registry::ComponentKind::Model) diff --git a/src/repl/test.rs b/src/repl/test.rs index 12d7531..2fc8686 100644 --- a/src/repl/test.rs +++ b/src/repl/test.rs @@ -239,6 +239,47 @@ fn error_on_unterminated_quoted_string() { assert!(matches!(err, TinyAgentsError::Parse { .. })); } +#[test] +fn parse_errors_report_real_positions_not_0_0() { + // `parse_command` always parses one line, so `line` is always 1; `column` + // must point at the offending token rather than always reporting (0, 0). + match parse_command("frobnicate something").unwrap_err() { + TinyAgentsError::Parse { line, column, .. } => { + assert_eq!(line, 1); + assert_eq!(column, 1, "unknown verb starts at column 1"); + } + other => panic!("expected Parse error, got {other:?}"), + } + + match parse_command("run my_graph").unwrap_err() { + TinyAgentsError::Parse { line, column, .. } => { + assert_eq!(line, 1); + // "run my_graph" is 12 chars; the missing second argument is + // reported at the end of input, not column 0. + assert_eq!(column, 13); + } + other => panic!("expected Parse error, got {other:?}"), + } + + match parse_command("call my_cap not-valid-json").unwrap_err() { + TinyAgentsError::Parse { line, column, .. } => { + assert_eq!(line, 1); + // "not-valid-json" starts right after "call my_cap ". + assert_eq!(column, "call my_cap ".chars().count() + 1); + } + other => panic!("expected Parse error, got {other:?}"), + } + + match parse_command(r#"load "unclosed"#).unwrap_err() { + TinyAgentsError::Parse { line, column, .. } => { + assert_eq!(line, 1); + // The unterminated string starts right after "load ". + assert_eq!(column, "load ".chars().count() + 1); + } + other => panic!("expected Parse error, got {other:?}"), + } +} + #[test] fn quoted_string_escape_sequences() { let cmd = parse_command(r#"set msg "hello \"world\"\nnewline""#).unwrap(); diff --git a/src/repl/types.rs b/src/repl/types.rs index ee04aef..24d2c81 100644 --- a/src/repl/types.rs +++ b/src/repl/types.rs @@ -197,7 +197,31 @@ impl CapabilityPolicy { /// An interactive REPL session holding session variables, a capability policy, /// and a command history. /// -/// `ReplSession` is the primary entry point for driving the `.ragsh` skeleton. +/// `ReplSession` is the primary entry point for driving the `.ragsh` skeleton +/// — the **line-oriented command REPL** (verbs like `set`/`get`/`run`/`call` +/// parsed from a single line; see the [`repl`](crate::repl) module docs for +/// the grammar). +/// +/// # Not to be confused with `repl::session::ReplSession` +/// +/// There are two distinct types named `ReplSession` in this crate, gated +/// differently and serving different layers of the `.ragsh` design: +/// +/// - **This type** (`repl::ReplSession`, always available) — the +/// line-oriented command skeleton documented here. +/// - [`crate::repl::session::ReplSession`] (feature `repl` only) — the +/// Rhai-backed scripting session: a persistent namespace evaluated one cell +/// (small script) at a time, with capability calls (`model_query`, +/// `tool_call`, `graph_run`, …) wired to live registries. +/// +/// Only **this** type is re-exported as `repl::ReplSession`; the scripting +/// session is deliberately *not* re-exported there to avoid shadowing it, and +/// must be reached via `repl::session::ReplSession`. With the `repl` feature +/// enabled, `crate::ReplSession` (the crate-root re-export) resolves to the +/// **scripting** session instead — the crate root and the `repl` module +/// re-export different types under the same final path segment, so always +/// check which path (`crate::ReplSession` vs. `crate::repl::ReplSession`) you +/// actually imported from. /// /// ## Execution model /// diff --git a/tests/e2e_graph_resolver_contracts.rs b/tests/e2e_graph_resolver_contracts.rs index 4c41e77..45ce9a3 100644 --- a/tests/e2e_graph_resolver_contracts.rs +++ b/tests/e2e_graph_resolver_contracts.rs @@ -8,7 +8,7 @@ use tinyagents::graph::{Command, GraphBuilder, GraphDefaults, NodeResult, Route, use tinyagents::harness::ids::GraphId; use tinyagents::harness::providers::MockModel; use tinyagents::harness::testkit::FakeTool; -use tinyagents::language::compiler::CapabilityResolver; +use tinyagents::language::capability_resolver::CapabilityResolver; use tinyagents::language::resolver::{Resolver, resolve_source}; use tinyagents::language::{Blueprint, ChannelSpec, EdgeSpec, Literal, NodeSpec, Routing, parser}; use tinyagents::registry::{CapabilityRegistry, ComponentKind}; diff --git a/tests/e2e_graph_support_contracts.rs b/tests/e2e_graph_support_contracts.rs index 9599ef1..c2fa45d 100644 --- a/tests/e2e_graph_support_contracts.rs +++ b/tests/e2e_graph_support_contracts.rs @@ -116,6 +116,61 @@ async fn file_checkpointer_persists_lists_copies_prunes_and_deletes_threads() { let _ = std::fs::remove_dir_all(&dir); } +#[tokio::test] +async fn file_checkpointer_state_history_walks_lineage_newest_first() { + // Exercises FileCheckpointer's single-pass `state_history` override: the + // whole thread file is read once and the parent lineage walked in memory, + // returning snapshots newest-first with the lineage spine intact. + let dir = temp_path("file-checkpointer-history"); + let checkpointer = std::sync::Arc::new(FileCheckpointer::::new(&dir)); + 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() + .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); + + let _ = std::fs::remove_dir_all(&dir); +} + #[tokio::test] async fn graph_testkit_helpers_drive_recorded_graphs_and_assertions() { let checkpointer = std::sync::Arc::new(InMemoryCheckpointer::::new()); diff --git a/tests/e2e_harness_provider_contracts.rs b/tests/e2e_harness_provider_contracts.rs index 082a8dc..4f402e3 100644 --- a/tests/e2e_harness_provider_contracts.rs +++ b/tests/e2e_harness_provider_contracts.rs @@ -362,14 +362,18 @@ fn model_profiles_stream_chunks_usage_and_cost_are_additive() { output_reasoning_per_token: Some(0.003), ..ModelPricing::default() }; + // cache_read_tokens (3) and reasoning_tokens (4) are subsets of + // input_tokens/output_tokens, not additions to them, so the standard + // rate only applies to the non-cached/non-reasoning remainder: + // (10-3)*0.001 = 0.007, (5-4)*0.002 = 0.002. let cost = estimate_cost(&pricing, &usage); - assert_eq!(cost.input_cost, 0.01); - assert_eq!(cost.output_cost, 0.01); + assert_eq!(cost.input_cost, 0.007); + assert_eq!(cost.output_cost, 0.002); assert_eq!(cost.cache_cost, 0.0013); assert_eq!(cost.reasoning_cost, 0.012); - assert!((cost.total_cost - 0.0333).abs() < f64::EPSILON); + assert!((cost.total_cost - 0.0223).abs() < f64::EPSILON); let combined = CostTotals::new() + cost + cost; - assert!((combined.total_cost - 0.0666).abs() < f64::EPSILON); + assert!((combined.total_cost - 0.0446).abs() < f64::EPSILON); } #[test] diff --git a/tests/e2e_language_contracts.rs b/tests/e2e_language_contracts.rs index 980133b..8b3d62c 100644 --- a/tests/e2e_language_contracts.rs +++ b/tests/e2e_language_contracts.rs @@ -57,7 +57,6 @@ graph review_flow_v2 { } node work { kind tool_executor - next END } node audit { kind model diff --git a/tests/e2e_middleware_parser_contracts.rs b/tests/e2e_middleware_parser_contracts.rs index 1055dcc..d9bb41f 100644 --- a/tests/e2e_middleware_parser_contracts.rs +++ b/tests/e2e_middleware_parser_contracts.rs @@ -156,6 +156,7 @@ async fn middleware_stack_runs_lifecycle_hooks_and_builtin_guards() { let mut delta = ModelDelta { call_id: "model-1".into(), content: "piece".into(), + reasoning: String::new(), tool_call: None, }; stack diff --git a/tests/e2e_misc_public_helpers.rs b/tests/e2e_misc_public_helpers.rs index cef4c4c..f0dccda 100644 --- a/tests/e2e_misc_public_helpers.rs +++ b/tests/e2e_misc_public_helpers.rs @@ -175,6 +175,8 @@ async fn graph_reducers_streams_observability_and_status_helpers_work() { journal_sink.emit(GraphEvent::CheckpointSaved { checkpoint_id: CheckpointId::new("cp-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); @@ -328,10 +330,8 @@ fn tool_schema_limits_ids_and_repl_contracts_cover_public_helpers() { .with_max_tool_calls(1) .with_max_wall_clock_ms(Some(1)) .with_max_retries_per_call(3) - .with_max_concurrency(Some(2)) .with_max_depth(9); assert_eq!(limits.max_retries_per_call, 3); - assert_eq!(limits.max_concurrency, Some(2)); assert_eq!(limits.max_depth, 9); let mut tracker = LimitTracker::new(limits); assert_eq!(tracker.model_calls(), 0); diff --git a/tests/e2e_observability.rs b/tests/e2e_observability.rs index b822dfd..678e3e7 100644 --- a/tests/e2e_observability.rs +++ b/tests/e2e_observability.rs @@ -96,6 +96,9 @@ async fn harness_run_journals_redacted_replayable_events() { // RedactingSink wraps the JournalSink: every event is masked *before* it is // persisted, so the durable journal never sees the secret. let journal_sink = JournalSink::new(journal.clone(), run_id.clone()); + // Persistence is asynchronous; keep a handle to flush the shared drain + // before reading the journal back. Clones share the same background worker. + let journal_flush = journal_sink.clone(); let redacting = RedactingSink::new(Arc::new(journal_sink), vec![SECRET.to_string()]); // Attach the sinks through RunContext.events, then drive the run in-context. @@ -128,6 +131,7 @@ async fn harness_run_journals_redacted_replayable_events() { assert!(stored.error.is_none()); // --- Journal replays the whole run from offset 0. --- + journal_flush.flush(); let all = journal .read_from(run_id.as_str(), 0) .await @@ -209,13 +213,16 @@ async fn harness_run_with_capture_exports_generation_and_tool_io() { let journal: Arc = Arc::new(InMemoryEventJournal::new()); let run_id = RunId::new("run-capture"); let ctx: RunContext<()> = RunContext::new(RunConfig::new(run_id.as_str()), ()); - ctx.events - .subscribe(Arc::new(JournalSink::new(journal.clone(), run_id.clone()))); + let journal_sink = JournalSink::new(journal.clone(), run_id.clone()); + let journal_flush = journal_sink.clone(); + ctx.events.subscribe(Arc::new(journal_sink)); harness .invoke_in_context_with_status(&(), ctx, vec![Message::user("please look up")]) .await .expect("run succeeds"); + // Persistence is asynchronous; block until the durable log catches up. + journal_flush.flush(); let observations = journal .read_from(run_id.as_str(), 0) @@ -263,9 +270,11 @@ async fn harness_run_with_capture_exports_generation_and_tool_io() { "all done" ); + // Tool observations are exported as `span-create` (a valid Langfuse + // ingestion type); `tool-create` is rejected by older/self-hosted Langfuse. let tool = events .iter() - .find(|e| e["type"] == "tool-create") + .find(|e| e["type"] == "span-create") .expect("a tool observation"); assert_eq!(tool["body"]["input"]["q"], "weather"); assert_eq!(tool["body"]["output"], "tool-output"); diff --git a/tests/e2e_rag_pipeline.rs b/tests/e2e_rag_pipeline.rs index dfdbe3d..85fe0f0 100644 --- a/tests/e2e_rag_pipeline.rs +++ b/tests/e2e_rag_pipeline.rs @@ -11,9 +11,8 @@ use std::sync::Arc; use tinyagents::graph::{END, NodeFuture}; -use tinyagents::language::compiler::{ - BoxedNode, CapabilityResolver, NodeFactory, bind_capabilities, build_graph, compile, -}; +use tinyagents::language::capability_resolver::{CapabilityResolver, bind_capabilities}; +use tinyagents::language::compiler::{BoxedNode, NodeFactory, build_graph, compile}; use tinyagents::language::parser::parse_str; use tinyagents::language::types::{END as LANG_END, NodeSpec, Routing}; use tinyagents::{Command, NodeContext, NodeResult, Result}; diff --git a/tests/e2e_registry_observability_contracts.rs b/tests/e2e_registry_observability_contracts.rs index d46731a..5c32ec3 100644 --- a/tests/e2e_registry_observability_contracts.rs +++ b/tests/e2e_registry_observability_contracts.rs @@ -139,8 +139,9 @@ fn component_metadata_and_event_kinds_are_stable_serializable_contracts() { let id = ComponentId::new("researcher"); assert_eq!(id.as_str(), "researcher"); assert_eq!(id.to_string(), "researcher"); - assert_eq!(ComponentKind::ALL.len(), 11); + assert_eq!(ComponentKind::ALL.len(), 12); assert_eq!(ComponentKind::Agent.as_str(), "agent"); + assert_eq!(ComponentKind::Script.as_str(), "script"); assert_eq!(ComponentKind::TaskStore.as_str(), "task_store"); assert_eq!(ComponentKind::Tool.to_string(), "tool"); @@ -434,6 +435,7 @@ async fn fanout_redaction_journal_and_jsonl_sinks_forward_best_effort_events() { let journal_sink = JournalSink::new(journal.clone(), RunId::new("child-run")) .with_lineage(Some(RunId::new("parent-run")), RunId::new("root-run")); journal_sink.on_event(&secret_record); + journal_sink.flush(); let observations = journal.read_from("child-run", 0).await.unwrap(); assert_eq!(observations.len(), 1); assert_eq!( @@ -452,6 +454,7 @@ async fn fanout_redaction_journal_and_jsonl_sinks_forward_best_effort_events() { let jsonl_store = JsonlAppendStore::new(root.clone()); let jsonl_sink = JsonlSink::new(jsonl_store.clone(), "events"); jsonl_sink.on_event(&secret_record); + jsonl_sink.flush(); let rows = jsonl_store.read_from("events", 0).await.unwrap(); assert_eq!(rows.len(), 1); assert_eq!(rows[0].0, 0); diff --git a/tests/language_pipeline.rs b/tests/language_pipeline.rs index 6edcc2b..c4b6358 100644 --- a/tests/language_pipeline.rs +++ b/tests/language_pipeline.rs @@ -8,9 +8,8 @@ use std::sync::Arc; use tinyagents::graph::{END, NodeFuture}; -use tinyagents::language::compiler::{ - BoxedNode, CapabilityResolver, NodeFactory, bind_capabilities, build_graph, compile, -}; +use tinyagents::language::capability_resolver::{CapabilityResolver, bind_capabilities}; +use tinyagents::language::compiler::{BoxedNode, NodeFactory, build_graph, compile}; use tinyagents::language::parser::parse_str; use tinyagents::language::types::{END as LANG_END, NodeSpec, Routing}; use tinyagents::{Command, NodeContext, NodeResult, Result, TinyAgentsError}; diff --git a/wiki b/wiki index 9321604..b02f582 160000 --- a/wiki +++ b/wiki @@ -1 +1 @@ -Subproject commit 9321604d35fa34b53dc5f0383f70ae1b7410ac15 +Subproject commit b02f58269a4fc8742571f92e191ec893e9890eb0