1use std::collections::{HashMap, HashSet};
17
18use serde::{Deserialize, Serialize};
19
20use super::server::LspServerConfig;
21use crate::error::{Error, Result};
22
23#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
31pub struct ServerId(String);
32
33impl ServerId {
34 #[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#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
82#[serde(rename_all = "snake_case")]
83pub enum ToolKind {
84 Hover,
86 Definition,
88 TypeDefinition,
90 Implementation,
92 References,
94 Diagnostics,
96 Rename,
98 Completions,
100 SignatureHelp,
102 DocumentSymbols,
104 WorkspaceSymbols,
106 FormatDocument,
108 CodeActions,
110 CallHierarchy,
112 InlayHints,
114}
115
116impl ToolKind {
117 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 #[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#[derive(Debug, Default)]
168struct LanguageRoutes {
169 explicit: HashMap<ToolKind, ServerId>,
171 default: Option<ServerId>,
174}
175
176#[derive(Debug, Default)]
183pub struct ToolRouter {
184 by_language: HashMap<String, LanguageRoutes>,
185 order: Vec<ServerId>,
189}
190
191impl ToolRouter {
192 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 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 #[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 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 #[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 #[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 #[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 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(®istered);
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(®istered);
576
577 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 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(®istered);
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}