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