fix(graph): auto-merge only schema-declared keys from llm output
state_updates::apply merged EVERY top-level key of an llm/agent node's JSON
output into graph state whenever the node had an output_schema, and
structured::extract never validated the output against that schema. Any
prompt text able to steer the model — including plan/diff text pasted from a
repo under review — could therefore plant an undeclared key (e.g.
verification_commands) into state and overwrite a caller-declared value
before a downstream script consumed it.
- apply() now takes the node's output_schema and merges only keys present in
schema.properties; undeclared keys are dropped. No schema, or a schema
without an object properties map, merges nothing.
- Nodes without output_schema are unchanged (None == the old false path);
explicit state_updates are unchanged and still see the full output via
{{output.<key>}}, which remains the author-controlled way to lift a key.
- Callers: llm.rs, agent.rs (x2), rag.rs (None).
- Sweep test bundled_graph_output_schemas_declare_every_key_relied_on: for
every bundled graph, every consumed state key has a declared writer
(initial_state, variable, state_updates target, map key, script-emitted key,
or an output_schema property), and no schema node's field list names a key
its schema omits — so this restriction cannot silently break a shipped
graph. It passes with no graph.yaml changes.
- graph.example.yaml documents the declared-keys-only merge; the residual
disclosed in 665fee9's README/docstring notes is now recorded as closed.
This commit is contained in:
1 parent
300c7aeff0
commit
27101909e0
9 files changed
+370
-43
No files matched your search
@@ -8267,6 +8267,12 @@ mod tests {
|
||||
adversary.contains("--agent-variable verification_commands '[\"cargo test --all\"]'"),
|
||||
"adversary README must show the CLI form"
|
||||
);
|
||||
assert!(
|
||||
adversary.contains(
|
||||
"merges only the keys declared under `parse`'s `output_schema.properties`"
|
||||
) && !adversary.contains("Known residual"),
|
||||
"adversary README must state the schema-declared merge closes the extra-key channel"
|
||||
);
|
||||
let gauntlet = read_to_string(readmes.join("review-gauntlet/README.md")).unwrap();
|
||||
assert!(
|
||||
!gauntlet.contains("Verification commands:"),
|
||||
@@ -8287,6 +8293,10 @@ mod tests {
|
||||
runner.contains("Trust boundary:") && runner.contains("never a field an LLM"),
|
||||
"run_checks.py docstring must state the declared-variable trust boundary"
|
||||
);
|
||||
assert!(
|
||||
!runner.contains("Known residual") && !gauntlet.contains("residual:"),
|
||||
"run_checks.py and the review-gauntlet README must not disclose the closed residual"
|
||||
);
|
||||
}
|
||||
|
||||
// The spawner's own `max_agent_depth` is checked against the child's
|
||||
@@ -8558,6 +8568,207 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
// `state_updates::apply` merges only the keys an llm/agent node's
|
||||
// `output_schema.properties` declares, so a bundled graph that relied on
|
||||
// an undeclared model-emitted key reaching state would now break at
|
||||
// runtime. Every key a graph consumes must have a declared writer, and a
|
||||
// schema node's field list must not name a state key its schema omits.
|
||||
#[test]
|
||||
#[serial]
|
||||
fn bundled_graph_output_schemas_declare_every_key_relied_on() {
|
||||
use crate::graph::types::{ConcurrencyCap, Node};
|
||||
use crate::graph::{GraphParser, NodeType};
|
||||
use fancy_regex::Regex;
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
fn idents(re: &Regex, text: &str) -> BTreeSet<String> {
|
||||
re.captures_iter(text)
|
||||
.flatten()
|
||||
.filter_map(|c| c.get(1).map(|m| m.as_str().to_string()))
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn state_updates_of(node: &Node) -> Option<&HashMap<String, String>> {
|
||||
match &node.node_type {
|
||||
NodeType::Llm(n) => n.state_updates.as_ref(),
|
||||
NodeType::Agent(n) => n.state_updates.as_ref(),
|
||||
NodeType::Rag(n) => n.state_updates.as_ref(),
|
||||
NodeType::Approval(n) => n.state_updates.as_ref(),
|
||||
NodeType::Input(n) => n.state_updates.as_ref(),
|
||||
NodeType::Script(n) => n.state_updates.as_ref(),
|
||||
NodeType::End(n) => n.state_updates.as_ref(),
|
||||
NodeType::Map(_) => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn templated_fields(node: &Node) -> Vec<&str> {
|
||||
let mut fields: Vec<&str> = match &node.node_type {
|
||||
NodeType::Llm(n) => {
|
||||
let mut v = vec![n.prompt.as_str()];
|
||||
v.extend(n.instructions.as_deref());
|
||||
v
|
||||
}
|
||||
NodeType::Agent(n) => {
|
||||
let mut v = vec![n.prompt.as_str()];
|
||||
v.extend(n.inputs.iter().flat_map(|m| m.values().map(String::as_str)));
|
||||
v
|
||||
}
|
||||
NodeType::Rag(n) => {
|
||||
let mut v: Vec<&str> = n.documents.iter().map(String::as_str).collect();
|
||||
v.extend(n.query.as_deref());
|
||||
v.extend(n.extractor_prompt.as_deref());
|
||||
v
|
||||
}
|
||||
NodeType::Approval(n) => vec![n.question.as_str()],
|
||||
NodeType::Input(n) => {
|
||||
let mut v = vec![n.question.as_str()];
|
||||
v.extend(n.default.as_deref());
|
||||
v
|
||||
}
|
||||
NodeType::End(n) => vec![n.output.as_str()],
|
||||
NodeType::Map(n) => {
|
||||
let mut v = vec![n.over.as_str()];
|
||||
if let Some(ConcurrencyCap::Template(t)) = &n.max_concurrency {
|
||||
v.push(t);
|
||||
}
|
||||
v
|
||||
}
|
||||
NodeType::Script(_) => Vec::new(),
|
||||
};
|
||||
fields.extend(
|
||||
state_updates_of(node)
|
||||
.into_iter()
|
||||
.flat_map(|m| m.values().map(String::as_str)),
|
||||
);
|
||||
fields
|
||||
}
|
||||
|
||||
fn schema_node_text(node: &Node) -> Option<(Option<&serde_json::Value>, String)> {
|
||||
match &node.node_type {
|
||||
NodeType::Llm(n) => Some((
|
||||
n.output_schema.as_ref(),
|
||||
format!("{}\n{}", n.instructions.as_deref().unwrap_or(""), n.prompt),
|
||||
)),
|
||||
NodeType::Agent(n) => Some((n.output_schema.as_ref(), n.prompt.clone())),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
let template_root = Regex::new(r"\{\{\s*([A-Za-z_][A-Za-z0-9_]*)").unwrap();
|
||||
let script_read =
|
||||
Regex::new(r#"state(?:\.get\(|\[)\s*["']([A-Za-z_][A-Za-z0-9_]*)["']"#).unwrap();
|
||||
let script_dict_key = Regex::new(r#"["']([A-Za-z_][A-Za-z0-9_]*)["']\s*:"#).unwrap();
|
||||
let script_subscript_assign =
|
||||
Regex::new(r#"\[["']([A-Za-z_][A-Za-z0-9_]*)["']\]\s*=[^=]"#).unwrap();
|
||||
let backticked_field = Regex::new(r"(?m)^\s*-\s*`([a-z][a-z0-9_]*)`").unwrap();
|
||||
|
||||
// Seeded by the engine rather than any graph author: `initial_prompt`
|
||||
// (dispatch) and the per-node scoped `output`/`choice`/`input` bindings.
|
||||
let engine_seeded = ["initial_prompt", "output", "choice", "input"];
|
||||
|
||||
let _guard = TestConfigDirGuard::new();
|
||||
Agent::install_builtin_agents(false).unwrap();
|
||||
|
||||
let mut checked = Vec::new();
|
||||
for entry in std::fs::read_dir(paths::agents_data_dir()).unwrap() {
|
||||
let dir = entry.unwrap().path();
|
||||
let graph_path = dir.join("graph.yaml");
|
||||
if !graph_path.exists() {
|
||||
continue;
|
||||
}
|
||||
let name = dir.file_name().unwrap().to_string_lossy().to_string();
|
||||
let graph = GraphParser::new(&dir)
|
||||
.load_from_file(&graph_path)
|
||||
.unwrap_or_else(|e| panic!("graph.yaml for '{name}' failed to parse: {e}"));
|
||||
|
||||
let mut scripts = String::new();
|
||||
if let Ok(entries) = std::fs::read_dir(dir.join("scripts")) {
|
||||
for script in entries.map(|e| e.unwrap().path()) {
|
||||
if script.extension().is_some_and(|ext| ext == "py") {
|
||||
scripts.push_str(&read_to_string(&script).unwrap());
|
||||
scripts.push('\n');
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let mut writers: BTreeSet<String> =
|
||||
engine_seeded.iter().map(|k| k.to_string()).collect();
|
||||
writers.extend(graph.initial_state.keys().cloned());
|
||||
writers.extend(graph.variables.iter().map(|v| v.name.clone()));
|
||||
writers.extend(idents(&script_dict_key, &scripts));
|
||||
writers.extend(idents(&script_subscript_assign, &scripts));
|
||||
|
||||
let mut consumed = idents(&script_read, &scripts);
|
||||
let mut declared: BTreeMap<&str, BTreeSet<String>> = BTreeMap::new();
|
||||
for (id, node) in &graph.nodes {
|
||||
for field in templated_fields(node) {
|
||||
consumed.extend(idents(&template_root, field));
|
||||
}
|
||||
writers.extend(
|
||||
state_updates_of(node)
|
||||
.into_iter()
|
||||
.flat_map(|m| m.keys().cloned()),
|
||||
);
|
||||
if let NodeType::Map(m) = &node.node_type {
|
||||
writers.extend([
|
||||
m.as_name.clone(),
|
||||
m.output_key.clone(),
|
||||
m.collect_into.clone(),
|
||||
]);
|
||||
}
|
||||
if let Some((Some(schema), _)) = schema_node_text(node) {
|
||||
let properties = schema
|
||||
.get("properties")
|
||||
.and_then(serde_json::Value::as_object)
|
||||
.unwrap_or_else(|| {
|
||||
panic!(
|
||||
"{name}/{id}: output_schema has no object `properties`, so the \
|
||||
engine would merge nothing from its output: {schema}"
|
||||
)
|
||||
});
|
||||
declared.insert(id, properties.keys().cloned().collect());
|
||||
}
|
||||
}
|
||||
let schema_keys: BTreeSet<&String> = declared.values().flatten().collect();
|
||||
|
||||
let unwritten: Vec<&String> = consumed
|
||||
.iter()
|
||||
.filter(|k| !writers.contains(*k) && !schema_keys.contains(k))
|
||||
.collect();
|
||||
assert!(
|
||||
unwritten.is_empty(),
|
||||
"{name}: state keys consumed with no declared writer (initial_state, variable, \
|
||||
state_updates target, map key, script-emitted key, or output_schema property): \
|
||||
{unwritten:?}"
|
||||
);
|
||||
|
||||
for (id, properties) in &declared {
|
||||
let (_, text) = schema_node_text(&graph.nodes[*id]).unwrap();
|
||||
let omitted: Vec<String> = idents(&backticked_field, &text)
|
||||
.into_iter()
|
||||
.filter(|k| {
|
||||
(consumed.contains(k) || schema_keys.contains(&k))
|
||||
&& !properties.contains(k)
|
||||
&& !writers.contains(k)
|
||||
})
|
||||
.collect();
|
||||
assert!(
|
||||
omitted.is_empty(),
|
||||
"{name}/{id}: the node's field list names state keys its output_schema does \
|
||||
not declare, so the engine would drop them when the model emits them: \
|
||||
{omitted:?}"
|
||||
);
|
||||
}
|
||||
checked.push(name);
|
||||
}
|
||||
for expected in ["adversary", "review-gauntlet"] {
|
||||
assert!(
|
||||
checked.iter().any(|n| n == expected),
|
||||
"expected bundled graph agent '{expected}' to be checked; found {checked:?}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
// ---- adversary suite-script regression tests ----
|
||||
//
|
||||
// Fault paths of the adversary's verdict/gate scripts, exercised by
|
||||
@@ -10452,8 +10663,14 @@ mod tests {
|
||||
.output_schema
|
||||
.as_ref()
|
||||
.unwrap_or_else(|| panic!("{name}: parse must have an output_schema"));
|
||||
let properties = schema["properties"].as_object().unwrap_or_else(|| {
|
||||
panic!(
|
||||
"{name}: parse output_schema must declare `properties`; the engine merges \
|
||||
only declared keys into state: {schema}"
|
||||
)
|
||||
});
|
||||
assert!(
|
||||
schema["properties"].get("verification_commands").is_none(),
|
||||
!properties.contains_key("verification_commands"),
|
||||
"{name}: parse must not extract verification_commands: {schema}"
|
||||
);
|
||||
assert!(
|
||||
|
||||
Reference in new issue
Block a user