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                // First iteration has no previous.  For Decision nodes
183                // produce an empty synthetic JSON so that string
184                // interpolation (e.g. ${reviewDecision.previous.comment})
185                // yields an empty string without changing the global
186                // null → "null" mapping.
187                if let Some(crate::WorkflowNode::Decision(_)) = ctx.def.nodes.get(node_id) {
188                    let empty = serde_json::json!({"by": null, "comment": ""});
189                    return walk_path(empty, &path, ref_spec);
190                }
191                return Ok(Value::Null);
192            }
193            let prev_iteration = loop_ctx.iteration - 1;
194
195            // For Decision nodes, read decision metadata from the loop
196            // iteration state rather than from activity outputs.  This
197            // supports `${reviewDecision.previous.comment}` even when
198            // the previous iteration was rejected (activityFailed, no
199            // output blob).
200            if let Some(crate::WorkflowNode::Decision(_)) = ctx.def.nodes.get(node_id) {
201                if let Some(iter_state) = ctx
202                    .snapshot
203                    .loops
204                    .as_ref()
205                    .and_then(|loops| loops.get(loop_ctx.loop_id))
206                    .and_then(|ls| {
207                        ls.iterations
208                            .iter()
209                            .find(|it| it.iteration == prev_iteration)
210                    })
211                {
212                    let decision_data = serde_json::json!({
213                        "by": iter_state.decision_by,
214                        "comment": iter_state.decision_comment,
215                    });
216                    return walk_path(decision_data, &path, ref_spec);
217                }
218                // No previous iteration data yet — for Decision
219                // nodes return empty synthetic JSON.
220                let empty = serde_json::json!({"by": null, "comment": ""});
221                return walk_path(empty, &path, ref_spec);
222            }
223
224            (
225                previous_loop_output_ref(node_id, ctx, loop_ctx.loop_id, prev_iteration)?
226                    .ok_or_else(|| anyhow::anyhow!(
227                        "$ref '{}' references node '{}' which has not produced a successful output yet",
228                        ref_spec, node_id
229                    ))?,
230                path,
231                node_id,
232            )
233        }
234    };
235
236    let blob = fs::read_to_string(&output_ref.output_path).with_context(|| {
237        format!(
238            "$ref '{}' failed to read output blob at {}",
239            ref_spec, output_ref.output_path
240        )
241    })?;
242    let raw: Value = serde_json::from_str(&blob)
243        .with_context(|| format!("$ref '{}' output blob is not valid JSON", ref_spec))?;
244    let logical_root = match ctx.def.nodes.get(node_id) {
245        Some(crate::WorkflowNode::HostExecutor(_)) => {
246            raw.get("output").cloned().unwrap_or(Value::Null)
247        }
248        _ => raw,
249    };
250    walk_path(logical_root, &path, ref_spec)
251}
252
253fn latest_output_ref<'a>(node_id: &'a str, ctx: &BindingContext<'a>) -> Result<WorkflowOutputRef> {
254    let run_id = &ctx.snapshot.run.run_id;
255    let plain = ctx
256        .snapshot
257        .outputs
258        .get(&format!("{}::work::{}", run_id, node_id));
259    if let Some(out) = plain {
260        return Ok(out.clone());
261    }
262    find_loop_output_ref(node_id, ctx, None, None)
263        .cloned()
264        .ok_or_else(|| {
265            anyhow::anyhow!(
266                "$ref targets node '{}' which has not produced a successful output yet",
267                node_id
268            )
269        })
270}
271
272fn previous_loop_output_ref<'a>(
273    node_id: &'a str,
274    ctx: &BindingContext<'a>,
275    loop_id: &'a str,
276    iteration: u64,
277) -> Result<Option<WorkflowOutputRef>> {
278    Ok(find_loop_output_ref(node_id, ctx, Some(loop_id), Some(iteration)).cloned())
279}
280
281fn find_loop_output_ref<'a>(
282    node_id: &'a str,
283    ctx: &'a BindingContext<'a>,
284    loop_id: Option<&'a str>,
285    iteration: Option<u64>,
286) -> Option<&'a WorkflowOutputRef> {
287    let node_def = ctx.def.nodes.get(node_id)?;
288    let expected_kind = match node_def {
289        crate::WorkflowNode::Decision(_) => "gate",
290        _ => "work",
291    };
292    let mut best: Option<(u64, &WorkflowOutputRef)> = None;
293    for (activity_id, output_ref) in &ctx.snapshot.outputs {
294        let Some(parsed) = parse_activity_id(activity_id) else {
295            continue;
296        };
297        if parsed.node_id != node_id || parsed.activity_kind != expected_kind {
298            continue;
299        }
300        if let Some(loop_id) = loop_id {
301            if parsed.loop_id.as_deref() != Some(loop_id) {
302                continue;
303            }
304        }
305        if let Some(iteration) = iteration {
306            if parsed.iteration != Some(iteration) {
307                continue;
308            }
309        }
310        let iter = parsed.iteration.unwrap_or(0);
311        if best.map(|(prev, _)| iter > prev).unwrap_or(true) {
312            best = Some((iter, output_ref));
313        }
314    }
315    best.map(|(_, output_ref)| output_ref)
316}
317
318#[derive(Debug, Clone, PartialEq, Eq)]
319struct ParsedActivityId<'a> {
320    node_id: &'a str,
321    activity_kind: &'a str,
322    loop_id: Option<&'a str>,
323    iteration: Option<u64>,
324}
325
326fn parse_activity_id<'a>(s: &'a str) -> Option<ParsedActivityId<'a>> {
327    if let Some(loop_idx) = s.find("::loop::") {
328        let run_id_end = loop_idx;
329        let after_loop = &s[loop_idx + "::loop::".len()..];
330        let iter_end = after_loop.find("::")?;
331        let loop_part = &after_loop[..iter_end];
332        let (loop_id, iter) = loop_part.rsplit_once('.')?;
333        let iteration = iter.parse().ok()?;
334        let after_iter = &after_loop[iter_end + 2..];
335        let mut segs = after_iter.splitn(2, "::");
336        let activity_kind = segs.next()?;
337        let node_id = segs.next()?;
338        let _ = run_id_end; // keep the parser symmetric with TS, run id is not needed here.
339        return Some(ParsedActivityId {
340            node_id,
341            activity_kind,
342            loop_id: Some(loop_id),
343            iteration: Some(iteration),
344        });
345    }
346    let mut parts = s.rsplitn(3, "::");
347    let node_id = parts.next()?;
348    let activity_kind = parts.next()?;
349    let _run_id = parts.next()?;
350    Some(ParsedActivityId {
351        node_id,
352        activity_kind,
353        loop_id: None,
354        iteration: None,
355    })
356}
357
358async fn load_run_params(ref_spec: &str, ctx: &BindingContext<'_>) -> Result<Value> {
359    let input_path = ctx
360        .snapshot
361        .run
362        .input
363        .as_ref()
364        .context(format!("$ref '{}' requires run input", ref_spec))?
365        .output_path
366        .clone();
367    let raw = fs::read_to_string(&input_path).with_context(|| {
368        format!(
369            "$ref '{}' failed to read run params at {}",
370            ref_spec, input_path
371        )
372    })?;
373    let parsed: Value = serde_json::from_str(&raw)
374        .with_context(|| format!("$ref '{}' run params blob is not valid JSON", ref_spec))?;
375    if !parsed.is_object() {
376        anyhow::bail!(
377            "$ref '{}' resolved run params to non-object input",
378            ref_spec
379        );
380    }
381    Ok(parsed)
382}
383
384fn walk_path(value: Value, segments: &[&str], ref_spec: &str) -> Result<Value> {
385    let mut cursor = value;
386    for seg in segments {
387        if cursor.is_null() {
388            anyhow::bail!("$ref '{}' hit null at '{}'", ref_spec, seg);
389        }
390        if let Some(arr) = cursor.as_array() {
391            let idx: usize = seg
392                .parse()
393                .with_context(|| format!("$ref '{}' array index '{}' invalid", ref_spec, seg))?;
394            cursor = arr.get(idx).cloned().with_context(|| {
395                format!("$ref '{}' array index '{}' out of bounds", ref_spec, seg)
396            })?;
397            continue;
398        }
399        let obj = cursor
400            .as_object()
401            .with_context(|| format!("$ref '{}' segment '{}' not found", ref_spec, seg))?;
402        cursor = obj
403            .get(*seg)
404            .cloned()
405            .with_context(|| format!("$ref '{}' segment '{}' not found", ref_spec, seg))?;
406    }
407    Ok(cursor)
408}
409
410async fn resolve_interpolated_string(value: &str, ctx: &BindingContext<'_>) -> Result<String> {
411    if !value.contains("${") {
412        return Ok(value.to_string());
413    }
414    let mut out = String::new();
415    let mut cursor = 0usize;
416    while let Some(start_rel) = value[cursor..].find("${") {
417        let start = cursor + start_rel;
418        out.push_str(&value[cursor..start]);
419        let end_rel = value[start + 2..].find('}').context(format!(
420            "unterminated string ref interpolation in '{}'",
421            value
422        ))?;
423        let end = start + 2 + end_rel;
424        let ref_spec = &value[start + 2..end];
425        if ref_spec.is_empty() {
426            anyhow::bail!("empty string ref interpolation in '{}'", value);
427        }
428        let resolved = resolve_output_ref(ref_spec, ctx).await?;
429        out.push_str(&stringify_interpolated_value(ref_spec, resolved)?);
430        cursor = end + 1;
431    }
432    out.push_str(&value[cursor..]);
433    Ok(out)
434}
435
436fn stringify_interpolated_value(ref_spec: &str, value: Value) -> Result<String> {
437    match value {
438        Value::Null => Ok("null".to_string()),
439        Value::String(s) => Ok(s),
440        Value::Number(n) => Ok(n.to_string()),
441        Value::Bool(b) => Ok(b.to_string()),
442        other => anyhow::bail!(
443            "string interpolation '${{{}}}' resolved to {}",
444            ref_spec,
445            describe_json_kind(&other)
446        ),
447    }
448}
449
450fn describe_json_kind(value: &Value) -> &'static str {
451    match value {
452        Value::Null => "null",
453        Value::Bool(_) => "bool",
454        Value::Number(_) => "number",
455        Value::String(_) => "string",
456        Value::Array(_) => "array",
457        Value::Object(_) => "object",
458    }
459}
460
461#[cfg(test)]
462mod tests {
463    use super::*;
464    use crate::workflow_definition::NodeBase;
465    use crate::workflow_snapshot::NodeStatus;
466    use crate::{RunState, RunStatus, WorkflowNode, WorkflowOutputRef};
467    use std::collections::BTreeMap;
468
469    fn ctx<'a>(
470        run_dir: &'a Path,
471        snapshot: &'a RunSnapshotDTO,
472        def: &'a WorkflowDefinition,
473    ) -> BindingContext<'a> {
474        BindingContext {
475            snapshot,
476            def,
477            run_dir,
478            loop_context: None,
479        }
480    }
481
482    #[tokio::test]
483    async fn resolve_bound_string_handles_params_and_output_refs() {
484        let temp = std::env::temp_dir().join(format!("beam-binding-{}", std::process::id()));
485        let _ = fs::remove_dir_all(&temp);
486        fs::create_dir_all(&temp).unwrap();
487        fs::write(temp.join("params.json"), r#"{"name":"beam"}"#).unwrap();
488        let snapshot = RunSnapshotDTO {
489            run_id: "run-1".to_string(),
490            run: RunState {
491                run_id: "run-1".to_string(),
492                status: RunStatus::Running,
493                workflow_id: Some("flow-a".to_string()),
494                revision_id: Some("rev-a".to_string()),
495                initiator: None,
496                input: Some(WorkflowOutputRef {
497                    output_hash: "sha256:params".to_string(),
498                    output_path: temp.join("params.json").display().to_string(),
499                    output_bytes: 17,
500                    output_schema_version: 1,
501                    content_type: Some("application/json".to_string()),
502                }),
503                output: None,
504                failed_node_id: None,
505                root_cause_event_id: None,
506                cancel_origin_event_id: None,
507                bot_snapshots: None,
508                cancelled_run_intent: None,
509                cancelled_node_intents: BTreeMap::new(),
510            },
511            last_seq: 1,
512            nodes: vec![crate::NodeState {
513                node_id: "a".to_string(),
514                status: NodeStatus::Succeeded,
515                activity_id: Some("run-1::work::a".to_string()),
516                retry_count: 0,
517                next_attempt_at: None,
518                error_class: None,
519                condition_event_id: None,
520                cancel_origin_event_id: None,
521            }],
522            activities: Vec::new(),
523            loops: None,
524            dangling: crate::DanglingSnapshot {
525                activities: Vec::new(),
526                effect_attempted: Vec::new(),
527                waits: Vec::new(),
528                wait_resolutions: Vec::new(),
529                cancels: Vec::new(),
530            },
531            outputs: BTreeMap::from([(
532                "run-1::work::a".to_string(),
533                WorkflowOutputRef {
534                    output_hash: "sha256:out".to_string(),
535                    output_path: temp.join("out.json").display().to_string(),
536                    output_bytes: 15,
537                    output_schema_version: 1,
538                    content_type: Some("application/json".to_string()),
539                },
540            )]),
541            attempt_io: BTreeMap::new(),
542            chat_binding: None,
543            updated_at: 1,
544        };
545        fs::write(temp.join("out.json"), r#"{"message":"ok"}"#).unwrap();
546        let def = WorkflowDefinition {
547            workflow_id: "flow-a".to_string(),
548            version: 1,
549            params: None,
550            defaults: None,
551            nodes: BTreeMap::from([(
552                "a".to_string(),
553                WorkflowNode::Subagent(crate::SubagentNode {
554                    base: NodeBase {
555                        description: None,
556                        depends: None,
557                        human_gate: None,
558                        retry_policy: None,
559                        timeout_ms: None,
560                        max_output_bytes: None,
561                        output_schema: None,
562                        unsafe_allow_ungated: None,
563                    },
564                    bot: "bot-a".to_string(),
565                    prompt: Value::String("hello ${params.name} ${a.output.message}".to_string()),
566                    working_dir: None,
567                    model_overrides: None,
568                    tool_policy: None,
569                }),
570            )]),
571        };
572        let resolved = resolve_bindings(
573            &Value::String("hello ${params.name} ${a.output.message}".to_string()),
574            &ctx(&temp, &snapshot, &def),
575        )
576        .await
577        .unwrap();
578        assert_eq!(resolved.as_str(), Some("hello beam ok"));
579        let _ = fs::remove_dir_all(&temp);
580    }
581
582    /// Null in string interpolation remains the literal "null".
583    #[tokio::test]
584    async fn string_interpolation_null_produces_literal_null() {
585        let temp = std::env::temp_dir().join(format!("beam-binding-null-{}", std::process::id()));
586        let _ = fs::remove_dir_all(&temp);
587        fs::create_dir_all(&temp).unwrap();
588
589        let snapshot = RunSnapshotDTO {
590            run_id: "run-null".to_string(),
591            run: RunState {
592                run_id: "run-null".to_string(),
593                status: RunStatus::Running,
594                workflow_id: None,
595                revision_id: None,
596                initiator: None,
597                input: None,
598                output: None,
599                failed_node_id: None,
600                root_cause_event_id: None,
601                cancel_origin_event_id: None,
602                bot_snapshots: None,
603                cancelled_run_intent: None,
604                cancelled_node_intents: BTreeMap::new(),
605            },
606            last_seq: 0,
607            nodes: vec![crate::NodeState {
608                node_id: "a".to_string(),
609                status: NodeStatus::Succeeded,
610                activity_id: Some("run-null::work::a".to_string()),
611                retry_count: 0,
612                next_attempt_at: None,
613                error_class: None,
614                condition_event_id: None,
615                cancel_origin_event_id: None,
616            }],
617            activities: Vec::new(),
618            loops: None,
619            dangling: crate::DanglingSnapshot {
620                activities: Vec::new(),
621                effect_attempted: Vec::new(),
622                waits: Vec::new(),
623                wait_resolutions: Vec::new(),
624                cancels: Vec::new(),
625            },
626            outputs: BTreeMap::from([(
627                "run-null::work::a".to_string(),
628                WorkflowOutputRef {
629                    output_hash: "sha256:null-out".to_string(),
630                    output_path: temp.join("null-out.json").display().to_string(),
631                    output_bytes: 16,
632                    output_schema_version: 1,
633                    content_type: Some("application/json".to_string()),
634                },
635            )]),
636            attempt_io: BTreeMap::new(),
637            chat_binding: None,
638            updated_at: 1,
639        };
640        // Output blob: { "val": null }
641        fs::write(temp.join("null-out.json"), r#"{"val":null}"#).unwrap();
642
643        let def = WorkflowDefinition {
644            workflow_id: "flow-null".to_string(),
645            version: 1,
646            params: None,
647            defaults: None,
648            nodes: BTreeMap::from([(
649                "a".to_string(),
650                WorkflowNode::Subagent(crate::SubagentNode {
651                    base: NodeBase {
652                        description: None,
653                        depends: None,
654                        human_gate: None,
655                        retry_policy: None,
656                        timeout_ms: None,
657                        max_output_bytes: None,
658                        output_schema: None,
659                        unsafe_allow_ungated: None,
660                    },
661                    bot: "bot-a".to_string(),
662                    prompt: Value::String("x".to_string()),
663                    working_dir: None,
664                    model_overrides: None,
665                    tool_policy: None,
666                }),
667            )]),
668        };
669
670        // ${a.output.val} — the val field is null.
671        let resolved = resolve_bindings(
672            &Value::String("before ${a.output.val} after".to_string()),
673            &ctx(&temp, &snapshot, &def),
674        )
675        .await
676        .unwrap();
677        assert_eq!(
678            resolved.as_str(),
679            Some("before null after"),
680            "null interpolation should produce literal 'null'"
681        );
682        let _ = fs::remove_dir_all(&temp);
683    }
684}