From 27101909e09f18b6dde8dbf4842bf6cffb2d3cc1 Mon Sep 17 00:00:00 2001 From: Dark-Alex-17 Date: Sun, 13 Sep 2026 16:43:08 -0600 Subject: [PATCH] fix(graph): auto-merge only schema-declared keys from llm output MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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.}}, 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. --- assets/agents/adversary/README.md | 7 +- assets/agents/adversary/scripts/run_checks.py | 8 +- assets/agents/review-gauntlet/README.md | 7 +- graph.example.yaml | 14 +- src/config/request_context.rs | 219 +++++++++++++++++- src/graph/agent.rs | 22 +- src/graph/llm.rs | 17 +- src/graph/rag.rs | 2 +- src/graph/state_updates.rs | 117 ++++++++-- 9 files changed, 370 insertions(+), 43 deletions(-) diff --git a/assets/agents/adversary/README.md b/assets/agents/adversary/README.md index 792bb9e..544bd8b 100644 --- a/assets/agents/adversary/README.md +++ b/assets/agents/adversary/README.md @@ -139,10 +139,9 @@ as a structured `inputs:` passthrough, not as prompt text. an invalid declaration: `run_checks` executes **nothing** and records a `PIPELINE-FAULT: verification_commands declaration invalid — …` marker, which forces DIVERGES — fail-closed, not degraded to the soft ENVIRONMENT marker. -- Known residual: the engine merges every top-level key of `parse`'s JSON output into state, not - only the keys in its `output_schema`, so an injected instruction that makes the parse LLM emit an - extra `verification_commands` key is not yet blocked at the engine level. Closing that requires - restricting the llm-node auto-merge to `output_schema.properties` (tracked as a follow-up). +- The engine merges only the keys declared under `parse`'s `output_schema.properties` into state, + so an injected instruction that makes the parse LLM emit an extra `verification_commands` key + cannot reach state: the undeclared key is dropped and the declared variable stands. ### Tools diff --git a/assets/agents/adversary/scripts/run_checks.py b/assets/agents/adversary/scripts/run_checks.py index 029719f..7d46bc2 100644 --- a/assets/agents/adversary/scripts/run_checks.py +++ b/assets/agents/adversary/scripts/run_checks.py @@ -18,10 +18,10 @@ shell=True. Variables land in state as strings, so a JSON-encoded list is accepted alongside a real list. A declaration that is not a JSON array of strings is a caller contract violation, not a runner hiccup: it executes NOTHING and records a PIPELINE-FAULT (fail-closed into DIVERGES) instead of -degrading to the soft ENVIRONMENT marker. Known residual: the engine merges -every top-level key of an llm node's JSON output into state, so a parse LLM -coaxed into emitting an extra `verification_commands` key could still -overwrite the declared value before this script runs. +degrading to the soft ENVIRONMENT marker. The engine merges only the keys an +llm node's `output_schema.properties` declares, so a parse LLM coaxed into +emitting an extra `verification_commands` key cannot overwrite the declared +value: the undeclared key is dropped before it reaches state. Budget: commands run sequentially, so TOTAL runtime is bounded by a deadline (TOTAL_DEADLINE_SECS = 3300s, the run_checks node's 3600s timeout minus diff --git a/assets/agents/review-gauntlet/README.md b/assets/agents/review-gauntlet/README.md index 897e9f9..82c698c 100644 --- a/assets/agents/review-gauntlet/README.md +++ b/assets/agents/review-gauntlet/README.md @@ -133,9 +133,10 @@ final report verbatim. and never extracted by `parse`. The adversary's `run_checks` stage executes those commands with a shell, which makes the declaration a trust boundary: the prompt also carries plan text pasted from the repo under review, so - commands must never be inferred from it (residual: the engine still merges - every top-level key of `parse`'s output into state, so an extra - `verification_commands` key emitted by the LLM is not yet blocked). + commands must never be inferred from it, and the engine merges only the keys + `parse`'s `output_schema.properties` declares, so an extra + `verification_commands` key emitted by the LLM is dropped before it reaches + state. Undeclared stays undeclared (`'[]'` ⇒ "none declared"); a declaration that is not a JSON array of strings executes nothing and records a `PIPELINE-FAULT`, which the diff --git a/graph.example.yaml b/graph.example.yaml index fa018b9..7a29cbd 100644 --- a/graph.example.yaml +++ b/graph.example.yaml @@ -213,8 +213,10 @@ nodes: # agent/graph/global default when omitted. skill_instructions: null # Per-node skill-hint text override; uses the built-in default when omitted. output_schema: # Optional JSON Schema. The output is parsed to JSON - type: object # and its top-level object keys auto-merge into state - properties: # (so `topic` / `needs_deep_dive` become {{topic}} etc). + type: object # and only the keys declared under `properties` auto-merge + properties: # into state (so `topic` / `needs_deep_dive` become {{topic}} + # etc). Undeclared keys the model emits are dropped; lift + # one explicitly via state_updates: `{{output.}}`. topic: { type: string } needs_deep_dive: { type: boolean } required: [topic, needs_deep_dive] @@ -284,8 +286,9 @@ nodes: type: array items: { type: string } required: [context, sources] - # When `output_schema` is set, top-level keys auto-merge into state, so - # `context` and `sources` are produced without needing `state_updates`. + # When `output_schema` is set, the keys declared under `properties` + # auto-merge into state, so `context` and `sources` are produced without + # needing `state_updates`. Keys the schema does not declare are dropped. next: synthesize # Joins with retrieve at `synthesize`. # --- script node (the diamond's join; also dispatches) ----------------- @@ -369,7 +372,8 @@ nodes: # and is retired when the chain ends. A message delivered # after the chain's last flagged agent step is accepted # but never read. - output_schema: # Optional. Same extraction as llm nodes + output_schema: # Optional. Same extraction as llm nodes: only the keys + # declared under `properties` auto-merge into state. type: object properties: summary: { type: string } diff --git a/src/config/request_context.rs b/src/config/request_context.rs index 79ef479..c564699 100644 --- a/src/config/request_context.rs +++ b/src/config/request_context.rs @@ -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 { + 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> { + 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 = + 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> = 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 = 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!( diff --git a/src/graph/agent.rs b/src/graph/agent.rs index 7faf413..201011f 100644 --- a/src/graph/agent.rs +++ b/src/graph/agent.rs @@ -71,7 +71,7 @@ fn outcome_from( state_updates::apply( state_manager, &Value::String(format!("Agent node failed: {e:#}")), - node.output_schema.is_some(), + node.output_schema.as_ref(), node.state_updates.as_ref(), ); Ok(AgentExecutionOutcome::FellBack(fb.clone())) @@ -309,7 +309,7 @@ fn apply_state_updates(node: &AgentNode, state_manager: &mut StateManager, outpu state_updates::apply( state_manager, output, - node.output_schema.is_some(), + node.output_schema.as_ref(), node.state_updates.as_ref(), ); } @@ -1090,7 +1090,11 @@ mod tests { #[test] fn output_schema_auto_merges_top_level_keys() { - let node = node_with_schema("hi", None, json!({"type": "object"})); + let node = node_with_schema( + "hi", + None, + json!({"type": "object", "properties": {"goal": {}, "summary": {}}}), + ); let mut state = manager_with(&[]); let output = json!({"goal": "do X", "summary": "details"}); @@ -1102,7 +1106,11 @@ mod tests { #[test] fn output_schema_preserves_nested_value_types() { - let node = node_with_schema("hi", None, json!({"type": "object"})); + let node = node_with_schema( + "hi", + None, + json!({"type": "object", "properties": {"tags": {}, "config": {}, "count": {}}}), + ); let mut state = manager_with(&[]); let output = json!({ "tags": ["a", "b"], @@ -1121,7 +1129,11 @@ mod tests { fn output_schema_explicit_state_updates_override_auto_merge() { let mut u = HashMap::new(); u.insert("goal".into(), "renamed-{{output.goal}}".into()); - let node = node_with_schema("hi", Some(u), json!({"type": "object"})); + let node = node_with_schema( + "hi", + Some(u), + json!({"type": "object", "properties": {"goal": {}}}), + ); let mut state = manager_with(&[]); let output = json!({"goal": "do X"}); diff --git a/src/graph/llm.rs b/src/graph/llm.rs index 614df7c..6e6b7b3 100644 --- a/src/graph/llm.rs +++ b/src/graph/llm.rs @@ -520,7 +520,7 @@ fn apply_state_updates_with_output( state_updates::apply( state_manager, output, - node.output_schema.is_some(), + node.output_schema.as_ref(), node.state_updates.as_ref(), ); } @@ -671,7 +671,10 @@ mod tests { #[test] fn output_schema_auto_merges_top_level_keys() { - let node = node_with_schema(None, json!({"type": "object"})); + let node = node_with_schema( + None, + json!({"type": "object", "properties": {"goal": {}, "summary": {}}}), + ); let mut state = manager_with(&[]); let output = json!({"goal": "do X", "summary": "details"}); @@ -683,7 +686,10 @@ mod tests { #[test] fn output_schema_preserves_nested_value_types() { - let node = node_with_schema(None, json!({"type": "object"})); + let node = node_with_schema( + None, + json!({"type": "object", "properties": {"tags": {}, "config": {}, "count": {}}}), + ); let mut state = manager_with(&[]); let output = json!({ "tags": ["a", "b"], @@ -702,7 +708,10 @@ mod tests { fn output_schema_explicit_state_updates_override_auto_merge() { let mut u = HashMap::new(); u.insert("goal".into(), "renamed-{{output.goal}}".into()); - let node = node_with_schema(Some(u), json!({"type": "object"})); + let node = node_with_schema( + Some(u), + json!({"type": "object", "properties": {"goal": {}}}), + ); let mut state = manager_with(&[]); let output = json!({"goal": "do X"}); diff --git a/src/graph/rag.rs b/src/graph/rag.rs index fbba2aa..ffae3e7 100644 --- a/src/graph/rag.rs +++ b/src/graph/rag.rs @@ -68,7 +68,7 @@ fn build_rag_output(context: String, sources_str: &str) -> Value { } fn apply_state_updates(node: &RagNode, state_manager: &mut StateManager, output: &Value) { - state_updates::apply(state_manager, output, false, node.state_updates.as_ref()); + state_updates::apply(state_manager, output, None, node.state_updates.as_ref()); } #[cfg(test)] diff --git a/src/graph/state_updates.rs b/src/graph/state_updates.rs index a4ea300..12fc388 100644 --- a/src/graph/state_updates.rs +++ b/src/graph/state_updates.rs @@ -4,19 +4,32 @@ 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 removed -/// again if it was absent. +/// Merge the `output` keys declared under `schema["properties"]` and apply `state_updates` with +/// `{{output}}` bound to the full, unfiltered `output` for the duration of interpolation only. +/// +/// The schema is the graph author's declaration of what the model may write to state; any other +/// key in the output is model- (and therefore prompt-) controlled and is dropped rather than +/// merged implicitly. No `schema`, or a schema without an object `properties`, merges nothing. +/// Authors who want an undeclared key can still lift it explicitly via +/// `state_updates: { key: '{{output.key}}' }`. +/// +/// An explicit `output` key in `updates` is honored; otherwise `output` is restored to its prior +/// value, or removed again if it was absent. pub(super) fn apply( state_manager: &mut StateManager, output: &Value, - has_schema: bool, + schema: Option<&Value>, updates: Option<&HashMap>, ) { - if has_schema && let Some(obj) = output.as_object() { + if let Some(declared) = schema + .and_then(|s| s.get("properties")) + .and_then(Value::as_object) + && let Some(obj) = output.as_object() + { for (k, v) in obj { - state_manager.state_mut().set(k.clone(), v.clone()); + if declared.contains_key(k) { + state_manager.state_mut().set(k.clone(), v.clone()); + } } } @@ -75,7 +88,7 @@ mod tests { let u = updates(&[("output", "{{output}}")]); let mut state = manager_with(&[("output", json!("stale"))]); - apply(&mut state, &json!("the completion"), false, Some(&u)); + apply(&mut state, &json!("the completion"), None, Some(&u)); assert_eq!( state.state().get(OUTPUT_KEY), @@ -88,7 +101,7 @@ mod tests { let u = updates(&[("output", "{{output}}"), ("summary", "got: {{output}}")]); let mut state = manager_with(&[]); - apply(&mut state, &json!("done"), false, Some(&u)); + apply(&mut state, &json!("done"), None, Some(&u)); assert_eq!(state.state().get(OUTPUT_KEY), Some(&json!("done"))); assert_eq!(state.state().get("summary"), Some(&json!("got: done"))); @@ -98,7 +111,7 @@ mod tests { fn no_updates_is_a_noop() { let mut state = manager_with(&[("keep", json!(1))]); - apply(&mut state, &json!("ignored"), false, None); + apply(&mut state, &json!("ignored"), None, None); assert_eq!(state.state().get("keep"), Some(&json!(1))); assert!(state.state().get(OUTPUT_KEY).is_none()); @@ -109,7 +122,7 @@ mod tests { let u = updates(&[("response", "{{output}}")]); let mut state = manager_with(&[("output", json!("preserved"))]); - apply(&mut state, &json!("new"), false, Some(&u)); + apply(&mut state, &json!("new"), None, Some(&u)); assert_eq!(state.state().get("response"), Some(&json!("new"))); assert_eq!(state.state().get(OUTPUT_KEY), Some(&json!("preserved"))); @@ -120,7 +133,7 @@ mod tests { let u = updates(&[("response", "{{output}}")]); let mut state = manager_with(&[]); - apply(&mut state, &json!("new"), false, Some(&u)); + apply(&mut state, &json!("new"), None, Some(&u)); assert_eq!(state.state().get("response"), Some(&json!("new"))); assert!(state.state().get(OUTPUT_KEY).is_none()); @@ -131,23 +144,95 @@ mod tests { let u = updates(&[("response", "{{output}}")]); let mut state = manager_with(&[("output", json!(null))]); - apply(&mut state, &json!("new"), false, Some(&u)); + apply(&mut state, &json!("new"), None, Some(&u)); assert_eq!(state.state().get("response"), Some(&json!("new"))); assert_eq!(state.state().get(OUTPUT_KEY), Some(&json!(null))); } + fn schema_with(keys: &[&str]) -> Value { + let props: serde_json::Map = keys + .iter() + .map(|k| ((*k).to_string(), json!({"type": "string"}))) + .collect(); + json!({"type": "object", "properties": props}) + } + #[test] - fn schema_merges_top_level_keys_when_has_schema() { + fn schema_merges_declared_top_level_keys() { + let schema = schema_with(&["summary", "score"]); let mut state = manager_with(&[]); - apply(&mut state, &json!({"summary": "s", "score": 7}), true, None); + apply( + &mut state, + &json!({"summary": "s", "score": 7}), + Some(&schema), + 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 schema_drops_undeclared_keys_and_preserves_existing_state() { + let schema = schema_with(&["summary"]); + let mut state = manager_with(&[("verification_commands", json!(["cargo test"]))]); + + apply( + &mut state, + &json!({"summary": "s", "verification_commands": ["rm -rf /"]}), + Some(&schema), + None, + ); + + assert_eq!(state.state().get("summary"), Some(&json!("s"))); + assert_eq!( + state.state().get("verification_commands"), + Some(&json!(["cargo test"])) + ); + } + + #[test] + fn none_schema_merges_nothing() { + let mut state = manager_with(&[]); + + apply(&mut state, &json!({"summary": "s"}), None, None); + + assert!(state.state().get("summary").is_none()); + assert!(state.state().get(OUTPUT_KEY).is_none()); + } + + #[test] + fn schema_without_properties_merges_nothing() { + let schema = json!({"type": "object"}); + let mut state = manager_with(&[]); + + apply(&mut state, &json!({"summary": "s"}), Some(&schema), None); + + assert!(state.state().get("summary").is_none()); + } + + #[test] + fn state_updates_still_see_undeclared_output_keys() { + let schema = schema_with(&["summary"]); + let u = updates(&[("lifted", "{{output.extra}}")]); + let mut state = manager_with(&[]); + + apply( + &mut state, + &json!({"summary": "s", "extra": "e"}), + Some(&schema), + Some(&u), + ); + + assert_eq!(state.state().get("summary"), Some(&json!("s"))); + assert!(state.state().get("extra").is_none()); + assert_eq!(state.state().get("lifted"), Some(&json!("e"))); + assert!(state.state().get(OUTPUT_KEY).is_none()); + } + #[test] fn object_output_is_not_merged_without_schema() { let u = updates(&[("ctx", "{{output.context}}")]); @@ -156,7 +241,7 @@ mod tests { apply( &mut state, &json!({"context": "c", "sources": ["a"]}), - false, + None, Some(&u), );