From b4df20a0fed6e6b5cd6c39ef3f5356db07617766 Mon Sep 17 00:00:00 2001 From: Riccardo Strina Date: Sat, 1 Aug 2026 19:59:59 +0200 Subject: [PATCH 1/7] feat(proxy): Coordinate decompilation --- Cargo.lock | 2 + proxy/Cargo.toml | 2 + proxy/src/completions.rs | 163 +++++- proxy/src/decompile.rs | 1159 +++++++++++++++++++++++++++++++++----- proxy/src/http.rs | 75 +-- proxy/src/log.rs | 22 +- proxy/src/main.rs | 547 +++++++++++++++--- proxy/src/output.rs | 122 ++++ proxy/src/pending.rs | 93 +++ 9 files changed, 1891 insertions(+), 294 deletions(-) create mode 100644 proxy/src/output.rs create mode 100644 proxy/src/pending.rs diff --git a/Cargo.lock b/Cargo.lock index 6498bee..2730aca 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -504,9 +504,11 @@ checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" name = "java-lsp-proxy" version = "6.8.23" dependencies = [ + "hex", "proxy-common", "serde", "serde_json", + "sha1", ] [[package]] diff --git a/proxy/Cargo.toml b/proxy/Cargo.toml index dad6fea..50a81fe 100644 --- a/proxy/Cargo.toml +++ b/proxy/Cargo.toml @@ -12,5 +12,7 @@ path = "src/main.rs" [dependencies] proxy-common.workspace = true +hex = "0.4" serde = { version = "1.0", features = ["derive"] } serde_json = "1.0" +sha1 = "0.10" diff --git a/proxy/src/completions.rs b/proxy/src/completions.rs index cdd1a79..c575fd4 100644 --- a/proxy/src/completions.rs +++ b/proxy/src/completions.rs @@ -1,19 +1,15 @@ use serde_json::Value; -/// 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() - }) -} - /// 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 default_insert_text_format = msg + .pointer("/result/itemDefaults/insertTextFormat") + .and_then(Value::as_u64); 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()), + Some(result) => result.get_mut("items").and_then(Value::as_array_mut), None => None, }; @@ -33,18 +29,31 @@ pub fn process_completions(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"); - } _ => {} } + + let insert_text_format = item + .get("insertTextFormat") + .and_then(Value::as_u64) + .or(default_insert_text_format); + if kind == 15 || insert_text_format == Some(2) { + sanitize_completion_item(item); + } + } +} + +fn sanitize_completion_item(item: &mut Value) { + strip_tm_selected_text(item, "textEditText"); + strip_tm_selected_text(item, "insertText"); + if let Some(new_text) = item.pointer("/textEdit/newText").and_then(Value::as_str) { + if new_text.contains("$TM_SELECTED_TEXT") { + item["textEdit"]["newText"] = Value::String(new_text.replace("$TM_SELECTED_TEXT", "")); + } } } fn strip_tm_selected_text(item: &mut Value, key: &str) { - if let Some(text) = item.get(key).and_then(|v| v.as_str()) { + if let Some(text) = item.get(key).and_then(Value::as_str) { if text.contains("$TM_SELECTED_TEXT") { item[key] = Value::String(text.replace("$TM_SELECTED_TEXT", "")); } @@ -56,15 +65,7 @@ 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", "")); - } - } + sanitize_completion_item(result); } fn count_params(detail: &str) -> usize { @@ -87,3 +88,119 @@ fn count_params(detail: &str) -> usize { } count } + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + #[test] + fn processes_array_completion_results() { + let mut response = json!({ + "result": [{ + "kind": 2, + "labelDetails": { "detail": "(String, List)" }, + "sortText": "method" + }] + }); + + process_completions(&mut response); + + assert_eq!(response["result"][0]["sortText"], json!("02method")); + } + + #[test] + fn applies_list_item_default_snippet_format() { + let mut response = json!({ + "result": { + "itemDefaults": { "insertTextFormat": 2 }, + "items": [{ + "kind": 2, + "insertText": "$TM_SELECTED_TEXT.trim()" + }] + } + }); + + process_completions(&mut response); + + assert_eq!( + response["result"]["items"][0]["insertText"], + json!(".trim()") + ); + } + + #[test] + fn sanitizes_insert_replace_edit_text() { + let mut response = json!({ + "result": [{ + "kind": 15, + "textEdit": { + "newText": "$TM_SELECTED_TEXT.var", + "insert": { + "start": { "line": 0, "character": 0 }, + "end": { "line": 0, "character": 0 } + }, + "replace": { + "start": { "line": 0, "character": 0 }, + "end": { "line": 0, "character": 3 } + } + } + }] + }); + + process_completions(&mut response); + + assert_eq!(response["result"][0]["textEdit"]["newText"], json!(".var")); + } + + #[test] + fn leaves_plain_text_completion_unchanged() { + let mut response = json!({ + "result": [{ + "kind": 1, + "insertTextFormat": 1, + "insertText": "$TM_SELECTED_TEXT" + }] + }); + + process_completions(&mut response); + + assert_eq!( + response["result"][0]["insertText"], + json!("$TM_SELECTED_TEXT") + ); + } + + #[test] + fn postfix_var_completion_keeps_jdtls_spacing() { + let insertion = "var name = \"hello world\";"; + let mut response = json!({ + "result": [{ + "label": ".var", + "kind": 15, + "insertTextFormat": 2, + "textEdit": { + "newText": insertion, + "range": { + "start": { "line": 0, "character": 0 }, + "end": { "line": 0, "character": 17 } + } + } + }] + }); + + process_completions(&mut response); + + assert_eq!( + response["result"][0]["textEdit"]["newText"], + json!(insertion) + ); + } + + #[test] + fn counts_nested_generic_parameters() { + assert_eq!(count_params("(Map>, int)"), 2); + assert_eq!(count_params("()"), 0); + assert_eq!(count_params("not-a-signature"), 0); + } +} diff --git a/proxy/src/decompile.rs b/proxy/src/decompile.rs index f8e079c..b217402 100644 --- a/proxy/src/decompile.rs +++ b/proxy/src/decompile.rs @@ -1,203 +1,1090 @@ +use crate::{lsp_error, pending::PendingResponses}; +use proxy_common::{encode_lsp, path_to_file_uri}; use serde_json::{json, Value}; +use sha1::{Digest, Sha1}; use std::{ - collections::{hash_map::DefaultHasher, HashMap}, - env, fs, - hash::{Hash, Hasher}, + collections::{HashMap, HashSet, VecDeque}, + env, + fs::{self, OpenOptions}, io::Write, - path::PathBuf, - sync::{mpsc, Arc, Mutex}, + path::{Path, PathBuf}, + sync::{ + atomic::{AtomicU64, Ordering}, + Arc, Condvar, Mutex, + }, + thread, + time::{Duration, Instant}, }; -use crate::{lsp_error, lsp_warn}; -use proxy_common::{encode_lsp, path_to_file_uri}; - const DECOMPILED_DIR: &str = "jdtls-decompiled"; +const FETCH_WORKERS: usize = 2; +const JOB_WORKERS: usize = 2; +const MAX_QUEUED_JOBS: usize = 64; +const FETCH_TIMEOUT: Duration = Duration::from_secs(5); +const NEGATIVE_CACHE_TTL: Duration = Duration::from_secs(2); + +pub type SharedWriter = Arc>>; + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum RewriteMode { + Locations, + Strings, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Priority { + Interactive, + Bulk, +} + +pub struct RewriteJob { + pub token: u64, + pub message: Value, + pub mode: RewriteMode, + pub priority: Priority, + pub deadline: Instant, + pub complete: Box, +} + +#[derive(Clone)] +pub struct DecompileCoordinator { + inner: Arc, +} + +struct Inner { + fetcher: Arc, + owned_id_prefix: String, + request_counter: AtomicU64, + max_bulk_jobs: usize, + state: Mutex, + work_available: Condvar, + state_changed: Condvar, +} + +#[derive(Default)] +struct State { + uris: HashMap, + interactive_uris: VecDeque, + bulk_uris: VecDeque, + interactive_jobs: VecDeque, + bulk_jobs: VecDeque, + canceled_jobs: HashSet, + latest_bulk_job: Option, + bulk_jobs_active: usize, + bulk_fetches: usize, + shutdown: bool, +} + +struct UriEntry { + status: UriStatus, + waiters: HashSet, +} + +enum UriStatus { + Queued(Priority), + InFlight { request_id: Value }, + Ready(String), + Failed(Instant), +} + +trait ClassContentFetcher: Send + Sync { + fn fetch(&self, uri: &str, request_id: Value) -> Option; + fn cancel(&self, request_id: &Value); +} + +struct JdtlsFetcher { + writer: SharedWriter, + pending: Arc, +} + +impl ClassContentFetcher for JdtlsFetcher { + fn fetch(&self, uri: &str, request_id: Value) -> Option { + let receiver = self.pending.register(request_id.clone()); + let request = encode_lsp(&json!({ + "jsonrpc": "2.0", + "id": request_id, + "method": "java/classFileContents", + "params": { "uri": uri } + })); + + let write_succeeded = { + let mut writer = self.writer.lock().unwrap(); + writer.write_all(request.as_bytes()).is_ok() && writer.flush().is_ok() + }; + if !write_succeeded { + self.pending.remove(&request_id); + return None; + } + + match receiver.recv_timeout(FETCH_TIMEOUT) { + Ok(response) => response + .get("result") + .and_then(Value::as_str) + .filter(|content| !content.is_empty()) + .map(str::to_string), + Err(_) => { + self.pending.remove(&request_id); + None + } + } + } + + fn cancel(&self, request_id: &Value) { + self.pending.remove(request_id); + let cancel = encode_lsp(&json!({ + "jsonrpc": "2.0", + "method": "$/cancelRequest", + "params": { "id": request_id } + })); + let mut writer = self.writer.lock().unwrap(); + let _ = writer.write_all(cancel.as_bytes()); + let _ = writer.flush(); + } +} + +impl DecompileCoordinator { + pub fn new( + writer: SharedWriter, + pending: Arc, + owned_id_prefix: String, + ) -> Self { + Self::with_fetcher( + Arc::new(JdtlsFetcher { writer, pending }), + owned_id_prefix, + FETCH_WORKERS, + JOB_WORKERS, + ) + } + + fn with_fetcher( + fetcher: Arc, + owned_id_prefix: String, + fetch_workers: usize, + job_workers: usize, + ) -> Self { + let coordinator = Self { + inner: Arc::new(Inner { + fetcher, + owned_id_prefix, + request_counter: AtomicU64::new(1), + max_bulk_jobs: job_workers.saturating_sub(1).max(1), + state: Mutex::new(State::default()), + work_available: Condvar::new(), + state_changed: Condvar::new(), + }), + }; + + for _ in 0..fetch_workers { + let inner = Arc::clone(&coordinator.inner); + thread::spawn(move || fetch_worker(inner)); + } + for _ in 0..job_workers { + let inner = Arc::clone(&coordinator.inner); + thread::spawn(move || job_worker(inner)); + } + + coordinator + } + + /// Enqueues without blocking the JDTLS stdout router. On saturation the + /// original job is returned so the caller can forward it unchanged. + pub fn submit(&self, job: RewriteJob) -> Result<(), RewriteJob> { + let mut state = self.inner.state.lock().unwrap(); + if state.shutdown || state.interactive_jobs.len() + state.bulk_jobs.len() >= MAX_QUEUED_JOBS + { + return Err(job); + } + + if job.priority == Priority::Bulk { + if let Some(previous) = state.latest_bulk_job.replace(job.token) { + if cancel_job_locked(&mut state, previous) { + state.canceled_jobs.remove(&previous); + } + } + state.bulk_jobs.push_back(job); + } else { + state.interactive_jobs.push_back(job); + } + self.inner.work_available.notify_all(); + Ok(()) + } + + pub fn cancel(&self, token: u64) { + let request_ids = { + let mut state = self.inner.state.lock().unwrap(); + if cancel_job_locked(&mut state, token) { + state.canceled_jobs.remove(&token); + } + orphaned_request_ids(&state) + }; + for request_id in request_ids { + self.inner.fetcher.cancel(&request_id); + } + self.inner.state_changed.notify_all(); + self.inner.work_available.notify_all(); + } + + pub fn consume_cancellation(&self, token: u64) -> bool { + self.inner + .state + .lock() + .unwrap() + .canceled_jobs + .remove(&token) + } + + pub fn is_canceled(&self, token: u64) -> bool { + self.inner + .state + .lock() + .unwrap() + .canceled_jobs + .contains(&token) + } + + pub fn shutdown(&self) { + let mut state = self.inner.state.lock().unwrap(); + state.shutdown = true; + state.interactive_jobs.clear(); + state.bulk_jobs.clear(); + self.inner.work_available.notify_all(); + self.inner.state_changed.notify_all(); + } +} + +fn cancel_job_locked(state: &mut State, token: u64) -> bool { + state.canceled_jobs.insert(token); + let queued_jobs = state.interactive_jobs.len() + state.bulk_jobs.len(); + state + .interactive_jobs + .retain(|queued| queued.token != token); + state.bulk_jobs.retain(|queued| queued.token != token); + for entry in state.uris.values_mut() { + entry.waiters.remove(&token); + } + queued_jobs != state.interactive_jobs.len() + state.bulk_jobs.len() +} + +fn orphaned_request_ids(state: &State) -> Vec { + state + .uris + .values() + .filter(|entry| entry.waiters.is_empty()) + .filter_map(|entry| match &entry.status { + UriStatus::InFlight { request_id, .. } => Some(request_id.clone()), + _ => None, + }) + .collect() +} + +fn job_worker(inner: Arc) { + loop { + let job = { + let mut state = inner.state.lock().unwrap(); + loop { + if state.shutdown { + return; + } + if let Some(job) = state.interactive_jobs.pop_front() { + break job; + } + if state.bulk_jobs_active < inner.max_bulk_jobs { + if let Some(job) = state.bulk_jobs.pop_front() { + state.bulk_jobs_active += 1; + break job; + } + } + state = inner.work_available.wait(state).unwrap(); + } + }; + + if is_canceled(&inner, job.token) { + finish_job(&inner, job.token, job.priority); + continue; + } + + let mut message = job.message; + rewrite_message( + &inner, + job.token, + &mut message, + job.mode, + job.priority, + job.deadline, + ); + let canceled = is_canceled(&inner, job.token); + finish_job(&inner, job.token, job.priority); + if !canceled { + (job.complete)(message); + } + } +} + +fn is_canceled(inner: &Inner, token: u64) -> bool { + inner.state.lock().unwrap().canceled_jobs.contains(&token) +} + +fn finish_job(inner: &Inner, token: u64, priority: Priority) { + let mut state = inner.state.lock().unwrap(); + if priority == Priority::Bulk { + state.bulk_jobs_active = state.bulk_jobs_active.saturating_sub(1); + } + state.canceled_jobs.remove(&token); + if state.latest_bulk_job == Some(token) { + state.latest_bulk_job = None; + } + for entry in state.uris.values_mut() { + entry.waiters.remove(&token); + } + let now = Instant::now(); + state.uris.retain(|_, entry| { + !entry.waiters.is_empty() + || matches!(entry.status, UriStatus::InFlight { .. }) + || matches!(entry.status, UriStatus::Failed(retry_after) if retry_after > now) + }); + inner.work_available.notify_all(); +} + +fn rewrite_message( + inner: &Inner, + token: u64, + message: &mut Value, + mode: RewriteMode, + priority: Priority, + deadline: Instant, +) { + let Some(result) = message.get_mut("result") else { + return; + }; + + let mut uris = Vec::new(); + match mode { + RewriteMode::Locations => collect_jdt_location_uris(result, &mut uris, &mut HashSet::new()), + RewriteMode::Strings => collect_jdt_uris(result, &mut uris, &mut HashSet::new()), + } + if uris.is_empty() { + return; + } + + let replacements = resolve_uris(inner, token, &uris, priority, deadline); + match mode { + RewriteMode::Locations => { + replace_jdt_location_uris(result, &replacements); + } + RewriteMode::Strings => replace_in_strings(result, &replacements), + } +} + +fn resolve_uris( + inner: &Inner, + token: u64, + uris: &[String], + priority: Priority, + deadline: Instant, +) -> HashMap { + let mut replacements = HashMap::new(); + let mut unresolved = Vec::new(); + + for uri in uris { + let path = cache_path(uri); + if path.is_file() { + replacements.insert(uri.clone(), path_to_file_uri(&path)); + } else { + unresolved.push(uri.clone()); + } + } + if unresolved.is_empty() { + return replacements; + } + + let mut state = inner.state.lock().unwrap(); + let now = Instant::now(); + state.uris.retain(|_, entry| { + !entry.waiters.is_empty() + || matches!(entry.status, UriStatus::InFlight { .. }) + || matches!(entry.status, UriStatus::Failed(retry_after) if retry_after > now) + }); + if state.canceled_jobs.contains(&token) { + return replacements; + } + + for uri in &unresolved { + let mut enqueue = None; + match state.uris.get_mut(uri) { + Some(entry) => { + entry.waiters.insert(token); + match entry.status { + UriStatus::Ready(ref file_uri) => { + replacements.insert(uri.clone(), file_uri.clone()); + } + UriStatus::Failed(retry_after) if retry_after <= Instant::now() => { + entry.status = UriStatus::Queued(priority); + enqueue = Some(priority); + } + UriStatus::Queued(Priority::Bulk) if priority == Priority::Interactive => { + entry.status = UriStatus::Queued(Priority::Interactive); + enqueue = Some(Priority::Interactive); + } + _ => {} + } + } + None => { + state.uris.insert( + uri.clone(), + UriEntry { + status: UriStatus::Queued(priority), + waiters: HashSet::from([token]), + }, + ); + enqueue = Some(priority); + } + } + if let Some(priority) = enqueue { + queue_uri(&mut state, uri.clone(), priority); + } + } + inner.work_available.notify_all(); + + loop { + let mut waiting = false; + for uri in &unresolved { + match state.uris.get(uri).map(|entry| &entry.status) { + Some(UriStatus::Ready(file_uri)) => { + replacements.insert(uri.clone(), file_uri.clone()); + } + Some(UriStatus::Queued(_) | UriStatus::InFlight { .. }) => waiting = true, + Some(UriStatus::Failed(_)) | None => {} + } + } + + if !waiting || state.canceled_jobs.contains(&token) || Instant::now() >= deadline { + break; + } + + let remaining = deadline.saturating_duration_since(Instant::now()); + let (next_state, _) = inner.state_changed.wait_timeout(state, remaining).unwrap(); + state = next_state; + } + + for uri in &unresolved { + if let Some(entry) = state.uris.get_mut(uri) { + entry.waiters.remove(&token); + } + } + let orphaned = orphaned_request_ids(&state); + drop(state); + for request_id in orphaned { + inner.fetcher.cancel(&request_id); + } + + replacements +} + +fn queue_uri(state: &mut State, uri: String, priority: Priority) { + match priority { + Priority::Interactive => state.interactive_uris.push_back(uri), + Priority::Bulk => state.bulk_uris.push_back(uri), + } +} + +fn fetch_worker(inner: Arc) { + loop { + let Some((uri, priority, request_id)) = take_uri_work(&inner) else { + return; + }; + + let content = inner.fetcher.fetch(&uri, request_id.clone()); + let resolved = content.and_then(|content| write_cached_source(&uri, content.as_bytes())); + + let mut state = inner.state.lock().unwrap(); + if priority == Priority::Bulk { + state.bulk_fetches = state.bulk_fetches.saturating_sub(1); + } + if let Some(entry) = state.uris.get_mut(&uri) { + let is_current = matches!( + &entry.status, + UriStatus::InFlight { + request_id: current, + .. + } if current == &request_id + ); + if is_current { + entry.status = match resolved { + Some(file_uri) => UriStatus::Ready(file_uri), + None => UriStatus::Failed(Instant::now() + NEGATIVE_CACHE_TTL), + }; + } + } + let remove_ready = state.uris.get(&uri).is_some_and(|entry| { + entry.waiters.is_empty() && matches!(entry.status, UriStatus::Ready(_)) + }); + if remove_ready { + state.uris.remove(&uri); + } + inner.state_changed.notify_all(); + inner.work_available.notify_all(); + } +} + +fn take_uri_work(inner: &Inner) -> Option<(String, Priority, Value)> { + let mut state = inner.state.lock().unwrap(); + loop { + if state.shutdown { + return None; + } + + let candidate = pop_valid_uri(&mut state, Priority::Interactive).or_else(|| { + if state.bulk_fetches == 0 { + pop_valid_uri(&mut state, Priority::Bulk) + } else { + None + } + }); + + if let Some((uri, priority)) = candidate { + let sequence = inner.request_counter.fetch_add(1, Ordering::Relaxed); + let request_id = + Value::String(format!("{}decompile-{sequence}", inner.owned_id_prefix)); + if priority == Priority::Bulk { + state.bulk_fetches += 1; + } + if let Some(entry) = state.uris.get_mut(&uri) { + entry.status = UriStatus::InFlight { + request_id: request_id.clone(), + }; + } + return Some((uri, priority, request_id)); + } + + state = inner.work_available.wait(state).unwrap(); + } +} + +fn pop_valid_uri(state: &mut State, priority: Priority) -> Option<(String, Priority)> { + let queue = match priority { + Priority::Interactive => &mut state.interactive_uris, + Priority::Bulk => &mut state.bulk_uris, + }; + while let Some(uri) = queue.pop_front() { + let valid = state.uris.get(&uri).is_some_and(|entry| { + !entry.waiters.is_empty() + && matches!(entry.status, UriStatus::Queued(current) if current == priority) + }); + if valid { + return Some((uri, priority)); + } + } + None +} fn cache_dir() -> PathBuf { env::temp_dir().join(DECOMPILED_DIR) } fn cache_path(uri: &str) -> PathBuf { - let mut hasher = DefaultHasher::new(); - uri.hash(&mut hasher); - let hex = format!("{:016x}", hasher.finish()); + cache_path_in(&cache_dir(), uri) +} - // jdt://contents/java.base/java.util/ArrayList.java?=.../%3Cjava.util%28ArrayList.class - // The class name is between the last %28 (URL-encoded '(') and .class at the end - let name = uri +fn cache_path_in(directory: &Path, uri: &str) -> PathBuf { + let digest = hex::encode(Sha1::digest(uri.as_bytes())); + let raw_name = uri .rsplit_once("%28") .and_then(|(_, rest)| rest.strip_suffix(".class")) .or_else(|| { uri.split('?') .next() .and_then(|path| path.rsplit('/').next()) - .and_then(|seg| seg.strip_suffix(".java").or(seg.strip_suffix(".class"))) + .and_then(|segment| { + segment + .strip_suffix(".java") + .or(segment.strip_suffix(".class")) + }) }) .unwrap_or("Decompiled"); + let name: String = raw_name + .chars() + .map(|character| { + if character.is_ascii_alphanumeric() || matches!(character, '-' | '_' | '.') { + character + } else { + '_' + } + }) + .collect(); - cache_dir().join(format!("{name}-{hex}.java")) -} - -/// Send `java/classFileContents` to JDTLS and wait for the response. -fn fetch_class_contents( - uri: &str, - writer: &Arc>, - pending: &Arc>>>, - request_id: Value, -) -> Option { - let (tx, rx) = mpsc::channel(); - pending.lock().unwrap().insert(request_id.clone(), tx); - - let req = encode_lsp(&json!({ - "jsonrpc": "2.0", - "id": request_id, - "method": "java/classFileContents", - "params": { "uri": uri } - })); - { - let mut w = writer.lock().unwrap(); - let _ = w.write_all(req.as_bytes()); - let _ = w.flush(); - } - - match rx.recv_timeout(std::time::Duration::from_secs(10)) { - Ok(resp) => { - let content = resp.get("result")?.as_str()?; - Some(content.to_string()) - } - Err(_) => { - lsp_warn!("[decompile] Timed out fetching class contents for {uri}"); - None - } - } + directory.join(format!("{name}-{digest}.java")) } -fn resolve_jdt_uri( - uri: &str, - writer: &Arc>, - pending: &Arc>>>, - request_id: Value, -) -> Option { - let path = cache_path(uri); - if path.exists() { - return Some(path_to_file_uri(&path)); +fn write_cached_source(uri: &str, content: &[u8]) -> Option { + if content.is_empty() { + return None; + } + + let directory = cache_dir(); + if let Err(error) = fs::create_dir_all(&directory) { + lsp_error!( + "[decompile] Failed to create {}: {error}", + directory.display() + ); + return None; } + let target = cache_path_in(&directory, uri); + if target.is_file() { + return Some(path_to_file_uri(&target)); + } + + static TEMP_COUNTER: AtomicU64 = AtomicU64::new(1); + let sequence = TEMP_COUNTER.fetch_add(1, Ordering::Relaxed); + let temporary = target.with_extension(format!("tmp-{}-{sequence}", std::process::id())); + let result = (|| { + let mut file = OpenOptions::new() + .write(true) + .create_new(true) + .open(&temporary)?; + file.write_all(content)?; + file.flush()?; + fs::rename(&temporary, &target) + })(); - let content = fetch_class_contents(uri, writer, pending, request_id)?; - let _ = fs::create_dir_all(cache_dir()); - match fs::write(&path, &content) { - Ok(_) => Some(path_to_file_uri(&path)), - Err(e) => { - lsp_error!("[decompile] Failed to write {}: {e}", path.display()); + match result { + Ok(()) => Some(path_to_file_uri(&target)), + Err(_error) if target.is_file() => { + let _ = fs::remove_file(&temporary); + Some(path_to_file_uri(&target)) + } + Err(error) => { + let _ = fs::remove_file(&temporary); + lsp_error!( + "[decompile] Failed to atomically write {}: {error}", + target.display() + ); None } } } -/// Rewrite any `jdt://` URIs in a definition/typeDefinition/implementation response. -/// Returns `true` if any URI was rewritten. -pub fn rewrite_jdt_locations( - msg: &mut Value, - writer: &Arc>, - pending: &Arc>>>, - next_id: &mut impl FnMut() -> Value, -) -> bool { - let results = match msg.get_mut("result") { - Some(Value::Array(arr)) => arr.iter_mut().collect::>(), - Some(obj @ Value::Object(_)) => vec![obj], - _ => return false, - }; +fn collect_jdt_location_uris(value: &Value, uris: &mut Vec, seen: &mut HashSet) { + match value { + Value::Array(values) => { + for value in values { + collect_jdt_location_uris(value, uris, seen); + } + } + Value::Object(object) => { + for (key, value) in object { + if matches!(key.as_str(), "uri" | "targetUri") { + if let Value::String(uri) = value { + if uri.starts_with("jdt://") && seen.insert(uri.clone()) { + uris.push(uri.clone()); + } + } + continue; + } + collect_jdt_location_uris(value, uris, seen); + } + } + _ => {} + } +} - let mut rewritten = false; - for loc in results { - for key in &["uri", "targetUri"] { - if let Some(Value::String(uri)) = loc.get(key) { - if uri.starts_with("jdt://") { - let jdt_uri = uri.clone(); - if let Some(file_uri) = resolve_jdt_uri(&jdt_uri, writer, pending, next_id()) { - loc[*key] = Value::String(file_uri); - rewritten = true; +fn replace_jdt_location_uris(value: &mut Value, replacements: &HashMap) -> bool { + match value { + Value::Array(values) => { + let mut rewritten = false; + for value in values { + rewritten |= replace_jdt_location_uris(value, replacements); + } + rewritten + } + Value::Object(object) => { + let mut rewritten = false; + for (key, value) in object { + if matches!(key.as_str(), "uri" | "targetUri") { + if let Value::String(uri) = value { + if let Some(file_uri) = replacements.get(uri) { + *uri = file_uri.clone(); + rewritten = true; + } } + continue; } + rewritten |= replace_jdt_location_uris(value, replacements); } + rewritten } + _ => false, } - rewritten } -/// A jdt:// URI in embedded markdown/text terminates at whitespace or any of these -/// delimiters commonly used in markdown links and JSON strings. The URI itself only -/// contains URL-encoded forms of these characters, so scanning until we hit one of -/// them is safe. -fn jdt_uri_end(s: &str) -> usize { - s.find(|c: char| c.is_whitespace() || matches!(c, ')' | ']' | '"' | '>' | '`' | '\'')) - .unwrap_or(s.len()) +fn jdt_uri_end(value: &str) -> usize { + value + .find(|character: char| { + character.is_whitespace() || matches!(character, ')' | ']' | '"' | '>' | '`' | '\'') + }) + .unwrap_or(value.len()) } -/// Extract all unique `jdt://` URIs appearing inside any string in `value`. -fn collect_jdt_uris(value: &Value, out: &mut Vec) { +fn collect_jdt_uris(value: &Value, uris: &mut Vec, seen: &mut HashSet) { match value { - Value::String(s) => { - let mut rest = s.as_str(); - while let Some(pos) = rest.find("jdt://") { - let tail = &rest[pos..]; + Value::String(string) => { + let mut rest = string.as_str(); + while let Some(position) = rest.find("jdt://") { + let tail = &rest[position..]; let end = jdt_uri_end(tail); let uri = tail[..end].to_string(); - if !out.contains(&uri) { - out.push(uri); + if seen.insert(uri.clone()) { + uris.push(uri); } rest = &tail[end..]; } } - Value::Array(arr) => arr.iter().for_each(|v| collect_jdt_uris(v, out)), - Value::Object(obj) => obj.values().for_each(|v| collect_jdt_uris(v, out)), + Value::Array(values) => { + for value in values { + collect_jdt_uris(value, uris, seen); + } + } + Value::Object(object) => { + for value in object.values() { + collect_jdt_uris(value, uris, seen); + } + } _ => {} } } -/// Replace all occurrences of any key in `map` with its value, inside every string -/// contained in `value` (recursively). -fn replace_in_strings(value: &mut Value, map: &HashMap) { +fn replace_in_strings(value: &mut Value, replacements: &HashMap) { match value { - Value::String(s) => { - for (from, to) in map { - if s.contains(from.as_str()) { - *s = s.replace(from.as_str(), to); + Value::String(string) => { + for (from, to) in replacements { + if string.contains(from) { + *string = string.replace(from, to); } } } - Value::Array(arr) => arr.iter_mut().for_each(|v| replace_in_strings(v, map)), - Value::Object(obj) => obj.values_mut().for_each(|v| replace_in_strings(v, map)), + Value::Array(values) => { + for value in values { + replace_in_strings(value, replacements); + } + } + Value::Object(object) => { + for value in object.values_mut() { + replace_in_strings(value, replacements); + } + } _ => {} } } -/// Scan a documentation response (hover, signatureHelp, completionItem/resolve, …) -/// for embedded `jdt://` URIs, resolve each one to a `file://` URI backed by a temp -/// file, and replace the URIs in-place in every string of `msg.result`. -pub fn rewrite_jdt_in_strings( - msg: &mut Value, - writer: &Arc>, - pending: &Arc>>>, - next_id: &mut impl FnMut() -> Value, -) { - let Some(result) = msg.get_mut("result") else { - return; - }; +#[cfg(test)] +mod tests { + use super::*; + use std::sync::atomic::{AtomicUsize, Ordering}; - let mut uris = Vec::new(); - collect_jdt_uris(result, &mut uris); - if uris.is_empty() { - return; + struct FakeFetcher { + calls: AtomicUsize, + cancellations: AtomicUsize, + active: AtomicUsize, + max_active: AtomicUsize, + delay: Duration, + succeeds: bool, } - let mut map = HashMap::new(); - for uri in uris { - if let Some(file_uri) = resolve_jdt_uri(&uri, writer, pending, next_id()) { - map.insert(uri, file_uri); + impl FakeFetcher { + fn new(delay: Duration) -> Self { + Self { + calls: AtomicUsize::new(0), + cancellations: AtomicUsize::new(0), + active: AtomicUsize::new(0), + max_active: AtomicUsize::new(0), + delay, + succeeds: true, + } + } + + fn failing(delay: Duration) -> Self { + Self { + succeeds: false, + ..Self::new(delay) + } } } - if !map.is_empty() { - replace_in_strings(result, &map); + + impl ClassContentFetcher for FakeFetcher { + fn fetch(&self, uri: &str, _request_id: Value) -> Option { + self.calls.fetch_add(1, Ordering::Relaxed); + let active = self.active.fetch_add(1, Ordering::SeqCst) + 1; + self.max_active.fetch_max(active, Ordering::SeqCst); + thread::sleep(self.delay); + self.active.fetch_sub(1, Ordering::SeqCst); + self.succeeds.then(|| format!("class {} {{}}", uri.len())) + } + + fn cancel(&self, _request_id: &Value) { + self.cancellations.fetch_add(1, Ordering::Relaxed); + } + } + + fn unique_uri(name: &str) -> String { + format!( + "jdt://contents/test/{name}-{}.class", + AtomicU64::new(1).fetch_add(1, Ordering::Relaxed) + std::process::id() as u64 + ) + } + + fn rewrite_for_test( + coordinator: &DecompileCoordinator, + token: u64, + uris: &[String], + priority: Priority, + ) -> HashMap { + resolve_uris( + &coordinator.inner, + token, + uris, + priority, + Instant::now() + Duration::from_secs(2), + ) + } + + #[test] + fn rewrites_nested_workspace_symbol_locations() { + let mut result = json!([ + { + "location": { + "uri": "jdt://contents/java.base/java.lang/String.class" + } + }, + { + "location": { + "uri": "file:///workspace/ProjectString.java" + } + } + ]); + let replacements = HashMap::from([( + "jdt://contents/java.base/java.lang/String.class".to_string(), + "file:///tmp/String.java".to_string(), + )]); + + assert!(replace_jdt_location_uris(&mut result, &replacements)); + assert_eq!( + result[0]["location"]["uri"], + json!("file:///tmp/String.java") + ); + assert_eq!( + result[1]["location"]["uri"], + json!("file:///workspace/ProjectString.java") + ); + } + + #[test] + fn deduplicates_location_uris_before_resolution() { + let uri = "jdt://contents/java.base/java.lang/String.class"; + let result = json!([{ "uri": uri }, { "targetUri": uri }]); + let mut uris = Vec::new(); + + collect_jdt_location_uris(&result, &mut uris, &mut HashSet::new()); + + assert_eq!(uris, vec![uri]); + } + + #[test] + fn ignores_jdt_uris_outside_location_fields() { + let result = json!({ + "documentation": "jdt://contents/java.base/java.lang/String.class", + "data": { + "sourceUri": "jdt://contents/java.base/java.lang/Object.class" + } + }); + let mut uris = Vec::new(); + + collect_jdt_location_uris(&result, &mut uris, &mut HashSet::new()); + + assert!(uris.is_empty()); + } + + #[test] + fn concurrent_waiters_share_one_fetch() { + let fetcher = Arc::new(FakeFetcher::new(Duration::from_millis(20))); + let coordinator = + DecompileCoordinator::with_fetcher(fetcher.clone(), "test-".to_string(), 2, 0); + let uri = unique_uri("single-flight"); + let _ = fs::remove_file(cache_path(&uri)); + + let first = { + let coordinator = coordinator.clone(); + let uri = uri.clone(); + thread::spawn(move || rewrite_for_test(&coordinator, 1, &[uri], Priority::Interactive)) + }; + let second = { + let coordinator = coordinator.clone(); + let uri = uri.clone(); + thread::spawn(move || rewrite_for_test(&coordinator, 2, &[uri], Priority::Interactive)) + }; + + assert_eq!(first.join().unwrap().len(), 1); + assert_eq!(second.join().unwrap().len(), 1); + assert_eq!(fetcher.calls.load(Ordering::Relaxed), 1); + let _ = fs::remove_file(cache_path(&uri)); + coordinator.shutdown(); + } + + #[test] + fn interactive_fetches_use_bounded_parallelism() { + let fetcher = Arc::new(FakeFetcher::new(Duration::from_millis(30))); + let coordinator = + DecompileCoordinator::with_fetcher(fetcher.clone(), "test-".to_string(), 2, 0); + let uris: Vec<_> = (0..4) + .map(|index| unique_uri(&format!("parallel-{index}"))) + .collect(); + for uri in &uris { + let _ = fs::remove_file(cache_path(uri)); + } + + assert_eq!( + rewrite_for_test(&coordinator, 1, &uris, Priority::Interactive).len(), + uris.len() + ); + assert_eq!(fetcher.max_active.load(Ordering::SeqCst), 2); + for uri in &uris { + let _ = fs::remove_file(cache_path(uri)); + } + coordinator.shutdown(); + } + + #[test] + fn deadline_returns_completed_replacements_only() { + let fetcher = Arc::new(FakeFetcher::new(Duration::from_millis(50))); + let coordinator = DecompileCoordinator::with_fetcher(fetcher, "test-".to_string(), 1, 0); + let uris = vec![unique_uri("deadline-a"), unique_uri("deadline-b")]; + for uri in &uris { + let _ = fs::remove_file(cache_path(uri)); + } + + let replacements = resolve_uris( + &coordinator.inner, + 1, + &uris, + Priority::Interactive, + Instant::now() + Duration::from_millis(70), + ); + + assert_eq!(replacements.len(), 1); + for uri in &uris { + let _ = fs::remove_file(cache_path(uri)); + } + coordinator.shutdown(); + } + + #[test] + fn canceling_one_waiter_keeps_shared_fetch_alive() { + let fetcher = Arc::new(FakeFetcher::new(Duration::from_millis(50))); + let coordinator = + DecompileCoordinator::with_fetcher(fetcher.clone(), "test-".to_string(), 1, 0); + let uri = unique_uri("shared-cancel"); + let _ = fs::remove_file(cache_path(&uri)); + + let first = { + let coordinator = coordinator.clone(); + let uri = uri.clone(); + thread::spawn(move || rewrite_for_test(&coordinator, 1, &[uri], Priority::Interactive)) + }; + let second = { + let coordinator = coordinator.clone(); + let uri = uri.clone(); + thread::spawn(move || rewrite_for_test(&coordinator, 2, &[uri], Priority::Interactive)) + }; + while { + let state = coordinator.inner.state.lock().unwrap(); + state + .uris + .get(&uri) + .is_none_or(|entry| entry.waiters.len() < 2) + } { + thread::yield_now(); + } + coordinator.cancel(1); + + assert!(first.join().unwrap().is_empty()); + assert_eq!(second.join().unwrap().len(), 1); + assert_eq!(fetcher.cancellations.load(Ordering::Relaxed), 0); + let _ = fs::remove_file(cache_path(&uri)); + coordinator.shutdown(); + } + + #[test] + fn negative_cache_prevents_immediate_retry_storms() { + let fetcher = Arc::new(FakeFetcher::failing(Duration::ZERO)); + let coordinator = + DecompileCoordinator::with_fetcher(fetcher.clone(), "test-".to_string(), 1, 0); + let uri = unique_uri("negative-cache"); + let _ = fs::remove_file(cache_path(&uri)); + + assert!(rewrite_for_test( + &coordinator, + 1, + std::slice::from_ref(&uri), + Priority::Interactive + ) + .is_empty()); + assert!(rewrite_for_test( + &coordinator, + 2, + std::slice::from_ref(&uri), + Priority::Interactive + ) + .is_empty()); + assert_eq!(fetcher.calls.load(Ordering::Relaxed), 1); + coordinator.shutdown(); + } + + #[test] + fn competing_cache_writes_never_expose_partial_content() { + let uri = unique_uri("atomic-cache"); + let path = cache_path(&uri); + let _ = fs::remove_file(&path); + let first_content = vec![b'a'; 32 * 1024]; + let second_content = vec![b'b'; 32 * 1024]; + + let first = { + let uri = uri.clone(); + let content = first_content.clone(); + thread::spawn(move || write_cached_source(&uri, &content)) + }; + let second = { + let uri = uri.clone(); + let content = second_content.clone(); + thread::spawn(move || write_cached_source(&uri, &content)) + }; + assert!(first.join().unwrap().is_some()); + assert!(second.join().unwrap().is_some()); + + let cached = fs::read(&path).unwrap(); + assert!(cached == first_content || cached == second_content); + let _ = fs::remove_file(path); + } + + #[test] + fn cache_names_use_stable_digest_and_sanitized_name() { + let directory = Path::new("/tmp/cache"); + let first = cache_path_in(directory, "jdt://contents/a/../../Bad Name.class"); + let second = cache_path_in(directory, "jdt://contents/a/../../Bad Name.class"); + + assert_eq!(first, second); + assert!(first + .file_name() + .unwrap() + .to_string_lossy() + .starts_with("Bad_Name-")); + assert_eq!(first.extension().unwrap(), "java"); + } + + #[test] + fn empty_sources_are_not_cached() { + let uri = unique_uri("empty"); + let path = cache_path(&uri); + let _ = fs::remove_file(&path); + + assert_eq!(write_cached_source(&uri, b""), None); + assert!(!path.exists()); } } diff --git a/proxy/src/http.rs b/proxy/src/http.rs index 33b4d3d..f4446b7 100644 --- a/proxy/src/http.rs +++ b/proxy/src/http.rs @@ -1,15 +1,15 @@ use serde::{Deserialize, Serialize}; use serde_json::Value; use std::{ - collections::HashMap, io::{BufRead, BufReader, Read, Write}, sync::{ atomic::{AtomicU64, Ordering}, - mpsc, Arc, Mutex, + Arc, Mutex, }, time::Duration, }; +use crate::pending::PendingResponses; use proxy_common::encode_lsp; pub const TIMEOUT: Duration = Duration::from_secs(5); @@ -31,9 +31,9 @@ struct LspRequest { pub fn handle_http( mut stream: std::net::TcpStream, writer: Arc>, - pending: Arc>>>, + pending: Arc, counter: Arc, - proxy_id: &str, + owned_id_prefix: &str, ) { let mut reader = BufReader::new(&stream); @@ -78,10 +78,8 @@ pub fn handle_http( }; let seq = counter.fetch_add(1, Ordering::Relaxed); - let id = Value::String(format!("{proxy_id}-{seq}")); - - let (tx, rx) = mpsc::channel(); - pending.lock().unwrap().insert(id.clone(), tx); + let id = Value::String(format!("{owned_id_prefix}http-{seq}")); + let rx = pending.register(id.clone()); let lsp_req = LspRequest { jsonrpc: "2.0", @@ -90,34 +88,45 @@ pub fn handle_http( params: req.params, }; let encoded = encode_lsp(&lsp_req); - { + let request_sent = { let mut w = writer.lock().unwrap(); - let _ = w.write_all(encoded.as_bytes()); - let _ = w.flush(); - } - - let response = match rx.recv_timeout(TIMEOUT) { - Ok(resp) => resp, - Err(_) => { - pending.lock().unwrap().remove(&id); - let cancel = encode_lsp(&serde_json::json!({ - "jsonrpc": "2.0", - "method": "$/cancelRequest", - "params": { "id": id } - })); - let mut w = writer.lock().unwrap(); - let _ = w.write_all(cancel.as_bytes()); - let _ = w.flush(); + w.write_all(encoded.as_bytes()).is_ok() && w.flush().is_ok() + }; - serde_json::json!({ - "jsonrpc": "2.0", - "id": id, - "error": { - "code": -32803, - "message": "Request to language server timed out after 5000ms." - } - }) + let response = if request_sent { + match rx.recv_timeout(TIMEOUT) { + Ok(resp) => resp, + Err(_) => { + pending.remove(&id); + let cancel = encode_lsp(&serde_json::json!({ + "jsonrpc": "2.0", + "method": "$/cancelRequest", + "params": { "id": id } + })); + let mut w = writer.lock().unwrap(); + let _ = w.write_all(cancel.as_bytes()); + let _ = w.flush(); + + serde_json::json!({ + "jsonrpc": "2.0", + "id": id, + "error": { + "code": -32803, + "message": "Request to language server timed out after 5000ms." + } + }) + } } + } else { + pending.remove(&id); + serde_json::json!({ + "jsonrpc": "2.0", + "id": id, + "error": { + "code": -32603, + "message": "Failed to write request to language server." + } + }) }; let resp_body = serde_json::to_vec(&response).unwrap(); diff --git a/proxy/src/log.rs b/proxy/src/log.rs index a93ae58..29b676d 100644 --- a/proxy/src/log.rs +++ b/proxy/src/log.rs @@ -1,7 +1,9 @@ use serde::Serialize; -use std::io::{self, Write}; +use std::sync::OnceLock; -use proxy_common::encode_lsp; +use crate::output::Output; + +static OUTPUT: OnceLock = OnceLock::new(); /// LSP `MessageType` constants as defined in the specification. /// https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#messageType @@ -30,8 +32,11 @@ struct LogMessageParams<'a> { /// Sends a `window/logMessage` LSP notification to stdout so that Zed /// displays the message in its Server Logs panel. /// -/// This locks stdout for the duration of the write to ensure the LSP -/// framing is not interleaved with other output. +/// The shared output broker serializes this with every other LSP frame. +pub fn init(output: Output) { + let _ = OUTPUT.set(output); +} + fn send_log_message(level: u8, message: &str) { let notification = LogMessageNotification { jsonrpc: "2.0", @@ -42,12 +47,9 @@ fn send_log_message(level: u8, message: &str) { }, }; - let encoded = encode_lsp(¬ification); - - let stdout = io::stdout(); - let mut w = stdout.lock(); - let _ = w.write_all(encoded.as_bytes()); - let _ = w.flush(); + if let Some(output) = OUTPUT.get() { + output.send_value(¬ification); + } } #[allow(dead_code)] diff --git a/proxy/src/main.rs b/proxy/src/main.rs index 1022cfe..9d2f676 100644 --- a/proxy/src/main.rs +++ b/proxy/src/main.rs @@ -2,14 +2,18 @@ mod completions; mod decompile; mod http; mod log; +mod output; +mod pending; -use completions::{is_completion_response, process_completions, sanitize_resolved_completion}; -use decompile::{rewrite_jdt_in_strings, rewrite_jdt_locations}; +use completions::{process_completions, sanitize_resolved_completion}; +use decompile::{DecompileCoordinator, Priority, RewriteJob, RewriteMode, SharedWriter}; use http::handle_http; +use output::Output; +use pending::PendingResponses; use proxy_common::{ - parse_lsp_content, raw_has_id, spawn_parent_monitor, write_raw, write_to_stdout, LspReader, + contains_subslice, parse_lsp_content, raw_has_id, spawn_parent_monitor, LspReader, }; -use serde_json::Value; +use serde_json::{json, Value}; use std::{ collections::HashMap, env, fs, @@ -19,18 +23,41 @@ use std::{ process::{self, Command, Stdio}, sync::{ atomic::{AtomicBool, AtomicU64, Ordering}, - mpsc, Arc, Mutex, + Arc, Mutex, OnceLock, }, thread, }; -#[derive(Clone, Copy)] -enum TrackedKind { - Definition, - Doc, +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum RewriteKind { + Locations, + Documentation, + Completion, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +struct TrackedRequest { + token: u64, + method: String, + rewrite: RewriteKind, + original_params: Option, +} + +impl TrackedRequest { + fn new(token: u64, method: &str, rewrite: RewriteKind, original_params: Option) -> Self { + Self { + token, + method: method.to_string(), + rewrite, + original_params, + } + } } fn main() { + let output = Output::start(); + log::init(output.clone()); + let args: Vec = env::args().skip(1).collect(); if args.len() < 2 { @@ -78,12 +105,17 @@ fn main() { lsp_info!("JDTLS process spawned (pid={})", child.id()); - let child_stdin = Arc::new(Mutex::new(child.stdin.take().unwrap())); + let child_stdin: SharedWriter = Arc::new(Mutex::new(Box::new(child.stdin.take().unwrap()))); let child_stdout = child.stdout.take().unwrap(); let alive = Arc::new(AtomicBool::new(true)); - let pending: Arc>>> = - Arc::new(Mutex::new(HashMap::new())); + let owned_id_prefix = format!("{proxy_id}-proxy-"); + let pending = Arc::new(PendingResponses::new(owned_id_prefix.clone())); + let decompile = DecompileCoordinator::new( + Arc::clone(&child_stdin), + Arc::clone(&pending), + owned_id_prefix.clone(), + ); let listener = TcpListener::bind("127.0.0.1:0").unwrap(); let port = listener.local_addr().unwrap().port(); @@ -95,43 +127,68 @@ fn main() { lsp_info!("HTTP server listening on 127.0.0.1:{port}"); let id_counter = Arc::new(AtomicU64::new(1)); + let job_counter = Arc::new(AtomicU64::new(1)); - // Track definition/typeDefinition/implementation and documentation request IDs - // so their responses can be intercepted and rewritten. - let tracked_ids: Arc>> = Arc::new(Mutex::new(HashMap::new())); + // Track requests whose responses may contain jdt:// URIs so they can be + // intercepted and rewritten. + let tracked_ids: Arc>> = + Arc::new(Mutex::new(HashMap::new())); + let active_rewrites: Arc>> = Arc::new(Mutex::new(HashMap::new())); // --- Thread 1: Zed stdin -> JDTLS stdin (track definition requests) --- let stdin_writer = Arc::clone(&child_stdin); let alive_stdin = Arc::clone(&alive); let tracked_in = Arc::clone(&tracked_ids); + let active_in = Arc::clone(&active_rewrites); + let jobs_in = Arc::clone(&job_counter); + let decompile_in = decompile.clone(); + let latest_workspace_job = Arc::new(Mutex::new(None::)); + let latest_workspace_in = Arc::clone(&latest_workspace_job); thread::spawn(move || { let stdin = io::stdin().lock(); let mut reader = LspReader::new(BufReader::new(stdin)); while alive_stdin.load(Ordering::Relaxed) { match reader.read_message() { Ok(Some(raw)) => { - // Only requests (not notifications) carry an `id`; skip the - // JSON parse entirely for high-volume notifications like - // textDocument/didChange. - if raw_has_id(&raw) { - if let Some(msg) = parse_lsp_content(&raw) { - if let Some(method) = msg.get("method").and_then(|m| m.as_str()) { - let kind = match method { - "textDocument/definition" - | "textDocument/typeDefinition" - | "textDocument/implementation" => { - Some(TrackedKind::Definition) - } - "textDocument/hover" - | "textDocument/signatureHelp" - | "completionItem/resolve" => Some(TrackedKind::Doc), - _ => None, - }; - if let Some(kind) = kind { - if let Some(id) = msg.get("id").cloned() { - tracked_in.lock().unwrap().insert(id, kind); + let should_parse = + raw_has_id(&raw) || contains_subslice(&raw, b"$/cancelRequest"); + if should_parse { + let Some(msg) = parse_lsp_content(&raw) else { + let mut writer = stdin_writer.lock().unwrap(); + if writer.write_all(&raw).is_err() || writer.flush().is_err() { + break; + } + continue; + }; + if msg.get("method").and_then(Value::as_str) == Some("$/cancelRequest") { + if let Some(id) = msg.pointer("/params/id") { + let request = tracked_in.lock().unwrap().get(id).cloned(); + let active_token = active_in.lock().unwrap().remove(id); + let token = request + .as_ref() + .map(|request| request.token) + .or(active_token); + if let Some(token) = token { + clear_latest_workspace(&latest_workspace_in, token); + decompile_in.cancel(token); + } + } + } else if raw_has_id(&raw) { + let token = jobs_in.fetch_add(1, Ordering::Relaxed); + if let Some((id, request)) = tracked_request_for(&msg, token) { + if request.method == "workspace/symbol" { + let previous = + latest_workspace_in.lock().unwrap().replace(token); + if let Some(previous) = previous { + remove_active_token(&active_in, previous); + decompile_in.cancel(previous); } } + let previous = tracked_in.lock().unwrap().insert(id, request); + if let Some(previous) = previous { + remove_active_token(&active_in, previous.token); + decompile_in.cancel(previous.token); + } } } } @@ -150,10 +207,10 @@ fn main() { let pending_out = Arc::clone(&pending); let alive_out = Arc::clone(&alive); let tracked_out = Arc::clone(&tracked_ids); - let decompile_writer = Arc::clone(&child_stdin); - let decompile_pending = Arc::clone(&pending); - let decompile_counter = Arc::clone(&id_counter); - let decompile_proxy_id = proxy_id.clone(); + let active_out = Arc::clone(&active_rewrites); + let decompile_out = decompile.clone(); + let output_router = output.clone(); + let latest_workspace_out = Arc::clone(&latest_workspace_job); thread::spawn(move || { let mut reader = LspReader::new(BufReader::new(child_stdout)); while alive_out.load(Ordering::Relaxed) { @@ -162,71 +219,122 @@ fn main() { // Fast path: notifications (no `id`) can't be responses we // need to intercept. Forward the raw bytes without parsing. if !raw_has_id(&raw) { - write_raw(&mut io::stdout().lock(), &raw); + output_router.send_raw(raw); continue; } let Some(mut msg) = parse_lsp_content(&raw) else { - write_raw(&mut io::stdout().lock(), &raw); + output_router.send_raw(raw); continue; }; // Route responses to pending HTTP requests - if let Some(id) = msg.get("id") { - if let Some(tx) = pending_out.lock().unwrap().remove(id) { - let _ = tx.send(msg); - continue; - } + if pending_out.route(&msg) { + continue; } - // Rewrite jdt:// URIs in definition or documentation responses. - // Spawns a thread so this loop stays unblocked and can route - // the java/classFileContents response back via `pending`. - if let Some(id) = msg.get("id").cloned() { - if let Some(kind) = tracked_out.lock().unwrap().remove(&id) { - let writer = Arc::clone(&decompile_writer); - let pending = Arc::clone(&decompile_pending); - let pid = decompile_proxy_id.clone(); - let counter = Arc::clone(&decompile_counter); - thread::spawn(move || { - let mut next_id = move || { - let seq = counter.fetch_add(1, Ordering::Relaxed); - Value::String(format!("{pid}-decompile-{seq}")) - }; - match kind { - TrackedKind::Definition => { - rewrite_jdt_locations( - &mut msg, - &writer, - &pending, - &mut next_id, - ); + // Rewrite jdt:// URIs in location or documentation responses. + // The bounded coordinator keeps this router free to deliver + // java/classFileContents responses through `pending`. + if msg.get("method").is_none() { + let Some(id) = msg.get("id").cloned() else { + output_router.send_raw(raw); + continue; + }; + let request = tracked_out.lock().unwrap().get(&id).cloned(); + if let Some(request) = request { + let canceled = decompile_out.is_canceled(request.token); + if canceled { + tracked_out.lock().unwrap().remove(&id); + decompile_out.consume_cancellation(request.token); + clear_latest_workspace(&latest_workspace_out, request.token); + continue; + } + if let Some(fallback) = completion_resolve_fallback(&msg, &request) { + tracked_out.lock().unwrap().remove(&id); + if should_log_completion_fallback() { + lsp_warn!( + "JDTLS completion resolution failed with -32603; \ + using the unresolved item, so documentation, imports, \ + commands, or additional edits may be missing" + ); + } + output_router.send_value(&fallback); + continue; + } + if msg.get("error").is_some() { + tracked_out.lock().unwrap().remove(&id); + clear_latest_workspace(&latest_workspace_out, request.token); + output_router.send_raw(raw); + continue; + } + if request.rewrite == RewriteKind::Completion { + tracked_out.lock().unwrap().remove(&id); + process_completions(&mut msg); + output_router.send_value(&msg); + continue; + } + + let output = output_router.clone(); + let sanitize_completion = request.method == "completionItem/resolve"; + let mode = match request.rewrite { + RewriteKind::Locations => RewriteMode::Locations, + RewriteKind::Documentation => RewriteMode::Strings, + RewriteKind::Completion => unreachable!(), + }; + let priority = if request.method == "workspace/symbol" { + Priority::Bulk + } else { + Priority::Interactive + }; + let deadline = + std::time::Instant::now() + rewrite_timeout(&request.method); + let workspace_job = (request.method == "workspace/symbol") + .then(|| (Arc::clone(&latest_workspace_out), request.token)); + active_out.lock().unwrap().insert(id.clone(), request.token); + tracked_out.lock().unwrap().remove(&id); + if decompile_out.is_canceled(request.token) { + active_out.lock().unwrap().remove(&id); + decompile_out.consume_cancellation(request.token); + clear_latest_workspace(&latest_workspace_out, request.token); + continue; + } + let active = Arc::clone(&active_out); + let active_id = id; + let completion_coordinator = decompile_out.clone(); + let job = RewriteJob { + token: request.token, + message: msg, + mode, + priority, + deadline, + complete: Box::new(move |mut message| { + if sanitize_completion { + sanitize_resolved_completion(&mut message); } - TrackedKind::Doc => { - rewrite_jdt_in_strings( - &mut msg, - &writer, - &pending, - &mut next_id, - ); - sanitize_resolved_completion(&mut msg); + let owns_response = active.lock().unwrap().remove(&active_id) + == Some(request.token); + completion_coordinator.consume_cancellation(request.token); + if owns_response { + output.send_value(&message); } - } - write_to_stdout(&msg); - }); + if let Some((latest, token)) = workspace_job { + clear_latest_workspace(&latest, token); + } + }), + }; + if let Err(job) = decompile_out.submit(job) { + let RewriteJob { + message, complete, .. + } = job; + complete(message); + } continue; } } - // Process completion responses (sort + sanitize) in a single pass - if is_completion_response(&msg) { - process_completions(&mut msg); - write_to_stdout(&msg); - continue; - } - // Passthrough - write_raw(&mut io::stdout().lock(), &raw); + output_router.send_raw(raw); } Ok(None) | Err(_) => break, } @@ -239,7 +347,7 @@ fn main() { let http_pending = Arc::clone(&pending); let http_alive = Arc::clone(&alive); let http_id_counter = Arc::clone(&id_counter); - let http_proxy_id = proxy_id.clone(); + let http_owned_id_prefix = owned_id_prefix.clone(); thread::spawn(move || { for stream in listener.incoming() { if !http_alive.load(Ordering::Relaxed) { @@ -249,10 +357,10 @@ fn main() { let writer = Arc::clone(&http_writer); let pend = Arc::clone(&http_pending); let counter = Arc::clone(&http_id_counter); - let pid = http_proxy_id.clone(); + let owned_id_prefix = http_owned_id_prefix.clone(); thread::spawn(move || { - handle_http(stream, writer, pend, counter, &pid); + handle_http(stream, writer, pend, counter, &owned_id_prefix); }); } }); @@ -260,10 +368,25 @@ fn main() { // --- Thread 4: Parent process monitor --- spawn_parent_monitor(Arc::clone(&alive), child.id()); - // Wait for child to exit - let status = child.wait(); + // Poll so broken editor/JDTLS transport can terminate the child and release + // all pending proxy jobs instead of waiting indefinitely in `Child::wait`. + let status = loop { + match child.try_wait() { + Ok(Some(status)) => break Ok(status), + Ok(None) if alive.load(Ordering::Relaxed) && !output.failed() => { + thread::sleep(std::time::Duration::from_millis(50)); + } + Ok(None) => { + let _ = child.kill(); + break child.wait(); + } + Err(error) => break Err(error), + } + }; lsp_info!("JDTLS process exited: {status:?}"); alive.store(false, Ordering::Relaxed); + decompile.shutdown(); + pending.clear(); let _ = fs::remove_file(&port_file); } @@ -272,3 +395,243 @@ fn main() { fn hex_encode(s: &str) -> String { s.as_bytes().iter().map(|b| format!("{b:02x}")).collect() } + +fn clear_latest_workspace(latest: &Mutex>, token: u64) { + let mut latest = latest.lock().unwrap(); + if *latest == Some(token) { + *latest = None; + } +} + +fn remove_active_token(active: &Mutex>, token: u64) { + active + .lock() + .unwrap() + .retain(|_, active_token| *active_token != token); +} + +fn should_log_completion_fallback() -> bool { + static LAST_WARNING: OnceLock>> = OnceLock::new(); + let now = std::time::Instant::now(); + let mut last = LAST_WARNING + .get_or_init(|| Mutex::new(None)) + .lock() + .unwrap(); + if last + .is_some_and(|previous| now.duration_since(previous) < std::time::Duration::from_secs(60)) + { + return false; + } + *last = Some(now); + true +} + +fn tracked_request_for(msg: &Value, token: u64) -> Option<(Value, TrackedRequest)> { + let method = msg.get("method")?.as_str()?; + let id = msg.get("id")?.clone(); + let rewrite = match method { + "textDocument/completion" => RewriteKind::Completion, + "textDocument/definition" + | "textDocument/declaration" + | "textDocument/typeDefinition" + | "textDocument/implementation" + | "textDocument/references" + | "textDocument/prepareCallHierarchy" + | "callHierarchy/incomingCalls" + | "callHierarchy/outgoingCalls" + | "textDocument/prepareTypeHierarchy" + | "typeHierarchy/supertypes" + | "typeHierarchy/subtypes" + | "workspace/symbol" => RewriteKind::Locations, + "textDocument/hover" | "textDocument/signatureHelp" | "completionItem/resolve" => { + RewriteKind::Documentation + } + _ => return None, + }; + + let original_params = if method == "completionItem/resolve" { + msg.get("params") + .filter(|params| params.is_object()) + .cloned() + } else { + None + }; + + Some(( + id, + TrackedRequest::new(token, method, rewrite, original_params), + )) +} + +fn completion_resolve_fallback(message: &Value, request: &TrackedRequest) -> Option { + if request.method != "completionItem/resolve" + || message.pointer("/error/code").and_then(Value::as_i64) != Some(-32603) + { + return None; + } + let item = request.original_params.as_ref()?; + Some(json!({ + "jsonrpc": "2.0", + "id": message.get("id")?.clone(), + "result": item, + })) +} + +fn rewrite_timeout(method: &str) -> std::time::Duration { + match method { + "textDocument/hover" | "textDocument/signatureHelp" | "completionItem/resolve" => { + std::time::Duration::from_millis(500) + } + "workspace/symbol" => std::time::Duration::from_secs(1), + _ => std::time::Duration::from_secs(2), + } +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + #[test] + fn tracks_location_response_methods() { + let methods = [ + "textDocument/definition", + "textDocument/declaration", + "textDocument/typeDefinition", + "textDocument/implementation", + "textDocument/references", + "textDocument/prepareCallHierarchy", + "callHierarchy/incomingCalls", + "callHierarchy/outgoingCalls", + "textDocument/prepareTypeHierarchy", + "typeHierarchy/supertypes", + "typeHierarchy/subtypes", + "workspace/symbol", + ]; + + for method in methods { + let request = json!({ + "jsonrpc": "2.0", + "id": 7, + "method": method, + "params": {} + }); + + assert_eq!( + tracked_request_for(&request, 42), + Some(( + json!(7), + TrackedRequest::new(42, method, RewriteKind::Locations, None) + )), + "{method} should rewrite location URIs" + ); + } + } + + #[test] + fn tracks_completion_resolve_documentation() { + let request = json!({ + "jsonrpc": "2.0", + "id": "resolve-1", + "method": "completionItem/resolve", + "params": { "label": "String" } + }); + + assert_eq!( + tracked_request_for(&request, 43), + Some(( + json!("resolve-1"), + TrackedRequest::new( + 43, + "completionItem/resolve", + RewriteKind::Documentation, + Some(json!({ "label": "String" })) + ) + )) + ); + } + + #[test] + fn ignores_untracked_requests_and_notifications() { + let request = json!({ + "jsonrpc": "2.0", + "id": 8, + "method": "textDocument/rename", + "params": {} + }); + let notification = json!({ + "jsonrpc": "2.0", + "method": "textDocument/didChange", + "params": {} + }); + + assert_eq!(tracked_request_for(&request, 1), None); + assert_eq!(tracked_request_for(¬ification, 2), None); + } + + #[test] + fn tracks_completion_responses_by_method() { + let request = json!({ + "jsonrpc": "2.0", + "id": 9, + "method": "textDocument/completion", + "params": {} + }); + + assert_eq!( + tracked_request_for(&request, 44), + Some(( + json!(9), + TrackedRequest::new(44, "textDocument/completion", RewriteKind::Completion, None) + )) + ); + } + + #[test] + fn falls_back_only_for_internal_completion_resolve_errors() { + let request = TrackedRequest::new( + 1, + "completionItem/resolve", + RewriteKind::Documentation, + Some(json!({ "label": "value", "kind": 6 })), + ); + let internal = json!({ + "jsonrpc": "2.0", + "id": 5, + "error": { "code": -32603, "message": "Invalid completion proposal" } + }); + let canceled = json!({ + "jsonrpc": "2.0", + "id": 5, + "error": { "code": -32800, "message": "Request cancelled" } + }); + + assert_eq!( + completion_resolve_fallback(&internal, &request), + Some(json!({ + "jsonrpc": "2.0", + "id": 5, + "result": { "label": "value", "kind": 6 } + })) + ); + assert_eq!(completion_resolve_fallback(&canceled, &request), None); + } + + #[test] + fn malformed_completion_resolve_params_do_not_fallback() { + let request = json!({ + "jsonrpc": "2.0", + "id": 5, + "method": "completionItem/resolve", + "params": null + }); + let (_, tracked) = tracked_request_for(&request, 1).unwrap(); + let error = json!({ + "jsonrpc": "2.0", + "id": 5, + "error": { "code": -32603, "message": "Internal error" } + }); + + assert_eq!(completion_resolve_fallback(&error, &tracked), None); + } +} diff --git a/proxy/src/output.rs b/proxy/src/output.rs new file mode 100644 index 0000000..674f2c8 --- /dev/null +++ b/proxy/src/output.rs @@ -0,0 +1,122 @@ +use proxy_common::encode_lsp; +use serde::Serialize; +use std::{ + io::{self, Write}, + sync::{ + atomic::{AtomicBool, Ordering}, + mpsc, Arc, + }, + thread, +}; + +#[derive(Clone)] +pub struct Output { + sender: mpsc::Sender>, + failed: Arc, +} + +impl Output { + pub fn start() -> Self { + Self::start_with_writer(io::stdout()) + } + + fn start_with_writer(writer: impl Write + Send + 'static) -> Self { + let (sender, receiver) = mpsc::channel(); + let failed = Arc::new(AtomicBool::new(false)); + let writer_failed = Arc::clone(&failed); + thread::spawn(move || write_loop(writer, receiver, writer_failed)); + Self { sender, failed } + } + + pub fn send_raw(&self, raw: Vec) -> bool { + let sent = self.sender.send(raw).is_ok(); + if !sent { + self.failed.store(true, Ordering::Relaxed); + } + sent + } + + pub fn send_value(&self, value: &impl Serialize) -> bool { + self.send_raw(encode_lsp(value).into_bytes()) + } + + pub fn failed(&self) -> bool { + self.failed.load(Ordering::Relaxed) + } +} + +fn write_loop(mut writer: impl Write, receiver: mpsc::Receiver>, failed: Arc) { + while let Ok(message) = receiver.recv() { + if writer.write_all(&message).is_err() || writer.flush().is_err() { + failed.store(true, Ordering::Relaxed); + break; + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::{Arc, Mutex}; + + #[derive(Clone, Default)] + struct SharedWriter(Arc>>); + + impl Write for SharedWriter { + fn write(&mut self, buf: &[u8]) -> io::Result { + self.0.lock().unwrap().extend_from_slice(buf); + Ok(buf.len()) + } + + fn flush(&mut self) -> io::Result<()> { + Ok(()) + } + } + + struct FailingWriter; + + impl Write for FailingWriter { + fn write(&mut self, _buf: &[u8]) -> io::Result { + Err(io::Error::new(io::ErrorKind::BrokenPipe, "closed")) + } + + fn flush(&mut self) -> io::Result<()> { + Ok(()) + } + } + + #[test] + fn serializes_complete_frames_in_submission_order() { + let writer = SharedWriter::default(); + let bytes = Arc::clone(&writer.0); + let output = Output::start_with_writer(writer); + + assert!(output.send_raw(b"first".to_vec())); + assert!(output.send_raw(b"second".to_vec())); + drop(output); + + for _ in 0..100 { + if bytes.lock().unwrap().len() == 11 { + break; + } + thread::sleep(std::time::Duration::from_millis(1)); + } + + assert_eq!(*bytes.lock().unwrap(), b"firstsecond"); + } + + #[test] + fn reports_writer_failure() { + let output = Output::start_with_writer(FailingWriter); + assert!(output.send_raw(b"message".to_vec())); + + for _ in 0..100 { + if output.failed() { + break; + } + thread::sleep(std::time::Duration::from_millis(1)); + } + + assert!(output.failed()); + } +} diff --git a/proxy/src/pending.rs b/proxy/src/pending.rs new file mode 100644 index 0000000..84bf043 --- /dev/null +++ b/proxy/src/pending.rs @@ -0,0 +1,93 @@ +use serde_json::Value; +use std::{ + collections::HashMap, + sync::{mpsc, Mutex}, +}; + +pub struct PendingResponses { + owned_id_prefix: String, + pending: Mutex>>, +} + +impl PendingResponses { + pub fn new(owned_id_prefix: String) -> Self { + Self { + owned_id_prefix, + pending: Mutex::new(HashMap::new()), + } + } + + pub fn register(&self, id: Value) -> mpsc::Receiver { + let (sender, receiver) = mpsc::channel(); + self.pending.lock().unwrap().insert(id, sender); + receiver + } + + pub fn remove(&self, id: &Value) { + self.pending.lock().unwrap().remove(id); + } + + /// Returns true when the response belongs to the proxy and must not be + /// forwarded to the editor, including responses that arrive after timeout. + pub fn route(&self, message: &Value) -> bool { + if message.get("method").is_some() { + return false; + } + let Some(id) = message.get("id") else { + return false; + }; + + let sender = self.pending.lock().unwrap().remove(id); + if let Some(sender) = sender { + let _ = sender.send(message.clone()); + return true; + } + + id.as_str() + .is_some_and(|id| id.starts_with(&self.owned_id_prefix)) + } + + pub fn clear(&self) { + self.pending.lock().unwrap().clear(); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + #[test] + fn routes_each_pending_response_once() { + let pending = PendingResponses::new("proxy-owned-".to_string()); + let id = json!("proxy-owned-1"); + let receiver = pending.register(id.clone()); + let response = json!({ "jsonrpc": "2.0", "id": id, "result": "ok" }); + + assert!(pending.route(&response)); + assert_eq!(receiver.recv().unwrap(), response); + assert!(pending.route(&response)); + } + + #[test] + fn swallows_late_owned_responses_but_not_client_responses() { + let pending = PendingResponses::new("proxy-owned-".to_string()); + + assert!(pending.route(&json!({ + "jsonrpc": "2.0", + "id": "proxy-owned-late", + "result": null + }))); + assert!(!pending.route(&json!({ + "jsonrpc": "2.0", + "id": 7, + "result": null + }))); + assert!(!pending.route(&json!({ + "jsonrpc": "2.0", + "id": "proxy-owned-server-request", + "method": "workspace/applyEdit", + "params": {} + }))); + } +} From a183762841034a393fdda1b92760b099c99a56f9 Mon Sep 17 00:00:00 2001 From: Riccardo Strina Date: Sat, 1 Aug 2026 23:05:16 +0200 Subject: [PATCH 2/7] fix(proxy): Harden decompilation lifecycle --- Cargo.lock | 1 + proxy-common/Cargo.toml | 1 + proxy-common/src/uri.rs | 93 ++++++-- proxy/src/decompile.rs | 466 ++++++++++++++++++++++++++++++++++------ proxy/src/main.rs | 442 +++++++++++++++++++++++++------------ proxy/src/output.rs | 128 +++++++++-- proxy/src/pending.rs | 73 +++++-- 7 files changed, 953 insertions(+), 251 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 2730aca..eb5d51b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -642,6 +642,7 @@ name = "proxy-common" version = "6.8.23" dependencies = [ "libc", + "percent-encoding", "serde", "serde_json", "tokio", diff --git a/proxy-common/Cargo.toml b/proxy-common/Cargo.toml index 08d80ce..4de15ac 100644 --- a/proxy-common/Cargo.toml +++ b/proxy-common/Cargo.toml @@ -14,6 +14,7 @@ path = "src/lib.rs" tokio = ["dep:tokio"] [dependencies] +percent-encoding = "2.3" serde = { version = "1.0", features = ["derive"] } serde_json = "1.0" tokio = { version = "1", default-features = false, features = ["io-util"], optional = true } diff --git a/proxy-common/src/uri.rs b/proxy-common/src/uri.rs index ffa0864..fead948 100644 --- a/proxy-common/src/uri.rs +++ b/proxy-common/src/uri.rs @@ -1,21 +1,84 @@ +use percent_encoding::{utf8_percent_encode, AsciiSet, NON_ALPHANUMERIC}; use std::path::Path; -/// Convert a filesystem path to a `file://` URI, matching how language servers' -/// `publishDiagnostics` and the editor key documents. -/// -/// On Unix the path already starts with `/`, so `file://` + path gives the -/// correct `file:///…` form with no extra work. -/// -/// On Windows the backslashes are replaced with `/` and an extra `/` is -/// prepended before the drive letter, so we get `file:///C:/…` rather than -/// `file://C:\…`. -#[cfg(unix)] +const PATH_ENCODE_SET: AsciiSet = NON_ALPHANUMERIC + .remove(b'/') + .remove(b':') + .remove(b'-') + .remove(b'.') + .remove(b'_') + .remove(b'~') + .remove(b'@'); + +/// Convert a filesystem path to an RFC 3986 percent-encoded `file://` URI. pub fn path_to_file_uri(path: &Path) -> String { - format!("file://{}", path.display()) + file_uri_from_path_string(&path.to_string_lossy(), cfg!(windows)) } -#[cfg(windows)] -pub fn path_to_file_uri(path: &Path) -> String { - let s = path.display().to_string().replace('\\', "/"); - format!("file:///{s}") +fn file_uri_from_path_string(path: &str, windows: bool) -> String { + let mut normalized = if windows { + path.replace('\\', "/") + } else { + path.to_string() + }; + if windows { + if let Some(unc) = normalized.strip_prefix("//?/UNC/") { + normalized = format!("//{unc}"); + } else if let Some(verbatim) = normalized.strip_prefix("//?/") { + normalized = verbatim.to_string(); + } + if let Some(unc) = normalized.strip_prefix("//") { + return format!("file://{}", utf8_percent_encode(unc, &PATH_ENCODE_SET)); + } + } + let prefix = if normalized.starts_with('/') { + "file://" + } else { + "file:///" + }; + format!( + "{prefix}{}", + utf8_percent_encode(&normalized, &PATH_ENCODE_SET) + ) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn encodes_reserved_characters_in_unix_paths() { + assert_eq!( + file_uri_from_path_string("/tmp/Java Sources/#1%?.java", false), + "file:///tmp/Java%20Sources/%231%25%3F.java" + ); + } + + #[test] + fn normalizes_and_encodes_windows_paths() { + assert_eq!( + file_uri_from_path_string(r"C:\Users\Jane Doe\A#1.java", true), + "file:///C:/Users/Jane%20Doe/A%231.java" + ); + } + + #[test] + fn preserves_windows_unc_authority() { + assert_eq!( + file_uri_from_path_string(r"\\server\share\A File.java", true), + "file://server/share/A%20File.java" + ); + } + + #[test] + fn normalizes_windows_verbatim_paths() { + assert_eq!( + file_uri_from_path_string(r"\\?\C:\Users\Jane Doe\A.java", true), + "file:///C:/Users/Jane%20Doe/A.java" + ); + assert_eq!( + file_uri_from_path_string(r"\\?\UNC\server\share\A.java", true), + "file://server/share/A.java" + ); + } } diff --git a/proxy/src/decompile.rs b/proxy/src/decompile.rs index b217402..87995bc 100644 --- a/proxy/src/decompile.rs +++ b/proxy/src/decompile.rs @@ -12,8 +12,8 @@ use std::{ atomic::{AtomicU64, Ordering}, Arc, Condvar, Mutex, }, - thread, - time::{Duration, Instant}, + thread::{self, JoinHandle}, + time::{Duration, Instant, SystemTime, UNIX_EPOCH}, }; const DECOMPILED_DIR: &str = "jdtls-decompiled"; @@ -54,11 +54,19 @@ pub struct DecompileCoordinator { struct Inner { fetcher: Arc, owned_id_prefix: String, + cache_dir: PathBuf, request_counter: AtomicU64, max_bulk_jobs: usize, state: Mutex, work_available: Condvar, state_changed: Condvar, + workers: Mutex>>, +} + +impl Drop for Inner { + fn drop(&mut self) { + let _ = fs::remove_dir_all(&self.cache_dir); + } } #[derive(Default)] @@ -69,7 +77,6 @@ struct State { interactive_jobs: VecDeque, bulk_jobs: VecDeque, canceled_jobs: HashSet, - latest_bulk_job: Option, bulk_jobs_active: usize, bulk_fetches: usize, shutdown: bool, @@ -162,26 +169,33 @@ impl DecompileCoordinator { fetch_workers: usize, job_workers: usize, ) -> Self { + let cache_dir = session_cache_dir(&owned_id_prefix); let coordinator = Self { inner: Arc::new(Inner { fetcher, owned_id_prefix, + cache_dir, request_counter: AtomicU64::new(1), max_bulk_jobs: job_workers.saturating_sub(1).max(1), state: Mutex::new(State::default()), work_available: Condvar::new(), state_changed: Condvar::new(), + workers: Mutex::new(Vec::new()), }), }; + let mut workers = Vec::with_capacity(fetch_workers + job_workers + 1); for _ in 0..fetch_workers { let inner = Arc::clone(&coordinator.inner); - thread::spawn(move || fetch_worker(inner)); + workers.push(thread::spawn(move || fetch_worker(inner))); } for _ in 0..job_workers { let inner = Arc::clone(&coordinator.inner); - thread::spawn(move || job_worker(inner)); + workers.push(thread::spawn(move || job_worker(inner))); } + let inner = Arc::clone(&coordinator.inner); + workers.push(thread::spawn(move || deadline_worker(inner))); + *coordinator.inner.workers.lock().unwrap() = workers; coordinator } @@ -196,11 +210,6 @@ impl DecompileCoordinator { } if job.priority == Priority::Bulk { - if let Some(previous) = state.latest_bulk_job.replace(job.token) { - if cancel_job_locked(&mut state, previous) { - state.canceled_jobs.remove(&previous); - } - } state.bulk_jobs.push_back(job); } else { state.interactive_jobs.push_back(job); @@ -215,7 +224,7 @@ impl DecompileCoordinator { if cancel_job_locked(&mut state, token) { state.canceled_jobs.remove(&token); } - orphaned_request_ids(&state) + detach_orphaned_fetches(&mut state) }; for request_id in request_ids { self.inner.fetcher.cancel(&request_id); @@ -243,12 +252,34 @@ impl DecompileCoordinator { } pub fn shutdown(&self) { - let mut state = self.inner.state.lock().unwrap(); - state.shutdown = true; - state.interactive_jobs.clear(); - state.bulk_jobs.clear(); + let queued = { + let mut state = self.inner.state.lock().unwrap(); + if state.shutdown { + Vec::new() + } else { + state.shutdown = true; + let mut queued: Vec<_> = state.interactive_jobs.drain(..).collect(); + queued.extend(state.bulk_jobs.drain(..)); + queued + } + }; self.inner.work_available.notify_all(); self.inner.state_changed.notify_all(); + + for job in queued { + if !self.consume_cancellation(job.token) { + (job.complete)(job.message); + } + } + + let workers = std::mem::take(&mut *self.inner.workers.lock().unwrap()); + for worker in workers { + let _ = worker.join(); + } + } + + pub fn cleanup_cache(&self) { + let _ = fs::remove_dir_all(&self.inner.cache_dir); } } @@ -265,16 +296,28 @@ fn cancel_job_locked(state: &mut State, token: u64) -> bool { queued_jobs != state.interactive_jobs.len() + state.bulk_jobs.len() } -fn orphaned_request_ids(state: &State) -> Vec { - state +fn detach_orphaned_fetches(state: &mut State) -> Vec { + let now = Instant::now(); + let orphaned: Vec<_> = state .uris - .values() - .filter(|entry| entry.waiters.is_empty()) - .filter_map(|entry| match &entry.status { - UriStatus::InFlight { request_id, .. } => Some(request_id.clone()), - _ => None, + .iter() + .filter(|(_, entry)| entry.waiters.is_empty()) + .filter(|(_, entry)| { + !matches!(entry.status, UriStatus::Failed(retry_after) if retry_after > now) }) - .collect() + .map(|(uri, _)| uri.clone()) + .collect(); + let mut request_ids = Vec::new(); + for uri in orphaned { + if let Some(UriEntry { + status: UriStatus::InFlight { request_id }, + .. + }) = state.uris.remove(&uri) + { + request_ids.push(request_id); + } + } + request_ids } fn job_worker(inner: Arc) { @@ -304,14 +347,16 @@ fn job_worker(inner: Arc) { } let mut message = job.message; - rewrite_message( - &inner, - job.token, - &mut message, - job.mode, - job.priority, - job.deadline, - ); + if Instant::now() < job.deadline { + rewrite_message( + &inner, + job.token, + &mut message, + job.mode, + job.priority, + job.deadline, + ); + } let canceled = is_canceled(&inner, job.token); finish_job(&inner, job.token, job.priority); if !canceled { @@ -320,6 +365,57 @@ fn job_worker(inner: Arc) { } } +fn deadline_worker(inner: Arc) { + loop { + let expired = { + let mut state = inner.state.lock().unwrap(); + loop { + if state.shutdown { + return; + } + + let now = Instant::now(); + let mut expired = take_expired_jobs(&mut state.interactive_jobs, now); + expired.extend(take_expired_jobs(&mut state.bulk_jobs, now)); + if !expired.is_empty() { + break expired; + } + + let next_deadline = state + .interactive_jobs + .iter() + .chain(state.bulk_jobs.iter()) + .map(|job| job.deadline) + .min(); + state = if let Some(deadline) = next_deadline { + let timeout = deadline.saturating_duration_since(Instant::now()); + inner.work_available.wait_timeout(state, timeout).unwrap().0 + } else { + inner.work_available.wait(state).unwrap() + }; + } + }; + + for job in expired { + (job.complete)(job.message); + } + } +} + +fn take_expired_jobs(queue: &mut VecDeque, now: Instant) -> Vec { + let mut expired = Vec::new(); + let queued = queue.len(); + for _ in 0..queued { + let job = queue.pop_front().unwrap(); + if job.deadline <= now { + expired.push(job); + } else { + queue.push_back(job); + } + } + expired +} + fn is_canceled(inner: &Inner, token: u64) -> bool { inner.state.lock().unwrap().canceled_jobs.contains(&token) } @@ -330,9 +426,6 @@ fn finish_job(inner: &Inner, token: u64, priority: Priority) { state.bulk_jobs_active = state.bulk_jobs_active.saturating_sub(1); } state.canceled_jobs.remove(&token); - if state.latest_bulk_job == Some(token) { - state.latest_bulk_job = None; - } for entry in state.uris.values_mut() { entry.waiters.remove(&token); } @@ -386,7 +479,7 @@ fn resolve_uris( let mut unresolved = Vec::new(); for uri in uris { - let path = cache_path(uri); + let path = cache_path_in(&inner.cache_dir, uri); if path.is_file() { replacements.insert(uri.clone(), path_to_file_uri(&path)); } else { @@ -457,7 +550,11 @@ fn resolve_uris( } } - if !waiting || state.canceled_jobs.contains(&token) || Instant::now() >= deadline { + if !waiting + || state.shutdown + || state.canceled_jobs.contains(&token) + || Instant::now() >= deadline + { break; } @@ -471,7 +568,7 @@ fn resolve_uris( entry.waiters.remove(&token); } } - let orphaned = orphaned_request_ids(&state); + let orphaned = detach_orphaned_fetches(&mut state); drop(state); for request_id in orphaned { inner.fetcher.cancel(&request_id); @@ -494,7 +591,8 @@ fn fetch_worker(inner: Arc) { }; let content = inner.fetcher.fetch(&uri, request_id.clone()); - let resolved = content.and_then(|content| write_cached_source(&uri, content.as_bytes())); + let resolved = content + .and_then(|content| write_cached_source(&inner.cache_dir, &uri, content.as_bytes())); let mut state = inner.state.lock().unwrap(); if priority == Priority::Bulk { @@ -577,12 +675,29 @@ fn pop_valid_uri(state: &mut State, priority: Priority) -> Option<(String, Prior None } -fn cache_dir() -> PathBuf { - env::temp_dir().join(DECOMPILED_DIR) -} - -fn cache_path(uri: &str) -> PathBuf { - cache_path_in(&cache_dir(), uri) +fn session_cache_dir(scope: &str) -> PathBuf { + static SESSION_COUNTER: AtomicU64 = AtomicU64::new(1); + let mut sequence = SESSION_COUNTER.fetch_add(1, Ordering::Relaxed); + let scope = hex::encode(Sha1::digest(scope.as_bytes())); + let started_at = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_nanos(); + let root = env::temp_dir().join(DECOMPILED_DIR); + let _ = fs::create_dir_all(&root); + loop { + let directory = root.join(format!( + "{scope}-{}-{started_at}-{sequence}", + std::process::id() + )); + match fs::create_dir(&directory) { + Ok(()) => return directory, + Err(error) if error.kind() == std::io::ErrorKind::AlreadyExists => { + sequence = SESSION_COUNTER.fetch_add(1, Ordering::Relaxed); + } + Err(_) => return directory, + } + } } fn cache_path_in(directory: &Path, uri: &str) -> PathBuf { @@ -615,20 +730,19 @@ fn cache_path_in(directory: &Path, uri: &str) -> PathBuf { directory.join(format!("{name}-{digest}.java")) } -fn write_cached_source(uri: &str, content: &[u8]) -> Option { +fn write_cached_source(directory: &Path, uri: &str, content: &[u8]) -> Option { if content.is_empty() { return None; } - let directory = cache_dir(); - if let Err(error) = fs::create_dir_all(&directory) { + if let Err(error) = fs::create_dir_all(directory) { lsp_error!( "[decompile] Failed to create {}: {error}", directory.display() ); return None; } - let target = cache_path_in(&directory, uri); + let target = cache_path_in(directory, uri); if target.is_file() { return Some(path_to_file_uri(&target)); } @@ -717,11 +831,25 @@ fn replace_jdt_location_uris(value: &mut Value, replacements: &HashMap usize { - value - .find(|character: char| { - character.is_whitespace() || matches!(character, ')' | ']' | '"' | '>' | '`' | '\'') - }) - .unwrap_or(value.len()) + if let Some(class_end) = value.find(".class") { + let class_end = class_end + ".class".len(); + if value[class_end..].starts_with('?') { + return value[class_end..] + .find(is_jdt_uri_terminator) + .map(|query_end| class_end + query_end) + .unwrap_or(value.len()); + } + return class_end; + } + value.find(is_jdt_uri_terminator).unwrap_or(value.len()) +} + +fn is_jdt_uri_terminator(character: char) -> bool { + character.is_whitespace() + || matches!( + character, + ')' | ']' | '}' | '"' | '>' | '`' | '\'' | ',' | ';' + ) } fn collect_jdt_uris(value: &Value, uris: &mut Vec, seen: &mut HashSet) { @@ -778,7 +906,10 @@ fn replace_in_strings(value: &mut Value, replacements: &HashMap) #[cfg(test)] mod tests { use super::*; - use std::sync::atomic::{AtomicUsize, Ordering}; + use std::sync::{ + atomic::{AtomicUsize, Ordering}, + mpsc, + }; struct FakeFetcher { calls: AtomicUsize, @@ -846,6 +977,10 @@ mod tests { ) } + fn cached_path(coordinator: &DecompileCoordinator, uri: &str) -> PathBuf { + cache_path_in(&coordinator.inner.cache_dir, uri) + } + #[test] fn rewrites_nested_workspace_symbol_locations() { let mut result = json!([ @@ -902,13 +1037,35 @@ mod tests { assert!(uris.is_empty()); } + #[test] + fn embedded_uri_extraction_stops_before_sentence_punctuation() { + let uri = "jdt://contents/java.base/java.lang/String.class"; + let result = json!(format!("See {uri}, then java.lang.Object;")); + let mut uris = Vec::new(); + + collect_jdt_uris(&result, &mut uris, &mut HashSet::new()); + + assert_eq!(uris, vec![uri]); + } + + #[test] + fn embedded_uri_extraction_preserves_query_components() { + let uri = "jdt://contents/java.base/java.lang/String.class?=java.base/String"; + let result = json!(format!("See {uri}, then continue.")); + let mut uris = Vec::new(); + + collect_jdt_uris(&result, &mut uris, &mut HashSet::new()); + + assert_eq!(uris, vec![uri]); + } + #[test] fn concurrent_waiters_share_one_fetch() { let fetcher = Arc::new(FakeFetcher::new(Duration::from_millis(20))); let coordinator = DecompileCoordinator::with_fetcher(fetcher.clone(), "test-".to_string(), 2, 0); let uri = unique_uri("single-flight"); - let _ = fs::remove_file(cache_path(&uri)); + let _ = fs::remove_file(cached_path(&coordinator, &uri)); let first = { let coordinator = coordinator.clone(); @@ -924,7 +1081,7 @@ mod tests { assert_eq!(first.join().unwrap().len(), 1); assert_eq!(second.join().unwrap().len(), 1); assert_eq!(fetcher.calls.load(Ordering::Relaxed), 1); - let _ = fs::remove_file(cache_path(&uri)); + let _ = fs::remove_file(cached_path(&coordinator, &uri)); coordinator.shutdown(); } @@ -937,7 +1094,7 @@ mod tests { .map(|index| unique_uri(&format!("parallel-{index}"))) .collect(); for uri in &uris { - let _ = fs::remove_file(cache_path(uri)); + let _ = fs::remove_file(cached_path(&coordinator, uri)); } assert_eq!( @@ -946,7 +1103,7 @@ mod tests { ); assert_eq!(fetcher.max_active.load(Ordering::SeqCst), 2); for uri in &uris { - let _ = fs::remove_file(cache_path(uri)); + let _ = fs::remove_file(cached_path(&coordinator, uri)); } coordinator.shutdown(); } @@ -957,7 +1114,7 @@ mod tests { let coordinator = DecompileCoordinator::with_fetcher(fetcher, "test-".to_string(), 1, 0); let uris = vec![unique_uri("deadline-a"), unique_uri("deadline-b")]; for uri in &uris { - let _ = fs::remove_file(cache_path(uri)); + let _ = fs::remove_file(cached_path(&coordinator, uri)); } let replacements = resolve_uris( @@ -970,8 +1127,56 @@ mod tests { assert_eq!(replacements.len(), 1); for uri in &uris { - let _ = fs::remove_file(cache_path(uri)); + let _ = fs::remove_file(cached_path(&coordinator, uri)); + } + coordinator.shutdown(); + } + + #[test] + fn queued_jobs_complete_when_their_deadline_expires() { + let fetcher = Arc::new(FakeFetcher::new(Duration::from_millis(150))); + let coordinator = + DecompileCoordinator::with_fetcher(fetcher.clone(), "deadline-".to_string(), 1, 1); + let first_uri = unique_uri("blocking-job"); + let second_uri = unique_uri("queued-deadline"); + let (first_tx, _first_rx) = mpsc::channel(); + assert!(coordinator + .submit(RewriteJob { + token: 1, + message: json!({ "result": [{ "uri": first_uri }] }), + mode: RewriteMode::Locations, + priority: Priority::Interactive, + deadline: Instant::now() + Duration::from_secs(1), + complete: Box::new(move |message| { + let _ = first_tx.send(message); + }), + }) + .is_ok()); + let worker_start_deadline = Instant::now() + Duration::from_secs(1); + while fetcher.active.load(Ordering::SeqCst) == 0 && Instant::now() < worker_start_deadline { + thread::sleep(Duration::from_millis(1)); } + assert_eq!(fetcher.active.load(Ordering::SeqCst), 1); + + let original = json!({ "result": [{ "uri": second_uri }] }); + let (second_tx, second_rx) = mpsc::channel(); + assert!(coordinator + .submit(RewriteJob { + token: 2, + message: original.clone(), + mode: RewriteMode::Locations, + priority: Priority::Interactive, + deadline: Instant::now() + Duration::from_millis(20), + complete: Box::new(move |message| { + let _ = second_tx.send(message); + }), + }) + .is_ok()); + + assert_eq!( + second_rx.recv_timeout(Duration::from_millis(500)).unwrap(), + original + ); coordinator.shutdown(); } @@ -981,7 +1186,7 @@ mod tests { let coordinator = DecompileCoordinator::with_fetcher(fetcher.clone(), "test-".to_string(), 1, 0); let uri = unique_uri("shared-cancel"); - let _ = fs::remove_file(cache_path(&uri)); + let _ = fs::remove_file(cached_path(&coordinator, &uri)); let first = { let coordinator = coordinator.clone(); @@ -1007,7 +1212,36 @@ mod tests { assert!(first.join().unwrap().is_empty()); assert_eq!(second.join().unwrap().len(), 1); assert_eq!(fetcher.cancellations.load(Ordering::Relaxed), 0); - let _ = fs::remove_file(cache_path(&uri)); + let _ = fs::remove_file(cached_path(&coordinator, &uri)); + coordinator.shutdown(); + } + + #[test] + fn timed_out_fetch_does_not_poison_immediate_retry() { + let fetcher = Arc::new(FakeFetcher::new(Duration::from_millis(60))); + let coordinator = + DecompileCoordinator::with_fetcher(fetcher.clone(), "retry-".to_string(), 2, 0); + let uri = unique_uri("retry-after-timeout"); + + let first = resolve_uris( + &coordinator.inner, + 1, + std::slice::from_ref(&uri), + Priority::Interactive, + Instant::now() + Duration::from_millis(10), + ); + let second = resolve_uris( + &coordinator.inner, + 2, + std::slice::from_ref(&uri), + Priority::Interactive, + Instant::now() + Duration::from_millis(250), + ); + + assert!(first.is_empty()); + assert_eq!(second.len(), 1); + assert_eq!(fetcher.calls.load(Ordering::Relaxed), 2); + assert_eq!(fetcher.cancellations.load(Ordering::Relaxed), 1); coordinator.shutdown(); } @@ -1017,7 +1251,7 @@ mod tests { let coordinator = DecompileCoordinator::with_fetcher(fetcher.clone(), "test-".to_string(), 1, 0); let uri = unique_uri("negative-cache"); - let _ = fs::remove_file(cache_path(&uri)); + let _ = fs::remove_file(cached_path(&coordinator, &uri)); assert!(rewrite_for_test( &coordinator, @@ -1037,10 +1271,103 @@ mod tests { coordinator.shutdown(); } + #[test] + fn coordinators_do_not_share_cache_across_server_sessions() { + let uri = unique_uri("session-cache"); + let first_fetcher = Arc::new(FakeFetcher::new(Duration::ZERO)); + let first = DecompileCoordinator::with_fetcher( + first_fetcher.clone(), + "same-server-scope-".to_string(), + 1, + 0, + ); + assert_eq!( + rewrite_for_test(&first, 1, std::slice::from_ref(&uri), Priority::Interactive).len(), + 1 + ); + let first_cache = first.inner.cache_dir.clone(); + first.shutdown(); + first.cleanup_cache(); + assert!(!first_cache.exists()); + + let second_fetcher = Arc::new(FakeFetcher::new(Duration::ZERO)); + let second = DecompileCoordinator::with_fetcher( + second_fetcher.clone(), + "same-server-scope-".to_string(), + 1, + 0, + ); + assert_ne!(first_cache, second.inner.cache_dir); + assert_eq!( + rewrite_for_test( + &second, + 2, + std::slice::from_ref(&uri), + Priority::Interactive + ) + .len(), + 1 + ); + assert_eq!(second_fetcher.calls.load(Ordering::Relaxed), 1); + second.shutdown(); + second.cleanup_cache(); + } + + #[test] + fn canceling_a_queued_bulk_job_retires_all_state() { + let fetcher = Arc::new(FakeFetcher::new(Duration::ZERO)); + let coordinator = + DecompileCoordinator::with_fetcher(fetcher, "queued-cancel-".to_string(), 0, 0); + assert!(coordinator + .submit(RewriteJob { + token: 7, + message: json!({ "result": [] }), + mode: RewriteMode::Locations, + priority: Priority::Bulk, + deadline: Instant::now() + Duration::from_secs(1), + complete: Box::new(|_| {}), + }) + .is_ok()); + + coordinator.cancel(7); + + let state = coordinator.inner.state.lock().unwrap(); + assert!(!state.canceled_jobs.contains(&7)); + assert!(state.bulk_jobs.is_empty()); + drop(state); + coordinator.shutdown(); + } + + #[test] + fn shutdown_completes_queued_jobs_unchanged() { + let fetcher = Arc::new(FakeFetcher::new(Duration::ZERO)); + let coordinator = + DecompileCoordinator::with_fetcher(fetcher, "shutdown-".to_string(), 0, 0); + let original = json!({ "result": [{ "uri": "jdt://contents/test/A.class" }] }); + let (sender, receiver) = mpsc::channel(); + assert!(coordinator + .submit(RewriteJob { + token: 8, + message: original.clone(), + mode: RewriteMode::Locations, + priority: Priority::Interactive, + deadline: Instant::now() + Duration::from_secs(1), + complete: Box::new(move |message| { + let _ = sender.send(message); + }), + }) + .is_ok()); + + coordinator.shutdown(); + + assert_eq!(receiver.recv().unwrap(), original); + } + #[test] fn competing_cache_writes_never_expose_partial_content() { let uri = unique_uri("atomic-cache"); - let path = cache_path(&uri); + let directory = session_cache_dir("atomic-cache-test"); + let path = cache_path_in(&directory, &uri); let _ = fs::remove_file(&path); let first_content = vec![b'a'; 32 * 1024]; let second_content = vec![b'b'; 32 * 1024]; @@ -1048,12 +1375,14 @@ mod tests { let first = { let uri = uri.clone(); let content = first_content.clone(); - thread::spawn(move || write_cached_source(&uri, &content)) + let directory = directory.clone(); + thread::spawn(move || write_cached_source(&directory, &uri, &content)) }; let second = { let uri = uri.clone(); let content = second_content.clone(); - thread::spawn(move || write_cached_source(&uri, &content)) + let directory = directory.clone(); + thread::spawn(move || write_cached_source(&directory, &uri, &content)) }; assert!(first.join().unwrap().is_some()); assert!(second.join().unwrap().is_some()); @@ -1081,10 +1410,11 @@ mod tests { #[test] fn empty_sources_are_not_cached() { let uri = unique_uri("empty"); - let path = cache_path(&uri); + let directory = session_cache_dir("empty-cache-test"); + let path = cache_path_in(&directory, &uri); let _ = fs::remove_file(&path); - assert_eq!(write_cached_source(&uri, b""), None); + assert_eq!(write_cached_source(&directory, &uri, b""), None); assert!(!path.exists()); } } diff --git a/proxy/src/main.rs b/proxy/src/main.rs index 9d2f676..839494e 100644 --- a/proxy/src/main.rs +++ b/proxy/src/main.rs @@ -11,11 +11,11 @@ use http::handle_http; use output::Output; use pending::PendingResponses; use proxy_common::{ - contains_subslice, parse_lsp_content, raw_has_id, spawn_parent_monitor, LspReader, + contains_subslice, encode_lsp, parse_lsp_content, raw_has_id, spawn_parent_monitor, LspReader, }; use serde_json::{json, Value}; use std::{ - collections::HashMap, + collections::{HashMap, HashSet}, env, fs, io::{self, BufReader, Write}, net::TcpListener, @@ -54,6 +54,21 @@ impl TrackedRequest { } } +#[derive(Default)] +struct SuppressedResponses { + ids: HashSet, +} + +impl SuppressedResponses { + fn insert(&mut self, id: Value) { + self.ids.insert(id); + } + + fn take(&mut self, id: &Value) -> bool { + self.ids.remove(id) + } +} + fn main() { let output = Output::start(); log::init(output.clone()); @@ -110,7 +125,7 @@ fn main() { let alive = Arc::new(AtomicBool::new(true)); let owned_id_prefix = format!("{proxy_id}-proxy-"); - let pending = Arc::new(PendingResponses::new(owned_id_prefix.clone())); + let pending = Arc::new(PendingResponses::new()); let decompile = DecompileCoordinator::new( Arc::clone(&child_stdin), Arc::clone(&pending), @@ -134,6 +149,7 @@ fn main() { let tracked_ids: Arc>> = Arc::new(Mutex::new(HashMap::new())); let active_rewrites: Arc>> = Arc::new(Mutex::new(HashMap::new())); + let suppressed_responses = Arc::new(Mutex::new(SuppressedResponses::default())); // --- Thread 1: Zed stdin -> JDTLS stdin (track definition requests) --- let stdin_writer = Arc::clone(&child_stdin); @@ -142,6 +158,8 @@ fn main() { let active_in = Arc::clone(&active_rewrites); let jobs_in = Arc::clone(&job_counter); let decompile_in = decompile.clone(); + let output_in = output.clone(); + let suppressed_in = Arc::clone(&suppressed_responses); let latest_workspace_job = Arc::new(Mutex::new(None::)); let latest_workspace_in = Arc::clone(&latest_workspace_job); thread::spawn(move || { @@ -162,15 +180,15 @@ fn main() { }; if msg.get("method").and_then(Value::as_str) == Some("$/cancelRequest") { if let Some(id) = msg.pointer("/params/id") { - let request = tracked_in.lock().unwrap().get(id).cloned(); - let active_token = active_in.lock().unwrap().remove(id); - let token = request - .as_ref() - .map(|request| request.token) - .or(active_token); - if let Some(token) = token { + let (tracked, active_token) = + take_request_for_cancellation(&tracked_in, &active_in, id); + if let Some(request) = tracked { + clear_latest_workspace(&latest_workspace_in, request.token); + } + if let Some(token) = active_token { clear_latest_workspace(&latest_workspace_in, token); decompile_in.cancel(token); + output_in.send_value(&request_canceled(id)); } } } else if raw_has_id(&raw) { @@ -180,14 +198,29 @@ fn main() { let previous = latest_workspace_in.lock().unwrap().replace(token); if let Some(previous) = previous { - remove_active_token(&active_in, previous); - decompile_in.cancel(previous); + let (suppressed, active) = + retire_request_token(&tracked_in, &active_in, previous); + if let Some(id) = suppressed { + suppressed_in.lock().unwrap().insert(id.clone()); + cancel_jdtls_request(&stdin_writer, &id); + } + if active { + decompile_in.cancel(previous); + } } } - let previous = tracked_in.lock().unwrap().insert(id, request); + let (previous, active_token) = { + let mut tracked = tracked_in.lock().unwrap(); + let mut active = active_in.lock().unwrap(); + let previous = tracked.insert(id.clone(), request); + let active_token = active.remove(&id); + (previous, active_token) + }; if let Some(previous) = previous { - remove_active_token(&active_in, previous.token); - decompile_in.cancel(previous.token); + clear_latest_workspace(&latest_workspace_in, previous.token); + } + if let Some(active_token) = active_token { + decompile_in.cancel(active_token); } } } @@ -210,134 +243,130 @@ fn main() { let active_out = Arc::clone(&active_rewrites); let decompile_out = decompile.clone(); let output_router = output.clone(); + let suppressed_out = Arc::clone(&suppressed_responses); let latest_workspace_out = Arc::clone(&latest_workspace_job); - thread::spawn(move || { + let stdout_thread = thread::spawn(move || { let mut reader = LspReader::new(BufReader::new(child_stdout)); - while alive_out.load(Ordering::Relaxed) { - match reader.read_message() { - Ok(Some(raw)) => { - // Fast path: notifications (no `id`) can't be responses we - // need to intercept. Forward the raw bytes without parsing. - if !raw_has_id(&raw) { + while let Ok(Some(raw)) = reader.read_message() { + // Fast path: notifications (no `id`) can't be responses we + // need to intercept. Forward the raw bytes without parsing. + if !raw_has_id(&raw) { + output_router.send_raw(raw); + continue; + } + + let Some(mut msg) = parse_lsp_content(&raw) else { + output_router.send_raw(raw); + continue; + }; + + // Route responses to pending HTTP requests + if pending_out.route(&msg) { + continue; + } + if msg + .get("id") + .is_some_and(|id| suppressed_out.lock().unwrap().take(id)) + { + continue; + } + + // Rewrite jdt:// URIs in location or documentation responses. + // The bounded coordinator keeps this router free to deliver + // java/classFileContents responses through `pending`. + if msg.get("method").is_none() { + let Some(id) = msg.get("id").cloned() else { + output_router.send_raw(raw); + continue; + }; + let request = tracked_out.lock().unwrap().get(&id).cloned(); + if let Some(request) = request { + if let Some(fallback) = completion_resolve_fallback(&msg, &request) { + remove_tracked_request(&tracked_out, &id, request.token); + if should_log_completion_fallback() { + lsp_warn!( + "JDTLS completion resolution failed with -32603; \ + using the unresolved item, so documentation, imports, \ + commands, or additional edits may be missing" + ); + } + output_router.send_value(&fallback); + continue; + } + if msg.get("error").is_some() { + remove_tracked_request(&tracked_out, &id, request.token); + clear_latest_workspace(&latest_workspace_out, request.token); output_router.send_raw(raw); continue; } + if request.rewrite == RewriteKind::Completion { + remove_tracked_request(&tracked_out, &id, request.token); + process_completions(&mut msg); + output_router.send_value(&msg); + continue; + } - let Some(mut msg) = parse_lsp_content(&raw) else { + let output = output_router.clone(); + let sanitize_completion = request.method == "completionItem/resolve"; + let mode = match request.rewrite { + RewriteKind::Locations => RewriteMode::Locations, + RewriteKind::Documentation => RewriteMode::Strings, + RewriteKind::Completion => unreachable!(), + }; + let priority = if request.method == "workspace/symbol" { + Priority::Bulk + } else { + Priority::Interactive + }; + let deadline = std::time::Instant::now() + rewrite_timeout(&request.method); + let workspace_job = (request.method == "workspace/symbol") + .then(|| (Arc::clone(&latest_workspace_out), request.token)); + if !activate_rewrite(&tracked_out, &active_out, &id, request.token) { output_router.send_raw(raw); continue; - }; - - // Route responses to pending HTTP requests - if pending_out.route(&msg) { + } + if decompile_out.is_canceled(request.token) { + active_out.lock().unwrap().remove(&id); + decompile_out.consume_cancellation(request.token); + clear_latest_workspace(&latest_workspace_out, request.token); continue; } - - // Rewrite jdt:// URIs in location or documentation responses. - // The bounded coordinator keeps this router free to deliver - // java/classFileContents responses through `pending`. - if msg.get("method").is_none() { - let Some(id) = msg.get("id").cloned() else { - output_router.send_raw(raw); - continue; - }; - let request = tracked_out.lock().unwrap().get(&id).cloned(); - if let Some(request) = request { - let canceled = decompile_out.is_canceled(request.token); - if canceled { - tracked_out.lock().unwrap().remove(&id); - decompile_out.consume_cancellation(request.token); - clear_latest_workspace(&latest_workspace_out, request.token); - continue; - } - if let Some(fallback) = completion_resolve_fallback(&msg, &request) { - tracked_out.lock().unwrap().remove(&id); - if should_log_completion_fallback() { - lsp_warn!( - "JDTLS completion resolution failed with -32603; \ - using the unresolved item, so documentation, imports, \ - commands, or additional edits may be missing" - ); - } - output_router.send_value(&fallback); - continue; + let active = Arc::clone(&active_out); + let active_id = id; + let completion_coordinator = decompile_out.clone(); + let job = RewriteJob { + token: request.token, + message: msg, + mode, + priority, + deadline, + complete: Box::new(move |mut message| { + if sanitize_completion { + sanitize_resolved_completion(&mut message); } - if msg.get("error").is_some() { - tracked_out.lock().unwrap().remove(&id); - clear_latest_workspace(&latest_workspace_out, request.token); - output_router.send_raw(raw); - continue; + let owns_response = + active.lock().unwrap().remove(&active_id) == Some(request.token); + completion_coordinator.consume_cancellation(request.token); + if owns_response { + output.send_value(&message); } - if request.rewrite == RewriteKind::Completion { - tracked_out.lock().unwrap().remove(&id); - process_completions(&mut msg); - output_router.send_value(&msg); - continue; + if let Some((latest, token)) = workspace_job { + clear_latest_workspace(&latest, token); } - - let output = output_router.clone(); - let sanitize_completion = request.method == "completionItem/resolve"; - let mode = match request.rewrite { - RewriteKind::Locations => RewriteMode::Locations, - RewriteKind::Documentation => RewriteMode::Strings, - RewriteKind::Completion => unreachable!(), - }; - let priority = if request.method == "workspace/symbol" { - Priority::Bulk - } else { - Priority::Interactive - }; - let deadline = - std::time::Instant::now() + rewrite_timeout(&request.method); - let workspace_job = (request.method == "workspace/symbol") - .then(|| (Arc::clone(&latest_workspace_out), request.token)); - active_out.lock().unwrap().insert(id.clone(), request.token); - tracked_out.lock().unwrap().remove(&id); - if decompile_out.is_canceled(request.token) { - active_out.lock().unwrap().remove(&id); - decompile_out.consume_cancellation(request.token); - clear_latest_workspace(&latest_workspace_out, request.token); - continue; - } - let active = Arc::clone(&active_out); - let active_id = id; - let completion_coordinator = decompile_out.clone(); - let job = RewriteJob { - token: request.token, - message: msg, - mode, - priority, - deadline, - complete: Box::new(move |mut message| { - if sanitize_completion { - sanitize_resolved_completion(&mut message); - } - let owns_response = active.lock().unwrap().remove(&active_id) - == Some(request.token); - completion_coordinator.consume_cancellation(request.token); - if owns_response { - output.send_value(&message); - } - if let Some((latest, token)) = workspace_job { - clear_latest_workspace(&latest, token); - } - }), - }; - if let Err(job) = decompile_out.submit(job) { - let RewriteJob { - message, complete, .. - } = job; - complete(message); - } - continue; - } + }), + }; + if let Err(job) = decompile_out.submit(job) { + let RewriteJob { + message, complete, .. + } = job; + complete(message); } - - // Passthrough - output_router.send_raw(raw); + continue; } - Ok(None) | Err(_) => break, } + + // Passthrough + output_router.send_raw(raw); } alive_out.store(false, Ordering::Relaxed); }); @@ -383,10 +412,13 @@ fn main() { Err(error) => break Err(error), } }; - lsp_info!("JDTLS process exited: {status:?}"); alive.store(false, Ordering::Relaxed); - decompile.shutdown(); + let _ = stdout_thread.join(); + lsp_info!("JDTLS process exited: {status:?}"); pending.clear(); + decompile.shutdown(); + output.shutdown(); + decompile.cleanup_cache(); let _ = fs::remove_file(&port_file); } @@ -403,11 +435,81 @@ fn clear_latest_workspace(latest: &Mutex>, token: u64) { } } -fn remove_active_token(active: &Mutex>, token: u64) { - active - .lock() - .unwrap() - .retain(|_, active_token| *active_token != token); +fn request_canceled(id: &Value) -> Value { + json!({ + "jsonrpc": "2.0", + "id": id, + "error": { + "code": -32800, + "message": "Request cancelled." + } + }) +} + +fn cancel_jdtls_request(writer: &SharedWriter, id: &Value) { + let cancel = encode_lsp(&json!({ + "jsonrpc": "2.0", + "method": "$/cancelRequest", + "params": { "id": id } + })); + let mut writer = writer.lock().unwrap(); + let _ = writer.write_all(cancel.as_bytes()); + let _ = writer.flush(); +} + +fn take_request_for_cancellation( + tracked: &Mutex>, + active: &Mutex>, + id: &Value, +) -> (Option, Option) { + let mut tracked = tracked.lock().unwrap(); + let mut active = active.lock().unwrap(); + (tracked.remove(id), active.remove(id)) +} + +fn retire_request_token( + tracked: &Mutex>, + active: &Mutex>, + token: u64, +) -> (Option, bool) { + let mut tracked = tracked.lock().unwrap(); + let mut active = active.lock().unwrap(); + let tracked_id = tracked + .iter() + .find_map(|(id, request)| (request.token == token).then(|| id.clone())); + if let Some(id) = &tracked_id { + tracked.remove(id); + } + let active_id = active + .iter() + .find_map(|(id, active_token)| (*active_token == token).then(|| id.clone())); + if let Some(id) = &active_id { + active.remove(id); + } + (tracked_id, active_id.is_some()) +} + +fn activate_rewrite( + tracked: &Mutex>, + active: &Mutex>, + id: &Value, + token: u64, +) -> bool { + let mut tracked = tracked.lock().unwrap(); + let mut active = active.lock().unwrap(); + if tracked.get(id).map(|request| request.token) != Some(token) { + return false; + } + tracked.remove(id); + active.insert(id.clone(), token); + true +} + +fn remove_tracked_request(tracked: &Mutex>, id: &Value, token: u64) { + let mut tracked = tracked.lock().unwrap(); + if tracked.get(id).map(|request| request.token) == Some(token) { + tracked.remove(id); + } } fn should_log_completion_fallback() -> bool { @@ -587,6 +689,78 @@ mod tests { ); } + #[test] + fn activating_rewrite_moves_request_atomically() { + let id = json!(9); + let tracked = Mutex::new(HashMap::from([( + id.clone(), + TrackedRequest::new(44, "textDocument/definition", RewriteKind::Locations, None), + )])); + let active = Mutex::new(HashMap::new()); + + assert!(activate_rewrite(&tracked, &active, &id, 44)); + assert!(tracked.lock().unwrap().is_empty()); + assert_eq!(active.lock().unwrap().get(&id), Some(&44)); + } + + #[test] + fn retiring_tracked_request_returns_id_for_late_suppression() { + let id = json!("workspace-1"); + let tracked = Mutex::new(HashMap::from([( + id.clone(), + TrackedRequest::new(45, "workspace/symbol", RewriteKind::Locations, None), + )])); + let active = Mutex::new(HashMap::new()); + + assert_eq!( + retire_request_token(&tracked, &active, 45), + (Some(id), false) + ); + assert!(tracked.lock().unwrap().is_empty()); + } + + #[test] + fn suppressed_response_ids_remain_owned_until_the_response_arrives() { + let mut suppressed = SuppressedResponses::default(); + for id in 0..=1024 { + suppressed.insert(json!(id)); + } + + assert!(suppressed.take(&json!(0))); + assert!(suppressed.take(&json!(1024))); + } + + #[test] + fn cancellation_response_preserves_request_id() { + assert_eq!( + request_canceled(&json!("request-1")), + json!({ + "jsonrpc": "2.0", + "id": "request-1", + "error": { + "code": -32800, + "message": "Request cancelled." + } + }) + ); + } + + #[test] + fn cancellation_removes_unanswered_tracking_state() { + let id = json!("hover-1"); + let tracked = Mutex::new(HashMap::from([( + id.clone(), + TrackedRequest::new(46, "textDocument/hover", RewriteKind::Documentation, None), + )])); + let active = Mutex::new(HashMap::new()); + + let (request, active_token) = take_request_for_cancellation(&tracked, &active, &id); + + assert_eq!(request.map(|request| request.token), Some(46)); + assert_eq!(active_token, None); + assert!(tracked.lock().unwrap().is_empty()); + } + #[test] fn falls_back_only_for_internal_completion_resolve_errors() { let request = TrackedRequest::new( diff --git a/proxy/src/output.rs b/proxy/src/output.rs index 674f2c8..91118b0 100644 --- a/proxy/src/output.rs +++ b/proxy/src/output.rs @@ -4,17 +4,31 @@ use std::{ io::{self, Write}, sync::{ atomic::{AtomicBool, Ordering}, - mpsc, Arc, + mpsc, Arc, Mutex, }, - thread, + thread::{self, JoinHandle}, + time::Duration, }; +const SHUTDOWN_DRAIN_TIMEOUT: Duration = Duration::from_millis(250); + #[derive(Clone)] pub struct Output { - sender: mpsc::Sender>, + inner: Arc, +} + +struct Inner { + sender: Mutex>>, + worker: Mutex>>, + done: Mutex>>, failed: Arc, } +enum Message { + Frame(Vec), + Shutdown, +} + impl Output { pub fn start() -> Self { Self::start_with_writer(io::stdout()) @@ -22,16 +36,33 @@ impl Output { fn start_with_writer(writer: impl Write + Send + 'static) -> Self { let (sender, receiver) = mpsc::channel(); + let (done_sender, done) = mpsc::channel(); let failed = Arc::new(AtomicBool::new(false)); let writer_failed = Arc::clone(&failed); - thread::spawn(move || write_loop(writer, receiver, writer_failed)); - Self { sender, failed } + let worker = thread::spawn(move || { + write_loop(writer, receiver, writer_failed); + let _ = done_sender.send(()); + }); + Self { + inner: Arc::new(Inner { + sender: Mutex::new(Some(sender)), + worker: Mutex::new(Some(worker)), + done: Mutex::new(Some(done)), + failed, + }), + } } pub fn send_raw(&self, raw: Vec) -> bool { - let sent = self.sender.send(raw).is_ok(); + let sent = self + .inner + .sender + .lock() + .unwrap() + .as_ref() + .is_some_and(|sender| sender.send(Message::Frame(raw)).is_ok()); if !sent { - self.failed.store(true, Ordering::Relaxed); + self.inner.failed.store(true, Ordering::Relaxed); } sent } @@ -41,15 +72,40 @@ impl Output { } pub fn failed(&self) -> bool { - self.failed.load(Ordering::Relaxed) + self.inner.failed.load(Ordering::Relaxed) + } + + /// Stops accepting frames and gives stdout a bounded interval to drain. + /// A blocked editor must not prevent the proxy process from terminating. + pub fn shutdown(&self) { + if let Some(sender) = self.inner.sender.lock().unwrap().take() { + let _ = sender.send(Message::Shutdown); + } + let drained = self + .inner + .done + .lock() + .unwrap() + .take() + .is_none_or(|done| done.recv_timeout(SHUTDOWN_DRAIN_TIMEOUT).is_ok()); + if let Some(worker) = self.inner.worker.lock().unwrap().take() { + if drained { + let _ = worker.join(); + } + } } } -fn write_loop(mut writer: impl Write, receiver: mpsc::Receiver>, failed: Arc) { +fn write_loop(mut writer: impl Write, receiver: mpsc::Receiver, failed: Arc) { while let Ok(message) = receiver.recv() { - if writer.write_all(&message).is_err() || writer.flush().is_err() { - failed.store(true, Ordering::Relaxed); - break; + match message { + Message::Frame(frame) => { + if writer.write_all(&frame).is_err() || writer.flush().is_err() { + failed.store(true, Ordering::Relaxed); + break; + } + } + Message::Shutdown => break, } } } @@ -85,6 +141,23 @@ mod tests { } } + struct BlockingWriter { + started: mpsc::Sender<()>, + release: mpsc::Receiver<()>, + } + + impl Write for BlockingWriter { + fn write(&mut self, buf: &[u8]) -> io::Result { + let _ = self.started.send(()); + let _ = self.release.recv(); + Ok(buf.len()) + } + + fn flush(&mut self) -> io::Result<()> { + Ok(()) + } + } + #[test] fn serializes_complete_frames_in_submission_order() { let writer = SharedWriter::default(); @@ -93,16 +166,10 @@ mod tests { assert!(output.send_raw(b"first".to_vec())); assert!(output.send_raw(b"second".to_vec())); - drop(output); - - for _ in 0..100 { - if bytes.lock().unwrap().len() == 11 { - break; - } - thread::sleep(std::time::Duration::from_millis(1)); - } + output.shutdown(); assert_eq!(*bytes.lock().unwrap(), b"firstsecond"); + assert!(!output.send_raw(b"third".to_vec())); } #[test] @@ -118,5 +185,26 @@ mod tests { } assert!(output.failed()); + output.shutdown(); + } + + #[test] + fn shutdown_does_not_wait_forever_for_blocked_stdout() { + let (started_sender, started_receiver) = mpsc::channel(); + let (release_sender, release_receiver) = mpsc::channel(); + let output = Output::start_with_writer(BlockingWriter { + started: started_sender, + release: release_receiver, + }); + assert!(output.send_raw(b"blocked".to_vec())); + started_receiver + .recv_timeout(Duration::from_secs(1)) + .unwrap(); + + let started = std::time::Instant::now(); + output.shutdown(); + + assert!(started.elapsed() < Duration::from_secs(1)); + let _ = release_sender.send(()); } } diff --git a/proxy/src/pending.rs b/proxy/src/pending.rs index 84bf043..176acdf 100644 --- a/proxy/src/pending.rs +++ b/proxy/src/pending.rs @@ -1,30 +1,37 @@ use serde_json::Value; use std::{ - collections::HashMap, + collections::{HashMap, HashSet}, sync::{mpsc, Mutex}, }; pub struct PendingResponses { - owned_id_prefix: String, - pending: Mutex>>, + state: Mutex, +} + +#[derive(Default)] +struct State { + pending: HashMap>, + owned: HashSet, } impl PendingResponses { - pub fn new(owned_id_prefix: String) -> Self { + pub fn new() -> Self { Self { - owned_id_prefix, - pending: Mutex::new(HashMap::new()), + state: Mutex::new(State::default()), } } pub fn register(&self, id: Value) -> mpsc::Receiver { let (sender, receiver) = mpsc::channel(); - self.pending.lock().unwrap().insert(id, sender); + let mut state = self.state.lock().unwrap(); + state.owned.insert(id.clone()); + state.pending.insert(id, sender); receiver } pub fn remove(&self, id: &Value) { - self.pending.lock().unwrap().remove(id); + let mut state = self.state.lock().unwrap(); + state.pending.remove(id); } /// Returns true when the response belongs to the proxy and must not be @@ -37,18 +44,26 @@ impl PendingResponses { return false; }; - let sender = self.pending.lock().unwrap().remove(id); + let sender = { + let mut state = self.state.lock().unwrap(); + let sender = state.pending.remove(id); + if sender.is_none() && state.owned.contains(id) { + return true; + } + sender + }; if let Some(sender) = sender { let _ = sender.send(message.clone()); return true; } - id.as_str() - .is_some_and(|id| id.starts_with(&self.owned_id_prefix)) + false } pub fn clear(&self) { - self.pending.lock().unwrap().clear(); + let mut state = self.state.lock().unwrap(); + state.pending.clear(); + state.owned.clear(); } } @@ -59,7 +74,7 @@ mod tests { #[test] fn routes_each_pending_response_once() { - let pending = PendingResponses::new("proxy-owned-".to_string()); + let pending = PendingResponses::new(); let id = json!("proxy-owned-1"); let receiver = pending.register(id.clone()); let response = json!({ "jsonrpc": "2.0", "id": id, "result": "ok" }); @@ -71,7 +86,10 @@ mod tests { #[test] fn swallows_late_owned_responses_but_not_client_responses() { - let pending = PendingResponses::new("proxy-owned-".to_string()); + let pending = PendingResponses::new(); + let retired = json!("proxy-owned-late"); + pending.register(retired.clone()); + pending.remove(&retired); assert!(pending.route(&json!({ "jsonrpc": "2.0", @@ -90,4 +108,31 @@ mod tests { "params": {} }))); } + + #[test] + fn does_not_claim_unregistered_prefixed_ids() { + let pending = PendingResponses::new(); + + assert!(!pending.route(&json!({ + "jsonrpc": "2.0", + "id": "proxy-owned-editor-request", + "result": null + }))); + } + + #[test] + fn ownership_is_retained_for_the_entire_session() { + let pending = PendingResponses::new(); + for id in 0..=1024 { + let id = json!(format!("proxy-owned-{id}")); + pending.register(id.clone()); + pending.remove(&id); + } + + assert!(pending.route(&json!({ + "jsonrpc": "2.0", + "id": "proxy-owned-0", + "result": null + }))); + } } From 9bc6d25b2c788c216be97c373e917ceffcfa8dbb Mon Sep 17 00:00:00 2001 From: Riccardo Strina Date: Sat, 1 Aug 2026 23:21:12 +0200 Subject: [PATCH 3/7] fix(proxy): Use valid decompiled filenames --- proxy/src/decompile.rs | 53 +++++++++++++++++++++++++++++++++++------- 1 file changed, 45 insertions(+), 8 deletions(-) diff --git a/proxy/src/decompile.rs b/proxy/src/decompile.rs index 87995bc..2529c7e 100644 --- a/proxy/src/decompile.rs +++ b/proxy/src/decompile.rs @@ -16,7 +16,7 @@ use std::{ time::{Duration, Instant, SystemTime, UNIX_EPOCH}, }; -const DECOMPILED_DIR: &str = "jdtls-decompiled"; +const DECOMPILED_DIR: &str = "jdtls_decompiled"; const FETCH_WORKERS: usize = 2; const JOB_WORKERS: usize = 2; const MAX_QUEUED_JOBS: usize = 64; @@ -687,7 +687,7 @@ fn session_cache_dir(scope: &str) -> PathBuf { let _ = fs::create_dir_all(&root); loop { let directory = root.join(format!( - "{scope}-{}-{started_at}-{sequence}", + "session_{scope}_{}_{started_at}_{sequence}", std::process::id() )); match fs::create_dir(&directory) { @@ -716,18 +716,29 @@ fn cache_path_in(directory: &Path, uri: &str) -> PathBuf { }) }) .unwrap_or("Decompiled"); - let name: String = raw_name + let mut name: String = raw_name .chars() - .map(|character| { - if character.is_ascii_alphanumeric() || matches!(character, '-' | '_' | '.') { + .enumerate() + .map(|(index, character)| { + let valid = if index == 0 { + character.is_ascii_alphabetic() || matches!(character, '_' | '$') + } else { + character.is_ascii_alphanumeric() || matches!(character, '_' | '$') + }; + if valid { character } else { '_' } }) .collect(); + if name.is_empty() { + name.push_str("Decompiled"); + } - directory.join(format!("{name}-{digest}.java")) + directory + .join(format!("uri_{digest}")) + .join(format!("{name}.java")) } fn write_cached_source(directory: &Path, uri: &str, content: &[u8]) -> Option { @@ -746,6 +757,12 @@ fn write_cached_source(directory: &Path, uri: &str, content: &[u8]) -> Option Date: Sat, 1 Aug 2026 23:36:39 +0200 Subject: [PATCH 4/7] fix(proxy): Avoid redundant cancellations --- proxy/src/main.rs | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/proxy/src/main.rs b/proxy/src/main.rs index 839494e..7a6a580 100644 --- a/proxy/src/main.rs +++ b/proxy/src/main.rs @@ -64,6 +64,10 @@ impl SuppressedResponses { self.ids.insert(id); } + fn contains(&self, id: &Value) -> bool { + self.ids.contains(id) + } + fn take(&mut self, id: &Value) -> bool { self.ids.remove(id) } @@ -168,6 +172,7 @@ fn main() { while alive_stdin.load(Ordering::Relaxed) { match reader.read_message() { Ok(Some(raw)) => { + let mut forward_to_jdtls = true; let should_parse = raw_has_id(&raw) || contains_subslice(&raw, b"$/cancelRequest"); if should_parse { @@ -180,16 +185,19 @@ fn main() { }; if msg.get("method").and_then(Value::as_str) == Some("$/cancelRequest") { if let Some(id) = msg.pointer("/params/id") { + let already_suppressed = suppressed_in.lock().unwrap().contains(id); let (tracked, active_token) = take_request_for_cancellation(&tracked_in, &active_in, id); if let Some(request) = tracked { clear_latest_workspace(&latest_workspace_in, request.token); } + let handled_locally = active_token.is_some(); if let Some(token) = active_token { clear_latest_workspace(&latest_workspace_in, token); decompile_in.cancel(token); output_in.send_value(&request_canceled(id)); } + forward_to_jdtls = !handled_locally && !already_suppressed; } } else if raw_has_id(&raw) { let token = jobs_in.fetch_add(1, Ordering::Relaxed); @@ -225,6 +233,9 @@ fn main() { } } } + if !forward_to_jdtls { + continue; + } let mut w = stdin_writer.lock().unwrap(); if w.write_all(&raw).is_err() || w.flush().is_err() { break; @@ -726,7 +737,9 @@ mod tests { suppressed.insert(json!(id)); } + assert!(suppressed.contains(&json!(0))); assert!(suppressed.take(&json!(0))); + assert!(!suppressed.contains(&json!(0))); assert!(suppressed.take(&json!(1024))); } From 03e845980406abb01ddfa2f7ed98d2178d80f26e Mon Sep 17 00:00:00 2001 From: Riccardo Strina Date: Mon, 3 Aug 2026 12:20:13 +0200 Subject: [PATCH 5/7] Refactor main to clear up the code managing input and output threads --- proxy/src/main.rs | 734 ++++++++++++++++++++++++++++++++-------------- 1 file changed, 506 insertions(+), 228 deletions(-) diff --git a/proxy/src/main.rs b/proxy/src/main.rs index 7a6a580..7d5191e 100644 --- a/proxy/src/main.rs +++ b/proxy/src/main.rs @@ -73,6 +73,47 @@ impl SuppressedResponses { } } +type TrackedRequests = Arc>>; +type ActiveRewrites = Arc>>; +type SharedSuppressedResponses = Arc>; +type LatestWorkspaceJob = Arc>>; + +struct StdinContext { + writer: SharedWriter, + alive: Arc, + tracked: TrackedRequests, + active: ActiveRewrites, + jobs: Arc, + decompile: DecompileCoordinator, + output: Output, + suppressed: SharedSuppressedResponses, + latest_workspace: LatestWorkspaceJob, +} + +struct StdoutContext { + pending: Arc, + alive: Arc, + tracked: TrackedRequests, + active: ActiveRewrites, + decompile: DecompileCoordinator, + output: Output, + suppressed: SharedSuppressedResponses, + latest_workspace: LatestWorkspaceJob, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum InputRoute { + Forward, + Consumed, +} + +#[derive(Debug, PartialEq)] +enum OutputRoute { + Raw(Vec), + Value(Value), + Consumed, +} + fn main() { let output = Output::start(); log::init(output.clone()); @@ -150,237 +191,37 @@ fn main() { // Track requests whose responses may contain jdt:// URIs so they can be // intercepted and rewritten. - let tracked_ids: Arc>> = - Arc::new(Mutex::new(HashMap::new())); - let active_rewrites: Arc>> = Arc::new(Mutex::new(HashMap::new())); + let tracked_ids: TrackedRequests = Arc::new(Mutex::new(HashMap::new())); + let active_rewrites: ActiveRewrites = Arc::new(Mutex::new(HashMap::new())); let suppressed_responses = Arc::new(Mutex::new(SuppressedResponses::default())); - - // --- Thread 1: Zed stdin -> JDTLS stdin (track definition requests) --- - let stdin_writer = Arc::clone(&child_stdin); - let alive_stdin = Arc::clone(&alive); - let tracked_in = Arc::clone(&tracked_ids); - let active_in = Arc::clone(&active_rewrites); - let jobs_in = Arc::clone(&job_counter); - let decompile_in = decompile.clone(); - let output_in = output.clone(); - let suppressed_in = Arc::clone(&suppressed_responses); let latest_workspace_job = Arc::new(Mutex::new(None::)); - let latest_workspace_in = Arc::clone(&latest_workspace_job); - thread::spawn(move || { - let stdin = io::stdin().lock(); - let mut reader = LspReader::new(BufReader::new(stdin)); - while alive_stdin.load(Ordering::Relaxed) { - match reader.read_message() { - Ok(Some(raw)) => { - let mut forward_to_jdtls = true; - let should_parse = - raw_has_id(&raw) || contains_subslice(&raw, b"$/cancelRequest"); - if should_parse { - let Some(msg) = parse_lsp_content(&raw) else { - let mut writer = stdin_writer.lock().unwrap(); - if writer.write_all(&raw).is_err() || writer.flush().is_err() { - break; - } - continue; - }; - if msg.get("method").and_then(Value::as_str) == Some("$/cancelRequest") { - if let Some(id) = msg.pointer("/params/id") { - let already_suppressed = suppressed_in.lock().unwrap().contains(id); - let (tracked, active_token) = - take_request_for_cancellation(&tracked_in, &active_in, id); - if let Some(request) = tracked { - clear_latest_workspace(&latest_workspace_in, request.token); - } - let handled_locally = active_token.is_some(); - if let Some(token) = active_token { - clear_latest_workspace(&latest_workspace_in, token); - decompile_in.cancel(token); - output_in.send_value(&request_canceled(id)); - } - forward_to_jdtls = !handled_locally && !already_suppressed; - } - } else if raw_has_id(&raw) { - let token = jobs_in.fetch_add(1, Ordering::Relaxed); - if let Some((id, request)) = tracked_request_for(&msg, token) { - if request.method == "workspace/symbol" { - let previous = - latest_workspace_in.lock().unwrap().replace(token); - if let Some(previous) = previous { - let (suppressed, active) = - retire_request_token(&tracked_in, &active_in, previous); - if let Some(id) = suppressed { - suppressed_in.lock().unwrap().insert(id.clone()); - cancel_jdtls_request(&stdin_writer, &id); - } - if active { - decompile_in.cancel(previous); - } - } - } - let (previous, active_token) = { - let mut tracked = tracked_in.lock().unwrap(); - let mut active = active_in.lock().unwrap(); - let previous = tracked.insert(id.clone(), request); - let active_token = active.remove(&id); - (previous, active_token) - }; - if let Some(previous) = previous { - clear_latest_workspace(&latest_workspace_in, previous.token); - } - if let Some(active_token) = active_token { - decompile_in.cancel(active_token); - } - } - } - } - if !forward_to_jdtls { - continue; - } - let mut w = stdin_writer.lock().unwrap(); - if w.write_all(&raw).is_err() || w.flush().is_err() { - break; - } - } - Ok(None) | Err(_) => break, - } - } - alive_stdin.store(false, Ordering::Relaxed); - }); - - // --- Thread 2: JDTLS stdout -> rewrite jdt:// URIs, modify completions -> Zed stdout / resolve pending --- - let pending_out = Arc::clone(&pending); - let alive_out = Arc::clone(&alive); - let tracked_out = Arc::clone(&tracked_ids); - let active_out = Arc::clone(&active_rewrites); - let decompile_out = decompile.clone(); - let output_router = output.clone(); - let suppressed_out = Arc::clone(&suppressed_responses); - let latest_workspace_out = Arc::clone(&latest_workspace_job); - let stdout_thread = thread::spawn(move || { - let mut reader = LspReader::new(BufReader::new(child_stdout)); - while let Ok(Some(raw)) = reader.read_message() { - // Fast path: notifications (no `id`) can't be responses we - // need to intercept. Forward the raw bytes without parsing. - if !raw_has_id(&raw) { - output_router.send_raw(raw); - continue; - } - - let Some(mut msg) = parse_lsp_content(&raw) else { - output_router.send_raw(raw); - continue; - }; - - // Route responses to pending HTTP requests - if pending_out.route(&msg) { - continue; - } - if msg - .get("id") - .is_some_and(|id| suppressed_out.lock().unwrap().take(id)) - { - continue; - } - // Rewrite jdt:// URIs in location or documentation responses. - // The bounded coordinator keeps this router free to deliver - // java/classFileContents responses through `pending`. - if msg.get("method").is_none() { - let Some(id) = msg.get("id").cloned() else { - output_router.send_raw(raw); - continue; - }; - let request = tracked_out.lock().unwrap().get(&id).cloned(); - if let Some(request) = request { - if let Some(fallback) = completion_resolve_fallback(&msg, &request) { - remove_tracked_request(&tracked_out, &id, request.token); - if should_log_completion_fallback() { - lsp_warn!( - "JDTLS completion resolution failed with -32603; \ - using the unresolved item, so documentation, imports, \ - commands, or additional edits may be missing" - ); - } - output_router.send_value(&fallback); - continue; - } - if msg.get("error").is_some() { - remove_tracked_request(&tracked_out, &id, request.token); - clear_latest_workspace(&latest_workspace_out, request.token); - output_router.send_raw(raw); - continue; - } - if request.rewrite == RewriteKind::Completion { - remove_tracked_request(&tracked_out, &id, request.token); - process_completions(&mut msg); - output_router.send_value(&msg); - continue; - } - - let output = output_router.clone(); - let sanitize_completion = request.method == "completionItem/resolve"; - let mode = match request.rewrite { - RewriteKind::Locations => RewriteMode::Locations, - RewriteKind::Documentation => RewriteMode::Strings, - RewriteKind::Completion => unreachable!(), - }; - let priority = if request.method == "workspace/symbol" { - Priority::Bulk - } else { - Priority::Interactive - }; - let deadline = std::time::Instant::now() + rewrite_timeout(&request.method); - let workspace_job = (request.method == "workspace/symbol") - .then(|| (Arc::clone(&latest_workspace_out), request.token)); - if !activate_rewrite(&tracked_out, &active_out, &id, request.token) { - output_router.send_raw(raw); - continue; - } - if decompile_out.is_canceled(request.token) { - active_out.lock().unwrap().remove(&id); - decompile_out.consume_cancellation(request.token); - clear_latest_workspace(&latest_workspace_out, request.token); - continue; - } - let active = Arc::clone(&active_out); - let active_id = id; - let completion_coordinator = decompile_out.clone(); - let job = RewriteJob { - token: request.token, - message: msg, - mode, - priority, - deadline, - complete: Box::new(move |mut message| { - if sanitize_completion { - sanitize_resolved_completion(&mut message); - } - let owns_response = - active.lock().unwrap().remove(&active_id) == Some(request.token); - completion_coordinator.consume_cancellation(request.token); - if owns_response { - output.send_value(&message); - } - if let Some((latest, token)) = workspace_job { - clear_latest_workspace(&latest, token); - } - }), - }; - if let Err(job) = decompile_out.submit(job) { - let RewriteJob { - message, complete, .. - } = job; - complete(message); - } - continue; - } - } - - // Passthrough - output_router.send_raw(raw); - } - alive_out.store(false, Ordering::Relaxed); - }); + // --- Thread 1: Zed stdin -> JDTLS stdin --- + let stdin_context = StdinContext { + writer: Arc::clone(&child_stdin), + alive: Arc::clone(&alive), + tracked: Arc::clone(&tracked_ids), + active: Arc::clone(&active_rewrites), + jobs: Arc::clone(&job_counter), + decompile: decompile.clone(), + output: output.clone(), + suppressed: Arc::clone(&suppressed_responses), + latest_workspace: Arc::clone(&latest_workspace_job), + }; + thread::spawn(move || run_zed_input(stdin_context)); + + // --- Thread 2: JDTLS stdout -> Zed stdout --- + let stdout_context = StdoutContext { + pending: Arc::clone(&pending), + alive: Arc::clone(&alive), + tracked: Arc::clone(&tracked_ids), + active: Arc::clone(&active_rewrites), + decompile: decompile.clone(), + output: output.clone(), + suppressed: Arc::clone(&suppressed_responses), + latest_workspace: Arc::clone(&latest_workspace_job), + }; + let stdout_thread = thread::spawn(move || run_jdtls_output(child_stdout, stdout_context)); // --- Thread 3: HTTP server for extension requests --- let http_writer = Arc::clone(&child_stdin); @@ -433,6 +274,251 @@ fn main() { let _ = fs::remove_file(&port_file); } +fn run_zed_input(context: StdinContext) { + let stdin = io::stdin().lock(); + let mut reader = LspReader::new(BufReader::new(stdin)); + while context.alive.load(Ordering::Relaxed) { + let raw = match reader.read_message() { + Ok(Some(raw)) => raw, + Ok(None) | Err(_) => break, + }; + if route_zed_message(&context, &raw) == InputRoute::Forward + && !write_to_jdtls(&context.writer, &raw) + { + break; + } + } + context.alive.store(false, Ordering::Relaxed); +} + +fn route_zed_message(context: &StdinContext, raw: &[u8]) -> InputRoute { + let has_id = raw_has_id(raw); + if !has_id && !contains_subslice(raw, b"$/cancelRequest") { + return InputRoute::Forward; + } + let Some(message) = parse_lsp_content(raw) else { + return InputRoute::Forward; + }; + if message.get("method").and_then(Value::as_str) == Some("$/cancelRequest") { + return route_zed_cancellation(context, &message); + } + if has_id { + track_zed_request(context, &message); + } + InputRoute::Forward +} + +fn route_zed_cancellation(context: &StdinContext, message: &Value) -> InputRoute { + let Some(id) = message.pointer("/params/id") else { + return InputRoute::Forward; + }; + let already_suppressed = context.suppressed.lock().unwrap().contains(id); + let (tracked, active_token) = + take_request_for_cancellation(&context.tracked, &context.active, id); + if let Some(request) = tracked { + clear_latest_workspace(&context.latest_workspace, request.token); + } + let handled_locally = active_token.is_some(); + if let Some(token) = active_token { + clear_latest_workspace(&context.latest_workspace, token); + context.decompile.cancel(token); + context.output.send_value(&request_canceled(id)); + } + + if handled_locally || already_suppressed { + InputRoute::Consumed + } else { + InputRoute::Forward + } +} + +fn track_zed_request(context: &StdinContext, message: &Value) { + let token = context.jobs.fetch_add(1, Ordering::Relaxed); + let Some((id, request)) = tracked_request_for(message, token) else { + return; + }; + if request.method == "workspace/symbol" { + supersede_workspace_job(context, token); + } + + let (previous, active_token) = { + let mut tracked = context.tracked.lock().unwrap(); + let mut active = context.active.lock().unwrap(); + let previous = tracked.insert(id.clone(), request); + let active_token = active.remove(&id); + (previous, active_token) + }; + if let Some(previous) = previous { + clear_latest_workspace(&context.latest_workspace, previous.token); + } + if let Some(active_token) = active_token { + context.decompile.cancel(active_token); + } +} + +fn supersede_workspace_job(context: &StdinContext, token: u64) { + let previous = context.latest_workspace.lock().unwrap().replace(token); + let Some(previous) = previous else { + return; + }; + let (suppressed, active) = retire_request_token(&context.tracked, &context.active, previous); + if let Some(id) = suppressed { + context.suppressed.lock().unwrap().insert(id.clone()); + cancel_jdtls_request(&context.writer, &id); + } + if active { + context.decompile.cancel(previous); + } +} + +fn write_to_jdtls(writer: &SharedWriter, raw: &[u8]) -> bool { + let mut writer = writer.lock().unwrap(); + writer.write_all(raw).is_ok() && writer.flush().is_ok() +} + +fn run_jdtls_output(reader: impl io::Read, context: StdoutContext) { + let mut reader = LspReader::new(BufReader::new(reader)); + while let Ok(Some(raw)) = reader.read_message() { + match route_jdtls_message(&context, raw) { + OutputRoute::Raw(raw) => { + context.output.send_raw(raw); + } + OutputRoute::Value(message) => { + context.output.send_value(&message); + } + OutputRoute::Consumed => {} + } + } + context.alive.store(false, Ordering::Relaxed); +} + +fn route_jdtls_message(context: &StdoutContext, raw: Vec) -> OutputRoute { + // Notifications cannot be responses that the proxy needs to intercept. + if !raw_has_id(&raw) { + return OutputRoute::Raw(raw); + } + let Some(message) = parse_lsp_content(&raw) else { + return OutputRoute::Raw(raw); + }; + if context.pending.route(&message) { + return OutputRoute::Consumed; + } + if message + .get("id") + .is_some_and(|id| context.suppressed.lock().unwrap().take(id)) + { + return OutputRoute::Consumed; + } + if message.get("method").is_some() { + return OutputRoute::Raw(raw); + } + let Some(id) = message.get("id").cloned() else { + return OutputRoute::Raw(raw); + }; + let request = context.tracked.lock().unwrap().get(&id).cloned(); + let Some(request) = request else { + return OutputRoute::Raw(raw); + }; + + route_tracked_response(context, raw, message, id, request) +} + +fn route_tracked_response( + context: &StdoutContext, + raw: Vec, + mut message: Value, + id: Value, + request: TrackedRequest, +) -> OutputRoute { + if let Some(fallback) = completion_resolve_fallback(&message, &request) { + remove_tracked_request(&context.tracked, &id, request.token); + if should_log_completion_fallback() { + lsp_warn!( + "JDTLS completion resolution failed with -32603; \ + using the unresolved item, so documentation, imports, \ + commands, or additional edits may be missing" + ); + } + return OutputRoute::Value(fallback); + } + if message.get("error").is_some() { + remove_tracked_request(&context.tracked, &id, request.token); + clear_latest_workspace(&context.latest_workspace, request.token); + return OutputRoute::Raw(raw); + } + if request.rewrite == RewriteKind::Completion { + remove_tracked_request(&context.tracked, &id, request.token); + process_completions(&mut message); + return OutputRoute::Value(message); + } + + queue_rewrite_response(context, raw, message, id, request) +} + +fn queue_rewrite_response( + context: &StdoutContext, + raw: Vec, + message: Value, + id: Value, + request: TrackedRequest, +) -> OutputRoute { + let sanitize_completion = request.method == "completionItem/resolve"; + let mode = match request.rewrite { + RewriteKind::Locations => RewriteMode::Locations, + RewriteKind::Documentation => RewriteMode::Strings, + RewriteKind::Completion => unreachable!(), + }; + let priority = if request.method == "workspace/symbol" { + Priority::Bulk + } else { + Priority::Interactive + }; + let deadline = std::time::Instant::now() + rewrite_timeout(&request.method); + let workspace_job = (request.method == "workspace/symbol") + .then(|| (Arc::clone(&context.latest_workspace), request.token)); + if !activate_rewrite(&context.tracked, &context.active, &id, request.token) { + return OutputRoute::Raw(raw); + } + if context.decompile.is_canceled(request.token) { + context.active.lock().unwrap().remove(&id); + context.decompile.consume_cancellation(request.token); + clear_latest_workspace(&context.latest_workspace, request.token); + return OutputRoute::Consumed; + } + + let token = request.token; + let output = context.output.clone(); + let active = Arc::clone(&context.active); + let completion_coordinator = context.decompile.clone(); + let job = RewriteJob { + token, + message, + mode, + priority, + deadline, + complete: Box::new(move |mut message| { + if sanitize_completion { + sanitize_resolved_completion(&mut message); + } + let owns_response = active.lock().unwrap().remove(&id) == Some(token); + completion_coordinator.consume_cancellation(token); + if owns_response { + output.send_value(&message); + } + if let Some((latest, token)) = workspace_job { + clear_latest_workspace(&latest, token); + } + }), + }; + if let Err(job) = context.decompile.submit(job) { + let RewriteJob { + message, complete, .. + } = job; + complete(message); + } + OutputRoute::Consumed +} + // --- Utilities --- fn hex_encode(s: &str) -> String { @@ -605,6 +691,198 @@ mod tests { use super::*; use serde_json::json; + struct RoutingFixture { + writer: SharedWriter, + alive: Arc, + jobs: Arc, + pending: Arc, + tracked: TrackedRequests, + active: ActiveRewrites, + suppressed: SharedSuppressedResponses, + latest_workspace: LatestWorkspaceJob, + decompile: DecompileCoordinator, + output: Output, + } + + impl RoutingFixture { + fn new() -> Self { + let writer: SharedWriter = Arc::new(Mutex::new(Box::new(Vec::::new()))); + let pending = Arc::new(PendingResponses::new()); + let decompile = DecompileCoordinator::new( + Arc::clone(&writer), + Arc::clone(&pending), + "routing-test-".to_string(), + ); + Self { + writer, + alive: Arc::new(AtomicBool::new(true)), + jobs: Arc::new(AtomicU64::new(1)), + pending, + tracked: Arc::new(Mutex::new(HashMap::new())), + active: Arc::new(Mutex::new(HashMap::new())), + suppressed: Arc::new(Mutex::new(SuppressedResponses::default())), + latest_workspace: Arc::new(Mutex::new(None)), + decompile, + output: Output::start(), + } + } + + fn stdin_context(&self) -> StdinContext { + StdinContext { + writer: Arc::clone(&self.writer), + alive: Arc::clone(&self.alive), + tracked: Arc::clone(&self.tracked), + active: Arc::clone(&self.active), + jobs: Arc::clone(&self.jobs), + decompile: self.decompile.clone(), + output: self.output.clone(), + suppressed: Arc::clone(&self.suppressed), + latest_workspace: Arc::clone(&self.latest_workspace), + } + } + + fn stdout_context(&self) -> StdoutContext { + StdoutContext { + pending: Arc::clone(&self.pending), + alive: Arc::clone(&self.alive), + tracked: Arc::clone(&self.tracked), + active: Arc::clone(&self.active), + decompile: self.decompile.clone(), + output: self.output.clone(), + suppressed: Arc::clone(&self.suppressed), + latest_workspace: Arc::clone(&self.latest_workspace), + } + } + } + + impl Drop for RoutingFixture { + fn drop(&mut self) { + self.decompile.shutdown(); + self.output.shutdown(); + } + } + + fn frame(value: &Value) -> Vec { + encode_lsp(value).into_bytes() + } + + #[test] + fn stdin_router_tracks_requests_and_forwards_cancellation_before_rewrite() { + let fixture = RoutingFixture::new(); + let context = fixture.stdin_context(); + let request = json!({ + "jsonrpc": "2.0", + "id": 7, + "method": "textDocument/definition", + "params": {} + }); + + assert_eq!( + route_zed_message(&context, &frame(&request)), + InputRoute::Forward + ); + assert_eq!( + fixture.tracked.lock().unwrap().get(&json!(7)), + Some(&TrackedRequest::new( + 1, + "textDocument/definition", + RewriteKind::Locations, + None + )) + ); + + let cancellation = json!({ + "jsonrpc": "2.0", + "method": "$/cancelRequest", + "params": { "id": 7 } + }); + assert_eq!( + route_zed_message(&context, &frame(&cancellation)), + InputRoute::Forward + ); + assert!(fixture.tracked.lock().unwrap().is_empty()); + } + + #[test] + fn stdout_router_returns_processed_completion_values() { + let fixture = RoutingFixture::new(); + let context = fixture.stdout_context(); + fixture.tracked.lock().unwrap().insert( + json!(9), + TrackedRequest::new(1, "textDocument/completion", RewriteKind::Completion, None), + ); + let response = json!({ + "jsonrpc": "2.0", + "id": 9, + "result": [{ + "kind": 15, + "textEditText": "$TM_SELECTED_TEXT.field" + }] + }); + + assert_eq!( + route_jdtls_message(&context, frame(&response)), + OutputRoute::Value(json!({ + "jsonrpc": "2.0", + "id": 9, + "result": [{ + "kind": 15, + "textEditText": ".field" + }] + })) + ); + assert!(fixture.tracked.lock().unwrap().is_empty()); + } + + #[test] + fn stdout_router_consumes_pending_and_suppressed_responses() { + let fixture = RoutingFixture::new(); + let context = fixture.stdout_context(); + let pending_id = json!("proxy-request"); + let receiver = fixture.pending.register(pending_id.clone()); + let pending_response = json!({ "jsonrpc": "2.0", "id": pending_id, "result": "ok" }); + + assert_eq!( + route_jdtls_message(&context, frame(&pending_response)), + OutputRoute::Consumed + ); + assert_eq!(receiver.recv().unwrap(), pending_response); + + let suppressed_id = json!("superseded-request"); + fixture + .suppressed + .lock() + .unwrap() + .insert(suppressed_id.clone()); + let suppressed_response = json!({ "jsonrpc": "2.0", "id": suppressed_id, "result": null }); + assert_eq!( + route_jdtls_message(&context, frame(&suppressed_response)), + OutputRoute::Consumed + ); + assert!(!fixture + .suppressed + .lock() + .unwrap() + .contains(&json!("superseded-request"))); + } + + #[test] + fn stdio_routers_preserve_unhandled_raw_frames() { + let fixture = RoutingFixture::new(); + let stdin_context = fixture.stdin_context(); + let stdout_context = fixture.stdout_context(); + let malformed = b"Content-Length: 6\r\n\r\n{\"id\":".to_vec(); + + assert_eq!( + route_zed_message(&stdin_context, &malformed), + InputRoute::Forward + ); + assert_eq!( + route_jdtls_message(&stdout_context, malformed.clone()), + OutputRoute::Raw(malformed) + ); + } + #[test] fn tracks_location_response_methods() { let methods = [ From 63a35a549143737c65b8b68821b6e0a9dfb9dd67 Mon Sep 17 00:00:00 2001 From: Riccardo Strina Date: Mon, 3 Aug 2026 12:20:52 +0200 Subject: [PATCH 6/7] Remove owned requests on removal --- proxy/src/pending.rs | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/proxy/src/pending.rs b/proxy/src/pending.rs index 176acdf..308c8c1 100644 --- a/proxy/src/pending.rs +++ b/proxy/src/pending.rs @@ -44,20 +44,18 @@ impl PendingResponses { return false; }; - let sender = { + let (sender, owned) = { let mut state = self.state.lock().unwrap(); let sender = state.pending.remove(id); - if sender.is_none() && state.owned.contains(id) { - return true; - } - sender + let owned = state.owned.remove(id); + (sender, owned) }; if let Some(sender) = sender { let _ = sender.send(message.clone()); return true; } - false + owned } pub fn clear(&self) { @@ -81,7 +79,7 @@ mod tests { assert!(pending.route(&response)); assert_eq!(receiver.recv().unwrap(), response); - assert!(pending.route(&response)); + assert!(!pending.route(&response)); } #[test] @@ -91,11 +89,13 @@ mod tests { pending.register(retired.clone()); pending.remove(&retired); - assert!(pending.route(&json!({ + let response = json!({ "jsonrpc": "2.0", "id": "proxy-owned-late", "result": null - }))); + }); + assert!(pending.route(&response)); + assert!(!pending.route(&response)); assert!(!pending.route(&json!({ "jsonrpc": "2.0", "id": 7, @@ -121,7 +121,7 @@ mod tests { } #[test] - fn ownership_is_retained_for_the_entire_session() { + fn ownership_is_retained_until_a_response_arrives() { let pending = PendingResponses::new(); for id in 0..=1024 { let id = json!(format!("proxy-owned-{id}")); From 76200b031b07306ac75532b881e7f639c9f802cc Mon Sep 17 00:00:00 2001 From: Riccardo Strina Date: Mon, 3 Aug 2026 12:21:28 +0200 Subject: [PATCH 7/7] Refactor completions clean up to avoid duplications --- proxy/src/completions.rs | 30 ++++++++++++++++++++---------- 1 file changed, 20 insertions(+), 10 deletions(-) diff --git a/proxy/src/completions.rs b/proxy/src/completions.rs index c575fd4..3dad734 100644 --- a/proxy/src/completions.rs +++ b/proxy/src/completions.rs @@ -43,19 +43,15 @@ pub fn process_completions(msg: &mut Value) { } fn sanitize_completion_item(item: &mut Value) { - strip_tm_selected_text(item, "textEditText"); - strip_tm_selected_text(item, "insertText"); - if let Some(new_text) = item.pointer("/textEdit/newText").and_then(Value::as_str) { - if new_text.contains("$TM_SELECTED_TEXT") { - item["textEdit"]["newText"] = Value::String(new_text.replace("$TM_SELECTED_TEXT", "")); - } - } + strip_tm_selected_text(item, "/textEditText"); + strip_tm_selected_text(item, "/insertText"); + strip_tm_selected_text(item, "/textEdit/newText"); } -fn strip_tm_selected_text(item: &mut Value, key: &str) { - if let Some(text) = item.get(key).and_then(Value::as_str) { +fn strip_tm_selected_text(item: &mut Value, pointer: &str) { + if let Some(Value::String(text)) = item.pointer_mut(pointer) { if text.contains("$TM_SELECTED_TEXT") { - item[key] = Value::String(text.replace("$TM_SELECTED_TEXT", "")); + *text = text.replace("$TM_SELECTED_TEXT", ""); } } } @@ -129,6 +125,20 @@ mod tests { ); } + #[test] + fn sanitizes_text_edit_text() { + let mut response = json!({ + "result": [{ + "kind": 15, + "textEditText": "$TM_SELECTED_TEXT.field" + }] + }); + + process_completions(&mut response); + + assert_eq!(response["result"][0]["textEditText"], json!(".field")); + } + #[test] fn sanitizes_insert_replace_edit_text() { let mut response = json!({