Skip to main content

liminal_cli/mcp/
serve.rs

1//! Blocking stdin and blocking tap readers share a single framed writer.
2
3use super::framing::{self, FrameError, INVALID_REQUEST, Incoming, PARSE_ERROR};
4use super::notifications::{self, Batch};
5use super::server::Server;
6use super::subscriptions::{Pending, Service};
7use super::writer::Writer;
8use crate::network::Rail;
9use serde_json::Value;
10use std::io::{self, BufRead, Write};
11use std::sync::Arc;
12use std::sync::atomic::Ordering;
13use std::thread::ScopedJoinHandle;
14
15type ReaderThread<'scope> = (Arc<Service>, ScopedJoinHandle<'scope, io::Result<()>>);
16
17impl<R: Rail> Server<R> {
18    /// Serves requests while tap readers send unsolicited logging lines.
19    /// EOF signals and joins every receiver before this returns.
20    ///
21    /// # Errors
22    /// Returns the input, output, framing or subscription-teardown failure.
23    pub fn serve<I: BufRead, O: Write + Send>(
24        &mut self,
25        input: &mut I,
26        output: &mut O,
27    ) -> io::Result<()> {
28        let writer = Arc::new(Writer::new(
29            output,
30            self.frame_bytes,
31            Arc::clone(&self.warning_only),
32            Arc::clone(&self.subscriptions),
33        ));
34        std::thread::scope(|scope| {
35            let mut readers = Vec::new();
36            let served = loop {
37                if let Err(error) = writer.check().and_then(|()| join_closed(&mut readers)) {
38                    break Err(error);
39                }
40                let incoming = framing::read_message(input, self.frame_bytes);
41                if let Err(error) = writer.check() {
42                    break Err(error);
43                }
44                let answer = match incoming {
45                    Ok(None) => break Ok(()),
46                    Ok(Some(Incoming::Notification { .. })) => continue,
47                    Ok(Some(Incoming::Request { id, method, params })) => {
48                        self.answer(&id, &method, &params)
49                    }
50                    Err(FrameError::Stream { message }) => break Err(io::Error::other(message)),
51                    Err(error @ (FrameError::Overlong { .. } | FrameError::NotRequest { .. })) => {
52                        framing::error(&Value::Null, INVALID_REQUEST, &error.to_string(), None)
53                    }
54                    Err(error @ FrameError::NotJson { .. }) => {
55                        framing::error(&Value::Null, PARSE_ERROR, &error.to_string(), None)
56                    }
57                };
58                if let Err(error) = join_closed(&mut readers) {
59                    break Err(error);
60                }
61                if let Err(error) = writer.write(answer, None, false) {
62                    break Err(error);
63                }
64                let pending = match self.subscriptions.pending() {
65                    Ok(pending) => pending,
66                    Err(words) => break Err(io::Error::other(words)),
67                };
68                let mut spawn_error = None;
69                for pending in pending {
70                    let writer = Arc::clone(&writer);
71                    let service = Arc::clone(&pending.service);
72                    match std::thread::Builder::new()
73                        .name("liminal-mcp-tap".to_owned())
74                        .spawn_scoped(scope, move || receive(&pending, &writer))
75                    {
76                        Ok(reader) => readers.push((service, reader)),
77                        Err(error) => {
78                            spawn_error = Some(error);
79                            break;
80                        }
81                    }
82                }
83                if let Some(error) = spawn_error {
84                    break Err(error);
85                }
86            };
87            let mut result = combine(
88                served,
89                self.subscriptions.close_all().map_err(io::Error::other),
90            );
91            for (_, reader) in readers {
92                result = combine(result, joined(reader));
93            }
94            combine(result, writer.check())
95        })
96    }
97}
98
99fn joined(reader: ScopedJoinHandle<'_, io::Result<()>>) -> io::Result<()> {
100    reader
101        .join()
102        .map_err(|_| io::Error::other("a tap receiver thread panicked"))?
103}
104
105fn join_closed(readers: &mut Vec<ReaderThread<'_>>) -> io::Result<()> {
106    let mut n = 0;
107    while n < readers.len() {
108        if readers[n].0.active.load(Ordering::SeqCst) {
109            n += 1;
110        } else {
111            joined(readers.swap_remove(n).1)?;
112        }
113    }
114    Ok(())
115}
116
117fn combine(first: io::Result<()>, second: io::Result<()>) -> io::Result<()> {
118    match (first, second) {
119        (Ok(()), result) | (result, Ok(())) => result,
120        (Err(first), Err(second)) => Err(io::Error::other(format!("{first}; {second}"))),
121    }
122}
123
124fn receive<O: Write>(pending: &Pending, writer: &Writer<O>) -> io::Result<()> {
125    while pending.service.active.load(Ordering::SeqCst) {
126        let batch = pending
127            .reader
128            .receive()
129            .map_err(|error| error.to_string())
130            .and_then(|record| Batch::read(record, &pending.subject));
131        let batch = match batch {
132            Ok(batch) => batch,
133            Err(reason) => {
134                if pending.service.active.swap(false, Ordering::SeqCst) {
135                    let closed = pending.service.close();
136                    let reason = match closed {
137                        Ok(()) => reason,
138                        Err(words) => format!("{reason}; subscription teardown failed: {words}"),
139                    };
140                    return writer.write(
141                        notifications::warning(
142                            &pending.service,
143                            &pending.subject,
144                            "ended",
145                            &format!("subscription to {} ended: {reason}", pending.service.name),
146                        ),
147                        Some(&pending.service),
148                        true,
149                    );
150                }
151                return Ok(());
152            }
153        };
154        writer.record(&pending.service, &pending.subject, &batch)?;
155    }
156    Ok(())
157}