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::{
AgentBlockInProcessSpawnerFactory, LuaInProcessSpawnerFactory, OperatorSpawnerFactory,
RustFnInProcessSpawnerFactory, SpawnerRegistry, SubprocessProcessSpawnerFactory,
};
use mlua_swarm::{
Compiler, Engine, EngineCfg, EnhanceApplication, EnhanceApplicationConfig, Role,
TaskLaunchService,
};
use mlua_swarm_server::{
build_blueprints_router_with_refs, build_enhance_log_router, build_enhance_settings_router,
build_issues_router, 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)]
inject_endpoint_for_worker: bool,
#[arg(long)]
long_hold_warn_ms: Option<u64>,
#[arg(long)]
default_agent_kind: Option<String>,
#[arg(long)]
sync_timeout_secs: Option<u64>,
#[arg(long)]
stale_run_sweep_secs: Option<u64>,
#[arg(long)]
engine_max_hold_ms: 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
},
inject_endpoint_for_worker: if args.inject_endpoint_for_worker {
Some(true)
} else {
None
},
long_hold_warn_ms: args.long_hold_warn_ms,
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,
stale_run_sweep_secs: args.stale_run_sweep_secs,
engine_max_hold_ms: args.engine_max_hold_ms,
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 engine = Engine::new_with_layers(
engine_cfg_from(&cfg),
mlua_swarm_server::default_layer_registry_with(mlua_swarm_server::LayerOptions {
long_hold_warn_ms: cfg.long_hold_warn_ms,
}),
);
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::<AgentBlockInProcessSpawnerFactory>(Arc::new(
AgentBlockInProcessSpawnerFactory::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()),
};
let run_trace_store: Arc<dyn mlua_swarm::store::trace::RunTraceStore> =
match &cfg.run_store_path {
Some(path) => {
eprintln!("mse serve: SqliteRunTraceStore at {}", path.display());
let (s, driver) = mlua_swarm::store::trace::SqliteRunTraceStore::open(path)
.await
.unwrap_or_else(|e| panic!("mse serve: SqliteRunTraceStore open failed: {e}"));
isle_drivers.push(driver);
Arc::new(s)
}
None => Arc::new(mlua_swarm::store::trace::InMemoryRunTraceStore::new()),
};
recover_interrupted_runs(&task_store, &run_store, &replay_store).await;
let base_url: Option<std::sync::Arc<str>> = if cfg.inject_endpoint_for_worker {
Some(format!("http://{}", cfg.bind).into())
} else {
None
};
let shutdown_task_store = task_store.clone();
let shutdown_run_store = run_store.clone();
let sweep_task_store = task_store.clone();
let sweep_run_store = run_store.clone();
let sweep_trace_store = run_trace_store.clone();
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,
base_url,
Some(task_store),
Some(run_store),
Some(replay_store),
Some(run_trace_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 stale_run_sweep_secs = cfg.stale_run_sweep_secs;
let stale_run_sweeper = if stale_run_sweep_secs == 0 {
eprintln!("mse serve: stale run sweep disabled (stale_run_sweep_secs = 0)");
None
} else {
eprintln!(
"mse serve: stale run sweep every {}s, idle threshold {stale_run_sweep_secs}s",
STALE_RUN_SWEEP_PERIOD.as_secs()
);
Some(tokio::spawn(async move {
let mut ticker = tokio::time::interval(STALE_RUN_SWEEP_PERIOD);
ticker.tick().await;
loop {
ticker.tick().await;
sweep_stale_running_runs(
&sweep_task_store,
&sweep_run_store,
&sweep_trace_store,
stale_run_sweep_secs,
)
.await;
}
}))
};
let doctor_info = DoctorInfo {
server_version: env!("CARGO_PKG_VERSION").to_string(),
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 shutdown_signal = async {
tokio::select! {
_ = tokio::signal::ctrl_c() => { eprintln!("mse serve: ctrl-c, shutting down"); }
_ = wait_sigterm() => { eprintln!("mse serve: SIGTERM, shutting down"); }
}
};
axum::serve(listener, app)
.with_graceful_shutdown(shutdown_signal)
.await
.expect("serve");
enhance_loop.abort();
if let Some(sweeper) = stale_run_sweeper {
sweeper.abort();
}
interrupt_running_on_shutdown(&shutdown_task_store, &shutdown_run_store).await;
for driver in isle_drivers {
if let Err(e) = driver.shutdown().await {
eprintln!("mse serve: isle driver shutdown error: {e}");
}
}
Ok(())
}
fn engine_cfg_from(cfg: &mlua_swarm_server::config::ResolvedConfig) -> EngineCfg {
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;
if let Some(ms) = cfg.engine_max_hold_ms {
c.max_hold_ms = ms as u128;
}
c
}
const STALE_RUN_SWEEP_PERIOD: Duration = Duration::from_secs(60);
async fn sweep_stale_running_runs(
task_store: &std::sync::Arc<dyn mlua_swarm::store::task::TaskStore>,
run_store: &std::sync::Arc<dyn mlua_swarm::store::run::RunStore>,
run_trace_store: &std::sync::Arc<dyn mlua_swarm::store::trace::RunTraceStore>,
threshold_secs: u64,
) -> usize {
if threshold_secs == 0 {
return 0;
}
let now = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map(|d| d.as_secs())
.unwrap_or(0);
let running = match run_store.list_running().await {
Ok(v) => v,
Err(e) => {
tracing::warn!(error = %e, "stale run sweep: list_running failed");
return 0;
}
};
let mut reaped = 0;
for run in running {
let idle_secs = now.saturating_sub(run.updated_at);
if idle_secs <= threshold_secs {
continue;
}
match run_store
.try_transition(
&run.id,
mlua_swarm::store::run::RunStatus::Running,
mlua_swarm::store::run::RunStatus::Interrupted,
)
.await
{
Ok(true) => {}
Ok(false) => {
tracing::debug!(
run_id = %run.id,
"stale run sweep: run left Running between the scan and the transition; skipped"
);
continue;
}
Err(e) => {
tracing::warn!(run_id = %run.id, error = %e, "stale run sweep: try_transition failed");
continue;
}
}
let envelope = serde_json::json!({ "error": format!("orphaned: no driver progress for {idle_secs}s") });
if let Err(e) = run_store.set_result(&run.id, envelope).await {
tracing::warn!(run_id = %run.id, error = %e, "stale run sweep: set_result failed");
}
if let Err(e) = task_store
.update_status(
&run.task_id,
mlua_swarm::store::task::TaskRecordStatus::Interrupted,
)
.await
{
tracing::warn!(task_id = %run.task_id, error = %e, "stale run sweep: task update_status failed");
}
mlua_swarm::store::trace::TraceHandle::new(run.id.clone(), run_trace_store.clone())
.append(
mlua_swarm::store::trace::kind::RUN_FINISHED,
None,
None,
serde_json::json!({ "status": "interrupted", "reason": "stale run sweep" }),
)
.await;
tracing::info!(
run_id = %run.id,
task_id = %run.task_id,
idle_secs,
threshold_secs,
resume_url = %format!("POST /v1/runs/{}/resume", run.id),
"stale run sweep: marked an orphaned run Interrupted"
);
reaped += 1;
}
reaped
}
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 {
mark_run_interrupted(task_store, run_store, &run, "server restart", "boot sweep").await;
match replay_store.list_by_run(&run.id).await {
Ok(entries) => {
let replayed_steps = entries.len();
if replayed_steps > 0 || run.input_json.is_some() {
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::info!(
run_id = %run.id,
task_id = %run.task_id,
"boot sweep: not resumable (no replay entries, no input snapshot)"
);
}
}
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"
);
}
}
}
}
async fn mark_run_interrupted(
task_store: &std::sync::Arc<dyn mlua_swarm::store::task::TaskStore>,
run_store: &std::sync::Arc<dyn mlua_swarm::store::run::RunStore>,
run: &mlua_swarm::store::run::RunRecord,
reason: &str,
context: &str,
) {
let envelope = serde_json::json!({ "error": reason });
if let Err(e) = run_store.set_result(&run.id, envelope).await {
eprintln!(
"mse serve: {context}: 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: {context}: 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: {context}: task {} update_status failed: {e}",
run.task_id
);
}
}
async fn interrupt_running_on_shutdown(
task_store: &std::sync::Arc<dyn mlua_swarm::store::task::TaskStore>,
run_store: &std::sync::Arc<dyn mlua_swarm::store::run::RunStore>,
) {
let running = match run_store.list_running().await {
Ok(v) => v,
Err(e) => {
eprintln!("mse serve: shutdown drain: list_running failed: {e}");
return;
}
};
for run in running {
mark_run_interrupted(
task_store,
run_store,
&run,
"server shutdown",
"shutdown drain",
)
.await;
}
}
#[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,
lints: 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,
lints: 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,
subprocesses: vec![],
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::store::trace::{
kind as trace_kind, InMemoryRunTraceStore, RunTraceStore, TraceQuery,
};
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());
}
#[tokio::test]
async fn interrupt_running_on_shutdown_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 running_task_id = TaskId::parse("T-shutdown-running").unwrap();
let done_task_id = TaskId::parse("T-shutdown-done").unwrap();
let running_run_id = RunId::parse("R-shutdown-running").unwrap();
let done_run_id = RunId::parse("R-shutdown-done").unwrap();
for (tid, status) in [
(&running_task_id, TaskRecordStatus::Running),
(&done_task_id, TaskRecordStatus::Done),
] {
task_store
.create(TaskRecord {
id: tid.clone(),
goal: "shutdown drain fixture".into(),
blueprint_ref: json!({}),
input_ctx: json!({}),
task_input_spec: None,
status,
created_at: 1,
updated_at: 1,
})
.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();
interrupt_running_on_shutdown(&task_store, &run_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 shutdown"}))
);
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);
}
fn now_secs() -> u64 {
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map(|d| d.as_secs())
.unwrap_or(0)
}
async fn seed_pair(
task_store: &Arc<dyn TaskStore>,
run_store: &Arc<dyn RunStore>,
key: &str,
status: RunStatus,
updated_at: u64,
) -> (TaskId, RunId) {
let task_id = TaskId::parse(format!("T-{key}")).unwrap();
let run_id = RunId::parse(format!("R-{key}")).unwrap();
task_store
.create(TaskRecord {
id: task_id.clone(),
goal: "stale sweep 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: run_id.clone(),
task_id: task_id.clone(),
status,
step_entries: vec![],
degradations: vec![],
operator_sid: None,
result_ref: None,
input_json: Some("{}".to_string()),
created_at: 1,
updated_at,
})
.await
.unwrap();
(task_id, run_id)
}
#[tokio::test]
async fn stale_run_sweep_reaps_the_orphan_and_leaves_the_fresh_run_running() {
let task_store: Arc<dyn TaskStore> = Arc::new(InMemoryTaskStore::new());
let run_store: Arc<dyn RunStore> = Arc::new(InMemoryRunStore::new());
let trace_store: Arc<dyn RunTraceStore> = Arc::new(InMemoryRunTraceStore::new());
let now = now_secs();
let (stale_task_id, stale_run_id) = seed_pair(
&task_store,
&run_store,
"stale",
RunStatus::Running,
now - 1_000,
)
.await;
let (fresh_task_id, fresh_run_id) =
seed_pair(&task_store, &run_store, "fresh", RunStatus::Running, now).await;
let reaped = sweep_stale_running_runs(&task_store, &run_store, &trace_store, 100).await;
assert_eq!(reaped, 1, "only the idle run is reaped");
let stale_run = run_store.get(&stale_run_id).await.unwrap();
assert_eq!(stale_run.status, RunStatus::Interrupted);
let err = stale_run
.result_ref
.as_ref()
.and_then(|r| r.get("error"))
.and_then(|e| e.as_str())
.expect("terminal envelope carries an `error` string");
assert!(
err.starts_with("orphaned: no driver progress for"),
"reason names the actual cause: {err}"
);
assert!(
stale_run.input_json.is_some(),
"resume precondition: Interrupted + a launch-input snapshot"
);
assert_eq!(
task_store.get(&stale_task_id).await.unwrap().status,
TaskRecordStatus::Interrupted
);
let events = trace_store
.list(&stale_run_id, &TraceQuery::default())
.await
.expect("trace list");
let finished: Vec<_> = events
.iter()
.filter(|e| e.kind == trace_kind::RUN_FINISHED)
.collect();
assert_eq!(finished.len(), 1);
assert_eq!(finished[0].payload["status"], json!("interrupted"));
let fresh_run = run_store.get(&fresh_run_id).await.unwrap();
assert_eq!(fresh_run.status, RunStatus::Running);
assert_eq!(fresh_run.result_ref, None);
assert_eq!(
task_store.get(&fresh_task_id).await.unwrap().status,
TaskRecordStatus::Running
);
assert!(trace_store
.list(&fresh_run_id, &TraceQuery::default())
.await
.expect("trace list")
.is_empty());
}
#[tokio::test]
async fn stale_run_sweep_threshold_zero_is_a_no_op() {
let task_store: Arc<dyn TaskStore> = Arc::new(InMemoryTaskStore::new());
let run_store: Arc<dyn RunStore> = Arc::new(InMemoryRunStore::new());
let trace_store: Arc<dyn RunTraceStore> = Arc::new(InMemoryRunTraceStore::new());
let (_, run_id) = seed_pair(
&task_store,
&run_store,
"disabled",
RunStatus::Running,
now_secs() - 100_000,
)
.await;
let reaped = sweep_stale_running_runs(&task_store, &run_store, &trace_store, 0).await;
assert_eq!(reaped, 0);
assert_eq!(
run_store.get(&run_id).await.unwrap().status,
RunStatus::Running
);
}
#[tokio::test]
async fn stale_run_sweep_cas_never_clobbers_a_terminal_run() {
let task_store: Arc<dyn TaskStore> = Arc::new(InMemoryTaskStore::new());
let run_store: Arc<dyn RunStore> = Arc::new(InMemoryRunStore::new());
let trace_store: Arc<dyn RunTraceStore> = Arc::new(InMemoryRunTraceStore::new());
let (_, done_run_id) = seed_pair(
&task_store,
&run_store,
"already-done",
RunStatus::Done,
now_secs() - 100_000,
)
.await;
run_store
.set_result(&done_run_id, json!({"out": "finished on its own"}))
.await
.unwrap();
let reaped = sweep_stale_running_runs(&task_store, &run_store, &trace_store, 100).await;
assert_eq!(reaped, 0, "a terminal run is not a sweep candidate");
let done_run = run_store.get(&done_run_id).await.unwrap();
assert_eq!(done_run.status, RunStatus::Done);
assert_eq!(
done_run.result_ref,
Some(json!({"out": "finished on its own"})),
"the run's own terminal result survives the sweep"
);
assert!(
!run_store
.try_transition(&done_run_id, RunStatus::Running, RunStatus::Interrupted)
.await
.unwrap(),
"Running -> Interrupted CAS must lose against a terminal status"
);
}
#[test]
fn engine_cfg_from_keeps_the_engine_default_max_hold_when_unset() {
let cfg = mlua_swarm_server::config::ResolvedConfig::default();
assert_eq!(cfg.engine_max_hold_ms, None);
assert_eq!(engine_cfg_from(&cfg).max_hold_ms, 50);
}
#[test]
fn engine_cfg_from_applies_the_configured_max_hold() {
let cfg = mlua_swarm_server::config::ResolvedConfig {
engine_max_hold_ms: Some(200),
..Default::default()
};
assert_eq!(engine_cfg_from(&cfg).max_hold_ms, 200);
}
}