Buffer streaming chunks to reduce rerenders

2026-01-01 15:43:06 +07:00
parent 63247f73b3
commit 0c4d1c3a58
5 changed files with 406 additions and 26 deletions
@@ -79,6 +79,9 @@ export function AgentStreamView({

  // Get ws for connection status
  const ws = useSessionStore((state) => state.sessions[resolvedServerId]?.ws);
+  const streamingBuffer = useSessionStore((state) =>
+    state.sessions[resolvedServerId]?.agentStreamingBuffer?.get(agentId)
+  );

  // Get methods for file operations
  const methods = useSessionStore((state) => state.sessions[resolvedServerId]?.methods);
@@ -269,6 +272,7 @@ export function AgentStreamView({
            message={item.text}
            timestamp={item.timestamp.getTime()}
            onInlinePathPress={handleInlinePathPress}
+            isStreaming={streamingItemId === item.id}
          />
        );
        break;
@@ -350,6 +354,19 @@ export function AgentStreamView({
    [pendingPermissions, agentId]
  );

+  const streamingItem = useMemo<StreamItem | null>(() => {
+    if (!streamingBuffer) {
+      return null;
+    }
+    return {
+      kind: "assistant_message",
+      id: streamingBuffer.id,
+      text: streamingBuffer.text,
+      timestamp: streamingBuffer.timestamp,
+    };
+  }, [streamingBuffer]);
+
+  const streamingItemId = streamingItem?.id ?? null;
  const showWorkingIndicator = agent.status === "running";

  const listHeaderComponent = useMemo(() => {
@@ -379,8 +396,12 @@ export function AgentStreamView({
  }, [pendingPermissionItems, showWorkingIndicator, wsOrInert]);

  const flatListData = useMemo(() => {
-    return [...streamItems].reverse();
-  }, [streamItems]);
+    const reversed = [...streamItems].reverse();
+    if (streamingItem) {
+      reversed.unshift(streamingItem);
+    }
+    return reversed;
+  }, [streamItems, streamingItem]);

  const flatListExtraData = useMemo(
    () => ({
@@ -5,7 +5,11 @@ import { useWebSocket, type UseWebSocketReturn } from "@/hooks/use-websocket";
 import { useDaemonRequest } from "@/hooks/use-daemon-request";
 import { useSessionRpc } from "@/hooks/use-session-rpc";
 import { useAudioPlayer } from "@/hooks/use-audio-player";
-import { reduceStreamUpdate, generateMessageId, hydrateStreamState } from "@/types/stream";
+import {
+  applyStreamEventWithBuffer,
+  generateMessageId,
+  hydrateStreamState,
+} from "@/types/stream";
 import type {
  ActivityLogPayload,
  AgentSnapshotPayload,
@@ -233,6 +237,8 @@ export function SessionProvider({ children, serverUrl, serverId }: SessionProvid
  const setMessages = useSessionStore((state) => state.setMessages);
  const setCurrentAssistantMessage = useSessionStore((state) => state.setCurrentAssistantMessage);
  const setAgentStreamState = useSessionStore((state) => state.setAgentStreamState);
+  const setAgentStreamingBuffer = useSessionStore((state) => state.setAgentStreamingBuffer);
+  const clearAgentStreamingBuffer = useSessionStore((state) => state.clearAgentStreamingBuffer);
  const setInitializingAgents = useSessionStore((state) => state.setInitializingAgents);
  const setHasHydratedAgents = useSessionStore((state) => state.setHasHydratedAgents);
  const setAgents = useSessionStore((state) => state.setAgents);
@@ -675,6 +681,24 @@ export function SessionProvider({ children, serverUrl, serverId }: SessionProvid

        return changed ? next : prev;
      });
+      setAgentStreamingBuffer(serverId, (prev) => {
+        if (prev.size === 0) {
+          return prev;
+        }
+
+        const validAgentIds = new Set(agentsList.map((snapshot) => snapshot.id));
+        let changed = false;
+        const next = new Map(prev);
+
+        for (const agentId of prev.keys()) {
+          if (!validAgentIds.has(agentId)) {
+            next.delete(agentId);
+            changed = true;
+          }
+        }
+
+        return changed ? next : prev;
+      });
      setInitializingAgents(serverId, (prev) => {
        if (prev.size === 0) {
          return prev;
@@ -736,28 +760,36 @@ export function SessionProvider({ children, serverUrl, serverId }: SessionProvid

      console.log("[Session] agent_stream", { agentId, eventType: event.type });

-      setAgentStreamState(serverId, (prev) => {
-        const currentStream = prev.get(agentId) || [];
-        const lastItem = currentStream[currentStream.length - 1];
-        const lastItemText = lastItem && "text" in lastItem ? lastItem.text : "(no text)";
-        console.log(`[AGENT_STREAM BEFORE] agentId=${agentId} currentStreamLength=${currentStream.length} lastText="${lastItemText}"`);
-        console.log(`[AGENT_STREAM CHUNK] incoming type=${event.type}`);
-
-        const newStream = reduceStreamUpdate(
-          currentStream,
-          event as AgentStreamEventPayload,
-          parsedTimestamp
-        );
-
-        const newLastItem = newStream[newStream.length - 1];
-        const newLastText = newLastItem && "text" in newLastItem ? newLastItem.text : "(no text)";
-        console.log(`[AGENT_STREAM AFTER] newStreamLength=${newStream.length} newLastText="${newLastText}"`);
-
-        const next = new Map(prev);
-        next.set(agentId, newStream);
-        return next;
+      const session = useSessionStore.getState().sessions[serverId];
+      const currentStream = session?.agentStreamState.get(agentId) ?? [];
+      const currentBuffer = session?.agentStreamingBuffer.get(agentId) ?? null;
+      const { stream, buffer, changedStream, changedBuffer } = applyStreamEventWithBuffer({
+        state: currentStream,
+        buffer: currentBuffer,
+        event: event as AgentStreamEventPayload,
+        timestamp: parsedTimestamp,
      });

+      if (changedStream) {
+        setAgentStreamState(serverId, (prev) => {
+          const next = new Map(prev);
+          next.set(agentId, stream);
+          return next;
+        });
+      }
+
+      if (changedBuffer) {
+        setAgentStreamingBuffer(serverId, (prev) => {
+          const next = new Map(prev);
+          if (buffer) {
+            next.set(agentId, buffer);
+          } else {
+            next.delete(agentId);
+          }
+          return next;
+        });
+      }
+
      setInitializingAgents(serverId, (prev) => {
        const currentState = prev.get(agentId);
        if (currentState === false) {
@@ -794,6 +826,7 @@ export function SessionProvider({ children, serverUrl, serverId }: SessionProvid
        next.set(agentId, hydrated);
        return next;
      });
+      clearAgentStreamingBuffer(serverId, agentId);

      setInitializingAgents(serverId, (prev) => {
        if (!prev.has(agentId)) {
@@ -1159,6 +1192,7 @@ export function SessionProvider({ children, serverUrl, serverId }: SessionProvid
        next.delete(agentId);
        return next;
      });
+      clearAgentStreamingBuffer(serverId, agentId);

      // Remove draft input
      saveDraftInput(agentId, { text: "", images: [] });
@@ -1218,7 +1252,7 @@ export function SessionProvider({ children, serverUrl, serverId }: SessionProvid
      unsubProviderModels();
      unsubAgentDeleted();
    };
-  }, [ws, audioPlayer, serverId, setIsPlayingAudio, setMessages, setCurrentAssistantMessage, setAgentStreamState, setInitializingAgents, setAgents, setAgentLastActivity, setCommands, setPendingPermissions, setGitDiffs, setFileExplorer, setProviderModels, setHasHydratedAgents, updateConnectionStatus, getSession, saveDraftInput]);
+  }, [ws, audioPlayer, serverId, setIsPlayingAudio, setMessages, setCurrentAssistantMessage, setAgentStreamState, setAgentStreamingBuffer, clearAgentStreamingBuffer, setInitializingAgents, setAgents, setAgentLastActivity, setCommands, setPendingPermissions, setGitDiffs, setFileExplorer, setProviderModels, setHasHydratedAgents, updateConnectionStatus, getSession, saveDraftInput]);

  // Auto-flush queued messages when agent transitions from running -> not running
  useEffect(() => {
@@ -1256,6 +1290,7 @@ export function SessionProvider({ children, serverUrl, serverId }: SessionProvid
      next.set(agentId, []);
      return next;
    });
+    clearAgentStreamingBuffer(serverId, agentId);

    initializeAgentRpc
      .send({ agentId })
@@ -1267,7 +1302,7 @@ export function SessionProvider({ children, serverUrl, serverId }: SessionProvid
          return next;
        });
      });
-  }, [serverId, initializeAgentRpc, setAgentStreamState, setInitializingAgents]);
+  }, [serverId, initializeAgentRpc, setAgentStreamState, setInitializingAgents, clearAgentStreamingBuffer]);

  const refreshAgent = useCallback(({ agentId, requestId }: { agentId: string; requestId?: string }) => {
    setInitializingAgents(serverId, (prev) => {
@@ -1281,6 +1316,7 @@ export function SessionProvider({ children, serverUrl, serverId }: SessionProvid
      next.set(agentId, []);
      return next;
    });
+    clearAgentStreamingBuffer(serverId, agentId);

    refreshAgentRequest
      .execute({ agentId }, { requestKeyOverride: agentId, dedupe: false })
@@ -1292,7 +1328,7 @@ export function SessionProvider({ children, serverUrl, serverId }: SessionProvid
          return next;
        });
      });
-  }, [serverId, refreshAgentRequest, setAgentStreamState, setInitializingAgents]);
+  }, [serverId, refreshAgentRequest, setAgentStreamState, setInitializingAgents, clearAgentStreamingBuffer]);

  const requestProviderModels = useCallback((provider: any, options?: { cwd?: string }) => {
    const requestId = generateMessageId();
@@ -225,6 +225,7 @@ export interface SessionState {

  // Stream state
  agentStreamState: Map<string, StreamItem[]>;
+  agentStreamingBuffer: Map<string, { id: string; text: string; timestamp: Date }>;

  // Initializing agents
  initializingAgents: Map<string, boolean>;
@@ -283,6 +284,15 @@ interface SessionStoreActions {

  // Stream state
  setAgentStreamState: (serverId: string, state: Map<string, StreamItem[]> | ((prev: Map<string, StreamItem[]>) => Map<string, StreamItem[]>)) => void;
+  setAgentStreamingBuffer: (
+    serverId: string,
+    state:
+      | Map<string, { id: string; text: string; timestamp: Date }>
+      | ((
+          prev: Map<string, { id: string; text: string; timestamp: Date }>
+        ) => Map<string, { id: string; text: string; timestamp: Date }>)
+  ) => void;
+  clearAgentStreamingBuffer: (serverId: string, agentId: string) => void;

  // Initializing agents
  setInitializingAgents: (serverId: string, state: Map<string, boolean> | ((prev: Map<string, boolean>) => Map<string, boolean>)) => void;
@@ -368,6 +378,7 @@ function createInitialSessionState(serverId: string, ws: UseWebSocketReturn, aud
    messages: [],
    currentAssistantMessage: "",
    agentStreamState: new Map(),
+    agentStreamingBuffer: new Map(),
    initializingAgents: new Map(),
    agents: new Map(),
    commands: new Map(),
@@ -551,6 +562,49 @@ export const useSessionStore = create<SessionStore>()(
      });
    },

+    setAgentStreamingBuffer: (serverId, state) => {
+      set((prev) => {
+        const session = prev.sessions[serverId];
+        if (!session) {
+          return prev;
+        }
+        const nextState = typeof state === "function" ? state(session.agentStreamingBuffer) : state;
+        if (session.agentStreamingBuffer === nextState) {
+          return prev;
+        }
+        logSessionStoreUpdate("setAgentStreamingBuffer", serverId, { agentCount: nextState.size });
+        return {
+          ...prev,
+          sessions: {
+            ...prev.sessions,
+            [serverId]: { ...session, agentStreamingBuffer: nextState },
+          },
+        };
+      });
+    },
+
+    clearAgentStreamingBuffer: (serverId, agentId) => {
+      set((prev) => {
+        const session = prev.sessions[serverId];
+        if (!session) {
+          return prev;
+        }
+        if (!session.agentStreamingBuffer.has(agentId)) {
+          return prev;
+        }
+        const nextBuffer = new Map(session.agentStreamingBuffer);
+        nextBuffer.delete(agentId);
+        logSessionStoreUpdate("clearAgentStreamingBuffer", serverId, { agentId });
+        return {
+          ...prev,
+          sessions: {
+            ...prev.sessions,
+            [serverId]: { ...session, agentStreamingBuffer: nextBuffer },
+          },
+        };
+      });
+    },
+
    // Initializing agents
    setInitializingAgents: (serverId, state) => {
      set((prev) => {
@@ -0,0 +1,157 @@
+import { describe, expect, it } from "vitest";
+import type { AgentStreamEventPayload } from "@server/server/messages";
+import type { StreamItem } from "@/types/stream";
+import { applyStreamEventWithBuffer, type StreamingBufferEntry } from "@/types/stream";
+
+const baseTimestamp = new Date(0);
+
+const assistantChunk = (text: string): AgentStreamEventPayload => ({
+  type: "timeline",
+  provider: "codex",
+  item: {
+    type: "assistant_message",
+    text,
+  },
+});
+
+const toolCallEvent = (): AgentStreamEventPayload => ({
+  type: "timeline",
+  provider: "codex",
+  item: {
+    type: "tool_call",
+    server: "local",
+    tool: "run",
+    status: "executing",
+  },
+});
+
+const completionEvent = (): AgentStreamEventPayload => ({
+  type: "turn_completed",
+  provider: "codex",
+});
+
+const permissionEvent = (): AgentStreamEventPayload => ({
+  type: "permission_requested",
+  provider: "codex",
+  request: {
+    id: "perm-1",
+    provider: "codex",
+    name: "test",
+    kind: "tool",
+  },
+});
+
+describe("applyStreamEventWithBuffer", () => {
+  it("buffers assistant chunks without changing stream", () => {
+    const stream: StreamItem[] = [];
+    const result = applyStreamEventWithBuffer({
+      state: stream,
+      buffer: null,
+      event: assistantChunk("Hel"),
+      timestamp: baseTimestamp,
+    });
+
+    expect(result.stream).toBe(stream);
+    expect(result.changedStream).toBe(false);
+    expect(result.buffer?.text).toBe("Hel");
+  });
+
+  it("appends assistant chunks to the buffer", () => {
+    const stream: StreamItem[] = [];
+    const initial = applyStreamEventWithBuffer({
+      state: stream,
+      buffer: null,
+      event: assistantChunk("Hel"),
+      timestamp: baseTimestamp,
+    });
+    const next = applyStreamEventWithBuffer({
+      state: stream,
+      buffer: initial.buffer,
+      event: assistantChunk("lo"),
+      timestamp: baseTimestamp,
+    });
+
+    expect(next.buffer?.text).toBe("Hello");
+    expect(next.changedStream).toBe(false);
+  });
+
+  it("commits buffered message on completion", () => {
+    const stream: StreamItem[] = [];
+    const buffered = applyStreamEventWithBuffer({
+      state: stream,
+      buffer: null,
+      event: assistantChunk("Hello"),
+      timestamp: baseTimestamp,
+    });
+    const result = applyStreamEventWithBuffer({
+      state: stream,
+      buffer: buffered.buffer,
+      event: completionEvent(),
+      timestamp: baseTimestamp,
+    });
+
+    expect(result.buffer).toBe(null);
+    expect(result.stream).toHaveLength(1);
+    expect(result.stream[0].kind).toBe("assistant_message");
+    expect((result.stream[0] as { text: string }).text).toBe("Hello");
+  });
+
+  it("commits buffer before non-assistant timeline items", () => {
+    const stream: StreamItem[] = [];
+    const buffered = applyStreamEventWithBuffer({
+      state: stream,
+      buffer: null,
+      event: assistantChunk("Hello"),
+      timestamp: baseTimestamp,
+    });
+    const result = applyStreamEventWithBuffer({
+      state: stream,
+      buffer: buffered.buffer,
+      event: toolCallEvent(),
+      timestamp: baseTimestamp,
+    });
+
+    expect(result.buffer).toBe(null);
+    expect(result.stream).toHaveLength(2);
+    expect(result.stream[0].kind).toBe("assistant_message");
+    expect(result.stream[1].kind).toBe("tool_call");
+  });
+
+  it("keeps stream reference for no-op events", () => {
+    const stream: StreamItem[] = [];
+    const result = applyStreamEventWithBuffer({
+      state: stream,
+      buffer: null,
+      event: permissionEvent(),
+      timestamp: baseTimestamp,
+    });
+
+    expect(result.stream).toBe(stream);
+    expect(result.changedStream).toBe(false);
+  });
+
+  it("avoids double-committing an already flushed buffer", () => {
+    const stream: StreamItem[] = [
+      {
+        kind: "assistant_message",
+        id: "dup",
+        text: "Hello",
+        timestamp: baseTimestamp,
+      },
+    ];
+    const buffer: StreamingBufferEntry = {
+      id: "dup",
+      text: "Hello",
+      timestamp: baseTimestamp,
+    };
+    const result = applyStreamEventWithBuffer({
+      state: stream,
+      buffer,
+      event: completionEvent(),
+      timestamp: baseTimestamp,
+    });
+
+    expect(result.stream).toBe(stream);
+    expect(result.buffer).toBe(null);
+  });
+});
@@ -841,3 +841,115 @@ export function hydrateStreamState(

  return finalizeActiveThoughts(hydrated);
 }
+
+export type StreamingBufferEntry = {
+  id: string;
+  text: string;
+  timestamp: Date;
+};
+
+const STREAM_COMPLETION_EVENTS = new Set<AgentStreamEventPayload["type"]>([
+  "turn_completed",
+  "turn_failed",
+  "turn_canceled",
+]);
+
+function appendCompletedAssistantMessage(
+  state: StreamItem[],
+  buffer: StreamingBufferEntry
+): StreamItem[] {
+  if (state.some((item) => item.id === buffer.id)) {
+    return state;
+  }
+
+  const entry: AssistantMessageItem = {
+    kind: "assistant_message",
+    id: buffer.id,
+    text: buffer.text,
+    timestamp: buffer.timestamp,
+  };
+
+  return [...state, entry];
+}
+
+export function applyStreamEventWithBuffer(params: {
+  state: StreamItem[];
+  buffer: StreamingBufferEntry | null;
+  event: AgentStreamEventPayload;
+  timestamp: Date;
+}): {
+  stream: StreamItem[];
+  buffer: StreamingBufferEntry | null;
+  changedStream: boolean;
+  changedBuffer: boolean;
+} {
+  const { state, buffer, event, timestamp } = params;
+  let nextStream = state;
+  let nextBuffer = buffer;
+  let changedStream = false;
+  let changedBuffer = false;
+
+  const commitBuffer = () => {
+    if (!nextBuffer) {
+      return;
+    }
+    const next = appendCompletedAssistantMessage(nextStream, nextBuffer);
+    if (next !== nextStream) {
+      nextStream = next;
+      changedStream = true;
+    }
+    nextBuffer = null;
+    changedBuffer = true;
+  };
+
+  if (event.type === "timeline") {
+    if (event.item.type === "assistant_message") {
+      const { chunk, hasContent } = normalizeChunk(event.item.text);
+      if (!chunk) {
+        return { stream: nextStream, buffer: nextBuffer, changedStream, changedBuffer };
+      }
+      if (!hasContent && !nextBuffer) {
+        return { stream: nextStream, buffer: nextBuffer, changedStream, changedBuffer };
+      }
+      if (nextBuffer) {
+        const updatedText = `${nextBuffer.text}${chunk}`;
+        if (updatedText !== nextBuffer.text) {
+          nextBuffer = { ...nextBuffer, text: updatedText };
+          changedBuffer = true;
+        }
+      } else {
+        const idSeed = chunk.trim() || chunk;
+        nextBuffer = {
+          id: createUniqueTimelineId(nextStream, "assistant", idSeed, timestamp),
+          text: chunk,
+          timestamp,
+        };
+        changedBuffer = true;
+      }
+      return { stream: nextStream, buffer: nextBuffer, changedStream, changedBuffer };
+    }
+
+    if (nextBuffer) {
+      commitBuffer();
+    }
+
+    const reduced = reduceStreamUpdate(nextStream, event, timestamp);
+    if (reduced !== nextStream) {
+      nextStream = reduced;
+      changedStream = true;
+    }
+    return { stream: nextStream, buffer: nextBuffer, changedStream, changedBuffer };
+  }
+
+  if (STREAM_COMPLETION_EVENTS.has(event.type)) {
+    commitBuffer();
+    return { stream: nextStream, buffer: nextBuffer, changedStream, changedBuffer };
+  }
+
+  const reduced = reduceStreamUpdate(nextStream, event, timestamp);
+  if (reduced !== nextStream) {
+    nextStream = reduced;
+    changedStream = true;
+  }
+  return { stream: nextStream, buffer: nextBuffer, changedStream, changedBuffer };
+}