Preserve the ErrorContext cause chain across a defer on the unwind path (B-611)

baml_language/crates/baml_compiler2_mir/src/lower.rs

@@ -5782,15 +5782,18 @@ impl LoweringContext<'_> {
 
             AstExpr::Throw { value } => {
                 let val_op = self.lower_throw_operand(value);
-                if let Some(catch_ctx) = &self.catch_context {
-                    // Inside a catch block: store the value into the error
-                    // local and jump to the handler instead of unwinding.
-                    let error_local = catch_ctx.error_local;
-                    let unwind_target = catch_ctx.unwind_target;
-                    self.builder
-                        .assign(Place::Local(error_local), Rvalue::Use(val_op));
-                    self.builder.goto(unwind_target);
-                } else if self.operand_is_marked_rethrow(&val_op) {
+                // Route every throw through the exception funnel (like
+                // `AstStmt::Throw`) rather than a static jump to
+                // `catch_context.unwind_target`. The funnel computes the
+                // BEP-042 cause chain (`find_cause_context`) and materializes
+                // the destination handler's `ErrorContext`; a static goto
+                // bypasses both, so a `throw` in expression position inside a
+                // `defer` region (or a `catch` arm/base) would drop its cause
+                // and leave a bound `ctx` unmaterialized (B-611). The exception
+                // table routes the throw to the same innermost handler the
+                // static jump targeted — its region covers this PC — so control
+                // flow is unchanged.
+                if self.operand_is_marked_rethrow(&val_op) {
                     self.builder.rethrow(val_op);
                 } else {
                     self.builder.throw(val_op);

baml_language/crates/baml_tests/tests/errorcontext.rs

@@ -237,3 +237,188 @@ function main() -> string {
     );
     assert_eq!(expect_string(output.result.unwrap()), "E");
 }
+
+/// B-611 — a non-throwing `defer` on the unwind frame must NOT wipe the
+/// propagating error's cause chain.
+///
+/// `mid` catches "origin" and throws "wrap" (so wrap.cause == origin). A
+/// `defer {}` armed in `mid` routes that expression-position throw through its
+/// landing pad, which replays the (empty) defer and then re-raises "wrap"
+/// unchanged. Before the fix the re-raise nulled the cause, so `root_cause`
+/// stopped at "wrap"; now the pad's transparent re-raise preserves the cause
+/// computed at wrap's throw site (`thrown_value_causes` in `vm.rs`), so
+/// `root_cause` still walks past "wrap" to the original "origin".
+#[tokio::test]
+async fn defer_on_unwind_preserves_cause_chain() {
+    let output = baml_test!(
+        r#"
+function fail_low() -> string { throw "origin" }
+
+function mid() -> string {
+  defer { }
+  fail_low() catch (e, ctx) {
+    _ => throw "wrap"
+  }
+}
+
+function main() -> string {
+  mid() catch (e, ctx) {
+    _ => {
+      match (ctx.root_cause().error) {
+        let s: string => s
+        _ => "LOST: cause chain wiped by the defer re-raise"
+      }
+    }
+  }
+}
+"#
+    );
+    assert_eq!(expect_string(output.result.unwrap()), "origin");
+}
+
+/// B-611 — the rendered chain must keep its "During handling of..." section
+/// across the defer re-raise. Same shape as above, but return `to_string`: the
+/// full chain (origin, then wrap) must survive, not just the surviving "wrap".
+#[tokio::test]
+async fn defer_on_unwind_preserves_to_string_chain() {
+    let output = baml_test!(
+        r#"
+function fail_low() -> string { throw "origin" }
+
+function mid() -> string {
+  defer { }
+  fail_low() catch (e, ctx) {
+    _ => throw "wrap"
+  }
+}
+
+function main() -> string {
+  mid() catch (e, ctx) {
+    _ => ctx.to_string()
+  }
+}
+"#
+    );
+    let rendered = expect_string(output.result.unwrap());
+    assert!(
+        rendered.contains("During handling of the above error, another error occurred"),
+        "missing chain separator; the defer re-raise dropped the cause:\n{rendered}"
+    );
+    assert!(
+        rendered.contains("origin"),
+        "missing the original error frame:\n{rendered}"
+    );
+    assert!(
+        rendered.contains("wrap"),
+        "missing the superseding error frame:\n{rendered}"
+    );
+}
+
+/// B-611 — a NON-empty (but non-throwing) defer body must behave the same: it
+/// is the pad's transparent re-raise, not empty-block elimination, that must
+/// preserve the cause. The defer mutates an outer local (a real, observable
+/// body) yet does not throw.
+#[tokio::test]
+async fn defer_with_body_preserves_cause_chain() {
+    let output = baml_test!(
+        r#"
+function fail_low() -> string { throw "origin" }
+
+function mid() -> string {
+  let marker = "start"
+  defer { marker = marker + "-cleanup" }
+  fail_low() catch (e, ctx) {
+    _ => throw "wrap"
+  }
+}
+
+function main() -> string {
+  mid() catch (e, ctx) {
+    _ => {
+      match (ctx.root_cause().error) {
+        let s: string => s
+        _ => "LOST: cause chain wiped by the defer re-raise"
+      }
+    }
+  }
+}
+"#
+    );
+    assert_eq!(expect_string(output.result.unwrap()), "origin");
+}
+
+/// B-611 — two stacked non-throwing defers on the unwind frame each add a
+/// transparent re-raise; the cause must survive all of them.
+#[tokio::test]
+async fn nested_defers_nonthrowing_preserve_cause() {
+    let output = baml_test!(
+        r#"
+function fail_low() -> string { throw "origin" }
+
+function mid() -> string {
+  defer { }
+  defer { }
+  fail_low() catch (e, ctx) {
+    _ => throw "wrap"
+  }
+}
+
+function main() -> string {
+  mid() catch (e, ctx) {
+    _ => {
+      match (ctx.root_cause().error) {
+        let s: string => s
+        _ => "LOST: cause chain wiped by a stacked defer re-raise"
+      }
+    }
+  }
+}
+"#
+    );
+    assert_eq!(expect_string(output.result.unwrap()), "origin");
+}
+
+/// B-611 regression guard — a defer whose body ITSELF throws must still build
+/// the full chain. `mid` catches "origin" and throws "wrap"; the armed defer
+/// then throws "cleanup" while "wrap" is unwinding. The result chains
+/// cleanup -> wrap -> origin (cleanup is "during handling of" wrap, wrap is
+/// "during handling of" origin). `root_cause` reaches "origin" and the middle
+/// "wrap" link is not dropped — confirming the fix preserves a pre-existing
+/// chain even when the defer contributes a fresh link on top of it.
+#[tokio::test]
+async fn defer_that_throws_still_chains_through_wrap() {
+    let output = baml_test!(
+        r#"
+function fail_low() -> string { throw "origin" }
+
+function mid() -> string {
+  defer { throw "cleanup" }
+  fail_low() catch (e, ctx) {
+    _ => throw "wrap"
+  }
+}
+
+function main() -> string {
+  mid() catch (e, ctx) {
+    _ => ctx.to_string()
+  }
+}
+"#
+    );
+    let rendered = expect_string(output.result.unwrap());
+    // Oldest -> newest: origin, wrap, cleanup — all three present, in order.
+    let origin_at = rendered.find("origin").expect("missing 'origin' frame");
+    let wrap_at = rendered.find("wrap").expect("missing 'wrap' frame");
+    let cleanup_at = rendered.find("cleanup").expect("missing 'cleanup' frame");
+    assert!(
+        origin_at < wrap_at && wrap_at < cleanup_at,
+        "cause chain out of order or a link dropped:\n{rendered}"
+    );
+    assert_eq!(
+        rendered
+            .matches("During handling of the above error, another error occurred")
+            .count(),
+        2,
+        "expected two chain separators (origin->wrap->cleanup):\n{rendered}"
+    );
+}

baml_language/crates/bex_vm/src/vm.rs

@@ -426,6 +426,7 @@ mod tests {
             pending_call_captures: Vec::new(),
             call_input_capture_hook: None,
             seen_throw_values: Vec::new(),
+            thrown_value_causes: Vec::new(),
             bex_ref_seed: None,
             id_overrides: Vec::new(),
             argv: Arc::from([]),
@@ -863,6 +864,17 @@ pub struct BexVm {
     /// instead of raw bits so GC forwarding can preserve rethrow identity.
     seen_throw_values: Vec<Value>,
 
+    /// BEP-042 cause chain across a transparent re-raise: the `cause` context
+    /// computed at each error value's original (fresh) throw site, keyed by the
+    /// thrown value. A later *rethrow* of the same value (a non-throwing
+    /// `defer` pad re-raising the in-flight error, the no-match fall-through,
+    /// or `throw <binding>`) reuses this instead of re-running the cause walk —
+    /// which from inside a handler body would self-link — so the chain survives
+    /// the re-raise. Stored as values (not raw bits) so GC forwarding preserves
+    /// both key and cause identity; deduplicated by value and cleared each
+    /// `finalize`, like `seen_throw_values`.
+    thrown_value_causes: Vec<(Value, Value)>,
+
     /// Constants for building BEX `CallRef`s on demand (the `$id` surface):
     /// `(process_euid, engine_id)`, set once by the engine when it attaches
     /// identity to this VM. Unconditional — `$id` works with profiling off.
@@ -1307,6 +1319,7 @@ impl BexVm {
             pending_call_captures: Vec::new(),
             call_input_capture_hook: None,
             seen_throw_values: Vec::new(),
+            thrown_value_causes: Vec::new(),
             bex_ref_seed: None,
             id_overrides: Vec::new(),
             argv,
@@ -1399,6 +1412,39 @@ impl BexVm {
         true
     }
 
+    /// Remember the `cause` context computed at `value`'s original (fresh)
+    /// throw site so a later rethrow of the same value can recover it. The
+    /// stored cause is the error `value` superseded (the enclosing handler's
+    /// context), NEVER `value`'s own materialized context, so it can never form
+    /// a self-link. Keyed by value identity (like `seen_throw_values`).
+    fn record_throw_cause(&mut self, value: Value, cause: Value) {
+        // A fresh throw with no enclosing handler carries no chain; skip it so
+        // the map only holds values that actually supersede an error (a missing
+        // entry looks up as `Value::NULL` anyway, preserving the deliberate
+        // null cause for genuine user/no-match rethrows).
+        if cause == Value::NULL {
+            return;
+        }
+        if let Some((_, existing)) = self
+            .thrown_value_causes
+            .iter_mut()
+            .find(|(v, _)| *v == value)
+        {
+            *existing = cause;
+        } else {
+            self.thrown_value_causes.push((value, cause));
+        }
+    }
+
+    /// The `cause` context recorded for `value` at its fresh throw site, or
+    /// `Value::NULL` if it never superseded an error (a genuine rethrow).
+    fn recorded_throw_cause(&self, value: Value) -> Value {
+        self.thrown_value_causes
+            .iter()
+            .find(|(v, _)| *v == value)
+            .map_or(Value::NULL, |(_, cause)| *cause)
+    }
+
     fn maybe_queue_call_error_origin(
         &mut self,
         call_id: u64,
@@ -2268,6 +2314,7 @@ impl BexVm {
         self.frames.clear();
         self.pending_call_captures.clear();
         self.seen_throw_values.clear();
+        self.thrown_value_causes.clear();
         self.pending_sysop_call_id = None;
         self.pending_sysop_capture_mask = VmCaptureMask::disabled();
     }
@@ -3374,23 +3421,30 @@ impl BexVm {
         // error becomes the new error's `cause`.
         //
         // A *rethrow* re-raises an already-thrown value: a bare re-raise inside
-        // a handler, the no-match fall-through that re-enters this funnel, or
-        // `ThrowIfPanic` (all pass `is_rethrow`). None is a *new* failure
-        // "during handling of" the caught error, so none may graft another link
-        // onto the chain — in particular the no-match fall-through re-raises
-        // from inside the catch's own handler body, which the cause walk would
-        // otherwise mis-read as a self-link. So we skip the walk and give the
-        // re-raised value a fresh context with `cause = null`.
+        // a handler, the no-match fall-through that re-enters this funnel, a
+        // non-throwing `defer` pad's transparent re-raise of the in-flight
+        // error, or `ThrowIfPanic` (all pass `is_rethrow`). None is a *new*
+        // failure "during handling of" the caught error, so none may graft
+        // another link onto the chain by re-running the cause walk here — in
+        // particular the no-match fall-through and the defer pad both re-raise
+        // from inside a handler body, which the walk would mis-read as a
+        // self-link.
         //
-        // Caveat: a prior chain carried by the rethrown value is NOT preserved
-        // across the re-raise (its next catch sees `cause = null`). Recovering
-        // it would need a value->context association we don't track; reusing the
-        // walk result instead is unsafe — rethrowing an *outer* binding from an
-        // inner handler would bind the inner handler's mismatched error.
+        // Instead we reuse the cause the value's *original* (fresh) throw site
+        // computed, keyed by the value in `thrown_value_causes`. This preserves
+        // a pre-existing chain across a transparent re-raise (the defer-pad
+        // case: `mid` throws B while handling A, so B's recorded cause is
+        // ErrorContext(A); the pad's re-raise of B recovers A instead of
+        // nulling it) while keeping the deliberate null cause for a genuine
+        // rethrow that never superseded an error (no recorded entry -> NULL).
+        // The recorded value is the superseded error, never the re-raised
+        // value's own context, so this can never form a self-link.
         let cause_context = if is_rethrow {
-            Value::NULL
+            self.recorded_throw_cause(exception_value)
         } else {
-            self.find_cause_context()
+            let cause = self.find_cause_context();
+            self.record_throw_cause(exception_value, cause);
+            cause
         };
 
         // Frames popped by this unwind close with a status derived from the
@@ -8002,6 +8056,11 @@ impl ::bex_vm_types::RootHaver for BexVm {
                 .iter()
                 .filter_map(Value::as_object_ptr),
         );
+        // Both key (thrown value) and cause context are heap pointers.
+        for (value, cause) in &self.thrown_value_causes {
+            roots.extend(value.as_object_ptr());
+            roots.extend(cause.as_object_ptr());
+        }
 
         // Frame function pointers (needed once closures are heap-allocated)
         roots.extend(self.collect_frame_roots());
@@ -8044,6 +8103,18 @@ impl ::bex_vm_types::RootHaver for BexVm {
                 *value = Value::object(new_ptr);
             }
         }
+        for (value, cause) in &mut self.thrown_value_causes {
+            if let Some(ptr) = value.as_object_ptr()
+                && let Some(&new_ptr) = roots.get(&ptr)
+            {
+                *value = Value::object(new_ptr);
+            }
+            if let Some(ptr) = cause.as_object_ptr()
+                && let Some(&new_ptr) = roots.get(&ptr)
+            {
+                *cause = Value::object(new_ptr);
+            }
+        }
 
         // Frame function pointers (needed once closures are heap-allocated)
         for frame in &mut self.frames {