collect actions for the right runtime only

Author: mischnic

crates/next-api/src/app.rs

@@ -55,7 +55,7 @@ use turbopack_core::{
     },
     file_source::FileSource,
     ident::{AssetIdent, Layer},
-    module::Module,
+    module::{Module, Modules},
     module_graph::{
         GraphEntries, ModuleGraph, SingleModuleGraph, VisitedModules,
         binding_usage_info::compute_binding_usage_info,
@@ -973,7 +973,7 @@ impl AppProject {
     pub async fn app_module_graph(
         &self,
         rsc_entry: ResolvedVc<Box<dyn Module>>,
-        server_action_loader_module: ResolvedVc<Box<dyn Module>>,
+        server_action_loader_modules: ResolvedVc<Modules>,
         client_shared_entries: Vc<EvaluatableAssets>,
         has_layout_segments: bool,
     ) -> Result<Vc<ModuleGraph>> {
@@ -991,8 +991,14 @@ impl AppProject {
             // Implements layout segment optimization to compute a graph "chain" for each layout
             // segment
             async move {
-                let rsc_entry_chunk_group =
-                    ChunkGroupEntry::Entry(vec![server_action_loader_module, rsc_entry]);
+                let rsc_entry_chunk_group = ChunkGroupEntry::Entry(
+                    server_action_loader_modules
+                        .await?
+                        .iter()
+                        .copied()
+                        .chain(std::iter::once(rsc_entry))
+                        .collect(),
+                );
 
                 let mut graphs = vec![];
                 let visited_modules = if has_layout_segments {
@@ -1279,15 +1285,36 @@ impl AppEndpoint {
     }
 
     #[turbo_tasks::function]
-    async fn server_action_loader_module(self: Vc<Self>) -> Result<Vc<Box<dyn Module>>> {
+    async fn server_action_loader_modules(self: Vc<Self>) -> Result<Vc<Modules>> {
         let this = self.await?;
-        // let app_entry = self.app_endpoint_entry().await?;
-        // let runtime = app_entry.config.await?.runtime.unwrap_or_default();
+        let app_entry = self.app_endpoint_entry().await?;
+        let runtime = app_entry.config.await?.runtime.unwrap_or_default();
 
-        Ok(Vc::upcast(ServerActionCollectModule::new(
-            rcstr!("next/server-actions"),
-            this.page.to_string().into(),
-        )))
+        Ok(Vc::cell(vec![
+            // Collect inline "use server" function inside of RSC modules
+            ResolvedVc::upcast(
+                ServerActionCollectModule::new(
+                    rcstr!("next/server-actions"),
+                    this.page.to_string().into(),
+                )
+                .to_resolved()
+                .await?,
+            ),
+            // Collect "use server" modules imported from client components.
+            // By only collecting the correct runtime, modules emitted for the other runtime are
+            // never even compiled.
+            ResolvedVc::upcast(
+                ServerActionCollectModule::new(
+                    match runtime {
+                        NextRuntime::NodeJs => rcstr!("next/server-actions/node"),
+                        NextRuntime::Edge => rcstr!("next/server-actions/edge"),
+                    },
+                    this.page.to_string().into(),
+                )
+                .to_resolved()
+                .await?,
+            ),
+        ]))
     }
 
     #[turbo_tasks::function]
@@ -1342,11 +1369,11 @@ impl AppEndpoint {
 
         let is_app_page = matches!(this.ty, AppEndpointType::Page { .. });
 
-        let server_action_loader = self.server_action_loader_module();
+        let server_action_loader_modules = self.server_action_loader_modules();
 
         let module_graph = this.app_project.app_module_graph(
             *rsc_entry,
-            server_action_loader,
+            server_action_loader_modules,
             // We only need the client runtime entries for pages not for Route Handlers
             if is_app_page {
                 this.app_project.client_runtime_entries()
@@ -1547,7 +1574,7 @@ impl AppEndpoint {
         if emit_rsc_manifests {
             server_assets.insert(
                 create_server_actions_manifest(
-                    self.server_action_loader_module(),
+                    self.server_action_loader_modules(),
                     node_root.clone(),
                     app_entry.original_name.clone(),
                     runtime,
@@ -1864,24 +1891,29 @@ impl AppEndpoint {
         let project = this.app_project.project();
         let app_entry = self.app_endpoint_entry().await?;
         let runtime = app_entry.config.await?.runtime.unwrap_or_default();
-        let server_action_loader = self.server_action_loader_module().to_resolved().await?;
+        let server_action_loaders = self.server_action_loader_modules().await?;
 
         let chunking_context = project.runtime_chunking_context(process_client_assets, runtime);
 
+        let entry_chunk_group = ChunkGroup::Entry(
+            server_action_loaders
+                .iter()
+                .copied()
+                .chain(std::iter::once(app_entry.rsc_entry))
+                .collect(),
+        );
+
         Ok(match runtime {
             NextRuntime::Edge => chunking_context.evaluated_chunk_group_assets(
                 app_entry.rsc_entry.ident(),
-                ChunkGroup::Entry(vec![server_action_loader, app_entry.rsc_entry]),
+                entry_chunk_group,
                 module_graph,
                 AvailabilityInfo::root(),
             ),
             NextRuntime::NodeJs => {
                 async {
                     let mut current_chunk_group = ChunkGroupResult::empty_resolved();
 
-                    let entry_chunk_group =
-                        ChunkGroup::Entry(vec![server_action_loader, app_entry.rsc_entry]);
-
                     let chunk_group_info = module_graph.chunk_group_info();
 
                     let client_references = client_references.await?;
@@ -2138,10 +2170,14 @@ impl Endpoint for AppEndpoint {
     async fn entries(self: Vc<Self>) -> Result<Vc<GraphEntries>> {
         let this = self.await?;
         Ok(Vc::cell(vec![
-            ChunkGroupEntry::Entry(vec![
-                self.server_action_loader_module().to_resolved().await?,
-                self.app_endpoint_entry().await?.rsc_entry,
-            ]),
+            ChunkGroupEntry::Entry(
+                self.server_action_loader_modules()
+                    .await?
+                    .iter()
+                    .copied()
+                    .chain(std::iter::once(self.app_endpoint_entry().await?.rsc_entry))
+                    .collect(),
+            ),
             ChunkGroupEntry::Entry(
                 this.app_project
                     .client_runtime_entries()
@@ -2158,12 +2194,12 @@ impl Endpoint for AppEndpoint {
     async fn module_graphs(self: Vc<Self>) -> Result<Vc<ModuleGraphs>> {
         let this = self.await?;
         let app_entry = self.app_endpoint_entry().await?;
-        let server_action_loader = self.server_action_loader_module();
+        let server_action_loader_modules = self.server_action_loader_modules();
         let module_graph = this
             .app_project
             .app_module_graph(
                 *app_entry.rsc_entry,
-                server_action_loader,
+                server_action_loader_modules,
                 this.app_project.client_runtime_entries(),
                 matches!(this.ty, AppEndpointType::Page { .. }),
             )

crates/next-api/src/server_actions.rs

@@ -1,4 +1,4 @@
-use anyhow::{Context, Result};
+use anyhow::{Context, Result, bail};
 use bincode::{Decode, Encode};
 use next_core::{
     next_manifests::{ActionLayer, ActionManifestWorkerEntry, ServerReferenceManifest},
@@ -10,7 +10,10 @@ use turbo_tasks_fs::{self, File, FileContent, FileSystem, FileSystemPath, Virtua
 use turbopack_core::{
     self,
     asset::AssetContent,
-    chunk::{AsyncModuleInfo, ChunkItem, ChunkableModule, ChunkingContext, EvaluatableAsset},
+    chunk::{
+        AsyncModuleInfo, ChunkItem, ChunkableModule, ChunkingContext, ChunkingType,
+        EvaluatableAsset,
+    },
     emit_collect::{CollectingModule, EmittedModuleReference},
     ident::AssetIdent,
     module::{Module, ModuleSideEffects, Modules},
@@ -31,14 +34,14 @@ use turbopack_ecmascript::chunk::{
 /// loader.
 #[turbo_tasks::function]
 pub(crate) async fn create_server_actions_manifest(
-    server_action_loader: Vc<Box<dyn Module>>,
+    server_action_loader_modules: Vc<Modules>,
     node_root: FileSystemPath,
     page_name: RcStr,
     runtime: NextRuntime,
     module_graph: Vc<ModuleGraph>,
     chunking_context: Vc<Box<dyn ChunkingContext>>,
 ) -> Vc<Box<dyn OutputAsset>> {
-    let actions = collect_actions(server_action_loader, module_graph);
+    let actions = collect_actions(server_action_loader_modules, module_graph);
     build_manifest(
         node_root,
         page_name,
@@ -51,29 +54,36 @@ pub(crate) async fn create_server_actions_manifest(
 
 #[turbo_tasks::function]
 async fn collect_actions(
-    server_action_loader: ResolvedVc<Box<dyn Module>>,
+    server_action_loader_modules: Vc<Modules>,
     module_graph: Vc<ModuleGraph>,
 ) -> Result<Vc<AllActions>> {
     // This mirrors what the ServerActionCollectModule ends up chunking into the chunk.
     let collected_modules = module_graph.collected_modules().await?;
-    // This can be none if there are no server actions
-    let actions =
-        collected_modules
-            .collected_references
-            .iter()
-            .find_map(|((entry, loader), actions)| {
-                if *entry == server_action_loader && *loader == server_action_loader {
-                    Some(actions)
-                } else {
-                    None
-                }
-            });
+    let server_action_loader_modules = server_action_loader_modules.await?;
+    assert_eq!(server_action_loader_modules.len(), 2);
+
+    let actions = collected_modules
+        .collected_references
+        .iter()
+        .filter(|((_entry_modules, loader), _)| {
+            // No need to check entry_modules. Each page (ChunkGroup::Entry) has its own loader
+            // module anyway.
+
+            // server_action_loader_modules contains only 2 modules, so converting that into a
+            // hashset for quicker lookup is not necessary.
+            server_action_loader_modules.contains(loader)
+        })
+        // Early stop the search. There can only ever be two matches
+        .take(2)
+        .flat_map(|(_, actions)| actions.iter());
 
     Ok(Vc::cell(
         actions
-            .into_iter()
-            .flatten()
             .map(async |(data, module, _)| {
+                let namespace = match &data.chunking_type {
+                    ChunkingType::Collected { merge_tag, .. } => merge_tag,
+                    _ => bail!("unexpected chunking type for collected reference"),
+                };
                 let data =
                     ResolvedVc::try_sidecast::<Box<dyn EmittedModuleReference>>(data.reference)
                         .context(
@@ -92,7 +102,13 @@ async fn collect_actions(
                 Ok((
                     hash.to_string(),
                     (
-                        ActionLayer::ActionBrowser, // TODO
+                        match namespace.as_str() {
+                            "next/server-actions" => ActionLayer::Rsc,
+                            "next/server-actions/edge" | "next/server-actions/node" => {
+                                ActionLayer::ActionBrowser
+                            }
+                            _ => bail!("unexpected namespace {namespace} for collected reference"),
+                        },
                         ActionMeta {
                             name: name.to_string(),
                             source_path: "".to_string(), // TODO

crates/next-custom-transforms/src/transforms/server_actions.rs

@@ -3276,13 +3276,13 @@ fn emit_server_action(
         } else {
             vec![
                 quote!(
-                    "$emit($req, {namespace:'next/server-actions', data: $data, with: { 'turbopack-transition': 'next-rsc' }});" as Stmt,
+                    "$emit($req, {namespace:'next/server-actions/node', data: $data, with: { 'turbopack-transition': 'next-rsc' }});" as Stmt,
                     emit = emit.clone(),
                     req: Expr = req.clone(),
                     data: Expr = data.clone()
                 ),
                 quote!(
-                    "$emit($req, {namespace:'next/server-actions', data: $data, with: { 'turbopack-transition': 'next-edge-rsc' }});" as Stmt,
+                    "$emit($req, {namespace:'next/server-actions/edge', data: $data, with: { 'turbopack-transition': 'next-edge-rsc' }});" as Stmt,
                     emit = emit,
                     req: Expr = req,
                     data: Expr = data,

crates/next-custom-transforms/tests/fixture/server-actions/client-graph/turbopack/1/output.js

@@ -2,14 +2,14 @@ import { createServerReference, callServer, findSourceMapURL } from "private-nex
 const $$RSC_SERVER_ACTION_0 = /*#__PURE__*/ createServerReference("00ab21efdafbe611287bc25c0462b1e0510d13e48b", callServer, void 0, findSourceMapURL, "foo");
 export { $$RSC_SERVER_ACTION_0 as foo };
 __turbopack_emit__("./item.js", {
-    namespace: 'next/server-actions',
+    namespace: 'next/server-actions/node',
     data: "00ab21efdafbe611287bc25c0462b1e0510d13e48b|foo",
     with: {
         'turbopack-transition': 'next-rsc'
     }
 });
 __turbopack_emit__("./item.js", {
-    namespace: 'next/server-actions',
+    namespace: 'next/server-actions/edge',
     data: "00ab21efdafbe611287bc25c0462b1e0510d13e48b|foo",
     with: {
         'turbopack-transition': 'next-edge-rsc'

crates/next-custom-transforms/tests/fixture/server-actions/client-graph/turbopack/2/output.js

@@ -1,14 +1,14 @@
 import { createServerReference, callServer, findSourceMapURL } from "private-next-rsc-action-client-wrapper";
 export default /*#__PURE__*/ createServerReference("00c18c215a6b7cdc64bf709f3a714ffdef1bf9651d", callServer, void 0, findSourceMapURL, "default");
 __turbopack_emit__("./item.js", {
-    namespace: 'next/server-actions',
+    namespace: 'next/server-actions/node',
     data: "00c18c215a6b7cdc64bf709f3a714ffdef1bf9651d|default",
     with: {
         'turbopack-transition': 'next-rsc'
     }
 });
 __turbopack_emit__("./item.js", {
-    namespace: 'next/server-actions',
+    namespace: 'next/server-actions/edge',
     data: "00c18c215a6b7cdc64bf709f3a714ffdef1bf9651d|default",
     with: {
         'turbopack-transition': 'next-edge-rsc'