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

Newer simpler stream algo

7 files changed, 240 insertions, 102 deletions

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);
+      }
+    });
   });
 }