agents 0.20.1 → 0.22.0

package/dist/chat/react.js

@@ -1,559 +1,39 @@
-import { i as transition, r as STREAM_RESUME_NONE_REASONS } from "../wire-types-CU9rLoeS.js";
+import { n as STREAM_RESUME_NONE_REASONS, r as transition } from "../protocol-B0nh6KNf.js";
+import "../wire-types-CnMt6_HR.js";
+import { t as WebSocketChatTransport } from "../ws-chat-transport-rWwta645.js";
 import { getToolName, isToolUIPart } from "ai";
 import { nanoid } from "nanoid";
 import { use, useCallback, useEffect, useMemo, useRef, useState } from "react";
 import { useChat } from "@ai-sdk/react";
-//#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.
+* Picks the throttle from an explicit caller value, the deprecated alias, or
+* the default — in that order. `false` turns throttling off.
 */
-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;
+function resolveChatThrottleMs(options) {
+	if (options.throttle === false) return void 0;
+	return options.throttle ?? options.experimental_throttle ?? 50;
+}
 /**
-* ChatTransport that sends messages over WebSocket and returns a
-* ReadableStream<UIMessageChunk> that the AI SDK's useChat consumes directly.
-* No fake fetch, no Response reconstruction, no double SSE parsing.
+* The throttle spelled under both option names, or neither name when it is off.
+*
+* The two names are not interchangeable across the peer range. `@ai-sdk/react`
+* v3 only reads `experimental_throttle`; v4 renamed it to `throttle` and reads
+* `throttle ?? experimental_throttle`. Our peer range allows both majors, so
+* sending one name would silently do nothing on the other. Unknown option keys
+* are ignored by both, so sending both names is safe.
+*
+* Turning throttling off omits both names rather than inventing a numeric
+* value. Both majors also treat an explicit `0` as unthrottled, so callers that
+* pass `0` keep that value under both spellings.
 */
-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;
-		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(() => readerController?.close());
-			return true;
-		};
-		this._detachResumeStream = detachResumeStreamRef;
-		return new ReadableStream({
-			start(controller) {
-				readerController = controller;
-				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(() => 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());
-				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;
-		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(() => streamController?.close());
-			return true;
-		};
-		this._detachResumeStream = detachResumeStream;
-		return 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: chunkController.signal });
-				agent.addEventListener("close", onClose, { signal: chunkController.signal });
-			},
-			cancel() {
-				if (transport.cancelOnClientAbort) {
-					transport.sendCancelFrame(requestId);
-					finish(() => {}, true);
-				} else finish(() => {}, false, false);
-			}
-		});
-	}
-};
+function chatThrottleOptions(options) {
+	const ms = resolveChatThrottleMs(options);
+	if (ms === void 0) return {};
+	return {
+		experimental_throttle: ms,
+		throttle: ms
+	};
+}
 //#endregion
 //#region src/chat/react.tsx
 /**
@@ -736,7 +216,7 @@ function detectToolsRequiringConfirmation(tools) {
 	return Object.entries(tools).filter(([_name, tool]) => !tool.execute).map(([name]) => name);
 }
 function useAgentChat(options) {
-	const { agent, getInitialMessages, messages: optionsInitialMessages, onToolCall, onData, experimental_automaticToolResolution, tools, toolsRequiringConfirmation: manualToolsRequiringConfirmation, autoContinueAfterToolResult = true, autoSendAfterAllConfirmationsResolved = true, resume = true, cancelOnClientAbort = false, syncMessagesToServer = true, body: bodyOption, prepareSendMessagesRequest, ...rest } = options;
+	const { agent, getInitialMessages, messages: optionsInitialMessages, onToolCall, onData, experimental_automaticToolResolution, tools, toolsRequiringConfirmation: manualToolsRequiringConfirmation, autoContinueAfterToolResult = true, autoSendAfterAllConfirmationsResolved = true, resume = true, cancelOnClientAbort = false, syncMessagesToServer = true, body: bodyOption, prepareSendMessagesRequest, throttle, experimental_throttle, ...rest } = options;
 	if (manualToolsRequiringConfirmation) warnDeprecated("useAgentChat.toolsRequiringConfirmation", "The 'toolsRequiringConfirmation' option is deprecated. Use needsApproval on server-side tools instead. Will be removed in the next major version.");
 	if (experimental_automaticToolResolution) warnDeprecated("useAgentChat.experimental_automaticToolResolution", "The 'experimental_automaticToolResolution' option is deprecated. Use the onToolCall callback instead. Will be removed in the next major version.");
 	if (options.autoSendAfterAllConfirmationsResolved !== void 0) warnDeprecated("useAgentChat.autoSendAfterAllConfirmationsResolved", "The 'autoSendAfterAllConfirmationsResolved' option is deprecated. Use sendAutomaticallyWhen from AI SDK instead. Will be removed in the next major version.");
@@ -862,6 +342,10 @@ function useAgentChat(options) {
 	const customTransport = customTransportRef.current;
 	const useChatHelpers = useChat({
 		...rest,
+		...chatThrottleOptions({
+			experimental_throttle,
+			throttle
+		}),
 		onData,
 		messages: initialMessages,
 		transport: customTransport,
@@ -965,8 +449,6 @@ function useAgentChat(options) {
 	const isResolvingToolsRef = useRef(false);
 	const [toolResolutionTrigger, setToolResolutionTrigger] = useState(0);
 	const [clientToolResults, setClientToolResults] = useState(/* @__PURE__ */ new Map());
