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            && parsed.loop_id != Some(loop_id)
302        {
303            continue;
304        }
305        if let Some(iteration) = iteration
306            && parsed.iteration != Some(iteration)
307        {
308            continue;
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 (activity_kind, node_id) = after_iter.split_once("::")?;
336
337        let _ = run_id_end; // keep the parser symmetric with TS, run id is not needed here.
338        return Some(ParsedActivityId {
339            node_id,
340            activity_kind,
341            loop_id: Some(loop_id),
342            iteration: Some(iteration),
343        });
344    }
345    let mut parts = s.rsplitn(3, "::");
346    let node_id = parts.next()?;
347    let activity_kind = parts.next()?;
348    let _run_id = parts.next()?;
349    Some(ParsedActivityId {
350        node_id,
351        activity_kind,
352        loop_id: None,
353        iteration: None,
354    })
355}
356
357async fn load_run_params(ref_spec: &str, ctx: &BindingContext<'_>) -> Result<Value> {
358    let input_path = ctx
359        .snapshot
360        .run
361        .input
362        .as_ref()
363        .context(format!("$ref '{}' requires run input", ref_spec))?
364        .output_path
365        .clone();
366    let raw = fs::read_to_string(&input_path).with_context(|| {
367        format!(
368            "$ref '{}' failed to read run params at {}",
369            ref_spec, input_path
370        )
371    })?;
372    let parsed: Value = serde_json::from_str(&raw)
373        .with_context(|| format!("$ref '{}' run params blob is not valid JSON", ref_spec))?;
374    if !parsed.is_object() {
375        anyhow::bail!(
376            "$ref '{}' resolved run params to non-object input",
377            ref_spec
378        );
379    }
380    Ok(parsed)
381}
382
383fn walk_path(value: Value, segments: &[&str], ref_spec: &str) -> Result<Value> {
384    let mut cursor = value;
385    for seg in segments {
386        if cursor.is_null() {
387            anyhow::bail!("$ref '{}' hit null at '{}'", ref_spec, seg);
388        }
389        if let Some(arr) = cursor.as_array() {
390            let idx: usize = seg
391                .parse()
392                .with_context(|| format!("$ref '{}' array index '{}' invalid", ref_spec, seg))?;
393            cursor = arr.get(idx).cloned().with_context(|| {
394                format!("$ref '{}' array index '{}' out of bounds", ref_spec, seg)
395            })?;
396            continue;
397        }
398        let obj = cursor
399            .as_object()
400            .with_context(|| format!("$ref '{}' segment '{}' not found", ref_spec, seg))?;
401        cursor = obj
402            .get(*seg)
403            .cloned()
404            .with_context(|| format!("$ref '{}' segment '{}' not found", ref_spec, seg))?;
405    }
406    Ok(cursor)
407}
408
409async fn resolve_interpolated_string(value: &str, ctx: &BindingContext<'_>) -> Result<String> {
410    if !value.contains("${") {
411        return Ok(value.to_string());
412    }
413    let mut out = String::new();
414    let mut cursor = 0usize;
415    while let Some(start_rel) = value[cursor..].find("${") {
416        let start = cursor + start_rel;
417        out.push_str(&value[cursor..start]);
418        let end_rel = value[start + 2..].find('}').context(format!(
419            "unterminated string ref interpolation in '{}'",
420            value
421        ))?;
422        let end = start + 2 + end_rel;
423        let ref_spec = &value[start + 2..end];
424        if ref_spec.is_empty() {
425            anyhow::bail!("empty string ref interpolation in '{}'", value);
426        }
427        let resolved = resolve_output_ref(ref_spec, ctx).await?;
428        out.push_str(&stringify_interpolated_value(ref_spec, resolved)?);
429        cursor = end + 1;
430    }
431    out.push_str(&value[cursor..]);
432    Ok(out)
433}
434
435fn stringify_interpolated_value(ref_spec: &str, value: Value) -> Result<String> {
436    match value {
437        Value::Null => Ok("null".to_string()),
438        Value::String(s) => Ok(s),
439        Value::Number(n) => Ok(n.to_string()),
440        Value::Bool(b) => Ok(b.to_string()),
441        other => anyhow::bail!(
442            "string interpolation '${{{}}}' resolved to {}",
443            ref_spec,
444            describe_json_kind(&other)
445        ),
446    }
447}
448
449fn describe_json_kind(value: &Value) -> &'static str {
450    match value {
451        Value::Null => "null",
452        Value::Bool(_) => "bool",
453        Value::Number(_) => "number",
454        Value::String(_) => "string",
455        Value::Array(_) => "array",
456        Value::Object(_) => "object",
457    }
458}
459
460#[cfg(test)]
461mod tests {
462    use super::*;
463    use crate::workflow_definition::NodeBase;
464    use crate::workflow_snapshot::NodeStatus;
465    use crate::{RunState, RunStatus, WorkflowNode, WorkflowOutputRef};
466    use std::collections::BTreeMap;
467
468    fn ctx<'a>(
469        run_dir: &'a Path,
470        snapshot: &'a RunSnapshotDTO,
471        def: &'a WorkflowDefinition,
472    ) -> BindingContext<'a> {
473        BindingContext {
474            snapshot,
475            def,
476            run_dir,
477            loop_context: None,
478        }
479    }
480
481    #[tokio::test]
482    async fn resolve_bound_string_handles_params_and_output_refs() {
483        let temp = std::env::temp_dir().join(format!("beam-binding-{}", std::process::id()));
484        let _ = fs::remove_dir_all(&temp);
485        fs::create_dir_all(&temp).unwrap();
486        fs::write(temp.join("params.json"), r#"{"name":"beam"}"#).unwrap();
487        let snapshot = RunSnapshotDTO {
488            run_id: "run-1".to_string(),
489            run: RunState {
490                run_id: "run-1".to_string(),
491                status: RunStatus::Running,
492                workflow_id: Some("flow-a".to_string()),
493                revision_id: Some("rev-a".to_string()),
494                initiator: None,
495                input: Some(WorkflowOutputRef {
496                    output_hash: "sha256:params".to_string(),
497                    output_path: temp.join("params.json").display().to_string(),
498                    output_bytes: 17,
499                    output_schema_version: 1,
500                    content_type: Some("application/json".to_string()),
501                }),
502                output: None,
503                failed_node_id: None,
504                root_cause_event_id: None,
505                cancel_origin_event_id: None,
506                bot_snapshots: None,
507                cancelled_run_intent: None,
508                cancelled_node_intents: BTreeMap::new(),
509            },
510            last_seq: 1,
511            nodes: vec![crate::NodeState {
512                node_id: "a".to_string(),
513                status: NodeStatus::Succeeded,
514                activity_id: Some("run-1::work::a".to_string()),
515                retry_count: 0,
516                next_attempt_at: None,
517                error_class: None,
518                condition_event_id: None,
519                cancel_origin_event_id: None,
520            }],
521            activities: Vec::new(),
522            loops: None,
523            dangling: crate::DanglingSnapshot {
524                activities: Vec::new(),
525                effect_attempted: Vec::new(),
526                waits: Vec::new(),
527                wait_resolutions: Vec::new(),
528                cancels: Vec::new(),
529            },
530            outputs: BTreeMap::from([(
531                "run-1::work::a".to_string(),
532                WorkflowOutputRef {
533                    output_hash: "sha256:out".to_string(),
534                    output_path: temp.join("out.json").display().to_string(),
535                    output_bytes: 15,
536                    output_schema_version: 1,
537                    content_type: Some("application/json".to_string()),
538                },
539            )]),
540            attempt_io: BTreeMap::new(),
541            chat_binding: None,
542            updated_at: 1,
543        };
544        fs::write(temp.join("out.json"), r#"{"message":"ok"}"#).unwrap();
545        let def = WorkflowDefinition {
546            workflow_id: "flow-a".to_string(),
547            version: 1,
548            params: None,
549            defaults: None,
550            nodes: BTreeMap::from([(
551                "a".to_string(),
552                WorkflowNode::Subagent(crate::SubagentNode {
553                    base: NodeBase {
554                        description: None,
555                        depends: None,
556                        human_gate: None,
557                        retry_policy: None,
558                        timeout_ms: None,
559                        max_output_bytes: None,
560                        output_schema: None,
561                        unsafe_allow_ungated: None,
562                    },
563                    bot: "bot-a".to_string(),
564                    prompt: Value::String("hello ${params.name} ${a.output.message}".to_string()),
565                    working_dir: None,
566                    model_overrides: None,
567                    tool_policy: None,
568                }),
569            )]),
570        };
571        let resolved = resolve_bindings(
572            &Value::String("hello ${params.name} ${a.output.message}".to_string()),
573            &ctx(&temp, &snapshot, &def),
574        )
575        .await
576        .unwrap();
577        assert_eq!(resolved.as_str(), Some("hello beam ok"));
578        let _ = fs::remove_dir_all(&temp);
579    }
580
581    /// Null in string interpolation remains the literal "null".
582    #[tokio::test]
583    async fn string_interpolation_null_produces_literal_null() {
584        let temp = std::env::temp_dir().join(format!("beam-binding-null-{}", std::process::id()));
585        let _ = fs::remove_dir_all(&temp);
586        fs::create_dir_all(&temp).unwrap();
587
588        let snapshot = RunSnapshotDTO {
589            run_id: "run-null".to_string(),
590            run: RunState {
591                run_id: "run-null".to_string(),
592                status: RunStatus::Running,
593                workflow_id: None,
594                revision_id: None,
595                initiator: None,
596                input: None,
597                output: None,
598                failed_node_id: None,
599                root_cause_event_id: None,
600                cancel_origin_event_id: None,
601                bot_snapshots: None,
602                cancelled_run_intent: None,
603                cancelled_node_intents: BTreeMap::new(),
604            },
605            last_seq: 0,
606            nodes: vec![crate::NodeState {
607                node_id: "a".to_string(),
608                status: NodeStatus::Succeeded,
609                activity_id: Some("run-null::work::a".to_string()),
610                retry_count: 0,
611                next_attempt_at: None,
612                error_class: None,
613                condition_event_id: None,
614                cancel_origin_event_id: None,
615            }],
616            activities: Vec::new(),
617            loops: None,
618            dangling: crate::DanglingSnapshot {
619                activities: Vec::new(),
620                effect_attempted: Vec::new(),
621                waits: Vec::new(),
622                wait_resolutions: Vec::new(),
623                cancels: Vec::new(),
624            },
625            outputs: BTreeMap::from([(
626                "run-null::work::a".to_string(),
627                WorkflowOutputRef {
628                    output_hash: "sha256:null-out".to_string(),
629                    output_path: temp.join("null-out.json").display().to_string(),
630                    output_bytes: 16,
631                    output_schema_version: 1,
632                    content_type: Some("application/json".to_string()),
633                },
634            )]),
635            attempt_io: BTreeMap::new(),
636            chat_binding: None,
637            updated_at: 1,
638        };
639        // Output blob: { "val": null }
640        fs::write(temp.join("null-out.json"), r#"{"val":null}"#).unwrap();
641
642        let def = WorkflowDefinition {
643            workflow_id: "flow-null".to_string(),
644            version: 1,
645            params: None,
646            defaults: None,
647            nodes: BTreeMap::from([(
648                "a".to_string(),
649                WorkflowNode::Subagent(crate::SubagentNode {
650                    base: NodeBase {
651                        description: None,
652                        depends: None,
653                        human_gate: None,
654                        retry_policy: None,
655                        timeout_ms: None,
656                        max_output_bytes: None,
657                        output_schema: None,
658                        unsafe_allow_ungated: None,
659                    },
660                    bot: "bot-a".to_string(),
661                    prompt: Value::String("x".to_string()),
662                    working_dir: None,
663                    model_overrides: None,
664                    tool_policy: None,
665                }),
666            )]),
667        };
668
669        // ${a.output.val} — the val field is null.
670        let resolved = resolve_bindings(
671            &Value::String("before ${a.output.val} after".to_string()),
672            &ctx(&temp, &snapshot, &def),
673        )
674        .await
675        .unwrap();
676        assert_eq!(
677            resolved.as_str(),
678            Some("before null after"),
679            "null interpolation should produce literal 'null'"
680        );
681        let _ = fs::remove_dir_all(&temp);
682    }
683}