@alexkroman1/aai 1.7.1 → 1.8.0

package/host/s2s.ts

@@ -8,12 +8,10 @@ import WsWebSocket from "ws";
 import { z } from "zod";
 import { WS_OPEN } from "../sdk/constants.ts";
 import type { ClientEvent } from "../sdk/protocol.ts";
+import { base64ToUint8, uint8ToBase64 } from "./_base64.ts";
 import type { Logger, S2SConfig } from "./runtime-config.ts";
 import { consoleLogger } from "./runtime-config.ts";
 
-const uint8ToBase64 = (bytes: Uint8Array): string => Buffer.from(bytes).toString("base64");
-const base64ToUint8 = (base64: string): Uint8Array => new Uint8Array(Buffer.from(base64, "base64"));
-
 export type S2sWebSocket = {
   readonly readyState: number;
   send(data: string): void;
@@ -32,8 +30,7 @@ export type CreateS2sWebSocket = (
   opts: { headers: Record<string, string> },
 ) => S2sWebSocket;
 
-// Node's native WebSocket doesn't support custom headers.
-// Use the `ws` package which accepts { headers } in the constructor.
+// Node's native WebSocket doesn't support custom headers; the `ws` package does.
 export const defaultCreateS2sWebSocket: CreateS2sWebSocket = (url, opts) =>
   new WsWebSocket(url, { headers: opts.headers }) as unknown as S2sWebSocket;
 
@@ -41,7 +38,12 @@ export const defaultCreateS2sWebSocket: CreateS2sWebSocket = (url, opts) =>
 
 const S2sMessageSchema = z.discriminatedUnion("type", [
   z.object({ type: z.literal("session.ready"), session_id: z.string() }).passthrough(),
-  z.object({ type: z.literal("session.updated") }).passthrough(),
+  z
+    .object({
+      type: z.literal("session.updated"),
+      config: z.object({ id: z.string().optional() }).passthrough().optional(),
+    })
+    .passthrough(),
   z.object({ type: z.literal("input.speech.started") }),
   z.object({ type: z.literal("input.speech.stopped") }),
   z.object({ type: z.literal("transcript.user"), item_id: z.string(), text: z.string() }),
@@ -86,12 +88,14 @@ export type S2sEvent = ClientEvent & { _interrupted?: boolean };
 type DispatchState = { speechActive: boolean };
 
 type DispatchContext = {
-  /** Logger used for diagnostic `S2S <<` arrival logs. */
   log: Logger;
-  /** Session id threaded through diagnostic logs; omitted when undefined. */
   sid?: string;
 };
 
+function sidFields(ctx: DispatchContext): { sid?: string } {
+  return ctx.sid !== undefined ? { sid: ctx.sid } : {};
+}
+
 function dispatchS2sMessage(
   callbacks: S2sCallbacks,
   msg: S2sServerMessage,
@@ -103,6 +107,10 @@ function dispatchS2sMessage(
       callbacks.onSessionReady(msg.session_id);
       break;
     case "session.updated":
+      // The S2S API conveys the session id via `config.id` in the success
+      // path (no separate `session.ready` is emitted); capturing it here is
+      // required for resume on transient close.
+      if (msg.config?.id !== undefined) callbacks.onSessionReady(msg.config.id);
       break;
     case "input.speech.started":
       if (!state.speechActive) {
@@ -129,20 +137,21 @@ function dispatchS2sMessage(
       callbacks.onToolCall(msg.call_id, msg.name, msg.args);
       break;
     case "reply.done":
-      // Log every raw reply.done arrival from the S2S service — one line per
-      // event, before any client-facing dedup — so we can cross-check which
-      // stalled sessions actually received reply.done for their turn.
+      // Log every raw reply.done before client-facing dedup so we can
+      // cross-check which stalled sessions actually got one.
       ctx.log.info("S2S << reply.done", {
-        ...(ctx.sid !== undefined ? { sid: ctx.sid } : {}),
+        ...sidFields(ctx),
         status: msg.status ?? "completed",
       });
-      if (msg.status === "interrupted") {
-        callbacks.onCancelled();
-      } else {
-        callbacks.onReplyDone();
-      }
+      if (msg.status === "interrupted") callbacks.onCancelled();
+      else callbacks.onReplyDone();
       break;
     case "session.error":
+      ctx.log.warn("S2S << session.error", {
+        ...sidFields(ctx),
+        code: msg.code,
+        message: msg.message,
+      });
       if (msg.code === "session_not_found" || msg.code === "session_forbidden") {
         callbacks.onSessionExpired();
       } else {
@@ -234,12 +243,10 @@ export function connectS2s(opts: ConnectS2sOptions): Promise<S2sHandle> {
         return;
       }
       const json = JSON.stringify(msg);
-      if (msg.type !== "input.audio") {
-        if (msg.type === "session.update") {
-          log.info(`S2S >> ${msg.type}`, { payload: json });
-        } else {
-          log.info(`S2S >> ${msg.type}`);
-        }
+      if (msg.type === "session.update") {
+        log.info(`S2S >> ${msg.type}`, { payload: json });
+      } else if (msg.type !== "input.audio") {
+        log.info(`S2S >> ${msg.type}`);
       }
       ws.send(json);
     }
@@ -281,42 +288,41 @@ export function connectS2s(opts: ConnectS2sOptions): Promise<S2sHandle> {
       resolve(handle);
     });
 
-    function tryParseJson(data: unknown): unknown | undefined {
-      try {
-        return JSON.parse(String(data));
-      } catch {
-        log.warn("S2S << invalid JSON", { data: String(data).slice(0, 200) });
-      }
-    }
-
-    function handleAudioFastPath(obj: { type?: unknown; data?: unknown }): boolean {
-      if (obj.type === "reply.audio" && typeof obj.data === "string") {
-        callbacks.onAudio(base64ToUint8(obj.data));
-        return true;
-      }
-      return false;
-    }
-
-    function logIncoming(obj: { type?: unknown }): void {
+    function logIncoming(type: unknown): void {
       // reply.audio and input.audio are ~95% of traffic — skip logging.
-      if (obj.type === "reply.audio" || obj.type === "input.audio") return;
-      // reply.done gets a richer log (sid + status) inside dispatch;
-      // skip the generic line here to avoid a duplicate.
-      if (obj.type === "reply.done") return;
-      log.info(`S2S << ${obj.type}`);
+      // reply.done and session.error get richer logs inside dispatch;
+      // skip here to avoid a duplicate line.
+      if (
+        type === "reply.audio" ||
+        type === "input.audio" ||
+        type === "reply.done" ||
+        type === "session.error"
+      ) {
+        return;
+      }
+      log.info(`S2S << ${type}`);
     }
 
     ws.addEventListener("message", (ev) => {
-      const raw = tryParseJson(ev.data);
-      if (raw === undefined) return;
+      let raw: unknown;
+      try {
+        raw = JSON.parse(String(ev.data));
+      } catch {
+        log.warn("S2S << invalid JSON", { data: String(ev.data).slice(0, 200) });
+        return;
+      }
 
       if (typeof raw !== "object" || raw === null || Array.isArray(raw)) {
         log.warn("S2S << non-object JSON message", { type: typeof raw });
         return;
       }
       const obj = raw as Record<string, unknown>;
-      logIncoming(obj);
-      if (handleAudioFastPath(obj)) return;
+      logIncoming(obj.type);
+
+      if (obj.type === "reply.audio" && typeof obj.data === "string") {
+        callbacks.onAudio(base64ToUint8(obj.data));
+        return;
+      }
 
       const parsed = parseS2sMessage(obj);
       if (!parsed) {

package/host/server-shutdown.test.ts

@@ -1,10 +1,4 @@
 // Copyright 2025 the AAI authors. MIT license.
-/**
- * Tests for server shutdown timeout behavior.
- *
- * Creates a mock runtime with controlled shutdown behavior, then exercises
- * the timeout and graceful paths in close().
- */
 
 import { afterEach, describe, expect, test, vi } from "vitest";
 import { silentLogger } from "./_test-utils.ts";
@@ -24,35 +18,25 @@ function createMockRuntime(): Runtime {
   };
 }
 
-describe("server shutdown timeout", () => {
-  let server: ReturnType<typeof createServer> | null = null;
+async function startServer(): Promise<ReturnType<typeof createServer>> {
+  const server = createServer({ runtime: createMockRuntime(), logger: silentLogger });
+  await server.listen(0);
+  return server;
+}
 
+describe("server shutdown timeout", () => {
   afterEach(() => {
     mockShutdown = vi.fn<() => Promise<void>>().mockResolvedValue(undefined);
-    server = null;
   });
 
   test("close calls runtime.shutdown()", async () => {
-    mockShutdown = vi.fn<() => Promise<void>>().mockResolvedValue(undefined);
-
-    server = createServer({
-      runtime: createMockRuntime(),
-      logger: silentLogger,
-    });
-    await server.listen(0);
-
+    const server = await startServer();
     await server.close();
     expect(mockShutdown).toHaveBeenCalledOnce();
   }, 10_000);
 
   test("close resolves quickly when runtime.shutdown() resolves", async () => {
-    mockShutdown = vi.fn<() => Promise<void>>().mockResolvedValue(undefined);
-
-    server = createServer({
-      runtime: createMockRuntime(),
-      logger: silentLogger,
-    });
-    await server.listen(0);
+    const server = await startServer();
 
     const start = Date.now();
     await server.close();
@@ -65,12 +49,7 @@ describe("server shutdown timeout", () => {
   test("close propagates when runtime.shutdown() rejects", async () => {
     mockShutdown = vi.fn<() => Promise<void>>().mockRejectedValue(new Error("boom"));
 
-    server = createServer({
-      runtime: createMockRuntime(),
-      logger: silentLogger,
-    });
-    await server.listen(0);
-
+    const server = await startServer();
     await expect(server.close()).rejects.toThrow("boom");
   }, 10_000);
 });

package/host/server.test.ts

@@ -36,16 +36,12 @@ describe("createServer", () => {
     const { runtime } = makeRuntime({ name: "health-agent" });
     server = createServer({ runtime, name: "health-agent", logger: silentLogger });
     await server.listen(0);
-    await server.close();
-    server = null;
   });
 
   test("listen and close lifecycle works", async () => {
     const { runtime } = makeRuntime();
     server = createServer({ runtime, logger: silentLogger });
     await server.listen(0);
-    await server.close();
-    server = null;
   });
 
   test("/ returns default HTML with escaped agent name", async () => {
@@ -77,11 +73,7 @@ describe("createServer", () => {
 
   test("/health returns JSON with agent name", async () => {
     const { runtime } = makeRuntime({ name: "my-agent" });
-    server = createServer({
-      runtime,
-      name: "my-agent",
-      logger: silentLogger,
-    });
+    server = createServer({ runtime, name: "my-agent", logger: silentLogger });
     await server.listen(0);
 
     const res = await fetch(`http://localhost:${server.port}/health`);
@@ -91,10 +83,7 @@ describe("createServer", () => {
 
   test("404 triggers error-level logging", async () => {
     const { runtime } = makeRuntime();
-    server = createServer({
-      runtime,
-      logger: silentLogger,
-    });
+    server = createServer({ runtime, logger: silentLogger });
     await server.listen(0);
 
     await fetch(`http://localhost:${server.port}/nonexistent-path`);

package/host/server.ts

@@ -51,7 +51,12 @@ export type AgentServer = {
   port: number | undefined;
 };
 
-// ── Static file serving ─────────────────────────────────────────────────
+const JSON_HEADERS = { "Content-Type": "application/json" } as const;
+
+function sendJson(res: http.ServerResponse, status: number, body: unknown): void {
+  res.writeHead(status, JSON_HEADERS);
+  res.end(JSON.stringify(body));
+}
 
 async function serveStatic(
   dir: string,
@@ -61,8 +66,8 @@ async function serveStatic(
   const url = req.url?.split("?")[0] ?? "/";
   const filePath = path.join(dir, url === "/" ? "index.html" : url);
 
-  // Prevent path traversal — use resolved dir + separator to avoid prefix
-  // collisions (e.g. dir="/app/static" matching "/app/static-secrets/…").
+  // Use resolved dir + separator to avoid prefix collisions
+  // (e.g. dir="/app/static" matching "/app/static-secrets/…").
   const resolved = path.resolve(dir);
   if (!filePath.startsWith(resolved + path.sep) && filePath !== resolved) return false;
 
@@ -79,77 +84,72 @@ async function serveStatic(
   }
 }
 
-// ── Server ──────────────────────────────────────────────────────────────
+async function readBody(req: http.IncomingMessage): Promise<string> {
+  let body = "";
+  for await (const chunk of req) body += chunk;
+  return body;
+}
 
-function handleVectorPost(
+async function handleVectorPost(
   vector: Vector,
   req: http.IncomingMessage,
   res: http.ServerResponse,
-): void {
-  let body = "";
-  req.on("data", (chunk) => {
-    body += chunk;
-  });
-  req.on("end", async () => {
-    try {
-      const json = JSON.parse(body);
-      const parsed = VectorRequestSchema.safeParse(json);
-      if (!parsed.success) {
-        res.statusCode = 400;
-        res.end(JSON.stringify({ error: parsed.error.message }));
-        return;
-      }
-      const op = parsed.data;
-      let result: unknown;
-      switch (op.op) {
-        case "upsert":
-          await vector.upsert(op.id, op.text, op.metadata);
-          result = "OK";
-          break;
-        case "query":
-          result = await vector.query(op.text, {
-            ...(op.topK !== undefined ? { topK: op.topK } : {}),
-            ...(op.filter !== undefined ? { filter: op.filter } : {}),
-          });
-          break;
-        case "delete":
-          await vector.delete(op.ids);
-          result = "OK";
-          break;
-        default:
-          break;
+): Promise<void> {
+  try {
+    const parsed = VectorRequestSchema.safeParse(JSON.parse(await readBody(req)));
+    if (!parsed.success) {
+      sendJson(res, 400, { error: parsed.error.message });
+      return;
+    }
+    const op = parsed.data;
+    let result: unknown;
+    switch (op.op) {
+      case "upsert":
+        await vector.upsert(op.id, op.text, op.metadata);
+        result = "OK";
+        break;
+      case "query":
+        result = await vector.query(op.text, {
+          ...(op.topK !== undefined ? { topK: op.topK } : {}),
+          ...(op.filter !== undefined ? { filter: op.filter } : {}),
+        });
+        break;
+      case "delete":
+        await vector.delete(op.ids);
+        result = "OK";
+        break;
+      default: {
+        const _exhaustive: never = op;
+        return _exhaustive;
       }
-      res.statusCode = 200;
-      res.end(JSON.stringify({ result }));
-    } catch (err) {
-      res.statusCode = 500;
-      res.end(JSON.stringify({ error: err instanceof Error ? err.message : String(err) }));
     }
-  });
+    sendJson(res, 200, { result });
+  } catch (err) {
+    sendJson(res, 500, { error: err instanceof Error ? err.message : String(err) });
+  }
 }
 
-function handleKvGet(kv: Kv, req: http.IncomingMessage, res: http.ServerResponse): void {
-  const fullUrl = new URL(req.url ?? "/", "http://localhost");
-  const key = fullUrl.searchParams.get("key");
+async function handleKvGet(
+  kv: Kv,
+  req: http.IncomingMessage,
+  res: http.ServerResponse,
+): Promise<void> {
+  const key = new URL(req.url ?? "/", "http://localhost").searchParams.get("key");
   if (!key) {
-    res.writeHead(400, { "Content-Type": "application/json" });
-    res.end(JSON.stringify({ error: "Missing key query parameter" }));
+    sendJson(res, 400, { error: "Missing key query parameter" });
     return;
   }
-  kv.get(key)
-    .then((value) => {
-      if (value === null) {
-        res.writeHead(404, { "Content-Type": "application/json" });
-        res.end("null");
-      } else {
-        res.writeHead(200, { "Content-Type": "application/json" });
-        res.end(JSON.stringify(value));
-      }
-    })
-    .catch(() => {
-      res.writeHead(500, { "Content-Type": "application/json" });
-      res.end(JSON.stringify({ error: "KV error" }));
-    });
+  try {
+    const value = await kv.get(key);
+    if (value === null) {
+      res.writeHead(404, JSON_HEADERS);
+      res.end("null");
+      return;
+    }
+    sendJson(res, 200, value);
+  } catch {
+    sendJson(res, 500, { error: "KV error" });
+  }
 }
 
 /**
@@ -165,67 +165,52 @@ export function createServer(options: ServerOptions): AgentServer {
     throw new Error("clientHtml and clientDir are mutually exclusive");
   }
 
-  // Pre-compute the default HTML page once (the agent name never changes).
-  const escapedName = escapeHtml(name);
   const defaultHtml =
     clientHtml ??
-    `<!DOCTYPE html><html><body><h1>${escapedName}</h1><p>Agent server running.</p></body></html>`;
+    `<!DOCTYPE html><html><body><h1>${escapeHtml(name)}</h1><p>Agent server running.</p></body></html>`;
+
+  async function handleRequest(
+    req: http.IncomingMessage,
+    res: http.ServerResponse,
+    url: string,
+    method: string,
+  ): Promise<void> {
+    if (clientDir && (await serveStatic(clientDir, req, res))) return;
+
+    if (method === "GET" && url === "/") {
+      res.writeHead(200, { "Content-Type": "text/html" });
+      res.end(defaultHtml);
+      return;
+    }
+
+    logger.error(`${method} ${url} 404`);
+    sendJson(res, 404, { error: "Not found" });
+  }
 
   const httpServer = http.createServer((req, res) => {
     const url = req.url?.split("?")[0] ?? "/";
     const method = req.method ?? "GET";
 
-    // Security headers
     res.setHeader("Content-Security-Policy", AGENT_CSP);
     res.setHeader("X-Content-Type-Options", "nosniff");
     res.setHeader("X-Frame-Options", "SAMEORIGIN");
 
-    // Health endpoint
     if (method === "GET" && url === "/health") {
-      res.writeHead(200, { "Content-Type": "application/json" });
-      res.end(JSON.stringify({ status: "ok", name }));
+      sendJson(res, 200, { status: "ok", name });
       return;
     }
-
-    // KV endpoint
     if (kv && method === "GET" && url === "/kv") {
-      handleKvGet(kv, req, res);
+      void handleKvGet(kv, req, res);
       return;
     }
-
-    // Vector endpoint
     if (vector && method === "POST" && url === "/vector") {
-      handleVectorPost(vector, req, res);
+      void handleVectorPost(vector, req, res);
       return;
     }
 
-    // Routes that may need async handling
     void handleRequest(req, res, url, method);
   });
 
-  async function handleRequest(
-    req: http.IncomingMessage,
-    res: http.ServerResponse,
-    url: string,
-    method: string,
-  ): Promise<void> {
-    // Static files from client dir
-    if (clientDir && (await serveStatic(clientDir, req, res))) return;
-
-    // Default HTML
-    if (method === "GET" && url === "/") {
-      res.writeHead(200, { "Content-Type": "text/html" });
-      res.end(defaultHtml);
-      return;
-    }
-
-    // 404
-    logger.error(`${method} ${url} 404`);
-    res.writeHead(404, { "Content-Type": "application/json" });
-    res.end(JSON.stringify({ error: "Not found" }));
-  }
-
-  // WebSocket upgrade via ws
   const wss = new WebSocketServer({ noServer: true, maxPayload: MAX_WS_PAYLOAD_BYTES });
 
   httpServer.on("upgrade", (req, socket, head) => {

package/host/session-core.test.ts

@@ -1,4 +1,5 @@
 import { describe, expect, test, vi } from "vitest";
+import type { AgentConfig } from "../sdk/_internal-types.ts";
 import type { ClientEvent, ClientSink } from "../sdk/protocol.ts";
 import { DEFAULT_SYSTEM_PROMPT } from "../sdk/types.ts";
 import { flush } from "./_test-utils.ts";
@@ -9,7 +10,7 @@ import type { Transport } from "./transports/types.ts";
 function makeSink(): {
   events: ClientEvent[];
   audioChunks: Uint8Array[];
-  audioDoneCount: number;
+  readonly audioDoneCount: number;
   sink: ClientSink;
 } {
   const events: ClientEvent[] = [];
@@ -22,9 +23,7 @@ function makeSink(): {
       return audioDoneCount;
     },
     sink: {
-      get open() {
-        return true;
-      },
+      open: true,
       event: (e) => {
         events.push(e);
       },
@@ -34,13 +33,13 @@ function makeSink(): {
       playAudioDone: () => {
         audioDoneCount++;
       },
-    } satisfies ClientSink,
+    },
   };
 }
 
-function makeTransport(): Transport & { starts: number; stops: number } {
-  let starts = 0,
-    stops = 0;
+function makeTransport(): Transport & { readonly starts: number; readonly stops: number } {
+  let starts = 0;
+  let stops = 0;
   return {
     start: async () => {
       starts++;
@@ -60,6 +59,10 @@ function makeTransport(): Transport & { starts: number; stops: number } {
   };
 }
 
+function makeAgentConfig(overrides: Partial<AgentConfig> = {}): AgentConfig {
+  return { name: "test", systemPrompt: DEFAULT_SYSTEM_PROMPT, greeting: "", ...overrides };
+}
+
 function makeCore(overrides: Partial<SessionCoreOptions> = {}): {
   core: SessionCore;
   sink: ReturnType<typeof makeSink>;
@@ -71,7 +74,7 @@ function makeCore(overrides: Partial<SessionCoreOptions> = {}): {
     id: "s-test",
     agent: "test-agent",
     client: sink.sink,
-    agentConfig: { name: "test", systemPrompt: DEFAULT_SYSTEM_PROMPT, greeting: "" },
+    agentConfig: makeAgentConfig(),
     executeTool: vi.fn(async () => "ok"),
     transport,
     ...overrides,
@@ -98,12 +101,7 @@ describe("createSessionCore — lifecycle", () => {
     vi.useFakeTimers();
     try {
       const { core, sink } = makeCore({
-        agentConfig: {
-          name: "test",
-          systemPrompt: DEFAULT_SYSTEM_PROMPT,
-          greeting: "",
-          idleTimeoutMs: 1000,
-        } as unknown as SessionCoreOptions["agentConfig"],
+        agentConfig: makeAgentConfig({ idleTimeoutMs: 1000 }),
       });
       await core.start();
       await core.stop();
@@ -190,10 +188,8 @@ describe("createSessionCore — tool call pending results", () => {
     await core.start();
     core.onReplyStarted("r1");
     core.onToolCall("cid", "my_tool", {});
-    // Let the async tool IIFE settle and push to pendingTools
     await flush();
     core.onReplyDone();
-    // Poll until tool results are forwarded and toolCallDone fires
     await vi.waitFor(() =>
       expect(transport.sendToolResult).toHaveBeenCalledWith("cid", "tool-output"),
     );
@@ -206,12 +202,7 @@ describe("createSessionCore — idle timeout", () => {
     vi.useFakeTimers();
     try {
       const { core, sink } = makeCore({
-        agentConfig: {
-          name: "t",
-          systemPrompt: DEFAULT_SYSTEM_PROMPT,
-          greeting: "",
-          idleTimeoutMs: 1000,
-        } as unknown as SessionCoreOptions["agentConfig"],
+        agentConfig: makeAgentConfig({ name: "t", idleTimeoutMs: 1000 }),
       });
       await core.start();
       expect(sink.events.filter((e) => e.type === "idle_timeout")).toHaveLength(0);
@@ -225,12 +216,7 @@ describe("createSessionCore — idle timeout", () => {
     vi.useFakeTimers();
     try {
       const { core, sink } = makeCore({
-        agentConfig: {
-          name: "t",
-          systemPrompt: DEFAULT_SYSTEM_PROMPT,
-          greeting: "",
-          idleTimeoutMs: 1000,
-        } as unknown as SessionCoreOptions["agentConfig"],
+        agentConfig: makeAgentConfig({ name: "t", idleTimeoutMs: 1000 }),
       });
       await core.start();
       vi.advanceTimersByTime(500);
@@ -251,7 +237,6 @@ describe("createSessionCore — history", () => {
     await core.start();
     core.onHistory([{ role: "user", content: "prior" }]);
     core.onUserTranscript("now");
-    // No direct introspection — but onReset clears history and replay should see no effect on subsequent behavior.
     core.onReset();
   });
 });