From 20493163645f17a69d48927a448142d8d060d3b3 Mon Sep 17 00:00:00 2001 From: Alex Clarke Date: Wed, 9 Sep 2026 18:05:52 -0600 Subject: [PATCH] feat(graph): accept a `{{template}}` for a map node's `max_concurrency` MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `MapNode.max_concurrency` becomes `Option`, where `ConcurrencyCap` is `Fixed(usize) | Template(String)`. Integers parse exactly as before (`4` -> `Fixed(4)`, `0` -> `Fixed(0)`, absent -> `None`); a string parses as a template that the map executor resolves against the parent state right before the fan-out via `StateManager::interpolate_raw`, mirroring `over:`. A template must resolve to an integer >= 1 or a numeric string (scripts often emit numbers as strings); anything else, including 0, is an error naming the map node, the template, and the resolved value/type. The fixed path keeps its `.max(1)` clamp for byte-compatibility. `Deserialize` is hand-written so a malformed value (`-1`, `4.5`, `true`, a list, a map) reports the offending value and the accepted forms rather than serde's opaque untagged-enum "did not match any variant". Validator: a string cap that contains no `{{ ... }}` pair is a load-time error (catches `"4"` typed for `4`); a templated cap joins the map node's read set so `validate_parallel_reads` flags a read of a sibling branch's write. The existing `Fixed(0)` rule and every other rule keep their order. Design: plans/map-subgraph-branches-design.md §3.5, §3.11 (C3, I3). --- src/graph/map.rs | 137 ++++++++++++++++++++++++++++++++++++++-- src/graph/types.rs | 123 +++++++++++++++++++++++++++++++++++- src/graph/validator.rs | 140 +++++++++++++++++++++++++++++++++++++++-- 3 files changed, 389 insertions(+), 11 deletions(-) diff --git a/src/graph/map.rs b/src/graph/map.rs index e541950..99325f9 100644 --- a/src/graph/map.rs +++ b/src/graph/map.rs @@ -3,7 +3,7 @@ use super::executor::StepContext; use super::llm::LlmNodeExecutor; use super::rag::RagNodeExecutor; use super::state::StateManager; -use super::types::{MapNode, NodeType}; +use super::types::{ConcurrencyCap, MapNode, NodeType}; use crate::config::{RenderMode, RequestContext}; use crate::graph::type_name; use crate::supervisor::mailbox::{Inbox, PeerAssignment, PeerRegistry, graph_agent_id}; @@ -49,10 +49,7 @@ impl MapNodeExecutor { })? .clone(); - let max_conc = node - .max_concurrency - .unwrap_or(step_ctx.max_concurrency) - .max(1); + let max_conc = resolve_max_concurrency(node, state, step_ctx.max_concurrency, node_id)?; let semaphore = Arc::new(Semaphore::new(max_conc)); let mut sub_tasks = Vec::with_capacity(items.len()); @@ -160,6 +157,42 @@ impl MapNodeExecutor { } } +/// A templated cap is resolved against the parent state right before the +/// fan-out. Scripts often emit numbers as strings, so a numeric string is +/// accepted; anything else — including 0 — is the author's bug and surfaces +/// as an error rather than being clamped. +fn resolve_max_concurrency( + node: &MapNode, + state: &StateManager, + default: usize, + node_id: &str, +) -> Result { + let template = match &node.max_concurrency { + None => return Ok(default.max(1)), + Some(ConcurrencyCap::Fixed(n)) => return Ok((*n).max(1)), + Some(ConcurrencyCap::Template(t)) => t, + }; + + let value = state + .interpolate_raw(template) + .with_context(|| format!("map node '{node_id}': evaluating `max_concurrency` template"))?; + + let resolved = match &value { + Value::Number(n) => n.as_u64().and_then(|n| usize::try_from(n).ok()), + Value::String(s) => s.trim().parse::().ok(), + _ => None, + }; + + match resolved { + Some(n) if n >= 1 => Ok(n), + _ => Err(anyhow!( + "map node '{node_id}': max_concurrency template \"{template}\" resolved to \ + {value} ({}); expected a positive integer", + type_name(&value) + )), + } +} + /// Pre-provision teammate identities before any branch starts, so a fast /// sibling can message one still waiting on the semaphore (the message /// queues in the pre-created inbox). A single-item fan-out has no peers, so @@ -238,4 +271,98 @@ mod tests { }); assert!(provision_map_peers(&branch, "shards", 3).is_none()); } + + fn map_with_cap(cap: Option) -> MapNode { + MapNode { + over: "{{items}}".into(), + as_name: "item".into(), + branch: "br".into(), + output_key: "output".into(), + collect_into: "results".into(), + max_concurrency: cap, + } + } + + fn state_with(key: &str, value: Value) -> StateManager { + StateManager::new(HashMap::from([(key.to_string(), value)])) + } + + fn resolve(cap: Option, state: &StateManager) -> Result { + resolve_max_concurrency(&map_with_cap(cap), state, 8, "m") + } + + #[test] + fn resolve_max_concurrency_fixed_uses_literal() { + let state = StateManager::new(HashMap::new()); + assert_eq!(resolve(Some(ConcurrencyCap::Fixed(4)), &state).unwrap(), 4); + } + + #[test] + fn resolve_max_concurrency_none_falls_back_to_step_default() { + let state = StateManager::new(HashMap::new()); + assert_eq!(resolve(None, &state).unwrap(), 8); + } + + #[test] + fn resolve_max_concurrency_template_accepts_integer_value() { + let state = state_with("budget", Value::from(3)); + let cap = Some(ConcurrencyCap::Template("{{budget}}".into())); + assert_eq!(resolve(cap, &state).unwrap(), 3); + } + + #[test] + fn resolve_max_concurrency_template_accepts_numeric_string() { + let state = state_with("budget", Value::from("3")); + let cap = Some(ConcurrencyCap::Template("{{budget}}".into())); + assert_eq!(resolve(cap, &state).unwrap(), 3); + } + + #[test] + fn resolve_max_concurrency_template_rejects_non_numeric_string() { + let state = state_with("budget", Value::from("high")); + let cap = Some(ConcurrencyCap::Template("{{budget}}".into())); + let msg = resolve(cap, &state).unwrap_err().to_string(); + assert!(msg.contains("map node 'm'"), "{msg}"); + assert!(msg.contains("\"{{budget}}\""), "{msg}"); + assert!(msg.contains("resolved to \"high\" (string)"), "{msg}"); + assert!(msg.contains("expected a positive integer"), "{msg}"); + } + + #[test] + fn resolve_max_concurrency_template_rejects_zero() { + let state = state_with("budget", Value::from(0)); + let cap = Some(ConcurrencyCap::Template("{{budget}}".into())); + let msg = resolve(cap, &state).unwrap_err().to_string(); + assert!(msg.contains("resolved to 0 (number)"), "{msg}"); + assert!(msg.contains("expected a positive integer"), "{msg}"); + } + + #[test] + fn resolve_max_concurrency_template_rejects_non_integer_number() { + let state = state_with("budget", Value::from(2.5)); + let cap = Some(ConcurrencyCap::Template("{{budget}}".into())); + let msg = resolve(cap, &state).unwrap_err().to_string(); + assert!(msg.contains("resolved to 2.5 (number)"), "{msg}"); + } + + #[test] + fn resolve_max_concurrency_template_rejects_array() { + let state = state_with("budget", Value::Array(vec![Value::from(3)])); + let cap = Some(ConcurrencyCap::Template("{{budget}}".into())); + let msg = resolve(cap, &state).unwrap_err().to_string(); + assert!(msg.contains("(array)"), "{msg}"); + } + + #[test] + fn resolve_max_concurrency_template_missing_key_names_map_node() { + let state = StateManager::new(HashMap::new()); + let cap = Some(ConcurrencyCap::Template("{{budget}}".into())); + let err = resolve(cap, &state).unwrap_err(); + let chain = format!("{err:#}"); + assert!( + chain.contains("map node 'm': evaluating `max_concurrency` template"), + "{chain}" + ); + assert!(chain.contains("budget"), "{chain}"); + } } diff --git a/src/graph/types.rs b/src/graph/types.rs index 7af4483..39f3143 100644 --- a/src/graph/types.rs +++ b/src/graph/types.rs @@ -426,13 +426,60 @@ pub struct MapNode { pub collect_into: String, #[serde(default, skip_serializing_if = "Option::is_none")] - pub max_concurrency: Option, + pub max_concurrency: Option, } fn default_map_output_key() -> String { "output".to_string() } +/// A map node's per-item concurrency cap: either a literal integer or a +/// `{{template}}` resolved against graph state when the map runs. +#[derive(Debug, Clone, Serialize, PartialEq)] +#[serde(untagged)] +pub enum ConcurrencyCap { + Fixed(usize), + Template(String), +} + +// Hand-written so a malformed value (`-1`, `4.5`, `true`) reports what was +// received and what is accepted, instead of serde's untagged-enum "did not +// match any variant". +impl<'de> Deserialize<'de> for ConcurrencyCap { + fn deserialize(deserializer: D) -> Result + where + D: serde::Deserializer<'de>, + { + struct CapVisitor; + + impl serde::de::Visitor<'_> for CapVisitor { + type Value = ConcurrencyCap; + + fn expecting(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { + f.write_str("a positive integer or a `{{template}}` string for max_concurrency") + } + + fn visit_u64(self, v: u64) -> Result { + usize::try_from(v) + .map(ConcurrencyCap::Fixed) + .map_err(|_| E::invalid_value(serde::de::Unexpected::Unsigned(v), &self)) + } + + fn visit_i64(self, v: i64) -> Result { + usize::try_from(v) + .map(ConcurrencyCap::Fixed) + .map_err(|_| E::invalid_value(serde::de::Unexpected::Signed(v), &self)) + } + + fn visit_str(self, v: &str) -> Result { + Ok(ConcurrencyCap::Template(v.to_owned())) + } + } + + deserializer.deserialize_any(CapVisitor) + } +} + #[derive(Debug, Clone, Copy, Deserialize, Serialize, PartialEq, Eq)] #[serde(rename_all = "lowercase")] pub enum Reducer { @@ -1232,7 +1279,7 @@ next: rank assert_eq!(map.branch, "research_subject"); assert_eq!(map.output_key, "research_result"); assert_eq!(map.collect_into, "research_results"); - assert_eq!(map.max_concurrency, Some(5)); + assert_eq!(map.max_concurrency, Some(ConcurrencyCap::Fixed(5))); } #[test] @@ -1256,6 +1303,78 @@ collect_into: results assert!(map.max_concurrency.is_none()); } + fn map_node_yaml_with_cap(cap: &str) -> String { + format!( + "id: fan_out\ntype: map\nover: \"{{{{items}}}}\"\nas: item\nbranch: process\n\ + collect_into: results\nmax_concurrency: {cap}\n" + ) + } + + fn parse_map_cap(cap: &str) -> Option { + let node: Node = serde_yaml::from_str(&map_node_yaml_with_cap(cap)).unwrap(); + match node.node_type { + NodeType::Map(m) => m.max_concurrency, + _ => panic!("expected Map variant"), + } + } + + #[test] + fn map_max_concurrency_integer_parses_as_fixed() { + assert_eq!(parse_map_cap("4"), Some(ConcurrencyCap::Fixed(4))); + } + + #[test] + fn map_max_concurrency_zero_parses_as_fixed_zero() { + assert_eq!(parse_map_cap("0"), Some(ConcurrencyCap::Fixed(0))); + } + + #[test] + fn map_max_concurrency_string_parses_as_template() { + assert_eq!( + parse_map_cap("\"{{budget}}\""), + Some(ConcurrencyCap::Template("{{budget}}".into())) + ); + } + + #[test] + fn map_max_concurrency_quoted_integer_parses_as_template() { + assert_eq!( + parse_map_cap("\"4\""), + Some(ConcurrencyCap::Template("4".into())) + ); + } + + #[test] + fn map_max_concurrency_malformed_values_name_accepted_forms() { + for bad in ["-1", "4.5", "true", "[4]", "{n: 4}"] { + let err = serde_yaml::from_str::(&map_node_yaml_with_cap(bad)) + .expect_err(&format!("max_concurrency: {bad} should fail to parse")); + let msg = err.to_string(); + assert!( + msg.contains("a positive integer or a `{{template}}` string for max_concurrency"), + "max_concurrency: {bad} produced an unhelpful error: {msg}" + ); + } + } + + #[test] + fn map_max_concurrency_negative_error_names_offending_value() { + let err = serde_yaml::from_str::(&map_node_yaml_with_cap("-1")).unwrap_err(); + assert!(err.to_string().contains("-1"), "{err}"); + } + + #[test] + fn concurrency_cap_serializes_as_bare_scalar() { + assert_eq!( + serde_json::to_value(ConcurrencyCap::Fixed(4)).unwrap(), + serde_json::json!(4) + ); + assert_eq!( + serde_json::to_value(ConcurrencyCap::Template("{{k}}".into())).unwrap(), + serde_json::json!("{{k}}") + ); + } + #[test] fn full_graph_with_all_new_phase_a_fields_parses() { let yaml = r#" diff --git a/src/graph/validator.rs b/src/graph/validator.rs index b900e14..1432d1a 100644 --- a/src/graph/validator.rs +++ b/src/graph/validator.rs @@ -1,5 +1,5 @@ use super::state::template_root_keys; -use super::types::{Graph, Node, NodeType}; +use super::types::{ConcurrencyCap, Graph, Node, NodeType}; use crate::client::{Model, ModelType}; use crate::config; use crate::config::{Agent, AppConfig, paths}; @@ -176,6 +176,7 @@ impl GraphValidator { self.validate_llm_nodes(graph, &mut result); self.validate_llm_skills(graph, &mut result); self.validate_max_concurrency(graph, &mut result); + self.validate_max_concurrency_template(graph, &mut result); self.validate_orchestration_limits(graph, &mut result); self.validate_map_branches(graph, &mut result); self.validate_parallel_user_interaction(graph, &mut result); @@ -519,7 +520,7 @@ impl GraphValidator { for (node_id, node) in &graph.nodes { if let NodeType::Map(m) = &node.node_type - && let Some(0) = m.max_concurrency + && let Some(ConcurrencyCap::Fixed(0)) = m.max_concurrency { result.error(ValidationError::with_node( node_id, @@ -530,6 +531,23 @@ impl GraphValidator { } } + // A string cap that never interpolates would only fail at run time, after the + // map's inputs were already computed; catch the typo (`"4"` for `4`) at load. + fn validate_max_concurrency_template(&self, graph: &Graph, result: &mut ValidationResult) { + for (node_id, node) in &graph.nodes { + if let NodeType::Map(m) = &node.node_type + && let Some(ConcurrencyCap::Template(t)) = &m.max_concurrency + && !contains_template(t) + { + result.error(ValidationError::with_node( + node_id, + "map node's `max_concurrency` is a string but not a template; write \ + an integer or a `{{key}}` template", + )); + } + } + } + fn validate_orchestration_limits(&self, graph: &Graph, result: &mut ValidationResult) { if graph.max_agent_depth == Some(0) { result.warning(ValidationError::new( @@ -940,11 +958,22 @@ fn primary_templated_fields(node: &Node) -> Vec { v } NodeType::End(n) => vec![n.output.clone()], - NodeType::Map(n) => vec![n.over.clone()], + NodeType::Map(n) => { + let mut v = vec![n.over.clone()]; + if let Some(ConcurrencyCap::Template(t)) = &n.max_concurrency { + v.push(t.clone()); + } + v + } NodeType::Script(_) => Vec::new(), } } +fn contains_template(s: &str) -> bool { + s.find("{{") + .is_some_and(|open| s[open + 2..].contains("}}")) +} + fn node_state_updates_map(node: &Node) -> Option<&std::collections::HashMap> { match &node.node_type { NodeType::Llm(n) => n.state_updates.as_ref(), @@ -2673,7 +2702,7 @@ mod tests { fn map_max_concurrency_zero_errors() { let mut map = map_node_basic("m", "br", Some("end")); if let NodeType::Map(ref mut mm) = map.node_type { - mm.max_concurrency = Some(0); + mm.max_concurrency = Some(ConcurrencyCap::Fixed(0)); } let branch = llm_with_state_updates("br", &[("output", "{{output}}")], None); let graph = graph_with( @@ -2714,6 +2743,109 @@ mod tests { ); } + fn map_with_cap(id: &str, branch: &str, next: Option<&str>, cap: &str) -> Node { + let mut map = map_node_basic(id, branch, next); + if let NodeType::Map(ref mut mm) = map.node_type { + mm.max_concurrency = Some(ConcurrencyCap::Template(cap.into())); + } + map + } + + #[test] + fn map_max_concurrency_non_template_string_errors() { + let map = map_with_cap("m", "br", Some("end"), "4"); + let branch = llm_with_state_updates("br", &[("output", "{{output}}")], None); + let graph = graph_with( + vec![("m", map), ("br", branch), ("end", end_node("end"))], + "m", + ); + + let result = validator().validate(&graph); + + assert!( + result.errors.iter().any(|e| e + .message + .contains("`max_concurrency` is a string but not a template") + && e.node_id.as_deref() == Some("m")), + "expected non-template string cap error: {:?}", + result.errors + ); + } + + #[test] + fn map_max_concurrency_template_string_is_valid() { + let map = map_with_cap("m", "br", Some("end"), "{{k}}"); + let branch = llm_with_state_updates("br", &[("output", "{{output}}")], None); + let graph = graph_with( + vec![("m", map), ("br", branch), ("end", end_node("end"))], + "m", + ); + + let result = validator().validate(&graph); + + assert!( + !result + .errors + .iter() + .any(|e| e.message.contains("max_concurrency")), + "templated cap should not error: {:?}", + result.errors + ); + } + + #[test] + fn map_max_concurrency_unclosed_template_errors() { + let map = map_with_cap("m", "br", Some("end"), "{{k"); + let branch = llm_with_state_updates("br", &[("output", "{{output}}")], None); + let graph = graph_with( + vec![("m", map), ("br", branch), ("end", end_node("end"))], + "m", + ); + + let result = validator().validate(&graph); + + assert!( + result + .errors + .iter() + .any(|e| e.message.contains("is a string but not a template")), + "expected unclosed-template cap error: {:?}", + result.errors + ); + } + + #[test] + fn map_max_concurrency_template_reading_sibling_write_errors() { + let mut start = end_node("start"); + start.next = Some(NextTargets::Many(vec!["m".into(), "worker_b".into()])); + let map = map_with_cap("m", "br", Some("end"), "{{summary}}"); + let branch = llm_with_state_updates("br", &[("output", "{{output}}")], None); + let writer = llm_with_state_updates("worker_b", &[("summary", "static")], Some("end")); + let graph = graph_with( + vec![ + ("start", start), + ("m", map), + ("br", branch), + ("worker_b", writer), + ("end", end_node("end")), + ], + "start", + ); + + let result = validator().validate(&graph); + + assert!( + result + .errors + .iter() + .any(|e| e.message.contains("reads state key(s) `summary`") + && e.message.contains("'worker_b'") + && e.node_id.as_deref() == Some("m")), + "expected cross-branch read error for templated cap: {:?}", + result.errors + ); + } + fn llm_with_prompt(id: &str, prompt: &str, next: Option<&str>) -> Node { let mut node = llm_node(id, None, next); if let NodeType::Llm(ref mut n) = node.node_type {