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