auto_lsp_server/
notification_registry.rs

1/*
2This file is part of auto-lsp.
3Copyright (C) 2025 CLAUZEL Adrien
4
5auto-lsp is free software: you can redistribute it and/or modify
6it under the terms of the GNU General Public License as published by
7the Free Software Foundation, either version 3 of the License, or
8(at your option) any later version.
9
10This program is distributed in the hope that it will be useful,
11but WITHOUT ANY WARRANTY; without even the implied warranty of
12MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
13GNU General Public License for more details.
14
15You should have received a copy of the GNU General Public License
16along with this program.  If not, see <http://www.gnu.org/licenses/>
17*/
18
19use super::{main_loop::Task, Session};
20use lsp_server::{Message, Notification};
21use serde::de::DeserializeOwned;
22use std::{collections::HashMap, panic::RefUnwindSafe, sync::Arc};
23
24/// Callback for parallelized notifications
25type Callback<Db> = Arc<
26    dyn Fn(&Db, serde_json::Value) -> anyhow::Result<()> + Send + Sync + RefUnwindSafe + 'static,
27>;
28
29/// Callback for synchronous mutable notifications
30type SyncMutCallback<Db> = Box<dyn Fn(&mut Session<Db>, serde_json::Value) -> anyhow::Result<()>>;
31
32/// A registry for LSP notifications.
33///
34/// This registry allows you to register handlers for LSP notifications.
35///
36/// The handlers can be executed in a separate thread or synchronously with mutable access to the session.
37///
38/// The handlers are registered using the `on` and `on_mut` methods.
39#[derive(Default)]
40pub struct NotificationRegistry<Db: salsa::Database> {
41    handlers: HashMap<String, Callback<Db>>,
42    sync_mut_handlers: HashMap<String, SyncMutCallback<Db>>,
43}
44
45impl<Db: salsa::Database + Clone + Send + RefUnwindSafe> NotificationRegistry<Db> {
46    /// Register a notification handler that will be pushed to the task pool.
47    ///
48    /// This handler is Cancelable and will be executed in a separate thread.
49    ///
50    /// Note that there is no retry mechanism for cancelled or failed notifications.
51    pub fn on<N, F>(&mut self, handler: F) -> &mut Self
52    where
53        N: lsp_types::notification::Notification,
54        N::Params: DeserializeOwned,
55        F: Fn(&Db, N::Params) -> anyhow::Result<()> + Send + Sync + RefUnwindSafe + 'static,
56    {
57        let method = N::METHOD.to_string();
58        let callback: Callback<Db> = Arc::new(move |session, params| {
59            let parsed_params: N::Params = serde_json::from_value(params)?;
60            handler(session, parsed_params)?;
61            Ok(())
62        });
63
64        self.handlers.insert(method, callback);
65        self
66    }
67
68    /// Register a notification handler that will be executed synchronously with a mutable access to [`Session`].
69    pub fn on_mut<N, F>(&mut self, handler: F) -> &mut Self
70    where
71        N: lsp_types::notification::Notification,
72        N::Params: DeserializeOwned,
73        F: Fn(&mut Session<Db>, N::Params) -> anyhow::Result<()> + Send + 'static,
74    {
75        let method = N::METHOD.to_string();
76        let callback: SyncMutCallback<Db> = Box::new(move |session, params| {
77            let parsed_params: N::Params = serde_json::from_value(params)?;
78            handler(session, parsed_params)?;
79            Ok(())
80        });
81
82        self.sync_mut_handlers.insert(method, callback);
83        self
84    }
85
86    pub(crate) fn get(&self, req: &Notification) -> Option<&Callback<Db>> {
87        self.handlers.get(&req.method)
88    }
89
90    pub(crate) fn get_sync_mut(&self, req: &Notification) -> Option<&SyncMutCallback<Db>> {
91        self.sync_mut_handlers.get(&req.method)
92    }
93
94    /// Push a notification handler to the task pool.
95    pub(crate) fn exec(session: &Session<Db>, callback: &Callback<Db>, not: Notification) {
96        let params = not.params;
97
98        let snapshot = session.snapshot();
99        let cb = Arc::clone(callback);
100        session
101            .task_pool
102            .spawn(move |sender| match snapshot.with_db(|db| cb(db, params)) {
103                Err(e) => log::warn!("Cancelled notification: {e}"),
104                Ok(result) => {
105                    if let Err(e) = result {
106                        sender.send(Task::NotificationError(e)).unwrap();
107                    }
108                }
109            });
110    }
111
112    /// Execute a synchronous mutable notification handler immediatly.
113    ///
114    /// Depending on the handler, this may cancel parallelized notifications.
115    pub(crate) fn exec_sync_mut(
116        session: &mut Session<Db>,
117        callback: &SyncMutCallback<Db>,
118        not: Notification,
119    ) -> anyhow::Result<()> {
120        if let Err(e) = callback(session, not.params) {
121            Self::handle_error(session, e)
122        } else {
123            Ok(())
124        }
125    }
126
127    pub(crate) fn handle_error(session: &Session<Db>, error: anyhow::Error) -> anyhow::Result<()> {
128        session
129            .connection
130            .sender
131            .send(Message::Notification(Notification {
132                method: "window/showMessage".to_string(),
133                params: serde_json::json!({
134                    "type": lsp_types::MessageType::ERROR,
135                    "message": error.to_string(),
136                }),
137            }))?;
138        Ok(())
139    }
140}