From 1a5824ce4f6973162d7c22a44308c6e4a0e41601 Mon Sep 17 00:00:00 2001 From: Alex Clarke Date: Thu, 10 Sep 2026 18:44:33 -0600 Subject: [PATCH] fix(graph): leave output absent after state_updates when it was absent before --- src/graph/agent.rs | 2 +- src/graph/executor.rs | 57 ++++++++++++++++++++++++++++++++++++++ src/graph/llm.rs | 2 +- src/graph/map.rs | 49 ++++++++++++++++++++++++++++++++ src/graph/state_updates.rs | 27 ++++++++++++++---- 5 files changed, 129 insertions(+), 8 deletions(-) diff --git a/src/graph/agent.rs b/src/graph/agent.rs index c6560dc..03b9830 100644 --- a/src/graph/agent.rs +++ b/src/graph/agent.rs @@ -291,7 +291,7 @@ mod tests { apply_state_updates(&node, &mut state, &json!("anything")); - assert_eq!(state.state().get("output"), Some(&Value::Null)); + assert!(state.state().get("output").is_none()); } #[test] diff --git a/src/graph/executor.rs b/src/graph/executor.rs index 2ddc2a3..31b6a67 100644 --- a/src/graph/executor.rs +++ b/src/graph/executor.rs @@ -1002,6 +1002,63 @@ nodes: assert!(strs.contains(&"beta"), "missing 'beta' in {strs:?}"); } + /// Two parallel agent branches whose `state_updates` only read `{{output}}` + /// leave the key as they found it (absent), so the join sees no write to + /// `output` from either side and needs no reducer for it. + #[tokio::test] + #[serial] + async fn parallel_state_updates_branches_do_not_contend_on_output() { + if !cmd_available("bash") || !cmd_available("python3") { + eprintln!("skipping: bash or python3 not available"); + return; + } + let _guard = TestConfigDirGuard::new(); + materialize_probe_agent(0.0); + let ws = TestWorkspace::new(); + ws.write_script("dispatcher.sh", "#!/bin/bash\necho '{}'\n"); + + let yaml = format!( + r#" +name: t +settings: + validate_before_run: false +start: dispatcher +nodes: + dispatcher: + type: script + script: dispatcher.sh + state_updates: {{}} + next: [a, b] + a: + type: agent + agent: {PROBE_AGENT} + prompt: "p" + state_updates: + note_a: "{{{{output}}}}" + next: join + b: + type: agent + agent: {PROBE_AGENT} + prompt: "p" + state_updates: + note_b: "{{{{output}}}}" + next: join + join: + type: end + output: "{{{{note_a}}}}|{{{{note_b}}}}" +"# + ); + let graph: Graph = serde_yaml::from_str(&yaml).unwrap(); + let mut ctx = make_ctx(); + let abort = create_abort_signal(); + let result = GraphExecutor::new(graph, &ws.dir) + .execute(&mut ctx, abort) + .await + .unwrap_or_else(|e| panic!("executor failed: {e:#}")); + + assert_eq!(result, "held|held"); + } + #[tokio::test] async fn map_over_list_collects_outputs_in_input_order() { if !cmd_available("python3") { diff --git a/src/graph/llm.rs b/src/graph/llm.rs index 414442f..4be1ef2 100644 --- a/src/graph/llm.rs +++ b/src/graph/llm.rs @@ -555,7 +555,7 @@ mod tests { apply_state_updates_with_output(&node, &mut state, &json!("anything")); - assert_eq!(state.state().get(OUTPUT_KEY), Some(&json!(null))); + assert!(state.state().get(OUTPUT_KEY).is_none()); } #[test] diff --git a/src/graph/map.rs b/src/graph/map.rs index 38ceb9f..7c15f9d 100644 --- a/src/graph/map.rs +++ b/src/graph/map.rs @@ -1922,6 +1922,55 @@ nodes: assert!(ctx.peer_registry.is_none() && ctx.peer_assignment.is_none()); } + /// An agent step whose `state_updates` name only other keys leaves the + /// default `output_key` unwritten, so the collector reports the missing + /// key instead of collecting a manufactured `null`. + #[tokio::test] + #[serial] + async fn map_branch_with_unrelated_state_updates_errors_instead_of_null() { + if !cmd_available("python3") { + eprintln!("skipping: python3 not available"); + return; + } + let _guard = TestConfigDirGuard::new(); + materialize_probe_agent(0.0); + let ws = TestWorkspace::new(); + let yaml = format!( + r#" +name: t +settings: + validate_before_run: false +initial_state: + items: [1] +start: fan_out +nodes: + fan_out: + type: map + over: "{{{{items}}}}" + as: item + branch: worker + collect_into: results + next: done + worker: + type: agent + agent: {PROBE_AGENT} + prompt: "p" + state_updates: + note: "{{{{output}}}}" + done: + type: end + output: "{{{{results}}}}" +"# + ); + + let chain = error_chain(run_graph(&yaml, &ws).await); + + assert!( + chain.contains("sub-branch [0] did not write output_key 'output'"), + "{chain}" + ); + } + /// The materialized agent's graph holds its script step open for seconds, /// so the 1s timer fires while that step is still pending: the agent /// future is dropped mid-flight and the chain runner alone retires the diff --git a/src/graph/state_updates.rs b/src/graph/state_updates.rs index d5f52ad..a4ea300 100644 --- a/src/graph/state_updates.rs +++ b/src/graph/state_updates.rs @@ -6,7 +6,8 @@ pub(super) const OUTPUT_KEY: &str = "output"; /// Merge `output_schema` top-level keys (when `has_schema`) and apply `state_updates` with /// `{{output}}` bound to `output` for the duration of interpolation only. An explicit `output` -/// key in `updates` is honored; otherwise `output` is restored to its prior value (or `null`). +/// key in `updates` is honored; otherwise `output` is restored to its prior value, or removed +/// again if it was absent. pub(super) fn apply( state_manager: &mut StateManager, output: &Value, @@ -37,9 +38,12 @@ pub(super) fn apply( }) .collect(); - state_manager - .state_mut() - .set(OUTPUT_KEY.into(), prev_output.unwrap_or(Value::Null)); + match prev_output { + Some(prev) => state_manager.state_mut().set(OUTPUT_KEY.into(), prev), + None => { + state_manager.state_mut().remove(OUTPUT_KEY); + } + } for (key, value) in computed { state_manager.state_mut().set(key, value); @@ -112,12 +116,23 @@ mod tests { } #[test] - fn output_is_restored_to_null_when_absent() { + fn output_is_removed_when_absent() { let u = updates(&[("response", "{{output}}")]); let mut state = manager_with(&[]); apply(&mut state, &json!("new"), false, Some(&u)); + assert_eq!(state.state().get("response"), Some(&json!("new"))); + assert!(state.state().get(OUTPUT_KEY).is_none()); + } + + #[test] + fn explicit_null_prior_output_is_restored_as_null() { + let u = updates(&[("response", "{{output}}")]); + let mut state = manager_with(&[("output", json!(null))]); + + apply(&mut state, &json!("new"), false, Some(&u)); + assert_eq!(state.state().get("response"), Some(&json!("new"))); assert_eq!(state.state().get(OUTPUT_KEY), Some(&json!(null))); } @@ -148,6 +163,6 @@ mod tests { assert!(state.state().get("context").is_none()); assert!(state.state().get("sources").is_none()); assert_eq!(state.state().get("ctx"), Some(&json!("c"))); - assert_eq!(state.state().get(OUTPUT_KEY), Some(&json!(null))); + assert!(state.state().get(OUTPUT_KEY).is_none()); } }