Fix LSP latency, mutex contention, and session correctness#3961
Draft
rossirpaulo wants to merge 1 commit into
Draft
Fix LSP latency, mutex contention, and session correctness#3961rossirpaulo wants to merge 1 commit into
rossirpaulo wants to merge 1 commit into
Conversation
|
The latest updates on your projects. Learn more about Vercel for GitHub.
|
Contributor
📝 WalkthroughWalkthroughThe change introduces a shared LSP runtime with bounded transport handling, negotiated source-position encoding, project runtime leasing and fencing, updated playground protocols, canonical VS Code project routing, Salsa query memoization, and testset/profiling lifecycle updates. ChangesLSP runtime and transport
Position encoding and project LSP
Playground runtime lifecycle
VS Code integration
Supporting behavior
Estimated code review effort: 5 (Critical) | ~120 minutes Possibly related PRs
Suggested reviewers: Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
⏭️ Performance benchmarks were skippedPerf benchmarks (CodSpeed) are opt-in on pull requests — they no longer run on every push. They always run automatically after merge to To run them on this PR, do any of the following, then push a commit (or re-run CI):
|
Binary size checks failed❌ 4 violations · ✅ 3 passed
Details & how to fixViolations:
Add/update baselines:
[artifacts.baml-cli]
file_bytes = 17152688
stripped_bytes = 17152736
gzip_bytes = 8294648
[artifacts.bridge_wasm]
file_bytes = 14917601
gzip_bytes = 4235047
[artifacts.baml-cli]
file_bytes = 18760192
stripped_bytes = 18760192
gzip_bytes = 8510483
[artifacts.baml-cli]
file_bytes = 22310784
stripped_bytes = 22310776
gzip_bytes = 9550926Generated by |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 10
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
baml_language/crates/baml_lsp_server/src/playground_server.rs (1)
1462-1494: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick winUpdate the document mirror only after successful delivery.
The mirror is advanced before budget reservation and broadcast. If either fails, the browser misses the edit, while later identical file events are suppressed as echoes.
Proposed fix
{ - let Ok(mut mirror) = doc_mirror.lock() else { + let Ok(mirror) = doc_mirror.lock() else { continue; }; if mirror.get(&canonical).map(String::as_str) == Some(content.as_str()) { continue; } - mirror.insert(canonical.clone(), content.clone()); } ... -let _ = lsp_out_tx.send(frame); +if lsp_out_tx.send(frame).is_err() { + continue; +} +if let Ok(mut mirror) = doc_mirror.lock() { + mirror.insert(canonical.clone(), content.clone()); +}🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@baml_language/crates/baml_lsp_server/src/playground_server.rs` around lines 1462 - 1494, Update the document mirror only after the notification frame has been successfully reserved and sent. In the disk-change handling block, move the existing `mirror.insert` in coordination with `lsp_out_budget.try_message` and `lsp_out_tx.send`, ensuring all saturation, oversized, serialization, and send failures leave the previous mirror value unchanged so a later event can be retried.typescript2/pkg-playground/src/ExecutionPanel.tsx (1)
2517-2562: 🚀 Performance & Scalability | 🟡 Minor | ⚡ Quick winAvoid issuing the initial collection request twice.
The first effect posts after setting
pendingTestTarget; the next effect observes that target with no tree and posts again. Let the centralized effect own collection requests.Proposed fix
setPendingTestTarget({ project: selectedProject, kind: initialTestName ? 'test' : 'testset', name: initialTestName ?? initialTestsetName!, }); - port.postMessage({ type: 'requestCollectTests', project: selectedProject });🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@typescript2/pkg-playground/src/ExecutionPanel.tsx` around lines 2517 - 2562, Remove the direct port.postMessage requestCollectTests call from the initial-target useEffect, leaving it responsible only for initializing state and pendingTestTarget. Let the pendingTestTarget collection useEffect handle the single centralized request after the target is set, while preserving its existing readiness and project guards.
🧹 Nitpick comments (1)
baml_language/crates/bex_engine/tests/collect_tests.rs (1)
172-206: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueConsider adding this as a unit test in
src/lib.rsinstead of an integration test.The new test helpers and
test_registry_expansion_can_be_rolled_back_before_publicationtest exercise the publicengine.call_functionAPI, which is equally accessible from#[cfg(test)]unit tests. The existing integration test structure is consistent, but the coding guideline prefers unit tests where possible for faster feedback and access to internal state.As per coding guidelines: "Prefer writing Rust unit tests over integration tests where possible" and "Always run
cargo test --libif you changed any Rust code". Ensurecargo test --libpasses for the changed Rust crates (bex_events,baml_lsp_server,bex_engine).Also applies to: 751-786
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@baml_language/crates/bex_engine/tests/collect_tests.rs` around lines 172 - 206, Move the rollback expansion test, including the helpers registry_is_expanded and rollback_testset_expansion and test_registry_expansion_can_be_rolled_back_before_publication, from the integration test module into the appropriate #[cfg(test)] unit-test section of src/lib.rs. Preserve coverage through the public engine.call_function API, adjust imports and setup as needed, and run cargo test --lib for bex_events, baml_lsp_server, and bex_engine.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@baml_language/crates/baml_lsp_server/src/lib.rs`:
- Around line 51-61: Update OutboundFrame payload storage so clones share the
underlying Message/Value via Arc instead of deep-cloning payload data while
sharing one OutboundCharge; adjust constructors and all payload accesses,
including the writer match, to use payload.as_ref(), and apply the same change
to the additional affected clone/delivery sites.
- Around line 227-267: The LSP header parsing loop can allocate unbounded memory
before validating Content-Length. In the header-reading logic, replace unbounded
read_line usage with a capped reader or bounded incremental reads, and track
cumulative header bytes against a fixed maximum header-size limit; return
StdioReadError::Framing when the limit is exceeded, including for unterminated
headers, while preserving CRLF and malformed-header validation.
In `@baml_language/crates/baml_lsp_server/src/lsp_runtime.rs`:
- Line 492: The lsp_runtime response branch currently discards all client
responses, leaving requests created by make_request unresolved. Update BexLsp’s
message handling and scheduler dispatch to correlate each
lsp_server::Message::Response with its pending request and deliver the result,
or remove the unsupported server-request/make_request functionality
consistently.
In `@baml_language/crates/bex_project/src/bex_lsp/mod.rs`:
- Around line 302-314: Make prepared-run registration RAII-backed by turning
PreparedFunctionRun and PreparedTestRun into non-Copy lease handles that own the
registered call_id. Implement Drop for the handle to invoke the equivalent of
finish_project_run and release runtime/cancellation state, ensuring cleanup on
normal completion, panic, or task abort. Update preparation and detached-task
call sites to retain and drop the lease instead of manually calling
finish_project_run.
In `@baml_language/crates/bex_project/src/bex_lsp/multi_project/notification.rs`:
- Around line 45-61: Preserve per-session ownership of shared document overlays
in the open-document handling and corresponding change/close paths (including
the regions around the referenced lines). Update the BexProject overlay identity
or ownership tracking so one LSP connection cannot overwrite another session’s
unsaved text or restore disk contents while another session remains open; only
remove/restore an overlay after its final owner closes. Add a test covering two
sessions opening, changing, and closing the same document.
In `@baml_language/crates/bex_project/src/lsp_ingress.rs`:
- Around line 717-721: Update the policy construction in the message-handling
logic to preserve the forced cancellation setting for protected methods. Combine
the forced and requested cancellation policies using the same precedence rule as
class strength, ensuring custom policies cannot downgrade `initialize`,
`shutdown`, or `workspace/executeCommand` from `NonCancellableWhenRunning`; use
the existing `MessagePolicy::for_message` and `stronger_class` logic as
references.
In `@typescript2/app-promptfiddle/src/playground/worker-backend.ts`:
- Around line 187-190: Update the worker initialization flow around
workerReadyPromise to reject when the Worker emits an error and to reject after
a finite timeout if no ready message arrives. Ensure connect propagates either
failure instead of hanging, and clean up the error listener and timeout when
workerReadyPromise resolves or rejects.
In `@typescript2/app-vscode-ext/src/extension.ts`:
- Around line 596-609: Preserve the startup error recorded by ensureClient when
every client fails to start. Update the surrounding startup flow and
updateAggregateServerState so an empty clients map does not overwrite an
existing error state with stopped; only transition to stopped when no startup
error is present.
In `@typescript2/pkg-playground/src/ports/WebSocketRuntimePort.ts`:
- Around line 579-587: Update the server-side
WsOutMessage::ControlFlowGraphResult variant and every producer to serialize
sessionEpoch, project, projectIncarnation, sourceRevision, generation,
derivedEpoch, functionName, and graph; ensure the TypeScript
controlFlowGraphResult handler receives these fields so its sessionEpoch fence
and identity metadata are valid.
- Around line 150-198: Preserve the latest desired runtime lease across
pre-handshake sends and reconnects instead of deleting it in connect() or hello
processing. Update the hello-success path to replay desiredRuntimeLease exactly
once after confirming playground compatibility and runtimeDemandSupported, then
clear or mark it as sent only after transmission. Ensure ensureProjectRuntime
and releaseProjectRuntime continue updating the lease while disconnected,
without entering the generic reconnect queue.
---
Outside diff comments:
In `@baml_language/crates/baml_lsp_server/src/playground_server.rs`:
- Around line 1462-1494: Update the document mirror only after the notification
frame has been successfully reserved and sent. In the disk-change handling
block, move the existing `mirror.insert` in coordination with
`lsp_out_budget.try_message` and `lsp_out_tx.send`, ensuring all saturation,
oversized, serialization, and send failures leave the previous mirror value
unchanged so a later event can be retried.
In `@typescript2/pkg-playground/src/ExecutionPanel.tsx`:
- Around line 2517-2562: Remove the direct port.postMessage requestCollectTests
call from the initial-target useEffect, leaving it responsible only for
initializing state and pendingTestTarget. Let the pendingTestTarget collection
useEffect handle the single centralized request after the target is set, while
preserving its existing readiness and project guards.
---
Nitpick comments:
In `@baml_language/crates/bex_engine/tests/collect_tests.rs`:
- Around line 172-206: Move the rollback expansion test, including the helpers
registry_is_expanded and rollback_testset_expansion and
test_registry_expansion_can_be_rolled_back_before_publication, from the
integration test module into the appropriate #[cfg(test)] unit-test section of
src/lib.rs. Preserve coverage through the public engine.call_function API,
adjust imports and setup as needed, and run cargo test --lib for bex_events,
baml_lsp_server, and bex_engine.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 8340b0e8-a911-4916-a12d-91db384d37fe
⛔ Files ignored due to path filters (2)
baml_language/Cargo.lockis excluded by!**/*.locktypescript2/pnpm-lock.yamlis excluded by!**/pnpm-lock.yaml
📒 Files selected for processing (70)
baml_language/crates/baml_builtins2/baml_std/testing/registry.bamlbaml_language/crates/baml_cli/src/paint.rsbaml_language/crates/baml_lsp2_actions/src/annotations.rsbaml_language/crates/baml_lsp2_actions/src/lib.rsbaml_language/crates/baml_lsp2_actions/src/tokens.rsbaml_language/crates/baml_lsp2_actions_tests/Cargo.tomlbaml_language/crates/baml_lsp2_actions_tests/src/lib.rsbaml_language/crates/baml_lsp2_actions_tests/src/memoization.rsbaml_language/crates/baml_lsp2_actions_tests/src/runner.rsbaml_language/crates/baml_lsp_server/src/lib.rsbaml_language/crates/baml_lsp_server/src/lsp_runtime.rsbaml_language/crates/baml_lsp_server/src/native_lsp_sender.rsbaml_language/crates/baml_lsp_server/src/playground_env.rsbaml_language/crates/baml_lsp_server/src/playground_runs.rsbaml_language/crates/baml_lsp_server/src/playground_sender.rsbaml_language/crates/baml_lsp_server/src/playground_server.rsbaml_language/crates/baml_lsp_server/src/playground_session.rsbaml_language/crates/baml_lsp_server/src/playground_ws.rsbaml_language/crates/baml_project/src/db.rsbaml_language/crates/baml_project/src/position.rsbaml_language/crates/bex_engine/src/lib.rsbaml_language/crates/bex_engine/tests/collect_tests.rsbaml_language/crates/bex_events/src/history/mod.rsbaml_language/crates/bex_events/src/run.rsbaml_language/crates/bex_project/Cargo.tomlbaml_language/crates/bex_project/src/bex_lsp/mod.rsbaml_language/crates/bex_project/src/bex_lsp/multi_project/diagnostics.rsbaml_language/crates/bex_project/src/bex_lsp/multi_project/mod.rsbaml_language/crates/bex_project/src/bex_lsp/multi_project/notification.rsbaml_language/crates/bex_project/src/bex_lsp/multi_project/request.rsbaml_language/crates/bex_project/src/bex_lsp/multi_project/wasm_helpers.rsbaml_language/crates/bex_project/src/bex_lsp/notification.rsbaml_language/crates/bex_project/src/bex_lsp/request.rsbaml_language/crates/bex_project/src/lib.rsbaml_language/crates/bex_project/src/lsp_ingress.rsbaml_language/crates/bex_project/src/project.rsbaml_language/crates/bridge_wasm/src/lib.rsbaml_language/crates/bridge_wasm/src/wasm_playground.rsbaml_language/crates/bridge_wasm/tests/host_callable.rsbaml_language/crates/bridge_wasm/tests/wasm_fs_glob.rstypescript2/app-promptfiddle/src/playground/baml-lsp-worker.tstypescript2/app-promptfiddle/src/playground/worker-backend.tstypescript2/app-vscode-ext/src/__tests__/projectRoots.test.tstypescript2/app-vscode-ext/src/extension.tstypescript2/app-vscode-ext/src/panels/WebviewPanel.tstypescript2/app-vscode-ext/src/panels/__tests__/sourcePosition.test.tstypescript2/app-vscode-ext/src/panels/sourcePosition.tstypescript2/app-vscode-ext/src/projectRoots.tstypescript2/app-vscode-webview/src/App.tsxtypescript2/app-vscode-webview/src/ExecutionPanel.strict-mode.test.tsxtypescript2/pkg-editor/package.jsontypescript2/pkg-editor/src/MonacoEditor.tsxtypescript2/pkg-editor/src/backend.tstypescript2/pkg-playground/package.jsontypescript2/pkg-playground/src/ExecutionPanel.tsxtypescript2/pkg-playground/src/FunctionSidebar.tsxtypescript2/pkg-playground/src/execution-store.test.tstypescript2/pkg-playground/src/execution-store.tstypescript2/pkg-playground/src/index.tstypescript2/pkg-playground/src/ports/WebSocketRuntimePort.test.tstypescript2/pkg-playground/src/ports/WebSocketRuntimePort.tstypescript2/pkg-playground/src/ports/WorkerRuntimePort.tstypescript2/pkg-playground/src/project-runtime-state.test.tstypescript2/pkg-playground/src/project-runtime-state.tstypescript2/pkg-playground/src/protocol.test.tstypescript2/pkg-playground/src/protocol.tstypescript2/pkg-playground/src/run-store-client.test.tstypescript2/pkg-playground/src/run-store-client.tstypescript2/pkg-playground/src/runtime-port.tstypescript2/pkg-playground/src/worker-protocol.ts
Comment on lines
+51
to
+61
| #[derive(Debug, Clone)] | ||
| enum OutboundPayload { | ||
| Message(lsp_server::Message), | ||
| RawJson(serde_json::Value), | ||
| } | ||
|
|
||
| #[derive(Debug, Clone)] | ||
| pub(crate) struct OutboundFrame { | ||
| payload: OutboundPayload, | ||
| _charge: Arc<OutboundCharge>, | ||
| } |
Contributor
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Share payload storage across OutboundFrame clones.
Broadcast delivery deep-clones Message/Value, while every clone shares one OutboundCharge. The 64 MiB budget therefore undercounts actual memory by the receiver count. Store the payload in an Arc, or charge each clone independently.
Proposed fix
pub(crate) struct OutboundFrame {
- payload: OutboundPayload,
+ payload: Arc<OutboundPayload>,
_charge: Arc<OutboundCharge>,
}
- match &self.payload {
+ match self.payload.as_ref() {
Ok(OutboundFrame {
- payload,
+ payload: Arc::new(payload),Use frame.payload.as_ref() in the writer match as well.
Also applies to: 139-145, 665-669
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@baml_language/crates/baml_lsp_server/src/lib.rs` around lines 51 - 61, Update
OutboundFrame payload storage so clones share the underlying Message/Value via
Arc instead of deep-cloning payload data while sharing one OutboundCharge;
adjust constructors and all payload accesses, including the writer match, to use
payload.as_ref(), and apply the same change to the additional affected
clone/delivery sites.
Comment on lines
+227
to
+267
| let mut content_length = None; | ||
| let mut header = String::new(); | ||
| loop { | ||
| header.clear(); | ||
| let read = input | ||
| .read_line(&mut header) | ||
| .map_err(|error| StdioReadError::Framing(error.to_string()))?; | ||
| if read == 0 { | ||
| return if content_length.is_none() { | ||
| Ok(None) | ||
| } else { | ||
| Err(StdioReadError::Framing( | ||
| "unexpected EOF in LSP headers".to_string(), | ||
| )) | ||
| }; | ||
| } | ||
| let Some(line) = header.strip_suffix("\r\n") else { | ||
| return Err(StdioReadError::Framing( | ||
| "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::<usize>() | ||
| .map_err(|_| StdioReadError::Framing("invalid Content-Length".to_string()))?; | ||
| if parsed > MAX_STDIO_FRAME_BYTES { | ||
| return Err(StdioReadError::Framing(format!( | ||
| "LSP frame exceeds {MAX_STDIO_FRAME_BYTES} bytes" | ||
| ))); | ||
| } | ||
| content_length = Some(parsed); | ||
| } | ||
| } |
Contributor
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Bound header reads before allocation.
read_line can grow indefinitely until CRLF arrives, so the 8 MiB body limit does not prevent an oversized or unterminated header from exhausting memory. Use a capped reader and enforce a total header-byte limit.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@baml_language/crates/baml_lsp_server/src/lib.rs` around lines 227 - 267, The
LSP header parsing loop can allocate unbounded memory before validating
Content-Length. In the header-reading logic, replace unbounded read_line usage
with a capped reader or bounded incremental reads, and track cumulative header
bytes against a fixed maximum header-size limit; return StdioReadError::Framing
when the limit is exceeded, including for unterminated headers, while preserving
CRLF and malformed-header validation.
| self.close_orphaned_document(&endpoint, &tracking); | ||
| } | ||
| } | ||
| lsp_server::Message::Response(_) => {} |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | 🏗️ Heavy lift
Route client responses instead of discarding them.
make_request can emit server-to-client requests, and the scheduler admits client responses, but this branch drops them all. Those requests can never complete. Add response correlation/dispatch to BexLsp, or remove unsupported server-request functionality.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@baml_language/crates/baml_lsp_server/src/lsp_runtime.rs` at line 492, The
lsp_runtime response branch currently discards all client responses, leaving
requests created by make_request unresolved. Update BexLsp’s message handling
and scheduler dispatch to correlate each lsp_server::Message::Response with its
pending request and deliver the result, or remove the unsupported
server-request/make_request functionality consistently.
Comment on lines
+302
to
+314
| #[derive(Clone)] | ||
| pub struct PreparedFunctionRun { | ||
| pub source_revision: u64, | ||
| pub generation: u64, | ||
| pub engine: Arc<dyn crate::Bex>, | ||
| } | ||
|
|
||
| #[derive(Debug, Clone, Copy, serde::Serialize, PartialEq, Eq)] | ||
| #[serde(rename_all = "camelCase")] | ||
| pub struct PreparedTestRun { | ||
| pub source_revision: u64, | ||
| pub generation: u64, | ||
| } |
Contributor
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy lift
Make prepared-run registration RAII-backed.
Preparation registers call_id, but callers only call finish_project_run at the tail of detached tasks. A panic or task abort skips cleanup and can leave runtime/cancellation state pinned. Return a non-Copy lease whose Drop performs cleanup instead of requiring manual pairing.
Based on learnings, registry cleanup should be encapsulated in resource handles’ Drop implementations rather than external cleanup paths.
Also applies to: 406-424
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@baml_language/crates/bex_project/src/bex_lsp/mod.rs` around lines 302 - 314,
Make prepared-run registration RAII-backed by turning PreparedFunctionRun and
PreparedTestRun into non-Copy lease handles that own the registered call_id.
Implement Drop for the handle to invoke the equivalent of finish_project_run and
release runtime/cancellation state, ensuring cleanup on normal completion,
panic, or task abort. Update preparation and detached-task call sites to retain
and drop the lease instead of manually calling finish_project_run.
Source: Learnings
Comment on lines
+45
to
+61
| let document_path = canonical_fs_path_identity(&crate::fs::FsPath::from_vfs(&path)); | ||
| let document = crate::project::OpenDocument { | ||
| client_uri: params.text_document.uri.clone(), | ||
| version: params.text_document.version, | ||
| text: params.text_document.text.clone(), | ||
| }; | ||
| project_handle.project.apply_all_sources_and_open_document( | ||
| &sources, | ||
| document_path.clone(), | ||
| params.text_document.uri, | ||
| params.text_document.version, | ||
| params.text_document.text, | ||
| ); | ||
| drop(in_memory_changes); | ||
|
|
||
| self.refresh_project(&project_root, ProjectRefreshMode::Full); | ||
| self.open_documents | ||
| .lock() | ||
| .unwrap() | ||
| .insert(document_path, document); |
Contributor
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | 🏗️ Heavy lift
Preserve overlays owned by other LSP sessions.
These calls mutate the shared BexProject overlay map, which is keyed only by canonical path. A second connection can overwrite the first connection’s unsaved text, and either connection’s didClose restores disk text while the other editor remains open. (raw.githubusercontent.com)
Include connection/session ownership in the project overlay identity—or retain remaining owners before closing—and add a two-session open/change/close test.
Also applies to: 134-149, 166-176
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@baml_language/crates/bex_project/src/bex_lsp/multi_project/notification.rs`
around lines 45 - 61, Preserve per-session ownership of shared document overlays
in the open-document handling and corresponding change/close paths (including
the regions around the referenced lines). Update the BexProject overlay identity
or ownership tracking so one LSP connection cannot overwrite another session’s
unsaved text or restore disk contents while another session remains open; only
remove/restore an overlay after its final owner closes. Add a test covering two
sessions opening, changing, and closing the same document.
Comment on lines
+717
to
+721
| let forced = MessagePolicy::for_message(&message); | ||
| let policy = MessagePolicy { | ||
| class: stronger_class(forced.class, requested_policy.class), | ||
| cancellation: requested_policy.cancellation, | ||
| }; |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Preserve forced cancellation behavior for protected methods.
A custom policy can downgrade initialize, shutdown, or workspace/executeCommand from NonCancellableWhenRunning, reopening canceled-response/committed-effect races.
Proposed fix
let forced = MessagePolicy::for_message(&message);
+let cancellation = if forced.class.is_reserved() {
+ forced.cancellation
+} else {
+ requested_policy.cancellation
+};
let policy = MessagePolicy {
class: stronger_class(forced.class, requested_policy.class),
- cancellation: requested_policy.cancellation,
+ cancellation,
};📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| let forced = MessagePolicy::for_message(&message); | |
| let policy = MessagePolicy { | |
| class: stronger_class(forced.class, requested_policy.class), | |
| cancellation: requested_policy.cancellation, | |
| }; | |
| let forced = MessagePolicy::for_message(&message); | |
| let cancellation = if forced.class.is_reserved() { | |
| forced.cancellation | |
| } else { | |
| requested_policy.cancellation | |
| }; | |
| let policy = MessagePolicy { | |
| class: stronger_class(forced.class, requested_policy.class), | |
| cancellation, | |
| }; |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@baml_language/crates/bex_project/src/lsp_ingress.rs` around lines 717 - 721,
Update the policy construction in the message-handling logic to preserve the
forced cancellation setting for protected methods. Combine the forced and
requested cancellation policies using the same precedence rule as class
strength, ensuring custom policies cannot downgrade `initialize`, `shutdown`, or
`workspace/executeCommand` from `NonCancellableWhenRunning`; use the existing
`MessagePolicy::for_message` and `stronger_class` logic as references.
Comment on lines
+187
to
+190
| // Do not expose the runtime port until the worker has created the WASM | ||
| // runtime. Otherwise ExecutionPanel's initial catalog-only requestState | ||
| // can race initialization and be silently ignored. | ||
| await workerReadyPromise; |
Contributor
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Add timeout and error handling to workerReadyPromise.
The new await workerReadyPromise will hang indefinitely if the worker fails to send ready (e.g., WASM compilation error, network failure fetching the worker script). There is no worker.addEventListener('error', ...) handler and no timeout on the promise, so connect never resolves or rejects.
🛡️ Proposed fix
const workerReadyPromise = new Promise<void>((resolve, reject) => {
const onMsg = (event: MessageEvent) => {
if (event.data?.type === 'ready') {
worker.removeEventListener('message', onMsg);
worker.removeEventListener('error', onErr);
resolve();
}
};
+ const onErr = (event: ErrorEvent) => {
+ worker.removeEventListener('message', onMsg);
+ worker.removeEventListener('error', onErr);
+ reject(new Error(`Worker failed to initialize: ${event.message}`));
+ };
worker.addEventListener('message', onMsg);
+ worker.addEventListener('error', onErr);
+
+ // Timeout to avoid hanging forever if the worker never sends 'ready'
+ setTimeout(() => {
+ worker.removeEventListener('message', onMsg);
+ worker.removeEventListener('error', onErr);
+ reject(new Error('Worker initialization timed out'));
+ }, 30_000);
});🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@typescript2/app-promptfiddle/src/playground/worker-backend.ts` around lines
187 - 190, Update the worker initialization flow around workerReadyPromise to
reject when the Worker emits an error and to reject after a finite timeout if no
ready message arrives. Ensure connect propagates either failure instead of
hanging, and clean up the error listener and timeout when workerReadyPromise
resolves or rejects.
Comment on lines
+596
to
+609
| 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(); |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Preserve the error state when all client starts fail.
ensureClient sets error, but the final updateAggregateServerState() sees an empty client map and immediately changes it to stopped, hiding the startup failure.
Proposed fix
+ let clientStartFailed = false;
for (const { owner, routablePatterns } of desiredRoots.values()) {
try {
await ensureClient(owner, [...routablePatterns.values()]);
} catch (error) {
+ clientStartFailed = true;
console.error(`Failed to start BAML language server for ${owner.fsPath}`, error);
}
}
- updateAggregateServerState();
+ if (clientStartFailed && clients.size === 0) {
+ updateStatusBar('error');
+ } else {
+ updateAggregateServerState();
+ }📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| 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(); | |
| let clientStartFailed = false; | |
| for (const { owner, routablePatterns } of desiredRoots.values()) { | |
| try { | |
| await ensureClient(owner, [...routablePatterns.values()]); | |
| } catch (error) { | |
| clientStartFailed = true; | |
| 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); | |
| } | |
| } | |
| if (clientStartFailed && clients.size === 0) { | |
| updateStatusBar('error'); | |
| } else { | |
| updateAggregateServerState(); | |
| } |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@typescript2/app-vscode-ext/src/extension.ts` around lines 596 - 609, Preserve
the startup error recorded by ensureClient when every client fails to start.
Update the surrounding startup flow and updateAggregateServerState so an empty
clients map does not overwrite an existing error state with stopped; only
transition to stopped when no startup error is present.
Comment on lines
+150
to
+198
| if (msg.type === 'ensureProjectRuntime') { | ||
| this.desiredRuntimeLease = msg; | ||
| } else if (msg.type === 'releaseProjectRuntime') { | ||
| if ( | ||
| this.desiredRuntimeLease?.project === msg.project && | ||
| this.desiredRuntimeLease.incarnation === msg.incarnation | ||
| ) { | ||
| this.desiredRuntimeLease = null; | ||
| } | ||
| } | ||
|
|
||
| // Every successful open issues one forced state request after the socket | ||
| // subscription exists, so pre-open callers do not need to accumulate | ||
| // duplicate requestState frames in the reconnect queue. | ||
| if ( | ||
| msg.type === 'requestState' && | ||
| (!this.ws || this.ws.readyState !== WebSocket.OPEN) | ||
| ) { | ||
| return; | ||
| } | ||
|
|
||
| const serverMsg = this.toServer(msg); | ||
| if (!serverMsg) return; | ||
| this.sendServerMessage(serverMsg); | ||
| } | ||
|
|
||
| private sendServerMessage(serverMsg: WebSocketInMessage): void { | ||
| const raw = JSON.stringify(serverMsg); | ||
| if (this.ws && this.ws.readyState === WebSocket.OPEN) { | ||
| this.ws.send(raw); | ||
| } else { | ||
| this.outQueue.push(raw); | ||
| // Lease controls describe desired session state, so stale transitions must | ||
| // never accumulate in the generic reconnect queue. `connect()` restores | ||
| // exactly the latest lease after requesting the catalog. | ||
| if ( | ||
| (msg.type === 'ensureProjectRuntime' || | ||
| msg.type === 'releaseProjectRuntime' || | ||
| msg.type === 'retryProjectRuntime') && | ||
| (!this.ws || | ||
| this.ws.readyState !== WebSocket.OPEN || | ||
| this.runtimeDemandSupported !== true) | ||
| ) { | ||
| return; | ||
| } | ||
| // Commands are session-scoped. Never queue them across a disconnect, and | ||
| // do not let them race the hello that establishes the new session. The | ||
| // catalog-only requestState frame is the sole pre-handshake exception. | ||
| if ( | ||
| !this.ws || | ||
| this.ws.readyState !== WebSocket.OPEN || | ||
| (!this.handshakeComplete && serverMsg.type !== 'requestState') || | ||
| (this.handshakeComplete && !this.playgroundCompatible) | ||
| ) { | ||
| return; | ||
| } | ||
| this.ws.send(JSON.stringify(serverMsg)); |
Contributor
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Preserve and replay the desired runtime lease after hello.
An ensureProjectRuntime received before handshake is saved at Line 151 but intentionally not sent; Line 419 then deletes the only copy. The same loss occurs across reconnects. Preserve the latest lease and send it exactly once after a compatible hello confirms runtime-demand support.
Also applies to: 419-443
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@typescript2/pkg-playground/src/ports/WebSocketRuntimePort.ts` around lines
150 - 198, Preserve the latest desired runtime lease across pre-handshake sends
and reconnects instead of deleting it in connect() or hello processing. Update
the hello-success path to replay desiredRuntimeLease exactly once after
confirming playground compatibility and runtimeDemandSupported, then clear or
mark it as sent only after transmission. Ensure ensureProjectRuntime and
releaseProjectRuntime continue updating the lease while disconnected, without
entering the generic reconnect queue.
Comment on lines
+579
to
+587
| if (raw.sessionEpoch !== this.serverSessionEpoch) return null; | ||
| return { | ||
| type: 'controlFlowGraphResult', | ||
| sessionEpoch: raw.sessionEpoch, | ||
| project: raw.project, | ||
| projectIncarnation: raw.projectIncarnation, | ||
| sourceRevision: raw.sourceRevision, | ||
| generation: raw.generation, | ||
| derivedEpoch: raw.derivedEpoch, |
Contributor
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
Align controlFlowGraphResult with the server wire contract.
The Rust WsOutMessage::ControlFlowGraphResult currently serializes only functionName and graph. Consequently, raw.sessionEpoch is undefined, this fence always rejects the frame, and all added identity fields are absent. Extend the server variant and its producers with these fields before enforcing this check.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@typescript2/pkg-playground/src/ports/WebSocketRuntimePort.ts` around lines
579 - 587, Update the server-side WsOutMessage::ControlFlowGraphResult variant
and every producer to serialize sessionEpoch, project, projectIncarnation,
sourceRevision, generation, derivedEpoch, functionName, and graph; ensure the
TypeScript controlFlowGraphResult handler receives these fields so its
sessionEpoch fence and identity metadata are valid.
This was referenced Jul 10, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Root cause
The debounced engine rebuild introduced a second thread that could hold the project database mutex while interactive requests were running. Request paths treated ordinary contention as an error, older build candidates could install after newer edits, and busy diagnostic reads could look like an empty result. Stdio, browser LSP, and playground delivery also used separate or lossy routing paths, which allowed stale session output and unbounded queues.
User impact
Interactive requests no longer fail with the database-mutex
-32001regression, final diagnostics are retained and retried, browser takeover cannot leak old-session output, and Run/Test actions only use the current source/runtime inputs. The playground renders immediately while the selected project prepares and keeps prior test discovery visible but disabled when stale.Performance comparison
Same-machine measurements used
/tmp/baml-inlay-reprowith 400 generated functions, comparing the installed 0.14.1 toolchain with this PR's optimizedbaml-cli. The final driver waits for the orderedinitializeresponse before sendinginitialized, as required by LSP.-32001failures-32800At 1,500 functions, the deliberately worst-case cold-typing workload remains expensive because every edit rewrites every generated function: p50 was about 1.15 s on 0.14.1 and 1.38 s on this PR. The reliability regression is still eliminated: approximately 115 mutex errors and 198 unanswered requests became zero errors and zero unanswered requests. Once source is unchanged, the 1,500-function blast completed 240/240 inlay, hover, and completion requests successfully.
Diagnostic delivery now intentionally uses a 150 ms trailing debounce. At 400 functions, final-edit diagnostics moved from about 59.5 ms to 185 ms while avoiding 94 redundant diagnostic and catalog rebuilds during continuous typing.
Validation
cargo clippy -p bex_project -p baml_lsp_server -p bridge_wasm -p baml_lsp2_actions --all-targets -- -D warningscargo check -p bridge_wasm --target wasm32-unknown-unknownbaml-clibuild.-32800.Summary by CodeRabbit
New Features
Bug Fixes