@alexkroman1/aai 0.12.2 → 1.0.2

package/host/runtime.ts

@@ -0,0 +1,308 @@
+// Copyright 2025 the AAI authors. MIT license.
+/**
+ * Agent runtime — the execution engine for voice agents.
+ *
+ * {@link createRuntime} builds the single execution engine used by both
+ * self-hosted servers and the platform sandbox. It wires up tool execution,
+ * lifecycle hooks, and session management.
+ */
+
+import pTimeout from "p-timeout";
+import { createStorage } from "unstorage";
+import { agentToolsToSchemas, type ToolSchema, toAgentConfig } from "../sdk/_internal-types.ts";
+import { DEFAULT_SHUTDOWN_TIMEOUT_MS } from "../sdk/constants.ts";
+import type { Kv } from "../sdk/kv.ts";
+import type { ClientSink } from "../sdk/protocol.ts";
+import { buildReadyConfig, type ReadyConfig } from "../sdk/protocol.ts";
+import type { AgentDef } from "../sdk/types.ts";
+import { toolError } from "../sdk/utils.ts";
+import { resolveAllBuiltins } from "./builtin-tools.ts";
+import type { Logger, S2SConfig } from "./runtime-config.ts";
+import { consoleLogger, DEFAULT_S2S_CONFIG } from "./runtime-config.ts";
+import type { CreateS2sWebSocket } from "./s2s.ts";
+import { createS2sSession, type Session } from "./session.ts";
+import { type ExecuteTool, executeToolCall } from "./tool-executor.ts";
+import { createUnstorageKv } from "./unstorage-kv.ts";
+import { type SessionWebSocket, wireSessionSocket } from "./ws-handler.ts";
+
+// ─── Runtime adapter (formerly adapter.ts) ──────────────────────────────────
+
+/** Per-session options passed to {@link AgentRuntime.startSession}. */
+export type SessionStartOptions = {
+  skipGreeting?: boolean;
+  resumeFrom?: string;
+  logContext?: Record<string, string>;
+  onOpen?: () => void;
+  onClose?: () => void;
+  /** Called with session ID after session cleanup, for guest state cleanup. */
+  onSessionEnd?: (sessionId: string) => void;
+};
+
+/**
+ * Common interface for agent runtimes.
+ *
+ * Implemented by {@link createRuntime} and the platform sandbox.
+ */
+export type AgentRuntime = {
+  startSession(ws: SessionWebSocket, opts?: SessionStartOptions): void;
+  shutdown(): Promise<void>;
+  readonly readyConfig: ReadyConfig;
+};
+
+// ─── Runtime implementation ──────────────────────────────────────────────────
+
+/** Create an in-memory KV store (default for self-hosted). */
+function createLocalKv(): Kv {
+  return createUnstorageKv({ storage: createStorage() });
+}
+
+/**
+ * Configuration for {@link createRuntime}.
+ *
+ * Configures the agent, environment, KV store, logging, and S2S connection.
+ *
+ * @public
+ */
+export type RuntimeOptions = {
+  // biome-ignore lint/suspicious/noExplicitAny: accepts any state type
+  agent: AgentDef<any>;
+  env: Record<string, string>;
+  kv?: Kv | undefined;
+  /** Custom WebSocket factory for the S2S connection (useful for testing). */
+  createWebSocket?: CreateS2sWebSocket | undefined;
+  logger?: Logger | undefined;
+  s2sConfig?: S2SConfig | undefined;
+  /**
+   * Timeout in ms for `session.start()` (S2S connection setup).
+   * Defaults to 10 000 (10 s).
+   */
+  sessionStartTimeoutMs?: number | undefined;
+  /**
+   * Maximum time in milliseconds to wait for sessions to stop during
+   * {@link AgentRuntime.shutdown | shutdown()}. Defaults to `30_000` (30 s).
+   */
+  shutdownTimeoutMs?: number | undefined;
+  /**
+   * Override tool execution. When provided, `createRuntime` skips building
+   * in-process tool definitions and uses this function instead. Used by the
+   * platform sandbox to RPC tool calls to the isolate.
+   */
+  executeTool?: ExecuteTool | undefined;
+  /**
+   * Override tool schemas sent to the S2S API. Required when `executeTool`
+   * is provided (the host doesn't have the tool definitions to derive schemas).
+   */
+  toolSchemas?: ToolSchema[] | undefined;
+  /** System prompt guidance for builtin tools. Passed through in sandbox mode. */
+  toolGuidance?: string[] | undefined;
+  /**
+   * Pre-resolved builtin tool definitions. When provided alongside `executeTool`
+   * and `toolSchemas`, skips calling `resolveAllBuiltins` on the host.
+   */
+  builtinDefs?: Record<string, import("../sdk/types.ts").ToolDef> | undefined;
+  /**
+   * Override the fetch implementation used by built-in tools (web_search,
+   * visit_webpage, fetch_json). Defaults to `globalThis.fetch`.
+   *
+   * In platform mode, pass an SSRF-safe fetch to prevent requests to
+   * private/internal networks. In self-hosted mode, users may provide
+   * their own fetch wrapper.
+   */
+  fetch?: typeof globalThis.fetch | undefined;
+};
+
+/**
+ * The agent runtime returned by {@link createRuntime}.
+ *
+ * Satisfies {@link AgentRuntime} for use by transport code, and also exposes
+ * lower-level helpers (`executeTool`, `hooks`, `toolSchemas`,
+ * `createSession`) for testing and advanced usage.
+ *
+ * @public
+ */
+export type Runtime = AgentRuntime & {
+  /** Execute a named tool with the given args, returning a JSON result string. */
+  executeTool: ExecuteTool;
+  /** Tool schemas registered with the S2S API (custom + built-in). */
+  toolSchemas: ToolSchema[];
+  /** Create a new voice session for a connected client (lower-level than startSession). */
+  createSession(opts: {
+    id: string;
+    agent: string;
+    client: ClientSink;
+    skipGreeting?: boolean;
+    resumeFrom?: string;
+  }): Session;
+};
+
+/**
+ * Create an agent runtime — the execution engine for a voice agent.
+ *
+ * Merges built-in and custom tool definitions, builds tool schemas for the
+ * S2S API, and wires up lifecycle hooks.
+ *
+ * @param opts - Runtime configuration. See {@link RuntimeOptions}.
+ * @returns A {@link Runtime} with tool execution, hook invocation,
+ *   schemas, and session management.
+ *
+ * @public
+ */
+export function createRuntime(opts: RuntimeOptions): Runtime {
+  const {
+    agent,
+    env,
+    kv = createLocalKv(),
+    createWebSocket,
+    logger = consoleLogger,
+    s2sConfig = DEFAULT_S2S_CONFIG,
+    sessionStartTimeoutMs,
+    shutdownTimeoutMs = DEFAULT_SHUTDOWN_TIMEOUT_MS,
+  } = opts;
+  const agentConfig = toAgentConfig(agent);
+  const sessions = new Map<string, Session>();
+  const readyConfig: ReadyConfig = buildReadyConfig(s2sConfig);
+
+  // When overrides are provided (sandbox mode), skip in-process tool setup
+  let executeTool: ExecuteTool;
+  let toolSchemas: ToolSchema[];
+  let toolGuidance: string[] = [];
+
+  const builtinFetchOpt = opts.fetch ? { fetch: opts.fetch } : undefined;
+
+  if (opts.executeTool && opts.toolSchemas) {
+    // Sandbox mode — custom tools are RPC-backed; builtins run host-side
+    const builtinDefs =
+      opts.builtinDefs ?? resolveAllBuiltins(agent.builtinTools ?? [], builtinFetchOpt).defs;
+    const rpcExecuteTool = opts.executeTool;
+    const frozenEnv = Object.freeze({ ...env });
+
+    executeTool = async (name, args, sessionId, messages) => {
+      // Handle builtins on the host (where SSRF-safe fetch lives)
+      if (builtinDefs[name]) {
+        const tool = builtinDefs[name];
+        return executeToolCall(name, args, {
+          tool,
+          env: frozenEnv,
+          sessionId: sessionId ?? "",
+          kv,
+          messages,
+          logger,
+        });
+      }
+      // Delegate custom tools to the isolate via RPC
+      return rpcExecuteTool(name, args, sessionId, messages);
+    };
+
+    toolSchemas = opts.toolSchemas;
+    toolGuidance = opts.toolGuidance ?? [];
+  } else {
+    // Self-hosted mode — in-process tool execution
+    const builtins = resolveAllBuiltins(agent.builtinTools ?? [], builtinFetchOpt);
+    const allTools: Record<string, AgentDef["tools"][string]> = {
+      ...builtins.defs,
+      ...agent.tools,
+    };
+    const customSchemas = agentToolsToSchemas(agent.tools ?? {});
+    toolSchemas = [...customSchemas, ...builtins.schemas];
+    toolGuidance = builtins.guidance;
+
+    const stateMap = new Map<string, Record<string, unknown>>();
+    const getState = (sid: string) => {
+      if (!stateMap.has(sid) && agent.state) stateMap.set(sid, agent.state());
+      return stateMap.get(sid) ?? {};
+    };
+    const frozenEnv = Object.freeze({ ...env });
+
+    executeTool = async (name, args, sessionId, messages) => {
+      const tool = allTools[name];
+      if (!tool) return toolError(`Unknown tool: ${name}`);
+      return executeToolCall(name, args, {
+        tool,
+        env: frozenEnv,
+        state: getState(sessionId ?? ""),
+        sessionId: sessionId ?? "",
+        kv,
+        messages,
+        logger,
+      });
+    };
+  }
+
+  function createSession(sessionOpts: {
+    id: string;
+    agent: string;
+    client: ClientSink;
+    skipGreeting?: boolean;
+    resumeFrom?: string;
+  }): Session {
+    const apiKey = env.ASSEMBLYAI_API_KEY ?? "";
+    return createS2sSession({
+      id: sessionOpts.id,
+      agent: sessionOpts.agent,
+      client: sessionOpts.client,
+      agentConfig,
+      toolSchemas,
+      toolGuidance,
+      apiKey,
+      s2sConfig,
+      executeTool,
+      ...(createWebSocket ? { createWebSocket } : {}),
+      skipGreeting: sessionOpts.skipGreeting ?? false,
+      logger,
+      ...(sessionOpts.resumeFrom ? { resumeFrom: sessionOpts.resumeFrom } : {}),
+    });
+  }
+
+  // ── AgentRuntime methods ──────────────────────────────────────────────
+
+  function startSession(ws: SessionWebSocket, startOpts?: SessionStartOptions): void {
+    const resumeFrom = startOpts?.resumeFrom;
+    wireSessionSocket(ws, {
+      sessions,
+      createSession: (sid, client) =>
+        createSession({
+          id: sid,
+          agent: agent.name,
+          client,
+          skipGreeting: startOpts?.skipGreeting ?? false,
+          ...(resumeFrom ? { resumeFrom } : {}),
+        }),
+      readyConfig,
+      logger,
+      ...(startOpts?.logContext ? { logContext: startOpts.logContext } : {}),
+      ...(startOpts?.onOpen ? { onOpen: startOpts.onOpen } : {}),
+      ...(startOpts?.onClose ? { onClose: startOpts.onClose } : {}),
+      ...(startOpts?.onSessionEnd ? { onSessionEnd: startOpts.onSessionEnd } : {}),
+      ...(sessionStartTimeoutMs !== undefined ? { sessionStartTimeoutMs } : {}),
+      ...(resumeFrom ? { resumeFrom } : {}),
+    });
+  }
+
+  async function shutdown(): Promise<void> {
+    if (sessions.size === 0) return;
+    try {
+      const results = await pTimeout(
+        Promise.allSettled([...sessions.values()].map((s) => s.stop())),
+        { milliseconds: shutdownTimeoutMs },
+      );
+      for (const r of results) {
+        if (r.status === "rejected")
+          logger.warn(`Session stop failed during shutdown: ${r.reason}`);
+      }
+    } catch {
+      logger.warn(
+        `Shutdown timeout (${shutdownTimeoutMs}ms) exceeded — force-closing ${sessions.size} remaining session(s)`,
+      );
+    }
+    sessions.clear();
+  }
+
+  return {
+    executeTool,
+    toolSchemas,
+    createSession,
+    startSession,
+    shutdown,
+    readyConfig,
+  };
+}

