1use 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::output_replay::page_limits;
25use crate::stream::Subscription;
26use crate::ResourceLimitDetails;
27
28#[derive(Debug, Clone)]
30pub(crate) enum ProcessOutcome {
31 Pending,
32 Failed { error: ClientError, rejected: bool },
33 Exited(i32),
34}
35
36impl ProcessOutcome {
37 pub(crate) fn completion(process_id: &str, exit_code: Option<i32>) -> Self {
38 match exit_code {
39 Some(code) => Self::Exited(code),
40 None => Self::Failed {
41 error: ClientError::TerminationFailed {
42 process_id: process_id.into(),
43 reason: "background execution completed without an exit status".into(),
44 },
45 rejected: false,
46 },
47 }
48 }
49
50 fn exit_code(&self) -> Option<i32> {
51 match self {
52 Self::Exited(code) => Some(*code),
53 _ => None,
54 }
55 }
56
57 fn rejected(&self) -> bool {
58 matches!(self, Self::Failed { rejected: true, .. })
59 }
60
61 fn reclaimable(&self) -> bool {
62 self.exit_code().is_some() || self.rejected()
63 }
64
65 fn wait_result(&self) -> Option<std::result::Result<i32, ClientError>> {
66 match self {
67 Self::Pending => None,
68 Self::Failed { error, .. } => Some(Err(error.clone())),
69 Self::Exited(code) => Some(Ok(*code)),
70 }
71 }
72
73 fn launch_failure(process_id: &str, error: ClientError) -> Self {
74 let rejected = matches!(
76 error,
77 ClientError::Kernel { .. } | ClientError::ResourceLimit { .. }
78 );
79 let error = if rejected {
80 error
81 } else {
82 ClientError::TerminationFailed {
83 process_id: process_id.into(),
84 reason: error.to_string(),
85 }
86 };
87 Self::Failed { error, rejected }
88 }
89}
90
91fn stdin_has_no_reader(pid: u32) -> ClientError {
93 ClientError::Kernel {
94 code: String::from("EPIPE"),
95 message: format!("process {pid} has exited, so its stdin has no reader"),
96 }
97}
98
99async fn wait_for_process_outcome(
100 mut rx: watch::Receiver<ProcessOutcome>,
101 process_id: &str,
102) -> std::result::Result<i32, ClientError> {
103 loop {
104 if let Some(result) = rx.borrow().wait_result() {
105 return result;
106 }
107 if rx.changed().await.is_err() {
108 return Err(ClientError::TerminationFailed {
109 process_id: process_id.into(),
110 reason: "exit observation channel closed before an exit event".into(),
111 });
112 }
113 }
114}
115
116async fn next_spawn_event(
117 events: &mut broadcast::Receiver<(wire::OwnershipScope, EventPayload)>,
118 ownership: &wire::OwnershipScope,
119 process_id: &str,
120 outcome: &watch::Sender<ProcessOutcome>,
121) -> Option<EventPayload> {
122 let mut replay_outcome = outcome.subscribe();
123 loop {
124 if matches!(*replay_outcome.borrow(), ProcessOutcome::Exited(_)) {
125 return None;
126 }
127 let event = tokio::select! {
128 event = events.recv() => event,
129 changed = replay_outcome.changed() => {
130 if let Err(error) = changed {
131 tracing::error!(%process_id, %error, "process outcome channel closed during observation");
132 return None;
133 }
134 continue;
135 }
136 };
137 match event {
138 Ok((scope, payload)) if scope == *ownership => {
139 if matches!(&payload, EventPayload::VmLifecycleEvent(event)
140 if matches!(event.state, wire::VmLifecycleState::Disposed | wire::VmLifecycleState::Failed))
141 {
142 let reason = "VM ended before a process exit was observed";
143 tracing::error!(%process_id, reason, "spawn observation failed");
144 outcome.send_if_modified(|state| {
145 if !matches!(state, ProcessOutcome::Pending) {
146 return false;
147 }
148 *state = ProcessOutcome::Failed {
149 error: ClientError::TerminationFailed {
150 process_id: process_id.into(),
151 reason: reason.into(),
152 },
153 rejected: false,
154 };
155 true
156 });
157 return None;
158 }
159 if let EventPayload::ProcessExitedEvent(exited) = &payload {
160 if exited.process_id == process_id {
161 outcome.send_replace(ProcessOutcome::Exited(exited.exit_code));
163 }
164 }
165 return Some(payload);
166 }
167 Ok(_) => continue,
168 Err(error) => {
169 let closed = matches!(error, broadcast::error::RecvError::Closed);
170 let reason =
171 format!("spawn event stream failed before exit was confirmed: {error}");
172 tracing::error!(%process_id, %reason, "spawn observation failed");
173 let error = if closed {
174 sidecar_exited_before_exit(process_id)
175 } else {
176 ClientError::TerminationFailed {
177 process_id: process_id.into(),
178 reason,
179 }
180 };
181 outcome.send_if_modified(|state| {
182 if !matches!(state, ProcessOutcome::Pending) {
183 return false;
184 }
185 *state = ProcessOutcome::Failed {
186 error,
187 rejected: false,
188 };
189 true
190 });
191 if closed {
192 return None;
193 }
194 }
196 }
197 }
198}
199
200pub(crate) fn sidecar_exited_before_exit(process_id: &str) -> ClientError {
203 ClientError::SidecarExited(format!("process {process_id} ended with its sidecar"))
204}
205
206fn reconcile_replayed_exit(outcome: &watch::Sender<ProcessOutcome>, exit_code: Option<i32>) {
207 if let Some(exit_code) = exit_code {
208 outcome.send_if_modified(|state| {
209 if matches!(state, ProcessOutcome::Exited(_)) {
210 return false;
211 }
212 *state = ProcessOutcome::Exited(exit_code);
213 true
214 });
215 }
216}
217
218const PROCESS_STREAM_CAPACITY: usize = 1024;
220
221const PROCESS_REGISTRY_LIMIT: usize = 1024;
223
224const OBSERVED_PROCESS_TIME_LIMIT: usize = 4096;
226
227const EXEC_OUTPUT_CAPTURE_LIMIT_BYTES: usize = 16 * 1024 * 1024;
229const EXEC_TERMINATION_CONFIRMATION_TIMEOUT: std::time::Duration =
230 std::time::Duration::from_secs(30);
231const EXEC_KILL_ACK_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(5);
232
233pub(crate) const DEFAULT_EXEC_CWD: &str = "/workspace";
235
236pub(crate) const SYNTHETIC_PID_BASE: u64 = 1_000_000;
239
240#[cfg_attr(feature = "contract", derive(ts_rs::TS))]
246#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
247#[serde(rename_all = "lowercase")]
248pub enum TimingMitigation {
249 #[default]
250 Off,
251 Freeze,
252}
253
254#[derive(Debug, Clone, PartialEq, Eq)]
256pub enum StdinInput {
257 Text(String),
258 Bytes(Vec<u8>),
259}
260
261pub type OutputCallback = Box<dyn FnMut(&[u8]) + Send>;
264
265struct ExecRunCleanup {
269 vm: AgentOs,
270 process_id: String,
271 launch: Option<JoinHandle<std::result::Result<wire::ProcessStartedResponse, ClientError>>>,
272 events: Option<broadcast::Receiver<(wire::OwnershipScope, EventPayload)>>,
273 armed: bool,
274}
275
276impl ExecRunCleanup {
277 fn events(&mut self) -> &mut broadcast::Receiver<(wire::OwnershipScope, EventPayload)> {
278 self.events.as_mut().expect("exec event receiver")
279 }
280
281 fn disarm(&mut self) {
282 self.armed = false;
283 }
284}
285
286impl Drop for ExecRunCleanup {
287 fn drop(&mut self) {
288 if !self.armed {
289 return;
290 }
291 let vm = self.vm.clone();
292 let process_id = self.process_id.clone();
293 let launch = self.launch.take();
294 let mut events = self.events.take().expect("exec event receiver");
295 let Ok(runtime) = tokio::runtime::Handle::try_current() else {
296 if let Some(launch) = launch {
297 launch.abort();
298 }
299 tracing::error!(%process_id, "cancelled exec has no runtime for cleanup; guest status is unknown");
300 return;
301 };
302 runtime.spawn(async move {
303 if let Some(mut launch) = launch {
304 match tokio::time::timeout(EXEC_TERMINATION_CONFIRMATION_TIMEOUT, &mut launch).await {
305 Ok(Ok(Ok(_))) => {}
306 Ok(Ok(Err(ClientError::Kernel { .. } | ClientError::ResourceLimit { .. }))) => {
307 return;
309 }
310 Ok(Ok(Err(error))) => {
311 tracing::warn!(?error, %process_id, "cancelled exec launch failed before cleanup");
312 }
313 Ok(Err(error)) => {
314 tracing::warn!(?error, %process_id, "cancelled exec launch task failed");
315 }
316 Err(_) => {
317 launch.abort();
320 tracing::error!(%process_id, "cancelled exec launch did not acknowledge within cleanup deadline; guest status is unknown");
321 }
322 }
323 }
324 if let Err(error) = vm.stop_and_confirm_exec(&process_id, &mut events).await {
325 tracing::error!(?error, %process_id, "cancelled exec could not confirm guest termination");
326 }
327 });
328 }
329}
330
331pub struct ExecOptions {
337 pub env: BTreeMap<String, String>,
338 pub cwd: Option<String>,
339 pub stdin: Option<StdinInput>,
340 pub timeout: Option<f64>,
341 pub on_stdout: Option<OutputCallback>,
342 pub on_stderr: Option<OutputCallback>,
343 pub capture_stdio: Option<bool>,
344 pub file_path: Option<String>,
345 pub cpu_time_limit_ms: Option<f64>,
346 pub timing_mitigation: Option<TimingMitigation>,
347}
348
349impl Default for ExecOptions {
350 fn default() -> Self {
351 Self {
352 env: BTreeMap::new(),
353 cwd: Some(DEFAULT_EXEC_CWD.to_string()),
354 stdin: None,
355 timeout: None,
356 on_stdout: None,
357 on_stderr: None,
358 capture_stdio: None,
359 file_path: None,
360 cpu_time_limit_ms: None,
361 timing_mitigation: None,
362 }
363 }
364}
365
366#[derive(Debug, Clone, PartialEq, Eq)]
368pub struct ExecResult {
369 pub exit_code: i32,
370 pub stdout: String,
371 pub stderr: String,
372}
373
374#[cfg_attr(feature = "contract", derive(ts_rs::TS))]
376#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
377#[serde(rename_all = "lowercase")]
378pub enum SpawnStdio {
379 #[default]
380 Pipe,
381 Inherit,
382}
383
384#[derive(Default)]
386pub struct SpawnOptions {
387 pub env: BTreeMap<String, String>,
388 pub cwd: Option<String>,
389 pub stdio: Option<SpawnStdio>,
390 pub stdin_fd: Option<i32>,
391 pub stdout_fd: Option<i32>,
392 pub stderr_fd: Option<i32>,
393 pub stream_stdin: Option<bool>,
396 pub retain_output: bool,
398}
399
400#[cfg_attr(feature = "contract", derive(ts_rs::TS))]
401#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
402#[serde(rename_all = "lowercase")]
403pub enum ProcessStream {
404 Stdout,
405 Stderr,
406}
407
408#[cfg_attr(feature = "contract", derive(ts_rs::TS))]
409#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
410pub struct ProcessOutput {
411 pub pid: u32,
412 pub stream: ProcessStream,
413 pub data: Vec<u8>,
414 pub sequence: Option<u64>,
415 #[serde(rename = "timestampMs")]
416 pub timestamp_ms: Option<i64>,
417}
418
419#[cfg_attr(feature = "contract", derive(ts_rs::TS))]
420#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
421pub struct ProcessOutputEvent {
422 pub pid: u32,
423 pub sequence: u64,
424 pub stream: ProcessStream,
425 pub data: Vec<u8>,
426 #[serde(rename = "timestampMs")]
427 pub timestamp_ms: i64,
428}
429
430#[cfg_attr(feature = "contract", derive(ts_rs::TS))]
431#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
432pub struct ProcessOutputReplay {
433 pub pid: u32,
434 pub events: Vec<ProcessOutputEvent>,
435 #[serde(rename = "nextCursor")]
436 pub next_cursor: Option<u64>,
437 #[serde(rename = "hasMore")]
438 pub has_more: bool,
439 pub truncated: bool,
440 #[serde(rename = "exitCode")]
442 pub exit_code: Option<i32>,
443}
444
445#[cfg_attr(feature = "contract", derive(ts_rs::TS))]
446#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
447pub struct ProcessExit {
448 pub pid: u32,
449 #[serde(rename = "exitCode")]
450 pub exit_code: i32,
451}
452
453#[cfg_attr(feature = "contract", derive(ts_rs::TS))]
455#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
456pub struct SpawnedProcessInfo {
457 pub pid: u32,
458 pub command: String,
459 pub args: Vec<String>,
460 pub running: bool,
461 #[serde(rename = "exitCode")]
462 pub exit_code: Option<i32>,
463 #[serde(rename = "startedAt")]
465 pub started_at: i64,
466}
467
468#[cfg_attr(feature = "contract", derive(ts_rs::TS))]
470#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
471pub struct SpawnHandle {
472 pub pid: u32,
473}
474
475#[cfg_attr(feature = "contract", derive(ts_rs::TS))]
477#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
478#[serde(rename_all = "lowercase")]
479pub enum ProcessStatus {
480 Running,
481 Exited,
482}
483
484#[cfg_attr(feature = "contract", derive(ts_rs::TS))]
486#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
487pub struct ProcessInfo {
488 pub pid: u32,
489 #[serde(skip)]
493 #[cfg_attr(feature = "contract", ts(skip))]
494 pub tracked_pid: Option<u32>,
495 pub ppid: u32,
496 pub pgid: u32,
497 pub sid: u32,
498 pub driver: String,
499 pub command: String,
500 pub args: Vec<String>,
501 pub cwd: String,
502 pub status: ProcessStatus,
503 #[serde(rename = "exitCode")]
504 pub exit_code: Option<i32>,
505 #[serde(rename = "startTime")]
506 pub start_time: f64,
507 #[serde(rename = "exitTime")]
508 pub exit_time: Option<f64>,
509}
510
511#[cfg_attr(feature = "contract", derive(ts_rs::TS))]
513#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
514pub struct ProcessTreeNode {
515 #[serde(flatten)]
516 pub info: ProcessInfo,
517 pub children: Vec<ProcessTreeNode>,
518}
519
520#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
521enum ProcessTreeKey {
522 Kernel(u32),
523 Tracked(u32),
524}
525
526struct ProcessTreeRecord {
527 info: ProcessInfo,
528 key: ProcessTreeKey,
529 parent: Option<ProcessTreeKey>,
530}
531
532impl AgentOs {
537 pub async fn exec_process(&self, command: &str, options: ExecOptions) -> Result<ExecResult> {
542 let (resolved_command, resolved_args) = resolve_exec_command(command)?;
546 self.exec_argv_process(&resolved_command, &resolved_args, options)
547 .await
548 }
549
550 pub async fn exec_argv_process(
556 &self,
557 command: &str,
558 args: &[String],
559 mut options: ExecOptions,
560 ) -> Result<ExecResult> {
561 let timeout_deadline = options
565 .timeout
566 .filter(|ms| ms.is_finite() && *ms >= 0.0)
567 .map(|ms| {
568 tokio::time::Instant::now() + std::time::Duration::from_secs_f64(ms / 1000.0)
569 });
570 let process_id = self.next_process_id();
571
572 let events = self.transport().subscribe_wire_events();
575
576 let launch_vm = self.clone();
577 let launch_id = process_id.clone();
578 let resolved_command = command.to_owned();
579 let resolved_args = args.to_vec();
580 let env = options.env.clone();
581 let cwd = options.cwd.clone();
582 let launch = tokio::spawn(async move {
583 launch_vm
584 .send_execute(
585 &launch_id,
586 Some(resolved_command),
587 resolved_args,
588 env,
589 cwd,
590 false,
591 )
592 .await
593 });
594 let mut cleanup = ExecRunCleanup {
595 vm: self.clone(),
596 process_id: process_id.clone(),
597 launch: Some(launch),
598 events: Some(events),
599 armed: true,
600 };
601 let launched = match timeout_deadline {
602 Some(deadline) => {
603 match tokio::time::timeout_at(
604 deadline,
605 cleanup.launch.as_mut().expect("launch task"),
606 )
607 .await
608 {
609 Ok(result) => result,
610 Err(_) => {
611 return Err(ClientError::TerminationFailed {
612 process_id,
613 reason: "Execute did not acknowledge before the execution deadline; cancellation cleanup is still attempting to stop it".to_owned(),
614 }
615 .into());
616 }
617 }
618 }
619 None => cleanup.launch.as_mut().expect("launch task").await,
620 };
621 cleanup.launch.take();
622 let started = match launched {
623 Ok(Ok(started)) => started,
624 Ok(Err(error @ (ClientError::Kernel { .. } | ClientError::ResourceLimit { .. }))) => {
625 cleanup.disarm();
628 return Err(error).context("exec: Execute request failed");
629 }
630 Ok(Err(error)) => return Err(error).context("exec: Execute request failed"),
631 Err(error) => {
632 return Err(
633 ClientError::Sidecar(format!("exec: launch task failed: {error}")).into(),
634 )
635 }
636 };
637 debug_assert_eq!(started.process_id, process_id);
638
639 let deliver_stdin = async {
642 if let Some(stdin) = options.stdin.take() {
643 let chunk = stdin_to_bytes(stdin);
644 if let Err(error) = self.write_wire_stdin(&process_id, chunk).await {
645 tracing::warn!(?error, %process_id, "exec stdin write failed");
646 }
647 }
648 if let Err(error) = self.close_wire_stdin(&process_id).await {
649 tracing::warn!(?error, %process_id, "exec stdin close failed");
650 }
651 };
652 if let Some(deadline) = timeout_deadline {
653 if tokio::time::timeout_at(deadline, deliver_stdin)
654 .await
655 .is_err()
656 {
657 let stopped = self
658 .stop_and_confirm_exec(&process_id, cleanup.events())
659 .await;
660 if stopped.is_ok() {
661 cleanup.disarm();
662 return Err(ClientError::ExecutionTimedOut { process_id }.into());
663 }
664 return Err(stopped.expect_err("failed confirmation").into());
665 }
666 } else {
667 deliver_stdin.await;
668 }
669
670 let mut on_stdout = options.on_stdout.take();
671 let mut on_stderr = options.on_stderr.take();
672
673 let capture_stdio = options.capture_stdio.unwrap_or(true);
674 let mut stdout = Vec::<u8>::new();
675 let mut stderr = Vec::<u8>::new();
676 let mut captured_output_bytes = 0usize;
677 let mut capture_error: Option<ClientError> = None;
678 let exit_code = loop {
679 let frame = match timeout_deadline {
680 Some(deadline) => {
681 match tokio::time::timeout_at(deadline, cleanup.events().recv()).await {
682 Ok(result) => result,
683 Err(_) => {
684 let stopped = self
685 .stop_and_confirm_exec(&process_id, cleanup.events())
686 .await;
687 if stopped.is_ok() {
688 cleanup.disarm();
689 return Err(ClientError::ExecutionTimedOut { process_id }.into());
690 }
691 return Err(stopped.expect_err("failed confirmation").into());
692 }
693 }
694 }
695 None => cleanup.events().recv().await,
696 };
697 let (ownership, payload) = match frame {
698 Ok(frame) => frame,
699 Err(broadcast::error::RecvError::Lagged(_)) => continue,
700 Err(broadcast::error::RecvError::Closed) => {
701 return Err(sidecar_exited_before_exit(&process_id).into());
702 }
703 };
704 if ownership != self.vm_scope() {
705 continue;
706 }
707 match payload {
708 EventPayload::ProcessOutputEvent(output) if output.process_id == process_id => {
709 match output.channel {
710 StreamChannel::Stdout => {
711 if let Some(cb) = on_stdout.as_mut() {
712 cb(&output.chunk);
713 }
714 if capture_stdio && capture_error.is_none() {
715 match append_exec_output(
716 &mut stdout,
717 &output.chunk,
718 &mut captured_output_bytes,
719 "stdout",
720 ) {
721 Ok(()) => {}
722 Err(error) => {
723 self.stop_and_confirm_exec(&process_id, cleanup.events())
724 .await?;
725 cleanup.disarm();
726 capture_error = Some(error);
727 break -1;
728 }
729 }
730 }
731 }
732 StreamChannel::Stderr => {
733 if let Some(cb) = on_stderr.as_mut() {
734 cb(&output.chunk);
735 }
736 if capture_stdio && capture_error.is_none() {
737 match append_exec_output(
738 &mut stderr,
739 &output.chunk,
740 &mut captured_output_bytes,
741 "stderr",
742 ) {
743 Ok(()) => {}
744 Err(error) => {
745 self.stop_and_confirm_exec(&process_id, cleanup.events())
746 .await?;
747 cleanup.disarm();
748 capture_error = Some(error);
749 break -1;
750 }
751 }
752 }
753 }
754 }
755 }
756 EventPayload::ProcessExitedEvent(exited) if exited.process_id == process_id => {
757 break exited.exit_code;
758 }
759 EventPayload::ProcessOutputEvent(_)
760 | EventPayload::ProcessExitedEvent(_)
761 | EventPayload::ExecutionOutputEvent(_)
762 | EventPayload::ExecutionCompletedEvent(_)
763 | EventPayload::VmLifecycleEvent(_)
764 | EventPayload::StructuredEvent(_)
765 | EventPayload::ExtEnvelope(_) => {}
766 }
767 };
768
769 cleanup.disarm();
770
771 if let Some(error) = capture_error {
772 return Err(error.into());
773 }
774
775 Ok(ExecResult {
776 exit_code,
777 stdout: String::from_utf8_lossy(&stdout).into_owned(),
778 stderr: String::from_utf8_lossy(&stderr).into_owned(),
779 })
780 }
781
782 async fn stop_and_confirm_exec(
785 &self,
786 process_id: &str,
787 events: &mut broadcast::Receiver<(wire::OwnershipScope, EventPayload)>,
788 ) -> std::result::Result<i32, ClientError> {
789 confirm_exec_exit(
790 process_id,
791 &self.vm_scope(),
792 events,
793 self.signal_wire_process(process_id, "SIGKILL"),
794 EXEC_KILL_ACK_TIMEOUT,
795 EXEC_TERMINATION_CONFIRMATION_TIMEOUT,
796 )
797 .await
798 }
799
800 pub fn spawn_process(
804 &self,
805 command: &str,
806 args: Vec<String>,
807 options: SpawnOptions,
808 ) -> Result<SpawnHandle> {
809 let reservation = self.reserve_process_registry_slot()?;
810
811 let pid = self
814 .inner()
815 .synthetic_pid_counter
816 .fetch_add(1, Ordering::SeqCst) as u32;
817 let process_id = format!("proc-{pid}-{}", uuid::Uuid::new_v4());
818
819 let (stdout_tx, _) = broadcast::channel::<Vec<u8>>(PROCESS_STREAM_CAPACITY);
820 let (stderr_tx, _) = broadcast::channel::<Vec<u8>>(PROCESS_STREAM_CAPACITY);
821 let (output_tx, _) = broadcast::channel::<ProcessOutput>(PROCESS_STREAM_CAPACITY);
822 let (exit_tx, _) = watch::channel(ProcessOutcome::Pending);
823 let (kernel_pid_tx, _) = watch::channel::<Option<u32>>(None);
826 let entry = ProcessEntry {
827 command: command.to_owned(),
828 args: args.clone(),
829 stdout_tx: stdout_tx.clone(),
830 stderr_tx: stderr_tx.clone(),
831 output_tx: output_tx.clone(),
832 exit_tx: exit_tx.clone(),
833 process_id: process_id.clone(),
834 kernel_pid: kernel_pid_tx.clone(),
835 output_tasks: Vec::new(),
836 retain_output: options.retain_output,
837 execution_id: None,
838 execution_generation: None,
839 started_at: epoch_ms_now() as i64,
840 };
841 reservation.commit(pid, entry)?;
844
845 let events = self.transport().subscribe_wire_events();
847
848 let this = self.clone();
849 let command = command.to_owned();
850 tokio::spawn(async move {
851 this.run_spawn(
852 pid,
853 process_id,
854 command,
855 args,
856 options,
857 events,
858 stdout_tx,
859 stderr_tx,
860 output_tx,
861 exit_tx,
862 kernel_pid_tx,
863 )
864 .await;
865 });
866
867 Ok(SpawnHandle { pid })
868 }
869
870 pub fn write_process_stdin(
872 &self,
873 pid: u32,
874 data: StdinInput,
875 ) -> std::result::Result<(), ClientError> {
876 let Some(process_id) = self.live_process_id(pid)? else {
877 return Err(stdin_has_no_reader(pid));
878 };
879 let chunk: Vec<u8> = stdin_to_bytes(data);
880 let this = self.clone();
881 tokio::spawn(async move {
882 if let Err(error) = this.write_wire_stdin(&process_id, chunk).await {
883 tracing::warn!(?error, pid, "write_process_stdin failed");
884 }
885 });
886 Ok(())
887 }
888
889 pub async fn write_process_stdin_awaited(
891 &self,
892 pid: u32,
893 data: StdinInput,
894 ) -> std::result::Result<(), ClientError> {
895 let Some(process_id) = self.live_process_id(pid)? else {
896 return Err(stdin_has_no_reader(pid));
897 };
898 let result = self
899 .write_wire_stdin(&process_id, stdin_to_bytes(data))
900 .await;
901 match result {
902 Err(_) if self.live_process_id(pid)?.is_none() => Err(stdin_has_no_reader(pid)),
904 result => result,
905 }
906 }
907
908 pub fn close_process_stdin(&self, pid: u32) -> std::result::Result<(), ClientError> {
910 let Some(process_id) = self.live_process_id(pid)? else {
911 return Ok(());
912 };
913 let this = self.clone();
914 tokio::spawn(async move {
915 if let Err(error) = this.close_wire_stdin(&process_id).await {
916 tracing::warn!(?error, pid, "close_process_stdin failed");
917 }
918 });
919 Ok(())
920 }
921
922 pub async fn close_process_stdin_awaited(
924 &self,
925 pid: u32,
926 ) -> std::result::Result<(), ClientError> {
927 let Some(process_id) = self.live_process_id(pid)? else {
928 return Ok(());
929 };
930 self.close_wire_stdin(&process_id).await
931 }
932
933 pub fn on_process_output(
935 &self,
936 pid: u32,
937 mut handler: impl FnMut(ProcessOutput) + Send + 'static,
938 ) -> std::result::Result<Subscription, ClientError> {
939 let mut rx = self
940 .inner()
941 .processes
942 .read(&pid, |_, entry| entry.output_tx.subscribe())
943 .ok_or(ClientError::ProcessNotFound(pid))?;
944 let task = tokio::spawn(async move {
945 loop {
946 match rx.recv().await {
947 Ok(event) => handler(event),
948 Err(broadcast::error::RecvError::Lagged(skipped)) => {
949 tracing::warn!(
950 pid,
951 skipped,
952 "process output subscriber lagged; recover with read_process_output"
953 );
954 }
955 Err(broadcast::error::RecvError::Closed) => return,
956 }
957 }
958 });
959 Ok(Subscription::new(move || task.abort()))
960 }
961
962 pub fn on_process_exit(
967 &self,
968 pid: u32,
969 handler: impl FnOnce(ProcessExit) + Send + 'static,
970 ) -> std::result::Result<Subscription, ClientError> {
971 let mut rx = self
972 .inner()
973 .processes
974 .read(&pid, |_, entry| entry.exit_tx.subscribe())
975 .ok_or(ClientError::ProcessNotFound(pid))?;
976
977 if let Some(code) = rx.borrow().exit_code() {
979 handler(ProcessExit {
980 pid,
981 exit_code: code,
982 });
983 return Ok(Subscription::noop());
984 }
985
986 if rx.borrow().rejected() {
987 return Err(rx
988 .borrow()
989 .wait_result()
990 .expect("rejected launch")
991 .unwrap_err());
992 }
993 let task = tokio::spawn(async move {
996 while rx.changed().await.is_ok() {
997 if let Some(code) = rx.borrow().exit_code() {
998 handler(ProcessExit {
999 pid,
1000 exit_code: code,
1001 });
1002 return;
1003 }
1004 if rx.borrow().rejected() {
1005 return;
1006 }
1007 }
1008 });
1009 Ok(Subscription::new(move || task.abort()))
1010 }
1011
1012 pub async fn wait_process(&self, pid: u32) -> std::result::Result<i32, ClientError> {
1015 let (rx, process_id) = self
1016 .inner()
1017 .processes
1018 .read(&pid, |_, entry| {
1019 (entry.exit_tx.subscribe(), entry.process_id.clone())
1020 })
1021 .ok_or(ClientError::ProcessNotFound(pid))?;
1022
1023 wait_for_process_outcome(rx, &process_id).await
1024 }
1025
1026 pub async fn read_process_output(
1028 &self,
1029 pid: u32,
1030 after: Option<u64>,
1031 max_events: Option<usize>,
1032 max_bytes: Option<usize>,
1033 ) -> std::result::Result<ProcessOutputReplay, ClientError> {
1034 let (
1035 process_id,
1036 retain_output,
1037 execution_id,
1038 execution_generation,
1039 mut started,
1040 mut outcome,
1041 ) = self
1042 .inner()
1043 .processes
1044 .read(&pid, |_, entry| {
1045 (
1046 entry.process_id.clone(),
1047 entry.retain_output,
1048 entry.execution_id.clone(),
1049 entry.execution_generation,
1050 entry.kernel_pid.subscribe(),
1051 entry.exit_tx.subscribe(),
1052 )
1053 })
1054 .ok_or(ClientError::ProcessNotFound(pid))?;
1055 loop {
1058 if started.borrow().is_some() {
1059 break;
1060 }
1061 match &*outcome.borrow() {
1062 ProcessOutcome::Failed { error, .. } => return Err(error.clone()),
1063 ProcessOutcome::Exited(_) => break,
1064 ProcessOutcome::Pending => {}
1065 }
1066 tokio::select! {
1067 changed = started.changed() => if changed.is_err() { break },
1068 changed = outcome.changed() => if changed.is_err() { break },
1069 }
1070 }
1071 if !retain_output {
1072 return Err(ClientError::Sidecar(format!(
1073 "process {pid} was not spawned with output retention enabled"
1074 )));
1075 }
1076 if let Some(execution_id) = execution_id {
1077 let (max_events, max_bytes) =
1078 page_limits(max_events, max_bytes, "process.output.read")?;
1079 let generation = execution_generation.ok_or_else(|| {
1080 ClientError::Sidecar(format!(
1081 "language process {pid} is missing its execution generation"
1082 ))
1083 })?;
1084 let response = self
1085 .transport()
1086 .request_wire(
1087 self.vm_scope(),
1088 wire::RequestPayload::ReadExecutionOutputRequest(
1089 wire::ReadExecutionOutputRequest {
1090 execution_id,
1091 cursor: after
1092 .map(|cursor| format!("{generation}:{}", cursor.saturating_add(1))),
1093 limit: Some(u32::try_from(max_events).map_err(|_| {
1094 ClientError::Sidecar(String::from(
1095 "process.output.read maxEvents exceeds the wire u32 range",
1096 ))
1097 })?),
1098 },
1099 ),
1100 )
1101 .await?;
1102 let page = match response {
1103 wire::ResponsePayload::ExecutionOutputPageResponse(page) => page,
1104 wire::ResponsePayload::RejectedResponse(rejected) => {
1105 return Err(ClientError::from_rejection(rejected));
1106 }
1107 other => {
1108 return Err(ClientError::Sidecar(format!(
1109 "ReadExecutionOutput: unexpected response {other:?}"
1110 )));
1111 }
1112 };
1113 let mut bytes = 0usize;
1114 let mut events = Vec::new();
1115 let mut has_more = page.has_more;
1116 for event in page.events {
1117 if bytes.saturating_add(event.chunk.len()) > max_bytes {
1118 if events.is_empty() {
1119 return Err(ClientError::ResourceLimit {
1120 code: String::from("ERR_AGENTOS_RESOURCE_LIMIT"),
1121 message: format!(
1122 "process.output.read next retained event requires {} bytes, exceeding maxBytes={max_bytes}",
1123 event.chunk.len()
1124 ),
1125 details: Box::new(ResourceLimitDetails {
1126 limit_name: Some(String::from("output_replay_page_bytes")),
1127 configured_limit: Some(max_bytes as u64),
1128 requested: Some(event.chunk.len() as u64),
1129 unit: Some(String::from("bytes")),
1130 scope: Some(String::from("vm")),
1131 operation: Some(String::from("process.output.read")),
1132 configuration_path: Some(String::from("maxBytes")),
1133 retryable: Some(true),
1134 ..ResourceLimitDetails::default()
1135 }),
1136 });
1137 }
1138 has_more = true;
1139 break;
1140 }
1141 bytes = bytes.saturating_add(event.chunk.len());
1142 events.push(ProcessOutputEvent {
1143 pid,
1144 sequence: event.sequence,
1145 stream: match event.channel {
1146 wire::ExecutionStreamChannel::Stdout => ProcessStream::Stdout,
1147 wire::ExecutionStreamChannel::Stderr
1148 | wire::ExecutionStreamChannel::Pty => ProcessStream::Stderr,
1149 },
1150 data: event.chunk,
1151 timestamp_ms: event.timestamp_ms.min(i64::MAX as u64) as i64,
1152 });
1153 }
1154 let next_cursor = events.last().map(|event| event.sequence).or(after);
1155 return Ok(ProcessOutputReplay {
1156 pid,
1157 events,
1158 next_cursor,
1159 has_more,
1160 truncated: page.truncated,
1161 exit_code: self.get_process(pid)?.exit_code,
1162 });
1163 }
1164 let (max_events, max_bytes) =
1165 crate::output_replay::wire_page_limits(max_events, max_bytes, "process.output.read")?;
1166 let response = self
1167 .transport()
1168 .request_wire(
1169 self.vm_scope(),
1170 wire::RequestPayload::ReadProcessOutputRequest(wire::ReadProcessOutputRequest {
1171 process_id: process_id.clone(),
1172 after,
1173 max_events: u32::try_from(max_events).map_err(|_| {
1174 ClientError::Sidecar(String::from(
1175 "process.output.read maxEvents exceeds the wire u32 range",
1176 ))
1177 })?,
1178 max_bytes: u32::try_from(max_bytes).map_err(|_| {
1179 ClientError::Sidecar(String::from(
1180 "process.output.read maxBytes exceeds the wire u32 range",
1181 ))
1182 })?,
1183 }),
1184 )
1185 .await?;
1186 let page = match response {
1187 wire::ResponsePayload::ProcessOutputPageResponse(page) => page,
1188 wire::ResponsePayload::RejectedResponse(rejected) => {
1189 return Err(ClientError::from_rejection(rejected));
1190 }
1191 other => {
1192 return Err(ClientError::Sidecar(format!(
1193 "ReadProcessOutput: unexpected response {other:?}"
1194 )));
1195 }
1196 };
1197 self.inner().processes.read(&pid, |_, entry| {
1201 if entry.process_id == process_id {
1202 reconcile_replayed_exit(&entry.exit_tx, page.exit_code);
1203 }
1204 });
1205 Ok(ProcessOutputReplay {
1206 pid,
1207 events: page
1208 .events
1209 .into_iter()
1210 .map(|event| ProcessOutputEvent {
1211 pid,
1212 sequence: event.sequence,
1213 stream: match event.channel {
1214 wire::StreamChannel::Stdout => ProcessStream::Stdout,
1215 wire::StreamChannel::Stderr => ProcessStream::Stderr,
1216 },
1217 data: event.chunk,
1218 timestamp_ms: event.timestamp_ms.min(i64::MAX as u64) as i64,
1219 })
1220 .collect(),
1221 next_cursor: page.next_cursor,
1222 has_more: page.has_more,
1223 truncated: page.truncated,
1224 exit_code: page.exit_code,
1225 })
1226 }
1227
1228 pub fn list_processes(&self) -> Vec<SpawnedProcessInfo> {
1230 let mut out = Vec::new();
1231 self.inner().processes.scan(|pid, entry| {
1232 let exit_code = entry.exit_tx.borrow().exit_code();
1233 out.push(SpawnedProcessInfo {
1234 pid: *pid,
1235 command: entry.command.clone(),
1236 args: entry.args.clone(),
1237 running: exit_code.is_none(),
1238 exit_code,
1239 started_at: entry.started_at,
1240 });
1241 });
1242 out
1243 }
1244
1245 pub async fn all_processes(&self) -> Result<Vec<ProcessInfo>> {
1253 Ok(self
1254 .process_tree_records()
1255 .await?
1256 .into_iter()
1257 .map(|record| record.info)
1258 .collect())
1259 }
1260
1261 async fn process_tree_records(&self) -> Result<Vec<ProcessTreeRecord>> {
1262 let ownership = self.vm_scope();
1263 let response = self
1264 .transport()
1265 .request_wire(ownership, wire::RequestPayload::GetProcessSnapshotRequest)
1266 .await
1267 .context("all_processes: GetProcessSnapshot request failed")?;
1268 let snapshot = match response {
1269 wire::ResponsePayload::ProcessSnapshotResponse(snapshot) => snapshot,
1270 wire::ResponsePayload::RejectedResponse(rejected) => {
1271 return Err(ClientError::from_rejection(rejected).into());
1272 }
1273 other => {
1274 return Err(ClientError::Sidecar(format!(
1275 "all_processes: unexpected response {other:?}"
1276 ))
1277 .into());
1278 }
1279 };
1280
1281 struct Tracked {
1285 display_pid: u32,
1286 exit_code: Option<i32>,
1287 command: String,
1288 args: Vec<String>,
1289 }
1290 let mut tracked_by_process_id: BTreeMap<String, Tracked> = BTreeMap::new();
1291 let mut display_pid_by_kernel_pid: BTreeMap<u32, u32> = BTreeMap::new();
1292 self.inner().processes.scan(|display_pid, entry| {
1293 let exit_code = entry.exit_tx.borrow().exit_code();
1294 if let Some(kernel_pid) = *entry.kernel_pid.borrow() {
1295 display_pid_by_kernel_pid.insert(kernel_pid, *display_pid);
1296 }
1297 tracked_by_process_id.insert(
1298 entry.process_id.clone(),
1299 Tracked {
1300 display_pid: *display_pid,
1301 exit_code,
1302 command: entry.command.clone(),
1303 args: entry.args.clone(),
1304 },
1305 );
1306 });
1307
1308 let now_ms = epoch_ms_now();
1309 let mut seen_process_ids: std::collections::BTreeSet<String> =
1310 std::collections::BTreeSet::new();
1311 let mut out: Vec<ProcessTreeRecord> = Vec::new();
1312
1313 for entry in &snapshot.processes {
1316 if let Some(tracked) = tracked_by_process_id.get(&entry.process_id) {
1317 display_pid_by_kernel_pid.insert(entry.pid, tracked.display_pid);
1318 }
1319 }
1320
1321 for entry in snapshot.processes {
1322 let tracked = tracked_by_process_id.get(&entry.process_id);
1323 let display_pid = tracked.map_or(entry.pid, |process| process.display_pid);
1324 let display_ppid = display_pid_by_kernel_pid
1325 .get(&entry.ppid)
1326 .copied()
1327 .unwrap_or(entry.ppid);
1328 let display_pgid = display_pid_by_kernel_pid
1329 .get(&entry.pgid)
1330 .copied()
1331 .unwrap_or(entry.pgid);
1332 let display_sid = display_pid_by_kernel_pid
1333 .get(&entry.sid)
1334 .copied()
1335 .unwrap_or(entry.sid);
1336
1337 let process_key = format!("{}:{}", entry.process_id, entry.pid);
1339 let start_time = self.observed_start_time(&process_key, now_ms);
1340
1341 let (status, exit_code) = match tracked {
1344 Some(t) => match t.exit_code {
1345 Some(code) => (ProcessStatus::Exited, Some(code)),
1346 None => (ProcessStatus::Running, entry.exit_code),
1347 },
1348 None => {
1349 let status = match entry.status {
1350 ProcessSnapshotStatus::Running | ProcessSnapshotStatus::Stopped => {
1351 ProcessStatus::Running
1352 }
1353 ProcessSnapshotStatus::Exited => ProcessStatus::Exited,
1354 };
1355 (status, entry.exit_code)
1356 }
1357 };
1358
1359 let exit_time = match (tracked, status) {
1361 (Some(_), ProcessStatus::Exited) => {
1362 Some(self.observed_exit_time(&entry.process_id, now_ms))
1363 }
1364 _ => None,
1365 };
1366
1367 let (command, args) = match tracked {
1368 Some(t) => (t.command.clone(), t.args.clone()),
1369 None => (entry.command, entry.args),
1370 };
1371
1372 seen_process_ids.insert(entry.process_id.clone());
1373 out.push(ProcessTreeRecord {
1374 key: ProcessTreeKey::Kernel(entry.pid),
1375 parent: Some(ProcessTreeKey::Kernel(entry.ppid)),
1376 info: ProcessInfo {
1377 pid: display_pid,
1378 tracked_pid: tracked.map(|process| process.display_pid),
1379 ppid: display_ppid,
1380 pgid: display_pgid,
1381 sid: display_sid,
1382 driver: entry.driver,
1383 command,
1384 args,
1385 cwd: entry.cwd,
1386 status,
1387 exit_code,
1388 start_time,
1389 exit_time,
1390 },
1391 });
1392 }
1393
1394 self.inner().processes.scan(|display_pid, entry| {
1397 if seen_process_ids.contains(&entry.process_id) {
1398 return;
1399 }
1400 let exit_code = entry.exit_tx.borrow().exit_code();
1401 let process_key = format!("{}:{}", entry.process_id, display_pid);
1402 let start_time = self.observed_start_time(&process_key, now_ms);
1403 let (status, exit_time) = match exit_code {
1404 Some(_) => (
1405 ProcessStatus::Exited,
1406 Some(self.observed_exit_time(&entry.process_id, now_ms)),
1407 ),
1408 None => (ProcessStatus::Running, None),
1409 };
1410 out.push(ProcessTreeRecord {
1411 key: ProcessTreeKey::Tracked(*display_pid),
1412 parent: None,
1413 info: ProcessInfo {
1414 pid: *display_pid,
1415 tracked_pid: Some(*display_pid),
1416 ppid: 0,
1417 pgid: *display_pid,
1418 sid: *display_pid,
1419 driver: String::new(),
1420 command: entry.command.clone(),
1421 args: entry.args.clone(),
1422 cwd: String::new(),
1423 status,
1424 exit_code,
1425 start_time,
1426 exit_time,
1427 },
1428 });
1429 });
1430
1431 out.sort_by_key(|record| record.info.pid);
1432 Ok(out)
1433 }
1434
1435 fn observed_start_time(&self, process_key: &str, now_ms: f64) -> f64 {
1438 let _guard = self.inner().observed_process_time_lock.lock();
1439 if let Some(existing) = self
1440 .inner()
1441 .observed_process_start_times
1442 .read(process_key, |_, value| *value)
1443 {
1444 return existing;
1445 }
1446 let _ = self
1447 .inner()
1448 .observed_process_start_times
1449 .insert(process_key.to_owned(), now_ms);
1450 prune_string_f64_map(
1451 &self.inner().observed_process_start_times,
1452 OBSERVED_PROCESS_TIME_LIMIT,
1453 );
1454 self.inner()
1457 .observed_process_start_times
1458 .read(process_key, |_, value| *value)
1459 .unwrap_or(now_ms)
1460 }
1461
1462 fn observed_exit_time(&self, process_id: &str, now_ms: f64) -> f64 {
1464 let _guard = self.inner().observed_process_time_lock.lock();
1465 if let Some(existing) = self
1466 .inner()
1467 .observed_process_exit_times
1468 .read(process_id, |_, value| *value)
1469 {
1470 return existing;
1471 }
1472 let _ = self
1473 .inner()
1474 .observed_process_exit_times
1475 .insert(process_id.to_owned(), now_ms);
1476 prune_string_f64_map(
1477 &self.inner().observed_process_exit_times,
1478 OBSERVED_PROCESS_TIME_LIMIT,
1479 );
1480 self.inner()
1481 .observed_process_exit_times
1482 .read(process_id, |_, value| *value)
1483 .unwrap_or(now_ms)
1484 }
1485
1486 pub async fn process_tree(&self) -> Result<Vec<ProcessTreeNode>> {
1488 let processes = self.process_tree_records().await?;
1489 Ok(build_process_forest(processes))
1490 }
1491
1492 pub fn get_process(&self, pid: u32) -> std::result::Result<SpawnedProcessInfo, ClientError> {
1494 self.inner()
1495 .processes
1496 .read(&pid, |pid, entry| {
1497 let exit_code = entry.exit_tx.borrow().exit_code();
1498 SpawnedProcessInfo {
1499 pid: *pid,
1500 command: entry.command.clone(),
1501 args: entry.args.clone(),
1502 running: exit_code.is_none(),
1503 exit_code,
1504 started_at: entry.started_at,
1505 }
1506 })
1507 .ok_or(ClientError::ProcessNotFound(pid))
1508 }
1509
1510 pub fn stop_process(&self, pid: u32) -> std::result::Result<(), ClientError> {
1512 self.signal_process(pid, "SIGTERM")
1513 }
1514
1515 pub fn kill_process(&self, pid: u32) -> std::result::Result<(), ClientError> {
1517 self.signal_process(pid, "SIGKILL")
1518 }
1519
1520 pub async fn signal_process_awaited(
1522 &self,
1523 pid: u32,
1524 signal: &str,
1525 ) -> std::result::Result<(), ClientError> {
1526 let (process_id, already_exited) = self
1527 .inner()
1528 .processes
1529 .read(&pid, |_, entry| {
1530 (
1531 entry.process_id.clone(),
1532 entry.exit_tx.borrow().reclaimable(),
1533 )
1534 })
1535 .ok_or(ClientError::ProcessNotFound(pid))?;
1536 if already_exited {
1537 return Ok(());
1538 }
1539 self.signal_wire_process(&process_id, signal).await
1540 }
1541
1542 pub async fn resize_process_pty_awaited(
1544 &self,
1545 pid: u32,
1546 cols: u16,
1547 rows: u16,
1548 ) -> std::result::Result<(), ClientError> {
1549 let process_id = self.lookup_process_id(pid)?;
1550 let response = self
1551 .transport()
1552 .request_wire(
1553 self.vm_scope(),
1554 wire::RequestPayload::ResizePtyRequest(wire::ResizePtyRequest {
1555 process_id,
1556 cols,
1557 rows,
1558 }),
1559 )
1560 .await?;
1561 match response {
1562 wire::ResponsePayload::PtyResizedResponse(_) => Ok(()),
1563 wire::ResponsePayload::RejectedResponse(rejected) => {
1564 Err(ClientError::from_rejection(rejected))
1565 }
1566 other => Err(ClientError::Sidecar(format!(
1567 "resize process PTY: unexpected response {other:?}"
1568 ))),
1569 }
1570 }
1571
1572 pub(crate) fn vm_scope(&self) -> wire::OwnershipScope {
1578 wire::OwnershipScope::VmOwnership(wire::VmOwnership {
1579 connection_id: self.connection_id().to_string(),
1580 session_id: self.wire_session_id().to_string(),
1581 vm_id: self.vm_id().to_string(),
1582 })
1583 }
1584
1585 fn next_process_id(&self) -> String {
1587 let n = self.inner().process_counter.fetch_add(1, Ordering::SeqCst);
1588 format!("proc-{n}-{}", uuid::Uuid::new_v4())
1589 }
1590
1591 fn live_process_id(&self, pid: u32) -> std::result::Result<Option<String>, ClientError> {
1594 self.inner()
1595 .processes
1596 .read(&pid, |_, entry| {
1597 entry
1598 .exit_tx
1599 .borrow()
1600 .exit_code()
1601 .is_none()
1602 .then(|| entry.process_id.clone())
1603 })
1604 .ok_or(ClientError::ProcessNotFound(pid))
1605 }
1606
1607 fn lookup_process_id(&self, pid: u32) -> std::result::Result<String, ClientError> {
1609 self.inner()
1610 .processes
1611 .read(&pid, |_, entry| entry.process_id.clone())
1612 .ok_or(ClientError::ProcessNotFound(pid))
1613 }
1614
1615 async fn send_execute(
1617 &self,
1618 process_id: &str,
1619 command: Option<String>,
1620 args: Vec<String>,
1621 env: BTreeMap<String, String>,
1622 cwd: Option<String>,
1623 retain_output: bool,
1624 ) -> std::result::Result<wire::ProcessStartedResponse, ClientError> {
1625 let ownership = self.vm_scope();
1626 let response = self
1627 .transport()
1628 .request_wire(
1629 ownership,
1630 wire::RequestPayload::ExecuteRequest(wire::ExecuteRequest {
1631 process_id: process_id.to_owned(),
1632 command,
1633 runtime: None,
1634 entrypoint: None,
1635 args,
1636 env: env.into_iter().collect(),
1637 cwd,
1638 wasm_permission_tier: None,
1639 retain_output,
1640 }),
1641 )
1642 .await?;
1643 match response {
1644 wire::ResponsePayload::ProcessStartedResponse(started) => Ok(started),
1645 wire::ResponsePayload::RejectedResponse(rejected) => {
1646 Err(ClientError::from_rejection(rejected))
1647 }
1648 other => Err(ClientError::Sidecar(format!(
1649 "Execute: unexpected response {other:?}"
1650 ))),
1651 }
1652 }
1653
1654 fn signal_process(&self, pid: u32, signal: &str) -> std::result::Result<(), ClientError> {
1657 let (process_id, already_exited) = self
1658 .inner()
1659 .processes
1660 .read(&pid, |_, entry| {
1661 (
1662 entry.process_id.clone(),
1663 entry.exit_tx.borrow().reclaimable(),
1664 )
1665 })
1666 .ok_or(ClientError::ProcessNotFound(pid))?;
1667 if already_exited {
1668 return Ok(());
1669 }
1670 let signal = signal.to_owned();
1671 let this = self.clone();
1672 tokio::spawn(async move {
1673 if let Err(error) = this.signal_wire_process(&process_id, &signal).await {
1674 tracing::warn!(?error, pid, %signal, "signal_process failed");
1675 }
1676 });
1677 Ok(())
1678 }
1679
1680 async fn write_wire_stdin(
1681 &self,
1682 process_id: &str,
1683 chunk: Vec<u8>,
1684 ) -> std::result::Result<(), ClientError> {
1685 let response = self
1686 .transport()
1687 .request_wire(
1688 self.vm_scope(),
1689 wire::RequestPayload::WriteStdinRequest(wire::WriteStdinRequest {
1690 process_id: process_id.to_owned(),
1691 chunk,
1692 }),
1693 )
1694 .await?;
1695 match response {
1696 wire::ResponsePayload::StdinWrittenResponse(_) => Ok(()),
1697 wire::ResponsePayload::RejectedResponse(rejected) => {
1698 Err(ClientError::from_rejection(rejected))
1699 }
1700 other => Err(ClientError::Sidecar(format!(
1701 "write process stdin: unexpected response {other:?}"
1702 ))),
1703 }
1704 }
1705
1706 async fn close_wire_stdin(&self, process_id: &str) -> std::result::Result<(), ClientError> {
1707 let response = self
1708 .transport()
1709 .request_wire(
1710 self.vm_scope(),
1711 wire::RequestPayload::CloseStdinRequest(wire::CloseStdinRequest {
1712 process_id: process_id.to_owned(),
1713 }),
1714 )
1715 .await?;
1716 match response {
1717 wire::ResponsePayload::StdinClosedResponse(_) => Ok(()),
1718 wire::ResponsePayload::RejectedResponse(rejected) => {
1719 Err(ClientError::from_rejection(rejected))
1720 }
1721 other => Err(ClientError::Sidecar(format!(
1722 "close process stdin: unexpected response {other:?}"
1723 ))),
1724 }
1725 }
1726
1727 async fn signal_wire_process(
1728 &self,
1729 process_id: &str,
1730 signal: &str,
1731 ) -> std::result::Result<(), ClientError> {
1732 let response = self
1733 .transport()
1734 .request_wire(
1735 self.vm_scope(),
1736 wire::RequestPayload::KillProcessRequest(wire::KillProcessRequest {
1737 process_id: process_id.to_owned(),
1738 signal: signal.to_owned(),
1739 }),
1740 )
1741 .await?;
1742 match response {
1743 wire::ResponsePayload::ProcessKilledResponse(_) => Ok(()),
1744 wire::ResponsePayload::RejectedResponse(rejected) => {
1745 Err(ClientError::from_rejection(rejected))
1746 }
1747 other => Err(ClientError::Sidecar(format!(
1748 "signal process: unexpected response {other:?}"
1749 ))),
1750 }
1751 }
1752
1753 fn process_registry_len_locked(&self) -> usize {
1754 let mut count = 0usize;
1755 self.inner().processes.scan(|_, _| {
1756 count += 1;
1757 });
1758 count
1759 }
1760
1761 pub(crate) fn reserve_process_registry_slot(
1765 &self,
1766 ) -> std::result::Result<ProcessRegistryReservation, ClientError> {
1767 let _guard = self.inner().process_registry_lock.lock();
1768 let pending = self
1769 .inner()
1770 .pending_process_registrations
1771 .load(Ordering::SeqCst);
1772 self.prune_exited_processes_locked(pending.saturating_add(1));
1773 let retained = self.process_registry_len_locked();
1774 if retained.saturating_add(pending) >= PROCESS_REGISTRY_LIMIT {
1775 return Err(process_registry_limit_error(
1776 retained.saturating_add(pending).saturating_add(1),
1777 ));
1778 }
1779 self.inner()
1780 .pending_process_registrations
1781 .fetch_add(1, Ordering::SeqCst);
1782 let admitted = retained.saturating_add(pending).saturating_add(1);
1783 if admitted >= PROCESS_REGISTRY_LIMIT * 4 / 5 {
1784 tracing::warn!(
1785 admitted,
1786 limit = PROCESS_REGISTRY_LIMIT,
1787 configuration_path = "PROCESS_REGISTRY_LIMIT",
1788 "process registry admission approaches its configured limit"
1789 );
1790 }
1791 Ok(ProcessRegistryReservation {
1792 client: self.clone(),
1793 active: true,
1794 })
1795 }
1796
1797 fn prune_exited_processes_locked(&self, reserve_slots: usize) {
1798 let mut entries = Vec::new();
1799 self.inner().processes.scan(|pid, entry| {
1800 entries.push((*pid, entry.exit_tx.borrow().reclaimable(), entry.started_at));
1801 });
1802 let target_len = PROCESS_REGISTRY_LIMIT.saturating_sub(reserve_slots);
1803 if entries.len() <= target_len {
1804 return;
1805 }
1806
1807 for pid in exited_pids_to_prune(entries, target_len) {
1808 self.remove_process_tracking_locked(pid);
1809 }
1810 }
1811
1812 fn remove_process_tracking_locked(&self, pid: u32) {
1813 if let Some((_, entry)) = self.inner().processes.remove(&pid) {
1814 let _time_guard = self.inner().observed_process_time_lock.lock();
1815 let _ = self
1816 .inner()
1817 .observed_process_exit_times
1818 .remove(&entry.process_id);
1819 let fallback_start_key = format!("{}:{pid}", entry.process_id);
1820 let _ = self
1821 .inner()
1822 .observed_process_start_times
1823 .remove(&fallback_start_key);
1824 if let Some(kernel_pid) = *entry.kernel_pid.borrow() {
1825 let start_key = format!("{}:{kernel_pid}", entry.process_id);
1826 let _ = self.inner().observed_process_start_times.remove(&start_key);
1827 }
1828 }
1829 }
1830
1831 #[allow(clippy::too_many_arguments)]
1836 async fn run_spawn(
1837 self,
1838 pid: u32,
1839 process_id: String,
1840 command: String,
1841 args: Vec<String>,
1842 options: SpawnOptions,
1843 mut events: broadcast::Receiver<(wire::OwnershipScope, EventPayload)>,
1844 stdout_tx: broadcast::Sender<Vec<u8>>,
1845 stderr_tx: broadcast::Sender<Vec<u8>>,
1846 output_tx: broadcast::Sender<ProcessOutput>,
1847 exit_tx: watch::Sender<ProcessOutcome>,
1848 kernel_pid_tx: watch::Sender<Option<u32>>,
1849 ) {
1850 match self
1851 .send_execute(
1852 &process_id,
1853 Some(command),
1854 args,
1855 options.env.clone(),
1856 options.cwd.clone(),
1857 options.retain_output,
1858 )
1859 .await
1860 {
1861 Ok(started) => {
1862 if let Some(kernel_pid) = started.pid {
1865 kernel_pid_tx.send_replace(Some(kernel_pid));
1866 }
1867 if options.stream_stdin == Some(false) {
1869 if let Err(error) = self.close_wire_stdin(&process_id).await {
1870 tracing::warn!(?error, pid, %process_id, "spawn: closing stdin failed");
1871 }
1872 }
1873 }
1874 Err(error) => {
1875 let message = format!("{error}\n");
1877 let bytes = message.into_bytes();
1878 let _ = stderr_tx.send(bytes.clone());
1879 let _ = output_tx.send(ProcessOutput {
1880 pid,
1881 stream: ProcessStream::Stderr,
1882 data: bytes,
1883 sequence: None,
1884 timestamp_ms: None,
1885 });
1886 tracing::error!(?error, pid, %process_id, "spawn: Execute request failed");
1887 let failure = ProcessOutcome::launch_failure(&process_id, error);
1888 let rejected = failure.rejected();
1889 exit_tx.send_replace(failure);
1890 if rejected {
1891 let _guard = self.inner().process_registry_lock.lock();
1892 self.prune_exited_processes_locked(0);
1893 return;
1894 }
1895 }
1898 }
1899
1900 let ownership = self.vm_scope();
1901 while let Some(payload) =
1902 next_spawn_event(&mut events, &ownership, &process_id, &exit_tx).await
1903 {
1904 match payload {
1905 EventPayload::ProcessOutputEvent(output) if output.process_id == process_id => {
1906 let bytes = output.chunk;
1907 let stream = match output.channel {
1908 StreamChannel::Stdout => ProcessStream::Stdout,
1909 StreamChannel::Stderr => ProcessStream::Stderr,
1910 };
1911 let _ = output_tx.send(ProcessOutput {
1912 pid,
1913 stream,
1914 data: bytes.clone(),
1915 sequence: output.sequence,
1916 timestamp_ms: output
1917 .timestamp_ms
1918 .map(|value| value.min(i64::MAX as u64) as i64),
1919 });
1920 match output.channel {
1921 StreamChannel::Stdout => {
1922 let _ = stdout_tx.send(bytes);
1923 }
1924 StreamChannel::Stderr => {
1925 let _ = stderr_tx.send(bytes);
1926 }
1927 }
1928 }
1929 EventPayload::ProcessExitedEvent(exited) if exited.process_id == process_id => {
1930 break;
1931 }
1932 EventPayload::ProcessOutputEvent(_)
1933 | EventPayload::ProcessExitedEvent(_)
1934 | EventPayload::ExecutionOutputEvent(_)
1935 | EventPayload::ExecutionCompletedEvent(_)
1936 | EventPayload::VmLifecycleEvent(_)
1937 | EventPayload::StructuredEvent(_)
1938 | EventPayload::ExtEnvelope(_) => {}
1939 }
1940 }
1941 let _guard = self.inner().process_registry_lock.lock();
1942 self.prune_exited_processes_locked(0);
1943 }
1944}
1945
1946pub(crate) struct ProcessRegistryReservation {
1947 client: AgentOs,
1948 active: bool,
1949}
1950
1951impl ProcessRegistryReservation {
1952 pub(crate) fn commit(
1953 mut self,
1954 pid: u32,
1955 entry: ProcessEntry,
1956 ) -> std::result::Result<(), ClientError> {
1957 let client = self.client.clone();
1958 let _guard = client.inner().process_registry_lock.lock();
1959 client.inner().processes.insert(pid, entry).map_err(|_| {
1960 ClientError::Sidecar(format!(
1961 "process registry already contains public pid {pid}; retry the spawn"
1962 ))
1963 })?;
1964 client
1965 .inner()
1966 .pending_process_registrations
1967 .fetch_sub(1, Ordering::SeqCst);
1968 self.active = false;
1969 Ok(())
1970 }
1971
1972 fn release(&mut self) {
1973 if self.active {
1974 self.client
1975 .inner()
1976 .pending_process_registrations
1977 .fetch_sub(1, Ordering::SeqCst);
1978 self.active = false;
1979 }
1980 }
1981}
1982
1983impl Drop for ProcessRegistryReservation {
1984 fn drop(&mut self) {
1985 self.release();
1986 }
1987}
1988
1989fn process_registry_limit_error(requested: usize) -> ClientError {
1990 ClientError::ResourceLimit {
1991 code: String::from("ERR_AGENTOS_RESOURCE_LIMIT"),
1992 message: format!(
1993 "process registry limit {PROCESS_REGISTRY_LIMIT} reached; wait for an exited process to be evicted or raise PROCESS_REGISTRY_LIMIT"
1994 ),
1995 details: Box::new(crate::ResourceLimitDetails {
1996 limit_name: Some(String::from("process_registry_entries")),
1997 configured_limit: Some(PROCESS_REGISTRY_LIMIT as u64),
1998 requested: Some(requested as u64),
1999 unit: Some(String::from("processes")),
2000 scope: Some(String::from("vm")),
2001 operation: Some(String::from("process.spawn")),
2002 configuration_path: Some(String::from("PROCESS_REGISTRY_LIMIT")),
2003 retryable: Some(true),
2004 ..Default::default()
2005 }),
2006 }
2007}
2008
2009async fn confirm_exec_exit(
2012 process_id: &str,
2013 ownership: &wire::OwnershipScope,
2014 events: &mut broadcast::Receiver<(wire::OwnershipScope, EventPayload)>,
2015 kill: impl std::future::Future<Output = std::result::Result<(), ClientError>>,
2016 kill_ack_timeout: std::time::Duration,
2017 confirmation_timeout: std::time::Duration,
2018) -> std::result::Result<i32, ClientError> {
2019 let deadline = tokio::time::sleep(confirmation_timeout);
2020 let kill = tokio::time::timeout(kill_ack_timeout, kill);
2021 tokio::pin!(deadline, kill);
2022 let mut kill_finished = false;
2023 let mut kill_error = None;
2024 loop {
2025 let failure = tokio::select! {
2026 biased;
2027 _ = &mut deadline => Some(format!("no exit event within {}ms", confirmation_timeout.as_millis())),
2028 result = &mut kill, if !kill_finished => {
2029 kill_finished = true;
2030 kill_error = match result {
2031 Ok(Ok(())) => None,
2032 Ok(Err(error)) => Some(error.to_string()),
2033 Err(_) => Some(format!("SIGKILL did not acknowledge within {}ms", kill_ack_timeout.as_millis())),
2034 };
2035 None
2036 }
2037 frame = events.recv() => match frame {
2038 Ok((scope, EventPayload::ProcessExitedEvent(exited)))
2039 if scope == *ownership && exited.process_id == process_id => return Ok(exited.exit_code),
2040 Ok(_) => None,
2041 Err(broadcast::error::RecvError::Lagged(skipped)) => {
2042 tracing::warn!(%process_id, skipped, "exec exit confirmation receiver lagged");
2043 None
2044 }
2045 Err(broadcast::error::RecvError::Closed) => {
2046 let mut error = sidecar_exited_before_exit(process_id);
2047 if let (ClientError::SidecarExited(message), Some(kill_error)) = (&mut error, &kill_error) {
2048 message.push_str(&format!("; SIGKILL failed: {kill_error}"));
2049 }
2050 return Err(error);
2051 }
2052 },
2053 };
2054 if let Some(mut reason) = failure {
2055 if let Some(kill_error) = kill_error {
2056 reason.push_str(&format!("; SIGKILL failed: {kill_error}"));
2057 }
2058 return Err(ClientError::TerminationFailed {
2059 process_id: process_id.to_owned(),
2060 reason,
2061 });
2062 }
2063 }
2064}
2065
2066fn build_process_forest(processes: Vec<ProcessTreeRecord>) -> Vec<ProcessTreeNode> {
2069 use std::collections::BTreeMap as Map;
2070
2071 let pids: std::collections::BTreeSet<ProcessTreeKey> =
2072 processes.iter().map(|record| record.key).collect();
2073 let mut children_of: Map<ProcessTreeKey, Vec<usize>> = Map::new();
2074 let mut roots: Vec<usize> = Vec::new();
2075 for (index, proc) in processes.iter().enumerate() {
2076 if let Some(parent) = proc.parent.filter(|parent| pids.contains(parent)) {
2077 children_of.entry(parent).or_default().push(index);
2078 } else {
2079 roots.push(index);
2080 }
2081 }
2082
2083 fn build_node(
2084 index: usize,
2085 processes: &[ProcessTreeRecord],
2086 children_of: &Map<ProcessTreeKey, Vec<usize>>,
2087 seen: &mut std::collections::BTreeSet<usize>,
2088 ) -> ProcessTreeNode {
2089 let record = &processes[index];
2090 let info = record.info.clone();
2091 seen.insert(index);
2092 let child_indices: Vec<usize> = children_of
2093 .get(&record.key)
2094 .map(|indices| {
2095 indices
2096 .iter()
2097 .copied()
2098 .filter(|child_index| !seen.contains(child_index))
2099 .collect()
2100 })
2101 .unwrap_or_default();
2102 let children = child_indices
2103 .into_iter()
2104 .map(|child_index| build_node(child_index, processes, children_of, seen))
2105 .collect();
2106 ProcessTreeNode { info, children }
2107 }
2108
2109 let mut seen = std::collections::BTreeSet::new();
2110 roots
2111 .into_iter()
2112 .map(|index| build_node(index, &processes, &children_of, &mut seen))
2113 .collect()
2114}
2115
2116fn stdin_to_bytes(input: StdinInput) -> Vec<u8> {
2119 match input {
2120 StdinInput::Text(text) => text.into_bytes(),
2121 StdinInput::Bytes(bytes) => bytes,
2122 }
2123}
2124
2125fn append_exec_output(
2126 buffer: &mut Vec<u8>,
2127 chunk: &[u8],
2128 captured_output_bytes: &mut usize,
2129 channel: &str,
2130) -> std::result::Result<(), ClientError> {
2131 let next_total = captured_output_bytes
2132 .checked_add(chunk.len())
2133 .ok_or_else(|| exec_output_limit_error(channel, usize::MAX))?;
2134 if next_total > EXEC_OUTPUT_CAPTURE_LIMIT_BYTES {
2135 return Err(exec_output_limit_error(channel, next_total));
2136 }
2137 buffer.extend_from_slice(chunk);
2138 *captured_output_bytes = next_total;
2139 Ok(())
2140}
2141
2142fn exec_output_limit_error(channel: &str, size: usize) -> ClientError {
2143 ClientError::Sidecar(format!(
2144 "exec {channel} capture is {size} bytes, limit is {EXEC_OUTPUT_CAPTURE_LIMIT_BYTES}"
2145 ))
2146}
2147
2148fn exited_pids_to_prune(mut entries: Vec<(u32, bool, i64)>, target_len: usize) -> Vec<u32> {
2149 if entries.len() <= target_len {
2150 return Vec::new();
2151 }
2152 let mut remove_count = entries.len() - target_len;
2153 entries.sort_by_key(|(pid, _, started_at)| (*started_at, *pid));
2154 let mut out = Vec::new();
2155 for (pid, exited, _) in entries {
2156 if remove_count == 0 {
2157 break;
2158 }
2159 if !exited {
2160 continue;
2161 }
2162 out.push(pid);
2163 remove_count -= 1;
2164 }
2165 out
2166}
2167
2168fn prune_string_f64_map(map: &SccHashMap<String, f64>, limit: usize) {
2169 let mut keys = Vec::new();
2170 map.scan(|key, _| {
2171 keys.push(key.clone());
2172 });
2173 if keys.len() <= limit {
2174 return;
2175 }
2176 let remove_count = keys.len() - limit;
2177 keys.sort();
2178 for key in keys.into_iter().take(remove_count) {
2179 let _ = map.remove(&key);
2180 }
2181}
2182
2183pub(crate) fn install_output_callback(
2192 tx: broadcast::Sender<Vec<u8>>,
2193 mut callback: OutputCallback,
2194) -> JoinHandle<()> {
2195 let mut rx = tx.subscribe();
2196 tokio::spawn(async move {
2197 loop {
2198 match rx.recv().await {
2199 Ok(chunk) => callback(&chunk),
2200 Err(broadcast::error::RecvError::Lagged(_)) => continue,
2201 Err(broadcast::error::RecvError::Closed) => break,
2202 }
2203 }
2204 })
2205}
2206
2207pub(crate) fn drain_process_output_tasks(processes: &SccHashMap<u32, ProcessEntry>) {
2212 let mut tasks = Vec::new();
2213 processes.retain(|_, entry| {
2214 entry.exit_tx.send_if_modified(|state| {
2215 if !matches!(state, ProcessOutcome::Pending) { return false; }
2216 tracing::error!(process_id = %entry.process_id, "VM shutdown before process exit was observed");
2217 *state = ProcessOutcome::Failed {
2218 error: ClientError::TerminationFailed {
2219 process_id: entry.process_id.clone(),
2220 reason: "VM shutdown before process exit was observed".into(),
2221 },
2222 rejected: false,
2223 };
2224 true
2225 });
2226 tasks.append(&mut entry.output_tasks);
2227 false
2228 });
2229 for task in tasks {
2230 task.abort();
2231 }
2232}
2233
2234fn epoch_ms_now() -> f64 {
2236 use std::time::{SystemTime, UNIX_EPOCH};
2237 SystemTime::now()
2238 .duration_since(UNIX_EPOCH)
2239 .map(|d| d.as_secs_f64() * 1000.0)
2240 .unwrap_or(0.0)
2241}
2242
2243#[cfg(test)]
2244mod tests {
2245 use super::{
2246 append_exec_output, drain_process_output_tasks, exited_pids_to_prune,
2247 install_output_callback, process_registry_limit_error, prune_string_f64_map, ExecOptions,
2248 OutputCallback, DEFAULT_EXEC_CWD, EXEC_OUTPUT_CAPTURE_LIMIT_BYTES, PROCESS_REGISTRY_LIMIT,
2249 };
2250 use super::{confirm_exec_exit, wire, ClientError, EventPayload};
2251 use crate::agent_os::ProcessEntry;
2252 use scc::HashMap as SccHashMap;
2253 use tokio::sync::{broadcast, watch};
2254
2255 #[tokio::test]
2256 async fn spawn_launch_rejection_is_typed_and_retained_for_late_waiters() {
2257 let (outcome, _) = watch::channel(super::ProcessOutcome::Pending);
2258 outcome.send_replace(super::ProcessOutcome::launch_failure(
2259 "p",
2260 ClientError::Kernel {
2261 code: "ENOENT".into(),
2262 message: "command missing".into(),
2263 },
2264 ));
2265 assert_eq!(outcome.borrow().exit_code(), None);
2266 assert!(outcome.borrow().reclaimable());
2267 assert!(
2268 matches!(super::wait_for_process_outcome(outcome.subscribe(), "p").await,
2269 Err(ClientError::Kernel { code, .. }) if code == "ENOENT")
2270 );
2271 }
2272
2273 #[tokio::test]
2274 async fn spawn_ambiguous_failure_keeps_observing_a_vm_scoped_real_exit() {
2275 let (outcome, _) = watch::channel(super::ProcessOutcome::Pending);
2276 outcome.send_replace(super::ProcessOutcome::launch_failure(
2277 "p",
2278 ClientError::Sidecar("lost acknowledgement".into()),
2279 ));
2280 assert_eq!(outcome.borrow().exit_code(), None);
2281 assert!(!outcome.borrow().reclaimable());
2282 assert!(matches!(
2283 super::wait_for_process_outcome(outcome.subscribe(), "p").await,
2284 Err(ClientError::TerminationFailed { .. })
2285 ));
2286
2287 let (sender, mut events) = broadcast::channel(2);
2288 let (ownership, event) = exec_exit_frame("vm", "p");
2289 let (foreign_scope, _) = exec_exit_frame("other-vm", "p");
2290 sender
2291 .send((
2292 foreign_scope,
2293 EventPayload::ProcessExitedEvent(wire::ProcessExitedEvent {
2294 process_id: "p".into(),
2295 exit_code: 0,
2296 }),
2297 ))
2298 .unwrap();
2299 sender.send((ownership.clone(), event)).unwrap();
2300 let delivered = super::next_spawn_event(&mut events, &ownership, "p", &outcome).await;
2301 assert!(matches!(
2302 delivered,
2303 Some(EventPayload::ProcessExitedEvent(_))
2304 ));
2305 assert_eq!(
2307 super::wait_for_process_outcome(outcome.subscribe(), "p")
2308 .await
2309 .unwrap(),
2310 137
2311 );
2312 assert!(outcome.borrow().reclaimable());
2313 }
2314
2315 #[tokio::test]
2316 async fn replayed_exit_recovers_a_missed_event_and_releases_the_observer() {
2317 for failed in [false, true] {
2318 let initial = if failed {
2319 super::ProcessOutcome::launch_failure(
2320 "p",
2321 ClientError::Sidecar("lost exit observation".into()),
2322 )
2323 } else {
2324 super::ProcessOutcome::Pending
2325 };
2326 let (outcome, _) = watch::channel(initial);
2327 let (_sender, mut events) = broadcast::channel(1);
2328 let (ownership, _) = exec_exit_frame("vm", "p");
2329 let observation = super::next_spawn_event(&mut events, &ownership, "p", &outcome);
2330 tokio::pin!(observation);
2331 tokio::select! {
2332 _ = &mut observation => panic!("no exit has been observed"),
2333 _ = tokio::task::yield_now() => {}
2334 }
2335 super::reconcile_replayed_exit(&outcome, None);
2336 assert_eq!(outcome.borrow().exit_code(), None);
2337 super::reconcile_replayed_exit(&outcome, Some(7));
2338 assert!(
2339 tokio::time::timeout(std::time::Duration::from_secs(1), observation)
2340 .await
2341 .unwrap()
2342 .is_none()
2343 );
2344 assert_eq!(
2345 super::wait_for_process_outcome(outcome.subscribe(), "p")
2346 .await
2347 .unwrap(),
2348 7
2349 );
2350 let receiver = outcome.subscribe();
2352 super::reconcile_replayed_exit(&outcome, Some(7));
2353 assert!(!receiver.has_changed().unwrap());
2354 }
2355 }
2356
2357 #[tokio::test]
2358 async fn spawn_event_stream_closure_and_lag_never_manufacture_success() {
2359 for lagged in [false, true] {
2360 let (outcome, _) = watch::channel(super::ProcessOutcome::Pending);
2361 let (sender, mut events) = broadcast::channel(1);
2362 if lagged {
2363 sender.send(exec_exit_frame("other-vm", "p")).unwrap();
2364 sender.send(exec_exit_frame("other-vm", "p")).unwrap();
2365 }
2366 drop(sender);
2367 let (ownership, _) = exec_exit_frame("vm", "p");
2368 assert!(
2369 super::next_spawn_event(&mut events, &ownership, "p", &outcome)
2370 .await
2371 .is_none()
2372 );
2373 assert_eq!(outcome.borrow().exit_code(), None);
2374 assert!(!outcome.borrow().reclaimable());
2375 let result = super::wait_for_process_outcome(outcome.subscribe(), "p").await;
2377 if lagged {
2378 assert!(
2379 matches!(result, Err(ClientError::TerminationFailed { .. })),
2380 "{result:?}"
2381 );
2382 } else {
2383 assert!(
2384 matches!(result, Err(ClientError::SidecarExited(_))),
2385 "{result:?}"
2386 );
2387 }
2388 }
2389 }
2390
2391 #[tokio::test]
2392 async fn spawn_closed_outcome_channel_is_a_typed_observation_failure() {
2393 let (outcome, rx) = watch::channel(super::ProcessOutcome::Pending);
2394 drop(outcome);
2395 assert!(matches!(
2396 super::wait_for_process_outcome(rx, "p").await,
2397 Err(ClientError::TerminationFailed { .. })
2398 ));
2399 }
2400
2401 #[tokio::test]
2402 async fn spawn_vm_disposal_ends_observation_without_a_synthetic_exit() {
2403 let (outcome, _) = watch::channel(super::ProcessOutcome::Pending);
2404 let (sender, mut events) = broadcast::channel(1);
2405 let (ownership, _) = exec_exit_frame("vm", "p");
2406 sender
2407 .send((
2408 ownership.clone(),
2409 EventPayload::VmLifecycleEvent(wire::VmLifecycleEvent {
2410 state: wire::VmLifecycleState::Disposed,
2411 }),
2412 ))
2413 .unwrap();
2414 assert!(
2415 super::next_spawn_event(&mut events, &ownership, "p", &outcome)
2416 .await
2417 .is_none()
2418 );
2419 assert_eq!(outcome.borrow().exit_code(), None);
2420 assert!(matches!(
2421 super::wait_for_process_outcome(outcome.subscribe(), "p").await,
2422 Err(ClientError::TerminationFailed { .. })
2423 ));
2424 }
2425
2426 #[test]
2427 fn background_completion_without_status_remains_a_failure() {
2428 assert_eq!(
2429 super::ProcessOutcome::completion("p", Some(0)).exit_code(),
2430 Some(0)
2431 );
2432 let missing = super::ProcessOutcome::completion("p", None);
2433 assert_eq!(missing.exit_code(), None);
2434 assert!(matches!(
2435 missing.wait_result(),
2436 Some(Err(ClientError::TerminationFailed { .. }))
2437 ));
2438 }
2439
2440 fn exec_exit_frame(vm_id: &str, process_id: &str) -> (wire::OwnershipScope, EventPayload) {
2441 (
2442 wire::OwnershipScope::VmOwnership(wire::VmOwnership {
2443 connection_id: "connection".into(),
2444 session_id: "session".into(),
2445 vm_id: vm_id.into(),
2446 }),
2447 EventPayload::ProcessExitedEvent(wire::ProcessExitedEvent {
2448 process_id: process_id.into(),
2449 exit_code: 137,
2450 }),
2451 )
2452 }
2453
2454 #[tokio::test]
2455 async fn exec_confirmation_requires_matching_vm_and_process() {
2456 let (sender, mut events) = broadcast::channel(8);
2457 let (ownership, _) = exec_exit_frame("vm", "process");
2458 sender.send(exec_exit_frame("other-vm", "process")).unwrap();
2459 sender.send(exec_exit_frame("vm", "other-process")).unwrap();
2460 let error = confirm_exec_exit(
2461 "process",
2462 &ownership,
2463 &mut events,
2464 async { Ok(()) },
2465 std::time::Duration::from_millis(5),
2466 std::time::Duration::from_millis(20),
2467 )
2468 .await
2469 .unwrap_err();
2470 assert!(matches!(error, ClientError::TerminationFailed { .. }));
2471 }
2472
2473 #[tokio::test]
2474 async fn exec_confirmation_observes_exit_without_waiting_for_kill_ack() {
2475 let (sender, mut events) = broadcast::channel(8);
2476 let frame = exec_exit_frame("vm", "process");
2477 let ownership = frame.0.clone();
2478 sender.send(frame).unwrap();
2479 let result = tokio::time::timeout(
2480 std::time::Duration::from_millis(100),
2481 confirm_exec_exit(
2482 "process",
2483 &ownership,
2484 &mut events,
2485 std::future::pending(),
2486 std::time::Duration::from_secs(5),
2487 std::time::Duration::from_secs(30),
2488 ),
2489 )
2490 .await
2491 .expect("an exit event is proof even while kill acknowledgement is pending");
2492 assert_eq!(result.unwrap(), 137);
2493 }
2494
2495 #[tokio::test]
2496 async fn exec_confirmation_preserves_signal_failure_on_closed_stream() {
2497 let (sender, mut events) = broadcast::channel(8);
2498 drop(sender);
2499 let (ownership, _) = exec_exit_frame("vm", "process");
2500 let error = confirm_exec_exit(
2501 "process",
2502 &ownership,
2503 &mut events,
2504 async { Err(ClientError::Sidecar("kill rejection".into())) },
2505 std::time::Duration::from_millis(5),
2506 std::time::Duration::from_millis(20),
2507 )
2508 .await
2509 .unwrap_err();
2510 assert!(matches!(error, ClientError::SidecarExited(_)), "{error:?}");
2511 assert!(error.to_string().contains("kill rejection"));
2512 }
2513
2514 #[test]
2515 fn process_forest_separates_colliding_display_pids_and_fallbacks() {
2516 use super::{
2517 build_process_forest, ProcessInfo, ProcessStatus, ProcessTreeKey, ProcessTreeRecord,
2518 };
2519 let record = |pid, tracked_pid, key, parent, command: &str| ProcessTreeRecord {
2520 info: ProcessInfo {
2521 pid,
2522 tracked_pid,
2523 ppid: 0,
2524 pgid: pid,
2525 sid: pid,
2526 driver: "test".into(),
2527 command: command.into(),
2528 args: Vec::new(),
2529 cwd: "/".into(),
2530 status: ProcessStatus::Running,
2531 exit_code: None,
2532 start_time: 1.0,
2533 exit_time: None,
2534 },
2535 key,
2536 parent,
2537 };
2538 let roots = build_process_forest(vec![
2539 record(
2540 1_000_000,
2541 Some(1_000_000),
2542 ProcessTreeKey::Kernel(7),
2543 None,
2544 "tracked",
2545 ),
2546 record(
2547 1_000_000,
2548 None,
2549 ProcessTreeKey::Kernel(1_000_000),
2550 Some(ProcessTreeKey::Kernel(7)),
2551 "guest",
2552 ),
2553 record(
2554 8,
2555 None,
2556 ProcessTreeKey::Kernel(8),
2557 Some(ProcessTreeKey::Kernel(1_000_000)),
2558 "grandchild",
2559 ),
2560 record(
2561 8,
2562 Some(8),
2563 ProcessTreeKey::Tracked(8),
2564 None,
2565 "pending tracked",
2566 ),
2567 ]);
2568 assert_eq!(roots.len(), 2);
2569 assert_eq!(roots[0].info.command, "tracked");
2570 assert_eq!(roots[0].children.len(), 1);
2571 let child = &roots[0].children[0];
2572 assert_eq!(child.info.pid, roots[0].info.pid);
2573 assert_eq!(child.info.tracked_pid, None);
2574 assert_eq!(child.info.command, "guest");
2575 assert_eq!(child.children.len(), 1);
2576 assert_eq!(child.children[0].info.command, "grandchild");
2577 assert_eq!(roots[1].info.command, "pending tracked");
2578 assert!(roots[1].children.is_empty());
2579 }
2580
2581 #[tokio::test]
2586 async fn drain_process_output_tasks_clears_registry_and_aborts_tasks() {
2587 let processes: SccHashMap<u32, ProcessEntry> = SccHashMap::new();
2588
2589 let (stdout_tx, _) = broadcast::channel::<Vec<u8>>(8);
2590 let (stderr_tx, _) = broadcast::channel::<Vec<u8>>(8);
2591 let (output_tx, _) = broadcast::channel(8);
2592 let (exit_tx, _) = watch::channel(super::ProcessOutcome::Pending);
2593 let exit_rx = exit_tx.subscribe();
2594 let (kernel_pid_tx, _) = watch::channel::<Option<u32>>(None);
2595
2596 let task = tokio::spawn(async {
2599 loop {
2600 tokio::time::sleep(std::time::Duration::from_secs(3600)).await;
2601 }
2602 });
2603 let abort_handle = task.abort_handle();
2604
2605 let entry = ProcessEntry {
2606 command: "sleep".to_string(),
2607 args: vec!["3600".to_string()],
2608 stdout_tx,
2609 stderr_tx,
2610 output_tx,
2611 exit_tx,
2612 process_id: "proc-test".to_string(),
2613 kernel_pid: kernel_pid_tx,
2614 output_tasks: vec![task],
2615 retain_output: false,
2616 execution_id: None,
2617 execution_generation: None,
2618 started_at: 0,
2619 };
2620 let _ = processes.insert(1, entry);
2621
2622 assert!(!abort_handle.is_finished(), "task should start alive");
2623
2624 drain_process_output_tasks(&processes);
2625
2626 assert!(matches!(
2627 super::wait_for_process_outcome(exit_rx, "proc-test").await,
2628 Err(ClientError::TerminationFailed { .. })
2629 ));
2630
2631 assert!(processes.is_empty(), "registry must be cleared on drain");
2632
2633 for _ in 0..100 {
2635 if abort_handle.is_finished() {
2636 break;
2637 }
2638 tokio::time::sleep(std::time::Duration::from_millis(5)).await;
2639 }
2640 assert!(
2641 abort_handle.is_finished(),
2642 "output task must be aborted after drain"
2643 );
2644 }
2645
2646 #[tokio::test]
2652 async fn install_output_callback_handle_is_captured_into_process_entry() {
2653 use std::sync::atomic::{AtomicUsize, Ordering};
2654 use std::sync::Arc;
2655
2656 let (stdout_tx, _) = broadcast::channel::<Vec<u8>>(8);
2657 let (stderr_tx, _) = broadcast::channel::<Vec<u8>>(8);
2658 let (output_tx, _) = broadcast::channel(8);
2659 let (exit_tx, _) = watch::channel(super::ProcessOutcome::Pending);
2660 let (kernel_pid_tx, _) = watch::channel::<Option<u32>>(None);
2661
2662 let calls = Arc::new(AtomicUsize::new(0));
2663 let calls_cb = Arc::clone(&calls);
2664 let cb: OutputCallback = Box::new(move |_chunk: &[u8]| {
2665 calls_cb.fetch_add(1, Ordering::SeqCst);
2666 });
2667
2668 let output_tasks = vec![install_output_callback(stdout_tx.clone(), cb)];
2670
2671 let entry = ProcessEntry {
2672 command: "sleep".to_string(),
2673 args: vec!["3600".to_string()],
2674 stdout_tx: stdout_tx.clone(),
2675 stderr_tx,
2676 output_tx,
2677 exit_tx,
2678 process_id: "proc-test".to_string(),
2679 kernel_pid: kernel_pid_tx,
2680 output_tasks,
2681 retain_output: false,
2682 execution_id: None,
2683 execution_generation: None,
2684 started_at: 0,
2685 };
2686
2687 assert_eq!(
2688 entry.output_tasks.len(),
2689 1,
2690 "the install_output_callback handle must be captured on the entry"
2691 );
2692
2693 stdout_tx
2695 .send(b"hello".to_vec())
2696 .expect("broadcast send to subscribed callback task");
2697 for _ in 0..100 {
2698 if calls.load(Ordering::SeqCst) > 0 {
2699 break;
2700 }
2701 tokio::time::sleep(std::time::Duration::from_millis(5)).await;
2702 }
2703 assert_eq!(
2704 calls.load(Ordering::SeqCst),
2705 1,
2706 "the stored handle must drive the registered callback"
2707 );
2708
2709 let processes: SccHashMap<u32, ProcessEntry> = SccHashMap::new();
2711 let _ = processes.insert(1, entry);
2712 drain_process_output_tasks(&processes);
2713 assert!(processes.is_empty(), "registry must be cleared on drain");
2714 }
2715
2716 #[test]
2717 fn exec_options_default_uses_workspace_cwd() {
2718 assert_eq!(
2719 ExecOptions::default().cwd.as_deref(),
2720 Some(DEFAULT_EXEC_CWD)
2721 );
2722 }
2723
2724 #[test]
2725 fn append_exec_output_rejects_capture_over_limit() {
2726 let mut buffer = vec![0u8; EXEC_OUTPUT_CAPTURE_LIMIT_BYTES - 1];
2727 let mut captured = buffer.len();
2728
2729 append_exec_output(&mut buffer, &[1], &mut captured, "stdout")
2730 .expect("chunk at limit should fit");
2731 assert_eq!(captured, EXEC_OUTPUT_CAPTURE_LIMIT_BYTES);
2732
2733 let error = append_exec_output(&mut buffer, &[2], &mut captured, "stdout")
2734 .expect_err("chunk over limit should fail");
2735 assert!(
2736 error.to_string().contains("exec stdout capture is"),
2737 "unexpected error: {error}"
2738 );
2739 assert_eq!(captured, EXEC_OUTPUT_CAPTURE_LIMIT_BYTES);
2740 assert_eq!(buffer.len(), EXEC_OUTPUT_CAPTURE_LIMIT_BYTES);
2741 }
2742
2743 #[test]
2744 fn exited_pid_pruning_keeps_live_entries_and_removes_oldest_exited() {
2745 let pids = exited_pids_to_prune(
2746 vec![(3, true, 30), (1, false, 10), (2, true, 40), (4, true, 20)],
2747 2,
2748 );
2749 assert_eq!(pids, vec![4, 3]);
2750 }
2751
2752 #[test]
2753 fn process_registry_limit_error_is_typed_and_actionable() {
2754 let error = process_registry_limit_error(PROCESS_REGISTRY_LIMIT + 1);
2755 let ClientError::ResourceLimit {
2756 code,
2757 details,
2758 message,
2759 } = error
2760 else {
2761 panic!("expected typed resource limit");
2762 };
2763 assert_eq!(code, "ERR_AGENTOS_RESOURCE_LIMIT");
2764 assert_eq!(
2765 details.limit_name.as_deref(),
2766 Some("process_registry_entries")
2767 );
2768 assert_eq!(
2769 details.configured_limit,
2770 Some(PROCESS_REGISTRY_LIMIT as u64)
2771 );
2772 assert_eq!(details.requested, Some((PROCESS_REGISTRY_LIMIT + 1) as u64));
2773 assert_eq!(details.operation.as_deref(), Some("process.spawn"));
2774 assert_eq!(
2775 details.configuration_path.as_deref(),
2776 Some("PROCESS_REGISTRY_LIMIT")
2777 );
2778 assert!(message.contains("raise PROCESS_REGISTRY_LIMIT"));
2779 }
2780
2781 #[test]
2782 fn observed_time_pruning_enforces_limit() {
2783 let map = SccHashMap::new();
2784 let _ = map.insert("b".to_string(), 2.0);
2785 let _ = map.insert("a".to_string(), 1.0);
2786 let _ = map.insert("c".to_string(), 3.0);
2787
2788 prune_string_f64_map(&map, 2);
2789
2790 assert!(map.read("a", |_, _| ()).is_none());
2791 assert!(map.read("b", |_, _| ()).is_some());
2792 assert!(map.read("c", |_, _| ()).is_some());
2793 }
2794}