diff --git a/crates/eryx-runtime/prebuilt/liberyx_runtime.so.zst b/crates/eryx-runtime/prebuilt/liberyx_runtime.so.zst index 088828a3..80542d59 100644 Binary files a/crates/eryx-runtime/prebuilt/liberyx_runtime.so.zst and b/crates/eryx-runtime/prebuilt/liberyx_runtime.so.zst differ diff --git a/crates/eryx-runtime/wit/runtime.wit b/crates/eryx-runtime/wit/runtime.wit index 518d3e24..e00a451b 100644 --- a/crates/eryx-runtime/wit/runtime.wit +++ b/crates/eryx-runtime/wit/runtime.wit @@ -73,7 +73,10 @@ world sandbox { record execution-options { /// Whether the guest should install Python tracing. python-tracing: bool, - /// Whether a fresh instance may reuse its pre-initialized empty callbacks. + /// Historical hint that a fresh instance may keep its pre-initialized + /// empty callbacks. The guest now compares the host's declarations with + /// what is installed and ignores this field; it is kept so existing + /// hosts stay compatible. reuse-empty-callbacks: bool, } diff --git a/crates/eryx-wasm-runtime/src/lib.rs b/crates/eryx-wasm-runtime/src/lib.rs index 87d4d9f7..a69a9a9b 100644 --- a/crates/eryx-wasm-runtime/src/lib.rs +++ b/crates/eryx-wasm-runtime/src/lib.rs @@ -719,13 +719,11 @@ fn call_invoke(name: &str, args_json: &str) -> Result { #[derive(Clone, Copy)] struct ExecutionOptions { python_tracing: bool, - reuse_empty_callbacks: bool, } impl ExecutionOptions { const CONSERVATIVE: Self = Self { python_tracing: true, - reuse_empty_callbacks: false, }; } @@ -745,12 +743,11 @@ fn call_get_execution_options(wit: Wit) -> ExecutionOptions { match cx.stack.pop() { Some(Value::Record(fields)) if fields.len() == 2 => { let mut fields = fields.into_iter(); + // The second field, `reuse-empty-callbacks`, is kept in the WIT + // for compatibility; the guest now tracks installed callbacks itself. match (fields.next(), fields.next()) { - (Some(Value::Bool(python_tracing)), Some(Value::Bool(reuse_empty_callbacks))) => { - ExecutionOptions { - python_tracing, - reuse_empty_callbacks, - } + (Some(Value::Bool(python_tracing)), Some(Value::Bool(_))) => { + ExecutionOptions { python_tracing } } _ => { eprintln!("call_get_execution_options: unexpected record fields"); @@ -1476,18 +1473,12 @@ pub fn do_tls_close(handle: u32) { /// - Session reuse: callbacks may change between executions /// - Error recovery: a previous failed setup won't prevent future attempts /// -/// Fresh instances can reuse empty callback infrastructure captured during -/// initialization. Persistent sessions must perform setup on every execution -/// because their callback state may have changed since the previous request. -fn initialize_callbacks(wit: Wit, reuse_empty_callbacks: bool) { +/// The host's callback set can change between executions of a persistent +/// session, so it is fetched every time; `setup_callbacks` only reinstalls +/// when it differs from what is already installed. +fn initialize_callbacks(wit: Wit) { let callbacks = call_list_callbacks(wit); - let skip_setup = - reuse_empty_callbacks && callbacks.is_empty() && python::callbacks_pre_initialized(); - if skip_setup { - return; - } - if let Err(e) = python::setup_callbacks(&callbacks) { eprintln!( "ERROR: Failed to set up callbacks ({} declarations): {e}", @@ -1543,7 +1534,7 @@ fn handle_export(wit: Wit, func_index: usize, cx: &mut EryxCall) -> HandleExport // Set up callbacks from the host's current per-request state. // This runs on every execute to stay in sync with the host. let options = call_get_execution_options(wit); - initialize_callbacks(wit, options.reuse_empty_callbacks); + initialize_callbacks(wit); // Execute Python with Wit handle available for callbacks let result = with_wit(wit, || { diff --git a/crates/eryx-wasm-runtime/src/python.rs b/crates/eryx-wasm-runtime/src/python.rs index 2f59a853..548a1d85 100644 --- a/crates/eryx-wasm-runtime/src/python.rs +++ b/crates/eryx-wasm-runtime/src/python.rs @@ -9,6 +9,7 @@ #![allow(missing_debug_implementations)] use std::ffi::c_char; +use std::sync::Mutex; use std::sync::atomic::{AtomicBool, Ordering}; // Re-export pyo3::ffi types and functions available in the stable ABI @@ -1313,13 +1314,20 @@ sys.modules['_eryx_async'] = _eryx_async /// Track whether we've initialized Python. static PYTHON_INITIALIZED: AtomicBool = AtomicBool::new(false); -/// True when initialization installed the empty-callback infrastructure. -/// A failed installation degrades to per-execution callback setup. -static CALLBACKS_PRE_INITIALIZED: AtomicBool = AtomicBool::new(false); - -/// Whether the pre-init snapshot contains the empty-callback infrastructure. -pub fn callbacks_pre_initialized() -> bool { - CALLBACKS_PRE_INITIALIZED.load(Ordering::SeqCst) +/// The callback declarations whose wrappers are currently installed, in the +/// JSON form handed to the setup script. +/// +/// `None` means the installation is unknown or may have been disturbed, so +/// the next [`setup_callbacks`] reinstalls unconditionally. Pre-initialization +/// installs the empty set, and that value is carried in the snapshot, so a +/// fresh instance with no callbacks never runs the setup script. +static INSTALLED_CALLBACKS: Mutex> = Mutex::new(None); + +/// Forget which callbacks are installed so the next setup reinstalls them. +fn invalidate_installed_callbacks() { + if let Ok(mut installed) = INSTALLED_CALLBACKS.lock() { + *installed = None; + } } /// Initialize Python interpreter. @@ -2677,15 +2685,10 @@ pub fn initialize_python() { PyErr_Clear(); } - // Install empty callback infrastructure once so fresh instances with no - // callbacks do not have to recreate it for every execution. Non-empty - // callback lists still reinstall their request-specific definitions. - let callbacks_ok = setup_callbacks(&[]).is_ok(); - CALLBACKS_PRE_INITIALIZED.store(callbacks_ok, Ordering::SeqCst); - if !callbacks_ok { - eprintln!( - "WARNING: pre-init setup_callbacks([]) failed; empty-callback fast path disabled" - ); + // Install the empty callback infrastructure once so fresh instances + // with no callbacks do not have to recreate it for every execution. + if let Err(e) = setup_callbacks(&[]) { + eprintln!("WARNING: pre-init setup_callbacks([]) failed: {e}"); } // Note: We do NOT call reset_wasi_state() here! @@ -3261,9 +3264,11 @@ del _eryx_restore_bytes, _eryx_restored_dict, _eryx_dill, _eryx_types, _eryx_reb let _ = PyRun_SimpleString(c"del _eryx_restore_bytes".as_ptr()); return Err(format!("Failed to restore state: {err}")); } - - Ok(()) } + + // The restored globals may have replaced callback wrappers. + invalidate_installed_callbacks(); + Ok(()) } /// Clear all user-defined state from `_eryx_user_globals`. @@ -3319,6 +3324,9 @@ del _eryx_keep, _eryx_should_keep, _eryx_to_delete, _k PyErr_Clear(); } } + + // The keep-list above is a heuristic; make the next execution reinstall. + invalidate_installed_callbacks(); } // ============================================================================= @@ -3340,15 +3348,29 @@ pub struct CallbackInfo { /// 2. A `list_callbacks()` function for introspection /// 3. Direct wrapper functions for each callback (e.g., `sleep(ms=100)`) /// 4. Namespace objects for dotted callbacks (e.g., `http.get(url="...")`) +/// +/// Installing is idempotent: the declarations are compared with the set that +/// is already installed and the setup script only runs when they differ. The +/// script compiles a few hundred lines of Python, which is several times the +/// cost of a short execution, so this is what keeps per-execution overhead +/// low for sessions and for sandboxes with callbacks. pub fn setup_callbacks(callbacks: &[CallbackInfo]) -> Result<(), String> { if !is_python_initialized() { return Err("Python not initialized".to_string()); } - unsafe { - // Serialize callbacks to JSON for Python to parse - let callbacks_json = serde_json_mini_serialize_callbacks(callbacks); + // Serialize callbacks to JSON for Python to parse + let callbacks_json = serde_json_mini_serialize_callbacks(callbacks); + + let already_installed = INSTALLED_CALLBACKS + .lock() + .is_ok_and(|installed| installed.as_deref() == Some(callbacks_json.as_str())); + if already_installed { + return Ok(()); + } + invalidate_installed_callbacks(); + unsafe { // Inject the callback setup code let setup_code = format!( r#" @@ -3563,9 +3585,12 @@ except NameError: let err = get_last_error_message(); return Err(format!("Failed to set up callbacks: {err}")); } + } - Ok(()) + if let Ok(mut installed) = INSTALLED_CALLBACKS.lock() { + *installed = Some(callbacks_json); } + Ok(()) } /// Simple JSON serialization for callbacks (avoiding serde dependency in WASM) diff --git a/crates/eryx/src/wasm.rs b/crates/eryx/src/wasm.rs index 4818c5c8..e5b36150 100644 --- a/crates/eryx/src/wasm.rs +++ b/crates/eryx/src/wasm.rs @@ -618,6 +618,10 @@ pub struct ExecutorState { pub(crate) suspended: Option, /// Whether this execution uses a fresh instance whose initialized empty /// callback state can be reused safely. + /// + /// Reported to the guest as `reuse-empty-callbacks`. Current guests ignore + /// it and reinstall callbacks only when the declarations change; older + /// guests use it to skip setup on fresh, callback-free instances. pub(crate) reuse_empty_callbacks: bool, } diff --git a/crates/eryx/tests/callback_setup_cache.rs b/crates/eryx/tests/callback_setup_cache.rs new file mode 100644 index 00000000..95e4acf8 --- /dev/null +++ b/crates/eryx/tests/callback_setup_cache.rs @@ -0,0 +1,235 @@ +//! Coverage for the guest's idempotent callback installation. +//! +//! The guest only runs its callback setup script when the host's callback +//! declarations differ from the ones already installed. These tests check +//! that the wrappers keep working across executions that reuse the same set, +//! that changes to the set are picked up, and that the installation is +//! refreshed after the operations that can disturb it. +#![cfg(feature = "embedded")] +#![allow(clippy::expect_used, clippy::unwrap_used)] + +use std::collections::HashMap; +use std::future::Future; +use std::pin::Pin; +use std::sync::{Arc, OnceLock}; + +use eryx::callback_handler::run_callback_handler; +use eryx::{ + Callback, CallbackError, JsonSchema, PythonExecutor, ResourceLimits, Sandbox, SessionExecutor, + TypedCallback, +}; +use serde::Deserialize; +use serde_json::{Value, json}; + +static EXECUTOR: OnceLock> = OnceLock::new(); + +fn executor() -> Arc { + EXECUTOR + .get_or_init(|| { + let resources = eryx::embedded::EmbeddedResources::get().unwrap(); + #[allow(unsafe_code)] + Arc::new( + unsafe { PythonExecutor::from_precompiled_file(resources.runtime()) } + .unwrap() + .with_python_stdlib(resources.stdlib()), + ) + }) + .clone() +} + +#[derive(Deserialize, JsonSchema)] +struct EchoArgs { + /// Data to echo back + data: Value, +} + +struct EchoCallback; + +impl TypedCallback for EchoCallback { + type Args = EchoArgs; + + fn name(&self) -> &str { + "echo" + } + + fn description(&self) -> &str { + "Echoes the input data back" + } + + fn invoke_typed( + &self, + args: EchoArgs, + ) -> Pin> + Send + '_>> { + Box::pin(async move { Ok(args.data) }) + } +} + +struct PingCallback; + +impl TypedCallback for PingCallback { + type Args = (); + + fn name(&self) -> &str { + "ping" + } + + fn description(&self) -> &str { + "Returns pong" + } + + fn invoke_typed( + &self, + _args: (), + ) -> Pin> + Send + '_>> { + Box::pin(async move { Ok(json!("pong")) }) + } +} + +fn echo_only() -> Vec> { + vec![Arc::new(EchoCallback)] +} + +fn echo_and_ping() -> Vec> { + vec![Arc::new(EchoCallback), Arc::new(PingCallback)] +} + +/// Run `code` on `session` with `callbacks` available, serving their +/// invocations with the same handler `Sandbox` uses. +async fn run(session: &mut SessionExecutor, callbacks: &[Arc], code: &str) -> String { + let (callback_tx, callback_rx) = tokio::sync::mpsc::channel(4); + let callbacks_map: HashMap> = callbacks + .iter() + .map(|cb| (cb.name().to_string(), Arc::clone(cb))) + .collect(); + let handler = tokio::spawn(run_callback_handler( + callback_rx, + Arc::new(callbacks_map), + ResourceLimits::unlimited(), + Arc::new(HashMap::new()), + )); + + let output = session + .execute(code) + .with_callbacks(callbacks, callback_tx) + .run() + .await + .expect("execution failed"); + handler.await.unwrap(); + output.stdout +} + +#[tokio::test] +async fn wrappers_keep_working_when_the_callback_set_is_unchanged() { + let callbacks = echo_only(); + let mut session = SessionExecutor::new(executor(), &callbacks).await.unwrap(); + + for i in 0..3 { + let code = format!("print(await echo(data={i}))"); + assert_eq!(run(&mut session, &callbacks, &code).await, i.to_string()); + } +} + +#[tokio::test] +async fn a_changed_callback_set_is_installed() { + let mut session = SessionExecutor::new(executor(), &echo_only()) + .await + .unwrap(); + + assert_eq!( + run( + &mut session, + &echo_only(), + "print(sorted(c['name'] for c in list_callbacks()))" + ) + .await, + "['echo']" + ); + + // Adding a callback exposes its wrapper and updates introspection. + assert_eq!( + run( + &mut session, + &echo_and_ping(), + "print(sorted(c['name'] for c in list_callbacks()), await ping())" + ) + .await, + "['echo', 'ping'] pong" + ); + + // Going back to the smaller set is picked up too. + assert_eq!( + run( + &mut session, + &echo_only(), + "print(sorted(c['name'] for c in list_callbacks()))" + ) + .await, + "['echo']" + ); +} + +#[tokio::test] +async fn callbacks_survive_clear_state() { + let callbacks = echo_only(); + let mut session = SessionExecutor::new(executor(), &callbacks).await.unwrap(); + + assert_eq!( + run( + &mut session, + &callbacks, + "x = 1\nprint(await echo(data='a'))" + ) + .await, + "a" + ); + session.clear_state().await.unwrap(); + assert_eq!( + run( + &mut session, + &callbacks, + "print('x' in globals(), await echo(data='b'))" + ) + .await, + "False b" + ); +} + +#[tokio::test] +async fn callbacks_survive_snapshot_and_restore() { + let callbacks = echo_only(); + let mut session = SessionExecutor::new(executor(), &callbacks).await.unwrap(); + + assert_eq!( + run( + &mut session, + &callbacks, + "x = 41\nprint(await echo(data=x))" + ) + .await, + "41" + ); + let snapshot = session.snapshot_state().await.unwrap(); + + let mut restored = SessionExecutor::new(executor(), &callbacks).await.unwrap(); + restored.restore_state(&snapshot).await.unwrap(); + assert_eq!( + run(&mut restored, &callbacks, "print(await echo(data=x + 1))").await, + "42" + ); +} + +#[tokio::test] +async fn stateless_sandboxes_with_callbacks_work_repeatedly() { + let sandbox = Sandbox::embedded() + .with_callback(EchoCallback) + .build() + .unwrap(); + + for i in 0..3 { + let output = sandbox + .execute(&format!("print(await echo(data={i}))")) + .await + .unwrap(); + assert_eq!(output.stdout, i.to_string()); + } +} diff --git a/mise.toml b/mise.toml index 5b341218..05d45d02 100644 --- a/mise.toml +++ b/mise.toml @@ -1,6 +1,6 @@ [tools] # Profile "default" includes rustc, rust-std, cargo, rust-docs, rustfmt, and clippy -rust = { version = "1.98", profile = "default" } +rust = { version = "1.98.1", profile = "default" } [tools."cargo:cargo-nextest"] version = "0.9.143"