fix(rsc): group client chunks by server-entry ownership

Author: SyMind

crates/rspack_binding_api/src/plugins/rsc.rs

@@ -77,7 +77,8 @@ impl From<&JsRscClientPluginOptions<'_>> for RscClientPluginOptions {
 #[napi(object, object_to_js = false)]
 pub struct JsRscServerPluginOptions<'a> {
   pub coordinator: ClassInstance<'a, JsCoordinator>,
-  pub on_server_component_changes: Option<Either3<Function<'static, (), ()>, Undefined, Null>>,
+  pub on_server_component_changes:
+    Option<Either3<Function<'static, (), Promise<()>>, Undefined, Null>>,
   pub on_manifest: Option<Either3<Function<'static, String, Promise<()>>, Undefined, Null>>,
 }
 
@@ -100,7 +101,14 @@ impl TryFrom<&JsRscServerPluginOptions<'_>> for RscServerPluginOptions {
         Some(Box::new(
           move || -> BoxFuture<'static, rspack_error::Result<()>> {
             let ts_fn = ts_fn.clone();
-            Box::pin(async move { ts_fn.call_async(()).await.to_rspack_result() })
+            Box::pin(async move {
+              ts_fn
+                .call_async(())
+                .await
+                .to_rspack_result()?
+                .await
+                .to_rspack_result()
+            })
           },
         ))
       }

crates/rspack_plugin_rsc/src/client_plugin.rs

