1use std::collections::{BTreeMap, HashMap, VecDeque};
9use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering};
10use std::sync::{Arc, Weak};
11use std::time::Duration;
12
13use scc::HashMap as SccHashMap;
14use serde::Deserialize;
15use serde_json::{Map, Value};
16use tokio::sync::{broadcast, watch};
17use tokio::task::JoinHandle;
18
19use agentos_sidecar_client::{wire, TransportError};
20use agentos_vm_config as vm_config;
21
22use crate::config::{
23 resolve_host_functions, AgentOsConfig, AgentOsLimits, MountConfig, ResolvedHostFunction,
24 ResolvedHostFunctions, RootFilesystemConfig, RootFilesystemKind,
25 RootFilesystemMode as ConfigRootFilesystemMode, RootLowerInput, SidecarJsBridgeCall,
26 SidecarJsBridgeCallback, SidecarSqliteCallback, TimerScheduleDriver,
27};
28use crate::cron::CronManager;
29use crate::error::ClientError;
30use crate::process::SYNTHETIC_PID_BASE;
31use crate::sidecar::{AgentOsSidecar, AgentOsSidecarPlacement, AgentOsSidecarVmLease};
32use crate::transport::{SidecarProcess, WireSidecarCallback};
33
34use once_cell::sync::OnceCell;
35
36pub(crate) struct ProcessEntry {
42 pub command: String,
43 pub args: Vec<String>,
44 #[allow(dead_code)]
45 pub stdout_tx: broadcast::Sender<Vec<u8>>,
46 #[allow(dead_code)]
47 pub stderr_tx: broadcast::Sender<Vec<u8>>,
48 pub output_tx: broadcast::Sender<crate::process::ProcessOutput>,
49 pub exit_tx: watch::Sender<crate::process::ProcessOutcome>,
51 pub process_id: String,
53 pub kernel_pid: watch::Sender<Option<u32>>,
57 pub output_tasks: Vec<JoinHandle<()>>,
61 pub retain_output: bool,
63 pub execution_id: Option<String>,
65 pub execution_generation: Option<u64>,
66 pub started_at: i64,
68}
69
70pub(crate) struct ShellEntry {
76 pub pid: u32,
77 pub data_tx: broadcast::Sender<Vec<u8>>,
78 pub stderr_tx: broadcast::Sender<Vec<u8>>,
79 pub event_tx: broadcast::Sender<crate::shell::TerminalOutputEvent>,
80 pub process_id: String,
82 pub spawned_tx: watch::Sender<Option<Result<(), ClientError>>>,
87 pub exit_tx: watch::Sender<Option<Result<i32, ClientError>>>,
90}
91
92#[derive(Clone)]
93pub(crate) struct ClosedShellEntry {
94 pub shell_id: String,
95 pub process_id: String,
96 pub pid: u32,
97 pub result: Result<i32, ClientError>,
98}
99
100pub(crate) struct TerminalEntry {
102 pub exit_task: JoinHandle<()>,
103}
104
105#[derive(Debug, Clone)]
115pub struct PackageDescriptor {
116 pub path: String,
117}
118
119#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
120pub struct SoftwareInfo {
121 #[serde(rename = "packageName")]
122 pub package_name: String,
123 pub commands: Vec<String>,
124}
125
126#[derive(Clone)]
128pub struct AgentOs {
129 inner: Arc<AgentOsInner>,
130}
131
132pub(crate) struct AgentOsInner {
133 pub(crate) transport: Arc<SidecarProcess>,
135 pub(crate) connection_id: String,
136 pub(crate) session_id: String,
137 pub(crate) vm_id: String,
138 pub(crate) projected_commands: parking_lot::Mutex<BTreeMap<String, String>>,
140 pub(crate) vm_configuration_operation: tokio::sync::Mutex<()>,
143 pub(crate) installed_software:
144 parking_lot::Mutex<BTreeMap<String, crate::software::InstalledSoftware>>,
145
146 pub(crate) process_registry_lock: parking_lot::Mutex<()>,
148 pub(crate) pending_process_registrations: AtomicUsize,
151 pub(crate) processes: SccHashMap<u32, ProcessEntry>,
152 pub(crate) process_counter: AtomicU64,
156 pub(crate) synthetic_pid_counter: AtomicU64,
159 pub(crate) observed_process_time_lock: parking_lot::Mutex<()>,
160 pub(crate) observed_process_start_times: SccHashMap<String, f64>,
164 pub(crate) observed_process_exit_times: SccHashMap<String, f64>,
167
168 pub(crate) shells: SccHashMap<String, ShellEntry>,
170 pub(crate) shell_counter: AtomicU64,
171 pub(crate) pending_shell_exits: SccHashMap<u64, JoinHandle<()>>,
172 pub(crate) closed_shells: parking_lot::Mutex<VecDeque<ClosedShellEntry>>,
176 pub(crate) terminals: SccHashMap<String, TerminalEntry>,
177 pub(crate) terminal_count: AtomicUsize,
178 pub(crate) terminal_lifecycle_lock: tokio::sync::Mutex<()>,
179
180 pub(crate) cron: Arc<CronManager>,
182
183 pub(crate) config: Arc<AgentOsConfig>,
185 pub(crate) sidecar: Arc<AgentOsSidecar>,
186 pub(crate) sidecar_lease: parking_lot::Mutex<Option<AgentOsSidecarVmLease>>,
187 pub(crate) dynamic_mounts: parking_lot::Mutex<Vec<wire::MountDescriptor>>,
188 pub(crate) disposed: AtomicBool,
189 shutdown_result: tokio::sync::Mutex<Option<Result<(), ClientError>>>,
190 vm_disposed: AtomicBool,
191}
192
193struct AgentOsCreationTask {
197 task: Option<tokio::task::JoinHandle<Result<AgentOs, ClientError>>>,
198}
199
200impl Drop for AgentOsCreationTask {
201 fn drop(&mut self) {
202 let Some(task) = self.task.take() else {
203 return;
204 };
205 let Ok(runtime) = tokio::runtime::Handle::try_current() else {
206 tracing::error!(
207 "VM creation waiter was dropped outside a Tokio runtime; process shutdown must reap the sidecar"
208 );
209 return;
210 };
211 runtime.spawn(async move {
212 match task.await {
213 Ok(Ok(vm)) => {
214 if let Err(error) = vm.shutdown().await {
215 tracing::error!(%error, "failed to dispose VM created after its caller was cancelled");
216 }
217 }
218 Ok(Err(error)) => {
219 tracing::debug!(%error, "cancelled VM creation finished with an error");
220 }
221 Err(error) => {
222 tracing::error!(%error, "cancelled VM creation task failed");
223 }
224 }
225 });
226 }
227}
228
229impl AgentOs {
230 pub async fn create(options: AgentOsConfig) -> Result<AgentOs, ClientError> {
234 let mut creation = AgentOsCreationTask {
235 task: Some(tokio::spawn(Self::create_owned(options))),
236 };
237 let result = creation
238 .task
239 .as_mut()
240 .expect("creation task is present")
241 .await;
242 creation.task.take();
245 match result {
246 Ok(result) => result,
247 Err(error) => Err(ClientError::Sidecar(format!(
248 "VM creation task failed: {error}"
249 ))),
250 }
251 }
252
253 async fn create_owned(options: AgentOsConfig) -> Result<AgentOs, ClientError> {
254 let config = Arc::new(options);
255
256 let sidecar = match &config.sidecar {
260 Some(crate::config::AgentOsSidecarConfig::Explicit { handle }) => handle.clone(),
261 Some(crate::config::AgentOsSidecarConfig::Shared { pool }) => {
262 AgentOs::get_shared_sidecar(pool.clone(), config.sidecar_binary_path.clone())
263 .await?
264 }
265 None => AgentOs::get_shared_sidecar(None, config.sidecar_binary_path.clone()).await?,
266 };
267 let mut lease = Some(sidecar.acquire_vm_lease().await?);
268 let mut created_vm = None;
269 let mut created_session_key = None;
270 let result = async {
271 let (transport, connection_id, _) = sidecar.ensure_connection().await?;
272
273 let session = match transport
275 .request_wire(
276 wire_connection_ownership(&connection_id),
277 wire::RequestPayload::OpenSessionRequest(wire::OpenSessionRequest {
278 placement: sidecar_wire_placement(&sidecar),
279 metadata: HashMap::new(),
280 }),
281 )
282 .await?
283 {
284 wire::ResponsePayload::SessionOpenedResponse(opened) => opened,
285 wire::ResponsePayload::RejectedResponse(rejected) => {
286 return Err(rejected_to_error(rejected));
287 }
288 wire::ResponsePayload::AuthenticatedResponse(_)
289 | wire::ResponsePayload::VmCreatedResponse(_)
290 | wire::ResponsePayload::VmDisposedResponse(_)
291 | wire::ResponsePayload::VmConfigComparedResponse(_)
292 | wire::ResponsePayload::RootFilesystemBootstrappedResponse(_)
293 | wire::ResponsePayload::VmConfiguredResponse(_)
294 | wire::ResponsePayload::HostCallbacksRegisteredResponse(_)
295 | wire::ResponsePayload::LayerCreatedResponse(_)
296 | wire::ResponsePayload::LayerSealedResponse(_)
297 | wire::ResponsePayload::SnapshotImportedResponse(_)
298 | wire::ResponsePayload::SnapshotExportedResponse(_)
299 | wire::ResponsePayload::OverlayCreatedResponse(_)
300 | wire::ResponsePayload::GuestFilesystemResultResponse(_)
301 | wire::ResponsePayload::RootFilesystemSnapshotResponse(_)
302 | wire::ResponsePayload::ProcessStartedResponse(_)
303 | wire::ResponsePayload::StdinWrittenResponse(_)
304 | wire::ResponsePayload::PtyResizedResponse(_)
305 | wire::ResponsePayload::StdinClosedResponse(_)
306 | wire::ResponsePayload::ProcessKilledResponse(_)
307 | wire::ResponsePayload::ProcessSnapshotResponse(_)
308 | wire::ResponsePayload::ListenerSnapshotResponse(_)
309 | wire::ResponsePayload::BoundUdpSnapshotResponse(_)
310 | wire::ResponsePayload::SignalStateResponse(_)
311 | wire::ResponsePayload::ZombieTimerCountResponse(_)
312 | wire::ResponsePayload::FilesystemResultResponse(_)
313 | wire::ResponsePayload::PermissionDecisionResponse(_)
314 | wire::ResponsePayload::PersistenceStateResponse(_)
315 | wire::ResponsePayload::PersistenceFlushedResponse(_)
316 | wire::ResponsePayload::VmFetchResponse(_)
317 | wire::ResponsePayload::ExtEnvelope(_)
318 | wire::ResponsePayload::GuestKernelResultResponse(_)
319 | wire::ResponsePayload::ResourceSnapshotResponse(_)
320 | wire::ResponsePayload::PackageLinkedResponse(_)
321 | wire::ResponsePayload::PackageUnlinkedResponse(_)
322 | wire::ResponsePayload::PackageAcquiredResponse(_)
323 | wire::ResponsePayload::PackageInstalledResponse(_)
324 | wire::ResponsePayload::PackageCacheStatsResponse(_)
325 | wire::ResponsePayload::ProvidedCommandsResponse(_)
326 | wire::ResponsePayload::ListMountsResponse(_)
327 | wire::ResponsePayload::ExecutionAcceptedResponse(_)
328 | wire::ResponsePayload::ExecutionCompletedResponse(_)
329 | wire::ResponsePayload::ExecutionEvaluationResponse(_)
330 | wire::ResponsePayload::TypeScriptCheckResponse(_)
331 | wire::ResponsePayload::ExecutionDescriptorResponse(_)
332 | wire::ResponsePayload::ExecutionListResponse(_)
333 | wire::ResponsePayload::ExecutionDeletedResponse(_)
334 | wire::ResponsePayload::ExecutionIoResponse(_)
335 | wire::ResponsePayload::ExecutionOutputPageResponse(_)
336 | wire::ResponsePayload::ProcessOutputPageResponse(_) => {
337 return Err(ClientError::Sidecar(
338 "unexpected open_session response".to_string(),
339 ));
340 }
341 };
342 let session_id = session.session_id;
343 created_session_key = Some(sidecar_session_key(&connection_id, &session_id));
344
345 let mut events = transport.subscribe_wire_events();
347 let create_vm_config = serialize_create_vm_config_for_sidecar(&config)?;
348 if let Some(callback) = config.sidecar_js_bridge_callback.clone() {
349 let _ = session_js_bridge_callbacks()
350 .insert(sidecar_session_key(&connection_id, &session_id), callback);
351 transport.register_wire_callback("js_bridge_call", js_bridge_call_callback());
352 }
353 if let Some(callback) = config.sidecar_sqlite_callback.clone() {
354 let _ = session_sqlite_callbacks()
355 .insert(sidecar_session_key(&connection_id, &session_id), callback);
356 transport.register_wire_callback("ext", sqlite_callback_callback());
357 }
358
359 let vm = match transport
361 .request_wire(
362 wire_session_ownership(&connection_id, &session_id),
363 wire::RequestPayload::CreateVmRequest(wire::CreateVmRequest {
364 runtime: wire::GuestRuntimeKind::JavaScript,
365 config: serde_json::to_string(&create_vm_config).map_err(|error| {
366 ClientError::Sidecar(format!(
367 "failed to serialize create VM config: {error}"
368 ))
369 })?,
370 }),
371 )
372 .await?
373 {
374 wire::ResponsePayload::VmCreatedResponse(created) => created,
375 wire::ResponsePayload::RejectedResponse(rejected) => {
376 return Err(rejected_to_error(rejected));
377 }
378 wire::ResponsePayload::AuthenticatedResponse(_)
379 | wire::ResponsePayload::SessionOpenedResponse(_)
380 | wire::ResponsePayload::VmDisposedResponse(_)
381 | wire::ResponsePayload::VmConfigComparedResponse(_)
382 | wire::ResponsePayload::RootFilesystemBootstrappedResponse(_)
383 | wire::ResponsePayload::VmConfiguredResponse(_)
384 | wire::ResponsePayload::HostCallbacksRegisteredResponse(_)
385 | wire::ResponsePayload::LayerCreatedResponse(_)
386 | wire::ResponsePayload::LayerSealedResponse(_)
387 | wire::ResponsePayload::SnapshotImportedResponse(_)
388 | wire::ResponsePayload::SnapshotExportedResponse(_)
389 | wire::ResponsePayload::OverlayCreatedResponse(_)
390 | wire::ResponsePayload::GuestFilesystemResultResponse(_)
391 | wire::ResponsePayload::RootFilesystemSnapshotResponse(_)
392 | wire::ResponsePayload::ProcessStartedResponse(_)
393 | wire::ResponsePayload::StdinWrittenResponse(_)
394 | wire::ResponsePayload::PtyResizedResponse(_)
395 | wire::ResponsePayload::StdinClosedResponse(_)
396 | wire::ResponsePayload::ProcessKilledResponse(_)
397 | wire::ResponsePayload::ProcessSnapshotResponse(_)
398 | wire::ResponsePayload::ListenerSnapshotResponse(_)
399 | wire::ResponsePayload::BoundUdpSnapshotResponse(_)
400 | wire::ResponsePayload::SignalStateResponse(_)
401 | wire::ResponsePayload::ZombieTimerCountResponse(_)
402 | wire::ResponsePayload::FilesystemResultResponse(_)
403 | wire::ResponsePayload::PermissionDecisionResponse(_)
404 | wire::ResponsePayload::PersistenceStateResponse(_)
405 | wire::ResponsePayload::PersistenceFlushedResponse(_)
406 | wire::ResponsePayload::VmFetchResponse(_)
407 | wire::ResponsePayload::ExtEnvelope(_)
408 | wire::ResponsePayload::GuestKernelResultResponse(_)
409 | wire::ResponsePayload::ResourceSnapshotResponse(_)
410 | wire::ResponsePayload::PackageLinkedResponse(_)
411 | wire::ResponsePayload::PackageUnlinkedResponse(_)
412 | wire::ResponsePayload::PackageAcquiredResponse(_)
413 | wire::ResponsePayload::PackageInstalledResponse(_)
414 | wire::ResponsePayload::PackageCacheStatsResponse(_)
415 | wire::ResponsePayload::ProvidedCommandsResponse(_)
416 | wire::ResponsePayload::ListMountsResponse(_)
417 | wire::ResponsePayload::ExecutionAcceptedResponse(_)
418 | wire::ResponsePayload::ExecutionCompletedResponse(_)
419 | wire::ResponsePayload::ExecutionEvaluationResponse(_)
420 | wire::ResponsePayload::TypeScriptCheckResponse(_)
421 | wire::ResponsePayload::ExecutionDescriptorResponse(_)
422 | wire::ResponsePayload::ExecutionListResponse(_)
423 | wire::ResponsePayload::ExecutionDeletedResponse(_)
424 | wire::ResponsePayload::ExecutionIoResponse(_)
425 | wire::ResponsePayload::ExecutionOutputPageResponse(_)
426 | wire::ResponsePayload::ProcessOutputPageResponse(_) => {
427 return Err(ClientError::Sidecar(
428 "unexpected create_vm response".to_string(),
429 ));
430 }
431 };
432 let vm_id = vm.vm_id;
433 created_vm = Some((
434 transport.clone(),
435 connection_id.clone(),
436 session_id.clone(),
437 vm_id.clone(),
438 ));
439
440 wait_for_vm_ready(&mut events, &vm_id, crate::VM_READY_TIMEOUT_MS).await?;
442
443 let packages = build_package_descriptors(&config);
446
447 let mounts = serialize_mounts(&config)?;
449 let configured_mounts = mounts.clone();
450
451 let projected_commands = match transport
455 .request_wire(
456 wire_vm_ownership(&connection_id, &session_id, &vm_id),
457 wire::RequestPayload::ConfigureVmRequest(wire::ConfigureVmRequest {
458 mounts,
459 software: Vec::new(),
462 permissions: None,
466 module_access_cwd: None,
470 instructions: Vec::new(),
471 projected_modules: Vec::new(),
472 command_permissions: HashMap::new(),
473 loopback_exempt_ports: config.loopback_exempt_ports.clone(),
474 packages,
475 packages_mount_at: config.packages_mount_at.clone().unwrap_or_default(),
476 bootstrap_commands: Vec::new(),
477 host_function_shim_commands: Vec::new(),
478 }),
479 )
480 .await?
481 {
482 wire::ResponsePayload::VmConfiguredResponse(configured) => configured
483 .projected_commands
484 .into_iter()
485 .map(|command| (command.name, command.guest_path))
486 .collect(),
487 wire::ResponsePayload::RejectedResponse(rejected) => {
488 return Err(rejected_to_error(rejected));
489 }
490 wire::ResponsePayload::AuthenticatedResponse(_)
491 | wire::ResponsePayload::SessionOpenedResponse(_)
492 | wire::ResponsePayload::VmCreatedResponse(_)
493 | wire::ResponsePayload::VmDisposedResponse(_)
494 | wire::ResponsePayload::VmConfigComparedResponse(_)
495 | wire::ResponsePayload::RootFilesystemBootstrappedResponse(_)
496 | wire::ResponsePayload::HostCallbacksRegisteredResponse(_)
497 | wire::ResponsePayload::LayerCreatedResponse(_)
498 | wire::ResponsePayload::LayerSealedResponse(_)
499 | wire::ResponsePayload::SnapshotImportedResponse(_)
500 | wire::ResponsePayload::SnapshotExportedResponse(_)
501 | wire::ResponsePayload::OverlayCreatedResponse(_)
502 | wire::ResponsePayload::GuestFilesystemResultResponse(_)
503 | wire::ResponsePayload::RootFilesystemSnapshotResponse(_)
504 | wire::ResponsePayload::ProcessStartedResponse(_)
505 | wire::ResponsePayload::StdinWrittenResponse(_)
506 | wire::ResponsePayload::PtyResizedResponse(_)
507 | wire::ResponsePayload::StdinClosedResponse(_)
508 | wire::ResponsePayload::ProcessKilledResponse(_)
509 | wire::ResponsePayload::ProcessSnapshotResponse(_)
510 | wire::ResponsePayload::ListenerSnapshotResponse(_)
511 | wire::ResponsePayload::BoundUdpSnapshotResponse(_)
512 | wire::ResponsePayload::SignalStateResponse(_)
513 | wire::ResponsePayload::ZombieTimerCountResponse(_)
514 | wire::ResponsePayload::FilesystemResultResponse(_)
515 | wire::ResponsePayload::PermissionDecisionResponse(_)
516 | wire::ResponsePayload::PersistenceStateResponse(_)
517 | wire::ResponsePayload::PersistenceFlushedResponse(_)
518 | wire::ResponsePayload::VmFetchResponse(_)
519 | wire::ResponsePayload::ExtEnvelope(_)
520 | wire::ResponsePayload::GuestKernelResultResponse(_)
521 | wire::ResponsePayload::ResourceSnapshotResponse(_)
522 | wire::ResponsePayload::PackageLinkedResponse(_)
523 | wire::ResponsePayload::PackageUnlinkedResponse(_)
524 | wire::ResponsePayload::PackageAcquiredResponse(_)
525 | wire::ResponsePayload::PackageInstalledResponse(_)
526 | wire::ResponsePayload::PackageCacheStatsResponse(_)
527 | wire::ResponsePayload::ProvidedCommandsResponse(_)
528 | wire::ResponsePayload::ListMountsResponse(_)
529 | wire::ResponsePayload::ExecutionAcceptedResponse(_)
530 | wire::ResponsePayload::ExecutionCompletedResponse(_)
531 | wire::ResponsePayload::ExecutionEvaluationResponse(_)
532 | wire::ResponsePayload::TypeScriptCheckResponse(_)
533 | wire::ResponsePayload::ExecutionDescriptorResponse(_)
534 | wire::ResponsePayload::ExecutionListResponse(_)
535 | wire::ResponsePayload::ExecutionDeletedResponse(_)
536 | wire::ResponsePayload::ExecutionIoResponse(_)
537 | wire::ResponsePayload::ExecutionOutputPageResponse(_)
538 | wire::ResponsePayload::ProcessOutputPageResponse(_) => {
539 return Err(ClientError::Sidecar(
540 "unexpected configure_vm response".to_string(),
541 ));
542 }
543 };
544
545 let resolved_host_functions = resolve_host_functions(&config.host_functions)
549 .map_err(ClientError::InvalidConfig)?;
550 if !resolved_host_functions.is_empty() {
551 let mut host_function_map: HashMap<String, ResolvedHostFunction> = HashMap::new();
552 for collection in &resolved_host_functions {
553 let mut host_functions = HashMap::new();
554 for host_function in &collection.functions {
555 host_functions.insert(
556 host_function.name.clone(),
557 wire::RegisteredHostCallbackDefinition {
558 description: host_function.description.clone(),
559 input_schema: json_utf8(
560 &host_function.input_schema,
561 "host callback input schema",
562 )?,
563 timeout_ms: host_function.timeout_ms,
564 examples: Vec::new(),
565 },
566 );
567 host_function_map.insert(
568 format!("{}:{}", collection.name, host_function.name),
569 host_function.clone(),
570 );
571 }
572 match transport
573 .request_wire(
574 wire_vm_ownership(&connection_id, &session_id, &vm_id),
575 wire::RequestPayload::RegisterHostCallbacksRequest(
576 wire::RegisterHostCallbacksRequest {
577 name: collection.name.clone(),
578 description: String::new(),
579 command_aliases: vec![format!("agentos-{}", collection.name)],
580 registry_command_aliases: vec![String::from("agentos")],
581 callbacks: host_functions,
582 },
583 ),
584 )
585 .await?
586 {
587 wire::ResponsePayload::HostCallbacksRegisteredResponse(_) => {}
588 wire::ResponsePayload::RejectedResponse(rejected) => {
589 return Err(rejected_to_error(rejected));
590 }
591 wire::ResponsePayload::AuthenticatedResponse(_)
592 | wire::ResponsePayload::SessionOpenedResponse(_)
593 | wire::ResponsePayload::VmCreatedResponse(_)
594 | wire::ResponsePayload::VmDisposedResponse(_)
595 | wire::ResponsePayload::VmConfigComparedResponse(_)
596 | wire::ResponsePayload::RootFilesystemBootstrappedResponse(_)
597 | wire::ResponsePayload::VmConfiguredResponse(_)
598 | wire::ResponsePayload::LayerCreatedResponse(_)
599 | wire::ResponsePayload::LayerSealedResponse(_)
600 | wire::ResponsePayload::SnapshotImportedResponse(_)
601 | wire::ResponsePayload::SnapshotExportedResponse(_)
602 | wire::ResponsePayload::OverlayCreatedResponse(_)
603 | wire::ResponsePayload::GuestFilesystemResultResponse(_)
604 | wire::ResponsePayload::RootFilesystemSnapshotResponse(_)
605 | wire::ResponsePayload::ProcessStartedResponse(_)
606 | wire::ResponsePayload::StdinWrittenResponse(_)
607 | wire::ResponsePayload::PtyResizedResponse(_)
608 | wire::ResponsePayload::StdinClosedResponse(_)
609 | wire::ResponsePayload::ProcessKilledResponse(_)
610 | wire::ResponsePayload::ProcessSnapshotResponse(_)
611 | wire::ResponsePayload::ListenerSnapshotResponse(_)
612 | wire::ResponsePayload::BoundUdpSnapshotResponse(_)
613 | wire::ResponsePayload::SignalStateResponse(_)
614 | wire::ResponsePayload::ZombieTimerCountResponse(_)
615 | wire::ResponsePayload::FilesystemResultResponse(_)
616 | wire::ResponsePayload::PermissionDecisionResponse(_)
617 | wire::ResponsePayload::PersistenceStateResponse(_)
618 | wire::ResponsePayload::PersistenceFlushedResponse(_)
619 | wire::ResponsePayload::VmFetchResponse(_)
620 | wire::ResponsePayload::ExtEnvelope(_)
621 | wire::ResponsePayload::GuestKernelResultResponse(_)
622 | wire::ResponsePayload::ResourceSnapshotResponse(_)
623 | wire::ResponsePayload::PackageLinkedResponse(_)
624 | wire::ResponsePayload::PackageUnlinkedResponse(_)
625 | wire::ResponsePayload::PackageAcquiredResponse(_)
626 | wire::ResponsePayload::PackageInstalledResponse(_)
627 | wire::ResponsePayload::PackageCacheStatsResponse(_)
628 | wire::ResponsePayload::ProvidedCommandsResponse(_)
629 | wire::ResponsePayload::ListMountsResponse(_)
630 | wire::ResponsePayload::ExecutionAcceptedResponse(_)
631 | wire::ResponsePayload::ExecutionCompletedResponse(_)
632 | wire::ResponsePayload::ExecutionEvaluationResponse(_)
633 | wire::ResponsePayload::TypeScriptCheckResponse(_)
634 | wire::ResponsePayload::ExecutionDescriptorResponse(_)
635 | wire::ResponsePayload::ExecutionListResponse(_)
636 | wire::ResponsePayload::ExecutionDeletedResponse(_)
637 | wire::ResponsePayload::ExecutionIoResponse(_)
638 | wire::ResponsePayload::ExecutionOutputPageResponse(_)
639 | wire::ResponsePayload::ProcessOutputPageResponse(_) => {
640 return Err(ClientError::Sidecar(
641 "unexpected register_host_callbacks response".to_string(),
642 ));
643 }
644 }
645 }
646 let _ = vm_host_functions().insert(
647 vm_id.clone(),
648 Arc::new(VmHostFunctionRegistry {
649 host_functions: resolved_host_functions.clone(),
650 host_function_map,
651 }),
652 );
653 transport.register_wire_callback("host_callback", host_callback_callback());
654 }
655
656 let driver = config
658 .schedule_driver
659 .clone()
660 .unwrap_or_else(|| Arc::new(TimerScheduleDriver::new()));
661 let cron = Arc::new(CronManager::new(driver));
662
663 let inner = AgentOsInner {
664 transport,
665 connection_id,
666 session_id,
667 vm_id,
668 projected_commands: parking_lot::Mutex::new(projected_commands),
669 vm_configuration_operation: tokio::sync::Mutex::new(()),
670 installed_software: parking_lot::Mutex::new(BTreeMap::new()),
671 process_registry_lock: parking_lot::Mutex::new(()),
672 pending_process_registrations: AtomicUsize::new(0),
673 processes: SccHashMap::new(),
674 process_counter: AtomicU64::new(1),
675 synthetic_pid_counter: AtomicU64::new(SYNTHETIC_PID_BASE),
676 observed_process_time_lock: parking_lot::Mutex::new(()),
677 observed_process_start_times: SccHashMap::new(),
678 observed_process_exit_times: SccHashMap::new(),
679 shells: SccHashMap::new(),
680 shell_counter: AtomicU64::new(0),
681 pending_shell_exits: SccHashMap::new(),
682 closed_shells: parking_lot::Mutex::new(VecDeque::new()),
683 terminals: SccHashMap::new(),
684 terminal_count: AtomicUsize::new(0),
685 terminal_lifecycle_lock: tokio::sync::Mutex::new(()),
686 cron,
687 config,
688 sidecar: sidecar.clone(),
689 sidecar_lease: parking_lot::Mutex::new(lease.take()),
690 dynamic_mounts: parking_lot::Mutex::new(configured_mounts),
691 disposed: AtomicBool::new(false),
692 shutdown_result: tokio::sync::Mutex::new(None),
693 vm_disposed: AtomicBool::new(false),
694 };
695
696 let client = AgentOs {
697 inner: Arc::new(inner),
698 };
699 let _ = vm_clients().insert(client.inner.vm_id.clone(), Arc::downgrade(&client.inner));
703 Ok(client)
704 }
705 .await;
706 if result.is_err() {
707 let mut cleanup_confirmed = true;
708 if let Some((transport, connection_id, session_id, vm_id)) = created_vm {
709 let cleanup = transport
710 .request_wire(
711 wire_vm_ownership(&connection_id, &session_id, &vm_id),
712 wire::RequestPayload::DisposeVmRequest(wire::DisposeVmRequest {
713 reason: wire::DisposeReason::Requested,
714 }),
715 )
716 .await
717 .map_err(ClientError::from)
718 .and_then(|response| vm_dispose_response(&vm_id, response));
719 if let Err(error) = cleanup {
720 cleanup_confirmed = false;
721 tracing::error!(%vm_id, %error, "failed to dispose partially initialized VM; retaining sidecar ownership until explicit worker shutdown");
722 }
723 vm_host_functions().remove(&vm_id);
724 vm_clients().remove(&vm_id);
725 }
726 if let Some(key) = created_session_key {
727 session_js_bridge_callbacks().remove(&key);
728 session_sqlite_callbacks().remove(&key);
729 }
730 if let Some(lease) = lease {
731 if cleanup_confirmed {
732 lease.dispose().await?;
733 } else {
734 lease.retain_until_worker_shutdown();
735 }
736 }
737 if let Err(error) = sidecar.dispose_if_unused().await {
738 tracing::error!(%error, "failed to stop sidecar after VM creation failed");
739 }
740 }
741 result
742 }
743
744 pub async fn link_software(&self, descriptor: PackageDescriptor) -> Result<(), ClientError> {
759 let vm_configuration = try_vm_mutation(&self.inner.vm_configuration_operation)?;
760 let package_id = format!("path:{}", descriptor.path);
761 self.link_software_path(&descriptor.path, &package_id, &vm_configuration)
762 .await
763 }
764
765 pub async fn install_software(
768 &self,
769 source: crate::software::PackageSource,
770 ) -> Result<crate::software::InstalledSoftware, ClientError> {
771 let _vm_configuration = try_vm_mutation(&self.inner.vm_configuration_operation)?;
772 let source = match source {
773 crate::software::PackageSource::Url {
774 url,
775 expected_digest,
776 } => wire::PackageAcquisitionSource::PackageUrlSource(wire::PackageUrlSource {
777 url,
778 expected_digest,
779 }),
780 crate::software::PackageSource::Path {
781 path,
782 expected_digest,
783 } => wire::PackageAcquisitionSource::PackagePathSource(wire::PackagePathSource {
784 path,
785 expected_digest,
786 }),
787 };
788 let options = self.inner.config.package_resolver.as_ref();
789 let response = self
790 .transport()
791 .request_wire(
792 wire_vm_ownership(
793 &self.inner.connection_id,
794 &self.inner.session_id,
795 &self.inner.vm_id,
796 ),
797 wire::RequestPayload::InstallPackageRequest(wire::InstallPackageRequest {
798 acquisition: wire::AcquirePackageRequest {
799 source,
800 advisory: false,
801 timeout_ms: None,
802 max_package_bytes: options.map(|options| options.max_package_bytes),
803 download_timeout_ms: options.map(|options| options.download_timeout_ms),
804 connect_timeout_ms: options.map(|options| options.connect_timeout_ms),
805 max_redirects: options.map(|options| options.max_redirects as u32),
806 allow_insecure_local_http: options
807 .is_some_and(|options| options.allow_insecure_local_http),
808 },
809 }),
810 )
811 .await?;
812 match response {
813 wire::ResponsePayload::PackageInstalledResponse(installed) => {
814 let package = installed.package;
815 let info = crate::software::InstalledSoftware {
816 package_id: package.package_id,
817 digest: package.digest,
818 size_bytes: package.size,
819 package_name: package.package_name,
820 version: package.version,
821 commands: package.commands,
822 };
823 let mut projected = self.inner.projected_commands.lock();
824 for command in installed.projected_commands {
825 projected.insert(command.name, command.guest_path);
826 }
827 self.inner
828 .installed_software
829 .lock()
830 .insert(info.package_id.clone(), info.clone());
831 Ok(info)
832 }
833 wire::ResponsePayload::RejectedResponse(rejected) => {
834 Err(ClientError::from_rejection(rejected))
835 }
836 other => Err(ClientError::Sidecar(format!(
837 "unexpected install_package response: {other:?}"
838 ))),
839 }
840 }
841
842 pub async fn uninstall_software(
845 &self,
846 package_id: &str,
847 ) -> Result<crate::software::InstalledSoftware, ClientError> {
848 let _vm_configuration = try_vm_mutation(&self.inner.vm_configuration_operation)?;
849 let installed = self
850 .inner
851 .installed_software
852 .lock()
853 .get(package_id)
854 .cloned()
855 .ok_or_else(|| ClientError::SoftwareNotFound(package_id.to_owned()))?;
856 let inner = self.inner();
857 let response = self
858 .transport()
859 .request_wire(
860 wire_vm_ownership(&inner.connection_id, &inner.session_id, &inner.vm_id),
861 wire::RequestPayload::UnlinkPackageRequest(wire::UnlinkPackageRequest {
862 package_id: package_id.to_owned(),
863 }),
864 )
865 .await?;
866 match response {
867 wire::ResponsePayload::PackageUnlinkedResponse(unlinked) => {
868 let mut commands = inner.projected_commands.lock();
869 for command in unlinked.removed_commands {
870 commands.remove(&command);
871 }
872 inner.installed_software.lock().remove(package_id);
873 Ok(installed)
874 }
875 wire::ResponsePayload::RejectedResponse(rejected) => Err(rejected_to_error(rejected)),
876 other => Err(ClientError::Sidecar(format!(
877 "unexpected unlink_package response: {other:?}"
878 ))),
879 }
880 }
881
882 pub fn installed_software(&self) -> Vec<crate::software::InstalledSoftware> {
886 self.inner
887 .installed_software
888 .lock()
889 .values()
890 .cloned()
891 .collect()
892 }
893
894 async fn link_software_path(
895 &self,
896 path: &str,
897 package_id: &str,
898 _vm_configuration: &tokio::sync::MutexGuard<'_, ()>,
899 ) -> Result<(), ClientError> {
900 let inner = self.inner();
901 let response = self
902 .transport()
903 .request_wire(
904 wire_vm_ownership(&inner.connection_id, &inner.session_id, &inner.vm_id),
905 wire::RequestPayload::LinkPackageRequest(wire::LinkPackageRequest {
906 package: wire::PackageDescriptor {
909 path: path.to_owned(),
910 },
911 package_id: package_id.to_owned(),
912 }),
913 )
914 .await?;
915 match response {
916 wire::ResponsePayload::PackageLinkedResponse(linked) => {
917 let mut guard = inner.projected_commands.lock();
918 for command in linked.projected_commands {
919 guard.insert(command.name, command.guest_path);
920 }
921 Ok(())
922 }
923 wire::ResponsePayload::RejectedResponse(rejected) => Err(rejected_to_error(rejected)),
924 other => Err(ClientError::Sidecar(format!(
925 "unexpected link_package response: {other:?}"
926 ))),
927 }
928 }
929
930 pub async fn list_software(&self) -> Result<Vec<SoftwareInfo>, ClientError> {
931 let inner = self.inner();
932 let response = self
933 .transport()
934 .request_wire(
935 wire_vm_ownership(&inner.connection_id, &inner.session_id, &inner.vm_id),
936 wire::RequestPayload::ProvidedCommandsRequest,
937 )
938 .await?;
939 match response {
940 wire::ResponsePayload::ProvidedCommandsResponse(provided) => Ok(provided
941 .packages
942 .into_iter()
943 .map(|package| SoftwareInfo {
944 package_name: package.package_name,
945 commands: package.commands,
946 })
947 .collect()),
948 wire::ResponsePayload::RejectedResponse(rejected) => Err(rejected_to_error(rejected)),
949 other => Err(ClientError::Sidecar(format!(
950 "unexpected list_software response: {other:?}"
951 ))),
952 }
953 }
954
955 pub async fn shutdown(&self) -> Result<(), ClientError> {
956 let mut shutdown_result = self.inner.shutdown_result.lock().await;
959 if let Some(result) = shutdown_result.as_ref() {
960 return result.clone();
961 }
962 self.inner.disposed.store(true, Ordering::SeqCst);
963
964 self.inner.cron.dispose();
969
970 crate::process::drain_process_output_tasks(&self.inner.processes);
973
974 let mut exit_tasks = Vec::new();
977 self.inner.pending_shell_exits.retain(|_, task| {
978 exit_tasks.push(std::mem::replace(task, tokio::spawn(async {})));
979 false
980 });
981
982 {
983 let _terminal_lifecycle_guard = self.inner.terminal_lifecycle_lock.lock().await;
984 let mut terminal_entries = Vec::new();
985 self.inner.terminals.retain(|process_id, entry| {
986 terminal_entries.push((
987 process_id.clone(),
988 std::mem::replace(&mut entry.exit_task, tokio::spawn(async {})),
989 ));
990 false
991 });
992 self.inner.terminal_count.store(0, Ordering::SeqCst);
993 for (process_id, _) in &terminal_entries {
994 let transport = self.transport().clone();
995 let ownership = wire::OwnershipScope::VmOwnership(wire::VmOwnership {
996 connection_id: self.inner.connection_id.clone(),
997 session_id: self.inner.session_id.clone(),
998 vm_id: self.inner.vm_id.clone(),
999 });
1000 let process_id = process_id.clone();
1001 exit_tasks.push(tokio::spawn(async move {
1002 let _ = transport
1003 .request_wire(
1004 ownership,
1005 wire::RequestPayload::KillProcessRequest(wire::KillProcessRequest {
1006 process_id,
1007 signal: String::from("SIGTERM"),
1008 }),
1009 )
1010 .await;
1011 }));
1012 }
1013 for (_, task) in terminal_entries {
1014 exit_tasks.push(task);
1015 }
1016 }
1017
1018 if !exit_tasks.is_empty() {
1019 let mut drain_tasks = exit_tasks;
1020 if tokio::time::timeout(
1021 Duration::from_millis(crate::SHELL_DISPOSE_TIMEOUT_MS),
1022 futures::future::join_all(drain_tasks.iter_mut()),
1023 )
1024 .await
1025 .is_err()
1026 {
1027 for task in drain_tasks {
1028 task.abort();
1029 }
1030 }
1031 }
1032
1033 let mut disposal_result = if self.inner.vm_disposed.load(Ordering::SeqCst)
1039 || self.transport().is_disconnected()
1040 {
1041 Ok(())
1042 } else {
1043 self.transport()
1044 .request_wire(
1045 wire::OwnershipScope::VmOwnership(wire::VmOwnership {
1046 connection_id: self.inner.connection_id.clone(),
1047 session_id: self.inner.session_id.clone(),
1048 vm_id: self.inner.vm_id.clone(),
1049 }),
1050 wire::RequestPayload::DisposeVmRequest(wire::DisposeVmRequest {
1051 reason: wire::DisposeReason::Requested,
1052 }),
1053 )
1054 .await
1055 .map_err(ClientError::from)
1056 .and_then(|response| vm_dispose_response(&self.inner.vm_id, response))
1057 };
1058 if disposal_result.is_ok() {
1059 self.inner.vm_disposed.store(true, Ordering::SeqCst);
1060 let _ = vm_host_functions().remove(&self.inner.vm_id);
1063 let _ = vm_clients().remove(&self.inner.vm_id);
1064 let _ = session_js_bridge_callbacks().remove(&sidecar_session_key(
1065 &self.inner.connection_id,
1066 &self.inner.session_id,
1067 ));
1068 let _ = session_sqlite_callbacks().remove(&sidecar_session_key(
1069 &self.inner.connection_id,
1070 &self.inner.session_id,
1071 ));
1072 }
1073 let sidecar = self.inner.sidecar.clone();
1074 let lease = if disposal_result.is_ok() {
1075 self.inner.sidecar_lease.lock().take()
1076 } else {
1077 None
1078 };
1079 if let Some(lease) = lease {
1080 retain_shutdown_failure(&mut disposal_result, lease.dispose().await);
1081 }
1082 retain_shutdown_failure(&mut disposal_result, sidecar.dispose_if_unused().await);
1083 if disposal_result.is_ok() {
1084 *shutdown_result = Some(Ok(()));
1085 }
1086 disposal_result
1087 }
1088
1089 pub(crate) fn inner(&self) -> &AgentOsInner {
1092 &self.inner
1093 }
1094
1095 pub(crate) fn transport(&self) -> &Arc<SidecarProcess> {
1096 &self.inner.transport
1097 }
1098
1099 pub(crate) fn connection_id(&self) -> &str {
1100 &self.inner.connection_id
1101 }
1102
1103 pub(crate) fn wire_session_id(&self) -> &str {
1104 &self.inner.session_id
1105 }
1106
1107 pub(crate) fn vm_id(&self) -> &str {
1108 &self.inner.vm_id
1109 }
1110
1111 pub(crate) fn cron(&self) -> &Arc<CronManager> {
1112 &self.inner.cron
1113 }
1114
1115 pub fn sidecar(&self) -> Arc<AgentOsSidecar> {
1118 self.inner.sidecar.clone()
1119 }
1120
1121 #[cfg(feature = "service-internals")]
1122 pub(crate) async fn vm_config_equivalent(
1123 &self,
1124 before: &AgentOsConfig,
1125 after: &AgentOsConfig,
1126 before_restart_identity: Vec<String>,
1127 after_restart_identity: Vec<String>,
1128 ) -> Result<bool, ClientError> {
1129 let before_mounts = serialize_mounts(before)?;
1130 let after_mounts = serialize_mounts(after)?;
1131 let before = serialize_create_vm_config_for_sidecar(before)?;
1132 let after = serialize_create_vm_config_for_sidecar(after)?;
1133 let response = self
1134 .transport()
1135 .request_wire(
1136 wire_session_ownership(self.connection_id(), self.wire_session_id()),
1137 wire::RequestPayload::CompareVmConfigRequest(wire::CompareVmConfigRequest {
1138 before: serde_json::to_string(&before).map_err(|error| {
1139 ClientError::Sidecar(format!("serialize before VM config: {error}"))
1140 })?,
1141 after: serde_json::to_string(&after).map_err(|error| {
1142 ClientError::Sidecar(format!("serialize after VM config: {error}"))
1143 })?,
1144 before_mounts,
1145 after_mounts,
1146 before_restart_identity,
1147 after_restart_identity,
1148 }),
1149 )
1150 .await?;
1151 match response {
1152 wire::ResponsePayload::VmConfigComparedResponse(result) => Ok(result.equivalent),
1153 wire::ResponsePayload::RejectedResponse(rejected) => Err(rejected_to_error(rejected)),
1154 _ => Err(ClientError::Sidecar(String::from(
1155 "unexpected compare VM config response",
1156 ))),
1157 }
1158 }
1159}
1160
1161pub(crate) fn sidecar_wire_placement(sidecar: &AgentOsSidecar) -> wire::SidecarPlacement {
1163 match &sidecar.placement {
1164 AgentOsSidecarPlacement::Shared { pool } => {
1165 wire::SidecarPlacement::SidecarPlacementShared(wire::SidecarPlacementShared {
1166 pool: pool.clone(),
1167 })
1168 }
1169 AgentOsSidecarPlacement::Explicit { sidecar_id } => {
1170 wire::SidecarPlacement::SidecarPlacementExplicit(wire::SidecarPlacementExplicit {
1171 sidecar_id: sidecar_id.clone(),
1172 })
1173 }
1174 }
1175}
1176
1177fn wire_connection_ownership(connection_id: &str) -> wire::OwnershipScope {
1178 wire::OwnershipScope::ConnectionOwnership(wire::ConnectionOwnership {
1179 connection_id: connection_id.to_string(),
1180 })
1181}
1182
1183fn wire_session_ownership(connection_id: &str, session_id: &str) -> wire::OwnershipScope {
1184 wire::OwnershipScope::SessionOwnership(wire::SessionOwnership {
1185 connection_id: connection_id.to_string(),
1186 session_id: session_id.to_string(),
1187 })
1188}
1189
1190fn wire_vm_ownership(connection_id: &str, session_id: &str, vm_id: &str) -> wire::OwnershipScope {
1191 wire::OwnershipScope::VmOwnership(wire::VmOwnership {
1192 connection_id: connection_id.to_string(),
1193 session_id: session_id.to_string(),
1194 vm_id: vm_id.to_string(),
1195 })
1196}
1197
1198fn serialize_create_vm_config_for_sidecar(
1199 config: &AgentOsConfig,
1200) -> Result<vm_config::CreateVmConfig, ClientError> {
1201 let (root_filesystem, native_root) =
1202 serialize_root_filesystem_config_for_sidecar(&config.root_filesystem)?;
1203 let mut create = vm_config::CreateVmConfig {
1204 defaults_profile: Some(vm_config::VmDefaultsProfile::AgentOs),
1205 database: config.database.clone(),
1206 cwd: None,
1207 env: config.environment.clone(),
1208 user: config.user.clone(),
1209 root_filesystem,
1210 permissions: permissions_policy_config(config),
1211 limits: serialize_limits_config_for_sidecar(config.limits.as_ref())?,
1212 dns: None,
1213 native_root,
1214 listen: None,
1215 loopback_exempt_ports: config.loopback_exempt_ports.clone(),
1216 js_runtime: (config.allowed_node_builtins.is_some()
1222 || config.high_resolution_time.is_some())
1223 .then(|| vm_config::JsRuntimeConfig {
1224 platform: vm_config::JsRuntimePlatform::default(),
1225 module_resolution: vm_config::JsModuleResolution::default(),
1226 allowed_builtins: config.allowed_node_builtins.clone(),
1227 high_resolution_time: config.high_resolution_time,
1228 }),
1229 bootstrap_commands: None,
1230 };
1231 create
1232 .normalize()
1233 .map_err(|error| ClientError::Sidecar(format!("invalid VM config: {error}")))?;
1234 Ok(create)
1235}
1236
1237pub(crate) fn validate_config(config: &AgentOsConfig) -> Result<(), ClientError> {
1238 if matches!(
1239 config.database,
1240 Some(vm_config::VmSqliteDescriptor::HostCallback { .. })
1241 ) && config.sidecar_sqlite_callback.is_none()
1242 {
1243 return Err(ClientError::Sidecar(String::from(
1244 "database type host_callback requires sidecar_sqlite_callback",
1245 )));
1246 }
1247 if let Some(options) = &config.package_resolver {
1248 options.validate()?;
1249 }
1250 let create = serialize_create_vm_config_for_sidecar(config)?;
1251 create
1252 .validate(wire::DEFAULT_MAX_FRAME_BYTES)
1253 .map_err(|error| ClientError::Sidecar(format!("invalid VM config: {error}")))?;
1254 serialize_mounts(config)?;
1255 Ok(())
1256}
1257
1258fn serialize_root_filesystem_config_for_sidecar(
1259 config: &RootFilesystemConfig,
1260) -> Result<
1261 (
1262 vm_config::RootFilesystemConfig,
1263 Option<vm_config::NativeRootFilesystemConfig>,
1264 ),
1265 ClientError,
1266> {
1267 let mode = match config.mode.unwrap_or(ConfigRootFilesystemMode::Ephemeral) {
1268 ConfigRootFilesystemMode::Ephemeral => vm_config::RootFilesystemMode::Ephemeral,
1269 ConfigRootFilesystemMode::ReadOnly => vm_config::RootFilesystemMode::ReadOnly,
1270 };
1271 match config.kind {
1272 RootFilesystemKind::Overlay => {
1273 if config.native_plugin.is_some() {
1274 return Err(ClientError::Sidecar(
1275 "rootFilesystem.nativePlugin requires type \"native\"".to_string(),
1276 ));
1277 }
1278 let lowers = config
1279 .lowers
1280 .iter()
1281 .map(serialize_root_lower_config_for_sidecar)
1282 .collect::<Result<Vec<_>, _>>()?;
1283 Ok((
1284 vm_config::RootFilesystemConfig {
1285 mode,
1286 disable_default_base_layer: config.disable_default_base_layer,
1287 lowers,
1288 bootstrap_entries: Vec::new(),
1289 },
1290 None,
1291 ))
1292 }
1293 RootFilesystemKind::Native => {
1294 if !config.lowers.is_empty() {
1295 return Err(ClientError::Sidecar(
1296 "native root filesystems do not support rootFilesystem.lowers".to_string(),
1297 ));
1298 }
1299 let plugin = config.native_plugin.as_ref().ok_or_else(|| {
1300 ClientError::Sidecar(
1301 "rootFilesystem.nativePlugin is required for type \"native\"".to_string(),
1302 )
1303 })?;
1304 Ok((
1305 vm_config::RootFilesystemConfig {
1306 mode,
1307 disable_default_base_layer: config.disable_default_base_layer,
1308 lowers: Vec::new(),
1309 bootstrap_entries: Vec::new(),
1310 },
1311 Some(vm_config::NativeRootFilesystemConfig {
1312 plugin: vm_config::MountPluginDescriptor {
1313 id: plugin.id.clone(),
1314 config: plugin
1315 .config
1316 .clone()
1317 .unwrap_or_else(|| serde_json::Value::Object(serde_json::Map::new())),
1318 },
1319 read_only: config.mode == Some(ConfigRootFilesystemMode::ReadOnly),
1320 }),
1321 ))
1322 }
1323 }
1324}
1325
1326fn serialize_root_lower_config_for_sidecar(
1327 lower: &RootLowerInput,
1328) -> Result<vm_config::RootFilesystemLowerDescriptor, ClientError> {
1329 match lower {
1330 RootLowerInput::BundledBaseFilesystem => {
1331 Ok(vm_config::RootFilesystemLowerDescriptor::BundledBaseFilesystem)
1332 }
1333 RootLowerInput::SnapshotExport(snapshot) => {
1334 let entries = snapshot
1335 .source
1336 .filesystem
1337 .entries
1338 .iter()
1339 .map(serialize_filesystem_entry_config_for_sidecar)
1340 .collect::<Result<Vec<_>, _>>()?;
1341 Ok(vm_config::RootFilesystemLowerDescriptor::Snapshot { entries })
1342 }
1343 }
1344}
1345
1346fn serialize_filesystem_entry_config_for_sidecar(
1347 entry: &crate::fs::FilesystemEntry,
1348) -> Result<vm_config::RootFilesystemEntry, ClientError> {
1349 let mode = u32::from_str_radix(entry.mode.trim_start_matches("0o"), 8).map_err(|error| {
1350 ClientError::Sidecar(format!(
1351 "invalid root filesystem mode {} for {}: {error}",
1352 entry.mode, entry.path
1353 ))
1354 })?;
1355 let kind = match entry.entry_type {
1356 crate::fs::DirEntryType::File => vm_config::RootFilesystemEntryKind::File,
1357 crate::fs::DirEntryType::Directory => vm_config::RootFilesystemEntryKind::Directory,
1358 crate::fs::DirEntryType::Symlink => vm_config::RootFilesystemEntryKind::Symlink,
1359 };
1360 let encoding = entry.encoding.map(|encoding| match encoding {
1361 crate::fs::FilesystemEntryEncoding::Utf8 => vm_config::RootFilesystemEntryEncoding::Utf8,
1362 crate::fs::FilesystemEntryEncoding::Base64 => {
1363 vm_config::RootFilesystemEntryEncoding::Base64
1364 }
1365 });
1366
1367 Ok(vm_config::RootFilesystemEntry {
1368 path: entry.path.clone(),
1369 kind,
1370 mode: Some(mode),
1371 uid: Some(entry.uid),
1372 gid: Some(entry.gid),
1373 content: entry.content.clone(),
1374 encoding,
1375 target: entry.target.clone(),
1376 executable: entry.entry_type == crate::fs::DirEntryType::File && (mode & 0o111) != 0,
1377 })
1378}
1379
1380fn serialize_limits_config_for_sidecar(
1381 limits: Option<&AgentOsLimits>,
1382) -> Result<Option<vm_config::VmLimitsConfig>, ClientError> {
1383 let Some(limits) = limits else {
1384 return Ok(None);
1385 };
1386 let value = serde_json::to_value(limits).map_err(|error| {
1387 ClientError::Sidecar(format!("failed to serialize VM limits config: {error}"))
1388 })?;
1389 serde_json::from_value(value).map(Some).map_err(|error| {
1390 ClientError::Sidecar(format!("failed to encode VM limits config: {error}"))
1391 })
1392}
1393
1394fn permissions_policy_config(config: &AgentOsConfig) -> Option<vm_config::PermissionsPolicy> {
1395 config
1396 .permissions
1397 .as_ref()
1398 .map(|permissions| vm_config::PermissionsPolicy {
1399 fs: permissions.fs.as_ref().map(serialize_fs_permissions_config),
1400 network: permissions
1401 .network
1402 .as_ref()
1403 .map(serialize_pattern_permissions_config),
1404 child_process: permissions
1405 .child_process
1406 .as_ref()
1407 .map(serialize_pattern_permissions_config),
1408 process: permissions
1409 .process
1410 .as_ref()
1411 .map(serialize_pattern_permissions_config),
1412 env: permissions
1413 .env
1414 .as_ref()
1415 .map(serialize_pattern_permissions_config),
1416 host_function: permissions
1417 .host_function
1418 .as_ref()
1419 .map(serialize_pattern_permissions_config),
1420 })
1421}
1422
1423fn serialize_fs_permissions_config(
1424 permissions: &crate::config::FsPermissions,
1425) -> vm_config::FsPermissionScope {
1426 match permissions {
1427 crate::config::FsPermissions::Mode(mode) => {
1428 vm_config::FsPermissionScope::Mode(serialize_permission_mode_config(*mode))
1429 }
1430 crate::config::FsPermissions::Rules(rules) => {
1431 vm_config::FsPermissionScope::Rules(vm_config::FsPermissionRuleSet {
1432 default: rules.default.map(serialize_permission_mode_config),
1433 rules: rules
1434 .rules
1435 .iter()
1436 .map(|rule| vm_config::FsPermissionRule {
1437 mode: serialize_permission_mode_config(rule.mode),
1438 operations: operation_wildcard_if_omitted(&rule.operations),
1439 paths: resource_wildcard_if_omitted(&rule.paths),
1440 })
1441 .collect(),
1442 })
1443 }
1444 }
1445}
1446
1447fn serialize_pattern_permissions_config(
1448 permissions: &crate::config::PatternPermissions,
1449) -> vm_config::PatternPermissionScope {
1450 match permissions {
1451 crate::config::PatternPermissions::Mode(mode) => {
1452 vm_config::PatternPermissionScope::Mode(serialize_permission_mode_config(*mode))
1453 }
1454 crate::config::PatternPermissions::Rules(rules) => {
1455 vm_config::PatternPermissionScope::Rules(vm_config::PatternPermissionRuleSet {
1456 default: rules.default.map(serialize_permission_mode_config),
1457 rules: rules
1458 .rules
1459 .iter()
1460 .map(|rule| vm_config::PatternPermissionRule {
1461 mode: serialize_permission_mode_config(rule.mode),
1462 operations: operation_wildcard_if_omitted(&rule.operations),
1463 patterns: resource_wildcard_if_omitted(&rule.patterns),
1464 })
1465 .collect(),
1466 })
1467 }
1468 }
1469}
1470
1471fn serialize_permission_mode_config(
1472 mode: crate::config::PermissionMode,
1473) -> vm_config::PermissionMode {
1474 match mode {
1475 crate::config::PermissionMode::Allow => vm_config::PermissionMode::Allow,
1476 crate::config::PermissionMode::Deny => vm_config::PermissionMode::Deny,
1477 }
1478}
1479
1480async fn wait_for_vm_ready(
1482 events: &mut broadcast::Receiver<(wire::OwnershipScope, wire::EventPayload)>,
1483 vm_id: &str,
1484 timeout_ms: u64,
1485) -> Result<(), ClientError> {
1486 let wait = async {
1487 loop {
1488 match events.recv().await {
1489 Ok((ownership, payload)) => match payload {
1490 wire::EventPayload::VmLifecycleEvent(event) => {
1491 if matches!(event.state, wire::VmLifecycleState::Ready)
1492 && wire_ownership_vm_id(&ownership) == Some(vm_id)
1493 {
1494 return Ok(());
1495 }
1496 }
1497 wire::EventPayload::ProcessOutputEvent(_)
1498 | wire::EventPayload::ProcessExitedEvent(_)
1499 | wire::EventPayload::ExecutionOutputEvent(_)
1500 | wire::EventPayload::ExecutionCompletedEvent(_)
1501 | wire::EventPayload::StructuredEvent(_)
1502 | wire::EventPayload::ExtEnvelope(_) => {}
1503 },
1504 Err(broadcast::error::RecvError::Lagged(_)) => {}
1505 Err(broadcast::error::RecvError::Closed) => {
1506 return Err(ClientError::SidecarExited(
1507 "the sidecar exited before the VM became ready".to_string(),
1508 ));
1509 }
1510 }
1511 }
1512 };
1513 tokio::time::timeout(Duration::from_millis(timeout_ms), wait)
1514 .await
1515 .map_err(|_| {
1516 ClientError::Sidecar("timed out waiting for the VM to become ready".to_string())
1517 })?
1518}
1519
1520static VM_HOST_FUNCTIONS: OnceCell<SccHashMap<String, Arc<VmHostFunctionRegistry>>> =
1523 OnceCell::new();
1524
1525#[derive(Clone)]
1526struct VmHostFunctionRegistry {
1527 host_functions: Vec<ResolvedHostFunctions>,
1528 host_function_map: HashMap<String, ResolvedHostFunction>,
1529}
1530
1531fn vm_host_functions() -> &'static SccHashMap<String, Arc<VmHostFunctionRegistry>> {
1532 VM_HOST_FUNCTIONS.get_or_init(SccHashMap::new)
1533}
1534
1535static VM_CLIENTS: OnceCell<SccHashMap<String, Weak<AgentOsInner>>> = OnceCell::new();
1539
1540fn vm_clients() -> &'static SccHashMap<String, Weak<AgentOsInner>> {
1541 VM_CLIENTS.get_or_init(SccHashMap::new)
1542}
1543
1544static SESSION_JS_BRIDGE_CALLBACKS: OnceCell<SccHashMap<String, SidecarJsBridgeCallback>> =
1549 OnceCell::new();
1550
1551fn session_js_bridge_callbacks() -> &'static SccHashMap<String, SidecarJsBridgeCallback> {
1552 SESSION_JS_BRIDGE_CALLBACKS.get_or_init(SccHashMap::new)
1553}
1554
1555static SESSION_SQLITE_CALLBACKS: OnceCell<SccHashMap<String, SidecarSqliteCallback>> =
1559 OnceCell::new();
1560
1561fn session_sqlite_callbacks() -> &'static SccHashMap<String, SidecarSqliteCallback> {
1562 SESSION_SQLITE_CALLBACKS.get_or_init(SccHashMap::new)
1563}
1564
1565fn sidecar_session_key(connection_id: &str, session_id: &str) -> String {
1566 format!("{connection_id}\0{session_id}")
1567}
1568
1569fn wire_ownership_session_key(ownership: &wire::OwnershipScope) -> Option<String> {
1570 match ownership {
1571 wire::OwnershipScope::SessionOwnership(ownership) => Some(sidecar_session_key(
1572 &ownership.connection_id,
1573 &ownership.session_id,
1574 )),
1575 wire::OwnershipScope::VmOwnership(ownership) => Some(sidecar_session_key(
1576 &ownership.connection_id,
1577 &ownership.session_id,
1578 )),
1579 wire::OwnershipScope::ConnectionOwnership(_) => None,
1580 }
1581}
1582
1583fn js_bridge_call_callback() -> WireSidecarCallback {
1584 Arc::new(|payload, ownership| {
1585 Box::pin(async move {
1586 let request = match payload {
1587 wire::SidecarRequestPayload::JsBridgeCallRequest(request) => request,
1588 wire::SidecarRequestPayload::HostCallbackRequest(_) => {
1589 return Ok(wire::SidecarResponsePayload::JsBridgeResultResponse(
1590 wire::JsBridgeResultResponse {
1591 call_id: "unknown".to_string(),
1592 result: None,
1593 error: Some(
1594 "js-bridge callback received a host callback request".to_string(),
1595 ),
1596 },
1597 ));
1598 }
1599 wire::SidecarRequestPayload::ExtEnvelope(_) => {
1600 return Ok(wire::SidecarResponsePayload::JsBridgeResultResponse(
1601 wire::JsBridgeResultResponse {
1602 call_id: "unknown".to_string(),
1603 result: None,
1604 error: Some(
1605 "js-bridge callback received an extension request".to_string(),
1606 ),
1607 },
1608 ));
1609 }
1610 };
1611 Ok(wire::SidecarResponsePayload::JsBridgeResultResponse(
1612 run_js_bridge_callback(&ownership, request).await,
1613 ))
1614 })
1615 })
1616}
1617
1618fn sqlite_callback_callback() -> WireSidecarCallback {
1619 Arc::new(|payload, ownership| {
1620 Box::pin(async move {
1621 let envelope = match payload {
1622 wire::SidecarRequestPayload::ExtEnvelope(envelope) => envelope,
1623 _ => {
1624 return Err(TransportError::Sidecar(String::from(
1625 "SQLite callback received a non-extension request",
1626 )))
1627 }
1628 };
1629 if envelope.namespace != vm_config::VM_SQLITE_CALLBACK_NAMESPACE {
1630 return Ok(wire::SidecarResponsePayload::ExtEnvelope(
1631 wire::ExtEnvelope {
1632 namespace: envelope.namespace,
1633 payload: serde_json::to_vec(&vm_config::VmSqliteCallbackResponse::Error {
1634 message: String::from("unsupported sidecar extension namespace"),
1635 })
1636 .expect("serialize SQLite extension error"),
1637 },
1638 ));
1639 }
1640 let response = match serde_json::from_slice::<vm_config::VmSqliteCallbackRequest>(
1641 &envelope.payload,
1642 ) {
1643 Ok(request) => {
1644 let callback = wire_ownership_session_key(&ownership).and_then(|key| {
1645 session_sqlite_callbacks().read(&key, |_, callback| callback.clone())
1646 });
1647 match callback {
1648 Some(callback) => match callback(request).await {
1649 Ok(response) => response,
1650 Err(message) => vm_config::VmSqliteCallbackResponse::Error { message },
1651 },
1652 None => vm_config::VmSqliteCallbackResponse::Error {
1653 message: String::from(
1654 "no SQLite callback registered for sidecar session",
1655 ),
1656 },
1657 }
1658 }
1659 Err(error) => vm_config::VmSqliteCallbackResponse::Error {
1660 message: format!("invalid SQLite callback request: {error}"),
1661 },
1662 };
1663 Ok(wire::SidecarResponsePayload::ExtEnvelope(
1664 wire::ExtEnvelope {
1665 namespace: vm_config::VM_SQLITE_CALLBACK_NAMESPACE.to_owned(),
1666 payload: serde_json::to_vec(&response).map_err(|error| {
1667 TransportError::Sidecar(format!(
1668 "serialize SQLite callback response: {error}"
1669 ))
1670 })?,
1671 },
1672 ))
1673 })
1674 })
1675}
1676
1677async fn run_js_bridge_callback(
1678 ownership: &wire::OwnershipScope,
1679 request: wire::JsBridgeCallRequest,
1680) -> wire::JsBridgeResultResponse {
1681 let call_id = request.call_id;
1682 let args = match serde_json::from_str::<Value>(&request.args) {
1683 Ok(args) => args,
1684 Err(error) => {
1685 return wire::JsBridgeResultResponse {
1686 call_id,
1687 result: None,
1688 error: Some(format!("Invalid js_bridge args: {error}")),
1689 };
1690 }
1691 };
1692 let callback = wire_ownership_session_key(ownership)
1693 .and_then(|key| session_js_bridge_callbacks().read(&key, |_, callback| callback.clone()));
1694 let Some(callback) = callback else {
1695 return wire::JsBridgeResultResponse {
1696 call_id,
1697 result: None,
1698 error: Some("No js_bridge callback registered for sidecar session".to_string()),
1699 };
1700 };
1701
1702 let call = SidecarJsBridgeCall {
1703 call_id: call_id.clone(),
1704 mount_id: request.mount_id,
1705 operation: request.operation,
1706 args,
1707 };
1708 match callback(call).await {
1709 Ok(result) => match result {
1710 Some(value) => match serde_json::to_string(&value) {
1711 Ok(result) => wire::JsBridgeResultResponse {
1712 call_id,
1713 result: Some(result),
1714 error: None,
1715 },
1716 Err(error) => wire::JsBridgeResultResponse {
1717 call_id,
1718 result: None,
1719 error: Some(format!("Invalid js_bridge result: {error}")),
1720 },
1721 },
1722 None => wire::JsBridgeResultResponse {
1723 call_id,
1724 result: None,
1725 error: None,
1726 },
1727 },
1728 Err(error) => wire::JsBridgeResultResponse {
1729 call_id,
1730 result: None,
1731 error: Some(error),
1732 },
1733 }
1734}
1735
1736fn host_callback_callback() -> WireSidecarCallback {
1738 Arc::new(|payload, ownership| {
1739 Box::pin(async move {
1740 let request = match payload {
1741 wire::SidecarRequestPayload::HostCallbackRequest(request) => request,
1742 wire::SidecarRequestPayload::JsBridgeCallRequest(_) => {
1743 return Ok(wire::SidecarResponsePayload::HostCallbackResultResponse(
1744 wire::HostCallbackResultResponse {
1745 invocation_id: "unknown".to_string(),
1746 result: None,
1747 error: Some(
1748 "host callback received a non-host-function request".to_string(),
1749 ),
1750 },
1751 ));
1752 }
1753 wire::SidecarRequestPayload::ExtEnvelope(envelope) => {
1754 return Ok(wire::SidecarResponsePayload::ExtEnvelope(
1755 wire::ExtEnvelope {
1756 namespace: envelope.namespace,
1757 payload: b"host-callback received an extension request".to_vec(),
1758 },
1759 ));
1760 }
1761 };
1762 Ok(wire::SidecarResponsePayload::HostCallbackResultResponse(
1763 run_host_callback(&ownership, request).await,
1764 ))
1765 })
1766 })
1767}
1768
1769async fn run_host_callback(
1772 ownership: &wire::OwnershipScope,
1773 request: wire::HostCallbackRequest,
1774) -> wire::HostCallbackResultResponse {
1775 let input = match serde_json::from_str::<Value>(&request.input) {
1776 Ok(input) => input,
1777 Err(error) => {
1778 return wire::HostCallbackResultResponse {
1779 invocation_id: request.invocation_id,
1780 result: None,
1781 error: Some(format!("Invalid host callback input: {error}")),
1782 };
1783 }
1784 };
1785 let vm_id = wire_ownership_vm_id(ownership).unwrap_or("");
1786 let registry = vm_host_functions().read(vm_id, |_, registry| registry.clone());
1787 let Some(registry) = registry else {
1788 return wire::HostCallbackResultResponse {
1789 invocation_id: request.invocation_id,
1790 result: None,
1791 error: Some(format!(
1792 "Unknown host function \"{}\"",
1793 request.callback_key
1794 )),
1795 };
1796 };
1797
1798 if request.callback_key == "agentos" {
1799 let Some(command) = parse_host_command_callback_input(&input) else {
1800 return wire::HostCallbackResultResponse {
1801 invocation_id: request.invocation_id,
1802 result: None,
1803 error: Some(String::from("Invalid registry callback for \"agentos\"")),
1804 };
1805 };
1806 if command.command != request.callback_key {
1807 return wire::HostCallbackResultResponse {
1808 invocation_id: request.invocation_id,
1809 result: None,
1810 error: Some(String::from("Invalid registry callback for \"agentos\"")),
1811 };
1812 }
1813 return match run_host_command_callback(ownership, registry.as_ref(), command).await {
1814 Ok(value) => match host_callback_json_result(value) {
1815 Ok(result) => wire::HostCallbackResultResponse {
1816 invocation_id: request.invocation_id,
1817 result: Some(result),
1818 error: None,
1819 },
1820 Err(error) => wire::HostCallbackResultResponse {
1821 invocation_id: request.invocation_id,
1822 result: None,
1823 error: Some(error),
1824 },
1825 },
1826 Err(error) => wire::HostCallbackResultResponse {
1827 invocation_id: request.invocation_id,
1828 result: None,
1829 error: Some(error),
1830 },
1831 };
1832 }
1833
1834 let host_function = registry
1835 .host_function_map
1836 .get(&request.callback_key)
1837 .cloned();
1838 let Some(host_function) = host_function else {
1839 return wire::HostCallbackResultResponse {
1840 invocation_id: request.invocation_id,
1841 result: None,
1842 error: Some(format!(
1843 "Unknown host function \"{}\"",
1844 request.callback_key
1845 )),
1846 };
1847 };
1848 let timeout = Duration::from_millis(request.timeout_ms.max(1));
1849 match tokio::time::timeout(timeout, (host_function.execute)(input)).await {
1850 Ok(Ok(value)) => match host_callback_json_result(value) {
1851 Ok(result) => wire::HostCallbackResultResponse {
1852 invocation_id: request.invocation_id,
1853 result: Some(result),
1854 error: None,
1855 },
1856 Err(error) => wire::HostCallbackResultResponse {
1857 invocation_id: request.invocation_id,
1858 result: None,
1859 error: Some(error),
1860 },
1861 },
1862 Ok(Err(error)) => wire::HostCallbackResultResponse {
1863 invocation_id: request.invocation_id,
1864 result: None,
1865 error: Some(error),
1866 },
1867 Err(_) => wire::HostCallbackResultResponse {
1868 invocation_id: request.invocation_id,
1869 result: None,
1870 error: Some(format!(
1871 "Host function \"{}\" timed out after {}ms",
1872 request.callback_key, request.timeout_ms
1873 )),
1874 },
1875 }
1876}
1877
1878#[derive(Debug, Deserialize)]
1879struct HostCommandCallbackInput {
1880 #[serde(rename = "type")]
1881 kind: String,
1882 command: String,
1883 #[serde(default)]
1884 args: Vec<String>,
1885 cwd: String,
1886}
1887
1888fn parse_host_command_callback_input(input: &Value) -> Option<HostCommandCallbackInput> {
1889 let command = serde_json::from_value::<HostCommandCallbackInput>(input.clone()).ok()?;
1890 if command.kind == "command" {
1891 Some(command)
1892 } else {
1893 None
1894 }
1895}
1896
1897async fn run_host_command_callback(
1898 ownership: &wire::OwnershipScope,
1899 registry: &VmHostFunctionRegistry,
1900 command: HostCommandCallbackInput,
1901) -> Result<Value, String> {
1902 if command.command == "agentos" {
1903 return handle_agentos_registry_command(ownership, registry, &command).await;
1904 }
1905 let Some(collection) = registry
1906 .host_functions
1907 .iter()
1908 .find(|collection| format!("agentos-{}", collection.name) == command.command)
1909 else {
1910 return Err(format!(
1911 "Unknown host callback command \"{}\"",
1912 command.command
1913 ));
1914 };
1915 handle_agentos_host_function_command(ownership, registry, &command, collection).await
1916}
1917
1918async fn handle_agentos_registry_command(
1919 ownership: &wire::OwnershipScope,
1920 registry: &VmHostFunctionRegistry,
1921 command: &HostCommandCallbackInput,
1922) -> Result<Value, String> {
1923 let Some(subcommand) = command.args.first() else {
1924 return Ok(json_object([(
1925 "usage",
1926 Value::String(String::from(
1927 "agentos <command>: list-host-functions [collection], <collection> --help, or <collection> <function> ...",
1928 )),
1929 )]));
1930 };
1931 if is_help_flag(subcommand) {
1932 return Ok(json_object([(
1933 "usage",
1934 Value::String(String::from(
1935 "agentos <command>: list-host-functions [collection], <collection> --help, or <collection> <function> ...",
1936 )),
1937 )]));
1938 }
1939 if subcommand == "list-host-functions" || subcommand == "list-bindings" {
1940 return match command.args.get(1) {
1941 Some(collection_name) => {
1942 describe_host_functions_payload(®istry.host_functions, collection_name)
1943 }
1944 None => Ok(list_host_functions_payload(®istry.host_functions)),
1945 };
1946 }
1947
1948 let Some(collection) = registry
1949 .host_functions
1950 .iter()
1951 .find(|collection| collection.name == *subcommand)
1952 else {
1953 return Err(format!(
1954 "No host function collection \"{subcommand}\". Available: {}",
1955 host_functions_names(®istry.host_functions)
1956 ));
1957 };
1958
1959 let Some(host_function_name) = command.args.get(1) else {
1960 return describe_host_functions_payload(®istry.host_functions, subcommand);
1961 };
1962 if is_help_flag(host_function_name) {
1963 return describe_host_functions_payload(®istry.host_functions, subcommand);
1964 }
1965 if command.args.get(2).is_some_and(|value| is_help_flag(value)) {
1966 return describe_host_function_payload(collection, host_function_name);
1967 }
1968 invoke_host_function(
1969 ownership,
1970 registry,
1971 collection,
1972 host_function_name,
1973 command.args.get(2..).unwrap_or_default(),
1974 &command.cwd,
1975 )
1976 .await
1977}
1978
1979async fn handle_agentos_host_function_command(
1980 ownership: &wire::OwnershipScope,
1981 registry: &VmHostFunctionRegistry,
1982 command: &HostCommandCallbackInput,
1983 collection: &ResolvedHostFunctions,
1984) -> Result<Value, String> {
1985 let Some(host_function_name) = command.args.first() else {
1986 return describe_host_functions_payload(®istry.host_functions, &collection.name);
1987 };
1988 if is_help_flag(host_function_name) {
1989 return describe_host_functions_payload(®istry.host_functions, &collection.name);
1990 }
1991 if command.args.get(1).is_some_and(|value| is_help_flag(value)) {
1992 return describe_host_function_payload(collection, host_function_name);
1993 }
1994 invoke_host_function(
1995 ownership,
1996 registry,
1997 collection,
1998 host_function_name,
1999 command.args.get(1..).unwrap_or_default(),
2000 &command.cwd,
2001 )
2002 .await
2003}
2004
2005async fn invoke_host_function(
2006 ownership: &wire::OwnershipScope,
2007 registry: &VmHostFunctionRegistry,
2008 collection: &ResolvedHostFunctions,
2009 host_function_name: &str,
2010 args: &[String],
2011 cwd: &str,
2012) -> Result<Value, String> {
2013 let callback_key = format!("{}:{host_function_name}", collection.name);
2014 let Some(host_function) = registry.host_function_map.get(&callback_key).cloned() else {
2015 return Err(format!(
2016 "No host function \"{host_function_name}\" in collection \"{}\". Available: {}",
2017 collection.name,
2018 host_function_names(collection)
2019 ));
2020 };
2021
2022 let input = parse_host_function_input(ownership, &host_function, args, cwd).await?;
2026 validate_host_function_input(&host_function.input_schema, &input)
2027 .map_err(|error| error.to_string())?;
2028
2029 let timeout = Duration::from_millis(host_function.timeout_ms.unwrap_or(30_000).max(1));
2030 match tokio::time::timeout(timeout, (host_function.execute)(input)).await {
2031 Ok(Ok(value)) => Ok(value),
2032 Ok(Err(error)) => Err(error),
2033 Err(_) => Err(format!(
2034 "Host function \"{callback_key}\" timed out after {}ms",
2035 host_function.timeout_ms.unwrap_or(30_000)
2036 )),
2037 }
2038}
2039
2040async fn parse_host_function_input(
2041 ownership: &wire::OwnershipScope,
2042 host_function: &ResolvedHostFunction,
2043 args: &[String],
2044 cwd: &str,
2045) -> Result<Value, String> {
2046 if args.first().is_some_and(|arg| arg == "--json") {
2047 let value = args
2048 .get(1)
2049 .ok_or_else(|| String::from("Flag --json requires a value"))?;
2050 return serde_json::from_str(value)
2051 .map_err(|error| format!("Invalid JSON for --json: {error}"));
2052 }
2053
2054 if args.first().is_some_and(|arg| arg == "--json-file") {
2055 let path = args
2056 .get(1)
2057 .ok_or_else(|| String::from("Flag --json-file requires a value"))?;
2058 let guest_path = normalize_guest_path(if path.starts_with('/') {
2059 path.clone()
2060 } else {
2061 format!("{cwd}/{path}")
2062 });
2063 let vm_id = wire_ownership_vm_id(ownership).unwrap_or("");
2064 let inner = vm_clients()
2065 .read(vm_id, |_, weak| weak.clone())
2066 .and_then(|weak| weak.upgrade())
2067 .ok_or_else(|| String::from("Invalid JSON file: VM is no longer available"))?;
2068 let bytes = AgentOs { inner }
2069 .read_file(&guest_path)
2070 .await
2071 .map_err(|error| format!("Invalid JSON file: {error}"))?;
2072 let text =
2073 String::from_utf8(bytes).map_err(|error| format!("Invalid JSON file: {error}"))?;
2074 return serde_json::from_str(&text).map_err(|error| format!("Invalid JSON file: {error}"));
2075 }
2076
2077 parse_host_function_argv(&host_function.input_schema, args)
2078}
2079
2080fn host_callback_json_result(value: Value) -> Result<String, String> {
2081 serde_json::to_string(&value).map_err(|error| format!("Invalid host callback result: {error}"))
2082}
2083
2084fn parse_host_function_argv(schema: &Value, argv: &[String]) -> Result<Value, String> {
2085 let properties = schema
2086 .get("properties")
2087 .and_then(Value::as_object)
2088 .cloned()
2089 .unwrap_or_default();
2090 let required = schema
2091 .get("required")
2092 .and_then(Value::as_array)
2093 .map(|items| {
2094 items
2095 .iter()
2096 .filter_map(Value::as_str)
2097 .map(str::to_owned)
2098 .collect::<std::collections::BTreeSet<_>>()
2099 })
2100 .unwrap_or_default();
2101
2102 let mut flag_to_field = BTreeMap::new();
2103 for (field_name, field_schema) in &properties {
2104 flag_to_field.insert(
2105 camel_to_kebab(field_name),
2106 (field_name.clone(), field_schema.clone()),
2107 );
2108 }
2109
2110 let mut input = Map::new();
2111 let mut index = 0;
2112 while index < argv.len() {
2113 let arg = &argv[index];
2114 if !arg.starts_with("--") {
2115 return Err(format!("Unexpected positional argument: \"{arg}\""));
2116 }
2117
2118 let raw_flag = &arg[2..];
2119 let (flag_name, negated) = raw_flag
2120 .strip_prefix("no-")
2121 .map(|name| (name, true))
2122 .unwrap_or((raw_flag, false));
2123 let Some((field_name, field_schema)) = flag_to_field.get(flag_name) else {
2124 return Err(format!("Unknown flag: --{raw_flag}"));
2125 };
2126 let field_type = json_schema_type(field_schema);
2127
2128 if negated {
2129 if field_type != Some("boolean") {
2130 return Err(format!("Unknown flag: --{raw_flag}"));
2131 }
2132 input.insert(field_name.clone(), Value::Bool(false));
2133 index += 1;
2134 continue;
2135 }
2136
2137 match field_type {
2138 Some("boolean") => {
2139 input.insert(field_name.clone(), Value::Bool(true));
2140 index += 1;
2141 }
2142 Some("number") | Some("integer") => {
2143 let value = argv
2144 .get(index + 1)
2145 .ok_or_else(|| format!("Flag --{raw_flag} requires a value"))?;
2146 let number = value
2147 .parse::<f64>()
2148 .map_err(|_| format!("Flag --{raw_flag} expects a number, got \"{value}\""))?;
2149 let number = serde_json::Number::from_f64(number).ok_or_else(|| {
2150 format!("Flag --{raw_flag} expects a finite number, got \"{value}\"")
2151 })?;
2152 input.insert(field_name.clone(), Value::Number(number));
2153 index += 2;
2154 }
2155 Some("array") => {
2156 let value = argv
2157 .get(index + 1)
2158 .ok_or_else(|| format!("Flag --{raw_flag} requires a value"))?;
2159 let item_type = field_schema.get("items").and_then(json_schema_type);
2160 let parsed_value = match item_type {
2161 Some("number") | Some("integer") => {
2162 let number = value.parse::<f64>().map_err(|_| {
2163 format!("Flag --{raw_flag} expects a number value, got \"{value}\"")
2164 })?;
2165 let number = serde_json::Number::from_f64(number).ok_or_else(|| {
2166 format!(
2167 "Flag --{raw_flag} expects a finite number value, got \"{value}\""
2168 )
2169 })?;
2170 Value::Number(number)
2171 }
2172 Some("boolean") => {
2173 let boolean = value.parse::<bool>().map_err(|_| {
2174 format!("Flag --{raw_flag} expects a boolean value, got \"{value}\"")
2175 })?;
2176 Value::Bool(boolean)
2177 }
2178 _ => Value::String(value.clone()),
2179 };
2180 input
2181 .entry(field_name.clone())
2182 .or_insert_with(|| Value::Array(Vec::new()))
2183 .as_array_mut()
2184 .expect("array field should always contain an array")
2185 .push(parsed_value);
2186 index += 2;
2187 }
2188 _ => {
2189 let value = argv
2190 .get(index + 1)
2191 .ok_or_else(|| format!("Flag --{raw_flag} requires a value"))?;
2192 input.insert(field_name.clone(), Value::String(value.clone()));
2193 index += 2;
2194 }
2195 }
2196 }
2197
2198 for field_name in required {
2199 if !input.contains_key(&field_name) {
2200 return Err(format!(
2201 "Missing required flag: --{}",
2202 camel_to_kebab(&field_name)
2203 ));
2204 }
2205 }
2206
2207 Ok(Value::Object(input))
2208}
2209
2210#[derive(Debug, Clone, PartialEq, Eq)]
2211struct HostFunctionInputSchemaViolation {
2212 path: String,
2213 expected: String,
2214 actual: String,
2215}
2216
2217impl HostFunctionInputSchemaViolation {
2218 fn new(
2219 path: impl Into<String>,
2220 expected: impl Into<String>,
2221 actual: impl Into<String>,
2222 ) -> Self {
2223 Self {
2224 path: path.into(),
2225 expected: expected.into(),
2226 actual: actual.into(),
2227 }
2228 }
2229}
2230
2231impl std::fmt::Display for HostFunctionInputSchemaViolation {
2232 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
2233 write!(
2234 f,
2235 "HostFunctionInputSchemaViolation at {}: expected {}, got {}",
2236 self.path, self.expected, self.actual
2237 )
2238 }
2239}
2240
2241fn validate_host_function_input(
2242 schema: &Value,
2243 input: &Value,
2244) -> Result<(), HostFunctionInputSchemaViolation> {
2245 validate_host_function_input_at_path(schema, input, "$")
2246}
2247
2248fn validate_host_function_input_at_path(
2249 schema: &Value,
2250 input: &Value,
2251 path: &str,
2252) -> Result<(), HostFunctionInputSchemaViolation> {
2253 if schema.is_null() || schema.as_object().is_some_and(|object| object.is_empty()) {
2254 return Ok(());
2255 }
2256 if let Some(branches) = schema.get("anyOf").and_then(Value::as_array) {
2257 return validate_schema_branches(branches, input, path, "anyOf");
2258 }
2259 if let Some(branches) = schema.get("oneOf").and_then(Value::as_array) {
2260 return validate_schema_branches(branches, input, path, "oneOf");
2261 }
2262 if let Some(enum_values) = schema.get("enum").and_then(Value::as_array) {
2263 if enum_values.iter().any(|candidate| candidate == input) {
2264 return Ok(());
2265 }
2266 return Err(HostFunctionInputSchemaViolation::new(
2267 path,
2268 format!(
2269 "one of {}",
2270 enum_values
2271 .iter()
2272 .map(compact_json)
2273 .collect::<Vec<_>>()
2274 .join(", ")
2275 ),
2276 describe_value(input),
2277 ));
2278 }
2279 if let Some(expected) = schema.get("const") {
2280 if expected == input {
2281 return Ok(());
2282 }
2283 return Err(HostFunctionInputSchemaViolation::new(
2284 path,
2285 format!("constant {}", compact_json(expected)),
2286 describe_value(input),
2287 ));
2288 }
2289
2290 match schema.get("type") {
2291 Some(Value::String(expected_type)) => {
2292 validate_typed_host_function_input(schema, input, path, expected_type)
2293 }
2294 Some(Value::Array(expected_types)) => {
2295 let mut first_error = None;
2296 for expected_type in expected_types.iter().filter_map(Value::as_str) {
2297 match validate_typed_host_function_input(schema, input, path, expected_type) {
2298 Ok(()) => return Ok(()),
2299 Err(error) if first_error.is_none() => first_error = Some(error),
2300 Err(_) => {}
2301 }
2302 }
2303 Err(first_error.unwrap_or_else(|| {
2304 HostFunctionInputSchemaViolation::new(
2305 path,
2306 describe_expected(schema),
2307 describe_value(input),
2308 )
2309 }))
2310 }
2311 Some(_) => Ok(()),
2312 None if has_object_keywords(schema) => {
2313 validate_typed_host_function_input(schema, input, path, "object")
2314 }
2315 None => Ok(()),
2316 }
2317}
2318
2319fn validate_schema_branches(
2320 branches: &[Value],
2321 input: &Value,
2322 path: &str,
2323 keyword: &str,
2324) -> Result<(), HostFunctionInputSchemaViolation> {
2325 let mut first_error = None;
2326 for branch in branches {
2327 match validate_host_function_input_at_path(branch, input, path) {
2328 Ok(()) => return Ok(()),
2329 Err(error) if first_error.is_none() => first_error = Some(error),
2330 Err(_) => {}
2331 }
2332 }
2333 Err(first_error.unwrap_or_else(|| {
2334 HostFunctionInputSchemaViolation::new(
2335 path,
2336 format!(
2337 "{keyword} branch ({})",
2338 branches
2339 .iter()
2340 .map(describe_expected)
2341 .collect::<Vec<_>>()
2342 .join(" | ")
2343 ),
2344 describe_value(input),
2345 )
2346 }))
2347}
2348
2349fn validate_typed_host_function_input(
2350 schema: &Value,
2351 input: &Value,
2352 path: &str,
2353 expected_type: &str,
2354) -> Result<(), HostFunctionInputSchemaViolation> {
2355 match expected_type {
2356 "null" if input.is_null() => Ok(()),
2357 "null" => Err(type_violation(path, expected_type, input)),
2358 "boolean" if input.is_boolean() => Ok(()),
2359 "boolean" => Err(type_violation(path, expected_type, input)),
2360 "string" => validate_string_host_function_input(schema, input, path),
2361 "number" => validate_number_host_function_input(schema, input, path, false),
2362 "integer" => validate_number_host_function_input(schema, input, path, true),
2363 "array" => validate_array_host_function_input(schema, input, path),
2364 "object" => validate_object_host_function_input(schema, input, path),
2365 _ => Ok(()),
2366 }
2367}
2368
2369fn validate_string_host_function_input(
2370 schema: &Value,
2371 input: &Value,
2372 path: &str,
2373) -> Result<(), HostFunctionInputSchemaViolation> {
2374 let Some(value) = input.as_str() else {
2375 return Err(type_violation(path, "string", input));
2376 };
2377 if let Some(min_length) = schema.get("minLength").and_then(Value::as_u64) {
2378 if value.chars().count() < min_length as usize {
2379 return Err(HostFunctionInputSchemaViolation::new(
2380 path,
2381 format!("string with minLength {min_length}"),
2382 format!("string length {}", value.chars().count()),
2383 ));
2384 }
2385 }
2386 if let Some(max_length) = schema.get("maxLength").and_then(Value::as_u64) {
2387 if value.chars().count() > max_length as usize {
2388 return Err(HostFunctionInputSchemaViolation::new(
2389 path,
2390 format!("string with maxLength {max_length}"),
2391 format!("string length {}", value.chars().count()),
2392 ));
2393 }
2394 }
2395 Ok(())
2396}
2397
2398fn validate_number_host_function_input(
2399 schema: &Value,
2400 input: &Value,
2401 path: &str,
2402 expect_integer: bool,
2403) -> Result<(), HostFunctionInputSchemaViolation> {
2404 let Some(number) = input.as_f64() else {
2405 return Err(type_violation(
2406 path,
2407 if expect_integer { "integer" } else { "number" },
2408 input,
2409 ));
2410 };
2411 if expect_integer && number.fract() != 0.0 {
2412 return Err(type_violation(path, "integer", input));
2413 }
2414 if let Some(minimum) = schema.get("minimum").and_then(Value::as_f64) {
2415 if number < minimum {
2416 return Err(HostFunctionInputSchemaViolation::new(
2417 path,
2418 format!(
2419 "{} >= {}",
2420 if expect_integer { "integer" } else { "number" },
2421 minimum
2422 ),
2423 compact_json(input),
2424 ));
2425 }
2426 }
2427 if let Some(minimum) = schema.get("exclusiveMinimum").and_then(Value::as_f64) {
2428 if number <= minimum {
2429 return Err(HostFunctionInputSchemaViolation::new(
2430 path,
2431 format!(
2432 "{} > {}",
2433 if expect_integer { "integer" } else { "number" },
2434 minimum
2435 ),
2436 compact_json(input),
2437 ));
2438 }
2439 }
2440 if let Some(maximum) = schema.get("maximum").and_then(Value::as_f64) {
2441 if number > maximum {
2442 return Err(HostFunctionInputSchemaViolation::new(
2443 path,
2444 format!(
2445 "{} <= {}",
2446 if expect_integer { "integer" } else { "number" },
2447 maximum
2448 ),
2449 compact_json(input),
2450 ));
2451 }
2452 }
2453 if let Some(maximum) = schema.get("exclusiveMaximum").and_then(Value::as_f64) {
2454 if number >= maximum {
2455 return Err(HostFunctionInputSchemaViolation::new(
2456 path,
2457 format!(
2458 "{} < {}",
2459 if expect_integer { "integer" } else { "number" },
2460 maximum
2461 ),
2462 compact_json(input),
2463 ));
2464 }
2465 }
2466 Ok(())
2467}
2468
2469fn validate_array_host_function_input(
2470 schema: &Value,
2471 input: &Value,
2472 path: &str,
2473) -> Result<(), HostFunctionInputSchemaViolation> {
2474 let Some(items) = input.as_array() else {
2475 return Err(type_violation(path, "array", input));
2476 };
2477 if let Some(min_items) = schema.get("minItems").and_then(Value::as_u64) {
2478 if items.len() < min_items as usize {
2479 return Err(HostFunctionInputSchemaViolation::new(
2480 path,
2481 format!("array with minItems {min_items}"),
2482 format!("array length {}", items.len()),
2483 ));
2484 }
2485 }
2486 if let Some(max_items) = schema.get("maxItems").and_then(Value::as_u64) {
2487 if items.len() > max_items as usize {
2488 return Err(HostFunctionInputSchemaViolation::new(
2489 path,
2490 format!("array with maxItems {max_items}"),
2491 format!("array length {}", items.len()),
2492 ));
2493 }
2494 }
2495 if let Some(item_schema) = schema.get("items") {
2496 for (index, item) in items.iter().enumerate() {
2497 validate_host_function_input_at_path(item_schema, item, &format!("{path}[{index}]"))?;
2498 }
2499 }
2500 Ok(())
2501}
2502
2503fn validate_object_host_function_input(
2504 schema: &Value,
2505 input: &Value,
2506 path: &str,
2507) -> Result<(), HostFunctionInputSchemaViolation> {
2508 let Some(object) = input.as_object() else {
2509 return Err(type_violation(path, "object", input));
2510 };
2511 let properties = schema
2512 .get("properties")
2513 .and_then(Value::as_object)
2514 .cloned()
2515 .unwrap_or_default();
2516 let required = schema
2517 .get("required")
2518 .and_then(Value::as_array)
2519 .cloned()
2520 .unwrap_or_default();
2521 for field in required.iter().filter_map(Value::as_str) {
2522 if !object.contains_key(field) {
2523 let field_path = format!("{path}.{field}");
2524 let expected = properties
2525 .get(field)
2526 .map(describe_expected)
2527 .unwrap_or_else(|| String::from("required value"));
2528 return Err(HostFunctionInputSchemaViolation::new(
2529 field_path,
2530 expected,
2531 "missing value",
2532 ));
2533 }
2534 }
2535 for (field, value) in object {
2536 let field_path = format!("{path}.{field}");
2537 if let Some(field_schema) = properties.get(field) {
2538 validate_host_function_input_at_path(field_schema, value, &field_path)?;
2539 continue;
2540 }
2541 match schema.get("additionalProperties") {
2542 Some(Value::Bool(false)) => {
2543 return Err(HostFunctionInputSchemaViolation::new(
2544 field_path,
2545 "no additional properties",
2546 describe_value(value),
2547 ));
2548 }
2549 Some(additional_schema) => {
2550 validate_host_function_input_at_path(additional_schema, value, &field_path)?;
2551 }
2552 None => {}
2553 }
2554 }
2555 Ok(())
2556}
2557
2558fn has_object_keywords(schema: &Value) -> bool {
2559 schema.get("properties").is_some()
2560 || schema.get("required").is_some()
2561 || schema.get("additionalProperties").is_some()
2562}
2563
2564fn type_violation(path: &str, expected: &str, input: &Value) -> HostFunctionInputSchemaViolation {
2565 HostFunctionInputSchemaViolation::new(path, expected, describe_value(input))
2566}
2567
2568fn describe_expected(schema: &Value) -> String {
2569 if let Some(enum_values) = schema.get("enum").and_then(Value::as_array) {
2570 return format!(
2571 "one of {}",
2572 enum_values
2573 .iter()
2574 .map(compact_json)
2575 .collect::<Vec<_>>()
2576 .join(", ")
2577 );
2578 }
2579 if let Some(expected) = schema.get("const") {
2580 return format!("constant {}", compact_json(expected));
2581 }
2582 match schema.get("type") {
2583 Some(Value::String(expected_type)) => expected_type.clone(),
2584 Some(Value::Array(expected_types)) => expected_types
2585 .iter()
2586 .filter_map(Value::as_str)
2587 .collect::<Vec<_>>()
2588 .join(" | "),
2589 _ if has_object_keywords(schema) => String::from("object"),
2590 _ => String::from("value"),
2591 }
2592}
2593
2594fn describe_value(value: &Value) -> String {
2595 match value {
2596 Value::Null => String::from("null"),
2597 Value::Bool(_) => String::from("boolean"),
2598 Value::Number(number) => {
2599 let is_integer = number.as_i64().is_some()
2600 || number.as_u64().is_some()
2601 || number.as_f64().is_some_and(|float| float.fract() == 0.0);
2602 if is_integer {
2603 String::from("integer")
2604 } else {
2605 String::from("number")
2606 }
2607 }
2608 Value::String(_) => String::from("string"),
2609 Value::Array(_) => String::from("array"),
2610 Value::Object(_) => String::from("object"),
2611 }
2612}
2613
2614fn compact_json(value: &Value) -> String {
2615 serde_json::to_string(value).unwrap_or_else(|_| String::from("<invalid json>"))
2616}
2617
2618fn list_host_functions_payload(host_functions: &[ResolvedHostFunctions]) -> Value {
2619 Value::Object(Map::from_iter([(
2620 String::from("hostFunctions"),
2621 Value::Array(
2622 host_functions
2623 .iter()
2624 .map(|collection| {
2625 json_object([
2626 ("name", Value::String(collection.name.clone())),
2627 (
2628 "functions",
2629 Value::Array(
2630 collection
2631 .functions
2632 .iter()
2633 .map(|host_function| Value::String(host_function.name.clone()))
2634 .collect(),
2635 ),
2636 ),
2637 ])
2638 })
2639 .collect(),
2640 ),
2641 )]))
2642}
2643
2644fn describe_host_functions_payload(
2645 host_functions: &[ResolvedHostFunctions],
2646 collection_name: &str,
2647) -> Result<Value, String> {
2648 let Some(collection) = host_functions
2649 .iter()
2650 .find(|collection| collection.name == collection_name)
2651 else {
2652 return Err(format!(
2653 "No host function collection \"{collection_name}\". Available: {}",
2654 host_functions_names(host_functions)
2655 ));
2656 };
2657 Ok(json_object([
2658 ("name", Value::String(collection.name.clone())),
2659 (
2660 "functions",
2661 Value::Object(Map::from_iter(collection.functions.iter().map(
2662 |host_function| {
2663 (
2664 host_function.name.clone(),
2665 json_object([
2666 (
2667 "description",
2668 Value::String(host_function.description.clone()),
2669 ),
2670 (
2671 "flags",
2672 Value::Array(describe_host_function_flags(
2673 &host_function.input_schema,
2674 )),
2675 ),
2676 ]),
2677 )
2678 },
2679 ))),
2680 ),
2681 ]))
2682}
2683
2684fn describe_host_function_payload(
2685 collection: &ResolvedHostFunctions,
2686 host_function_name: &str,
2687) -> Result<Value, String> {
2688 let Some(host_function) = collection
2689 .functions
2690 .iter()
2691 .find(|host_function| host_function.name == host_function_name)
2692 else {
2693 return Err(format!(
2694 "No host function \"{host_function_name}\" in collection \"{}\". Available: {}",
2695 collection.name,
2696 host_function_names(collection)
2697 ));
2698 };
2699 Ok(json_object([
2700 ("collection", Value::String(collection.name.clone())),
2701 ("function", Value::String(host_function_name.to_string())),
2702 (
2703 "description",
2704 Value::String(host_function.description.clone()),
2705 ),
2706 (
2707 "flags",
2708 Value::Array(describe_host_function_flags(&host_function.input_schema)),
2709 ),
2710 ("examples", Value::Array(Vec::new())),
2711 ]))
2712}
2713
2714fn describe_host_function_flags(schema: &Value) -> Vec<Value> {
2715 let properties = schema
2716 .get("properties")
2717 .and_then(Value::as_object)
2718 .cloned()
2719 .unwrap_or_default();
2720 let required = schema
2721 .get("required")
2722 .and_then(Value::as_array)
2723 .map(|items| {
2724 items
2725 .iter()
2726 .filter_map(Value::as_str)
2727 .map(str::to_owned)
2728 .collect::<std::collections::BTreeSet<_>>()
2729 })
2730 .unwrap_or_default();
2731 properties
2732 .into_iter()
2733 .map(|(field_name, field_schema)| {
2734 json_object([
2735 (
2736 "name",
2737 Value::String(format!("--{}", camel_to_kebab(&field_name))),
2738 ),
2739 (
2740 "type",
2741 Value::String(describe_host_function_flag_type(&field_schema)),
2742 ),
2743 ("required", Value::Bool(required.contains(&field_name))),
2744 ])
2745 })
2746 .collect()
2747}
2748
2749fn describe_host_function_flag_type(schema: &Value) -> String {
2750 match json_schema_type(schema) {
2751 Some("array") => {
2752 let item_type = schema
2753 .get("items")
2754 .and_then(json_schema_type)
2755 .unwrap_or("string");
2756 format!("{item_type}[]")
2757 }
2758 Some("string") => schema
2759 .get("enum")
2760 .and_then(Value::as_array)
2761 .map(|values| values.iter().filter_map(Value::as_str).collect::<Vec<_>>())
2762 .filter(|values| !values.is_empty())
2763 .map(|values| values.join("|"))
2764 .unwrap_or_else(|| String::from("string")),
2765 Some(other) => other.to_string(),
2766 None => String::from("string"),
2767 }
2768}
2769
2770fn host_functions_names(host_functions: &[ResolvedHostFunctions]) -> String {
2771 host_functions
2772 .iter()
2773 .map(|collection| collection.name.clone())
2774 .collect::<Vec<_>>()
2775 .join(", ")
2776}
2777
2778fn host_function_names(collection: &ResolvedHostFunctions) -> String {
2779 collection
2780 .functions
2781 .iter()
2782 .map(|host_function| host_function.name.clone())
2783 .collect::<Vec<_>>()
2784 .join(", ")
2785}
2786
2787fn is_help_flag(value: &str) -> bool {
2788 matches!(value, "--help" | "-h")
2789}
2790
2791fn json_schema_type(schema: &Value) -> Option<&str> {
2792 schema.get("type").and_then(Value::as_str)
2793}
2794
2795fn camel_to_kebab(value: &str) -> String {
2796 let mut output = String::new();
2797 for (index, ch) in value.chars().enumerate() {
2798 if ch.is_ascii_uppercase() && index > 0 {
2799 output.push('-');
2800 }
2801 output.push(ch.to_ascii_lowercase());
2802 }
2803 output
2804}
2805
2806fn normalize_guest_path(path: String) -> String {
2807 let absolute = path.starts_with('/');
2808 let mut parts = Vec::new();
2809 for part in path.split('/') {
2810 match part {
2811 "" | "." => {}
2812 ".." => {
2813 parts.pop();
2814 }
2815 _ => parts.push(part),
2816 }
2817 }
2818 let normalized = parts.join("/");
2819 if absolute {
2820 format!("/{normalized}")
2821 } else {
2822 normalized
2823 }
2824}
2825
2826fn json_object<const N: usize>(entries: [(&str, Value); N]) -> Value {
2827 Value::Object(Map::from_iter(
2828 entries
2829 .into_iter()
2830 .map(|(key, value)| (key.to_string(), value)),
2831 ))
2832}
2833
2834pub(crate) fn build_package_descriptors(config: &AgentOsConfig) -> Vec<wire::PackageDescriptor> {
2837 config
2838 .packages
2839 .iter()
2840 .map(|package| wire::PackageDescriptor {
2841 path: package.path.clone(),
2842 })
2843 .collect()
2844}
2845
2846pub(crate) fn serialize_mounts(
2847 config: &AgentOsConfig,
2848) -> Result<Vec<wire::MountDescriptor>, ClientError> {
2849 config
2850 .mounts
2851 .iter()
2852 .map(|mount| match mount {
2853 MountConfig::Native {
2854 path,
2855 plugin,
2856 guest_source,
2857 guest_fstype,
2858 read_only,
2859 } => {
2860 let plugin_config = plugin
2861 .config
2862 .clone()
2863 .unwrap_or_else(|| serde_json::Value::Object(Default::default()));
2864 Ok(wire::MountDescriptor {
2865 guest_path: path.clone(),
2866 guest_source: guest_source.clone().unwrap_or_else(|| plugin.id.clone()),
2867 guest_fstype: guest_fstype.clone().unwrap_or_else(|| plugin.id.clone()),
2868 read_only: *read_only,
2869 plugin: wire::MountPluginDescriptor {
2870 id: plugin.id.clone(),
2871 config: json_utf8(&plugin_config, "native mount plugin config")?,
2872 },
2873 })
2874 }
2875 MountConfig::Plain { .. } => Err(ClientError::Sidecar(
2876 "plain mounts cannot be configured during Rust client VM creation".to_string(),
2877 )),
2878 MountConfig::Overlay { .. } => Err(ClientError::Sidecar(
2879 "overlay mounts cannot be configured during Rust client VM creation".to_string(),
2880 )),
2881 })
2882 .collect()
2883}
2884
2885fn json_utf8(value: &serde_json::Value, context: &str) -> Result<String, ClientError> {
2886 serde_json::to_string(value)
2887 .map_err(|error| ClientError::Sidecar(format!("failed to serialize {context}: {error}")))
2888}
2889
2890fn operation_wildcard_if_omitted(values: &Option<Vec<String>>) -> Vec<String> {
2891 values.clone().unwrap_or_else(|| vec!["*".to_string()])
2892}
2893
2894fn resource_wildcard_if_omitted(values: &Option<Vec<String>>) -> Vec<String> {
2895 values.clone().unwrap_or_else(|| vec!["**".to_string()])
2896}
2897
2898fn wire_ownership_vm_id(ownership: &wire::OwnershipScope) -> Option<&str> {
2900 match ownership {
2901 wire::OwnershipScope::VmOwnership(ownership) => Some(ownership.vm_id.as_str()),
2902 wire::OwnershipScope::ConnectionOwnership(_)
2903 | wire::OwnershipScope::SessionOwnership(_) => None,
2904 }
2905}
2906
2907fn vm_dispose_response(vm_id: &str, response: wire::ResponsePayload) -> Result<(), ClientError> {
2908 match response {
2909 wire::ResponsePayload::VmDisposedResponse(disposed) if disposed.vm_id == vm_id => Ok(()),
2910 wire::ResponsePayload::RejectedResponse(rejected) => Err(rejected_to_error(rejected)),
2911 other => Err(ClientError::Sidecar(format!(
2912 "unexpected dispose_vm response: {other:?}"
2913 ))),
2914 }
2915}
2916
2917fn retain_shutdown_failure(result: &mut Result<(), ClientError>, cleanup: Result<(), ClientError>) {
2918 if let Err(error) = cleanup {
2919 eprintln!("agentOS secondary shutdown cleanup failed: {error}");
2920 if result.is_ok() {
2921 *result = Err(error);
2922 }
2923 }
2924}
2925
2926fn concurrent_vm_mutation_error() -> ClientError {
2927 ClientError::Sidecar(String::from(
2928 "invalid_state: another VM mount or software mutation is already in progress; wait for it to finish before retrying",
2929 ))
2930}
2931
2932fn try_vm_mutation(
2933 operation: &tokio::sync::Mutex<()>,
2934) -> Result<tokio::sync::MutexGuard<'_, ()>, ClientError> {
2935 operation
2936 .try_lock()
2937 .map_err(|_| concurrent_vm_mutation_error())
2938}
2939
2940fn rejected_to_error(rejected: wire::RejectedResponse) -> ClientError {
2942 ClientError::from_rejection(rejected)
2943}
2944
2945#[cfg(test)]
2946mod tests {
2947 use std::collections::{BTreeMap, HashMap};
2948 use std::sync::{Arc, Mutex};
2949
2950 use super::{
2951 run_host_callback, serialize_create_vm_config_for_sidecar,
2952 serialize_root_filesystem_config_for_sidecar, try_vm_mutation, vm_host_functions,
2953 wire_vm_ownership, AgentOsCreationTask, VmHostFunctionRegistry,
2954 };
2955 use crate::config::{
2956 AgentOsConfig, AgentOsLimits, FsPermissionRule, FsPermissions, HostFunctionLimits,
2957 HttpLimits, JsRuntimeLimits, MountPlugin, PatternPermissions, PermissionMode, Permissions,
2958 PythonLimits, ResolvedHostFunction, ResolvedHostFunctions, ResourceLimits,
2959 RootFilesystemConfig, RootFilesystemKind, RootFilesystemMode, RootLowerInput,
2960 RulePermissions, WasmLimits,
2961 };
2962 use crate::fs::{
2963 DirEntryType, FilesystemEntry, FilesystemEntryEncoding, FilesystemSnapshotEntries,
2964 FilesystemSnapshotExport, RootSnapshotExport, SnapshotExportKind,
2965 };
2966 use agentos_vm_config::{
2967 FsPermissionScope, PatternPermissionScope, PermissionMode as ConfigPermissionMode,
2968 RootFilesystemEntryKind, RootFilesystemLowerDescriptor,
2969 RootFilesystemMode as ConfigRootFilesystemMode, VmDefaultsProfile,
2970 };
2971 use serde_json::json;
2972
2973 #[tokio::test]
2974 async fn command_shaped_input_cannot_redirect_an_authorized_host_callback() {
2975 let vm_id = format!("host-callback-test-{}", uuid::Uuid::new_v4());
2976 let calls = Arc::new(Mutex::new(Vec::new()));
2977 let allowed_calls = Arc::clone(&calls);
2978 let danger_calls = Arc::clone(&calls);
2979 let allowed = ResolvedHostFunction {
2980 name: String::from("allowed"),
2981 description: String::new(),
2982 input_schema: json!({ "type": "object" }),
2983 timeout_ms: None,
2984 execute: Arc::new(move |_| {
2985 allowed_calls.lock().expect("calls lock").push("allowed");
2986 Box::pin(async { Ok(json!("allowed")) })
2987 }),
2988 };
2989 let danger = ResolvedHostFunction {
2990 name: String::from("danger"),
2991 description: String::new(),
2992 input_schema: json!({ "type": "object" }),
2993 timeout_ms: None,
2994 execute: Arc::new(move |_| {
2995 danger_calls.lock().expect("calls lock").push("danger");
2996 Box::pin(async { Ok(json!("danger")) })
2997 }),
2998 };
2999 let host_function_map = HashMap::from([
3000 (String::from("math:allowed"), allowed.clone()),
3001 (String::from("math:danger"), danger.clone()),
3002 ]);
3003 let _ = vm_host_functions().insert(
3004 vm_id.clone(),
3005 Arc::new(VmHostFunctionRegistry {
3006 host_functions: vec![ResolvedHostFunctions {
3007 name: String::from("math"),
3008 functions: vec![allowed, danger],
3009 }],
3010 host_function_map,
3011 }),
3012 );
3013
3014 let response = run_host_callback(
3015 &wire_vm_ownership("connection", "session", &vm_id),
3016 super::wire::HostCallbackRequest {
3017 invocation_id: String::from("invocation"),
3018 callback_key: String::from("math:allowed"),
3019 input: json!({
3020 "type": "command",
3021 "command": "agentos-math",
3022 "args": ["danger"],
3023 "cwd": "/workspace"
3024 })
3025 .to_string(),
3026 timeout_ms: 1_000,
3027 },
3028 )
3029 .await;
3030 vm_host_functions().remove(&vm_id);
3031
3032 assert_eq!(response.error, None);
3033 assert_eq!(response.result.as_deref(), Some("\"allowed\""));
3034 assert_eq!(*calls.lock().expect("calls lock"), vec!["allowed"]);
3035 }
3036
3037 #[tokio::test]
3038 async fn cancelled_creation_waiter_keeps_the_owned_task_running() {
3039 let (release, wait_for_release) = tokio::sync::oneshot::channel();
3040 let (completed, wait_for_completion) = tokio::sync::oneshot::channel();
3041 let guard = AgentOsCreationTask {
3042 task: Some(tokio::spawn(async move {
3043 let _ = wait_for_release.await;
3044 let _ = completed.send(());
3045 Err(crate::ClientError::Sidecar(String::from(
3046 "injected creation failure",
3047 )))
3048 })),
3049 };
3050
3051 drop(guard);
3052 release.send(()).unwrap();
3053 tokio::time::timeout(std::time::Duration::from_secs(1), wait_for_completion)
3054 .await
3055 .expect("detached creation task remains owned")
3056 .expect("creation task reports completion");
3057 }
3058
3059 #[tokio::test]
3060 async fn concurrent_vm_mutations_fail_without_queueing() {
3061 let operation = tokio::sync::Mutex::new(());
3062 let active = try_vm_mutation(&operation).expect("first mutation is admitted");
3063 let error = try_vm_mutation(&operation).expect_err("second mutation is rejected");
3064 assert!(error.to_string().contains("invalid_state"));
3065 drop(active);
3066 assert!(try_vm_mutation(&operation).is_ok());
3067 }
3068
3069 #[test]
3070 fn vm_dispose_preserves_timeout_details_and_original_error_after_cleanup() {
3071 use super::{retain_shutdown_failure, vm_dispose_response, wire, ClientError};
3072
3073 let mut result = vm_dispose_response(
3074 "vm-timeout",
3075 wire::ResponsePayload::RejectedResponse(wire::RejectedResponse {
3076 code: "timeout".into(),
3077 message: "SQLite close is unconfirmed; raise limits.reactor.shutdownDeadlineMs"
3078 .into(),
3079 limit_name: Some("reactor.shutdownDeadlineMs".into()),
3080 configured_limit: Some(5_000),
3081 current_usage: None,
3082 requested: None,
3083 unit: Some("milliseconds".into()),
3084 scope: Some("vm".into()),
3085 vm_id: Some("vm-timeout".into()),
3086 session_generation: None,
3087 capability_id: None,
3088 operation: Some("vm.dispose".into()),
3089 configuration_path: Some("limits.reactor.shutdownDeadlineMs".into()),
3090 retryable: Some(false),
3091 errno: Some("ETIMEDOUT".into()),
3092 }),
3093 );
3094 retain_shutdown_failure(
3095 &mut result,
3096 Err(ClientError::Sidecar("secondary cleanup".into())),
3097 );
3098 let ClientError::OperationTimedOut { message, details } = result.unwrap_err() else {
3099 panic!("expected original typed timeout");
3100 };
3101 assert!(message.contains("SQLite close is unconfirmed"));
3102 assert_eq!(
3103 details.limit_name.as_deref(),
3104 Some("reactor.shutdownDeadlineMs")
3105 );
3106 assert_eq!(details.configured_limit, Some(5_000));
3107 assert_eq!(details.unit.as_deref(), Some("milliseconds"));
3108 assert_eq!(details.vm_id.as_deref(), Some("vm-timeout"));
3109 assert_eq!(details.operation.as_deref(), Some("vm.dispose"));
3110 assert_eq!(
3111 details.configuration_path.as_deref(),
3112 Some("limits.reactor.shutdownDeadlineMs")
3113 );
3114 assert_eq!(details.retryable, Some(false));
3115 assert_eq!(details.errno.as_deref(), Some("ETIMEDOUT"));
3116
3117 assert!(vm_dispose_response(
3118 "vm-ok",
3119 wire::ResponsePayload::VmDisposedResponse(wire::VmDisposedResponse {
3120 vm_id: "vm-ok".into()
3121 })
3122 )
3123 .is_ok());
3124 assert!(vm_dispose_response(
3125 "vm-ok",
3126 wire::ResponsePayload::VmDisposedResponse(wire::VmDisposedResponse {
3127 vm_id: "other-vm".into()
3128 })
3129 )
3130 .is_err());
3131 let mut result = Ok(());
3132 retain_shutdown_failure(&mut result, Err(ClientError::Sidecar("cleanup".into())));
3133 assert!(result.is_err());
3134 }
3135
3136 #[test]
3137 fn create_vm_omits_permissions_so_sidecar_applies_product_defaults() {
3138 let config = serialize_create_vm_config_for_sidecar(&AgentOsConfig::default()).unwrap();
3139 assert_eq!(config.defaults_profile, Some(VmDefaultsProfile::AgentOs));
3140 assert_eq!(config.permissions, None);
3141 }
3142
3143 #[test]
3144 fn create_vm_sends_only_explicit_permission_overrides() {
3145 let config = serialize_create_vm_config_for_sidecar(&AgentOsConfig {
3146 permissions: Some(Permissions {
3147 network: Some(PatternPermissions::Mode(PermissionMode::Deny)),
3148 ..Default::default()
3149 }),
3150 ..Default::default()
3151 })
3152 .unwrap();
3153 let policy = config.permissions.expect("explicit policy");
3154 assert_eq!(
3155 policy.network,
3156 Some(PatternPermissionScope::Mode(ConfigPermissionMode::Deny))
3157 );
3158 assert_eq!(policy.child_process, None);
3159 }
3160
3161 #[test]
3162 fn create_vm_expands_omitted_rule_fields_to_domain_wildcards() {
3163 let config = serialize_create_vm_config_for_sidecar(&AgentOsConfig {
3164 permissions: Some(Permissions {
3165 fs: Some(FsPermissions::Rules(RulePermissions {
3166 default: Some(PermissionMode::Deny),
3167 rules: vec![FsPermissionRule {
3168 mode: PermissionMode::Allow,
3169 operations: None,
3170 paths: Some(vec!["/workspace/**".to_string()]),
3171 }],
3172 })),
3173 ..Default::default()
3174 }),
3175 ..Default::default()
3176 })
3177 .unwrap();
3178 let policy = config.permissions.expect("explicit fs policy");
3179 let Some(FsPermissionScope::Rules(rules)) = policy.fs else {
3180 panic!("expected fs rule set");
3181 };
3182 assert_eq!(rules.default, Some(ConfigPermissionMode::Deny));
3183 assert_eq!(rules.rules[0].operations, vec!["*"]);
3184 assert_eq!(rules.rules[0].paths, vec!["/workspace/**"]);
3185
3186 let config = serialize_create_vm_config_for_sidecar(&AgentOsConfig {
3187 permissions: Some(Permissions {
3188 network: Some(PatternPermissions::Rules(RulePermissions {
3189 default: Some(PermissionMode::Allow),
3190 rules: vec![crate::config::PatternPermissionRule {
3191 mode: PermissionMode::Deny,
3192 operations: None,
3193 patterns: None,
3194 }],
3195 })),
3196 ..Default::default()
3197 }),
3198 ..Default::default()
3199 })
3200 .unwrap();
3201 let policy = config.permissions.expect("explicit network policy");
3202 let Some(PatternPermissionScope::Rules(rules)) = policy.network else {
3203 panic!("expected network rule set");
3204 };
3205 assert_eq!(rules.default, Some(ConfigPermissionMode::Allow));
3206 assert_eq!(rules.rules[0].operations, vec!["*"]);
3207 assert_eq!(rules.rules[0].patterns, vec!["**"]);
3208 }
3209
3210 #[test]
3211 fn root_filesystem_serializer_preserves_configured_descriptor() {
3212 let (descriptor, native_root) =
3213 serialize_root_filesystem_config_for_sidecar(&RootFilesystemConfig {
3214 mode: Some(RootFilesystemMode::ReadOnly),
3215 disable_default_base_layer: true,
3216 lowers: vec![
3217 RootLowerInput::BundledBaseFilesystem,
3218 RootLowerInput::SnapshotExport(RootSnapshotExport {
3219 kind: SnapshotExportKind::SnapshotExport,
3220 source: FilesystemSnapshotExport {
3221 format: "agentos-filesystem-snapshot-v1".to_string(),
3222 filesystem: FilesystemSnapshotEntries {
3223 entries: vec![
3224 FilesystemEntry {
3225 path: "/bin/run".to_string(),
3226 entry_type: DirEntryType::File,
3227 mode: "0755".to_string(),
3228 uid: 1000,
3229 gid: 1000,
3230 content: Some("#!/bin/sh".to_string()),
3231 encoding: Some(FilesystemEntryEncoding::Utf8),
3232 target: None,
3233 },
3234 FilesystemEntry {
3235 path: "/link".to_string(),
3236 entry_type: DirEntryType::Symlink,
3237 mode: "0777".to_string(),
3238 uid: 0,
3239 gid: 0,
3240 content: None,
3241 encoding: None,
3242 target: Some("/bin/run".to_string()),
3243 },
3244 ],
3245 },
3246 },
3247 }),
3248 ],
3249 ..Default::default()
3250 })
3251 .expect("serialize root filesystem");
3252
3253 assert!(native_root.is_none());
3254 assert_eq!(descriptor.mode, ConfigRootFilesystemMode::ReadOnly);
3255 assert!(descriptor.disable_default_base_layer);
3256 assert_eq!(descriptor.bootstrap_entries, Vec::new());
3257 assert!(matches!(
3258 descriptor.lowers[0],
3259 RootFilesystemLowerDescriptor::BundledBaseFilesystem
3260 ));
3261
3262 let RootFilesystemLowerDescriptor::Snapshot { entries } = &descriptor.lowers[1] else {
3263 panic!("expected snapshot lower");
3264 };
3265 assert_eq!(entries[0].path, "/bin/run");
3266 assert_eq!(entries[0].kind, RootFilesystemEntryKind::File);
3267 assert_eq!(entries[0].mode, Some(0o755));
3268 assert!(entries[0].executable);
3269 assert_eq!(entries[1].kind, RootFilesystemEntryKind::Symlink);
3270 assert_eq!(entries[1].target.as_deref(), Some("/bin/run"));
3271 }
3272
3273 #[test]
3274 fn create_vm_config_preserves_native_root_config() {
3275 let config = serialize_create_vm_config_for_sidecar(&AgentOsConfig {
3276 root_filesystem: RootFilesystemConfig {
3277 kind: RootFilesystemKind::Native,
3278 mode: Some(RootFilesystemMode::ReadOnly),
3279 native_plugin: Some(MountPlugin {
3280 id: "sqlite_vfs".to_string(),
3281 config: Some(serde_json::json!({
3282 "databasePath": "/tmp/agentos-root.sqlite"
3283 })),
3284 }),
3285 ..Default::default()
3286 },
3287 ..Default::default()
3288 })
3289 .expect("serialize create VM config");
3290 let native_root = config.native_root.expect("native root config");
3291
3292 assert_eq!(native_root.plugin.id, "sqlite_vfs");
3293 assert_eq!(
3294 native_root.plugin.config,
3295 serde_json::json!({ "databasePath": "/tmp/agentos-root.sqlite" })
3296 );
3297 assert!(native_root.read_only);
3298 }
3299
3300 #[test]
3301 fn create_vm_config_preserves_environment_and_timer_policy() {
3302 let environment = BTreeMap::from([
3303 (String::from("EMPTY"), String::new()),
3304 (String::from("PATH"), String::from("/opt/agentos/bin")),
3305 ]);
3306 let config = serialize_create_vm_config_for_sidecar(&AgentOsConfig {
3307 environment: Some(environment.clone()),
3308 high_resolution_time: Some(false),
3309 ..Default::default()
3310 })
3311 .expect("serialize create VM config");
3312
3313 assert_eq!(config.env, Some(environment));
3314 let js_runtime = config.js_runtime.expect("explicit timer policy");
3315 assert_eq!(js_runtime.allowed_builtins, None);
3316 assert_eq!(js_runtime.high_resolution_time, Some(false));
3317
3318 let empty = serialize_create_vm_config_for_sidecar(&AgentOsConfig {
3319 environment: Some(BTreeMap::new()),
3320 ..Default::default()
3321 })
3322 .expect("serialize empty environment");
3323 assert_eq!(empty.env, Some(BTreeMap::new()));
3324
3325 let defaulted = serialize_create_vm_config_for_sidecar(&AgentOsConfig::default())
3326 .expect("serialize default environment");
3327 assert_eq!(defaulted.env, None);
3328 }
3329
3330 #[test]
3331 fn create_vm_config_preserves_typed_limits() {
3332 let config = serialize_create_vm_config_for_sidecar(&AgentOsConfig {
3333 limits: Some(AgentOsLimits {
3334 agentos_packages: Some(crate::config::AgentOsPackageLimits {
3335 max_mounts: Some(8192),
3336 }),
3337 resources: Some(ResourceLimits {
3338 max_processes: Some(7),
3339 max_filesystem_bytes: Some(4096),
3340 ..Default::default()
3341 }),
3342 http: Some(HttpLimits {
3343 max_fetch_response_bytes: Some(1024),
3344 }),
3345 tls: Some(crate::TlsLimits {
3346 max_buffered_bytes: Some(2048),
3347 }),
3348 execution: Some(crate::ExecutionLimits {
3349 completed_ttl_ms: Some(60_000),
3350 max_completed_executions: Some(128),
3351 live_execution_warning_threshold: Some(32),
3352 }),
3353 host_functions: Some(HostFunctionLimits {
3354 default_timeout_ms: Some(500),
3355 max_registered_functions_per_vm: Some(12),
3356 ..Default::default()
3357 }),
3358 js_runtime: Some(JsRuntimeLimits {
3359 v8_heap_limit_mb: Some(64),
3360 sync_rpc_wait_timeout_ms: Some(2_000),
3361 cpu_time_limit_ms: Some(30_000),
3362 wall_clock_limit_ms: Some(0),
3363 import_cache_materialize_timeout_ms: Some(30_000),
3364 ..Default::default()
3365 }),
3366 python: Some(PythonLimits {
3367 max_old_space_mb: Some(256),
3368 ..Default::default()
3369 }),
3370 wasm: Some(WasmLimits {
3371 prewarm_timeout_ms: Some(30_000),
3372 runner_heap_limit_mb: Some(2_048),
3373 runner_cpu_time_limit_ms: Some(60_000),
3374 ..Default::default()
3375 }),
3376 ..Default::default()
3377 }),
3378 ..Default::default()
3379 })
3380 .expect("serialize create VM config");
3381 let limits = config.limits.expect("limits config");
3382
3383 assert_eq!(
3384 limits.tls.expect("TLS limits").max_buffered_bytes,
3385 Some(2048)
3386 );
3387 let execution = limits.execution.expect("execution limits");
3388 assert_eq!(execution.completed_ttl_ms, Some(60_000));
3389 assert_eq!(execution.max_completed_executions, Some(128));
3390 assert_eq!(execution.live_execution_warning_threshold, Some(32));
3391 assert_eq!(
3392 limits.agentos_packages.expect("package limits").max_mounts,
3393 Some(8192)
3394 );
3395 let resources = limits.resources.expect("resource limits");
3396 assert_eq!(resources.max_processes, Some(7));
3397 assert_eq!(resources.max_filesystem_bytes, Some(4096));
3398 assert_eq!(
3399 limits.http.expect("http limits").max_fetch_response_bytes,
3400 Some(1024)
3401 );
3402 assert_eq!(
3403 limits
3404 .host_functions
3405 .as_ref()
3406 .expect("host_function limits")
3407 .default_timeout_ms,
3408 Some(500)
3409 );
3410 assert_eq!(
3411 limits
3412 .host_functions
3413 .expect("host_function limits")
3414 .max_registered_functions_per_vm,
3415 Some(12)
3416 );
3417 assert_eq!(
3418 limits
3419 .js_runtime
3420 .as_ref()
3421 .expect("js runtime limits")
3422 .v8_heap_limit_mb,
3423 Some(64)
3424 );
3425 let js_runtime = limits.js_runtime.expect("js runtime limits");
3426 assert_eq!(js_runtime.sync_rpc_wait_timeout_ms, Some(2_000));
3427 assert_eq!(js_runtime.cpu_time_limit_ms, Some(30_000));
3428 assert_eq!(js_runtime.wall_clock_limit_ms, Some(0));
3429 assert_eq!(js_runtime.import_cache_materialize_timeout_ms, Some(30_000));
3430 assert_eq!(
3431 limits.python.expect("python limits").max_old_space_mb,
3432 Some(256)
3433 );
3434 let wasm = limits.wasm.expect("wasm limits");
3435 assert_eq!(wasm.prewarm_timeout_ms, Some(30_000));
3436 assert_eq!(wasm.runner_heap_limit_mb, Some(2_048));
3437 assert_eq!(wasm.runner_cpu_time_limit_ms, Some(60_000));
3438 }
3439
3440 #[test]
3441 fn tls_execution_limit_validation_stays_in_shared_vm_config() {
3442 let defaulted = serialize_create_vm_config_for_sidecar(&AgentOsConfig::default()).unwrap();
3443 assert!(defaulted.limits.is_none());
3444 let limits = AgentOsLimits {
3445 tls: Some(crate::TlsLimits::default()),
3446 execution: Some(crate::ExecutionLimits::default()),
3447 ..Default::default()
3448 };
3449 let config = AgentOsConfig {
3450 limits: Some(limits),
3451 ..Default::default()
3452 };
3453 super::validate_config(&config).unwrap();
3454 let empty = serialize_create_vm_config_for_sidecar(&config)
3455 .unwrap()
3456 .limits
3457 .unwrap();
3458 assert_eq!(empty.tls.unwrap().max_buffered_bytes, None);
3459 assert_eq!(
3460 empty.execution.unwrap(),
3461 agentos_vm_config::ExecutionLimitsConfig::default()
3462 );
3463
3464 for (group, field) in [
3465 ("tls", "maxBufferedBytes"),
3466 ("execution", "completedTtlMs"),
3467 ("execution", "maxCompletedExecutions"),
3468 ("execution", "liveExecutionWarningThreshold"),
3469 ] {
3470 for value in [0, u64::MAX] {
3471 let mut limits: AgentOsLimits =
3472 serde_json::from_value(serde_json::json!({group: {field: value}})).unwrap();
3473 limits.resources = Some(ResourceLimits {
3476 max_socket_buffered_bytes: Some(4096),
3477 ..Default::default()
3478 });
3479 let config = AgentOsConfig {
3480 limits: Some(limits),
3481 ..Default::default()
3482 };
3483 let error = super::validate_config(&config).unwrap_err().to_string();
3484 assert!(
3485 error.contains(&format!("limits.{group}.{field}")),
3486 "{error}"
3487 );
3488 }
3489 }
3490 }
3491}