feat(sdkgen): embed compiled BAML more efficiently

.github/workflows/verify-rust-sdk.reusable.yaml

@@ -276,7 +276,7 @@ jobs:
           # GNU/Linux invocation; on macOS the functional smokes above already
           # prove the ABI (the consumer calls through every symbol).
           if [[ "$RUNNER_OS" == "Linux" ]]; then
-            for sym in version free_buffer create_baml_runtime initialize_runtime_from_bytecode \
+            for sym in version free_buffer create_baml_runtime initialize_runtime_from_blob \
                        register_callback new_function_call call_function; do
               nm -D "${{ steps.art.outputs.dylib }}" | grep -qw "$sym" || {
                 echo "::error::dylib missing exported symbol $sym"; exit 1; }

baml_language/Cargo.toml

@@ -233,6 +233,8 @@ fancy-regex = { version = "0.19" }
 # gitignore-aware directory walking (workspace project discovery in the LSP)
 ignore = { version = "0.4" }
 base64 = "0.22"
+# LZ4 frames compress the bytecode that generated SDKs embed.
+lz4_flex = { version = "0.13" }
 mime_guess = "2"
 bytecount = "0.6"
 reqwest = { version = "0.13.1", default-features = false, features = [

baml_language/crates/baml_artifact/Cargo.toml

@@ -23,7 +23,9 @@ workspace = true
 
 [dependencies]
 baml_version = { workspace = true }
+base64 = { workspace = true }
 borsh = { workspace = true }
 log = { workspace = true }
+lz4_flex = { workspace = true }
 sha2 = { workspace = true }
 thiserror = { workspace = true }

baml_language/crates/baml_artifact/src/lib.rs

@@ -5,8 +5,12 @@
 //! the release metadata check at bridge startup useful while adding an ABI
 //! check that runs even when generated metadata is unavailable.
 
-use std::fmt;
+use std::{
+    fmt,
+    io::{Read as _, Write as _},
+};
 
+use base64::{Engine as _, engine::general_purpose::STANDARD as BASE64};
 use borsh::{BorshDeserialize, BorshSerialize};
 use sha2::{Digest, Sha256};
 
@@ -171,6 +175,54 @@ pub enum Error {
     Encode { kind: ArtifactKind, message: String },
     #[error("failed to decode {kind}: {message}")]
     Decode { kind: ArtifactKind, message: String },
+    #[error("failed to decode embedded {kind}: {message}")]
+    Embedded { kind: ArtifactKind, message: String },
+}
+
+/// Encode an artifact for embedding as a single-line string literal in
+/// generated SDK source: an LZ4 frame, then standard padded base64.
+/// [`decode_embedded`] reverses this; bridges call it so host languages never
+/// decode the payload themselves. The encoding cannot begin with [`MAGIC`], so
+/// bridges can accept either form.
+pub fn encode_embedded(artifact: &[u8]) -> String {
+    let mut encoder = lz4_flex::frame::FrameEncoder::new(Vec::new());
+    encoder
+        .write_all(artifact)
+        .expect("writing to a Vec cannot fail");
+    BASE64.encode(encoder.finish().expect("writing to a Vec cannot fail"))
+}
+
+/// Largest artifact [`decode_embedded`] will inflate. Generated programs are a
+/// few MB; the cap keeps a malformed, highly compressible payload from
+/// allocating unbounded memory before artifact validation can reject it.
+pub const MAX_EMBEDDED_ARTIFACT_BYTES: u64 = 512 * 1024 * 1024;
+
+/// Decode an [`encode_embedded`] payload back into the artifact bytes.
+pub fn decode_embedded(kind: ArtifactKind, encoded: &[u8]) -> Result<Vec<u8>, Error> {
+    decode_embedded_within(kind, encoded, MAX_EMBEDDED_ARTIFACT_BYTES)
+}
+
+fn decode_embedded_within(
+    kind: ArtifactKind,
+    encoded: &[u8],
+    max_artifact_bytes: u64,
+) -> Result<Vec<u8>, Error> {
+    let embedded_error = |message: String| Error::Embedded { kind, message };
+    let compressed = BASE64
+        .decode(encoded)
+        .map_err(|error| embedded_error(format!("invalid base64: {error}")))?;
+    let mut artifact = Vec::new();
+    // Read one byte past the cap so an oversized artifact is detectable.
+    lz4_flex::frame::FrameDecoder::new(compressed.as_slice())
+        .take(max_artifact_bytes.saturating_add(1))
+        .read_to_end(&mut artifact)
+        .map_err(|error| embedded_error(format!("invalid LZ4 frame: {error}")))?;
+    if artifact.len() as u64 > max_artifact_bytes {
+        return Err(embedded_error(format!(
+            "decoded artifact exceeds {max_artifact_bytes} bytes"
+        )));
+    }
+    Ok(artifact)
 }
 
 /// Serialize `value` and wrap it in a versioned artifact envelope.
@@ -366,6 +418,46 @@ pub fn decode_payload(kind: ArtifactKind, bytes: &[u8]) -> Result<&[u8], Error>
 mod tests {
     use super::*;
 
+    #[test]
+    fn embedded_encoding_round_trips() {
+        let artifact = encode(ArtifactKind::Program, &vec![7_u32; 1024]).unwrap();
+        let encoded = encode_embedded(&artifact);
+        assert!(encoded.len() < artifact.len());
+        assert!(!encoded.contains('\n'));
+        assert!(!encoded.as_bytes().starts_with(MAGIC));
+        let restored = decode_embedded(ArtifactKind::Program, encoded.as_bytes()).unwrap();
+        assert_eq!(restored, artifact);
+        let decoded: Vec<u32> = decode(ArtifactKind::Program, &restored).unwrap();
+        assert_eq!(decoded, vec![7_u32; 1024]);
+    }
+
+    #[test]
+    fn oversized_embedded_artifacts_are_rejected() {
+        let artifact = encode(ArtifactKind::Program, &vec![0_u8; 4096]).unwrap();
+        let encoded = encode_embedded(&artifact);
+        let limit = artifact.len() as u64;
+        assert_eq!(
+            decode_embedded_within(ArtifactKind::Program, encoded.as_bytes(), limit).unwrap(),
+            artifact
+        );
+        let error = decode_embedded_within(ArtifactKind::Program, encoded.as_bytes(), limit - 1)
+            .unwrap_err();
+        assert!(
+            matches!(&error, Error::Embedded { message, .. } if message.contains("exceeds")),
+            "{error}"
+        );
+    }
+
+    #[test]
+    fn corrupt_embedded_payloads_are_embedded_errors() {
+        let encoded = encode_embedded(&encode(ArtifactKind::Program, &vec![1_u32; 64]).unwrap());
+        for corrupt in ["not base64!", &encoded[..(encoded.len() / 2) & !3]] {
+            let corrupt = corrupt.as_bytes();
+            let error = decode_embedded(ArtifactKind::Program, corrupt).unwrap_err();
+            assert!(matches!(error, Error::Embedded { .. }), "{error}");
+        }
+    }
+
     #[test]
     fn round_trips_payload() {
         let artifact = encode(ArtifactKind::Program, &vec![1_u32, 2, 3]).unwrap();

baml_language/crates/baml_sdkgen_types/Cargo.toml

@@ -17,6 +17,7 @@ doctest = false
 workspace = true
 
 [dependencies]
+baml_artifact = { workspace = true }
 baml_base = { workspace = true }
 baml_type = { workspace = true }
 serde = { workspace = true }

baml_language/crates/baml_sdkgen_types/src/embedded_bytecode.rs

@@ -0,0 +1,15 @@
+//! The bytecode payload that generated SDKs embed in source.
+
+/// Encode compiled BAML bytecode for embedding as a single-line string
+/// literal (see [`baml_artifact::encode_embedded`]). Generated SDKs hand the
+/// string to their bridge's embedded-bytecode initializer unchanged; decoding
+/// happens natively, never in the host language.
+///
+/// Empty bytecode stays empty so the bridge still reports a missing payload
+/// rather than a corrupt one.
+pub fn embedded_bytecode_base64(bytecode: &[u8]) -> String {
+    if bytecode.is_empty() {
+        return String::new();
+    }
+    baml_artifact::encode_embedded(bytecode)
+}

baml_language/crates/baml_sdkgen_types/src/lib.rs

@@ -2,6 +2,7 @@
 //! Attributes "don't exist"
 //! Union / Optional / Null invariance guaranteed
 
+mod embedded_bytecode;
 mod errors;
 mod generator_fields;
 mod output_writer;
@@ -10,6 +11,7 @@ mod string_literal;
 mod symbols;
 mod ty;
 mod type_queries;
+pub use embedded_bytecode::embedded_bytecode_base64;
 pub use errors::*;
 pub use generator_fields::*;
 pub use output_writer::*;

baml_language/crates/bridge_cffi/Cargo.toml

@@ -21,6 +21,7 @@ rustls = [ "sys_native/rustls" ]
 bundle-http = [ "sys_native/bundle-http" ]
 
 [dependencies]
+baml_artifact = { workspace = true }
 baml_version = { workspace = true }
 bex_project = { workspace = true }
 bridge_ctypes = { workspace = true }

baml_language/crates/bridge_cffi/cbindgen.toml

@@ -73,7 +73,7 @@ include = [
   "BamlHostReleaseCallback",
   "BamlUnhandledSpawnErrorCallback",
   "BamlVersionFn",
-  "BamlInitializeRuntimeFromBytecodeFn",
+  "BamlInitializeRuntimeFromBlobFn",
   "BamlFreeBufferFn",
   "BamlRegisterCallbackFn",
   "BamlCallFunctionFn",
@@ -118,7 +118,7 @@ exclude = [
   "flush_events",
   "version",
   "create_baml_runtime",
-  "initialize_runtime_from_bytecode",
+  "initialize_runtime_from_blob",
   "destroy_baml_runtime",
   "invoke_runtime_cli",
 ]

baml_language/crates/bridge_cffi/include/baml_cffi.h

@@ -178,8 +178,7 @@ typedef struct BamlBuffer {
 
 typedef struct BamlBuffer (*BamlVersionFn)(void);
 
-typedef struct BamlBuffer (*BamlInitializeRuntimeFromBytecodeFn)(const uint8_t *bytecode,
-                                                                 size_t length);
+typedef struct BamlBuffer (*BamlInitializeRuntimeFromBlobFn)(const uint8_t *bytecode, size_t length);
 
 typedef void (*BamlFreeBufferFn)(struct BamlBuffer buffer);
 
@@ -298,9 +297,9 @@ typedef void (*BamlRegisterUnhandledSpawnErrorCallbackFn)(BamlUnhandledSpawnErro
 
 typedef struct BamlBuffer (*BamlShutdownRuntimeFn)(void);
 
-typedef struct BamlBuffer (*BamlInitializeRuntimeFromBytecodeWithMetadataFn)(const uint8_t *bytecode,
-                                                                             size_t length,
-                                                                             const char *baml_toml);
+typedef struct BamlBuffer (*BamlInitializeRuntimeFromBlobWithMetadataFn)(const uint8_t *bytecode,
+                                                                         size_t length,
+                                                                         const char *baml_toml);
 
 /**
  * First version of the shared BAML C API.
@@ -350,7 +349,7 @@ typedef struct BamlApiV1 {
    * calls are serialized only while replacing the global runtime; calls
    * already in progress retain their previous runtime instance.
    */
-  BamlInitializeRuntimeFromBytecodeFn initialize_runtime_from_bytecode;
+  BamlInitializeRuntimeFromBlobFn initialize_runtime_from_blob;
   /**
    * Release exactly one runtime-owned buffer returned through this table.
    *
@@ -483,7 +482,7 @@ typedef struct BamlApiV1 {
   /**
    * Replace the runtime from bytecode after validating embedded generation metadata.
    */
-  BamlInitializeRuntimeFromBytecodeWithMetadataFn initialize_runtime_from_bytecode_with_metadata;
+  BamlInitializeRuntimeFromBlobWithMetadataFn initialize_runtime_from_blob_with_metadata;
 } BamlApiV1;
 
 typedef const struct BamlApiV1 *(*BamlGetApiV1Fn)(void);
@@ -515,9 +514,9 @@ BamlCffiStatus __testonly_seed_heap_handle(uint64_t slab_key,
  * Initialize generated bytecode after validating its embedded `baml.toml`.
  */
 BAML_CFFI_API
-struct BamlBuffer initialize_runtime_from_bytecode_with_metadata(const uint8_t *bytecode,
-                                                                 size_t length,
-                                                                 const char *baml_toml);
+struct BamlBuffer initialize_runtime_from_blob_with_metadata(const uint8_t *bytecode,
+                                                             size_t length,
+                                                             const char *baml_toml);
 
 #ifdef __cplusplus
 }  // extern "C"

baml_language/crates/bridge_cffi/src/api.rs

@@ -95,9 +95,9 @@ pub type BamlUnhandledSpawnErrorCallback =
     extern "C" fn(content: *const i8, length: usize, cancelled: i32);
 
 pub type BamlVersionFn = extern "C" fn() -> Buffer;
-pub type BamlInitializeRuntimeFromBytecodeFn =
+pub type BamlInitializeRuntimeFromBlobFn =
     extern "C" fn(bytecode: *const u8, length: usize) -> Buffer;
-pub type BamlInitializeRuntimeFromBytecodeWithMetadataFn =
+pub type BamlInitializeRuntimeFromBlobWithMetadataFn =
     extern "C" fn(bytecode: *const u8, length: usize, baml_toml: *const libc::c_char) -> Buffer;
 pub type BamlFreeBufferFn = extern "C" fn(buffer: Buffer);
 pub type BamlRegisterCallbackFn = extern "C" fn(callback: BamlResultCallback);
@@ -167,7 +167,7 @@ pub struct BamlApiV1 {
     /// failure and must always be passed once to `free_buffer`. Concurrent
     /// calls are serialized only while replacing the global runtime; calls
     /// already in progress retain their previous runtime instance.
-    pub initialize_runtime_from_bytecode: BamlInitializeRuntimeFromBytecodeFn,
+    pub initialize_runtime_from_blob: BamlInitializeRuntimeFromBlobFn,
     /// Release exactly one runtime-owned buffer returned through this table.
     ///
     /// Do not use a different library instance's function, release a copied
@@ -258,15 +258,14 @@ pub struct BamlApiV1 {
     /// Wait for spawned work, report unreachable errors, and release the runtime.
     pub shutdown_runtime: BamlShutdownRuntimeFn,
     /// Replace the runtime from bytecode after validating embedded generation metadata.
-    pub initialize_runtime_from_bytecode_with_metadata:
-        BamlInitializeRuntimeFromBytecodeWithMetadataFn,
+    pub initialize_runtime_from_blob_with_metadata: BamlInitializeRuntimeFromBlobWithMetadataFn,
 }
 
 static BAML_API_V1: BamlApiV1 = BamlApiV1 {
     abi_version: BAML_API_V1_ABI_VERSION,
     struct_size: std::mem::size_of::<BamlApiV1>(),
     version: crate::version,
-    initialize_runtime_from_bytecode: crate::initialize_runtime_from_bytecode_ffi,
+    initialize_runtime_from_blob: crate::initialize_runtime_from_blob_ffi,
     free_buffer: crate::free_buffer,
     register_callback: crate::register_callback,
     call_function: crate::call_function,
@@ -287,8 +286,7 @@ static BAML_API_V1: BamlApiV1 = BamlApiV1 {
     register_bridge: crate::register_bridge_ffi,
     register_unhandled_spawn_error_callback: crate::register_unhandled_spawn_error_callback,
     shutdown_runtime: crate::shutdown_runtime_ffi,
-    initialize_runtime_from_bytecode_with_metadata:
-        crate::initialize_runtime_from_bytecode_with_metadata,
+    initialize_runtime_from_blob_with_metadata: crate::initialize_runtime_from_blob_with_metadata,
 };
 
 /// Return the immutable version-1 BAML C API function table.
@@ -328,8 +326,8 @@ mod tests {
         let api = unsafe { &*baml_get_api_v1() };
         assert_same_function!(api.version, crate::version);
         assert_same_function!(
-            api.initialize_runtime_from_bytecode,
-            crate::initialize_runtime_from_bytecode_ffi
+            api.initialize_runtime_from_blob,
+            crate::initialize_runtime_from_blob_ffi
         );
         assert_same_function!(api.free_buffer, crate::free_buffer);
         assert_same_function!(api.register_callback, crate::register_callback);
@@ -361,16 +359,16 @@ mod tests {
         );
         assert_same_function!(api.shutdown_runtime, crate::shutdown_runtime_ffi);
         assert_same_function!(
-            api.initialize_runtime_from_bytecode_with_metadata,
-            crate::initialize_runtime_from_bytecode_with_metadata
+            api.initialize_runtime_from_blob_with_metadata,
+            crate::initialize_runtime_from_blob_with_metadata
         );
     }
 
     #[test]
     fn every_v1_field_retains_its_declared_function_type() {
         let api = unsafe { &*baml_get_api_v1() };
         let _: BamlVersionFn = api.version;
-        let _: BamlInitializeRuntimeFromBytecodeFn = api.initialize_runtime_from_bytecode;
+        let _: BamlInitializeRuntimeFromBlobFn = api.initialize_runtime_from_blob;
         let _: BamlFreeBufferFn = api.free_buffer;
         let _: BamlRegisterCallbackFn = api.register_callback;
         let _: BamlCallFunctionFn = api.call_function;
@@ -392,8 +390,8 @@ mod tests {
         let _: BamlRegisterUnhandledSpawnErrorCallbackFn =
             api.register_unhandled_spawn_error_callback;
         let _: BamlShutdownRuntimeFn = api.shutdown_runtime;
-        let _: BamlInitializeRuntimeFromBytecodeWithMetadataFn =
-            api.initialize_runtime_from_bytecode_with_metadata;
+        let _: BamlInitializeRuntimeFromBlobWithMetadataFn =
+            api.initialize_runtime_from_blob_with_metadata;
         let _: BamlGetApiV1Fn = baml_get_api_v1;
     }