liminal_cli/mcp/

serve.rs

//! Blocking stdin and blocking tap readers share a single framed writer.

use super::framing::{self, FrameError, INVALID_REQUEST, Incoming, PARSE_ERROR};
use super::notifications::{self, Batch};
use super::server::Server;
use super::subscriptions::{Pending, Service};
use super::writer::Writer;
use crate::network::Rail;
use serde_json::Value;
use std::io::{self, BufRead, Write};
use std::sync::Arc;
use std::sync::atomic::Ordering;
use std::thread::ScopedJoinHandle;

type ReaderThread<'scope> = (Arc<Service>, ScopedJoinHandle<'scope, io::Result<()>>);

impl<R: Rail> Server<R> {
    /// Serves requests while tap readers send unsolicited logging lines.
    /// EOF signals and joins every receiver before this returns.
    ///
    /// # Errors
    /// Returns the input, output, framing or subscription-teardown failure.
    pub fn serve<I: BufRead, O: Write + Send>(
        &mut self,
        input: &mut I,
        output: &mut O,
    ) -> io::Result<()> {
        if let Err(error) = self.start_listens() {
            return combine(
                Err(error),
                self.channel_source
                    .as_ref()
                    .map_or(Ok(()), |source| source.revoke().map_err(io::Error::other)),
            );
        }
        let writer = Arc::new(Writer::new(
            output,
            self.frame_bytes,
            Arc::clone(&self.warning_only),
            Arc::clone(&self.subscriptions),
            Arc::clone(&self.listens),
        ));
        let served = std::thread::scope(|scope| {
            let mut readers = Vec::new();
            let mut channels = Vec::new();
            let served = self.requests(scope, input, &writer, &mut readers, &mut channels);
            let mut result = combine(
                served,
                self.subscriptions.close_all().map_err(io::Error::other),
            );
            result = combine(result, self.listens.close_all().map_err(io::Error::other));
            for (_, reader) in readers {
                result = combine(result, joined(reader));
            }
            for (_, reader) in channels {
                result = combine(result, super::channel_serve::joined(reader, &self.listens));
            }
            combine(result, writer.check())
        });
        combine(
            served,
            self.channel_source
                .as_ref()
                .map_or(Ok(()), |source| source.revoke().map_err(io::Error::other)),
        )
    }

    fn requests<'scope, 'env, I: BufRead, O: Write + Send + 'scope>(
        &mut self,
        scope: &'scope std::thread::Scope<'scope, 'env>,
        input: &mut I,
        writer: &Arc<Writer<O>>,
        readers: &mut Vec<ReaderThread<'scope>>,
        channels: &mut Vec<super::channel_serve::Reader<'scope>>,
    ) -> io::Result<()> {
        let context = super::channel_serve::Context {
            config: self.listen_config.clone(),
            device: self.rail().server_address().to_owned(),
            record_bytes: self
                .channel_source
                .as_ref()
                .map_or(self.frame_bytes, |source| source.record_bytes),
        };
        super::channel_serve::spawn(scope, channels, &self.listens, writer, &context)?;
        loop {
            writer.check()?;
            join_closed(readers)?;
            super::channel_serve::join_closed(channels, &self.listens)?;
            let incoming = framing::read_message(input, self.frame_bytes);
            writer.check()?;
            let answer = match incoming {
                Ok(None) => return Ok(()),
                Ok(Some(Incoming::Notification { .. })) => continue,
                Ok(Some(Incoming::Request { id, method, params })) => {
                    self.answer(&id, &method, &params)
                }
                Err(FrameError::Stream { message }) => return Err(io::Error::other(message)),
                Err(error @ (FrameError::Overlong { .. } | FrameError::NotRequest { .. })) => {
                    framing::error(&Value::Null, INVALID_REQUEST, &error.to_string(), None)
                }
                Err(error @ FrameError::NotJson { .. }) => {
                    framing::error(&Value::Null, PARSE_ERROR, &error.to_string(), None)
                }
            };
            join_closed(readers)?;
            super::channel_serve::join_closed(channels, &self.listens)?;
            writer.write(answer, None, false)?;
            spawn_taps(
                scope,
                self.subscriptions.pending().map_err(io::Error::other)?,
                writer,
                readers,
            )?;
            super::channel_serve::spawn(scope, channels, &self.listens, writer, &context)?;
        }
    }
}

fn joined(reader: ScopedJoinHandle<'_, io::Result<()>>) -> io::Result<()> {
    reader
        .join()
        .map_err(|_| io::Error::other("a tap receiver thread panicked"))?
}

fn spawn_taps<'scope, 'env, O: Write + Send + 'scope>(
    scope: &'scope std::thread::Scope<'scope, 'env>,
    pending: Vec<Pending>,
    writer: &Arc<Writer<O>>,
    readers: &mut Vec<ReaderThread<'scope>>,
) -> io::Result<()> {
    for pending in pending {
        let writer = Arc::clone(writer);
        let service = Arc::clone(&pending.service);
        let reader = std::thread::Builder::new()
            .name("liminal-mcp-tap".to_owned())
            .spawn_scoped(scope, move || receive(&pending, &writer))?;
        readers.push((service, reader));
    }
    Ok(())
}

fn join_closed(readers: &mut Vec<ReaderThread<'_>>) -> io::Result<()> {
    let mut n = 0;
    while n < readers.len() {
        if readers[n].0.active.load(Ordering::SeqCst) {
            n += 1;
        } else {
            joined(readers.swap_remove(n).1)?;
        }
    }
    Ok(())
}

fn combine(first: io::Result<()>, second: io::Result<()>) -> io::Result<()> {
    match (first, second) {
        (Ok(()), result) | (result, Ok(())) => result,
        (Err(first), Err(second)) => Err(io::Error::other(format!("{first}; {second}"))),
    }
}

fn receive<O: Write>(pending: &Pending, writer: &Writer<O>) -> io::Result<()> {
    while pending.service.active.load(Ordering::SeqCst) {
        let batch = pending
            .reader
            .receive()
            .map_err(|error| error.to_string())
            .and_then(|record| Batch::read(record, &pending.subject));
        let batch = match batch {
            Ok(batch) => batch,
            Err(reason) => {
                if pending.service.active.swap(false, Ordering::SeqCst) {
                    let closed = pending.service.close();
                    let reason = match closed {
                        Ok(()) => reason,
                        Err(words) => format!("{reason}; subscription teardown failed: {words}"),
                    };
                    return writer.write(
                        notifications::warning(
                            &pending.service,
                            &pending.subject,
                            "ended",
                            &format!("subscription to {} ended: {reason}", pending.service.name),
                        ),
                        Some(&pending.service),
                        true,
                    );
                }
                return Ok(());
            }
        };
        writer.record(&pending.service, &pending.subject, &batch)?;
    }
    Ok(())
}