Skip to main content

agentos_client/
process.rs

1//! Process execution & management methods + supporting types.
2//!
3//! Ported from `packages/core/src/agent-os.ts` (process methods) and `runtime-compat.ts`
4//! (`ExecOptions`, `ExecResult`, `ProcessInfo`, etc.).
5//!
6//! Two distinct process views: SDK-spawned processes (`processes` map, keyed by user-facing pid)
7//! back `spawn` + the stdin/stdout/stderr/exit subscriptions + `wait/list/get/stop/kill`; the kernel
8//! process table backs `exec`, `all_processes`, `process_tree`.
9
10use std::collections::BTreeMap;
11use std::sync::atomic::Ordering;
12
13use anyhow::{Context, Result};
14use scc::HashMap as SccHashMap;
15use serde::{Deserialize, Serialize};
16use tokio::sync::{broadcast, watch};
17use tokio::task::JoinHandle;
18
19use agentos_sidecar_client::wire::{self, EventPayload, ProcessSnapshotStatus, StreamChannel};
20
21use crate::agent_os::{AgentOs, ProcessEntry};
22use crate::command_line::resolve_exec_command;
23use crate::error::ClientError;
24use crate::stream::Subscription;
25
26/// Broadcast channel capacity for a spawned process's stdout/stderr fan-out.
27const PROCESS_STREAM_CAPACITY: usize = 1024;
28
29/// Maximum SDK-spawned process entries retained per VM.
30const PROCESS_REGISTRY_LIMIT: usize = 1024;
31
32/// Maximum first-observed process timestamp entries retained per VM.
33const OBSERVED_PROCESS_TIME_LIMIT: usize = 4096;
34
35/// Maximum bytes captured by `exec` across stdout and stderr.
36const EXEC_OUTPUT_CAPTURE_LIMIT_BYTES: usize = 16 * 1024 * 1024;
37
38/// Default guest working directory for `exec`/`spawn`, matching the TS sidecar client.
39pub(crate) const DEFAULT_EXEC_CWD: &str = "/workspace";
40
41/// Base value for the synthetic display-pid sequence used by `spawn` (TS `SYNTHETIC_PID_BASE`). The
42/// first spawned process is assigned exactly this value.
43pub(crate) const SYNTHETIC_PID_BASE: u64 = 1_000_000;
44
45// ---------------------------------------------------------------------------
46// Supporting types
47// ---------------------------------------------------------------------------
48
49/// Timing-mitigation mode for an execution.
50#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
51#[serde(rename_all = "lowercase")]
52pub enum TimingMitigation {
53    #[default]
54    Off,
55    Freeze,
56}
57
58/// `stdin` value: a string or raw bytes.
59#[derive(Debug, Clone, PartialEq, Eq)]
60pub enum StdinInput {
61    Text(String),
62    Bytes(Vec<u8>),
63}
64
65/// A raw-byte streaming callback for stdout/stderr (TS `(data: Uint8Array) => void`). Invoked once
66/// per output chunk as it arrives. Never assume UTF-8: chunks are delivered as raw bytes.
67pub type OutputCallback = Box<dyn FnMut(&[u8]) + Send>;
68
69/// Base options shared by `exec` and `spawn`.
70///
71/// `on_stdout`/`on_stderr` mirror the TS `ExecOptions.onStdout`/`onStderr` raw-byte streaming
72/// callbacks. For `exec` they fire for the duration of the call; for `spawn` they are seeded into the
73/// stdout/stderr fan-out at spawn time (matching the TS initial-handler-set behavior).
74pub struct ExecOptions {
75    pub env: BTreeMap<String, String>,
76    pub cwd: Option<String>,
77    pub stdin: Option<StdinInput>,
78    pub timeout: Option<f64>,
79    pub on_stdout: Option<OutputCallback>,
80    pub on_stderr: Option<OutputCallback>,
81    pub capture_stdio: Option<bool>,
82    pub file_path: Option<String>,
83    pub cpu_time_limit_ms: Option<f64>,
84    pub timing_mitigation: Option<TimingMitigation>,
85}
86
87impl Default for ExecOptions {
88    fn default() -> Self {
89        Self {
90            env: BTreeMap::new(),
91            cwd: Some(DEFAULT_EXEC_CWD.to_string()),
92            stdin: None,
93            timeout: None,
94            on_stdout: None,
95            on_stderr: None,
96            capture_stdio: None,
97            file_path: None,
98            cpu_time_limit_ms: None,
99            timing_mitigation: None,
100        }
101    }
102}
103
104/// Result of `exec`.
105#[derive(Debug, Clone, PartialEq, Eq)]
106pub struct ExecResult {
107    pub exit_code: i32,
108    pub stdout: String,
109    pub stderr: String,
110}
111
112/// `stdio` mode for a spawn.
113#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
114#[serde(rename_all = "lowercase")]
115pub enum SpawnStdio {
116    #[default]
117    Pipe,
118    Inherit,
119}
120
121/// Callback-free options for portable `spawn`.
122#[derive(Default)]
123pub struct SpawnOptions {
124    pub env: BTreeMap<String, String>,
125    pub cwd: Option<String>,
126    pub stdio: Option<SpawnStdio>,
127    pub stdin_fd: Option<i32>,
128    pub stdout_fd: Option<i32>,
129    pub stderr_fd: Option<i32>,
130    pub stream_stdin: Option<bool>,
131}
132
133#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
134#[serde(rename_all = "lowercase")]
135pub enum ProcessStream {
136    Stdout,
137    Stderr,
138}
139
140#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
141pub struct ProcessOutput {
142    pub pid: u32,
143    pub stream: ProcessStream,
144    pub data: Vec<u8>,
145}
146
147#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
148pub struct ProcessExit {
149    pub pid: u32,
150    #[serde(rename = "exitCode")]
151    pub exit_code: i32,
152}
153
154/// Public JSON info for SDK-spawned processes.
155#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
156pub struct SpawnedProcessInfo {
157    pub pid: u32,
158    pub command: String,
159    pub args: Vec<String>,
160    pub running: bool,
161    #[serde(rename = "exitCode")]
162    pub exit_code: Option<i32>,
163    /// Epoch milliseconds when `spawn` registered the process.
164    #[serde(rename = "startedAt")]
165    pub started_at: i64,
166}
167
168/// The pid returned by `spawn`.
169#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
170pub struct SpawnHandle {
171    pub pid: u32,
172}
173
174/// Process status from the kernel process table.
175#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
176#[serde(rename_all = "lowercase")]
177pub enum ProcessStatus {
178    Running,
179    Exited,
180}
181
182/// Full kernel process info (TS `KernelProcessInfo`).
183#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
184pub struct ProcessInfo {
185    pub pid: u32,
186    pub ppid: u32,
187    pub pgid: u32,
188    pub sid: u32,
189    pub driver: String,
190    pub command: String,
191    pub args: Vec<String>,
192    pub cwd: String,
193    pub status: ProcessStatus,
194    #[serde(rename = "exitCode")]
195    pub exit_code: Option<i32>,
196    #[serde(rename = "startTime")]
197    pub start_time: f64,
198    #[serde(rename = "exitTime")]
199    pub exit_time: Option<f64>,
200}
201
202/// A node in the process forest (`ProcessInfo` + children).
203#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
204pub struct ProcessTreeNode {
205    #[serde(flatten)]
206    pub info: ProcessInfo,
207    pub children: Vec<ProcessTreeNode>,
208}
209
210// ---------------------------------------------------------------------------
211// Methods
212// ---------------------------------------------------------------------------
213
214impl AgentOs {
215    /// Run a command to completion. The wire `Execute` request starts the process and returns a
216    /// process id immediately; stdout/stderr are accumulated and the call resolves once the matching
217    /// `ProcessExited` event arrives. This mirrors the TS pass-through to `kernel.exec` semantically:
218    /// the result is the full captured stdout/stderr plus exit code.
219    pub async fn exec_process(&self, command: &str, options: ExecOptions) -> Result<ExecResult> {
220        // Parse the command line into a `(command, args)` pair the same way the sidecar's
221        // child_process path does: shell-free argv lists spawn directly (preserving the command's
222        // real exit code), while shell syntax or a builtin head runs under `sh -c <line>`.
223        let (resolved_command, resolved_args) = resolve_exec_command(command)?;
224        self.exec_argv_process(&resolved_command, &resolved_args, options)
225            .await
226    }
227
228    /// Run a command to completion from an already-structured `(command, args)` argv, bypassing the
229    /// `exec` command-line parser. Each `args` element is sent verbatim as a distinct argv element —
230    /// no whitespace re-splitting, no shell metacharacter detection, and no routing through
231    /// `sh -c`. Callers that already hold a structured argv (for example the cron `Exec` action)
232    /// must use this so the structured-argv contract is preserved end to end.
233    pub async fn exec_argv_process(
234        &self,
235        command: &str,
236        args: &[String],
237        mut options: ExecOptions,
238    ) -> Result<ExecResult> {
239        let process_id = self.next_process_id();
240
241        // Subscribe to events BEFORE issuing the request so no output/exit is missed between the
242        // request landing and the subscription being installed.
243        let mut events = self.transport().subscribe_wire_events();
244
245        let resolved_command = command.to_owned();
246        let resolved_args = args.to_vec();
247        let started = self
248            .send_execute(
249                &process_id,
250                Some(resolved_command),
251                resolved_args,
252                options.env.clone(),
253                options.cwd.clone(),
254            )
255            .await
256            .context("exec: Execute request failed")?;
257        debug_assert_eq!(started.process_id, process_id);
258
259        // Deliver any provided stdin, then close stdin so a non-interactive run observes EOF. This
260        // mirrors the TS `runAndCapture` path (`proc.writeStdin(options.stdin); proc.closeStdin()`).
261        if let Some(stdin) = options.stdin.take() {
262            let chunk = stdin_to_bytes(stdin);
263            let ownership = self.vm_scope();
264            let _ = self
265                .transport()
266                .request_wire(
267                    ownership,
268                    wire::RequestPayload::WriteStdinRequest(wire::WriteStdinRequest {
269                        process_id: process_id.clone(),
270                        chunk,
271                    }),
272                )
273                .await;
274        }
275        {
276            let ownership = self.vm_scope();
277            let _ = self
278                .transport()
279                .request_wire(
280                    ownership,
281                    wire::RequestPayload::CloseStdinRequest(wire::CloseStdinRequest {
282                        process_id: process_id.clone(),
283                    }),
284                )
285                .await;
286        }
287
288        let mut on_stdout = options.on_stdout.take();
289        let mut on_stderr = options.on_stderr.take();
290
291        // A `timeout` (ms) bounds the run: when it elapses, SIGKILL the process and keep draining
292        // until the exit event lands. This mirrors the TS `runAndCapture` timeout race that kills the
293        // process and then awaits its exit code.
294        let timeout_deadline = options
295            .timeout
296            .filter(|ms| ms.is_finite() && *ms >= 0.0)
297            .map(|ms| {
298                tokio::time::Instant::now() + std::time::Duration::from_secs_f64(ms / 1000.0)
299            });
300        let mut killed_for_timeout = false;
301
302        let capture_stdio = options.capture_stdio.unwrap_or(true);
303        let mut stdout = Vec::<u8>::new();
304        let mut stderr = Vec::<u8>::new();
305        let mut captured_output_bytes = 0usize;
306        let mut capture_error: Option<ClientError> = None;
307        let exit_code = loop {
308            let recv = events.recv();
309            let frame = match timeout_deadline {
310                Some(deadline) => {
311                    tokio::select! {
312                        result = recv => result,
313                        _ = tokio::time::sleep_until(deadline), if !killed_for_timeout => {
314                            killed_for_timeout = true;
315                            self.kill_wire_process(&process_id, "SIGKILL");
316                            continue;
317                        }
318                    }
319                }
320                None => recv.await,
321            };
322            let (_, payload) = match frame {
323                Ok(frame) => frame,
324                Err(broadcast::error::RecvError::Lagged(_)) => continue,
325                Err(broadcast::error::RecvError::Closed) => {
326                    return Err(ClientError::Sidecar(
327                        "exec: event stream closed before process exit".to_owned(),
328                    )
329                    .into());
330                }
331            };
332            match payload {
333                EventPayload::ProcessOutputEvent(output) if output.process_id == process_id => {
334                    match output.channel {
335                        StreamChannel::Stdout => {
336                            if let Some(cb) = on_stdout.as_mut() {
337                                cb(&output.chunk);
338                            }
339                            if capture_stdio && capture_error.is_none() {
340                                match append_exec_output(
341                                    &mut stdout,
342                                    &output.chunk,
343                                    &mut captured_output_bytes,
344                                    "stdout",
345                                ) {
346                                    Ok(()) => {}
347                                    Err(error) => {
348                                        self.kill_wire_process(&process_id, "SIGKILL");
349                                        capture_error = Some(error);
350                                    }
351                                }
352                            }
353                        }
354                        StreamChannel::Stderr => {
355                            if let Some(cb) = on_stderr.as_mut() {
356                                cb(&output.chunk);
357                            }
358                            if capture_stdio && capture_error.is_none() {
359                                match append_exec_output(
360                                    &mut stderr,
361                                    &output.chunk,
362                                    &mut captured_output_bytes,
363                                    "stderr",
364                                ) {
365                                    Ok(()) => {}
366                                    Err(error) => {
367                                        self.kill_wire_process(&process_id, "SIGKILL");
368                                        capture_error = Some(error);
369                                    }
370                                }
371                            }
372                        }
373                    }
374                }
375                EventPayload::ProcessExitedEvent(exited) if exited.process_id == process_id => {
376                    break exited.exit_code;
377                }
378                EventPayload::ProcessOutputEvent(_)
379                | EventPayload::ProcessExitedEvent(_)
380                | EventPayload::ExecutionOutputEvent(_)
381                | EventPayload::ExecutionCompletedEvent(_)
382                | EventPayload::VmLifecycleEvent(_)
383                | EventPayload::StructuredEvent(_)
384                | EventPayload::ExtEnvelope(_) => {}
385            }
386        };
387
388        if let Some(error) = capture_error {
389            return Err(error.into());
390        }
391
392        Ok(ExecResult {
393            exit_code,
394            stdout: String::from_utf8_lossy(&stdout).into_owned(),
395            stderr: String::from_utf8_lossy(&stderr).into_owned(),
396        })
397    }
398
399    /// Spawn a process. SYNC; returns `{ pid }` only. Installs stdout/stderr fan-out over broadcast
400    /// channels and wires exit via a background event-pump task. The user-facing `pid` is the
401    /// SDK-allocated map key (the wire `process_id` is held inside the [`ProcessEntry`]).
402    pub fn spawn_process(
403        &self,
404        command: &str,
405        args: Vec<String>,
406        options: SpawnOptions,
407    ) -> Result<SpawnHandle> {
408        let registry_guard = self.inner().process_registry_lock.lock();
409        self.prune_exited_processes_locked(1);
410        if self.process_registry_len_locked() >= PROCESS_REGISTRY_LIMIT {
411            return Err(ClientError::Sidecar(format!(
412                "process registry limit exceeded: at most {PROCESS_REGISTRY_LIMIT} processes can be tracked per VM"
413            ))
414            .into());
415        }
416
417        // Draw the public pid from the dedicated synthetic-pid space (TS `nextSyntheticPid`), seeded
418        // at `SYNTHETIC_PID_BASE`. `exec` uses a separate counter so it never perturbs this sequence.
419        let pid = self
420            .inner()
421            .synthetic_pid_counter
422            .fetch_add(1, Ordering::SeqCst) as u32;
423        let process_id = format!("proc-{pid}-{}", uuid::Uuid::new_v4());
424
425        let (stdout_tx, _) = broadcast::channel::<Vec<u8>>(PROCESS_STREAM_CAPACITY);
426        let (stderr_tx, _) = broadcast::channel::<Vec<u8>>(PROCESS_STREAM_CAPACITY);
427        let (output_tx, _) = broadcast::channel::<ProcessOutput>(PROCESS_STREAM_CAPACITY);
428        // Seeded `None`; the already-exited branch of `on_process_exit` fires immediately once this
429        // watch holds `Some(code)`.
430        let (exit_tx, _) = watch::channel::<Option<i32>>(None);
431        // Seeded `None`; filled with the kernel pid once the `Execute` response lands so
432        // `all_processes`/`process_tree` can remap the kernel snapshot back to this display pid.
433        let (kernel_pid_tx, _) = watch::channel::<Option<u32>>(None);
434
435        let entry = ProcessEntry {
436            command: command.to_owned(),
437            args: args.clone(),
438            stdout_tx: stdout_tx.clone(),
439            stderr_tx: stderr_tx.clone(),
440            output_tx: output_tx.clone(),
441            exit_tx: exit_tx.clone(),
442            process_id: process_id.clone(),
443            kernel_pid: kernel_pid_tx.clone(),
444            output_tasks: Vec::new(),
445            started_at: epoch_ms_now() as i64,
446        };
447        // `spawn` is documented as overwriting any prior entry for a freshly allocated pid; the pid
448        // is monotonic so a collision is not expected.
449        let _ = self.inner().processes.insert(pid, entry);
450        drop(registry_guard);
451
452        // Subscribe to events before issuing the request so the pump sees everything.
453        let events = self.transport().subscribe_wire_events();
454
455        let this = self.clone();
456        let command = command.to_owned();
457        tokio::spawn(async move {
458            this.run_spawn(
459                pid,
460                process_id,
461                command,
462                args,
463                options,
464                events,
465                stdout_tx,
466                stderr_tx,
467                output_tx,
468                exit_tx,
469                kernel_pid_tx,
470            )
471            .await;
472        });
473
474        Ok(SpawnHandle { pid })
475    }
476
477    /// Write to a spawned process's stdin. SYNC. Errors with `ProcessNotFound`.
478    pub fn write_process_stdin(
479        &self,
480        pid: u32,
481        data: StdinInput,
482    ) -> std::result::Result<(), ClientError> {
483        let process_id = self.lookup_process_id(pid)?;
484        let chunk: Vec<u8> = stdin_to_bytes(data);
485        let this = self.clone();
486        // Fire-and-forget: the TS API is synchronous and does not surface a write error.
487        tokio::spawn(async move {
488            let ownership = this.vm_scope();
489            let _ = this
490                .transport()
491                .request_wire(
492                    ownership,
493                    wire::RequestPayload::WriteStdinRequest(wire::WriteStdinRequest {
494                        process_id,
495                        chunk,
496                    }),
497                )
498                .await;
499        });
500        Ok(())
501    }
502
503    /// Close a spawned process's stdin. SYNC. Errors with `ProcessNotFound`.
504    pub fn close_process_stdin(&self, pid: u32) -> std::result::Result<(), ClientError> {
505        let process_id = self.lookup_process_id(pid)?;
506        let this = self.clone();
507        tokio::spawn(async move {
508            let ownership = this.vm_scope();
509            let _ = this
510                .transport()
511                .request_wire(
512                    ownership,
513                    wire::RequestPayload::CloseStdinRequest(wire::CloseStdinRequest { process_id }),
514                )
515                .await;
516        });
517        Ok(())
518    }
519
520    /// Subscribe to the unified stdout/stderr event stream for a process.
521    pub fn on_process_output(
522        &self,
523        pid: u32,
524        mut handler: impl FnMut(ProcessOutput) + Send + 'static,
525    ) -> std::result::Result<Subscription, ClientError> {
526        let mut rx = self
527            .inner()
528            .processes
529            .read(&pid, |_, entry| entry.output_tx.subscribe())
530            .ok_or(ClientError::ProcessNotFound(pid))?;
531        let task = tokio::spawn(async move {
532            loop {
533                match rx.recv().await {
534                    Ok(event) => handler(event),
535                    Err(broadcast::error::RecvError::Lagged(_)) => continue,
536                    Err(broadcast::error::RecvError::Closed) => return,
537                }
538            }
539        });
540        Ok(Subscription::new(move || task.abort()))
541    }
542
543    /// Register a once-only exit handler. If the process has already exited, the handler fires
544    /// immediately and synchronously and a no-op unsubscribe is returned (the `watch` already holds
545    /// `Some(code)`). Otherwise the handler fires once when the exit code lands. The exit code is
546    /// `i32`, never null.
547    pub fn on_process_exit(
548        &self,
549        pid: u32,
550        handler: impl FnOnce(ProcessExit) + Send + 'static,
551    ) -> std::result::Result<Subscription, ClientError> {
552        let mut rx = self
553            .inner()
554            .processes
555            .read(&pid, |_, entry| entry.exit_tx.subscribe())
556            .ok_or(ClientError::ProcessNotFound(pid))?;
557
558        // Already-exited branch: fire immediately + synchronously, return a no-op unsubscribe.
559        if let Some(code) = *rx.borrow() {
560            handler(ProcessExit {
561                pid,
562                exit_code: code,
563            });
564            return Ok(Subscription::noop());
565        }
566
567        // Otherwise wait for the watch to transition to `Some(code)` and fire exactly once. The
568        // returned `Subscription` cancels the waiting task on drop (= unsubscribe).
569        let task = tokio::spawn(async move {
570            while rx.changed().await.is_ok() {
571                if let Some(code) = *rx.borrow() {
572                    handler(ProcessExit {
573                        pid,
574                        exit_code: code,
575                    });
576                    return;
577                }
578            }
579        });
580        Ok(Subscription::new(move || task.abort()))
581    }
582
583    /// Await a spawned process's exit code. Unknown-pid lookup errors (synchronously in TS; here the
584    /// lookup error is returned before any awaiting begins).
585    pub async fn wait_process(&self, pid: u32) -> std::result::Result<i32, ClientError> {
586        let mut rx = self
587            .inner()
588            .processes
589            .read(&pid, |_, entry| entry.exit_tx.subscribe())
590            .ok_or(ClientError::ProcessNotFound(pid))?;
591
592        if let Some(code) = *rx.borrow() {
593            return Ok(code);
594        }
595        while rx.changed().await.is_ok() {
596            if let Some(code) = *rx.borrow() {
597                return Ok(code);
598            }
599        }
600        Err(ClientError::Sidecar(format!(
601            "wait_process: exit channel closed before process {pid} reported an exit code"
602        )))
603    }
604
605    /// List SDK-spawned processes only. `running = exit_code.is_none()`.
606    pub fn list_processes(&self) -> Vec<SpawnedProcessInfo> {
607        let mut out = Vec::new();
608        self.inner().processes.scan(|pid, entry| {
609            let exit_code = *entry.exit_tx.borrow();
610            out.push(SpawnedProcessInfo {
611                pid: *pid,
612                command: entry.command.clone(),
613                args: entry.args.clone(),
614                running: exit_code.is_none(),
615                exit_code,
616                started_at: entry.started_at,
617            });
618        });
619        out
620    }
621
622    /// List ALL kernel processes (native sidecar process snapshot).
623    ///
624    /// The kernel snapshot keys processes by their raw kernel pid. SDK-spawned root processes carry a
625    /// synthetic display pid (the `spawn` return value); this remaps each snapshot entry's
626    /// pid/ppid/pgid/sid back to that display pid via the per-process `kernel_pid` watch, so a caller
627    /// can correlate `spawn()` with `all_processes()`/`process_tree()`. Results are sorted ascending
628    /// by display pid (TS `snapshotProcesses` `.sort((l,r) => l.pid - r.pid)`).
629    pub async fn all_processes(&self) -> Result<Vec<ProcessInfo>> {
630        let ownership = self.vm_scope();
631        let response = self
632            .transport()
633            .request_wire(ownership, wire::RequestPayload::GetProcessSnapshotRequest)
634            .await
635            .context("all_processes: GetProcessSnapshot request failed")?;
636        let snapshot = match response {
637            wire::ResponsePayload::ProcessSnapshotResponse(snapshot) => snapshot,
638            wire::ResponsePayload::RejectedResponse(rejected) => {
639                return Err(ClientError::from_rejection(rejected).into());
640            }
641            other => {
642                return Err(ClientError::Sidecar(format!(
643                    "all_processes: unexpected response {other:?}"
644                ))
645                .into());
646            }
647        };
648
649        // Snapshot the SDK process registry, keyed by wire `process_id`, capturing exit code,
650        // command, and args. This mirrors the TS `trackedProcessesById` lookup used to build
651        // `displayPidByKernelPid` and override fields.
652        struct Tracked {
653            exit_code: Option<i32>,
654            command: String,
655            args: Vec<String>,
656        }
657        let mut tracked_by_process_id: BTreeMap<String, Tracked> = BTreeMap::new();
658        let mut display_pid_by_kernel_pid: BTreeMap<u32, u32> = BTreeMap::new();
659        self.inner().processes.scan(|display_pid, entry| {
660            let exit_code = *entry.exit_tx.borrow();
661            if let Some(kernel_pid) = *entry.kernel_pid.borrow() {
662                display_pid_by_kernel_pid.insert(kernel_pid, *display_pid);
663            }
664            tracked_by_process_id.insert(
665                entry.process_id.clone(),
666                Tracked {
667                    exit_code,
668                    command: entry.command.clone(),
669                    args: entry.args.clone(),
670                },
671            );
672        });
673
674        let now_ms = epoch_ms_now();
675        let mut seen_display_pids: std::collections::BTreeSet<u32> =
676            std::collections::BTreeSet::new();
677        let mut out: Vec<ProcessInfo> = Vec::new();
678
679        for entry in snapshot.processes {
680            let tracked = tracked_by_process_id.get(&entry.process_id);
681            let display_pid = display_pid_by_kernel_pid
682                .get(&entry.pid)
683                .copied()
684                .unwrap_or(entry.pid);
685            let display_ppid = display_pid_by_kernel_pid
686                .get(&entry.ppid)
687                .copied()
688                .unwrap_or(entry.ppid);
689            let display_pgid = display_pid_by_kernel_pid
690                .get(&entry.pgid)
691                .copied()
692                .unwrap_or(entry.pgid);
693            let display_sid = display_pid_by_kernel_pid
694                .get(&entry.sid)
695                .copied()
696                .unwrap_or(entry.sid);
697
698            // First-observed start time, keyed by `"<process_id>:<kernel_pid>"` (TS `processKey`).
699            let process_key = format!("{}:{}", entry.process_id, entry.pid);
700            let start_time = self.observed_start_time(&process_key, now_ms);
701
702            // Status/exit code: a tracked process whose SDK exit code is known is `exited`; otherwise
703            // a tracked process is `running`; an untracked process uses the snapshot status.
704            let (status, exit_code) = match tracked {
705                Some(t) => match t.exit_code {
706                    Some(code) => (ProcessStatus::Exited, Some(code)),
707                    None => (ProcessStatus::Running, entry.exit_code),
708                },
709                None => {
710                    let status = match entry.status {
711                        ProcessSnapshotStatus::Running | ProcessSnapshotStatus::Stopped => {
712                            ProcessStatus::Running
713                        }
714                        ProcessSnapshotStatus::Exited => ProcessStatus::Exited,
715                    };
716                    (status, entry.exit_code)
717                }
718            };
719
720            // Exit time: only tracked-and-exited processes carry one (TS `tracked?.exitTime`).
721            let exit_time = match (tracked, status) {
722                (Some(_), ProcessStatus::Exited) => {
723                    Some(self.observed_exit_time(&entry.process_id, now_ms))
724                }
725                _ => None,
726            };
727
728            let (command, args) = match tracked {
729                Some(t) => (t.command.clone(), t.args.clone()),
730                None => (entry.command, entry.args),
731            };
732
733            seen_display_pids.insert(display_pid);
734            out.push(ProcessInfo {
735                pid: display_pid,
736                ppid: display_ppid,
737                pgid: display_pgid,
738                sid: display_sid,
739                driver: entry.driver,
740                command,
741                args,
742                cwd: entry.cwd,
743                status,
744                exit_code,
745                start_time,
746                exit_time,
747            });
748        }
749
750        // Tracked processes not yet present in the snapshot (the spawn `Execute` has not surfaced in
751        // the kernel table yet). TS fills these with `ppid:0, pgid/sid = pid`.
752        self.inner().processes.scan(|display_pid, entry| {
753            if seen_display_pids.contains(display_pid) {
754                return;
755            }
756            let exit_code = *entry.exit_tx.borrow();
757            let process_key = format!("{}:{}", entry.process_id, display_pid);
758            let start_time = self.observed_start_time(&process_key, now_ms);
759            let (status, exit_time) = match exit_code {
760                Some(_) => (
761                    ProcessStatus::Exited,
762                    Some(self.observed_exit_time(&entry.process_id, now_ms)),
763                ),
764                None => (ProcessStatus::Running, None),
765            };
766            out.push(ProcessInfo {
767                pid: *display_pid,
768                ppid: 0,
769                pgid: *display_pid,
770                sid: *display_pid,
771                driver: String::new(),
772                command: entry.command.clone(),
773                args: entry.args.clone(),
774                cwd: String::new(),
775                status,
776                exit_code,
777                start_time,
778                exit_time,
779            });
780        });
781
782        out.sort_by_key(|info| info.pid);
783        Ok(out)
784    }
785
786    /// Return the first-observed start time for a process key, recording `now` the first time it is
787    /// seen so later snapshots report a stable timestamp (TS `observedProcessStartTimes`).
788    fn observed_start_time(&self, process_key: &str, now_ms: f64) -> f64 {
789        let _guard = self.inner().observed_process_time_lock.lock();
790        if let Some(existing) = self
791            .inner()
792            .observed_process_start_times
793            .read(process_key, |_, value| *value)
794        {
795            return existing;
796        }
797        let _ = self
798            .inner()
799            .observed_process_start_times
800            .insert(process_key.to_owned(), now_ms);
801        prune_string_f64_map(
802            &self.inner().observed_process_start_times,
803            OBSERVED_PROCESS_TIME_LIMIT,
804        );
805        // Re-read to honor a racing insert that may have won; either value is a valid first-observed
806        // timestamp.
807        self.inner()
808            .observed_process_start_times
809            .read(process_key, |_, value| *value)
810            .unwrap_or(now_ms)
811    }
812
813    /// Return the first-observed exit time for an SDK process id, recording `now` on first sight.
814    fn observed_exit_time(&self, process_id: &str, now_ms: f64) -> f64 {
815        let _guard = self.inner().observed_process_time_lock.lock();
816        if let Some(existing) = self
817            .inner()
818            .observed_process_exit_times
819            .read(process_id, |_, value| *value)
820        {
821            return existing;
822        }
823        let _ = self
824            .inner()
825            .observed_process_exit_times
826            .insert(process_id.to_owned(), now_ms);
827        prune_string_f64_map(
828            &self.inner().observed_process_exit_times,
829            OBSERVED_PROCESS_TIME_LIMIT,
830        );
831        self.inner()
832            .observed_process_exit_times
833            .read(process_id, |_, value| *value)
834            .unwrap_or(now_ms)
835    }
836
837    /// Build the process forest from `all_processes`, linked by `ppid`.
838    pub async fn process_tree(&self) -> Result<Vec<ProcessTreeNode>> {
839        let processes = self.all_processes().await?;
840        Ok(build_process_forest(processes))
841    }
842
843    /// Get a single SDK-spawned process's info. Errors (not None) when not found.
844    pub fn get_process(&self, pid: u32) -> std::result::Result<SpawnedProcessInfo, ClientError> {
845        self.inner()
846            .processes
847            .read(&pid, |pid, entry| {
848                let exit_code = *entry.exit_tx.borrow();
849                SpawnedProcessInfo {
850                    pid: *pid,
851                    command: entry.command.clone(),
852                    args: entry.args.clone(),
853                    running: exit_code.is_none(),
854                    exit_code,
855                    started_at: entry.started_at,
856                }
857            })
858            .ok_or(ClientError::ProcessNotFound(pid))
859    }
860
861    /// SIGTERM a spawned process. No-op if already exited; errors if unknown.
862    pub fn stop_process(&self, pid: u32) -> std::result::Result<(), ClientError> {
863        self.signal_process(pid, "SIGTERM")
864    }
865
866    /// SIGKILL a spawned process. No-op if already exited; errors if unknown.
867    pub fn kill_process(&self, pid: u32) -> std::result::Result<(), ClientError> {
868        self.signal_process(pid, "SIGKILL")
869    }
870
871    // -----------------------------------------------------------------------
872    // Internal helpers
873    // -----------------------------------------------------------------------
874
875    /// Build the VM-scoped ownership for a wire request.
876    fn vm_scope(&self) -> wire::OwnershipScope {
877        wire::OwnershipScope::VmOwnership(wire::VmOwnership {
878            connection_id: self.connection_id().to_string(),
879            session_id: self.wire_session_id().to_string(),
880            vm_id: self.vm_id().to_string(),
881        })
882    }
883
884    /// Allocate a fresh wire `process_id` (used by `exec`, which does not register in the SDK map).
885    fn next_process_id(&self) -> String {
886        let n = self.inner().process_counter.fetch_add(1, Ordering::SeqCst);
887        format!("proc-{n}-{}", uuid::Uuid::new_v4())
888    }
889
890    /// Resolve the wire `process_id` for an SDK pid, erroring with `ProcessNotFound` if unknown.
891    fn lookup_process_id(&self, pid: u32) -> std::result::Result<String, ClientError> {
892        self.inner()
893            .processes
894            .read(&pid, |_, entry| entry.process_id.clone())
895            .ok_or(ClientError::ProcessNotFound(pid))
896    }
897
898    /// Send the `Execute` wire request, mapping a rejection into [`ClientError::Kernel`].
899    async fn send_execute(
900        &self,
901        process_id: &str,
902        command: Option<String>,
903        args: Vec<String>,
904        env: BTreeMap<String, String>,
905        cwd: Option<String>,
906    ) -> std::result::Result<wire::ProcessStartedResponse, ClientError> {
907        let ownership = self.vm_scope();
908        let response = self
909            .transport()
910            .request_wire(
911                ownership,
912                wire::RequestPayload::ExecuteRequest(wire::ExecuteRequest {
913                    process_id: process_id.to_owned(),
914                    command,
915                    runtime: None,
916                    entrypoint: None,
917                    args,
918                    env: env.into_iter().collect(),
919                    cwd,
920                    wasm_permission_tier: None,
921                }),
922            )
923            .await?;
924        match response {
925            wire::ResponsePayload::ProcessStartedResponse(started) => Ok(started),
926            wire::ResponsePayload::RejectedResponse(rejected) => {
927                Err(ClientError::from_rejection(rejected))
928            }
929            other => Err(ClientError::Sidecar(format!(
930                "Execute: unexpected response {other:?}"
931            ))),
932        }
933    }
934
935    /// Fire-and-forget kill of a wire process by its `process_id` (used by `exec` timeout). The TS
936    /// timeout path calls `proc.kill(9)`, which maps to a `SIGKILL` kill request.
937    fn kill_wire_process(&self, process_id: &str, signal: &str) {
938        let process_id = process_id.to_owned();
939        let signal = signal.to_owned();
940        let this = self.clone();
941        tokio::spawn(async move {
942            let ownership = this.vm_scope();
943            let _ = this
944                .transport()
945                .request_wire(
946                    ownership,
947                    wire::RequestPayload::KillProcessRequest(wire::KillProcessRequest {
948                        process_id,
949                        signal,
950                    }),
951                )
952                .await;
953        });
954    }
955
956    /// Send a kill signal for an SDK pid. No-op if already exited; errors with `ProcessNotFound` if
957    /// the pid is unknown.
958    fn signal_process(&self, pid: u32, signal: &str) -> std::result::Result<(), ClientError> {
959        let (process_id, already_exited) = self
960            .inner()
961            .processes
962            .read(&pid, |_, entry| {
963                (entry.process_id.clone(), entry.exit_tx.borrow().is_some())
964            })
965            .ok_or(ClientError::ProcessNotFound(pid))?;
966        if already_exited {
967            return Ok(());
968        }
969        let signal = signal.to_owned();
970        let this = self.clone();
971        tokio::spawn(async move {
972            let ownership = this.vm_scope();
973            let _ = this
974                .transport()
975                .request_wire(
976                    ownership,
977                    wire::RequestPayload::KillProcessRequest(wire::KillProcessRequest {
978                        process_id,
979                        signal,
980                    }),
981                )
982                .await;
983        });
984        Ok(())
985    }
986
987    fn process_registry_len_locked(&self) -> usize {
988        let mut count = 0usize;
989        self.inner().processes.scan(|_, _| {
990            count += 1;
991        });
992        count
993    }
994
995    fn prune_exited_processes_locked(&self, reserve_slots: usize) {
996        let mut entries = Vec::new();
997        self.inner().processes.scan(|pid, entry| {
998            entries.push((*pid, entry.exit_tx.borrow().is_some()));
999        });
1000        let target_len = PROCESS_REGISTRY_LIMIT.saturating_sub(reserve_slots);
1001        if entries.len() <= target_len {
1002            return;
1003        }
1004
1005        for pid in exited_pids_to_prune(entries, target_len) {
1006            self.remove_process_tracking_locked(pid);
1007        }
1008    }
1009
1010    fn remove_process_tracking_locked(&self, pid: u32) {
1011        if let Some((_, entry)) = self.inner().processes.remove(&pid) {
1012            let _time_guard = self.inner().observed_process_time_lock.lock();
1013            let _ = self
1014                .inner()
1015                .observed_process_exit_times
1016                .remove(&entry.process_id);
1017            let fallback_start_key = format!("{}:{pid}", entry.process_id);
1018            let _ = self
1019                .inner()
1020                .observed_process_start_times
1021                .remove(&fallback_start_key);
1022            if let Some(kernel_pid) = *entry.kernel_pid.borrow() {
1023                let start_key = format!("{}:{kernel_pid}", entry.process_id);
1024                let _ = self.inner().observed_process_start_times.remove(&start_key);
1025            }
1026        }
1027    }
1028
1029    /// Background pump for a spawned process: issue the `Execute` request, then fan kernel
1030    /// `ProcessOutput`/`ProcessExited` events for this process id into the per-process broadcast and
1031    /// watch channels. Exited entries are retained for post-exit inspection, then pruned oldest-first
1032    /// under registry pressure.
1033    #[allow(clippy::too_many_arguments)]
1034    async fn run_spawn(
1035        self,
1036        pid: u32,
1037        process_id: String,
1038        command: String,
1039        args: Vec<String>,
1040        options: SpawnOptions,
1041        mut events: broadcast::Receiver<(wire::OwnershipScope, EventPayload)>,
1042        stdout_tx: broadcast::Sender<Vec<u8>>,
1043        stderr_tx: broadcast::Sender<Vec<u8>>,
1044        output_tx: broadcast::Sender<ProcessOutput>,
1045        exit_tx: watch::Sender<Option<i32>>,
1046        kernel_pid_tx: watch::Sender<Option<u32>>,
1047    ) {
1048        match self
1049            .send_execute(
1050                &process_id,
1051                Some(command),
1052                args,
1053                options.env.clone(),
1054                options.cwd.clone(),
1055            )
1056            .await
1057        {
1058            Ok(started) => {
1059                // Seed the kernel pid so `all_processes`/`process_tree` can remap this process's
1060                // kernel-snapshot entry back to its display pid.
1061                if let Some(kernel_pid) = started.pid {
1062                    let _ = kernel_pid_tx.send(Some(kernel_pid));
1063                }
1064            }
1065            Err(error) => {
1066                // The native TS launch-failure path emits the error message (plus a trailing
1067                // newline) on stderr and resolves the wait with exit code 1 (`startTrackedProcess`
1068                // catch -> stderr handlers + `finishProcess(entry, 1)`).
1069                let message = format!("{error}\n");
1070                let bytes = message.into_bytes();
1071                let _ = stderr_tx.send(bytes.clone());
1072                let _ = output_tx.send(ProcessOutput {
1073                    pid,
1074                    stream: ProcessStream::Stderr,
1075                    data: bytes,
1076                });
1077                tracing::error!(?error, pid, %process_id, "spawn: Execute request failed");
1078                let _ = exit_tx.send(Some(1));
1079                let _guard = self.inner().process_registry_lock.lock();
1080                self.prune_exited_processes_locked(0);
1081                return;
1082            }
1083        }
1084
1085        loop {
1086            let (_, payload) = match events.recv().await {
1087                Ok(frame) => frame,
1088                Err(broadcast::error::RecvError::Lagged(_)) => continue,
1089                Err(broadcast::error::RecvError::Closed) => {
1090                    // The event stream closed before an exit event landed. The TS fallback treats a
1091                    // process that has fully disappeared from the VM snapshot as reaped with exit
1092                    // code 0; mirror that terminal value so waiters resolve instead of hanging.
1093                    let _ = exit_tx.send(Some(0));
1094                    break;
1095                }
1096            };
1097            match payload {
1098                EventPayload::ProcessOutputEvent(output) if output.process_id == process_id => {
1099                    let bytes = output.chunk;
1100                    let _ = output_tx.send(ProcessOutput {
1101                        pid,
1102                        stream: match output.channel {
1103                            StreamChannel::Stdout => ProcessStream::Stdout,
1104                            StreamChannel::Stderr => ProcessStream::Stderr,
1105                        },
1106                        data: bytes.clone(),
1107                    });
1108                    match output.channel {
1109                        StreamChannel::Stdout => {
1110                            let _ = stdout_tx.send(bytes);
1111                        }
1112                        StreamChannel::Stderr => {
1113                            let _ = stderr_tx.send(bytes);
1114                        }
1115                    }
1116                }
1117                EventPayload::ProcessExitedEvent(exited) if exited.process_id == process_id => {
1118                    let _ = exit_tx.send(Some(exited.exit_code));
1119                    break;
1120                }
1121                EventPayload::ProcessOutputEvent(_)
1122                | EventPayload::ProcessExitedEvent(_)
1123                | EventPayload::ExecutionOutputEvent(_)
1124                | EventPayload::ExecutionCompletedEvent(_)
1125                | EventPayload::VmLifecycleEvent(_)
1126                | EventPayload::StructuredEvent(_)
1127                | EventPayload::ExtEnvelope(_) => {}
1128            }
1129        }
1130        let _guard = self.inner().process_registry_lock.lock();
1131        self.prune_exited_processes_locked(0);
1132    }
1133}
1134
1135/// Assemble a process forest from a flat process list, linking children by `ppid`.
1136///
1137/// Mirrors the TS `processTree` `nodeMap` algorithm exactly: a process is a root iff its `ppid` is
1138/// NOT present among the listed pids. A self-parented process (`ppid == pid`) finds itself as its
1139/// parent, so it is attached as its own child and is excluded from the roots (effectively dropped
1140/// from the output tree). A `seen` guard prevents the self-cycle from recursing forever.
1141fn build_process_forest(processes: Vec<ProcessInfo>) -> Vec<ProcessTreeNode> {
1142    use std::collections::BTreeMap as Map;
1143
1144    let pids: std::collections::BTreeSet<u32> = processes.iter().map(|p| p.pid).collect();
1145    // Children adjacency keyed by parent pid, preserving input (sorted) order.
1146    let mut children_of: Map<u32, Vec<usize>> = Map::new();
1147    let mut roots: Vec<usize> = Vec::new();
1148    for (index, proc) in processes.iter().enumerate() {
1149        if pids.contains(&proc.ppid) {
1150            children_of.entry(proc.ppid).or_default().push(index);
1151        } else {
1152            roots.push(index);
1153        }
1154    }
1155
1156    fn build_node(
1157        index: usize,
1158        processes: &[ProcessInfo],
1159        children_of: &Map<u32, Vec<usize>>,
1160        seen: &mut std::collections::BTreeSet<usize>,
1161    ) -> ProcessTreeNode {
1162        let info = processes[index].clone();
1163        seen.insert(index);
1164        let child_indices: Vec<usize> = children_of
1165            .get(&info.pid)
1166            .map(|indices| {
1167                indices
1168                    .iter()
1169                    .copied()
1170                    .filter(|child_index| !seen.contains(child_index))
1171                    .collect()
1172            })
1173            .unwrap_or_default();
1174        let children = child_indices
1175            .into_iter()
1176            .map(|child_index| build_node(child_index, processes, children_of, seen))
1177            .collect();
1178        ProcessTreeNode { info, children }
1179    }
1180
1181    let mut seen = std::collections::BTreeSet::new();
1182    roots
1183        .into_iter()
1184        .map(|index| build_node(index, &processes, &children_of, &mut seen))
1185        .collect()
1186}
1187
1188/// Convert a [`StdinInput`] to raw bytes. A string is delivered as its UTF-8 bytes; raw bytes are
1189/// delivered verbatim (binary-safe, never lossy).
1190fn stdin_to_bytes(input: StdinInput) -> Vec<u8> {
1191    match input {
1192        StdinInput::Text(text) => text.into_bytes(),
1193        StdinInput::Bytes(bytes) => bytes,
1194    }
1195}
1196
1197fn append_exec_output(
1198    buffer: &mut Vec<u8>,
1199    chunk: &[u8],
1200    captured_output_bytes: &mut usize,
1201    channel: &str,
1202) -> std::result::Result<(), ClientError> {
1203    let next_total = captured_output_bytes
1204        .checked_add(chunk.len())
1205        .ok_or_else(|| exec_output_limit_error(channel, usize::MAX))?;
1206    if next_total > EXEC_OUTPUT_CAPTURE_LIMIT_BYTES {
1207        return Err(exec_output_limit_error(channel, next_total));
1208    }
1209    buffer.extend_from_slice(chunk);
1210    *captured_output_bytes = next_total;
1211    Ok(())
1212}
1213
1214fn exec_output_limit_error(channel: &str, size: usize) -> ClientError {
1215    ClientError::Sidecar(format!(
1216        "exec {channel} capture is {size} bytes, limit is {EXEC_OUTPUT_CAPTURE_LIMIT_BYTES}"
1217    ))
1218}
1219
1220fn exited_pids_to_prune(mut entries: Vec<(u32, bool)>, target_len: usize) -> Vec<u32> {
1221    if entries.len() <= target_len {
1222        return Vec::new();
1223    }
1224    let mut remove_count = entries.len() - target_len;
1225    entries.sort_by_key(|(pid, _)| *pid);
1226    let mut out = Vec::new();
1227    for (pid, exited) in entries {
1228        if remove_count == 0 {
1229            break;
1230        }
1231        if !exited {
1232            continue;
1233        }
1234        out.push(pid);
1235        remove_count -= 1;
1236    }
1237    out
1238}
1239
1240fn prune_string_f64_map(map: &SccHashMap<String, f64>, limit: usize) {
1241    let mut keys = Vec::new();
1242    map.scan(|key, _| {
1243        keys.push(key.clone());
1244    });
1245    if keys.len() <= limit {
1246        return;
1247    }
1248    let remove_count = keys.len() - limit;
1249    keys.sort();
1250    for key in keys.into_iter().take(remove_count) {
1251        let _ = map.remove(&key);
1252    }
1253}
1254
1255/// Drive a caller-supplied output callback from a fresh subscription on the given broadcast channel.
1256/// Each chunk delivered to the channel is forwarded to `callback` as raw bytes. The task ends when
1257/// the channel closes (process exit), matching the TS handler-set lifetime.
1258///
1259/// Returns the spawned task's handle so the owner can abort it on teardown: a [`ProcessEntry`]
1260/// retains its own `stdout_tx`/`stderr_tx` clone for late subscribers, so the broadcast channel
1261/// never closes (and this task never observes `Closed`) until the entry is dropped. `shutdown`
1262/// drains the registry and aborts these handles rather than waiting on the channel close.
1263pub(crate) fn install_output_callback(
1264    tx: broadcast::Sender<Vec<u8>>,
1265    mut callback: OutputCallback,
1266) -> JoinHandle<()> {
1267    let mut rx = tx.subscribe();
1268    tokio::spawn(async move {
1269        loop {
1270            match rx.recv().await {
1271                Ok(chunk) => callback(&chunk),
1272                Err(broadcast::error::RecvError::Lagged(_)) => continue,
1273                Err(broadcast::error::RecvError::Closed) => break,
1274            }
1275        }
1276    })
1277}
1278
1279/// Drain the SDK-spawned process registry, dropping each entry's retained sender clones and aborting
1280/// its per-process output-callback tasks. Called from `shutdown` so the output tasks (which would
1281/// otherwise await a `Closed` that never fires, see [`install_output_callback`]) cannot outlive the
1282/// disposed VM. Mirrors the `pending_shell_exits` / ACP-terminal drain in `shutdown`.
1283pub(crate) fn drain_process_output_tasks(processes: &SccHashMap<u32, ProcessEntry>) {
1284    let mut tasks = Vec::new();
1285    processes.retain(|_, entry| {
1286        tasks.append(&mut entry.output_tasks);
1287        false
1288    });
1289    for task in tasks {
1290        task.abort();
1291    }
1292}
1293
1294/// Current wall-clock time as epoch milliseconds (TS `Date.now()`).
1295fn epoch_ms_now() -> f64 {
1296    use std::time::{SystemTime, UNIX_EPOCH};
1297    SystemTime::now()
1298        .duration_since(UNIX_EPOCH)
1299        .map(|d| d.as_secs_f64() * 1000.0)
1300        .unwrap_or(0.0)
1301}
1302
1303#[cfg(test)]
1304mod tests {
1305    use super::{
1306        append_exec_output, drain_process_output_tasks, exited_pids_to_prune,
1307        install_output_callback, prune_string_f64_map, ExecOptions, OutputCallback,
1308        DEFAULT_EXEC_CWD, EXEC_OUTPUT_CAPTURE_LIMIT_BYTES,
1309    };
1310    use crate::agent_os::ProcessEntry;
1311    use scc::HashMap as SccHashMap;
1312    use tokio::sync::{broadcast, watch};
1313
1314    /// Regression for the per-process output-callback leak (H3): a `ProcessEntry` retains clones of
1315    /// its `stdout_tx`/`stderr_tx`, so the output tasks never observe the broadcast `Closed` and hang
1316    /// forever unless teardown aborts them. `drain_process_output_tasks` must empty the registry and
1317    /// abort every retained output task.
1318    #[tokio::test]
1319    async fn drain_process_output_tasks_clears_registry_and_aborts_tasks() {
1320        let processes: SccHashMap<u32, ProcessEntry> = SccHashMap::new();
1321
1322        let (stdout_tx, _) = broadcast::channel::<Vec<u8>>(8);
1323        let (stderr_tx, _) = broadcast::channel::<Vec<u8>>(8);
1324        let (output_tx, _) = broadcast::channel(8);
1325        let (exit_tx, _) = watch::channel::<Option<i32>>(None);
1326        let (kernel_pid_tx, _) = watch::channel::<Option<u32>>(None);
1327
1328        // A task that never completes on its own, standing in for an output-callback task that is
1329        // waiting on a `Closed` that the retained sender clone prevents.
1330        let task = tokio::spawn(async {
1331            loop {
1332                tokio::time::sleep(std::time::Duration::from_secs(3600)).await;
1333            }
1334        });
1335        let abort_handle = task.abort_handle();
1336
1337        let entry = ProcessEntry {
1338            command: "sleep".to_string(),
1339            args: vec!["3600".to_string()],
1340            stdout_tx,
1341            stderr_tx,
1342            output_tx,
1343            exit_tx,
1344            process_id: "proc-test".to_string(),
1345            kernel_pid: kernel_pid_tx,
1346            output_tasks: vec![task],
1347            started_at: 0,
1348        };
1349        let _ = processes.insert(1, entry);
1350
1351        assert!(!abort_handle.is_finished(), "task should start alive");
1352
1353        drain_process_output_tasks(&processes);
1354
1355        assert!(processes.is_empty(), "registry must be cleared on drain");
1356
1357        // The abort is asynchronous; give the runtime a bounded window to reap the cancelled task.
1358        for _ in 0..100 {
1359            if abort_handle.is_finished() {
1360                break;
1361            }
1362            tokio::time::sleep(std::time::Duration::from_millis(5)).await;
1363        }
1364        assert!(
1365            abort_handle.is_finished(),
1366            "output task must be aborted after drain"
1367        );
1368    }
1369
1370    /// Regression for the H3 wiring (not just the drain helper): `spawn`/`spawn_inner` must capture
1371    /// the `JoinHandle` returned by `install_output_callback` into `ProcessEntry::output_tasks`. If a
1372    /// refactor forgot to push the handle, the callback task would be unreachable and
1373    /// `drain_process_output_tasks` would have nothing to abort, re-leaking the task. This reproduces
1374    /// that exact seam and asserts the stored handle is the live callback task.
1375    #[tokio::test]
1376    async fn install_output_callback_handle_is_captured_into_process_entry() {
1377        use std::sync::atomic::{AtomicUsize, Ordering};
1378        use std::sync::Arc;
1379
1380        let (stdout_tx, _) = broadcast::channel::<Vec<u8>>(8);
1381        let (stderr_tx, _) = broadcast::channel::<Vec<u8>>(8);
1382        let (output_tx, _) = broadcast::channel(8);
1383        let (exit_tx, _) = watch::channel::<Option<i32>>(None);
1384        let (kernel_pid_tx, _) = watch::channel::<Option<u32>>(None);
1385
1386        let calls = Arc::new(AtomicUsize::new(0));
1387        let calls_cb = Arc::clone(&calls);
1388        let cb: OutputCallback = Box::new(move |_chunk: &[u8]| {
1389            calls_cb.fetch_add(1, Ordering::SeqCst);
1390        });
1391
1392        // The exact seam from `spawn_inner`: capture the returned handle in `output_tasks`.
1393        let output_tasks = vec![install_output_callback(stdout_tx.clone(), cb)];
1394
1395        let entry = ProcessEntry {
1396            command: "sleep".to_string(),
1397            args: vec!["3600".to_string()],
1398            stdout_tx: stdout_tx.clone(),
1399            stderr_tx,
1400            output_tx,
1401            exit_tx,
1402            process_id: "proc-test".to_string(),
1403            kernel_pid: kernel_pid_tx,
1404            output_tasks,
1405            started_at: 0,
1406        };
1407
1408        assert_eq!(
1409            entry.output_tasks.len(),
1410            1,
1411            "the install_output_callback handle must be captured on the entry"
1412        );
1413
1414        // Prove the captured handle is the live callback task: a chunk on the channel runs it.
1415        stdout_tx
1416            .send(b"hello".to_vec())
1417            .expect("broadcast send to subscribed callback task");
1418        for _ in 0..100 {
1419            if calls.load(Ordering::SeqCst) > 0 {
1420                break;
1421            }
1422            tokio::time::sleep(std::time::Duration::from_millis(5)).await;
1423        }
1424        assert_eq!(
1425            calls.load(Ordering::SeqCst),
1426            1,
1427            "the stored handle must drive the registered callback"
1428        );
1429
1430        // And it is the handle `drain_process_output_tasks` aborts on teardown.
1431        let processes: SccHashMap<u32, ProcessEntry> = SccHashMap::new();
1432        let _ = processes.insert(1, entry);
1433        drain_process_output_tasks(&processes);
1434        assert!(processes.is_empty(), "registry must be cleared on drain");
1435    }
1436
1437    #[test]
1438    fn exec_options_default_uses_workspace_cwd() {
1439        assert_eq!(
1440            ExecOptions::default().cwd.as_deref(),
1441            Some(DEFAULT_EXEC_CWD)
1442        );
1443    }
1444
1445    #[test]
1446    fn append_exec_output_rejects_capture_over_limit() {
1447        let mut buffer = vec![0u8; EXEC_OUTPUT_CAPTURE_LIMIT_BYTES - 1];
1448        let mut captured = buffer.len();
1449
1450        append_exec_output(&mut buffer, &[1], &mut captured, "stdout")
1451            .expect("chunk at limit should fit");
1452        assert_eq!(captured, EXEC_OUTPUT_CAPTURE_LIMIT_BYTES);
1453
1454        let error = append_exec_output(&mut buffer, &[2], &mut captured, "stdout")
1455            .expect_err("chunk over limit should fail");
1456        assert!(
1457            error.to_string().contains("exec stdout capture is"),
1458            "unexpected error: {error}"
1459        );
1460        assert_eq!(captured, EXEC_OUTPUT_CAPTURE_LIMIT_BYTES);
1461        assert_eq!(buffer.len(), EXEC_OUTPUT_CAPTURE_LIMIT_BYTES);
1462    }
1463
1464    #[test]
1465    fn exited_pid_pruning_keeps_live_entries_and_removes_oldest_exited() {
1466        let pids = exited_pids_to_prune(vec![(3, true), (1, false), (2, true), (4, true)], 2);
1467        assert_eq!(pids, vec![2, 3]);
1468    }
1469
1470    #[test]
1471    fn observed_time_pruning_enforces_limit() {
1472        let map = SccHashMap::new();
1473        let _ = map.insert("b".to_string(), 2.0);
1474        let _ = map.insert("a".to_string(), 1.0);
1475        let _ = map.insert("c".to_string(), 3.0);
1476
1477        prune_string_f64_map(&map, 2);
1478
1479        assert!(map.read("a", |_, _| ()).is_none());
1480        assert!(map.read("b", |_, _| ()).is_some());
1481        assert!(map.read("c", |_, _| ()).is_some());
1482    }
1483}