klieo-mcp-server 2.2.0

//! Per-transport frame-write abstraction over JSON-RPC envelopes.
//!
//! Stdio writes newline-delimited frames to stdout. HTTP writes
//! frames into a per-session mpsc whose receiver feeds the
//! `GET /mcp` SSE body. This trait keeps `OutboundRequests`
//! transport-agnostic.

use async_trait::async_trait;
use serde_json::Value;
use tokio::io::AsyncWriteExt;

use crate::outbound::SharedWriter;

/// Per-transport sink that accepts one outbound JSON-RPC frame at a
/// time and pushes it to the peer.
///
/// Implementations are responsible for any framing the transport
/// requires (newline delimiters for stdio, channel-bounded enqueue
/// for HTTP/SSE). Returning [`OutboundSinkError::TransportClosed`]
/// signals the caller's correlation table to drop the pending entry
/// and surface `TransportClosed` to the originating tool.
///
/// The frame is shared ownership: callers wrap the
/// [`serde_json::Value`] in [`std::sync::Arc`] once at construction
/// and hand the same pointer to every implementor. Implementors read
/// through the [`std::sync::Arc`] (refcount bump on share, no defensive
/// deep clone of the JSON tree) and may pass it onward into ring/replay
/// storage by cloning the [`std::sync::Arc`] itself.
#[async_trait]
pub trait OutboundFrameSink: Send + Sync {
    /// Push one outbound JSON-RPC frame to the peer.
    async fn send_frame(&self, frame: std::sync::Arc<Value>) -> Result<(), OutboundSinkError>;
}

