feat: use AsyncLocalStorage when working with websockets

[banjo]

Describe the feature you'd like to request

This might already be a feature, but I am not getting it to work currently.

I want to use AsyncLocalStorage in WebSockets. I currently use them for requestId purposes, and they work great with HTTP, but I cannot get it to work with WebSockets when using TRPC.

With Express I only need to add a simple middleware like so to make it work with TRPC:

const setupContext = (req, res, next) => {
    context.run(store, () => {
        const requestId = uuid();
        setRequestId(requestId);
        next();
    });
};

app.use(setupContext);

I want the same type of experience with WebSockets. This is my current setup:

const wss = new WebSocketServer({
    port: WS_PORT,
});

const handler = applyWSSHandler({
    wss,
    router: appRouter,
    createContext: opts => {
        const req = opts?.req;
        const res = opts?.res;

        return createTRPCContext({ req, res });
    },
});

The applyWSSHandler function handles all the connection logic, so I am not able to wrap it in a AsyncLocalStorage context.

Describe the solution you'd like to see

With Express and HTTP, adding a middleware is super simple, but I am not aware of a similar solution that can be used for WebSockets.

I guess one would need to wrap the whole connection event in WSS in a AsyncLocalStorage.

Something like:

wss.on("connection", () => {
    setupContext(() => {
        // TRPC logic
    });
});

And once again, the applyWSSHandler function handles most of that logic which makes it difficult to modify.

Describe alternate solutions

Perhaps there are other, more specific TRPC ways to handle this, but I am not aware enough of the internals to know about it.

I tried to add the context in the createContext callback, but that does not work. Perhaps TRPC middleware can be used somehow?

const handler = applyWSSHandler({
    wss,
    router: appRouter,
    createContext: opts => {
        setupContext(() => {
            const req = opts?.req;
            const res = opts?.res;

            return createTRPCContext({ req, res });
        });
    },
});

Additional information

AsyncLocalStorage is a very node specific API, and I assume there are more environments than Node that needs to be supported, which makes this a bit more difficult.

👨‍👧‍👦 Contributing

• [ ] 🙋‍♂️ Yes, I'd be down to file a PR implementing this feature!

Funding

• You can sponsor this specific effort via a Polar.sh pledge below
• We receive the pledge once the issue is completed & verified

[KATT]

You can already make a middleware that creates an async storage, does that suffice?

[KATT]

https://github.com/trpc/trpc/blob/84bfa8b8206570d35e1cb3a75be97d1eb3bc8aa2/packages/tests/server/validators.test.ts#L442-L447

[banjo]

Amazing, it worked!

I had some problems with the types when trying before but your example helped me get on the right track. Thank you and glad midsommar!

[github-actions[bot]]

This issue has been locked because we are very unlikely to see comments on closed issues. If you are running into a similar issue, please create a new issue. Thank you.