Adapt to the new version of Cursor API

Author: ktiays

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

@@ -1,4 +1,3 @@
-// pub mod session;
 pub mod token;
 
 use std::future::IntoFuture;
@@ -16,10 +15,13 @@ use uuid::Uuid;
 use wasm_bindgen::prelude::*;
 use wasm_bindgen_futures::future_to_promise;
 
+const AUTH_TOKEN_KEY: &str = "auth_token";
+
 use crate::{
     bindings::{progress_location::ProgressLocation, progress_options::ProgressOptions},
     context::get_extension_context,
     request::make_request,
+    storage::GlobalStorage,
 };
 
 use self::token::Token;
@@ -53,6 +55,14 @@ where
 
 #[wasm_bindgen(js_name = signIn)]
 pub async fn sign_in() {
+    let context = get_extension_context();
+    let storage = context.storage();
+    if storage.get(AUTH_TOKEN_KEY).is_some() {
+        // If there is already an authentication token, it means that the user has logged in
+        // and does not need to log in again.
+        return;
+    }
+
     let uuid = Uuid::new_v4().to_string();
     let verifier = base64_encode(random_bytes());
     let challenge = base64_encode(sha256(verifier.clone()));
@@ -62,7 +72,6 @@ pub async fn sign_in() {
         uuid.clone()
     );
 
-    let context = get_extension_context();
     // The API of VSCode does not allow us to obtain the execution result of the 'vscode.open' command,
     // so we cannot determine whether the user has confirmed to open url.
     context
@@ -73,16 +82,18 @@ pub async fn sign_in() {
         .with_progress(
             ProgressOptions {
                 location: ProgressLocation::Notification,
-                title: Some("Signing in...".to_owned()),
+                title: Some("Signing In/Up...".to_owned()),
                 cancellable: true,
             },
             closure!(|abort_signal: AbortSignal| {
                 let uuid = uuid.clone();
                 let verifier = verifier.clone();
+                let storage: GlobalStorage = storage.clone().into();
                 future_to_promise(async move {
-                    polling(&uuid, &verifier, abort_signal)
-                        .await
-                        .map(Into::into)
+                    if let Some(token) = polling(&uuid, &verifier, abort_signal).await? {
+                        storage.update(AUTH_TOKEN_KEY, Some(&token));
+                    }
+                    Ok(JsValue::null())
                 })
             })
             .into_js_value()
@@ -151,13 +162,47 @@ async fn polling(
     }
 }
 
-pub fn sign_out() {}
+#[wasm_bindgen(js_name = signOut)]
+pub fn sign_out() {
+    get_extension_context()
+        .storage()
+        .update(AUTH_TOKEN_KEY, None);
+}
+
+pub fn account_token() -> Option<Token> {
+    get_extension_context()
+        .storage()
+        .get(AUTH_TOKEN_KEY)
+        .and_then(|token| serde_json::from_str(&token).ok())
+}
 
 #[cfg(test)]
 mod test {
+    use super::*;
+
     #[test]
     fn test_random_bytes() {
         let bytes = super::random_bytes();
         assert_eq!(bytes.len(), 32);
     }
+
+    #[test]
+    fn test_base64_encode() {
+        let bytes = vec![
+            0xa9, 0x1e, 0x74, 0x36, 0x4a, 0x57, 0xb6, 0x40, 0xcf, 0x25, 0x37, 0xf7, 0x20, 0x26,
+            0x7a, 0x2e, 0x94, 0x90, 0x03, 0x85, 0x5b, 0xb8, 0xd0, 0x92, 0x37, 0xdc, 0xb3, 0xd9,
+            0x0a, 0x4d, 0xd4, 0xc5,
+        ];
+        let encoded = base64_encode(bytes);
+        assert_eq!(encoded, "qR50NkpXtkDPJTf3ICZ6LpSQA4VbuNCSN9yz2QpN1MU");
+    }
+
+    #[test]
+    fn test_sha256() {
+        let v = "qR50NkpXtkDPJTf3ICZ6LpSQA4VbuNCSN9yz2QpN1MU";
+        assert_eq!(
+            base64_encode(sha256(v)),
+            "ddiNacYgAjUZTDf6Pza1wRlSjuWIQRz5Z1Jc2Bj4DII"
+        );
+    }
 }

crates/cursor-core/src/auth/session.rs

@@ -1,6 +1,6 @@
 use serde::{Deserialize, Serialize};
 
-#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
+#[derive(Debug, Clone, Serialize, Deserialize)]
 #[serde(rename_all = "camelCase")]
 pub struct Token {
     pub access_token: String,

crates/cursor-core/src/auth/token.rs

@@ -17,7 +17,7 @@ extern "C" {
     #[wasm_bindgen(typescript_type = "IExtensionContext")]
     pub type ExtensionContext;
 
-    #[wasm_bindgen(method, structural)]
+    #[wasm_bindgen(method, structural, getter)]
     pub fn storage(this: &ExtensionContext) -> GlobalStorage;
 
     #[wasm_bindgen(method, structural, js_name = executeCommand)]

crates/cursor-core/src/context.rs

@@ -69,7 +69,7 @@ impl Session {
 
     pub async fn send_message(&mut self, input: &GenerateInput) -> Result<(), JsValue> {
         let request_body = self.body_with_input(input);
-        let mut state = make_conversation_request("/conversation", request_body, true).await?;
+        let mut state = make_conversation_request("/conversation", request_body).await?;
 
         let mut message: String = "".to_owned();
 

crates/cursor-core/src/conversation/chat/session.rs

@@ -6,100 +6,44 @@ use futures::{
 };
 use node_bridge::futures::Defer;
 use node_bridge::prelude::*;
-use uuid::Uuid;
 use wasm_bindgen::prelude::*;
 
 use crate::GenerateInput;
 
 use super::{
     make_conversation_request,
-    models::{BotMessage, MessageType, RequestBody},
+    models::{MessageType, RequestBody},
 };
 
 async fn generate_code_inner(input: &GenerateInput) -> Result<(), JsValue> {
-    let file_path = input.file_path();
     let selection = input.selection_range();
 
     let message_type = if selection.is_empty() {
         MessageType::Generate
     } else {
         MessageType::Edit
     };
-    let mut request_body = RequestBody::new_with_input(input, message_type);
-
+    let request_body = RequestBody::new_with_input(input, message_type);
     let result_stream = input.result_stream();
 
-    let mut conversation_id: Option<String> = None;
-    // The last message received from the server.
-    let mut previous_message: String = "".to_owned();
-    let mut last_token = "".to_owned();
-
-    // If the conversation was interrupted, we need to send a "continue" request.
-    let mut interrupted = false;
-
-    loop {
-        #[cfg(debug_assertions)]
-        console::log_str(&serde_json::to_string(&request_body).unwrap());
-
-        let mut state = make_conversation_request(
-            if interrupted {
-                "/continue/"
-            } else {
-                "/conversation"
-            },
-            &request_body,
-            !interrupted,
-        )
-        .await?;
-        interrupted = false;
-
-        #[cfg(debug_assertions)]
-        console::log_str("response received");
-
-        let mut data_stream = state.data_stream();
-        while let Some(mut data) = data_stream.next().await {
-            if data.contains("<|END_interrupt|>") {
-                interrupted = true;
-                last_token = data.clone();
-                // `END_interrupt` is included in valid data,
-                // we cannot discard it.
-                data = data.replace("<|END_interrupt|>", "");
-            }
-            previous_message.push_str(&data);
-            result_stream.write(&data);
-
-            #[cfg(debug_assertions)]
-            console::log_str(&format!("wrote: {}", &data));
-        }
-        drop(data_stream);
+    #[cfg(debug_assertions)]
+    console::log_str(&serde_json::to_string(&request_body).unwrap());
 
-        // Make sure the response is fully received without errors.
-        state.complete().await?;
+    let mut state = make_conversation_request("/conversation", &request_body).await?;
 
-        if !interrupted {
-            break;
-        }
+    #[cfg(debug_assertions)]
+    console::log_str("response received");
 
+    let mut data_stream = state.data_stream();
+    while let Some(data) = data_stream.next().await {
+        result_stream.write(&data);
         #[cfg(debug_assertions)]
-        console::log_str("generation interrupted");
-
-        // Generate an UUID as conversation ID.
-        if conversation_id.is_none() {
-            conversation_id = Some(Uuid::new_v4().to_string());
-        }
-        let bot_message = BotMessage::new(
-            conversation_id.clone().unwrap(),
-            message_type,
-            previous_message.clone(),
-            last_token.clone(),
-            file_path.to_owned(),
-            false,
-        );
-        request_body.bot_messages = vec![bot_message];
+        console::log_str(&format!("wrote: {}", &data));
     }
+    drop(data_stream);
 
-    console::log_str("generate done");
-
+    // Make sure the response is fully received without errors.
+    state.complete().await?;
     result_stream.end();
     Ok(())
 }