Merge pull request #155 from continuedev/newer-simpler-stream-algo

Newer simpler stream algo

13 files changed, 491 insertions, 332 deletions

diff --git a/continuedev/src/continuedev/core/autopilot.py b/continuedev/src/continuedev/core/autopilot.py
index b9308409..3b2b65db 100644
--- a/continuedev/src/continuedev/core/autopilot.py
+++ b/continuedev/src/continuedev/core/autopilot.py
@@ -141,7 +141,7 @@ class Autopilot(ContinueBaseModel):
         # If a parent step is deleted/cancelled, don't run this step
         last_depth = self._step_depth
         i = self.history.current_index
-        while i >= 0 and self.history.timeline[i].depth == last_depth + 1:
+        while i >= 0 and self.history.timeline[i].depth == last_depth - 1:
             if self.history.timeline[i].deleted:
                 return None
             last_depth = self.history.timeline[i].depth
diff --git a/continuedev/src/continuedev/core/sdk.py b/continuedev/src/continuedev/core/sdk.py
index 62361250..d929a612 100644
--- a/continuedev/src/continuedev/core/sdk.py
+++ b/continuedev/src/continuedev/core/sdk.py
@@ -210,3 +210,6 @@ class ContinueSDK(AbstractContinueSDK):
 
     async def clear_history(self):
         await self.__autopilot.clear_history()
+
+    def current_step_was_deleted(self):
+        return self.history.timeline[self.history.current_index].deleted
diff --git a/continuedev/src/continuedev/server/ide.py b/continuedev/src/continuedev/server/ide.py
index c83fbc8a..c2ebaccf 100644
--- a/continuedev/src/continuedev/server/ide.py
+++ b/continuedev/src/continuedev/server/ide.py
@@ -146,8 +146,10 @@ class IdeProtocolServer(AbstractIdeProtocolServer):
 
     # ------------------------------- #
     # Request actions in IDE, doesn't matter which Session
-    def showSuggestion():
-        pass
+    async def showSuggestion(self, file_edit: FileEdit):
+        await self._send_json("showSuggestion", {
+            "edit": file_edit.dict()
+        })
 
     async def setFileOpen(self, filepath: str, open: bool = True):
         # Autopilot needs access to this.
diff --git a/continuedev/src/continuedev/server/ide_protocol.py b/continuedev/src/continuedev/server/ide_protocol.py
index 2dcedc30..79820c36 100644
--- a/continuedev/src/continuedev/server/ide_protocol.py
+++ b/continuedev/src/continuedev/server/ide_protocol.py
@@ -12,7 +12,7 @@ class AbstractIdeProtocolServer(ABC):
         """Handle a json message"""
 
     @abstractmethod
-    def showSuggestion():
+    def showSuggestion(self, file_edit: FileEdit):
         """Show a suggestion to the user"""
 
     @abstractmethod
diff --git a/continuedev/src/continuedev/steps/chat.py b/continuedev/src/continuedev/steps/chat.py
index 54d9c657..50e0f905 100644
--- a/continuedev/src/continuedev/steps/chat.py
+++ b/continuedev/src/continuedev/steps/chat.py
@@ -47,9 +47,13 @@ class AddFileStep(Step):
         except FileNotFoundError:
             self.description = f"File {self.filename} does not exist."
             return
-        currently_open_file = (await sdk.ide.getOpenFiles())[0]
+
         await sdk.ide.setFileOpen(os.path.join(sdk.ide.workspace_directory, self.filename))
-        await sdk.ide.setFileOpen(currently_open_file)
+
+        open_files = await sdk.ide.getOpenFiles()
+        if len(open_files) > 0:
+            currently_open_file = (await sdk.ide.getOpenFiles())[0]
+            await sdk.ide.setFileOpen(currently_open_file)
 
 
 class DeleteFileStep(Step):
diff --git a/continuedev/src/continuedev/steps/core/core.py b/continuedev/src/continuedev/steps/core/core.py
index 0d82b228..9545e9c7 100644
--- a/continuedev/src/continuedev/steps/core/core.py
+++ b/continuedev/src/continuedev/steps/core/core.py
@@ -2,14 +2,14 @@
 import os
 import subprocess
 from textwrap import dedent
-from typing import Coroutine, List, Union
+from typing import Coroutine, List, Literal, Union
 
 from ...models.main import Range
 from ...libs.llm.prompt_utils import MarkdownStyleEncoderDecoder
 from ...models.filesystem_edit import EditDiff, FileEdit, FileEditWithFullContents, FileSystemEdit
 from ...models.filesystem import FileSystem, RangeInFile, RangeInFileWithContents
 from ...core.observation import Observation, TextObservation, TracebackObservation, UserInputObservation
-from ...core.main import Step, SequentialStep
+from ...core.main import ChatMessage, Step, SequentialStep
 from ...libs.util.count_tokens import MAX_TOKENS_FOR_MODEL, DEFAULT_MAX_TOKENS
 import difflib
 
@@ -159,6 +159,238 @@ class DefaultModelEditCodeStep(Step):
         self.name = await models.gpt35.complete(f"Write a very short title to describe this requested change: '{self.user_input}'. This is the title:")
         return f"`{self.user_input}`\n\n" + description
 