@@ -570,7 +570,9 @@ async fn make(&self, compilation: &mut Compilation) -> Result<()> {
     if !client_modules.is_empty() || entry_state.has_css_imports_by_server_entry() {
       let dependency = Box::new(RscEntryDependency::new(
         entry_name.clone(),
-        client_modules.clone(),
+        entry_state.isolated_client_entries.clone(),
+        entry_state.root_client_entries.clone(),
+        entry_state.client_entries_by_server_entry.clone(),
         entry_state.css_imports_by_server_entry(),
         false,
       ));

crates/rspack_plugin_rsc/src/component_info.rs

@@ -19,11 +19,17 @@ use crate::{
 
 // { [request to inject into client compilation]: [exported names] }
 pub type ClientComponentImports = FxIndexMap<String, FxIndexSet<Atom>>;
+pub type ClientComponentImportsByServerEntry = FxIndexMap<String, ClientComponentImports>;
 
 #[derive(Debug, Default)]
 pub struct ComponentInfo {
   pub should_inject_ssr_modules: bool,
+  /// All client component imports reached from this RSC entry.
   pub client_component_imports: ClientComponentImports,
+  /// Client component imports reached without a `use server-entry` owner.
+  pub root_client_component_imports: ClientComponentImports,
+  /// Client component imports reached under each `use server-entry` owner.
+  pub client_component_imports_by_server_entry: ClientComponentImportsByServerEntry,
   pub css_imports_by_server_entry: CssImportsByServerEntry,
   pub root_css_imports: RootCssImports,
   pub action_imports: Vec<(String, Vec<ActionIdNamePair>)>,
@@ -155,12 +161,12 @@ fn filter_client_components(
       .client_component_imports
       .contains_key(resource.as_ref())
     {
-      add_client_import(
+      add_client_import_for_owners(
         module,
-        &resource,
+        resource.as_ref(),
         imported_identifiers,
-        false,
-        &mut component_info.client_component_imports,
+        server_entries,
+        component_info,
       );
     }
     return;
@@ -188,20 +194,12 @@ fn filter_client_components(
 
   let module_graph = compilation.get_module_graph();
   if is_client_component_entry_module(module) {
-    if !component_info
-      .client_component_imports
-      .contains_key(resource.as_ref())
-    {
-      component_info
-        .client_component_imports
-        .insert(resource.to_string(), Default::default());
-    }
-    add_client_import(
+    add_client_import_for_owners(
       module,
       resource.as_ref(),
       imported_identifiers,
-      true,
-      &mut component_info.client_component_imports,
+      server_entries,
+      component_info,
     );
     return;
   }
@@ -251,6 +249,63 @@ fn filter_client_components(
   }
 }
 
+fn add_client_import_for_owners(
+  module: &dyn Module,
+  resource: &str,
+  imported_identifiers: &[Atom],
+  server_entries: &[String],
+  component_info: &mut ComponentInfo,
+) {
+  add_client_import_to_map(
+    module,
+    resource,
+    imported_identifiers,
+    &mut component_info.client_component_imports,
+  );
+
+  if server_entries.is_empty() {
+    add_client_import_to_map(
+      module,
+      resource,
+      imported_identifiers,
+      &mut component_info.root_client_component_imports,
+    );
+    return;
+  }
+
+  for server_entry in server_entries {
+    let client_component_imports = component_info
+      .client_component_imports_by_server_entry
+      .entry(server_entry.clone())
+      .or_default();
+    add_client_import_to_map(
+      module,
+      resource,
+      imported_identifiers,
+      client_component_imports,
+    );
+  }
+}
+
+fn add_client_import_to_map(
+  module: &dyn Module,
+  mod_request: &str,
+  imported_identifiers: &[Atom],
+  client_component_imports: &mut ClientComponentImports,
+) {
+  let is_first_visit_module = !client_component_imports.contains_key(mod_request);
+  if is_first_visit_module {
+    client_component_imports.insert(mod_request.to_string(), Default::default());
+  }
+  add_client_import(
+    module,
+    mod_request,
+    imported_identifiers,
+    is_first_visit_module,
+    client_component_imports,
+  );
+}
+
 fn add_client_import(
   module: &dyn Module,
   mod_request: &str,

crates/rspack_plugin_rsc/src/plugin_state.rs

@@ -16,6 +16,7 @@ use crate::reference_manifest::{
 
 pub type ActionIdNamePair = (Atom, Atom);
 pub type CssImportsByServerEntry = FxHashMap<String, FxIndexSet<String>>;
+pub type ClientModulesByServerEntry = FxHashMap<String, Vec<ClientModuleImport>>;
 pub type RootCssImports = FxIndexSet<String>;
 
 /// Structured info about a client module to inject into the client compiler.
@@ -39,7 +40,14 @@ pub struct ServerEntryState {
 #[derive(Debug, Default)]
 pub struct EntryState {
   pub server_entries: FxHashMap<String, ServerEntryState>,
+  /// All client modules discovered from the RSC graph for this entry.
   pub injected_client_entries: Vec<ClientModuleImport>,
+  /// Client modules used by multiple owners and kept as individual async chunks.
+  pub isolated_client_entries: Vec<ClientModuleImport>,
+  /// Client modules used only by the root RSC tree.
+  pub root_client_entries: Vec<ClientModuleImport>,
+  /// Client modules used only by one `use server-entry` subtree.
+  pub client_entries_by_server_entry: ClientModulesByServerEntry,
   pub client_modules: FxHashMap<String, ManifestExport>,
   /// Root CSS import paths reached through a parent chain without `use server-entry`.
   /// These are attached directly to the matching client compiler entry.

crates/rspack_plugin_rsc/src/rsc_entry_dependency.rs

@@ -9,14 +9,21 @@ use rspack_core::{
   DependencyType, FactorizeInfo, ModuleDependency, ResourceIdentifier,
 };
 
-use crate::plugin_state::{ClientModuleImport, CssImportsByServerEntry};
+use crate::plugin_state::{
+  ClientModuleImport, ClientModulesByServerEntry, CssImportsByServerEntry,
+};
 
 #[cacheable]
 #[derive(Debug, Clone)]
 pub struct RscEntryDependency {
   id: DependencyId,
   pub name: Arc<str>,
+  /// Client modules that should keep the existing one-module-per-async-block behavior.
   pub client_modules: Vec<ClientModuleImport>,
+  /// Client modules owned by the root RSC entry.
+  pub root_client_modules: Vec<ClientModuleImport>,
+  #[cacheable(with=AsMap<AsCacheable, AsVec>)]
+  pub client_modules_by_server_entry: ClientModulesByServerEntry,
   #[cacheable(with=AsMap<AsCacheable, AsVec>)]
   pub css_imports_by_server_entry: CssImportsByServerEntry,
   /// When true, client modules are loaded eagerly (not as code-split points).
@@ -30,6 +37,8 @@ impl RscEntryDependency {
   pub fn new(
     name: Arc<str>,
     client_modules: Vec<ClientModuleImport>,
+    root_client_modules: Vec<ClientModuleImport>,
+    client_modules_by_server_entry: ClientModulesByServerEntry,
     css_imports_by_server_entry: CssImportsByServerEntry,
     is_server_side_rendering: bool,
   ) -> Self {
@@ -38,6 +47,8 @@ impl RscEntryDependency {
       id: DependencyId::new(),
       name,
       client_modules,
+      root_client_modules,
+      client_modules_by_server_entry,
       css_imports_by_server_entry,
       is_server_side_rendering,
       resource_identifier,