-	const messagesRef = useRef(chatMessages);
-	messagesRef.current = chatMessages;
 	const initialMessagesRef = useRef(initialMessages);
 	initialMessagesRef.current = initialMessages;
 	const seededInitialMessagesKeyRef = useRef(null);
@@ -986,7 +468,7 @@ function useAgentChat(options) {
 	]);
 	const localResponseMessageIdsRef = useRef(/* @__PURE__ */ new Map());
 	const protectedStreamingAssistantRef = useRef(null);
-	const preserveProtectedStreamingAssistant = useCallback((messages) => {
+	const preserveProtectedStreamingAssistant = useCallback((messages, currentMessages) => {
 		const protection = protectedStreamingAssistantRef.current;
 		if (!protection) return [...messages];
 		const protectedIndex = messages.findIndex((message) => message.id === protection.assistantId);
@@ -994,22 +476,20 @@ function useAgentChat(options) {
 			protectedStreamingAssistantRef.current = null;
 			return [...messages];
 		}
-		const protectedAssistant = messagesRef.current.find((message) => message.id === protection.assistantId) ?? messages.find((message) => message.id === protection.assistantId);
+		const protectedAssistant = currentMessages.find((message) => message.id === protection.assistantId) ?? messages.find((message) => message.id === protection.assistantId);
 		if (!protectedAssistant) return [...messages];
 		return [...messages.filter((message) => message.id !== protection.assistantId), protectedAssistant];
 	}, []);
 	const protectStreamingAssistantTail = useCallback(() => {
 		if (statusRef.current !== "streaming") return;
-		const assistantInfo = findLastAssistantMessage(messagesRef.current);
-		if (!assistantInfo) return;
-		if (protectedStreamingAssistantRef.current?.assistantId !== assistantInfo.message.id) protectedStreamingAssistantRef.current = {
-			assistantId: assistantInfo.message.id,
-			anchorMessageId: messagesRef.current[assistantInfo.index - 1]?.id ?? null
-		};
-		setMessages((prevMessages) => {
-			const protection = protectedStreamingAssistantRef.current;
-			if (!protection) return prevMessages;
-			return moveMessageToEnd(prevMessages, protection.assistantId);
+		setMessages((currentMessages) => {
+			const assistantInfo = findLastAssistantMessage(currentMessages);
+			if (!assistantInfo) return currentMessages;
+			if (protectedStreamingAssistantRef.current?.assistantId !== assistantInfo.message.id) protectedStreamingAssistantRef.current = {
+				assistantId: assistantInfo.message.id,
+				anchorMessageId: currentMessages[assistantInfo.index - 1]?.id ?? null
+			};
+			return moveMessageToEnd(currentMessages, assistantInfo.message.id);
 		});
 	}, [setMessages]);
 	const restoreProtectedStreamingAssistant = useCallback((assistantId) => {
@@ -1290,17 +770,18 @@ function useAgentChat(options) {
 				case "cf_agent_chat_recovering":
 					setIsRecovering(Boolean(data.recovering));
 					break;
-				case "cf_agent_chat_messages": {
-					let next = preserveProtectedStreamingAssistant(data.messages);
-					const observed = streamStateRef.current;
-					if (observed.status === "observing" && observed.accumulator.parts.length > 0) {
-						const snapshotIdx = next.findIndex((m) => m.id === observed.accumulator.messageId);
-						const snapshotParts = snapshotIdx >= 0 ? next[snapshotIdx].parts.length : 0;
-						if (observed.accumulator.parts.length >= snapshotParts) next = observed.accumulator.mergeInto(next);
-					}
-					setMessages(next);
+				case "cf_agent_chat_messages":
+					setMessages((currentMessages) => {
+						let next = preserveProtectedStreamingAssistant(data.messages, currentMessages);
+						const observed = streamStateRef.current;
+						if (observed.status === "observing" && observed.accumulator.parts.length > 0) {
+							const snapshotIdx = next.findIndex((m) => m.id === observed.accumulator.messageId);
+							const snapshotParts = snapshotIdx >= 0 ? next[snapshotIdx].parts.length : 0;
+							if (observed.accumulator.parts.length >= snapshotParts) next = observed.accumulator.mergeInto(next);
+						}
+						return next;
+					});
 					break;
-				}
 				case "cf_agent_message_updated":
 					setMessages((prevMessages) => {
 						const updatedMessage = data.message;
@@ -1368,26 +849,26 @@ function useAgentChat(options) {
 						if (data.body?.trim()) try {
 							const chunkData = JSON.parse(data.body);
 							if (chunkData.type === "start" && typeof chunkData.messageId === "string") {
-								localResponseIds.set(data.id, chunkData.messageId);
+								const messageId = chunkData.messageId;
+								localResponseIds.set(data.id, messageId);
 								if (!data.continuation) {
-									if (protectedStreamingAssistantRef.current?.assistantId !== chunkData.messageId) {
-										const msgs = messagesRef.current;
-										const idx = msgs.findIndex((m) => m.id === chunkData.messageId);
-										const anchorMessageId = idx >= 0 ? msgs[idx - 1]?.id ?? null : msgs[msgs.length - 1]?.id ?? null;
+									if (protectedStreamingAssistantRef.current?.assistantId !== messageId) setMessages((currentMessages) => {
+										const idx = currentMessages.findIndex((message) => message.id === messageId);
 										protectedStreamingAssistantRef.current = {
-											assistantId: chunkData.messageId,
-											anchorMessageId
+											assistantId: messageId,
+											anchorMessageId: idx >= 0 ? currentMessages[idx - 1]?.id ?? null : currentMessages[currentMessages.length - 1]?.id ?? null
 										};
-									}
+										return currentMessages;
+									});
 								}
 								if (data.replay && !data.continuation && !resumingToolContinuationRef.current && observedToolContinuationRequestIdRef.current !== data.id) {
 									pendingReplayResumeRequestIdsRef.current.delete(data.id);
-									resetMatchingHydratedAssistantForReplay(chunkData.messageId);
+									resetMatchingHydratedAssistantForReplay(messageId);
 								}
 							}
 						} catch {}
-						if (data.done || data.replayComplete) pendingReplayResumeRequestIdsRef.current.delete(data.id);
-						if (data.done) {
+						if (data.done || data.error || data.replayComplete) pendingReplayResumeRequestIdsRef.current.delete(data.id);
+						if (data.done || data.error) {
 							if (streamStateRef.current.status === "observing" && streamStateRef.current.streamId === data.id) {
 								streamStateRef.current = { status: "idle" };
 								setIsServerStreaming(false);
@@ -1403,6 +884,19 @@ function useAgentChat(options) {
 					}
 					let chunkData;
 					if (data.replay && streamStateRef.current.status !== "observing" && !pendingReplayResumeRequestIdsRef.current.has(data.id)) return;
+					if (data.error) {
+						pendingReplayResumeRequestIdsRef.current.delete(data.id);
+						customTransport.handleServerTurnCompleted(data.id);
+						fallbackAckedResumeRequestIdsRef.current.delete(data.id);
+						setIsRecovering(false);
+						if (streamStateRef.current.status === "idle" || streamStateRef.current.streamId === data.id) {
+							const result = transition(streamStateRef.current, { type: "clear" });
+							streamStateRef.current = result.state;
+							setIsServerStreaming(result.isStreaming);
+						}
+						if (observedToolContinuationRequestIdRef.current === data.id) resetToolContinuation();
+						break;
+					}
 					if (data.body?.trim()) try {
 						chunkData = JSON.parse(data.body);
 						if (data.replay && !data.continuation && !resumingToolContinuationRef.current && observedToolContinuationRequestIdRef.current !== data.id && typeof chunkData.messageId === "string" && chunkData.type === "start") {
@@ -1429,8 +923,7 @@ function useAgentChat(options) {
 						error: data.error,
 						replay: data.replay,
 						replayComplete: data.replayComplete,
-						continuation: data.continuation,
-						currentMessages: data.continuation ? messagesRef.current : void 0
+						continuation: data.continuation
 					});
 					streamStateRef.current = result.state;
 					if (result.messagesUpdate) setMessages(result.messagesUpdate);
@@ -1550,13 +1043,13 @@ function useAgentChat(options) {
 	const addToolApprovalResponseAndNotifyServer = (args) => {
 		const { id: approvalId, approved } = args;
 		let toolCallId;
-		for (const msg of messagesRef.current) {
-			for (const part of msg.parts) if ("toolCallId" in part && "approval" in part && part.approval?.id === approvalId) {
+		setMessages((currentMessages) => {
+			for (const msg of currentMessages) for (const part of msg.parts) if ("toolCallId" in part && "approval" in part && part.approval?.id === approvalId) {
 				toolCallId = part.toolCallId;
-				break;
+				return currentMessages;
 			}
-			if (toolCallId) break;
-		}
+			return currentMessages;
+		});
 		if (toolCallId) sendToolApprovalToServer(toolCallId, approved);
 		else console.warn(`[useAgentChat] addToolApprovalResponse: Could not find toolCallId for approval ID "${approvalId}". Server will not be notified, which may cause duplicate messages.`);
 		addToolApprovalResponse(args);
@@ -1655,15 +1148,15 @@ function useAgentChat(options) {
 			agent.send(JSON.stringify({ type: "cf_agent_chat_clear" }));
 		},
 		setMessages: (messagesOrUpdater) => {
-			let resolvedMessages;
-			if (typeof messagesOrUpdater === "function") resolvedMessages = messagesOrUpdater(messagesRef.current);
-			else resolvedMessages = messagesOrUpdater;
-			if (resolvedMessages.length === 0) markInitialMessagesSeeded();
-			setMessages(resolvedMessages);
-			if (syncMessagesToServer) agent.send(JSON.stringify({
-				messages: resolvedMessages,
-				type: "cf_agent_chat_messages"
-			}));
+			setMessages((currentMessages) => {
+				const resolvedMessages = typeof messagesOrUpdater === "function" ? messagesOrUpdater(currentMessages) : messagesOrUpdater;
+				if (resolvedMessages.length === 0) markInitialMessagesSeeded();
+				if (syncMessagesToServer) agent.send(JSON.stringify({
+					messages: resolvedMessages,
+					type: "cf_agent_chat_messages"
+				}));
+				return resolvedMessages;
+			});
 		}
 	};
 }