use std::collections::{BTreeMap, BTreeSet};
use mlua_flow_ir::{Expr, Node};
use crate::blueprint::Blueprint;
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct StepNameEntry {
pub canonical: String,
pub aliases: BTreeSet<String>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct StepNamingWarning {
pub name: String,
pub first_step_ref: String,
pub second_step_ref: String,
}
#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
#[error(
"StepNaming collision: name '{name}' is claimed by both step '{first_step_ref}' and step \
'{second_step_ref}' ({reason})"
)]
pub struct StepNamingError {
pub name: String,
pub first_step_ref: String,
pub second_step_ref: String,
pub reason: String,
}
#[derive(Debug, Clone, Default)]
pub struct StepNaming {
by_ref: BTreeMap<String, String>,
by_name: BTreeMap<String, String>,
entries: BTreeMap<String, StepNameEntry>,
}
impl StepNaming {
pub fn resolve(&self, name: &str) -> Option<&str> {
self.by_name.get(name).map(String::as_str)
}
pub fn canonical_of_producer(&self, ref_name: &str) -> Option<&str> {
self.by_ref.get(ref_name).map(String::as_str)
}
pub fn names(&self) -> impl Iterator<Item = &str> {
self.entries.keys().map(String::as_str)
}
pub fn entries(&self) -> impl Iterator<Item = &StepNameEntry> {
self.entries.values()
}
pub fn from_blueprint(
bp: &Blueprint,
) -> Result<(StepNaming, Vec<StepNamingWarning>), StepNamingError> {
let mut occurrences: Vec<(String, Option<(String, bool)>)> = Vec::new();
collect_steps(&bp.flow, &mut occurrences);
let mut order: Vec<String> = Vec::new();
let mut out_tops: BTreeMap<String, BTreeSet<(String, bool)>> = BTreeMap::new();
for (ref_, top) in occurrences {
let tops = out_tops.entry(ref_.clone()).or_default();
if let Some(top) = top {
tops.insert(top);
}
if !order.contains(&ref_) {
order.push(ref_);
}
}
let declared: BTreeMap<&str, &str> = bp
.agents
.iter()
.filter_map(|ad| {
let name = ad.meta.as_ref()?.projection_name.as_deref()?;
Some((ad.name.as_str(), name))
})
.collect();
let mut plans: Vec<StepClaims> = Vec::with_capacity(order.len());
for ref_ in &order {
let is_declared = declared.contains_key(ref_.as_str());
let canonical = declared
.get(ref_.as_str())
.map(|s| s.to_string())
.unwrap_or_else(|| ref_.clone());
let mut strong_aliases: BTreeSet<String> = BTreeSet::new();
strong_aliases.insert(ref_.clone());
let mut weak_aliases: BTreeSet<String> = BTreeSet::new();
for (top, nested) in out_tops.remove(ref_).unwrap_or_default() {
if nested {
weak_aliases.insert(top);
} else {
strong_aliases.insert(top);
}
}
weak_aliases.retain(|n| n != &canonical && !strong_aliases.contains(n));
plans.push(StepClaims {
ref_: ref_.clone(),
canonical,
is_declared,
strong_aliases,
weak_aliases,
});
}
let mut claimants: BTreeMap<&str, usize> = BTreeMap::new();
for plan in &plans {
for name in plan.claimed_strong().chain(plan.weak_aliases.iter()) {
*claimants.entry(name.as_str()).or_default() += 1;
}
}
let mut naming = StepNaming::default();
let mut warnings = Vec::new();
let mut claims: BTreeMap<String, (String, bool)> = BTreeMap::new();
for plan in &plans {
let StepClaims {
ref_,
canonical,
is_declared,
strong_aliases,
weak_aliases,
} = plan;
let (is_declared, canonical) = (*is_declared, canonical.clone());
let mut aliases: BTreeSet<String> = strong_aliases.clone();
for name in weak_aliases {
if claimants.get(name.as_str()).copied().unwrap_or(0) > 1 {
tracing::debug!(
name = %name,
step_ref = %ref_,
"StepNaming: dropping a contested weak (nesting-root) alias claim; \
address this step by its ref or projection_name instead"
);
continue;
}
aliases.insert(name.clone());
}
let mut claimed: BTreeSet<String> = aliases.clone();
claimed.insert(canonical.clone());
for name in &claimed {
match claims.get(name).cloned() {
None => {
claims.insert(name.clone(), (ref_.clone(), is_declared));
naming.by_name.insert(name.clone(), canonical.clone());
}
Some((other_ref, other_declared)) => {
if is_declared || other_declared {
return Err(StepNamingError {
name: name.clone(),
first_step_ref: other_ref,
second_step_ref: ref_.clone(),
reason: collision_reason(other_declared, is_declared),
});
}
warnings.push(StepNamingWarning {
name: name.clone(),
first_step_ref: other_ref.clone(),
second_step_ref: ref_.clone(),
});
if ref_ == name && &other_ref != name {
claims.insert(name.clone(), (ref_.clone(), false));
naming.by_name.insert(name.clone(), canonical.clone());
}
}
}
}
naming.by_ref.insert(ref_.clone(), canonical.clone());
naming
.entries
.insert(canonical.clone(), StepNameEntry { canonical, aliases });
}
Ok((naming, warnings))
}
}
struct StepClaims {
ref_: String,
canonical: String,
is_declared: bool,
strong_aliases: BTreeSet<String>,
weak_aliases: BTreeSet<String>,
}
impl StepClaims {
fn claimed_strong(&self) -> impl Iterator<Item = &String> {
self.strong_aliases.iter().chain(
std::iter::once(&self.canonical).filter(|c| !self.strong_aliases.contains(c.as_str())),
)
}
}
fn collision_reason(other_declared: bool, is_declared: bool) -> String {
match (other_declared, is_declared) {
(true, true) => "both sides declare projection_name".to_string(),
(true, false) => "the first step declares projection_name".to_string(),
(false, true) => "the second step declares projection_name".to_string(),
(false, false) => {
unreachable!("hard StepNamingError requires at least one declared side")
}
}
}
fn collect_steps(node: &Node, out: &mut Vec<(String, Option<(String, bool)>)>) {
match node {
Node::Step {
ref_,
out: out_expr,
..
} => {
out.push((ref_.clone(), out_alias(out_expr)));
}
Node::Seq { children } => {
for child in children {
collect_steps(child, out);
}
}
Node::Branch { then_, else_, .. } => {
collect_steps(then_, out);
collect_steps(else_, out);
}
Node::Fanout { body, .. } => collect_steps(body, out),
Node::Loop { body, .. } => collect_steps(body, out),
Node::Try { body, catch, .. } => {
collect_steps(body, out);
collect_steps(catch, out);
}
Node::Assign { .. } => {} }
}
fn out_alias(expr: &Expr) -> Option<(String, bool)> {
let Expr::Path { at } = expr else {
return None;
};
let rendered = at.to_string();
let trimmed = rendered
.strip_prefix("$.")
.or_else(|| rendered.strip_prefix('$'))?;
let mut segments = trimmed.split('.').filter(|s| !s.is_empty());
let top = segments.next()?.to_string();
Some((top, segments.next().is_some()))
}
#[cfg(test)]
mod tests {
use super::*;
use crate::blueprint::{
current_schema_version, AgentDef, AgentKind, AgentMeta, BlueprintMetadata, CompilerHints,
CompilerStrategy,
};
use mlua_flow_ir::JoinMode;
use serde_json::json;
fn path(s: &str) -> Expr {
Expr::Path {
at: s.parse().expect("literal test path"),
}
}
fn step(ref_: &str, out: &str) -> Node {
Node::Step {
ref_: ref_.to_string(),
in_: path("$.in"),
out: path(out),
}
}
fn agent(name: &str, projection_name: Option<&str>) -> AgentDef {
AgentDef {
name: name.to_string(),
kind: AgentKind::RustFn,
spec: json!({ "fn_id": name }),
profile: None,
meta: Some(AgentMeta {
projection_name: projection_name.map(str::to_string),
..Default::default()
}),
runner: None,
runner_ref: None,
verdict: None,
}
}
fn bp(flow: Node, agents: Vec<AgentDef>) -> Blueprint {
Blueprint {
schema_version: current_schema_version(),
id: "step-naming-ut".into(),
flow,
agents,
operators: vec![],
metas: vec![],
hints: CompilerHints::default(),
strategy: CompilerStrategy::default(),
metadata: BlueprintMetadata::default(),
spawner_hints: Default::default(),
default_agent_kind: AgentKind::Operator,
default_operator_kind: None,
default_init_ctx: None,
default_agent_ctx: None,
default_context_policy: None,
projection_placement: None,
audits: vec![],
degradation_policy: None,
runners: vec![],
default_runner: None,
subprocesses: vec![],
check_policy: None,
blueprint_ref_includes: Vec::new(),
}
}
#[test]
fn declared_step_canonical_is_projection_name_aliases_are_ref_and_out_top() {
let flow = step("planner", "$.plan");
let bp = bp(flow, vec![agent("planner", Some("plan-out"))]);
let (naming, warnings) = StepNaming::from_blueprint(&bp).expect("no collision");
assert!(warnings.is_empty());
assert_eq!(naming.canonical_of_producer("planner"), Some("plan-out"));
let entry = naming
.entries()
.find(|e| e.canonical == "plan-out")
.expect("entry present");
assert_eq!(
entry.aliases,
BTreeSet::from(["planner".to_string(), "plan".to_string()])
);
}
#[test]
fn undeclared_step_canonical_is_ref_aliases_are_ref_and_out_top() {
let flow = step("worker", "$.result");
let bp = bp(flow, vec![agent("worker", None)]);
let (naming, warnings) = StepNaming::from_blueprint(&bp).expect("no collision");
assert!(warnings.is_empty());
assert_eq!(naming.canonical_of_producer("worker"), Some("worker"));
let entry = naming
.entries()
.find(|e| e.canonical == "worker")
.expect("entry present");
assert_eq!(
entry.aliases,
BTreeSet::from(["worker".to_string(), "result".to_string()])
);
}
#[test]
fn ref_equal_to_out_top_collapses_to_single_alias_and_is_not_a_collision() {
let flow = step("scout", "$.scout");
let bp = bp(flow, vec![agent("scout", None)]);
let (naming, warnings) = StepNaming::from_blueprint(&bp).expect("no collision");
assert!(warnings.is_empty());
let entry = naming
.entries()
.find(|e| e.canonical == "scout")
.expect("entry present");
assert_eq!(entry.aliases, BTreeSet::from(["scout".to_string()]));
}
#[test]
fn declared_name_colliding_with_another_steps_ref_is_a_hard_error() {
let flow = Node::Seq {
children: vec![step("a", "$.a_out"), step("b", "$.b_out")],
};
let bp = bp(flow, vec![agent("a", Some("b")), agent("b", None)]);
let err = StepNaming::from_blueprint(&bp).expect_err("declared collision must reject");
assert_eq!(err.name, "b");
assert!(
err.reason.contains("declare"),
"reason should explain which side declared: {}",
err.reason
);
}
#[test]
fn undeclared_collision_is_ok_with_a_warning_and_data_plane_priority() {
let flow = Node::Seq {
children: vec![step("foo", "$.bar"), step("bar", "$.baz")],
};
let bp = bp(flow, vec![agent("foo", None), agent("bar", None)]);
let (naming, warnings) = StepNaming::from_blueprint(&bp).expect("soft collision is Ok");
assert_eq!(warnings.len(), 1);
assert_eq!(warnings[0].name, "bar");
assert_eq!(naming.resolve("bar"), Some("bar"));
}
#[test]
fn walk_covers_seq_branch_fanout_loop_and_try_nesting() {
let flow = Node::Seq {
children: vec![
step("in-seq", "$.a"),
Node::Branch {
cond: Expr::Lit { value: json!(true) },
then_: Box::new(step("in-then", "$.b")),
else_: Box::new(step("in-else", "$.c")),
},
Node::Fanout {
items: path("$.items"),
bind: path("$.item"),
body: Box::new(step("in-fanout", "$.d")),
join: JoinMode::All,
out: path("$.results"),
},
Node::Loop {
counter: path("$.n"),
cond: Expr::Lit { value: json!(true) },
body: Box::new(step("in-loop", "$.e")),
max: 3,
},
Node::Try {
body: Box::new(step("in-try", "$.f")),
catch: Box::new(step("in-catch", "$.g")),
err_at: None,
},
Node::Assign {
at: path("$.h"),
value: Expr::Lit { value: json!(1) },
},
],
};
let agents = vec![
"in-seq",
"in-then",
"in-else",
"in-fanout",
"in-loop",
"in-try",
"in-catch",
]
.into_iter()
.map(|n| agent(n, None))
.collect();
let bp = bp(flow, agents);
let (naming, warnings) = StepNaming::from_blueprint(&bp).expect("no collision");
assert!(warnings.is_empty());
let mut names: Vec<&str> = naming.names().collect();
names.sort_unstable();
assert_eq!(
names,
vec![
"in-catch",
"in-else",
"in-fanout",
"in-loop",
"in-seq",
"in-then",
"in-try",
]
);
}
#[test]
fn resolve_returns_canonical_for_alias_lookup() {
let flow = step("planner", "$.plan");
let bp = bp(flow, vec![agent("planner", Some("plan-out"))]);
let (naming, _) = StepNaming::from_blueprint(&bp).expect("no collision");
assert_eq!(naming.resolve("plan-out"), Some("plan-out"));
assert_eq!(naming.resolve("planner"), Some("plan-out"));
assert_eq!(naming.resolve("plan"), Some("plan-out"));
assert_eq!(naming.resolve("does-not-exist"), None);
}
#[test]
fn shared_nesting_root_is_claimed_by_nobody_and_warns_nowhere() {
let flow = Node::Seq {
children: vec![
step("lane-a", "$.r.a"),
step("lane-b", "$.r.b"),
step("lane-c", "$.r.c"),
],
};
let bp = bp(
flow,
vec![
agent("lane-a", None),
agent("lane-b", None),
agent("lane-c", None),
],
);
let (naming, warnings) = StepNaming::from_blueprint(&bp).expect("no collision");
assert!(
warnings.is_empty(),
"a shared nesting root is not an ambiguity: {warnings:?}"
);
assert_eq!(naming.resolve("r"), None);
for ref_ in ["lane-a", "lane-b", "lane-c"] {
let entry = naming
.entries()
.find(|e| e.canonical == ref_)
.expect("entry present");
assert_eq!(
entry.aliases,
BTreeSet::from([ref_.to_string()]),
"{ref_} must keep its own ref and drop the contested root"
);
assert_eq!(naming.resolve(ref_), Some(ref_));
}
}
#[test]
fn two_steps_writing_the_identical_root_still_warn() {
let flow = Node::Seq {
children: vec![step("first", "$.r"), step("second", "$.r")],
};
let bp = bp(flow, vec![agent("first", None), agent("second", None)]);
let (naming, warnings) = StepNaming::from_blueprint(&bp).expect("soft collision is Ok");
assert_eq!(warnings.len(), 1, "warnings: {warnings:?}");
assert_eq!(warnings[0].name, "r");
assert_eq!(naming.resolve("r"), Some("first"));
}
#[test]
fn strong_ref_claim_beats_a_nesting_sibling_without_warning() {
let flow = Node::Seq {
children: vec![step("r", "$.r_out"), step("lane", "$.r.x")],
};
let bp = bp(flow, vec![agent("r", None), agent("lane", None)]);
let (naming, warnings) = StepNaming::from_blueprint(&bp).expect("no collision");
assert!(warnings.is_empty(), "warnings: {warnings:?}");
assert_eq!(naming.resolve("r"), Some("r"));
let lane = naming
.entries()
.find(|e| e.canonical == "lane")
.expect("entry present");
assert_eq!(lane.aliases, BTreeSet::from(["lane".to_string()]));
}
#[test]
fn declared_name_against_a_nesting_sibling_compiles() {
let flow = Node::Seq {
children: vec![step("declarer", "$.declarer_out"), step("lane", "$.r.x")],
};
let bp = bp(
flow,
vec![agent("declarer", Some("r")), agent("lane", None)],
);
let (naming, warnings) =
StepNaming::from_blueprint(&bp).expect("a yielding weak claim must not hard-error");
assert!(warnings.is_empty(), "warnings: {warnings:?}");
assert_eq!(naming.resolve("r"), Some("r"));
assert_eq!(naming.canonical_of_producer("declarer"), Some("r"));
}
#[test]
fn uncontested_nesting_root_stays_an_alias() {
let flow = step("only", "$.r.a");
let bp = bp(flow, vec![agent("only", None)]);
let (naming, warnings) = StepNaming::from_blueprint(&bp).expect("no collision");
assert!(warnings.is_empty());
let entry = naming
.entries()
.find(|e| e.canonical == "only")
.expect("entry present");
assert_eq!(
entry.aliases,
BTreeSet::from(["only".to_string(), "r".to_string()])
);
assert_eq!(naming.resolve("r"), Some("only"));
}
#[test]
fn same_ref_dispatched_twice_unions_out_top_aliases_without_self_collision() {
let flow = Node::Seq {
children: vec![step("worker", "$.first"), step("worker", "$.second")],
};
let bp = bp(flow, vec![agent("worker", None)]);
let (naming, warnings) = StepNaming::from_blueprint(&bp).expect("no collision");
assert!(warnings.is_empty());
let entry = naming
.entries()
.find(|e| e.canonical == "worker")
.expect("entry present");
assert_eq!(
entry.aliases,
BTreeSet::from([
"worker".to_string(),
"first".to_string(),
"second".to_string()
])
);
}
}