From ae894194217db69aa1ea9347f6e5bb9bf870c9c8 Mon Sep 17 00:00:00 2001
From: Beatrix <beatrix@sourcegraph.com>
Date: Tue, 28 Jan 2025 22:48:25 +0800
Subject: [PATCH 1/3] feat: add support for <think> tags in ChatMessageContent

This PR adds support for rendering <think> tags in the ChatMessageContent component. The <think> content is displayed in a collapsible details element, allowing users to view the AI's internal thought process.

The MarkdownFromCody component is also updated to allow the <think> element.
---
 .../ChatMessageContent/ChatMessageContent.tsx | 31 ++++++++++++++++++-
 .../webviews/components/MarkdownFromCody.tsx  |  1 +
 2 files changed, 31 insertions(+), 1 deletion(-)

diff --git a/vscode/webviews/chat/ChatMessageContent/ChatMessageContent.tsx b/vscode/webviews/chat/ChatMessageContent/ChatMessageContent.tsx
index ff5600a4ebdc..51ba6bc08d72 100644
--- a/vscode/webviews/chat/ChatMessageContent/ChatMessageContent.tsx
+++ b/vscode/webviews/chat/ChatMessageContent/ChatMessageContent.tsx
@@ -39,6 +39,22 @@ interface ChatMessageContentProps {
     className?: string
 }
 
+const extractThinkContent = (content: string): { displayContent: string; thinkContent: string } => {
+    const thinkRegex = /<think>([\s\S]*?)<\/think>/g
+    const thinkMatches = [...content.matchAll(thinkRegex)]
+
+    // Collect all think content
+    const thinkContent = thinkMatches
+        .map(match => match[1].trim())
+        .filter(Boolean)
+        .join('\n\n')
+
+    // Remove think tags from display content
+    const displayContent = content.replace(thinkRegex, '')
+
+    return { displayContent, thinkContent }
+}
+
 /**
  * A component presenting the content of a chat message.
  */
@@ -204,10 +220,23 @@ export const ChatMessageContent: React.FunctionComponent<ChatMessageContentProps
         smartApplyStates,
     ])
 
+    const { displayContent, thinkContent } = useMemo(
+        () => extractThinkContent(displayMarkdown),
+        [displayMarkdown]
+    )
+
     return (
         <div ref={rootRef} data-testid="chat-message-content">
+            {thinkContent && (
+                <details className="tw-container tw-my-4 tw-border tw-border-muted-foreground tw-rounded-md tw-p-2">
+                    <summary className="tw-px-2">Thinking...</summary>
+                    <MarkdownFromCody className="tw-my-2 tw-text-muted-foreground tw-p-2">
+                        {thinkContent}
+                    </MarkdownFromCody>
+                </details>
+            )}
             <MarkdownFromCody className={clsx(styles.content, className)}>
-                {displayMarkdown}
+                {displayContent}
             </MarkdownFromCody>
         </div>
     )
