@tanstack/ai-client 0.23.2 → 0.24.0

package/dist/esm/connection-adapters.d.ts

@@ -49,6 +49,32 @@ export interface ReconnectOptions {
     /** Delay between reconnect attempts, in ms, to avoid hammering. Default 250. */
     delayMs?: number;
 }
+/**
+ * Reconnect bookkeeping shared by every resumable-stream driver: de-dupes
+ * offsets, tracks the last acknowledged offset, honors the SSE empty-id reset
+ * convention, and bounds consecutive no-progress reconnects behind a
+ * throttling delay. Extracted out of {@link resumableStream} so a WebSocket
+ * reconnect driver can reuse the exact same semantics.
+ */
+export interface ReconnectTracker {
+    /** The most recently accepted (non-duplicate, non-empty) offset, if any. */
+    readonly lastEventId: string | undefined;
+    /**
+     * Record an incoming offset. Returns `'reset'` for an empty id (SSE's
+     * resume-cursor reset — clears the de-dupe set and `lastEventId`),
+     * `'duplicate'` for an already-seen id, and `'new'` otherwise (including
+     * `undefined`, which is untracked — no offset to remember).
+     */
+    note: (id: string | undefined) => 'new' | 'duplicate' | 'reset';
+    /**
+     * Throttle before a reconnect attempt. Resets the no-progress counter when
+     * `madeProgress` is true; otherwise increments it and throws
+     * {@link StreamReconnectLimitError} once it exceeds the configured ceiling.
+     */
+    waitBeforeReconnect: (madeProgress: boolean, signal?: AbortSignal) => Promise<void>;
+}
+/** Create a {@link ReconnectTracker} bound to the given reconnect bounds. */
+export declare function createReconnectTracker(options?: ReconnectOptions): ReconnectTracker;
 /**
  * Per-send context provided by the chat client to the connection adapter.
  * The adapter combines this with serialized messages to build a full
@@ -327,6 +353,29 @@ export declare function xhrServerSentEvents(url: string | (() => string), option
  * (bare-line) response is a single plain request, exactly as before.
  */
 export declare function xhrHttpStream(url: string | (() => string), options?: XhrConnectionOptionsResolver): ResumableConnectConnectionAdapter;
