use axum::{
extract::{Query, State},
http::{header, header::AUTHORIZATION, HeaderMap, StatusCode},
Json,
};
use mlua_swarm::core::agent_context::StepPointer;
use mlua_swarm::core::step_naming::StepNaming;
use mlua_swarm::store::run::{DegradationEntry, RunStatus, RunStoreError};
use mlua_swarm::{CapToken, ContentRef, OutputEvent, RunId, StepId, WorkerPayload};
use mlua_swarm_schema::ContextPolicy;
use serde::Deserialize;
use serde_json::Value;
use crate::projection::McpQueryAdapter;
use crate::{ApiError, AppState};
#[derive(Debug, Deserialize)]
pub struct PromptQuery {
pub task_id: StepId,
}
pub async fn worker_prompt(
State(state): State<AppState>,
headers: HeaderMap,
Query(q): Query<PromptQuery>,
) -> Result<Json<WorkerPayload>, ApiError> {
let task_id = q.task_id;
let bearer = extract_bearer_raw(&headers)?;
let mut payload = if let Some(handle) = parse_worker_handle(&bearer) {
let resolved = state
.engine
.task_id_from_handle(handle)
.await
.map_err(|e| ApiError::engine(format!("task_id_from_handle: {e}")))?;
if resolved != task_id {
return Err(ApiError::bad_request(format!(
"handle {handle} is bound to task {resolved}, not {task_id}"
)));
}
state
.engine
.fetch_worker_payload_trusted(&task_id)
.await
.map_err(|e| ApiError::engine(format!("fetch_worker_payload_trusted: {e}")))?
} else {
let token = CapToken::decode(bearer.trim())
.map_err(|e| ApiError::bad_request(format!("invalid token: {e}")))?;
state
.engine
.fetch_worker_payload(&token, &task_id)
.await
.map_err(|e| ApiError::engine(format!("fetch_worker_payload: {e}")))?
};
assemble_step_pointers(&state, &mut payload).await;
Ok(Json(payload))
}
async fn assemble_step_pointers(state: &AppState, payload: &mut WorkerPayload) {
let Some(context) = payload.context.as_mut() else {
return;
};
let Some(run_id_str) = context.run_id.clone() else {
return;
};
let Ok(run_id) = RunId::parse(run_id_str) else {
return;
};
let adapter = McpQueryAdapter::new(
state.data_store.clone(),
state.run_store.clone(),
state.engine.clone(),
);
let Ok((run, resolved_steps)) = adapter.list_steps_by_run_id(&run_id).await else {
return;
};
let naming = state.engine.step_naming_for(&payload.task_id).await;
let policy = state
.engine
.context_policy_for(&payload.task_id, payload.attempt)
.await;
let self_canonical = naming
.as_deref()
.and_then(|n| n.canonical_of_producer(&payload.agent))
.map(str::to_string)
.unwrap_or_else(|| payload.agent.clone());
let mut pointers = Vec::new();
for step in &resolved_steps {
if step.name == self_canonical
|| !allows_step_canonical(&policy, naming.as_deref(), &step.name)
{
continue;
}
if let Some((size_bytes, file_path, content_url, sha256)) =
crate::projection::resolve_step_pointer_fields(state, &run, step).await
{
pointers.push(StepPointer {
name: step.name.clone(),
size_bytes,
file_path,
content_url,
sha256,
});
}
}
context.steps = pointers;
}
fn allows_step_canonical(
policy: &ContextPolicy,
naming: Option<&StepNaming>,
canonical_name: &str,
) -> bool {
let resolves_to = |raw: &str| -> bool {
match naming {
Some(n) => n
.resolve(raw)
.map(|c| c == canonical_name)
.unwrap_or(raw == canonical_name),
None => raw == canonical_name,
}
};
if policy
.steps_exclude
.iter()
.any(|excluded| resolves_to(excluded))
{
return false;
}
match &policy.steps {
None => true,
Some(list) => list.iter().any(|included| resolves_to(included)),
}
}
#[derive(Debug, Deserialize)]
pub struct WorkerResultReq {
pub task_id: StepId,
pub value: Value,
#[serde(default = "default_ok_true")]
pub ok: bool,
#[serde(default)]
pub attempt: Option<u32>,
}
fn default_ok_true() -> bool {
true
}
pub async fn worker_result(
State(state): State<AppState>,
headers: HeaderMap,
Json(req): Json<WorkerResultReq>,
) -> Result<StatusCode, ApiError> {
let token = decode_worker_bearer(&headers)?;
let task_id = req.task_id.clone();
let attempt = match req.attempt {
Some(n) => n,
None => state
.engine
.task_attempt(&task_id)
.await
.map_err(|e| ApiError::engine(format!("task_attempt: {e}")))?,
};
let event = OutputEvent::Final {
content: ContentRef::Inline {
value: req.value.clone(),
},
ok: req.ok,
};
state
.engine
.submit_output(&token, &task_id, attempt, event)
.await
.map_err(|e| ApiError::engine(format!("submit_output: {e}")))?;
state
.engine
.post_result(&token, &task_id, req.value)
.await
.map_err(|e| ApiError::engine(format!("post_result: {e}")))?;
Ok(StatusCode::NO_CONTENT)
}
#[derive(Debug, Deserialize, Default)]
pub struct SubmitQuery {
#[serde(default)]
pub ok: Option<bool>,
}
pub async fn worker_submit(
State(state): State<AppState>,
headers: HeaderMap,
Query(q): Query<SubmitQuery>,
body: axum::body::Bytes,
) -> Result<StatusCode, ApiError> {
let bearer = extract_bearer_raw(&headers)?;
let task_id = if let Some(handle) = parse_worker_handle(&bearer) {
state
.engine
.task_id_from_handle(handle)
.await
.map_err(|e| ApiError::engine(format!("task_id_from_handle: {e}")))?
} else {
let token = CapToken::decode(bearer.trim())
.map_err(|e| ApiError::bad_request(format!("invalid token: {e}")))?;
state
.engine
.task_id_from_token(&token)
.await
.map_err(|e| ApiError::engine(format!("task_id_from_token: {e}")))?
};
let attempt = state
.engine
.task_attempt(&task_id)
.await
.map_err(|e| ApiError::engine(format!("task_attempt: {e}")))?;
reject_if_run_terminal(&state, &task_id, attempt).await?;
let body_str = String::from_utf8_lossy(&body).trim_end().to_string();
let value = Value::String(body_str);
let ok = q.ok.unwrap_or(true);
state
.engine
.submit_worker_result_trusted(&task_id, attempt, value, ok)
.await
.map_err(|e| ApiError::engine(format!("submit_worker_result_trusted: {e}")))?;
Ok(StatusCode::NO_CONTENT)
}
#[derive(Debug, Deserialize)]
pub struct ArtifactQuery {
pub name: String,
}
pub async fn worker_artifact(
State(state): State<AppState>,
headers: HeaderMap,
Query(q): Query<ArtifactQuery>,
body: axum::body::Bytes,
) -> Result<StatusCode, ApiError> {
let name = q.name.trim();
if name.is_empty() {
return Err(ApiError::bad_request("name must not be empty".into()));
}
let name = name.to_string();
let bearer = extract_bearer_raw(&headers)?;
let task_id = if let Some(handle) = parse_worker_handle(&bearer) {
state
.engine
.task_id_from_handle(handle)
.await
.map_err(|e| ApiError::engine(format!("task_id_from_handle: {e}")))?
} else {
let token = CapToken::decode(bearer.trim())
.map_err(|e| ApiError::bad_request(format!("invalid token: {e}")))?;
state
.engine
.task_id_from_token(&token)
.await
.map_err(|e| ApiError::engine(format!("task_id_from_token: {e}")))?
};
let attempt = state
.engine
.task_attempt(&task_id)
.await
.map_err(|e| ApiError::engine(format!("task_attempt: {e}")))?;
reject_if_run_terminal(&state, &task_id, attempt).await?;
let body_str = String::from_utf8_lossy(&body).trim_end().to_string();
let value = Value::String(body_str);
state
.engine
.stage_worker_artifact_trusted(&task_id, attempt, name, value)
.await
.map_err(|e| ApiError::engine(format!("stage_worker_artifact_trusted: {e}")))?;
Ok(StatusCode::NO_CONTENT)
}
#[derive(Debug, Deserialize)]
pub struct DegradationBody {
pub tool: String,
pub error: String,
pub fallback: String,
#[serde(default)]
pub note: Option<String>,
}
pub async fn worker_degradation(
State(state): State<AppState>,
headers: HeaderMap,
Json(body): Json<DegradationBody>,
) -> Result<StatusCode, ApiError> {
let bearer = extract_bearer_raw(&headers)?;
let task_id = if let Some(handle) = parse_worker_handle(&bearer) {
state
.engine
.task_id_from_handle(handle)
.await
.map_err(|e| ApiError::engine(format!("task_id_from_handle: {e}")))?
} else {
let token = CapToken::decode(bearer.trim())
.map_err(|e| ApiError::bad_request(format!("invalid token: {e}")))?;
state
.engine
.task_id_from_token(&token)
.await
.map_err(|e| ApiError::engine(format!("task_id_from_token: {e}")))?
};
let attempt = state
.engine
.task_attempt(&task_id)
.await
.map_err(|e| ApiError::engine(format!("task_attempt: {e}")))?;
reject_if_run_terminal(&state, &task_id, attempt).await?;
let tid = task_id.clone();
let (run_id_str, agent) = match state
.engine
.with_state("worker_degradation_run_lookup", move |s| {
s.agent_ctx.get(&(tid, attempt)).and_then(|e| {
e.view
.run_id
.clone()
.map(|run_id| (run_id, e.view.agent.clone()))
})
})
.await
{
Ok(Some(pair)) => pair,
_ => {
tracing::warn!(%task_id, "worker_degradation: no run linkage for this task; entry dropped");
return Ok(StatusCode::NO_CONTENT);
}
};
let Ok(run_id) = RunId::parse(run_id_str) else {
tracing::warn!(%task_id, "worker_degradation: run_id failed to parse; entry dropped");
return Ok(StatusCode::NO_CONTENT);
};
let entry = DegradationEntry {
tool: body.tool,
error: body.error,
fallback: body.fallback,
note: body.note,
step_ref: Some(agent),
attempt: Some(attempt),
at: crate::tasks::now_secs(),
};
match state.run_store.append_degradation(&run_id, entry).await {
Ok(()) => Ok(StatusCode::NO_CONTENT),
Err(RunStoreError::NotFound(_)) => {
tracing::warn!(%task_id, %run_id, "worker_degradation: run not found in run_store; entry dropped");
Ok(StatusCode::NO_CONTENT)
}
Err(e) => Err(ApiError::engine(format!("append_degradation: {e}"))),
}
}
async fn reject_if_run_terminal(
state: &AppState,
task_id: &StepId,
attempt: u32,
) -> Result<(), ApiError> {
let tid = task_id.clone();
let run_id_str = match state
.engine
.with_state("worker_terminal_run_guard", move |s| {
s.agent_ctx
.get(&(tid, attempt))
.and_then(|e| e.view.run_id.clone())
})
.await
{
Ok(Some(rid)) => rid,
_ => return Ok(()),
};
let Ok(run_id) = RunId::parse(run_id_str) else {
return Ok(());
};
let Ok(rec) = state.run_store.get(&run_id).await else {
return Ok(());
};
match rec.status {
RunStatus::Done | RunStatus::Failed | RunStatus::Interrupted => {
Err(ApiError::gone(format!(
"run {run_id} is already terminal ({:?}): this attempt's output cannot be \
delivered to a flow context; re-kick the task (POST /v1/tasks/:id/runs) and \
fetch a fresh prompt",
rec.status
)))
}
RunStatus::Pending | RunStatus::Running => Ok(()),
}
}
#[derive(Debug, Deserialize)]
pub struct PromptSystemQuery {
pub task_id: StepId,
pub attempt: u32,
}
pub async fn worker_prompt_system(
State(state): State<AppState>,
headers: HeaderMap,
Query(q): Query<PromptSystemQuery>,
) -> Result<impl axum::response::IntoResponse, ApiError> {
let task_id = q.task_id;
let attempt = q.attempt;
let bearer = extract_bearer_raw(&headers)?;
if let Some(handle) = parse_worker_handle(&bearer) {
let resolved = state
.engine
.task_id_from_handle(handle)
.await
.map_err(|e| ApiError::engine(format!("task_id_from_handle: {e}")))?;
if resolved != task_id {
return Err(ApiError::bad_request(format!(
"handle {handle} is bound to task {resolved}, not {task_id}"
)));
}
} else {
let token = CapToken::decode(bearer.trim())
.map_err(|e| ApiError::bad_request(format!("invalid token: {e}")))?;
state
.engine
.verify_token_for_task(&token, mlua_swarm::Verb::FetchPrompt, &task_id)
.await
.map_err(|e| ApiError::engine(format!("verify_token_for_task: {e}")))?;
}
let system = state
.engine
.raw_system_prompt(&task_id, attempt)
.await
.map_err(|e| ApiError::engine(format!("raw_system_prompt: {e}")))?
.ok_or_else(|| {
ApiError::not_found(format!(
"no baked system prompt for task {task_id} attempt {attempt}"
))
})?;
Ok((
[(header::CONTENT_TYPE, "text/plain; charset=utf-8")],
system,
))
}
#[derive(Debug, serde::Serialize)]
pub struct AgentRenderSizeResponse {
pub agent: String,
pub last_rendered_bytes: Option<usize>,
}
pub async fn agent_render_size(
State(state): State<AppState>,
axum::extract::Path(name): axum::extract::Path<String>,
) -> Json<AgentRenderSizeResponse> {
let last_rendered_bytes = state.engine.agent_last_rendered_size(&name).await;
Json(AgentRenderSizeResponse {
agent: name,
last_rendered_bytes,
})
}
fn extract_bearer_raw(headers: &HeaderMap) -> Result<String, ApiError> {
let v = headers
.get(AUTHORIZATION)
.ok_or_else(|| ApiError::bad_request("missing Authorization header".into()))?
.to_str()
.map_err(|_| ApiError::bad_request("invalid Authorization header encoding".into()))?;
let s = v
.strip_prefix("Bearer ")
.ok_or_else(|| ApiError::bad_request("Authorization must be 'Bearer <token>'".into()))?
.trim();
if s.is_empty() {
return Err(ApiError::bad_request("Bearer is empty".into()));
}
Ok(s.to_string())
}
fn parse_worker_handle(s: &str) -> Option<&str> {
let s = s.trim();
if s.starts_with("wh-")
&& s.len() >= 5
&& s.len() <= 64
&& s[3..].chars().all(|c| c.is_ascii_alphanumeric())
{
Some(s)
} else {
None
}
}
fn decode_worker_bearer(headers: &HeaderMap) -> Result<CapToken, ApiError> {
let v = headers
.get(AUTHORIZATION)
.ok_or_else(|| ApiError::bad_request("missing Authorization header".into()))?
.to_str()
.map_err(|_| ApiError::bad_request("invalid Authorization header encoding".into()))?;
let encoded = v
.strip_prefix("Bearer ")
.ok_or_else(|| ApiError::bad_request("Authorization must be 'Bearer <token>'".into()))?
.trim();
if encoded.is_empty() {
return Err(ApiError::bad_request("Bearer token is empty".into()));
}
CapToken::decode(encoded).map_err(|e| ApiError::bad_request(format!("invalid token: {e}")))
}
#[cfg(test)]
mod tests {
use super::*;
use axum::response::IntoResponse;
use mlua_swarm::core::agent_context::AgentContextView;
use mlua_swarm::core::config::EngineCfg;
use mlua_swarm::core::engine::Engine;
use mlua_swarm::store::output::{InMemoryOutputStore, OutputStore};
use mlua_swarm::store::run::{InMemoryRunStore, RunRecord, RunStatus, RunStore, StepEntry};
use mlua_swarm::store::task::InMemoryTaskStore;
use mlua_swarm::{RunId, StepId, TaskId};
use serde_json::json;
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::Mutex;
fn test_state(data_store: Arc<dyn OutputStore>, run_store: Arc<dyn RunStore>) -> AppState {
let engine = Engine::new(EngineCfg::default());
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,
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,
base_url: None,
sync_timeout_secs: 300,
}
}
async fn append_final(
data_store: &Arc<dyn OutputStore>,
task_id: &str,
producer: &str,
value: Value,
) {
data_store
.append(
task_id,
1,
producer,
OutputEvent::Final {
content: ContentRef::Inline { value },
ok: true,
},
vec![],
)
.await
.expect("append final");
}
fn step_entry(step_id: &StepId, step_ref: &str) -> StepEntry {
StepEntry {
step_id: step_id.clone(),
step_ref: Some(step_ref.to_string()),
status: Some("passed".to_string()),
at: 0,
}
}
fn run_record(task_id: &TaskId, run_id: &RunId, step_entries: Vec<StepEntry>) -> RunRecord {
RunRecord {
id: run_id.clone(),
task_id: task_id.clone(),
status: RunStatus::Running,
step_entries,
degradations: Vec::new(),
operator_sid: None,
result_ref: None,
created_at: 0,
updated_at: 0,
}
}
fn consumer_payload(consumer_step_id: &StepId, run_id: &RunId) -> WorkerPayload {
WorkerPayload {
task_id: consumer_step_id.clone(),
attempt: 1,
agent: "consumer".to_string(),
system: None,
prompt: String::new(),
context: Some(AgentContextView {
task_id: consumer_step_id.to_string(),
agent: "consumer".to_string(),
attempt: 1,
run_id: Some(run_id.to_string()),
..Default::default()
}),
system_ref: None,
}
}
#[tokio::test]
async fn context_policy_unspecified_yields_every_submitted_step() {
let data_store: Arc<dyn OutputStore> = Arc::new(InMemoryOutputStore::new());
let run_store: Arc<dyn RunStore> = Arc::new(InMemoryRunStore::new());
let task_id = TaskId::new();
let run_id = RunId::new();
let planner_id = StepId::new();
let coder_id = StepId::new();
append_final(
&data_store,
planner_id.as_str(),
"planner",
json!({"plan": "x"}),
)
.await;
append_final(
&data_store,
coder_id.as_str(),
"coder",
json!({"code": "y"}),
)
.await;
run_store
.create(run_record(
&task_id,
&run_id,
vec![
step_entry(&planner_id, "planner"),
step_entry(&coder_id, "coder"),
],
))
.await
.expect("create run");
let state = test_state(data_store, run_store);
let consumer_id = StepId::new();
let mut payload = consumer_payload(&consumer_id, &run_id);
assemble_step_pointers(&state, &mut payload).await;
let names: Vec<&str> = payload
.context
.as_ref()
.expect("context")
.steps
.iter()
.map(|p| p.name.as_str())
.collect();
assert!(names.contains(&"planner"), "names: {names:?}");
assert!(names.contains(&"coder"), "names: {names:?}");
}
#[tokio::test]
async fn context_policy_steps_include_list_filters_to_named_steps() {
let data_store: Arc<dyn OutputStore> = Arc::new(InMemoryOutputStore::new());
let run_store: Arc<dyn RunStore> = Arc::new(InMemoryRunStore::new());
let task_id = TaskId::new();
let run_id = RunId::new();
let planner_id = StepId::new();
let coder_id = StepId::new();
append_final(&data_store, planner_id.as_str(), "planner", json!("x")).await;
append_final(&data_store, coder_id.as_str(), "coder", json!("y")).await;
run_store
.create(run_record(
&task_id,
&run_id,
vec![
step_entry(&planner_id, "planner"),
step_entry(&coder_id, "coder"),
],
))
.await
.expect("create run");
let state = test_state(data_store, run_store);
let consumer_id = StepId::new();
state
.engine
.with_state("test.seed_policy", {
let consumer_id = consumer_id.clone();
move |s| {
s.agent_ctx.insert(
(consumer_id, 1),
mlua_swarm::core::state::AgentCtxEntry {
policy: mlua_swarm_schema::ContextPolicy {
steps: Some(vec!["planner".to_string()]),
..Default::default()
},
..Default::default()
},
);
}
})
.await
.expect("seed policy");
let mut payload = consumer_payload(&consumer_id, &run_id);
assemble_step_pointers(&state, &mut payload).await;
let names: Vec<&str> = payload
.context
.as_ref()
.expect("context")
.steps
.iter()
.map(|p| p.name.as_str())
.collect();
assert_eq!(names, vec!["planner"], "names: {names:?}");
}
#[tokio::test]
async fn context_policy_steps_empty_list_yields_no_pointers() {
let data_store: Arc<dyn OutputStore> = Arc::new(InMemoryOutputStore::new());
let run_store: Arc<dyn RunStore> = Arc::new(InMemoryRunStore::new());
let task_id = TaskId::new();
let run_id = RunId::new();
let planner_id = StepId::new();
append_final(&data_store, planner_id.as_str(), "planner", json!("x")).await;
run_store
.create(run_record(
&task_id,
&run_id,
vec![step_entry(&planner_id, "planner")],
))
.await
.expect("create run");
let state = test_state(data_store, run_store);
let consumer_id = StepId::new();
state
.engine
.with_state("test.seed_policy", {
let consumer_id = consumer_id.clone();
move |s| {
s.agent_ctx.insert(
(consumer_id, 1),
mlua_swarm::core::state::AgentCtxEntry {
policy: mlua_swarm_schema::ContextPolicy {
steps: Some(vec![]),
..Default::default()
},
..Default::default()
},
);
}
})
.await
.expect("seed policy");
let mut payload = consumer_payload(&consumer_id, &run_id);
assemble_step_pointers(&state, &mut payload).await;
assert!(payload.context.expect("context").steps.is_empty());
}
#[tokio::test]
async fn context_policy_steps_exclude_wins_over_steps() {
let data_store: Arc<dyn OutputStore> = Arc::new(InMemoryOutputStore::new());
let run_store: Arc<dyn RunStore> = Arc::new(InMemoryRunStore::new());
let task_id = TaskId::new();
let run_id = RunId::new();
let planner_id = StepId::new();
let coder_id = StepId::new();
append_final(&data_store, planner_id.as_str(), "planner", json!("x")).await;
append_final(&data_store, coder_id.as_str(), "coder", json!("y")).await;
run_store
.create(run_record(
&task_id,
&run_id,
vec![
step_entry(&planner_id, "planner"),
step_entry(&coder_id, "coder"),
],
))
.await
.expect("create run");
let state = test_state(data_store, run_store);
let consumer_id = StepId::new();
state
.engine
.with_state("test.seed_policy", {
let consumer_id = consumer_id.clone();
move |s| {
s.agent_ctx.insert(
(consumer_id, 1),
mlua_swarm::core::state::AgentCtxEntry {
policy: mlua_swarm_schema::ContextPolicy {
steps: Some(vec!["planner".to_string(), "coder".to_string()]),
steps_exclude: vec!["planner".to_string()],
..Default::default()
},
..Default::default()
},
);
}
})
.await
.expect("seed policy");
let mut payload = consumer_payload(&consumer_id, &run_id);
assemble_step_pointers(&state, &mut payload).await;
let names: Vec<&str> = payload
.context
.as_ref()
.expect("context")
.steps
.iter()
.map(|p| p.name.as_str())
.collect();
assert_eq!(names, vec!["coder"], "names: {names:?}");
}
#[tokio::test]
async fn in_flight_step_output_is_visible_before_run_finalizes() {
let data_store: Arc<dyn OutputStore> = Arc::new(InMemoryOutputStore::new());
let run_store: Arc<dyn RunStore> = Arc::new(InMemoryRunStore::new());
let task_id = TaskId::new();
let run_id = RunId::new();
let step1_id = StepId::new();
append_final(
&data_store,
step1_id.as_str(),
"step1",
json!({"step1_out": "hi"}),
)
.await;
let mut run = run_record(&task_id, &run_id, vec![step_entry(&step1_id, "step1")]);
run.status = RunStatus::Running;
run.result_ref = None; run_store.create(run).await.expect("create run");
let state = test_state(data_store, run_store);
let consumer_id = StepId::new();
let mut payload = consumer_payload(&consumer_id, &run_id);
assemble_step_pointers(&state, &mut payload).await;
let steps = &payload.context.expect("context").steps;
assert_eq!(steps.len(), 1);
assert_eq!(steps[0].name, "step1");
}
#[tokio::test]
async fn self_agent_name_is_always_excluded() {
let data_store: Arc<dyn OutputStore> = Arc::new(InMemoryOutputStore::new());
let run_store: Arc<dyn RunStore> = Arc::new(InMemoryRunStore::new());
let task_id = TaskId::new();
let run_id = RunId::new();
let planner_id = StepId::new();
let consumer_prior_id = StepId::new();
append_final(&data_store, planner_id.as_str(), "planner", json!("x")).await;
append_final(
&data_store,
consumer_prior_id.as_str(),
"consumer",
json!("self"),
)
.await;
run_store
.create(run_record(
&task_id,
&run_id,
vec![
step_entry(&planner_id, "planner"),
step_entry(&consumer_prior_id, "consumer"),
],
))
.await
.expect("create run");
let state = test_state(data_store, run_store);
let consumer_id = StepId::new();
let mut payload = consumer_payload(&consumer_id, &run_id);
assemble_step_pointers(&state, &mut payload).await;
let names: Vec<&str> = payload
.context
.as_ref()
.expect("context")
.steps
.iter()
.map(|p| p.name.as_str())
.collect();
assert!(!names.contains(&"consumer"), "names: {names:?}");
assert!(names.contains(&"planner"), "names: {names:?}");
}
#[tokio::test]
async fn step_pointer_serializes_with_no_preview_or_content_bytes() {
let data_store: Arc<dyn OutputStore> = Arc::new(InMemoryOutputStore::new());
let run_store: Arc<dyn RunStore> = Arc::new(InMemoryRunStore::new());
let task_id = TaskId::new();
let run_id = RunId::new();
let planner_id = StepId::new();
append_final(
&data_store,
planner_id.as_str(),
"planner",
json!({"plan": "do the thing, at length".repeat(50)}),
)
.await;
run_store
.create(run_record(
&task_id,
&run_id,
vec![step_entry(&planner_id, "planner")],
))
.await
.expect("create run");
let state = test_state(data_store, run_store);
let consumer_id = StepId::new();
let mut payload = consumer_payload(&consumer_id, &run_id);
assemble_step_pointers(&state, &mut payload).await;
let steps = &payload.context.expect("context").steps;
assert_eq!(steps.len(), 1);
let json_value = serde_json::to_value(&steps[0]).expect("serialize StepPointer");
let obj = json_value.as_object().expect("object");
for forbidden in ["preview", "content", "value", "bytes"] {
assert!(
!obj.contains_key(forbidden),
"StepPointer must not carry a {forbidden:?} field: {obj:?}"
);
}
assert!(obj.contains_key("name"));
assert!(obj.contains_key("size_bytes"));
assert!(obj.contains_key("content_url"));
assert!(obj.contains_key("sha256"));
}
fn declared_name_bp() -> mlua_swarm::blueprint::Blueprint {
use mlua_flow_ir::{Expr, Node};
use mlua_swarm::blueprint::{
current_schema_version, AgentDef, AgentKind, AgentMeta, Blueprint, BlueprintMetadata,
CompilerHints, CompilerStrategy,
};
Blueprint {
schema_version: current_schema_version(),
id: "worker-test-declared-name-bp".into(),
flow: Node::Step {
ref_: "planner".to_string(),
in_: Expr::Path {
at: "$.in".parse().expect("literal test path: $.in"),
},
out: Expr::Path {
at: "$.plan".parse().expect("literal test path: $.plan"),
},
},
agents: vec![AgentDef {
name: "planner".to_string(),
kind: AgentKind::RustFn,
spec: json!({"fn_id": "planner"}),
profile: None,
meta: Some(AgentMeta {
projection_name: Some("plan-out".to_string()),
..Default::default()
}),
}],
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,
}
}
#[tokio::test]
async fn declared_projection_name_pointer_name_is_canonical_and_policy_matches_it() {
let data_store: Arc<dyn OutputStore> = Arc::new(InMemoryOutputStore::new());
let run_store: Arc<dyn RunStore> = Arc::new(InMemoryRunStore::new());
let task_id = TaskId::new();
let run_id = RunId::new();
let planner_id = StepId::new();
append_final(
&data_store,
planner_id.as_str(),
"plan-out",
json!({"plan": "x"}),
)
.await;
run_store
.create(run_record(
&task_id,
&run_id,
vec![step_entry(&planner_id, "planner")],
))
.await
.expect("create run");
let state = test_state(data_store, run_store);
let (naming, _warnings) =
mlua_swarm::core::step_naming::StepNaming::from_blueprint(&declared_name_bp())
.expect("no collision");
let naming = Arc::new(naming);
let consumer_id = StepId::new();
state
.engine
.with_state("test.seed_step_naming", {
let naming = naming.clone();
let planner_id = planner_id.clone();
let consumer_id = consumer_id.clone();
move |s| {
s.step_namings.insert(planner_id, naming.clone());
s.step_namings.insert(consumer_id, naming);
}
})
.await
.expect("seed step naming");
state
.engine
.with_state("test.seed_policy", {
let consumer_id = consumer_id.clone();
move |s| {
s.agent_ctx.insert(
(consumer_id, 1),
mlua_swarm::core::state::AgentCtxEntry {
policy: mlua_swarm_schema::ContextPolicy {
steps: Some(vec!["plan-out".to_string()]),
..Default::default()
},
..Default::default()
},
);
}
})
.await
.expect("seed policy");
let mut payload = consumer_payload(&consumer_id, &run_id);
assemble_step_pointers(&state, &mut payload).await;
let steps = &payload.context.expect("context").steps;
assert_eq!(steps.len(), 1, "steps: {steps:?}");
assert_eq!(
steps[0].name, "plan-out",
"StepPointer.name must be the canonical name"
);
}
async fn seed_task_with_handle(
state: &AppState,
task_id: &StepId,
agent: &str,
attempt: u32,
system: Option<String>,
) -> String {
let handle = format!("wh-{}", mlua_swarm::types::secure_hex(4));
let task_id = task_id.clone();
let agent = agent.to_string();
let handle_clone = handle.clone();
state
.engine
.with_state("test.seed_task_with_handle", move |s| {
let mut task = mlua_swarm::core::state::TaskState::new(
task_id.clone(),
mlua_swarm::core::state::TaskSpec {
agent: agent.clone(),
initial_directive: json!("x"),
step_ctx: None,
},
);
task.attempt = attempt;
s.tasks.insert(task_id.clone(), task);
s.systems.insert((task_id.clone(), attempt), system);
let token = CapToken {
agent_id: agent,
role: mlua_swarm::Role::Worker,
scopes: vec!["*".to_string()],
issued_at: 0,
expire_at: u64::MAX,
max_uses: None,
nonce: format!("test-nonce-{task_id}"),
sig_hex: String::new(),
};
let fp = token.fingerprint();
s.tokens.insert(
fp.clone(),
mlua_swarm::core::state::CapTokenRecord {
token,
uses_left: None,
revoked: false,
task_id: Some(task_id),
},
);
s.worker_handles.insert(handle_clone, fp);
})
.await
.expect("seed_task_with_handle");
handle
}
fn bearer_headers(handle: &str) -> HeaderMap {
let mut headers = HeaderMap::new();
headers.insert(
AUTHORIZATION,
format!("Bearer {handle}").parse().expect("header value"),
);
headers
}
#[tokio::test]
async fn worker_prompt_system_returns_raw_bytes_for_baked_system() {
let data_store: Arc<dyn OutputStore> = Arc::new(InMemoryOutputStore::new());
let run_store: Arc<dyn RunStore> = Arc::new(InMemoryRunStore::new());
let state = test_state(data_store, run_store);
let task_id = StepId::new();
let rendered = "# Hello\n\nThis is the baked system prompt.".to_string();
let handle =
seed_task_with_handle(&state, &task_id, "planner", 1, Some(rendered.clone())).await;
let resp = worker_prompt_system(
State(state.clone()),
bearer_headers(&handle),
Query(PromptSystemQuery {
task_id: task_id.clone(),
attempt: 1,
}),
)
.await
.expect("worker_prompt_system")
.into_response();
assert_eq!(resp.status(), StatusCode::OK);
let content_type = resp
.headers()
.get(header::CONTENT_TYPE)
.expect("content-type header")
.to_str()
.expect("ascii");
assert_eq!(content_type, "text/plain; charset=utf-8");
let body_bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
.await
.expect("body bytes");
assert_eq!(body_bytes.as_ref(), rendered.as_bytes());
}
#[tokio::test]
async fn worker_prompt_system_404s_when_no_baked_system() {
let data_store: Arc<dyn OutputStore> = Arc::new(InMemoryOutputStore::new());
let run_store: Arc<dyn RunStore> = Arc::new(InMemoryRunStore::new());
let state = test_state(data_store, run_store);
let task_id = StepId::new();
let handle = seed_task_with_handle(&state, &task_id, "planner", 1, None).await;
let result = worker_prompt_system(
State(state.clone()),
bearer_headers(&handle),
Query(PromptSystemQuery {
task_id: task_id.clone(),
attempt: 1,
}),
)
.await;
let err = match result {
Ok(_) => panic!("expected 404 ApiError, got Ok"),
Err(e) => e,
};
assert_eq!(err.into_response().status(), StatusCode::NOT_FOUND);
}
#[tokio::test]
async fn worker_prompt_system_rejects_handle_task_mismatch() {
let data_store: Arc<dyn OutputStore> = Arc::new(InMemoryOutputStore::new());
let run_store: Arc<dyn RunStore> = Arc::new(InMemoryRunStore::new());
let state = test_state(data_store, run_store);
let task_id = StepId::new();
let other_task_id = StepId::new();
let handle =
seed_task_with_handle(&state, &task_id, "planner", 1, Some("x".to_string())).await;
let result = worker_prompt_system(
State(state.clone()),
bearer_headers(&handle),
Query(PromptSystemQuery {
task_id: other_task_id,
attempt: 1,
}),
)
.await;
let err = match result {
Ok(_) => panic!("expected 400 ApiError for task mismatch, got Ok"),
Err(e) => e,
};
assert_eq!(err.into_response().status(), StatusCode::BAD_REQUEST);
}
#[tokio::test]
async fn agent_render_size_returns_null_for_unknown_agent() {
let data_store: Arc<dyn OutputStore> = Arc::new(InMemoryOutputStore::new());
let run_store: Arc<dyn RunStore> = Arc::new(InMemoryRunStore::new());
let state = test_state(data_store, run_store);
let Json(body) = agent_render_size(
State(state.clone()),
axum::extract::Path("never-dispatched".to_string()),
)
.await;
assert_eq!(body.agent, "never-dispatched");
assert_eq!(body.last_rendered_bytes, None);
}
#[tokio::test]
async fn agent_render_size_reports_last_rendered_bytes() {
let data_store: Arc<dyn OutputStore> = Arc::new(InMemoryOutputStore::new());
let run_store: Arc<dyn RunStore> = Arc::new(InMemoryRunStore::new());
let state = test_state(data_store, run_store);
let task_id = StepId::new();
state
.engine
.with_state("test.seed_agent_ctx_for_bake", {
let task_id = task_id.clone();
move |s| {
s.tasks.insert(
task_id.clone(),
mlua_swarm::core::state::TaskState::new(
task_id,
mlua_swarm::core::state::TaskSpec {
agent: "coder".to_string(),
initial_directive: json!("x"),
step_ctx: None,
},
),
);
}
})
.await
.expect("seed task");
state
.engine
.bake_worker_system_prompt(&task_id, 1, Some("z".repeat(42)))
.await
.expect("bake_worker_system_prompt");
let Json(body) = agent_render_size(
State(state.clone()),
axum::extract::Path("coder".to_string()),
)
.await;
assert_eq!(body.agent, "coder");
assert_eq!(body.last_rendered_bytes, Some(42));
}
#[tokio::test]
async fn worker_artifact_stages_and_204s_for_valid_request() {
let data_store: Arc<dyn OutputStore> = Arc::new(InMemoryOutputStore::new());
let run_store: Arc<dyn RunStore> = Arc::new(InMemoryRunStore::new());
let state = test_state(data_store, run_store);
let task_id = StepId::new();
let handle = seed_task_with_handle(&state, &task_id, "planner", 1, None).await;
let status = worker_artifact(
State(state.clone()),
bearer_headers(&handle),
Query(ArtifactQuery {
name: "summary".to_string(),
}),
axum::body::Bytes::from_static(b"hello artifact\n"),
)
.await
.expect("worker_artifact");
assert_eq!(status, StatusCode::NO_CONTENT);
let tail = state.engine.output_tail(&task_id, 1).await;
assert_eq!(tail.len(), 1, "tail: {tail:?}");
match &tail[0] {
OutputEvent::Artifact { name, content } => {
assert_eq!(name, "summary");
match content {
ContentRef::Inline { value } => {
assert_eq!(value, &json!("hello artifact"));
}
other => panic!("expected Inline content, got {other:?}"),
}
}
other => panic!("expected Artifact event, got {other:?}"),
}
}
#[tokio::test]
async fn worker_artifact_rejects_blank_name() {
let data_store: Arc<dyn OutputStore> = Arc::new(InMemoryOutputStore::new());
let run_store: Arc<dyn RunStore> = Arc::new(InMemoryRunStore::new());
let state = test_state(data_store, run_store);
let task_id = StepId::new();
let handle = seed_task_with_handle(&state, &task_id, "planner", 1, None).await;
let result = worker_artifact(
State(state.clone()),
bearer_headers(&handle),
Query(ArtifactQuery {
name: " ".to_string(),
}),
axum::body::Bytes::from_static(b"x"),
)
.await;
let err = match result {
Ok(_) => panic!("expected 400 ApiError for blank name, got Ok"),
Err(e) => e,
};
assert_eq!(err.into_response().status(), StatusCode::BAD_REQUEST);
assert!(state.engine.output_tail(&task_id, 1).await.is_empty());
}
#[tokio::test]
async fn worker_artifact_staging_same_name_twice_appends_both_events_in_order() {
let data_store: Arc<dyn OutputStore> = Arc::new(InMemoryOutputStore::new());
let run_store: Arc<dyn RunStore> = Arc::new(InMemoryRunStore::new());
let state = test_state(data_store, run_store);
let task_id = StepId::new();
let handle = seed_task_with_handle(&state, &task_id, "planner", 1, None).await;
for body in [b"first".as_slice(), b"second".as_slice()] {
worker_artifact(
State(state.clone()),
bearer_headers(&handle),
Query(ArtifactQuery {
name: "a".to_string(),
}),
axum::body::Bytes::copy_from_slice(body),
)
.await
.expect("worker_artifact");
}
let tail = state.engine.output_tail(&task_id, 1).await;
assert_eq!(tail.len(), 2, "tail: {tail:?}");
let values: Vec<&str> = tail
.iter()
.map(|ev| match ev {
OutputEvent::Artifact {
content: ContentRef::Inline { value },
..
} => value.as_str().expect("string value"),
other => panic!("expected Artifact/Inline event, got {other:?}"),
})
.collect();
assert_eq!(values, vec!["first", "second"]);
}
async fn link_task_to_run(state: &AppState, task_id: &StepId, attempt: u32, run_id: &RunId) {
let tid = task_id.clone();
let rid_str = run_id.to_string();
state
.engine
.with_state("test.link_task_to_run", move |s| {
let mut entry = mlua_swarm::core::state::AgentCtxEntry::default();
entry.view.run_id = Some(rid_str);
s.agent_ctx.insert((tid, attempt), entry);
})
.await
.expect("link_task_to_run");
}
#[tokio::test]
async fn submit_and_artifact_against_terminal_run_return_410() {
let data_store: Arc<dyn OutputStore> = Arc::new(InMemoryOutputStore::new());
let run_store: Arc<dyn RunStore> = Arc::new(InMemoryRunStore::new());
let state = test_state(data_store, run_store.clone());
let task_id = StepId::new();
let handle = seed_task_with_handle(&state, &task_id, "planner", 1, None).await;
let owner_task = TaskId::new();
let run_id = RunId::new();
let mut rec = run_record(&owner_task, &run_id, vec![]);
rec.status = RunStatus::Failed;
run_store.create(rec).await.expect("run create");
link_task_to_run(&state, &task_id, 1, &run_id).await;
let err = worker_submit(
State(state.clone()),
bearer_headers(&handle),
Query(SubmitQuery { ok: None }),
axum::body::Bytes::from_static(b"LATE OUTPUT"),
)
.await
.expect_err("a submit against a Failed run must be rejected");
assert_eq!(err.status, StatusCode::GONE);
assert!(
err.message.contains(&run_id.to_string()),
"the 410 must name the terminal run: {}",
err.message
);
let err = worker_artifact(
State(state.clone()),
bearer_headers(&handle),
Query(ArtifactQuery {
name: "part.md".to_string(),
}),
axum::body::Bytes::from_static(b"LATE PART"),
)
.await
.expect_err("an artifact staged against a Failed run must be rejected");
assert_eq!(err.status, StatusCode::GONE);
let tail = state.engine.output_tail(&task_id, 1).await;
assert!(tail.is_empty(), "rejected submits must not land: {tail:?}");
}
#[tokio::test]
async fn terminal_run_guard_is_fail_open() {
let data_store: Arc<dyn OutputStore> = Arc::new(InMemoryOutputStore::new());
let run_store: Arc<dyn RunStore> = Arc::new(InMemoryRunStore::new());
let state = test_state(data_store, run_store.clone());
let task_id = StepId::new();
seed_task_with_handle(&state, &task_id, "planner", 1, None).await;
reject_if_run_terminal(&state, &task_id, 1)
.await
.expect("no linkage must fail open");
let unknown_run = RunId::new();
link_task_to_run(&state, &task_id, 1, &unknown_run).await;
reject_if_run_terminal(&state, &task_id, 1)
.await
.expect("unknown run must fail open");
let owner_task = TaskId::new();
let live_run = RunId::new();
run_store
.create(run_record(&owner_task, &live_run, vec![]))
.await
.expect("run create");
link_task_to_run(&state, &task_id, 1, &live_run).await;
reject_if_run_terminal(&state, &task_id, 1)
.await
.expect("a Running run must pass the guard");
}
fn degradation_body(tool: &str, note: Option<&str>) -> DegradationBody {
DegradationBody {
tool: tool.to_string(),
error: "boom".to_string(),
fallback: "used cached value".to_string(),
note: note.map(str::to_string),
}
}
async fn link_task_to_run_with_agent(
state: &AppState,
task_id: &StepId,
attempt: u32,
run_id: &RunId,
agent: &str,
) {
let tid = task_id.clone();
let rid_str = run_id.to_string();
let agent = agent.to_string();
state
.engine
.with_state("test.link_task_to_run_with_agent", move |s| {
let mut entry = mlua_swarm::core::state::AgentCtxEntry::default();
entry.view.run_id = Some(rid_str);
entry.view.agent = agent;
s.agent_ctx.insert((tid, attempt), entry);
})
.await
.expect("link_task_to_run_with_agent");
}
#[tokio::test]
async fn worker_degradation_persists_entry_when_run_tracked() {
let data_store: Arc<dyn OutputStore> = Arc::new(InMemoryOutputStore::new());
let run_store: Arc<dyn RunStore> = Arc::new(InMemoryRunStore::new());
let state = test_state(data_store, run_store.clone());
let task_id = StepId::new();
let handle = seed_task_with_handle(&state, &task_id, "planner", 1, None).await;
let owner_task = TaskId::new();
let run_id = RunId::new();
run_store
.create(run_record(&owner_task, &run_id, vec![]))
.await
.expect("run create");
link_task_to_run_with_agent(&state, &task_id, 1, &run_id, "planner").await;
let status = worker_degradation(
State(state.clone()),
bearer_headers(&handle),
Json(degradation_body("web_search", Some("rate limited"))),
)
.await
.expect("worker_degradation");
assert_eq!(status, StatusCode::NO_CONTENT);
let rec = run_store.get(&run_id).await.expect("run get");
assert_eq!(
rec.degradations.len(),
1,
"degradations: {:?}",
rec.degradations
);
let entry = &rec.degradations[0];
assert_eq!(entry.tool, "web_search");
assert_eq!(entry.error, "boom");
assert_eq!(entry.fallback, "used cached value");
assert_eq!(entry.note.as_deref(), Some("rate limited"));
assert_eq!(entry.step_ref.as_deref(), Some("planner"));
assert_eq!(entry.attempt, Some(1));
assert!(entry.at > 0, "at must be a real timestamp: {}", entry.at);
}
#[tokio::test]
async fn worker_degradation_appends_in_order() {
let data_store: Arc<dyn OutputStore> = Arc::new(InMemoryOutputStore::new());
let run_store: Arc<dyn RunStore> = Arc::new(InMemoryRunStore::new());
let state = test_state(data_store, run_store.clone());
let task_id = StepId::new();
let handle = seed_task_with_handle(&state, &task_id, "planner", 1, None).await;
let owner_task = TaskId::new();
let run_id = RunId::new();
run_store
.create(run_record(&owner_task, &run_id, vec![]))
.await
.expect("run create");
link_task_to_run(&state, &task_id, 1, &run_id).await;
for tool in ["first_tool", "second_tool"] {
worker_degradation(
State(state.clone()),
bearer_headers(&handle),
Json(degradation_body(tool, None)),
)
.await
.expect("worker_degradation");
}
let rec = run_store.get(&run_id).await.expect("run get");
let tools: Vec<&str> = rec.degradations.iter().map(|e| e.tool.as_str()).collect();
assert_eq!(tools, vec!["first_tool", "second_tool"]);
}
#[tokio::test]
async fn worker_degradation_silent_ok_when_no_run_tracked() {
let data_store: Arc<dyn OutputStore> = Arc::new(InMemoryOutputStore::new());
let run_store: Arc<dyn RunStore> = Arc::new(InMemoryRunStore::new());
let state = test_state(data_store, run_store);
let task_id = StepId::new();
let handle = seed_task_with_handle(&state, &task_id, "planner", 1, None).await;
let status = worker_degradation(
State(state.clone()),
bearer_headers(&handle),
Json(degradation_body("some_tool", None)),
)
.await
.expect("worker_degradation must not error on missing run linkage");
assert_eq!(status, StatusCode::NO_CONTENT);
}
#[tokio::test]
async fn worker_degradation_rejects_terminal_run() {
let data_store: Arc<dyn OutputStore> = Arc::new(InMemoryOutputStore::new());
let run_store: Arc<dyn RunStore> = Arc::new(InMemoryRunStore::new());
let state = test_state(data_store, run_store.clone());
let task_id = StepId::new();
let handle = seed_task_with_handle(&state, &task_id, "planner", 1, None).await;
let owner_task = TaskId::new();
let run_id = RunId::new();
let mut rec = run_record(&owner_task, &run_id, vec![]);
rec.status = RunStatus::Done;
run_store.create(rec).await.expect("run create");
link_task_to_run(&state, &task_id, 1, &run_id).await;
let err = worker_degradation(
State(state.clone()),
bearer_headers(&handle),
Json(degradation_body("some_tool", None)),
)
.await
.expect_err("a degradation against a Done run must be rejected");
assert_eq!(err.status, StatusCode::GONE);
let rec = run_store.get(&run_id).await.expect("run get");
assert!(
rec.degradations.is_empty(),
"rejected degradation must not land: {:?}",
rec.degradations
);
}
}