1use axum::{
68 extract::{Query, State},
69 http::{header, header::AUTHORIZATION, HeaderMap, StatusCode},
70 Json,
71};
72use mlua_swarm::core::agent_context::StepPointer;
73use mlua_swarm::core::step_naming::StepNaming;
74use mlua_swarm::store::run::{DegradationEntry, RunStatus, RunStoreError};
75use mlua_swarm::{CapToken, ContentRef, OutputEvent, RunId, StepId, WorkerPayload};
76use mlua_swarm_schema::ContextPolicy;
77use serde::Deserialize;
78use serde_json::Value;
79
80use crate::projection::McpQueryAdapter;
81use crate::{ApiError, AppState};
82
83#[derive(Debug, Deserialize)]
85pub struct PromptQuery {
86 pub task_id: StepId,
90}
91
92pub async fn worker_prompt(
98 State(state): State<AppState>,
99 headers: HeaderMap,
100 Query(q): Query<PromptQuery>,
101) -> Result<Json<WorkerPayload>, ApiError> {
102 let task_id = q.task_id;
103 let bearer = extract_bearer_raw(&headers)?;
104 let mut payload = if let Some(handle) = parse_worker_handle(&bearer) {
105 let resolved = state
107 .engine
108 .task_id_from_handle(handle)
109 .await
110 .map_err(|e| ApiError::engine(format!("task_id_from_handle: {e}")))?;
111 if resolved != task_id {
112 return Err(ApiError::bad_request(format!(
113 "handle {handle} is bound to task {resolved}, not {task_id}"
114 )));
115 }
116 state
117 .engine
118 .fetch_worker_payload_trusted(&task_id)
119 .await
120 .map_err(|e| ApiError::engine(format!("fetch_worker_payload_trusted: {e}")))?
121 } else {
122 let token = CapToken::decode(bearer.trim())
124 .map_err(|e| ApiError::bad_request(format!("invalid token: {e}")))?;
125 state
126 .engine
127 .fetch_worker_payload(&token, &task_id)
128 .await
129 .map_err(|e| ApiError::engine(format!("fetch_worker_payload: {e}")))?
130 };
131 assemble_step_pointers(&state, &mut payload).await;
132 Ok(Json(payload))
133}
134
135async fn assemble_step_pointers(state: &AppState, payload: &mut WorkerPayload) {
163 let Some(context) = payload.context.as_mut() else {
164 return;
165 };
166 let Some(run_id_str) = context.run_id.clone() else {
167 return;
168 };
169 let Ok(run_id) = RunId::parse(run_id_str) else {
170 return;
171 };
172
173 let adapter = McpQueryAdapter::new(
174 state.data_store.clone(),
175 state.run_store.clone(),
176 state.engine.clone(),
177 );
178 let Ok((run, resolved_steps)) = adapter.list_steps_by_run_id(&run_id).await else {
179 return;
180 };
181
182 let naming = state.engine.step_naming_for(&payload.task_id).await;
183 let policy = state
184 .engine
185 .context_policy_for(&payload.task_id, payload.attempt)
186 .await;
187 let self_canonical = naming
188 .as_deref()
189 .and_then(|n| n.canonical_of_producer(&payload.agent))
190 .map(str::to_string)
191 .unwrap_or_else(|| payload.agent.clone());
192
193 let mut pointers = Vec::new();
194 for step in &resolved_steps {
195 if step.name == self_canonical
196 || !allows_step_canonical(&policy, naming.as_deref(), &step.name)
197 {
198 continue;
199 }
200 if let Some((size_bytes, file_path, content_url, sha256)) =
201 crate::projection::resolve_step_pointer_fields(state, &run, step).await
202 {
203 pointers.push(StepPointer {
204 name: step.name.clone(),
205 size_bytes,
206 file_path,
207 content_url,
208 sha256,
209 });
210 }
211 }
212 context.steps = pointers;
213}
214
215fn allows_step_canonical(
230 policy: &ContextPolicy,
231 naming: Option<&StepNaming>,
232 canonical_name: &str,
233) -> bool {
234 let resolves_to = |raw: &str| -> bool {
235 match naming {
236 Some(n) => n
237 .resolve(raw)
238 .map(|c| c == canonical_name)
239 .unwrap_or(raw == canonical_name),
240 None => raw == canonical_name,
241 }
242 };
243 if policy
244 .steps_exclude
245 .iter()
246 .any(|excluded| resolves_to(excluded))
247 {
248 return false;
249 }
250 match &policy.steps {
251 None => true,
252 Some(list) => list.iter().any(|included| resolves_to(included)),
253 }
254}
255
256#[derive(Debug, Deserialize)]
258pub struct WorkerResultReq {
259 pub task_id: StepId,
262 pub value: Value,
264 #[serde(default = "default_ok_true")]
268 pub ok: bool,
269 #[serde(default)]
272 pub attempt: Option<u32>,
273}
274
275fn default_ok_true() -> bool {
276 true
277}
278
279pub async fn worker_result(
282 State(state): State<AppState>,
283 headers: HeaderMap,
284 Json(req): Json<WorkerResultReq>,
285) -> Result<StatusCode, ApiError> {
286 let token = decode_worker_bearer(&headers)?;
287 let task_id = req.task_id.clone();
288
289 let attempt = match req.attempt {
291 Some(n) => n,
292 None => state
293 .engine
294 .task_attempt(&task_id)
295 .await
296 .map_err(|e| ApiError::engine(format!("task_attempt: {e}")))?,
297 };
298
299 let event = OutputEvent::Final {
300 content: ContentRef::Inline {
301 value: req.value.clone(),
302 },
303 ok: req.ok,
304 };
305 state
306 .engine
307 .submit_output(&token, &task_id, attempt, event)
308 .await
309 .map_err(|e| ApiError::engine(format!("submit_output: {e}")))?;
310 state
311 .engine
312 .post_result(&token, &task_id, req.value)
313 .await
314 .map_err(|e| ApiError::engine(format!("post_result: {e}")))?;
315 Ok(StatusCode::NO_CONTENT)
316}
317
318#[derive(Debug, Deserialize, Default)]
332pub struct SubmitQuery {
333 #[serde(default)]
337 pub ok: Option<bool>,
338}
339
340pub async fn worker_submit(
346 State(state): State<AppState>,
347 headers: HeaderMap,
348 Query(q): Query<SubmitQuery>,
349 body: axum::body::Bytes,
350) -> Result<StatusCode, ApiError> {
351 let bearer = extract_bearer_raw(&headers)?;
354 let task_id = if let Some(handle) = parse_worker_handle(&bearer) {
355 state
356 .engine
357 .task_id_from_handle(handle)
358 .await
359 .map_err(|e| ApiError::engine(format!("task_id_from_handle: {e}")))?
360 } else {
361 let token = CapToken::decode(bearer.trim())
362 .map_err(|e| ApiError::bad_request(format!("invalid token: {e}")))?;
363 state
364 .engine
365 .task_id_from_token(&token)
366 .await
367 .map_err(|e| ApiError::engine(format!("task_id_from_token: {e}")))?
368 };
369 let attempt = state
370 .engine
371 .task_attempt(&task_id)
372 .await
373 .map_err(|e| ApiError::engine(format!("task_attempt: {e}")))?;
374 reject_if_run_terminal(&state, &task_id, attempt).await?;
377 let body_str = String::from_utf8_lossy(&body).trim_end().to_string();
382 let value = Value::String(body_str);
383
384 let ok = q.ok.unwrap_or(true);
390 state
391 .engine
392 .submit_worker_result_trusted(&task_id, attempt, value, ok)
393 .await
394 .map_err(|e| ApiError::engine(format!("submit_worker_result_trusted: {e}")))?;
395 Ok(StatusCode::NO_CONTENT)
396}
397
398#[derive(Debug, Deserialize)]
400pub struct ArtifactQuery {
401 pub name: String,
408}
409
410pub async fn worker_artifact(
432 State(state): State<AppState>,
433 headers: HeaderMap,
434 Query(q): Query<ArtifactQuery>,
435 body: axum::body::Bytes,
436) -> Result<StatusCode, ApiError> {
437 let name = q.name.trim();
438 if name.is_empty() {
439 return Err(ApiError::bad_request("name must not be empty".into()));
440 }
441 let name = name.to_string();
442
443 let bearer = extract_bearer_raw(&headers)?;
444 let task_id = if let Some(handle) = parse_worker_handle(&bearer) {
445 state
446 .engine
447 .task_id_from_handle(handle)
448 .await
449 .map_err(|e| ApiError::engine(format!("task_id_from_handle: {e}")))?
450 } else {
451 let token = CapToken::decode(bearer.trim())
452 .map_err(|e| ApiError::bad_request(format!("invalid token: {e}")))?;
453 state
454 .engine
455 .task_id_from_token(&token)
456 .await
457 .map_err(|e| ApiError::engine(format!("task_id_from_token: {e}")))?
458 };
459 let attempt = state
460 .engine
461 .task_attempt(&task_id)
462 .await
463 .map_err(|e| ApiError::engine(format!("task_attempt: {e}")))?;
464 reject_if_run_terminal(&state, &task_id, attempt).await?;
467 let body_str = String::from_utf8_lossy(&body).trim_end().to_string();
468 let value = Value::String(body_str);
469
470 state
471 .engine
472 .stage_worker_artifact_trusted(&task_id, attempt, name, value)
473 .await
474 .map_err(|e| ApiError::engine(format!("stage_worker_artifact_trusted: {e}")))?;
475 Ok(StatusCode::NO_CONTENT)
476}
477
478#[derive(Debug, Deserialize)]
480pub struct DegradationBody {
481 pub tool: String,
483 pub error: String,
485 pub fallback: String,
487 #[serde(default)]
489 pub note: Option<String>,
490}
491
492pub async fn worker_degradation(
520 State(state): State<AppState>,
521 headers: HeaderMap,
522 Json(body): Json<DegradationBody>,
523) -> Result<StatusCode, ApiError> {
524 let bearer = extract_bearer_raw(&headers)?;
525 let task_id = if let Some(handle) = parse_worker_handle(&bearer) {
526 state
527 .engine
528 .task_id_from_handle(handle)
529 .await
530 .map_err(|e| ApiError::engine(format!("task_id_from_handle: {e}")))?
531 } else {
532 let token = CapToken::decode(bearer.trim())
533 .map_err(|e| ApiError::bad_request(format!("invalid token: {e}")))?;
534 state
535 .engine
536 .task_id_from_token(&token)
537 .await
538 .map_err(|e| ApiError::engine(format!("task_id_from_token: {e}")))?
539 };
540 let attempt = state
541 .engine
542 .task_attempt(&task_id)
543 .await
544 .map_err(|e| ApiError::engine(format!("task_attempt: {e}")))?;
545 reject_if_run_terminal(&state, &task_id, attempt).await?;
548
549 let tid = task_id.clone();
554 let (run_id_str, agent) = match state
555 .engine
556 .with_state("worker_degradation_run_lookup", move |s| {
557 s.agent_ctx.get(&(tid, attempt)).and_then(|e| {
558 e.view
559 .run_id
560 .clone()
561 .map(|run_id| (run_id, e.view.agent.clone()))
562 })
563 })
564 .await
565 {
566 Ok(Some(pair)) => pair,
567 _ => {
568 tracing::warn!(%task_id, "worker_degradation: no run linkage for this task; entry dropped");
569 return Ok(StatusCode::NO_CONTENT);
570 }
571 };
572 let Ok(run_id) = RunId::parse(run_id_str) else {
573 tracing::warn!(%task_id, "worker_degradation: run_id failed to parse; entry dropped");
574 return Ok(StatusCode::NO_CONTENT);
575 };
576
577 let entry = DegradationEntry {
578 tool: body.tool,
579 error: body.error,
580 fallback: body.fallback,
581 note: body.note,
582 step_ref: Some(agent),
583 attempt: Some(attempt),
584 at: crate::tasks::now_secs(),
585 };
586 match state.run_store.append_degradation(&run_id, entry).await {
587 Ok(()) => Ok(StatusCode::NO_CONTENT),
588 Err(RunStoreError::NotFound(_)) => {
589 tracing::warn!(%task_id, %run_id, "worker_degradation: run not found in run_store; entry dropped");
590 Ok(StatusCode::NO_CONTENT)
591 }
592 Err(e) => Err(ApiError::engine(format!("append_degradation: {e}"))),
593 }
594}
595
596async fn reject_if_run_terminal(
613 state: &AppState,
614 task_id: &StepId,
615 attempt: u32,
616) -> Result<(), ApiError> {
617 let tid = task_id.clone();
618 let run_id_str = match state
619 .engine
620 .with_state("worker_terminal_run_guard", move |s| {
621 s.agent_ctx
622 .get(&(tid, attempt))
623 .and_then(|e| e.view.run_id.clone())
624 })
625 .await
626 {
627 Ok(Some(rid)) => rid,
628 _ => return Ok(()),
629 };
630 let Ok(run_id) = RunId::parse(run_id_str) else {
631 return Ok(());
632 };
633 let Ok(rec) = state.run_store.get(&run_id).await else {
634 return Ok(());
635 };
636 match rec.status {
637 RunStatus::Done | RunStatus::Failed | RunStatus::Interrupted => {
638 Err(ApiError::gone(format!(
639 "run {run_id} is already terminal ({:?}): this attempt's output cannot be \
640 delivered to a flow context; re-kick the task (POST /v1/tasks/:id/runs) and \
641 fetch a fresh prompt",
642 rec.status
643 )))
644 }
645 RunStatus::Pending | RunStatus::Running => Ok(()),
646 }
647}
648
649#[derive(Debug, Deserialize)]
654pub struct PromptSystemQuery {
655 pub task_id: StepId,
658 pub attempt: u32,
660}
661
662pub async fn worker_prompt_system(
672 State(state): State<AppState>,
673 headers: HeaderMap,
674 Query(q): Query<PromptSystemQuery>,
675) -> Result<impl axum::response::IntoResponse, ApiError> {
676 let task_id = q.task_id;
677 let attempt = q.attempt;
678 let bearer = extract_bearer_raw(&headers)?;
679 if let Some(handle) = parse_worker_handle(&bearer) {
680 let resolved = state
681 .engine
682 .task_id_from_handle(handle)
683 .await
684 .map_err(|e| ApiError::engine(format!("task_id_from_handle: {e}")))?;
685 if resolved != task_id {
686 return Err(ApiError::bad_request(format!(
687 "handle {handle} is bound to task {resolved}, not {task_id}"
688 )));
689 }
690 } else {
691 let token = CapToken::decode(bearer.trim())
692 .map_err(|e| ApiError::bad_request(format!("invalid token: {e}")))?;
693 state
694 .engine
695 .verify_token_for_task(&token, mlua_swarm::Verb::FetchPrompt, &task_id)
696 .await
697 .map_err(|e| ApiError::engine(format!("verify_token_for_task: {e}")))?;
698 }
699 let system = state
700 .engine
701 .raw_system_prompt(&task_id, attempt)
702 .await
703 .map_err(|e| ApiError::engine(format!("raw_system_prompt: {e}")))?
704 .ok_or_else(|| {
705 ApiError::not_found(format!(
706 "no baked system prompt for task {task_id} attempt {attempt}"
707 ))
708 })?;
709 Ok((
710 [(header::CONTENT_TYPE, "text/plain; charset=utf-8")],
711 system,
712 ))
713}
714
715#[derive(Debug, serde::Serialize)]
717pub struct AgentRenderSizeResponse {
718 pub agent: String,
720 pub last_rendered_bytes: Option<usize>,
724}
725
726pub async fn agent_render_size(
736 State(state): State<AppState>,
737 axum::extract::Path(name): axum::extract::Path<String>,
738) -> Json<AgentRenderSizeResponse> {
739 let last_rendered_bytes = state.engine.agent_last_rendered_size(&name).await;
740 Json(AgentRenderSizeResponse {
741 agent: name,
742 last_rendered_bytes,
743 })
744}
745
746fn extract_bearer_raw(headers: &HeaderMap) -> Result<String, ApiError> {
750 let v = headers
751 .get(AUTHORIZATION)
752 .ok_or_else(|| ApiError::bad_request("missing Authorization header".into()))?
753 .to_str()
754 .map_err(|_| ApiError::bad_request("invalid Authorization header encoding".into()))?;
755 let s = v
756 .strip_prefix("Bearer ")
757 .ok_or_else(|| ApiError::bad_request("Authorization must be 'Bearer <token>'".into()))?
758 .trim();
759 if s.is_empty() {
760 return Err(ApiError::bad_request("Bearer is empty".into()));
761 }
762 Ok(s.to_string())
763}
764
765fn parse_worker_handle(s: &str) -> Option<&str> {
769 let s = s.trim();
770 if s.starts_with("wh-")
771 && s.len() >= 5
772 && s.len() <= 64
773 && s[3..].chars().all(|c| c.is_ascii_alphanumeric())
774 {
775 Some(s)
776 } else {
777 None
778 }
779}
780
781fn decode_worker_bearer(headers: &HeaderMap) -> Result<CapToken, ApiError> {
785 let v = headers
786 .get(AUTHORIZATION)
787 .ok_or_else(|| ApiError::bad_request("missing Authorization header".into()))?
788 .to_str()
789 .map_err(|_| ApiError::bad_request("invalid Authorization header encoding".into()))?;
790 let encoded = v
791 .strip_prefix("Bearer ")
792 .ok_or_else(|| ApiError::bad_request("Authorization must be 'Bearer <token>'".into()))?
793 .trim();
794 if encoded.is_empty() {
795 return Err(ApiError::bad_request("Bearer token is empty".into()));
796 }
797 CapToken::decode(encoded).map_err(|e| ApiError::bad_request(format!("invalid token: {e}")))
798}
799
800#[cfg(test)]
805mod tests {
806 use super::*;
807 use axum::response::IntoResponse;
808 use mlua_swarm::core::agent_context::AgentContextView;
809 use mlua_swarm::core::config::EngineCfg;
810 use mlua_swarm::core::engine::Engine;
811 use mlua_swarm::store::output::{InMemoryOutputStore, OutputStore};
812 use mlua_swarm::store::run::{InMemoryRunStore, RunRecord, RunStatus, RunStore, StepEntry};
813 use mlua_swarm::store::task::InMemoryTaskStore;
814 use mlua_swarm::{RunId, StepId, TaskId};
815 use serde_json::json;
816 use std::collections::HashMap;
817 use std::sync::Arc;
818 use tokio::sync::Mutex;
819
820 fn test_state(data_store: Arc<dyn OutputStore>, run_store: Arc<dyn RunStore>) -> AppState {
826 let engine = Engine::new(EngineCfg::default());
827 let compiler = mlua_swarm::Compiler::new(crate::default_registry());
828 let launch = Arc::new(mlua_swarm::TaskLaunchService::new(engine.clone(), compiler));
829 AppState {
830 engine,
831 sessions: Arc::new(Mutex::new(crate::SessionStore::default())),
832 task_app: Arc::new(mlua_swarm::TaskApplication::new_inline_only(launch)),
833 ws_operator_factory: None,
834 data_store,
835 operator_sessions: Arc::new(Mutex::new(HashMap::new())),
836 roles_to_sid: Arc::new(Mutex::new(HashMap::new())),
837 task_store: Arc::new(InMemoryTaskStore::new()),
838 run_store,
839 base_url: None,
840 sync_timeout_secs: 300,
841 }
842 }
843
844 async fn append_final(
845 data_store: &Arc<dyn OutputStore>,
846 task_id: &str,
847 producer: &str,
848 value: Value,
849 ) {
850 data_store
851 .append(
852 task_id,
853 1,
854 producer,
855 OutputEvent::Final {
856 content: ContentRef::Inline { value },
857 ok: true,
858 },
859 vec![],
860 )
861 .await
862 .expect("append final");
863 }
864
865 fn step_entry(step_id: &StepId, step_ref: &str) -> StepEntry {
866 StepEntry {
867 step_id: step_id.clone(),
868 step_ref: Some(step_ref.to_string()),
869 status: Some("passed".to_string()),
870 at: 0,
871 }
872 }
873
874 fn run_record(task_id: &TaskId, run_id: &RunId, step_entries: Vec<StepEntry>) -> RunRecord {
875 RunRecord {
876 id: run_id.clone(),
877 task_id: task_id.clone(),
878 status: RunStatus::Running,
879 step_entries,
880 degradations: Vec::new(),
881 operator_sid: None,
882 result_ref: None,
883 created_at: 0,
884 updated_at: 0,
885 }
886 }
887
888 fn consumer_payload(consumer_step_id: &StepId, run_id: &RunId) -> WorkerPayload {
889 WorkerPayload {
890 task_id: consumer_step_id.clone(),
891 attempt: 1,
892 agent: "consumer".to_string(),
893 system: None,
894 prompt: String::new(),
895 context: Some(AgentContextView {
896 task_id: consumer_step_id.to_string(),
897 agent: "consumer".to_string(),
898 attempt: 1,
899 run_id: Some(run_id.to_string()),
900 ..Default::default()
901 }),
902 system_ref: None,
903 }
904 }
905
906 #[tokio::test]
911 async fn context_policy_unspecified_yields_every_submitted_step() {
912 let data_store: Arc<dyn OutputStore> = Arc::new(InMemoryOutputStore::new());
913 let run_store: Arc<dyn RunStore> = Arc::new(InMemoryRunStore::new());
914 let task_id = TaskId::new();
915 let run_id = RunId::new();
916 let planner_id = StepId::new();
917 let coder_id = StepId::new();
918
919 append_final(
920 &data_store,
921 planner_id.as_str(),
922 "planner",
923 json!({"plan": "x"}),
924 )
925 .await;
926 append_final(
927 &data_store,
928 coder_id.as_str(),
929 "coder",
930 json!({"code": "y"}),
931 )
932 .await;
933 run_store
934 .create(run_record(
935 &task_id,
936 &run_id,
937 vec![
938 step_entry(&planner_id, "planner"),
939 step_entry(&coder_id, "coder"),
940 ],
941 ))
942 .await
943 .expect("create run");
944
945 let state = test_state(data_store, run_store);
946 let consumer_id = StepId::new();
947 let mut payload = consumer_payload(&consumer_id, &run_id);
948 assemble_step_pointers(&state, &mut payload).await;
949
950 let names: Vec<&str> = payload
951 .context
952 .as_ref()
953 .expect("context")
954 .steps
955 .iter()
956 .map(|p| p.name.as_str())
957 .collect();
958 assert!(names.contains(&"planner"), "names: {names:?}");
959 assert!(names.contains(&"coder"), "names: {names:?}");
960 }
961
962 #[tokio::test]
964 async fn context_policy_steps_include_list_filters_to_named_steps() {
965 let data_store: Arc<dyn OutputStore> = Arc::new(InMemoryOutputStore::new());
966 let run_store: Arc<dyn RunStore> = Arc::new(InMemoryRunStore::new());
967 let task_id = TaskId::new();
968 let run_id = RunId::new();
969 let planner_id = StepId::new();
970 let coder_id = StepId::new();
971 append_final(&data_store, planner_id.as_str(), "planner", json!("x")).await;
972 append_final(&data_store, coder_id.as_str(), "coder", json!("y")).await;
973 run_store
974 .create(run_record(
975 &task_id,
976 &run_id,
977 vec![
978 step_entry(&planner_id, "planner"),
979 step_entry(&coder_id, "coder"),
980 ],
981 ))
982 .await
983 .expect("create run");
984
985 let state = test_state(data_store, run_store);
986 let consumer_id = StepId::new();
987 state
988 .engine
989 .with_state("test.seed_policy", {
990 let consumer_id = consumer_id.clone();
991 move |s| {
992 s.agent_ctx.insert(
993 (consumer_id, 1),
994 mlua_swarm::core::state::AgentCtxEntry {
995 policy: mlua_swarm_schema::ContextPolicy {
996 steps: Some(vec!["planner".to_string()]),
997 ..Default::default()
998 },
999 ..Default::default()
1000 },
1001 );
1002 }
1003 })
1004 .await
1005 .expect("seed policy");
1006
1007 let mut payload = consumer_payload(&consumer_id, &run_id);
1008 assemble_step_pointers(&state, &mut payload).await;
1009
1010 let names: Vec<&str> = payload
1011 .context
1012 .as_ref()
1013 .expect("context")
1014 .steps
1015 .iter()
1016 .map(|p| p.name.as_str())
1017 .collect();
1018 assert_eq!(names, vec!["planner"], "names: {names:?}");
1019 }
1020
1021 #[tokio::test]
1023 async fn context_policy_steps_empty_list_yields_no_pointers() {
1024 let data_store: Arc<dyn OutputStore> = Arc::new(InMemoryOutputStore::new());
1025 let run_store: Arc<dyn RunStore> = Arc::new(InMemoryRunStore::new());
1026 let task_id = TaskId::new();
1027 let run_id = RunId::new();
1028 let planner_id = StepId::new();
1029 append_final(&data_store, planner_id.as_str(), "planner", json!("x")).await;
1030 run_store
1031 .create(run_record(
1032 &task_id,
1033 &run_id,
1034 vec![step_entry(&planner_id, "planner")],
1035 ))
1036 .await
1037 .expect("create run");
1038
1039 let state = test_state(data_store, run_store);
1040 let consumer_id = StepId::new();
1041 state
1042 .engine
1043 .with_state("test.seed_policy", {
1044 let consumer_id = consumer_id.clone();
1045 move |s| {
1046 s.agent_ctx.insert(
1047 (consumer_id, 1),
1048 mlua_swarm::core::state::AgentCtxEntry {
1049 policy: mlua_swarm_schema::ContextPolicy {
1050 steps: Some(vec![]),
1051 ..Default::default()
1052 },
1053 ..Default::default()
1054 },
1055 );
1056 }
1057 })
1058 .await
1059 .expect("seed policy");
1060
1061 let mut payload = consumer_payload(&consumer_id, &run_id);
1062 assemble_step_pointers(&state, &mut payload).await;
1063
1064 assert!(payload.context.expect("context").steps.is_empty());
1065 }
1066
1067 #[tokio::test]
1069 async fn context_policy_steps_exclude_wins_over_steps() {
1070 let data_store: Arc<dyn OutputStore> = Arc::new(InMemoryOutputStore::new());
1071 let run_store: Arc<dyn RunStore> = Arc::new(InMemoryRunStore::new());
1072 let task_id = TaskId::new();
1073 let run_id = RunId::new();
1074 let planner_id = StepId::new();
1075 let coder_id = StepId::new();
1076 append_final(&data_store, planner_id.as_str(), "planner", json!("x")).await;
1077 append_final(&data_store, coder_id.as_str(), "coder", json!("y")).await;
1078 run_store
1079 .create(run_record(
1080 &task_id,
1081 &run_id,
1082 vec![
1083 step_entry(&planner_id, "planner"),
1084 step_entry(&coder_id, "coder"),
1085 ],
1086 ))
1087 .await
1088 .expect("create run");
1089
1090 let state = test_state(data_store, run_store);
1091 let consumer_id = StepId::new();
1092 state
1093 .engine
1094 .with_state("test.seed_policy", {
1095 let consumer_id = consumer_id.clone();
1096 move |s| {
1097 s.agent_ctx.insert(
1098 (consumer_id, 1),
1099 mlua_swarm::core::state::AgentCtxEntry {
1100 policy: mlua_swarm_schema::ContextPolicy {
1101 steps: Some(vec!["planner".to_string(), "coder".to_string()]),
1102 steps_exclude: vec!["planner".to_string()],
1103 ..Default::default()
1104 },
1105 ..Default::default()
1106 },
1107 );
1108 }
1109 })
1110 .await
1111 .expect("seed policy");
1112
1113 let mut payload = consumer_payload(&consumer_id, &run_id);
1114 assemble_step_pointers(&state, &mut payload).await;
1115
1116 let names: Vec<&str> = payload
1117 .context
1118 .as_ref()
1119 .expect("context")
1120 .steps
1121 .iter()
1122 .map(|p| p.name.as_str())
1123 .collect();
1124 assert_eq!(names, vec!["coder"], "names: {names:?}");
1125 }
1126
1127 #[tokio::test]
1136 async fn in_flight_step_output_is_visible_before_run_finalizes() {
1137 let data_store: Arc<dyn OutputStore> = Arc::new(InMemoryOutputStore::new());
1138 let run_store: Arc<dyn RunStore> = Arc::new(InMemoryRunStore::new());
1139 let task_id = TaskId::new();
1140 let run_id = RunId::new();
1141 let step1_id = StepId::new();
1142 append_final(
1143 &data_store,
1144 step1_id.as_str(),
1145 "step1",
1146 json!({"step1_out": "hi"}),
1147 )
1148 .await;
1149 let mut run = run_record(&task_id, &run_id, vec![step_entry(&step1_id, "step1")]);
1150 run.status = RunStatus::Running;
1151 run.result_ref = None; run_store.create(run).await.expect("create run");
1153
1154 let state = test_state(data_store, run_store);
1155 let consumer_id = StepId::new();
1156 let mut payload = consumer_payload(&consumer_id, &run_id);
1157 assemble_step_pointers(&state, &mut payload).await;
1158
1159 let steps = &payload.context.expect("context").steps;
1160 assert_eq!(steps.len(), 1);
1161 assert_eq!(steps[0].name, "step1");
1162 }
1163
1164 #[tokio::test]
1168 async fn self_agent_name_is_always_excluded() {
1169 let data_store: Arc<dyn OutputStore> = Arc::new(InMemoryOutputStore::new());
1170 let run_store: Arc<dyn RunStore> = Arc::new(InMemoryRunStore::new());
1171 let task_id = TaskId::new();
1172 let run_id = RunId::new();
1173 let planner_id = StepId::new();
1174 let consumer_prior_id = StepId::new();
1175 append_final(&data_store, planner_id.as_str(), "planner", json!("x")).await;
1176 append_final(
1177 &data_store,
1178 consumer_prior_id.as_str(),
1179 "consumer",
1180 json!("self"),
1181 )
1182 .await;
1183 run_store
1184 .create(run_record(
1185 &task_id,
1186 &run_id,
1187 vec![
1188 step_entry(&planner_id, "planner"),
1189 step_entry(&consumer_prior_id, "consumer"),
1190 ],
1191 ))
1192 .await
1193 .expect("create run");
1194
1195 let state = test_state(data_store, run_store);
1196 let consumer_id = StepId::new();
1197 let mut payload = consumer_payload(&consumer_id, &run_id);
1198 assemble_step_pointers(&state, &mut payload).await;
1199
1200 let names: Vec<&str> = payload
1201 .context
1202 .as_ref()
1203 .expect("context")
1204 .steps
1205 .iter()
1206 .map(|p| p.name.as_str())
1207 .collect();
1208 assert!(!names.contains(&"consumer"), "names: {names:?}");
1209 assert!(names.contains(&"planner"), "names: {names:?}");
1210 }
1211
1212 #[tokio::test]
1216 async fn step_pointer_serializes_with_no_preview_or_content_bytes() {
1217 let data_store: Arc<dyn OutputStore> = Arc::new(InMemoryOutputStore::new());
1218 let run_store: Arc<dyn RunStore> = Arc::new(InMemoryRunStore::new());
1219 let task_id = TaskId::new();
1220 let run_id = RunId::new();
1221 let planner_id = StepId::new();
1222 append_final(
1223 &data_store,
1224 planner_id.as_str(),
1225 "planner",
1226 json!({"plan": "do the thing, at length".repeat(50)}),
1227 )
1228 .await;
1229 run_store
1230 .create(run_record(
1231 &task_id,
1232 &run_id,
1233 vec![step_entry(&planner_id, "planner")],
1234 ))
1235 .await
1236 .expect("create run");
1237
1238 let state = test_state(data_store, run_store);
1239 let consumer_id = StepId::new();
1240 let mut payload = consumer_payload(&consumer_id, &run_id);
1241 assemble_step_pointers(&state, &mut payload).await;
1242
1243 let steps = &payload.context.expect("context").steps;
1244 assert_eq!(steps.len(), 1);
1245 let json_value = serde_json::to_value(&steps[0]).expect("serialize StepPointer");
1246 let obj = json_value.as_object().expect("object");
1247 for forbidden in ["preview", "content", "value", "bytes"] {
1248 assert!(
1249 !obj.contains_key(forbidden),
1250 "StepPointer must not carry a {forbidden:?} field: {obj:?}"
1251 );
1252 }
1253 assert!(obj.contains_key("name"));
1254 assert!(obj.contains_key("size_bytes"));
1255 assert!(obj.contains_key("content_url"));
1256 assert!(obj.contains_key("sha256"));
1257 }
1258
1259 fn declared_name_bp() -> mlua_swarm::blueprint::Blueprint {
1267 use mlua_flow_ir::{Expr, Node};
1268 use mlua_swarm::blueprint::{
1269 current_schema_version, AgentDef, AgentKind, AgentMeta, Blueprint, BlueprintMetadata,
1270 CompilerHints, CompilerStrategy,
1271 };
1272 Blueprint {
1273 schema_version: current_schema_version(),
1274 id: "worker-test-declared-name-bp".into(),
1275 flow: Node::Step {
1276 ref_: "planner".to_string(),
1277 in_: Expr::Path {
1278 at: "$.in".to_string(),
1279 },
1280 out: Expr::Path {
1281 at: "$.plan".to_string(),
1282 },
1283 },
1284 agents: vec![AgentDef {
1285 name: "planner".to_string(),
1286 kind: AgentKind::RustFn,
1287 spec: json!({"fn_id": "planner"}),
1288 profile: None,
1289 meta: Some(AgentMeta {
1290 projection_name: Some("plan-out".to_string()),
1291 ..Default::default()
1292 }),
1293 }],
1294 operators: vec![],
1295 metas: vec![],
1296 hints: CompilerHints::default(),
1297 strategy: CompilerStrategy::default(),
1298 metadata: BlueprintMetadata::default(),
1299 spawner_hints: Default::default(),
1300 default_agent_kind: AgentKind::Operator,
1301 default_operator_kind: None,
1302 default_init_ctx: None,
1303 default_agent_ctx: None,
1304 default_context_policy: None,
1305 projection_placement: None,
1306 audits: vec![],
1307 degradation_policy: None,
1308 }
1309 }
1310
1311 #[tokio::test]
1317 async fn declared_projection_name_pointer_name_is_canonical_and_policy_matches_it() {
1318 let data_store: Arc<dyn OutputStore> = Arc::new(InMemoryOutputStore::new());
1319 let run_store: Arc<dyn RunStore> = Arc::new(InMemoryRunStore::new());
1320 let task_id = TaskId::new();
1321 let run_id = RunId::new();
1322 let planner_id = StepId::new();
1323
1324 append_final(
1327 &data_store,
1328 planner_id.as_str(),
1329 "plan-out",
1330 json!({"plan": "x"}),
1331 )
1332 .await;
1333 run_store
1334 .create(run_record(
1335 &task_id,
1336 &run_id,
1337 vec![step_entry(&planner_id, "planner")],
1338 ))
1339 .await
1340 .expect("create run");
1341
1342 let state = test_state(data_store, run_store);
1343
1344 let (naming, _warnings) =
1350 mlua_swarm::core::step_naming::StepNaming::from_blueprint(&declared_name_bp())
1351 .expect("no collision");
1352 let naming = Arc::new(naming);
1353 let consumer_id = StepId::new();
1354 state
1355 .engine
1356 .with_state("test.seed_step_naming", {
1357 let naming = naming.clone();
1358 let planner_id = planner_id.clone();
1359 let consumer_id = consumer_id.clone();
1360 move |s| {
1361 s.step_namings.insert(planner_id, naming.clone());
1362 s.step_namings.insert(consumer_id, naming);
1363 }
1364 })
1365 .await
1366 .expect("seed step naming");
1367 state
1368 .engine
1369 .with_state("test.seed_policy", {
1370 let consumer_id = consumer_id.clone();
1371 move |s| {
1372 s.agent_ctx.insert(
1373 (consumer_id, 1),
1374 mlua_swarm::core::state::AgentCtxEntry {
1375 policy: mlua_swarm_schema::ContextPolicy {
1376 steps: Some(vec!["plan-out".to_string()]),
1377 ..Default::default()
1378 },
1379 ..Default::default()
1380 },
1381 );
1382 }
1383 })
1384 .await
1385 .expect("seed policy");
1386
1387 let mut payload = consumer_payload(&consumer_id, &run_id);
1388 assemble_step_pointers(&state, &mut payload).await;
1389
1390 let steps = &payload.context.expect("context").steps;
1391 assert_eq!(steps.len(), 1, "steps: {steps:?}");
1392 assert_eq!(
1393 steps[0].name, "plan-out",
1394 "StepPointer.name must be the canonical name"
1395 );
1396 }
1397
1398 async fn seed_task_with_handle(
1408 state: &AppState,
1409 task_id: &StepId,
1410 agent: &str,
1411 attempt: u32,
1412 system: Option<String>,
1413 ) -> String {
1414 let handle = format!("wh-{}", mlua_swarm::types::secure_hex(4));
1415 let task_id = task_id.clone();
1416 let agent = agent.to_string();
1417 let handle_clone = handle.clone();
1418 state
1419 .engine
1420 .with_state("test.seed_task_with_handle", move |s| {
1421 let mut task = mlua_swarm::core::state::TaskState::new(
1422 task_id.clone(),
1423 mlua_swarm::core::state::TaskSpec {
1424 agent: agent.clone(),
1425 initial_directive: json!("x"),
1426 step_ctx: None,
1427 },
1428 );
1429 task.attempt = attempt;
1430 s.tasks.insert(task_id.clone(), task);
1431 s.systems.insert((task_id.clone(), attempt), system);
1432 let token = CapToken {
1433 agent_id: agent,
1434 role: mlua_swarm::Role::Worker,
1435 scopes: vec!["*".to_string()],
1436 issued_at: 0,
1437 expire_at: u64::MAX,
1438 max_uses: None,
1439 nonce: format!("test-nonce-{task_id}"),
1440 sig_hex: String::new(),
1441 };
1442 let fp = token.fingerprint();
1443 s.tokens.insert(
1444 fp.clone(),
1445 mlua_swarm::core::state::CapTokenRecord {
1446 token,
1447 uses_left: None,
1448 revoked: false,
1449 task_id: Some(task_id),
1450 },
1451 );
1452 s.worker_handles.insert(handle_clone, fp);
1453 })
1454 .await
1455 .expect("seed_task_with_handle");
1456 handle
1457 }
1458
1459 fn bearer_headers(handle: &str) -> HeaderMap {
1460 let mut headers = HeaderMap::new();
1461 headers.insert(
1462 AUTHORIZATION,
1463 format!("Bearer {handle}").parse().expect("header value"),
1464 );
1465 headers
1466 }
1467
1468 #[tokio::test]
1472 async fn worker_prompt_system_returns_raw_bytes_for_baked_system() {
1473 let data_store: Arc<dyn OutputStore> = Arc::new(InMemoryOutputStore::new());
1474 let run_store: Arc<dyn RunStore> = Arc::new(InMemoryRunStore::new());
1475 let state = test_state(data_store, run_store);
1476 let task_id = StepId::new();
1477 let rendered = "# Hello\n\nThis is the baked system prompt.".to_string();
1478 let handle =
1479 seed_task_with_handle(&state, &task_id, "planner", 1, Some(rendered.clone())).await;
1480
1481 let resp = worker_prompt_system(
1482 State(state.clone()),
1483 bearer_headers(&handle),
1484 Query(PromptSystemQuery {
1485 task_id: task_id.clone(),
1486 attempt: 1,
1487 }),
1488 )
1489 .await
1490 .expect("worker_prompt_system")
1491 .into_response();
1492
1493 assert_eq!(resp.status(), StatusCode::OK);
1494 let content_type = resp
1495 .headers()
1496 .get(header::CONTENT_TYPE)
1497 .expect("content-type header")
1498 .to_str()
1499 .expect("ascii");
1500 assert_eq!(content_type, "text/plain; charset=utf-8");
1501 let body_bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
1502 .await
1503 .expect("body bytes");
1504 assert_eq!(body_bytes.as_ref(), rendered.as_bytes());
1505 }
1506
1507 #[tokio::test]
1510 async fn worker_prompt_system_404s_when_no_baked_system() {
1511 let data_store: Arc<dyn OutputStore> = Arc::new(InMemoryOutputStore::new());
1512 let run_store: Arc<dyn RunStore> = Arc::new(InMemoryRunStore::new());
1513 let state = test_state(data_store, run_store);
1514 let task_id = StepId::new();
1515 let handle = seed_task_with_handle(&state, &task_id, "planner", 1, None).await;
1516
1517 let result = worker_prompt_system(
1518 State(state.clone()),
1519 bearer_headers(&handle),
1520 Query(PromptSystemQuery {
1521 task_id: task_id.clone(),
1522 attempt: 1,
1523 }),
1524 )
1525 .await;
1526 let err = match result {
1527 Ok(_) => panic!("expected 404 ApiError, got Ok"),
1528 Err(e) => e,
1529 };
1530 assert_eq!(err.into_response().status(), StatusCode::NOT_FOUND);
1531 }
1532
1533 #[tokio::test]
1536 async fn worker_prompt_system_rejects_handle_task_mismatch() {
1537 let data_store: Arc<dyn OutputStore> = Arc::new(InMemoryOutputStore::new());
1538 let run_store: Arc<dyn RunStore> = Arc::new(InMemoryRunStore::new());
1539 let state = test_state(data_store, run_store);
1540 let task_id = StepId::new();
1541 let other_task_id = StepId::new();
1542 let handle =
1543 seed_task_with_handle(&state, &task_id, "planner", 1, Some("x".to_string())).await;
1544
1545 let result = worker_prompt_system(
1546 State(state.clone()),
1547 bearer_headers(&handle),
1548 Query(PromptSystemQuery {
1549 task_id: other_task_id,
1550 attempt: 1,
1551 }),
1552 )
1553 .await;
1554 let err = match result {
1555 Ok(_) => panic!("expected 400 ApiError for task mismatch, got Ok"),
1556 Err(e) => e,
1557 };
1558 assert_eq!(err.into_response().status(), StatusCode::BAD_REQUEST);
1559 }
1560
1561 #[tokio::test]
1565 async fn agent_render_size_returns_null_for_unknown_agent() {
1566 let data_store: Arc<dyn OutputStore> = Arc::new(InMemoryOutputStore::new());
1567 let run_store: Arc<dyn RunStore> = Arc::new(InMemoryRunStore::new());
1568 let state = test_state(data_store, run_store);
1569
1570 let Json(body) = agent_render_size(
1571 State(state.clone()),
1572 axum::extract::Path("never-dispatched".to_string()),
1573 )
1574 .await;
1575 assert_eq!(body.agent, "never-dispatched");
1576 assert_eq!(body.last_rendered_bytes, None);
1577 }
1578
1579 #[tokio::test]
1582 async fn agent_render_size_reports_last_rendered_bytes() {
1583 let data_store: Arc<dyn OutputStore> = Arc::new(InMemoryOutputStore::new());
1584 let run_store: Arc<dyn RunStore> = Arc::new(InMemoryRunStore::new());
1585 let state = test_state(data_store, run_store);
1586 let task_id = StepId::new();
1587 state
1588 .engine
1589 .with_state("test.seed_agent_ctx_for_bake", {
1590 let task_id = task_id.clone();
1591 move |s| {
1592 s.tasks.insert(
1593 task_id.clone(),
1594 mlua_swarm::core::state::TaskState::new(
1595 task_id,
1596 mlua_swarm::core::state::TaskSpec {
1597 agent: "coder".to_string(),
1598 initial_directive: json!("x"),
1599 step_ctx: None,
1600 },
1601 ),
1602 );
1603 }
1604 })
1605 .await
1606 .expect("seed task");
1607 state
1608 .engine
1609 .bake_worker_system_prompt(&task_id, 1, Some("z".repeat(42)))
1610 .await
1611 .expect("bake_worker_system_prompt");
1612
1613 let Json(body) = agent_render_size(
1614 State(state.clone()),
1615 axum::extract::Path("coder".to_string()),
1616 )
1617 .await;
1618 assert_eq!(body.agent, "coder");
1619 assert_eq!(body.last_rendered_bytes, Some(42));
1620 }
1621
1622 #[tokio::test]
1630 async fn worker_artifact_stages_and_204s_for_valid_request() {
1631 let data_store: Arc<dyn OutputStore> = Arc::new(InMemoryOutputStore::new());
1632 let run_store: Arc<dyn RunStore> = Arc::new(InMemoryRunStore::new());
1633 let state = test_state(data_store, run_store);
1634 let task_id = StepId::new();
1635 let handle = seed_task_with_handle(&state, &task_id, "planner", 1, None).await;
1636
1637 let status = worker_artifact(
1638 State(state.clone()),
1639 bearer_headers(&handle),
1640 Query(ArtifactQuery {
1641 name: "summary".to_string(),
1642 }),
1643 axum::body::Bytes::from_static(b"hello artifact\n"),
1644 )
1645 .await
1646 .expect("worker_artifact");
1647 assert_eq!(status, StatusCode::NO_CONTENT);
1648
1649 let tail = state.engine.output_tail(&task_id, 1).await;
1650 assert_eq!(tail.len(), 1, "tail: {tail:?}");
1651 match &tail[0] {
1652 OutputEvent::Artifact { name, content } => {
1653 assert_eq!(name, "summary");
1654 match content {
1655 ContentRef::Inline { value } => {
1656 assert_eq!(value, &json!("hello artifact"));
1657 }
1658 other => panic!("expected Inline content, got {other:?}"),
1659 }
1660 }
1661 other => panic!("expected Artifact event, got {other:?}"),
1662 }
1663 }
1664
1665 #[tokio::test]
1672 async fn worker_artifact_rejects_blank_name() {
1673 let data_store: Arc<dyn OutputStore> = Arc::new(InMemoryOutputStore::new());
1674 let run_store: Arc<dyn RunStore> = Arc::new(InMemoryRunStore::new());
1675 let state = test_state(data_store, run_store);
1676 let task_id = StepId::new();
1677 let handle = seed_task_with_handle(&state, &task_id, "planner", 1, None).await;
1678
1679 let result = worker_artifact(
1680 State(state.clone()),
1681 bearer_headers(&handle),
1682 Query(ArtifactQuery {
1683 name: " ".to_string(),
1684 }),
1685 axum::body::Bytes::from_static(b"x"),
1686 )
1687 .await;
1688 let err = match result {
1689 Ok(_) => panic!("expected 400 ApiError for blank name, got Ok"),
1690 Err(e) => e,
1691 };
1692 assert_eq!(err.into_response().status(), StatusCode::BAD_REQUEST);
1693
1694 assert!(state.engine.output_tail(&task_id, 1).await.is_empty());
1696 }
1697
1698 #[tokio::test]
1704 async fn worker_artifact_staging_same_name_twice_appends_both_events_in_order() {
1705 let data_store: Arc<dyn OutputStore> = Arc::new(InMemoryOutputStore::new());
1706 let run_store: Arc<dyn RunStore> = Arc::new(InMemoryRunStore::new());
1707 let state = test_state(data_store, run_store);
1708 let task_id = StepId::new();
1709 let handle = seed_task_with_handle(&state, &task_id, "planner", 1, None).await;
1710
1711 for body in [b"first".as_slice(), b"second".as_slice()] {
1712 worker_artifact(
1713 State(state.clone()),
1714 bearer_headers(&handle),
1715 Query(ArtifactQuery {
1716 name: "a".to_string(),
1717 }),
1718 axum::body::Bytes::copy_from_slice(body),
1719 )
1720 .await
1721 .expect("worker_artifact");
1722 }
1723
1724 let tail = state.engine.output_tail(&task_id, 1).await;
1725 assert_eq!(tail.len(), 2, "tail: {tail:?}");
1726 let values: Vec<&str> = tail
1727 .iter()
1728 .map(|ev| match ev {
1729 OutputEvent::Artifact {
1730 content: ContentRef::Inline { value },
1731 ..
1732 } => value.as_str().expect("string value"),
1733 other => panic!("expected Artifact/Inline event, got {other:?}"),
1734 })
1735 .collect();
1736 assert_eq!(values, vec!["first", "second"]);
1737 }
1738
1739 async fn link_task_to_run(state: &AppState, task_id: &StepId, attempt: u32, run_id: &RunId) {
1747 let tid = task_id.clone();
1748 let rid_str = run_id.to_string();
1749 state
1750 .engine
1751 .with_state("test.link_task_to_run", move |s| {
1752 let mut entry = mlua_swarm::core::state::AgentCtxEntry::default();
1753 entry.view.run_id = Some(rid_str);
1754 s.agent_ctx.insert((tid, attempt), entry);
1755 })
1756 .await
1757 .expect("link_task_to_run");
1758 }
1759
1760 #[tokio::test]
1765 async fn submit_and_artifact_against_terminal_run_return_410() {
1766 let data_store: Arc<dyn OutputStore> = Arc::new(InMemoryOutputStore::new());
1767 let run_store: Arc<dyn RunStore> = Arc::new(InMemoryRunStore::new());
1768 let state = test_state(data_store, run_store.clone());
1769 let task_id = StepId::new();
1770 let handle = seed_task_with_handle(&state, &task_id, "planner", 1, None).await;
1771
1772 let owner_task = TaskId::new();
1773 let run_id = RunId::new();
1774 let mut rec = run_record(&owner_task, &run_id, vec![]);
1775 rec.status = RunStatus::Failed;
1776 run_store.create(rec).await.expect("run create");
1777 link_task_to_run(&state, &task_id, 1, &run_id).await;
1778
1779 let err = worker_submit(
1780 State(state.clone()),
1781 bearer_headers(&handle),
1782 Query(SubmitQuery { ok: None }),
1783 axum::body::Bytes::from_static(b"LATE OUTPUT"),
1784 )
1785 .await
1786 .expect_err("a submit against a Failed run must be rejected");
1787 assert_eq!(err.status, StatusCode::GONE);
1788 assert!(
1789 err.message.contains(&run_id.to_string()),
1790 "the 410 must name the terminal run: {}",
1791 err.message
1792 );
1793
1794 let err = worker_artifact(
1795 State(state.clone()),
1796 bearer_headers(&handle),
1797 Query(ArtifactQuery {
1798 name: "part.md".to_string(),
1799 }),
1800 axum::body::Bytes::from_static(b"LATE PART"),
1801 )
1802 .await
1803 .expect_err("an artifact staged against a Failed run must be rejected");
1804 assert_eq!(err.status, StatusCode::GONE);
1805
1806 let tail = state.engine.output_tail(&task_id, 1).await;
1808 assert!(tail.is_empty(), "rejected submits must not land: {tail:?}");
1809 }
1810
1811 #[tokio::test]
1815 async fn terminal_run_guard_is_fail_open() {
1816 let data_store: Arc<dyn OutputStore> = Arc::new(InMemoryOutputStore::new());
1817 let run_store: Arc<dyn RunStore> = Arc::new(InMemoryRunStore::new());
1818 let state = test_state(data_store, run_store.clone());
1819 let task_id = StepId::new();
1820 seed_task_with_handle(&state, &task_id, "planner", 1, None).await;
1821
1822 reject_if_run_terminal(&state, &task_id, 1)
1824 .await
1825 .expect("no linkage must fail open");
1826
1827 let unknown_run = RunId::new();
1829 link_task_to_run(&state, &task_id, 1, &unknown_run).await;
1830 reject_if_run_terminal(&state, &task_id, 1)
1831 .await
1832 .expect("unknown run must fail open");
1833
1834 let owner_task = TaskId::new();
1836 let live_run = RunId::new();
1837 run_store
1838 .create(run_record(&owner_task, &live_run, vec![]))
1839 .await
1840 .expect("run create");
1841 link_task_to_run(&state, &task_id, 1, &live_run).await;
1842 reject_if_run_terminal(&state, &task_id, 1)
1843 .await
1844 .expect("a Running run must pass the guard");
1845 }
1846
1847 fn degradation_body(tool: &str, note: Option<&str>) -> DegradationBody {
1852 DegradationBody {
1853 tool: tool.to_string(),
1854 error: "boom".to_string(),
1855 fallback: "used cached value".to_string(),
1856 note: note.map(str::to_string),
1857 }
1858 }
1859
1860 async fn link_task_to_run_with_agent(
1866 state: &AppState,
1867 task_id: &StepId,
1868 attempt: u32,
1869 run_id: &RunId,
1870 agent: &str,
1871 ) {
1872 let tid = task_id.clone();
1873 let rid_str = run_id.to_string();
1874 let agent = agent.to_string();
1875 state
1876 .engine
1877 .with_state("test.link_task_to_run_with_agent", move |s| {
1878 let mut entry = mlua_swarm::core::state::AgentCtxEntry::default();
1879 entry.view.run_id = Some(rid_str);
1880 entry.view.agent = agent;
1881 s.agent_ctx.insert((tid, attempt), entry);
1882 })
1883 .await
1884 .expect("link_task_to_run_with_agent");
1885 }
1886
1887 #[tokio::test]
1892 async fn worker_degradation_persists_entry_when_run_tracked() {
1893 let data_store: Arc<dyn OutputStore> = Arc::new(InMemoryOutputStore::new());
1894 let run_store: Arc<dyn RunStore> = Arc::new(InMemoryRunStore::new());
1895 let state = test_state(data_store, run_store.clone());
1896 let task_id = StepId::new();
1897 let handle = seed_task_with_handle(&state, &task_id, "planner", 1, None).await;
1898
1899 let owner_task = TaskId::new();
1900 let run_id = RunId::new();
1901 run_store
1902 .create(run_record(&owner_task, &run_id, vec![]))
1903 .await
1904 .expect("run create");
1905 link_task_to_run_with_agent(&state, &task_id, 1, &run_id, "planner").await;
1906
1907 let status = worker_degradation(
1908 State(state.clone()),
1909 bearer_headers(&handle),
1910 Json(degradation_body("web_search", Some("rate limited"))),
1911 )
1912 .await
1913 .expect("worker_degradation");
1914 assert_eq!(status, StatusCode::NO_CONTENT);
1915
1916 let rec = run_store.get(&run_id).await.expect("run get");
1917 assert_eq!(
1918 rec.degradations.len(),
1919 1,
1920 "degradations: {:?}",
1921 rec.degradations
1922 );
1923 let entry = &rec.degradations[0];
1924 assert_eq!(entry.tool, "web_search");
1925 assert_eq!(entry.error, "boom");
1926 assert_eq!(entry.fallback, "used cached value");
1927 assert_eq!(entry.note.as_deref(), Some("rate limited"));
1928 assert_eq!(entry.step_ref.as_deref(), Some("planner"));
1929 assert_eq!(entry.attempt, Some(1));
1930 assert!(entry.at > 0, "at must be a real timestamp: {}", entry.at);
1931 }
1932
1933 #[tokio::test]
1935 async fn worker_degradation_appends_in_order() {
1936 let data_store: Arc<dyn OutputStore> = Arc::new(InMemoryOutputStore::new());
1937 let run_store: Arc<dyn RunStore> = Arc::new(InMemoryRunStore::new());
1938 let state = test_state(data_store, run_store.clone());
1939 let task_id = StepId::new();
1940 let handle = seed_task_with_handle(&state, &task_id, "planner", 1, None).await;
1941
1942 let owner_task = TaskId::new();
1943 let run_id = RunId::new();
1944 run_store
1945 .create(run_record(&owner_task, &run_id, vec![]))
1946 .await
1947 .expect("run create");
1948 link_task_to_run(&state, &task_id, 1, &run_id).await;
1949
1950 for tool in ["first_tool", "second_tool"] {
1951 worker_degradation(
1952 State(state.clone()),
1953 bearer_headers(&handle),
1954 Json(degradation_body(tool, None)),
1955 )
1956 .await
1957 .expect("worker_degradation");
1958 }
1959
1960 let rec = run_store.get(&run_id).await.expect("run get");
1961 let tools: Vec<&str> = rec.degradations.iter().map(|e| e.tool.as_str()).collect();
1962 assert_eq!(tools, vec!["first_tool", "second_tool"]);
1963 }
1964
1965 #[tokio::test]
1969 async fn worker_degradation_silent_ok_when_no_run_tracked() {
1970 let data_store: Arc<dyn OutputStore> = Arc::new(InMemoryOutputStore::new());
1971 let run_store: Arc<dyn RunStore> = Arc::new(InMemoryRunStore::new());
1972 let state = test_state(data_store, run_store);
1973 let task_id = StepId::new();
1974 let handle = seed_task_with_handle(&state, &task_id, "planner", 1, None).await;
1975
1976 let status = worker_degradation(
1977 State(state.clone()),
1978 bearer_headers(&handle),
1979 Json(degradation_body("some_tool", None)),
1980 )
1981 .await
1982 .expect("worker_degradation must not error on missing run linkage");
1983 assert_eq!(status, StatusCode::NO_CONTENT);
1984 }
1985
1986 #[tokio::test]
1989 async fn worker_degradation_rejects_terminal_run() {
1990 let data_store: Arc<dyn OutputStore> = Arc::new(InMemoryOutputStore::new());
1991 let run_store: Arc<dyn RunStore> = Arc::new(InMemoryRunStore::new());
1992 let state = test_state(data_store, run_store.clone());
1993 let task_id = StepId::new();
1994 let handle = seed_task_with_handle(&state, &task_id, "planner", 1, None).await;
1995
1996 let owner_task = TaskId::new();
1997 let run_id = RunId::new();
1998 let mut rec = run_record(&owner_task, &run_id, vec![]);
1999 rec.status = RunStatus::Done;
2000 run_store.create(rec).await.expect("run create");
2001 link_task_to_run(&state, &task_id, 1, &run_id).await;
2002
2003 let err = worker_degradation(
2004 State(state.clone()),
2005 bearer_headers(&handle),
2006 Json(degradation_body("some_tool", None)),
2007 )
2008 .await
2009 .expect_err("a degradation against a Done run must be rejected");
2010 assert_eq!(err.status, StatusCode::GONE);
2011
2012 let rec = run_store.get(&run_id).await.expect("run get");
2013 assert!(
2014 rec.degradations.is_empty(),
2015 "rejected degradation must not land: {:?}",
2016 rec.degradations
2017 );
2018 }
2019}