From 5cd9bca97f8c99dc7f1272d39514d32705529889 Mon Sep 17 00:00:00 2001
From: Shantur Rathore <i@shantur.com>
Date: Mon, 3 Nov 2025 16:23:51 +0000
Subject: [PATCH] Disable syntax highlighting while tool calls stream

---
 src/components/markdown.tsx  | 24 ++++++++++++++++++++++++
 src/components/tool-call.tsx |  7 ++++++-
 src/lib/markdown.ts          | 33 ++++++++++++++++++++++++++++-----
 3 files changed, 58 insertions(+), 6 deletions(-)

diff --git a/src/components/markdown.tsx b/src/components/markdown.tsx
index 4e54db99..742e49c6 100644
--- a/src/components/markdown.tsx
+++ b/src/components/markdown.tsx
@@ -6,6 +6,7 @@ interface MarkdownProps {
   part: TextPart
   isDark?: boolean
   size?: "base" | "sm" | "tight"
+  disableHighlight?: boolean
 }
 
 export function Markdown(props: MarkdownProps) {
@@ -19,11 +20,30 @@ export function Markdown(props: MarkdownProps) {
     const text = decodeHtmlEntities(rawText)
     const dark = Boolean(props.isDark)
     const themeKey = dark ? "dark" : "light"
+    const highlightEnabled = !props.disableHighlight
 
     latestRequestedText = text
 
     await initMarkdown(dark)
 
+    if (!highlightEnabled) {
+      part.renderCache = undefined
+
+      try {
+        const rendered = await renderMarkdown(text, { suppressHighlight: true })
+
+        if (latestRequestedText === text) {
+          setHtml(rendered)
+        }
+      } catch (error) {
+        console.error("Failed to render markdown:", error)
+        if (latestRequestedText === text) {
+          setHtml(text)
+        }
+      }
+      return
+    }
+
     const cache = part.renderCache
     if (cache && cache.text === text && cache.theme === themeKey) {
       setHtml(cache.html)
@@ -72,6 +92,10 @@ export function Markdown(props: MarkdownProps) {
 
     // Register listener for language loading completion
     const cleanupLanguageListener = onLanguagesLoaded(async () => {
+      if (props.disableHighlight) {
+        return
+      }
+
       const part = props.part
       const rawText = typeof part.text === "string" ? part.text : ""
       const text = decodeHtmlEntities(rawText)
diff --git a/src/components/tool-call.tsx b/src/components/tool-call.tsx
index 48454110..05b9be7f 100644
--- a/src/components/tool-call.tsx
+++ b/src/components/tool-call.tsx
@@ -290,10 +290,15 @@ export default function ToolCall(props: ToolCallProps) {
 
     const isLarge = toolName === "edit" || toolName === "write" || toolName === "patch"
     const messageClass = `message-text tool-call-markdown${isLarge ? " tool-call-markdown-large" : ""}`
+    const disableHighlight = state?.status === "running"
 
     return (
       <div class={messageClass}>
-        <Markdown part={{ type: "text", text: content }} isDark={isDark()} />
+        <Markdown
+          part={{ type: "text", text: content }}
+          isDark={isDark()}
+          disableHighlight={disableHighlight}
+        />
       </div>
     )
   }
diff --git a/src/lib/markdown.ts b/src/lib/markdown.ts
index 8f1a48f8..ac31d5fc 100644
--- a/src/lib/markdown.ts
+++ b/src/lib/markdown.ts
@@ -5,6 +5,7 @@ let highlighter: Highlighter | null = null
 let highlighterPromise: Promise<Highlighter> | null = null
 let currentTheme: "light" | "dark" = "light"
 let isInitialized = false
+let highlightSuppressed = false
 
 // Track loaded languages and queue for on-demand loading
 const loadedLanguages = new Set<string>()
@@ -81,6 +82,9 @@ function resolveLanguage(token: string): { canonical: string | null; raw: string
 }
 
 async function ensureLanguages(content: string) {
+  if (highlightSuppressed) {
+    return
+  }
   // Parse code fences to extract language tokens
   // Updated regex to capture optional language tokens and handle trailing annotations
   const codeBlockRegex = /```[ \t]*([A-Za-z0-9_.+#-]+)?[^`]*?```/g
@@ -228,6 +232,10 @@ function setupRenderer(isDark: boolean) {
 </div>
 `.trim()
 
+    if (highlightSuppressed) {
+      return `<div class="markdown-code-block" data-language="${escapedLang}" data-code="${encodedCode}">${header}<pre><code class="language-${escapedLang}">${escapeHtml(decodedCode)}</code></pre></div>`
+    }
+
     // Skip highlighting for "text" language or when highlighter is not available
     if (resolvedLang === "text" || !highlighter) {
       return `<div class="markdown-code-block" data-language="${escapedLang}" data-code="${encodedCode}">${header}<pre><code>${escapeHtml(decodedCode)}</code></pre></div>`
@@ -281,18 +289,33 @@ export function isMarkdownReady(): boolean {
   return isInitialized && highlighter !== null
 }
 
-export async function renderMarkdown(content: string): Promise<string> {
+export async function renderMarkdown(
+  content: string,
+  options?: {
+    suppressHighlight?: boolean
+  },
+): Promise<string> {
   if (!isInitialized) {
     await initMarkdown(currentTheme === "dark")
   }
 
+  const suppressHighlight = options?.suppressHighlight ?? false
   const decoded = decodeHtmlEntities(content)
 
-  // Queue language loading but don't wait for it to complete
-  await ensureLanguages(decoded)
+  if (!suppressHighlight) {
+    // Queue language loading but don't wait for it to complete
+    await ensureLanguages(decoded)
+  }
 
-  // Proceed to parse immediately - highlighting will be available on next render
-  return marked.parse(decoded) as Promise<string>
+  const previousSuppressed = highlightSuppressed
+  highlightSuppressed = suppressHighlight
+
+  try {
+    // Proceed to parse immediately - highlighting will be available on next render
+    return marked.parse(decoded) as Promise<string>
+  } finally {
+    highlightSuppressed = previousSuppressed
+  }
 }
 
 export async function getSharedHighlighter(): Promise<Highlighter> {