Extend runtime-context to support key-based multi-context storage and retrieval

Author: oamm

package.json

@@ -1,6 +1,6 @@
 {
   "name": "@oamm/runtime-context",
-  "version": "1.0.1",
+  "version": "1.0.2",
   "description": "Shared, injectable request-scoped runtime context for Node.js and Edge runtimes",
   "type": "module",
   "main": "./dist/index.js",

src/index.ts

@@ -39,29 +39,71 @@ function getStorage<TContext>(): ContextStorage<TContext> {
   return storage!;
 }
 
-export function getContext<TContext = DefaultContext>(): TContext | undefined {
-  return getStorage<TContext>().getStore();
+export function getContext<TContext = DefaultContext>(key?: any): TContext | undefined {
+  const store = getStorage<any>().getStore();
+  if (key !== undefined) {
+    if (store instanceof Map) {
+      return store.get(key);
+    }
+    if (typeof store === 'object' && store !== null) {
+      return (store as any)[key];
+    }
+    return undefined;
+  }
+  return store as TContext | undefined;
 }
 
-export function requireContext<TContext = DefaultContext>(): TContext {
-  const ctx = getContext<TContext>();
+export function requireContext<TContext = DefaultContext>(key?: any): TContext {
+  const ctx = getContext<TContext>(key);
   if (ctx === undefined) {
-    throw new Error('Runtime context is missing. Ensure you are running within runWithContext.');
+    throw new Error(
+      `Runtime context${key ? ` for ${String(key)}` : ''} is missing. Ensure you are running within runWithContext.`
+    );
   }
   return ctx;
 }
 
 export function runWithContext<TContext, R>(
   ctx: TContext,
   fn: () => R
-): R {
-  return getStorage<TContext>().run(ctx, fn);
+): R;
+export function runWithContext<TContext, R>(
+  key: any,
+  ctx: TContext,
+  fn: () => R
+): R;
+export function runWithContext<TContext, R>(...args: any[]): R {
+  const storage = getStorage();
+  if (args.length === 3) {
+    const [key, ctx, fn] = args;
+    const parent = storage.getStore();
+    const map = parent instanceof Map ? new Map(parent) : new Map();
+    map.set(key, ctx);
+    return storage.run(map, fn);
+  }
+  const [ctx, fn] = args;
+  return storage.run(ctx, fn);
 }
 
 export function ensureContext<TContext, R>(
   create: () => TContext,
   fn: () => Promise<R> | R
-): Promise<R> | R {
+): Promise<R> | R;
+export function ensureContext<TContext, R>(
+  key: any,
+  create: () => TContext,
+  fn: () => Promise<R> | R
+): Promise<R> | R;
+export function ensureContext<TContext, R>(...args: any[]): Promise<R> | R {
+  if (args.length === 3) {
+    const [key, create, fn] = args;
+    const existing = getContext<TContext>(key);
+    if (existing !== undefined) {
+      return fn();
+    }
+    return runWithContext(key, create(), fn);
+  }
+  const [create, fn] = args;
   const existing = getContext<TContext>();
   if (existing !== undefined) {
     return fn();
@@ -73,30 +115,48 @@ export function ensureContext<TContext, R>(
 
 export function setValue<TContext extends object = DefaultContext, K extends keyof TContext = keyof TContext>(
   key: K,
-  value: TContext[K]
+  value: TContext[K],
+  contextKey?: any
 ): void {
-  const ctx = requireContext<TContext>();
+  const ctx = requireContext<any>(contextKey);
+  if (ctx instanceof Map) {
+    ctx.set(key, value);
+    return;
+  }
   if (typeof ctx !== 'object' || ctx === null) {
-    throw new Error('Context must be an object to use setValue.');
+    throw new Error('Context must be an object or Map to use setValue.');
   }
   (ctx as any)[key] = value;
 }
 
 export function getValue<TContext extends object = DefaultContext, K extends keyof TContext = keyof TContext>(
-  key: K
+  key: K,
+  contextKey?: any
 ): TContext[K] | undefined {
-  const ctx = getContext<TContext>();
+  const ctx = getContext<any>(contextKey);
   if (!ctx) return undefined;
+  if (ctx instanceof Map) {
+    return ctx.get(key);
+  }
   if (typeof ctx !== 'object') {
-    throw new Error('Context must be an object to use getValue.');
+    throw new Error('Context must be an object or Map to use getValue.');
   }
   return (ctx as any)[key];
 }
 
-export function mergeContext<TContext extends object = DefaultContext>(partial: Partial<TContext>): void {
-  const ctx = requireContext<TContext>();
+export function mergeContext<TContext extends object = DefaultContext>(
+  partial: Partial<TContext>,
+  contextKey?: any
+): void {
+  const ctx = requireContext<any>(contextKey);
+  if (ctx instanceof Map) {
+    for (const [key, value] of Object.entries(partial)) {
+      ctx.set(key, value);
+    }
+    return;
+  }
   if (typeof ctx !== 'object' || ctx === null) {
-    throw new Error('Context must be an object to use mergeContext.');
+    throw new Error('Context must be an object or Map to use mergeContext.');
   }
   Object.assign(ctx, partial);
 }

tests/context.test.ts

@@ -120,4 +120,104 @@ describe('Runtime Context', () => {
     expect(() => runWithContext({}, () => {})).toThrow(/Runtime context storage is missing/);
     vi.restoreAllMocks();
   });
+
+  describe('Map-based Multi-Context', () => {
+    it('should support multiple contexts using keys', () => {
+      const tokenKitData = { token: 'abc' };
+      const sessionData = { user: 'admin' };
+
+      runWithContext('token-kit', tokenKitData, () => {
+        expect(getContext('token-kit')).toBe(tokenKitData);
+        expect(requireContext('token-kit')).toBe(tokenKitData);
+
+        runWithContext('session', sessionData, () => {
+          expect(getContext('token-kit')).toBe(tokenKitData);
+          expect(getContext('session')).toBe(sessionData);
+
+          // Without key, it should return the Map (the whole store)
+          const store = getContext();
+          expect(store).toBeInstanceOf(Map);
+          expect((store as any as Map<any, any>).get('token-kit')).toBe(tokenKitData);
+          expect((store as any as Map<any, any>).get('session')).toBe(sessionData);
+        });
+
+        expect(getContext('session')).toBeUndefined();
+        expect(getContext('token-kit')).toBe(tokenKitData);
+      });
+    });
+
+    it('should support classes as keys', () => {
+      class TokenKit {
+        constructor(public token: string) {}
+      }
+      class Session {
+        constructor(public user: string) {}
+      }
+
+      const tk = new TokenKit('abc');
+      const s = new Session('admin');
+
+      runWithContext(TokenKit, tk, () => {
+        runWithContext(Session, s, () => {
+          expect(getContext(TokenKit)).toBe(tk);
+          expect(getContext(Session)).toBe(s);
+        });
+      });
+    });
+
+    it('should support ensureContext with keys', async () => {
+      const key = Symbol('key');
+      const data1 = { a: 1 };
+      const data2 = { a: 2 };
+
+      const result = await ensureContext(
+        key,
+        () => data1,
+        async () => {
+          expect(getContext(key)).toBe(data1);
+
+          return await ensureContext(key, () => data2, () => {
+            expect(getContext(key)).toBe(data1); // Should reuse existing
+            return 'done';
+          });
+        }
+      );
+      expect(result).toBe('done');
+    });
+
+    it('should support setValue/getValue/mergeContext with context keys', () => {
+      const key = 'my-context';
+      const data = { a: 1 } as any;
+
+      runWithContext(key, data, () => {
+        setValue('b', 2, key);
+        expect(getValue('b', key)).toBe(2);
+        expect(data.b).toBe(2);
+
+        mergeContext({ c: 3 }, key);
+        expect(data.c).toBe(3);
+      });
+    });
+
+    it('should support setValue/getValue/mergeContext on Map context', () => {
+      const key = 'my-map';
+      const data = new Map();
+
+      runWithContext(key, data, () => {
+        setValue('foo', 'bar', key);
+        expect(getValue('foo', key)).toBe('bar');
+        expect(data.get('foo')).toBe('bar');
+
+        mergeContext({ hello: 'world' }, key);
+        expect(data.get('hello')).toBe('world');
+      });
+    });
+
+    it('should fall back to object property if not a Map', () => {
+      const ctx = { foo: 'bar' };
+      runWithContext(ctx, () => {
+        expect(getContext('foo')).toBe('bar');
+      });
+    });
+  });
 });