highlighted code improvements

6 files changed, 46 insertions, 24 deletions

diff --git a/extension/react-app/src/components/ComboBox.tsx b/extension/react-app/src/components/ComboBox.tsx
index 3e1f3e16..97f5b57e 100644
--- a/extension/react-app/src/components/ComboBox.tsx
+++ b/extension/react-app/src/components/ComboBox.tsx
@@ -11,7 +11,12 @@ import CodeBlock from "./CodeBlock";
 import { RangeInFile } from "../../../src/client";
 import PillButton from "./PillButton";
 import HeaderButtonWithText from "./HeaderButtonWithText";
-import { Trash, LockClosed, LockOpen } from "@styled-icons/heroicons-outline";
+import {
+  Trash,
+  LockClosed,
+  LockOpen,
+  Plus,
+} from "@styled-icons/heroicons-outline";
 
 // #region styled components
 const mainInputFontSize = 16;
@@ -100,6 +105,8 @@ interface ComboBoxProps {
   highlightedCodeSections: (RangeInFile & { contents: string })[];
   deleteContextItems: (indices: number[]) => void;
   onTogglePin: () => void;
+  onToggleAddContext: () => void;
+  addingHighlightedCode: boolean;
 }
 
 const ComboBox = React.forwardRef((props: ComboBoxProps, ref) => {
@@ -249,6 +256,19 @@ const ComboBox = React.forwardRef((props: ComboBoxProps, ref) => {
         </Ul>
       </div>
       <div className="px-2 flex gap-2 items-center flex-wrap">
+        {highlightedCodeSections.length === 0 && (
+          <HeaderButtonWithText
+            text={
+              props.addingHighlightedCode ? "Adding Context" : "Add Context"
+            }
+            onClick={() => {
+              props.onToggleAddContext();
+            }}
+            inverted={props.addingHighlightedCode}
+          >
+            <Plus size="1.6em" />
+          </HeaderButtonWithText>
+        )}
         {highlightedCodeSections.length > 0 && (
           <>
             <HeaderButtonWithText
@@ -305,9 +325,8 @@ const ComboBox = React.forwardRef((props: ComboBoxProps, ref) => {
         ))}
 
         <span className="text-trueGray-400 ml-auto mr-4 text-xs">
-          Highlight code to include as context.{" "}
-          {highlightedCodeSections.length === 0 &&
-            "Otherwise using entire currently open file."}
+          Highlight code to include as context. Currently open file included by
+          default. {highlightedCodeSections.length === 0 && ""}
         </span>
       </div>
       <ContextDropdown
diff --git a/extension/react-app/src/components/PillButton.tsx b/extension/react-app/src/components/PillButton.tsx
index 2352c3ad..5a02c6b2 100644
--- a/extension/react-app/src/components/PillButton.tsx
+++ b/extension/react-app/src/components/PillButton.tsx
@@ -15,6 +15,8 @@ const Button = styled.button`
     background-color: white;
     color: black;
   }
+
+  cursor: pointer;
 `;
 
 interface PillButtonProps {
@@ -39,26 +41,13 @@ const PillButton = (props: PillButtonProps) => {
           props.onHover(false);
         }
       }}
+      onClick={() => {
+        if (props.onDelete) {
+          props.onDelete();
+        }
+      }}
     >
-      <div
-        style={{ display: "grid", gridTemplateColumns: "1fr auto", gap: "4px" }}
-      >
-        <span
-          style={{
-            cursor: "pointer",
-            color: "red",
-            borderRight: "1px solid black",
-            paddingRight: "4px",
-          }}
-          onClick={() => {
-            props.onDelete?.();
-            props.onHover?.(false);
-          }}
-        >
-          <XMark style={{ padding: "0px" }} size="1.2em" strokeWidth="2px" />
-        </span>
-        <span>{props.title}</span>
-      </div>
+      {props.title}
     </Button>
   );
 };
diff --git a/extension/react-app/src/components/UserInputContainer.tsx b/extension/react-app/src/components/UserInputContainer.tsx
index 44fdba38..a72f6098 100644
--- a/extension/react-app/src/components/UserInputContainer.tsx
+++ b/extension/react-app/src/components/UserInputContainer.tsx
@@ -14,7 +14,7 @@ interface UserInputContainerProps {
   historyNode: HistoryNode;
 }
 
-const StyledDiv = styled.div`
+const StyledDiv = styled.div<{ hidden: boolean }>`
   background-color: rgb(50 50 50);
   padding: 8px;
   padding-left: 16px;
@@ -24,6 +24,8 @@ const StyledDiv = styled.div`
   font-size: 13px;
   display: flex;
   align-items: center;
+  visibility: ${(props) => (props.hidden ? "hidden" : "visible")};
+  height: ${(props) => (props.hidden ? "0px" : "auto")};
 `;
 
 const UserInputContainer = (props: UserInputContainerProps) => {
diff --git a/extension/react-app/src/hooks/ContinueGUIClientProtocol.ts b/extension/react-app/src/hooks/ContinueGUIClientProtocol.ts
index 96ea7ab3..f123bb2b 100644
--- a/extension/react-app/src/hooks/ContinueGUIClientProtocol.ts
+++ b/extension/react-app/src/hooks/ContinueGUIClientProtocol.ts
@@ -22,6 +22,8 @@ abstract class AbstractContinueGUIClientProtocol {
   abstract deleteAtIndex(index: number): void;
 
   abstract deleteContextAtIndices(indices: number[]): void;
+
+  abstract toggleAddingHighlightedCode(): void;
 }
 
 export default AbstractContinueGUIClientProtocol;
diff --git a/extension/react-app/src/hooks/useContinueGUIProtocol.ts b/extension/react-app/src/hooks/useContinueGUIProtocol.ts
index e950387c..49f200ae 100644
--- a/extension/react-app/src/hooks/useContinueGUIProtocol.ts
+++ b/extension/react-app/src/hooks/useContinueGUIProtocol.ts
@@ -74,6 +74,10 @@ class ContinueGUIClientProtocol extends AbstractContinueGUIClientProtocol {
   deleteContextAtIndices(indices: number[]) {
     this.messenger.send("delete_context_at_indices", { indices });
   }
+
+  toggleAddingHighlightedCode(): void {
+    this.messenger.send("toggle_adding_highlighted_code", {});
+  }
 }
 
 export default ContinueGUIClientProtocol;
diff --git a/extension/react-app/src/tabs/gui.tsx b/extension/react-app/src/tabs/gui.tsx
index bbf0b126..851045d5 100644
--- a/extension/react-app/src/tabs/gui.tsx
+++ b/extension/react-app/src/tabs/gui.tsx
@@ -71,6 +71,7 @@ function GUI(props: GUIProps) {
   const [waitingForSteps, setWaitingForSteps] = useState(false);
   const [userInputQueue, setUserInputQueue] = useState<string[]>([]);
   const [highlightedRanges, setHighlightedRanges] = useState([]);
+  const [addingHighlightedCode, setAddingHighlightedCode] = useState(false);
   const [availableSlashCommands, setAvailableSlashCommands] = useState<
     { name: string; description: string }[]
   >([]);
@@ -157,6 +158,7 @@ function GUI(props: GUIProps) {
       setHistory(state.history);
       setHighlightedRanges(state.highlighted_ranges);
       setUserInputQueue(state.user_input_queue);
+      setAddingHighlightedCode(state.adding_highlighted_code);
       setAvailableSlashCommands(
         state.slash_commands.map((c: any) => {
           return {
@@ -361,6 +363,10 @@ function GUI(props: GUIProps) {
           onTogglePin={() => {
             setPinned((prev: boolean) => !prev);
           }}
+          onToggleAddContext={() => {
+            client?.toggleAddingHighlightedCode();
+          }}
+          addingHighlightedCode={addingHighlightedCode}
         />
         <ContinueButton onClick={onMainTextInput} />
       </TopGUIDiv>