use std::collections::{HashMap, HashSet};
use serde::{Deserialize, Serialize};
use super::server::LspServerConfig;
use crate::error::{Error, Result};
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct ServerId(String);
impl ServerId {
#[must_use]
pub fn as_str(&self) -> &str {
&self.0
}
}
impl std::fmt::Display for ServerId {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str(&self.0)
}
}
impl From<String> for ServerId {
fn from(id: String) -> Self {
Self(id)
}
}
impl From<&str> for ServerId {
fn from(id: &str) -> Self {
Self(id.to_string())
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ToolKind {
Hover,
Definition,
TypeDefinition,
Implementation,
References,
Diagnostics,
Rename,
Completions,
SignatureHelp,
DocumentSymbols,
WorkspaceSymbols,
FormatDocument,
CodeActions,
CallHierarchy,
InlayHints,
}
impl ToolKind {
pub const ALL: [Self; 15] = [
Self::Hover,
Self::Definition,
Self::TypeDefinition,
Self::Implementation,
Self::References,
Self::Diagnostics,
Self::Rename,
Self::Completions,
Self::SignatureHelp,
Self::DocumentSymbols,
Self::WorkspaceSymbols,
Self::FormatDocument,
Self::CodeActions,
Self::CallHierarchy,
Self::InlayHints,
];
#[must_use]
pub const fn as_str(&self) -> &'static str {
match self {
Self::Hover => "hover",
Self::Definition => "definition",
Self::TypeDefinition => "type_definition",
Self::Implementation => "implementation",
Self::References => "references",
Self::Diagnostics => "diagnostics",
Self::Rename => "rename",
Self::Completions => "completions",
Self::SignatureHelp => "signature_help",
Self::DocumentSymbols => "document_symbols",
Self::WorkspaceSymbols => "workspace_symbols",
Self::FormatDocument => "format_document",
Self::CodeActions => "code_actions",
Self::CallHierarchy => "call_hierarchy",
Self::InlayHints => "inlay_hints",
}
}
}
impl std::fmt::Display for ToolKind {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str(self.as_str())
}
}
fn describe_entry(cfg: &LspServerConfig) -> String {
if cfg.args.is_empty() {
format!("language '{}', command '{}'", cfg.language_id, cfg.command)
} else {
format!(
"language '{}', command '{}', args {:?}",
cfg.language_id, cfg.command, cfg.args
)
}
}
#[derive(Debug, Default)]
struct LanguageRoutes {
explicit: HashMap<ToolKind, ServerId>,
default: Option<ServerId>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum NoServerReason {
NothingRegistered,
NoClaimant,
}
#[derive(Debug, Default)]
pub struct ToolRouter {
by_language: HashMap<String, LanguageRoutes>,
order: Vec<ServerId>,
}
impl ToolRouter {
pub fn from_configs<'a, I>(cfgs: I) -> Result<Self>
where
I: IntoIterator<Item = &'a LspServerConfig>,
{
let mut by_language: HashMap<String, LanguageRoutes> = HashMap::new();
let mut order: Vec<ServerId> = Vec::new();
let mut seen_ids: HashMap<ServerId, String> = HashMap::new();
for cfg in cfgs {
let id = cfg.id();
if let Some(prev_description) = seen_ids.get(&id) {
return Err(Error::InvalidConfig(format!(
"duplicate server id '{id}' in this workspace (used by both an entry with \
{prev_description} and one with {}); add a unique `name` to each \
`[[lsp_servers]]` entry",
describe_entry(cfg)
)));
}
seen_ids.insert(id.clone(), describe_entry(cfg));
order.push(id.clone());
let routes = by_language.entry(cfg.language_id.clone()).or_default();
match &cfg.handles {
None => {
if let Some(existing) = &routes.default {
return Err(Error::InvalidConfig(format!(
"language '{}' has two catch-all servers ('{existing}' and '{id}'); \
at most one server per language may omit `handles`",
cfg.language_id
)));
}
routes.default = Some(id);
}
Some(tools) => {
for tool in tools {
if let Some(existing) = routes.explicit.get(tool) {
return Err(Error::InvalidConfig(format!(
"tool '{tool}' for language '{}' is claimed by both \
'{existing}' and '{id}'",
cfg.language_id
)));
}
routes.explicit.insert(*tool, id.clone());
}
}
}
}
for (language, routes) in &by_language {
if routes.default.is_none() {
let uncovered: Vec<&str> = ToolKind::ALL
.iter()
.filter(|t| !routes.explicit.contains_key(t))
.map(ToolKind::as_str)
.collect();
if !uncovered.is_empty() {
tracing::warn!(
"language '{language}' has no catch-all server and does not claim: {}",
uncovered.join(", ")
);
}
}
}
Ok(Self { by_language, order })
}
#[must_use]
pub fn catch_all<I>(entries: I) -> Self
where
I: IntoIterator<Item = (ServerId, String)>,
{
let mut by_language: HashMap<String, LanguageRoutes> = HashMap::new();
let mut order = Vec::new();
for (id, language) in entries {
order.push(id.clone());
by_language.entry(language).or_default().default = Some(id);
}
Self { by_language, order }
}
pub fn rebind_to_registered(&mut self, registered: &HashSet<ServerId>) {
for (language, routes) in &mut self.by_language {
let live_catch_all = routes.default.clone().filter(|id| registered.contains(id));
let mut dead: HashMap<ServerId, Vec<ToolKind>> = HashMap::new();
for (tool, id) in &routes.explicit {
if !registered.contains(id) {
dead.entry(id.clone()).or_default().push(*tool);
}
}
for (dead_id, tools) in dead {
let tool_names: Vec<&str> = tools.iter().map(ToolKind::as_str).collect();
if let Some(catch_all_id) = &live_catch_all {
for tool in &tools {
routes.explicit.insert(*tool, catch_all_id.clone());
}
tracing::warn!(
"language '{language}': server '{dead_id}' failed to spawn; \
rebinding [{}] to catch-all '{catch_all_id}'",
tool_names.join(", ")
);
} else {
for tool in &tools {
routes.explicit.remove(tool);
}
tracing::warn!(
"language '{language}': server '{dead_id}' failed to spawn and no \
live catch-all is available; [{}] will report no server available",
tool_names.join(", ")
);
}
}
if let Some(dead_catch_all) = routes
.default
.as_ref()
.filter(|id| !registered.contains(*id))
.cloned()
{
routes.default = None;
tracing::warn!(
"language '{language}': catch-all server '{dead_catch_all}' failed to \
spawn; every tool it wasn't already explicitly rebound above will report \
no server available"
);
}
}
self.order.retain(|id| registered.contains(id));
}
#[must_use]
pub fn resolve(&self, language_id: &str, tool: ToolKind) -> Option<&ServerId> {
let routes = self.by_language.get(language_id)?;
routes.explicit.get(&tool).or(routes.default.as_ref())
}
pub fn resolve_any(&self, tool: ToolKind) -> std::result::Result<&ServerId, NoServerReason> {
let claims_explicitly = |id: &ServerId| {
self.by_language
.values()
.any(|r| r.explicit.get(&tool) == Some(id))
};
let is_catch_all = |id: &ServerId| {
self.by_language
.values()
.any(|r| r.default.as_ref() == Some(id))
};
self.order
.iter()
.find(|id| claims_explicitly(id))
.or_else(|| self.order.iter().find(|id| is_catch_all(id)))
.ok_or(if self.order.is_empty() {
NoServerReason::NothingRegistered
} else {
NoServerReason::NoClaimant
})
}
#[must_use]
pub fn has_language(&self, language_id: &str) -> bool {
self.by_language
.get(language_id)
.is_some_and(|r| r.default.is_some() || !r.explicit.is_empty())
}
}
#[cfg(test)]
#[allow(clippy::unwrap_used)]
mod tests {
use super::*;
fn cfg(
language_id: &str,
name: Option<&str>,
handles: Option<Vec<ToolKind>>,
) -> LspServerConfig {
LspServerConfig {
language_id: language_id.to_string(),
command: "cmd".to_string(),
args: vec![],
env: HashMap::new(),
file_patterns: vec![],
initialization_options: None,
timeout_seconds: 30,
request_timeout_seconds: 30,
heuristics: None,
name: name.map(str::to_string),
handles,
}
}
#[test]
fn test_resolve_explicit_wins_over_catch_all() {
let configs = vec![
cfg("python", Some("pyright"), Some(vec![ToolKind::Hover])),
cfg("python", Some("pylsp"), None),
];
let router = ToolRouter::from_configs(&configs).unwrap();
assert_eq!(
router.resolve("python", ToolKind::Hover),
Some(&ServerId::from("pyright"))
);
assert_eq!(
router.resolve("python", ToolKind::Diagnostics),
Some(&ServerId::from("pylsp"))
);
}
#[test]
fn test_resolve_no_catch_all_unclaimed_is_none() {
let configs = vec![cfg("python", Some("pyright"), Some(vec![ToolKind::Hover]))];
let router = ToolRouter::from_configs(&configs).unwrap();
assert_eq!(router.resolve("python", ToolKind::Diagnostics), None);
}
#[test]
fn test_resolve_any_explicit_claimer_beats_catch_all_declared_first() {
let configs = vec![
cfg("python", Some("python-narrow"), Some(vec![ToolKind::Hover])),
cfg("rust", Some("rust-catch-all"), None),
];
let router = ToolRouter::from_configs(&configs).unwrap();
assert_eq!(
router.resolve_any(ToolKind::WorkspaceSymbols),
Ok(&ServerId::from("rust-catch-all"))
);
}
#[test]
fn test_resolve_any_prefers_explicit_claimer_over_catch_all() {
let configs = vec![
cfg("rust", Some("rust-catch-all"), None),
cfg(
"python",
Some("python-explicit"),
Some(vec![ToolKind::WorkspaceSymbols]),
),
];
let router = ToolRouter::from_configs(&configs).unwrap();
assert_eq!(
router.resolve_any(ToolKind::WorkspaceSymbols),
Ok(&ServerId::from("python-explicit"))
);
}
#[test]
fn test_from_configs_rejects_duplicate_server_id_across_languages() {
let configs = vec![
cfg("python", None, None),
cfg("typescript", Some("python"), None),
];
let err = ToolRouter::from_configs(&configs).unwrap_err();
assert!(matches!(err, Error::InvalidConfig(_)));
}
#[test]
fn test_from_configs_duplicate_server_id_error_distinguishes_entries() {
let configs = vec![
LspServerConfig {
language_id: "rust".to_string(),
command: "rust-analyzer".to_string(),
args: vec![],
env: HashMap::new(),
file_patterns: vec![],
initialization_options: None,
timeout_seconds: 30,
request_timeout_seconds: 30,
heuristics: None,
name: None,
handles: None,
},
LspServerConfig {
language_id: "rust".to_string(),
command: "rust-analyzer".to_string(),
args: vec!["--dummy-second-instance".to_string()],
env: HashMap::new(),
file_patterns: vec![],
initialization_options: None,
timeout_seconds: 30,
request_timeout_seconds: 30,
heuristics: None,
name: None,
handles: None,
},
];
let err = ToolRouter::from_configs(&configs).unwrap_err();
let Error::InvalidConfig(msg) = err else {
panic!("expected InvalidConfig, got {err:?}");
};
assert!(!msg.contains("entry #"), "message was: {msg}");
assert!(msg.contains("rust-analyzer"), "message was: {msg}");
assert!(
msg.contains("--dummy-second-instance"),
"message was: {msg}"
);
}
#[test]
fn test_from_configs_duplicate_server_id_error_identical_entries_still_reports() {
let configs = vec![cfg("rust", None, None), cfg("rust", None, None)];
let err = ToolRouter::from_configs(&configs).unwrap_err();
let Error::InvalidConfig(msg) = err else {
panic!("expected InvalidConfig, got {err:?}");
};
assert!(!msg.contains("entry #"), "message was: {msg}");
assert!(
msg.contains("duplicate server id 'rust'"),
"message was: {msg}"
);
}
#[test]
fn test_from_configs_rejects_two_catch_alls() {
let configs = vec![
cfg("python", Some("a"), None),
cfg("python", Some("b"), None),
];
let err = ToolRouter::from_configs(&configs).unwrap_err();
assert!(matches!(err, Error::InvalidConfig(_)));
}
#[test]
fn test_from_configs_rejects_duplicate_tool_claim() {
let configs = vec![
cfg("python", Some("a"), Some(vec![ToolKind::Hover])),
cfg("python", Some("b"), Some(vec![ToolKind::Hover])),
];
let err = ToolRouter::from_configs(&configs).unwrap_err();
assert!(matches!(err, Error::InvalidConfig(_)));
}
#[test]
fn test_rebind_to_registered_dead_server_with_live_catch_all() {
let configs = vec![
cfg("python", Some("pyright"), Some(vec![ToolKind::Hover])),
cfg("python", Some("pylsp"), None),
];
let mut router = ToolRouter::from_configs(&configs).unwrap();
let registered: HashSet<ServerId> = HashSet::from([ServerId::from("pylsp")]);
router.rebind_to_registered(®istered);
assert_eq!(
router.resolve("python", ToolKind::Hover),
Some(&ServerId::from("pylsp"))
);
}
#[test]
fn test_rebind_to_registered_dead_server_no_catch_all_drops_route() {
let configs = vec![
cfg("python", Some("pyright"), Some(vec![ToolKind::Hover])),
cfg("python", Some("pylsp"), Some(vec![ToolKind::Diagnostics])),
];
let mut router = ToolRouter::from_configs(&configs).unwrap();
let registered: HashSet<ServerId> = HashSet::from([ServerId::from("pylsp")]);
router.rebind_to_registered(®istered);
assert_eq!(router.resolve("python", ToolKind::Hover), None);
assert_eq!(
router.resolve("python", ToolKind::Diagnostics),
Some(&ServerId::from("pylsp"))
);
}
#[test]
fn test_rebind_to_registered_all_failed_drops_everything() {
let configs = vec![cfg("rust", None, None)];
let mut router = ToolRouter::from_configs(&configs).unwrap();
router.rebind_to_registered(&HashSet::new());
assert_eq!(router.resolve("rust", ToolKind::Hover), None);
assert_eq!(
router.resolve_any(ToolKind::Hover),
Err(NoServerReason::NothingRegistered)
);
assert!(!router.has_language("rust"));
}
#[test]
fn test_rebind_prunes_order_for_resolve_any() {
let configs = vec![cfg("rust", Some("a"), None), cfg("python", Some("b"), None)];
let mut router = ToolRouter::from_configs(&configs).unwrap();
let registered: HashSet<ServerId> = HashSet::from([ServerId::from("b")]);
router.rebind_to_registered(®istered);
assert_eq!(
router.resolve_any(ToolKind::Hover),
Ok(&ServerId::from("b"))
);
}
#[test]
fn test_resolve_any_no_claimant_does_not_fall_back_to_arbitrary_server() {
let configs = vec![cfg("python", Some("pyright"), Some(vec![ToolKind::Hover]))];
let router = ToolRouter::from_configs(&configs).unwrap();
assert_eq!(
router.resolve_any(ToolKind::WorkspaceSymbols),
Err(NoServerReason::NoClaimant)
);
}
#[test]
fn test_has_language() {
let configs = vec![cfg("rust", None, None)];
let router = ToolRouter::from_configs(&configs).unwrap();
assert!(router.has_language("rust"));
assert!(!router.has_language("python"));
}
#[test]
fn test_catch_all_helper_registers_two_entries() {
let router = ToolRouter::catch_all([
(ServerId::from("ts"), "typescript".to_string()),
(ServerId::from("tsx"), "typescriptreact".to_string()),
]);
assert_eq!(
router.resolve("typescript", ToolKind::Hover),
Some(&ServerId::from("ts"))
);
assert_eq!(
router.resolve("typescriptreact", ToolKind::Hover),
Some(&ServerId::from("tsx"))
);
}
#[test]
fn test_tool_kind_as_str_and_all_len() {
assert_eq!(ToolKind::Hover.as_str(), "hover");
assert_eq!(ToolKind::CallHierarchy.as_str(), "call_hierarchy");
assert_eq!(ToolKind::ALL.len(), 15);
}
#[test]
fn test_server_id_display_and_as_str() {
let id = ServerId::from("pyright");
assert_eq!(id.as_str(), "pyright");
assert_eq!(id.to_string(), "pyright");
}
}