[after] fix: execute revalidates added within unstable_after() (#70458)

### What?

Execute revalidations (written into
`staticGenerationStore.{revalidatedTags,pendingRevalidates,pendingRevalidateWrites}`) that were
added during `unstable_after` callbacks.

### Why?

Previously, if `revalidatePath`/`revalidateTag` were called in an
`unstable_after` callback, nothing would happen.

I missed the fact that other codepaths (app-render, action-handler,
app-route) all [manually call
`staticGenerationStore.incrementalCache.revalidateTag`](https://github.com/vercel/next.js/blob/79f4490b0abda3fa7129c49b402dbadf6eadd79e/packages/next/src/server/app-render/app-render.tsx#L1072-L1080)
to actually perform revalidations scheduled via
`revalidatePath`/`revalidateTag`, and `after-context` wasn't doing that,
so nothing was happening.

packages/next/src/server/after/after-context.ts

@@ -8,6 +8,8 @@ import type { RequestLifecycleOpts } from '../base-server'
 import type { AfterCallback, AfterTask } from './after'
 import { InvariantError } from '../../shared/lib/invariant-error'
 import { isThenable } from '../../shared/lib/is-thenable'
+import { staticGenerationAsyncStorage } from '../../client/components/static-generation-async-storage.external'
+import { withExecuteRevalidates } from './revalidation-utils'
 
 export type AfterContextOpts = {
   waitUntil: RequestLifecycleOpts['waitUntil'] | undefined
@@ -105,10 +107,14 @@ export class AfterContext {
     const readonlyRequestStore: RequestStore =
       wrapRequestStoreForAfterCallbacks(requestStore)
 
-    return requestAsyncStorage.run(readonlyRequestStore, () => {
-      this.callbackQueue.start()
-      return this.callbackQueue.onIdle()
-    })
+    const staticGenerationStore = staticGenerationAsyncStorage.getStore()
+
+    return withExecuteRevalidates(staticGenerationStore, () =>
+      requestAsyncStorage.run(readonlyRequestStore, async () => {
+        this.callbackQueue.start()
+        await this.callbackQueue.onIdle()
+      })
+    )
   }
 }
 

packages/next/src/server/after/revalidation-utils.ts

@@ -0,0 +1,77 @@
+import type { StaticGenerationStore } from '../../client/components/static-generation-async-storage.external'
+
+/** Run a callback, and execute any *new* revalidations added during its runtime. */
+export async function withExecuteRevalidates<T>(
+  store: StaticGenerationStore | undefined,
+  callback: () => Promise<T>
+): Promise<T> {
+  if (!store) {
+    return callback()
+  }
+  // If we executed any revalidates during the request, then we don't want to execute them again.
+  // save the state so we can check if anything changed after we're done running callbacks.
+  const savedRevalidationState = cloneRevalidationState(store)
+  try {
+    return await callback()
+  } finally {
+    // Check if we have any new revalidates, and if so, wait until they are all resolved.
+    const newRevalidates = diffRevalidationState(
+      savedRevalidationState,
+      cloneRevalidationState(store)
+    )
+    await executeRevalidates(store, newRevalidates)
+  }
+}
+
+type RevalidationState = Required<
+  Pick<
+    StaticGenerationStore,
+    'revalidatedTags' | 'pendingRevalidates' | 'pendingRevalidateWrites'
+  >
+>
+
+function cloneRevalidationState(
+  store: StaticGenerationStore
+): RevalidationState {
+  return {
+    revalidatedTags: store.revalidatedTags ? [...store.revalidatedTags] : [],
+    pendingRevalidates: { ...store.pendingRevalidates },
+    pendingRevalidateWrites: store.pendingRevalidateWrites
+      ? [...store.pendingRevalidateWrites]
+      : [],
+  }
+}
+
+function diffRevalidationState(
+  prev: RevalidationState,
+  curr: RevalidationState
+): RevalidationState {
+  const prevTags = new Set(prev.revalidatedTags)
+  const prevRevalidateWrites = new Set(prev.pendingRevalidateWrites)
+  return {
+    revalidatedTags: curr.revalidatedTags.filter((tag) => !prevTags.has(tag)),
+    pendingRevalidates: Object.fromEntries(
+      Object.entries(curr.pendingRevalidates).filter(
+        ([key]) => !(key in prev.pendingRevalidates)
+      )
+    ),
+    pendingRevalidateWrites: curr.pendingRevalidateWrites.filter(
+      (promise) => !prevRevalidateWrites.has(promise)
+    ),
+  }
+}
+
+async function executeRevalidates(
+  staticGenerationStore: StaticGenerationStore,
+  {
+    revalidatedTags,
+    pendingRevalidates,
+    pendingRevalidateWrites,
+  }: RevalidationState
+) {
+  return Promise.all([
+    staticGenerationStore.incrementalCache?.revalidateTag(revalidatedTags),
+    ...Object.values(pendingRevalidates),
+    ...pendingRevalidateWrites,
+  ])
+}

test/e2e/app-dir/next-after-app-deploy/app/edge/dynamic-page/page.js

@@ -0,0 +1 @@
+export { default } from '../../nodejs/dynamic-page/page'

test/e2e/app-dir/next-after-app-deploy/app/edge/layout.js

@@ -0,0 +1,3 @@
+export { default } from '../nodejs/layout'
+
+export const runtime = 'edge'

test/e2e/app-dir/next-after-app-deploy/app/edge/middleware/page.js

@@ -0,0 +1 @@
+export { default } from '../../nodejs/middleware/page'

test/e2e/app-dir/next-after-app-deploy/app/edge/route/route.js

@@ -0,0 +1,4 @@
+export { GET } from '../../nodejs/route/route'
+
+export const runtime = 'edge'
+export const dynamic = 'force-dynamic'

test/e2e/app-dir/next-after-app-deploy/app/edge/server-action/page.js

@@ -0,0 +1 @@
+export { default } from '../../nodejs/server-action/page'

test/e2e/app-dir/next-after-app-deploy/app/layout.js

@@ -0,0 +1,10 @@
+export default function AppLayout({ children }) {
+  return (
+    <html>
+      <head>
+        <title>after</title>
+      </head>
+      <body>{children}</body>
+    </html>
+  )
+}

test/e2e/app-dir/next-after-app-deploy/app/nodejs/dynamic-page/page.js

@@ -0,0 +1,11 @@
+import { unstable_after as after } from 'next/server'
+import { revalidateTimestampPage } from '../../timestamp/revalidate'
+import { pathPrefix } from '../../path-prefix'
+
+export default function Page() {
+  after(async () => {
+    await revalidateTimestampPage(pathPrefix + `/dynamic-page`)
+  })
+
+  return <div>Page with after()</div>
+}

test/e2e/app-dir/next-after-app-deploy/app/nodejs/layout.js

@@ -0,0 +1,5 @@
+export const runtime = 'nodejs'
+
+export default function Layout({ children }) {
+  return <>{children}</>
+}

test/e2e/app-dir/next-after-app-deploy/app/nodejs/middleware/page.js

@@ -0,0 +1,3 @@
+export default function Page() {
+  return <div>Redirect</div>
+}

test/e2e/app-dir/next-after-app-deploy/app/nodejs/route/route.js

@@ -0,0 +1,15 @@
+import { unstable_after as after } from 'next/server'
+import { revalidateTimestampPage } from '../../timestamp/revalidate'
+import { pathPrefix } from '../../path-prefix'
+
+export const runtime = 'nodejs'
+export const dynamic = 'force-dynamic'
+
+export async function GET() {
+  const data = { message: 'Hello, world!' }
+  after(async () => {
+    await revalidateTimestampPage(pathPrefix + `/route`)
+  })
+
+  return Response.json({ data })
+}

test/e2e/app-dir/next-after-app-deploy/app/nodejs/server-action/page.js

@@ -0,0 +1,20 @@
+import { unstable_after as after } from 'next/server'
+import { revalidateTimestampPage } from '../../timestamp/revalidate'
+import { pathPrefix } from '../../path-prefix'
+
+export default function Page() {
+  return (
+    <div>
+      <form
+        action={async () => {
+          'use server'
+          after(async () => {
+            await revalidateTimestampPage(pathPrefix + `/server-action`)
+          })
+        }}
+      >
+        <button type="submit">Submit</button>
+      </form>
+    </div>
+  )
+}

test/e2e/app-dir/next-after-app-deploy/app/path-prefix.js

@@ -0,0 +1 @@
+export const pathPrefix = '/' + process.env.NEXT_RUNTIME

test/e2e/app-dir/next-after-app-deploy/app/timestamp/key/[key]/page.js

@@ -0,0 +1,36 @@
+import Link from 'next/link'
+
+export const dynamic = 'error'
+export const revalidate = 3600 // arbitrarily long, just so that it doesn't happen during a test run
+export const dynamicParams = true
+
+export async function generateStaticParams() {
+  return []
+}
+
+export default async function Page({ params }) {
+  const { key } = await params
+  const data = {
+    key,
+    timestamp: Date.now(),
+  }
+  console.log('/timestamp/key/[key] rendered', data)
+
+  let path = null
+  try {
+    const decoded = decodeURIComponent(key)
+    new URL(decoded, 'http://__n')
+    path = decoded
+  } catch (err) {}
+
+  return (
+    <>
+      {path !== null && (
+        <Link prefetch={false} href={path}>
+          Go to {path}
+        </Link>
+      )}
+      <div id="page-info">{JSON.stringify(data)}</div>
+    </>
+  )
+}

test/e2e/app-dir/next-after-app-deploy/app/timestamp/revalidate.js

@@ -0,0 +1,33 @@
+import { revalidatePath } from 'next/cache'
+
+export async function revalidateTimestampPage(/** @type {string} */ key) {
+  const path = `/timestamp/key/${encodeURIComponent(key)}`
+
+  const sleepDuration = getSleepDuration()
+  if (sleepDuration > 0) {
+    console.log(`revalidateTimestampPage :: sleeping for ${sleepDuration} ms`)
+    await sleep(sleepDuration)
+  }
+
+  console.log('revalidateTimestampPage :: revalidating', path)
+  revalidatePath(path)
+}
+
+const WAIT_BEFORE_REVALIDATING_DEFAULT = 1000
+
+function getSleepDuration() {
+  const raw = process.env.WAIT_BEFORE_REVALIDATING
+  if (!raw) return WAIT_BEFORE_REVALIDATING_DEFAULT
+
+  const parsed = Number.parseInt(raw)
+  if (Number.isNaN(parsed)) {
+    throw new Error(
+      `WAIT_BEFORE_REVALIDATING must be a valid number, got: ${JSON.stringify(raw)}`
+    )
+  }
+  return parsed
+}
+
+function sleep(ms) {
+  return new Promise((resolve) => setTimeout(resolve, ms))
+}

test/e2e/app-dir/next-after-app-deploy/app/timestamp/trigger-revalidate/route.js

@@ -0,0 +1,14 @@
+import { revalidateTimestampPage } from '../revalidate'
+
+export async function POST(/** @type {Request} */ request) {
+  // we can't call revalidatePath from middleware, so we need to do it from here instead
+  const path = new URL(request.url).searchParams.get('path')
+  if (!path) {
+    return Response.json(
+      { message: 'Missing "path" search param' },
+      { status: 400 }
+    )
+  }
+  await revalidateTimestampPage(path)
+  return Response.json({})
+}