@b9g/async-context 0.1.1 → 0.1.2

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@b9g/async-context",
-  "version": "0.1.1",
+  "version": "0.1.2",
   "description": "Lightweight AsyncContext polyfill for JavaScript runtimes. Implements TC39 AsyncContext proposal using AsyncLocalStorage.",
   "keywords": [
     "asynccontext",
@@ -16,7 +16,7 @@
   ],
   "dependencies": {},
   "devDependencies": {
-    "@b9g/libuild": "^0.1.11",
+    "@b9g/libuild": "^0.1.18",
     "@types/node": "^20.0.0",
     "bun-types": "latest"
   },

package/src/index.d.ts

@@ -7,6 +7,7 @@
  * This implementation uses Node.js AsyncLocalStorage under the hood
  * to provide async context propagation across promise chains and async callbacks.
  */
+import { AsyncLocalStorage } from "node:async_hooks";
 /**
  * Options for creating an AsyncContext.Variable
  */
@@ -27,7 +28,7 @@ export interface AsyncVariableOptions<T> {
  *
  * @example
  * ```ts
- * const userContext = new AsyncVariable<User>();
+ * const userContext = new AsyncContext.Variable<User>();
  *
  * userContext.run(currentUser, async () => {
  *   await someAsyncOperation();
@@ -44,9 +45,10 @@ export declare class AsyncVariable<T> {
      *
      * @param value - The context value to set
      * @param fn - The function to execute with the context
+     * @param args - Additional arguments to pass to fn
      * @returns The return value of fn
      */
-    run<R>(value: T, fn: () => R): R;
+    run<R, Args extends unknown[]>(value: T, fn: (...args: Args) => R, ...args: Args): R;
     /**
      * Get the current context value
      * Returns the default value if no context is set
@@ -65,15 +67,78 @@ export declare class AsyncVariable<T> {
      * Get the name of this variable (for debugging)
      */
     get name(): string | undefined;
+    /**
+     * Internal: Get the underlying storage (used by Snapshot)
+     * @internal
+     */
+    _getStorage(): AsyncLocalStorage<T>;
 }
 /**
- * Namespace matching the TC39 AsyncContext proposal
+ * AsyncContext.Snapshot - captures the current values of all Variables
+ *
+ * A Snapshot captures the state of all AsyncContext.Variable instances at the
+ * time of construction. Later, calling `run()` restores that state for the
+ * duration of the callback.
+ *
+ * @example
+ * ```ts
+ * const userVar = new AsyncContext.Variable<string>();
+ *
+ * userVar.run("alice", () => {
+ *   const snapshot = new AsyncContext.Snapshot();
+ *
+ *   // Later, in a different context...
+ *   userVar.run("bob", () => {
+ *     console.log(userVar.get()); // "bob"
+ *
+ *     snapshot.run(() => {
+ *       console.log(userVar.get()); // "alice"
+ *     });
+ *   });
+ * });
+ * ```
  */
-export declare namespace AsyncContext {
+export declare class AsyncSnapshot {
+    #private;
+    constructor();
     /**
-     * AsyncContext.Variable - stores a value that propagates through async operations
+     * Execute a function with the captured context values
+     *
+     * @param fn - The function to execute
+     * @param args - Additional arguments to pass to fn
+     * @returns The return value of fn
      */
-    class Variable<T> extends AsyncVariable<T> {
-    }
+    run<R, Args extends unknown[]>(fn: (...args: Args) => R, ...args: Args): R;
+    /**
+     * Wrap a function to capture the current context
+     *
+     * Creates a new function that, when called, will execute with the
+     * context values that were active when wrap() was called.
+     *
+     * @param fn - The function to wrap
+     * @returns A wrapped function that preserves context
+     *
+     * @example
+     * ```ts
+     * const userVar = new AsyncContext.Variable<string>();
+     *
+     * const wrappedFn = userVar.run("alice", () => {
+     *   return AsyncContext.Snapshot.wrap(() => {
+     *     return userVar.get();
+     *   });
+     * });
+     *
+     * // Later, even outside the run() context:
+     * wrappedFn(); // returns "alice"
+     * ```
+     */
+    static wrap<T, A extends unknown[], R>(fn: (this: T, ...args: A) => R): (this: T, ...args: A) => R;
 }
+/**
+ * AsyncContext object matching the TC39 AsyncContext proposal
+ */
+export declare const AsyncContext: {
+    Variable: typeof AsyncVariable;
+    Snapshot: typeof AsyncSnapshot;
+};
 export default AsyncContext;

package/src/index.js

@@ -1,6 +1,8 @@
 /// <reference types="./index.d.ts" />
 // src/index.ts
 import { AsyncLocalStorage } from "node:async_hooks";
+var variableRegistry = /* @__PURE__ */ new Set();
+var NO_VALUE = Symbol("NO_VALUE");
 var AsyncVariable = class {
   #storage;
   #defaultValue;
@@ -9,6 +11,7 @@ var AsyncVariable = class {
     this.#storage = new AsyncLocalStorage();
     this.#defaultValue = options?.defaultValue;
     this.#name = options?.name;
+    variableRegistry.add(this);
   }
   /**
    * Execute a function with a context value
@@ -16,10 +19,11 @@ var AsyncVariable = class {
    *
    * @param value - The context value to set
    * @param fn - The function to execute with the context
+   * @param args - Additional arguments to pass to fn
    * @returns The return value of fn
    */
-  run(value, fn) {
-    return this.#storage.run(value, fn);
+  run(value, fn, ...args) {
+    return this.#storage.run(value, fn, ...args);
   }
   /**
    * Get the current context value
@@ -46,16 +50,77 @@ var AsyncVariable = class {
   get name() {
     return this.#name;
   }
+  /**
+   * Internal: Get the underlying storage (used by Snapshot)
+   * @internal
+   */
+  _getStorage() {
+    return this.#storage;
+  }
 };
-var AsyncContext;
-((AsyncContext2) => {
-  class Variable extends AsyncVariable {
+var AsyncSnapshot = class _AsyncSnapshot {
+  #captured;
+  constructor() {
+    this.#captured = /* @__PURE__ */ new Map();
+    for (const variable of variableRegistry) {
+      const value = variable.getStore();
+      this.#captured.set(variable, value !== void 0 ? value : NO_VALUE);
+    }
+  }
+  /**
+   * Execute a function with the captured context values
+   *
+   * @param fn - The function to execute
+   * @param args - Additional arguments to pass to fn
+   * @returns The return value of fn
+   */
+  run(fn, ...args) {
+    let result = () => fn(...args);
+    for (const [variable, value] of this.#captured) {
+      const prev = result;
+      const actualValue = value === NO_VALUE ? void 0 : value;
+      result = () => variable._getStorage().run(actualValue, prev);
+    }
+    return result();
+  }
+  /**
+   * Wrap a function to capture the current context
+   *
+   * Creates a new function that, when called, will execute with the
+   * context values that were active when wrap() was called.
+   *
+   * @param fn - The function to wrap
+   * @returns A wrapped function that preserves context
+   *
+   * @example
+   * ```ts
+   * const userVar = new AsyncContext.Variable<string>();
+   *
+   * const wrappedFn = userVar.run("alice", () => {
+   *   return AsyncContext.Snapshot.wrap(() => {
+   *     return userVar.get();
+   *   });
+   * });
+   *
+   * // Later, even outside the run() context:
+   * wrappedFn(); // returns "alice"
+   * ```
+   */
+  static wrap(fn) {
+    const snapshot = new _AsyncSnapshot();
+    return function(...args) {
+      return snapshot.run(() => fn.apply(this, args));
+    };
   }
-  AsyncContext2.Variable = Variable;
-})(AsyncContext || (AsyncContext = {}));
+};
+var AsyncContext = {
+  Variable: AsyncVariable,
+  Snapshot: AsyncSnapshot
+};
 var src_default = AsyncContext;
 export {
   AsyncContext,
+  AsyncSnapshot,
   AsyncVariable,
   src_default as default
 };