acp_utils/client/
tokio_agent.rs

1//! Tokio-native parent-side ACP transport.
2//!
3//! `agent_client_protocol::AcpAgent` spawns the child via smol's
4//! `async_process::Command`, which wraps stdio in `blocking::Unblock`. Inside a
5//! tokio runtime that causes a busy loop. This avoids the issue by spawning stdio agents with `tokio::process::Command`
6
7use agent_client_protocol::util::internal_error;
8use agent_client_protocol::{
9    AcpAgent, AcpAgentConfig, ByteStreams, ConnectTo, Error, INCOMING_TRANSPORT_CLOSED_REASON, Role,
10    is_incoming_transport_closed,
11};
12use std::path::PathBuf;
13use std::process::{ExitStatus, Stdio};
14use std::str::FromStr;
15use tokio::io::{AsyncBufReadExt, BufReader};
16use tokio::process::Command;
17use tokio::sync::oneshot;
18use tokio::time::{Duration, timeout};
19use tokio_util::compat::{TokioAsyncReadCompatExt, TokioAsyncWriteCompatExt};
20
21pub struct TokioAcpAgent {
22    config: AcpAgentConfig,
23}
24
25impl TokioAcpAgent {
26    pub fn from_command(command: impl Into<PathBuf>, args: Vec<String>) -> Self {
27        Self { config: AcpAgentConfig::new(command).args(args) }
28    }
29
30    pub fn config(&self) -> &AcpAgentConfig {
31        &self.config
32    }
33}
34
35impl<T: Role> ConnectTo<T> for TokioAcpAgent {
36    async fn connect_to(self, client: impl ConnectTo<T::Counterpart>) -> Result<(), Error> {
37        connect_stdio::<T>(self.config, client).await
38    }
39}
40
41impl FromStr for TokioAcpAgent {
42    type Err = Error;
43
44    fn from_str(s: &str) -> Result<Self, Self::Err> {
45        Ok(Self { config: AcpAgent::from_str(s)?.into_config() })
46    }
47}
48
49async fn connect_stdio<T: Role>(config: AcpAgentConfig, client: impl ConnectTo<T::Counterpart>) -> Result<(), Error> {
50    let (stdin, stdout, stderr, mut child) = {
51        let mut cmd = Command::new(config.command());
52        cmd.args(config.arguments());
53        for (name, value) in config.environment() {
54            cmd.env(name, value);
55        }
56
57        let mut child = cmd
58            .stdin(Stdio::piped())
59            .stdout(Stdio::piped())
60            .stderr(Stdio::piped())
61            .kill_on_drop(true)
62            .spawn()
63            .map_err(Error::into_internal_error)?;
64
65        let stdin = child.stdin.take().ok_or_else(|| internal_error("missing child stdin"))?;
66        let stdout = child.stdout.take().ok_or_else(|| internal_error("missing child stdout"))?;
67        let stderr = child.stderr.take().ok_or_else(|| internal_error("missing child stderr"))?;
68        (stdin, stdout, stderr, child)
69    };
70
71    let (stderr_tx, stderr_rx) = oneshot::channel::<String>();
72    tokio::spawn(async move {
73        let mut lines = BufReader::new(stderr).lines();
74        let mut buf = String::new();
75        while let Ok(Some(line)) = lines.next_line().await {
76            if !buf.is_empty() {
77                buf.push('\n');
78            }
79            buf.push_str(&line);
80        }
81        let _ = stderr_tx.send(buf);
82    });
83
84    let child_fut = async move {
85        let status = child.wait().await.map_err(Error::into_internal_error)?;
86        finish_child_exit(status, stderr_rx).await
87    };
88
89    let bytes = ByteStreams::new(stdin.compat_write(), stdout.compat());
90    let protocol_fut = ConnectTo::<T>::connect_to(bytes, client);
91    tokio::pin!(child_fut);
92
93    tokio::select! {
94        result = &mut child_fut => result,
95        result = protocol_fut => match result {
96            Ok(()) => timeout(SHUTDOWN_GRACE_PERIOD, &mut child_fut).await.unwrap_or(Ok(())),
97            Err(protocol_error) if has_incoming_transport_closed(&protocol_error) => {
98                match timeout(SHUTDOWN_GRACE_PERIOD, &mut child_fut).await {
99                    Ok(Err(child_error)) => Err(child_error),
100                    _ => Err(protocol_error),
101                }
102            }
103            Err(error) => Err(error),
104        },
105    }
106}
107
108fn has_incoming_transport_closed(error: &Error) -> bool {
109    fn data_has_reason(data: &serde_json::Value) -> bool {
110        data.get("reason").and_then(serde_json::Value::as_str) == Some(INCOMING_TRANSPORT_CLOSED_REASON)
111            || data.get("data").is_some_and(data_has_reason)
112    }
113
114    is_incoming_transport_closed(error) || error.data.as_ref().is_some_and(data_has_reason)
115}
116
117async fn finish_child_exit(status: ExitStatus, stderr_rx: oneshot::Receiver<String>) -> Result<(), Error> {
118    if status.success() {
119        return Ok(());
120    }
121
122    let stderr = match timeout(SHUTDOWN_GRACE_PERIOD, stderr_rx).await {
123        Ok(Ok(stderr)) => stderr,
124        _ => String::new(),
125    };
126    let message = if stderr.is_empty() {
127        format!("agent process exited ({status})")
128    } else {
129        format!("agent process exited ({status}): {stderr}")
130    };
131
132    Err(internal_error(message))
133}
134
135const SHUTDOWN_GRACE_PERIOD: Duration = Duration::from_secs(1);
acp_utils/client/tokio_agent.rs

acp_utils/client/
tokio_agent.rs
