Add snippet sanitization for Zed completion items

proxy/src/completions.rs

@@ -1,26 +1,30 @@
 use serde_json::Value;
 
-pub fn should_sort_completions(msg: &Value) -> bool {
+/// Returns true if the message contains a completion response with items.
+pub fn is_completion_response(msg: &Value) -> bool {
     msg.get("result").is_some_and(|result| {
         result.get("items").is_some_and(|v| v.is_array()) || result.is_array()
     })
 }
 
-pub fn sort_completions_by_param_count(msg: &mut Value) {
-    let items = if let Some(result) = msg.get_mut("result") {
-        if result.is_array() {
-            result.as_array_mut()
-        } else {
-            result.get_mut("items").and_then(|v| v.as_array_mut())
-        }
-    } else {
-        None
+/// Single-pass processing of completion items:
+/// - Sorts methods/functions by parameter count (prepends count to sortText)
+/// - Strips unsupported VS Code snippet variables ($TM_SELECTED_TEXT) from snippets
+pub fn process_completions(msg: &mut Value) {
+    let items = match msg.get_mut("result") {
+        Some(result) if result.is_array() => result.as_array_mut(),
+        Some(result) => result.get_mut("items").and_then(|v| v.as_array_mut()),
+        None => None,
     };
 
-    if let Some(items) = items {
-        for item in items.iter_mut() {
-            let kind = item.get("kind").and_then(|v| v.as_u64()).unwrap_or(0);
-            if kind == 2 || kind == 3 {
+    let Some(items) = items else { return };
+
+    for item in items.iter_mut() {
+        let kind = item.get("kind").and_then(|v| v.as_u64()).unwrap_or(0);
+
+        match kind {
+            // Method (2) or Function (3): prepend param count to sortText
+            2 | 3 => {
                 let detail = item
                     .pointer("/labelDetails/detail")
                     .and_then(|v| v.as_str())
@@ -29,29 +33,55 @@ pub fn sort_completions_by_param_count(msg: &mut Value) {
                 let existing = item.get("sortText").and_then(|v| v.as_str()).unwrap_or("");
                 item["sortText"] = Value::String(format!("{count:02}{existing}"));
             }
+            // Snippet (15): strip $TM_SELECTED_TEXT
+            15 => {
+                strip_tm_selected_text(item, "textEditText");
+                strip_tm_selected_text(item, "insertText");
+            }
+            _ => {}
         }
     }
 }
 
-fn count_params(detail: &str) -> usize {
-    if detail.is_empty() || detail == "()" {
-        return 0;
+fn strip_tm_selected_text(item: &mut Value, key: &str) {
+    if let Some(text) = item.get(key).and_then(|v| v.as_str()) {
+        if text.contains("$TM_SELECTED_TEXT") {
+            item[key] = Value::String(text.replace("$TM_SELECTED_TEXT", ""));
+        }
     }
-    let inner = detail
-        .strip_prefix('(')
-        .and_then(|s| s.strip_suffix(')'))
-        .unwrap_or(detail)
-        .trim();
+}
+
+/// Sanitize a single resolved completion item (completionItem/resolve response).
+pub fn sanitize_resolved_completion(msg: &mut Value) {
+    let Some(result) = msg.get_mut("result") else {
+        return;
+    };
+    strip_tm_selected_text(result, "textEditText");
+    strip_tm_selected_text(result, "insertText");
+    // Also check inside textEdit.newText
+    if let Some(new_text) = result.pointer("/textEdit/newText").and_then(|v| v.as_str()) {
+        if new_text.contains("$TM_SELECTED_TEXT") {
+            result["textEdit"]["newText"] =
+                Value::String(new_text.replace("$TM_SELECTED_TEXT", ""));
+        }
+    }
+}
+
+fn count_params(detail: &str) -> usize {
+    let inner = match detail.strip_prefix('(').and_then(|s| s.strip_suffix(')')) {
+        Some(s) => s.trim(),
+        None => return 0,
+    };
     if inner.is_empty() {
         return 0;
     }
     let mut count = 1usize;
     let mut depth = 0i32;
-    for ch in inner.chars() {
+    for ch in inner.bytes() {
         match ch {
-            '<' => depth += 1,
-            '>' => depth -= 1,
-            ',' if depth == 0 => count += 1,
+            b'<' => depth += 1,
+            b'>' => depth -= 1,
+            b',' if depth == 0 => count += 1,
             _ => {}
         }
     }

proxy/src/main.rs

@@ -5,7 +5,7 @@ mod log;
 mod lsp;
 mod platform;
 
-use completions::{should_sort_completions, sort_completions_by_param_count};
+use completions::{is_completion_response, process_completions, sanitize_resolved_completion};
 use decompile::{rewrite_jdt_in_strings, rewrite_jdt_locations};
 use http::handle_http;
 use lsp::{parse_lsp_content, raw_has_id, write_raw, write_to_stdout, LspReader};
@@ -209,6 +209,7 @@ fn main() {
                                             &pending,
                                             &mut next_id,
                                         );
+                                        sanitize_resolved_completion(&mut msg);
                                     }
                                 }
                                 write_to_stdout(&msg);
@@ -217,9 +218,9 @@ fn main() {
                         }
                     }
 
-                    // Sort completion responses by param count
-                    if should_sort_completions(&msg) {
-                        sort_completions_by_param_count(&mut msg);
+                    // Process completion responses (sort + sanitize) in a single pass
+                    if is_completion_response(&msg) {
+                        process_completions(&mut msg);
                         write_to_stdout(&msg);
                         continue;
                     }