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).
This commit is contained in:
Dark-Alex-17 committed 2026-09-11 10:56:26 -06:00
1 parent 1df82b50b6
commit 24088a48a1
5 files changed
+172 -87

No files matched your search

+7 -32
View File
@@ -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()
&& let Some(obj) = output.as_object()
{
for (k, v) in obj {
state_manager.state_mut().set(k.clone(), v.clone());
}
}
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);
}
}
state_updates::apply(
state_manager,
output,
node.output_schema.is_some(),
node.state_updates.as_ref(),
);
}
#[cfg(test)]