@alexkroman1/aai 1.0.6 → 1.2.0

package/host/session.test.ts

@@ -83,14 +83,14 @@ describe("createS2sSession", () => {
     const { session, client } = setup();
     await session.start();
     session.onCancel();
-    expect(client.events).toContainEqual(expect.objectContaining({ type: "cancelled" }));
+    expect(client.events).toContainEvent("cancelled");
   });
 
   test("onReset clears state and emits reset event", async () => {
     const { session, client, mockHandle } = setup();
     await session.start();
     session.onReset();
-    expect(client.events).toContainEqual(expect.objectContaining({ type: "reset" }));
+    expect(client.events).toContainEvent("reset");
     expect(mockHandle.close).toHaveBeenCalled();
   });
 
@@ -119,10 +119,7 @@ describe("createS2sSession", () => {
     mockHandle._fire("event", { type: "user_transcript", text: "Hello there" });
     await flush();
 
-    expect(client.events).toContainEqual({
-      type: "user_transcript",
-      text: "Hello there",
-    });
+    expect(client.events).toContainEvent("user_transcript", { text: "Hello there" });
   });
 
   test("audio event forwards audio to client", async () => {
@@ -145,10 +142,7 @@ describe("createS2sSession", () => {
       _interrupted: false,
     });
 
-    expect(client.events).toContainEqual({
-      type: "agent_transcript",
-      text: "Full response",
-    });
+    expect(client.events).toContainEvent("agent_transcript", { text: "Full response" });
   });
 
   test("speech_started and speech_stopped events are forwarded", async () => {
@@ -158,8 +152,8 @@ describe("createS2sSession", () => {
     mockHandle._fire("event", { type: "speech_started" });
     mockHandle._fire("event", { type: "speech_stopped" });
 
-    expect(client.events).toContainEqual({ type: "speech_started" });
-    expect(client.events).toContainEqual({ type: "speech_stopped" });
+    expect(client.events).toContainEvent("speech_started");
+    expect(client.events).toContainEvent("speech_stopped");
   });
 
   test("reply_started resets tool call count", async () => {
@@ -177,7 +171,7 @@ describe("createS2sSession", () => {
     mockHandle._fire("event", { type: "reply_done" });
 
     expect(client.audioDoneCount).toBe(1);
-    expect(client.events).toContainEqual({ type: "reply_done" });
+    expect(client.events).toContainEvent("reply_done");
   });
 
   test("cancelled event emits cancelled", async () => {
@@ -186,7 +180,7 @@ describe("createS2sSession", () => {
 
     mockHandle._fire("event", { type: "cancelled" });
 
-    expect(client.events).toContainEqual({ type: "cancelled" });
+    expect(client.events).toContainEvent("cancelled");
   });
 
   test("error event emits error to client and closes handle", async () => {
@@ -195,11 +189,7 @@ describe("createS2sSession", () => {
 
     mockHandle._fire("error", new Error("Something broke"));
 
-    expect(client.events).toContainEqual({
-      type: "error",
-      code: "internal",
-      message: "Something broke",
-    });
+    expect(client.events).toContainEvent("error", { code: "internal", message: "Something broke" });
     expect(mockHandle.close).toHaveBeenCalled();
   });
 
@@ -227,14 +217,12 @@ describe("createS2sSession", () => {
       "session-1",
       expect.any(Array),
     );
-    expect(client.events).toContainEqual({
-      type: "tool_call",
+    expect(client.events).toContainEvent("tool_call", {
       toolCallId: "call-1",
       toolName: "my_tool",
       args: { key: "value" },
     });
-    expect(client.events).toContainEqual({
-      type: "tool_call_done",
+    expect(client.events).toContainEvent("tool_call_done", {
       toolCallId: "call-1",
       result: "tool-output",
     });
@@ -318,13 +306,7 @@ describe("createS2sSession", () => {
 
     await session.start();
 
-    expect(client.events).toContainEqual(
-      expect.objectContaining({
-        type: "error",
-        code: "internal",
-        message: "connect failed",
-      }),
-    );
+    expect(client.events).toContainEvent("error", { code: "internal", message: "connect failed" });
 
     spy.mockRestore();
   });
@@ -503,10 +485,6 @@ describe("createS2sSession", () => {
 
   // ─── Idle timeout tests ──────────────────────────────────────────────
 
-  function hasIdleTimeout(events: unknown[]): boolean {
-    return events.some((e) => (e as Record<string, unknown>).type === "idle_timeout");
-  }
-
   test("idle timeout fires after configured period of inactivity", async () => {
     vi.useFakeTimers();
     const { session, client, mockHandle } = setup({
@@ -519,7 +497,7 @@ describe("createS2sSession", () => {
     });
     await session.start();
     vi.advanceTimersByTime(10_000);
-    expect(hasIdleTimeout(client.events)).toBe(true);
+    expect(client.events).toContainEvent("idle_timeout");
     expect(mockHandle.close).toHaveBeenCalled();
     vi.useRealTimers();
   });
@@ -538,9 +516,9 @@ describe("createS2sSession", () => {
     vi.advanceTimersByTime(8000);
     session.onAudio(new Uint8Array([1, 2, 3]));
     vi.advanceTimersByTime(8000);
-    expect(hasIdleTimeout(client.events)).toBe(false);
+    expect(client.events).not.toContainEvent("idle_timeout");
     vi.advanceTimersByTime(2000);
-    expect(hasIdleTimeout(client.events)).toBe(true);
+    expect(client.events).toContainEvent("idle_timeout");
     vi.useRealTimers();
   });
 
@@ -556,7 +534,7 @@ describe("createS2sSession", () => {
     });
     await session.start();
     vi.advanceTimersByTime(600_000);
-    expect(hasIdleTimeout(client.events)).toBe(false);
+    expect(client.events).not.toContainEvent("idle_timeout");
     vi.useRealTimers();
   });
 
@@ -573,7 +551,7 @@ describe("createS2sSession", () => {
     await session.start();
     await session.stop();
     vi.advanceTimersByTime(20_000);
-    expect(hasIdleTimeout(client.events)).toBe(false);
+    expect(client.events).not.toContainEvent("idle_timeout");
     vi.useRealTimers();
   });
 
@@ -582,9 +560,9 @@ describe("createS2sSession", () => {
     const { session, client } = setup();
     await session.start();
     vi.advanceTimersByTime(240_000);
-    expect(hasIdleTimeout(client.events)).toBe(false);
+    expect(client.events).not.toContainEvent("idle_timeout");
     vi.advanceTimersByTime(60_000);
-    expect(hasIdleTimeout(client.events)).toBe(true);
+    expect(client.events).toContainEvent("idle_timeout");
     vi.useRealTimers();
   });
 });

package/host/tool-executor.test.ts

@@ -121,4 +121,40 @@ describe("executeToolCall", () => {
     expect(result).toBe(JSON.stringify({ error: 'Tool "slow" timed out after 30000ms' }));
     vi.useRealTimers();
   });
+
+  test("ctx.send calls the send callback", async () => {
+    const sends: Array<{ event: string; data: unknown }> = [];
+    const tool: ToolDef = {
+      description: "sends an event",
+      parameters: z.object({}),
+      execute: (_args, ctx) => {
+        ctx.send("game_state", { hp: 10 });
+        return "ok";
+      },
+    };
+    const result = await executeToolCall(
+      "sender",
+      {},
+      {
+        tool,
+        env: {},
+        send: (event, data) => sends.push({ event, data }),
+      },
+    );
+    expect(result).toBe("ok");
+    expect(sends).toEqual([{ event: "game_state", data: { hp: 10 } }]);
+  });
+
+  test("ctx.send is a no-op when no send callback provided", async () => {
+    const tool: ToolDef = {
+      description: "sends an event",
+      parameters: z.object({}),
+      execute: (_args, ctx) => {
+        ctx.send("test", {});
+        return "ok";
+      },
+    };
+    const result = await executeToolCall("sender", {}, { tool, env: {} });
+    expect(result).toBe("ok");
+  });
 });

package/host/tool-executor.ts

@@ -27,6 +27,7 @@ export type ExecuteToolCallOptions = {
   kv?: Kv | undefined;
   messages?: readonly Message[] | undefined;
   logger?: Logger | undefined;
+  send?: ((event: string, data: unknown) => void) | undefined;
 };
 
 function buildToolContext(opts: ExecuteToolCallOptions): ToolContext {
@@ -40,6 +41,9 @@ function buildToolContext(opts: ExecuteToolCallOptions): ToolContext {
     },
     messages: messages ?? [],
     sessionId: sessionId ?? "",
+    send(event: string, data: unknown): void {
+      opts.send?.(event, data);
+    },
   };
 }
 

package/host/ws-handler.ts

@@ -47,6 +47,8 @@ export type WsSessionOptions = {
   onClose?: () => void;
   /** Callback invoked with the session ID after session cleanup. */
   onSessionEnd?: (sessionId: string) => void;
+  /** Callback invoked with the session ID and client sink after session setup. */
+  onSinkCreated?: (sessionId: string, sink: ClientSink) => void;
   /** Logger instance. Defaults to console. */
   logger?: Logger;
   /** Timeout in ms for session.start(). Defaults to 10 000 (10s). */
@@ -184,6 +186,7 @@ export function wireSessionSocket(ws: SessionWebSocket, opts: WsSessionOptions):
     const client = createClientSink(ws, log);
     session = opts.createSession(sessionId, client);
     sessions.set(sessionId, session);
+    opts.onSinkCreated?.(sessionId, client);
 
     // Send config immediately — zero RTT. Include sessionId so the client
     // can reconnect with ?sessionId=<id> to resume a persisted session.

package/index.ts

@@ -8,6 +8,7 @@
 
 // biome-ignore-all lint/performance/noReExportAll: barrel file by design
 
+export * from "./sdk/allowed-hosts.ts";
 export * from "./sdk/constants.ts";
 export * from "./sdk/define.ts";
 export * from "./sdk/kv.ts";

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@alexkroman1/aai",
-  "version": "1.0.6",
+  "version": "1.2.0",
   "type": "module",
   "exports": {
     ".": {

package/sdk/__snapshots__/exports.test.ts.snap

@@ -0,0 +1,79 @@
+// Vitest Snapshot v1, https://vitest.dev/guide/snapshot.html
+
+exports[`export surface stability > @alexkroman1/aai main export 1`] = `
+[
+  "AGENT_CSP",
+  "BuiltinToolSchema",
+  "DEFAULT_GREETING",
+  "DEFAULT_IDLE_TIMEOUT_MS",
+  "DEFAULT_MAX_HISTORY",
+  "DEFAULT_SESSION_START_TIMEOUT_MS",
+  "DEFAULT_SHUTDOWN_TIMEOUT_MS",
+  "DEFAULT_STT_SAMPLE_RATE",
+  "DEFAULT_SYSTEM_PROMPT",
+  "DEFAULT_TTS_SAMPLE_RATE",
+  "FETCH_TIMEOUT_MS",
+  "MAX_HTML_BYTES",
+  "MAX_MESSAGE_BUFFER_SIZE",
+  "MAX_PAGE_CHARS",
+  "MAX_TOOL_RESULT_CHARS",
+  "MAX_VALUE_SIZE",
+  "MAX_WS_PAYLOAD_BYTES",
+  "RUN_CODE_TIMEOUT_MS",
+  "TOOL_EXECUTION_TIMEOUT_MS",
+  "ToolChoiceSchema",
+  "WS_OPEN",
+  "agent",
+  "errorDetail",
+  "errorMessage",
+  "matchesAllowedHost",
+  "parseWsUpgradeParams",
+  "tool",
+  "toolError",
+  "validateAllowedHostPattern",
+]
+`;
+
+exports[`export surface stability > @alexkroman1/aai/manifest export 1`] = `
+[
+  "AgentConfigSchema",
+  "EMPTY_PARAMS",
+  "ToolSchemaSchema",
+  "agentToolsToSchemas",
+  "toAgentConfig",
+]
+`;
+
+exports[`export surface stability > @alexkroman1/aai/protocol export 1`] = `
+[
+  "ClientEventSchema",
+  "ClientMessageSchema",
+  "KvDelSchema",
+  "KvGetSchema",
+  "KvRequestSchema",
+  "KvSetSchema",
+  "ReadyConfigSchema",
+  "ServerMessageSchema",
+  "SessionErrorCodeSchema",
+  "buildReadyConfig",
+  "lenientParse",
+]
+`;
+
+exports[`export surface stability > @alexkroman1/aai/runtime export 1`] = `
+[
+  "DEFAULT_S2S_CONFIG",
+  "_internals",
+  "buildCtx",
+  "consoleLogger",
+  "createRuntime",
+  "createS2sSession",
+  "createServer",
+  "createUnstorageKv",
+  "executeInIsolate",
+  "executeToolCall",
+  "jsonLogger",
+  "resolveAllBuiltins",
+  "wireSessionSocket",
+]
+`;

package/sdk/__snapshots__/schema-shapes.test.ts.snap

@@ -0,0 +1,187 @@
+// Vitest Snapshot v1, https://vitest.dev/guide/snapshot.html
+
+exports[`manifest schema shapes > AgentConfigSchema shape 1`] = `
+[
+  "builtinTools",
+  "greeting",
+  "idleTimeoutMs",
+  "maxSteps",
+  "name",
+  "sttPrompt",
+  "systemPrompt",
+  "toolChoice",
+]
+`;
+
+exports[`manifest schema shapes > ToolSchemaSchema shape 1`] = `
+[
+  "description",
+  "name",
+  "parameters",
+]
+`;
+
+exports[`protocol schema shapes > ClientEventSchema option shapes 1`] = `
+{
+  "agent_transcript": [
+    "text",
+    "type",
+  ],
+  "cancelled": [
+    "type",
+  ],
+  "custom_event": [
+    "data",
+    "event",
+    "type",
+  ],
+  "error": [
+    "code",
+    "message",
+    "type",
+  ],
+  "idle_timeout": [
+    "type",
+  ],
+  "reply_done": [
+    "type",
+  ],
+  "reset": [
+    "type",
+  ],
+  "speech_started": [
+    "type",
+  ],
+  "speech_stopped": [
+    "type",
+  ],
+  "tool_call": [
+    "args",
+    "toolCallId",
+    "toolName",
+    "type",
+  ],
+  "tool_call_done": [
+    "result",
+    "toolCallId",
+    "type",
+  ],
+  "user_transcript": [
+    "text",
+    "turnOrder",
+    "type",
+  ],
+}
+`;
+
+exports[`protocol schema shapes > ClientMessageSchema option shapes 1`] = `
+{
+  "audio_ready": [
+    "type",
+  ],
+  "cancel": [
+    "type",
+  ],
+  "history": [
+    "messages",
+    "type",
+  ],
+  "reset": [
+    "type",
+  ],
+}
+`;
+
+exports[`protocol schema shapes > KvDelSchema shape 1`] = `
+[
+  "key",
+  "op",
+]
+`;
+
+exports[`protocol schema shapes > KvGetSchema shape 1`] = `
+[
+  "key",
+  "op",
+]
+`;
+
+exports[`protocol schema shapes > KvSetSchema shape 1`] = `
+[
+  "expireIn",
+  "key",
+  "op",
+  "value",
+]
+`;
+
+exports[`protocol schema shapes > ReadyConfigSchema shape 1`] = `
+[
+  "audioFormat",
+  "sampleRate",
+  "ttsSampleRate",
+]
+`;
+
+exports[`protocol schema shapes > ServerMessageSchema option shapes 1`] = `
+{
+  "agent_transcript": [
+    "text",
+    "type",
+  ],
+  "audio_done": [
+    "type",
+  ],
+  "cancelled": [
+    "type",
+  ],
+  "config": [
+    "audioFormat",
+    "sampleRate",
+    "sessionId",
+    "ttsSampleRate",
+    "type",
+  ],
+  "custom_event": [
+    "data",
+    "event",
+    "type",
+  ],
+  "error": [
+    "code",
+    "message",
+    "type",
+  ],
+  "idle_timeout": [
+    "type",
+  ],
+  "reply_done": [
+    "type",
+  ],
+  "reset": [
+    "type",
+  ],
+  "speech_started": [
+    "type",
+  ],
+  "speech_stopped": [
+    "type",
+  ],
+  "tool_call": [
+    "args",
+    "toolCallId",
+    "toolName",
+    "type",
+  ],
+  "tool_call_done": [
+    "result",
+    "toolCallId",
+    "type",
+  ],
+  "user_transcript": [
+    "text",
+    "turnOrder",
+    "type",
+  ],
+}
+`;

package/sdk/_test-matchers.test.ts

@@ -0,0 +1,75 @@
+// Copyright 2025 the AAI authors. MIT license.
+
+import { describe, expect, it } from "vitest";
+import "./_test-matchers.ts";
+
+describe("toBeValidClientEvent", () => {
+  it("passes for a valid event", () => {
+    expect({ type: "speech_started" }).toBeValidClientEvent();
+  });
+
+  it("passes for a valid event with fields", () => {
+    expect({
+      type: "user_transcript",
+      text: "hello world",
+    }).toBeValidClientEvent();
+  });
+
+  it("fails for an invalid event", () => {
+    expect(() => {
+      expect({ type: "not_a_real_event" }).toBeValidClientEvent();
+    }).toThrow(/expected value to be a valid ClientEvent/);
+  });
+
+  it("fails for a missing type field", () => {
+    expect(() => {
+      expect({ text: "no type" }).toBeValidClientEvent();
+    }).toThrow(/expected value to be a valid ClientEvent/);
+  });
+});
+
+describe("toContainEvent", () => {
+  const events = [
+    { type: "speech_started" },
+    { type: "user_transcript", text: "hello" },
+    { type: "tool_call", toolCallId: "tc1", toolName: "search", args: { q: "test" } },
+    { type: "reply_done" },
+  ];
+
+  it("finds a matching event by type", () => {
+    expect(events).toContainEvent("speech_started");
+  });
+
+  it("finds a matching event with fields", () => {
+    expect(events).toContainEvent("tool_call", { toolName: "search" });
+  });
+
+  it("matches a subset of fields", () => {
+    expect(events).toContainEvent("tool_call", {
+      toolName: "search",
+      args: { q: "test" },
+    });
+  });
+
+  it("fails when event type is not found", () => {
+    expect(() => {
+      expect(events).toContainEvent("cancelled");
+    }).toThrow(/expected array to contain event of type "cancelled"/);
+  });
+
+  it("fails when fields do not match", () => {
+    expect(() => {
+      expect(events).toContainEvent("tool_call", { toolName: "visit_webpage" });
+    }).toThrow(/expected array to contain event of type "tool_call"/);
+  });
+
+  it("fails when received value is not an array", () => {
+    expect(() => {
+      expect("not-an-array").toContainEvent("speech_started");
+    }).toThrow(/expected an array of events/);
+  });
+
+  it("supports .not negation", () => {
+    expect(events).not.toContainEvent("cancelled");
+  });
+});

package/sdk/_test-matchers.ts

@@ -0,0 +1,73 @@
+// Copyright 2025 the AAI authors. MIT license.
+/**
+ * Custom Vitest matchers for AAI domain types.
+ *
+ * Registered via `expect.extend()` — add this file to the vitest `setupFiles`
+ * for the `aai` project so matchers are available in every test.
+ */
+
+import { isDeepStrictEqual } from "node:util";
+import { expect } from "vitest";
+import { ClientEventSchema } from "./protocol.ts";
+
+type MatcherResult = { pass: boolean; message: () => string };
+
+// ─── Matcher implementations ────────────────────────────────────────────────
+
+function toBeValidClientEvent(received: unknown): MatcherResult {
+  const result = ClientEventSchema.safeParse(received);
+  return {
+    pass: result.success,
+    message: () =>
+      result.success
+        ? "expected value NOT to be a valid ClientEvent, but it parsed successfully"
+        : `expected value to be a valid ClientEvent\n\nZod errors:\n${result.error.issues.map((i) => `  - ${i.path.join(".")}: ${i.message}`).join("\n")}`,
+  };
+}
+
+function toContainEvent(
+  received: unknown,
+  type: string,
+  fields?: Record<string, unknown>,
+): MatcherResult {
+  if (!Array.isArray(received)) {
+    return {
+      pass: false,
+      message: () => `expected an array of events, but received ${typeof received}`,
+    };
+  }
+
+  const match = received.some((event: Record<string, unknown>) => {
+    if (event?.type !== type) return false;
+    if (!fields) return true;
+    return Object.entries(fields).every(([key, value]) => isDeepStrictEqual(event[key], value));
+  });
+
+  return {
+    pass: match,
+    message: () =>
+      match
+        ? `expected array NOT to contain event of type "${type}"${fields ? ` with fields ${JSON.stringify(fields)}` : ""}`
+        : `expected array to contain event of type "${type}"${fields ? ` with fields ${JSON.stringify(fields)}` : ""}\n\nReceived event types: [${received.map((e: Record<string, unknown>) => `"${e?.type}"`).join(", ")}]`,
+  };
+}
+
+// ─── Register matchers ──────────────────────────────────────────────────────
+
+expect.extend({
+  toBeValidClientEvent,
+  toContainEvent,
+});
+
+// ─── Type augmentation ──────────────────────────────────────────────────────
+
+declare module "vitest" {
+  interface Assertion<T> {
+    toBeValidClientEvent(): void;
+    toContainEvent(type: string, fields?: Record<string, unknown>): void;
+  }
+  interface AsymmetricMatchersContaining {
+    toBeValidClientEvent(): void;
+    toContainEvent(type: string, fields?: Record<string, unknown>): void;
+  }
+}