From 8e65efd82b53e700cf57cd90da19eedf38de338e Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Fri, 3 Jul 2026 12:20:56 -0700 Subject: [PATCH 001/106] fix(observability): RedactingSink fails closed on redaction failure Treat the RedactingSink as a security boundary: if the event cannot be serialized, redacted, and rebuilt, drop it rather than forward a record that may still contain unredacted secrets. Add a fast path that forwards the original record unchanged when no secrets are configured. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/harness/observability/mod.rs | 13 +++++++++---- src/harness/observability/test.rs | 28 ++++++++++++++++++++++++++++ 2 files changed, 37 insertions(+), 4 deletions(-) diff --git a/src/harness/observability/mod.rs b/src/harness/observability/mod.rs index d789558..0cb82ec 100644 --- a/src/harness/observability/mod.rs +++ b/src/harness/observability/mod.rs @@ -384,16 +384,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 { diff --git a/src/harness/observability/test.rs b/src/harness/observability/test.rs index ae378cf..6ac288b 100644 --- a/src/harness/observability/test.rs +++ b/src/harness/observability/test.rs @@ -389,6 +389,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()); From 9b1d88540c49a1814cd011439ef77a64bbe4fc70 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Fri, 3 Jul 2026 12:22:09 -0700 Subject: [PATCH 002/106] fix(tool): enforce schema required fields without properties The `required` check was nested inside the `properties` branch, so a schema declaring `required` without listing `properties` skipped validation and accepted calls that omitted mandatory arguments. Lift the `required` enforcement out so it fails closed independently of `properties`. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/harness/tool/mod.rs | 25 ++++++++++++++++--------- src/harness/tool/test.rs | 28 ++++++++++++++++++++++++++++ 2 files changed, 44 insertions(+), 9 deletions(-) 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( From 2112eb3c5f62087db971cf2dccaf5618da31a291 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Fri, 3 Jul 2026 12:24:35 -0700 Subject: [PATCH 003/106] fix(workspace): anchor path gate to defeat leading `..` sibling escape The lexical path gate dropped leading `..` components, so a relative candidate like `ws/../../ws/secret` collapsed back onto the root name and was admitted even though it resolves into a same-named sibling outside the workspace. Anchor relative candidates and roots to the current directory and preserve escaping `..` during normalization so the gate fails closed. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/harness/workspace/test.rs | 15 ++++++++++++ src/harness/workspace/types.rs | 43 +++++++++++++++++++++++++++++----- 2 files changed, 52 insertions(+), 6 deletions(-) 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..8c70ce3 100644 --- a/src/harness/workspace/types.rs +++ b/src/harness/workspace/types.rs @@ -70,12 +70,18 @@ impl WorkspaceDescriptor { /// /// 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. + /// 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 candidate = normalize(path); + let Some(candidate) = anchored_normalize(path) else { + return false; + }; std::iter::once(&self.root) .chain(self.trusted_roots.iter()) - .any(|root| candidate.starts_with(normalize(root))) + .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 @@ -97,16 +103,41 @@ impl WorkspaceDescriptor { } } +/// 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 => { - out.pop(); - } + 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()), } From 2ceefe6a193d1c1bbb54ebe9ab54c0bcbafb4f83 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Fri, 3 Jul 2026 12:35:37 -0700 Subject: [PATCH 004/106] fix(language): fail closed on subagent/repl_agent references via one shared policy MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The strict blueprint gate (`CapabilityResolver::bind_blueprint`) routed every non-subgraph/router node through a model check, so `subagent` `agent` references were never validated and `repl_agent` `script` references were validated by no path at all — both admitted unregistered capabilities. Extract one shared kind-to-reference policy (`CapabilityResolver::classify_reference` + `reference_allowed`) and route all three binding gates through it — the compiler blueprint gate and both `Resolver` paths — so they cannot drift. Move the agent allowlist into `CapabilityResolver`, add a `scripts` allowlist and a `ComponentKind::Script` so scripts can be registered and bound. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- .../implementation-status.md | 7 +- docs/modules/registry/README.md | 7 +- src/language/compiler.rs | 175 ++++++++++++++---- src/language/resolver.rs | 174 +++++++---------- src/language/test.rs | 42 +++++ src/registry/component/mod.rs | 4 +- src/registry/component/types.rs | 2 + tests/e2e_registry_observability_contracts.rs | 3 +- 8 files changed, 261 insertions(+), 153 deletions(-) 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/registry/README.md b/docs/modules/registry/README.md index 090d496..01bf3db 100644 --- a/docs/modules/registry/README.md +++ b/docs/modules/registry/README.md @@ -112,9 +112,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/src/language/compiler.rs b/src/language/compiler.rs index f415100..fb9df3a 100644 --- a/src/language/compiler.rs +++ b/src/language/compiler.rs @@ -437,12 +437,61 @@ pub struct CapabilityResolver { 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 { @@ -485,6 +534,8 @@ impl CapabilityResolver { 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(), } } @@ -519,6 +570,19 @@ impl CapabilityResolver { 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 @@ -555,6 +619,59 @@ impl CapabilityResolver { 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 { @@ -569,15 +686,17 @@ impl CapabilityResolver { /// - 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; + /// `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, or reducer reference. + /// 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) { @@ -587,39 +706,23 @@ impl CapabilityResolver { ))); } - 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 - ))); - } - } + // 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 { diff --git a/src/language/resolver.rs b/src/language/resolver.rs index ba6ec19..af7cb05 100644 --- a/src/language/resolver.rs +++ b/src/language/resolver.rs @@ -31,11 +31,9 @@ //! 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::compiler::{CapabilityResolver, DEFAULT_NODE_KINDS, ReferenceClass, compile}; use crate::language::diagnostic::Diagnostic; use crate::language::parser::parse_str; use crate::language::source::SourceFile; @@ -49,6 +47,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 +61,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 +94,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 +145,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 +268,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 +324,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 index 403e385..dc9c3aa 100644 --- a/src/language/test.rs +++ b/src/language/test.rs @@ -630,6 +630,8 @@ 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") @@ -641,6 +643,46 @@ fn extended_kinds_bind_against_a_resolver() { 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::compiler::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}"); +} + // --------------------------------------------------------------------------- // Capability binding // --------------------------------------------------------------------------- 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/tests/e2e_registry_observability_contracts.rs b/tests/e2e_registry_observability_contracts.rs index d46731a..6335741 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"); From 13d8a5bb909878b810d157d4e6c2cacb87aca384 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Fri, 3 Jul 2026 12:37:01 -0700 Subject: [PATCH 005/106] fix(language): return parse error for empty token slice `parse(&[])` underflowed `tokens.len() - 1` in the token cursor and panicked. Reject an empty token slice up front with a parse error, since a well-formed stream always ends with an `Eof` sentinel. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/language/parser.rs | 12 +++++++++++- src/language/test.rs | 14 ++++++++++++++ 2 files changed, 25 insertions(+), 1 deletion(-) 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/test.rs b/src/language/test.rs index dc9c3aa..49bb364 100644 --- a/src/language/test.rs +++ b/src/language/test.rs @@ -231,6 +231,20 @@ fn parse_error_without_source_renders_plain() { } } +#[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:?}"), + } +} + // --------------------------------------------------------------------------- // Parser // --------------------------------------------------------------------------- From f0e42667875811fa321785101917e3b20d51ef52 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Fri, 3 Jul 2026 12:41:58 -0700 Subject: [PATCH 006/106] fix(language): clamp diagnostic caret range for out-of-bounds spans Diagnostic::render panicked when a label span pointed past the end of the source: `caret_start` could exceed `line_end`, tripping `clamp`'s `min <= max` precondition (and risking an out-of-bounds slice). Clamp the caret range into the line's byte range, mirroring `SourceFile::snippet`. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/language/diagnostic.rs | 9 +++++++-- src/language/test.rs | 16 ++++++++++++++++ 2 files changed, 23 insertions(+), 2 deletions(-) diff --git a/src/language/diagnostic.rs b/src/language/diagnostic.rs index 8271309..455e8b9 100644 --- a/src/language/diagnostic.rs +++ b/src/language/diagnostic.rs @@ -241,8 +241,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/test.rs b/src/language/test.rs index 49bb364..73c8767 100644 --- a/src/language/test.rs +++ b/src/language/test.rs @@ -189,6 +189,22 @@ fn diagnostic_renders_caret_under_primary_span() { ); } +#[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 severity_labels_are_lowercase() { assert_eq!(Severity::Error.label(), "error"); From 26e7153a8ff8eee2bdeb1d49cbee67c94a4b3f13 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Fri, 3 Jul 2026 12:53:00 -0700 Subject: [PATCH 007/106] fix(openai): decode SSE across UTF-8 and EOF boundaries Buffer the raw SSE byte stream as `Vec` and only lossily decode complete newline-terminated lines, so a multi-byte UTF-8 character split across two network chunks is reassembled before decoding instead of being corrupted into replacement characters. Drain any leftover buffered bytes at end-of-stream so a final `data:` event sent without a trailing newline (and without a `[DONE]` sentinel) is still processed. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/harness/providers/openai/mod.rs | 82 +++++++++++++++++++--------- src/harness/providers/openai/test.rs | 66 +++++++++++++++++++++- 2 files changed, 119 insertions(+), 29 deletions(-) diff --git a/src/harness/providers/openai/mod.rs b/src/harness/providers/openai/mod.rs index c101f18..f7f6a52 100644 --- a/src/harness/providers/openai/mod.rs +++ b/src/harness/providers/openai/mod.rs @@ -848,8 +848,11 @@ impl OpenAiStreamAcc { 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, + /// 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. + buf: Vec, /// Parsed items waiting to be yielded, in order. pending: VecDeque, /// Provider-side response reconstruction. @@ -868,30 +871,52 @@ struct SseState { } 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. + /// 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) { - 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; - } + while let Some(pos) = self.buf.iter().position(|&b| b == b'\n') { + let line_bytes: Vec = self.buf.drain(..=pos).collect(); + // A complete line (bounded by the ASCII `\n`) is whole UTF-8, so a + // lossy decode here can no longer straddle a chunk boundary. + let line = String::from_utf8_lossy(&line_bytes).into_owned(); + self.process_line(&line); + } + } + + /// 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. + 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; } } } @@ -928,7 +953,7 @@ async fn sse_next(mut state: SseState) -> Option<(ModelStreamItem, SseState)> { } match state.bytes.next().await { Some(Ok(chunk)) => { - state.buf.push_str(&String::from_utf8_lossy(&chunk)); + state.buf.extend_from_slice(&chunk); state.drain_lines(); } Some(Err(error)) => { @@ -944,6 +969,9 @@ async fn sse_next(mut state: SseState) -> Option<(ModelStreamItem, SseState)> { 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; } } @@ -1059,7 +1087,7 @@ impl ChatModel for OpenAiModel { let state = SseState { bytes: Box::pin(bytes), - buf: String::new(), + buf: Vec::new(), pending: VecDeque::new(), acc: OpenAiStreamAcc::default(), provider: self.provider.clone(), diff --git a/src/harness/providers/openai/test.rs b/src/harness/providers/openai/test.rs index fe6a495..8c9694d 100644 --- a/src/harness/providers/openai/test.rs +++ b/src/harness/providers/openai/test.rs @@ -443,7 +443,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(), @@ -512,7 +512,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 +545,68 @@ 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_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"); +} + #[test] fn from_env_errors_when_api_key_missing() { // Snapshot and clear the key so the missing-key path is exercised From 94eac5ec6b39b5d8c3824e3794bc9dc13e5fb34a Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Fri, 3 Jul 2026 12:53:56 -0700 Subject: [PATCH 008/106] fix(openai): surface mid-stream error payloads instead of swallowing them A mid-stream `{"error": ...}` SSE payload deserializes cleanly as an all-defaults `ChatCompletionChunk`, so it was folded in as an empty chunk and silently dropped. Parse each `data:` payload as a generic value first and, when it carries an `error` object, emit a terminal `ProviderFailed` (with the provider message and code) and stop the stream. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/harness/providers/openai/mod.rs | 40 +++++++++++++++++++++++++++- src/harness/providers/openai/test.rs | 39 +++++++++++++++++++++++++++ 2 files changed, 78 insertions(+), 1 deletion(-) diff --git a/src/harness/providers/openai/mod.rs b/src/harness/providers/openai/mod.rs index f7f6a52..790a51c 100644 --- a/src/harness/providers/openai/mod.rs +++ b/src/harness/providers/openai/mod.rs @@ -913,12 +913,50 @@ impl SseState { return; } // Ignore keepalives / unparseable lines rather than failing the run. - if let Ok(chunk) = serde_json::from_str::(payload) { + 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`]. diff --git a/src/harness/providers/openai/test.rs b/src/harness/providers/openai/test.rs index 8c9694d..abc12c5 100644 --- a/src/harness/providers/openai/test.rs +++ b/src/harness/providers/openai/test.rs @@ -607,6 +607,45 @@ async fn sse_stream_drains_final_line_without_trailing_newline() { 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")); +} + #[test] fn from_env_errors_when_api_key_missing() { // Snapshot and clear the key so the missing-key path is exercised From 47502bb653a87136eff513344122c393119d6403 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Fri, 3 Jul 2026 12:56:01 -0700 Subject: [PATCH 009/106] fix(stream): adjacently tag ModelStreamItem and StreamChunk so every variant round-trips MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Both enums were internally tagged (`tag = "type"`). Internal tagging cannot serialize newtype variants wrapping a non-map payload: `ModelStreamItem::Failed` (a `String`) errored, and `StreamChunk::Values(json!(null))` and other scalar / array `Value` payloads corrupted (`null` became `{}`). Switch both to adjacent tagging (`{"type": …, "content": …}`), which keeps the ergonomic tuple variants and all existing call sites unchanged while making every variant round-trip. Add serde round-trip tests covering every variant of both enums. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/harness/model/test.rs | 42 +++++++++++++++++++++++++++++++++++++ src/harness/model/types.rs | 9 +++++++- src/harness/stream/test.rs | 29 +++++++++++++++++++++++++ src/harness/stream/types.rs | 10 ++++++++- 4 files changed, 88 insertions(+), 2 deletions(-) diff --git a/src/harness/model/test.rs b/src/harness/model/test.rs index 8116938..e804cd1 100644 --- a/src/harness/model/test.rs +++ b/src/harness/model/test.rs @@ -610,3 +610,45 @@ fn stream_accumulator_collects_reasoning_side_channel() { let response = acc.finish().unwrap(); assert_eq!(response.text(), "visible answer"); } + +/// 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..b6b7930 100644 --- a/src/harness/model/types.rs +++ b/src/harness/model/types.rs @@ -480,8 +480,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/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), From 3e325e6e83bf98f714f91e7b164ae9643559e14d Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Fri, 3 Jul 2026 13:21:59 -0700 Subject: [PATCH 010/106] fix(openai): correlate index-less streamed tool-call fragments by id Some OpenAI-compatible backends omit the tool-call `index` field, which collapsed every parallel call onto slot 0. Make `index` optional and, when absent, correlate fragments by id (new id opens a slot, known id reuses it, id-less continuations append to the last slot). Also enumerate slots before filtering in `into_response` and key the synthetic `tool-{slot}` fallback id to the slot position so streamed delta ids always match the final call ids. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/harness/providers/openai/mod.rs | 64 ++++++++++++++++----- src/harness/providers/openai/test.rs | 80 +++++++++++++++++++++++++++ src/harness/providers/openai/types.rs | 6 +- 3 files changed, 134 insertions(+), 16 deletions(-) diff --git a/src/harness/providers/openai/mod.rs b/src/harness/providers/openai/mod.rs index 790a51c..afedc61 100644 --- a/src/harness/providers/openai/mod.rs +++ b/src/harness/providers/openai/mod.rs @@ -696,6 +696,17 @@ fn parse_response(value: Value) -> Result { }) } +/// 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. +fn tool_call_id(slot: usize, id: &str) -> String { + if id.is_empty() { + format!("tool-{slot}") + } else { + id.to_string() + } +} + 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!( @@ -774,10 +785,7 @@ impl OpenAiStreamAcc { })); } 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 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; @@ -788,11 +796,7 @@ impl OpenAiStreamAcc { } 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() - }; + let call_id = tool_call_id(idx, &slot.id); pending.push_back(ModelStreamItem::ToolCallDelta(ToolDelta { call_id, content: args, @@ -803,23 +807,53 @@ impl OpenAiStreamAcc { } } + /// 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() - .filter(|b| !b.name.is_empty() || !b.args.is_empty()) .enumerate() + .filter(|(_, b)| !b.name.is_empty() || !b.args.is_empty()) .map(|(idx, b)| { - let id = if b.id.is_empty() { - format!("tool-{idx}") - } else { - b.id.clone() - }; + let id = tool_call_id(idx, &b.id); Ok(ToolCall { id: id.clone(), name: b.name.clone(), diff --git a/src/harness/providers/openai/test.rs b/src/harness/providers/openai/test.rs index abc12c5..eae31fe 100644 --- a/src/harness/providers/openai/test.rs +++ b/src/harness/providers/openai/test.rs @@ -646,6 +646,86 @@ async fn sse_stream_surfaces_mid_stream_error_payload() { 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 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/types.rs b/src/harness/providers/openai/types.rs index 5aa5a50..f1c1d84 100644 --- a/src/harness/providers/openai/types.rs +++ b/src/harness/providers/openai/types.rs @@ -106,8 +106,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, From 50312b12f91439698d879c26fade61292026377d Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Fri, 3 Jul 2026 13:22:49 -0700 Subject: [PATCH 011/106] fix(openai): map empty tool arguments to an empty object Zero-argument tool calls from OpenAI-compatible backends arrive with an empty arguments string. Treat a blank arguments payload as `{}` in both the unary and streamed reconstruction paths instead of failing it as malformed JSON. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/harness/providers/openai/mod.rs | 6 ++++ src/harness/providers/openai/test.rs | 51 ++++++++++++++++++++++++++++ 2 files changed, 57 insertions(+) diff --git a/src/harness/providers/openai/mod.rs b/src/harness/providers/openai/mod.rs index afedc61..1a8d29a 100644 --- a/src/harness/providers/openai/mod.rs +++ b/src/harness/providers/openai/mod.rs @@ -708,6 +708,12 @@ fn tool_call_id(slot: usize, id: &str) -> String { } 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:?}" diff --git a/src/harness/providers/openai/test.rs b/src/harness/providers/openai/test.rs index eae31fe..88a2d0e 100644 --- a/src/harness/providers/openai/test.rs +++ b/src/harness/providers/openai/test.rs @@ -344,6 +344,57 @@ fn parses_text_only_response_without_usage_details() { assert_eq!(usage.cache_read_tokens, 0); } +#[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": [] }); From c1918b25f4d4d910351ec77f9698f5d2c6ea69d1 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Fri, 3 Jul 2026 13:24:11 -0700 Subject: [PATCH 012/106] fix(openai): wire request and client timeouts The provider never set any reqwest timeout and ignored ModelRequest::timeout_ms, so a stalled endpoint could hang a call forever. Give the shared client a sane connect timeout, honor an explicit per-request timeout_ms on both unary and streaming calls, and apply a default overall timeout to unary calls (streaming stays uncapped so long streams are not truncated). Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/harness/providers/openai/mod.rs | 73 +++++++++++++++++++--------- src/harness/providers/openai/test.rs | 24 +++++++++ 2 files changed, 74 insertions(+), 23 deletions(-) diff --git a/src/harness/providers/openai/mod.rs b/src/harness/providers/openai/mod.rs index 1a8d29a..41af48b 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,6 +57,13 @@ 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"; +/// 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; /// A [`ChatModel`] backed by the hosted OpenAI Chat Completions API. /// @@ -114,7 +122,10 @@ impl OpenAiModel { /// (`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(), + 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(), @@ -507,6 +518,20 @@ impl OpenAiModel { } } +/// 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). +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 @@ -1075,18 +1100,19 @@ impl ChatModel for OpenAiModel { let body = self.translate_request(&request)?; let url = format!("{}/chat/completions", self.base_url); - let response = self + let mut builder = 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)) - })?; + .json(&body); + if let Some(timeout) = request_timeout(request.timeout_ms, false) { + builder = builder.timeout(timeout); + } + let response = builder.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| { @@ -1129,22 +1155,23 @@ impl ChatModel for OpenAiModel { body.stream_options = Some(json!({ "include_usage": true })); let url = format!("{}/chat/completions", self.base_url); - let response = self + let mut builder = 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)) - })?; + .json(&body); + if let Some(timeout) = request_timeout(request.timeout_ms, true) { + builder = builder.timeout(timeout); + } + let response = builder.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() { diff --git a/src/harness/providers/openai/test.rs b/src/harness/providers/openai/test.rs index 88a2d0e..c468d70 100644 --- a/src/harness/providers/openai/test.rs +++ b/src/harness/providers/openai/test.rs @@ -777,6 +777,30 @@ async fn sse_stream_indexless_fallback_ids_match_between_delta_and_final() { assert_eq!(calls[0].arguments, json!({ "n": 7 })); } +#[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 From 23d2fe940a2b55f01e51da4d8844b3cff3236779 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Fri, 3 Jul 2026 13:25:53 -0700 Subject: [PATCH 013/106] fix(openai): send max_completion_tokens to o-series models o-series reasoning models (o1/o3/o4) reject `max_tokens` and require `max_completion_tokens`. Route the request's output-token cap to whichever field the target model accepts, add the new wire field, and reserve it from provider_options passthrough. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/harness/providers/openai/mod.rs | 32 ++++++++++++++++++++------- src/harness/providers/openai/test.rs | 21 ++++++++++++++++++ src/harness/providers/openai/types.rs | 8 ++++++- 3 files changed, 52 insertions(+), 9 deletions(-) diff --git a/src/harness/providers/openai/mod.rs b/src/harness/providers/openai/mod.rs index 41af48b..bff5f82 100644 --- a/src/harness/providers/openai/mod.rs +++ b/src/harness/providers/openai/mod.rs @@ -85,6 +85,13 @@ pub struct OpenAiModel { profile: ModelProfile, } +/// Returns `true` for OpenAI o-series reasoning models (`o1`/`o3`/`o4`), which +/// reject `max_tokens` and require `max_completion_tokens` instead. +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), @@ -92,12 +99,8 @@ pub struct OpenAiModel { /// 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"); + 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()), @@ -443,15 +446,27 @@ impl OpenAiModel { .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: request.model.clone().unwrap_or_else(|| self.model.clone()), + model, messages, tools, tool_choice, response_format, temperature: request.temperature, top_p: request.top_p, - max_tokens: request.max_tokens, + max_tokens, + max_completion_tokens, stop: request.stop_sequences.clone(), seed: request.seed, stream: false, @@ -556,6 +571,7 @@ fn provider_extra_options(options: &Value) -> Result> { "temperature", "top_p", "max_tokens", + "max_completion_tokens", "stop", "seed", "stream", diff --git a/src/harness/providers/openai/test.rs b/src/harness/providers/openai/test.rs index c468d70..a7ad242 100644 --- a/src/harness/providers/openai/test.rs +++ b/src/harness/providers/openai/test.rs @@ -777,6 +777,27 @@ async fn sse_stream_indexless_fallback_ids_match_between_delta_and_final() { assert_eq!(calls[0].arguments, json!({ "n": 7 })); } +#[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. diff --git a/src/harness/providers/openai/types.rs b/src/harness/providers/openai/types.rs index f1c1d84..02a7ee7 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, From f8605f86832338c95ed3d2adde84e1168e0118bb Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Fri, 3 Jul 2026 13:27:25 -0700 Subject: [PATCH 014/106] fix(openai): send image content parts instead of dropping them translate_message flattened every message to its text, so image and JSON blocks were silently dropped even though the profile advertises image_in. Render user messages with non-text blocks as OpenAI content-parts (text + image_url), serialize JSON blocks into text, and fail closed with a validation error on provider-extension blocks that have no OpenAI representation. Text-only messages keep their plain-string wire shape. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/harness/providers/openai/mod.rs | 75 +++++++++++++++++++++++++-- src/harness/providers/openai/test.rs | 49 +++++++++++++++++ src/harness/providers/openai/types.rs | 43 +++++++++++++-- 3 files changed, 159 insertions(+), 8 deletions(-) diff --git a/src/harness/providers/openai/mod.rs b/src/harness/providers/openai/mod.rs index bff5f82..96327c5 100644 --- a/src/harness/providers/openai/mod.rs +++ b/src/harness/providers/openai/mod.rs @@ -586,17 +586,22 @@ fn provider_extra_options(options: &Value) -> Result> { } /// 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. fn translate_message(message: &Message) -> Result { let wire = match message { Message::System(_) => ChatMessageWire { role: "system".to_string(), - content: Some(message.text()), + content: Some(MessageContentWire::Text(message.text())), tool_calls: Vec::new(), tool_call_id: None, }, - Message::User(_) => ChatMessageWire { + Message::User(user) => ChatMessageWire { role: "user".to_string(), - content: Some(message.text()), + content: Some(translate_user_content(&user.content)?), tool_calls: Vec::new(), tool_call_id: None, }, @@ -606,7 +611,7 @@ fn translate_message(message: &Message) -> Result { let content = if text.is_empty() && !assistant.tool_calls.is_empty() { None } else { - Some(text) + Some(MessageContentWire::Text(text)) }; let tool_calls = assistant .tool_calls @@ -632,7 +637,7 @@ fn translate_message(message: &Message) -> Result { } Message::Tool(tool) => ChatMessageWire { role: "tool".to_string(), - content: Some(message.text()), + content: Some(MessageContentWire::Text(message.text())), tool_calls: Vec::new(), tool_call_id: Some(tool.tool_call_id.clone()), }, @@ -640,6 +645,66 @@ fn translate_message(message: &Message) -> Result { 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. +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. +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. fn translate_tool_choice(choice: &ToolChoice) -> Value { match choice { diff --git a/src/harness/providers/openai/test.rs b/src/harness/providers/openai/test.rs index a7ad242..92b5cdd 100644 --- a/src/harness/providers/openai/test.rs +++ b/src/harness/providers/openai/test.rs @@ -777,6 +777,55 @@ async fn sse_stream_indexless_fallback_ids_match_between_delta_and_final() { 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 diff --git a/src/harness/providers/openai/types.rs b/src/harness/providers/openai/types.rs index 02a7ee7..ee27f05 100644 --- a/src/harness/providers/openai/types.rs +++ b/src/harness/providers/openai/types.rs @@ -137,15 +137,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, From eb3c032f6320807a14d854a5bb5b01df09d08322 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Fri, 3 Jul 2026 13:28:29 -0700 Subject: [PATCH 015/106] fix(model): stop StreamAccumulator::finish clobbering message usage MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit When a completed response carried usage only on its message (top-level usage None) and no UsageDelta streamed, finish overwrote message.usage with None. Merge the response, message, and streamed usage instead — prefer a present value, never overwrite a known usage with None — and keep both fields in sync. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/harness/model/mod.rs | 11 +++++++---- src/harness/model/test.rs | 36 ++++++++++++++++++++++++++++++++++++ 2 files changed, 43 insertions(+), 4 deletions(-) diff --git a/src/harness/model/mod.rs b/src/harness/model/mod.rs index f82cca6..fed88b4 100644 --- a/src/harness/model/mod.rs +++ b/src/harness/model/mod.rs @@ -673,10 +673,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 e804cd1..7169882 100644 --- a/src/harness/model/test.rs +++ b/src/harness/model/test.rs @@ -611,6 +611,42 @@ fn stream_accumulator_collects_reasoning_side_channel() { 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) { From 22addf2051af7cd6bdff13f563d95bb757991177 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Fri, 3 Jul 2026 13:38:55 -0700 Subject: [PATCH 016/106] fix(graph): collision-free checkpoint/run ids across restarts Checkpoint and run ids were built from a process-local AtomicU64 that restarts at 0 in every new process, so a resumed thread restarted after a crash re-minted ids it had already used. That corrupts the parent-lineage map (prune could delete a live record's ancestor) and can make ResumeTarget::Checkpoint load the wrong record. Add a process nonce to harness/ids and expose new_run_id/new_checkpoint_id (`--`), collision-free across restarts while next_seq keeps them ordered within a process. The graph executor now mints ids through those helpers instead of its own bare counter. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/graph/compiled/mod.rs | 24 ++++++++++++++++------- src/harness/ids/mod.rs | 40 +++++++++++++++++++++++++++++++++++++++ src/harness/ids/test.rs | 25 ++++++++++++++++++++++++ 3 files changed, 82 insertions(+), 7 deletions(-) diff --git a/src/graph/compiled/mod.rs b/src/graph/compiled/mod.rs index 01ecb92..ead97b3 100644 --- a/src/graph/compiled/mod.rs +++ b/src/graph/compiled/mod.rs @@ -99,8 +99,18 @@ 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 @@ -685,7 +695,7 @@ where }; let completed_tasks: Vec = as_node.iter().cloned().collect(); - let checkpoint_id = next_id("ckpt"); + 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(), @@ -751,7 +761,7 @@ where )) })?; let step = source.to_metadata().step; - let checkpoint_id = next_id("ckpt"); + 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(), @@ -778,7 +788,7 @@ where thread_id: Option, resume_map: HashMap, ) -> Result> { - let run_id = RunId::new(next_id("run")); + 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 @@ -1242,7 +1252,7 @@ where }; let checkpoint = Checkpoint { thread_id: thread.to_string(), - checkpoint_id: next_id("ckpt"), + checkpoint_id: next_checkpoint_id(), run_id: Some(run_id.to_string()), parent_checkpoint_id: parent, namespace: self.namespace.clone(), @@ -1735,7 +1745,7 @@ where let (Some(checkpointer), Some(thread)) = (&self.checkpointer, thread_id) else { return Ok(None); }; - let checkpoint_id = next_id("ckpt"); + let checkpoint_id = next_checkpoint_id(); let checkpoint = Checkpoint { thread_id: thread.to_string(), checkpoint_id, diff --git a/src/harness/ids/mod.rs b/src/harness/ids/mod.rs index fa2d13d..1b66b11 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,44 @@ 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) + }) +} + +/// 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..cdeae05 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!( From 9628befa6ff46e19b66a64c3c901d1b253c98544 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Fri, 3 Jul 2026 13:40:14 -0700 Subject: [PATCH 017/106] fix(graph): pin duplicate checkpoint-id lookup to last-write-wins MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The three checkpointer backends disagreed on `get(thread, Some(id))` when an id appeared more than once: the in-memory backend returned the first match while file and sqlite returned the last. Pin one semantic — last-write-wins, matching the append-only history and `get(None)` — fix the in-memory backend to `rfind`, and assert it in the checkpointer conformance suite so all backends stay interchangeable. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/graph/checkpoint/mod.rs | 6 +++++- src/graph/testkit/conformance.rs | 19 +++++++++++++++++++ 2 files changed, 24 insertions(+), 1 deletion(-) diff --git a/src/graph/checkpoint/mod.rs b/src/graph/checkpoint/mod.rs index 6841ea5..0c073c0 100644 --- a/src/graph/checkpoint/mod.rs +++ b/src/graph/checkpoint/mod.rs @@ -284,8 +284,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/testkit/conformance.rs b/src/graph/testkit/conformance.rs index bb22a7b..0e8dfa4 100644 --- a/src/graph/testkit/conformance.rs +++ b/src/graph/testkit/conformance.rs @@ -65,6 +65,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!( From ec89cedc4c82e3caa22aed72389d209b44ae38f1 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Fri, 3 Jul 2026 13:43:20 -0700 Subject: [PATCH 018/106] fix(graph): key command routing per activation, not per node MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The step's goto map was keyed by node id, so two Send activations of the same node that each returned a Command::goto collided: the second clobbered the first, and both branches then routed to the survivor's target — losing one branch's routing and duplicating the other. Key the map by active-set index instead, and pass each activation's own goto targets into `route`. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/graph/compiled/mod.rs | 34 ++++++++++++-------- src/graph/compiled/test.rs | 63 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 85 insertions(+), 12 deletions(-) diff --git a/src/graph/compiled/mod.rs b/src/graph/compiled/mod.rs index ead97b3..6f991fb 100644 --- a/src/graph/compiled/mod.rs +++ b/src/graph/compiled/mod.rs @@ -143,8 +143,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 @@ -686,7 +691,7 @@ where // 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)? + .route(node, None, &new_state)? .into_iter() .map(|t| t.node().clone()) .filter(|n| n.as_str() != END) @@ -1090,9 +1095,10 @@ where let completed = active.clone(); let mut next: Vec = Vec::new(); let mut next_seen: HashSet = HashSet::new(); - for activation in &completed { + for (index, activation) in completed.iter().enumerate() { let node_id = &activation.node; - let targets = self.route(node_id, &goto_map, &state)?; + 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 { @@ -1387,7 +1393,7 @@ where step: usize, result: NodeResult, updates: &mut Vec, - goto_map: &mut HashMap>, + goto_map: &mut HashMap>, visited: &mut Vec, ) -> Option<(usize, Interrupt)> { visited.push(node_id.clone()); @@ -1408,7 +1414,7 @@ where }); } if !command.goto.is_empty() { - goto_map.insert(node_id.clone(), command.goto); + goto_map.insert(index, command.goto); } } NodeResult::Interrupt(emitted) => { @@ -1444,7 +1450,7 @@ where child_runs: &ChildRunSink, ) -> Result> { let mut updates: Vec = Vec::new(); - let mut goto_map: HashMap> = HashMap::new(); + let mut goto_map: HashMap> = HashMap::new(); let mut interrupt: Option<(usize, Interrupt)> = None; let mut failure: Option = None; @@ -1615,7 +1621,7 @@ where // Fold in deterministic active-set index order. let mut updates: Vec = Vec::new(); - let mut goto_map: HashMap> = HashMap::new(); + let mut goto_map: HashMap> = HashMap::new(); let mut interrupt: Option<(usize, Interrupt)> = None; let mut failure: Option = None; @@ -1668,15 +1674,19 @@ where /// /// 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. fn route( &self, node_id: &NodeId, - goto_map: &HashMap>, + goto: Option<&[RouteTarget]>, state: &State, ) -> Result> { - if let Some(targets) = goto_map.get(node_id) { + if let Some(targets) = goto { self.validate_route_targets(node_id, targets)?; - return Ok(targets.clone()); + return Ok(targets.to_vec()); } if let Some(target) = self.edges.get(node_id) { return Ok(vec![RouteTarget::Node(target.clone())]); diff --git a/src/graph/compiled/test.rs b/src/graph/compiled/test.rs index 27a10c8..0886b0b 100644 --- a/src/graph/compiled/test.rs +++ b/src/graph/compiled/test.rs @@ -952,6 +952,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() From bd6238c6399cf5d70a15bb395094933ea2212b29 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Fri, 3 Jul 2026 14:00:15 -0700 Subject: [PATCH 019/106] fix(graph): persist pending activations, completed successors, and barriers MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Interrupt and failure boundaries lost durable state three ways: - B2: the checkpoint only stored next_nodes (Vec), so every pending Send worker re-ran on resume with ctx.send_arg == None — breaking the map-reduce fanout the resumable-failure feature exists for. - B3: successors of branches that completed before the pause were never scheduled, so in parallel [a->x, b] where b interrupts, x silently never ran after resume. - B4: barrier (waiting-edge) arrivals were run-local, so a join whose first predecessor arrived before an interrupt could never satisfy its precondition after resume. Add pending_activations (node + send_arg) and barrier_arrivals to the Checkpoint, both #[serde(default)] for back-compat (old checkpoints load and fall back to next_nodes / an empty barrier set). At the interrupt/failure boundary the executor now routes the completed lower-index branches into their successors, appends the pending tail with args preserved, and persists the accumulated barrier arrivals; resume rebuilds activations and seeds the barrier map from the checkpoint. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/graph/checkpoint/mod.rs | 4 +- src/graph/checkpoint/test.rs | 57 ++++++ src/graph/checkpoint/types.rs | 45 +++++ src/graph/compiled/mod.rs | 329 ++++++++++++++++++++++--------- src/graph/compiled/test.rs | 200 +++++++++++++++++++ src/graph/mod.rs | 5 +- src/graph/testkit/conformance.rs | 2 + src/lib.rs | 13 +- 8 files changed, 551 insertions(+), 104 deletions(-) diff --git a/src/graph/checkpoint/mod.rs b/src/graph/checkpoint/mod.rs index 0c073c0..93364ff 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}; 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/mod.rs b/src/graph/compiled/mod.rs index 6f991fb..c8b1cdd 100644 --- a/src/graph/compiled/mod.rs +++ b/src/graph/compiled/mod.rs @@ -77,7 +77,8 @@ 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::{ @@ -162,14 +163,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, } @@ -196,6 +195,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 { @@ -401,6 +436,7 @@ where vec![Activation::node(self.entry.clone())], None, HashMap::new(), + HashMap::new(), ) .await } @@ -419,7 +455,8 @@ where inputs: impl IntoIterator, ) -> Result> { let active = self.initial_inputs(inputs)?; - self.execute(state, active, None, HashMap::new()).await + self.execute(state, active, None, HashMap::new(), HashMap::new()) + .await } /// Runs the graph under a thread id, persisting checkpoints at every @@ -434,6 +471,7 @@ where vec![Activation::node(self.entry.clone())], Some(thread_id.into()), HashMap::new(), + HashMap::new(), ) .await } @@ -448,8 +486,14 @@ where inputs: impl IntoIterator, ) -> Result> { let active = self.initial_inputs(inputs)?; - self.execute(state, active, Some(thread_id.into()), HashMap::new()) - .await + self.execute( + state, + active, + Some(thread_id.into()), + HashMap::new(), + HashMap::new(), + ) + .await } /// Resumes an interrupted run from its latest checkpoint, re-running the @@ -529,7 +573,18 @@ where checkpoint_id: CheckpointId::new(checkpoint.checkpoint_id.clone()), }); - let active = checkpoint.next_nodes.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(), @@ -538,16 +593,21 @@ where let mut resume_map = HashMap::new(); if let Some(value) = command.resume { - for node in &active { - resume_map.insert(node.clone(), value.clone()); + 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); + self.execute( checkpoint.state, - active.into_iter().map(Activation::node).collect(), + active, Some(thread_id), resume_map, + initial_barriers, ) .await } @@ -699,6 +759,15 @@ where 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)); @@ -713,6 +782,8 @@ where 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?; @@ -779,6 +850,8 @@ where 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?; @@ -792,6 +865,7 @@ where initial_active: Vec, thread_id: Option, resume_map: HashMap, + initial_barriers: HashMap>, ) -> Result> { let run_id = crate::harness::ids::new_run_id(); // When a durable journal is configured, run against a clone whose event @@ -801,11 +875,25 @@ where // 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 + this.execute_run( + run_id, + state, + initial_active, + thread_id, + resume_map, + initial_barriers, + ) + .await } else { - self.execute_run(run_id, state, initial_active, thread_id, resume_map) - .await + self.execute_run( + run_id, + state, + initial_active, + thread_id, + resume_map, + initial_barriers, + ) + .await } } @@ -845,6 +933,7 @@ where initial_active: Vec, thread_id: Option, mut resume_map: HashMap, + initial_barriers: HashMap>, ) -> Result> { let started_at = SystemTime::now(); let mut visited: Vec = Vec::new(); @@ -901,7 +990,9 @@ where 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(); + // 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(), @@ -1008,18 +1099,31 @@ where // checkpoint is a no-op and the run aborts exactly as before. if let Some(fail) = failure { let StepFailure { - failed_node, - completed, - pending, + 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 = self.route_completed( + &active[..failed_index], + &goto_map, + &state, + &mut barrier_arrivals, + )?; + let mut pending = successors; + pending.extend(active[failed_index..].iter().cloned()); + let completed_nodes = activation_nodes(&active[..failed_index]); let checkpoint_id = self .persist_failure_checkpoint( &thread_id, &run_id, &state, &pending, - &completed, + &completed_nodes, + &barrier_arrivals, parent_checkpoint.clone(), steps, &failed_node, @@ -1040,16 +1144,27 @@ where 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. + // 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 pending: Vec = activation_nodes(&active[index..]); + let successors = self.route_completed( + &active[..index], + &goto_map, + &state, + &mut barrier_arrivals, + )?; + 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 = self .persist_checkpoint( @@ -1059,6 +1174,7 @@ where &pending, &activation_nodes(&active[..index]), vec![emitted.clone()], + &barrier_arrivals, parent_checkpoint.clone(), steps, "loop", @@ -1070,7 +1186,7 @@ where 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.active_nodes = pending_nodes; status.pending_interrupts = vec![interrupt_id]; status.checkpoint_id = checkpoint_id.clone(); self.save_status(status.clone()).await; @@ -1091,56 +1207,14 @@ where } // 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 (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, - }); - } - } - } - - // 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); + // edges, evaluated against the freshly-committed state. Barrier + // arrivals accumulate into `barrier_arrivals` (persisted below). + let completed_nodes = activation_nodes(&active); + let next = self.route_completed(&active, &goto_map, &state, &mut barrier_arrivals)?; + + // 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, @@ -1150,9 +1224,10 @@ where &thread_id, &run_id, &state, - &next_nodes, + &next, &completed_nodes, Vec::new(), + &barrier_arrivals, parent_checkpoint.clone(), steps, "loop", @@ -1244,8 +1319,9 @@ where thread_id: &Option, run_id: &RunId, state: &State, - next_nodes: &[NodeId], + pending: &[Activation], completed_tasks: &[NodeId], + barrier_arrivals: &HashMap>, parent: Option, step: usize, failed_node: &NodeId, @@ -1263,10 +1339,12 @@ where parent_checkpoint_id: parent, namespace: self.namespace.clone(), state: state.clone(), - next_nodes: next_nodes.to_vec(), + 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, @@ -1494,12 +1572,11 @@ where 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. + // 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_node: node_id.clone(), - completed: activation_nodes(&active[..index]), - pending: activation_nodes(&active[index..]), + failed_index: index, error, }); break; @@ -1637,11 +1714,10 @@ where }); // 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. + // their successors plus this branch and the rest for a + // resumable retry. failure = Some(StepFailure { - failed_node: node_id.clone(), - completed: activation_nodes(&active[..index]), - pending: activation_nodes(&active[index..]), + failed_index: index, error, }); break; @@ -1670,6 +1746,68 @@ where }) } + /// 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). + 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 @@ -1743,9 +1881,10 @@ where thread_id: &Option, run_id: &RunId, state: &State, - next_nodes: &[NodeId], + pending: &[Activation], completed_tasks: &[NodeId], interrupts: Vec, + barrier_arrivals: &HashMap>, parent: Option, step: usize, source: &str, @@ -1763,9 +1902,11 @@ where parent_checkpoint_id: parent, namespace: self.namespace.clone(), state: state.clone(), - next_nodes: next_nodes.to_vec(), + 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, diff --git a/src/graph/compiled/test.rs b/src/graph/compiled/test.rs index 0886b0b..36d9ebb 100644 --- a/src/graph/compiled/test.rs +++ b/src/graph/compiled/test.rs @@ -879,6 +879,206 @@ 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 status_snapshot_reports_run() { let graph = adding_graph(); diff --git a/src/graph/mod.rs b/src/graph/mod.rs index f369847..7e07144 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/testkit/conformance.rs b/src/graph/testkit/conformance.rs index 0e8dfa4..b025cbf 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 }), } } diff --git a/src/lib.rs b/src/lib.rs index 00906a9..8028938 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -146,12 +146,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) --- From b0996ecbb7c780dc78bfd34743007ee289f61fec Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Fri, 3 Jul 2026 14:05:33 -0700 Subject: [PATCH 020/106] fix(graph): keep parent-checkpoint lineage connected across resume execute_run always started with parent_checkpoint = None, so the first boundary checkpoint written after a resume had no parent. That orphaned the pre-interrupt history: get_state_history stopped at the resume point and prune(keep_last) deleted the ancestors a live record still depended on. Thread the loaded checkpoint id into the run as the initial parent so the lineage spine stays connected. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/graph/compiled/mod.rs | 21 ++++++++++- src/graph/compiled/test.rs | 75 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 94 insertions(+), 2 deletions(-) diff --git a/src/graph/compiled/mod.rs b/src/graph/compiled/mod.rs index c8b1cdd..ecf1a6f 100644 --- a/src/graph/compiled/mod.rs +++ b/src/graph/compiled/mod.rs @@ -437,6 +437,7 @@ where None, HashMap::new(), HashMap::new(), + None, ) .await } @@ -455,7 +456,7 @@ where inputs: impl IntoIterator, ) -> Result> { let active = self.initial_inputs(inputs)?; - self.execute(state, active, None, HashMap::new(), HashMap::new()) + self.execute(state, active, None, HashMap::new(), HashMap::new(), None) .await } @@ -472,6 +473,7 @@ where Some(thread_id.into()), HashMap::new(), HashMap::new(), + None, ) .await } @@ -492,6 +494,7 @@ where Some(thread_id.into()), HashMap::new(), HashMap::new(), + None, ) .await } @@ -601,6 +604,9 @@ where // 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, @@ -608,6 +614,7 @@ where Some(thread_id), resume_map, initial_barriers, + initial_parent, ) .await } @@ -859,6 +866,7 @@ where Ok(config) } + #[allow(clippy::too_many_arguments)] async fn execute( &self, state: State, @@ -866,6 +874,7 @@ where 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 @@ -882,6 +891,7 @@ where thread_id, resume_map, initial_barriers, + initial_parent, ) .await } else { @@ -892,6 +902,7 @@ where thread_id, resume_map, initial_barriers, + initial_parent, ) .await } @@ -926,6 +937,7 @@ where } } + #[allow(clippy::too_many_arguments)] async fn execute_run( &self, run_id: RunId, @@ -934,12 +946,17 @@ where 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; - let mut parent_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 diff --git a/src/graph/compiled/test.rs b/src/graph/compiled/test.rs index 36d9ebb..de80797 100644 --- a/src/graph/compiled/test.rs +++ b/src/graph/compiled/test.rs @@ -329,6 +329,81 @@ async fn interrupt_then_resume_reruns_node() { assert_eq!(resumed.status.status, ExecutionStatus::Completed); } +#[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() From dccf8a227f25407be3cd9bd2a5366c8c73c6377d Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Fri, 3 Jul 2026 14:09:55 -0700 Subject: [PATCH 021/106] fix(graph): propagate subgraph interrupts to the parent run An embedded child graph that paused on a human-approval interrupt had its partial state treated as a completed output: the subgraph node returned NodeResult::Update(child.state) and the parent kept executing, so the interrupt was unreachable through the parent's resume API. Detect an interrupted child execution and re-emit its interrupt as the parent node's interrupt so the parent pauses too. On a resumed activation the child is now resumed from its own checkpoint (rather than re-run from scratch), carrying the parent's resume value down. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/graph/subgraph/mod.rs | 46 +++++++++++++++++++++++++++++++------ src/graph/subgraph/test.rs | 47 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 86 insertions(+), 7 deletions(-) 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..2cd7f39 100644 --- a/src/graph/subgraph/test.rs +++ b/src/graph/subgraph/test.rs @@ -229,6 +229,53 @@ 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_child_run_distinct_and_shares_root() { // A parent embedding one child: the parent run records exactly one child run From 3d525056fff81e9fb1f6b0fab51e7e58218c24b9 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Fri, 3 Jul 2026 14:14:30 -0700 Subject: [PATCH 022/106] fix(graph): make checkpoint lookups namespace-aware A subgraph child runs under the parent's thread id but a distinct namespace, yet no backend filtered lookups on the stored namespace: a parent resume (or state inspection) could load the child's checkpoint, resuming a same-named node with the wrong role or failing with MissingNode. Add a default get_scoped(thread, id, namespace) to the Checkpointer trait (composed from list+get, so all three backends inherit it), route get_tuple through it, and scope the executor's resume/update/fork lookups to the graph's own namespace. This also completes nested resume: resuming the parent now resumes the child from its own namespaced checkpoint. Pinned in the checkpointer conformance suite. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/graph/checkpoint/mod.rs | 42 +++++++++++++++++++++++++++++++- src/graph/compiled/mod.rs | 17 +++++++------ src/graph/subgraph/test.rs | 33 +++++++++++++++++++++++++ src/graph/testkit/conformance.rs | 30 +++++++++++++++++++++++ 4 files changed, 114 insertions(+), 8 deletions(-) diff --git a/src/graph/checkpoint/mod.rs b/src/graph/checkpoint/mod.rs index 93364ff..795b96b 100644 --- a/src/graph/checkpoint/mod.rs +++ b/src/graph/checkpoint/mod.rs @@ -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); diff --git a/src/graph/compiled/mod.rs b/src/graph/compiled/mod.rs index ecf1a6f..c574b15 100644 --- a/src/graph/compiled/mod.rs +++ b/src/graph/compiled/mod.rs @@ -562,7 +562,7 @@ where ResumeTarget::Checkpoint(id) => Some(id.as_str()), }; let checkpoint = checkpointer - .get(thread_id.as_str(), checkpoint_id) + .get_scoped(thread_id.as_str(), checkpoint_id, &self.namespace) .await? .ok_or_else(|| match &target { ResumeTarget::Latest => { @@ -746,11 +746,14 @@ where 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 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)?; @@ -836,7 +839,7 @@ where ) -> Result { let checkpointer = self.require_checkpointer()?; let source = checkpointer - .get(source_thread, source_checkpoint_id) + .get_scoped(source_thread, source_checkpoint_id, &self.namespace) .await? .ok_or_else(|| { TinyAgentsError::Checkpoint(format!( diff --git a/src/graph/subgraph/test.rs b/src/graph/subgraph/test.rs index 2cd7f39..af2fdd2 100644 --- a/src/graph/subgraph/test.rs +++ b/src/graph/subgraph/test.rs @@ -276,6 +276,39 @@ async fn subgraph_interrupt_propagates_to_parent() { ); } +#[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/conformance.rs b/src/graph/testkit/conformance.rs index b025cbf..d1f4ad3 100644 --- a/src/graph/testkit/conformance.rs +++ b/src/graph/testkit/conformance.rs @@ -95,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"); From ff659dc16d365da16db3e03151cef1707b264402 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Fri, 3 Jul 2026 14:19:04 -0700 Subject: [PATCH 023/106] fix(graph): fail the run on reducer/route/persist errors at the boundary MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Errors raised at the step boundary — a reducer merge, route resolution, or checkpoint persist — used `?` and unwound out of the run loop without calling fail_run, so observers saw the run stuck in Running forever. Route each boundary error through a fail_and_return helper that records the Failed status first. A failure-boundary persist error no longer replaces the original node error either: it just drops the resumable checkpoint reference and keeps reporting the node error. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/graph/compiled/mod.rs | 122 +++++++++++++++++++++++++++++-------- src/graph/compiled/test.rs | 33 ++++++++++ 2 files changed, 130 insertions(+), 25 deletions(-) diff --git a/src/graph/compiled/mod.rs b/src/graph/compiled/mod.rs index c574b15..7ac3c17 100644 --- a/src/graph/compiled/mod.rs +++ b/src/graph/compiled/mod.rs @@ -1097,9 +1097,18 @@ where } }; - // Apply collected updates through the reducer at the boundary. + // 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 = self.reducer.apply(state, update)?; + 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 @@ -1127,15 +1136,25 @@ where // 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 = self.route_completed( + 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, @@ -1151,7 +1170,8 @@ where &recursion_meta, &child_runs_meta, ) - .await?; + .await + .unwrap_or(None); self.fail_run( &run_id, &thread_id, @@ -1176,17 +1196,24 @@ where .await; return Err(err); } - let successors = self.route_completed( + 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 = self + let checkpoint_id = match self .persist_checkpoint( &thread_id, &run_id, @@ -1201,7 +1228,15 @@ where &recursion_meta, &child_runs_meta, ) - .await?; + .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; @@ -1230,7 +1265,15 @@ where // edges, evaluated against the freshly-committed state. Barrier // arrivals accumulate into `barrier_arrivals` (persisted below). let completed_nodes = activation_nodes(&active); - let next = self.route_completed(&active, &goto_map, &state, &mut barrier_arrivals)?; + 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 @@ -1240,21 +1283,30 @@ where DurabilityMode::Sync | DurabilityMode::Async => true, }; let checkpoint_id = if persist_now { - 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? + 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 }; @@ -1322,6 +1374,26 @@ where 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. /// diff --git a/src/graph/compiled/test.rs b/src/graph/compiled/test.rs index de80797..3982738 100644 --- a/src/graph/compiled/test.rs +++ b/src/graph/compiled/test.rs @@ -1154,6 +1154,39 @@ async fn barrier_arrivals_survive_interrupt_and_resume() { 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(); From fe020f91e9aeb3cddb7084301ce77cdd09a76f13 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Fri, 3 Jul 2026 14:23:02 -0700 Subject: [PATCH 024/106] fix(graph): emit CheckpointRestored (not CheckpointSaved) on resume load resume_from emitted GraphEvent::CheckpointSaved when it *loaded* a checkpoint, so durability observers counting saves double-counted every resume as a persist. Add a CheckpointRestored event for the read side and emit it on load instead. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/graph/compiled/mod.rs | 5 +++- src/graph/compiled/test.rs | 52 ++++++++++++++++++++++++++++++++++++++ src/graph/stream/types.rs | 6 +++++ 3 files changed, 62 insertions(+), 1 deletion(-) diff --git a/src/graph/compiled/mod.rs b/src/graph/compiled/mod.rs index 7ac3c17..711acf3 100644 --- a/src/graph/compiled/mod.rs +++ b/src/graph/compiled/mod.rs @@ -572,7 +572,10 @@ where "no checkpoint `{id}` found for thread `{thread_id}`" )), })?; - self.emit(GraphEvent::CheckpointSaved { + // 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()), }); diff --git a/src/graph/compiled/test.rs b/src/graph/compiled/test.rs index 3982738..63d6af5 100644 --- a/src/graph/compiled/test.rs +++ b/src/graph/compiled/test.rs @@ -329,6 +329,58 @@ 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 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", From 584e2dffb7695a52eaae39a71d847c57c4753f21 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Fri, 3 Jul 2026 18:57:02 -0700 Subject: [PATCH 025/106] fix(graph): reject update_state as a command node A command node routes dynamically via the Command it returns at runtime and has no static successors, so route() resolves it as a sink. Attributing a manual update_state write to such a node persisted an empty next_nodes and silently rendered the thread non-resumable. Reject it at write time with a clear Graph error instead. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/graph/compiled/mod.rs | 25 ++++++++++++++++++++----- src/graph/compiled/test.rs | 37 +++++++++++++++++++++++++++++++++++++ 2 files changed, 57 insertions(+), 5 deletions(-) diff --git a/src/graph/compiled/mod.rs b/src/graph/compiled/mod.rs index 711acf3..4ce32b2 100644 --- a/src/graph/compiled/mod.rs +++ b/src/graph/compiled/mod.rs @@ -733,7 +733,10 @@ where /// 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 + /// 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( @@ -743,10 +746,22 @@ where 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())); + 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 diff --git a/src/graph/compiled/test.rs b/src/graph/compiled/test.rs index 63d6af5..72505d2 100644 --- a/src/graph/compiled/test.rs +++ b/src/graph/compiled/test.rs @@ -664,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; From 110c8e8d26f40cc781dcf76e8b826ff4c35229e7 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Fri, 3 Jul 2026 19:00:07 -0700 Subject: [PATCH 026/106] fix(graph): map_reduce FailFast returns first error in input order Items complete out of order under buffer_unordered, so FailFast broke on the first error to *complete* and could return a later item's error instead of the first failure in input order (as documented). Track the lowest-index error and only stop once every earlier item has resolved, then cancel remaining work. Also corrects the stale module doc that claimed ordering came from `buffered`. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/graph/parallel/mod.rs | 31 +++++++++++++++++++++++++------ src/graph/parallel/test.rs | 27 +++++++++++++++++++++++++++ 2 files changed, 52 insertions(+), 6 deletions(-) diff --git a/src/graph/parallel/mod.rs b/src/graph/parallel/mod.rs index d711778..03fbf66 100644 --- a/src/graph/parallel/mod.rs +++ b/src/graph/parallel/mod.rs @@ -9,8 +9,10 @@ //! 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; @@ -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)); From 3b76789c0bad5d62606e8bf2d9eed218274f1269 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Fri, 3 Jul 2026 19:17:11 -0700 Subject: [PATCH 027/106] docs(harness): add module READMEs for middleware, openai provider, agent_loop, observability These modules had no README.md despite being complex, multi-file public surfaces (onion middleware composition, provider translation/SSE decoding, the default agent loop, durable observability journals/sinks). Document design, public surface, and operational constraints for each, consistent with the repo's per-module README convention. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/harness/agent_loop/README.md | 92 ++++++++++++++++++++ src/harness/middleware/README.md | 111 +++++++++++++++++++++++++ src/harness/observability/README.md | 80 ++++++++++++++++++ src/harness/providers/openai/README.md | 98 ++++++++++++++++++++++ 4 files changed, 381 insertions(+) create mode 100644 src/harness/agent_loop/README.md create mode 100644 src/harness/middleware/README.md create mode 100644 src/harness/observability/README.md create mode 100644 src/harness/providers/openai/README.md diff --git a/src/harness/agent_loop/README.md b/src/harness/agent_loop/README.md new file mode 100644 index 0000000..97786e7 --- /dev/null +++ b/src/harness/agent_loop/README.md @@ -0,0 +1,92 @@ +# 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` | `AgentHarness::invoke` / `invoke_with_status` and the loop body. | +| `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/middleware/README.md b/src/harness/middleware/README.md new file mode 100644 index 0000000..a4f06f8 --- /dev/null +++ b/src/harness/middleware/README.md @@ -0,0 +1,111 @@ +# 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. + +## 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, built-in `Middleware` impls. | +| `library/` | Constructors and additional impls for the built-in middleware. | +| `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/observability/README.md b/src/harness/observability/README.md new file mode 100644 index 0000000..3a27fa3 --- /dev/null +++ b/src/harness/observability/README.md @@ -0,0 +1,80 @@ +# 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 with +`futures::executor::block_on`, and treat persistence as **best-effort**: a +backend error is logged/dropped and never aborts the run. 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.rs` | `LangfuseClient` and payload helpers (`clean_nulls`, `iso_ms`) shared with `graph::observability::langfuse`. | +| `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/providers/openai/README.md b/src/harness/providers/openai/README.md new file mode 100644 index 0000000..7218f4e --- /dev/null +++ b/src/harness/providers/openai/README.md @@ -0,0 +1,98 @@ +# 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`; `mod.rs` owns only the +translation logic and the HTTP transport, 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` / +`SseAccumulator`, roughly lines 845+ in `mod.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 and transport failures are normalized through +`parse_error_body` / `provider_error` / `provider_failure_message` into +`TinyAgentsError::Model`; 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` | `OpenAiModel`, translation logic, HTTP transport, SSE decoding. | +| `types.rs` | Wire (de)serialization shapes (`ModelListWire`, `ModelListing`, request/response bodies). | +| `test.rs` | Unit tests (SSE boundary decoding, tool-call correlation, error mapping, presets). | From a7d42bb2654b5d405a6d77ed62fa1ac4eb1773fe Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Fri, 3 Jul 2026 19:18:59 -0700 Subject: [PATCH 028/106] docs(graph): add module READMEs for compiled, checkpoint, channel Document the superstep executor's parallel/retry/resumable-failure semantics, the Checkpointer trait and its file/sqlite/in-memory backends, and the channel-per-field reducer model, matching the repo's per-module README convention. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/graph/channel/README.md | 81 +++++++++++++++++++ src/graph/checkpoint/README.md | 95 ++++++++++++++++++++++ src/graph/compiled/README.md | 140 +++++++++++++++++++++++++++++++++ 3 files changed, 316 insertions(+) create mode 100644 src/graph/channel/README.md create mode 100644 src/graph/checkpoint/README.md create mode 100644 src/graph/compiled/README.md 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/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/compiled/README.md b/src/graph/compiled/README.md new file mode 100644 index 0000000..5f9ad8a --- /dev/null +++ b/src/graph/compiled/README.md @@ -0,0 +1,140 @@ +# 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` | The superstep loop, retry/resumable-failure machinery, run/resume/state APIs. | +| `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. From 3d30ac0b82f0b30fe412d60e1a53ac0d5e6b6a0f Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Fri, 3 Jul 2026 19:21:13 -0700 Subject: [PATCH 029/106] docs(graph): add module READMEs for observability, orchestration, subgraph, testkit Document durable graph observability (journals/sinks/Langfuse export), the managed child-work orchestration tool surface, graph-in-graph subgraph embedding, and the graph-level test doubles/assertions/conformance suites. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/graph/observability/README.md | 75 +++++++++++++++++++++++++ src/graph/orchestration/README.md | 92 +++++++++++++++++++++++++++++++ src/graph/subgraph/README.md | 77 ++++++++++++++++++++++++++ src/graph/testkit/README.md | 89 ++++++++++++++++++++++++++++++ 4 files changed, 333 insertions(+) create mode 100644 src/graph/observability/README.md create mode 100644 src/graph/orchestration/README.md create mode 100644 src/graph/subgraph/README.md create mode 100644 src/graph/testkit/README.md diff --git a/src/graph/observability/README.md b/src/graph/observability/README.md new file mode 100644 index 0000000..bd20ea2 --- /dev/null +++ b/src/graph/observability/README.md @@ -0,0 +1,75 @@ +# 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 with `futures::executor::block_on`, and treats persistence +as **best-effort**: a backend error never aborts the run. 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/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/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/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. From 38e66d669ac05f56d4397269a64ae18e6d6321b1 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Fri, 3 Jul 2026 19:22:04 -0700 Subject: [PATCH 030/106] docs(graph): add module README for src/graph root Document the graph runtime's module map and how authoring, execution, persistence, observability, recursion, and testing pieces fit together, linking to each submodule's own README. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/graph/README.md | 83 +++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 83 insertions(+) create mode 100644 src/graph/README.md 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. From 1b06c8f4ef7fe8fd4c8906de04cf521cb44a552f Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Fri, 3 Jul 2026 19:26:27 -0700 Subject: [PATCH 031/106] docs(repl): cross-reference the two ReplSession types repl::ReplSession (line-oriented command skeleton) and repl::session::ReplSession (Rhai scripting session, feature `repl`) share a bare name across module paths, and only the latter is re-exported at the crate root when the feature is enabled. Add rustdoc on both types pointing at each other so callers can tell which one they imported. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/repl/session/mod.rs | 18 ++++++++++++++++++ src/repl/types.rs | 26 +++++++++++++++++++++++++- 2 files changed, 43 insertions(+), 1 deletion(-) diff --git a/src/repl/session/mod.rs b/src/repl/session/mod.rs index 7543585..62c0cf7 100644 --- a/src/repl/session/mod.rs +++ b/src/repl/session/mod.rs @@ -159,6 +159,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, 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 /// From 35af4936bbc1fbe0e1037a4c9b8a040a46d6457f Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Fri, 3 Jul 2026 19:45:20 -0700 Subject: [PATCH 032/106] fix(harness): unify model/tool call limits and cap fallback loops Two termination bugs in the agent loop: - `RunPolicy::limits` and the `RunConfig`-derived `LimitTracker` on `RunContext` disagreed on the model/tool call cap. A policy configured with a higher cap than `RunConfig`'s default (25) would trip the context tracker first, but the mapped error still reported the policy's (wrong) limit. `LimitTracker::sync_call_limits` reconciles the two into one enforced source at the start of each run, and the agent loop now reports whichever limit actually tripped. - A fallback chain containing a repeated model name (e.g. `[primary, backup, primary]`) alternated between the same two models forever, since `FallbackPolicy::next_after` always resolves from the first occurrence of the current name. The loop now tracks visited model names per call and never retries one already tried, and a non-retryable `Timeout` no longer feeds back into the fallback chain. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/harness/agent_loop/mod.rs | 82 ++++++++++++++++++++-------------- src/harness/agent_loop/test.rs | 73 ++++++++++++++++++++++++++++++ src/harness/limits/mod.rs | 15 +++++++ 3 files changed, 137 insertions(+), 33 deletions(-) diff --git a/src/harness/agent_loop/mod.rs b/src/harness/agent_loop/mod.rs index 3d2efcf..d574bfd 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 //! @@ -306,6 +310,15 @@ impl AgentHarness { 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; status.mark_running(HarnessPhase::Middleware); @@ -340,14 +353,14 @@ impl AgentHarness { ctx.run_id() ))); } - if run.model_calls >= self.policy.limits.max_model_calls { + // 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(format!( - "max model calls ({}) reached", - self.policy.limits.max_model_calls - ))); + return Err(TinyAgentsError::LimitExceeded(err.to_string())); } // Build the request from the working transcript, tool schemas, and @@ -367,15 +380,6 @@ impl AgentHarness { .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 @@ -555,21 +559,16 @@ impl AgentHarness { ctx.run_id() ))); } - if run.tool_calls >= self.policy.limits.max_tool_calls { + // 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(format!( - "max tool calls ({}) reached", - self.policy.limits.max_tool_calls - ))); + return Err(TinyAgentsError::LimitExceeded(err.to_string())); } - 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) @@ -802,6 +801,12 @@ impl AgentHarness { 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. @@ -858,16 +863,27 @@ impl AgentHarness { 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. + // 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); + .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()), diff --git a/src/harness/agent_loop/test.rs b/src/harness/agent_loop/test.rs index 30db8bf..3b9fe83 100644 --- a/src/harness/agent_loop/test.rs +++ b/src/harness/agent_loop/test.rs @@ -399,6 +399,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(); @@ -701,6 +737,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}; diff --git a/src/harness/limits/mod.rs b/src/harness/limits/mod.rs index 11bd695..4c47e1b 100644 --- a/src/harness/limits/mod.rs +++ b/src/harness/limits/mod.rs @@ -190,6 +190,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)] From 61feba46abef4837d5ad3f62c06aa18723b9e031 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Fri, 3 Jul 2026 20:01:55 -0700 Subject: [PATCH 033/106] fix(repl): enforce ReplPolicy::timeout and max_iterations fail-closed MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `ReplPolicy::timeout` was parsed into the policy but nothing ever read it, so a hanging capability call (or a runaway script loop) could block a `.ragsh` session forever. Two enforcement points now cover it: - The engine's `on_progress` hook checks the per-cell deadline between Rhai statements/operations, terminating a runaway script loop with a `Timeout` error. - `bridge_block_on` (the sync/async adapter every `model_query`, `tool_call`, `agent_query`, and their batched variants go through) now races the capability future against the same deadline and drops it on elapse, which cancels the underlying request rather than blocking indefinitely, since providers are built on cancel-safe `reqwest`/`futures`. Also addressed the two dead `ReplPolicy`/`ReplCapabilities` fields flagged alongside it: - `max_iterations` is now enforced: each `eval_cell` call is one CodeAct-style iteration, counted and capped fail-closed. - `ReplCapabilities::stores` was removed rather than left half-wired — no built-in ever read or wrote through it. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/repl/session/builtins.rs | 205 ++++++++++++++++++++++++++++------- src/repl/session/mod.rs | 61 ++++++++++- src/repl/session/test.rs | 50 +++++++++ src/repl/session/types.rs | 27 ++--- 4 files changed, 285 insertions(+), 58 deletions(-) diff --git a/src/repl/session/builtins.rs b/src/repl/session/builtins.rs index 1d9a834..172a7b7 100644 --- a/src/repl/session/builtins.rs +++ b/src/repl/session/builtins.rs @@ -110,10 +110,58 @@ pub(super) struct HostContext { 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) +/// 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, @@ -312,8 +360,8 @@ fn model_query_impl( .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 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(); @@ -349,7 +397,8 @@ fn tool_call_impl( arguments: arguments.clone(), }; let start = Instant::now(); - let result = bridge_block_on(tool.call(&ctx.state, call)).map_err(|err| raise(ctx, err))?; + 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, @@ -398,8 +447,8 @@ fn agent_query_impl( 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))?; + let output = bridge_block_on(ctx.buffers.deadline(), agent.run(input, ctx.events.clone())) + .map_err(|err| raise(ctx, err))?; record( ctx, ReplCallKind::Agent, @@ -703,22 +752,24 @@ fn model_query_batched_impl( } 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 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 { @@ -759,7 +810,7 @@ fn tool_call_batched_impl( let concurrency = ctx.policy.max_concurrency.max(1); let results: Vec< Result<(String, crate::harness::tool::ToolResult, Duration), TinyAgentsError>, - > = bridge_block_on(async { + > = bridge_block_on_raw(ctx.buffers.deadline(), async { stream::iter(prepared.iter().map(|(name, tool, arguments)| { let name = name.clone(); let call = ToolCall { @@ -776,7 +827,8 @@ fn tool_call_batched_impl( .buffered(concurrency) .collect() .await - }); + }) + .map_err(|err| raise(ctx, err))?; let mut out = Array::with_capacity(results.len()); for result in results { @@ -827,19 +879,21 @@ fn agent_query_batched_impl( } 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 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 { @@ -864,6 +918,12 @@ fn graph_run_batched_impl( // ── 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. /// @@ -874,6 +934,22 @@ pub(super) fn build_engine(ctx: Arc= 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)); @@ -1007,3 +1083,52 @@ pub(super) fn build_engine(ctx: Arc(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 62c0cf7..4281739 100644 --- a/src/repl/session/mod.rs +++ b/src/repl/session/mod.rs @@ -72,6 +72,13 @@ 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 persistent variable namespace of a session. @@ -213,6 +220,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 { @@ -266,6 +277,7 @@ impl ReplSession { drafts: Arc::new(Mutex::new(BTreeMap::new())), engine: Engine::new(), buffers, + iterations: 0, }; session.rebuild_engine(); session @@ -350,13 +362,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 {}", @@ -365,8 +391,15 @@ 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(); + self.buffers + .arm_deadline(self.policy.timeout.map(|d| start + d)); let before = self.variables.snapshot(); // Expose the pre-cell namespace to `show_vars()`. @@ -439,6 +472,19 @@ 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; + } + + /// Arms the per-cell wall-clock deadline, replacing any previous one. + fn arm_deadline(&self, deadline: Option) { + *self.deadline.lock().expect("deadline poisoned") = deadline; + } + + /// 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 { @@ -498,6 +544,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..7c4f170 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; @@ -62,6 +64,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 { diff --git a/src/repl/session/types.rs b/src/repl/session/types.rs index e622790..d75f41d 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; @@ -293,19 +292,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 +319,6 @@ impl ReplCapabilities { pub fn new(registry: Arc>) -> Self { Self { registry, - stores: StoreRegistry::new(), language: None, } } @@ -326,12 +329,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) From 47295a75cdccd95ca3164a982229a8cba4ee743d Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Fri, 3 Jul 2026 20:12:29 -0700 Subject: [PATCH 034/106] fix(harness): wire RunLimits::max_retries_per_call and bound tool calls by deadline `RunLimits::max_retries_per_call` and `max_concurrency` were parsed and exposed via builder methods but never consulted anywhere: - `max_retries_per_call` now caps the agent loop's retry decision via `RetryPolicy::max_attempts_capped_at`, so it acts as a hard ceiling a looser `RetryPolicy::max_attempts` cannot exceed, matching the field's "hard limits enforced fail-closed" doc contract. - `max_concurrency` had no real enforcement point anywhere in the harness (tool calls run serially; no other in-flight-call fan-out exists to bound) and is unrelated to the separate, already-wired `graph::parallel`/`ReplPolicy` concurrency knobs. Removed rather than left advertising a capability the engine never provided. Also, the run's remaining wall-clock budget was previously enforced only around model calls (`with_call_budget`); a hanging tool call could block the run past its deadline. `with_call_budget` is now generic over the call's result type and wraps the tool-call future the same way, canceling it on elapse. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/harness/agent_loop/mod.rs | 43 ++++++++++++------ src/harness/agent_loop/test.rs | 76 ++++++++++++++++++++++++++++++++ src/harness/limits/mod.rs | 9 +--- src/harness/limits/types.rs | 10 +++-- src/harness/retry/mod.rs | 14 ++++++ src/harness/retry/test.rs | 18 ++++++++ tests/e2e_misc_public_helpers.rs | 2 - 7 files changed, 146 insertions(+), 26 deletions(-) diff --git a/src/harness/agent_loop/mod.rs b/src/harness/agent_loop/mod.rs index d574bfd..3772760 100644 --- a/src/harness/agent_loop/mod.rs +++ b/src/harness/agent_loop/mod.rs @@ -649,11 +649,14 @@ impl AgentHarness { // 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. + // 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 mut result = self - .middleware - .run_wrapped_tool(ctx, state, call, &base) + 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(); @@ -830,15 +833,22 @@ impl AgentHarness { 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 + 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(), fut).await + Self::with_call_budget(remaining, run_id.as_str(), "model call", fut).await }; match attempt_result { Ok(response) => break Ok(response), Err(error) => { - if is_retryable(&error) && self.policy.retry.should_retry(attempt) { + // `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); + if is_retryable(&error) && attempt + 1 < max_attempts { attempt += 1; ctx.emit(AgentEvent::RetryScheduled { call_id: call_id.clone(), @@ -930,26 +940,33 @@ impl AgentHarness { } } - /// Awaits a single model-call future, optionally bounded by `budget`. + /// 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 request) and a + /// 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. - async fn with_call_budget(budget: Option, run_id: &str, fut: F) -> F::Output + /// approaches. `what` names the kind of call in the timeout message (e.g. + /// `"model call"`, `"tool call"`). + async fn with_call_budget( + budget: Option, + run_id: &str, + what: &str, + fut: F, + ) -> Result where - F: Future>, + 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 \ + "{what} for run `{run_id}` exceeded its remaining wall-clock budget \ ({} ms)", budget.as_millis() ))), diff --git a/src/harness/agent_loop/test.rs b/src/harness/agent_loop/test.rs index 3b9fe83..b1f4fc3 100644 --- a/src/harness/agent_loop/test.rs +++ b/src/harness/agent_loop/test.rs @@ -64,6 +64,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 { @@ -716,6 +740,31 @@ 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 non_retryable_or_exhausted_without_fallback_errors() { let mut harness: AgentHarness<()> = AgentHarness::new(); @@ -1070,6 +1119,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] diff --git a/src/harness/limits/mod.rs b/src/harness/limits/mod.rs index 4c47e1b..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; 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/retry/mod.rs b/src/harness/retry/mod.rs index e8906d8..312f84f 100644 --- a/src/harness/retry/mod.rs +++ b/src/harness/retry/mod.rs @@ -102,6 +102,20 @@ impl RetryPolicy { attempt + 1 < self.max_attempts } + /// 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`. /// diff --git a/src/harness/retry/test.rs b/src/harness/retry/test.rs index c0fafcc..b5e5087 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] diff --git a/tests/e2e_misc_public_helpers.rs b/tests/e2e_misc_public_helpers.rs index cef4c4c..75c1c50 100644 --- a/tests/e2e_misc_public_helpers.rs +++ b/tests/e2e_misc_public_helpers.rs @@ -328,10 +328,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); From b90cf9b1bc1aa3551e8af9ed37c36ab9551ebd1c Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Fri, 3 Jul 2026 20:18:52 -0700 Subject: [PATCH 035/106] fix(harness): route SubAgentTool::call_with_context depth limit to a tool error `call_with_context` propagated a `child_config` depth-limit rejection raw via `?`, bypassing the limit-to-tool-error conversion that `SubAgentTool::call` already applies to the same error and aborting the whole parent run instead of surfacing a delegated-agent limit signal the model could react to. `call_with_context` is the path the agent loop actually drives tools through, so this was the live bug. The depth-limit rejection now routes through the same `limit_result_for_parent` conversion used for a `run_child` failure. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/harness/subagent/mod.rs | 19 +++++++++++++++-- src/harness/subagent/test.rs | 40 +++++++++++++++++++++++++++++++++++- 2 files changed, 56 insertions(+), 3 deletions(-) diff --git a/src/harness/subagent/mod.rs b/src/harness/subagent/mod.rs index 47c82d1..c1c8434 100644 --- a/src/harness/subagent/mod.rs +++ b/src/harness/subagent/mod.rs @@ -557,11 +557,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( From 0e5dd1cc618530107d7f43099dcad4cb7ae63a2d Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Fri, 3 Jul 2026 20:19:08 -0700 Subject: [PATCH 036/106] Support reasoning stream deltas --- docs/modules/harness/model.md | 6 ++ docs/modules/harness/streaming.md | 1 + src/harness/agent_loop/mod.rs | 1 + src/harness/agent_loop/test.rs | 70 +++++++++++++++++++++++- src/harness/model/types.rs | 3 + src/harness/providers/openai/mod.rs | 54 +++++++++++++++++- src/harness/providers/openai/test.rs | 45 ++++++++++++++- src/harness/providers/openai/types.rs | 25 +++++++++ tests/e2e_middleware_parser_contracts.rs | 1 + wiki | 2 +- 10 files changed, 203 insertions(+), 5 deletions(-) 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/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/src/harness/agent_loop/mod.rs b/src/harness/agent_loop/mod.rs index 3772760..321497c 100644 --- a/src/harness/agent_loop/mod.rs +++ b/src/harness/agent_loop/mod.rs @@ -1039,6 +1039,7 @@ impl AgentHarness { 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(), }; self.middleware diff --git a/src/harness/agent_loop/test.rs b/src/harness/agent_loop/test.rs index b1f4fc3..c873374 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}; @@ -856,6 +857,7 @@ fn agent_run_default_is_empty() { struct DeltaRecorder { count: Arc>, texts: Arc>>, + reasonings: Arc>>, } #[async_trait] @@ -871,6 +873,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(()) } } @@ -881,6 +887,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( @@ -890,6 +897,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 @@ -912,6 +920,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] diff --git a/src/harness/model/types.rs b/src/harness/model/types.rs index b6b7930..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, diff --git a/src/harness/providers/openai/mod.rs b/src/harness/providers/openai/mod.rs index 96327c5..63e05d2 100644 --- a/src/harness/providers/openai/mod.rs +++ b/src/harness/providers/openai/mod.rs @@ -837,10 +837,54 @@ fn convert_usage(wire: UsageWire) -> Usage { .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. +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. +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 +} + // --------------------------------------------------------------------------- // Streaming (SSE) machinery // --------------------------------------------------------------------------- @@ -884,10 +928,18 @@ impl OpenAiStreamAcc { self.usage = Some(usage); pending.push_back(ModelStreamItem::UsageDelta(usage)); } - for choice in chunk.choices { + 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 { diff --git a/src/harness/providers/openai/test.rs b/src/harness/providers/openai/test.rs index 92b5cdd..8185a5d 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)); @@ -550,6 +552,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; diff --git a/src/harness/providers/openai/types.rs b/src/harness/providers/openai/types.rs index ee27f05..66fdd3c 100644 --- a/src/harness/providers/openai/types.rs +++ b/src/harness/providers/openai/types.rs @@ -103,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, @@ -272,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, @@ -292,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`]. @@ -302,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/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/wiki b/wiki index 9321604..b02f582 160000 --- a/wiki +++ b/wiki @@ -1 +1 @@ -Subproject commit 9321604d35fa34b53dc5f0383f70ae1b7410ac15 +Subproject commit b02f58269a4fc8742571f92e191ec893e9890eb0 From b1cbc22db0cfa1efe2c84eada74ceda015d14610 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Fri, 3 Jul 2026 20:31:31 -0700 Subject: [PATCH 037/106] fix(harness): correct retry backoff schedule and classify provider retryability Two related retry-correctness bugs: - Both retry sites (the agent loop's fallback/retry core and `RetryMiddleware`) computed the sleep duration 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 exponent higher than `RetryPolicy::backoff_for_attempt` documents. Both now sleep on the pre-increment attempt number. Added deterministic `start_paused` tests asserting the actual sleep gaps between attempts, which the previous sleep-off-by-default tests couldn't have caught. - The OpenAI provider's unary and streaming paths flattened a structured `ProviderError` (HTTP status, provider code, and a `retryable` flag derived from status) into a plain `TinyAgentsError::Model(String)`, so `is_retryable` could not tell a retryable 429 from a non-retryable 401 and retried both. Added `TinyAgentsError::Provider(Box)` (boxed to avoid inflating every `Result` in the crate) and routed the two HTTP-status error paths through it; `is_retryable` now reads `ProviderError::retryable` for this variant instead of assuming every provider failure is transient. `TinyAgentsError::Model` remains for transport-level failures that have no such structure to preserve. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/error.rs | 21 ++++ src/harness/agent_loop/mod.rs | 10 +- src/harness/agent_loop/test.rs | 144 +++++++++++++++++++++++++ src/harness/middleware/library/mod.rs | 10 +- src/harness/middleware/library/test.rs | 44 ++++++++ src/harness/model/mod.rs | 26 +++++ src/harness/providers/openai/README.md | 12 ++- src/harness/providers/openai/mod.rs | 12 +-- src/harness/providers/openai/test.rs | 26 +++++ src/harness/retry/mod.rs | 9 +- src/harness/retry/test.rs | 35 ++++++ 11 files changed, 333 insertions(+), 16 deletions(-) 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/harness/agent_loop/mod.rs b/src/harness/agent_loop/mod.rs index 321497c..aae30ab 100644 --- a/src/harness/agent_loop/mod.rs +++ b/src/harness/agent_loop/mod.rs @@ -849,6 +849,14 @@ impl AgentHarness { .retry .max_attempts_capped_at(self.policy.limits.max_retries_per_call); if is_retryable(&error) && attempt + 1 < max_attempts { + // 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(), @@ -857,7 +865,7 @@ impl AgentHarness { // 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; + self.policy.retry.sleep_backoff(backoff_attempt).await; continue; } break Err(error); diff --git a/src/harness/agent_loop/test.rs b/src/harness/agent_loop/test.rs index c873374..3fd7357 100644 --- a/src/harness/agent_loop/test.rs +++ b/src/harness/agent_loop/test.rs @@ -241,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; @@ -716,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(); @@ -766,6 +861,55 @@ async fn run_limits_max_retries_per_call_caps_a_looser_retry_policy() { 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(); diff --git a/src/harness/middleware/library/mod.rs b/src/harness/middleware/library/mod.rs index 14265a9..53e68ac 100644 --- a/src/harness/middleware/library/mod.rs +++ b/src/harness/middleware/library/mod.rs @@ -94,12 +94,20 @@ impl ModelMiddleware for Retry Ok(outcome) => return Ok(outcome), Err(error) => { if is_retryable(&error) && self.policy.should_retry(attempt) { + // 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(attempt).await; + self.policy.sleep_backoff(backoff_attempt).await; continue; } return Err(error); diff --git a/src/harness/middleware/library/test.rs b/src/harness/middleware/library/test.rs index f859067..de8c632 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(); diff --git a/src/harness/model/mod.rs b/src/harness/model/mod.rs index fed88b4..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 { diff --git a/src/harness/providers/openai/README.md b/src/harness/providers/openai/README.md index 7218f4e..c8dbbc3 100644 --- a/src/harness/providers/openai/README.md +++ b/src/harness/providers/openai/README.md @@ -71,9 +71,15 @@ Streaming responses are decoded by a small state machine (`SseState` / ## Error handling -Non-2xx responses and transport failures are normalized through -`parse_error_body` / `provider_error` / `provider_failure_message` into -`TinyAgentsError::Model`; malformed JSON bodies surface as +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 diff --git a/src/harness/providers/openai/mod.rs b/src/harness/providers/openai/mod.rs index 63e05d2..2b9251f 100644 --- a/src/harness/providers/openai/mod.rs +++ b/src/harness/providers/openai/mod.rs @@ -273,9 +273,7 @@ impl OpenAiModel { if !status.is_success() { let error = self.parse_error_body(status.as_u16(), &text); - return Err(TinyAgentsError::Model( - self.provider_failure_message(&error), - )); + return Err(TinyAgentsError::Provider(Box::new(error))); } let listing: ModelListWire = serde_json::from_str(&text)?; @@ -1254,9 +1252,7 @@ impl ChatModel for OpenAiModel { if !status.is_success() { let error = self.parse_error_body(status.as_u16(), &text); - return Err(TinyAgentsError::Model( - self.provider_failure_message(&error), - )); + return Err(TinyAgentsError::Provider(Box::new(error))); } let value: Value = serde_json::from_str(&text)?; @@ -1310,9 +1306,7 @@ impl ChatModel for OpenAiModel { 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), - )); + return Err(TinyAgentsError::Provider(Box::new(error))); } // Map each raw byte chunk onto an owned `Vec` so the boxed stream's diff --git a/src/harness/providers/openai/test.rs b/src/harness/providers/openai/test.rs index 8185a5d..d105445 100644 --- a/src/harness/providers/openai/test.rs +++ b/src/harness/providers/openai/test.rs @@ -404,6 +404,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"); diff --git a/src/harness/retry/mod.rs b/src/harness/retry/mod.rs index 312f84f..4bbb693 100644 --- a/src/harness/retry/mod.rs +++ b/src/harness/retry/mod.rs @@ -157,7 +157,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. | @@ -166,7 +167,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 b5e5087..5d174eb 100644 --- a/src/harness/retry/test.rs +++ b/src/harness/retry/test.rs @@ -131,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] From 71319061e26d2f17892d6a45dce635e31a36ada1 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Fri, 3 Jul 2026 20:43:45 -0700 Subject: [PATCH 038/106] fix(harness): correct usage/cost accounting for omitted and cached tokens MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Usage::add_assign summed raw total_tokens fields directly, so a record with total_tokens:0 (provider omitted it) silently dropped its own contribution when merged with a record that did report a total; fix by combining effective_total() on both sides. The OpenAI provider now also back-fills total_tokens as prompt+completion when the wire payload omits it, so a single record's total_tokens field is never a misleading zero. estimate_cost priced cache_read_tokens and reasoning_tokens both at their dedicated rate and at the standard input/output rate, since OpenAI reports them as a subset of input_tokens/output_tokens rather than an addition — double-counting cost for cached and reasoning-heavy calls. Standard-rate cost is now computed on the non-cached/non-reasoning remainder. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/harness/cost/mod.rs | 14 ++++++++++++-- src/harness/cost/test.rs | 29 +++++++++++++++++++++++++--- src/harness/providers/openai/mod.rs | 11 ++++++++++- src/harness/providers/openai/test.rs | 26 +++++++++++++++++++++++++ src/harness/usage/mod.rs | 8 +++++++- src/harness/usage/test.rs | 17 ++++++++++++++++ 6 files changed, 98 insertions(+), 7 deletions(-) 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/providers/openai/mod.rs b/src/harness/providers/openai/mod.rs index 2b9251f..ac82444 100644 --- a/src/harness/providers/openai/mod.rs +++ b/src/harness/providers/openai/mod.rs @@ -827,10 +827,19 @@ fn parse_tool_arguments(context: &str, call_id: &str, name: &str, raw: &str) -> /// Converts an OpenAI [`UsageWire`] into the harness-neutral [`Usage`]. 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: wire.total_tokens, + total_tokens, cache_read_tokens: wire .prompt_tokens_details .map(|d| d.cached_tokens) diff --git a/src/harness/providers/openai/test.rs b/src/harness/providers/openai/test.rs index d105445..b8f457b 100644 --- a/src/harness/providers/openai/test.rs +++ b/src/harness/providers/openai/test.rs @@ -346,6 +346,32 @@ 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 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(); From 0a2e20c335de66f8e2185b6209042731d206a3a7 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Fri, 3 Jul 2026 21:08:47 -0700 Subject: [PATCH 039/106] fix(harness): make BudgetTracker reservations atomic, per-run, and fail-closed BudgetMiddleware's preflight reservation checked capacity and wrote the reservation in two separate lock acquisitions, so concurrent runs sharing a tracker could all observe spare capacity and reserve past it, collectively overshooting the budget. The reservation amount also lived in one shared slot on BudgetSpend, so concurrent runs clobbered each other's in-flight reservation, and a response with usage:None never released its reservation, permanently starving later calls. Fix: check-and-reserve now happens under a single lock (aggregated into a new reserved_input_total), each BudgetMiddleware instance tracks its own run's outstanding reservation locally, and after_model always releases the reservation regardless of whether usage came back. snapshot()/record() now recover a poisoned mutex's last-written value instead of defaulting to a zero spend, so a panic while holding the lock can no longer make the budget enforcer fail open. Warn-once is now a single check-and-set under lock instead of a racy read-then-write. Also fixes a self-deadlock introduced while implementing the above: the warn-once block re-locked the same (non-reentrant) mutex via tracker.snapshot() while still holding the guard on the no-warning path. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/harness/middleware/library/mod.rs | 153 ++++++++++++++++-------- src/harness/middleware/library/test.rs | 114 ++++++++++++++++++ src/harness/middleware/library/types.rs | 16 ++- tests/e2e_harness_provider_contracts.rs | 12 +- 4 files changed, 236 insertions(+), 59 deletions(-) diff --git a/src/harness/middleware/library/mod.rs b/src/harness/middleware/library/mod.rs index 53e68ac..abd1eb3 100644 --- a/src/harness/middleware/library/mod.rs +++ b/src/harness/middleware/library/mod.rs @@ -338,9 +338,23 @@ impl BudgetTracker { 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.inner.lock().map(|g| *g).unwrap_or_default() + *self.lock_recovering() } /// Folds a model call's usage and estimated cost into the tracker. @@ -349,10 +363,9 @@ impl BudgetTracker { usage: crate::harness::usage::Usage, cost: crate::harness::cost::CostTotals, ) { - if let Ok(mut guard) = self.inner.lock() { - guard.usage += usage; - guard.cost += cost; - } + let mut guard = self.lock_recovering(); + guard.usage += usage; + guard.cost += cost; } } @@ -457,6 +470,7 @@ impl BudgetMiddleware { limits, tracker: BudgetTracker::new(), pricing: std::collections::HashMap::new(), + pending_reservation: std::sync::Mutex::new(0), } } @@ -508,40 +522,58 @@ impl Middleware for BudgetMidd _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. + // 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); - 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; + 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, }); @@ -554,17 +586,25 @@ impl Middleware for BudgetMidd _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); - // 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, @@ -576,18 +616,27 @@ impl Middleware for BudgetMidd 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) - { + // 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 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 { diff --git a/src/harness/middleware/library/test.rs b/src/harness/middleware/library/test.rs index de8c632..99c5f7a 100644 --- a/src/harness/middleware/library/test.rs +++ b/src/harness/middleware/library/test.rs @@ -643,6 +643,120 @@ 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); +} + +#[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] diff --git a/src/harness/middleware/library/types.rs b/src/harness/middleware/library/types.rs index 62df59b..874a1e8 100644 --- a/src/harness/middleware/library/types.rs +++ b/src/harness/middleware/library/types.rs @@ -204,9 +204,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 +238,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 ────────────────────────────────────────────────────── 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] From 804202806864c9623539cf06be3b742e7405fe96 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Fri, 3 Jul 2026 21:18:55 -0700 Subject: [PATCH 040/106] fix(repl): enforce agent-call limit independently of the model-call limit bump_agent compared the agent-call counter against ReplPolicy::max_model_calls (and quoted that same number in its error message), so a session had no real per-session cap on agent_query calls distinct from model_query calls. Since each agent_query itself drives one or more model calls, a session's combined model spend could reach roughly twice the configured max_model_calls before anything failed closed. Add a real ReplPolicy::max_agent_calls (default 32) and enforce it in bump_agent instead of max_model_calls. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/repl/session/builtins.rs | 4 +-- src/repl/session/test.rs | 67 ++++++++++++++++++++++++++++++++++++ src/repl/session/types.rs | 9 +++++ 3 files changed, 78 insertions(+), 2 deletions(-) diff --git a/src/repl/session/builtins.rs b/src/repl/session/builtins.rs index 172a7b7..a5f0407 100644 --- a/src/repl/session/builtins.rs +++ b/src/repl/session/builtins.rs @@ -282,12 +282,12 @@ fn bump_graph(ctx: &HostContext) -> Result<(), Box(ctx: &HostContext) -> Result<(), Box> { let mut counters = ctx.counters.lock().expect("counters poisoned"); - if counters.agent >= ctx.policy.max_model_calls { + if counters.agent >= ctx.policy.max_agent_calls { return Err(raise( ctx, TinyAgentsError::LimitExceeded(format!( "agent call limit ({}) exceeded", - ctx.policy.max_model_calls + ctx.policy.max_agent_calls )), )); } diff --git a/src/repl/session/test.rs b/src/repl/session/test.rs index 7c4f170..8808903 100644 --- a/src/repl/session/test.rs +++ b/src/repl/session/test.rs @@ -203,6 +203,73 @@ fn output_byte_limit_fails_closed() { assert!(matches!(err, TinyAgentsError::LimitExceeded(_))); } +/// 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 d75f41d..ce416c9 100644 --- a/src/repl/session/types.rs +++ b/src/repl/session/types.rs @@ -84,6 +84,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. @@ -109,6 +117,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, From 5daf35b76d2b4e7dee84435631d43e9e13f293ec Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Fri, 3 Jul 2026 21:27:47 -0700 Subject: [PATCH 041/106] fix(harness): release budget reservation on model-call error BudgetMiddleware::before_model adds the estimated input-token reservation to the shared tracker's reserved_input_total, but the only release path was after_model. A model call that ultimately fails (retries/fallback exhausted, hard provider error, middleware timeout) short-circuits with `?` before after_model runs, so the reservation leaked permanently on a process-lifetime-shared BudgetTracker, eventually starving every future call regardless of real usage. Add an on_error hook that releases the pending reservation the same way after_model does, since the harness always calls on_error for every middleware when a run fails. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/harness/middleware/library/mod.rs | 21 +++++++++++++ src/harness/middleware/library/test.rs | 41 ++++++++++++++++++++++++++ 2 files changed, 62 insertions(+) diff --git a/src/harness/middleware/library/mod.rs b/src/harness/middleware/library/mod.rs index abd1eb3..5d030da 100644 --- a/src/harness/middleware/library/mod.rs +++ b/src/harness/middleware/library/mod.rs @@ -646,6 +646,27 @@ impl Middleware for BudgetMidd } 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(()) + } } // ── ToolPolicyMiddleware ────────────────────────────────────────────────────── diff --git a/src/harness/middleware/library/test.rs b/src/harness/middleware/library/test.rs index 99c5f7a..8b25fce 100644 --- a/src/harness/middleware/library/test.rs +++ b/src/harness/middleware/library/test.rs @@ -673,6 +673,47 @@ async fn reservation_is_released_even_when_response_carries_no_usage() { 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 From d954ae1a9a4598cfc63bdbbd04ae7157282c7645 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Fri, 3 Jul 2026 21:38:42 -0700 Subject: [PATCH 042/106] fix(language): stop Literal::as_display saturating huge floats to i64::MAX A finite f64 with a zero fractional part but outside i64's range (or NaN / infinity) was cast with `as i64`, silently saturating to i64::MAX/MIN. Guard the integer-render fast path with is_finite() and an i64-range check so huge or non-finite numeric literals fall through to float formatting instead. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/language/ast.rs | 6 +++++- src/language/test.rs | 15 +++++++++++++++ 2 files changed, 20 insertions(+), 1 deletion(-) 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/test.rs b/src/language/test.rs index 73c8767..98cba59 100644 --- a/src/language/test.rs +++ b/src/language/test.rs @@ -106,6 +106,21 @@ fn invalid_escape_is_a_parse_error() { 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"); +} + // --------------------------------------------------------------------------- // Spans, source map, and diagnostics // --------------------------------------------------------------------------- From 3116b4a4114474c8d78b28f44b8f382575cd67b4 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Fri, 3 Jul 2026 21:39:19 -0700 Subject: [PATCH 043/106] fix(language): diagnose conflicting routing sources and duplicate channels/graph ids compile_graph previously resolved a node's routing by silent precedence (routes > next > command.goto > top-level edge), so e.g. `next b` plus `command { goto c }` on the same node quietly routed to `b` with the command's goto discarded, and a second top-level edge from the same source stayed in blueprint.edges even though only the first was ever used for routing. Both are now explicit compile errors naming the conflicting sources. Duplicate `channel` declarations within a graph and duplicate top-level `graph` ids across a program also compiled silently, each collapsing to whichever declaration was seen first; both are now rejected. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/language/compiler.rs | 64 +++++++++++++++++++++++++++++++++ src/language/test.rs | 44 +++++++++++++++++++++++ tests/e2e_language_contracts.rs | 1 - 3 files changed, 108 insertions(+), 1 deletion(-) diff --git a/src/language/compiler.rs b/src/language/compiler.rs index fb9df3a..50f91c2 100644 --- a/src/language/compiler.rs +++ b/src/language/compiler.rs @@ -54,6 +54,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 +83,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 +131,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 +165,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 { diff --git a/src/language/test.rs b/src/language/test.rs index 98cba59..8a5a7b2 100644 --- a/src/language/test.rs +++ b/src/language/test.rs @@ -458,6 +458,50 @@ fn duplicate_route_label_is_a_compile_error() { 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}" + ); +} + // --------------------------------------------------------------------------- // Extended grammar (H2): channels+policy, command, send/join, subgraph, // subagent, repl_agent, interrupt, io shape, checkpoint/interrupt policy. 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 From 3a1e355160907e86d9fecfbf88afe43ac03b2a71 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Fri, 3 Jul 2026 21:40:11 -0700 Subject: [PATCH 044/106] fix(language): make blueprint_diff cover command/sends/retry/metadata and graph-level fields MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit blueprint_diff previously only compared a node's kind/model/prompt/tools/ routing/agent/subgraph/script/input/join_sources/options/checkpoint/timeout, and only compared graph-level graph_id/start plus channels/nodes/edges. A revision that changed only a node's `command`, `sends`, `retry`, or `metadata`, or only the graph's `defaults`, `input`, `output`, `checkpoint`, `interrupt`, or `joins`, rendered as "no changes" — defeating the review gate blueprint_diff exists to serve. Add the missing node-field comparisons and new graph-level defaults_changed/input_changed/output_changed/checkpoint_changed/ interrupt_changed/joins_changed diff fields, rendered in Display and covered by is_empty(). Add a test asserting a command-only revision is reported as a change. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/language/diff.rs | 154 ++++++++++++++++++++++++++++++++++++++++++- src/language/test.rs | 68 +++++++++++++++++++ 2 files changed, 221 insertions(+), 1 deletion(-) 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/test.rs b/src/language/test.rs index 8a5a7b2..62b249e 100644 --- a/src/language/test.rs +++ b/src/language/test.rs @@ -1330,6 +1330,74 @@ fn blueprint_diff_of_identical_blueprints_is_empty() { 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); From 789587550d58f30e7c2b3f1730abbcc8a2df73c4 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Fri, 3 Jul 2026 21:44:53 -0700 Subject: [PATCH 045/106] fix(repl): report real 1-based line/column in .ragsh parse errors parse_command always built TinyAgentsError::Parse with a hardcoded (0, 0) position, so every parse failure (unknown verb, missing argument, unterminated quoted string, invalid call JSON) pointed nowhere. Thread the original trimmed line through split_token/require_token and compute a 1-based character column from the pointer offset of the failing subslice; line is always 1 since parse_command parses a single command line. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/repl/mod.rs | 97 +++++++++++++++++++++++++++++++++++------------- src/repl/test.rs | 41 ++++++++++++++++++++ 2 files changed, 112 insertions(+), 26 deletions(-) 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/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(); From dd11e3ea2a317689cba830c0f806b0683490c896 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Fri, 3 Jul 2026 21:55:04 -0700 Subject: [PATCH 046/106] fix(repl): enforce max_output_bytes inside push_stdout_line, not only post-hoc eval_cell previously only checked the accumulated stdout length against ReplPolicy::max_output_bytes after the whole script finished running, so a print-heavy runaway script (e.g. a tight loop with no other limit hit) could buffer unbounded output for the entire cell before the check ever ran. Arm the byte budget on CellBuffers at the start of every eval_cell call and enforce it directly inside push_stdout_line: once appending a line would exceed the budget, the line is dropped and a LimitExceeded error is stashed. Since on_print/on_debug can't themselves fail a Rhai script, wire the engine's existing on_progress poll to abort promptly once a host error is pending, and have eval_cell check for a stashed host error on the success path too (previously only checked when the script itself errored). Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/repl/session/builtins.rs | 18 +++++++--- src/repl/session/mod.rs | 65 ++++++++++++++++++++++++++++++++++-- src/repl/session/test.rs | 22 ++++++++++++ 3 files changed, 99 insertions(+), 6 deletions(-) diff --git a/src/repl/session/builtins.rs b/src/repl/session/builtins.rs index a5f0407..977771a 100644 --- a/src/repl/session/builtins.rs +++ b/src/repl/session/builtins.rs @@ -943,11 +943,21 @@ pub(super) fn build_engine(ctx: Arc= deadline => { - Some(Dynamic::from(DEADLINE_EXCEEDED_TOKEN.to_string())) + 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, } - _ => None, }); // ── stdout capture ── diff --git a/src/repl/session/mod.rs b/src/repl/session/mod.rs index 4281739..790b3b5 100644 --- a/src/repl/session/mod.rs +++ b/src/repl/session/mod.rs @@ -79,6 +79,12 @@ pub(super) struct CellBuffers { /// 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. @@ -400,6 +406,7 @@ impl ReplSession { self.buffers.reset(); self.buffers .arm_deadline(self.policy.timeout.map(|d| start + d)); + self.buffers.arm_output_limit(self.policy.max_output_bytes); let before = self.variables.snapshot(); // Expose the pre-cell namespace to `show_vars()`. @@ -420,7 +427,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 @@ -473,6 +490,10 @@ impl CellBuffers { *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. @@ -480,6 +501,15 @@ impl CellBuffers { *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`]). @@ -510,13 +540,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); diff --git a/src/repl/session/test.rs b/src/repl/session/test.rs index 8808903..bdbe355 100644 --- a/src/repl/session/test.rs +++ b/src/repl/session/test.rs @@ -203,6 +203,28 @@ 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:?}"); +} + /// A trivial [`HarnessAgent`] that returns a fixed response, for exercising /// `agent_query` without a real model/harness run. struct StubAgent; From 9a1b336bb5a3cf8cd57477d7734e103a15082e64 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Fri, 3 Jul 2026 22:01:02 -0700 Subject: [PATCH 047/106] fix(repl): don't consume a graph_define slot on a failed draft graph_define incremented the max_graph_definitions counter before parsing or compiling the source, so a request that failed to parse, failed to compile, or named a graph not present in the source still spent a definition slot even though no draft was ever recorded. Move the increment to just before the draft is inserted, after a successful parse, compile, and name lookup, re-checking the limit under the same lock to guard against a race with a concurrent call. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/repl/session/builtins.rs | 39 +++++++++++++++++++++++---------- src/repl/session/test.rs | 42 ++++++++++++++++++++++++++++++++++++ 2 files changed, 70 insertions(+), 11 deletions(-) diff --git a/src/repl/session/builtins.rs b/src/repl/session/builtins.rs index 977771a..66ca684 100644 --- a/src/repl/session/builtins.rs +++ b/src/repl/session/builtins.rs @@ -517,18 +517,18 @@ fn graph_define_impl( 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 { - 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; + 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( @@ -560,6 +560,23 @@ fn graph_define_impl( ) })?; + // 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, diff --git a/src/repl/session/test.rs b/src/repl/session/test.rs index bdbe355..d01bf0b 100644 --- a/src/repl/session/test.rs +++ b/src/repl/session/test.rs @@ -225,6 +225,48 @@ fn output_byte_limit_bounds_intra_cell_buffering_in_a_print_loop() { 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 trivial [`HarnessAgent`] that returns a fixed response, for exercising /// `agent_query` without a real model/harness run. struct StubAgent; From caee6c3bf8992266aa864bea43155e1b6d169fb6 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Fri, 3 Jul 2026 22:06:32 -0700 Subject: [PATCH 048/106] fix(repl): stop double-seeding the `answer` reserved name into scope `answer` was listed in both RESERVED_VARIABLES and RESERVED_FUNCTIONS, so `reserved_names()` yielded it twice and `ReplVariables::seeded` pushed two scope entries for the same name. `answer` is only ever a capability function (`answer(content)`), never a readable session variable like `context`/`state`, so remove it from RESERVED_VARIABLES. Also make `reserved_names()` dedupe defensively so a future accidental overlap degrades gracefully instead of silently double-seeding again. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/repl/session/test.rs | 26 ++++++++++++++++++++++++++ src/repl/session/types.rs | 15 +++++++++++---- 2 files changed, 37 insertions(+), 4 deletions(-) diff --git a/src/repl/session/test.rs b/src/repl/session/test.rs index d01bf0b..882f299 100644 --- a/src/repl/session/test.rs +++ b/src/repl/session/test.rs @@ -148,6 +148,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(); diff --git a/src/repl/session/types.rs b/src/repl/session/types.rs index ce416c9..2058bb1 100644 --- a/src/repl/session/types.rs +++ b/src/repl/session/types.rs @@ -23,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. /// @@ -55,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 ────────────────────────────────────────────────────────────────── From 27fd45b7d25789f8045c3c2f727a9d90bc2b6bea Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Fri, 3 Jul 2026 22:12:02 -0700 Subject: [PATCH 049/106] docs(repl): fix broken intra-doc link to nonexistent crate::repl::codeact app_state()'s doc comment referenced [\`crate::repl::codeact\`], a module that doesn't exist anywhere in the crate. Reword to describe the CodeAct driver loop in prose instead of a dangling intra-doc link. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/repl/session/mod.rs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/repl/session/mod.rs b/src/repl/session/mod.rs index 790b3b5..35b9ab7 100644 --- a/src/repl/session/mod.rs +++ b/src/repl/session/mod.rs @@ -338,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() } From cd98867744a6f345704a6dca24a441e0d98e3308 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Fri, 3 Jul 2026 22:19:47 -0700 Subject: [PATCH 050/106] fix(repl): stop tool_call_batched from discarding successes on one item's error MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Previously, as soon as any item in a tool_call_batched batch came back with a tool-reported error (ToolResult.error, distinct from a harness/transport Err), the whole call returned Err and every other item's already-computed successful result was thrown away — a batch of N independent tool calls lost N-1 successes because one item failed. This also contradicted the documented contract ("records per-item failures without losing successful results"). Each item's outcome is now reported independently as `#{ ok, content }` or `#{ ok: false, error }`, matching model_query_batched's per-item design. A harness/transport-level failure (no results to preserve) still aborts the whole call, as before. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/repl/session/builtins.rs | 23 ++++++-- src/repl/session/test.rs | 100 +++++++++++++++++++++++++++++++++++ 2 files changed, 120 insertions(+), 3 deletions(-) diff --git a/src/repl/session/builtins.rs b/src/repl/session/builtins.rs index 66ca684..f01f9c6 100644 --- a/src/repl/session/builtins.rs +++ b/src/repl/session/builtins.rs @@ -847,6 +847,13 @@ fn tool_call_batched_impl( }) .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))?; @@ -857,10 +864,20 @@ fn tool_call_batched_impl( json!({ "chars": tool_result.content.len() }), elapsed, ); - if let Some(error) = tool_result.error { - return Err(raise(ctx, TinyAgentsError::Tool(error))); + 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)); + } } - out.push(Dynamic::from(tool_result.content)); } Ok(Dynamic::from_array(out)) } diff --git a/src/repl/session/test.rs b/src/repl/session/test.rs index 882f299..e57f72f 100644 --- a/src/repl/session/test.rs +++ b/src/repl/session/test.rs @@ -293,6 +293,106 @@ fn graph_define_does_not_consume_the_limit_on_a_failed_draft() { ); } +/// 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; From f37d43b707cbf03bc64357296523fbbf2ebd21c0 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Fri, 3 Jul 2026 22:26:01 -0700 Subject: [PATCH 051/106] fix(language): honor a diagnostic's stored line/col when byte offsets are absent MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit into_parse_error always resolved line/column from the primary span's byte offset whenever a SourceFile was supplied, even for a back-compat span built with Span::new(line, column) — which leaves start/end at 0 because no real offset was ever available. Resolving offset 0 against the file always produced 1:1, silently discarding the real position the caller had anchored the span at. Detect the offsets-absent case (start == 0 && end == 0) and fall back to the span's own stored line/column and the source-free rendering, exactly as already happens when no SourceFile is passed at all. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/language/diagnostic.rs | 19 +++++++++------ src/language/test.rs | 50 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 62 insertions(+), 7 deletions(-) diff --git a/src/language/diagnostic.rs b/src/language/diagnostic.rs index 455e8b9..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, diff --git a/src/language/test.rs b/src/language/test.rs index 62b249e..d6c9974 100644 --- a/src/language/test.rs +++ b/src/language/test.rs @@ -220,6 +220,56 @@ fn diagnostic_renders_span_past_end_of_source_without_panic() { 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"); From 9cb8fd7c9be5e5b94ba95b5e23b780ad8c148fca Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Fri, 3 Jul 2026 22:40:38 -0700 Subject: [PATCH 052/106] perf(harness): drop per-delta middleware event bracketing The streaming `on_model_delta` stack runner emitted a `MiddlewareStarted`/`MiddlewareCompleted` pair per middleware per streamed token, each cloning `mw.name()` and taking the recorder mutex. On the hot path (50-200 deltas/s) this dominated the stream loop for zero observability value. Drop the bracketing on the delta hook only (all other hooks keep it), and build the middleware-facing `ModelDelta` before moving the event payload so the token is cloned once instead of twice. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/harness/agent_loop/mod.rs | 16 ++++++++------ src/harness/middleware/mod.rs | 25 ++++++++++----------- src/harness/middleware/test.rs | 40 ++++++++++++++++++++++++++++++++++ 3 files changed, 61 insertions(+), 20 deletions(-) diff --git a/src/harness/agent_loop/mod.rs b/src/harness/agent_loop/mod.rs index aae30ab..5d86928 100644 --- a/src/harness/agent_loop/mod.rs +++ b/src/harness/agent_loop/mod.rs @@ -1037,19 +1037,21 @@ impl AgentHarness { }; 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; - + // 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?; diff --git a/src/harness/middleware/mod.rs b/src/harness/middleware/mod.rs index 5e8b323..03b53c9 100644 --- a/src/harness/middleware/mod.rs +++ b/src/harness/middleware/mod.rs @@ -209,6 +209,15 @@ impl MiddlewareStack { /// 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, @@ -216,19 +225,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(()) diff --git a/src/harness/middleware/test.rs b/src/harness/middleware/test.rs index 8b9b10b..ea937af 100644 --- a/src/harness/middleware/test.rs +++ b/src/harness/middleware/test.rs @@ -194,6 +194,46 @@ async fn emits_started_and_completed_events() { ); } +#[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)); From bc43a9bf5b666bf5aa8a69fbcdd009d31825baa2 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Fri, 3 Jul 2026 22:42:56 -0700 Subject: [PATCH 053/106] perf(store): cache JSONL append offsets and make FileStore writes atomic MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `JsonlAppendStore::append` re-read and re-parsed the entire stream file on every append to learn the next offset — O(n²) per stream, run synchronously inside an async fn. Cache the next offset per stream (lazily initialised from disk once), hold the guard across the write so concurrent appends stay ordered, and run the blocking file I/O on `spawn_blocking` when a runtime is present so it no longer stalls a tokio worker. `FileStore::put` used a plain `fs::write` despite the type docs promising atomic-write semantics; a crash or concurrent read mid-write could expose a truncated file. Write to a uniquely named temp file in the same directory and rename over the destination, cleaning up the temp file on a failed rename. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/harness/store/mod.rs | 94 ++++++++++++++++++++++++++++++-------- src/harness/store/test.rs | 55 ++++++++++++++++++++++ src/harness/store/types.rs | 23 +++++++--- 3 files changed, 145 insertions(+), 27 deletions(-) diff --git a/src/harness/store/mod.rs b/src/harness/store/mod.rs index 4a5a50a..497a252 100644 --- a/src/harness/store/mod.rs +++ b/src/harness/store/mod.rs @@ -37,6 +37,10 @@ pub use types::*; use crate::error::{Result, TinyAgentsError}; +/// 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 ───────────────────────────────────────────────────────────── impl InMemoryStore { @@ -164,8 +168,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 +280,7 @@ impl JsonlAppendStore { pub fn new(root_dir: impl Into) -> Self { Self { root_dir: root_dir.into(), + offsets: Default::default(), } } @@ -291,26 +313,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> { 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 ──────────────────────────────────────────────────────────── From 31ece7f70f72f97b202195344dd8fe302ad379cf Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Fri, 3 Jul 2026 22:51:59 -0700 Subject: [PATCH 054/106] perf(observability): drain durable harness sinks off the run thread MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `JournalSink` and `JsonlSink` bridged the synchronous `EventListener::on_event` hook to their async journal/store backends with `futures::executor::block_on`, blocking a tokio worker for the whole append (a file write in the JSONL case) on the run's critical path and risking a deadlock on a current-thread runtime. Introduce `AppendWorker`: a bounded-channel drain backed by a dedicated thread with its own runtime. `on_event` now hands the payload off without blocking; persistence is best-effort with an explicit, documented policy — a full queue drops and counts the overflow (never stalls the run) and append errors are reported to stderr rather than silently discarded. Add `flush()` to both sinks to block until the durable log has caught up, and use it in the tests that read a sink back immediately after emitting. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/harness/observability/README.md | 14 +- src/harness/observability/mod.rs | 60 ++++-- src/harness/observability/test.rs | 48 +++++ src/harness/observability/types.rs | 30 +-- src/harness/observability/worker.rs | 177 ++++++++++++++++++ tests/e2e_observability.rs | 11 +- tests/e2e_registry_observability_contracts.rs | 2 + 7 files changed, 311 insertions(+), 31 deletions(-) create mode 100644 src/harness/observability/worker.rs diff --git a/src/harness/observability/README.md b/src/harness/observability/README.md index 3a27fa3..1fe52b1 100644 --- a/src/harness/observability/README.md +++ b/src/harness/observability/README.md @@ -45,11 +45,15 @@ surface and operational constraints. ## Persistence bridge Persisting sinks (`JournalSink`, `JsonlSink`) 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 is logged/dropped and never aborts the run. Do not rely on a -sink for delivery guarantees stronger than "usually persisted, never -run-blocking." +`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 diff --git a/src/harness/observability/mod.rs b/src/harness/observability/mod.rs index 0cb82ec..66c9d7a 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 @@ -444,11 +449,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, } } @@ -458,6 +471,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 { @@ -468,8 +490,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); } } @@ -483,10 +505,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(); } } @@ -495,8 +533,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 6ac288b..135bb41 100644 --- a/src/harness/observability/test.rs +++ b/src/harness/observability/test.rs @@ -6,6 +6,7 @@ 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::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>, diff --git a/src/harness/observability/types.rs b/src/harness/observability/types.rs index 007e41d..735fff8 100644 --- a/src/harness/observability/types.rs +++ b/src/harness/observability/types.rs @@ -20,7 +20,9 @@ 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::store::AppendStore; + +use super::worker::AppendWorker; // --------------------------------------------------------------------------- // AgentObservation @@ -300,37 +302,39 @@ 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, + /// Background drain that appends records without blocking the run. + pub(crate) worker: Arc>, } /// Returns the current time in Unix-epoch milliseconds, saturating at `0`. 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/tests/e2e_observability.rs b/tests/e2e_observability.rs index b822dfd..ad99876 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) diff --git a/tests/e2e_registry_observability_contracts.rs b/tests/e2e_registry_observability_contracts.rs index 6335741..5c32ec3 100644 --- a/tests/e2e_registry_observability_contracts.rs +++ b/tests/e2e_registry_observability_contracts.rs @@ -435,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!( @@ -453,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); From 12fed70a6411a7761574108fef6124941c57209b Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Fri, 3 Jul 2026 23:00:23 -0700 Subject: [PATCH 055/106] perf(graph): drain JournalGraphSink off the executor thread `JournalGraphSink::emit` bridged the synchronous `GraphEventSink::emit` hook to its async journal with `futures::executor::block_on`, blocking the executor for the whole append on the step critical path and risking a deadlock on a current-thread runtime. Route observations through the shared `AppendWorker` background drain: `emit` now hands off without blocking, with the same bounded-queue drop policy and error reporting as the harness sinks. Add a default-no-op `flush` to the `GraphEventSink` trait (durable sinks override it); the executor flushes after a terminal run event so a caller that reads the journal right after the run returns still sees a complete log. Tests that emit directly then read now flush first. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/graph/compiled/mod.rs | 10 ++++++++++ src/graph/observability/README.md | 12 ++++++++---- src/graph/observability/mod.rs | 30 +++++++++++++++++++++++++----- src/graph/observability/test.rs | 3 +++ src/graph/observability/types.rs | 12 ++++++++---- src/graph/stream/mod.rs | 8 ++++++++ tests/e2e_misc_public_helpers.rs | 2 ++ 7 files changed, 64 insertions(+), 13 deletions(-) diff --git a/src/graph/compiled/mod.rs b/src/graph/compiled/mod.rs index 4ce32b2..a3acdf1 100644 --- a/src/graph/compiled/mod.rs +++ b/src/graph/compiled/mod.rs @@ -416,7 +416,17 @@ impl CompiledGraph { fn emit(&self, event: GraphEvent) { if let Some(sink) = &self.event_sink { + // Durable sinks persist asynchronously off the executor thread. On a + // terminal run event, flush so a caller that reads the journal right + // after the run returns sees a complete log. + let terminal = matches!( + event, + GraphEvent::RunCompleted { .. } | GraphEvent::RunFailed { .. } + ); sink.emit(event); + if terminal { + sink.flush(); + } } } } diff --git a/src/graph/observability/README.md b/src/graph/observability/README.md index bd20ea2..0a218c8 100644 --- a/src/graph/observability/README.md +++ b/src/graph/observability/README.md @@ -47,10 +47,14 @@ so existing runs are unchanged. ## Persistence bridge `JournalGraphSink` 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. Do not rely on the -sink for delivery guarantees stronger than "usually persisted, never -run-blocking." +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 diff --git a/src/graph/observability/mod.rs b/src/graph/observability/mod.rs index 9fedcfb..b955e48 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; @@ -54,6 +58,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, 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 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/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/tests/e2e_misc_public_helpers.rs b/tests/e2e_misc_public_helpers.rs index 75c1c50..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); From 4914afe56c3b33dac438014bcc2745e7c7b51f82 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Fri, 3 Jul 2026 23:00:33 -0700 Subject: [PATCH 056/106] fix(middleware): insert compression summary after the system prompt MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `ContextCompressionMiddleware` prepended the summary of elided older turns at index 0, displacing a leading system prompt out of first position — dropping its persistent instructions from their privileged slot and churning the cacheable prefix on every compaction. `plan` returns `to_keep` as `[system prompts..., recent turns...]`; insert the summary after the leading system prompts instead, so the real system prompt stays at position 0 and the summary sits in chronological order before the kept recent turns. Add a compression-with-system-prompt test. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/harness/middleware/mod.rs | 18 ++++++++++-- src/harness/middleware/test.rs | 52 ++++++++++++++++++++++++++++++++++ 2 files changed, 67 insertions(+), 3 deletions(-) diff --git a/src/harness/middleware/mod.rs b/src/harness/middleware/mod.rs index 03b53c9..7c34939 100644 --- a/src/harness/middleware/mod.rs +++ b/src/harness/middleware/mod.rs @@ -690,7 +690,7 @@ impl Middleware for ContextCom return Ok(()); } - let (to_summarize, to_keep) = self.policy.plan(&request.messages); + 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() { @@ -700,9 +700,21 @@ impl Middleware for ContextCom 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); + // `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(to_keep); + new_messages.extend(recent); let to_tokens = total_message_tokens(&new_messages); self.records diff --git a/src/harness/middleware/test.rs b/src/harness/middleware/test.rs index ea937af..0a04f8a 100644 --- a/src/harness/middleware/test.rs +++ b/src/harness/middleware/test.rs @@ -356,6 +356,58 @@ async fn context_compression_compresses_at_or_above_threshold() { assert!(compressed[0].1 > 0); } +#[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 From a4fbc11b2d3e85bfaede29167d32fe8850c3386f Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Fri, 3 Jul 2026 23:03:55 -0700 Subject: [PATCH 057/106] perf(channel): dedup Messages merge via id->index map MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The `Messages` channel merged incoming writes by scanning the whole existing list for each incoming message — O(existing x incoming) per step, which bit at a few thousand messages. Build an id->index map over the existing list once, then upsert each incoming message in O(1). Semantics are unchanged: keyed messages replace the same id in place (or append), unkeyed messages always append, and existing order is preserved. The `merge` clone-per-step of the current value is inherent to the public `Channel::merge(&self, current: Option<&Value>, ...)` signature; avoiding it would require an owned-`current` signature change (semver-major on the exported `Channel` trait) and is left as a separate follow-up. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/graph/channel/mod.rs | 33 ++++++++++++++++++++++++++------- src/graph/channel/test.rs | 38 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 64 insertions(+), 7 deletions(-) diff --git a/src/graph/channel/mod.rs b/src/graph/channel/mod.rs index d79635b..ddc1e2b 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; @@ -144,13 +144,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), } } diff --git a/src/graph/channel/test.rs b/src/graph/channel/test.rs index 37f480e..c89f3a6 100644 --- a/src/graph/channel/test.rs +++ b/src/graph/channel/test.rs @@ -59,6 +59,44 @@ 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; From 2ec6771729ef33a7e6c9bf812896db293729a3f0 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Fri, 3 Jul 2026 23:10:01 -0700 Subject: [PATCH 058/106] perf(checkpoint): single-pass state history and lean file get MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `get_state_history` walked the parent lineage by calling `get_tuple` per hop, and `FileCheckpointer::get` re-read and fully deserialized the entire thread file on every lookup — so history over a file-backed thread was O(H²). Add a `Checkpointer::state_history` method (default: the existing per-hop walk, behavior-preserving for in-memory/sqlite) and route `get_state_history` through it. Override it in `FileCheckpointer` to read the thread once and walk the lineage in memory (O(H)). Also make `FileCheckpointer::get` stream lines and fully decode only the single selected record instead of deserializing every record's state, keeping the previous last-match / latest selection semantics. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/graph/checkpoint/file.rs | 123 +++++++++++++++++++++++++-- src/graph/checkpoint/mod.rs | 42 +++++++++ src/graph/compiled/mod.rs | 27 ++---- tests/e2e_graph_support_contracts.rs | 55 ++++++++++++ 4 files changed, 221 insertions(+), 26 deletions(-) diff --git a/src/graph/checkpoint/file.rs b/src/graph/checkpoint/file.rs index 8715334..5d70f9f 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, @@ -178,12 +211,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 +256,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 795b96b..a07651f 100644 --- a/src/graph/checkpoint/mod.rs +++ b/src/graph/checkpoint/mod.rs @@ -130,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`, diff --git a/src/graph/compiled/mod.rs b/src/graph/compiled/mod.rs index a3acdf1..49d41fd 100644 --- a/src/graph/compiled/mod.rs +++ b/src/graph/compiled/mod.rs @@ -712,26 +712,13 @@ where 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) + // 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 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()); From 544060eae62939c6427a7b2abe0a7a67dbba3729 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Fri, 3 Jul 2026 23:12:31 -0700 Subject: [PATCH 059/106] perf(agent-loop): build tool schemas once per run, not per turn `ToolRegistry::schemas()` re-collects every tool's `schema()`, sorts, and allocates a fresh vec; the run loop called it on every turn (every model call). The tool set is fixed for a run, so hoist the sorted schema vec out of the loop and clone it per request. This removes the per-turn re-collect and re-sort; the remaining per-turn clone is inherent to rebuilding `ModelRequest`, and fully eliminating the request/history clone family would require an `invoke(&ModelRequest)` signature on the public `ChatModel` trait (a semver-major change), left as a separate follow-up. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/harness/agent_loop/mod.rs | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/src/harness/agent_loop/mod.rs b/src/harness/agent_loop/mod.rs index 5d86928..6d7bc62 100644 --- a/src/harness/agent_loop/mod.rs +++ b/src/harness/agent_loop/mod.rs @@ -321,6 +321,11 @@ impl AgentHarness { 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?; @@ -366,7 +371,7 @@ impl AgentHarness { // 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()); + 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()); } From 50e47901727b731bdc5f0b1b236045deb3114433 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Sat, 4 Jul 2026 01:12:10 -0700 Subject: [PATCH 060/106] perf(graph): roll max_concurrency window instead of join_all chunks Replace the fixed `join_all` chunking in parallel-branch execution with a rolling `select_all` window that starts a new branch as soon as any in-flight one finishes, matching parallel/mod.rs. Fixed chunks let one slow branch head-of-line block its whole chunk. Results are re-ordered back into deterministic order for the fold. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/graph/compiled/mod.rs | 48 +++++++++++++++++---- src/graph/compiled/test.rs | 85 +++++++++++++++++++++++++++++++++++++- 2 files changed, 123 insertions(+), 10 deletions(-) diff --git a/src/graph/compiled/mod.rs b/src/graph/compiled/mod.rs index 49d41fd..9dfae72 100644 --- a/src/graph/compiled/mod.rs +++ b/src/graph/compiled/mod.rs @@ -1780,25 +1780,55 @@ where ); let handler = node.handler.clone(); let owned_node = node_id.clone(); - futures.push(async move { + // 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 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; + 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); } - out.extend(futures::future::join_all(chunk).await); } - out + slots + .into_iter() + .map(|slot| slot.expect("every branch produced a result")) + .collect::>() } _ => futures::future::join_all(futures).await, }; diff --git a/src/graph/compiled/test.rs b/src/graph/compiled/test.rs index 72505d2..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)] @@ -1759,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)] From 487e03549fe25e894d19a4b6b76b50c087b87f94 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Sat, 4 Jul 2026 01:15:09 -0700 Subject: [PATCH 061/106] perf(memory): atomic bulk replace instead of clear-then-append Add `ChatHistory::replace`, overridden by the in-memory and store-backed backends to rewrite a thread in a single write. `ShortTermMemory::save` now delegates to it instead of clear-then-N-appends, which was O(n) writes and destroyed a thread's history if any append failed after the clear. A failing write now leaves the prior history intact. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/harness/memory/mod.rs | 36 +++++++++++++++---- src/harness/memory/test.rs | 70 +++++++++++++++++++++++++++++++++++++ src/harness/memory/types.rs | 15 ++++++++ 3 files changed, 114 insertions(+), 7 deletions(-) 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. From ba022aa19befba0f2e2f85e4e2a01e1feebc32a1 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Sat, 4 Jul 2026 01:19:13 -0700 Subject: [PATCH 062/106] perf(cache): LRU-bound response cache, stream-hash key, skip multi-turn - Bound InMemoryResponseCache with an LRU eviction policy (default 1024 entries) so a long-lived cache no longer grows without limit. - Stream cache_key's canonical JSON straight into the SHA-256 digest instead of allocating an intermediate Vec. - Skip caching requests whose transcript already contains a prior assistant or tool turn: those keys are unique per call and can never be re-served, so caching them only burns hashing cost and fills the cache with dead entries. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/harness/agent_loop/mod.rs | 12 +++++ src/harness/agent_loop/test.rs | 38 +++++++++++++++ src/harness/cache/mod.rs | 84 +++++++++++++++++++++++++++------- src/harness/cache/test.rs | 29 ++++++++++++ src/harness/cache/types.rs | 23 ++++++++-- 5 files changed, 166 insertions(+), 20 deletions(-) diff --git a/src/harness/agent_loop/mod.rs b/src/harness/agent_loop/mod.rs index 6d7bc62..fb5af74 100644 --- a/src/harness/agent_loop/mod.rs +++ b/src/harness/agent_loop/mod.rs @@ -727,6 +727,18 @@ impl AgentHarness { 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 3fd7357..7fd452e 100644 --- a/src/harness/agent_loop/test.rs +++ b/src/harness/agent_loop/test.rs @@ -1421,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 ───────────────────────────────────────────────────────── From 9142c3394d540497f8e4908d7574bda877e095b2 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Sat, 4 Jul 2026 01:22:03 -0700 Subject: [PATCH 063/106] perf(summarization): linear token trimming without per-message String allocs - Add `Message::char_len` to count text-block characters without allocating the concatenated string, and use it for token estimation. - Rewrite `trim_messages(MaxTokens)` to precompute per-message estimates once and drop from the front by advancing an index, replacing the O(n^2) re-sum-plus-remove(0) loop. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/harness/message/mod.rs | 20 ++++++++++++ src/harness/message/test.rs | 19 +++++++++++ src/harness/summarization/mod.rs | 54 +++++++++++++++++++------------- 3 files changed, 71 insertions(+), 22 deletions(-) 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/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 } } From 5978894fbb1aca1ab0dfd665c8743424acf11a2b Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Sat, 4 Jul 2026 01:28:55 -0700 Subject: [PATCH 064/106] perf(channel): merge takes current by value to fold aggregates in place Change `Channel::merge` to take `current: Option` instead of `Option<&Value>` so Topic/Messages/Barrier/NamedBarrier/BinaryAggregate reuse the accumulated Vec/Map in place rather than cloning the entire accumulated value on every write (was O(existing) allocation per merge). `apply_update` moves the current value out of the set and into the merge; a rejected merge leaves the channel unset, matching the executor's reduce-error contract. BREAKING: `Channel::merge` signature changed (custom Channel impls must update). Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/graph/channel/mod.rs | 55 ++++++++++++++++++++----------- src/graph/channel/test.rs | 66 ++++++++++++++++++++++++++++---------- src/graph/channel/types.rs | 7 +++- 3 files changed, 92 insertions(+), 36 deletions(-) diff --git a/src/graph/channel/mod.rs b/src/graph/channel/mod.rs index ddc1e2b..f92847a 100644 --- a/src/graph/channel/mod.rs +++ b/src/graph/channel/mod.rs @@ -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(), @@ -190,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) } @@ -208,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) } @@ -233,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 { @@ -276,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(), @@ -339,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), } } @@ -404,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 c89f3a6..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, @@ -74,7 +74,7 @@ fn messages_merge_dedup_map_preserves_order_and_appends_unkeyed() { .unwrap(); let v = c .merge( - Some(&v), + Some(v), json!([ {"id": "a", "text": "1-updated"}, {"id": "c", "text": "3"}, @@ -84,7 +84,7 @@ fn messages_merge_dedup_map_preserves_order_and_appends_unkeyed() { .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"})) + .merge(Some(v), json!({"id": "c", "text": "3-updated"})) .unwrap(); assert_eq!( v, @@ -100,7 +100,7 @@ fn messages_merge_dedup_map_preserves_order_and_appends_unkeyed() { #[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()); } @@ -118,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()); } @@ -133,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)); } @@ -143,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))); } @@ -152,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()); } @@ -186,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 From 96898e0dc052a33265aa209c9500ec9757b1846d Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Sat, 4 Jul 2026 01:31:47 -0700 Subject: [PATCH 065/106] perf(openai): batch SSE line draining to avoid per-line alloc and O(n^2) shifts Rewrite SseState::drain_lines to scan the byte buffer with a moving offset and drain the consumed prefix once, instead of `drain(..=pos).collect()` per line which allocated a Vec per line and shifted the remaining buffer down on every line (O(n^2) over a chunk carrying many SSE events). Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/harness/providers/openai/mod.rs | 19 +++++++++++++---- src/harness/providers/openai/test.rs | 32 ++++++++++++++++++++++++++++ 2 files changed, 47 insertions(+), 4 deletions(-) diff --git a/src/harness/providers/openai/mod.rs b/src/harness/providers/openai/mod.rs index ac82444..39c8c15 100644 --- a/src/harness/providers/openai/mod.rs +++ b/src/harness/providers/openai/mod.rs @@ -1082,13 +1082,24 @@ impl SseState { /// 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) { - while let Some(pos) = self.buf.iter().position(|&b| b == b'\n') { - let line_bytes: Vec = self.buf.drain(..=pos).collect(); + // 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. - let line = String::from_utf8_lossy(&line_bytes).into_owned(); + // 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 diff --git a/src/harness/providers/openai/test.rs b/src/harness/providers/openai/test.rs index b8f457b..759815c 100644 --- a/src/harness/providers/openai/test.rs +++ b/src/harness/providers/openai/test.rs @@ -734,6 +734,38 @@ async fn sse_stream_reassembles_multibyte_char_split_across_chunks() { 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 From 717e6e7e9743a88278b696f8e149007c58119887 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Sat, 4 Jul 2026 01:39:45 -0700 Subject: [PATCH 066/106] perf(checkpoint): run durable put IO on the blocking pool Wrap the synchronous filesystem append (FileCheckpointer) and rusqlite insert (SqliteCheckpointer) in tokio::task::spawn_blocking so the step-critical checkpoint write no longer blocks a tokio worker thread. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/graph/checkpoint/file.rs | 30 +++++++++++------ src/graph/checkpoint/sqlite.rs | 61 +++++++++++++++++++++------------- 2 files changed, 56 insertions(+), 35 deletions(-) diff --git a/src/graph/checkpoint/file.rs b/src/graph/checkpoint/file.rs index 5d70f9f..feba4f0 100644 --- a/src/graph/checkpoint/file.rs +++ b/src/graph/checkpoint/file.rs @@ -191,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) } 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) } From 0d51b4c7d5cb9ab5529f128947d1df03484e5f7f Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Sat, 4 Jul 2026 01:43:05 -0700 Subject: [PATCH 067/106] fix(langfuse): slim metadata, valid span type, surface 207 failures - Observation metadata now carries only run lineage, offset, and event kind instead of the full event payload, which duplicated input/output already in the body and grew batch bytes O(turns^2) toward the ~3.5MB cap. - Emit tool observations as `span-create`; `tool-create` is not a valid Langfuse ingestion type and is rejected by older/self-hosted Langfuse, silently dropping every tool observation. - Surface `207 Multi-Status` partial failures (non-empty `errors`) as an error instead of reporting success. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/harness/observability/langfuse.rs | 37 ++++++++++++++++++++++++--- 1 file changed, 34 insertions(+), 3 deletions(-) diff --git a/src/harness/observability/langfuse.rs b/src/harness/observability/langfuse.rs index ba47698..183b009 100644 --- a/src/harness/observability/langfuse.rs +++ b/src/harness/observability/langfuse.rs @@ -227,6 +227,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 +320,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 +362,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, @@ -569,9 +590,19 @@ mod tests { 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]["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] From f78c9b91b0d46e014e2d246d9c8141c259fc2cb3 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Sat, 4 Jul 2026 01:47:26 -0700 Subject: [PATCH 068/106] perf(repl): drop redundant namespace snapshot clone per cell eval_cell snapshotted the pre-cell namespace and then cloned the whole name->debug-string map into the shared vars_snapshot, a third O(namespace-bytes) pass per cell. Move the single baseline snapshot into vars_snapshot and read it back for the change diff, so a cell pays one baseline snapshot plus the after snapshot rather than snapshot + clone + after. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/repl/session/mod.rs | 20 ++++++++++++++++---- src/repl/session/test.rs | 22 ++++++++++++++++++++++ 2 files changed, 38 insertions(+), 4 deletions(-) diff --git a/src/repl/session/mod.rs b/src/repl/session/mod.rs index 35b9ab7..6730143 100644 --- a/src/repl/session/mod.rs +++ b/src/repl/session/mod.rs @@ -408,13 +408,16 @@ impl ReplSession { .arm_deadline(self.policy.timeout.map(|d| start + d)); self.buffers.arm_output_limit(self.policy.max_output_bytes); - let before = self.variables.snapshot(); - // Expose the pre-cell namespace to `show_vars()`. + // 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. @@ -470,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, diff --git a/src/repl/session/test.rs b/src/repl/session/test.rs index e57f72f..2d59cc3 100644 --- a/src/repl/session/test.rs +++ b/src/repl/session/test.rs @@ -43,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. From 29089f3c4fc0a553189cbe92499cfd0172a02dc6 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Sat, 4 Jul 2026 01:50:09 -0700 Subject: [PATCH 069/106] test(langfuse): assert tool observations export as span-create Follow-up to the Langfuse observation-type fix: the harness e2e export test now looks for the valid `span-create` tool observation type instead of the rejected `tool-create`. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- tests/e2e_observability.rs | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/e2e_observability.rs b/tests/e2e_observability.rs index ad99876..678e3e7 100644 --- a/tests/e2e_observability.rs +++ b/tests/e2e_observability.rs @@ -270,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"); From 139b10a5dd40935f3a277c9b1c22d7a0afc6e6a7 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Sat, 4 Jul 2026 02:09:57 -0700 Subject: [PATCH 070/106] fix(middleware): bound TracingMiddleware/ContextCompressionMiddleware/PromptCacheGuardMiddleware recorders These middlewares accumulated PhaseTrace/SummaryRecord/CacheLayoutEvent entries in an unbounded Vec for the lifetime of a long-running agent loop. Switch each recorder to a bounded VecDeque ring buffer (default cap 1024, configurable via with_max_records/with_max_events) so memory stays flat regardless of run length. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/harness/middleware/library/mod.rs | 39 ++++++++++--- src/harness/middleware/library/test.rs | 32 ++++++++++ src/harness/middleware/library/types.rs | 15 ++++- src/harness/middleware/mod.rs | 77 ++++++++++++++++++++----- src/harness/middleware/test.rs | 54 +++++++++++++++++ src/harness/middleware/types.rs | 15 ++++- 6 files changed, 208 insertions(+), 24 deletions(-) diff --git a/src/harness/middleware/library/mod.rs b/src/harness/middleware/library/mod.rs index 5d030da..24e9bc9 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}; @@ -1291,14 +1291,35 @@ impl TracingMiddleware { pub fn with_label(label: &'static str) -> Self { Self { label, - records: Mutex::new(Vec::new()), + 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").clone() + self.records + .lock() + .expect("records mutex poisoned") + .iter() + .cloned() + .collect() } /// Returns a snapshot of the per-phase begin counts. @@ -1307,10 +1328,14 @@ impl TracingMiddleware { } fn push(&self, phase: &'static str, boundary: TraceBoundary) { - self.records - .lock() - .expect("records mutex poisoned") - .push(PhaseTrace { phase, boundary }); + 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 }); } } diff --git a/src/harness/middleware/library/test.rs b/src/harness/middleware/library/test.rs index 8b25fce..a54071d 100644 --- a/src/harness/middleware/library/test.rs +++ b/src/harness/middleware/library/test.rs @@ -1219,3 +1219,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/types.rs b/src/harness/middleware/library/types.rs index 874a1e8..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}; @@ -464,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. /// @@ -474,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 7c34939..bea5fff 100644 --- a/src/harness/middleware/mod.rs +++ b/src/harness/middleware/mod.rs @@ -657,19 +657,39 @@ impl ContextCompressionMiddleware { label: "context_compression", policy, summarizer, - records: std::sync::Mutex::new(Vec::new()), + 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. Each record - /// carries the compression provenance for one compaction. + /// 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").clone() + self.records + .lock() + .expect("records mutex poisoned") + .iter() + .cloned() + .collect() } } @@ -717,10 +737,15 @@ impl Middleware for ContextCom new_messages.extend(recent); let to_tokens = total_message_tokens(&new_messages); - self.records - .lock() - .expect("records mutex poisoned") - .push(record); + { + 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 { @@ -740,13 +765,34 @@ impl PromptCacheGuardMiddleware { Self { label: "prompt_cache_guard", previous: std::sync::Mutex::new(None), - events: std::sync::Mutex::new(Vec::new()), + 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").clone() + self.events + .lock() + .expect("events mutex poisoned") + .iter() + .cloned() + .collect() } } @@ -774,10 +820,13 @@ impl Middleware for PromptCach && !prev.is_prefix_stable_against(&layout) { let event = CacheLayoutEvent::new(prev, &layout); - self.events - .lock() - .expect("events mutex poisoned") - .push(event); + 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/test.rs b/src/harness/middleware/test.rs index 0a04f8a..4ef7896 100644 --- a/src/harness/middleware/test.rs +++ b/src/harness/middleware/test.rs @@ -356,6 +356,40 @@ 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 @@ -495,6 +529,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 40ca649..fba1b49 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, } // ── PromptCacheGuardMiddleware ──────────────────────────────────────────────── @@ -502,10 +508,15 @@ pub struct ContextCompressionMiddleware { /// [`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 ───────────────────────────────────────────────── From 0628cb5db169d6935ed214156d16ba1dfeff0b0f Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Sat, 4 Jul 2026 02:17:00 -0700 Subject: [PATCH 071/106] fix(observability): bound InMemoryEventJournal and InMemoryStatusStore retention Both in-memory backends grew a run_id-keyed map forever with no eviction, so a long-lived process journaling/tracking many runs leaks memory without bound. Add a configurable max_runs cap (default 10_000): InMemoryEventJournal evicts the oldest run's stream on a new run once exceeded; InMemoryStatusStore evicts the oldest terminal run first, never evicting an active (Pending/Running/Interrupted) run to make room. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/harness/observability/mod.rs | 152 ++++++++++++++++++++++------- src/harness/observability/test.rs | 63 +++++++++++- src/harness/observability/types.rs | 71 ++++++++++++-- 3 files changed, 243 insertions(+), 43 deletions(-) diff --git a/src/harness/observability/mod.rs b/src/harness/observability/mod.rs index 66c9d7a..60ffcb5 100644 --- a/src/harness/observability/mod.rs +++ b/src/harness/observability/mod.rs @@ -37,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; @@ -149,16 +149,28 @@ 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()) .unwrap_or(0) @@ -168,27 +180,52 @@ 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 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 { + if let Some(oldest) = state.order.pop_front() { + state.streams.remove(&oldest); + } else { + break; + } + } + } + } + let entries = state.streams.entry(run_id).or_default(); let offset = entries.len() as u64; entries.push(obs); Ok(offset) } async fn read_from(&self, run_id: &str, offset: u64) -> Result> { - let runs = self - .runs + let state = self + .state .lock() .map_err(|e| poisoned("InMemoryEventJournal", e))?; - let Some(entries) = runs.get(run_id) else { + let Some(entries) = state.streams.get(run_id) else { return Ok(Vec::new()); }; Ok(entries.iter().skip(offset as usize).cloned().collect()) @@ -233,17 +270,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() } @@ -256,28 +315,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 @@ -289,11 +377,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() @@ -301,21 +390,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()) } diff --git a/src/harness/observability/test.rs b/src/harness/observability/test.rs index 135bb41..c6afe46 100644 --- a/src/harness/observability/test.rs +++ b/src/harness/observability/test.rs @@ -5,7 +5,7 @@ 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, @@ -488,3 +488,64 @@ 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 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 735fff8..35e60ba 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; @@ -195,14 +195,42 @@ 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; + +/// 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, +} + /// 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`]. @@ -257,13 +285,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, + } + } } // --------------------------------------------------------------------------- From 52826540c002302fec361ba3c092ebf643860141 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Sat, 4 Jul 2026 02:34:09 -0700 Subject: [PATCH 072/106] fix(observability): preserve durable offsets across journal eviction InMemoryEventJournal::append evicted a run's entire stream by FIFO insertion order with no regard for whether that run was still actively appending, unlike the sibling InMemoryStatusStore fix which never evicts in-flight runs. Because the evicted run's Vec was fully removed, a later append for the same run_id restarted numbering at offset 0, so a consumer resuming with a previously saved durable offset had read_from silently skip past its new entries instead of erroring. Track the next offset for evicted runs (bounded the same way `streams` is) so a resumed stream continues numbering, and make read_from return a clear error when a requested offset falls in an evicted range instead of silently truncating. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/harness/observability/mod.rs | 60 ++++++++++++++++++++++++------ src/harness/observability/test.rs | 54 +++++++++++++++++++++++++++ src/harness/observability/types.rs | 23 +++++++++++- 3 files changed, 125 insertions(+), 12 deletions(-) diff --git a/src/harness/observability/mod.rs b/src/harness/observability/mod.rs index 60ffcb5..8acbbd1 100644 --- a/src/harness/observability/mod.rs +++ b/src/harness/observability/mod.rs @@ -172,7 +172,7 @@ impl InMemoryEventJournal { .expect("InMemoryEventJournal lock poisoned") .streams .get(run_id) - .map(|v| v.len()) + .map(|s| s.entries.len()) .unwrap_or(0) } @@ -206,17 +206,35 @@ impl HarnessEventJournal for InMemoryEventJournal { // map without bound. `max_runs == 0` disables the cap. if self.max_runs > 0 { while state.order.len() > self.max_runs { - if let Some(oldest) = state.order.pop_front() { - state.streams.remove(&oldest); - } else { + 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 entries = state.streams.entry(run_id).or_default(); - let offset = entries.len() as u64; - entries.push(obs); + 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) } @@ -225,10 +243,30 @@ impl HarnessEventJournal for InMemoryEventJournal { .state .lock() .map_err(|e| poisoned("InMemoryEventJournal", e))?; - let Some(entries) = state.streams.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()) + } + } } } diff --git a/src/harness/observability/test.rs b/src/harness/observability/test.rs index c6afe46..dfc2b80 100644 --- a/src/harness/observability/test.rs +++ b/src/harness/observability/test.rs @@ -517,6 +517,60 @@ async fn in_memory_journal_evicts_oldest_run_once_max_runs_exceeded() { 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 diff --git a/src/harness/observability/types.rs b/src/harness/observability/types.rs index 35e60ba..ae23aff 100644 --- a/src/harness/observability/types.rs +++ b/src/harness/observability/types.rs @@ -199,14 +199,35 @@ pub trait HarnessEventJournal: Send + Sync { /// 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>, + 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`. From eac08972a20a77a6ec4ce7cc50989add59c19fd7 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Sat, 4 Jul 2026 02:35:33 -0700 Subject: [PATCH 073/106] docs: remove phantom openai feature and fix README metadata drift lib.rs claimed hosted providers live behind a nonexistent `openai` Cargo feature; they are compiled in unconditionally (only `sqlite` and `repl` are real optional features). README repeated the same claim, pointed at version 0.1 and the old rustagents clone URL, and was missing the goals_and_todos/subconscious_loop examples that exist on disk. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- README.md | 18 ++++++++++++------ src/lib.rs | 7 +++++-- 2 files changed, 17 insertions(+), 8 deletions(-) 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/src/lib.rs b/src/lib.rs index 8028938..a16dc4a 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 //! From 9f81efb865a8962a4b52c4475c677619f858dba7 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Sat, 4 Jul 2026 02:36:13 -0700 Subject: [PATCH 074/106] docs: fix AGENTS.md project structure to match current tree The Project Structure section still described src/chat.rs, src/model.rs, src/tool.rs, and src/graph.rs as flat modules, and the tests/ description said "currently focused on serialization behavior" though it now covers dozens of e2e/live suites. Rewrite to describe the five real module directories, the sqlite/repl Cargo features, docs/modules/, and the wiki submodule. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- AGENTS.md | 30 ++++++++++++++++++++++-------- 1 file changed, 22 insertions(+), 8 deletions(-) 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 From 167ceed1619e1b51d8296b6cfefa794b431a840f Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Sat, 4 Jul 2026 02:37:57 -0700 Subject: [PATCH 075/106] docs: fix stale RustAgents naming in audit.md Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- docs/audit.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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. From 4205019f5b89244b85d9ed9bdc4402b0d1bcb0bc Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Sat, 4 Jul 2026 02:37:57 -0700 Subject: [PATCH 076/106] docs: refresh ROADMAP.md against what has actually shipped ROADMAP.md still described chat/model/tool primitives and a state graph as aspirational "Current Foundation" and called the crate pre-1.0, even though v1.5.0 ships the full harness, graph, registry, .rag, and .ragsh surfaces. Rewrite each section to describe shipped functionality and real near-term work instead of the original bootstrap plan. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- ROADMAP.md | 96 +++++++++++++++++++++++------------------------------- 1 file changed, 40 insertions(+), 56 deletions(-) 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. From b787eb4dc7e2d845145745f291315a5bcc441b9a Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Sat, 4 Jul 2026 02:38:04 -0700 Subject: [PATCH 077/106] docs: split docs/spec/README.md and fix stale Package Layout/Milestones docs/spec/README.md had grown to 1121 lines, well past the repo's own 500-line limit, and its Package Layout described a flat src/chat.rs-style layout that hasn't existed for a long time; its Milestones and Open Questions sections described planning-stage work that has since shipped in full (v1.5.0). Extract the inline Harness/Graph/Expressive-Language module specs into docs/spec/harness-spec.md, docs/spec/graph-spec.md, and docs/spec/expressive-language-spec.md (linked from the README), and rewrite Package Layout to match the real src/ module-directory layout and Cargo features, and Milestones/Open Questions to reflect what has shipped. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- docs/spec/README.md | 1042 ++----------------------- docs/spec/expressive-language-spec.md | 137 ++++ docs/spec/graph-spec.md | 338 ++++++++ docs/spec/harness-spec.md | 441 +++++++++++ 4 files changed, 987 insertions(+), 971 deletions(-) create mode 100644 docs/spec/expressive-language-spec.md create mode 100644 docs/spec/graph-spec.md create mode 100644 docs/spec/harness-spec.md 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 + From ce123a668db7fc372f19fdc68cf30eb375d10492 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Sat, 4 Jul 2026 02:38:37 -0700 Subject: [PATCH 078/106] docs: relocate goal.md to docs/sdk-gaps.md goal.md lived at the repo root looking like a top-level planning document, but it is an internal OpenHuman-to-TinyAgents migration backlog. Move it under docs/ alongside the rest of the design notes and flag it clearly as an internal working document, pointing readers at ROADMAP.md for the public roadmap instead. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- goal.md => docs/sdk-gaps.md | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename goal.md => docs/sdk-gaps.md (100%) diff --git a/goal.md b/docs/sdk-gaps.md similarity index 100% rename from goal.md rename to docs/sdk-gaps.md From 4f30e6e5e1de3781e0d2feefe372633101fc4c4a Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Sat, 4 Jul 2026 02:38:51 -0700 Subject: [PATCH 079/106] docs: flag docs/sdk-gaps.md as an internal migration backlog Add a note clarifying this is an internal OpenHuman-to-TinyAgents migration backlog, not the public roadmap, and point readers at ROADMAP.md instead. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- docs/sdk-gaps.md | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/docs/sdk-gaps.md b/docs/sdk-gaps.md index 94fea22..85287be 100644 --- a/docs/sdk-gaps.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. From 1a185fef8f8f19a74bd73baa309c7b4e0b9ca0b3 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Sat, 4 Jul 2026 02:39:54 -0700 Subject: [PATCH 080/106] docs: split docs/modules/registry/design.md under the 500-line limit design.md had grown to 970 lines. Split the store/checkpointer/listener registration and event model/bus/filters into events.md, and the lifecycle/discovery/error-model/testkit/milestones sections into operations.md, linked from design.md and the module README. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- docs/modules/registry/README.md | 4 + docs/modules/registry/design.md | 509 +--------------------------- docs/modules/registry/events.md | 231 +++++++++++++ docs/modules/registry/operations.md | 289 ++++++++++++++++ 4 files changed, 528 insertions(+), 505 deletions(-) create mode 100644 docs/modules/registry/events.md create mode 100644 docs/modules/registry/operations.md diff --git a/docs/modules/registry/README.md b/docs/modules/registry/README.md index 01bf3db..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 ``` 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 From b396ff15ab953619819ae26dbad66dcb35f06d93 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Sat, 4 Jul 2026 02:40:30 -0700 Subject: [PATCH 081/106] docs: split docs/modules/harness/README.md under the 500-line limit README.md had grown to 941 lines. Split the tool registry/agent loop/middleware/memory sections into runtime.md, and structured output/events/errors/testkit/milestones into observability-overview.md, linked from the module README. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- docs/modules/harness/README.md | 481 +----------------- .../modules/harness/observability-overview.md | 157 ++++++ docs/modules/harness/runtime.md | 328 ++++++++++++ 3 files changed, 494 insertions(+), 472 deletions(-) create mode 100644 docs/modules/harness/observability-overview.md create mode 100644 docs/modules/harness/runtime.md 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/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. + From aaaf870739b3f76b3d08533638415039ff4b4e5d Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Sat, 4 Jul 2026 02:40:51 -0700 Subject: [PATCH 082/106] docs: split docs/modules/expressive-language/README.md under 500 lines README.md had grown to 814 lines. Split node kinds, binding to Rust, the state model, routes, policies, safety, examples, formatting, testkit, and milestones into reference.md, linked from README.md. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- docs/modules/expressive-language/README.md | 429 +---------------- docs/modules/expressive-language/reference.md | 432 ++++++++++++++++++ 2 files changed, 436 insertions(+), 425 deletions(-) create mode 100644 docs/modules/expressive-language/reference.md 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/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 From f23f9e1cbd0054070547d9ababa15967abd1d7d7 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Sat, 4 Jul 2026 02:41:18 -0700 Subject: [PATCH 083/106] docs: split docs/modules/repl-language/design.md under 500 lines design.md had grown to 807 lines. Split the RLM feature map, CodeAct loop, examples, Rhai embedding plan, Python compatibility backend, safety, events, diagnostics, testkit, and milestones into operations.md, linked from design.md and the module README. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- docs/modules/repl-language/README.md | 1 + docs/modules/repl-language/design.md | 346 +--------------------- docs/modules/repl-language/operations.md | 351 +++++++++++++++++++++++ 3 files changed, 356 insertions(+), 342 deletions(-) create mode 100644 docs/modules/repl-language/operations.md 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 From 5c3d4f8c2be6c2f2ff8a16a9237dce85339f4224 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Sat, 4 Jul 2026 02:41:57 -0700 Subject: [PATCH 084/106] docs: use https for the wiki submodule URL and document it .gitmodules pointed at an SSH URL (git@github.com:...), which requires SSH key setup even for read-only clones/CI. Switch to https and document the wiki/ submodule (what it is, how to fetch it, and that it should not be touched incidentally) in CONTRIBUTING.md. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- .gitmodules | 2 +- CONTRIBUTING.md | 25 +++++++++++++++++++++++++ 2 files changed, 26 insertions(+), 1 deletion(-) 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/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: From 3840112471b72c18e43ce4654a1c3fb7f82f1148 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Sat, 4 Jul 2026 02:50:46 -0700 Subject: [PATCH 085/106] fix(registry): re-export ModelCatalog/AliasBinding, fix docs, escape DOT ids - registry::AliasBinding was defined and used in the public RegistrySnapshot::aliases field but never re-exported from registry::mod or the crate root, so downstream code could observe it but not name its type. Re-export it alongside the rest of the registry surface. - Re-export the ModelCatalog family (ModelCatalog, ModelCatalogEntry, ModelCatalogSnapshot, ModelCatalogSource, ModelCapabilities, ModelPricing) at the crate root for parity with registry::mod, which already exposed them. - register_descriptor's doc comment claimed it was "reserved" for Store and Agent kinds, but Agent has its own dedicated register_agent method that does not call register_descriptor; the doc now names the actual fallback kinds (Store, Script, Middleware, Checkpointer, TaskStore, Listener). - RegistrySnapshot::to_dot interpolated component ids directly into Graphviz DOT quoted strings without escaping backslashes or double quotes, so a component id containing a `"` could break out of its quoted string and inject arbitrary DOT syntax. Escape both characters and add a regression test. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/lib.rs | 5 +++-- src/registry/capability/mod.rs | 11 ++++++++-- src/registry/diagnostics.rs | 39 ++++++++++++++++++++++++++++++++-- src/registry/mod.rs | 2 +- 4 files changed, 50 insertions(+), 7 deletions(-) diff --git a/src/lib.rs b/src/lib.rs index a16dc4a..f29ffcb 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -79,8 +79,9 @@ 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 --- diff --git a/src/registry/capability/mod.rs b/src/registry/capability/mod.rs index 92844bd..4d779c0 100644 --- a/src/registry/capability/mod.rs +++ b/src/registry/capability/mod.rs @@ -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/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}; From 4e292be3ccc5e3b20d131a026f6b37e40febc69e Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Sat, 4 Jul 2026 03:07:45 -0700 Subject: [PATCH 086/106] refactor(retry): share one retry-decision engine; fallback skips non-retryable Fold the per-attempt retry decision into a single `RetryPolicy::should_retry_error` (classification + attempt cap) used by both the agent loop and RetryMiddleware, so the retry logic that previously carried the same off-by-one in two places now lives in exactly one place and cannot drift. ModelFallbackMiddleware now stops switching models on a non-retryable error (auth/validation/schema) instead of burning every backend on a failure that will recur identically. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/harness/agent_loop/mod.rs | 12 +++++++++--- src/harness/middleware/library/mod.rs | 10 +++++++++- src/harness/middleware/library/test.rs | 26 ++++++++++++++++++++++++++ src/harness/retry/mod.rs | 16 ++++++++++++++++ src/harness/retry/test.rs | 17 +++++++++++++++++ 5 files changed, 77 insertions(+), 4 deletions(-) diff --git a/src/harness/agent_loop/mod.rs b/src/harness/agent_loop/mod.rs index fb5af74..51b5026 100644 --- a/src/harness/agent_loop/mod.rs +++ b/src/harness/agent_loop/mod.rs @@ -82,7 +82,6 @@ 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}; @@ -803,7 +802,8 @@ impl AgentHarness { /// 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 + /// 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 @@ -865,7 +865,13 @@ impl AgentHarness { .policy .retry .max_attempts_capped_at(self.policy.limits.max_retries_per_call); - if is_retryable(&error) && attempt + 1 < max_attempts { + // 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` diff --git a/src/harness/middleware/library/mod.rs b/src/harness/middleware/library/mod.rs index 24e9bc9..ca68cfb 100644 --- a/src/harness/middleware/library/mod.rs +++ b/src/harness/middleware/library/mod.rs @@ -93,7 +93,7 @@ impl ModelMiddleware for Retry match next.run(ctx, state, request.clone()).await { Ok(outcome) => return Ok(outcome), Err(error) => { - if is_retryable(&error) && self.policy.should_retry(attempt) { + 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` @@ -190,6 +190,14 @@ impl ModelMiddleware for Model 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(), diff --git a/src/harness/middleware/library/test.rs b/src/harness/middleware/library/test.rs index a54071d..c8e9696 100644 --- a/src/harness/middleware/library/test.rs +++ b/src/harness/middleware/library/test.rs @@ -260,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(); diff --git a/src/harness/retry/mod.rs b/src/harness/retry/mod.rs index 4bbb693..aea92c7 100644 --- a/src/harness/retry/mod.rs +++ b/src/harness/retry/mod.rs @@ -102,6 +102,22 @@ 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 diff --git a/src/harness/retry/test.rs b/src/harness/retry/test.rs index 5d174eb..adfcd7d 100644 --- a/src/harness/retry/test.rs +++ b/src/harness/retry/test.rs @@ -264,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)); +} From c3d0076b71066d69e205c3577aaa185945b565df Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Sat, 4 Jul 2026 03:26:10 -0700 Subject: [PATCH 087/106] refactor(subagent): one shared child-depth guard across recursion surfaces Extract the `depth + 1` / `max_depth` fail-closed check into a single `RunConfig::checked_child_depth`. SubAgent, SubAgentTool, and the REPL sub-run builtin previously hand-rolled the same guard three times; consolidating removes the drift that let one path bypass the depth limit. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/harness/context/mod.rs | 16 ++++++++++++++++ src/harness/context/test.rs | 15 +++++++++++++++ src/harness/subagent/mod.rs | 10 ++-------- src/repl/session/builtins.rs | 13 +++++-------- 4 files changed, 38 insertions(+), 16 deletions(-) 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/subagent/mod.rs b/src/harness/subagent/mod.rs index c1c8434..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 diff --git a/src/repl/session/builtins.rs b/src/repl/session/builtins.rs index f01f9c6..a8cb571 100644 --- a/src/repl/session/builtins.rs +++ b/src/repl/session/builtins.rs @@ -302,14 +302,11 @@ fn bump_agent(ctx: &HostContext) -> Result<(), Box(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(()) + // 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 ──────────────────────────────────────────────────────── From 96cfb71676a624fcee8d09502c5b26da37c6c2df Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Sat, 4 Jul 2026 03:28:24 -0700 Subject: [PATCH 088/106] refactor(ids): fold graph id counter into harness::ids single source Remove the graph-local `SEQ` AtomicU64 in `compiled` and draw every graph id (graph-*, subagent-*, testkit run ids) from `harness::ids::next_seq`, the same process-unique counter todos/goals/orchestration already use. Checkpoint ids keep delegating to `new_checkpoint_id` (restart-collision-free), so the durability fix is preserved. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/graph/builder/mod.rs | 2 +- src/graph/compiled/mod.rs | 8 -------- src/graph/subagent_node/mod.rs | 2 +- src/graph/testkit/mod.rs | 2 +- 4 files changed, 3 insertions(+), 11 deletions(-) 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/compiled/mod.rs b/src/graph/compiled/mod.rs index 9dfae72..b1a8347 100644 --- a/src/graph/compiled/mod.rs +++ b/src/graph/compiled/mod.rs @@ -70,7 +70,6 @@ pub use types::{CompiledGraph, GraphExecution, GraphInput, ResumeTarget, StateSn use std::collections::{HashMap, HashSet}; use std::sync::Arc; -use std::sync::atomic::{AtomicU64, Ordering}; use std::time::{Duration, SystemTime}; use crate::graph::builder::{ @@ -93,13 +92,6 @@ 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) -} - /// Allocates a fresh checkpoint id (string form) that is collision-free across /// process restarts. /// 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/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, From 06f7ec634137059243099931341de69d2df385e6 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Sat, 4 Jul 2026 03:38:54 -0700 Subject: [PATCH 089/106] refactor(time): single now_ms helper in harness::ids Collapse five identical `now_ms()` copies (goals store, graph + harness observability, langfuse, store) into one `harness::ids::now_ms`, alongside the existing process-time primitives. Every timestamp now flows through one epoch/`unwrap_or(0)` implementation. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/graph/goals/continuation.rs | 2 +- src/graph/goals/store.rs | 11 +---------- src/graph/observability/mod.rs | 10 +--------- src/harness/ids/mod.rs | 14 ++++++++++++++ src/harness/ids/test.rs | 12 ++++++++++++ src/harness/observability/langfuse.rs | 9 +-------- src/harness/observability/types.rs | 10 +--------- src/harness/store/mod.rs | 10 +--------- 8 files changed, 32 insertions(+), 46 deletions(-) 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/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/observability/mod.rs b/src/graph/observability/mod.rs index b955e48..3c047bf 100644 --- a/src/graph/observability/mod.rs +++ b/src/graph/observability/mod.rs @@ -57,7 +57,7 @@ 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; @@ -519,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/harness/ids/mod.rs b/src/harness/ids/mod.rs index 1b66b11..6384797 100644 --- a/src/harness/ids/mod.rs +++ b/src/harness/ids/mod.rs @@ -105,6 +105,20 @@ pub fn process_nonce() -> u64 { }) } +/// 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 { diff --git a/src/harness/ids/test.rs b/src/harness/ids/test.rs index cdeae05..2ecd603 100644 --- a/src/harness/ids/test.rs +++ b/src/harness/ids/test.rs @@ -66,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/observability/langfuse.rs b/src/harness/observability/langfuse.rs index 183b009..2ea18b6 100644 --- a/src/harness/observability/langfuse.rs +++ b/src/harness/observability/langfuse.rs @@ -9,6 +9,7 @@ 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; @@ -433,14 +434,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; diff --git a/src/harness/observability/types.rs b/src/harness/observability/types.rs index ae23aff..410bc57 100644 --- a/src/harness/observability/types.rs +++ b/src/harness/observability/types.rs @@ -19,7 +19,7 @@ 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::ids::{CallId, EventId, RunId, now_ms}; use crate::harness::store::AppendStore; use super::worker::AppendWorker; @@ -414,11 +414,3 @@ pub struct JsonlSink { /// Background drain that appends records without blocking the run. pub(crate) worker: Arc>, } - -/// 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) -} diff --git a/src/harness/store/mod.rs b/src/harness/store/mod.rs index 497a252..9c24610 100644 --- a/src/harness/store/mod.rs +++ b/src/harness/store/mod.rs @@ -36,6 +36,7 @@ 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. @@ -382,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 { From 271f3dca8625f9594e30995022a33f588a637f09 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Sat, 4 Jul 2026 03:50:10 -0700 Subject: [PATCH 090/106] refactor(openai): extract shared post_json/send transport helpers Collapse the duplicated HTTP tails of invoke, stream, and list_models into `authorized` (bearer header), `send_checked` (send + transport-error + non-2xx status decoding, returning the raw response), and `post_json` (chat POST with timeout selection). Auth, timeout, and error/status handling now live in one place instead of being hand-repeated three times. Behavior is unchanged. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/harness/providers/openai/mod.rs | 128 +++++++++++++++------------- 1 file changed, 67 insertions(+), 61 deletions(-) diff --git a/src/harness/providers/openai/mod.rs b/src/harness/providers/openai/mod.rs index 39c8c15..07544dd 100644 --- a/src/harness/providers/openai/mod.rs +++ b/src/harness/providers/openai/mod.rs @@ -255,27 +255,13 @@ impl OpenAiModel { 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)) - })?; + .send_checked(self.authorized(self.client.get(&url)), "request", &url) + .await?; - 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::Provider(Box::new(error))); - } - let listing: ModelListWire = serde_json::from_str(&text)?; Ok(listing.data) } @@ -473,6 +459,65 @@ impl OpenAiModel { }) } + /// 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, @@ -1249,32 +1294,15 @@ impl ChatModel for OpenAiModel { /// [`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 mut builder = self - .client - .post(&url) - .header("Authorization", format!("Bearer {}", self.api_key)) - .json(&body); - if let Some(timeout) = request_timeout(request.timeout_ms, false) { - builder = builder.timeout(timeout); - } - let response = builder.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 response = self + .post_json(&body, request.timeout_ms, false, "request") + .await?; - 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::Provider(Box::new(error))); - } - let value: Value = serde_json::from_str(&text)?; parse_response(value) } @@ -1302,32 +1330,10 @@ impl ChatModel for OpenAiModel { 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 mut builder = self - .client - .post(&url) - .header("Authorization", format!("Bearer {}", self.api_key)) - .json(&body); - if let Some(timeout) = request_timeout(request.timeout_ms, true) { - builder = builder.timeout(timeout); - } - let response = builder.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::Provider(Box::new(error))); - } + 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. From 3926ca16d2ae2679d8c41f504082944e9849d35d Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Sat, 4 Jul 2026 03:57:39 -0700 Subject: [PATCH 091/106] refactor(middleware): factor stack-runners; balance Started/Completed on error MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replace seven near-identical stack-runner bodies (before/after agent, model, tool + on_tool_delta) with one `run_stack_hook!` macro that centralizes the event bracketing and on_error fan-out. The macro — and the two wrap-onion handlers — now emit `MiddlewareCompleted` on the error path too, so a hook that returns `Err` can no longer leave a dangling `MiddlewareStarted` with no matching `Completed` in the event stream. Adds a regression test and documents the balance guarantee in the module README. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/harness/middleware/README.md | 6 +- src/harness/middleware/mod.rs | 194 ++++++++++--------------------- src/harness/middleware/test.rs | 41 +++++++ 3 files changed, 105 insertions(+), 136 deletions(-) diff --git a/src/harness/middleware/README.md b/src/harness/middleware/README.md index a4f06f8..588209f 100644 --- a/src/harness/middleware/README.md +++ b/src/harness/middleware/README.md @@ -40,7 +40,11 @@ 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. +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 diff --git a/src/harness/middleware/mod.rs b/src/harness/middleware/mod.rs index bea5fff..9f6dca5 100644 --- a/src/harness/middleware/mod.rs +++ b/src/harness/middleware/mod.rs @@ -50,6 +50,36 @@ 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 ──────────────────────────────────────────────────────────────── impl AgentRun { @@ -134,23 +164,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 @@ -161,23 +176,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 @@ -188,23 +188,8 @@ 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 @@ -241,23 +226,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 @@ -268,23 +238,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 @@ -295,23 +250,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 @@ -322,23 +262,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, @@ -424,14 +349,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?, @@ -455,14 +380,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?, diff --git a/src/harness/middleware/test.rs b/src/harness/middleware/test.rs index 4ef7896..6fe1862 100644 --- a/src/harness/middleware/test.rs +++ b/src/harness/middleware/test.rs @@ -194,6 +194,47 @@ 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 From ae532d3bc6b86644cf6200266fcb1feedd2ac81f Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Sat, 4 Jul 2026 04:13:30 -0700 Subject: [PATCH 092/106] refactor(graph): split compiled/mod.rs into executor/state_api/routing compiled/mod.rs had grown to ~2k lines mixing the superstep executor, state-inspection/time-travel API, and step-routing logic in one impl block. Split mechanically into executor.rs (run/resume/execute loop), state_api.rs (get_state/update_state/fork_state), and routing.rs (route/route_completed/interrupt-durability), keeping mod.rs for shared helpers and the builder/configuration impl. No behavior change. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/graph/compiled/README.md | 5 +- src/graph/compiled/executor.rs | 1319 ++++++++++++++++++++++++ src/graph/compiled/mod.rs | 1650 +------------------------------ src/graph/compiled/routing.rs | 131 +++ src/graph/compiled/state_api.rs | 224 +++++ 5 files changed, 1681 insertions(+), 1648 deletions(-) create mode 100644 src/graph/compiled/executor.rs create mode 100644 src/graph/compiled/routing.rs create mode 100644 src/graph/compiled/state_api.rs diff --git a/src/graph/compiled/README.md b/src/graph/compiled/README.md index 5f9ad8a..0b878d6 100644 --- a/src/graph/compiled/README.md +++ b/src/graph/compiled/README.md @@ -122,7 +122,10 @@ Accessors: `graph_id()`, `name()`, `namespace()`. | File | Role | | --- | --- | | `types.rs` | `CompiledGraph`, `GraphExecution`, `GraphInput`, `StateSnapshot`, `ResumeTarget`. | -| `mod.rs` | The superstep loop, retry/resumable-failure machinery, run/resume/state APIs. | +| `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 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 b1a8347..32aef11 100644 --- a/src/graph/compiled/mod.rs +++ b/src/graph/compiled/mod.rs @@ -64,6 +64,9 @@ //! [`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}; @@ -423,1652 +426,5 @@ impl CompiledGraph { } } -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. - 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()?; - // 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) - } - - #[allow(clippy::too_many_arguments)] - 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). - 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. - 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(()) - } - - 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, - 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 - } -} - #[cfg(test)] mod test; 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) + } +} From 5c1730929cb3ab3f7eb87aa84b5fe7c734d76873 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Sat, 4 Jul 2026 04:22:24 -0700 Subject: [PATCH 093/106] refactor(harness): split middleware/library/mod.rs by concern library/mod.rs mixed 14 built-in middleware across 4 unrelated concerns in ~1450 lines. Split mechanically into resilience.rs (retry/timeout/fallback/rate-limit), budget.rs (token/cost tracking and enforcement), tool_policy.rs (allowlisting, policy, dynamic/ contextual selection, human approval), and observe.rs (structured- output validation, dynamic prompt, redaction, tracing). No behavior change. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/harness/middleware/README.md | 2 +- src/harness/middleware/library/budget.rs | 346 ++++ src/harness/middleware/library/mod.rs | 1397 +---------------- src/harness/middleware/library/observe.rs | 351 +++++ src/harness/middleware/library/resilience.rs | 256 +++ src/harness/middleware/library/tool_policy.rs | 469 ++++++ 6 files changed, 1427 insertions(+), 1394 deletions(-) create mode 100644 src/harness/middleware/library/budget.rs create mode 100644 src/harness/middleware/library/observe.rs create mode 100644 src/harness/middleware/library/resilience.rs create mode 100644 src/harness/middleware/library/tool_policy.rs diff --git a/src/harness/middleware/README.md b/src/harness/middleware/README.md index 588209f..f5fe5c9 100644 --- a/src/harness/middleware/README.md +++ b/src/harness/middleware/README.md @@ -101,7 +101,7 @@ cause or replace it with a different error. | --- | --- | | `types.rs` | Every public type: traits, `MiddlewareStack`, built-in middleware structs. | | `mod.rs` | Behavioral code: `AgentRun` helpers, the stack runner, built-in `Middleware` impls. | -| `library/` | Constructors and additional impls for the built-in middleware. | +| `library/` | Constructors and additional impls for the 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), `observe.rs` (structured-output validation, dynamic prompt, redaction, tracing). | | `test.rs` | Unit tests (ordering, short-circuiting, each built-in middleware). | ## Operational constraints 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/mod.rs b/src/harness/middleware/library/mod.rs index ca68cfb..9634e3a 100644 --- a/src/harness/middleware/library/mod.rs +++ b/src/harness/middleware/library/mod.rs @@ -49,1399 +49,10 @@ 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 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 - } -} - -// ── 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() - } - - /// 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. - 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(()) - } -} - -// ── 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(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(()) - } -} +mod budget; +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..e3df4b4 --- /dev/null +++ b/src/harness/middleware/library/observe.rs @@ -0,0 +1,351 @@ +//! 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::*; + +// ── 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(()) + } +} 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/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(()) + } +} From 02590351ce22da9b3d5d012ab5aeffab3cb3f055 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Sat, 4 Jul 2026 04:30:50 -0700 Subject: [PATCH 094/106] refactor(harness): split agent_loop/mod.rs into entry/run_loop/model_call agent_loop/mod.rs had grown to ~1160 lines mixing public invoke* entry points, the ~400-line run_loop body, and model-invocation dispatch in one impl block. Split mechanically into entry.rs (invoke*/invoke_streaming*/drive), run_loop.rs (the core loop body and response-cache decision), and model_call.rs (cache-aware retry/ fallback dispatch plus the ModelBaseCall/ToolBaseCall impls). No behavior change. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/harness/agent_loop/README.md | 5 +- src/harness/agent_loop/entry.rs | 210 +++++ src/harness/agent_loop/mod.rs | 1072 +------------------------- src/harness/agent_loop/model_call.rs | 426 ++++++++++ src/harness/agent_loop/run_loop.rs | 462 +++++++++++ 5 files changed, 1105 insertions(+), 1070 deletions(-) create mode 100644 src/harness/agent_loop/entry.rs create mode 100644 src/harness/agent_loop/model_call.rs create mode 100644 src/harness/agent_loop/run_loop.rs diff --git a/src/harness/agent_loop/README.md b/src/harness/agent_loop/README.md index 97786e7..b4da19d 100644 --- a/src/harness/agent_loop/README.md +++ b/src/harness/agent_loop/README.md @@ -77,7 +77,10 @@ error surfaced by a model, tool, middleware, or structured-output extraction. | File | Role | | --- | --- | -| `mod.rs` | `AgentHarness::invoke` / `invoke_with_status` and the loop body. | +| `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). | 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 51b5026..45a63ef 100644 --- a/src/harness/agent_loop/mod.rs +++ b/src/harness/agent_loop/mod.rs @@ -88,1075 +88,9 @@ 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); - - // 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. - 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))) - } - - /// 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. - 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"`). - 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. -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))) + } +} From 30607852bc0f68348fc5b620dbd3445e87181fc9 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Sat, 4 Jul 2026 04:42:08 -0700 Subject: [PATCH 095/106] refactor(harness): split openai provider mod.rs by concern providers/openai/mod.rs had grown to ~1360 lines mixing HTTP transport, request/response conversion, and SSE stream parsing. Split mechanically into transport.rs (OpenAiModel construction, presets, request building, the ChatModel impl), convert.rs (message/response translation), and sse.rs (SseState/OpenAiStreamAcc/sse_next). Public API (OpenAiModel) re-exported from mod.rs unchanged. No behavior change. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/harness/providers/openai/README.md | 15 +- src/harness/providers/openai/convert.rs | 322 +++++ src/harness/providers/openai/mod.rs | 1299 +-------------------- src/harness/providers/openai/sse.rs | 343 ++++++ src/harness/providers/openai/transport.rs | 655 +++++++++++ 5 files changed, 1337 insertions(+), 1297 deletions(-) create mode 100644 src/harness/providers/openai/convert.rs create mode 100644 src/harness/providers/openai/sse.rs create mode 100644 src/harness/providers/openai/transport.rs diff --git a/src/harness/providers/openai/README.md b/src/harness/providers/openai/README.md index c8dbbc3..a0e711b 100644 --- a/src/harness/providers/openai/README.md +++ b/src/harness/providers/openai/README.md @@ -21,9 +21,10 @@ It: `AssistantMessage`, `ToolCall`s, `Usage`, and finish reason — or, for streaming, a `ModelStream` of `ModelStreamItem`s. -The wire (de)serialization shapes live in `types.rs`; `mod.rs` owns only the -translation logic and the HTTP transport, keeping OpenAI-specific JSON out of -the rest of the harness. +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 @@ -56,8 +57,7 @@ Together, Groq, OpenRouter, ...); returned ids can be fed straight into ## Streaming (SSE) Streaming responses are decoded by a small state machine (`SseState` / -`SseAccumulator`, roughly lines 845+ in `mod.rs`) built on -`futures::stream::unfold`: +`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 @@ -99,6 +99,9 @@ via `provider_failure_message`. Malformed JSON bodies surface as | File | Role | | --- | --- | -| `mod.rs` | `OpenAiModel`, translation logic, HTTP transport, SSE decoding. | +| `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 07544dd..68078a4 100644 --- a/src/harness/providers/openai/mod.rs +++ b/src/harness/providers/openai/mod.rs @@ -65,1299 +65,16 @@ const DEFAULT_CONNECT_TIMEOUT_SECS: u64 = 30; /// [`ModelRequest::timeout_ms`]. Streaming calls get no overall cap by default. const DEFAULT_REQUEST_TIMEOUT_SECS: u64 = 600; -/// 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, -} +mod convert; +mod sse; +mod transport; -/// Returns `true` for OpenAI o-series reasoning models (`o1`/`o3`/`o4`), which -/// reject `max_tokens` and require `max_completion_tokens` instead. -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") -} +pub use transport::OpenAiModel; -/// 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 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. - 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 - ) - } - - 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). -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`. -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()) -} - -/// 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. -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. -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. -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. -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, - }) -} - -/// 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. -fn tool_call_id(slot: usize, id: &str) -> String { - if id.is_empty() { - format!("tool-{slot}") - } else { - id.to_string() - } -} - -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`]. -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. -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. -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 -} - -// --------------------------------------------------------------------------- -// 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 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. -struct SseState { - /// Raw response byte chunks (errors already mapped onto the crate error). - 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. - buf: Vec, - /// 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 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`]. -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; - } - } - } -} - -#[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))) - } -} +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/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))) + } +} From ab3ad68b86b4727612452470c1da0c892f37bae5 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Sat, 4 Jul 2026 04:47:22 -0700 Subject: [PATCH 096/106] refactor(repl): split session/builtins.rs into capabilities/batched/authoring builtins.rs had grown to ~1175 lines covering five jobs (engine construction, single-shot capability calls, batched variants, graph authoring, and shared helpers). Converted to a builtins/ module directory and split mechanically into capabilities.rs (model_query/ tool_call/agent_query/graph_run), batched.rs (the *_batched variants), and authoring.rs (graph_define/validate/compile/diff/register), with mod.rs keeping engine construction and shared helpers. External path (session::builtins) unchanged. No behavior change. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/repl/session/builtins/authoring.rs | 231 +++++++ src/repl/session/builtins/batched.rs | 228 +++++++ src/repl/session/builtins/capabilities.rs | 168 +++++ .../session/{builtins.rs => builtins/mod.rs} | 609 +----------------- 4 files changed, 633 insertions(+), 603 deletions(-) create mode 100644 src/repl/session/builtins/authoring.rs create mode 100644 src/repl/session/builtins/batched.rs create mode 100644 src/repl/session/builtins/capabilities.rs rename src/repl/session/{builtins.rs => builtins/mod.rs} (51%) diff --git a/src/repl/session/builtins/authoring.rs b/src/repl/session/builtins/authoring.rs new file mode 100644 index 0000000..388f62e --- /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 ───────────────────────────────────────── + +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. +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)) +} diff --git a/src/repl/session/builtins/batched.rs b/src/repl/session/builtins/batched.rs new file mode 100644 index 0000000..62f8887 --- /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() +} + +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)) +} + +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)) +} + +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)) +} + +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..73e7c36 --- /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 ─────────────────────────────────────── + +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), + )) +} + +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)) + } +} + +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. +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.rs b/src/repl/session/builtins/mod.rs similarity index 51% rename from src/repl/session/builtins.rs rename to src/repl/session/builtins/mod.rs index a8cb571..5f624cd 100644 --- a/src/repl/session/builtins.rs +++ b/src/repl/session/builtins/mod.rs @@ -342,610 +342,13 @@ fn model_value(text: String, finish_reason: Option, structured: bool) -> } } -// ── Single capability implementations ─────────────────────────────────────── +mod authoring; +mod batched; +mod capabilities; -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), - )) -} - -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)) - } -} - -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. -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`"))?; - - // 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. -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_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)) -} - -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)) -} - -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)) -} - -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)) -} +use authoring::*; +use batched::*; +use capabilities::*; // ── Engine construction ───────────────────────────────────────────────────── From b8a6484921456753e51b833c6bf9b4ff35377024 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Sat, 4 Jul 2026 04:53:12 -0700 Subject: [PATCH 097/106] refactor(language): move CapabilityResolver out of compiler.rs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit CapabilityResolver, ReferenceClass, PrimaryReference, DEFAULT_NODE_KINDS, bind_capabilities, and bind_capabilities_with_registry move from compiler.rs into a new capability_resolver.rs sibling of resolver.rs, which already shares this binding policy (Resolver wraps a CapabilityResolver internally) — the binding-policy triplication the audit flagged is easier to see and keep in sync when the type lives next to its other consumer instead of buried in the compiler. Public API path updated at the crate root (lib.rs re-exports from the new module); internal callers (resolver.rs, registry/capability, tests, examples) updated to the new path. No behavior change. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- examples/openai_self_blueprint.rs | 5 +- examples/rag_blueprint.rs | 3 +- src/language/capability_resolver.rs | 437 ++++++++++++++++++++++++++ src/language/compiler.rs | 423 +------------------------ src/language/mod.rs | 1 + src/language/resolver.rs | 5 +- src/language/test.rs | 14 +- src/lib.rs | 6 +- src/registry/capability/mod.rs | 2 +- tests/e2e_graph_resolver_contracts.rs | 2 +- tests/e2e_rag_pipeline.rs | 5 +- tests/language_pipeline.rs | 5 +- 12 files changed, 462 insertions(+), 446 deletions(-) create mode 100644 src/language/capability_resolver.rs 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/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 50f91c2..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, @@ -444,428 +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, - /// 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) -} - /// 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/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/resolver.rs b/src/language/resolver.rs index af7cb05..9e57b9b 100644 --- a/src/language/resolver.rs +++ b/src/language/resolver.rs @@ -33,7 +33,10 @@ use crate::error::{Result, TinyAgentsError}; use crate::language::ast::{ChannelDecl, GraphDecl, NodeDecl, Program}; -use crate::language::compiler::{CapabilityResolver, DEFAULT_NODE_KINDS, ReferenceClass, 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; diff --git a/src/language/test.rs b/src/language/test.rs index d6c9974..cfbf006 100644 --- a/src/language/test.rs +++ b/src/language/test.rs @@ -4,9 +4,8 @@ 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::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}; @@ -775,7 +774,7 @@ fn extended_kinds_bind_against_a_resolver() { .allow_reducer("aggregate") .allow_reducer("barrier") .with_node_kinds( - crate::language::compiler::DEFAULT_NODE_KINDS + crate::language::capability_resolver::DEFAULT_NODE_KINDS .iter() .map(|k| k.to_string()), ); @@ -789,7 +788,7 @@ fn bind_blueprint_rejects_unregistered_subagent_and_script() { // 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::compiler::DEFAULT_NODE_KINDS + crate::language::capability_resolver::DEFAULT_NODE_KINDS .iter() .map(|k| k.to_string()) }; @@ -972,9 +971,8 @@ async fn build_graph_handles_linear_terminal() { // Registry-backed capability binding (registry → language binding) // --------------------------------------------------------------------------- -use crate::language::compiler::{ - DEFAULT_NODE_KINDS, bind_capabilities_with_registry, compile_source, -}; +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, diff --git a/src/lib.rs b/src/lib.rs index f29ffcb..ae389d0 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -88,10 +88,10 @@ pub use registry::{ // 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 diff --git a/src/registry/capability/mod.rs b/src/registry/capability/mod.rs index 4d779c0..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::*; 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_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/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}; From e140ea1a81862790680b40c70d04757499413984 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Sat, 4 Jul 2026 04:56:49 -0700 Subject: [PATCH 098/106] refactor(language): split test.rs into per-phase test modules MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit test.rs had grown to ~1568 lines covering the whole pipeline in one file. Converted to a test/ module directory and split by pipeline phase: lexer, diagnostics, parser, compiler, extended_grammar, capability_binding, graph_materialisation, registry_binding, resolver, and provenance_diff_testkit, with mod.rs keeping the shared SUPPORT_AGENT fixture and module declarations. No behavior or coverage change — same 76 tests, same assertions. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/language/test.rs | 1568 ------------------ src/language/test/capability_binding.rs | 45 + src/language/test/compiler.rs | 156 ++ src/language/test/diagnostics.rs | 210 +++ src/language/test/extended_grammar.rs | 277 ++++ src/language/test/graph_materialisation.rs | 113 ++ src/language/test/lexer.rs | 81 + src/language/test/mod.rs | 59 + src/language/test/parser.rs | 80 + src/language/test/provenance_diff_testkit.rs | 299 ++++ src/language/test/registry_binding.rs | 182 ++ src/language/test/resolver.rs | 138 ++ 12 files changed, 1640 insertions(+), 1568 deletions(-) delete mode 100644 src/language/test.rs create mode 100644 src/language/test/capability_binding.rs create mode 100644 src/language/test/compiler.rs create mode 100644 src/language/test/diagnostics.rs create mode 100644 src/language/test/extended_grammar.rs create mode 100644 src/language/test/graph_materialisation.rs create mode 100644 src/language/test/lexer.rs create mode 100644 src/language/test/mod.rs create mode 100644 src/language/test/parser.rs create mode 100644 src/language/test/provenance_diff_testkit.rs create mode 100644 src/language/test/registry_binding.rs create mode 100644 src/language/test/resolver.rs diff --git a/src/language/test.rs b/src/language/test.rs deleted file mode 100644 index cfbf006..0000000 --- a/src/language/test.rs +++ /dev/null @@ -1,1568 +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::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 - } -} -"#; - -// --------------------------------------------------------------------------- -// 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"); -} - -// --------------------------------------------------------------------------- -// 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:?}"), - } -} - -// --------------------------------------------------------------------------- -// 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}"); -} - -#[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}" - ); -} - -// --------------------------------------------------------------------------- -// 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}"); -} - -// --------------------------------------------------------------------------- -// 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::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. -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"); -} - -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. -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}"); +} From 69656bfdcba7c8e1c6160a87b7d200d5c6efe4c7 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Sat, 4 Jul 2026 05:04:25 -0700 Subject: [PATCH 099/106] refactor(harness,graph): convert no_progress.rs and langfuse.rs to module dirs Both flat files violated the dir/{mod.rs,types.rs,test.rs} convention. - harness/no_progress.rs -> no_progress/{mod.rs,types.rs,test.rs} (ToolAttempt/NoProgress/LadderState/NoProgressTracker type defs split from the escalation-ladder impl). - harness/observability/langfuse.rs -> langfuse/{mod.rs,types.rs,test.rs} (LangfuseAuth/LangfuseTraceConfig/LangfuseClient type defs split from the client impl and payload helpers). - graph/observability/langfuse/tests.rs renamed to test.rs to unify naming with the harness module and the repo-wide test.rs convention. Public API paths (re-exported from harness::no_progress and harness::observability) unchanged. No behavior change. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/graph/observability/langfuse.rs | 2 +- .../langfuse/{tests.rs => test.rs} | 0 .../{no_progress.rs => no_progress/mod.rs} | 208 +---------------- src/harness/no_progress/test.rs | 139 +++++++++++ src/harness/no_progress/types.rs | 65 ++++++ src/harness/observability/README.md | 2 +- .../{langfuse.rs => langfuse/mod.rs} | 221 +----------------- src/harness/observability/langfuse/test.rs | 163 +++++++++++++ src/harness/observability/langfuse/types.rs | 60 +++++ 9 files changed, 438 insertions(+), 422 deletions(-) rename src/graph/observability/langfuse/{tests.rs => test.rs} (100%) rename src/harness/{no_progress.rs => no_progress/mod.rs} (54%) create mode 100644 src/harness/no_progress/test.rs create mode 100644 src/harness/no_progress/types.rs rename src/harness/observability/{langfuse.rs => langfuse/mod.rs} (65%) create mode 100644 src/harness/observability/langfuse/test.rs create mode 100644 src/harness/observability/langfuse/types.rs diff --git a/src/graph/observability/langfuse.rs b/src/graph/observability/langfuse.rs index 13ba672..083b086 100644 --- a/src/graph/observability/langfuse.rs +++ b/src/graph/observability/langfuse.rs @@ -484,4 +484,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/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 index 1fe52b1..b993ec3 100644 --- a/src/harness/observability/README.md +++ b/src/harness/observability/README.md @@ -69,7 +69,7 @@ silently excluded from the rollup rather than reported with a bogus duration. | --- | --- | | `types.rs` | Every public type: `AgentObservation`, journal/status-store traits and in-memory impls, sink structs, latency types. | | `mod.rs` | Behavioral code: latency rollups, journal/store/sink impls. | -| `langfuse.rs` | `LangfuseClient` and payload helpers (`clean_nulls`, `iso_ms`) shared with `graph::observability::langfuse`. | +| `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 diff --git a/src/harness/observability/langfuse.rs b/src/harness/observability/langfuse/mod.rs similarity index 65% rename from src/harness/observability/langfuse.rs rename to src/harness/observability/langfuse/mod.rs index 2ea18b6..a260e4c 100644 --- a/src/harness/observability/langfuse.rs +++ b/src/harness/observability/langfuse/mod.rs @@ -4,7 +4,6 @@ //! 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}; @@ -13,58 +12,9 @@ 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, - }, -} - -/// 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, -} +mod types; -/// 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. @@ -460,169 +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"], "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"); - } -} +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, +} From 71428d0da5299e806928bb3d7a745ac3037e40b2 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Sat, 4 Jul 2026 05:07:27 -0700 Subject: [PATCH 100/106] refactor(harness): move MockModel impl out of providers/mod.rs providers/mod.rs carried ~275 lines of MockModel constructors, the ChatModel impl, and token-estimation helpers alongside the module's own wiring/doc comment. Moved into mock.rs, leaving mod.rs with just module declarations and shared imports. MockModel's type definition already lived in types.rs per convention. No behavior change. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/harness/providers/mock.rs | 282 ++++++++++++++++++++++++++++++++++ src/harness/providers/mod.rs | 274 +-------------------------------- 2 files changed, 283 insertions(+), 273 deletions(-) create mode 100644 src/harness/providers/mock.rs 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; From 380b7812969ebd0d8f237daf14709fa89f465792 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Sat, 4 Jul 2026 05:12:19 -0700 Subject: [PATCH 101/106] refactor(harness): move built-in middleware impls from mod.rs into library/ middleware/mod.rs hosted five full middleware impls (Logging, MessageTrim, ContextCompression, PromptCacheGuard, UsageAccounting) alongside the module's own stack-runner wiring, while every other built-in middleware already lived in library/. Moved MessageTrim/ContextCompression/PromptCacheGuard into a new library/context.rs and Logging/UsageAccounting into library/observe.rs (they fit the existing observation concern). mod.rs now holds only AgentRun helpers and the stack runner. No behavior change. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/harness/middleware/README.md | 4 +- src/harness/middleware/library/context.rs | 239 +++++++++++++ src/harness/middleware/library/mod.rs | 1 + src/harness/middleware/library/observe.rs | 193 ++++++++++ src/harness/middleware/mod.rs | 416 ---------------------- 5 files changed, 435 insertions(+), 418 deletions(-) create mode 100644 src/harness/middleware/library/context.rs diff --git a/src/harness/middleware/README.md b/src/harness/middleware/README.md index f5fe5c9..72f0d31 100644 --- a/src/harness/middleware/README.md +++ b/src/harness/middleware/README.md @@ -100,8 +100,8 @@ cause or replace it with a different error. | File | Role | | --- | --- | | `types.rs` | Every public type: traits, `MiddlewareStack`, built-in middleware structs. | -| `mod.rs` | Behavioral code: `AgentRun` helpers, the stack runner, built-in `Middleware` impls. | -| `library/` | Constructors and additional impls for the 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), `observe.rs` (structured-output validation, dynamic prompt, redaction, tracing). | +| `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 diff --git a/src/harness/middleware/library/context.rs b/src/harness/middleware/library/context.rs new file mode 100644 index 0000000..9caf209 --- /dev/null +++ b/src/harness/middleware/library/context.rs @@ -0,0 +1,239 @@ +//! 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, 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(()) + } +} + +// ── 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 9634e3a..339bfd0 100644 --- a/src/harness/middleware/library/mod.rs +++ b/src/harness/middleware/library/mod.rs @@ -50,6 +50,7 @@ use crate::harness::structured::{StructuredExtractor, StructuredStrategy}; use crate::harness::tool::{ToolCall, ToolDelta, ToolResult, ToolSchema}; mod budget; +mod context; mod observe; mod resilience; mod tool_policy; diff --git a/src/harness/middleware/library/observe.rs b/src/harness/middleware/library/observe.rs index e3df4b4..a2d410c 100644 --- a/src/harness/middleware/library/observe.rs +++ b/src/harness/middleware/library/observe.rs @@ -5,6 +5,10 @@ //! full built-in middleware library overview. use super::*; +use crate::harness::middleware::{ + AgentRun, HookCounts, LoggingMiddleware, UsageAccountingMiddleware, +}; +use crate::harness::usage::UsageTotals; // ── StructuredOutputValidatorMiddleware ─────────────────────────────────────── @@ -349,3 +353,192 @@ impl Middleware for TracingMid 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/mod.rs b/src/harness/middleware/mod.rs index 9f6dca5..294c704 100644 --- a/src/harness/middleware/mod.rs +++ b/src/harness/middleware/mod.rs @@ -37,18 +37,10 @@ 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::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 @@ -395,413 +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(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(()) - } -} - -// ── 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(()) - } -} - -// ── 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; From 95e9a3a55370ce0646252e91f6a20c6206ca3e72 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Sat, 4 Jul 2026 05:14:54 -0700 Subject: [PATCH 102/106] refactor(harness): move embeddings OpenAI provider out of inline mod embeddings/mod.rs embedded a full HTTP provider (OpenAiEmbeddingModel) as an inline `mod openai { ... }` block instead of its own file. Moved to openai.rs, declared via `mod openai;`, mirroring the convention used by harness::providers::openai. No behavior change. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/harness/embeddings/mod.rs | 180 +------------------------------ src/harness/embeddings/openai.rs | 179 ++++++++++++++++++++++++++++++ 2 files changed, 180 insertions(+), 179 deletions(-) create mode 100644 src/harness/embeddings/openai.rs 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 + } +} From 7dd08a8415eb99fd3c611811f14e72c9fe78501c Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Sat, 4 Jul 2026 05:18:01 -0700 Subject: [PATCH 103/106] refactor(harness): move workspace path-gating policy out of types.rs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit workspace/types.rs mixed plain type definitions with the allows/enforce path-gating logic and its lexical path-normalization helpers — the security-relevant behavior the module exists to provide. Moved into policy.rs, leaving types.rs with just WorkspaceDescriptor's fields and the WorkspaceIsolation trait. No behavior change. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/harness/workspace/mod.rs | 1 + src/harness/workspace/policy.rs | 121 ++++++++++++++++++++++++++++++++ src/harness/workspace/types.rs | 111 +---------------------------- 3 files changed, 123 insertions(+), 110 deletions(-) create mode 100644 src/harness/workspace/policy.rs 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/types.rs b/src/harness/workspace/types.rs index 8c70ce3..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,115 +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. 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 -} - /// Prepares and tears down per-agent execution environments. /// /// Implementations create a worktree/sandbox for one agent run and clean it up From d6a00a4c666dc3b280297a0536b5238e9941f88b Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Sat, 4 Jul 2026 05:21:12 -0700 Subject: [PATCH 104/106] refactor(graph): move goals prompt-rendering out of types.rs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit goals/types.rs mixed the ThreadGoal/ThreadGoalStatus/GoalProgress data model with active_goal_context_block, a prompt-rendering function — presentation, not state. Moved to prompt.rs. No behavior change. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/graph/goals/README.md | 3 ++- src/graph/goals/mod.rs | 6 +++--- src/graph/goals/prompt.rs | 35 +++++++++++++++++++++++++++++++++++ src/graph/goals/test.rs | 1 + src/graph/goals/types.rs | 28 ---------------------------- 5 files changed, 41 insertions(+), 32 deletions(-) create mode 100644 src/graph/goals/prompt.rs 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/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/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, - } -} From 8723fcb9b402bbf1d83cab518e32cb22494c2218 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Sat, 4 Jul 2026 05:22:55 -0700 Subject: [PATCH 105/106] refactor(graph): replace glob re-export in parallel/mod.rs with explicit list parallel/mod.rs did `pub use types::*;`; replaced with an explicit list (FailurePolicy, ItemOutcome, ParallelOptions, ParallelOutcome) so the module's public surface is discoverable without reading types.rs. No behavior change. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/graph/parallel/mod.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/graph/parallel/mod.rs b/src/graph/parallel/mod.rs index 03fbf66..d168c8e 100644 --- a/src/graph/parallel/mod.rs +++ b/src/graph/parallel/mod.rs @@ -16,7 +16,7 @@ mod types; -pub use types::*; +pub use types::{FailurePolicy, ItemOutcome, ParallelOptions, ParallelOutcome}; use futures::stream::StreamExt; use std::future::Future; From 2919cc47922d8ef9f1c4ed3bd2043384d73eab51 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Sat, 4 Jul 2026 05:45:07 -0700 Subject: [PATCH 106/106] fix(repl): make split builtins helpers pub(super) so the repl feature compiles The builtins split left the shared *_impl helpers private in the new capabilities/batched/authoring submodules; under --all-features (repl) the cross-module calls in mod.rs and batched.rs failed to resolve. Local gates ran without --all-features, so CI's clippy-all-features job caught it. Claude-Session: https://claude.ai/code/session_01Ca2v3JeHcZvpYYsEX4WpcN --- src/repl/session/builtins/authoring.rs | 12 ++++++------ src/repl/session/builtins/batched.rs | 8 ++++---- src/repl/session/builtins/capabilities.rs | 8 ++++---- 3 files changed, 14 insertions(+), 14 deletions(-) diff --git a/src/repl/session/builtins/authoring.rs b/src/repl/session/builtins/authoring.rs index 388f62e..f9ef98f 100644 --- a/src/repl/session/builtins/authoring.rs +++ b/src/repl/session/builtins/authoring.rs @@ -9,7 +9,7 @@ use super::*; // ── Graph-authoring implementations ───────────────────────────────────────── -fn graph_define_impl( +pub(super) fn graph_define_impl( ctx: &HostContext, params: &Map, ) -> Result> { @@ -119,7 +119,7 @@ fn draft_descriptor(handle: &GraphBlueprintHandle) -> Dynamic { } /// Looks up a graph draft by the `name` field of a descriptor map. -fn lookup_draft( +pub(super) fn lookup_draft( ctx: &HostContext, descriptor: &Map, func: &str, @@ -134,7 +134,7 @@ fn lookup_draft( .ok_or_else(|| invalid(ctx, format!("{func}: no graph draft named `{name}`"))) } -fn graph_validate_impl( +pub(super) fn graph_validate_impl( ctx: &HostContext, descriptor: &Map, ) -> Result> { @@ -148,7 +148,7 @@ fn graph_validate_impl( Ok(Dynamic::from_array(array)) } -fn graph_compile_impl( +pub(super) fn graph_compile_impl( ctx: &HostContext, descriptor: &Map, ) -> Result> { @@ -174,7 +174,7 @@ fn graph_compile_impl( Ok(draft_descriptor(&handle)) } -fn graph_diff_handles( +pub(super) fn graph_diff_handles( ctx: &HostContext, old: &Blueprint, new: &Blueprint, @@ -185,7 +185,7 @@ fn graph_diff_handles( Ok(repl_value_to_dynamic(&json_to_repl_value(&value))) } -fn graph_register_impl( +pub(super) fn graph_register_impl( ctx: &HostContext, params: &Map, ) -> Result> { diff --git a/src/repl/session/builtins/batched.rs b/src/repl/session/builtins/batched.rs index 62f8887..f47c176 100644 --- a/src/repl/session/builtins/batched.rs +++ b/src/repl/session/builtins/batched.rs @@ -24,7 +24,7 @@ fn batch_items( .collect() } -fn model_query_batched_impl( +pub(super) fn model_query_batched_impl( ctx: &HostContext, items: &Array, ) -> Result> { @@ -81,7 +81,7 @@ fn model_query_batched_impl( Ok(Dynamic::from_array(out)) } -fn tool_call_batched_impl( +pub(super) fn tool_call_batched_impl( ctx: &HostContext, items: &Array, ) -> Result> { @@ -159,7 +159,7 @@ fn tool_call_batched_impl( Ok(Dynamic::from_array(out)) } -fn agent_query_batched_impl( +pub(super) fn agent_query_batched_impl( ctx: &HostContext, items: &Array, ) -> Result> { @@ -215,7 +215,7 @@ fn agent_query_batched_impl( Ok(Dynamic::from_array(out)) } -fn graph_run_batched_impl( +pub(super) fn graph_run_batched_impl( ctx: &HostContext, items: &Array, ) -> Result> { diff --git a/src/repl/session/builtins/capabilities.rs b/src/repl/session/builtins/capabilities.rs index 73e7c36..64a9557 100644 --- a/src/repl/session/builtins/capabilities.rs +++ b/src/repl/session/builtins/capabilities.rs @@ -8,7 +8,7 @@ use super::*; // ── Single capability implementations ─────────────────────────────────────── -fn model_query_impl( +pub(super) fn model_query_impl( ctx: &HostContext, params: &Map, ) -> Result> { @@ -40,7 +40,7 @@ fn model_query_impl( )) } -fn tool_call_impl( +pub(super) fn tool_call_impl( ctx: &HostContext, params: &Map, ) -> Result> { @@ -85,7 +85,7 @@ fn tool_call_impl( } } -fn agent_query_impl( +pub(super) fn agent_query_impl( ctx: &HostContext, params: &Map, ) -> Result> { @@ -129,7 +129,7 @@ fn agent_query_impl( /// `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( +pub(super) fn graph_run_impl( ctx: &HostContext, params: &Map, ) -> Result> {