Function restoreMessageInCache

packages/kilo-chat-hooks/src/use-messages.ts:224–253

(
  queryClient: ReturnType<typeof useQueryClient>,
  queryKey: readonly unknown[],
  snapshot: Message,
  shouldRestore?: RestoreMessageGuard
)

 * optimistic mutations can pass a guard so newer live events are not replaced
 * by stale snapshots.
 */
export function restoreMessageInCache(
  queryClient: ReturnType<typeof useQueryClient>,
  queryKey: readonly unknown[],
  snapshot: Message,
  shouldRestore?: RestoreMessageGuard
): boolean {
  let restored = false;
  queryClient.setQueryData<MessageInfiniteData>(queryKey, old => {
    if (!old) return old;
    for (let pageIndex = 0; pageIndex < old.pages.length; pageIndex++) {
      const page = old.pages[pageIndex];
      if (!page) continue;
      const messageIndex = page.messages.findIndex(msg => msg.id === snapshot.id);
      if (messageIndex === -1) continue;

      const current = page.messages[messageIndex];
      if (!current) return old;
      if (shouldRestore && !shouldRestore(current)) return old;

      const pages = old.pages.slice();
      const updatedMessages = page.messages.slice();
      updatedMessages[messageIndex] = snapshot;
      pages[pageIndex] = withPageMessages(page, updatedMessages);
      restored = true;
      return { ...old, pages };
    }
    return old;
  });
  return restored;
}

/**