fix: next/root-params erroring when rerendering after action (vercel#82326)

When implementing the check that disallows using `next/root-params` in
server actions, i forgot that when a page is rerendered after an action,
the render is still wrapped with an `actionAsyncStorage` with
`actionStore.isAction === true`. The fix is to also check that we're in
the `'action'` phase, i.e. we haven't switched to rendering yet.

Closes vercel#82302

Author: lubieowoce

packages/next/errors.json

@@ -771,5 +771,6 @@
   "770": "createParamsFromClient should not be called in a runtime prerender.",
   "771": "\\`%s\\` was called during a runtime prerender. Next.js should be preventing %s from being included in server components statically, but did not in this case.",
   "772": "FetchStrategy.PPRRuntime should never be used when `experimental.clientSegmentCache` is disabled",
-  "773": "Missing workStore in createPrerenderParamsForClientSegment"
+  "773": "Missing workStore in createPrerenderParamsForClientSegment",
+  "774": "Route %s used %s outside of a Server Component. This is not allowed."
 }

packages/next/src/server/request/root-params.ts

@@ -203,6 +203,13 @@ export function getRootParam(paramName: string): Promise<ParamValue> {
     throw new InvariantError(`Missing workStore in ${apiName}`)
   }
 
+  const workUnitStore = workUnitAsyncStorage.getStore()
+  if (!workUnitStore) {
+    throw new Error(
+      `Route ${workStore.route} used ${apiName} outside of a Server Component. This is not allowed.`
+    )
+  }
+
   const actionStore = actionAsyncStorage.getStore()
   if (actionStore) {
     if (actionStore.isAppRoute) {
@@ -211,23 +218,17 @@ export function getRootParam(paramName: string): Promise<ParamValue> {
         `Route ${workStore.route} used ${apiName} inside a Route Handler. Support for this API in Route Handlers is planned for a future version of Next.js.`
       )
     }
-    if (actionStore.isAction) {
+    if (actionStore.isAction && workUnitStore.phase === 'action') {
       // Actions are not fundamentally tied to a route (even if they're always submitted from some page),
       // so root params would be inconsistent if an action is called from multiple roots.
+      // Make sure we check if the phase is "action" - we should not error in the rerender
+      // after an action revalidates or updates cookies (which will still have `actionStore.isAction === true`)
       throw new Error(
         `${apiName} was used inside a Server Action. This is not supported. Functions from 'next/root-params' can only be called in the context of a route.`
       )
     }
   }
 
-  const workUnitStore = workUnitAsyncStorage.getStore()
-
-  if (!workUnitStore) {
-    throw new Error(
-      `Route ${workStore.route} used ${apiName} in Pages Router. This API is only available within App Router.`
-    )
-  }
-
   switch (workUnitStore.type) {
     case 'unstable-cache':
     case 'cache': {

test/e2e/app-dir/app-root-params-getters/fixtures/simple/app/[lang]/[locale]/rerender-after-server-action/page.tsx

@@ -0,0 +1,37 @@
+import { lang, locale } from 'next/root-params'
+import { connection } from 'next/server'
+import { cookies } from 'next/headers'
+import { Suspense } from 'react'
+
+export default async function Page() {
+  const currentLang = await lang()
+  const currentLocale = await locale()
+  return (
+    <main>
+      <div>
+        Root params are{' '}
+        <span id="root-params">
+          {currentLang} {currentLocale}
+        </span>
+      </div>
+      <Suspense fallback="Loading...">
+        <Timestamp />
+      </Suspense>
+      <form
+        action={async () => {
+          'use server'
+          // rerender the page and return it alongside the action result
+          const cookieStore = await cookies()
+          cookieStore.set('my-cookie', Date.now() + '')
+        }}
+      >
+        <button type="submit">Submit form</button>
+      </form>
+    </main>
+  )
+}
+
+async function Timestamp() {
+  await connection()
+  return <div id="timestamp">{Date.now()}</div>
+}

test/e2e/app-dir/app-root-params-getters/simple.test.ts

@@ -6,6 +6,23 @@ import { outdent } from 'outdent'
 import { createRequestTracker } from '../../../lib/e2e-utils/request-tracker'
 
 describe('app-root-param-getters - simple', () => {
+  let currentCliOutputIndex = 0
+  beforeEach(() => {
+    resetCliOutput()
+  })
+
+  const getCliOutput = () => {
+    if (next.cliOutput.length < currentCliOutputIndex) {
+      // cliOutput shrank since we started the test, so something (like a `sandbox`) reset the logs
+      currentCliOutputIndex = 0
+    }
+    return next.cliOutput.slice(currentCliOutputIndex)
+  }
+
+  const resetCliOutput = () => {
+    currentCliOutputIndex = next.cliOutput.length
+  }
+
   const { next, isNextDev, isTurbopack, isNextDeploy } = nextTestSetup({
     files: join(__dirname, 'fixtures', 'simple'),
   })
@@ -109,12 +126,55 @@ describe('app-root-param-getters - simple', () => {
     )
     expect(response.status()).toBe(500)
     if (!isNextDeploy) {
-      expect(next.cliOutput).toInclude(
+      expect(getCliOutput()).toInclude(
         "`import('next/root-params').lang()` was used inside a Server Action. This is not supported. Functions from 'next/root-params' can only be called in the context of a route."
       )
     }
   })
 
+  it('should not error when rerendering the page after a server action', async () => {
+    const params = { lang: 'en', locale: 'us' }
+    const browser = await next.browser(
+      `/${params.lang}/${params.locale}/rerender-after-server-action`
+    )
+    expect(await browser.elementById('root-params').text()).toBe(
+      `${params.lang} ${params.locale}`
+    )
+    const initialDate = await browser.elementById('timestamp')
+
+    // Run a server action and rerender the page
+    const tracker = createRequestTracker(browser)
+    const [, response] = await tracker.captureResponse(
+      async () => {
+        await browser.elementByCss('button[type="submit"]').click()
+      },
+      {
+        request: {
+          method: 'POST',
+          pathname: `/${params.lang}/${params.locale}/rerender-after-server-action`,
+        },
+      }
+    )
+    // We're using lang() outside of an action, so we should see no errors
+    expect(response.status()).toBe(200)
+    if (!isNextDeploy) {
+      expect(getCliOutput()).not.toInclude(
+        "`import('next/root-params').lang()` was used inside a Server Action. This is not supported. Functions from 'next/root-params' can only be called in the context of a route."
+      )
+    }
+
+    await retry(async () => {
+      // The page should've been rerendered because of the cookie update
+      const updatedDate = await browser.elementById('timestamp')
+      expect(initialDate).not.toEqual(updatedDate)
+    })
+
+    // It should still display correct root params
+    expect(await browser.elementById('root-params').text()).toBe(
+      `${params.lang} ${params.locale}`
+    )
+  })
+
   // TODO(root-params): add support for route handlers
   it('should error when used in a route handler (until we implement it)', async () => {
     const params = { lang: 'en', locale: 'us' }
@@ -123,7 +183,7 @@ describe('app-root-param-getters - simple', () => {
     )
     expect(response.status).toBe(500)
     if (!isNextDeploy) {
-      expect(next.cliOutput).toInclude(
+      expect(getCliOutput()).toInclude(
         "Route /[lang]/[locale]/route-handler used `import('next/root-params').lang()` inside a Route Handler. Support for this API in Route Handlers is planned for a future version of Next.js."
       )
     }