+    async def get_prompt_parts(self, rif: RangeInFileWithContents, sdk: ContinueSDK, full_file_contents: str):
+        # If using 3.5 and overflows, upgrade to 3.5.16k
+        model_to_use = sdk.models.default
+        if model_to_use.name == "gpt-3.5-turbo":
+            if sdk.models.gpt35.count_tokens(full_file_contents) > MAX_TOKENS_FOR_MODEL["gpt-3.5-turbo"]:
+                model_to_use = sdk.models.gpt3516k
+
+        # Remove tokens from the end first, and then the start to clear space
+        # This part finds the start and end lines
+        full_file_contents_lst = full_file_contents.split("\n")
+        max_start_line = rif.range.start.line
+        min_end_line = rif.range.end.line
+        cur_start_line = 0
+        cur_end_line = len(full_file_contents_lst) - 1
+
+        total_tokens = model_to_use.count_tokens(
+            full_file_contents + self._prompt)
+
+        if total_tokens > MAX_TOKENS_FOR_MODEL[model_to_use.name]:
+            while cur_end_line > min_end_line:
+                total_tokens -= model_to_use.count_tokens(
+                    full_file_contents_lst[cur_end_line])
+                cur_end_line -= 1
+                if total_tokens < MAX_TOKENS_FOR_MODEL[model_to_use.name]:
+                    return cur_start_line, cur_end_line
+
+            if total_tokens > MAX_TOKENS_FOR_MODEL[model_to_use.name]:
+                while cur_start_line < max_start_line:
+                    cur_start_line += 1
+                    total_tokens -= model_to_use.count_tokens(
+                        full_file_contents_lst[cur_end_line])
+                    if total_tokens < MAX_TOKENS_FOR_MODEL[model_to_use.name]:
+                        return cur_start_line, cur_end_line
+
+        # Now use the found start/end lines to get the prefix and suffix strings
+        file_prefix = "\n".join(
+            full_file_contents_lst[cur_start_line:max_start_line])
+        file_suffix = "\n".join(
+            full_file_contents_lst[min_end_line:cur_end_line - 1])
+
+        # Move any surrounding blank line in rif.contents to the prefix/suffix
+        # TODO: Keep track of start line of the range, because it's needed below for offset stuff
+        rif_start_line = rif.range.start.line
+        if len(rif.contents) > 0:
+            first_line = rif.contents.splitlines(keepends=True)[0]
+            while first_line.strip() == "":
+                file_prefix += first_line
+                rif.contents = rif.contents[len(first_line):]
+                first_line = rif.contents.splitlines(keepends=True)[0]
+
+            last_line = rif.contents.splitlines(keepends=True)[-1]
+            while last_line.strip() == "":
+                file_suffix = last_line + file_suffix
+                rif.contents = rif.contents[:len(
+                    rif.contents) - len(last_line)]
+                last_line = rif.contents.splitlines(keepends=True)[-1]
+
+            while rif.contents.startswith("\n"):
+                file_prefix += "\n"
+                rif.contents = rif.contents[1:]
+            while rif.contents.endswith("\n"):
+                file_suffix = "\n" + file_suffix
+                rif.contents = rif.contents[:-1]
+
+        return file_prefix, rif.contents, file_suffix, model_to_use
+
+    def compile_prompt(self, file_prefix: str, contents: str, file_suffix: str, sdk: ContinueSDK) -> str:
+        prompt = self._prompt
+        if file_prefix.strip() != "":
+            prompt += dedent(f"""
+<file_prefix>
+{file_prefix}
+</file_prefix>""")
+        prompt += dedent(f"""
+<code_to_edit>
+{contents}
+</code_to_edit>""")
+        if file_suffix.strip() != "":
+            prompt += dedent(f"""
+<file_suffix>
+{file_suffix}
+</file_suffix>""")
+        prompt += dedent(f"""
+<user_request>
+{self.user_input}
+</user_request>
+<modified_code_to_edit>
+""")
+
+        return prompt
+
+    def is_end_line(self, line: str) -> bool:
+        return "</modified_code_to_edit>" in line
+
+    def line_to_be_ignored(self, line: str) -> bool:
+        return "```" in line or "<modified_code_to_edit>" in line or "<file_prefix>" in line or "</file_prefix>" in line or "<file_suffix>" in line or "</file_suffix>" in line or "<user_request>" in line or "</user_request>" in line or "<code_to_edit>" in line or "</code_to_edit>" in line
+
+    async def stream_rif(self, rif: RangeInFileWithContents, sdk: ContinueSDK):
+        full_file_contents = await sdk.ide.readFile(rif.filepath)
+
+        file_prefix, contents, file_suffix, model_to_use = await self.get_prompt_parts(
+            rif, sdk, full_file_contents)
+        prompt = self.compile_prompt(file_prefix, contents, file_suffix, sdk)
+
+        full_file_contents_lines = full_file_contents.split("\n")
+        original_lines = rif.contents.split("\n")
+        completion_lines_covered = 0
+        # In the actual file, as it is with blocks and such
+        current_line_in_file = rif.range.start.line
+
+        current_block_lines = []
+        original_lines_below_previous_blocks = original_lines
+        current_block_start = -1
+        offset_from_blocks = 0
+
+        lines_of_prefix_copied = 0
+        repeating_file_suffix = False
+        line_below_highlighted_range = file_suffix.lstrip().split("\n")[0]
+        lines = []
+        unfinished_line = ""
+
+        async def handle_generated_line(line: str):
+            nonlocal lines, current_block_start, current_line_in_file, original_lines, original_lines_below_previous_blocks, current_block_lines, offset_from_blocks
+
+            # Highlight the line to show progress
+            await sdk.ide.highlightCode(RangeInFile(filepath=rif.filepath, range=Range.from_shorthand(
+                current_line_in_file, 0, current_line_in_file, 0)), "#FFFFFF22" if len(current_block_lines) == 0 else "#FFFF0022")
+
+            if len(current_block_lines) == 0:
+                if len(original_lines_below_previous_blocks) == 0 or line != original_lines_below_previous_blocks[0]:
+                    current_block_lines.append(line)
+                    current_block_start = current_line_in_file
+
+                else:
+                    original_lines_below_previous_blocks = original_lines_below_previous_blocks[
+                        1:]
+                return
+
+            # We are in a block currently, and checking for whether it should be ended
+            for i in range(len(original_lines_below_previous_blocks)):
+                og_line = original_lines_below_previous_blocks[i]
+                if og_line == line and len(og_line.strip()):
+                    # Gather the lines to insert/replace for the suggestion
+                    lines_to_replace = current_block_lines[:i]
+                    original_lines_below_previous_blocks = original_lines_below_previous_blocks[
+                        i + 1:]
+
+                    # Insert the suggestion
+                    await sdk.ide.showSuggestion(FileEdit(
+                        filepath=rif.filepath,
+                        range=Range.from_shorthand(
+                            current_block_start, 0, current_block_start + i, 0),
+                        replacement="\n".join(current_block_lines)
+                    ))
+
+                    # Reset current block
+                    offset_from_blocks += len(current_block_lines)
+                    current_block_lines = []
+                    current_block_start = -1
+                    return
+
+            current_block_lines.append(line)
+
+        messages = await sdk.get_chat_context()
+        messages.append(ChatMessage(
+            role="user",
+            content=prompt,
+            summary=self.user_input
+        ))
+        async for chunk in model_to_use.stream_chat(messages, temperature=0):
+            # Stop early if it is repeating the file_suffix or the step was deleted
+            if repeating_file_suffix:
+                break
+            if sdk.current_step_was_deleted():
+                return
+
+            # Accumulate lines
+            if "content" not in chunk:
+                continue
+            chunk = chunk["content"]
+            chunk_lines = chunk.split("\n")
+            chunk_lines[0] = unfinished_line + chunk_lines[0]
+            if chunk.endswith("\n"):
+                unfinished_line = ""
+                chunk_lines.pop()  # because this will be an empty string
+            else:
+                unfinished_line = chunk_lines.pop()
+            lines.extend(chunk_lines)
+
+            # Deal with newly accumulated lines
+            for line in chunk_lines:
+                # Lines that should signify the end of generation
+                if self.is_end_line(line):
+                    break
+                # Lines that should be ignored, like the <> tags
+                elif self.line_to_be_ignored(line):
+                    continue
+                # Check if we are currently just copying the prefix
+                elif (lines_of_prefix_copied > 0 or completion_lines_covered == 0) and lines_of_prefix_copied < len(file_prefix.splitlines()) and line == full_file_contents_lines[lines_of_prefix_copied]:
+                    # This is a sketchy way of stopping it from repeating the file_prefix. Is a bug if output happens to have a matching line
+                    lines_of_prefix_copied += 1
+                    continue
+                # Because really short lines might be expected to be repeated, this is only a !heuristic!
+                # Stop when it starts copying the file_suffix
+                elif line.strip() == line_below_highlighted_range.strip() and len(line.strip()) > 4:
+                    repeating_file_suffix = True
+                    break
+
+                # If none of the above, insert the line!
+                await handle_generated_line(line)
+                completion_lines_covered += 1
+                current_line_in_file += 1
+
+        # Add the unfinished line
+        if unfinished_line != "" and not self.line_to_be_ignored(unfinished_line) and not self.is_end_line(unfinished_line):
+            lines.append(unfinished_line)
+            await handle_generated_line(unfinished_line)
+            completion_lines_covered += 1
+
+        # If the current block isn't empty, add that suggestion
+        if len(current_block_lines) > 0:
+            await sdk.ide.showSuggestion(FileEdit(
+                filepath=rif.filepath,
+                range=Range.from_shorthand(
+                    current_block_start, 0, current_block_start + len(original_lines_below_previous_blocks), 0),
+                replacement="\n".join(current_block_lines)
+            ))
+
+        # Record the completion
+        completion = "\n".join(lines)
+        self._prompt_and_completion += prompt + completion
+
     async def run(self, sdk: ContinueSDK) -> Coroutine[Observation, None, None]:
         self.description = f"{self.user_input}"
         await sdk.update_ui()
