aether_cli/acp/
server.rs

1use super::state::{AcpState, ClientGuard};
2use super::{AcpArgs, AcpRunError, create_acp_state};
3use crate::output::OutputFormat;
4use crate::prompt::prompt_or_stdin;
5use acp_utils::websocket::WebSocketTransport;
6use std::fs::canonicalize;
7use std::future::Future;
8use std::io;
9use std::net::SocketAddr;
10use std::path::PathBuf;
11use std::sync::Arc;
12use thiserror::Error;
13use tokio::net::{TcpListener, TcpStream};
14#[cfg(unix)]
15use tokio::signal::unix::{SignalKind, signal};
16use tokio::task::JoinSet;
17use tokio_tungstenite::tungstenite::{
18    handshake::server::{ErrorResponse, Request},
19    http::StatusCode,
20};
21use tokio_util::sync::CancellationToken;
22use tracing::{info, warn};
23
24#[derive(clap::Args, Debug)]
25pub struct ServerArgs {
26    /// Address to listen on. The raw server is unauthenticated; use a private network or authenticating proxy.
27    #[arg(long, default_value = "127.0.0.1:8765")]
28    pub listen: SocketAddr,
29
30    /// Server workspace used to resolve settings and as the default session directory.
31    #[arg(short = 'C', long, default_value = ".")]
32    pub cwd: PathBuf,
33
34    /// Initial prompt to run.
35    #[arg(long)]
36    pub prompt: Option<String>,
37
38    #[command(flatten)]
39    pub detached: DetachedArgs,
40
41    #[command(flatten)]
42    pub acp: AcpArgs,
43}
44
45#[derive(clap::Args, Clone, Debug, Default)]
46pub struct DetachedArgs {
47    /// Output format for when no clients are attached.
48    #[arg(long)]
49    pub output: Option<OutputFormat>,
50
51    /// How many seconds to wait after the agent becomes idle before running the --on-idle command.
52    #[arg(long, requires = "on_idle")]
53    pub idle_after: Option<u64>,
54
55    /// Shell command to spawn when the agent becomes idle.
56    #[arg(long, requires = "idle_after")]
57    pub on_idle: Option<String>,
58}
59
60#[derive(Debug, Error)]
61pub enum ServerRunError {
62    #[error("Invalid server workspace {path}: {source}")]
63    Workspace { path: PathBuf, source: io::Error },
64    #[error(transparent)]
65    Initialization(#[from] AcpRunError),
66    #[error("Failed to bind ACP listener at {address}: {source}")]
67    Bind { address: SocketAddr, source: io::Error },
68    #[error("ACP listener failed: {0}")]
69    Accept(#[source] io::Error),
70    #[error("Failed to listen for shutdown signals: {0}")]
71    Signal(#[source] io::Error),
72    #[error("Failed to read the initial prompt from stdin: {0}")]
73    PromptStdin(#[source] io::Error),
74    #[error("Failed to start the initial session: {0}")]
75    InitialSession(#[source] agent_client_protocol::Error),
76}
77
78pub async fn run_server(args: ServerArgs) -> Result<(), ServerRunError> {
79    let cwd = canonicalize(&args.cwd)
80        .and_then(|cwd| {
81            if cwd.is_dir() {
82                Ok(cwd)
83            } else {
84                Err(io::Error::new(io::ErrorKind::NotADirectory, "workspace must be a directory"))
85            }
86        })
87        .map_err(|source| ServerRunError::Workspace { path: args.cwd, source })?;
88
89    let prompt = prompt_or_stdin(args.prompt).map_err(ServerRunError::PromptStdin)?;
90    let state = Arc::new(create_acp_state(args.acp, &cwd, args.detached)?);
91    let server = AcpServer::bind(args.listen, state.clone()).await?;
92    info!(address = %args.listen, cwd = %cwd.display(), "Starting Aether ACP WebSocket server");
93
94    let result = async {
95        if let Some(prompt) = prompt {
96            let session_id = state.start_session(prompt).await.map_err(ServerRunError::InitialSession)?;
97            info!(session_id = %session_id.0, "Started initial server session");
98        }
99        server.run_until(async { shutdown_signal().await.map_err(ServerRunError::Signal) }).await
100    }
101    .await;
102    state.shutdown_all().await;
103    result
104}
105
106/// Owns networking, not session lifetime. Drop closes the listener and aborts
107/// connection tasks; explicit shutdown also waits for output detachment.
108pub(crate) struct AcpServer {
109    listener: TcpListener,
110    state: Arc<AcpState>,
111    connections: JoinSet<()>,
112    stop: CancellationToken,
113}
114
115impl AcpServer {
116    pub(crate) async fn bind(address: SocketAddr, state: Arc<AcpState>) -> Result<Self, ServerRunError> {
117        let listener = TcpListener::bind(address).await.map_err(|source| ServerRunError::Bind { address, source })?;
118        Ok(Self { listener, stop: state.stop_token().child_token(), state, connections: JoinSet::new() })
119    }
120
121    #[cfg(any(test, feature = "testing"))]
122    pub(crate) fn local_addr(&self) -> io::Result<SocketAddr> {
123        self.listener.local_addr()
124    }
125
126    pub(crate) async fn run_until(
127        mut self,
128        stop: impl Future<Output = Result<(), ServerRunError>>,
129    ) -> Result<(), ServerRunError> {
130        let result = tokio::select! {
131            result = self.run() => result,
132            result = stop => result,
133        };
134        self.shutdown().await;
135        result
136    }
137
138    async fn run(&mut self) -> Result<(), ServerRunError> {
139        loop {
140            tokio::select! {
141                biased;
142                () = self.stop.cancelled() => return Ok(()),
143                result = self.connections.join_next(), if !self.connections.is_empty() => {
144                    if let Some(Err(error)) = result {
145                        warn!(%error, "ACP connection task failed");
146                    }
147                }
148                accepted = self.listener.accept() => {
149                    let (socket, peer) = accepted.map_err(ServerRunError::Accept)?;
150                    self.accept(socket, peer);
151                }
152            }
153        }
154    }
155
156    /// Close admission and join all networking, leaving the host running.
157    pub(crate) async fn shutdown(self) {
158        self.stop.cancel();
159        drop(self.listener);
160        let mut connections = self.connections;
161        while let Some(result) = connections.join_next().await {
162            if let Err(error) = result
163                && !error.is_cancelled()
164            {
165                warn!(%error, "ACP connection task failed during shutdown");
166            }
167        }
168    }
169
170    fn accept(&mut self, socket: TcpStream, peer: SocketAddr) {
171        let stop = self.stop.clone();
172        let guard = self.state.try_attach();
173        self.connections.spawn(serve_connection(socket, peer, guard, stop));
174    }
175}
176
177#[expect(clippy::result_large_err, reason = "tungstenite's handshake callback requires an HTTP error response")]
178async fn serve_connection(socket: TcpStream, peer: SocketAddr, guard: Option<ClientGuard>, stop: CancellationToken) {
179    let handshake = tokio::select! {
180        biased;
181        () = stop.cancelled() => {
182            if let Some(guard) = guard {
183                guard.release().await;
184            }
185            return;
186        },
187        result = tokio_tungstenite::accept_hdr_async(socket, |_request: &Request, response| {
188            if guard.is_some() {
189                Ok(response)
190            } else {
191                let mut error = ErrorResponse::new(Some("client already attached".to_string()));
192                *error.status_mut() = StatusCode::CONFLICT;
193                Err(error)
194            }
195        }) => result,
196    };
197    let socket = match handshake {
198        Ok(socket) => socket,
199        Err(error) => {
200            warn!(%peer, %error, "ACP WebSocket handshake failed");
201            if let Some(guard) = guard {
202                guard.release().await;
203            }
204            return;
205        }
206    };
207    if let Err(error) =
208        guard.expect("successful handshake owns the client").serve(WebSocketTransport::new(socket), stop).await
209    {
210        warn!(%peer, %error, "ACP connection failed");
211    }
212}
213
214async fn shutdown_signal() -> io::Result<()> {
215    #[cfg(unix)]
216    {
217        let mut terminate = signal(SignalKind::terminate())?;
218        tokio::select! {
219            result = tokio::signal::ctrl_c() => result,
220            _ = terminate.recv() => Ok(()),
221        }
222    }
223    #[cfg(not(unix))]
224    tokio::signal::ctrl_c().await
225}
aether_cli/acp/server.rs

aether_cli/acp/
server.rs
