@alexkroman1/aai 0.3.0

package/sdk/kv.ts

@@ -0,0 +1,183 @@
+// Copyright 2025 the AAI authors. MIT license.
+/**
+ * Key-value storage interface and in-memory implementation.
+ *
+ * @module
+ */
+
+/**
+ * A single key-value entry returned by {@linkcode Kv.list}.
+ *
+ * @typeParam T The type of the stored value. Defaults to `unknown`.
+ */
+export type KvEntry<T = unknown> = {
+  /** The key under which the value is stored. */
+  key: string;
+  /** The deserialized value. */
+  value: T;
+};
+
+/**
+ * Options for listing keys from the KV store.
+ *
+ * Used with {@linkcode Kv.list} to control result ordering and pagination.
+ */
+export type KvListOptions = {
+  /** Maximum number of entries to return. */
+  limit?: number;
+  /** Return entries in reverse key order. */
+  reverse?: boolean;
+};
+
+/**
+ * Async key-value store interface used by agents.
+ *
+ * Agents access the KV store via {@linkcode ToolContext.kv} or
+ * {@linkcode HookContext.kv}. Values are JSON-serialized and stored as
+ * strings with an optional TTL.
+ *
+ * @example
+ * ```ts
+ * // Inside a tool execute function:
+ * const myTool = {
+ *   description: "Save and retrieve data",
+ *   execute: async (_args: unknown, ctx: { kv: Kv }) => {
+ *     await ctx.kv.set("user:name", "Alice", { expireIn: 60_000 });
+ *     const name = await ctx.kv.get<string>("user:name");
+ *     return name; // "Alice"
+ *   },
+ * };
+ * ```
+ */
+export type Kv = {
+  /**
+   * Get a value by key, or `null` if not found.
+   *
+   * @typeParam T The expected type of the stored value.
+   * @param key The key to look up.
+   * @returns The deserialized value, or `null` if the key does not exist
+   *   or has expired.
+   */
+  get<T = unknown>(key: string): Promise<T | null>;
+  /**
+   * Set a value, optionally with a TTL in milliseconds.
+   *
+   * @param key The key to store the value under.
+   * @param value The value to store. Must be JSON-serializable.
+   * @param options Optional settings.
+   * @param options.expireIn Time-to-live in milliseconds. The entry is
+   *   automatically removed after this duration.
+   * @throws {Error} If the serialized value exceeds 65,536 bytes.
+   */
+  set(key: string, value: unknown, options?: { expireIn?: number }): Promise<void>;
+  /**
+   * Delete a key.
+   *
+   * @param key The key to remove. No-op if the key does not exist.
+   */
+  delete(key: string): Promise<void>;
+  /**
+   * List entries whose keys start with the given prefix.
+   *
+   * Results are sorted by key in ascending lexicographic order by default.
+   *
+   * @typeParam T The expected type of the stored values.
+   * @param prefix Key prefix to filter by. Use `""` to list all entries.
+   * @param options Optional pagination and ordering settings.
+   * @returns An array of matching {@linkcode KvEntry} objects.
+   */
+  list<T = unknown>(prefix: string, options?: KvListOptions): Promise<KvEntry<T>[]>;
+};
+
+export const MAX_VALUE_SIZE = 65_536;
+
+/** Sort entries by key and apply reverse/limit options. Mutates the array. */
+export function sortAndPaginate<T extends { key: string }>(
+  entries: T[],
+  options?: { limit?: number; reverse?: boolean },
+): T[] {
+  entries.sort((a, b) => (a.key < b.key ? -1 : a.key > b.key ? 1 : 0));
+  if (options?.reverse) entries.reverse();
+  if (options?.limit && options.limit > 0) {
+    entries.length = Math.min(entries.length, options.limit);
+  }
+  return entries;
+}
+
+/**
+ * Create an in-memory KV store (useful for testing and local development).
+ *
+ * Data is stored in a plain `Map` and does not persist across restarts.
+ * TTL expiration is checked lazily on reads and list operations.
+ *
+ * @returns A {@linkcode Kv} instance backed by in-memory storage.
+ *
+ * @example
+ * ```ts
+ * import { createMemoryKv } from "./kv.ts";
+ *
+ * const kv = createMemoryKv();
+ * await kv.set("greeting", "hello");
+ * const value = await kv.get<string>("greeting"); // "hello"
+ * ```
+ *
+ * @example With TTL
+ * ```ts
+ * import { createMemoryKv } from "./kv.ts";
+ *
+ * const kv = createMemoryKv();
+ * await kv.set("temp", "expires soon", { expireIn: 5000 });
+ * ```
+ */
+export function createMemoryKv(): Kv {
+  const store = new Map<string, { raw: string; expiresAt?: number }>();
+
+  function isExpired(entry: { expiresAt?: number }): boolean {
+    return entry.expiresAt !== undefined && entry.expiresAt <= Date.now();
+  }
+
+  return {
+    get<T = unknown>(key: string): Promise<T | null> {
+      const entry = store.get(key);
+      if (!entry || isExpired(entry)) {
+        if (entry) store.delete(key);
+        return Promise.resolve(null);
+      }
+      return Promise.resolve(JSON.parse(entry.raw) as T);
+    },
+
+    set(key: string, value: unknown, options?: { expireIn?: number }): Promise<void> {
+      const raw = JSON.stringify(value);
+      if (raw.length > MAX_VALUE_SIZE) {
+        throw new Error(`Value exceeds max size of ${MAX_VALUE_SIZE} bytes`);
+      }
+      const expireIn = options?.expireIn;
+      const entry: { raw: string; expiresAt?: number } = { raw };
+      if (expireIn && expireIn > 0) {
+        entry.expiresAt = Date.now() + expireIn;
+      }
+      store.set(key, entry);
+      return Promise.resolve();
+    },
+
+    delete(key: string): Promise<void> {
+      store.delete(key);
+      return Promise.resolve();
+    },
+
+    list<T = unknown>(prefix: string, options?: KvListOptions): Promise<KvEntry<T>[]> {
+      const now = Date.now();
+      const entries: KvEntry<T>[] = [];
+      for (const [key, entry] of store) {
+        if (entry.expiresAt && entry.expiresAt <= now) {
+          store.delete(key);
+          continue;
+        }
+        if (key.startsWith(prefix)) {
+          entries.push({ key, value: JSON.parse(entry.raw) as T });
+        }
+      }
+      return Promise.resolve(sortAndPaginate(entries, options));
+    },
+  };
+}

