From afd28d713392d7cbe7d6788bde8df8e7ff899a3f Mon Sep 17 00:00:00 2001
From: Carson Full <carson_full@tsco.org>
Date: Thu, 15 May 2025 18:41:30 -0500
Subject: [PATCH] Swap AsyncResource with ALS.snapshot

---
 src/core/graphql/driver.ts | 13 ++++++++-----
 1 file changed, 8 insertions(+), 5 deletions(-)

diff --git a/src/core/graphql/driver.ts b/src/core/graphql/driver.ts
index 7fbf2ac037..55dbba72c6 100644
--- a/src/core/graphql/driver.ts
+++ b/src/core/graphql/driver.ts
@@ -13,7 +13,7 @@ import {
   type YogaServerInstance,
   type YogaServerOptions,
 } from 'graphql-yoga';
-import { AsyncResource } from 'node:async_hooks';
+import { AsyncLocalStorage } from 'node:async_hooks';
 import type { WebSocket } from 'ws';
 import { type GqlContextType } from '~/common';
 import { HttpAdapter, type IRequest } from '../http';
@@ -105,7 +105,10 @@ export class Driver extends AbstractDriver<DriverConfig> {
    *    So this allows our "yoga" plugins to be executed.
    */
   private makeWsHandler(options: DriverConfig) {
-    const asyncContextBySocket = new WeakMap<WebSocket, AsyncResource>();
+    const asyncContextBySocket = new WeakMap<
+      WebSocket,
+      <R>(fn: () => R) => R
+    >();
     interface WsExecutionArgs extends ExecutionArgs {
       socket: WebSocket;
       envelop: ReturnType<ReturnType<typeof envelop>>;
@@ -125,7 +128,7 @@ export class Driver extends AbstractDriver<DriverConfig> {
       // unique envelop (yoga) instance per request.
       execute: (wsArgs) => {
         const { envelop, socket, ...args } = wsArgs as WsExecutionArgs;
-        return asyncContextBySocket.get(socket)!.runInAsyncScope(() => {
+        return asyncContextBySocket.get(socket)!(() => {
           return envelop.execute(args);
         });
       },
@@ -134,7 +137,7 @@ export class Driver extends AbstractDriver<DriverConfig> {
         // Because this is called via socket.onmessage, we don't have
         // the same async context we started with.
         // Grab and resume it.
-        return asyncContextBySocket.get(socket)!.runInAsyncScope(() => {
+        return asyncContextBySocket.get(socket)!(() => {
           return envelop.subscribe(args);
         });
       },
@@ -174,7 +177,7 @@ export class Driver extends AbstractDriver<DriverConfig> {
 
     const wsHandler: FastifyRoute['wsHandler'] = function (socket, req) {
       // Save a reference to the current async context, so we can resume it.
-      asyncContextBySocket.set(socket, new AsyncResource('graphql-ws'));
+      asyncContextBySocket.set(socket, AsyncLocalStorage.snapshot());
       return fastifyWsHandler.call(this, socket, req);
     };
     return wsHandler;