From 3b5e35ef685c101af0c256d0f29e55593b4cbc97 Mon Sep 17 00:00:00 2001
From: Robert Soriano <sorianorobertc@gmail.com>
Date: Tue, 24 May 2022 09:46:15 -0700
Subject: [PATCH] feat: add composable for custom client headers

---
 package.json                    |  2 +
 playground/pages/index.vue      | 17 ++++++--
 playground/server/trpc/index.ts |  4 +-
 pnpm-lock.yaml                  | 69 +++++++++++++++++++++++++++++++++
 src/module.ts                   |  4 ++
 src/runtime/client.ts           | 24 ++++++++++++
 src/runtime/plugin.ts           | 22 +++++++----
 7 files changed, 130 insertions(+), 12 deletions(-)

diff --git a/package.json b/package.json
index b507c5d..453c8f6 100644
--- a/package.json
+++ b/package.json
@@ -35,6 +35,8 @@
     "@nuxt/kit": "^3.0.0-rc.3",
     "@trpc/client": "^9.23.3",
     "@trpc/server": "^9.23.2",
+    "@vueuse/core": "^8.5.0",
+    "@vueuse/nuxt": "^8.5.0",
     "defu": "^6.0.0",
     "h3": "^0.7.8",
     "pathe": "^0.3.0",
diff --git a/playground/pages/index.vue b/playground/pages/index.vue
index 198b5fa..adeef95 100644
--- a/playground/pages/index.vue
+++ b/playground/pages/index.vue
@@ -1,12 +1,20 @@
 <script setup lang="ts">
-const { $client } = useNuxtApp()
+const client = useClient()
+const headers = useClientHeader({ authorization: 'asdada' })
 const { data: todos, pending, error, refresh } = await useAsyncQuery(['getTodos'])
 
