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, EngineError, OutputEvent, RunId, StepId, WorkerPayload};
use mlua_swarm_schema::{ContextPolicy, VerdictChannel};
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,
};
map_completion_result(
state
.engine
.submit_output(&token, &task_id, attempt, event)
.await,
"submit_output",
)?;
state
.engine
.post_result(&token, &task_id, req.value)
.await
.map_err(|e| ApiError::engine(format!("post_result: {e}")))?;
Ok(StatusCode::NO_CONTENT)
}
const FILE_SENTINEL_PREFIX: &str = "@file:";
const FILE_SENTINEL_MAX_BYTES: u64 = 2 * 1024 * 1024;
const FILE_SENTINEL_ALLOW_KEY: &str = "allow_file_submit";
async fn resolve_file_sentinel(
state: &AppState,
task_id: &StepId,
attempt: u32,
body_str: String,
) -> Result<String, ApiError> {
let Some(rest) = body_str.strip_prefix(FILE_SENTINEL_PREFIX) else {
return Ok(body_str);
};
let path_str = rest.trim();
if path_str.is_empty() {
return Err(ApiError::bad_request(
"@file: sentinel: empty path".to_string(),
));
}
if path_str.contains('\n') || path_str.contains('\r') {
return Err(ApiError::bad_request(
"@file: sentinel: path must be a single line".to_string(),
));
}
let path = std::path::Path::new(path_str);
if !path.is_absolute() {
return Err(ApiError::bad_request(format!(
"@file: sentinel: path must be absolute (got {path_str:?})"
)));
}
let view = state
.engine
.agent_context_for(task_id, attempt)
.await
.ok_or_else(|| {
ApiError::bad_request(
"@file: sentinel: no AgentContextView for this task/attempt \
(spawn must run through AgentContextMiddleware to enable \
sentinel resolution)"
.to_string(),
)
})?;
if view.extra.get(FILE_SENTINEL_ALLOW_KEY) != Some(&Value::Bool(true)) {
return Err(ApiError::bad_request(format!(
"@file: sentinel: file submission is not allowed for this step \
(declare `{FILE_SENTINEL_ALLOW_KEY}: true` via `$step_meta` / \
`AgentMeta.ctx` / `Blueprint.metas`; strict boolean `true` \
required)"
)));
}
let work_dir = view.work_dir.ok_or_else(|| {
ApiError::bad_request("@file: sentinel: task has no resolved work_dir".to_string())
})?;
let work_dir_canon = tokio::fs::canonicalize(&work_dir).await.map_err(|e| {
ApiError::engine(format!(
"@file: sentinel: canonicalize work_dir {work_dir:?}: {e}"
))
})?;
let path_canon = match tokio::fs::canonicalize(path).await {
Ok(p) => p,
Err(e) if e.kind() == std::io::ErrorKind::NotFound => {
return Err(ApiError::not_found(format!(
"@file: sentinel: file not found: {path_str}"
)));
}
Err(e) => {
return Err(ApiError::engine(format!(
"@file: sentinel: canonicalize {path_str:?}: {e}"
)));
}
};
if !path_canon.starts_with(&work_dir_canon) {
return Err(ApiError::bad_request(format!(
"@file: sentinel: path {} is not under work_dir {} (canonicalized: {} vs {})",
path_str,
work_dir,
path_canon.display(),
work_dir_canon.display(),
)));
}
let meta = tokio::fs::metadata(&path_canon)
.await
.map_err(|e| ApiError::engine(format!("@file: sentinel: metadata {path_str:?}: {e}")))?;
if meta.len() > FILE_SENTINEL_MAX_BYTES {
return Err(ApiError::payload_too_large(format!(
"@file: sentinel: file size {} exceeds limit {}",
meta.len(),
FILE_SENTINEL_MAX_BYTES
)));
}
let bytes = tokio::fs::read(&path_canon)
.await
.map_err(|e| ApiError::engine(format!("@file: sentinel: read {path_str:?}: {e}")))?;
Ok(String::from_utf8_lossy(&bytes).trim_end().to_string())
}
async fn check_verdict_contract(
state: &AppState,
task_id: &StepId,
channel: VerdictChannel,
value: &str,
) -> Result<(), ApiError> {
let Some(contract) = state.engine.verdict_contract_for_task(task_id).await else {
return Ok(());
};
if contract.channel != channel {
return Ok(());
}
if contract.values.iter().any(|v| v == value) {
return Ok(());
}
Err(ApiError::unprocessable(format!(
"verdict contract violation: {value:?} is not a member of the declared values {:?}",
contract.values
)))
}
fn map_completion_result<T>(result: Result<T, EngineError>, context: &str) -> Result<T, ApiError> {
result.map_err(|e| match e {
EngineError::VerdictValueRejected { value, allowed } => ApiError::unprocessable(format!(
"verdict contract violation: {value:?} is not a member of the declared values {allowed:?}"
)),
EngineError::VerdictPartMissing { allowed } => ApiError::unprocessable(format!(
"verdict contract violation: no staged \"verdict\" part found for this attempt; declared values {allowed:?}"
)),
other => ApiError::engine(format!("{context}: {other}")),
})
}
#[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 body_str = resolve_file_sentinel(&state, &task_id, attempt, body_str).await?;
let value = Value::String(body_str);
let ok = q.ok.unwrap_or(true);
map_completion_result(
state
.engine
.submit_worker_result_trusted(&task_id, attempt, value, ok)
.await,
"submit_worker_result_trusted",
)?;
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 body_str = resolve_file_sentinel(&state, &task_id, attempt, body_str).await?;
if name == "verdict" {
check_verdict_contract(&state, &task_id, VerdictChannel::Part, &body_str).await?;
}
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()
}),
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,
}
}
#[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
);
}
async fn seed_work_dir(
state: &AppState,
task_id: &StepId,
attempt: u32,
work_dir: &str,
allow_file_submit: Option<Value>,
) {
let tid = task_id.clone();
let work_dir = work_dir.to_string();
state
.engine
.with_state("test.seed_work_dir", move |s| {
let mut entry = mlua_swarm::core::state::AgentCtxEntry::default();
entry.view.work_dir = Some(work_dir);
if let Some(v) = allow_file_submit {
entry
.view
.extra
.insert(FILE_SENTINEL_ALLOW_KEY.to_string(), v);
}
s.agent_ctx.insert((tid, attempt), entry);
})
.await
.expect("seed_work_dir");
}
#[tokio::test]
async fn worker_submit_resolves_file_sentinel_under_work_dir() {
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.clone(), run_store);
let task_id = StepId::new();
let handle = seed_task_with_handle(&state, &task_id, "planner", 1, None).await;
let tmp = tempfile::tempdir().expect("tempdir");
let work_dir = tmp.path().to_path_buf();
seed_work_dir(
&state,
&task_id,
1,
work_dir.to_str().expect("work_dir utf-8"),
Some(Value::Bool(true)),
)
.await;
let payload_path = work_dir.join("scout.md");
let payload = "## Context Package (broad)\n\nlarge body content\n";
tokio::fs::write(&payload_path, payload)
.await
.expect("write payload");
let body = format!(
"@file:{}",
payload_path.to_str().expect("payload path utf-8")
);
let status = worker_submit(
State(state.clone()),
bearer_headers(&handle),
Query(SubmitQuery { ok: None }),
axum::body::Bytes::from(body),
)
.await
.expect("worker_submit sentinel");
assert_eq!(status, StatusCode::NO_CONTENT);
let tid = task_id.clone();
let value = state
.engine
.with_state("test.inspect_output_store", move |s| {
s.output_store.get(&(tid.clone(), 1)).and_then(|evs| {
evs.iter().find_map(|ev| match ev {
OutputEvent::Final {
content: ContentRef::Inline { value },
..
} => Some(value.clone()),
_ => None,
})
})
})
.await
.expect("with_state")
.expect("Final event present");
assert_eq!(value, Value::String(payload.trim_end().to_string()));
}
#[tokio::test]
async fn worker_submit_passes_non_sentinel_body_unchanged() {
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.clone(), run_store);
let task_id = StepId::new();
let handle = seed_task_with_handle(&state, &task_id, "planner", 1, None).await;
let status = worker_submit(
State(state.clone()),
bearer_headers(&handle),
Query(SubmitQuery { ok: None }),
axum::body::Bytes::from_static(b"DONE yes=1 maybe=0 no=0"),
)
.await
.expect("worker_submit inline");
assert_eq!(status, StatusCode::NO_CONTENT);
let tid = task_id.clone();
let value = state
.engine
.with_state("test.inspect_output_store", move |s| {
s.output_store.get(&(tid.clone(), 1)).and_then(|evs| {
evs.iter().find_map(|ev| match ev {
OutputEvent::Final {
content: ContentRef::Inline { value },
..
} => Some(value.clone()),
_ => None,
})
})
})
.await
.expect("with_state")
.expect("Final event present");
assert_eq!(value, Value::String("DONE yes=1 maybe=0 no=0".to_string()));
}
#[tokio::test]
async fn worker_submit_rejects_sentinel_path_outside_work_dir() {
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 allowed = tempfile::tempdir().expect("allowed tempdir");
let outside = tempfile::tempdir().expect("outside tempdir");
seed_work_dir(
&state,
&task_id,
1,
allowed.path().to_str().expect("utf-8"),
Some(Value::Bool(true)),
)
.await;
let outside_file = outside.path().join("leak.md");
tokio::fs::write(&outside_file, b"outside content")
.await
.expect("write outside");
let body = format!(
"@file:{}",
outside_file.to_str().expect("outside path utf-8")
);
let err = worker_submit(
State(state.clone()),
bearer_headers(&handle),
Query(SubmitQuery { ok: None }),
axum::body::Bytes::from(body),
)
.await
.expect_err("outside-work_dir sentinel must be rejected");
assert_eq!(err.status, StatusCode::BAD_REQUEST);
}
#[tokio::test]
async fn worker_submit_rejects_sentinel_missing_file() {
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 tmp = tempfile::tempdir().expect("tempdir");
seed_work_dir(
&state,
&task_id,
1,
tmp.path().to_str().expect("utf-8"),
Some(Value::Bool(true)),
)
.await;
let missing = tmp.path().join("does-not-exist.md");
let body = format!("@file:{}", missing.to_str().expect("utf-8"));
let err = worker_submit(
State(state.clone()),
bearer_headers(&handle),
Query(SubmitQuery { ok: None }),
axum::body::Bytes::from(body),
)
.await
.expect_err("missing-file sentinel must be rejected");
assert_eq!(err.status, StatusCode::NOT_FOUND);
}
#[tokio::test]
async fn worker_submit_rejects_sentinel_relative_path() {
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 err = worker_submit(
State(state.clone()),
bearer_headers(&handle),
Query(SubmitQuery { ok: None }),
axum::body::Bytes::from_static(b"@file:relative/path.md"),
)
.await
.expect_err("relative-path sentinel must be rejected");
assert_eq!(err.status, StatusCode::BAD_REQUEST);
}
#[tokio::test]
async fn worker_submit_rejects_sentinel_without_agent_context_view() {
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 err = worker_submit(
State(state.clone()),
bearer_headers(&handle),
Query(SubmitQuery { ok: None }),
axum::body::Bytes::from_static(b"@file:/tmp/anywhere.md"),
)
.await
.expect_err("missing AgentContextView must reject sentinel");
assert_eq!(err.status, StatusCode::BAD_REQUEST);
}
#[tokio::test]
async fn worker_artifact_resolves_file_sentinel_under_work_dir() {
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 tmp = tempfile::tempdir().expect("tempdir");
seed_work_dir(
&state,
&task_id,
1,
tmp.path().to_str().expect("utf-8"),
Some(Value::Bool(true)),
)
.await;
let payload_path = tmp.path().join("part.md");
let payload = "artifact part body\n";
tokio::fs::write(&payload_path, payload)
.await
.expect("write payload");
let body = format!("@file:{}", payload_path.to_str().expect("utf-8"));
let status = worker_artifact(
State(state.clone()),
bearer_headers(&handle),
Query(ArtifactQuery {
name: "scout".to_string(),
}),
axum::body::Bytes::from(body),
)
.await
.expect("worker_artifact sentinel");
assert_eq!(status, StatusCode::NO_CONTENT);
}
#[tokio::test]
async fn worker_submit_rejects_sentinel_without_allow_flag() {
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 tmp = tempfile::tempdir().expect("tempdir");
seed_work_dir(
&state,
&task_id,
1,
tmp.path().to_str().expect("utf-8"),
None,
)
.await;
let payload_path = tmp.path().join("out.md");
tokio::fs::write(&payload_path, b"resolvable body")
.await
.expect("write payload");
let body = format!("@file:{}", payload_path.to_str().expect("utf-8"));
let err = worker_submit(
State(state.clone()),
bearer_headers(&handle),
Query(SubmitQuery { ok: None }),
axum::body::Bytes::from(body),
)
.await
.expect_err("missing opt-in must reject sentinel");
assert_eq!(err.status, StatusCode::BAD_REQUEST);
assert!(
err.message.contains("not allowed"),
"rejection must name the opt-in guard, got: {}",
err.message
);
}
#[tokio::test]
async fn worker_submit_rejects_sentinel_with_non_true_allow_values() {
for allow in [Value::Bool(false), Value::String("true".to_string())] {
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 tmp = tempfile::tempdir().expect("tempdir");
seed_work_dir(
&state,
&task_id,
1,
tmp.path().to_str().expect("utf-8"),
Some(allow.clone()),
)
.await;
let payload_path = tmp.path().join("out.md");
tokio::fs::write(&payload_path, b"resolvable body")
.await
.expect("write payload");
let body = format!("@file:{}", payload_path.to_str().expect("utf-8"));
let err = worker_submit(
State(state.clone()),
bearer_headers(&handle),
Query(SubmitQuery { ok: None }),
axum::body::Bytes::from(body),
)
.await
.expect_err("non-true opt-in value must reject sentinel");
assert_eq!(err.status, StatusCode::BAD_REQUEST, "value: {allow:?}");
}
}
fn body_verdict_contract(values: &[&str]) -> mlua_swarm_schema::VerdictContract {
mlua_swarm_schema::VerdictContract {
channel: VerdictChannel::Body,
values: values.iter().map(|v| v.to_string()).collect(),
}
}
fn part_verdict_contract(values: &[&str]) -> mlua_swarm_schema::VerdictContract {
mlua_swarm_schema::VerdictContract {
channel: VerdictChannel::Part,
values: values.iter().map(|v| v.to_string()).collect(),
}
}
#[tokio::test]
async fn worker_submit_rejects_body_outside_contract_values_with_422() {
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, "gate", 1, None).await;
state.engine.register_verdict_contracts(HashMap::from([(
"gate".to_string(),
body_verdict_contract(&["PASS", "BLOCKED"]),
)]));
let err = worker_submit(
State(state.clone()),
bearer_headers(&handle),
Query(SubmitQuery { ok: None }),
axum::body::Bytes::from("UNKNOWN"),
)
.await
.expect_err("value outside declared values must reject");
assert_eq!(err.status, StatusCode::UNPROCESSABLE_ENTITY);
assert!(
err.message.contains("PASS") && err.message.contains("BLOCKED"),
"rejection must echo the declared values, got: {}",
err.message
);
}
#[tokio::test]
async fn worker_submit_accepts_body_inside_contract_values() {
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, "gate", 1, None).await;
state.engine.register_verdict_contracts(HashMap::from([(
"gate".to_string(),
body_verdict_contract(&["PASS", "BLOCKED"]),
)]));
let status = worker_submit(
State(state.clone()),
bearer_headers(&handle),
Query(SubmitQuery { ok: None }),
axum::body::Bytes::from("PASS"),
)
.await
.expect("value inside declared values must succeed");
assert_eq!(status, StatusCode::NO_CONTENT);
}
#[tokio::test]
async fn worker_submit_without_a_declared_contract_is_unaffected() {
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, "undeclared-agent", 1, None).await;
let status = worker_submit(
State(state.clone()),
bearer_headers(&handle),
Query(SubmitQuery { ok: None }),
axum::body::Bytes::from("anything at all, no contract to violate"),
)
.await
.expect("no contract declared must never reject");
assert_eq!(status, StatusCode::NO_CONTENT);
}
#[tokio::test]
async fn worker_artifact_verdict_part_rejects_value_outside_contract_with_422() {
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, "gate", 1, None).await;
state.engine.register_verdict_contracts(HashMap::from([(
"gate".to_string(),
part_verdict_contract(&["PASS", "BLOCKED"]),
)]));
let err = worker_artifact(
State(state.clone()),
bearer_headers(&handle),
Query(ArtifactQuery {
name: "verdict".to_string(),
}),
axum::body::Bytes::from("UNKNOWN"),
)
.await
.expect_err("value outside declared values must reject");
assert_eq!(err.status, StatusCode::UNPROCESSABLE_ENTITY);
}
#[tokio::test]
async fn worker_artifact_non_verdict_part_skips_the_gate() {
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, "gate", 1, None).await;
state.engine.register_verdict_contracts(HashMap::from([(
"gate".to_string(),
part_verdict_contract(&["PASS", "BLOCKED"]),
)]));
let status = worker_artifact(
State(state.clone()),
bearer_headers(&handle),
Query(ArtifactQuery {
name: "notes".to_string(),
}),
axum::body::Bytes::from("anything at all"),
)
.await
.expect("non-verdict part name must never be gated");
assert_eq!(status, StatusCode::NO_CONTENT);
}
}