use axum::body::Body;
use axum::http::{Request, StatusCode};
use http_body_util::BodyExt;
use jamjet_agents::InMemoryAgentRegistry;
use jamjet_api::{routes::build_router_with_opts, state::AppState};
use jamjet_audit::{AuditEnricher, NoopAuditBackend};
use jamjet_core::workflow::{ExecutionId, WorkflowExecution, WorkflowStatus};
use jamjet_state::backend::{StateBackend, WorkItem, WorkflowDefinition};
use jamjet_state::event::EventKind;
use jamjet_state::{Event, SqliteBackend, DEFAULT_TENANT};
use serde_json::Value;
use std::sync::Arc;
use std::time::Duration;
use tower::ServiceExt;
use uuid::Uuid;
fn make_state(backend: Arc<dyn StateBackend>) -> AppState {
let backend_for_fn = backend.clone();
let audit: Arc<dyn jamjet_audit::AuditBackend> = Arc::new(NoopAuditBackend);
let enricher = Arc::new(AuditEnricher::new(Arc::clone(&audit)));
AppState {
backend: backend.clone(),
backend_for_fn: Arc::new(move |_tenant_id: &jamjet_state::TenantId| backend_for_fn.clone()),
agents: Arc::new(InMemoryAgentRegistry::new()),
audit,
enricher,
protocols: jamjet_api::state::default_protocol_registry(),
cron_store: None,
}
}
async fn body_json(body: Body) -> Value {
let bytes = body.collect().await.unwrap().to_bytes();
serde_json::from_slice(&bytes).unwrap()
}
fn java_tool_workflow_ir() -> Value {
let condition = |id: &str| {
serde_json::json!({
"id": id,
"kind": { "type": "condition", "branches": [] },
"retry_policy": null,
"node_timeout_secs": null,
"description": null,
"labels": {}
})
};
serde_json::json!({
"workflow_id": "java-tool-e2e",
"version": "0.1.0",
"name": null,
"description": null,
"state_schema": "",
"start_node": "gate_in",
"nodes": {
"gate_in": condition("gate_in"),
"java_tool_node": {
"id": "java_tool_node",
"kind": {
"type": "java_fn",
"class_name": "com.example.tools.WeatherTool",
"method": "getWeather",
"output_schema": ""
},
"retry_policy": "no_retry",
"node_timeout_secs": null,
"description": null,
"labels": {}
},
"gate_out": condition("gate_out")
},
"edges": [
{ "from": "gate_in", "to": "java_tool_node", "condition": null },
{ "from": "java_tool_node", "to": "gate_out", "condition": null },
{ "from": "gate_out", "to": "end", "condition": null }
],
"retry_policies": {},
"timeouts": {},
"models": {},
"tools": {},
"mcp_servers": {},
"remote_agents": {},
"labels": {}
})
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn java_tool_claim_complete_round_trip_advances_durably() {
let db_path = std::env::temp_dir().join(format!("jjtest-java-{}.db", Uuid::new_v4()));
let url = format!("sqlite://{}", db_path.display());
let backend: Arc<dyn StateBackend> = Arc::new(
SqliteBackend::open(&url)
.await
.expect("open sqlite backend"),
);
backend
.store_workflow(WorkflowDefinition {
workflow_id: "java-tool-e2e".into(),
version: "0.1.0".into(),
ir: java_tool_workflow_ir(),
created_at: chrono::Utc::now(),
tenant_id: DEFAULT_TENANT.into(),
})
.await
.expect("store_workflow");
let execution_id = ExecutionId::new();
let now = chrono::Utc::now();
backend
.create_execution(WorkflowExecution {
execution_id: execution_id.clone(),
workflow_id: "java-tool-e2e".into(),
workflow_version: "0.1.0".into(),
status: WorkflowStatus::Running,
initial_input: serde_json::json!({}),
current_state: serde_json::json!({}),
started_at: now,
updated_at: now,
completed_at: None,
session_type: None,
parent_execution_id: None,
segment_number: 0,
})
.await
.expect("create_execution");
backend
.append_event(Event::new(
execution_id.clone(),
1,
EventKind::WorkflowStarted {
workflow_id: "java-tool-e2e".into(),
workflow_version: "0.1.0".into(),
initial_input: serde_json::json!({}),
},
))
.await
.expect("append WorkflowStarted");
backend
.append_event(Event::new(
execution_id.clone(),
2,
EventKind::NodeScheduled {
node_id: "gate_in".into(),
queue_type: "general".into(),
},
))
.await
.expect("append NodeScheduled");
backend
.enqueue_work_item(WorkItem {
id: Uuid::new_v4(),
execution_id: execution_id.clone(),
node_id: "gate_in".into(),
queue_type: "general".into(),
payload: serde_json::json!({
"workflow_id": "java-tool-e2e",
"workflow_version": "0.1.0",
}),
attempt: 0,
max_attempts: 3,
created_at: now,
lease_expires_at: None,
worker_id: None,
lease_fence: 0,
tenant_id: DEFAULT_TENANT.into(),
})
.await
.expect("enqueue_work_item");
let scheduler = jamjet_scheduler::Scheduler::new(backend.clone())
.with_poll_interval(Duration::from_millis(25));
let sched_handle = tokio::spawn(async move { scheduler.run().await });
let worker_handles = jamjet_worker::default_pool(backend.clone()).spawn();
let state = make_state(backend.clone());
let claimed = tokio::time::timeout(Duration::from_secs(15), async {
loop {
let body = serde_json::json!({
"worker_id": "java-tool-worker-0",
"queue_types": ["java_tool"],
});
let resp = build_router_with_opts(state.clone(), true)
.oneshot(
Request::post("/work-items/claim")
.header("content-type", "application/json")
.body(Body::from(serde_json::to_vec(&body).unwrap()))
.unwrap(),
)
.await
.unwrap();
assert_eq!(resp.status(), StatusCode::OK, "claim must not error");
let json = body_json(resp.into_body()).await;
if json["claimed"] == true {
return json["work_item"].clone();
}
tokio::time::sleep(Duration::from_millis(50)).await;
}
})
.await
.expect("a java_tool work item must become claimable");
assert_eq!(claimed["queue_type"], "java_tool");
assert_eq!(claimed["node_id"], "java_tool_node");
assert_eq!(
claimed["payload"]["class"], "com.example.tools.WeatherTool",
"claimed payload must carry the class name"
);
assert_eq!(
claimed["payload"]["method"], "getWeather",
"claimed payload must carry the method name"
);
assert!(
claimed["payload"]["input"].is_object(),
"claimed payload.input must be the accumulated workflow state"
);
let work_item_id = claimed["id"].as_str().expect("work item id").to_string();
let lease_fence = claimed["lease_fence"]
.as_i64()
.expect("claim response must carry lease_fence");
assert!(
lease_fence > 0,
"a claimed item must have a non-zero lease fence"
);
let complete_body = serde_json::json!({
"execution_id": execution_id.to_string(),
"node_id": "java_tool_node",
"output": { "temp_c": 18 },
"state_patch": { "java_result": "sunny" },
"duration_ms": 7,
"lease_fence": lease_fence
});
let resp = build_router_with_opts(state.clone(), true)
.oneshot(
Request::post(format!("/work-items/{work_item_id}/complete"))
.header("content-type", "application/json")
.body(Body::from(serde_json::to_vec(&complete_body).unwrap()))
.unwrap(),
)
.await
.unwrap();
assert_eq!(resp.status(), StatusCode::OK);
assert_eq!(body_json(resp.into_body()).await["completed"], true);
let final_status = tokio::time::timeout(Duration::from_secs(15), async {
loop {
let exec = backend
.get_execution(&execution_id)
.await
.expect("get_execution")
.expect("execution exists");
if matches!(
exec.status,
WorkflowStatus::Completed | WorkflowStatus::Failed
) {
return exec.status;
}
tokio::time::sleep(Duration::from_millis(50)).await;
}
})
.await;
sched_handle.abort();
for h in worker_handles {
h.abort();
}
assert_eq!(
final_status.ok(),
Some(WorkflowStatus::Completed),
"completing the java_tool work item must drive the execution to completion"
);
let exec = backend
.get_execution(&execution_id)
.await
.expect("get_execution")
.expect("execution exists");
assert_eq!(
exec.current_state["java_result"], "sunny",
"the java_tool result must be committed to execution state"
);
let events = backend.get_events(&execution_id).await.expect("get_events");
assert!(
events.iter().any(|e| matches!(
&e.kind,
EventKind::NodeCompleted { node_id, .. } if node_id == "java_tool_node"
)),
"a NodeCompleted event must be committed for the java_tool node"
);
assert!(
events.iter().any(|e| matches!(
&e.kind,
EventKind::NodeScheduled { node_id, .. } if node_id == "gate_out"
)),
"the next node must be scheduled after the java_tool node completes"
);
let _ = std::fs::remove_file(&db_path);
}
fn temp_sqlite_url() -> (std::path::PathBuf, String) {
let db_path = std::env::temp_dir().join(format!("jjtest-fence-{}.db", Uuid::new_v4()));
let url = format!("sqlite://{}", db_path.display());
(db_path, url)
}
async fn seed_claimable_java_item(backend: &Arc<dyn StateBackend>) -> ExecutionId {
let execution_id = ExecutionId::new();
let now = chrono::Utc::now();
let mut ir = java_tool_workflow_ir();
ir["workflow_id"] = serde_json::json!("java-tool-fence");
backend
.store_workflow(WorkflowDefinition {
workflow_id: "java-tool-fence".into(),
version: "0.1.0".into(),
ir,
created_at: now,
tenant_id: DEFAULT_TENANT.into(),
})
.await
.expect("store_workflow");
backend
.create_execution(WorkflowExecution {
execution_id: execution_id.clone(),
workflow_id: "java-tool-fence".into(),
workflow_version: "0.1.0".into(),
status: WorkflowStatus::Running,
initial_input: serde_json::json!({}),
current_state: serde_json::json!({}),
started_at: now,
updated_at: now,
completed_at: None,
session_type: None,
parent_execution_id: None,
segment_number: 0,
})
.await
.expect("create_execution");
backend
.append_event(Event::new(
execution_id.clone(),
1,
EventKind::WorkflowStarted {
workflow_id: "java-tool-fence".into(),
workflow_version: "0.1.0".into(),
initial_input: serde_json::json!({}),
},
))
.await
.expect("append WorkflowStarted");
backend
.enqueue_work_item(WorkItem {
id: Uuid::new_v4(),
execution_id: execution_id.clone(),
node_id: "java_tool_node".into(),
queue_type: "java_tool".into(),
payload: serde_json::json!({
"workflow_id": "java-tool-fence",
"workflow_version": "0.1.0",
"node_id": "java_tool_node",
"class": "com.example.tools.WeatherTool",
"method": "getWeather",
"input": {},
}),
attempt: 0,
max_attempts: 3,
created_at: now,
lease_expires_at: None,
worker_id: None,
lease_fence: 0,
tenant_id: DEFAULT_TENANT.into(),
})
.await
.expect("enqueue_work_item");
execution_id
}
async fn http_claim_java(state: &AppState) -> Value {
let body = serde_json::json!({
"worker_id": "java-fence-worker",
"queue_types": ["java_tool"],
});
let resp = build_router_with_opts(state.clone(), true)
.oneshot(
Request::post("/work-items/claim")
.header("content-type", "application/json")
.body(Body::from(serde_json::to_vec(&body).unwrap()))
.unwrap(),
)
.await
.unwrap();
assert_eq!(resp.status(), StatusCode::OK, "claim must not error");
let json = body_json(resp.into_body()).await;
assert_eq!(json["claimed"], true, "an item must be claimable");
json["work_item"].clone()
}
async fn http_complete(state: &AppState, item_id: &str, body: Value) -> (StatusCode, Value) {
let resp = build_router_with_opts(state.clone(), true)
.oneshot(
Request::post(format!("/work-items/{item_id}/complete"))
.header("content-type", "application/json")
.body(Body::from(serde_json::to_vec(&body).unwrap()))
.unwrap(),
)
.await
.unwrap();
let status = resp.status();
(status, body_json(resp.into_body()).await)
}
async fn count_node_completed(
backend: &Arc<dyn StateBackend>,
eid: &ExecutionId,
node_id: &str,
) -> usize {
backend
.get_events(eid)
.await
.expect("get_events")
.iter()
.filter(|e| matches!(&e.kind, EventKind::NodeCompleted { node_id: n, .. } if n == node_id))
.count()
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn complete_with_matching_fence_succeeds() {
let (db_path, url) = temp_sqlite_url();
let backend: Arc<dyn StateBackend> =
Arc::new(SqliteBackend::open(&url).await.expect("open sqlite"));
let execution_id = seed_claimable_java_item(&backend).await;
let state = make_state(backend.clone());
let claimed = http_claim_java(&state).await;
let fence = claimed["lease_fence"]
.as_i64()
.expect("claim response must carry lease_fence");
assert!(fence > 0, "a claimed item must have a non-zero lease fence");
let item_id = claimed["id"].as_str().unwrap().to_string();
let (status, body) = http_complete(
&state,
&item_id,
serde_json::json!({
"execution_id": execution_id.to_string(),
"node_id": "java_tool_node",
"output": { "temp_c": 18 },
"state_patch": { "java_result": "sunny" },
"duration_ms": 7,
"lease_fence": fence
}),
)
.await;
assert_eq!(status, StatusCode::OK, "matching fence must be accepted");
assert_eq!(body["completed"], true);
assert_eq!(
count_node_completed(&backend, &execution_id, "java_tool_node").await,
1,
"exactly one NodeCompleted must be committed"
);
let exec = backend.get_execution(&execution_id).await.unwrap().unwrap();
assert_eq!(
exec.current_state["java_result"], "sunny",
"the result must land in durable state"
);
let _ = std::fs::remove_file(&db_path);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn complete_with_stale_or_forged_fence_is_rejected() {
let (db_path, url) = temp_sqlite_url();
let backend: Arc<dyn StateBackend> =
Arc::new(SqliteBackend::open(&url).await.expect("open sqlite"));
let execution_id = seed_claimable_java_item(&backend).await;
let state = make_state(backend.clone());
let first = http_claim_java(&state).await;
let item_id = first["id"].as_str().unwrap().to_string();
let fence_1 = first["lease_fence"].as_i64().expect("F1 fence");
let wi_id = Uuid::parse_str(&item_id).unwrap();
let past = (chrono::Utc::now() - chrono::Duration::seconds(1)).to_rfc3339();
let parked = backend
.park_work_item(wi_id, fence_1, &past, 1)
.await
.expect("park_work_item");
assert!(parked, "parking with the held fence must succeed");
let second = http_claim_java(&state).await;
assert_eq!(
second["id"].as_str().unwrap(),
item_id,
"the reclaim must be the same item"
);
let fence_2 = second["lease_fence"].as_i64().expect("F2 fence");
assert!(
fence_2 > fence_1,
"a reclaim must mint a strictly-greater fence (F2={fence_2} > F1={fence_1})"
);
let (status, body) = http_complete(
&state,
&item_id,
serde_json::json!({
"execution_id": execution_id.to_string(),
"node_id": "java_tool_node",
"output": { "stale": true },
"state_patch": { "java_result": "STALE" },
"duration_ms": 1,
"lease_fence": fence_1
}),
)
.await;
assert_eq!(
status,
StatusCode::CONFLICT,
"a stale fence must be rejected with 409"
);
assert_eq!(
body["completed"], false,
"the stale completion must report completed=false"
);
assert_eq!(
count_node_completed(&backend, &execution_id, "java_tool_node").await,
0,
"a stale completion must NOT append a NodeCompleted event"
);
let (forged_status, forged_body) = http_complete(
&state,
&item_id,
serde_json::json!({
"execution_id": execution_id.to_string(),
"node_id": "java_tool_node",
"output": { "forged": true },
"state_patch": { "java_result": "FORGED" },
"duration_ms": 1,
"lease_fence": 999_999_999_i64
}),
)
.await;
assert_eq!(
forged_status,
StatusCode::CONFLICT,
"a forged fence must be rejected"
);
assert_eq!(forged_body["completed"], false);
assert_eq!(
count_node_completed(&backend, &execution_id, "java_tool_node").await,
0,
"a forged completion must NOT append a NodeCompleted event"
);
let exec = backend.get_execution(&execution_id).await.unwrap().unwrap();
assert!(
exec.current_state.get("java_result").is_none(),
"no rejected state_patch may be applied"
);
let (ok_status, ok_body) = http_complete(
&state,
&item_id,
serde_json::json!({
"execution_id": execution_id.to_string(),
"node_id": "java_tool_node",
"output": { "ok": true },
"state_patch": { "java_result": "fresh" },
"duration_ms": 2,
"lease_fence": fence_2
}),
)
.await;
assert_eq!(
ok_status,
StatusCode::OK,
"the current fence holder must settle"
);
assert_eq!(ok_body["completed"], true);
assert_eq!(
count_node_completed(&backend, &execution_id, "java_tool_node").await,
1,
"the legitimate completion appends exactly one NodeCompleted"
);
let _ = std::fs::remove_file(&db_path);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn complete_without_fence_is_backward_compatible() {
let (db_path, url) = temp_sqlite_url();
let backend: Arc<dyn StateBackend> =
Arc::new(SqliteBackend::open(&url).await.expect("open sqlite"));
let execution_id = seed_claimable_java_item(&backend).await;
let state = make_state(backend.clone());
let claimed = http_claim_java(&state).await;
let item_id = claimed["id"].as_str().unwrap().to_string();
let (status, body) = http_complete(
&state,
&item_id,
serde_json::json!({
"execution_id": execution_id.to_string(),
"node_id": "java_tool_node",
"output": { "legacy": true },
"state_patch": { "java_result": "legacy" },
"duration_ms": 3
}),
)
.await;
assert_eq!(status, StatusCode::OK, "absent fence must keep working");
assert_eq!(body["completed"], true);
assert_eq!(
count_node_completed(&backend, &execution_id, "java_tool_node").await,
1,
"the unfenced path still emits NodeCompleted"
);
let _ = std::fs::remove_file(&db_path);
}