winx_code_agent/daemon/
control.rs

1use std::io::ErrorKind;
2use std::os::unix::fs::PermissionsExt;
3use std::path::{Path, PathBuf};
4use std::sync::Arc;
5
6use tokio::net::{UnixListener, UnixStream};
7
8use super::lifecycle::{GuardianLifecycle, GuardianLimits};
9use super::protocol::{
10    read_json_frame, write_json_frame, ConfigureSessionParams, HelloResult, JournalReadParams,
11    PruneParams, RpcError, RpcRequest, RpcResponse, RunActionParams, SessionInfo, SessionParams,
12    MAX_FRAME_BYTES, PROTOCOL_MAJOR, PROTOCOL_MINOR,
13};
14use crate::errors::{Result, WinxError};
15use crate::types::normalize_thread_id;
16
17/// Stable control plane. Each logical session is owned by a separate guardian
18/// process, so restarting this process does not close any PTY master.
19pub struct ControlServer {
20    listener: UnixListener,
21    socket_path: PathBuf,
22    lifecycle: Arc<GuardianLifecycle>,
23    epoch: String,
24}
25
26impl ControlServer {
27    pub async fn bind(socket_path: impl AsRef<Path>) -> Result<Self> {
28        let socket_path = socket_path.as_ref().to_path_buf();
29        let parent = socket_path.parent().ok_or_else(|| {
30            WinxError::ConfigurationError("daemon socket must have a parent directory".to_string())
31        })?;
32        tokio::fs::create_dir_all(parent).await?;
33        tokio::fs::set_permissions(parent, std::fs::Permissions::from_mode(0o700)).await?;
34        let guardian_dir = parent.join("guardians");
35        tokio::fs::create_dir_all(&guardian_dir).await?;
36        tokio::fs::set_permissions(&guardian_dir, std::fs::Permissions::from_mode(0o700)).await?;
37
38        if tokio::fs::try_exists(&socket_path).await? {
39            if UnixStream::connect(&socket_path).await.is_ok() {
40                return Err(WinxError::ConfigurationError(format!(
41                    "a winxd instance is already listening at {}",
42                    socket_path.display()
43                )));
44            }
45            tokio::fs::remove_file(&socket_path).await?;
46        }
47        let listener = UnixListener::bind(&socket_path)?;
48        tokio::fs::set_permissions(&socket_path, std::fs::Permissions::from_mode(0o600)).await?;
49        let lifecycle = Arc::new(GuardianLifecycle::new(
50            guardian_dir,
51            guardian_binary()?,
52            GuardianLimits::from_env()?,
53        ));
54
55        Ok(Self {
56            listener,
57            socket_path,
58            lifecycle,
59            epoch: format!("{:016x}", rand::random::<u64>()),
60        })
61    }
62
63    pub async fn serve(self) -> Result<()> {
64        self.lifecycle.clone().spawn_sweeper();
65        loop {
66            let (stream, _) = self.listener.accept().await?;
67            if !same_uid(&stream)? {
68                tracing::warn!("Rejected winxd control connection from a different uid");
69                continue;
70            }
71            let lifecycle = self.lifecycle.clone();
72            let epoch = self.epoch.clone();
73            tokio::spawn(async move {
74                if let Err(error) = serve_connection(stream, lifecycle, epoch).await {
75                    tracing::debug!("winxd control client disconnected: {error}");
76                }
77            });
78        }
79    }
80}
81
82impl Drop for ControlServer {
83    fn drop(&mut self) {
84        let _ = std::fs::remove_file(&self.socket_path);
85    }
86}
87
88async fn serve_connection(
89    mut stream: UnixStream,
90    lifecycle: Arc<GuardianLifecycle>,
91    epoch: String,
92) -> Result<()> {
93    loop {
94        let request: RpcRequest = match read_json_frame(&mut stream).await {
95            Ok(request) => request,
96            Err(error) if error.kind() == ErrorKind::UnexpectedEof => return Ok(()),
97            Err(error) => return Err(error.into()),
98        };
99        let response = dispatch(request, &lifecycle, &epoch).await;
100        write_json_frame(&mut stream, &response).await?;
101    }
102}
103
104#[allow(clippy::too_many_lines)]
105async fn dispatch(request: RpcRequest, lifecycle: &GuardianLifecycle, epoch: &str) -> RpcResponse {
106    if request.jsonrpc != "2.0" {
107        return rpc_error(request.id, -32600, "JSON-RPC version must be 2.0");
108    }
109    if request.method == "winx.hello" {
110        return rpc_result(
111            request.id,
112            &HelloResult {
113                protocol_major: PROTOCOL_MAJOR,
114                protocol_minor: PROTOCOL_MINOR,
115                capabilities: vec![
116                    "per_session_guardians".to_string(),
117                    "planned_control_restart".to_string(),
118                    "guardian_quota".to_string(),
119                    "guardian_idle_ttl".to_string(),
120                    "guardian_activity_clock".to_string(),
121                    "unused_session_ttl".to_string(),
122                    "quota_reclaims_unused".to_string(),
123                    "session.configure".to_string(),
124                    "shell.run_action".to_string(),
125                    "session.list".to_string(),
126                    "session.read_output".to_string(),
127                    "session.kill".to_string(),
128                    "session.prune".to_string(),
129                    "multi_consumer_cursors".to_string(),
130                    "idempotency".to_string(),
131                ],
132                max_frame_bytes: MAX_FRAME_BYTES,
133                daemon_epoch: epoch.to_string(),
134                daemon_pid: std::process::id(),
135            },
136        );
137    }
138
139    if request.method == "session.list" {
140        return match list_sessions(lifecycle.guardian_dir()).await {
141            Ok(sessions) => rpc_result(request.id, &sessions),
142            Err(error) => rpc_error(request.id, -32000, &error.to_string()),
143        };
144    }
145
146    if request.method == "session.prune" {
147        let params = match decode::<PruneParams>(request.params) {
148            Ok(params) => params,
149            Err(error) => return rpc_error(request.id, -32602, &error.to_string()),
150        };
151        return match lifecycle.prune(params.idle_seconds).await {
152            Ok(result) => rpc_result(request.id, &result),
153            Err(error) => rpc_error(request.id, -32003, &error.to_string()),
154        };
155    }
156
157    let thread_id = match request_thread_id(&request) {
158        Ok(thread_id) if !thread_id.is_empty() => thread_id,
159        Ok(_) => return rpc_error(request.id, -32602, "explicit thread_id is required"),
160        Err(error) => return rpc_error(request.id, -32602, &error.to_string()),
161    };
162    let guardian_socket = lifecycle.socket_for(&thread_id);
163    if matches!(request.method.as_str(), "session.configure" | "shell.run_action") {
164        if let Err(error) = lifecycle.ensure_guardian(&thread_id, &guardian_socket).await {
165            return rpc_error(request.id, -32001, &error.to_string());
166        }
167    }
168    if request.method == "session.kill" {
169        match tokio::fs::try_exists(&guardian_socket).await {
170            Ok(true) => {}
171            Ok(false) => return rpc_result(request.id, &false),
172            Err(error) => return rpc_error(request.id, -32002, &error.to_string()),
173        }
174    }
175
176    let kill_after = request.method == "session.kill";
177    let request_id = request.id;
178    let (response, hello) = match relay(&guardian_socket, request).await {
179        Ok(result) => result,
180        Err(error) => return rpc_error(request_id, -32002, &error.to_string()),
181    };
182
183    if kill_after && response.error.is_none() {
184        if let Err(error) = lifecycle.finish_kill(&guardian_socket, hello.daemon_pid).await {
185            return rpc_error(
186                request_id,
187                -32003,
188                &format!("session state was removed, but guardian cleanup failed: {error}"),
189            );
190        }
191    } else {
192        lifecycle.note_activity(&guardian_socket, &thread_id, hello.daemon_pid).await;
193    }
194    response
195}
196
197fn request_thread_id(request: &RpcRequest) -> Result<String> {
198    let thread_id = match request.method.as_str() {
199        "session.configure" => {
200            decode::<ConfigureSessionParams>(request.params.clone())?.snapshot.chat_id
201        }
202        "shell.run_action" => decode::<RunActionParams>(request.params.clone())?.command.thread_id,
203        "session.read_output" => decode::<JournalReadParams>(request.params.clone())?.thread_id,
204        "session.info" | "session.kill" | "session.interrupt" => {
205            decode::<SessionParams>(request.params.clone())?.thread_id
206        }
207        _ => {
208            return Err(WinxError::InvalidInput(format!(
209                "unsupported control method {}",
210                request.method
211            )))
212        }
213    };
214    Ok(normalize_thread_id(&thread_id))
215}
216
217async fn relay(socket: &Path, request: RpcRequest) -> Result<(RpcResponse, HelloResult)> {
218    let mut stream = UnixStream::connect(socket).await.map_err(|error| {
219        WinxError::ShellInitializationError(format!(
220            "guardian at {} is unavailable: {error}",
221            socket.display()
222        ))
223    })?;
224    let hello_id = rand::random::<u64>();
225    write_json_frame(
226        &mut stream,
227        &RpcRequest {
228            jsonrpc: "2.0".to_string(),
229            id: hello_id,
230            method: "winx.hello".to_string(),
231            params: serde_json::json!({ "protocol_major": PROTOCOL_MAJOR }),
232        },
233    )
234    .await?;
235    let hello_response: RpcResponse = read_json_frame(&mut stream).await?;
236    let hello: HelloResult = decode(hello_response.result.ok_or_else(|| {
237        WinxError::ConfigurationError("guardian handshake omitted result".to_string())
238    })?)?;
239    if hello.protocol_major != PROTOCOL_MAJOR {
240        return Err(WinxError::ConfigurationError(format!(
241            "guardian protocol major {} is incompatible with control {}",
242            hello.protocol_major, PROTOCOL_MAJOR
243        )));
244    }
245    write_json_frame(&mut stream, &request).await?;
246    let response = read_json_frame(&mut stream).await?;
247    Ok((response, hello))
248}
249
250async fn list_sessions(guardian_dir: &Path) -> Result<Vec<SessionInfo>> {
251    let mut sessions = Vec::new();
252    let mut entries = tokio::fs::read_dir(guardian_dir).await?;
253    while let Some(entry) = entries.next_entry().await? {
254        let path = entry.path();
255        if path.extension().and_then(|extension| extension.to_str()) != Some("sock") {
256            continue;
257        }
258        let request = RpcRequest {
259            jsonrpc: "2.0".to_string(),
260            id: rand::random::<u64>(),
261            method: "session.list".to_string(),
262            params: serde_json::json!({}),
263        };
264        let Ok((response, _)) = relay(&path, request).await else { continue };
265        let Some(result) = response.result else { continue };
266        if let Ok(mut guardian_sessions) = serde_json::from_value::<Vec<SessionInfo>>(result) {
267            sessions.append(&mut guardian_sessions);
268        }
269    }
270    sessions.sort_by(|left, right| left.thread_id.cmp(&right.thread_id));
271    Ok(sessions)
272}
273
274fn decode<T: serde::de::DeserializeOwned>(value: serde_json::Value) -> Result<T> {
275    serde_json::from_value(value)
276        .map_err(|error| WinxError::DeserializationError(error.to_string()))
277}
278
279fn guardian_binary() -> Result<PathBuf> {
280    if let Some(path) = std::env::var_os("WINX_GUARDIAN_BIN") {
281        return Ok(PathBuf::from(path));
282    }
283    let executable = std::env::current_exe()?;
284    let sibling = executable.with_file_name("winx-guardian");
285    if sibling.is_file() {
286        Ok(sibling)
287    } else {
288        Err(WinxError::ConfigurationError(format!(
289            "winx-guardian not found beside {} (set WINX_GUARDIAN_BIN)",
290            executable.display()
291        )))
292    }
293}
294
295fn rpc_result<T: serde::Serialize>(id: u64, result: &T) -> RpcResponse {
296    match serde_json::to_value(result) {
297        Ok(result) => {
298            RpcResponse { jsonrpc: "2.0".to_string(), id, result: Some(result), error: None }
299        }
300        Err(error) => rpc_error(id, -32603, &error.to_string()),
301    }
302}
303
304fn rpc_error(id: u64, code: i32, message: &str) -> RpcResponse {
305    RpcResponse {
306        jsonrpc: "2.0".to_string(),
307        id,
308        result: None,
309        error: Some(RpcError { code, message: message.to_string() }),
310    }
311}
312
313fn same_uid(stream: &UnixStream) -> Result<bool> {
314    let credentials = stream.peer_cred()?;
315    Ok(credentials.uid() == crate::os::unix::effective_uid())
316}
winx_code_agent/daemon/control.rs

winx_code_agent/daemon/
control.rs
