From 4c79545bd22e99225aef59bcff4bae1ed79a57e9 Mon Sep 17 00:00:00 2001 From: Natanael Mojica Date: Mon, 6 Jul 2026 11:01:55 -0600 Subject: [PATCH] fix: embed host compaction baml --- Dockerfile | 5 +- .../baml-agent-runner/src/host_compaction.rs | 69 ++++++++++--- crates/baml-rt-a2a/src/host_compaction/mod.rs | 34 +++++-- crates/baml-rt-a2a/src/lib.rs | 2 +- .../tests/host_compaction_baml_test.rs | 56 +++++++---- .../baml-rt-quickjs/src/baml/schema_invoke.rs | 96 ++++++++++++------- crates/baml-rt-quickjs/src/baml_execution.rs | 27 +++++- .../tests/load_schema_from_files_test.rs | 53 ++++++++++ 8 files changed, 265 insertions(+), 77 deletions(-) create mode 100644 crates/baml-rt-quickjs/tests/load_schema_from_files_test.rs diff --git a/Dockerfile b/Dockerfile index 23bbb310e..a92ff5390 100644 --- a/Dockerfile +++ b/Dockerfile @@ -20,6 +20,7 @@ WORKDIR /build # Copy manifests first for better layer caching. COPY Cargo.toml Cargo.lock rust-toolchain.toml ./ COPY crates/ crates/ +COPY baml_src/host baml_src/host # Override to stable (rust-toolchain.toml pins nightly for local dev). ENV RUSTUP_TOOLCHAIN=stable @@ -57,10 +58,6 @@ RUN mkdir -p /data && chown -R agentium:agentium /data COPY --from=builder /build/target/release/agentium /usr/local/bin/ -# Host-owned context compaction BAML (SummarizeConversationPrefix). -COPY baml_src/host /opt/agentium/baml_src/host -ENV BAML_HOST_SCHEMA_DIR=/opt/agentium - # ONNX models for embedding/drift detection (git-lfs tracked). # Fail fast if LFS pointers were not resolved (e.g. missing `git lfs pull`). COPY models/fastembed /models/fastembed diff --git a/crates/baml-agent-runner/src/host_compaction.rs b/crates/baml-agent-runner/src/host_compaction.rs index e13c7b253..870cc186b 100644 --- a/crates/baml-agent-runner/src/host_compaction.rs +++ b/crates/baml-agent-runner/src/host_compaction.rs @@ -6,7 +6,7 @@ use std::{path::PathBuf, sync::Arc}; -use baml_rt_a2a::{HostCompactionConfig, HostCompactionEngine}; +use baml_rt_a2a::{HostCompactionConfig, HostCompactionEngine, HostCompactionSource}; use baml_rt_core::{Result, bus::EffectEmitter}; use baml_rt_llm_config::{StaticResolver, load_stored_config}; use baml_rt_provenance::ConversationCompactionSummarizer; @@ -14,16 +14,29 @@ use baml_rt_quickjs::llm_resolver_adapter::SecretResolverToLlmAdapter; use crate::config::ProvenanceConfig; -fn default_host_schema_dir() -> PathBuf { - if let Ok(dir) = std::env::var("BAML_HOST_SCHEMA_DIR") { - return PathBuf::from(dir); - } - const CONTAINER_DEFAULT: &str = "/opt/agentium"; - let container_root = PathBuf::from(CONTAINER_DEFAULT); - if container_root.join("baml_src").is_dir() { - return container_root; +const HOST_CLIENTS_BAML: &str = include_str!("../../../baml_src/host/clients.baml"); +const HOST_COMPACTION_BAML: &str = include_str!("../../../baml_src/host/context_compaction.baml"); + +const EMBEDDED_HOST_BAML: &[(&str, &str)] = &[ + ("baml_src/host/clients.baml", HOST_CLIENTS_BAML), + ( + "baml_src/host/context_compaction.baml", + HOST_COMPACTION_BAML, + ), +]; + +fn host_compaction_source_from(env_override: Option) -> HostCompactionSource { + match env_override { + Some(dir) => HostCompactionSource::Dir(PathBuf::from(dir)), + None => HostCompactionSource::Embedded { + root: "baml_src", + files: EMBEDDED_HOST_BAML, + }, } - PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("../..") +} + +fn host_compaction_source() -> HostCompactionSource { + host_compaction_source_from(std::env::var("BAML_HOST_SCHEMA_DIR").ok()) } pub(crate) async fn boot_compaction_summarizer( @@ -47,7 +60,7 @@ pub(crate) async fn boot_compaction_summarizer( )); let engine = HostCompactionEngine::boot( HostCompactionConfig { - schema_dir: default_host_schema_dir(), + source: host_compaction_source(), }, effect_emitter, llm_resolver, @@ -56,3 +69,37 @@ pub(crate) async fn boot_compaction_summarizer( .await?; Ok(Arc::new(engine)) } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn host_compaction_source_uses_dir_override() { + match host_compaction_source_from(Some("/tmp/host-baml".into())) { + HostCompactionSource::Dir(path) => assert_eq!(path, PathBuf::from("/tmp/host-baml")), + HostCompactionSource::Embedded { .. } => panic!("expected dir source"), + } + } + + #[test] + fn host_compaction_source_defaults_to_embedded() { + match host_compaction_source_from(None) { + HostCompactionSource::Embedded { root, files } => { + assert_eq!(root, "baml_src"); + assert_eq!(files.len(), 2); + assert!( + files + .iter() + .any(|(path, _)| *path == "baml_src/host/clients.baml") + ); + assert!( + files + .iter() + .any(|(path, _)| *path == "baml_src/host/context_compaction.baml") + ); + } + HostCompactionSource::Dir(_) => panic!("expected embedded source"), + } + } +} diff --git a/crates/baml-rt-a2a/src/host_compaction/mod.rs b/crates/baml-rt-a2a/src/host_compaction/mod.rs index 08d090d25..f419a9bdc 100644 --- a/crates/baml-rt-a2a/src/host_compaction/mod.rs +++ b/crates/baml-rt-a2a/src/host_compaction/mod.rs @@ -4,7 +4,7 @@ //! Host-owned BAML compaction engine (implements [`ConversationCompactionSummarizer`]). -use std::sync::Arc; +use std::{path::PathBuf, sync::Arc}; use async_trait::async_trait; use baml_rt_core::{Result, bus::EffectEmitter, context::RuntimeScope}; @@ -17,9 +17,18 @@ use baml_rt_quickjs::{BamlRuntimeManager, llm_client_registry::LlmSecretResolver use serde_json::{Value, json}; use tokio::sync::RwLock; -/// Directory whose `baml_src/` subtree contains host compaction BAML (repo root or package root). +pub enum HostCompactionSource { + /// Load host BAML from embedded, compiled-in content. + Embedded { + root: &'static str, + files: &'static [(&'static str, &'static str)], + }, + /// Load from an on-disk `/baml_src` (dev/operator override). + Dir(PathBuf), +} + pub struct HostCompactionConfig { - pub schema_dir: std::path::PathBuf, + pub source: HostCompactionSource, } /// Shared host BAML runtime loaded with `SummarizeConversationPrefix`. @@ -39,12 +48,19 @@ impl HostCompactionEngine { .build()?; manager.set_llm_client_resolver(llm_resolver); manager.set_effect_emitter(effect_emitter); - let schema_path = config.schema_dir.to_str().ok_or_else(|| { - baml_rt_core::BamlRtError::InvalidArgument( - "host compaction schema path is not valid UTF-8".into(), - ) - })?; - manager.load_schema(schema_path)?; + match config.source { + HostCompactionSource::Embedded { root, files } => { + manager.load_schema_from_files(root, files)?; + } + HostCompactionSource::Dir(schema_dir) => { + let schema_path = schema_dir.to_str().ok_or_else(|| { + baml_rt_core::BamlRtError::InvalidArgument( + "host compaction schema path is not valid UTF-8".into(), + ) + })?; + manager.load_schema(schema_path)?; + } + } Ok(Self { runtime: Arc::new(RwLock::new(manager)), }) diff --git a/crates/baml-rt-a2a/src/lib.rs b/crates/baml-rt-a2a/src/lib.rs index 9dd8df37c..28bd7622d 100644 --- a/crates/baml-rt-a2a/src/lib.rs +++ b/crates/baml-rt-a2a/src/lib.rs @@ -45,4 +45,4 @@ pub use baml_rt_core::{A2aJsChatHost, A2aRequestHandler}; pub use conversation_context::ProjectingConversationContextProvider; pub use dispatch_port::RegistryDispatchPort; pub use event_dispatcher::EventDispatcher; -pub use host_compaction::{HostCompactionConfig, HostCompactionEngine}; +pub use host_compaction::{HostCompactionConfig, HostCompactionEngine, HostCompactionSource}; diff --git a/crates/baml-rt-a2a/tests/host_compaction_baml_test.rs b/crates/baml-rt-a2a/tests/host_compaction_baml_test.rs index 8b0301489..71c659062 100644 --- a/crates/baml-rt-a2a/tests/host_compaction_baml_test.rs +++ b/crates/baml-rt-a2a/tests/host_compaction_baml_test.rs @@ -6,7 +6,7 @@ use std::{path::PathBuf, sync::Arc}; -use baml_rt_a2a::{HostCompactionConfig, HostCompactionEngine}; +use baml_rt_a2a::{HostCompactionConfig, HostCompactionEngine, HostCompactionSource}; use baml_rt_core::bus::BusWithEffects; use baml_rt_llm_config::{LlmClientConfig, StaticResolver, test_model_default}; use baml_rt_provenance::{ @@ -20,6 +20,17 @@ use baml_rt_tools::{ }; use test_support::testing::provenance_fixtures::{provenance_agent_id, provenance_context_id}; +const HOST_CLIENTS_BAML: &str = include_str!("../../../baml_src/host/clients.baml"); +const HOST_COMPACTION_BAML: &str = include_str!("../../../baml_src/host/context_compaction.baml"); + +const EMBEDDED_HOST_BAML: &[(&str, &str)] = &[ + ("baml_src/host/clients.baml", HOST_CLIENTS_BAML), + ( + "baml_src/host/context_compaction.baml", + HOST_COMPACTION_BAML, + ), +]; + struct EnvSecretResolver; impl LlmSecretResolver for EnvSecretResolver { @@ -32,6 +43,23 @@ fn host_schema_dir() -> PathBuf { PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("../..") } +async fn boot_host_compaction_engine(source: HostCompactionSource) -> HostCompactionEngine { + let effect_emitter: Arc = Arc::new(BusWithEffects::new()); + let llm_config = LlmClientConfig::sensible_default(); + let llm_resolver = Arc::new(StaticResolver::new( + Arc::new(llm_config), + Arc::new(baml_rt_llm_config::FnoxFileSecretResolver::default_path_resolver()), + )); + HostCompactionEngine::boot( + HostCompactionConfig { source }, + effect_emitter, + llm_resolver, + Arc::new(EnvSecretResolver), + ) + .await + .expect("boot host compaction engine") +} + fn compaction_ref_table_with_at3_and_hash2() -> Arc { let table = Arc::new(RefTable::new()); table.insert_virtual_archive( @@ -52,6 +80,11 @@ fn compaction_ref_table_with_at3_and_hash2() -> Arc { table } +#[tokio::test] +async fn host_compaction_boots_from_dir_override() { + let _engine = boot_host_compaction_engine(HostCompactionSource::Dir(host_schema_dir())).await; +} + #[tokio::test] async fn host_compaction_baml_validates_wire_refs_after_finalize() { if std::env::var("OPENROUTER_API_KEY").is_err() { @@ -62,23 +95,12 @@ async fn host_compaction_baml_validates_wire_refs_after_finalize() { std::env::set_var("BAML_TEST_MODEL", test_model_default()); } - let effect_emitter: Arc = Arc::new(BusWithEffects::new()); - let llm_config = LlmClientConfig::sensible_default(); - let llm_resolver = Arc::new(StaticResolver::new( - Arc::new(llm_config), - Arc::new(baml_rt_llm_config::FnoxFileSecretResolver::default_path_resolver()), - )); let engine = Arc::new( - HostCompactionEngine::boot( - HostCompactionConfig { - schema_dir: host_schema_dir(), - }, - effect_emitter, - llm_resolver, - Arc::new(EnvSecretResolver), - ) - .await - .expect("boot host compaction engine"), + boot_host_compaction_engine(HostCompactionSource::Embedded { + root: "baml_src", + files: EMBEDDED_HOST_BAML, + }) + .await, ); let request = CompactionRequest { diff --git a/crates/baml-rt-quickjs/src/baml/schema_invoke.rs b/crates/baml-rt-quickjs/src/baml/schema_invoke.rs index 8e0073c8f..b88568cbe 100644 --- a/crates/baml-rt-quickjs/src/baml/schema_invoke.rs +++ b/crates/baml-rt-quickjs/src/baml/schema_invoke.rs @@ -42,7 +42,7 @@ impl BamlRuntimeManager { // Resolve LLM secrets from fnox resolver and inject as env_vars so BAML schema's // `api_key env.X` references resolve without relying on std::env::var. let env_vars = self.resolve_secrets_as_env_vars(); - let mut executor = BamlExecutor::load_il(&baml_src_dir, env_vars)?; + let executor = BamlExecutor::load_il(&baml_src_dir, env_vars)?; // Load the rendered agent-wide tool schema catalog produced by the builder. This is // the JSON-shape catalog text emitted via BAML's `{{ ctx.output_format }}` renderer @@ -50,7 +50,7 @@ impl BamlRuntimeManager { // `_baml_runtime.baml` source. Old packages without the sidecar simply load no // prelude (graceful degradation; no source dump fallback). let catalog_path = baml_src_dir.join(baml_rt_tools::TOOL_SCHEMA_CATALOG_SIDECAR_FILE); - self.state.tool_schema_prelude = match std::fs::read_to_string(&catalog_path) { + let tool_schema_prelude = match std::fs::read_to_string(&catalog_path) { Ok(s) => Some(std::sync::Arc::::from(s)), Err(e) => { tracing::debug!( @@ -61,6 +61,67 @@ impl BamlRuntimeManager { None } }; + self.finish_schema_load(executor, tool_schema_prelude); + + self.state.session_plan_functions = runtime_io::load_build_manifest::< + SessionPlanFunctionsMap, + >(project_root, "session_plan_functions.json"); + self.state.tool_step_executors = runtime_io::load_build_manifest::< + std::collections::HashMap, + >(project_root, "tool_step_executors.json"); + self.state.unified_step_executor_functions = + runtime_io::load_build_manifest::( + project_root, + "unified_step_executor_functions.json", + ); + + tracing::debug!( + function_count = self.state.function_registry.len(), + session_plan_manifest = self + .state + .session_plan_functions + .as_ref() + .map(|m| m.len()) + .unwrap_or(0), + unified_step_executor_roots = self + .state + .unified_step_executor_functions + .as_ref() + .map(|m| m.len()) + .unwrap_or(0), + "Loaded BAML IL" + ); + + Ok(()) + } + + /// Load BAML schema from embedded/in-memory files, without disk sidecars. + pub fn load_schema_from_files( + &mut self, + root_path: &str, + files: &[(&str, &str)], + ) -> Result<()> { + tracing::debug!( + root_path, + file_count = files.len(), + "Loading BAML IL from files" + ); + let env_vars = self.resolve_secrets_as_env_vars(); + let files_map: HashMap<&str, &str> = files.iter().copied().collect(); + let executor = BamlExecutor::load_il_from_content(root_path, &files_map, env_vars)?; + self.finish_schema_load(executor, None); + self.state.session_plan_functions = None; + self.state.tool_step_executors = None; + self.state.unified_step_executor_functions = None; + Ok(()) + } + + fn finish_schema_load( + &mut self, + mut executor: BamlExecutor, + tool_schema_prelude: Option>, + ) { + self.state.tool_schema_prelude = tool_schema_prelude; // Set effect emitter if available if let Some(ref emitter) = self.state.effect_emitter { @@ -97,37 +158,6 @@ impl BamlRuntimeManager { } self.state.executor = Some(executor); - - self.state.session_plan_functions = runtime_io::load_build_manifest::< - SessionPlanFunctionsMap, - >(project_root, "session_plan_functions.json"); - self.state.tool_step_executors = runtime_io::load_build_manifest::< - std::collections::HashMap, - >(project_root, "tool_step_executors.json"); - self.state.unified_step_executor_functions = - runtime_io::load_build_manifest::( - project_root, - "unified_step_executor_functions.json", - ); - - tracing::debug!( - function_count = self.state.function_registry.len(), - session_plan_manifest = self - .state - .session_plan_functions - .as_ref() - .map(|m| m.len()) - .unwrap_or(0), - unified_step_executor_roots = self - .state - .unified_step_executor_functions - .as_ref() - .map(|m| m.len()) - .unwrap_or(0), - "Loaded BAML IL" - ); - - Ok(()) } /// Get the signature of a function by name diff --git a/crates/baml-rt-quickjs/src/baml_execution.rs b/crates/baml-rt-quickjs/src/baml_execution.rs index 96ca4f7c9..37a698311 100644 --- a/crates/baml-rt-quickjs/src/baml_execution.rs +++ b/crates/baml-rt-quickjs/src/baml_execution.rs @@ -261,7 +261,30 @@ impl BamlExecutor { let runtime = BamlRuntime::from_directory(baml_src_dir, env_vars, feature_flags) .map_err(|e| BamlRtError::RuntimeLoadFailed { source: e })?; - Ok(Self { + Ok(Self::from_runtime(runtime)) + } + + /// Load BAML IL from in-memory source (path -> content), no filesystem access. + pub fn load_il_from_content( + root_path: &str, + files: &HashMap<&str, &str>, + env_vars: HashMap, + ) -> Result { + tracing::debug!( + root_path, + file_count = files.len(), + "Loading BAML runtime from content" + ); + + let feature_flags = internal_baml_core::feature_flags::FeatureFlags::default(); + let runtime = BamlRuntime::from_file_content(root_path, files, env_vars, feature_flags) + .map_err(|e| BamlRtError::RuntimeLoadFailed { source: e })?; + + Ok(Self::from_runtime(runtime)) + } + + fn from_runtime(runtime: BamlRuntime) -> Self { + Self { runtime: Arc::new(runtime), effect_emitter: None, conversation_context_provider: None, @@ -269,7 +292,7 @@ impl BamlExecutor { llm_secret_resolver: None, llm_client_resolver: None, llm_fallback_client_resolver: None, - }) + } } /// Set the LLM secret resolver for ClientRegistry-based API key injection. diff --git a/crates/baml-rt-quickjs/tests/load_schema_from_files_test.rs b/crates/baml-rt-quickjs/tests/load_schema_from_files_test.rs new file mode 100644 index 000000000..f1978be13 --- /dev/null +++ b/crates/baml-rt-quickjs/tests/load_schema_from_files_test.rs @@ -0,0 +1,53 @@ +// SPDX-FileCopyrightText: 2026 Semiotic AI, Inc. +// +// SPDX-License-Identifier: Apache-2.0 + +use baml_rt_provenance::HOST_COMPACTION_BAML_FUNCTION; +use baml_rt_quickjs::BamlRuntimeManager; + +const HOST_CLIENTS_BAML: &str = include_str!("../../../baml_src/host/clients.baml"); +const HOST_COMPACTION_BAML: &str = include_str!("../../../baml_src/host/context_compaction.baml"); + +const EMBEDDED_HOST_BAML: &[(&str, &str)] = &[ + ("baml_src/host/clients.baml", HOST_CLIENTS_BAML), + ( + "baml_src/host/context_compaction.baml", + HOST_COMPACTION_BAML, + ), +]; + +#[test] +fn load_schema_from_files_registers_embedded_host_compaction_function() { + let mut manager = BamlRuntimeManager::new().expect("manager"); + manager + .load_schema_from_files("baml_src", EMBEDDED_HOST_BAML) + .expect("load embedded host BAML"); + + assert!( + manager + .get_function_signature(HOST_COMPACTION_BAML_FUNCTION) + .is_some() + ); +} + +#[test] +fn load_schema_from_files_does_not_require_disk_root() { + let files = &[ + ("/nonexistent/baml_src/host/clients.baml", HOST_CLIENTS_BAML), + ( + "/nonexistent/baml_src/host/context_compaction.baml", + HOST_COMPACTION_BAML, + ), + ]; + + let mut manager = BamlRuntimeManager::new().expect("manager"); + manager + .load_schema_from_files("/nonexistent/baml_src", files) + .expect("load embedded host BAML from nonexistent root"); + + assert!( + manager + .get_function_signature(HOST_COMPACTION_BAML_FUNCTION) + .is_some() + ); +}