@alexkroman1/aai 1.7.1 → 1.8.0

package/host/transports/s2s-transport.test.ts

@@ -1,7 +1,7 @@
 import { afterEach, describe, expect, test, vi } from "vitest";
 import { makeMockHandle, silentLogger } from "../_test-utils.ts";
-import type { S2sCallbacks, S2sHandle } from "../s2s.ts";
-import { _internals, createS2sTransport } from "./s2s-transport.ts";
+import type { ConnectS2sOptions, S2sCallbacks, S2sHandle, S2sWebSocket } from "../s2s.ts";
+import { _internals, createS2sTransport, type S2sTransportOptions } from "./s2s-transport.ts";
 import type { TransportCallbacks } from "./types.ts";
 
 function makeCallbacks(): TransportCallbacks {
@@ -21,34 +21,37 @@ function makeCallbacks(): TransportCallbacks {
   };
 }
 
+function makeTransportOptions(overrides: Partial<S2sTransportOptions> = {}): S2sTransportOptions {
+  return {
+    apiKey: "k",
+    s2sConfig: { wssUrl: "wss://fake", inputSampleRate: 16_000, outputSampleRate: 24_000 },
+    sessionConfig: { systemPrompt: "test", tools: [] },
+    toolSchemas: [],
+    callbacks: makeCallbacks(),
+    sid: "sid-1",
+    agent: "a",
+    logger: silentLogger,
+    ...overrides,
+  };
+}
+
 describe("S2sTransport", () => {
   test("start() opens an S2S connection and sends session.update", async () => {
     const send = vi.fn();
     const close = vi.fn();
-    const ws = Object.assign(new EventTarget(), {
+    const target = new EventTarget();
+    const ws = Object.assign(target, {
       readyState: 0,
       send,
       close,
-      addEventListener: EventTarget.prototype.addEventListener as unknown as (
-        type: string,
-        listener: EventListener,
-      ) => void,
-    }) as unknown as import("../s2s.ts").S2sWebSocket;
+      addEventListener: target.addEventListener.bind(target),
+    }) as unknown as S2sWebSocket;
     setTimeout(() => {
       (ws as unknown as { readyState: number }).readyState = 1;
-      (ws as unknown as EventTarget).dispatchEvent(new Event("open"));
+      target.dispatchEvent(new Event("open"));
     }, 0);
 
-    const t = createS2sTransport({
-      apiKey: "k",
-      s2sConfig: { wssUrl: "wss://fake", inputSampleRate: 16_000, outputSampleRate: 24_000 },
-      sessionConfig: { systemPrompt: "test", tools: [] },
-      toolSchemas: [],
-      callbacks: makeCallbacks(),
-      sid: "sid-1",
-      agent: "a",
-      createWebSocket: () => ws,
-    });
+    const t = createS2sTransport(makeTransportOptions({ createWebSocket: () => ws }));
     await t.start();
     expect(send).toHaveBeenCalled();
     const firstSend = JSON.parse(send.mock.calls[0]?.[0] as string);
@@ -58,31 +61,27 @@ describe("S2sTransport", () => {
   });
 });
 
-// ─── Reconnect tests ────────────────────────────────────────────────────────
-
 /** Capture the S2sCallbacks that the transport hands to connectS2s. */
 function setupSpiedTransport(): {
   callbacks: TransportCallbacks;
   handles: S2sHandle[];
   capturedCallbacks: S2sCallbacks[];
-  spy: ReturnType<typeof vi.spyOn>;
 } {
   const handles: S2sHandle[] = [];
   const capturedCallbacks: S2sCallbacks[] = [];
-  const spy = vi
-    .spyOn(_internals, "connectS2s")
-    .mockImplementation(async (opts: import("../s2s.ts").ConnectS2sOptions) => {
-      capturedCallbacks.push(opts.callbacks);
-      const h = makeMockHandle();
-      handles.push(h);
-      return h;
-    });
-  return {
-    callbacks: makeCallbacks(),
-    handles,
-    capturedCallbacks,
-    spy,
-  };
+  vi.spyOn(_internals, "connectS2s").mockImplementation(async (opts: ConnectS2sOptions) => {
+    capturedCallbacks.push(opts.callbacks);
+    const h = makeMockHandle();
+    handles.push(h);
+    return h;
+  });
+  return { callbacks: makeCallbacks(), handles, capturedCallbacks };
+}
+
+function expectAt<T>(arr: T[], index: number, label: string): T {
+  const value = arr[index];
+  if (!value) throw new Error(`expected ${label} at index ${index}`);
+  return value;
 }
 
 describe("S2sTransport reconnect", () => {
@@ -92,66 +91,37 @@ describe("S2sTransport reconnect", () => {
 
   test("attempts session.resume on transient close (1005) inside the resume window", async () => {
     const { callbacks, handles, capturedCallbacks } = setupSpiedTransport();
-
-    const t = createS2sTransport({
-      apiKey: "k",
-      s2sConfig: { wssUrl: "wss://fake", inputSampleRate: 16_000, outputSampleRate: 24_000 },
-      sessionConfig: { systemPrompt: "test", tools: [] },
-      toolSchemas: [],
-      callbacks,
-      sid: "sid-1",
-      agent: "a",
-      logger: silentLogger,
-    });
+    const t = createS2sTransport(makeTransportOptions({ callbacks }));
     await t.start();
 
-    // Establish session, start a reply, then drop the socket.
-    const cb1 = capturedCallbacks[0];
-    if (!cb1) throw new Error("expected first callbacks");
+    const cb1 = expectAt(capturedCallbacks, 0, "first callbacks");
     cb1.onSessionReady("sess_abc");
     cb1.onReplyStarted("rep_1");
     cb1.onClose(1005, "");
 
-    // Wait for the async resume() to fire connectS2s a second time.
     await vi.waitFor(() => {
       expect(handles.length).toBe(2);
     });
 
-    // The new handle should have received resumeSession with the prior id.
-    const newHandle = handles[1];
-    if (!newHandle) throw new Error("expected new handle");
+    const newHandle = expectAt(handles, 1, "new handle");
     expect(newHandle.resumeSession).toHaveBeenCalledWith("sess_abc");
 
-    // The in-flight reply was unblocked via onCancelled, NOT a fatal error.
     expect(callbacks.onCancelled).toHaveBeenCalledOnce();
     expect(callbacks.onError).not.toHaveBeenCalled();
   });
 
   test("does NOT reconnect on fatal close codes (1008 unauthorized)", async () => {
     const { callbacks, handles, capturedCallbacks } = setupSpiedTransport();
-
-    const t = createS2sTransport({
-      apiKey: "k",
-      s2sConfig: { wssUrl: "wss://fake", inputSampleRate: 16_000, outputSampleRate: 24_000 },
-      sessionConfig: { systemPrompt: "test", tools: [] },
-      toolSchemas: [],
-      callbacks,
-      sid: "sid-1",
-      agent: "a",
-      logger: silentLogger,
-    });
+    const t = createS2sTransport(makeTransportOptions({ callbacks }));
     await t.start();
 
-    const cb1 = capturedCallbacks[0];
-    if (!cb1) throw new Error("expected first callbacks");
+    const cb1 = expectAt(capturedCallbacks, 0, "first callbacks");
     cb1.onSessionReady("sess_abc");
     cb1.onReplyStarted("rep_1");
     cb1.onClose(1008, "unauthorized");
 
-    // No reconnect — only one connectS2s call total.
     await new Promise((resolve) => setTimeout(resolve, 5));
     expect(handles.length).toBe(1);
-    // Fatal error surfaces, since a reply was in flight.
     expect(callbacks.onError).toHaveBeenCalledWith(
       "connection",
       expect.stringContaining("S2S closed mid-reply"),
@@ -160,26 +130,14 @@ describe("S2sTransport reconnect", () => {
 
   test("does NOT reconnect when stop() was called", async () => {
     const { callbacks, handles, capturedCallbacks } = setupSpiedTransport();
-
-    const t = createS2sTransport({
-      apiKey: "k",
-      s2sConfig: { wssUrl: "wss://fake", inputSampleRate: 16_000, outputSampleRate: 24_000 },
-      sessionConfig: { systemPrompt: "test", tools: [] },
-      toolSchemas: [],
-      callbacks,
-      sid: "sid-1",
-      agent: "a",
-      logger: silentLogger,
-    });
+    const t = createS2sTransport(makeTransportOptions({ callbacks }));
     await t.start();
 
-    const cb1 = capturedCallbacks[0];
-    if (!cb1) throw new Error("expected first callbacks");
+    const cb1 = expectAt(capturedCallbacks, 0, "first callbacks");
     cb1.onSessionReady("sess_abc");
     await t.stop();
 
-    // Simulate the upstream's close arriving after stop() — it should be
-    // treated as a clean shutdown, not a transient drop worth resuming.
+    // Upstream close after stop() must be treated as clean shutdown, not a transient drop.
     cb1.onClose(1005, "");
 
     await new Promise((resolve) => setTimeout(resolve, 5));
@@ -189,28 +147,17 @@ describe("S2sTransport reconnect", () => {
 
   test("surfaces resume failure when the resumed socket also closes", async () => {
     const { callbacks, handles, capturedCallbacks } = setupSpiedTransport();
-
-    const t = createS2sTransport({
-      apiKey: "k",
-      s2sConfig: { wssUrl: "wss://fake", inputSampleRate: 16_000, outputSampleRate: 24_000 },
-      sessionConfig: { systemPrompt: "test", tools: [] },
-      toolSchemas: [],
-      callbacks,
-      sid: "sid-1",
-      agent: "a",
-      logger: silentLogger,
-    });
+    const t = createS2sTransport(makeTransportOptions({ callbacks }));
     await t.start();
 
-    capturedCallbacks[0]?.onSessionReady("sess_abc");
-    capturedCallbacks[0]?.onReplyStarted("rep_1");
-    capturedCallbacks[0]?.onClose(1005, "");
+    const cb1 = expectAt(capturedCallbacks, 0, "first callbacks");
+    cb1.onSessionReady("sess_abc");
+    cb1.onReplyStarted("rep_1");
+    cb1.onClose(1005, "");
 
     await vi.waitFor(() => expect(handles.length).toBe(2));
 
-    // The resume socket also drops before its session.ready arrives.
-    const cb2 = capturedCallbacks[1];
-    if (!cb2) throw new Error("expected resume callbacks");
+    const cb2 = expectAt(capturedCallbacks, 1, "resume callbacks");
     cb2.onClose(1006, "");
 
     expect(callbacks.onError).toHaveBeenCalledWith(
@@ -221,25 +168,17 @@ describe("S2sTransport reconnect", () => {
 
   test("surfaces resume failure when server reports session_not_found", async () => {
     const { callbacks, handles, capturedCallbacks } = setupSpiedTransport();
-
-    const t = createS2sTransport({
-      apiKey: "k",
-      s2sConfig: { wssUrl: "wss://fake", inputSampleRate: 16_000, outputSampleRate: 24_000 },
-      sessionConfig: { systemPrompt: "test", tools: [] },
-      toolSchemas: [],
-      callbacks,
-      sid: "sid-1",
-      agent: "a",
-      logger: silentLogger,
-    });
+    const t = createS2sTransport(makeTransportOptions({ callbacks }));
     await t.start();
 
-    capturedCallbacks[0]?.onSessionReady("sess_abc");
-    capturedCallbacks[0]?.onClose(1005, "");
+    const cb1 = expectAt(capturedCallbacks, 0, "first callbacks");
+    cb1.onSessionReady("sess_abc");
+    cb1.onClose(1005, "");
 
     await vi.waitFor(() => expect(handles.length).toBe(2));
 
-    capturedCallbacks[1]?.onSessionExpired();
+    const cb2 = expectAt(capturedCallbacks, 1, "resume callbacks");
+    cb2.onSessionExpired();
 
     expect(callbacks.onError).toHaveBeenCalledWith(
       "connection",
@@ -249,28 +188,20 @@ describe("S2sTransport reconnect", () => {
 
   test("after a successful resume, a later transient drop also resumes", async () => {
     const { callbacks, handles, capturedCallbacks } = setupSpiedTransport();
-
-    const t = createS2sTransport({
-      apiKey: "k",
-      s2sConfig: { wssUrl: "wss://fake", inputSampleRate: 16_000, outputSampleRate: 24_000 },
-      sessionConfig: { systemPrompt: "test", tools: [] },
-      toolSchemas: [],
-      callbacks,
-      sid: "sid-1",
-      agent: "a",
-      logger: silentLogger,
-    });
+    const t = createS2sTransport(makeTransportOptions({ callbacks }));
     await t.start();
 
-    // First connection establishes, drops, resumes, becomes ready again.
-    capturedCallbacks[0]?.onSessionReady("sess_abc");
-    capturedCallbacks[0]?.onClose(1005, "");
+    const cb1 = expectAt(capturedCallbacks, 0, "first callbacks");
+    cb1.onSessionReady("sess_abc");
+    cb1.onClose(1005, "");
     await vi.waitFor(() => expect(handles.length).toBe(2));
-    capturedCallbacks[1]?.onSessionReady("sess_abc");
 
-    // Second drop — should trigger another resume attempt.
-    capturedCallbacks[1]?.onClose(1006, "");
+    const cb2 = expectAt(capturedCallbacks, 1, "resume callbacks");
+    cb2.onSessionReady("sess_abc");
+    cb2.onClose(1006, "");
     await vi.waitFor(() => expect(handles.length).toBe(3));
-    expect(handles[2]?.resumeSession).toHaveBeenCalledWith("sess_abc");
+    expect(expectAt(handles, 2, "second resume handle").resumeSession).toHaveBeenCalledWith(
+      "sess_abc",
+    );
   });
 });

package/host/transports/s2s-transport.ts

@@ -31,9 +31,9 @@ export type S2sTransportOptions = {
 
 /**
  * Close codes worth attempting `session.resume` on. These are network/server
- * blips, not protocol or auth violations. Per AssemblyAI's docs, sessions are
- * preserved for 30 s after disconnect, so resume is bounded by the window in
- * `RESUME_WINDOW_MS` below.
+ * blips, not protocol or auth violations. AssemblyAI keeps the session
+ * available for 30 s after disconnect; reconnect runs immediately on close,
+ * so the resume request reliably lands inside that window.
  */
 const TRANSIENT_CLOSE_CODES = new Set<number>([
   1005, // No Status Received (abnormal close, no frame)
@@ -42,46 +42,28 @@ const TRANSIENT_CLOSE_CODES = new Set<number>([
   3005, // Session Cancelled (unknown server error)
 ]);
 
-/**
- * AssemblyAI keeps the session alive for 30 s after disconnect; we leave a
- * little headroom so the resume request still fits inside that window after
- * the new WebSocket finishes opening.
- */
-const RESUME_WINDOW_MS = 25_000;
-
 export function createS2sTransport(opts: S2sTransportOptions): Transport {
   const log = opts.logger ?? consoleLogger;
   const createWs = opts.createWebSocket ?? defaultCreateS2sWebSocket;
   let handle: S2sHandle | null = null;
   let currentReplyId: string | null = null;
-  /** Most recent `session.ready` ID — present once the upstream session is established. */
   let providerSessionId: string | null = null;
-  /** When the current session became ready; bounds the resume window. */
-  let sessionReadyAt = 0;
-  /** Set by `stop()` so a deliberate close doesn't trigger a reconnect. */
   let closing = false;
-  /**
-   * True while a `session.resume` round-trip is in flight (between sending
-   * resume and the next `session.ready`). Used to distinguish a resume failure
-   * (close before ready) from a normal close.
-   */
+  // True between sending `session.resume` and the next `session.ready`.
+  // Distinguishes a resume failure (close before ready) from a normal close
+  // and prevents back-to-back reconnect loops if the resumed socket also drops.
   let reconnecting = false;
-  /**
-   * Set when a reconnect attempt is kicked off, cleared once the resumed
-   * session's `session.ready` arrives. Prevents back-to-back reconnect loops
-   * when the freshly-resumed socket also drops before fully recovering.
-   */
-  let reconnectInFlight = false;
 
   function buildCallbacks(): S2sCallbacks {
     return {
       onSessionReady: (id) => {
+        const isFirstReady = providerSessionId === null;
         providerSessionId = id;
-        sessionReadyAt = Date.now();
         if (reconnecting) {
           reconnecting = false;
-          reconnectInFlight = false;
           log.info("S2S resumed", { sid: opts.sid, sessionId: id });
+        } else if (isFirstReady) {
+          log.info("S2S session ready", { sid: opts.sid, sessionId: id });
         }
         opts.callbacks.onSessionReady?.(id);
       },
@@ -104,12 +86,10 @@ export function createS2sTransport(opts: S2sTransportOptions): Transport {
       onSpeechStarted: opts.callbacks.onSpeechStarted,
       onSpeechStopped: opts.callbacks.onSpeechStopped,
       onSessionExpired: () => {
-        // The server told us the session no longer exists (most likely
-        // session_not_found in response to our resume). Surface as fatal
-        // rather than retrying — there's nothing left to resume.
+        // Server reports session no longer exists (likely session_not_found
+        // in response to our resume). Surface as fatal — nothing to resume.
         if (reconnecting) {
           reconnecting = false;
-          reconnectInFlight = false;
           log.warn("S2S resume rejected: session expired", { sid: opts.sid });
           opts.callbacks.onError("connection", "S2S resume failed: session expired");
           return;
@@ -123,17 +103,13 @@ export function createS2sTransport(opts: S2sTransportOptions): Transport {
   }
 
   function canResumeAfter(code: number): boolean {
-    if (!TRANSIENT_CLOSE_CODES.has(code)) return false;
-    if (providerSessionId === null) return false;
-    if (reconnectInFlight) return false;
-    return sessionReadyAt > 0 && Date.now() - sessionReadyAt < RESUME_WINDOW_MS;
+    return TRANSIENT_CLOSE_CODES.has(code) && providerSessionId !== null && !reconnecting;
   }
 
   function emitFatalClose(code: number, reason: string, wasReconnecting: boolean): void {
     if (wasReconnecting) {
       // Fresh resume socket closed before session.ready — resume failed.
       reconnecting = false;
-      reconnectInFlight = false;
       opts.callbacks.onError("connection", `S2S resume failed (code=${code})`);
       return;
     }
@@ -152,7 +128,6 @@ export function createS2sTransport(opts: S2sTransportOptions): Transport {
   }
 
   function startResume(prevId: string, code: number, reason: string): void {
-    reconnectInFlight = true;
     reconnecting = true;
     log.warn("S2S unexpected close — attempting resume", {
       sid: opts.sid,
@@ -168,7 +143,6 @@ export function createS2sTransport(opts: S2sTransportOptions): Transport {
     }
     void resume(prevId).catch((err: unknown) => {
       reconnecting = false;
-      reconnectInFlight = false;
       const msg = err instanceof Error ? err.message : String(err);
       log.warn("S2S resume failed", { sid: opts.sid, error: msg });
       opts.callbacks.onError("connection", `S2S resume failed: ${msg}`);
@@ -181,13 +155,11 @@ export function createS2sTransport(opts: S2sTransportOptions): Transport {
       return;
     }
     const wasReconnecting = reconnecting;
-    if (!canResumeAfter(code)) {
+    const prevId = providerSessionId;
+    if (!canResumeAfter(code) || prevId === null) {
       emitFatalClose(code, reason, wasReconnecting);
       return;
     }
-    // canResumeAfter ensures providerSessionId !== null; capture as const.
-    const prevId = providerSessionId;
-    if (prevId === null) return;
     startResume(prevId, code, reason);
   }
 
@@ -197,7 +169,7 @@ export function createS2sTransport(opts: S2sTransportOptions): Transport {
       config: opts.s2sConfig,
       createWebSocket: createWs,
       logger: log,
-      ...(opts.sid !== undefined ? { sid: opts.sid } : {}),
+      sid: opts.sid,
       callbacks: buildCallbacks(),
     });
     if (closing) {

package/host/transports/types.test.ts

@@ -3,20 +3,16 @@ import type { Transport, TransportCallbacks } from "./types.ts";
 
 describe("Transport types", () => {
   test("file compiles", () => {
-    // Types only; runtime check is trivial.
+    const noop = (): void => undefined;
     const stub: Transport = {
       start: () => Promise.resolve(),
       stop: () => Promise.resolve(),
-      // biome-ignore lint/suspicious/noEmptyBlockStatements: intentional no-op stub
-      sendUserAudio: () => {},
-      // biome-ignore lint/suspicious/noEmptyBlockStatements: intentional no-op stub
-      sendToolResult: () => {},
-      // biome-ignore lint/suspicious/noEmptyBlockStatements: intentional no-op stub
-      cancelReply: () => {},
+      sendUserAudio: noop,
+      sendToolResult: noop,
+      cancelReply: noop,
     };
     expect(stub).toBeDefined();
 
-    // Ensure TransportCallbacks is referenced (type-only check).
     type _CB = TransportCallbacks;
   });
 });

package/host/unstorage-kv.test.ts

@@ -5,8 +5,9 @@ import { describe, expect, test } from "vitest";
 import { createUnstorageKv } from "./unstorage-kv.ts";
 
 function makeKv(prefix?: string) {
-  const opts = prefix != null ? { storage: createStorage(), prefix } : { storage: createStorage() };
-  return createUnstorageKv(opts);
+  return createUnstorageKv(
+    prefix === undefined ? { storage: createStorage() } : { storage: createStorage(), prefix },
+  );
 }
 
 describe("createUnstorageKv", () => {

package/host/unstorage-kv.ts

@@ -1,40 +1,14 @@
 // Copyright 2025 the AAI authors. MIT license.
-/**
- * Key-value store backed by unstorage.
- *
- * Works with any unstorage driver (memory, fs, S3/R2, etc.).
- */
 
-import { prefixStorage, type Storage } from "unstorage";
+import { prefixStorage, type Storage, type StorageValue } from "unstorage";
 import { MAX_VALUE_SIZE } from "../sdk/constants.ts";
 import type { Kv } from "../sdk/kv.ts";
 
-/**
- * Options for creating an unstorage-backed KV store.
- */
 export type UnstorageKvOptions = {
-  /** Configured unstorage Storage instance. */
   storage: Storage;
-  /** Key prefix prepended to all operations (e.g. `"agents/my-agent/kv"`). */
   prefix?: string;
 };
 
-/**
- * Create a KV store backed by any unstorage driver.
- *
- * @param options - See {@link UnstorageKvOptions}.
- * @returns A {@link Kv} instance.
- *
- * @example
- * ```ts
- * import { createStorage } from "unstorage";
- * import { createUnstorageKv } from "@alexkroman1/aai/unstorage-kv";
- *
- * const kv = createUnstorageKv({ storage: createStorage() });
- * await kv.set("greeting", "hello");
- * const value = await kv.get<string>("greeting"); // "hello"
- * ```
- */
 export function createUnstorageKv(options: UnstorageKvOptions): Kv {
   const store = options.prefix ? prefixStorage(options.storage, options.prefix) : options.storage;
 
@@ -45,16 +19,12 @@ export function createUnstorageKv(options: UnstorageKvOptions): Kv {
     },
 
     async set(key: string, value: unknown, setOptions?: { expireIn?: number }): Promise<void> {
-      const serialized = JSON.stringify(value);
-      if (serialized.length > MAX_VALUE_SIZE) {
+      if (JSON.stringify(value).length > MAX_VALUE_SIZE) {
         throw new Error(`Value exceeds max size of ${MAX_VALUE_SIZE} bytes`);
       }
-      const storable = value as import("unstorage").StorageValue;
-      if (setOptions?.expireIn && setOptions.expireIn > 0) {
-        await store.setItem(key, storable, { ttl: Math.ceil(setOptions.expireIn / 1000) });
-      } else {
-        await store.setItem(key, storable);
-      }
+      const expireIn = setOptions?.expireIn;
+      const ttlOption = expireIn && expireIn > 0 ? { ttl: Math.ceil(expireIn / 1000) } : undefined;
+      await store.setItem(key, value as StorageValue, ttlOption);
     },
 
     async delete(keys: string | string[]): Promise<void> {