diff --git a/src/function/agents.rs b/src/function/agents.rs index d991d73..d8a98ba 100644 --- a/src/function/agents.rs +++ b/src/function/agents.rs @@ -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> + 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>, ) -> Pin> + 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>, ) -> Result { 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<()> { diff --git a/src/graph/agent.rs b/src/graph/agent.rs index 34e1087..e2b9608 100644 --- a/src/graph/agent.rs +++ b/src/graph/agent.rs @@ -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>> { + 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 = { diff --git a/src/graph/dispatch.rs b/src/graph/dispatch.rs index a65b36f..5cda4ec 100644 --- a/src/graph/dispatch.rs +++ b/src/graph/dispatch.rs @@ -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 { 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 { + 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>, ) -> Result { 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, + prompt: &str, + inputs: Option>, +) { + 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 { + 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)); + } +} diff --git a/src/graph/executor.rs b/src/graph/executor.rs index e96ad88..66a4822 100644 --- a/src/graph/executor.rs +++ b/src/graph/executor.rs @@ -660,6 +660,7 @@ mod tests { state_updates: None, output_schema: None, timeout: None, + inputs: None, teammates, }), next: None, diff --git a/src/graph/map.rs b/src/graph/map.rs index a2ee287..e069b3c 100644 --- a/src/graph/map.rs +++ b/src/graph/map.rs @@ -421,6 +421,7 @@ mod tests { state_updates: None, output_schema: None, timeout: None, + inputs: None, teammates, }) } diff --git a/src/graph/mod.rs b/src/graph/mod.rs index 1706c59..0c5830f 100644 --- a/src/graph/mod.rs +++ b/src/graph/mod.rs @@ -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; diff --git a/src/graph/types.rs b/src/graph/types.rs index ddc4ac8..3b36c6e 100644 --- a/src/graph/types.rs +++ b/src/graph/types.rs @@ -240,6 +240,9 @@ pub struct AgentNode { #[serde(default, skip_serializing_if = "Option::is_none")] pub timeout: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub inputs: Option>, + #[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#" diff --git a/src/graph/validator.rs b/src/graph/validator.rs index 51d2093..01816bd 100644 --- a/src/graph/validator.rs +++ b/src/graph/validator.rs @@ -1790,6 +1790,7 @@ mod tests { state_updates: None, output_schema: None, timeout: None, + inputs: None, teammates: false, }), next: next.map(NextTargets::from),