fix(graph): honor an explicit output key in state_updates

Hoist the three copies of apply_state_updates (llm, agent, rag) into
graph::state_updates::apply(state, output, has_schema, updates). The
helper computes every state_updates value first, restores the transient
`output` binding, and only then writes the computed updates, so a
`state_updates: { output: "{{output}}" }` entry is no longer clobbered
by the restore step.

The per-file wrappers stay as one-line calls so their existing tests are
unchanged; new direct helper tests live in state_updates.rs, including
the regression test for the clobber.

Design: plans/map-subgraph-branches-design.md §2.6, §3.7 (C5, D3).

src/graph/agent.rs

@@ -1,4 +1,5 @@
 use super::state::StateManager;
+use super::state_updates;
 use super::structured;
 use super::types::AgentNode;
 use crate::config::RequestContext;
@@ -8,7 +9,6 @@ use serde_json::Value;
 use std::time::Duration;
 use tokio::time::timeout;
 
-const OUTPUT_KEY: &str = "output";
 const DEFAULT_TIMEOUT_SECS: u64 = 300;
 
 pub struct AgentNodeExecutor;
@@ -58,37 +58,12 @@ impl AgentNodeExecutor {
 }
 
 fn apply_state_updates(node: &AgentNode, state_manager: &mut StateManager, output: &Value) {
-    if node.output_schema.is_some()
+    state_updates::apply(
-        && let Some(obj) = output.as_object()
+        state_manager,
-    {
+        output,
-        for (k, v) in obj {
+        node.output_schema.is_some(),
-            state_manager.state_mut().set(k.clone(), v.clone());
+        node.state_updates.as_ref(),
-        }
+    );
-    }
-
-    let Some(updates) = &node.state_updates else {
-        return;
-    };
-    let prev_output = state_manager.state().get(OUTPUT_KEY).cloned();
-    state_manager
-        .state_mut()
-        .set(OUTPUT_KEY.into(), output.clone());
-
-    for (key, template) in updates {
-        let value = state_manager.interpolate_lenient(template);
-        state_manager
-            .state_mut()
-            .set(key.clone(), Value::String(value));
-    }
-
-    match prev_output {
-        Some(v) => state_manager.state_mut().set(OUTPUT_KEY.into(), v),
-        None => {
-            state_manager
-                .state_mut()
-                .set(OUTPUT_KEY.into(), Value::Null);
-        }
-    }
 }
 
 #[cfg(test)]

src/graph/llm.rs

@@ -1,4 +1,7 @@
 use super::state::StateManager;
+use super::state_updates;
+#[cfg(test)]
+use super::state_updates::OUTPUT_KEY;
 use super::structured;
 use super::types::LlmNode;
 use crate::client::{Model, ModelType, call_chat_completions};
@@ -18,8 +21,6 @@ use std::sync::Arc;
 use std::time::Duration;
 use tokio::time::timeout;
 
