Support aborting code generation

crates/cursor-core/src/generate/mod.rs

@@ -1,9 +1,16 @@
 mod models;
 
+use std::future::IntoFuture;
+
 use crate::GenerateInput;
-use futures::StreamExt;
+use futures::{
+    future::{select, Either},
+    StreamExt,
+};
 use models::*;
+use node_bridge::futures::Defer;
 use node_bridge::http_client::{HttpMethod, HttpRequest};
+use node_bridge::prelude::*;
 use wasm_bindgen::prelude::*;
 
 // Split the code into chunks of 20 line blocks.
@@ -24,14 +31,14 @@ fn split_code_into_blocks(code: &str) -> Vec<String> {
     blocks
 }
 
-#[wasm_bindgen(js_name = generateCode)]
-pub async fn generate_code(input: &GenerateInput) -> Result<(), JsValue> {
+async fn generate_code_inner(input: &GenerateInput) -> Result<(), JsValue> {
     let file_path = input.file_path();
     let file_dir = file_path
         .split("/")
         .take(file_path.split("/").count() - 1)
         .collect::<Vec<&str>>()
         .join("/");
+    #[cfg(debug_assertions)]
     node_bridge::bindings::console::log_str(&format!("file_dir: {}", file_dir));
     let workspace_directory = input.workspace_directory();
     let selection = input.selection_range();
@@ -99,7 +106,8 @@ pub async fn generate_code(input: &GenerateInput) -> Result<(), JsValue> {
             request_body.bot_messages = vec![bot_message];
         }
 
-        node_bridge::bindings::console::log_str(&serde_json::to_string(&request_body).unwrap());
+        #[cfg(debug_assertions)]
+        console::log_str(&serde_json::to_string(&request_body).unwrap());
 
         let request = HttpRequest::new(&format!(
             "https://aicursor.com/{}",
@@ -121,7 +129,8 @@ pub async fn generate_code(input: &GenerateInput) -> Result<(), JsValue> {
         let body = response.body();
         while let Some(chunk) = body.next().await {
             let chunk = chunk.to_string("utf-8");
-            node_bridge::bindings::console::log_str(&chunk);
+            #[cfg(debug_assertions)]
+            console::log_str(&chunk);
             let lines = chunk.split("\n").filter(|l| l.len() > 0);
             let mut message_ended = false;
             for line in lines {
@@ -169,8 +178,33 @@ pub async fn generate_code(input: &GenerateInput) -> Result<(), JsValue> {
         response.await?;
     }
 
-    node_bridge::bindings::console::log_str("done");
+    node_bridge::bindings::console::log_str("generate done");
 
     result_stream.end();
     Ok(())
 }
+
+#[wasm_bindgen(js_name = generateCode)]
+pub async fn generate_code(input: &GenerateInput) -> Result<(), JsValue> {
+    let defer_abort = Defer::new();
+    let defer_abort_clone = defer_abort.clone();
+    let abort_signal = input.abort_signal();
+    abort_signal.add_event_listener(
+        "abort",
+        closure_once!(|| {
+            defer_abort_clone.resolve(JsValue::null());
+        })
+        .into_js_value(),
+    );
+
+    let fut = generate_code_inner(input);
+
+    match select(defer_abort.into_future(), Box::pin(fut)).await {
+        Either::Left(_) => {
+            return Ok(());
+        }
+        Either::Right((res, _)) => {
+            return res;
+        }
+    }
+}

crates/cursor-core/src/lib.rs

@@ -1,5 +1,6 @@
 mod generate;
 
+use node_bridge::bindings::AbortSignal;
 use wasm_bindgen::prelude::*;
 
 #[wasm_bindgen(typescript_custom_section)]
@@ -27,6 +28,7 @@ interface IGenerateInput {
     get workspaceDirectory(): string | null;
     get selectionRange(): ISelectionRange;
     get resultStream(): IResultStream;
+    get abortSignal(): AbortSignal;
 }
 "#;
 
@@ -76,4 +78,7 @@ extern "C" {
 
     #[wasm_bindgen(method, getter, structural, js_name = resultStream)]
     pub fn result_stream(this: &GenerateInput) -> ResultStream;
+
+    #[wasm_bindgen(method, getter, structural, js_name = abortSignal)]
+    pub fn abort_signal(this: &GenerateInput) -> AbortSignal;
 }

crates/node-bridge/src/bindings/mod.rs

@@ -4,4 +4,5 @@ pub mod console;
 pub mod https;
 pub mod uuid;
 
+pub use abort_signal::AbortSignal;
 pub use buffer::Buffer;

crates/node-bridge/src/http_client.rs

@@ -233,6 +233,9 @@ impl HttpResponse {
 
 impl Drop for HttpResponse {
     fn drop(&mut self) {
+        #[cfg(debug_assertions)]
+        console::log2(&"response dropped: ".into(), &self.resp);
+
         self.resp.destroy(None);
     }
 }

src/generate/generate.ts

@@ -3,60 +3,9 @@ import path = require("path");
 import { ResultStream } from "./result-stream";
 import {
     generateCode as rustGenerateCode,
-    IGenerateInput,
-    IResultStream,
     ISelectionRange,
 } from "../../crates/cursor-core/pkg";
 
-class GenerateInput implements IGenerateInput {
-    private _prompt: string;
-    private _documentText: string;
-    private _filePath: string;
-    private _workspaceDirectory: string | null;
-    private _selectionRange: ISelectionRange;
-    private _resultStream: IResultStream;
-
-    constructor(
-        prompt: string,
-        documentText: string,
-        filePath: string,
-        workspaceDirectory: string | null,
-        selectionRange: ISelectionRange,
-        resultStream: IResultStream
-    ) {
-        this._prompt = prompt;
-        this._documentText = documentText;
-        this._filePath = filePath;
-        this._workspaceDirectory = workspaceDirectory;
-        this._selectionRange = selectionRange;
-        this._resultStream = resultStream;
-    }
-
-    get prompt(): string {
-        return this._prompt;
-    }
-
-    get documentText(): string {
-        return this._documentText;
-    }
-
-    get filePath(): string {
-        return this._filePath;
-    }
-
-    get workspaceDirectory(): string | null {
-        return this._workspaceDirectory;
-    }
-
-    get selectionRange(): ISelectionRange {
-        return this._selectionRange;
-    }
-
-    get resultStream(): IResultStream {
-        return this._resultStream;
-    }
-}
-
 class SelectionRange implements ISelectionRange {
     private _offset: number;
     private _length: number;
@@ -81,27 +30,31 @@ export async function generateCode(
     cancellationToken: vscode.CancellationToken,
     resultStream: ResultStream<String>
 ): Promise<void> {
-    const document = editor.document;
+    const { document, selection } = editor;
     const filePath = document.uri.fsPath;
     const workspaceDirectory =
-        vscode.workspace.workspaceFolders?.[0].uri.fsPath ?? null;
-    const selection = editor.selection;
-    const text = document.getText();
+        vscode.workspace.getWorkspaceFolder(document.uri)?.uri.fsPath ?? null;
+    const documentText = document.getText();
     const selectionStartOffset = document.offsetAt(selection.start);
     const selectionEndOffset = document.offsetAt(selection.end);
     const selectionRange = new SelectionRange(
         selectionStartOffset,
         selectionEndOffset - selectionStartOffset
     );
 
-    await rustGenerateCode(
-        new GenerateInput(
-            prompt,
-            text,
-            filePath,
-            workspaceDirectory,
-            selectionRange,
-            resultStream
-        )
-    );
+    const abortController = new AbortController();
+    cancellationToken.onCancellationRequested(() => {
+        abortController.abort();
+    });
+    const { signal: abortSignal } = abortController;
+
+    await rustGenerateCode({
+        prompt,
+        documentText,
+        filePath,
+        workspaceDirectory,
+        selectionRange,
+        resultStream,
+        abortSignal,
+    });
 }