@alexkroman1/aai 0.3.0

package/sdk/capnweb.ts

@@ -0,0 +1,341 @@
+// Copyright 2025 the AAI authors. MIT license.
+/**
+ * MessagePort RPC + WebSocket bridge for capnweb sandbox communication.
+ *
+ * Provides bidirectional RPC over MessagePort/Worker and WebSocket
+ * bridging for both standard WebSocket (client connections) and
+ * S2sWebSocket (.on()-style API for S2S connections).
+ *
+ * @module
+ */
+
+import type { S2sWebSocket } from "./s2s.ts";
+
+// ─── MessagePort RPC ─────────────────────────────────────────────────────────
+
+/** Minimal port interface for CapnwebEndpoint. Works with Worker, MessagePort, or worker self. */
+export type CapnwebPort = {
+  postMessage(msg: unknown, transfer?: Transferable[]): void;
+  onmessage: ((ev: MessageEvent) => any) | null; // eslint-disable-line @typescript-eslint/no-explicit-any -- varies across Worker/MessagePort/self
+};
+
+/** RPC handler function. Receives call arguments and any transferred ports. */
+export type RpcHandler = (
+  args: unknown[],
+  ports: readonly MessagePort[],
+) => unknown | Promise<unknown>;
+
+/**
+ * Bidirectional RPC endpoint over MessagePort or Worker.
+ *
+ * Both sides can send calls and handle incoming calls on the same channel.
+ *
+ * Wire protocol:
+ * - Call: `{$: 0, id: number, m: string, a: unknown[]}`
+ * - Result: `{$: 1, id: number, v?: unknown, e?: string}`
+ * - Notify (fire-and-forget): id = -1, no response sent
+ */
+export class CapnwebEndpoint {
+  private nextId = 1;
+  private pending = new Map<
+    number,
+    { resolve: (v: unknown) => void; reject: (e: Error) => void }
+  >();
+  private handlers = new Map<string, RpcHandler>();
+  private port: CapnwebPort;
+
+  constructor(port: CapnwebPort) {
+    this.port = port;
+    port.onmessage = (ev: MessageEvent) => this.onMessage(ev);
+  }
+
+  /** Register an RPC handler for incoming calls with the given method name. */
+  handle(method: string, handler: RpcHandler): void {
+    this.handlers.set(method, handler);
+  }
+
+  /** Call a remote method and wait for the result. */
+  call(method: string, args: unknown[], transfer?: Transferable[]): Promise<unknown> {
+    const id = this.nextId++;
+    return new Promise((resolve, reject) => {
+      this.pending.set(id, { resolve, reject });
+      this.send({ $: 0, id, m: method, a: args }, transfer);
+    });
+  }
+
+  /** Fire-and-forget: call a remote method without waiting for a response. */
+  notify(method: string, args: unknown[], transfer?: Transferable[]): void {
+    this.send({ $: 0, id: -1, m: method, a: args }, transfer);
+  }
+
+  private send(msg: unknown, transfer?: Transferable[]): void {
+    if (transfer && transfer.length > 0) {
+      this.port.postMessage(msg, transfer);
+    } else {
+      this.port.postMessage(msg);
+    }
+  }
+
+  private onMessage(ev: MessageEvent): void {
+    const msg = ev.data;
+    if (!msg || typeof msg !== "object" || !("$" in msg)) return;
+
+    if (msg.$ === 0) {
+      // Incoming call
+      const { id, m, a } = msg as { id: number; m: string; a: unknown[] };
+      const handler = this.handlers.get(m);
+      if (!handler) {
+        if (id >= 0) this.send({ $: 1, id, e: `No handler for ${m}` });
+        return;
+      }
+      const ports = [...ev.ports];
+      Promise.resolve()
+        .then(() => handler(a, ports))
+        .then((v) => {
+          if (id >= 0) this.send({ $: 1, id, v });
+        })
+        .catch((err: unknown) => {
+          if (id >= 0) {
+            this.send({
+              $: 1,
+              id,
+              e: err instanceof Error ? err.message : String(err),
+            });
+          }
+        });
+    } else if (msg.$ === 1) {
+      // Incoming result
+      const { id, v, e } = msg as { id: number; v?: unknown; e?: string };
+      const pending = this.pending.get(id);
+      if (!pending) return;
+      this.pending.delete(id);
+      if (e !== undefined) {
+        pending.reject(new Error(e));
+      } else {
+        pending.resolve(v);
+      }
+    }
+  }
+}
+
+// ─── WebSocket Bridge Protocol ───────────────────────────────────────────────
+
+/**
+ * Bridge message format for WebSocket-over-MessagePort:
+ * - `{k:0, d:string}` — text frame
+ * - `{k:1, d:ArrayBuffer}` — binary frame (transferred zero-copy)
+ * - `{k:2, code?, reason?}` — close
+ * - `{k:3}` — open
+ * - `{k:4, m:string}` — error
+ */
+type BridgeMsg =
+  | { k: 0; d: string }
+  | { k: 1; d: ArrayBuffer }
+  | { k: 2; code?: number; reason?: string }
+  | { k: 3 }
+  | { k: 4; m: string };
+
+// ─── BridgedWebSocket (standard EventTarget-based) ───────────────────────────
+
+/**
+ * Wraps a MessagePort as a standard WebSocket (extends EventTarget).
+ * Used in the worker for client connections bridged from the host.
+ */
+export class BridgedWebSocket extends EventTarget {
+  readyState = 0;
+  private port: MessagePort;
+
+  constructor(port: MessagePort) {
+    super();
+    this.port = port;
+    port.onmessage = (ev: MessageEvent) => {
+      const msg = ev.data as BridgeMsg;
+      switch (msg.k) {
+        case 0:
+          this.dispatchEvent(new MessageEvent("message", { data: msg.d }));
+          break;
+        case 1:
+          this.dispatchEvent(new MessageEvent("message", { data: msg.d }));
+          break;
+        case 2:
+          this.readyState = 3;
+          this.dispatchEvent(
+            new CloseEvent("close", {
+              ...(msg.code !== undefined ? { code: msg.code } : {}),
+              ...(msg.reason !== undefined ? { reason: msg.reason } : {}),
+            }),
+          );
+          break;
+        case 3:
+          this.readyState = 1;
+          this.dispatchEvent(new Event("open"));
+          break;
+        case 4:
+          this.dispatchEvent(new ErrorEvent("error", { message: msg.m }));
+          break;
+      }
+    };
+  }
+
+  send(data: string | ArrayBuffer | Uint8Array): void {
+    if (this.readyState !== 1) return;
+    if (typeof data === "string") {
+      this.port.postMessage({ k: 0, d: data });
+    } else {
+      const ab =
+        data instanceof ArrayBuffer
+          ? data
+          : data.buffer.slice(data.byteOffset, data.byteOffset + data.byteLength);
+      this.port.postMessage({ k: 1, d: ab }, [ab]);
+    }
+  }
+
+  close(code?: number, reason?: string): void {
+    if (this.readyState >= 2) return;
+    this.readyState = 2;
+    this.port.postMessage({ k: 2, code, reason });
+  }
+}
+
+// ─── BridgedS2sWebSocket (.on()-style API) ───────────────────────────────────
+
+/**
+ * Wraps a MessagePort as an {@linkcode S2sWebSocket} (.on() event API).
+ * Used in the worker for S2S connections bridged from the host.
+ */
+export function createBridgedS2sWebSocket(port: MessagePort): S2sWebSocket {
+  let readyState = 0;
+  const handlers = new Map<string, ((...args: unknown[]) => void)[]>();
+
+  function emit(event: string, ...args: unknown[]): void {
+    for (const h of handlers.get(event) ?? []) h(...args);
+  }
+
+  port.onmessage = (ev: MessageEvent) => {
+    const msg = ev.data as BridgeMsg;
+    switch (msg.k) {
+      case 0:
+        emit("message", msg.d);
+        break;
+      case 2:
+        readyState = 3;
+        emit("close", msg.code, msg.reason);
+        break;
+      case 3:
+        readyState = 1;
+        emit("open");
+        break;
+      case 4:
+        emit("error", new Error(msg.m));
+        break;
+    }
+  };
+
+  return {
+    get readyState() {
+      return readyState;
+    },
+    send(data: string): void {
+      if (readyState !== 1) return;
+      port.postMessage({ k: 0, d: data });
+    },
+    close(): void {
+      if (readyState >= 2) return;
+      readyState = 2;
+      port.postMessage({ k: 2 });
+    },
+    on(event: string, handler: (...args: unknown[]) => void): void {
+      if (!handlers.has(event)) handlers.set(event, []);
+      handlers.get(event)!.push(handler);
+    },
+  };
+}
+
+// ─── Host-side bridges ───────────────────────────────────────────────────────
+
+/**
+ * Bridges a standard WebSocket (e.g. from `Deno.upgradeWebSocket`) to a
+ * MessagePort. Used on the host side for client connections.
+ */
+export function bridgeWebSocketToPort(ws: WebSocket, port: MessagePort): void {
+  ws.binaryType = "arraybuffer";
+
+  ws.addEventListener("open", () => {
+    port.postMessage({ k: 3 });
+  });
+
+  ws.addEventListener("message", (ev: Event) => {
+    const { data } = ev as MessageEvent;
+    if (typeof data === "string") {
+      port.postMessage({ k: 0, d: data });
+    } else if (data instanceof ArrayBuffer) {
+      port.postMessage({ k: 1, d: data }, [data]);
+    }
+  });
+
+  ws.addEventListener("close", (ev: Event) => {
+    const e = ev as CloseEvent;
+    port.postMessage({ k: 2, code: e.code, reason: e.reason });
+  });
+
+  ws.addEventListener("error", (ev: Event) => {
+    const msg = ev instanceof ErrorEvent ? ev.message : "WebSocket error";
+    port.postMessage({ k: 4, m: msg });
+  });
+
+  // Messages from worker → real WebSocket
+  port.onmessage = (ev: MessageEvent) => {
+    const msg = ev.data as BridgeMsg;
+    switch (msg.k) {
+      case 0:
+        if (ws.readyState === 1) ws.send(msg.d);
+        break;
+      case 1:
+        if (ws.readyState === 1) ws.send(msg.d);
+        break;
+      case 2:
+        ws.close(msg.code, msg.reason);
+        break;
+    }
+  };
+}
+
+/**
+ * Bridges a ws-style {@linkcode S2sWebSocket} to a MessagePort.
+ * Used on the host side for S2S connections to AssemblyAI.
+ */
+export function bridgeS2sWebSocketToPort(ws: S2sWebSocket, port: MessagePort): void {
+  ws.on("open", () => {
+    port.postMessage({ k: 3 });
+  });
+
+  ws.on("message", (data: unknown) => {
+    port.postMessage({ k: 0, d: String(data) });
+  });
+
+  ws.on("close", (code: unknown, reason: unknown) => {
+    port.postMessage({
+      k: 2,
+      code: typeof code === "number" ? code : undefined,
+      reason: String(reason ?? ""),
+    });
+  });
+
+  ws.on("error", (err: unknown) => {
+    const msg = err instanceof Error ? err.message : String(err);
+    port.postMessage({ k: 4, m: msg });
+  });
+
+  // Messages from worker → real S2S WebSocket
+  port.onmessage = (ev: MessageEvent) => {
+    const msg = ev.data as BridgeMsg;
+    switch (msg.k) {
+      case 0:
+        if (ws.readyState === 1) ws.send(msg.d);
+        break;
+      case 2:
+        ws.close();
+        break;
+    }
+  };
+}

