agents/dist/ws-chat-transport-rWwta645.js

A home for your AI agents

import "./wire-types-CnMt6_HR.js";
import { nanoid } from "nanoid";
//#region src/chat/replay-batch.ts
/**
* Batches chunks replayed during stream resume.
*
* Replaying every chunk separately can exceed React's update limit. Adjacent
* deltas for the same message part are merged until the replay pauses or ends.
*/
var ReplayChunkBatch = class {
	constructor(controller, closeWindow = (flush) => {
		const timer = setTimeout(flush, 0);
		return () => clearTimeout(timer);
	}) {
		this.controller = controller;
		this.closeWindow = closeWindow;
		this.chunks = [];
		this.tailMergeKey = null;
	}
	get bufferedStartMessageId() {
		return this.startMessageId;
	}
	get isEmpty() {
		return this.chunks.length === 0;
	}
	push(chunk) {
		const key = mergeKeyOf(chunk);
		const tail = this.chunks[this.chunks.length - 1];
		if (tail !== void 0 && key !== null && key === this.tailMergeKey) {
			appendDeltaText(tail, chunk);
			return;
		}
		if (this.chunks.length === 0) this.cancelWindow = this.closeWindow(() => {
			this.cancelWindow = void 0;
			this.flush();
		});
		if (chunk.type === "start") this.startMessageId = chunk.messageId;
		this.chunks.push(chunk);
		this.tailMergeKey = key;
	}
	discard() {
		this.cancelWindow?.();
		this.cancelWindow = void 0;
		this.chunks = [];
		this.tailMergeKey = null;
		this.startMessageId = void 0;
	}
	flush() {
		this.cancelWindow?.();
		this.cancelWindow = void 0;
		if (this.chunks.length === 0) return;
		const batch = this.chunks;
		this.chunks = [];
		this.tailMergeKey = null;
		this.startMessageId = void 0;
		try {
			for (const chunk of batch) this.controller.enqueue(chunk);
		} catch {}
	}
	enqueueNow(chunk) {
		this.controller.enqueue(chunk);
	}
	closeNow() {
		this.controller.close();
	}
	errorNow(error) {
		this.controller.error(error);
	}
};
/** Buffers replay chunks and flushes them before live or boundary frames. */
function applyChatResponseFrame(batch, frame) {
	const endsReplayBurst = frame.replay !== true || frame.done === true || frame.replayComplete === true;
	const body = frame.body?.trim();
	if (!body) {
		if (endsReplayBurst) batch.flush();
		return;
	}
	let chunk;
	try {
		chunk = JSON.parse(body);
	} catch {
		if (endsReplayBurst) batch.flush();
		return;
	}
	if (endsReplayBurst) {
		batch.flush();
		batch.enqueueNow(chunk);
		return;
	}
	if (supersedesBufferedPass(batch, chunk, frame)) batch.discard();
	batch.push(chunk);
}
function supersedesBufferedPass(batch, chunk, frame) {
	return chunk.type === "start" && frame.continuation !== true && chunk.messageId !== void 0 && chunk.messageId === batch.bufferedStartMessageId;
}
/**
* Reports an error after buffered content. `controller.error()` would discard
* chunks that have not reached the consumer.
*/
function failChatStream(batch, message) {
	if (batch.isEmpty) {
		batch.errorNow(new Error(message));
		return;
	}
	batch.flush();
	batch.enqueueNow({
		errorText: message,
		type: "error"
	});
	batch.closeNow();
}
function mergeKeyOf(chunk) {
	switch (chunk.type) {
		case "text-delta":
		case "reasoning-delta": return chunk.providerMetadata === void 0 ? `${chunk.type}\u0000${chunk.id}` : null;
		case "tool-input-delta": return `tool-input-delta\u0000${chunk.toolCallId}`;
		default: return null;
	}
}
function appendDeltaText(target, source) {
	if (target.type === "tool-input-delta" && source.type === "tool-input-delta") {
		target.inputTextDelta += source.inputTextDelta;
		return;
	}
	if ((target.type === "text-delta" || target.type === "reasoning-delta") && (source.type === "text-delta" || source.type === "reasoning-delta")) target.delta += source.delta;
}
//#endregion
//#region src/chat/ws-chat-transport.ts
/**
* Short safety-net timeout for a resume probe when the server has said nothing.
* Under normal operation the server answers a `STREAM_RESUME_REQUEST` with
* `STREAM_RESUMING`, `STREAM_RESUME_NONE`, or `STREAM_PENDING` well before this.
*/
const RESUME_PROBE_TIMEOUT_MS = 5e3;
/**
* Extended backstop applied once the server says a turn is pending
* (`STREAM_PENDING`, #1784). The pre-stream window (queueing, MCP setup,
* debounce, model latency) can exceed the short probe timeout, and the server
* guarantees a follow-up `STREAM_RESUMING` or `STREAM_RESUME_NONE` — so we wait
* much longer (refreshed on every keep-waiting frame) but still cap it so a
* dropped follow-up degrades to a null resolve instead of hanging forever.
*/
const RESUME_PENDING_TIMEOUT_MS = 6e4;
/**
* ChatTransport that sends messages over WebSocket and returns a
* ReadableStream<UIMessageChunk> that the AI SDK's useChat consumes directly.
*
* This low-level transport handles new and regenerated request streams plus
* cancellation. Higher-level protocol coordination such as automatic reconnect
* resume, cross-tab transcript synchronization, and client-tool continuations
* is provided by integrations such as `useAgentChat`.
*/
var WebSocketChatTransport = class {
	constructor(options) {
		this._resumeResolver = null;
		this._resumeNoneResolver = null;
		this._onStreamPending = null;
		this._retryResumeProbe = null;
		this._expectToolContinuation = false;
		this._abortToolContinuation = null;
		this._activeServerTurnId = null;
		this._cancelAttachedStream = null;
		this._detachResumeStream = null;
		this.agent = options.agent;
		this.prepareBody = options.prepareBody;
		this.activeRequestIds = options.activeRequestIds;
		this.cancelOnClientAbort = options.cancelOnClientAbort ?? false;
	}
	/**
	* Point the singleton transport at a new Agent connection. A pending resolver
	* belongs to the old Chat/socket generation and must settle before messages
	* from the replacement connection can be consumed (#1914 review).
	*/
	setAgent(agent) {
		if (this.agent === agent) return;
		this.resetResumeState();
		this.agent = agent;
	}
	setCancelOnClientAbort(cancelOnClientAbort) {
		this.cancelOnClientAbort = cancelOnClientAbort;
	}
	/**
	* Explicitly cancel the active server turn, if any.
	* This is separate from generic client-side abort/cancel lifecycle so
	* clients can detach locally without stopping server work.
	*/
	cancelActiveServerTurn() {
		const requestId = this._activeServerTurnId;
		let cancelledRequest = false;
		if (requestId) {
			this.sendCancelFrame(requestId);
			this._cancelAttachedStream?.();
			this.clearActiveServerTurn(requestId);
			cancelledRequest = true;
		}
		const cancelledToolContinuation = this.abortActiveToolContinuation();
		return cancelledRequest || cancelledToolContinuation;
	}
	sendCancelFrame(requestId) {
		try {
			this.agent.send(JSON.stringify({
				id: requestId,
				type: "cf_agent_chat_request_cancel"
			}));
		} catch {}
	}
	setActiveServerTurn(requestId, cancelAttachedStream) {
		this._activeServerTurnId = requestId;
		this._cancelAttachedStream = cancelAttachedStream;
	}
	clearActiveServerTurn(requestId) {
		if (this._activeServerTurnId === requestId) {
			this._activeServerTurnId = null;
			this._cancelAttachedStream = null;
		}
	}
	/**
	* Mark that the next reconnectToStream() call should attach to a
	* server-initiated tool continuation rather than a page-load resume.
	*/
	expectToolContinuation() {
		this._expectToolContinuation = true;
	}
	/**
	* Abort the active client-side tool continuation stream, if one is attached
	* to a server request id.
	*/
	abortActiveToolContinuation() {
		return this._abortToolContinuation?.() ?? false;
	}
	/**
	* True when the transport is waiting for a resume handshake.
	*/
	isAwaitingResume() {
		return this._resumeResolver !== null || this._resumeNoneResolver !== null;
	}
	/**
	* Settle and detach the current handshake without interpreting it as a
	* server-idle response. Used when the owning hook/agent generation changes.
	*/
	cancelPendingResume() {
		const resolveNone = this._resumeNoneResolver;
		if (!resolveNone) return false;
		return resolveNone({});
	}
	/**
	* Invalidate all client-side resume state for an obsolete hook/agent
	* generation without cancelling its durable server turn.
	*/
	resetResumeState() {
		this._expectToolContinuation = false;
		this.cancelPendingResume();
		this._detachResumeStream?.();
	}
	/**
	* Re-send the active handshake request on the latest socket generation. This
	* preserves one AI SDK resume operation while recovering a request/reply lost
	* with the previous WebSocket.
	*/
	retryPendingResume() {
		const retry = this._retryResumeProbe;
		if (!retry) return false;
		retry();
		return true;
	}
	/**
	* Called by onAgentMessage when it receives CF_AGENT_STREAM_RESUMING.
	* If reconnectToStream is waiting, this handles the resume handshake
	* (ACK + stream creation) and returns true. Otherwise returns false
	* so the caller can use its own fallback path.
	*/
	handleStreamResuming(data) {
		if (!this._resumeResolver) return false;
		this._resumeResolver(data);
		return true;
	}
	/**
	* Called by onAgentMessage when it receives CF_AGENT_STREAM_RESUME_NONE.
	* If reconnectToStream is waiting, resolves the promise with null
	* immediately (no 5-second timeout). Returns true if handled.
	*/
	handleStreamResumeNone(data = {}) {
		if (!this._resumeNoneResolver) return false;
		return this._resumeNoneResolver(data);
	}
	/**
	* Called by onAgentMessage when it receives CF_AGENT_STREAM_PENDING (#1784):
	* the server accepted a turn but its stream has not started yet. If a resume
	* path is awaiting, extend its probe timeout (so it keeps waiting for the
	* eventual STREAM_RESUMING / STREAM_RESUME_NONE instead of resolving null
	* after the short window). Returns true if a waiting path consumed it.
	*/
	handleStreamPending() {
		if (!this._onStreamPending) return false;
		this._onStreamPending();
		return true;
	}
	/**
	* Called by the hook's shared message handler when a server turn finishes
	* outside the currently attached transport stream, such as after local-only
	* client cleanup.
	*/
	handleServerTurnCompleted(requestId) {
		this.clearActiveServerTurn(requestId);
	}
	/**
	* Register a server turn that is being rendered outside a transport-owned
	* stream, such as the hook's fallback cross-tab/resume observer path.
	*/
	observeServerTurn(requestId) {
		this.setActiveServerTurn(requestId, null);
	}
	async sendMessages(options) {
		const requestId = nanoid(8);
		const abortController = new AbortController();
		let completed = false;
		let requestSent = false;
		let extraBody = {};
		if (this.prepareBody) extraBody = await this.prepareBody({
			messages: options.messages,
			trigger: options.trigger,
			messageId: options.messageId
		});
		if (options.body) extraBody = {
			...extraBody,
			...options.body
		};
		const bodyPayload = JSON.stringify({
			messages: options.messages,
			trigger: options.trigger,
			...extraBody
		});
		this.activeRequestIds?.add(requestId);
		const agent = this.agent;
		const activeIds = this.activeRequestIds;
		const finish = (action, keepId = false, clearServerTurn = true) => {
			if (completed) return;
			completed = true;
			if (clearServerTurn) this.clearActiveServerTurn(requestId);
			try {
				action();
			} catch {}
			if (!keepId) activeIds?.delete(requestId);
			abortController.abort();
		};
		const abortError = /* @__PURE__ */ new Error("Aborted");
		abortError.name = "AbortError";
		const cancelActiveRequest = () => {
			if (completed) return false;
			finish(() => streamController.error(abortError), true);
			return true;
		};
		this.setActiveServerTurn(requestId, cancelActiveRequest);
		const onAbort = () => {
			if (completed) return;
			if (this.cancelOnClientAbort) {
				if (requestSent) this.sendCancelFrame(requestId);
				finish(() => streamController.error(abortError), requestSent);
			} else finish(() => streamController.error(abortError), false, !requestSent);
		};
		let streamController;
		const stream = new ReadableStream({
			start(controller) {
				streamController = controller;
				const onMessage = (event) => {
					try {
						const data = JSON.parse(event.data);
						if (data.type !== "cf_agent_use_chat_response") return;
						if (data.id !== requestId) return;
						if (data.error) {
							finish(() => controller.error(new Error(data.body || "Stream error")));
							return;
						}
						if (data.body?.trim()) try {
							const chunk = JSON.parse(data.body);
							controller.enqueue(chunk);
						} catch {}
						if (data.done) finish(() => controller.close());
					} catch {}
				};
				const onClose = () => {
					finish(() => controller.close(), false, false);
				};
				agent.addEventListener("message", onMessage, { signal: abortController.signal });
				agent.addEventListener("close", onClose, { signal: abortController.signal });
			},
			cancel() {
				onAbort();
			}
		});
		if (options.abortSignal) {
			options.abortSignal.addEventListener("abort", onAbort, { once: true });
			if (options.abortSignal.aborted) onAbort();
		}
		if (completed) return stream;
		requestSent = true;
		agent.send(JSON.stringify({
			id: requestId,
			init: {
				method: "POST",
				body: bodyPayload
			},
			type: "cf_agent_use_chat_request"
		}));
		return stream;
	}
	async reconnectToStream(_options) {
		if (this.isAwaitingResume()) return null;
		if (this._expectToolContinuation) {
			this._expectToolContinuation = false;
			return this._createToolContinuationStream();
		}
		const activeIds = this.activeRequestIds;
		const probeId = nanoid(8);
		return new Promise((resolve) => {
			let resolved = false;
			let timeout;
			let resumeResolver = null;
			let resumeNoneResolver = null;
			let onStreamPending = null;
			let retryResumeProbe = null;
			const clearOwnedCallbacks = () => {
				if (resumeResolver && this._resumeResolver === resumeResolver) this._resumeResolver = null;
				if (resumeNoneResolver && this._resumeNoneResolver === resumeNoneResolver) this._resumeNoneResolver = null;
				if (onStreamPending && this._onStreamPending === onStreamPending) this._onStreamPending = null;
				if (retryResumeProbe && this._retryResumeProbe === retryResumeProbe) this._retryResumeProbe = null;
			};
			const done = (value) => {
				if (resolved) return;
				resolved = true;
				clearOwnedCallbacks();
				if (timeout) clearTimeout(timeout);
				resolve(value);
			};
			const armTimeout = (delay) => {
				if (timeout) clearTimeout(timeout);
				timeout = setTimeout(() => done(null), delay);
			};
			onStreamPending = () => {
				if (resolved) return;
				armTimeout(RESUME_PENDING_TIMEOUT_MS);
			};
			resumeNoneResolver = (data) => {
				if (data.probeId && data.probeId !== probeId) return false;
				done(null);
				return true;
			};
			resumeResolver = (data) => {
				const requestId = data.id;
				activeIds?.add(requestId);
				const stream = this._createResumeStream(requestId);
				this.agent.send(JSON.stringify({
					type: "cf_agent_stream_resume_ack",
					id: requestId
				}));
				done(stream);
			};
			retryResumeProbe = () => {
				if (resolved) return;
				armTimeout(RESUME_PROBE_TIMEOUT_MS);
				try {
					this.agent.send(JSON.stringify({
						type: "cf_agent_stream_resume_request",
						probeId
					}));
				} catch {}
			};
			this._onStreamPending = onStreamPending;
			this._resumeNoneResolver = resumeNoneResolver;
			this._resumeResolver = resumeResolver;
			this._retryResumeProbe = retryResumeProbe;
			retryResumeProbe();
		});
	}
	/**
	* Creates a deferred ReadableStream for client-side tool continuations.
	* The stream is returned immediately so AI SDK status becomes "submitted"
	* right after addToolOutput()/addToolApprovalResponse(), then it waits for
	* the server to announce the continuation via STREAM_RESUMING.
	*/
	_createToolContinuationStream() {
		const agent = this.agent;
		const activeIds = this.activeRequestIds;
		const streamController = new AbortController();
		const abortError = /* @__PURE__ */ new Error("Aborted");
		abortError.name = "AbortError";
		let completed = false;
		let requestId = null;
		let readerController = null;
		let replayBatch = null;
		const probeId = nanoid(8);
		let onResumeRef = null;
		let onResumeNoneRef = null;
		let onStreamPendingRef = null;
		let retryResumeProbeRef = null;
		let abortToolContinuationRef = null;
		let detachResumeStreamRef = null;
		const clearOwnedHandshake = () => {
			if (onResumeRef && this._resumeResolver === onResumeRef) this._resumeResolver = null;
			if (onResumeNoneRef && this._resumeNoneResolver === onResumeNoneRef) this._resumeNoneResolver = null;
			if (onStreamPendingRef && this._onStreamPending === onStreamPendingRef) this._onStreamPending = null;
			if (retryResumeProbeRef && this._retryResumeProbe === retryResumeProbeRef) this._retryResumeProbe = null;
		};
		const finish = (action, keepRequestId = false) => {
			if (completed) return;
			completed = true;
			if (this._abortToolContinuation === abortToolContinuationRef) this._abortToolContinuation = null;
			if (this._detachResumeStream === detachResumeStreamRef) this._detachResumeStream = null;
			clearOwnedHandshake();
			try {
				action();
			} catch {}
			if (requestId && !keepRequestId) activeIds?.delete(requestId);
			streamController.abort();
		};
		const transport = this;
		abortToolContinuationRef = () => {
			if (completed) return false;
			if (requestId === null) {
				finish(() => readerController?.error(abortError));
				return true;
			}
			try {
				agent.send(JSON.stringify({
					type: "cf_agent_chat_request_cancel",
					id: requestId
				}));
			} catch {}
			finish(() => readerController?.error(abortError), true);
			return true;
		};
		this._abortToolContinuation = abortToolContinuationRef;
		detachResumeStreamRef = () => {
			if (completed) return false;
			finish(() => {
				replayBatch?.flush();
				readerController?.close();
			});
			return true;
		};
		this._detachResumeStream = detachResumeStreamRef;
		return new ReadableStream({
			start(controller) {
				readerController = controller;
				const batch = new ReplayChunkBatch(controller);
				replayBatch = batch;
				let timeout;
				const armTimeout = (delay) => {
					if (timeout) clearTimeout(timeout);
					timeout = setTimeout(() => finish(() => controller.close()), delay);
				};
				const onResumeNone = (data) => {
					if (data.probeId && data.probeId !== probeId) return false;
					finish(() => controller.close());
					return true;
				};
				const onResume = (data) => {
					if (requestId) return;
					requestId = data.id;
					activeIds?.add(requestId);
					clearOwnedHandshake();
					if (timeout) clearTimeout(timeout);
					agent.send(JSON.stringify({
						type: "cf_agent_stream_resume_ack",
						id: requestId
					}));
				};
				const onStreamPending = () => {
					if (completed) return;
					armTimeout(RESUME_PENDING_TIMEOUT_MS);
				};
				const retryResumeProbe = () => {
					if (completed || requestId !== null) return;
					armTimeout(RESUME_PROBE_TIMEOUT_MS);
					try {
						transport.agent.send(JSON.stringify({
							type: "cf_agent_stream_resume_request",
							probeId
						}));
					} catch {}
				};
				onResumeRef = onResume;
				onResumeNoneRef = onResumeNone;
				onStreamPendingRef = onStreamPending;
				retryResumeProbeRef = retryResumeProbe;
				transport._resumeResolver = onResume;
				transport._resumeNoneResolver = onResumeNone;
				transport._onStreamPending = onStreamPending;
				transport._retryResumeProbe = retryResumeProbe;
				const onMessage = (event) => {
					try {
						const data = JSON.parse(event.data);
						if (data.type !== "cf_agent_use_chat_response" || requestId == null || data.id !== requestId) return;
						if (data.error) {
							finish(() => failChatStream(batch, data.body || "Stream error"));
							return;
						}
						applyChatResponseFrame(batch, data);
						if (data.done) finish(() => controller.close());
					} catch {}
				};
				const onClose = () => finish(() => {
					batch.flush();
					controller.close();
				});
				agent.addEventListener("message", onMessage, { signal: streamController.signal });
				agent.addEventListener("close", onClose, { signal: streamController.signal });
				retryResumeProbe();
			},
			cancel() {
				if (requestId && transport.cancelOnClientAbort) {
					transport.sendCancelFrame(requestId);
					finish(() => {}, true);
				} else finish(() => {});
			}
		});
	}
	/**
	* Creates a ReadableStream that receives resumed stream chunks
	* and forwards them to useChat as UIMessageChunk objects.
	*/
	_createResumeStream(requestId) {
		const agent = this.agent;
		const activeIds = this.activeRequestIds;
		const chunkController = new AbortController();
		const abortError = /* @__PURE__ */ new Error("Aborted");
		abortError.name = "AbortError";
		let completed = false;
		let detachResumeStream = null;
		const finish = (action, keepId = false, clearServerTurn = true) => {
			if (completed) return;
			completed = true;
			if (clearServerTurn) this.clearActiveServerTurn(requestId);
			if (this._detachResumeStream === detachResumeStream) this._detachResumeStream = null;
			try {
				action();
			} catch {}
			if (!keepId) activeIds?.delete(requestId);
			chunkController.abort();
		};
		let streamController = null;
		let replayBatch = null;
		const cancelActiveRequest = () => {
			if (completed) return false;
			finish(() => streamController?.error(abortError), true);
			return true;
		};
		this.setActiveServerTurn(requestId, cancelActiveRequest);
		const transport = this;
		detachResumeStream = () => {
			if (completed) return false;
			finish(() => {
				replayBatch?.flush();
				streamController?.close();
			});
			return true;
		};
		this._detachResumeStream = detachResumeStream;
		return new ReadableStream({
			start(controller) {
				streamController = controller;
				const batch = new ReplayChunkBatch(controller);
				replayBatch = batch;
				const onMessage = (event) => {
					try {
						const data = JSON.parse(event.data);
						if (data.type !== "cf_agent_use_chat_response") return;
						if (data.id !== requestId) return;
						if (data.error) {
							finish(() => failChatStream(batch, data.body || "Stream error"));
							return;
						}
						applyChatResponseFrame(batch, data);
						if (data.done) finish(() => controller.close());
					} catch {}
				};
				const onClose = () => {
					finish(() => {
						batch.flush();
						controller.close();
					}, false, false);
				};
				agent.addEventListener("message", onMessage, { signal: chunkController.signal });
				agent.addEventListener("close", onClose, { signal: chunkController.signal });
			},
			cancel() {
				if (transport.cancelOnClientAbort) {
					transport.sendCancelFrame(requestId);
					finish(() => {}, true);
				} else finish(() => {}, false, false);
			}
		});
	}
};
//#endregion
export { WebSocketChatTransport as t };

//# sourceMappingURL=ws-chat-transport-rWwta645.js.map