perf(plugin-chatbot): stop rebuilding the chat transport on every render (#4187)

.changeset/chat-transport-memo-4187.md

@@ -0,0 +1,27 @@
+---
+'@object-ui/plugin-chatbot': patch
+---
+
+`useObjectChat` no longer rebuilds its `DefaultChatTransport` on every render
+(objectui#4187).
+
+The transport `useMemo` listed the caller's `body` and `headers` in its dep list.
+Both are object props and every caller passes a fresh literal each render — the AI
+page's chat pane builds its `body.context` inline — so the memo never hit and a
+transport was constructed on every render of every chat surface, which during a
+streaming turn is once per token batch.
+
+`body` and `headers` are now read through refs inside
+`prepareSendMessagesRequest`, the idiom this hook already uses for the live model
+(`modelRef`) and the handoff conversation id (`parentConvRef`), and they are gone
+from the dep list. Unlike memoizing at each call site, a future caller cannot
+undo it.
+
+No user-visible behaviour changes: `@ai-sdk/react` keeps the transport in a ref
+and re-keys its `Chat` only on `chat`/`id` (verified against the installed
+4.0.68), which `useObjectChat` passes neither of, so the message thread was never
+at risk — the rebuild was pure waste. The one real difference is *when* the two
+values are sampled: a send now reads them at send time, so it observes the values
+of the most recent render instead of those of the last render that happened to
+rebuild the transport. That is never staler than before, and it is pinned by
+`useObjectChat.transportIdentity.test.tsx`.

packages/plugin-chatbot/src/__tests__/useObjectChat.transportIdentity.test.tsx

@@ -0,0 +1,198 @@
+/**
+ * ObjectUI
+ * Copyright (c) 2024-present ObjectStack Inc.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ *
+ * objectui#4187 — the transport `useMemo` in `useObjectChat` used to list the
+ * caller's `body`/`headers` as deps. Every caller passes a fresh object literal
+ * each render, so the memo never hit and a `DefaultChatTransport` was built on
+ * every render of every chat surface — once per token batch while streaming.
+ *
+ * These tests pin BOTH halves of the fix:
+ * 1. the transport is built once, however often the caller re-renders;
+ * 2. what a send actually observes afterwards — `body`/`headers` are now read
+ * through refs at SEND time, so the values are those of the most recent
+ * render rather than of the last render that happened to rebuild the
+ * transport. That timing shift is the one real behavioural difference
+ * between this fix and call-site memoization, and it is not visible by
+ * reading the diff.
+ */
+import { renderHook, act, waitFor } from '@testing-library/react';
+import { vi, describe, it, expect, beforeEach, afterEach } from 'vitest';
+import { useObjectChat } from '../useObjectChat';
+
+/** Counts `new DefaultChatTransport(...)` across a test. */
+const transportSpy = vi.hoisted(() => ({ constructions: 0 }));
+
+// Count constructions without changing behaviour: a construct-trap Proxy around
+// the real class, so the hook still talks to the genuine transport.
+vi.mock('ai', async (importOriginal) => {
+ const actual = (await importOriginal()) as { [key: string]: unknown };
+ const Real = actual.DefaultChatTransport as new (...args: never[]) => object;
+ const Counting = new Proxy(Real, {
+ construct(target, args) {
+ transportSpy.constructions += 1;
+ return Reflect.construct(target, args);
+ },
+ });
+ return { ...actual, DefaultChatTransport: Counting };
+});
+
+const API = 'https://example.test/api/v1/ai/agents/build/chat';
+
+type AnyRecord = { [key: string]: unknown };
+type FetchMock = { mock: { calls: unknown[][] } };
+
+/** A minimal, well-formed Vercel AI UI-message data stream so a send completes. */
+function dataStreamResponse(): Response {
+ const body =
+ 'data: {"type":"start"}\n\n' + 'data: {"type":"finish"}\n\n' + 'data: [DONE]\n\n';
+ return new Response(body, {
+ status: 200,
+ headers: {
+ 'Content-Type': 'text/event-stream',
+ 'x-vercel-ai-ui-message-stream': 'v1',
+ },
+ });
+}
+
+/** The JSON body of a captured chat POST. */
+function bodyOf(fetchMock: FetchMock, callIndex: number): AnyRecord {
+ const init = fetchMock.mock.calls[callIndex]?.[1] as { body?: string } | undefined;
+ return JSON.parse(init?.body ?? '{}') as AnyRecord;
+}
+
+/** One request header off a captured chat POST, however the SDK shaped it. */
+function headerOf(fetchMock: FetchMock, callIndex: number, name: string): string | null {
+ const init = fetchMock.mock.calls[callIndex]?.[1] as { headers?: HeadersInit } | undefined;
+ return new Headers(init?.headers ?? {}).get(name);
+}
+
+beforeEach(() => {
+ transportSpy.constructions = 0;
+});
+
+afterEach(() => {
+ vi.unstubAllGlobals();
+ vi.restoreAllMocks();
+});
+
+describe('useObjectChat — transport identity (#4187)', () => {
+ it('builds ONE transport across renders that pass fresh body/headers literals', () => {
+ const { rerender } = renderHook(
+ ({ tick }: { tick: number }) =>
+ useObjectChat({
+ api: API,
+ conversationId: 'build_1',
+ // Fresh identity on every render — the shape every caller uses today
+ // (the AI page's chat pane builds its `body.context` inline).
+ body: { context: { activeApp: 'AI', tick } },
+ headers: { 'x-oui-tick': String(tick) },
+ }),
+ { initialProps: { tick: 0 } },
+ );
+
+ expect(transportSpy.constructions).toBe(1);
+
+ rerender({ tick: 1 });
+ rerender({ tick: 2 });
+ rerender({ tick: 3 });
+
+ // Without the fix this is 4 — one construction per render.
+ expect(transportSpy.constructions).toBe(1);
+ });
+
+ it('still rebuilds the transport when a memoized dep really changes', () => {
+ const { rerender } = renderHook(
+ ({ conversationId }: { conversationId: string }) =>
+ useObjectChat({ api: API, conversationId, body: { context: {} } }),
+ { initialProps: { conversationId: 'build_1' } },
+ );
+
+ expect(transportSpy.constructions).toBe(1);
+ rerender({ conversationId: 'build_2' });
+ expect(transportSpy.constructions).toBe(2);
+ });
+
+ it('a send carries the body/headers of the MOST RECENT render', async () => {
+ const fetchMock = vi.fn(async () => dataStreamResponse());
+ vi.stubGlobal('fetch', fetchMock);
+
+ const { result, rerender } = renderHook(
+ ({ tick }: { tick: number }) =>
+ useObjectChat({
+ api: API,
+ conversationId: 'build_1',
+ body: { context: { tick } },
+ headers: { 'x-oui-tick': String(tick) },
+ }),
+ { initialProps: { tick: 1 } },
+ );
+
+ rerender({ tick: 2 });
+ rerender({ tick: 3 });
+
+ await act(async () => {
+ result.current.sendMessage('hello');
+ });
+ await waitFor(() => expect(fetchMock).toHaveBeenCalledTimes(1));
+
+ expect((bodyOf(fetchMock as unknown as FetchMock, 0).context as AnyRecord).tick).toBe(3);
+ expect(headerOf(fetchMock as unknown as FetchMock, 0, 'x-oui-tick')).toBe('3');
+ });
+
+ it('reads `body` at SEND time, not at transport-construction time', async () => {
+ const fetchMock = vi.fn(async () => dataStreamResponse());
+ vi.stubGlobal('fetch', fetchMock);
+
+ // One STABLE object identity for the whole test. With a stable identity the
+ // memo never re-ran even before this change, so the value a send observes is
+ // decided purely by WHEN it is read: the old code spread `body` into the
+ // transport at CONSTRUCTION time and froze `tick: 1`; the ref is spread at
+ // SEND time. (Mutating a prop is not endorsed here — it is simply the only
+ // externally observable probe of the read timing.)
+ const stableBody: { tick: number } = { tick: 1 };
+
+ const { result } = renderHook(() =>
+ useObjectChat({ api: API, conversationId: 'build_1', body: stableBody }),
+ );
+
+ stableBody.tick = 2;
+
+ await act(async () => {
+ result.current.sendMessage('hello');
+ });
+ await waitFor(() => expect(fetchMock).toHaveBeenCalledTimes(1));
+
+ expect(bodyOf(fetchMock as unknown as FetchMock, 0).tick).toBe(2);
+ });
+
+ it('keeps the hook-owned keys winning over the caller body', async () => {
+ const fetchMock = vi.fn(async () => dataStreamResponse());
+ vi.stubGlobal('fetch', fetchMock);
+
+ const { result } = renderHook(() =>
+ useObjectChat({
+ api: API,
+ conversationId: 'build_1',
+ model: 'claude-sonnet',
+ systemPrompt: 'be brief',
+ // A caller body that tries to shadow every hook-owned key.
+ body: { conversationId: 'spoofed', model: 'spoofed', systemPrompt: 'spoofed', stream: false },
+ }),
+ );
+
+ await act(async () => {
+ result.current.sendMessage('hello');
+ });
+ await waitFor(() => expect(fetchMock).toHaveBeenCalledTimes(1));
+
+ const sent = bodyOf(fetchMock as unknown as FetchMock, 0);
+ expect(sent.conversationId).toBe('build_1');
+ expect(sent.model).toBe('claude-sonnet');
+ expect(sent.systemPrompt).toBe('be brief');
+ expect(sent.stream).toBe(true);
+ });
+});

packages/plugin-chatbot/src/useObjectChat.ts

@@ -441,6 +441,26 @@ export function useObjectChat(options: UseObjectChatOptions = {}): UseObjectChat
  if (parentConversationId && !prev) parentConvRef.current = parentConversationId;
  }, [parentConversationId]);
 
