use serde_json::{Value, json};
use tokio::io::{
AsyncBufRead, AsyncBufReadExt as _, AsyncRead, AsyncWrite, AsyncWriteExt as _, BufReader,
};
use tokio::sync::mpsc;
use tokio::task::JoinHandle;
#[derive(Clone, Copy, serde::Serialize, serde::Deserialize)]
#[serde(rename_all = "snake_case")]
pub(super) enum ToolStatus {
Pending,
InProgress,
Completed,
Failed,
Interrupted,
}
impl ToolStatus {
pub const fn terminal(self) -> bool {
matches!(self, Self::Completed | Self::Failed | Self::Interrupted)
}
}
pub(super) struct RpcError {
pub code: i64,
pub message: String,
}
pub(super) enum Frame {
Request {
id: Value,
method: String,
params: Value,
},
Notification {
method: String,
params: Value,
},
Response {
id: Value,
result: Result<Value, RpcError>,
},
}
pub(super) async fn read<R: AsyncBufRead + Unpin>(
reader: &mut R,
bound: usize,
) -> anyhow::Result<Option<Frame>> {
let mut line = Vec::new();
loop {
let buffer = reader
.fill_buf()
.await
.map_err(|error| anyhow::anyhow!("ACP read failed: {}", error.kind()))?;
if buffer.is_empty() {
anyhow::ensure!(line.is_empty(), "ACP frame ended before newline");
return Ok(None);
}
let count = buffer
.iter()
.position(|byte| *byte == b'\n')
.map_or(buffer.len(), |index| index + 1);
anyhow::ensure!(
count <= bound.saturating_sub(line.len()),
"ACP frame exceeds acp_frame_bytes {bound}"
);
let complete = buffer[count - 1] == b'\n';
line.extend_from_slice(&buffer[..count]);
reader.consume(count);
if complete {
break;
}
}
let value: Value =
serde_json::from_slice(&line).map_err(|_| anyhow::anyhow!("ACP frame is not JSON"))?;
parse(value).map(Some)
}
fn parse(mut value: Value) -> anyhow::Result<Frame> {
let object = value
.as_object_mut()
.ok_or_else(|| anyhow::anyhow!("ACP frame is not an object"))?;
anyhow::ensure!(
object.get("jsonrpc").and_then(Value::as_str) == Some("2.0"),
"ACP frame requires jsonrpc 2.0"
);
let id = object.remove("id");
anyhow::ensure!(
id.as_ref()
.is_none_or(|id| id.is_string() || id.as_i64().is_some() || id.as_u64().is_some()),
"ACP id is not an integer or string"
);
if let Some(method) = object.remove("method") {
let method = method
.as_str()
.ok_or_else(|| anyhow::anyhow!("ACP method is not text"))?
.to_owned();
anyhow::ensure!(
!object.contains_key("result") && !object.contains_key("error"),
"ACP request carries a response"
);
let params = object.remove("params").unwrap_or_else(|| json!({}));
anyhow::ensure!(params.is_object(), "ACP params is not an object");
return Ok(match id {
Some(id) => Frame::Request { id, method, params },
None => Frame::Notification { method, params },
});
}
let id = id.ok_or_else(|| anyhow::anyhow!("ACP response has no id"))?;
let result = match (object.remove("result"), object.remove("error")) {
(Some(result), None) => Ok(result),
(None, Some(error)) => Err(RpcError {
code: error
.get("code")
.and_then(Value::as_i64)
.ok_or_else(|| anyhow::anyhow!("ACP error has no numeric code"))?,
message: error
.get("message")
.and_then(Value::as_str)
.ok_or_else(|| anyhow::anyhow!("ACP error has no message"))?
.to_owned(),
}),
_ => anyhow::bail!("ACP response requires exactly one result or error"),
};
Ok(Frame::Response { id, result })
}
#[derive(Clone)]
pub(super) struct Output {
sender: mpsc::Sender<PendingWrite>,
bound: usize,
}
impl Output {
pub fn send(&self, message: Value) -> anyhow::Result<()> {
anyhow::ensure!(
serde_json::to_vec(&message)?.len().saturating_add(1) <= self.bound,
"ACP output exceeds acp_frame_bytes {}",
self.bound
);
self.sender
.try_send(PendingWrite {
value: message,
written: None,
})
.map_err(|error| match error {
mpsc::error::TrySendError::Full(_) => anyhow::anyhow!("acp_write_queue_full"),
mpsc::error::TrySendError::Closed(_) => anyhow::anyhow!("acp_writer_closed"),
})
}
pub fn response_tracked(&self, id: Value, result: Value, key: String) -> anyhow::Result<()> {
let value = envelope([("id", id), ("result", result)]);
anyhow::ensure!(
serde_json::to_vec(&value)?.len().saturating_add(1) <= self.bound,
"ACP output exceeds acp_frame_bytes"
);
self.sender
.try_send(PendingWrite {
value,
written: Some(key),
})
.map_err(|error| match error {
mpsc::error::TrySendError::Full(_) => anyhow::anyhow!("acp_write_queue_full"),
mpsc::error::TrySendError::Closed(_) => anyhow::anyhow!("acp_writer_closed"),
})
}
pub fn cancelled(&self, id: Value) -> anyhow::Result<()> {
self.send(envelope([
("id", id),
("result", json!({"outcome":{"outcome":"cancelled"}})),
]))
}
pub fn request(&self, id: u64, method: &str, params: Value) -> anyhow::Result<()> {
self.send(envelope([
("id", Value::from(id)),
("method", Value::from(method)),
("params", params),
]))
}
pub fn notification(&self, method: &str, params: Value) -> anyhow::Result<()> {
self.send(envelope([
("method", Value::from(method)),
("params", params),
]))
}
#[cfg(test)]
pub fn response(&self, id: Value, result: Value) -> anyhow::Result<()> {
self.send(envelope([("id", id), ("result", result)]))
}
pub fn refuse(&self, id: Value, code: i64, words: &str) -> anyhow::Result<()> {
self.send(envelope([
("id", id),
("error", json!({"code": code, "message": words})),
]))
}
}
struct PendingWrite {
value: Value,
written: Option<String>,
}
pub(super) struct Written {
pub key: String,
pub delivered: bool,
}
pub(super) struct Pipe {
pub written: mpsc::Receiver<Written>,
pub input: mpsc::Receiver<anyhow::Result<Option<Frame>>>,
pub output: Output,
pub reader: JoinHandle<()>,
pub writer: JoinHandle<anyhow::Result<()>>,
}
impl Pipe {
pub fn start<R, W>(
read_half: R,
mut write_half: W,
bound: usize,
queue: usize,
read_queue: usize,
) -> Self
where
R: AsyncRead + Unpin + Send + 'static,
W: AsyncWrite + Unpin + Send + 'static,
{
let (sender, mut receiver) = mpsc::channel::<PendingWrite>(queue);
let (completed, written) = mpsc::channel(queue);
let writer = tokio::spawn(async move {
while let Some(message) = receiver.recv().await {
let mut bytes = serde_json::to_vec(&message.value)?;
bytes.push(b'\n');
let result = async {
write_half
.write_all(&bytes)
.await
.map_err(|error| anyhow::anyhow!("ACP write failed: {}", error.kind()))?;
write_half
.flush()
.await
.map_err(|error| anyhow::anyhow!("ACP flush failed: {}", error.kind()))
}
.await;
if let Some(key) = message.written {
completed
.try_send(Written {
key,
delivered: result.is_ok(),
})
.map_err(|error| match error {
mpsc::error::TrySendError::Full(_) => {
anyhow::anyhow!("acp_written_receipts_full")
}
mpsc::error::TrySendError::Closed(_) => {
anyhow::anyhow!("acp_written_receipts_closed")
}
})?;
}
result?;
}
write_half
.shutdown()
.await
.map_err(|error| anyhow::anyhow!("ACP shutdown failed: {}", error.kind()))
});
let (events, input) = mpsc::channel(read_queue);
let reader = tokio::spawn(async move {
let mut reader = BufReader::new(read_half);
loop {
let frame = read(&mut reader, bound).await;
let ended = !matches!(&frame, Ok(Some(_)));
if events.send(frame).await.is_err() || ended {
return;
}
}
});
Self {
written,
input,
output: Output { sender, bound },
reader,
writer,
}
}
}
fn envelope<const N: usize>(members: [(&str, Value); N]) -> Value {
let mut object = serde_json::Map::with_capacity(N + 1);
object.insert("jsonrpc".to_owned(), Value::from("2.0"));
for (name, value) in members {
object.insert(name.to_owned(), value);
}
Value::Object(object)
}
#[cfg(test)]
#[path = "wire_status_tests.rs"]
mod status_tests;