From 150b793962dad1650b71ed08c62df4a49556cb09 Mon Sep 17 00:00:00 2001 From: Nathan Flurry Date: Mon, 21 Sep 2026 00:16:52 -0700 Subject: [PATCH] refactor(agentos): drop the host function wrappers, names, and descriptions Host functions are now a record of collections. The keys name everything and a function needs only `inputSchema` and `execute`. hostFunctions: { store: { listOrders: { inputSchema, execute }, }, } Removed from the public API: the `hostFunction()` and `hostFunctions()` wrappers, `HostFunctions.name`, `HostFunctions.description`, the `functions:` nesting key, `HostFunction.description`, `MAX_HOST_FUNCTION_DESCRIPTION_LENGTH`, `validateHostFunctions`, and `HostFunctionDefinition.commandAliases`. The description now comes from the input schema's `.describe()`, so it is one optional field instead of two required ones. The sidecar keeps the length cap and no longer requires a description at all. `NodeRuntime` now takes the same shape. The host function module moves down into runtime-core so both layers share one implementation, which is a breaking change for guest code: a function is reached as `agentos- ` rather than as a top-level command named after the function. Co-Authored-By: Claude Opus 5 --- crates/client/src/agent_os.rs | 44 +-- crates/client/src/config.rs | 143 ++++++++- crates/client/src/error.rs | 6 + crates/client/src/lib.rs | 2 +- crates/client/src/session.rs | 23 +- crates/client/tests/os_instructions_e2e.rs | 45 +-- .../native-sidecar-core/src/host_functions.rs | 14 +- docs/content/docs/host-functions.mdx | 11 +- examples/agent-to-agent/server.ts | 28 +- .../crash-course/agent-to-agent-server.ts | 69 ++--- examples/crash-course/sandbox-stubs.d.ts | 4 +- examples/crash-course/sandbox.ts | 2 +- examples/embedded/agent-to-agent.ts | 38 ++- examples/embedded/host-functions.ts | 31 +- examples/host-functions/exec-bash.ts | 27 +- examples/host-functions/exec-javascript.ts | 27 +- examples/host-functions/exec-python.ts | 27 +- examples/host-functions/server.ts | 55 ++-- examples/js-code-mode/src/index.ts | 31 +- examples/quickstart/host-functions/index.ts | 65 ++-- .../tests/vm-integration.test.ts | 14 +- packages/agentos/src/actor.ts | 15 +- packages/core/CLAUDE.md | 2 +- packages/core/src/agent-os.ts | 126 +++++--- packages/core/src/host-functions.ts | 86 ------ packages/core/src/index.ts | 55 ++-- packages/core/src/options-schema.ts | 25 +- packages/core/src/sandbox.ts | 289 +++++++++--------- packages/core/src/types.ts | 9 +- .../core/tests/acp-reactor-regression.test.ts | 35 +-- .../tests/host-function-permissions.test.ts | 208 ++++++------- .../tests/host-function-reference.test.ts | 52 ++-- .../core/tests/host-functions-zod.test.ts | 8 +- packages/core/tests/host-functions.test.ts | 162 +++++----- packages/core/tests/migration-parity.test.ts | 34 +-- packages/core/tests/options-schema.test.ts | 20 +- .../core/tests/public-api-exports.test.ts | 18 +- .../core/tests/sandbox-integration.test.ts | 22 +- ...car-host-function-dispatch.nightly.test.ts | 24 +- packages/runtime-core/package.json | 13 +- .../src/host-functions-zod.ts | 15 +- packages/runtime-core/src/host-functions.ts | 161 ++++++++++ .../src/node-runtime-options-schema.ts | 19 +- packages/runtime-core/src/node-runtime.ts | 65 ++-- packages/runtime-core/src/test-runtime.ts | 130 ++++---- pnpm-lock.yaml | 3 + .../docs/content/docs/api-reference.mdx | 1 - .../docs/content/docs/host-functions.mdx | 19 +- .../examples/host-functions/src/index.ts | 42 ++- secure-exec/src/index.ts | 37 ++- secure-exec/src/runtime.ts | 34 ++- secure-exec/src/typescript.ts | 42 ++- secure-exec/tests/secure-exec.test.ts | 21 +- 53 files changed, 1383 insertions(+), 1115 deletions(-) delete mode 100644 packages/core/src/host-functions.ts rename packages/{core => runtime-core}/src/host-functions-zod.ts (96%) create mode 100644 packages/runtime-core/src/host-functions.ts diff --git a/crates/client/src/agent_os.rs b/crates/client/src/agent_os.rs index ab7c2c9b5a..5391458f2a 100644 --- a/crates/client/src/agent_os.rs +++ b/crates/client/src/agent_os.rs @@ -25,7 +25,8 @@ use agentos_sidecar_client::wire; use agentos_vm_config as vm_config; use crate::config::{ - AgentOsConfig, AgentOsLimits, HostFunction, HostFunctions, MountConfig, RootFilesystemConfig, + resolve_host_functions, AgentOsConfig, AgentOsLimits, MountConfig, ResolvedHostFunction, + ResolvedHostFunctions, RootFilesystemConfig, RootFilesystemKind, RootFilesystemMode as ConfigRootFilesystemMode, RootLowerInput, SidecarJsBridgeCall, SidecarJsBridgeCallback, TimerScheduleDriver, }; @@ -211,6 +212,8 @@ pub(crate) struct AgentOsInner { // Config / lifecycle. pub(crate) config: Arc, + /// `config.host_functions` resolved to command names once, at create. + pub(crate) host_functions: Vec, pub(crate) sidecar: Arc, pub(crate) sidecar_lease: parking_lot::Mutex>, pub(crate) dynamic_mounts: parking_lot::Mutex>, @@ -487,9 +490,11 @@ impl AgentOs { // 6b. Register host-function collections (if any): forward each definition via `register_host_callbacks`, // record the host execute callbacks in the per-VM registry, and install the shared // host-callback that routes guest host_function calls back to the host by VM. - if !config.host_functions.is_empty() { - let mut host_function_map: HashMap = HashMap::new(); - for collection in &config.host_functions { + let resolved_host_functions = + resolve_host_functions(&config.host_functions).map_err(ClientError::InvalidConfig)?; + if !resolved_host_functions.is_empty() { + let mut host_function_map: HashMap = HashMap::new(); + for collection in &resolved_host_functions { let mut host_functions = HashMap::new(); for host_function in &collection.functions { host_functions.insert( @@ -515,7 +520,7 @@ impl AgentOs { wire::RequestPayload::RegisterHostCallbacksRequest( wire::RegisterHostCallbacksRequest { name: collection.name.clone(), - description: collection.description.clone(), + description: String::new(), command_aliases: vec![format!("agentos-{}", collection.name)], registry_command_aliases: vec![String::from("agentos")], callbacks: host_functions, @@ -580,7 +585,7 @@ impl AgentOs { let _ = vm_host_functions().insert( vm_id.clone(), Arc::new(VmHostFunctionRegistry { - host_functions: config.host_functions.clone(), + host_functions: resolved_host_functions.clone(), host_function_map, }), ); @@ -626,6 +631,7 @@ impl AgentOs { durable_agent_exit_tx: broadcast::channel(64).0, cron, config, + host_functions: resolved_host_functions, sidecar, sidecar_lease: parking_lot::Mutex::new(Some(lease)), dynamic_mounts: parking_lot::Mutex::new(configured_mounts), @@ -874,6 +880,10 @@ impl AgentOs { &self.inner.config } + pub(crate) fn host_functions(&self) -> &[ResolvedHostFunctions] { + &self.inner.host_functions + } + pub(crate) fn cron(&self) -> &Arc { &self.inner.cron } @@ -1371,8 +1381,8 @@ static VM_HOST_FUNCTIONS: OnceCell, - host_function_map: HashMap, + host_functions: Vec, + host_function_map: HashMap, } fn vm_host_functions() -> &'static SccHashMap> { @@ -2493,7 +2503,7 @@ async fn handle_agentos_host_function_command( ownership: &wire::OwnershipScope, registry: &VmHostFunctionRegistry, command: &HostCommandCallbackInput, - collection: &HostFunctions, + collection: &ResolvedHostFunctions, ) -> Result { let Some(host_function_name) = command.args.first() else { return describe_host_functions_payload(®istry.host_functions, &collection.name); @@ -2518,7 +2528,7 @@ async fn handle_agentos_host_function_command( async fn invoke_host_function( ownership: &wire::OwnershipScope, registry: &VmHostFunctionRegistry, - collection: &HostFunctions, + collection: &ResolvedHostFunctions, host_function_name: &str, args: &[String], cwd: &str, @@ -2552,7 +2562,7 @@ async fn invoke_host_function( async fn parse_host_function_input( ownership: &wire::OwnershipScope, - host_function: &HostFunction, + host_function: &ResolvedHostFunction, args: &[String], cwd: &str, ) -> Result { @@ -3128,7 +3138,7 @@ fn compact_json(value: &Value) -> String { serde_json::to_string(value).unwrap_or_else(|_| String::from("")) } -fn list_host_functions_payload(host_functions: &[HostFunctions]) -> Value { +fn list_host_functions_payload(host_functions: &[ResolvedHostFunctions]) -> Value { Value::Object(Map::from_iter([( String::from("hostFunctions"), Value::Array( @@ -3137,7 +3147,6 @@ fn list_host_functions_payload(host_functions: &[HostFunctions]) -> Value { .map(|collection| { json_object([ ("name", Value::String(collection.name.clone())), - ("description", Value::String(collection.description.clone())), ( "functions", Value::Array( @@ -3156,7 +3165,7 @@ fn list_host_functions_payload(host_functions: &[HostFunctions]) -> Value { } fn describe_host_functions_payload( - host_functions: &[HostFunctions], + host_functions: &[ResolvedHostFunctions], collection_name: &str, ) -> Result { let Some(collection) = host_functions @@ -3170,7 +3179,6 @@ fn describe_host_functions_payload( }; Ok(json_object([ ("name", Value::String(collection.name.clone())), - ("description", Value::String(collection.description.clone())), ( "functions", Value::Object(Map::from_iter(collection.functions.iter().map( @@ -3197,7 +3205,7 @@ fn describe_host_functions_payload( } fn describe_host_function_payload( - collection: &HostFunctions, + collection: &ResolvedHostFunctions, host_function_name: &str, ) -> Result { let Some(host_function) = collection @@ -3282,7 +3290,7 @@ fn describe_host_function_flag_type(schema: &Value) -> String { } } -fn host_functions_names(host_functions: &[HostFunctions]) -> String { +fn host_functions_names(host_functions: &[ResolvedHostFunctions]) -> String { host_functions .iter() .map(|collection| collection.name.clone()) @@ -3290,7 +3298,7 @@ fn host_functions_names(host_functions: &[HostFunctions]) -> String { .join(", ") } -fn host_function_names(collection: &HostFunctions) -> String { +fn host_function_names(collection: &ResolvedHostFunctions) -> String { collection .functions .iter() diff --git a/crates/client/src/config.rs b/crates/client/src/config.rs index fdd849b98f..3600fc9dd9 100644 --- a/crates/client/src/config.rs +++ b/crates/client/src/config.rs @@ -8,6 +8,7 @@ //! only and become `Arc` trait objects; they cannot cross the wire and are gated exactly as //! the actor layer gates them. +use std::collections::BTreeMap; use std::sync::Arc; use serde::{Deserialize, Serialize}; @@ -48,8 +49,8 @@ pub struct AgentOsConfig { pub additional_instructions: Option, /// Schedule driver used by the cron manager. Default: [`TimerScheduleDriver`]. pub schedule_driver: Option>, - /// HostFunction collections to register. - pub host_functions: Vec, + /// Host function collections to register, keyed by collection name. + pub host_functions: HostFunctionCollections, /// Rust-only sidecar callback handler for `js_bridge`-style plugin requests. pub sidecar_js_bridge_callback: Option, /// Permission policy. Default: allow-all. @@ -126,7 +127,7 @@ impl AgentOsConfigBuilder { self } - pub fn host_functions(mut self, host_functions: Vec) -> Self { + pub fn host_functions(mut self, host_functions: HostFunctionCollections) -> Self { self.config.host_functions = host_functions; self } @@ -232,11 +233,10 @@ pub type SidecarJsBridgeCallback = Arc< + Sync, >; -/// A single host function within a [`HostFunctions`]. +/// A single host function. The key it is registered under names it, and the +/// input schema's `description` documents it for the agent. #[derive(Clone)] pub struct HostFunction { - pub name: String, - pub description: String, /// JSON Schema for the host function input (forwarded to the sidecar `register_host_callbacks` definition). pub input_schema: serde_json::Value, pub timeout_ms: Option, @@ -244,14 +244,135 @@ pub struct HostFunction { pub execute: HostFunctionCallback, } -/// A registered host function collection (in-process; implementations stay host-side). Functions are exposed to the -/// guest as `:` and dispatched back to [`HostFunction::execute`] via the sidecar -/// host-callback channel. +/// One collection of host functions, keyed by function name. Each key becomes a +/// subcommand of the collection's CLI binary and a method on its guest global. +pub type HostFunctionCollection = BTreeMap; + +/// Host function collections, keyed by collection name (in-process; +/// implementations stay host-side). Each key becomes the CLI binary +/// `agentos-{name}` and a frozen guest global; functions are exposed to the +/// guest as `:` and dispatched back to +/// [`HostFunction::execute`] via the sidecar host-callback channel. +pub type HostFunctionCollections = BTreeMap; + +/// A host function resolved to the names the VM uses. #[derive(Clone)] -pub struct HostFunctions { +pub struct ResolvedHostFunction { + /// Kebab-case command name. Becomes the CLI subcommand. pub name: String, + /// Taken from the input schema's `description`. pub description: String, - pub functions: Vec, + pub input_schema: serde_json::Value, + pub timeout_ms: Option, + pub execute: HostFunctionCallback, +} + +/// A collection resolved to the names the VM uses. Keys arrive as identifiers +/// and are converted once here, so the rest of the client and the sidecar only +/// ever see kebab-case command names. +#[derive(Clone)] +pub struct ResolvedHostFunctions { + /// Kebab-case collection name. Becomes the CLI suffix: `agentos-{name}`. + pub name: String, + pub functions: Vec, +} + +/// Convert a registration key to its command name. `listOrders` and +/// `list-orders` both become `list-orders`, so the guest sees one spelling +/// whichever the caller wrote. Mirrors `hostFunctionCommandName` in the +/// TypeScript client. +pub fn host_function_command_name(key: &str) -> String { + let mut out = String::with_capacity(key.len() + 4); + let chars: Vec = key.chars().collect(); + for (index, character) in chars.iter().enumerate() { + if character.is_ascii_uppercase() && index > 0 { + let previous = chars[index - 1]; + let next_is_lower = chars + .get(index + 1) + .is_some_and(|next| next.is_ascii_lowercase()); + if previous.is_ascii_lowercase() + || previous.is_ascii_digit() + || (previous.is_ascii_uppercase() && next_is_lower) + { + out.push('-'); + } + } + out.extend(character.to_lowercase()); + } + out +} + +fn is_command_name(name: &str) -> bool { + !name.is_empty() + && !name.starts_with('-') + && !name.ends_with('-') + && !name.contains("--") + && name + .chars() + .all(|character| character.is_ascii_lowercase() || character.is_ascii_digit() || character == '-') +} + +fn to_command_name(kind: &str, key: &str) -> Result { + let name = host_function_command_name(key); + if is_command_name(&name) { + Ok(name) + } else { + Err(format!( + "{kind} name \"{key}\" must be alphanumeric, written in camelCase or with single hyphen separators" + )) + } +} + +/// The description the agent sees, taken from the input schema's `description`. +fn schema_description(input_schema: &serde_json::Value) -> String { + input_schema + .get("description") + .and_then(serde_json::Value::as_str) + .unwrap_or_default() + .to_string() +} + +/// Resolve the caller's collections into the shape the client and sidecar use. +/// Fails on a key that cannot become a command name, and on two keys that +/// collide once converted. +pub fn resolve_host_functions( + collections: &HostFunctionCollections, +) -> Result, String> { + let mut resolved = Vec::with_capacity(collections.len()); + let mut seen_collections: BTreeMap = BTreeMap::new(); + + for (collection_key, collection) in collections { + let name = to_command_name("Host function collection", collection_key)?; + if let Some(collided) = seen_collections.get(&name) { + return Err(format!( + "Host function collections \"{collided}\" and \"{collection_key}\" both resolve to the command name \"{name}\"" + )); + } + seen_collections.insert(name.clone(), collection_key.clone()); + + let mut functions = Vec::with_capacity(collection.len()); + let mut seen_functions: BTreeMap = BTreeMap::new(); + for (function_key, definition) in collection { + let function_name = to_command_name("Host function", function_key)?; + if let Some(collided) = seen_functions.get(&function_name) { + return Err(format!( + "Host functions \"{collided}\" and \"{function_key}\" in collection \"{collection_key}\" both resolve to the command name \"{function_name}\"" + )); + } + seen_functions.insert(function_name.clone(), function_key.clone()); + functions.push(ResolvedHostFunction { + name: function_name, + description: schema_description(&definition.input_schema), + input_schema: definition.input_schema.clone(), + timeout_ms: definition.timeout_ms, + execute: definition.execute.clone(), + }); + } + + resolved.push(ResolvedHostFunctions { name, functions }); + } + + Ok(resolved) } // --------------------------------------------------------------------------- diff --git a/crates/client/src/error.rs b/crates/client/src/error.rs index 2e9eeb4ac5..ddede7245c 100644 --- a/crates/client/src/error.rs +++ b/crates/client/src/error.rs @@ -89,6 +89,11 @@ pub enum ClientError { #[error("ACP operation [{code}]: {message}")] AcpOperation { code: String, message: String }, + /// A caller-supplied config value was rejected before the VM was created, + /// for example a host function key that cannot become a command name. + #[error("invalid config: {0}")] + InvalidConfig(String), + /// A cron schedule string could not be parsed/validated. #[error("invalid schedule: {0}")] InvalidSchedule(String), @@ -176,6 +181,7 @@ impl ClientError { ClientError::PathNotAbsolute(_) | ClientError::PathNotNormalized(_) | ClientError::PathReadOnly(_) + | ClientError::InvalidConfig(_) | ClientError::ProcessNotFound(_) | ClientError::ShellNotFound(_) | ClientError::SessionNotFound(_) diff --git a/crates/client/src/lib.rs b/crates/client/src/lib.rs index f64a7d9a12..1c48c75b49 100644 --- a/crates/client/src/lib.rs +++ b/crates/client/src/lib.rs @@ -67,7 +67,7 @@ pub use stream::{ByteStream, Subscription}; pub use config::{ node_modules_mount, AcpLimits, AgentOsConfig, AgentOsConfigBuilder, AgentOsLimits, AgentOsSidecarConfig, FsPermissionRule, FsPermissions, HostFunction, HostFunctionCallback, - HostFunctionLimits, HostFunctions, HttpLimits, JsRuntimeLimits, MountConfig, MountPlugin, + HostFunctionCollection, HostFunctionCollections, HostFunctionLimits, HttpLimits, JsRuntimeLimits, MountConfig, MountPlugin, OverlayMountConfig, PackageRef, PatternPermissionRule, PatternPermissions, PermissionMode, Permissions, PluginLimits, PythonLimits, ResourceLimits, RootFilesystemConfig, RootFilesystemKind, RootFilesystemMode, RootLowerInput, RulePermissions, ScheduleCallback, diff --git a/crates/client/src/session.rs b/crates/client/src/session.rs index ce5764c643..f854015d19 100644 --- a/crates/client/src/session.rs +++ b/crates/client/src/session.rs @@ -37,7 +37,7 @@ use agentos_protocol::ACP_EXTENSION_NAMESPACE; use agentos_sidecar_client::wire; use crate::agent_os::AgentOs; -use crate::config::HostFunctions; +use crate::config::ResolvedHostFunctions; use crate::error::ClientError; use crate::stream::Subscription; pub type DurableSessionEventStream = Pin< @@ -578,7 +578,7 @@ fn combine_instructions(additional: Option<&str>, host_function_reference: &str) } } -fn build_host_function_reference(host_functions: &[HostFunctions]) -> String { +fn build_host_function_reference(host_functions: &[ResolvedHostFunctions]) -> String { if host_functions.is_empty() { return String::new(); } @@ -593,8 +593,6 @@ fn build_host_function_reference(host_functions: &[HostFunctions]) -> String { for collection in host_functions { lines.push(format!("### {}", collection.name)); lines.push(String::new()); - lines.push(collection.description.clone()); - lines.push(String::new()); for host_function in &collection.functions { let signature = build_host_function_flag_signature(&host_function.input_schema); let suffix = if signature.is_empty() { @@ -602,10 +600,17 @@ fn build_host_function_reference(host_functions: &[HostFunctions]) -> String { } else { format!(" {signature}") }; - lines.push(format!( - "- `agentos-{} {}{}` - {}", - collection.name, host_function.name, suffix, host_function.description - )); + lines.push(if host_function.description.is_empty() { + format!( + "- `agentos-{} {}{}`", + collection.name, host_function.name, suffix + ) + } else { + format!( + "- `agentos-{} {}{}` - {}", + collection.name, host_function.name, suffix, host_function.description + ) + }); } lines.push(String::new()); lines.push(format!( @@ -803,7 +808,7 @@ impl AgentOs { .filter(|value| !value.is_empty()) .collect::>() .join("\n\n"); - let host_function_reference = build_host_function_reference(&self.config().host_functions); + let host_function_reference = build_host_function_reference(self.host_functions()); let additional_instructions = combine_instructions( (!caller_instructions.is_empty()).then_some(caller_instructions.as_str()), &host_function_reference, diff --git a/crates/client/tests/os_instructions_e2e.rs b/crates/client/tests/os_instructions_e2e.rs index 17c9cb3b1c..8b78ac7145 100644 --- a/crates/client/tests/os_instructions_e2e.rs +++ b/crates/client/tests/os_instructions_e2e.rs @@ -13,8 +13,8 @@ use std::path::Path; use std::sync::Arc; use agentos_client::config::{ - node_modules_mount, AgentOsConfig, AgentOsSidecarConfig, FsPermissions, HostFunction, - HostFunctions, PackageRef, PatternPermissions, PermissionMode, Permissions, + node_modules_mount, AgentOsConfig, AgentOsSidecarConfig, FsPermissions, HostFunction, HostFunctionCollection, + HostFunctionCollections, PackageRef, PatternPermissions, PermissionMode, Permissions, }; use agentos_client::{AgentOs, OpenSessionInput}; use agentos_vm_config::VmSqliteDescriptor; @@ -103,12 +103,12 @@ fn write_mock_pi_adapter(module_root: &std::path::Path) -> std::path::PathBuf { } async fn launch_pi_session_and_read_prompt(options: OpenSessionInput) -> String { - launch_pi_session_with_tools_and_read_prompt(options, Vec::new()).await + launch_pi_session_with_tools_and_read_prompt(options, HostFunctionCollections::new()).await } async fn launch_pi_session_with_tools_and_read_prompt( options: OpenSessionInput, - host_functions: Vec, + host_functions: HostFunctionCollections, ) -> String { let module_access_dir = std::env::temp_dir().join(format!("agentos-client-os-instructions-{}", Uuid::new_v4())); @@ -124,7 +124,7 @@ async fn run_session( module_access_dir: &Path, package_dir: &Path, options: OpenSessionInput, - host_functions: Vec, + host_functions: HostFunctionCollections, ) -> String { let os = AgentOs::create(AgentOsConfig { database: Some(VmSqliteDescriptor::SqliteFile { @@ -220,23 +220,24 @@ async fn create_session_injects_host_function_reference_from_client_config() { skip_os_instructions: None, additional_instructions: None, }, - vec![HostFunctions { - name: "weather".to_string(), - description: "Weather lookup tools.".to_string(), - functions: vec![HostFunction { - name: "forecast".to_string(), - description: "Get a forecast.".to_string(), - input_schema: json!({ - "type": "object", - "properties": { - "zipCode": { "type": "string" }, - }, - "required": ["zipCode"], - }), - timeout_ms: None, - execute: Arc::new(|_input| Box::pin(async { Ok(json!({ "ok": true })) })), - }], - }], + HostFunctionCollections::from([( + "weather".to_string(), + HostFunctionCollection::from([( + "forecast".to_string(), + HostFunction { + input_schema: json!({ + "type": "object", + "description": "Get a forecast.", + "properties": { + "zipCode": { "type": "string" }, + }, + "required": ["zipCode"], + }), + timeout_ms: None, + execute: Arc::new(|_input| Box::pin(async { Ok(json!({ "ok": true })) })), + }, + )]), + )]), ) .await; diff --git a/crates/native-sidecar-core/src/host_functions.rs b/crates/native-sidecar-core/src/host_functions.rs index c3b261c9f9..bdfd44f340 100644 --- a/crates/native-sidecar-core/src/host_functions.rs +++ b/crates/native-sidecar-core/src/host_functions.rs @@ -37,12 +37,8 @@ pub fn validate_host_functions_registration( payload: &RegisterHostCallbacksRequest, ) -> Result<(), HostFunctionRegistrationError> { validate_collection_name(&payload.name)?; - if payload.description.is_empty() { - return Err(HostFunctionRegistrationError::InvalidState(format!( - "collection {} is missing a description", - payload.name - ))); - } + // Descriptions are optional: a collection has none of its own, and a + // function's comes from its input schema, which need not carry one. validate_description_length( &format!("Host function collection \"{}\"", payload.name), &payload.description, @@ -71,12 +67,6 @@ pub fn validate_host_functions_registration( } for (host_function_name, host_function) in &payload.callbacks { validate_host_function_name(host_function_name)?; - if host_function.description.is_empty() { - return Err(HostFunctionRegistrationError::InvalidState(format!( - "host function {} in collection {} is missing a description", - host_function_name, payload.name - ))); - } validate_description_length( &format!("Host function \"{}/{}\"", payload.name, host_function_name), &host_function.description, diff --git a/docs/content/docs/host-functions.mdx b/docs/content/docs/host-functions.mdx index a281d5203a..697e979405 100644 --- a/docs/content/docs/host-functions.mdx +++ b/docs/content/docs/host-functions.mdx @@ -8,7 +8,7 @@ Expose your host JavaScript functions (defined with Zod input schemas) to agents ## Getting started -Define a host-function collection with Zod input schemas and pass it to `agentOS({ hostFunctions })`. Each function becomes a CLI subcommand inside the VM. +Pass `agentOS({ hostFunctions })` a record of collections. The keys name everything: the collection key becomes the CLI binary `agentos-{name}` and a guest global, and each function key becomes a subcommand. A function needs only an `inputSchema` and an `execute` handler. `.describe()` on the schema is what the agent reads. @@ -30,11 +30,11 @@ Zod schema fields are converted to CLI flags automatically. Field names are conv | `z.enum(["a","b"])` | `--name a` | `--format json` | | `z.array(z.string())` | `--name a --name b` | `--tags foo --tags bar` | -Optional fields (via `.optional()`) become optional flags. Required fields are enforced at validation time. Use `.describe()` on Zod fields to generate useful `--help` output. +Optional fields (via `.optional()`) become optional flags. Required fields are enforced at validation time. Use `.describe()` on Zod fields to generate useful `--help` output, and `.describe()` on the whole schema to describe the function itself. ### What the agent sees -When host functions are registered, CLI shims are installed at `/bin/agentos-{name}` inside the VM and the function list is injected into the agent's [system prompt](/agentos/docs/system-prompt), so keep descriptions concise to save tokens. +When host functions are registered, CLI shims are installed at `/bin/agentos-{name}` inside the VM and the function list is injected into the agent's [system prompt](/agentos/docs/system-prompt), so keep each schema's `.describe()` concise to save tokens. The agent interacts with host functions as shell commands: @@ -74,8 +74,9 @@ Missing required flag: --city Inline JavaScript and TypeScript get each collection as a frozen global, and each host function as an async function, so generated code calls your tools like any -other API. Names become camelCase identifiers: `agentos-order-store list-orders` -is `orderStore.listOrders(input)`. +other API. Keys written in camelCase become kebab-case commands and camelCase guest +identifiers, so `orderStore.listOrders` is `agentos-order-store list-orders` on +the CLI and `orderStore.listOrders(input)` in guest JavaScript. diff --git a/examples/agent-to-agent/server.ts b/examples/agent-to-agent/server.ts index 5809c3daf5..64449fdda5 100644 --- a/examples/agent-to-agent/server.ts +++ b/examples/agent-to-agent/server.ts @@ -43,24 +43,22 @@ async function reviewCode(code: string): Promise { // The writer agent gets a `review` host-function collection. When the writer runs // `agentos-review submit`, the bridge above executes on the host. const writer = agentOS({ - hostFunctions: [ - { - name: "review", - description: "Send code to the reviewer agent and get back a review.", - functions: { - submit: { - description: - "Submit the full contents of a file to the reviewer agent for review. Returns the reviewer's feedback as text.", - inputSchema: z.object({ + hostFunctions: { + review: { + submit: { + inputSchema: z + .object({ code: z.string().describe("The full source code to review."), - }), - execute: async (input: { code: string }) => ({ - review: await reviewCode(input.code), - }), - }, + }) + .describe( + "Submit the full contents of a file to the reviewer agent for review. Returns the reviewer's feedback as text.", + ), + execute: async (input: { code: string }) => ({ + review: await reviewCode(input.code), + }), }, }, - ], + }, }); export const registry = setup({ use: { writer, reviewer } }); diff --git a/examples/crash-course/agent-to-agent-server.ts b/examples/crash-course/agent-to-agent-server.ts index 68e24b6e2c..0b304ba656 100644 --- a/examples/crash-course/agent-to-agent-server.ts +++ b/examples/crash-course/agent-to-agent-server.ts @@ -1,7 +1,7 @@ +import pi from "@agentos-software/pi"; import { agentOS, setup } from "@rivet-dev/agentos"; import { createClient } from "@rivet-dev/agentos/client"; import { z } from "zod"; -import pi from "@agentos-software/pi"; // The reviewer is its own isolated agent VM. const reviewer = agentOS({ software: [pi] }); @@ -10,44 +10,41 @@ const reviewer = agentOS({ software: [pi] }); // coder's VM into the reviewer's VM and asks the reviewer to review it. const coder = agentOS({ software: [pi], - hostFunctions: [ - { - name: "review", - description: "Send a file to the reviewer agent and get back a review.", - functions: { - submit: { - description: "Submit a file path for review by the reviewer agent.", - inputSchema: z.object({ path: z.string() }), - execute: async ({ path }: { path: string }) => { - const client = createClient({ - endpoint: "http://localhost:6420", - }); - const content = await client.coder - .getOrCreate("feature-auth") - .filesystem.readFile(path); - const reviewerHandle = client.reviewer.getOrCreate("feature-auth"); - await reviewerHandle.filesystem.writeFile(path, content); - await reviewerHandle.sessions.open({ - agent: "pi", - env: { ANTHROPIC_API_KEY: process.env.ANTHROPIC_API_KEY! }, - }); - const result = await reviewerHandle.sessions.prompt({ - content: [ - { type: "text", text: `Review ${path} for security issues` }, - ], - }); - return { - review: - result.message?.content - .filter((block) => block.type === "text") - .map((block) => block.text) - .join("") ?? "", - }; - }, + hostFunctions: { + review: { + submit: { + inputSchema: z + .object({ path: z.string() }) + .describe("Submit a file path for review by the reviewer agent."), + execute: async ({ path }: { path: string }) => { + const client = createClient({ + endpoint: "http://localhost:6420", + }); + const content = await client.coder + .getOrCreate("feature-auth") + .filesystem.readFile(path); + const reviewerHandle = client.reviewer.getOrCreate("feature-auth"); + await reviewerHandle.filesystem.writeFile(path, content); + await reviewerHandle.sessions.open({ + agent: "pi", + env: { ANTHROPIC_API_KEY: process.env.ANTHROPIC_API_KEY! }, + }); + const result = await reviewerHandle.sessions.prompt({ + content: [ + { type: "text", text: `Review ${path} for security issues` }, + ], + }); + return { + review: + result.message?.content + .filter((block) => block.type === "text") + .map((block) => block.text) + .join("") ?? "", + }; }, }, }, - ], + }, }); export const registry = setup({ use: { coder, reviewer } }); diff --git a/examples/crash-course/sandbox-stubs.d.ts b/examples/crash-course/sandbox-stubs.d.ts index d62bf14767..1ae6d863b5 100644 --- a/examples/crash-course/sandbox-stubs.d.ts +++ b/examples/crash-course/sandbox-stubs.d.ts @@ -25,8 +25,8 @@ declare module "sandbox-agent/docker" { declare module "@rivet-dev/agentos-sandbox" { import type { + HostFunctionCollection, NativeMountPluginDescriptor, - HostFunctions, } from "@rivet-dev/agentos"; import type { SandboxAgent } from "sandbox-agent"; @@ -54,5 +54,5 @@ declare module "@rivet-dev/agentos-sandbox" { /** Build a host-function collection that exposes the sandbox's process management. */ export function createSandboxHostFunctions( options: SandboxHostFunctionsOptions, - ): HostFunctions; + ): HostFunctionCollection; } diff --git a/examples/crash-course/sandbox.ts b/examples/crash-course/sandbox.ts index 1d6165f1ff..5bda496838 100644 --- a/examples/crash-course/sandbox.ts +++ b/examples/crash-course/sandbox.ts @@ -10,7 +10,7 @@ const sandbox = await SandboxAgent.start({ sandbox: docker() }); const vm = agentOS({ // Host functions let the agent control the sandbox. - hostFunctions: [createSandboxHostFunctions({ client: sandbox })], + hostFunctions: { sandbox: createSandboxHostFunctions({ client: sandbox }) }, // Mounts let the agent read the sandbox filesystem (optional) mounts: [ { diff --git a/examples/embedded/agent-to-agent.ts b/examples/embedded/agent-to-agent.ts index 8ff0988438..33ea744f6b 100644 --- a/examples/embedded/agent-to-agent.ts +++ b/examples/embedded/agent-to-agent.ts @@ -1,5 +1,5 @@ import pi from "@agentos-software/pi"; -import { AgentOs, type HostFunctions } from "@rivet-dev/agentos-core"; +import { AgentOs } from "@rivet-dev/agentos-core"; import { z } from "zod"; const apiKey = process.env.ANTHROPIC_API_KEY; @@ -13,29 +13,27 @@ await reviewer.sessions.open({ env: { ANTHROPIC_API_KEY: apiKey }, }); -const review: HostFunctions = { - name: "review", - description: "Ask the reviewer agent for feedback", - functions: { - draft: { - description: "Review a draft", - inputSchema: z.object({ draft: z.string() }), - execute: async ({ draft }: { draft: string }) => { - const response = await reviewer.sessions.prompt({ - content: [{ type: "text", text: `Review this draft:\n\n${draft}` }], - }); - const feedback = - response.message?.content - .filter((block) => block.type === "text") - .map((block) => block.text) - .join("") ?? ""; - return { feedback }; - }, +const review = { + draft: { + inputSchema: z.object({ draft: z.string() }).describe("Review a draft"), + execute: async ({ draft }: { draft: string }) => { + const response = await reviewer.sessions.prompt({ + content: [{ type: "text", text: `Review this draft:\n\n${draft}` }], + }); + const feedback = + response.message?.content + .filter((block) => block.type === "text") + .map((block) => block.text) + .join("") ?? ""; + return { feedback }; }, }, }; -const writer = await AgentOs.create({ software: [pi], hostFunctions: [review] }); +const writer = await AgentOs.create({ + software: [pi], + hostFunctions: { review: review }, +}); try { await writer.sessions.open({ diff --git a/examples/embedded/host-functions.ts b/examples/embedded/host-functions.ts index 8817a72c6d..858994d7d2 100644 --- a/examples/embedded/host-functions.ts +++ b/examples/embedded/host-functions.ts @@ -1,24 +1,21 @@ -import { AgentOs, type HostFunctions } from "@rivet-dev/agentos-core"; +import { AgentOs } from "@rivet-dev/agentos-core"; import { z } from "zod"; -// Host-function collections are defined exactly as they are for the actor. Pass them to -// AgentOs.create() and `execute` runs in this host process. -const weather: HostFunctions = { - name: "weather", - description: "Weather data functions", - functions: { - forecast: { - description: "Get the weather forecast for a city", - inputSchema: z.object({ city: z.string().describe("City name") }), - execute: async (input: { city: string }) => ({ - city: input.city, - temperature: 22, - }), +// Host functions are defined exactly as they are for the actor. Pass them to +// AgentOs.create() and `execute` runs in this host process, with its input typed +// by its own schema. +const vm = await AgentOs.create({ + hostFunctions: { + weather: { + forecast: { + inputSchema: z + .object({ city: z.string().describe("City name") }) + .describe("Get the weather forecast for a city"), + execute: async ({ city }) => ({ city, temperature: 22 }), + }, }, }, -}; - -const vm = await AgentOs.create({ hostFunctions: [weather] }); +}); // The agent calls it as `agentos-weather forecast --city Paris`. const result = await vm.process.exec("agentos-weather forecast --city Paris"); diff --git a/examples/host-functions/exec-bash.ts b/examples/host-functions/exec-bash.ts index e120dcfe2e..141500c504 100644 --- a/examples/host-functions/exec-bash.ts +++ b/examples/host-functions/exec-bash.ts @@ -1,27 +1,24 @@ -import { AgentOs, type HostFunctions } from "@rivet-dev/agentos-core"; +import { AgentOs } from "@rivet-dev/agentos-core"; import { z } from "zod"; // The handler runs on the host, so the API key never enters the VM. -const weather: HostFunctions = { - name: "weather", - description: "Weather data functions", - functions: { - forecast: { - description: "Get the weather forecast for a city", - inputSchema: z.object({ city: z.string() }), - execute: async ({ city }: { city: string }) => { - const res = await fetch( - `https://api.weather.example/forecast?city=${city}&key=${process.env.WEATHER_API_KEY}`, - ); - return res.json(); - }, +const weather = { + forecast: { + inputSchema: z + .object({ city: z.string() }) + .describe("Get the weather forecast for a city"), + execute: async ({ city }: { city: string }) => { + const res = await fetch( + `https://api.weather.example/forecast?city=${city}&key=${process.env.WEATHER_API_KEY}`, + ); + return res.json(); }, }, }; // The collection is projected into the VM as an `agentos-weather` command, so // it composes with pipes and redirects like any other program. -const runtime = await AgentOs.create({ hostFunctions: [weather] }); +const runtime = await AgentOs.create({ hostFunctions: { weather: weather } }); try { const result = await runtime.process.exec( diff --git a/examples/host-functions/exec-javascript.ts b/examples/host-functions/exec-javascript.ts index bae12bca2a..dd547877d2 100644 --- a/examples/host-functions/exec-javascript.ts +++ b/examples/host-functions/exec-javascript.ts @@ -1,26 +1,23 @@ -import { AgentOs, type HostFunctions } from "@rivet-dev/agentos-core"; +import { AgentOs } from "@rivet-dev/agentos-core"; import { z } from "zod"; // The handler runs on the host, so the API key never enters the VM. -const weather: HostFunctions = { - name: "weather", - description: "Weather data functions", - functions: { - forecast: { - description: "Get the weather forecast for a city", - inputSchema: z.object({ city: z.string() }), - execute: async ({ city }: { city: string }) => { - const res = await fetch( - `https://api.weather.example/forecast?city=${city}&key=${process.env.WEATHER_API_KEY}`, - ); - return res.json(); - }, +const weather = { + forecast: { + inputSchema: z + .object({ city: z.string() }) + .describe("Get the weather forecast for a city"), + execute: async ({ city }: { city: string }) => { + const res = await fetch( + `https://api.weather.example/forecast?city=${city}&key=${process.env.WEATHER_API_KEY}`, + ); + return res.json(); }, }, }; // The collection is projected into the VM as an `agentos-weather` command. -const runtime = await AgentOs.create({ hostFunctions: [weather] }); +const runtime = await AgentOs.create({ hostFunctions: { weather: weather } }); try { const result = await runtime.javascript.execute( diff --git a/examples/host-functions/exec-python.ts b/examples/host-functions/exec-python.ts index 962ae86753..ebaed28fa1 100644 --- a/examples/host-functions/exec-python.ts +++ b/examples/host-functions/exec-python.ts @@ -1,26 +1,23 @@ -import { AgentOs, type HostFunctions } from "@rivet-dev/agentos-core"; +import { AgentOs } from "@rivet-dev/agentos-core"; import { z } from "zod"; // The handler runs on the host, so the API key never enters the VM. -const weather: HostFunctions = { - name: "weather", - description: "Weather data functions", - functions: { - forecast: { - description: "Get the weather forecast for a city", - inputSchema: z.object({ city: z.string() }), - execute: async ({ city }: { city: string }) => { - const res = await fetch( - `https://api.weather.example/forecast?city=${city}&key=${process.env.WEATHER_API_KEY}`, - ); - return res.json(); - }, +const weather = { + forecast: { + inputSchema: z + .object({ city: z.string() }) + .describe("Get the weather forecast for a city"), + execute: async ({ city }: { city: string }) => { + const res = await fetch( + `https://api.weather.example/forecast?city=${city}&key=${process.env.WEATHER_API_KEY}`, + ); + return res.json(); }, }, }; // The collection is projected into the VM as an `agentos-weather` command. -const runtime = await AgentOs.create({ hostFunctions: [weather] }); +const runtime = await AgentOs.create({ hostFunctions: { weather: weather } }); try { const result = await runtime.python.execute( diff --git a/examples/host-functions/server.ts b/examples/host-functions/server.ts index d3b1a7ac1a..068b19766b 100644 --- a/examples/host-functions/server.ts +++ b/examples/host-functions/server.ts @@ -1,37 +1,36 @@ import { agentOS, setup } from "@rivet-dev/agentos"; import { z } from "zod"; -// Define a collection of host functions. Each function has a Zod input -// schema and an `execute` handler that runs on the host. The group is exposed to -// the agent as a CLI command at /bin/agentos-{name} inside the VM. -const weatherFunctions = { - name: "weather", - description: "Weather data functions", - functions: { - forecast: { - description: "Get the weather forecast for a city", - inputSchema: z.object({ - city: z.string().describe("City name"), - days: z.number().optional().describe("Number of days"), - }), - execute: async (input: { city: string; days?: number }) => { - const res = await fetch( - `https://api.weather.example/forecast?city=${input.city}&days=${input.days ?? 3}`, - ); - return res.json(); - }, - examples: [ - { - description: "3-day forecast for Paris", - input: { city: "Paris", days: 3 }, +// Host functions are a record of collections. The keys name everything: the +// collection key becomes the CLI command /bin/agentos-{name} inside the VM, and +// each function key becomes one of its subcommands. A function needs only a Zod +// input schema and an `execute` handler that runs on the host; `.describe()` on +// the schema is what the agent reads. +const vm = agentOS({ + hostFunctions: { + weather: { + forecast: { + inputSchema: z + .object({ + city: z.string().describe("City name"), + days: z.number().optional().describe("Number of days"), + }) + .describe("Get the weather forecast for a city"), + execute: async ({ city, days }) => { + const res = await fetch( + `https://api.weather.example/forecast?city=${city}&days=${days ?? 3}`, + ); + return res.json(); }, - ], + examples: [ + { + description: "3-day forecast for Paris", + input: { city: "Paris", days: 3 }, + }, + ], + }, }, }, -}; - -const vm = agentOS({ - hostFunctions: [weatherFunctions], }); export const registry = setup({ use: { vm } }); diff --git a/examples/js-code-mode/src/index.ts b/examples/js-code-mode/src/index.ts index 0294e34c24..6a3c0e795f 100644 --- a/examples/js-code-mode/src/index.ts +++ b/examples/js-code-mode/src/index.ts @@ -1,27 +1,26 @@ // docs:start hostFunctions -import { AgentOs, hostFunction, hostFunctions } from "@rivet-dev/agentos-core"; +import { AgentOs } from "@rivet-dev/agentos-core"; import { z } from "zod"; -const toolFunctions = hostFunctions({ - name: "tools", - description: "Curated host capabilities for generated code.", - functions: { - weather: hostFunction({ - description: "Look up a city's temperature.", - inputSchema: z.object({ city: z.string() }), - execute: ({ city }) => ({ - city, - tempF: city === "San Francisco" ? 61 : 75, - }), - }), +const runtime = await AgentOs.create({ + hostFunctions: { + tools: { + weather: { + inputSchema: z + .object({ city: z.string() }) + .describe("Look up a city's temperature."), + execute: ({ city }) => ({ + city, + tempF: city === "San Francisco" ? 61 : 75, + }), + }, + }, }, }); - -const runtime = await AgentOs.create({ hostFunctions: [toolFunctions] }); // docs:end hostFunctions // docs:start generated-code -// Each binding collection is a global inside the VM, and each binding is an +// Each collection is a global inside the VM, and each host function is an // async function, so generated code calls your tools like any other API. const llmGeneratedExpression = `(async () => { const [sf, tokyo] = await Promise.all([ diff --git a/examples/quickstart/host-functions/index.ts b/examples/quickstart/host-functions/index.ts index 815003d16d..52f1424890 100644 --- a/examples/quickstart/host-functions/index.ts +++ b/examples/quickstart/host-functions/index.ts @@ -1,42 +1,35 @@ -import { AgentOs, hostFunction, hostFunctions } from "@rivet-dev/agentos-core"; +import { AgentOs } from "@rivet-dev/agentos-core"; import { z } from "zod"; -const weatherFunctions = hostFunctions({ - name: "weather", - description: "Look up weather information for cities.", - functions: { - get: hostFunction({ - description: "Get the current weather for a city.", - inputSchema: z.object({ - city: z.string().describe("City name (e.g. 'London')."), - }), - execute: async ({ city }) => ({ - city, - temperature: 18, - conditions: "partly cloudy", - humidity: 65, - }), - examples: [ - { description: "Get London weather", input: { city: "London" } }, - ], - }), - }, -}); - -const calculatorFunctions = hostFunctions({ - name: "calc", - description: "Simple calculator operations.", - functions: { - add: hostFunction({ - description: "Add two numbers.", - inputSchema: z.object({ a: z.number(), b: z.number() }), - execute: ({ a, b }) => ({ result: a + b }), - }), - }, -}); - const vm = await AgentOs.create({ - hostFunctions: [weatherFunctions, calculatorFunctions], + hostFunctions: { + weather: { + get: { + inputSchema: z + .object({ + city: z.string().describe("City name (e.g. 'London')."), + }) + .describe("Get the current weather for a city."), + execute: async ({ city }) => ({ + city, + temperature: 18, + conditions: "partly cloudy", + humidity: 65, + }), + examples: [ + { description: "Get London weather", input: { city: "London" } }, + ], + }, + }, + calc: { + add: { + inputSchema: z + .object({ a: z.number(), b: z.number() }) + .describe("Add two numbers."), + execute: ({ a, b }) => ({ result: a + b }), + }, + }, + }, permissions: { fs: "allow", network: "allow", diff --git a/packages/agentos-sandbox/tests/vm-integration.test.ts b/packages/agentos-sandbox/tests/vm-integration.test.ts index 233eeb4c1c..c3fdf8dee7 100644 --- a/packages/agentos-sandbox/tests/vm-integration.test.ts +++ b/packages/agentos-sandbox/tests/vm-integration.test.ts @@ -109,7 +109,7 @@ describe("VM integration", () => { it("should execute the run-command hostFunction directly via the hostFunction collection", async () => { const tk = createSandboxHostFunctions({ client: sandbox.client }); - const result = await tk.hostFunctions["run-command"].execute({ + const result = await tk.runCommand.execute({ command: "echo", args: ["hello", "from", "sandbox"], }); @@ -122,26 +122,28 @@ describe("VM integration", () => { const tk = createSandboxHostFunctions({ client: sandbox.client }); // Confirm the sandbox hostFunction collection runs commands successfully. - const result = await tk.hostFunctions["run-command"].execute({ + const result = await tk.runCommand.execute({ command: "echo", args: ["hello from sandbox hostFunction collection"], }); expect(result.exitCode).toBe(0); - expect(result.stdout).toContain("hello from sandbox hostFunction collection"); + expect(result.stdout).toContain( + "hello from sandbox hostFunction collection", + ); // Create a process and list it. - const proc = await tk.hostFunctions["create-process"].execute({ + const proc = await tk.createProcess.execute({ command: "sleep", args: ["60"], }); expect(proc.status).toBe("running"); - const listed = await tk.hostFunctions["list-processes"].execute({}); + const listed = await tk.listProcesses.execute({}); const found = listed.processes.find( (p: { id: string }) => p.id === proc.id, ); expect(found).toBeDefined(); - await tk.hostFunctions["kill-process"].execute({ id: proc.id }); + await tk.killProcess.execute({ id: proc.id }); }); }); diff --git a/packages/agentos/src/actor.ts b/packages/agentos/src/actor.ts index 3b789ca9a7..e1e827e52d 100644 --- a/packages/agentos/src/actor.ts +++ b/packages/agentos/src/actor.ts @@ -8,6 +8,7 @@ import { type CodeExecutionResult, type CronEvent, type DynamicMountDescriptor, + type HostFunctionSchemas, type OpenSessionInput, type PackageDescriptor, type ProcessDescriptor, @@ -2020,7 +2021,9 @@ export type AgentOsActorDefinition = ActorDefinition< AgentOsActions >; -export interface AgentOsActorExtras extends AgentOsOptions { +export interface AgentOsActorExtras< + THostFunctions extends HostFunctionSchemas = HostFunctionSchemas, +> extends AgentOsOptions { /** * Resolve trusted VM options from actor state immediately before the VM's * first boot. Evaluated once per wake. The actor-owned root filesystem and @@ -2075,6 +2078,7 @@ export type AgentOsActorConfigInput< TEvents, TQueues > = Record, + THostFunctions extends HostFunctionSchemas = HostFunctionSchemas, > = DistributiveOmit< ActorConfigInput< TState, @@ -2089,7 +2093,7 @@ export type AgentOsActorConfigInput< >, "db" > & - AgentOsActorExtras & + AgentOsActorExtras & AgentOsEventHooks< ActorContext< TState, @@ -2252,6 +2256,7 @@ export function createAgentOS< TEvents, TQueues > = Record, + THostFunctions extends HostFunctionSchemas = HostFunctionSchemas, >( config: AgentOsActorConfigInput< TState, @@ -2261,7 +2266,8 @@ export function createAgentOS< TInput, TEvents, TQueues, - TUserActions + TUserActions, + THostFunctions > = {} as AgentOsActorConfigInput< TState, TConnParams, @@ -2270,7 +2276,8 @@ export function createAgentOS< TInput, TEvents, TQueues, - TUserActions + TUserActions, + THostFunctions >, ): ActorDefinition< TState, diff --git a/packages/core/CLAUDE.md b/packages/core/CLAUDE.md index fa728d36b8..b542029353 100644 --- a/packages/core/CLAUDE.md +++ b/packages/core/CLAUDE.md @@ -23,7 +23,7 @@ - If a file must be visible to both `vm.readFile()` and guest shell commands, it cannot live only in a local compat mount. Put it on a real sidecar-visible path or mount, and keep any read-only guarantees enforced below the TypeScript proxy layer. - Host-function registration is split across the boundary: TypeScript converts Zod schemas to JSON Schema, generates prompt markdown, validates sidecar host-function invocations, and runs the local `execute()` callbacks, while the sidecar owns CLI flag parsing and `agentos` command dispatch via `registerHostCallbacks` / `RegisterHostCallbacks`. - Host-function `inputSchema` conversion in `src/host-functions-zod.ts` is intentionally fail-closed. Support only the Zod subset that round-trips cleanly into the sidecar-facing JSON Schema contract; if a schema would degrade semantics or emit `$ref`/`$defs` (`discriminatedUnion`, `intersection`, `tuple`, `record`, `date`, `bigint`, custom refinements, metadata `id`, etc.), throw `HostFunctionSchemaConversionError` with the offending field path instead of coercing it to `{ type: "string" }`. -- The host-function description limit is a cross-boundary contract: keep the 200-character maximum aligned between `src/host-functions.ts` and Rust `RegisterHostCallbacks` validation in `crates/native-sidecar/src/host_functions.rs`, with boundary tests on both sides when changing it. +- Host functions carry no name or description fields. The registration key names a collection or function, and the description is the input schema's `.describe()` text (`description` on the JSON Schema the Rust client takes). Both are optional; the sidecar owns the length cap in `crates/native-sidecar-core/src/host_functions.rs` and the client does not duplicate it. - `src/sidecar/rpc-client.ts` is the consolidated home for framed sidecar I/O, compat proxy helpers, and sidecar descriptor serializers. Keep shared/explicit sidecar pool and VM lease bookkeeping in `src/agent-os.ts` rather than reintroducing another sidecar lifecycle layer. - In `src/agent-os.ts`, shell teardown is two-phase: public `_shells` entries can disappear immediately on `closeShell()`, but `dispose()` must still await the separate pending shell-exit set before dropping the sidecar event listener, or late shell stdout/exit delivery can race into a closed bridge. - The native sidecar framed stdio path now defaults to the BARE payload codec. Keep any JSON payload support behind explicit migration-only opts such as `payloadCodec: "json"`, and remember that BARE structs need every positional field serialized explicitly across the Rust/TypeScript boundary rather than relying on JSON-style `skip_serializing_if` omissions. diff --git a/packages/core/src/agent-os.ts b/packages/core/src/agent-os.ts index e6d1ddb46a..f1b9a4bebf 100644 --- a/packages/core/src/agent-os.ts +++ b/packages/core/src/agent-os.ts @@ -18,8 +18,15 @@ import type { CreateVmConfig, VmUserConfig, } from "@rivet-dev/agentos-runtime-core/vm-config"; -import { type HostFunction, type HostFunctions, validateHostFunctions } from "./host-functions.js"; -import { zodToJsonSchema } from "./host-functions-zod.js"; +import { + type HostFunction, + type HostFunctionCollections, + type HostFunctionSchemas, + hostFunctionDescription, + type ResolvedHostFunctions, + resolveHostFunctions, +} from "@rivet-dev/agentos-runtime-core/host-functions"; +import { zodToJsonSchema } from "@rivet-dev/agentos-runtime-core/host-functions-zod"; import type { JsonRpcNotification, JsonRpcRequest, @@ -541,7 +548,7 @@ interface AgentOsVmAdmin extends InProcessSidecarVmAdmin { sidecarSession: AuthenticatedSession; sidecarVm: CreatedVm; snapshotRootFilesystem?: (maxBytes: number) => Promise; - hostFunctions: HostFunctions[]; + hostFunctions: ResolvedHostFunctions[]; hostFunctionReference: string; } @@ -821,7 +828,9 @@ export type LimitWarningHandler = (warning: LimitWarning) => void; * `packages/core/src/options-schema.ts::agentOsOptionsSchema`. The TypeScript * Rivet actor accepts this surface directly alongside ordinary actor options. */ -export interface AgentOsOptions { +export interface AgentOsOptions< + HOST_FUNCTIONS extends HostFunctionSchemas = HostFunctionSchemas, +> { /** Initial virtual Linux credentials and account record. Defaults to `1000:1000` (`agentos`). */ user?: VmUserConfig; /** @@ -863,7 +872,7 @@ export interface AgentOsOptions { /** Custom schedule driver for cron jobs. Defaults to TimerScheduleDriver. */ scheduleDriver?: ScheduleDriver; /** Host functions available to agents inside the VM. */ - hostFunctions?: HostFunctions[]; + hostFunctions?: HostFunctionCollections; /** * Permission policy for the kernel. By default the guest behaves like a * sandboxed machine: its virtual filesystem, processes, environment, listeners, @@ -1703,14 +1712,18 @@ function collectSidecarMountPlan(options: { mounts?: MountConfig[] }): { return { sidecarMounts, hostMounts, hostPathMappings }; } -function collectHostFunctionBootstrapCommands(hostFunctions: HostFunctions[]): string[] { +function collectHostFunctionBootstrapCommands( + hostFunctions: ResolvedHostFunctions[], +): string[] { if (hostFunctions.length === 0) { return []; } return [ "agentos", - ...hostFunctions.map((hostFunctionCollection) => `agentos-${hostFunctionCollection.name}`), + ...hostFunctions.map( + (hostFunctionCollection) => `agentos-${hostFunctionCollection.name}`, + ), ]; } @@ -1740,7 +1753,7 @@ function hostFunctionToSidecarDefinition( definition: HostFunction, ): SidecarRegisteredHostCallbackDefinition { return { - description: definition.description, + description: hostFunctionDescription(definition), inputSchema: zodToJsonSchema(definition.inputSchema), ...(definition.timeout !== undefined ? { timeoutMs: definition.timeout } @@ -1769,7 +1782,9 @@ function combineInstructions( return parts.join("\n\n"); } -function buildHostFunctionReference(hostFunctions: HostFunctions[]): string { +function buildHostFunctionReference( + hostFunctions: ResolvedHostFunctions[], +): string { if (hostFunctions.length === 0) { return ""; } @@ -1784,16 +1799,19 @@ function buildHostFunctionReference(hostFunctions: HostFunctions[]): string { for (const hostFunctionCollection of hostFunctions) { lines.push(`### ${hostFunctionCollection.name}`); lines.push(""); - lines.push(hostFunctionCollection.description); - lines.push(""); for (const [functionName, definition] of Object.entries( hostFunctionCollection.functions, )) { const sidecarHostFunction = hostFunctionToSidecarDefinition(definition); - const signature = buildHostFunctionFlagSignature(sidecarHostFunction.inputSchema); + const signature = buildHostFunctionFlagSignature( + sidecarHostFunction.inputSchema, + ); const suffix = signature.length > 0 ? ` ${signature}` : ""; + const description = hostFunctionDescription(definition); lines.push( - `- \`agentos-${hostFunctionCollection.name} ${functionName}${suffix}\` - ${definition.description}`, + description.length > 0 + ? `- \`agentos-${hostFunctionCollection.name} ${functionName}${suffix}\` - ${description}` + : `- \`agentos-${hostFunctionCollection.name} ${functionName}${suffix}\``, ); } lines.push(""); @@ -1989,13 +2007,18 @@ async function handleHostCallback( } } -function buildHostFunctionMap(hostFunctions: HostFunctions[]): Map { +function buildHostFunctionMap( + hostFunctions: ResolvedHostFunctions[], +): Map { const hostFunctionMap = new Map(); for (const hostFunctionCollection of hostFunctions) { for (const [functionName, definition] of Object.entries( hostFunctionCollection.functions, )) { - hostFunctionMap.set(`${hostFunctionCollection.name}:${functionName}`, definition); + hostFunctionMap.set( + `${hostFunctionCollection.name}:${functionName}`, + definition, + ); } } return hostFunctionMap; @@ -2009,7 +2032,7 @@ interface HostCommandCallbackInput { } interface HostCallbackContext { - hostFunctions: HostFunctions[]; + hostFunctions: ResolvedHostFunctions[]; hostFunctionMap: ReadonlyMap; readFile(path: string): Promise; } @@ -2223,7 +2246,11 @@ async function handleHostCommandCallback( return handleAgentOsRegistryCommand(command, context); } if (directHostFunctions) { - return handleAgentOsHostFunctionCommand(command, context, directHostFunctions); + return handleAgentOsHostFunctionCommand( + command, + context, + directHostFunctions, + ); } throw new Error(`Unknown host callback command "${command.command}"`); } @@ -2273,11 +2300,14 @@ async function handleAgentOsRegistryCommand( async function handleAgentOsHostFunctionCommand( command: HostCommandCallbackInput, context: HostCallbackContext, - hostFunctionCollection: HostFunctions, + hostFunctionCollection: ResolvedHostFunctions, ): Promise { const [functionName, helpOrFirstArg, ...rest] = command.args; if (!functionName || isHelpFlag(functionName)) { - return describeHostFunctionsPayload(context.hostFunctions, hostFunctionCollection.name); + return describeHostFunctionsPayload( + context.hostFunctions, + hostFunctionCollection.name, + ); } if (helpOrFirstArg && isHelpFlag(helpOrFirstArg)) { return describeHostFunctionPayload(hostFunctionCollection, functionName); @@ -2300,7 +2330,7 @@ async function invokeHostFunction({ cwd, context, }: { - hostFunctionCollection: HostFunctions; + hostFunctionCollection: ResolvedHostFunctions; functionName: string; args: string[]; cwd: string; @@ -2467,18 +2497,19 @@ function parseHostFunctionArgv( return input; } -function listHostFunctionsPayload(hostFunctions: HostFunctions[]): unknown { +function listHostFunctionsPayload( + hostFunctions: ResolvedHostFunctions[], +): unknown { return { hostFunctions: hostFunctions.map((hostFunctionCollection) => ({ name: hostFunctionCollection.name, - description: hostFunctionCollection.description, functions: Object.keys(hostFunctionCollection.functions), })), }; } function describeHostFunctionsPayload( - hostFunctions: HostFunctions[], + hostFunctions: ResolvedHostFunctions[], collectionName: string, ): unknown { const hostFunctionCollection = hostFunctions.find( @@ -2491,13 +2522,12 @@ function describeHostFunctionsPayload( } return { name: hostFunctionCollection.name, - description: hostFunctionCollection.description, functions: Object.fromEntries( Object.entries(hostFunctionCollection.functions).map( ([functionName, definition]) => [ functionName, { - description: definition.description, + description: hostFunctionDescription(definition), flags: describeHostFunctionFlags( hostFunctionToSidecarDefinition(definition).inputSchema, ), @@ -2509,7 +2539,7 @@ function describeHostFunctionsPayload( } function describeHostFunctionPayload( - hostFunctionCollection: HostFunctions, + hostFunctionCollection: ResolvedHostFunctions, functionName: string, ): unknown { const definition = hostFunctionCollection.functions[functionName]; @@ -2521,7 +2551,7 @@ function describeHostFunctionPayload( return { collection: hostFunctionCollection.name, function: functionName, - description: definition.description, + description: hostFunctionDescription(definition), flags: describeHostFunctionFlags( hostFunctionToSidecarDefinition(definition).inputSchema, ), @@ -2533,13 +2563,15 @@ function describeHostFunctionPayload( }; } -function hostFunctionsNames(hostFunctions: HostFunctions[]): string { +function hostFunctionsNames(hostFunctions: ResolvedHostFunctions[]): string { return hostFunctions .map((hostFunctionCollection) => hostFunctionCollection.name) .join(", "); } -function hostFunctionNames(hostFunctionCollection: HostFunctions): string { +function hostFunctionNames( + hostFunctionCollection: ResolvedHostFunctions, +): string { return Object.keys(hostFunctionCollection.functions).join(", "); } @@ -2556,7 +2588,7 @@ async function registerHostFunctionsOnSidecar( client: SidecarProcess, session: AuthenticatedSession, vm: CreatedVm, - hostFunctions: HostFunctions[], + hostFunctions: ResolvedHostFunctions[], ): Promise { if (hostFunctions.length === 0) { return ""; @@ -2565,7 +2597,7 @@ async function registerHostFunctionsOnSidecar( for (const hostFunctionCollection of hostFunctions) { await client.registerHostCallbacks(session, vm, { name: hostFunctionCollection.name, - description: hostFunctionCollection.description, + description: "", commandAliases: [`agentos-${hostFunctionCollection.name}`], registryCommandAliases: ["agentos"], callbacks: Object.fromEntries( @@ -2932,7 +2964,7 @@ export class AgentOs { private _acpTerminalCounter = 0; private _softwareRoots: SoftwareRoot[]; private _cronManager!: CronManager; - private _hostFunctions: HostFunctions[] = []; + private _hostFunctions: ResolvedHostFunctions[] = []; private _hostFunctionReference = ""; private _hostMounts: HostMountInfo[]; private _env: Record; @@ -3121,8 +3153,15 @@ export class AgentOs { return getSharedAgentOsSidecarInternal(options); } - static async create(options?: AgentOsOptions): Promise { - options = parseAgentOsOptions(options); + static async create( + callerOptions?: AgentOsOptions, + ): Promise { + // The generic exists only so each `execute` infers its input from its own + // `inputSchema`. Past this point the concrete schemas carry no meaning, so + // the rest of `create()` works with the plain option type. + let options: AgentOsOptions | undefined = parseAgentOsOptions( + callerOptions as AgentOsOptions | undefined, + ); // Default software is FULLY DYNAMIC: this package's own NON-agent // @agentos-software/* dependencies (e.g. common), each default-exporting // its registry-built descriptor. Agent packages are NOT projected here — @@ -3159,25 +3198,26 @@ export class AgentOs { // sidecar owns agent resolution, agent enumeration, and agent snapshot // bundle loading from the projected package dirs. const localMounts = await resolveCompatLocalMounts(options?.mounts); - if (options?.hostFunctions && options.hostFunctions.length > 0) { - validateHostFunctions(options.hostFunctions); - } + // Keys become command names, so resolve and check them before anything + // else in `create()` allocates a sidecar or a sandbox. + const resolvedHostFunctions = options?.hostFunctions + ? resolveHostFunctions(options.hostFunctions) + : []; // Resolve the sidecar handle before starting an external sandbox so option // validation failures cannot leak provider resources. const sidecar = resolveAgentOsSidecar(options?.sidecar); options = await resolveSandboxOptions(options); const sandboxDisposeHooks = getSandboxDisposeHooks(options); - const hostFunctions = options.hostFunctions; + const hostFunctions = resolvedHostFunctions; const createVmAdmin = async (): Promise => { // The `/opt/agentos` projection is built by the sidecar from the // forwarded `packages` (it owns the staging dir + read-only mount, and // runtime `linkSoftware` appends to that live dir). The client no longer // stages packages host-side. - const hostFunctionBootstrapCommands = collectHostFunctionBootstrapCommands( - hostFunctions ?? [], - ); + const hostFunctionBootstrapCommands = + collectHostFunctionBootstrapCommands(hostFunctions); const bootstrapCommands = [ ...RUNTIME_BOOTSTRAP_COMMANDS, ...hostFunctionBootstrapCommands, @@ -3291,7 +3331,7 @@ export class AgentOs { for (const command of configuredVm.projectedCommands) { commandGuestPaths.set(command.name, command.guestPath); } - if (hostFunctions && hostFunctions.length > 0) { + if (hostFunctions.length > 0) { hostFunctionReference = await registerHostFunctionsOnSidecar( client, session, @@ -3364,7 +3404,7 @@ export class AgentOs { ), ), ), - hostFunctions: hostFunctions ?? [], + hostFunctions, hostFunctionReference, async dispose() { if (kernel) { diff --git a/packages/core/src/host-functions.ts b/packages/core/src/host-functions.ts deleted file mode 100644 index 7a53093f87..0000000000 --- a/packages/core/src/host-functions.ts +++ /dev/null @@ -1,86 +0,0 @@ -import type { ZodType } from "zod"; - -/** Maximum length for host function descriptions (characters). */ -export const MAX_HOST_FUNCTION_DESCRIPTION_LENGTH = 200; - -/** - * A single function that executes on the host. - */ -export interface HostFunction { - /** Description shown to the agent in --help and prompt docs. Max 200 characters. */ - description: string; - /** Zod schema for the input. Drives CLI flag generation and validation. */ - inputSchema: ZodType; - /** Runs on the host when the agent invokes the function. */ - execute: (input: INPUT) => Promise | OUTPUT; - /** Examples included in auto-generated prompt docs. */ - examples?: HostFunctionExample[]; - /** Timeout in ms. Default: 30000. */ - timeout?: number; -} - -export interface HostFunctionExample { - /** Human description of what this example does. */ - description: string; - /** The input args for the example. */ - input: INPUT; -} - -/** - * A named collection of host functions. Becomes a CLI binary: agentos-{name}. - */ -export interface HostFunctions { - /** Collection name. Must be lowercase alphanumeric + hyphens. Becomes the CLI suffix: agentos-{name}. */ - name: string; - /** Description shown in `agentos list-host-functions` and prompt docs. */ - description: string; - /** The functions in this collection. Keys become subcommands. */ - functions: Record; -} - -/** Helper to create a host function with type inference. */ -export function hostFunction( - def: HostFunction, -): HostFunction { - return def; -} - -/** Helper to create a named host function collection. */ -export function hostFunctions(def: HostFunctions): HostFunctions { - return def; -} - -const HOST_FUNCTION_COMMAND_NAME_RE = /^[a-z0-9]+(?:-[a-z0-9]+)*$/; - -function validateHostFunctionCommandName(kind: string, name: string): void { - if (HOST_FUNCTION_COMMAND_NAME_RE.test(name)) { - return; - } - throw new Error( - `${kind} name "${name}" must be lowercase alphanumeric with optional single hyphen separators`, - ); -} - -/** - * Validate every host function collection and function. - */ -export function validateHostFunctions(collections: HostFunctions[]): void { - for (const collection of collections) { - validateHostFunctionCommandName("Host function collection", collection.name); - if (collection.description.length > MAX_HOST_FUNCTION_DESCRIPTION_LENGTH) { - throw new Error( - `Host function collection "${collection.name}" description is ${collection.description.length} characters, max is ${MAX_HOST_FUNCTION_DESCRIPTION_LENGTH}`, - ); - } - for (const [functionName, definition] of Object.entries( - collection.functions, - )) { - validateHostFunctionCommandName("Host function", functionName); - if (definition.description.length > MAX_HOST_FUNCTION_DESCRIPTION_LENGTH) { - throw new Error( - `Host function "${collection.name}/${functionName}" description is ${definition.description.length} characters, max is ${MAX_HOST_FUNCTION_DESCRIPTION_LENGTH}`, - ); - } - } - } -} diff --git a/packages/core/src/index.ts b/packages/core/src/index.ts index d4ab118a1d..41e12f24ad 100644 --- a/packages/core/src/index.ts +++ b/packages/core/src/index.ts @@ -1,7 +1,19 @@ // @rivet-dev/agentos +export { + SidecarProcessError, + SidecarProcessExited, + SidecarRejectedError, + type SidecarRejectionDetail, + SidecarSilenceTimeout, +} from "@rivet-dev/agentos-runtime-core/sidecar-errors"; export { AgentOs, AgentOsSidecar } from "./agent-os.js"; -export type * from "./language-execution.js"; +export { + isPackageDescriptor, + OPT_AGENTOS_BIN, + OPT_AGENTOS_ROOT, + tryReadAgentosPackageManifest, +} from "./agentos-package.js"; export { CronManager, InvalidScheduleError, @@ -13,18 +25,28 @@ export { hostDirMount, nodeModulesMount, } from "./host-dir-mount.js"; +export type { + HostFunction, + HostFunctionCollection, + HostFunctionCollections, + HostFunctionExample, + HostFunctionSchemas, + ResolvedHostFunctions, +} from "@rivet-dev/agentos-runtime-core/host-functions"; export { - hostFunction, - MAX_HOST_FUNCTION_DESCRIPTION_LENGTH, - hostFunctions, - validateHostFunctions, -} from "./host-functions.js"; -export type { HostFunction, HostFunctionExample, HostFunctions } from "./host-functions.js"; + hostFunctionCommandName, + hostFunctionDescription, + resolveHostFunctions, +} from "@rivet-dev/agentos-runtime-core/host-functions"; +export type * from "./language-execution.js"; +export { createSnapshotExport } from "./layers.js"; export { agentOsLimitsSchema, agentOsOptionFieldSchemas, agentOsOptionsSchema, + hostFunctionCollectionSchema, hostFunctionSchema, + hostFunctionsSchema, mountConfigSchema, nativeMountConfigSchema, parseAgentOsOptions, @@ -32,25 +54,9 @@ export { rootFilesystemConfigSchema, sharedSidecarConfigSchema, sidecarConfigSchema, - hostFunctionsSchema, sidecarRuntimeConfigSchema, } from "./options-schema.js"; -export { createSnapshotExport } from "./layers.js"; export { defineSoftware } from "./packages.js"; -export { - isPackageDescriptor, - OPT_AGENTOS_BIN, - OPT_AGENTOS_ROOT, - tryReadAgentosPackageManifest, -} from "./agentos-package.js"; -export { KernelError } from "./runtime-compat.js"; -export { - SidecarProcessError, - SidecarProcessExited, - SidecarRejectedError, - type SidecarRejectionDetail, - SidecarSilenceTimeout, -} from "@rivet-dev/agentos-runtime-core/sidecar-errors"; export type { ExecOptions, ExecResult, @@ -60,9 +66,10 @@ export type { VirtualDirEntry, VirtualStat, } from "./runtime.js"; +export { KernelError } from "./runtime-compat.js"; export { - createSandboxHostFunctions, createSandboxFs, + createSandboxHostFunctions, getSandboxDisposeHooks, resolveSandboxOptions, } from "./sandbox.js"; diff --git a/packages/core/src/options-schema.ts b/packages/core/src/options-schema.ts index 41c6b196e5..1910125a06 100644 --- a/packages/core/src/options-schema.ts +++ b/packages/core/src/options-schema.ts @@ -6,7 +6,11 @@ import type { LimitWarningHandler, NativeMountConfig, } from "./agent-os.js"; -import type { HostFunction, HostFunctions } from "./host-functions.js"; +import type { + HostFunction, + HostFunctionCollection, + HostFunctionCollections, +} from "@rivet-dev/agentos-runtime-core/host-functions"; const stringArray = z.array(z.string()); const nonNegativeInteger = z.number().int().nonnegative(); @@ -363,7 +367,6 @@ const hostFunctionExampleSchema = z export const hostFunctionSchema = z .object({ - description: z.string(), inputSchema: z.custom( (value) => typeof value === "object" && value !== null, { @@ -376,13 +379,15 @@ export const hostFunctionSchema = z }) .strict() as z.ZodType; -export const hostFunctionsSchema = z - .object({ - name: z.string(), - description: z.string(), - functions: z.record(z.string(), hostFunctionSchema), - }) - .strict() as z.ZodType; +export const hostFunctionCollectionSchema = z.record( + z.string(), + hostFunctionSchema, +) as z.ZodType; + +export const hostFunctionsSchema = z.record( + z.string(), + hostFunctionCollectionSchema, +) as z.ZodType; /** * Shared AgentOsOptions field schemas. @@ -428,7 +433,7 @@ export const agentOsOptionFieldSchemas = { message: "Expected schedule driver object", }) .optional(), - hostFunctions: z.array(hostFunctionsSchema).optional(), + hostFunctions: hostFunctionsSchema.optional(), permissions: permissionsSchema.optional(), sidecar: sidecarConfigSchema.optional(), limits: agentOsLimitsSchema.optional(), diff --git a/packages/core/src/sandbox.ts b/packages/core/src/sandbox.ts index 93826d2aa9..71a94968f6 100644 --- a/packages/core/src/sandbox.ts +++ b/packages/core/src/sandbox.ts @@ -1,10 +1,14 @@ +import type { ZodType } from "zod"; import { z } from "zod"; import type { MountConfig, MountConfigJsonObject, NativeMountPluginDescriptor, } from "./agent-os.js"; -import type { HostFunction, HostFunctions } from "./host-functions.js"; +import type { + HostFunction, + HostFunctionCollections, +} from "@rivet-dev/agentos-runtime-core/host-functions"; export interface AgentOsSandboxProcessResult { stdout?: string; @@ -106,7 +110,7 @@ type SandboxDisposeHook = () => void | Promise; export type AgentOsSandboxExpandedOptions = { mounts?: MountConfig[]; - hostFunctions?: HostFunctions[]; + hostFunctions?: HostFunctionCollections; [sandboxDisposeHooks]?: SandboxDisposeHook[]; }; @@ -129,12 +133,6 @@ interface SerializableSandboxClient { defaultHeaders?: RequestInit["headers"]; } -function hostFunction( - def: HostFunction, -): HostFunction { - return def; -} - function normalizeHeaders( headers: RequestInit["headers"] | undefined, ): Record | undefined { @@ -196,21 +194,28 @@ export function createSandboxFs( }; } +/** + * Identity helper so each `execute` below infers its input from its own + * `inputSchema`. Callers writing an object literal get that from + * `AgentOs.create()`; code that builds a collection programmatically has no + * contextual type to infer from, so it needs this. + */ +function hostFunction( + definition: HostFunction, +): HostFunction { + return definition; +} + export function createSandboxHostFunctions( input: ResolvedSandboxOptions | AgentOsSandboxClientOptions, -): HostFunctions { +) { const options = input; const { client } = options; return { - name: "sandbox", - description: - "Execute commands and manage processes in a remote sandbox environment.", - functions: { - "run-command": hostFunction({ - description: - "Run a command synchronously in the sandbox and return its stdout, stderr, and exit code.", - inputSchema: z.object({ + runCommand: hostFunction({ + inputSchema: z + .object({ command: z .string() .describe("The command to execute (e.g. 'ls', 'python3')."), @@ -218,129 +223,139 @@ export function createSandboxHostFunctions( cwd: z.string().optional(), env: z.record(z.string(), z.string()).optional(), timeoutMs: z.number().optional(), - }), - timeout: 120_000, - execute: async (input) => { - const result = await client.runProcess(input); - return { - stdout: result.stdout, - stderr: result.stderr, - exitCode: result.exitCode, - timedOut: result.timedOut, - durationMs: result.durationMs, - }; - }, - }), - - "create-process": hostFunction({ - description: - "Start a long-running background process in the sandbox. Returns a process ID for later management.", - inputSchema: z.object({ + }) + .describe( + "Run a command synchronously in the sandbox and return its stdout, stderr, and exit code.", + ), + timeout: 120_000, + execute: async (input) => { + const result = await client.runProcess(input); + return { + stdout: result.stdout, + stderr: result.stderr, + exitCode: result.exitCode, + timedOut: result.timedOut, + durationMs: result.durationMs, + }; + }, + }), + + createProcess: hostFunction({ + inputSchema: z + .object({ command: z.string(), args: z.array(z.string()).optional(), cwd: z.string().optional(), env: z.record(z.string(), z.string()).optional(), - }), - execute: async (input) => { - const proc = await client.createProcess(input); - return { - id: proc.id, - command: proc.command, - args: proc.args, - status: proc.status, - pid: proc.pid, - }; - }, - }), - - "list-processes": hostFunction({ - description: "List all processes running in the sandbox.", - inputSchema: z.object({}), - execute: async () => { - const result = await client.listProcesses(); - return { - processes: result.processes.map((p) => ({ - id: p.id, - command: p.command, - args: p.args, - status: p.status, - exitCode: p.exitCode, - pid: p.pid, - })), - }; - }, - }), - - "stop-process": hostFunction({ - description: "Gracefully stop a running process in the sandbox.", - inputSchema: z.object({ id: z.string() }), - execute: async (input) => { - const proc = await client.stopProcess(input.id); - return { - id: proc.id, - status: proc.status, - exitCode: proc.exitCode, - }; - }, - }), - - "kill-process": hostFunction({ - description: "Forcefully kill a running process in the sandbox.", - inputSchema: z.object({ id: z.string() }), - execute: async (input) => { - const proc = await client.killProcess(input.id); - return { - id: proc.id, - status: proc.status, - exitCode: proc.exitCode, - }; - }, - }), - - "get-process-logs": hostFunction({ - description: "Get stdout/stderr logs from a sandbox process.", - inputSchema: z.object({ + }) + .describe( + "Start a long-running background process in the sandbox. Returns a process ID for later management.", + ), + execute: async (input) => { + const proc = await client.createProcess(input); + return { + id: proc.id, + command: proc.command, + args: proc.args, + status: proc.status, + pid: proc.pid, + }; + }, + }), + + listProcesses: hostFunction({ + inputSchema: z + .object({}) + .describe("List all processes running in the sandbox."), + execute: async () => { + const result = await client.listProcesses(); + return { + processes: result.processes.map((p) => ({ + id: p.id, + command: p.command, + args: p.args, + status: p.status, + exitCode: p.exitCode, + pid: p.pid, + })), + }; + }, + }), + + stopProcess: hostFunction({ + inputSchema: z + .object({ id: z.string() }) + .describe("Gracefully stop a running process in the sandbox."), + execute: async (input) => { + const proc = await client.stopProcess(input.id); + return { + id: proc.id, + status: proc.status, + exitCode: proc.exitCode, + }; + }, + }), + + killProcess: hostFunction({ + inputSchema: z + .object({ id: z.string() }) + .describe("Forcefully kill a running process in the sandbox."), + execute: async (input) => { + const proc = await client.killProcess(input.id); + return { + id: proc.id, + status: proc.status, + exitCode: proc.exitCode, + }; + }, + }), + + getProcessLogs: hostFunction({ + inputSchema: z + .object({ id: z.string(), stream: z.enum(["stdout", "stderr", "combined"]).optional(), tail: z.number().optional(), - }), - execute: async (input) => { - const result = await client.getProcessLogs(input.id, { - stream: input.stream, - tail: input.tail, - }); - return { - logs: result.entries.map((e) => { - const data = - e.encoding === "base64" - ? Buffer.from(e.data, "base64").toString("utf-8") - : e.data; - return { - data, - stream: e.stream, - timestampMs: e.timestampMs, - }; - }), - }; - }, - }), - - "send-input": hostFunction({ - description: - "Send text input to an interactive sandbox process via stdin.", - inputSchema: z.object({ + }) + .describe("Get stdout/stderr logs from a sandbox process."), + execute: async (input) => { + const result = await client.getProcessLogs(input.id, { + stream: input.stream, + tail: input.tail, + }); + return { + logs: result.entries.map((e) => { + const data = + e.encoding === "base64" + ? Buffer.from(e.data, "base64").toString("utf-8") + : e.data; + return { + data, + stream: e.stream, + timestampMs: e.timestampMs, + }; + }), + }; + }, + }), + + sendInput: hostFunction({ + inputSchema: z + .object({ id: z.string(), data: z.string(), - }), - execute: async (input) => { - await client.sendProcessInput(input.id, { - data: Buffer.from(input.data, "utf-8").toString("base64"), - encoding: "base64", - }); - return { sent: true }; - }, - }), - }, + }) + .describe( + "Send text input to an interactive sandbox process via stdin.", + ), + execute: async (input) => { + await client.sendProcessInput(input.id, { + data: Buffer.from(input.data, "utf-8").toString("base64"), + encoding: "base64", + }); + return { sent: true }; + }, + }), }; } @@ -434,7 +449,7 @@ export async function resolveSandboxOptions< ): Promise< Omit & { mounts?: MountConfig[]; - hostFunctions?: HostFunctions[]; + hostFunctions?: HostFunctionCollections; } > { const { sandbox, ...rest } = options; @@ -447,7 +462,7 @@ export async function resolveSandboxOptions< const sandboxOptions = normalizedSandbox.options; const expanded = rest as Omit & { mounts?: MountConfig[]; - hostFunctions?: HostFunctions[]; + hostFunctions?: HostFunctionCollections; }; const mountPath = sandboxOptions.mountPath ?? "/mnt/sandbox"; const mounts = [ @@ -458,10 +473,10 @@ export async function resolveSandboxOptions< readOnly: sandboxOptions.readOnly, }, ]; - const hostFunctions = [ - ...(expanded.hostFunctions ?? []), - createSandboxHostFunctions(sandboxOptions), - ]; + const hostFunctions = { + ...expanded.hostFunctions, + sandbox: createSandboxHostFunctions(sandboxOptions), + }; return attachSandboxDisposeHooks( { diff --git a/packages/core/src/types.ts b/packages/core/src/types.ts index e3290847aa..89b341893e 100644 --- a/packages/core/src/types.ts +++ b/packages/core/src/types.ts @@ -85,7 +85,6 @@ export type { PackageRef, SoftwarePackageRef, } from "./agentos-package.js"; -export type { HostFunction, HostFunctionExample, HostFunctions } from "./host-functions.js"; export type { CronAction, CronActionInfo, @@ -103,6 +102,14 @@ export type { HostDirMountPluginConfig, NodeModulesMountConfig, } from "./host-dir-mount.js"; +export type { + HostFunction, + HostFunctionCollection, + HostFunctionCollections, + HostFunctionExample, + HostFunctionSchemas, + ResolvedHostFunctions, +} from "@rivet-dev/agentos-runtime-core/host-functions"; export type { FilesystemSnapshotExport, LayerHandle, diff --git a/packages/core/tests/acp-reactor-regression.test.ts b/packages/core/tests/acp-reactor-regression.test.ts index 57bd964d78..09c223caf5 100644 --- a/packages/core/tests/acp-reactor-regression.test.ts +++ b/packages/core/tests/acp-reactor-regression.test.ts @@ -299,27 +299,24 @@ describe("ACP adapter reactor regression", () => { test("routes a delayed host-tool response past 256 ordinary updates and keeps the session reusable", async () => { let hostToolCalls = 0; const hostToolInputs: Array<{ a: number; b: number }> = []; - const mathFunctions = hostFunctions({ - name: "math", - description: "Math utilities", - functions: { - add: hostFunction({ - description: "Add two numbers", - inputSchema: z.object({ + const mathFunctions = { + add: { + inputSchema: z + .object({ a: z.number(), b: z.number(), - }), - execute: async ({ a, b }) => { - hostToolCalls += 1; - hostToolInputs.push({ a, b }); - await new Promise((resolveDelay) => - setTimeout(resolveDelay, 50), - ); - return { sum: a + b }; - }, - }), + }) + .describe("Add two numbers"), + execute: async ({ a, b }) => { + hostToolCalls += 1; + hostToolInputs.push({ a, b }); + await new Promise((resolveDelay) => + setTimeout(resolveDelay, 50), + ); + return { sum: a + b }; + }, }, - }); + }; const agentPackage = createProjectedAgentPackage({ name: "acp-reactor-regression", adapterScript: ACP_REACTOR_ADAPTER, @@ -332,7 +329,7 @@ describe("ACP adapter reactor regression", () => { mounts: moduleAccessMounts(MODULE_ACCESS_CWD), defaultSoftware: false, software: [coreutils, agentPackage.software], - hostFunctions: [mathFunctions], + hostFunctions: { math: mathFunctions }, permissions: { fs: "allow", childProcess: "allow", diff --git a/packages/core/tests/host-function-permissions.test.ts b/packages/core/tests/host-function-permissions.test.ts index ac7428ada9..516933d1e8 100644 --- a/packages/core/tests/host-function-permissions.test.ts +++ b/packages/core/tests/host-function-permissions.test.ts @@ -93,35 +93,29 @@ function commandHostCallbackFrame(command: string, args: string[]) { }; } -const mathFunctions = hostFunctions({ - name: "math", - description: "Math utilities", - functions: { - add: hostFunction({ - description: "Add two numbers", - inputSchema: z.object({ +const mathFunctions = { + add: { + inputSchema: z + .object({ a: z.number(), b: z.number(), - }), - execute: ({ a, b }) => ({ sum: a + b }), - }), + }) + .describe("Add two numbers"), + execute: ({ a, b }) => ({ sum: a + b }), }, -}); +}; -const duplicateMathFunctions = hostFunctions({ - name: "math", - description: "Duplicate math utilities", - functions: { - multiply: hostFunction({ - description: "Multiply two numbers", - inputSchema: z.object({ +const duplicateMathFunctions = { + multiply: { + inputSchema: z + .object({ a: z.number(), b: z.number(), - }), - execute: ({ a, b }) => ({ product: a * b }), - }), + }) + .describe("Multiply two numbers"), + execute: ({ a, b }) => ({ product: a * b }), }, -}); +}; async function runCommand(vm: AgentOs, command: string, args: string[]) { const stdoutChunks: string[] = []; @@ -150,18 +144,18 @@ describe("hostFunction collection permissions", () => { vm = null; }); - test("rejects duplicate hostFunction collection registration with a conflict", async () => { + test("rejects two collection keys that resolve to the same command name", async () => { await expect( AgentOs.create({ - hostFunctions: [mathFunctions, duplicateMathFunctions], + hostFunctions: { math: mathFunctions, Math: duplicateMathFunctions }, }), - ).rejects.toThrow(/conflict: hostFunction collection already registered: math/); + ).rejects.toThrow(/both resolve to the command name "math"/); }); test("allows hostFunction collection invocation with default permissions", async () => { vm = await AgentOs.create({ software: [common], - hostFunctions: [mathFunctions], + hostFunctions: { math: mathFunctions }, }); const result = await runCommand(vm, "agentos-math", [ @@ -181,7 +175,7 @@ describe("hostFunction collection permissions", () => { test("denies hostFunction collection invocation by default until hostFunction permissions are granted", async () => { vm = await AgentOs.create({ software: [common], - hostFunctions: [mathFunctions], + hostFunctions: { math: mathFunctions }, permissions: { fs: "allow", childProcess: "allow", @@ -204,7 +198,7 @@ describe("hostFunction collection permissions", () => { test("allows hostFunction collection invocation when a matching hostFunction permission is granted", async () => { vm = await AgentOs.create({ software: [common], - hostFunctions: [mathFunctions], + hostFunctions: { math: mathFunctions }, permissions: { fs: "allow", childProcess: "allow", @@ -247,26 +241,23 @@ describe("host-function collection permissions: raw host_callback RPC path", () // N-001 (J.1/J.2): host_callback RPC must honor hostFunction.invoke deny. test("denies host_callback RPC hostFunction invocation when hostFunction.invoke policy is deny (not just the CLI path)", async () => { const executed: unknown[] = []; - const spyFunctions = hostFunctions({ - name: "math", - description: "Math utilities", - functions: { - add: hostFunction({ - description: "Add two numbers", - inputSchema: z.object({ a: z.number(), b: z.number() }), - execute: ({ a, b }) => { - executed.push({ a, b }); - return { sum: a + b }; - }, - }), + const spyFunctions = { + add: { + inputSchema: z + .object({ a: z.number(), b: z.number() }) + .describe("Add two numbers"), + execute: ({ a, b }) => { + executed.push({ a, b }); + return { sum: a + b }; + }, }, - }); + }; const created = await createVmCapturingHandler({ // No `software` needed: this exercises the raw host_callback RPC // handler directly (the guest-controlled path), which does not spawn // any in-VM CLI. Keeping the VM minimal makes the safeguard fast. - hostFunctions: [spyFunctions], + hostFunctions: { math: spyFunctions }, permissions: { fs: "allow", childProcess: "allow", @@ -286,37 +277,33 @@ describe("host-function collection permissions: raw host_callback RPC path", () expect(response.type).toBe("host_callback_result"); expect(response.result).toBeUndefined(); expect(typeof response.error).toBe("string"); - expect(response.error).toMatch(/hostFunction\.invoke|EACCES|denied|permission/i); + expect(response.error).toMatch( + /hostFunction\.invoke|EACCES|denied|permission/i, + ); }); // N-002 (J.2): host_callback RPC must respect hostFunction.invoke pattern scope. test("host_callback RPC respects hostFunction.invoke pattern scope and denies a non-matching hostFunction", async () => { const executed: string[] = []; - const dangerFunctions = hostFunctions({ - name: "math", - description: "Math utilities with a dangerous hostFunction", - functions: { - safe: hostFunction({ - description: "Safe op", - inputSchema: z.object({ x: z.number() }), - execute: ({ x }) => { - executed.push("safe"); - return { x }; - }, - }), - danger: hostFunction({ - description: "Dangerous op", - inputSchema: z.object({ x: z.number() }), - execute: ({ x }) => { - executed.push("danger"); - return { x }; - }, - }), + const dangerFunctions = { + safe: { + inputSchema: z.object({ x: z.number() }).describe("Safe op"), + execute: ({ x }) => { + executed.push("safe"); + return { x }; + }, }, - }); + danger: { + inputSchema: z.object({ x: z.number() }).describe("Dangerous op"), + execute: ({ x }) => { + executed.push("danger"); + return { x }; + }, + }, + }; const created = await createVmCapturingHandler({ - hostFunctions: [dangerFunctions], + hostFunctions: { math: dangerFunctions }, permissions: { fs: "allow", childProcess: "allow", @@ -339,7 +326,9 @@ describe("host-function collection permissions: raw host_callback RPC path", () expect(response.type).toBe("host_callback_result"); expect(response.result).toBeUndefined(); expect(typeof response.error).toBe("string"); - expect(response.error).toMatch(/hostFunction\.invoke|EACCES|denied|permission/i); + expect(response.error).toMatch( + /hostFunction\.invoke|EACCES|denied|permission/i, + ); }); // AOSFS-1 (P1, J.1/J.2): the raw host_callback RPC path is fully @@ -352,25 +341,22 @@ describe("host-function collection permissions: raw host_callback RPC path", () // pollution may occur. Asserts the system strips the hostile/extra keys. test("host_callback strips hostile/extra input keys; execute receives only validated Zod data and no prototype pollution", async () => { const seen: unknown[] = []; - const collection = hostFunctions({ - name: "math", - description: "Math utilities", - functions: { - add: hostFunction({ - description: "Add two numbers", - inputSchema: z.object({ a: z.number(), b: z.number() }), - execute: (input) => { - // Capture exactly what execute is handed. - seen.push(input); - const { a, b } = input; - return { sum: a + b }; - }, - }), + const collection = { + add: { + inputSchema: z + .object({ a: z.number(), b: z.number() }) + .describe("Add two numbers"), + execute: (input) => { + // Capture exactly what execute is handed. + seen.push(input); + const { a, b } = input; + return { sum: a + b }; + }, }, - }); + }; const created = await createVmCapturingHandler({ - hostFunctions: [collection], + hostFunctions: { math: collection }, permissions: { fs: "allow", childProcess: "allow", @@ -412,9 +398,7 @@ describe("host-function collection permissions: raw host_callback RPC path", () // No prototype pollution of Object.prototype on the host. expect(({} as Record).polluted).toBeUndefined(); expect(({} as Record).polluted2).toBeUndefined(); - expect( - Object.prototype.hasOwnProperty.call(Object.prototype, "polluted"), - ).toBe(false); + expect(Object.hasOwn(Object.prototype, "polluted")).toBe(false); }); // AOSFS-2 (P2): a guest can send schema-failing input on the raw host_callback @@ -423,23 +407,20 @@ describe("host-function collection permissions: raw host_callback RPC path", () // safeParse and return a validation error WITHOUT invoking execute. test("host_callback rejects schema-failing input without invoking execute", async () => { const executed: unknown[] = []; - const collection = hostFunctions({ - name: "math", - description: "Math utilities", - functions: { - add: hostFunction({ - description: "Add two numbers", - inputSchema: z.object({ a: z.number(), b: z.number() }), - execute: ({ a, b }) => { - executed.push({ a, b }); - return { sum: a + b }; - }, - }), + const collection = { + add: { + inputSchema: z + .object({ a: z.number(), b: z.number() }) + .describe("Add two numbers"), + execute: ({ a, b }) => { + executed.push({ a, b }); + return { sum: a + b }; + }, }, - }); + }; const created = await createVmCapturingHandler({ - hostFunctions: [collection], + hostFunctions: { math: collection }, permissions: { fs: "allow", childProcess: "allow", @@ -473,23 +454,20 @@ describe("host-function collection permissions: raw host_callback RPC path", () // re-discovery — assert the gate holds on this branch.) test("forged {type:'command'} host_callback is denied by hostFunction.invoke on the command dispatch branch", async () => { const executed: unknown[] = []; - const spyFunctions = hostFunctions({ - name: "math", - description: "Math utilities", - functions: { - add: hostFunction({ - description: "Add two numbers", - inputSchema: z.object({ a: z.number(), b: z.number() }), - execute: ({ a, b }) => { - executed.push({ a, b }); - return { sum: a + b }; - }, - }), + const spyFunctions = { + add: { + inputSchema: z + .object({ a: z.number(), b: z.number() }) + .describe("Add two numbers"), + execute: ({ a, b }) => { + executed.push({ a, b }); + return { sum: a + b }; + }, }, - }); + }; const created = await createVmCapturingHandler({ - hostFunctions: [spyFunctions], + hostFunctions: { math: spyFunctions }, permissions: { fs: "allow", childProcess: "allow", @@ -510,6 +488,8 @@ describe("host-function collection permissions: raw host_callback RPC path", () expect(response.type).toBe("host_callback_result"); expect(response.result).toBeUndefined(); expect(typeof response.error).toBe("string"); - expect(response.error).toMatch(/hostFunction\.invoke|EACCES|denied|permission/i); + expect(response.error).toMatch( + /hostFunction\.invoke|EACCES|denied|permission/i, + ); }); }); diff --git a/packages/core/tests/host-function-reference.test.ts b/packages/core/tests/host-function-reference.test.ts index 90c31088ec..536d39e900 100644 --- a/packages/core/tests/host-function-reference.test.ts +++ b/packages/core/tests/host-function-reference.test.ts @@ -1,6 +1,6 @@ import { afterEach, beforeEach, describe, expect, test } from "vitest"; import { z } from "zod"; -import { AgentOs, hostFunction, hostFunctions } from "../src/index.js"; +import { AgentOs } from "../src/index.js"; import { createProjectedAgentPackage, type ProjectedAgentPackage, @@ -46,26 +46,23 @@ process.stdin.on('data', (chunk) => { }); `; -const mathFunctions = hostFunctions({ - name: "math", - description: "Math utilities", - functions: { - add: hostFunction({ - description: "Add two numbers", - inputSchema: z.object({ +const mathFunctions = { + add: { + inputSchema: z + .object({ a: z.number(), b: z.number(), - }), - execute: ({ a, b }) => ({ sum: a + b }), - examples: [ - { - description: "Add 1 and 2", - input: { a: 1, b: 2 }, - }, - ], - }), + }) + .describe("Add two numbers"), + execute: ({ a, b }) => ({ sum: a + b }), + examples: [ + { + description: "Add 1 and 2", + input: { a: 1, b: 2 }, + }, + ], }, -}); +}; describe("hostFunction reference registration", () => { let vm: AgentOs; @@ -79,7 +76,7 @@ describe("hostFunction reference registration", () => { vm = await AgentOs.create({ defaultSoftware: false, software: [agentPackage.software], - hostFunctions: [mathFunctions], + hostFunctions: { math: mathFunctions }, }); }); @@ -89,19 +86,18 @@ describe("hostFunction reference registration", () => { }); test("stores generated hostFunction reference markdown on the VM", () => { - const bindingReference = (vm as unknown as { _bindingReference: string }) - ._bindingReference; + const reference = (vm as unknown as { _hostFunctionReference: string }) + ._hostFunctionReference; - expect(bindingReference).toContain("## Available Host Functions"); - expect(bindingReference).toContain( + expect(reference).toContain("## Available Host Functions"); + expect(reference).toContain( "Run `agentos list-host-functions` to see all available host functions.", ); - expect(bindingReference).toContain("### math"); - expect(bindingReference).toContain("Math utilities"); - expect(bindingReference).toContain( - "`agentos-math add --a --b `", + expect(reference).toContain("### math"); + expect(reference).toContain( + "`agentos-math add --a --b ` - Add two numbers", ); - expect(bindingReference).toContain("Add 1 and 2"); + expect(reference).toContain("Add 1 and 2"); }); test("openSession injects the registered hostFunction reference into the system prompt", async () => { diff --git a/packages/core/tests/host-functions-zod.test.ts b/packages/core/tests/host-functions-zod.test.ts index 54dfb1ca03..90d0510907 100644 --- a/packages/core/tests/host-functions-zod.test.ts +++ b/packages/core/tests/host-functions-zod.test.ts @@ -1,10 +1,10 @@ -import { describe, expect, test } from "vitest"; -import { z } from "zod"; -import { z as z3 } from "zod3"; import { HostFunctionSchemaConversionError, zodToJsonSchema, -} from "../src/host-functions-zod.js"; +} from "@rivet-dev/agentos-runtime-core/host-functions-zod"; +import { describe, expect, test } from "vitest"; +import { z } from "zod"; +import { z as z3 } from "zod3"; describe("zodToJsonSchema", () => { test("converts objects with supported scalar constraints", () => { diff --git a/packages/core/tests/host-functions.test.ts b/packages/core/tests/host-functions.test.ts index db86fc9a85..62b403a736 100644 --- a/packages/core/tests/host-functions.test.ts +++ b/packages/core/tests/host-functions.test.ts @@ -1,128 +1,102 @@ -import { describe, expect, test } from "vitest"; -import { z } from "zod"; import { HostFunctionSchemaConversionError, zodToJsonSchema, -} from "../src/host-functions-zod.js"; +} from "@rivet-dev/agentos-runtime-core/host-functions-zod"; +import { describe, expect, test } from "vitest"; +import { z } from "zod"; import { - MAX_HOST_FUNCTION_DESCRIPTION_LENGTH, - hostFunction, - hostFunctions, - validateHostFunctions, + hostFunctionCommandName, + hostFunctionDescription, + resolveHostFunctions, } from "../src/index.js"; -describe("host-function description limits", () => { - test("accepts collection and function descriptions at the exported limit", () => { - const description = "a".repeat(MAX_HOST_FUNCTION_DESCRIPTION_LENGTH); +const screenshot = { + inputSchema: z.object({ url: z.string() }).describe("Take a screenshot"), + execute: () => ({ ok: true }), +}; - expect(() => - validateHostFunctions([ - hostFunctions({ - name: "browser", - description, - functions: { - screenshot: hostFunction({ - description, - inputSchema: z.object({ url: z.string() }), - execute: () => ({ ok: true }), - }), - }, - }), - ]), - ).not.toThrow(); +describe("host-function names", () => { + test("converts camelCase keys to kebab-case command names", () => { + expect(hostFunctionCommandName("orderStore")).toBe("order-store"); + expect(hostFunctionCommandName("listOpenOrders")).toBe("list-open-orders"); + expect(hostFunctionCommandName("order-store")).toBe("order-store"); + expect(hostFunctionCommandName("orders")).toBe("orders"); }); - test("rejects collection descriptions longer than the exported limit", () => { + test("resolves both key spellings to the same command names", () => { + expect( + resolveHostFunctions({ orderStore: { listOrders: screenshot } }), + ).toEqual([ + { name: "order-store", functions: { "list-orders": screenshot } }, + ]); + }); + + test("rejects collection keys that cannot become command names", () => { expect(() => - validateHostFunctions([ - hostFunctions({ - name: "browser", - description: "a".repeat(MAX_HOST_FUNCTION_DESCRIPTION_LENGTH + 1), - functions: { - screenshot: hostFunction({ - description: "Take a screenshot", - inputSchema: z.object({ url: z.string() }), - execute: () => ({ ok: true }), - }), - }, - }), - ]), + resolveHostFunctions({ Browser_Host_Functions: { screenshot } }), ).toThrow( - `Host function collection "browser" description is ${MAX_HOST_FUNCTION_DESCRIPTION_LENGTH + 1} characters, max is ${MAX_HOST_FUNCTION_DESCRIPTION_LENGTH}`, + 'Host function collection name "Browser_Host_Functions" must be alphanumeric, written in camelCase or with single hyphen separators', ); }); - test("rejects function descriptions longer than the exported limit", () => { + test("rejects function keys that cannot become subcommands", () => { expect(() => - validateHostFunctions([ - hostFunctions({ - name: "browser", - description: "Browser automation", - functions: { - screenshot: hostFunction({ - description: "a".repeat(MAX_HOST_FUNCTION_DESCRIPTION_LENGTH + 1), - inputSchema: z.object({ url: z.string() }), - execute: () => ({ ok: true }), - }), - }, - }), - ]), + resolveHostFunctions({ browser: { screenshot_now: screenshot } }), ).toThrow( - `Host function "browser/screenshot" description is ${MAX_HOST_FUNCTION_DESCRIPTION_LENGTH + 1} characters, max is ${MAX_HOST_FUNCTION_DESCRIPTION_LENGTH}`, + 'Host function name "screenshot_now" must be alphanumeric, written in camelCase or with single hyphen separators', ); }); - test("rejects collection names that cannot become stable command names", () => { + test("rejects two collection keys that resolve to the same command name", () => { expect(() => - validateHostFunctions([ - hostFunctions({ - name: "Browser_Host_Functions", - description: "Browser automation", - functions: { - screenshot: hostFunction({ - description: "Take a screenshot", - inputSchema: z.object({ url: z.string() }), - execute: () => ({ ok: true }), - }), - }, - }), - ]), + resolveHostFunctions({ + orderStore: { screenshot }, + "order-store": { screenshot }, + }), ).toThrow( - 'Host function collection name "Browser_Host_Functions" must be lowercase alphanumeric with optional single hyphen separators', + 'Host function collections "orderStore" and "order-store" both resolve to the command name "order-store"', ); }); - test("rejects function names that cannot become stable subcommands", () => { + test("rejects two function keys that resolve to the same command name", () => { expect(() => - validateHostFunctions([ - hostFunctions({ - name: "browser-host-functions", - description: "Browser automation", - functions: { - screenshot_now: hostFunction({ - description: "Take a screenshot", - inputSchema: z.object({ url: z.string() }), - execute: () => ({ ok: true }), - }), - }, - }), - ]), + resolveHostFunctions({ + browser: { screenshotNow: screenshot, "screenshot-now": screenshot }, + }), ).toThrow( - 'Host function name "screenshot_now" must be lowercase alphanumeric with optional single hyphen separators', + 'Host functions "screenshotNow" and "screenshot-now" in collection "browser" both resolve to the command name "screenshot-now"', ); }); +}); - test("fails loudly when a host-function input schema uses an unsupported discriminated union", () => { - const definition = hostFunction({ - description: "Inspect a variant payload", - inputSchema: z.object({ - payload: z.discriminatedUnion("kind", [ - z.object({ kind: z.literal("text"), value: z.string() }), - z.object({ kind: z.literal("code"), status: z.number() }), - ]), +describe("host-function descriptions", () => { + test("reads the description from the input schema", () => { + expect(hostFunctionDescription(screenshot)).toBe("Take a screenshot"); + }); + + test("is empty when the schema carries no description", () => { + expect( + hostFunctionDescription({ + inputSchema: z.object({ url: z.string() }), + execute: () => ({ ok: true }), }), + ).toBe(""); + }); +}); + +describe("host-function schemas", () => { + test("fails loudly when a host-function input schema uses an unsupported discriminated union", () => { + const definition = { + inputSchema: z + .object({ + payload: z.discriminatedUnion("kind", [ + z.object({ kind: z.literal("text"), value: z.string() }), + z.object({ kind: z.literal("code"), status: z.number() }), + ]), + }) + .describe("Inspect a variant payload"), execute: () => ({ ok: true }), - }); + }; try { zodToJsonSchema(definition.inputSchema); diff --git a/packages/core/tests/migration-parity.test.ts b/packages/core/tests/migration-parity.test.ts index 336010c17e..8b448423e6 100644 --- a/packages/core/tests/migration-parity.test.ts +++ b/packages/core/tests/migration-parity.test.ts @@ -1,10 +1,10 @@ import { createServer, type IncomingMessage } from "node:http"; import { resolve } from "node:path"; -import { moduleAccessMounts } from "./helpers/node-modules-mount.js"; import { afterEach, describe, expect, test } from "vitest"; import { z } from "zod"; import { AgentOs, hostFunction, hostFunctions } from "../src/index.js"; import type { SessionStreamEntry } from "../src/session-api.js"; +import { moduleAccessMounts } from "./helpers/node-modules-mount.js"; import { createProjectedAgentPackage } from "./helpers/projected-agent-package.js"; import { promptResultText } from "./helpers/session-result.js"; @@ -123,29 +123,24 @@ process.stdin.on("data", (chunk) => { }); `.trim(); -const mathFunctions = hostFunctions({ - name: "math", - description: "Math utilities", - functions: { - add: hostFunction({ - description: "Add two numbers", - inputSchema: z.object({ +const mathFunctions = { + add: { + inputSchema: z + .object({ a: z.number(), b: z.number(), - }), - execute: ({ a, b }) => ({ sum: a + b }), - }), + }) + .describe("Add two numbers"), + execute: ({ a, b }) => ({ sum: a + b }), }, -}); +}; function assertNativeSidecar(vm: AgentOs): void { expect(vm.sidecar.describe()).toMatchObject({ state: "ready", }); expect("kernel" in (vm as unknown as Record)).toBe(false); - expect( - (vm as unknown as Record).kernel, - ).toBeUndefined(); + expect((vm as unknown as Record).kernel).toBeUndefined(); } async function runSpawnedProcess( @@ -235,7 +230,7 @@ describe("native sidecar migration parity gate", () => { test("covers registered host functions through guest command dispatch on the Rust sidecar path", async () => { const vm = await AgentOs.create({ defaultSoftware: false, - hostFunctions: [mathFunctions], + hostFunctions: { math: mathFunctions }, permissions: { fs: "allow", childProcess: "allow", @@ -247,7 +242,9 @@ describe("native sidecar migration parity gate", () => { }); assertNativeSidecar(vm); - const listed = await runSpawnedProcess(vm, "agentos", ["list-host-functions"]); + const listed = await runSpawnedProcess(vm, "agentos", [ + "list-host-functions", + ]); expect(listed.exitCode).toBe(0); expect(JSON.parse(listed.stdout)).toEqual({ ok: true, @@ -255,8 +252,7 @@ describe("native sidecar migration parity gate", () => { hostFunctions: [ { name: "math", - description: "Math utilities", - hostFunctions: ["add"], + functions: ["add"], }, ], }, diff --git a/packages/core/tests/options-schema.test.ts b/packages/core/tests/options-schema.test.ts index cae0f3f6f7..13b67032bc 100644 --- a/packages/core/tests/options-schema.test.ts +++ b/packages/core/tests/options-schema.test.ts @@ -60,13 +60,7 @@ describe("AgentOsOptions validation", () => { test("accepts hostFunctions as the public name for host-function collections", () => { expect( agentOsOptionsSchema.safeParse({ - hostFunctions: [ - { - name: "weather", - description: "Weather functions", - functions: {}, - }, - ], + hostFunctions: { weather: {} }, }).success, ).toBe(true); }); @@ -130,7 +124,7 @@ describe("AgentOsOptions validation", () => { } as never); expect(options).not.toHaveProperty("sandbox"); expect(options.mounts?.[0]?.path).toBe("/mnt/sandbox"); - expect(options.hostFunctions?.[0]?.name).toBe("sandbox"); + expect(Object.keys(options.hostFunctions ?? {})).toContain("sandbox"); for (const hook of getSandboxDisposeHooks(options)) { await hook(); @@ -187,15 +181,9 @@ describe("AgentOsOptions validation", () => { }, }, }, - hostFunctions: [ - { - name: "INVALID", - description: "Invalid hostFunction collection", - functions: {}, - }, - ], + hostFunctions: { INVALID_NAME: {} }, }), - ).rejects.toThrow(/must be lowercase alphanumeric/); + ).rejects.toThrow(/must be alphanumeric, written in camelCase/); expect(started).toBe(0); expect(disposed).toBe(0); }); diff --git a/packages/core/tests/public-api-exports.test.ts b/packages/core/tests/public-api-exports.test.ts index 1b4034a5b3..9f80df09c6 100644 --- a/packages/core/tests/public-api-exports.test.ts +++ b/packages/core/tests/public-api-exports.test.ts @@ -7,10 +7,6 @@ import { type AgentOsSidecarRuntimeConfig, agentOsLimitsSchema, agentOsOptionsSchema, - hostFunction, - hostFunctionSchema, - hostFunctions, - hostFunctionsSchema, type ContextDescriptor, CronManager, createHostDirBackend, @@ -18,6 +14,10 @@ import { defineSoftware, type ExecOptions, type HostDirMountPluginConfig, + hostFunctionCommandName, + hostFunctionDescription, + hostFunctionSchema, + hostFunctionsSchema, InvalidScheduleError, isPackageDescriptor, KernelError, @@ -25,7 +25,6 @@ import { type KernelExecResult, type KernelSpawnOptions, type LanguageSpawnOptions, - MAX_HOST_FUNCTION_DESCRIPTION_LENGTH, type MountConfigJsonPrimitive, mountConfigSchema, type NodeModulesMountConfig, @@ -40,6 +39,7 @@ import { type ProcessExit, type PromptResult, parseAgentOsOptions, + resolveHostFunctions, rootFilesystemConfigSchema, type SessionCapabilities, type SessionInfo, @@ -49,7 +49,6 @@ import { sidecarRuntimeConfigSchema, TimerScheduleDriver, type TimingMitigation, - validateHostFunctions, } from "../src/index.js"; describe("root public API exports", () => { @@ -63,10 +62,9 @@ describe("root public API exports", () => { expect(CronManager).toBeTypeOf("function"); expect(TimerScheduleDriver).toBeTypeOf("function"); expect(createHostDirBackend).toBeTypeOf("function"); - expect(hostFunction).toBeTypeOf("function"); - expect(hostFunctions).toBeTypeOf("function"); - expect(validateHostFunctions).toBeTypeOf("function"); - expect(MAX_HOST_FUNCTION_DESCRIPTION_LENGTH).toBeGreaterThan(0); + expect(hostFunctionCommandName).toBeTypeOf("function"); + expect(hostFunctionDescription).toBeTypeOf("function"); + expect(resolveHostFunctions).toBeTypeOf("function"); expect(agentOsLimitsSchema.safeParse({}).success).toBe(true); expect( agentOsLimitsSchema.safeParse({ diff --git a/packages/core/tests/sandbox-integration.test.ts b/packages/core/tests/sandbox-integration.test.ts index 3f6d641086..0f49e1288e 100644 --- a/packages/core/tests/sandbox-integration.test.ts +++ b/packages/core/tests/sandbox-integration.test.ts @@ -1,9 +1,9 @@ import common from "@agentos-software/common"; +import { afterAll, afterEach, beforeAll, describe, expect, test } from "vitest"; import { createSandboxFs, createSandboxHostFunctions, } from "../../agentos-sandbox/src/index.js"; -import { afterAll, afterEach, beforeAll, describe, expect, test } from "vitest"; import { AgentOs } from "../src/index.js"; import type { MockSandboxAgentHandle } from "../src/test/sandbox-agent.js"; import { startMockSandboxAgent } from "../src/test/sandbox-agent.js"; @@ -60,7 +60,9 @@ describe("sandbox quickstart truth test", () => { plugin: createSandboxFs({ client: sandbox.client }), }, ], - hostFunctions: [createSandboxHostFunctions({ client: sandbox.client })], + hostFunctions: { + sandbox: createSandboxHostFunctions({ client: sandbox.client }), + }, }); await sandbox.client.writeFsFile( @@ -70,8 +72,12 @@ describe("sandbox quickstart truth test", () => { const content = await vm.readFile(SANDBOX_FILE_PATH); expect(new TextDecoder().decode(content)).toBe(SANDBOX_FILE_CONTENT); - const hostFunctions = createSandboxHostFunctions({ client: sandbox.client }); - const runCommandResponse = (await hostFunctions.hostFunctions["run-command"].execute({ + const hostFunctions = createSandboxHostFunctions({ + client: sandbox.client, + }); + const runCommandResponse = (await hostFunctions.hostFunctions[ + "run-command" + ].execute({ command: "echo", args: ["hello from sandbox"], })) as { @@ -83,7 +89,9 @@ describe("sandbox quickstart truth test", () => { expect(runCommandResponse.stderr).toBe(""); expect(runCommandResponse.stdout.trim()).toBe("hello from sandbox"); - const createdProcess = (await hostFunctions.hostFunctions["create-process"].execute({ + const createdProcess = (await hostFunctions.hostFunctions[ + "create-process" + ].execute({ command: "sleep", args: ["60"], })) as { @@ -110,6 +118,8 @@ describe("sandbox quickstart truth test", () => { ), ).toBe(true); - await hostFunctions.hostFunctions["kill-process"].execute({ id: createdProcess.id }); + await hostFunctions.hostFunctions["kill-process"].execute({ + id: createdProcess.id, + }); }, 150_000); }); diff --git a/packages/core/tests/sidecar-host-function-dispatch.nightly.test.ts b/packages/core/tests/sidecar-host-function-dispatch.nightly.test.ts index f95367a855..4d243d2c86 100644 --- a/packages/core/tests/sidecar-host-function-dispatch.nightly.test.ts +++ b/packages/core/tests/sidecar-host-function-dispatch.nightly.test.ts @@ -4,20 +4,17 @@ import { z } from "zod"; import { AgentOs, hostFunction, hostFunctions } from "../src/index.js"; import { ALLOW_ALL_VM_PERMISSIONS } from "./helpers/permissions.js"; -const mathFunctions = hostFunctions({ - name: "math", - description: "Math utilities", - functions: { - add: hostFunction({ - description: "Add two numbers", - inputSchema: z.object({ +const mathFunctions = { + add: { + inputSchema: z + .object({ a: z.number(), b: z.number(), - }), - execute: ({ a, b }) => ({ sum: a + b }), - }), + }) + .describe("Add two numbers"), + execute: ({ a, b }) => ({ sum: a + b }), }, -}); +}; async function runCommand(vm: AgentOs, command: string, args: string[]) { const stdoutChunks: string[] = []; @@ -44,7 +41,7 @@ describe("native sidecar hostFunction dispatch", () => { beforeEach(async () => { vm = await AgentOs.create({ software: [common], - hostFunctions: [mathFunctions], + hostFunctions: { math: mathFunctions }, permissions: ALLOW_ALL_VM_PERMISSIONS, }); }, 20_000); @@ -62,8 +59,7 @@ describe("native sidecar hostFunction dispatch", () => { hostFunctions: [ { name: "math", - description: "Math utilities", - hostFunctions: ["add"], + functions: ["add"], }, ], }, diff --git a/packages/runtime-core/package.json b/packages/runtime-core/package.json index 7f8b0b40c8..8de18cecd4 100644 --- a/packages/runtime-core/package.json +++ b/packages/runtime-core/package.json @@ -21,6 +21,16 @@ "import": "./dist/generated-protocol.js", "default": "./dist/generated-protocol.js" }, + "./host-functions": { + "types": "./dist/host-functions.d.ts", + "import": "./dist/host-functions.js", + "default": "./dist/host-functions.js" + }, + "./host-functions-zod": { + "types": "./dist/host-functions-zod.d.ts", + "import": "./dist/host-functions-zod.js", + "default": "./dist/host-functions-zod.js" + }, "./vm-config": { "types": "./dist/vm-config.d.ts", "import": "./dist/vm-config.js", @@ -199,7 +209,8 @@ "dependencies": { "@rivet-dev/agentos-runtime-sidecar": "workspace:*", "@rivetkit/bare-ts": "^0.6.2", - "zod": "^4.1.11" + "zod": "^4.1.11", + "zod-to-json-schema": "^3.25.2" }, "devDependencies": { "@types/node": "^22.10.2", diff --git a/packages/core/src/host-functions-zod.ts b/packages/runtime-core/src/host-functions-zod.ts similarity index 96% rename from packages/core/src/host-functions-zod.ts rename to packages/runtime-core/src/host-functions-zod.ts index 56e1dffd7a..c4ce2c6865 100644 --- a/packages/core/src/host-functions-zod.ts +++ b/packages/runtime-core/src/host-functions-zod.ts @@ -94,6 +94,11 @@ function displayTypeName(typeName: string): string { } } +/** The `.describe()` text on a schema, if it has one. */ +export function schemaDescription(schema: ZodType): string | undefined { + return getDescription(schema); +} + function getDescription(schema: ZodType): string | undefined { const def = getSchemaDef(schema); if (typeof def.description === "string") { @@ -155,11 +160,17 @@ function validateSchema(schema: ZodType, path: string) { } if (UNSUPPORTED_TYPES.has(typeName)) { - throw new HostFunctionSchemaConversionError(path, displayTypeName(typeName)); + throw new HostFunctionSchemaConversionError( + path, + displayTypeName(typeName), + ); } if (typeName === "discriminatedunion") { - throw new HostFunctionSchemaConversionError(path, displayTypeName(typeName)); + throw new HostFunctionSchemaConversionError( + path, + displayTypeName(typeName), + ); } if (TRANSPARENT_WRAPPER_TYPES.has(typeName)) { diff --git a/packages/runtime-core/src/host-functions.ts b/packages/runtime-core/src/host-functions.ts new file mode 100644 index 0000000000..cd518e68fc --- /dev/null +++ b/packages/runtime-core/src/host-functions.ts @@ -0,0 +1,161 @@ +import type { ZodType, z } from "zod"; +import { schemaDescription } from "./host-functions-zod.js"; + +/** + * A single function that executes on the host. + * + * The function carries no name or description of its own: the key it is + * registered under names it, and `inputSchema.describe()` documents it for the + * agent. + */ +export interface HostFunction { + /** + * Zod schema for the input. Drives CLI flag generation and validation. + * `.describe()` on the schema becomes the function's description in `--help` + * and in the agent's system prompt; `.describe()` on a field documents that + * field's flag. + */ + inputSchema: SCHEMA; + /** + * Runs on the host when the agent invokes the function. Its input is the + * type `inputSchema` describes, inferred when the collection is written + * inline in the call that takes it. + */ + execute: (input: HostFunctionInput) => Promise | OUTPUT; + /** Examples included in auto-generated prompt docs. */ + examples?: HostFunctionExample>[]; + /** Timeout in ms. Default: 30000. */ + timeout?: number; +} + +/** + * The input type a function's `execute` receives. A concrete schema gives the + * type it describes; a collection built outside the call that takes it has no + * schema to infer from, so it falls back to `any` rather than making every + * handler annotate its parameter. + */ +export type HostFunctionInput = ZodType extends SCHEMA + ? any + : z.infer; + +export interface HostFunctionExample { + /** Human description of what this example does. */ + description: string; + /** The input args for the example. */ + input: INPUT; +} + +/** + * One collection of host functions, keyed by function name. Each key becomes a + * subcommand of the collection's CLI binary and a method on the collection's + * guest global. + */ +export type HostFunctionCollection< + T extends Record = Record, +> = { [name in keyof T]: HostFunction }; + +/** + * The input schemas behind a set of collections. Only used to give each + * `execute` its input type: `AgentOs.create()` infers this from the literal, so + * `execute` sees the type its own `inputSchema` describes without a wrapper + * call to hang the inference on. + */ +export type HostFunctionSchemas = Record>; + +/** + * Host function collections, keyed by collection name. Each key becomes the CLI + * binary `agentos-{name}` and a frozen guest global. + * + * ```ts + * hostFunctions: { + * store: { + * listOrders: { inputSchema, execute }, + * }, + * } + * ``` + */ +export type HostFunctionCollections< + T extends HostFunctionSchemas = HostFunctionSchemas, +> = { [collection in keyof T]: HostFunctionCollection }; + +/** + * A collection resolved to the names the VM uses. Keys arrive as JavaScript + * identifiers and are converted once here, so the rest of the client and the + * sidecar only ever see kebab-case command names. + */ +export interface ResolvedHostFunctions { + /** Kebab-case collection name. Becomes the CLI suffix: `agentos-{name}`. */ + name: string; + /** Functions keyed by kebab-case command name. */ + functions: Record; +} + +const HOST_FUNCTION_COMMAND_NAME_RE = /^[a-z0-9]+(?:-[a-z0-9]+)*$/; + +/** + * Convert a registration key to its command name. `listOrders` and + * `list-orders` both become `list-orders`, so the guest sees one spelling + * whichever the caller wrote. + */ +export function hostFunctionCommandName(key: string): string { + return key + .replace(/([a-z0-9])([A-Z])/g, "$1-$2") + .replace(/([A-Z]+)([A-Z][a-z])/g, "$1-$2") + .toLowerCase(); +} + +function toCommandName(kind: string, key: string): string { + const name = hostFunctionCommandName(key); + if (!HOST_FUNCTION_COMMAND_NAME_RE.test(name)) { + throw new Error( + `${kind} name "${key}" must be alphanumeric, written in camelCase or with single hyphen separators`, + ); + } + return name; +} + +/** The description the agent sees, taken from the input schema's `.describe()`. */ +export function hostFunctionDescription(definition: HostFunction): string { + return schemaDescription(definition.inputSchema) ?? ""; +} + +/** + * Resolve the caller's collections into the shape the client and sidecar use. + * Throws on a key that cannot become a command name, and on two keys that + * collide once converted. + */ +export function resolveHostFunctions( + collections: HostFunctionCollections, +): ResolvedHostFunctions[] { + const resolved: ResolvedHostFunctions[] = []; + const seenCollections = new Map(); + + for (const [collectionKey, collection] of Object.entries(collections)) { + const name = toCommandName("Host function collection", collectionKey); + const collidedCollection = seenCollections.get(name); + if (collidedCollection !== undefined) { + throw new Error( + `Host function collections "${collidedCollection}" and "${collectionKey}" both resolve to the command name "${name}"`, + ); + } + seenCollections.set(name, collectionKey); + + const functions: Record = {}; + const seenFunctions = new Map(); + for (const [functionKey, definition] of Object.entries(collection)) { + const functionName = toCommandName("Host function", functionKey); + const collidedFunction = seenFunctions.get(functionName); + if (collidedFunction !== undefined) { + throw new Error( + `Host functions "${collidedFunction}" and "${functionKey}" in collection "${collectionKey}" both resolve to the command name "${functionName}"`, + ); + } + seenFunctions.set(functionName, functionKey); + functions[functionName] = definition; + } + + resolved.push({ name, functions }); + } + + return resolved; +} diff --git a/packages/runtime-core/src/node-runtime-options-schema.ts b/packages/runtime-core/src/node-runtime-options-schema.ts index 1e4aaca658..7115e70c87 100644 --- a/packages/runtime-core/src/node-runtime-options-schema.ts +++ b/packages/runtime-core/src/node-runtime-options-schema.ts @@ -74,7 +74,10 @@ const patternRulePermissionsSchema = z }) .strict(); -const fsPermissionsSchema = z.union([permissionModeSchema, fsRulePermissionsSchema]); +const fsPermissionsSchema = z.union([ + permissionModeSchema, + fsRulePermissionsSchema, +]); const patternPermissionsSchema = z.union([ permissionModeSchema, patternRulePermissionsSchema, @@ -129,15 +132,13 @@ const hostFunctionExampleSchema = z const hostFunctionDefinitionSchema = z .object({ - description: z.string(), inputSchema: z.custom( (value: unknown) => typeof value === "object" && value !== null, - { message: "Expected JSON Schema object" }, + { message: "Expected Zod schema object" }, ), - timeoutMs: z.number().int().nonnegative().optional(), + timeout: z.number().int().nonnegative().optional(), examples: z.array(hostFunctionExampleSchema).optional(), - commandAliases: stringArray.optional(), - handler: z.custom<(input: unknown) => unknown | Promise>( + execute: z.custom<(input: unknown) => unknown | Promise>( (value: unknown) => typeof value === "function", { message: "Expected function" }, ), @@ -181,10 +182,10 @@ export const nodeRuntimeCreateOptionsSchema = z .optional(), mounts: z.array(hostDirectoryMountSchema).optional(), nodeModules: z.union([z.string(), nodeModulesMountSchema]).optional(), - hostFunctions: z.record(z.string(), hostFunctionDefinitionSchema).optional(), - loopbackExemptPorts: z - .array(z.number().int().min(0).max(65535)) + hostFunctions: z + .record(z.string(), hostFunctionDefinitionSchema) .optional(), + loopbackExemptPorts: z.array(z.number().int().min(0).max(65535)).optional(), jsRuntime: jsRuntimeSchema.optional(), }) .strict() as z.ZodType; diff --git a/packages/runtime-core/src/node-runtime.ts b/packages/runtime-core/src/node-runtime.ts index d058890245..5ac58e5224 100644 --- a/packages/runtime-core/src/node-runtime.ts +++ b/packages/runtime-core/src/node-runtime.ts @@ -21,28 +21,30 @@ import { existsSync } from "node:fs"; import path from "node:path"; import { fileURLToPath } from "node:url"; +import type { JsRuntimeConfig } from "./generated/JsRuntimeConfig.js"; +import type { VmUserConfig } from "./generated/VmUserConfig.js"; +import type { + HostFunctionCollections, + HostFunctionSchemas, +} from "./host-functions.js"; +import { parseNodeRuntimeCreateOptions } from "./node-runtime-options-schema.js"; +import type { SidecarProcess } from "./sidecar-process.js"; import type { ExecResult, - HostFunctionDefinition, Kernel, KernelBootTiming, Permissions, VirtualDirEntry, VirtualFileSystem, } from "./test-runtime.js"; -import type { JsRuntimeConfig } from "./generated/JsRuntimeConfig.js"; -import type { VmUserConfig } from "./generated/VmUserConfig.js"; -import type { SidecarProcess } from "./sidecar-process.js"; import { createKernel, createNodeRuntime, createWasmVmRuntime, NodeFileSystem, } from "./test-runtime.js"; -import { parseNodeRuntimeCreateOptions } from "./node-runtime-options-schema.js"; export type { - HostFunctionDefinition, HostFunctionExample, VirtualDirEntry, } from "./test-runtime.js"; @@ -119,7 +121,9 @@ export function resolveNodeRuntimeCommandsDir(explicit?: string): string { * Options that translate into sidecar VM JSON must also stay aligned with * `crates/vm-config/src/lib.rs::CreateVmConfig`. */ -export interface NodeRuntimeCreateOptions { +export interface NodeRuntimeCreateOptions< + HOST_FUNCTIONS extends HostFunctionSchemas = HostFunctionSchemas, +> { /** * Caller-owned filesystem used only by this low-level compatibility runtime. * AgentOS clients do not create a TypeScript filesystem implicitly; normal @@ -231,32 +235,35 @@ export interface NodeRuntimeCreateOptions { */ nodeModules?: string | NodeModulesMount; /** - * Host-side hostFunctions the guest can invoke as shell commands. Each entry is - * registered as a named guest command; when the guest runs it, the - * invocation round-trips back to the host and runs the hostFunction's `handler`, - * whose return value is delivered back to the guest. This is how you give - * sandboxed guest code controlled, named host capabilities (the kind an AI - * agent calls as tools) without granting it the underlying access directly. + * Host-side functions the guest can invoke as shell commands, as a record of + * collections. The keys name everything: the collection key becomes the guest + * command `agentos-{name}` and each function key becomes one of its + * subcommands. When the guest runs it the invocation round-trips back to the + * host, runs the function's `execute`, and its return value is delivered back + * to the guest. This is how you give sandboxed guest code controlled, named + * host capabilities (the kind an AI agent calls as tools) without granting it + * the underlying access directly. * - * The guest invokes a hostFunction by name with JSON input: + * This is the same shape `AgentOs.create()` takes. A function needs only a + * Zod `inputSchema` and an `execute` handler; `.describe()` on the schema is + * what the agent reads. * * ```ts * const rt = await NodeRuntime.create({ * hostFunctions: { - * add: { - * description: "Add two numbers", - * inputSchema: { - * type: "object", - * properties: { a: { type: "number" }, b: { type: "number" } }, - * required: ["a", "b"], + * math: { + * add: { + * inputSchema: z + * .object({ a: z.number(), b: z.number() }) + * .describe("Add two numbers"), + * execute: ({ a, b }) => ({ sum: a + b }), * }, - * handler: ({ a, b }: { a: number; b: number }) => ({ sum: a + b }), * }, * }, * }); * await rt.exec(` * import { execFileSync } from "node:child_process"; - * const out = execFileSync("add", ["add", "--json", JSON.stringify({ a: 2, b: 3 })]); + * const out = execFileSync("agentos-math", ["add", "--json", JSON.stringify({ a: 2, b: 3 })]); * console.log(out.toString()); * `); * ``` @@ -264,7 +271,7 @@ export interface NodeRuntimeCreateOptions { * The `hostFunction` permission scope is allowed by default; pass your own * `permissions.hostFunction` policy to gate individual host functions. */ - hostFunctions?: Record; + hostFunctions?: HostFunctionCollections; /** * Guest-bound ports that may accept non-loopback connections. By default a * guest server is reachable only over loopback inside the VM; listing a port @@ -567,10 +574,14 @@ export class NodeRuntime { * session, creates the VM with a bootstrapped root filesystem, mounts the * shell and Node runtimes, and waits for the VM to report ready. */ - static async create( - options: NodeRuntimeCreateOptions, + static async create( + callerOptions: NodeRuntimeCreateOptions, ): Promise { - options = parseNodeRuntimeCreateOptions(options); + // The generic exists only so each `execute` infers its input from its own + // `inputSchema`; past this point the concrete schemas carry no meaning. + const options: NodeRuntimeCreateOptions = parseNodeRuntimeCreateOptions( + callerOptions as NodeRuntimeCreateOptions, + ); const commandsDir = resolveNodeRuntimeCommandsDir(options.commandsDir); // Seed caller-provided files into the VM's in-memory filesystem before @@ -1088,7 +1099,7 @@ export class NodeRuntime { * functions are invocable unless the runtime's policy restricts them. */ async registerHostFunctions( - hostFunctions: Record, + hostFunctions: HostFunctionCollections, ): Promise { await this.kernel.registerHostFunctions(hostFunctions); } diff --git a/packages/runtime-core/src/test-runtime.ts b/packages/runtime-core/src/test-runtime.ts index 69b8306235..21254c5a83 100644 --- a/packages/runtime-core/src/test-runtime.ts +++ b/packages/runtime-core/src/test-runtime.ts @@ -322,14 +322,6 @@ export interface Permissions { hostFunction?: HostFunctionPermissions; } -/** A worked example shown alongside a registered host function. */ -export interface HostFunctionExample { - /** What this example demonstrates. */ - description: string; - /** Example input matching the host function's input schema. */ - input: unknown; -} - /** * A host-side function that guest code can invoke as a shell command. The guest * runs the host function by name and the invocation round-trips back to the host JS @@ -338,26 +330,19 @@ export interface HostFunctionExample { * sandboxed guest code controlled, named capabilities (the kind AI agents call * as tools). */ -export interface HostFunctionDefinition { - /** Human-readable description of what the host function does. */ - description: string; - /** JSON Schema describing the host function's input. */ - inputSchema: object; - /** Abort the invocation after this many milliseconds. */ - timeoutMs?: number; - /** Worked examples shown alongside the hostFunction. */ - examples?: HostFunctionExample[]; - /** - * Extra command names the guest can use to invoke this hostFunction, in addition - * to the key it is registered under. - */ - commandAliases?: string[]; - /** - * Host handler invoked when guest code runs the hostFunction. Receives the parsed - * input and returns a JSON-serializable result delivered back to the guest. - */ - handler: (input: unknown) => unknown | Promise; -} +import { + type HostFunctionCollections, + hostFunctionDescription, + resolveHostFunctions, +} from "./host-functions.js"; +import { zodToJsonSchema } from "./host-functions-zod.js"; + +export type { + HostFunction, + HostFunctionCollection, + HostFunctionCollections, + HostFunctionExample, +} from "./host-functions.js"; export interface ResourceBudgets { maxOutputBytes?: number; @@ -513,7 +498,7 @@ export interface Kernel extends KernelInterface { streamId?: string; maxBytes?: number; }): Promise; - registerHostFunctions(hostFunctions: Record): Promise; + registerHostFunctions(hostFunctions: HostFunctionCollections): Promise; getResourceSnapshot(): Promise<{ runningProcesses: number; exitedProcesses: number; @@ -561,6 +546,12 @@ export interface HostFunctionTree { export type HostFunctionHandler = (...args: unknown[]) => unknown; +/** The description the agent sees, taken from the input schema. */ +function jsonSchemaDescription(inputSchema: object): string { + const description = (inputSchema as { description?: unknown }).description; + return typeof description === "string" ? description : ""; +} + export interface ModuleAccessOptions { cwd?: string; } @@ -3085,7 +3076,7 @@ class NativeKernel implements Kernel { } async registerHostFunctions( - hostFunctions: Record, + hostFunctions: HostFunctionCollections, ): Promise { await this.ensureReady(); if (!this.client || !this.session || !this.vm) { @@ -3102,38 +3093,44 @@ class NativeKernel implements Kernel { this.hostFunctionRequestHandlerInstalled = true; } - for (const [name, hostFunction] of Object.entries(hostFunctions)) { - this.hostFunctionHandlers.set(name, hostFunction.handler); - const definition: SidecarRegisteredHostCallbackDefinition = { - description: hostFunction.description, - inputSchema: hostFunction.inputSchema, - ...(hostFunction.timeoutMs !== undefined - ? { timeoutMs: hostFunction.timeoutMs } - : {}), - ...(hostFunction.examples && hostFunction.examples.length > 0 - ? { - examples: hostFunction.examples.map((example) => ({ - description: example.description, - input: example.input, - })), - } - : {}), - }; - // Register each host function as its own single-callback host-function collection so the guest - // can invoke it directly by name (or by any caller-provided alias). The - // sidecar exposes the host-function collection name as a guest command; the single - // callback carries the host function's schema and gates the `hostFunction` - // permission. + for (const collection of resolveHostFunctions(hostFunctions)) { + const callbacks: Record = + {}; + for (const [functionName, definition] of Object.entries( + collection.functions, + )) { + this.hostFunctionHandlers.set( + `${collection.name}:${functionName}`, + definition.execute, + ); + callbacks[functionName] = { + description: hostFunctionDescription(definition), + inputSchema: zodToJsonSchema(definition.inputSchema) as object, + ...(definition.timeout !== undefined + ? { timeoutMs: definition.timeout } + : {}), + ...(definition.examples && definition.examples.length > 0 + ? { + examples: definition.examples.map((example) => ({ + description: example.description, + input: example.input, + })), + } + : {}), + }; + } + // One registration per collection, exposed to the guest as the command + // `agentos-` with each function as a subcommand. This is the + // same shape `AgentOs.create()` registers. + const command = `agentos-${collection.name}`; await this.client.registerHostCallbacks(this.session, this.vm, { - name, - description: hostFunction.description, - commandAliases: [name, ...(hostFunction.commandAliases ?? [])], - callbacks: { [name]: definition }, + name: collection.name, + description: "", + commandAliases: [command], + registryCommandAliases: ["agentos"], + callbacks, }); - this.commands.set(name, "wasmvm"); - for (const alias of hostFunction.commandAliases ?? []) { - this.commands.set(alias, "wasmvm"); - } + this.commands.set(command, "wasmvm"); } } @@ -3146,17 +3143,10 @@ class NativeKernel implements Kernel { `unsupported sidecar request for host functions: ${payload.type}`, ); } - // Callback keys arrive as `:` for collection invocations - // and as the bare command name otherwise. The collection and hostFunction name are - // the same here, so the registered hostFunction name is the segment after the last - // colon (or the whole key when no colon is present). + // Callback keys arrive as `:`, the key each handler + // was registered under. const callbackKey = payload.callback_key; - const hostFunctionName = callbackKey.includes(":") - ? callbackKey.slice(callbackKey.lastIndexOf(":") + 1) - : callbackKey; - const handler = - this.hostFunctionHandlers.get(hostFunctionName) ?? - this.hostFunctionHandlers.get(callbackKey); + const handler = this.hostFunctionHandlers.get(callbackKey); if (!handler) { return { type: "host_callback_result", diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 18017687b8..41f5e6fdd3 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -3296,6 +3296,9 @@ importers: zod: specifier: ^4.1.11 version: 4.3.6 + zod-to-json-schema: + specifier: ^3.25.2 + version: 3.25.2(zod@4.3.6) devDependencies: '@types/node': specifier: ^22.10.2 diff --git a/secure-exec/docs/content/docs/api-reference.mdx b/secure-exec/docs/content/docs/api-reference.mdx index 3629ce9077..faaa894860 100644 --- a/secure-exec/docs/content/docs/api-reference.mdx +++ b/secure-exec/docs/content/docs/api-reference.mdx @@ -14,7 +14,6 @@ skill: true | `createVm` | `createVm(vmOptions?)` → `Vm` | Create a VM that lives until you dispose it | | `init` | `init()` | Start the shared sidecar ahead of time | | `shutdown` | `shutdown()` | Stop the sidecar and every VM in it | -| `hostFunction`, `hostFunctions` | | Define [host functions](/secure-exec/docs/host-functions) | | `hostDirMount` | `hostDirMount(path, hostPath, { readOnly? })` | Mount a host directory | | `nodeModulesMount` | `nodeModulesMount(hostPath, { readOnly? })` | Mount a host `node_modules` | diff --git a/secure-exec/docs/content/docs/host-functions.mdx b/secure-exec/docs/content/docs/host-functions.mdx index b5bc91565b..789cdc1254 100644 --- a/secure-exec/docs/content/docs/host-functions.mdx +++ b/secure-exec/docs/content/docs/host-functions.mdx @@ -12,23 +12,24 @@ writes one program that chains your tools. ## Define host functions -Host functions are grouped into named collections with `hostFunctions`. Each -function -has a Zod input schema, so arguments from the guest are validated before your -code runs. +Host functions are a record of collections, keyed by name. Each function needs +only a Zod `inputSchema` and an `execute` handler; the keys name the collection +and the function, and `.describe()` on the schema documents it for the model. +Arguments from the guest are validated before your code runs. + +`hostFunctions` is a VM option. Pass it on a call, or on `createVm`. ## Call them from guest code -`hostFunctions` is a VM option. Pass it on a call, or on `createVm`. Inside the -VM, -each collection is a global object and each function is async. +Inside the VM each collection is a global object and each function is async. -- Names become camelCase identifiers: a collection named `order-store` with a - function `list-orders` is `orderStore.listOrders(input)`. +- Keys round-trip: `orderStore.listOrders` in your options is + `orderStore.listOrders(input)` in the guest and `agentos-order-store + list-orders` on the CLI. - Each function takes one input object and resolves to what your `execute` returned. A schema violation, a thrown error, or a timeout rejects the promise. - The globals are frozen. A collection whose name is already a global, such as diff --git a/secure-exec/examples/host-functions/src/index.ts b/secure-exec/examples/host-functions/src/index.ts index badb204b06..48aaab934c 100644 --- a/secure-exec/examples/host-functions/src/index.ts +++ b/secure-exec/examples/host-functions/src/index.ts @@ -1,25 +1,6 @@ -import { evaluate, hostFunction, hostFunctions } from "secure-exec"; +import { evaluate } from "secure-exec"; import { z } from "zod"; -// docs:start define -// Host functions run in your process, with your credentials. The guest only -// sees their inputs and outputs. -const orders = hostFunctions({ - name: "orders", - description: "Look up customer orders.", - functions: { - list: hostFunction({ - description: "List a customer's orders.", - inputSchema: z.object({ customer: z.string() }), - execute: ({ customer }) => [ - { customer, amount: 40 }, - { customer, amount: 2 }, - ], - }), - }, -}); -// docs:end define - // docs:start call // Inside the VM each collection is a global, and each function is async. This // is the code a model would write. @@ -27,12 +8,29 @@ const generated = `(async () => { const list = await orders.list({ customer: inputs.customer }); return list.reduce((sum, order) => sum + order.amount, 0); })()`; +// docs:end call +// docs:start define +// Host functions run in your process, with your credentials. The guest only +// sees their inputs and outputs. The keys name the collection and the function, +// and `execute` receives the input its own schema describes. const total = await evaluate(generated, { - hostFunctions: [orders], + hostFunctions: { + orders: { + list: { + inputSchema: z + .object({ customer: z.string() }) + .describe("List a customer's orders."), + execute: ({ customer }) => [ + { customer, amount: 40 }, + { customer, amount: 2 }, + ], + }, + }, + }, inputs: { customer: "customer_123" }, timeoutMs: 5_000, output: { capture: "stderr" }, }); console.log(total.outcome === "succeeded" ? total.value : total.stderr); // 42 -// docs:end call +// docs:end define diff --git a/secure-exec/src/index.ts b/secure-exec/src/index.ts index a5c1dd9edd..3255ae5112 100644 --- a/secure-exec/src/index.ts +++ b/secure-exec/src/index.ts @@ -1,6 +1,7 @@ import type { CodeEvaluationResult, CodeExecutionResult, + HostFunctionSchemas, JavaScriptEvaluationOptions, JavaScriptExecutionOptions, JsonValue, @@ -14,7 +15,8 @@ export type { ExecutionErrorData, ExecutionOutputOptions, HostFunction, - HostFunctions, + HostFunctionCollection, + HostFunctionCollections, HttpRequest, HttpResponse, JsonValue, @@ -27,8 +29,6 @@ export type { SidecarRejectionDetail, } from "@rivet-dev/agentos-core"; export { - hostFunction, - hostFunctions, createHostDirBackend, hostDirMount, KernelError, @@ -51,14 +51,22 @@ export { // that is disposed when the call finishes. For anything that should persist, // create a VM with `createVm()` and call the same methods on it. -export type ExecuteOptions = OneShot; -export type EvaluateOptions = OneShot; -export type ExecuteFileOptions = OneShot; +export type ExecuteOptions< + HOST_FUNCTIONS extends HostFunctionSchemas = HostFunctionSchemas, +> = OneShot; +export type EvaluateOptions< + HOST_FUNCTIONS extends HostFunctionSchemas = HostFunctionSchemas, +> = OneShot; +export type ExecuteFileOptions< + HOST_FUNCTIONS extends HostFunctionSchemas = HostFunctionSchemas, +> = OneShot; /** Run JavaScript for its side effects and captured output. */ -export function execute( +export function execute< + HOST_FUNCTIONS extends HostFunctionSchemas = HostFunctionSchemas, +>( source: string, - options?: ExecuteOptions, + options?: ExecuteOptions, ): Promise { return run(options, (vm, operationOptions) => vm.javascript.execute(source, operationOptions), @@ -66,9 +74,12 @@ export function execute( } /** Evaluate one JavaScript expression and return its JSON value. */ -export function evaluate( +export function evaluate< + T = JsonValue, + HOST_FUNCTIONS extends HostFunctionSchemas = HostFunctionSchemas, +>( source: string, - options?: EvaluateOptions, + options?: EvaluateOptions, ): Promise> { return run(options, (vm, operationOptions) => vm.javascript.evaluate(source, operationOptions), @@ -76,9 +87,11 @@ export function evaluate( } /** Run a JavaScript file from a mount, by its path inside the VM. */ -export function executeFile( +export function executeFile< + HOST_FUNCTIONS extends HostFunctionSchemas = HostFunctionSchemas, +>( path: string, - options?: ExecuteFileOptions, + options?: ExecuteFileOptions, ): Promise { return run(options, (vm, operationOptions) => vm.javascript.executeFile(path, operationOptions), diff --git a/secure-exec/src/runtime.ts b/secure-exec/src/runtime.ts index 45a38fc432..ea348227dc 100644 --- a/secure-exec/src/runtime.ts +++ b/secure-exec/src/runtime.ts @@ -1,8 +1,19 @@ -import { AgentOs, type AgentOsOptions } from "@rivet-dev/agentos-core"; +import { + AgentOs, + type AgentOsOptions, + type HostFunctionSchemas, +} from "@rivet-dev/agentos-core"; import { type Context, createContext } from "./context.js"; -/** Options for the VM that runs the code: permissions, limits, mounts, and so on. */ -export type VmOptions = AgentOsOptions; +/** + * Options for the VM that runs the code: permissions, limits, mounts, and so on. + * + * The type parameter exists only so each host function's `execute` infers its + * input from its own `inputSchema`; callers never write it. + */ +export type VmOptions< + HOST_FUNCTIONS extends HostFunctionSchemas = HostFunctionSchemas, +> = AgentOsOptions; // Keep in sync with `AgentOsOptions`. The record type makes a missing or // unknown key a compile error. @@ -45,7 +56,9 @@ export interface Vm extends AsyncDisposable { dispose(): Promise; } -export async function createVm(options: VmOptions = {}): Promise { +export async function createVm( + options: VmOptions = {}, +): Promise { const vm = await AgentOs.create(options); const { npm, ...javascript } = vm.javascript; const dispose = () => vm.dispose(); @@ -63,7 +76,10 @@ export async function createVm(options: VmOptions = {}): Promise { } /** Options for a one-shot call: the operation's own options plus VM options. */ -export type OneShot = Omit & VmOptions; +export type OneShot< + O, + HOST_FUNCTIONS extends HostFunctionSchemas = HostFunctionSchemas, +> = Omit & VmOptions; /** Start the shared sidecar process now so the first operation does not pay for it. */ export async function init(): Promise { @@ -83,8 +99,12 @@ export async function shutdown(): Promise { } /** Run `operation` in a fresh VM built from the VM options, then dispose it. */ -export async function run( - options: (O & VmOptions) | undefined, +export async function run< + O extends object, + R, + HOST_FUNCTIONS extends HostFunctionSchemas = HostFunctionSchemas, +>( + options: (O & VmOptions) | undefined, operation: (vm: Vm, options: O) => Promise, ): Promise { const vmOptions: Record = {}; diff --git a/secure-exec/src/typescript.ts b/secure-exec/src/typescript.ts index 9191022ced..b05e02aa87 100644 --- a/secure-exec/src/typescript.ts +++ b/secure-exec/src/typescript.ts @@ -1,6 +1,7 @@ import type { CodeEvaluationResult, CodeExecutionResult, + HostFunctionSchemas, JsonValue, TypeScriptCheckOptions, TypeScriptCheckResult, @@ -18,15 +19,25 @@ export type { // One-shot conveniences, like the main entry point. On a VM from `createVm()`, // use `vm.typescript`. -export type ExecuteOptions = OneShot; -export type EvaluateOptions = OneShot; -export type ExecuteFileOptions = OneShot; -export type CheckOptions = OneShot; +export type ExecuteOptions< + HOST_FUNCTIONS extends HostFunctionSchemas = HostFunctionSchemas, +> = OneShot; +export type EvaluateOptions< + HOST_FUNCTIONS extends HostFunctionSchemas = HostFunctionSchemas, +> = OneShot; +export type ExecuteFileOptions< + HOST_FUNCTIONS extends HostFunctionSchemas = HostFunctionSchemas, +> = OneShot; +export type CheckOptions< + HOST_FUNCTIONS extends HostFunctionSchemas = HostFunctionSchemas, +> = OneShot; /** Run TypeScript for its side effects and captured output. Types are stripped, not checked. */ -export function execute( +export function execute< + HOST_FUNCTIONS extends HostFunctionSchemas = HostFunctionSchemas, +>( source: string, - options?: ExecuteOptions, + options?: ExecuteOptions, ): Promise { return run(options, (vm, operationOptions) => vm.typescript.execute(source, operationOptions), @@ -34,9 +45,12 @@ export function execute( } /** Evaluate one TypeScript expression and return its JSON value. Types are stripped, not checked. */ -export function evaluate( +export function evaluate< + T = JsonValue, + HOST_FUNCTIONS extends HostFunctionSchemas = HostFunctionSchemas, +>( source: string, - options?: EvaluateOptions, + options?: EvaluateOptions, ): Promise> { return run(options, (vm, operationOptions) => vm.typescript.evaluate(source, operationOptions), @@ -44,9 +58,11 @@ export function evaluate( } /** Run a TypeScript file from a mount, by its path inside the VM. Types are stripped, not checked. */ -export function executeFile( +export function executeFile< + HOST_FUNCTIONS extends HostFunctionSchemas = HostFunctionSchemas, +>( path: string, - options?: ExecuteFileOptions, + options?: ExecuteFileOptions, ): Promise { return run(options, (vm, operationOptions) => vm.typescript.executeFile(path, operationOptions), @@ -54,9 +70,11 @@ export function executeFile( } /** Type-check TypeScript without running it. */ -export function check( +export function check< + HOST_FUNCTIONS extends HostFunctionSchemas = HostFunctionSchemas, +>( source: string, - options?: CheckOptions, + options?: CheckOptions, ): Promise { return run(options, (vm, operationOptions) => vm.typescript.check(source, operationOptions), diff --git a/secure-exec/tests/secure-exec.test.ts b/secure-exec/tests/secure-exec.test.ts index 4de47f6573..394d44491b 100644 --- a/secure-exec/tests/secure-exec.test.ts +++ b/secure-exec/tests/secure-exec.test.ts @@ -6,8 +6,6 @@ import { AgentOs } from "@rivet-dev/agentos-core"; import { afterAll, describe, expect, test } from "vitest"; import { z } from "zod"; import { - hostFunction, - hostFunctions, createVm, evaluate, execute, @@ -160,20 +158,17 @@ describe("one-shot calls", () => { }); test("call host functions as guest globals", async () => { - const math = hostFunctions({ - name: "math", - description: "Arithmetic on the host.", - functions: { - add: hostFunction({ - description: "Add two numbers.", - inputSchema: z.object({ a: z.number(), b: z.number() }), - execute: ({ a, b }) => a + b, - }), + const math = { + add: { + inputSchema: z + .object({ a: z.number(), b: z.number() }) + .describe("Add two numbers."), + execute: ({ a, b }) => a + b, }, - }); + }; // Each collection is a guest global, and each function an async function. const result = await evaluate("math.add({ a: 40, b: 2 })", { - hostFunctions: [math], + hostFunctions: { math }, ...bare, }); expect(result).toMatchObject({ outcome: "succeeded", value: 42 });