@alexkroman1/aai 0.12.3 → 1.0.2

package/sdk/protocol-snapshot.test.ts

@@ -0,0 +1,199 @@
+// Copyright 2025 the AAI authors. MIT license.
+/**
+ * Wire format snapshot tests for the WebSocket protocol.
+ *
+ * These ensure that changes to Zod schemas in protocol.ts don't
+ * accidentally alter the wire format. If a snapshot breaks, it
+ * signals a potentially breaking protocol change.
+ */
+import { describe, expect, test } from "vitest";
+import {
+  DEFAULT_STT_SAMPLE_RATE,
+  DEFAULT_TTS_SAMPLE_RATE,
+  TOOL_EXECUTION_TIMEOUT_MS,
+} from "./constants.ts";
+import type { ClientEvent, ClientMessage, ServerMessage } from "./protocol.ts";
+import {
+  AUDIO_FORMAT,
+  ClientEventSchema,
+  ClientMessageSchema,
+  KvRequestSchema,
+  SessionErrorCodeSchema,
+} from "./protocol.ts";
+
+// ── Constants ────────────────────────────────────────────────────────────
+
+describe("protocol constants", () => {
+  test("audio format", () => {
+    expect(AUDIO_FORMAT).toMatchInlineSnapshot(`"pcm16"`);
+  });
+
+  test("sample rates", () => {
+    expect(DEFAULT_STT_SAMPLE_RATE).toMatchInlineSnapshot("16000");
+    expect(DEFAULT_TTS_SAMPLE_RATE).toMatchInlineSnapshot("24000");
+  });
+
+  test("timeout constants", () => {
+    expect(TOOL_EXECUTION_TIMEOUT_MS).toMatchInlineSnapshot("30000");
+  });
+
+  test("error codes", () => {
+    expect(SessionErrorCodeSchema.options).toMatchInlineSnapshot(`
+      [
+        "stt",
+        "llm",
+        "tts",
+        "tool",
+        "protocol",
+        "connection",
+        "audio",
+        "internal",
+      ]
+    `);
+  });
+});
+
+// ── Server → Client events (ClientEventSchema) ──────────────────────────
+
+describe("server→client event wire format", () => {
+  const valid: [string, ClientEvent][] = [
+    ["speech_started", { type: "speech_started" }],
+    ["speech_stopped", { type: "speech_stopped" }],
+    ["user_transcript", { type: "user_transcript", text: "hello" }],
+    ["user_transcript (with order)", { type: "user_transcript", text: "hello", turnOrder: 1 }],
+    ["agent_transcript", { type: "agent_transcript", text: "response" }],
+    [
+      "tool_call",
+      {
+        type: "tool_call",
+        toolCallId: "tc1",
+        toolName: "web_search",
+        args: { query: "weather" },
+      },
+    ],
+    ["tool_call_done", { type: "tool_call_done", toolCallId: "tc1", result: "72F" }],
+    ["reply_done", { type: "reply_done" }],
+    ["cancelled", { type: "cancelled" }],
+    ["reset", { type: "reset" }],
+    ["idle_timeout", { type: "idle_timeout" }],
+    ["error", { type: "error", code: "stt", message: "Speech recognition failed" }],
+  ];
+
+  test.each(valid)("%s parses successfully", (_label, event) => {
+    const result = ClientEventSchema.safeParse(event);
+    expect(result.success).toBe(true);
+  });
+
+  test("rejects unknown event type", () => {
+    expect(ClientEventSchema.safeParse({ type: "bogus" }).success).toBe(false);
+  });
+
+  test("rejects invalid error code", () => {
+    expect(
+      ClientEventSchema.safeParse({ type: "error", code: "invalid_code", message: "x" }).success,
+    ).toBe(false);
+  });
+
+  test("rejects tool_call_done with oversized result", () => {
+    expect(
+      ClientEventSchema.safeParse({
+        type: "tool_call_done",
+        toolCallId: "tc1",
+        result: "x".repeat(4001),
+      }).success,
+    ).toBe(false);
+  });
+});
+
+// ── Client → Server messages (ClientMessageSchema) ──────────────────────
+
+describe("client→server message wire format", () => {
+  const valid: [string, ClientMessage][] = [
+    ["audio_ready", { type: "audio_ready" }],
+    ["cancel", { type: "cancel" }],
+    ["reset", { type: "reset" }],
+    [
+      "history",
+      {
+        type: "history",
+        messages: [
+          { role: "user", content: "Hello" },
+          { role: "assistant", content: "Hi" },
+        ],
+      },
+    ],
+  ];
+
+  test.each(valid)("%s parses successfully", (_label, msg) => {
+    const result = ClientMessageSchema.safeParse(msg);
+    expect(result.success).toBe(true);
+  });
+
+  test("rejects unknown message type", () => {
+    expect(ClientMessageSchema.safeParse({ type: "bogus" }).success).toBe(false);
+  });
+
+  test("rejects history with invalid role", () => {
+    expect(
+      ClientMessageSchema.safeParse({
+        type: "history",
+        messages: [{ role: "system", text: "nope" }],
+      }).success,
+    ).toBe(false);
+  });
+
+  test("rejects history exceeding 200 messages", () => {
+    const messages = Array.from({ length: 201 }, (_, i) => ({
+      role: "user" as const,
+      text: `msg ${i}`,
+    }));
+    expect(ClientMessageSchema.safeParse({ type: "history", messages }).success).toBe(false);
+  });
+});
+
+// ── ServerMessage union (type check) ────────────────────────────────────
+
+describe("ServerMessage type covers all variants", () => {
+  test("config message shape", () => {
+    const msg: ServerMessage = {
+      type: "config",
+      audioFormat: "pcm16",
+      sampleRate: 16_000,
+      ttsSampleRate: 24_000,
+    };
+    expect(msg.type).toBe("config");
+  });
+
+  test("audio_done message shape", () => {
+    const msg: ServerMessage = { type: "audio_done" };
+    expect(msg.type).toBe("audio_done");
+  });
+
+  test("ClientEvent is a valid ServerMessage", () => {
+    const msg: ServerMessage = { type: "speech_started" };
+    expect(msg.type).toBe("speech_started");
+  });
+});
+
+// ── KvRequestSchema ─────────────────────────────────────────────────────
+
+describe("KvRequest wire format", () => {
+  const valid = [
+    ["get", { op: "get", key: "k1" }],
+    ["set", { op: "set", key: "k1", value: "v1" }],
+    ["set with expireIn", { op: "set", key: "k1", value: "v1", expireIn: 60_000 }],
+    ["del", { op: "del", key: "k1" }],
+  ] as const;
+
+  test.each(valid)("%s parses successfully", (_label, req) => {
+    expect(KvRequestSchema.safeParse(req).success).toBe(true);
+  });
+
+  test("rejects unknown op", () => {
+    expect(KvRequestSchema.safeParse({ op: "update", key: "k1" }).success).toBe(false);
+  });
+
+  test("rejects empty key for get", () => {
+    expect(KvRequestSchema.safeParse({ op: "get", key: "" }).success).toBe(false);
+  });
+});

