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