From 8a3e432dd0bca2d8aeaaede0690590076606a2d6 Mon Sep 17 00:00:00 2001 From: Bohdan Triapitsyn Date: Sat, 17 Jan 2026 02:32:22 +0200 Subject: [PATCH] perf: optimize message loading based on streaming state Decreased message limit to 90 for historical views when not streaming. Add fetch buffer to preload additional messages for efficient syncing. --- packages/ui/src/hooks/useMessageSync.ts | 16 +++++++---- packages/ui/src/lib/opencode/client.ts | 5 ++-- packages/ui/src/stores/messageStore.ts | 29 +++++++++++--------- packages/ui/src/stores/types/sessionTypes.ts | 2 ++ 4 files changed, 31 insertions(+), 21 deletions(-) diff --git a/packages/ui/src/hooks/useMessageSync.ts b/packages/ui/src/hooks/useMessageSync.ts index 86fe55c3..1e57831a 100644 --- a/packages/ui/src/hooks/useMessageSync.ts +++ b/packages/ui/src/hooks/useMessageSync.ts @@ -1,6 +1,7 @@ import React from 'react'; import type { AssistantMessage, Message, Part } from '@opencode-ai/sdk/v2'; -import { useSessionStore, MEMORY_LIMITS } from '@/stores/useSessionStore'; +import { useSessionStore } from '@/stores/useSessionStore'; +import { MEMORY_LIMITS } from '@/stores/types/sessionTypes'; import { opencodeClient } from '@/lib/opencode/client'; import { readSessionCursor } from '@/lib/messageCursorPersistence'; import { extractTextFromPart } from '@/stores/utils/messageUtils'; @@ -124,7 +125,10 @@ export const useMessageSync = () => { const currentMessages = (messages.get(currentSessionId) || []) as SessionMessageRecord[]; - const latestMessages = (await opencodeClient.getSessionMessages(currentSessionId)) as SessionMessageRecord[]; + const memoryState = useSessionStore.getState().sessionMemoryState.get(currentSessionId); + const targetLimit = memoryState?.isStreaming ? MEMORY_LIMITS.VIEWPORT_MESSAGES : MEMORY_LIMITS.HISTORICAL_MESSAGES; + const fetchLimit = targetLimit + MEMORY_LIMITS.FETCH_BUFFER; + const latestMessages = (await opencodeClient.getSessionMessages(currentSessionId, fetchLimit)) as SessionMessageRecord[]; const cursorRecord = await readSessionCursor(currentSessionId); if (!latestMessages) return; @@ -167,7 +171,7 @@ export const useMessageSync = () => { } else { if (isUserMessageInfo(lastLocalMessage.info)) { - const messagesToLoad = latestMessages.slice(-MEMORY_LIMITS.VIEWPORT_MESSAGES); + const messagesToLoad = latestMessages.slice(-targetLimit); console.log('[SYNC] Local user message missing by ID; merging latest messages for deduplication'); const { syncMessages } = useSessionStore.getState(); syncMessages(currentSessionId, messagesToLoad); @@ -181,7 +185,7 @@ export const useMessageSync = () => { if (cursorIndex !== -1) { if (cursorIndex < latestMessages.length - 1) { const newMessages = latestMessages.slice(cursorIndex + 1); - const limited = newMessages.slice(-MEMORY_LIMITS.VIEWPORT_MESSAGES); + const limited = newMessages.slice(-targetLimit); if (limited.length > 0) { console.log(`[SYNC] Restoring ${limited.length} messages after cursor`); const { syncMessages } = useSessionStore.getState(); @@ -190,13 +194,13 @@ export const useMessageSync = () => { } } else if (latestMessages.length > 0) { console.log('[SYNC] Cursor not found on server response, loading recent messages'); - const messagesToLoad = latestMessages.slice(-MEMORY_LIMITS.VIEWPORT_MESSAGES); + const messagesToLoad = latestMessages.slice(-targetLimit); const { syncMessages } = useSessionStore.getState(); syncMessages(currentSessionId, messagesToLoad); } } else if (latestMessages.length > 0) { - const messagesToLoad = latestMessages.slice(-MEMORY_LIMITS.VIEWPORT_MESSAGES); + const messagesToLoad = latestMessages.slice(-targetLimit); console.log(`[SYNC] Loading last ${messagesToLoad.length} messages`); const { syncMessages } = useSessionStore.getState(); syncMessages(currentSessionId, messagesToLoad); diff --git a/packages/ui/src/lib/opencode/client.ts b/packages/ui/src/lib/opencode/client.ts index 0a44acbf..5dfe5d8e 100644 --- a/packages/ui/src/lib/opencode/client.ts +++ b/packages/ui/src/lib/opencode/client.ts @@ -354,10 +354,11 @@ class OpencodeService { return response.data; } - async getSessionMessages(id: string): Promise<{ info: Message; parts: Part[] }[]> { + async getSessionMessages(id: string, limit?: number): Promise<{ info: Message; parts: Part[] }[]> { const response = await this.client.session.messages({ sessionID: id, - ...(this.currentDirectory ? { directory: this.currentDirectory } : {}) + ...(this.currentDirectory ? { directory: this.currentDirectory } : {}), + ...(typeof limit === 'number' ? { limit } : {}), }); return response.data || []; } diff --git a/packages/ui/src/stores/messageStore.ts b/packages/ui/src/stores/messageStore.ts index 9bda2cc2..ceb1ff46 100644 --- a/packages/ui/src/stores/messageStore.ts +++ b/packages/ui/src/stores/messageStore.ts @@ -385,8 +385,11 @@ export const useMessageStore = create()( pendingAssistantHeaderSessions: new Set(), pendingUserMessageMetaBySession: new Map(), - loadMessages: async (sessionId: string, limit: number = MEMORY_LIMITS.VIEWPORT_MESSAGES) => { - const allMessages = await executeWithSessionDirectory(sessionId, () => opencodeClient.getSessionMessages(sessionId)); + loadMessages: async (sessionId: string, limit: number = MEMORY_LIMITS.HISTORICAL_MESSAGES) => { + const isStreaming = get().sessionMemoryState.get(sessionId)?.isStreaming; + const targetLimit = isStreaming ? MEMORY_LIMITS.VIEWPORT_MESSAGES : limit; + const fetchLimit = isStreaming ? undefined : targetLimit + MEMORY_LIMITS.FETCH_BUFFER; + const allMessages = await executeWithSessionDirectory(sessionId, () => opencodeClient.getSessionMessages(sessionId, fetchLimit)); // Filter out reverted messages first const revertMessageId = getSessionRevertMessageId(sessionId); @@ -401,7 +404,7 @@ export const useMessageStore = create()( return isIdNewer(messageId, watermark); }) : messagesWithoutReverted; - const messagesToKeep = afterWatermark.slice(-limit); + const messagesToKeep = afterWatermark.slice(-targetLimit); set((state) => { const newMessages = new Map(state.messages); @@ -2418,28 +2421,28 @@ export const useMessageStore = create()( const allMessages = await executeWithSessionDirectory(sessionId, () => opencodeClient.getSessionMessages(sessionId)); if (direction === "up" && currentMessages.length > 0) { - + const dedupedMessages = dedupeMessagesById(allMessages); const firstCurrentMessage = currentMessages[0]; - const indexInAll = allMessages.findIndex((m) => m.info.id === firstCurrentMessage.info.id); + const indexInAll = dedupedMessages.findIndex((message) => message.info.id === firstCurrentMessage.info.id); if (indexInAll > 0) { - const loadCount = Math.min(MEMORY_LIMITS.VIEWPORT_MESSAGES, indexInAll); - const newMessages = allMessages.slice(indexInAll - loadCount, indexInAll); + const newMessages = dedupedMessages.slice(indexInAll - loadCount, indexInAll); set((state) => { const updatedMessages = [...newMessages, ...currentMessages]; - const dedupedMessages = dedupeMessagesById(updatedMessages); - const newMessagesMap = new Map(state.messages); - newMessagesMap.set(sessionId, dedupedMessages); + const mergedMessages = dedupeMessagesById(updatedMessages); + const addedCount = Math.max(0, mergedMessages.length - currentMessages.length); + + const newMessagesMap = new Map(state.messages); + newMessagesMap.set(sessionId, mergedMessages); - const addedCount = Math.max(0, dedupedMessages.length - currentMessages.length); const newMemoryState = new Map(state.sessionMemoryState); newMemoryState.set(sessionId, { ...memoryState, viewportAnchor: memoryState.viewportAnchor + addedCount, hasMoreAbove: indexInAll - loadCount > 0, - totalAvailableMessages: allMessages.length, + totalAvailableMessages: dedupedMessages.length, }); return { @@ -2448,12 +2451,12 @@ export const useMessageStore = create()( }; }); } else if (indexInAll === 0) { - set((state) => { const newMemoryState = new Map(state.sessionMemoryState); newMemoryState.set(sessionId, { ...memoryState, hasMoreAbove: false, + totalAvailableMessages: dedupedMessages.length, }); return { sessionMemoryState: newMemoryState }; }); diff --git a/packages/ui/src/stores/types/sessionTypes.ts b/packages/ui/src/stores/types/sessionTypes.ts index e2f557f1..93e41771 100644 --- a/packages/ui/src/stores/types/sessionTypes.ts +++ b/packages/ui/src/stores/types/sessionTypes.ts @@ -54,6 +54,8 @@ export interface SessionContextUsage { export const MEMORY_LIMITS = { MAX_SESSIONS: 3, VIEWPORT_MESSAGES: 120, + HISTORICAL_MESSAGES: 90, + FETCH_BUFFER: 20, STREAMING_BUFFER: Infinity, BACKGROUND_STREAMING_BUFFER: 120, ZOMBIE_TIMEOUT: 10 * 60 * 1000,