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, AtomicU64, Ordering};
6use std::sync::Arc;
7use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
8
9use tokio::net::{UnixListener, UnixStream};
10use tokio::sync::{Mutex, Notify};
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 TimedEntry<T> {
28    value: T,
29    last_seen: Instant,
30}
31
32struct DaemonSession {
33    state: SharedState,
34    completed: Mutex<HashMap<String, TimedEntry<RunActionResult>>>,
35    journal: Mutex<OutputJournal>,
36    observed: Mutex<HashMap<String, (u64, String)>>,
37    background_ids: Mutex<HashSet<String>>,
38    command_id: Mutex<Option<String>>,
39    delivery_cursors: Mutex<HashMap<String, TimedEntry<Arc<Mutex<ShellDeliveryCursor>>>>>,
40    created_at_unix_ms: u64,
41    last_activity_unix_ms: AtomicU64,
42    last_command_at_unix_ms: AtomicU64,
43    ever_ran_command: AtomicBool,
44    drainer_started: AtomicBool,
45    activity: Notify,
46}
47
48impl DaemonSession {
49    fn new() -> Self {
50        let now = unix_ms();
51        Self {
52            state: Arc::new(Mutex::new(None)),
53            completed: Mutex::new(HashMap::new()),
54            journal: Mutex::new(OutputJournal::default()),
55            observed: Mutex::new(HashMap::new()),
56            background_ids: Mutex::new(HashSet::new()),
57            command_id: Mutex::new(None),
58            delivery_cursors: Mutex::new(HashMap::new()),
59            created_at_unix_ms: now,
60            last_activity_unix_ms: AtomicU64::new(now),
61            last_command_at_unix_ms: AtomicU64::new(0),
62            ever_ran_command: AtomicBool::new(false),
63            drainer_started: AtomicBool::new(false),
64            activity: Notify::new(),
65        }
66    }
67
68    fn note_activity(&self, command: bool) {
69        let now = unix_ms();
70        self.last_activity_unix_ms.store(now, Ordering::Release);
71        if command {
72            self.last_command_at_unix_ms.store(now, Ordering::Release);
73            self.ever_ran_command.store(true, Ordering::Release);
74        }
75        self.activity.notify_one();
76    }
77}
78
79const MAX_JOURNAL_BYTES: usize = 4 * 1024 * 1024;
80const MAX_SESSION_CACHE_ENTRIES: usize = 256;
81const MAX_IDENTIFIER_BYTES: usize = 256;
82const SESSION_CACHE_TTL: Duration = Duration::from_secs(60 * 60);
83const DRAIN_CHANGED_INTERVAL: Duration = Duration::from_millis(20);
84const DRAIN_ACTIVE_INTERVAL: Duration = Duration::from_millis(100);
85const DRAIN_IDLE_INTERVAL: Duration = Duration::from_secs(1);
86
87#[derive(Clone, Copy, Default)]
88struct CaptureState {
89    changed: bool,
90    active: bool,
91}
92
93impl CaptureState {
94    fn merge(&mut self, other: Self) {
95        self.changed |= other.changed;
96        self.active |= other.active;
97    }
98
99    fn next_delay(self) -> Duration {
100        if self.changed {
101            DRAIN_CHANGED_INTERVAL
102        } else if self.active {
103            DRAIN_ACTIVE_INTERVAL
104        } else {
105            DRAIN_IDLE_INTERVAL
106        }
107    }
108}
109
110#[derive(Default)]
111struct OutputJournal {
112    chunks: VecDeque<(u64, String)>,
113    cursors: HashMap<String, TimedEntry<u64>>,
114    next_seq: u64,
115    bytes: usize,
116}
117
118fn prune_timed_entries<T>(entries: &mut HashMap<String, TimedEntry<T>>, now: Instant) {
119    entries.retain(|_, entry| now.duration_since(entry.last_seen) <= SESSION_CACHE_TTL);
120}
121
122fn evict_oldest_if_full<T>(entries: &mut HashMap<String, TimedEntry<T>>, incoming: &str) {
123    if entries.contains_key(incoming) || entries.len() < MAX_SESSION_CACHE_ENTRIES {
124        return;
125    }
126    if let Some(oldest) =
127        entries.iter().min_by_key(|(_, entry)| entry.last_seen).map(|(key, _)| key.clone())
128    {
129        entries.remove(&oldest);
130    }
131}
132
133fn validate_identifier(label: &str, value: &str) -> Result<()> {
134    if value.len() > MAX_IDENTIFIER_BYTES {
135        return Err(WinxError::InvalidInput(format!(
136            "{label} exceeds the {MAX_IDENTIFIER_BYTES}-byte limit"
137        )));
138    }
139    Ok(())
140}
141
142impl OutputJournal {
143    fn append(&mut self, output: String) {
144        if output.is_empty() {
145            return;
146        }
147        self.next_seq = self.next_seq.saturating_add(1);
148        self.bytes = self.bytes.saturating_add(output.len());
149        self.chunks.push_back((self.next_seq, output));
150        while self.bytes > MAX_JOURNAL_BYTES && self.chunks.len() > 1 {
151            if let Some((_, removed)) = self.chunks.pop_front() {
152                self.bytes = self.bytes.saturating_sub(removed.len());
153            }
154        }
155    }
156
157    fn read(&mut self, consumer_id: &str) -> JournalRead {
158        let now = Instant::now();
159        prune_timed_entries(&mut self.cursors, now);
160        evict_oldest_if_full(&mut self.cursors, consumer_id);
161
162        let oldest = self.chunks.front().map_or(self.next_seq.saturating_add(1), |(seq, _)| *seq);
163        let cursor = self.cursors.get(consumer_id).map_or(0, |entry| entry.value);
164        let gap = cursor.saturating_add(1) < oldest;
165        let mut output = String::new();
166        let mut next_seq = cursor;
167        for (seq, chunk) in self.chunks.iter().filter(|(seq, _)| *seq > cursor) {
168            output.push_str(chunk);
169            next_seq = *seq;
170        }
171        self.cursors
172            .insert(consumer_id.to_string(), TimedEntry { value: next_seq, last_seen: now });
173        JournalRead { output, next_seq, gap }
174    }
175}
176
177/// Long-lived JSON-RPC shell owner listening on a Unix-domain socket.
178pub struct DaemonServer {
179    listener: UnixListener,
180    socket_path: PathBuf,
181    sessions: Arc<Mutex<HashMap<String, Arc<DaemonSession>>>>,
182    epoch: String,
183}
184
185impl DaemonServer {
186    pub async fn bind(socket_path: impl AsRef<Path>) -> Result<Self> {
187        let socket_path = socket_path.as_ref().to_path_buf();
188        let parent = socket_path.parent().ok_or_else(|| {
189            WinxError::ConfigurationError("daemon socket must have a parent directory".to_string())
190        })?;
191        tokio::fs::create_dir_all(parent).await?;
192        tokio::fs::set_permissions(parent, std::fs::Permissions::from_mode(0o700)).await?;
193
194        if tokio::fs::try_exists(&socket_path).await? {
195            if UnixStream::connect(&socket_path).await.is_ok() {
196                return Err(WinxError::ConfigurationError(format!(
197                    "a winxd instance is already listening at {}",
198                    socket_path.display()
199                )));
200            }
201            tokio::fs::remove_file(&socket_path).await?;
202        }
203
204        let listener = UnixListener::bind(&socket_path)?;
205        tokio::fs::set_permissions(&socket_path, std::fs::Permissions::from_mode(0o600)).await?;
206        let epoch = format!("{:016x}", rand::random::<u64>());
207        Ok(Self { listener, socket_path, sessions: Arc::new(Mutex::new(HashMap::new())), epoch })
208    }
209
210    pub async fn serve(self) -> Result<()> {
211        loop {
212            let (stream, _) = self.listener.accept().await?;
213            if !same_uid(&stream)? {
214                tracing::warn!("Rejected winxd connection from a different uid");
215                continue;
216            }
217            let sessions = self.sessions.clone();
218            let epoch = self.epoch.clone();
219            tokio::spawn(async move {
220                if let Err(error) = serve_connection(stream, sessions, epoch).await {
221                    tracing::debug!("winxd client disconnected: {error}");
222                }
223            });
224        }
225    }
226}
227
228impl Drop for DaemonServer {
229    fn drop(&mut self) {
230        let _ = std::fs::remove_file(&self.socket_path);
231    }
232}
233
234async fn serve_connection(
235    mut stream: UnixStream,
236    sessions: Arc<Mutex<HashMap<String, Arc<DaemonSession>>>>,
237    epoch: String,
238) -> Result<()> {
239    loop {
240        let request: RpcRequest = match read_json_frame(&mut stream).await {
241            Ok(request) => request,
242            Err(error) if error.kind() == ErrorKind::UnexpectedEof => return Ok(()),
243            Err(error) => return Err(error.into()),
244        };
245        let response = dispatch(request, &sessions, &epoch).await;
246        write_json_frame(&mut stream, &response).await?;
247    }
248}
249
250#[allow(clippy::too_many_lines)]
251async fn dispatch(
252    request: RpcRequest,
253    sessions: &Arc<Mutex<HashMap<String, Arc<DaemonSession>>>>,
254    epoch: &str,
255) -> RpcResponse {
256    if request.jsonrpc != "2.0" {
257        return rpc_error(request.id, -32600, "JSON-RPC version must be 2.0");
258    }
259
260    match request.method.as_str() {
261        "winx.hello" => rpc_result(
262            request.id,
263            &HelloResult {
264                protocol_major: PROTOCOL_MAJOR,
265                protocol_minor: PROTOCOL_MINOR,
266                capabilities: vec![
267                    "session.configure".to_string(),
268                    "shell.run_action".to_string(),
269                    "session.list".to_string(),
270                    "session.read_output".to_string(),
271                    "session.kill".to_string(),
272                    "multi_consumer_cursors".to_string(),
273                    "idempotency".to_string(),
274                    "session_activity_timestamps".to_string(),
275                    "attach_or_create".to_string(),
276                ],
277                max_frame_bytes: MAX_FRAME_BYTES,
278                daemon_epoch: epoch.to_string(),
279                daemon_pid: std::process::id(),
280            },
281        ),
282        "session.configure" => {
283            let params: ConfigureSessionParams = match serde_json::from_value(request.params) {
284                Ok(params) => params,
285                Err(error) => return rpc_error(request.id, -32602, &error.to_string()),
286            };
287            match configure_session(sessions, params).await {
288                Ok(result) => rpc_result(request.id, &result),
289                Err(error) => rpc_error(request.id, -32603, &error.to_string()),
290            }
291        }
292        "shell.run_action" => {
293            let params: RunActionParams = match serde_json::from_value(request.params) {
294                Ok(params) => params,
295                Err(error) => return rpc_error(request.id, -32602, &error.to_string()),
296            };
297            match run_action(sessions, params).await {
298                Ok(result) => rpc_result(request.id, &result),
299                Err(error) => rpc_error(request.id, -32603, &error.to_string()),
300            }
301        }
302        "session.list" => {
303            let entries = sessions.lock().await.values().cloned().collect::<Vec<_>>();
304            let mut infos = Vec::with_capacity(entries.len());
305            for session in entries {
306                if let Some(info) = session_info(&session).await {
307                    infos.push(info);
308                }
309            }
310            infos.sort_by(|left, right| left.thread_id.cmp(&right.thread_id));
311            rpc_result(request.id, &infos)
312        }
313        "session.info" => {
314            let params: SessionParams = match serde_json::from_value(request.params) {
315                Ok(params) => params,
316                Err(error) => return rpc_error(request.id, -32602, &error.to_string()),
317            };
318            let session =
319                sessions.lock().await.get(&normalize_thread_id(&params.thread_id)).cloned();
320            match session {
321                Some(session) => match session_info(&session).await {
322                    Some(info) => rpc_result(request.id, &info),
323                    None => rpc_error(request.id, -32004, "session not initialized"),
324                },
325                None => rpc_error(request.id, -32004, "session not found"),
326            }
327        }
328        "session.read_output" => {
329            let params: JournalReadParams = match serde_json::from_value(request.params) {
330                Ok(params) => params,
331                Err(error) => return rpc_error(request.id, -32602, &error.to_string()),
332            };
333            if let Err(error) = validate_identifier("consumer_id", &params.consumer_id) {
334                return rpc_error(request.id, -32602, &error.to_string());
335            }
336            let thread_id = normalize_thread_id(&params.thread_id);
337            let session = sessions.lock().await.get(&thread_id).cloned();
338            match session {
339                Some(session) => {
340                    session.note_activity(false);
341                    capture_session_outputs(&session).await;
342                    let read = session.journal.lock().await.read(&params.consumer_id);
343                    rpc_result(request.id, &read)
344                }
345                None => rpc_error(request.id, -32004, "session not found"),
346            }
347        }
348        "session.kill" => {
349            let params: SessionParams = match serde_json::from_value(request.params) {
350                Ok(params) => params,
351                Err(error) => return rpc_error(request.id, -32602, &error.to_string()),
352            };
353            let session = sessions.lock().await.remove(&normalize_thread_id(&params.thread_id));
354            match session {
355                Some(session) => {
356                    let background_ids = session.background_ids.lock().await.clone();
357                    {
358                        let mut store = lock_session_store();
359                        for id in background_ids {
360                            store.remove_shell(&id);
361                        }
362                    }
363                    *session.state.lock().await = None;
364                    rpc_result(request.id, &true)
365                }
366                None => rpc_result(request.id, &false),
367            }
368        }
369        "session.interrupt" => {
370            let params: SessionParams = match serde_json::from_value(request.params) {
371                Ok(params) => params,
372                Err(error) => return rpc_error(request.id, -32602, &error.to_string()),
373            };
374            let session =
375                sessions.lock().await.get(&normalize_thread_id(&params.thread_id)).cloned();
376            match session {
377                Some(session) => match interrupt_session(&session).await {
378                    Ok(()) => rpc_result(request.id, &true),
379                    Err(error) => rpc_error(request.id, -32005, &error.to_string()),
380                },
381                None => rpc_error(request.id, -32004, "session not found"),
382            }
383        }
384        _ => rpc_error(request.id, -32601, "method not found"),
385    }
386}
387
388async fn configure_session(
389    sessions: &Arc<Mutex<HashMap<String, Arc<DaemonSession>>>>,
390    params: ConfigureSessionParams,
391) -> Result<ConfigureSessionResult> {
392    let thread_id = normalize_thread_id(&params.snapshot.chat_id);
393    if thread_id.is_empty() {
394        return Err(WinxError::ThreadIdMismatch(
395            "daemon session configuration requires an explicit thread_id".to_string(),
396        ));
397    }
398
399    let (session, entry_created) = {
400        let mut sessions = sessions.lock().await;
401        match params.transition {
402            ConfigureSessionTransition::FirstCall => match sessions.entry(thread_id) {
403                std::collections::hash_map::Entry::Occupied(entry) => (entry.get().clone(), false),
404                std::collections::hash_map::Entry::Vacant(entry) => {
405                    let session = Arc::new(DaemonSession::new());
406                    entry.insert(session.clone());
407                    (session, true)
408                }
409            },
410            _ => (
411                sessions.get(&thread_id).cloned().ok_or(WinxError::BashStateNotInitialized)?,
412                false,
413            ),
414        }
415    };
416    ensure_drainer(&session);
417    session.note_activity(false);
418
419    let (attach_hint, snapshot, attached_existing) = {
420        let mut guard = session.state.lock().await;
421        let mut attached_existing = false;
422        match params.transition {
423            ConfigureSessionTransition::FirstCall => {
424                if entry_created || guard.is_none() {
425                    let mut state = BashState::new();
426                    state.apply_snapshot(&params.snapshot);
427                    if state.cwd.exists() {
428                        state.init_pty_shell().await?;
429                    }
430                    *guard = Some(state);
431                } else {
432                    // Attach-or-create: preserve the guardian-owned PTY, cwd, mode,
433                    // output journal, and command state across adapter reconnects.
434                    attached_existing = true;
435                }
436            }
437            ConfigureSessionTransition::ModeChange => {
438                let state = guard.as_mut().ok_or(WinxError::BashStateNotInitialized)?;
439                let daemon_cwd = state.cwd.clone();
440                state.apply_snapshot(&params.snapshot);
441                state.cwd = daemon_cwd;
442            }
443            ConfigureSessionTransition::Reset => {
444                let state = guard.as_mut().ok_or(WinxError::BashStateNotInitialized)?;
445                let daemon_cwd = state.cwd.clone();
446                state.apply_snapshot(&params.snapshot);
447                state.cwd = daemon_cwd;
448                state.init_pty_shell().await?;
449            }
450            ConfigureSessionTransition::WorkspaceChange => {
451                let state = guard.as_mut().ok_or(WinxError::BashStateNotInitialized)?;
452                state.apply_snapshot(&params.snapshot);
453            }
454        }
455        let state = guard.as_ref().ok_or(WinxError::BashStateNotInitialized)?;
456        let attach_hint = {
457            let shell = state.pty_shell.lock().await;
458            shell.as_ref().and_then(|shell| shell.attach_hint.clone())
459        };
460        (attach_hint, state.snapshot(), attached_existing)
461    };
462
463    if matches!(params.transition, ConfigureSessionTransition::Reset)
464        || (matches!(params.transition, ConfigureSessionTransition::FirstCall)
465            && !attached_existing)
466    {
467        *session.command_id.lock().await = None;
468        session.completed.lock().await.clear();
469        session.delivery_cursors.lock().await.clear();
470        session.observed.lock().await.remove("main");
471    }
472    capture_session_outputs(&session).await;
473    session.activity.notify_one();
474    Ok(ConfigureSessionResult { attach_hint, snapshot: Some(snapshot), attached_existing })
475}
476
477async fn run_action(
478    sessions: &Arc<Mutex<HashMap<String, Arc<DaemonSession>>>>,
479    params: RunActionParams,
480) -> Result<RunActionResult> {
481    validate_identifier("request_key", &params.request_key)?;
482    validate_identifier("consumer_id", &params.consumer_id)?;
483    let thread_id = normalize_thread_id(&params.command.thread_id);
484    if thread_id.is_empty() {
485        return Err(WinxError::ThreadIdMismatch(
486            "daemon requests require an explicit thread_id".to_string(),
487        ));
488    }
489
490    let is_command =
491        matches!(&params.command.action_json, crate::types::BashCommandAction::Command { .. });
492    let session = {
493        let mut sessions = sessions.lock().await;
494        sessions.entry(thread_id.clone()).or_insert_with(|| Arc::new(DaemonSession::new())).clone()
495    };
496    ensure_drainer(&session);
497    session.note_activity(is_command);
498
499    let now = Instant::now();
500    let cached = {
501        let mut completed = session.completed.lock().await;
502        prune_timed_entries(&mut completed, now);
503        completed.get_mut(&params.request_key).map(|entry| {
504            entry.last_seen = now;
505            entry.value.clone()
506        })
507    };
508    if let Some(cached) = cached {
509        return Ok(cached);
510    }
511
512    {
513        let mut guard = session.state.lock().await;
514        let state = if let Some(state) = guard.as_mut() {
515            let daemon_cwd = state.cwd.clone();
516            state.apply_snapshot(&params.snapshot);
517            state.cwd = daemon_cwd;
518            state
519        } else {
520            let mut state = BashState::new();
521            state.apply_snapshot(&params.snapshot);
522            guard.insert(state)
523        };
524        let needs_shell = state.pty_shell.lock().await.is_none();
525        if needs_shell && state.cwd.exists() {
526            state.init_pty_shell().await?;
527        }
528    }
529
530    if is_command {
531        *session.command_id.lock().await = Some(params.request_key.clone());
532    }
533    let consumer_id = if params.consumer_id.is_empty() {
534        "legacy-adapter".to_string()
535    } else {
536        params.consumer_id
537    };
538    let cursor_key = delivery_cursor_key(&consumer_id, &params.command.action_json);
539    let cursor = {
540        let now = Instant::now();
541        let mut cursors = session.delivery_cursors.lock().await;
542        prune_timed_entries(&mut cursors, now);
543        evict_oldest_if_full(&mut cursors, &cursor_key);
544        let entry = cursors.entry(cursor_key).or_insert_with(|| TimedEntry {
545            value: Arc::new(Mutex::new(ShellDeliveryCursor::default())),
546            last_seen: now,
547        });
548        entry.last_seen = now;
549        entry.value.clone()
550    };
551    let outcome = crate::tools::bash_command::handle_embedded_tool_call_with_cursor(
552        &session.state,
553        params.command,
554        &cursor,
555    )
556    .await;
557    if let Ok(output) = &outcome {
558        if let Some(id) = output.lines().find_map(|line| line.strip_prefix("bg_command_id = ")) {
559            session.background_ids.lock().await.insert(id.to_string());
560        }
561    }
562    capture_session_outputs(&session).await;
563    let snapshot =
564        session.state.lock().await.as_ref().ok_or(WinxError::BashStateNotInitialized)?.snapshot();
565    let result = match outcome {
566        Ok(output) => RunActionResult { output: Some(output), snapshot, error: None },
567        Err(error) => RunActionResult { output: None, snapshot, error: Some(to_wire_error(error)) },
568    };
569
570    let mut completed = session.completed.lock().await;
571    let now = Instant::now();
572    prune_timed_entries(&mut completed, now);
573    evict_oldest_if_full(&mut completed, &params.request_key);
574    completed.insert(params.request_key, TimedEntry { value: result.clone(), last_seen: now });
575    drop(completed);
576    session.activity.notify_one();
577    Ok(result)
578}
579
580fn delivery_cursor_key(consumer_id: &str, action: &crate::types::BashCommandAction) -> String {
581    let background_id = match action {
582        crate::types::BashCommandAction::StatusCheck { bg_command_id, .. }
583        | crate::types::BashCommandAction::SendText { bg_command_id, .. }
584        | crate::types::BashCommandAction::SendSpecials { bg_command_id, .. }
585        | crate::types::BashCommandAction::SendAscii { bg_command_id, .. }
586        | crate::types::BashCommandAction::Screen { bg_command_id, .. }
587        | crate::types::BashCommandAction::WaitForTurn { bg_command_id, .. } => {
588            bg_command_id.as_deref()
589        }
590        crate::types::BashCommandAction::Command { .. } => None,
591    };
592    match background_id {
593        Some(id) => format!("{consumer_id}:background:{id}"),
594        None => format!("{consumer_id}:main"),
595    }
596}
597
598fn ensure_drainer(session: &Arc<DaemonSession>) {
599    if session.drainer_started.swap(true, Ordering::AcqRel) {
600        session.activity.notify_one();
601        return;
602    }
603    let session = session.clone();
604    tokio::spawn(async move {
605        loop {
606            if Arc::strong_count(&session) == 1 {
607                return;
608            }
609            let capture = capture_session_outputs(&session).await;
610            tokio::select! {
611                () = tokio::time::sleep(capture.next_delay()) => {}
612                () = session.activity.notified() => {}
613            }
614        }
615    });
616}
617
618async fn capture_session_outputs(session: &Arc<DaemonSession>) -> CaptureState {
619    let (thread_id, main) = {
620        let state = session.state.lock().await;
621        let Some(state) = state.as_ref() else { return CaptureState::default() };
622        (state.current_thread_id.clone(), state.pty_shell.clone())
623    };
624    let mut capture = capture_shell_output(session, "main", &main).await;
625
626    let background_ids = session.background_ids.lock().await.clone();
627    for id in background_ids {
628        let shell = {
629            let store = lock_session_store();
630            store.resolve(&thread_id, &ShellTarget::Background(id.clone()))
631        };
632        if let Some(shell) = shell {
633            capture.merge(capture_shell_output(session, &id, &shell).await);
634        }
635    }
636    capture
637}
638
639async fn capture_shell_output(
640    session: &Arc<DaemonSession>,
641    terminal_id: &str,
642    shell: &SharedPtyShell,
643) -> CaptureState {
644    let (generation, snapshot, active) = {
645        let mut guard = shell.lock().await;
646        let Some(shell) = guard.as_mut() else { return CaptureState::default() };
647        shell.poll_output_nonblocking();
648        (shell.command_generation(), shell.output_snapshot(), shell.command_running)
649    };
650    let delta = {
651        let mut observed = session.observed.lock().await;
652        let (previous_generation, previous) =
653            observed.entry(terminal_id.to_string()).or_insert_with(|| (generation, String::new()));
654        if *previous_generation != generation {
655            *previous_generation = generation;
656            previous.clear();
657        }
658        let delta = snapshot
659            .strip_prefix(previous.as_str())
660            .map_or_else(|| snapshot.clone(), str::to_string);
661        previous.clone_from(&snapshot);
662        delta
663    };
664    let changed = !delta.is_empty();
665    if changed {
666        session.last_activity_unix_ms.store(unix_ms(), Ordering::Release);
667        session.journal.lock().await.append(delta);
668    }
669    CaptureState { changed, active }
670}
671
672async fn session_info(session: &Arc<DaemonSession>) -> Option<SessionInfo> {
673    let (thread_id, cwd, main) = {
674        let state = session.state.lock().await;
675        let state = state.as_ref()?;
676        (
677            state.current_thread_id.clone(),
678            state.cwd.to_string_lossy().into_owned(),
679            state.pty_shell.clone(),
680        )
681    };
682    let (shell_pid, running) = {
683        let guard = main.lock().await;
684        guard.as_ref().map_or((None, false), |shell| (shell.process_id(), shell.command_running))
685    };
686    let background_command_ids = active_background_command_ids(session, &thread_id).await;
687    let last_command_at_unix_ms = session.last_command_at_unix_ms.load(Ordering::Acquire);
688    Some(SessionInfo {
689        thread_id,
690        cwd,
691        shell_pid,
692        command_id: session.command_id.lock().await.clone(),
693        running,
694        background_command_ids,
695        created_at_unix_ms: Some(session.created_at_unix_ms),
696        last_activity_unix_ms: Some(session.last_activity_unix_ms.load(Ordering::Acquire)),
697        last_command_at_unix_ms: (last_command_at_unix_ms > 0).then_some(last_command_at_unix_ms),
698        ever_ran_command: session.ever_ran_command.load(Ordering::Acquire),
699    })
700}
701
702async fn active_background_command_ids(
703    session: &Arc<DaemonSession>,
704    thread_id: &str,
705) -> Vec<String> {
706    let candidates = session.background_ids.lock().await.clone();
707    let mut active = Vec::new();
708    for id in candidates {
709        let shell = {
710            let store = lock_session_store();
711            store.resolve(thread_id, &ShellTarget::Background(id.clone()))
712        };
713        let Some(shell) = shell else { continue };
714        // Capture the final queued bytes before deciding that a completed shell is
715        // no longer active. Otherwise pruning the id here could make the journal
716        // drainer stop one poll too early.
717        capture_shell_output(session, &id, &shell).await;
718        let running = {
719            let guard = shell.lock().await;
720            guard.as_ref().is_some_and(|shell| shell.command_running)
721        };
722        if running {
723            active.push(id);
724        }
725    }
726
727    let active_set = active.iter().cloned().collect::<HashSet<_>>();
728    session.background_ids.lock().await.retain(|id| active_set.contains(id));
729    active.sort();
730    active
731}
732
733async fn interrupt_session(session: &Arc<DaemonSession>) -> Result<()> {
734    session.note_activity(false);
735    let shell = {
736        let state = session.state.lock().await;
737        state.as_ref().map(|state| state.pty_shell.clone())
738    };
739    let Some(shell) = shell else { return Ok(()) };
740    {
741        let mut guard = shell.lock().await;
742        if let Some(shell) = guard.as_mut() {
743            shell.send_interrupt().map_err(|error| {
744                WinxError::CommandExecutionError(format!("failed to interrupt shell: {error}"))
745            })?;
746        }
747    }
748    let deadline = std::time::Instant::now() + Duration::from_secs(3);
749    while std::time::Instant::now() < deadline {
750        let recovered = {
751            let mut guard = shell.lock().await;
752            match guard.as_mut() {
753                Some(shell) => shell.poll_output_nonblocking() || !shell.command_running,
754                None => true,
755            }
756        };
757        if recovered {
758            capture_session_outputs(session).await;
759            return Ok(());
760        }
761        tokio::time::sleep(Duration::from_millis(50)).await;
762    }
763    Err(WinxError::CommandExecutionError(
764        "interrupted shell did not return to a prompt within 3 seconds".to_string(),
765    ))
766}
767
768fn unix_ms() -> u64 {
769    SystemTime::now()
770        .duration_since(UNIX_EPOCH)
771        .ok()
772        .and_then(|duration| u64::try_from(duration.as_millis()).ok())
773        .unwrap_or(0)
774}
775
776fn to_wire_error(error: WinxError) -> WireShellError {
777    match error {
778        WinxError::BashStateNotInitialized => WireShellError::BashStateNotInitialized,
779        WinxError::ShellInitializationError(message) => {
780            WireShellError::ShellInitialization(message)
781        }
782        WinxError::CommandExecutionError(message) => WireShellError::CommandExecution(message),
783        WinxError::NoActiveCommand(message) => WireShellError::NoActiveCommand(message),
784        WinxError::BackgroundSessionNotFound(message) => {
785            WireShellError::BackgroundSessionNotFound(message)
786        }
787        WinxError::EmptyInteractiveInput { action } => {
788            WireShellError::EmptyInteractiveInput(action)
789        }
790        WinxError::InteractiveTargetNotRunning(message) => {
791            WireShellError::InteractiveTargetNotRunning(message)
792        }
793        WinxError::CommandAlreadyRunning { current_command, duration_seconds } => {
794            WireShellError::CommandAlreadyRunning { current_command, duration_seconds }
795        }
796        WinxError::CommandNotAllowed(message) => WireShellError::CommandNotAllowed(message),
797        WinxError::ThreadIdMismatch(message) => WireShellError::ThreadIdMismatch(message),
798        WinxError::InvalidInput(message) | WinxError::ArgumentParseError(message) => {
799            WireShellError::InvalidInput(message)
800        }
801        other => WireShellError::Other(other.to_string()),
802    }
803}
804
805fn rpc_result<T: serde::Serialize>(id: u64, result: &T) -> RpcResponse {
806    match serde_json::to_value(result) {
807        Ok(result) => {
808            RpcResponse { jsonrpc: "2.0".to_string(), id, result: Some(result), error: None }
809        }
810        Err(error) => rpc_error(id, -32603, &error.to_string()),
811    }
812}
813
814fn rpc_error(id: u64, code: i32, message: &str) -> RpcResponse {
815    RpcResponse {
816        jsonrpc: "2.0".to_string(),
817        id,
818        result: None,
819        error: Some(RpcError { code, message: message.to_string() }),
820    }
821}
822
823fn same_uid(stream: &UnixStream) -> Result<bool> {
824    let credentials = stream.peer_cred()?;
825    Ok(credentials.uid() == crate::os::unix::effective_uid())
826}
827
828#[cfg(test)]
829mod tests {
830    use super::{validate_identifier, CaptureState, OutputJournal, MAX_SESSION_CACHE_ENTRIES};
831
832    #[test]
833    fn late_consumer_is_told_when_journal_head_was_dropped() {
834        let mut journal = OutputJournal::default();
835        journal.append("a".repeat(3 * 1024 * 1024));
836        journal.append("b".repeat(3 * 1024 * 1024));
837
838        let read = journal.read("late");
839        assert!(read.gap);
840        assert!(read.output.starts_with('b'));
841    }
842
843    #[test]
844    fn journal_consumer_cursors_are_lru_bounded() {
845        let mut journal = OutputJournal::default();
846        journal.append("output".to_string());
847        for index in 0..(MAX_SESSION_CACHE_ENTRIES + 100) {
848            let _ = journal.read(&format!("consumer-{index}"));
849        }
850        assert_eq!(journal.cursors.len(), MAX_SESSION_CACHE_ENTRIES);
851    }
852
853    #[test]
854    fn daemon_identifiers_have_a_small_memory_bound() {
855        assert!(validate_identifier("consumer_id", &"x".repeat(256)).is_ok());
856        assert!(validate_identifier("consumer_id", &"x".repeat(257)).is_err());
857    }
858
859    #[test]
860    fn drainer_backoff_tracks_output_and_command_activity() {
861        assert!(
862            CaptureState { changed: true, active: false }.next_delay()
863                < CaptureState { changed: false, active: true }.next_delay()
864        );
865        assert!(
866            CaptureState { changed: false, active: true }.next_delay()
867                < CaptureState { changed: false, active: false }.next_delay()
868        );
869    }
870}
winx_code_agent/daemon/server.rs

winx_code_agent/daemon/
server.rs