+ // objectui#4187 - the caller's `body`/`headers` are read through refs at SEND
+ // time instead of being closed over by the transport, so they are NOT memo
+ // deps below. Every caller passes a fresh object literal each render (the AI
+ // page's chat pane rebuilds `body.context` inline), so listing them rebuilt
+ // `DefaultChatTransport` on every render of every chat surface - once per
+ // token batch during a streaming turn. Same idiom as `modelRef` above, and
+ // the only one of the two candidate fixes a future caller cannot silently
+ // undo by forgetting its own `useMemo`.
+ //
+ // SAMPLING CONTRACT (the one real behavioural change): a send serializes
+ // whatever these refs hold when `prepareSendMessagesRequest` runs, i.e. the
+ // values from the most recent render, spread at SEND time. Before, the values
+ // were spread into the transport at CONSTRUCTION time, so a send observed the
+ // last render that happened to rebuild it. The ref read is never staler than
+ // that and is now unconditional - see `useObjectChat.transportIdentity.test`.
+ const bodyRef = useRef(body);
+ bodyRef.current = body;
+ const headersRef = useRef(headers);
+ headersRef.current = headers;
+
  // Build a transport for API mode that posts to the configured endpoint and
  // forwards conversation/system/model metadata in the request body.
  // Note: conversationId is sent in the body (not a header) to avoid CORS
