1use axum::{
45 extract::{Query, State},
46 http::{header::AUTHORIZATION, HeaderMap, StatusCode},
47 Json,
48};
49use mlua_swarm::core::agent_context::StepPointer;
50use mlua_swarm::{CapToken, ContentRef, OutputEvent, RunId, StepId, WorkerPayload};
51use serde::Deserialize;
52use serde_json::Value;
53
54use crate::projection::McpQueryAdapter;
55use crate::{ApiError, AppState};
56
57#[derive(Debug, Deserialize)]
59pub struct PromptQuery {
60 pub task_id: StepId,
64}
65
66pub async fn worker_prompt(
72 State(state): State<AppState>,
73 headers: HeaderMap,
74 Query(q): Query<PromptQuery>,
75) -> Result<Json<WorkerPayload>, ApiError> {
76 let task_id = q.task_id;
77 let bearer = extract_bearer_raw(&headers)?;
78 let mut payload = if let Some(handle) = parse_worker_handle(&bearer) {
79 let resolved = state
81 .engine
82 .task_id_from_handle(handle)
83 .await
84 .map_err(|e| ApiError::engine(format!("task_id_from_handle: {e}")))?;
85 if resolved != task_id {
86 return Err(ApiError::bad_request(format!(
87 "handle {handle} is bound to task {resolved}, not {task_id}"
88 )));
89 }
90 state
91 .engine
92 .fetch_worker_payload_trusted(&task_id)
93 .await
94 .map_err(|e| ApiError::engine(format!("fetch_worker_payload_trusted: {e}")))?
95 } else {
96 let token = CapToken::decode(bearer.trim())
98 .map_err(|e| ApiError::bad_request(format!("invalid token: {e}")))?;
99 state
100 .engine
101 .fetch_worker_payload(&token, &task_id)
102 .await
103 .map_err(|e| ApiError::engine(format!("fetch_worker_payload: {e}")))?
104 };
105 assemble_step_pointers(&state, &mut payload).await;
106 Ok(Json(payload))
107}
108
109async fn assemble_step_pointers(state: &AppState, payload: &mut WorkerPayload) {
124 let Some(context) = payload.context.as_mut() else {
125 return;
126 };
127 let Some(run_id_str) = context.run_id.clone() else {
128 return;
129 };
130 let Ok(run_id) = RunId::parse(run_id_str) else {
131 return;
132 };
133
134 let adapter = McpQueryAdapter::new(state.data_store.clone(), state.run_store.clone());
135 let Ok((run, resolved_steps)) = adapter.list_steps_by_run_id(&run_id).await else {
136 return;
137 };
138
139 let policy = state
140 .engine
141 .context_policy_for(&payload.task_id, payload.attempt)
142 .await;
143 let self_name = payload.agent.clone();
144
145 let mut pointers = Vec::new();
146 for step in &resolved_steps {
147 if step.name == self_name || !policy.allows_step(&step.name) {
148 continue;
149 }
150 if let Some((size_bytes, file_path, content_url, sha256)) =
151 crate::projection::resolve_step_pointer_fields(state, &run, step).await
152 {
153 pointers.push(StepPointer {
154 name: step.name.clone(),
155 size_bytes,
156 file_path,
157 content_url,
158 sha256,
159 });
160 }
161 }
162 context.steps = pointers;
163}
164
165#[derive(Debug, Deserialize)]
167pub struct WorkerResultReq {
168 pub task_id: StepId,
171 pub value: Value,
173 #[serde(default = "default_ok_true")]
177 pub ok: bool,
178 #[serde(default)]
181 pub attempt: Option<u32>,
182}
183
184fn default_ok_true() -> bool {
185 true
186}
187
188pub async fn worker_result(
191 State(state): State<AppState>,
192 headers: HeaderMap,
193 Json(req): Json<WorkerResultReq>,
194) -> Result<StatusCode, ApiError> {
195 let token = decode_worker_bearer(&headers)?;
196 let task_id = req.task_id.clone();
197
198 let attempt = match req.attempt {
200 Some(n) => n,
201 None => state
202 .engine
203 .task_attempt(&task_id)
204 .await
205 .map_err(|e| ApiError::engine(format!("task_attempt: {e}")))?,
206 };
207
208 let event = OutputEvent::Final {
209 content: ContentRef::Inline {
210 value: req.value.clone(),
211 },
212 ok: req.ok,
213 };
214 state
215 .engine
216 .submit_output(&token, &task_id, attempt, event)
217 .await
218 .map_err(|e| ApiError::engine(format!("submit_output: {e}")))?;
219 state
220 .engine
221 .post_result(&token, &task_id, req.value)
222 .await
223 .map_err(|e| ApiError::engine(format!("post_result: {e}")))?;
224 Ok(StatusCode::NO_CONTENT)
225}
226
227#[derive(Debug, Deserialize, Default)]
241pub struct SubmitQuery {
242 #[serde(default)]
246 pub ok: Option<bool>,
247}
248
249pub async fn worker_submit(
255 State(state): State<AppState>,
256 headers: HeaderMap,
257 Query(q): Query<SubmitQuery>,
258 body: axum::body::Bytes,
259) -> Result<StatusCode, ApiError> {
260 let bearer = extract_bearer_raw(&headers)?;
263 let task_id = if let Some(handle) = parse_worker_handle(&bearer) {
264 state
265 .engine
266 .task_id_from_handle(handle)
267 .await
268 .map_err(|e| ApiError::engine(format!("task_id_from_handle: {e}")))?
269 } else {
270 let token = CapToken::decode(bearer.trim())
271 .map_err(|e| ApiError::bad_request(format!("invalid token: {e}")))?;
272 state
273 .engine
274 .task_id_from_token(&token)
275 .await
276 .map_err(|e| ApiError::engine(format!("task_id_from_token: {e}")))?
277 };
278 let attempt = state
279 .engine
280 .task_attempt(&task_id)
281 .await
282 .map_err(|e| ApiError::engine(format!("task_attempt: {e}")))?;
283 let body_str = String::from_utf8_lossy(&body).trim_end().to_string();
288 let value = Value::String(body_str);
289
290 let ok = q.ok.unwrap_or(true);
296 state
297 .engine
298 .submit_worker_result_trusted(&task_id, attempt, value, ok)
299 .await
300 .map_err(|e| ApiError::engine(format!("submit_worker_result_trusted: {e}")))?;
301 Ok(StatusCode::NO_CONTENT)
302}
303
304fn extract_bearer_raw(headers: &HeaderMap) -> Result<String, ApiError> {
308 let v = headers
309 .get(AUTHORIZATION)
310 .ok_or_else(|| ApiError::bad_request("missing Authorization header".into()))?
311 .to_str()
312 .map_err(|_| ApiError::bad_request("invalid Authorization header encoding".into()))?;
313 let s = v
314 .strip_prefix("Bearer ")
315 .ok_or_else(|| ApiError::bad_request("Authorization must be 'Bearer <token>'".into()))?
316 .trim();
317 if s.is_empty() {
318 return Err(ApiError::bad_request("Bearer is empty".into()));
319 }
320 Ok(s.to_string())
321}
322
323fn parse_worker_handle(s: &str) -> Option<&str> {
327 let s = s.trim();
328 if s.starts_with("wh-")
329 && s.len() >= 5
330 && s.len() <= 64
331 && s[3..].chars().all(|c| c.is_ascii_alphanumeric())
332 {
333 Some(s)
334 } else {
335 None
336 }
337}
338
339fn decode_worker_bearer(headers: &HeaderMap) -> Result<CapToken, ApiError> {
343 let v = headers
344 .get(AUTHORIZATION)
345 .ok_or_else(|| ApiError::bad_request("missing Authorization header".into()))?
346 .to_str()
347 .map_err(|_| ApiError::bad_request("invalid Authorization header encoding".into()))?;
348 let encoded = v
349 .strip_prefix("Bearer ")
350 .ok_or_else(|| ApiError::bad_request("Authorization must be 'Bearer <token>'".into()))?
351 .trim();
352 if encoded.is_empty() {
353 return Err(ApiError::bad_request("Bearer token is empty".into()));
354 }
355 CapToken::decode(encoded).map_err(|e| ApiError::bad_request(format!("invalid token: {e}")))
356}
357
358#[cfg(test)]
363mod tests {
364 use super::*;
365 use mlua_swarm::core::agent_context::AgentContextView;
366 use mlua_swarm::core::config::EngineCfg;
367 use mlua_swarm::core::engine::Engine;
368 use mlua_swarm::store::output::{InMemoryOutputStore, OutputStore};
369 use mlua_swarm::store::run::{InMemoryRunStore, RunRecord, RunStatus, RunStore, StepEntry};
370 use mlua_swarm::store::task::InMemoryTaskStore;
371 use mlua_swarm::{RunId, StepId, TaskId};
372 use serde_json::json;
373 use std::collections::HashMap;
374 use std::sync::Arc;
375 use tokio::sync::Mutex;
376
377 fn test_state(data_store: Arc<dyn OutputStore>, run_store: Arc<dyn RunStore>) -> AppState {
383 let engine = Engine::new(EngineCfg::default());
384 let compiler = mlua_swarm::Compiler::new(crate::default_registry());
385 let launch = Arc::new(mlua_swarm::TaskLaunchService::new(engine.clone(), compiler));
386 AppState {
387 engine,
388 sessions: Arc::new(Mutex::new(crate::SessionStore::default())),
389 task_app: Arc::new(mlua_swarm::TaskApplication::new_inline_only(launch)),
390 ws_operator_factory: None,
391 data_store,
392 operator_sessions: Arc::new(Mutex::new(HashMap::new())),
393 roles_to_sid: Arc::new(Mutex::new(HashMap::new())),
394 task_store: Arc::new(InMemoryTaskStore::new()),
395 run_store,
396 base_url: None,
397 }
398 }
399
400 async fn append_final(
401 data_store: &Arc<dyn OutputStore>,
402 task_id: &str,
403 producer: &str,
404 value: Value,
405 ) {
406 data_store
407 .append(
408 task_id,
409 1,
410 producer,
411 OutputEvent::Final {
412 content: ContentRef::Inline { value },
413 ok: true,
414 },
415 vec![],
416 )
417 .await
418 .expect("append final");
419 }
420
421 fn step_entry(step_id: &StepId, step_ref: &str) -> StepEntry {
422 StepEntry {
423 step_id: step_id.clone(),
424 step_ref: Some(step_ref.to_string()),
425 status: Some("passed".to_string()),
426 at: 0,
427 }
428 }
429
430 fn run_record(task_id: &TaskId, run_id: &RunId, step_entries: Vec<StepEntry>) -> RunRecord {
431 RunRecord {
432 id: run_id.clone(),
433 task_id: task_id.clone(),
434 status: RunStatus::Running,
435 step_entries,
436 operator_sid: None,
437 result_ref: None,
438 created_at: 0,
439 updated_at: 0,
440 }
441 }
442
443 fn consumer_payload(consumer_step_id: &StepId, run_id: &RunId) -> WorkerPayload {
444 WorkerPayload {
445 task_id: consumer_step_id.clone(),
446 attempt: 1,
447 agent: "consumer".to_string(),
448 system: None,
449 prompt: String::new(),
450 context: Some(AgentContextView {
451 task_id: consumer_step_id.to_string(),
452 agent: "consumer".to_string(),
453 attempt: 1,
454 run_id: Some(run_id.to_string()),
455 ..Default::default()
456 }),
457 }
458 }
459
460 #[tokio::test]
465 async fn context_policy_unspecified_yields_every_submitted_step() {
466 let data_store: Arc<dyn OutputStore> = Arc::new(InMemoryOutputStore::new());
467 let run_store: Arc<dyn RunStore> = Arc::new(InMemoryRunStore::new());
468 let task_id = TaskId::new();
469 let run_id = RunId::new();
470 let planner_id = StepId::new();
471 let coder_id = StepId::new();
472
473 append_final(
474 &data_store,
475 planner_id.as_str(),
476 "planner",
477 json!({"plan": "x"}),
478 )
479 .await;
480 append_final(
481 &data_store,
482 coder_id.as_str(),
483 "coder",
484 json!({"code": "y"}),
485 )
486 .await;
487 run_store
488 .create(run_record(
489 &task_id,
490 &run_id,
491 vec![
492 step_entry(&planner_id, "planner"),
493 step_entry(&coder_id, "coder"),
494 ],
495 ))
496 .await
497 .expect("create run");
498
499 let state = test_state(data_store, run_store);
500 let consumer_id = StepId::new();
501 let mut payload = consumer_payload(&consumer_id, &run_id);
502 assemble_step_pointers(&state, &mut payload).await;
503
504 let names: Vec<&str> = payload
505 .context
506 .as_ref()
507 .expect("context")
508 .steps
509 .iter()
510 .map(|p| p.name.as_str())
511 .collect();
512 assert!(names.contains(&"planner"), "names: {names:?}");
513 assert!(names.contains(&"coder"), "names: {names:?}");
514 }
515
516 #[tokio::test]
518 async fn context_policy_steps_include_list_filters_to_named_steps() {
519 let data_store: Arc<dyn OutputStore> = Arc::new(InMemoryOutputStore::new());
520 let run_store: Arc<dyn RunStore> = Arc::new(InMemoryRunStore::new());
521 let task_id = TaskId::new();
522 let run_id = RunId::new();
523 let planner_id = StepId::new();
524 let coder_id = StepId::new();
525 append_final(&data_store, planner_id.as_str(), "planner", json!("x")).await;
526 append_final(&data_store, coder_id.as_str(), "coder", json!("y")).await;
527 run_store
528 .create(run_record(
529 &task_id,
530 &run_id,
531 vec![
532 step_entry(&planner_id, "planner"),
533 step_entry(&coder_id, "coder"),
534 ],
535 ))
536 .await
537 .expect("create run");
538
539 let state = test_state(data_store, run_store);
540 let consumer_id = StepId::new();
541 state
542 .engine
543 .with_state("test.seed_policy", {
544 let consumer_id = consumer_id.clone();
545 move |s| {
546 s.context_policies.insert(
547 (consumer_id, 1),
548 mlua_swarm_schema::ContextPolicy {
549 steps: Some(vec!["planner".to_string()]),
550 ..Default::default()
551 },
552 );
553 }
554 })
555 .await
556 .expect("seed policy");
557
558 let mut payload = consumer_payload(&consumer_id, &run_id);
559 assemble_step_pointers(&state, &mut payload).await;
560
561 let names: Vec<&str> = payload
562 .context
563 .as_ref()
564 .expect("context")
565 .steps
566 .iter()
567 .map(|p| p.name.as_str())
568 .collect();
569 assert_eq!(names, vec!["planner"], "names: {names:?}");
570 }
571
572 #[tokio::test]
574 async fn context_policy_steps_empty_list_yields_no_pointers() {
575 let data_store: Arc<dyn OutputStore> = Arc::new(InMemoryOutputStore::new());
576 let run_store: Arc<dyn RunStore> = Arc::new(InMemoryRunStore::new());
577 let task_id = TaskId::new();
578 let run_id = RunId::new();
579 let planner_id = StepId::new();
580 append_final(&data_store, planner_id.as_str(), "planner", json!("x")).await;
581 run_store
582 .create(run_record(
583 &task_id,
584 &run_id,
585 vec![step_entry(&planner_id, "planner")],
586 ))
587 .await
588 .expect("create run");
589
590 let state = test_state(data_store, run_store);
591 let consumer_id = StepId::new();
592 state
593 .engine
594 .with_state("test.seed_policy", {
595 let consumer_id = consumer_id.clone();
596 move |s| {
597 s.context_policies.insert(
598 (consumer_id, 1),
599 mlua_swarm_schema::ContextPolicy {
600 steps: Some(vec![]),
601 ..Default::default()
602 },
603 );
604 }
605 })
606 .await
607 .expect("seed policy");
608
609 let mut payload = consumer_payload(&consumer_id, &run_id);
610 assemble_step_pointers(&state, &mut payload).await;
611
612 assert!(payload.context.expect("context").steps.is_empty());
613 }
614
615 #[tokio::test]
617 async fn context_policy_steps_exclude_wins_over_steps() {
618 let data_store: Arc<dyn OutputStore> = Arc::new(InMemoryOutputStore::new());
619 let run_store: Arc<dyn RunStore> = Arc::new(InMemoryRunStore::new());
620 let task_id = TaskId::new();
621 let run_id = RunId::new();
622 let planner_id = StepId::new();
623 let coder_id = StepId::new();
624 append_final(&data_store, planner_id.as_str(), "planner", json!("x")).await;
625 append_final(&data_store, coder_id.as_str(), "coder", json!("y")).await;
626 run_store
627 .create(run_record(
628 &task_id,
629 &run_id,
630 vec![
631 step_entry(&planner_id, "planner"),
632 step_entry(&coder_id, "coder"),
633 ],
634 ))
635 .await
636 .expect("create run");
637
638 let state = test_state(data_store, run_store);
639 let consumer_id = StepId::new();
640 state
641 .engine
642 .with_state("test.seed_policy", {
643 let consumer_id = consumer_id.clone();
644 move |s| {
645 s.context_policies.insert(
646 (consumer_id, 1),
647 mlua_swarm_schema::ContextPolicy {
648 steps: Some(vec!["planner".to_string(), "coder".to_string()]),
649 steps_exclude: vec!["planner".to_string()],
650 ..Default::default()
651 },
652 );
653 }
654 })
655 .await
656 .expect("seed policy");
657
658 let mut payload = consumer_payload(&consumer_id, &run_id);
659 assemble_step_pointers(&state, &mut payload).await;
660
661 let names: Vec<&str> = payload
662 .context
663 .as_ref()
664 .expect("context")
665 .steps
666 .iter()
667 .map(|p| p.name.as_str())
668 .collect();
669 assert_eq!(names, vec!["coder"], "names: {names:?}");
670 }
671
672 #[tokio::test]
681 async fn in_flight_step_output_is_visible_before_run_finalizes() {
682 let data_store: Arc<dyn OutputStore> = Arc::new(InMemoryOutputStore::new());
683 let run_store: Arc<dyn RunStore> = Arc::new(InMemoryRunStore::new());
684 let task_id = TaskId::new();
685 let run_id = RunId::new();
686 let step1_id = StepId::new();
687 append_final(
688 &data_store,
689 step1_id.as_str(),
690 "step1",
691 json!({"step1_out": "hi"}),
692 )
693 .await;
694 let mut run = run_record(&task_id, &run_id, vec![step_entry(&step1_id, "step1")]);
695 run.status = RunStatus::Running;
696 run.result_ref = None; run_store.create(run).await.expect("create run");
698
699 let state = test_state(data_store, run_store);
700 let consumer_id = StepId::new();
701 let mut payload = consumer_payload(&consumer_id, &run_id);
702 assemble_step_pointers(&state, &mut payload).await;
703
704 let steps = &payload.context.expect("context").steps;
705 assert_eq!(steps.len(), 1);
706 assert_eq!(steps[0].name, "step1");
707 }
708
709 #[tokio::test]
713 async fn self_agent_name_is_always_excluded() {
714 let data_store: Arc<dyn OutputStore> = Arc::new(InMemoryOutputStore::new());
715 let run_store: Arc<dyn RunStore> = Arc::new(InMemoryRunStore::new());
716 let task_id = TaskId::new();
717 let run_id = RunId::new();
718 let planner_id = StepId::new();
719 let consumer_prior_id = StepId::new();
720 append_final(&data_store, planner_id.as_str(), "planner", json!("x")).await;
721 append_final(
722 &data_store,
723 consumer_prior_id.as_str(),
724 "consumer",
725 json!("self"),
726 )
727 .await;
728 run_store
729 .create(run_record(
730 &task_id,
731 &run_id,
732 vec![
733 step_entry(&planner_id, "planner"),
734 step_entry(&consumer_prior_id, "consumer"),
735 ],
736 ))
737 .await
738 .expect("create run");
739
740 let state = test_state(data_store, run_store);
741 let consumer_id = StepId::new();
742 let mut payload = consumer_payload(&consumer_id, &run_id);
743 assemble_step_pointers(&state, &mut payload).await;
744
745 let names: Vec<&str> = payload
746 .context
747 .as_ref()
748 .expect("context")
749 .steps
750 .iter()
751 .map(|p| p.name.as_str())
752 .collect();
753 assert!(!names.contains(&"consumer"), "names: {names:?}");
754 assert!(names.contains(&"planner"), "names: {names:?}");
755 }
756
757 #[tokio::test]
761 async fn step_pointer_serializes_with_no_preview_or_content_bytes() {
762 let data_store: Arc<dyn OutputStore> = Arc::new(InMemoryOutputStore::new());
763 let run_store: Arc<dyn RunStore> = Arc::new(InMemoryRunStore::new());
764 let task_id = TaskId::new();
765 let run_id = RunId::new();
766 let planner_id = StepId::new();
767 append_final(
768 &data_store,
769 planner_id.as_str(),
770 "planner",
771 json!({"plan": "do the thing, at length".repeat(50)}),
772 )
773 .await;
774 run_store
775 .create(run_record(
776 &task_id,
777 &run_id,
778 vec![step_entry(&planner_id, "planner")],
779 ))
780 .await
781 .expect("create run");
782
783 let state = test_state(data_store, run_store);
784 let consumer_id = StepId::new();
785 let mut payload = consumer_payload(&consumer_id, &run_id);
786 assemble_step_pointers(&state, &mut payload).await;
787
788 let steps = &payload.context.expect("context").steps;
789 assert_eq!(steps.len(), 1);
790 let json_value = serde_json::to_value(&steps[0]).expect("serialize StepPointer");
791 let obj = json_value.as_object().expect("object");
792 for forbidden in ["preview", "content", "value", "bytes"] {
793 assert!(
794 !obj.contains_key(forbidden),
795 "StepPointer must not carry a {forbidden:?} field: {obj:?}"
796 );
797 }
798 assert!(obj.contains_key("name"));
799 assert!(obj.contains_key("size_bytes"));
800 assert!(obj.contains_key("content_url"));
801 assert!(obj.contains_key("sha256"));
802 }
803}