Threads 10: Store ordered thread snapshots in ChatJS - #240
Threads 10: Store ordered thread snapshots in ChatJS#240FranciscoMoretti wants to merge 4 commits into
Conversation
|
Bugbot is not enabled for your account, so this pull request was not reviewed. Enable Bugbot in the Cursor dashboard to get automatic reviews on future PRs. |
![sourcery-ai[bot]](https://avatars.githubusercontent.com/in/48477?s=60&v=4){kind=small}
There was a problem hiding this comment.
Sorry @FranciscoMoretti, you have reached your weekly rate limit of 500000 diff characters.
Please try again later or upgrade to continue using Sourcery
|
The latest updates on your projects. Learn more about Vercel for GitHub.
|
📝 WalkthroughWalkthroughThe thread store now uses ChangesThread snapshot refactor
Estimated code review effort: 4 (Complex) | ~45 minutes Possibly related issues
Possibly related PRs
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 71d3dfc64a
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
03bcd2a to
61f6927
Compare
71d3dfc to
5a30d6f
Compare
![cubic-dev-ai[bot]](https://avatars.githubusercontent.com/in/1082092?s=60&v=4){kind=small}
There was a problem hiding this comment.
All reported issues were addressed
Reply with feedback, questions, or to request a fix.
Re-trigger cubic
Greptile SummaryThis PR refactors the ChatJS store to carry a canonical
Confidence Score: 3/5The refactor is behaviorally correct in the happy path but setAllMessages fires two separate Zustand set calls in the non-streaming path, leaving a window where messages has advanced while allMessages, childrenMap, and _throttledMessages still reflect the old tree. The double-set in setAllMessages is the clearest concrete defect introduced: any subscriber combining messages with childrenMap can render with a new message list but stale sibling ordering until the second set lands. Several previously flagged gaps compound the risk in this same file, and the test suite does not exercise torn-state or streaming-interaction edge cases. Files Needing Attention: apps/chat/lib/stores/with-threads.ts — specifically the setAllMessages non-streaming path and setTreeSnapshot. Important Files Changed
Reviews (27): Last reviewed commit: "refactor(chat): derive indexes from orde..." | Re-trigger Greptile |
579d626 to
1212bbc
Compare
fcf9e9d to
3fc1ed6
Compare
1212bbc to
8668264
Compare
bb2e6db to
f9d5553
Compare
8668264 to
7078da9
Compare
7078da9 to
7f16e34
Compare
f9d5553 to
68fb205
Compare
7f16e34 to
c12565e
Compare
68fb205 to
41c5189
Compare
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 5
🧹 Nitpick comments (9)
apps/chat/lib/stores/with-threads.ts (7)
504-518: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDuplicated metadata-preservation logic; reuse
mergeMessageIntoMap.Lines 504-518 replicate the branch in
mergeMessageIntoMap(lines 263-281) verbatim. Extract the shared "keep existing metadata when the update omits it" rule into one helper so the two paths can't drift.♻️ Sketch
function withPreservedMetadata<UM extends UIMessage>( existing: UM | undefined, message: UM ): UM { const existingMetadata = (existing as (UM & MessageNode) | undefined)?.metadata; if ( existing && (message as UM & MessageNode).metadata === undefined && existingMetadata !== undefined ) { return { ...message, metadata: { ...existingMetadata } } as UM; } return message; }
mergeMessageIntoMapbecomesmerged.set(message.id, withPreservedMetadata(merged.get(message.id), message)), andaddMessageToTreeusesnext[idx] = withPreservedMetadata(existing, message).🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@apps/chat/lib/stores/with-threads.ts` around lines 504 - 518, Extract the duplicated metadata-preservation branch into a shared helper such as withPreservedMetadata near mergeMessageIntoMap. Update mergeMessageIntoMap and addMessageToTree to call this helper, preserving existing metadata only when the incoming message omits it and otherwise returning the incoming message unchanged.
362-366: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winSnapshot is built three times during initialization.
buildTreeSnapshotFromMessages(base.messages)runs at line 362, again insidegetSnapshotSignatureat line 363-365, and a third time viarebuildMap's default parameter at line 366. Build once and reuse.♻️ Proposed refactor
+ const initialSnapshot = buildTreeSnapshotFromMessages(base.messages); return { ...base, threadEpoch: 0, threadInitialMessages: base.messages, allMessages: base.messages, - treeSnapshot: buildTreeSnapshotFromMessages(base.messages), - treeSnapshotSignature: getSnapshotSignature( - buildTreeSnapshotFromMessages(base.messages) - ), - childrenMap: rebuildMap(base.messages), + treeSnapshot: initialSnapshot, + treeSnapshotSignature: getSnapshotSignature(initialSnapshot), + childrenMap: rebuildMap(base.messages, initialSnapshot),Note that
initialSnapshotmust be declared before the returned object literal to avoid a TDZ access.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@apps/chat/lib/stores/with-threads.ts` around lines 362 - 366, In the initialization flow surrounding the returned object, compute buildTreeSnapshotFromMessages(base.messages) once in an initialSnapshot variable declared before the object literal, then reuse it for treeSnapshot and getSnapshotSignature. Update rebuildMap to receive the same snapshot explicitly instead of triggering its default snapshot construction.
248-252: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winFull
JSON.stringifyof the tree as a change signature is expensive and order-sensitive.This serializes every message (all parts and metadata) and runs on every store mutation —
setMessagesWithEpoch,setTreeSnapshot,setAllMessages,addMessageToTree,setMessages— i.e. per streaming chunk. It's also key-order dependent, so two structurally identical snapshots whose message objects were built with different property insertion order produce different signatures, defeating the dedupe at line 420. A compact derived signature (e.g.cursorIdplusid:parentIdpairs and a per-message revision marker) would be cheaper and stable.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@apps/chat/lib/stores/with-threads.ts` around lines 248 - 252, Replace the full JSON.stringify implementation in getSnapshotSignature with a compact, stable signature derived from cursorId and each message’s id, parentId, and revision marker. Avoid serializing message parts or metadata, and ensure the derived ordering is deterministic so structurally identical snapshots deduplicate consistently at the existing call sites.
113-127: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winAvoid re-allocating the sibling array on every insert.
Lines 119-122 rebuild each children array per message (
O(k²)copies per parent bucket).buildTreeSnapshotFromMessagesruns on nearly every store mutation, including per streaming update, so push into the existing array instead.♻️ Proposed refactor
for (const message of messages) { const metadataParentId = getMetadataParentId(message); const parentId = metadataParentId && messagesById.has(metadataParentId) ? metadataParentId : null; - childrenByParentId.set(parentId, [ - ...(childrenByParentId.get(parentId) ?? []), - message, - ]); + const siblings = childrenByParentId.get(parentId); + if (siblings) { + siblings.push(message); + } else { + childrenByParentId.set(parentId, [message]); + } }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@apps/chat/lib/stores/with-threads.ts` around lines 113 - 127, Update the message-grouping loop in buildTreeSnapshotFromMessages to retrieve or initialize each parent’s sibling array once, then append the current message with push instead of creating a copied array on every insert. Preserve the existing parentId assignment and subsequent compareSiblingMessages sorting.
332-353: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
parentByIddefaulting to{}makes fallback-metadata derivation a silent no-op.
setAllMessages(line 451) calls this with three arguments, soaddFallbackMetadataToMessagesresolves everyparentIdtonulland never fills anything; onlysetTreeSnapshotpasses the map. If that asymmetry is intentional it deserves a comment; otherwise make the parameter required (or derive it from the merged messages) so the behavior can't be lost by omission.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@apps/chat/lib/stores/with-threads.ts` around lines 332 - 353, Make the parentById input to mergeTreeMessages required instead of defaulting to an empty object, and update setAllMessages to provide the appropriate parent map (or derive it from the merged messages) before addFallbackMetadataToMessages runs. Preserve setTreeSnapshot’s existing map flow so fallback metadata is populated consistently and cannot silently no-op when the argument is omitted.
156-175: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick win
getSnapshotIndexesis rebuilt on every call, including per recursion level.It walks all snapshot nodes and (line 166) copies each child array per insert. Callers make this quadratic on hot paths:
findLeafDfsToRightFromSnapshot(line 235) rebuilds the full index at every recursion depth.getMessageSiblingInfo(line 545) andgetParallelGroupInfo(line 567) rebuild it per message lookup, and these run per rendered message.Consider computing the indexes once per snapshot — e.g. cache keyed on
treeSnapshotSignature, or pass the indexes into the traversal helpers instead of the snapshot.♻️ Minimal fix for the per-insert copy plus a snapshot-keyed cache
- for (const { message, parentId } of snapshot.nodes) { - const key = parentKey(parentId); - childrenByParentId[key] = [...(childrenByParentId[key] ?? []), message.id]; + for (const { message, parentId } of snapshot.nodes) { + const key = parentKey(parentId); + (childrenByParentId[key] ??= []).push(message.id); messagesById[message.id] = message; parentById[message.id] = parentId; if (parentId === null) { rootIds.push(message.id); } }Plus a module-scope memo so repeated calls with the same snapshot object are free:
const snapshotIndexCache = new WeakMap< MessageTreeSnapshot<UIMessage>, ReturnType<typeof computeSnapshotIndexes> >();🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@apps/chat/lib/stores/with-threads.ts` around lines 156 - 175, Refactor getSnapshotIndexes to avoid rebuilding indexes during repeated lookups: add a module-scope WeakMap cache keyed by the MessageTreeSnapshot object, return the cached result when available, and store newly computed indexes before returning. While computing, append message IDs to childrenByParentId arrays without copying the existing array on every insertion, and ensure findLeafDfsToRightFromSnapshot, getMessageSiblingInfo, and getParallelGroupInfo reuse the cached indexes.
409-414: 🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick winAlign
setTreeSnapshotvisible-message preservation withsetAllMessages.
setTreeSnapshotmergescurrentVisibleMessagesas[], so any message on the visible path that has not yet been persisted intoallMessagesis discarded when the snapshot is applied. If this can run duringstreaming/submitted, preservestate.messagesthere, matchingsetAllMessages; otherwise add a comment why snapshots are never applied mid-stream.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@apps/chat/lib/stores/with-threads.ts` around lines 409 - 414, Update setTreeSnapshot’s mergeTreeMessages call to pass the current visible messages from state.messages instead of an empty array, preserving unpersisted messages during streaming or submitted states and matching setAllMessages behavior; if snapshots are guaranteed not to apply mid-stream, document that invariant at this call site instead.apps/chat/lib/stores/with-threads.test.ts (2)
366-381: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winCoverage gap: roots whose
metadata.parentMessageIdpoints outside the tree.All roots here have
parentMessageId: null, sogetMessageSiblingInfo'sparentById[...] ?? getMetadataParentId(...)fallback happens to resolve tonulleither way. A message whose metadata parent is absent fromallMessagesis treated as a root bybuildTreeSnapshotFromMessagesbut the fallback then resolves to the dangling id, yielding empty siblings (see the store comment on lines 544-546). Add a case for that shape.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@apps/chat/lib/stores/with-threads.test.ts` around lines 366 - 381, The test setup currently covers only roots with null metadata parents; add a case in the sibling-navigation test using a message whose metadata.parentMessageId references an ID absent from allMessages. Verify buildTreeSnapshotFromMessages and getMessageSiblingInfo treat that message as a root and return the expected root siblings rather than an empty list, preserving the existing root assertions.
49-57: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winStub omits
status, so the streaming code path is never exercised.
setAllMessagesbranches onstate.status === "streaming" | "submitted"to keep the visible path authoritative and preserve in-flight messages — a central behavior of this refactor — but the stub has nostatus, so every test runs the non-streaming branch. Addstatusto the stub and a case that syncs stale server data while streaming, assertingmessagesis untouched whileallMessages/treeSnapshotabsorb the server nodes.♻️ Stub addition
({ _messageIndex: { update: () => undefined }, _memoizedSelectors: new Map(), _throttledMessages: initialMessages, messages: initialMessages, + status: "ready", setMessages: (messages: ChatMessage[]) => set({ messages }), }) as unknown as BaseChatStoreState<ChatMessage>🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@apps/chat/lib/stores/with-threads.test.ts` around lines 49 - 57, Update the BaseChatStoreState stub in with-threads.test.ts to include a status field, then add a streaming-status test for setAllMessages that supplies stale server data and verifies visible messages remain unchanged while allMessages and treeSnapshot incorporate the server nodes.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@apps/chat/lib/stores/with-threads.ts`:
- Around line 477-480: Update the nextVisibleThread selection around
currentLeafId and buildThreadFromSnapshot so an empty currentVisibleMessages
path falls back to the snapshot’s rightmost leaf and adopts the merged server
history. Preserve the existing currentLeafId-based behavior when visible
messages are present, and keep the empty result only when no suitable snapshot
leaf exists.
- Around line 544-546: In both parent lookups within
apps/chat/lib/stores/with-threads.ts at lines 544-546 and 567-569, replace
nullish-coalescing fallback with a presence check against
getSnapshotIndexes(state.treeSnapshot).parentById, preserving stored null values
for genuine roots and only falling back to getMetadataParentId(message) or
metadata?.parentMessageId when the message ID is absent; preferably centralize
this logic in one shared helper.
- Around line 210-226: Update buildThreadFromSnapshot to remove the hard-coded
100-iteration limit, track visited message IDs (or bound traversal by the
snapshot node count) for cycle safety, and preserve every acyclic ancestor.
Replace thread.unshift with push during traversal, then reverse the collected
thread before returning; use a descriptive named bound if a snapshot-size limit
is chosen.
- Line 488: Update the setAllMessages flow so it does not assign
threadInitialMessages; preserve that field as a remount seed updated only during
thread switches or store initialization. Keep setAllMessages responsible for
current message state, and retain the existing thread-switch assignment to
nextVisibleThread.
- Around line 424-430: Update setTreeSnapshot’s nextVisibleThread logic so a
null mergedSnapshot.cursorId does not clear the existing visible messages when
mergedMessages still contains thread content; preserve the current
messages/_throttledMessages or derive the visible path from the merged snapshot.
Add or adjust the emitsNullCursorWithMessages test to verify allMessages and the
visible thread remain consistent.
---
Nitpick comments:
In `@apps/chat/lib/stores/with-threads.test.ts`:
- Around line 366-381: The test setup currently covers only roots with null
metadata parents; add a case in the sibling-navigation test using a message
whose metadata.parentMessageId references an ID absent from allMessages. Verify
buildTreeSnapshotFromMessages and getMessageSiblingInfo treat that message as a
root and return the expected root siblings rather than an empty list, preserving
the existing root assertions.
- Around line 49-57: Update the BaseChatStoreState stub in with-threads.test.ts
to include a status field, then add a streaming-status test for setAllMessages
that supplies stale server data and verifies visible messages remain unchanged
while allMessages and treeSnapshot incorporate the server nodes.
In `@apps/chat/lib/stores/with-threads.ts`:
- Around line 504-518: Extract the duplicated metadata-preservation branch into
a shared helper such as withPreservedMetadata near mergeMessageIntoMap. Update
mergeMessageIntoMap and addMessageToTree to call this helper, preserving
existing metadata only when the incoming message omits it and otherwise
returning the incoming message unchanged.
- Around line 362-366: In the initialization flow surrounding the returned
object, compute buildTreeSnapshotFromMessages(base.messages) once in an
initialSnapshot variable declared before the object literal, then reuse it for
treeSnapshot and getSnapshotSignature. Update rebuildMap to receive the same
snapshot explicitly instead of triggering its default snapshot construction.
- Around line 248-252: Replace the full JSON.stringify implementation in
getSnapshotSignature with a compact, stable signature derived from cursorId and
each message’s id, parentId, and revision marker. Avoid serializing message
parts or metadata, and ensure the derived ordering is deterministic so
structurally identical snapshots deduplicate consistently at the existing call
sites.
- Around line 113-127: Update the message-grouping loop in
buildTreeSnapshotFromMessages to retrieve or initialize each parent’s sibling
array once, then append the current message with push instead of creating a
copied array on every insert. Preserve the existing parentId assignment and
subsequent compareSiblingMessages sorting.
- Around line 332-353: Make the parentById input to mergeTreeMessages required
instead of defaulting to an empty object, and update setAllMessages to provide
the appropriate parent map (or derive it from the merged messages) before
addFallbackMetadataToMessages runs. Preserve setTreeSnapshot’s existing map flow
so fallback metadata is populated consistently and cannot silently no-op when
the argument is omitted.
- Around line 156-175: Refactor getSnapshotIndexes to avoid rebuilding indexes
during repeated lookups: add a module-scope WeakMap cache keyed by the
MessageTreeSnapshot object, return the cached result when available, and store
newly computed indexes before returning. While computing, append message IDs to
childrenByParentId arrays without copying the existing array on every insertion,
and ensure findLeafDfsToRightFromSnapshot, getMessageSiblingInfo, and
getParallelGroupInfo reuse the cached indexes.
- Around line 409-414: Update setTreeSnapshot’s mergeTreeMessages call to pass
the current visible messages from state.messages instead of an empty array,
preserving unpersisted messages during streaming or submitted states and
matching setAllMessages behavior; if snapshots are guaranteed not to apply
mid-stream, document that invariant at this call site instead.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 685a148b-4d14-480e-a8b9-f43ab6def433
⛔ Files ignored due to path filters (1)
bun.lockis excluded by!**/*.lock
📒 Files selected for processing (3)
apps/chat/lib/stores/with-threads.test.tsapps/chat/lib/stores/with-threads.tsapps/chat/package.json
|
Superseded by the same branch rebased onto the new package compatibility layer. Review continues in #267; this PR remains available for its existing discussion history. |
1 participant
Summary
useThread.Behavior
Refactor only. No intended user-visible behavior change.
Verification
Review focus
Snapshot synchronization and preservation of current store semantics.
Summary by cubic
Stores a canonical
MessageTreeSnapshotfrom@chatjs/threadin the ChatJS store and derives navigation, sibling order, and indexes from ordered snapshots. Adds ordered snapshot merging and visible-path hydration, and finalizes assistant continuation, native regeneration, and restore-safe tool/approval ownership, with no user-visible changes.Refactors
treeSnapshot,treeSnapshotSignature, andsetTreeSnapshot; merge incoming snapshots with existing state; rebuild the visible path fromcursorId(empty whennull); no-op when unchanged.childrenMap, parent links, and sibling order (byparallelIndex, then time); recover invalid/cyclic topology as extra roots; preserve paths of any depth; hydrate an empty visible path from server trees.activeStreamId); dedupe root user siblings.streaming/submitted, update only the snapshot/indexes;setMessagesWithEpoch/setAllMessages/addMessageToTreerebuild the snapshot and indexes, reset selector caches and_throttledMessages, update_messageIndex.switchToSibling/switchToMessage;getParallelGroupInfois available before any assistant exists.New Features
sendMessage()on a selected assistant continues that node; passing an explicit assistant message streams into the same message ID; branching still requires a new input.regenerate({ messageId })uses AI SDK’sregenerate-messagetrigger, creates a sibling replacement (including root assistants), and only auto-follows when the cursor still targets the message; rejects assistant-with-assistant parent until an SDK fix.Written for commit edf3932. Summary will update on new commits.
Summary by CodeRabbit
New Features
Bug Fixes
Stack