1use async_trait::async_trait;
14use mlua_swarm::{
15 CapToken, Ctx, Operator, SeniorBridge, SessionId, SpawnHook, StepId, WorkerBinding,
16 WorkerError, WorkerResult,
17};
18use serde_json::Value;
19use std::collections::HashMap;
20use tokio::sync::{mpsc, oneshot, Mutex};
21
22use super::protocol::{current_parent_req_id, PendingReply, ServerMsg};
23
24pub struct WSOperatorSession {
26 sid: SessionId,
27 tx: Mutex<Option<mpsc::UnboundedSender<ServerMsg>>>,
30 pending: Mutex<HashMap<String, oneshot::Sender<PendingReply>>>,
33 base_url: Option<std::sync::Arc<str>>,
39}
40
41impl WSOperatorSession {
42 pub(super) fn new_with_base_url(
52 sid: SessionId,
53 tx: mpsc::UnboundedSender<ServerMsg>,
54 base_url: Option<std::sync::Arc<str>>,
55 ) -> Self {
56 Self {
57 sid,
58 tx: Mutex::new(Some(tx)),
59 pending: Mutex::new(HashMap::new()),
60 base_url,
61 }
62 }
63
64 pub(super) async fn replace_tx(&self, new_tx: mpsc::UnboundedSender<ServerMsg>) {
66 *self.tx.lock().await = Some(new_tx);
67 }
68
69 pub(crate) async fn clear_tx(&self) {
71 *self.tx.lock().await = None;
72 }
73
74 pub(super) async fn resolve_pending(&self, req_id: &str, reply: PendingReply) {
77 if let Some(otx) = self.pending.lock().await.remove(req_id) {
78 let _ = otx.send(reply);
79 }
80 }
81
82 async fn send_and_await(&self, req_id: String, msg: ServerMsg) -> Result<PendingReply, String> {
86 let (otx, orx) = oneshot::channel::<PendingReply>();
87 self.pending.lock().await.insert(req_id.clone(), otx);
88
89 let send_result = {
91 let guard = self.tx.lock().await;
92 match guard.as_ref() {
93 Some(tx) => tx
94 .send(msg)
95 .map_err(|_| "ws send channel closed".to_string()),
96 None => Err("ws operator disconnected".to_string()),
97 }
98 };
99 if let Err(e) = send_result {
100 self.pending.lock().await.remove(&req_id);
101 return Err(e);
102 }
103
104 orx.await
105 .map_err(|_| "ws operator: oneshot cancelled (= reply path closed)".to_string())
106 }
107
108 async fn send_oneway(&self, msg: ServerMsg) -> Result<(), String> {
110 let guard = self.tx.lock().await;
111 match guard.as_ref() {
112 Some(tx) => tx
113 .send(msg)
114 .map_err(|_| "ws send channel closed".to_string()),
115 None => Err("ws operator disconnected".to_string()),
116 }
117 }
118}
119
120#[async_trait]
121impl SeniorBridge for WSOperatorSession {
122 async fn ask(&self, task_id: &StepId, question: Value) -> Result<Value, String> {
123 let req_id = format!("{}-ask-{}", self.sid, uuid::Uuid::new_v4());
124 let msg = ServerMsg::Ask {
125 req_id: req_id.clone(),
126 parent_req_id: current_parent_req_id(),
127 task_id: task_id.clone(),
128 question,
129 };
130 match self.send_and_await(req_id, msg).await? {
131 PendingReply::Answer(v) => Ok(v),
132 PendingReply::HookAck { .. } => {
133 Err("ws operator: unexpected hook_ack reply to ask".into())
134 }
135 PendingReply::SpawnAck { .. } => {
136 Err("ws operator: unexpected spawn_ack reply to ask".into())
137 }
138 PendingReply::SpawnHalt { .. } => {
139 Err("ws operator: unexpected spawn_halt reply to ask".into())
140 }
141 }
142 }
143}
144
145#[async_trait]
146impl SpawnHook for WSOperatorSession {
147 async fn before(&self, ctx: &Ctx) -> Result<(), String> {
148 let req_id = format!("{}-hb-{}", self.sid, uuid::Uuid::new_v4());
149 let msg = ServerMsg::HookBefore {
150 req_id: req_id.clone(),
151 parent_req_id: current_parent_req_id(),
152 task_id: ctx.task_id.clone(),
153 agent: ctx.agent.clone(),
154 attempt: ctx.attempt,
155 };
156 match self.send_and_await(req_id, msg).await? {
157 PendingReply::HookAck { ok: true, .. } => Ok(()),
158 PendingReply::HookAck { ok: false, reason } => {
159 Err(reason.unwrap_or_else(|| "ws operator: spawn rejected".into()))
160 }
161 PendingReply::Answer(_) => {
162 Err("ws operator: unexpected answer reply to hook_before".into())
163 }
164 PendingReply::SpawnAck { .. } => {
165 Err("ws operator: unexpected spawn_ack reply to hook_before".into())
166 }
167 PendingReply::SpawnHalt { .. } => {
168 Err("ws operator: unexpected spawn_halt reply to hook_before".into())
169 }
170 }
171 }
172
173 async fn after(&self, ctx: &Ctx, result: &Value) -> Result<(), String> {
174 let req_id = format!("{}-ha-{}", self.sid, uuid::Uuid::new_v4());
175 let msg = ServerMsg::HookAfter {
176 req_id,
177 parent_req_id: current_parent_req_id(),
178 task_id: ctx.task_id.clone(),
179 agent: ctx.agent.clone(),
180 attempt: ctx.attempt,
181 result: result.clone(),
182 };
183 let _ = self.send_oneway(msg).await;
185 Ok(())
186 }
187}
188
189#[async_trait]
190impl Operator for WSOperatorSession {
191 async fn execute(
211 &self,
212 ctx: &Ctx,
213 _system: Option<String>,
214 _prompt: String,
215 worker: Option<WorkerBinding>,
216 worker_token: CapToken,
217 ) -> Result<WorkerResult, WorkerError> {
218 let Some(worker) = worker else {
219 return Err(WorkerError::Failed(format!(
220 "agent '{}' has no worker_binding; WS thin-path requires one \
221 (Blueprint AgentDef.profile.worker_binding)",
222 ctx.agent
223 )));
224 };
225 let req_id = format!("{}-spawn-{}", self.sid, uuid::Uuid::new_v4());
226 let worker_handle = ctx
227 .meta
228 .runtime
229 .get("worker_handle")
230 .and_then(|v| v.as_str())
231 .map(|s| s.to_string());
232 let project_name_alias = ctx
233 .meta
234 .runtime
235 .get("project_name_alias")
236 .and_then(|v| v.as_str());
237 let data_sink_endpoint = ctx
238 .meta
239 .runtime
240 .get("data_sink_endpoint")
241 .and_then(|v| v.as_str());
242 let run_id = ctx.meta.runtime.get("run_id").and_then(|v| v.as_str());
247 let directive = default_spawn_directive(
248 &ctx.agent,
249 ctx.task_id.as_str(),
250 &worker.variant,
251 project_name_alias,
252 data_sink_endpoint,
253 self.base_url.as_deref(),
254 run_id,
255 );
256 let msg = ServerMsg::Spawn {
257 req_id: req_id.clone(),
258 parent_req_id: current_parent_req_id(),
259 task_id: ctx.task_id.clone(),
260 agent: ctx.agent.clone(),
261 attempt: ctx.attempt,
262 capability_token: worker_token.encode(),
263 worker_handle,
264 worker: Some(worker),
265 directive,
266 };
267 match self.send_and_await(req_id, msg).await {
268 Ok(PendingReply::SpawnAck {
269 value,
270 ok,
271 error: None,
272 }) => Ok(WorkerResult { value, ok }),
273 Ok(PendingReply::SpawnAck {
274 error: Some(msg), ..
275 }) => Err(WorkerError::Failed(msg)),
276 Ok(PendingReply::SpawnHalt { value, reason }) => {
283 let marker = serde_json::json!({
284 "halted": true,
285 "reason": reason,
286 "value": value,
287 });
288 Ok(WorkerResult {
289 value: marker,
290 ok: true,
291 })
292 }
293 Ok(_) => Err(WorkerError::Failed(
294 "ws operator: unexpected non-spawn reply".into(),
295 )),
296 Err(e) => Err(WorkerError::Failed(format!("ws operator spawn: {e}"))),
297 }
298 }
299
300 fn requires_worker_binding(&self) -> bool {
301 true
302 }
303}
304
305pub(super) fn default_spawn_directive(
352 agent: &str,
353 task_id: &str,
354 subagent_type: &str,
355 project_name_alias: Option<&str>,
356 data_sink_endpoint: Option<&str>,
357 base_url: Option<&str>,
358 run_id: Option<&str>,
359) -> String {
360 let project_alias_line = match project_name_alias {
364 Some(a) => format!("project_name_alias: {a}\n"),
365 None => String::new(),
366 };
367 let data_endpoint_block = match data_sink_endpoint {
374 Some(base) => format!(
375 "\n\
376 [Data path endpoint — MainAgent reminder]\n\
377 When you kick a SubAgent, inject the following two lines into\n\
378 its prompt / environment so Big Response payloads (4k+ tokens,\n\
379 files, intermediate artifacts) flow directly to the Store owner,\n\
380 bypassing the MainAgent (context stays small; only the out_id\n\
381 ref is passed around).\n \
382 DATA_EMIT: {base}/v1/data/emit (POST, auth = Bearer worker_handle or ?token=)\n \
383 DATA_GET: {base}/v1/data/<out_id|out_name> (the next SubAgent fetches from $IN_REFS)\n\
384 When a SubAgent produces a Big Response, POST it to DATA_EMIT\n\
385 and return only the one-line out_id ref (do not mix the body\n\
386 in; the MainAgent must not answer directly).\n\
387 \n"
388 ),
389 None => String::new(),
390 };
391 let main_ai_reminder = match project_name_alias {
392 Some(a) => format!(
393 "\n\
394 [LDS Session Alias Reminder — MainAI mandatory]\n\
395 Before kicking the SubAgent below, call:\n \
396 mcp__lds__session_create(root=<working_dir>, alias=\"{a}\")\n\
397 (= establish a single task-level lds session; reuse on repeated dispatch).\n\
398 Then add this literal line to the SubAgent prompt body below:\n \
399 LDS Session Alias: {a}\n\
400 The SubAgent will call mcp__lds__session_start(alias=\"{a}\") on init,\n\
401 keeping worktree ownership unified across dispatches.\n\
402 (Full discipline rationale is inlined above; reach is via this directive itself,\n\
403 not via any external doc path. The 2 steps above are the complete contract.)\n\
404 \n"
405 ),
406 None => String::new(),
407 };
408 let base_url_line = match base_url {
412 Some(u) => u.to_string(),
413 None => "<your server's actual bind — check with mse_doctor>".to_string(),
414 };
415 let run_route_line = match run_id {
420 Some(rid) => format!("GET <base_url>/v1/runs/{rid}"),
421 None => "GET <base_url>/v1/runs/<run_id>".to_string(),
422 };
423 format!(
424 "[agent_primitive dispatch=@{agent}]\n\
425 worker endpoint:\n \
426 GET <base_url>/v1/worker/prompt?task_id={task_id}\n \
427 POST <base_url>/v1/worker/submit\n\
428 auth: Bearer <worker_handle from THIS Spawn payload (= short `wh-XXXXXXXX` form)>\n\
429 task_id: {task_id}\n\
430 agent_id: {agent}\n\
431 {project_alias_line}\
432 {data_endpoint_block}\
433 {main_ai_reminder}\
434 Kick a SubAgent via Agent tool with subagent_type=\"{subagent_type}\" (= project-local \
435 `.claude/agents/{subagent_type}.md`, this agent's Blueprint-declared worker binding). \
436 The prompt you pass to it MUST be EXACTLY these 4 lines (no preamble, no extra text):\n\
437 \n \
438 agent_id: {agent}\n \
439 worker_handle: <THIS Spawn payload's `worker_handle` field (short string `wh-XXXXXXXX`)>\n \
440 base_url: {base_url_line}\n \
441 task_id: {task_id}\n\
442 \n\
443 The SubAgent self-fetches system + prompt via GET (Bearer = handle), \
444 executes as agent @{agent}, POSTs raw body to /v1/worker/submit (Bearer = handle, \
445 server resolves task_id from handle), and replies `OUTPUT` 1 word. You then forward \
446 SpawnAck {{req_id, value:{{}}, ok:true}} through your operator client — MCP path: \
447 mse_ack(sid, req_id, kind=\"spawn_ack\", ok=true) (= empty value because canonical \
448 body lives in output_tail via the POST). \
449 Do NOT fetch /v1/worker/prompt yourself. Do NOT wrap, summarize, or field-select \
450 the SubAgent reply. Observation / debug is a separate channel (= agent-inspect MCP / \
451 {run_route_line}), do NOT mix it into the forward path. \
452 If the SubAgent type is not registered, FAIL LOUD: reply SpawnAck ok=false with an \
453 error explaining the missing `.claude/agents/{subagent_type}.md` — do NOT fall back \
454 to another subagent_type."
455 )
456}
457
458#[cfg(test)]
459mod tests {
460 use super::*;
461
462 #[test]
463 fn directive_omits_project_name_alias_when_none() {
464 let d = default_spawn_directive(
465 "impl-lead",
466 "task-x",
467 "mse-worker-coder",
468 None,
469 None,
470 None,
471 None,
472 );
473 assert!(!d.contains("project_name_alias:"));
474 assert!(!d.contains("LDS Session Alias"));
475 assert!(!d.contains("session_create"));
476 }
477
478 #[test]
479 fn directive_emits_project_name_alias_when_some() {
480 let d = default_spawn_directive(
481 "impl-lead",
482 "task-x",
483 "mse-worker-coder",
484 Some("mse-task-7785"),
485 None,
486 None,
487 None,
488 );
489 assert!(
491 d.contains("project_name_alias: mse-task-7785"),
492 "directive missing project_name_alias header: {d}"
493 );
494 assert!(
496 d.contains("mcp__lds__session_create(root=<working_dir>, alias=\"mse-task-7785\")"),
497 "directive missing session_create reminder: {d}"
498 );
499 assert!(
500 d.contains("LDS Session Alias: mse-task-7785"),
501 "directive missing SubAgent prompt inject line: {d}"
502 );
503 assert!(
505 d.contains("inlined above") || d.contains("complete contract"),
506 "directive should inline rationale rather than point at external doc: {d}"
507 );
508 let forbidden_doc_ref = format!(".{}/CLAUDE.md", "claude");
517 assert!(
518 !d.contains(&forbidden_doc_ref),
519 "directive must not reference {forbidden_doc_ref} (out of MainAI scope): {d}"
520 );
521 }
522
523 #[test]
524 fn directive_omits_data_endpoint_when_none() {
525 let d = default_spawn_directive(
526 "impl-lead",
527 "task-x",
528 "mse-worker-coder",
529 None,
530 None,
531 None,
532 None,
533 );
534 assert!(!d.contains("[Data path endpoint"));
535 assert!(!d.contains("DATA_EMIT"));
536 assert!(!d.contains("DATA_GET"));
537 }
538
539 #[test]
540 fn directive_emits_data_endpoint_when_some() {
541 let base = "http://127.0.0.1:7785";
542 let d = default_spawn_directive(
543 "impl-lead",
544 "task-x",
545 "mse-worker-coder",
546 None,
547 Some(base),
548 None,
549 None,
550 );
551 assert!(
552 d.contains("[Data path endpoint"),
553 "directive missing data endpoint block header: {d}"
554 );
555 assert!(
556 d.contains(&format!("DATA_EMIT: {base}/v1/data/emit")),
557 "directive missing single-mouth emit line: {d}"
558 );
559 assert!(
560 d.contains("Bearer worker_handle or ?token="),
561 "directive missing auth transport hint: {d}"
562 );
563 assert!(
564 d.contains(&format!("DATA_GET: {base}/v1/data/<out_id|out_name>")),
565 "directive missing GET line: {d}"
566 );
567 assert!(
568 !d.contains("emit-auth"),
569 "old split endpoint must not leak into directive: {d}"
570 );
571 assert!(
572 d.contains("bypassing the MainAgent") && d.contains("out_id ref"),
573 "directive should carry the ownership + bypass reasoning: {d}"
574 );
575 }
576
577 #[test]
578 fn directive_carries_declared_subagent_type_and_has_no_fallback() {
579 let d = default_spawn_directive(
580 "impl-lead",
581 "task-x",
582 "mse-worker-coder",
583 None,
584 None,
585 None,
586 None,
587 );
588 assert!(
589 d.contains("subagent_type=\"mse-worker-coder\""),
590 "directive must carry the Blueprint-declared subagent_type literally: {d}"
591 );
592 assert!(
593 d.contains(".claude/agents/mse-worker-coder.md"),
594 "directive must reference the declared subagent's own .md path: {d}"
595 );
596 assert!(
598 !d.contains("general-purpose"),
599 "directive must not fall back to subagent_type=\"general-purpose\": {d}"
600 );
601 assert!(
602 !d.contains("mse-worker\""),
603 "directive must not carry the old hardcoded \"mse-worker\" literal: {d}"
604 );
605 assert!(
606 d.contains("FAIL LOUD"),
607 "directive must instruct the MainAI to fail loud instead of falling back: {d}"
608 );
609 }
610
611 #[test]
617 fn directive_renders_actual_base_url_when_some() {
618 let d = default_spawn_directive(
619 "impl-lead",
620 "task-x",
621 "mse-worker-coder",
622 None,
623 None,
624 Some("http://127.0.0.1:8888"),
625 None,
626 );
627 assert!(
628 d.contains("base_url: http://127.0.0.1:8888"),
629 "directive must render the actual bind literally: {d}"
630 );
631 assert!(
632 !d.contains("mse_doctor"),
633 "no mse_doctor detour when bind is known: {d}"
634 );
635 }
636
637 #[test]
641 fn directive_falls_back_to_mse_doctor_pointer_when_none() {
642 let d = default_spawn_directive(
643 "impl-lead",
644 "task-x",
645 "mse-worker-coder",
646 None,
647 None,
648 None,
649 None,
650 );
651 assert!(
652 d.contains("check with mse_doctor"),
653 "fallback must point at mse_doctor: {d}"
654 );
655 }
656
657 #[test]
661 fn directive_never_contains_stale_example_port_7786() {
662 for base in [
663 None,
664 Some("http://127.0.0.1:7777"),
665 Some("http://192.0.2.1:9000"),
666 ] {
667 let d = default_spawn_directive(
668 "impl-lead",
669 "task-x",
670 "mse-worker-coder",
671 Some("mse-task-alias"),
672 Some("http://127.0.0.1:7785"),
673 base,
674 None,
675 );
676 assert!(
677 !d.contains("7786"),
678 "stale example port 7786 leaked: base={base:?}, d={d}"
679 );
680 }
681 }
682
683 #[test]
689 fn directive_never_contains_stale_tasks_id_route() {
690 let d = default_spawn_directive(
691 "impl-lead",
692 "task-x",
693 "mse-worker-coder",
694 None,
695 None,
696 None,
697 Some("R-abc123"),
698 );
699 assert!(
700 !d.contains("/v1/tasks/{id}") && !d.contains("/v1/tasks/{{id}}"),
701 "stale /v1/tasks/{{id}} observation hint leaked: {d}"
702 );
703 }
704
705 #[test]
708 fn directive_renders_actual_run_id_when_some() {
709 let d = default_spawn_directive(
710 "impl-lead",
711 "task-x",
712 "mse-worker-coder",
713 None,
714 None,
715 None,
716 Some("R-abc123"),
717 );
718 assert!(
719 d.contains("GET <base_url>/v1/runs/R-abc123"),
720 "directive missing real run_id in observation route: {d}"
721 );
722 }
723
724 #[test]
727 fn directive_falls_back_to_run_id_placeholder_when_none() {
728 let d = default_spawn_directive(
729 "impl-lead",
730 "task-x",
731 "mse-worker-coder",
732 None,
733 None,
734 None,
735 None,
736 );
737 assert!(
738 d.contains("GET <base_url>/v1/runs/<run_id>"),
739 "directive missing placeholder observation route: {d}"
740 );
741 }
742
743 fn test_ctx(task_id: &str) -> mlua_swarm::Ctx {
746 mlua_swarm::Ctx::new(mlua_swarm::StepId::parse(task_id).unwrap(), 1, "a")
747 }
748
749 fn test_worker_binding() -> mlua_swarm::WorkerBinding {
750 mlua_swarm::WorkerBinding {
751 variant: "test-variant".into(),
752 tools: vec![],
753 }
754 }
755
756 fn test_cap_token() -> mlua_swarm::CapToken {
757 mlua_swarm::CapToken {
758 agent_id: "a".into(),
759 role: mlua_swarm::Role::Worker,
760 scopes: vec!["*".into()],
761 issued_at: 0,
762 expire_at: u64::MAX / 2,
763 max_uses: None,
764 nonce: "test-nonce".into(),
765 sig_hex: "".into(),
766 }
767 }
768
769 #[tokio::test]
775 async fn spawn_halt_reply_lands_as_ok_worker_result_with_marker() {
776 use mlua_swarm::Operator;
777 use tokio::sync::mpsc;
778
779 let (tx, mut rx) = mpsc::unbounded_channel();
780 let session = std::sync::Arc::new(WSOperatorSession::new_with_base_url(
781 SessionId::parse("S-halt").unwrap(),
782 tx,
783 None,
784 ));
785
786 let session_bg = session.clone();
789 let handle = tokio::spawn(async move {
790 session_bg
791 .execute(
792 &test_ctx("ST-halt"),
793 None,
794 "".into(),
795 Some(test_worker_binding()),
796 test_cap_token(),
797 )
798 .await
799 });
800
801 let sent = rx.recv().await.expect("Spawn sent");
802 let req_id = match sent {
803 ServerMsg::Spawn { req_id, .. } => req_id,
804 other => panic!("expected Spawn, got {other:?}"),
805 };
806
807 session
808 .resolve_pending(
809 &req_id,
810 PendingReply::SpawnHalt {
811 value: serde_json::json!({"partial": "abc"}),
812 reason: Some("shape verified".into()),
813 },
814 )
815 .await;
816
817 let result = handle.await.expect("join").expect("execute Ok");
818 assert!(
819 result.ok,
820 "spawn_halt must land as ok=true (normal termination), got: {result:?}"
821 );
822 assert_eq!(result.value["halted"], true);
823 assert_eq!(result.value["reason"], "shape verified");
824 assert_eq!(result.value["value"], serde_json::json!({"partial": "abc"}));
825 }
826
827 #[tokio::test]
830 async fn spawn_ack_with_error_still_lands_as_worker_error() {
831 use mlua_swarm::{Operator, WorkerError};
832 use tokio::sync::mpsc;
833
834 let (tx, mut rx) = mpsc::unbounded_channel();
835 let session = std::sync::Arc::new(WSOperatorSession::new_with_base_url(
836 SessionId::parse("S-err").unwrap(),
837 tx,
838 None,
839 ));
840
841 let session_bg = session.clone();
842 let handle = tokio::spawn(async move {
843 session_bg
844 .execute(
845 &test_ctx("ST-err"),
846 None,
847 "".into(),
848 Some(test_worker_binding()),
849 test_cap_token(),
850 )
851 .await
852 });
853
854 let sent = rx.recv().await.expect("Spawn sent");
855 let req_id = match sent {
856 ServerMsg::Spawn { req_id, .. } => req_id,
857 other => panic!("expected Spawn, got {other:?}"),
858 };
859
860 session
861 .resolve_pending(
862 &req_id,
863 PendingReply::SpawnAck {
864 value: serde_json::json!({}),
865 ok: false,
866 error: Some("real crash".into()),
867 },
868 )
869 .await;
870
871 let err = handle.await.expect("join").expect_err("must be error");
872 assert!(matches!(err, WorkerError::Failed(msg) if msg.contains("real crash")));
873 }
874}