From d729cef8b76026f6650fd0089ce883ccdc6bdfd8 Mon Sep 17 00:00:00 2001 From: Ross <144740362+ross0x01@users.noreply.github.com> Date: Wed, 2 Sep 2026 22:54:22 -0500 Subject: [PATCH 1/2] Cut ask-route preflight latency and batch client data-stream appends Port the transferable perf patterns from t3code: - /api/chat ran ~10 serial awaits before the model call with no Promise.all. Overlap user customization with the message fetch, project context with the extra-usage config, and start the notes fetch early so it no longer adds a round-trip right before streamText. - Add preflight.duration_ms and stream.first_chunk_ms to the chat wide event so a slow first token can be attributed to preflight vs provider. - Batch dataStream appends into one state update per 100ms window. Every data part (including each terminal output chunk) used to copy the whole array and re-run every consumer effect. Chat resets drop pending parts. Co-Authored-By: Claude Fable 5.1 --- app/components/chat.tsx | 13 +- .../useBatchedDataStreamAppend.test.ts | 126 ++++++++++++++++++ app/hooks/useBatchedDataStreamAppend.ts | 62 +++++++++ lib/__tests__/wide-event-timing.test.ts | 48 +++++++ .../chat-stream-helpers-notes.test.ts | 50 +++++++ lib/api/chat-handler.ts | 83 +++++++----- lib/api/chat-logger.ts | 7 + lib/api/chat-stream-helpers.ts | 15 ++- lib/logger.ts | 30 +++++ 9 files changed, 391 insertions(+), 43 deletions(-) create mode 100644 app/hooks/__tests__/useBatchedDataStreamAppend.test.ts create mode 100644 app/hooks/useBatchedDataStreamAppend.ts create mode 100644 lib/__tests__/wide-event-timing.test.ts diff --git a/app/components/chat.tsx b/app/components/chat.tsx index 49116467d..3a792e034 100644 --- a/app/components/chat.tsx +++ b/app/components/chat.tsx @@ -102,6 +102,7 @@ import { useAutoContinue } from "../hooks/useAutoContinue"; import { findActiveTimelineAnchorMessageId } from "./message-timeline-rows"; import { useLatestRef } from "../hooks/useLatestRef"; import { useDataStreamDispatch } from "./DataStreamProvider"; +import { useBatchedDataStreamAppend } from "@/app/hooks/useBatchedDataStreamAppend"; import { markSidebarTaskVisited, removeDraft, @@ -522,6 +523,8 @@ export const Chat = ({ autoResume }: { autoResume: boolean }) => { const computerDialogRef = useRef(null); const computerDialogPreviousFocusRef = useRef(null); const { setDataStream, setIsAutoResuming } = useDataStreamDispatch(); + const { appendDataPart, clearDataStream } = + useBatchedDataStreamAppend(setDataStream); const { isLoading: isConvexAuthLoading, isAuthenticated: isConvexAuthenticated, @@ -1005,7 +1008,7 @@ export const Chat = ({ autoResume }: { autoResume: boolean }) => { return; } agentLongHasVisibleProgressRef.current = true; - setDataStream((ds) => [...ds, { ...dataPart, __chatId: chatId }]); + appendDataPart({ ...dataPart, __chatId: chatId }); switch (dataPart.type) { case "data-agent-approval-session": { const approvalData = dataPart.data as { @@ -1262,10 +1265,10 @@ export const Chat = ({ autoResume }: { autoResume: boolean }) => { ) { stopRef.current(); } - setDataStream([]); + clearDataStream(); setIsAutoResuming(false); }, - [setDataStream, setIsAutoResuming], + [clearDataStream, setIsAutoResuming], ); useEffect(() => { @@ -1388,10 +1391,10 @@ export const Chat = ({ autoResume }: { autoResume: boolean }) => { agentLongRunFallbackAllowedRef.current = true; setAgentLongRunId(null); agentLongHasVisibleProgressRef.current = false; - setDataStream([]); + clearDataStream(); setIsAutoResuming(false); dispatchStreaming({ type: "RESET_ON_CHAT_CHANGE" }); - }, [chatId, setDataStream, setIsAutoResuming]); + }, [chatId, clearDataStream, setIsAutoResuming]); useEffect(() => { return () => { diff --git a/app/hooks/__tests__/useBatchedDataStreamAppend.test.ts b/app/hooks/__tests__/useBatchedDataStreamAppend.test.ts new file mode 100644 index 000000000..16a10c477 --- /dev/null +++ b/app/hooks/__tests__/useBatchedDataStreamAppend.test.ts @@ -0,0 +1,126 @@ +import { describe, it, expect, beforeEach, afterEach } from "@jest/globals"; +import { renderHook, act } from "@testing-library/react"; +import type { ScopedDataUIPart } from "@/app/components/DataStreamProvider"; +import { + DATA_STREAM_BATCH_WINDOW_MS, + useBatchedDataStreamAppend, +} from "../useBatchedDataStreamAppend"; + +type Updater = + ScopedDataUIPart[] | ((current: ScopedDataUIPart[]) => ScopedDataUIPart[]); + +function createStore() { + let state: ScopedDataUIPart[] = []; + const setDataStream = jest.fn((updater: Updater) => { + state = typeof updater === "function" ? updater(state) : updater; + }); + return { + setDataStream, + get state() { + return state; + }, + }; +} + +const part = (id: string): ScopedDataUIPart => + ({ type: "data-terminal", id, data: { terminal: id } }) as ScopedDataUIPart; + +describe("useBatchedDataStreamAppend", () => { + beforeEach(() => { + jest.useFakeTimers(); + }); + + afterEach(() => { + jest.useRealTimers(); + }); + + it("appends every part from one window in a single state update, in order", () => { + const store = createStore(); + const { result } = renderHook(() => + useBatchedDataStreamAppend(store.setDataStream), + ); + + act(() => { + result.current.appendDataPart(part("a")); + result.current.appendDataPart(part("b")); + result.current.appendDataPart(part("c")); + }); + expect(store.setDataStream).not.toHaveBeenCalled(); + + act(() => { + jest.advanceTimersByTime(DATA_STREAM_BATCH_WINDOW_MS); + }); + expect(store.setDataStream).toHaveBeenCalledTimes(1); + expect(store.state.map((p) => p.id)).toEqual(["a", "b", "c"]); + }); + + it("opens a new window after a flush so later parts still arrive", () => { + const store = createStore(); + const { result } = renderHook(() => + useBatchedDataStreamAppend(store.setDataStream), + ); + + act(() => { + result.current.appendDataPart(part("a")); + jest.advanceTimersByTime(DATA_STREAM_BATCH_WINDOW_MS); + }); + act(() => { + result.current.appendDataPart(part("b")); + jest.advanceTimersByTime(DATA_STREAM_BATCH_WINDOW_MS); + }); + + expect(store.setDataStream).toHaveBeenCalledTimes(2); + expect(store.state.map((p) => p.id)).toEqual(["a", "b"]); + }); + + it("clear drops pending parts so a reset chat cannot resurrect them", () => { + const store = createStore(); + const { result } = renderHook(() => + useBatchedDataStreamAppend(store.setDataStream), + ); + + act(() => { + result.current.appendDataPart(part("stale")); + result.current.clearDataStream(); + jest.advanceTimersByTime(DATA_STREAM_BATCH_WINDOW_MS); + }); + + expect(store.setDataStream).toHaveBeenCalledTimes(1); + expect(store.state).toEqual([]); + }); + + it("flush applies pending parts immediately", () => { + const store = createStore(); + const { result } = renderHook(() => + useBatchedDataStreamAppend(store.setDataStream), + ); + + act(() => { + result.current.appendDataPart(part("a")); + result.current.flushDataStream(); + }); + + expect(store.state.map((p) => p.id)).toEqual(["a"]); + act(() => { + jest.advanceTimersByTime(DATA_STREAM_BATCH_WINDOW_MS); + }); + expect(store.setDataStream).toHaveBeenCalledTimes(1); + }); + + it("does not update state after unmount", () => { + const store = createStore(); + const { result, unmount } = renderHook(() => + useBatchedDataStreamAppend(store.setDataStream), + ); + + act(() => { + result.current.appendDataPart(part("a")); + }); + unmount(); + act(() => { + jest.advanceTimersByTime(DATA_STREAM_BATCH_WINDOW_MS); + }); + + expect(store.setDataStream).not.toHaveBeenCalled(); + }); +}); diff --git a/app/hooks/useBatchedDataStreamAppend.ts b/app/hooks/useBatchedDataStreamAppend.ts new file mode 100644 index 000000000..4a7be5017 --- /dev/null +++ b/app/hooks/useBatchedDataStreamAppend.ts @@ -0,0 +1,62 @@ +"use client"; + +import { useCallback, useEffect, useRef } from "react"; +import type { ScopedDataUIPart } from "@/app/components/DataStreamProvider"; + +/** + * Data parts that arrive within this window are appended in one state update. + * Terminal output alone can emit dozens of parts a second, and every append + * used to copy the whole array and re-run every consumer effect. + */ +export const DATA_STREAM_BATCH_WINDOW_MS = 100; + +type SetDataStream = React.Dispatch>; + +/** + * Throttle-first batching for `dataStream` appends: the first part opens a + * window, later parts are absorbed into it, and one `setDataStream` runs when + * the window closes. Nothing is dropped and order is preserved. `clear` + * discards pending parts as well so a stale chat cannot resurrect after reset. + */ +export function useBatchedDataStreamAppend( + setDataStream: SetDataStream, + windowMs: number = DATA_STREAM_BATCH_WINDOW_MS, +) { + const pendingRef = useRef([]); + const timerRef = useRef | null>(null); + + const cancelTimer = () => { + if (timerRef.current !== null) { + clearTimeout(timerRef.current); + timerRef.current = null; + } + }; + + const flushDataStream = useCallback(() => { + cancelTimer(); + const pending = pendingRef.current; + if (pending.length === 0) return; + pendingRef.current = []; + setDataStream((current) => current.concat(pending)); + }, [setDataStream]); + + const appendDataPart = useCallback( + (part: ScopedDataUIPart) => { + pendingRef.current.push(part); + if (timerRef.current === null) { + timerRef.current = setTimeout(flushDataStream, windowMs); + } + }, + [flushDataStream, windowMs], + ); + + const clearDataStream = useCallback(() => { + cancelTimer(); + pendingRef.current = []; + setDataStream([]); + }, [setDataStream]); + + useEffect(() => cancelTimer, []); + + return { appendDataPart, flushDataStream, clearDataStream }; +} diff --git a/lib/__tests__/wide-event-timing.test.ts b/lib/__tests__/wide-event-timing.test.ts new file mode 100644 index 000000000..919e2bfe8 --- /dev/null +++ b/lib/__tests__/wide-event-timing.test.ts @@ -0,0 +1,48 @@ +import { createWideEventBuilder } from "../logger"; + +describe("wide event preflight and first-chunk timing", () => { + beforeEach(() => { + jest.useFakeTimers({ now: 1_000_000 }); + }); + + afterEach(() => { + jest.useRealTimers(); + }); + + it("records preflight, first chunk, and stream durations", () => { + const builder = createWideEventBuilder("chat-1", "/api/chat"); + + jest.advanceTimersByTime(120); + builder.startStream(); + jest.advanceTimersByTime(300); + builder.markFirstChunk(); + jest.advanceTimersByTime(500); + // Later chunks must not move the first-chunk measurement. + builder.markFirstChunk(); + builder.setStreamResult({ + finishReason: "stop", + wasAborted: false, + wasPreemptiveTimeout: false, + hadSummarization: false, + }); + builder.setSuccess(); + + const event = builder.build(); + expect(event.preflight).toEqual({ duration_ms: 120 }); + expect(event.stream?.first_chunk_ms).toBe(300); + expect(event.stream?.duration_ms).toBe(800); + }); + + it("omits first_chunk_ms when no chunk arrived", () => { + const builder = createWideEventBuilder("chat-1", "/api/chat"); + builder.startStream(); + builder.setStreamResult({ + wasAborted: true, + wasPreemptiveTimeout: false, + hadSummarization: false, + }); + builder.setAborted(); + + expect(builder.build().stream).not.toHaveProperty("first_chunk_ms"); + }); +}); diff --git a/lib/api/__tests__/chat-stream-helpers-notes.test.ts b/lib/api/__tests__/chat-stream-helpers-notes.test.ts index dda39cdc5..2844b21f7 100644 --- a/lib/api/__tests__/chat-stream-helpers-notes.test.ts +++ b/lib/api/__tests__/chat-stream-helpers-notes.test.ts @@ -7,6 +7,7 @@ */ import { + injectNotesIntoMessages, replaceNotesBlock, refreshNotesInModelMessages, } from "@/lib/api/chat-stream-helpers"; @@ -347,3 +348,52 @@ describe("refreshNotesInModelMessages", () => { expect(text).not.toContain("Old Note"); }); }); + +// ── injectNotesIntoMessages preload ───────────────────────────────────────── + +describe("injectNotesIntoMessages", () => { + beforeEach(() => { + mockGetNotes.mockReset(); + }); + + const userMessage = { + id: "u1", + role: "user" as const, + parts: [{ type: "text" as const, text: "hello" }], + }; + const note = { + note_id: "note_1", + title: "Preloaded", + content: "some content", + tags: ["general"], + updated_at: Date.parse("2024-01-15T00:00:00Z"), + }; + + it("uses preloaded notes instead of fetching again", async () => { + const result = await injectNotesIntoMessages([userMessage], { + userId: "user-1", + subscription: "pro", + shouldIncludeNotes: true, + preloadedNotes: Promise.resolve([note] as never), + }); + + expect(mockGetNotes).not.toHaveBeenCalled(); + const text = (result[0].parts[0] as { text: string }).text; + expect(text).toContain(""); + expect(text).toContain("Preloaded"); + }); + + it("falls back to fetching when no preload is provided", async () => { + mockGetNotes.mockResolvedValue([note]); + + const result = await injectNotesIntoMessages([userMessage], { + userId: "user-1", + subscription: "pro", + shouldIncludeNotes: true, + }); + + expect(mockGetNotes).toHaveBeenCalledTimes(1); + const text = (result[0].parts[0] as { text: string }).text; + expect(text).toContain("Preloaded"); + }); +}); diff --git a/lib/api/chat-handler.ts b/lib/api/chat-handler.ts index 6b4a09293..6e2d22568 100644 --- a/lib/api/chat-handler.ts +++ b/lib/api/chat-handler.ts @@ -100,14 +100,15 @@ import { import { geolocation } from "@vercel/functions"; import { NextRequest } from "next/server"; import { - handleInitialChatAndUserMessage, - saveMessage, - updateChat, - updateChatTitle, getMessagesByChatId, + getNotes, getUserCustomization, + handleInitialChatAndUserMessage, prepareForNewStream, + saveMessage, startStream, + updateChat, + updateChatTitle, } from "@/lib/db/actions"; import { createCancellationSubscriber, @@ -347,25 +348,46 @@ export const createChatHandler = () => { }); } - const userCustomization = await getUserCustomization({ userId }); - - const fetched = await getMessagesByChatId({ - chatId, - userId, - subscription, - newMessages: requestMessages, - regenerate, - mode, - useClientMessagesForRegenerate, - }); + // These reads only depend on the authenticated user, so overlap them + // instead of paying one Convex round-trip after another before the + // model call. Mirrors the Trigger agent route's preflight. + const [userCustomization, fetched] = await Promise.all([ + getUserCustomization({ userId }), + getMessagesByChatId({ + chatId, + userId, + subscription, + newMessages: requestMessages, + regenerate, + mode, + useClientMessagesForRegenerate, + }), + ]); const { chat, isNewChat, fileTokens } = fetched; - const projectContext = await resolveProjectExecutionContext({ - chat, - requestedProjectId, - userId, - mode, - sandboxPreference, - }); + + // Notes are injected right before streaming. Start the fetch now so it + // overlaps the remaining preflight instead of adding a serial + // round-trip at the end. getNotes never rejects (it returns [] on error). + const shouldIncludeNotes = userCustomization?.include_notes ?? true; + const preloadedNotes = shouldIncludeNotes + ? getNotes({ userId, subscription }) + : undefined; + + const [projectContext, baseExtraUsageConfig] = await Promise.all([ + resolveProjectExecutionContext({ + chat, + requestedProjectId, + userId, + mode, + sandboxPreference, + }), + buildExtraUsageConfig({ + userId, + subscription, + userCustomization, + organizationId, + }), + ]); const truncatedMessages = subscription === "free" ? stripImageAttachments(fetched.truncatedMessages) @@ -375,13 +397,6 @@ export const createChatHandler = () => { (chat?.todos as unknown as Todo[]) || [], { regenerate }, ); - - const baseExtraUsageConfig = await buildExtraUsageConfig({ - userId, - subscription, - userCustomization, - organizationId, - }); const extraUsageAvailable = canUseExtraUsage(baseExtraUsageConfig); selectedModelOverride = normalizeMaxModelForSubscription(selectedModelOverride, subscription, { @@ -937,16 +952,15 @@ export const createChatHandler = () => { // Inject notes into messages instead of system prompt // to keep the system prompt stable for prompt caching - const shouldIncludeNotes = userCustomization?.include_notes ?? true; const noteInjectionOpts = { userId, subscription, shouldIncludeNotes, }; - finalMessages = await injectNotesIntoMessages( - finalMessages, - noteInjectionOpts, - ); + finalMessages = await injectNotesIntoMessages(finalMessages, { + ...noteInjectionOpts, + preloadedNotes, + }); // Mutable stream state — updated in-place by the shared runner. const state = initAgentStreamState( @@ -1414,6 +1428,7 @@ export const createChatHandler = () => { ctxSystemTokens, ctxMaxTokens, streamStartTime, + onModelChunk: () => chatLogger?.markFirstChunk(), contextUsageOn, isReasoningModel, platformAuthorized, diff --git a/lib/api/chat-logger.ts b/lib/api/chat-logger.ts index 9c2943419..0fd61b9ac 100644 --- a/lib/api/chat-logger.ts +++ b/lib/api/chat-logger.ts @@ -608,6 +608,13 @@ export function createChatLogger(config: ChatLoggerConfig) { builder.startStream(); }, + /** + * Record the first model chunk (first call wins within a request) + */ + markFirstChunk() { + builder.markFirstChunk(); + }, + /** * Set sandbox execution info */ diff --git a/lib/api/chat-stream-helpers.ts b/lib/api/chat-stream-helpers.ts index 4da378cff..6f7f90fda 100644 --- a/lib/api/chat-stream-helpers.ts +++ b/lib/api/chat-stream-helpers.ts @@ -1123,15 +1123,22 @@ export async function injectNotesIntoMessages( userId: string; subscription: SubscriptionTier; shouldIncludeNotes: boolean; + /** + * Notes fetch started earlier in the request so it overlaps other + * preflight work instead of adding a round-trip right before the model + * call. Falls back to fetching here when absent. + */ + preloadedNotes?: Promise>>; }, ): Promise { if (!opts.shouldIncludeNotes) return messages; try { - const notes = await getNotes({ - userId: opts.userId, - subscription: opts.subscription, - }); + const notes = await (opts.preloadedNotes ?? + getNotes({ + userId: opts.userId, + subscription: opts.subscription, + })); const notesContent = generateNotesSection(notes); if (!notesContent) return messages; diff --git a/lib/logger.ts b/lib/logger.ts index 5402abc84..399948c20 100644 --- a/lib/logger.ts +++ b/lib/logger.ts @@ -134,9 +134,19 @@ export interface ChatWideEvent { }; }; + // Time from wide-event creation (request parsed) to the first provider + // call. Everything before the stream — auth, Convex reads, rate limits, + // tokenizer passes — lands here, so a slow first token can be attributed + // to preflight versus the provider. + preflight?: { + duration_ms: number; + }; + // Stream execution stream?: { duration_ms: number; + // Stream start to the first model chunk. Absent when no chunk arrived. + first_chunk_ms?: number; finish_reason?: string; was_aborted: boolean; was_preemptive_timeout: boolean; @@ -237,7 +247,9 @@ export interface ChatWideEvent { export class WideEventBuilder { private event: Partial; private toolCalls: Array<{ name: string; sandbox_type?: string }> = []; + private readonly createdAtMs = Date.now(); private streamStartTime?: number; + private firstChunkTime?: number; private anthropicPromptRepairCount = 0; constructor(requestId: string, chatId: string, endpoint: ChatApiEndpoint) { @@ -421,6 +433,20 @@ export class WideEventBuilder { */ startStream(): this { this.streamStartTime = Date.now(); + this.event.preflight = { + duration_ms: this.streamStartTime - this.createdAtMs, + }; + return this; + } + + /** + * Record the first model chunk. First call wins so provider retries and + * fallbacks do not move the measurement. + */ + markFirstChunk(): this { + if (this.firstChunkTime === undefined) { + this.firstChunkTime = Date.now(); + } return this; } @@ -469,6 +495,10 @@ export class WideEventBuilder { }): this { this.event.stream = { duration_ms: this.streamStartTime ? Date.now() - this.streamStartTime : 0, + ...(this.streamStartTime !== undefined && + this.firstChunkTime !== undefined && { + first_chunk_ms: this.firstChunkTime - this.streamStartTime, + }), finish_reason: result.finishReason, was_aborted: result.wasAborted, was_preemptive_timeout: result.wasPreemptiveTimeout, From 74d18ccc785216cba96661f230a1cc5fdae07d16 Mon Sep 17 00:00:00 2001 From: Ross <144740362+ross0x01@users.noreply.github.com> Date: Thu, 3 Sep 2026 08:44:18 -0500 Subject: [PATCH 2/2] Clear pre-existing exhaustive-deps warnings in chat.tsx The two flagged setters are referentially stable (a useState setter and an empty-deps useCallback from GlobalState), so listing them changes nothing at runtime. The queue-clear effect reads the latest-value ref at cleanup on purpose; move the disable comment onto that read and list the stable ref object as a dependency. Co-Authored-By: Claude Fable 5.1 --- app/components/chat.tsx | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/app/components/chat.tsx b/app/components/chat.tsx index 3a792e034..0ba092348 100644 --- a/app/components/chat.tsx +++ b/app/components/chat.tsx @@ -1608,6 +1608,7 @@ export const Chat = ({ autoResume }: { autoResume: boolean }) => { return () => setChatReset(null); }, [ setChatReset, + setHasUserDismissedRateLimitWarning, setMessages, setStreamedTitle, setTodos, @@ -1739,7 +1740,7 @@ export const Chat = ({ autoResume }: { autoResume: boolean }) => { if (coerced) { setSelectedModel(coerced); } - }, [chatData, isExistingChat, chatId]); + }, [chatData, isExistingChat, chatId, setSelectedModel]); // Persist picker preferences (model + mode) when the user toggles them. // Debounced so quick toggles don't spam Convex; baseline is seeded from the @@ -1920,12 +1921,12 @@ export const Chat = ({ autoResume }: { autoResume: boolean }) => { // Intentionally reads messageQueueRef at cleanup time (latest value). useEffect(() => { return () => { + // eslint-disable-next-line react-hooks/exhaustive-deps if (messageQueueRef.current.length > 0) { clearQueue(); } }; - // eslint-disable-next-line react-hooks/exhaustive-deps - }, [chatId, clearQueue]); + }, [chatId, clearQueue, messageQueueRef]); // Document-level drag and drop listeners encapsulated in a hook useDocumentDragAndDrop({