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, PruneParams,
8    PruneResult, RpcRequest, RpcResponse, RunActionParams, RunActionResult, SessionInfo,
9    SessionParams, WireShellError, PROTOCOL_MAJOR,
10};
11use crate::errors::{Result, WinxError};
12use crate::runtime::{
13    ShellRuntime, ShellRuntimeConfigureFuture, ShellRuntimeFuture, ShellRuntimeUnitFuture,
14    ShellSessionConfiguration, 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 prune_sessions(&self, idle_seconds: Option<u64>) -> Result<PruneResult> {
134        let mut stream = self.connected().await?;
135        DaemonShellRuntime::request(
136            &mut stream,
137            rand::random::<u64>(),
138            "session.prune",
139            serde_json::to_value(PruneParams { idle_seconds })
140                .map_err(|error| WinxError::SerializationError(error.to_string()))?,
141        )
142        .await
143    }
144
145    pub async fn interrupt_session(&self, thread_id: &str) -> Result<()> {
146        let mut stream = self.connected().await?;
147        let _: bool = DaemonShellRuntime::request(
148            &mut stream,
149            rand::random::<u64>(),
150            "session.interrupt",
151            serde_json::to_value(SessionParams { thread_id: thread_id.to_string() })
152                .map_err(|error| WinxError::SerializationError(error.to_string()))?,
153        )
154        .await?;
155        Ok(())
156    }
157}
158
159impl DaemonShellRuntime {
160    pub fn new(socket_path: impl AsRef<Path>) -> Self {
161        Self::with_consumer_id(socket_path, format!("adapter-{:016x}", rand::random::<u64>()))
162    }
163
164    pub fn with_consumer_id(socket_path: impl AsRef<Path>, consumer_id: impl Into<String>) -> Self {
165        Self { socket_path: socket_path.as_ref().to_path_buf(), consumer_id: consumer_id.into() }
166    }
167
168    pub fn socket_path(&self) -> &Path {
169        &self.socket_path
170    }
171
172    async fn request<T: serde::de::DeserializeOwned>(
173        stream: &mut UnixStream,
174        id: u64,
175        method: &str,
176        params: serde_json::Value,
177    ) -> Result<T> {
178        write_json_frame(
179            stream,
180            &RpcRequest { jsonrpc: "2.0".to_string(), id, method: method.to_string(), params },
181        )
182        .await?;
183        let response: RpcResponse = read_json_frame(stream).await?;
184        if response.id != id || response.jsonrpc != "2.0" {
185            return Err(WinxError::ParseError("mismatched JSON-RPC response".to_string()));
186        }
187        if let Some(error) = response.error {
188            return Err(WinxError::CommandExecutionError(format!(
189                "winxd JSON-RPC error {}: {}",
190                error.code, error.message
191            )));
192        }
193        serde_json::from_value(
194            response.result.ok_or_else(|| {
195                WinxError::ParseError("JSON-RPC response omitted result".to_string())
196            })?,
197        )
198        .map_err(|error| WinxError::DeserializationError(error.to_string()))
199    }
200
201    async fn run_remote(
202        &self,
203        bash_state: &std::sync::Arc<tokio::sync::Mutex<Option<BashState>>>,
204        command: BashCommand,
205    ) -> Result<String> {
206        let snapshot =
207            bash_state.lock().await.as_ref().ok_or(WinxError::BashStateNotInitialized)?.snapshot();
208        let mut stream = DaemonClient::new(&self.socket_path).connected().await?;
209
210        let request_key = format!("{:016x}", rand::random::<u64>());
211        let id = rand::random::<u64>();
212        let result: RunActionResult = Self::request(
213            &mut stream,
214            id,
215            "shell.run_action",
216            serde_json::to_value(RunActionParams {
217                snapshot,
218                command,
219                request_key,
220                consumer_id: self.consumer_id.clone(),
221            })
222            .map_err(|error| WinxError::SerializationError(error.to_string()))?,
223        )
224        .await?;
225
226        if let Some(state) = bash_state.lock().await.as_mut() {
227            state.apply_snapshot(&result.snapshot);
228        }
229        match (result.output, result.error) {
230            (Some(output), None) => Ok(output),
231            (_, Some(error)) => Err(from_wire_error(error)),
232            _ => Err(WinxError::ParseError(
233                "winxd action response had neither output nor error".to_string(),
234            )),
235        }
236    }
237
238    async fn configure_remote(
239        &self,
240        bash_state: &mut BashState,
241        transition: ShellSessionTransition,
242    ) -> Result<ShellSessionConfiguration> {
243        let transition = match transition {
244            ShellSessionTransition::FirstCall => ConfigureSessionTransition::FirstCall,
245            ShellSessionTransition::ModeChange => ConfigureSessionTransition::ModeChange,
246            ShellSessionTransition::Reset => ConfigureSessionTransition::Reset,
247            ShellSessionTransition::WorkspaceChange => ConfigureSessionTransition::WorkspaceChange,
248        };
249        let mut stream = DaemonClient::new(&self.socket_path).connected().await?;
250        let result: ConfigureSessionResult = Self::request(
251            &mut stream,
252            rand::random::<u64>(),
253            "session.configure",
254            serde_json::to_value(ConfigureSessionParams {
255                snapshot: bash_state.snapshot(),
256                transition,
257            })
258            .map_err(|error| WinxError::SerializationError(error.to_string()))?,
259        )
260        .await?;
261        if let Some(snapshot) = result.snapshot {
262            bash_state.apply_snapshot(&snapshot);
263        }
264        Ok(ShellSessionConfiguration {
265            attach_hint: result.attach_hint,
266            attached_existing: result.attached_existing,
267        })
268    }
269}
270
271impl ShellRuntime for DaemonShellRuntime {
272    fn configure_session<'a>(
273        &'a self,
274        bash_state: &'a mut BashState,
275        transition: ShellSessionTransition,
276    ) -> ShellRuntimeConfigureFuture<'a> {
277        Box::pin(self.configure_remote(bash_state, transition))
278    }
279
280    fn run_action<'a>(
281        &'a self,
282        bash_state: &'a std::sync::Arc<tokio::sync::Mutex<Option<BashState>>>,
283        command: BashCommand,
284    ) -> ShellRuntimeFuture<'a> {
285        Box::pin(self.run_remote(bash_state, command))
286    }
287
288    fn interrupt<'a>(
289        &'a self,
290        bash_state: &'a std::sync::Arc<tokio::sync::Mutex<Option<BashState>>>,
291    ) -> ShellRuntimeUnitFuture<'a> {
292        Box::pin(async move {
293            let thread_id = bash_state
294                .lock()
295                .await
296                .as_ref()
297                .ok_or(WinxError::BashStateNotInitialized)?
298                .current_thread_id
299                .clone();
300            DaemonClient::new(&self.socket_path).interrupt_session(&thread_id).await
301        })
302    }
303
304    fn terminate_session<'a>(&'a self, thread_id: &'a str) -> ShellRuntimeUnitFuture<'a> {
305        Box::pin(async move {
306            let _ = DaemonClient::new(&self.socket_path).kill_session(thread_id).await?;
307            Ok(())
308        })
309    }
310}
311
312fn from_wire_error(error: WireShellError) -> WinxError {
313    match error {
314        WireShellError::BashStateNotInitialized => WinxError::BashStateNotInitialized,
315        WireShellError::ShellInitialization(message) => {
316            WinxError::ShellInitializationError(message)
317        }
318        WireShellError::CommandExecution(message) | WireShellError::Other(message) => {
319            WinxError::CommandExecutionError(message)
320        }
321        WireShellError::NoActiveCommand(message) => WinxError::NoActiveCommand(message),
322        WireShellError::BackgroundSessionNotFound(message) => {
323            WinxError::BackgroundSessionNotFound(message)
324        }
325        WireShellError::EmptyInteractiveInput(action) => {
326            WinxError::EmptyInteractiveInput { action }
327        }
328        WireShellError::InteractiveTargetNotRunning(message) => {
329            WinxError::InteractiveTargetNotRunning(message)
330        }
331        WireShellError::CommandAlreadyRunning { current_command, duration_seconds } => {
332            WinxError::CommandAlreadyRunning { current_command, duration_seconds }
333        }
334        WireShellError::CommandNotAllowed(message) => WinxError::CommandNotAllowed(message),
335        WireShellError::ThreadIdMismatch(message) => WinxError::ThreadIdMismatch(message),
336        WireShellError::InvalidInput(message) => WinxError::InvalidInput(message),
337    }
338}