1use super::*;
2
3use lsp_types::{notification::Notification as Notif, request::Request as Req, *};
4
5type PureHandler<S, T> = fn(srv: &mut S, args: T) -> LspResult<()>;
6
7impl<S: 'static> TypedLspClient<S> {
8 pub fn send_lsp_request<R: Req>(
11 &self,
12 params: R::Params,
13 handler: impl FnOnce(&mut S, lsp::Response) + Send + Sync + 'static,
14 ) {
15 let caster = self.caster.clone();
16 self.client
17 .send_lsp_request_::<R>(params, move |s, resp| handler(caster(s), resp))
18 }
19}
20
21impl LspClient {
22 pub fn send_lsp_request_<R: Req>(
24 &self,
25 params: R::Params,
26 handler: impl FnOnce(&mut dyn Any, lsp::Response) + Send + Sync + 'static,
27 ) {
28 let mut req_queue = self.req_queue.lock();
29 let request = req_queue.outgoing.register(
30 R::METHOD.to_owned(),
31 params,
32 Box::new(|s, resp| handler(s, resp.try_into().unwrap())),
33 );
34
35 self.sender.send_message(request.into());
36 }
37
38 pub fn respond_lsp(&self, response: lsp::Response) {
40 self.respond(response.id.clone(), response.into())
41 }
42
43 pub fn send_notification<N: Notif>(&self, params: &N::Params) {
45 self.send_notification_(lsp::Notification::new(N::METHOD.to_owned(), params));
46 }
47
48 pub fn send_notification_(&self, notif: lsp::Notification) {
50 self.sender.send_message(notif.into());
51 }
52}
53
54impl<Args: Initializer> LsBuilder<LspMessage, Args>
55where
56 Args::S: 'static,
57{
58 pub fn with_command_(
60 mut self,
61 cmd: &'static str,
62 handler: RawHandler<Args::S, Vec<JsonValue>>,
63 ) -> Self {
64 self.command_handlers.insert(cmd, Box::new(handler));
65 self
66 }
67
68 pub fn with_command<R: Serialize + 'static>(
70 mut self,
71 cmd: &'static str,
72 handler: AsyncHandler<Args::S, Vec<JsonValue>, R>,
73 ) -> Self {
74 self.command_handlers.insert(
75 cmd,
76 Box::new(move |s, req| erased_response(handler(s, req))),
77 );
78 self
79 }
80
81 pub fn with_notification_<R: Notif>(
83 mut self,
84 handler: PureHandler<Args::S, JsonValue>,
85 ) -> Self {
86 self.notif_handlers.insert(R::METHOD, Box::new(handler));
87 self
88 }
89
90 pub fn with_notification<R: Notif>(mut self, handler: PureHandler<Args::S, R::Params>) -> Self {
92 self.notif_handlers.insert(
93 R::METHOD,
94 Box::new(move |s, req| handler(s, from_json(req)?)),
95 );
96 self
97 }
98
99 pub fn with_raw_request<R: Req>(mut self, handler: RawHandler<Args::S, JsonValue>) -> Self {
102 self.req_handlers.insert(R::METHOD, Box::new(handler));
103 self
104 }
105
106 pub fn with_request_<R: Req>(
110 mut self,
111 handler: fn(&mut Args::S, R::Params) -> ScheduleResult,
112 ) -> Self {
113 self.req_handlers.insert(
114 R::METHOD,
115 Box::new(move |s, req| handler(s, from_json(req)?)),
116 );
117 self
118 }
119
120 pub fn with_request<R: Req>(
122 mut self,
123 handler: AsyncHandler<Args::S, R::Params, R::Result>,
124 ) -> Self {
125 self.req_handlers.insert(
126 R::METHOD,
127 Box::new(move |s, req| erased_response(handler(s, from_json(req)?))),
128 );
129 self
130 }
131}
132
133impl<Args: Initializer> LsDriver<LspMessage, Args>
134where
135 Args::S: 'static,
136{
137 #[cfg(feature = "system")]
145 pub fn start(
146 &mut self,
147 inbox: TConnectionRx<LspMessage>,
148 is_replay: bool,
149 ) -> anyhow::Result<()> {
150 let res = self.start_(inbox);
151
152 if is_replay {
153 let client = self.client.clone();
154 let _ = std::thread::spawn(move || {
155 let since = tinymist_std::time::Instant::now();
156 let timeout = std::env::var("REPLAY_TIMEOUT")
157 .ok()
158 .and_then(|s| s.parse().ok())
159 .unwrap_or(60);
160 client.handle.block_on(async {
161 while client.has_pending_requests() {
162 if since.elapsed().as_secs() > timeout {
163 log::error!("replay timeout reached, {timeout}s");
164 client.begin_panic();
165 }
166
167 tokio::time::sleep(std::time::Duration::from_millis(10)).await;
168 }
169 })
170 })
171 .join();
172 }
173
174 res
175 }
176
177 #[cfg(feature = "system")]
179 pub fn start_(&mut self, inbox: TConnectionRx<LspMessage>) -> anyhow::Result<()> {
180 use EventOrMessage::*;
181 while let Ok(msg) = inbox.recv() {
202 const EXIT_METHOD: &str = notification::Exit::METHOD;
203 let loop_start = Instant::now();
204 match msg {
205 Evt(event) => {
206 let Some(event_handler) = self.events.get(&event.as_ref().type_id()) else {
207 log::warn!("unhandled event: {:?}", event.as_ref().type_id());
208 continue;
209 };
210
211 let s = match &mut self.state {
212 State::Uninitialized(u) => ServiceState::Uninitialized(u.as_deref_mut()),
213 State::Initializing(s) | State::Ready(s) => ServiceState::Ready(s),
214 State::ShuttingDown => {
215 log::warn!("server is shutting down");
216 continue;
217 }
218 };
219
220 event_handler(s, &self.client, event)?;
221 }
222 Msg(LspMessage::Request(req)) => {
223 let client = self.client.clone();
224 let req_id = req.id.clone();
225 client.register_request(&req.method, &req_id, loop_start);
226 let fut =
227 client.schedule_tail(req_id, self.on_lsp_request(&req.method, req.params));
228 self.client.handle.spawn(fut);
229 }
230 Msg(LspMessage::Notification(not)) => {
231 let is_exit = not.method == EXIT_METHOD;
232 self.client.hook.start_notification(¬.method);
233 let result = self.on_notification(¬.method, not.params);
234 self.client
235 .hook
236 .stop_notification(¬.method, loop_start, result);
237 if is_exit {
238 return Ok(());
239 }
240 }
241 Msg(LspMessage::Response(resp)) => {
242 let s = match &mut self.state {
243 State::Ready(s) => s,
244 _ => {
245 log::warn!("server is not ready yet");
246 continue;
247 }
248 };
249
250 self.client.clone().complete_lsp_request(s, resp)
251 }
252 }
253 }
254
255 log::warn!("client exited without proper shutdown sequence");
256 Ok(())
257 }
258
259 #[cfg(feature = "web")]
261 pub fn on_server_event(&mut self, event_id: u32) {
262 let evt = match &self.client.sender {
263 TransportHost::Js(sender) => sender.events.lock().remove(&event_id),
264 TransportHost::System(_) => {
265 panic!("cannot send server event in system transport");
266 }
267 };
268
269 if let Some(event) = evt {
270 let Some(event_handler) = self.events.get(&event.as_ref().type_id()) else {
271 log::warn!("unhandled event: {:?}", event.as_ref().type_id());
272 return;
273 };
274
275 let s = match &mut self.state {
276 State::Uninitialized(u) => ServiceState::Uninitialized(u.as_deref_mut()),
277 State::Initializing(s) | State::Ready(s) => ServiceState::Ready(s),
278 State::ShuttingDown => {
279 log::warn!("server is shutting down");
280 return;
281 }
282 };
283
284 let res = event_handler(s, &self.client, event);
285 if let Err(err) = res {
286 log::error!("failed to handle server event {event_id}: {err}");
287 }
288 }
289 }
290
291 pub fn on_lsp_request(&mut self, method: &str, params: JsonValue) -> ScheduleResult {
294 match (&mut self.state, method) {
295 (State::Uninitialized(args), request::Initialize::METHOD) => {
296 let params = serde_json::from_value::<Args::I>(params);
298 match params {
299 Ok(params) => {
300 let args = args.take().expect("already initialized");
301 let (s, res) = args.initialize(params);
302 self.state = State::Initializing(s);
303 res
304 }
305 Err(e) => just_result(Err(invalid_request(e))),
306 }
307 }
308 (State::Uninitialized(..) | State::Initializing(..), _) => {
309 just_result(Err(not_initialized()))
310 }
311 (_, request::Initialize::METHOD) => {
312 just_result(Err(invalid_request("server is already initialized")))
313 }
314 (State::Ready(..), request::ExecuteCommand::METHOD) => self.on_execute_command(params),
316 (State::Ready(s), method) => 'serve_req: {
317 let is_shutdown = method == request::Shutdown::METHOD;
318
319 let Some(handler) = self.requests.get(method) else {
320 log::warn!("unhandled lsp request: {method}");
321 break 'serve_req just_result(Err(method_not_found()));
322 };
323
324 let resp = handler(s, params);
325
326 if is_shutdown {
327 self.state = State::ShuttingDown;
328 }
329
330 resp
331 }
332 (State::ShuttingDown, _) => {
333 just_result(Err(invalid_request("server is shutting down")))
334 }
335 }
336 }
337
338 fn on_execute_command(&mut self, params: JsonValue) -> ScheduleResult {
340 let s = self.state.opt_mut().ok_or_else(not_initialized)?;
341
342 let params = from_value::<ExecuteCommandParams>(params)
343 .map_err(|e| invalid_params(e.to_string()))?;
344
345 let ExecuteCommandParams {
346 command, arguments, ..
347 } = params;
348
349 if command == "tinymist.getResources" {
351 self.get_resources(arguments)
352 } else {
353 let Some(handler) = self.commands.get(command.as_str()) else {
354 log::error!("asked to execute unknown command: {command}");
355 return Err(method_not_found());
356 };
357 handler(s, arguments)
358 }
359 }
360
361 pub fn on_notification(&mut self, method: &str, params: JsonValue) -> LspResult<()> {
363 let handle = |s, method: &str, params: JsonValue| {
364 let Some(handler) = self.notifications.get(method) else {
365 log::warn!("unhandled notification: {method}");
366 return Ok(());
367 };
368
369 handler(s, params)
370 };
371
372 match (&mut self.state, method) {
373 (state, notification::Initialized::METHOD) => {
374 let mut s = State::ShuttingDown;
375 std::mem::swap(state, &mut s);
376 match s {
377 State::Initializing(s) => {
378 *state = State::Ready(s);
379 }
380 _ => {
381 std::mem::swap(state, &mut s);
382 }
383 }
384
385 let s = match state {
386 State::Ready(s) => s,
387 _ => {
388 log::warn!("server is not ready yet");
389 return Ok(());
390 }
391 };
392 handle(s, method, params)
393 }
394 (State::Ready(state), method) => handle(state, method, params),
395 (State::Uninitialized(..) | State::Initializing(..), method) => {
397 log::warn!("server is not ready yet, while received notification {method}");
398 Ok(())
399 }
400 (State::ShuttingDown, method) => {
401 log::warn!("server is shutting down, while received notification {method}");
402 Ok(())
403 }
404 }
405 }
406}