tenant.openagents/omega

1use std::str;
2use std::sync::Arc;
3
4use anyhow::{Context as _, Result};
5use collections::HashMap;
6use futures::{
7    AsyncBufReadExt, AsyncRead, AsyncReadExt as _, SinkExt as _,
8    channel::mpsc::{Receiver, Sender, channel},
9    io::BufReader,
10};
11use gpui::{BackgroundExecutor, Task};
12use log::warn;
13use parking_lot::Mutex;
14
15use crate::{
16    AnyResponse, CONTENT_LEN_HEADER, IoHandler, IoKind, NotificationOrRequest, RequestId,
17    ResponseHandler,
18};
19
20const HEADER_DELIMITER: &[u8; 4] = b"\r\n\r\n";
21
22/// Bounds the number of incoming LSP messages buffered between the background
23/// reader and the foreground dispatcher. When the queue is full, the reader
24/// stops reading the server's stdout, letting the OS pipe apply backpressure
25/// to the server instead of buffering messages in memory without limit while
26/// the foreground thread is unresponsive.
27pub(crate) const INCOMING_MESSAGE_QUEUE_CAPACITY: usize = 128;
28
29/// Handler for stdout of language server.
30pub struct LspStdoutHandler {
31    pub(super) loop_handle: Task<Result<()>>,
32    pub(super) incoming_messages: Receiver<NotificationOrRequest>,
33}
34
35async fn read_headers<Stdout>(reader: &mut BufReader<Stdout>, buffer: &mut Vec<u8>) -> Result<()>
36where
37    Stdout: AsyncRead + Unpin + Send + 'static,
38{
39    loop {
40        if buffer.len() >= HEADER_DELIMITER.len()
41            && buffer[(buffer.len() - HEADER_DELIMITER.len())..] == HEADER_DELIMITER[..]
42        {
43            return Ok(());
44        }
45
46        if reader.read_until(b'\n', buffer).await? == 0 {
47            anyhow::bail!("cannot read LSP message headers");
48        }
49    }
50}
51
52impl LspStdoutHandler {
53    pub fn new<Input>(
54        stdout: Input,
55        response_handlers: Arc<Mutex<Option<HashMap<RequestId, ResponseHandler>>>>,
56        io_handlers: Arc<Mutex<HashMap<i32, IoHandler>>>,
57        cx: BackgroundExecutor,
58    ) -> Self
59    where
60        Input: AsyncRead + Unpin + Send + 'static,
61    {
62        let (tx, notifications_channel) = channel(INCOMING_MESSAGE_QUEUE_CAPACITY);
63        let loop_handle = cx.spawn(Self::handler(stdout, tx, response_handlers, io_handlers));
64        Self {
65            loop_handle,
66            incoming_messages: notifications_channel,
67        }
68    }
69
70    async fn handler<Input>(
71        stdout: Input,
72        mut notifications_sender: Sender<NotificationOrRequest>,
73        response_handlers: Arc<Mutex<Option<HashMap<RequestId, ResponseHandler>>>>,
74        io_handlers: Arc<Mutex<HashMap<i32, IoHandler>>>,
75    ) -> anyhow::Result<()>
76    where
77        Input: AsyncRead + Unpin + Send + 'static,
78    {
79        let mut stdout = BufReader::new(stdout);
80
81        let mut buffer = Vec::new();
82
83        loop {
84            buffer.clear();
85
86            read_headers(&mut stdout, &mut buffer).await?;
87
88            let headers = std::str::from_utf8(&buffer)?;
89
90            let message_len = headers
91                .split('\n')
92                .find(|line| line.starts_with(CONTENT_LEN_HEADER))
93                .and_then(|line| line.strip_prefix(CONTENT_LEN_HEADER))
94                .with_context(|| format!("invalid LSP message header {headers:?}"))?
95                .trim_end()
96                .parse()?;
97
98            buffer.resize(message_len, 0);
99            stdout.read_exact(&mut buffer).await?;
100
101            if let Ok(message) = str::from_utf8(&buffer) {
102                log::trace!("incoming message: {message}");
103                for handler in io_handlers.lock().values_mut() {
104                    handler(IoKind::StdOut, message);
105                }
106            }
107
108            if let Ok(msg) = serde_json::from_slice::<NotificationOrRequest>(&buffer) {
109                notifications_sender.send(msg).await?;
110            } else if let Ok(AnyResponse {
111                id, error, result, ..
112            }) = serde_json::from_slice(&buffer)
113            {
114                let handler = {
115                    response_handlers
116                        .lock()
117                        .as_mut()
118                        .and_then(|handlers| handlers.remove(&id))
119                };
120                if let Some(handler) = handler {
121                    if let Some(error) = error {
122                        handler(Err(error)).await;
123                    } else if let Some(result) = result {
124                        handler(Ok(result.get().into())).await;
125                    } else {
126                        handler(Ok("null".into())).await;
127                    }
128                }
129            } else {
130                warn!(
131                    "failed to deserialize LSP message:\n{}",
132                    std::str::from_utf8(&buffer)?
133                );
134            }
135        }
136    }
137}
138
139#[cfg(test)]
140mod tests {
141    use super::*;
142    use futures::{AsyncWriteExt as _, StreamExt as _};
143    use gpui::TestAppContext;
144
145    #[gpui::test]
146    async fn test_backpressure_when_messages_are_not_consumed(cx: &mut TestAppContext) {
147        let total_messages = INCOMING_MESSAGE_QUEUE_CAPACITY * 4;
148        let (mut writer, reader) = async_pipe::pipe();
149        let mut handler = LspStdoutHandler::new(
150            reader,
151            Arc::new(Mutex::new(Some(HashMap::default()))),
152            Arc::new(Mutex::new(HashMap::default())),
153            cx.background_executor.clone(),
154        );
155
156        cx.background_executor
157            .spawn(async move {
158                let payload = r#"{"jsonrpc":"2.0","method":"test/notification","params":{}}"#;
159                let message = format!("Content-Length: {}\r\n\r\n{}", payload.len(), payload);
160                for _ in 0..total_messages {
161                    writer.write_all(message.as_bytes()).await.unwrap();
162                }
163            })
164            .detach();
165
166        cx.run_until_parked();
167        let mut received = 0;
168        while handler.incoming_messages.try_recv().is_ok() {
169            received += 1;
170        }
171        assert!(
172            received < total_messages,
173            "the reader buffered all {total_messages} messages while the consumer was wedged"
174        );
175        assert!(
176            received <= INCOMING_MESSAGE_QUEUE_CAPACITY + 2,
177            "expected at most {} buffered messages, got {received}",
178            INCOMING_MESSAGE_QUEUE_CAPACITY + 2
179        );
180
181        while received < total_messages {
182            assert!(
183                handler.incoming_messages.next().await.is_some(),
184                "the message stream ended after {received} of {total_messages} messages"
185            );
186            received += 1;
187        }
188    }
189
190    #[gpui::test]
191    async fn test_read_headers() {
192        let mut buf = Vec::new();
193        let mut reader = BufReader::new(b"Content-Length: 123\r\n\r\n" as &[u8]);
194        read_headers(&mut reader, &mut buf).await.unwrap();
195        assert_eq!(buf, b"Content-Length: 123\r\n\r\n");
196
197        let mut buf = Vec::new();
198        let mut reader = BufReader::new(b"Content-Type: application/vscode-jsonrpc\r\nContent-Length: 1235\r\n\r\n{\"somecontent\":123}" as &[u8]);
199        read_headers(&mut reader, &mut buf).await.unwrap();
200        assert_eq!(
201            buf,
202            b"Content-Type: application/vscode-jsonrpc\r\nContent-Length: 1235\r\n\r\n"
203        );
204
205        let mut buf = Vec::new();
206        let mut reader = BufReader::new(b"Content-Length: 1235\r\nContent-Type: application/vscode-jsonrpc\r\n\r\n{\"somecontent\":true}" as &[u8]);
207        read_headers(&mut reader, &mut buf).await.unwrap();
208        assert_eq!(
209            buf,
210            b"Content-Length: 1235\r\nContent-Type: application/vscode-jsonrpc\r\n\r\n"
211        );
212    }
213}
214