UNPKG

agents

Version:

A home for your AI agents

825 lines (824 loc) • 27.3 kB
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/replay-dedupe.ts /** * Requests the ledger remembers. A request that finishes while the client is * disconnected never delivers the terminal frame that would forget it. */ const MAX_TRACKED_REQUESTS = 32; /** Highest chunk `seq` this client applied, per request id. */ var AppliedChunkLedger = class { constructor() { this.applied = /* @__PURE__ */ new Map(); } record(requestId, seq) { if (seq === void 0) return; const applied = Math.max(seq, this.applied.get(requestId) ?? -1); this.applied.delete(requestId); this.applied.set(requestId, applied); if (this.applied.size > MAX_TRACKED_REQUESTS) { const oldest = this.applied.keys().next().value; if (oldest !== void 0) this.applied.delete(oldest); } } /** Whether `frame` replays a continuation chunk this client already applied. */ isAppliedReplay(requestId, frame) { return frame.replay === true && frame.continuation === true && frame.seq !== void 0 && frame.seq <= (this.applied.get(requestId) ?? -1); } forget(requestId) { this.applied.delete(requestId); } clear() { this.applied.clear(); } }; /** * Filters one resumed stream's frames for a consumer that needs every part * opened before its deltas arrive (the AI SDK UI message stream). Applied * replay chunks are dropped; a part that was still open when they end is * re-opened with its original start chunk, so the new deltas have a part to * land in. The re-opened part renders as a second, adjacent part until the * server's final message replaces it. */ var ContinuationReplayFilter = class { constructor(ledger, requestId) { this.ledger = ledger; this.requestId = requestId; this.openParts = /* @__PURE__ */ new Map(); } /** The frames to apply in place of `frame`. */ frames(frame) { if (this.ledger.isAppliedReplay(this.requestId, frame)) { this.trackOpenPart(frame.body); return frame.body?.trim() ? [{ ...frame, body: "" }] : [frame]; } this.ledger.record(this.requestId, frame.seq); if (this.openParts.size === 0 || !frame.body?.trim()) return [frame]; const reopened = [...this.openParts.values()].map((body) => ({ ...frame, body, seq: void 0 })); this.openParts.clear(); return [...reopened, frame]; } trackOpenPart(body) { let chunk; try { chunk = JSON.parse(body ?? ""); } catch { return; } switch (chunk.type) { case "text-start": case "reasoning-start": this.openParts.set(`${chunk.type}:${chunk.id}`, body ?? ""); break; case "text-end": this.openParts.delete(`text-start:${chunk.id}`); break; case "reasoning-end": this.openParts.delete(`reasoning-start:${chunk.id}`); break; case "tool-input-start": this.openParts.set(`tool:${chunk.toolCallId}`, body ?? ""); break; case "tool-input-available": case "tool-input-error": this.openParts.delete(`tool:${chunk.toolCallId}`); break; } } }; //#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; const SOCKET_CLOSED_MID_STREAM = "WebSocket closed mid-stream"; /** * Ends a stream whose socket closed before the terminal `done` frame, so the * AI SDK reports an interrupted turn instead of a completed one (#2013). * Uses an error chunk because `controller.error()` would discard chunks the * consumer has not read yet. */ function interruptChatStream(controller, batch) { batch?.flush(); controller.enqueue({ type: "error", errorText: SOCKET_CLOSED_MID_STREAM }); controller.close(); } /** * 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.appliedChunks = new AppliedChunkLedger(); 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; this.onRequestBuffered = options.onRequestBuffered; } /** * 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; } /** * The server turn this transport last attached to that has not reported a * terminal frame. A socket close ends the local stream but keeps this set, * since the server keeps running the turn. */ get activeServerTurnId() { return this._activeServerTurnId; } 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(() => interruptChatStream(controller), 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; if (agent.send(JSON.stringify({ id: requestId, init: { method: "POST", body: bodyPayload }, type: "cf_agent_use_chat_request" })) === false && options.trigger === "submit-message") { const lastUserMessage = [...options.messages].reverse().find((message) => message.role === "user"); this.onRequestBuffered?.(lastUserMessage?.id); } 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 replayFilter = null; 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) { transport.appliedChunks.forget(requestId); finish(() => failChatStream(batch, data.body || "Stream error")); return; } replayFilter ??= new ContinuationReplayFilter(transport.appliedChunks, requestId); for (const frame of replayFilter.frames(data)) applyChatResponseFrame(batch, frame); if (data.done) { transport.appliedChunks.forget(requestId); finish(() => controller.close()); } } catch {} }; const onClose = () => finish(() => requestId === null ? controller.close() : interruptChatStream(controller, batch)); 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 replayFilter = new ContinuationReplayFilter(transport.appliedChunks, requestId); 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) { transport.appliedChunks.forget(requestId); finish(() => failChatStream(batch, data.body || "Stream error")); return; } for (const frame of replayFilter.frames(data)) applyChatResponseFrame(batch, frame); if (data.done) { transport.appliedChunks.forget(requestId); finish(() => controller.close()); } } catch {} }; const onClose = () => { finish(() => interruptChatStream(controller, batch), 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-cqMKq3IY.js.map