@@ -179,228 +411,8 @@ class DefaultModelEditCodeStep(Step):
 
         for rif in rif_with_contents:
             await sdk.ide.setFileOpen(rif.filepath)
-
-            model_to_use = sdk.models.default
-
-            full_file_contents = await sdk.ide.readFile(rif.filepath)
-
-            full_file_contents_lst = full_file_contents.split("\n")
-
-            max_start_line = rif.range.start.line
-            min_end_line = rif.range.end.line
-            cur_start_line = 0
-            cur_end_line = len(full_file_contents_lst) - 1
-
-            def cut_context(model_to_use, total_tokens, cur_start_line, cur_end_line):
-
-                if total_tokens > MAX_TOKENS_FOR_MODEL[model_to_use.name]:
-                    while cur_end_line > min_end_line:
-                        total_tokens -= model_to_use.count_tokens(
-                            full_file_contents_lst[cur_end_line])
-                        cur_end_line -= 1
-                        if total_tokens < MAX_TOKENS_FOR_MODEL[model_to_use.name]:
-                            return cur_start_line, cur_end_line
-
-                    if total_tokens > MAX_TOKENS_FOR_MODEL[model_to_use.name]:
-                        while cur_start_line < max_start_line:
-                            cur_start_line += 1
-                            total_tokens -= model_to_use.count_tokens(
-                                full_file_contents_lst[cur_end_line])
-                            if total_tokens < MAX_TOKENS_FOR_MODEL[model_to_use.name]:
-                                return cur_start_line, cur_end_line
-
-                return cur_start_line, cur_end_line
-
-            # We don't know here all of the functions being passed in.
-            # We care because if this prompt itself goes over the limit, then the entire message will have to be cut from the completion.
-            # Overflow won't happen, but prune_chat_messages in count_tokens.py will cut out this whole thing, instead of us cutting out only as many lines as we need.
-            BUFFER_FOR_FUNCTIONS = 200
-            total_tokens = model_to_use.count_tokens(
-                full_file_contents + self._prompt + self.user_input) + DEFAULT_MAX_TOKENS + BUFFER_FOR_FUNCTIONS
-
-            model_to_use = sdk.models.default
-            if model_to_use.name == "gpt-3.5-turbo":
-                if total_tokens > MAX_TOKENS_FOR_MODEL["gpt-3.5-turbo"]:
-                    model_to_use = sdk.models.gpt3516k
-
-            cur_start_line, cur_end_line = cut_context(
-                model_to_use, total_tokens, cur_start_line, cur_end_line)
-
-            code_before = "\n".join(
-                full_file_contents_lst[cur_start_line:max_start_line])
-            code_after = "\n".join(
-                full_file_contents_lst[min_end_line:cur_end_line - 1])
-
-            segs = [code_before, code_after]
-            if segs[0].strip() == "":
-                segs[0] = segs[0].strip()
-            if segs[1].strip() == "":
-                segs[1] = segs[1].strip()
-
-            # Move any surrounding blank line in rif.contents to the prefix/suffix
-            if len(rif.contents) > 0:
-                first_line = rif.contents.splitlines(keepends=True)[0]
-                while first_line.strip() == "":
-                    segs[0] += first_line
-                    rif.contents = rif.contents[len(first_line):]
-                    first_line = rif.contents.splitlines(keepends=True)[0]
-
-                last_line = rif.contents.splitlines(keepends=True)[-1]
-                while last_line.strip() == "":
-                    segs[1] = last_line + segs[1]
-                    rif.contents = rif.contents[:len(
-                        rif.contents) - len(last_line)]
-                    last_line = rif.contents.splitlines(keepends=True)[-1]
-
-                while rif.contents.startswith("\n"):
-                    segs[0] += "\n"
-                    rif.contents = rif.contents[1:]
-                while rif.contents.endswith("\n"):
-                    segs[1] = "\n" + segs[1]
-                    rif.contents = rif.contents[:-1]
-
-            # .format(code=rif.contents, user_request=self.user_input, file_prefix=segs[0], file_suffix=segs[1])
-            prompt = self._prompt
-            if segs[0].strip() != "":
-                prompt += dedent(f"""
-<file_prefix>
-{segs[0]}
-</file_prefix>""")
-            prompt += dedent(f"""
-<code_to_edit>
-{rif.contents}
-</code_to_edit>""")
-            if segs[1].strip() != "":
-                prompt += dedent(f"""
-<file_suffix>
-{segs[1]}
-</file_suffix>""")
-            prompt += dedent(f"""
-<user_request>
-{self.user_input}
-</user_request>
-<modified_code_to_edit>
-""")
-
-            lines = []
-            unfinished_line = ""
-            i = 0
-            original_lines = rif.contents.split("\n")
-
-            async def add_line(i: int, line: str):
-                if i == 0:
-                    # First line indentation, because the model will assume that it is replacing in this way
-                    line = original_lines[0].replace(
-                        original_lines[0].strip(), "") + line
-
-                if i < len(original_lines):
-                    # Replace original line
-                    range = Range.from_shorthand(
-                        rif.range.start.line + i, rif.range.start.character if i == 0 else 0, rif.range.start.line + i + 1, 0)
-                else:
-                    # Insert a line
-                    range = Range.from_shorthand(
-                        rif.range.start.line + i, 0, rif.range.start.line + i, 0)
-
-                await sdk.ide.applyFileSystemEdit(FileEdit(
-                    filepath=rif.filepath,
-                    range=range,
-                    replacement=line + "\n"
-                ))
-
-            lines_of_prefix_copied = 0
-            line_below_highlighted_range = segs[1].lstrip().split("\n")[0]
-            should_stop = False
-            async for chunk in model_to_use.stream_complete(prompt, with_history=await sdk.get_chat_context(), temperature=0):
-                if should_stop:
-                    break
-                chunk_lines = chunk.split("\n")
-                chunk_lines[0] = unfinished_line + chunk_lines[0]
-                if chunk.endswith("\n"):
-                    unfinished_line = ""
-                    chunk_lines.pop()  # because this will be an empty string
-                else:
-                    unfinished_line = chunk_lines.pop()
-                lines.extend(chunk_lines)
-
-                for line in chunk_lines:
-                    if "</modified_code_to_edit>" in line:
-                        break
-                    elif "```" in line or "<modified_code_to_edit>" in line or "<file_prefix>" in line or "</file_prefix>" in line or "<file_suffix>" in line or "</file_suffix>" in line or "<user_request>" in line or "</user_request>" in line or "<code_to_edit>" in line or "</code_to_edit>" in line:
-                        continue
-                    elif (lines_of_prefix_copied > 0 or i == 0) and lines_of_prefix_copied < len(segs[0].splitlines()) and line == full_file_contents_lst[lines_of_prefix_copied]:
-                        # This is a sketchy way of stopping it from repeating the file_prefix. Is a bug if output happens to have a matching line
-                        lines_of_prefix_copied += 1
-                        continue
-                    elif i < len(original_lines) and line == original_lines[i]:
-                        i += 1
-                        continue
-                    # Because really short lines might be expected to be repeated !heuristic!
-                    elif line.strip() == line_below_highlighted_range.strip() and len(line.strip()) > 4:
-                        should_stop = True
-                        break
-                    await add_line(i, line)
-                    i += 1
-
-            # Add the unfinished line
-            if unfinished_line != "":
-                unfinished_line = unfinished_line.replace(
-                    "</modified_code_to_edit>", "").replace("</code_to_edit>", "").replace("```", "").replace("</file_suffix>", "").replace("</file_prefix", "").replace(
-                    "<modified_code_to_edit>", "").replace("<code_to_edit>", "").replace("<file_suffix>", "").replace("<file_prefix", "")
-                if not i < len(original_lines) or not unfinished_line == original_lines[i]:
-                    await add_line(i, unfinished_line)
-                lines.append(unfinished_line)
-                i += 1
-
-            # Remove the leftover original lines
-            while i < len(original_lines):
-                range = Range.from_shorthand(
-                    rif.range.start.line + i, rif.range.start.character, rif.range.start.line + i, len(original_lines[i]) + 1)
-                await sdk.ide.applyFileSystemEdit(FileEdit(
-                    filepath=rif.filepath,
-                    range=range,
-                    replacement=""
-                ))
-                i += 1
-
-            completion = "\n".join(lines)
-
-            self._prompt_and_completion += prompt + completion
-
-            diff = list(difflib.ndiff(rif.contents.splitlines(
-                keepends=True), completion.splitlines(keepends=True)))
-
-            lines_to_highlight = set()
-            index = 0
-            for line in diff:
-                if line.startswith("-"):
-                    pass
-                elif line.startswith("+"):
-                    lines_to_highlight.add(index + rif.range.start.line)
-                    index += 1
-                elif line.startswith(" "):
-                    index += 1
-
-            current_hl_start = None
-            last_hl = None
-            rifs_to_highlight = []
-            for line in lines_to_highlight:
-                if current_hl_start is None:
-                    current_hl_start = line
-                elif line != last_hl + 1:
-                    rifs_to_highlight.append(RangeInFile(
-                        filepath=rif.filepath, range=Range.from_shorthand(current_hl_start, 0, last_hl, 0)))
-                    current_hl_start = line
-                last_hl = line
-
-            if current_hl_start is not None:
-                rifs_to_highlight.append(RangeInFile(
-                    filepath=rif.filepath, range=Range.from_shorthand(current_hl_start, 0, last_hl, 0)))
-
-            for rif_to_hl in rifs_to_highlight:
-                await sdk.ide.highlightCode(rif_to_hl)
-
-            await sdk.ide.saveFile(rif.filepath)
+            await self.stream_rif(rif, sdk)
+            # await sdk.ide.saveFile(rif.filepath)
 
 
 class EditFileStep(Step):
