use mlua_swarm::blueprint::{
current_schema_version, AgentDef, AgentKind, Blueprint, BlueprintMetadata, CompilerHints,
CompilerStrategy,
};
use mlua_swarm::core::config::EngineCfg;
use mlua_swarm::core::engine::Engine;
use mlua_swarm::store::replay::{InMemoryReplayStore, ReplayEntry, ReplayStore};
use mlua_swarm::store::run::{InMemoryRunStore, RunRecord, RunStatus, RunStore, StepEntry};
use mlua_swarm::types::StepId;
use mlua_swarm::{RunId, TaskId};
use serde_json::json;
use std::sync::Arc;
use std::time::Duration;
fn two_step_blueprint() -> Blueprint {
Blueprint {
schema_version: current_schema_version(),
id: "rerun-from-test-bp".into(),
flow: serde_json::from_value(json!({
"kind": "seq",
"children": [
{
"kind": "step",
"ref": "agent-a",
"in": {"op": "lit", "value": "hello"},
"out": {"op": "path", "at": "$.a"},
},
{
"kind": "step",
"ref": "agent-b",
"in": {"op": "path", "at": "$.a"},
"out": {"op": "path", "at": "$.b"},
},
],
}))
.expect("flow parse"),
agents: vec![
AgentDef {
name: "agent-a".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,
},
AgentDef {
name: "agent-b".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::default(),
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(),
}
}
fn two_step_blueprint_with_unbound_operator() -> Blueprint {
let mut bp = two_step_blueprint();
bp.agents[1] = AgentDef {
name: "agent-b".into(),
kind: AgentKind::Operator,
spec: json!({ "operator_ref": "nonexistent-operator" }),
profile: None,
meta: None,
runner: None,
runner_ref: None,
verdict: None,
};
bp
}
async fn spawn_server(run_store: Arc<dyn RunStore>, replay_store: Arc<dyn ReplayStore>) -> String {
let engine = Engine::new_with_layers(
EngineCfg::default(),
mlua_swarm_server::default_layer_registry(),
);
let router = mlua_swarm_server::build_router_full(
engine,
mlua_swarm_server::default_registry(),
None,
None,
None,
None,
None,
Some(run_store),
Some(replay_store),
300,
);
let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
.await
.expect("bind ephemeral port");
let addr = listener.local_addr().expect("local addr");
tokio::spawn(async move {
let _ = axum::serve(listener, router).await;
});
format!("http://{addr}")
}
fn seed_run(
run_id: &RunId,
task_id: &TaskId,
status: RunStatus,
input_json: Option<String>,
) -> RunRecord {
RunRecord {
id: run_id.clone(),
task_id: task_id.clone(),
status,
step_entries: vec![],
degradations: vec![],
operator_sid: None,
result_ref: None,
input_json,
created_at: 0,
updated_at: 0,
}
}
fn snapshot_json_for(bp: &Blueprint) -> String {
json!({
"blueprint": { "kind": "inline", "value": bp },
"operator_id": "test-op",
"role": "operator",
"ttl": { "secs": 30, "nanos": 0 },
"init_ctx": { "in": "hello" },
"operator_kind": null,
"bridge_id": null,
"hook_id": null,
"operator_backend_id": null,
"operator_kind_overrides": {},
"task_input": null,
"check_policy": null,
})
.to_string()
}
#[tokio::test]
async fn rerun_from_unknown_run_returns_404() {
let run_store: Arc<dyn RunStore> = Arc::new(InMemoryRunStore::new());
let replay_store: Arc<dyn ReplayStore> = Arc::new(InMemoryReplayStore::new());
let base = spawn_server(run_store, replay_store).await;
let resp = reqwest::Client::new()
.post(format!("{base}/v1/runs/{}/rerun-from", RunId::new()))
.json(&json!({ "from_step": "agent-b" }))
.send()
.await
.expect("request");
assert_eq!(resp.status(), reqwest::StatusCode::NOT_FOUND);
}
#[tokio::test]
async fn rerun_from_running_run_returns_409() {
let run_store: Arc<dyn RunStore> = Arc::new(InMemoryRunStore::new());
let replay_store: Arc<dyn ReplayStore> = Arc::new(InMemoryReplayStore::new());
let run_id = RunId::new();
let task_id = TaskId::new();
run_store
.create(seed_run(
&run_id,
&task_id,
RunStatus::Running,
Some("{}".to_string()),
))
.await
.expect("seed run");
let base = spawn_server(run_store, replay_store).await;
let resp = reqwest::Client::new()
.post(format!("{base}/v1/runs/{run_id}/rerun-from"))
.json(&json!({ "from_step": "agent-b" }))
.send()
.await
.expect("request");
assert_eq!(resp.status(), reqwest::StatusCode::CONFLICT);
}
#[tokio::test]
async fn rerun_from_done_run_without_input_returns_422() {
let run_store: Arc<dyn RunStore> = Arc::new(InMemoryRunStore::new());
let replay_store: Arc<dyn ReplayStore> = Arc::new(InMemoryReplayStore::new());
let run_id = RunId::new();
let task_id = TaskId::new();
run_store
.create(seed_run(&run_id, &task_id, RunStatus::Done, None))
.await
.expect("seed run");
let base = spawn_server(run_store.clone(), replay_store).await;
let resp = reqwest::Client::new()
.post(format!("{base}/v1/runs/{run_id}/rerun-from"))
.json(&json!({ "from_step": "agent-b" }))
.send()
.await
.expect("request");
assert_eq!(resp.status(), reqwest::StatusCode::UNPROCESSABLE_ENTITY);
let after = run_store.get(&run_id).await.expect("run present");
assert_eq!(after.status, RunStatus::Done);
}
#[tokio::test]
async fn rerun_from_missing_step_returns_422() {
let run_store: Arc<dyn RunStore> = Arc::new(InMemoryRunStore::new());
let replay_store: Arc<dyn ReplayStore> = Arc::new(InMemoryReplayStore::new());
let run_id = RunId::new();
let task_id = TaskId::new();
let bp = two_step_blueprint();
run_store
.create(seed_run(
&run_id,
&task_id,
RunStatus::Done,
Some(snapshot_json_for(&bp)),
))
.await
.expect("seed run");
let ctx = mlua_swarm::core::ctx::Ctx::new(mlua_swarm::types::StepId::new(), 1, "agent-a");
replay_store
.append(
ReplayEntry::from_completion(
run_id.clone(),
"agent-a",
"h",
0,
&ctx,
&json!({ "v": 1 }),
)
.expect("entry build"),
)
.await
.expect("seed replay");
let base = spawn_server(run_store.clone(), replay_store.clone()).await;
let resp = reqwest::Client::new()
.post(format!("{base}/v1/runs/{run_id}/rerun-from"))
.json(&json!({ "from_step": "nonexistent-step" }))
.send()
.await
.expect("request");
assert_eq!(resp.status(), reqwest::StatusCode::UNPROCESSABLE_ENTITY);
let after = run_store.get(&run_id).await.expect("run present");
assert_eq!(after.status, RunStatus::Done);
let entries_after = replay_store.list_by_run(&run_id).await.expect("list");
assert_eq!(
entries_after.len(),
1,
"replay log must be untouched on 422"
);
}
#[tokio::test]
async fn rerun_from_empty_from_step_returns_400() {
let run_store: Arc<dyn RunStore> = Arc::new(InMemoryRunStore::new());
let replay_store: Arc<dyn ReplayStore> = Arc::new(InMemoryReplayStore::new());
let run_id = RunId::new();
let task_id = TaskId::new();
run_store
.create(seed_run(
&run_id,
&task_id,
RunStatus::Done,
Some("{}".to_string()),
))
.await
.expect("seed run");
let base = spawn_server(run_store, replay_store).await;
let resp = reqwest::Client::new()
.post(format!("{base}/v1/runs/{run_id}/rerun-from"))
.json(&json!({ "from_step": "" }))
.send()
.await
.expect("request");
assert_eq!(resp.status(), reqwest::StatusCode::BAD_REQUEST);
}
#[tokio::test]
async fn rerun_from_happy_path_truncates_and_completes() {
let run_store: Arc<dyn RunStore> = Arc::new(InMemoryRunStore::new());
let replay_store: Arc<dyn ReplayStore> = Arc::new(InMemoryReplayStore::new());
let base = spawn_server(run_store.clone(), replay_store.clone()).await;
let client = reqwest::Client::new();
let launch = client
.post(format!("{base}/v1/tasks"))
.json(&json!({
"blueprint": { "kind": "inline", "value": two_step_blueprint() },
"init_ctx": {},
"goal": "rerun-from happy path",
}))
.send()
.await
.expect("launch request");
assert_eq!(
launch.status(),
reqwest::StatusCode::OK,
"launch body: {}",
launch.text().await.unwrap_or_default()
);
let launched: serde_json::Value = launch.json().await.expect("launch json");
let run_id =
RunId::parse(launched["run_id"].as_str().expect("run_id string")).expect("run_id parse");
let before_entries = replay_store
.list_by_run(&run_id)
.await
.expect("list before rerun");
let refs_before: Vec<String> = before_entries.iter().map(|e| e.step_ref.clone()).collect();
assert_eq!(
refs_before,
vec!["agent-a".to_string(), "agent-b".to_string()],
"seeded replay log shape: {refs_before:?}"
);
let run_before = run_store.get(&run_id).await.expect("run get");
assert_eq!(run_before.status, RunStatus::Done);
let resp = client
.post(format!("{base}/v1/runs/{run_id}/rerun-from"))
.json(&json!({ "from_step": "agent-b" }))
.send()
.await
.expect("rerun request");
assert_eq!(resp.status(), reqwest::StatusCode::ACCEPTED);
let body: serde_json::Value = resp.json().await.expect("rerun json");
assert_eq!(
body["run_id"].as_str(),
Some(run_id.to_string().as_str()),
"rerun-from must not mint a new run_id"
);
assert_eq!(
body["replayed_steps"].as_u64(),
Some(1),
"one entry (agent-a) survives the cut"
);
assert_eq!(
body["dropped_steps"].as_u64(),
Some(1),
"one entry (agent-b) is dropped"
);
let mut terminal = None;
for _ in 0..50 {
let rec = run_store.get(&run_id).await.expect("run get");
if !matches!(rec.status, RunStatus::Pending | RunStatus::Running) {
terminal = Some(rec);
break;
}
tokio::time::sleep(Duration::from_millis(100)).await;
}
let rec = terminal.expect("rerun run reached a terminal status within ~5s");
assert_eq!(rec.status, RunStatus::Done, "rerun must complete to Done");
let after_entries = replay_store
.list_by_run(&run_id)
.await
.expect("list after rerun");
let refs_after: Vec<String> = after_entries.iter().map(|e| e.step_ref.clone()).collect();
assert_eq!(
refs_after,
vec!["agent-a".to_string(), "agent-b".to_string()],
"post-rerun replay log carries the fresh agent-b row, not the ghost: {refs_after:?}"
);
}
#[tokio::test]
async fn rerun_from_compile_fail_leaves_replay_intact() {
let run_store: Arc<dyn RunStore> = Arc::new(InMemoryRunStore::new());
let replay_store: Arc<dyn ReplayStore> = Arc::new(InMemoryReplayStore::new());
let run_id = RunId::new();
let task_id = TaskId::new();
let bp = two_step_blueprint_with_unbound_operator();
run_store
.create(seed_run(
&run_id,
&task_id,
RunStatus::Done,
Some(snapshot_json_for(&bp)),
))
.await
.expect("seed run");
let ctx_a = mlua_swarm::core::ctx::Ctx::new(StepId::new(), 1, "agent-a");
replay_store
.append(
ReplayEntry::from_completion(run_id.clone(), "agent-a", "h", 0, &ctx_a, &json!({}))
.expect("entry build"),
)
.await
.expect("seed replay a");
let ctx_b = mlua_swarm::core::ctx::Ctx::new(StepId::new(), 1, "agent-b");
replay_store
.append(
ReplayEntry::from_completion(run_id.clone(), "agent-b", "h", 0, &ctx_b, &json!({}))
.expect("entry build"),
)
.await
.expect("seed replay b");
let base = spawn_server(run_store.clone(), replay_store.clone()).await;
let resp = reqwest::Client::new()
.post(format!("{base}/v1/runs/{run_id}/rerun-from"))
.json(&json!({ "from_step": "agent-b" }))
.send()
.await
.expect("request");
assert_eq!(resp.status(), reqwest::StatusCode::UNPROCESSABLE_ENTITY);
let body = resp.text().await.unwrap_or_default();
assert!(
body.contains("fails to compile"),
"422 body must name the compile failure so the caller can fix the BP: {body}"
);
let after = run_store.get(&run_id).await.expect("run present");
assert_eq!(after.status, RunStatus::Done);
let entries_after = replay_store.list_by_run(&run_id).await.expect("list");
let refs_after: Vec<String> = entries_after.iter().map(|e| e.step_ref.clone()).collect();
assert_eq!(
refs_after,
vec!["agent-a".to_string(), "agent-b".to_string()],
"compile-fail 422 must not truncate the replay log: {refs_after:?}"
);
}
#[tokio::test]
async fn rerun_from_after_consumed_log_returns_helpful_422() {
let run_store: Arc<dyn RunStore> = Arc::new(InMemoryRunStore::new());
let replay_store: Arc<dyn ReplayStore> = Arc::new(InMemoryReplayStore::new());
let run_id = RunId::new();
let task_id = TaskId::new();
let bp = two_step_blueprint();
run_store
.create(seed_run(
&run_id,
&task_id,
RunStatus::Done,
Some(snapshot_json_for(&bp)),
))
.await
.expect("seed run");
for name in ["agent-a", "agent-b"] {
run_store
.append_step_entry(
&run_id,
StepEntry {
step_id: StepId::new(),
step_ref: Some(name.into()),
status: Some("passed".into()),
binding_digest: None,
at: 0,
},
)
.await
.expect("seed step entry");
}
let base = spawn_server(run_store.clone(), replay_store.clone()).await;
let resp = reqwest::Client::new()
.post(format!("{base}/v1/runs/{run_id}/rerun-from"))
.json(&json!({ "from_step": "agent-b" }))
.send()
.await
.expect("request");
assert_eq!(resp.status(), reqwest::StatusCode::UNPROCESSABLE_ENTITY);
let body = resp.text().await.unwrap_or_default();
assert!(
body.contains("consumed by a prior rerun-from"),
"422 body must name the consumed-log condition rather than the \
generic \"not present\" message: {body}"
);
let after = run_store.get(&run_id).await.expect("run present");
assert_eq!(after.status, RunStatus::Done);
}