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