diff --git a/extension/package.json b/extension/package.json
index ae55a96b..ceba8698 100644
--- a/extension/package.json
+++ b/extension/package.json
@@ -65,12 +65,68 @@
         }
       }
     },
-    "commands": [],
+    "commands": [
+      {
+        "command": "continue.suggestionDown",
+        "category": "Continue",
+        "title": "Suggestion Down"
+      },
+      {
+        "command": "continue.suggestionUp",
+        "category": "Continue",
+        "title": "Suggestion Up"
+      },
+      {
+        "command": "continue.acceptSuggestion",
+        "category": "Continue",
+        "title": "Accept Suggestion"
+      },
+      {
+        "command": "continue.rejectSuggestion",
+        "category": "Continue",
+        "title": "Reject Suggestion"
+      },
+      {
+        "command": "continue.acceptAllSuggestions",
+        "category": "Continue",
+        "title": "Accept All Suggestions"
+      },
+      {
+        "command": "continue.rejectAllSuggestions",
+        "category": "Continue",
+        "title": "Reject All Suggestions"
+      }
+    ],
     "keybindings": [
       {
         "command": "continue.focusContinueInput",
         "mac": "cmd+k",
         "key": "ctrl+k"
+      },
+      {
+        "command": "continue.suggestionDown",
+        "mac": "shift+ctrl+down",
+        "key": "shift+ctrl+down"
+      },
+      {
+        "command": "continue.suggestionUp",
+        "mac": "shift+ctrl+up",
+        "key": "shift+ctrl+up"
+      },
+      {
+        "command": "continue.acceptSuggestion",
+        "mac": "shift+ctrl+enter",
+        "key": "shift+ctrl+enter"
+      },
+      {
+        "command": "continue.acceptAllSuggestions",
+        "mac": "shift+cmd+enter",
+        "key": "shift+ctrl+enter"
+      },
+      {
+        "command": "continue.rejectAllSuggestions",
+        "mac": "shift+cmd+backspace",
+        "key": "shift+ctrl+backspace"
       }
     ],
     "menus": {
diff --git a/extension/src/activation/activate.ts b/extension/src/activation/activate.ts
index 0c92f095..df8b6871 100644
--- a/extension/src/activation/activate.ts
+++ b/extension/src/activation/activate.ts
@@ -24,8 +24,19 @@ export async function activateExtension(
   registerAllCodeLensProviders(context);
   registerAllCommands(context);
 
-  // vscode.window.registerWebviewViewProvider("continue.continueGUIView", setupDebugPanel);
-  await startContinuePythonServer();
+  await new Promise((resolve, reject) => {
+    vscode.window.withProgress(
+      {
+        location: vscode.ProgressLocation.Notification,
+        title: "Starting Continue Server...",
+        cancellable: false,
+      },
+      async (progress, token) => {
+        await startContinuePythonServer();
+        resolve(null);
+      }
+    );
+  });
   const serverUrl = getContinueServerUrl();
 
   ideProtocolClient = new IdeProtocolClient(
diff --git a/extension/src/commands.ts b/extension/src/commands.ts
index 77273343..8072353b 100644
--- a/extension/src/commands.ts
+++ b/extension/src/commands.ts
@@ -9,6 +9,8 @@ import {
   rejectSuggestionCommand,
   suggestionDownCommand,
   suggestionUpCommand,
+  acceptAllSuggestionsCommand,
+  rejectAllSuggestionsCommand,
 } from "./suggestions";
 import * as bridge from "./bridge";
 import { debugPanelWebview } from "./debugPanel";
@@ -49,6 +51,8 @@ const commandsMap: { [command: string]: (...args: any) => any } = {
   "continue.suggestionUp": suggestionUpCommand,
   "continue.acceptSuggestion": acceptSuggestionCommand,
   "continue.rejectSuggestion": rejectSuggestionCommand,
+  "continue.acceptAllSuggestions": acceptAllSuggestionsCommand,
+  "continue.rejectAllSuggestions": rejectAllSuggestionsCommand,
   "continue.focusContinueInput": async () => {
     vscode.commands.executeCommand("continue.continueGUIView.focus");
     debugPanelWebview?.postMessage({
diff --git a/extension/src/continueIdeClient.ts b/extension/src/continueIdeClient.ts
index 08a0b74d..8ab3e075 100644
--- a/extension/src/continueIdeClient.ts
+++ b/extension/src/continueIdeClient.ts
@@ -1,5 +1,8 @@
 // import { ShowSuggestionRequest } from "../schema/ShowSuggestionRequest";
-import { showSuggestion, SuggestionRanges } from "./suggestions";
+import {
+  showSuggestion as showSuggestionInEditor,
+  SuggestionRanges,
+} from "./suggestions";
 import { openEditorAndRevealRange, getRightViewColumn } from "./util/vscode";
 import { FileEdit } from "../schema/FileEdit";
 import { RangeInFile } from "../schema/RangeInFile";
@@ -114,7 +117,10 @@ class IdeProtocolClient {
         break;
       case "setFileOpen":
         this.openFile(data.filepath);
-        // TODO: Close file
+        // TODO: Close file if False
+        break;
+      case "showSuggestion":
+        this.showSuggestion(data.edit);
         break;
       case "openGUI":
       case "connected":
@@ -138,6 +144,7 @@ class IdeProtocolClient {
   // ------------------------------------ //
   // On message handlers
 
+  private _lastDecorationType: vscode.TextEditorDecorationType | null = null;
   async highlightCode(rangeInFile: RangeInFile, color: string) {
     const range = new vscode.Range(
       rangeInFile.range.start.line,
@@ -157,24 +164,25 @@ class IdeProtocolClient {
       });
       editor.setDecorations(decorationType, [range]);
 
-      // Listen for changes to cursor position and then remove the decoration (but keep for at least 2 seconds)
-      const allowRemoveHighlight = () => {
-        const cursorDisposable = vscode.window.onDidChangeTextEditorSelection(
-          (event) => {
-            if (event.textEditor.document.uri.fsPath === rangeInFile.filepath) {
-              cursorDisposable.dispose();
-              editor.setDecorations(decorationType, []);
-            }
+      const cursorDisposable = vscode.window.onDidChangeTextEditorSelection(
+        (event) => {
+          if (event.textEditor.document.uri.fsPath === rangeInFile.filepath) {
+            cursorDisposable.dispose();
+            editor.setDecorations(decorationType, []);
           }
-        );
-      };
-      setTimeout(allowRemoveHighlight, 2000);
+        }
+      );
+
+      if (this._lastDecorationType) {
+        editor.setDecorations(this._lastDecorationType, []);
+      }
+      this._lastDecorationType = decorationType;
     }
   }
 
   showSuggestion(edit: FileEdit) {
     // showSuggestion already exists
-    showSuggestion(
+    showSuggestionInEditor(
       edit.filepath,
       new vscode.Range(
         edit.range.start.line,
diff --git a/extension/src/lang-server/codeLens.ts b/extension/src/lang-server/codeLens.ts
index 26528d96..1f352797 100644
--- a/extension/src/lang-server/codeLens.ts
+++ b/extension/src/lang-server/codeLens.ts
@@ -12,21 +12,25 @@ class SuggestionsCodeLensProvider implements vscode.CodeLensProvider {
     }
 
     let codeLenses: vscode.CodeLens[] = [];
-    for (let suggestion of suggestions) {
-      let range = new vscode.Range(
+    for (const suggestion of suggestions) {
+      const range = new vscode.Range(
         suggestion.oldRange.start,
         suggestion.newRange.end
       );
       codeLenses.push(
         new vscode.CodeLens(range, {
-          title: "Accept",
+          title: "Accept ✅",
           command: "continue.acceptSuggestion",
           arguments: [suggestion],
         }),
         new vscode.CodeLens(range, {
-          title: "Reject",
+          title: "Reject ❌",
           command: "continue.rejectSuggestion",
           arguments: [suggestion],
+        }),
+        new vscode.CodeLens(range, {
+          title: "(⌘⇧↩/⌘⇧⌫ to accept/reject all)",
+          command: "",
         })
       );
     }
@@ -53,12 +57,13 @@ class SuggestionsCodeLensProvider implements vscode.CodeLensProvider {
 
 const allCodeLensProviders: { [langauge: string]: vscode.CodeLensProvider[] } =
   {
-    python: [new SuggestionsCodeLensProvider()],
+    // python: [new SuggestionsCodeLensProvider(), new PytestCodeLensProvider()],
+    "*": [new SuggestionsCodeLensProvider()],
   };
 
 export function registerAllCodeLensProviders(context: vscode.ExtensionContext) {
-  for (let language in allCodeLensProviders) {
-    for (let codeLensProvider of allCodeLensProviders[language]) {
+  for (const language in allCodeLensProviders) {
+    for (const codeLensProvider of allCodeLensProviders[language]) {
       context.subscriptions.push(
         vscode.languages.registerCodeLensProvider(language, codeLensProvider)
       );
diff --git a/extension/src/suggestions.ts b/extension/src/suggestions.ts
index c66fad86..209bf8b2 100644
--- a/extension/src/suggestions.ts
+++ b/extension/src/suggestions.ts
@@ -14,7 +14,7 @@ export const editorToSuggestions: Map<
   string, // URI of file
   SuggestionRanges[]
 > = new Map();
-export let currentSuggestion: Map<string, number> = new Map(); // Map from editor URI to index of current SuggestionRanges in editorToSuggestions
+export const currentSuggestion: Map<string, number> = new Map(); // Map from editor URI to index of current SuggestionRanges in editorToSuggestions
 
 // When tab is reopened, rerender the decorations:
 vscode.window.onDidChangeActiveTextEditor((editor) => {
@@ -25,16 +25,16 @@ vscode.workspace.onDidOpenTextDocument((doc) => {
   rerenderDecorations(doc.uri.toString());
 });
 
-let newDecorationType = vscode.window.createTextEditorDecorationType({
+const newDecorationType = vscode.window.createTextEditorDecorationType({
   backgroundColor: "rgb(0, 255, 0, 0.1)",
   isWholeLine: true,
 });
-let oldDecorationType = vscode.window.createTextEditorDecorationType({
+const oldDecorationType = vscode.window.createTextEditorDecorationType({
   backgroundColor: "rgb(255, 0, 0, 0.1)",
   isWholeLine: true,
   cursor: "pointer",
 });
-let newSelDecorationType = vscode.window.createTextEditorDecorationType({
+const newSelDecorationType = vscode.window.createTextEditorDecorationType({
   backgroundColor: "rgb(0, 255, 0, 0.25)",
   isWholeLine: true,
   after: {
@@ -42,7 +42,7 @@ let newSelDecorationType = vscode.window.createTextEditorDecorationType({
     margin: "0 0 0 1em",
   },
 });
-let oldSelDecorationType = vscode.window.createTextEditorDecorationType({
+const oldSelDecorationType = vscode.window.createTextEditorDecorationType({
   backgroundColor: "rgb(255, 0, 0, 0.25)",
   isWholeLine: true,
   after: {
@@ -52,19 +52,44 @@ let oldSelDecorationType = vscode.window.createTextEditorDecorationType({
 });
 
 export function rerenderDecorations(editorUri: string) {
-  let suggestions = editorToSuggestions.get(editorUri);
-  let idx = currentSuggestion.get(editorUri);
-  let editor = vscode.window.visibleTextEditors.find(
+  const suggestions = editorToSuggestions.get(editorUri);
+  const idx = currentSuggestion.get(editorUri);
+  const editor = vscode.window.visibleTextEditors.find(
     (editor) => editor.document.uri.toString() === editorUri
   );
   if (!suggestions || !editor) return;
 
-  let olds: vscode.Range[] = [],
-    news: vscode.Range[] = [],
-    oldSels: vscode.Range[] = [],
-    newSels: vscode.Range[] = [];
+  const rangesWithoutEmptyLastLine = (ranges: vscode.Range[]) => {
+    const newRanges: vscode.Range[] = [];
+    for (let i = 0; i < ranges.length; i++) {
+      const range = ranges[i];
+      if (
+        range.start.line === range.end.line &&
+        range.start.character === 0 &&
+        range.end.character === 0
+      ) {
+        // Empty range, don't show it
+        continue;
+      }
+      newRanges.push(
+        new vscode.Range(
+          range.start.line,
+          range.start.character,
+          // Don't include the last line if it is empty
+          range.end.line - (range.end.character === 0 ? 1 : 0),
+          range.end.character
+        )
+      );
+    }
+    return newRanges;
+  };
+
+  let olds: vscode.Range[] = [];
+  let news: vscode.Range[] = [];
+  let oldSels: vscode.Range[] = [];
+  let newSels: vscode.Range[] = [];
   for (let i = 0; i < suggestions.length; i++) {
-    let suggestion = suggestions[i];
+    const suggestion = suggestions[i];
     if (typeof idx != "undefined" && idx === i) {
       if (suggestion.newSelected) {
         olds.push(suggestion.oldRange);
@@ -78,6 +103,13 @@ export function rerenderDecorations(editorUri: string) {
       news.push(suggestion.newRange);
     }
   }
+
+  // Don't highlight the last line if it is empty
+  olds = rangesWithoutEmptyLastLine(olds);
+  news = rangesWithoutEmptyLastLine(news);
+  oldSels = rangesWithoutEmptyLastLine(oldSels);
+  newSels = rangesWithoutEmptyLastLine(newSels);
+
   editor.setDecorations(oldDecorationType, olds);
   editor.setDecorations(newDecorationType, news);
   editor.setDecorations(oldSelDecorationType, oldSels);
@@ -92,14 +124,14 @@ export function rerenderDecorations(editorUri: string) {
 }
 
 export function suggestionDownCommand() {
-  let editor = vscode.window.activeTextEditor;
+  const editor = vscode.window.activeTextEditor;
   if (!editor) return;
-  let editorUri = editor.document.uri.toString();
-  let suggestions = editorToSuggestions.get(editorUri);
-  let idx = currentSuggestion.get(editorUri);
+  const editorUri = editor.document.uri.toString();
+  const suggestions = editorToSuggestions.get(editorUri);
+  const idx = currentSuggestion.get(editorUri);
   if (!suggestions || idx === undefined) return;
 
-  let suggestion = suggestions[idx];
+  const suggestion = suggestions[idx];
   if (!suggestion.newSelected) {
     suggestion.newSelected = true;
   } else if (idx + 1 < suggestions.length) {
@@ -109,14 +141,14 @@ export function suggestionDownCommand() {
 }
 
 export function suggestionUpCommand() {
-  let editor = vscode.window.activeTextEditor;
+  const editor = vscode.window.activeTextEditor;
   if (!editor) return;
-  let editorUri = editor.document.uri.toString();
-  let suggestions = editorToSuggestions.get(editorUri);
-  let idx = currentSuggestion.get(editorUri);
+  const editorUri = editor.document.uri.toString();
+  const suggestions = editorToSuggestions.get(editorUri);
+  const idx = currentSuggestion.get(editorUri);
   if (!suggestions || idx === undefined) return;
 
-  let suggestion = suggestions[idx];
+  const suggestion = suggestions[idx];
   if (suggestion.newSelected) {
     suggestion.newSelected = false;
   } else if (idx > 0) {
@@ -130,10 +162,10 @@ function selectSuggestion(
   accept: SuggestionSelectionOption,
   key: SuggestionRanges | null = null
 ) {
-  let editor = vscode.window.activeTextEditor;
+  const editor = vscode.window.activeTextEditor;
   if (!editor) return;
-  let editorUri = editor.document.uri.toString();
-  let suggestions = editorToSuggestions.get(editorUri);
+  const editorUri = editor.document.uri.toString();
+  const suggestions = editorToSuggestions.get(editorUri);
 
   if (!suggestions) return;
 
@@ -174,7 +206,7 @@ function selectSuggestion(
 
   rangeToDelete = new vscode.Range(
     rangeToDelete.start,
-    new vscode.Position(rangeToDelete.end.line + 1, 0)
+    new vscode.Position(rangeToDelete.end.line, 0)
   );
   editor.edit((edit) => {
     edit.delete(rangeToDelete);
@@ -206,6 +238,26 @@ export function acceptSuggestionCommand(key: SuggestionRanges | null = null) {
   selectSuggestion("selected", key);
 }
 
+function handleAllSuggestions(accept: boolean) {
+  const editor = vscode.window.activeTextEditor;
+  if (!editor) return;
+  const editorUri = editor.document.uri.toString();
+  const suggestions = editorToSuggestions.get(editorUri);
+  if (!suggestions) return;
+
+  while (suggestions.length > 0) {
+    selectSuggestion(accept ? "new" : "old", suggestions[0]);
+  }
+}
+
+export function acceptAllSuggestionsCommand() {
+  handleAllSuggestions(true);
+}
+
+export function rejectAllSuggestionsCommand() {
+  handleAllSuggestions(false);
+}
+
 export async function rejectSuggestionCommand(
   key: SuggestionRanges | null = null
 ) {
@@ -218,62 +270,62 @@ export async function showSuggestion(
   range: vscode.Range,
   suggestion: string
 ): Promise<boolean> {
-  let existingCode = await readFileAtRange(
-    new vscode.Range(range.start, range.end),
-    editorFilename
-  );
+  // const existingCode = await readFileAtRange(
+  //   new vscode.Range(range.start, range.end),
+  //   editorFilename
+  // );
 
   // If any of the outside lines are the same, don't repeat them in the suggestion
-  let slines = suggestion.split("\n");
-  let elines = existingCode.split("\n");
-  let linesRemovedBefore = 0;
-  let linesRemovedAfter = 0;
-  while (slines.length > 0 && elines.length > 0 && slines[0] === elines[0]) {
-    slines.shift();
-    elines.shift();
-    linesRemovedBefore++;
-  }
+  // const slines = suggestion.split("\n");
+  // const elines = existingCode.split("\n");
+  // let linesRemovedBefore = 0;
+  // let linesRemovedAfter = 0;
+  // while (slines.length > 0 && elines.length > 0 && slines[0] === elines[0]) {
+  //   slines.shift();
+  //   elines.shift();
+  //   linesRemovedBefore++;
+  // }
 
-  while (
-    slines.length > 0 &&
-    elines.length > 0 &&
-    slines[slines.length - 1] === elines[elines.length - 1]
-  ) {
-    slines.pop();
-    elines.pop();
-    linesRemovedAfter++;
-  }
+  // while (
+  //   slines.length > 0 &&
+  //   elines.length > 0 &&
+  //   slines[slines.length - 1] === elines[elines.length - 1]
+  // ) {
+  //   slines.pop();
+  //   elines.pop();
+  //   linesRemovedAfter++;
+  // }
 
-  suggestion = slines.join("\n");
-  if (suggestion === "") return Promise.resolve(false); // Don't even make a suggestion if they are exactly the same
+  // suggestion = slines.join("\n");
+  // if (suggestion === "") return Promise.resolve(false); // Don't even make a suggestion if they are exactly the same
 
-  range = new vscode.Range(
-    new vscode.Position(range.start.line + linesRemovedBefore, 0),
-    new vscode.Position(
-      range.end.line - linesRemovedAfter,
-      elines.at(-1)?.length || 0
-    )
-  );
+  // range = new vscode.Range(
+  //   new vscode.Position(range.start.line + linesRemovedBefore, 0),
+  //   new vscode.Position(
+  //     range.end.line - linesRemovedAfter,
+  //     elines.at(-1)?.length || 0
+  //   )
+  // );
 
-  let editor = await openEditorAndRevealRange(editorFilename, range);
+  const editor = await openEditorAndRevealRange(editorFilename, range);
   if (!editor) return Promise.resolve(false);
 
   return new Promise((resolve, reject) => {
     editor!
-      .edit((edit) => {
-        if (range.end.line + 1 >= editor.document.lineCount) {
-          suggestion = "\n" + suggestion;
-        }
-        edit.insert(
-          new vscode.Position(range.end.line + 1, 0),
-          suggestion + "\n"
-        );
-      })
+      .edit(
+        (edit) => {
+          edit.insert(
+            new vscode.Position(range.end.line, 0),
+            suggestion + "\n"
+          );
+        },
+        { undoStopBefore: false, undoStopAfter: false }
+      )
       .then(
         (success) => {
           if (success) {
             let suggestionRange = new vscode.Range(
-              new vscode.Position(range.end.line + 1, 0),
+              new vscode.Position(range.end.line, 0),
               new vscode.Position(
                 range.end.line + suggestion.split("\n").length,
                 0
diff --git a/extension/src/util/vscode.ts b/extension/src/util/vscode.ts
index a76b53c7..3110d589 100644
--- a/extension/src/util/vscode.ts
+++ b/extension/src/util/vscode.ts
@@ -118,9 +118,11 @@ export async function readFileAtRange(
             )
           );
         } else {
-          let firstLine = lines[range.start.line].slice(range.start.character);
-          let lastLine = lines[range.end.line].slice(0, range.end.character);
-          let middleLines = lines.slice(range.start.line + 1, range.end.line);
+          const firstLine = lines[range.start.line].slice(
+            range.start.character
+          );
+          const lastLine = lines[range.end.line].slice(0, range.end.character);
+          const middleLines = lines.slice(range.start.line + 1, range.end.line);
           resolve([firstLine, ...middleLines, lastLine].join("\n"));
         }
       }
@@ -144,7 +146,7 @@ export function openEditorAndRevealRange(
             setInterval(() => {
               resolve(null);
             }, 200);
-          })
+          });
         }
         showTextDocumentInProcess = true;
         vscode.window
@@ -158,10 +160,10 @@ export function openEditorAndRevealRange(
             }
             resolve(editor);
             showTextDocumentInProcess = false;
-          })
-        } catch (err) {
-          console.log(err);
-        }
-      });
+          });
+      } catch (err) {
+        console.log(err);
+      }
+    });
   });
 }