package/sdk/protocol.test.ts

@@ -0,0 +1,170 @@
+import { describe, expect, test } from "vitest";
+import {
+  DEFAULT_STT_SAMPLE_RATE,
+  DEFAULT_TTS_SAMPLE_RATE,
+  TOOL_EXECUTION_TIMEOUT_MS,
+} from "./constants.ts";
+import {
+  AUDIO_FORMAT,
+  buildReadyConfig,
+  ClientEventSchema,
+  ClientMessageSchema,
+  KvRequestSchema,
+  SessionErrorCodeSchema,
+} from "./protocol.ts";
+
+describe("protocol constants", () => {
+  test("DEFAULT_STT_SAMPLE_RATE is 16000", () => {
+    expect(DEFAULT_STT_SAMPLE_RATE).toBe(16_000);
+  });
+
+  test("DEFAULT_TTS_SAMPLE_RATE is 24000", () => {
+    expect(DEFAULT_TTS_SAMPLE_RATE).toBe(24_000);
+  });
+
+  test('AUDIO_FORMAT is "pcm16"', () => {
+    expect(AUDIO_FORMAT).toBe("pcm16");
+  });
+
+  test("TOOL_EXECUTION_TIMEOUT_MS is 30000", () => {
+    expect(TOOL_EXECUTION_TIMEOUT_MS).toBe(30_000);
+  });
+});
+
+describe("KvRequestSchema", () => {
+  test("accepts valid get request", () => {
+    const result = KvRequestSchema.safeParse({
+      op: "get",
+      key: "my-key",
+    });
+    expect(result.success).toBe(true);
+  });
+
+  test("accepts valid set request with ttl", () => {
+    const result = KvRequestSchema.safeParse({
+      op: "set",
+      key: "my-key",
+      value: "my-value",
+      ttl: 3600,
+    });
+    expect(result.success).toBe(true);
+  });
+
+  // The Kv interface accepts `value: unknown` — ensure the schema matches.
+  test.each([
+    { nested: true },
+    [1, 2, 3],
+    null,
+    42,
+  ])("set accepts non-string value: %j", (value) => {
+    const result = KvRequestSchema.safeParse({ op: "set", key: "k", value });
+    expect(result.success).toBe(true);
+  });
+
+  test("rejects empty key on get", () => {
+    const result = KvRequestSchema.safeParse({
+      op: "get",
+      key: "",
+    });
+    expect(result.success).toBe(false);
+  });
+});
+
+describe("SessionErrorCodeSchema", () => {
+  test.each([
+    "stt",
+    "llm",
+    "tts",
+    "tool",
+    "protocol",
+    "connection",
+    "audio",
+    "internal",
+  ])("accepts valid code: %s", (code) => {
+    expect(SessionErrorCodeSchema.safeParse(code).success).toBe(true);
+  });
+
+  test("rejects invalid code", () => {
+    expect(SessionErrorCodeSchema.safeParse("not_a_real_code").success).toBe(false);
+  });
+});
+
+describe("ClientEventSchema", () => {
+  test("accepts speech_started", () => {
+    const result = ClientEventSchema.safeParse({ type: "speech_started" });
+    expect(result.success).toBe(true);
+  });
+
+  test("accepts user_transcript", () => {
+    const result = ClientEventSchema.safeParse({
+      type: "user_transcript",
+      text: "hello world",
+    });
+    expect(result.success).toBe(true);
+  });
+
+  test("accepts error event", () => {
+    const result = ClientEventSchema.safeParse({
+      type: "error",
+      code: "internal",
+      message: "something went wrong",
+    });
+    expect(result.success).toBe(true);
+  });
+
+  test("rejects unknown type", () => {
+    const result = ClientEventSchema.safeParse({
+      type: "unknown_event_type",
+    });
+    expect(result.success).toBe(false);
+  });
+});
+
+describe("ClientMessageSchema", () => {
+  test("accepts audio_ready", () => {
+    const result = ClientMessageSchema.safeParse({ type: "audio_ready" });
+    expect(result.success).toBe(true);
+  });
+
+  test("accepts cancel", () => {
+    const result = ClientMessageSchema.safeParse({ type: "cancel" });
+    expect(result.success).toBe(true);
+  });
+
+  test("accepts reset", () => {
+    const result = ClientMessageSchema.safeParse({ type: "reset" });
+    expect(result.success).toBe(true);
+  });
+
+  test("accepts history with messages", () => {
+    const result = ClientMessageSchema.safeParse({
+      type: "history",
+      messages: [{ role: "user", content: "hello" }],
+    });
+    expect(result.success).toBe(true);
+  });
+
+  test("rejects unknown type", () => {
+    const result = ClientMessageSchema.safeParse({
+      type: "unknown_message_type",
+    });
+    expect(result.success).toBe(false);
+  });
+});
+
+describe("buildReadyConfig", () => {
+  test("builds config from sample rates", () => {
+    const config = buildReadyConfig({ inputSampleRate: 16_000, outputSampleRate: 24_000 });
+    expect(config).toEqual({
+      audioFormat: AUDIO_FORMAT,
+      sampleRate: 16_000,
+      ttsSampleRate: 24_000,
+    });
+  });
+
+  test("uses custom sample rates", () => {
+    const config = buildReadyConfig({ inputSampleRate: 8000, outputSampleRate: 48_000 });
+    expect(config.sampleRate).toBe(8000);
+    expect(config.ttsSampleRate).toBe(48_000);
+  });
+});

