Skip to main content

winx_code_agent/daemon/
client.rs

1use std::path::{Path, PathBuf};
2
3use tokio::net::UnixStream;
4
5use super::protocol::{
6    read_json_frame, write_json_frame, ConfigureSessionParams, ConfigureSessionResult,
7    ConfigureSessionTransition, HelloResult, JournalRead, JournalReadParams, RpcRequest,
8    RpcResponse, RunActionParams, RunActionResult, SessionInfo, SessionParams, WireShellError,
9    PROTOCOL_MAJOR,
10};
11use crate::errors::{Result, WinxError};
12use crate::runtime::{
13    ShellRuntime, ShellRuntimeConfigureFuture, ShellRuntimeFuture, ShellRuntimeUnitFuture,
14    ShellSessionTransition,
15};
16use crate::state::bash_state::BashState;
17use crate::types::BashCommand;
18
19/// Shell runtime backed by a `winxd` Unix-domain socket.
20#[derive(Clone, Debug)]
21pub struct DaemonShellRuntime {
22    socket_path: PathBuf,
23    consumer_id: String,
24}
25
26/// Control-plane client used by the CLI and reconnection tests.
27#[derive(Clone, Debug)]
28pub struct DaemonClient {
29    socket_path: PathBuf,
30}
31
32impl DaemonClient {
33    pub fn new(socket_path: impl AsRef<Path>) -> Self {
34        Self { socket_path: socket_path.as_ref().to_path_buf() }
35    }
36
37    async fn connected(&self) -> Result<UnixStream> {
38        let mut stream = UnixStream::connect(&self.socket_path).await.map_err(|error| {
39            WinxError::ShellInitializationError(format!(
40                "cannot connect to winxd at {}: {error}",
41                self.socket_path.display()
42            ))
43        })?;
44        let hello: HelloResult = DaemonShellRuntime::request(
45            &mut stream,
46            rand::random::<u64>(),
47            "winx.hello",
48            serde_json::json!({ "protocol_major": PROTOCOL_MAJOR }),
49        )
50        .await?;
51        if hello.protocol_major != PROTOCOL_MAJOR {
52            return Err(WinxError::ConfigurationError(format!(
53                "winxd protocol major {} is incompatible with client {}",
54                hello.protocol_major, PROTOCOL_MAJOR
55            )));
56        }
57        Ok(stream)
58    }
59
60    pub async fn hello(&self) -> Result<HelloResult> {
61        let mut stream = UnixStream::connect(&self.socket_path).await.map_err(|error| {
62            WinxError::ShellInitializationError(format!(
63                "cannot connect to winxd at {}: {error}",
64                self.socket_path.display()
65            ))
66        })?;
67        let hello: HelloResult = DaemonShellRuntime::request(
68            &mut stream,
69            rand::random::<u64>(),
70            "winx.hello",
71            serde_json::json!({ "protocol_major": PROTOCOL_MAJOR }),
72        )
73        .await?;
74        if hello.protocol_major != PROTOCOL_MAJOR {
75            return Err(WinxError::ConfigurationError(format!(
76                "winxd protocol major {} is incompatible with client {}",
77                hello.protocol_major, PROTOCOL_MAJOR
78            )));
79        }
80        Ok(hello)
81    }
82
83    pub async fn list_sessions(&self) -> Result<Vec<SessionInfo>> {
84        let mut stream = self.connected().await?;
85        DaemonShellRuntime::request(
86            &mut stream,
87            rand::random::<u64>(),
88            "session.list",
89            serde_json::json!({}),
90        )
91        .await
92    }
93
94    pub async fn session_info(&self, thread_id: &str) -> Result<SessionInfo> {
95        let mut stream = self.connected().await?;
96        DaemonShellRuntime::request(
97            &mut stream,
98            rand::random::<u64>(),
99            "session.info",
100            serde_json::to_value(SessionParams { thread_id: thread_id.to_string() })
101                .map_err(|error| WinxError::SerializationError(error.to_string()))?,
102        )
103        .await
104    }
105
106    pub async fn read_output(&self, thread_id: &str, consumer_id: &str) -> Result<JournalRead> {
107        let mut stream = self.connected().await?;
108        DaemonShellRuntime::request(
109            &mut stream,
110            rand::random::<u64>(),
111            "session.read_output",
112            serde_json::to_value(JournalReadParams {
113                thread_id: thread_id.to_string(),
114                consumer_id: consumer_id.to_string(),
115            })
116            .map_err(|error| WinxError::SerializationError(error.to_string()))?,
117        )
118        .await
119    }
120
121    pub async fn kill_session(&self, thread_id: &str) -> Result<bool> {
122        let mut stream = self.connected().await?;
123        DaemonShellRuntime::request(
124            &mut stream,
125            rand::random::<u64>(),
126            "session.kill",
127            serde_json::to_value(SessionParams { thread_id: thread_id.to_string() })
128                .map_err(|error| WinxError::SerializationError(error.to_string()))?,
129        )
130        .await
131    }
132
133    pub async fn interrupt_session(&self, thread_id: &str) -> Result<()> {
134        let mut stream = self.connected().await?;
135        let _: bool = DaemonShellRuntime::request(
136            &mut stream,
137            rand::random::<u64>(),
138            "session.interrupt",
139            serde_json::to_value(SessionParams { thread_id: thread_id.to_string() })
140                .map_err(|error| WinxError::SerializationError(error.to_string()))?,
141        )
142        .await?;
143        Ok(())
144    }
145}
146
147impl DaemonShellRuntime {
148    pub fn new(socket_path: impl AsRef<Path>) -> Self {
149        Self::with_consumer_id(socket_path, format!("adapter-{:016x}", rand::random::<u64>()))
150    }
151
152    pub fn with_consumer_id(socket_path: impl AsRef<Path>, consumer_id: impl Into<String>) -> Self {
153        Self { socket_path: socket_path.as_ref().to_path_buf(), consumer_id: consumer_id.into() }
154    }
155
156    pub fn socket_path(&self) -> &Path {
157        &self.socket_path
158    }
159
160    async fn request<T: serde::de::DeserializeOwned>(
161        stream: &mut UnixStream,
162        id: u64,
163        method: &str,
164        params: serde_json::Value,
165    ) -> Result<T> {
166        write_json_frame(
167            stream,
168            &RpcRequest { jsonrpc: "2.0".to_string(), id, method: method.to_string(), params },
169        )
170        .await?;
171        let response: RpcResponse = read_json_frame(stream).await?;
172        if response.id != id || response.jsonrpc != "2.0" {
173            return Err(WinxError::ParseError("mismatched JSON-RPC response".to_string()));
174        }
175        if let Some(error) = response.error {
176            return Err(WinxError::CommandExecutionError(format!(
177                "winxd JSON-RPC error {}: {}",
178                error.code, error.message
179            )));
180        }
181        serde_json::from_value(
182            response.result.ok_or_else(|| {
183                WinxError::ParseError("JSON-RPC response omitted result".to_string())
184            })?,
185        )
186        .map_err(|error| WinxError::DeserializationError(error.to_string()))
187    }
188
189    async fn run_remote(
190        &self,
191        bash_state: &std::sync::Arc<tokio::sync::Mutex<Option<BashState>>>,
192        command: BashCommand,
193    ) -> Result<String> {
194        let snapshot =
195            bash_state.lock().await.as_ref().ok_or(WinxError::BashStateNotInitialized)?.snapshot();
196        let mut stream = DaemonClient::new(&self.socket_path).connected().await?;
197
198        let request_key = format!("{:016x}", rand::random::<u64>());
199        let id = rand::random::<u64>();
200        let result: RunActionResult = Self::request(
201            &mut stream,
202            id,
203            "shell.run_action",
204            serde_json::to_value(RunActionParams {
205                snapshot,
206                command,
207                request_key,
208                consumer_id: self.consumer_id.clone(),
209            })
210            .map_err(|error| WinxError::SerializationError(error.to_string()))?,
211        )
212        .await?;
213
214        if let Some(state) = bash_state.lock().await.as_mut() {
215            state.apply_snapshot(&result.snapshot);
216        }
217        match (result.output, result.error) {
218            (Some(output), None) => Ok(output),
219            (_, Some(error)) => Err(from_wire_error(error)),
220            _ => Err(WinxError::ParseError(
221                "winxd action response had neither output nor error".to_string(),
222            )),
223        }
224    }
225
226    async fn configure_remote(
227        &self,
228        bash_state: &mut BashState,
229        transition: ShellSessionTransition,
230    ) -> Result<Option<String>> {
231        let transition = match transition {
232            ShellSessionTransition::FirstCall => ConfigureSessionTransition::FirstCall,
233            ShellSessionTransition::ModeChange => ConfigureSessionTransition::ModeChange,
234            ShellSessionTransition::Reset => ConfigureSessionTransition::Reset,
235            ShellSessionTransition::WorkspaceChange => ConfigureSessionTransition::WorkspaceChange,
236        };
237        let mut stream = DaemonClient::new(&self.socket_path).connected().await?;
238        let result: ConfigureSessionResult = Self::request(
239            &mut stream,
240            rand::random::<u64>(),
241            "session.configure",
242            serde_json::to_value(ConfigureSessionParams {
243                snapshot: bash_state.snapshot(),
244                transition,
245            })
246            .map_err(|error| WinxError::SerializationError(error.to_string()))?,
247        )
248        .await?;
249        Ok(result.attach_hint)
250    }
251}
252
253impl ShellRuntime for DaemonShellRuntime {
254    fn configure_session<'a>(
255        &'a self,
256        bash_state: &'a mut BashState,
257        transition: ShellSessionTransition,
258    ) -> ShellRuntimeConfigureFuture<'a> {
259        Box::pin(self.configure_remote(bash_state, transition))
260    }
261
262    fn run_action<'a>(
263        &'a self,
264        bash_state: &'a std::sync::Arc<tokio::sync::Mutex<Option<BashState>>>,
265        command: BashCommand,
266    ) -> ShellRuntimeFuture<'a> {
267        Box::pin(self.run_remote(bash_state, command))
268    }
269
270    fn interrupt<'a>(
271        &'a self,
272        bash_state: &'a std::sync::Arc<tokio::sync::Mutex<Option<BashState>>>,
273    ) -> ShellRuntimeUnitFuture<'a> {
274        Box::pin(async move {
275            let thread_id = bash_state
276                .lock()
277                .await
278                .as_ref()
279                .ok_or(WinxError::BashStateNotInitialized)?
280                .current_thread_id
281                .clone();
282            DaemonClient::new(&self.socket_path).interrupt_session(&thread_id).await
283        })
284    }
285}
286
287fn from_wire_error(error: WireShellError) -> WinxError {
288    match error {
289        WireShellError::BashStateNotInitialized => WinxError::BashStateNotInitialized,
290        WireShellError::ShellInitialization(message) => {
291            WinxError::ShellInitializationError(message)
292        }
293        WireShellError::CommandExecution(message) => WinxError::CommandExecutionError(message),
294        WireShellError::CommandNotAllowed(message) => WinxError::CommandNotAllowed(message),
295        WireShellError::ThreadIdMismatch(message) => WinxError::ThreadIdMismatch(message),
296        WireShellError::Other(message) => WinxError::CommandExecutionError(message),
297    }
298}