fix(chat): preserve literal code around source chips

Author: waleedlatif1

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

@@ -6,6 +6,48 @@ import { sanitizeChatDisplayContent } from '@/app/workspace/[workspaceId]/home/c
 import { scalingRatioOver4x } from '@/app/workspace/[workspaceId]/home/components/message-content/components/scaling-test-helpers'
 
 describe('sanitizeChatDisplayContent', () => {
+  it.each(['source', 'workspace_resource'])(
+    'unwraps %s JSON that mentions the other chip tag',
+    (name) => {
+      const otherTag = name === 'source' ? 'workspace_resource' : 'source'
+      const tag = `<${name}>${JSON.stringify({ title: `Use <${otherTag}>` })}</${name}>`
+
+      expect(sanitizeChatDisplayContent(`\`${tag}\``)).toBe(tag)
+    }
+  )
+
+  it.each([2, 3, 4])('preserves a %i-backtick code span containing a chip', (length) => {
+    const delimiter = '`'.repeat(length)
+    const tag = '<source>{"url":"https://example.com","title":"Use `config`"}</source>'
+    const content = `${delimiter}${tag}${delimiter}`
+
+    expect(sanitizeChatDisplayContent(content)).toBe(content)
+    expect(sanitizeChatDisplayContent(`${delimiter}json\n\`${tag}\`\n${delimiter}`)).toBe(
+      `${delimiter}json\n\`${tag}\`\n${delimiter}`
+    )
+    expect(sanitizeChatDisplayContent(`${content} then \`${tag}\``)).toBe(`${content} then ${tag}`)
+  })
+
+  it('does not let an unmatched backtick run suppress later citations', () => {
+    const prefix = 'Use `` for two backticks.\n'
+    const tag = '<source>{"url":"https://example.com"}</source>'
+
+    expect(sanitizeChatDisplayContent(`${prefix}\`${tag}\``)).toBe(`${prefix}${tag}`)
+  })
+
+  it('preserves fences closed by a longer run and unwraps citations after them', () => {
+    const tag = '<source>{"url":"https://example.com"}</source>'
+    const block = `\`\`\`json\n\`${tag}\`\n\`\`\`\`\n`
+
+    expect(sanitizeChatDisplayContent(`${block}\`${tag}\``)).toBe(`${block}${tag}`)
+  })
+
+  it('leaves an unclosed streaming fence literal', () => {
+    const content = '```json\n`<source>{"url":"https://example.com"}</source>`'
+
+    expect(sanitizeChatDisplayContent(content)).toBe(content)
+  })
+
   it('unwraps workspace resource tags from inline code spans', () => {
     const content =
       '`I updated <workspace_resource>{"type":"workflow","id":"wf-1","title":"Workflow"}</workspace_resource>.`'
@@ -159,4 +201,24 @@ describe('sanitizeChatDisplayContent', () => {
       '<workspace_resource>{"type":"file","path":"a.md","title":"a"}</workspace_resource> done'
     )
   })
+
+  it.each(['source', 'workspace_resource'])(
+    'stays linear on repeated %s tags with unterminated JSON strings',
+    (name) => {
+      expect(
+        scalingRatioOver4x(sanitizeChatDisplayContent, (times) =>
+          `<${name}>{${String.fromCharCode(92, 34)}`.repeat(times)
+        )
+      ).toBeLessThan(8)
+    }
+  )
+
+  it.each(['<source>"', '<source>{"key":"'])(
+    'stays linear on repeated quoted payload prefix %s',
+    (prefix) => {
+      expect(
+        scalingRatioOver4x(sanitizeChatDisplayContent, (times) => prefix.repeat(times))
+      ).toBeLessThan(8)
+    }
+  )
 })

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

@@ -1,66 +1,108 @@
 const HIDDEN_INLINE_REFERENCE_PATTERN =
   /`[^`\n]*(?:internal\/tool-results\/|internal\/blocktips\/|components\/integrations\/[^`\n]*README)[^`\n]*`/g
 
-/** JSON strings own their escaped quotes, backticks, and any quoted tag markers. */
-const JSON_STRING_SOURCE = String.raw`"(?:[^"\\\r\n]|\\[^\r\n])*"`
+/** JSON strings own their escaped quotes, backticks, and quoted tag markers. */
+const JSON_STRING_SOURCE = '"(?:\\\\(?:["\\\\/bfnrt]|u[0-9a-fA-F]{4})|[^"\\\\\\r\\n])*"'
 
-/**
- * Complete chip tags consume JSON strings atomically. Outside strings, a new
- * opener or backtick ends the candidate, so prose mentions cannot join into a
- * tag and repeated unclosed openers cannot repeatedly scan the same suffix.
- */
-const COMPLETE_TAG_SOURCE = `<(?<chipTag>workspace_resource|source)>\\s*\\{(?:${JSON_STRING_SOURCE}|[^"\`<])*?\\}\\s*</\\k<chipTag>>`
+/** Unquoted openers, backticks, and invalid backslashes bound failed payload scans. */
+const COMPLETE_TAG_SOURCE = `<(?<chipTag>workspace_resource|source)>\\s*\\{(?:${JSON_STRING_SOURCE}|[^"\`<\\\\])*?\\}\\s*</\\k<chipTag>>`
 
-const CHIP_OR_CODE_DELIMITER = new RegExp(`${COMPLETE_TAG_SOURCE}|\`|\n`, 'g')
+const INLINE_CHIP_OR_DELIMITER = new RegExp(`${COMPLETE_TAG_SOURCE}|\`+|\\n`, 'g')
 
