Skip to main content

mcpls_core/config/
routing.rs

1//! Explicit per-tool routing (#174).
2//!
3//! `language_id` alone is not a unique server identity: two servers can
4//! share one language (e.g. pyright and pylsp both for `python`), each
5//! handling a different subset of MCP tools. This module defines the typed
6//! vocabulary for that routing — [`ServerId`], [`ToolKind`] — and
7//! [`ToolRouter`], which resolves `(language, tool)` to the server that
8//! should handle it.
9//!
10//! `ToolKind` lives here, in `config`, rather than in `mcp` (which is where
11//! its variants are semantically drawn from) to keep `config` a leaf module:
12//! `mcp` and `bridge` both depend on `config`, so putting `ToolKind` in `mcp`
13//! would create a `config -> mcp -> bridge -> config` cycle. When a new
14//! routable MCP tool is added, extend [`ToolKind::ALL`] here.
15
16use std::collections::{HashMap, HashSet};
17
18use serde::{Deserialize, Serialize};
19
20use super::server::LspServerConfig;
21use crate::error::{Error, Result};
22
23/// Unique identity of a configured LSP server within a workspace.
24///
25/// Derived from [`LspServerConfig::id`]: a server's explicit `name` if set,
26/// otherwise its `language_id`. This is the key used throughout the bridge
27/// layer (`Translator::lsp_clients`, `lsp_servers`, notification receivers)
28/// instead of a raw language string, so two servers sharing a language no
29/// longer silently overwrite each other.
30#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
31pub struct ServerId(String);
32
33impl ServerId {
34    /// Borrow the identity as a plain string, e.g. for log messages or map
35    /// lookups against external APIs that expect `&str`.
36    #[must_use]
37    pub fn as_str(&self) -> &str {
38        &self.0
39    }
40}
41
42impl std::fmt::Display for ServerId {
43    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
44        f.write_str(&self.0)
45    }
46}
47
48impl From<String> for ServerId {
49    fn from(id: String) -> Self {
50        Self(id)
51    }
52}
53
54impl From<&str> for ServerId {
55    fn from(id: &str) -> Self {
56        Self(id.to_string())
57    }
58}
59
60/// A routable MCP tool: every MCP tool that dispatches a request to a
61/// specific LSP server via [`ToolRouter`].
62///
63/// Cache-only tools (`get_cached_diagnostics`, `get_server_logs`,
64/// `get_server_messages`) are deliberately excluded — they never reach a
65/// client directly, so they have nothing to route.
66///
67/// `CallHierarchy` covers `prepare`, `incoming_calls`, and `outgoing_calls`
68/// as a single route: the opaque item returned by `prepare` is only
69/// meaningful to the server that produced it, and the incoming/outgoing
70/// handlers never call `ensure_open` themselves — they rely on `prepare`
71/// having already synced the document to the *same* server.
72///
73/// # Examples
74///
75/// ```
76/// use mcpls_core::config::ToolKind;
77///
78/// assert_eq!(ToolKind::Hover.as_str(), "hover");
79/// assert_eq!(ToolKind::ALL.len(), 15);
80/// ```
81#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
82#[serde(rename_all = "snake_case")]
83pub enum ToolKind {
84    /// `textDocument/hover`.
85    Hover,
86    /// `textDocument/definition`.
87    Definition,
88    /// `textDocument/typeDefinition`.
89    TypeDefinition,
90    /// `textDocument/implementation`.
91    Implementation,
92    /// `textDocument/references`.
93    References,
94    /// `textDocument/diagnostic` (pull) and the `publishDiagnostics` cache filter.
95    Diagnostics,
96    /// `textDocument/rename`.
97    Rename,
98    /// `textDocument/completion`.
99    Completions,
100    /// `textDocument/signatureHelp`.
101    SignatureHelp,
102    /// `textDocument/documentSymbol`.
103    DocumentSymbols,
104    /// `workspace/symbol`.
105    WorkspaceSymbols,
106    /// `textDocument/formatting`.
107    FormatDocument,
108    /// `textDocument/codeAction`.
109    CodeActions,
110    /// `textDocument/prepareCallHierarchy`, `callHierarchy/incomingCalls`, `callHierarchy/outgoingCalls`.
111    CallHierarchy,
112    /// `textDocument/inlayHint`.
113    InlayHints,
114}
115
116impl ToolKind {
117    /// Every routable tool, in a fixed order. Used to compute the §5
118    /// coverage warning and to build error messages that enumerate tools.
119    pub const ALL: [Self; 15] = [
120        Self::Hover,
121        Self::Definition,
122        Self::TypeDefinition,
123        Self::Implementation,
124        Self::References,
125        Self::Diagnostics,
126        Self::Rename,
127        Self::Completions,
128        Self::SignatureHelp,
129        Self::DocumentSymbols,
130        Self::WorkspaceSymbols,
131        Self::FormatDocument,
132        Self::CodeActions,
133        Self::CallHierarchy,
134        Self::InlayHints,
135    ];
136
137    /// The `snake_case` name used in config `handles` lists and error messages.
138    #[must_use]
139    pub const fn as_str(&self) -> &'static str {
140        match self {
141            Self::Hover => "hover",
142            Self::Definition => "definition",
143            Self::TypeDefinition => "type_definition",
144            Self::Implementation => "implementation",
145            Self::References => "references",
146            Self::Diagnostics => "diagnostics",
147            Self::Rename => "rename",
148            Self::Completions => "completions",
149            Self::SignatureHelp => "signature_help",
150            Self::DocumentSymbols => "document_symbols",
151            Self::WorkspaceSymbols => "workspace_symbols",
152            Self::FormatDocument => "format_document",
153            Self::CodeActions => "code_actions",
154            Self::CallHierarchy => "call_hierarchy",
155            Self::InlayHints => "inlay_hints",
156        }
157    }
158}
159
160impl std::fmt::Display for ToolKind {
161    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
162        f.write_str(self.as_str())
163    }
164}
165
166/// Per-language routing table: which server handles which tool.
167#[derive(Debug, Default)]
168struct LanguageRoutes {
169    /// Tools explicitly claimed via a server's `handles` list.
170    explicit: HashMap<ToolKind, ServerId>,
171    /// The single server (if any) that omitted `handles` — serves every
172    /// tool not explicitly claimed by another server for this language.
173    default: Option<ServerId>,
174}
175
176/// Resolves `(language, tool)` to the [`ServerId`] that should handle it.
177///
178/// Built once at startup by [`Self::from_configs`] over the *applicable*
179/// (post-heuristics) server configs, then rebound once at registration time
180/// by [`Self::rebind_to_registered`] so that no route ever points at a
181/// server that failed to spawn.
182#[derive(Debug, Default)]
183pub struct ToolRouter {
184    by_language: HashMap<String, LanguageRoutes>,
185    /// Config declaration order, used by `resolve_any` for a deterministic
186    /// choice among candidates. Pruned to registered servers by
187    /// `rebind_to_registered`.
188    order: Vec<ServerId>,
189}
190
191impl ToolRouter {
192    /// Build a router from the configs applicable in this workspace,
193    /// enforcing the workspace-scoped validation rules:
194    ///
195    /// 1. No two applicable servers (in any language) may share a
196    ///    [`ServerId`] — it is the key of every map keyed by server identity.
197    /// 2. No two applicable servers for one language may both omit `handles`
198    ///    (two catch-alls).
199    /// 3. No tool may be claimed via `handles` by two applicable servers of
200    ///    the same language.
201    ///
202    /// Also emits a `tracing::warn!` for any language whose union of
203    /// `handles` claims is partial and has no catch-all server, naming the
204    /// tools nobody will serve.
205    ///
206    /// # Errors
207    ///
208    /// Returns `Error::InvalidConfig` naming the conflicting entries if any
209    /// of the three rules above is violated.
210    pub fn from_configs<'a, I>(cfgs: I) -> Result<Self>
211    where
212        I: IntoIterator<Item = &'a LspServerConfig>,
213    {
214        let mut by_language: HashMap<String, LanguageRoutes> = HashMap::new();
215        let mut order: Vec<ServerId> = Vec::new();
216        let mut seen_ids: HashMap<ServerId, String> = HashMap::new();
217
218        for cfg in cfgs {
219            let id = cfg.id();
220
221            if let Some(prev_language) = seen_ids.get(&id) {
222                return Err(Error::InvalidConfig(format!(
223                    "duplicate server id '{id}' in this workspace (used by both the \
224                     '{prev_language}' and '{}' language entries); add a unique `name` to \
225                     each `[[lsp_servers]]` entry",
226                    cfg.language_id
227                )));
228            }
229            seen_ids.insert(id.clone(), cfg.language_id.clone());
230            order.push(id.clone());
231
232            let routes = by_language.entry(cfg.language_id.clone()).or_default();
233
234            match &cfg.handles {
235                None => {
236                    if let Some(existing) = &routes.default {
237                        return Err(Error::InvalidConfig(format!(
238                            "language '{}' has two catch-all servers ('{existing}' and '{id}'); \
239                             at most one server per language may omit `handles`",
240                            cfg.language_id
241                        )));
242                    }
243                    routes.default = Some(id);
244                }
245                Some(tools) => {
246                    for tool in tools {
247                        if let Some(existing) = routes.explicit.get(tool) {
248                            return Err(Error::InvalidConfig(format!(
249                                "tool '{tool}' for language '{}' is claimed by both \
250                                 '{existing}' and '{id}'",
251                                cfg.language_id
252                            )));
253                        }
254                        routes.explicit.insert(*tool, id.clone());
255                    }
256                }
257            }
258        }
259
260        // Deliberately untested (M4): asserting on `tracing` output would
261        // need a subscriber/capture dev-dependency this crate doesn't
262        // otherwise pull in. Verified by inspection instead; the `uncovered`
263        // computation itself is exercised indirectly by every `resolve`
264        // test above that checks an unclaimed tool returns `None`.
265        for (language, routes) in &by_language {
266            if routes.default.is_none() {
267                let uncovered: Vec<&str> = ToolKind::ALL
268                    .iter()
269                    .filter(|t| !routes.explicit.contains_key(t))
270                    .map(ToolKind::as_str)
271                    .collect();
272                if !uncovered.is_empty() {
273                    tracing::warn!(
274                        "language '{language}' has no catch-all server and does not claim: {}",
275                        uncovered.join(", ")
276                    );
277                }
278            }
279        }
280
281        Ok(Self { by_language, order })
282    }
283
284    /// Build a router where every entry is a catch-all for its language.
285    ///
286    /// Test helper: takes `(id, language)` pairs rather than a single entry
287    /// because some tests (e.g. the `typescript`/`typescriptreact` exact-match
288    /// preference) need two catch-alls registered at once.
289    #[must_use]
290    pub fn catch_all<I>(entries: I) -> Self
291    where
292        I: IntoIterator<Item = (ServerId, String)>,
293    {
294        let mut by_language: HashMap<String, LanguageRoutes> = HashMap::new();
295        let mut order = Vec::new();
296        for (id, language) in entries {
297            order.push(id.clone());
298            by_language.entry(language).or_default().default = Some(id);
299        }
300        Self { by_language, order }
301    }
302
303    /// Rebind every route pointing at a server that did not register — i.e.
304    /// failed to spawn — to that language's live catch-all, or drop the
305    /// route entirely if no catch-all is live.
306    ///
307    /// A dead route is never rebound to a *narrowly-scoped* live server: a
308    /// server that declared `handles = [...]` has explicitly declined every
309    /// other tool, and conscripting it would override that declaration (and,
310    /// via the diagnostics cache filter, start caching diagnostics the user
311    /// deliberately routed away).
312    ///
313    /// # Preconditions
314    ///
315    /// Call this exactly once, after all spawn attempts for a `serve_with`
316    /// invocation have completed and before any request can observe the
317    /// router. This is sound only because `LspServer::spawn_batch` is a
318    /// sequential loop that produces one `ServerInitResult` registered under
319    /// a single lock — registration is one atomic all-or-nothing event, so
320    /// no request can observe a half-rebound router. If server registration
321    /// is ever made incremental (servers registering as they finish spawning,
322    /// rather than all together), an early rebind here would permanently
323    /// steal a slow server's routes with no way back; this function would
324    /// need to be replaced with a design that derives the active table on
325    /// each lookup instead of mutating it once.
326    pub fn rebind_to_registered(&mut self, registered: &HashSet<ServerId>) {
327        for (language, routes) in &mut self.by_language {
328            let live_catch_all = routes.default.clone().filter(|id| registered.contains(id));
329
330            let mut dead: HashMap<ServerId, Vec<ToolKind>> = HashMap::new();
331            for (tool, id) in &routes.explicit {
332                if !registered.contains(id) {
333                    dead.entry(id.clone()).or_default().push(*tool);
334                }
335            }
336
337            for (dead_id, tools) in dead {
338                let tool_names: Vec<&str> = tools.iter().map(ToolKind::as_str).collect();
339                if let Some(catch_all_id) = &live_catch_all {
340                    for tool in &tools {
341                        routes.explicit.insert(*tool, catch_all_id.clone());
342                    }
343                    tracing::warn!(
344                        "language '{language}': server '{dead_id}' failed to spawn; \
345                         rebinding [{}] to catch-all '{catch_all_id}'",
346                        tool_names.join(", ")
347                    );
348                } else {
349                    for tool in &tools {
350                        routes.explicit.remove(tool);
351                    }
352                    tracing::warn!(
353                        "language '{language}': server '{dead_id}' failed to spawn and no \
354                         live catch-all is available; [{}] will report no server available",
355                        tool_names.join(", ")
356                    );
357                }
358            }
359
360            if let Some(dead_catch_all) = routes
361                .default
362                .as_ref()
363                .filter(|id| !registered.contains(*id))
364                .cloned()
365            {
366                routes.default = None;
367                tracing::warn!(
368                    "language '{language}': catch-all server '{dead_catch_all}' failed to \
369                     spawn; every tool it wasn't already explicitly rebound above will report \
370                     no server available"
371                );
372            }
373        }
374
375        self.order.retain(|id| registered.contains(id));
376    }
377
378    /// Resolve the server that should handle `tool` for `language_id`.
379    ///
380    /// Explicit claims win over the language's catch-all; if neither exists,
381    /// returns `None`.
382    #[must_use]
383    pub fn resolve(&self, language_id: &str, tool: ToolKind) -> Option<&ServerId> {
384        let routes = self.by_language.get(language_id)?;
385        routes.explicit.get(&tool).or(routes.default.as_ref())
386    }
387
388    /// Resolve a server for `tool` without a specific language — used for
389    /// workspace-wide tools like `workspace_symbol_search` that have no
390    /// document to detect a language from.
391    ///
392    /// Resolves in three tiers, in config declaration order:
393    /// 1. the first server that explicitly claims `tool`;
394    /// 2. else the first catch-all server;
395    /// 3. else the first server at all.
396    ///
397    /// Tier 2 exists so that a narrowly-scoped server declared before a
398    /// catch-all cannot win a tool it explicitly declined: a catch-all
399    /// claims every tool *implicitly*, so it must still lose to an *explicit*
400    /// claimer in tier 1, but it must beat tier 3's arbitrary "any server"
401    /// fallback.
402    #[must_use]
403    pub fn resolve_any(&self, tool: ToolKind) -> Option<&ServerId> {
404        let claims_explicitly = |id: &ServerId| {
405            self.by_language
406                .values()
407                .any(|r| r.explicit.get(&tool) == Some(id))
408        };
409        let is_catch_all = |id: &ServerId| {
410            self.by_language
411                .values()
412                .any(|r| r.default.as_ref() == Some(id))
413        };
414
415        self.order
416            .iter()
417            .find(|id| claims_explicitly(id))
418            .or_else(|| self.order.iter().find(|id| is_catch_all(id)))
419            .or_else(|| self.order.first())
420    }
421
422    /// Whether `language_id` currently has at least one live-or-configured
423    /// route (a catch-all or an explicit claim), used to distinguish
424    /// `NoServerForTool` (some server handles this language, just not this
425    /// tool) from `NoServerForLanguage` (nothing does).
426    ///
427    /// Deliberately checks route *contents*, not just map-key presence: after
428    /// `rebind_to_registered` drops every route for a language whose sole
429    /// server failed to spawn, this must go back to `false` so that language
430    /// reports `NoServerForLanguage` exactly as it did before per-tool
431    /// routing existed, not `NoServerForTool`.
432    #[must_use]
433    pub fn has_language(&self, language_id: &str) -> bool {
434        self.by_language
435            .get(language_id)
436            .is_some_and(|r| r.default.is_some() || !r.explicit.is_empty())
437    }
438}
439
440#[cfg(test)]
441#[allow(clippy::unwrap_used)]
442mod tests {
443    use super::*;
444
445    fn cfg(
446        language_id: &str,
447        name: Option<&str>,
448        handles: Option<Vec<ToolKind>>,
449    ) -> LspServerConfig {
450        LspServerConfig {
451            language_id: language_id.to_string(),
452            command: "cmd".to_string(),
453            args: vec![],
454            env: HashMap::new(),
455            file_patterns: vec![],
456            initialization_options: None,
457            timeout_seconds: 30,
458            heuristics: None,
459            name: name.map(str::to_string),
460            handles,
461        }
462    }
463
464    #[test]
465    fn test_resolve_explicit_wins_over_catch_all() {
466        let configs = vec![
467            cfg("python", Some("pyright"), Some(vec![ToolKind::Hover])),
468            cfg("python", Some("pylsp"), None),
469        ];
470        let router = ToolRouter::from_configs(&configs).unwrap();
471        assert_eq!(
472            router.resolve("python", ToolKind::Hover),
473            Some(&ServerId::from("pyright"))
474        );
475        assert_eq!(
476            router.resolve("python", ToolKind::Diagnostics),
477            Some(&ServerId::from("pylsp"))
478        );
479    }
480
481    #[test]
482    fn test_resolve_no_catch_all_unclaimed_is_none() {
483        let configs = vec![cfg("python", Some("pyright"), Some(vec![ToolKind::Hover]))];
484        let router = ToolRouter::from_configs(&configs).unwrap();
485        assert_eq!(router.resolve("python", ToolKind::Diagnostics), None);
486    }
487
488    #[test]
489    fn test_resolve_any_explicit_claimer_beats_catch_all_declared_first() {
490        let configs = vec![
491            cfg("python", Some("python-narrow"), Some(vec![ToolKind::Hover])),
492            cfg("rust", Some("rust-catch-all"), None),
493        ];
494        let router = ToolRouter::from_configs(&configs).unwrap();
495        // Neither server explicitly claims WorkspaceSymbols, so the rust
496        // catch-all must win over the narrowly-scoped python server, even
497        // though python was declared first.
498        assert_eq!(
499            router.resolve_any(ToolKind::WorkspaceSymbols),
500            Some(&ServerId::from("rust-catch-all"))
501        );
502    }
503
504    #[test]
505    fn test_resolve_any_prefers_explicit_claimer_over_catch_all() {
506        let configs = vec![
507            cfg("rust", Some("rust-catch-all"), None),
508            cfg(
509                "python",
510                Some("python-explicit"),
511                Some(vec![ToolKind::WorkspaceSymbols]),
512            ),
513        ];
514        let router = ToolRouter::from_configs(&configs).unwrap();
515        assert_eq!(
516            router.resolve_any(ToolKind::WorkspaceSymbols),
517            Some(&ServerId::from("python-explicit"))
518        );
519    }
520
521    #[test]
522    fn test_from_configs_rejects_duplicate_server_id_across_languages() {
523        let configs = vec![
524            cfg("python", None, None),
525            cfg("typescript", Some("python"), None),
526        ];
527        let err = ToolRouter::from_configs(&configs).unwrap_err();
528        assert!(matches!(err, Error::InvalidConfig(_)));
529    }
530
531    #[test]
532    fn test_from_configs_rejects_two_catch_alls() {
533        let configs = vec![
534            cfg("python", Some("a"), None),
535            cfg("python", Some("b"), None),
536        ];
537        let err = ToolRouter::from_configs(&configs).unwrap_err();
538        assert!(matches!(err, Error::InvalidConfig(_)));
539    }
540
541    #[test]
542    fn test_from_configs_rejects_duplicate_tool_claim() {
543        let configs = vec![
544            cfg("python", Some("a"), Some(vec![ToolKind::Hover])),
545            cfg("python", Some("b"), Some(vec![ToolKind::Hover])),
546        ];
547        let err = ToolRouter::from_configs(&configs).unwrap_err();
548        assert!(matches!(err, Error::InvalidConfig(_)));
549    }
550
551    #[test]
552    fn test_rebind_to_registered_dead_server_with_live_catch_all() {
553        let configs = vec![
554            cfg("python", Some("pyright"), Some(vec![ToolKind::Hover])),
555            cfg("python", Some("pylsp"), None),
556        ];
557        let mut router = ToolRouter::from_configs(&configs).unwrap();
558        let registered: HashSet<ServerId> = HashSet::from([ServerId::from("pylsp")]);
559        router.rebind_to_registered(&registered);
560
561        assert_eq!(
562            router.resolve("python", ToolKind::Hover),
563            Some(&ServerId::from("pylsp"))
564        );
565    }
566
567    #[test]
568    fn test_rebind_to_registered_dead_server_no_catch_all_drops_route() {
569        let configs = vec![
570            cfg("python", Some("pyright"), Some(vec![ToolKind::Hover])),
571            cfg("python", Some("pylsp"), Some(vec![ToolKind::Diagnostics])),
572        ];
573        let mut router = ToolRouter::from_configs(&configs).unwrap();
574        let registered: HashSet<ServerId> = HashSet::from([ServerId::from("pylsp")]);
575        router.rebind_to_registered(&registered);
576
577        // pyright died, no catch-all exists, and pylsp never claimed Hover:
578        // the route must drop rather than conscript pylsp.
579        assert_eq!(router.resolve("python", ToolKind::Hover), None);
580        assert_eq!(
581            router.resolve("python", ToolKind::Diagnostics),
582            Some(&ServerId::from("pylsp"))
583        );
584    }
585
586    #[test]
587    fn test_rebind_to_registered_all_failed_drops_everything() {
588        let configs = vec![cfg("rust", None, None)];
589        let mut router = ToolRouter::from_configs(&configs).unwrap();
590        router.rebind_to_registered(&HashSet::new());
591        assert_eq!(router.resolve("rust", ToolKind::Hover), None);
592        assert_eq!(router.resolve_any(ToolKind::Hover), None);
593        // A single-server-per-language config whose server fails to spawn
594        // must report NoServerForLanguage upstream, not NoServerForTool --
595        // has_language must go back to false once every route is dropped.
596        assert!(!router.has_language("rust"));
597    }
598
599    #[test]
600    fn test_rebind_prunes_order_for_resolve_any() {
601        let configs = vec![cfg("rust", Some("a"), None), cfg("python", Some("b"), None)];
602        let mut router = ToolRouter::from_configs(&configs).unwrap();
603        let registered: HashSet<ServerId> = HashSet::from([ServerId::from("b")]);
604        router.rebind_to_registered(&registered);
605        assert_eq!(
606            router.resolve_any(ToolKind::Hover),
607            Some(&ServerId::from("b"))
608        );
609    }
610
611    #[test]
612    fn test_has_language() {
613        let configs = vec![cfg("rust", None, None)];
614        let router = ToolRouter::from_configs(&configs).unwrap();
615        assert!(router.has_language("rust"));
616        assert!(!router.has_language("python"));
617    }
618
619    #[test]
620    fn test_catch_all_helper_registers_two_entries() {
621        let router = ToolRouter::catch_all([
622            (ServerId::from("ts"), "typescript".to_string()),
623            (ServerId::from("tsx"), "typescriptreact".to_string()),
624        ]);
625        assert_eq!(
626            router.resolve("typescript", ToolKind::Hover),
627            Some(&ServerId::from("ts"))
628        );
629        assert_eq!(
630            router.resolve("typescriptreact", ToolKind::Hover),
631            Some(&ServerId::from("tsx"))
632        );
633    }
634
635    #[test]
636    fn test_tool_kind_as_str_and_all_len() {
637        assert_eq!(ToolKind::Hover.as_str(), "hover");
638        assert_eq!(ToolKind::CallHierarchy.as_str(), "call_hierarchy");
639        assert_eq!(ToolKind::ALL.len(), 15);
640    }
641
642    #[test]
643    fn test_server_id_display_and_as_str() {
644        let id = ServerId::from("pyright");
645        assert_eq!(id.as_str(), "pyright");
646        assert_eq!(id.to_string(), "pyright");
647    }
648}