@trpc/server 11.0.0-alpha-tmp-app-router-example.388 → 11.0.0-alpha-tmp-issues-5851-take-two.448

package/src/adapters/ws.ts

@@ -12,19 +12,26 @@ import {
   transformTRPCResponse,
   TRPCError,
 } from '../@trpc/server';
-import type { BaseHandlerOptions } from '../@trpc/server/http';
+import type { TRPCRequestInfo } from '../@trpc/server/http';
+import { toURL, type BaseHandlerOptions } from '../@trpc/server/http';
 import { parseTRPCMessage } from '../@trpc/server/rpc';
 // @trpc/server/rpc
 import type {
-  JSONRPC2,
   TRPCClientOutgoingMessage,
+  TRPCConnectionParamsMessage,
   TRPCReconnectNotification,
   TRPCResponseMessage,
 } from '../@trpc/server/rpc';
+import { parseConnectionParamsFromUnknown } from '../http';
 import { isObservable } from '../observable';
-import type { Unsubscribable } from '../observable';
+import { observableToAsyncIterable } from '../observable/observable';
 // eslint-disable-next-line no-restricted-imports
-import type { MaybePromise } from '../unstable-core-do-not-import';
+import {
+  isAsyncIterable,
+  isObject,
+  run,
+  type MaybePromise,
+} from '../unstable-core-do-not-import';
 import type { NodeHTTPCreateContextFnOptions } from './node-http';
 
 /**
@@ -36,9 +43,9 @@ const WEBSOCKET_OPEN = 1; /* ws.WebSocket.OPEN */
 /**
  * @public
  */
-export type CreateWSSContextFnOptions = Omit<
-  NodeHTTPCreateContextFnOptions<IncomingMessage, ws.WebSocket>,
-  'info'
+export type CreateWSSContextFnOptions = NodeHTTPCreateContextFnOptions<
+  IncomingMessage,
+  ws.WebSocket
 >;
 
 /**
@@ -81,6 +88,7 @@ export type WSSHandlerOptions<TRouter extends AnyRouter> =
     };
   };
 
+const unsetContextSymbol = Symbol('unsetContext');
 export function getWSConnectionHandler<TRouter extends AnyRouter>(
   opts: WSConnectionHandlerOptions<TRouter>,
 ) {
@@ -88,7 +96,7 @@ export function getWSConnectionHandler<TRouter extends AnyRouter>(
   const { transformer } = router._def._config;
 
   return async (client: ws.WebSocket, req: IncomingMessage) => {
-    const clientSubscriptions = new Map<number | string, Unsubscribable>();
+    const clientSubscriptions = new Map<number | string, AbortController>();
 
     function respond(untransformedJSON: TRPCResponseMessage) {
       client.send(
@@ -98,22 +106,65 @@ export function getWSConnectionHandler<TRouter extends AnyRouter>(
       );
     }
 
-    function stopSubscription(
-      subscription: Unsubscribable,
-      { id, jsonrpc }: JSONRPC2.BaseEnvelope & { id: JSONRPC2.RequestId },
-    ) {
-      subscription.unsubscribe();
-
-      respond({
-        id,
-        jsonrpc,
-        result: {
-          type: 'stopped',
-        },
+    function createCtxPromise(
+      getConnectionParams: () => TRPCRequestInfo['connectionParams'],
+    ): Promise<inferRouterContext<TRouter>> {
+      return run(async () => {
+        ctx = await createContext?.({
+          req,
+          res: client,
+          info: {
+            connectionParams: getConnectionParams(),
+            calls: [],
+            isBatchCall: false,
+            accept: null,
+            type: 'unknown',
+          },
+        });
+
+        return ctx;
+      }).catch((cause) => {
+        const error = getTRPCErrorFromUnknown(cause);
+        opts.onError?.({
+          error,
+          path: undefined,
+          type: 'unknown',
+          ctx,
+          req,
+          input: undefined,
+        });
+        respond({
+          id: null,
+          error: getErrorShape({
+            config: router._def._config,
+            error,
+            type: 'unknown',
+            path: undefined,
+            input: undefined,
+            ctx,
+          }),
+        });
+
+        // close in next tick
+        (global.setImmediate ?? global.setTimeout)(() => {
+          client.close();
+        });
+
+        throw error;
       });
     }
 
-    const ctxPromise = createContext?.({ req, res: client });
+    /**
+     * promise for initializing the context
+     *
+     * - the context promise will be created immediately on connection if no connectionParams are expected
+     * - if connection params are expected, they will be created once received
+     */
+    let ctxPromise =
+      toURL(req.url ?? '').searchParams.get('connectionParams') === '1'
+        ? unsetContextSymbol
+        : createCtxPromise(() => null);
+
     let ctx: inferRouterContext<TRouter> | undefined = undefined;
 
     async function handleRequest(msg: TRPCClientOutgoingMessage) {
@@ -126,11 +177,7 @@ export function getWSConnectionHandler<TRouter extends AnyRouter>(
         });
       }
       if (msg.method === 'subscription.stop') {
-        const sub = clientSubscriptions.get(id);
-        if (sub) {
-          stopSubscription(sub, { id, jsonrpc });
-        }
-        clientSubscriptions.delete(id);
+        clientSubscriptions.get(id)?.abort();
         return;
       }
       const { path, input } = msg.params;