package/sdk/define_agent.ts

@@ -0,0 +1,70 @@
+// Copyright 2025 the AAI authors. MIT license.
+/**
+ * Agent definition factory.
+ *
+ * @module
+ */
+
+import {
+  type AgentDef,
+  type AgentOptions,
+  DEFAULT_GREETING,
+  DEFAULT_INSTRUCTIONS,
+  normalizeTransport,
+} from "./types.ts";
+
+/**
+ * Create an agent definition from the given options, applying sensible defaults.
+ *
+ * This is the main entry point for defining a voice agent. The returned
+ * {@linkcode AgentDef} is consumed by the AAI server at deploy time.
+ *
+ * @param options Configuration for the agent including name, instructions,
+ *   tools, hooks, and transport settings.
+ * @returns A fully resolved agent definition with all defaults applied.
+ *
+ * @example Basic agent with a custom tool
+ * ```ts
+ * import { defineAgent } from "aai";
+ * import { z } from "zod";
+ *
+ * export default defineAgent({
+ *   name: "greeter",
+ *   instructions: "You greet people warmly.",
+ *   tools: {
+ *     greet: {
+ *       description: "Greet a user by name",
+ *       parameters: z.object({ name: z.string() }),
+ *       execute: ({ name }) => `Hello, ${name}!`,
+ *     },
+ *   },
+ * });
+ * ```
+ */
+export function defineAgent<S>(options: AgentOptions<S>): AgentDef {
+  // AgentDef erases the S generic (it's a runtime artifact consumed by the
+  // server which doesn't need the compile-time state type). The cast is safe
+  // because AgentDef's hooks/tools use the same shapes with `any`/`unknown`.
+  return {
+    name: options.name,
+    env: options.env ?? ["ASSEMBLYAI_API_KEY"],
+    transport: normalizeTransport(options.transport),
+    mode: options.mode ?? "s2s",
+    instructions: options.instructions ?? DEFAULT_INSTRUCTIONS,
+    greeting: options.greeting ?? DEFAULT_GREETING,
+    voice: options.voice ?? "",
+    ...(options.sttPrompt !== undefined && { sttPrompt: options.sttPrompt }),
+    maxSteps: options.maxSteps ?? 5,
+    ...(options.toolChoice !== undefined && { toolChoice: options.toolChoice }),
+    ...(options.builtinTools !== undefined && { builtinTools: options.builtinTools }),
+    ...(options.activeTools !== undefined && { activeTools: options.activeTools }),
+    tools: options.tools ?? {},
+    ...(options.state !== undefined && { state: options.state }),
+    ...(options.onConnect !== undefined && { onConnect: options.onConnect }),
+    ...(options.onDisconnect !== undefined && { onDisconnect: options.onDisconnect }),
+    ...(options.onError !== undefined && { onError: options.onError }),
+    ...(options.onTurn !== undefined && { onTurn: options.onTurn }),
+    ...(options.onStep !== undefined && { onStep: options.onStep }),
+    ...(options.onBeforeStep !== undefined && { onBeforeStep: options.onBeforeStep }),
+  } as AgentDef;
+}

