From 617b050d587e05ee6d259bd204b95672f803bc91 Mon Sep 17 00:00:00 2001
From: Pavel Jbanov <pavelj@google.com>
Date: Tue, 16 Jul 2024 19:29:44 -0400
Subject: [PATCH 01/12] feat: made registry hierarchical, allowing ephemeral
 child registries for testing

---
 js/core/src/registry.ts        | 340 +++++++++++++++++++++------------
 js/core/tests/registry_test.ts | 217 ++++++++++++++++++++-
 2 files changed, 430 insertions(+), 127 deletions(-)

diff --git a/js/core/src/registry.ts b/js/core/src/registry.ts
index b46b636baf..830278bbd3 100644
--- a/js/core/src/registry.ts
+++ b/js/core/src/registry.ts
@@ -14,6 +14,7 @@
  * limitations under the License.
  */
 
+import { AsyncLocalStorage } from 'node:async_hooks';
 import * as z from 'zod';
 import { Action } from './action.js';
 import { FlowStateStore } from './flowTypes.js';
@@ -25,45 +26,7 @@ import { TraceStore } from './tracing/types.js';
 
 export type AsyncProvider<T> = () => Promise<T>;
 
-const ACTIONS_BY_ID = 'genkit__ACTIONS_BY_ID';
-const TRACE_STORES_BY_ENV = 'genkit__TRACE_STORES_BY_ENV';
-const FLOW_STATE_STORES_BY_ENV = 'genkit__FLOW_STATE_STORES_BY_ENV';
-const PLUGINS_BY_NAME = 'genkit__PLUGINS_BY_NAME';
-const SCHEMAS_BY_NAME = 'genkit__SCHEMAS_BY_NAME';
-
-function actionsById(): Record<string, Action<z.ZodTypeAny, z.ZodTypeAny>> {
-  if (global[ACTIONS_BY_ID] === undefined) {
-    global[ACTIONS_BY_ID] = {};
-  }
-  return global[ACTIONS_BY_ID];
-}
-function traceStoresByEnv(): Record<string, AsyncProvider<TraceStore>> {
-  if (global[TRACE_STORES_BY_ENV] === undefined) {
-    global[TRACE_STORES_BY_ENV] = {};
-  }
-  return global[TRACE_STORES_BY_ENV];
-}
-function flowStateStoresByEnv(): Record<string, AsyncProvider<FlowStateStore>> {
-  if (global[FLOW_STATE_STORES_BY_ENV] === undefined) {
-    global[FLOW_STATE_STORES_BY_ENV] = {};
-  }
-  return global[FLOW_STATE_STORES_BY_ENV];
-}
-function pluginsByName(): Record<string, PluginProvider> {
-  if (global[PLUGINS_BY_NAME] === undefined) {
-    global[PLUGINS_BY_NAME] = {};
-  }
-  return global[PLUGINS_BY_NAME];
-}
-function schemasByName(): Record<
-  string,
-  { schema?: z.ZodTypeAny; jsonSchema?: JSONSchema }
-> {
-  if (global[SCHEMAS_BY_NAME] === undefined) {
-    global[SCHEMAS_BY_NAME] = {};
-  }
-  return global[SCHEMAS_BY_NAME];
-}
+const REGISTRY_KEY = 'genkit__REGISTRY';
 
 /**
  * Type of a runnable action.
@@ -82,17 +45,12 @@ export type ActionType =
 /**
  * Looks up a registry key (action type and key) in the registry.
  */
