use axum::{
extract::{Path, Query, State},
http::StatusCode,
Json,
};
use mlua_swarm::application::{TaskApplicationError, TaskApplicationInput, TaskApplicationOutput};
use mlua_swarm::store::run::{RunContext, RunRecord, RunStatus, RunStoreError};
use mlua_swarm::store::task::{TaskRecord, TaskRecordStatus, TaskStoreError};
use mlua_swarm::{Role, RunId, TaskId};
use serde::{Deserialize, Serialize};
use std::time::Duration;
use crate::{ApiError, AppState};
pub(crate) fn now_secs() -> u64 {
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map(|d| d.as_secs())
.unwrap_or(0)
}
pub(crate) async fn finalize_run(
state: &AppState,
task_id: &TaskId,
run_id: &RunId,
outcome: Result<TaskApplicationOutput, TaskApplicationError>,
) -> Result<TaskApplicationOutput, TaskApplicationError> {
match &outcome {
Ok(out) => {
if let Err(e) = state
.run_store
.set_result(run_id, out.final_ctx.clone())
.await
{
tracing::warn!(%run_id, error = %e, "finalize_run: set_result failed");
}
if let Err(e) = state.run_store.update_status(run_id, RunStatus::Done).await {
tracing::warn!(%run_id, error = %e, "finalize_run: run update_status(Done) failed");
}
if let Err(e) = state
.task_store
.update_status(task_id, TaskRecordStatus::Done)
.await
{
tracing::warn!(%task_id, error = %e, "finalize_run: task update_status(Done) failed");
}
}
Err(e) => {
if let Err(store_err) = state
.run_store
.update_status(run_id, RunStatus::Failed)
.await
{
tracing::warn!(%run_id, error = %store_err, "finalize_run: run update_status(Failed) failed");
}
if let Err(store_err) = state
.task_store
.update_status(task_id, TaskRecordStatus::Failed)
.await
{
tracing::warn!(%task_id, error = %store_err, "finalize_run: task update_status(Failed) failed");
}
tracing::warn!(%task_id, %run_id, error = %e, "finalize_run: dispatch failed");
}
}
outcome
}
#[derive(Debug, Deserialize, Default)]
pub struct TasksListQuery {
#[serde(default)]
pub limit: Option<usize>,
}
pub async fn tasks_list(
State(state): State<AppState>,
Query(q): Query<TasksListQuery>,
) -> Result<Json<Vec<TaskRecord>>, ApiError> {
let mut records = state.task_store.list().await.map_err(ApiError::engine)?;
if let Some(limit) = q.limit {
records.truncate(limit);
}
Ok(Json(records))
}
#[derive(Debug, Clone, Serialize)]
pub struct TaskDetailResponse {
pub task: TaskRecord,
pub runs: Vec<RunRecord>,
}
pub async fn task_get(
State(state): State<AppState>,
Path(id): Path<String>,
) -> Result<Json<TaskDetailResponse>, ApiError> {
let task_id =
TaskId::parse(id).map_err(|e| ApiError::bad_request(format!("invalid task id: {e}")))?;
let task = state
.task_store
.get(&task_id)
.await
.map_err(map_task_store_err)?;
let runs = state
.run_store
.list_by_task(&task_id)
.await
.map_err(ApiError::engine)?;
Ok(Json(TaskDetailResponse { task, runs }))
}
#[derive(Debug, Clone, Serialize)]
pub struct RunKickResponse {
pub task_id: TaskId,
pub run_id: RunId,
}
pub async fn task_rekick(
State(state): State<AppState>,
Path(id): Path<String>,
) -> Result<(StatusCode, Json<RunKickResponse>), ApiError> {
let task_id =
TaskId::parse(id).map_err(|e| ApiError::bad_request(format!("invalid task id: {e}")))?;
let task = state
.task_store
.get(&task_id)
.await
.map_err(map_task_store_err)?;
let blueprint: mlua_swarm::application::BlueprintRef =
serde_json::from_value(task.blueprint_ref.clone()).map_err(|e| {
ApiError::bad_request(format!(
"task {task_id}: stored blueprint_ref failed to decode: {e}"
))
})?;
let run_id = RunId::new();
let now = now_secs();
state
.task_store
.update_status(&task_id, TaskRecordStatus::Running)
.await
.map_err(ApiError::engine)?;
state
.run_store
.create(RunRecord {
id: run_id.clone(),
task_id: task_id.clone(),
status: RunStatus::Running,
step_entries: Vec::new(),
operator_sid: None,
result_ref: None,
created_at: now,
updated_at: now,
})
.await
.map_err(ApiError::engine)?;
let input = TaskApplicationInput::automate(
blueprint,
"http-run",
Role::Operator,
Duration::from_secs(crate::default_run_ttl()),
task.input_ctx.clone(),
);
let run_ctx = RunContext {
run_id: run_id.clone(),
run_store: state.run_store.clone(),
};
let outcome = state.task_app.handle_with_run(input, Some(run_ctx)).await;
finalize_run(&state, &task_id, &run_id, outcome)
.await
.map_err(|e| ApiError::bad_request(format!("run: {e}")))?;
Ok((
StatusCode::CREATED,
Json(RunKickResponse { task_id, run_id }),
))
}
pub async fn run_get(
State(state): State<AppState>,
Path(id): Path<String>,
) -> Result<Json<RunRecord>, ApiError> {
let run_id =
RunId::parse(id).map_err(|e| ApiError::bad_request(format!("invalid run id: {e}")))?;
let run = state
.run_store
.get(&run_id)
.await
.map_err(map_run_store_err)?;
Ok(Json(run))
}
fn map_task_store_err(e: TaskStoreError) -> ApiError {
match e {
TaskStoreError::NotFound(id) => ApiError::not_found(format!("task not found: {id}")),
other => ApiError::engine(other),
}
}
fn map_run_store_err(e: RunStoreError) -> ApiError {
match e {
RunStoreError::NotFound(id) => ApiError::not_found(format!("run not found: {id}")),
other => ApiError::engine(other),
}
}
#[cfg(test)]
mod tests {
use super::*;
use mlua_swarm::application::BlueprintRef;
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::output::InMemoryOutputStore;
use mlua_swarm::store::run::InMemoryRunStore;
use mlua_swarm::store::task::InMemoryTaskStore;
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::Mutex;
fn identity_blueprint() -> Blueprint {
Blueprint {
schema_version: current_schema_version(),
id: "tasks-test-bp".into(),
flow: serde_json::from_value(serde_json::json!({
"kind": "step",
"ref": mlua_swarm::worker::baseline::AG_IDENTITY,
"in": {"op": "lit", "value": "hello"},
"out": {"op": "path", "at": "$.out"},
}))
.expect("flow parse"),
agents: vec![AgentDef {
name: mlua_swarm::worker::baseline::AG_IDENTITY.into(),
kind: AgentKind::RustFn,
spec: serde_json::json!({"fn_id": mlua_swarm::worker::baseline::AG_IDENTITY}),
profile: None,
meta: None,
}],
operators: vec![],
hints: CompilerHints::default(),
strategy: CompilerStrategy::default(),
metadata: BlueprintMetadata::default(),
spawner_hints: Default::default(),
default_agent_kind: AgentKind::Operator,
default_operator_kind: None,
}
}
fn test_state() -> AppState {
let engine = Engine::new_with_layers(EngineCfg::default(), crate::default_layer_registry());
let compiler = mlua_swarm::Compiler::new(crate::default_registry());
let launch = Arc::new(mlua_swarm::TaskLaunchService::new(engine.clone(), compiler));
AppState {
engine,
sessions: Arc::new(Mutex::new(crate::SessionStore::default())),
task_app: Arc::new(mlua_swarm::TaskApplication::new_inline_only(launch)),
ws_operator_factory: None,
data_store: Arc::new(InMemoryOutputStore::new()),
operator_sessions: Arc::new(Mutex::new(HashMap::new())),
roles_to_sid: Arc::new(Mutex::new(HashMap::new())),
task_store: Arc::new(InMemoryTaskStore::new()),
run_store: Arc::new(InMemoryRunStore::new()),
base_url: None,
}
}
fn post_tasks_req(goal: &str) -> crate::FlowTasksReq {
crate::FlowTasksReq {
blueprint: BlueprintRef::Inline {
value: Box::new(identity_blueprint()),
},
init_ctx: serde_json::json!({"in": "hello"}),
ttl_secs: None,
operator: None,
operator_sid: None,
goal: Some(goal.to_string()),
}
}
#[test]
fn task_id_serializes_as_bare_string() {
let v = serde_json::to_value(TaskId::parse("T-abc").unwrap()).expect("serialize");
assert_eq!(v, serde_json::json!("T-abc"));
}
#[tokio::test]
async fn post_then_get_drill_down() {
let state = test_state();
let posted = crate::tasks_start(State(state.clone()), Json(post_tasks_req("smoke goal")))
.await
.expect("tasks_start")
.0;
let task_id = posted.task_id.clone();
let run_id = posted.run_id.clone();
let list = tasks_list(State(state.clone()), Query(TasksListQuery { limit: None }))
.await
.expect("tasks_list")
.0;
assert!(
list.iter().any(|t| t.id == task_id),
"task {task_id} missing from list of {} tasks",
list.len()
);
let detail = task_get(State(state.clone()), Path(task_id.to_string()))
.await
.expect("task_get")
.0;
assert_eq!(detail.task.id, task_id);
assert_eq!(detail.task.goal, "smoke goal");
assert_eq!(detail.task.status, TaskRecordStatus::Done);
assert_eq!(detail.runs.len(), 1);
assert_eq!(detail.runs[0].id, run_id);
assert_eq!(detail.runs[0].status, RunStatus::Done);
let run = run_get(State(state.clone()), Path(run_id.to_string()))
.await
.expect("run_get")
.0;
assert_eq!(run.id, run_id);
assert_eq!(run.task_id, task_id);
assert_eq!(run.result_ref, Some(posted.final_ctx));
assert_eq!(
run.step_entries.len(),
1,
"expected one step_entry for the 1-step identity Blueprint, got {:?}",
run.step_entries
);
assert_eq!(
run.step_entries[0].step_ref,
Some(mlua_swarm::worker::baseline::AG_IDENTITY.to_string())
);
assert_eq!(run.step_entries[0].status, Some("passed".to_string()));
}
#[tokio::test]
async fn rekick_adds_a_second_run_to_the_same_task() {
let state = test_state();
let posted = crate::tasks_start(State(state.clone()), Json(post_tasks_req("rekick goal")))
.await
.expect("tasks_start")
.0;
let task_id = posted.task_id.clone();
let first_run_id = posted.run_id.clone();
let (status, rekicked) = task_rekick(State(state.clone()), Path(task_id.to_string()))
.await
.expect("task_rekick");
assert_eq!(status, StatusCode::CREATED);
let second_run_id = rekicked.0.run_id.clone();
assert_ne!(first_run_id, second_run_id);
let detail = task_get(State(state.clone()), Path(task_id.to_string()))
.await
.expect("task_get")
.0;
assert_eq!(
detail.runs.len(),
2,
"expected 2 runs, got {:?}",
detail.runs
);
let ids: Vec<&RunId> = detail.runs.iter().map(|r| &r.id).collect();
assert!(ids.contains(&&first_run_id));
assert!(ids.contains(&&second_run_id));
let first_run = detail
.runs
.iter()
.find(|r| r.id == first_run_id)
.expect("first run present in detail.runs");
let second_run = detail
.runs
.iter()
.find(|r| r.id == second_run_id)
.expect("second run present in detail.runs");
assert_eq!(
first_run.step_entries.len(),
1,
"first run step_entries: {:?}",
first_run.step_entries
);
assert_eq!(
second_run.step_entries.len(),
1,
"second run step_entries: {:?}",
second_run.step_entries
);
assert_eq!(
first_run.step_entries[0].step_ref,
Some(mlua_swarm::worker::baseline::AG_IDENTITY.to_string())
);
assert_eq!(
second_run.step_entries[0].step_ref,
Some(mlua_swarm::worker::baseline::AG_IDENTITY.to_string())
);
assert_eq!(first_run.step_entries[0].status, Some("passed".to_string()));
assert_eq!(
second_run.step_entries[0].status,
Some("passed".to_string())
);
assert_ne!(
first_run.step_entries[0].step_id, second_run.step_entries[0].step_id,
"each kick dispatches its own StepId — runs must not share step_entries"
);
}
#[tokio::test]
async fn rekick_unknown_task_returns_404() {
let state = test_state();
match task_rekick(State(state), Path("T-does-not-exist".to_string())).await {
Ok(_) => panic!("expected 404 for an unknown task"),
Err(e) => assert_eq!(e.status, StatusCode::NOT_FOUND),
}
}
#[tokio::test]
async fn run_get_unknown_id_returns_404() {
let state = test_state();
match run_get(State(state), Path("R-does-not-exist".to_string())).await {
Ok(_) => panic!("expected 404 for an unknown run"),
Err(e) => assert_eq!(e.status, StatusCode::NOT_FOUND),
}
}
#[tokio::test]
async fn task_get_unknown_id_returns_404() {
let state = test_state();
match task_get(State(state), Path("T-does-not-exist".to_string())).await {
Ok(_) => panic!("expected 404 for an unknown task"),
Err(e) => assert_eq!(e.status, StatusCode::NOT_FOUND),
}
}
}