use clap::Parser;
use mlua_swarm::blueprint::store::{
blueprint_version, BlueprintId, BlueprintStore, CommitMetadata, Git2BlueprintStore,
};
use mlua_swarm::blueprint::{
current_schema_version, AgentDef, AgentKind, Blueprint, BlueprintMetadata, BlueprintOrigin,
CompilerHints, CompilerStrategy,
};
use mlua_swarm::store::enhance_log::{
EnhanceLogStore, InMemoryEnhanceLogStore, SqliteEnhanceLogStore,
};
use mlua_swarm::store::enhance_setting::{
EnhanceSettingId, EnhanceSettingStore, InMemoryEnhanceSettingStore, SqliteEnhanceSettingStore,
};
use mlua_swarm::store::issue::{InMemoryIssueStore, IssueStore, SqliteIssueStore};
use mlua_swarm::store::output::{InMemoryOutputStore, OutputStore, SqliteOutputStore};
use mlua_swarm::store::replay::{InMemoryReplayStore, ReplayStore, SqliteReplayStore};
use mlua_swarm::store::run::{InMemoryRunStore, RunStore, SqliteRunStore};
use mlua_swarm::store::task::{InMemoryTaskStore, SqliteTaskStore, TaskStore};
use mlua_swarm::{
Compiler, Engine, EngineCfg, EnhanceApplication, EnhanceApplicationConfig, Role,
TaskLaunchService,
};
use mlua_swarm::{
LuaInProcessSpawnerFactory, OperatorSpawnerFactory, RustFnInProcessSpawnerFactory,
SpawnerRegistry, SubprocessProcessSpawnerFactory,
};
use mlua_swarm_server::{
build_blueprints_router_with_refs, build_enhance_log_router, build_enhance_settings_router,
build_issues_router, default_layer_registry, default_registry_with_enhance_flow,
doctor::{build_doctor_router, DoctorInfo},
};
use serde_json::json;
use std::sync::Arc;
use std::time::Duration;
#[derive(Parser, Debug)]
#[command(about = "Run the HTTP server (mse serve).")]
pub struct Args {
#[arg(long)]
config: Option<std::path::PathBuf>,
#[arg(long)]
bind: Option<String>,
#[arg(long)]
token_secret: Option<String>,
#[arg(long)]
seed_blueprint_id: Option<String>,
#[arg(long)]
git_store_path: Option<std::path::PathBuf>,
#[arg(long)]
issue_store_path: Option<std::path::PathBuf>,
#[arg(long)]
enhance_setting_store_path: Option<std::path::PathBuf>,
#[arg(long)]
enhance_log_store_path: Option<std::path::PathBuf>,
#[arg(long)]
output_store_path: Option<std::path::PathBuf>,
#[arg(long)]
task_store_path: Option<std::path::PathBuf>,
#[arg(long)]
run_store_path: Option<std::path::PathBuf>,
#[arg(long)]
replay_store_path: Option<std::path::PathBuf>,
#[arg(long)]
enable_enhance_flow: bool,
#[arg(long, value_name = "allow|reject")]
legacy_worker_binding_policy: Option<String>,
#[arg(long)]
blueprint_ref_base: Option<std::path::PathBuf>,
#[arg(long = "include", action = clap::ArgAction::Append, value_name = "DIR")]
blueprint_ref_includes: Vec<std::path::PathBuf>,
#[arg(long)]
blueprint_strict_embed: bool,
#[arg(long)]
default_agent_kind: Option<String>,
#[arg(long)]
sync_timeout_secs: Option<u64>,
#[arg(long)]
ephemeral: bool,
#[arg(long)]
check_policy: Option<String>,
}
fn parse_agent_kind_cli(s: &str) -> Result<mlua_swarm::blueprint::AgentKind, String> {
serde_json::from_value(serde_json::Value::String(s.to_string()))
.map_err(|e| format!("invalid --default-agent-kind {s:?}: {e}"))
}
fn parse_check_policy_cli(s: &str) -> Result<mlua_swarm::core::config::CheckPolicy, String> {
serde_json::from_value(serde_json::Value::String(s.to_string()))
.map_err(|e| format!("invalid --check-policy {s:?}: {e}"))
}
fn parse_legacy_worker_binding_policy_cli(
s: &str,
) -> Result<mlua_swarm::LegacyWorkerBindingPolicy, String> {
serde_json::from_value(serde_json::Value::String(s.to_string()))
.map_err(|e| format!("invalid --legacy-worker-binding-policy {s:?}: {e}"))
}
pub async fn run(args: Args) -> anyhow::Result<()> {
let config_path = args
.config
.clone()
.unwrap_or_else(mlua_swarm_server::config::default_config_path);
let file_config = mlua_swarm_server::config::load_file_config(&config_path)
.unwrap_or_else(|e| panic!("mse serve: config load failed: {e}"));
let cli_overrides = mlua_swarm_server::config::CliOverrides {
bind: args.bind.clone(),
enable_enhance_flow: if args.enable_enhance_flow {
Some(true)
} else {
None
},
legacy_worker_binding_policy: args.legacy_worker_binding_policy.as_ref().map(|s| {
parse_legacy_worker_binding_policy_cli(s).unwrap_or_else(|e| panic!("mse serve: {e}"))
}),
blueprint_ref_base: args.blueprint_ref_base.clone(),
blueprint_ref_includes: args.blueprint_ref_includes.clone(),
blueprint_strict_embed: if args.blueprint_strict_embed {
Some(true)
} else {
None
},
git_store_path: args.git_store_path.clone(),
issue_store_path: args.issue_store_path.clone(),
enhance_setting_store_path: args.enhance_setting_store_path.clone(),
enhance_log_store_path: args.enhance_log_store_path.clone(),
output_store_path: args.output_store_path.clone(),
task_store_path: args.task_store_path.clone(),
run_store_path: args.run_store_path.clone(),
replay_store_path: args.replay_store_path.clone(),
ephemeral: if args.ephemeral { Some(true) } else { None },
seed_blueprint_id: args.seed_blueprint_id.clone(),
default_agent_kind: args.default_agent_kind.clone(),
token_secret: args.token_secret.clone(),
sync_timeout_secs: args.sync_timeout_secs,
check_policy: args
.check_policy
.as_ref()
.map(|s| parse_check_policy_cli(s).unwrap_or_else(|e| panic!("mse serve: {e}"))),
};
let cfg = mlua_swarm_server::config::resolve(cli_overrides, file_config)
.unwrap_or_else(|e| panic!("mse serve: config resolve failed: {e}"));
let default_agent_kind: Option<mlua_swarm::blueprint::AgentKind> = cfg
.default_agent_kind
.as_ref()
.map(|s| parse_agent_kind_cli(s).unwrap_or_else(|e| panic!("mse serve: {e}")));
eprintln!("mse serve: config loaded from {}", config_path.display());
let make_cfg = || {
let mut c = EngineCfg::default();
if let Some(hex_secret) = &cfg.token_secret {
c.token_secret = hex::decode(hex_secret).expect("token-secret must be hex");
}
c.check_policy = cfg.check_policy;
c
};
let engine = Engine::new_with_layers(make_cfg(), default_layer_registry());
let store: Arc<dyn BlueprintStore> = {
let bp_root = cfg.git_store_path.join("blueprints");
let s = Git2BlueprintStore::open_or_init(&bp_root).expect("git store open_or_init");
eprintln!(
"mse serve: blueprint store = Git2 root={} (per-id repos)",
bp_root.display()
);
Arc::new(s)
};
let id = BlueprintId::new(cfg.seed_blueprint_id.clone());
let need_seed = store.read_head(&id).await.is_err();
if need_seed {
let bp = seed_blueprint(&cfg.seed_blueprint_id);
let v0 = blueprint_version(&bp).expect("blueprint_version");
let now_ms = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map(|d| d.as_millis() as i64)
.unwrap_or(0);
store
.write_new(&id, &bp, &[], CommitMetadata::seed(id.clone(), v0, now_ms))
.await
.expect("seed write");
eprintln!("mse serve: seeded blueprint_id={}", id.as_str());
} else {
eprintln!("mse serve: existing head found, skip seed");
}
let op_factory = Arc::new(OperatorSpawnerFactory::new());
let make_registry = || -> SpawnerRegistry {
let mut reg = if cfg.enable_enhance_flow {
default_registry_with_enhance_flow()
} else {
let rustfn_factory = mlua_swarm::worker::baseline::extend_with_baseline(
RustFnInProcessSpawnerFactory::new(),
);
let mut r = SpawnerRegistry::new();
r.register::<SubprocessProcessSpawnerFactory>(Arc::new(
SubprocessProcessSpawnerFactory,
));
r.register::<RustFnInProcessSpawnerFactory>(Arc::new(rustfn_factory));
r.register::<LuaInProcessSpawnerFactory>(Arc::new(LuaInProcessSpawnerFactory::new()));
r.register::<OperatorSpawnerFactory>(op_factory.clone());
r
};
reg.register::<OperatorSpawnerFactory>(op_factory.clone());
reg
};
let mut isle_drivers: Vec<rusqlite_isle::AsyncIsleDriver> = Vec::new();
let issue_store: Arc<dyn IssueStore> = match &cfg.issue_store_path {
Some(path) => {
eprintln!("mse serve: SqliteIssueStore at {}", path.display());
let (s, driver) = SqliteIssueStore::open(path)
.await
.unwrap_or_else(|e| panic!("mse serve: SqliteIssueStore open failed: {e}"));
isle_drivers.push(driver);
Arc::new(s)
}
None => Arc::new(InMemoryIssueStore::new()),
};
let setting_store: Arc<dyn EnhanceSettingStore> = match &cfg.enhance_setting_store_path {
Some(path) => {
eprintln!("mse serve: SqliteEnhanceSettingStore at {}", path.display());
let (s, driver) = SqliteEnhanceSettingStore::open(path)
.await
.unwrap_or_else(|e| {
panic!("mse serve: SqliteEnhanceSettingStore open failed: {e}")
});
isle_drivers.push(driver);
Arc::new(s)
}
None => Arc::new(InMemoryEnhanceSettingStore::new()),
};
let log_store: Arc<dyn EnhanceLogStore> = match &cfg.enhance_log_store_path {
Some(path) => {
eprintln!("mse serve: SqliteEnhanceLogStore at {}", path.display());
let (s, driver) = SqliteEnhanceLogStore::open(path)
.await
.unwrap_or_else(|e| panic!("mse serve: SqliteEnhanceLogStore open failed: {e}"));
isle_drivers.push(driver);
Arc::new(s)
}
None => Arc::new(InMemoryEnhanceLogStore::new()),
};
let output_store: Option<Arc<dyn OutputStore>> = match &cfg.output_store_path {
Some(path) => {
eprintln!("mse serve: SqliteOutputStore at {}", path.display());
let (s, driver) = SqliteOutputStore::open(path)
.await
.unwrap_or_else(|e| panic!("mse serve: SqliteOutputStore open failed: {e}"));
isle_drivers.push(driver);
Some(Arc::new(s))
}
None => Some(Arc::new(InMemoryOutputStore::new())),
};
let task_store: Arc<dyn TaskStore> = match &cfg.task_store_path {
Some(path) => {
eprintln!("mse serve: SqliteTaskStore at {}", path.display());
let (s, driver) = SqliteTaskStore::open(path)
.await
.unwrap_or_else(|e| panic!("mse serve: SqliteTaskStore open failed: {e}"));
isle_drivers.push(driver);
Arc::new(s)
}
None => Arc::new(InMemoryTaskStore::new()),
};
let run_store: Arc<dyn RunStore> = match &cfg.run_store_path {
Some(path) => {
eprintln!("mse serve: SqliteRunStore at {}", path.display());
let (s, driver) = SqliteRunStore::open(path)
.await
.unwrap_or_else(|e| panic!("mse serve: SqliteRunStore open failed: {e}"));
isle_drivers.push(driver);
Arc::new(s)
}
None => Arc::new(InMemoryRunStore::new()),
};
let replay_store: Arc<dyn ReplayStore> = match &cfg.replay_store_path {
Some(path) => {
eprintln!("mse serve: SqliteReplayStore at {}", path.display());
let (s, driver) = SqliteReplayStore::open(path)
.await
.unwrap_or_else(|e| panic!("mse serve: SqliteReplayStore open failed: {e}"));
isle_drivers.push(driver);
Arc::new(s)
}
None => Arc::new(InMemoryReplayStore::new()),
};
recover_interrupted_runs(&task_store, &run_store, &replay_store).await;
let base_url: std::sync::Arc<str> = format!("http://{}", cfg.bind).into();
let mut app = mlua_swarm_server::build_router_full_with_legacy_worker_binding_policy(
engine.clone(),
make_registry(),
Some(store.clone()),
Some(op_factory.clone()),
output_store,
Some(base_url),
Some(task_store),
Some(run_store),
Some(replay_store),
cfg.sync_timeout_secs,
cfg.legacy_worker_binding_policy,
);
let compiler = Compiler::new(make_registry());
let launch_enhance = Arc::new(
TaskLaunchService::new(engine.clone(), compiler)
.with_legacy_worker_binding_policy(cfg.legacy_worker_binding_policy),
);
let enhance_app = Arc::new(EnhanceApplication::new(
EnhanceApplicationConfig {
name: "enhance".into(),
setting_id: EnhanceSettingId::default_id(),
operator_id: "mse-enhance".into(),
role: Role::Operator,
},
issue_store.clone(),
setting_store.clone(),
store.clone(),
log_store.clone(),
launch_enhance,
));
let enhance_loop = tokio::spawn(enhance_app.clone().run_forever(Duration::from_millis(100)));
let doctor_info = DoctorInfo {
bind: cfg.bind.to_string(),
blueprint_backend: "git2".into(),
blueprint_store_root: Some(cfg.git_store_path.join("blueprints").display().to_string()),
blueprint_ref_base: cfg
.blueprint_ref_base
.as_ref()
.map(|p| p.display().to_string()),
enhance_flow_enabled: cfg.enable_enhance_flow,
legacy_worker_binding_policy: cfg.legacy_worker_binding_policy,
seed_blueprint_id: cfg.seed_blueprint_id.clone(),
check_policy: cfg.check_policy,
};
app = app
.merge(build_issues_router(issue_store.clone()))
.merge(build_blueprints_router_with_refs(
store.clone(),
cfg.blueprint_ref_base.clone(),
cfg.blueprint_ref_includes.clone(),
default_agent_kind,
cfg.blueprint_strict_embed,
cfg.legacy_worker_binding_policy,
))
.merge(build_enhance_log_router(log_store.clone()))
.merge(build_enhance_settings_router(
setting_store.clone(),
store.clone(),
))
.merge(build_doctor_router(doctor_info, store.clone()));
let _ = id;
eprintln!(
"mse serve: combined mode (task+enhance+operator) listening on http://{}",
cfg.bind
);
let listener = tokio::net::TcpListener::bind(cfg.bind).await.expect("bind");
let serve = axum::serve(listener, app);
tokio::select! {
r = serve => { r.expect("serve"); }
_ = tokio::signal::ctrl_c() => { eprintln!("mse serve: ctrl-c, shutting down"); }
_ = wait_sigterm() => { eprintln!("mse serve: SIGTERM, shutting down"); }
}
enhance_loop.abort();
for driver in isle_drivers {
if let Err(e) = driver.shutdown().await {
eprintln!("mse serve: isle driver shutdown error: {e}");
}
}
Ok(())
}
async fn recover_interrupted_runs(
task_store: &std::sync::Arc<dyn mlua_swarm::store::task::TaskStore>,
run_store: &std::sync::Arc<dyn mlua_swarm::store::run::RunStore>,
replay_store: &std::sync::Arc<dyn mlua_swarm::store::replay::ReplayStore>,
) {
let running = match run_store.list_running().await {
Ok(v) => v,
Err(e) => {
eprintln!("mse serve: boot sweep: list_running failed: {e}");
return;
}
};
for run in running {
let reason = serde_json::json!({"error": "server restart"});
if let Err(e) = run_store.set_result(&run.id, reason).await {
eprintln!(
"mse serve: boot sweep: run {} set_result failed: {e}",
run.id
);
}
if let Err(e) = run_store
.update_status(&run.id, mlua_swarm::store::run::RunStatus::Interrupted)
.await
{
eprintln!(
"mse serve: boot sweep: run {} update_status failed: {e}",
run.id
);
}
if let Err(e) = task_store
.update_status(
&run.task_id,
mlua_swarm::store::task::TaskRecordStatus::Interrupted,
)
.await
{
eprintln!(
"mse serve: boot sweep: task {} update_status failed: {e}",
run.task_id
);
}
match replay_store.list_by_run(&run.id).await {
Ok(entries) => {
let replayed_steps = entries.len();
if replayed_steps > 0 {
tracing::info!(
run_id = %run.id,
task_id = %run.task_id,
replayed_steps,
resume_url = %format!("POST /v1/runs/{}/resume", run.id),
"boot sweep: resumable Interrupted run"
);
} else {
tracing::debug!(
run_id = %run.id,
task_id = %run.task_id,
"boot sweep: not resumable, no replay entries"
);
}
}
Err(e) => {
tracing::warn!(
run_id = %run.id,
task_id = %run.task_id,
error = %e,
"boot sweep: replay_store list_by_run failed; skipping resumable classification"
);
}
}
}
}
#[cfg(unix)]
async fn wait_sigterm() {
use tokio::signal::unix::{signal, SignalKind};
match signal(SignalKind::terminate()) {
Ok(mut sig) => {
sig.recv().await;
}
Err(e) => {
eprintln!("mse serve: failed to install SIGTERM handler: {e}");
std::future::pending::<()>().await;
}
}
}
#[cfg(not(unix))]
async fn wait_sigterm() {
std::future::pending::<()>().await;
}
fn seed_blueprint(id: &str) -> Blueprint {
Blueprint {
schema_version: current_schema_version(),
id: id.into(),
flow: serde_json::from_value(json!({
"kind": "step",
"ref": mlua_swarm::worker::baseline::AG_IDENTITY,
"in": {"op": "lit", "value": "hello"},
"out": {"op": "path", "at": "$.out"},
}))
.unwrap(),
agents: vec![AgentDef {
name: mlua_swarm::worker::baseline::AG_IDENTITY.into(),
kind: AgentKind::RustFn,
spec: json!({"fn_id": mlua_swarm::worker::baseline::AG_IDENTITY}),
profile: None,
meta: None,
runner: None,
runner_ref: None,
verdict: None,
}],
operators: vec![],
metas: vec![],
hints: CompilerHints::default(),
strategy: CompilerStrategy::default(),
metadata: BlueprintMetadata {
description: Some("mse serve enhance seed".into()),
origin: BlueprintOrigin::Inline,
tags: vec![],
version_label: Some("0.1.0".into()),
project_name_alias: None,
default_run_ttl_secs: None,
strict_verdict_handling: None,
},
spawner_hints: Default::default(),
default_agent_kind: AgentKind::Operator,
default_operator_kind: None,
default_init_ctx: None,
default_agent_ctx: None,
default_context_policy: None,
projection_placement: None,
audits: vec![],
degradation_policy: None,
runners: vec![],
default_runner: None,
check_policy: None,
blueprint_ref_includes: Vec::new(),
}
}
#[cfg(test)]
mod tests {
use super::*;
use mlua_swarm::store::replay::ReplayEntry;
use mlua_swarm::store::run::{RunRecord, RunStatus};
use mlua_swarm::store::task::{TaskRecord, TaskRecordStatus};
use mlua_swarm::types::{RunId, TaskId};
#[tokio::test]
async fn recover_interrupted_runs_marks_running_as_interrupted() {
let task_store: Arc<dyn TaskStore> = Arc::new(InMemoryTaskStore::new());
let run_store: Arc<dyn RunStore> = Arc::new(InMemoryRunStore::new());
let replay_store: Arc<dyn ReplayStore> = Arc::new(InMemoryReplayStore::new());
let running_task_id = TaskId::parse("T-running").unwrap();
let done_task_id = TaskId::parse("T-done").unwrap();
let running_run_id = RunId::parse("R-running").unwrap();
let done_run_id = RunId::parse("R-done").unwrap();
task_store
.create(TaskRecord {
id: running_task_id.clone(),
goal: "resolve issue #35".into(),
blueprint_ref: json!({}),
input_ctx: json!({}),
task_input_spec: None,
status: TaskRecordStatus::Running,
created_at: 1,
updated_at: 1,
})
.await
.unwrap();
task_store
.create(TaskRecord {
id: done_task_id.clone(),
goal: "unrelated done task".into(),
blueprint_ref: json!({}),
input_ctx: json!({}),
task_input_spec: None,
status: TaskRecordStatus::Done,
created_at: 2,
updated_at: 2,
})
.await
.unwrap();
run_store
.create(RunRecord {
id: running_run_id.clone(),
task_id: running_task_id.clone(),
status: RunStatus::Running,
step_entries: vec![],
degradations: vec![],
operator_sid: None,
result_ref: None,
input_json: None,
created_at: 1,
updated_at: 1,
})
.await
.unwrap();
run_store
.create(RunRecord {
id: done_run_id.clone(),
task_id: done_task_id.clone(),
status: RunStatus::Done,
step_entries: vec![],
degradations: vec![],
operator_sid: None,
result_ref: None,
input_json: None,
created_at: 2,
updated_at: 2,
})
.await
.unwrap();
recover_interrupted_runs(&task_store, &run_store, &replay_store).await;
let running_run = run_store.get(&running_run_id).await.unwrap();
assert_eq!(running_run.status, RunStatus::Interrupted);
assert_eq!(
running_run.result_ref,
Some(json!({"error": "server restart"}))
);
let running_task = task_store.get(&running_task_id).await.unwrap();
assert_eq!(running_task.status, TaskRecordStatus::Interrupted);
let done_run = run_store.get(&done_run_id).await.unwrap();
assert_eq!(done_run.status, RunStatus::Done);
assert_eq!(done_run.result_ref, None);
let done_task = task_store.get(&done_task_id).await.unwrap();
assert_eq!(done_task.status, TaskRecordStatus::Done);
}
#[tokio::test]
async fn recover_interrupted_runs_classifies_by_replay_entries() {
let task_store: Arc<dyn TaskStore> = Arc::new(InMemoryTaskStore::new());
let run_store: Arc<dyn RunStore> = Arc::new(InMemoryRunStore::new());
let replay_store: Arc<dyn ReplayStore> = Arc::new(InMemoryReplayStore::new());
let task_with_replay = TaskId::parse("T-resumable").unwrap();
let task_without_replay = TaskId::parse("T-not-resumable").unwrap();
let run_with_replay = RunId::parse("R-resumable").unwrap();
let run_without_replay = RunId::parse("R-not-resumable").unwrap();
for (tid, rid) in [
(&task_with_replay, &run_with_replay),
(&task_without_replay, &run_without_replay),
] {
task_store
.create(TaskRecord {
id: tid.clone(),
goal: "resume classification fixture".into(),
blueprint_ref: json!({}),
input_ctx: json!({}),
task_input_spec: None,
status: TaskRecordStatus::Running,
created_at: 1,
updated_at: 1,
})
.await
.unwrap();
run_store
.create(RunRecord {
id: rid.clone(),
task_id: tid.clone(),
status: RunStatus::Running,
step_entries: vec![],
degradations: vec![],
operator_sid: None,
result_ref: None,
input_json: Some("{}".to_string()),
created_at: 1,
updated_at: 1,
})
.await
.unwrap();
}
replay_store
.append(ReplayEntry {
run_id: run_with_replay.clone(),
step_ref: "step-a".into(),
input_hash: "hash-a".into(),
occurrence: 0,
ctx_snapshot_json: "{}".into(),
step_output_json: "{}".into(),
created_at: 1,
})
.await
.unwrap();
recover_interrupted_runs(&task_store, &run_store, &replay_store).await;
let with_replay = run_store.get(&run_with_replay).await.unwrap();
assert_eq!(with_replay.status, RunStatus::Interrupted);
let without_replay = run_store.get(&run_without_replay).await.unwrap();
assert_eq!(without_replay.status, RunStatus::Interrupted);
let entries = replay_store.list_by_run(&run_with_replay).await.unwrap();
assert_eq!(entries.len(), 1);
let empty = replay_store.list_by_run(&run_without_replay).await.unwrap();
assert!(empty.is_empty());
}
}