Skip to main content

WebSocketChatTransport

还没翻译译自英文(ai-sdk.dev)本页 36 段都还没译

A ChatTransport implementation that sends chat requests and receives UIMessageChunk streams over one persistent WebSocket connection.

import { useChat } from '@ai-sdk/react';
import { WebSocketChatTransport } from 'ai';
import { useMemo } from 'react';

export default function Chat() {
  const transport = useMemo(
    () => new WebSocketChatTransport({ url: '/api/chat' }),
    [],
  );

  const { messages, sendMessage, stop } = useChat({ transport });

  // ... render the chat
}

Import

import { WebSocketChatTransport } from "ai"

Constructor

  • url (string): The WebSocket endpoint. Relative URLs are resolved against the browser location. HTTP and HTTPS URLs are converted to WS and WSS.
  • protocols? (string | string[]): WebSocket subprotocols sent during the handshake.
  • params? (Resolvable<Record<string, string>>): Query parameters resolved and appended whenever a new connection is opened.
  • headers? (Resolvable<Record<string, string> | Headers>): Headers included inside each application-level send or resume frame. These are not WebSocket handshake headers.
  • body? (Resolvable<object>): Additional application data merged with per-request body data and included in send and resume frames.
  • webSocket? (WebSocketConstructor): A custom WebSocket constructor. Use this for Node.js WebSocket implementations or tests. Defaults to globalThis.WebSocket.
  • prepareSendMessagesRequest? (PrepareWebSocketChatTransportSendMessagesRequest): Transforms the messages, headers, or body before a send frame is transmitted.
  • prepareReconnectToStreamRequest? (PrepareWebSocketChatTransportReconnectToStreamRequest): Transforms the headers or body before a resume frame is transmitted.

headers supplied to sendMessage, regenerate, or an automatic tool follow-up override transport-level headers with the same lower-cased name. Per-request body properties override transport-level body properties. Preparation callbacks receive the merged values; values they return replace the corresponding merged value.

Methods

sendMessages()

Opens the socket when necessary, sends a send frame, and returns the correlated ReadableStream<UIMessageChunk>. Multiple requests can share the same connection and are routed by requestId.

reconnectToStream()

Sends a resume frame. The promise resolves as follows:

  • start, the first chunk, or end resolves to a readable stream.
  • no-active resolves to null.
  • error, cancellation, or a connection failure rejects the promise if no stream has been returned yet.

Servers should send start as soon as they confirm that a resumable stream is active, then replay the retained response from sequence zero. The transport omits lastSequence because Chat starts a fresh message parser when resuming: text, reasoning, and tool deltas need their preceding chunks. Retain and replay the original UI start chunk with a stable messageId so Chat replaces the partial assistant message instead of appending another one. The transport rejects gaps and ignores duplicate sequence numbers within each response stream, including the replay.

close()

Closes the current socket and errors active streams. A later send or resume request opens a new connection.

The code that creates a transport owns its lifecycle. React and Vue useChat do not close caller-supplied transports during replacement or cleanup, because one transport can be shared by multiple chats. Call close() after its final consumer is gone.

If a directly managed Chat instance exclusively owns the transport, call dispose():

await chat.dispose();

For a shared transport, stop or discard individual chats without disposing their shared transport, then call transport.close() after all consumers are finished.

Wire Protocol

Each frame is a JSON object with a type and a correlated requestId.

Client-to-server frames

TypePurpose
sendStarts a new or regenerated turn and includes id, messages, trigger, messageId, request headers, body, and metadata.
resumeRequests the retained response for a chat id, with request headers, body, and metadata. lastSequence is omitted to request a complete replay.
abortCancels the operation identified by requestId.

Use WebSocketChatTransportRequest<UI_MESSAGE>, WebSocketChatTransportSendRequest, WebSocketChatTransportResumeRequest, and WebSocketChatTransportAbortRequest to type these frames.

Server-to-client frames

TypePurpose
startConfirms that a resume request has an active stream.
chunkCarries a monotonic sequence and one valid UIMessageChunk in chunk.
endCloses the correlated response stream.
errorErrors the correlated stream with the optional errorText.
no-activeResolves a resume request to null; for a send request it closes the stream.

Use WebSocketChatTransportResponse to type server frames.

Responses for unknown or completed request IDs are ignored. Malformed frames, unknown response types, and invalid UI message chunks are treated as protocol errors: the connection is closed and all active streams are errored.

Validate every untrusted client frame before using its fields:

const validated = await safeValidateWebSocketChatTransportRequest({
  value: parsed.value,
});

if (!validated.success) {
  socket.close(1008, 'Invalid frame');
  return;
}

const frame = validated.data;

Server Example

The server owns authentication, authorization, message validation, model execution, resumable stream storage, and cancellation. The following Node.js outline forwards a streamText UI message stream:

import { WebSocket, WebSocketServer } from 'ws';
import {
  convertToModelMessages,
  safeValidateWebSocketChatTransportRequest,
  streamText,
  toUIMessageStream,
  validateUIMessages,
  type UIMessage,
  type UIMessageChunk,
  type WebSocketChatTransportResponse,
} from 'ai';
import { generateId, safeParseJSON } from '@ai-sdk/provider-utils';

