fix(px): when+return short-circuits procedure execution

crates/pluresdb-px/src/px/executor.rs

@@ -430,7 +430,23 @@ fn execute_when(
     let mut nested_results = Vec::new();
     for (i, nested) in nested_steps.iter().enumerate() {
         let result = execute_step(nested, i, vars, handler)?;
+        let is_return = result.kind == "return";
         nested_results.push(result);
+        if is_return {
+            break;
+        }
+    }
+
+    // Propagate a return step so the procedure executor short-circuits.
+    if let Some(last) = nested_results.last() {
+        if last.kind == "return" {
+            return Ok(StepResult {
+                index,
+                kind: "return".into(),
+                output: last.output.clone(),
+                skipped: false,
+            });
+        }
     }
 
     // Return the last nested result as the when step's output
@@ -3563,6 +3579,103 @@ mod tests {
         assert!(result.step_results[0].skipped);
     }
 
+    #[test]
+    fn when_return_short_circuits_procedure() {
+        // A `return` inside a `when` block must stop all subsequent steps.
+        let handler = MockHandler::new().with_result("side_effect", json!("ran"));
+
+        let procedure = json!({
+            "type": "procedure",
+            "name": "guard_test",
+            "steps": [
+                {
+                    "kind": "when",
+                    "condition": "input == empty",
+                    "steps": [
+                        { "kind": "return", "value": "no_input" }
+                    ]
+                },
+                // This step must NOT execute when the when-return fires.
+                { "kind": "call", "name": "side_effect", "params": {} }
+            ]
+        });
+
+        let vars = HashMap::from([("input".to_string(), json!("empty"))]);
+        let result = execute_with_vars(&procedure, &handler, vars).unwrap();
+
+        // Only the when step should appear; the call after must be absent.
+        assert_eq!(result.step_results.len(), 1);
+        assert_eq!(result.step_results[0].kind, "return");
+        assert_eq!(result.step_results[0].output, Some(json!("no_input")));
+    }
+
+    #[test]
+    fn when_condition_false_does_not_short_circuit() {
+        // When the when-condition is false the return is skipped and the rest continues.
+        let handler = MockHandler::new().with_result("next_step", json!("did_run"));
+
+        let procedure = json!({
+            "type": "procedure",
+            "name": "no_guard",
+            "steps": [
+                {
+                    "kind": "when",
+                    "condition": "input == empty",
+                    "steps": [
+                        { "kind": "return", "value": "no_input" }
+                    ]
+                },
+                { "kind": "call", "name": "next_step", "params": {} }
+            ]
+        });
+
+        let vars = HashMap::from([("input".to_string(), json!("hello"))]);
+        let result = execute_with_vars(&procedure, &handler, vars).unwrap();
+
+        // Both the (skipped) when and the call step should be present.
+        assert_eq!(result.step_results.len(), 2);
+        assert!(result.step_results[0].skipped); // when was skipped
+        assert_eq!(result.step_results[1].output, Some(json!("did_run")));
+    }
+
+    #[test]
+    fn when_return_nested_when_also_short_circuits() {
+        // A return inside a nested when should propagate all the way up.
+        let handler = MockHandler::new().with_result("outer_step", json!("outer_ran"));
+
+        let procedure = json!({
+            "type": "procedure",
+            "name": "nested_guard",
+            "steps": [
+                {
+                    "kind": "when",
+                    "condition": "outer == true",
+                    "steps": [
+                        {
+                            "kind": "when",
+                            "condition": "inner == true",
+                            "steps": [
+                                { "kind": "return", "value": "inner_return" }
+                            ]
+                        }
+                    ]
+                },
+                { "kind": "call", "name": "outer_step", "params": {} }
+            ]
+        });
+
+        let vars = HashMap::from([
+            ("outer".to_string(), json!(true)),
+            ("inner".to_string(), json!(true)),
+        ]);
+        let result = execute_with_vars(&procedure, &handler, vars).unwrap();
+
+        // The outer when propagates the return; outer_step never runs.
+        assert_eq!(result.step_results.len(), 1);
+        assert_eq!(result.step_results[0].kind, "return");
+        assert_eq!(result.step_results[0].output, Some(json!("inner_return")));
+    }
+
     #[test]
     fn execute_match_selects_first_true_arm() {
         let handler = MockHandler::new();

crates/pluresdb-px/src/px/scenario_runner.rs

@@ -23,7 +23,7 @@ use std::collections::HashMap;
 use serde::{Deserialize, Serialize};
 use serde_json::Value;
 
-use super::executor::{ActionHandler, ExecutionError};
+use super::executor::{ActionHandler, ExecutionError, execute_with_vars};
 
 // ── Types ─────────────────────────────────────────────────────────────────────
 
@@ -349,18 +349,38 @@ pub fn run_scenario(
         vars.extend(run_params);
 
         if let Some(proc_data) = procedures.get(&proc_name) {
-            if let Some(steps) = proc_data.get("steps").and_then(|v| v.as_array()) {
-                for step in steps {
-                    if let Err(e) = execute_step(step, &handler, &mut vars) {
-                        return ScenarioResult {
-                            name,
-                            given,
-                            passed: false,
-                            expectations: vec![],
-                            error: Some(format!("procedure '{proc_name}' failed: {e}")),
-                            duration_ms: start.elapsed().as_millis() as u64,
-                        };
+            match execute_with_vars(proc_data, &handler, vars) {
+                Ok(result) => {
+                    // The executor's `emit` step stores events in `result.variables["emit"]`
+                    // rather than calling `handler.call("emit", ...)`.  Replay them through
+                    // the handler so ScenarioActionHandler captures them in emitted_events.
+                    if let Some(Value::Array(events)) =
+                        result.variables.get("emit").cloned()
+                    {
+                        for event in &events {
+                            if let Err(e) = handler.call("emit", event) {
+                                return ScenarioResult {
+                                    name,
+                                    given,
+                                    passed: false,
+                                    expectations: vec![],
+                                    error: Some(format!("emit failed: {e}")),
+                                    duration_ms: start.elapsed().as_millis() as u64,
+                                };
+                            }
+                        }
                     }
+                    vars = result.variables;
+                }
+                Err(e) => {
+                    return ScenarioResult {
+                        name,
+                        given,
+                        passed: false,
+                        expectations: vec![],
+                        error: Some(format!("procedure '{proc_name}' failed: {e}")),
+                        duration_ms: start.elapsed().as_millis() as u64,
+                    };
                 }
             }
         } else {
@@ -957,4 +977,83 @@ mod tests {
         let result = run_scenario(&scenario, &procedures, &BuiltinChecker);
         assert!(result.passed, "expected pass, got: {:?}", result);
     }
+
+    #[test]
+    fn when_return_short_circuits_procedure_in_scenario() {
+        // A `return` inside a `when` block must stop procedure execution.
+        // The entry must NOT be stored when the guard fires.
+        let mut procedures = HashMap::new();
+        procedures.insert(
+            "guard_proc".to_string(),
+            json!({
+                "name": "guard_proc",
+                "steps": [
+                    // Bind result to the return value via echo so var_equals can check it.
+                    // The when fires (input == empty), returns "no_input".
+                    {
+                        "kind": "when",
+                        "condition": "input == empty",
+                        "steps": [
+                            {"kind": "call", "name": "echo", "params": "no_input", "output_var": "result"},
+                            {"kind": "return", "value": "no_input"}
+                        ]
+                    },
+                    // These steps must NOT run.
+                    {"kind": "call", "name": "put_entry", "params": {"key": "should_not_exist", "value": 1}},
+                    {"kind": "call", "name": "echo", "params": "input_present", "output_var": "result"}
+                ]
+            }),
+        );
+
+        let scenario = json!({
+            "name": "guard_short_circuit",
+            "setup": [
+                {"kind": "call", "name": "echo", "params": "empty", "output_var": "input"}
+            ],
+            "run": "guard_proc",
+            "expectations": [
+                {"check": "has_entry", "params": {"key": "should_not_exist"}, "negated": true},
+                {"check": "var_equals", "params": {"var": "result", "value": "no_input"}, "negated": false}
+            ]
+        });
+
+        let result = run_scenario(&scenario, &procedures, &BuiltinChecker);
+        assert!(result.passed, "expected pass, got: {:?}", result);
+    }
+
+    #[test]
+    fn when_condition_false_continues_execution() {
+        // When the when-condition is false the return is NOT triggered and execution continues.
+        let scenario = json!({
+            "name": "no_guard_trigger",
+            "setup": [
+                {"kind": "call", "name": "echo", "params": "hello", "output_var": "input"}
+            ],
+            "run": "guard_proc",
+            "expectations": [
+                {"check": "has_entry", "params": {"key": "did_run"}, "negated": false}
+            ]
+        });
+
+        let mut procedures = HashMap::new();
+        procedures.insert(
+            "guard_proc".to_string(),
+            json!({
+                "name": "guard_proc",
+                "steps": [
+                    {
+                        "kind": "when",
+                        "condition": "input == empty",
+                        "steps": [
+                            {"kind": "return", "value": "no_input"}
+                        ]
+                    },
+                    {"kind": "call", "name": "put_entry", "params": {"key": "did_run", "value": 1}}
+                ]
+            }),
+        );
+
+        let result = run_scenario(&scenario, &procedures, &BuiltinChecker);
+        assert!(result.passed, "expected pass, got: {:?}", result);
+    }
 }