lsp_max/composition/
server.rs

1//! Composition state and ComposedServer core (constructor + route_request).
2
3use std::collections::HashMap;
4use std::sync::Arc;
5
6use lsp_types_max::Diagnostic;
7use serde_json::{json, Value};
8use tokio::sync::{broadcast, Mutex};
9
10use crate::jsonrpc::{Error, Request, Result};
11use crate::Client;
12
13use super::capability_tracker::CapabilityTracker;
14use super::edit_gate::TransactionEditGate;
15use super::strategy::method_strategy;
16use super::upstream::PersistentUpstream;
17use super::version_tracker::DocumentVersionTracker;
18
19/// State manager for language server composition.
20#[derive(Debug)]
21pub struct CompositionState {
22    /// Tracks upstream capability changes and merges them.
23    pub capability_tracker: CapabilityTracker,
24    /// Tracks document open, change, and version states.
25    pub doc_tracker: DocumentVersionTracker,
26    /// Transactional guard to validate mutations.
27    pub edit_gate: TransactionEditGate,
28    /// Upstream request timeout config in milliseconds.
29    pub upstream_timeout_ms: u64,
30    /// Currently reported diagnostics mapped by source.
31    pub diagnostics: HashMap<String, HashMap<String, Value>>,
32    /// Monotonically increasing counter of requests.
33    pub request_counter: u64,
34    /// Trace history of requests.
35    pub request_traces: Arc<std::sync::Mutex<Vec<Value>>>,
36}
37
38impl CompositionState {
39    /// Creates a new composition state configured with the given upstreams.
40    pub fn new(upstream_addresses: Vec<(String, String)>) -> Self {
41        use super::strategy::UpstreamSource;
42        let mut tracker = CapabilityTracker::new();
43        for (id, addr) in upstream_addresses {
44            tracker.add_source(UpstreamSource::new(id, addr));
45        }
46        let upstream_timeout_ms = std::env::var("LSP_MAX_TIMEOUT")
47            .ok()
48            .and_then(|v| v.parse::<u64>().ok())
49            .unwrap_or(150);
50        Self {
51            capability_tracker: tracker,
52            doc_tracker: DocumentVersionTracker::new(),
53            edit_gate: TransactionEditGate::new(),
54            upstream_timeout_ms,
55            diagnostics: HashMap::new(),
56            request_counter: 0,
57            request_traces: Arc::new(std::sync::Mutex::new(Vec::new())),
58        }
59    }
60}
61
62/// Shared thread-safe handle to the composition state.
63pub type SharedCompositionState = Arc<Mutex<CompositionState>>;
64
65/// Composed LSP server that coordinates multiple upstream language servers.
66#[derive(Debug)]
67pub struct ComposedServer {
68    pub(super) _client: Client,
69    pub(super) state: SharedCompositionState,
70    pub(super) upstreams: Arc<dashmap::DashMap<String, PersistentUpstream>>,
71}
72
73impl ComposedServer {
74    /// Gets the shared composition state.
75    pub fn state(&self) -> &SharedCompositionState {
76        &self.state
77    }
78
79    /// Creates a new `ComposedServer` with a client handle and a list of upstream addresses.
80    pub fn new(client: Client, upstream_addresses: Vec<(String, String)>) -> Self {
81        let state = Arc::new(Mutex::new(CompositionState::new(
82            upstream_addresses.clone(),
83        )));
84        let upstreams = Arc::new(dashmap::DashMap::new());
85        let (unsolicited_tx, mut unsolicited_rx) = broadcast::channel(128);
86
87        let _upstreams_clone = upstreams.clone();
88        for (id, addr) in upstream_addresses {
89            let id_clone = id.clone();
90            let addr_clone = addr.clone();
91            let unsolicited_tx = unsolicited_tx.clone();
92            let upstreams_connections = upstreams.clone();
93            tokio::spawn(async move {
94                if let Some(conn) =
95                    PersistentUpstream::connect(id_clone.clone(), addr_clone, 1000, unsolicited_tx)
96                        .await
97                {
98                    upstreams_connections.insert(id_clone, conn);
99                }
100            });
101        }
102
103        let client_clone = client.clone();
104        let state_clone = state.clone();
105        let upstreams_connections_clone = upstreams.clone();
106        tokio::spawn(async move {
107            loop {
108                match unsolicited_rx.recv().await {
109                    Ok(notification) => {
110                        let msg = notification.message;
111                        let source_id = notification.source_id;
112
113                        let method = match msg.get("method").and_then(|v| v.as_str()) {
114                            Some(m) => m,
115                            None => continue,
116                        };
117
118                        if method == "client/registerCapability" {
119                            let mut valid_registrations = Vec::new();
120                            if let Some(params) = msg.get("params") {
121                                if let Some(registrations) =
122                                    params.get("registrations").and_then(|v| v.as_array())
123                                {
124                                    let mut s = state_clone.lock().await;
125                                    for reg in registrations {
126                                        if let (Some(id), Some(m)) = (
127                                            reg.get("id").and_then(|v| v.as_str()),
128                                            reg.get("method").and_then(|v| v.as_str()),
129                                        ) {
130                                            let options = reg
131                                                .get("registerOptions")
132                                                .cloned()
133                                                .unwrap_or(Value::Null);
134                                            if s.capability_tracker
135                                                .register_dynamic(id, m, &source_id, options)
136                                            {
137                                                valid_registrations.push(reg.clone());
138                                            }
139                                        }
140                                    }
141                                }
142                            }
143                            let msg_id_opt = msg.get("id").cloned();
144                            let client = client_clone.clone();
145                            let upstreams_connections = upstreams_connections_clone.clone();
146                            let source_id = source_id.clone();
147                            if valid_registrations.is_empty() {
148                                if let Some(msg_id) = msg_id_opt {
149                                    tokio::spawn(async move {
150                                        let resp = json!({ "jsonrpc": "2.0", "id": msg_id, "result": null });
151                                        if let Some(conn) = upstreams_connections.get(&source_id) {
152                                            let _ = conn.send_raw(resp).await;
153                                        }
154                                    });
155                                }
156                            } else {
157                                let params = json!({ "registrations": valid_registrations });
158                                tokio::spawn(async move {
159                                    let req_id = client.next_request_id();
160                                    let request =
161                                        Request::build("client/registerCapability".to_string())
162                                            .id(req_id.clone())
163                                            .params(params)
164                                            .finish();
165                                    let mut client_clone = client.clone();
166                                    use tower::Service;
167                                    let resp_res = client_clone.call(request).await;
168                                    if let Some(msg_id) = msg_id_opt {
169                                        let resp = match resp_res {
170                                            Ok(Some(response)) => {
171                                                let (_, result) = response.into_parts();
172                                                match result {
173                                                    Ok(res) => {
174                                                        json!({ "jsonrpc": "2.0", "id": msg_id, "result": res })
175                                                    }
176                                                    Err(err) => {
177                                                        json!({ "jsonrpc": "2.0", "id": msg_id, "error": err })
178                                                    }
179                                                }
180                                            }
181                                            _ => {
182                                                json!({ "jsonrpc": "2.0", "id": msg_id, "error": { "code": -32603, "message": "Internal error" } })
183                                            }
184                                        };
185                                        if let Some(conn) = upstreams_connections.get(&source_id) {
186                                            let _ = conn.send_raw(resp).await;
187                                        }
188                                    }
189                                });
190                            }
191                        } else if method == "client/unregisterCapability" {
192                            let mut valid_unregistrations = Vec::new();
193                            if let Some(params) = msg.get("params") {
194                                let unregs_opt = params
195                                    .get("unregisterations")
196                                    .and_then(|v| v.as_array())
197                                    .or_else(|| {
198                                        params.get("unregistrations").and_then(|v| v.as_array())
199                                    });
200                                if let Some(unregistrations) = unregs_opt {
201                                    let mut s = state_clone.lock().await;
202                                    for unreg in unregistrations {
203                                        if let Some(id) = unreg.get("id").and_then(|v| v.as_str()) {
204                                            if s.capability_tracker.unregister_dynamic(id) {
205                                                valid_unregistrations.push(unreg.clone());
206                                            }
207                                        }
208                                    }
209                                }
210                            }
211                            let msg_id_opt = msg.get("id").cloned();
212                            let client = client_clone.clone();
213                            let upstreams_connections = upstreams_connections_clone.clone();
214                            let source_id = source_id.clone();
215                            if valid_unregistrations.is_empty() {
216                                if let Some(msg_id) = msg_id_opt {
217                                    tokio::spawn(async move {
218                                        let resp = json!({ "jsonrpc": "2.0", "id": msg_id, "result": null });
219                                        if let Some(conn) = upstreams_connections.get(&source_id) {
220                                            let _ = conn.send_raw(resp).await;
221                                        }
222                                    });
223                                }
224                            } else {
225                                let params = json!({ "unregistrations": valid_unregistrations });
226                                tokio::spawn(async move {
227                                    let req_id = client.next_request_id();
228                                    let request =
229                                        Request::build("client/unregisterCapability".to_string())
230                                            .id(req_id.clone())
231                                            .params(params)
232                                            .finish();
233                                    let mut client_clone = client.clone();
234                                    use tower::Service;
235                                    let resp_res = client_clone.call(request).await;
236                                    if let Some(msg_id) = msg_id_opt {
237                                        let resp = match resp_res {
238                                            Ok(Some(response)) => {
239                                                let (_, result) = response.into_parts();
240                                                match result {
241                                                    Ok(res) => {
242                                                        json!({ "jsonrpc": "2.0", "id": msg_id, "result": res })
243                                                    }
244                                                    Err(err) => {
245                                                        json!({ "jsonrpc": "2.0", "id": msg_id, "error": err })
246                                                    }
247                                                }
248                                            }
249                                            _ => {
250                                                json!({ "jsonrpc": "2.0", "id": msg_id, "error": { "code": -32603, "message": "Internal error" } })
251                                            }
252                                        };
253                                        if let Some(conn) = upstreams_connections.get(&source_id) {
254                                            let _ = conn.send_raw(resp).await;
255                                        }
256                                    }
257                                });
258                            }
259                        } else if method == "textDocument/publishDiagnostics" {
260                            if let Some(params) = msg.get("params") {
261                                if let Some(uri) = params.get("uri").and_then(|v| v.as_str()) {
262                                    let diagnostics_val =
263                                        params.get("diagnostics").cloned().unwrap_or(json!([]));
264                                    let (merged_diags, version) = {
265                                        let mut s = state_clone.lock().await;
266                                        s.diagnostics
267                                            .entry(uri.to_string())
268                                            .or_default()
269                                            .insert(source_id.clone(), diagnostics_val);
270                                        let mut flat_diags = Vec::new();
271                                        if let Some(uri_diags) = s.diagnostics.get(uri) {
272                                            for (src_id, diags_value) in uri_diags {
273                                                if let Some(src) =
274                                                    s.capability_tracker.sources.get(src_id)
275                                                {
276                                                    if src.is_routable() {
277                                                        if let Some(diags_arr) =
278                                                            diags_value.as_array()
279                                                        {
280                                                            for diag in diags_arr {
281                                                                if src_id == "ggen-lsp" {
282                                                                    let has_source_id = diag
283                                                                        .get("data")
284                                                                        .and_then(|data| {
285                                                                            data.get("source_id")
286                                                                        })
287                                                                        .is_some();
288                                                                    if !has_source_id {
289                                                                        continue;
290                                                                    }
291                                                                }
292                                                                let mut diag_clone = diag.clone();
293                                                                if let Some(obj) =
294                                                                    diag_clone.as_object_mut()
295                                                                {
296                                                                    obj.insert(
297                                                                        "source".to_string(),
298                                                                        Value::String(
299                                                                            src_id.clone(),
300                                                                        ),
301                                                                    );
302                                                                }
303                                                                flat_diags.push(diag_clone);
304                                                            }
305                                                        }
306                                                    }
307                                                }
308                                            }
309                                        }
310                                        let version = s.doc_tracker.current_version(uri);
311                                        (flat_diags, version)
312                                    };
313                                    if let Ok(diag_list) = serde_json::from_value::<Vec<Diagnostic>>(
314                                        Value::Array(merged_diags),
315                                    ) {
316                                        if let Ok(parsed_uri) =
317                                            <lsp_types_max::Uri as std::str::FromStr>::from_str(uri)
318                                        {
319                                            client_clone
320                                                .publish_diagnostics(parsed_uri, diag_list, version)
321                                                .await;
322                                        }
323                                    }
324                                }
325                            }
326                        }
327                    }
328                    Err(broadcast::error::RecvError::Lagged(n)) => {
329                        tracing::warn!("Unsolicited message channel lagged by {} messages", n);
330                    }
331                    Err(broadcast::error::RecvError::Closed) => {
332                        break;
333                    }
334                }
335            }
336        });
337
338        Self {
339            _client: client,
340            state,
341            upstreams,
342        }
343    }
344
345    pub(super) async fn route_request<P, R>(&self, method: &str, params: P) -> Result<Option<R>>
346    where
347        P: serde::Serialize,
348        R: serde::de::DeserializeOwned,
349    {
350        tracing::trace!("--- route_request [{}] start", method);
351        let params_val = serde_json::to_value(params).unwrap_or(Value::Null);
352
353        tracing::trace!("--- route_request [{}] locking state at start", method);
354        let (request_id, strategy_str, uri_opt, doc_version) = {
355            let mut s = self.state.lock().await;
356            tracing::trace!(
357                "--- route_request [{}] locked state at start inside block",
358                method
359            );
360            s.request_counter += 1;
361            let counter = s.request_counter;
362            let strategy = method_strategy(method);
363            let uri = params_val
364                .get("textDocument")
365                .and_then(|td| td.get("uri"))
366                .and_then(|u| u.as_str())
367                .map(|s| s.to_string())
368                .or_else(|| {
369                    params_val
370                        .get("uri")
371                        .and_then(|u| u.as_str())
372                        .map(|s| s.to_string())
373                });
374            let ver = uri.as_ref().and_then(|u| s.doc_tracker.current_version(u));
375            (
376                format!("req_{}", counter),
377                format!("{:?}", strategy),
378                uri,
379                ver,
380            )
381        };
382        tracing::trace!("--- route_request [{}] strategy: {}", method, strategy_str);
383
384        let sources_contacted = Arc::new(std::sync::Mutex::new(Vec::<String>::new()));
385        let sources_returned = Arc::new(std::sync::Mutex::new(Vec::<String>::new()));
386        let source_health = Arc::new(std::sync::Mutex::new(std::collections::HashMap::<
387            String,
388            String,
389        >::new()));
390        let gate_outcome = Arc::new(std::sync::Mutex::new(None::<String>));
391        let staleness_outcome = Arc::new(std::sync::Mutex::new(None::<String>));
392
393        let res = self
394            .route_request_internal(
395                method,
396                params_val,
397                sources_contacted.clone(),
398                sources_returned.clone(),
399                source_health.clone(),
400                gate_outcome.clone(),
401                staleness_outcome.clone(),
402            )
403            .await;
404
405        let final_response_val = match &res {
406            Ok(Some(v)) => v.clone(),
407            _ => Value::Null,
408        };
409
410        // Record the trace
411        {
412            let s = self.state.lock().await;
413            let trace = json!({
414                "request_id": request_id,
415                "method": method,
416                "document_uri": uri_opt,
417                "document_version": doc_version,
418                "strategy": strategy_str,
419                "sources_contacted": *sources_contacted.lock().unwrap(),
420                "sources_returned": *sources_returned.lock().unwrap(),
421                "source_health": *source_health.lock().unwrap(),
422                "gate_outcome": *gate_outcome.lock().unwrap(),
423                "staleness_outcome": *staleness_outcome.lock().unwrap(),
424                "final_response": final_response_val,
425            });
426            s.request_traces.lock().unwrap().push(trace);
427        }
428
429        match res {
430            Ok(Some(v)) => match serde_json::from_value::<R>(v) {
431                Ok(r) => Ok(Some(r)),
432                Err(e) => Err(Error::invalid_params(format!(
433                    "Failed to deserialize response: {}",
434                    e
435                ))),
436            },
437            Ok(None) => Ok(None),
438            Err(e) => Err(e),
439        }
440    }
441}
lsp_max/composition/server.rs

lsp_max/composition/
server.rs
