feat: add subscription types

playground/pages/index.vue

@@ -28,7 +28,6 @@ const { data: todos, pending, error, refresh } = await $client.todo.getTodos.que
   trpc: {
     abortOnUnmount: true,
   },
-  immediate: true,
 })
 </script>
 

src/client/index.ts

@@ -80,7 +80,7 @@ export function createTRPCNuxtProxyClient<TRouter extends AnyRouter>(opts: Creat
 
   const decoratedClient = createFlatProxy((key) => {
     return createNuxtProxyDecoration(key, client)
-  }) as DecoratedProcedureRecord<TRouter['_def']['record']>
+  }) as DecoratedProcedureRecord<TRouter['_def']['record'], TRouter>
 
   return decoratedClient
 }

src/client/types.ts

@@ -1,13 +1,17 @@
 import type { TRPCClientErrorLike, TRPCRequestOptions } from '@trpc/client'
+import type { TRPCSubscriptionObserver } from '@trpc/client/dist/internals/TRPCClient'
 import type {
   AnyMutationProcedure,
   AnyProcedure,
   AnyQueryProcedure,
   AnyRouter,
+  AnySubscriptionProcedure,
+  ProcedureArgs,
   ProcedureRouterRecord,
   inferProcedureInput,
   inferProcedureOutput,
 } from '@trpc/server'
+import type { Unsubscribable, inferObservableValue } from '@trpc/server/observable'
 import type {
   AsyncData,
   AsyncDataOptions,
@@ -16,17 +20,34 @@ import type {
   _Transform,
 } from 'nuxt/dist/app/composables/asyncData'
 
-// Inspired by trpc/react-query client types
+// Modified @trpc/client and @trpc/react-query types
+// https://github.com/trpc/trpc/blob/next/packages/client/src/createTRPCClientProxy.ts
 // https://github.com/trpc/trpc/blob/next/packages/react-query/src/createTRPCReact.tsx
 
 interface TRPCOptions extends TRPCRequestOptions {
   abortOnUnmount?: boolean
 }
 
-/**
- * @internal
- */
-export type DecorateProcedure<TProcedure extends AnyProcedure> = TProcedure extends AnyQueryProcedure
+type SubscriptionResolver<
+  TProcedure extends AnyProcedure,
+  TRouter extends AnyRouter,
+> = (
+  ...args: [
+    input: ProcedureArgs<TProcedure['_def']>[0],
+    opts: ProcedureArgs<TProcedure['_def']>[1] &
+    Partial<
+        TRPCSubscriptionObserver<
+          inferObservableValue<inferProcedureOutput<TProcedure>>,
+          TRPCClientErrorLike<TRouter>
+        >
+      >,
+  ]
+) => Unsubscribable
+
+type DecorateProcedure<
+  TProcedure extends AnyProcedure,
+  TRouter extends AnyRouter,
+> = TProcedure extends AnyQueryProcedure
   ? {
       query: <
       TData = inferProcedureOutput<TProcedure>,
@@ -45,15 +66,20 @@ export type DecorateProcedure<TProcedure extends AnyProcedure> = TProcedure exte
         input: inferProcedureInput<TProcedure>,
         opts?: AsyncDataOptions<TData, Transform, PickKeys> & { trpc: TRPCOptions },
       ) => AsyncData<PickFrom<ReturnType<Transform>, PickKeys>, TRPCClientErrorLike<TProcedure>>
+    } : TProcedure extends AnySubscriptionProcedure ? {
+      subscribe: SubscriptionResolver<TProcedure, TRouter>
     } : never
 
 /**
 * @internal
 */
-export type DecoratedProcedureRecord<TProcedures extends ProcedureRouterRecord> = {
+export type DecoratedProcedureRecord<
+  TProcedures extends ProcedureRouterRecord,
+  TRouter extends AnyRouter,
+> = {
   [TKey in keyof TProcedures]: TProcedures[TKey] extends AnyRouter
-    ? DecoratedProcedureRecord<TProcedures[TKey]['_def']['record']>
+    ? DecoratedProcedureRecord<TProcedures[TKey]['_def']['record'], TRouter>
     : TProcedures[TKey] extends AnyProcedure
-      ? DecorateProcedure<TProcedures[TKey]>
+      ? DecorateProcedure<TProcedures[TKey], TRouter>
       : never;
 }