From 0c4d1c3a58953bd7ba21c1be6bfdbc4832c5c3c8 Mon Sep 17 00:00:00 2001 From: Mohamed Boudra Date: Thu, 1 Jan 2026 15:43:06 +0700 Subject: [PATCH] Buffer streaming chunks to reduce rerenders --- .../app/src/components/agent-stream-view.tsx | 25 ++- packages/app/src/contexts/session-context.tsx | 84 +++++++--- packages/app/src/stores/session-store.ts | 54 ++++++ packages/app/src/types/stream-buffer.test.ts | 157 ++++++++++++++++++ packages/app/src/types/stream.ts | 112 +++++++++++++ 5 files changed, 406 insertions(+), 26 deletions(-) create mode 100644 packages/app/src/types/stream-buffer.test.ts diff --git a/packages/app/src/components/agent-stream-view.tsx b/packages/app/src/components/agent-stream-view.tsx index e1245d091..e7a06fd7a 100644 --- a/packages/app/src/components/agent-stream-view.tsx +++ b/packages/app/src/components/agent-stream-view.tsx @@ -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(() => { + 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( () => ({ diff --git a/packages/app/src/contexts/session-context.tsx b/packages/app/src/contexts/session-context.tsx index b9368de95..f391e0fd6 100644 --- a/packages/app/src/contexts/session-context.tsx +++ b/packages/app/src/contexts/session-context.tsx @@ -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(); diff --git a/packages/app/src/stores/session-store.ts b/packages/app/src/stores/session-store.ts index bdc67a589..9edd09f65 100644 --- a/packages/app/src/stores/session-store.ts +++ b/packages/app/src/stores/session-store.ts @@ -225,6 +225,7 @@ export interface SessionState { // Stream state agentStreamState: Map; + agentStreamingBuffer: Map; // Initializing agents initializingAgents: Map; @@ -283,6 +284,15 @@ interface SessionStoreActions { // Stream state setAgentStreamState: (serverId: string, state: Map | ((prev: Map) => Map)) => void; + setAgentStreamingBuffer: ( + serverId: string, + state: + | Map + | (( + prev: Map + ) => Map) + ) => void; + clearAgentStreamingBuffer: (serverId: string, agentId: string) => void; // Initializing agents setInitializingAgents: (serverId: string, state: Map | ((prev: Map) => Map)) => 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()( }); }, + 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) => { diff --git a/packages/app/src/types/stream-buffer.test.ts b/packages/app/src/types/stream-buffer.test.ts new file mode 100644 index 000000000..52966f560 --- /dev/null +++ b/packages/app/src/types/stream-buffer.test.ts @@ -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); + }); +}); diff --git a/packages/app/src/types/stream.ts b/packages/app/src/types/stream.ts index a2b703775..a9ab2a72b 100644 --- a/packages/app/src/types/stream.ts +++ b/packages/app/src/types/stream.ts @@ -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([ + "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 }; +}