-/**
- * Pair Markdown delimiters outside chip payloads in one forward pass. A pair
- * containing a chip is unwrapped; a lone delimiter is removed only when flush
- * against a chip. Neighbouring code spans and multiline fences keep their pairs.
- */
-export function sanitizeChatDisplayContent(content: string): string {
-  const removedDelimiters: number[] = []
-  let openingTick = -1
-  let containsChip = false
-  let adjacentToChip = false
-  let lastChipEnd = -1
+interface OpenCodeSpan {
+  index: number
+  containsChip: boolean
+  touchesChip: boolean
+}
 
-  for (const match of content.matchAll(CHIP_OR_CODE_DELIMITER)) {
-    const index = match.index
-    if (match.groups?.chipTag) {
-      if (openingTick !== -1) {
-        containsChip = true
-        adjacentToChip ||= index === openingTick + 1
-      }
-      lastChipEnd = index + match[0].length
-      continue
+/** Only matched multi-backtick runs are code; an unmatched run remains ordinary prose. */
+function unwrapInlineChips(content: string): string {
+  const remainingRuns = new Map<number, number>()
+  for (const [value] of content.matchAll(INLINE_CHIP_OR_DELIMITER)) {
+    if (value.startsWith('`') && value.length > 1) {
+      remainingRuns.set(value.length, (remainingRuns.get(value.length) ?? 0) + 1)
     }
+  }
+  const removedDelimiters: number[] = []
+  let openSpan: OpenCodeSpan | null = null
+  let previousChipEnd = -1
+  let protectedRunLength: number | null = null
 
-    if (match[0] === '\n') {
-      if (openingTick !== -1 && adjacentToChip) removedDelimiters.push(openingTick)
-      openingTick = -1
-      lastChipEnd = -1
-      continue
-    }
+  const finishLine = () => {
+    if (openSpan?.touchesChip) removedDelimiters.push(openSpan.index)
+    openSpan = null
+  }
 
-    if (openingTick === -1) {
-      openingTick = index
-      containsChip = false
-      adjacentToChip = lastChipEnd === index
+  for (const token of content.matchAll(INLINE_CHIP_OR_DELIMITER)) {
+    const [value] = token
+    const index = token.index
+    if (value === '\n') {
+      finishLine()
+      previousChipEnd = -1
+    } else if (value.startsWith('`')) {
+      if (value.length > 1) {
+        remainingRuns.set(value.length, (remainingRuns.get(value.length) ?? 1) - 1)
+      }
+      if (protectedRunLength !== null) {
+        if (value.length === protectedRunLength) protectedRunLength = null
+        continue
+      }
+      if (value.length > 1) {
+        if (!openSpan && remainingRuns.get(value.length)) protectedRunLength = value.length
+        continue
+      }
+      if (openSpan) {
+        if (openSpan.containsChip) removedDelimiters.push(openSpan.index, index)
+        openSpan = null
+      } else {
+        openSpan = { index, containsChip: false, touchesChip: previousChipEnd === index }
+      }
     } else {
-      if (containsChip) removedDelimiters.push(openingTick, index)
-      openingTick = -1
+      if (openSpan) {
+        openSpan.containsChip = true
+        openSpan.touchesChip ||= index === openSpan.index + 1
+      }
+      previousChipEnd = index + value.length
     }
   }
-
-  if (openingTick !== -1 && adjacentToChip) removedDelimiters.push(openingTick)
+  finishLine()
 
   const parts: string[] = []
-  let start = 0
+  let cursor = 0
   for (const index of removedDelimiters) {
-    parts.push(content.slice(start, index))
-    start = index + 1
+    parts.push(content.slice(cursor, index))
+    cursor = index + 1
+  }
+  parts.push(content.slice(cursor))
+  return parts.join('')
+}
+
+/** Fenced blocks are literal, including unclosed streaming fences and longer closing runs. */
+export function sanitizeChatDisplayContent(content: string): string {
+  const parts: string[] = []
+  let cursor = 0
+  let fenceStart: number | null = null
+  let fenceLength = 0
+
+  for (const line of content.matchAll(/^ {0,3}(`{3,})([^`\n]*)(?:\n|$)/gm)) {
+    if (fenceStart === null) {
+      fenceStart = line.index
+      fenceLength = line[1].length
+    } else if (line[1].length >= fenceLength && line[2].trim() === '') {
+      const end = line.index + line[0].length
+      parts.push(
+        unwrapInlineChips(content.slice(cursor, fenceStart)),
+        content.slice(fenceStart, end)
+      )
+      cursor = end
+      fenceStart = null
+    }
+  }
+
+  if (fenceStart === null) {
+    parts.push(unwrapInlineChips(content.slice(cursor)))
+  } else {
+    parts.push(unwrapInlineChips(content.slice(cursor, fenceStart)), content.slice(fenceStart))
   }
-  parts.push(content.slice(start))
   return parts.join('').replace(HIDDEN_INLINE_REFERENCE_PATTERN, '')
 }

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

@@ -27,13 +27,16 @@ function fastest(run: (content: string) => void, content: string): number {
  * through at the single size it happens to sample. Quadratic costs ~16x for 4x
  * the input; linear costs ~4x.
  */
-export function scalingRatioOver4x(run: (content: string) => void): number {
+export function scalingRatioOver4x(
+  run: (content: string) => void,
+  buildContent: (times: number) => string = buildRepeatedTagMentions
+): number {
   // Warm up first — the JIT would otherwise charge the whole compile to the
   // small sample and flatter the ratio.
-  fastest(run, buildRepeatedTagMentions(2_000))
+  fastest(run, buildContent(2_000))
 
-  const small = fastest(run, buildRepeatedTagMentions(2_000))
-  const large = fastest(run, buildRepeatedTagMentions(8_000))
+  const small = fastest(run, buildContent(2_000))
+  const large = fastest(run, buildContent(8_000))
 
   return large / small
 }