package/host/s2s-fixtures.test.ts

@@ -0,0 +1,237 @@
+import { describe, expect, test, vi } from "vitest";
+import { loadFixture, silentLogger } from "./_test-utils.ts";
+import type { S2sWebSocket } from "./s2s.ts";
+import { connectS2s } from "./s2s.ts";
+
+/** EventTarget-based WebSocket stub (standard API, no `.on()` adapter needed). */
+function createWebSocketStub() {
+  const target = new EventTarget();
+  return Object.assign(target, {
+    readyState: 0,
+    send: vi.fn(),
+    close: vi.fn(),
+    addEventListener: target.addEventListener.bind(target) as S2sWebSocket["addEventListener"],
+    /** Simulate a server-side event for testing. */
+    emit(event: string, ...args: unknown[]) {
+      const builders: Record<string, () => Event> = {
+        open: () => new Event("open"),
+        message: () => new MessageEvent("message", { data: args[0] }),
+        close: () => {
+          const ev = new Event("close");
+          if (typeof args[0] === "number") Object.assign(ev, { code: args[0] });
+          if (typeof args[1] === "string") Object.assign(ev, { reason: args[1] });
+          return ev;
+        },
+        error: () => {
+          const msg = args[0] instanceof Error ? args[0].message : String(args[0]);
+          const ev = new Event("error");
+          Object.defineProperty(ev, "message", { value: msg });
+          return ev;
+        },
+      };
+      const build = builders[event];
+      if (build) target.dispatchEvent(build());
+    },
+  });
+}
+
+const s2sConfig = { wssUrl: "wss://fake", inputSampleRate: 16_000, outputSampleRate: 16_000 };
+
+function createTestS2s() {
+  const raw = createWebSocketStub();
+  const createWebSocket = () => {
+    setTimeout(() => {
+      raw.readyState = 1;
+      raw.emit("open");
+    }, 0);
+    return raw;
+  };
+  return { raw, createWebSocket, logger: { ...silentLogger } };
+}
+
+async function setupHandle() {
+  const { raw, createWebSocket, logger } = createTestS2s();
+  const handle = await connectS2s({
+    apiKey: "test-key",
+    config: s2sConfig,
+    createWebSocket,
+    logger,
+  });
+  return { raw, handle, logger };
+}
+
+// ─── Fixture-based tests (real API responses from Kokoro TTS audio) ─────
+
+describe("real API fixtures", () => {
+  /**
+   * Replay all fixture messages through the S2S handle and collect events.
+   * Events are collected from both the 'ready'/'replyStarted'/'sessionExpired' special
+   * events and the unified 'event' emitter, tagged with their source type.
+   */
+  async function replayFixture(fixtureName: string) {
+    const { raw, handle } = await setupHandle();
+    const events: { type: string; payload: unknown }[] = [];
+
+    // Special events that are NOT in the 'event' emitter
+    handle.on("ready", (p) => events.push({ type: "ready", payload: p }));
+    handle.on("replyStarted", (p) => events.push({ type: "replyStarted", payload: p }));
+    handle.on("sessionExpired", () => events.push({ type: "sessionExpired", payload: undefined }));
+    handle.on("audio", (p) => events.push({ type: "audio", payload: p }));
+    handle.on("error", (p) => events.push({ type: "error", payload: p }));
+
+    // All protocol-shaped events via the unified 'event' emitter
+    handle.on("event", (event) => events.push({ type: event.type, payload: event }));
+
+    const fixtures = loadFixture<Record<string, unknown>[]>(fixtureName);
+    for (const msg of fixtures) {
+      raw.emit("message", Buffer.from(JSON.stringify(msg)));
+    }
+
+    return { events, fixtures, raw, handle };
+  }
+
+  // ── Session lifecycle ──────────────────────────────────────────────────
+
+  test("parses real session.ready messages with extra fields (timestamp, config)", async () => {
+    const { events } = await replayFixture("session-ready.json");
+
+    const readyEvents = events.filter((e) => e.type === "ready");
+    expect(readyEvents.length).toBeGreaterThan(0);
+    expect((readyEvents[0]?.payload as { sessionId: string }).sessionId).toMatch(/^sess_/);
+  });
+
+  test("parses real session.updated messages — they are now silently dropped", async () => {
+    const { events } = await replayFixture("session-updated.json");
+
+    // session.updated is no longer dispatched — it is dropped in s2s.ts
+    const updatedEvents = events.filter((e) => e.type === "session.updated");
+    expect(updatedEvents.length).toBe(0);
+  });
+
+  // ── Greeting session ───────────────────────────────────────────────────
+
+  test("greeting session produces correct event sequence", async () => {
+    const { events } = await replayFixture("greeting-session-sequence.json");
+
+    const types = events.map((e) => e.type);
+    // session.updated is dropped, so first non-audio event is 'ready'
+    expect(types[0]).toBe("ready");
+    expect(types[1]).toBe("replyStarted");
+    expect(types.filter((t) => t === "agent_transcript").length).toBeGreaterThan(0);
+    expect(types).toContain("agent_transcript");
+    expect(types.at(-1)).toBe("reply_done");
+  });
+
+  // ── Reply lifecycle ────────────────────────────────────────────────────
+
+  test("real transcript.agent has _interrupted field", async () => {
+    const { events } = await replayFixture("reply-lifecycle.json");
+
+    const transcripts = events.filter((e) => e.type === "agent_transcript");
+    expect(transcripts.length).toBe(1);
+    const payload = transcripts[0]?.payload as {
+      type: string;
+      text: string;
+      _interrupted: boolean;
+    };
+    expect(payload._interrupted).toBe(false);
+  });
+
+  // ── Audio ──────────────────────────────────────────────────────────────
+
+  test("real reply.audio messages decode to Uint8Array", async () => {
+    const { events } = await replayFixture("reply-audio-samples.json");
+
+    const audioEvents = events.filter((e) => e.type === "audio");
+    expect(audioEvents.length).toBeGreaterThan(0);
+    for (const e of audioEvents) {
+      expect((e.payload as { audio: Uint8Array }).audio).toBeInstanceOf(Uint8Array);
+    }
+  });
+
+  // ── User speech recognition (from Kokoro TTS audio) ────────────────────
+
+  test("user speech events from real STT (Kokoro-generated audio)", async () => {
+    const { events } = await replayFixture("user-speech-recognition.json");
+
+    const types = events.map((e) => e.type);
+    expect(types).toContain("speech_started");
+    expect(types).toContain("speech_stopped");
+    expect(types).toContain("user_transcript");
+
+    // Verify the STT correctly transcribed the Kokoro audio
+    const transcripts = events.filter((e) => e.type === "user_transcript");
+    const texts = transcripts.map((e) => (e.payload as { text: string }).text);
+    expect(texts.some((t) => t.toLowerCase().includes("space"))).toBe(true);
+    expect(texts.some((t) => t.toLowerCase().includes("weather"))).toBe(true);
+  });
+
+  // ── Simple question flow ───────────────────────────────────────────────
+
+  test("simple question: greeting → user speech → agent response", async () => {
+    const { events } = await replayFixture("simple-question-sequence.json");
+
+    const types = events.map((e) => e.type);
+
+    // Session setup: session.updated is dropped; first events are ready + replyStarted
+    expect(types[0]).toBe("ready");
+    expect(types[1]).toBe("replyStarted");
+
+    // Greeting reply
+    expect(types).toContain("replyStarted");
+
+    // User speech recognition
+    expect(types).toContain("speech_started");
+    expect(types).toContain("user_transcript");
+
+    // Agent response
+    const agentTranscripts = events.filter((e) => e.type === "agent_transcript");
+    expect(agentTranscripts.length).toBe(2); // greeting + answer
+
+    // Two complete reply cycles (greeting + answer)
+    expect(types.filter((t) => t === "reply_done").length).toBe(2);
+  });
+
+  // ── Tool call flow ─────────────────────────────────────────────────────
+
+  test("tool call: user asks weather → tool_call event dispatched with parsed args", async () => {
+    const { events } = await replayFixture("tool-calls.json");
+
+    const toolCallEvents = events.filter((e) => e.type === "tool_call");
+    expect(toolCallEvents.length).toBe(1);
+    const tc = toolCallEvents[0]?.payload as {
+      toolCallId: string;
+      toolName: string;
+      args: Record<string, unknown>;
+    };
+    expect(tc.toolName).toBe("get_weather");
+    expect(tc.args.city).toBe("San Francisco");
+    expect(tc.toolCallId).toMatch(/^chatcmpl-tool-/);
+  });
+
+  test("tool call sequence: greeting → user speech → tool call → agent response", async () => {
+    const { events } = await replayFixture("tool-call-sequence.json");
+
+    const types = events.map((e) => e.type);
+
+    // Session setup: session.updated dropped; first events are ready + replyStarted
+    expect(types[0]).toBe("ready");
+    expect(types[1]).toBe("replyStarted");
+
+    // User speech was recognized
+    expect(types).toContain("user_transcript");
+    const userTx = events.find((e) => e.type === "user_transcript");
+    expect((userTx?.payload as { text: string }).text.toLowerCase()).toContain("weather");
+
+    // Tool was called
+    expect(types).toContain("tool_call");
+    const toolCall = events.find((e) => e.type === "tool_call");
+    expect((toolCall?.payload as { toolName: string }).toolName).toBe("get_weather");
+
+    // Agent responded after tool result
+    const agentTxs = events.filter((e) => e.type === "agent_transcript");
+    expect(agentTxs.length).toBe(2); // greeting + tool response
+    const toolResponse = agentTxs.at(-1)?.payload as { text: string };
+    expect(toolResponse.text.toLowerCase()).toContain("san francisco");
+  });
+});