[build] Pass context to invoke functions (breadboard-ai#1518)

.changeset/strong-planets-brake.md

@@ -0,0 +1,5 @@
+---
+"@breadboard-ai/build": patch
+---
+
+Pass NodeHandlerContext to invoke functions

packages/build/src/internal/define/define.ts

@@ -7,6 +7,7 @@
 
 import type {
   NodeDescriberContext,
+  NodeHandlerContext,
   NodeHandlerMetadata,
 } from "@google-labs/breadboard";
 import type { CountUnion, Expand, MaybePromise } from "../common/type-util.js";
@@ -179,10 +180,7 @@ export function defineNodeType<
     params.outputs["*"],
     primary(params.inputs),
     primary(params.outputs),
-    params.invoke as Function as (
-      staticParams: Record<string, JsonSerializable>,
-      dynamicParams: Record<string, JsonSerializable>
-    ) => { [K: string]: JsonSerializable },
+    params.invoke as Function as VeryLooseInvokeFn,
     params.describe as LooseDescribeFn
   );
   return Object.assign(impl.instantiate.bind(impl), {
@@ -245,19 +243,24 @@ type LooseInvokeFn<I extends Record<string, InputPortConfig>> = Expand<
   (
     staticParams: Expand<StaticInvokeParams<I>>,
     dynamicParams: Expand<DynamicInvokeParams<I>>
-  ) =>
-    | { [K: string]: JsonSerializable }
-    | Promise<{ [K: string]: JsonSerializable }>
+  ) => MaybePromise<{ [K: string]: JsonSerializable }>
 >;
 
+export type VeryLooseInvokeFn = (
+  staticParams: Record<string, JsonSerializable>,
+  dynamicParams: Record<string, JsonSerializable>,
+  context: NodeHandlerContext
+) => { [K: string]: JsonSerializable };
+
 type StrictInvokeFn<
   I extends Record<string, InputPortConfig>,
   O extends Record<string, OutputPortConfig>,
   F extends LooseInvokeFn<I>,
 > = (
   staticInputs: Expand<StaticInvokeParams<I>>,
-  dynamicInputs: Expand<DynamicInvokeParams<I>>
-) => StrictInvokeFnReturn<I, O, F> | Promise<StrictInvokeFnReturn<I, O, F>>;
+  dynamicInputs: Expand<DynamicInvokeParams<I>>,
+  context: NodeHandlerContext
+) => MaybePromise<StrictInvokeFnReturn<I, O, F>>;
 
 type StrictInvokeFnReturn<
   I extends Record<string, InputPortConfig>,

packages/build/src/internal/define/definition.ts

@@ -8,6 +8,7 @@ import type {
   InputValues,
   NodeDescriberContext,
   NodeDescriberResult,
+  NodeHandlerContext,
   OutputValues,
   Schema,
 } from "@google-labs/breadboard";
@@ -26,7 +27,11 @@ import type {
   StaticInputPortConfig,
   StaticOutputPortConfig,
 } from "./config.js";
-import type { DynamicInputPorts, LooseDescribeFn } from "./define.js";
+import type {
+  DynamicInputPorts,
+  LooseDescribeFn,
+  VeryLooseInvokeFn,
+} from "./define.js";
 import { Instance } from "./instance.js";
 import { portConfigMapToJSONSchema } from "./json-schema.js";
 
@@ -71,11 +76,7 @@ export class DefinitionImpl<
   readonly #reflective: boolean;
   readonly #primaryInput: string | undefined;
   readonly #primaryOutput: string | undefined;
-  // TODO(aomarks) Support promises
-  readonly #invoke: (
-    staticParams: Record<string, JsonSerializable>,
-    dynamicParams: Record<string, JsonSerializable>
-  ) => { [K: string]: JsonSerializable };
+  readonly #invoke: VeryLooseInvokeFn;
   readonly #describe?: LooseDescribeFn;
 
   constructor(
@@ -86,10 +87,7 @@ export class DefinitionImpl<
     dynamicOutputs: DynamicOutputPortConfig | undefined,
     primaryInput: string | undefined,
     primaryOutput: string | undefined,
-    invoke: (
-      staticParams: Record<string, JsonSerializable>,
-      dynamicParams: Record<string, JsonSerializable>
-    ) => { [K: string]: JsonSerializable },
+    invoke: VeryLooseInvokeFn,
     describe?: LooseDescribeFn
   ) {
     this.#name = name;
@@ -130,10 +128,13 @@ export class DefinitionImpl<
     );
   }
 
-  invoke(values: InputValues): Promise<OutputValues> {
+  invoke(
+    values: InputValues,
+    context: NodeHandlerContext
+  ): Promise<OutputValues> {
     const { staticValues, dynamicValues } =
       this.#applyDefaultsAndPartitionRuntimeInputValues(values);
-    return Promise.resolve(this.#invoke(staticValues, dynamicValues));
+    return Promise.resolve(this.#invoke(staticValues, dynamicValues, context));
   }
 
   /**

packages/build/src/test/invoke_test.ts

@@ -0,0 +1,28 @@
+/**
+ * @license
+ * Copyright 2024 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import type { NodeHandlerContext } from "@google-labs/breadboard";
+import assert from "node:assert/strict";
+import { test } from "node:test";
+import { defineNodeType } from "../internal/define/define.js";
+
+test("invoke receives context", async () => {
+  const expected: NodeHandlerContext = {
+    base: new URL("http://example.com/"),
+    outerGraph: { nodes: [], edges: [] },
+  };
+  let actual: NodeHandlerContext | undefined;
+  defineNodeType({
+    name: "foo",
+    inputs: {},
+    outputs: {},
+    invoke: (_staticInputs, _dynamicInputs, context) => {
+      actual = context;
+      return {};
+    },
+  }).invoke({}, expected);
+  assert.deepEqual(actual, expected);
+});