package/sdk/direct_executor.ts

@@ -0,0 +1,195 @@
+// Copyright 2025 the AAI authors. MIT license.
+/**
+ * Direct tool execution for self-hosted mode.
+ *
+ * In self-hosted mode, agent code is trusted (you're running your own code).
+ * Tools execute directly in-process — no sandbox, no RPC.
+ *
+ * @module
+ */
+
+import { type AgentConfig, agentToolsToSchemas, type ToolSchema } from "./_internal_types.ts";
+import { getBuiltinToolDefs, getBuiltinToolSchemas } from "./builtin_tools.ts";
+import type { Kv } from "./kv.ts";
+import { createMemoryKv } from "./kv.ts";
+import type { ClientSink } from "./protocol.ts";
+import type { Logger, Metrics, S2SConfig } from "./runtime.ts";
+import { consoleLogger, DEFAULT_S2S_CONFIG, noopMetrics } from "./runtime.ts";
+import type { CreateS2sWebSocket } from "./s2s.ts";
+import { createS2sSession, type HookInvoker, type Session } from "./session.ts";
+import type { AgentDef, HookContext, StepInfo } from "./types.ts";
+import type { ExecuteTool } from "./worker_entry.ts";
+import { executeToolCall } from "./worker_entry.ts";
+
+export type DirectExecutorOptions = {
+  agent: AgentDef;
+  env: Record<string, string>;
+  kv?: Kv | undefined;
+  vectorSearch?: ((query: string, topK: number) => Promise<string>) | undefined;
+  createWebSocket?: CreateS2sWebSocket | undefined;
+  logger?: Logger | undefined;
+  metrics?: Metrics | undefined;
+  s2sConfig?: S2SConfig | undefined;
+};
+
+export type DirectExecutor = {
+  executeTool: ExecuteTool;
+  hookInvoker: HookInvoker;
+  toolSchemas: ToolSchema[];
+  createSession(opts: {
+    id: string;
+    agent: string;
+    client: ClientSink;
+    skipGreeting?: boolean;
+  }): Session;
+};
+
+/** Build a serializable AgentConfig from an AgentDef. */
+export function buildAgentConfig(agent: AgentDef): AgentConfig {
+  const config: AgentConfig = {
+    name: agent.name,
+    instructions: agent.instructions,
+    greeting: agent.greeting,
+    voice: agent.voice,
+  };
+  if (agent.sttPrompt !== undefined) config.sttPrompt = agent.sttPrompt;
+  if (typeof agent.maxSteps !== "function") config.maxSteps = agent.maxSteps;
+  if (agent.toolChoice !== undefined) config.toolChoice = agent.toolChoice;
+  if (agent.builtinTools) config.builtinTools = [...agent.builtinTools];
+  if (agent.activeTools) config.activeTools = [...agent.activeTools];
+  return config;
+}
+
+/** Create a direct (in-process) tool executor and hook invoker for an agent. */
+export function createDirectExecutor(opts: DirectExecutorOptions): DirectExecutor {
+  const {
+    agent,
+    env,
+    kv = createMemoryKv(),
+    vectorSearch,
+    createWebSocket,
+    logger = consoleLogger,
+    metrics = noopMetrics,
+    s2sConfig = DEFAULT_S2S_CONFIG,
+  } = opts;
+  const agentConfig = buildAgentConfig(agent);
+
+  // Merge custom + builtin tool definitions
+  const builtinDefs = getBuiltinToolDefs(
+    agent.builtinTools ?? [],
+    vectorSearch ? { vectorSearch } : undefined,
+  );
+  const allTools: Record<string, AgentDef["tools"][string]> = {
+    ...builtinDefs,
+    ...agent.tools,
+  };
+
+  // Build tool schemas for the S2S API
+  const customSchemas = agentToolsToSchemas(agent.tools ?? {});
+  const builtinSchemas = getBuiltinToolSchemas(agent.builtinTools ?? []);
+  const toolSchemas: ToolSchema[] = [...customSchemas, ...builtinSchemas];
+
+  // Per-session mutable state
+  const sessionState = new Map<string, unknown>();
+  const frozenEnv = Object.freeze({ ...env });
+
+  function getState(sessionId: string): unknown {
+    if (!sessionState.has(sessionId) && agent.state) {
+      sessionState.set(sessionId, agent.state());
+    }
+    return sessionState.get(sessionId) ?? {};
+  }
+
+  function makeHookContext(sessionId: string): HookContext {
+    return {
+      sessionId,
+      env: frozenEnv,
+      state: getState(sessionId) as Record<string, unknown>,
+      get kv() {
+        return kv;
+      },
+    };
+  }
+
+  const executeTool: ExecuteTool = async (name, args, sessionId, messages) => {
+    const tool = allTools[name];
+    if (!tool) return JSON.stringify({ error: `Unknown tool: ${name}` });
+
+    return executeToolCall(name, args, {
+      tool,
+      env: frozenEnv,
+      sessionId,
+      state: getState(sessionId ?? ""),
+      kv,
+      messages,
+    });
+  };
+
+  const hookInvoker: HookInvoker = {
+    async onConnect(sessionId) {
+      await agent.onConnect?.(makeHookContext(sessionId));
+    },
+    async onDisconnect(sessionId) {
+      await agent.onDisconnect?.(makeHookContext(sessionId));
+      sessionState.delete(sessionId);
+    },
+    async onTurn(sessionId, text) {
+      await agent.onTurn?.(text, makeHookContext(sessionId));
+    },
+    async onError(sessionId, error) {
+      await agent.onError?.(new Error(error.message), makeHookContext(sessionId));
+    },
+    async onStep(sessionId, step: StepInfo) {
+      await agent.onStep?.(step, makeHookContext(sessionId));
+    },
+    async resolveTurnConfig(sessionId) {
+      const ctx = makeHookContext(sessionId);
+      let maxSteps: number | undefined;
+      let activeTools: string[] | undefined;
+
+      if (typeof agent.maxSteps === "function") {
+        maxSteps = (await agent.maxSteps(ctx)) ?? undefined;
+      }
+
+      if (agent.onBeforeStep) {
+        const result = await agent.onBeforeStep(0, ctx);
+        activeTools = result?.activeTools;
+      }
+
+      if (maxSteps === undefined && activeTools === undefined) return null;
+      const config: { maxSteps?: number; activeTools?: string[] } = {};
+      if (maxSteps !== undefined) config.maxSteps = maxSteps;
+      if (activeTools !== undefined) config.activeTools = activeTools;
+      return config;
+    },
+  };
+
+  function createSession(sessionOpts: {
+    id: string;
+    agent: string;
+    client: ClientSink;
+    skipGreeting?: boolean;
+  }): Session {
+    if (!createWebSocket) {
+      throw new Error("createWebSocket not provided — pass it in DirectExecutorOptions");
+    }
+    const apiKey = frozenEnv.ASSEMBLYAI_API_KEY ?? "";
+    return createS2sSession({
+      id: sessionOpts.id,
+      agent: sessionOpts.agent,
+      client: sessionOpts.client,
+      agentConfig,
+      toolSchemas,
+      apiKey,
+      s2sConfig,
+      executeTool,
+      createWebSocket,
+      hookInvoker,
+      skipGreeting: sessionOpts.skipGreeting ?? false,
+      logger,
+      metrics,
+    });
+  }
+
+  return { executeTool, hookInvoker, toolSchemas, createSession };
+}