highlighted code improvements

1 files changed, 7 insertions, 0 deletions

diff --git a/continuedev/src/continuedev/server/gui.py b/continuedev/src/continuedev/server/gui.py
index 4e960f7c..fa573b37 100644
--- a/continuedev/src/continuedev/server/gui.py
+++ b/continuedev/src/continuedev/server/gui.py
@@ -85,6 +85,8 @@ class GUIProtocolServer(AbstractGUIProtocolServer):
                 self.on_delete_at_index(data["index"])
             elif message_type == "delete_context_at_indices":
                 self.on_delete_context_at_indices(data["indices"])
+            elif message_type == "toggle_adding_highlighted_code":
+                self.on_toggle_adding_highlighted_code()
         except Exception as e:
             print(e)
 
@@ -128,6 +130,11 @@ class GUIProtocolServer(AbstractGUIProtocolServer):
             self.session.autopilot.delete_context_at_indices(indices)
         )
 
+    def on_toggle_adding_highlighted_code(self):
+        asyncio.create_task(
+            self.session.autopilot.toggle_adding_highlighted_code()
+        )
+
 
 @router.websocket("/ws")
 async def websocket_endpoint(websocket: WebSocket, session: Session = Depends(websocket_session)):