Allow async in client props

packages/reactDom/src/ReactServerDOM.ml

@@ -43,9 +43,9 @@ module Stream = struct
         Lwt.return ());
     index
 
-  let make ~initial_index =
+  let make ?(initial_index = 0) ?(pending = 0) () =
     let stream, push, close = Push_stream.make () in
-    (stream, { push; close; pending = 0; index = initial_index })
+    (stream, { push; close; pending; index = initial_index })
 end
 
 (* Resources module maintains insertion order while deduplicating based on src/href *)
@@ -388,7 +388,7 @@ module Model = struct
     List.map (fun (name, value) -> (name, client_value_to_json ~context ~to_chunk ~env value)) props
 
   let render ?(env = `Dev) ?(debug = false) ?subscribe element =
-    let stream, context = Stream.make ~initial_index:0 in
+    let stream, context = Stream.make () in
     let to_root_chunk element id =
       let payload = element_to_payload ~debug ~context ~to_chunk ~env element in
       to_chunk (Value payload) id
@@ -407,7 +407,7 @@ module Model = struct
         let digest = stack |> Hashtbl.hash |> Int.to_string in
         Lwt.return (React.Error { message; stack; env = "Server"; digest })
     in
-    let stream, context = Stream.make ~initial_index:0 in
+    let stream, context = Stream.make () in
     let to_root_chunk value id =
       let payload = client_value_to_json ~debug ~context ~to_chunk ~env value in
       to_chunk (Value payload) id
@@ -501,11 +501,9 @@ let rec client_to_html ~(fiber : Fiber.t) (element : React.element) =
         | output -> client_to_html ~fiber output
       in
       wait_for_suspense_to_resolve ()
-  | Async_component (_, _component) ->
-      (* async components can't be interleaved in client components, for now *)
-      raise
-        (Invalid_argument
-           "async components can't be part of a client component. This should never raise, the ppx should catch it")
+  | Async_component (_, component) ->
+      let%lwt element = component () in
+      client_to_html ~fiber element
   | Suspense { key = _; children; fallback } ->
       (* TODO: Do we need to care if there's Any_promise raising ? *)
       let%lwt fallback = client_to_html ~fiber fallback in
@@ -771,8 +769,13 @@ let render_html ?(skipRoot = false) ?(env = `Dev) ?debug:(_ = false) ?bootstrapS
         | None -> [])
   in
   (* Since we don't push the root_data_payload to the stream but return it immediately with the initial HTML, 
-     the stream's initial index starts at 1, with index 0 reserved for the root_data_payload. *)
-  let stream, context = Stream.make ~initial_index:1 in
+     the stream's initial index starts at 1, with index 0 reserved for the root_data_payload.
+
+     The root is also treated as a pending segment that must complete before the stream can be closed, 
+     as we don't push_async it to the stream, the pending counter starts at 1.
+     Similar on how react does: https://github.com/facebook/react/blob/7d9f876cbc7e9363092e60436704cf8ae435b969/packages/react-server/src/ReactFizzServer.js#L572-L581
+     *)
+  let stream, context = Stream.make ~initial_index:1 ~pending:1 () in
   let fiber : Fiber.t =
     {
       context;
@@ -789,8 +792,10 @@ let render_html ?(skipRoot = false) ?(env = `Dev) ?debug:(_ = false) ?bootstrapS
   let%lwt root_html, root_model = render_element_to_html ~fiber element in
   (* To return the model value immediately, we don't push it to the stream but return it as a payload script together with the user_scripts *)
   let root_data_payload = model_to_chunk (Value root_model) 0 in
+  (* Decrement the pending counter to signal that the root data payload is complete. *)
+  context.pending <- context.pending - 1;
   (* In case of not having any task pending, we can close the stream *)
-  (match context.pending = 0 with true -> context.close () | false -> ());
+  if context.pending = 0 then context.close ();
   let bootstrap_script_content =
     match bootstrapScriptContent with
     | None -> Html.null

packages/reactDom/test/test_RSC_html.ml

@@ -349,6 +349,36 @@ let client_with_element_props () =
        '>window.srr_stream.push()</script>";
     ]
 
+let client_component_with_async_component () =
+  let children =
+    React.Async_component
+      ( __FUNCTION__,
+        fun () ->
+          let%lwt () = sleep ~ms:10 in
+          Lwt.return (React.string "Async Component") )
+  in
+  let app ~children =
+    React.Upper_case_component
+      ( "app",
+        fun () ->
+          React.Client_component
+            {
+              import_module = "./client.js";
+              import_name = "Client";
+              props = [ ("children", React.Element children) ];
+              client = children;
+            } )
+  in
+  assert_html (app ~children)
+    ~shell:
+      "Async Component<script data-payload='0:[\"$\",\"$3\",null,{\"children\":\"$2\"},null,[],{}]\n\
+       '>window.srr_stream.push()</script>"
+    [
+      "<script data-payload='2:\"$L1\"\n'>window.srr_stream.push()</script>";
+      "<script data-payload='1:\"Async Component\"\n'>window.srr_stream.push()</script>";
+      "<script data-payload='3:I[\"./client.js\",[],\"Client\"]\n'>window.srr_stream.push()</script>";
+    ]
+
 let suspense_with_error () =
   let app () =
     React.Suspense.make ~fallback:(React.string "Loading...")
@@ -641,6 +671,7 @@ let tests =
     test "suspense_without_promise" suspense_without_promise;
     test "with_sleepy_promise" with_sleepy_promise;
     test "client_with_promise_props" client_with_promise_props;
+    test "client_component_with_async_component" client_component_with_async_component;
     test "async_component_with_promise" async_component_with_promise;
     test "suspense_with_error" suspense_with_error;
     test "suspense_with_error_in_async" suspense_with_error_in_async;

packages/reactDom/test/test_RSC_model.ml

@@ -813,6 +813,37 @@ let client_component_with_resources_metadata () =
   Alcotest.(check bool) "should have head with resources" true has_head_with_resources;
   Lwt.return ()
 
+let client_component_with_async_component () =
+  let async_component =
+    React.Async_component
+      ( __FUNCTION__,
+        fun () ->
+          let%lwt () = sleep ~ms:10 in
+          Lwt.return (React.string "Async Component") )
+  in
+  let app ~children =
+    React.Upper_case_component
+      ( "app",
+        fun () ->
+          React.Client_component
+            {
+              import_module = "./client.js";
+              import_name = "Client";
+              props = [ ("children", React.Element children) ];
+              client = children;
+            } )
+  in
+  let output, subscribe = capture_stream () in
+  let%lwt () = ReactServerDOM.render_model ~subscribe (app ~children:async_component) in
+  assert_list_of_strings !output
+    [
+      "1:I[\"./client.js\",[],\"Client\"]\n";
+      "3:\"$L2\"\n";
+      "0:[\"$\",\"$1\",null,{\"children\":\"$3\"},null,[],{}]\n";
+      "2:\"Async Component\"\n";
+    ];
+  Lwt.return ()
+
 let page_with_hoisted_resources () =
   (* Test that resources like scripts and styles are properly hoisted *)
   let app () =
@@ -953,6 +984,7 @@ let tests =
     test "client_without_props" client_without_props;
     test "client_with_element_props" client_with_element_props;
     test "client_with_server_children" client_with_server_children;
+    test "client_component_with_async_component" client_component_with_async_component;
     test "act_with_simple_response" act_with_simple_response;
     test "env_development_adds_debug_info" env_development_adds_debug_info;
     test "act_with_error" act_with_error;