package/sdk/mod.ts

@@ -0,0 +1,35 @@
+// Copyright 2025 the AAI authors. MIT license.
+/**
+ * AAI SDK — build voice agents powered by STT, LLM, and TTS.
+ *
+ * @example
+ * ```ts
+ * import { defineAgent } from "aai";
+ * import { z } from "zod";
+ *
+ * export default defineAgent({
+ *   name: "my-agent",
+ *   instructions: "You are a helpful voice assistant.",
+ *   tools: {
+ *     greet: {
+ *       description: "Greet the user by name",
+ *       parameters: z.object({ name: z.string() }),
+ *       execute: ({ name }) => `Hello, ${name}!`,
+ *     },
+ *   },
+ * });
+ * ```
+ *
+ * @module
+ */
+
+export { defineAgent } from "./define_agent.ts";
+export type {
+  AgentOptions,
+  BeforeStepResult,
+  HookContext,
+  Message,
+  PipelineMode,
+  ToolContext,
+  ToolDef,
+} from "./types.ts";

package/sdk/protocol.ts

@@ -0,0 +1,313 @@
+// 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.
+ *
+ * @module
+ */
+
+import { z } from "zod";
+import type { WorkerConfig } from "./_internal_types.ts";
+import type { Message, StepInfo } from "./types.ts";
+
+/**
+ * Current protocol version for client-server compatibility checks.
+ *
+ * Increment this when making breaking changes to the wire protocol.
+ */
+export const PROTOCOL_VERSION = 1;
+
+/**
+ * Default sample rate for speech-to-text audio in Hz.
+ *
+ * This is the sample rate expected by the STT provider (AssemblyAI).
+ */
+export const DEFAULT_STT_SAMPLE_RATE = 16_000;
+
+/**
+ * Default sample rate for text-to-speech audio in Hz.
+ *
+ * This is the sample rate produced by the TTS provider (Rime).
+ */
+export const DEFAULT_TTS_SAMPLE_RATE = 24_000;
+
+/**
+ * Audio codec identifier used in the wire protocol.
+ *
+ * All audio frames are 16-bit signed PCM, little-endian, mono.
+ */
+export const AUDIO_FORMAT = "pcm16" as const;
+
+/**
+ * Binary audio frame specification. All audio exchanged over the WebSocket as
+ * binary frames MUST conform to this spec. Any change here is a breaking
+ * protocol change.
+ */
+const _bitsPerSample = 16 as const;
+const _channels = 1 as const;
+
+/** Specification for binary audio frames exchanged over WebSocket. */
+export const AudioFrameSpec = {
+  /** Audio codec identifier sent in the `ready` message. */
+  format: AUDIO_FORMAT,
+  /** Signed 16-bit integer samples. */
+  bitsPerSample: _bitsPerSample,
+  /** Little-endian byte order. */
+  endianness: "little" as const,
+  /** Mono audio. */
+  channels: _channels,
+  /** Bytes per sample — derived from bitsPerSample and channels. */
+  bytesPerSample: (_bitsPerSample / 8) * _channels,
+} as const;
+
+/**
+ * KV operation request sent from the worker to the host via postMessage RPC.
+ *
+ * This is a discriminated union on the `op` field, representing the four
+ * key-value store operations available to sandboxed agent workers.
+ */
+export type KvRequest =
+  | { op: "get"; key: string }
+  | { op: "set"; key: string; value: string; ttl?: number | undefined }
+  | { op: "del"; key: string }
+  | {
+      op: "list";
+      prefix: string;
+      limit?: number | undefined;
+      reverse?: boolean | undefined;
+    };
+
+/** Zod schema for {@linkcode KvRequest}. */
+export const KvRequestBaseSchema: z.ZodType<KvRequest> = z.discriminatedUnion("op", [
+  z.object({ op: z.literal("get"), key: z.string().min(1) }),
+  z.object({
+    op: z.literal("set"),
+    key: z.string().min(1),
+    value: z.string(),
+    ttl: z.number().int().positive().optional(),
+  }),
+  z.object({ op: z.literal("del"), key: z.string().min(1) }),
+  z.object({
+    op: z.literal("list"),
+    prefix: z.string(),
+    limit: z.number().int().positive().optional(),
+    reverse: z.boolean().optional(),
+  }),
+]);
+
+// ─── Timeout constants ─────────────────────────────────────────────────────
+
+/** Default timeout for agent lifecycle hooks (onConnect, onTurn, etc). */
+export const HOOK_TIMEOUT_MS = 5_000;
+
+/** Default timeout for tool execution in the worker. */
+export const TOOL_EXECUTION_TIMEOUT_MS = 30_000;
+
+// ─── Error codes ───────────────────────────────────────────────────────────
+
+/** Error codes for categorizing session errors on the wire. */
+export type SessionErrorCode =
+  | "stt"
+  | "llm"
+  | "tts"
+  | "tool"
+  | "protocol"
+  | "connection"
+  | "audio"
+  | "internal";
+
+/** Zod schema for {@linkcode SessionErrorCode}. */
+export const SessionErrorCodeSchema: z.ZodType<SessionErrorCode> = z.enum([
+  "stt",
+  "llm",
+  "tts",
+  "tool",
+  "protocol",
+  "connection",
+  "audio",
+  "internal",
+]);
+
+// ─── Client events ─────────────────────────────────────────────────────────
+
+/**
+ * Discriminated union of all server→client session events.
+ *
+ * Sent via a single `event()` RPC method instead of one method per type.
+ */
+export type ClientEvent =
+  | { type: "speech_started" }
+  | { type: "speech_stopped" }
+  | { type: "transcript"; text: string; isFinal: false }
+  | {
+      type: "transcript";
+      text: string;
+      isFinal: true;
+      turnOrder?: number | undefined;
+    }
+  | { type: "turn"; text: string; turnOrder?: number | undefined }
+  | { type: "chat"; text: string }
+  | {
+      type: "tool_call_start";
+      toolCallId: string;
+      toolName: string;
+      args: Record<string, unknown>;
+    }
+  | { type: "tool_call_done"; toolCallId: string; result: string }
+  | { type: "tts_done" }
+  | { type: "cancelled" }
+  | { type: "reset" }
+  | { type: "error"; code: SessionErrorCode; message: string };
+
+/** Zod schema for a transcript event (partial or final). */
+const TranscriptEventSchema = z.object({
+  type: z.literal("transcript"),
+  text: z.string(),
+  isFinal: z.boolean(),
+  turnOrder: z.number().int().nonnegative().optional(),
+});
+
+/** Zod schema for {@linkcode ClientEvent}. */
+export const ClientEventSchema: z.ZodType<ClientEvent> = z.discriminatedUnion("type", [
+  z.object({ type: z.literal("speech_started") }),
+  z.object({ type: z.literal("speech_stopped") }),
+  TranscriptEventSchema,
+  z.object({
+    type: z.literal("turn"),
+    text: z.string(),
+    turnOrder: z.number().int().nonnegative().optional(),
+  }),
+  z.object({ type: z.literal("chat"), text: z.string() }),
+  z.object({
+    type: z.literal("tool_call_start"),
+    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(4000),
+  }),
+  z.object({ type: z.literal("tts_done") }),
+  z.object({ type: z.literal("cancelled") }),
+  z.object({ type: z.literal("reset") }),
+  z.object({
+    type: z.literal("error"),
+    code: SessionErrorCodeSchema,
+    message: z.string(),
+  }),
+]);
+
+/**
+ * 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";
+
+/** Protocol-level session config returned to the client on connect. */
+export type ReadyConfig = {
+  protocolVersion: number;
+  audioFormat: AudioFormatId;
+  sampleRate: number;
+  ttsSampleRate: number;
+  mode?: "s2s";
+};
+
+/** Client→server text messages (binary frames carry raw PCM16 audio). */
+export type ClientMessage =
+  | { type: "audio_ready" }
+  | { type: "cancel" }
+  | { type: "reset" }
+  | {
+      type: "history";
+      messages: readonly { role: "user" | "assistant"; text: string }[];
+    };
+
+/** Server→client text messages (binary frames carry raw PCM16 audio). */
+export type ServerMessage =
+  | ({ type: "config" } & ReadyConfig)
+  | { type: "audio_done" }
+  | ClientEvent;
+
+/** Zod schema for {@linkcode ClientMessage}. */
+export const ClientMessageSchema: z.ZodType<ClientMessage> = z.discriminatedUnion("type", [
+  z.object({ type: z.literal("audio_ready") }),
+  z.object({ type: z.literal("cancel") }),
+  z.object({ type: z.literal("reset") }),
+  z.object({
+    type: z.literal("history"),
+    messages: z
+      .array(
+        z.object({
+          role: z.enum(["user", "assistant"]),
+          text: z.string().max(100_000),
+        }),
+      )
+      .max(200),
+  }),
+]);
+
+// ─── Worker RPC interfaces ─────────────────────────────────────────────────
+
+/**
+ * API shape the host process exposes to the sandboxed worker.
+ *
+ * Since workers run with all permissions denied, they use this interface
+ * to proxy network requests and KV operations back to the host.
+ */
+export type HostApi = {
+  fetch(req: {
+    url: string;
+    method: string;
+    headers: Readonly<Record<string, string>>;
+    body: string | null;
+  }): Promise<{
+    status: number;
+    statusText: string;
+    headers: Record<string, string>;
+    body: string;
+  }>;
+  kv(req: KvRequest): Promise<{ result: unknown }>;
+  vectorSearch(req: { query: string; topK: number }): Promise<string>;
+};
+
+/** Combined turn configuration resolved from the worker before a turn starts. */
+export type TurnConfig = {
+  maxSteps?: number;
+  activeTools?: string[];
+};
+
+/** Worker-side RPC target interface (host calls these methods). */
+export interface WorkerRpcApi {
+  withEnv(env: Record<string, string>): WorkerRpcApi;
+  getConfig(): Promise<WorkerConfig>;
+  executeTool(
+    name: string,
+    args: Readonly<Record<string, unknown>>,
+    sessionId: string | undefined,
+    messages: readonly Message[] | undefined,
+  ): Promise<string>;
+  onConnect(sessionId: string): Promise<void>;
+  onDisconnect(sessionId: string): Promise<void>;
+  onTurn(sessionId: string, text: string): Promise<void>;
+  onError(sessionId: string, error: string): void;
+  onStep(sessionId: string, step: StepInfo): Promise<void>;
+  resolveTurnConfig(sessionId: string): Promise<TurnConfig | null>;
+}

