From b529c9f14f402064dadde710bc7728630c657802 Mon Sep 17 00:00:00 2001
From: Beatrix <beatrix@sourcegraph.com>
Date: Mon, 6 Jan 2025 20:41:32 -0800
Subject: [PATCH] clean up & apply feedback

---
 vscode/src/chat/agentic/DeepCody.ts           |  8 +++++--
 vscode/src/chat/chat-view/ChatsController.ts  |  2 ++
 .../human/editor/toolbar/AddContextButton.tsx |  8 ++++++-
 .../human/editor/toolbar/ToolboxButton.tsx    | 24 ++++++++++++-------
 4 files changed, 30 insertions(+), 12 deletions(-)

diff --git a/vscode/src/chat/agentic/DeepCody.ts b/vscode/src/chat/agentic/DeepCody.ts
index cef8473cf0e9..c7aba8626a74 100644
--- a/vscode/src/chat/agentic/DeepCody.ts
+++ b/vscode/src/chat/agentic/DeepCody.ts
@@ -41,6 +41,9 @@ import { ACTIONS_TAGS, CODYAGENT_PROMPTS } from './prompts'
  */
 export class DeepCodyAgent {
     public static readonly id = 'deep-cody'
+    /**
+     * NOTE: Currently A/B test to default to 3.5 Haiku / 3.5 Sonnet for the review step.
+     */
     public static model: string | undefined = undefined
 
     protected readonly multiplexer = new BotResponseMultiplexer()
@@ -289,7 +292,7 @@ export class DeepCodyAgent {
             await this.multiplexer.notifyTurnComplete()
         }
 
-        return accumulated.toString()
+        return accumulated.consumeAndClear()
     }
 
     protected getPrompter(items: ContextItem[]): DefaultPrompter {
@@ -313,7 +316,8 @@ export class RawTextProcessor {
         this.parts.push(str)
     }
 
-    public toString(): string {
+    // Destructive read that clears state
+    public consumeAndClear(): string {
         const joined = this.parts.join('')
         this.reset()
         return joined
diff --git a/vscode/src/chat/chat-view/ChatsController.ts b/vscode/src/chat/chat-view/ChatsController.ts
index e3a346dc9a74..8ad025d94785 100644
--- a/vscode/src/chat/chat-view/ChatsController.ts
+++ b/vscode/src/chat/chat-view/ChatsController.ts
@@ -28,6 +28,7 @@ import {
     handleCodeFromInsertAtCursor,
     handleCodeFromSaveToNewFile,
 } from '../../services/utils/codeblock-action-tracker'
+import { CodyToolProvider } from '../agentic/CodyToolProvider'
 import type { ChatIntentAPIClient } from '../context/chatIntentAPIClient'
 import type { SmartApplyResult } from '../protocol'
 import {
@@ -579,6 +580,7 @@ export class ChatsController implements vscode.Disposable {
 
     public dispose(): void {
         this.disposeAllChats()
+        CodyToolProvider.dispose()
         vscode.Disposable.from(...this.disposables).dispose()
     }
 }
diff --git a/vscode/webviews/chat/cells/messageCell/human/editor/toolbar/AddContextButton.tsx b/vscode/webviews/chat/cells/messageCell/human/editor/toolbar/AddContextButton.tsx
index 0d4446b66507..321fa2af14a3 100644
--- a/vscode/webviews/chat/cells/messageCell/human/editor/toolbar/AddContextButton.tsx
+++ b/vscode/webviews/chat/cells/messageCell/human/editor/toolbar/AddContextButton.tsx
@@ -10,7 +10,13 @@ export const AddContextButton: FunctionComponent<{
 }> = ({ onClick, className }) => (
     <Tooltip>
         <TooltipTrigger asChild>
-            <Button variant="ghost" onClick={onClick} aria-label="Add context" className={className}>
+            <Button
+                variant="ghost"
+                size="none"
+                onClick={onClick}
+                aria-label="Add context"
+                className={className}
+            >
                 <AtSignIcon className="tw-w-8 tw-h-8" strokeWidth={1.25} />
             </Button>
         </TooltipTrigger>
diff --git a/vscode/webviews/chat/cells/messageCell/human/editor/toolbar/ToolboxButton.tsx b/vscode/webviews/chat/cells/messageCell/human/editor/toolbar/ToolboxButton.tsx
index f44bc5a25268..0a57808a658d 100644
--- a/vscode/webviews/chat/cells/messageCell/human/editor/toolbar/ToolboxButton.tsx
+++ b/vscode/webviews/chat/cells/messageCell/human/editor/toolbar/ToolboxButton.tsx
@@ -10,7 +10,6 @@ import { useTelemetryRecorder } from '../../../../../../utils/telemetry'
 interface ToolboxButtonProps {
     api: WebviewToExtensionAPI
     settings: AgentToolboxSettings
-    className?: string
 }
 
 const ToolboxOptionText = {
@@ -44,7 +43,7 @@ export const ToolboxButton: FC<ToolboxButtonProps> = memo(({ settings, api }) =>
     const onSubmit = useCallback(
         (close: () => void) => {
             setIsLoading(true)
-            api.updateToolboxSettings(settingsForm).subscribe({
+            const subscription = api.updateToolboxSettings(settingsForm).subscribe({
                 next: () => {
                     setIsLoading(false)
                     close()
@@ -54,7 +53,13 @@ export const ToolboxButton: FC<ToolboxButtonProps> = memo(({ settings, api }) =>
                     setSettingsForm(settings)
                     setIsLoading(false)
                 },
+                complete: () => {
+                    setIsLoading(false)
+                },
             })
+            return () => {
+                subscription.unsubscribe()
+            }
         },
         [api.updateToolboxSettings, settingsForm, settings]
     )
@@ -149,13 +154,14 @@ export const ToolboxButton: FC<ToolboxButtonProps> = memo(({ settings, api }) =>
                     },
                 }}
             >
-                <Button variant="ghost" className="!tw-p-2 tw-relative">
-                    <BrainIcon size={16} strokeWidth={1.25} className="tw-w-8 tw-h-8" />
-                    {settings.agent?.name ? (
-                        <span className="tw-absolute tw-top-0 tw-right-0 tw-w-2 tw-h-2 tw-rounded-full tw-bg-green-500 tw-animate-[pulse_5s_ease-in-out_infinite]" />
-                    ) : (
-                        <span className="tw-absolute tw-top-0 tw-right-0 tw-w-2 tw-h-2 tw-rounded-full tw-bg-red-500 tw-animate-[pulse_5s_ease-in-out_infinite]" />
-                    )}
+                <Button variant="ghost" size="none">
+                    <BrainIcon
+                        size={16}
+                        strokeWidth={1.25}
+                        className={`tw-w-8 tw-h-8 ${
+                            settings.agent?.name ? 'tw-text-green-500' : 'tw-text-muted-foreground'
+                        }`}
+                    />
                 </Button>
             </ToolbarPopoverItem>
         </div>