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(&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(&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(
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::VmLifecycleEvent(_)
381                | EventPayload::StructuredEvent(_)
382                | EventPayload::ExtEnvelope(_) => {}
383            }
384        };
385
386        if let Some(error) = capture_error {
387            return Err(error.into());
388        }
389
390        Ok(ExecResult {
391            exit_code,
392            stdout: String::from_utf8_lossy(&stdout).into_owned(),
393            stderr: String::from_utf8_lossy(&stderr).into_owned(),
394        })
395    }
396
397    /// Spawn a process. SYNC; returns `{ pid }` only. Installs stdout/stderr fan-out over broadcast
398    /// channels and wires exit via a background event-pump task. The user-facing `pid` is the
399    /// SDK-allocated map key (the wire `process_id` is held inside the [`ProcessEntry`]).
400    pub fn spawn(
401        &self,
402        command: &str,
403        args: Vec<String>,
404        options: SpawnOptions,
405    ) -> Result<SpawnHandle> {
406        let registry_guard = self.inner().process_registry_lock.lock();
407        self.prune_exited_processes_locked(1);
408        if self.process_registry_len_locked() >= PROCESS_REGISTRY_LIMIT {
409            return Err(ClientError::Sidecar(format!(
410                "process registry limit exceeded: at most {PROCESS_REGISTRY_LIMIT} processes can be tracked per VM"
411            ))
412            .into());
413        }
414
415        // Draw the public pid from the dedicated synthetic-pid space (TS `nextSyntheticPid`), seeded
416        // at `SYNTHETIC_PID_BASE`. `exec` uses a separate counter so it never perturbs this sequence.
417        let pid = self
418            .inner()
419            .synthetic_pid_counter
420            .fetch_add(1, Ordering::SeqCst) as u32;
421        let process_id = format!("proc-{pid}-{}", uuid::Uuid::new_v4());
422
423        let (stdout_tx, _) = broadcast::channel::<Vec<u8>>(PROCESS_STREAM_CAPACITY);
424        let (stderr_tx, _) = broadcast::channel::<Vec<u8>>(PROCESS_STREAM_CAPACITY);
425        let (output_tx, _) = broadcast::channel::<ProcessOutput>(PROCESS_STREAM_CAPACITY);
426        // Seeded `None`; the already-exited branch of `on_process_exit` fires immediately once this
427        // watch holds `Some(code)`.
428        let (exit_tx, _) = watch::channel::<Option<i32>>(None);
429        // Seeded `None`; filled with the kernel pid once the `Execute` response lands so
430        // `all_processes`/`process_tree` can remap the kernel snapshot back to this display pid.
431        let (kernel_pid_tx, _) = watch::channel::<Option<u32>>(None);
432
433        let entry = ProcessEntry {
434            command: command.to_owned(),
435            args: args.clone(),
436            stdout_tx: stdout_tx.clone(),
437            stderr_tx: stderr_tx.clone(),
438            output_tx: output_tx.clone(),
439            exit_tx: exit_tx.clone(),
440            process_id: process_id.clone(),
441            kernel_pid: kernel_pid_tx.clone(),
442            output_tasks: Vec::new(),
443            started_at: epoch_ms_now() as i64,
444        };
445        // `spawn` is documented as overwriting any prior entry for a freshly allocated pid; the pid
446        // is monotonic so a collision is not expected.
447        let _ = self.inner().processes.insert(pid, entry);
448        drop(registry_guard);
449
450        // Subscribe to events before issuing the request so the pump sees everything.
451        let events = self.transport().subscribe_wire_events();
452
453        let this = self.clone();
454        let command = command.to_owned();
455        tokio::spawn(async move {
456            this.run_spawn(
457                pid,
458                process_id,
459                command,
460                args,
461                options,
462                events,
463                stdout_tx,
464                stderr_tx,
465                output_tx,
466                exit_tx,
467                kernel_pid_tx,
468            )
469            .await;
470        });
471
472        Ok(SpawnHandle { pid })
473    }
474
475    /// Write to a spawned process's stdin. SYNC. Errors with `ProcessNotFound`.
476    pub fn write_process_stdin(
477        &self,
478        pid: u32,
479        data: StdinInput,
480    ) -> std::result::Result<(), ClientError> {
481        let process_id = self.lookup_process_id(pid)?;
482        let chunk: Vec<u8> = stdin_to_bytes(data);
483        let this = self.clone();
484        // Fire-and-forget: the TS API is synchronous and does not surface a write error.
485        tokio::spawn(async move {
486            let ownership = this.vm_scope();
487            let _ = this
488                .transport()
489                .request_wire(
490                    ownership,
491                    wire::RequestPayload::WriteStdinRequest(wire::WriteStdinRequest {
492                        process_id,
493                        chunk,
494                    }),
495                )
496                .await;
497        });
498        Ok(())
499    }
500
501    /// Close a spawned process's stdin. SYNC. Errors with `ProcessNotFound`.
502    pub fn close_process_stdin(&self, pid: u32) -> std::result::Result<(), ClientError> {
503        let process_id = self.lookup_process_id(pid)?;
504        let this = self.clone();
505        tokio::spawn(async move {
506            let ownership = this.vm_scope();
507            let _ = this
508                .transport()
509                .request_wire(
510                    ownership,
511                    wire::RequestPayload::CloseStdinRequest(wire::CloseStdinRequest { process_id }),
512                )
513                .await;
514        });
515        Ok(())
516    }
517
518    /// Subscribe to the unified stdout/stderr event stream for a process.
519    pub fn on_process_output(
520        &self,
521        pid: u32,
522        mut handler: impl FnMut(ProcessOutput) + Send + 'static,
523    ) -> std::result::Result<Subscription, ClientError> {
524        let mut rx = self
525            .inner()
526            .processes
527            .read(&pid, |_, entry| entry.output_tx.subscribe())
528            .ok_or(ClientError::ProcessNotFound(pid))?;
529        let task = tokio::spawn(async move {
530            loop {
531                match rx.recv().await {
532                    Ok(event) => handler(event),
533                    Err(broadcast::error::RecvError::Lagged(_)) => continue,
534                    Err(broadcast::error::RecvError::Closed) => return,
535                }
536            }
537        });
538        Ok(Subscription::new(move || task.abort()))
539    }
540
541    /// Register a once-only exit handler. If the process has already exited, the handler fires
542    /// immediately and synchronously and a no-op unsubscribe is returned (the `watch` already holds
543    /// `Some(code)`). Otherwise the handler fires once when the exit code lands. The exit code is
544    /// `i32`, never null.
545    pub fn on_process_exit(
546        &self,
547        pid: u32,
548        handler: impl FnOnce(ProcessExit) + Send + 'static,
549    ) -> std::result::Result<Subscription, ClientError> {
550        let mut rx = self
551            .inner()
552            .processes
553            .read(&pid, |_, entry| entry.exit_tx.subscribe())
554            .ok_or(ClientError::ProcessNotFound(pid))?;
555
556        // Already-exited branch: fire immediately + synchronously, return a no-op unsubscribe.
557        if let Some(code) = *rx.borrow() {
558            handler(ProcessExit {
559                pid,
560                exit_code: code,
561            });
562            return Ok(Subscription::noop());
563        }
564
565        // Otherwise wait for the watch to transition to `Some(code)` and fire exactly once. The
566        // returned `Subscription` cancels the waiting task on drop (= unsubscribe).
567        let task = tokio::spawn(async move {
568            while rx.changed().await.is_ok() {
569                if let Some(code) = *rx.borrow() {
570                    handler(ProcessExit {
571                        pid,
572                        exit_code: code,
573                    });
574                    return;
575                }
576            }
577        });
578        Ok(Subscription::new(move || task.abort()))
579    }
580
581    /// Await a spawned process's exit code. Unknown-pid lookup errors (synchronously in TS; here the
582    /// lookup error is returned before any awaiting begins).
583    pub async fn wait_process(&self, pid: u32) -> std::result::Result<i32, ClientError> {
584        let mut rx = self
585            .inner()
586            .processes
587            .read(&pid, |_, entry| entry.exit_tx.subscribe())
588            .ok_or(ClientError::ProcessNotFound(pid))?;
589
590        if let Some(code) = *rx.borrow() {
591            return Ok(code);
592        }
593        while rx.changed().await.is_ok() {
594            if let Some(code) = *rx.borrow() {
595                return Ok(code);
596            }
597        }
598        Err(ClientError::Sidecar(format!(
599            "wait_process: exit channel closed before process {pid} reported an exit code"
600        )))
601    }
602
603    /// List SDK-spawned processes only. `running = exit_code.is_none()`.
604    pub fn list_processes(&self) -> Vec<SpawnedProcessInfo> {
605        let mut out = Vec::new();
606        self.inner().processes.scan(|pid, entry| {
607            let exit_code = *entry.exit_tx.borrow();
608            out.push(SpawnedProcessInfo {
609                pid: *pid,
610                command: entry.command.clone(),
611                args: entry.args.clone(),
612                running: exit_code.is_none(),
613                exit_code,
614                started_at: entry.started_at,
615            });
616        });
617        out
618    }
619
620    /// List ALL kernel processes (native sidecar process snapshot).
621    ///
622    /// The kernel snapshot keys processes by their raw kernel pid. SDK-spawned root processes carry a
623    /// synthetic display pid (the `spawn` return value); this remaps each snapshot entry's
624    /// pid/ppid/pgid/sid back to that display pid via the per-process `kernel_pid` watch, so a caller
625    /// can correlate `spawn()` with `all_processes()`/`process_tree()`. Results are sorted ascending
626    /// by display pid (TS `snapshotProcesses` `.sort((l,r) => l.pid - r.pid)`).
627    pub async fn all_processes(&self) -> Result<Vec<ProcessInfo>> {
628        let ownership = self.vm_scope();
629        let response = self
630            .transport()
631            .request_wire(ownership, wire::RequestPayload::GetProcessSnapshotRequest)
632            .await
633            .context("all_processes: GetProcessSnapshot request failed")?;
634        let snapshot = match response {
635            wire::ResponsePayload::ProcessSnapshotResponse(snapshot) => snapshot,
636            wire::ResponsePayload::RejectedResponse(rejected) => {
637                return Err(ClientError::from_rejection(rejected).into());
638            }
639            other => {
640                return Err(ClientError::Sidecar(format!(
641                    "all_processes: unexpected response {other:?}"
642                ))
643                .into());
644            }
645        };
646
647        // Snapshot the SDK process registry, keyed by wire `process_id`, capturing exit code,
648        // command, and args. This mirrors the TS `trackedProcessesById` lookup used to build
649        // `displayPidByKernelPid` and override fields.
650        struct Tracked {
651            exit_code: Option<i32>,
652            command: String,
653            args: Vec<String>,
654        }
655        let mut tracked_by_process_id: BTreeMap<String, Tracked> = BTreeMap::new();
656        let mut display_pid_by_kernel_pid: BTreeMap<u32, u32> = BTreeMap::new();
657        self.inner().processes.scan(|display_pid, entry| {
658            let exit_code = *entry.exit_tx.borrow();
659            if let Some(kernel_pid) = *entry.kernel_pid.borrow() {
660                display_pid_by_kernel_pid.insert(kernel_pid, *display_pid);
661            }
662            tracked_by_process_id.insert(
663                entry.process_id.clone(),
664                Tracked {
665                    exit_code,
666                    command: entry.command.clone(),
667                    args: entry.args.clone(),
668                },
669            );
670        });
671
672        let now_ms = epoch_ms_now();
673        let mut seen_display_pids: std::collections::BTreeSet<u32> =
674            std::collections::BTreeSet::new();
675        let mut out: Vec<ProcessInfo> = Vec::new();
676
677        for entry in snapshot.processes {
678            let tracked = tracked_by_process_id.get(&entry.process_id);
679            let display_pid = display_pid_by_kernel_pid
680                .get(&entry.pid)
681                .copied()
682                .unwrap_or(entry.pid);
683            let display_ppid = display_pid_by_kernel_pid
684                .get(&entry.ppid)
685                .copied()
686                .unwrap_or(entry.ppid);
687            let display_pgid = display_pid_by_kernel_pid
688                .get(&entry.pgid)
689                .copied()
690                .unwrap_or(entry.pgid);
691            let display_sid = display_pid_by_kernel_pid
692                .get(&entry.sid)
693                .copied()
694                .unwrap_or(entry.sid);
695
696            // First-observed start time, keyed by `"<process_id>:<kernel_pid>"` (TS `processKey`).
697            let process_key = format!("{}:{}", entry.process_id, entry.pid);
698            let start_time = self.observed_start_time(&process_key, now_ms);
699
700            // Status/exit code: a tracked process whose SDK exit code is known is `exited`; otherwise
701            // a tracked process is `running`; an untracked process uses the snapshot status.
702            let (status, exit_code) = match tracked {
703                Some(t) => match t.exit_code {
704                    Some(code) => (ProcessStatus::Exited, Some(code)),
705                    None => (ProcessStatus::Running, entry.exit_code),
706                },
707                None => {
708                    let status = match entry.status {
709                        ProcessSnapshotStatus::Running | ProcessSnapshotStatus::Stopped => {
710                            ProcessStatus::Running
711                        }
712                        ProcessSnapshotStatus::Exited => ProcessStatus::Exited,
713                    };
714                    (status, entry.exit_code)
715                }
716            };
717
718            // Exit time: only tracked-and-exited processes carry one (TS `tracked?.exitTime`).
719            let exit_time = match (tracked, status) {
720                (Some(_), ProcessStatus::Exited) => {
721                    Some(self.observed_exit_time(&entry.process_id, now_ms))
722                }
723                _ => None,
724            };
725
726            let (command, args) = match tracked {
727                Some(t) => (t.command.clone(), t.args.clone()),
728                None => (entry.command, entry.args),
729            };
730
731            seen_display_pids.insert(display_pid);
732            out.push(ProcessInfo {
733                pid: display_pid,
734                ppid: display_ppid,
735                pgid: display_pgid,
736                sid: display_sid,
737                driver: entry.driver,
738                command,
739                args,
740                cwd: entry.cwd,
741                status,
742                exit_code,
743                start_time,
744                exit_time,
745            });
746        }
747
748        // Tracked processes not yet present in the snapshot (the spawn `Execute` has not surfaced in
749        // the kernel table yet). TS fills these with `ppid:0, pgid/sid = pid`.
750        self.inner().processes.scan(|display_pid, entry| {
751            if seen_display_pids.contains(display_pid) {
752                return;
753            }
754            let exit_code = *entry.exit_tx.borrow();
755            let process_key = format!("{}:{}", entry.process_id, display_pid);
756            let start_time = self.observed_start_time(&process_key, now_ms);
757            let (status, exit_time) = match exit_code {
758                Some(_) => (
759                    ProcessStatus::Exited,
760                    Some(self.observed_exit_time(&entry.process_id, now_ms)),
761                ),
762                None => (ProcessStatus::Running, None),
763            };
764            out.push(ProcessInfo {
765                pid: *display_pid,
766                ppid: 0,
767                pgid: *display_pid,
768                sid: *display_pid,
769                driver: String::new(),
770                command: entry.command.clone(),
771                args: entry.args.clone(),
772                cwd: String::new(),
773                status,
774                exit_code,
775                start_time,
776                exit_time,
777            });
778        });
779
780        out.sort_by_key(|info| info.pid);
781        Ok(out)
782    }
783
784    /// Return the first-observed start time for a process key, recording `now` the first time it is
785    /// seen so later snapshots report a stable timestamp (TS `observedProcessStartTimes`).
786    fn observed_start_time(&self, process_key: &str, now_ms: f64) -> f64 {
787        let _guard = self.inner().observed_process_time_lock.lock();
788        if let Some(existing) = self
789            .inner()
790            .observed_process_start_times
791            .read(process_key, |_, value| *value)
792        {
793            return existing;
794        }
795        let _ = self
796            .inner()
797            .observed_process_start_times
798            .insert(process_key.to_owned(), now_ms);
799        prune_string_f64_map(
800            &self.inner().observed_process_start_times,
801            OBSERVED_PROCESS_TIME_LIMIT,
802        );
803        // Re-read to honor a racing insert that may have won; either value is a valid first-observed
804        // timestamp.
805        self.inner()
806            .observed_process_start_times
807            .read(process_key, |_, value| *value)
808            .unwrap_or(now_ms)
809    }
810
811    /// Return the first-observed exit time for an SDK process id, recording `now` on first sight.
812    fn observed_exit_time(&self, process_id: &str, now_ms: f64) -> f64 {
813        let _guard = self.inner().observed_process_time_lock.lock();
814        if let Some(existing) = self
815            .inner()
816            .observed_process_exit_times
817            .read(process_id, |_, value| *value)
818        {
819            return existing;
820        }
821        let _ = self
822            .inner()
823            .observed_process_exit_times
824            .insert(process_id.to_owned(), now_ms);
825        prune_string_f64_map(
826            &self.inner().observed_process_exit_times,
827            OBSERVED_PROCESS_TIME_LIMIT,
828        );
829        self.inner()
830            .observed_process_exit_times
831            .read(process_id, |_, value| *value)
832            .unwrap_or(now_ms)
833    }
834
835    /// Build the process forest from `all_processes`, linked by `ppid`.
836    pub async fn process_tree(&self) -> Result<Vec<ProcessTreeNode>> {
837        let processes = self.all_processes().await?;
838        Ok(build_process_forest(processes))
839    }
840
841    /// Get a single SDK-spawned process's info. Errors (not None) when not found.
842    pub fn get_process(&self, pid: u32) -> std::result::Result<SpawnedProcessInfo, ClientError> {
843        self.inner()
844            .processes
845            .read(&pid, |pid, entry| {
846                let exit_code = *entry.exit_tx.borrow();
847                SpawnedProcessInfo {
848                    pid: *pid,
849                    command: entry.command.clone(),
850                    args: entry.args.clone(),
851                    running: exit_code.is_none(),
852                    exit_code,
853                    started_at: entry.started_at,
854                }
855            })
856            .ok_or(ClientError::ProcessNotFound(pid))
857    }
858
859    /// SIGTERM a spawned process. No-op if already exited; errors if unknown.
860    pub fn stop_process(&self, pid: u32) -> std::result::Result<(), ClientError> {
861        self.signal_process(pid, "SIGTERM")
862    }
863
864    /// SIGKILL a spawned process. No-op if already exited; errors if unknown.
865    pub fn kill_process(&self, pid: u32) -> std::result::Result<(), ClientError> {
866        self.signal_process(pid, "SIGKILL")
867    }
868
869    // -----------------------------------------------------------------------
870    // Internal helpers
871    // -----------------------------------------------------------------------
872
873    /// Build the VM-scoped ownership for a wire request.
874    fn vm_scope(&self) -> wire::OwnershipScope {
875        wire::OwnershipScope::VmOwnership(wire::VmOwnership {
876            connection_id: self.connection_id().to_string(),
877            session_id: self.wire_session_id().to_string(),
878            vm_id: self.vm_id().to_string(),
879        })
880    }
881
882    /// Allocate a fresh wire `process_id` (used by `exec`, which does not register in the SDK map).
883    fn next_process_id(&self) -> String {
884        let n = self.inner().process_counter.fetch_add(1, Ordering::SeqCst);
885        format!("proc-{n}-{}", uuid::Uuid::new_v4())
886    }
887
888    /// Resolve the wire `process_id` for an SDK pid, erroring with `ProcessNotFound` if unknown.
889    fn lookup_process_id(&self, pid: u32) -> std::result::Result<String, ClientError> {
890        self.inner()
891            .processes
892            .read(&pid, |_, entry| entry.process_id.clone())
893            .ok_or(ClientError::ProcessNotFound(pid))
894    }
895
896    /// Send the `Execute` wire request, mapping a rejection into [`ClientError::Kernel`].
897    async fn send_execute(
898        &self,
899        process_id: &str,
900        command: Option<String>,
901        args: Vec<String>,
902        env: BTreeMap<String, String>,
903        cwd: Option<String>,
904    ) -> std::result::Result<wire::ProcessStartedResponse, ClientError> {
905        let ownership = self.vm_scope();
906        let response = self
907            .transport()
908            .request_wire(
909                ownership,
910                wire::RequestPayload::ExecuteRequest(wire::ExecuteRequest {
911                    process_id: process_id.to_owned(),
912                    command,
913                    runtime: None,
914                    entrypoint: None,
915                    args,
916                    env: env.into_iter().collect(),
917                    cwd,
918                    wasm_permission_tier: None,
919                }),
920            )
921            .await?;
922        match response {
923            wire::ResponsePayload::ProcessStartedResponse(started) => Ok(started),
924            wire::ResponsePayload::RejectedResponse(rejected) => {
925                Err(ClientError::from_rejection(rejected))
926            }
927            other => Err(ClientError::Sidecar(format!(
928                "Execute: unexpected response {other:?}"
929            ))),
930        }
931    }
932
933    /// Fire-and-forget kill of a wire process by its `process_id` (used by `exec` timeout). The TS
934    /// timeout path calls `proc.kill(9)`, which maps to a `SIGKILL` kill request.
935    fn kill_wire_process(&self, process_id: &str, signal: &str) {
936        let process_id = process_id.to_owned();
937        let signal = signal.to_owned();
938        let this = self.clone();
939        tokio::spawn(async move {
940            let ownership = this.vm_scope();
941            let _ = this
942                .transport()
943                .request_wire(
944                    ownership,
945                    wire::RequestPayload::KillProcessRequest(wire::KillProcessRequest {
946                        process_id,
947                        signal,
948                    }),
949                )
950                .await;
951        });
952    }
953
954    /// Send a kill signal for an SDK pid. No-op if already exited; errors with `ProcessNotFound` if
955    /// the pid is unknown.
956    fn signal_process(&self, pid: u32, signal: &str) -> std::result::Result<(), ClientError> {
957        let (process_id, already_exited) = self
958            .inner()
959            .processes
960            .read(&pid, |_, entry| {
961                (entry.process_id.clone(), entry.exit_tx.borrow().is_some())
962            })
963            .ok_or(ClientError::ProcessNotFound(pid))?;
964        if already_exited {
965            return Ok(());
966        }
967        let signal = signal.to_owned();
968        let this = self.clone();
969        tokio::spawn(async move {
970            let ownership = this.vm_scope();
971            let _ = this
972                .transport()
973                .request_wire(
974                    ownership,
975                    wire::RequestPayload::KillProcessRequest(wire::KillProcessRequest {
976                        process_id,
977                        signal,
978                    }),
979                )
980                .await;
981        });
982        Ok(())
983    }
984
985    fn process_registry_len_locked(&self) -> usize {
986        let mut count = 0usize;
987        self.inner().processes.scan(|_, _| {
988            count += 1;
989        });
990        count
991    }
992
993    fn prune_exited_processes_locked(&self, reserve_slots: usize) {
994        let mut entries = Vec::new();
995        self.inner().processes.scan(|pid, entry| {
996            entries.push((*pid, entry.exit_tx.borrow().is_some()));
997        });
998        let target_len = PROCESS_REGISTRY_LIMIT.saturating_sub(reserve_slots);
999        if entries.len() <= target_len {
1000            return;
1001        }
1002
1003        for pid in exited_pids_to_prune(entries, target_len) {
1004            self.remove_process_tracking_locked(pid);
1005        }
1006    }
1007
1008    fn remove_process_tracking_locked(&self, pid: u32) {
1009        if let Some((_, entry)) = self.inner().processes.remove(&pid) {
1010            let _time_guard = self.inner().observed_process_time_lock.lock();
1011            let _ = self
1012                .inner()
1013                .observed_process_exit_times
1014                .remove(&entry.process_id);
1015            let fallback_start_key = format!("{}:{pid}", entry.process_id);
1016            let _ = self
1017                .inner()
1018                .observed_process_start_times
1019                .remove(&fallback_start_key);
1020            if let Some(kernel_pid) = *entry.kernel_pid.borrow() {
1021                let start_key = format!("{}:{kernel_pid}", entry.process_id);
1022                let _ = self.inner().observed_process_start_times.remove(&start_key);
1023            }
1024        }
1025    }
1026
1027    /// Background pump for a spawned process: issue the `Execute` request, then fan kernel
1028    /// `ProcessOutput`/`ProcessExited` events for this process id into the per-process broadcast and
1029    /// watch channels. Exited entries are retained for post-exit inspection, then pruned oldest-first
1030    /// under registry pressure.
1031    #[allow(clippy::too_many_arguments)]
1032    async fn run_spawn(
1033        self,
1034        pid: u32,
1035        process_id: String,
1036        command: String,
1037        args: Vec<String>,
1038        options: SpawnOptions,
1039        mut events: broadcast::Receiver<(wire::OwnershipScope, EventPayload)>,
1040        stdout_tx: broadcast::Sender<Vec<u8>>,
1041        stderr_tx: broadcast::Sender<Vec<u8>>,
1042        output_tx: broadcast::Sender<ProcessOutput>,
1043        exit_tx: watch::Sender<Option<i32>>,
1044        kernel_pid_tx: watch::Sender<Option<u32>>,
1045    ) {
1046        match self
1047            .send_execute(
1048                &process_id,
1049                Some(command),
1050                args,
1051                options.env.clone(),
1052                options.cwd.clone(),
1053            )
1054            .await
1055        {
1056            Ok(started) => {
1057                // Seed the kernel pid so `all_processes`/`process_tree` can remap this process's
1058                // kernel-snapshot entry back to its display pid.
1059                if let Some(kernel_pid) = started.pid {
1060                    let _ = kernel_pid_tx.send(Some(kernel_pid));
1061                }
1062            }
1063            Err(error) => {
1064                // The native TS launch-failure path emits the error message (plus a trailing
1065                // newline) on stderr and resolves the wait with exit code 1 (`startTrackedProcess`
1066                // catch -> stderr handlers + `finishProcess(entry, 1)`).
1067                let message = format!("{error}\n");
1068                let bytes = message.into_bytes();
1069                let _ = stderr_tx.send(bytes.clone());
1070                let _ = output_tx.send(ProcessOutput {
1071                    pid,
1072                    stream: ProcessStream::Stderr,
1073                    data: bytes,
1074                });
1075                tracing::error!(?error, pid, %process_id, "spawn: Execute request failed");
1076                let _ = exit_tx.send(Some(1));
1077                let _guard = self.inner().process_registry_lock.lock();
1078                self.prune_exited_processes_locked(0);
1079                return;
1080            }
1081        }
1082
1083        loop {
1084            let (_, payload) = match events.recv().await {
1085                Ok(frame) => frame,
1086                Err(broadcast::error::RecvError::Lagged(_)) => continue,
1087                Err(broadcast::error::RecvError::Closed) => {
1088                    // The event stream closed before an exit event landed. The TS fallback treats a
1089                    // process that has fully disappeared from the VM snapshot as reaped with exit
1090                    // code 0; mirror that terminal value so waiters resolve instead of hanging.
1091                    let _ = exit_tx.send(Some(0));
1092                    break;
1093                }
1094            };
1095            match payload {
1096                EventPayload::ProcessOutputEvent(output) if output.process_id == process_id => {
1097                    let bytes = output.chunk;
1098                    let _ = output_tx.send(ProcessOutput {
1099                        pid,
1100                        stream: match output.channel {
1101                            StreamChannel::Stdout => ProcessStream::Stdout,
1102                            StreamChannel::Stderr => ProcessStream::Stderr,
1103                        },
1104                        data: bytes.clone(),
1105                    });
1106                    match output.channel {
1107                        StreamChannel::Stdout => {
1108                            let _ = stdout_tx.send(bytes);
1109                        }
1110                        StreamChannel::Stderr => {
1111                            let _ = stderr_tx.send(bytes);
1112                        }
1113                    }
1114                }
1115                EventPayload::ProcessExitedEvent(exited) if exited.process_id == process_id => {
1116                    let _ = exit_tx.send(Some(exited.exit_code));
1117                    break;
1118                }
1119                EventPayload::ProcessOutputEvent(_)
1120                | EventPayload::ProcessExitedEvent(_)
1121                | EventPayload::VmLifecycleEvent(_)
1122                | EventPayload::StructuredEvent(_)
1123                | EventPayload::ExtEnvelope(_) => {}
1124            }
1125        }
1126        let _guard = self.inner().process_registry_lock.lock();
1127        self.prune_exited_processes_locked(0);
1128    }
1129}
1130
1131/// Assemble a process forest from a flat process list, linking children by `ppid`.
1132///
1133/// Mirrors the TS `processTree` `nodeMap` algorithm exactly: a process is a root iff its `ppid` is
1134/// NOT present among the listed pids. A self-parented process (`ppid == pid`) finds itself as its
1135/// parent, so it is attached as its own child and is excluded from the roots (effectively dropped
1136/// from the output tree). A `seen` guard prevents the self-cycle from recursing forever.
1137fn build_process_forest(processes: Vec<ProcessInfo>) -> Vec<ProcessTreeNode> {
1138    use std::collections::BTreeMap as Map;
1139
1140    let pids: std::collections::BTreeSet<u32> = processes.iter().map(|p| p.pid).collect();
1141    // Children adjacency keyed by parent pid, preserving input (sorted) order.
1142    let mut children_of: Map<u32, Vec<usize>> = Map::new();
1143    let mut roots: Vec<usize> = Vec::new();
1144    for (index, proc) in processes.iter().enumerate() {
1145        if pids.contains(&proc.ppid) {
1146            children_of.entry(proc.ppid).or_default().push(index);
1147        } else {
1148            roots.push(index);
1149        }
1150    }
1151
1152    fn build_node(
1153        index: usize,
1154        processes: &[ProcessInfo],
1155        children_of: &Map<u32, Vec<usize>>,
1156        seen: &mut std::collections::BTreeSet<usize>,
1157    ) -> ProcessTreeNode {
1158        let info = processes[index].clone();
1159        seen.insert(index);
1160        let child_indices: Vec<usize> = children_of
1161            .get(&info.pid)
1162            .map(|indices| {
1163                indices
1164                    .iter()
1165                    .copied()
1166                    .filter(|child_index| !seen.contains(child_index))
1167                    .collect()
1168            })
1169            .unwrap_or_default();
1170        let children = child_indices
1171            .into_iter()
1172            .map(|child_index| build_node(child_index, processes, children_of, seen))
1173            .collect();
1174        ProcessTreeNode { info, children }
1175    }
1176
1177    let mut seen = std::collections::BTreeSet::new();
1178    roots
1179        .into_iter()
1180        .map(|index| build_node(index, &processes, &children_of, &mut seen))
1181        .collect()
1182}
1183
1184/// Convert a [`StdinInput`] to raw bytes. A string is delivered as its UTF-8 bytes; raw bytes are
1185/// delivered verbatim (binary-safe, never lossy).
1186fn stdin_to_bytes(input: StdinInput) -> Vec<u8> {
1187    match input {
1188        StdinInput::Text(text) => text.into_bytes(),
1189        StdinInput::Bytes(bytes) => bytes,
1190    }
1191}
1192
1193fn append_exec_output(
1194    buffer: &mut Vec<u8>,
1195    chunk: &[u8],
1196    captured_output_bytes: &mut usize,
1197    channel: &str,
1198) -> std::result::Result<(), ClientError> {
1199    let next_total = captured_output_bytes
1200        .checked_add(chunk.len())
1201        .ok_or_else(|| exec_output_limit_error(channel, usize::MAX))?;
1202    if next_total > EXEC_OUTPUT_CAPTURE_LIMIT_BYTES {
1203        return Err(exec_output_limit_error(channel, next_total));
1204    }
1205    buffer.extend_from_slice(chunk);
1206    *captured_output_bytes = next_total;
1207    Ok(())
1208}
1209
1210fn exec_output_limit_error(channel: &str, size: usize) -> ClientError {
1211    ClientError::Sidecar(format!(
1212        "exec {channel} capture is {size} bytes, limit is {EXEC_OUTPUT_CAPTURE_LIMIT_BYTES}"
1213    ))
1214}
1215
1216fn exited_pids_to_prune(mut entries: Vec<(u32, bool)>, target_len: usize) -> Vec<u32> {
1217    if entries.len() <= target_len {
1218        return Vec::new();
1219    }
1220    let mut remove_count = entries.len() - target_len;
1221    entries.sort_by_key(|(pid, _)| *pid);
1222    let mut out = Vec::new();
1223    for (pid, exited) in entries {
1224        if remove_count == 0 {
1225            break;
1226        }
1227        if !exited {
1228            continue;
1229        }
1230        out.push(pid);
1231        remove_count -= 1;
1232    }
1233    out
1234}
1235
1236fn prune_string_f64_map(map: &SccHashMap<String, f64>, limit: usize) {
1237    let mut keys = Vec::new();
1238    map.scan(|key, _| {
1239        keys.push(key.clone());
1240    });
1241    if keys.len() <= limit {
1242        return;
1243    }
1244    let remove_count = keys.len() - limit;
1245    keys.sort();
1246    for key in keys.into_iter().take(remove_count) {
1247        let _ = map.remove(&key);
1248    }
1249}
1250
1251/// Drive a caller-supplied output callback from a fresh subscription on the given broadcast channel.
1252/// Each chunk delivered to the channel is forwarded to `callback` as raw bytes. The task ends when
1253/// the channel closes (process exit), matching the TS handler-set lifetime.
1254///
1255/// Returns the spawned task's handle so the owner can abort it on teardown: a [`ProcessEntry`]
1256/// retains its own `stdout_tx`/`stderr_tx` clone for late subscribers, so the broadcast channel
1257/// never closes (and this task never observes `Closed`) until the entry is dropped. `shutdown`
1258/// drains the registry and aborts these handles rather than waiting on the channel close.
1259pub(crate) fn install_output_callback(
1260    tx: broadcast::Sender<Vec<u8>>,
1261    mut callback: OutputCallback,
1262) -> JoinHandle<()> {
1263    let mut rx = tx.subscribe();
1264    tokio::spawn(async move {
1265        loop {
1266            match rx.recv().await {
1267                Ok(chunk) => callback(&chunk),
1268                Err(broadcast::error::RecvError::Lagged(_)) => continue,
1269                Err(broadcast::error::RecvError::Closed) => break,
1270            }
1271        }
1272    })
1273}
1274
1275/// Drain the SDK-spawned process registry, dropping each entry's retained sender clones and aborting
1276/// its per-process output-callback tasks. Called from `shutdown` so the output tasks (which would
1277/// otherwise await a `Closed` that never fires, see [`install_output_callback`]) cannot outlive the
1278/// disposed VM. Mirrors the `pending_shell_exits` / ACP-terminal drain in `shutdown`.
1279pub(crate) fn drain_process_output_tasks(processes: &SccHashMap<u32, ProcessEntry>) {
1280    let mut tasks = Vec::new();
1281    processes.retain(|_, entry| {
1282        tasks.append(&mut entry.output_tasks);
1283        false
1284    });
1285    for task in tasks {
1286        task.abort();
1287    }
1288}
1289
1290/// Current wall-clock time as epoch milliseconds (TS `Date.now()`).
1291fn epoch_ms_now() -> f64 {
1292    use std::time::{SystemTime, UNIX_EPOCH};
1293    SystemTime::now()
1294        .duration_since(UNIX_EPOCH)
1295        .map(|d| d.as_secs_f64() * 1000.0)
1296        .unwrap_or(0.0)
1297}
1298
1299#[cfg(test)]
1300mod tests {
1301    use super::{
1302        append_exec_output, drain_process_output_tasks, exited_pids_to_prune,
1303        install_output_callback, prune_string_f64_map, ExecOptions, OutputCallback,
1304        DEFAULT_EXEC_CWD, EXEC_OUTPUT_CAPTURE_LIMIT_BYTES,
1305    };
1306    use crate::agent_os::ProcessEntry;
1307    use scc::HashMap as SccHashMap;
1308    use tokio::sync::{broadcast, watch};
1309
1310    /// Regression for the per-process output-callback leak (H3): a `ProcessEntry` retains clones of
1311    /// its `stdout_tx`/`stderr_tx`, so the output tasks never observe the broadcast `Closed` and hang
1312    /// forever unless teardown aborts them. `drain_process_output_tasks` must empty the registry and
1313    /// abort every retained output task.
1314    #[tokio::test]
1315    async fn drain_process_output_tasks_clears_registry_and_aborts_tasks() {
1316        let processes: SccHashMap<u32, ProcessEntry> = SccHashMap::new();
1317
1318        let (stdout_tx, _) = broadcast::channel::<Vec<u8>>(8);
1319        let (stderr_tx, _) = broadcast::channel::<Vec<u8>>(8);
1320        let (output_tx, _) = broadcast::channel(8);
1321        let (exit_tx, _) = watch::channel::<Option<i32>>(None);
1322        let (kernel_pid_tx, _) = watch::channel::<Option<u32>>(None);
1323
1324        // A task that never completes on its own, standing in for an output-callback task that is
1325        // waiting on a `Closed` that the retained sender clone prevents.
1326        let task = tokio::spawn(async {
1327            loop {
1328                tokio::time::sleep(std::time::Duration::from_secs(3600)).await;
1329            }
1330        });
1331        let abort_handle = task.abort_handle();
1332
1333        let entry = ProcessEntry {
1334            command: "sleep".to_string(),
1335            args: vec!["3600".to_string()],
1336            stdout_tx,
1337            stderr_tx,
1338            output_tx,
1339            exit_tx,
1340            process_id: "proc-test".to_string(),
1341            kernel_pid: kernel_pid_tx,
1342            output_tasks: vec![task],
1343            started_at: 0,
1344        };
1345        let _ = processes.insert(1, entry);
1346
1347        assert!(!abort_handle.is_finished(), "task should start alive");
1348
1349        drain_process_output_tasks(&processes);
1350
1351        assert!(processes.is_empty(), "registry must be cleared on drain");
1352
1353        // The abort is asynchronous; give the runtime a bounded window to reap the cancelled task.
1354        for _ in 0..100 {
1355            if abort_handle.is_finished() {
1356                break;
1357            }
1358            tokio::time::sleep(std::time::Duration::from_millis(5)).await;
1359        }
1360        assert!(
1361            abort_handle.is_finished(),
1362            "output task must be aborted after drain"
1363        );
1364    }
1365
1366    /// Regression for the H3 wiring (not just the drain helper): `spawn`/`spawn_inner` must capture
1367    /// the `JoinHandle` returned by `install_output_callback` into `ProcessEntry::output_tasks`. If a
1368    /// refactor forgot to push the handle, the callback task would be unreachable and
1369    /// `drain_process_output_tasks` would have nothing to abort, re-leaking the task. This reproduces
1370    /// that exact seam and asserts the stored handle is the live callback task.
1371    #[tokio::test]
1372    async fn install_output_callback_handle_is_captured_into_process_entry() {
1373        use std::sync::atomic::{AtomicUsize, Ordering};
1374        use std::sync::Arc;
1375
1376        let (stdout_tx, _) = broadcast::channel::<Vec<u8>>(8);
1377        let (stderr_tx, _) = broadcast::channel::<Vec<u8>>(8);
1378        let (output_tx, _) = broadcast::channel(8);
1379        let (exit_tx, _) = watch::channel::<Option<i32>>(None);
1380        let (kernel_pid_tx, _) = watch::channel::<Option<u32>>(None);
1381
1382        let calls = Arc::new(AtomicUsize::new(0));
1383        let calls_cb = Arc::clone(&calls);
1384        let cb: OutputCallback = Box::new(move |_chunk: &[u8]| {
1385            calls_cb.fetch_add(1, Ordering::SeqCst);
1386        });
1387
1388        // The exact seam from `spawn_inner`: capture the returned handle in `output_tasks`.
1389        let output_tasks = vec![install_output_callback(stdout_tx.clone(), cb)];
1390
1391        let entry = ProcessEntry {
1392            command: "sleep".to_string(),
1393            args: vec!["3600".to_string()],
1394            stdout_tx: stdout_tx.clone(),
1395            stderr_tx,
1396            output_tx,
1397            exit_tx,
1398            process_id: "proc-test".to_string(),
1399            kernel_pid: kernel_pid_tx,
1400            output_tasks,
1401            started_at: 0,
1402        };
1403
1404        assert_eq!(
1405            entry.output_tasks.len(),
1406            1,
1407            "the install_output_callback handle must be captured on the entry"
1408        );
1409
1410        // Prove the captured handle is the live callback task: a chunk on the channel runs it.
1411        stdout_tx
1412            .send(b"hello".to_vec())
1413            .expect("broadcast send to subscribed callback task");
1414        for _ in 0..100 {
1415            if calls.load(Ordering::SeqCst) > 0 {
1416                break;
1417            }
1418            tokio::time::sleep(std::time::Duration::from_millis(5)).await;
1419        }
1420        assert_eq!(
1421            calls.load(Ordering::SeqCst),
1422            1,
1423            "the stored handle must drive the registered callback"
1424        );
1425
1426        // And it is the handle `drain_process_output_tasks` aborts on teardown.
1427        let processes: SccHashMap<u32, ProcessEntry> = SccHashMap::new();
1428        let _ = processes.insert(1, entry);
1429        drain_process_output_tasks(&processes);
1430        assert!(processes.is_empty(), "registry must be cleared on drain");
1431    }
1432
1433    #[test]
1434    fn exec_options_default_uses_workspace_cwd() {
1435        assert_eq!(
1436            ExecOptions::default().cwd.as_deref(),
1437            Some(DEFAULT_EXEC_CWD)
1438        );
1439    }
1440
1441    #[test]
1442    fn append_exec_output_rejects_capture_over_limit() {
1443        let mut buffer = vec![0u8; EXEC_OUTPUT_CAPTURE_LIMIT_BYTES - 1];
1444        let mut captured = buffer.len();
1445
1446        append_exec_output(&mut buffer, &[1], &mut captured, "stdout")
1447            .expect("chunk at limit should fit");
1448        assert_eq!(captured, EXEC_OUTPUT_CAPTURE_LIMIT_BYTES);
1449
1450        let error = append_exec_output(&mut buffer, &[2], &mut captured, "stdout")
1451            .expect_err("chunk over limit should fail");
1452        assert!(
1453            error.to_string().contains("exec stdout capture is"),
1454            "unexpected error: {error}"
1455        );
1456        assert_eq!(captured, EXEC_OUTPUT_CAPTURE_LIMIT_BYTES);
1457        assert_eq!(buffer.len(), EXEC_OUTPUT_CAPTURE_LIMIT_BYTES);
1458    }
1459
1460    #[test]
1461    fn exited_pid_pruning_keeps_live_entries_and_removes_oldest_exited() {
1462        let pids = exited_pids_to_prune(vec![(3, true), (1, false), (2, true), (4, true)], 2);
1463        assert_eq!(pids, vec![2, 3]);
1464    }
1465
1466    #[test]
1467    fn observed_time_pruning_enforces_limit() {
1468        let map = SccHashMap::new();
1469        let _ = map.insert("b".to_string(), 2.0);
1470        let _ = map.insert("a".to_string(), 1.0);
1471        let _ = map.insert("c".to_string(), 3.0);
1472
1473        prune_string_f64_map(&map, 2);
1474
1475        assert!(map.read("a", |_, _| ()).is_none());
1476        assert!(map.read("b", |_, _| ()).is_some());
1477        assert!(map.read("c", |_, _| ()).is_some());
1478    }
1479}