Skip to main content

mcpls_core/bridge/translator/
mod.rs

1//! MCP to LSP translation layer.
2//!
3//! `Translator` owns the LSP client/server registries and dispatches MCP
4//! tool calls to per-domain handler modules. This module defines the
5//! `Translator` struct itself plus setup/lifecycle methods (construction,
6//! registration, shutdown); actual tool-call handling lives in the sibling
7//! modules below, grouped by domain.
8
9use std::collections::{HashMap, HashSet};
10use std::path::{Path, PathBuf};
11use std::sync::{Arc, Mutex as StdMutex};
12
13use tokio::sync::Mutex;
14
15use self::clock::{Clock, SystemClock};
16use self::encoding_ctx::EncodingCtx;
17use self::respawn::RespawnBackoff;
18use crate::bridge::encoding::PositionEncoding;
19use crate::bridge::state::ResourceLimits;
20use crate::bridge::{DocumentTracker, NotificationCache, lock_std};
21use crate::config::{ServerId, ToolKind, ToolRouter};
22use crate::lsp::{LspClient, LspServer, ServerInitConfig};
23
24mod assist;
25mod call_hierarchy;
26#[cfg(test)]
27mod characterization;
28mod clock;
29mod diagnostics;
30mod dto;
31mod edits;
32mod encoding_ctx;
33mod navigation;
34mod respawn;
35mod routing;
36mod symbols;
37#[cfg(test)]
38#[allow(clippy::unwrap_used, clippy::expect_used)]
39mod testing;
40
41pub use dto::*;
42pub use routing::validate_path_against_roots;
43
44/// Translator handles MCP tool calls by converting them to LSP requests.
45///
46/// All fields use interior mutability so `Translator` can be shared via a
47/// plain `Arc<Translator>` with no outer lock: every LSP tool call would
48/// otherwise serialize behind a single mutex for its entire round trip
49/// (including the LSP request timeout), which is the root cause fixed here.
50/// Each field is locked independently and only for the short, synchronous
51/// section that touches it. In particular, the actual LSP request/response
52/// round trip (`client.request(...)`) always runs with no lock held.
53///
54/// `document_tracker` is no exception: `DocumentTracker` locks its own state
55/// per-path internally (see its docs), so `prepare_document`'s call into
56/// `ensure_open` never holds a lock shared across unrelated paths or
57/// languages while it does that document's disk I/O and
58/// `textDocument/didOpen`/`didChange` notify.
59#[derive(Debug)]
60pub struct Translator {
61    /// LSP clients indexed by routing identity. Locked only for the map
62    /// lookup/insert itself, never across an LSP request.
63    lsp_clients: Arc<StdMutex<HashMap<ServerId, LspClient>>>,
64    /// LSP servers indexed by routing identity (held for lifetime management).
65    lsp_servers: Arc<StdMutex<HashMap<ServerId, LspServer>>>,
66    /// Document state tracker. Locks its own state internally, per path.
67    document_tracker: Arc<DocumentTracker>,
68    /// Resource limits `document_tracker` was last built with. Kept
69    /// alongside `document_tracker` so [`Self::with_extensions`] and
70    /// [`Self::with_resource_limits`] can each rebuild the tracker from
71    /// whichever of (limits, extension map) the other has already set,
72    /// regardless of call order -- see [`Self::with_resource_limits`].
73    resource_limits: ResourceLimits,
74    /// Allowed workspace roots for path validation. Read-only after `serve()`
75    /// setup, so no lock is needed.
76    workspace_roots: Arc<Vec<PathBuf>>,
77    /// Custom file extension to language ID mappings. Read-only after
78    /// `serve()` setup, so no lock is needed.
79    extension_map: Arc<HashMap<String, String>>,
80    /// Servers that are configured + applicable but may not have finished
81    /// initializing yet (background init). Used to return a clear "still
82    /// initializing" error instead of "no server configured".
83    expected_servers: Arc<StdMutex<HashSet<ServerId>>>,
84    /// Per-tool routing table: resolves `(language, tool)` to a `ServerId`.
85    /// Locked independently so `rebind_router` (called from a background
86    /// task once registration completes) never contends with an in-flight
87    /// LSP round trip.
88    router: Arc<StdMutex<ToolRouter>>,
89    /// Configs needed to respawn a server if its process dies later, keyed
90    /// by routing identity. Populated once per server right after a
91    /// successful spawn (see [`Self::register_server_config`]); the respawn
92    /// path ([`Self::respawn_if_dead`]) is the only reader.
93    server_configs: Arc<StdMutex<HashMap<ServerId, ServerInitConfig>>>,
94    /// Per-server single-flight lock so concurrent callers that both observe
95    /// a dead process don't race to respawn it independently -- the loser
96    /// waits for the winner's attempt to finish (success or failure) and
97    /// then re-reads whatever ended up registered. See
98    /// [`Self::respawn_if_dead`].
99    respawn_locks: Arc<StdMutex<HashMap<ServerId, Arc<Mutex<()>>>>>,
100    /// Consecutive respawn failures and last-attempt time per server, so a
101    /// crash-looping server backs off instead of eating a fresh
102    /// `timeout_seconds` on every tool call that arrives while it is down.
103    /// See [`Self::respawn_if_dead`].
104    respawn_backoffs: Arc<StdMutex<HashMap<ServerId, RespawnBackoff>>>,
105    /// Diagnostics cache, shared with `serve_with`'s notification pump.
106    ///
107    /// `None` for a `Translator` built without [`Self::with_notification_cache`]
108    /// (e.g. most unit tests). When present, [`Self::respawn_if_dead`] uses
109    /// it to invalidate a respawned server's stale cached diagnostics --
110    /// see that method's docs for why that matters.
111    notification_cache: Option<Arc<Mutex<NotificationCache>>>,
112    /// `AbortHandle` for the currently-running lifecycle-lane forwarding
113    /// task spawned by the most recent [`Self::respawn_if_dead`] call for
114    /// each server, keyed by routing identity. Under a fast crash loop, an
115    /// earlier respawn's forwarder can still be alive (or have a buffered
116    /// notification in flight) when a later respawn for the same `id`
117    /// resets the cache -- aborting the previous handle before installing a
118    /// new one bounds that stale write instead of letting it run
119    /// indefinitely. See [`Self::respawn_if_dead`].
120    ///
121    /// Scope: this only covers forwarder-vs-forwarder races across
122    /// consecutive respawns. The *original* `diagnostics_pump` task from a
123    /// server's initial spawn (`serve_with`'s scope, not `Translator`'s) has
124    /// the same theoretical backlog risk on the *first* respawn, but fixing
125    /// that is out of scope here -- same accepted trade-off already
126    /// documented at `NotificationCache::push_degraded`'s `#249` reference.
127    lifecycle_forwarders: Arc<StdMutex<HashMap<ServerId, tokio::task::AbortHandle>>>,
128    /// Time source for respawn-backoff bookkeeping ([`respawn`](self::respawn)).
129    /// Always [`SystemClock`] in production; overridden via
130    /// [`Self::with_clock`] in tests so backoff-window tests can advance
131    /// time deterministically instead of sleeping in real time.
132    clock: Arc<dyn Clock>,
133    /// Bound `Self::wait_for_indexing_ready` waits for a routed server to
134    /// report indexing readiness. Defaults to `navigation::INDEXING_READY_TIMEOUT`;
135    /// overridable via [`Self::with_indexing_ready_timeout`], wired from
136    /// `workspace.indexing_ready_timeout_seconds` in `mcpls.toml` (#424).
137    indexing_ready_timeout: std::time::Duration,
138}
139
140/// Upper bound on how long [`Translator::shutdown_servers`] waits for a
141/// single LSP server's graceful `shutdown`/`exit` handshake before giving up
142/// and letting `kill_on_drop` terminate it instead.
143const SERVER_SHUTDOWN_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(10);
144
145impl Translator {
146    /// Create a new translator.
147    ///
148    /// Starts with an empty router: nothing is routable until [`Self::with_router`]
149    /// installs one, which matches having no servers registered.
150    ///
151    /// Also starts with no workspace roots, which makes every path-taking
152    /// operation fail closed with `Error::NoWorkspaceRoots` -- embedders
153    /// MUST call [`Self::set_workspace_roots`] before serving any
154    /// path-taking request.
155    #[must_use]
156    pub fn new() -> Self {
157        Self {
158            lsp_clients: Arc::new(StdMutex::new(HashMap::new())),
159            lsp_servers: Arc::new(StdMutex::new(HashMap::new())),
160            document_tracker: Arc::new(DocumentTracker::new(
161                ResourceLimits::default(),
162                HashMap::new(),
163            )),
164            resource_limits: ResourceLimits::default(),
165            workspace_roots: Arc::new(Vec::new()),
166            extension_map: Arc::new(HashMap::new()),
167            expected_servers: Arc::new(StdMutex::new(HashSet::new())),
168            router: Arc::new(StdMutex::new(ToolRouter::default())),
169            server_configs: Arc::new(StdMutex::new(HashMap::new())),
170            respawn_locks: Arc::new(StdMutex::new(HashMap::new())),
171            respawn_backoffs: Arc::new(StdMutex::new(HashMap::new())),
172            notification_cache: None,
173            lifecycle_forwarders: Arc::new(StdMutex::new(HashMap::new())),
174            clock: Arc::new(SystemClock),
175            indexing_ready_timeout: navigation::INDEXING_READY_TIMEOUT,
176        }
177    }
178
179    /// Override the time source used by respawn-backoff bookkeeping.
180    ///
181    /// Test-only: production always uses [`SystemClock`]. Lets
182    /// backoff-window tests advance a `FakeClock` deterministically instead
183    /// of sleeping in real time.
184    #[cfg(test)]
185    #[must_use]
186    fn with_clock(mut self, clock: Arc<dyn Clock>) -> Self {
187        self.clock = clock;
188        self
189    }
190
191    /// Set the workspace roots for path validation.
192    ///
193    /// Only called during single-owner setup, before the translator is
194    /// shared, so this replaces the `Arc` wholesale rather than locking.
195    ///
196    /// Mandatory for any embedder that will serve path-taking requests:
197    /// leaving `roots` empty (or never calling this) makes every such
198    /// operation reject with `Error::NoWorkspaceRoots` instead of allowing
199    /// unrestricted access.
200    pub fn set_workspace_roots(&mut self, roots: Vec<PathBuf>) {
201        self.workspace_roots = Arc::new(roots);
202    }
203
204    /// Give the translator a handle to the shared diagnostics cache, so the
205    /// respawn path can invalidate a respawned server's stale entries.
206    ///
207    /// Only called during single-owner setup (mirrors [`Self::with_router`]),
208    /// before the translator is shared -- `serve_with` passes the same
209    /// `Arc<Mutex<NotificationCache>>` used by the notification pump tasks.
210    #[must_use]
211    pub fn with_notification_cache(mut self, cache: Arc<Mutex<NotificationCache>>) -> Self {
212        self.notification_cache = Some(cache);
213        self
214    }
215
216    /// Override the bound `Self::wait_for_indexing_ready` waits for a routed
217    /// server to report indexing readiness, in place of the built-in
218    /// `navigation::INDEXING_READY_TIMEOUT` default.
219    ///
220    /// Only called during single-owner setup (mirrors [`Self::with_notification_cache`]),
221    /// before the translator is shared. `serve()` wires this from
222    /// `workspace.indexing_ready_timeout_seconds`, already range-checked by
223    /// [`crate::config::ServerConfig::validate`].
224    #[must_use]
225    pub const fn with_indexing_ready_timeout(mut self, timeout: std::time::Duration) -> Self {
226        self.indexing_ready_timeout = timeout;
227        self
228    }
229
230    /// Mark the set of servers that are expected (configured + applicable)
231    /// but may still be initializing in the background.
232    pub fn set_expected_servers(&self, servers: HashSet<ServerId>) {
233        *lock_std(&self.expected_servers) = servers;
234    }
235
236    /// Clear the expected-servers set (e.g. after background init failed).
237    pub fn clear_expected_servers(&self) {
238        lock_std(&self.expected_servers).clear();
239    }
240
241    /// Install the per-tool routing table built from the applicable configs.
242    ///
243    /// Only called during single-owner setup, before the translator is
244    /// shared, so this replaces the `Arc`-wrapped router wholesale.
245    #[must_use]
246    pub fn with_router(mut self, router: ToolRouter) -> Self {
247        self.router = Arc::new(StdMutex::new(router));
248        self
249    }
250
251    /// Rebind the routing table to the set of servers that actually
252    /// registered, dropping or redirecting routes to servers that failed to
253    /// spawn. See `ToolRouter::rebind_to_registered` for the full semantics.
254    pub fn rebind_router(&self, registered: &HashSet<ServerId>) {
255        lock_std(&self.router).rebind_to_registered(registered);
256    }
257
258    /// Whether `id` is the server the router currently resolves
259    /// `ToolKind::Diagnostics` to for `language_id`.
260    ///
261    /// Purpose-built for `register_servers`, which needs this to compute the
262    /// diagnostics-cache filter passed into each pump task, without exposing
263    /// the router's lock guard outside this module.
264    #[must_use]
265    pub fn is_diagnostics_route(&self, language_id: &str, id: &ServerId) -> bool {
266        lock_std(&self.router).resolve(language_id, ToolKind::Diagnostics) == Some(id)
267    }
268
269    /// Negotiated [`PositionEncoding`] of the registered server `id`, or the
270    /// LSP spec's own default (UTF-16) if `id` is not currently registered.
271    ///
272    /// Note this falls back to UTF-16, not [`PositionEncoding::default`]
273    /// (UTF-8): UTF-16 is what an absent/unrecognized negotiation means per
274    /// the LSP spec and what [`crate::lsp::LspServer::spawn`] itself falls
275    /// back to, so this must match rather than use the bridge type's own
276    /// default, which exists only for `PositionEncoding`'s own internal use.
277    #[must_use]
278    pub(crate) fn position_encoding_for(&self, server_id: &ServerId) -> PositionEncoding {
279        lock_std(&self.lsp_servers)
280            .get(server_id)
281            .and_then(|server| PositionEncoding::from_lsp(server.position_encoding().as_str()))
282            .unwrap_or(PositionEncoding::Utf16)
283    }
284
285    /// Build the [`EncodingCtx`] for converting positions/ranges in
286    /// responses from the registered server `id`.
287    fn encoding_ctx(&self, server_id: &ServerId) -> EncodingCtx {
288        EncodingCtx {
289            encoding: self.position_encoding_for(server_id),
290            tracker: self.document_tracker.clone(),
291            workspace_roots: self.workspace_roots.clone(),
292            line_cache: encoding_ctx::new_line_cache(),
293        }
294    }
295
296    /// Rebuilds `document_tracker` from `self.resource_limits` and
297    /// `self.extension_map`, whatever the two are currently set to.
298    ///
299    /// Called by every builder that touches either input ([`Self::with_extensions`],
300    /// [`Self::with_resource_limits`]), so each one only needs to set its own
301    /// field and call this -- it always reads *both* current values, so the
302    /// builders remain order-independent (see [`Self::with_resource_limits`])
303    /// without each one needing to know the other's field. A future builder
304    /// that adds a third tracker input should follow the same pattern:
305    /// update its own field, then call this.
306    fn rebuild_document_tracker(&mut self) {
307        self.document_tracker = Arc::new(DocumentTracker::new(
308            self.resource_limits,
309            (*self.extension_map).clone(),
310        ));
311    }
312
313    /// Configure custom file extension mappings.
314    ///
315    /// This method sets the extension map and updates the document tracker
316    /// to use the same mappings for language detection.
317    ///
318    /// Only called during single-owner setup, before the translator is
319    /// shared, so this replaces the `Arc`-wrapped fields wholesale.
320    #[must_use]
321    pub fn with_extensions(mut self, extension_map: HashMap<String, String>) -> Self {
322        self.extension_map = Arc::new(extension_map);
323        self.rebuild_document_tracker();
324        self
325    }
326
327    /// Configure resource limits (max open documents, max file size) for the
328    /// document tracker.
329    ///
330    /// Only called during single-owner setup, before the translator is
331    /// shared. This builder and [`Self::with_extensions`] may be called in
332    /// either order -- each rebuilds `document_tracker` from *both* of
333    /// `self.resource_limits`/`self.extension_map`'s current values,
334    /// instead of one of them starting fresh from
335    /// `ResourceLimits::default()`/an empty extension map, which previously
336    /// meant whichever builder ran last silently discarded the other's
337    /// effect.
338    #[must_use]
339    pub fn with_resource_limits(mut self, limits: ResourceLimits) -> Self {
340        self.resource_limits = limits;
341        self.rebuild_document_tracker();
342        self
343    }
344
345    /// Register an LSP client under its routing identity.
346    ///
347    /// Only called once per server, from `register_servers` during initial
348    /// background init. The respawn path does not reuse this method: it
349    /// needs the previous client back (to fail its pending requests) and
350    /// must also reset `document_tracker` for the swapped-in server, neither
351    /// of which this method does.
352    pub fn register_client(&self, id: impl Into<ServerId>, client: LspClient) {
353        lock_std(&self.lsp_clients).insert(id.into(), client);
354    }
355
356    /// Register an LSP server under its routing identity.
357    pub fn register_server(&self, id: impl Into<ServerId>, server: LspServer) {
358        lock_std(&self.lsp_servers).insert(id.into(), server);
359    }
360
361    /// Store the config needed to respawn `id` if its process dies later.
362    ///
363    /// Called once per server, right after a successful spawn (see the
364    /// crate-root `register_servers`); [`Self::respawn_if_dead`] is the only
365    /// reader.
366    pub(crate) fn register_server_config(&self, id: impl Into<ServerId>, config: ServerInitConfig) {
367        lock_std(&self.server_configs).insert(id.into(), config);
368    }
369
370    /// Number of currently registered LSP servers.
371    ///
372    /// Test-only: `lsp_servers` is private, so this is the one way a test
373    /// outside this module (e.g. `crate::tests`, exercising
374    /// [`Translator::shutdown_servers`] indirectly through `serve_with`'s
375    /// shutdown sequence) can observe that a registered server was actually
376    /// drained.
377    #[cfg(test)]
378    pub(crate) fn registered_server_count(&self) -> usize {
379        lock_std(&self.lsp_servers).len()
380    }
381
382    /// Snapshot of currently open document paths, used for MCP resource listing.
383    #[must_use]
384    pub fn open_document_paths(&self) -> Vec<PathBuf> {
385        self.document_tracker.open_paths()
386    }
387
388    /// Whether a document is currently tracked as open.
389    #[must_use]
390    pub fn is_document_open(&self, path: &Path) -> bool {
391        self.document_tracker.is_open(path)
392    }
393
394    /// The document tracker, shared with [`EncodingCtx`] so a cache-only
395    /// caller (e.g. `get_cached_diagnostics`) can still prefer tracked
396    /// in-memory content over a disk read when converting positions.
397    #[must_use]
398    pub(crate) const fn document_tracker(&self) -> &Arc<DocumentTracker> {
399        &self.document_tracker
400    }
401
402    /// Gracefully shut down every registered LSP server.
403    ///
404    /// Drains the registered LSP servers and, for each one concurrently,
405    /// sends the LSP `shutdown` request and `exit` notification via
406    /// [`LspServer::shutdown`], bounded by a fixed per-server timeout. A
407    /// server that errors or fails to respond in time is simply dropped
408    /// instead: its child process handle is `kill_on_drop(true)`, so the
409    /// process is killed rather than left running. Call this once, from the
410    /// top-level shutdown path, after the MCP transport has stopped
411    /// accepting new requests.
412    ///
413    /// # Limitations
414    ///
415    /// This only runs on the normal shutdown path (stdio EOF, `SIGTERM`/
416    /// `SIGINT`, or the HTTP transport's own graceful shutdown). This crate's
417    /// workspace `[profile.release]` builds with `panic = "abort"`, so a
418    /// panic reachable from a request handler or background pump task in a
419    /// release build still terminates the process without unwinding — this
420    /// method never runs, and spawned LSP children are orphaned exactly as
421    /// before this fix. Making that path safe would need process-group
422    /// isolation (`kill_on_drop` alone doesn't help, since no `Drop` runs
423    /// either); tracked separately, out of scope here.
424    ///
425    /// `pub(crate)` rather than `pub`: this is meant for exactly one call
426    /// site (`serve_with`'s post-transport shutdown sequence), after the MCP
427    /// transport is already down. An external caller invoking it mid-session
428    /// would drain `lsp_servers` while `lsp_clients` (routing table) still
429    /// points at the now-shut-down servers, so in-flight tool calls would
430    /// resolve to a client whose server is gone.
431    pub(crate) async fn shutdown_servers(&self) {
432        let servers: Vec<(ServerId, LspServer)> = lock_std(&self.lsp_servers).drain().collect();
433        if servers.is_empty() {
434            return;
435        }
436
437        let mut tasks = tokio::task::JoinSet::new();
438        for (id, server) in servers {
439            tasks.spawn(async move {
440                match tokio::time::timeout(SERVER_SHUTDOWN_TIMEOUT, server.shutdown()).await {
441                    Ok(Ok(())) => tracing::debug!(%id, "LSP server shut down gracefully"),
442                    Ok(Err(e)) => tracing::warn!(
443                        %id, error = %e,
444                        "LSP server shutdown handshake failed, killing process instead"
445                    ),
446                    Err(_) => tracing::warn!(
447                        %id, timeout = ?SERVER_SHUTDOWN_TIMEOUT,
448                        "LSP server did not shut down in time, killing process instead"
449                    ),
450                }
451            });
452        }
453        tasks.join_all().await;
454    }
455}
456
457impl Default for Translator {
458    /// Same as [`Translator::new`]: no workspace roots configured, so every
459    /// path-taking operation fails closed with `Error::NoWorkspaceRoots`
460    /// until [`Translator::set_workspace_roots`] is called.
461    fn default() -> Self {
462        Self::new()
463    }
464}
465
466#[cfg(test)]
467#[allow(clippy::unwrap_used, clippy::expect_used)]
468mod tests {
469    use std::collections::{HashMap, HashSet};
470    use std::path::PathBuf;
471
472    use tempfile::TempDir;
473    use tokio::time::Duration;
474
475    use super::*;
476    use crate::bridge::state::detect_language;
477    use crate::config::{ServerId, ToolKind, ToolRouter};
478    use crate::error::Error;
479    use crate::test_lsp::fake_lsp_client;
480
481    #[test]
482    fn test_translator_new() {
483        let translator = Translator::new();
484        assert_eq!(translator.workspace_roots.len(), 0);
485        assert_eq!(lock_std(&translator.lsp_clients).len(), 0);
486        assert_eq!(lock_std(&translator.lsp_servers).len(), 0);
487    }
488
489    #[test]
490    fn test_set_workspace_roots() {
491        let mut translator = Translator::new();
492        let roots = vec![PathBuf::from("/test/root1"), PathBuf::from("/test/root2")];
493        translator.set_workspace_roots(roots.clone());
494        assert_eq!(*translator.workspace_roots, roots);
495    }
496
497    #[test]
498    fn test_register_server() {
499        let translator = Translator::new();
500
501        // Initial state: no servers registered
502        assert_eq!(lock_std(&translator.lsp_servers).len(), 0);
503
504        // The register_server method exists and is callable
505        // Full integration testing with real LspServer is done in integration tests
506        // This unit test verifies the method signature and basic functionality
507
508        // Note: We can't easily construct an LspServer in a unit test without async
509        // and a real LSP server process. The actual registration functionality is
510        // tested in integration tests (see rust_analyzer_tests.rs).
511        // This test verifies the data structure is properly initialized.
512    }
513
514    /// #241: `shutdown_servers` on an empty registry must return immediately
515    /// rather than blocking (e.g. on a `JoinSet` that's never populated).
516    #[tokio::test]
517    async fn test_shutdown_servers_empty_registry_returns_promptly() {
518        let translator = Translator::new();
519
520        let result =
521            tokio::time::timeout(Duration::from_secs(1), translator.shutdown_servers()).await;
522
523        assert!(
524            result.is_ok(),
525            "shutdown_servers must return promptly when no servers are registered"
526        );
527    }
528
529    /// #241: `shutdown_servers` must drain every registered `LspServer` —
530    /// this is the core behavior the issue is about (orphaned LSP children
531    /// on shutdown). Uses `fake_lsp_server()` (an inert in-memory
532    /// transport plus a real `LspServer`, see `lsp::lifecycle`), which
533    /// won't answer the LSP `shutdown` handshake — proving the drain
534    /// completes, via the error fallback path, without hanging on
535    /// non-responsive servers.
536    #[tokio::test]
537    async fn test_shutdown_servers_drains_registered_servers() {
538        let translator = Translator::new();
539        translator.register_server("server-a", crate::lsp::fake_lsp_server());
540        translator.register_server("server-b", crate::lsp::fake_lsp_server());
541        assert_eq!(lock_std(&translator.lsp_servers).len(), 2);
542
543        // Bounded well above `SERVER_SHUTDOWN_TIMEOUT` (10s) so a genuine
544        // regression (a hang) still fails the test instead of the harness
545        // itself timing out ambiguously.
546        let result =
547            tokio::time::timeout(Duration::from_secs(20), translator.shutdown_servers()).await;
548
549        assert!(
550            result.is_ok(),
551            "shutdown_servers must not hang against non-responsive mock servers"
552        );
553        assert_eq!(
554            lock_std(&translator.lsp_servers).len(),
555            0,
556            "all registered servers must be drained"
557        );
558    }
559
560    #[test]
561    fn test_clear_expected_servers_reverts_to_no_server_after_all_routes_dropped() {
562        // Mirrors the real `serve_with` flow: `rebind_router` (called from
563        // `register_servers`/the all-failed path) drops routes to servers
564        // that never registered, then `clear_expected_servers` runs under
565        // the same lock. Subsequent lookups must fall back to
566        // NoServerForLanguage rather than keep implying the server is still
567        // on its way.
568        let path = PathBuf::from("/ws/Assets/Scripts/Player.cs");
569        let lang = detect_language(&path, &HashMap::new());
570        let id = ServerId::from(lang.clone());
571
572        let translator = Translator::new().with_router(ToolRouter::catch_all([(id.clone(), lang)]));
573        let mut expected = HashSet::new();
574        expected.insert(id);
575        translator.set_expected_servers(expected);
576
577        translator.rebind_router(&HashSet::new());
578        translator.clear_expected_servers();
579
580        let err = translator
581            .client_for_file(&path, ToolKind::Hover)
582            .unwrap_err();
583        assert!(matches!(err, Error::NoServerForLanguage(_)));
584    }
585
586    #[test]
587    fn test_translator_with_custom_extensions() {
588        let mut extension_map = HashMap::new();
589        extension_map.insert("nu".to_string(), "nushell".to_string());
590        extension_map.insert("customext".to_string(), "customlang".to_string());
591
592        let translator = Translator::new().with_extensions(extension_map.clone());
593
594        assert_eq!(translator.extension_map.len(), 2);
595        assert_eq!(
596            translator.extension_map.get("nu"),
597            Some(&"nushell".to_string())
598        );
599        assert_eq!(
600            translator.extension_map.get("customext"),
601            Some(&"customlang".to_string())
602        );
603    }
604
605    /// `with_resource_limits` called before `with_extensions` (the order
606    /// `serve()` uses) must reach `document_tracker`. With `max_documents:
607    /// 1` and neither document locked, the second `ensure_open` evicts the
608    /// first (#495) rather than failing -- `document_tracker.len()` staying
609    /// at 1 is what proves the limit actually reached the tracker. Goes
610    /// through `ensure_open` (not the raw `open`) so the first document is
611    /// disk-verified and therefore actually evictable (#495 S4).
612    #[tokio::test]
613    async fn test_with_resource_limits_applies_before_with_extensions() {
614        let limits = ResourceLimits {
615            max_documents: 1,
616            max_file_size: 0,
617        };
618        let translator = Translator::new()
619            .with_resource_limits(limits)
620            .with_extensions(HashMap::new());
621
622        let dir = TempDir::new().unwrap();
623        let path_a = dir.path().join("a.rs");
624        std::fs::write(&path_a, "a").unwrap();
625        let path_b = dir.path().join("b.rs");
626        std::fs::write(&path_b, "b").unwrap();
627
628        let (client, _server) = fake_lsp_client();
629        let server_id = ServerId::from("rust");
630
631        translator
632            .document_tracker
633            .ensure_open(&path_a, &server_id, &client)
634            .await
635            .unwrap();
636        translator
637            .document_tracker
638            .ensure_open(&path_b, &server_id, &client)
639            .await
640            .unwrap();
641        assert_eq!(translator.document_tracker.len(), 1);
642        assert!(!translator.document_tracker.is_open(&path_a));
643        assert!(translator.document_tracker.is_open(&path_b));
644    }
645
646    /// `with_resource_limits` called *after* `with_extensions` (the reverse
647    /// of `serve()`'s order) must still reach `document_tracker` -- the two
648    /// builders must not clobber each other regardless of call order. See
649    /// `Translator::with_resource_limits`'s docs.
650    ///
651    /// Uses a non-empty extension map (unlike the "before" test above) and
652    /// asserts it survived `with_resource_limits`'s rebuild by checking the
653    /// tracked document's resolved `language_id` -- a bug that dropped the
654    /// extension map (e.g. rebuilding from `HashMap::new()` instead of
655    /// `self.extension_map`) would leave `max_documents` correct but the
656    /// extension map silently empty, which the "before" test alone cannot
657    /// detect.
658    #[tokio::test]
659    async fn test_with_resource_limits_applies_after_with_extensions() {
660        let limits = ResourceLimits {
661            max_documents: 1,
662            max_file_size: 0,
663        };
664        let translator = Translator::new()
665            .with_extensions(HashMap::from([("rs".to_string(), "rust".to_string())]))
666            .with_resource_limits(limits);
667
668        let dir = TempDir::new().unwrap();
669        let path_a = dir.path().join("a.rs");
670        std::fs::write(&path_a, "a").unwrap();
671        let path_b = dir.path().join("b.rs");
672        std::fs::write(&path_b, "b").unwrap();
673
674        let (client, _server) = fake_lsp_client();
675        let server_id = ServerId::from("rust");
676
677        translator
678            .document_tracker
679            .ensure_open(&path_a, &server_id, &client)
680            .await
681            .unwrap();
682        translator
683            .document_tracker
684            .ensure_open(&path_b, &server_id, &client)
685            .await
686            .unwrap();
687        assert_eq!(translator.document_tracker.len(), 1);
688        assert!(!translator.document_tracker.is_open(&path_a));
689
690        let state = translator.document_tracker.close(&path_b).unwrap();
691        assert_eq!(
692            state.language_id(),
693            "rust",
694            "extension map must have survived with_resource_limits's rebuild"
695        );
696    }
697}