-const OUTPUT_KEY: &str = "output";
-
 #[derive(Debug, PartialEq, Eq)]
 pub(super) enum LlmExecutionOutcome {
     Continue,
@@ -456,37 +457,12 @@ fn apply_state_updates_with_output(
     state_manager: &mut StateManager,
     output: &Value,
 ) {
-    if node.output_schema.is_some()
+    state_updates::apply(
-        && let Some(obj) = output.as_object()
+        state_manager,
-    {
+        output,
-        for (k, v) in obj {
+        node.output_schema.is_some(),
-            state_manager.state_mut().set(k.clone(), v.clone());
+        node.state_updates.as_ref(),
-        }
+    );
-    }
-
-    let Some(updates) = &node.state_updates else {
-        return;
-    };
-    let prev_output = state_manager.state().get(OUTPUT_KEY).cloned();
-    state_manager
-        .state_mut()
-        .set(OUTPUT_KEY.into(), output.clone());
-
-    for (key, template) in updates {
-        let value = state_manager.interpolate_lenient(template);
-        state_manager
-            .state_mut()
-            .set(key.clone(), Value::String(value));
-    }
-
-    match prev_output {
-        Some(v) => state_manager.state_mut().set(OUTPUT_KEY.into(), v),
-        None => {
-            state_manager
-                .state_mut()
-                .set(OUTPUT_KEY.into(), Value::Null);
-        }
-    }
 }
 
 fn format_schema_hint(schema: &Value) -> String {

src/graph/mod.rs

@@ -10,6 +10,7 @@ pub mod reducer;
 pub mod script;
 pub mod staging;
 pub mod state;
+pub mod state_updates;
 pub mod structured;
 pub mod types;
 pub mod user_interaction;

src/graph/rag.rs

@@ -1,4 +1,5 @@
 use super::state::StateManager;
+use super::state_updates;
 use super::types::RagNode;
 use crate::config::RequestContext;
 use crate::utils::create_abort_signal;
@@ -7,7 +8,6 @@ use serde_json::{Map, Value};
 use std::time::Duration;
 use tokio::time::timeout;
 
-const OUTPUT_KEY: &str = "output";
 const DEFAULT_QUERY: &str = "{{initial_prompt}}";
 const DEFAULT_RAG_TIMEOUT_SECS: u64 = 120;
 
@@ -70,27 +70,7 @@ fn build_rag_output(context: String, sources_str: &str) -> Value {
 }
 
 fn apply_state_updates(node: &RagNode, state_manager: &mut StateManager, output: &Value) {
-    let Some(updates) = &node.state_updates else {
+    state_updates::apply(state_manager, output, false, node.state_updates.as_ref());
-        return;
-    };
-    let prev_output = state_manager.state().get(OUTPUT_KEY).cloned();
-    state_manager
-        .state_mut()
-        .set(OUTPUT_KEY.into(), output.clone());
-
-    for (key, template) in updates {
-        let value = state_manager.interpolate_lenient(template);
-        state_manager
-            .state_mut()
-            .set(key.clone(), Value::String(value));
-    }
-
-    match prev_output {
-        Some(v) => state_manager.state_mut().set(OUTPUT_KEY.into(), v),
-        None => state_manager
-            .state_mut()
-            .set(OUTPUT_KEY.into(), Value::Null),
-    }
 }
 
 #[cfg(test)]

src/graph/state_updates.rs

@@ -0,0 +1,153 @@
+use super::state::StateManager;
+use serde_json::Value;
+use std::collections::HashMap;
+
+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`).
+pub(super) fn apply(
+    state_manager: &mut StateManager,
+    output: &Value,
+    has_schema: bool,
+    updates: Option<&HashMap<String, String>>,
+) {
+    if has_schema && let Some(obj) = output.as_object() {
+        for (k, v) in obj {
+            state_manager.state_mut().set(k.clone(), v.clone());
+        }
+    }
+
+    let Some(updates) = updates else {
+        return;
+    };
+    let prev_output = state_manager.state().get(OUTPUT_KEY).cloned();
+    state_manager
+        .state_mut()
+        .set(OUTPUT_KEY.into(), output.clone());
+
+    let computed: Vec<(String, Value)> = updates
+        .iter()
+        .map(|(key, template)| {
+            (
+                key.clone(),
+                Value::String(state_manager.interpolate_lenient(template)),
+            )
+        })
+        .collect();
+
+    state_manager
+        .state_mut()
+        .set(OUTPUT_KEY.into(), prev_output.unwrap_or(Value::Null));
+
+    for (key, value) in computed {
+        state_manager.state_mut().set(key, value);
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use serde_json::json;
+
+    fn manager_with(pairs: &[(&str, Value)]) -> StateManager {
+        let mut map = HashMap::new();
+        for (k, v) in pairs {
+            map.insert((*k).into(), v.clone());
+        }
+        StateManager::new(map)
+    }
+
+    fn updates(pairs: &[(&str, &str)]) -> HashMap<String, String> {
+        pairs
+            .iter()
+            .map(|(k, v)| ((*k).to_string(), (*v).to_string()))
+            .collect()
+    }
+
+    #[test]
+    fn explicit_output_state_update_is_honored() {
+        let u = updates(&[("output", "{{output}}")]);
+        let mut state = manager_with(&[("output", json!("stale"))]);
+
+        apply(&mut state, &json!("the completion"), false, Some(&u));
+
+        assert_eq!(
+            state.state().get(OUTPUT_KEY),
+            Some(&json!("the completion"))
+        );
+    }
+
+    #[test]
+    fn explicit_output_is_honored_alongside_other_updates() {
+        let u = updates(&[("output", "{{output}}"), ("summary", "got: {{output}}")]);
+        let mut state = manager_with(&[]);
+
+        apply(&mut state, &json!("done"), false, Some(&u));
+
+        assert_eq!(state.state().get(OUTPUT_KEY), Some(&json!("done")));
+        assert_eq!(state.state().get("summary"), Some(&json!("got: done")));
+    }
+
+    #[test]
+    fn no_updates_is_a_noop() {
+        let mut state = manager_with(&[("keep", json!(1))]);
+
+        apply(&mut state, &json!("ignored"), false, None);
+
+        assert_eq!(state.state().get("keep"), Some(&json!(1)));
+        assert!(state.state().get(OUTPUT_KEY).is_none());
+    }
+
+    #[test]
+    fn output_is_restored_to_previous_value() {
+        let u = updates(&[("response", "{{output}}")]);
+        let mut state = manager_with(&[("output", json!("preserved"))]);
+
+        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!("preserved")));
+    }
+
+    #[test]
+    fn output_is_restored_to_null_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_eq!(state.state().get(OUTPUT_KEY), Some(&json!(null)));
+    }
+
+    #[test]
+    fn schema_merges_top_level_keys_when_has_schema() {
+        let mut state = manager_with(&[]);
+
+        apply(&mut state, &json!({"summary": "s", "score": 7}), true, None);
+
+        assert_eq!(state.state().get("summary"), Some(&json!("s")));
+        assert_eq!(state.state().get("score"), Some(&json!(7)));
+        assert!(state.state().get(OUTPUT_KEY).is_none());
+    }
+
+    #[test]
+    fn object_output_is_not_merged_without_schema() {
+        let u = updates(&[("ctx", "{{output.context}}")]);
+        let mut state = manager_with(&[]);
+
+        apply(
+            &mut state,
+            &json!({"context": "c", "sources": ["a"]}),
+            false,
+            Some(&u),
+        );
+
+        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)));
+    }
+}