lspkit_server/
dispatch.rs

1//! JSON-RPC method dispatcher.
2//!
3//! Handlers are registered by method name. Every request is associated with a
4//! [`CancellationToken`] that fires when the client sends `$/cancelRequest`.
5
6use std::collections::HashMap;
7use std::pin::Pin;
8use std::sync::{Arc, RwLock};
9
10use serde_json::Value;
11use tokio_util::sync::CancellationToken;
12
13use crate::jsonrpc::{JsonRpcError, RequestId};
14
15type HandlerFuture = Pin<Box<dyn std::future::Future<Output = HandlerResult> + Send>>;
16type HandlerFn = Arc<dyn Fn(Value, CancellationToken) -> HandlerFuture + Send + Sync>;
17
18/// The two outcomes a request handler can produce.
19#[non_exhaustive]
20#[derive(Debug)]
21pub enum HandlerResult {
22    /// Success — value is JSON-serialized into the response's `result`.
23    Ok(Value),
24    /// Failure — error is serialized into the response's `error`.
25    Err(JsonRpcError),
26}
27
28/// Errors from dispatching a request.
29#[non_exhaustive]
30#[derive(Debug, thiserror::Error)]
31pub enum DispatchError {
32    /// No handler registered for the method.
33    #[error("method not found: {0}")]
34    UnknownMethod(String),
35}
36
37/// Routes inbound JSON-RPC method calls to async handlers.
38#[derive(Clone, Default)]
39pub struct Dispatcher {
40    handlers: Arc<RwLock<HashMap<String, HandlerFn>>>,
41    cancellations: Arc<RwLock<HashMap<RequestId, CancellationToken>>>,
42}
43
44impl Dispatcher {
45    /// New empty dispatcher.
46    #[must_use]
47    pub fn new() -> Self {
48        Self::default()
49    }
50
51    /// Register a handler for `method`. Replaces any previously registered handler.
52    pub fn register<F, Fut>(&self, method: impl Into<String>, handler: F)
53    where
54        F: Fn(Value, CancellationToken) -> Fut + Send + Sync + 'static,
55        Fut: std::future::Future<Output = HandlerResult> + Send + 'static,
56    {
57        let wrapped: HandlerFn = Arc::new(move |params, cancel| Box::pin(handler(params, cancel)));
58        if let Ok(mut guard) = self.handlers.write() {
59            guard.insert(method.into(), wrapped);
60        }
61    }
62
63    /// Dispatch a request. Returns the handler's result.
64    ///
65    /// The supplied `id` is used to register a cancellation token so a later
66    /// `$/cancelRequest` for the same id can trip the handler.
67    ///
68    /// # Errors
69    /// Returns [`DispatchError::UnknownMethod`] if no handler is registered.
70    pub async fn dispatch(
71        &self,
72        id: &RequestId,
73        method: &str,
74        params: Value,
75    ) -> Result<HandlerResult, DispatchError> {
76        let handler = self
77            .handlers
78            .read()
79            .ok()
80            .and_then(|guard| guard.get(method).cloned())
81            .ok_or_else(|| DispatchError::UnknownMethod(method.to_owned()))?;
82        let cancel = CancellationToken::new();
83        if let Ok(mut guard) = self.cancellations.write() {
84            guard.insert(id.clone(), cancel.clone());
85        }
86        let result = handler(params, cancel).await;
87        if let Ok(mut guard) = self.cancellations.write() {
88            guard.remove(id);
89        }
90        Ok(result)
91    }
92
93    /// Trip the cancellation token for an in-flight request. No-op if unknown.
94    pub fn cancel(&self, id: &RequestId) {
95        if let Ok(guard) = self.cancellations.read() {
96            if let Some(token) = guard.get(id) {
97                token.cancel();
98            }
99        }
100    }
101}
102
103impl std::fmt::Debug for Dispatcher {
104    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
105        let handlers = self.handlers.read().map_or(0, |g| g.len());
106        let in_flight = self.cancellations.read().map_or(0, |g| g.len());
107        f.debug_struct("Dispatcher")
108            .field("handlers", &handlers)
109            .field("in_flight", &in_flight)
110            .finish()
111    }
112}
lspkit_server/dispatch.rs

lspkit_server/
dispatch.rs
