winx_code_agent/daemon/
server.rs

1use std::collections::{HashMap, HashSet, VecDeque};
2use std::io::ErrorKind;
3use std::os::unix::fs::PermissionsExt;
4use std::path::{Path, PathBuf};
5use std::sync::atomic::{AtomicBool, Ordering};
6use std::sync::Arc;
7use std::time::Duration;
8
9use tokio::net::{UnixListener, UnixStream};
10use tokio::sync::Mutex;
11
12use super::protocol::{
13    read_json_frame, write_json_frame, ConfigureSessionParams, ConfigureSessionResult,
14    ConfigureSessionTransition, HelloResult, JournalRead, JournalReadParams, RpcError, RpcRequest,
15    RpcResponse, RunActionParams, RunActionResult, SessionInfo, SessionParams, WireShellError,
16    MAX_FRAME_BYTES, PROTOCOL_MAJOR, PROTOCOL_MINOR,
17};
18use crate::errors::{Result, WinxError};
19use crate::runtime::{lock_session_store, ShellTarget};
20use crate::state::bash_state::BashState;
21use crate::state::pty::SharedPtyShell;
22use crate::tools::bash_command::ShellDeliveryCursor;
23use crate::types::normalize_thread_id;
24
25type SharedState = Arc<Mutex<Option<BashState>>>;
26
27struct DaemonSession {
28    state: SharedState,
29    completed: Mutex<HashMap<String, RunActionResult>>,
30    journal: Mutex<OutputJournal>,
31    observed: Mutex<HashMap<String, (u64, String)>>,
32    background_ids: Mutex<HashSet<String>>,
33    command_id: Mutex<Option<String>>,
34    delivery_cursors: Mutex<HashMap<String, Arc<Mutex<ShellDeliveryCursor>>>>,
35    drainer_started: AtomicBool,
36}
37
38impl DaemonSession {
39    fn new() -> Self {
40        Self {
41            state: Arc::new(Mutex::new(None)),
42            completed: Mutex::new(HashMap::new()),
43            journal: Mutex::new(OutputJournal::default()),
44            observed: Mutex::new(HashMap::new()),
45            background_ids: Mutex::new(HashSet::new()),
46            command_id: Mutex::new(None),
47            delivery_cursors: Mutex::new(HashMap::new()),
48            drainer_started: AtomicBool::new(false),
49        }
50    }
51}
52
53const MAX_JOURNAL_BYTES: usize = 4 * 1024 * 1024;
54
55#[derive(Default)]
56struct OutputJournal {
57    chunks: VecDeque<(u64, String)>,
58    cursors: HashMap<String, u64>,
59    next_seq: u64,
60    bytes: usize,
61}
62
63impl OutputJournal {
64    fn append(&mut self, output: String) {
65        if output.is_empty() {
66            return;
67        }
68        self.next_seq = self.next_seq.saturating_add(1);
69        self.bytes = self.bytes.saturating_add(output.len());
70        self.chunks.push_back((self.next_seq, output));
71        while self.bytes > MAX_JOURNAL_BYTES && self.chunks.len() > 1 {
72            if let Some((_, removed)) = self.chunks.pop_front() {
73                self.bytes = self.bytes.saturating_sub(removed.len());
74            }
75        }
76    }
77
78    fn read(&mut self, consumer_id: &str) -> JournalRead {
79        let oldest = self.chunks.front().map_or(self.next_seq.saturating_add(1), |(seq, _)| *seq);
80        let cursor = self.cursors.get(consumer_id).copied().unwrap_or(0);
81        let gap = cursor.saturating_add(1) < oldest;
82        let mut output = String::new();
83        let mut next_seq = cursor;
84        for (seq, chunk) in self.chunks.iter().filter(|(seq, _)| *seq > cursor) {
85            output.push_str(chunk);
86            next_seq = *seq;
87        }
88        self.cursors.insert(consumer_id.to_string(), next_seq);
89        JournalRead { output, next_seq, gap }
90    }
91}
92
93/// Long-lived JSON-RPC shell owner listening on a Unix-domain socket.
94pub struct DaemonServer {
95    listener: UnixListener,
96    socket_path: PathBuf,
97    sessions: Arc<Mutex<HashMap<String, Arc<DaemonSession>>>>,
98    epoch: String,
99}
100
101impl DaemonServer {
102    pub async fn bind(socket_path: impl AsRef<Path>) -> Result<Self> {
103        let socket_path = socket_path.as_ref().to_path_buf();
104        let parent = socket_path.parent().ok_or_else(|| {
105            WinxError::ConfigurationError("daemon socket must have a parent directory".to_string())
106        })?;
107        tokio::fs::create_dir_all(parent).await?;
108        tokio::fs::set_permissions(parent, std::fs::Permissions::from_mode(0o700)).await?;
109
110        if tokio::fs::try_exists(&socket_path).await? {
111            if UnixStream::connect(&socket_path).await.is_ok() {
112                return Err(WinxError::ConfigurationError(format!(
113                    "a winxd instance is already listening at {}",
114                    socket_path.display()
115                )));
116            }
117            tokio::fs::remove_file(&socket_path).await?;
118        }
119
120        let listener = UnixListener::bind(&socket_path)?;
121        tokio::fs::set_permissions(&socket_path, std::fs::Permissions::from_mode(0o600)).await?;
122        let epoch = format!("{:016x}", rand::random::<u64>());
123        Ok(Self { listener, socket_path, sessions: Arc::new(Mutex::new(HashMap::new())), epoch })
124    }
125
126    pub async fn serve(self) -> Result<()> {
127        loop {
128            let (stream, _) = self.listener.accept().await?;
129            if !same_uid(&stream)? {
130                tracing::warn!("Rejected winxd connection from a different uid");
131                continue;
132            }
133            let sessions = self.sessions.clone();
134            let epoch = self.epoch.clone();
135            tokio::spawn(async move {
136                if let Err(error) = serve_connection(stream, sessions, epoch).await {
137                    tracing::debug!("winxd client disconnected: {error}");
138                }
139            });
140        }
141    }
142}
143
144impl Drop for DaemonServer {
145    fn drop(&mut self) {
146        let _ = std::fs::remove_file(&self.socket_path);
147    }
148}
149
150async fn serve_connection(
151    mut stream: UnixStream,
152    sessions: Arc<Mutex<HashMap<String, Arc<DaemonSession>>>>,
153    epoch: String,
154) -> Result<()> {
155    loop {
156        let request: RpcRequest = match read_json_frame(&mut stream).await {
157            Ok(request) => request,
158            Err(error) if error.kind() == ErrorKind::UnexpectedEof => return Ok(()),
159            Err(error) => return Err(error.into()),
160        };
161        let response = dispatch(request, &sessions, &epoch).await;
162        write_json_frame(&mut stream, &response).await?;
163    }
164}
165
166async fn dispatch(
167    request: RpcRequest,
168    sessions: &Arc<Mutex<HashMap<String, Arc<DaemonSession>>>>,
169    epoch: &str,
170) -> RpcResponse {
171    if request.jsonrpc != "2.0" {
172        return rpc_error(request.id, -32600, "JSON-RPC version must be 2.0");
173    }
174
175    match request.method.as_str() {
176        "winx.hello" => rpc_result(
177            request.id,
178            &HelloResult {
179                protocol_major: PROTOCOL_MAJOR,
180                protocol_minor: PROTOCOL_MINOR,
181                capabilities: vec![
182                    "session.configure".to_string(),
183                    "shell.run_action".to_string(),
184                    "session.list".to_string(),
185                    "session.read_output".to_string(),
186                    "session.kill".to_string(),
187                    "multi_consumer_cursors".to_string(),
188                    "idempotency".to_string(),
189                ],
190                max_frame_bytes: MAX_FRAME_BYTES,
191                daemon_epoch: epoch.to_string(),
192                daemon_pid: std::process::id(),
193            },
194        ),
195        "session.configure" => {
196            let params: ConfigureSessionParams = match serde_json::from_value(request.params) {
197                Ok(params) => params,
198                Err(error) => return rpc_error(request.id, -32602, &error.to_string()),
199            };
200            match configure_session(sessions, params).await {
201                Ok(result) => rpc_result(request.id, &result),
202                Err(error) => rpc_error(request.id, -32603, &error.to_string()),
203            }
204        }
205        "shell.run_action" => {
206            let params: RunActionParams = match serde_json::from_value(request.params) {
207                Ok(params) => params,
208                Err(error) => return rpc_error(request.id, -32602, &error.to_string()),
209            };
210            match run_action(sessions, params).await {
211                Ok(result) => rpc_result(request.id, &result),
212                Err(error) => rpc_error(request.id, -32603, &error.to_string()),
213            }
214        }
215        "session.list" => {
216            let entries = sessions.lock().await.values().cloned().collect::<Vec<_>>();
217            let mut infos = Vec::with_capacity(entries.len());
218            for session in entries {
219                if let Some(info) = session_info(&session).await {
220                    infos.push(info);
221                }
222            }
223            infos.sort_by(|left, right| left.thread_id.cmp(&right.thread_id));
224            rpc_result(request.id, &infos)
225        }
226        "session.info" => {
227            let params: SessionParams = match serde_json::from_value(request.params) {
228                Ok(params) => params,
229                Err(error) => return rpc_error(request.id, -32602, &error.to_string()),
230            };
231            let session =
232                sessions.lock().await.get(&normalize_thread_id(&params.thread_id)).cloned();
233            match session {
234                Some(session) => match session_info(&session).await {
235                    Some(info) => rpc_result(request.id, &info),
236                    None => rpc_error(request.id, -32004, "session not initialized"),
237                },
238                None => rpc_error(request.id, -32004, "session not found"),
239            }
240        }
241        "session.read_output" => {
242            let params: JournalReadParams = match serde_json::from_value(request.params) {
243                Ok(params) => params,
244                Err(error) => return rpc_error(request.id, -32602, &error.to_string()),
245            };
246            let thread_id = normalize_thread_id(&params.thread_id);
247            let session = sessions.lock().await.get(&thread_id).cloned();
248            match session {
249                Some(session) => {
250                    capture_session_outputs(&session).await;
251                    let read = session.journal.lock().await.read(&params.consumer_id);
252                    rpc_result(request.id, &read)
253                }
254                None => rpc_error(request.id, -32004, "session not found"),
255            }
256        }
257        "session.kill" => {
258            let params: SessionParams = match serde_json::from_value(request.params) {
259                Ok(params) => params,
260                Err(error) => return rpc_error(request.id, -32602, &error.to_string()),
261            };
262            let session = sessions.lock().await.remove(&normalize_thread_id(&params.thread_id));
263            match session {
264                Some(session) => {
265                    let background_ids = session.background_ids.lock().await.clone();
266                    {
267                        let mut store = lock_session_store();
268                        for id in background_ids {
269                            store.remove_shell(&id);
270                        }
271                    }
272                    *session.state.lock().await = None;
273                    rpc_result(request.id, &true)
274                }
275                None => rpc_result(request.id, &false),
276            }
277        }
278        "session.interrupt" => {
279            let params: SessionParams = match serde_json::from_value(request.params) {
280                Ok(params) => params,
281                Err(error) => return rpc_error(request.id, -32602, &error.to_string()),
282            };
283            let session =
284                sessions.lock().await.get(&normalize_thread_id(&params.thread_id)).cloned();
285            match session {
286                Some(session) => match interrupt_session(&session).await {
287                    Ok(()) => rpc_result(request.id, &true),
288                    Err(error) => rpc_error(request.id, -32005, &error.to_string()),
289                },
290                None => rpc_error(request.id, -32004, "session not found"),
291            }
292        }
293        _ => rpc_error(request.id, -32601, "method not found"),
294    }
295}
296
297async fn configure_session(
298    sessions: &Arc<Mutex<HashMap<String, Arc<DaemonSession>>>>,
299    params: ConfigureSessionParams,
300) -> Result<ConfigureSessionResult> {
301    let thread_id = normalize_thread_id(&params.snapshot.chat_id);
302    if thread_id.is_empty() {
303        return Err(WinxError::ThreadIdMismatch(
304            "daemon session configuration requires an explicit thread_id".to_string(),
305        ));
306    }
307
308    let session = {
309        let mut sessions = sessions.lock().await;
310        match params.transition {
311            ConfigureSessionTransition::FirstCall => {
312                sessions.entry(thread_id).or_insert_with(|| Arc::new(DaemonSession::new())).clone()
313            }
314            _ => sessions.get(&thread_id).cloned().ok_or(WinxError::BashStateNotInitialized)?,
315        }
316    };
317    ensure_drainer(&session);
318
319    let attach_hint = {
320        let mut guard = session.state.lock().await;
321        match params.transition {
322            ConfigureSessionTransition::FirstCall => {
323                let mut state = BashState::new();
324                state.apply_snapshot(&params.snapshot);
325                if state.cwd.exists() {
326                    state.init_pty_shell().await?;
327                }
328                *guard = Some(state);
329            }
330            ConfigureSessionTransition::ModeChange => {
331                let state = guard.as_mut().ok_or(WinxError::BashStateNotInitialized)?;
332                let daemon_cwd = state.cwd.clone();
333                state.apply_snapshot(&params.snapshot);
334                state.cwd = daemon_cwd;
335            }
336            ConfigureSessionTransition::Reset => {
337                let state = guard.as_mut().ok_or(WinxError::BashStateNotInitialized)?;
338                let daemon_cwd = state.cwd.clone();
339                state.apply_snapshot(&params.snapshot);
340                state.cwd = daemon_cwd;
341                state.init_pty_shell().await?;
342            }
343            ConfigureSessionTransition::WorkspaceChange => {
344                let state = guard.as_mut().ok_or(WinxError::BashStateNotInitialized)?;
345                state.apply_snapshot(&params.snapshot);
346            }
347        }
348        let state = guard.as_ref().ok_or(WinxError::BashStateNotInitialized)?;
349        let shell = state.pty_shell.lock().await;
350        shell.as_ref().and_then(|shell| shell.attach_hint.clone())
351    };
352
353    if matches!(
354        params.transition,
355        ConfigureSessionTransition::FirstCall | ConfigureSessionTransition::Reset
356    ) {
357        *session.command_id.lock().await = None;
358        session.completed.lock().await.clear();
359        session.delivery_cursors.lock().await.clear();
360        session.observed.lock().await.remove("main");
361    }
362    capture_session_outputs(&session).await;
363    Ok(ConfigureSessionResult { attach_hint })
364}
365
366async fn run_action(
367    sessions: &Arc<Mutex<HashMap<String, Arc<DaemonSession>>>>,
368    params: RunActionParams,
369) -> Result<RunActionResult> {
370    let thread_id = normalize_thread_id(&params.command.thread_id);
371    if thread_id.is_empty() {
372        return Err(WinxError::ThreadIdMismatch(
373            "daemon requests require an explicit thread_id".to_string(),
374        ));
375    }
376
377    let session = {
378        let mut sessions = sessions.lock().await;
379        sessions.entry(thread_id.clone()).or_insert_with(|| Arc::new(DaemonSession::new())).clone()
380    };
381    ensure_drainer(&session);
382
383    if let Some(cached) = session.completed.lock().await.get(&params.request_key).cloned() {
384        return Ok(cached);
385    }
386
387    {
388        let mut guard = session.state.lock().await;
389        match guard.as_mut() {
390            Some(state) => {
391                let daemon_cwd = state.cwd.clone();
392                state.apply_snapshot(&params.snapshot);
393                state.cwd = daemon_cwd;
394            }
395            None => {
396                let mut state = BashState::new();
397                state.apply_snapshot(&params.snapshot);
398                *guard = Some(state);
399            }
400        }
401    }
402
403    if matches!(&params.command.action_json, crate::types::BashCommandAction::Command { .. }) {
404        *session.command_id.lock().await = Some(params.request_key.clone());
405    }
406    let consumer_id = if params.consumer_id.is_empty() {
407        "legacy-adapter".to_string()
408    } else {
409        params.consumer_id
410    };
411    let cursor_key = delivery_cursor_key(&consumer_id, &params.command.action_json);
412    let cursor = {
413        let mut cursors = session.delivery_cursors.lock().await;
414        if cursors.len() >= 256 && !cursors.contains_key(&cursor_key) {
415            if let Some(oldest) = cursors.keys().next().cloned() {
416                cursors.remove(&oldest);
417            }
418        }
419        cursors
420            .entry(cursor_key)
421            .or_insert_with(|| Arc::new(Mutex::new(ShellDeliveryCursor::default())))
422            .clone()
423    };
424    let outcome = crate::tools::bash_command::handle_embedded_tool_call_with_cursor(
425        &session.state,
426        params.command,
427        &cursor,
428    )
429    .await;
430    if let Ok(output) = &outcome {
431        if let Some(id) = output.lines().find_map(|line| line.strip_prefix("bg_command_id = ")) {
432            session.background_ids.lock().await.insert(id.to_string());
433        }
434    }
435    capture_session_outputs(&session).await;
436    let snapshot =
437        session.state.lock().await.as_ref().ok_or(WinxError::BashStateNotInitialized)?.snapshot();
438    let result = match outcome {
439        Ok(output) => RunActionResult { output: Some(output), snapshot, error: None },
440        Err(error) => RunActionResult { output: None, snapshot, error: Some(to_wire_error(error)) },
441    };
442
443    let mut completed = session.completed.lock().await;
444    if completed.len() >= 256 {
445        completed.clear();
446    }
447    completed.insert(params.request_key, result.clone());
448    Ok(result)
449}
450
451fn delivery_cursor_key(consumer_id: &str, action: &crate::types::BashCommandAction) -> String {
452    let background_id = match action {
453        crate::types::BashCommandAction::StatusCheck { bg_command_id, .. }
454        | crate::types::BashCommandAction::SendText { bg_command_id, .. }
455        | crate::types::BashCommandAction::SendSpecials { bg_command_id, .. }
456        | crate::types::BashCommandAction::SendAscii { bg_command_id, .. }
457        | crate::types::BashCommandAction::Screen { bg_command_id, .. }
458        | crate::types::BashCommandAction::WaitForTurn { bg_command_id, .. } => {
459            bg_command_id.as_deref()
460        }
461        crate::types::BashCommandAction::Command { .. } => None,
462    };
463    match background_id {
464        Some(id) => format!("{consumer_id}:background:{id}"),
465        None => format!("{consumer_id}:main"),
466    }
467}
468
469fn ensure_drainer(session: &Arc<DaemonSession>) {
470    if session.drainer_started.swap(true, Ordering::AcqRel) {
471        return;
472    }
473    let session = session.clone();
474    tokio::spawn(async move {
475        loop {
476            tokio::time::sleep(Duration::from_millis(20)).await;
477            if Arc::strong_count(&session) == 1 {
478                return;
479            }
480            capture_session_outputs(&session).await;
481        }
482    });
483}
484
485async fn capture_session_outputs(session: &Arc<DaemonSession>) {
486    let (thread_id, main) = {
487        let state = session.state.lock().await;
488        let Some(state) = state.as_ref() else { return };
489        (state.current_thread_id.clone(), state.pty_shell.clone())
490    };
491    capture_shell_output(session, "main", &main).await;
492
493    let background_ids = session.background_ids.lock().await.clone();
494    for id in background_ids {
495        let shell = {
496            let store = lock_session_store();
497            store.resolve(&thread_id, &ShellTarget::Background(id.clone()))
498        };
499        if let Some(shell) = shell {
500            capture_shell_output(session, &id, &shell).await;
501        }
502    }
503}
504
505async fn capture_shell_output(
506    session: &Arc<DaemonSession>,
507    terminal_id: &str,
508    shell: &SharedPtyShell,
509) {
510    let (generation, snapshot) = {
511        let mut guard = shell.lock().await;
512        let Some(shell) = guard.as_mut() else { return };
513        shell.poll_output_nonblocking();
514        (shell.command_generation(), shell.output_snapshot())
515    };
516    let delta = {
517        let mut observed = session.observed.lock().await;
518        let (previous_generation, previous) =
519            observed.entry(terminal_id.to_string()).or_insert_with(|| (generation, String::new()));
520        if *previous_generation != generation {
521            *previous_generation = generation;
522            previous.clear();
523        }
524        let delta = snapshot
525            .strip_prefix(previous.as_str())
526            .map_or_else(|| snapshot.clone(), str::to_string);
527        previous.clone_from(&snapshot);
528        delta
529    };
530    session.journal.lock().await.append(delta);
531}
532
533async fn session_info(session: &Arc<DaemonSession>) -> Option<SessionInfo> {
534    let (thread_id, cwd, main) = {
535        let state = session.state.lock().await;
536        let state = state.as_ref()?;
537        (
538            state.current_thread_id.clone(),
539            state.cwd.to_string_lossy().into_owned(),
540            state.pty_shell.clone(),
541        )
542    };
543    let (shell_pid, running) = {
544        let guard = main.lock().await;
545        guard.as_ref().map_or((None, false), |shell| (shell.process_id(), shell.command_running))
546    };
547    let mut background_command_ids =
548        session.background_ids.lock().await.iter().cloned().collect::<Vec<_>>();
549    background_command_ids.sort();
550    Some(SessionInfo {
551        thread_id,
552        cwd,
553        shell_pid,
554        command_id: session.command_id.lock().await.clone(),
555        running,
556        background_command_ids,
557    })
558}
559
560async fn interrupt_session(session: &Arc<DaemonSession>) -> Result<()> {
561    let shell = {
562        let state = session.state.lock().await;
563        state.as_ref().map(|state| state.pty_shell.clone())
564    };
565    let Some(shell) = shell else { return Ok(()) };
566    {
567        let mut guard = shell.lock().await;
568        if let Some(shell) = guard.as_mut() {
569            shell.send_interrupt().map_err(|error| {
570                WinxError::CommandExecutionError(format!("failed to interrupt shell: {error}"))
571            })?;
572        }
573    }
574    let deadline = std::time::Instant::now() + Duration::from_secs(3);
575    while std::time::Instant::now() < deadline {
576        let recovered = {
577            let mut guard = shell.lock().await;
578            match guard.as_mut() {
579                Some(shell) => shell.poll_output_nonblocking() || !shell.command_running,
580                None => true,
581            }
582        };
583        if recovered {
584            capture_session_outputs(session).await;
585            return Ok(());
586        }
587        tokio::time::sleep(Duration::from_millis(50)).await;
588    }
589    Err(WinxError::CommandExecutionError(
590        "interrupted shell did not return to a prompt within 3 seconds".to_string(),
591    ))
592}
593
594fn to_wire_error(error: WinxError) -> WireShellError {
595    match error {
596        WinxError::BashStateNotInitialized => WireShellError::BashStateNotInitialized,
597        WinxError::ShellInitializationError(message) => {
598            WireShellError::ShellInitialization(message)
599        }
600        WinxError::CommandExecutionError(message) => WireShellError::CommandExecution(message),
601        WinxError::CommandNotAllowed(message) => WireShellError::CommandNotAllowed(message),
602        WinxError::ThreadIdMismatch(message) => WireShellError::ThreadIdMismatch(message),
603        other => WireShellError::Other(other.to_string()),
604    }
605}
606
607fn rpc_result<T: serde::Serialize>(id: u64, result: &T) -> RpcResponse {
608    match serde_json::to_value(result) {
609        Ok(result) => {
610            RpcResponse { jsonrpc: "2.0".to_string(), id, result: Some(result), error: None }
611        }
612        Err(error) => rpc_error(id, -32603, &error.to_string()),
613    }
614}
615
616fn rpc_error(id: u64, code: i32, message: &str) -> RpcResponse {
617    RpcResponse {
618        jsonrpc: "2.0".to_string(),
619        id,
620        result: None,
621        error: Some(RpcError { code, message: message.to_string() }),
622    }
623}
624
625fn same_uid(stream: &UnixStream) -> Result<bool> {
626    let credentials = stream.peer_cred()?;
627    // SAFETY: geteuid has no preconditions and reads process credentials only.
628    Ok(credentials.uid() == unsafe { libc::geteuid() })
629}
630
631#[cfg(test)]
632mod tests {
633    use super::OutputJournal;
634
635    #[test]
636    fn late_consumer_is_told_when_journal_head_was_dropped() {
637        let mut journal = OutputJournal::default();
638        journal.append("a".repeat(3 * 1024 * 1024));
639        journal.append("b".repeat(3 * 1024 * 1024));
640
641        let read = journal.read("late");
642        assert!(read.gap);
643        assert!(read.output.starts_with('b'));
644    }
645}
winx_code_agent/daemon/server.rs

winx_code_agent/daemon/
server.rs
