3 files changed, 62 insertions, 0 deletions
diff --git a/extension/src/activation/activate.ts b/extension/src/activation/activate.ts
index cd885b12..5c6ffa02 100644
--- a/extension/src/activation/activate.ts
+++ b/extension/src/activation/activate.ts
@@ -10,6 +10,7 @@ import {
   startContinuePythonServer,
 } from "./environmentSetup";
 import fetch from "node-fetch";
+import registerQuickFixProvider from "../lang-server/codeActions";
 // import { CapturedTerminal } from "../terminal/terminalEmulator";
 
 const PACKAGE_JSON_RAW_GITHUB_URL =
@@ -55,6 +56,7 @@ export async function activateExtension(context: vscode.ExtensionContext) {
   sendTelemetryEvent(TelemetryEvent.ExtensionActivated);
   registerAllCodeLensProviders(context);
   registerAllCommands(context);
+  registerQuickFixProvider();
 
   // Initialize IDE Protocol Client
   const serverUrl = getContinueServerUrl();
diff --git a/extension/src/commands.ts b/extension/src/commands.ts
index 888f01ed..2b7f4c0c 100644
--- a/extension/src/commands.ts
+++ b/extension/src/commands.ts
@@ -34,6 +34,13 @@ const commandsMap: { [command: string]: (...args: any) => any } = {
   "continue.rejectDiff": rejectDiffCommand,
   "continue.acceptAllSuggestions": acceptAllSuggestionsCommand,
   "continue.rejectAllSuggestions": rejectAllSuggestionsCommand,
+  "continue.quickFix": async (message: string, code: string, edit: boolean) => {
+    ideProtocolClient.sendMainUserInput(
+      `${
+        edit ? "/edit " : ""
+      }${code}\n\nHow do I fix this problem in the above code?: ${message}`
+    );
+  },
   "continue.focusContinueInput": async () => {
     if (focusedOnContinueInput) {
       vscode.commands.executeCommand("workbench.action.focusActiveEditorGroup");
diff --git a/extension/src/lang-server/codeActions.ts b/extension/src/lang-server/codeActions.ts
new file mode 100644
index 00000000..07cf5f4e
--- /dev/null
+++ b/extension/src/lang-server/codeActions.ts
@@ -0,0 +1,53 @@
+import * as vscode from "vscode";
+
+class ContinueQuickFixProvider implements vscode.CodeActionProvider {
+  public static readonly providedCodeActionKinds = [
+    vscode.CodeActionKind.QuickFix,
+  ];
+
+  provideCodeActions(
+    document: vscode.TextDocument,
+    range: vscode.Range | vscode.Selection,
+    context: vscode.CodeActionContext,
+    token: vscode.CancellationToken
+  ): vscode.ProviderResult<(vscode.Command | vscode.CodeAction)[]> {
+    if (context.diagnostics.length === 0) {
+      return [];
+    }
+
+    const createQuickFix = (edit: boolean) => {
+      const diagnostic = context.diagnostics[0];
+      const quickFix = new vscode.CodeAction(
+        edit ? "Fix with Continue" : "Ask Continue",
+        vscode.CodeActionKind.QuickFix
+      );
+      quickFix.isPreferred = false;
+      const surroundingRange = new vscode.Range(
+        range.start.translate(-3, 0),
+        range.end.translate(3, 0)
+      );
+      quickFix.command = {
+        command: "continue.quickFix",
+        title: "Continue Quick Fix",
+        arguments: [
+          diagnostic.message,
+          document.getText(surroundingRange),
+          edit,
+        ],
+      };
+      return quickFix;
+    };
+    return [createQuickFix(true), createQuickFix(false)];
+  }
+}
+
+export default function registerQuickFixProvider() {
+  // In your extension's activate function:
+  vscode.languages.registerCodeActionsProvider(
+    { language: "*" },
+    new ContinueQuickFixProvider(),
+    {
+      providedCodeActionKinds: ContinueQuickFixProvider.providedCodeActionKinds,
+    }
+  );
+}