/// Errors a [`OutboundFrameSink`] may surface to its caller.
#[derive(Debug, thiserror::Error)]
#[non_exhaustive]
pub enum OutboundSinkError {
    /// The underlying transport is closed or cannot accept further
    /// frames (broken pipe on stdio, dropped receiver on HTTP).
    #[error("transport closed")]
    TransportClosed,
    /// A frame could not be serialised. Indicates an internal
    /// invariant violation rather than a transport failure.
    #[error("frame serialisation failed: {0}")]
    Serialisation(#[source] serde_json::Error),
}

/// Stdio-flavoured sink: serialises each frame and writes it as a
/// newline-delimited JSON payload through the shared transport
/// writer, holding the writer's mutex across the byte-write +
/// delimiter + flush so concurrent senders never interleave.
pub(crate) struct StdioFrameSink {
    writer: SharedWriter,
}

impl StdioFrameSink {
    pub(crate) fn new(writer: SharedWriter) -> Self {
        Self { writer }
    }
}

#[async_trait]
impl OutboundFrameSink for StdioFrameSink {
    async fn send_frame(&self, frame: std::sync::Arc<Value>) -> Result<(), OutboundSinkError> {
        let bytes = match serde_json::to_vec(&*frame) {
            Ok(bytes) => bytes,
            // serde_json::Value is always serializable; this arm is a safety net
            // for future Value representation changes.
            Err(err) => {
                tracing::error!(
                    target: "klieo::mcp::stdio",
                    error = %err,
                    "outbound frame serialisation failed",
                );
                return Err(OutboundSinkError::Serialisation(err));
            }
        };
        let mut w = self.writer.lock().await;
        if w.write_all(&bytes).await.is_err()
            || w.write_all(b"\n").await.is_err()
            || w.flush().await.is_err()
        {
            return Err(OutboundSinkError::TransportClosed);
        }
        Ok(())
    }
}

/// Capacity of the per-session outbound ring. Sized large enough that
/// drop-oldest cannot fire under normal MCP-notification bursts; only
/// a slow SSE consumer can push the ring to capacity.
#[cfg(all(feature = "http", not(feature = "test-fixtures")))]
pub(crate) const OUTBOUND_QUEUE_CAPACITY: usize = 1024;

/// Same constant exposed publicly so integration tests can name the
/// ring capacity when driving the drop-oldest path past it. Gated on
/// `test-fixtures` so production callers cannot reach the value.
#[cfg(all(feature = "http", feature = "test-fixtures"))]
pub const OUTBOUND_QUEUE_CAPACITY: usize = 1024;

/// Outbound sink writing JSON-RPC frames into a per-session
/// drop-oldest ring whose receiver feeds the `GET /mcp` SSE body.
/// Constructed in the GET handler; the correlation table sees only
/// the [`OutboundFrameSink`] trait so the same `OutboundRequests`
/// primitive works against either transport.
///
/// Holds a [`std::sync::Weak`] handle to the owning session so the
/// registry remains the sole strong owner: DELETE, the idle reaper,
/// and SSE-body drop can each release the session, dropping the
/// outbound ring sender and terminating the live SSE body, even while
/// the sink itself is still reachable through `OutboundRequests`.
#[cfg(feature = "http")]
pub(crate) struct HttpFrameSink {
    session: std::sync::Weak<crate::session::Session>,
    tx: crate::outbound_ring::RingSender<(u64, std::sync::Arc<Value>)>,
    sse_replay_capacity: usize,
}

#[cfg(feature = "http")]
impl HttpFrameSink {
    pub(crate) fn new(
        session: std::sync::Weak<crate::session::Session>,
        tx: crate::outbound_ring::RingSender<(u64, std::sync::Arc<Value>)>,
        sse_replay_capacity: usize,
    ) -> Self {
        Self {
            session,
            tx,
            sse_replay_capacity,
        }
    }
}

#[cfg(feature = "http")]
#[async_trait]
impl OutboundFrameSink for HttpFrameSink {
    /// Push one outbound JSON-RPC frame onto the per-session ring,
    /// allocating a monotonic SSE `event_id` for it and (when the
    /// replay buffer is enabled) appending the `(event_id, frame)`
    /// pair to the buffer's drop-oldest deque under the same lock.
    ///
    /// The session's `sse_replay_buffer` Mutex is the per-session
    /// outbound serialization point: id allocation, buffer push, and
    /// ring push all happen under the guard so concurrent producers
    /// cannot observe id allocation out of ring/buffer order. The
    /// guard is taken even when `sse_replay_capacity == 0` (no
    /// buffer write) to preserve ring-publish ordering, then dropped
    /// before any tracing/metric work outside the critical section.
    /// The buffer entry shares the same `Arc<Value>` payload as the
    /// ring entry: `Arc::clone` inside the guard is a refcount bump,
    /// not a JSON tree walk.
    ///
    /// The Mutex is `parking_lot::Mutex`: the guard is acquired
    /// synchronously and released by an explicit `drop(buffer)`
    /// before any tracing or metric emission, so the surrounding
    /// `async fn` never holds a sync lock across an `.await` point.
    async fn send_frame(&self, frame: std::sync::Arc<Value>) -> Result<(), OutboundSinkError> {
        if self.tx.is_receiver_dropped() {
            return Err(OutboundSinkError::TransportClosed);
        }
        let Some(session) = self.session.upgrade() else {
            return Err(OutboundSinkError::TransportClosed);
        };
        let buffered = (self.sse_replay_capacity > 0).then(|| std::sync::Arc::clone(&frame));
        let dropped = {
            let mut buffer = session.sse_replay_buffer.lock();
            let event_id = session
                .next_event_id
                .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
            if let Some(buffered) = buffered {
                if buffer.len() >= self.sse_replay_capacity {
                    buffer.pop_front();
                }
                buffer.push_back((event_id, buffered));
            }
            let dropped = self.tx.push((event_id, frame));
            drop(buffer);
            dropped
        };
        if dropped > 0 {
            tracing::warn!(
                target = "klieo::mcp::outbound",
                policy = "drop_oldest",
                dropped = dropped,
                "outbound ring full; dropped oldest frame(s)"
            );
            metrics::counter!(
                "klieo_mcp_outbound_dropped_total",
                "policy" => "oldest"
            )
            .increment(dropped as u64);
        }
        Ok(())
    }
}

/// In-memory writer used exclusively by `bench_stdio_sink`.
/// Not part of the public API.
#[cfg(feature = "bench")]
struct BenchWriter(std::sync::Arc<tokio::sync::Mutex<Vec<u8>>>);

#[cfg(feature = "bench")]
impl tokio::io::AsyncWrite for BenchWriter {
    fn poll_write(
        self: std::pin::Pin<&mut Self>,
        _cx: &mut std::task::Context<'_>,
        buf: &[u8],
    ) -> std::task::Poll<std::io::Result<usize>> {
        if let Ok(mut guard) = self.0.try_lock() {
            guard.extend_from_slice(buf);
        }
        std::task::Poll::Ready(Ok(buf.len()))
    }

    fn poll_flush(
        self: std::pin::Pin<&mut Self>,
        _cx: &mut std::task::Context<'_>,
    ) -> std::task::Poll<std::io::Result<()>> {
        std::task::Poll::Ready(Ok(()))
    }

