use crate::worker::adapter::{InProcSpawner, WorkerError, WorkerInvocation, WorkerResult};
use agent_block_core::bus::dispatcher::Handler;
use agent_block_core::host::{PromptSource, ScriptSource};
use agent_block_core::{run, BlockConfig};
use agent_block_types::error::BlockError;
use async_trait::async_trait;
use serde_json::Value;
use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::{Arc, Mutex};
use std::time::Duration;
use tokio::sync::oneshot;
struct WorkerResultCaptor {
tx: Mutex<Option<oneshot::Sender<WorkerResult>>>,
}
impl WorkerResultCaptor {
fn extract(payload: &Value) -> (Value, bool) {
let ok = payload.get("ok").and_then(|v| v.as_bool()).unwrap_or(true);
let value = payload
.get("content")
.cloned()
.or_else(|| payload.get("response").cloned())
.unwrap_or_else(|| payload.clone());
(value, ok)
}
}
#[async_trait]
impl Handler for WorkerResultCaptor {
async fn call(
&self,
_kind: String,
_id: String,
payload: Value,
_meta: Value,
) -> Result<Value, BlockError> {
let (value, ok) = Self::extract(&payload);
let wr = WorkerResult { value, ok };
if let Ok(mut guard) = self.tx.lock() {
if let Some(tx) = guard.take() {
let _ = tx.send(wr);
}
}
Ok(Value::Null)
}
}
#[derive(Clone)]
struct AgentBlockSettings {
script: ScriptSource,
project_root: PathBuf,
mcp_rpc_timeout: Duration,
profile_context: Option<String>,
}
async fn run_agent_block_worker(
settings: Arc<AgentBlockSettings>,
inv: WorkerInvocation,
) -> Result<WorkerResult, WorkerError> {
let (tx, rx) = oneshot::channel();
let captor: Arc<dyn Handler> = Arc::new(WorkerResultCaptor {
tx: Mutex::new(Some(tx)),
});
let shutdown_token = inv.cancel_token.clone().unwrap_or_default();
let config = BlockConfig {
script: settings.script.clone(),
project_root: settings.project_root.clone(),
relay_url: None,
secret_key: None,
mcp_rpc_timeout: settings.mcp_rpc_timeout,
prompt: Some(PromptSource::Inline(inv.prompt)),
context: settings.profile_context.clone().map(PromptSource::Inline),
host_handlers: HashMap::new(),
host_handler: Some(captor),
auto_serve_bus: true,
shutdown_token: Some(shutdown_token.clone()),
};
let run_handle = tokio::spawn(run(config));
let run_result = run_handle
.await
.map_err(|e| WorkerError::Failed(format!("agent-block task join: {e}")))?;
run_result.map_err(|e| WorkerError::Failed(format!("agent-block run failed: {e}")))?;
rx.await.map_err(|_| {
WorkerError::Failed("agent-block script finished without emitting result via bus".into())
})
}
pub fn resolve_needed_mcp_servers(
profile_tools: &[String],
spec_mcp_servers: &[Value],
) -> Vec<Value> {
use std::collections::HashSet;
let needed: HashSet<&str> = profile_tools
.iter()
.filter_map(|t| {
let rest = t.strip_prefix("mcp__")?;
let idx = rest.find("__")?;
Some(&rest[..idx])
})
.collect();
spec_mcp_servers
.iter()
.filter(|cfg| {
cfg.get("name")
.and_then(|n| n.as_str())
.map(|name| needed.contains(name))
.unwrap_or(false)
})
.cloned()
.collect()
}
pub fn build_inline_agent_invoker(mcp_servers: &[Value]) -> ScriptSource {
let mcp_lua = json_array_to_lua_literal(mcp_servers);
let source = format!(
r##"local agent = require("agent")
local mcp_servers = {mcp_lua}
local r = agent.run({{
prompt = _PROMPT,
system = _CONTEXT,
mcp_servers = mcp_servers,
}})
bus.emit("agent_result", r)
"##
);
ScriptSource::Inline {
source,
name: "mlua_swarm_engine_default_agent_invoker.lua".into(),
}
}
fn json_to_lua_literal(v: &Value) -> String {
match v {
Value::Null => "nil".to_string(),
Value::Bool(b) => b.to_string(),
Value::Number(n) => n.to_string(),
Value::String(s) => format!("{s:?}"),
Value::Array(arr) => {
let items: Vec<String> = arr.iter().map(json_to_lua_literal).collect();
format!("{{{}}}", items.join(", "))
}
Value::Object(map) => {
let items: Vec<String> = map
.iter()
.map(|(k, v)| format!("[{k:?}]={}", json_to_lua_literal(v)))
.collect();
format!("{{{}}}", items.join(", "))
}
}
}
fn json_array_to_lua_literal(arr: &[Value]) -> String {
if arr.is_empty() {
return "{}".to_string();
}
let items: Vec<String> = arr.iter().map(json_to_lua_literal).collect();
format!("{{{}}}", items.join(", "))
}
pub struct AgentBlockInProcessSpawnerFactory;
impl Default for AgentBlockInProcessSpawnerFactory {
fn default() -> Self {
Self
}
}
impl AgentBlockInProcessSpawnerFactory {
pub fn new() -> Self {
Self
}
}
impl crate::blueprint::compiler::SpawnerFactoryKind for AgentBlockInProcessSpawnerFactory {
const KIND: crate::blueprint::AgentKind = crate::blueprint::AgentKind::AgentBlock;
type Worker = AgentBlockWorker;
}
impl crate::blueprint::compiler::SpawnerFactory for AgentBlockInProcessSpawnerFactory {
fn build(
&self,
agent_def: &crate::blueprint::AgentDef,
_hint: Option<&Value>,
) -> Result<
Arc<dyn crate::worker::adapter::SpawnerAdapter>,
crate::blueprint::compiler::CompileError,
> {
let agent_name = agent_def.name.clone();
let spec = &agent_def.spec;
let profile_tools: Vec<String> = agent_def
.profile
.as_ref()
.map(|p| p.tools.clone())
.unwrap_or_default();
let spec_mcp_servers: Vec<Value> = spec
.get("mcp_servers")
.and_then(|v| v.as_array())
.cloned()
.unwrap_or_default();
let needed_mcp_servers = resolve_needed_mcp_servers(&profile_tools, &spec_mcp_servers);
let script = match spec.get("script_path").and_then(|v| v.as_str()) {
Some(s) => ScriptSource::Path(PathBuf::from(s)),
None => build_inline_agent_invoker(&needed_mcp_servers),
};
let project_root = match spec.get("project_root").and_then(|v| v.as_str()) {
Some(s) => PathBuf::from(s),
None => std::env::current_dir().unwrap_or_else(|_| PathBuf::from(".")),
};
let mcp_rpc_timeout = match spec.get("mcp_rpc_timeout_ms").and_then(|v| v.as_u64()) {
Some(ms) => Duration::from_millis(ms),
None => Duration::from_secs(30),
};
let profile_context = agent_def.profile.as_ref().map(|p| p.system_prompt.clone());
let settings = Arc::new(AgentBlockSettings {
script,
project_root,
mcp_rpc_timeout,
profile_context,
});
let worker_fn: crate::worker::adapter::WorkerFn = Arc::new(move |inv| {
let settings = settings.clone();
Box::pin(run_agent_block_worker(settings, inv))
});
let mut sp: InProcSpawner<AgentBlockWorker> = InProcSpawner::<AgentBlockWorker>::typed();
sp.registry.insert(agent_name, worker_fn);
Ok(Arc::new(sp))
}
}
pub struct AgentBlockWorker {
pub handler: crate::worker::WorkerJoinHandler,
}
impl From<crate::worker::WorkerJoinHandler> for AgentBlockWorker {
fn from(handler: crate::worker::WorkerJoinHandler) -> Self {
Self { handler }
}
}
#[async_trait]
impl crate::worker::Worker for AgentBlockWorker {
fn id(&self) -> &crate::types::WorkerId {
&self.handler.worker_id
}
fn cancel_token(&self) -> tokio_util::sync::CancellationToken {
self.handler.cancel.clone()
}
async fn join(self: Box<Self>) -> Result<(), WorkerError> {
self.handler.await_completion().await
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn resolve_needed_mcp_servers_filters_by_tool_prefix() {
let tools = vec![
"mcp__semantic-scholar__search_papers".to_string(),
"mcp__semantic-scholar__get_paper".to_string(),
"Read".to_string(),
"mcp__outline__list_docs".to_string(),
"WebSearch".to_string(),
];
let spec_servers = vec![
serde_json::json!({"name": "semantic-scholar", "command": "ss-mcp", "args": []}),
serde_json::json!({"name": "outline", "command": "outline-mcp", "args": []}),
serde_json::json!({"name": "unused", "command": "nope", "args": []}),
];
let needed = resolve_needed_mcp_servers(&tools, &spec_servers);
assert_eq!(needed.len(), 2, "got: {needed:?}");
let names: Vec<&str> = needed
.iter()
.filter_map(|c| c.get("name").and_then(|n| n.as_str()))
.collect();
assert!(names.contains(&"semantic-scholar"));
assert!(names.contains(&"outline"));
assert!(!names.contains(&"unused"), "unused server is filtered out");
}
#[test]
fn resolve_needed_mcp_servers_returns_empty_when_no_mcp_tools() {
let tools = vec!["Read".to_string(), "WebSearch".to_string()];
let spec_servers =
vec![serde_json::json!({"name": "outline", "command": "outline-mcp", "args": []})];
let needed = resolve_needed_mcp_servers(&tools, &spec_servers);
assert!(
needed.is_empty(),
"no mcp__-prefixed tools → empty result, got: {needed:?}"
);
}
#[test]
fn build_inline_agent_invoker_embeds_mcp_servers_as_lua_literal() {
let servers =
vec![serde_json::json!({"name": "outline", "command": "outline-mcp", "args": []})];
let script = build_inline_agent_invoker(&servers);
match script {
ScriptSource::Inline { source, name } => {
assert!(name.ends_with(".lua"));
assert!(source.contains("require(\"agent\")"));
assert!(source.contains("mcp_servers = mcp_servers"));
assert!(source.contains("bus.emit(\"agent_result\""));
assert!(source.contains("[\"name\"]=\"outline\""));
assert!(source.contains("[\"command\"]=\"outline-mcp\""));
assert!(source.contains("[\"args\"]={}"), "args empty array literal");
}
other => panic!("expected Inline, got: {other:?}"),
}
}
#[test]
fn build_inline_agent_invoker_with_empty_servers_still_valid() {
let script = build_inline_agent_invoker(&[]);
match script {
ScriptSource::Inline { source, .. } => {
assert!(source.contains("local mcp_servers = {}"));
}
other => panic!("expected Inline, got: {other:?}"),
}
}
#[test]
fn json_to_lua_literal_handles_primitives_and_nested() {
assert_eq!(json_to_lua_literal(&serde_json::json!(null)), "nil");
assert_eq!(json_to_lua_literal(&serde_json::json!(true)), "true");
assert_eq!(json_to_lua_literal(&serde_json::json!(42)), "42");
assert_eq!(json_to_lua_literal(&serde_json::json!("hi")), "\"hi\"");
assert_eq!(
json_to_lua_literal(&serde_json::json!(["a", "b"])),
"{\"a\", \"b\"}"
);
assert_eq!(
json_to_lua_literal(&serde_json::json!({"k": 1})),
"{[\"k\"]=1}"
);
}
#[test]
fn extract_prefers_content_then_response_then_whole() {
let p = serde_json::json!({
"content": "Water boils at 100°C",
"messages": [{"role": "assistant"}],
"usage": {"input_tokens": 67, "output_tokens": 29},
"ok": true,
});
let (value, ok) = WorkerResultCaptor::extract(&p);
assert_eq!(value, serde_json::json!("Water boils at 100°C"));
assert!(ok);
let p = serde_json::json!({ "ok": false, "response": {"patch": "..."} });
let (value, ok) = WorkerResultCaptor::extract(&p);
assert_eq!(value, serde_json::json!({"patch": "..."}));
assert!(!ok);
let p = serde_json::json!({ "custom_field": 42 });
let (value, ok) = WorkerResultCaptor::extract(&p);
assert_eq!(value, serde_json::json!({"custom_field": 42}));
assert!(ok); }
#[tokio::test]
async fn captor_emits_worker_result_from_payload() {
let (tx, rx) = oneshot::channel();
let captor = WorkerResultCaptor {
tx: Mutex::new(Some(tx)),
};
let payload = serde_json::json!({ "ok": true, "response": "hello" });
let ack = captor
.call("worker_result".into(), "evt-1".into(), payload, Value::Null)
.await
.expect("handler ack");
assert_eq!(ack, Value::Null);
let wr = rx.await.expect("recv");
assert!(wr.ok);
assert_eq!(wr.value, serde_json::json!("hello"));
}
#[tokio::test]
async fn factory_builds_prompt_based_agent_when_script_path_absent() {
use crate::blueprint::compiler::SpawnerFactory;
use crate::blueprint::{AgentDef, AgentKind, AgentProfile};
let factory = AgentBlockInProcessSpawnerFactory::new();
let ad = AgentDef {
name: "writer".into(),
kind: AgentKind::AgentBlock,
spec: serde_json::json!({}),
profile: Some(AgentProfile {
system_prompt: "You are writer.".into(),
..Default::default()
}),
meta: None,
};
let _spawner = factory.build(&ad, None).expect("factory build");
}
#[tokio::test]
async fn factory_builds_script_based_agent_when_script_path_present() {
use crate::blueprint::compiler::SpawnerFactory;
use crate::blueprint::{AgentDef, AgentKind, AgentProfile};
let factory = AgentBlockInProcessSpawnerFactory::new();
let ad = AgentDef {
name: "patch-spawner".into(),
kind: AgentKind::AgentBlock,
spec: serde_json::json!({
"script_path": "assets/operator_scripts/blueprint_patch_spawner.lua",
"project_root": ".",
}),
profile: Some(AgentProfile {
system_prompt: "Patch generator.".into(),
..Default::default()
}),
meta: None,
};
let _spawner = factory.build(&ad, None).expect("factory build");
}
}