async_lsp/
server.rs

1//! Language Server lifecycle.
2//!
3//! *Only applies to Language Servers.*
4//!
5//! This middleware handles
6//! [the lifecycle of Language Servers](https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#lifeCycleMessages),
7//! specifically:
8//! - Exit the main loop with `ControlFlow::Break(Ok(()))` on `exit` notification.
9//! - Responds unrelated requests with errors and ignore unrelated notifications during
10//!   initialization and shutting down.
11use std::future::{ready, Future, Ready};
12use std::ops::ControlFlow;
13use std::pin::Pin;
14use std::task::{Context, Poll};
15
16use futures::future::Either;
17use lsp_types::notification::{self, Notification};
18use lsp_types::request::{self, Request};
19use pin_project_lite::pin_project;
20use tower_layer::Layer;
21use tower_service::Service;
22
23use crate::{
24    AnyEvent, AnyNotification, AnyRequest, Error, ErrorCode, LspService, ResponseError, Result,
25};
26
27#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
28enum State {
29    #[default]
30    Uninitialized,
31    Initializing,
32    Ready,
33    ShuttingDown,
34}
35
36/// The middleware handling Language Server lifecycle.
37///
38/// See [module level documentations](self) for details.
39#[derive(Debug, Default)]
40pub struct Lifecycle<S> {
41    service: S,
42    state: State,
43}
44
45define_getters!(impl[S] Lifecycle<S>, service: S);
46
47impl<S> Lifecycle<S> {
48    /// Creating the `Lifecycle` middleware in uninitialized state.
49    #[must_use]
50    pub fn new(service: S) -> Self {
51        Self {
52            service,
53            state: State::Uninitialized,
54        }
55    }
56}
57
58impl<S: LspService> Service<AnyRequest> for Lifecycle<S>
59where
60    S::Error: From<ResponseError>,
61{
62    type Response = S::Response;
63    type Error = S::Error;
64    type Future = ResponseFuture<S::Future>;
65
66    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
67        self.service.poll_ready(cx)
68    }
69
70    fn call(&mut self, req: AnyRequest) -> Self::Future {
71        let inner = match (self.state, &*req.method) {
72            (State::Uninitialized, request::Initialize::METHOD) => {
73                self.state = State::Initializing;
74                Either::Left(self.service.call(req))
75            }
76            (State::Uninitialized | State::Initializing, _) => {
77                Either::Right(ready(Err(ResponseError {
78                    code: ErrorCode::SERVER_NOT_INITIALIZED,
79                    message: "Server is not initialized yet".into(),
80                    data: None,
81                }
82                .into())))
83            }
84            (_, request::Initialize::METHOD) => Either::Right(ready(Err(ResponseError {
85                code: ErrorCode::INVALID_REQUEST,
86                message: "Server is already initialized".into(),
87                data: None,
88            }
89            .into()))),
90            (State::Ready, _) => {
91                if req.method == request::Shutdown::METHOD {
92                    self.state = State::ShuttingDown;
93                }
94                Either::Left(self.service.call(req))
95            }
96            (State::ShuttingDown, _) => Either::Right(ready(Err(ResponseError {
97                code: ErrorCode::INVALID_REQUEST,
98                message: "Server is shutting down".into(),
99                data: None,
100            }
101            .into()))),
102        };
103        ResponseFuture { inner }
104    }
105}
106
107impl<S: LspService> LspService for Lifecycle<S>
108where
109    S::Error: From<ResponseError>,
110{
111    fn notify(&mut self, notif: AnyNotification) -> ControlFlow<Result<()>> {
112        match &*notif.method {
113            notification::Exit::METHOD => {
114                self.service.notify(notif)?;
115                ControlFlow::Break(Ok(()))
116            }
117            notification::Initialized::METHOD => {
118                if self.state != State::Initializing {
119                    return ControlFlow::Break(Err(Error::Protocol(format!(
120                        "Unexpected initialized notification on state {:?}",
121                        self.state
122                    ))));
123                }
124                self.state = State::Ready;
125                self.service.notify(notif)?;
126                ControlFlow::Continue(())
127            }
128            _ => self.service.notify(notif),
129        }
130    }
131
132    fn emit(&mut self, event: AnyEvent) -> ControlFlow<Result<()>> {
133        self.service.emit(event)
134    }
135}
136
137pin_project! {
138    /// The [`Future`] type used by the [`Lifecycle`] middleware.
139    pub struct ResponseFuture<Fut: Future> {
140        #[pin]
141        inner: Either<Fut, Ready<Fut::Output>>,
142    }
143}
144
145impl<Fut: Future> Future for ResponseFuture<Fut> {
146    type Output = Fut::Output;
147
148    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
149        self.project().inner.poll(cx)
150    }
151}
152
153/// A [`tower_layer::Layer`] which builds [`Lifecycle`].
154#[must_use]
155#[derive(Clone, Default)]
156pub struct LifecycleLayer {
157    _private: (),
158}
159
160impl<S> Layer<S> for LifecycleLayer {
161    type Service = Lifecycle<S>;
162
163    fn layer(&self, inner: S) -> Self::Service {
164        Lifecycle::new(inner)
165    }
166}