@@ -453,18 +473,39 @@ export function useObjectChat(options: UseObjectChatOptions = {}): UseObjectChat
  // `notSent` so the composer can restore the input and show a clear error
  // instead of silently dropping the message (see sendAwareFetch).
  fetch: sendAwareFetch,
- headers: { ...headers },
+ // No `headers` here (objectui#4187): the caller's headers are applied
+ // per-send in prepareSendMessagesRequest below. `reconnectToStream` does
+ // NOT run that hook — it reads this constructor's `headers` — so the first
+ // consumer to wire up stream resumption must merge `headersRef.current`
+ // into `prepareReconnectToStreamRequest` too, or it will resume without
+ // them. Nothing in this repo resumes a stream today.
  body: {
- ...body,
  ...(conversationId ? { conversationId } : {}),
  ...(model ? { model } : {}),
  ...(systemPrompt ? { systemPrompt } : {}),
  ...(streamingEnabled !== undefined ? { stream: streamingEnabled } : {}),
  },
  // Stamp a stable per-turn idempotency key (ADR-0013 D1). See withTurnId —
  // it reconstructs the full default body (incl. messages) + adds turnId.
- prepareSendMessagesRequest: ({ id, body: reqBody, messages, trigger, messageId }) => {
- const req = withTurnId({ id, body: reqBody, messages, trigger, messageId });
+ prepareSendMessagesRequest: ({
+ id,
+ body: reqBody,
+ messages,
+ trigger,
+ messageId,
+ headers: reqHeaders,
+ }) => {
+ // #4187: the caller's live `body` goes in FIRST, so the fixed keys above
+ // (conversationId/model/systemPrompt/stream) and any per-send body still
+ // win - the same precedence as when it was spread into the transport's
+ // own `body` option.
+ const req = withTurnId({
+ id,
+ body: { ...bodyRef.current, ...reqBody },
+ messages,
+ trigger,
+ messageId,
+ });
  // ADR-0028: always send the CURRENTLY selected model (see modelRef above)
  // so a mid-session picker switch routes, despite the cached transport.
  if (modelRef.current) (req.body as Record<string, unknown>).model = modelRef.current;
@@ -476,10 +517,16 @@ export function useObjectChat(options: UseObjectChatOptions = {}): UseObjectChat
  req.body = withHandoffContext(req.body as Record<string, unknown>, parentConvRef.current);
  parentConvRef.current = undefined;
  }
- return req;
+ // #4187: the SDK hands us its merged base headers and REPLACES them with
+ // whatever we return here, so re-merge instead of replacing. The caller's
+ // live headers go in first so a per-send header still overrides them,
+ // exactly as the transport's own `headers` option behaved.
+ const sendHeaders = new Headers(headersRef.current ?? {});
+ new Headers(reqHeaders ?? {}).forEach((value, key) => sendHeaders.set(key, value));
+ return { ...req, headers: sendHeaders };
  },
  });
- }, [isApiMode, api, headers, body, model, systemPrompt, streamingEnabled, conversationId]);
+ }, [isApiMode, api, model, systemPrompt, streamingEnabled, conversationId]);
 
  // --- @ai-sdk/react useChat (always called to satisfy Rules of Hooks, but only active in API mode) ---
  // Ref so `onError` (fired later, async) can reach the live setMessages/messages