const server = new WebSocketServer({ port: 8080 });

type RetainedStream = {
  abortController: AbortController;
  chunks: Array<{ sequence: number; chunk: UIMessageChunk }>;
  subscribers: Map<WebSocket, string>;
  ended: boolean;
};

// Scope retained streams to an authenticated user. Use durable storage when
// reconnects can land in another process.
const streamsByUser = new Map<string, Map<string, RetainedStream>>();

server.on('connection', (socket, request) => {
  const userId = authenticateAndValidateOrigin(request);
  const streams =
    streamsByUser.get(userId) ?? new Map<string, RetainedStream>();
  streamsByUser.set(userId, streams);

  socket.on('message', async data => {
    const parsed = await safeParseJSON({ text: data.toString() });
    if (!parsed.success) {
      socket.close(1008, 'Invalid frame');
      return;
    }

    const validated =
      await safeValidateWebSocketChatTransportRequest<UIMessage>({
        value: parsed.value,
      });
    if (!validated.success) {
      socket.close(1008, 'Invalid frame');
      return;
    }
    const frame = validated.data;

    const send = (response: WebSocketChatTransportResponse) => {
      socket.send(JSON.stringify(response));
    };

    if (frame.type === 'abort') {
      for (const stream of streams.values()) {
        if (stream.subscribers.get(socket) === frame.requestId) {
          stream.subscribers.delete(socket);
          stream.abortController.abort();
        }
      }
      return;
    }

    if (frame.type === 'resume') {
      const stream = streams.get(frame.id);
      if (stream == null) {
        send({ type: 'no-active', requestId: frame.requestId });
        return;
      }

      send({ type: 'start', requestId: frame.requestId });
      if (!stream.ended) {
        stream.subscribers.set(socket, frame.requestId);
      }
      for (const stored of stream.chunks) {
        if (stored.sequence > (frame.lastSequence ?? -1)) {
          send({
            type: 'chunk',
            requestId: frame.requestId,
            ...stored,
          });
        }
      }
      if (stream.ended) {
        send({ type: 'end', requestId: frame.requestId });
      }
      return;
    }

    if (streams.get(frame.id)?.ended === false) {
      send({
        type: 'error',
        requestId: frame.requestId,
        errorText: 'A response is already active for this chat.',
      });
      return;
    }

    const abortController = new AbortController();
    const stream: RetainedStream = {
      abortController,
      chunks: [],
      subscribers: new Map([[socket, frame.requestId]]),
      ended: false,
    };
    streams.set(frame.id, stream);

    try {
      const messages = await validateUIMessages({
        messages: frame.messages,
      });
      const result = streamText({
        model: 'anthropic/claude-sonnet-5.5',
        messages: await convertToModelMessages(messages),
        abortSignal: abortController.signal,
      });

      for await (const chunk of toUIMessageStream({
        stream: result.stream,
        originalMessages: messages,
        generateMessageId: generateId,
      })) {
        const stored = { sequence: stream.chunks.length, chunk };
        stream.chunks.push(stored);
        for (const [subscriber, requestId] of stream.subscribers) {
          subscriber.send(
            JSON.stringify({ type: 'chunk', requestId, ...stored }),
          );
        }
      }
      stream.ended = true;
      for (const [subscriber, requestId] of stream.subscribers) {
        subscriber.send(JSON.stringify({ type: 'end', requestId }));
      }
    } catch (error) {
      // Failed responses cannot resume. Remove the retained entry so a retry
      // can start a new response and resume requests receive no-active.
      streams.delete(frame.id);
      // Log the internal error server-side, but do not expose it to the client.
      console.error(error);
      for (const [subscriber, requestId] of stream.subscribers) {
        subscriber.send(
          JSON.stringify({
            type: 'error',
            requestId,
            errorText: 'Chat request failed.',
          }),
        );
      }
    } finally {
      stream.subscribers.clear();
    }
  });

  socket.on('close', () => {
    for (const stream of streams.values()) {
      stream.subscribers.delete(socket);
    }
  });
});

This outline keeps completed chunks in memory to make the replay behavior clear. Production servers should expire completed streams and store resumable state durably when reconnects can reach another process. If response setup or streaming throws, the outline removes that response from retained state before reporting the error. A later resume returns no-active, and a new send can retry the chat.

Authentication and Deployment

  • Use WSS in production.
  • Browser WebSockets cannot set arbitrary handshake headers. Prefer secure cookies, short-lived query tokens, or subprotocols.
  • Validate the handshake origin, authorize chat IDs, validate every frame and UI message, limit frame sizes, and rate-limit connections.
  • Do not depend on process memory for reconnection when your deployment can restart or move the connection to another process.
  • The transport does not automatically replay interrupted send requests, because replay can duplicate model output or tool side effects.

WebSocketChatTransport is separate from the experimental Realtime APIs. Realtime transports implement provider realtime sessions; this transport carries the provider-independent AI SDK UI message protocol used by useChat.

Client-side Tool Follow-ups

When sendAutomaticallyWhen triggers after client-side tool output is added, the next turn reuses the same WebSocket. This avoids another HTTP connection, but the current ChatTransport contract remains turn-based: the tool output is not injected into a model generation that is still active.