a8e-acp 0.3.0

Articulate (a8e): The sovereign AI operator for your terminal.
//! Shared adapter classes for converting mpsc channels to AsyncRead/AsyncWrite streams
//! Used by both HTTP and WebSocket transports

use std::{
    pin::Pin,
    task::{Context, Poll},
};
use tokio::sync::mpsc;
use tracing::error;

/// Converts an mpsc::Receiver<String> to AsyncRead
/// Each message is terminated with a newline for JSON-RPC framing
pub(crate) struct ReceiverToAsyncRead {
    rx: mpsc::Receiver<String>,
    buffer: Vec<u8>,
    pos: usize,
}

impl ReceiverToAsyncRead {
    pub(crate) fn new(rx: mpsc::Receiver<String>) -> Self {
        Self {
            rx,
            buffer: Vec::new(),
            pos: 0,
        }
    }
}

impl tokio::io::AsyncRead for ReceiverToAsyncRead {
    fn poll_read(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &mut tokio::io::ReadBuf<'_>,
    ) -> Poll<std::io::Result<()>> {
        if self.pos < self.buffer.len() {
            let remaining = &self.buffer[self.pos..];
            let to_copy = remaining.len().min(buf.remaining());
            buf.put_slice(&remaining[..to_copy]);
            self.pos += to_copy;
            if self.pos >= self.buffer.len() {
                self.buffer.clear();
                self.pos = 0;
            }
            return Poll::Ready(Ok(()));
        }

        match Pin::new(&mut self.rx).poll_recv(cx) {
            Poll::Ready(Some(msg)) => {
                let bytes = format!("{}\n", msg).into_bytes();
                let to_copy = bytes.len().min(buf.remaining());
                buf.put_slice(&bytes[..to_copy]);
                if to_copy < bytes.len() {
                    self.buffer = bytes[to_copy..].to_vec();
                    self.pos = 0;
                }
                Poll::Ready(Ok(()))
            }
            Poll::Ready(None) => Poll::Ready(Ok(())),
            Poll::Pending => Poll::Pending,
        }
    }
}

/// Converts an mpsc::Sender<String> to AsyncWrite
/// Splits incoming data on newlines for JSON-RPC framing
pub(crate) struct SenderToAsyncWrite {
    tx: mpsc::Sender<String>,
    buffer: Vec<u8>,
}

impl SenderToAsyncWrite {
    pub(crate) fn new(tx: mpsc::Sender<String>) -> Self {
        Self {
            tx,
            buffer: Vec::new(),
        }
    }
}

impl tokio::io::AsyncWrite for SenderToAsyncWrite {
    fn poll_write(
        mut self: Pin<&mut Self>,
        _cx: &mut Context<'_>,
        buf: &[u8],
    ) -> Poll<std::io::Result<usize>> {
        self.buffer.extend_from_slice(buf);

        while let Some(pos) = self.buffer.iter().position(|&b| b == b'\n') {
            let line = String::from_utf8_lossy(&self.buffer[..pos]).to_string();
            self.buffer.drain(..=pos);

            if !line.is_empty() {
                if let Err(e) = self.tx.try_send(line.clone()) {
                    match e {
                        mpsc::error::TrySendError::Full(_) => {
                            let truncated: String = line.chars().take(100).collect();
                            error!(
                                "Channel full, dropping message (backpressure): {}",
                                truncated
                            );
                        }
                        mpsc::error::TrySendError::Closed(_) => {
                            return Poll::Ready(Err(std::io::Error::new(
                                std::io::ErrorKind::BrokenPipe,
                                "Channel closed",
                            )));
                        }
                    }
                }
            }
        }

        Poll::Ready(Ok(buf.len()))
    }

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

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