-export async function lookupAction<
+export function lookupAction<
   I extends z.ZodTypeAny,
   O extends z.ZodTypeAny,
   R extends Action<I, O>,
 >(key: string): Promise<R> {
-  // If we don't see the key in the registry we try to initialize the plugin first.
-  const pluginName = parsePluginName(key);
-  if (!actionsById()[key] && pluginName) {
-    await initializePlugin(pluginName);
-  }
-  return actionsById()[key] as R;
+  return getRegistryInstance().lookupAction(key);
 }
 
 function parsePluginName(registryKey: string) {
@@ -110,14 +68,7 @@ export function registerAction<I extends z.ZodTypeAny, O extends z.ZodTypeAny>(
   type: ActionType,
   action: Action<I, O>
 ) {
-  logger.info(`Registering ${type}: ${action.__action.name}`);
-  const key = `/${type}/${action.__action.name}`;
-  if (actionsById().hasOwnProperty(key)) {
-    logger.warn(
-      `WARNING: ${key} already has an entry in the registry. Overwriting.`
-    );
-  }
-  actionsById()[key] = action;
+  return getRegistryInstance().registerAction(type, action);
 }
 
 type ActionsRecord = Record<string, Action<z.ZodTypeAny, z.ZodTypeAny>>;
@@ -125,20 +76,8 @@ type ActionsRecord = Record<string, Action<z.ZodTypeAny, z.ZodTypeAny>>;
 /**
  * Returns all actions in the registry.
  */
-export async function listActions(): Promise<ActionsRecord> {
-  await initializeAllPlugins();
-  return Object.assign({}, actionsById());
-}
-
-let allPluginsInitialized = false;
-export async function initializeAllPlugins() {
-  if (allPluginsInitialized) {
-    return;
-  }
-  for (const pluginName of Object.keys(pluginsByName())) {
-    await initializePlugin(pluginName);
-  }
-  allPluginsInitialized = true;
+export function listActions(): Promise<ActionsRecord> {
+  return getRegistryInstance().listActions();
 }
 
 /**
@@ -148,27 +87,14 @@ export function registerTraceStore(
   env: string,
   traceStoreProvider: AsyncProvider<TraceStore>
 ) {
-  traceStoresByEnv()[env] = traceStoreProvider;
+  return getRegistryInstance().registerTraceStore(env, traceStoreProvider);
 }
 
-const traceStoresByEnvCache: Record<any, Promise<TraceStore>> = {};
-
 /**
  * Looks up the trace store for the given environment.
  */
-export async function lookupTraceStore(
-  env: string
-): Promise<TraceStore | undefined> {
-  if (!traceStoresByEnv()[env]) {
-    return undefined;
-  }
-  const cached = traceStoresByEnvCache[env];
-  if (!cached) {
-    const newStore = traceStoresByEnv()[env]();
-    traceStoresByEnvCache[env] = newStore;
-    return newStore;
-  }
-  return cached;
+export function lookupTraceStore(env: string): Promise<TraceStore | undefined> {
+  return getRegistryInstance().lookupTraceStore(env);
 }
 
 /**
@@ -178,71 +104,48 @@ export function registerFlowStateStore(
   env: string,
   flowStateStoreProvider: AsyncProvider<FlowStateStore>
 ) {
-  flowStateStoresByEnv()[env] = flowStateStoreProvider;
+  return getRegistryInstance().registerFlowStateStore(
+    env,
+    flowStateStoreProvider
+  );
 }
 
-const flowStateStoresByEnvCache: Record<any, Promise<FlowStateStore>> = {};
 /**
  * Looks up the flow state store for the given environment.
  */
 export async function lookupFlowStateStore(
   env: string
 ): Promise<FlowStateStore | undefined> {
-  if (!flowStateStoresByEnv()[env]) {
-    return undefined;
-  }
-  const cached = flowStateStoresByEnvCache[env];
-  if (!cached) {
-    const newStore = flowStateStoresByEnv()[env]();
-    flowStateStoresByEnvCache[env] = newStore;
-    return newStore;
-  }
-  return cached;
+  return getRegistryInstance().lookupFlowStateStore(env);
 }
 
 /**
  * Registers a flow state store for the given environment.
  */
 export function registerPluginProvider(name: string, provider: PluginProvider) {
-  allPluginsInitialized = false;
-  let cached;
-  let isInitialized = false;
-  pluginsByName()[name] = {
-    name: provider.name,
-    initializer: () => {
-      if (isInitialized) {
-        return cached;
-      }
-      cached = provider.initializer();
-      isInitialized = true;
-      return cached;
-    },
-  };
+  return getRegistryInstance().registerPluginProvider(name, provider);
 }
 
 export function lookupPlugin(name: string) {
-  return pluginsByName()[name];
+  return getRegistryInstance().lookupFlowStateStore(name);
 }
 
 /**
- *
+ * Initialize plugin -- calls the plugin initialization function.
  */
 export async function initializePlugin(name: string) {
-  if (pluginsByName()[name]) {
-    return await pluginsByName()[name].initializer();
-  }
-  return undefined;
+  return getRegistryInstance().initializePlugin(name);
 }
 
 export function registerSchema(
   name: string,
   data: { schema?: z.ZodTypeAny; jsonSchema?: JSONSchema }
 ) {
-  schemasByName()[name] = data;
+  return getRegistryInstance().registerSchema(name, data);
 }
 
 export function lookupSchema(name: string) {
-  return schemasByName()[name];
+  return getRegistryInstance().lookupSchema(name);
 }
 
 /**
@@ -253,14 +156,199 @@ if (process.env.GENKIT_ENV === 'dev') {
 }
 
 export function __hardResetRegistryForTesting() {
-  delete global[ACTIONS_BY_ID];
-  delete global[TRACE_STORES_BY_ENV];
-  delete global[FLOW_STATE_STORES_BY_ENV];
-  delete global[PLUGINS_BY_NAME];
-  deleteAll(flowStateStoresByEnvCache);
-  deleteAll(traceStoresByEnvCache);
+  delete global[REGISTRY_KEY];
+  global[REGISTRY_KEY] = new Registry();
 }
 
-function deleteAll(map: Record<any, any>) {
-  Object.keys(map).forEach((key) => delete map[key]);
+export class Registry {
+  private actionsById: Record<string, Action<z.ZodTypeAny, z.ZodTypeAny>> = {};
+  private traceStoresByEnv: Record<string, AsyncProvider<TraceStore>> = {};
+  private flowStateStoresByEnv: Record<string, AsyncProvider<FlowStateStore>> =
+    {};
+  private pluginsByName: Record<string, PluginProvider> = {};
+  private schemasByName: Record<
+    string,
+    { schema?: z.ZodTypeAny; jsonSchema?: JSONSchema }
+  > = {};
+
+  private traceStoresByEnvCache: Record<any, Promise<TraceStore>> = {};
+  private flowStateStoresByEnvCache: Record<any, Promise<FlowStateStore>> = {};
+  private allPluginsInitialized = false;
+
+  constructor(public parent?: Registry) {}
+
+  static withParent(parent: Registry) {
+    return new Registry(parent);
+  }
+
+  async lookupAction<
+    I extends z.ZodTypeAny,
+    O extends z.ZodTypeAny,
+    R extends Action<I, O>,
+  >(key: string): Promise<R> {
+    return (await this._lookupAction(key)) || this.parent?.lookupAction(key);
+  }
+
+  private async _lookupAction<
+    I extends z.ZodTypeAny,
+    O extends z.ZodTypeAny,
+    R extends Action<I, O>,
+  >(key: string): Promise<R> {
+    // If we don't see the key in the registry we try to initialize the plugin first.
+    const pluginName = parsePluginName(key);
+    if (!this.actionsById[key] && pluginName) {
+      await this.initializePlugin(pluginName);
+    }
+    return this.actionsById[key] as R;
+  }
+
+  registerAction<I extends z.ZodTypeAny, O extends z.ZodTypeAny>(
+    type: ActionType,
+    action: Action<I, O>
+  ) {
+    logger.info(`Registering ${type}: ${action.__action.name}`);
+    const key = `/${type}/${action.__action.name}`;
+    if (this.actionsById.hasOwnProperty(key)) {
+      logger.warn(
+        `WARNING: ${key} already has an entry in the registry. Overwriting.`
+      );
+    }
+    this.actionsById[key] = action;
+  }
+
+  async listActions(): Promise<ActionsRecord> {
+    await this.initializeAllPlugins();
+    return {
+      ...(await this.parent?.listActions()),
+      ...this.actionsById,
+    };
+  }
+
+  async initializeAllPlugins() {
+    if (this.allPluginsInitialized) {
+      return;
+    }
+    for (const pluginName of Object.keys(this.pluginsByName)) {
+      await initializePlugin(pluginName);
+    }
+    this.allPluginsInitialized = true;
+  }
+  
+
+  registerTraceStore(
+    env: string,
+    traceStoreProvider: AsyncProvider<TraceStore>
+  ) {
+    this.traceStoresByEnv[env] = traceStoreProvider;
+  }
+
+  async lookupTraceStore(env: string): Promise<TraceStore | undefined> {
+    return (
+      (await this._lookupTraceStore(env)) || this.parent?.lookupTraceStore(env)
+    );
+  }
+
+  private async _lookupTraceStore(
+    env: string
+  ): Promise<TraceStore | undefined> {
+    if (!this.traceStoresByEnv[env]) {
+      return undefined;
+    }
+    const cached = this.traceStoresByEnvCache[env];
+    if (!cached) {
+      const newStore = this.traceStoresByEnv[env]();
+      this.traceStoresByEnvCache[env] = newStore;
+      return newStore;
+    }
+    return cached;
+  }
+
+  registerFlowStateStore(
+    env: string,
+    flowStateStoreProvider: AsyncProvider<FlowStateStore>
+  ) {
+    this.flowStateStoresByEnv[env] = flowStateStoreProvider;
+  }
+
+  async lookupFlowStateStore(env: string): Promise<FlowStateStore | undefined> {
+    return (
+      (await this._lookupFlowStateStore(env)) ||
+      this.parent?.lookupFlowStateStore(env)
+    );
+  }
+
+  private async _lookupFlowStateStore(
+    env: string
+  ): Promise<FlowStateStore | undefined> {
+    if (!this.flowStateStoresByEnv[env]) {
+      return undefined;
+    }
+    const cached = this.flowStateStoresByEnvCache[env];
+    if (!cached) {
+      const newStore = this.flowStateStoresByEnv[env]();
+      this.flowStateStoresByEnvCache[env] = newStore;
+      return newStore;
+    }
+    return cached;
+  }
+
+  registerPluginProvider(name: string, provider: PluginProvider) {
+    this.allPluginsInitialized = false;
+    let cached;
+    let isInitialized = false;
+    this.pluginsByName[name] = {
+      name: provider.name,
+      initializer: () => {
+        if (isInitialized) {
+          return cached;
+        }
+        cached = provider.initializer();
+        isInitialized = true;
+        return cached;
+      },
+    };
+    }
+
+  lookupPlugin(name: string) {
+    return this.pluginsByName[name] || this.parent?.lookupPlugin(name);
+  }
+
+  async initializePlugin(name: string) {
+    if (this.pluginsByName[name]) {
+      return await this.pluginsByName[name].initializer();
+    }
+    return undefined;
+  }
+
+  registerSchema(
+    name: string,
+    data: { schema?: z.ZodTypeAny; jsonSchema?: JSONSchema }
+  ) {
+    this.schemasByName[name] = data;
+  }
+
+  lookupSchema(name: string) {
+    return this.schemasByName[name] || this.parent?.lookupSchema(name);
+  }
+}
+
+// global regustry instance
+global[REGISTRY_KEY] = new Registry();
+
+const registryAls = new AsyncLocalStorage<Registry>();
+
+/**
+ * Executes provided function with within the provided registry.
+ */
+export function runWithRegistry<O>(registry: Registry, fn: () => O): O {
+  return registryAls.run(registry, fn);
+}
+
+/**
+ * Returns the current registry instance:
+ *  - if running within `runWithRegistry` then return that registry
+ *  - else return the global registry.
+ */
+export function getRegistryInstance(): Registry {
+  return registryAls.getStore() || global[REGISTRY_KEY];
 }
diff --git a/js/core/tests/registry_test.ts b/js/core/tests/registry_test.ts
index c969eba290..08fe05de58 100644
--- a/js/core/tests/registry_test.ts
+++ b/js/core/tests/registry_test.ts
@@ -18,6 +18,7 @@ import assert from 'node:assert';
 import { beforeEach, describe, it } from 'node:test';
 import { action } from '../src/action.js';
 import {
+  Registry,
   __hardResetRegistryForTesting,
   listActions,
   lookupAction,
@@ -25,7 +26,7 @@ import {
   registerPluginProvider,
 } from '../src/registry.js';
 
-describe('registry', () => {
+describe('global registry', () => {
   beforeEach(__hardResetRegistryForTesting);
 
   describe('listActions', () => {
@@ -169,3 +170,217 @@ describe('registry', () => {
     assert.strictEqual(await lookupAction('/model/foo/something'), undefined);
   });
 });
+
+describe('registry class', () => {
+  var registry: Registry;
+  beforeEach(() => {
+    registry = new Registry();
+  });
+
+  describe('listActions', () => {
+    it('returns all registered actions', async () => {
+      const fooSomethingAction = action(
+        { name: 'foo_something' },
+        async () => null
+      );
+      registry.registerAction('model', fooSomethingAction);
+      const barSomethingAction = action(
+        { name: 'bar_something' },
+        async () => null
+      );
+      registry.registerAction('model', barSomethingAction);
+
+      assert.deepEqual(await registry.listActions(), {
+        '/model/foo_something': fooSomethingAction,
+        '/model/bar_something': barSomethingAction,
+      });
+    });
+
+    it('returns all registered actions by plugins', async () => {
+      registry.registerPluginProvider('foo', {
+        name: 'foo',
+        async initializer() {
+          registry.registerAction('model', fooSomethingAction);
+          return {};
+        },
+      });
+      const fooSomethingAction = action(
+        {
+          name: {
+            pluginId: 'foo',
+            actionId: 'something',
+          },
+        },
+        async () => null
+      );
+      registry.registerPluginProvider('bar', {
+        name: 'bar',
+        async initializer() {
+          registry.registerAction('model', barSomethingAction);
+          return {};
+        },
+      });
+      const barSomethingAction = action(
+        {
+          name: {
+            pluginId: 'bar',
+            actionId: 'something',
+          },
+        },
+        async () => null
+      );
+
+      assert.deepEqual(await registry.listActions(), {
+        '/model/foo/something': fooSomethingAction,
+        '/model/bar/something': barSomethingAction,
+      });
+    });
+
+    it('returns all registered actions, including parent', async () => {
+      const child = Registry.withParent(registry);
+
+      const fooSomethingAction = action(
+        { name: 'foo_something' },
+        async () => null
+      );
+      registry.registerAction('model', fooSomethingAction);
+      const barSomethingAction = action(
+        { name: 'bar_something' },
+        async () => null
+      );
+      child.registerAction('model', barSomethingAction);
+
+      assert.deepEqual(await child.listActions(), {
+        '/model/foo_something': fooSomethingAction,
+        '/model/bar_something': barSomethingAction,
+      });
+      assert.deepEqual(await registry.listActions(), {
+        '/model/foo_something': fooSomethingAction,
+      });
+    });
+  });
+
+  describe('lookupAction', () => {
+    it('initializes plugin for action first', async () => {
+      let fooInitialized = false;
+      registry.registerPluginProvider('foo', {
+        name: 'foo',
+        async initializer() {
+          fooInitialized = true;
+          return {};
+        },
+      });
+      let barInitialized = false;
+      registry.registerPluginProvider('bar', {
+        name: 'bar',
+        async initializer() {
+          barInitialized = true;
+          return {};
+        },
+      });
+
+      await registry.lookupAction('/model/foo/something');
+
+      assert.strictEqual(fooInitialized, true);
+      assert.strictEqual(barInitialized, false);
+
+      await registry.lookupAction('/model/bar/something');
+
+      assert.strictEqual(fooInitialized, true);
+      assert.strictEqual(barInitialized, true);
+    });
+
+    it('returns registered action', async () => {
+      const fooSomethingAction = action(
+        { name: 'foo_something' },
+        async () => null
+      );
+      registry.registerAction('model', fooSomethingAction);
+      const barSomethingAction = action(
+        { name: 'bar_something' },
+        async () => null
+      );
+      registry.registerAction('model', barSomethingAction);
+
+      assert.strictEqual(
+        await registry.lookupAction('/model/foo_something'),
+        fooSomethingAction
+      );
+      assert.strictEqual(
+        await registry.lookupAction('/model/bar_something'),
+        barSomethingAction
+      );
+    });
+
+    it('returns action registered by plugin', async () => {
+      registry.registerPluginProvider('foo', {
+        name: 'foo',
+        async initializer() {
+          registry.registerAction('model', somethingAction);
+          return {};
+        },
+      });
+      const somethingAction = action(
+        {
+          name: {
+            pluginId: 'foo',
+            actionId: 'something',
+          },
+        },
+        async () => null
+      );
+
+      assert.strictEqual(
+        await registry.lookupAction('/model/foo/something'),
+        somethingAction
+      );
+    });
+
+    it('returns undefined for unknown action', async () => {
+      assert.strictEqual(
+        await registry.lookupAction('/model/foo/something'),
+        undefined
+      );
+    });
+
+    it('should lookup parent registry when child missing action', async () => {
+      const childRegistry = new Registry(registry);
+
+      const fooSomethingAction = action(
+        { name: 'foo_something' },
+        async () => null
+      );
+      registry.registerAction('model', fooSomethingAction);
+
+      assert.strictEqual(
+        await registry.lookupAction('/model/foo_something'),
+        fooSomethingAction
+      );
+      assert.strictEqual(
+        await childRegistry.lookupAction('/model/foo_something'),
+        fooSomethingAction
+      );
+    });
+
+    it('registration on the child registry should not modify parent', async () => {
+      const childRegistry = Registry.withParent(registry);
+
+      assert.strictEqual(childRegistry.parent, registry);
+
+      const fooSomethingAction = action(
+        { name: 'foo_something' },
+        async () => null
+      );
+      childRegistry.registerAction('model', fooSomethingAction);
+
+      assert.strictEqual(
+        await registry.lookupAction('/model/foo_something'),
+        undefined
+      );
+      assert.strictEqual(
+        await childRegistry.lookupAction('/model/foo_something'),
+        fooSomethingAction
+      );
+    });
+  });
+});

From 5a5a80cdf74b69ba63f09bbfc01e5d4d741be8f2 Mon Sep 17 00:00:00 2001
From: Pavel Jbanov <pavelj@google.com>
Date: Tue, 16 Jul 2024 19:38:55 -0400
Subject: [PATCH 02/12] runWithRegistry tests

---
 js/core/src/registry.ts        |  6 +++++-
 js/core/tests/registry_test.ts | 36 ++++++++++++++++++++++++++++++++--
 2 files changed, 39 insertions(+), 3 deletions(-)

diff --git a/js/core/src/registry.ts b/js/core/src/registry.ts
index 830278bbd3..cfaa18986e 100644
--- a/js/core/src/registry.ts
+++ b/js/core/src/registry.ts
@@ -177,6 +177,10 @@ export class Registry {
 
   constructor(public parent?: Registry) {}
 
+  static withCurrent() {
+    return new Registry(getRegistryInstance());
+  }
+
   static withParent(parent: Registry) {
     return new Registry(parent);
   }
@@ -349,6 +353,6 @@ export function runWithRegistry<O>(registry: Registry, fn: () => O): O {
  *  - if running within `runWithRegistry` then return that registry
  *  - else return the global registry.
  */
-export function getRegistryInstance(): Registry {
+function getRegistryInstance(): Registry {
   return registryAls.getStore() || global[REGISTRY_KEY];
 }
diff --git a/js/core/tests/registry_test.ts b/js/core/tests/registry_test.ts
index 08fe05de58..63402ac146 100644
--- a/js/core/tests/registry_test.ts
+++ b/js/core/tests/registry_test.ts
@@ -18,12 +18,13 @@ import assert from 'node:assert';
 import { beforeEach, describe, it } from 'node:test';
 import { action } from '../src/action.js';
 import {
-  Registry,
-  __hardResetRegistryForTesting,
   listActions,
   lookupAction,
   registerAction,
   registerPluginProvider,
+  Registry,
+  runWithRegistry,
+  __hardResetRegistryForTesting,
 } from '../src/registry.js';
 
 describe('global registry', () => {
@@ -383,4 +384,35 @@ describe('registry class', () => {
       );
     });
   });
+
+  describe('runWithRegistry', () => {
+    it('should lookup parent registry when child missing action', async () => {
+      const childRegistry = new Registry(registry);
+
+      const fooSomethingAction = action(
+        { name: 'foo_something' },
+        async () => null
+      );
+
+      runWithRegistry(childRegistry, () => {
+        registerAction('model', fooSomethingAction);
+      });
+
+      assert.strictEqual(
+        await registry.lookupAction('/model/foo_something'),
+        undefined
+      );
+      assert.strictEqual(
+        await childRegistry.lookupAction('/model/foo_something'),
+        fooSomethingAction
+      );
+      assert.strictEqual(await lookupAction('/model/foo_something'), undefined);
+      await runWithRegistry(childRegistry, async () => {
+        assert.strictEqual(
+          await lookupAction('/model/foo_something'),
+          fooSomethingAction
+        );
+      });
+    });
+  });
 });

From 9e81f90aa39ee60dfff5090da40e0d5dbcdc8699 Mon Sep 17 00:00:00 2001
From: Pavel Jbanov <pavelj@google.com>
Date: Tue, 16 Jul 2024 19:50:38 -0400
Subject: [PATCH 03/12] runInEphemeralRegistry

---
 js/core/src/registry.ts        | 8 ++++++++
 js/core/tests/registry_test.ts | 4 ++--
 2 files changed, 10 insertions(+), 2 deletions(-)

diff --git a/js/core/src/registry.ts b/js/core/src/registry.ts
index cfaa18986e..697aa7d196 100644
--- a/js/core/src/registry.ts
+++ b/js/core/src/registry.ts
@@ -348,6 +348,14 @@ export function runWithRegistry<O>(registry: Registry, fn: () => O): O {
   return registryAls.run(registry, fn);
 }
 
+/**
+ * Executes provided function with within a child/ephemeral registry.
+ */
+export function runInEphemeralRegistry<O>(fn: (registry: Registry) => O): O {
+  const registry = Registry.withCurrent();
+  return registryAls.run(registry, () => fn(registry));
+}
+
 /**
  * Returns the current registry instance:
  *  - if running within `runWithRegistry` then return that registry
diff --git a/js/core/tests/registry_test.ts b/js/core/tests/registry_test.ts
index 63402ac146..4855b365ed 100644
--- a/js/core/tests/registry_test.ts
+++ b/js/core/tests/registry_test.ts
@@ -18,13 +18,13 @@ import assert from 'node:assert';
 import { beforeEach, describe, it } from 'node:test';
 import { action } from '../src/action.js';
 import {
+  Registry,
+  __hardResetRegistryForTesting,
   listActions,
   lookupAction,
   registerAction,
   registerPluginProvider,
-  Registry,
   runWithRegistry,
-  __hardResetRegistryForTesting,
 } from '../src/registry.js';
 
 describe('global registry', () => {

From 5d202b8d0de49637ff5ab6c3d347f2ddbffe21f3 Mon Sep 17 00:00:00 2001
From: Pavel Jbanov <pavelj@google.com>
Date: Wed, 17 Jul 2024 13:23:34 -0400
Subject: [PATCH 04/12] feedback

---
 js/core/src/registry.ts        | 21 +++++---------
 js/core/tests/registry_test.ts | 50 +++++++++++++++++++++++++---------
 2 files changed, 44 insertions(+), 27 deletions(-)

diff --git a/js/core/src/registry.ts b/js/core/src/registry.ts
index 697aa7d196..cdf304f18c 100644
--- a/js/core/src/registry.ts
+++ b/js/core/src/registry.ts
@@ -189,21 +189,13 @@ export class Registry {
     I extends z.ZodTypeAny,
     O extends z.ZodTypeAny,
     R extends Action<I, O>,
-  >(key: string): Promise<R> {
-    return (await this._lookupAction(key)) || this.parent?.lookupAction(key);
-  }
-
-  private async _lookupAction<
-    I extends z.ZodTypeAny,
-    O extends z.ZodTypeAny,
-    R extends Action<I, O>,
   >(key: string): Promise<R> {
     // If we don't see the key in the registry we try to initialize the plugin first.
     const pluginName = parsePluginName(key);
     if (!this.actionsById[key] && pluginName) {
       await this.initializePlugin(pluginName);
     }
-    return this.actionsById[key] as R;
+    return (this.actionsById[key] as R) || this.parent?.lookupAction(key);
   }
 
   registerAction<I extends z.ZodTypeAny, O extends z.ZodTypeAny>(
@@ -342,16 +334,17 @@ global[REGISTRY_KEY] = new Registry();
 const registryAls = new AsyncLocalStorage<Registry>();
 
 /**
- * Executes provided function with within the provided registry.
+ * Executes provided function with within an isolated registry that has no visibility into global namespace.
  */
-export function runWithRegistry<O>(registry: Registry, fn: () => O): O {
-  return registryAls.run(registry, fn);
+export function runInIsolatedRegistry<O>(fn: (registry: Registry) => O): O {
+  const registry = new Registry();
+  return registryAls.run(registry, () => fn(registry));
 }
 
 /**
- * Executes provided function with within a child/ephemeral registry.
+ * Executes provided function with within a temporary registry that inherits from global/parent.
  */
-export function runInEphemeralRegistry<O>(fn: (registry: Registry) => O): O {
+export function runInTempRegistry<O>(fn: (registry: Registry) => O): O {
   const registry = Registry.withCurrent();
   return registryAls.run(registry, () => fn(registry));
 }
diff --git a/js/core/tests/registry_test.ts b/js/core/tests/registry_test.ts
index 4855b365ed..507afe42e4 100644
--- a/js/core/tests/registry_test.ts
+++ b/js/core/tests/registry_test.ts
@@ -24,7 +24,8 @@ import {
   lookupAction,
   registerAction,
   registerPluginProvider,
-  runWithRegistry,
+  runInIsolatedRegistry,
+  runInTempRegistry,
 } from '../src/registry.js';
 
 describe('global registry', () => {
@@ -385,34 +386,57 @@ describe('registry class', () => {
     });
   });
 
-  describe('runWithRegistry', () => {
+  describe('runInTempRegistry', () => {
     it('should lookup parent registry when child missing action', async () => {
-      const childRegistry = new Registry(registry);
-
       const fooSomethingAction = action(
         { name: 'foo_something' },
         async () => null
       );
+      const barSomethingAction = action(
+        { name: 'bar_something' },
+        async () => null
+      );
+      registerAction('model', barSomethingAction);
 
-      runWithRegistry(childRegistry, () => {
+      await runInTempRegistry(async () => {
         registerAction('model', fooSomethingAction);
+        assert.strictEqual(
+          await lookupAction('/model/bar_something'),
+          barSomethingAction
+        );
+        assert.strictEqual(
+          await lookupAction('/model/foo_something'),
+          fooSomethingAction
+        );
       });
+      assert.strictEqual(await lookupAction('/model/foo_something'), undefined);
+    });
+  });
 
-      assert.strictEqual(
-        await registry.lookupAction('/model/foo_something'),
-        undefined
+  describe('runInIsolatedRegistry', () => {
+    it('should lookup parent registry when child missing action', async () => {
+      const fooSomethingAction = action(
+        { name: 'foo_something' },
+        async () => null
       );
-      assert.strictEqual(
-        await childRegistry.lookupAction('/model/foo_something'),
-        fooSomethingAction
+      const barSomethingAction = action(
+        { name: 'bar_something' },
+        async () => null
       );
-      assert.strictEqual(await lookupAction('/model/foo_something'), undefined);
-      await runWithRegistry(childRegistry, async () => {
+      registerAction('model', barSomethingAction);
+
+      await runInIsolatedRegistry(async () => {
+        registerAction('model', fooSomethingAction);
+        assert.strictEqual(
+          await lookupAction('/model/bar_something'),
+          undefined
+        );
         assert.strictEqual(
           await lookupAction('/model/foo_something'),
           fooSomethingAction
         );
       });
+      assert.strictEqual(await lookupAction('/model/foo_something'), undefined);
     });
   });
 });

From ab847c14677a5bac8243334c3186472514a14f18 Mon Sep 17 00:00:00 2001
From: Pavel Jbanov <pavelj@google.com>
Date: Wed, 17 Jul 2024 13:26:56 -0400
Subject: [PATCH 05/12] doc comment cleanup

---
 js/core/src/registry.ts | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/js/core/src/registry.ts b/js/core/src/registry.ts
index cdf304f18c..899fe728b7 100644
--- a/js/core/src/registry.ts
+++ b/js/core/src/registry.ts
@@ -334,7 +334,8 @@ global[REGISTRY_KEY] = new Registry();
 const registryAls = new AsyncLocalStorage<Registry>();
 
 /**
- * Executes provided function with within an isolated registry that has no visibility into global namespace.
+ * Executes provided function with within an isolated registry that has no
+ * visibility into global namespace.
  */
 export function runInIsolatedRegistry<O>(fn: (registry: Registry) => O): O {
   const registry = new Registry();
@@ -342,7 +343,9 @@ export function runInIsolatedRegistry<O>(fn: (registry: Registry) => O): O {
 }
 
 /**
- * Executes provided function with within a temporary registry that inherits from global/parent.
+ * Executes provided function within a temporary registry overlaid onto the
+ * provided registry that will be immediately discarded at the end with no
+ * changes to the original.
  */
 export function runInTempRegistry<O>(fn: (registry: Registry) => O): O {
   const registry = Registry.withCurrent();

From 9adf3205c8b038dbdb49a34445f8be5d2caa62a9 Mon Sep 17 00:00:00 2001
From: Pavel Jbanov <pavelj@google.com>
Date: Wed, 17 Jul 2024 13:30:36 -0400
Subject: [PATCH 06/12] brought back runInRegistry

---
 js/core/src/registry.ts | 14 ++++++++++++--
 1 file changed, 12 insertions(+), 2 deletions(-)

diff --git a/js/core/src/registry.ts b/js/core/src/registry.ts
index 899fe728b7..5e7973aeb7 100644
--- a/js/core/src/registry.ts
+++ b/js/core/src/registry.ts
@@ -333,13 +333,23 @@ global[REGISTRY_KEY] = new Registry();
 
 const registryAls = new AsyncLocalStorage<Registry>();
 
+/**
+ * Executes provided function with within the provided registry.
+ */
+export function runInRegistry<O>(
+  registry: Registry,
+  fn: (registry: Registry) => O
+): O {
+  return registryAls.run(registry, () => fn(registry));
+}
+
 /**
  * Executes provided function with within an isolated registry that has no
  * visibility into global namespace.
  */
 export function runInIsolatedRegistry<O>(fn: (registry: Registry) => O): O {
   const registry = new Registry();
-  return registryAls.run(registry, () => fn(registry));
+  return runInRegistry(registry, fn);
 }
 
 /**
@@ -349,7 +359,7 @@ export function runInIsolatedRegistry<O>(fn: (registry: Registry) => O): O {
  */
 export function runInTempRegistry<O>(fn: (registry: Registry) => O): O {
   const registry = Registry.withCurrent();
-  return registryAls.run(registry, () => fn(registry));
+  return runInRegistry(registry, fn);
 }
 
 /**

From 08f0f876ce61f1ccc75863125057a0090b23f713 Mon Sep 17 00:00:00 2001
From: Pavel Jbanov <pavelj@google.com>
Date: Wed, 17 Jul 2024 14:31:59 -0400
Subject: [PATCH 07/12] test cleanup

---
 js/core/tests/registry_test.ts | 109 +++++++++++++++------------------
 1 file changed, 48 insertions(+), 61 deletions(-)

diff --git a/js/core/tests/registry_test.ts b/js/core/tests/registry_test.ts
index 507afe42e4..5c089a949d 100644
--- a/js/core/tests/registry_test.ts
+++ b/js/core/tests/registry_test.ts
@@ -25,6 +25,7 @@ import {
   registerAction,
   registerPluginProvider,
   runInIsolatedRegistry,
+  runInRegistry,
   runInTempRegistry,
 } from '../src/registry.js';
 
@@ -348,19 +349,13 @@ describe('registry class', () => {
     it('should lookup parent registry when child missing action', async () => {
       const childRegistry = new Registry(registry);
 
-      const fooSomethingAction = action(
-        { name: 'foo_something' },
-        async () => null
-      );
-      registry.registerAction('model', fooSomethingAction);
+      const fooAction = action({ name: 'foo' }, async () => null);
+      registry.registerAction('model', fooAction);
 
+      assert.strictEqual(await registry.lookupAction('/model/foo'), fooAction);
       assert.strictEqual(
-        await registry.lookupAction('/model/foo_something'),
-        fooSomethingAction
-      );
-      assert.strictEqual(
-        await childRegistry.lookupAction('/model/foo_something'),
-        fooSomethingAction
+        await childRegistry.lookupAction('/model/foo'),
+        fooAction
       );
     });
 
@@ -369,74 +364,66 @@ describe('registry class', () => {
 
       assert.strictEqual(childRegistry.parent, registry);
 
-      const fooSomethingAction = action(
-        { name: 'foo_something' },
-        async () => null
-      );
-      childRegistry.registerAction('model', fooSomethingAction);
+      const fooAction = action({ name: 'foo' }, async () => null);
+      childRegistry.registerAction('model', fooAction);
 
+      assert.strictEqual(await registry.lookupAction('/model/foo'), undefined);
       assert.strictEqual(
-        await registry.lookupAction('/model/foo_something'),
-        undefined
-      );
-      assert.strictEqual(
-        await childRegistry.lookupAction('/model/foo_something'),
-        fooSomethingAction
+        await childRegistry.lookupAction('/model/foo'),
+        fooAction
       );
     });
   });
 
   describe('runInTempRegistry', () => {
     it('should lookup parent registry when child missing action', async () => {
-      const fooSomethingAction = action(
-        { name: 'foo_something' },
-        async () => null
-      );
-      const barSomethingAction = action(
-        { name: 'bar_something' },
-        async () => null
-      );
-      registerAction('model', barSomethingAction);
+      const fooAction = action({ name: 'foo' }, async () => null);
+      const barAction = action({ name: 'bar' }, async () => null);
+      registerAction('model', barAction);
 
       await runInTempRegistry(async () => {
-        registerAction('model', fooSomethingAction);
-        assert.strictEqual(
-          await lookupAction('/model/bar_something'),
-          barSomethingAction
-        );
-        assert.strictEqual(
-          await lookupAction('/model/foo_something'),
-          fooSomethingAction
-        );
+        registerAction('model', fooAction);
+        assert.strictEqual(await lookupAction('/model/bar'), barAction);
+        assert.strictEqual(await lookupAction('/model/foo'), fooAction);
       });
-      assert.strictEqual(await lookupAction('/model/foo_something'), undefined);
+      assert.strictEqual(await lookupAction('/model/foo'), undefined);
     });
   });
 
   describe('runInIsolatedRegistry', () => {
-    it('should lookup parent registry when child missing action', async () => {
-      const fooSomethingAction = action(
-        { name: 'foo_something' },
-        async () => null
-      );
-      const barSomethingAction = action(
-        { name: 'bar_something' },
-        async () => null
-      );
-      registerAction('model', barSomethingAction);
+    it('should not lookup parent registry when child missing action', async () => {
+      const fooAction = action({ name: 'foo' }, async () => null);
+      const barAction = action({ name: 'bar' }, async () => null);
+      registerAction('model', barAction);
 
       await runInIsolatedRegistry(async () => {
-        registerAction('model', fooSomethingAction);
-        assert.strictEqual(
-          await lookupAction('/model/bar_something'),
-          undefined
-        );
-        assert.strictEqual(
-          await lookupAction('/model/foo_something'),
-          fooSomethingAction
-        );
+        registerAction('model', fooAction);
+        assert.strictEqual(await lookupAction('/model/bar'), undefined);
+        assert.strictEqual(await lookupAction('/model/foo'), fooAction);
+      });
+      assert.strictEqual(await lookupAction('/model/foo'), undefined);
+    });
+  });
+
+  describe('runInRegistry', () => {
+    it('should use the provided registry', async () => {
+      const childRegistry = new Registry();
+
+      const fooAction = action({ name: 'foo' }, async () => null);
+
+      runInRegistry(childRegistry, () => {
+        registerAction('model', fooAction);
+      });
+
+      assert.strictEqual(await registry.lookupAction('/model/foo'), undefined);
+      assert.strictEqual(
+        await childRegistry.lookupAction('/model/foo'),
+        fooAction
+      );
+      assert.strictEqual(await lookupAction('/model/foo'), undefined);
+      await runInRegistry(childRegistry, async () => {
+        assert.strictEqual(await lookupAction('/model/foo'), fooAction);
       });
-      assert.strictEqual(await lookupAction('/model/foo_something'), undefined);
     });
   });
 });

From c1376420816d4d5487eaa4688b7c01da3eedd21b Mon Sep 17 00:00:00 2001
From: Pavel Jbanov <pavelj@google.com>
Date: Wed, 17 Jul 2024 14:36:48 -0400
Subject: [PATCH 08/12] Update js/core/src/registry.ts

Co-authored-by: Alex Pascal <apascal07@gmail.com>
---
 js/core/src/registry.ts | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/js/core/src/registry.ts b/js/core/src/registry.ts
index 5e7973aeb7..70aba66444 100644
--- a/js/core/src/registry.ts
+++ b/js/core/src/registry.ts
@@ -244,7 +244,7 @@ export class Registry {
     );
   }
 
-  private async _lookupTraceStore(
+  private async lookupOverlaidTraceStore(
     env: string
   ): Promise<TraceStore | undefined> {
     if (!this.traceStoresByEnv[env]) {

From 7e0bd97cb59bbf62502d7104d00d0bb2f6a0fc75 Mon Sep 17 00:00:00 2001
From: Pavel Jbanov <pavelj@google.com>
Date: Wed, 17 Jul 2024 14:36:54 -0400
Subject: [PATCH 09/12] Update js/core/src/registry.ts

Co-authored-by: Alex Pascal <apascal07@gmail.com>
---
 js/core/src/registry.ts | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/js/core/src/registry.ts b/js/core/src/registry.ts
index 70aba66444..bf09c56d02 100644
--- a/js/core/src/registry.ts
+++ b/js/core/src/registry.ts
@@ -273,7 +273,7 @@ export class Registry {
     );
   }
 
-  private async _lookupFlowStateStore(
+  private async lookupOverlaidFlowStateStore(
     env: string
   ): Promise<FlowStateStore | undefined> {
     if (!this.flowStateStoresByEnv[env]) {

From d658e1b1ca86220bee28575535e7ca61ad2b7c28 Mon Sep 17 00:00:00 2001
From: Pavel Jbanov <pavelj@google.com>
Date: Wed, 17 Jul 2024 14:39:17 -0400
Subject: [PATCH 10/12] fix

---
 js/core/src/registry.ts | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/js/core/src/registry.ts b/js/core/src/registry.ts
index bf09c56d02..2174bae67b 100644
--- a/js/core/src/registry.ts
+++ b/js/core/src/registry.ts
@@ -240,7 +240,8 @@ export class Registry {
 
   async lookupTraceStore(env: string): Promise<TraceStore | undefined> {
     return (
-      (await this._lookupTraceStore(env)) || this.parent?.lookupTraceStore(env)
+      (await this.lookupOverlaidTraceStore(env)) ||
+      this.parent?.lookupTraceStore(env)
     );
   }
 
@@ -268,7 +269,7 @@ export class Registry {
 
   async lookupFlowStateStore(env: string): Promise<FlowStateStore | undefined> {
     return (
-      (await this._lookupFlowStateStore(env)) ||
+      (await this.lookupOverlaidFlowStateStore(env)) ||
       this.parent?.lookupFlowStateStore(env)
     );
   }

From 7f47bce1c43a2ac574caca0ce9340f3c0bb2122c Mon Sep 17 00:00:00 2001
From: Pavel Jbanov <pavelj@google.com>
Date: Mon, 12 Aug 2024 14:29:58 -0400
Subject: [PATCH 11/12] cleanup

---
 js/core/src/registry.ts        | 46 ++++--------------------
 js/core/tests/registry_test.ts | 64 ++++------------------------------
 2 files changed, 13 insertions(+), 97 deletions(-)

diff --git a/js/core/src/registry.ts b/js/core/src/registry.ts
index 2174bae67b..b8fb1fd398 100644
--- a/js/core/src/registry.ts
+++ b/js/core/src/registry.ts
@@ -14,7 +14,6 @@
  * limitations under the License.
  */
 
-import { AsyncLocalStorage } from 'node:async_hooks';
 import * as z from 'zod';
 import { Action } from './action.js';
 import { FlowStateStore } from './flowTypes.js';
@@ -229,7 +228,6 @@ export class Registry {
     }
     this.allPluginsInitialized = true;
   }
-  
 
   registerTraceStore(
     env: string,
@@ -304,7 +302,7 @@ export class Registry {
         return cached;
       },
     };
-    }
+  }
 
   lookupPlugin(name: string) {
     return this.pluginsByName[name] || this.parent?.lookupPlugin(name);
@@ -332,42 +330,12 @@ export class Registry {
 // global regustry instance
 global[REGISTRY_KEY] = new Registry();
 
-const registryAls = new AsyncLocalStorage<Registry>();
-
-/**
- * Executes provided function with within the provided registry.
- */
-export function runInRegistry<O>(
-  registry: Registry,
-  fn: (registry: Registry) => O
-): O {
-  return registryAls.run(registry, () => fn(registry));
+/** Returns the current registry instance. */
+export function getRegistryInstance(): Registry {
+  return global[REGISTRY_KEY];
 }
 
-/**
- * Executes provided function with within an isolated registry that has no
- * visibility into global namespace.
- */
-export function runInIsolatedRegistry<O>(fn: (registry: Registry) => O): O {
-  const registry = new Registry();
-  return runInRegistry(registry, fn);
-}
-
-/**
- * Executes provided function within a temporary registry overlaid onto the
- * provided registry that will be immediately discarded at the end with no
- * changes to the original.
- */
-export function runInTempRegistry<O>(fn: (registry: Registry) => O): O {
-  const registry = Registry.withCurrent();
-  return runInRegistry(registry, fn);
-}
-
-/**
- * Returns the current registry instance:
- *  - if running within `runWithRegistry` then return that registry
- *  - else return the global registry.
- */
-function getRegistryInstance(): Registry {
-  return registryAls.getStore() || global[REGISTRY_KEY];
+/** Sets global registry instance. */
+export function setRegistryInstance(reg: Registry) {
+  global[REGISTRY_KEY] = reg;
 }
diff --git a/js/core/tests/registry_test.ts b/js/core/tests/registry_test.ts
index 5c089a949d..23c165c6eb 100644
--- a/js/core/tests/registry_test.ts
+++ b/js/core/tests/registry_test.ts
@@ -15,7 +15,7 @@
  */
 
 import assert from 'node:assert';
-import { beforeEach, describe, it } from 'node:test';
+import { afterEach, beforeEach, describe, it } from 'node:test';
 import { action } from '../src/action.js';
 import {
   Registry,
@@ -24,13 +24,11 @@ import {
   lookupAction,
   registerAction,
   registerPluginProvider,
-  runInIsolatedRegistry,
-  runInRegistry,
-  runInTempRegistry,
 } from '../src/registry.js';
 
 describe('global registry', () => {
   beforeEach(__hardResetRegistryForTesting);
+  afterEach(__hardResetRegistryForTesting);
 
   describe('listActions', () => {
     it('returns all registered actions', async () => {
@@ -216,6 +214,7 @@ describe('registry class', () => {
         },
         async () => null
       );
+      registry.registerAction('custom', fooSomethingAction);
       registry.registerPluginProvider('bar', {
         name: 'bar',
         async initializer() {
@@ -232,10 +231,11 @@ describe('registry class', () => {
         },
         async () => null
       );
+      registry.registerAction('custom', barSomethingAction);
 
       assert.deepEqual(await registry.listActions(), {
-        '/model/foo/something': fooSomethingAction,
-        '/model/bar/something': barSomethingAction,
+        '/custom/foo/something': fooSomethingAction,
+        '/custom/bar/something': barSomethingAction,
       });
     });
 
@@ -374,56 +374,4 @@ describe('registry class', () => {
       );
     });
   });
-
-  describe('runInTempRegistry', () => {
-    it('should lookup parent registry when child missing action', async () => {
-      const fooAction = action({ name: 'foo' }, async () => null);
-      const barAction = action({ name: 'bar' }, async () => null);
-      registerAction('model', barAction);
-
-      await runInTempRegistry(async () => {
-        registerAction('model', fooAction);
-        assert.strictEqual(await lookupAction('/model/bar'), barAction);
-        assert.strictEqual(await lookupAction('/model/foo'), fooAction);
-      });
-      assert.strictEqual(await lookupAction('/model/foo'), undefined);
-    });
-  });
-
-  describe('runInIsolatedRegistry', () => {
-    it('should not lookup parent registry when child missing action', async () => {
-      const fooAction = action({ name: 'foo' }, async () => null);
-      const barAction = action({ name: 'bar' }, async () => null);
-      registerAction('model', barAction);
-
-      await runInIsolatedRegistry(async () => {
-        registerAction('model', fooAction);
-        assert.strictEqual(await lookupAction('/model/bar'), undefined);
-        assert.strictEqual(await lookupAction('/model/foo'), fooAction);
-      });
-      assert.strictEqual(await lookupAction('/model/foo'), undefined);
-    });
-  });
-
-  describe('runInRegistry', () => {
-    it('should use the provided registry', async () => {
-      const childRegistry = new Registry();
-
-      const fooAction = action({ name: 'foo' }, async () => null);
-
-      runInRegistry(childRegistry, () => {
-        registerAction('model', fooAction);
-      });
-
-      assert.strictEqual(await registry.lookupAction('/model/foo'), undefined);
-      assert.strictEqual(
-        await childRegistry.lookupAction('/model/foo'),
-        fooAction
-      );
-      assert.strictEqual(await lookupAction('/model/foo'), undefined);
-      await runInRegistry(childRegistry, async () => {
-        assert.strictEqual(await lookupAction('/model/foo'), fooAction);
-      });
-    });
-  });
 });

From 557a6d40b538dea6f49056cdf348f3cfe8d66b52 Mon Sep 17 00:00:00 2001
From: Pavel Jbanov <pavelj@google.com>
Date: Mon, 12 Aug 2024 14:34:00 -0400
Subject: [PATCH 12/12] initializeAllPlugins

---
 js/core/src/registry.ts | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/js/core/src/registry.ts b/js/core/src/registry.ts
index b8fb1fd398..7924a002ef 100644
--- a/js/core/src/registry.ts
+++ b/js/core/src/registry.ts
@@ -72,6 +72,13 @@ export function registerAction<I extends z.ZodTypeAny, O extends z.ZodTypeAny>(
 
 type ActionsRecord = Record<string, Action<z.ZodTypeAny, z.ZodTypeAny>>;
 
+/**
+ * Initialize all plugins in the registry.
+ */
+export async function initializeAllPlugins() {
+  await getRegistryInstance().initializeAllPlugins();
+}
+
 /**
  * Returns all actions in the registry.
  */