@@ -146,14 +193,7 @@ export function getWSConnectionHandler<TRouter extends AnyRouter>(
           type,
         });
 
-        if (type === 'subscription') {
-          if (!isObservable(result)) {
-            throw new TRPCError({
-              message: `Subscription ${path} did not return an observable`,
-              code: 'INTERNAL_SERVER_ERROR',
-            });
-          }
-        } else {
+        if (type !== 'subscription') {
           // send the value as data if the method is not a subscription
           respond({
             id,
@@ -166,62 +206,112 @@ export function getWSConnectionHandler<TRouter extends AnyRouter>(
           return;
         }
 
-        const observable = result;
-        const sub = observable.subscribe({
-          next(data) {
-            respond({
-              id,
-              jsonrpc,
-              result: {
-                type: 'data',
-                data,
-              },
-            });
-          },
-          error(err) {
-            const error = getTRPCErrorFromUnknown(err);
-            opts.onError?.({ error, path, type, ctx, req, input });
-            respond({
-              id,
-              jsonrpc,
-              error: getErrorShape({
-                config: router._def._config,
-                error,
-                type,
-                path,
-                input,
-                ctx,
-              }),
-            });
-          },
-          complete() {
-            respond({
-              id,
-              jsonrpc,
-              result: {
-                type: 'stopped',
-              },
-            });
-          },
-        });
+        if (!isObservable(result) && !isAsyncIterable(result)) {
+          throw new TRPCError({
+            message: `Subscription ${path} did not return an observable or a AsyncGenerator`,
+            code: 'INTERNAL_SERVER_ERROR',
+          });
+        }
+
         /* istanbul ignore next -- @preserve */
         if (client.readyState !== WEBSOCKET_OPEN) {
           // if the client got disconnected whilst initializing the subscription
           // no need to send stopped message if the client is disconnected
-          sub.unsubscribe();
+
           return;
         }
 
         /* istanbul ignore next -- @preserve */
         if (clientSubscriptions.has(id)) {
           // duplicate request ids for client
-          stopSubscription(sub, { id, jsonrpc });
+
           throw new TRPCError({
             message: `Duplicate id ${id}`,
             code: 'BAD_REQUEST',
           });
         }
-        clientSubscriptions.set(id, sub);
+
+        const iterable = isObservable(result)
+          ? observableToAsyncIterable(result)
+          : result;
+
+        const iterator: AsyncIterator<unknown> =
+          iterable[Symbol.asyncIterator]();
+        const abortController = new AbortController();
+
+        const abortPromise = new Promise<'abort'>((resolve) => {
+          abortController.signal.onabort = () => resolve('abort');
+        });
+
+        run(async () => {
+          while (true) {
+            const next = await Promise.race([
+              iterator.next().catch(getTRPCErrorFromUnknown),
+              abortPromise,
+            ]);
+
+            if (next === 'abort') {
+              await iterator.return?.();
+              break;
+            }
+            if (next instanceof Error) {
+              const error = getTRPCErrorFromUnknown(next);
+              opts.onError?.({ error, path, type, ctx, req, input });
+              respond({
+                id,
+                jsonrpc,
+                error: getErrorShape({
+                  config: router._def._config,
+                  error,
+                  type,
+                  path,
+                  input,
+                  ctx,
+                }),
+              });
+              break;
+            }
+            if (next.done) {
+              break;
+            }
+
+            respond({
+              id,
+              jsonrpc,
+              result: {
+                type: 'data',
+                data: next.value,
+              },
+            });
+          }
+
+          await iterator.return?.();
+          respond({
+            id,
+            jsonrpc,
+            result: {
+              type: 'stopped',
+            },
+          });
+          clientSubscriptions.delete(id);
+        }).catch((cause) => {
+          const error = getTRPCErrorFromUnknown(cause);
+          opts.onError?.({ error, path, type, ctx, req, input });
+          respond({
+            id,
+            jsonrpc,
+            error: getErrorShape({
+              config: router._def._config,
+              error,
+              type,
+              path,
+              input,
+              ctx,
+            }),
+          });
+          abortController.abort();
+        });
+        clientSubscriptions.set(id, abortController);
 
         respond({
           id,
@@ -249,8 +339,31 @@ export function getWSConnectionHandler<TRouter extends AnyRouter>(
       }
     }
     client.on('message', async (message) => {
+      if (ctxPromise === unsetContextSymbol) {
+        // If the ctxPromise wasn't created immediately, we're expecting the first message to be a TRPCConnectionParamsMessage
+        ctxPromise = createCtxPromise(() => {
+          let msg;
+          try {
+            msg = JSON.parse(message.toString()) as TRPCConnectionParamsMessage;
+
+            if (!isObject(msg)) {
+              throw new Error('Message was not an object');
+            }
+          } catch (cause) {
+            throw new TRPCError({
+              code: 'PARSE_ERROR',
+              message: `Malformed TRPCConnectionParamsMessage`,
+              cause,
+            });
+          }
+
+          const connectionParams = parseConnectionParamsFromUnknown(msg.data);
+
+          return connectionParams;
+        });
+        return;
+      }
       try {
-        // eslint-disable-next-line @typescript-eslint/no-base-to-string
         const msgJSON: unknown = JSON.parse(message.toString());
         const msgs: unknown[] = Array.isArray(msgJSON) ? msgJSON : [msgJSON];
         const promises = msgs
@@ -294,42 +407,15 @@ export function getWSConnectionHandler<TRouter extends AnyRouter>(
 
     client.once('close', () => {
       for (const sub of clientSubscriptions.values()) {
-        sub.unsubscribe();
+        sub.abort();
       }
       clientSubscriptions.clear();
     });
-    async function createContextAsync() {
-      try {
-        ctx = await ctxPromise;
-      } catch (cause) {
-        const error = getTRPCErrorFromUnknown(cause);
-        opts.onError?.({
-          error,
-          path: undefined,
-          type: 'unknown',
-          ctx,
-          req,
-          input: undefined,
-        });
-        respond({
-          id: null,
-          error: getErrorShape({
-            config: router._def._config,
-            error,
-            type: 'unknown',
-            path: undefined,
-            input: undefined,
-            ctx,
-          }),
-        });
 
-        // close in next tick
-        (global.setImmediate ?? global.setTimeout)(() => {
-          client.close();
-        });
-      }
+    if (ctxPromise !== unsetContextSymbol) {
+      // prevent unhandled promise rejection errors
+      await ctxPromise.catch(() => null);
     }
-    await createContextAsync();
   };
 }
 

package/src/observable/observable.ts

@@ -4,6 +4,7 @@ import type {
   OperatorFunction,
   TeardownLogic,
   UnaryFunction,
+  Unsubscribable,
 } from './types';
 
 /** @public */
@@ -136,3 +137,65 @@ export function observableToPromise<TValue>(
     abort: abort!,
   };
 }
+
+/**
+ * @internal
+ */
+function observableToReadableStream<TValue>(
+  observable: Observable<TValue, unknown>,
+): ReadableStream<TValue> {
+  let unsub: Unsubscribable | null = null;
+  return new ReadableStream<TValue>({
+    start(controller) {
+      unsub = observable.subscribe({
+        next(data) {
+          controller.enqueue(data);
+        },
+        error(err) {
+          controller.error(err);
+        },
+        complete() {
+          controller.close();
+        },
+      });
+    },
+    cancel() {
+      unsub?.unsubscribe();
+    },
+  });
+}
+
+export function observableToAsyncIterable<TValue>(
+  observable: Observable<TValue, unknown>,
+): AsyncIterable<TValue> {
+  const stream = observableToReadableStream(observable);
+
+  const reader = stream.getReader();
+  const iterator: AsyncIterator<TValue> = {
+    async next() {
+      const value = await reader.read();
+      if (value.done) {
+        return {
+          value: undefined,
+          done: true,
+        };
+      }
+      return {
+        value: value.value,
+        done: false,
+      };
+    },
+    async return() {
+      await reader.cancel();
+      return {
+        value: undefined,
+        done: true,
+      };
+    },
+  };
+  return {
+    [Symbol.asyncIterator]() {
+      return iterator;
+    },
+  };
+}

package/src/unstable-core-do-not-import/createProxy.ts

@@ -1,6 +1,6 @@
 interface ProxyCallbackOptions {
-  path: string[];
-  args: unknown[];
+  path: readonly string[];
+  args: readonly unknown[];
 }
 type ProxyCallback = (opts: ProxyCallbackOptions) => unknown;
 
@@ -8,15 +8,27 @@ const noop = () => {
   // noop
 };
 
-function createInnerProxy(callback: ProxyCallback, path: string[]) {
-  const proxy: unknown = new Proxy(noop, {
+const freezeIfAvailable = (obj: object) => {
+  if (Object.freeze) {
+    Object.freeze(obj);
+  }
+};
+
+function createInnerProxy(
+  callback: ProxyCallback,
+  path: readonly string[],
+  memo: Record<string, unknown>,
+) {
+  const cacheKey = path.join('.');
+
+  memo[cacheKey] ??= new Proxy(noop, {
     get(_obj, key) {
       if (typeof key !== 'string' || key === 'then') {
         // special case for if the proxy is accidentally treated
         // like a PromiseLike (like in `Promise.resolve(proxy)`)
         return undefined;
       }
-      return createInnerProxy(callback, [...path, key]);
+      return createInnerProxy(callback, [...path, key], memo);
     },
     apply(_1, _2, args) {
       const lastOfPath = path[path.length - 1];
@@ -34,11 +46,13 @@ function createInnerProxy(callback: ProxyCallback, path: string[]) {
           path: path.slice(0, -1),
         };
       }
+      freezeIfAvailable(opts.args);
+      freezeIfAvailable(opts.path);
       return callback(opts);
     },
   });
 
-  return proxy;
+  return memo[cacheKey];
 }
 
 /**
@@ -46,8 +60,9 @@ function createInnerProxy(callback: ProxyCallback, path: string[]) {
  *
  * @internal
  */
-export const createRecursiveProxy = (callback: ProxyCallback) =>
-  createInnerProxy(callback, []);
+export const createRecursiveProxy = <TFaux = unknown>(
+  callback: ProxyCallback,
+): TFaux => createInnerProxy(callback, [], Object.create(null)) as TFaux;
 
 /**
  * Used in place of `new Proxy` where each handler will map 1 level deep to another value.