Skip to content

WIP2: maybe chat optimizations #28715

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
chrisnojima-zoom wants to merge 10 commits into
base: master
Choose a base branch
from
Open

WIP2: maybe chat optimizations #28715

Changes from all commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
ebd3a04
WIP
chrisnojima-zoom Dec 18, 2025
527f734
WIP
chrisnojima-zoom Dec 18, 2025
da73633
WIP
chrisnojima-zoom Dec 18, 2025
c53a77a
WIP
chrisnojima-zoom Dec 19, 2025
7e62dc3
mm
chrisnojima-zoom Dec 19, 2025
2eb93e7
Merge remote-tracking branch 'origin/master' into nojima/ZCLIENT-chat…
chrisnojima Dec 19, 2025
7b4c8d1
WIP
chrisnojima Dec 19, 2025
eefa59a
WIP
chrisnojima-zoom Dec 19, 2025
7b0f1a1
WIP
chrisnojima Dec 19, 2025
fa8edc6
WIP
chrisnojima Dec 19, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
139 changes: 108 additions & 31 deletions shared/constants/chat2/convostate.tsx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@ import * as Z from '@/util/zustand'
import {makeActionForOpenPathInFilesTab} from '@/constants/fs/util'
import HiddenString from '@/util/hidden-string'
import isEqual from 'lodash/isEqual'
import sortedIndexBy from 'lodash/sortedIndexBy'
import logger from '@/logger'
import throttle from 'lodash/throttle'
import type {DebouncedFunc} from 'lodash'
Expand Down Expand Up @@ -115,9 +116,10 @@ type ConvoStore = T.Immutable<{
loaded: boolean // did we ever load this thread yet
markedAsUnread: T.Chat.Ordinal
messageCenterOrdinal?: T.Chat.CenterOrdinal // ordinals to center threads on,
messageTypeMap: Map<T.Chat.Ordinal, T.Chat.RenderMessageType> // messages T.Chat to help the thread, text is never used
messageOrdinals?: ReadonlyArray<T.Chat.Ordinal> // ordered ordinals in a thread,
messageIDToOrdinalMap: Map<T.Chat.MessageID, T.Chat.Ordinal> // reverse lookup for O(1) messageID -> ordinal
messageMap: Map<T.Chat.Ordinal, T.Chat.Message> // messages in a thread,
messageOrdinals?: ReadonlyArray<T.Chat.Ordinal> // ordered ordinals in a thread,
messageTypeMap: Map<T.Chat.Ordinal, T.Chat.RenderMessageType> // messages T.Chat to help the thread, text is never used
meta: T.Chat.ConversationMeta // metadata about a thread, There is a special node for the pending conversation,
moreToLoadBack: boolean
moreToLoadForward: boolean
Expand Down Expand Up @@ -153,6 +155,7 @@ const initialConvoStore: ConvoStore = {
loaded: false,
markedAsUnread: T.Chat.numberToOrdinal(0),
messageCenterOrdinal: undefined,
messageIDToOrdinalMap: new Map(),
messageMap: new Map(),
messageOrdinals: undefined,
messageTypeMap: new Map(),
Expand Down Expand Up @@ -338,8 +341,20 @@ const makeAttachmentViewInfo = (): T.Chat.AttachmentViewInfo => ({
const messageIDToOrdinal = (
map: ConvoState['messageMap'],
pendingOutboxToOrdinal: ConvoState['pendingOutboxToOrdinal'] | undefined,
messageID: T.Chat.MessageID
messageID: T.Chat.MessageID,
messageIDToOrdinalMap?: ConvoState['messageIDToOrdinalMap']
) => {
// Fast path: use reverse lookup map if available
if (messageIDToOrdinalMap) {
const cachedOrdinal = messageIDToOrdinalMap.get(messageID)
if (cachedOrdinal !== undefined) {
const m = map.get(cachedOrdinal)
if (m?.id !== 0 && m?.id === messageID) {
return cachedOrdinal
}
}
}

// A message we didn't send in this session?
let m = map.get(T.Chat.numberToOrdinal(messageID))
if (m?.id !== 0 && m?.id === messageID) {
Expand Down Expand Up @@ -446,6 +461,25 @@ const createSlice: Z.ImmerStateCreator<ConvoState> = (set, get) => {
}

const syncMessageDerived = (s: Z.WritableDraft<ConvoState>) => {
const currentSize = s.messageOrdinals?.length ?? 0
const mapSize = s.messageMap.size

// Early exit: if sizes match and we have ordinals, check if we actually need to recalculate
if (currentSize === mapSize && s.messageOrdinals) {
// Quick check: verify all ordinals in messageOrdinals still exist and are regular messages
let needsRecalc = false
for (const ord of s.messageOrdinals) {
const m = s.messageMap.get(ord)
if (!m || m.conversationMessage === false) {
needsRecalc = true
break
}
}
if (!needsRecalc) {
return
}
}

const mo = [...s.messageMap]
.filter(([, m]) => {
const regularMessage = m.conversationMessage !== false
Expand Down Expand Up @@ -523,6 +557,10 @@ const createSlice: Z.ImmerStateCreator<ConvoState> = (set, get) => {
const regularMessage = m.conversationMessage !== false

if (regularMessage && m.type === 'deleted') {
const oldMessage = s.messageMap.get(m.ordinal)
if (oldMessage && oldMessage.id !== 0) {
s.messageIDToOrdinalMap.delete(oldMessage.id)
}
s.messageMap.delete(m.ordinal)
s.messageTypeMap.delete(m.ordinal)
} else {
Expand All @@ -547,7 +585,15 @@ const createSlice: Z.ImmerStateCreator<ConvoState> = (set, get) => {
m.ordinal = mapOrdinal
}

const oldMessage = s.messageMap.get(mapOrdinal)
if (oldMessage && oldMessage.id !== 0 && oldMessage.id !== m.id) {
s.messageIDToOrdinalMap.delete(oldMessage.id)
}

s.messageMap.set(mapOrdinal, T.castDraft(m))
if (regularMessage && m.id !== 0) {
s.messageIDToOrdinalMap.set(m.id, mapOrdinal)
}
if (
regularMessage &&
m.outboxID &&
Expand Down Expand Up @@ -763,7 +809,8 @@ const createSlice: Z.ImmerStateCreator<ConvoState> = (set, get) => {
const ordinal = messageIDToOrdinal(
get().messageMap,
get().pendingOutboxToOrdinal,
T.Chat.numberToMessageID(msgID)
T.Chat.numberToMessageID(msgID),
get().messageIDToOrdinalMap
)
if (!ordinal) {
logger.info(`downloadComplete: no ordinal found: conversationIDKey: ${get().id} msgID: ${msgID}`)
Expand All @@ -788,7 +835,8 @@ const createSlice: Z.ImmerStateCreator<ConvoState> = (set, get) => {
const ordinal = messageIDToOrdinal(
get().messageMap,
get().pendingOutboxToOrdinal,
T.Chat.numberToMessageID(msgID)
T.Chat.numberToMessageID(msgID),
get().messageIDToOrdinalMap
)
if (!ordinal) {
logger.info(`downloadProgress: no ordinal found: conversationIDKey: ${get().id} msgID: ${msgID}`)
Expand Down Expand Up @@ -903,6 +951,9 @@ const createSlice: Z.ImmerStateCreator<ConvoState> = (set, get) => {
const existing = get().messageMap.get(toDelOrdinal)
if (existing) {
set(s => {
if (existing.id !== 0) {
s.messageIDToOrdinalMap.delete(existing.id)
}
s.messageMap.delete(toDelOrdinal)
syncMessageDerived(s)
})
Expand Down Expand Up @@ -956,7 +1007,8 @@ const createSlice: Z.ImmerStateCreator<ConvoState> = (set, get) => {
const ordinal = messageIDToOrdinal(
get().messageMap,
get().pendingOutboxToOrdinal,
T.Chat.numberToMessageID(placeholderID)
T.Chat.numberToMessageID(placeholderID),
get().messageIDToOrdinalMap
)
const existing = ordinal ? get().messageMap.get(ordinal) : undefined
if (ordinal && existing) {
Expand Down Expand Up @@ -1426,15 +1478,22 @@ const createSlice: Z.ImmerStateCreator<ConvoState> = (set, get) => {
if (m) {
// conversationMessage is used to tell if its this gallery load or not but if we
// load a message we already have we don't want to overwrite that it really belongs
const message = {...m, conversationMessage: get().messageMap.has(m.ordinal)}
const message: T.Chat.Message = {...m, conversationMessage: get().messageMap.has(m.ordinal)}
set(s => {
const info = mapGetEnsureValue(
s.attachmentViewMap,
viewType,
T.castDraft(makeAttachmentViewInfo())
)
if (!info.messages.find(item => item.id === message.id)) {
info.messages = info.messages.concat(T.castDraft(message)).sort((l, r) => r.id - l.id)
// Use lodash sortedIndexBy with reversed comparator for descending order
// sortedIndexBy assumes ascending, so we negate the ID to reverse the sort
const insertIndex = sortedIndexBy(
info.messages,
message,
m => -T.Chat.messageIDToNumber(m.id)
)
info.messages.splice(insertIndex, 0, T.castDraft(message))
}
})
// inject them into the message map
Expand Down Expand Up @@ -2016,7 +2075,9 @@ const createSlice: Z.ImmerStateCreator<ConvoState> = (set, get) => {
ignorePromise(f())
},
messagesClear: () => {
convIDCache.clear()
set(s => {
s.messageIDToOrdinalMap.clear()
s.pendingOutboxToOrdinal.clear()
s.loaded = false
s.messageMap.clear()
Expand All @@ -2028,7 +2089,12 @@ const createSlice: Z.ImmerStateCreator<ConvoState> = (set, get) => {
logger.info(`messagesExploded: exploding ${messageIDs.length} messages`)
set(s => {
messageIDs.forEach(mid => {
const ordinal = messageIDToOrdinal(s.messageMap, s.pendingOutboxToOrdinal, mid)
const ordinal = messageIDToOrdinal(
s.messageMap,
s.pendingOutboxToOrdinal,
mid,
s.messageIDToOrdinalMap
)
const m = ordinal && s.messageMap.get(ordinal)
if (!m) return
m.exploded = true
Expand All @@ -2050,33 +2116,38 @@ const createSlice: Z.ImmerStateCreator<ConvoState> = (set, get) => {
ordinals = [],
upToMessageID = null,
} = p
const {pendingOutboxToOrdinal, messageMap} = get()
const {messageIDToOrdinalMap, pendingOutboxToOrdinal, messageMap} = get()

const allOrdinals = new Set<T.Chat.Ordinal>()

// Add explicit ordinals
ordinals.forEach(ord => {
if (ord) allOrdinals.add(ord)
})

let upToOrdinals: Array<T.Chat.Ordinal> = []
// Add ordinals from messageIDs using reverse lookup map (O(1) per lookup)
messageIDs.forEach(messageID => {
const ordinal =
messageIDToOrdinalMap.get(messageID) ??
messageIDToOrdinal(messageMap, pendingOutboxToOrdinal, messageID, messageIDToOrdinalMap)
if (ordinal) allOrdinals.add(ordinal)
})

// Single pass: collect upToOrdinals and add to set in one iteration
if (upToMessageID) {
upToOrdinals = [...messageMap.entries()].reduce((arr, [ordinal, m]) => {
for (const [ordinal, m] of messageMap.entries()) {
if (m.id < upToMessageID && deletableMessageTypes.has(m.type)) {
arr.push(ordinal)
allOrdinals.add(ordinal)
}
return arr
}, new Array<T.Chat.Ordinal>())
}

const allOrdinals = new Set(
[
...ordinals,
...messageIDs.map(messageID => messageIDToOrdinal(messageMap, pendingOutboxToOrdinal, messageID)),
...upToOrdinals,
].reduce<Array<T.Chat.Ordinal>>((arr, n) => {
if (n) {
arr.push(n)
}
return arr
}, [])
)
}
}

set(s => {
allOrdinals.forEach(ordinal => {
const m = s.messageMap.get(ordinal)
if (m && m.id !== 0) {
s.messageIDToOrdinalMap.delete(m.id)
}
s.messageMap.delete(ordinal)
})
syncMessageDerived(s)
Expand Down Expand Up @@ -2347,7 +2418,12 @@ const createSlice: Z.ImmerStateCreator<ConvoState> = (set, get) => {
if (!message) {
return
}
const ordinal = messageIDToOrdinal(get().messageMap, get().pendingOutboxToOrdinal, messageID)
const ordinal = messageIDToOrdinal(
get().messageMap,
get().pendingOutboxToOrdinal,
messageID,
get().messageIDToOrdinalMap
)
set(s => {
const existing = ordinal ? s.messageMap.get(ordinal) : undefined
if (existing) {
Expand Down Expand Up @@ -3166,7 +3242,8 @@ const createSlice: Z.ImmerStateCreator<ConvoState> = (set, get) => {
const targetOrdinal = messageIDToOrdinal(
get().messageMap,
get().pendingOutboxToOrdinal,
u.targetMsgID
u.targetMsgID,
get().messageIDToOrdinalMap
)
if (!targetOrdinal) {
logger.info(
Expand Down