fix(chat): keep find stable across footnotes and pending chats, and match mention chips next to punctuation

Author: waleedlatif1

apps/sim/app/workspace/[workspaceId]/home/components/message-content/components/chat-content/chat-content.tsx

@@ -366,7 +366,15 @@ const MARKDOWN_COMPONENTS = {
       </div>
     )
   },
-  a({ children, href }: { children?: React.ReactNode; href?: string }) {
+  a({
+    children,
+    href,
+    'data-footnote-ref': footnoteRef,
+  }: {
+    children?: React.ReactNode
+    href?: string
+    'data-footnote-ref'?: boolean
+  }) {
     if (href?.startsWith(SOURCE_LINK_PREFIX)) {
       return (
         <SourceReference index={Number(href.slice(SOURCE_LINK_PREFIX.length))}>
@@ -393,7 +401,13 @@ const MARKDOWN_COMPONENTS = {
       )
     }
     return (
-      <a href={href} className={PROSE_LINK_CLASS} target='_blank' rel='noopener noreferrer'>
+      <a
+        href={href}
+        data-footnote-ref={footnoteRef}
+        className={PROSE_LINK_CLASS}
+        target='_blank'
+        rel='noopener noreferrer'
+      >
         {children}
       </a>
     )

apps/sim/app/workspace/[workspaceId]/home/components/mothership-chat/chat-find-text.ts

@@ -22,6 +22,7 @@ export function getChatFindText(markdown: string): string {
       return node.value
     if (node.type === 'image' || node.type === 'imageReference') return '\uffff'
     if (node.type === 'break') return '\n'
+    if (node.type === 'footnoteDefinition') return ''
     if (!('children' in node)) return ''
     return node.children
       .map(text)

apps/sim/app/workspace/[workspaceId]/home/components/mothership-chat/use-chat-find.test.tsx

@@ -0,0 +1,128 @@
+/**
+ * @vitest-environment jsdom
+ */
+import { act, useRef } from 'react'
+import type { Virtualizer } from '@tanstack/react-virtual'
+import { createRoot } from 'react-dom/client'
+import { renderToStaticMarkup } from 'react-dom/server'
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
+import { ChatContent } from '@/app/workspace/[workspaceId]/home/components/message-content/components/chat-content/chat-content'
+import { useChatFind } from '@/app/workspace/[workspaceId]/home/components/mothership-chat/use-chat-find'
+import type { ChatMessage } from '@/app/workspace/[workspaceId]/home/types'
+
+type ChatFind = ReturnType<typeof useChatFind>
+
+interface HarnessProps {
+  chatId?: string
+  messages: ChatMessage[]
+  onRender: (find: ChatFind) => void
+}
+
+// double-cast-allowed: the hook only calls scrollToIndex on the virtualizer
+const virtualizer = { scrollToIndex: () => {} } as unknown as Virtualizer<HTMLDivElement, Element>
+
+function Harness({ chatId, messages, onRender }: HarnessProps) {
+  const containerRef = useRef<HTMLDivElement>(null)
+  const scrollElementRef = useRef<HTMLDivElement>(null)
+  const find = useChatFind({
+    chatId,
+    messages,
+    hiddenUserByIndex: [],
+    containerRef,
+    scrollElementRef,
+    virtualizer,
+  })
+  onRender(find)
+  return (
+    <div ref={containerRef} tabIndex={-1} onKeyDown={find.onKeyDown}>
+      <div ref={scrollElementRef}>
+        {messages.map((message, index) => (
+          <div
+            key={message.id}
+            data-index={index}
+            dangerouslySetInnerHTML={{
+              __html: renderToStaticMarkup(<ChatContent content={message.content} />),
+            }}
+          />
+        ))}
+      </div>
+    </div>
+  )
+}
+
+const cleanups: (() => void)[] = []
+
+function renderFind(props: Omit<HarnessProps, 'onRender'>) {
+  let find: ChatFind | undefined
+  const host = document.createElement('div')
+  document.body.append(host)
+  const root = createRoot(host)
+  const render = (next: Omit<HarnessProps, 'onRender'>) =>
+    act(() => root.render(<Harness {...next} onRender={(value) => (find = value)} />))
+  cleanups.push(() => {
+    act(() => root.unmount())
+    host.remove()
+  })
+  render(props)
+  return {
+    get current() {
+      if (!find) throw new Error('Hook was not rendered')
+      return find
+    },
+    rerender: render,
+    search(query: string) {
+      act(() => {
+        host.firstElementChild?.dispatchEvent(
+          new KeyboardEvent('keydown', { key: 'f', metaKey: true, bubbles: true })
+        )
+      })
+      act(() => find?.onQueryChange(query))
+    },
+  }
+}
+
+function assistant(id: string, content: string): ChatMessage {
+  return { id, role: 'assistant', content }
+}
+
+class TestHighlight extends Set<Range> {}
+
+beforeEach(() => {
+  vi.stubGlobal('Highlight', TestHighlight)
+  vi.stubGlobal('CSS', { highlights: new Map<string, TestHighlight>() })
+  Range.prototype.getBoundingClientRect = () => new DOMRect()
+})
+
+afterEach(() => {
+  for (const cleanup of cleanups.splice(0)) cleanup()
+  vi.unstubAllGlobals()
+  Reflect.deleteProperty(Range.prototype, 'getBoundingClientRect')
+})
+
+describe('useChatFind', () => {
+  it('keeps the open search when a pending chat is persisted, and resets on a chat switch', () => {
+    const messages = [assistant('a', 'The answer is here.')]
+    const find = renderFind({ messages })
+    find.search('answer')
+    expect(find.current).toMatchObject({ isOpen: true, query: 'answer', count: 1 })
+
+    find.rerender({ chatId: 'chat-1', messages })
+    expect(find.current).toMatchObject({ isOpen: true, query: 'answer', count: 1 })
+
+    find.rerender({ chatId: 'chat-2', messages })
+    expect(find.current).toMatchObject({ isOpen: false, query: '', count: 0 })
+  })
+
+  it('counts the same footnoted matches in the index as it highlights in the rendered message', async () => {
+    const find = renderFind({
+      chatId: 'chat',
+      messages: [assistant('a', 'Chapter 1 cites this[^1].\n\n[^1]: Page 1.')],
+    })
+    find.search('1')
+    await act(() => new Promise<void>((resolve) => requestAnimationFrame(() => resolve())))
+
+    const highlighted = CSS.highlights.get('chat-find')
+    expect(find.current.count).toBe(1)
+    expect(highlighted?.size).toBe(find.current.count)
+  })
+})

apps/sim/app/workspace/[workspaceId]/home/components/mothership-chat/use-chat-find.ts

@@ -23,7 +23,7 @@ import type { ChatMessage } from '@/app/workspace/[workspaceId]/home/types'
 
 const MATCH_LIMIT = 500
 const EXCLUDED_CONTENT =
-  'button, [aria-hidden="true"], [data-agent-group], [data-chat-activity], [data-interaction-card], [data-chat-find-ignore]'
+  'button, [aria-hidden="true"], [data-agent-group], [data-chat-activity], [data-interaction-card], [data-chat-find-ignore], [data-footnote-ref], [data-footnotes]'
 const TEXT_BLOCK = 'p, h1, h2, h3, h4, h5, h6, pre, li, td, th'
 
 interface UseChatFindProps {
@@ -127,9 +127,11 @@ export function useChatFind({
   const [index, setIndex] = useState(0)
   if (scope !== chatId) {
     setScope(chatId)
-    setIsOpen(false)
-    setQuery('')
-    setIndex(0)
+    if (scope !== undefined) {
+      setIsOpen(false)
+      setQuery('')
+      setIndex(0)
+    }
   }
   const deferredQuery = useDeferredValue(query)
   const term = isOpen ? deferredQuery.trim() : ''

apps/sim/app/workspace/[workspaceId]/home/components/user-message-content/utils.test.ts

@@ -0,0 +1,34 @@
+import { describe, expect, it, vi } from 'vitest'
+import { computeMentionRanges } from '@/app/workspace/[workspaceId]/home/components/user-message-content/utils'
+import type { ChatMessageContext } from '@/app/workspace/[workspaceId]/home/types'
+
+vi.mock('@/blocks/integration-matcher', () => ({
+  getIntegrationMatcher: () => ({ regex: null, byName: new Map() }),
+}))
+
+function spans(text: string, contexts: ChatMessageContext[]): string[] {
+  return computeMentionRanges(text, contexts).map((range) => text.slice(range.start, range.end))
+}
+
+describe('computeMentionRanges', () => {
+  const workflow: ChatMessageContext = { kind: 'workflow', label: 'Workflow' }
+
+  it('matches a mention followed by punctuation', () => {
+    expect(spans('Run @Workflow, then stop.', [workflow])).toEqual(['@Workflow'])
+    expect(spans('(see @Workflow)', [workflow])).toEqual(['@Workflow'])
+  })
+
+  it('matches every repeat of a mention separated by one space', () => {
+    expect(spans('@Workflow @Workflow', [workflow])).toEqual(['@Workflow', '@Workflow'])
+  })
+
+  it('does not match a mention that is a prefix of a longer name', () => {
+    expect(spans('@Workflow-2', [workflow])).toEqual([])
+  })
+
+  it('matches slash commands by their slash prefix', () => {
+    expect(spans('/research the topic', [{ kind: 'slash_command', label: 'research' }])).toEqual([
+      '/research',
+    ])
+  })
+})

apps/sim/app/workspace/[workspaceId]/home/components/user-message-content/utils.ts

@@ -28,9 +28,10 @@ export function computeMentionRanges(text: string, contexts: ChatMessageContext[
   for (const rawCtx of contexts) {
     if (!rawCtx.label) continue
     const ctx = withResolvedBlockType(rawCtx)
-    const prefix = ctx.kind === 'skill' || ctx.kind === 'mcp' ? '/' : '@'
+    const prefix =
+      ctx.kind === 'skill' || ctx.kind === 'mcp' || ctx.kind === 'slash_command' ? '/' : '@'
     const token = `${prefix}${ctx.label}`
-    const pattern = new RegExp(`(^|\\s)(${escapeRegExp(token)})(\\s|$)`, 'g')
+    const pattern = new RegExp(`(^|\\s)(${escapeRegExp(token)})(?=[\\s.,;:!?)\\]]|$)`, 'g')
     let match: RegExpExecArray | null
     while ((match = pattern.exec(text)) !== null) {
       const leadingSpace = match[1]