Cut ask-route preflight latency and batch client data-stream appends

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<HTMLDivElement>(null);
   const computerDialogPreviousFocusRef = useRef<HTMLElement | null>(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 () => {
@@ -1605,6 +1608,7 @@ export const Chat = ({ autoResume }: { autoResume: boolean }) => {
     return () => setChatReset(null);
   }, [
     setChatReset,
+    setHasUserDismissedRateLimitWarning,
     setMessages,
     setStreamedTitle,
     setTodos,
@@ -1736,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
@@ -1917,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({

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();
+  });
+});

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<React.SetStateAction<ScopedDataUIPart[]>>;
+
+/**
+ * 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<ScopedDataUIPart[]>([]);
+  const timerRef = useRef<ReturnType<typeof setTimeout> | 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 };
+}

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");
+  });
+});

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("<notes>");
+    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");
+  });
+});