diff --git a/vscode/webviews/components/MarkdownFromCody.tsx b/vscode/webviews/components/MarkdownFromCody.tsx
index dc952a29961f..f43fcd744445 100644
--- a/vscode/webviews/components/MarkdownFromCody.tsx
+++ b/vscode/webviews/components/MarkdownFromCody.tsx
@@ -53,6 +53,7 @@ const ALLOWED_ELEMENTS = [
     'h5',
     'h6',
     'br',
+    'think',
 ]
 
 function defaultUrlProcessor(url: string): string {

From be2347cd32524c466d1b1793d78ea794c71d4d1d Mon Sep 17 00:00:00 2001
From: arafatkatze <arafat.da.khan@gmail.com>
Date: Fri, 31 Jan 2025 08:39:58 +0000
Subject: [PATCH 2/3] Adding a fix for think rendering

---
 .../ChatMessageContent/ChatMessageContent.tsx | 43 ++++++++++++++-----
 1 file changed, 33 insertions(+), 10 deletions(-)

diff --git a/vscode/webviews/chat/ChatMessageContent/ChatMessageContent.tsx b/vscode/webviews/chat/ChatMessageContent/ChatMessageContent.tsx
index 51ba6bc08d72..19ec9721f8ba 100644
--- a/vscode/webviews/chat/ChatMessageContent/ChatMessageContent.tsx
+++ b/vscode/webviews/chat/ChatMessageContent/ChatMessageContent.tsx
@@ -39,20 +39,43 @@ interface ChatMessageContentProps {
     className?: string
 }
 
-const extractThinkContent = (content: string): { displayContent: string; thinkContent: string } => {
+interface StreamingContent {
+    displayContent: string
+    thinkContent: string
+    hasThinkTag: boolean
+}
+
+const extractThinkContent = (content: string): StreamingContent => {
     const thinkRegex = /<think>([\s\S]*?)<\/think>/g
     const thinkMatches = [...content.matchAll(thinkRegex)]
-
-    // Collect all think content
-    const thinkContent = thinkMatches
+    
+    // Check if content has an unclosed think tag
+    const hasOpenThinkTag = content.includes('<think>') && 
+        (content.lastIndexOf('<think>') > content.lastIndexOf('</think>'))
+    
+    // Collect all think content, including partial content from unclosed tag
+    let thinkContent = thinkMatches
         .map(match => match[1].trim())
         .filter(Boolean)
         .join('\n\n')
+    
+    if (hasOpenThinkTag) {
+        const lastThinkContent = content.slice(content.lastIndexOf('<think>') + 7)
+        thinkContent = thinkContent ? `${thinkContent}\n\n${lastThinkContent}` : lastThinkContent
+    }
 
-    // Remove think tags from display content
-    const displayContent = content.replace(thinkRegex, '')
+    // Remove complete think tags from display content
+    let displayContent = content.replace(thinkRegex, '')
+    // Remove any unclosed think tag and its content
+    if (hasOpenThinkTag) {
+        displayContent = displayContent.slice(0, displayContent.lastIndexOf('<think>'))
+    }
 
-    return { displayContent, thinkContent }
+    return { 
+        displayContent, 
+        thinkContent,
+        hasThinkTag: thinkMatches.length > 0 || hasOpenThinkTag
+    }
 }
 
 /**
@@ -220,15 +243,15 @@ export const ChatMessageContent: React.FunctionComponent<ChatMessageContentProps
         smartApplyStates,
     ])
 
-    const { displayContent, thinkContent } = useMemo(
+    const { displayContent, thinkContent, hasThinkTag } = useMemo(
         () => extractThinkContent(displayMarkdown),
         [displayMarkdown]
     )
 
     return (
         <div ref={rootRef} data-testid="chat-message-content">
-            {thinkContent && (
-                <details className="tw-container tw-my-4 tw-border tw-border-muted-foreground tw-rounded-md tw-p-2">
+            {hasThinkTag && (
+                <details open className="tw-container tw-my-4 tw-border tw-border-muted-foreground tw-rounded-md tw-p-2">
                     <summary className="tw-px-2">Thinking...</summary>
                     <MarkdownFromCody className="tw-my-2 tw-text-muted-foreground tw-p-2">
                         {thinkContent}

From da034b970018e16424dfddee416cd45a40eaddb5 Mon Sep 17 00:00:00 2001
From: arafatkatze <arafat.da.khan@gmail.com>
Date: Fri, 31 Jan 2025 08:42:35 +0000
Subject: [PATCH 3/3] Aesthetic

---
 .../ChatMessageContent/ChatMessageContent.tsx | 49 +++++++++++++++++--
 1 file changed, 44 insertions(+), 5 deletions(-)

diff --git a/vscode/webviews/chat/ChatMessageContent/ChatMessageContent.tsx b/vscode/webviews/chat/ChatMessageContent/ChatMessageContent.tsx
index 19ec9721f8ba..e98d677a4a20 100644
--- a/vscode/webviews/chat/ChatMessageContent/ChatMessageContent.tsx
+++ b/vscode/webviews/chat/ChatMessageContent/ChatMessageContent.tsx
@@ -251,11 +251,50 @@ export const ChatMessageContent: React.FunctionComponent<ChatMessageContentProps
     return (
         <div ref={rootRef} data-testid="chat-message-content">
             {hasThinkTag && (
-                <details open className="tw-container tw-my-4 tw-border tw-border-muted-foreground tw-rounded-md tw-p-2">
-                    <summary className="tw-px-2">Thinking...</summary>
-                    <MarkdownFromCody className="tw-my-2 tw-text-muted-foreground tw-p-2">
-                        {thinkContent}
-                    </MarkdownFromCody>
+                <details 
+                    open 
+                    className={clsx(
+                        "tw-container tw-mb-4",
+                        "tw-border tw-border-gray-500/20 dark:tw-border-gray-600/40",
+                        "tw-rounded-lg tw-overflow-hidden",
+                        "tw-bg-gray-50/50 dark:tw-bg-gray-800/50",
+                        "tw-backdrop-blur-sm"
+                    )}
+                >
+                    <summary className={clsx(
+                        "tw-flex tw-items-center tw-gap-2 tw-px-3 tw-py-2",
+                        "tw-bg-gray-100/50 dark:tw-bg-gray-800/80",
+                        "tw-cursor-pointer hover:tw-bg-gray-200/50 dark:hover:tw-bg-gray-700/50",
+                        "tw-select-none tw-transition-colors"
+                    )}>
+                        <svg 
+                            className="tw-w-4 tw-h-4 tw-text-gray-500 dark:tw-text-gray-400 tw-animate-pulse" 
+                            viewBox="0 0 24 24" 
+                            fill="none" 
+                            stroke="currentColor" 
+                            strokeWidth="2"
+                        >
+                            <path 
+                                strokeLinecap="round" 
+                                strokeLinejoin="round" 
+                                d="M12 3v1m0 16v1m9-9h-1M4 12H3m15.364 6.364l-.707-.707M6.343 6.343l-.707-.707m12.728 0l-.707.707M6.343 17.657l-.707.707"
+                            />
+                        </svg>
+                        <span className="tw-text-sm tw-font-medium tw-text-gray-600 dark:tw-text-gray-300">
+                            Thinking...
+                        </span>
+                    </summary>
+                    <div className="tw-px-4 tw-py-3">
+                        <MarkdownFromCody 
+                            className={clsx(
+                                "tw-text-sm tw-text-gray-600 dark:tw-text-gray-300",
+                                "tw-prose dark:tw-prose-invert tw-max-w-none",
+                                "tw-leading-relaxed"
+                            )}
+                        >
+                            {thinkContent}
+                        </MarkdownFromCody>
+                    </div>
                 </details>
             )}
             <MarkdownFromCody className={clsx(styles.content, className)}>