+export interface WebSocketConnectionOptions {
+    protocols?: string | Array<string>;
+    body?: Record<string, unknown>;
+    reconnect?: ReconnectOptions;
+    /** Override the WebSocket implementation (tests / non-browser runtimes). */
+    WebSocketImpl?: typeof WebSocket;
+}
+/**
+ * Full-duplex, conversation-scoped WebSocket connection adapter. Pairs with the
+ * server `toWebSocketResponse` / `toWebSocketStream`. `send()` writes a
+ * RunAgentInput frame; `subscribe()` yields inbound chunks.
+ *
+ * Resumable: `send()` establishes a run session backed by a
+ * {@link createReconnectTracker}. If the socket closes before a terminal
+ * (`RUN_FINISHED`/`RUN_ERROR`) chunk is seen and the run is durable
+ * (offset-tagged `{ id, chunk }` envelopes), the socket is reopened at
+ * `?runId=&offset=<lastEventId>`, de-duping the replayed boundary. A drop with
+ * no offset ever observed (non-durable) surfaces {@link StreamReadError}
+ * instead of reconnecting — there is nothing to resume from.
+ */
+export declare function webSocket(url: string | (() => string), options?: WebSocketConnectionOptions): SubscribeConnectionAdapter & {
+    joinRun: (runId: string, abortSignal?: AbortSignal) => AsyncIterable<StreamChunk>;
+};
 /**
  * Optional persistence handlers for the lightweight adapters (`stream()`,
  * `rpcStream()`). These are one-shot, request-scoped calls with no built-in

package/dist/esm/connection-adapters.js

@@ -68,6 +68,38 @@ function resolveReconnectOptions(options) {
 		delayMs
 	};
 }
+/** Create a {@link ReconnectTracker} bound to the given reconnect bounds. */
+function createReconnectTracker(options) {
+	const reconnect = resolveReconnectOptions(options);
+	const seen = /* @__PURE__ */ new Set();
+	let lastEventId;
+	let reconnectAttempts = 0;
+	return {
+		get lastEventId() {
+			return lastEventId;
+		},
+		note(id) {
+			if (id === void 0) return "new";
+			if (id === "") {
+				lastEventId = void 0;
+				seen.clear();
+				return "reset";
+			}
+			if (seen.has(id)) return "duplicate";
+			seen.add(id);
+			lastEventId = id;
+			return "new";
+		},
+		async waitBeforeReconnect(madeProgress, signal) {
+			if (madeProgress) reconnectAttempts = 0;
+			else {
+				reconnectAttempts += 1;
+				if (reconnectAttempts > reconnect.maxAttempts) throw new StreamReconnectLimitError(reconnect.maxAttempts);
+			}
+			await abortableDelay(reconnect.delayMs, signal);
+		}
+	};
+}
 /** Resolve after `ms`, or immediately once `signal` aborts. Never rejects. */
 function abortableDelay(ms, signal) {
 	if (ms <= 0 || signal?.aborted) return Promise.resolve();
@@ -350,50 +382,32 @@ function fetchEventSource(fetchClient, url, requestInit, parseResponse) {
 * only difference being the {@link StreamEventSource} they pass in.
 */
 async function* resumableStream(openEventSource, abortSignal, reconnectOptions) {
-	const seen = /* @__PURE__ */ new Set();
-	let lastEventId;
-	const reconnect = resolveReconnectOptions(reconnectOptions);
-	let reconnectAttempts = 0;
-	async function waitBeforeReconnect(madeProgress) {
-		if (madeProgress) reconnectAttempts = 0;
-		else {
-			reconnectAttempts += 1;
-			if (reconnectAttempts > reconnect.maxAttempts) throw new StreamReconnectLimitError(reconnect.maxAttempts);
-		}
-		await abortableDelay(reconnect.delayMs, abortSignal);
-	}
+	const tracker = createReconnectTracker(reconnectOptions);
 	for (;;) {
 		if (abortSignal?.aborted) return;
-		const extraHeaders = lastEventId !== void 0 ? { "Last-Event-ID": lastEventId } : {};
+		const extraHeaders = tracker.lastEventId !== void 0 ? { "Last-Event-ID": tracker.lastEventId } : {};
 		let sawTerminal = false;
 		let progressed = false;
 		try {
 			for await (const { chunk, id } of openEventSource(extraHeaders, abortSignal)) {
-				if (id !== void 0) if (id === "") {
-					lastEventId = void 0;
-					seen.clear();
-				} else {
-					if (seen.has(id)) continue;
-					seen.add(id);
-					lastEventId = id;
-				}
+				if (tracker.note(id) === "duplicate") continue;
 				progressed = true;
 				if (chunk.type === "RUN_FINISHED" || chunk.type === "RUN_ERROR") sawTerminal = true;
 				yield chunk;
 			}
 		} catch (error) {
 			if (abortSignal?.aborted) return;
-			if ((error instanceof StreamTruncatedError || error instanceof StreamReadError) && lastEventId !== void 0) {
-				await waitBeforeReconnect(progressed);
+			if ((error instanceof StreamTruncatedError || error instanceof StreamReadError) && tracker.lastEventId !== void 0) {
+				await tracker.waitBeforeReconnect(progressed, abortSignal);
 				continue;
 			}
 			throw error;
 		}
 		if (abortSignal?.aborted) return;
 		if (sawTerminal) return;
-		if (lastEventId !== void 0) {
+		if (tracker.lastEventId !== void 0) {
 			if (progressed) {
-				await waitBeforeReconnect(true);
+				await tracker.waitBeforeReconnect(true, abortSignal);
 				continue;
 			}
 			throw new DurableStreamIncompleteError();
@@ -900,6 +914,232 @@ function xhrHttpStream(url, options = {}) {
 		}
 	};
 }
+function runIdQuery(url, runId) {
+	return runId ? withSearchParams(url, { runId }) : url;
+}
+function isPingFrame(parsed) {
+	return typeof parsed === "object" && parsed !== null && parsed.type === "ping";
+}
+/**
+* A push→pull bridge from socket callbacks to an async iterable: chunks queue
+* until the consumer pulls, a recorded failure rejects the iterator, and
+* `end()` (or the abort signal) finishes it cleanly. Shared by `webSocket()`'s
+* `subscribe()` and `joinRun()`.
+*/
+function createChunkPipe(abortSignal, onFinally) {
+	const queue = [];
+	const waiters = [];
+	let failure;
+	let ended = false;
+	const wake = () => waiters.shift()?.(null);
+	const push = (chunk) => {
+		const w = waiters.shift();
+		if (w) w(chunk);
+		else queue.push(chunk);
+	};
+	const fail = (error) => {
+		failure = error;
+		wake();
+	};
+	const end = () => {
+		ended = true;
+		wake();
+	};
+	const onAbort = () => wake();
+	abortSignal?.addEventListener("abort", onAbort);
+	return {
+		push,
+		fail,
+		end,
+		iterable: (async function* () {
+			try {
+				while (!abortSignal?.aborted) {
+					const buffered = queue.shift();
+					if (buffered !== void 0) {
+						yield buffered;
+						continue;
+					}
+					if (failure !== void 0) throw failure;
+					if (ended) return;
+					const chunk = await new Promise((r) => waiters.push(r));
+					if (failure !== void 0) throw failure;
+					if (chunk === null) return;
+					yield chunk;
+				}
+			} finally {
+				abortSignal?.removeEventListener("abort", onAbort);
+				onFinally();
+			}
+		})()
+	};
+}
+/**
+* Full-duplex, conversation-scoped WebSocket connection adapter. Pairs with the
+* server `toWebSocketResponse` / `toWebSocketStream`. `send()` writes a
+* RunAgentInput frame; `subscribe()` yields inbound chunks.
+*
+* Resumable: `send()` establishes a run session backed by a
+* {@link createReconnectTracker}. If the socket closes before a terminal
+* (`RUN_FINISHED`/`RUN_ERROR`) chunk is seen and the run is durable
+* (offset-tagged `{ id, chunk }` envelopes), the socket is reopened at
+* `?runId=&offset=<lastEventId>`, de-duping the replayed boundary. A drop with
+* no offset ever observed (non-durable) surfaces {@link StreamReadError}
+* instead of reconnecting — there is nothing to resume from.
+*/
+function webSocket(url, options = {}) {
+	const Impl = options.WebSocketImpl ?? WebSocket;
+	let socket;
+	let socketMode;
+	let openPromise;
+	const listeners = /* @__PURE__ */ new Set();
+	let currentSession;
+	function failAll(error) {
+		for (const l of listeners) l.fail(error);
+	}
+	function openOnce(target, mode) {
+		if (socket && socket.readyState <= 1 && mode === "run" && socketMode === "run") return socket;
+		const prior = socket;
+		const ws = options.protocols ? new Impl(target, options.protocols) : new Impl(target);
+		socket = ws;
+		socketMode = mode;
+		openPromise = new Promise((resolve, reject) => {
+			ws.onopen = () => resolve();
+			ws.onerror = (e) => reject(new StreamReadError(e));
+		});
+		openPromise.catch(() => {});
+		ws.onmessage = (event) => {
+			if (ws !== socket) return;
+			let parsed;
+			try {
+				parsed = JSON.parse(String(event.data));
+			} catch (error) {
+				failAll(new StreamReadError(error));
+				return;
+			}
+			if (isPingFrame(parsed)) return;
+			const envelopeId = isNdjsonEnvelope(parsed) ? parsed.id : void 0;
+			const chunk = isNdjsonEnvelope(parsed) ? parsed.chunk : parsed;
+			const session = currentSession;
+			if (session) {
+				if (session.tracker.note(envelopeId) === "duplicate") return;
+				session.progressed = true;
+				if (session.runId === void 0) session.runId = getChunkRunId(chunk);
+				if (chunk.type === "RUN_FINISHED" || chunk.type === "RUN_ERROR") session.sawTerminal = true;
+			}
+			for (const l of listeners) l.push(chunk);
+		};
+		ws.onclose = () => {
+			if (ws !== socket) return;
+			const session = currentSession;
+			if (!session) {
+				failAll(new StreamReadError(/* @__PURE__ */ new Error("WebSocket connection closed")));
+				return;
+			}
+			if (session.signal?.aborted || session.sawTerminal) return;
+			const lastEventId = session.tracker.lastEventId;
+			if (lastEventId === void 0) {
+				currentSession = void 0;
+				failAll(new StreamReadError(/* @__PURE__ */ new Error("WebSocket connection closed")));
+				return;
+			}
+			reconnect(session, lastEventId);
+		};
+		if (prior && prior.readyState <= 1) prior.close();
+		return ws;
+	}
+	async function reconnect(session, offset) {
+		try {
+			await session.tracker.waitBeforeReconnect(session.progressed, session.signal);
+		} catch (error) {
+			if (currentSession === session) currentSession = void 0;
+			failAll(error);
+			return;
+		}
+		if (session.signal?.aborted) return;
+		if (currentSession !== session) return;
+		if (socket && socket.readyState <= 1) return;
+		session.progressed = false;
+		openOnce(withSearchParams(typeof url === "function" ? url() : url, {
+			...session.runId !== void 0 ? { runId: session.runId } : {},
+			offset
+		}), "resume");
+	}
+	function waitOpen(ws) {
+		if (ws.readyState === 1) return Promise.resolve();
+		return openPromise ?? Promise.resolve();
+	}
+	return {
+		subscribe(abortSignal) {
+			const pipe = createChunkPipe(abortSignal, () => listeners.delete(sink));
+			const sink = {
+				push: pipe.push,
+				fail: pipe.fail
+			};
+			listeners.add(sink);
+			return pipe.iterable;
+		},
+		async send(messages, data, abortSignal, runContext) {
+			const ws = openOnce(runIdQuery(typeof url === "function" ? url() : url, runContext?.runId), "run");
+			await waitOpen(ws);
+			if (!currentSession || currentSession.runId !== runContext?.runId) currentSession = {
+				runId: runContext?.runId,
+				tracker: createReconnectTracker(options.reconnect),
+				sawTerminal: false,
+				progressed: false,
+				signal: abortSignal
+			};
+			else {
+				currentSession.signal = abortSignal;
+				currentSession.sawTerminal = false;
+				currentSession.progressed = false;
+			}
+			const session = currentSession;
+			abortSignal?.addEventListener("abort", () => {
+				const abortRunId = session.runId;
+				const live = socket;
+				if (abortRunId === void 0 || session.sawTerminal || socketMode !== "run" || live === void 0 || live.readyState !== 1) return;
+				try {
+					live.send(JSON.stringify({
+						type: "abort",
+						runId: abortRunId
+					}));
+				} catch {}
+			}, { once: true });
+			const body = buildRunAgentInputBody(messages, data, runContext, { body: options.body });
+			ws.send(JSON.stringify(body));
+		},
+		joinRun(runId, abortSignal) {
+			const target = withSearchParams(typeof url === "function" ? url() : url, {
+				offset: "-1",
+				runId
+			});
+			const ws = options.protocols ? new Impl(target, options.protocols) : new Impl(target);
+			const pipe = createChunkPipe(abortSignal, () => {
+				if (ws.readyState <= 1) ws.close();
+			});
+			ws.onmessage = (event) => {
+				let parsed;
+				try {
+					parsed = JSON.parse(String(event.data));
+				} catch (error) {
+					pipe.fail(new StreamReadError(error));
+					return;
+				}
+				if (isPingFrame(parsed)) return;
+				pipe.push(isNdjsonEnvelope(parsed) ? parsed.chunk : parsed);
+			};
+			ws.onclose = (event) => {
+				if (event?.code === 1e3) {
+					pipe.end();
+					return;
+				}
+				const detail = event ? `${event.code}${event.reason ? `: ${event.reason}` : ""}` : "unknown";
+				pipe.fail(new StreamReadError(/* @__PURE__ */ new Error(`WebSocket connection closed (${detail})`)));
+			};
+			return pipe.iterable;
+		}
+	};
+}
 /**
 * Create a direct stream connection adapter (for server functions or direct streams)
 *
@@ -1021,6 +1261,6 @@ function rpcStream(rpcCall, handlers) {
 	};
 }
 //#endregion
-export { DurableStreamIncompleteError, StreamReconnectLimitError, StreamTruncatedError, fetchHttpStream, fetchServerSentEvents, fetcherToConnectionAdapter, getChunkRunId, normalizeConnectionAdapter, rpcStream, stream, xhrHttpStream, xhrServerSentEvents };
+export { DurableStreamIncompleteError, StreamReconnectLimitError, StreamTruncatedError, createReconnectTracker, fetchHttpStream, fetchServerSentEvents, fetcherToConnectionAdapter, getChunkRunId, normalizeConnectionAdapter, rpcStream, stream, webSocket, xhrHttpStream, xhrServerSentEvents };
 
 //# sourceMappingURL=connection-adapters.js.map