package/sdk/runtime.ts

@@ -0,0 +1,65 @@
+// Copyright 2025 the AAI authors. MIT license.
+/**
+ * Pluggable interfaces for cross-runtime concerns.
+ *
+ * @module
+ */
+
+import { DEFAULT_TTS_SAMPLE_RATE } from "./protocol.ts";
+
+/** Runtime-agnostic structured logger. */
+export type Logger = {
+  info(msg: string, ctx?: Record<string, unknown>): void;
+  warn(msg: string, ctx?: Record<string, unknown>): void;
+  error(msg: string, ctx?: Record<string, unknown>): void;
+  debug(msg: string, ctx?: Record<string, unknown>): void;
+};
+
+/** Runtime-agnostic session metrics. */
+export type Metrics = {
+  sessionsTotal: { inc(labels: Record<string, string>): void };
+  sessionsActive: {
+    inc(labels: Record<string, string>): void;
+    dec(labels: Record<string, string>): void;
+  };
+};
+
+/** Console-based logger that works in all runtimes. */
+export const consoleLogger: Logger = {
+  info(msg, ctx) {
+    if (ctx) console.log(msg, ctx);
+    else console.log(msg);
+  },
+  warn(msg, ctx) {
+    if (ctx) console.warn(msg, ctx);
+    else console.warn(msg);
+  },
+  error(msg, ctx) {
+    if (ctx) console.error(msg, ctx);
+    else console.error(msg);
+  },
+  debug(msg, ctx) {
+    if (ctx) console.debug(msg, ctx);
+    else console.debug(msg);
+  },
+};
+
+/** No-op metrics implementation for environments without monitoring. */
+export const noopMetrics: Metrics = {
+  sessionsTotal: { inc() {} },
+  sessionsActive: { inc() {}, dec() {} },
+};
+
+/** Configuration for the AssemblyAI Speech-to-Speech connection. */
+export type S2SConfig = {
+  wssUrl: string;
+  inputSampleRate: number;
+  outputSampleRate: number;
+};
+
+/** Default S2S configuration pointing to AssemblyAI's production endpoint. */
+export const DEFAULT_S2S_CONFIG: S2SConfig = {
+  wssUrl: "wss://speech-to-speech.us.assemblyai.com/v1/realtime",
+  inputSampleRate: DEFAULT_TTS_SAMPLE_RATE,
+  outputSampleRate: DEFAULT_TTS_SAMPLE_RATE,
+};