@b9g/async-context 0.1.1 → 0.1.3

package/README.md

@@ -12,7 +12,7 @@ The TC39 AsyncContext proposal aims to standardize async context propagation in
 
 This package provides a **lightweight, maintainable polyfill** that:
 
-✅ Implements the TC39 `AsyncContext.Variable` API
+✅ Implements the TC39 `AsyncContext.Variable` and `AsyncContext.Snapshot` APIs
 ✅ Uses battle-tested `AsyncLocalStorage` under the hood
 ✅ Zero dependencies (beyond Node.js built-ins)
 ✅ Full TypeScript support
@@ -121,19 +121,17 @@ console.log(themeContext.get()); // "light"
 ### Classes
 
 - `AsyncVariable<T>` - Main class for creating async context variables
-- `AsyncContext.Variable<T>` - Alias matching TC39 proposal namespace
+- `AsyncSnapshot` - Captures and restores context state
+- `AsyncContext.Variable<T>` - Alias for AsyncVariable (TC39 API)
+- `AsyncContext.Snapshot` - Alias for AsyncSnapshot (TC39 API)
 
 ### Types
 
 - `AsyncVariableOptions<T>` - Options for AsyncVariable constructor (defaultValue, name)
 
-### Namespaces
-
-- `AsyncContext` - Namespace containing Variable class (TC39 API)
-
 ### Default Export
 
-- `AsyncContext` - The AsyncContext namespace
+- `AsyncContext` - Object containing Variable and Snapshot classes
 
 ## API
 
@@ -147,13 +145,14 @@ Options:
 - `defaultValue?: T` - Default value when no context is set
 - `name?: string` - Optional name for debugging
 
-#### `run<R>(value: T, fn: () => R): R`
+#### `run<R>(value: T, fn: (...args) => R, ...args): R`
 
 Execute a function with a context value. The value is available via `get()` throughout the entire async execution of `fn`.
 
 **Parameters:**
 - `value: T` - The context value to set
-- `fn: () => R` - Function to execute (can be sync or async)
+- `fn: (...args) => R` - Function to execute (can be sync or async)
+- `...args` - Additional arguments to pass to fn
 
 **Returns:** The return value of `fn`
 
@@ -165,14 +164,39 @@ Get the current context value. Returns `defaultValue` if no context is set.
 
 Get the name of this variable (for debugging).
 
-### `AsyncContext.Variable<T>`
+### `AsyncSnapshot`
+
+Captures the current values of all Variables at construction time. Use `run()` to restore that state later.
+
+#### `constructor()`
+
+Creates a snapshot of all current Variable values.
 
-Alias for `AsyncVariable<T>` that matches the TC39 proposal namespace.
+#### `run<R>(fn: (...args) => R, ...args): R`
+
+Execute a function with the captured context values restored.
+
+**Parameters:**
+- `fn: (...args) => R` - Function to execute
+- `...args` - Additional arguments to pass to fn
+
+**Returns:** The return value of `fn`
+
+#### `static wrap<F>(fn: F): F`
+
+Wrap a function to preserve the current context. When the wrapped function is called later, it will execute with the context values that were active when `wrap()` was called.
 
 ```typescript
-import { AsyncContext } from "@b9g/async-context";
+const userVar = new AsyncContext.Variable<string>();
 
-const ctx = 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"
 ```
 
 ## How It Works
@@ -210,18 +234,14 @@ This package works in any JavaScript runtime that supports `AsyncLocalStorage`:
 
 ## Differences from TC39 Proposal
 
-This polyfill currently implements:
-
-- ✅ `AsyncContext.Variable`
-- ✅ `.run(value, fn)` method
-- ✅ `.get()` method
+This polyfill implements the core TC39 AsyncContext API:
 
-Not yet implemented (future additions):
+- ✅ `AsyncContext.Variable` - context variables with `run()` and `get()`
+- ✅ `AsyncContext.Snapshot` - context capture with `run()` and `wrap()`
 
-- ⏳ `AsyncContext.Snapshot`
-- ⏳ `AsyncContext.Mapping`
+The implementation uses Node.js `AsyncLocalStorage` rather than the pure-JS reference implementation, which means async context propagation works natively without monkey-patching `Promise.prototype.then`.
 
-These may be added in future versions as the proposal evolves.
+Test suite adapted from the [TC39 proposal repository](https://github.com/tc39/proposal-async-context).
 
 ## Migration Path
 

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@b9g/async-context",
-  "version": "0.1.1",
+  "version": "0.1.3",
   "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
 };