perf(preinit): bake callback declarations into the snapshot

book/src/guide/packages.md

@@ -114,6 +114,29 @@ Some packages with native extensions have WASI-compiled versions available:
 
 Packages with native C extensions (like `numpy`, `pandas`) require WASI-compiled wheels. Check if WASI builds are available for your specific package.
 
+## Baking Callbacks into the Factory
+
+Registering callbacks installs a Python wrapper for each of them inside the guest. A fresh sandbox starts from the factory's snapshot, and the snapshot only knows the callbacks it was created with — none, by default — so every `create_sandbox()` re-runs that installation on its first execution (a few milliseconds for a handful of callbacks, more than creating the sandbox itself).
+
+If every sandbox from a factory registers the same callbacks, pass them to the factory instead. Their declarations (name, description, parameter schema) are baked into the snapshot, and `create_sandbox()` / `create_session()` register them unless given callbacks of their own:
+
+```python
+import eryx
+
+def get_time():
+    import time
+    return {"timestamp": time.time()}
+
+factory = eryx.SandboxFactory(
+    callbacks=[{"name": "get_time", "fn": get_time, "description": "Returns current time"}],
+)
+
+sandbox = factory.create_sandbox()  # get_time() is ready; no per-sandbox setup
+result = sandbox.execute("print((await get_time())['timestamp'] > 0)")
+```
+
+A sandbox that is given a different set of callbacks still works; it installs them itself as before. Setup code cannot invoke the callbacks, since nothing answers them while the snapshot is being built. A saved factory keeps the baked declarations but not the Python functions, so pass the same `callbacks=` to `SandboxFactory.load()`.
+
 ## Package Dependencies
 
 When installing packages, you must include all dependencies:

book/src/guide/precompile.md

@@ -148,6 +148,25 @@ eryx-precompile compile runtime.wasm -o jinja2.cwasm \
 
 Each sandbox created from the resulting artifact starts with `env` already defined. Per-request code only needs to handle the request-specific work (deserializing data, compiling the template, rendering).
 
+### Callbacks
+
+Registering callbacks on a sandbox installs a Python wrapper for each of them inside the guest. A persistent session pays for that once, but a fresh sandbox per request starts from the snapshot, and the snapshot only knows the callbacks it was taken with — none, by default — so every request re-runs the installation (a few milliseconds of Python for a handful of callbacks, more than the instantiation itself).
+
+If the set of callbacks is fixed per deployment, declare it at pre-initialization and the wrappers are part of the snapshot. A sandbox that registers callbacks with the same names, descriptions and parameter schemas (in any order) then skips the installation; one that registers a different set still works, it just installs as before.
+
+```bash
+# callbacks.json: what the sandboxes will register at runtime
+# [
+#   {"name": "get_time", "description": "Current time", "parameters": {"type": "object", "properties": {}}},
+#   {"name": "fetch", "description": "HTTP GET", "parameters": {"type": "object", "properties": {"url": {"type": "string"}}}}
+# ]
+eryx-precompile compile runtime.wasm -o runtime.cwasm \
+  --preinit --stdlib ./python-stdlib \
+  --callbacks callbacks.json
+```
+
+`parameters` defaults to `{}` when omitted. Setup code cannot invoke the callbacks: there is no host to answer them during pre-initialization. In Rust, pass `PreInitOptions::callbacks` (build the declarations with `eryx::preinit::callback_declaration`); in Python, pass `callbacks=` to `SandboxFactory`.
+
 ### Verification
 
 By default, `compile` verifies the output by creating a test sandbox. You can add custom verification:

crates/eryx-precompile/Cargo.toml

@@ -17,6 +17,8 @@ clap = { workspace = true } #unified
 eryx = { workspace = true, features = ["preinit"] }
 reqwest.workspace = true
 rustls = { workspace = true, features = ["aws-lc-rs"] }
+serde.workspace = true      # --callbacks file
+serde_json.workspace = true
 tokio = { workspace = true, features = ["rt-multi-thread", "macros"] }
 tracing.workspace = true
 tracing-subscriber = { version = "0.3", features = ["env-filter", "registry"] }

crates/eryx-precompile/src/main.rs

@@ -36,7 +36,7 @@
 //!   --verify-code "import numpy; print(numpy.array(\[1,2,3\]).sum())"
 //! ```
 
-use std::path::PathBuf;
+use std::path::{Path, PathBuf};
 use std::time::Instant;
 
 use anyhow::{Context, Result};
@@ -193,6 +193,19 @@ struct CompileArgs {
     )]
     setup_file: Option<PathBuf>,
 
+    /// JSON file declaring the callbacks sandboxes will register (baked into snapshot)
+    ///
+    /// A list of objects with "name", "description" and optional "parameters"
+    /// (a JSON Schema object, `{}` if omitted). Their Python wrappers are
+    /// installed during pre-initialization, so a fresh instance whose host
+    /// registers callbacks with the same names, descriptions and schemas skips
+    /// the per-execution callback setup. Requires --preinit.
+    ///
+    /// Example: `--callbacks callbacks.json` with
+    /// `[{"name": "get_time", "description": "Current time", "parameters": {"type": "object"}}]`
+    #[arg(long, value_name = "PATH", requires = "preinit")]
+    callbacks: Option<PathBuf>,
+
     /// Python code to execute during verification
     ///
     /// Runs after the standard import verification. Useful for testing that
@@ -378,6 +391,40 @@ async fn run_setup(args: SetupArgs) -> Result<()> {
     Ok(())
 }
 
