diff --git a/AGENTS.md b/AGENTS.md
index 725007dd7b..681b0feb1d 100644
--- a/AGENTS.md
+++ b/AGENTS.md
@@ -12,6 +12,32 @@ make && make test
 make format && make checkformat
 ```
 
+## Container Tooling Setup
+
+The base container image does not ship with the OCaml toolchain or `dune`.
+Install them once per container session before running any compiler builds or
+tests:
+
+```bash
+# Install the OCaml 4.14 toolchain and findlib support
+apt-get update
+apt-get install -y ocaml-nox ocaml-native-compilers ocaml-findlib m4 pkg-config
+
+# Build and install dune ≥3.20 from source
+curl -LO https://github.com/ocaml/dune/releases/download/3.20.2/dune-3.20.2.tbz
+tar -xf dune-3.20.2.tbz
+cd dune-3.20.2 && make release
+install -m 0755 _build/default/src/dune.exe /usr/local/bin/dune
+cd .. && rm -rf dune-3.20.2 dune-3.20.2.tbz
+
+# Verify dune is available
+dune --version
+```
+
+These steps match the manual bootstrapping previously performed in the
+container, ensuring subsequent `make`/`dune` commands succeed without relying
+on external package mirrors such as `opam.ocaml.org`.
+
 ## ⚠️ Critical Guidelines & Common Pitfalls
 
 - **We are NOT bound by OCaml compatibility** - The ReScript compiler originated as a fork of the OCaml compiler, but we maintain our own AST and can make breaking changes. Focus on what's best for ReScript's JavaScript compilation target.
diff --git a/tests/ounit_tests/dune b/tests/ounit_tests/dune
index 56ac24bd1f..dc9e6a7c48 100644
--- a/tests/ounit_tests/dune
+++ b/tests/ounit_tests/dune
@@ -11,4 +11,4 @@
   (<> %{profile} browser))
  (flags
   (:standard -w +a-4-9-30-40-41-42-48-70))
- (libraries bsb bsb_helper core ounit2 analysis))
+ (libraries bsb bsb_helper core frontend ounit2 analysis))
diff --git a/tests/ounit_tests/ounit_ast_return_tests.ml b/tests/ounit_tests/ounit_ast_return_tests.ml
index b0345acb4b..5f8e2b89f2 100644
--- a/tests/ounit_tests/ounit_ast_return_tests.ml
+++ b/tests/ounit_tests/ounit_ast_return_tests.ml
@@ -9,8 +9,7 @@ let parse_structure source =
   in
   result.parsetree
 
-let test_simple_return _ctx =
-  let structure = parse_structure "let f = x => x + 1\n" in
+let extract_fun_body structure =
   match structure with
   | [
    {
@@ -21,74 +20,67 @@ let test_simple_return _ctx =
          );
      _;
    };
-  ] -> (
-    OUnit.assert_bool "function body should be marked as a return"
-      (Res_return_marker.expression_is_return body);
-    match body.pexp_desc with
-    | Pexp_apply {funct; args; _} ->
-      OUnit.assert_bool "callee should not be marked"
-        (not (Res_return_marker.expression_is_return funct));
-      List.iter
-        (fun (_, arg) ->
-          OUnit.assert_bool "arguments should not be marked"
-            (not (Res_return_marker.expression_is_return arg)))
-        args
-    | _ -> OUnit.assert_failure "expected application in function body")
+  ] -> body
   | _ -> OUnit.assert_failure "unexpected structure for simple function"
 
+let test_simple_return _ctx =
+  let structure = parse_structure "let f = x => x + 1\n" in
+  let body = extract_fun_body structure in
+  OUnit.assert_bool "function body should be marked as a return"
+    (Res_return_marker.expression_is_return body);
+  match body.pexp_desc with
+  | Pexp_apply {funct; args; _} ->
+    OUnit.assert_bool "callee should not be marked"
+      (not (Res_return_marker.expression_is_return funct));
+    List.iter
+      (fun (_, arg) ->
+        OUnit.assert_bool "arguments should not be marked"
+          (not (Res_return_marker.expression_is_return arg)))
+      args
+  | _ -> OUnit.assert_failure "expected application in function body"
+
 let test_sequence_return _ctx =
   let structure = parse_structure "let f = x => {let y = x; y}\n" in
-  match structure with
-  | [
-   {
-     Parsetree.pstr_desc =
-       Pstr_value
-         ( _,
-           [{pvb_expr = {Parsetree.pexp_desc = Pexp_fun {rhs = body; _}; _}; _}]
-         );
-     _;
-   };
-  ] -> (
-    OUnit.assert_bool "function body should be marked"
-      (Res_return_marker.expression_is_return body);
-    match body.pexp_desc with
-    | Pexp_sequence (first, second) ->
-      OUnit.assert_bool "last expression in sequence should be marked"
-        (Res_return_marker.expression_is_return second);
-      OUnit.assert_bool "first expression in sequence should not be marked"
-        (not (Res_return_marker.expression_is_return first))
-    | Pexp_let (_, _, inner) ->
-      OUnit.assert_bool "inner expression should be marked"
-        (Res_return_marker.expression_is_return inner)
-    | _ -> OUnit.assert_failure "expected sequence or let in body")
-  | _ -> OUnit.assert_failure "unexpected structure for block body"
+  let body = extract_fun_body structure in
+  OUnit.assert_bool "function body should be marked"
+    (Res_return_marker.expression_is_return body);
+  match body.pexp_desc with
+  | Pexp_sequence (first, second) ->
+    OUnit.assert_bool "last expression in sequence should be marked"
+      (Res_return_marker.expression_is_return second);
+    OUnit.assert_bool "first expression in sequence should not be marked"
+      (not (Res_return_marker.expression_is_return first))
+  | Pexp_let (_, _, inner) ->
+    OUnit.assert_bool "inner expression should be marked"
+      (Res_return_marker.expression_is_return inner)
+  | _ -> OUnit.assert_failure "expected sequence or let in body"
 
 let test_switch_returns _ctx =
   let structure =
     parse_structure "let f = x => switch x { | 0 => 1 | _ => x }\n"
   in
-  match structure with
-  | [
-   {
-     Parsetree.pstr_desc =
-       Pstr_value
-         ( _,
-           [{pvb_expr = {Parsetree.pexp_desc = Pexp_fun {rhs = body; _}; _}; _}]
-         );
-     _;
-   };
-  ] -> (
-    OUnit.assert_bool "switch expression should be marked"
-      (Res_return_marker.expression_is_return body);
-    match body.pexp_desc with
-    | Pexp_match (_, cases) ->
-      List.iter
-        (fun {pc_rhs; _} ->
-          OUnit.assert_bool "case rhs should be marked"
-            (Res_return_marker.expression_is_return pc_rhs))
-        cases
-    | _ -> OUnit.assert_failure "expected switch in function body")
-  | _ -> OUnit.assert_failure "unexpected structure for switch body"
+  let body = extract_fun_body structure in
+  OUnit.assert_bool "switch expression should be marked"
+    (Res_return_marker.expression_is_return body);
+  match body.pexp_desc with
+  | Pexp_match (_, cases) ->
+    List.iter
+      (fun {pc_rhs; _} ->
+        OUnit.assert_bool "case rhs should be marked"
+          (Res_return_marker.expression_is_return pc_rhs))
+      cases
+  | _ -> OUnit.assert_failure "expected switch in function body"
+
+let test_rewrite_loses_return_flag _ctx =
+  let source = "let f = x => g(x)\n" in
+  let structure = parse_structure source in
+  let body = extract_fun_body structure in
+  OUnit.assert_bool "pre-rewrite body should be marked"
+    (Res_return_marker.expression_is_return body);
+  let rewritten = Ppx_entry.rewrite_implementation structure in
+  let rewritten_body = extract_fun_body rewritten in
+  OUnit.assert_bool "rewritten body currently loses its return flag"
+    (not (Res_return_marker.expression_is_return rewritten_body))
 
 let suites =
   "return marker"
@@ -96,4 +88,5 @@ let suites =
          "function body" >:: test_simple_return;
          "sequence tail" >:: test_sequence_return;
          "switch cases" >:: test_switch_returns;
+         "rewrite loses return flag" >:: test_rewrite_loses_return_flag;
        ]