fix(after): broken error handling for unstable_after(promise) (vercel…

…#71373)

In vercel#71231 we've added `onAfterTaskError`, which is intended for handling
errors that happened within `unstable_after` (currently only used to
fail the build if anything is thrown during prerendering). It was being
called correctly for the callback form, `unstable_after(() => { ... })`,
but we missed the fact that `unstable_after(promise)` is also valid, and
weren't calling it for rejected promises.

(Apparently we also weren't printing any kind of error message for
those, and it's been that way since the start. my bad.)

This PR unifies the error handling for callbacks and promises.

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

@@ -379,18 +379,22 @@ describe('AfterContext', () => {
     const promise3 = new DetachedPromise<string>()
     const afterCallback3 = jest.fn(() => promise3.promise)
 
+    const thrownFromPromise4 = new Error('4')
+    const promise4 = Promise.reject(thrownFromPromise4)
+
     workAsyncStorage.run(workStore, () =>
       workUnitAsyncStorage.run(createMockWorkUnitStore(), () => {
         after(afterCallback1)
         after(afterCallback2)
         after(afterCallback3)
+        after(promise4)
       })
     )
 
     expect(afterCallback1).not.toHaveBeenCalled()
     expect(afterCallback2).not.toHaveBeenCalled()
     expect(afterCallback3).not.toHaveBeenCalled()
-    expect(waitUntil).toHaveBeenCalledTimes(1)
+    expect(waitUntil).toHaveBeenCalledTimes(1 + 1) // once for callbacks, once for the promise
 
     // the response is done.
     onCloseCallback!()
@@ -399,8 +403,9 @@ describe('AfterContext', () => {
     expect(afterCallback1).toHaveBeenCalledTimes(1)
     expect(afterCallback2).toHaveBeenCalledTimes(1)
     expect(afterCallback3).toHaveBeenCalledTimes(1)
-    expect(waitUntil).toHaveBeenCalledTimes(1)
+    expect(waitUntil).toHaveBeenCalledTimes(1 + 1) // once for callbacks, once for the promise
 
+    // resolve any pending promises we have
     const thrownFromCallback1 = new Error('1')
     promise1.reject(thrownFromCallback1)
     promise3.resolve('3')
@@ -409,6 +414,7 @@ describe('AfterContext', () => {
     expect(results).toEqual([undefined])
     expect(onTaskError).toHaveBeenCalledWith(thrownFromCallback2)
     expect(onTaskError).toHaveBeenCalledWith(thrownFromCallback1)
+    expect(onTaskError).toHaveBeenCalledWith(thrownFromPromise4)
   })
 
   it('throws from after() if waitUntil is not provided', async () => {

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

@@ -37,11 +37,10 @@ export class AfterContext {
 
   public after(task: AfterTask): void {
     if (isThenable(task)) {
-      task.catch(() => {}) // avoid unhandled rejection crashes
       if (!this.waitUntil) {
         errorWaitUntilNotAvailable()
       }
-      this.waitUntil(task)
+      this.waitUntil(task.catch((error) => this.reportTaskError(error)))
     } else if (typeof task === 'function') {
       // TODO(after): implement tracing
       this.addCallback(task)
@@ -85,14 +84,8 @@ export class AfterContext {
     const wrappedCallback = bindSnapshot(async () => {
       try {
         await callback()
-      } catch (err) {
-        // TODO(after): this is fine for now, but will need better intergration with our error reporting.
-        // TODO(after): should we log this if we have a onTaskError callback?
-        console.error(
-          'An error occurred in a function passed to `unstable_after()`:',
-          err
-        )
-        this.onTaskError?.(err)
+      } catch (error) {
+        this.reportTaskError(error)
       }
     })
 
@@ -121,6 +114,30 @@ export class AfterContext {
       return this.callbackQueue.onIdle()
     })
   }
+
+  private reportTaskError(error: unknown) {
+    // TODO(after): this is fine for now, but will need better intergration with our error reporting.
+    // TODO(after): should we log this if we have a onTaskError callback?
+    console.error(
+      'An error occurred in a function passed to `unstable_after()`:',
+      error
+    )
+    if (this.onTaskError) {
+      // this is very defensive, but we really don't want anything to blow up in an error handler
+      try {
+        this.onTaskError?.(error)
+      } catch (handlerError) {
+        console.error(
+          new InvariantError(
+            '`onTaskError` threw while handling an error thrown from an `unstable_after` task',
+            {
+              cause: handlerError,
+            }
+          )
+        )
+      }
+    }
+  }
 }
 
 function errorWaitUntilNotAvailable(): never {

test/e2e/app-dir/next-after-app-static/build-time-error/app/page-throws-in-after/page.tsx → test/e2e/app-dir/next-after-app-static/build-time-error/app/page-throws-in-after/callback/page.tsx

@@ -7,7 +7,9 @@ export const dynamic = 'error'
 export default function Index() {
   after(async () => {
     await setTimeout(500)
-    throw new Error('Error thrown from unstable_after: /page-throws-in-after')
+    throw new Error(
+      'My cool error thrown inside unstable_after on route "/page-throws-in-after/callback"'
+    )
   })
   return <div>Page with after()</div>
 }

test/e2e/app-dir/next-after-app-static/build-time-error/app/page-throws-in-after/promise/page.tsx

@@ -0,0 +1,16 @@
+import * as React from 'react'
+import { unstable_after as after } from 'next/server'
+import { setTimeout } from 'timers/promises'
+
+export const dynamic = 'error'
+
+export default function Index() {
+  const promise = (async () => {
+    await setTimeout(500)
+    throw new Error(
+      'My cool error thrown inside unstable_after on route "/page-throws-in-after/promise"'
+    )
+  })()
+  after(promise)
+  return <div>Page with after()</div>
+}

test/e2e/app-dir/next-after-app-static/build-time-error/app/route-throws-in-after/route.ts → test/e2e/app-dir/next-after-app-static/build-time-error/app/route-throws-in-after/callback/route.ts

@@ -6,7 +6,9 @@ export const dynamic = 'error'
 export async function GET() {
   after(async () => {
     await setTimeout(500)
-    throw new Error('Error thrown from unstable_after: /route-throws-in-after')
+    throw new Error(
+      'My cool error thrown inside unstable_after on route "/route-throws-in-after/callback"'
+    )
   })
   return new Response()
 }

test/e2e/app-dir/next-after-app-static/build-time-error/app/route-throws-in-after/promise/route.ts

@@ -0,0 +1,15 @@
+import { unstable_after as after } from 'next/server'
+import { setTimeout } from 'timers/promises'
+
+export const dynamic = 'error'
+
+export async function GET() {
+  const promise = (async () => {
+    await setTimeout(500)
+    throw new Error(
+      'My cool error thrown inside unstable_after on route "/route-throws-in-after/promise"'
+    )
+  })()
+  after(promise)
+  return new Response()
+}

test/e2e/app-dir/next-after-app-static/build-time-error/build-time-error.test.ts

@@ -17,11 +17,44 @@ _describe('unstable_after() in static pages - thrown errors', () => {
     const buildResult = await next.build()
     expect(buildResult?.exitCode).toBe(1)
 
-    expect(next.cliOutput).toContain(
-      'Error thrown from unstable_after: /page-throws-in-after'
-    )
-    expect(next.cliOutput).toContain(
-      'Error thrown from unstable_after: /route-throws-in-after'
-    )
+    {
+      const path = '/page-throws-in-after/callback'
+      expect(next.cliOutput).toContain(
+        `Error occurred prerendering page "${path}"`
+      )
+      expect(next.cliOutput).toContain(
+        `My cool error thrown inside unstable_after on route "${path}"`
+      )
+    }
+
+    {
+      const path = '/page-throws-in-after/promise'
+      expect(next.cliOutput).toContain(
+        `Error occurred prerendering page "${path}"`
+      )
+      expect(next.cliOutput).toContain(
+        `My cool error thrown inside unstable_after on route "${path}"`
+      )
+    }
+
+    {
+      const path = '/route-throws-in-after/callback'
+      expect(next.cliOutput).toContain(
+        `Error occurred prerendering page "${path}"`
+      )
+      expect(next.cliOutput).toContain(
+        `My cool error thrown inside unstable_after on route "${path}"`
+      )
+    }
+
+    {
+      const path = '/route-throws-in-after/promise'
+      expect(next.cliOutput).toContain(
+        `Error occurred prerendering page "${path}"`
+      )
+      expect(next.cliOutput).toContain(
+        `My cool error thrown inside unstable_after on route "${path}"`
+      )
+    }
   })
 })