feat(graph): seed a child graph agent's state from an agent node's inputs map
Agent nodes gain an optional `inputs:` map of templates resolved against
the parent state and overlaid on the child graph agent's initial_state
before its executor starts. A lone {{key}} keeps the state value's JSON
type; mixed text renders to a string; initial_prompt is still owned by the
dispatcher. run_child_agent and run_active_agent_graph keep their
signatures as wrappers over the new *_with_inputs variants, and
run_agent_for_graph rejects inputs aimed at a config-only agent or
containing initial_prompt.
Design: plans/map-subgraph-branches-design.md §3.12 (C18), task T5.
This commit is contained in:
8 files changed
+283
-8
No files matched your search
+33
-3
@@ -18,6 +18,7 @@ use indexmap::IndexMap;
|
||||
use log::{debug, warn};
|
||||
use parking_lot::RwLock;
|
||||
use serde_json::{Value, json};
|
||||
use std::collections::HashMap;
|
||||
use std::pin::Pin;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
@@ -535,16 +536,29 @@ pub async fn handle_agent_tool(
|
||||
}
|
||||
|
||||
pub fn run_child_agent(
|
||||
mut child_ctx: RequestContext,
|
||||
child_ctx: RequestContext,
|
||||
initial_input: Input,
|
||||
abort_signal: AbortSignal,
|
||||
) -> Pin<Box<dyn Future<Output = Result<String>> + Send>> {
|
||||
run_child_agent_with_graph_inputs(child_ctx, initial_input, abort_signal, None)
|
||||
}
|
||||
|
||||
/// Runs a child agent to completion. When the child is a graph agent,
|
||||
/// `graph_inputs` is overlaid on its `initial_state` before the graph
|
||||
/// starts; it is ignored for config-only agents.
|
||||
fn run_child_agent_with_graph_inputs(
|
||||
mut child_ctx: RequestContext,
|
||||
initial_input: Input,
|
||||
abort_signal: AbortSignal,
|
||||
graph_inputs: Option<HashMap<String, Value>>,
|
||||
) -> Pin<Box<dyn Future<Output = Result<String>> + Send>> {
|
||||
Box::pin(async move {
|
||||
if graph::active_agent_graph_name(&child_ctx).is_some() {
|
||||
return graph::run_active_agent_graph(
|
||||
return graph::run_active_agent_graph_with_inputs(
|
||||
&mut child_ctx,
|
||||
&initial_input.text(),
|
||||
abort_signal,
|
||||
graph_inputs,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
@@ -685,10 +699,13 @@ fn effective_max_agent_depth(parent_ctx: &RequestContext) -> usize {
|
||||
/// output. This is similar to `handle_spawn` but runs the child agent in the
|
||||
/// current task (no tokio::spawn, no supervisor handle registration) so the
|
||||
/// graph executor can sequence agent nodes directly.
|
||||
/// `graph_inputs` seeds the child's graph state and is only accepted when the
|
||||
/// target is a graph agent.
|
||||
pub async fn run_agent_for_graph(
|
||||
parent_ctx: &mut RequestContext,
|
||||
agent_name: &str,
|
||||
prompt: &str,
|
||||
graph_inputs: Option<HashMap<String, Value>>,
|
||||
) -> Result<String> {
|
||||
let peer_assignment = parent_ctx.peer_assignment.take();
|
||||
let assigned_peer = peer_assignment.is_some();
|
||||
@@ -759,6 +776,19 @@ pub async fn run_agent_for_graph(
|
||||
);
|
||||
child_ctx.rag = agent.rag();
|
||||
child_ctx.agent = Some(agent);
|
||||
if graph_inputs.is_some() && graph::active_agent_graph_name(&child_ctx).is_none() {
|
||||
bail!(
|
||||
"agent '{agent_name}' has no graph.yaml; `inputs:` is only valid on agent nodes that target a graph agent"
|
||||
);
|
||||
}
|
||||
if graph_inputs
|
||||
.as_ref()
|
||||
.is_some_and(|inputs| inputs.contains_key("initial_prompt"))
|
||||
{
|
||||
bail!(
|
||||
"agent '{agent_name}': `inputs.initial_prompt` is reserved (the dispatcher seeds it from `prompt:`)"
|
||||
);
|
||||
}
|
||||
if assigned_peer {
|
||||
child_ctx.peer_registry = peer_registry.clone();
|
||||
}
|
||||
@@ -786,7 +816,7 @@ pub async fn run_agent_for_graph(
|
||||
|
||||
debug!("Spawning agent '{agent_name}' for graph node as '{agent_id}'");
|
||||
|
||||
run_child_agent(child_ctx, input, child_abort).await
|
||||
run_child_agent_with_graph_inputs(child_ctx, input, child_abort, graph_inputs).await
|
||||
}
|
||||
|
||||
async fn populate_agent_mcp_runtime(ctx: &mut RequestContext, server_ids: &[String]) -> Result<()> {
|
||||
|
||||
+97
-1
@@ -5,7 +5,9 @@ use super::types::AgentNode;
|
||||
use crate::config::RequestContext;
|
||||
use crate::function::agents::run_agent_for_graph;
|
||||
use anyhow::{Context, Result};
|
||||
use log::debug;
|
||||
use serde_json::Value;
|
||||
use std::collections::HashMap;
|
||||
use std::time::Duration;
|
||||
use tokio::time::timeout;
|
||||
|
||||
@@ -24,6 +26,13 @@ impl AgentNodeExecutor {
|
||||
.interpolate(&node.prompt)
|
||||
.with_context(|| format!("Failed to interpolate prompt for agent '{}'", node.agent))?;
|
||||
|
||||
let graph_inputs = resolve_inputs(node, state_manager)?;
|
||||
if let Some(inputs) = &graph_inputs {
|
||||
let mut keys: Vec<&String> = inputs.keys().collect();
|
||||
keys.sort_unstable();
|
||||
debug!("Agent '{}' graph inputs: {keys:?}", node.agent);
|
||||
}
|
||||
|
||||
let timeout_dur = Duration::from_secs(node.timeout.unwrap_or(DEFAULT_TIMEOUT_SECS));
|
||||
|
||||
// run_agent_for_graph takes the identity off the ctx; keep a handle so
|
||||
@@ -39,7 +48,7 @@ impl AgentNodeExecutor {
|
||||
|
||||
let raw_result = timeout(
|
||||
timeout_dur,
|
||||
run_agent_for_graph(parent_ctx, &node.agent, &prompt),
|
||||
run_agent_for_graph(parent_ctx, &node.agent, &prompt, graph_inputs),
|
||||
)
|
||||
.await;
|
||||
if retire_peer_on_return && let Some((registry, id)) = &peer {
|
||||
@@ -73,6 +82,30 @@ impl AgentNodeExecutor {
|
||||
}
|
||||
}
|
||||
|
||||
/// Resolves each `inputs` template against the parent state. A lone
|
||||
/// `{{key}}` yields the state value as-is (numbers, arrays, objects and
|
||||
/// `null` all survive); anything else renders to a string. Every key is
|
||||
/// strict: a missing reference fails the node before the child starts.
|
||||
fn resolve_inputs(
|
||||
node: &AgentNode,
|
||||
state_manager: &StateManager,
|
||||
) -> Result<Option<HashMap<String, Value>>> {
|
||||
let Some(inputs) = &node.inputs else {
|
||||
return Ok(None);
|
||||
};
|
||||
let mut resolved = HashMap::with_capacity(inputs.len());
|
||||
for (key, template) in inputs {
|
||||
let value = state_manager.interpolate_raw(template).with_context(|| {
|
||||
format!(
|
||||
"Failed to interpolate inputs.{key} for agent '{}'",
|
||||
node.agent
|
||||
)
|
||||
})?;
|
||||
resolved.insert(key.clone(), value);
|
||||
}
|
||||
Ok(Some(resolved))
|
||||
}
|
||||
|
||||
fn apply_state_updates(node: &AgentNode, state_manager: &mut StateManager, output: &Value) {
|
||||
state_updates::apply(
|
||||
state_manager,
|
||||
@@ -107,6 +140,7 @@ mod tests {
|
||||
state_updates: updates,
|
||||
output_schema: None,
|
||||
timeout: None,
|
||||
inputs: None,
|
||||
teammates: false,
|
||||
}
|
||||
}
|
||||
@@ -155,6 +189,68 @@ mod tests {
|
||||
assert!(!registry.is_finished(&id));
|
||||
}
|
||||
|
||||
fn node_with_inputs(pairs: &[(&str, &str)]) -> AgentNode {
|
||||
let mut node = node_with("hi", None);
|
||||
node.inputs = Some(
|
||||
pairs
|
||||
.iter()
|
||||
.map(|(k, v)| ((*k).to_string(), (*v).to_string()))
|
||||
.collect(),
|
||||
);
|
||||
node
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_inputs_is_none_when_node_has_no_inputs() {
|
||||
let node = node_with("hi", None);
|
||||
let state = manager_with(&[("n", json!(3))]);
|
||||
|
||||
assert_eq!(resolve_inputs(&node, &state).unwrap(), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_inputs_lone_reference_keeps_the_array_value() {
|
||||
let node = node_with_inputs(&[("items", "{{list}}")]);
|
||||
let state = manager_with(&[("list", json!(["a", "b"]))]);
|
||||
|
||||
let inputs = resolve_inputs(&node, &state).unwrap().unwrap();
|
||||
|
||||
assert_eq!(inputs.get("items"), Some(&json!(["a", "b"])));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_inputs_mixed_text_renders_to_a_string() {
|
||||
let node = node_with_inputs(&[("label", "n={{n}}")]);
|
||||
let state = manager_with(&[("n", json!(3))]);
|
||||
|
||||
let inputs = resolve_inputs(&node, &state).unwrap().unwrap();
|
||||
|
||||
assert_eq!(inputs.get("label"), Some(&json!("n=3")));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_inputs_lone_reference_to_null_yields_null() {
|
||||
let node = node_with_inputs(&[("x", "{{k}}")]);
|
||||
let state = manager_with(&[("k", Value::Null)]);
|
||||
|
||||
let inputs = resolve_inputs(&node, &state).unwrap().unwrap();
|
||||
|
||||
assert_eq!(inputs.get("x"), Some(&Value::Null));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_inputs_missing_key_errors_naming_the_input() {
|
||||
let node = node_with_inputs(&[("width", "{{nope}}")]);
|
||||
let state = manager_with(&[]);
|
||||
|
||||
let err = resolve_inputs(&node, &state).expect_err("missing key must fail");
|
||||
|
||||
let chain = format!("{err:#}");
|
||||
assert!(chain.contains("inputs.width"), "{chain}");
|
||||
assert!(chain.contains("for agent 'test_agent'"), "{chain}");
|
||||
assert!(chain.contains("'nope' not found in state"), "{chain}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn state_updates_use_output_placeholder() {
|
||||
let node = {
|
||||
|
||||
+102
-3
@@ -5,6 +5,7 @@ use crate::utils::AbortSignal;
|
||||
use anyhow::{Context, Result, anyhow};
|
||||
use log::info;
|
||||
use serde_json::Value;
|
||||
use std::collections::HashMap;
|
||||
|
||||
pub fn active_agent_graph_name(ctx: &RequestContext) -> Option<String> {
|
||||
let name = ctx.agent.as_ref()?.name().to_string();
|
||||
@@ -15,6 +16,18 @@ pub async fn run_active_agent_graph(
|
||||
ctx: &mut RequestContext,
|
||||
prompt: &str,
|
||||
abort_signal: AbortSignal,
|
||||
) -> Result<String> {
|
||||
run_active_agent_graph_with_inputs(ctx, prompt, abort_signal, None).await
|
||||
}
|
||||
|
||||
/// Runs the active agent's graph with `inputs` overlaid on its
|
||||
/// `initial_state` before the executor starts. `None` seeds exactly what
|
||||
/// `run_active_agent_graph` does: the prompt under `initial_prompt`.
|
||||
pub async fn run_active_agent_graph_with_inputs(
|
||||
ctx: &mut RequestContext,
|
||||
prompt: &str,
|
||||
abort_signal: AbortSignal,
|
||||
inputs: Option<HashMap<String, Value>>,
|
||||
) -> Result<String> {
|
||||
let agent_name =
|
||||
active_agent_graph_name(ctx).ok_or_else(|| anyhow!("Active agent has no graph.yaml"))?;
|
||||
@@ -29,9 +42,7 @@ pub async fn run_active_agent_graph(
|
||||
.load_from_file(&graph_path)
|
||||
.with_context(|| format!("Failed to load graph.yaml for agent '{agent_name}'"))?;
|
||||
|
||||
graph
|
||||
.initial_state
|
||||
.insert("initial_prompt".into(), Value::String(prompt.to_string()));
|
||||
seed_initial_state(&mut graph.initial_state, prompt, inputs);
|
||||
|
||||
let executor = GraphExecutor::new(graph, agent_dir);
|
||||
let output = executor
|
||||
@@ -45,3 +56,91 @@ pub async fn run_active_agent_graph(
|
||||
|
||||
Ok(output)
|
||||
}
|
||||
|
||||
/// Seeds a child graph's `initial_state`, later writers winning: the graph's
|
||||
/// own YAML defaults, then the caller's `inputs`, then `initial_prompt`,
|
||||
/// which the dispatcher always owns.
|
||||
pub fn seed_initial_state(
|
||||
initial_state: &mut HashMap<String, Value>,
|
||||
prompt: &str,
|
||||
inputs: Option<HashMap<String, Value>>,
|
||||
) {
|
||||
if let Some(inputs) = inputs {
|
||||
initial_state.extend(inputs);
|
||||
}
|
||||
initial_state.insert("initial_prompt".into(), Value::String(prompt.to_string()));
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use serde_json::json;
|
||||
|
||||
fn yaml_defaults() -> HashMap<String, Value> {
|
||||
HashMap::from([
|
||||
("width".to_string(), json!(4)),
|
||||
("mode".to_string(), json!("fast")),
|
||||
])
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn seed_initial_state_without_inputs_only_adds_initial_prompt() {
|
||||
let mut state = yaml_defaults();
|
||||
|
||||
seed_initial_state(&mut state, "hello", None);
|
||||
|
||||
let mut expected = yaml_defaults();
|
||||
expected.insert("initial_prompt".into(), json!("hello"));
|
||||
assert_eq!(state, expected);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn seed_initial_state_without_inputs_on_empty_state_inserts_exactly_one_key() {
|
||||
let mut state = HashMap::new();
|
||||
|
||||
seed_initial_state(&mut state, "hello", None);
|
||||
|
||||
assert_eq!(
|
||||
state,
|
||||
HashMap::from([("initial_prompt".to_string(), json!("hello"))])
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn seed_initial_state_inputs_override_yaml_defaults_and_keep_the_rest() {
|
||||
let mut state = yaml_defaults();
|
||||
let inputs = HashMap::from([
|
||||
("width".to_string(), json!(2)),
|
||||
("extra".to_string(), json!(["a", "b"])),
|
||||
]);
|
||||
|
||||
seed_initial_state(&mut state, "hello", Some(inputs));
|
||||
|
||||
assert_eq!(state.get("width"), Some(&json!(2)));
|
||||
assert_eq!(state.get("mode"), Some(&json!("fast")));
|
||||
assert_eq!(state.get("extra"), Some(&json!(["a", "b"])));
|
||||
assert_eq!(state.get("initial_prompt"), Some(&json!("hello")));
|
||||
assert_eq!(state.len(), 4);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn seed_initial_state_dispatcher_prompt_beats_an_initial_prompt_input() {
|
||||
let mut state = HashMap::from([("initial_prompt".to_string(), json!("from yaml"))]);
|
||||
let inputs = HashMap::from([("initial_prompt".to_string(), json!("from inputs"))]);
|
||||
|
||||
seed_initial_state(&mut state, "from dispatcher", Some(inputs));
|
||||
|
||||
assert_eq!(state.get("initial_prompt"), Some(&json!("from dispatcher")));
|
||||
assert_eq!(state.len(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn seed_initial_state_null_input_overlays_a_default() {
|
||||
let mut state = yaml_defaults();
|
||||
let inputs = HashMap::from([("width".to_string(), Value::Null)]);
|
||||
|
||||
seed_initial_state(&mut state, "hello", Some(inputs));
|
||||
|
||||
assert_eq!(state.get("width"), Some(&Value::Null));
|
||||
}
|
||||
}
|
||||
@@ -660,6 +660,7 @@ mod tests {
|
||||
state_updates: None,
|
||||
output_schema: None,
|
||||
timeout: None,
|
||||
inputs: None,
|
||||
teammates,
|
||||
}),
|
||||
next: None,
|
||||
|
||||
@@ -421,6 +421,7 @@ mod tests {
|
||||
state_updates: None,
|
||||
output_schema: None,
|
||||
timeout: None,
|
||||
inputs: None,
|
||||
teammates,
|
||||
})
|
||||
}
|
||||
|
||||
+3
-1
@@ -16,7 +16,9 @@ pub mod types;
|
||||
pub mod user_interaction;
|
||||
pub mod validator;
|
||||
|
||||
pub use dispatch::{active_agent_graph_name, run_active_agent_graph};
|
||||
pub use dispatch::{
|
||||
active_agent_graph_name, run_active_agent_graph, run_active_agent_graph_with_inputs,
|
||||
};
|
||||
pub use executor::GraphExecutor;
|
||||
pub use parser::{GraphParser, agent_has_graph};
|
||||
use serde_json::Value;
|
||||
|
||||
@@ -240,6 +240,9 @@ pub struct AgentNode {
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub timeout: Option<u64>,
|
||||
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub inputs: Option<HashMap<String, String>>,
|
||||
|
||||
#[serde(default)]
|
||||
pub teammates: bool,
|
||||
}
|
||||
@@ -634,6 +637,48 @@ nodes:
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn agent_node_inputs_parse_optionally_and_empty_map_is_some() {
|
||||
let yaml = r#"
|
||||
name: g
|
||||
start: with_inputs
|
||||
nodes:
|
||||
with_inputs:
|
||||
id: with_inputs
|
||||
type: agent
|
||||
agent: helper
|
||||
prompt: hi
|
||||
inputs:
|
||||
width: "{{n}}"
|
||||
note: "Repo: {{repo}}"
|
||||
next: without_inputs
|
||||
without_inputs:
|
||||
id: without_inputs
|
||||
type: agent
|
||||
agent: helper
|
||||
prompt: hi
|
||||
next: empty_inputs
|
||||
empty_inputs:
|
||||
id: empty_inputs
|
||||
type: agent
|
||||
agent: helper
|
||||
prompt: hi
|
||||
inputs: {}
|
||||
"#;
|
||||
let graph: Graph = serde_yaml::from_str(yaml).unwrap();
|
||||
let inputs_of = |id: &str| match &graph.get_node(id).unwrap().node_type {
|
||||
NodeType::Agent(n) => n.inputs.clone(),
|
||||
_ => panic!("expected Agent variant"),
|
||||
};
|
||||
|
||||
let with = inputs_of("with_inputs").expect("inputs should parse");
|
||||
assert_eq!(with.len(), 2);
|
||||
assert_eq!(with.get("width").map(String::as_str), Some("{{n}}"));
|
||||
assert_eq!(with.get("note").map(String::as_str), Some("Repo: {{repo}}"));
|
||||
assert_eq!(inputs_of("without_inputs"), None);
|
||||
assert_eq!(inputs_of("empty_inputs"), Some(HashMap::new()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn deserializes_mcp_tools_at_graph_and_node_level() {
|
||||
let yaml = r#"
|
||||
|
||||
@@ -1790,6 +1790,7 @@ mod tests {
|
||||
state_updates: None,
|
||||
output_schema: None,
|
||||
timeout: None,
|
||||
inputs: None,
|
||||
teammates: false,
|
||||
}),
|
||||
next: next.map(NextTargets::from),
|
||||
|
||||
Reference in new issue
Block a user