2 files changed, 55 insertions, 20 deletions

diff --git a/extension/react-app/src/components/ComboBox.tsx b/extension/react-app/src/components/ComboBox.tsx
index 9b15713e..14e04a84 100644
--- a/extension/react-app/src/components/ComboBox.tsx
+++ b/extension/react-app/src/components/ComboBox.tsx
@@ -244,7 +244,9 @@ const ComboBox = React.forwardRef((props: ComboBoxProps, ref) => {
       (item) => item === document.activeElement
     );
     console.log(focusedItemIndex, focusableItems);
-    if (focusedItemIndex !== -1) {
+    if (focusedItemIndex === focusableItemsArray.length - 1) {
+      inputRef.current?.focus();
+    } else if (focusedItemIndex !== -1) {
       const nextItem =
         focusableItemsArray[
           (focusedItemIndex + 1) % focusableItemsArray.length
@@ -256,15 +258,27 @@ const ComboBox = React.forwardRef((props: ComboBoxProps, ref) => {
     }
   };
 
+  useEffect(() => {
+    if (typeof window !== "undefined") {
+      const listener = (e: any) => {
+        if (e.key === "Tab") {
+          e.preventDefault();
+          handleTabPressed();
+        }
+      };
+      window.addEventListener("keydown", listener);
+      return () => {
+        window.removeEventListener("keydown", listener);
+      };
+    }
+  }, []);
+
   const [metaKeyPressed, setMetaKeyPressed] = useState(false);
   const [focused, setFocused] = useState(false);
   useEffect(() => {
     const handleKeyDown = (e: KeyboardEvent) => {
       if (e.key === "Meta") {
         setMetaKeyPressed(true);
-      } else if (e.key === "Tab") {
-        e.preventDefault();
-        handleTabPressed();
       }
     };
     const handleKeyUp = (e: KeyboardEvent) => {
@@ -319,6 +333,10 @@ const ComboBox = React.forwardRef((props: ComboBoxProps, ref) => {
               }
               addingHighlightedCode={props.addingHighlightedCode}
               index={idx}
+              onDelete={() => {
+                client?.deleteContextWithIds([item.description.id]);
+                inputRef.current?.focus();
+              }}
             />
           );
         })}
@@ -337,7 +355,7 @@ const ComboBox = React.forwardRef((props: ComboBoxProps, ref) => {
               onClick={() => {
                 props.onToggleAddContext();
               }}
-              className="pill-button"
+              className="pill-button focus:outline-none focus:border-red-600 focus:border focus:border-solid"
               onKeyDown={(e: KeyboardEvent) => {
                 e.preventDefault();
                 if (e.key === "Enter") {
@@ -395,8 +413,6 @@ const ComboBox = React.forwardRef((props: ComboBoxProps, ref) => {
                 event.preventDefault();
               } else if (event.key === "Tab") {
                 (event.nativeEvent as any).preventDownshiftDefault = true;
-                event.preventDefault();
-                handleTabPressed();
               } else if (
                 (event.key === "ArrowUp" || event.key === "ArrowDown") &&
                 event.currentTarget.value.split("\n").length > 1
diff --git a/extension/react-app/src/components/PillButton.tsx b/extension/react-app/src/components/PillButton.tsx
index 5895a91c..0dcd43eb 100644
--- a/extension/react-app/src/components/PillButton.tsx
+++ b/extension/react-app/src/components/PillButton.tsx
@@ -84,8 +84,34 @@ interface PillButtonProps {
   index: number;
   addingHighlightedCode?: boolean;
   areMultipleItems?: boolean;
+  onDelete?: () => void;
 }
 
+interface StyledButtonProps {
+  warning: string;
+  editing?: boolean;
+  areMultipleItems?: boolean;
+}
+
+const StyledButton = styled(Button)<StyledButtonProps>`
+  position: relative;
+  border-color: ${(props) =>
+    props.warning
+      ? "red"
+      : props.editing && props.areMultipleItems
+      ? vscForeground
+      : "transparent"};
+  border-width: 1px;
+  border-style: solid;
+
+  &:focus {
+    outline: none;
+    border-color: red;
+    border-width: 1px;
+    border-style: solid;
+  }
+`;
+
 const PillButton = (props: PillButtonProps) => {
   const [isHovered, setIsHovered] = useState(false);
   const client = useContext(GUIClientContext);
@@ -132,17 +158,10 @@ const PillButton = (props: PillButtonProps) => {
   return (
     <>
       <div style={{ position: "relative" }}>
-        <Button
-          style={{
-            position: "relative",
-            borderColor: props.warning
-              ? "red"
-              : props.item.editing && props.areMultipleItems
-              ? vscForeground
-              : "transparent",
-            borderWidth: "1px",
-            borderStyle: "solid",
-          }}
+        <StyledButton
+          areMultipleItems={props.areMultipleItems}
+          warning={props.warning || ""}
+          editing={props.item.editing}
           onMouseEnter={() => {
             setIsHovered(true);
             if (props.onHover) {
@@ -158,7 +177,7 @@ const PillButton = (props: PillButtonProps) => {
           className="pill-button"
           onKeyDown={(e) => {
             if (e.key === "Backspace") {
-              client?.deleteContextWithIds([props.item.description.id]);
+              props.onDelete?.();
             }
           }}
         >
@@ -202,7 +221,7 @@ const PillButton = (props: PillButtonProps) => {
             </GridDiv>
           )}
           {props.item.description.name}
-        </Button>
+        </StyledButton>
         <StyledTooltip id={`edit-${props.index}`}>
           {props.item.editing
             ? "Editing this section (with entire file as context)"