Skip to main content

aether_cli/
client.rs

1use acp_utils::websocket::WebSocketTransport;
2use agent_client_protocol::schema::v2::SessionId;
3use std::path::PathBuf;
4use thiserror::Error;
5use tokio_tungstenite::connect_async;
6use tokio_tungstenite::tungstenite;
7use tokio_tungstenite::tungstenite::client::IntoClientRequest;
8use tokio_tungstenite::tungstenite::http::{HeaderName, HeaderValue, Request, StatusCode, Uri};
9use wisp::run_remote_tui;
10use wisp::settings::load_or_create_settings;
11
12#[derive(clap::Args)]
13pub struct ClientArgs {
14    /// Remote Aether WebSocket URL (use wss:// for an authenticating TLS proxy).
15    #[arg(default_value = "ws://127.0.0.1:8765")]
16    pub url: String,
17
18    /// Saved session to resume instead of the server's current session.
19    #[arg(long, value_name = "ID")]
20    pub session: Option<String>,
21
22    /// HTTP header for a gateway/proxy, in 'Name: value' format. Repeatable.
23    #[arg(short = 'H', long = "header", value_name = "HEADER")]
24    pub headers: Vec<String>,
25
26    /// Directory for client UI logs (default: /tmp/wisp-logs).
27    #[arg(long)]
28    pub log_dir: Option<PathBuf>,
29}
30
31#[derive(Debug, Error)]
32pub enum ClientRunError {
33    #[error("Invalid remote URL; expected ws:// or wss:// with a host")]
34    InvalidUrl,
35    #[error("Invalid HTTP header #{index}; expected 'Name: value' with a valid HTTP name and value")]
36    InvalidHeader { index: usize },
37    #[error("The server already has a client attached; disconnect it before connecting again")]
38    ServerOccupied,
39    #[error("Failed to establish the remote WebSocket connection; check the address, network, and proxy configuration")]
40    Handshake,
41    #[error(transparent)]
42    Tui(#[from] wisp::error::AppError),
43}
44
45impl ClientArgs {
46    /// Validate all connection inputs before performing any network I/O.
47    pub fn connection_request(&self) -> Result<Request<()>, ClientRunError> {
48        let uri: Uri = self.url.parse().map_err(|_| ClientRunError::InvalidUrl)?;
49        if !matches!(uri.scheme_str(), Some("ws" | "wss")) || uri.host().is_none_or(str::is_empty) {
50            return Err(ClientRunError::InvalidUrl);
51        }
52        let mut request = uri.into_client_request().map_err(|_| ClientRunError::InvalidUrl)?;
53        for (index, header) in self.headers.iter().enumerate() {
54            let invalid = || ClientRunError::InvalidHeader { index: index + 1 };
55            let (name, value) = header.split_once(':').ok_or_else(invalid)?;
56            let name = HeaderName::from_bytes(name.as_bytes()).map_err(|_| invalid())?;
57            let mut value = HeaderValue::from_str(value.trim_matches([' ', '\t'])).map_err(|_| invalid())?;
58            value.set_sensitive(true);
59            request.headers_mut().append(name, value);
60        }
61        Ok(request)
62    }
63}
64
65pub async fn run_client(args: ClientArgs) -> Result<(), ClientRunError> {
66    let request = args.connection_request()?;
67    let (socket, _) = connect_async(request).await.map_err(|error| match error {
68        tungstenite::Error::Http(response) if response.status() == StatusCode::CONFLICT => {
69            ClientRunError::ServerOccupied
70        }
71        _ => ClientRunError::Handshake,
72    })?;
73
74    let settings = load_or_create_settings();
75    run_remote_tui(
76        WebSocketTransport::new(socket),
77        args.session.map(SessionId::new),
78        settings,
79        args.log_dir.as_deref(),
80    )
81    .await?;
82
83    Ok(())
84}