From f9e134d52cb534346e82f215b6f8b90a9b33c7dd Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?David=20Ov=C4=8Da=C4=8D=C3=ADk?= <david.ovcacik@gmail.com>
Date: Mon, 24 May 2021 20:55:05 +0200
Subject: [PATCH] feat(ssr): Implement afterNuxtRender context method

---
 packages/types/app/index.d.ts                     |  1 +
 packages/vue-app/template/index.js                |  1 +
 packages/vue-app/template/server.js               | 11 +++++++++--
 packages/vue-app/template/utils.js                |  1 +
 test/dev/basic.ssr.test.js                        |  6 ++++++
 test/fixtures/basic/pages/special-state-after.vue | 15 +++++++++++++++
 6 files changed, 33 insertions(+), 2 deletions(-)
 create mode 100644 test/fixtures/basic/pages/special-state-after.vue

diff --git a/packages/types/app/index.d.ts b/packages/types/app/index.d.ts
index 677d005bacf3..2cafa7633bf3 100644
--- a/packages/types/app/index.d.ts
+++ b/packages/types/app/index.d.ts
@@ -72,6 +72,7 @@ export interface Context {
     redirected: boolean
     next: NextFunction
     beforeRenderFns: Array<() => any>
+    afterRenderFns: Array<() => any>
     fetchCounters: Record<string, number>
     nuxt: {
       layout: string
diff --git a/packages/vue-app/template/index.js b/packages/vue-app/template/index.js
index 13fdeca4933b..b85b6ea59c74 100644
--- a/packages/vue-app/template/index.js
+++ b/packages/vue-app/template/index.js
@@ -183,6 +183,7 @@ async function createApp(ssrContext, config = {}) {
     req: ssrContext ? ssrContext.req : undefined,
     res: ssrContext ? ssrContext.res : undefined,
     beforeRenderFns: ssrContext ? ssrContext.beforeRenderFns : undefined,
+    afterRenderFns: ssrContext ? ssrContext.afterRenderFns : undefined,
     ssrContext
   })
 
diff --git a/packages/vue-app/template/server.js b/packages/vue-app/template/server.js
index 1d62d425d1b6..c300d25213ef 100644
--- a/packages/vue-app/template/server.js
+++ b/packages/vue-app/template/server.js
@@ -85,6 +85,8 @@ export default async (ssrContext) => {
   ssrContext.next = createNext(ssrContext)
   // Used for beforeNuxtRender({ Components, nuxtState })
   ssrContext.beforeRenderFns = []
+  // Used for afterNuxtRender({ Components, nuxtState })
+  ssrContext.afterRenderFns = []
   // Nuxt object (window.{{globals.context}}, defaults to window.__NUXT__)
   ssrContext.nuxt = { <% if (features.layouts) { %>layout: 'default', <% } %>data: [], <% if (features.fetch) { %>fetch: {}, <% } %>error: null<%= (store ? ', state: null' : '') %>, serverRendered: true, routePath: '' }
   <% if (features.fetch) { %>
@@ -120,16 +122,21 @@ export default async (ssrContext) => {
   const beforeRender = async () => {
     // Call beforeNuxtRender() methods
     await Promise.all(ssrContext.beforeRenderFns.map(fn => promisify(fn, { Components, nuxtState: ssrContext.nuxt })))
-    <% if (store) { %>
+
     ssrContext.rendered = () => {
+      // Call afterNuxtRender() methods
+      ssrContext.afterRenderFns.forEach(fn => fn({ Components, nuxtState: ssrContext.nuxt }))
+
+      <% if (store) { %>
       // Add the state from the vuex store
       ssrContext.nuxt.state = store.state
+      <% } %>
+
       <% if (isFullStatic && store) { %>
       // Stop recording store mutations
       ssrContext.unsetMutationObserver()
       <% } %>
     }
-    <% } %>
   }
 
   const renderErrorPage = async () => {
diff --git a/packages/vue-app/template/utils.js b/packages/vue-app/template/utils.js
index 6d6e9978eb72..2f97e8eb7aa9 100644
--- a/packages/vue-app/template/utils.js
+++ b/packages/vue-app/template/utils.js
@@ -259,6 +259,7 @@ export async function setContext (app, context) {
     }
     if (process.server) {
       app.context.beforeNuxtRender = fn => context.beforeRenderFns.push(fn)
+      app.context.afterNuxtRender = fn => context.afterRenderFns.push(fn)
     }
     if (process.client) {
       app.context.nuxtState = window.<%= globals.context %>
diff --git a/test/dev/basic.ssr.test.js b/test/dev/basic.ssr.test.js
index acdc078ae870..12685692d64f 100644
--- a/test/dev/basic.ssr.test.js
+++ b/test/dev/basic.ssr.test.js
@@ -207,6 +207,12 @@ describe('basic ssr', () => {
     expect(window.__NUXT__.test).toBe(true)
   })
 
+  test('/special-state-after -> check window.__NUXT__.testAfter = true', async () => {
+    const window = await nuxt.server.renderAndGetWindow(url('/special-state-after'))
+    expect(window.document.title).toBe('Nuxt')
+    expect(window.__NUXT__.testAfter).toBe(true)
+  })
+
   test('/error', async () => {
     await expect(nuxt.server.renderRoute('/error', { req: {}, res: {} }))
       .rejects.toThrow('Error mouahahah')
diff --git a/test/fixtures/basic/pages/special-state-after.vue b/test/fixtures/basic/pages/special-state-after.vue
new file mode 100644
index 000000000000..58d899e4ea64
--- /dev/null
+++ b/test/fixtures/basic/pages/special-state-after.vue
@@ -0,0 +1,15 @@
+<template>
+  <h1>Special state in `window.__NUXT__`</h1>
+</template>
+
+<script>
+export default {
+  middleware ({ afterNuxtRender }) {
+    if (process.server) {
+      afterNuxtRender(({ nuxtState }) => {
+        nuxtState.testAfter = true
+      })
+    }
+  }
+}
+</script>