explicit context pill buttons

3 files changed, 21 insertions, 2 deletions

diff --git a/continuedev/src/continuedev/server/gui.py b/continuedev/src/continuedev/server/gui.py
index c0178920..9a33fb6c 100644
--- a/continuedev/src/continuedev/server/gui.py
+++ b/continuedev/src/continuedev/server/gui.py
@@ -83,6 +83,8 @@ class GUIProtocolServer(AbstractGUIProtocolServer):
                 self.on_clear_history()
             elif message_type == "delete_at_index":
                 self.on_delete_at_index(data["index"])
+            elif message_type == "delete_context_item_at_index":
+                self.on_delete_context_item_at_index(data["index"])
         except Exception as e:
             print(e)
 
@@ -127,6 +129,11 @@ class GUIProtocolServer(AbstractGUIProtocolServer):
     def on_delete_at_index(self, index: int):
         asyncio.create_task(self.session.autopilot.delete_at_index(index))
 
+    def on_delete_context_item_at_index(self, index: int):
+        asyncio.create_task(
+            self.session.autopilot.delete_context_item_at_index(index)
+        )
+
 
 @router.websocket("/ws")
 async def websocket_endpoint(websocket: WebSocket, session: Session = Depends(websocket_session)):
diff --git a/continuedev/src/continuedev/server/ide.py b/continuedev/src/continuedev/server/ide.py
index e1f19447..f3deecdb 100644
--- a/continuedev/src/continuedev/server/ide.py
+++ b/continuedev/src/continuedev/server/ide.py
@@ -8,7 +8,7 @@ from fastapi import WebSocket, Body, APIRouter
 from uvicorn.main import Server
 
 from ..libs.util.queue import AsyncSubscriptionQueue
-from ..models.filesystem import FileSystem, RangeInFile, EditDiff, RealFileSystem
+from ..models.filesystem import FileSystem, RangeInFile, EditDiff, RangeInFileWithContents, RealFileSystem
 from ..models.filesystem_edit import AddDirectory, AddFile, DeleteDirectory, DeleteFile, FileSystemEdit, FileEdit, FileEditWithFullContents, RenameDirectory, RenameFile, SequentialFileSystemEdit
 from pydantic import BaseModel
 from .gui import SessionManager, session_manager
@@ -139,6 +139,9 @@ class IdeProtocolServer(AbstractIdeProtocolServer):
             fileEdits = list(
                 map(lambda d: FileEditWithFullContents.parse_obj(d), data["fileEdits"]))
             self.onFileEdits(fileEdits)
+        elif message_type == "highlightedCodePush":
+            self.onHighlightedCodeUpdate(
+                [RangeInFileWithContents(**rif) for rif in data["highlightedCode"]])
         elif message_type == "commandOutput":
             output = data["output"]
             self.onCommandOutput(output)
@@ -229,6 +232,11 @@ class IdeProtocolServer(AbstractIdeProtocolServer):
             asyncio.create_task(
                 session.autopilot.handle_command_output(output))
 
+    def onHighlightedCodeUpdate(self, range_in_files: List[RangeInFileWithContents]):
+        for _, session in self.session_manager.sessions.items():
+            asyncio.create_task(
+                session.autopilot.handle_highlighted_code(range_in_files))
+
     # Request information. Session doesn't matter.
     async def getOpenFiles(self) -> List[str]:
         resp = await self._send_and_receive_json({}, OpenFilesResponse, "openFiles")
diff --git a/continuedev/src/continuedev/server/ide_protocol.py b/continuedev/src/continuedev/server/ide_protocol.py
index d2dafa9a..17a09c3d 100644
--- a/continuedev/src/continuedev/server/ide_protocol.py
+++ b/continuedev/src/continuedev/server/ide_protocol.py
@@ -3,7 +3,7 @@ from abc import ABC, abstractmethod, abstractproperty
 
 from ..models.main import Traceback
 from ..models.filesystem_edit import FileEdit, FileSystemEdit, EditDiff
-from ..models.filesystem import RangeInFile
+from ..models.filesystem import RangeInFile, RangeInFileWithContents
 
 
 class AbstractIdeProtocolServer(ABC):
@@ -91,6 +91,10 @@ class AbstractIdeProtocolServer(ABC):
     async def runCommand(self, command: str) -> str:
         """Run a command"""
 
+    @abstractmethod
+    def onHighlightedCodeUpdate(self, range_in_files: List[RangeInFileWithContents]):
+        """Called when highlighted code is updated"""
+
     @abstractproperty
     def workspace_directory(self) -> str:
         """Get the workspace directory"""