From cb9b8ae4868bf0dad8a3e5f593bb48bb10e83c1b Mon Sep 17 00:00:00 2001 From: rossirpaulo Date: Fri, 10 Jul 2026 17:15:00 -0700 Subject: [PATCH 1/2] feat(lsp): Phase 1 - bounded ingress, cancellation ownership, client gating Rust (stacked on the Phase 0 core): - IngressScheduler (transport-neutral admission/lifecycle/cancellation/ budget state machine) relocated into baml_lsp_server, with fixes: forced cancellation of protected methods can no longer be downgraded; permanent oversize is contained per-message (-32803) instead of closing the session; oversized $/cancelRequest no longer livelocks the control lane; pipelined `initialized` is admitted FIFO and accepted at dispatch. - New LspRuntime: stdio and browser LSP sessions dispatch through one scheduler; cancellation drained on the transport thread claims responses while handlers run ($/cancelRequest -> -32800); browser takeover revokes the old session's sink and replays document overlays; bounded outbound budgets with serialize-once shared frames. - Bounded stdio framing: 16 KiB header cap; oversized bodies discarded in chunks with a recoverable typed error; abnormal exit is nonzero. - bex_project kept minimal: LspError::RequestCanceled (-32800) at the one serialization boundary (I7); dispatch-scoped cancellation token checked at handler entry and after the source guard (never Salsa unwinding); connection-scoped LSP sessions (fresh encoding negotiation + workspace roots per connection, shared project registry). TypeScript: - D6 playground run-gating on is_bex_current: Run/test/preview disabled with a "Preparing current build..." state during rebuild windows; projectNotReady rejections render as that transient state and clear on the next current ProjectUpdate (no protocol version bump). - WebSocketRuntimePort reconnect correctness: fail-closed hello, bounded first-connect queue, session-scoped commands dropped on disconnect with synthesized rejections, desired-runtime lease resent after every hello. - Track A: projectRoots.ts canonical multi-root ownership + routing in the VS Code extension; closing the last document no longer stops the owner server; aggregate status preserves startup failures. Verification: baml_lsp_server 67/67, bex_project 45/45 (+1 pre-existing ignored), clippy -D warnings clean, wasm check clean; playground 125/125, extension 13/13, webview 19/19, all typechecks green. Harness (400 fns): $/cancelRequest now answers -32800; typing inlay p50 2.0 ms, p95 51.7 ms, zero errors, zero unanswered requests. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01VEMxti5WQ8bCgJUuX9AQvo --- baml_language/Cargo.lock | 1 + .../crates/baml_lsp_server/Cargo.toml | 1 + .../crates/baml_lsp_server/src/lib.rs | 584 +++- .../crates/baml_lsp_server/src/lsp_ingress.rs | 2475 +++++++++++++++++ .../crates/baml_lsp_server/src/lsp_runtime.rs | 1405 ++++++++++ .../baml_lsp_server/src/native_lsp_sender.rs | 107 +- .../baml_lsp_server/src/playground_server.rs | 283 +- .../crates/bex_project/src/bex_lsp/mod.rs | 45 +- .../src/bex_lsp/multi_project/mod.rs | 140 +- .../bex_project/src/bex_lsp/notification.rs | 14 +- .../crates/bex_project/src/bex_lsp/request.rs | 157 +- .../src/bex_lsp/request_cancellation.rs | 80 + .../src/playground/baml-lsp-worker.ts | 6 + .../src/__tests__/projectRoots.test.ts | 184 ++ typescript2/app-vscode-ext/src/extension.ts | 561 +++- .../app-vscode-ext/src/projectRoots.ts | 311 +++ .../src/ExecutionPanel.gating.test.tsx | 163 ++ .../pkg-playground/src/ExecutionPanel.tsx | 132 +- .../pkg-playground/src/FunctionSidebar.tsx | 16 +- .../src/ports/WebSocketRuntimePort.test.ts | 404 +++ .../src/ports/WebSocketRuntimePort.ts | 236 +- .../pkg-playground/src/run-gating.test.ts | 92 + typescript2/pkg-playground/src/run-gating.ts | 65 + .../src/run-store-client.test.ts | 36 +- .../pkg-playground/src/run-store-client.ts | 27 +- .../pkg-playground/src/worker-protocol.ts | 27 + 26 files changed, 7272 insertions(+), 280 deletions(-) create mode 100644 baml_language/crates/baml_lsp_server/src/lsp_ingress.rs create mode 100644 baml_language/crates/baml_lsp_server/src/lsp_runtime.rs create mode 100644 baml_language/crates/bex_project/src/bex_lsp/request_cancellation.rs create mode 100644 typescript2/app-vscode-ext/src/__tests__/projectRoots.test.ts create mode 100644 typescript2/app-vscode-ext/src/projectRoots.ts create mode 100644 typescript2/app-vscode-webview/src/ExecutionPanel.gating.test.tsx create mode 100644 typescript2/pkg-playground/src/ports/WebSocketRuntimePort.test.ts create mode 100644 typescript2/pkg-playground/src/run-gating.test.ts create mode 100644 typescript2/pkg-playground/src/run-gating.ts diff --git a/baml_language/Cargo.lock b/baml_language/Cargo.lock index d51461f7beb..61d9f4749ba 100644 --- a/baml_language/Cargo.lock +++ b/baml_language/Cargo.lock @@ -1043,6 +1043,7 @@ dependencies = [ "sys_native", "sys_ops", "sys_types", + "thiserror 2.0.18", "tokio", "tokio-tungstenite", "tower-http", diff --git a/baml_language/crates/baml_lsp_server/Cargo.toml b/baml_language/crates/baml_lsp_server/Cargo.toml index 49356935b55..913c51fa2da 100644 --- a/baml_language/crates/baml_lsp_server/Cargo.toml +++ b/baml_language/crates/baml_lsp_server/Cargo.toml @@ -54,6 +54,7 @@ reqwest = { workspace = true } rustls = { workspace = true, optional = true } serde = { workspace = true } serde_json = { workspace = true } +thiserror = { workspace = true } tokio = { workspace = true, features = [ "sync", "rt-multi-thread", "net", "time", "macros" ] } tokio-tungstenite = { workspace = true } tower-http = { workspace = true } diff --git a/baml_language/crates/baml_lsp_server/src/lib.rs b/baml_language/crates/baml_lsp_server/src/lib.rs index 1c9ef8bd59c..e89cd30d249 100644 --- a/baml_language/crates/baml_lsp_server/src/lib.rs +++ b/baml_language/crates/baml_lsp_server/src/lib.rs @@ -33,6 +33,8 @@ //! `--all-features` (both enabled); prefer one when building the LSP binary. mod deadlock_watchdog; +pub mod lsp_ingress; +mod lsp_runtime; mod native_lsp_sender; mod native_vfs; pub mod playground_env; @@ -47,6 +49,7 @@ pub mod playground_ws; use std::{ collections::BTreeMap, fs, + io::{BufRead, Write}, path::{Path, PathBuf}, sync::Arc, time::{Duration, UNIX_EPOCH}, @@ -60,6 +63,292 @@ use playground_session::PlaygroundSessionStore; use playground_ws::WsOutMessage; use tokio::net::TcpListener; +// --------------------------------------------------------------------------- +// Bounded outbound frames (B2: no transport hides an unbounded writer queue) +// --------------------------------------------------------------------------- + +const OUTBOUND_QUEUE_BYTES: usize = 64 * 1024 * 1024; +pub(crate) const MAX_OUTBOUND_FRAME_BYTES: usize = 4 * 1024 * 1024; + +/// One serialized outbound JSON-RPC frame. +/// +/// The body is serialized exactly once — with the `jsonrpc` member included, +/// so the same bytes are valid for stdio Content-Length framing and for a +/// WebSocket text frame — into shared `Arc<[u8]>` storage carrying one budget +/// charge. Clones, including the per-receiver clones made by +/// `tokio::sync::broadcast`, share the allocation *and* the charge, so the +/// budget accounts real memory exactly instead of deep-cloning the payload +/// per receiver while charging once. +#[derive(Debug, Clone)] +pub(crate) struct OutboundFrame { + bytes: Arc<[u8]>, + is_response: bool, + _charge: Arc, +} + +impl OutboundFrame { + pub(crate) fn bytes(&self) -> &[u8] { + &self.bytes + } + + /// Whether the frame is a JSON-RPC response. Response routing is owned by + /// the per-session runtime path; broadcast consumers must skip these. + pub(crate) fn is_response(&self) -> bool { + self.is_response + } +} + +#[derive(Debug)] +pub(crate) struct OutboundBudget { + used: std::sync::atomic::AtomicUsize, + limit: usize, + max_frame: usize, +} + +#[derive(Debug)] +struct OutboundCharge { + budget: Arc, + bytes: usize, +} + +impl Drop for OutboundCharge { + fn drop(&mut self) { + self.budget + .used + .fetch_sub(self.bytes, std::sync::atomic::Ordering::AcqRel); + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) enum OutboundReserveError { + Serialization, + Oversized, + Saturated, +} + +impl OutboundBudget { + pub(crate) fn new() -> Arc { + Arc::new(Self { + used: std::sync::atomic::AtomicUsize::new(0), + limit: OUTBOUND_QUEUE_BYTES, + max_frame: MAX_OUTBOUND_FRAME_BYTES, + }) + } + + pub(crate) fn try_message( + self: &Arc, + message: lsp_server::Message, + ) -> Result { + let is_response = matches!(&message, lsp_server::Message::Response(_)); + let bytes = + serialize_jsonrpc_message(&message).map_err(|_| OutboundReserveError::Serialization)?; + self.try_reserve(bytes, is_response) + } + + /// Raw pre-built JSON (transport-level null-ID protocol errors). + fn try_raw(self: &Arc, value: &serde_json::Value) -> Option { + let bytes = serde_json::to_vec(value).ok()?; + self.try_reserve(bytes, true).ok() + } + + fn try_reserve( + self: &Arc, + bytes: Vec, + is_response: bool, + ) -> Result { + let len = bytes.len(); + if len > self.max_frame { + return Err(OutboundReserveError::Oversized); + } + let mut used = self.used.load(std::sync::atomic::Ordering::Acquire); + loop { + if used.saturating_add(len) > self.limit { + return Err(OutboundReserveError::Saturated); + } + match self.used.compare_exchange_weak( + used, + used + len, + std::sync::atomic::Ordering::AcqRel, + std::sync::atomic::Ordering::Acquire, + ) { + Ok(_) => break, + Err(observed) => used = observed, + } + } + Ok(OutboundFrame { + bytes: bytes.into(), + is_response, + _charge: Arc::new(OutboundCharge { + budget: self.clone(), + bytes: len, + }), + }) + } +} + +/// Serialize with the `jsonrpc` member (plain `serde_json::to_vec` of an +/// `lsp_server::Message` omits it; only the crate's stdio writer adds it). +fn serialize_jsonrpc_message(message: &lsp_server::Message) -> serde_json::Result> { + let mut value = serde_json::to_value(message)?; + if let Some(object) = value.as_object_mut() { + object.insert( + "jsonrpc".to_string(), + serde_json::Value::String("2.0".to_string()), + ); + } + serde_json::to_vec(&value) +} + +fn write_frame(output: &mut impl Write, frame: &OutboundFrame) -> std::io::Result<()> { + write!(output, "Content-Length: {}\r\n\r\n", frame.bytes().len())?; + output.write_all(frame.bytes())?; + output.flush() +} + +// --------------------------------------------------------------------------- +// Bounded stdio framing (B2 transport adapter: parse/frame only) +// --------------------------------------------------------------------------- + +/// Per-message body cap. Larger than every ingress class budget, so any frame +/// that passes here is judged (and, if needed, rejected per-message) by the +/// scheduler's admission rather than by the transport. +const MAX_STDIO_FRAME_BYTES: usize = 8 * 1024 * 1024; + +/// Total header-block cap. Headers are read through a `Take` so a peer that +/// never sends CRLF cannot grow an unbounded header line in memory. +const MAX_STDIO_HEADER_BYTES: u64 = 16 * 1024; + +#[derive(Debug, PartialEq, Eq)] +enum StdioReadError { + Parse(String), + InvalidRequest(String), + Framing(String), + /// `Content-Length` exceeded [`MAX_STDIO_FRAME_BYTES`]. The body has + /// already been consumed and discarded in bounded chunks, so the stream + /// stays in sync and the session stays alive (recoverable per-message). + OversizedBody { + content_length: usize, + }, +} + +fn decode_lsp_message(value: serde_json::Value) -> Result { + let Some(object) = value.as_object() else { + return Err("JSON-RPC envelope must be an object".to_string()); + }; + if object.get("jsonrpc").and_then(serde_json::Value::as_str) != Some("2.0") { + return Err("JSON-RPC envelope must contain jsonrpc: \"2.0\"".to_string()); + } + let has_method = object + .get("method") + .is_some_and(serde_json::Value::is_string); + let has_result = object.contains_key("result"); + let has_error = object.contains_key("error"); + if has_method { + if has_result || has_error { + return Err("JSON-RPC request/notification cannot contain result or error".to_string()); + } + if object + .get("id") + .is_some_and(|id| !(id.is_string() || id.as_i64().is_some() || id.as_u64().is_some())) + { + return Err("JSON-RPC request id must be a string or integer".to_string()); + } + } else { + if has_result == has_error { + return Err( + "JSON-RPC response must contain exactly one of result or error".to_string(), + ); + } + let Some(id) = object.get("id") else { + return Err("JSON-RPC response is missing id".to_string()); + }; + if !(id.is_string() || id.as_i64().is_some() || id.as_u64().is_some()) { + return Err("JSON-RPC response id must be a string or integer".to_string()); + } + } + serde_json::from_value(value).map_err(|error| format!("Invalid request: {error}")) +} + +fn read_lsp_message( + input: &mut impl BufRead, +) -> Result, StdioReadError> { + let mut content_length = None; + let mut header = String::new(); + let mut read_any_header = false; + let mut remaining_header_bytes = MAX_STDIO_HEADER_BYTES; + loop { + header.clear(); + let read = std::io::Read::take(&mut *input, remaining_header_bytes) + .read_line(&mut header) + .map_err(|error| StdioReadError::Framing(error.to_string()))?; + remaining_header_bytes -= read as u64; + if read == 0 { + return if read_any_header { + Err(StdioReadError::Framing(if remaining_header_bytes == 0 { + format!("LSP header block exceeds {MAX_STDIO_HEADER_BYTES} bytes") + } else { + "unexpected EOF in LSP headers".to_string() + })) + } else { + Ok(None) + }; + } + read_any_header = true; + let Some(line) = header.strip_suffix("\r\n") else { + // Either malformed framing or the header cap truncated mid-line. + return Err(StdioReadError::Framing(if remaining_header_bytes == 0 { + format!("LSP header block exceeds {MAX_STDIO_HEADER_BYTES} bytes") + } else { + "LSP header must end in CRLF".to_string() + })); + }; + if line.is_empty() { + break; + } + let Some((name, value)) = line.split_once(": ") else { + return Err(StdioReadError::Framing(format!( + "malformed LSP header: {line}" + ))); + }; + if name.eq_ignore_ascii_case("Content-Length") { + let parsed = value + .parse::() + .map_err(|_| StdioReadError::Framing("invalid Content-Length".to_string()))?; + content_length = Some(parsed); + } + } + let content_length = content_length + .ok_or_else(|| StdioReadError::Framing("missing Content-Length".to_string()))?; + if content_length > MAX_STDIO_FRAME_BYTES { + // Consume the body in bounded chunks so the next frame parses from a + // clean boundary; the caller answers with a per-message error and the + // session survives. + discard_exact(input, content_length)?; + return Err(StdioReadError::OversizedBody { content_length }); + } + let mut body = vec![0; content_length]; + input + .read_exact(&mut body) + .map_err(|error| StdioReadError::Framing(error.to_string()))?; + let value: serde_json::Value = serde_json::from_slice(&body) + .map_err(|error| StdioReadError::Parse(format!("Parse error: {error}")))?; + decode_lsp_message(value) + .map(Some) + .map_err(StdioReadError::InvalidRequest) +} + +fn discard_exact(input: &mut impl BufRead, mut remaining: usize) -> Result<(), StdioReadError> { + let mut chunk = [0u8; 64 * 1024]; + while remaining > 0 { + let take = remaining.min(chunk.len()); + input + .read_exact(&mut chunk[..take]) + .map_err(|error| StdioReadError::Framing(error.to_string()))?; + remaining -= take; + } + Ok(()) +} + pub fn version() -> &'static str { baml_version::CANONICAL_VERSION } @@ -214,18 +503,25 @@ fn run_server_inner( Arc::new(Box::new(native_vfs::NativeVfs::new())); let baml_vfs = bex_project::BamlVFS::new(vfs); - // Stdio sender (LSP client sender) - let (writer_tx, writer_rx) = crossbeam_channel::unbounded::(); + // Stdio sender (LSP client sender): bounded frames charged against one + // process outbound budget (B2 — no unbounded writer queue). + let (writer_tx, writer_rx) = crossbeam_channel::bounded::(512); let writer_tx = Arc::new(writer_tx); - let lsp_sender: Arc = - Arc::new(native_lsp_sender::NativeLspSender::new(&writer_tx)); + let writer_budget = OutboundBudget::new(); + let lsp_sender: Arc = Arc::new( + native_lsp_sender::NativeLspSender::new(&writer_tx, &writer_budget), + ); // Browser-mode LSP transport: the standalone playground has no stdio LSP // client, so in browser mode we forward all LSP output (responses + // publishDiagnostics) to the `/api/lsp` WebSocket instead of stdout. A // bridge thread (spawned only in browser mode) drains `writer_rx` into // this broadcast channel. - let (lsp_out_tx, _lsp_out_rx) = tokio::sync::broadcast::channel::(256); + let (lsp_out_tx, _lsp_out_rx) = tokio::sync::broadcast::channel::(256); + + // Process-owned ingress runtime: one bounded scheduler + one dispatch + // worker shared by the stdio loop and every `/api/lsp` browser session. + let lsp_runtime = lsp_runtime::LspRuntime::new()?; // Mirror of the content the browser editor currently has per file. Shared // between the `/api/lsp` bridge (writes it on didOpen/didChange) and the disk @@ -344,8 +640,8 @@ fn run_server_inner( std::thread::Builder::new() .name("lsp-ws-bridge".into()) .spawn(move || { - while let Ok(msg) = writer_rx.recv() { - let _ = lsp_out_tx_bridge.send(msg); + while let Ok(frame) = writer_rx.recv() { + let _ = lsp_out_tx_bridge.send(frame); } })?; @@ -355,6 +651,7 @@ fn run_server_inner( let _disk_watcher = playground_server::spawn_disk_watcher( &workspace_roots, lsp_out_tx.clone(), + writer_budget.clone(), doc_mirror.clone(), ); @@ -367,6 +664,7 @@ fn run_server_inner( runs, playground_dir, lsp_out_tx, + lsp_runtime.clone(), doc_mirror, workspace_roots.clone(), session_token, @@ -375,6 +673,7 @@ fn run_server_inner( } let workspace_roots_for_server = workspace_roots.clone(); + let lsp_runtime_for_server = lsp_runtime.clone(); tokio_runtime.spawn(async move { if let Err(e) = playground_server::run( listener, @@ -385,6 +684,7 @@ fn run_server_inner( runs, playground_dir, lsp_out_tx, + lsp_runtime_for_server, doc_mirror, workspace_roots_for_server, session_token, @@ -399,60 +699,133 @@ fn run_server_inner( anyhow::bail!("Could not start playground server"); } + // The stdio session: a bounded sink into the writer channel. Saturation + // is backpressure (the response stays reserved and is retried), never + // silent loss; a disconnected writer closes the session. + let stdio_closed = Arc::new(std::sync::atomic::AtomicBool::new(false)); + let stdio_sink_tx = writer_tx.clone(); + let stdio_sink_budget = writer_budget.clone(); + let stdio_sink: lsp_runtime::Sink = Arc::new(move |message| { + let frame = match stdio_sink_budget.try_message(message) { + Ok(frame) => frame, + Err(OutboundReserveError::Saturated) => { + return lsp_runtime::SinkDelivery::Saturated; + } + Err(OutboundReserveError::Oversized | OutboundReserveError::Serialization) => { + return lsp_runtime::SinkDelivery::Oversized; + } + }; + match stdio_sink_tx.try_send(frame) { + Ok(()) => lsp_runtime::SinkDelivery::Sent, + Err(crossbeam_channel::TrySendError::Full(_)) => lsp_runtime::SinkDelivery::Saturated, + Err(crossbeam_channel::TrySendError::Disconnected(_)) => { + lsp_runtime::SinkDelivery::Closed + } + } + }); + let stdio_closed_for_endpoint = stdio_closed.clone(); + let stdio_close: lsp_runtime::Close = Arc::new(move || { + stdio_closed_for_endpoint.store(true, std::sync::atomic::Ordering::Release); + }); + let stdio_session = lsp_runtime + .open_session( + lsp_ingress::TransportKind::Stdio, + bex.clone(), + stdio_sink, + stdio_close, + None, + ) + .session_id; + // Spawn the stdout writer thread. std::thread::Builder::new() .name("lsp-stdout-writer".into()) .spawn(move || { let stdout = std::io::stdout(); let mut stdout = stdout.lock(); - while let Ok(msg) = writer_rx.recv() { - if msg.write(&mut stdout).is_err() { + while let Ok(frame) = writer_rx.recv() { + if write_frame(&mut stdout, &frame).is_err() { break; } } })?; - // Main event loop: read from stdin, dispatch to bex_project. + // Main event loop: bounded framing (capped headers, per-message body + // rejection) feeding the shared ingress runtime. Lifecycle — including + // shutdown/exit — is owned by the scheduler; no transport shortcuts. let stdin = std::io::stdin(); let mut stdin = stdin.lock(); - // Main event loop — forward all messages to bex_project. - // The `initialize` handshake is handled by `bex_project` via `handle_request`. + let mut abnormal_exit = false; loop { - let msg = match lsp_server::Message::read(&mut stdin) { + let msg = match read_lsp_message(&mut stdin) { Ok(Some(msg)) => msg, Ok(None) => break, - Err(e) => { - tracing::error!("Failed to read LSP message: {e}"); - break; + Err(error) => { + let (code, message, recoverable) = match error { + StdioReadError::Parse(message) => (-32700, message, true), + StdioReadError::InvalidRequest(message) => (-32600, message, true), + // The body was discarded without buffering, so the id is + // unknown: a null-ID error is the best per-message answer + // available. Frames within MAX_STDIO_FRAME_BYTES but over + // an ingress class budget do carry their id and get a + // typed per-request rejection from admission instead. + StdioReadError::OversizedBody { content_length } => ( + -32803, + format!( + "dropped LSP frame of {content_length} bytes \ + (limit {MAX_STDIO_FRAME_BYTES}); session stays open" + ), + true, + ), + StdioReadError::Framing(message) => (-32700, message, false), + }; + let response = serde_json::json!({ + "jsonrpc": "2.0", + "id": serde_json::Value::Null, + "error": { "code": code, "message": message }, + }); + let queued = writer_budget + .try_raw(&response) + .is_some_and(|frame| writer_tx.try_send(frame).is_ok()); + if !queued || !recoverable { + break; + } + continue; } }; - match msg { - lsp_server::Message::Notification(notification) => { - tracing::debug!("<<< notification: {}", notification.method); - if notification.method == "exit" { + let mut terminate = false; + loop { + match lsp_runtime.submit(stdio_session, msg.clone()) { + lsp_runtime::SubmitResult::Accepted | lsp_runtime::SubmitResult::Dropped => break, + lsp_runtime::SubmitResult::Backpressure => { + // Reads are rejected under overload (D3); only mutation/ + // lifecycle reserve pressure stalls the reader briefly. + std::thread::sleep(Duration::from_millis(2)); + } + lsp_runtime::SubmitResult::Exited { normal } => { + abnormal_exit = !normal; + terminate = true; break; } - bex.handle_notification(notification); - } - lsp_server::Message::Request(request) => { - tracing::debug!("<<< request: {} (id={})", request.method, request.id); - if request.method == "shutdown" { - let response = lsp_server::Response { - id: request.id, - result: Some(serde_json::Value::Null), - error: None, - }; - let _ = writer_tx.send(lsp_server::Message::Response(response)); - continue; + lsp_runtime::SubmitResult::Closed => { + terminate = true; + break; } - bex.handle_request(request); - } - lsp_server::Message::Response(response) => { - tracing::debug!("<<< response from client: {:?}", response.id); } } + if terminate || stdio_closed.load(std::sync::atomic::Ordering::Acquire) { + break; + } + } + + lsp_runtime.close_session(stdio_session); + + if abnormal_exit { + // B2 lifecycle: `exit` before a completed shutdown is an abnormal + // termination (nonzero for stdio). + anyhow::bail!("LSP client sent exit before completing shutdown"); } tracing::info!("LSP server shutting down"); @@ -627,6 +1000,145 @@ fn file_signature(metadata: &fs::Metadata) -> FileSignature { mod tests { use super::*; + fn framed(body: &str) -> Vec { + format!("Content-Length: {}\r\n\r\n{body}", body.len()).into_bytes() + } + + #[test] + fn stdio_parser_distinguishes_parse_and_invalid_request() { + let mut malformed = std::io::Cursor::new(framed("{")); + assert!(matches!( + read_lsp_message(&mut malformed), + Err(StdioReadError::Parse(_)) + )); + + let mut invalid = std::io::Cursor::new(framed(r#"{"jsonrpc":"2.0","wat":true}"#)); + assert!(matches!( + read_lsp_message(&mut invalid), + Err(StdioReadError::InvalidRequest(_)) + )); + + let mut request = std::io::Cursor::new(framed( + r#"{"jsonrpc":"2.0","id":1,"method":"initialize","params":{}}"#, + )); + assert!(matches!( + read_lsp_message(&mut request), + Ok(Some(lsp_server::Message::Request(_))) + )); + } + + /// Defect containment: an unbounded header line (no CRLF ever) must fail + /// at the header cap instead of growing an unbounded String. + #[test] + fn stdio_header_block_is_capped() { + let mut endless = std::io::Cursor::new(vec![b'X'; 64 * 1024]); + let error = read_lsp_message(&mut endless).unwrap_err(); + let StdioReadError::Framing(message) = error else { + panic!("expected a framing error, got {error:?}"); + }; + assert!(message.contains("header block exceeds"), "{message}"); + + // Many small headers also hit the cap. + let mut headers = String::new(); + for index in 0..2000 { + headers.push_str(&format!("X-Filler-{index}: value\r\n")); + } + headers.push_str("\r\n"); + let mut many = std::io::Cursor::new(headers.into_bytes()); + assert!(matches!( + read_lsp_message(&mut many), + Err(StdioReadError::Framing(_)) + )); + } + + /// Defect containment: an oversized body is discarded in bounded chunks + /// and reading continues with the next frame — the session stays alive. + #[test] + fn oversized_stdio_body_is_recoverable_per_message() { + let huge_length = MAX_STDIO_FRAME_BYTES + 1; + let mut stream = format!("Content-Length: {huge_length}\r\n\r\n").into_bytes(); + stream.extend(std::iter::repeat_n(b'x', huge_length)); + stream.extend(framed( + r#"{"jsonrpc":"2.0","id":7,"method":"shutdown","params":null}"#, + )); + let mut input = std::io::Cursor::new(stream); + + assert!(matches!( + read_lsp_message(&mut input), + Err(StdioReadError::OversizedBody { content_length }) if content_length == huge_length + )); + // The stream is still in sync: the next frame parses normally. + let Ok(Some(lsp_server::Message::Request(request))) = read_lsp_message(&mut input) else { + panic!("the frame after an oversized body must parse"); + }; + assert_eq!(request.method, "shutdown"); + } + + #[test] + fn outbound_frames_share_bytes_and_one_budget_charge() { + let budget = OutboundBudget::new(); + let frame = budget + .try_message(lsp_server::Message::Notification( + lsp_server::Notification::new( + "window/logMessage".to_string(), + serde_json::json!({ "type": 3, "message": "hello" }), + ), + )) + .unwrap(); + let used_with_one = budget.used.load(std::sync::atomic::Ordering::Acquire); + assert_eq!(used_with_one, frame.bytes().len()); + assert!(!frame.is_response()); + // Serialized once, with the jsonrpc member for both transports. + let value: serde_json::Value = serde_json::from_slice(frame.bytes()).unwrap(); + assert_eq!(value["jsonrpc"], "2.0"); + + // Broadcast-style clones share the allocation and the charge. + let clone = frame.clone(); + assert_eq!( + budget.used.load(std::sync::atomic::Ordering::Acquire), + used_with_one + ); + drop(frame); + assert_eq!( + budget.used.load(std::sync::atomic::Ordering::Acquire), + used_with_one + ); + drop(clone); + assert_eq!(budget.used.load(std::sync::atomic::Ordering::Acquire), 0); + } + + #[test] + fn oversized_outbound_frame_is_rejected_by_the_budget() { + let budget = OutboundBudget::new(); + let oversized = lsp_server::Message::Notification(lsp_server::Notification::new( + "test/oversized".to_string(), + serde_json::Value::String("x".repeat(MAX_OUTBOUND_FRAME_BYTES + 1)), + )); + assert_eq!( + budget.try_message(oversized).unwrap_err(), + OutboundReserveError::Oversized + ); + assert_eq!(budget.used.load(std::sync::atomic::Ordering::Acquire), 0); + } + + #[test] + fn raw_stdio_error_uses_null_id_and_content_length() { + let value = serde_json::json!({ + "jsonrpc": "2.0", + "id": serde_json::Value::Null, + "error": { "code": -32700, "message": "Parse error" }, + }); + let budget = OutboundBudget::new(); + let frame = budget.try_raw(&value).unwrap(); + let mut output = Vec::new(); + write_frame(&mut output, &frame).unwrap(); + let output = String::from_utf8(output).unwrap(); + assert!(output.starts_with("Content-Length: ")); + assert!(output.contains("\r\n\r\n")); + assert!(output.contains(r#""id":null"#)); + assert!(output.contains(r#""code":-32700"#)); + } + #[test] fn playground_banner_shows_url_project_root_and_ssh_hint() { let banner = format_playground_banner(4265, &[PathBuf::from("/home/dev/my-app")], None); diff --git a/baml_language/crates/baml_lsp_server/src/lsp_ingress.rs b/baml_language/crates/baml_lsp_server/src/lsp_ingress.rs new file mode 100644 index 00000000000..6343584039e --- /dev/null +++ b/baml_language/crates/baml_lsp_server/src/lsp_ingress.rs @@ -0,0 +1,2475 @@ +//! Process-scoped, transport-neutral LSP ingress scheduling. +//! +//! The scheduler is deliberately synchronous: callers put one instance behind +//! their process-owned mutex and let both stdio and browser adapters submit the +//! same [`lsp_server::Message`] representation. It owns admission ordering, +//! lifecycle validation, request response ownership, and bounded memory +//! accounting; handlers still own the actual LSP operation. +//! +//! No cancellation token returned by this module may be connected to Salsa's +//! unwind-based local cancellation. It is an operation-owned signal intended +//! for checks at safe handler boundaries only. + +use std::{ + collections::{BTreeMap, HashMap, HashSet, VecDeque}, + sync::Arc, +}; + +use lsp_server::{Message, RequestId}; +use sys_types::CancellationToken; + +/// Monotonic identity for one transport connection. +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct SessionId(u64); + +impl SessionId { + /// The process-local numeric identity. This is useful for logging only. + #[must_use] + pub const fn get(self) -> u64 { + self.0 + } +} + +/// Request response identity. Request IDs may be reused by a replacement +/// browser connection, so a bare [`RequestId`] is never sufficient. +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct ResponseToken { + session_id: SessionId, + request_id: RequestId, +} + +impl ResponseToken { + #[must_use] + pub const fn session_id(&self) -> SessionId { + self.session_id + } + + #[must_use] + pub const fn request_id(&self) -> &RequestId { + &self.request_id + } +} + +/// Sequence assigned to each admitted non-cancellation message. +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct IngressSequence(u64); + +impl IngressSequence { + #[must_use] + pub const fn get(self) -> u64 { + self.0 + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum TransportKind { + Stdio, + Browser, +} + +/// Per-connection lifecycle. `Superseded` is reported in takeover effects; +/// superseded sessions themselves are removed so reconnect churn cannot grow +/// an unbounded tombstone map. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum LifecycleState { + PreInitialize, + InitializeResponding, + InitializeResponded, + Initialized, + ShutdownResponding, + Shutdown, + Exited, + Superseded, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum AdmissionClass { + /// Read-only editor queries and formatting. These use the smaller read + /// budget and are rejected rather than backpressuring mutation traffic. + Read, + /// A lifecycle barrier such as initialize, initialized, or shutdown. + Lifecycle, + /// A source/workspace mutation barrier. + Mutation, + /// A request whose handler may commit an external side effect. + SideEffecting, + /// An ordinary notification that is neither a mutation nor lifecycle. + Other, +} + +impl AdmissionClass { + const fn is_reserved(self) -> bool { + matches!(self, Self::Lifecycle | Self::Mutation | Self::SideEffecting) + } +} + +/// What a running handler should do when cancellation wins response ownership. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum CancellationBehavior { + /// Claim the response and signal the operation-owned token. + SignalAtSafePoints, + /// Claim the response, but let already-started effects finish normally. + ClaimResponseOnly, + /// Once running, cancellation is a no-op. Queued work is still removable. + NonCancellableWhenRunning, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct MessagePolicy { + pub class: AdmissionClass, + pub cancellation: CancellationBehavior, +} + +impl MessagePolicy { + #[must_use] + pub const fn read() -> Self { + Self { + class: AdmissionClass::Read, + cancellation: CancellationBehavior::SignalAtSafePoints, + } + } + + #[must_use] + pub const fn side_effecting(cancellation: CancellationBehavior) -> Self { + Self { + class: AdmissionClass::SideEffecting, + cancellation, + } + } + + /// Shared default classification used by both transports. Integrators may + /// override ordinary custom requests, but known lifecycle and mutation + /// methods are always upgraded to their protected class. + #[must_use] + pub fn for_message(message: &Message) -> Self { + match message { + Message::Request(request) => match request.method.as_str() { + "initialize" | "shutdown" => Self { + class: AdmissionClass::Lifecycle, + // Once dispatched these handlers freeze session config or + // apply shutdown side effects. Queued cancellation still + // works, but running lifecycle barriers are committed. + cancellation: CancellationBehavior::NonCancellableWhenRunning, + }, + "workspace/executeCommand" => { + // Dispatch is the linearization point for command side + // effects. A cancellation that wins while the command is + // queued removes it and owns the response; once dispatch + // begins, cancellation must not report RequestCanceled + // while the handler can still commit its effect. + Self::side_effecting(CancellationBehavior::NonCancellableWhenRunning) + } + _ => Self::read(), + }, + Message::Notification(notification) => { + let class = match notification.method.as_str() { + "initialized" | "exit" => AdmissionClass::Lifecycle, + "textDocument/didOpen" + | "textDocument/didChange" + | "textDocument/didClose" + | "textDocument/didSave" + | "workspace/didChangeConfiguration" + | "workspace/didChangeWatchedFiles" + | "workspace/didChangeWorkspaceFolders" + | "workspace/didCreateFiles" + | "workspace/didRenameFiles" + | "workspace/didDeleteFiles" => AdmissionClass::Mutation, + _ => AdmissionClass::Other, + }; + Self { + class, + cancellation: CancellationBehavior::NonCancellableWhenRunning, + } + } + Message::Response(_) => Self { + class: AdmissionClass::Other, + cancellation: CancellationBehavior::NonCancellableWhenRunning, + }, + } + } +} + +/// Independent normal, reserve, control, read, and outbound-response budgets. +/// A request reserves `response_reservation_bytes` until its response has been +/// ordered onto the owning session's sink. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct IngressLimits { + pub normal_items: usize, + pub normal_bytes: usize, + pub reserved_items: usize, + pub reserved_bytes: usize, + pub read_items: usize, + pub read_bytes: usize, + pub control_items: usize, + pub control_bytes: usize, + pub outbound_response_items: usize, + pub outbound_response_bytes: usize, + pub response_reservation_bytes: usize, +} + +impl Default for IngressLimits { + fn default() -> Self { + Self { + normal_items: 256, + normal_bytes: 4 * 1024 * 1024, + reserved_items: 64, + reserved_bytes: 2 * 1024 * 1024, + read_items: 128, + read_bytes: 2 * 1024 * 1024, + control_items: 64, + control_bytes: 64 * 1024, + outbound_response_items: 256, + outbound_response_bytes: 4 * 1024 * 1024, + response_reservation_bytes: 16 * 1024, + } + } +} + +#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)] +pub enum LimitsError { + #[error("all ingress capacities and the response reservation must be non-zero")] + ZeroCapacity, + #[error("read capacity must fit inside normal capacity")] + ReadExceedsNormal, + #[error("one response reservation must fit in the outbound byte capacity")] + ResponseReservationTooLarge, + #[error("combined normal and reserved capacity overflowed usize")] + CapacityOverflow, +} + +impl IngressLimits { + fn validate(self) -> Result { + if self.normal_items == 0 + || self.normal_bytes == 0 + || self.reserved_items == 0 + || self.reserved_bytes == 0 + || self.read_items == 0 + || self.read_bytes == 0 + || self.control_items == 0 + || self.control_bytes == 0 + || self.outbound_response_items == 0 + || self.outbound_response_bytes == 0 + || self.response_reservation_bytes == 0 + { + return Err(LimitsError::ZeroCapacity); + } + if self.read_items > self.normal_items || self.read_bytes > self.normal_bytes { + return Err(LimitsError::ReadExceedsNormal); + } + if self.response_reservation_bytes > self.outbound_response_bytes { + return Err(LimitsError::ResponseReservationTooLarge); + } + self.normal_items + .checked_add(self.reserved_items) + .ok_or(LimitsError::CapacityOverflow)?; + self.normal_bytes + .checked_add(self.reserved_bytes) + .ok_or(LimitsError::CapacityOverflow)?; + Ok(self) + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct QueueUsage { + pub items: usize, + pub bytes: usize, + pub read_items: usize, + pub read_bytes: usize, + pub control_items: usize, + pub control_bytes: usize, + pub outbound_response_items: usize, + pub outbound_response_bytes: usize, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ProtocolErrorKind { + ParseError, + InvalidRequest, + ServerNotInitialized, + MethodNotFound, + InvalidParams, + InternalError, + RequestCanceled, + ContentModified, + RequestFailed, +} + +impl ProtocolErrorKind { + #[must_use] + pub const fn json_rpc_code(self) -> i32 { + match self { + Self::ParseError => -32700, + Self::InvalidRequest => -32600, + Self::ServerNotInitialized => -32002, + Self::MethodNotFound => -32601, + Self::InvalidParams => -32602, + Self::InternalError => -32603, + Self::RequestCanceled => -32800, + Self::ContentModified => -32801, + Self::RequestFailed => -32803, + } + } +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ProtocolError { + pub kind: ProtocolErrorKind, + pub message: String, +} + +impl ProtocolError { + fn new(kind: ProtocolErrorKind, message: impl Into) -> Self { + Self { + kind, + message: message.into(), + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum BackpressureReason { + NormalCapacity, + ReservedCapacity, + ControlCapacity, + OutboundResponseCapacity, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum DropReason { + NotificationBeforeInitialized, + InitializedInWrongState, + NotificationAfterShutdown, + /// The serialized message alone exceeds its class's byte capacity: it can + /// never be admitted, so backpressure would livelock the transport. The + /// notification (or client response) is dropped; the session stays alive. + OversizedMessage, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum AdmissionResult { + Admitted { + sequence: IngressSequence, + response_token: Option, + coalesced: bool, + }, + Rejected { + response_token: ResponseToken, + error: ProtocolError, + }, + Dropped(DropReason), + Backpressure(BackpressureReason), + /// `$/cancelRequest` must enter [`IngressScheduler::admit_cancel`]. + UseControlPath, + Exited(SessionTermination), + UnknownSession, + DuplicateRequestId(ResponseToken), +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ControlAdmission { + Admitted, + Duplicate, + Backpressure, + /// The control alone exceeds the control byte budget and can never be + /// admitted; retrying would livelock. Drop it and keep the session alive. + Oversized, + UnknownSession, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ResponseOutcome { + Success, + Error, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ResponseClaim { + Claimed, + AlreadyClaimed, + UnknownOrSuperseded, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ResponseOrder { + Ordered, + UnknownOrSuperseded, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ResponseSizeError { + UnknownOrSuperseded, + ExceedsReservation, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum OutstandingState { + Queued, + Running, + Responded, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct OwnedDocument { + pub uri: String, + pub version: Option, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum TerminationReason { + NormalExit, + EarlyExit, + TransportClosed, + BrowserSuperseded, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct SessionTermination { + pub session_id: SessionId, + pub prior_lifecycle: LifecycleState, + pub final_lifecycle: LifecycleState, + pub reason: TerminationReason, + pub dropped_normal_items: usize, + pub dropped_control_items: usize, + pub revoked_response_tokens: Vec, + pub owned_documents: Vec, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct OpenSession { + pub session_id: SessionId, + /// Present only when a new browser session atomically revoked the previous + /// browser session. The adapter must close these document overlays before + /// exposing `session_id` to the replacement transport. + pub takeover: Option, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum CancelEffect { + Respond { + response_token: ResponseToken, + was_running: bool, + signaled_operation: bool, + error: ProtocolError, + }, + Noop, +} + +#[derive(Debug)] +pub enum SchedulerEvent { + Control(CancelEffect), + Dispatch(DispatchItem), +} + +#[derive(Debug)] +pub struct DispatchItem { + pub sequence: IngressSequence, + pub session_id: SessionId, + pub class: AdmissionClass, + pub message: Message, + pub response_token: Option, + pub cancellation: Option, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum SessionConfigError { + UnknownSession, + WrongLifecycle(LifecycleState), + AlreadySet, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum RequestRole { + Initialize, + Shutdown, + Ordinary, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum ResponseOwner { + Normal(ResponseOutcome), + Canceled, + ImmediateError, +} + +struct OutstandingRequest { + state: OutstandingState, + cancellation: CancellationToken, + behavior: CancellationBehavior, + cancel_committed: bool, + role: RequestRole, + response_owner: Option, +} + +struct Session { + transport: TransportKind, + lifecycle: LifecycleState, + capabilities: Option>, + documents: BTreeMap>, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +struct FullChangeKey { + uri: String, +} + +struct QueuedItem { + sequence: IngressSequence, + session_id: SessionId, + class: AdmissionClass, + charge_bytes: usize, + read_charge: bool, + message: Message, + response_token: Option, + full_change: Option, + /// An `initialized` notification admitted while `initialize` was still in + /// flight (a pipelined client). Its acceptance is decided at dispatch: + /// it is valid only once the initialize response has been ordered + /// (`InitializeResponded`), and is dropped in any other state (B2). + initialized_gate: bool, +} + +struct QueuedCancel { + token: ResponseToken, + charge_bytes: usize, +} + +/// Process-owned scheduler shared by every LSP transport adapter. +pub struct IngressScheduler { + limits: IngressLimits, + next_session_id: u64, + next_sequence: u64, + sessions: HashMap, + active_browser: Option, + normal: VecDeque, + control: VecDeque, + pending_cancels: HashSet, + outstanding: HashMap, + normal_items: usize, + normal_bytes: usize, + read_items: usize, + read_bytes: usize, + control_bytes: usize, + outbound_response_items: usize, + outbound_response_bytes: usize, +} + +impl IngressScheduler { + pub fn new(limits: IngressLimits) -> Result { + let limits = limits.validate()?; + Ok(Self { + limits, + next_session_id: 1, + next_sequence: 1, + sessions: HashMap::new(), + active_browser: None, + normal: VecDeque::new(), + control: VecDeque::new(), + pending_cancels: HashSet::new(), + outstanding: HashMap::new(), + normal_items: 0, + normal_bytes: 0, + read_items: 0, + read_bytes: 0, + control_bytes: 0, + outbound_response_items: 0, + outbound_response_bytes: 0, + }) + } + + /// Opens a fresh connection. Opening a browser connection atomically + /// revokes the previous browser session and returns the cleanup work that + /// must be applied before the replacement socket is exposed. + pub fn open_session(&mut self, transport: TransportKind) -> OpenSession { + let takeover = if transport == TransportKind::Browser { + self.active_browser + .and_then(|old| self.terminate(old, TerminationReason::BrowserSuperseded)) + } else { + None + }; + let session_id = SessionId(self.next_session_id); + self.next_session_id = self.next_session_id.saturating_add(1); + self.sessions.insert( + session_id, + Session { + transport, + lifecycle: LifecycleState::PreInitialize, + capabilities: None, + documents: BTreeMap::new(), + }, + ); + if transport == TransportKind::Browser { + self.active_browser = Some(session_id); + } + OpenSession { + session_id, + takeover, + } + } + + #[must_use] + pub fn lifecycle(&self, session_id: SessionId) -> Option { + self.sessions.get(&session_id).map(|s| s.lifecycle) + } + + #[must_use] + pub fn transport(&self, session_id: SessionId) -> Option { + self.sessions.get(&session_id).map(|s| s.transport) + } + + #[must_use] + pub fn active_browser(&self) -> Option { + self.active_browser + } + + pub fn set_negotiated_capabilities( + &mut self, + session_id: SessionId, + capabilities: serde_json::Value, + ) -> Result<(), SessionConfigError> { + let Some(session) = self.sessions.get_mut(&session_id) else { + return Err(SessionConfigError::UnknownSession); + }; + if !matches!( + session.lifecycle, + LifecycleState::InitializeResponding | LifecycleState::InitializeResponded + ) { + return Err(SessionConfigError::WrongLifecycle(session.lifecycle)); + } + if session.capabilities.is_some() { + return Err(SessionConfigError::AlreadySet); + } + session.capabilities = Some(Arc::new(capabilities)); + Ok(()) + } + + #[must_use] + pub fn negotiated_capabilities(&self, session_id: SessionId) -> Option> { + self.sessions + .get(&session_id) + .and_then(|session| session.capabilities.clone()) + } + + /// Records an applied didOpen. Call this from the mutation handler's commit + /// point, not merely when the notification is admitted. + pub fn record_document_open( + &mut self, + session_id: SessionId, + uri: impl Into, + version: Option, + ) -> bool { + let Some(session) = self.sessions.get_mut(&session_id) else { + return false; + }; + session.documents.insert(uri.into(), version); + true + } + + pub fn record_document_version( + &mut self, + session_id: SessionId, + uri: &str, + version: Option, + ) -> bool { + let Some(session) = self.sessions.get_mut(&session_id) else { + return false; + }; + let Some(current) = session.documents.get_mut(uri) else { + return false; + }; + *current = version; + true + } + + pub fn record_document_close(&mut self, session_id: SessionId, uri: &str) -> bool { + self.sessions + .get_mut(&session_id) + .is_some_and(|session| session.documents.remove(uri).is_some()) + } + + #[must_use] + pub fn usage(&self) -> QueueUsage { + QueueUsage { + items: self.normal_items, + bytes: self.normal_bytes, + read_items: self.read_items, + read_bytes: self.read_bytes, + control_items: self.control.len(), + control_bytes: self.control_bytes, + outbound_response_items: self.outbound_response_items, + outbound_response_bytes: self.outbound_response_bytes, + } + } + + pub fn admit_message( + &mut self, + session_id: SessionId, + message: Message, + serialized_bytes: usize, + ) -> AdmissionResult { + let policy = MessagePolicy::for_message(&message); + self.admit_message_with_policy(session_id, message, serialized_bytes, policy) + } + + pub fn admit_message_with_policy( + &mut self, + session_id: SessionId, + message: Message, + serialized_bytes: usize, + requested_policy: MessagePolicy, + ) -> AdmissionResult { + let Some(session) = self.sessions.get(&session_id) else { + return AdmissionResult::UnknownSession; + }; + if matches!( + &message, + Message::Notification(notification) if notification.method == "$/cancelRequest" + ) { + return AdmissionResult::UseControlPath; + } + let forced = MessagePolicy::for_message(&message); + let policy = MessagePolicy { + class: stronger_class(forced.class, requested_policy.class), + cancellation: stronger_cancellation(forced.cancellation, requested_policy.cancellation), + }; + + let lifecycle_action = match lifecycle_action(session.lifecycle, &message) { + LifecycleAction::Reject(kind, text) => { + let Message::Request(request) = &message else { + unreachable!("only requests receive lifecycle rejections") + }; + return self.reject_request( + session_id, + request.id.clone(), + ProtocolError::new(kind, text), + ); + } + LifecycleAction::Drop(reason) => return AdmissionResult::Dropped(reason), + LifecycleAction::Exit(normal) => { + let reason = if normal { + TerminationReason::NormalExit + } else { + TerminationReason::EarlyExit + }; + let termination = self + .terminate(session_id, reason) + .expect("session existence checked above"); + return AdmissionResult::Exited(termination); + } + action => action, + }; + + let request_data = match &message { + Message::Request(request) => { + let token = ResponseToken { + session_id, + request_id: request.id.clone(), + }; + if self.outstanding.contains_key(&token) { + return AdmissionResult::DuplicateRequestId(token); + } + if !self.reserve_response() { + return AdmissionResult::Backpressure( + BackpressureReason::OutboundResponseCapacity, + ); + } + let role = match request.method.as_str() { + "initialize" => RequestRole::Initialize, + "shutdown" => RequestRole::Shutdown, + _ => RequestRole::Ordinary, + }; + Some((token, role)) + } + Message::Notification(_) => None, + Message::Response(_) => None, + }; + + let full_change = full_change_key(&message); + if let Some(key) = &full_change + && self.try_coalesce_tail( + session_id, + key, + message.clone(), + serialized_bytes, + policy.class, + ) + { + debug_assert!(request_data.is_none(), "didChange is a notification"); + let sequence = self.normal.back().expect("coalesced tail exists").sequence; + return AdmissionResult::Admitted { + sequence, + response_token: None, + coalesced: true, + }; + } + + let read_charge = policy.class == AdmissionClass::Read; + match self.has_normal_capacity(policy.class, serialized_bytes, read_charge) { + Ok(()) => {} + Err(reason) => { + if request_data.is_some() { + self.release_response(); + } + // A message whose size alone exceeds its class's byte + // capacity can never be admitted: report it once instead of + // backpressuring the transport forever. The session stays + // alive; a request receives a typed RequestFailed response + // and notifications/client responses are dropped. FULL-sync + // didChange recovers on the next (smaller) full text. + if self.permanently_oversized(policy.class, serialized_bytes) { + return match message { + Message::Request(request) => self.reject_request( + session_id, + request.id, + ProtocolError::new( + ProtocolErrorKind::RequestFailed, + "message exceeds the LSP ingress byte capacity for its class", + ), + ), + Message::Notification(_) | Message::Response(_) => { + AdmissionResult::Dropped(DropReason::OversizedMessage) + } + }; + } + if policy.class == AdmissionClass::Read { + let Message::Request(request) = message else { + return AdmissionResult::Backpressure(reason); + }; + return self.reject_request( + session_id, + request.id, + ProtocolError::new( + ProtocolErrorKind::RequestFailed, + "LSP ingress read budget is saturated", + ), + ); + } + return AdmissionResult::Backpressure(reason); + } + } + + let sequence = IngressSequence(self.next_sequence); + self.next_sequence = self.next_sequence.saturating_add(1); + let response_token = request_data.as_ref().map(|(token, _)| token.clone()); + if let Some((token, role)) = request_data { + self.outstanding.insert( + token, + OutstandingRequest { + state: OutstandingState::Queued, + cancellation: CancellationToken::new(), + behavior: policy.cancellation, + cancel_committed: false, + role, + response_owner: None, + }, + ); + } + self.normal_items += 1; + self.normal_bytes += serialized_bytes; + if read_charge { + self.read_items += 1; + self.read_bytes += serialized_bytes; + } + self.normal.push_back(QueuedItem { + sequence, + session_id, + class: policy.class, + charge_bytes: serialized_bytes, + read_charge, + message, + response_token: response_token.clone(), + full_change, + initialized_gate: lifecycle_action == LifecycleAction::QueueInitialized, + }); + self.apply_lifecycle_admission(session_id, lifecycle_action); + AdmissionResult::Admitted { + sequence, + response_token, + coalesced: false, + } + } + + /// Enqueues cancellation on a separate bounded path. [`Self::next_event`] + /// always drains this path before normal work. + pub fn admit_cancel( + &mut self, + session_id: SessionId, + request_id: RequestId, + serialized_bytes: usize, + ) -> ControlAdmission { + if !self.sessions.contains_key(&session_id) { + return ControlAdmission::UnknownSession; + } + let token = ResponseToken { + session_id, + request_id, + }; + if self.pending_cancels.contains(&token) { + return ControlAdmission::Duplicate; + } + if serialized_bytes > self.limits.control_bytes { + return ControlAdmission::Oversized; + } + if self.control.len() + 1 > self.limits.control_items + || self.control_bytes.saturating_add(serialized_bytes) > self.limits.control_bytes + { + return ControlAdmission::Backpressure; + } + self.control.push_back(QueuedCancel { + token: token.clone(), + charge_bytes: serialized_bytes, + }); + self.pending_cancels.insert(token); + self.control_bytes += serialized_bytes; + ControlAdmission::Admitted + } + + /// Controls are returned before normal work. Popping a request transitions + /// it atomically from queued to running and returns its safe-point token. + pub fn next_event(&mut self) -> Option { + if let Some(effect) = self.next_control() { + return Some(SchedulerEvent::Control(effect)); + } + loop { + let item = self.normal.pop_front()?; + self.remove_normal_charge(&item); + if !self.sessions.contains_key(&item.session_id) { + continue; + } + if item.initialized_gate { + // Pipelined `initialized`: accepted only if the initialize + // response has been ordered by now (FIFO dispatch runs it + // after the initialize dispatch). A response still pending on + // a saturated sink — or a failed initialize — means the + // client cannot have seen the response, so the premature + // notification is dropped ("initialized in any other state + // is ignored", B2). + let session = self + .sessions + .get_mut(&item.session_id) + .expect("session existence checked above"); + if session.lifecycle != LifecycleState::InitializeResponded { + continue; + } + session.lifecycle = LifecycleState::Initialized; + } + let cancellation = item.response_token.as_ref().and_then(|token| { + let outstanding = self.outstanding.get_mut(token)?; + if outstanding.state != OutstandingState::Queued { + return None; + } + outstanding.state = OutstandingState::Running; + Some(outstanding.cancellation.clone()) + }); + return Some(SchedulerEvent::Dispatch(DispatchItem { + sequence: item.sequence, + session_id: item.session_id, + class: item.class, + message: item.message, + response_token: item.response_token, + cancellation, + })); + } + } + + /// Drains one cancellation/close control without touching the normal FIFO. + /// + /// Transport runtimes call this directly after admitting cancellation so a + /// control can claim the response of a handler currently blocking the + /// dispatch worker. Calling [`Self::next_event`] alone would defer that + /// control until the synchronous handler returned. + pub fn next_control(&mut self) -> Option { + let control = self.control.pop_front()?; + self.control_bytes -= control.charge_bytes; + self.pending_cancels.remove(&control.token); + Some(self.apply_cancel(&control.token)) + } + + /// After a side-effecting handler crosses its documented commit point, + /// cancellation becomes a no-op and the normal result retains ownership. + pub fn mark_non_cancellable(&mut self, token: &ResponseToken) -> bool { + let Some(outstanding) = self.outstanding.get_mut(token) else { + return false; + }; + if outstanding.state != OutstandingState::Running { + return false; + } + outstanding.cancel_committed = true; + true + } + + pub fn claim_response( + &mut self, + token: &ResponseToken, + outcome: ResponseOutcome, + ) -> ResponseClaim { + if !self.sessions.contains_key(&token.session_id) { + return ResponseClaim::UnknownOrSuperseded; + } + let Some(outstanding) = self.outstanding.get_mut(token) else { + return ResponseClaim::UnknownOrSuperseded; + }; + if outstanding.response_owner.is_some() { + return ResponseClaim::AlreadyClaimed; + } + outstanding.response_owner = Some(ResponseOwner::Normal(outcome)); + outstanding.state = OutstandingState::Responded; + ResponseClaim::Claimed + } + + /// Verifies that the serialized response fits the bytes reserved at request + /// admission. Oversize responses must close the transport; they cannot be + /// dropped or placed on an unbounded writer queue. + pub fn validate_response_size( + &self, + token: &ResponseToken, + serialized_bytes: usize, + ) -> Result<(), ResponseSizeError> { + if !self.outstanding.contains_key(token) || !self.sessions.contains_key(&token.session_id) { + return Err(ResponseSizeError::UnknownOrSuperseded); + } + if serialized_bytes > self.limits.response_reservation_bytes { + return Err(ResponseSizeError::ExceedsReservation); + } + Ok(()) + } + + /// Releases the response reservation only after the response has been + /// ordered onto the owning sink. Initialize/shutdown lifecycle transitions + /// happen at this exact boundary. + pub fn response_ordered(&mut self, token: &ResponseToken) -> ResponseOrder { + if !self.sessions.contains_key(&token.session_id) { + return ResponseOrder::UnknownOrSuperseded; + } + let Some(outstanding) = self.outstanding.remove(token) else { + return ResponseOrder::UnknownOrSuperseded; + }; + let Some(owner) = outstanding.response_owner else { + self.outstanding.insert(token.clone(), outstanding); + return ResponseOrder::UnknownOrSuperseded; + }; + self.release_response(); + if let Some(session) = self.sessions.get_mut(&token.session_id) { + match (outstanding.role, owner) { + (RequestRole::Initialize, ResponseOwner::Normal(ResponseOutcome::Success)) => { + session.lifecycle = LifecycleState::InitializeResponded; + } + (RequestRole::Initialize, _) => { + // Initialize is accepted at most once even if its handler + // fails or a queued cancellation owns the response. + session.lifecycle = LifecycleState::Exited; + session.capabilities = None; + } + (RequestRole::Shutdown, ResponseOwner::Normal(ResponseOutcome::Success)) => { + session.lifecycle = LifecycleState::Shutdown; + } + (RequestRole::Shutdown, _) => { + // Shutdown handlers may already have committed process + // state; never reopen an initialized session on failure. + session.lifecycle = LifecycleState::Shutdown; + } + (RequestRole::Ordinary, _) => {} + } + } + ResponseOrder::Ordered + } + + #[must_use] + pub fn outstanding_state(&self, token: &ResponseToken) -> Option { + self.outstanding.get(token).map(|request| request.state) + } + + pub fn close_session(&mut self, session_id: SessionId) -> Option { + self.terminate(session_id, TerminationReason::TransportClosed) + } + + fn reject_request( + &mut self, + session_id: SessionId, + request_id: RequestId, + error: ProtocolError, + ) -> AdmissionResult { + let token = ResponseToken { + session_id, + request_id, + }; + if self.outstanding.contains_key(&token) { + return AdmissionResult::DuplicateRequestId(token); + } + if !self.reserve_response() { + return AdmissionResult::Backpressure(BackpressureReason::OutboundResponseCapacity); + } + self.outstanding.insert( + token.clone(), + OutstandingRequest { + state: OutstandingState::Responded, + cancellation: CancellationToken::new(), + behavior: CancellationBehavior::NonCancellableWhenRunning, + cancel_committed: true, + role: RequestRole::Ordinary, + response_owner: Some(ResponseOwner::ImmediateError), + }, + ); + AdmissionResult::Rejected { + response_token: token, + error, + } + } + + fn reserve_response(&mut self) -> bool { + if self.outbound_response_items + 1 > self.limits.outbound_response_items + || self + .outbound_response_bytes + .saturating_add(self.limits.response_reservation_bytes) + > self.limits.outbound_response_bytes + { + return false; + } + self.outbound_response_items += 1; + self.outbound_response_bytes += self.limits.response_reservation_bytes; + true + } + + fn release_response(&mut self) { + self.outbound_response_items -= 1; + self.outbound_response_bytes -= self.limits.response_reservation_bytes; + } + + /// True when `bytes` alone exceeds the class's byte capacity — such a + /// message cannot be admitted even from an empty queue, so retrying is a + /// livelock, not backpressure. + fn permanently_oversized(&self, class: AdmissionClass, bytes: usize) -> bool { + let byte_limit = if class.is_reserved() { + self.limits.normal_bytes + self.limits.reserved_bytes + } else { + self.limits.normal_bytes + }; + let class_limit = if class == AdmissionClass::Read { + byte_limit.min(self.limits.read_bytes) + } else { + byte_limit + }; + bytes > class_limit + } + + fn has_normal_capacity( + &self, + class: AdmissionClass, + bytes: usize, + read_charge: bool, + ) -> Result<(), BackpressureReason> { + let (item_limit, byte_limit, reason) = if class.is_reserved() { + ( + self.limits.normal_items + self.limits.reserved_items, + self.limits.normal_bytes + self.limits.reserved_bytes, + BackpressureReason::ReservedCapacity, + ) + } else { + ( + self.limits.normal_items, + self.limits.normal_bytes, + BackpressureReason::NormalCapacity, + ) + }; + if self.normal_items + 1 > item_limit + || self.normal_bytes.saturating_add(bytes) > byte_limit + { + return Err(reason); + } + if read_charge + && (self.read_items + 1 > self.limits.read_items + || self.read_bytes.saturating_add(bytes) > self.limits.read_bytes) + { + return Err(BackpressureReason::NormalCapacity); + } + Ok(()) + } + + fn try_coalesce_tail( + &mut self, + session_id: SessionId, + key: &FullChangeKey, + message: Message, + bytes: usize, + class: AdmissionClass, + ) -> bool { + let Some(tail) = self.normal.back() else { + return false; + }; + if tail.session_id != session_id || tail.full_change.as_ref() != Some(key) { + return false; + } + let new_total_bytes = self + .normal_bytes + .saturating_sub(tail.charge_bytes) + .saturating_add(bytes); + let byte_limit = if class.is_reserved() { + self.limits.normal_bytes + self.limits.reserved_bytes + } else { + self.limits.normal_bytes + }; + if new_total_bytes > byte_limit { + return false; + } + let sequence = IngressSequence(self.next_sequence); + self.next_sequence = self.next_sequence.saturating_add(1); + let tail = self.normal.back_mut().expect("tail checked above"); + self.normal_bytes = new_total_bytes; + tail.sequence = sequence; + tail.charge_bytes = bytes; + tail.message = message; + true + } + + fn remove_normal_charge(&mut self, item: &QueuedItem) { + self.normal_items -= 1; + self.normal_bytes -= item.charge_bytes; + if item.read_charge { + self.read_items -= 1; + self.read_bytes -= item.charge_bytes; + } + } + + fn apply_lifecycle_admission(&mut self, session_id: SessionId, action: LifecycleAction) { + let Some(session) = self.sessions.get_mut(&session_id) else { + return; + }; + match action { + LifecycleAction::BeginInitialize => { + session.lifecycle = LifecycleState::InitializeResponding; + } + LifecycleAction::Initialized => { + session.lifecycle = LifecycleState::Initialized; + } + LifecycleAction::BeginShutdown => { + session.lifecycle = LifecycleState::ShutdownResponding; + } + LifecycleAction::NoChange + | LifecycleAction::QueueInitialized + | LifecycleAction::Reject(_, _) + | LifecycleAction::Drop(_) + | LifecycleAction::Exit(_) => {} + } + } + + fn apply_cancel(&mut self, token: &ResponseToken) -> CancelEffect { + let Some(state) = self.outstanding.get(token).map(|request| request.state) else { + return CancelEffect::Noop; + }; + if state == OutstandingState::Responded { + return CancelEffect::Noop; + } + if state == OutstandingState::Queued { + if let Some(index) = self + .normal + .iter() + .position(|item| item.response_token.as_ref() == Some(token)) + { + let removed = self.normal.remove(index).expect("index came from queue"); + self.remove_normal_charge(&removed); + } + let outstanding = self + .outstanding + .get_mut(token) + .expect("state read from outstanding map"); + outstanding.state = OutstandingState::Responded; + outstanding.response_owner = Some(ResponseOwner::Canceled); + return CancelEffect::Respond { + response_token: token.clone(), + was_running: false, + signaled_operation: false, + error: ProtocolError::new( + ProtocolErrorKind::RequestCanceled, + "request canceled before dispatch", + ), + }; + } + + let outstanding = self + .outstanding + .get_mut(token) + .expect("state read from outstanding map"); + if outstanding.cancel_committed + || outstanding.behavior == CancellationBehavior::NonCancellableWhenRunning + { + return CancelEffect::Noop; + } + let signaled = outstanding.behavior == CancellationBehavior::SignalAtSafePoints; + if signaled { + outstanding.cancellation.cancel(); + } + outstanding.state = OutstandingState::Responded; + outstanding.response_owner = Some(ResponseOwner::Canceled); + CancelEffect::Respond { + response_token: token.clone(), + was_running: true, + signaled_operation: signaled, + error: ProtocolError::new( + ProtocolErrorKind::RequestCanceled, + "request canceled while running", + ), + } + } + + fn terminate( + &mut self, + session_id: SessionId, + reason: TerminationReason, + ) -> Option { + let session = self.sessions.remove(&session_id)?; + if self.active_browser == Some(session_id) { + self.active_browser = None; + } + let final_lifecycle = if reason == TerminationReason::BrowserSuperseded { + LifecycleState::Superseded + } else { + LifecycleState::Exited + }; + + let mut dropped_normal_items = 0; + let mut retained = VecDeque::with_capacity(self.normal.len()); + while let Some(item) = self.normal.pop_front() { + if item.session_id == session_id { + self.remove_normal_charge(&item); + dropped_normal_items += 1; + } else { + retained.push_back(item); + } + } + self.normal = retained; + + let mut dropped_control_items = 0; + let mut retained_control = VecDeque::with_capacity(self.control.len()); + while let Some(control) = self.control.pop_front() { + if control.token.session_id == session_id { + self.control_bytes -= control.charge_bytes; + self.pending_cancels.remove(&control.token); + dropped_control_items += 1; + } else { + retained_control.push_back(control); + } + } + self.control = retained_control; + + let mut tokens: Vec<_> = self + .outstanding + .keys() + .filter(|token| token.session_id == session_id) + .cloned() + .collect(); + tokens.sort(); + for token in &tokens { + if let Some(outstanding) = self.outstanding.remove(token) { + outstanding.cancellation.cancel(); + self.release_response(); + } + } + + Some(SessionTermination { + session_id, + prior_lifecycle: session.lifecycle, + final_lifecycle, + reason, + dropped_normal_items, + dropped_control_items, + revoked_response_tokens: tokens, + owned_documents: session + .documents + .into_iter() + .map(|(uri, version)| OwnedDocument { uri, version }) + .collect(), + }) + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum LifecycleAction { + NoChange, + BeginInitialize, + Initialized, + /// Pipelined `initialized` while `initialize` is still responding: admit + /// FIFO, decide acceptance at dispatch (see [`QueuedItem::initialized_gate`]). + QueueInitialized, + BeginShutdown, + Reject(ProtocolErrorKind, &'static str), + Drop(DropReason), + Exit(bool), +} + +fn lifecycle_action(state: LifecycleState, message: &Message) -> LifecycleAction { + match message { + Message::Request(request) if request.method == "initialize" => { + if state == LifecycleState::PreInitialize { + LifecycleAction::BeginInitialize + } else { + LifecycleAction::Reject( + ProtocolErrorKind::InvalidRequest, + "initialize may be requested exactly once", + ) + } + } + Message::Request(request) if request.method == "shutdown" => { + if state == LifecycleState::Initialized { + LifecycleAction::BeginShutdown + } else if matches!( + state, + LifecycleState::PreInitialize + | LifecycleState::InitializeResponding + | LifecycleState::InitializeResponded + ) { + LifecycleAction::Reject( + ProtocolErrorKind::ServerNotInitialized, + "shutdown requested before initialized", + ) + } else { + LifecycleAction::Reject( + ProtocolErrorKind::InvalidRequest, + "shutdown is not valid in the current lifecycle state", + ) + } + } + Message::Request(_) => match state { + LifecycleState::Initialized => LifecycleAction::NoChange, + LifecycleState::PreInitialize + | LifecycleState::InitializeResponding + | LifecycleState::InitializeResponded => LifecycleAction::Reject( + ProtocolErrorKind::ServerNotInitialized, + "request received before initialized", + ), + LifecycleState::ShutdownResponding + | LifecycleState::Shutdown + | LifecycleState::Exited + | LifecycleState::Superseded => LifecycleAction::Reject( + ProtocolErrorKind::InvalidRequest, + "request received after shutdown began", + ), + }, + Message::Notification(notification) if notification.method == "initialized" => { + match state { + LifecycleState::InitializeResponded => LifecycleAction::Initialized, + // A pipelined client may send `initialized` before it could + // have read the initialize response. Queue it FIFO behind the + // in-flight initialize; dispatch decides acceptance. + LifecycleState::InitializeResponding => LifecycleAction::QueueInitialized, + _ => LifecycleAction::Drop(DropReason::InitializedInWrongState), + } + } + Message::Notification(notification) if notification.method == "exit" => { + LifecycleAction::Exit(state == LifecycleState::Shutdown) + } + Message::Notification(_) => match state { + LifecycleState::Initialized => LifecycleAction::NoChange, + LifecycleState::PreInitialize + | LifecycleState::InitializeResponding + | LifecycleState::InitializeResponded => { + LifecycleAction::Drop(DropReason::NotificationBeforeInitialized) + } + LifecycleState::ShutdownResponding + | LifecycleState::Shutdown + | LifecycleState::Exited + | LifecycleState::Superseded => { + LifecycleAction::Drop(DropReason::NotificationAfterShutdown) + } + }, + Message::Response(_) => match state { + LifecycleState::Initialized => LifecycleAction::NoChange, + LifecycleState::PreInitialize + | LifecycleState::InitializeResponding + | LifecycleState::InitializeResponded => { + LifecycleAction::Drop(DropReason::NotificationBeforeInitialized) + } + LifecycleState::ShutdownResponding + | LifecycleState::Shutdown + | LifecycleState::Exited + | LifecycleState::Superseded => { + LifecycleAction::Drop(DropReason::NotificationAfterShutdown) + } + }, + } +} + +fn stronger_class(forced: AdmissionClass, requested: AdmissionClass) -> AdmissionClass { + if forced.is_reserved() { + forced + } else if requested.is_reserved() { + requested + } else if forced == AdmissionClass::Read || requested == AdmissionClass::Read { + AdmissionClass::Read + } else { + AdmissionClass::Other + } +} + +/// Callers may harden an ordinary message's cancellation behavior but can +/// never downgrade the forced protection of lifecycle and side-effecting +/// methods: reopening cancellation for a running `initialize`/`shutdown`/ +/// `workspace/executeCommand` would reintroduce the canceled-response vs +/// committed-effect race. +fn stronger_cancellation( + forced: CancellationBehavior, + requested: CancellationBehavior, +) -> CancellationBehavior { + const fn strength(behavior: CancellationBehavior) -> u8 { + match behavior { + CancellationBehavior::SignalAtSafePoints => 0, + CancellationBehavior::ClaimResponseOnly => 1, + CancellationBehavior::NonCancellableWhenRunning => 2, + } + } + if strength(requested) >= strength(forced) { + requested + } else { + forced + } +} + +fn full_change_key(message: &Message) -> Option { + let Message::Notification(notification) = message else { + return None; + }; + if notification.method != "textDocument/didChange" { + return None; + } + let uri = notification.params.pointer("/textDocument/uri")?.as_str()?; + notification + .params + .pointer("/textDocument/version")? + .as_i64()?; + let changes = notification.params.pointer("/contentChanges")?.as_array()?; + let [change] = changes.as_slice() else { + return None; + }; + if change + .get("text") + .and_then(serde_json::Value::as_str) + .is_none() + || change.get("range").is_some_and(|range| !range.is_null()) + || change + .get("rangeLength") + .is_some_and(|length| !length.is_null()) + { + return None; + } + Some(FullChangeKey { + uri: uri.to_string(), + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn limits() -> IngressLimits { + IngressLimits { + normal_items: 2, + normal_bytes: 200, + reserved_items: 2, + reserved_bytes: 200, + read_items: 1, + read_bytes: 100, + control_items: 2, + control_bytes: 100, + outbound_response_items: 4, + outbound_response_bytes: 400, + response_reservation_bytes: 100, + } + } + + fn request(id: i32, method: &str) -> Message { + Message::Request(lsp_server::Request { + id: RequestId::from(id), + method: method.to_string(), + params: serde_json::Value::Null, + }) + } + + fn notification(method: &str, params: serde_json::Value) -> Message { + Message::Notification(lsp_server::Notification { + method: method.to_string(), + params, + }) + } + + fn full_change(uri: &str, version: i32, text: &str) -> Message { + notification( + "textDocument/didChange", + serde_json::json!({ + "textDocument": { "uri": uri, "version": version }, + "contentChanges": [{ "text": text }], + }), + ) + } + + fn initialize(scheduler: &mut IngressScheduler, session: SessionId, id: i32) { + let AdmissionResult::Admitted { + response_token: Some(token), + .. + } = scheduler.admit_message(session, request(id, "initialize"), 20) + else { + panic!("initialize must be admitted"); + }; + assert_eq!( + scheduler.lifecycle(session), + Some(LifecycleState::InitializeResponding) + ); + let SchedulerEvent::Dispatch(dispatch) = scheduler.next_event().unwrap() else { + panic!("expected initialize dispatch"); + }; + assert_eq!(dispatch.response_token.as_ref(), Some(&token)); + assert_eq!( + scheduler.claim_response(&token, ResponseOutcome::Success), + ResponseClaim::Claimed + ); + assert_eq!( + scheduler.lifecycle(session), + Some(LifecycleState::InitializeResponding), + "claiming is not the outbound ordering boundary" + ); + assert_eq!(scheduler.response_ordered(&token), ResponseOrder::Ordered); + assert_eq!( + scheduler.lifecycle(session), + Some(LifecycleState::InitializeResponded) + ); + assert!(matches!( + scheduler.admit_message( + session, + notification("initialized", serde_json::Value::Null), + 10, + ), + AdmissionResult::Admitted { .. } + )); + assert_eq!( + scheduler.lifecycle(session), + Some(LifecycleState::Initialized) + ); + let _ = scheduler.next_event(); + } + + #[test] + fn lifecycle_advances_only_after_ordered_responses() { + let mut scheduler = IngressScheduler::new(limits()).unwrap(); + let session = scheduler.open_session(TransportKind::Stdio).session_id; + + assert!(matches!( + scheduler.admit_message(session, request(1, "textDocument/hover"), 20), + AdmissionResult::Rejected { + error: ProtocolError { + kind: ProtocolErrorKind::ServerNotInitialized, + .. + }, + .. + } + )); + initialize(&mut scheduler, session, 2); + + let AdmissionResult::Admitted { + response_token: Some(shutdown), + .. + } = scheduler.admit_message(session, request(3, "shutdown"), 20) + else { + panic!("shutdown must be admitted"); + }; + assert_eq!( + scheduler.lifecycle(session), + Some(LifecycleState::ShutdownResponding) + ); + assert!(matches!( + scheduler.admit_message(session, request(4, "textDocument/hover"), 20), + AdmissionResult::Rejected { + error: ProtocolError { + kind: ProtocolErrorKind::InvalidRequest, + .. + }, + .. + } + )); + let _ = scheduler.next_event(); + assert_eq!( + scheduler.claim_response(&shutdown, ResponseOutcome::Success), + ResponseClaim::Claimed + ); + assert_eq!( + scheduler.response_ordered(&shutdown), + ResponseOrder::Ordered + ); + assert_eq!(scheduler.lifecycle(session), Some(LifecycleState::Shutdown)); + assert!(matches!( + scheduler.admit_message(session, notification("exit", serde_json::Value::Null), 5), + AdmissionResult::Exited(SessionTermination { + reason: TerminationReason::NormalExit, + .. + }) + )); + } + + #[test] + fn duplicate_initialize_and_wrong_state_initialized_are_deterministic() { + let mut scheduler = IngressScheduler::new(limits()).unwrap(); + let session = scheduler.open_session(TransportKind::Browser).session_id; + assert_eq!( + scheduler.admit_message( + session, + notification("initialized", serde_json::Value::Null), + 5 + ), + AdmissionResult::Dropped(DropReason::InitializedInWrongState) + ); + assert!(matches!( + scheduler.admit_message(session, request(1, "initialize"), 20), + AdmissionResult::Admitted { .. } + )); + assert!(matches!( + scheduler.admit_message(session, request(2, "initialize"), 20), + AdmissionResult::Rejected { + error: ProtocolError { + kind: ProtocolErrorKind::InvalidRequest, + .. + }, + .. + } + )); + } + + #[test] + fn read_overload_rejects_new_request_while_reserve_admits_mutation() { + let mut scheduler = IngressScheduler::new(limits()).unwrap(); + let session = scheduler.open_session(TransportKind::Stdio).session_id; + initialize(&mut scheduler, session, 1); + + assert!(matches!( + scheduler.admit_message(session, request(2, "textDocument/hover"), 80), + AdmissionResult::Admitted { .. } + )); + let rejected = scheduler.admit_message(session, request(3, "textDocument/inlayHint"), 20); + assert!(matches!( + rejected, + AdmissionResult::Rejected { + error: ProtocolError { + kind: ProtocolErrorKind::RequestFailed, + .. + }, + .. + } + )); + assert!(matches!( + scheduler.admit_message(session, full_change("file:///a.baml", 2, "new"), 150), + AdmissionResult::Admitted { .. } + )); + assert_eq!(scheduler.usage().items, 2); + } + + #[test] + fn protected_capacity_backpressures_without_evicting_admitted_work() { + let mut tight = limits(); + tight.normal_items = 1; + tight.reserved_items = 1; + let mut scheduler = IngressScheduler::new(tight).unwrap(); + let session = scheduler.open_session(TransportKind::Stdio).session_id; + initialize(&mut scheduler, session, 1); + + assert!(matches!( + scheduler.admit_message(session, full_change("file:///a.baml", 2, "a"), 20), + AdmissionResult::Admitted { .. } + )); + assert!(matches!( + scheduler.admit_message(session, full_change("file:///b.baml", 2, "b"), 20), + AdmissionResult::Admitted { .. } + )); + assert_eq!( + scheduler.admit_message(session, full_change("file:///c.baml", 2, "c"), 20), + AdmissionResult::Backpressure(BackpressureReason::ReservedCapacity) + ); + assert_eq!(scheduler.usage().items, 2); + } + + #[test] + fn only_adjacent_same_uri_full_changes_coalesce() { + let mut scheduler = IngressScheduler::new(limits()).unwrap(); + let session = scheduler.open_session(TransportKind::Stdio).session_id; + initialize(&mut scheduler, session, 1); + + let first = scheduler.admit_message(session, full_change("file:///a.baml", 2, "old"), 80); + let second = scheduler.admit_message(session, full_change("file:///a.baml", 3, "new"), 40); + assert!(matches!( + first, + AdmissionResult::Admitted { + coalesced: false, + .. + } + )); + assert!(matches!( + second, + AdmissionResult::Admitted { + coalesced: true, + .. + } + )); + assert_eq!(scheduler.usage().items, 1); + assert_eq!(scheduler.usage().bytes, 40); + + let SchedulerEvent::Dispatch(item) = scheduler.next_event().unwrap() else { + panic!("expected coalesced change"); + }; + let Message::Notification(notification) = item.message else { + panic!("expected notification"); + }; + assert_eq!(notification.params["textDocument"]["version"], 3); + assert_eq!(notification.params["contentChanges"][0]["text"], "new"); + + assert!(matches!( + scheduler.admit_message(session, full_change("file:///a.baml", 4, "four"), 20), + AdmissionResult::Admitted { + coalesced: false, + .. + } + )); + assert!(matches!( + scheduler.admit_message(session, full_change("file:///b.baml", 4, "other"), 20), + AdmissionResult::Admitted { + coalesced: false, + .. + } + )); + assert!(matches!( + scheduler.admit_message(session, full_change("file:///a.baml", 5, "five"), 20), + AdmissionResult::Admitted { + coalesced: false, + .. + } + )); + assert_eq!(scheduler.usage().items, 3); + } + + #[test] + fn request_and_incremental_change_are_coalescing_barriers_and_fifo_is_preserved() { + let mut roomy = limits(); + roomy.normal_items = 8; + roomy.normal_bytes = 800; + roomy.reserved_items = 8; + roomy.reserved_bytes = 800; + roomy.read_items = 4; + roomy.read_bytes = 400; + roomy.outbound_response_items = 8; + roomy.outbound_response_bytes = 800; + let mut scheduler = IngressScheduler::new(roomy).unwrap(); + let session = scheduler.open_session(TransportKind::Stdio).session_id; + initialize(&mut scheduler, session, 1); + + let first = scheduler.admit_message(session, full_change("file:///a.baml", 2, "first"), 20); + let request = scheduler.admit_message(session, request(2, "textDocument/hover"), 20); + let after_request = + scheduler.admit_message(session, full_change("file:///a.baml", 3, "second"), 20); + let incremental = notification( + "textDocument/didChange", + serde_json::json!({ + "textDocument": { "uri": "file:///a.baml", "version": 4 }, + "contentChanges": [{ + "range": { + "start": { "line": 0, "character": 0 }, + "end": { "line": 0, "character": 1 } + }, + "text": "x" + }] + }), + ); + let incremental = scheduler.admit_message(session, incremental, 20); + assert!(matches!( + first, + AdmissionResult::Admitted { + coalesced: false, + .. + } + )); + assert!(matches!(request, AdmissionResult::Admitted { .. })); + assert!(matches!( + after_request, + AdmissionResult::Admitted { + coalesced: false, + .. + } + )); + assert!(matches!( + incremental, + AdmissionResult::Admitted { + coalesced: false, + .. + } + )); + + let mut sequences = Vec::new(); + for _ in 0..4 { + let SchedulerEvent::Dispatch(item) = scheduler.next_event().unwrap() else { + panic!("normal messages should dispatch in FIFO order"); + }; + sequences.push(item.sequence.get()); + } + assert!(sequences.windows(2).all(|pair| pair[0] < pair[1])); + } + + #[test] + fn outbound_and_control_paths_enforce_item_and_byte_bounds() { + let mut bounded = limits(); + bounded.outbound_response_items = 1; + bounded.outbound_response_bytes = 100; + bounded.response_reservation_bytes = 100; + bounded.control_items = 1; + bounded.control_bytes = 10; + let mut scheduler = IngressScheduler::new(bounded).unwrap(); + let session = scheduler.open_session(TransportKind::Stdio).session_id; + initialize(&mut scheduler, session, 1); + + let AdmissionResult::Admitted { + response_token: Some(token), + .. + } = scheduler.admit_message(session, request(2, "textDocument/hover"), 20) + else { + panic!("first response slot must be reserved"); + }; + assert_eq!( + scheduler.admit_message(session, request(3, "textDocument/hover"), 20), + AdmissionResult::Backpressure(BackpressureReason::OutboundResponseCapacity) + ); + assert_eq!( + scheduler.validate_response_size(&token, 101), + Err(ResponseSizeError::ExceedsReservation) + ); + assert_eq!( + scheduler.admit_cancel(session, RequestId::from(2), 10), + ControlAdmission::Admitted + ); + assert_eq!( + scheduler.admit_cancel(session, RequestId::from(3), 1), + ControlAdmission::Backpressure + ); + // A control that alone exceeds the control byte budget can never be + // admitted: it is dropped, not retried forever. + assert_eq!( + scheduler.admit_cancel(session, RequestId::from(4), 11), + ControlAdmission::Oversized + ); + } + + #[test] + fn browser_replacement_starts_with_fresh_capabilities() { + let mut scheduler = IngressScheduler::new(limits()).unwrap(); + let old = scheduler.open_session(TransportKind::Browser).session_id; + assert!(matches!( + scheduler.admit_message(old, request(1, "initialize"), 20), + AdmissionResult::Admitted { .. } + )); + scheduler + .set_negotiated_capabilities(old, serde_json::json!({ "positionEncoding": "utf-8" })) + .unwrap(); + assert!(scheduler.negotiated_capabilities(old).is_some()); + + let replacement = scheduler.open_session(TransportKind::Browser); + assert!(replacement.takeover.is_some()); + assert!( + scheduler + .negotiated_capabilities(replacement.session_id) + .is_none() + ); + assert_eq!( + scheduler.lifecycle(replacement.session_id), + Some(LifecycleState::PreInitialize) + ); + } + + #[test] + fn client_responses_share_fifo_ingress_after_initialize() { + let mut scheduler = IngressScheduler::new(limits()).unwrap(); + let session = scheduler.open_session(TransportKind::Stdio).session_id; + initialize(&mut scheduler, session, 1); + let response = Message::Response(lsp_server::Response { + id: RequestId::from(99), + result: Some(serde_json::Value::Null), + error: None, + }); + assert!(matches!( + scheduler.admit_message(session, response, 20), + AdmissionResult::Admitted { + response_token: None, + .. + } + )); + let SchedulerEvent::Dispatch(item) = scheduler.next_event().unwrap() else { + panic!("client response must be dispatched"); + }; + assert!(matches!(item.message, Message::Response(_))); + } + + #[test] + fn queued_cancel_removes_work_and_normal_result_loses_response_race() { + let mut scheduler = IngressScheduler::new(limits()).unwrap(); + let session = scheduler.open_session(TransportKind::Stdio).session_id; + initialize(&mut scheduler, session, 1); + let AdmissionResult::Admitted { + response_token: Some(token), + .. + } = scheduler.admit_message(session, request(2, "textDocument/hover"), 20) + else { + panic!("request must be admitted"); + }; + assert_eq!( + scheduler.admit_cancel(session, RequestId::from(2), 10), + ControlAdmission::Admitted + ); + assert_eq!( + scheduler.admit_cancel(session, RequestId::from(2), 10), + ControlAdmission::Duplicate + ); + let SchedulerEvent::Control(CancelEffect::Respond { + response_token, + was_running: false, + .. + }) = scheduler.next_event().unwrap() + else { + panic!("expected queued cancel response"); + }; + assert_eq!(response_token, token); + assert_eq!(scheduler.usage().items, 0); + assert_eq!( + scheduler.claim_response(&token, ResponseOutcome::Success), + ResponseClaim::AlreadyClaimed + ); + assert_eq!(scheduler.response_ordered(&token), ResponseOrder::Ordered); + } + + #[test] + fn running_cancel_signals_once_and_late_result_is_discarded() { + let mut scheduler = IngressScheduler::new(limits()).unwrap(); + let session = scheduler.open_session(TransportKind::Stdio).session_id; + initialize(&mut scheduler, session, 1); + let AdmissionResult::Admitted { + response_token: Some(token), + .. + } = scheduler.admit_message(session, request(2, "textDocument/hover"), 20) + else { + panic!("request must be admitted"); + }; + let SchedulerEvent::Dispatch(dispatch) = scheduler.next_event().unwrap() else { + panic!("expected request dispatch"); + }; + let cancellation = dispatch.cancellation.unwrap(); + assert!(!cancellation.is_cancelled()); + assert_eq!( + scheduler.admit_cancel(session, RequestId::from(2), 10), + ControlAdmission::Admitted + ); + let SchedulerEvent::Control(CancelEffect::Respond { + was_running: true, + signaled_operation: true, + .. + }) = scheduler.next_event().unwrap() + else { + panic!("expected running cancel response"); + }; + assert!(cancellation.is_cancelled()); + assert_eq!( + scheduler.claim_response(&token, ResponseOutcome::Success), + ResponseClaim::AlreadyClaimed + ); + assert_eq!( + scheduler.admit_cancel(session, RequestId::from(2), 10), + ControlAdmission::Admitted + ); + assert!(matches!( + scheduler.next_event(), + Some(SchedulerEvent::Control(CancelEffect::Noop)) + )); + } + + #[test] + fn committed_side_effect_ignores_running_cancel() { + let mut scheduler = IngressScheduler::new(limits()).unwrap(); + let session = scheduler.open_session(TransportKind::Stdio).session_id; + initialize(&mut scheduler, session, 1); + let AdmissionResult::Admitted { + response_token: Some(token), + .. + } = scheduler.admit_message_with_policy( + session, + request(2, "custom/commit"), + 20, + MessagePolicy::side_effecting(CancellationBehavior::ClaimResponseOnly), + ) + else { + panic!("request must be admitted"); + }; + let _ = scheduler.next_event(); + assert!(scheduler.mark_non_cancellable(&token)); + assert_eq!( + scheduler.admit_cancel(session, RequestId::from(2), 10), + ControlAdmission::Admitted + ); + assert!(matches!( + scheduler.next_event(), + Some(SchedulerEvent::Control(CancelEffect::Noop)) + )); + assert_eq!( + scheduler.claim_response(&token, ResponseOutcome::Success), + ResponseClaim::Claimed + ); + } + + /// A custom policy must not reopen the canceled-response vs + /// committed-effect race on protected methods: the forced + /// `NonCancellableWhenRunning` survives an attempted downgrade. + #[test] + fn forced_cancellation_of_protected_methods_cannot_be_downgraded() { + let mut scheduler = IngressScheduler::new(limits()).unwrap(); + let session = scheduler.open_session(TransportKind::Stdio).session_id; + initialize(&mut scheduler, session, 1); + + let AdmissionResult::Admitted { + response_token: Some(token), + .. + } = scheduler.admit_message_with_policy( + session, + request(2, "workspace/executeCommand"), + 20, + MessagePolicy::side_effecting(CancellationBehavior::SignalAtSafePoints), + ) + else { + panic!("request must be admitted"); + }; + let _ = scheduler.next_event(); + + // Running: cancellation must be a no-op despite the requested policy. + assert_eq!( + scheduler.admit_cancel(session, RequestId::from(2), 10), + ControlAdmission::Admitted + ); + assert!(matches!( + scheduler.next_event(), + Some(SchedulerEvent::Control(CancelEffect::Noop)) + )); + assert_eq!( + scheduler.claim_response(&token, ResponseOutcome::Success), + ResponseClaim::Claimed + ); + } + + #[test] + fn execute_command_cancel_is_linearized_at_dispatch() { + let mut scheduler = IngressScheduler::new(limits()).unwrap(); + let session = scheduler.open_session(TransportKind::Stdio).session_id; + initialize(&mut scheduler, session, 1); + + let AdmissionResult::Admitted { + response_token: Some(queued_token), + .. + } = scheduler.admit_message(session, request(2, "workspace/executeCommand"), 20) + else { + panic!("queued command must be admitted"); + }; + assert_eq!( + scheduler.admit_cancel(session, RequestId::from(2), 10), + ControlAdmission::Admitted + ); + let Some(SchedulerEvent::Control(CancelEffect::Respond { + response_token, + was_running: false, + signaled_operation: false, + error, + })) = scheduler.next_event() + else { + panic!("cancellation before command dispatch must own the response"); + }; + assert_eq!(response_token, queued_token); + assert_eq!(error.kind, ProtocolErrorKind::RequestCanceled); + assert_eq!( + scheduler.claim_response(&queued_token, ResponseOutcome::Success), + ResponseClaim::AlreadyClaimed + ); + + let AdmissionResult::Admitted { + response_token: Some(running_token), + .. + } = scheduler.admit_message(session, request(3, "workspace/executeCommand"), 20) + else { + panic!("running command must be admitted"); + }; + let Some(SchedulerEvent::Dispatch(dispatch)) = scheduler.next_event() else { + panic!("command must transition to running at dispatch"); + }; + assert_eq!(dispatch.response_token.as_ref(), Some(&running_token)); + assert!(!dispatch.cancellation.unwrap().is_cancelled()); + assert_eq!( + scheduler.admit_cancel(session, RequestId::from(3), 10), + ControlAdmission::Admitted + ); + assert!( + matches!( + scheduler.next_event(), + Some(SchedulerEvent::Control(CancelEffect::Noop)) + ), + "cancellation after dispatch must not report a canceled side effect" + ); + assert_eq!( + scheduler.claim_response(&running_token, ResponseOutcome::Success), + ResponseClaim::Claimed + ); + } + + #[test] + fn browser_takeover_revokes_old_work_docs_and_response_identity() { + let mut scheduler = IngressScheduler::new(limits()).unwrap(); + let old = scheduler.open_session(TransportKind::Browser).session_id; + initialize(&mut scheduler, old, 1); + assert!(scheduler.record_document_open(old, "file:///old.baml", Some(7))); + let AdmissionResult::Admitted { + response_token: Some(old_token), + .. + } = scheduler.admit_message(old, request(9, "textDocument/hover"), 20) + else { + panic!("old request must be admitted"); + }; + let _ = scheduler.next_event(); + + let replacement = scheduler.open_session(TransportKind::Browser); + let takeover = replacement + .takeover + .expect("old browser must be superseded"); + assert_eq!(takeover.session_id, old); + assert_eq!(takeover.final_lifecycle, LifecycleState::Superseded); + assert_eq!(takeover.owned_documents[0].uri, "file:///old.baml"); + assert_eq!(scheduler.active_browser(), Some(replacement.session_id)); + assert_eq!( + scheduler.claim_response(&old_token, ResponseOutcome::Success), + ResponseClaim::UnknownOrSuperseded + ); + + // Reusing request id 9 in the replacement is safe because the session + // component of ResponseToken differs. + assert!(matches!( + scheduler.admit_message(replacement.session_id, request(9, "initialize"), 20), + AdmissionResult::Admitted { + response_token: Some(_), + .. + } + )); + } + + #[test] + fn control_path_remains_available_when_normal_queue_is_saturated() { + let mut scheduler = IngressScheduler::new(limits()).unwrap(); + let session = scheduler.open_session(TransportKind::Stdio).session_id; + initialize(&mut scheduler, session, 1); + let AdmissionResult::Admitted { + response_token: Some(_), + .. + } = scheduler.admit_message(session, request(2, "textDocument/hover"), 80) + else { + panic!("read must be admitted"); + }; + assert_eq!( + scheduler.admit_cancel(session, RequestId::from(2), 10), + ControlAdmission::Admitted + ); + assert!(matches!( + scheduler.next_event(), + Some(SchedulerEvent::Control(CancelEffect::Respond { .. })) + )); + } + + /// A pipelined client (initialize + initialized in one burst, without + /// waiting for the response) must still reach `Initialized`: the early + /// `initialized` queues FIFO behind the in-flight initialize and is + /// accepted at dispatch, strictly after the response was ordered. + #[test] + fn pipelined_initialized_is_accepted_after_the_response_is_ordered() { + let mut scheduler = IngressScheduler::new(limits()).unwrap(); + let session = scheduler.open_session(TransportKind::Stdio).session_id; + + let AdmissionResult::Admitted { + response_token: Some(token), + .. + } = scheduler.admit_message(session, request(1, "initialize"), 20) + else { + panic!("initialize must be admitted"); + }; + assert!(matches!( + scheduler.admit_message( + session, + notification("initialized", serde_json::Value::Null), + 10, + ), + AdmissionResult::Admitted { .. } + )); + // Still not initialized: requests between the response and a *valid* + // initialized are rejected, and the state has not advanced. + assert_eq!( + scheduler.lifecycle(session), + Some(LifecycleState::InitializeResponding) + ); + assert!(matches!( + scheduler.admit_message(session, request(2, "textDocument/hover"), 20), + AdmissionResult::Rejected { + error: ProtocolError { + kind: ProtocolErrorKind::ServerNotInitialized, + .. + }, + .. + } + )); + + // FIFO: initialize dispatches first; its response is ordered. + let SchedulerEvent::Dispatch(dispatch) = scheduler.next_event().unwrap() else { + panic!("expected initialize dispatch"); + }; + assert_eq!(dispatch.response_token.as_ref(), Some(&token)); + assert_eq!( + scheduler.claim_response(&token, ResponseOutcome::Success), + ResponseClaim::Claimed + ); + assert_eq!(scheduler.response_ordered(&token), ResponseOrder::Ordered); + + // The queued initialized now dispatches and advances the lifecycle. + let SchedulerEvent::Dispatch(dispatch) = scheduler.next_event().unwrap() else { + panic!("expected the queued initialized to dispatch"); + }; + assert!(matches!( + dispatch.message, + Message::Notification(ref n) if n.method == "initialized" + )); + assert_eq!( + scheduler.lifecycle(session), + Some(LifecycleState::Initialized) + ); + } + + /// If the initialize response was never ordered (failed handler → the + /// session is torn down; saturated sink → still responding), a queued + /// pipelined `initialized` is dropped at dispatch instead of advancing + /// the lifecycle past an undelivered response. + #[test] + fn pipelined_initialized_is_dropped_when_the_response_was_not_ordered() { + let mut scheduler = IngressScheduler::new(limits()).unwrap(); + let session = scheduler.open_session(TransportKind::Stdio).session_id; + + assert!(matches!( + scheduler.admit_message(session, request(1, "initialize"), 20), + AdmissionResult::Admitted { .. } + )); + assert!(matches!( + scheduler.admit_message( + session, + notification("initialized", serde_json::Value::Null), + 10, + ), + AdmissionResult::Admitted { .. } + )); + let Some(SchedulerEvent::Dispatch(_)) = scheduler.next_event() else { + panic!("expected initialize dispatch"); + }; + // The response is never claimed/ordered: the queued initialized is + // skipped and the state stays InitializeResponding. + assert!(scheduler.next_event().is_none()); + assert_eq!( + scheduler.lifecycle(session), + Some(LifecycleState::InitializeResponding) + ); + } + + /// Defect containment: a message that can never fit its class's byte + /// budget must not close the session or backpressure forever — a request + /// gets one typed `RequestFailed`, a notification/client response is + /// dropped, and later traffic on the same session is admitted normally. + #[test] + fn permanently_oversized_messages_are_contained_per_message() { + let mut scheduler = IngressScheduler::new(limits()).unwrap(); + let session = scheduler.open_session(TransportKind::Stdio).session_id; + initialize(&mut scheduler, session, 1); + + // Side-effecting request over normal+reserve bytes (200+200). + assert!(matches!( + scheduler.admit_message(session, request(2, "workspace/executeCommand"), 500), + AdmissionResult::Rejected { + error: ProtocolError { + kind: ProtocolErrorKind::RequestFailed, + .. + }, + .. + } + )); + // Read request over the read byte budget (100). + assert!(matches!( + scheduler.admit_message(session, request(3, "textDocument/inlayHint"), 150), + AdmissionResult::Rejected { + error: ProtocolError { + kind: ProtocolErrorKind::RequestFailed, + .. + }, + .. + } + )); + // Mutation notification over normal+reserve bytes: dropped, not + // backpressured (a Backpressure here would spin the transport). + assert_eq!( + scheduler.admit_message(session, full_change("file:///a.baml", 2, "huge"), 500), + AdmissionResult::Dropped(DropReason::OversizedMessage) + ); + // Client response over the ordinary byte budget: dropped. + let response = Message::Response(lsp_server::Response { + id: RequestId::from(77), + result: Some(serde_json::Value::Null), + error: None, + }); + assert_eq!( + scheduler.admit_message(session, response, 500), + AdmissionResult::Dropped(DropReason::OversizedMessage) + ); + + // The session survives: ordinary traffic is still admitted. + assert_eq!( + scheduler.lifecycle(session), + Some(LifecycleState::Initialized) + ); + assert!(matches!( + scheduler.admit_message(session, full_change("file:///a.baml", 3, "ok"), 20), + AdmissionResult::Admitted { .. } + )); + } + + /// Transient saturation of the same class still backpressures: the + /// permanent-oversize path only triggers when the size alone can never + /// fit, never for a queue that merely happens to be full. + #[test] + fn transient_saturation_still_backpressures_admissible_sizes() { + let mut tight = limits(); + tight.normal_items = 1; + tight.reserved_items = 1; + let mut scheduler = IngressScheduler::new(tight).unwrap(); + let session = scheduler.open_session(TransportKind::Stdio).session_id; + initialize(&mut scheduler, session, 1); + + assert!(matches!( + scheduler.admit_message(session, full_change("file:///a.baml", 2, "a"), 20), + AdmissionResult::Admitted { .. } + )); + assert!(matches!( + scheduler.admit_message(session, full_change("file:///b.baml", 2, "b"), 20), + AdmissionResult::Admitted { .. } + )); + assert_eq!( + scheduler.admit_message(session, full_change("file:///c.baml", 2, "c"), 20), + AdmissionResult::Backpressure(BackpressureReason::ReservedCapacity) + ); + } + + #[test] + fn protocol_error_codes_match_i7_without_unknown_error_fallback() { + assert_eq!(ProtocolErrorKind::ParseError.json_rpc_code(), -32700); + assert_eq!(ProtocolErrorKind::InvalidRequest.json_rpc_code(), -32600); + assert_eq!( + ProtocolErrorKind::ServerNotInitialized.json_rpc_code(), + -32002 + ); + assert_eq!(ProtocolErrorKind::MethodNotFound.json_rpc_code(), -32601); + assert_eq!(ProtocolErrorKind::InvalidParams.json_rpc_code(), -32602); + assert_eq!(ProtocolErrorKind::InternalError.json_rpc_code(), -32603); + assert_eq!(ProtocolErrorKind::RequestCanceled.json_rpc_code(), -32800); + assert_eq!(ProtocolErrorKind::ContentModified.json_rpc_code(), -32801); + assert_eq!(ProtocolErrorKind::RequestFailed.json_rpc_code(), -32803); + } +} diff --git a/baml_language/crates/baml_lsp_server/src/lsp_runtime.rs b/baml_language/crates/baml_lsp_server/src/lsp_runtime.rs new file mode 100644 index 00000000000..90ce67d09ea --- /dev/null +++ b/baml_language/crates/baml_lsp_server/src/lsp_runtime.rs @@ -0,0 +1,1405 @@ +//! Process-owned, transport-neutral LSP ingress runtime (design §4 Track B). +//! +//! One [`IngressScheduler`] behind one process-owned mutex serves every LSP +//! transport: the stdio loop and each `/api/lsp` browser socket submit the +//! same [`lsp_server::Message`] representation and share admission, lifecycle +//! (B2), cancellation, and outbound-response accounting. Dispatch runs on one +//! dedicated worker thread in FIFO order; `$/cancelRequest` is processed on +//! the *submitting transport thread* so it can claim a response while the +//! worker is inside a slow synchronous handler. +//! +//! Cancellation tokens handed to handlers are operation-owned and observed +//! only at safe boundaries — they are never connected to Salsa's unwind-based +//! cancellation (abort-profile invariant). + +use std::{ + collections::{HashMap, HashSet, VecDeque}, + path::{Component, Path, PathBuf}, + sync::{Arc, Mutex, mpsc}, + time::Duration, +}; + +use crate::lsp_ingress::{ + AdmissionResult, CancelEffect, ControlAdmission, IngressLimits, IngressScheduler, + LifecycleState, ProtocolError, ResponseClaim, ResponseOutcome, ResponseSizeError, + ResponseToken, SchedulerEvent, SessionId, SessionTermination, TerminationReason, TransportKind, +}; + +/// Result of a non-blocking transport enqueue attempt. +/// +/// Saturation is deliberately distinct from closure: a full bounded queue is +/// a transient backpressure signal and must not tombstone the session sink. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) enum SinkDelivery { + Sent, + Saturated, + Oversized, + Closed, +} + +pub(crate) type Sink = Arc SinkDelivery + Send + Sync>; +pub(crate) type Close = Arc; +pub(crate) type NotificationHook = Arc; + +/// Responder for one server-initiated request's client response (see +/// [`LspRuntime::register_server_request`]). +pub(crate) type ServerRequestResponder = Box; + +/// Connection-owned LSP output. The sink lives behind a tombstone so an +/// asynchronous tail retaining an old `BexLsp` session clone cannot write +/// after browser takeover (D2) or transport close. +pub(crate) struct RevocableSessionSender { + sink: Mutex>, +} + +impl RevocableSessionSender { + fn new(sink: Sink) -> Self { + Self { + sink: Mutex::new(Some(sink)), + } + } + + fn revoke(&self) { + if let Ok(mut sink) = self.sink.lock() { + sink.take(); + } + } + + #[cfg(test)] + fn is_closed(&self) -> bool { + self.sink.lock().map_or(true, |sink| sink.is_none()) + } + + fn send(&self, message: lsp_server::Message) -> Result<(), bex_project::LspError> { + let mut sink = self + .sink + .lock() + .map_err(|_| bex_project::LspError::ClientClosed)?; + let Some(active) = sink.as_ref() else { + return Err(bex_project::LspError::ClientClosed); + }; + match active(message) { + SinkDelivery::Sent => Ok(()), + SinkDelivery::Saturated => Err(bex_project::LspError::OutboundSaturated), + SinkDelivery::Oversized => Err(bex_project::LspError::OutboundOversized), + SinkDelivery::Closed => { + sink.take(); + Err(bex_project::LspError::ClientClosed) + } + } + } +} + +impl bex_project::LspClientSenderTrait for RevocableSessionSender { + fn send_notification( + &self, + notification: lsp_server::Notification, + ) -> Result<(), bex_project::LspError> { + self.send(lsp_server::Message::Notification(notification)) + } + + fn send_response_impl( + &self, + response: lsp_server::Response, + ) -> Result<(), bex_project::LspError> { + self.send(lsp_server::Message::Response(response)) + } + + fn make_request(&self, request: lsp_server::Request) -> Result<(), bex_project::LspError> { + self.send(lsp_server::Message::Request(request)) + } +} + +#[derive(Clone)] +struct Endpoint { + bex: Arc, + outbound: Arc, + sink: Sink, + close: Close, + after_notification: Option, +} + +struct PendingResponse { + token: ResponseToken, + message: lsp_server::Message, +} + +#[derive(Debug)] +pub(crate) enum SubmitResult { + Accepted, + Dropped, + Backpressure, + Exited { normal: bool }, + Closed, +} + +pub(crate) struct OpenedSession { + pub session_id: SessionId, +} + +/// One scheduler and one dispatcher thread for every LSP transport in the +/// process. Endpoints are bounded sinks owned by the transport adapters. +pub(crate) struct LspRuntime { + scheduler: Mutex, + endpoints: Mutex>, + document_overlays: Mutex>, + pending_responses: Mutex>, + /// Correlation seam for server-initiated requests (workspace/ + /// configuration and friends). This server currently never issues such + /// requests, so no dispatch path populates the map today; it exists so + /// client→server responses have a real route instead of being admitted + /// and then silently lost. Uncorrelated responses are logged and dropped. + server_requests: Mutex>, + delivery_gate: Mutex<()>, + wake_tx: mpsc::SyncSender<()>, +} + +impl LspRuntime { + pub(crate) fn new() -> anyhow::Result> { + let (runtime, wake_rx) = Self::with_limits(Self::default_limits())?; + let worker = runtime.clone(); + std::thread::Builder::new() + .name("baml-lsp-ingress".to_string()) + .spawn(move || { + loop { + match wake_rx.recv_timeout(Duration::from_millis(5)) { + Ok(()) => worker.drain(), + Err(mpsc::RecvTimeoutError::Timeout) => { + if worker.has_pending_responses() { + worker.drain_pending_responses(); + } + } + Err(mpsc::RecvTimeoutError::Disconnected) => return, + } + } + })?; + Ok(runtime) + } + + fn default_limits() -> IngressLimits { + IngressLimits { + outbound_response_bytes: 64 * 1024 * 1024, + response_reservation_bytes: 4 * 1024 * 1024, + ..IngressLimits::default() + } + } + + fn with_limits(limits: IngressLimits) -> anyhow::Result<(Arc, mpsc::Receiver<()>)> { + let scheduler = IngressScheduler::new(limits)?; + let (wake_tx, wake_rx) = mpsc::sync_channel(1); + let runtime = Arc::new(Self { + scheduler: Mutex::new(scheduler), + endpoints: Mutex::new(HashMap::new()), + document_overlays: Mutex::new(HashMap::new()), + pending_responses: Mutex::new(VecDeque::new()), + server_requests: Mutex::new(HashMap::new()), + delivery_gate: Mutex::new(()), + wake_tx, + }); + Ok((runtime, wake_rx)) + } + + pub(crate) fn open_session( + &self, + transport: TransportKind, + bex: Arc, + sink: Sink, + close: Close, + after_notification: Option, + ) -> OpenedSession { + let _delivery = self.delivery_gate.lock().unwrap(); + let opened = self.scheduler.lock().unwrap().open_session(transport); + if let Some(takeover) = opened.takeover.clone() { + // Revoke the old sink and reconcile its overlays before exposing + // the replacement endpoint (D2). A surviving stdio owner is + // restored instead of being erased by the synthetic didClose. + self.finish_termination(takeover); + } + let outbound = Arc::new(RevocableSessionSender::new(sink.clone())); + let bex = bex.new_lsp_session(outbound.clone()); + self.endpoints.lock().unwrap().insert( + opened.session_id, + Endpoint { + bex, + outbound, + sink, + close, + after_notification, + }, + ); + OpenedSession { + session_id: opened.session_id, + } + } + + pub(crate) fn submit( + &self, + session_id: SessionId, + message: lsp_server::Message, + ) -> SubmitResult { + let serialized_bytes = serde_json::to_vec(&message).map_or(1, |bytes| bytes.len()); + if let lsp_server::Message::Notification(notification) = &message + && notification.method == "$/cancelRequest" + { + let Some(request_id) = notification + .params + .get("id") + .cloned() + .and_then(|value| serde_json::from_value::(value).ok()) + else { + return SubmitResult::Dropped; + }; + let result = self.scheduler.lock().unwrap().admit_cancel( + session_id, + request_id, + serialized_bytes, + ); + return match result { + ControlAdmission::Admitted => { + // Cancellation is a true control path: process it on the + // submitting transport thread so it can claim a response + // while the sole FIFO dispatch worker is inside a slow + // synchronous handler. + self.drain_controls(); + self.wake(); + SubmitResult::Accepted + } + ControlAdmission::Duplicate => SubmitResult::Dropped, + ControlAdmission::Backpressure => SubmitResult::Backpressure, + ControlAdmission::Oversized => { + tracing::warn!( + session = session_id.get(), + "dropping $/cancelRequest larger than the control byte budget" + ); + SubmitResult::Dropped + } + ControlAdmission::UnknownSession => SubmitResult::Closed, + }; + } + + let admission = + self.scheduler + .lock() + .unwrap() + .admit_message(session_id, message, serialized_bytes); + match admission { + AdmissionResult::Admitted { .. } => { + self.wake(); + SubmitResult::Accepted + } + AdmissionResult::Rejected { + response_token, + error, + } => { + self.send_owned_error(&response_token, error); + SubmitResult::Accepted + } + AdmissionResult::Dropped(reason) => { + tracing::warn!( + session = session_id.get(), + ?reason, + "dropping LSP message (lifecycle state or permanent oversize)" + ); + SubmitResult::Dropped + } + AdmissionResult::Backpressure(_) => SubmitResult::Backpressure, + AdmissionResult::UseControlPath => SubmitResult::Dropped, + AdmissionResult::Exited(termination) => { + let normal = termination.reason == TerminationReason::NormalExit; + let _delivery = self.delivery_gate.lock().unwrap(); + self.finish_termination(termination); + SubmitResult::Exited { normal } + } + AdmissionResult::UnknownSession => SubmitResult::Closed, + AdmissionResult::DuplicateRequestId(token) => { + tracing::warn!( + session = token.session_id().get(), + request_id = %token.request_id(), + "duplicate outstanding LSP request id" + ); + SubmitResult::Dropped + } + } + } + + /// Registers interest in the client's response to a server-initiated + /// request. Call this *before* writing the request to the session sink; + /// the responder runs on the dispatch worker when the response arrives. + /// (Correlation seam for a server that starts issuing + /// `workspace/configuration`-style requests — currently unused.) + #[cfg_attr(not(test), allow(dead_code))] + pub(crate) fn register_server_request( + &self, + session_id: SessionId, + request_id: lsp_server::RequestId, + responder: ServerRequestResponder, + ) { + self.server_requests + .lock() + .unwrap() + .insert((session_id, request_id), responder); + } + + pub(crate) fn close_session(&self, session_id: SessionId) { + let _delivery = self.delivery_gate.lock().unwrap(); + self.close_session_under_delivery_gate(session_id); + } + + fn close_session_under_delivery_gate(&self, session_id: SessionId) { + if let Some(termination) = self.scheduler.lock().unwrap().close_session(session_id) { + self.finish_termination(termination); + } + } + + fn finish_termination(&self, termination: SessionTermination) { + self.pending_responses + .lock() + .unwrap() + .retain(|pending| pending.token.session_id() != termination.session_id); + self.server_requests + .lock() + .unwrap() + .retain(|(session_id, _), _| *session_id != termination.session_id); + // Drop the endpoint-map guard before replaying a surviving overlay: + // `restore_surviving_overlay` resolves that survivor through the same + // map. An `if let` scrutinee temporary otherwise lives through the + // whole body and deadlocks on browser takeover. + let endpoint = { + self.endpoints + .lock() + .unwrap() + .remove(&termination.session_id) + }; + if let Some(endpoint) = endpoint { + endpoint.outbound.revoke(); + for document in termination.owned_documents { + let identity = canonical_document_identity_str(&document.uri); + if let Some(identity) = identity.as_ref() { + self.document_overlays + .lock() + .unwrap() + .remove(&(termination.session_id, identity.clone())); + } + if !identity + .as_ref() + .is_some_and(|identity| self.restore_surviving_overlay(identity)) + && let Ok(uri) = lsp_types::Url::parse(&document.uri) + { + endpoint + .bex + .handle_notification(lsp_server::Notification::new( + "textDocument/didClose".to_string(), + lsp_types::DidCloseTextDocumentParams { + text_document: lsp_types::TextDocumentIdentifier { uri }, + }, + )); + } + } + (endpoint.close)(); + } + } + + fn restore_surviving_overlay(&self, identity: &Path) -> bool { + let survivor = self.document_overlays.lock().ok().and_then(|overlays| { + overlays + .iter() + .filter(|((_, candidate_identity), _)| candidate_identity == identity) + .max_by_key(|((session_id, _), _)| session_id.get()) + .map(|((session_id, _), document)| (*session_id, document.clone())) + }); + let Some((session_id, document)) = survivor else { + return false; + }; + let Some(endpoint) = self.endpoint(session_id) else { + return false; + }; + endpoint + .bex + .handle_notification(lsp_server::Notification::new( + "textDocument/didOpen".to_string(), + lsp_types::DidOpenTextDocumentParams { + text_document: document, + }, + )); + true + } + + fn wake(&self) { + let _ = self.wake_tx.try_send(()); + } + + fn drain_controls(&self) { + loop { + let effect = self.scheduler.lock().unwrap().next_control(); + match effect { + Some(CancelEffect::Respond { + response_token, + error, + .. + }) => self.send_owned_error(&response_token, error), + Some(CancelEffect::Noop) => {} + None => return, + } + } + } + + fn endpoint(&self, session_id: SessionId) -> Option { + self.endpoints.lock().ok()?.get(&session_id).cloned() + } + + fn drain(&self) { + self.drain_pending_responses(); + loop { + let event = self.scheduler.lock().unwrap().next_event(); + let Some(event) = event else { + return; + }; + match event { + SchedulerEvent::Control(CancelEffect::Respond { + response_token, + error, + .. + }) => self.send_owned_error(&response_token, error), + SchedulerEvent::Control(CancelEffect::Noop) => {} + SchedulerEvent::Dispatch(item) => { + let Some(endpoint) = self.endpoint(item.session_id) else { + continue; + }; + match item.message { + lsp_server::Message::Request(request) => { + let Some(token) = item.response_token else { + continue; + }; + let runtime = self; + endpoint.bex.handle_request_with_cancellation( + request, + item.cancellation.as_ref(), + &mut move |_request_id, result| { + runtime.send_handler_result(&token, result); + }, + ); + } + lsp_server::Message::Notification(notification) => { + let tracking = notification.clone(); + let applied = + endpoint.bex.handle_notification_with_status(notification); + if !applied { + continue; + } + if self.record_applied_document(item.session_id, &tracking) { + if let Some(hook) = &endpoint.after_notification { + hook(&tracking); + } + if tracking.method == "textDocument/didClose" + && let Some(uri) = tracking + .params + .pointer("/textDocument/uri") + .and_then(serde_json::Value::as_str) + && let Some(identity) = canonical_document_identity_str(uri) + { + self.restore_surviving_overlay(&identity); + } + } else { + // Browser takeover/transport close can win + // while a synchronous mutation is running. + // Close the just-applied old-session overlay + // before the FIFO worker can dispatch a new + // session's didOpen for the same URI. + self.close_orphaned_document(&endpoint, &tracking); + } + } + lsp_server::Message::Response(response) => { + // Client response to a server-initiated request: + // route through the correlation seam; without a + // registration it is logged, never silently lost. + let responder = self + .server_requests + .lock() + .unwrap() + .remove(&(item.session_id, response.id.clone())); + match responder { + Some(respond) => respond(response), + None => tracing::warn!( + session = item.session_id.get(), + id = %response.id, + "dropping client response with no matching \ + server-initiated request" + ), + } + } + } + } + } + } + } + + fn send_handler_result( + &self, + token: &ResponseToken, + result: Result, + ) { + let outcome = if result.is_ok() { + ResponseOutcome::Success + } else { + ResponseOutcome::Error + }; + let claim = self + .scheduler + .lock() + .unwrap() + .claim_response(token, outcome); + if claim != ResponseClaim::Claimed { + return; + } + if let Ok(value) = &result { + // One lock scope: read the lifecycle and record the negotiated + // initialize result under the same guard. + let mut scheduler = self.scheduler.lock().unwrap(); + if scheduler.lifecycle(token.session_id()) == Some(LifecycleState::InitializeResponding) + { + let _ = scheduler.set_negotiated_capabilities(token.session_id(), value.clone()); + } + } + let response = match result { + Ok(value) => lsp_server::Response { + id: token.request_id().clone(), + result: Some(value), + error: None, + }, + Err(error) => lsp_server::Response { + id: token.request_id().clone(), + result: None, + // The one request error-code mapping boundary (I7). + error: Some(error.to_response_error()), + }, + }; + self.order_response(token, response); + } + + fn send_owned_error(&self, token: &ResponseToken, error: ProtocolError) { + let response = lsp_server::Response { + id: token.request_id().clone(), + result: None, + error: Some(lsp_server::ResponseError { + code: error.kind.json_rpc_code(), + message: error.message, + data: None, + }), + }; + self.order_response(token, response); + } + + fn order_response(&self, token: &ResponseToken, response: lsp_server::Response) { + let _delivery = self.delivery_gate.lock().unwrap(); + let mut message = lsp_server::Message::Response(response); + let bytes = serde_json::to_vec(&message).map_or(usize::MAX, |value| value.len()); + match self + .scheduler + .lock() + .unwrap() + .validate_response_size(token, bytes) + { + Ok(()) => {} + Err(ResponseSizeError::UnknownOrSuperseded) => return, + Err(ResponseSizeError::ExceedsReservation) => { + // An oversized payload fails only its own request: the + // replacement error is a few hundred bytes and always fits + // the reservation, so the session survives instead of being + // deterministically closed. + tracing::warn!( + session = token.session_id().get(), + request_id = %token.request_id(), + bytes, + "replacing an LSP response that exceeds its reservation" + ); + message = lsp_server::Message::Response(lsp_server::Response { + id: token.request_id().clone(), + result: None, + error: Some(lsp_server::ResponseError { + code: lsp_server::ErrorCode::RequestFailed as i32, + message: "response exceeds the transport's per-response reservation" + .to_string(), + data: None, + }), + }); + } + } + let mut pending = self.pending_responses.lock().unwrap(); + if pending + .iter() + .any(|response| response.token.session_id() == token.session_id()) + { + pending.push_back(PendingResponse { + token: token.clone(), + message, + }); + return; + } + drop(pending); + let Some(endpoint) = self.endpoint(token.session_id()) else { + return; + }; + match (endpoint.sink)(message.clone()) { + SinkDelivery::Sent => { + let _ = self.scheduler.lock().unwrap().response_ordered(token); + } + SinkDelivery::Saturated => { + self.pending_responses + .lock() + .unwrap() + .push_back(PendingResponse { + token: token.clone(), + message, + }); + } + SinkDelivery::Oversized | SinkDelivery::Closed => { + self.close_session_under_delivery_gate(token.session_id()); + } + } + } + + fn has_pending_responses(&self) -> bool { + !self.pending_responses.lock().unwrap().is_empty() + } + + /// Retry each session's oldest response once without allowing one slow + /// transport to head-of-line block unrelated sessions. A response keeps + /// its scheduler reservation until it is actually accepted by the sink. + fn drain_pending_responses(&self) { + let _delivery = self.delivery_gate.lock().unwrap(); + let pending = { + let mut queue = self.pending_responses.lock().unwrap(); + std::mem::take(&mut *queue) + }; + if pending.is_empty() { + return; + } + + let mut blocked_sessions = HashSet::new(); + let mut retry = VecDeque::new(); + for pending in pending { + let session_id = pending.token.session_id(); + if blocked_sessions.contains(&session_id) { + retry.push_back(pending); + continue; + } + let Some(endpoint) = self.endpoint(session_id) else { + continue; + }; + match (endpoint.sink)(pending.message.clone()) { + SinkDelivery::Sent => { + let _ = self + .scheduler + .lock() + .unwrap() + .response_ordered(&pending.token); + } + SinkDelivery::Saturated => { + blocked_sessions.insert(session_id); + retry.push_back(pending); + } + SinkDelivery::Oversized | SinkDelivery::Closed => { + self.close_session_under_delivery_gate(session_id); + } + } + } + self.pending_responses.lock().unwrap().extend(retry); + } + + fn record_applied_document( + &self, + session_id: SessionId, + notification: &lsp_server::Notification, + ) -> bool { + let mut scheduler = self.scheduler.lock().unwrap(); + if scheduler.lifecycle(session_id).is_none() { + return false; + } + match notification.method.as_str() { + "textDocument/didOpen" => { + if let Ok(params) = serde_json::from_value::( + notification.params.clone(), + ) && let Some(identity) = canonical_document_identity(¶ms.text_document.uri) + { + self.document_overlays + .lock() + .unwrap() + .insert((session_id, identity), params.text_document.clone()); + scheduler.record_document_open( + session_id, + params.text_document.uri.to_string(), + Some(params.text_document.version), + ); + } + } + "textDocument/didChange" => { + if let Ok(params) = serde_json::from_value::( + notification.params.clone(), + ) && let Some(identity) = canonical_document_identity(¶ms.text_document.uri) + { + if let [change] = params.content_changes.as_slice() + && change.range.is_none() + && let Some(document) = self + .document_overlays + .lock() + .unwrap() + .get_mut(&(session_id, identity)) + { + document.version = params.text_document.version; + document.text.clone_from(&change.text); + } + scheduler.record_document_version( + session_id, + params.text_document.uri.as_str(), + Some(params.text_document.version), + ); + } + } + "textDocument/didClose" => { + if let Ok(params) = serde_json::from_value::( + notification.params.clone(), + ) && let Some(identity) = canonical_document_identity(¶ms.text_document.uri) + { + self.document_overlays + .lock() + .unwrap() + .remove(&(session_id, identity)); + scheduler.record_document_close(session_id, params.text_document.uri.as_str()); + } + } + _ => {} + } + true + } + + fn close_orphaned_document( + &self, + endpoint: &Endpoint, + notification: &lsp_server::Notification, + ) { + if !matches!( + notification.method.as_str(), + "textDocument/didOpen" | "textDocument/didChange" + ) { + return; + } + let Some(uri) = notification + .params + .pointer("/textDocument/uri") + .and_then(serde_json::Value::as_str) + .and_then(|uri| lsp_types::Url::parse(uri).ok()) + else { + return; + }; + if canonical_document_identity(&uri) + .as_deref() + .is_some_and(|identity| self.restore_surviving_overlay(identity)) + { + return; + } + endpoint + .bex + .handle_notification(lsp_server::Notification::new( + "textDocument/didClose".to_string(), + lsp_types::DidCloseTextDocumentParams { + text_document: lsp_types::TextDocumentIdentifier { uri }, + }, + )); + } +} + +fn canonical_document_identity(uri: &lsp_types::Url) -> Option { + let path = uri.to_file_path().ok()?; + Some(canonical_physical_path(&path)) +} + +fn canonical_document_identity_str(uri: &str) -> Option { + canonical_document_identity(&lsp_types::Url::parse(uri).ok()?) +} + +/// Resolve URI aliases to one physical identity while still retaining the +/// client's exact URI in `TextDocumentItem` for replay and publications. +fn canonical_physical_path(path: &Path) -> PathBuf { + let absolute = if path.is_absolute() { + path.to_path_buf() + } else { + std::env::current_dir().unwrap_or_default().join(path) + }; + let mut normalized = PathBuf::new(); + for component in absolute.components() { + match component { + Component::CurDir => {} + Component::ParentDir => { + normalized.pop(); + } + other => normalized.push(other.as_os_str()), + } + } + let mut existing = normalized.as_path(); + let mut missing = Vec::new(); + while !existing.exists() { + let Some(name) = existing.file_name() else { + break; + }; + missing.push(name.to_os_string()); + let Some(parent) = existing.parent() else { + break; + }; + existing = parent; + } + let mut canonical = std::fs::canonicalize(existing).unwrap_or_else(|_| existing.to_path_buf()); + for component in missing.into_iter().rev() { + canonical.push(component); + } + #[cfg(windows)] + { + canonical = PathBuf::from(canonical.to_string_lossy().to_lowercase()); + } + canonical +} + +#[cfg(test)] +mod tests { + use bex_project::LspClientSenderTrait; + use lsp_server::{Message, Notification, Request, RequestId}; + + use super::*; + use crate::lsp_ingress::{ProtocolErrorKind, ResponseOrder}; + + #[derive(Default)] + struct NoopPlaygroundSender; + + impl bex_project::PlaygroundSender for NoopPlaygroundSender { + fn send_playground_notification(&self, _notification: bex_project::PlaygroundNotification) { + } + } + + fn capturing_sink(messages: Arc>>) -> Sink { + Arc::new(move |message| { + messages.lock().unwrap().push(message); + SinkDelivery::Sent + }) + } + + fn scripted_sink( + deliveries: impl IntoIterator, + messages: Arc>>, + ) -> Sink { + let deliveries = Arc::new(Mutex::new(deliveries.into_iter().collect::>())); + Arc::new(move |message| { + messages.lock().unwrap().push(message); + deliveries + .lock() + .unwrap() + .pop_front() + .unwrap_or(SinkDelivery::Sent) + }) + } + + fn test_bex(global_messages: Arc>>) -> Arc { + let global_sender = Arc::new(RevocableSessionSender::new(capturing_sink(global_messages))); + let fs: Arc> = + Arc::new(Box::new(crate::native_vfs::NativeVfs::new())); + Arc::new(bex_project::new_lsp( + Arc::new(|_| Arc::new(sys_ops::SysOpsBuilder::new().build())), + global_sender, + Arc::new(NoopPlaygroundSender), + bex_project::BamlVFS::new(fs), + bex_project::BackgroundSpawner::new(), + )) + } + + fn runtime_without_worker() -> Arc { + runtime_without_worker_with_limits(LspRuntime::default_limits()) + } + + fn runtime_without_worker_with_limits(limits: IngressLimits) -> Arc { + LspRuntime::with_limits(limits).unwrap().0 + } + + fn request(id: i32, method: &str) -> Message { + Message::Request(Request { + id: RequestId::from(id), + method: method.to_string(), + params: serde_json::Value::Null, + }) + } + + fn initialize(runtime: &LspRuntime, session_id: SessionId) { + let mut scheduler = runtime.scheduler.lock().unwrap(); + let AdmissionResult::Admitted { + response_token: Some(token), + .. + } = scheduler.admit_message(session_id, request(1, "initialize"), 20) + else { + panic!("initialize must be admitted"); + }; + let Some(SchedulerEvent::Dispatch(_)) = scheduler.next_event() else { + panic!("initialize must dispatch"); + }; + assert_eq!( + scheduler.claim_response(&token, ResponseOutcome::Success), + ResponseClaim::Claimed + ); + assert_eq!(scheduler.response_ordered(&token), ResponseOrder::Ordered); + assert!(matches!( + scheduler.admit_message( + session_id, + Message::Notification(Notification::new( + "initialized".to_string(), + serde_json::Value::Null, + )), + 10, + ), + AdmissionResult::Admitted { .. } + )); + let Some(SchedulerEvent::Dispatch(_)) = scheduler.next_event() else { + panic!("initialized must dispatch"); + }; + } + + fn admit_command(runtime: &LspRuntime, session_id: SessionId, id: i32) -> ResponseToken { + let AdmissionResult::Admitted { + response_token: Some(token), + .. + } = runtime.scheduler.lock().unwrap().admit_message( + session_id, + request(id, "workspace/executeCommand"), + 20, + ) + else { + panic!("executeCommand must be admitted"); + }; + token + } + + fn cancel(id: i32) -> Message { + Message::Notification(Notification::new( + "$/cancelRequest".to_string(), + serde_json::json!({ "id": id }), + )) + } + + #[test] + fn browser_takeover_tombstones_old_outbound_without_leaking_to_new_or_stdio() { + let runtime = runtime_without_worker(); + let global_messages = Arc::new(Mutex::new(Vec::new())); + let bex = test_bex(global_messages.clone()); + let stdio_messages = Arc::new(Mutex::new(Vec::new())); + let old_browser_messages = Arc::new(Mutex::new(Vec::new())); + let new_browser_messages = Arc::new(Mutex::new(Vec::new())); + let close: Close = Arc::new(|| {}); + + let stdio = runtime.open_session( + TransportKind::Stdio, + bex.clone(), + capturing_sink(stdio_messages.clone()), + close.clone(), + None, + ); + let old_browser = runtime.open_session( + TransportKind::Browser, + bex.clone(), + capturing_sink(old_browser_messages.clone()), + close.clone(), + None, + ); + let shared_uri = "file:///tmp/shared-overlay.baml".to_string(); + let shared_identity = canonical_document_identity_str(&shared_uri).unwrap(); + { + let mut scheduler = runtime.scheduler.lock().unwrap(); + scheduler.record_document_open(stdio.session_id, shared_uri.clone(), Some(3)); + scheduler.record_document_open(old_browser.session_id, shared_uri.clone(), Some(7)); + } + runtime.document_overlays.lock().unwrap().insert( + (stdio.session_id, shared_identity.clone()), + lsp_types::TextDocumentItem { + uri: lsp_types::Url::parse(&shared_uri).unwrap(), + language_id: "baml".to_string(), + version: 3, + text: "function stdio() -> int { 3 }".to_string(), + }, + ); + runtime.document_overlays.lock().unwrap().insert( + (old_browser.session_id, shared_identity.clone()), + lsp_types::TextDocumentItem { + uri: lsp_types::Url::parse(&shared_uri).unwrap(), + language_id: "baml".to_string(), + version: 7, + text: "function browser() -> int { 7 }".to_string(), + }, + ); + let stale_dispatcher = runtime.endpoint(old_browser.session_id).unwrap().bex; + let new_browser = runtime.open_session( + TransportKind::Browser, + bex, + capturing_sink(new_browser_messages.clone()), + close, + None, + ); + let stdio_messages_after_restore = stdio_messages.lock().unwrap().len(); + + // Unsupported incoming notifications cause BexLsp to emit a + // window/logMessage notification through the dispatcher's sender. + // The retained old dispatcher models a late diagnostic/tail clone. + stale_dispatcher.handle_notification(Notification::new( + "test/late-old-browser".to_string(), + serde_json::Value::Null, + )); + runtime + .endpoint(stdio.session_id) + .unwrap() + .bex + .handle_notification(Notification::new( + "test/stdio".to_string(), + serde_json::Value::Null, + )); + runtime + .endpoint(new_browser.session_id) + .unwrap() + .bex + .handle_notification(Notification::new( + "test/new-browser".to_string(), + serde_json::Value::Null, + )); + + assert!(old_browser_messages.lock().unwrap().is_empty()); + assert_eq!( + stdio_messages.lock().unwrap().len(), + stdio_messages_after_restore + 1 + ); + assert_eq!(new_browser_messages.lock().unwrap().len(), 1); + assert!(global_messages.lock().unwrap().is_empty()); + let overlays = runtime.document_overlays.lock().unwrap(); + assert!(overlays.contains_key(&(stdio.session_id, shared_identity.clone()))); + assert!(!overlays.contains_key(&(old_browser.session_id, shared_identity))); + } + + #[cfg(unix)] + #[test] + fn document_overlay_identity_resolves_symlink_uri_aliases() { + let root = std::env::temp_dir().join(format!( + "baml_lsp_overlay_alias_{}_{}", + std::process::id(), + std::thread::current().name().unwrap_or("test") + )); + let real = root.join("real"); + let alias = root.join("alias"); + let _ = std::fs::remove_dir_all(&root); + std::fs::create_dir_all(&real).unwrap(); + std::os::unix::fs::symlink(&real, &alias).unwrap(); + let real_uri = lsp_types::Url::from_file_path(real.join("main.baml")).unwrap(); + let alias_uri = lsp_types::Url::from_file_path(alias.join("main.baml")).unwrap(); + + assert_ne!(real_uri, alias_uri); + assert_eq!( + canonical_document_identity(&real_uri), + canonical_document_identity(&alias_uri), + "different client URIs for one physical path must share overlay ownership" + ); + std::fs::remove_dir_all(&root).unwrap(); + } + + #[test] + fn transient_sink_saturation_does_not_revoke_session_sender() { + let messages = Arc::new(Mutex::new(Vec::new())); + let sender = RevocableSessionSender::new(scripted_sink( + [SinkDelivery::Saturated, SinkDelivery::Sent], + messages.clone(), + )); + + let saturated = sender.send_notification(Notification::new( + "test/first".to_string(), + serde_json::Value::Null, + )); + assert!(matches!( + saturated, + Err(bex_project::LspError::OutboundSaturated) + )); + assert!(!sender.is_closed()); + assert!( + sender + .send_notification(Notification::new( + "test/second".to_string(), + serde_json::Value::Null, + )) + .is_ok(), + "the live sink must remain installed after transient backpressure" + ); + assert_eq!(messages.lock().unwrap().len(), 2); + } + + #[test] + fn closed_sink_revokes_session_sender() { + let messages = Arc::new(Mutex::new(Vec::new())); + let sender = RevocableSessionSender::new(scripted_sink( + [SinkDelivery::Closed, SinkDelivery::Sent], + messages.clone(), + )); + + assert!(matches!( + sender.send_notification(Notification::new( + "test/closed".to_string(), + serde_json::Value::Null, + )), + Err(bex_project::LspError::ClientClosed) + )); + assert!(sender.is_closed()); + assert!(matches!( + sender.send_notification(Notification::new( + "test/late".to_string(), + serde_json::Value::Null, + )), + Err(bex_project::LspError::ClientClosed) + )); + assert_eq!( + messages.lock().unwrap().len(), + 1, + "a closed sink is tombstoned after the first failed delivery" + ); + } + + #[test] + fn saturated_response_remains_reserved_until_retry_is_ordered() { + let runtime = runtime_without_worker(); + let global_messages = Arc::new(Mutex::new(Vec::new())); + let bex = test_bex(global_messages); + let messages = Arc::new(Mutex::new(Vec::new())); + let opened = runtime.open_session( + TransportKind::Stdio, + bex, + scripted_sink( + [SinkDelivery::Saturated, SinkDelivery::Sent], + messages.clone(), + ), + Arc::new(|| {}), + None, + ); + let token = { + let mut scheduler = runtime.scheduler.lock().unwrap(); + let AdmissionResult::Admitted { + response_token: Some(token), + .. + } = scheduler.admit_message(opened.session_id, request(1, "initialize"), 20) + else { + panic!("initialize must be admitted"); + }; + assert!(matches!( + scheduler.next_event(), + Some(SchedulerEvent::Dispatch(_)) + )); + assert_eq!( + scheduler.claim_response(&token, ResponseOutcome::Success), + ResponseClaim::Claimed + ); + token + }; + + runtime.order_response( + &token, + lsp_server::Response { + id: token.request_id().clone(), + result: Some(serde_json::json!({ "capabilities": {} })), + error: None, + }, + ); + assert_eq!(runtime.pending_responses.lock().unwrap().len(), 1); + assert_eq!( + runtime + .scheduler + .lock() + .unwrap() + .lifecycle(opened.session_id), + Some(LifecycleState::InitializeResponding), + "lifecycle must not advance before the response reaches the sink" + ); + assert!(runtime.endpoint(opened.session_id).is_some()); + + runtime.drain_pending_responses(); + + assert!(runtime.pending_responses.lock().unwrap().is_empty()); + assert_eq!( + runtime + .scheduler + .lock() + .unwrap() + .lifecycle(opened.session_id), + Some(LifecycleState::InitializeResponded) + ); + assert_eq!(messages.lock().unwrap().len(), 2); + } + + /// Defect containment: a response payload exceeding the per-response + /// reservation is replaced by a `RequestFailed` error for that request + /// only — the session stays open and the reservation is released. + #[test] + fn oversized_response_payload_fails_only_its_own_request() { + let limits = IngressLimits { + response_reservation_bytes: 256, + ..IngressLimits::default() + }; + let runtime = runtime_without_worker_with_limits(limits); + let global_messages = Arc::new(Mutex::new(Vec::new())); + let bex = test_bex(global_messages); + let messages = Arc::new(Mutex::new(Vec::new())); + let opened = runtime.open_session( + TransportKind::Stdio, + bex, + capturing_sink(messages.clone()), + Arc::new(|| {}), + None, + ); + initialize(&runtime, opened.session_id); + let AdmissionResult::Admitted { + response_token: Some(token), + .. + } = runtime.scheduler.lock().unwrap().admit_message( + opened.session_id, + request(2, "textDocument/hover"), + 20, + ) + else { + panic!("hover must be admitted"); + }; + assert!(matches!( + runtime.scheduler.lock().unwrap().next_event(), + Some(SchedulerEvent::Dispatch(_)) + )); + + runtime.send_handler_result(&token, Ok(serde_json::json!("x".repeat(1024)))); + + let messages = messages.lock().unwrap(); + let Some(Message::Response(response)) = messages.last() else { + panic!("the replacement error response must reach the sink"); + }; + assert_eq!(response.id, RequestId::from(2)); + let error = response.error.as_ref().expect("must be an error response"); + assert_eq!(error.code, ProtocolErrorKind::RequestFailed.json_rpc_code()); + // Session alive, reservation released (request no longer outstanding). + assert!(runtime.endpoint(opened.session_id).is_some()); + assert_eq!( + runtime.scheduler.lock().unwrap().outstanding_state(&token), + None + ); + } + + /// Defect containment: client responses to server-initiated requests + /// route through the correlation seam instead of being admitted and then + /// silently dropped. + #[test] + fn client_response_routes_through_the_server_request_seam() { + let runtime = runtime_without_worker(); + let global_messages = Arc::new(Mutex::new(Vec::new())); + let bex = test_bex(global_messages); + let messages = Arc::new(Mutex::new(Vec::new())); + let opened = runtime.open_session( + TransportKind::Stdio, + bex, + capturing_sink(messages), + Arc::new(|| {}), + None, + ); + initialize(&runtime, opened.session_id); + + let observed = Arc::new(Mutex::new(None)); + let observed_in_responder = observed.clone(); + runtime.register_server_request( + opened.session_id, + RequestId::from(41), + Box::new(move |response| { + *observed_in_responder.lock().unwrap() = Some(response); + }), + ); + + // An uncorrelated response is logged and dropped, not fatal. + assert!(matches!( + runtime.submit( + opened.session_id, + Message::Response(lsp_server::Response { + id: RequestId::from(99), + result: Some(serde_json::Value::Null), + error: None, + }), + ), + SubmitResult::Accepted + )); + // The registered response reaches its responder via dispatch. + assert!(matches!( + runtime.submit( + opened.session_id, + Message::Response(lsp_server::Response { + id: RequestId::from(41), + result: Some(serde_json::json!({ "answer": 42 })), + error: None, + }), + ), + SubmitResult::Accepted + )); + runtime.drain(); + + let observed = observed.lock().unwrap(); + let response = observed.as_ref().expect("responder must run"); + assert_eq!(response.id, RequestId::from(41)); + assert_eq!(response.result, Some(serde_json::json!({ "answer": 42 })),); + assert!(runtime.server_requests.lock().unwrap().is_empty()); + } + + #[test] + fn execute_command_cancel_before_dispatch_owns_response() { + let runtime = runtime_without_worker(); + let session_id = runtime + .scheduler + .lock() + .unwrap() + .open_session(TransportKind::Stdio) + .session_id; + initialize(&runtime, session_id); + let token = admit_command(&runtime, session_id, 2); + + assert!(matches!( + runtime.submit(session_id, cancel(2)), + SubmitResult::Accepted + )); + assert_eq!( + runtime + .scheduler + .lock() + .unwrap() + .claim_response(&token, ResponseOutcome::Success), + ResponseClaim::AlreadyClaimed + ); + } + + #[test] + fn execute_command_cancel_after_dispatch_leaves_normal_response_owner() { + let runtime = runtime_without_worker(); + let session_id = runtime + .scheduler + .lock() + .unwrap() + .open_session(TransportKind::Stdio) + .session_id; + initialize(&runtime, session_id); + let token = admit_command(&runtime, session_id, 2); + let Some(SchedulerEvent::Dispatch(dispatch)) = + runtime.scheduler.lock().unwrap().next_event() + else { + panic!("executeCommand must transition to running"); + }; + assert!(!dispatch.cancellation.unwrap().is_cancelled()); + + assert!(matches!( + runtime.submit(session_id, cancel(2)), + SubmitResult::Accepted + )); + assert_eq!( + runtime + .scheduler + .lock() + .unwrap() + .claim_response(&token, ResponseOutcome::Success), + ResponseClaim::Claimed, + "a running side effect must never be reported as canceled" + ); + } +} diff --git a/baml_language/crates/baml_lsp_server/src/native_lsp_sender.rs b/baml_language/crates/baml_lsp_server/src/native_lsp_sender.rs index 450d88ad3de..4cfca90512a 100644 --- a/baml_language/crates/baml_lsp_server/src/native_lsp_sender.rs +++ b/baml_language/crates/baml_lsp_server/src/native_lsp_sender.rs @@ -1,52 +1,113 @@ //! Native LSP client sender over crossbeam channel to stdio. //! -//! Implements `bex_project::LspClientSenderTrait` by writing -//! `lsp_server::Message` frames to a `crossbeam::channel::Sender`. +//! Implements `bex_project::LspClientSenderTrait` by reserving bounded +//! [`crate::OutboundFrame`]s and enqueueing them for the stdio writer thread. +//! Saturation and oversize map to their typed `LspError`s instead of blocking +//! or silently dropping (B2: no unbounded writer queue). use std::sync::Weak; use bex_project::LspError; use crossbeam_channel::Sender; -use lsp_server::Message; + +use crate::{OutboundBudget, OutboundFrame, OutboundReserveError}; #[derive(Clone)] pub struct NativeLspSender { - weak: Weak>, + weak: Weak>, + budget: Weak, } impl NativeLspSender { - pub fn new(sender: &std::sync::Arc>) -> Self { + pub fn new( + sender: &std::sync::Arc>, + budget: &std::sync::Arc, + ) -> Self { Self { weak: std::sync::Arc::downgrade(sender), + budget: std::sync::Arc::downgrade(budget), + } + } + + fn send(&self, message: lsp_server::Message) -> Result<(), LspError> { + let sender = self.weak.upgrade().ok_or(LspError::ClientClosed)?; + let budget = self.budget.upgrade().ok_or(LspError::ClientClosed)?; + let frame = budget.try_message(message).map_err(|error| match error { + OutboundReserveError::Saturated => LspError::OutboundSaturated, + OutboundReserveError::Oversized | OutboundReserveError::Serialization => { + LspError::OutboundOversized + } + })?; + match sender.try_send(frame) { + Ok(()) => Ok(()), + Err(crossbeam_channel::TrySendError::Full(_)) => Err(LspError::OutboundSaturated), + Err(crossbeam_channel::TrySendError::Disconnected(_)) => Err(LspError::ClientClosed), } } } impl bex_project::LspClientSenderTrait for NativeLspSender { fn send_notification(&self, notif: lsp_server::Notification) -> Result<(), LspError> { - let Some(sender) = self.weak.upgrade() else { - return Err(LspError::ClientClosed); - }; - sender - .send(Message::Notification(notif)) - .map_err(|_| LspError::ClientClosed) + self.send(lsp_server::Message::Notification(notif)) } fn send_response_impl(&self, response: lsp_server::Response) -> Result<(), LspError> { - let Some(sender) = self.weak.upgrade() else { - return Err(LspError::ClientClosed); - }; - sender - .send(Message::Response(response)) - .map_err(|_| LspError::ClientClosed) + self.send(lsp_server::Message::Response(response)) } fn make_request(&self, req: lsp_server::Request) -> Result<(), LspError> { - let Some(sender) = self.weak.upgrade() else { - return Err(LspError::ClientClosed); - }; - sender - .send(Message::Request(req)) - .map_err(|_| LspError::ClientClosed) + self.send(lsp_server::Message::Request(req)) + } +} + +#[cfg(test)] +mod tests { + use bex_project::LspClientSenderTrait; + + use super::*; + + fn notification(method: &str) -> lsp_server::Notification { + lsp_server::Notification::new(method.to_string(), serde_json::Value::Null) + } + + #[test] + fn full_native_queue_is_backpressure_not_transport_closure() { + let (tx, _rx) = crossbeam_channel::bounded(1); + let tx = std::sync::Arc::new(tx); + let budget = OutboundBudget::new(); + let sender = NativeLspSender::new(&tx, &budget); + + assert!(sender.send_notification(notification("test/first")).is_ok()); + assert!(matches!( + sender.send_notification(notification("test/full")), + Err(LspError::OutboundSaturated) + )); + + drop(tx); + assert!(matches!( + sender.send_notification(notification("test/closed")), + Err(LspError::ClientClosed) + )); + } + + #[test] + fn oversized_native_frame_is_distinct_from_temporary_backpressure() { + let (tx, _rx) = crossbeam_channel::bounded(1); + let tx = std::sync::Arc::new(tx); + let budget = OutboundBudget::new(); + let sender = NativeLspSender::new(&tx, &budget); + let oversized = lsp_server::Notification::new( + "test/oversized".to_string(), + serde_json::Value::String("x".repeat(crate::MAX_OUTBOUND_FRAME_BYTES + 1)), + ); + + assert!(matches!( + sender.send_notification(oversized), + Err(LspError::OutboundOversized) + )); + assert!( + sender.send_notification(notification("test/after")).is_ok(), + "an oversized frame must not poison the sender" + ); } } diff --git a/baml_language/crates/baml_lsp_server/src/playground_server.rs b/baml_language/crates/baml_lsp_server/src/playground_server.rs index 3b9bc4dce9f..a2f72c50caa 100644 --- a/baml_language/crates/baml_lsp_server/src/playground_server.rs +++ b/baml_language/crates/baml_lsp_server/src/playground_server.rs @@ -489,8 +489,12 @@ struct WsState { history_store: Arc, _history_observer_registration: Arc, value_store: LiveValueStore, - /// LSP output (responses + publishDiagnostics) destined for `/api/lsp`. - lsp_out_tx: broadcast::Sender, + /// Broadcast LSP output (root-dispatcher notifications, disk-change + /// pushes) destined for `/api/lsp`. Responses never travel here: they are + /// routed per-session by the ingress runtime. + lsp_out_tx: broadcast::Sender, + /// Process-owned ingress runtime shared with the stdio transport (B2). + lsp_runtime: Arc, /// What the browser currently has per file (for disk-watcher echo avoidance). doc_mirror: DocMirror, /// Workspace roots that browser-mode LSP saves are allowed to write under. @@ -960,7 +964,7 @@ struct UpdateSourceFileResponse { /// Start the playground server on the given listener. #[allow(clippy::too_many_arguments)] -pub async fn run( +pub(crate) async fn run( listener: TcpListener, bex: Arc, broadcast_tx: broadcast::Sender, @@ -968,7 +972,8 @@ pub async fn run( io_state: Arc, run_store: Arc, playground_dir_override: Option, - lsp_out_tx: broadcast::Sender, + lsp_out_tx: broadcast::Sender, + lsp_runtime: Arc, doc_mirror: DocMirror, workspace_roots: Vec, session_token: Option>, @@ -984,6 +989,7 @@ pub async fn run( run_store, playground_dir_override, lsp_out_tx, + lsp_runtime, doc_mirror, Arc::new(workspace_roots), access_guard, @@ -1005,7 +1011,8 @@ fn build_router( io_state: Arc, run_store: Arc, playground_dir_override: Option, - lsp_out_tx: broadcast::Sender, + lsp_out_tx: broadcast::Sender, + lsp_runtime: Arc, doc_mirror: DocMirror, workspace_roots: Arc>, access_guard: PlaygroundAccessGuard, @@ -1026,6 +1033,7 @@ fn build_router( _history_observer_registration: history_observer_registration, value_store, lsp_out_tx, + lsp_runtime, doc_mirror, workspace_roots, current_open_target, @@ -1097,64 +1105,98 @@ async fn lsp_ws_handler(State(state): State, ws: WebSocketUpgrade) -> R ws.on_upgrade(move |socket| lsp_ws_session(socket, state)) } -/// Serialize an `lsp_server::Message` as a JSON-RPC text frame (adds the -/// `jsonrpc` field that `lsp_server` only writes via its stdio framing). -fn lsp_message_to_ws_text(msg: &lsp_server::Message) -> Option { - let mut value = match serde_json::to_value(msg) { - Ok(v) => v, +/// One pre-serialized outbound frame as a WS text message. The bytes already +/// carry the `jsonrpc` member (see `crate::serialize_jsonrpc_message`). +fn frame_to_ws_text(frame: &crate::OutboundFrame) -> Option { + match std::str::from_utf8(frame.bytes()) { + Ok(text) => Some(AxumWsMsg::Text(text.to_string().into())), Err(e) => { - tracing::error!("LSP WS: failed to serialize message: {e}"); - return None; - } - }; - if let Some(obj) = value.as_object_mut() { - obj.insert( - "jsonrpc".to_string(), - serde_json::Value::String("2.0".to_string()), - ); - } - match serde_json::to_string(&value) { - Ok(text) => Some(AxumWsMsg::Text(text.into())), - Err(e) => { - tracing::error!("LSP WS: failed to stringify message: {e}"); + tracing::error!("LSP WS: outbound frame is not UTF-8: {e}"); None } } } +/// A browser session that cannot drain within this deadline is +/// deterministically closed (B2: slow readers cannot grow an unbounded +/// writer queue). +async fn send_lsp_ws_message( + sink: &mut futures::stream::SplitSink, + message: AxumWsMsg, +) -> bool { + matches!( + tokio::time::timeout(std::time::Duration::from_secs(5), sink.send(message)).await, + Ok(Ok(())) + ) +} + async fn lsp_ws_session(socket: WebSocket, state: WsState) { tracing::info!("Playground: LSP WS session started"); let (mut sink, mut stream) = socket.split(); let mut lsp_rx = state.lsp_out_tx.subscribe(); - // Dispatch `bex` LSP work (which can recompile the project on every keystroke - // via didChange) on a DEDICATED thread instead of the async WS task. Calling - // the synchronous, CPU-heavy `bex.handle_*` inline would block this task's - // select! — so a save (and its disk write-through) would queue behind the - // backlog of didChange recompiles and only land once the user stops typing. - // The thread processes messages in order; the WS loop stays responsive. - let (dispatch_tx, dispatch_rx) = std::sync::mpsc::channel::(); - let bex = state.bex.clone(); - let _dispatch_thread = std::thread::Builder::new() - .name("lsp-ws-dispatch".into()) - .spawn(move || { - while let Ok(msg) = dispatch_rx.recv() { - match msg { - lsp_server::Message::Request(req) => bex.handle_request(req), - lsp_server::Message::Notification(notif) => bex.handle_notification(notif), - lsp_server::Message::Response(_) => {} - } + // The session's response channel: the ingress runtime orders responses + // (and session-scoped notifications) into this bounded, budgeted queue; + // the WS loop forwards them. Saturation is backpressure — the runtime + // retries with the response still reserved — never loss. + let (response_tx, mut response_rx) = tokio::sync::mpsc::channel(256); + let response_budget = crate::OutboundBudget::new(); + let response_sink_budget = response_budget.clone(); + let response_sink: crate::lsp_runtime::Sink = Arc::new(move |message| { + let frame = match response_sink_budget.try_message(message) { + Ok(frame) => frame, + Err(crate::OutboundReserveError::Saturated) => { + return crate::lsp_runtime::SinkDelivery::Saturated; } - }); - - // Latest full text per document URI, captured from didOpen/didChange so we - // can write it through to disk on didSave. Disk writes happen ONLY on an - // explicit save (Cmd+S → didSave) — never on didChange — so the browser and - // any other editor (e.g. VS Code) don't race to write the same file. The - // disk watcher still pushes external edits to the browser live; applying one - // fires a didChange that is NOT written back, so there's no edit loop. - let mut pending_text: std::collections::HashMap = - std::collections::HashMap::new(); + Err( + crate::OutboundReserveError::Oversized | crate::OutboundReserveError::Serialization, + ) => return crate::lsp_runtime::SinkDelivery::Oversized, + }; + match response_tx.try_send(frame) { + Ok(()) => crate::lsp_runtime::SinkDelivery::Sent, + Err(tokio::sync::mpsc::error::TrySendError::Full(_)) => { + crate::lsp_runtime::SinkDelivery::Saturated + } + Err(tokio::sync::mpsc::error::TrySendError::Closed(_)) => { + crate::lsp_runtime::SinkDelivery::Closed + } + } + }); + // Browser takeover (D2) closes the superseded socket through this signal. + let (close_tx, mut close_rx) = tokio::sync::watch::channel(false); + let close_endpoint: crate::lsp_runtime::Close = Arc::new(move || { + let _ = close_tx.send(true); + }); + // Latest full text per document URI, captured from applied didOpen/ + // didChange so didSave can write it through to disk. The hook runs on the + // dispatch worker strictly after the mutation was accepted, so persisted + // bytes always correspond to an applied overlay. + let pending_text = Arc::new(std::sync::Mutex::new(std::collections::HashMap::< + String, + String, + >::new())); + let hook_pending_text = pending_text.clone(); + let hook_doc_mirror = state.doc_mirror.clone(); + let hook_workspace_roots = state.workspace_roots.clone(); + let after_notification: crate::lsp_runtime::NotificationHook = Arc::new(move |notification| { + let Ok(mut pending_text) = hook_pending_text.lock() else { + return; + }; + track_and_persist_lsp_notification( + notification, + &mut pending_text, + &hook_doc_mirror, + &hook_workspace_roots, + ); + }); + let opened = state.lsp_runtime.open_session( + crate::lsp_ingress::TransportKind::Browser, + state.bex.clone(), + response_sink, + close_endpoint, + Some(after_notification), + ); + let session_id = opened.session_id; loop { tokio::select! { @@ -1164,11 +1206,9 @@ async fn lsp_ws_session(socket: WebSocket, state: WsState) { let text_str: &str = &text; handle_lsp_client_text( text_str, - &dispatch_tx, - &state.doc_mirror, - &state.workspace_roots, + &state.lsp_runtime, + session_id, &mut sink, - &mut pending_text, ).await; } Some(Ok(AxumWsMsg::Close(_))) | None => break, @@ -1177,77 +1217,93 @@ async fn lsp_ws_session(socket: WebSocket, state: WsState) { } out_msg = lsp_rx.recv() => { match out_msg { - Ok(msg) => { - if let Some(ws_msg) = lsp_message_to_ws_text(&msg) - && sink.send(ws_msg).await.is_err() + // Responses are per-session (routed via `response_rx`); + // a lossy broadcast is not a response transport. + Ok(frame) if !frame.is_response() => { + if let Some(ws_msg) = frame_to_ws_text(&frame) + && !send_lsp_ws_message(&mut sink, ws_msg).await { break; } } + Ok(_) => {} Err(broadcast::error::RecvError::Lagged(n)) => { + // Overflow closes only this session; the process and + // other sessions continue. tracing::warn!("LSP WS: broadcast lagged by {n} messages"); + break; } Err(broadcast::error::RecvError::Closed) => break, } } + response = response_rx.recv() => { + match response { + Some(frame) => { + if let Some(ws_msg) = frame_to_ws_text(&frame) + && !send_lsp_ws_message(&mut sink, ws_msg).await + { + break; + } + } + None => break, + } + } + changed = close_rx.changed() => { + if changed.is_err() || *close_rx.borrow() { + break; + } + } } } - // Dropping `dispatch_tx` ends the dispatch thread (its recv() returns Err). - drop(dispatch_tx); + state.lsp_runtime.close_session(session_id); tracing::debug!("LSP WS session ended"); } async fn handle_lsp_client_text( text: &str, - dispatch_tx: &std::sync::mpsc::Sender, - doc_mirror: &DocMirror, - workspace_roots: &[PathBuf], + runtime: &Arc, + session_id: crate::lsp_ingress::SessionId, sink: &mut futures::stream::SplitSink, - pending_text: &mut std::collections::HashMap, ) { - let msg = match serde_json::from_str::(text) { - Ok(m) => m, - Err(e) => { - tracing::warn!("LSP WS: invalid JSON-RPC message: {e}"); + // Malformed JSON and invalid envelopes get the same null-ID protocol + // errors the stdio transport produces (I7: transport-identical traces). + let value = match serde_json::from_str::(text) { + Ok(value) => value, + Err(error) => { + tracing::warn!("LSP WS: malformed JSON: {error}"); + let parse_error = serde_json::json!({ + "jsonrpc": "2.0", + "id": serde_json::Value::Null, + "error": { "code": -32700, "message": format!("Parse error: {error}") }, + }); + let _ = + send_lsp_ws_message(sink, AxumWsMsg::Text(parse_error.to_string().into())).await; return; } }; - - match msg { - lsp_server::Message::Request(req) => { - // Mirror the stdio loop: answer `shutdown` directly without - // dispatching it into BexLsp. - if req.method == "shutdown" { - let response = lsp_server::Response { - id: req.id, - result: Some(serde_json::Value::Null), - error: None, - }; - if let Some(ws_msg) = - lsp_message_to_ws_text(&lsp_server::Message::Response(response)) - { - let _ = sink.send(ws_msg).await; - } - return; - } - let _ = dispatch_tx.send(lsp_server::Message::Request(req)); + let msg = match crate::decode_lsp_message(value) { + Ok(message) => message, + Err(error) => { + let invalid_request = serde_json::json!({ + "jsonrpc": "2.0", + "id": serde_json::Value::Null, + "error": { "code": -32600, "message": error }, + }); + let _ = send_lsp_ws_message(sink, AxumWsMsg::Text(invalid_request.to_string().into())) + .await; + return; } - lsp_server::Message::Notification(notif) => { - if notif.method == "exit" { - return; + }; + loop { + match runtime.submit(session_id, msg.clone()) { + crate::lsp_runtime::SubmitResult::Accepted + | crate::lsp_runtime::SubmitResult::Dropped + | crate::lsp_runtime::SubmitResult::Exited { .. } + | crate::lsp_runtime::SubmitResult::Closed => break, + crate::lsp_runtime::SubmitResult::Backpressure => { + tokio::time::sleep(std::time::Duration::from_millis(2)).await; } - // Capture the latest text and, on an explicit save (didSave), write - // it through to disk RIGHT HERE on the WS read loop — not the dispatch - // thread — so the save is immediate and independent of the recompile - // backlog. didChange never writes to disk. BexLsp still sees the - // notification (via the dispatch thread) for live diagnostics. - track_and_persist_lsp_notification(¬if, pending_text, doc_mirror, workspace_roots); - let _ = dispatch_tx.send(lsp_server::Message::Notification(notif)); - } - lsp_server::Message::Response(_) => { - // Responses to server-initiated requests; the stdio loop ignores - // these as well. } } } @@ -1392,9 +1448,10 @@ fn update_doc_mirror(doc_mirror: &DocMirror, uri: &str, text: &str) { /// a change whose content already matches `doc_mirror` (what the browser has, or /// just wrote through) is NOT pushed back. The returned watcher must be kept /// alive for as long as watching should continue. -pub fn spawn_disk_watcher( +pub(crate) fn spawn_disk_watcher( roots: &[PathBuf], - lsp_out_tx: broadcast::Sender, + lsp_out_tx: broadcast::Sender, + lsp_out_budget: Arc, doc_mirror: DocMirror, ) -> Option { use notify::{EventKind, RecursiveMode, Watcher}; @@ -1429,7 +1486,25 @@ pub fn spawn_disk_watcher( method: DISK_CHANGE_NOTIFICATION.to_string(), params: serde_json::json!({ "uri": url.to_string(), "text": content }), }; - let _ = lsp_out_tx.send(lsp_server::Message::Notification(notif)); + let message = lsp_server::Message::Notification(notif); + // Disk pushes are best-effort refresh hints: budget exhaustion + // drops the hint (the browser re-reads on save/reload) instead of + // growing an unbounded queue. + let frame = match lsp_out_budget.try_message(message) { + Ok(frame) => frame, + Err(crate::OutboundReserveError::Saturated) => { + tracing::warn!("Disk watcher: LSP outbound byte budget is saturated"); + continue; + } + Err( + crate::OutboundReserveError::Oversized + | crate::OutboundReserveError::Serialization, + ) => { + tracing::warn!("Disk watcher: LSP outbound frame is not deliverable"); + continue; + } + }; + let _ = lsp_out_tx.send(frame); tracing::debug!( "Disk watcher: pushed external change for {}", canonical.display() diff --git a/baml_language/crates/bex_project/src/bex_lsp/mod.rs b/baml_language/crates/bex_project/src/bex_lsp/mod.rs index cf5c2d217fb..664d2da8354 100644 --- a/baml_language/crates/bex_project/src/bex_lsp/mod.rs +++ b/baml_language/crates/bex_project/src/bex_lsp/mod.rs @@ -7,6 +7,7 @@ mod request; mod multi_project; mod position_codec; mod protocol; +pub(crate) mod request_cancellation; use std::{path::PathBuf, sync::Arc}; @@ -35,6 +36,16 @@ pub enum LspError { #[error("Client closed")] ClientClosed, + /// The connection-owned outbound sink cannot accept more frames right + /// now (bounded transport backpressure; LSP `RequestFailed`, `-32803`). + #[error("LSP outbound sink is saturated")] + OutboundSaturated, + + /// The serialized outbound frame exceeds the transport's frame limit + /// (LSP `RequestFailed`, `-32803`). + #[error("LSP outbound frame exceeds the transport limit")] + OutboundOversized, + #[error("Root path not found: {}: {}", .0.as_str(), .1)] ProjectRootNotFound(vfs::VfsPath, String), @@ -64,6 +75,12 @@ pub enum LspError { #[error("No projects found")] NoProjectsFound, + /// Explicit client cancellation won response ownership while the request + /// was queued or running (LSP `RequestCanceled`, `-32800`). Observed only + /// at safe handler boundaries — never through Salsa unwinding. + #[error("Request canceled: {0}")] + RequestCanceled(String), + /// The request's view became stale: sources were mutated while it waited /// (LSP `ContentModified`, `-32801`). #[error("Content modified: {0}")] @@ -111,17 +128,23 @@ impl LspError { // Request before initialize completed. LspError::ServerNotInitialized(_) => ErrorCode::ServerNotInitialized, + // Explicit cancellation claimed the response (B2). + LspError::RequestCanceled(_) => ErrorCode::RequestCanceled, + // The request's view became stale under an applied source change. LspError::ContentModified(_) => ErrorCode::ContentModified, // Valid request that cannot be served: unknown project/file, - // same-revision busy timeout, overload. + // same-revision busy timeout, overload, saturated/oversized + // outbound transport. LspError::ProjectRootNotFound(..) | LspError::ProjectNotFound(_) | LspError::FileNotFound(_) | LspError::InvalidPath { .. } | LspError::InvalidVFSPath { .. } | LspError::NoProjectsFound + | LspError::OutboundSaturated + | LspError::OutboundOversized | LspError::RequestFailed(_) => ErrorCode::RequestFailed, // Internal failures: poison, violated invariants, serialization, @@ -304,6 +327,17 @@ pub struct PreparedRun { // state (e.g. in playground_server's WsState), which must be Clone + Send + Sync. #[async_trait] pub trait BexLsp: Send + Sync + notification::BexLspNotification + request::BexLspRequest { + /// Create a connection-scoped LSP dispatcher (I7): it shares the + /// process-owned project registry but owns a fresh position-encoding + /// negotiation (C1) and fresh initialize workspace roots, and writes only + /// through `sender` — the connection's revocable outbound sink. Browser + /// takeover (D2) revokes that sink, so a retained clone of the old + /// session can no longer leak output into the replacement. + fn new_lsp_session( + &self, + sender: Arc, + ) -> Arc; + fn get_bex_for_project( &self, project_root: &crate::fs::FsPath, @@ -444,6 +478,10 @@ mod tests { fn typed_errors_map_to_spec_codes() { use lsp_server::ErrorCode; + assert_eq!( + code(&LspError::RequestCanceled("client canceled".into())), + ErrorCode::RequestCanceled as i32, + ); assert_eq!( code(&LspError::ContentModified("edit raced".into())), ErrorCode::ContentModified as i32, @@ -483,6 +521,8 @@ mod tests { LspError::ProjectNotFound(root.clone()), LspError::FileNotFound(root), LspError::NoProjectsFound, + LspError::OutboundSaturated, + LspError::OutboundOversized, ] { assert_eq!(code(&e), ErrorCode::RequestFailed as i32, "{e}"); } @@ -511,11 +551,14 @@ mod tests { message: "m".into(), }, LspError::NoProjectsFound, + LspError::RequestCanceled("m".into()), LspError::ContentModified("m".into()), LspError::RequestFailed("m".into()), LspError::Internal("m".into()), LspError::InvalidParams("m".into()), LspError::ServerNotInitialized("m".into()), + LspError::OutboundSaturated, + LspError::OutboundOversized, ]; for e in &all { assert_ne!(code(e), LEGACY_UNKNOWN, "{e} must not emit -32001"); diff --git a/baml_language/crates/bex_project/src/bex_lsp/multi_project/mod.rs b/baml_language/crates/bex_project/src/bex_lsp/multi_project/mod.rs index 11d1df3535a..6c4bce45e8a 100644 --- a/baml_language/crates/bex_project/src/bex_lsp/multi_project/mod.rs +++ b/baml_language/crates/bex_project/src/bex_lsp/multi_project/mod.rs @@ -240,10 +240,22 @@ pub(super) fn db_read_error_to_lsp(e: DbReadError) -> LspError { } /// Bounded request-lane read of a project's source gate. +/// +/// Safe point (B2): right after the gate is acquired — the request may have +/// waited up to the bounded deadline — the ambient dispatch cancellation is +/// checked, so a request whose cancellation already owns the response stops +/// here instead of paying for the database read. The token is observed, not +/// unwound (abort-profile invariant). pub(super) fn read_for_request(project: &BexProject) -> Result, LspError> { - project + let guard = project .read_source_for_request() - .map_err(db_read_error_to_lsp) + .map_err(db_read_error_to_lsp)?; + if crate::bex_lsp::request_cancellation::current_request_is_cancelled() { + return Err(LspError::RequestCanceled( + "request canceled after acquiring the source gate".to_string(), + )); + } + Ok(guard) } enum ProjectRefreshMode { @@ -279,6 +291,24 @@ impl BexMulitProject { } } + /// Connection-scoped dispatcher (I7): shares the process-owned project + /// registry (and everything hanging off it) but owns a fresh + /// position-encoding negotiation (C1) and fresh initialize workspace + /// roots, and writes only through the connection's revocable `sender`. + /// After browser takeover (D2) revokes that sender, a retained clone of + /// this session fails `send_*` with `ClientClosed` instead of leaking + /// into the replacement session. + fn connection_scoped_lsp_session( + &self, + sender: std::sync::Arc, + ) -> Self { + let mut session = self.clone(); + session.sender = sender; + session.negotiated_encoding = std::sync::Arc::new(OnceLock::new()); + session.workspace_roots = std::sync::Arc::new(std::sync::Mutex::new(Vec::new())); + session + } + /// Encoding for request handlers: requests before `initialize` completes /// are a protocol error and must not freeze a default (C1). fn encoding_for_request(&self) -> Result { @@ -1667,6 +1697,13 @@ fn ensure_source_belongs_to_project( #[async_trait::async_trait] impl super::BexLsp for BexMulitProject { + fn new_lsp_session( + &self, + sender: Arc, + ) -> Arc { + Arc::new(self.connection_scoped_lsp_session(sender)) + } + fn get_bex_for_project( &self, project_root: &crate::fs::FsPath, @@ -2234,4 +2271,103 @@ mod tests { names.sort_unstable(); assert_eq!(names, vec!["/manifest_proj", "/proj"]); } + + /// The B2 safe point: a request whose cancellation already claimed the + /// response stops right after acquiring the source gate with a typed + /// `RequestCanceled` instead of paying for the database read. + #[test] + fn read_for_request_cancels_at_the_source_gate_safe_point() { + let ws = TempWorkspace::new("cancel_safe_point"); + ws.file("proj/baml_src/main.baml", "// main"); + let root = ws.vfs_path("proj"); + let project = crate::project::BexProject::new( + &root, + std::sync::Arc::new(sys_ops::SysOpsBuilder::new().build()), + ); + + let token = sys_types::CancellationToken::new(); + let _scope = crate::bex_lsp::request_cancellation::RequestCancellationScope::enter(Some( + token.clone(), + )); + assert!(read_for_request(&project).is_ok()); + + token.cancel(); + assert!(matches!( + read_for_request(&project), + Err(LspError::RequestCanceled(_)) + )); + } + + struct RecordingSender { + notifications: std::sync::Mutex>, + } + + impl LspClientSenderTrait for RecordingSender { + fn send_notification(&self, msg: lsp_server::Notification) -> Result<(), LspError> { + self.notifications.lock().unwrap().push(msg); + Ok(()) + } + + fn send_response_impl(&self, _msg: lsp_server::Response) -> Result<(), LspError> { + Ok(()) + } + + fn make_request(&self, _msg: lsp_server::Request) -> Result<(), LspError> { + Ok(()) + } + } + + /// A connection-scoped session shares the project registry but owns its + /// own encoding negotiation and workspace roots (I7/C1) and routes output + /// through its own sender. + #[test] + fn connection_scoped_session_is_fresh_but_shares_projects() { + use crate::bex_lsp::notification::BexLspNotification as _; + + let sender = Arc::new(RecordingSender { + notifications: std::sync::Mutex::new(Vec::new()), + }); + let root = BexMulitProject::new( + Arc::new(|_: &vfs::VfsPath| Arc::new(sys_ops::SysOpsBuilder::new().build())), + sender.clone(), + Arc::new(NoopPlaygroundSender), + crate::fs::BamlVFS::new(Arc::new(Box::new(vfs::PhysicalFS::new("/")))), + BackgroundSpawner::new(), + ); + let _ = root + .negotiated_encoding + .set(crate::bex_lsp::position_codec::PositionEncoding::UTF8); + root.workspace_roots + .lock() + .unwrap() + .push(vfs::VfsPath::new(vfs::MemoryFS::new())); + + let session_sender = Arc::new(RecordingSender { + notifications: std::sync::Mutex::new(Vec::new()), + }); + let session = root.connection_scoped_lsp_session(session_sender.clone()); + + // Fresh negotiation and roots; shared project registry. + assert!(session.negotiated_encoding.get().is_none()); + assert!(session.workspace_roots.lock().unwrap().is_empty()); + assert!(Arc::ptr_eq(&session.projects, &root.projects)); + + // Output routes only through the session's own sender. + session.handle_notification(lsp_server::Notification::new( + "test/unsupported".to_string(), + serde_json::Value::Null, + )); + assert_eq!(session_sender.notifications.lock().unwrap().len(), 1); + assert!(sender.notifications.lock().unwrap().is_empty()); + } + + struct NoopPlaygroundSender; + + impl crate::bex_lsp::PlaygroundSender for NoopPlaygroundSender { + fn send_playground_notification( + &self, + _notification: crate::bex_lsp::PlaygroundNotification, + ) { + } + } } diff --git a/baml_language/crates/bex_project/src/bex_lsp/notification.rs b/baml_language/crates/bex_project/src/bex_lsp/notification.rs index ecbd59923f2..2c4073b14bc 100644 --- a/baml_language/crates/bex_project/src/bex_lsp/notification.rs +++ b/baml_language/crates/bex_project/src/bex_lsp/notification.rs @@ -27,6 +27,17 @@ macro_rules! define_lsp_notification_trait { pub trait BexLspNotification { /// Dispatch an incoming notification to the appropriate handler method. fn handle_notification(&self, notif: lsp_server::Notification) { + let _ = self.handle_notification_with_status(notif); + } + + /// Dispatch and report whether the handler accepted the + /// notification. Transport ownership/persistence bookkeeping + /// (session document maps, save write-through hooks) must + /// happen only after this returns `true`. + fn handle_notification_with_status( + &self, + notif: lsp_server::Notification, + ) -> bool { let is_log_message = notif.method.as_str() == "window/logMessage"; let result = match notif.method.as_str() { $( @@ -40,7 +51,7 @@ macro_rules! define_lsp_notification_trait { other => Err(LspError::NotificationNotSupported(other.to_string())), }; match result { - Ok(()) => (), + Ok(()) => true, Err(err) => { if !is_log_message { let _ = self.send_notification_log_message(lsp_types::LogMessageParams { @@ -48,6 +59,7 @@ macro_rules! define_lsp_notification_trait { message: err.to_string(), }); } + false } } } diff --git a/baml_language/crates/bex_project/src/bex_lsp/request.rs b/baml_language/crates/bex_project/src/bex_lsp/request.rs index 5f939bd255b..47940119189 100644 --- a/baml_language/crates/bex_project/src/bex_lsp/request.rs +++ b/baml_language/crates/bex_project/src/bex_lsp/request.rs @@ -33,14 +33,67 @@ macro_rules! define_lsp_request_trait { pub trait BexLspRequest { fn handle_request(&self, notif: lsp_server::Request) { let sender = self.request_sender(); + let mut responder = move |id, result| { + let _ = sender(id, result); + }; + self.handle_request_with_responder(notif, &mut responder); + } + + /// Dispatch with a connection-owned responder. Transport + /// runtimes use this to race handler completion against + /// cancellation and to route reused request IDs through a + /// session-scoped response token instead of the shared + /// sender (I7). + fn handle_request_with_responder( + &self, + notif: lsp_server::Request, + responder: &mut dyn FnMut( + lsp_server::RequestId, + Result, + ), + ) { + self.handle_request_with_cancellation(notif, None, responder); + } + + /// Dispatch with an operation-owned cancellation token (B2). + /// The token is observed only at safe handler boundaries — + /// entry, source-gate acquisition (see + /// `multi_project::read_for_request`), and completion — and + /// is never connected to Salsa's unwind-based cancellation + /// (abort-profile invariant). + fn handle_request_with_cancellation( + &self, + notif: lsp_server::Request, + cancellation: Option<&sys_types::CancellationToken>, + responder: &mut dyn FnMut( + lsp_server::RequestId, + Result, + ), + ) { let id = notif.id.clone(); + if cancellation.is_some_and(sys_types::CancellationToken::is_cancelled) { + responder( + id, + Err(LspError::RequestCanceled( + "request canceled before handler entry".to_string(), + )), + ); + return; + } + // Install the ambient token so the bounded source-gate + // read can observe it right after acquisition without + // threading it through every handler signature. + let _cancellation_scope = + crate::bex_lsp::request_cancellation::RequestCancellationScope::enter( + cancellation.cloned(), + ); match notif.method.as_str() { $( lsp_request!($lsp_method) => { let (id, params) = match lsp_request_extract!(notif, $lsp_method) { Ok(extracted) => extracted, Err(err) => { - let _ = sender(id, Err(err)); + responder(id, Err(err)); return; } }; @@ -48,11 +101,20 @@ macro_rules! define_lsp_request_trait { Ok(result) => serde_json::to_value(result).map_err(LspError::RequestSerializeError), Err(err) => Err(err), }; - let _ = sender(id, result); + if cancellation.is_some_and(sys_types::CancellationToken::is_cancelled) { + responder( + id, + Err(LspError::RequestCanceled( + "request canceled during handler execution".to_string(), + )), + ); + return; + } + responder(id, result); } ),*, other => { - let _ = sender(notif.id, Err(LspError::RequestNotSupported(other.to_string()))); + responder(notif.id, Err(LspError::RequestNotSupported(other.to_string()))); return; } } @@ -146,3 +208,92 @@ define_lsp_request_trait! { "typeHierarchy/supertypes" => type_hierarchy_supertypes, "workspaceSymbol/resolve" => workspace_symbol_resolve, } + +#[cfg(test)] +mod tests { + use super::*; + + /// Minimal handler: `shutdown` succeeds and (optionally) cancels its own + /// token mid-handler, modelling a cancel that lands while the handler + /// runs on the dispatch worker. + struct ShutdownHandler { + cancel_during_handler: Option, + } + + impl BexLspRequest for ShutdownHandler { + fn request_sender( + &self, + ) -> Box< + dyn Fn( + lsp_server::RequestId, + Result, + ) -> Result<(), LspError> + + '_, + > { + Box::new(|_, _| Ok(())) + } + + fn on_request_shutdown(&self, _params: ()) -> Result<(), LspError> { + if let Some(token) = &self.cancel_during_handler { + token.cancel(); + } + Ok(()) + } + } + + fn shutdown_request() -> lsp_server::Request { + lsp_server::Request { + id: lsp_server::RequestId::from(1), + method: "shutdown".to_string(), + params: serde_json::Value::Null, + } + } + + fn dispatch( + handler: &ShutdownHandler, + token: Option<&sys_types::CancellationToken>, + ) -> Result { + let mut observed = None; + handler.handle_request_with_cancellation(shutdown_request(), token, &mut |_, result| { + observed = Some(result); + }); + observed.expect("dispatch must respond exactly once") + } + + #[test] + fn uncancelled_dispatch_returns_the_handler_result() { + let handler = ShutdownHandler { + cancel_during_handler: None, + }; + let token = sys_types::CancellationToken::new(); + assert!(matches!( + dispatch(&handler, Some(&token)), + Ok(serde_json::Value::Null) + )); + } + + #[test] + fn cancellation_before_entry_short_circuits_the_handler() { + let handler = ShutdownHandler { + cancel_during_handler: None, + }; + let token = sys_types::CancellationToken::new(); + token.cancel(); + assert!(matches!( + dispatch(&handler, Some(&token)), + Err(LspError::RequestCanceled(_)) + )); + } + + #[test] + fn cancellation_during_the_handler_is_observed_at_completion() { + let token = sys_types::CancellationToken::new(); + let handler = ShutdownHandler { + cancel_during_handler: Some(token.clone()), + }; + assert!(matches!( + dispatch(&handler, Some(&token)), + Err(LspError::RequestCanceled(_)) + )); + } +} diff --git a/baml_language/crates/bex_project/src/bex_lsp/request_cancellation.rs b/baml_language/crates/bex_project/src/bex_lsp/request_cancellation.rs new file mode 100644 index 00000000000..0fde61eb2a2 --- /dev/null +++ b/baml_language/crates/bex_project/src/bex_lsp/request_cancellation.rs @@ -0,0 +1,80 @@ +//! Ambient, dispatch-scoped request cancellation (design §4 Track B, B2). +//! +//! The ingress runtime dispatches one request at a time per worker thread and +//! hands the operation-owned [`sys_types::CancellationToken`] to +//! `handle_request_with_cancellation`. Handler signatures stay stable: the +//! token is installed thread-locally for the duration of exactly one dispatch +//! and observed only at *safe points* — handler entry, immediately after the +//! bounded source-gate acquisition, and handler completion. +//! +//! Under the abort-profile invariant this token is never connected to Salsa's +//! unwind-based local cancellation: observing it returns a typed +//! `LspError::RequestCanceled`, it does not unwind an in-flight query. + +use std::cell::RefCell; + +thread_local! { + static CURRENT_REQUEST_CANCELLATION: RefCell> = + const { RefCell::new(None) }; +} + +/// Installs a token as the current dispatch's cancellation for this thread. +/// The previous value is restored on drop, so a nested dispatch (e.g. +/// `handle_request` delegating through `handle_request_with_cancellation`) +/// stays balanced. +pub(crate) struct RequestCancellationScope { + previous: Option, +} + +impl RequestCancellationScope { + pub(crate) fn enter(token: Option) -> Self { + let previous = CURRENT_REQUEST_CANCELLATION.with(|current| current.replace(token)); + Self { previous } + } +} + +impl Drop for RequestCancellationScope { + fn drop(&mut self) { + CURRENT_REQUEST_CANCELLATION.with(|current| { + *current.borrow_mut() = self.previous.take(); + }); + } +} + +/// Safe-point check for the ambient dispatch token. `false` when no token is +/// installed (plain `handle_request`, notifications, WASM single-thread path). +pub(crate) fn current_request_is_cancelled() -> bool { + CURRENT_REQUEST_CANCELLATION.with(|current| { + current + .borrow() + .as_ref() + .is_some_and(sys_types::CancellationToken::is_cancelled) + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn scope_installs_observes_and_restores_the_ambient_token() { + assert!(!current_request_is_cancelled()); + let outer = sys_types::CancellationToken::new(); + let scope = RequestCancellationScope::enter(Some(outer.clone())); + assert!(!current_request_is_cancelled()); + + { + let inner = sys_types::CancellationToken::new(); + inner.cancel(); + let _nested = RequestCancellationScope::enter(Some(inner)); + assert!(current_request_is_cancelled()); + } + + // The outer (uncancelled) token is restored after the nested scope. + assert!(!current_request_is_cancelled()); + outer.cancel(); + assert!(current_request_is_cancelled()); + drop(scope); + assert!(!current_request_is_cancelled()); + } +} diff --git a/typescript2/app-promptfiddle/src/playground/baml-lsp-worker.ts b/typescript2/app-promptfiddle/src/playground/baml-lsp-worker.ts index 0d5dc1e9a0e..76dc5b03a98 100644 --- a/typescript2/app-promptfiddle/src/playground/baml-lsp-worker.ts +++ b/typescript2/app-promptfiddle/src/playground/baml-lsp-worker.ts @@ -1024,6 +1024,12 @@ self.onmessage = async (event: MessageEvent) => { case "selectProject": return; + // Runtime-demand leases are only meaningful for the remote (WebSocket) + // transport; the in-worker runtime is always resident. + case "ensureProjectRuntime": + case "releaseProjectRuntime": + return; + case "requestState": runtime?.requestPlaygroundState(); postOut({ type: "buildTime", value: getBuildTime() }); diff --git a/typescript2/app-vscode-ext/src/__tests__/projectRoots.test.ts b/typescript2/app-vscode-ext/src/__tests__/projectRoots.test.ts new file mode 100644 index 00000000000..c13b4c5e38e --- /dev/null +++ b/typescript2/app-vscode-ext/src/__tests__/projectRoots.test.ts @@ -0,0 +1,184 @@ +import { + mkdirSync, + mkdtempSync, + realpathSync, + rmSync, + symlinkSync, + writeFileSync, +} from 'fs'; +import { tmpdir } from 'os'; +import * as path from 'path'; +import { afterEach, beforeEach, describe, expect, it } from 'vitest'; +import { + canonicalAncestorDirectories, + canonicalPathIdentity, + pathIdentity, + resolveBamlProjectRoots, + resolveOwnershipRoot, + routableOwnershipPattern, + routableOwnershipRoot, + resolveSemanticProjectRoot, +} from '../projectRoots'; + +describe('BAML project root resolution', () => { + let sandbox: string; + + beforeEach(() => { + sandbox = mkdtempSync(path.join(tmpdir(), 'baml-project-roots-')); + }); + + afterEach(() => { + rmSync(sandbox, { recursive: true, force: true }); + }); + + it('uses the nearest marker for semantics and the outermost marker for ownership', () => { + const outer = path.join(sandbox, 'outer'); + const nested = path.join(outer, 'packages', 'nested'); + const outerModel = path.join(outer, 'outer.baml'); + const model = path.join(nested, 'baml_src', 'models', 'main.baml'); + mkdirSync(path.dirname(model), { recursive: true }); + writeFileSync(path.join(outer, 'baml.toml'), ''); + writeFileSync(outerModel, 'class Outer {}'); + writeFileSync(model, 'class Main {}'); + + const roots = resolveBamlProjectRoots(model, 'file'); + + expect(roots.semanticRoot?.fsPath).toBe(realpathSync(nested)); + expect(roots.ownershipRoot?.fsPath).toBe(realpathSync(outer)); + expect(resolveSemanticProjectRoot(model, 'file')?.key).toBe(roots.semanticRoot?.key); + expect(resolveOwnershipRoot(model, 'file')?.key).toBe(roots.ownershipRoot?.key); + expect(resolveOwnershipRoot(outerModel, 'file')?.key).toBe(roots.ownershipRoot?.key); + expect(resolveOwnershipRoot(nested, 'directory')?.key).toBe(roots.ownershipRoot?.key); + }); + + it('keeps sibling top-level projects in separate ownership domains', () => { + const leftFile = path.join(sandbox, 'left', 'baml_src', 'left.baml'); + const rightFile = path.join(sandbox, 'right', 'baml_src', 'right.baml'); + mkdirSync(path.dirname(leftFile), { recursive: true }); + mkdirSync(path.dirname(rightFile), { recursive: true }); + writeFileSync(leftFile, 'class Left {}'); + writeFileSync(rightFile, 'class Right {}'); + + const leftOwner = resolveOwnershipRoot(leftFile, 'file'); + const rightOwner = resolveOwnershipRoot(rightFile, 'file'); + + expect(leftOwner?.fsPath).toBe(realpathSync(path.join(sandbox, 'left'))); + expect(rightOwner?.fsPath).toBe(realpathSync(path.join(sandbox, 'right'))); + expect(leftOwner?.key).not.toBe(rightOwner?.key); + }); + + it('does not invent an owner for an unmarked standalone file', () => { + const model = path.join(sandbox, 'standalone', 'main.baml'); + mkdirSync(path.dirname(model), { recursive: true }); + writeFileSync(model, 'class Main {}'); + + const roots = resolveBamlProjectRoots(model, 'file'); + + expect(roots.semanticRoot).toBeUndefined(); + expect(roots.ownershipRoot).toBeUndefined(); + }); + + it('requires baml_src to be a directory', () => { + const project = path.join(sandbox, 'project'); + const model = path.join(project, 'main.baml'); + mkdirSync(project, { recursive: true }); + writeFileSync(path.join(project, 'baml_src'), 'not a directory'); + writeFileSync(model, 'class Main {}'); + + expect(resolveOwnershipRoot(model, 'file')).toBeUndefined(); + }); + + it('observes marker additions without a stale resolver cache', () => { + const project = path.join(sandbox, 'project'); + const model = path.join(project, 'main.baml'); + mkdirSync(project, { recursive: true }); + writeFileSync(model, 'class Main {}'); + expect(resolveOwnershipRoot(model, 'file')).toBeUndefined(); + + writeFileSync(path.join(project, 'baml.toml'), ''); + + expect(resolveOwnershipRoot(model, 'file')?.fsPath).toBe(realpathSync(project)); + + rmSync(path.join(project, 'baml.toml')); + + expect(resolveOwnershipRoot(model, 'file')).toBeUndefined(); + }); + + it.skipIf(process.platform === 'win32')( + 'resolves symlinks through the nearest existing ancestor and preserves spaces', + () => { + const project = path.join(sandbox, 'project with spaces'); + const source = path.join(project, 'src'); + const model = path.join(source, 'main.baml'); + const alias = path.join(sandbox, 'alias'); + const sourceAlias = path.join(sandbox, 'source-alias'); + const fileAlias = path.join(sandbox, 'direct-main.baml'); + mkdirSync(source, { recursive: true }); + writeFileSync(path.join(project, 'baml.toml'), ''); + writeFileSync(model, 'class Main {}'); + symlinkSync(project, alias, 'dir'); + symlinkSync(source, sourceAlias, 'dir'); + symlinkSync(model, fileAlias, 'file'); + + const direct = canonicalPathIdentity(model); + const aliased = canonicalPathIdentity(path.join(alias, 'src', '..', 'src', 'main.baml')); + const missingOwner = resolveOwnershipRoot( + path.join(alias, 'generated', 'not-created-yet.baml'), + 'file', + ); + + expect(aliased.key).toBe(direct.key); + expect(aliased.fsPath).toBe(direct.fsPath); + expect(missingOwner?.key).toBe(canonicalPathIdentity(project).key); + + const aliasRoots = resolveBamlProjectRoots( + path.join(alias, 'src', 'main.baml'), + 'file', + ); + expect( + routableOwnershipRoot(path.join(alias, 'src', 'main.baml'), aliasRoots), + ).toBe(alias); + + const sourceAliasRoots = resolveBamlProjectRoots( + path.join(sourceAlias, 'main.baml'), + 'file', + ); + expect( + routableOwnershipRoot( + path.join(sourceAlias, 'main.baml'), + sourceAliasRoots, + ), + ).toBe(sourceAlias); + + const fileAliasRoots = resolveBamlProjectRoots(fileAlias, 'file'); + expect(routableOwnershipPattern(fileAlias, fileAliasRoots)).toEqual({ + basePath: sandbox, + pattern: 'direct-main.baml', + }); + expect( + routableOwnershipPattern( + path.join(sourceAlias, 'main.baml'), + sourceAliasRoots, + ), + ).toEqual({ basePath: sourceAlias, pattern: '**/*.baml' }); + }, + ); + + it('builds a component-by-component ancestor chain through the volume root', () => { + const nested = path.join(sandbox, 'one', 'two'); + mkdirSync(nested, { recursive: true }); + + const ancestors = canonicalAncestorDirectories(nested); + + expect(ancestors[0]?.fsPath).toBe(realpathSync(nested)); + expect(ancestors.at(-1)?.fsPath).toBe(path.parse(nested).root); + expect(new Set(ancestors.map((ancestor) => ancestor.key)).size).toBe(ancestors.length); + }); + + it('folds identity case only for case-insensitive filesystems', () => { + const mixedCase = path.join(path.parse(sandbox).root, 'Users', 'Example', 'Project'); + + expect(pathIdentity(mixedCase, true)).toBe(path.normalize(mixedCase)); + expect(pathIdentity(mixedCase, false)).toBe(path.normalize(mixedCase).toLowerCase()); + }); +}); diff --git a/typescript2/app-vscode-ext/src/extension.ts b/typescript2/app-vscode-ext/src/extension.ts index 4d65e5cfa29..f775096d7a2 100644 --- a/typescript2/app-vscode-ext/src/extension.ts +++ b/typescript2/app-vscode-ext/src/extension.ts @@ -16,12 +16,27 @@ import { isProtocolCompatible, type BamlServerMetadata, } from './compat'; +import { + canonicalPathIdentity, + resolveBamlProjectRoots, + resolveOwnershipRoot, + routableOwnershipPattern, + type BamlProjectRoots, + type CanonicalPath, + type RoutableOwnershipPattern, +} from './projectRoots'; const clients = new Map(); +const clientFileWatchers = new Map(); +const clientRouteSignatures = new Map(); +/** Owners whose most recent start attempt failed. Sticky until a successful + * start so the aggregate status can stay "error" when every start failed. */ +const failedClientStarts = new Set(); const knownProjects = new Map(); let currentServerState: 'starting' | 'running' | 'stopped' | 'error' = 'starting'; let statusBarItem: vscode.StatusBarItem | undefined; let extensionContext: vscode.ExtensionContext | undefined; +let ownershipCoordinator: OwnershipCoordinator | undefined; let playgroundDir: string | undefined; let wrapperPath = 'baml'; @@ -31,8 +46,10 @@ function getExtVersion(): string { /** Short display name: last path component (e.g. "/Users/x/repos/myapp/baml_src" → "myapp/baml_src") */ function projectLabel(fullPath: string): string { - const parts = fullPath.replace(/\/$/, '').split('/'); - return parts.length >= 2 ? parts.slice(-2).join('/') : parts[parts.length - 1] ?? fullPath; + const normalized = path.normalize(fullPath); + const name = path.basename(normalized); + const parent = path.basename(path.dirname(normalized)); + return parent && parent !== name ? `${parent}/${name}` : name || fullPath; } function buildStatusTooltip(serverState: 'starting' | 'running' | 'stopped' | 'error'): vscode.MarkdownString { @@ -107,42 +124,37 @@ function activeDocumentUri(): vscode.Uri | undefined { return vscode.workspace.textDocuments.find((doc) => doc.languageId === 'baml' && doc.uri.scheme === 'file')?.uri; } +/** + * Documents in an unmarked directory tree (no `baml.toml`/`baml_src`) still + * get a server: their workspace folder (or containing directory) acts as the + * ownership root, matching the extension's historical behavior. + */ +function fallbackOwnershipRoot(uri: vscode.Uri): CanonicalPath { + const workspaceRoot = vscode.workspace.getWorkspaceFolder(uri)?.uri.fsPath; + return canonicalPathIdentity(workspaceRoot ?? path.dirname(uri.fsPath)); +} + +function ownershipRootForUri(uri: vscode.Uri): CanonicalPath { + return resolveOwnershipRoot(uri.fsPath, 'file') ?? fallbackOwnershipRoot(uri); +} + function activeClient(): LanguageClient | undefined { const uri = activeDocumentUri(); if (!uri) { return clients.values().next().value; } - const root = findBamlProjectRoot(uri); - return clients.get(root); + const ownerKey = ownershipCoordinator?.ownerKeyForUri(uri) ?? ownershipRootForUri(uri).key; + return clients.get(ownerKey); } -function findBamlProjectRoot(uri: vscode.Uri): string { - let dir = fs.statSync(uri.fsPath).isDirectory() ? uri.fsPath : path.dirname(uri.fsPath); - const workspaceFolder = vscode.workspace.getWorkspaceFolder(uri); - const workspaceRoot = workspaceFolder?.uri.fsPath; - - while (true) { - if (fs.existsSync(path.join(dir, 'baml.toml'))) { - return dir; - } - if (workspaceRoot && dir === workspaceRoot) { - return workspaceRoot; - } - const parent = path.dirname(dir); - if (parent === dir) { - break; - } - dir = parent; - } - - return workspaceRoot ?? path.dirname(uri.fsPath); -} - -function projectKeyForPath(projectPath: string): string { +function projectKeyForPath(projectPath: string): string | undefined { try { - return findBamlProjectRoot(vscode.Uri.file(projectPath)); + return ( + resolveOwnershipRoot(projectPath, 'auto') ?? + fallbackOwnershipRoot(vscode.Uri.file(projectPath)) + ).key; } catch { - return projectPath; + return undefined; } } @@ -203,11 +215,52 @@ function openPlaygroundInBrowserTerminal(projectPath?: string): void { terminal.sendText(command); } -async function ensureClient(projectRoot: string): Promise { - const existing = clients.get(projectRoot); +async function ensureClient( + ownershipRoot: CanonicalPath, + routablePatterns: readonly RoutableOwnershipPattern[], +): Promise { + const ownerKey = ownershipRoot.key; + const projectRoot = ownershipRoot.fsPath; + const canonicalPattern: RoutableOwnershipPattern = { + basePath: projectRoot, + pattern: '**/*.baml', + }; + const patternByKey = new Map(); + for (const route of [canonicalPattern, ...routablePatterns]) { + const normalized = { + basePath: path.normalize(route.basePath), + pattern: route.pattern, + }; + patternByKey.set( + `${normalized.basePath}\u0000${normalized.pattern}`, + normalized, + ); + } + const routePatterns = [...patternByKey.values()].sort((left, right) => + `${left.basePath}\u0000${left.pattern}`.localeCompare( + `${right.basePath}\u0000${right.pattern}`, + ), + ); + const routeSignature = routePatterns + .map((route) => `${route.basePath}\u0000${route.pattern}`) + .join('\u0001'); + let existing = clients.get(ownerKey); + if (existing && clientRouteSignatures.get(ownerKey) !== routeSignature) { + // Keep one process for the canonical owner while refreshing the immutable + // LanguageClient selector to include a newly observed symlink spelling. + await removeClient(ownerKey); + existing = undefined; + } if (existing) { if (existing.state === State.Stopped) { - await existing.start(); + try { + await existing.start(); + failedClientStarts.delete(ownerKey); + } catch (error) { + failedClientStarts.add(ownerKey); + updateStatusBar('error'); + throw error; + } } return existing; } @@ -223,10 +276,41 @@ async function ensureClient(projectRoot: string): Promise { }, }; + const ownerFolder: vscode.WorkspaceFolder = { + uri: vscode.Uri.file(projectRoot), + name: path.basename(projectRoot), + index: 0, + }; + const bamlFiles = routePatterns.map( + ({ basePath, pattern }) => + new vscode.RelativePattern( + basePath === projectRoot ? ownerFolder : basePath, + pattern, + ), + ); + const fileWatchers = [ + ...bamlFiles.map((pattern) => + vscode.workspace.createFileSystemWatcher(pattern), + ), + vscode.workspace.createFileSystemWatcher( + new vscode.RelativePattern(ownerFolder, WORKSPACE_MARKER_GLOB), + ), + ]; + // vscode-languageclient 9 types this as the protocol selector (whose + // pattern is string-only), while its runtime passes the selector directly + // to vscode.languages.match, which supports RelativePattern. + const documentSelector = bamlFiles.map((pattern) => ({ + language: 'baml', + scheme: 'file', + pattern, + })) satisfies vscode.DocumentSelector; const clientOptions: LanguageClientOptions = { - documentSelector: [{ language: 'baml', scheme: 'file' }], + workspaceFolder: ownerFolder, + documentSelector: documentSelector as unknown as LanguageClientOptions['documentSelector'], synchronize: { - fileEvents: vscode.workspace.createFileSystemWatcher('**/*.baml'), + // Marker events must reach retained clients as didChangeWatchedFiles; + // the coordinator's topology watchers separately decide ownership. + fileEvents: fileWatchers, }, initializationOptions: { bamlClient: { @@ -241,18 +325,69 @@ async function ensureClient(projectRoot: string): Promise { }; const client = new LanguageClient( - `baml:${projectRoot}`, + `baml:${ownerKey}`, `BAML Language Server (${projectLabel(projectRoot)})`, serverOptions, clientOptions, ); - clients.set(projectRoot, client); - wireClient(projectRoot, client); - await client.start(); - return client; + clients.set(ownerKey, client); + clientFileWatchers.set(ownerKey, fileWatchers); + clientRouteSignatures.set(ownerKey, routeSignature); + wireClient(ownerKey, client); + try { + await client.start(); + failedClientStarts.delete(ownerKey); + return client; + } catch (error) { + clients.delete(ownerKey); + clientFileWatchers.delete(ownerKey); + clientRouteSignatures.delete(ownerKey); + for (const watcher of fileWatchers) watcher.dispose(); + failedClientStarts.add(ownerKey); + updateStatusBar('error'); + throw error; + } +} + +async function removeClient(ownerKey: string): Promise { + const client = clients.get(ownerKey); + if (!client) return; + + // Remove command routing before waiting for shutdown. A replacement owner is + // only started after this promise resolves. + clients.delete(ownerKey); + knownProjects.delete(ownerKey); + try { + if (client.state !== State.Stopped) { + await client.stop(); + } + } finally { + knownProjects.delete(ownerKey); + for (const watcher of clientFileWatchers.get(ownerKey) ?? []) { + watcher.dispose(); + } + clientFileWatchers.delete(ownerKey); + clientRouteSignatures.delete(ownerKey); + refreshTooltip(); + } } -function wireClient(projectRoot: string, client: LanguageClient) { +function updateAggregateServerState(): void { + const states = Array.from(clients.values(), (client) => client.state); + if (states.some((state) => state === State.Running)) { + updateStatusBar('running'); + } else if (states.some((state) => state === State.Starting)) { + updateStatusBar('starting'); + } else if (failedClientStarts.size > 0) { + // Every start attempt failed: stay in "error" so the user sees the + // failure instead of a neutral "stopped" from the final recompute. + updateStatusBar('error'); + } else { + updateStatusBar('stopped'); + } +} + +function wireClient(ownerKey: string, client: LanguageClient) { if (!extensionContext) { return; } @@ -260,15 +395,15 @@ function wireClient(projectRoot: string, client: LanguageClient) { client.onDidChangeState((e) => { switch (e.newState) { case State.Starting: - updateStatusBar('starting'); + updateAggregateServerState(); break; case State.Running: - updateStatusBar('running'); + updateAggregateServerState(); validateServerCompatibility(client); break; case State.Stopped: - knownProjects.delete(projectRoot); - updateStatusBar('stopped'); + knownProjects.delete(ownerKey); + updateAggregateServerState(); break; } }); @@ -294,7 +429,7 @@ function wireClient(projectRoot: string, client: LanguageClient) { client.onNotification( 'baml/listProjects', (params: { projects: string[] }) => { - knownProjects.set(projectRoot, params.projects ?? []); + knownProjects.set(ownerKey, params.projects ?? []); refreshTooltip(); }, ); @@ -313,6 +448,307 @@ function validateServerCompatibility(client: LanguageClient) { } } +interface TrackedDocument { + uri: vscode.Uri; + resolution: BamlProjectRoots | undefined; + ancestorKeys: Set; +} + +interface RefCountedTopologyWatcher { + refCount: number; + watcher: vscode.FileSystemWatcher; + subscriptions: vscode.Disposable[]; +} + +interface WorkspaceTopologyWatcher { + watcher: vscode.FileSystemWatcher; + subscriptions: vscode.Disposable[]; +} + +const EXACT_MARKER_GLOB = '{baml.toml,baml_src}'; +const WORKSPACE_MARKER_GLOB = '**/{baml.toml,baml_src}'; + +/** + * Owns document-to-client routing and marker topology. Language clients only + * watch BAML files inside their non-overlapping owner; marker changes are + * observed once here and migrated in a serialized lane. + * + * Clients are NOT stopped when their last document closes — a running server + * may still serve a playground webview. A client is only removed when marker + * topology reassigns its documents to a different owner (or on shutdown). + */ +class OwnershipCoordinator implements vscode.Disposable { + private readonly documents = new Map(); + private readonly documentOwners = new Map(); + private readonly ancestorWatchers = new Map(); + private readonly workspaceWatchers = new Map(); + private queue: Promise = Promise.resolve(); + private disposed = false; + + constructor() { + this.syncWorkspaceWatchers(); + } + + ownerKeyForUri(uri: vscode.Uri): string | undefined { + return this.documentOwners.get(uri.toString()); + } + + trackDocument(uri: vscode.Uri): Promise { + return this.trackDocuments([uri]); + } + + trackDocuments(uris: readonly vscode.Uri[]): Promise { + return this.enqueue(async () => { + for (const uri of uris) { + if (uri.scheme !== 'file') continue; + const documentKey = uri.toString(); + let tracked = this.documents.get(documentKey); + if (tracked) { + tracked.uri = uri; + } else { + tracked = { + uri, + resolution: undefined, + ancestorKeys: new Set(), + }; + this.documents.set(documentKey, tracked); + } + this.refreshDocumentResolution(tracked); + } + await this.reconcileClients(); + }); + } + + untrackDocument(uri: vscode.Uri): Promise { + return this.enqueue(async () => { + const documentKey = uri.toString(); + const tracked = this.documents.get(documentKey); + if (!tracked) return; + + // Routing for the closed document goes away, but its owner's client is + // deliberately retained (see class docs): stopping it here would kill + // the playground webview that server may be serving. + this.documentOwners.delete(documentKey); + for (const ancestorKey of tracked.ancestorKeys) { + this.releaseAncestorWatcher(ancestorKey); + } + this.documents.delete(documentKey); + await this.reconcileClients(); + }); + } + + topologyChanged(): Promise { + return this.enqueue(async () => { + await this.reconcile(); + }); + } + + workspaceFoldersChanged(): Promise { + return this.enqueue(async () => { + this.syncWorkspaceWatchers(); + await this.reconcile(); + }); + } + + private enqueue(operation: () => Promise): Promise { + this.queue = this.queue + .then(async () => { + if (!this.disposed) await operation(); + }) + .catch((error: unknown) => { + console.error('BAML ownership migration failed', error); + updateStatusBar('error'); + }); + return this.queue; + } + + private async reconcile(): Promise { + for (const tracked of this.documents.values()) { + this.refreshDocumentResolution(tracked); + } + await this.reconcileClients(); + } + + private refreshDocumentResolution(tracked: TrackedDocument): void { + const resolution = resolveBamlProjectRoots(tracked.uri.fsPath, 'file'); + const desiredAncestors = new Map( + resolution.ancestors.map((ancestor) => [ancestor.key, ancestor] as const), + ); + + // Install new observers before releasing obsolete ones so symlink topology + // changes cannot leave a gap in marker coverage. + for (const [ancestorKey, ancestor] of desiredAncestors) { + if (!tracked.ancestorKeys.has(ancestorKey)) { + this.retainAncestorWatcher(ancestor); + } + } + for (const ancestorKey of tracked.ancestorKeys) { + if (!desiredAncestors.has(ancestorKey)) { + this.releaseAncestorWatcher(ancestorKey); + } + } + + tracked.resolution = resolution; + tracked.ancestorKeys = new Set(desiredAncestors.keys()); + } + + private async reconcileClients(): Promise { + const desiredRoots = new Map< + string, + { + owner: CanonicalPath; + routablePatterns: Map; + } + >(); + const desiredRoutes = new Map(); + + for (const [documentKey, tracked] of this.documents) { + // Unmarked documents fall back to their workspace folder so a plain + // .baml file still gets language services. + const owner = + tracked.resolution?.ownershipRoot ?? fallbackOwnershipRoot(tracked.uri); + const desired = desiredRoots.get(owner.key) ?? { + owner, + routablePatterns: new Map(), + }; + const route = tracked.resolution + ? routableOwnershipPattern(tracked.uri.fsPath, tracked.resolution) + : undefined; + if (route) { + desired.routablePatterns.set( + `${path.normalize(route.basePath)}\u0000${route.pattern}`, + route, + ); + } + desiredRoots.set(owner.key, desired); + desiredRoutes.set(documentKey, owner.key); + } + + // A client is removed only when ownership genuinely migrated: one of the + // documents it used to own is still tracked but now resolves to a + // different owner (marker topology changed). Owners that merely lost + // their last tracked document keep running — their server may be serving + // a playground webview. + const migratedFrom = new Set(); + for (const [documentKey, newOwnerKey] of desiredRoutes) { + const previousOwner = this.documentOwners.get(documentKey); + if (previousOwner && previousOwner !== newOwnerKey) { + migratedFrom.add(previousOwner); + } + } + + // Routing is detached before any old process is stopped. Since replacement + // processes start only after all obsolete owners stop, a document is never + // selected by both clients during marker migration. + this.documentOwners.clear(); + for (const ownerKey of Array.from(clients.keys())) { + if (!desiredRoots.has(ownerKey) && migratedFrom.has(ownerKey)) { + await removeClient(ownerKey); + } + } + + for (const { owner, routablePatterns } of desiredRoots.values()) { + try { + await ensureClient(owner, [...routablePatterns.values()]); + } catch (error) { + console.error(`Failed to start BAML language server for ${owner.fsPath}`, error); + } + } + + for (const [documentKey, ownerKey] of desiredRoutes) { + if (clients.has(ownerKey)) { + this.documentOwners.set(documentKey, ownerKey); + } + } + updateAggregateServerState(); + } + + private retainAncestorWatcher(directory: CanonicalPath): void { + const existing = this.ancestorWatchers.get(directory.key); + if (existing) { + existing.refCount += 1; + return; + } + + const watcher = vscode.workspace.createFileSystemWatcher( + new vscode.RelativePattern(directory.fsPath, EXACT_MARKER_GLOB), + ); + this.ancestorWatchers.set(directory.key, { + refCount: 1, + watcher, + subscriptions: this.listenForTopologyChanges(watcher), + }); + } + + private releaseAncestorWatcher(directoryKey: string): void { + const entry = this.ancestorWatchers.get(directoryKey); + if (!entry) return; + entry.refCount -= 1; + if (entry.refCount > 0) return; + + this.disposeWatcher(entry); + this.ancestorWatchers.delete(directoryKey); + } + + private syncWorkspaceWatchers(): void { + const folders = vscode.workspace.workspaceFolders ?? []; + const desired = new Map(folders.map((folder) => [folder.uri.toString(), folder] as const)); + + for (const [folderKey, entry] of this.workspaceWatchers) { + if (!desired.has(folderKey)) { + this.disposeWatcher(entry); + this.workspaceWatchers.delete(folderKey); + } + } + + for (const [folderKey, folder] of desired) { + if (this.workspaceWatchers.has(folderKey)) continue; + const watcher = vscode.workspace.createFileSystemWatcher( + new vscode.RelativePattern(folder, WORKSPACE_MARKER_GLOB), + ); + this.workspaceWatchers.set(folderKey, { + watcher, + subscriptions: this.listenForTopologyChanges(watcher), + }); + } + } + + private listenForTopologyChanges(watcher: vscode.FileSystemWatcher): vscode.Disposable[] { + const onTopologyChange = () => { + void this.topologyChanged(); + }; + return [ + watcher.onDidCreate(onTopologyChange), + watcher.onDidDelete(onTopologyChange), + watcher.onDidChange(onTopologyChange), + ]; + } + + private disposeWatcher(entry: WorkspaceTopologyWatcher): void { + for (const subscription of entry.subscriptions) subscription.dispose(); + entry.watcher.dispose(); + } + + dispose(): void { + if (this.disposed) return; + this.disposed = true; + this.documentOwners.clear(); + for (const entry of this.ancestorWatchers.values()) this.disposeWatcher(entry); + for (const entry of this.workspaceWatchers.values()) this.disposeWatcher(entry); + this.ancestorWatchers.clear(); + this.workspaceWatchers.clear(); + this.documents.clear(); + } + + async shutdown(): Promise { + this.dispose(); + await this.queue; + for (const ownerKey of Array.from(clients.keys())) { + await removeClient(ownerKey); + } + } +} + export async function activate(context: vscode.ExtensionContext) { extensionContext = context; const config = vscode.workspace.getConfiguration('baml'); @@ -325,13 +761,17 @@ export async function activate(context: vscode.ExtensionContext) { statusBarItem.show(); context.subscriptions.push(statusBarItem); + const coordinator = new OwnershipCoordinator(); + ownershipCoordinator = coordinator; + context.subscriptions.push(coordinator); + const startForUri = async (uri: vscode.Uri | undefined) => { if (uri) { - await ensureClient(findBamlProjectRoot(uri)); + await coordinator.trackDocument(uri); } }; context.subscriptions.push(vscode.window.onDidChangeActiveTextEditor((editor) => { - if (editor?.document.languageId === 'baml') { + if (editor?.document.languageId === 'baml' && editor.document.uri.scheme === 'file') { void startForUri(editor.document.uri); } })); @@ -340,6 +780,14 @@ export async function activate(context: vscode.ExtensionContext) { void startForUri(document.uri); } })); + context.subscriptions.push(vscode.workspace.onDidCloseTextDocument((document) => { + if (document.languageId === 'baml' && document.uri.scheme === 'file') { + void coordinator.untrackDocument(document.uri); + } + })); + context.subscriptions.push(vscode.workspace.onDidChangeWorkspaceFolders(() => { + void coordinator.workspaceFoldersChanged(); + })); // ── Commands ──────────────────────────────────────────────────────── @@ -384,8 +832,9 @@ export async function activate(context: vscode.ExtensionContext) { // NativePlaygroundSender can decide how to open it (and attach the port). context.subscriptions.push( vscode.commands.registerCommand('baml.openPlayground', async (projectPath?: string) => { + const projectOwnerKey = projectPath ? projectKeyForPath(projectPath) : undefined; const client = projectPath - ? clients.get(projectKeyForPath(projectPath)) ?? activeClient() + ? (projectOwnerKey !== undefined ? clients.get(projectOwnerKey) : undefined) ?? activeClient() : activeClient(); if (!client || client.state !== State.Running) { vscode.window.showWarningMessage('BAML Language Server is not running.'); @@ -408,11 +857,21 @@ export async function activate(context: vscode.ExtensionContext) { }), ); - await startForUri(activeDocumentUri()); + const openBamlDocuments = vscode.workspace.textDocuments.filter( + (document) => document.languageId === 'baml' && document.uri.scheme === 'file', + ); + await coordinator.trackDocuments(openBamlDocuments.map((document) => document.uri)); } export async function deactivate() { - for (const client of clients.values()) { - await client.stop(); + const coordinator = ownershipCoordinator; + ownershipCoordinator = undefined; + if (coordinator) { + await coordinator.shutdown(); + } else { + for (const ownerKey of Array.from(clients.keys())) { + await removeClient(ownerKey); + } } + extensionContext = undefined; } diff --git a/typescript2/app-vscode-ext/src/projectRoots.ts b/typescript2/app-vscode-ext/src/projectRoots.ts new file mode 100644 index 00000000000..fe9b1c72228 --- /dev/null +++ b/typescript2/app-vscode-ext/src/projectRoots.ts @@ -0,0 +1,311 @@ +import * as fs from 'fs'; +import * as path from 'path'; + +export interface CanonicalPath { + /** Resolved path used for filesystem access. */ + fsPath: string; + /** Filesystem-aware identity used as a map key. */ + key: string; +} + +export interface BamlProjectRoots { + canonicalInput: CanonicalPath; + ancestors: CanonicalPath[]; + semanticRoot: CanonicalPath | undefined; + ownershipRoot: CanonicalPath | undefined; +} + +export interface RoutableOwnershipPattern { + basePath: string; + pattern: string; +} + +export type InputPathKind = 'file' | 'directory' | 'auto'; + +function toggleAsciiCase(value: string): string | undefined { + for (let index = 0; index < value.length; index += 1) { + const code = value.charCodeAt(index); + if (code >= 65 && code <= 90) { + return `${value.slice(0, index)}${value[index]!.toLowerCase()}${value.slice(index + 1)}`; + } + if (code >= 97 && code <= 122) { + return `${value.slice(0, index)}${value[index]!.toUpperCase()}${value.slice(index + 1)}`; + } + } + return undefined; +} + +function sameFile(left: fs.Stats, right: fs.Stats): boolean { + return left.dev === right.dev && left.ino === right.ino; +} + +/** + * Determine case behavior without writing a probe file into the user's project. + * On case-insensitive filesystems, changing the case of an existing component + * resolves to the same inode. + */ +function isCaseSensitive(existingPath: string): boolean { + let current = existingPath; + while (true) { + const parent = path.dirname(current); + if (parent === current) { + return process.platform !== 'win32'; + } + + const alternateName = toggleAsciiCase(path.basename(current)); + if (alternateName !== undefined) { + const alternate = path.join(parent, alternateName); + try { + return !sameFile(fs.statSync(current), fs.statSync(alternate)); + } catch (error) { + if ((error as NodeJS.ErrnoException).code === 'ENOENT') { + return true; + } + } + } + + current = parent; + } +} + +export function pathIdentity(resolvedPath: string, caseSensitive: boolean): string { + const normalized = path.normalize(resolvedPath); + return caseSensitive ? normalized : normalized.toLowerCase(); +} + +function nearestExistingAncestor(absolutePath: string): { + existingPath: string | undefined; + missingComponents: string[]; +} { + let candidate = absolutePath; + const missingComponents: string[] = []; + + while (!fs.existsSync(candidate)) { + const parent = path.dirname(candidate); + if (parent === candidate) { + return { existingPath: undefined, missingComponents }; + } + missingComponents.unshift(path.basename(candidate)); + candidate = parent; + } + + return { existingPath: candidate, missingComponents }; +} + +/** + * Produce one path identity for routing. Callers should pass `Uri.fsPath`, which + * is already URI-decoded by VS Code. Symlinks are resolved through the nearest + * existing ancestor so newly-created documents receive the same identity as + * existing documents beneath the same alias. + */ +export function canonicalPathIdentity(inputPath: string): CanonicalPath { + const absolutePath = path.resolve(inputPath); + const { existingPath, missingComponents } = nearestExistingAncestor(absolutePath); + + if (existingPath === undefined) { + const normalized = path.normalize(absolutePath); + return { + fsPath: normalized, + key: pathIdentity(normalized, process.platform !== 'win32'), + }; + } + + let realExistingPath: string; + try { + realExistingPath = fs.realpathSync.native(existingPath); + } catch { + realExistingPath = fs.realpathSync(existingPath); + } + + const resolvedPath = path.normalize(path.join(realExistingPath, ...missingComponents)); + return { + fsPath: resolvedPath, + key: pathIdentity(resolvedPath, isCaseSensitive(realExistingPath)), + }; +} + +function startDirectory(input: CanonicalPath, kind: InputPathKind): CanonicalPath { + if (kind === 'directory') { + return input; + } + if (kind === 'file') { + return canonicalPathIdentity(path.dirname(input.fsPath)); + } + + try { + if (fs.statSync(input.fsPath).isDirectory()) { + return input; + } + } catch { + // Missing command paths are file-like. Their nearest existing ancestor was + // still canonicalized above. + } + return canonicalPathIdentity(path.dirname(input.fsPath)); +} + +export function canonicalAncestorDirectories(directoryPath: string): CanonicalPath[] { + const result: CanonicalPath[] = []; + let current = canonicalPathIdentity(directoryPath); + + while (true) { + result.push(current); + const parentPath = path.dirname(current.fsPath); + if (parentPath === current.fsPath) { + break; + } + current = canonicalPathIdentity(parentPath); + } + + return result; +} + +function isFile(filePath: string): boolean { + try { + return fs.statSync(filePath).isFile(); + } catch { + return false; + } +} + +function isDirectory(directoryPath: string): boolean { + try { + return fs.statSync(directoryPath).isDirectory(); + } catch { + return false; + } +} + +export function hasBamlProjectMarker(directoryPath: string): boolean { + return isFile(path.join(directoryPath, 'baml.toml')) + || isDirectory(path.join(directoryPath, 'baml_src')); +} + +export function resolveBamlProjectRoots( + inputPath: string, + kind: InputPathKind = 'auto', +): BamlProjectRoots { + const canonicalInput = canonicalPathIdentity(inputPath); + const directory = startDirectory(canonicalInput, kind); + const ancestors = canonicalAncestorDirectories(directory.fsPath); + let semanticRoot: CanonicalPath | undefined; + let ownershipRoot: CanonicalPath | undefined; + + for (const ancestor of ancestors) { + if (!hasBamlProjectMarker(ancestor.fsPath)) { + continue; + } + semanticRoot ??= ancestor; + ownershipRoot = ancestor; + } + + return { + canonicalInput, + ancestors, + semanticRoot, + ownershipRoot, + }; +} + +export function resolveSemanticProjectRoot( + inputPath: string, + kind: InputPathKind = 'auto', +): CanonicalPath | undefined { + return resolveBamlProjectRoots(inputPath, kind).semanticRoot; +} + +export function resolveOwnershipRoot( + inputPath: string, + kind: InputPathKind = 'auto', +): CanonicalPath | undefined { + return resolveBamlProjectRoots(inputPath, kind).ownershipRoot; +} + +/** + * Recover the path spelling of an ownership root from the path spelling of a + * document. Ownership identity stays canonical, but VS Code's + * `RelativePattern` matches URI paths and therefore also needs the symlink + * alias through which the document was opened. + */ +export function routableOwnershipRoot( + inputPath: string, + roots: BamlProjectRoots, +): string | undefined { + return routableOwnershipMatch(inputPath, roots)?.basePath; +} + +function routableOwnershipMatch( + inputPath: string, + roots: BamlProjectRoots, +): { basePath: string; matchedSuffixComponents: number } | undefined { + const owner = roots.ownershipRoot; + if (!owner) return undefined; + + const relative = path.relative(owner.fsPath, roots.canonicalInput.fsPath); + if ( + relative === '' || + path.isAbsolute(relative) || + relative === '..' || + relative.startsWith(`..${path.sep}`) + ) { + return { basePath: owner.fsPath, matchedSuffixComponents: 1 }; + } + + const relativeComponents = relative.split(path.sep).filter(Boolean); + let candidate = path.resolve(inputPath); + let matchedSuffixComponents = 0; + for ( + let index = relativeComponents.length - 1; + index >= 0 && + path.basename(candidate) === relativeComponents[index]; + index -= 1 + ) { + candidate = path.dirname(candidate); + matchedSuffixComponents += 1; + } + + // A symlink may enter below the canonical owner (for example `alias-src` + // points at `project/src`). In that case only the shared suffix is safe as a + // selector base; ascending the full canonical relative path would broaden + // the selector above the alias and could overlap another owner's client. + return { + basePath: + matchedSuffixComponents > 0 ? candidate : path.dirname(inputPath), + matchedSuffixComponents, + }; +} + +function exactGlobSegment(value: string): string { + return value.replace(/[\[\]*?{}]/g, (character) => `[${character}]`); +} + +/** + * Build the narrowest RelativePattern needed for the URI spelling that was + * opened. Directory aliases can safely own their subtree. A direct file + * symlink (or an alias with no canonical suffix match) must use an exact file + * pattern; broadening to its parent would overlap sibling project clients. + */ +export function routableOwnershipPattern( + inputPath: string, + roots: BamlProjectRoots, +): RoutableOwnershipPattern | undefined { + const match = routableOwnershipMatch(inputPath, roots); + if (!match) return undefined; + + let directFileSymlink = false; + try { + directFileSymlink = + fs.lstatSync(path.resolve(inputPath)).isSymbolicLink() && + fs.statSync(path.resolve(inputPath)).isFile(); + } catch { + // Missing documents are covered by the suffix-based fallback below. + } + + if (directFileSymlink || match.matchedSuffixComponents === 0) { + return { + basePath: path.dirname(inputPath), + pattern: exactGlobSegment(path.basename(inputPath)), + }; + } + + return { basePath: match.basePath, pattern: '**/*.baml' }; +} diff --git a/typescript2/app-vscode-webview/src/ExecutionPanel.gating.test.tsx b/typescript2/app-vscode-webview/src/ExecutionPanel.gating.test.tsx new file mode 100644 index 00000000000..6ae753ff355 --- /dev/null +++ b/typescript2/app-vscode-webview/src/ExecutionPanel.gating.test.tsx @@ -0,0 +1,163 @@ +import { act, fireEvent, render, screen, waitFor } from '@testing-library/react'; +import { beforeAll, describe, expect, it, vi } from 'vitest'; +import { + ExecutionPanel, + type ProjectUpdate, + type RuntimePort, + type WorkerInMessage, + type WorkerOutMessage, +} from '@b/pkg-playground'; + +beforeAll(() => { + HTMLElement.prototype.scrollTo ??= vi.fn(); +}); + +const PREPARING_TEXT = 'Preparing current build…'; + +function projectUpdate(overrides: Partial = {}): ProjectUpdate { + return { + isBexCurrent: true, + functions: [{ name: 'ReadNote', kind: 'expr', origin: 'userDefined' }], + diagnostics: [], + ...overrides, + }; +} + +function announceProject(port: FakeRuntimePort, update: ProjectUpdate): void { + act(() => { + port.emit({ + type: 'playgroundNotification', + notification: { type: 'listProjects', projects: ['project'] }, + }); + port.emit({ + type: 'playgroundNotification', + notification: { type: 'updateProject', project: 'project', update }, + }); + }); +} + +async function selectReadNote(): Promise { + fireEvent.click(await screen.findByRole('button', { name: 'Functions (1)' })); + fireEvent.click(await screen.findByRole('button', { name: 'ReadNote' })); +} + +describe('ExecutionPanel run gating (fail-closed server)', () => { + it('disables Run and shows the preparing state while the build is stale, keeping the catalog visible', async () => { + const port = new FakeRuntimePort(); + render(); + + announceProject(port, projectUpdate({ isBexCurrent: false })); + await selectReadNote(); + + // The previous function listing stays visible… + expect(screen.getByRole('button', { name: 'ReadNote' })).toBeInTheDocument(); + // …but runtime-derived controls are gated behind the preparing state. + expect(await screen.findByText(PREPARING_TEXT)).toBeInTheDocument(); + expect(screen.getByRole('button', { name: 'Run' })).toBeDisabled(); + + fireEvent.click(screen.getByRole('button', { name: 'Run' })); + expect(port.sent.some((msg) => msg.type === 'startRun')).toBe(false); + + // The next current update re-enables Run automatically. + announceProject(port, projectUpdate({ isBexCurrent: true })); + await waitFor(() => { + expect(screen.getByRole('button', { name: 'Run' })).toBeEnabled(); + }); + expect(screen.queryByText(PREPARING_TEXT)).not.toBeInTheDocument(); + }); + + it('renders a projectNotReady run rejection as the transient preparing state, not a raw error', async () => { + const port = new FakeRuntimePort(); + render(); + + announceProject(port, projectUpdate({ isBexCurrent: true })); + await selectReadNote(); + fireEvent.click(await screen.findByRole('button', { name: 'Run' })); + + const startRun = await waitFor(() => { + const msg = port.sent.find( + (candidate): candidate is Extract => + candidate.type === 'startRun', + ); + expect(msg).toBeDefined(); + return msg!; + }); + + // The fail-closed server refuses the run while a rebuild is pending. + act(() => { + port.emit({ + type: 'commandError', + requestId: startRun.requestId, + code: 'projectNotReady', + message: 'Cannot start run: rebuild pending', + }); + }); + + expect(await screen.findByText(PREPARING_TEXT)).toBeInTheDocument(); + // No raw error text or code leaks into the panel. + expect(screen.queryByText(/projectNotReady/)).not.toBeInTheDocument(); + expect(screen.queryByText(/rebuild pending/)).not.toBeInTheDocument(); + expect(screen.getByRole('button', { name: 'Run' })).toBeDisabled(); + + // The next ProjectUpdate with a current build clears the state… + announceProject(port, projectUpdate({ isBexCurrent: true })); + await waitFor(() => { + expect(screen.queryByText(PREPARING_TEXT)).not.toBeInTheDocument(); + }); + + // …and Run works again. + const sentBefore = port.sent.filter((msg) => msg.type === 'startRun').length; + fireEvent.click(screen.getByRole('button', { name: 'Run' })); + await waitFor(() => { + expect( + port.sent.filter((msg) => msg.type === 'startRun').length, + ).toBeGreaterThan(sentBefore); + }); + }); + + it('shows the diagnostics banner (not the preparing spinner) for compile errors', async () => { + const port = new FakeRuntimePort(); + render(); + + announceProject( + port, + projectUpdate({ + isBexCurrent: false, + diagnostics: [{ severity: 'error', message: 'missing }' }], + }), + ); + await selectReadNote(); + + expect( + await screen.findByText(/1 error — current build unavailable/), + ).toBeInTheDocument(); + expect(screen.queryByText(PREPARING_TEXT)).not.toBeInTheDocument(); + expect(screen.getByRole('button', { name: 'Run' })).toBeDisabled(); + }); +}); + +class FakeRuntimePort implements RuntimePort { + sent: WorkerInMessage[] = []; + private handlers = new Set<(msg: WorkerOutMessage) => void>(); + + postMessage(msg: WorkerInMessage): void { + this.sent.push(msg); + } + + onMessage(handler: (msg: WorkerOutMessage) => void): () => void { + this.handlers.add(handler); + return () => { + this.handlers.delete(handler); + }; + } + + emit(msg: WorkerOutMessage): void { + for (const handler of this.handlers) { + handler(msg); + } + } + + dispose(): void { + this.handlers.clear(); + } +} diff --git a/typescript2/pkg-playground/src/ExecutionPanel.tsx b/typescript2/pkg-playground/src/ExecutionPanel.tsx index f7307caadcb..8b592718e22 100644 --- a/typescript2/pkg-playground/src/ExecutionPanel.tsx +++ b/typescript2/pkg-playground/src/ExecutionPanel.tsx @@ -13,7 +13,14 @@ import type { ChangeEvent, FC, RefObject } from 'react'; import { useCallback, useEffect, useMemo, useRef, useState } from 'react'; import { encodeRunArgs } from '@b/pkg-proto'; import type { BamlJsValue } from '@b/pkg-proto'; -import { KeyRound, PanelLeft, Play, Settings, Square } from 'lucide-react'; +import { + KeyRound, + Loader2, + PanelLeft, + Play, + Settings, + Square, +} from 'lucide-react'; import { Button } from './components/ui/button'; import { Tabs, TabsContent, TabsList, TabsTrigger } from './components/ui/tabs'; import { Input } from './components/ui/input'; @@ -69,7 +76,17 @@ import { GraphView } from './graph/GraphView'; import { FunctionSidebar } from './FunctionSidebar'; import { companionFunctionName } from './shared/companion-functions'; import { createExecutionStore, type ExecutionStore } from './execution-store'; -import { createRunStoreClient } from './run-store-client'; +import { + createRunStoreClient, + isProjectNotReadyError, +} from './run-store-client'; +import { + NO_NOT_READY_PROJECTS, + applyProjectUpdateToGating, + isRunGated, + markProjectNotReady, + type NotReadyProjects, +} from './run-gating'; import { createValueBodyCache } from './value-body-cache'; import type { ValueBodyCache } from './value-body-cache'; import type { ExecutionStoreSnapshot } from './execution-store'; @@ -739,6 +756,13 @@ export const ExecutionPanel: FC = ({ const [projectUpdates, setProjectUpdates] = useState< Record >({}); + // Projects whose last run/preview was refused with `projectNotReady`. The + // fail-closed server (D7) rejects runs while a rebuild is pending; the UI + // renders that as the transient "Preparing current build…" state and clears + // it when the next current ProjectUpdate arrives. + const [notReadyProjects, setNotReadyProjects] = useState( + NO_NOT_READY_PROJECTS, + ); const [testTree, setTestTree] = useState(null); const [collectionCallId, setCollectionCallId] = useState(null); const [generation, setGeneration] = useState(0); @@ -1235,6 +1259,11 @@ export const ExecutionPanel: FC = ({ break; case 'updateProject': setProjectUpdates((prev) => ({ ...prev, [n.project]: n.update })); + // A current build re-enables run controls automatically after a + // fail-closed `projectNotReady` rejection. + setNotReadyProjects((prev) => + applyProjectUpdateToGating(prev, n.project, n.update), + ); break; case 'testCollectionResult': { try { @@ -1551,6 +1580,20 @@ export const ExecutionPanel: FC = ({ setArgsJson(baseArgsFor(selectedFn)); }, [selectedFn, baseArgsFor]); + // Whether the fail-closed server is (re)building the selected project's + // runtime: either the latest ProjectUpdate is stale (isBexCurrent false) or + // a run/preview was refused with `projectNotReady`. Derived here — above + // the run/preview callbacks that must consult it. + const runtimePreparing = isRunGated( + notReadyProjects, + selectedProject, + selectedProject ? projectUpdates[selectedProject] : undefined, + ); + + const markSelectedProjectNotReady = useCallback((project: string) => { + setNotReadyProjects((prev) => markProjectNotReady(prev, project)); + }, []); + // Auto-refresh prompt/curl preview when args change while tab is active useEffect(() => { if (activeTab !== 'prompt' && activeTab !== 'curl') return; @@ -1590,6 +1633,14 @@ export const ExecutionPanel: FC = ({ return; } + // While the server prepares the current build, previews would only be + // refused with `projectNotReady`. Skip issuing them; this effect re-runs + // when the next ProjectUpdate flips `runtimePreparing` back off. + if (runtimePreparing) { + setPreviewLoading(false); + return; + } + setPreviewLoading(true); let cancelled = false; @@ -1659,6 +1710,15 @@ export const ExecutionPanel: FC = ({ setPreviewLoading(false); } catch (e) { if (cancelled) return; + if (isProjectNotReadyError(e)) { + // Transient fail-closed rejection — surface it as the "Preparing + // current build…" state (not a raw error) and keep the last valid + // preview visible. Cleared by the next current ProjectUpdate. + markSelectedProjectNotReady(selectedProject); + setError(null); + setPreviewLoading(false); + return; + } const errMsg = e instanceof Error ? e.message : String(e); // Don't clear result — keep last valid prompt visible with error banner above setError(errMsg); @@ -1681,6 +1741,8 @@ export const ExecutionPanel: FC = ({ executionStore, valueBodyCache, projectUpdateVersion, + runtimePreparing, + markSelectedProjectNotReady, ]); // Single write path for args edits (form and raw): the prompt/cURL preview @@ -1871,6 +1933,13 @@ export const ExecutionPanel: FC = ({ }); await waitForTerminalRun(boundaryId); } catch (e) { + if (isProjectNotReadyError(e)) { + // Fail-closed rebuild window: show the transient preparing state + // instead of a per-test error; controls re-enable on the next + // current ProjectUpdate. + markSelectedProjectNotReady(selectedProject); + return; + } setTestStartErrors( (prev) => new Map(prev).set( @@ -1880,7 +1949,13 @@ export const ExecutionPanel: FC = ({ ); } }, - [executionStore, generation, selectedProject, waitForTerminalRun], + [ + executionStore, + generation, + selectedProject, + waitForTerminalRun, + markSelectedProjectNotReady, + ], ); useEffect(() => { @@ -2006,7 +2081,6 @@ export const ExecutionPanel: FC = ({ ? functions : functions.filter((fn) => !isInternalFunction(fn)); const functionNames = visibleFunctions.map((f) => f.name); - const engineStale = currentUpdate ? !currentUpdate.isBexCurrent : false; const diags = currentUpdate?.diagnostics ?? []; const selectedFnInfo = visibleFunctions.find((f) => f.name === selectedFn); @@ -2159,6 +2233,13 @@ export const ExecutionPanel: FC = ({ [boundaryId]: runArgsJson, })); } catch (e) { + if (isProjectNotReadyError(e)) { + // Fail-closed rebuild window: render the transient "Preparing current + // build…" state instead of a raw error. Run re-enables automatically + // when the next ProjectUpdate reports a current build. + markSelectedProjectNotReady(selectedProject); + return; + } const errMsg = e instanceof Error ? e.message : String(e); setRunValidationError(errMsg); } @@ -2172,6 +2253,7 @@ export const ExecutionPanel: FC = ({ isRunning, executionStore, updateArgsJson, + markSelectedProjectNotReady, ]); // Names of LLM functions — only these have a meaningful raw (un-parsed LLM // output) vs parsed distinction, so the Parsed/Raw toggle is shown only for @@ -2415,6 +2497,10 @@ export const ExecutionPanel: FC = ({ const errors = diags.filter((d) => d.severity === 'error'); const warnings = diags.filter((d) => d.severity === 'warning'); const hasErrors = errors.length > 0; + // The fail-closed server refuses runs/previews over compile errors and + // while a rebuild is pending; keep runtime-derived controls disabled for + // both so users see one consistent gate instead of raw rejections. + const runtimeControlsDisabled = hasErrors || runtimePreparing; // Whether any known-required keys are missing — proactive, not just reactive to pending requests const hasMissingKeys = [...knownRequiredKeys].some((k) => !envVars[k]); @@ -2477,7 +2563,7 @@ export const ExecutionPanel: FC = ({ // keystroke when nothing will run. if (viewingCollection || viewingTestRun) return; e.preventDefault(); - if (!hasErrors) void onRunFunction(); + if (!runtimeControlsDisabled) void onRunFunction(); }} > {/* ──── Combined top bar ──── */} @@ -2574,7 +2660,9 @@ export const ExecutionPanel: FC = ({ size="icon-xs" className="h-7 w-7" aria-label="Run" - disabled={hasErrors || isRunning || !selectedProject} + disabled={ + runtimeControlsDisabled || isRunning || !selectedProject + } onClick={onRunFunction} > @@ -2701,8 +2789,23 @@ export const ExecutionPanel: FC = ({ )} + {/* Preparing banner. The sidebar keeps showing the previous + function/test catalog, but runtime-derived controls stay disabled + until the fail-closed server reports a current build. */} + {selectedProject && runtimePreparing && !hasErrors && ( +
+ + + Preparing current build… + +
+ )} + {/* Diagnostics banner */} - {(hasErrors || engineStale) && ( + {hasErrors && (
{diags.length > 0 ? ( <> @@ -2730,7 +2833,7 @@ export const ExecutionPanel: FC = ({ {warnings.length > 0 ? `${warnings.length} warning${warnings.length !== 1 ? 's' : ''}` : ''} - {' — using last successful build'} + {' — current build unavailable'} {diagsExpanded && ( @@ -2754,11 +2857,7 @@ export const ExecutionPanel: FC = ({
)} - ) : ( -
- Build is stale — using last successful build -
- )} + ) : null} )} @@ -2783,6 +2882,7 @@ export const ExecutionPanel: FC = ({ showInternalFunctions={showInternalFunctions} internalFunctionCount={internalFunctionCount} isLoadingProject={isLoadingProject} + runtimeControlsDisabled={runtimeControlsDisabled} testTree={testTree} previewTests={previewTests} selectedPreviewTestKey={selectedPreviewTestKey} @@ -3213,7 +3313,11 @@ export const ExecutionPanel: FC = ({ size="xs" className="mx-1 my-0.5 shrink-0 text-[11px] font-semibold" aria-label={isRunning ? 'Running' : 'Run'} - disabled={hasErrors || isRunning || !selectedProject} + disabled={ + runtimeControlsDisabled || + isRunning || + !selectedProject + } onClick={onRunFunction} > {isRunning ? ( diff --git a/typescript2/pkg-playground/src/FunctionSidebar.tsx b/typescript2/pkg-playground/src/FunctionSidebar.tsx index 5862f1cbfbf..62ff44799b2 100644 --- a/typescript2/pkg-playground/src/FunctionSidebar.tsx +++ b/typescript2/pkg-playground/src/FunctionSidebar.tsx @@ -40,6 +40,8 @@ import { interface TestTreeNodeProps { def: SerializedTestDef; depth?: number; + /** Disable run/retry actions while the runtime is unavailable. */ + disabled?: boolean; onRunTest?: (name: string) => void; testRunResults?: Map; failedExpands?: Set; @@ -49,6 +51,7 @@ interface TestTreeNodeProps { function TestTreeNode({ def, depth = 0, + disabled = false, onRunTest, testRunResults, failedExpands, @@ -85,7 +88,8 @@ function TestTreeNode({ {isFailed && onRetryExpand && (