@alexkroman1/aai 0.9.3 → 0.10.1

package/dist/s2s.js

@@ -1,110 +1,116 @@
 import { consoleLogger } from "./runtime.js";
-import { z } from "zod";
-import { WebSocket } from "ws";
+import { s2sConnectionDuration, s2sErrorCounter, tracer } from "./telemetry.js";
 import { createNanoEvents } from "nanoevents";
+import { WebSocket } from "ws";
 //#region s2s.ts
 const uint8ToBase64 = (bytes) => Buffer.from(bytes).toString("base64");
 const base64ToUint8 = (base64) => new Uint8Array(Buffer.from(base64, "base64"));
-/** WebSocket readyState constant for OPEN. */
 const WS_OPEN = 1;
-/** Default S2S WebSocket factory using the `ws` package (Node-only). */
 const defaultCreateS2sWebSocket = (url, opts) => new WebSocket(url, { headers: opts.headers });
-const S2sServerMessageSchema = z.discriminatedUnion("type", [
-	z.object({
-		type: z.literal("session.ready"),
-		session_id: z.string()
-	}),
-	z.object({ type: z.literal("session.updated") }).passthrough(),
-	z.object({ type: z.literal("input.speech.started") }),
-	z.object({ type: z.literal("input.speech.stopped") }),
-	z.object({
-		type: z.literal("transcript.user.delta"),
-		text: z.string()
-	}),
-	z.object({
-		type: z.literal("transcript.user"),
-		item_id: z.string(),
-		text: z.string()
-	}),
-	z.object({
-		type: z.literal("reply.started"),
-		reply_id: z.string()
-	}),
-	z.object({
-		type: z.literal("transcript.agent.delta"),
-		delta: z.string()
-	}).passthrough(),
-	z.object({
-		type: z.literal("transcript.agent"),
-		text: z.string()
-	}),
-	z.object({ type: z.literal("reply.content_part.started") }).passthrough(),
-	z.object({ type: z.literal("reply.content_part.done") }).passthrough(),
-	z.object({
-		type: z.literal("tool.call"),
-		call_id: z.string(),
-		name: z.string(),
-		args: z.record(z.string(), z.unknown()).optional().default({})
-	}),
-	z.object({
-		type: z.literal("reply.done"),
-		status: z.string().optional()
-	}),
-	z.object({
-		type: z.literal("session.error"),
-		code: z.string(),
-		message: z.string()
-	}),
-	z.object({
-		type: z.literal("error"),
-		message: z.string()
-	})
+function hasStringFields(obj, ...keys) {
+	for (const k of keys) if (typeof obj[k] !== "string") return false;
+	return true;
+}
+function parseAgentTranscript(obj) {
+	if (typeof obj.text !== "string") return;
+	return {
+		type: "transcript.agent",
+		text: obj.text,
+		reply_id: typeof obj.reply_id === "string" ? obj.reply_id : "",
+		item_id: typeof obj.item_id === "string" ? obj.item_id : "",
+		interrupted: obj.interrupted === true
+	};
+}
+function parseToolCall(obj) {
+	if (typeof obj.call_id !== "string" || typeof obj.name !== "string") return;
+	const args = obj.args != null && typeof obj.args === "object" && !Array.isArray(obj.args) ? obj.args : {};
+	return {
+		type: "tool.call",
+		call_id: obj.call_id,
+		name: obj.name,
+		args
+	};
+}
+function passthrough(obj) {
+	return obj;
+}
+function requireFields(...keys) {
+	return (obj) => hasStringFields(obj, ...keys) ? obj : void 0;
+}
+const MESSAGE_VALIDATORS = new Map([
+	["session.ready", requireFields("session_id")],
+	["session.updated", passthrough],
+	["input.speech.started", passthrough],
+	["input.speech.stopped", passthrough],
+	["reply.content_part.started", passthrough],
+	["reply.content_part.done", passthrough],
+	["transcript.user.delta", requireFields("text")],
+	["transcript.user", requireFields("item_id", "text")],
+	["reply.started", requireFields("reply_id")],
+	["transcript.agent.delta", requireFields("delta")],
+	["transcript.agent", parseAgentTranscript],
+	["tool.call", parseToolCall],
+	["reply.done", (obj) => ({
+		type: "reply.done",
+		...typeof obj.status === "string" ? { status: obj.status } : {}
+	})],
+	["session.error", requireFields("code", "message")],
+	["error", requireFields("message")]
 ]);
-/** Dispatch a parsed S2S server message to the emitter. */
+function parseS2sMessage(obj) {
+	const type = obj.type;
+	if (typeof type !== "string") return;
+	return MESSAGE_VALIDATORS.get(type)?.(obj);
+}
 function dispatchS2sMessage(emitter, msg) {
 	switch (msg.type) {
 		case "session.ready":
-			emitter.emit("ready", { session_id: msg.session_id });
+			emitter.emit("ready", { sessionId: msg.session_id });
 			break;
 		case "session.updated":
-			emitter.emit("session_updated", msg);
+			emitter.emit("sessionUpdated", msg);
 			break;
 		case "input.speech.started":
-			emitter.emit("speech_started");
+			emitter.emit("speechStarted");
 			break;
 		case "input.speech.stopped":
-			emitter.emit("speech_stopped");
+			emitter.emit("speechStopped");
 			break;
 		case "transcript.user.delta":
-			emitter.emit("user_transcript_delta", { text: msg.text });
+			emitter.emit("userTranscriptDelta", { text: msg.text });
 			break;
 		case "transcript.user":
-			emitter.emit("user_transcript", {
-				item_id: msg.item_id,
+			emitter.emit("userTranscript", {
+				itemId: msg.item_id,
 				text: msg.text
 			});
 			break;
 		case "reply.started":
-			emitter.emit("reply_started", { reply_id: msg.reply_id });
+			emitter.emit("replyStarted", { replyId: msg.reply_id });
 			break;
 		case "transcript.agent.delta":
-			emitter.emit("agent_transcript_delta", { text: msg.delta });
+			emitter.emit("agentTranscriptDelta", { text: msg.delta });
 			break;
 		case "transcript.agent":
-			emitter.emit("agent_transcript", { text: msg.text });
+			emitter.emit("agentTranscript", {
+				text: msg.text,
+				replyId: msg.reply_id,
+				itemId: msg.item_id,
+				interrupted: msg.interrupted
+			});
 			break;
 		case "tool.call":
-			emitter.emit("tool_call", {
-				call_id: msg.call_id,
+			emitter.emit("toolCall", {
+				callId: msg.call_id,
 				name: msg.name,
 				args: msg.args
 			});
 			break;
 		case "reply.done":
-			emitter.emit("reply_done", { ...msg.status ? { status: msg.status } : {} });
+			emitter.emit("replyDone", msg.status ? { status: msg.status } : {});
 			break;
 		case "session.error":
-			if (msg.code === "session_not_found" || msg.code === "session_forbidden") emitter.emit("session_expired", {
+			if (msg.code === "session_not_found" || msg.code === "session_forbidden") emitter.emit("sessionExpired", {
 				code: msg.code,
 				message: msg.message
 			});
@@ -121,26 +127,23 @@ function dispatchS2sMessage(emitter, msg) {
 			break;
 		case "reply.content_part.started":
 		case "reply.content_part.done": break;
+		default: break;
 	}
 }
-/**
-* Connect to AssemblyAI's Speech-to-Speech WebSocket API.
-*
-* Returns an {@link S2sHandle} with a typed `on()` method.
-* Consumers listen for events: `ready`, `speech_started`, `speech_stopped`,
-* `user_transcript_delta`, `user_transcript`, `reply_started`,
-* `reply_done`, `audio`, `agent_transcript`, `tool_call`,
-* `session_expired`, `error`, `close`.
-*/
 function connectS2s(opts) {
 	const { apiKey, config, createWebSocket, logger: log = consoleLogger } = opts;
 	return new Promise((resolve, reject) => {
 		log.info("S2S connecting", { url: config.wssUrl });
+		const connectionSpan = tracer.startSpan("s2s.connection", { attributes: { "aai.s2s.url": config.wssUrl } });
+		const connectStart = performance.now();
 		const ws = createWebSocket(config.wssUrl, { headers: { Authorization: `Bearer ${apiKey}` } });
 		const emitter = createNanoEvents();
 		let opened = false;
 		function send(msg) {
-			if (ws.readyState !== WS_OPEN) return;
+			if (ws.readyState !== WS_OPEN) {
+				log.debug("S2S send dropped: socket not open", { type: msg.type });
+				return;
+			}
 			const json = JSON.stringify(msg);
 			if (msg.type !== "input.audio") log.info(`S2S >> ${msg.type}`, msg.type === "session.update" ? { payload: json } : void 0);
 			ws.send(json);
@@ -148,7 +151,10 @@ function connectS2s(opts) {
 		const handle = {
 			on: emitter.on.bind(emitter),
 			sendAudio(audio) {
-				if (ws.readyState !== WS_OPEN) return;
+				if (ws.readyState !== WS_OPEN) {
+					log.debug("S2S sendAudio dropped: socket not open");
+					return;
+				}
 				ws.send(`{"type":"input.audio","audio":"${uint8ToBase64(audio)}"}`);
 			},
 			sendToolResult(callId, result) {
@@ -164,9 +170,13 @@ function connectS2s(opts) {
 				send(msg);
 			},
 			updateSession(sessionConfig) {
+				const { systemPrompt, ...rest } = sessionConfig;
 				send({
 					type: "session.update",
-					session: sessionConfig
+					session: {
+						system_prompt: systemPrompt,
+						...rest
+					}
 				});
 			},
 			resumeSession(sessionId) {
@@ -183,6 +193,7 @@ function connectS2s(opts) {
 		ws.addEventListener("open", () => {
 			opened = true;
 			log.info("S2S WebSocket open");
+			connectionSpan.addEvent("ws.open");
 			resolve(handle);
 		});
 		function tryParseJson(data) {
@@ -190,7 +201,6 @@ function connectS2s(opts) {
 				return JSON.parse(String(data));
 			} catch {
 				log.warn("S2S << invalid JSON", { data: String(data).slice(0, 200) });
-				return;
 			}
 		}
 		function handleAudioFastPath(obj) {
@@ -208,15 +218,19 @@ function connectS2s(opts) {
 		function handleS2sMessage(ev) {
 			const raw = tryParseJson(ev.data);
 			if (raw === void 0) return;
+			if (typeof raw !== "object" || raw === null || Array.isArray(raw)) {
+				log.warn("S2S << non-object JSON message", { type: typeof raw });
+				return;
+			}
 			const obj = raw;
 			logIncoming(obj);
 			if (handleAudioFastPath(obj)) return;
-			const parsed = S2sServerMessageSchema.safeParse(raw);
-			if (!parsed.success) {
+			const parsed = parseS2sMessage(obj);
+			if (!parsed) {
 				log.warn(`S2S << unrecognised message type: ${obj.type ?? JSON.stringify(raw).slice(0, 200)}`);
 				return;
 			}
-			dispatchS2sMessage(emitter, parsed.data);
+			dispatchS2sMessage(emitter, parsed);
 		}
 		ws.addEventListener("message", handleS2sMessage);
 		ws.addEventListener("close", (ev) => {
@@ -224,14 +238,30 @@ function connectS2s(opts) {
 				code: ev.code ?? 0,
 				reason: ev.reason ?? ""
 			});
+			const elapsed = (performance.now() - connectStart) / 1e3;
+			s2sConnectionDuration.record(elapsed);
+			connectionSpan.addEvent("ws.closed", {
+				"ws.close.code": ev.code ?? 0,
+				"ws.close.reason": ev.reason ?? ""
+			});
+			connectionSpan.end();
+			if (!opened) reject(/* @__PURE__ */ new Error(`WebSocket closed before open (code: ${ev.code ?? 0})`));
 			emitter.emit("close");
 		});
 		ws.addEventListener("error", (ev) => {
 			const message = typeof ev.message === "string" ? ev.message : "WebSocket error";
 			const errObj = new Error(message);
 			log.error("S2S WebSocket error", { error: errObj.message });
-			if (!opened) reject(errObj);
-			else emitter.emit("error", {
+			s2sErrorCounter.add(1);
+			connectionSpan.setStatus({
+				code: 2,
+				message: errObj.message
+			});
+			connectionSpan.recordException(errObj);
+			if (!opened) {
+				connectionSpan.end();
+				reject(errObj);
+			} else emitter.emit("error", {
 				code: "ws_error",
 				message: errObj.message
 			});

package/dist/server.d.ts

@@ -3,18 +3,26 @@
  *
  * `createServer()` returns a server with `listen()` for HTTP + WebSocket.
  * Calls `createDirectExecutor` + `wireSessionSocket` directly — no
- * intermediate WintercServer layer.
+ * intermediary needed.
  */
 import type { Kv } from "./kv.ts";
 import type { Logger, S2SConfig } from "./runtime.ts";
 import type { AgentDef } from "./types.ts";
+import type { VectorStore } from "./vector.ts";
+/**
+ * Configuration for a self-hosted agent server created by {@link createServer}.
+ *
+ * @public
+ */
 export type ServerOptions = {
     /** The agent definition returned by `defineAgent()`. */
-    agent: AgentDef;
+    agent: AgentDef<any>;
     /** Environment variables. Defaults to `process.env`. */
     env?: Record<string, string>;
-    /** KV store. Defaults to in-memory. */
+    /** KV store. Defaults to SQLite-backed (`.aai/local.db`). */
     kv?: Kv;
+    /** Vector store. Defaults to SQLite-backed (`.aai/local.db`). */
+    vector?: VectorStore;
     /** HTML to serve at `GET /`. */
     clientHtml?: string;
     /** Directory containing built client files (index.html + assets/). */
@@ -23,11 +31,31 @@ export type ServerOptions = {
     logger?: Logger;
     /** S2S configuration. Defaults to AssemblyAI production. */
     s2sConfig?: S2SConfig;
+    /**
+     * Timeout in ms for `session.start()` (S2S connection setup).
+     * Defaults to 10 000 (10 s). If the session doesn't initialize within
+     * this window the connection is cleaned up.
+     */
+    sessionStartTimeoutMs?: number;
+    /**
+     * Maximum time in milliseconds to wait for sessions to stop during
+     * {@link AgentServer.close | close()}. Sessions still running after this
+     * deadline are force-closed. Defaults to `30_000` (30 seconds).
+     */
+    shutdownTimeoutMs?: number;
 };
+/**
+ * Handle returned by {@link createServer} with lifecycle methods to start
+ * and stop the HTTP + WebSocket server.
+ *
+ * @public
+ */
 export type AgentServer = {
     /** Start listening on the given port. */
     listen(port?: number): Promise<void>;
     /** Stop the server. */
     close(): Promise<void>;
+    /** The port the server is listening on, or `undefined` before `listen()`. */
+    port: number | undefined;
 };
 export declare function createServer(options: ServerOptions): AgentServer;

package/dist/server.js

@@ -1,56 +1,115 @@
-import { AUDIO_FORMAT } from "./protocol.js";
+import { filterEnv } from "./_utils.js";
+import { t as createDirectExecutor } from "./direct-executor-Ca0wt5H0.js";
+import { buildReadyConfig } from "./protocol.js";
 import { DEFAULT_S2S_CONFIG, consoleLogger } from "./runtime.js";
-import { createDirectExecutor } from "./direct-executor.js";
+import { createSqliteKv } from "./sqlite-kv.js";
 import { wireSessionSocket } from "./ws-handler.js";
+import { mkdirSync } from "node:fs";
+import { WebSocketServer } from "ws";
 import { serve } from "@hono/node-server";
 import { serveStatic } from "@hono/node-server/serve-static";
 import { Hono } from "hono";
-import { WebSocketServer } from "ws";
 //#region server.ts
 /**
 * Self-hostable agent server.
 *
 * `createServer()` returns a server with `listen()` for HTTP + WebSocket.
 * Calls `createDirectExecutor` + `wireSessionSocket` directly — no
-* intermediate WintercServer layer.
+* intermediary needed.
 */
-function resolveEnv(env) {
-	return Object.fromEntries(Object.entries(env).filter(([, v]) => v !== void 0));
+/** Escape HTML special characters to prevent XSS. */
+function escapeHtml(s) {
+	return s.replace(/[&<>"']/g, (ch) => ({
+		"&": "&amp;",
+		"<": "&lt;",
+		">": "&gt;",
+		"\"": "&quot;",
+		"'": "&#39;"
+	})[ch]);
+}
+/**
+* Create an HTTP + WebSocket server for self-hosted agent deployments.
+*
+* Sets up a Hono HTTP server with a `/health` endpoint and WebSocket upgrade
+* handling. Agent tools execute directly in-process via {@link createDirectExecutor}.
+*
+* @param options - Server configuration including the agent definition, optional
+*   KV store, client assets, logger, and S2S config. See {@link ServerOptions}.
+* @returns An {@link AgentServer} with `listen()` and `close()` lifecycle methods.
+*
+* @example
+* ```ts
+* import { defineAgent } from "@alexkroman1/aai";
+* import { createServer } from "@alexkroman1/aai/server";
+*
+* const agent = defineAgent({ name: "my-agent" });
+* const server = createServer({ agent });
+* await server.listen(3000);
+* ```
+*
+* @public
+*/
+async function drainSessions(sessions, shutdownTimeoutMs, logger) {
+	if (sessions.size === 0) return;
+	let timer;
+	const timeout = new Promise((resolve) => {
+		timer = setTimeout(resolve, shutdownTimeoutMs, "timeout");
+	});
+	const graceful = Promise.allSettled([...sessions.values()].map((s) => s.stop())).then((results) => {
+		for (const r of results) if (r.status === "rejected") logger.warn(`Session stop failed during close: ${r.reason}`);
+		return "done";
+	});
+	const outcome = await Promise.race([graceful, timeout]);
+	if (timer) clearTimeout(timer);
+	if (outcome === "timeout") logger.warn(`Shutdown timeout (${shutdownTimeoutMs}ms) exceeded — force-closing ${sessions.size} remaining session(s)`);
+	sessions.clear();
 }
 function createServer(options) {
-	const { agent, kv, clientHtml, clientDir, logger = consoleLogger, s2sConfig = DEFAULT_S2S_CONFIG } = options;
+	if (options.clientHtml && options.clientDir) throw new Error("ServerOptions: clientHtml and clientDir are mutually exclusive — provide one or the other, not both.");
+	const { agent, kv, vector, clientHtml, clientDir, logger = consoleLogger, s2sConfig = DEFAULT_S2S_CONFIG, shutdownTimeoutMs = 3e4 } = options;
+	const env = filterEnv(options.env ?? (typeof process !== "undefined" ? process.env : {}));
+	const resolvedKv = kv ?? (() => {
+		mkdirSync(".aai", { recursive: true });
+		return createSqliteKv();
+	})();
 	const executor = createDirectExecutor({
 		agent,
-		env: resolveEnv(options.env ?? (typeof process !== "undefined" ? process.env : {})),
-		...kv ? { kv } : {},
+		env,
+		kv: resolvedKv,
+		...vector ? { vector } : {},
 		logger,
 		s2sConfig
 	});
 	const sessions = /* @__PURE__ */ new Map();
-	const readyConfig = {
-		audioFormat: AUDIO_FORMAT,
-		sampleRate: s2sConfig.inputSampleRate,
-		ttsSampleRate: s2sConfig.outputSampleRate
-	};
-	function handleWs(ws, skipGreeting) {
+	const readyConfig = buildReadyConfig(s2sConfig);
+	const safeAgentName = escapeHtml(agent.name);
+	function handleWs(ws, skipGreeting, resumeFrom) {
 		wireSessionSocket(ws, {
 			sessions,
 			createSession: (sid, client) => executor.createSession({
 				id: sid,
 				agent: agent.name,
 				client,
-				skipGreeting
+				skipGreeting,
+				...resumeFrom ? { resumeFrom } : {}
 			}),
 			readyConfig,
-			logger
+			logger,
+			...options.sessionStartTimeoutMs !== void 0 ? { sessionStartTimeoutMs: options.sessionStartTimeoutMs } : {},
+			...resumeFrom ? { resumeFrom } : {}
 		});
 	}
 	let serverHandle = null;
+	let listenPort;
 	return {
+		get port() {
+			return listenPort;
+		},
 		async listen(port = 3e3) {
+			if (serverHandle) throw new Error("Server is already listening");
 			const app = new Hono();
 			app.onError((err, c) => {
-				logger.error(`${c.req.method} ${new URL(c.req.url).pathname} error: ${err.message}`);
+				logger.error(`${c.req.method} ${c.req.path} error: ${err.message}`);
 				return c.json({ error: "Internal Server Error" }, 500);
 			});
 			app.use("/*", async (c, next) => {
@@ -59,7 +118,7 @@ function createServer(options) {
 				const ms = Date.now() - start;
 				const { status } = c.res;
 				const method = c.req.method;
-				const path = new URL(c.req.url).pathname;
+				const path = c.req.path;
 				if (status >= 400) logger.error(`${method} ${path} ${status} ${ms}ms`);
 				else logger.info(`${method} ${path} ${status} ${ms}ms`);
 			});
@@ -67,37 +126,52 @@ function createServer(options) {
 				status: "ok",
 				name: agent.name
 			}));
+			app.get("/kv", async (c) => {
+				const key = c.req.query("key");
+				if (!key) return c.json({ error: "Missing key query parameter" }, 400);
+				const value = await resolvedKv.get(key);
+				if (value === null) return c.json(null, 404);
+				return c.json(value);
+			});
 			if (clientDir) app.use("/*", serveStatic({ root: clientDir }));
+			const csp = "default-src 'self'; script-src 'self' blob:; style-src 'self' 'unsafe-inline' https://fonts.googleapis.com; connect-src 'self' wss: ws:; img-src 'self' data:; font-src 'self' https://fonts.gstatic.com; object-src 'none'; base-uri 'self'";
 			app.get("/", (c) => {
-				if (clientHtml) return c.html(clientHtml);
-				return c.html(`<!DOCTYPE html><html><body><h1>${agent.name}</h1><p>Agent server running.</p></body></html>`);
+				if (clientHtml) return c.html(clientHtml, 200, { "Content-Security-Policy": csp });
+				return c.html(`<!DOCTYPE html><html><body><h1>${safeAgentName}</h1><p>Agent server running.</p></body></html>`, 200, { "Content-Security-Policy": csp });
 			});
 			const nodeServer = serve({
 				fetch: app.fetch,
 				port
 			});
-			await new Promise((resolve) => {
+			await new Promise((resolve, reject) => {
 				nodeServer.on("listening", resolve);
+				nodeServer.on("error", reject);
 			});
+			const addr = nodeServer.address();
+			listenPort = typeof addr === "object" && addr ? addr.port : port;
 			const wss = new WebSocketServer({ noServer: true });
 			nodeServer.on("upgrade", (req, socket, head) => {
-				const reqUrl = new URL(req.url ?? "/", `http://localhost:${port}`);
-				const resume = reqUrl.searchParams.has("resume");
-				logger.info(`WS upgrade ${reqUrl.pathname}${resume ? " (resume)" : ""}`);
+				const reqUrl = new URL(req.url ?? "/", `http://localhost:${listenPort}`);
+				const resumeFrom = reqUrl.searchParams.get("sessionId") ?? void 0;
+				const skipGreeting = reqUrl.searchParams.has("resume") || resumeFrom !== void 0;
+				logger.info(`WS upgrade ${reqUrl.pathname}${skipGreeting ? " (resume)" : ""}`);
 				wss.handleUpgrade(req, socket, head, (ws) => {
-					handleWs(ws, resume);
+					handleWs(ws, skipGreeting, resumeFrom);
 				});
 			});
 			serverHandle = { async shutdown() {
+				await new Promise((resolve, reject) => {
+					wss.close((err) => err ? reject(err) : resolve());
+				});
 				await new Promise((resolve, reject) => {
 					nodeServer.close((err) => err ? reject(err) : resolve());
 				});
 			} };
 		},
 		async close() {
-			for (const session of sessions.values()) await session.stop();
-			sessions.clear();
+			await drainSessions(sessions, shutdownTimeoutMs, logger);
 			await serverHandle?.shutdown();
+			listenPort = void 0;
 		}
 	};
 }