Runner augmentation: register agents + engines via a3t/augment

orchestrator-v2/packages/agent-3-task/package.json

@@ -9,6 +9,7 @@
   "type": "module",
   "exports": {
     ".": "./dist/index.mjs",
+    "./augment": "./dist/augment.mjs",
     "./package.json": "./package.json"
   },
   "publishConfig": {
@@ -23,7 +24,8 @@
   },
   "dependencies": {
     "@bifrost-ai/engine": "workspace:*",
-    "@bifrost-ai/interfaces-work": "workspace:*"
+    "@bifrost-ai/interfaces-work": "workspace:*",
+    "@bifrost-ai/runner": "workspace:*"
   },
   "devDependencies": {
     "@types/node": "catalog:",

orchestrator-v2/packages/agent-3-task/src/augment.spec.ts

@@ -0,0 +1,37 @@
+import { describe, expect } from "vite-plus/test";
+import test from "vitest-gwt";
+import { Runner } from "@bifrost-ai/runner";
+
+import "./augment.js";
+
+describe("agent-3-task/augment on a bare Runner", () => {
+  test("registers a task agent without pre-wiring a data registry", {
+    given: { a_bare_runner },
+    when: { a_task_agent_is_registered },
+    then: { the_handler_is_available },
+  });
+});
+
+type Context = {
+  runner: Runner;
+};
+
+function a_bare_runner(this: Context) {
+  // A bare `new Runner()` starts with an empty data registry; the augment must
+  // lazily create the agentDefinition registry (regression: it used to throw).
+  this.runner = new Runner();
+}
+
+function a_task_agent_is_registered(this: Context) {
+  this.runner.registerTaskAgent({
+    name: "greeter",
+    description: "greets",
+    promptBody: "say hi",
+    template: { parameters: {} },
+    tools: [],
+  });
+}
+
+function the_handler_is_available(this: Context) {
+  expect(this.runner.hasWorkItemHandler("task", "greeter")).toBe(true);
+}

orchestrator-v2/packages/agent-3-task/src/augment.ts

@@ -0,0 +1,34 @@
+import type { AgentDefinition, Engine } from "@bifrost-ai/engine";
+import { Runner } from "@bifrost-ai/runner";
+
+import { createTaskAgent } from "./create-task-agent.js";
+import {
+  AGENT_DEFINITION_DATA_TYPE,
+  ENGINE_DATA_TYPE,
+  isAgentDefinition,
+  isEngine,
+} from "./types.js";
+
+declare module "@bifrost-ai/runner" {
+  // oxlint-disable-next-line typescript/consistent-type-definitions -- module augmentation
+  interface Runner {
+    registerTaskAgent(agent: AgentDefinition): void;
+    registerEngine(name: string, engine: Engine): void;
+  }
+}
+
+Runner.prototype.registerTaskAgent = function registerTaskAgent(
+  this: Runner,
+  agent: AgentDefinition,
+): void {
+  this.data.ensure(AGENT_DEFINITION_DATA_TYPE, isAgentDefinition).register(agent.name, agent);
+  this.registerWorkItemHandler(createTaskAgent(agent));
+};
+
+Runner.prototype.registerEngine = function registerEngine(
+  this: Runner,
+  name: string,
+  engine: Engine,
+): void {
+  this.data.ensure(ENGINE_DATA_TYPE, isEngine).register(name, engine);
+};

orchestrator-v2/packages/agent-3-task/vite.config.ts

@@ -1,3 +1,11 @@
+import { resolve } from "node:path";
 import { defineConfig } from "vite-plus";
 
-export default defineConfig({});
+export default defineConfig({
+  pack: {
+    entry: {
+      index: resolve(__dirname, "src/index.ts"),
+      augment: resolve(__dirname, "src/augment.ts"),
+    },
+  },
+});

orchestrator-v2/packages/interfaces-work/src/types.ts

@@ -103,6 +103,10 @@ export type DataRegistry<T extends Record<string, unknown>> = {
 
 export type MutableDataRegistry<T extends Record<string, unknown>> = {
   get<K extends keyof T & string>(type: K): Registry<T[K]>;
+  ensure<K extends keyof T & string>(
+    type: K,
+    guard: (value: unknown) => value is T[K],
+  ): Registry<T[K]>;
 };
 
 export type WorkItemHandlerRegistry = {

orchestrator-v2/packages/runner/src/data-registry.ts

@@ -15,8 +15,20 @@ export function createDataRegistry<T extends Record<string, unknown>>(
 ): MutableDataRegistry<T> {
   const registries = new Map<keyof T & string, Registry<unknown>>();
 
+  const ensureRegistry = <K extends keyof T & string>(
+    type: K,
+    guard: (value: unknown) => value is T[K],
+  ): Registry<T[K]> => {
+    let registry = registries.get(type);
+    if (registry === undefined) {
+      registry = createGuardedRegistry(guard);
+      registries.set(type, registry);
+    }
+    return registry as Registry<T[K]>;
+  };
+
   for (const type of Object.keys(guards) as (keyof T & string)[]) {
-    registries.set(type, createGuardedRegistry(guards[type]));
+    ensureRegistry(type, guards[type]);
   }
 
   return {
@@ -27,6 +39,7 @@ export function createDataRegistry<T extends Record<string, unknown>>(
       }
       return registry as Registry<T[K]>;
     },
+    ensure: ensureRegistry,
   };
 }