feat(js/core): DAP expansion: Dynamic Actions are now expanded in the… (#3927)

Author: ifielker

js/core/src/dynamic-action-provider.ts

@@ -16,7 +16,8 @@
 
 import type * as z from 'zod';
 import { Action, ActionMetadata, defineAction } from './action.js';
-import { ActionType, Registry } from './registry.js';
+import { GenkitError } from './error.js';
+import { ActionMetadataRecord, ActionType, Registry } from './registry.js';
 
 type DapValue = {
   [K in ActionType]?: Action<z.ZodTypeAny, z.ZodTypeAny, z.ZodTypeAny>[];
@@ -42,7 +43,12 @@ class SimpleCache {
       : config.cacheConfig?.ttlMillis;
   }
 
-  async getOrFetch(): Promise<DapValue> {
+  /**
+   * Gets or fetches the DAP data.
+   * @param skipTrace Don't run the action. i.e. don't create a trace log.
+   * @returns The DAP data
+   */
+  async getOrFetch(params?: { skipTrace?: boolean }): Promise<DapValue> {
     const isStale =
       !this.value ||
       !this.expiresAt ||
@@ -59,9 +65,13 @@ class SimpleCache {
           this.value = await this.dapFn(); // this returns the actual actions
           this.expiresAt = Date.now() + this.ttlMillis;
 
-          // Also run the action
-          this.dap.run(this.value); // This returns metadata and shows up in dev UI
-
+          if (!params?.skipTrace) {
+            // Also run the action
+            // This returns metadata and shows up in dev UI
+            // It does not change what we return, it just makes
+            // the content of the DAP visible in the trace.
+            await this.dap.run(this.value);
+          }
           return this.value;
         } catch (error) {
           console.error('Error fetching Dynamic Action Provider value:', error);
@@ -92,6 +102,7 @@ export interface DynamicRegistry {
     actionType: string,
     actionName: string
   ): Promise<ActionMetadata[]>;
+  getActionMetadataRecord(dapPrefix: string): Promise<ActionMetadataRecord>;
 }
 
 export type DynamicActionProviderAction = Action<
@@ -210,4 +221,29 @@ function implementDap(
     // Single match or empty array
     return metadata.filter((m) => m.name == actionName);
   };
+
+  // This is called by listResolvableActions which is used by the
+  // reflection API.
+  dap.getActionMetadataRecord = async (dapPrefix: string) => {
+    const dapActions = {} as ActionMetadataRecord;
+    // We want to skip traces so we don't get a new action trace
+    // every time the DevUI requests the list of actions.
+    // This is ok, because the DevUI will show the actions, so
+    // not having them in the trace is fine.
+    const result = await dap.__cache.getOrFetch({ skipTrace: true });
+    for (const [actionType, actions] of Object.entries(result)) {
+      const metadataList = actions.map((a) => a.__action);
+      for (const metadata of metadataList) {
+        if (!metadata.name) {
+          throw new GenkitError({
+            status: 'INVALID_ARGUMENT',
+            message: `Invalid metadata when listing dynamic actions from ${dapPrefix} - name required`,
+          });
+        }
+        const key = `${dapPrefix}:${actionType}/${metadata.name}`;
+        dapActions[key] = metadata;
+      }
+    }
+    return dapActions;
+  };
 }

js/core/src/registry.ts

@@ -353,7 +353,7 @@ export class Registry {
    * @returns All resolvable action metadata as a map of <key, action metadata>.
    */
   async listResolvableActions(): Promise<ActionMetadataRecord> {
-    const resolvableActions = {} as ActionMetadataRecord;
+    let resolvableActions = {} as ActionMetadataRecord;
     // We listActions for all plugins in parallel.
     await Promise.all(
       Object.entries(this.pluginsByName).map(async ([pluginName, plugin]) => {
@@ -380,9 +380,22 @@ export class Registry {
         }
       })
     );
-    // Also add actions that are already registered.
+    // Also add actions that are already registered, and expand DAP actions
     for (const [key, action] of Object.entries(await this.listActions())) {
       resolvableActions[key] = action.__action;
+      if (isDynamicActionProvider(action)) {
+        try {
+          // Include the dynamic actions
+          const dapPrefix = `${action.__action.actionType}/${action.__action.name}`;
+          const dapMetadataRecord =
+            await action.getActionMetadataRecord(dapPrefix);
+          resolvableActions = { ...resolvableActions, ...dapMetadataRecord };
+        } catch (e) {
+          logger.error(
+            `Error listing actions for Dynamic Action Provider ${action.__action.name}`
+          );
+        }
+      }
     }
     return {
       ...(await this.parent?.listResolvableActions()),

js/core/tests/dynamic-action-provider_test.ts

@@ -171,6 +171,25 @@ describe('dynamic action provider', () => {
     assert.strictEqual(callCount, 1);
   });
 
+  it('gets action metadata record', async () => {
+    let callCount = 0;
+    const dap = defineDynamicActionProvider(registry, 'my-dap', async () => {
+      callCount++;
+      return {
+        tool: [tool1, tool2],
+        flow: [tool1],
+      };
+    });
+
+    const record = await dap.getActionMetadataRecord('dap/my-dap');
+    assert.deepStrictEqual(record, {
+      'dap/my-dap:tool/tool1': tool1.__action,
+      'dap/my-dap:tool/tool2': tool2.__action,
+      'dap/my-dap:flow/tool1': tool1.__action,
+    });
+    assert.strictEqual(callCount, 1);
+  });
+
   it('handles concurrent requests', async () => {
     let callCount = 0;
     const dap = defineDynamicActionProvider(registry, 'my-dap', async () => {
@@ -224,6 +243,33 @@ describe('dynamic action provider', () => {
     });
   });
 
+  it('skips trace when requested', async () => {
+    let callCount = 0;
+    const dap = defineDynamicActionProvider(registry, 'my-dap', async () => {
+      callCount++;
+      return {
+        tool: [tool1, tool2],
+      };
+    });
+
+    const originalRun = dap.run.bind(dap);
+    let runCalled = false;
+    dap.run = async (input, options) => {
+      runCalled = true;
+      return originalRun(input, options);
+    };
+
+    await dap.__cache.getOrFetch({ skipTrace: true });
+    assert.strictEqual(runCalled, false);
+    assert.strictEqual(callCount, 1);
+
+    dap.invalidateCache();
+
+    await dap.__cache.getOrFetch();
+    assert.strictEqual(runCalled, true);
+    assert.strictEqual(callCount, 2);
+  });
+
   it('identifies dynamic action providers', async () => {
     const dap = defineDynamicActionProvider(registry, 'my-dap', async () => {
       return {};

js/core/tests/registry_test.ts

@@ -21,10 +21,11 @@ import {
   defineAction,
   runInActionRuntimeContext,
 } from '../src/action.js';
-import { initNodeAsyncContext } from '../src/node-async-context.js';
+import { defineDynamicActionProvider } from '../src/dynamic-action-provider.js';
+import { initNodeFeatures } from '../src/node.js';
 import { Registry } from '../src/registry.js';
 
-initNodeAsyncContext();
+initNodeFeatures();
 
 describe('registry class', () => {
   var registry: Registry;
@@ -344,6 +345,37 @@ describe('registry class', () => {
         },
       });
     });
+    it('expands actions from dynamic action providers', async () => {
+      const tool1 = action(
+        { name: 'fs/tool1', actionType: 'tool' },
+        async () => {}
+      );
+      const tool2 = action(
+        { name: 'fs/tool2', actionType: 'tool' },
+        async () => {}
+      );
+      const resource1 = action(
+        { name: 'abc/res1', actionType: 'resource' },
+        async () => {}
+      );
+      const resource2 = action(
+        { name: 'abc/res2', actionType: 'resource' },
+        async () => {}
+      );
+      const dap = defineDynamicActionProvider(registry, 'my-dap', async () => ({
+        tool: [tool1, tool2],
+        resource: [resource1, resource2],
+      }));
+
+      const resolvableActions = await registry.listResolvableActions();
+      assert.deepStrictEqual(resolvableActions, {
+        '/dynamic-action-provider/my-dap': dap.__action,
+        'dynamic-action-provider/my-dap:tool/fs/tool1': tool1.__action,
+        'dynamic-action-provider/my-dap:tool/fs/tool2': tool2.__action,
+        'dynamic-action-provider/my-dap:resource/abc/res1': resource1.__action,
+        'dynamic-action-provider/my-dap:resource/abc/res2': resource2.__action,
+      });
+    });
   });
 
   describe('lookupAction', () => {