feat: allow async in client props (#315)

Author: pedrobslisboa

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;