1use crate::core::agent_context::AgentContextView;
54use crate::core::ctx::Ctx;
55use crate::core::engine::Engine;
56use crate::types::{CapToken, StepId, WorkerId};
57use crate::worker::adapter::{SpawnError, SpawnerAdapter, WorkerError, WorkerResult};
58use crate::worker::output::{ContentRef, OutputEvent};
59use crate::worker::{Worker, WorkerJoinHandler};
60use async_trait::async_trait;
61use mlua_swarm_schema::SubprocessOutput;
62use serde_json::Value;
63use std::collections::BTreeMap;
64use std::process::Stdio;
65use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
66use tokio::process::Command;
67use tokio::sync::oneshot;
68use tokio_util::sync::CancellationToken;
69
70pub const EMBED_PLACEHOLDERS: [&str; 8] = [
77 "system",
78 "system_file",
79 "prompt",
80 "model",
81 "tools_csv",
82 "work_dir",
83 "task_id",
84 "attempt",
85];
86
87#[derive(Debug, Clone, Default)]
98pub struct EmbedTemplate {
99 pub argv: Vec<String>,
102 pub stdin: Option<String>,
105 pub env: BTreeMap<String, String>,
108 pub cwd: Option<String>,
114 pub output: Option<SubprocessOutput>,
116 pub system_prompt: Option<String>,
120 pub model: Option<String>,
122 pub tools_csv: String,
125 pub timeout: Option<std::time::Duration>,
128}
129
130#[derive(Debug, Clone)]
134pub enum StreamMode {
135 NdjsonLines,
137 SseEvents,
143 LengthPrefixed,
147}
148
149#[derive(Debug)]
153pub struct ProcessSpawner {
154 pub program: String,
157 pub args: Vec<String>,
159 pub use_stdin: bool,
163 pub stream_mode: Option<StreamMode>,
165 pub embed: Option<EmbedTemplate>,
169}
170
171impl ProcessSpawner {
172 pub fn new(program: impl Into<String>) -> Self {
175 Self {
176 program: program.into(),
177 args: Vec::new(),
178 use_stdin: true,
179 stream_mode: None,
180 embed: None,
181 }
182 }
183
184 pub fn arg(mut self, a: impl Into<String>) -> Self {
186 self.args.push(a.into());
187 self
188 }
189
190 pub fn args(mut self, args: impl IntoIterator<Item = impl Into<String>>) -> Self {
192 self.args.extend(args.into_iter().map(|a| a.into()));
193 self
194 }
195
196 pub fn use_stdin(mut self, v: bool) -> Self {
199 self.use_stdin = v;
200 self
201 }
202
203 pub fn stream_mode(mut self, mode: StreamMode) -> Self {
205 self.stream_mode = Some(mode);
206 self
207 }
208
209 pub fn plain(mut self) -> Self {
211 self.stream_mode = None;
212 self
213 }
214
215 pub fn ndjson(mut self, v: bool) -> Self {
219 self.stream_mode = if v {
220 Some(StreamMode::NdjsonLines)
221 } else {
222 None
223 };
224 self
225 }
226
227 pub fn run(cmd: impl Into<String>) -> Self {
229 Self {
230 program: "sh".into(),
231 args: vec!["-c".into(), cmd.into()],
232 use_stdin: true,
233 stream_mode: None,
234 embed: None,
235 }
236 }
237
238 pub fn cmd(program: impl Into<String>) -> Self {
241 Self {
242 program: program.into(),
243 args: Vec::new(),
244 use_stdin: true,
245 stream_mode: None,
246 embed: None,
247 }
248 }
249
250 pub fn embed(mut self, template: EmbedTemplate) -> Self {
252 self.embed = Some(template);
253 self
254 }
255}
256
257#[async_trait]
258impl SpawnerAdapter for ProcessSpawner {
259 async fn spawn(
260 &self,
261 engine: &Engine,
262 ctx: &Ctx,
263 task_id: StepId,
264 attempt: u32,
265 token: CapToken,
266 ) -> Result<Box<dyn Worker>, SpawnError> {
267 if let Some(embed) = &self.embed {
271 return self
272 .spawn_embed(embed, engine, ctx, task_id, attempt, token)
273 .await;
274 }
275 let directive = engine
283 .fetch_prompt(&token, &task_id)
284 .await
285 .map_err(|e| SpawnError::Internal(format!("fetch_prompt: {e}")))?;
286 let directive = crate::core::engine::render_directive_to_string(&directive);
289
290 let mut cmd = Command::new(&self.program);
291 cmd.args(&self.args)
292 .env("MSE_TOKEN_AGENT_ID", &token.agent_id)
293 .env("MSE_TOKEN_NONCE", &token.nonce)
294 .env("MSE_TASK_ID", task_id.as_str())
295 .env("MSE_ATTEMPT", attempt.to_string())
296 .env("MSE_CTX_AGENT", &ctx.agent)
297 .stdin(Stdio::piped())
298 .stdout(Stdio::piped())
299 .stderr(Stdio::piped());
300
301 if !self.use_stdin {
302 cmd.arg(&directive);
303 }
304
305 let mut child = cmd
306 .spawn()
307 .map_err(|e| SpawnError::Internal(format!("spawn failed: {e}")))?;
308
309 if self.use_stdin {
310 if let Some(mut stdin) = child.stdin.take() {
311 stdin
312 .write_all(directive.as_bytes())
313 .await
314 .map_err(|e| SpawnError::Internal(format!("stdin write: {e}")))?;
315 drop(stdin); }
317 }
318
319 let cancel = CancellationToken::new();
320 let cancel_inner = cancel.clone();
321 let worker_id = WorkerId::new();
322 tracing::debug!(worker_id = %worker_id, step_id = %task_id, "worker spawned (subprocess)");
324 let (tx, rx) = oneshot::channel();
325 let engine_for_emit = engine.clone();
329 let token_for_emit = token.clone();
330 let task_id_for_emit = task_id.clone();
331 let stream_mode = self.stream_mode.clone();
332
333 tokio::spawn(async move {
334 let result: Result<WorkerResult, WorkerError> = if let Some(mode) = stream_mode {
335 run_streaming_mode(
339 mode,
340 child,
341 &engine_for_emit,
342 &token_for_emit,
343 &task_id_for_emit,
344 attempt,
345 cancel_inner,
346 )
347 .await
348 } else {
349 let result = tokio::select! {
352 output = child.wait_with_output() => {
353 match output {
354 Ok(out) => {
355 let stdout = String::from_utf8_lossy(&out.stdout).to_string();
356 let value: Value = serde_json::from_str(stdout.trim())
357 .unwrap_or_else(|_| serde_json::json!({
358 "raw": stdout.trim_end(),
359 "stderr": String::from_utf8_lossy(&out.stderr).to_string(),
360 }));
361 Ok(WorkerResult {
362 value,
363 ok: out.status.success(),
364 stats: Some(subprocess_base_stats(&out.status, None)),
365 })
366 }
367 Err(e) => Err(WorkerError::Failed(format!("wait_with_output: {e}"))),
368 }
369 }
370 _ = cancel_inner.cancelled() => Err(WorkerError::Cancelled),
371 };
372 if let Ok(wr) = &result {
373 if let Some(stats) = wr.stats.clone() {
374 engine_for_emit
375 .record_worker_stats(&task_id_for_emit, attempt, stats)
376 .await;
377 }
378 let ev = OutputEvent::Final {
379 content: ContentRef::Inline {
380 value: wr.value.clone(),
381 },
382 ok: wr.ok,
383 };
384 let _ = engine_for_emit
385 .submit_output(&token_for_emit, &task_id_for_emit, attempt, ev)
386 .await;
387 }
388 result
389 };
390 let signal: Result<(), WorkerError> = result.map(|_| ());
392 let _ = tx.send(signal);
393 });
394
395 Ok(Box::new(ProcessWorker {
396 handler: WorkerJoinHandler {
397 worker_id,
398 cancel,
399 completion: rx,
400 },
401 }))
402 }
403}
404
405struct EmbedVars {
411 system: String,
412 system_file: Option<String>,
413 prompt: String,
414 model: Option<String>,
415 tools_csv: String,
416 work_dir: Option<String>,
417 task_id: String,
418 attempt: String,
419}
420
421impl EmbedVars {
422 fn lookup(&self, token: &str) -> Result<Option<&str>, SpawnError> {
427 let (value, cause): (Option<&str>, &str) = match token {
428 "system" => (Some(self.system.as_str()), ""),
429 "system_file" => (
430 self.system_file.as_deref(),
431 "no system prompt was baked for this attempt (agent has no profile.system_prompt?)",
432 ),
433 "prompt" => (Some(self.prompt.as_str()), ""),
434 "model" => (
435 self.model.as_deref(),
436 "no model declared (set profile.model or Runner::Subprocess overrides.model)",
437 ),
438 "tools_csv" => (Some(self.tools_csv.as_str()), ""),
439 "work_dir" => (
440 self.work_dir.as_deref(),
441 "no work_dir/project_root in the agent context view and no overrides.cwd",
442 ),
443 "task_id" => (Some(self.task_id.as_str()), ""),
444 "attempt" => (Some(self.attempt.as_str()), ""),
445 _ => return Ok(None),
446 };
447 match value {
448 Some(v) => Ok(Some(v)),
449 None => Err(SpawnError::Internal(format!(
450 "placeholder {{{token}}}: {cause}"
451 ))),
452 }
453 }
454
455 fn render(&self, tmpl: &str) -> Result<String, SpawnError> {
464 let mut out = String::with_capacity(tmpl.len());
465 let mut rest = tmpl;
466 while let Some(start) = rest.find('{') {
467 out.push_str(&rest[..start]);
468 let after = &rest[start + 1..];
469 let Some(end) = after.find('}') else {
470 out.push_str(&rest[start..]);
472 return Ok(out);
473 };
474 let token = &after[..end];
475 match self.lookup(token)? {
476 Some(value) => {
477 out.push_str(value);
478 rest = &after[end + 1..];
479 }
480 None => {
481 out.push('{');
486 rest = after;
487 }
488 }
489 }
490 out.push_str(rest);
491 Ok(out)
492 }
493}
494
495fn embed_references(embed: &EmbedTemplate, token: &str) -> bool {
497 embed.argv.iter().any(|a| a.contains(token))
498 || embed.stdin.as_deref().is_some_and(|s| s.contains(token))
499 || embed.env.values().any(|v| v.contains(token))
500 || embed.cwd.as_deref().is_some_and(|c| c.contains(token))
501}
502
503fn subprocess_base_stats(
509 status: &std::process::ExitStatus,
510 model: Option<&str>,
511) -> crate::store::trace::WorkerStats {
512 crate::store::trace::WorkerStats {
513 worker_kind: Some("subprocess".to_string()),
514 model: model.map(str::to_string),
515 usage: None,
516 num_turns: None,
517 adapter_data: Some(serde_json::json!({ "exit_code": status.code() })),
518 }
519}
520
521fn timed_out_result(limit: std::time::Duration, model: Option<&str>) -> WorkerResult {
525 let timeout_ms = u64::try_from(limit.as_millis()).unwrap_or(u64::MAX);
526 WorkerResult {
527 value: serde_json::json!({
528 "error": format!("subprocess timed out after {timeout_ms} ms and was killed"),
529 "timed_out": true,
530 "timeout_ms": timeout_ms,
531 }),
532 ok: false,
533 stats: Some(crate::store::trace::WorkerStats {
534 worker_kind: Some("subprocess".to_string()),
535 model: model.map(str::to_string),
536 usage: None,
537 num_turns: None,
538 adapter_data: Some(serde_json::json!({ "exit_code": null, "timed_out": true })),
539 }),
540 }
541}
542
543fn enrich_declared_stats(
548 stats: &mut crate::store::trace::WorkerStats,
549 parsed: &Value,
550 decl: &SubprocessOutput,
551) {
552 let Some(sd) = &decl.stats else { return };
553 if let Some(ptr) = sd.model_ptr.as_deref() {
554 if let Some(m) = parsed.pointer(ptr).and_then(|v| v.as_str()) {
555 stats.model = Some(m.to_string());
556 }
557 }
558 if let Some(ptr) = sd.num_turns_ptr.as_deref() {
559 if let Some(n) = parsed.pointer(ptr).and_then(|v| v.as_u64()) {
560 stats.num_turns = Some(n as u32);
561 }
562 }
563 if let Some(ptr) = sd.usage_ptr.as_deref() {
564 if let Some(u) = parsed.pointer(ptr) {
565 let input = u
566 .get("input_tokens")
567 .or_else(|| u.get("prompt_tokens"))
568 .and_then(|v| v.as_u64());
569 let output = u
570 .get("output_tokens")
571 .or_else(|| u.get("completion_tokens"))
572 .and_then(|v| v.as_u64());
573 let total = u.get("total_tokens").and_then(|v| v.as_u64());
574 if let Some(usage) = crate::store::trace::TokenUsage::from_parts(input, output, total) {
578 stats.usage = Some(usage);
579 if let Some(Value::Object(ad)) = stats.adapter_data.as_mut() {
582 ad.insert("usage_raw".to_string(), u.clone());
583 }
584 }
585 }
586 }
587}
588
589fn normalize_plain_output(
596 out: &std::process::Output,
597 decl: Option<&SubprocessOutput>,
598 model: Option<&str>,
599) -> WorkerResult {
600 let stdout = String::from_utf8_lossy(&out.stdout).to_string();
601 let stderr = || String::from_utf8_lossy(&out.stderr).to_string();
602 let exit_ok = out.status.success();
603 let mut stats = subprocess_base_stats(&out.status, model);
604
605 let Some(decl) = decl else {
606 let value: Value = serde_json::from_str(stdout.trim()).unwrap_or_else(|_| {
607 serde_json::json!({
608 "raw": stdout.trim_end(),
609 "stderr": stderr(),
610 })
611 });
612 return WorkerResult {
613 value,
614 ok: exit_ok,
615 stats: Some(stats),
616 };
617 };
618
619 let parsed: Result<Value, _> = serde_json::from_str(stdout.trim());
620 let parsed = match parsed {
621 Ok(v) => v,
622 Err(e) => {
623 if decl.format.as_deref() == Some("json") {
624 return WorkerResult {
627 value: serde_json::json!({
628 "raw": stdout.trim_end(),
629 "stderr": stderr(),
630 "parse_error": e.to_string(),
631 }),
632 ok: false,
633 stats: Some(stats),
634 };
635 }
636 return WorkerResult {
639 value: serde_json::json!({
640 "raw": stdout.trim_end(),
641 "stderr": stderr(),
642 }),
643 ok: exit_ok,
644 stats: Some(stats),
645 };
646 }
647 };
648
649 enrich_declared_stats(&mut stats, &parsed, decl);
650
651 let value = match decl.result_ptr.as_deref() {
652 Some(ptr) => match parsed.pointer(ptr) {
653 Some(v) => v.clone(),
654 None => {
655 return WorkerResult {
656 value: serde_json::json!({
657 "error": format!("result_ptr '{ptr}' not found in stdout JSON"),
658 "raw": parsed,
659 "stderr": stderr(),
660 }),
661 ok: false,
662 stats: Some(stats),
663 }
664 }
665 },
666 None => parsed.clone(),
667 };
668
669 let ok = match decl.ok_from.as_deref() {
670 None | Some("exit_code") => exit_ok,
671 Some(ptr) => parsed.pointer(ptr) == Some(&Value::Bool(true)),
672 };
673
674 WorkerResult {
675 value,
676 ok,
677 stats: Some(stats),
678 }
679}
680
681impl ProcessSpawner {
682 async fn spawn_embed(
688 &self,
689 embed: &EmbedTemplate,
690 engine: &Engine,
691 ctx: &Ctx,
692 task_id: StepId,
693 attempt: u32,
694 token: CapToken,
695 ) -> Result<Box<dyn Worker>, SpawnError> {
696 let prompt_value = engine
700 .fetch_prompt(&token, &task_id)
701 .await
702 .map_err(|e| SpawnError::Internal(format!("fetch_prompt: {e}")))?;
703 let system = match embed.system_prompt.as_deref() {
704 Some(tmpl) => {
705 let slots = crate::operator::render::slots_from_prompt(&prompt_value);
706 let rendered = crate::operator::render::render_system(tmpl, &slots)
707 .map_err(|e| SpawnError::Internal(format!("render system_prompt: {e}")))?;
708 Some(rendered)
709 }
710 None => None,
711 };
712 engine
713 .bake_worker_system_prompt(&task_id, attempt, system.clone())
714 .await
715 .map_err(|e| SpawnError::Internal(format!("bake system_prompt: {e}")))?;
716
717 let payload = engine
723 .fetch_worker_payload(&token, &task_id)
724 .await
725 .map_err(|e| SpawnError::Internal(format!("fetch_worker_payload: {e}")))?;
726
727 let system_file = if embed_references(embed, "{system_file}") {
730 match engine
731 .materialize_system_file(&task_id, attempt)
732 .await
733 .map_err(|e| SpawnError::Internal(format!("materialize system_file: {e}")))?
734 {
735 Some(path) => Some(path.display().to_string()),
736 None => {
737 return Err(SpawnError::Internal(
738 "placeholder {system_file}: no system prompt was baked for this \
739 attempt (agent has no profile.system_prompt?)"
740 .into(),
741 ))
742 }
743 }
744 } else {
745 None
746 };
747
748 let view = AgentContextView::materialized_or_from_ctx(ctx);
751 let work_dir = view.work_dir.clone().or_else(|| view.project_root.clone());
752
753 let vars = EmbedVars {
754 system: system.unwrap_or_default(),
755 system_file,
756 prompt: payload.prompt.clone(),
757 model: embed.model.clone(),
758 tools_csv: embed.tools_csv.clone(),
759 work_dir,
760 task_id: task_id.to_string(),
761 attempt: attempt.to_string(),
762 };
763
764 if embed.argv.is_empty() {
766 return Err(SpawnError::Internal(
767 "embed template: argv must not be empty".into(),
768 ));
769 }
770 let mut rendered_argv = Vec::with_capacity(embed.argv.len());
771 for a in &embed.argv {
772 rendered_argv.push(vars.render(a)?);
773 }
774 let mut rendered_env = Vec::with_capacity(embed.env.len());
775 for (k, v) in &embed.env {
776 rendered_env.push((k.clone(), vars.render(v)?));
777 }
778 let rendered_cwd = match embed.cwd.as_deref() {
779 Some(c) => Some(vars.render(c)?),
780 None => None,
781 };
782 let rendered_stdin = match embed.stdin.as_deref() {
783 Some(s) => Some(vars.render(s)?),
784 None => None,
785 };
786
787 let mut cmd = Command::new(&rendered_argv[0]);
788 cmd.args(&rendered_argv[1..])
789 .env("MSE_TOKEN_AGENT_ID", &token.agent_id)
790 .env("MSE_TOKEN_NONCE", &token.nonce)
791 .env("MSE_TASK_ID", task_id.as_str())
792 .env("MSE_ATTEMPT", attempt.to_string())
793 .env("MSE_CTX_AGENT", &ctx.agent)
794 .stdin(Stdio::piped())
795 .stdout(Stdio::piped())
796 .stderr(Stdio::piped())
797 .kill_on_drop(true);
800 for (k, v) in &rendered_env {
801 cmd.env(k, v);
802 }
803 if let Some(cwd) = &rendered_cwd {
804 cmd.current_dir(cwd);
805 }
806
807 let mut child = cmd
808 .spawn()
809 .map_err(|e| SpawnError::Internal(format!("spawn failed: {e}")))?;
810
811 if let Some(stdin_body) = rendered_stdin {
812 if let Some(mut stdin) = child.stdin.take() {
813 stdin
814 .write_all(stdin_body.as_bytes())
815 .await
816 .map_err(|e| SpawnError::Internal(format!("stdin write: {e}")))?;
817 drop(stdin); }
819 } else {
820 drop(child.stdin.take());
822 }
823
824 let cancel = CancellationToken::new();
825 let cancel_inner = cancel.clone();
826 let worker_id = WorkerId::new();
827 tracing::debug!(worker_id = %worker_id, step_id = %task_id, "worker spawned (subprocess embed)");
828 let (tx, rx) = oneshot::channel();
829 let engine_for_emit = engine.clone();
830 let token_for_emit = token.clone();
831 let task_id_for_emit = task_id.clone();
832 let stream_mode = self.stream_mode.clone();
833 let output_decl = embed.output.clone();
834 let model_for_stats = embed.model.clone();
835 let timeout = embed.timeout;
836
837 tokio::spawn(async move {
838 let result: Result<WorkerResult, WorkerError> = if let Some(mode) = stream_mode {
839 let streaming = run_streaming_mode(
842 mode,
843 child,
844 &engine_for_emit,
845 &token_for_emit,
846 &task_id_for_emit,
847 attempt,
848 cancel_inner,
849 );
850 match timeout {
851 None => streaming.await,
852 Some(limit) => match tokio::time::timeout(limit, streaming).await {
853 Ok(result) => result,
854 Err(_) => {
855 let wr = timed_out_result(limit, model_for_stats.as_deref());
858 if let Some(stats) = wr.stats.clone() {
859 engine_for_emit
860 .record_worker_stats(&task_id_for_emit, attempt, stats)
861 .await;
862 }
863 let ev = OutputEvent::Final {
864 content: ContentRef::Inline {
865 value: wr.value.clone(),
866 },
867 ok: false,
868 };
869 let _ = engine_for_emit
870 .submit_output(&token_for_emit, &task_id_for_emit, attempt, ev)
871 .await;
872 Ok(wr)
873 }
874 },
875 }
876 } else {
877 let deadline = async {
878 match timeout {
879 Some(limit) => tokio::time::sleep(limit).await,
880 None => std::future::pending::<()>().await,
881 }
882 };
883 let result = tokio::select! {
884 output = child.wait_with_output() => {
885 match output {
886 Ok(out) => Ok(normalize_plain_output(
887 &out,
888 output_decl.as_ref(),
889 model_for_stats.as_deref(),
890 )),
891 Err(e) => Err(WorkerError::Failed(format!("wait_with_output: {e}"))),
892 }
893 }
894 _ = cancel_inner.cancelled() => Err(WorkerError::Cancelled),
895 _ = deadline => Ok(timed_out_result(
896 timeout.unwrap_or_default(),
897 model_for_stats.as_deref(),
898 )),
899 };
900 if let Ok(wr) = &result {
901 if let Some(stats) = wr.stats.clone() {
902 engine_for_emit
903 .record_worker_stats(&task_id_for_emit, attempt, stats)
904 .await;
905 }
906 let ev = OutputEvent::Final {
907 content: ContentRef::Inline {
908 value: wr.value.clone(),
909 },
910 ok: wr.ok,
911 };
912 let _ = engine_for_emit
913 .submit_output(&token_for_emit, &task_id_for_emit, attempt, ev)
914 .await;
915 }
916 result
917 };
918 let signal: Result<(), WorkerError> = result.map(|_| ());
919 let _ = tx.send(signal);
920 });
921
922 Ok(Box::new(ProcessWorker {
923 handler: WorkerJoinHandler {
924 worker_id,
925 cancel,
926 completion: rx,
927 },
928 }))
929 }
930}
931
932pub struct ProcessWorker {
936 pub handler: WorkerJoinHandler,
939}
940
941#[async_trait]
942impl Worker for ProcessWorker {
943 fn id(&self) -> &WorkerId {
944 &self.handler.worker_id
945 }
946 fn cancel_token(&self) -> CancellationToken {
947 self.handler.cancel.clone()
948 }
949 async fn join(self: Box<Self>) -> Result<(), WorkerError> {
950 self.handler.await_completion().await
951 }
952}
953
954async fn run_streaming_mode(
960 mode: StreamMode,
961 mut child: tokio::process::Child,
962 engine: &Engine,
963 token: &CapToken,
964 task_id: &StepId,
965 attempt: u32,
966 cancel: CancellationToken,
967) -> Result<WorkerResult, WorkerError> {
968 let stdout = child
969 .stdout
970 .take()
971 .ok_or_else(|| WorkerError::Failed("streaming: stdout pipe missing".into()))?;
972
973 let last_final = match mode {
974 StreamMode::NdjsonLines => {
975 read_ndjson(stdout, engine, token, task_id, attempt, cancel.clone()).await?
976 }
977 StreamMode::SseEvents => {
978 read_sse(stdout, engine, token, task_id, attempt, cancel.clone()).await?
979 }
980 StreamMode::LengthPrefixed => {
981 read_length_prefixed(stdout, engine, token, task_id, attempt, cancel.clone()).await?
982 }
983 };
984
985 let status = child
986 .wait()
987 .await
988 .map_err(|e| WorkerError::Failed(format!("streaming wait: {e}")))?;
989
990 match last_final {
991 Some((value, ok)) => Ok(WorkerResult {
992 value,
993 ok: ok && status.success(),
994 stats: Some(subprocess_base_stats(&status, None)),
997 }),
998 None => {
999 let value = serde_json::json!({
1001 "raw": "",
1002 "note": "streaming mode: no Final event received",
1003 "exit_success": status.success(),
1004 });
1005 let _ = engine
1006 .submit_output(
1007 token,
1008 task_id,
1009 attempt,
1010 OutputEvent::Final {
1011 content: ContentRef::Inline {
1012 value: value.clone(),
1013 },
1014 ok: false,
1015 },
1016 )
1017 .await;
1018 Ok(WorkerResult {
1019 value,
1020 ok: false,
1021 stats: Some(subprocess_base_stats(&status, None)),
1022 })
1023 }
1024 }
1025}
1026
1027async fn forward_event(
1030 engine: &Engine,
1031 token: &CapToken,
1032 task_id: &StepId,
1033 attempt: u32,
1034 ev: OutputEvent,
1035 last_final: &mut Option<(Value, bool)>,
1036) {
1037 if let OutputEvent::Final { content, ok } = &ev {
1038 let value = match content {
1039 ContentRef::Inline { value } => value.clone(),
1040 ContentRef::FileRef {
1041 path,
1042 mime,
1043 size_hint,
1044 } => serde_json::json!({
1045 "file_ref": path.to_string_lossy(),
1046 "mime": mime,
1047 "size_hint": size_hint,
1048 }),
1049 };
1050 *last_final = Some((value, *ok));
1051 }
1052 let _ = engine.submit_output(token, task_id, attempt, ev).await;
1053}
1054
1055async fn read_ndjson(
1058 stdout: tokio::process::ChildStdout,
1059 engine: &Engine,
1060 token: &CapToken,
1061 task_id: &StepId,
1062 attempt: u32,
1063 cancel: CancellationToken,
1064) -> Result<Option<(Value, bool)>, WorkerError> {
1065 let mut reader = BufReader::new(stdout).lines();
1066 let mut last_final = None;
1067 loop {
1068 tokio::select! {
1069 line_res = reader.next_line() => match line_res {
1070 Ok(Some(line)) => {
1071 let trimmed = line.trim();
1072 if trimmed.is_empty() { continue; }
1073 if let Ok(ev) = serde_json::from_str::<OutputEvent>(trimmed) {
1074 forward_event(engine, token, task_id, attempt, ev, &mut last_final).await;
1075 }
1076 }
1077 Ok(None) => break,
1078 Err(e) => return Err(WorkerError::Failed(format!("ndjson read: {e}"))),
1079 },
1080 _ = cancel.cancelled() => return Err(WorkerError::Cancelled),
1081 }
1082 }
1083 Ok(last_final)
1084}
1085
1086async fn read_sse(
1090 stdout: tokio::process::ChildStdout,
1091 engine: &Engine,
1092 token: &CapToken,
1093 task_id: &StepId,
1094 attempt: u32,
1095 cancel: CancellationToken,
1096) -> Result<Option<(Value, bool)>, WorkerError> {
1097 let mut reader = BufReader::new(stdout).lines();
1098 let mut last_final = None;
1099 let mut data_buf = String::new();
1100 loop {
1101 tokio::select! {
1102 line_res = reader.next_line() => match line_res {
1103 Ok(Some(line)) => {
1104 if line.is_empty() {
1105 if !data_buf.is_empty() {
1107 if let Ok(ev) = serde_json::from_str::<OutputEvent>(data_buf.trim()) {
1108 forward_event(engine, token, task_id, attempt, ev, &mut last_final).await;
1109 }
1110 data_buf.clear();
1111 }
1112 } else if let Some(rest) = line.strip_prefix("data:") {
1113 let payload = rest.strip_prefix(' ').unwrap_or(rest);
1115 if !data_buf.is_empty() {
1116 data_buf.push('\n');
1117 }
1118 data_buf.push_str(payload);
1119 }
1120 }
1122 Ok(None) => {
1123 if !data_buf.is_empty() {
1125 if let Ok(ev) = serde_json::from_str::<OutputEvent>(data_buf.trim()) {
1126 forward_event(engine, token, task_id, attempt, ev, &mut last_final).await;
1127 }
1128 }
1129 break;
1130 }
1131 Err(e) => return Err(WorkerError::Failed(format!("sse read: {e}"))),
1132 },
1133 _ = cancel.cancelled() => return Err(WorkerError::Cancelled),
1134 }
1135 }
1136 Ok(last_final)
1137}
1138
1139async fn read_length_prefixed(
1142 mut stdout: tokio::process::ChildStdout,
1143 engine: &Engine,
1144 token: &CapToken,
1145 task_id: &StepId,
1146 attempt: u32,
1147 cancel: CancellationToken,
1148) -> Result<Option<(Value, bool)>, WorkerError> {
1149 use tokio::io::AsyncReadExt;
1150 let mut last_final = None;
1151 loop {
1152 let mut len_buf = [0u8; 4];
1154 let read_fut = stdout.read_exact(&mut len_buf);
1155 let read_res = tokio::select! {
1156 r = read_fut => r,
1157 _ = cancel.cancelled() => return Err(WorkerError::Cancelled),
1158 };
1159 match read_res {
1160 Ok(_) => {}
1161 Err(e) if e.kind() == std::io::ErrorKind::UnexpectedEof => break, Err(e) => return Err(WorkerError::Failed(format!("len read: {e}"))),
1163 }
1164 let len = u32::from_be_bytes(len_buf) as usize;
1165 if len == 0 || len > 16 * 1024 * 1024 {
1166 break;
1168 }
1169 let mut payload = vec![0u8; len];
1170 let read_fut = stdout.read_exact(&mut payload);
1171 let read_res = tokio::select! {
1172 r = read_fut => r,
1173 _ = cancel.cancelled() => return Err(WorkerError::Cancelled),
1174 };
1175 if read_res.is_err() {
1176 break;
1177 }
1178 if let Ok(ev) = serde_json::from_slice::<OutputEvent>(&payload) {
1179 forward_event(engine, token, task_id, attempt, ev, &mut last_final).await;
1180 }
1181 }
1182 Ok(last_final)
1183}
1184
1185#[cfg(test)]
1186mod embed_tests {
1187 use super::*;
1188
1189 fn vars() -> EmbedVars {
1190 EmbedVars {
1191 system: "SYS".to_string(),
1192 system_file: Some("/tmp/sys.md".to_string()),
1193 prompt: "do the task".to_string(),
1194 model: Some("small".to_string()),
1195 tools_csv: "Read,Grep".to_string(),
1196 work_dir: Some("/tmp/wd".to_string()),
1197 task_id: "ST-1".to_string(),
1198 attempt: "1".to_string(),
1199 }
1200 }
1201
1202 #[test]
1203 fn render_substitutes_every_closed_set_token() {
1204 let v = vars();
1205 let out = v
1206 .render("{system}|{system_file}|{prompt}|{model}|{tools_csv}|{work_dir}|{task_id}|{attempt}")
1207 .expect("render");
1208 assert_eq!(
1209 out,
1210 "SYS|/tmp/sys.md|do the task|small|Read,Grep|/tmp/wd|ST-1|1"
1211 );
1212 }
1213
1214 #[test]
1215 fn render_leaves_non_placeholder_braces_alone() {
1216 let v = vars();
1217 let out = v
1218 .render(r#"echo '{"result": "{prompt}"}'"#)
1219 .expect("render");
1220 assert_eq!(out, r#"echo '{"result": "do the task"}'"#);
1221 }
1222
1223 #[test]
1224 fn render_fails_loud_when_model_referenced_but_absent() {
1225 let mut v = vars();
1226 v.model = None;
1227 let err = v.render("--model {model}").unwrap_err();
1228 let msg = format!("{err:?}");
1229 assert!(msg.contains("{model}"), "actionable token in error: {msg}");
1230 assert!(msg.contains("profile.model"), "actionable cause: {msg}");
1231 }
1232
1233 #[test]
1234 fn render_fails_loud_when_work_dir_referenced_but_absent() {
1235 let mut v = vars();
1236 v.work_dir = None;
1237 let err = v.render("{work_dir}").unwrap_err();
1238 let msg = format!("{err:?}");
1239 assert!(msg.contains("{work_dir}"), "actionable token: {msg}");
1240 }
1241
1242 #[test]
1243 fn render_empty_system_is_legal() {
1244 let mut v = vars();
1245 v.system = String::new();
1246 assert_eq!(v.render("[{system}]").expect("render"), "[]");
1247 }
1248
1249 #[test]
1254 fn render_never_substitutes_inside_substituted_values() {
1255 let mut v = vars();
1256 v.prompt = "please mention {model} and {work_dir} literally".to_string();
1257 v.model = None; v.work_dir = None;
1259 let out = v.render("task: {prompt}").expect("render");
1260 assert_eq!(out, "task: please mention {model} and {work_dir} literally");
1261 }
1262
1263 #[test]
1266 fn render_substitutes_placeholder_nested_in_literal_braces() {
1267 let v = vars();
1268 let out = v.render(r#"{"task": "{prompt}", "n": 1}"#).expect("render");
1269 assert_eq!(out, r#"{"task": "do the task", "n": 1}"#);
1270 }
1271
1272 #[test]
1273 fn embed_references_scans_argv_stdin_env_cwd() {
1274 let mut t = EmbedTemplate {
1275 argv: vec!["cat".to_string()],
1276 ..Default::default()
1277 };
1278 assert!(!embed_references(&t, "{system_file}"));
1279 t.stdin = Some("{system_file}".to_string());
1280 assert!(embed_references(&t, "{system_file}"));
1281 t.stdin = None;
1282 t.env.insert("X".to_string(), "{system_file}".to_string());
1283 assert!(embed_references(&t, "{system_file}"));
1284 t.env.clear();
1285 t.cwd = Some("{system_file}".to_string());
1286 assert!(embed_references(&t, "{system_file}"));
1287 }
1288
1289 fn fake_output(stdout: &str, stderr: &str, code: i32) -> std::process::Output {
1290 #[cfg(unix)]
1291 use std::os::unix::process::ExitStatusExt;
1292 #[cfg(windows)]
1293 use std::os::windows::process::ExitStatusExt;
1294 use std::process::ExitStatus;
1295 #[cfg(unix)]
1296 let status = ExitStatus::from_raw(code << 8);
1297 #[cfg(windows)]
1298 let status = ExitStatus::from_raw(code as u32);
1299 std::process::Output {
1300 status,
1301 stdout: stdout.as_bytes().to_vec(),
1302 stderr: stderr.as_bytes().to_vec(),
1303 }
1304 }
1305
1306 #[test]
1310 fn normalize_without_decl_is_historical_json_or_raw() {
1311 let out = fake_output(r#"{"a": 1}"#, "", 0);
1312 let wr = normalize_plain_output(&out, None, None);
1313 assert_eq!(wr.value, serde_json::json!({"a": 1}));
1314 assert!(wr.ok);
1315
1316 let out = fake_output("plain text\n", "warned", 0);
1317 let wr = normalize_plain_output(&out, None, None);
1318 assert_eq!(
1319 wr.value,
1320 serde_json::json!({"raw": "plain text", "stderr": "warned"})
1321 );
1322 assert!(wr.ok);
1323
1324 let out = fake_output("boom", "", 1);
1325 let wr = normalize_plain_output(&out, None, None);
1326 assert!(!wr.ok, "non-zero exit is a failed step");
1327 }
1328
1329 #[test]
1330 fn normalize_result_ptr_extracts_declared_value() {
1331 let decl = SubprocessOutput {
1332 format: Some("json".to_string()),
1333 result_ptr: Some("/result".to_string()),
1334 ok_from: Some("exit_code".to_string()),
1335 stats: None,
1336 };
1337 let out = fake_output(r#"{"result": {"answer": 42}, "noise": true}"#, "", 0);
1338 let wr = normalize_plain_output(&out, Some(&decl), None);
1339 assert_eq!(wr.value, serde_json::json!({"answer": 42}));
1340 assert!(wr.ok);
1341 }
1342
1343 #[test]
1344 fn normalize_declared_json_unparsable_stdout_fails_loud() {
1345 let decl = SubprocessOutput {
1346 format: Some("json".to_string()),
1347 result_ptr: None,
1348 ok_from: None,
1349 stats: None,
1350 };
1351 let out = fake_output("not json at all", "stderr text", 0);
1352 let wr = normalize_plain_output(&out, Some(&decl), None);
1353 assert!(!wr.ok, "declared-JSON unparsable stdout is a failed step");
1354 assert_eq!(wr.value["raw"], "not json at all");
1355 assert_eq!(wr.value["stderr"], "stderr text");
1356 assert!(wr.value["parse_error"].is_string());
1357 }
1358
1359 #[test]
1360 fn normalize_missing_result_ptr_fails_loud_with_actionable_value() {
1361 let decl = SubprocessOutput {
1362 format: Some("json".to_string()),
1363 result_ptr: Some("/missing".to_string()),
1364 ok_from: None,
1365 stats: None,
1366 };
1367 let out = fake_output(r#"{"present": 1}"#, "", 0);
1368 let wr = normalize_plain_output(&out, Some(&decl), None);
1369 assert!(!wr.ok);
1370 assert!(wr.value["error"]
1371 .as_str()
1372 .expect("actionable error message")
1373 .contains("/missing"));
1374 }
1375
1376 #[test]
1377 fn normalize_ok_from_pointer_reads_boolean() {
1378 let decl = SubprocessOutput {
1379 format: Some("json".to_string()),
1380 result_ptr: Some("/result".to_string()),
1381 ok_from: Some("/ok".to_string()),
1382 stats: None,
1383 };
1384 let out = fake_output(r#"{"result": "r", "ok": true}"#, "", 0);
1386 let wr = normalize_plain_output(&out, Some(&decl), None);
1387 assert!(wr.ok);
1388 let out = fake_output(r#"{"result": "r", "ok": false}"#, "", 0);
1390 let wr = normalize_plain_output(&out, Some(&decl), None);
1391 assert!(!wr.ok);
1392 let out = fake_output(r#"{"result": "r", "ok": "yes"}"#, "", 0);
1394 let wr = normalize_plain_output(&out, Some(&decl), None);
1395 assert!(!wr.ok);
1396 }
1397
1398 #[test]
1402 fn declared_usage_ptr_accepts_a_partial_usage_object() {
1403 let decl = SubprocessOutput {
1404 format: Some("json".to_string()),
1405 result_ptr: Some("/result".to_string()),
1406 ok_from: None,
1407 stats: Some(mlua_swarm_schema::SubprocessStats {
1408 usage_ptr: Some("/usage".to_string()),
1409 model_ptr: None,
1410 num_turns_ptr: None,
1411 }),
1412 };
1413
1414 let out = fake_output(r#"{"result": "r", "usage": {"total_tokens": 512}}"#, "", 0);
1416 let usage = normalize_plain_output(&out, Some(&decl), None)
1417 .stats
1418 .and_then(|s| s.usage)
1419 .expect("a total-only usage must be recorded");
1420 assert_eq!(usage.total_tokens, 512);
1421 assert_eq!(usage.input_tokens, 0);
1422
1423 let out = fake_output(
1425 r#"{"result": "r", "usage": {"prompt_tokens": 10, "completion_tokens": 4}}"#,
1426 "",
1427 0,
1428 );
1429 let usage = normalize_plain_output(&out, Some(&decl), None)
1430 .stats
1431 .and_then(|s| s.usage)
1432 .expect("splits-only usage must be recorded");
1433 assert_eq!(usage.total_tokens, 14);
1434
1435 let out = fake_output(r#"{"result": "r", "usage": {"cached": 3}}"#, "", 0);
1437 assert!(normalize_plain_output(&out, Some(&decl), None)
1438 .stats
1439 .and_then(|s| s.usage)
1440 .is_none());
1441 }
1442}