diff --git a/.changeset/chat-transport-memo-4187.md b/.changeset/chat-transport-memo-4187.md new file mode 100644 index 0000000000..93978f70c6 --- /dev/null +++ b/.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`. diff --git a/packages/plugin-chatbot/src/__tests__/useObjectChat.transportIdentity.test.tsx b/packages/plugin-chatbot/src/__tests__/useObjectChat.transportIdentity.test.tsx new file mode 100644 index 0000000000..5ac0c3b42f --- /dev/null +++ b/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); + }); +}); diff --git a/packages/plugin-chatbot/src/useObjectChat.ts b/packages/plugin-chatbot/src/useObjectChat.ts index 1262875482..1343a8d7d5 100644 --- a/packages/plugin-chatbot/src/useObjectChat.ts +++ b/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,9 +473,13 @@ 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 } : {}), @@ -463,8 +487,25 @@ export function useObjectChat(options: UseObjectChatOptions = {}): UseObjectChat }, // 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).model = modelRef.current; @@ -476,10 +517,16 @@ export function useObjectChat(options: UseObjectChatOptions = {}): UseObjectChat req.body = withHandoffContext(req.body as Record, 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