feat: polyfill async context (#123)

.changeset/happy-seals-jog.md

@@ -0,0 +1,5 @@
+---
+"@llamaindex/workflow-core": minor
+---
+
+feat: polyfill async context

packages/core/package.json

@@ -204,5 +204,8 @@
   },
   "publishConfig": {
     "access": "public"
+  },
+  "dependencies": {
+    "simple-async-context": "^1.0.4"
   }
 }

packages/core/src/async-context/index.browser.ts

@@ -1,41 +1,3 @@
-export const createAsyncContext = <T>() => {
-  let currentStore: T | null = null;
-  return {
-    /**
-     * You must call `getContext()` in the top level of the workflow,
-     *  otherwise we will lose the async context of the workflow.
-     *
-     * @example
-     * ```
-     * workflow.handle([startEvent], async () => {
-     *   const { stream } = getContext(); // ✅ this is ok
-     *   await fetchData();
-     * });
-     *
-     * workflow.handle([startEvent], async () => {
-     *   await fetchData();
-     *   const { stream } = getContext(); // ❌ this is not ok
-     *   // we have no way
-     *   to know this code was originally part of the workflow
-     *   // w/o AsyncContext
-     * });
-     * ```
-     */
-    getStore: () => {
-      if (currentStore === null) {
-        console.warn(
-          "Woops! Looks like you are calling `getContext` after `await fn()`. Please move `getContext` to top level of handler.",
-        );
-      }
-      return currentStore;
-    },
-    run<R>(store: T, fn: () => R) {
-      currentStore = store;
-      try {
-        return fn();
-      } finally {
-        currentStore = null;
-      }
-    },
-  };
-};
+import { AsyncContext } from "simple-async-context";
+
+export { AsyncContext };

packages/core/src/async-context/index.ts

@@ -1,11 +1,17 @@
 import { AsyncLocalStorage } from "node:async_hooks";
 
-export const createAsyncContext = <T>() => {
-  const als = new AsyncLocalStorage<T>();
-  return {
-    getStore: () => als.getStore(),
-    run<R>(store: T, fn: () => R) {
-      return als.run(store, fn);
-    },
-  };
-};
+class AsyncVariable<T> {
+  als = new AsyncLocalStorage<T>();
+
+  get(): T | undefined {
+    return this.als.getStore();
+  }
+
+  run<R>(value: T, fn: () => R): R {
+    return this.als.run(value, fn);
+  }
+}
+
+export class AsyncContext {
+  static Variable = AsyncVariable;
+}

packages/core/src/core/context.ts

@@ -10,7 +10,7 @@ import {
   isPromiseLike,
   type Subscribable,
 } from "./utils";
-import { createAsyncContext } from "@llamaindex/workflow-core/async-context";
+import { AsyncContext } from "@llamaindex/workflow-core/async-context";
 import { WorkflowStream } from "./stream";
 
 export type Handler<
@@ -74,17 +74,19 @@ export type WorkflowContext = {
   >;
 };
 
-export const _executorAsyncLocalStorage = createAsyncContext<WorkflowContext>();
+export const _executorAsyncLocalStorage =
+  new AsyncContext.Variable<WorkflowContext>();
 
 export function getContext(): WorkflowContext {
-  const context = _executorAsyncLocalStorage.getStore();
+  const context = _executorAsyncLocalStorage.get();
   if (!context) {
     throw new Error("No current context found");
   }
   return context;
 }
 
-const handlerContextAsyncLocalStorage = createAsyncContext<HandlerContext>();
+const handlerContextAsyncLocalStorage =
+  new AsyncContext.Variable<HandlerContext>();
 
 const eventContextWeakMap = new WeakMap<
   WorkflowEventData<any>,

packages/core/src/mcp.ts

@@ -1,15 +1,16 @@
-import { createAsyncContext } from "@llamaindex/workflow-core/async-context";
+import { AsyncContext } from "@llamaindex/workflow-core/async-context";
 import { z, type ZodRawShape, type ZodTypeAny } from "zod";
 import type { Workflow, WorkflowEvent } from "@llamaindex/workflow-core";
 import { runWorkflow } from "./stream/run";
 import type { RequestHandlerExtra } from "@modelcontextprotocol/sdk/shared/protocol.js";
 import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
 
-const requestHandlerExtraAsyncLocalStorage =
-  createAsyncContext<RequestHandlerExtra<any, any>>();
+const requestHandlerExtraAsyncLocalStorage = new AsyncContext.Variable<
+  RequestHandlerExtra<any, any>
+>();
 
 export const getReqHandlerExtra = () => {
-  const extra = requestHandlerExtraAsyncLocalStorage.getStore();
+  const extra = requestHandlerExtraAsyncLocalStorage.get();
   if (!extra) {
     throw new Error("Request handler extra not set");
   }

pnpm-lock.yaml

@@ -190,6 +190,10 @@ importers:
   docs: {}
 
   packages/core:
+    dependencies:
+      simple-async-context:
+        specifier: ^1.0.4
+        version: 1.0.4
     devDependencies:
       '@modelcontextprotocol/sdk':
         specifier: ^1.13.1
@@ -3618,6 +3622,9 @@ packages:
     resolution: {integrity: sha512-bzyZ1e88w9O1iNJbKnOlvYTrWPDl46O1bG0D3XInv+9tkPrxrN8jUUTiFlDkkmKWgn1M6CfIA13SuGqOa9Korw==}
     engines: {node: '>=14'}
 
+  simple-async-context@1.0.4:
+    resolution: {integrity: sha512-OPH9mLzjbamUjRYDf2Xk/KA72nbhzraJrc9NZl0eLC5/npJR2dpC4tGAXrPbZw7WPjoAlv45N27t9uL6WtwV+A==}
+
   simple-swizzle@0.2.2:
     resolution: {integrity: sha512-JA//kQgZtbuY83m+xT+tXJkmJncGMTFT+C+g2h2R9uxkYIrE2yy9sgmcLhCnw57/WSD+Eh3J97FPEDFnbXnDUg==}
 
@@ -7392,6 +7399,8 @@ snapshots:
 
   signal-exit@4.1.0: {}
 
+  simple-async-context@1.0.4: {}
+
   simple-swizzle@0.2.2:
     dependencies:
       is-arrayish: 0.3.2