    fn poll_shutdown(
        self: std::pin::Pin<&mut Self>,
        _cx: &mut std::task::Context<'_>,
    ) -> std::task::Poll<std::io::Result<()>> {
        std::task::Poll::Ready(Ok(()))
    }
}

/// Returns a `StdioFrameSink` as `Arc<dyn OutboundFrameSink>` writing
/// into a shared `Vec<u8>` buffer. Available only under the `bench` feature.
#[cfg(feature = "bench")]
pub fn bench_stdio_sink() -> (
    std::sync::Arc<dyn OutboundFrameSink>,
    std::sync::Arc<tokio::sync::Mutex<Vec<u8>>>,
) {
    let buf = std::sync::Arc::new(tokio::sync::Mutex::new(Vec::new()));
    let writer: crate::outbound::SharedWriter =
        std::sync::Arc::new(tokio::sync::Mutex::new(BenchWriter(buf.clone())));
    (std::sync::Arc::new(StdioFrameSink::new(writer)), buf)
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::Arc;
    use tokio::sync::Mutex;

    #[tokio::test]
    async fn stdio_sink_writes_newline_delimited_payload() {
        let buffer: Arc<Mutex<Vec<u8>>> = Arc::new(Mutex::new(Vec::new()));
        let writer: SharedWriter = Arc::new(Mutex::new(CapturingWriter(buffer.clone())));
        let sink = StdioFrameSink::new(writer);

        sink.send_frame(std::sync::Arc::new(
            serde_json::json!({"jsonrpc": "2.0", "id": 7}),
        ))
        .await
        .expect("write succeeds against in-memory writer");

        let written = buffer.lock().await.clone();
        let text = String::from_utf8(written).expect("frame is utf-8");
        assert!(
            text.ends_with('\n'),
            "stdio framing requires a trailing newline, got {text:?}"
        );
        let parsed: serde_json::Value =
            serde_json::from_str(text.trim_end()).expect("frame is valid json");
        assert_eq!(parsed["jsonrpc"], "2.0");
        assert_eq!(parsed["id"], 7);
    }

    #[tokio::test]
    async fn stdio_sink_returns_transport_closed_on_write_error() {
        let writer: SharedWriter = Arc::new(Mutex::new(BrokenWriter));
        let sink = StdioFrameSink::new(writer);

        let outcome = sink
            .send_frame(std::sync::Arc::new(serde_json::json!({"jsonrpc": "2.0"})))
            .await;
        assert!(matches!(outcome, Err(OutboundSinkError::TransportClosed)));
    }

    struct CapturingWriter(Arc<Mutex<Vec<u8>>>);

    impl tokio::io::AsyncWrite for CapturingWriter {
        fn poll_write(
            self: std::pin::Pin<&mut Self>,
            _cx: &mut std::task::Context<'_>,
            buf: &[u8],
        ) -> std::task::Poll<std::io::Result<usize>> {
            // Lock without awaiting: this is a test-only writer and the
            // mutex is uncontended within a single send_frame call.
            let mut guard = self.0.try_lock().expect("test writer is uncontended");
            guard.extend_from_slice(buf);
            std::task::Poll::Ready(Ok(buf.len()))
        }

        fn poll_flush(
            self: std::pin::Pin<&mut Self>,
            _cx: &mut std::task::Context<'_>,
        ) -> std::task::Poll<std::io::Result<()>> {
            std::task::Poll::Ready(Ok(()))
        }

        fn poll_shutdown(
            self: std::pin::Pin<&mut Self>,
            _cx: &mut std::task::Context<'_>,
        ) -> std::task::Poll<std::io::Result<()>> {
            std::task::Poll::Ready(Ok(()))
        }
    }

    struct BrokenWriter;

    impl tokio::io::AsyncWrite for BrokenWriter {
        fn poll_write(
            self: std::pin::Pin<&mut Self>,
            _cx: &mut std::task::Context<'_>,
            _buf: &[u8],
        ) -> std::task::Poll<std::io::Result<usize>> {
            std::task::Poll::Ready(Err(std::io::Error::new(
                std::io::ErrorKind::BrokenPipe,
                "test broken pipe",
            )))
        }

        fn poll_flush(
            self: std::pin::Pin<&mut Self>,
            _cx: &mut std::task::Context<'_>,
        ) -> std::task::Poll<std::io::Result<()>> {
            std::task::Poll::Ready(Ok(()))
        }

        fn poll_shutdown(
            self: std::pin::Pin<&mut Self>,
            _cx: &mut std::task::Context<'_>,
        ) -> std::task::Poll<std::io::Result<()>> {
            std::task::Poll::Ready(Ok(()))
        }
    }

    #[cfg(feature = "http")]
    mod http_sink_tests {
        use super::*;
        use crate::outbound_ring::bounded_ring;

        #[tokio::test]
        async fn http_send_delivers_through_ring() {
            let session = std::sync::Arc::new(crate::session::Session::new_stdio());
            let (tx, mut rx) = bounded_ring::<(u64, std::sync::Arc<Value>)>(4);
            let sink = HttpFrameSink::new(std::sync::Arc::downgrade(&session), tx, 8);
            sink.send_frame(std::sync::Arc::new(
                serde_json::json!({"jsonrpc": "2.0", "id": 1}),
            ))
            .await
            .expect("send_frame succeeds against open ring");
            let (event_id, received) = rx.recv().await.expect("ring delivers value");
            assert_eq!(event_id, 1, "first frame gets event id 1");
            assert_eq!(received["id"], 1);
        }

        #[tokio::test]
        async fn http_send_returns_transport_closed_when_rx_dropped() {
            let session = std::sync::Arc::new(crate::session::Session::new_stdio());
            let (tx, rx) = bounded_ring::<(u64, std::sync::Arc<Value>)>(4);
            drop(rx);
            let sink = HttpFrameSink::new(std::sync::Arc::downgrade(&session), tx, 8);
            let outcome = sink
                .send_frame(std::sync::Arc::new(serde_json::json!({"id": 1})))
                .await;
            assert!(matches!(outcome, Err(OutboundSinkError::TransportClosed)));
        }

        /// `HttpFrameSink` holds a `Weak<Session>` so the registry
        /// (DELETE, idle reaper, SSE-body drop) can release the strong
        /// owner while the sink itself is still reachable through
        /// `OutboundRequests`. After the strong owner drops, the
        /// `upgrade()` inside `send_frame` returns `None` and the sink
        /// surfaces `TransportClosed` instead of touching a freed
        /// session's atomics or buffer.
        #[tokio::test]
        async fn http_send_returns_transport_closed_when_session_dropped() {
            let session = std::sync::Arc::new(crate::session::Session::new_stdio());
            let weak = std::sync::Arc::downgrade(&session);
            let (tx, _rx) = bounded_ring::<(u64, std::sync::Arc<Value>)>(4);
            let sink = HttpFrameSink::new(weak, tx, 8);
            drop(session);
            let outcome = sink
                .send_frame(std::sync::Arc::new(serde_json::json!({"id": 1})))
                .await;
            assert!(matches!(outcome, Err(OutboundSinkError::TransportClosed)));
        }

        #[tokio::test]
        async fn http_send_drops_oldest_when_ring_full() {
            let session = std::sync::Arc::new(crate::session::Session::new_stdio());
            let (tx, mut rx) = bounded_ring::<(u64, std::sync::Arc<Value>)>(2);
            let sink = HttpFrameSink::new(std::sync::Arc::downgrade(&session), tx.clone(), 8);
            for id in 1..=3 {
                sink.send_frame(std::sync::Arc::new(serde_json::json!({"id": id})))
                    .await
                    .expect("ring accepts each push");
            }
            assert_eq!(tx.dropped_oldest_count(), 1);
            let (_, first) = rx.recv().await.unwrap();
            let (_, second) = rx.recv().await.unwrap();
            assert_eq!(first["id"], 2);
            assert_eq!(second["id"], 3);
        }

        /// `sse_replay_capacity > 0` populates the session's SSE
        /// replay buffer with `(event_id, frame)` entries in send
        /// order. A reconnecting client snapshots this buffer to
        /// replay frames the original SSE body never delivered.
        #[tokio::test]
        async fn http_send_writes_to_sse_replay_buffer() {
            let session = std::sync::Arc::new(crate::session::Session::new_stdio());
            let (tx, _rx) = bounded_ring::<(u64, std::sync::Arc<Value>)>(4);
            let sink = HttpFrameSink::new(std::sync::Arc::downgrade(&session), tx, 8);
            for i in 1..=3 {
                sink.send_frame(std::sync::Arc::new(serde_json::json!({"id": i})))
                    .await
                    .expect("send_frame ok");
            }
            let buffer = session.sse_replay_buffer.lock();
            assert_eq!(buffer.len(), 3);
            let ids: Vec<u64> = buffer.iter().map(|(id, _)| *id).collect();
            assert_eq!(ids, vec![1, 2, 3]);
        }

        /// `sse_replay_capacity == 0` disables resumption: the sink
        /// skips the buffer-lock + push entirely, leaving the deque
        /// empty no matter how many frames flow through.
        #[tokio::test]
        async fn http_send_skips_sse_replay_buffer_when_disabled() {
            let session = std::sync::Arc::new(crate::session::Session::new_stdio());
            let (tx, _rx) = bounded_ring::<(u64, std::sync::Arc<Value>)>(4);
            let sink = HttpFrameSink::new(std::sync::Arc::downgrade(&session), tx, 0);
            sink.send_frame(std::sync::Arc::new(serde_json::json!({"id": 1})))
                .await
                .expect("send_frame ok");
            let buffer = session.sse_replay_buffer.lock();
            assert!(buffer.is_empty(), "sse_replay_capacity=0 disables writes");
        }

        #[test]
        fn capacity_constant_is_named_and_nonzero() {
            // Sanity: cap must be > 0 (zero would drop every push) and
            // at least 256 (smaller risks dropping under normal MCP-
            // notification bursts). Compile-time check.
            const _: () = assert!(OUTBOUND_QUEUE_CAPACITY > 0);
            const _: () = assert!(OUTBOUND_QUEUE_CAPACITY >= 256);
        }
    }
}