diff --git a/src/modules/blocks/blocks.hooks.ts b/src/modules/blocks/blocks.hooks.ts index 72afe14c4..355281084 100644 --- a/src/modules/blocks/blocks.hooks.ts +++ b/src/modules/blocks/blocks.hooks.ts @@ -1,6 +1,7 @@ import * as React from 'react'; import type { Diff as SanityTextDiff } from '@sanity/diff-match-patch'; +import { agiId } from '~/common/util/idUtils'; import { shallowEquals } from '~/common/util/hooks/useShallowObject'; import type { RenderBlockInputs } from './blocks.types'; @@ -43,42 +44,48 @@ export function useTextCollapser(origText: string, enable: boolean) { } +// Helper function to compare blocks without considering their IDs +function areBlocksEqualIdIgnored(block1: RenderBlockInputs[number] | undefined, block2: RenderBlockInputs[number] | undefined): boolean { + if (!block1 || !block2) + return false; + const { bkId: _, ...rest1 } = block1; + const { bkId: __, ...rest2 } = block2; + return shallowEquals(rest1, rest2); +} + export function useAutoBlocksMemo(text: string, forceCodeWithTitle: string | undefined, forceMarkdown: boolean, forceSanityTextDiffs: SanityTextDiff[] | undefined): RenderBlockInputs { // state - previous blocks, to stabilize objects const prevBlocksRef = React.useRef([]); + const prevTextRef = React.useRef(''); return React.useMemo(() => { - - // follow outside direction, or activate the auto-splitter based on content - const newBlocks: RenderBlockInputs = []; + let newBlocks: RenderBlockInputs; if (forceCodeWithTitle !== undefined) - newBlocks.push({ bkt: 'code-bk', title: forceCodeWithTitle, code: text, isPartial: false }); + newBlocks = [{ bkt: 'code-bk', title: forceCodeWithTitle, code: text, isPartial: false }]; else if (forceMarkdown) - newBlocks.push({ bkt: 'md-bk', content: text }); + newBlocks = [{ bkt: 'md-bk', content: text }]; else if (forceSanityTextDiffs && forceSanityTextDiffs.length >= 1) - newBlocks.push({ bkt: 'txt-diffs-bk', sanityTextDiffs: forceSanityTextDiffs }); + newBlocks = [{ bkt: 'txt-diffs-bk', sanityTextDiffs: forceSanityTextDiffs }]; else - newBlocks.push(...parseBlocksFromText(text)); + newBlocks = parseBlocksFromText(text); - // recycle the previous blocks if they are the same, for stable references to React - const recycledBlocks: RenderBlockInputs = []; - for (let i = 0; i < newBlocks.length; i++) { - const newBlock = newBlocks[i]; - const prevBlock = prevBlocksRef.current[i] ?? undefined; + const recycledBlocks: RenderBlockInputs = newBlocks.map((newBlock, index) => { + const prevBlock = prevBlocksRef.current[index] ?? undefined; + const isLastBlock = index === newBlocks.length - 1; + const isStreaming = isLastBlock && text.startsWith(prevTextRef.current); - // Check if the new block can be replaced by the previous block to maintain reference stability - if (prevBlock && shallowEquals(prevBlock, newBlock)) { - recycledBlocks.push(prevBlock); - } else { - // Once a block doesn't match, we use the new blocks from this point forward. - recycledBlocks.push(...newBlocks.slice(i)); - break; - } - } + if (areBlocksEqualIdIgnored(prevBlock, newBlock)) + return prevBlock; + + if (isStreaming && prevBlock?.bkt === newBlock.bkt) + return { ...newBlock, bkId: prevBlock.bkId }; + + return { ...newBlock, bkId: agiId('chat-block') }; + }); - // Update prevBlocksRef with the current blocks for the next render prevBlocksRef.current = recycledBlocks; + prevTextRef.current = text; return recycledBlocks; }, [forceCodeWithTitle, forceMarkdown, forceSanityTextDiffs, text]); diff --git a/src/modules/blocks/blocks.types.ts b/src/modules/blocks/blocks.types.ts index 82dcde45b..e02d1150e 100644 --- a/src/modules/blocks/blocks.types.ts +++ b/src/modules/blocks/blocks.types.ts @@ -6,6 +6,9 @@ export type RenderBlockInputs = BlockInput[]; // In order of priority from the most frequent to the least type BlockInput = { + bkId?: string; + /* Other fields remain the same */ +} & ({ /* Rendered as markdown or plain text */ bkt: 'md-bk'; content: string; @@ -28,4 +31,4 @@ type BlockInput = { /* Rendered as red/green text diffs */ bkt: 'txt-diffs-bk'; sanityTextDiffs: SanityTextDiff[]; -}; \ No newline at end of file +}); \ No newline at end of file