+console.log(headers.value)
+
+const addHeader = () => {
+  // headers.value.cookie = 'counter=69'
+  console.log(headers.value)
+}
+
 const addTodo = async () => {
   const title = Math.random().toString(36).slice(2, 7)
 
   try {
-    const result = await $client.mutation('addTodo', {
+    const result = await client.mutation('addTodo', {
       id: Date.now(),
       userId: 69,
       title,
@@ -27,7 +35,7 @@ const addTodo = async () => {
   <div v-else-if="error?.data?.code">
     Error: {{ error.data.code }}
   </div>
-  <div v-else>
+  <div v-else-if="todos">
     <ul>
       <li v-for="t in todos.slice(0, 10)" :key="t.id">
         <NuxtLink :class="{ completed: t.completed }" :to="`/todo/${t.id}`">
@@ -41,6 +49,9 @@ const addTodo = async () => {
     <button @click="() => refresh()">
       Refresh
     </button>
+    <button @click="addHeader">
+      Add header
+    </button>
   </div>
 </template>
 
diff --git a/playground/server/trpc/index.ts b/playground/server/trpc/index.ts
index 0fab987..5472076 100644
--- a/playground/server/trpc/index.ts
+++ b/playground/server/trpc/index.ts
@@ -42,8 +42,8 @@ export async function createContext(event: CompatibilityEvent) {
   // Will be available as `ctx` in all your resolvers
 
   // This is just an example of something you'd might want to do in your ctx fn
-  const x = useCookies(event)
-  console.log(x)
+  // const x = useCookies(event)
+  console.log(event.req.headers)
 
   return {
 
diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml
index f35ec8e..6135e56 100644
--- a/pnpm-lock.yaml
+++ b/pnpm-lock.yaml
@@ -10,6 +10,8 @@ importers:
       '@nuxt/module-builder': latest
       '@trpc/client': ^9.23.3
       '@trpc/server': ^9.23.2
+      '@vueuse/core': ^8.5.0
+      '@vueuse/nuxt': ^8.5.0
       bumpp: ^7.1.1
       defu: ^6.0.0
       eslint: ^8.14.0
@@ -26,6 +28,8 @@ importers:
       '@nuxt/kit': 3.0.0-rc.3
       '@trpc/client': 9.23.4_@trpc+server@9.23.4
       '@trpc/server': 9.23.4
+      '@vueuse/core': 8.5.0
+      '@vueuse/nuxt': 8.5.0
       defu: 6.0.0
       h3: 0.7.8
       pathe: 0.3.0
@@ -1264,6 +1268,22 @@ packages:
     resolution: {integrity: sha512-UBc1Pg1T3yZ97vsA2ueER0F6GbJebLHYlEi4ou1H5YL4KWvMOOWwpYo9/QpWq93wxKG6Wo13IY74Hcn/f7c7Bg==}
     dev: true
 
+  /@vueuse/core/8.5.0:
+    resolution: {integrity: sha512-VEJ6sGNsPlUp0o9BGda2YISvDZbhWJSOJu5zlp2TufRGVrLcYUKr31jyFEOj6RXzG3k/H4aCYeZyjpItfU8glw==}
+    peerDependencies:
+      '@vue/composition-api': ^1.1.0
+      vue: ^2.6.0 || ^3.2.0
+    peerDependenciesMeta:
+      '@vue/composition-api':
+        optional: true
+      vue:
+        optional: true
+    dependencies:
+      '@vueuse/metadata': 8.5.0
+      '@vueuse/shared': 8.5.0
+      vue-demi: 0.12.5
+    dev: false
+
   /@vueuse/head/0.7.6_vue@3.2.33:
     resolution: {integrity: sha512-cOWqCkT3WiF5oEpw+VVEWUJd9RLD5rc7DmnFp3cePsejp+t7686uKD9Z9ZU7Twb7R/BI8iexKTmXo9D/F3v6UA==}
     peerDependencies:
@@ -1272,6 +1292,42 @@ packages:
       vue: 3.2.33
     dev: true
 
+  /@vueuse/metadata/8.5.0:
+    resolution: {integrity: sha512-WxsD+Cd+bn+HcjpY6Dl9FJ8ywTRTT9pTwk3bCQpzEhXVYAyNczKDSahk50fCfIJKeWHhyI4B2+/ZEOxQAkUr0g==}
+    dev: false
+
+  /@vueuse/nuxt/8.5.0:
+    resolution: {integrity: sha512-riGrDwlTQbjSDxyw46oc1catXpZwzzRrEk+PTy8NQZG6uevgJ8wNuAhPVx/5oG0jYG/t2orIFfc9xGEA6uytSg==}
+    dependencies:
+      '@nuxt/kit': 3.0.0-rc.3
+      '@vueuse/core': 8.5.0
+      '@vueuse/metadata': 8.5.0
+      local-pkg: 0.4.1
+      vue-demi: 0.12.5
+    transitivePeerDependencies:
+      - '@vue/composition-api'
+      - esbuild
+      - rollup
+      - supports-color
+      - vite
+      - vue
+      - webpack
+    dev: false
+
+  /@vueuse/shared/8.5.0:
+    resolution: {integrity: sha512-qKG+SZb44VvGD4dU5cQ63z4JE2Yk39hQUecR0a9sEdJA01cx+XrxAvFKJfPooxwoiqalAVw/ktWK6xbyc/jS3g==}
+    peerDependencies:
+      '@vue/composition-api': ^1.1.0
+      vue: ^2.6.0 || ^3.2.0
+    peerDependenciesMeta:
+      '@vue/composition-api':
+        optional: true
+      vue:
+        optional: true
+    dependencies:
+      vue-demi: 0.12.5
+    dev: false
+
   /abbrev/1.1.1:
     resolution: {integrity: sha512-nne9/IiQ/hzIhY6pdDnbBtz7DjPTKrY00P/zvPSm5pOFkl6xuGrGnXn/VtTNNfNtAfZ9/1RtehkszU9qcTii0Q==}
     dev: true
@@ -6561,6 +6617,19 @@ packages:
       bundle-runner: 0.0.1
     dev: true
 
+  /vue-demi/0.12.5:
+    resolution: {integrity: sha512-BREuTgTYlUr0zw0EZn3hnhC3I6gPWv+Kwh4MCih6QcAeaTlaIX0DwOVN0wHej7hSvDPecz4jygy/idsgKfW58Q==}
+    engines: {node: '>=12'}
+    hasBin: true
+    requiresBuild: true
+    peerDependencies:
+      '@vue/composition-api': ^1.0.0-rc.1
+      vue: ^3.0.0-0 || ^2.6.0
+    peerDependenciesMeta:
+      '@vue/composition-api':
+        optional: true
+    dev: false
+
   /vue-eslint-parser/8.3.0_eslint@8.15.0:
     resolution: {integrity: sha512-dzHGG3+sYwSf6zFBa0Gi9ZDshD7+ad14DGOdTLjruRVgZXe2J+DcZ9iUhyR48z5g1PqRa20yt3Njna/veLJL/g==}
     engines: {node: ^12.22.0 || ^14.17.0 || >=16.0.0}
diff --git a/src/module.ts b/src/module.ts
index e97494a..1ede328 100644
--- a/src/module.ts
+++ b/src/module.ts
@@ -25,6 +25,9 @@ export default defineNuxtModule<ModuleOptions>({
     const handlerPath = join(nuxt.options.buildDir, 'trpc-handler.ts')
     const trpcOptionsPath = join(nuxt.options.rootDir, 'server/trpc')
 
+    // Add vueuse
+    nuxt.options.modules.push('@vueuse/nuxt')
+
     // Final resolved configuration
     const finalConfig = nuxt.options.runtimeConfig.public.trpc = defu(nuxt.options.runtimeConfig.public.trpc, {
       baseURL: options.baseURL,
@@ -35,6 +38,7 @@ export default defineNuxtModule<ModuleOptions>({
       imports.push(
         { name: 'useClient', from: join(runtimeDir, 'client') },
         { name: 'useAsyncQuery', from: join(runtimeDir, 'client') },
+        { name: 'useClientHeader', from: join(runtimeDir, 'client') },
       )
     })
 
diff --git a/src/runtime/client.ts b/src/runtime/client.ts
index 1c8e67f..03126a1 100644
--- a/src/runtime/client.ts
+++ b/src/runtime/client.ts
@@ -8,6 +8,8 @@ import type {
 import type { ProcedureRecord, inferHandlerInput, inferProcedureInput, inferProcedureOutput } from '@trpc/server'
 import type { TRPCClient, TRPCClientErrorLike } from '@trpc/client'
 import { objectHash } from 'ohash'
+import type { MaybeRef } from '@vueuse/core'
+import { useStorage } from '@vueuse/core'
 import { useAsyncData, useNuxtApp, useState } from '#app'
 import type { router } from '~/server/trpc'
 
@@ -27,6 +29,13 @@ export type TError = TRPCClientErrorLike<AppRouter>
 
 export type TQueryValues = inferProcedures<AppRouter['_def']['queries']>
 
+/**
+ * Additional header properties to pass to tRPC client.
+ *
+ * @see https://trpc.io/docs/vanilla
+ * @param pathAndInput tRPC client path and input.
+ * @param options Options to pass to useAsyncData.
+ */
 export async function useAsyncQuery<
   TPath extends keyof TQueryValues & string,
   TOutput extends TQueryValues[TPath]['output'] = TQueryValues[TPath]['output'],
@@ -59,7 +68,22 @@ export async function useAsyncQuery<
   } as any
 }
 
+/**
+ * tRPC Client.
+ *
+ * @see https://trpc.io/docs/vanilla
+ */
 export function useClient(): TRPCClient<AppRouter> {
   const { $client } = useNuxtApp()
   return $client
 }
+
+/**
+ * Additional header properties to pass to tRPC client.
+ *
+ * @see https://github.com/trpc/trpc/discussions/1686
+ * @param initialValue
+ */
+export function useClientHeaders(initialValue: MaybeRef<Record<string, any>>) {
+  return useStorage('trpc-nuxt-header', initialValue || {})
+}
diff --git a/src/runtime/plugin.ts b/src/runtime/plugin.ts
index 4874256..e7ada9b 100644
--- a/src/runtime/plugin.ts
+++ b/src/runtime/plugin.ts
@@ -4,18 +4,26 @@ import type { router } from '~/server/trpc'
 
 declare type AppRouter = typeof router
 
-export default defineNuxtPlugin(() => {
+export default defineNuxtPlugin((nuxtApp) => {
   const config = useRuntimeConfig().public.trpc
+  const headers = useRequestHeaders()
   const client = trpc.createTRPCClient<AppRouter>({
     url: `${config.baseURL}${config.trpcURL}`,
-    headers: useRequestHeaders(),
+    headers: () => {
+      let otherHeaders = {}
+      if (!process.server) {
+        const key = 'trpc-nuxt-header'
+        otherHeaders = JSON.parse(localStorage.getItem(key) || JSON.stringify({}))
+      }
+
+      return {
+        ...otherHeaders,
+        ...headers,
+      }
+    },
   })
 
-  return {
-    provide: {
-      client,
-    },
-  }
+  nuxtApp.provide('client', client)
 })
 
 declare module '#app' {