use super::super::wire::{self, Frame};
use serde_json::json;
use tokio::io::{AsyncBufRead, AsyncWrite, AsyncWriteExt as _};
pub(super) async fn serve(
mut input: impl AsyncBufRead + Unpin,
mut output: impl AsyncWrite + Unpin,
bound: usize,
session: String,
) -> anyhow::Result<()> {
let mut initialized = false;
while let Some(frame) = wire::read(&mut input, bound).await? {
let (id, method, params) = match frame {
Frame::Request { id, method, params } => (id, method, params),
Frame::Notification { method, .. }
if method == "initialized" || method.starts_with('_') =>
{
continue;
}
_ => anyhow::bail!("lys_stdio_unexpected_frame"),
};
let result = match method.as_str() {
"initialize" if !initialized => {
let request = serde_json::from_value::<
agent_client_protocol_schema::v1::InitializeRequest,
>(params);
match request {
Ok(request)
if request.protocol_version
== agent_client_protocol_schema::ProtocolVersion::V1 =>
{
initialized = true;
Ok(
json!({"protocolVersion":1,"agentCapabilities":{"loadSession":false},"agentInfo":{"name":"haematite-assistant","version":env!("CARGO_PKG_VERSION")},"authMethods":[]}),
)
}
_ => Err((-32602, "lys_initialize_invalid_or_unsupported")),
}
}
"initialize" => Err((-32600, "lys_initialize_repeated")),
_ if !initialized => Err((-32002, "lys_initialize_required")),
"session/new" => {
match serde_json::from_value::<agent_client_protocol_schema::v1::NewSessionRequest>(
params,
) {
Ok(request) if request.cwd.is_absolute() && request.mcp_servers.is_empty() => {
Ok(json!({"sessionId":session}))
}
_ => Err((-32602, "lys_session_new_invalid_or_mcp_not_profile_bound")),
}
}
_ => Err((-32601, "lys_request_unsupported_use_the_store_door")),
};
let value = match result {
Ok(result) => json!({"jsonrpc":"2.0","id":id,"result":result}),
Err((code, message)) => {
json!({"jsonrpc":"2.0","id":id,"error":{"code":code,"message":message}})
}
};
let mut bytes = serde_json::to_vec(&value)?;
bytes.push(b'\n');
anyhow::ensure!(bytes.len() <= bound, "lys_stdio_answer_exceeds_frame_bound");
output
.write_all(&bytes)
.await
.map_err(|error| anyhow::anyhow!("lys_stdio_write_failed: {}", error.kind()))?;
output
.flush()
.await
.map_err(|error| anyhow::anyhow!("lys_stdio_flush_failed: {}", error.kind()))?;
}
Ok(())
}
use std::{
fs::File,
io,
pin::Pin,
task::{Context, Poll},
};
use tokio::io::{AsyncRead, ReadBuf, unix::AsyncFd};
pub(super) struct FileIo(AsyncFd<File>);
impl FileIo {
pub fn open(path: &str, write: bool) -> anyhow::Result<Self> {
use std::os::unix::fs::OpenOptionsExt as _;
let file = std::fs::OpenOptions::new()
.read(!write)
.write(write)
.custom_flags(i32::try_from(rustix::fs::OFlags::NONBLOCK.bits())?)
.open(path)
.map_err(|error| anyhow::anyhow!("lys_stdio_open_failed: {}", error.kind()))?;
Ok(Self(AsyncFd::new(file).map_err(|error| {
anyhow::anyhow!("lys_stdio_registration_failed: {}", error.kind())
})?))
}
}
impl AsyncRead for FileIo {
fn poll_read(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buffer: &mut ReadBuf<'_>,
) -> Poll<io::Result<()>> {
loop {
let mut ready = std::task::ready!(self.0.poll_read_ready(cx))?;
if let Ok(result) =
ready.try_io(|fd| super::read_once(fd.get_ref(), buffer.initialize_unfilled()))
{
buffer.advance(result?);
return Poll::Ready(Ok(()));
}
}
}
}
impl AsyncWrite for FileIo {
fn poll_write(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
bytes: &[u8],
) -> Poll<io::Result<usize>> {
loop {
let mut ready = std::task::ready!(self.0.poll_write_ready(cx))?;
if let Ok(result) =
ready.try_io(|fd| rustix::io::write(fd.get_ref(), bytes).map_err(io::Error::from))
{
return Poll::Ready(result);
}
}
}
fn poll_flush(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<io::Result<()>> {
Poll::Ready(Ok(()))
}
fn poll_shutdown(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<io::Result<()>> {
Poll::Ready(Ok(()))
}
}
#[cfg(test)]
#[path = "input_protocol_tests.rs"]
mod tests;