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

package/dist/adapters/ws.mjs

@@ -1,15 +1,19 @@
-import { TRPCError, getTRPCErrorFromUnknown } from '../unstable-core-do-not-import/error/TRPCError.mjs';
 import { getErrorShape } from '../unstable-core-do-not-import/error/getErrorShape.mjs';
+import { TRPCError, getTRPCErrorFromUnknown } from '../unstable-core-do-not-import/error/TRPCError.mjs';
+import { run, isAsyncIterable, isObject } from '../unstable-core-do-not-import/utils.mjs';
+import { parseConnectionParamsFromUnknown } from '../unstable-core-do-not-import/http/parseConnectionParams.mjs';
 import { parseTRPCMessage } from '../unstable-core-do-not-import/rpc/parseTRPCMessage.mjs';
-import { callProcedure } from '../unstable-core-do-not-import/router.mjs';
+import { isObservable, observableToAsyncIterable } from '../observable/observable.mjs';
 import { transformTRPCResponse } from '../unstable-core-do-not-import/transformer.mjs';
+import { toURL } from '../unstable-core-do-not-import/http/toURL.mjs';
 import '../unstable-core-do-not-import/rootConfig.mjs';
-import { isObservable } from '../observable/observable.mjs';
+import { callProcedure } from '../unstable-core-do-not-import/router.mjs';
 
 /**
  * Importing ws causes a build error
  * @link https://github.com/trpc/trpc/pull/5279
  */ const WEBSOCKET_OPEN = 1; /* ws.WebSocket.OPEN */ 
+const unsetContextSymbol = Symbol('unsetContext');
 function getWSConnectionHandler(opts) {
     const { createContext , router  } = opts;
     const { transformer  } = router._def._config;
@@ -18,20 +22,54 @@ function getWSConnectionHandler(opts) {
         function respond(untransformedJSON) {
             client.send(JSON.stringify(transformTRPCResponse(router._def._config, untransformedJSON)));
         }
-        function stopSubscription(subscription, { id , jsonrpc  }) {
-            subscription.unsubscribe();
-            respond({
-                id,
-                jsonrpc,
-                result: {
-                    type: 'stopped'
-                }
+        function createCtxPromise(getConnectionParams) {
+            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 = undefined;
         async function handleRequest(msg) {
             const { id , jsonrpc  } = msg;
@@ -42,14 +80,7 @@ function getWSConnectionHandler(opts) {
                 });
             }
             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;
@@ -63,14 +94,7 @@ function getWSConnectionHandler(opts) {
                     ctx,
                     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,
@@ -82,69 +106,110 @@ function getWSConnectionHandler(opts) {
                     });
                     return;
                 }
-                const observable = result;
-                const sub1 = 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
-                    sub1.unsubscribe();
                     return;
                 }
                 /* istanbul ignore next -- @preserve */ if (clientSubscriptions.has(id)) {
                     // duplicate request ids for client
-                    stopSubscription(sub1, {
-                        id,
-                        jsonrpc
-                    });
                     throw new TRPCError({
                         message: `Duplicate id ${id}`,
                         code: 'BAD_REQUEST'
                     });
                 }
-                clientSubscriptions.set(id, sub1);
+                const iterable = isObservable(result) ? observableToAsyncIterable(result) : result;
+                const iterator = iterable[Symbol.asyncIterator]();
+                const abortController = new AbortController();
+                const abortPromise = new Promise((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,
                     jsonrpc,
@@ -178,8 +243,28 @@ function getWSConnectionHandler(opts) {
             }
         }
         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());
+                        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 = JSON.parse(message.toString());
                 const msgs = Array.isArray(msgJSON) ? msgJSON : [
                     msgJSON
@@ -220,41 +305,14 @@ function getWSConnectionHandler(opts) {
         });
         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();
     };
 }
 /**