From b1dbde1ed065c0cda28ca2d6d5bd2a269801041a Mon Sep 17 00:00:00 2001
From: Felix Kling <felix@felix-kling.de>
Date: Sat, 21 Dec 2024 15:38:30 +0100
Subject: [PATCH] feat(onebox): Use new prompt editor when onebox is enabled
 (#6288)

This updates HumanMessageEditor to use the new prompt editor when onebox
is enabled.

## Test plan

Tested locally by enabling the feature flag for myself on S2.
---
 lib/shared/src/experimentation/FeatureFlagProvider.ts |  2 ++
 vscode/webviews/AppWrapperForTest.tsx                 | 11 ++++++++++-
 .../messageCell/human/editor/HumanMessageEditor.tsx   |  7 ++++++-
 3 files changed, 18 insertions(+), 2 deletions(-)

diff --git a/lib/shared/src/experimentation/FeatureFlagProvider.ts b/lib/shared/src/experimentation/FeatureFlagProvider.ts
index dd74c4dbb1a9..b042b2d62761 100644
--- a/lib/shared/src/experimentation/FeatureFlagProvider.ts
+++ b/lib/shared/src/experimentation/FeatureFlagProvider.ts
@@ -78,6 +78,8 @@ export enum FeatureFlag {
     CodyExperimentalOneBox = 'cody-experimental-one-box',
     /** Enable debug mode for One Box feature in Cody */
     CodyExperimentalOneBoxDebug = 'cody-experimental-one-box-debug',
+    /** Enable use of new prosemirror prompt editor */
+    CodyExperimentalPromptEditor = 'cody-experimental-prompt-editor',
 
     /** Show Edit Code option in the Cody message submit dropdown */
     CodyExperimentalShowEditCodeIntent = 'cody-experimental-show-edit-code-intent',
diff --git a/vscode/webviews/AppWrapperForTest.tsx b/vscode/webviews/AppWrapperForTest.tsx
index b0f6c686c736..8cc967a8051b 100644
--- a/vscode/webviews/AppWrapperForTest.tsx
+++ b/vscode/webviews/AppWrapperForTest.tsx
@@ -7,6 +7,7 @@ import {
     type ContextItemSymbol,
     EMPTY,
     FILE_CONTEXT_MENTION_PROVIDER,
+    FeatureFlag,
     type ModelsData,
     type ResolvedConfiguration,
     SYMBOL_CONTEXT_MENTION_PROVIDER,
@@ -79,7 +80,15 @@ export const AppWrapperForTest: FunctionComponent<{ children: ReactNode }> = ({
                                             ].filter(f => f.uri.path.includes(queryTextLower)),
                             }
                         }),
-                    evaluatedFeatureFlag: _flag => Observable.of(true),
+                    evaluatedFeatureFlag: flag => {
+                        switch (flag) {
+                            case FeatureFlag.CodyExperimentalPromptEditor:
+                                // Do not enable the experimental prompt editor in tests (yet).
+                                return Observable.of(false)
+                            default:
+                                return Observable.of(true)
+                        }
+                    },
                     prompts: makePromptsAPIWithData({
                         prompts: FIXTURE_PROMPTS,
                         commands: FIXTURE_COMMANDS,
diff --git a/vscode/webviews/chat/cells/messageCell/human/editor/HumanMessageEditor.tsx b/vscode/webviews/chat/cells/messageCell/human/editor/HumanMessageEditor.tsx
index 6dd30c7a2111..9683d4e35ed2 100644
--- a/vscode/webviews/chat/cells/messageCell/human/editor/HumanMessageEditor.tsx
+++ b/vscode/webviews/chat/cells/messageCell/human/editor/HumanMessageEditor.tsx
@@ -1,6 +1,7 @@
 import {
     type ChatMessage,
     FAST_CHAT_INPUT_TOKEN_BUDGET,
+    FeatureFlag,
     type Model,
     ModelTag,
     type SerializedPromptEditorState,
@@ -12,6 +13,7 @@ import {
 import {
     PromptEditor,
     type PromptEditorRefAPI,
+    PromptEditorV2,
     useDefaultContextForChat,
     useExtensionAPI,
 } from '@sourcegraph/prompt-editor'
@@ -31,6 +33,7 @@ import { type ClientActionListener, useClientActionListener } from '../../../../
 import { promptModeToIntent } from '../../../../../prompts/PromptsTab'
 import { useTelemetryRecorder } from '../../../../../utils/telemetry'
 import { useExperimentalOneBox } from '../../../../../utils/useExperimentalOneBox'
+import { useFeatureFlag } from '../../../../../utils/useFeatureFlags'
 import styles from './HumanMessageEditor.module.css'
 import type { SubmitButtonState } from './toolbar/SubmitButton'
 import { Toolbar } from './toolbar/Toolbar'
@@ -118,6 +121,7 @@ export const HumanMessageEditor: FunctionComponent<{
           ? 'emptyEditorValue'
           : 'submittable'
 
+    const experimentalPromptEditorEnabled = useFeatureFlag(FeatureFlag.CodyExperimentalPromptEditor)
     const experimentalOneBoxEnabled = useExperimentalOneBox()
     const [submitIntent, setSubmitIntent] = useState<ChatMessage['intent'] | undefined>(
         initialIntent || (experimentalOneBoxEnabled ? undefined : 'chat')
@@ -417,6 +421,7 @@ export const HumanMessageEditor: FunctionComponent<{
         currentChatModel?.contextWindow?.context?.user ||
         currentChatModel?.contextWindow?.input ||
         FAST_CHAT_INPUT_TOKEN_BUDGET
+    const Editor = experimentalPromptEditorEnabled ? PromptEditorV2 : PromptEditor
 
     return (
         // biome-ignore lint/a11y/useKeyWithClickEvents: only relevant to click areas
@@ -436,7 +441,7 @@ export const HumanMessageEditor: FunctionComponent<{
             onFocus={onFocus}
             onBlur={onBlur}
         >
-            <PromptEditor
+            <Editor
                 seamless={true}
                 placeholder={placeholder}
                 initialEditorState={initialEditorState}