1use std::collections::HashSet;
11use std::sync::Arc;
12use std::time::Duration;
13
14use futures::StreamExt;
15use futures::future::BoxFuture;
16use schemars::JsonSchema;
17use serde::Serialize;
18use tokio::sync::mpsc;
19use tokio::task::AbortHandle;
20
21use super::Translator;
22use super::respawn::BackoffPolicy;
23use super::servers::{Backend, Phase};
24use crate::bridge::indexing::IndexingReset;
25use crate::bridge::{DiagnosticsKey, IndexingState};
26use crate::config::{LanguageId, ServerId};
27use crate::error::{Error, Result};
28use crate::lsp::{ExitGrace, LspNotification, ServerInitConfig};
29use crate::redaction::{Redactions, ServerText};
30use crate::util::lock_std;
31
32const RESTART_COOLDOWN: Duration = Duration::from_secs(5);
34
35const RESTART_HANDSHAKE_BUDGET: Duration = Duration::from_secs(3);
39
40const RESTART_EXIT_GRACE: Duration = crate::lsp::SHUTDOWN_TIMEOUT;
44
45pub const MAX_RESTART_SERVER_IDS: usize = 64;
47
48const RESTART_CONCURRENCY: usize = 4;
50
51pub const MAX_SERVER_ID_BYTES: usize = 256;
53
54#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)]
56pub enum ServerIdsError {
57 #[error("`servers` must not be empty")]
59 Empty,
60 #[error("`servers` may name at most {MAX_RESTART_SERVER_IDS} servers")]
62 TooMany,
63 #[error("server ids must not be blank")]
65 Blank,
66 #[error("server ids must be at most {MAX_SERVER_ID_BYTES} bytes")]
68 TooLong,
69}
70
71#[derive(Debug, Clone, PartialEq, Eq)]
85pub struct ServerIds(Vec<ServerId>);
86
87impl ServerIds {
88 pub fn try_new(ids: Vec<ServerId>) -> std::result::Result<Self, ServerIdsError> {
95 if ids.len() > MAX_RESTART_SERVER_IDS {
96 return Err(ServerIdsError::TooMany);
97 }
98 if ids.iter().any(|id| id.as_str().len() > MAX_SERVER_ID_BYTES) {
99 return Err(ServerIdsError::TooLong);
100 }
101 let mut seen = HashSet::new();
102 let unique: Vec<ServerId> = ids
103 .into_iter()
104 .filter(|id| seen.insert(id.clone()))
105 .collect();
106 if unique.is_empty() {
107 Err(ServerIdsError::Empty)
108 } else {
109 Ok(Self(unique))
110 }
111 }
112
113 pub fn parse(ids: Vec<String>) -> std::result::Result<Self, ServerIdsError> {
130 if ids.len() > MAX_RESTART_SERVER_IDS {
131 return Err(ServerIdsError::TooMany);
132 }
133 let ids = ids
134 .into_iter()
135 .map(|id| ServerId::new(id).map_err(|_| ServerIdsError::Blank))
136 .collect::<std::result::Result<Vec<_>, _>>()?;
137 Self::try_new(ids)
138 }
139
140 #[must_use]
142 pub fn as_slice(&self) -> &[ServerId] {
143 &self.0
144 }
145}
146
147#[derive(Debug, Clone, PartialEq, Eq)]
159pub enum RestartTarget {
160 All,
162 Servers(ServerIds),
164}
165
166#[derive(Debug, Clone, PartialEq, Eq, Serialize, JsonSchema)]
171#[serde(tag = "kind", rename_all = "snake_case")]
172pub enum RestartFailure {
173 SpawnFailed {
175 message: String,
177 },
178 InitializeFailed {
180 message: String,
182 },
183 ShuttingDown,
185}
186
187impl RestartFailure {
188 fn from_error(error: &Error) -> Self {
189 let message = error.to_string();
190 match error {
191 Error::ServerNotFound { .. } | Error::ServerSpawnFailed { .. } => {
192 Self::SpawnFailed { message }
193 }
194 Error::LspInitFailed { .. }
195 | Error::ServerExitedDuringInit { .. }
196 | Error::ServerTerminated
197 | Error::TaskFailed { .. }
198 | Error::StdioCapture(_)
199 | Error::Timeout(_)
200 | Error::LspServerError { .. }
201 | Error::LspProtocolError(_)
202 | Error::ShutdownTimeout
203 | Error::Io(_)
204 | Error::Json(_)
205 | Error::McpServerStart(_)
208 | Error::PathToUri(_)
209 | Error::HttpBind { .. }
210 | Error::InvalidClientPath(_)
211 | Error::MalformedPath { .. }
212 | Error::DocumentNotFound(_)
213 | Error::NoServerForLanguage { .. }
214 | Error::NoServerForTool { .. }
215 | Error::ServerFailedToStart(_)
216 | Error::ServerInitializing { .. }
217 | Error::ServerRestarted { .. }
218 | Error::SymbolResolution(_)
219 | Error::UnknownServers { .. }
220 | Error::WorkspaceServersInitializing
221 | Error::NoServerConfigured
222 | Error::NoServerForWorkspaceTool { .. }
223 | Error::ConfigNotFound(_)
224 | Error::ConfigInsideWorkspace { .. }
225 | Error::Config(_)
226 | Error::TomlDe(_)
227 | Error::TomlSer(_)
228 | Error::NoResolvableListenUris
229 | Error::ResourceUri(_)
230 | Error::InvalidPositionInput(_)
231 | Error::InvalidRangeInput(_)
232 | Error::InvalidHierarchyItemInput(_)
233 | Error::PositionBeyondDocument { .. }
234 | Error::ServerUnavailable { .. }
235 | Error::InvalidToolParams(_)
236 | Error::FileIo { .. }
237 | Error::PathOutsideWorkspace(_)
238 | Error::NoWorkspaceRoots(_)
239 | Error::DocumentLimitExceeded { .. }
240 | Error::SubscriptionLimitReached { .. }
241 | Error::ListenStreamsExhausted { .. }
242 | Error::ListenFilterTooLarge { .. }
243 | Error::FileSizeLimitExceeded { .. }
244 | Error::NotARegularFile(_)
245 | Error::AllServersFailedToInit { .. }
246 | Error::CapabilityNotSupported { .. }
247 | Error::WorkspaceIndexing { .. } => Self::InitializeFailed { message },
248 }
249 }
250}
251
252#[derive(Debug, Clone, PartialEq, Eq, Serialize, JsonSchema)]
263#[serde(tag = "status", rename_all = "snake_case")]
264pub enum RestartOutcome {
265 Restarted {
267 coalesced: bool,
270 indexing_state: IndexingState,
272 push_notifications_degraded: bool,
276 },
277 Failed {
279 reason: RestartFailure,
281 },
282 Throttled {
284 retry_in_ms: u64,
286 },
287 Initializing,
289 NotRunning {
292 message: String,
294 },
295}
296
297#[derive(Debug, Clone, PartialEq, Eq, Serialize, JsonSchema)]
299pub struct ServerRestartEntry {
300 pub server_id: ServerId,
302 #[serde(flatten)]
304 pub outcome: RestartOutcome,
305}
306
307#[derive(Debug, Clone, PartialEq, Eq, Serialize, JsonSchema)]
309pub struct RestartServerResult {
310 pub servers: Vec<ServerRestartEntry>,
312}
313
314#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
316pub(super) struct RestartGeneration(u64);
317
318impl RestartGeneration {
319 const fn bump(&mut self) {
320 self.0 = self.0.saturating_add(1);
321 }
322}
323
324#[derive(Debug)]
326pub struct NotificationReceivers {
327 pub(crate) notifications: crate::lsp::NotificationInbox,
329 pub(crate) lifecycle: mpsc::Receiver<LspNotification>,
331 pub(crate) pinned_tsserver: Option<std::path::PathBuf>,
334}
335
336#[derive(Clone, Copy)]
338pub(super) enum NotificationRouting<'a> {
339 Discard,
342 Pump(&'a dyn NotificationWiring),
344}
345
346#[derive(Debug, Clone, Copy, PartialEq, Eq)]
356pub enum DiagnosticsRole {
357 Authoritative,
360 Secondary,
364}
365
366impl DiagnosticsRole {
367 pub(crate) const fn from_route(is_diagnostics_route: bool) -> Self {
369 if is_diagnostics_route {
370 Self::Authoritative
371 } else {
372 Self::Secondary
373 }
374 }
375}
376
377pub trait NotificationWiring: std::fmt::Debug + Send + Sync {
381 fn spawn_pump(
384 &self,
385 id: ServerId,
386 receivers: NotificationReceivers,
387 role: DiagnosticsRole,
388 ) -> AbortHandle;
389
390 fn publish_invalidated<'a>(&'a self, cleared: &'a [DiagnosticsKey]) -> BoxFuture<'a, ()>;
394
395 fn has_subscriptions(&self) -> BoxFuture<'_, bool> {
399 Box::pin(async { true })
400 }
401
402 fn publish_changed<'a>(&'a self, file: &'a lsp_types::Uri) -> BoxFuture<'a, ()>;
406}
407
408#[derive(Debug)]
411pub struct StartupGuard<'a>(&'a Translator);
412
413impl Drop for StartupGuard<'_> {
414 fn drop(&mut self) {
415 lock_std(&self.0.phase).finish_startup();
416 }
417}
418
419struct Deregistered<'a> {
422 translator: &'a Translator,
423 id: ServerId,
424 backend: Option<Backend>,
425}
426
427impl<'a> Deregistered<'a> {
428 fn take(translator: &'a Translator, id: &ServerId) -> Option<Self> {
432 let backend = lock_std(&translator.servers).take_for_restart(id)?;
433 let held = Self {
434 translator,
435 id: id.clone(),
436 backend: Some(backend),
437 };
438 held.backend
439 .as_ref()
440 .is_some_and(|backend| backend.server().is_some())
441 .then_some(held)
442 }
443
444 async fn terminate(&mut self) -> Option<ServerInitConfig> {
447 let backend = self.backend.as_mut()?;
448 backend.client().mark_restarted(self.id.clone());
449 let server = backend.server_mut()?;
450 let config = server.init_config().clone();
451 let stopped = server
452 .terminate(
453 RESTART_HANDSHAKE_BUDGET,
454 ExitGrace::IfAnswered(RESTART_EXIT_GRACE),
455 )
456 .await;
457 if let Err(error) = stopped {
458 tracing::warn!(id = %self.id, %error, "old LSP server did not stop cleanly; continuing with restart");
459 }
460 Some(config)
461 }
462
463 fn restore(&mut self) {
465 if let Some(backend) = self.backend.take() {
466 let refused = lock_std(&self.translator.servers).restore(&self.id, backend);
467 if refused.is_some() {
468 tracing::warn!(id = %self.id, "restored server dropped: its slot was settled meanwhile");
469 }
470 drop(refused);
471 }
472 }
473}
474
475impl Drop for Deregistered<'_> {
476 fn drop(&mut self) {
477 self.restore();
478 }
479}
480
481impl Translator {
482 pub async fn restart_servers(&self, target: RestartTarget) -> Result<RestartServerResult> {
513 let known = self.known_server_ids();
514 let ids = match target {
515 RestartTarget::All => known,
516 RestartTarget::Servers(requested) => {
517 let unknown: Vec<ServerId> = requested
518 .as_slice()
519 .iter()
520 .filter(|id| !known.contains(id))
521 .cloned()
522 .collect();
523 if !unknown.is_empty() {
524 return Err(Error::UnknownServers {
525 unknown,
526 configured: known,
527 });
528 }
529 requested.0
530 }
531 };
532
533 let requests: Vec<(ServerId, RestartGeneration)> = ids
536 .into_iter()
537 .map(|id| {
538 let generation = self.restart_generation(&id);
539 (id, generation)
540 })
541 .collect();
542 let mut servers: Vec<ServerRestartEntry> = futures::stream::iter(requests)
543 .map(|(id, generation)| async move {
544 let outcome = Box::pin(self.restart_one(&id, generation)).await;
545 ServerRestartEntry {
546 server_id: id,
547 outcome,
548 }
549 })
550 .buffer_unordered(RESTART_CONCURRENCY)
551 .collect()
552 .await;
553 servers.sort_by(|a, b| a.server_id.as_str().cmp(b.server_id.as_str()));
554 Ok(RestartServerResult { servers })
555 }
556
557 fn known_server_ids(&self) -> Vec<ServerId> {
560 let mut ids: Vec<ServerId> = lock_std(&self.servers).ids().cloned().collect();
561 ids.sort_by(|a, b| a.as_str().cmp(b.as_str()));
562 ids
563 }
564
565 fn restart_generation(&self, id: &ServerId) -> RestartGeneration {
566 lock_std(&self.servers)
567 .get(id)
568 .map_or_default(|slot| slot.restart.generation)
569 }
570
571 pub(crate) fn install_wiring(&self, wiring: Arc<dyn NotificationWiring>) {
578 if self.wiring.set(wiring).is_err() {
579 tracing::debug!("notification wiring already installed");
580 }
581 }
582
583 pub(crate) fn set_notification_task(&self, id: &ServerId, handle: AbortHandle) {
586 if let Some(slot) = lock_std(&self.servers).get_mut(id) {
587 slot.notification_task = Some(handle);
588 }
589 }
590
591 pub(crate) fn begin_startup(&self) -> StartupGuard<'_> {
596 lock_std(&self.phase).begin_startup();
597 StartupGuard(self)
598 }
599
600 pub(crate) fn begin_shutdown(&self) {
602 lock_std(&self.phase).begin_shutdown();
603 }
604
605 fn is_shutting_down(&self) -> bool {
606 *lock_std(&self.phase) == Phase::ShuttingDown
607 }
608
609 fn restart_cooldown_remaining(&self, id: &ServerId) -> Option<Duration> {
611 let last = lock_std(&self.servers).get(id)?.restart.last_attempt?;
612 let elapsed = self.clock.now().saturating_duration_since(last);
613 RESTART_COOLDOWN.checked_sub(elapsed)
614 }
615
616 async fn restarted_outcome(&self, id: &ServerId, coalesced: bool) -> RestartOutcome {
617 let (indexing_state, push_notifications_degraded) = match &self.notification_cache {
618 Some(cache) => {
619 let cache = cache.lock().await;
620 (cache.indexing_state(id), cache.is_push_degraded(id))
621 }
622 None => (IndexingState::Unknown, false),
623 };
624 RestartOutcome::Restarted {
625 coalesced,
626 indexing_state,
627 push_notifications_degraded,
628 }
629 }
630
631 fn restart_permit(
636 &self,
637 id: &ServerId,
638 ) -> std::result::Result<&Arc<dyn NotificationWiring>, RestartOutcome> {
639 let phase = *lock_std(&self.phase);
640 if phase == Phase::ShuttingDown {
641 return Err(RestartOutcome::Failed {
642 reason: RestartFailure::ShuttingDown,
643 });
644 }
645 let not_running = {
646 let servers = lock_std(&self.servers);
647 let outcome = servers
648 .server(id)
649 .is_none()
650 .then(|| match servers.failure(id) {
651 Some(failure) => RestartOutcome::NotRunning {
652 message: Error::ServerFailedToStart(Box::new(failure.clone())).to_string(),
653 },
654 None if servers.is_expected(id) => RestartOutcome::Initializing,
655 None => RestartOutcome::NotRunning {
656 message: format!("LSP server '{id}' is not running"),
657 },
658 });
659 drop(servers);
660 outcome
661 };
662 if let Some(outcome) = not_running {
663 return Err(outcome);
664 }
665 if phase == Phase::InitPanicked {
666 return Err(RestartOutcome::NotRunning {
667 message: "startup was interrupted by a panic, so servers cannot be restarted; restart mcpls".to_string(),
668 });
669 }
670 let wiring = match (phase, self.wiring.get()) {
671 (Phase::Settling, _) | (_, None) => return Err(RestartOutcome::Initializing),
672 (_, Some(wiring)) => wiring,
673 };
674 self.restart_cooldown_remaining(id)
675 .map_or(Ok(wiring), |remaining| {
676 Err(RestartOutcome::Throttled {
677 retry_in_ms: u64::try_from(remaining.as_millis()).unwrap_or(u64::MAX),
678 })
679 })
680 }
681
682 async fn restart_one(&self, id: &ServerId, seen: RestartGeneration) -> RestartOutcome {
683 let lock = self.respawn_lock(id);
684 let _serialized = lock.lock().await;
685
686 if self.restart_generation(id) != seen {
687 return self.restarted_outcome(id, true).await;
688 }
689 let wiring = match self.restart_permit(id) {
690 Ok(wiring) => wiring,
691 Err(outcome) => return outcome,
692 };
693
694 self.note_restart_attempt(id);
695 tracing::info!(%id, "restarting LSP server on request");
696
697 let Some(mut old) = Deregistered::take(self, id) else {
698 return RestartOutcome::NotRunning {
699 message: format!("LSP server '{id}' is not running"),
700 };
701 };
702 let config = old.terminate().await;
703 drop(old);
704 let Some(config) = config else {
705 return RestartOutcome::NotRunning {
706 message: format!("LSP server '{id}' is not running"),
707 };
708 };
709 let language_id = config.server_config().language_id.clone();
710
711 if self.is_shutting_down() {
712 self.invalidate_stopped_server(id, &language_id).await;
713 return RestartOutcome::Failed {
714 reason: RestartFailure::ShuttingDown,
715 };
716 }
717 let respawned = self
718 .respawn_locked(
719 id,
720 config,
721 BackoffPolicy::Bypass,
722 NotificationRouting::Pump(wiring.as_ref()),
723 )
724 .await;
725 self.note_restart_attempt(id);
727 if let Err(error) = respawned {
728 tracing::warn!(%id, %error, "LSP server restart failed");
729 self.invalidate_stopped_server(id, &language_id).await;
730 return RestartOutcome::Failed {
731 reason: RestartFailure::from_error(&error),
732 };
733 }
734
735 if self.is_shutting_down() {
736 self.discard_registered(id).await;
737 return RestartOutcome::Failed {
738 reason: RestartFailure::ShuttingDown,
739 };
740 }
741 if let Some(slot) = lock_std(&self.servers).get_mut(id) {
742 slot.restart.generation.bump();
743 }
744 tracing::info!(%id, "LSP server restarted");
745 self.restarted_outcome(id, false).await
746 }
747
748 fn note_restart_attempt(&self, id: &ServerId) {
749 if let Some(slot) = lock_std(&self.servers).get_mut(id) {
750 slot.restart.last_attempt = Some(self.clock.now());
751 }
752 }
753
754 async fn invalidate_stopped_server(&self, id: &ServerId, language_id: &LanguageId) {
757 let stale_pump = lock_std(&self.servers)
761 .get_mut(id)
762 .and_then(|slot| slot.notification_task.take());
763 if let Some(pump) = stale_pump {
764 pump.abort();
765 tokio::task::yield_now().await;
766 }
767 let mut cleared = Vec::new();
768 if let Some(cache) = &self.notification_cache {
769 let mut cache = cache.lock().await;
770 cache.reset_indexing_state(id, IndexingReset::Forget);
771 if self.is_diagnostics_route(language_id, id) {
772 cleared = cache.clear_server_diagnostics(id);
773 cache.mark_push_degraded(id);
774 }
775 }
776 if let Some(wiring) = self.wiring.get() {
777 wiring.publish_invalidated(&cleared).await;
778 }
779 }
780
781 async fn discard_registered(&self, id: &ServerId) {
784 let server = lock_std(&self.servers).remove_server(id);
785 if let Some(server) = server
786 && let Err(error) = server.shutdown().await
787 {
788 tracing::warn!(%id, %error, "replacement LSP server shutdown failed");
789 }
790 }
791}
792
793impl ServerText for RestartFailure {
794 fn redact_server_text(&mut self, redactions: &Redactions) {
795 match self {
796 Self::SpawnFailed { message } | Self::InitializeFailed { message } => {
797 redactions.redact_in_place(message);
798 }
799 Self::ShuttingDown => {}
800 }
801 }
802}
803
804impl ServerText for RestartOutcome {
805 fn redact_server_text(&mut self, redactions: &Redactions) {
806 match self {
807 Self::Restarted {
808 coalesced: _,
809 indexing_state: _,
810 push_notifications_degraded: _,
811 }
812 | Self::Throttled { retry_in_ms: _ }
813 | Self::Initializing => {}
814 Self::Failed { reason } => reason.redact_server_text(redactions),
815 Self::NotRunning { message } => redactions.redact_in_place(message),
816 }
817 }
818}
819
820impl ServerText for ServerRestartEntry {
821 fn redact_server_text(&mut self, redactions: &Redactions) {
822 let Self {
823 server_id: _,
824 outcome,
825 } = self;
826 outcome.redact_server_text(redactions);
827 }
828}
829
830impl ServerText for RestartServerResult {
831 fn redact_server_text(&mut self, redactions: &Redactions) {
832 let Self { servers } = self;
833 servers.redact_server_text(redactions);
834 }
835}
836
837#[cfg(test)]
838mod tests {
839 use std::assert_matches;
840
841 use super::*;
842
843 #[test]
844 fn server_ids_drop_repeats_keep_order_and_reject_empty() {
845 let ids = ServerIds::try_new(vec![
846 ServerId::from_static("b"),
847 ServerId::from_static("a"),
848 ServerId::from_static("b"),
849 ])
850 .unwrap();
851 assert_eq!(
852 ids.as_slice(),
853 [ServerId::from_static("b"), ServerId::from_static("a")]
854 );
855 assert_eq!(ServerIds::try_new(Vec::new()), Err(ServerIdsError::Empty));
856 let ids = |count: usize| {
857 (0..count)
858 .map(|i| ServerId::new(format!("s{i}")).unwrap())
859 .collect()
860 };
861 assert!(ServerIds::try_new(ids(MAX_RESTART_SERVER_IDS)).is_ok());
862 assert_eq!(
863 ServerIds::try_new(ids(MAX_RESTART_SERVER_IDS + 1)),
864 Err(ServerIdsError::TooMany)
865 );
866 assert_eq!(
867 ServerIds::parse(vec![" ".to_owned()]),
868 Err(ServerIdsError::Blank)
869 );
870 assert_eq!(
871 ServerIds::try_new(vec![
872 ServerId::new("x".repeat(MAX_SERVER_ID_BYTES + 1)).unwrap()
873 ]),
874 Err(ServerIdsError::TooLong)
875 );
876 assert_eq!(
877 ServerIds::parse(vec![" ".to_owned(); MAX_RESTART_SERVER_IDS + 1]),
878 Err(ServerIdsError::TooMany)
879 );
880 assert_eq!(
881 ServerIds::parse(vec!["x".repeat(MAX_SERVER_ID_BYTES + 1), " ".to_owned()]),
882 Err(ServerIdsError::Blank)
883 );
884 }
885
886 #[test]
887 fn restart_entry_flattens_the_outcome_next_to_the_server_id() {
888 let entry = ServerRestartEntry {
889 server_id: ServerId::from_static("rust"),
890 outcome: RestartOutcome::Failed {
891 reason: RestartFailure::SpawnFailed {
892 message: "gone".to_string(),
893 },
894 },
895 };
896 assert_eq!(
897 serde_json::to_value(entry).unwrap(),
898 serde_json::json!({
899 "server_id": "rust",
900 "status": "failed",
901 "reason": {"kind": "spawn_failed", "message": "gone"},
902 })
903 );
904 }
905
906 #[test]
907 fn spawn_errors_map_to_spawn_failed_and_the_rest_to_initialize_failed() {
908 let not_found = Error::ServerNotFound {
909 command: crate::config::ServerCommand::from_static("x"),
910 source: std::io::Error::from(std::io::ErrorKind::NotFound),
911 };
912 assert_matches!(
913 RestartFailure::from_error(¬_found),
914 RestartFailure::SpawnFailed { .. }
915 );
916 assert_matches!(
917 RestartFailure::from_error(&Error::Timeout(Duration::from_secs(5))),
918 RestartFailure::InitializeFailed { .. }
919 );
920 }
921
922 #[cfg(unix)]
923 mod restart_tests {
924 use std::fs;
925 use std::path::{Path, PathBuf};
926
927 use tempfile::TempDir;
928 use tokio::sync::Mutex;
929 use tokio_util::sync::CancellationToken;
930
931 use super::*;
932 use crate::bridge::translator::clock::{Clock, FakeClock};
933 use crate::bridge::translator::testing::{
934 stub_server_config, write_crash_after_init_script, write_protocol_server_script,
935 write_responder_script, write_slow_exit_server_script,
936 };
937 use crate::bridge::{NotificationCache, WorkspaceRoots};
938 use crate::config::{LanguageId, ServerCommand, ToolRouter};
939 use crate::error::{ServerSpawnFailure, StartupFailure};
940 use crate::mcp::SubscriptionRegistry;
941 use crate::runtime::pump::{PumpShared, PumpWiring};
942
943 struct Fixture {
944 translator: Arc<Translator>,
945 cache: Arc<Mutex<NotificationCache>>,
946 clock: Arc<FakeClock>,
947 id: ServerId,
948 _dir: TempDir,
949 }
950
951 fn server_ids(names: &[&str]) -> RestartTarget {
952 RestartTarget::Servers(
953 ServerIds::try_new(names.iter().map(|n| ServerId::new(*n).unwrap()).collect())
954 .unwrap(),
955 )
956 }
957
958 async fn fixture(dir: TempDir, script: &Path, wired: bool) -> Fixture {
959 let id = ServerId::from_static("rust");
960 let cache = Arc::new(Mutex::new(NotificationCache::new()));
961 let clock = Arc::new(FakeClock::new());
962 let mut translator = Translator::new()
963 .with_router(ToolRouter::catch_all([(
964 id.clone(),
965 LanguageId::from_static("rust"),
966 )]))
967 .with_notification_cache(Arc::clone(&cache))
968 .with_clock(clock.clone());
969 translator.set_workspace_roots(
970 WorkspaceRoots::from_paths(&[dir.path().to_path_buf()]).unwrap(),
971 );
972
973 let mut server = crate::lsp::LspServer::spawn(stub_server_config("rust", script))
974 .await
975 .unwrap();
976 let receivers = NotificationReceivers {
977 notifications: server.take_notification_inbox(),
978 lifecycle: server.take_lifecycle_rx(),
979 pinned_tsserver: None,
980 };
981 translator.register_server_complete(server);
982
983 let wiring = PumpWiring::new(
984 PumpShared {
985 roles: crate::runtime::pump::DiagnosticsRoles::default(),
986 notification_cache: Arc::clone(&cache),
987 subs: SubscriptionRegistry::new(),
988 workspace_roots: translator.workspace_roots.clone(),
989 },
990 CancellationToken::new(),
991 );
992 let pump = wiring.spawn_pump(id.clone(), receivers, DiagnosticsRole::Authoritative);
993 translator.set_notification_task(&id, pump);
994 if wired {
995 translator.install_wiring(Arc::new(wiring));
996 }
997 Fixture {
998 translator: Arc::new(translator),
999 cache,
1000 clock,
1001 id,
1002 _dir: dir,
1003 }
1004 }
1005
1006 async fn protocol_fixture(publish: Option<&str>) -> (Fixture, PathBuf) {
1007 let dir = TempDir::new().unwrap();
1008 let log = dir.path().join("server.log");
1009 let script = write_protocol_server_script(dir.path(), &log, publish);
1010 (fixture(dir, &script, true).await, log)
1011 }
1012
1013 fn read_log(log: &Path) -> String {
1014 fs::read_to_string(log).unwrap_or_default()
1015 }
1016
1017 fn only_outcome(result: &RestartServerResult) -> &RestartOutcome {
1018 assert_eq!(result.servers.len(), 1, "{result:?}");
1019 &result.servers[0].outcome
1020 }
1021
1022 impl Fixture {
1023 fn registered(&self) -> bool {
1024 let servers = lock_std(&self.translator.servers);
1025 servers.server(&self.id).is_some() && servers.client(&self.id).is_some()
1026 }
1027
1028 fn expected(&self) -> bool {
1029 lock_std(&self.translator.servers).is_expected(&self.id)
1030 }
1031
1032 fn is_dead(&self) -> bool {
1033 lock_std(&self.translator.servers)
1034 .server_mut(&self.id)
1035 .unwrap()
1036 .is_dead()
1037 .unwrap()
1038 }
1039
1040 async fn wait_until_deregistered(&self) {
1041 tokio::time::timeout(Duration::from_secs(5), async {
1042 while !self.expected() {
1043 tokio::time::sleep(Duration::from_millis(10)).await;
1044 }
1045 })
1046 .await
1047 .expect("the restart must deregister the server");
1048 }
1049 }
1050
1051 #[tokio::test]
1052 async fn test_restart_replaces_the_process_after_a_graceful_stop() {
1053 let (fx, log) = protocol_fixture(None).await;
1054
1055 let result = fx
1056 .translator
1057 .restart_servers(RestartTarget::All)
1058 .await
1059 .unwrap();
1060
1061 assert_matches!(
1062 only_outcome(&result),
1063 RestartOutcome::Restarted {
1064 coalesced: false,
1065 ..
1066 },
1067 "{result:?}"
1068 );
1069 let log = read_log(&log);
1070 assert_eq!(log.matches("started").count(), 2, "{log}");
1071 assert!(log.contains(r#""method":"shutdown""#), "{log}");
1072 assert!(log.contains(r#""method":"exit""#), "{log}");
1073 assert!(fx.registered() && !fx.expected() && !fx.is_dead());
1074 }
1075
1076 #[tokio::test]
1077 async fn test_restart_of_a_wedged_server_completes_despite_the_terminate_error() {
1078 let dir = TempDir::new().unwrap();
1079 let script = write_responder_script(dir.path(), 600);
1080 let fx = fixture(dir, &script, true).await;
1081
1082 let result = fx
1083 .translator
1084 .restart_servers(RestartTarget::All)
1085 .await
1086 .unwrap();
1087
1088 assert_matches!(
1089 only_outcome(&result),
1090 RestartOutcome::Restarted { .. },
1091 "{result:?}"
1092 );
1093 assert!(fx.registered() && !fx.expected() && !fx.is_dead());
1094 }
1095
1096 #[tokio::test]
1097 async fn test_cancelled_restart_restores_the_registry() {
1098 let dir = TempDir::new().unwrap();
1099 let script = write_responder_script(dir.path(), 600);
1100 let fx = fixture(dir, &script, true).await;
1101
1102 let translator = Arc::clone(&fx.translator);
1103 let task =
1104 tokio::spawn(async move { translator.restart_servers(RestartTarget::All).await });
1105 fx.wait_until_deregistered().await;
1106 assert!(!fx.registered());
1107
1108 task.abort();
1109 assert!(task.await.unwrap_err().is_cancelled());
1110
1111 assert!(fx.registered() && !fx.expected());
1112 }
1113
1114 #[tokio::test]
1115 async fn test_restart_of_a_startup_failed_server_reports_not_running() {
1116 let id = ServerId::from_static("rust");
1117 let translator = Translator::new().with_router(ToolRouter::catch_all([(
1118 id.clone(),
1119 LanguageId::from_static("rust"),
1120 )]));
1121 translator.record_startup_failures(&[ServerSpawnFailure {
1122 server_id: id.clone(),
1123 language_id: LanguageId::from_static("rust"),
1124 command: ServerCommand::from_static("missing"),
1125 reason: StartupFailure::InitTaskPanicked,
1126 }]);
1127 translator.rebind_router(&HashSet::new());
1128 translator.clear_expected_servers();
1129
1130 let result = translator
1131 .restart_servers(server_ids(&["rust"]))
1132 .await
1133 .unwrap();
1134
1135 assert_matches!(
1136 only_outcome(&result),
1137 RestartOutcome::NotRunning { .. },
1138 "{result:?}"
1139 );
1140 }
1141
1142 #[tokio::test]
1143 async fn test_restart_all_reports_a_refused_server_as_not_running() {
1144 let config = crate::config::LspServerConfig::rust_analyzer();
1145 let translator =
1146 Translator::new().with_router(ToolRouter::from_configs([&config]).unwrap());
1147 translator.record_refusals(&[ServerSpawnFailure {
1148 server_id: config.id(),
1149 language_id: config.language_id.clone(),
1150 command: config.command.server_command().clone(),
1151 reason: StartupFailure::RefusedUntrustedWorkspace(
1152 crate::error::UntrustedRefusal::NotAllowed { builtin: None },
1153 ),
1154 }]);
1155
1156 let result = translator
1157 .restart_servers(RestartTarget::All)
1158 .await
1159 .unwrap();
1160
1161 let outcome = only_outcome(&result);
1162 assert_matches!(
1163 outcome,
1164 RestartOutcome::NotRunning { message } if message.contains("--allow-server rust"),
1165 "{result:?}"
1166 );
1167 }
1168
1169 #[tokio::test]
1170 async fn test_restart_of_a_settled_server_while_startup_settles_reports_initializing() {
1171 let dir = TempDir::new().unwrap();
1172 let log = dir.path().join("server.log");
1173 let script = write_protocol_server_script(dir.path(), &log, None);
1174 let fx = fixture(dir, &script, true).await;
1175 let startup = fx.translator.begin_startup();
1176
1177 let during = fx
1178 .translator
1179 .restart_servers(server_ids(&["rust"]))
1180 .await
1181 .unwrap();
1182 assert_eq!(*only_outcome(&during), RestartOutcome::Initializing);
1183
1184 drop(startup);
1185 let after = fx
1186 .translator
1187 .restart_servers(server_ids(&["rust"]))
1188 .await
1189 .unwrap();
1190 assert_matches!(
1191 only_outcome(&after),
1192 RestartOutcome::Restarted { .. },
1193 "{after:?}"
1194 );
1195 }
1196
1197 #[tokio::test]
1198 async fn test_restart_permit_classifies_in_order() {
1199 let id = ServerId::from_static("rust");
1200 let translator = Translator::new();
1201 assert_matches!(
1202 translator.restart_permit(&id),
1203 Err(RestartOutcome::NotRunning { .. })
1204 );
1205
1206 translator.set_expected_servers(HashSet::from([id.clone()]));
1207 assert_matches!(
1208 translator.restart_permit(&id),
1209 Err(RestartOutcome::Initializing)
1210 );
1211
1212 translator.clear_expected_servers();
1213 translator.record_startup_failures(&[ServerSpawnFailure {
1214 server_id: id.clone(),
1215 language_id: crate::config::LanguageId::from_static("rust"),
1216 command: ServerCommand::from_static("missing"),
1217 reason: StartupFailure::InitTaskPanicked,
1218 }]);
1219 assert_matches!(
1220 translator.restart_permit(&id),
1221 Err(RestartOutcome::NotRunning { .. })
1222 );
1223
1224 let config = crate::config::LspServerConfig::rust_analyzer();
1225 let running = Translator::new();
1226 running
1227 .register_server_complete(crate::lsp::fake_lsp_server_with_config(config.clone()));
1228 let rid = config.id();
1229 assert_matches!(
1230 running.restart_permit(&rid),
1231 Err(RestartOutcome::Initializing)
1232 );
1233
1234 lock_std(&running.phase).init_panicked();
1235 assert_matches!(
1236 running.restart_permit(&rid),
1237 Err(RestartOutcome::NotRunning { .. })
1238 );
1239
1240 running.begin_shutdown();
1241 assert_matches!(
1242 running.restart_permit(&rid),
1243 Err(RestartOutcome::Failed {
1244 reason: RestartFailure::ShuttingDown
1245 })
1246 );
1247 }
1248
1249 #[test]
1250 fn test_dropping_the_startup_guard_clears_the_flag() {
1251 let translator = Translator::new();
1252 let startup = translator.begin_startup();
1253 assert_eq!(*lock_std(&translator.phase), Phase::Settling);
1254 drop(startup);
1255 assert_eq!(*lock_std(&translator.phase), Phase::Settled);
1256 }
1257
1258 #[tokio::test]
1259 async fn test_restart_of_a_startup_failed_server_while_startup_settles_reports_not_running()
1260 {
1261 let id = ServerId::from_static("rust");
1262 let translator = Translator::new().with_router(ToolRouter::catch_all([(
1263 id.clone(),
1264 LanguageId::from_static("rust"),
1265 )]));
1266 translator.record_startup_failures(&[ServerSpawnFailure {
1267 server_id: id,
1268 language_id: LanguageId::from_static("rust"),
1269 command: ServerCommand::from_static("missing"),
1270 reason: StartupFailure::InitTaskPanicked,
1271 }]);
1272 let _startup = translator.begin_startup();
1273
1274 let result = translator
1275 .restart_servers(server_ids(&["rust"]))
1276 .await
1277 .unwrap();
1278
1279 assert_matches!(
1280 only_outcome(&result),
1281 RestartOutcome::NotRunning { .. },
1282 "{result:?}"
1283 );
1284 }
1285
1286 #[tokio::test]
1287 async fn test_restart_of_a_still_expected_server_reports_initializing() {
1288 let id = ServerId::from_static("rust");
1289 let translator = Translator::new().with_router(ToolRouter::catch_all([(
1290 id.clone(),
1291 LanguageId::from_static("rust"),
1292 )]));
1293 translator.set_expected_servers(HashSet::from([id]));
1294
1295 let result = translator
1296 .restart_servers(server_ids(&["rust"]))
1297 .await
1298 .unwrap();
1299
1300 assert_eq!(*only_outcome(&result), RestartOutcome::Initializing);
1301 }
1302
1303 #[tokio::test]
1304 async fn test_restart_before_the_wiring_is_installed_reports_initializing() {
1305 let dir = TempDir::new().unwrap();
1306 let log = dir.path().join("server.log");
1307 let script = write_protocol_server_script(dir.path(), &log, None);
1308 let fx = fixture(dir, &script, false).await;
1309
1310 let result = fx
1311 .translator
1312 .restart_servers(RestartTarget::All)
1313 .await
1314 .unwrap();
1315
1316 assert_eq!(*only_outcome(&result), RestartOutcome::Initializing);
1317 assert_eq!(read_log(&log).matches("started").count(), 1);
1318 assert!(fx.registered());
1319 }
1320
1321 #[tokio::test]
1322 async fn test_restart_rejects_unknown_ids_and_restarts_nothing() {
1323 let (fx, log) = protocol_fixture(None).await;
1324
1325 let err = fx
1326 .translator
1327 .restart_servers(server_ids(&["rust", "nope"]))
1328 .await
1329 .unwrap_err();
1330
1331 let Error::UnknownServers {
1332 unknown,
1333 configured,
1334 } = err
1335 else {
1336 panic!("expected UnknownServers, got {err:?}");
1337 };
1338 assert_eq!(unknown, [ServerId::from_static("nope")]);
1339 assert_eq!(configured, [ServerId::from_static("rust")]);
1340 assert_eq!(read_log(&log).matches("started").count(), 1);
1341 }
1342
1343 #[tokio::test]
1344 async fn test_restart_is_throttled_until_the_cooldown_elapses() {
1345 let (fx, _log) = protocol_fixture(None).await;
1346 let first = fx
1347 .translator
1348 .restart_servers(RestartTarget::All)
1349 .await
1350 .unwrap();
1351 assert_matches!(only_outcome(&first), RestartOutcome::Restarted { .. });
1352
1353 let second = fx
1354 .translator
1355 .restart_servers(RestartTarget::All)
1356 .await
1357 .unwrap();
1358 let RestartOutcome::Throttled { retry_in_ms } = only_outcome(&second) else {
1359 panic!("expected Throttled, got {second:?}");
1360 };
1361 assert!((1..=5000).contains(retry_in_ms));
1362
1363 fx.clock.advance(RESTART_COOLDOWN + Duration::from_secs(1));
1364 let third = fx
1365 .translator
1366 .restart_servers(RestartTarget::All)
1367 .await
1368 .unwrap();
1369 assert_matches!(only_outcome(&third), RestartOutcome::Restarted { .. });
1370 }
1371
1372 #[tokio::test]
1373 async fn test_concurrent_restarts_of_one_server_coalesce() {
1374 let (fx, log) = protocol_fixture(None).await;
1375
1376 let (a, b) = tokio::join!(
1377 fx.translator.restart_servers(RestartTarget::All),
1378 fx.translator.restart_servers(RestartTarget::All),
1379 );
1380
1381 let outcomes = [a.unwrap(), b.unwrap()];
1382 let coalesced = outcomes
1383 .iter()
1384 .filter(|result| {
1385 matches!(
1386 only_outcome(result),
1387 RestartOutcome::Restarted {
1388 coalesced: true,
1389 ..
1390 }
1391 )
1392 })
1393 .count();
1394 assert_eq!(coalesced, 1, "{outcomes:?}");
1395 assert_eq!(read_log(&log).matches("started").count(), 2);
1396 }
1397
1398 #[tokio::test]
1399 async fn test_restart_clears_degradation_and_keeps_the_early_diagnostics_push() {
1400 let dir = TempDir::new().unwrap();
1401 let file = dunce::canonicalize(dir.path()).unwrap().join("main.rs");
1402 fs::write(&file, "fn main() {}").unwrap();
1403 let uri = crate::bridge::try_path_to_uri(&file).unwrap();
1404 let log = dir.path().join("server.log");
1405 let script = write_protocol_server_script(dir.path(), &log, Some(uri.as_ref()));
1406 let fx = fixture(dir, &script, true).await;
1407 {
1408 let mut cache = fx.cache.lock().await;
1409 cache.mark_push_degraded(&fx.id);
1410 cache.clear_server_diagnostics(&fx.id);
1411 }
1412
1413 let result = fx
1414 .translator
1415 .restart_servers(RestartTarget::All)
1416 .await
1417 .unwrap();
1418
1419 assert_matches!(
1420 only_outcome(&result),
1421 RestartOutcome::Restarted {
1422 push_notifications_degraded: false,
1423 ..
1424 },
1425 "{result:?}"
1426 );
1427 tokio::time::timeout(Duration::from_secs(5), async {
1428 while !fx.cache.lock().await.has_diagnostics(&uri) {
1429 tokio::time::sleep(Duration::from_millis(20)).await;
1430 }
1431 })
1432 .await
1433 .expect("the diagnostics pushed during the new handshake must reach the cache");
1434 }
1435
1436 #[tokio::test]
1437 async fn test_restart_fails_requests_in_flight_with_server_restarted() {
1438 let (fx, _log) = protocol_fixture(None).await;
1439 let client = lock_std(&fx.translator.servers).client(&fx.id).unwrap();
1440 let in_flight = tokio::spawn(async move {
1441 client
1442 .request::<_, serde_json::Value>(
1443 "textDocument/hover",
1444 serde_json::json!({}),
1445 Duration::from_secs(30),
1446 )
1447 .await
1448 });
1449 tokio::time::sleep(Duration::from_millis(300)).await;
1450
1451 fx.translator
1452 .restart_servers(RestartTarget::All)
1453 .await
1454 .unwrap();
1455
1456 let failed = tokio::time::timeout(Duration::from_secs(5), in_flight)
1457 .await
1458 .expect("the in-flight request must fail promptly")
1459 .unwrap();
1460 assert_matches!(&failed, Err(Error::ServerRestarted { server_id }) if *server_id == fx.id, "{failed:?}");
1461 }
1462
1463 #[tokio::test]
1464 async fn test_restart_during_shutdown_starts_nothing() {
1465 let (fx, log) = protocol_fixture(None).await;
1466 fx.translator.begin_shutdown();
1467
1468 let result = fx
1469 .translator
1470 .restart_servers(RestartTarget::All)
1471 .await
1472 .unwrap();
1473
1474 assert_eq!(
1475 *only_outcome(&result),
1476 RestartOutcome::Failed {
1477 reason: RestartFailure::ShuttingDown
1478 }
1479 );
1480 assert_eq!(read_log(&log).matches("started").count(), 1);
1481 assert!(fx.registered());
1482 }
1483
1484 #[tokio::test]
1485 async fn test_shutdown_beginning_while_the_old_server_stops_prevents_the_respawn() {
1486 let dir = TempDir::new().unwrap();
1487 let script = write_responder_script(dir.path(), 600);
1488 let fx = fixture(dir, &script, true).await;
1489
1490 let translator = Arc::clone(&fx.translator);
1491 let task =
1492 tokio::spawn(async move { translator.restart_servers(RestartTarget::All).await });
1493 fx.wait_until_deregistered().await;
1494 fx.translator.begin_shutdown();
1495
1496 let result = task.await.unwrap().unwrap();
1497
1498 assert_eq!(
1499 *only_outcome(&result),
1500 RestartOutcome::Failed {
1501 reason: RestartFailure::ShuttingDown
1502 }
1503 );
1504 assert!(fx.registered() && !fx.expected() && fx.is_dead());
1505 }
1506
1507 #[tokio::test]
1508 async fn test_failed_replacement_leaves_the_server_registered_and_reports_spawn_failed() {
1509 let dir = TempDir::new().unwrap();
1510 let log = dir.path().join("server.log");
1511 let script = write_protocol_server_script(dir.path(), &log, None);
1512 let fx = fixture(dir, &script, true).await;
1513 let mut broken = stub_server_config("rust", &script);
1514 broken.server_config_mut().command =
1515 ServerCommand::from_static("mcpls-test-missing-server").into();
1516 lock_std(&fx.translator.servers)
1517 .server_mut(&fx.id)
1518 .unwrap()
1519 .set_init_config(broken);
1520
1521 let result = fx
1522 .translator
1523 .restart_servers(RestartTarget::All)
1524 .await
1525 .unwrap();
1526
1527 assert_matches!(
1528 only_outcome(&result),
1529 RestartOutcome::Failed {
1530 reason: RestartFailure::SpawnFailed { .. }
1531 },
1532 "{result:?}"
1533 );
1534 assert!(fx.registered() && !fx.expected() && fx.is_dead());
1535 }
1536
1537 #[tokio::test]
1538 async fn test_take_leaves_the_server_expected_in_the_same_step_and_restore_runs_it_again() {
1539 let config = crate::config::LspServerConfig::rust_analyzer();
1540 let id = config.id();
1541 let translator = Translator::new();
1542 translator.register_server_complete(crate::lsp::fake_lsp_server_with_config(config));
1543
1544 let held = Deregistered::take(&translator, &id).unwrap();
1545 {
1546 let servers = lock_std(&translator.servers);
1547 assert!(
1548 servers.is_expected(&id),
1549 "the id must never be in neither state"
1550 );
1551 assert!(servers.server(&id).is_none());
1552 drop(servers);
1553 }
1554 drop(held);
1555
1556 let servers = lock_std(&translator.servers);
1557 assert!(!servers.is_expected(&id));
1558 assert!(servers.server(&id).is_some());
1559 drop(servers);
1560 }
1561
1562 #[tokio::test]
1566 async fn test_overlapping_requests_over_more_than_four_servers_restart_each_once() {
1567 let dir = TempDir::new().unwrap();
1568 let cache = Arc::new(Mutex::new(NotificationCache::new()));
1569 let mut translator = Translator::new().with_notification_cache(Arc::clone(&cache));
1570 translator.set_workspace_roots(
1571 WorkspaceRoots::from_paths(&[dir.path().to_path_buf()]).unwrap(),
1572 );
1573 let wiring = PumpWiring::new(
1574 PumpShared {
1575 roles: crate::runtime::pump::DiagnosticsRoles::default(),
1576 notification_cache: Arc::clone(&cache),
1577 subs: SubscriptionRegistry::new(),
1578 workspace_roots: translator.workspace_roots.clone(),
1579 },
1580 CancellationToken::new(),
1581 );
1582 let mut logs = Vec::new();
1583 for index in 0..6 {
1584 let name = format!("s{index}");
1585 let log = dir.path().join(format!("{name}.log"));
1586 let script_dir = dir.path().join(&name);
1587 fs::create_dir(&script_dir).unwrap();
1588 let script = write_protocol_server_script(&script_dir, &log, None);
1589 let mut server = crate::lsp::LspServer::spawn(stub_server_config(&name, &script))
1590 .await
1591 .unwrap();
1592 let receivers = NotificationReceivers {
1593 notifications: server.take_notification_inbox(),
1594 lifecycle: server.take_lifecycle_rx(),
1595 pinned_tsserver: None,
1596 };
1597 translator.register_server_complete(server);
1598 let id = ServerId::new(name.as_str()).unwrap();
1599 let pump = wiring.spawn_pump(id.clone(), receivers, DiagnosticsRole::Secondary);
1600 translator.set_notification_task(&id, pump);
1601 logs.push(log);
1602 }
1603 translator.install_wiring(Arc::new(wiring));
1604
1605 let (a, b) = tokio::join!(
1606 translator.restart_servers(RestartTarget::All),
1607 translator.restart_servers(RestartTarget::All),
1608 );
1609 let (a, b) = (a.unwrap(), b.unwrap());
1610
1611 assert_eq!((a.servers.len(), b.servers.len()), (6, 6));
1612 for (first, second) in a.servers.iter().zip(&b.servers) {
1613 let coalesced = |entry: &ServerRestartEntry| match entry.outcome {
1614 RestartOutcome::Restarted { coalesced, .. } => coalesced,
1615 ref other => panic!("{}: unexpected outcome {other:?}", entry.server_id),
1616 };
1617 assert_ne!(
1618 coalesced(first),
1619 coalesced(second),
1620 "{}: exactly one request restarts it",
1621 first.server_id
1622 );
1623 }
1624 for log in logs {
1625 assert_eq!(read_log(&log).matches("started").count(), 2, "{log:?}");
1626 }
1627 }
1628
1629 #[tokio::test]
1630 async fn test_an_automatic_respawn_does_not_count_towards_coalescing() {
1631 let dir = TempDir::new().unwrap();
1632 let script = write_crash_after_init_script(dir.path());
1633 let fx = fixture(dir, &script, true).await;
1634 tokio::time::timeout(Duration::from_secs(3), async {
1635 while !fx.is_dead() {
1636 tokio::time::sleep(Duration::from_millis(20)).await;
1637 }
1638 })
1639 .await
1640 .expect("the seed server crashes after initialize");
1641 let before = fx.translator.restart_generation(&fx.id);
1642 fx.translator.respawn_if_dead(&fx.id).await.unwrap();
1643 assert_eq!(fx.translator.restart_generation(&fx.id), before);
1644
1645 let result = fx
1646 .translator
1647 .restart_servers(RestartTarget::All)
1648 .await
1649 .unwrap();
1650
1651 assert_matches!(
1652 only_outcome(&result),
1653 RestartOutcome::Restarted {
1654 coalesced: false,
1655 ..
1656 },
1657 "{result:?}"
1658 );
1659 }
1660
1661 #[tokio::test]
1662 async fn test_an_unregistered_server_is_never_throttled() {
1663 let id = ServerId::from_static("rust");
1664 let clock = Arc::new(FakeClock::new());
1665 let translator = Translator::new().with_clock(clock.clone());
1666 translator.set_expected_servers(HashSet::from([id.clone()]));
1667 lock_std(&translator.servers)
1668 .get_mut(&id)
1669 .unwrap()
1670 .restart
1671 .last_attempt = Some(clock.now());
1672 let starting = translator
1673 .restart_servers(server_ids(&["rust"]))
1674 .await
1675 .unwrap();
1676 assert_eq!(*only_outcome(&starting), RestartOutcome::Initializing);
1677
1678 translator.clear_expected_servers();
1679 translator.record_startup_failures(&[ServerSpawnFailure {
1680 server_id: id,
1681 language_id: LanguageId::from_static("rust"),
1682 command: ServerCommand::from_static("missing"),
1683 reason: StartupFailure::InitTaskPanicked,
1684 }]);
1685 let failed = translator
1686 .restart_servers(server_ids(&["rust"]))
1687 .await
1688 .unwrap();
1689 assert_matches!(
1690 only_outcome(&failed),
1691 RestartOutcome::NotRunning { .. },
1692 "{failed:?}"
1693 );
1694 }
1695
1696 #[tokio::test]
1697 async fn test_restart_after_an_init_panic_without_wiring_reports_not_running() {
1698 let dir = TempDir::new().unwrap();
1699 let log = dir.path().join("server.log");
1700 let script = write_protocol_server_script(dir.path(), &log, None);
1701 let fx = fixture(dir, &script, false).await;
1702 fx.translator.settle_after_init_panic(&[]).await;
1703
1704 let result = fx
1705 .translator
1706 .restart_servers(RestartTarget::All)
1707 .await
1708 .unwrap();
1709
1710 assert_matches!(
1711 only_outcome(&result),
1712 RestartOutcome::NotRunning { .. },
1713 "{result:?}"
1714 );
1715 }
1716
1717 #[tokio::test]
1718 async fn test_restart_after_an_init_panic_with_wiring_installed_reports_not_running() {
1719 let dir = TempDir::new().unwrap();
1720 let log = dir.path().join("server.log");
1721 let script = write_protocol_server_script(dir.path(), &log, None);
1722 let fx = fixture(dir, &script, true).await;
1723 fx.translator.settle_after_init_panic(&[]).await;
1724
1725 let result = fx
1726 .translator
1727 .restart_servers(RestartTarget::All)
1728 .await
1729 .unwrap();
1730
1731 assert_matches!(
1732 only_outcome(&result),
1733 RestartOutcome::NotRunning { .. },
1734 "{result:?}"
1735 );
1736 }
1737
1738 #[tokio::test]
1739 async fn test_a_failed_restart_invalidates_the_stopped_servers_cached_diagnostics() {
1740 let dir = TempDir::new().unwrap();
1741 let file = dunce::canonicalize(dir.path()).unwrap().join("main.rs");
1742 fs::write(&file, "fn main() {}").unwrap();
1743 let uri = crate::bridge::try_path_to_uri(&file).unwrap();
1744 let log = dir.path().join("server.log");
1745 let script = write_protocol_server_script(dir.path(), &log, Some(uri.as_ref()));
1746 let fx = fixture(dir, &script, true).await;
1747 tokio::time::timeout(Duration::from_secs(5), async {
1748 while !fx.cache.lock().await.has_diagnostics(&uri) {
1749 tokio::time::sleep(Duration::from_millis(20)).await;
1750 }
1751 })
1752 .await
1753 .expect("the seed server's diagnostics reach the cache");
1754 let mut broken = stub_server_config("rust", &script);
1755 broken.server_config_mut().command =
1756 ServerCommand::from_static("mcpls-test-missing-server").into();
1757 lock_std(&fx.translator.servers)
1758 .server_mut(&fx.id)
1759 .unwrap()
1760 .set_init_config(broken);
1761
1762 let result = fx
1763 .translator
1764 .restart_servers(RestartTarget::All)
1765 .await
1766 .unwrap();
1767
1768 assert_matches!(only_outcome(&result), RestartOutcome::Failed { .. });
1769 let (stale, degraded) = {
1770 let cache = fx.cache.lock().await;
1771 (cache.has_diagnostics(&uri), cache.is_push_degraded(&fx.id))
1772 };
1773 assert!(!stale, "stale diagnostics must go");
1774 assert!(degraded);
1775 assert!(
1776 lock_std(&fx.translator.servers)
1777 .get(&fx.id)
1778 .is_none_or(|slot| slot.notification_task.is_none()),
1779 "the dead server's pump handle must be dropped"
1780 );
1781 }
1782
1783 #[tokio::test]
1784 async fn test_cooldown_counts_from_the_completion_of_a_slow_restart() {
1785 let dir = TempDir::new().unwrap();
1786 let script = write_responder_script(dir.path(), 600);
1787 let fx = fixture(dir, &script, true).await;
1788 let translator = Arc::clone(&fx.translator);
1789 let restart =
1790 tokio::spawn(async move { translator.restart_servers(RestartTarget::All).await });
1791 fx.wait_until_deregistered().await;
1792 fx.clock.advance(Duration::from_secs(30));
1793
1794 restart.await.unwrap().unwrap();
1795 let again = fx
1796 .translator
1797 .restart_servers(RestartTarget::All)
1798 .await
1799 .unwrap();
1800
1801 assert_matches!(
1802 only_outcome(&again),
1803 RestartOutcome::Throttled { .. },
1804 "{again:?}"
1805 );
1806 }
1807
1808 #[tokio::test]
1809 async fn test_a_request_queued_behind_a_discarded_restart_is_not_reported_coalesced() {
1810 let dir = TempDir::new().unwrap();
1811 let script = write_responder_script(dir.path(), 600);
1812 let fx = fixture(dir, &script, true).await;
1813 let (first_tx, second_tx) = (Arc::clone(&fx.translator), Arc::clone(&fx.translator));
1814 let first =
1815 tokio::spawn(async move { first_tx.restart_servers(RestartTarget::All).await });
1816 fx.wait_until_deregistered().await;
1817 let second =
1818 tokio::spawn(async move { second_tx.restart_servers(RestartTarget::All).await });
1819 tokio::time::sleep(Duration::from_millis(100)).await;
1820 fx.translator.begin_shutdown();
1821
1822 for result in [
1823 first.await.unwrap().unwrap(),
1824 second.await.unwrap().unwrap(),
1825 ] {
1826 assert_eq!(
1827 *only_outcome(&result),
1828 RestartOutcome::Failed {
1829 reason: RestartFailure::ShuttingDown
1830 }
1831 );
1832 }
1833 }
1834
1835 #[tokio::test]
1838 async fn test_a_healthy_server_slow_to_exit_is_not_killed_at_the_handshake_budget() {
1839 let dir = TempDir::new().unwrap();
1840 let log = dir.path().join("server.log");
1841 let delay = u32::try_from(RESTART_HANDSHAKE_BUDGET.as_secs()).unwrap() + 1;
1842 let script = write_slow_exit_server_script(dir.path(), &log, delay);
1843 let fx = fixture(dir, &script, true).await;
1844
1845 let result = fx
1846 .translator
1847 .restart_servers(RestartTarget::All)
1848 .await
1849 .unwrap();
1850
1851 assert_matches!(
1852 only_outcome(&result),
1853 RestartOutcome::Restarted { .. },
1854 "{result:?}"
1855 );
1856 assert!(
1857 read_log(&log).contains("exiting"),
1858 "the old server must have exited by itself, not been killed"
1859 );
1860 }
1861
1862 #[tokio::test]
1863 async fn test_restart_recovers_a_crashed_server() {
1864 let dir = TempDir::new().unwrap();
1865 let script = write_crash_after_init_script(dir.path());
1866 let fx = fixture(dir, &script, true).await;
1867 tokio::time::timeout(Duration::from_secs(3), async {
1868 while !fx.is_dead() {
1869 tokio::time::sleep(Duration::from_millis(20)).await;
1870 }
1871 })
1872 .await
1873 .expect("the seed server crashes after initialize");
1874
1875 let result = fx
1876 .translator
1877 .restart_servers(RestartTarget::All)
1878 .await
1879 .unwrap();
1880
1881 assert_matches!(
1882 only_outcome(&result),
1883 RestartOutcome::Restarted { .. },
1884 "{result:?}"
1885 );
1886 }
1887
1888 #[tokio::test]
1892 async fn test_restart_of_a_signalling_server_gates_until_its_signal() {
1893 let (fx, _log) = protocol_fixture(None).await;
1894 fx.cache.lock().await.observe_indexing_signal(
1895 &fx.id,
1896 "experimental/serverStatus",
1897 Some(&serde_json::json!({"quiescent": true})),
1898 );
1899
1900 let result = fx
1901 .translator
1902 .restart_servers(RestartTarget::All)
1903 .await
1904 .unwrap();
1905 assert_matches!(
1906 only_outcome(&result),
1907 RestartOutcome::Restarted {
1908 indexing_state: IndexingState::Loading,
1909 ..
1910 },
1911 "{result:?}"
1912 );
1913
1914 let translator = Arc::clone(&fx.translator);
1915 let id = fx.id.clone();
1916 let gated = tokio::spawn(async move { translator.wait_for_indexing_ready(&id).await });
1917 tokio::time::sleep(Duration::from_millis(150)).await;
1918 assert!(!gated.is_finished(), "the gate must hold before the signal");
1919
1920 fx.cache.lock().await.observe_indexing_signal(
1921 &fx.id,
1922 "experimental/serverStatus",
1923 Some(&serde_json::json!({"quiescent": true})),
1924 );
1925 tokio::time::timeout(Duration::from_secs(5), gated)
1926 .await
1927 .expect("the gate releases once the replacement is ready")
1928 .unwrap()
1929 .unwrap();
1930 }
1931
1932 #[tokio::test]
1935 async fn test_restart_of_a_silent_server_reports_unknown() {
1936 let (fx, _log) = protocol_fixture(None).await;
1937
1938 let result = fx
1939 .translator
1940 .restart_servers(RestartTarget::All)
1941 .await
1942 .unwrap();
1943
1944 assert_matches!(
1945 only_outcome(&result),
1946 RestartOutcome::Restarted {
1947 indexing_state: IndexingState::Unknown,
1948 ..
1949 },
1950 "{result:?}"
1951 );
1952 fx.translator.wait_for_indexing_ready(&fx.id).await.unwrap();
1953 }
1954 }
1955}
1956
1957#[cfg(test)]
1958mod server_text_tests {
1959 use super::*;
1960
1961 #[test]
1962 fn test_restart_messages_are_redacted() {
1963 let secret = "SuperSecretValue123";
1964 let set = Redactions::new([("API_TOKEN".to_owned(), secret.to_owned())]);
1965 let mut result = RestartServerResult {
1966 servers: vec![
1967 ServerRestartEntry {
1968 server_id: ServerId::from_static("a"),
1969 outcome: RestartOutcome::Failed {
1970 reason: RestartFailure::SpawnFailed {
1971 message: format!("spawn {secret}"),
1972 },
1973 },
1974 },
1975 ServerRestartEntry {
1976 server_id: ServerId::from_static("b"),
1977 outcome: RestartOutcome::Failed {
1978 reason: RestartFailure::InitializeFailed {
1979 message: format!("init {secret}"),
1980 },
1981 },
1982 },
1983 ServerRestartEntry {
1984 server_id: ServerId::from_static("c"),
1985 outcome: RestartOutcome::NotRunning {
1986 message: format!("gone {secret}"),
1987 },
1988 },
1989 ],
1990 };
1991
1992 result.redact_server_text(&set);
1993
1994 let json = serde_json::to_string(&result).unwrap();
1995 assert!(!json.contains(secret), "{json}");
1996 assert_eq!(json.matches("[redacted:API_TOKEN]").count(), 3, "{json}");
1997 }
1998}