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}
126
127#[derive(Debug, Clone)]
131pub enum StreamMode {
132 NdjsonLines,
134 SseEvents,
140 LengthPrefixed,
144}
145
146#[derive(Debug)]
150pub struct ProcessSpawner {
151 pub program: String,
154 pub args: Vec<String>,
156 pub use_stdin: bool,
160 pub stream_mode: Option<StreamMode>,
162 pub embed: Option<EmbedTemplate>,
166}
167
168impl ProcessSpawner {
169 pub fn new(program: impl Into<String>) -> Self {
172 Self {
173 program: program.into(),
174 args: Vec::new(),
175 use_stdin: true,
176 stream_mode: None,
177 embed: None,
178 }
179 }
180
181 pub fn arg(mut self, a: impl Into<String>) -> Self {
183 self.args.push(a.into());
184 self
185 }
186
187 pub fn args(mut self, args: impl IntoIterator<Item = impl Into<String>>) -> Self {
189 self.args.extend(args.into_iter().map(|a| a.into()));
190 self
191 }
192
193 pub fn use_stdin(mut self, v: bool) -> Self {
196 self.use_stdin = v;
197 self
198 }
199
200 pub fn stream_mode(mut self, mode: StreamMode) -> Self {
202 self.stream_mode = Some(mode);
203 self
204 }
205
206 pub fn plain(mut self) -> Self {
208 self.stream_mode = None;
209 self
210 }
211
212 pub fn ndjson(mut self, v: bool) -> Self {
216 self.stream_mode = if v {
217 Some(StreamMode::NdjsonLines)
218 } else {
219 None
220 };
221 self
222 }
223
224 pub fn run(cmd: impl Into<String>) -> Self {
226 Self {
227 program: "sh".into(),
228 args: vec!["-c".into(), cmd.into()],
229 use_stdin: true,
230 stream_mode: None,
231 embed: None,
232 }
233 }
234
235 pub fn cmd(program: impl Into<String>) -> Self {
238 Self {
239 program: program.into(),
240 args: Vec::new(),
241 use_stdin: true,
242 stream_mode: None,
243 embed: None,
244 }
245 }
246
247 pub fn embed(mut self, template: EmbedTemplate) -> Self {
249 self.embed = Some(template);
250 self
251 }
252}
253
254#[async_trait]
255impl SpawnerAdapter for ProcessSpawner {
256 async fn spawn(
257 &self,
258 engine: &Engine,
259 ctx: &Ctx,
260 task_id: StepId,
261 attempt: u32,
262 token: CapToken,
263 ) -> Result<Box<dyn Worker>, SpawnError> {
264 if let Some(embed) = &self.embed {
268 return self
269 .spawn_embed(embed, engine, ctx, task_id, attempt, token)
270 .await;
271 }
272 let directive = engine
280 .fetch_prompt(&token, &task_id)
281 .await
282 .map_err(|e| SpawnError::Internal(format!("fetch_prompt: {e}")))?;
283 let directive = crate::core::engine::render_directive_to_string(&directive);
286
287 let mut cmd = Command::new(&self.program);
288 cmd.args(&self.args)
289 .env("MSE_TOKEN_AGENT_ID", &token.agent_id)
290 .env("MSE_TOKEN_NONCE", &token.nonce)
291 .env("MSE_TASK_ID", task_id.as_str())
292 .env("MSE_ATTEMPT", attempt.to_string())
293 .env("MSE_CTX_AGENT", &ctx.agent)
294 .stdin(Stdio::piped())
295 .stdout(Stdio::piped())
296 .stderr(Stdio::piped());
297
298 if !self.use_stdin {
299 cmd.arg(&directive);
300 }
301
302 let mut child = cmd
303 .spawn()
304 .map_err(|e| SpawnError::Internal(format!("spawn failed: {e}")))?;
305
306 if self.use_stdin {
307 if let Some(mut stdin) = child.stdin.take() {
308 stdin
309 .write_all(directive.as_bytes())
310 .await
311 .map_err(|e| SpawnError::Internal(format!("stdin write: {e}")))?;
312 drop(stdin); }
314 }
315
316 let cancel = CancellationToken::new();
317 let cancel_inner = cancel.clone();
318 let worker_id = WorkerId::new();
319 tracing::debug!(worker_id = %worker_id, step_id = %task_id, "worker spawned (subprocess)");
321 let (tx, rx) = oneshot::channel();
322 let engine_for_emit = engine.clone();
326 let token_for_emit = token.clone();
327 let task_id_for_emit = task_id.clone();
328 let stream_mode = self.stream_mode.clone();
329
330 tokio::spawn(async move {
331 let result: Result<WorkerResult, WorkerError> = if let Some(mode) = stream_mode {
332 run_streaming_mode(
336 mode,
337 child,
338 &engine_for_emit,
339 &token_for_emit,
340 &task_id_for_emit,
341 attempt,
342 cancel_inner,
343 )
344 .await
345 } else {
346 let result = tokio::select! {
349 output = child.wait_with_output() => {
350 match output {
351 Ok(out) => {
352 let stdout = String::from_utf8_lossy(&out.stdout).to_string();
353 let value: Value = serde_json::from_str(stdout.trim())
354 .unwrap_or_else(|_| serde_json::json!({
355 "raw": stdout.trim_end(),
356 "stderr": String::from_utf8_lossy(&out.stderr).to_string(),
357 }));
358 Ok(WorkerResult {
359 value,
360 ok: out.status.success(),
361 stats: Some(subprocess_base_stats(&out.status, None)),
362 })
363 }
364 Err(e) => Err(WorkerError::Failed(format!("wait_with_output: {e}"))),
365 }
366 }
367 _ = cancel_inner.cancelled() => Err(WorkerError::Cancelled),
368 };
369 if let Ok(wr) = &result {
370 if let Some(stats) = wr.stats.clone() {
371 engine_for_emit
372 .record_worker_stats(&task_id_for_emit, attempt, stats)
373 .await;
374 }
375 let ev = OutputEvent::Final {
376 content: ContentRef::Inline {
377 value: wr.value.clone(),
378 },
379 ok: wr.ok,
380 };
381 let _ = engine_for_emit
382 .submit_output(&token_for_emit, &task_id_for_emit, attempt, ev)
383 .await;
384 }
385 result
386 };
387 let signal: Result<(), WorkerError> = result.map(|_| ());
389 let _ = tx.send(signal);
390 });
391
392 Ok(Box::new(ProcessWorker {
393 handler: WorkerJoinHandler {
394 worker_id,
395 cancel,
396 completion: rx,
397 },
398 }))
399 }
400}
401
402struct EmbedVars {
408 system: String,
409 system_file: Option<String>,
410 prompt: String,
411 model: Option<String>,
412 tools_csv: String,
413 work_dir: Option<String>,
414 task_id: String,
415 attempt: String,
416}
417
418impl EmbedVars {
419 fn lookup(&self, token: &str) -> Result<Option<&str>, SpawnError> {
424 let (value, cause): (Option<&str>, &str) = match token {
425 "system" => (Some(self.system.as_str()), ""),
426 "system_file" => (
427 self.system_file.as_deref(),
428 "no system prompt was baked for this attempt (agent has no profile.system_prompt?)",
429 ),
430 "prompt" => (Some(self.prompt.as_str()), ""),
431 "model" => (
432 self.model.as_deref(),
433 "no model declared (set profile.model or Runner::Subprocess overrides.model)",
434 ),
435 "tools_csv" => (Some(self.tools_csv.as_str()), ""),
436 "work_dir" => (
437 self.work_dir.as_deref(),
438 "no work_dir/project_root in the agent context view and no overrides.cwd",
439 ),
440 "task_id" => (Some(self.task_id.as_str()), ""),
441 "attempt" => (Some(self.attempt.as_str()), ""),
442 _ => return Ok(None),
443 };
444 match value {
445 Some(v) => Ok(Some(v)),
446 None => Err(SpawnError::Internal(format!(
447 "placeholder {{{token}}}: {cause}"
448 ))),
449 }
450 }
451
452 fn render(&self, tmpl: &str) -> Result<String, SpawnError> {
461 let mut out = String::with_capacity(tmpl.len());
462 let mut rest = tmpl;
463 while let Some(start) = rest.find('{') {
464 out.push_str(&rest[..start]);
465 let after = &rest[start + 1..];
466 let Some(end) = after.find('}') else {
467 out.push_str(&rest[start..]);
469 return Ok(out);
470 };
471 let token = &after[..end];
472 match self.lookup(token)? {
473 Some(value) => {
474 out.push_str(value);
475 rest = &after[end + 1..];
476 }
477 None => {
478 out.push('{');
483 rest = after;
484 }
485 }
486 }
487 out.push_str(rest);
488 Ok(out)
489 }
490}
491
492fn embed_references(embed: &EmbedTemplate, token: &str) -> bool {
494 embed.argv.iter().any(|a| a.contains(token))
495 || embed.stdin.as_deref().is_some_and(|s| s.contains(token))
496 || embed.env.values().any(|v| v.contains(token))
497 || embed.cwd.as_deref().is_some_and(|c| c.contains(token))
498}
499
500fn subprocess_base_stats(
506 status: &std::process::ExitStatus,
507 model: Option<&str>,
508) -> crate::store::trace::WorkerStats {
509 crate::store::trace::WorkerStats {
510 worker_kind: Some("subprocess".to_string()),
511 model: model.map(str::to_string),
512 usage: None,
513 num_turns: None,
514 adapter_data: Some(serde_json::json!({ "exit_code": status.code() })),
515 }
516}
517
518fn enrich_declared_stats(
523 stats: &mut crate::store::trace::WorkerStats,
524 parsed: &Value,
525 decl: &SubprocessOutput,
526) {
527 let Some(sd) = &decl.stats else { return };
528 if let Some(ptr) = sd.model_ptr.as_deref() {
529 if let Some(m) = parsed.pointer(ptr).and_then(|v| v.as_str()) {
530 stats.model = Some(m.to_string());
531 }
532 }
533 if let Some(ptr) = sd.num_turns_ptr.as_deref() {
534 if let Some(n) = parsed.pointer(ptr).and_then(|v| v.as_u64()) {
535 stats.num_turns = Some(n as u32);
536 }
537 }
538 if let Some(ptr) = sd.usage_ptr.as_deref() {
539 if let Some(u) = parsed.pointer(ptr) {
540 let input = u
541 .get("input_tokens")
542 .or_else(|| u.get("prompt_tokens"))
543 .and_then(|v| v.as_u64());
544 let output = u
545 .get("output_tokens")
546 .or_else(|| u.get("completion_tokens"))
547 .and_then(|v| v.as_u64());
548 let total = u.get("total_tokens").and_then(|v| v.as_u64());
549 if let Some(usage) = crate::store::trace::TokenUsage::from_parts(input, output, total) {
553 stats.usage = Some(usage);
554 if let Some(Value::Object(ad)) = stats.adapter_data.as_mut() {
557 ad.insert("usage_raw".to_string(), u.clone());
558 }
559 }
560 }
561 }
562}
563
564fn normalize_plain_output(
571 out: &std::process::Output,
572 decl: Option<&SubprocessOutput>,
573 model: Option<&str>,
574) -> WorkerResult {
575 let stdout = String::from_utf8_lossy(&out.stdout).to_string();
576 let stderr = || String::from_utf8_lossy(&out.stderr).to_string();
577 let exit_ok = out.status.success();
578 let mut stats = subprocess_base_stats(&out.status, model);
579
580 let Some(decl) = decl else {
581 let value: Value = serde_json::from_str(stdout.trim()).unwrap_or_else(|_| {
582 serde_json::json!({
583 "raw": stdout.trim_end(),
584 "stderr": stderr(),
585 })
586 });
587 return WorkerResult {
588 value,
589 ok: exit_ok,
590 stats: Some(stats),
591 };
592 };
593
594 let parsed: Result<Value, _> = serde_json::from_str(stdout.trim());
595 let parsed = match parsed {
596 Ok(v) => v,
597 Err(e) => {
598 if decl.format.as_deref() == Some("json") {
599 return WorkerResult {
602 value: serde_json::json!({
603 "raw": stdout.trim_end(),
604 "stderr": stderr(),
605 "parse_error": e.to_string(),
606 }),
607 ok: false,
608 stats: Some(stats),
609 };
610 }
611 return WorkerResult {
614 value: serde_json::json!({
615 "raw": stdout.trim_end(),
616 "stderr": stderr(),
617 }),
618 ok: exit_ok,
619 stats: Some(stats),
620 };
621 }
622 };
623
624 enrich_declared_stats(&mut stats, &parsed, decl);
625
626 let value = match decl.result_ptr.as_deref() {
627 Some(ptr) => match parsed.pointer(ptr) {
628 Some(v) => v.clone(),
629 None => {
630 return WorkerResult {
631 value: serde_json::json!({
632 "error": format!("result_ptr '{ptr}' not found in stdout JSON"),
633 "raw": parsed,
634 "stderr": stderr(),
635 }),
636 ok: false,
637 stats: Some(stats),
638 }
639 }
640 },
641 None => parsed.clone(),
642 };
643
644 let ok = match decl.ok_from.as_deref() {
645 None | Some("exit_code") => exit_ok,
646 Some(ptr) => parsed.pointer(ptr) == Some(&Value::Bool(true)),
647 };
648
649 WorkerResult {
650 value,
651 ok,
652 stats: Some(stats),
653 }
654}
655
656impl ProcessSpawner {
657 async fn spawn_embed(
663 &self,
664 embed: &EmbedTemplate,
665 engine: &Engine,
666 ctx: &Ctx,
667 task_id: StepId,
668 attempt: u32,
669 token: CapToken,
670 ) -> Result<Box<dyn Worker>, SpawnError> {
671 let prompt_value = engine
675 .fetch_prompt(&token, &task_id)
676 .await
677 .map_err(|e| SpawnError::Internal(format!("fetch_prompt: {e}")))?;
678 let system = match embed.system_prompt.as_deref() {
679 Some(tmpl) => {
680 let slots = crate::operator::render::slots_from_prompt(&prompt_value);
681 let rendered = crate::operator::render::render_system(tmpl, &slots)
682 .map_err(|e| SpawnError::Internal(format!("render system_prompt: {e}")))?;
683 Some(rendered)
684 }
685 None => None,
686 };
687 engine
688 .bake_worker_system_prompt(&task_id, attempt, system.clone())
689 .await
690 .map_err(|e| SpawnError::Internal(format!("bake system_prompt: {e}")))?;
691
692 let payload = engine
698 .fetch_worker_payload(&token, &task_id)
699 .await
700 .map_err(|e| SpawnError::Internal(format!("fetch_worker_payload: {e}")))?;
701
702 let system_file = if embed_references(embed, "{system_file}") {
705 match engine
706 .materialize_system_file(&task_id, attempt)
707 .await
708 .map_err(|e| SpawnError::Internal(format!("materialize system_file: {e}")))?
709 {
710 Some(path) => Some(path.display().to_string()),
711 None => {
712 return Err(SpawnError::Internal(
713 "placeholder {system_file}: no system prompt was baked for this \
714 attempt (agent has no profile.system_prompt?)"
715 .into(),
716 ))
717 }
718 }
719 } else {
720 None
721 };
722
723 let view = AgentContextView::materialized_or_from_ctx(ctx);
726 let work_dir = view.work_dir.clone().or_else(|| view.project_root.clone());
727
728 let vars = EmbedVars {
729 system: system.unwrap_or_default(),
730 system_file,
731 prompt: payload.prompt.clone(),
732 model: embed.model.clone(),
733 tools_csv: embed.tools_csv.clone(),
734 work_dir,
735 task_id: task_id.to_string(),
736 attempt: attempt.to_string(),
737 };
738
739 if embed.argv.is_empty() {
741 return Err(SpawnError::Internal(
742 "embed template: argv must not be empty".into(),
743 ));
744 }
745 let mut rendered_argv = Vec::with_capacity(embed.argv.len());
746 for a in &embed.argv {
747 rendered_argv.push(vars.render(a)?);
748 }
749 let mut rendered_env = Vec::with_capacity(embed.env.len());
750 for (k, v) in &embed.env {
751 rendered_env.push((k.clone(), vars.render(v)?));
752 }
753 let rendered_cwd = match embed.cwd.as_deref() {
754 Some(c) => Some(vars.render(c)?),
755 None => None,
756 };
757 let rendered_stdin = match embed.stdin.as_deref() {
758 Some(s) => Some(vars.render(s)?),
759 None => None,
760 };
761
762 let mut cmd = Command::new(&rendered_argv[0]);
763 cmd.args(&rendered_argv[1..])
764 .env("MSE_TOKEN_AGENT_ID", &token.agent_id)
765 .env("MSE_TOKEN_NONCE", &token.nonce)
766 .env("MSE_TASK_ID", task_id.as_str())
767 .env("MSE_ATTEMPT", attempt.to_string())
768 .env("MSE_CTX_AGENT", &ctx.agent)
769 .stdin(Stdio::piped())
770 .stdout(Stdio::piped())
771 .stderr(Stdio::piped());
772 for (k, v) in &rendered_env {
773 cmd.env(k, v);
774 }
775 if let Some(cwd) = &rendered_cwd {
776 cmd.current_dir(cwd);
777 }
778
779 let mut child = cmd
780 .spawn()
781 .map_err(|e| SpawnError::Internal(format!("spawn failed: {e}")))?;
782
783 if let Some(stdin_body) = rendered_stdin {
784 if let Some(mut stdin) = child.stdin.take() {
785 stdin
786 .write_all(stdin_body.as_bytes())
787 .await
788 .map_err(|e| SpawnError::Internal(format!("stdin write: {e}")))?;
789 drop(stdin); }
791 } else {
792 drop(child.stdin.take());
794 }
795
796 let cancel = CancellationToken::new();
797 let cancel_inner = cancel.clone();
798 let worker_id = WorkerId::new();
799 tracing::debug!(worker_id = %worker_id, step_id = %task_id, "worker spawned (subprocess embed)");
800 let (tx, rx) = oneshot::channel();
801 let engine_for_emit = engine.clone();
802 let token_for_emit = token.clone();
803 let task_id_for_emit = task_id.clone();
804 let stream_mode = self.stream_mode.clone();
805 let output_decl = embed.output.clone();
806 let model_for_stats = embed.model.clone();
807
808 tokio::spawn(async move {
809 let result: Result<WorkerResult, WorkerError> = if let Some(mode) = stream_mode {
810 run_streaming_mode(
813 mode,
814 child,
815 &engine_for_emit,
816 &token_for_emit,
817 &task_id_for_emit,
818 attempt,
819 cancel_inner,
820 )
821 .await
822 } else {
823 let result = tokio::select! {
824 output = child.wait_with_output() => {
825 match output {
826 Ok(out) => Ok(normalize_plain_output(
827 &out,
828 output_decl.as_ref(),
829 model_for_stats.as_deref(),
830 )),
831 Err(e) => Err(WorkerError::Failed(format!("wait_with_output: {e}"))),
832 }
833 }
834 _ = cancel_inner.cancelled() => Err(WorkerError::Cancelled),
835 };
836 if let Ok(wr) = &result {
837 if let Some(stats) = wr.stats.clone() {
838 engine_for_emit
839 .record_worker_stats(&task_id_for_emit, attempt, stats)
840 .await;
841 }
842 let ev = OutputEvent::Final {
843 content: ContentRef::Inline {
844 value: wr.value.clone(),
845 },
846 ok: wr.ok,
847 };
848 let _ = engine_for_emit
849 .submit_output(&token_for_emit, &task_id_for_emit, attempt, ev)
850 .await;
851 }
852 result
853 };
854 let signal: Result<(), WorkerError> = result.map(|_| ());
855 let _ = tx.send(signal);
856 });
857
858 Ok(Box::new(ProcessWorker {
859 handler: WorkerJoinHandler {
860 worker_id,
861 cancel,
862 completion: rx,
863 },
864 }))
865 }
866}
867
868pub struct ProcessWorker {
872 pub handler: WorkerJoinHandler,
875}
876
877#[async_trait]
878impl Worker for ProcessWorker {
879 fn id(&self) -> &WorkerId {
880 &self.handler.worker_id
881 }
882 fn cancel_token(&self) -> CancellationToken {
883 self.handler.cancel.clone()
884 }
885 async fn join(self: Box<Self>) -> Result<(), WorkerError> {
886 self.handler.await_completion().await
887 }
888}
889
890async fn run_streaming_mode(
896 mode: StreamMode,
897 mut child: tokio::process::Child,
898 engine: &Engine,
899 token: &CapToken,
900 task_id: &StepId,
901 attempt: u32,
902 cancel: CancellationToken,
903) -> Result<WorkerResult, WorkerError> {
904 let stdout = child
905 .stdout
906 .take()
907 .ok_or_else(|| WorkerError::Failed("streaming: stdout pipe missing".into()))?;
908
909 let last_final = match mode {
910 StreamMode::NdjsonLines => {
911 read_ndjson(stdout, engine, token, task_id, attempt, cancel.clone()).await?
912 }
913 StreamMode::SseEvents => {
914 read_sse(stdout, engine, token, task_id, attempt, cancel.clone()).await?
915 }
916 StreamMode::LengthPrefixed => {
917 read_length_prefixed(stdout, engine, token, task_id, attempt, cancel.clone()).await?
918 }
919 };
920
921 let status = child
922 .wait()
923 .await
924 .map_err(|e| WorkerError::Failed(format!("streaming wait: {e}")))?;
925
926 match last_final {
927 Some((value, ok)) => Ok(WorkerResult {
928 value,
929 ok: ok && status.success(),
930 stats: Some(subprocess_base_stats(&status, None)),
933 }),
934 None => {
935 let value = serde_json::json!({
937 "raw": "",
938 "note": "streaming mode: no Final event received",
939 "exit_success": status.success(),
940 });
941 let _ = engine
942 .submit_output(
943 token,
944 task_id,
945 attempt,
946 OutputEvent::Final {
947 content: ContentRef::Inline {
948 value: value.clone(),
949 },
950 ok: false,
951 },
952 )
953 .await;
954 Ok(WorkerResult {
955 value,
956 ok: false,
957 stats: Some(subprocess_base_stats(&status, None)),
958 })
959 }
960 }
961}
962
963async fn forward_event(
966 engine: &Engine,
967 token: &CapToken,
968 task_id: &StepId,
969 attempt: u32,
970 ev: OutputEvent,
971 last_final: &mut Option<(Value, bool)>,
972) {
973 if let OutputEvent::Final { content, ok } = &ev {
974 let value = match content {
975 ContentRef::Inline { value } => value.clone(),
976 ContentRef::FileRef {
977 path,
978 mime,
979 size_hint,
980 } => serde_json::json!({
981 "file_ref": path.to_string_lossy(),
982 "mime": mime,
983 "size_hint": size_hint,
984 }),
985 };
986 *last_final = Some((value, *ok));
987 }
988 let _ = engine.submit_output(token, task_id, attempt, ev).await;
989}
990
991async fn read_ndjson(
994 stdout: tokio::process::ChildStdout,
995 engine: &Engine,
996 token: &CapToken,
997 task_id: &StepId,
998 attempt: u32,
999 cancel: CancellationToken,
1000) -> Result<Option<(Value, bool)>, WorkerError> {
1001 let mut reader = BufReader::new(stdout).lines();
1002 let mut last_final = None;
1003 loop {
1004 tokio::select! {
1005 line_res = reader.next_line() => match line_res {
1006 Ok(Some(line)) => {
1007 let trimmed = line.trim();
1008 if trimmed.is_empty() { continue; }
1009 if let Ok(ev) = serde_json::from_str::<OutputEvent>(trimmed) {
1010 forward_event(engine, token, task_id, attempt, ev, &mut last_final).await;
1011 }
1012 }
1013 Ok(None) => break,
1014 Err(e) => return Err(WorkerError::Failed(format!("ndjson read: {e}"))),
1015 },
1016 _ = cancel.cancelled() => return Err(WorkerError::Cancelled),
1017 }
1018 }
1019 Ok(last_final)
1020}
1021
1022async fn read_sse(
1026 stdout: tokio::process::ChildStdout,
1027 engine: &Engine,
1028 token: &CapToken,
1029 task_id: &StepId,
1030 attempt: u32,
1031 cancel: CancellationToken,
1032) -> Result<Option<(Value, bool)>, WorkerError> {
1033 let mut reader = BufReader::new(stdout).lines();
1034 let mut last_final = None;
1035 let mut data_buf = String::new();
1036 loop {
1037 tokio::select! {
1038 line_res = reader.next_line() => match line_res {
1039 Ok(Some(line)) => {
1040 if line.is_empty() {
1041 if !data_buf.is_empty() {
1043 if let Ok(ev) = serde_json::from_str::<OutputEvent>(data_buf.trim()) {
1044 forward_event(engine, token, task_id, attempt, ev, &mut last_final).await;
1045 }
1046 data_buf.clear();
1047 }
1048 } else if let Some(rest) = line.strip_prefix("data:") {
1049 let payload = rest.strip_prefix(' ').unwrap_or(rest);
1051 if !data_buf.is_empty() {
1052 data_buf.push('\n');
1053 }
1054 data_buf.push_str(payload);
1055 }
1056 }
1058 Ok(None) => {
1059 if !data_buf.is_empty() {
1061 if let Ok(ev) = serde_json::from_str::<OutputEvent>(data_buf.trim()) {
1062 forward_event(engine, token, task_id, attempt, ev, &mut last_final).await;
1063 }
1064 }
1065 break;
1066 }
1067 Err(e) => return Err(WorkerError::Failed(format!("sse read: {e}"))),
1068 },
1069 _ = cancel.cancelled() => return Err(WorkerError::Cancelled),
1070 }
1071 }
1072 Ok(last_final)
1073}
1074
1075async fn read_length_prefixed(
1078 mut stdout: tokio::process::ChildStdout,
1079 engine: &Engine,
1080 token: &CapToken,
1081 task_id: &StepId,
1082 attempt: u32,
1083 cancel: CancellationToken,
1084) -> Result<Option<(Value, bool)>, WorkerError> {
1085 use tokio::io::AsyncReadExt;
1086 let mut last_final = None;
1087 loop {
1088 let mut len_buf = [0u8; 4];
1090 let read_fut = stdout.read_exact(&mut len_buf);
1091 let read_res = tokio::select! {
1092 r = read_fut => r,
1093 _ = cancel.cancelled() => return Err(WorkerError::Cancelled),
1094 };
1095 match read_res {
1096 Ok(_) => {}
1097 Err(e) if e.kind() == std::io::ErrorKind::UnexpectedEof => break, Err(e) => return Err(WorkerError::Failed(format!("len read: {e}"))),
1099 }
1100 let len = u32::from_be_bytes(len_buf) as usize;
1101 if len == 0 || len > 16 * 1024 * 1024 {
1102 break;
1104 }
1105 let mut payload = vec![0u8; len];
1106 let read_fut = stdout.read_exact(&mut payload);
1107 let read_res = tokio::select! {
1108 r = read_fut => r,
1109 _ = cancel.cancelled() => return Err(WorkerError::Cancelled),
1110 };
1111 if read_res.is_err() {
1112 break;
1113 }
1114 if let Ok(ev) = serde_json::from_slice::<OutputEvent>(&payload) {
1115 forward_event(engine, token, task_id, attempt, ev, &mut last_final).await;
1116 }
1117 }
1118 Ok(last_final)
1119}
1120
1121#[cfg(test)]
1122mod embed_tests {
1123 use super::*;
1124
1125 fn vars() -> EmbedVars {
1126 EmbedVars {
1127 system: "SYS".to_string(),
1128 system_file: Some("/tmp/sys.md".to_string()),
1129 prompt: "do the task".to_string(),
1130 model: Some("small".to_string()),
1131 tools_csv: "Read,Grep".to_string(),
1132 work_dir: Some("/tmp/wd".to_string()),
1133 task_id: "ST-1".to_string(),
1134 attempt: "1".to_string(),
1135 }
1136 }
1137
1138 #[test]
1139 fn render_substitutes_every_closed_set_token() {
1140 let v = vars();
1141 let out = v
1142 .render("{system}|{system_file}|{prompt}|{model}|{tools_csv}|{work_dir}|{task_id}|{attempt}")
1143 .expect("render");
1144 assert_eq!(
1145 out,
1146 "SYS|/tmp/sys.md|do the task|small|Read,Grep|/tmp/wd|ST-1|1"
1147 );
1148 }
1149
1150 #[test]
1151 fn render_leaves_non_placeholder_braces_alone() {
1152 let v = vars();
1153 let out = v
1154 .render(r#"echo '{"result": "{prompt}"}'"#)
1155 .expect("render");
1156 assert_eq!(out, r#"echo '{"result": "do the task"}'"#);
1157 }
1158
1159 #[test]
1160 fn render_fails_loud_when_model_referenced_but_absent() {
1161 let mut v = vars();
1162 v.model = None;
1163 let err = v.render("--model {model}").unwrap_err();
1164 let msg = format!("{err:?}");
1165 assert!(msg.contains("{model}"), "actionable token in error: {msg}");
1166 assert!(msg.contains("profile.model"), "actionable cause: {msg}");
1167 }
1168
1169 #[test]
1170 fn render_fails_loud_when_work_dir_referenced_but_absent() {
1171 let mut v = vars();
1172 v.work_dir = None;
1173 let err = v.render("{work_dir}").unwrap_err();
1174 let msg = format!("{err:?}");
1175 assert!(msg.contains("{work_dir}"), "actionable token: {msg}");
1176 }
1177
1178 #[test]
1179 fn render_empty_system_is_legal() {
1180 let mut v = vars();
1181 v.system = String::new();
1182 assert_eq!(v.render("[{system}]").expect("render"), "[]");
1183 }
1184
1185 #[test]
1190 fn render_never_substitutes_inside_substituted_values() {
1191 let mut v = vars();
1192 v.prompt = "please mention {model} and {work_dir} literally".to_string();
1193 v.model = None; v.work_dir = None;
1195 let out = v.render("task: {prompt}").expect("render");
1196 assert_eq!(out, "task: please mention {model} and {work_dir} literally");
1197 }
1198
1199 #[test]
1202 fn render_substitutes_placeholder_nested_in_literal_braces() {
1203 let v = vars();
1204 let out = v.render(r#"{"task": "{prompt}", "n": 1}"#).expect("render");
1205 assert_eq!(out, r#"{"task": "do the task", "n": 1}"#);
1206 }
1207
1208 #[test]
1209 fn embed_references_scans_argv_stdin_env_cwd() {
1210 let mut t = EmbedTemplate {
1211 argv: vec!["cat".to_string()],
1212 ..Default::default()
1213 };
1214 assert!(!embed_references(&t, "{system_file}"));
1215 t.stdin = Some("{system_file}".to_string());
1216 assert!(embed_references(&t, "{system_file}"));
1217 t.stdin = None;
1218 t.env.insert("X".to_string(), "{system_file}".to_string());
1219 assert!(embed_references(&t, "{system_file}"));
1220 t.env.clear();
1221 t.cwd = Some("{system_file}".to_string());
1222 assert!(embed_references(&t, "{system_file}"));
1223 }
1224
1225 fn fake_output(stdout: &str, stderr: &str, code: i32) -> std::process::Output {
1226 #[cfg(unix)]
1227 use std::os::unix::process::ExitStatusExt;
1228 #[cfg(windows)]
1229 use std::os::windows::process::ExitStatusExt;
1230 use std::process::ExitStatus;
1231 #[cfg(unix)]
1232 let status = ExitStatus::from_raw(code << 8);
1233 #[cfg(windows)]
1234 let status = ExitStatus::from_raw(code as u32);
1235 std::process::Output {
1236 status,
1237 stdout: stdout.as_bytes().to_vec(),
1238 stderr: stderr.as_bytes().to_vec(),
1239 }
1240 }
1241
1242 #[test]
1246 fn normalize_without_decl_is_historical_json_or_raw() {
1247 let out = fake_output(r#"{"a": 1}"#, "", 0);
1248 let wr = normalize_plain_output(&out, None, None);
1249 assert_eq!(wr.value, serde_json::json!({"a": 1}));
1250 assert!(wr.ok);
1251
1252 let out = fake_output("plain text\n", "warned", 0);
1253 let wr = normalize_plain_output(&out, None, None);
1254 assert_eq!(
1255 wr.value,
1256 serde_json::json!({"raw": "plain text", "stderr": "warned"})
1257 );
1258 assert!(wr.ok);
1259
1260 let out = fake_output("boom", "", 1);
1261 let wr = normalize_plain_output(&out, None, None);
1262 assert!(!wr.ok, "non-zero exit is a failed step");
1263 }
1264
1265 #[test]
1266 fn normalize_result_ptr_extracts_declared_value() {
1267 let decl = SubprocessOutput {
1268 format: Some("json".to_string()),
1269 result_ptr: Some("/result".to_string()),
1270 ok_from: Some("exit_code".to_string()),
1271 stats: None,
1272 };
1273 let out = fake_output(r#"{"result": {"answer": 42}, "noise": true}"#, "", 0);
1274 let wr = normalize_plain_output(&out, Some(&decl), None);
1275 assert_eq!(wr.value, serde_json::json!({"answer": 42}));
1276 assert!(wr.ok);
1277 }
1278
1279 #[test]
1280 fn normalize_declared_json_unparsable_stdout_fails_loud() {
1281 let decl = SubprocessOutput {
1282 format: Some("json".to_string()),
1283 result_ptr: None,
1284 ok_from: None,
1285 stats: None,
1286 };
1287 let out = fake_output("not json at all", "stderr text", 0);
1288 let wr = normalize_plain_output(&out, Some(&decl), None);
1289 assert!(!wr.ok, "declared-JSON unparsable stdout is a failed step");
1290 assert_eq!(wr.value["raw"], "not json at all");
1291 assert_eq!(wr.value["stderr"], "stderr text");
1292 assert!(wr.value["parse_error"].is_string());
1293 }
1294
1295 #[test]
1296 fn normalize_missing_result_ptr_fails_loud_with_actionable_value() {
1297 let decl = SubprocessOutput {
1298 format: Some("json".to_string()),
1299 result_ptr: Some("/missing".to_string()),
1300 ok_from: None,
1301 stats: None,
1302 };
1303 let out = fake_output(r#"{"present": 1}"#, "", 0);
1304 let wr = normalize_plain_output(&out, Some(&decl), None);
1305 assert!(!wr.ok);
1306 assert!(wr.value["error"]
1307 .as_str()
1308 .expect("actionable error message")
1309 .contains("/missing"));
1310 }
1311
1312 #[test]
1313 fn normalize_ok_from_pointer_reads_boolean() {
1314 let decl = SubprocessOutput {
1315 format: Some("json".to_string()),
1316 result_ptr: Some("/result".to_string()),
1317 ok_from: Some("/ok".to_string()),
1318 stats: None,
1319 };
1320 let out = fake_output(r#"{"result": "r", "ok": true}"#, "", 0);
1322 let wr = normalize_plain_output(&out, Some(&decl), None);
1323 assert!(wr.ok);
1324 let out = fake_output(r#"{"result": "r", "ok": false}"#, "", 0);
1326 let wr = normalize_plain_output(&out, Some(&decl), None);
1327 assert!(!wr.ok);
1328 let out = fake_output(r#"{"result": "r", "ok": "yes"}"#, "", 0);
1330 let wr = normalize_plain_output(&out, Some(&decl), None);
1331 assert!(!wr.ok);
1332 }
1333
1334 #[test]
1338 fn declared_usage_ptr_accepts_a_partial_usage_object() {
1339 let decl = SubprocessOutput {
1340 format: Some("json".to_string()),
1341 result_ptr: Some("/result".to_string()),
1342 ok_from: None,
1343 stats: Some(mlua_swarm_schema::SubprocessStats {
1344 usage_ptr: Some("/usage".to_string()),
1345 model_ptr: None,
1346 num_turns_ptr: None,
1347 }),
1348 };
1349
1350 let out = fake_output(r#"{"result": "r", "usage": {"total_tokens": 512}}"#, "", 0);
1352 let usage = normalize_plain_output(&out, Some(&decl), None)
1353 .stats
1354 .and_then(|s| s.usage)
1355 .expect("a total-only usage must be recorded");
1356 assert_eq!(usage.total_tokens, 512);
1357 assert_eq!(usage.input_tokens, 0);
1358
1359 let out = fake_output(
1361 r#"{"result": "r", "usage": {"prompt_tokens": 10, "completion_tokens": 4}}"#,
1362 "",
1363 0,
1364 );
1365 let usage = normalize_plain_output(&out, Some(&decl), None)
1366 .stats
1367 .and_then(|s| s.usage)
1368 .expect("splits-only usage must be recorded");
1369 assert_eq!(usage.total_tokens, 14);
1370
1371 let out = fake_output(r#"{"result": "r", "usage": {"cached": 3}}"#, "", 0);
1373 assert!(normalize_plain_output(&out, Some(&decl), None)
1374 .stats
1375 .and_then(|s| s.usage)
1376 .is_none());
1377 }
1378}