Skip to main content

beam_core/
workflow_binding.rs

1use std::fs;
2use std::future::Future;
3use std::path::Path;
4use std::pin::Pin;
5
6use anyhow::{Context, Result};
7use serde_json::Value;
8
9use crate::{RunSnapshotDTO, WorkflowDefinition, WorkflowOutputRef};
10
11#[derive(Debug, Clone, PartialEq, Eq)]
12pub struct BindingError(pub String);
13
14impl std::fmt::Display for BindingError {
15    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
16        f.write_str(&self.0)
17    }
18}
19
20impl std::error::Error for BindingError {}
21
22#[derive(Debug, Clone, PartialEq)]
23pub struct BindingContext<'a> {
24    pub snapshot: &'a RunSnapshotDTO,
25    pub def: &'a WorkflowDefinition,
26    pub run_dir: &'a Path,
27    pub loop_context: Option<LoopContext<'a>>,
28}
29
30#[derive(Debug, Clone, Copy, PartialEq)]
31pub struct LoopContext<'a> {
32    pub loop_id: &'a str,
33    pub iteration: u64,
34}
35
36#[derive(Debug, Clone, PartialEq, Eq)]
37enum ParsedRef<'a> {
38    Output {
39        node_id: &'a str,
40        path: Vec<&'a str>,
41    },
42    Params {
43        path: Vec<&'a str>,
44    },
45    Previous {
46        node_id: &'a str,
47        path: Vec<&'a str>,
48    },
49}
50
51const REF_MARKER: &str = ".output.";
52const PREVIOUS_MARKER: &str = ".previous.";
53const PARAMS_PREFIX: &str = "params.";
54const FORBIDDEN_SEGMENTS: [&str; 3] = ["__proto__", "prototype", "constructor"];
55
56pub fn resolve_bindings<'a>(
57    value: &'a Value,
58    ctx: &'a BindingContext<'a>,
59) -> Pin<Box<dyn Future<Output = Result<Value>> + Send + 'a>> {
60    Box::pin(async move {
61        if let Some(ref_spec) = output_ref_spec(value) {
62            return resolve_output_ref(ref_spec, ctx).await;
63        }
64        if let Some(s) = value.as_str() {
65            return resolve_interpolated_string(s, ctx).await.map(Value::String);
66        }
67        if let Some(arr) = value.as_array() {
68            let mut out = Vec::with_capacity(arr.len());
69            for item in arr {
70                out.push(resolve_bindings(item, ctx).await?);
71            }
72            return Ok(Value::Array(out));
73        }
74        if let Some(obj) = value.as_object() {
75            let mut out = serde_json::Map::new();
76            for (k, v) in obj {
77                out.insert(k.clone(), resolve_bindings(v, ctx).await?);
78            }
79            return Ok(Value::Object(out));
80        }
81        Ok(value.clone())
82    })
83}
84
85pub fn resolve_bound_string<'a>(
86    value: &'a Value,
87    ctx: &'a BindingContext<'a>,
88) -> Pin<Box<dyn Future<Output = Result<String>> + Send + 'a>> {
89    Box::pin(async move {
90        let resolved = resolve_bindings(value, ctx).await?;
91        match resolved {
92            Value::String(s) => Ok(s),
93            other => anyhow::bail!(
94                "bound string field resolved to {}",
95                describe_json_kind(&other)
96            ),
97        }
98    })
99}
100
101fn output_ref_spec(value: &Value) -> Option<&str> {
102    let obj = value.as_object()?;
103    if obj.len() != 1 {
104        return None;
105    }
106    obj.get("$ref")?.as_str()
107}
108
109fn parse_ref<'a>(ref_spec: &'a str) -> Result<ParsedRef<'a>> {
110    if let Some(rest) = ref_spec.strip_prefix(PARAMS_PREFIX) {
111        return Ok(ParsedRef::Params {
112            path: parse_segments(rest, ref_spec)?,
113        });
114    }
115    if let Some(idx) = ref_spec.find(PREVIOUS_MARKER) {
116        let node_id = &ref_spec[..idx];
117        if node_id.is_empty() {
118            anyhow::bail!("$ref '{}' has empty nodeId before '.previous.'", ref_spec);
119        }
120        return Ok(ParsedRef::Previous {
121            node_id,
122            path: parse_segments(&ref_spec[idx + PREVIOUS_MARKER.len()..], ref_spec)?,
123        });
124    }
125    let Some(idx) = ref_spec.find(REF_MARKER) else {
126        anyhow::bail!(
127            "$ref '{}' missing '.output.' separator (expected '<nodeId>.output.<path>', '<nodeId>.previous.<path>', or 'params.<path>')",
128            ref_spec
129        );
130    };
131    let node_id = &ref_spec[..idx];
132    if node_id.is_empty() {
133        anyhow::bail!("$ref '{}' has empty nodeId before '.output.'", ref_spec);
134    }
135    Ok(ParsedRef::Output {
136        node_id,
137        path: parse_segments(&ref_spec[idx + REF_MARKER.len()..], ref_spec)?,
138    })
139}
140
141fn parse_segments<'a>(raw_path: &'a str, ref_spec: &str) -> Result<Vec<&'a str>> {
142    if raw_path.is_empty() {
143        anyhow::bail!("$ref '{}' has empty path", ref_spec);
144    }
145    let mut out = Vec::new();
146    for seg in raw_path.split('.') {
147        if seg.is_empty() {
148            anyhow::bail!("$ref '{}' has empty path segment", ref_spec);
149        }
150        if FORBIDDEN_SEGMENTS.contains(&seg) {
151            anyhow::bail!("$ref '{}' uses forbidden segment '{}'", ref_spec, seg);
152        }
153        if !seg
154            .chars()
155            .all(|ch| ch.is_ascii_alphanumeric() || matches!(ch, '_' | '-'))
156        {
157            anyhow::bail!(
158                "$ref '{}' has invalid segment '{}' (must match [A-Za-z0-9_-]+)",
159                ref_spec,
160                seg
161            );
162        }
163        out.push(seg);
164    }
165    Ok(out)
166}
167
168async fn resolve_output_ref(ref_spec: &str, ctx: &BindingContext<'_>) -> Result<Value> {
169    let parsed = parse_ref(ref_spec)?;
170    let (output_ref, path, node_id) = match parsed {
171        ParsedRef::Params { path } => {
172            let params = load_run_params(ref_spec, ctx).await?;
173            return walk_path(params, &path, ref_spec);
174        }
175        ParsedRef::Output { node_id, path } => (latest_output_ref(node_id, ctx)?, path, node_id),
176        ParsedRef::Previous { node_id, path } => {
177            let loop_ctx = ctx.loop_context.context(format!(
178                "$ref '{}' uses '.previous.' outside a loop iteration context",
179                ref_spec
180            ))?;
181            if loop_ctx.iteration <= 1 {
182                anyhow::bail!(
183                    "$ref '{}' has no previous iteration for iteration 1",
184                    ref_spec
185                );
186            }
187            (
188                previous_loop_output_ref(node_id, ctx, loop_ctx.loop_id, loop_ctx.iteration - 1)?
189                    .ok_or_else(|| anyhow::anyhow!(
190                        "$ref '{}' references node '{}' which has not produced a successful output yet",
191                        ref_spec, node_id
192                    ))?,
193                path,
194                node_id,
195            )
196        }
197    };
198
199    let blob = fs::read_to_string(&output_ref.output_path).with_context(|| {
200        format!(
201            "$ref '{}' failed to read output blob at {}",
202            ref_spec, output_ref.output_path
203        )
204    })?;
205    let raw: Value = serde_json::from_str(&blob)
206        .with_context(|| format!("$ref '{}' output blob is not valid JSON", ref_spec))?;
207    let logical_root = match ctx.def.nodes.get(node_id) {
208        Some(crate::WorkflowNode::HostExecutor(_)) => {
209            raw.get("output").cloned().unwrap_or(Value::Null)
210        }
211        _ => raw,
212    };
213    walk_path(logical_root, &path, ref_spec)
214}
215
216fn latest_output_ref<'a>(node_id: &'a str, ctx: &BindingContext<'a>) -> Result<WorkflowOutputRef> {
217    let run_id = &ctx.snapshot.run.run_id;
218    let plain = ctx
219        .snapshot
220        .outputs
221        .get(&format!("{}::work::{}", run_id, node_id));
222    if let Some(out) = plain {
223        return Ok(out.clone());
224    }
225    find_loop_output_ref(node_id, ctx, None, None)
226        .cloned()
227        .ok_or_else(|| {
228            anyhow::anyhow!(
229                "$ref targets node '{}' which has not produced a successful output yet",
230                node_id
231            )
232        })
233}
234
235fn previous_loop_output_ref<'a>(
236    node_id: &'a str,
237    ctx: &BindingContext<'a>,
238    loop_id: &'a str,
239    iteration: u64,
240) -> Result<Option<WorkflowOutputRef>> {
241    Ok(find_loop_output_ref(node_id, ctx, Some(loop_id), Some(iteration)).cloned())
242}
243
244fn find_loop_output_ref<'a>(
245    node_id: &'a str,
246    ctx: &'a BindingContext<'a>,
247    loop_id: Option<&'a str>,
248    iteration: Option<u64>,
249) -> Option<&'a WorkflowOutputRef> {
250    let node_def = ctx.def.nodes.get(node_id)?;
251    let expected_kind = match node_def {
252        crate::WorkflowNode::Decision(_) => "gate",
253        _ => "work",
254    };
255    let mut best: Option<(u64, &WorkflowOutputRef)> = None;
256    for (activity_id, output_ref) in &ctx.snapshot.outputs {
257        let Some(parsed) = parse_activity_id(activity_id) else {
258            continue;
259        };
260        if parsed.node_id != node_id || parsed.activity_kind != expected_kind {
261            continue;
262        }
263        if let Some(loop_id) = loop_id {
264            if parsed.loop_id.as_deref() != Some(loop_id) {
265                continue;
266            }
267        }
268        if let Some(iteration) = iteration {
269            if parsed.iteration != Some(iteration) {
270                continue;
271            }
272        }
273        let iter = parsed.iteration.unwrap_or(0);
274        if best.map(|(prev, _)| iter > prev).unwrap_or(true) {
275            best = Some((iter, output_ref));
276        }
277    }
278    best.map(|(_, output_ref)| output_ref)
279}
280
281#[derive(Debug, Clone, PartialEq, Eq)]
282struct ParsedActivityId<'a> {
283    node_id: &'a str,
284    activity_kind: &'a str,
285    loop_id: Option<&'a str>,
286    iteration: Option<u64>,
287}
288
289fn parse_activity_id<'a>(s: &'a str) -> Option<ParsedActivityId<'a>> {
290    if let Some(loop_idx) = s.find("::loop::") {
291        let run_id_end = loop_idx;
292        let after_loop = &s[loop_idx + "::loop::".len()..];
293        let iter_end = after_loop.find("::")?;
294        let loop_part = &after_loop[..iter_end];
295        let (loop_id, iter) = loop_part.rsplit_once('.')?;
296        let iteration = iter.parse().ok()?;
297        let after_iter = &after_loop[iter_end + 2..];
298        let mut segs = after_iter.splitn(2, "::");
299        let activity_kind = segs.next()?;
300        let node_id = segs.next()?;
301        let _ = run_id_end; // keep the parser symmetric with TS, run id is not needed here.
302        return Some(ParsedActivityId {
303            node_id,
304            activity_kind,
305            loop_id: Some(loop_id),
306            iteration: Some(iteration),
307        });
308    }
309    let mut parts = s.rsplitn(3, "::");
310    let node_id = parts.next()?;
311    let activity_kind = parts.next()?;
312    let _run_id = parts.next()?;
313    Some(ParsedActivityId {
314        node_id,
315        activity_kind,
316        loop_id: None,
317        iteration: None,
318    })
319}
320
321async fn load_run_params(ref_spec: &str, ctx: &BindingContext<'_>) -> Result<Value> {
322    let input_path = ctx
323        .snapshot
324        .run
325        .input
326        .as_ref()
327        .context(format!("$ref '{}' requires run input", ref_spec))?
328        .output_path
329        .clone();
330    let raw = fs::read_to_string(&input_path).with_context(|| {
331        format!(
332            "$ref '{}' failed to read run params at {}",
333            ref_spec, input_path
334        )
335    })?;
336    let parsed: Value = serde_json::from_str(&raw)
337        .with_context(|| format!("$ref '{}' run params blob is not valid JSON", ref_spec))?;
338    if !parsed.is_object() {
339        anyhow::bail!(
340            "$ref '{}' resolved run params to non-object input",
341            ref_spec
342        );
343    }
344    Ok(parsed)
345}
346
347fn walk_path(value: Value, segments: &[&str], ref_spec: &str) -> Result<Value> {
348    let mut cursor = value;
349    for seg in segments {
350        if cursor.is_null() {
351            anyhow::bail!("$ref '{}' hit null at '{}'", ref_spec, seg);
352        }
353        if let Some(arr) = cursor.as_array() {
354            let idx: usize = seg
355                .parse()
356                .with_context(|| format!("$ref '{}' array index '{}' invalid", ref_spec, seg))?;
357            cursor = arr.get(idx).cloned().with_context(|| {
358                format!("$ref '{}' array index '{}' out of bounds", ref_spec, seg)
359            })?;
360            continue;
361        }
362        let obj = cursor
363            .as_object()
364            .with_context(|| format!("$ref '{}' segment '{}' not found", ref_spec, seg))?;
365        cursor = obj
366            .get(*seg)
367            .cloned()
368            .with_context(|| format!("$ref '{}' segment '{}' not found", ref_spec, seg))?;
369    }
370    Ok(cursor)
371}
372
373async fn resolve_interpolated_string(value: &str, ctx: &BindingContext<'_>) -> Result<String> {
374    if !value.contains("${") {
375        return Ok(value.to_string());
376    }
377    let mut out = String::new();
378    let mut cursor = 0usize;
379    while let Some(start_rel) = value[cursor..].find("${") {
380        let start = cursor + start_rel;
381        out.push_str(&value[cursor..start]);
382        let end_rel = value[start + 2..].find('}').context(format!(
383            "unterminated string ref interpolation in '{}'",
384            value
385        ))?;
386        let end = start + 2 + end_rel;
387        let ref_spec = &value[start + 2..end];
388        if ref_spec.is_empty() {
389            anyhow::bail!("empty string ref interpolation in '{}'", value);
390        }
391        let resolved = resolve_output_ref(ref_spec, ctx).await?;
392        out.push_str(&stringify_interpolated_value(ref_spec, resolved)?);
393        cursor = end + 1;
394    }
395    out.push_str(&value[cursor..]);
396    Ok(out)
397}
398
399fn stringify_interpolated_value(ref_spec: &str, value: Value) -> Result<String> {
400    match value {
401        Value::Null => Ok("null".to_string()),
402        Value::String(s) => Ok(s),
403        Value::Number(n) => Ok(n.to_string()),
404        Value::Bool(b) => Ok(b.to_string()),
405        other => anyhow::bail!(
406            "string interpolation '${{{}}}' resolved to {}",
407            ref_spec,
408            describe_json_kind(&other)
409        ),
410    }
411}
412
413fn describe_json_kind(value: &Value) -> &'static str {
414    match value {
415        Value::Null => "null",
416        Value::Bool(_) => "bool",
417        Value::Number(_) => "number",
418        Value::String(_) => "string",
419        Value::Array(_) => "array",
420        Value::Object(_) => "object",
421    }
422}
423
424#[cfg(test)]
425mod tests {
426    use super::*;
427    use crate::workflow_definition::NodeBase;
428    use crate::workflow_snapshot::NodeStatus;
429    use crate::{RunState, RunStatus, WorkflowNode, WorkflowOutputRef};
430    use std::collections::BTreeMap;
431
432    fn ctx<'a>(
433        run_dir: &'a Path,
434        snapshot: &'a RunSnapshotDTO,
435        def: &'a WorkflowDefinition,
436    ) -> BindingContext<'a> {
437        BindingContext {
438            snapshot,
439            def,
440            run_dir,
441            loop_context: None,
442        }
443    }
444
445    #[tokio::test]
446    async fn resolve_bound_string_handles_params_and_output_refs() {
447        let temp = std::env::temp_dir().join(format!("beam-binding-{}", std::process::id()));
448        let _ = fs::remove_dir_all(&temp);
449        fs::create_dir_all(&temp).unwrap();
450        fs::write(temp.join("params.json"), r#"{"name":"beam"}"#).unwrap();
451        let snapshot = RunSnapshotDTO {
452            run_id: "run-1".to_string(),
453            run: RunState {
454                run_id: "run-1".to_string(),
455                status: RunStatus::Running,
456                workflow_id: Some("flow-a".to_string()),
457                revision_id: Some("rev-a".to_string()),
458                initiator: None,
459                input: Some(WorkflowOutputRef {
460                    output_hash: "sha256:params".to_string(),
461                    output_path: temp.join("params.json").display().to_string(),
462                    output_bytes: 17,
463                    output_schema_version: 1,
464                    content_type: Some("application/json".to_string()),
465                }),
466                output: None,
467                failed_node_id: None,
468                root_cause_event_id: None,
469                cancel_origin_event_id: None,
470                bot_snapshots: None,
471                cancelled_run_intent: None,
472                cancelled_node_intents: BTreeMap::new(),
473            },
474            last_seq: 1,
475            nodes: vec![crate::NodeState {
476                node_id: "a".to_string(),
477                status: NodeStatus::Succeeded,
478                activity_id: Some("run-1::work::a".to_string()),
479                retry_count: 0,
480                next_attempt_at: None,
481                error_class: None,
482                condition_event_id: None,
483                cancel_origin_event_id: None,
484            }],
485            activities: Vec::new(),
486            loops: None,
487            dangling: crate::DanglingSnapshot {
488                activities: Vec::new(),
489                effect_attempted: Vec::new(),
490                waits: Vec::new(),
491                cancels: Vec::new(),
492            },
493            outputs: BTreeMap::from([(
494                "run-1::work::a".to_string(),
495                WorkflowOutputRef {
496                    output_hash: "sha256:out".to_string(),
497                    output_path: temp.join("out.json").display().to_string(),
498                    output_bytes: 15,
499                    output_schema_version: 1,
500                    content_type: Some("application/json".to_string()),
501                },
502            )]),
503            attempt_io: BTreeMap::new(),
504            chat_binding: None,
505            updated_at: 1,
506        };
507        fs::write(temp.join("out.json"), r#"{"message":"ok"}"#).unwrap();
508        let def = WorkflowDefinition {
509            workflow_id: "flow-a".to_string(),
510            version: 1,
511            params: None,
512            defaults: None,
513            nodes: BTreeMap::from([(
514                "a".to_string(),
515                WorkflowNode::Subagent(crate::SubagentNode {
516                    base: NodeBase {
517                        description: None,
518                        depends: None,
519                        human_gate: None,
520                        retry_policy: None,
521                        timeout_ms: None,
522                        max_output_bytes: None,
523                        output_schema: None,
524                        unsafe_allow_ungated: None,
525                    },
526                    bot: "bot-a".to_string(),
527                    prompt: Value::String("hello ${params.name} ${a.output.message}".to_string()),
528                    working_dir: None,
529                    model_overrides: None,
530                    tool_policy: None,
531                }),
532            )]),
533        };
534        let resolved = resolve_bindings(
535            &Value::String("hello ${params.name} ${a.output.message}".to_string()),
536            &ctx(&temp, &snapshot, &def),
537        )
538        .await
539        .unwrap();
540        assert_eq!(resolved.as_str(), Some("hello beam ok"));
541        let _ = fs::remove_dir_all(&temp);
542    }
543}