4 files changed, 59 insertions, 24 deletions

diff --git a/extension/react-app/src/components/ComboBox.tsx b/extension/react-app/src/components/ComboBox.tsx
index 9aab4e93..da559383 100644
--- a/extension/react-app/src/components/ComboBox.tsx
+++ b/extension/react-app/src/components/ComboBox.tsx
@@ -162,33 +162,43 @@ const ComboBox = React.forwardRef((props: ComboBoxProps, ref) => {
       }
     },
     onInputValueChange({ inputValue, highlightedIndex }) {
-      if (!inputValue) return;
+      if (!inputValue) {
+        setItems([]);
+        return;
+      }
       props.onInputValueChange(inputValue);
 
       if (inputValue.endsWith("@") || currentlyInContextQuery) {
-        setCurrentlyInContextQuery(true);
-
         const segs = inputValue.split("@");
-        const providerAndQuery = segs[segs.length - 1];
-        const [provider, query] = providerAndQuery.split(" ");
-        searchClient
-          .index(SEARCH_INDEX_NAME)
-          .search(providerAndQuery)
-          .then((res) => {
-            setItems(
-              res.hits.map((hit) => {
-                return {
-                  name: hit.name,
-                  description: hit.description,
-                  id: hit.id,
-                };
-              })
-            );
-          })
-          .catch(() => {
-            // Swallow errors, because this simply is not supported on Windows at the moment
-          });
-        return;
+
+        if (segs.length > 1) {
+          // Get search results and return
+          setCurrentlyInContextQuery(true);
+          const providerAndQuery = segs[segs.length - 1];
+          const [provider, query] = providerAndQuery.split(" ");
+          searchClient
+            .index(SEARCH_INDEX_NAME)
+            .search(providerAndQuery)
+            .then((res) => {
+              setItems(
+                res.hits.map((hit) => {
+                  return {
+                    name: hit.name,
+                    description: hit.description,
+                    id: hit.id,
+                  };
+                })
+              );
+            })
+            .catch(() => {
+              // Swallow errors, because this simply is not supported on Windows at the moment
+            });
+          return;
+        } else {
+          // Exit the '@' context menu
+          setCurrentlyInContextQuery(false);
+          setItems;
+        }
       }
       setItems(
         props.items.filter((item) =>
@@ -262,7 +272,7 @@ const ComboBox = React.forwardRef((props: ComboBoxProps, ref) => {
               key={`${item.description.id.item_id}${idx}`}
               item={item}
               warning={
-                false && item.content.length > 4000 && item.editing
+                item.content.length > 4000 && item.editing
                   ? "Editing such a large range may be slow"
                   : undefined
               }
diff --git a/extension/react-app/src/hooks/ContinueGUIClientProtocol.ts b/extension/react-app/src/hooks/ContinueGUIClientProtocol.ts
index b8019664..5a5d4c30 100644
--- a/extension/react-app/src/hooks/ContinueGUIClientProtocol.ts
+++ b/extension/react-app/src/hooks/ContinueGUIClientProtocol.ts
@@ -17,6 +17,13 @@ class ContinueGUIClientProtocol extends AbstractContinueGUIClientProtocol {
     this.messenger = useVscodeMessagePassing
       ? new VscodeMessenger(serverUrlWithSessionId)
       : new WebsocketMessenger(serverUrlWithSessionId);
+
+    this.messenger.onClose(() => {
+      console.log("GUI -> IDE websocket closed");
+    });
+    this.messenger.onError((error) => {
+      console.log("GUI -> IDE websocket error", error);
+    });
   }
 
   sendMainInput(input: string) {
diff --git a/extension/react-app/src/hooks/messenger.ts b/extension/react-app/src/hooks/messenger.ts
index 00ce1fbb..ecf646c7 100644
--- a/extension/react-app/src/hooks/messenger.ts
+++ b/extension/react-app/src/hooks/messenger.ts
@@ -13,6 +13,8 @@ export abstract class Messenger {
   abstract onClose(callback: () => void): void;
 
   abstract sendAndReceive(messageType: string, data: any): Promise<any>;
+
+  abstract onError(callback: (error: any) => void): void;
 }
 
 export class WebsocketMessenger extends Messenger {
@@ -20,6 +22,7 @@ export class WebsocketMessenger extends Messenger {
   private onMessageListeners: {
     [messageType: string]: ((data: object) => void)[];
   } = {};
+  private onErrorListeners: ((error: any) => void)[] = [];
   private onOpenListeners: (() => void)[] = [];
   private onCloseListeners: (() => void)[] = [];
   private serverUrl: string;
@@ -37,6 +40,9 @@ export class WebsocketMessenger extends Messenger {
         this.onMessageType(messageType, listener);
       }
     }
+    for (const listener of this.onErrorListeners) {
+      this.onError(listener);
+    }
     return newWebsocket;
   }
 
@@ -95,4 +101,8 @@ export class WebsocketMessenger extends Messenger {
   onClose(callback: () => void): void {
     this.websocket.addEventListener("close", callback);
   }
+
+  onError(callback: (error: any) => void): void {
+    this.websocket.addEventListener("error", callback);
+  }
 }
diff --git a/extension/react-app/src/hooks/vscodeMessenger.ts b/extension/react-app/src/hooks/vscodeMessenger.ts
index ba19586b..13f5092b 100644
--- a/extension/react-app/src/hooks/vscodeMessenger.ts
+++ b/extension/react-app/src/hooks/vscodeMessenger.ts
@@ -38,6 +38,14 @@ export class VscodeMessenger extends Messenger {
     });
   }
 
+  onError(callback: (error: any) => void): void {
+    window.addEventListener("message", (event: any) => {
+      if (event.data.type === "websocketForwardingError") {
+        callback(event.data.error);
+      }
+    });
+  }
+
   sendAndReceive(messageType: string, data: any): Promise<any> {
     return new Promise((resolve) => {
       const handler = (event: any) => {