use std::collections::HashMap;
use std::collections::HashSet;
use std::process::Stdio;
use std::sync::Arc;
use std::sync::Mutex as StdMutex;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::Duration;
use super::AgentRequest;
use serde::Serialize;
use serde::de::DeserializeOwned;
use serde_json::Value;
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
use tokio::process::{Child, ChildStdin, ChildStdout, Command};
use tokio::sync::{Mutex, mpsc, oneshot};
use super::error::AcpPeerError;
use super::reply::AcpReply;
use super::stream::{AgentRequestStream, NotificationStream};
use crate::domain::acp::method::AcpMethod;
use crate::domain::acp::notification::SessionNotification;
#[derive(Debug, Clone)]
pub struct SpawnSpec {
pub program: String,
pub args: Vec<String>,
pub timeout: Option<Duration>,
}
impl SpawnSpec {
pub fn new(program: impl Into<String>, args: Vec<String>) -> Self {
Self {
program: program.into(),
args,
timeout: None,
}
}
#[must_use]
pub const fn with_timeout(mut self, timeout: Duration) -> Self {
self.timeout = Some(timeout);
self
}
}
struct Pending {
tx: oneshot::Sender<Result<AcpReply, AcpPeerError>>,
method: AcpMethod,
}
impl std::fmt::Debug for Pending {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter
.debug_struct("Pending")
.field("method", &self.method)
.finish_non_exhaustive()
}
}
pub struct AcpStdioPeer {
pending: Arc<Mutex<HashMap<u64, Pending>>>,
next_id: AtomicU64,
write_tx: mpsc::Sender<Vec<u8>>,
notifications_rx: Arc<tokio::sync::Mutex<mpsc::Receiver<SessionNotification>>>,
sent_ids: Arc<StdMutex<HashSet<u64>>>,
agent_requests_rx: Arc<tokio::sync::Mutex<mpsc::Receiver<AgentRequest>>>,
timeout: Duration,
#[expect(
dead_code,
reason = "the field is held to keep the `Child` alive; the `kill_on_drop` machinery reads it on drop"
)]
child: Child,
}
impl std::fmt::Debug for AcpStdioPeer {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter
.debug_struct("AcpStdioPeer")
.field("next_id", &self.next_id)
.field("timeout", &self.timeout)
.finish_non_exhaustive()
}
}
impl AcpStdioPeer {
pub fn spawn(spec: &SpawnSpec) -> Result<Self, AcpPeerError> {
let mut command = Command::new(&spec.program);
command
.args(&spec.args)
.stdin(Stdio::piped())
.stdout(Stdio::piped())
.stderr(Stdio::piped())
.kill_on_drop(true);
let mut child: Child = command.spawn().map_err(AcpPeerError::Spawn)?;
let stdin: ChildStdin = child
.stdin
.take()
.ok_or_else(|| AcpPeerError::Spawn(io_other("stdin was not piped")))?;
let stdout: ChildStdout = child
.stdout
.take()
.ok_or_else(|| AcpPeerError::Spawn(io_other("stdout was not piped")))?;
let pending: Arc<Mutex<HashMap<u64, Pending>>> = Arc::new(Mutex::new(HashMap::new()));
let (write_tx, write_rx) = mpsc::channel::<Vec<u8>>(32);
let (notifications_tx, notifications_rx) = mpsc::channel::<SessionNotification>(32);
let notifications_rx = Arc::new(tokio::sync::Mutex::new(notifications_rx));
let (agent_requests_tx, agent_requests_rx) = mpsc::channel::<AgentRequest>(8);
let agent_requests_rx = Arc::new(tokio::sync::Mutex::new(agent_requests_rx));
let sent_ids: Arc<StdMutex<HashSet<u64>>> = Arc::new(StdMutex::new(HashSet::new()));
spawn_reader(
BufReader::new(stdout),
Arc::clone(&pending),
notifications_tx,
agent_requests_tx,
Arc::clone(&sent_ids),
);
spawn_writer(stdin, write_rx);
Ok(Self {
pending,
next_id: AtomicU64::new(1),
write_tx,
notifications_rx,
sent_ids,
agent_requests_rx,
timeout: spec.timeout.unwrap_or(Duration::from_secs(30)),
child,
})
}
pub async fn send_request<P>(
&self,
method: AcpMethod,
params: P,
) -> Result<AcpReply, AcpPeerError>
where
P: Serialize,
{
let id = self.next_id.fetch_add(1, Ordering::Relaxed);
self.sent_ids
.lock()
.map_err(|error| {
AcpPeerError::Decode(serde_json::Error::io(io_other(&format!(
"sent_ids lock poisoned: {error}"
))))
})?
.insert(id);
let params_value: Value = serde_json::to_value(¶ms).map_err(AcpPeerError::Decode)?;
let envelope = serde_json::json!({
"jsonrpc": "2.0",
"id": id,
"method": method.as_wire_name(),
"params": params_value,
});
let bytes = serde_json::to_vec(&envelope).map_err(AcpPeerError::Decode)?;
let mut framed = bytes;
framed.push(b'\n');
self.write_tx
.send(framed)
.await
.map_err(|_| AcpPeerError::Write(io_other("the writer task is gone")))?;
let (tx, rx) = oneshot::channel();
self.pending.lock().await.insert(id, Pending { tx, method });
match tokio::time::timeout(self.timeout, rx).await {
Ok(Ok(result)) => {
if let Ok(mut sent_ids) = self.sent_ids.lock() {
sent_ids.remove(&id);
}
result
}
Ok(Err(_oneshot_dropped)) => {
if let Ok(mut sent_ids) = self.sent_ids.lock() {
sent_ids.remove(&id);
}
Err(AcpPeerError::Timeout(self.timeout))
}
Err(_elapsed) => {
self.pending.lock().await.remove(&id);
if let Ok(mut sent_ids) = self.sent_ids.lock() {
sent_ids.remove(&id);
}
Err(AcpPeerError::Timeout(self.timeout))
}
}
}
pub async fn send_typed_request<Req, Resp>(
&self,
method: AcpMethod,
params: Req,
) -> Result<Resp, AcpPeerError>
where
Req: Serialize,
Resp: DeserializeOwned,
{
let reply = self.send_request(method, params).await?;
serde_json::from_value(reply.result).map_err(AcpPeerError::Decode)
}
pub async fn send_notification<Params>(
&self,
method: AcpMethod,
params: Params,
) -> Result<(), AcpPeerError>
where
Params: Serialize,
{
let params_value: Value = serde_json::to_value(¶ms).map_err(AcpPeerError::Decode)?;
let envelope = serde_json::json!({
"jsonrpc": "2.0",
"method": method.as_wire_name(),
"params": params_value,
});
let bytes = serde_json::to_vec(&envelope).map_err(AcpPeerError::Decode)?;
let mut framed = bytes;
framed.push(b'\n');
self.write_tx
.send(framed)
.await
.map_err(|_| AcpPeerError::Write(io_other("the writer task is gone")))?;
Ok(())
}
pub async fn recv_notification(&self) -> Option<SessionNotification> {
let mut guard = self.notifications_rx.lock().await;
guard.recv().await
}
pub async fn recv_agent_request(&self) -> Option<AgentRequest> {
let mut guard = self.agent_requests_rx.lock().await;
guard.recv().await
}
pub fn notifications(&self) -> NotificationStream<'_> {
NotificationStream::new(Arc::clone(&self.notifications_rx))
}
pub fn agent_requests(&self) -> AgentRequestStream<'_> {
AgentRequestStream::new(Arc::clone(&self.agent_requests_rx))
}
pub async fn send_reply(&self, id: Value, result: Value) -> Result<(), AcpPeerError> {
let envelope = serde_json::json!({
"jsonrpc": "2.0",
"id": id,
"result": result,
});
let bytes = serde_json::to_vec(&envelope).map_err(AcpPeerError::Decode)?;
let mut framed = bytes;
framed.push(b'\n');
self.write_tx
.send(framed)
.await
.map_err(|_| AcpPeerError::Write(io_other("the writer task is gone")))?;
Ok(())
}
}
fn spawn_reader(
mut reader: BufReader<ChildStdout>,
pending: Arc<Mutex<HashMap<u64, Pending>>>,
notifications_tx: mpsc::Sender<SessionNotification>,
agent_requests_tx: mpsc::Sender<AgentRequest>,
sent_ids: Arc<StdMutex<HashSet<u64>>>,
) {
tokio::spawn(async move {
let mut line = String::new();
loop {
line.clear();
match reader.read_line(&mut line).await {
Ok(0) | Err(_) => return,
Ok(_) => {}
}
let trimmed = line.trim_end_matches('\n');
if trimmed.is_empty() {
continue;
}
let value: Value = match serde_json::from_str(trimmed) {
Ok(value) => value,
Err(_) => continue,
};
if value.get("id").is_some() {
let (id, result) = parse_and_dispatch_value(value.clone());
let is_client_reply = id != 0 && sent_ids.lock().is_ok_and(|set| set.contains(&id));
if is_client_reply {
let entry = pending.lock().await.remove(&id);
if let Some(entry) = entry {
let _ = entry.tx.send(result);
}
} else if id != 0
&& value.get("method").is_some()
&& let Some(request) = parse_agent_request(&value)
{
let _ = agent_requests_tx.send(request).await;
}
} else {
if value.get("method").and_then(Value::as_str) == Some("session/update")
&& let Some(params) = value.get("params")
&& let Ok(notification) =
serde_json::from_value::<SessionNotification>(params.clone())
{
let _ = notifications_tx.send(notification).await;
}
}
}
});
}
fn parse_agent_request(value: &Value) -> Option<AgentRequest> {
let id = value.get("id")?.clone();
let method = value.get("method").and_then(Value::as_str)?.to_owned();
let params = value.get("params").cloned().unwrap_or(Value::Null);
Some(AgentRequest { id, method, params })
}
fn parse_and_dispatch_value(value: Value) -> (u64, Result<AcpReply, AcpPeerError>) {
let id = value.get("id").and_then(Value::as_u64).unwrap_or(0);
if value.get("error").is_some() {
let code = value
.get("error")
.and_then(|error| error.get("code"))
.and_then(Value::as_i64)
.unwrap_or(0);
let message = value
.get("error")
.and_then(|error| error.get("message"))
.and_then(Value::as_str)
.unwrap_or("")
.to_owned();
return (id, Err(AcpPeerError::JsonRpc { code, message }));
}
match serde_json::from_value::<AcpReply>(value) {
Ok(reply) => (id, Ok(reply)),
Err(error) => (id, Err(AcpPeerError::Decode(error))),
}
}
fn spawn_writer(stdin: ChildStdin, mut write_rx: mpsc::Receiver<Vec<u8>>) {
tokio::spawn(async move {
let mut stdin = stdin;
while let Some(bytes) = write_rx.recv().await {
if stdin.write_all(&bytes).await.is_err() {
return;
}
if stdin.flush().await.is_err() {
return;
}
}
});
}
fn io_other(message: &str) -> std::io::Error {
std::io::Error::other(message)
}
#[cfg(test)]
mod tests {
use super::*;
#[derive(Debug)]
enum Failure {
WrongResult(&'static str),
}
impl std::fmt::Display for Failure {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::WrongResult(expected) => write!(formatter, "{expected}"),
}
}
}
impl std::error::Error for Failure {}
impl From<serde_json::Error> for Failure {
fn from(_error: serde_json::Error) -> Self {
Self::WrongResult("a JSON parse failed")
}
}
#[test]
fn parse_and_dispatch_routes_a_result_by_id() -> Result<(), Failure> {
let value: Value =
serde_json::from_str(r#"{"jsonrpc":"2.0","id":7,"result":{"answer":42}}"#)
.map_err(|_| Failure::WrongResult("valid JSON"))?;
let (id, result) = parse_and_dispatch_value(value);
assert_eq!(id, 7);
let Ok(reply) = result else {
return Err(Failure::WrongResult("Ok reply"));
};
assert_eq!(reply.id, serde_json::json!(7));
assert_eq!(reply.result, serde_json::json!({"answer": 42}));
Ok(())
}
#[test]
fn parse_and_dispatch_routes_an_error_by_id() -> Result<(), Failure> {
let value: Value = serde_json::from_str(
r#"{"jsonrpc":"2.0","id":7,"error":{"code":-32601,"message":"Method not found"}}"#,
)
.map_err(|_| Failure::WrongResult("valid JSON"))?;
let (id, result) = parse_and_dispatch_value(value);
assert_eq!(id, 7);
match result {
Err(AcpPeerError::JsonRpc { code, message }) => {
assert_eq!(code, -32601);
assert_eq!(message, "Method not found");
Ok(())
}
Err(_) => Err(Failure::WrongResult("JsonRpc variant")),
Ok(_) => Err(Failure::WrongResult("Err variant")),
}
}
#[test]
fn parse_and_dispatch_treats_missing_id_as_zero() -> Result<(), Failure> {
let value: Value = serde_json::from_str(
r#"{"jsonrpc":"2.0","error":{"code":-32600,"message":"Invalid Request"}}"#,
)
.map_err(|_| Failure::WrongResult("valid JSON"))?;
let (id, _result) = parse_and_dispatch_value(value);
assert_eq!(id, 0);
Ok(())
}
#[cfg(unix)]
#[tokio::test]
async fn the_peer_spawns_and_drops_cleanly() -> Result<(), Failure> {
let spec = SpawnSpec::new("cat", Vec::new());
let Ok(peer) = AcpStdioPeer::spawn(&spec) else {
return Err(Failure::WrongResult("spawn succeeds on Unix"));
};
drop(peer);
Ok(())
}
#[test]
fn a_spawn_spec_carries_the_program_and_args() {
let spec = SpawnSpec::new("node", vec!["agent.js".to_owned()])
.with_timeout(Duration::from_secs(5));
assert_eq!(spec.program, "node");
assert_eq!(spec.args, vec!["agent.js".to_owned()]);
assert_eq!(spec.timeout, Some(Duration::from_secs(5)));
}
}