+/// One entry of the `--callbacks` file.
+#[derive(serde::Deserialize)]
+struct CallbackDeclarationFile {
+    name: String,
+    #[serde(default)]
+    description: String,
+    /// JSON Schema for the arguments; an empty object when omitted.
+    #[serde(default = "empty_object")]
+    parameters: serde_json::Value,
+}
+
+fn empty_object() -> serde_json::Value {
+    serde_json::Value::Object(serde_json::Map::new())
+}
+
+/// Read the callback declarations to bake into the snapshot from `path`.
+fn read_callback_declarations(path: &Path) -> Result<Vec<eryx::preinit::CallbackDeclaration>> {
+    let text = std::fs::read_to_string(path)
+        .with_context(|| format!("Failed to read callbacks file: {}", path.display()))?;
+    let entries: Vec<CallbackDeclarationFile> = serde_json::from_str(&text)
+        .with_context(|| format!("Invalid callbacks file: {}", path.display()))?;
+    entries
+        .into_iter()
+        .map(|entry| {
+            Ok(eryx::preinit::CallbackDeclaration {
+                name: entry.name,
+                description: entry.description,
+                parameters_schema_json: serde_json::to_string(&entry.parameters)
+                    .context("Failed to serialize callback parameters schema")?,
+            })
+        })
+        .collect()
+}
+
 async fn run_compile(args: CompileArgs) -> Result<()> {
     // Set up tracing
     let filter = if args.verbose {
@@ -460,6 +507,21 @@ async fn run_compile(args: CompileArgs) -> Result<()> {
                 .map_or("inline code".to_string(), |p| p.display().to_string())
         );
     }
+
+    let callbacks = match &args.callbacks {
+        Some(path) => read_callback_declarations(path)?,
+        None => Vec::new(),
+    };
+    if !callbacks.is_empty() {
+        println!(
+            "Callbacks: {} (baked into snapshot)",
+            callbacks
+                .iter()
+                .map(|cb| cb.name.as_str())
+                .collect::<Vec<_>>()
+                .join(", ")
+        );
+    }
     println!();
 
     // Read input WASM
@@ -497,22 +559,23 @@ async fn run_compile(args: CompileArgs) -> Result<()> {
             println!("Native extensions: {}", extensions.len());
         }
 
-        // Convert imports to &str references
-        let import_refs: Vec<&str> = args.imports.iter().map(|s| s.as_str()).collect();
-
         println!();
         println!("Step 1: Pre-initializing Python...");
         let start = Instant::now();
 
-        let preinit_bytes = eryx::preinit::pre_initialize(
-            stdlib,
-            final_site_packages.as_deref(),
-            &import_refs,
-            &extensions,
-            setup_code.as_deref(),
-        )
-        .await
-        .context("Failed to pre-initialize Python")?;
+        let mut options = eryx::preinit::PreInitOptions::new(stdlib)
+            .imports(args.imports.iter().cloned())
+            .extensions(extensions)
+            .callbacks(callbacks);
+        if let Some(path) = &final_site_packages {
+            options = options.site_packages(path);
+        }
+        if let Some(code) = &setup_code {
+            options = options.setup_code(code.clone());
+        }
+        let preinit_bytes = eryx::preinit::pre_initialize_with_options(options)
+            .await
+            .context("Failed to pre-initialize Python")?;
 
         let elapsed = start.elapsed();
         println!(

crates/eryx-python/python/eryx/_eryx.pyi

@@ -646,6 +646,7 @@ class SandboxFactory:
         packages: Optional[Sequence[PathLike]] = None,
         imports: Optional[Sequence[str]] = None,
         setup_code: Optional[str] = None,
+        callbacks: Optional[Union[CallbackRegistry, Sequence[CallbackDict]]] = None,
         cache: bool = True,
     ) -> None:
         """Create a new sandbox factory with custom packages.
@@ -664,6 +665,13 @@ class SandboxFactory:
                 ``SandboxedEnvironment``) so every sandbox starts with them already
                 in memory. Each sandbox gets its own copy-on-write clone, preserving
                 full isolation.
+            callbacks: Optional callbacks (a ``CallbackRegistry`` or a list of callback
+                dicts) whose declarations are baked into the snapshot, so sandboxes
+                created from this factory skip the per-sandbox callback setup (a few
+                milliseconds per sandbox). ``create_sandbox()`` and ``create_session()``
+                register these callbacks unless given their own; a sandbox given a
+                different set still works, it just installs them itself. Setup code
+                cannot invoke the callbacks.
             cache: Whether to cache the pre-compiled component in the process-global
                 cache. When enabled, a BLAKE3 content hash is computed once during
                 factory construction and subsequent ``create_sandbox()`` calls skip
@@ -700,6 +708,7 @@ class SandboxFactory:
         path: PathLike,
         *,
         site_packages: Optional[PathLike] = None,
+        callbacks: Optional[Union[CallbackRegistry, Sequence[CallbackDict]]] = None,
         cache: bool = True,
     ) -> SandboxFactory:
         """Load a sandbox factory from a file.
@@ -711,6 +720,9 @@ class SandboxFactory:
             path: Path to the saved factory file.
             site_packages: Optional path to site-packages directory.
                 Required if the factory was saved without embedded packages.
+            callbacks: The callbacks the factory was created with, if any. The file
+                holds their baked declarations but not the Python callables, so pass
+                the same callbacks here to keep the setup-free fast path.
             cache: Whether to cache the pre-compiled component in the process-global
                 cache. When enabled, a BLAKE3 content hash is computed once during
                 loading and subsequent ``create_sandbox()`` calls skip component