1use lsp_types::{notification::Notification as Notif, request::Request as Req, *};
2
3use super::*;
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
65 .insert(cmd, Box::new(move |s, _req_id, args| handler(s, args)));
66 self
67 }
68
69 pub fn with_command<R: Serialize + 'static>(
71 mut self,
72 cmd: &'static str,
73 handler: AsyncHandler<Args::S, Vec<JsonValue>, R>,
74 ) -> Self {
75 self.command_handlers.insert(
76 cmd,
77 Box::new(move |s, _req_id, req| erased_response(handler(s, req))),
78 );
79 self
80 }
81
82 pub fn with_notification_<R: Notif>(
84 mut self,
85 handler: PureHandler<Args::S, JsonValue>,
86 ) -> Self {
87 self.notif_handlers.insert(R::METHOD, Box::new(handler));
88 self
89 }
90
91 pub fn with_notification<R: Notif>(mut self, handler: PureHandler<Args::S, R::Params>) -> Self {
93 self.notif_handlers.insert(
94 R::METHOD,
95 Box::new(move |s, req| handler(s, from_json(req)?)),
96 );
97 self
98 }
99
100 pub fn with_raw_request<R: Req>(mut self, handler: RawHandler<Args::S, JsonValue>) -> Self {
103 self.req_handlers
104 .insert(R::METHOD, Box::new(move |s, _req_id, req| handler(s, req)));
105 self
106 }
107
108 pub fn with_request_<R: Req>(
112 mut self,
113 handler: fn(&mut Args::S, R::Params) -> ScheduleResult,
114 ) -> Self {
115 self.req_handlers.insert(
116 R::METHOD,
117 Box::new(move |s, _req_id, req| handler(s, from_json(req)?)),
118 );
119 self
120 }
121
122 pub fn with_request<R: Req>(
124 mut self,
125 handler: AsyncHandler<Args::S, R::Params, R::Result>,
126 ) -> Self {
127 self.req_handlers.insert(
128 R::METHOD,
129 Box::new(move |s, _req_id, req| erased_response(handler(s, from_json(req)?))),
130 );
131 self
132 }
133}
134
135impl<Args: Initializer> LsDriver<LspMessage, Args>
136where
137 Args::S: 'static,
138{
139 #[cfg(feature = "system")]
147 pub fn start(
148 &mut self,
149 inbox: TConnectionRx<LspMessage>,
150 is_replay: bool,
151 ) -> anyhow::Result<()> {
152 let res = self.start_(inbox);
153
154 if is_replay {
155 let client = self.client.clone();
156 let _ = std::thread::spawn(move || {
157 let since = tinymist_std::time::Instant::now();
158 let timeout = std::env::var("REPLAY_TIMEOUT")
159 .ok()
160 .and_then(|s| s.parse().ok())
161 .unwrap_or(60);
162 client.handle.block_on(async {
163 while client.has_pending_requests() {
164 if since.elapsed().as_secs() > timeout {
165 log::error!("replay timeout reached, {timeout}s");
166 client.begin_panic();
167 }
168
169 tokio::time::sleep(std::time::Duration::from_millis(10)).await;
170 }
171 })
172 })
173 .join();
174 }
175
176 res
177 }
178
179 #[cfg(feature = "system")]
181 pub fn start_(&mut self, inbox: TConnectionRx<LspMessage>) -> anyhow::Result<()> {
182 use EventOrMessage::*;
183 while let Ok(msg) = inbox.recv() {
204 const EXIT_METHOD: &str = notification::Exit::METHOD;
205 let loop_start = tinymist_std::time::now();
206 match msg {
207 Evt(event) => {
208 let Some(event_handler) = self.events.get(&event.as_ref().type_id()) else {
209 log::warn!("unhandled event: {:?}", event.as_ref().type_id());
210 continue;
211 };
212
213 let s = match &mut self.state {
214 State::Uninitialized(u) => ServiceState::Uninitialized(u.as_deref_mut()),
215 State::Initializing(s) | State::Ready(s) => ServiceState::Ready(s),
216 State::ShuttingDown => {
217 log::warn!("server is shutting down");
218 continue;
219 }
220 };
221
222 event_handler(s, &self.client, event)?;
223 }
224 Msg(LspMessage::Request(req)) => {
225 let client = self.client.clone();
226 let req_id = req.id.clone();
227 client.register_request(&req.method, &req_id, loop_start);
228 let fut = client.schedule_tail(
229 req_id.clone(),
230 self.on_lsp_request(&req.method, req_id, req.params),
231 );
232 self.client.handle.spawn(fut);
233 }
234 Msg(LspMessage::Notification(not)) => {
235 let is_exit = not.method == EXIT_METHOD;
236 let track_id = self
237 .next_not_id
238 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
239 self.client.hook.start_notification(track_id, ¬.method);
240 let result = self.on_notification(¬.method, not.params);
241 self.client
242 .hook
243 .stop_notification(track_id, ¬.method, loop_start, result);
244 if is_exit {
245 return Ok(());
246 }
247 }
248 Msg(LspMessage::Response(resp)) => {
249 let s = match &mut self.state {
250 State::Ready(s) => s,
251 _ => {
252 log::warn!("server is not ready yet");
253 continue;
254 }
255 };
256
257 self.client.clone().complete_lsp_request(s, resp)
258 }
259 }
260 }
261
262 log::warn!("client exited without proper shutdown sequence");
263 Ok(())
264 }
265
266 #[cfg(feature = "web")]
268 pub fn on_server_event(&mut self, event_id: u32) {
269 let evt = match &self.client.sender {
270 TransportHost::Js { events, .. } => events.lock().remove(&event_id),
271 TransportHost::System(_) => {
272 panic!("cannot send server event in system transport");
273 }
274 };
275
276 if let Some(event) = evt {
277 let Some(event_handler) = self.events.get(&event.as_ref().type_id()) else {
278 log::warn!("unhandled event: {:?}", event.as_ref().type_id());
279 return;
280 };
281
282 let s = match &mut self.state {
283 State::Uninitialized(u) => ServiceState::Uninitialized(u.as_deref_mut()),
284 State::Initializing(s) | State::Ready(s) => ServiceState::Ready(s),
285 State::ShuttingDown => {
286 log::warn!("server is shutting down");
287 return;
288 }
289 };
290
291 let res = event_handler(s, &self.client, event);
292 if let Err(err) = res {
293 log::error!("failed to handle server event {event_id}: {err}");
294 }
295 }
296 }
297
298 pub fn on_lsp_request(
301 &mut self,
302 method: &str,
303 req_id: RequestId,
304 params: JsonValue,
305 ) -> ScheduleResult {
306 match (&mut self.state, method) {
307 (State::Uninitialized(args), request::Initialize::METHOD) => {
308 let params = serde_json::from_value::<Args::I>(params);
310 match params {
311 Ok(params) => {
312 let args = args.take().expect("already initialized");
313 let (s, res) = args.initialize(params);
314 self.state = State::Initializing(s);
315 res
316 }
317 Err(e) => just_result(Err(invalid_request(e))),
318 }
319 }
320 (State::Uninitialized(..) | State::Initializing(..), _) => {
321 just_result(Err(not_initialized()))
322 }
323 (_, request::Initialize::METHOD) => {
324 just_result(Err(invalid_request("server is already initialized")))
325 }
326 (State::Ready(..), request::ExecuteCommand::METHOD) => {
328 self.on_execute_command(req_id, params)
329 }
330 (State::Ready(s), method) => 'serve_req: {
331 let is_shutdown = method == request::Shutdown::METHOD;
332
333 let Some(handler) = self.requests.get(method) else {
334 log::warn!("unhandled lsp request: {method}");
335 break 'serve_req just_result(Err(method_not_found()));
336 };
337
338 let resp = handler(s, req_id, params);
339
340 if is_shutdown {
341 self.state = State::ShuttingDown;
342 }
343
344 resp
345 }
346 (State::ShuttingDown, _) => {
347 just_result(Err(invalid_request("server is shutting down")))
348 }
349 }
350 }
351
352 fn on_execute_command(&mut self, req_id: RequestId, params: JsonValue) -> ScheduleResult {
354 let s = self.state.opt_mut().ok_or_else(not_initialized)?;
355
356 let params = from_value::<ExecuteCommandParams>(params)
357 .map_err(|e| invalid_params(e.to_string()))?;
358
359 let ExecuteCommandParams {
360 command, arguments, ..
361 } = params;
362
363 if command == "tinymist.getResources" {
365 self.get_resources(req_id, arguments)
366 } else {
367 let Some(handler) = self.commands.get(command.as_str()) else {
368 log::error!("asked to execute unknown command: {command}");
369 return Err(method_not_found());
370 };
371 handler(s, req_id, arguments)
372 }
373 }
374
375 pub fn on_notification(&mut self, method: &str, params: JsonValue) -> LspResult<()> {
377 let handle = |s, method: &str, params: JsonValue| {
378 let Some(handler) = self.notifications.get(method) else {
379 log::warn!("unhandled notification: {method}");
380 return Ok(());
381 };
382
383 handler(s, params)
384 };
385
386 match (&mut self.state, method) {
387 (state, notification::Initialized::METHOD) => {
388 let mut s = State::ShuttingDown;
389 std::mem::swap(state, &mut s);
390 match s {
391 State::Initializing(s) => {
392 *state = State::Ready(s);
393 }
394 _ => {
395 std::mem::swap(state, &mut s);
396 }
397 }
398
399 let s = match state {
400 State::Ready(s) => s,
401 _ => {
402 log::warn!("server is not ready yet");
403 return Ok(());
404 }
405 };
406 handle(s, method, params)
407 }
408 (State::Ready(state), method) => handle(state, method, params),
409 (State::Uninitialized(..) | State::Initializing(..), method) => {
411 log::warn!("server is not ready yet, while received notification {method}");
412 Ok(())
413 }
414 (State::ShuttingDown, method) => {
415 log::warn!("server is shutting down, while received notification {method}");
416 Ok(())
417 }
418 }
419 }
420
421 pub fn on_lsp_response(&mut self, resp: lsp::Response) {
423 let client = self.client.clone();
424 let Some(s) = self.state_mut() else {
425 log::warn!("server is not ready yet, while received response");
426 return;
427 };
428
429 client.complete_lsp_request(s, resp)
430 }
431}