use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex};
use std::time::Duration;
use tokio::sync::{mpsc, RwLock};
use mlua_isle::{AsyncIsle, AsyncIsleDriver};
use tracing::{info, info_span, warn};
use crate::bridge;
use crate::bus::{Event, EventBus, Handler};
use agent_block_mcp::McpManager;
use agent_block_types::error::{BlockError, BlockResult};
use tokio_util::sync::CancellationToken;
const EMBEDDED_BLOCKS: &[(&str, &str)] = &[
("agent", include_str!("../blocks/agent/init.lua")),
(
"compile_loop",
include_str!("../blocks/tools/compile_loop/init.lua"),
),
(
"coding_agent",
include_str!("../blocks/tools/coding_agent/init.lua"),
),
];
const EMBEDDED_LIBS: &[(&str, &str)] = &[
("session", include_str!("../blocks/lib/session/init.lua")),
(
"llm_proto",
include_str!("../blocks/lib/llm_proto/init.lua"),
),
(
"llm_proto.openai",
include_str!("../blocks/lib/llm_proto/openai.lua"),
),
(
"llm_proto.anthropic",
include_str!("../blocks/lib/llm_proto/anthropic.lua"),
),
("lshape", include_str!("../blocks/lib/lshape/init.lua")),
("lshape.t", include_str!("../blocks/lib/lshape/t.lua")),
(
"lshape.check",
include_str!("../blocks/lib/lshape/check.lua"),
),
(
"lshape.reflect",
include_str!("../blocks/lib/lshape/reflect.lua"),
),
(
"lshape.luacats",
include_str!("../blocks/lib/lshape/luacats.lua"),
),
(
"mcp_tools",
include_str!("../blocks/lib/mcp_tools/init.lua"),
),
("knl", include_str!("../blocks/lib/knl/init.lua")),
(
"knl_adapter",
include_str!("../blocks/lib/knl_adapter/init.lua"),
),
("policy", include_str!("../blocks/lib/policy/init.lua")),
(
"supervisor",
include_str!("../blocks/lib/supervisor/init.lua"),
),
];
const SEALED: &[&str] = &[
"knl",
"knl_adapter",
"knl_types",
"lshape",
"lshape.t",
"lshape.check",
"lshape.reflect",
"lshape.luacats",
];
const EMBEDDED_ALIAS_PREFIX: &str = "embedded.";
const DEFAULT_AGENT_INVOKER: &str = r#"
local agent = require("agent")
local r = agent.run({
prompt = _PROMPT,
system = _CONTEXT,
})
bus.emit("_", r)
"#;
#[derive(Debug, Clone)]
pub enum ScriptSource {
Path(PathBuf),
Inline {
source: String,
name: String,
},
DefaultAgent,
}
#[derive(Debug, Clone)]
pub enum PromptSource {
Inline(String),
File(PathBuf),
}
#[derive(Debug, Clone)]
pub enum SecretKeySource {
Inline(String),
Env(String),
}
#[async_trait::async_trait]
pub trait ToolHandler: Send + Sync + 'static {
async fn call(&self, input: serde_json::Value) -> Result<serde_json::Value, BlockError>;
}
#[derive(Clone)]
pub struct HostToolSpec {
pub name: String,
pub description: String,
pub input_schema: serde_json::Value,
pub group: Option<String>,
pub handler: Arc<dyn ToolHandler>,
}
impl std::fmt::Debug for HostToolSpec {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("HostToolSpec")
.field("name", &self.name)
.field("description", &self.description)
.field("input_schema", &self.input_schema)
.field("group", &self.group)
.field("handler", &"<dyn ToolHandler>")
.finish()
}
}
#[derive(Debug, Clone)]
pub struct ToolMeta {
pub name: String,
pub description: String,
pub group: Option<String>,
pub source: ToolSource,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ToolSource {
HostRust,
EmbeddedBlock,
}
pub fn inspect_tools(config: &BlockConfig) -> Vec<ToolMeta> {
let mut out = Vec::new();
for t in &config.host_tools {
out.push(ToolMeta {
name: t.name.clone(),
description: t.description.clone(),
group: t.group.clone(),
source: ToolSource::HostRust,
});
}
for (name, _src) in EMBEDDED_BLOCKS {
out.push(ToolMeta {
name: (*name).to_string(),
description: format!("Embedded StdPkg block (require(\"{name}\"))"),
group: None,
source: ToolSource::EmbeddedBlock,
});
}
out
}
pub fn lib_roots(project_root: &Path) -> Vec<PathBuf> {
let mut out = Vec::new();
let project_lib = project_root.join("lib");
if project_lib.is_dir() {
out.push(project_lib);
}
if let Ok(home) = crate::bridge::config::base_dir() {
let user_lib = home.join("lib");
if user_lib.is_dir() {
out.push(user_lib);
}
}
out
}
pub fn block_roots(project_root: &Path) -> Vec<PathBuf> {
let mut out = Vec::new();
let project_blocks = project_root.join("blocks");
if project_blocks.is_dir() {
out.push(project_blocks);
}
if let Ok(home) = crate::bridge::config::base_dir() {
let user_blocks = home.join("blocks");
if user_blocks.is_dir() {
out.push(user_blocks);
}
}
out
}
fn package_path_prefix(roots: &[PathBuf]) -> String {
let mut out = String::new();
for root in roots {
let r = root.to_string_lossy();
out.push_str(&format!("{r}/?.lua;{r}/?/init.lua;"));
}
out
}
fn require_candidates(root: &Path, name: &str) -> [PathBuf; 2] {
let relative = name.replace('.', "/");
[
root.join(format!("{relative}.lua")),
root.join(format!("{relative}/init.lua")),
]
}
fn check_sealed_modules(roots: &[PathBuf]) -> BlockResult<()> {
let unsealed = std::env::var("AGENT_BLOCK_UNSEAL").is_ok_and(|v| v == "1");
for root in roots {
for name in SEALED {
for path in require_candidates(root, name) {
if !path.is_file() {
continue;
}
if unsealed {
warn!(
module = %name,
path = %path.display(),
"AGENT_BLOCK_UNSEAL=1: a sealed module is being replaced by a filesystem copy"
);
continue;
}
return Err(BlockError::Runtime(sealed_refusal(name, &path)));
}
}
}
Ok(())
}
fn sealed_refusal(name: &str, path: &Path) -> String {
format!(
"sealed module `{name}` cannot be shadowed: {} would replace the embedded one. \
The kernel is one thing across Rust and Lua — `knl` / `knl_adapter` / `knl_types` \
and the `lshape` they are declared in are held together by declaration tests, and a \
Lua-side replacement passes those tests while meaning something else. Change it \
upstream. To read the embedded module (wrapping it is fine, replacing it is not), \
`require(\"{EMBEDDED_ALIAS_PREFIX}{name}\")`. Set AGENT_BLOCK_UNSEAL=1 to downgrade \
this refusal to a warning — for work on the kernel itself, not for shipping.",
path.display()
)
}
#[non_exhaustive]
pub struct BlockConfig {
pub script: ScriptSource,
pub project_root: PathBuf,
pub relay_url: Option<String>,
pub secret_key: Option<SecretKeySource>,
pub mcp_rpc_timeout: Duration,
pub prompt: Option<PromptSource>,
pub context: Option<PromptSource>,
pub host_handlers: HashMap<String, Arc<dyn Handler>>,
pub host_handler: Option<Arc<dyn Handler>>,
pub host_tools: Vec<HostToolSpec>,
pub http_client: Option<reqwest::Client>,
pub sql_path: Option<PathBuf>,
pub kv_path: Option<PathBuf>,
pub ts_path: Option<PathBuf>,
pub extra_globals: HashMap<String, serde_json::Value>,
pub auto_serve_bus: bool,
pub shutdown_token: Option<CancellationToken>,
}
impl BlockConfig {
pub fn builder(script: ScriptSource, project_root: PathBuf) -> BlockConfigBuilder {
BlockConfigBuilder::new(script, project_root)
}
}
pub struct BlockConfigBuilder {
script: ScriptSource,
project_root: PathBuf,
relay_url: Option<String>,
secret_key: Option<SecretKeySource>,
mcp_rpc_timeout: Duration,
prompt: Option<PromptSource>,
context: Option<PromptSource>,
host_handlers: HashMap<String, Arc<dyn Handler>>,
host_handler: Option<Arc<dyn Handler>>,
host_tools: Vec<HostToolSpec>,
http_client: Option<reqwest::Client>,
sql_path: Option<PathBuf>,
kv_path: Option<PathBuf>,
ts_path: Option<PathBuf>,
extra_globals: HashMap<String, serde_json::Value>,
auto_serve_bus: bool,
shutdown_token: Option<CancellationToken>,
}
impl BlockConfigBuilder {
fn new(script: ScriptSource, project_root: PathBuf) -> Self {
Self {
script,
project_root,
relay_url: None,
secret_key: None,
mcp_rpc_timeout: agent_block_mcp::DEFAULT_RPC_TIMEOUT,
prompt: None,
context: None,
host_handlers: HashMap::new(),
host_handler: None,
host_tools: Vec::new(),
http_client: None,
sql_path: None,
kv_path: None,
ts_path: None,
extra_globals: HashMap::new(),
auto_serve_bus: false,
shutdown_token: None,
}
}
pub fn script(mut self, script: ScriptSource) -> Self {
self.script = script;
self
}
pub fn project_root(mut self, project_root: impl Into<PathBuf>) -> Self {
self.project_root = project_root.into();
self
}
pub fn relay_url(mut self, relay_url: impl Into<String>) -> Self {
self.relay_url = Some(relay_url.into());
self
}
pub fn secret_key(mut self, secret_key: SecretKeySource) -> Self {
self.secret_key = Some(secret_key);
self
}
pub fn mcp_rpc_timeout(mut self, mcp_rpc_timeout: Duration) -> Self {
self.mcp_rpc_timeout = mcp_rpc_timeout;
self
}
pub fn prompt(mut self, prompt: PromptSource) -> Self {
self.prompt = Some(prompt);
self
}
pub fn context(mut self, context: PromptSource) -> Self {
self.context = Some(context);
self
}
pub fn host_handlers(mut self, host_handlers: HashMap<String, Arc<dyn Handler>>) -> Self {
self.host_handlers = host_handlers;
self
}
pub fn host_handler(mut self, host_handler: Arc<dyn Handler>) -> Self {
self.host_handler = Some(host_handler);
self
}
pub fn host_tools(mut self, host_tools: Vec<HostToolSpec>) -> Self {
self.host_tools = host_tools;
self
}
pub fn http_client(mut self, http_client: reqwest::Client) -> Self {
self.http_client = Some(http_client);
self
}
pub fn sql_path(mut self, sql_path: impl Into<PathBuf>) -> Self {
self.sql_path = Some(sql_path.into());
self
}
pub fn kv_path(mut self, kv_path: impl Into<PathBuf>) -> Self {
self.kv_path = Some(kv_path.into());
self
}
pub fn ts_path(mut self, ts_path: impl Into<PathBuf>) -> Self {
self.ts_path = Some(ts_path.into());
self
}
pub fn extra_globals(mut self, extra_globals: HashMap<String, serde_json::Value>) -> Self {
self.extra_globals = extra_globals;
self
}
pub fn auto_serve_bus(mut self, auto_serve_bus: bool) -> Self {
self.auto_serve_bus = auto_serve_bus;
self
}
pub fn shutdown_token(mut self, shutdown_token: CancellationToken) -> Self {
self.shutdown_token = Some(shutdown_token);
self
}
pub fn build(self) -> BlockConfig {
BlockConfig {
script: self.script,
project_root: self.project_root,
relay_url: self.relay_url,
secret_key: self.secret_key,
mcp_rpc_timeout: self.mcp_rpc_timeout,
prompt: self.prompt,
context: self.context,
host_handlers: self.host_handlers,
host_handler: self.host_handler,
host_tools: self.host_tools,
http_client: self.http_client,
sql_path: self.sql_path,
kv_path: self.kv_path,
ts_path: self.ts_path,
extra_globals: self.extra_globals,
auto_serve_bus: self.auto_serve_bus,
shutdown_token: self.shutdown_token,
}
}
}
#[cfg(feature = "sqlite")]
#[derive(Clone)]
pub struct SqliteConn {
pub conn: Arc<Mutex<rusqlite::Connection>>,
pub interrupt: Arc<rusqlite::InterruptHandle>,
}
#[derive(Clone)]
pub struct HostContext {
pub project_root: PathBuf,
#[cfg(feature = "mesh")]
pub mesh_agent: Option<Arc<agent_mesh_sdk::MeshAgent>>,
pub mcp_manager: Arc<RwLock<McpManager>>,
pub http_client: reqwest::Client,
#[cfg(feature = "sqlite")]
pub sql_conn: SqliteConn,
#[cfg(feature = "sqlite")]
pub kv_conn: SqliteConn,
#[cfg(feature = "sqlite")]
pub ts_isle: rusqlite_isle::AsyncIsle,
#[allow(dead_code)]
pub isle: Arc<AsyncIsle>,
pub handler_isle: Arc<AsyncIsle>,
#[allow(dead_code)]
pub bus_tx: mpsc::Sender<Event>,
pub event_bus: Arc<Mutex<Option<EventBus>>>,
pub fs_snapshots: crate::bridge::fs::SnapshotStore,
pub knl_drivers: crate::knl::IsleDrivers,
pub knl_store: PathBuf,
}
impl HostContext {
#[cfg(feature = "mesh")]
pub fn mesh_agent_id(&self) -> Option<String> {
self.mesh_agent.as_ref().map(|a| a.agent_id().to_string())
}
#[cfg(not(feature = "mesh"))]
pub fn mesh_agent_id(&self) -> Option<String> {
None
}
}
#[cfg(feature = "sqlite")]
fn prepare_sqlite_dir(path: &Path, label: &'static str) -> BlockResult<bool> {
let is_memory = crate::bridge::config::is_memory_sql(path);
if !is_memory {
if let Some(parent) = path.parent() {
std::fs::create_dir_all(parent)
.map_err(|e| BlockError::Runtime(format!("{label} dir create: {e}")))?;
}
}
Ok(is_memory)
}
fn prepare_knl_dir(path: &Path) -> BlockResult<()> {
let Some(parent) = path.parent().filter(|p| !p.as_os_str().is_empty()) else {
return Ok(());
};
std::fs::create_dir_all(parent)
.map_err(|e| BlockError::Runtime(format!("knl dir create {}: {e}", parent.display())))
}
#[cfg(feature = "sqlite")]
async fn open_sqlite_isle<F>(
path: &Path,
label: &'static str,
init: F,
) -> BlockResult<(rusqlite_isle::AsyncIsle, rusqlite_isle::AsyncIsleDriver)>
where
F: FnOnce(&mut rusqlite::Connection) -> Result<(), rusqlite::Error> + Send + 'static,
{
let is_memory = prepare_sqlite_dir(path, label)?;
let busy = crate::bridge::config::sql_busy_timeout();
let journal = crate::bridge::config::sql_journal_mode();
let (isle, driver) = rusqlite_isle::AsyncIsle::builder()
.thread_name(label)
.busy_timeout(busy)
.spawn(path, move |conn| {
if !is_memory {
conn.pragma_update(None, "journal_mode", &journal)?;
}
init(conn)
})
.await
.map_err(|e| BlockError::Runtime(format!("sqlite open {}: {e}", path.display())))?;
info!(label, path = %path.display(), busy_ms = busy.as_millis() as i64, "sqlite initialized");
Ok((isle, driver))
}
#[cfg(feature = "sqlite")]
fn open_sqlite_conn<F>(path: &Path, label: &'static str, init: F) -> BlockResult<SqliteConn>
where
F: FnOnce(&rusqlite::Connection) -> Result<(), rusqlite::Error>,
{
let is_memory = prepare_sqlite_dir(path, label)?;
let busy = crate::bridge::config::sql_busy_timeout();
let journal = crate::bridge::config::sql_journal_mode();
let open = || -> Result<rusqlite::Connection, rusqlite::Error> {
let conn = rusqlite::Connection::open(path)?;
conn.busy_timeout(busy)?;
conn.pragma_update(None, "synchronous", "NORMAL")?;
if !is_memory {
conn.pragma_update(None, "journal_mode", &journal)?;
}
init(&conn)?;
Ok(conn)
};
let conn =
open().map_err(|e| BlockError::Runtime(format!("sqlite open {}: {e}", path.display())))?;
let interrupt = Arc::new(conn.get_interrupt_handle());
info!(label, path = %path.display(), busy_ms = busy.as_millis() as i64, "sqlite initialized");
Ok(SqliteConn {
conn: Arc::new(Mutex::new(conn)),
interrupt,
})
}
fn build_isle_init(
script_name: String,
script_dir: String,
lib_paths: String,
lib_roots: Vec<PathBuf>,
prompt: Option<String>,
context: Option<String>,
extra_globals: HashMap<String, serde_json::Value>,
) -> impl FnOnce(&mlua::Lua) -> mlua::Result<()> + Send + 'static {
move |lua| {
lua.globals().set("_SCRIPT_NAME", script_name.as_str())?;
if let Some(ref p) = prompt {
lua.globals().set("_PROMPT", p.as_str())?;
}
if let Some(ref c) = context {
lua.globals().set("_CONTEXT", c.as_str())?;
}
mlua_batteries::register_all(lua, "std")?;
mlua_batteries::async_overrides::register_by_name(lua, "std")?;
for (name, value) in &extra_globals {
let lua_value = crate::bridge::json_to_lua(lua, value.clone())
.map_err(|e| mlua::Error::external(format!("extra_globals[{name}]: {e}")))?;
lua.globals().set(name.as_str(), lua_value)?;
}
let package: mlua::Table = lua.globals().get("package")?;
let current_path: String = package.get("path")?;
let new_path =
format!("{script_dir}/?.lua;{script_dir}/?/init.lua;{lib_paths}{current_path}");
package.set("path", new_path)?;
let mut registry = mlua_pkg::Registry::new();
let mut aliases = mlua_pkg::resolvers::MemoryResolver::new();
for (name, source) in EMBEDDED_BLOCKS.iter().chain(EMBEDDED_LIBS.iter()) {
aliases = aliases.add(format!("{EMBEDDED_ALIAS_PREFIX}{name}"), *source);
}
aliases = aliases.add(
format!("{EMBEDDED_ALIAS_PREFIX}knl_types"),
crate::bridge::knl::lshape_module_source(),
);
registry.add(aliases);
let mut fs_roots: Vec<PathBuf> = vec![PathBuf::from(&script_dir)];
fs_roots.extend(lib_roots.iter().cloned());
for root in fs_roots {
match mlua_pkg::resolvers::FsResolver::new_symlink_aware(root.clone()) {
Ok(resolver) => {
registry.add(resolver);
}
Err(e) => {
warn!(root = %root.display(), error = %e, "FsResolver init skipped");
}
}
}
let mut memory = mlua_pkg::resolvers::MemoryResolver::new();
for (name, source) in EMBEDDED_BLOCKS.iter().chain(EMBEDDED_LIBS.iter()) {
memory = memory.add(*name, *source);
}
memory = memory.add("knl_types", crate::bridge::knl::lshape_module_source());
registry.add(memory);
registry
.install(lua)
.map_err(|e| mlua::Error::external(format!("require registry install failed: {e}")))?;
Ok(())
}
}
async fn spawn_handler_isle(
script_name: String,
script_dir: String,
lib_paths: String,
lib_roots: Vec<PathBuf>,
prompt: Option<String>,
context: Option<String>,
extra_globals: HashMap<String, serde_json::Value>,
) -> BlockResult<(Arc<AsyncIsle>, AsyncIsleDriver)> {
let init = build_isle_init(
script_name,
script_dir,
lib_paths,
lib_roots,
prompt,
context,
extra_globals,
);
let (isle, driver) = AsyncIsle::builder()
.thread_name("agent-block-handler-isle")
.spawn(init)
.await
.map_err(|e| BlockError::Runtime(format!("handler isle spawn failed: {e}")))?;
info!(
thread_name = "agent-block-handler-isle",
"handler Isle spawned"
);
Ok((Arc::new(isle), driver))
}
#[cfg(feature = "mesh")]
fn hex_decode_32(s: &str) -> Result<[u8; 32], String> {
let s = s.trim();
if s.len() != 64 {
return Err(format!("expected 64 hex chars, got {}", s.len()));
}
let mut out = [0u8; 32];
for (i, byte) in out.iter_mut().enumerate() {
let hi = u8::from_str_radix(&s[2 * i..2 * i + 1], 16)
.map_err(|e| format!("invalid hex at position {}: {e}", 2 * i))?;
let lo = u8::from_str_radix(&s[2 * i + 1..2 * i + 2], 16)
.map_err(|e| format!("invalid hex at position {}: {e}", 2 * i + 1))?;
*byte = (hi << 4) | lo;
}
Ok(out)
}
struct ResolvedSources {
script_source: String,
script_name: String,
script_dir: PathBuf,
prompt: Option<String>,
context: Option<String>,
secret_key: Option<String>,
}
fn resolve_sources(config: &BlockConfig) -> BlockResult<ResolvedSources> {
let (script_source, script_name, script_dir) = match &config.script {
ScriptSource::Path(p) => {
let source = std::fs::read_to_string(p)
.map_err(|e| BlockError::Script(format!("{}: {e}", p.display())))?;
let name = p
.file_name()
.map(|n| n.to_string_lossy().to_string())
.unwrap_or_else(|| "unknown".to_string());
let dir = p
.parent()
.map(|d| d.to_path_buf())
.unwrap_or_else(|| PathBuf::from("."));
(source, name, dir)
}
ScriptSource::Inline { source, name } => {
(source.clone(), name.clone(), config.project_root.clone())
}
ScriptSource::DefaultAgent => (
DEFAULT_AGENT_INVOKER.to_string(),
"default_agent_invoker.lua".to_string(),
config.project_root.clone(),
),
};
let prompt: Option<String> = match &config.prompt {
Some(PromptSource::Inline(s)) => Some(s.clone()),
Some(PromptSource::File(p)) => Some(
std::fs::read_to_string(p)
.map_err(|e| BlockError::Script(format!("prompt file {}: {e}", p.display())))?,
),
None => None,
};
let context: Option<String> = match &config.context {
Some(PromptSource::Inline(s)) => Some(s.clone()),
Some(PromptSource::File(p)) => Some(
std::fs::read_to_string(p)
.map_err(|e| BlockError::Script(format!("context file {}: {e}", p.display())))?,
),
None => None,
};
let secret_key: Option<String> = match &config.secret_key {
Some(SecretKeySource::Inline(s)) => Some(s.clone()),
Some(SecretKeySource::Env(var)) => std::env::var(var).ok(),
None => None,
};
Ok(ResolvedSources {
script_source,
script_name,
script_dir,
prompt,
context,
secret_key,
})
}
fn load_dotenv(project_root: &Path) {
let env_path = project_root.join(".env");
match dotenvy::from_path(&env_path) {
Ok(()) => info!(path = %env_path.display(), ".env loaded"),
Err(dotenvy::Error::Io(_)) => {} Err(e) => tracing::warn!(path = %env_path.display(), error = %e, ".env parse error"),
}
}
type AutoServeState = Option<(tokio::task::JoinHandle<()>, CancellationToken)>;
struct BusSetup {
event_bus: Arc<Mutex<Option<EventBus>>>,
bus_tx: mpsc::Sender<Event>,
auto_serve_state: AutoServeState,
}
fn setup_event_bus(config: &BlockConfig) -> BlockResult<BusSetup> {
let bus_capacity = crate::bridge::config::bus_capacity();
let (bus_tx, bus_rx) = mpsc::channel::<Event>(bus_capacity);
let event_bus = Arc::new(Mutex::new(Some(EventBus::new(bus_rx))));
let has_kind_handlers = !config.host_handlers.is_empty();
let has_any_handler = config.host_handler.is_some();
if has_kind_handlers || has_any_handler {
let mut guard = event_bus
.lock()
.map_err(|_| BlockError::Bus("event_bus mutex poisoned".into()))?;
let bus = guard
.as_mut()
.ok_or_else(|| BlockError::Bus("event_bus already taken".into()))?;
for (kind, handler) in &config.host_handlers {
bus.on(kind.clone(), Arc::clone(handler))
.map_err(|e| BlockError::Bus(format!("host_handlers on({kind}): {e}")))?;
}
if let Some(any_handler) = &config.host_handler {
bus.on_any(Arc::clone(any_handler))
.map_err(|e| BlockError::Bus(format!("host_handler on_any: {e}")))?;
}
info!(
kind_handlers = config.host_handlers.len(),
any_handler = has_any_handler,
"host handlers pre-installed"
);
}
let auto_serve = config.auto_serve_bus && (has_kind_handlers || has_any_handler);
let auto_serve_state: AutoServeState = if auto_serve {
let bus = {
let mut guard = event_bus
.lock()
.map_err(|_| BlockError::Bus("event_bus mutex poisoned".into()))?;
guard
.take()
.ok_or_else(|| BlockError::Bus("event_bus already taken".into()))?
};
let token = CancellationToken::new();
let token_for_task = token.clone();
let handle = tokio::spawn(async move {
let mut bus = bus;
if let Err(e) = bus.run(token_for_task).await {
tracing::error!(error = %e, "auto-serve: dispatcher loop returned error");
}
});
info!("auto-serve: dispatcher spawned");
Some((handle, token))
} else {
None
};
Ok(BusSetup {
event_bus,
bus_tx,
auto_serve_state,
})
}
#[cfg(feature = "mesh")]
async fn connect_mesh(
relay_url: Option<&String>,
secret_key: Option<&String>,
bus_tx: &mpsc::Sender<Event>,
) -> BlockResult<Option<Arc<agent_mesh_sdk::MeshAgent>>> {
let Some(relay_url) = relay_url else {
return Ok(None);
};
let keypair = match secret_key {
Some(hex_str) => {
let bytes = hex_decode_32(hex_str)
.map_err(|e| BlockError::Runtime(format!("secret-key: {e}")))?;
agent_mesh_core::identity::AgentKeypair::from_bytes(&bytes)
}
None => agent_mesh_core::identity::AgentKeypair::generate(),
};
info!(agent_id = %keypair.agent_id(), "mesh identity");
let acl = agent_mesh_core::acl::AclPolicy {
default_deny: false,
rules: vec![],
};
let handler: Arc<dyn agent_mesh_sdk::RequestHandler> =
Arc::new(BusRelayHandler::new(bus_tx.clone()));
let url = relay_url.clone();
let agent = agent_mesh_sdk::MeshAgent::connect(keypair, &url, acl, handler)
.await
.map_err(|e| BlockError::Mesh(format!("connect to {relay_url} failed: {e}")))?;
info!(relay_url = %relay_url, "mesh connected");
Ok(Some(Arc::new(agent)))
}
#[cfg(feature = "sqlite")]
struct SqliteConns {
sql: SqliteConn,
kv: SqliteConn,
ts_isle: rusqlite_isle::AsyncIsle,
drivers: SqliteDrivers,
}
#[cfg(feature = "sqlite")]
struct SqliteDrivers {
ts: rusqlite_isle::AsyncIsleDriver,
}
#[cfg(feature = "sqlite")]
async fn init_sqlite(config: &BlockConfig) -> BlockResult<SqliteConns> {
let sql_path = match &config.sql_path {
Some(p) => p.clone(),
None => crate::bridge::config::sql_path().map_err(BlockError::Runtime)?,
};
let sql = open_sqlite_conn(&sql_path, "sql", |_| Ok(()))?;
let kv_path = match &config.kv_path {
Some(p) => p.clone(),
None => crate::bridge::config::kv_path().map_err(BlockError::Runtime)?,
};
let kv = open_sqlite_conn(&kv_path, "kv", mlua_batteries_sqlite::kv::init_schema)?;
let ts_path = match &config.ts_path {
Some(p) => p.clone(),
None => crate::bridge::config::ts_path().map_err(BlockError::Runtime)?,
};
let (ts_isle, ts_driver) = open_sqlite_isle(&ts_path, "ts", |conn| {
conn.execute_batch(crate::bridge::ts::SCHEMA_DDL)
})
.await?;
Ok(SqliteConns {
sql,
kv,
ts_isle,
drivers: SqliteDrivers { ts: ts_driver },
})
}
struct SpawnedIsles {
isle: Arc<AsyncIsle>,
driver: AsyncIsleDriver,
handler_isle: Arc<AsyncIsle>,
handler_driver: AsyncIsleDriver,
}
async fn spawn_isles(
script_name: &str,
script_dir: &str,
lib_paths: &str,
lib_roots: &[PathBuf],
prompt: Option<String>,
context: Option<String>,
extra_globals: &HashMap<String, serde_json::Value>,
) -> BlockResult<SpawnedIsles> {
let (isle, driver) = AsyncIsle::spawn(build_isle_init(
script_name.to_string(),
script_dir.to_string(),
lib_paths.to_string(),
lib_roots.to_vec(),
prompt.clone(),
context.clone(),
extra_globals.clone(),
))
.await
.map_err(|e| BlockError::Runtime(format!("AsyncIsle spawn failed: {e}")))?;
let isle = Arc::new(isle);
let (handler_isle, handler_driver) = spawn_handler_isle(
script_name.to_string(),
script_dir.to_string(),
lib_paths.to_string(),
lib_roots.to_vec(),
prompt,
context,
extra_globals.clone(),
)
.await?;
Ok(SpawnedIsles {
isle,
driver,
handler_isle,
handler_driver,
})
}
async fn register_bridges(
ctx: &HostContext,
isle: &Arc<AsyncIsle>,
handler_isle: &Arc<AsyncIsle>,
) -> BlockResult<()> {
{
let ctx = ctx.clone();
isle.exec(move |lua| {
bridge::register_all(lua, &ctx)
.map_err(|e| mlua_isle::IsleError::Lua(format!("bridge register failed: {e}")))?;
Ok(String::new())
})
.await
.map_err(|e| BlockError::Runtime(format!("bridge register: {e}")))?;
}
{
let ctx = ctx.clone();
handler_isle
.exec(move |lua| {
bridge::register_all_handler_side(lua, &ctx).map_err(|e| {
mlua_isle::IsleError::Lua(format!("handler bridge register failed: {e}"))
})?;
Ok(String::new())
})
.await
.map_err(|e| BlockError::Runtime(format!("handler bridge register: {e}")))?;
}
Ok(())
}
async fn inject_host_tools(isle: &Arc<AsyncIsle>, host_tools: &[HostToolSpec]) -> BlockResult<()> {
if host_tools.is_empty() {
return Ok(());
}
let host_tools = host_tools.to_vec();
let tool_count = host_tools.len();
isle.exec(move |lua| {
let registry: mlua::Table = lua
.globals()
.get("_TOOL_REGISTRY")
.map_err(|e| mlua_isle::IsleError::Lua(format!("get _TOOL_REGISTRY: {e}")))?;
for tool in host_tools {
let entry = lua
.create_table()
.map_err(|e| mlua_isle::IsleError::Lua(format!("create entry: {e}")))?;
entry
.set("name", tool.name.as_str())
.map_err(|e| mlua_isle::IsleError::Lua(format!("set name: {e}")))?;
let schema = lua
.create_table()
.map_err(|e| mlua_isle::IsleError::Lua(format!("create schema: {e}")))?;
schema
.set("description", tool.description.as_str())
.map_err(|e| mlua_isle::IsleError::Lua(format!("set description: {e}")))?;
let input_schema_lua = crate::bridge::json_to_lua(lua, tool.input_schema.clone())
.map_err(|e| mlua_isle::IsleError::Lua(format!("input_schema: {e}")))?;
schema
.set("input_schema", input_schema_lua)
.map_err(|e| mlua_isle::IsleError::Lua(format!("set input_schema: {e}")))?;
entry
.set("schema", schema)
.map_err(|e| mlua_isle::IsleError::Lua(format!("set schema: {e}")))?;
if let Some(group) = &tool.group {
entry
.set("group", group.as_str())
.map_err(|e| mlua_isle::IsleError::Lua(format!("set group: {e}")))?;
}
let handler_arc = Arc::clone(&tool.handler);
let handler_fn = lua
.create_async_function(move |lua, input: mlua::Value| {
let handler = Arc::clone(&handler_arc);
async move {
let input_json = crate::bridge::lua_to_json(&lua, input)?;
let result = handler
.call(input_json)
.await
.map_err(mlua::Error::external)?;
crate::bridge::json_to_lua(&lua, result)
}
})
.map_err(|e| mlua_isle::IsleError::Lua(format!("create handler: {e}")))?;
entry
.set("handler", handler_fn)
.map_err(|e| mlua_isle::IsleError::Lua(format!("set handler: {e}")))?;
registry
.set(tool.name.as_str(), entry)
.map_err(|e| mlua_isle::IsleError::Lua(format!("registry set: {e}")))?;
}
Ok(String::new())
})
.await
.map_err(|e| BlockError::Runtime(format!("host_tools inject: {e}")))?;
info!(count = tool_count, "host tools injected into Lua registry");
Ok(())
}
async fn execute_script(
isle: &Arc<AsyncIsle>,
script_source: &str,
script_name: &str,
shutdown_token: Option<&CancellationToken>,
) -> BlockResult<String> {
let _exec_span = info_span!("execute", script = %script_name);
let mut task = isle.spawn_coroutine_eval(script_source);
let task_cancel = task.cancel_token().clone();
match shutdown_token {
Some(token) => {
tokio::select! {
biased;
_ = token.cancelled() => {
task_cancel.cancel();
let _ = (&mut task).await;
info!("shutdown_token: cancelled by caller");
Err(BlockError::Cancelled)
}
res = &mut task => res.map_err(|e| BlockError::Script(format!("{e}"))),
}
}
None => (&mut task)
.await
.map_err(|e| BlockError::Script(format!("{e}"))),
}
}
async fn drain_auto_serve(auto_serve_state: AutoServeState) {
if let Some((handle, token)) = auto_serve_state {
let grace_ms = crate::bridge::config::task_grace_ms();
let grace = Duration::from_millis(grace_ms);
tokio::time::sleep(grace).await;
token.cancel();
match tokio::time::timeout(grace, handle).await {
Ok(Ok(())) => info!("auto-serve: dispatcher shut down cleanly"),
Ok(Err(join_err)) => {
tracing::error!(error = %join_err, "auto-serve: dispatcher task join error");
}
Err(_) => {
tracing::warn!(
grace_ms,
"auto-serve: dispatcher join timed out after cancel; forcing exit"
);
}
}
}
}
async fn shutdown(
mcp_manager: &Arc<RwLock<McpManager>>,
driver: AsyncIsleDriver,
handler_driver: AsyncIsleDriver,
knl_drivers: crate::knl::IsleDrivers,
#[cfg(feature = "sqlite")] sqlite_drivers: SqliteDrivers,
) -> BlockResult<()> {
let _shutdown_span = info_span!("shutdown");
let mut failure: Option<BlockError> = None;
let mut record = |e: BlockError| {
tracing::error!(error = %e, "shutdown step failed; the teardown continues");
if failure.is_none() {
failure = Some(e);
}
};
if let Err(e) = mcp_manager.write().await.disconnect_all().await {
record(e);
}
if let Err(e) = driver.shutdown().await {
record(BlockError::Runtime(format!(
"AsyncIsle shutdown failed: {e}"
)));
}
match handler_driver.shutdown().await {
Ok(()) => info!(
thread_name = "agent-block-handler-isle",
"handler Isle shut down"
),
Err(e) => tracing::error!(
error = %e,
thread_name = "agent-block-handler-isle",
"handler Isle shutdown failed"
),
}
{
let count = knl_drivers.len();
let failures = knl_drivers.shutdown().await;
if failures.is_empty() {
if count > 0 {
info!(count, "knl session connection threads shut down");
}
} else {
for e in &failures {
tracing::error!(error = %e, "knl session connection thread shutdown failed");
}
}
}
#[cfg(feature = "sqlite")]
{
let SqliteDrivers { ts } = sqlite_drivers;
match ts.shutdown().await {
Ok(()) => info!(label = "ts", "sqlite connection thread shut down"),
Err(e) => {
tracing::error!(error = %e, label = "ts", "sqlite shutdown failed")
}
}
}
match failure {
Some(e) => Err(e),
None => Ok(()),
}
}
pub async fn run(config: BlockConfig) -> BlockResult<()> {
run_capture(config).await.map(|_| ())
}
pub async fn run_capture(config: BlockConfig) -> BlockResult<String> {
let ResolvedSources {
script_source,
script_name,
script_dir: script_dir_pathbuf,
prompt: prompt_resolved,
context: context_resolved,
secret_key: secret_key_resolved,
} = resolve_sources(&config)?;
let _root_span = info_span!("agent_block", script = %script_name);
load_dotenv(&config.project_root);
let _init_span = info_span!("init");
let BusSetup {
event_bus,
bus_tx,
auto_serve_state,
} = setup_event_bus(&config)?;
#[cfg(feature = "mesh")]
let mesh_agent = connect_mesh(
config.relay_url.as_ref(),
secret_key_resolved.as_ref(),
&bus_tx,
)
.await?;
#[cfg(not(feature = "mesh"))]
let _ = (&secret_key_resolved, &config.relay_url);
let mcp_manager = Arc::new(RwLock::new(McpManager::with_rpc_timeout(
config.mcp_rpc_timeout,
)?));
let project_root = config
.project_root
.canonicalize()
.or_else(|_| std::env::current_dir().map(|cwd| cwd.join(&config.project_root)))?;
let http_client = config.http_client.clone().unwrap_or_default();
#[cfg(feature = "sqlite")]
let SqliteConns {
sql: sql_conn,
kv: kv_conn,
ts_isle,
drivers: sqlite_drivers,
} = init_sqlite(&config).await?;
let knl_store = crate::bridge::config::knl_path(&project_root).map_err(BlockError::Runtime)?;
prepare_knl_dir(&knl_store)?;
let script_dir = script_dir_pathbuf.to_string_lossy().to_string();
let lib_roots = lib_roots(&project_root);
let lib_paths = package_path_prefix(&lib_roots);
let require_roots: Vec<PathBuf> = std::iter::once(script_dir_pathbuf.clone())
.chain(lib_roots.iter().cloned())
.collect();
check_sealed_modules(&require_roots)?;
let prompt = prompt_resolved.clone();
let context = context_resolved.clone();
let SpawnedIsles {
isle,
driver,
handler_isle,
handler_driver,
} = spawn_isles(
&script_name,
&script_dir,
&lib_paths,
&lib_roots,
prompt,
context,
&config.extra_globals,
)
.await?;
{
let mut mgr = mcp_manager.write().await;
mgr.set_handler_isle(Arc::clone(&handler_isle));
mgr.set_main_isle(Arc::clone(&isle));
}
let ctx = HostContext {
project_root,
#[cfg(feature = "mesh")]
mesh_agent,
mcp_manager: Arc::clone(&mcp_manager),
http_client,
#[cfg(feature = "sqlite")]
sql_conn,
#[cfg(feature = "sqlite")]
kv_conn,
#[cfg(feature = "sqlite")]
ts_isle,
isle: Arc::clone(&isle),
handler_isle: Arc::clone(&handler_isle),
bus_tx: bus_tx.clone(),
event_bus: Arc::clone(&event_bus),
fs_snapshots: Default::default(),
knl_drivers: crate::knl::IsleDrivers::new(),
knl_store,
};
let knl_drivers = ctx.knl_drivers.clone();
register_bridges(&ctx, &isle, &handler_isle).await?;
inject_host_tools(&isle, &config.host_tools).await?;
drop(_init_span);
let script_result = execute_script(
&isle,
&script_source,
&script_name,
config.shutdown_token.as_ref(),
)
.await;
drain_auto_serve(auto_serve_state).await;
shutdown(
&mcp_manager,
driver,
handler_driver,
knl_drivers,
#[cfg(feature = "sqlite")]
sqlite_drivers,
)
.await?;
script_result
}
#[cfg(feature = "mesh")]
struct BusRelayHandler {
tx: mpsc::Sender<Event>,
}
#[cfg(feature = "mesh")]
impl BusRelayHandler {
fn new(tx: mpsc::Sender<Event>) -> Self {
Self { tx }
}
}
#[cfg(feature = "mesh")]
const BUS_ACK_TIMEOUT: Duration = Duration::from_secs(30);
#[cfg(feature = "mesh")]
#[async_trait::async_trait]
impl agent_mesh_sdk::RequestHandler for BusRelayHandler {
async fn handle(
&self,
from: &agent_mesh_core::identity::AgentId,
payload: &serde_json::Value,
_cancel: agent_mesh_sdk::CancelToken,
) -> serde_json::Value {
let id = uuid::Uuid::new_v4().to_string();
let meta = serde_json::json!({"from": from.to_string()});
let (ack_tx, ack_rx) = tokio::sync::oneshot::channel();
let event = Event {
kind: "mesh".into(),
id: id.clone(),
payload: payload.clone(),
meta,
ack_tx: Some(ack_tx),
};
if let Err(e) = self.tx.send(event).await {
tracing::error!(error = %e, id = %id, "bus channel closed; rejecting mesh request");
return serde_json::json!({"error": "bus channel closed"});
}
match tokio::time::timeout(BUS_ACK_TIMEOUT, ack_rx).await {
Ok(Ok(Ok(v))) => v,
Ok(Ok(Err(e))) => {
tracing::error!(id = %id, error = %e, "mesh handler returned error");
serde_json::json!({"error": e.to_string()})
}
Ok(Err(e)) => {
tracing::error!(id = %id, error = %e, "mesh ack receiver dropped");
serde_json::json!({"error": "ack dropped"})
}
Err(_) => {
tracing::error!(id = %id, timeout_secs = BUS_ACK_TIMEOUT.as_secs(), "mesh handler timeout");
serde_json::json!({"error": "handler timeout"})
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn sealed_list_matches_the_readme() {
assert_eq!(
SEALED,
[
"knl",
"knl_adapter",
"knl_types",
"lshape",
"lshape.t",
"lshape.check",
"lshape.reflect",
"lshape.luacats",
]
.as_slice()
);
}
#[test]
fn every_sealed_name_is_an_embedded_module() {
for name in SEALED {
let embedded = *name == "knl_types"
|| EMBEDDED_BLOCKS
.iter()
.chain(EMBEDDED_LIBS.iter())
.any(|(n, _)| n == name);
assert!(embedded, "sealed name `{name}` is not an embedded module");
}
}
#[test]
fn require_candidates_match_the_fs_resolver_layout() {
let root = Path::new("/p/blocks");
assert_eq!(
require_candidates(root, "knl"),
[
PathBuf::from("/p/blocks/knl.lua"),
PathBuf::from("/p/blocks/knl/init.lua"),
]
);
assert_eq!(
require_candidates(root, "lshape.t"),
[
PathBuf::from("/p/blocks/lshape/t.lua"),
PathBuf::from("/p/blocks/lshape/t/init.lua"),
]
);
}
#[test]
fn an_empty_root_is_not_a_shadow() {
let tmp = tempfile::tempdir().expect("tempdir");
check_sealed_modules(&[tmp.path().to_path_buf()]).expect("nothing to seal against");
}
#[test]
fn a_shadowing_file_is_refused_by_name_and_path() {
let tmp = tempfile::tempdir().expect("tempdir");
let knl = tmp.path().join("knl");
std::fs::create_dir_all(&knl).expect("mkdir");
std::fs::write(knl.join("init.lua"), "return {}").expect("write");
let err = check_sealed_modules(&[tmp.path().to_path_buf()])
.expect_err("a sealed module was shadowed");
let msg = err.to_string();
assert!(msg.contains("knl"), "the module is not named: {msg}");
assert!(
msg.contains(&knl.join("init.lua").display().to_string()),
"the file is not named: {msg}"
);
assert!(
msg.contains("AGENT_BLOCK_UNSEAL"),
"the escape hatch is not mentioned: {msg}"
);
}
#[test]
fn a_shadowing_sub_module_is_refused() {
let tmp = tempfile::tempdir().expect("tempdir");
let lshape = tmp.path().join("lshape");
std::fs::create_dir_all(&lshape).expect("mkdir");
std::fs::write(lshape.join("t.lua"), "return {}").expect("write");
let err = check_sealed_modules(&[tmp.path().to_path_buf()])
.expect_err("a sealed sub-module was shadowed");
assert!(
err.to_string().contains("lshape.t"),
"the sub-module is not named: {err}"
);
}
#[test]
fn shadowing_an_unsealed_block_is_allowed() {
let tmp = tempfile::tempdir().expect("tempdir");
let agent = tmp.path().join("agent");
std::fs::create_dir_all(&agent).expect("mkdir");
std::fs::write(agent.join("init.lua"), "return {}").expect("write");
check_sealed_modules(&[tmp.path().to_path_buf()]).expect("`agent` is not sealed");
}
}