package/sdk/protocol.ts

@@ -0,0 +1,223 @@
+// Copyright 2025 the AAI authors. MIT license.
+/**
+ * WebSocket wire-format types shared by server and client.
+ *
+ * Note: this module is for internal use only and should not be used directly.
+ */
+
+import { z } from "zod";
+
+import { MAX_TOOL_RESULT_CHARS } from "./constants.ts";
+
+/**
+ * Audio codec identifier used in the wire protocol.
+ *
+ * All audio frames are 16-bit signed PCM, little-endian, mono.
+ */
+export const AUDIO_FORMAT = "pcm16";
+
+/**
+ * Minimal envelope schema for two-phase message parsing.
+ *
+ * When a strict schema (ServerMessageSchema / ClientMessageSchema) rejects a
+ * message, this schema determines whether the message is a valid but
+ * *unrecognised* type (safe to ignore during rolling upgrades) or genuinely
+ * malformed (should be warned about).
+ */
+export const MessageEnvelopeSchema = z.object({ type: z.string() }).passthrough();
+
+/**
+ * Two-phase message parse: tries the strict schema first, then falls back to
+ * the envelope to distinguish unknown-but-valid types (safe to ignore during
+ * rolling upgrades) from genuinely malformed messages.
+ *
+ * Return value when `ok: false`:
+ * - `malformed: true` — message doesn't even have a `{ type: string }` shape;
+ *   likely corrupt data, should warn
+ * - `malformed: false` — has a valid `type` field but the type is unrecognised;
+ *   safe to ignore (e.g. new message type from a newer server version)
+ */
+export function lenientParse<T>(
+  schema: z.ZodType<T>,
+  json: unknown,
+): { ok: true; data: T } | { ok: false; malformed: boolean; error: string } {
+  const result = schema.safeParse(json);
+  if (result.success) return { ok: true, data: result.data };
+  const malformed = !MessageEnvelopeSchema.safeParse(json).success;
+  return { ok: false, malformed, error: result.error.message };
+}
+
+/** Zod schema for the KV "get" operation. */
+export const KvGetSchema = z.object({ op: z.literal("get"), key: z.string().min(1) });
+
+/** Zod schema for the KV "set" operation. */
+export const KvSetSchema = z.object({
+  op: z.literal("set"),
+  key: z.string().min(1),
+  value: z.unknown(),
+  /** Time-to-live in **milliseconds**. */
+  expireIn: z.number().int().positive().optional(),
+});
+
+/** Zod schema for the KV "del" operation. */
+export const KvDelSchema = z.object({ op: z.literal("del"), key: z.string().min(1) });
+
+/** Zod schema for KV operation requests from the worker to the host. */
+export const KvRequestSchema = z.discriminatedUnion("op", [KvGetSchema, KvSetSchema, KvDelSchema]);
+
+/** KV operation request — discriminated union on the `op` field. */
+export type KvRequest = z.infer<typeof KvRequestSchema>;
+
+// ─── Error codes ───────────────────────────────────────────────────────────
+
+/**
+ * Zod schema for session error codes.
+ * @public
+ */
+export const SessionErrorCodeSchema = z.enum([
+  "stt",
+  "llm",
+  "tts",
+  "tool",
+  "protocol",
+  "connection",
+  "audio",
+  "internal",
+]);
+
+/**
+ * Error codes for categorizing session errors on the wire.
+ *
+ * @public
+ */
+export type SessionErrorCode = z.infer<typeof SessionErrorCodeSchema>;
+
+// ─── Client events ─────────────────────────────────────────────────────────
+
+/** Helper: simple event with only a type field. */
+const ev = <T extends string>(t: T) => z.object({ type: z.literal(t) });
+
+const turnOrder = z.number().int().nonnegative().optional();
+
+/** Zod schema for {@link ClientEvent}. */
+export const ClientEventSchema = z.discriminatedUnion("type", [
+  ev("speech_started"),
+  ev("speech_stopped"),
+  z.object({
+    type: z.literal("user_transcript"),
+    text: z.string(),
+    turnOrder,
+  }),
+  z.object({
+    type: z.literal("agent_transcript"),
+    text: z.string(),
+  }),
+  z.object({
+    type: z.literal("tool_call"),
+    toolCallId: z.string(),
+    toolName: z.string(),
+    args: z.record(z.string(), z.unknown()),
+  }),
+  z.object({
+    type: z.literal("tool_call_done"),
+    toolCallId: z.string(),
+    result: z.string().max(MAX_TOOL_RESULT_CHARS),
+  }),
+  ev("reply_done"),
+  ev("cancelled"),
+  ev("reset"),
+  ev("idle_timeout"),
+  z.object({ type: z.literal("error"), code: SessionErrorCodeSchema, message: z.string() }),
+]);
+
+/** Discriminated union of all server→client session events. */
+export type ClientEvent = z.infer<typeof ClientEventSchema>;
+
+/**
+ * Typed interface for pushing session events to a connected client.
+ *
+ * For WebSocket sessions this sends JSON text frames and binary audio frames.
+ */
+export interface ClientSink {
+  /** Whether the underlying connection is open and accepting calls. */
+  readonly open: boolean;
+  /** Push a session event to the client. */
+  event(e: ClientEvent): void;
+  /** Send a single TTS audio chunk to the client. */
+  playAudioChunk(chunk: Uint8Array): void;
+  /** Signal that TTS audio is complete. */
+  playAudioDone(): void;
+}
+
+// ─── WebSocket message types ────────────────────────────────────────────────
+
+/** Supported audio formats for the wire protocol. */
+export type AudioFormatId = "pcm16";
+
+/** Zod schema for {@link ReadyConfig}. */
+export const ReadyConfigSchema = z.object({
+  audioFormat: z.enum(["pcm16"]),
+  sampleRate: z.number().int().positive(),
+  ttsSampleRate: z.number().int().positive(),
+});
+
+/** Protocol-level session config returned to the client on connect. */
+export type ReadyConfig = z.infer<typeof ReadyConfigSchema>;
+
+/** Zod schema for server→client text messages. */
+export const ServerMessageSchema = z.discriminatedUnion("type", [
+  z.object({
+    type: z.literal("config"),
+    audioFormat: z.string(),
+    sampleRate: z.number(),
+    ttsSampleRate: z.number(),
+    /** Session ID for this connection. Clients can reconnect with
+     *  `?sessionId=<id>` to resume a persisted session. */
+    sessionId: z.string().optional(),
+  }),
+  ev("audio_done"),
+  ...ClientEventSchema.options,
+]);
+
+/** Server→client text messages (binary frames carry raw PCM16 audio). */
+export type ServerMessage = z.infer<typeof ServerMessageSchema>;
+
+/** Zod schema for client→server text messages. */
+export const ClientMessageSchema = z.discriminatedUnion("type", [
+  ev("audio_ready"),
+  ev("cancel"),
+  ev("reset"),
+  z.object({
+    type: z.literal("history"),
+    messages: z
+      .array(z.object({ role: z.enum(["user", "assistant"]), content: z.string().max(100_000) }))
+      .max(200),
+  }),
+]);
+
+/** Client→server text messages (binary frames carry raw PCM16 audio). */
+export type ClientMessage = z.infer<typeof ClientMessageSchema>;
+
+// ─── Ready config builder ───────────────────────────────────────────────────
+
+/** Build the protocol-level session config from S2S sample rates. */
+export function buildReadyConfig(s2sConfig: {
+  inputSampleRate: number;
+  outputSampleRate: number;
+}): ReadyConfig {
+  return {
+    audioFormat: AUDIO_FORMAT,
+    sampleRate: s2sConfig.inputSampleRate,
+    ttsSampleRate: s2sConfig.outputSampleRate,
+  };
+}
+
+// ─── Worker RPC interfaces ─────────────────────────────────────────────────
+
+/** Zod schema for {@link TurnConfig}. */
+export const TurnConfigSchema = z.object({
+  maxSteps: z.number().int().positive().optional(),
+});
+
+/** Combined turn configuration resolved from the worker before a turn starts. */
+export type TurnConfig = z.infer<typeof TurnConfigSchema>;