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 if let (Some(i), Some(o)) = (input, output) {
549 let total = u
550 .get("total_tokens")
551 .and_then(|v| v.as_u64())
552 .unwrap_or(i + o);
553 stats.usage = Some(crate::store::trace::TokenUsage {
554 input_tokens: i,
555 output_tokens: o,
556 total_tokens: total,
557 });
558 if let Some(Value::Object(ad)) = stats.adapter_data.as_mut() {
561 ad.insert("usage_raw".to_string(), u.clone());
562 }
563 }
564 }
565 }
566}
567
568fn normalize_plain_output(
575 out: &std::process::Output,
576 decl: Option<&SubprocessOutput>,
577 model: Option<&str>,
578) -> WorkerResult {
579 let stdout = String::from_utf8_lossy(&out.stdout).to_string();
580 let stderr = || String::from_utf8_lossy(&out.stderr).to_string();
581 let exit_ok = out.status.success();
582 let mut stats = subprocess_base_stats(&out.status, model);
583
584 let Some(decl) = decl else {
585 let value: Value = serde_json::from_str(stdout.trim()).unwrap_or_else(|_| {
586 serde_json::json!({
587 "raw": stdout.trim_end(),
588 "stderr": stderr(),
589 })
590 });
591 return WorkerResult {
592 value,
593 ok: exit_ok,
594 stats: Some(stats),
595 };
596 };
597
598 let parsed: Result<Value, _> = serde_json::from_str(stdout.trim());
599 let parsed = match parsed {
600 Ok(v) => v,
601 Err(e) => {
602 if decl.format.as_deref() == Some("json") {
603 return WorkerResult {
606 value: serde_json::json!({
607 "raw": stdout.trim_end(),
608 "stderr": stderr(),
609 "parse_error": e.to_string(),
610 }),
611 ok: false,
612 stats: Some(stats),
613 };
614 }
615 return WorkerResult {
618 value: serde_json::json!({
619 "raw": stdout.trim_end(),
620 "stderr": stderr(),
621 }),
622 ok: exit_ok,
623 stats: Some(stats),
624 };
625 }
626 };
627
628 enrich_declared_stats(&mut stats, &parsed, decl);
629
630 let value = match decl.result_ptr.as_deref() {
631 Some(ptr) => match parsed.pointer(ptr) {
632 Some(v) => v.clone(),
633 None => {
634 return WorkerResult {
635 value: serde_json::json!({
636 "error": format!("result_ptr '{ptr}' not found in stdout JSON"),
637 "raw": parsed,
638 "stderr": stderr(),
639 }),
640 ok: false,
641 stats: Some(stats),
642 }
643 }
644 },
645 None => parsed.clone(),
646 };
647
648 let ok = match decl.ok_from.as_deref() {
649 None | Some("exit_code") => exit_ok,
650 Some(ptr) => parsed.pointer(ptr) == Some(&Value::Bool(true)),
651 };
652
653 WorkerResult {
654 value,
655 ok,
656 stats: Some(stats),
657 }
658}
659
660impl ProcessSpawner {
661 async fn spawn_embed(
667 &self,
668 embed: &EmbedTemplate,
669 engine: &Engine,
670 ctx: &Ctx,
671 task_id: StepId,
672 attempt: u32,
673 token: CapToken,
674 ) -> Result<Box<dyn Worker>, SpawnError> {
675 let prompt_value = engine
679 .fetch_prompt(&token, &task_id)
680 .await
681 .map_err(|e| SpawnError::Internal(format!("fetch_prompt: {e}")))?;
682 let system = match embed.system_prompt.as_deref() {
683 Some(tmpl) => {
684 let slots = crate::operator::render::slots_from_prompt(&prompt_value);
685 let rendered = crate::operator::render::render_system(tmpl, &slots)
686 .map_err(|e| SpawnError::Internal(format!("render system_prompt: {e}")))?;
687 Some(rendered)
688 }
689 None => None,
690 };
691 engine
692 .bake_worker_system_prompt(&task_id, attempt, system.clone())
693 .await
694 .map_err(|e| SpawnError::Internal(format!("bake system_prompt: {e}")))?;
695
696 let payload = engine
702 .fetch_worker_payload(&token, &task_id)
703 .await
704 .map_err(|e| SpawnError::Internal(format!("fetch_worker_payload: {e}")))?;
705
706 let system_file = if embed_references(embed, "{system_file}") {
709 match engine
710 .materialize_system_file(&task_id, attempt)
711 .await
712 .map_err(|e| SpawnError::Internal(format!("materialize system_file: {e}")))?
713 {
714 Some(path) => Some(path.display().to_string()),
715 None => {
716 return Err(SpawnError::Internal(
717 "placeholder {system_file}: no system prompt was baked for this \
718 attempt (agent has no profile.system_prompt?)"
719 .into(),
720 ))
721 }
722 }
723 } else {
724 None
725 };
726
727 let view = AgentContextView::materialized_or_from_ctx(ctx);
730 let work_dir = view.work_dir.clone().or_else(|| view.project_root.clone());
731
732 let vars = EmbedVars {
733 system: system.unwrap_or_default(),
734 system_file,
735 prompt: payload.prompt.clone(),
736 model: embed.model.clone(),
737 tools_csv: embed.tools_csv.clone(),
738 work_dir,
739 task_id: task_id.to_string(),
740 attempt: attempt.to_string(),
741 };
742
743 if embed.argv.is_empty() {
745 return Err(SpawnError::Internal(
746 "embed template: argv must not be empty".into(),
747 ));
748 }
749 let mut rendered_argv = Vec::with_capacity(embed.argv.len());
750 for a in &embed.argv {
751 rendered_argv.push(vars.render(a)?);
752 }
753 let mut rendered_env = Vec::with_capacity(embed.env.len());
754 for (k, v) in &embed.env {
755 rendered_env.push((k.clone(), vars.render(v)?));
756 }
757 let rendered_cwd = match embed.cwd.as_deref() {
758 Some(c) => Some(vars.render(c)?),
759 None => None,
760 };
761 let rendered_stdin = match embed.stdin.as_deref() {
762 Some(s) => Some(vars.render(s)?),
763 None => None,
764 };
765
766 let mut cmd = Command::new(&rendered_argv[0]);
767 cmd.args(&rendered_argv[1..])
768 .env("MSE_TOKEN_AGENT_ID", &token.agent_id)
769 .env("MSE_TOKEN_NONCE", &token.nonce)
770 .env("MSE_TASK_ID", task_id.as_str())
771 .env("MSE_ATTEMPT", attempt.to_string())
772 .env("MSE_CTX_AGENT", &ctx.agent)
773 .stdin(Stdio::piped())
774 .stdout(Stdio::piped())
775 .stderr(Stdio::piped());
776 for (k, v) in &rendered_env {
777 cmd.env(k, v);
778 }
779 if let Some(cwd) = &rendered_cwd {
780 cmd.current_dir(cwd);
781 }
782
783 let mut child = cmd
784 .spawn()
785 .map_err(|e| SpawnError::Internal(format!("spawn failed: {e}")))?;
786
787 if let Some(stdin_body) = rendered_stdin {
788 if let Some(mut stdin) = child.stdin.take() {
789 stdin
790 .write_all(stdin_body.as_bytes())
791 .await
792 .map_err(|e| SpawnError::Internal(format!("stdin write: {e}")))?;
793 drop(stdin); }
795 } else {
796 drop(child.stdin.take());
798 }
799
800 let cancel = CancellationToken::new();
801 let cancel_inner = cancel.clone();
802 let worker_id = WorkerId::new();
803 tracing::debug!(worker_id = %worker_id, step_id = %task_id, "worker spawned (subprocess embed)");
804 let (tx, rx) = oneshot::channel();
805 let engine_for_emit = engine.clone();
806 let token_for_emit = token.clone();
807 let task_id_for_emit = task_id.clone();
808 let stream_mode = self.stream_mode.clone();
809 let output_decl = embed.output.clone();
810 let model_for_stats = embed.model.clone();
811
812 tokio::spawn(async move {
813 let result: Result<WorkerResult, WorkerError> = if let Some(mode) = stream_mode {
814 run_streaming_mode(
817 mode,
818 child,
819 &engine_for_emit,
820 &token_for_emit,
821 &task_id_for_emit,
822 attempt,
823 cancel_inner,
824 )
825 .await
826 } else {
827 let result = tokio::select! {
828 output = child.wait_with_output() => {
829 match output {
830 Ok(out) => Ok(normalize_plain_output(
831 &out,
832 output_decl.as_ref(),
833 model_for_stats.as_deref(),
834 )),
835 Err(e) => Err(WorkerError::Failed(format!("wait_with_output: {e}"))),
836 }
837 }
838 _ = cancel_inner.cancelled() => Err(WorkerError::Cancelled),
839 };
840 if let Ok(wr) = &result {
841 if let Some(stats) = wr.stats.clone() {
842 engine_for_emit
843 .record_worker_stats(&task_id_for_emit, attempt, stats)
844 .await;
845 }
846 let ev = OutputEvent::Final {
847 content: ContentRef::Inline {
848 value: wr.value.clone(),
849 },
850 ok: wr.ok,
851 };
852 let _ = engine_for_emit
853 .submit_output(&token_for_emit, &task_id_for_emit, attempt, ev)
854 .await;
855 }
856 result
857 };
858 let signal: Result<(), WorkerError> = result.map(|_| ());
859 let _ = tx.send(signal);
860 });
861
862 Ok(Box::new(ProcessWorker {
863 handler: WorkerJoinHandler {
864 worker_id,
865 cancel,
866 completion: rx,
867 },
868 }))
869 }
870}
871
872pub struct ProcessWorker {
876 pub handler: WorkerJoinHandler,
879}
880
881#[async_trait]
882impl Worker for ProcessWorker {
883 fn id(&self) -> &WorkerId {
884 &self.handler.worker_id
885 }
886 fn cancel_token(&self) -> CancellationToken {
887 self.handler.cancel.clone()
888 }
889 async fn join(self: Box<Self>) -> Result<(), WorkerError> {
890 self.handler.await_completion().await
891 }
892}
893
894async fn run_streaming_mode(
900 mode: StreamMode,
901 mut child: tokio::process::Child,
902 engine: &Engine,
903 token: &CapToken,
904 task_id: &StepId,
905 attempt: u32,
906 cancel: CancellationToken,
907) -> Result<WorkerResult, WorkerError> {
908 let stdout = child
909 .stdout
910 .take()
911 .ok_or_else(|| WorkerError::Failed("streaming: stdout pipe missing".into()))?;
912
913 let last_final = match mode {
914 StreamMode::NdjsonLines => {
915 read_ndjson(stdout, engine, token, task_id, attempt, cancel.clone()).await?
916 }
917 StreamMode::SseEvents => {
918 read_sse(stdout, engine, token, task_id, attempt, cancel.clone()).await?
919 }
920 StreamMode::LengthPrefixed => {
921 read_length_prefixed(stdout, engine, token, task_id, attempt, cancel.clone()).await?
922 }
923 };
924
925 let status = child
926 .wait()
927 .await
928 .map_err(|e| WorkerError::Failed(format!("streaming wait: {e}")))?;
929
930 match last_final {
931 Some((value, ok)) => Ok(WorkerResult {
932 value,
933 ok: ok && status.success(),
934 stats: Some(subprocess_base_stats(&status, None)),
937 }),
938 None => {
939 let value = serde_json::json!({
941 "raw": "",
942 "note": "streaming mode: no Final event received",
943 "exit_success": status.success(),
944 });
945 let _ = engine
946 .submit_output(
947 token,
948 task_id,
949 attempt,
950 OutputEvent::Final {
951 content: ContentRef::Inline {
952 value: value.clone(),
953 },
954 ok: false,
955 },
956 )
957 .await;
958 Ok(WorkerResult {
959 value,
960 ok: false,
961 stats: Some(subprocess_base_stats(&status, None)),
962 })
963 }
964 }
965}
966
967async fn forward_event(
970 engine: &Engine,
971 token: &CapToken,
972 task_id: &StepId,
973 attempt: u32,
974 ev: OutputEvent,
975 last_final: &mut Option<(Value, bool)>,
976) {
977 if let OutputEvent::Final { content, ok } = &ev {
978 let value = match content {
979 ContentRef::Inline { value } => value.clone(),
980 ContentRef::FileRef {
981 path,
982 mime,
983 size_hint,
984 } => serde_json::json!({
985 "file_ref": path.to_string_lossy(),
986 "mime": mime,
987 "size_hint": size_hint,
988 }),
989 };
990 *last_final = Some((value, *ok));
991 }
992 let _ = engine.submit_output(token, task_id, attempt, ev).await;
993}
994
995async fn read_ndjson(
998 stdout: tokio::process::ChildStdout,
999 engine: &Engine,
1000 token: &CapToken,
1001 task_id: &StepId,
1002 attempt: u32,
1003 cancel: CancellationToken,
1004) -> Result<Option<(Value, bool)>, WorkerError> {
1005 let mut reader = BufReader::new(stdout).lines();
1006 let mut last_final = None;
1007 loop {
1008 tokio::select! {
1009 line_res = reader.next_line() => match line_res {
1010 Ok(Some(line)) => {
1011 let trimmed = line.trim();
1012 if trimmed.is_empty() { continue; }
1013 if let Ok(ev) = serde_json::from_str::<OutputEvent>(trimmed) {
1014 forward_event(engine, token, task_id, attempt, ev, &mut last_final).await;
1015 }
1016 }
1017 Ok(None) => break,
1018 Err(e) => return Err(WorkerError::Failed(format!("ndjson read: {e}"))),
1019 },
1020 _ = cancel.cancelled() => return Err(WorkerError::Cancelled),
1021 }
1022 }
1023 Ok(last_final)
1024}
1025
1026async fn read_sse(
1030 stdout: tokio::process::ChildStdout,
1031 engine: &Engine,
1032 token: &CapToken,
1033 task_id: &StepId,
1034 attempt: u32,
1035 cancel: CancellationToken,
1036) -> Result<Option<(Value, bool)>, WorkerError> {
1037 let mut reader = BufReader::new(stdout).lines();
1038 let mut last_final = None;
1039 let mut data_buf = String::new();
1040 loop {
1041 tokio::select! {
1042 line_res = reader.next_line() => match line_res {
1043 Ok(Some(line)) => {
1044 if line.is_empty() {
1045 if !data_buf.is_empty() {
1047 if let Ok(ev) = serde_json::from_str::<OutputEvent>(data_buf.trim()) {
1048 forward_event(engine, token, task_id, attempt, ev, &mut last_final).await;
1049 }
1050 data_buf.clear();
1051 }
1052 } else if let Some(rest) = line.strip_prefix("data:") {
1053 let payload = rest.strip_prefix(' ').unwrap_or(rest);
1055 if !data_buf.is_empty() {
1056 data_buf.push('\n');
1057 }
1058 data_buf.push_str(payload);
1059 }
1060 }
1062 Ok(None) => {
1063 if !data_buf.is_empty() {
1065 if let Ok(ev) = serde_json::from_str::<OutputEvent>(data_buf.trim()) {
1066 forward_event(engine, token, task_id, attempt, ev, &mut last_final).await;
1067 }
1068 }
1069 break;
1070 }
1071 Err(e) => return Err(WorkerError::Failed(format!("sse read: {e}"))),
1072 },
1073 _ = cancel.cancelled() => return Err(WorkerError::Cancelled),
1074 }
1075 }
1076 Ok(last_final)
1077}
1078
1079async fn read_length_prefixed(
1082 mut stdout: tokio::process::ChildStdout,
1083 engine: &Engine,
1084 token: &CapToken,
1085 task_id: &StepId,
1086 attempt: u32,
1087 cancel: CancellationToken,
1088) -> Result<Option<(Value, bool)>, WorkerError> {
1089 use tokio::io::AsyncReadExt;
1090 let mut last_final = None;
1091 loop {
1092 let mut len_buf = [0u8; 4];
1094 let read_fut = stdout.read_exact(&mut len_buf);
1095 let read_res = tokio::select! {
1096 r = read_fut => r,
1097 _ = cancel.cancelled() => return Err(WorkerError::Cancelled),
1098 };
1099 match read_res {
1100 Ok(_) => {}
1101 Err(e) if e.kind() == std::io::ErrorKind::UnexpectedEof => break, Err(e) => return Err(WorkerError::Failed(format!("len read: {e}"))),
1103 }
1104 let len = u32::from_be_bytes(len_buf) as usize;
1105 if len == 0 || len > 16 * 1024 * 1024 {
1106 break;
1108 }
1109 let mut payload = vec![0u8; len];
1110 let read_fut = stdout.read_exact(&mut payload);
1111 let read_res = tokio::select! {
1112 r = read_fut => r,
1113 _ = cancel.cancelled() => return Err(WorkerError::Cancelled),
1114 };
1115 if read_res.is_err() {
1116 break;
1117 }
1118 if let Ok(ev) = serde_json::from_slice::<OutputEvent>(&payload) {
1119 forward_event(engine, token, task_id, attempt, ev, &mut last_final).await;
1120 }
1121 }
1122 Ok(last_final)
1123}
1124
1125#[cfg(test)]
1126mod embed_tests {
1127 use super::*;
1128
1129 fn vars() -> EmbedVars {
1130 EmbedVars {
1131 system: "SYS".to_string(),
1132 system_file: Some("/tmp/sys.md".to_string()),
1133 prompt: "do the task".to_string(),
1134 model: Some("small".to_string()),
1135 tools_csv: "Read,Grep".to_string(),
1136 work_dir: Some("/tmp/wd".to_string()),
1137 task_id: "ST-1".to_string(),
1138 attempt: "1".to_string(),
1139 }
1140 }
1141
1142 #[test]
1143 fn render_substitutes_every_closed_set_token() {
1144 let v = vars();
1145 let out = v
1146 .render("{system}|{system_file}|{prompt}|{model}|{tools_csv}|{work_dir}|{task_id}|{attempt}")
1147 .expect("render");
1148 assert_eq!(
1149 out,
1150 "SYS|/tmp/sys.md|do the task|small|Read,Grep|/tmp/wd|ST-1|1"
1151 );
1152 }
1153
1154 #[test]
1155 fn render_leaves_non_placeholder_braces_alone() {
1156 let v = vars();
1157 let out = v
1158 .render(r#"echo '{"result": "{prompt}"}'"#)
1159 .expect("render");
1160 assert_eq!(out, r#"echo '{"result": "do the task"}'"#);
1161 }
1162
1163 #[test]
1164 fn render_fails_loud_when_model_referenced_but_absent() {
1165 let mut v = vars();
1166 v.model = None;
1167 let err = v.render("--model {model}").unwrap_err();
1168 let msg = format!("{err:?}");
1169 assert!(msg.contains("{model}"), "actionable token in error: {msg}");
1170 assert!(msg.contains("profile.model"), "actionable cause: {msg}");
1171 }
1172
1173 #[test]
1174 fn render_fails_loud_when_work_dir_referenced_but_absent() {
1175 let mut v = vars();
1176 v.work_dir = None;
1177 let err = v.render("{work_dir}").unwrap_err();
1178 let msg = format!("{err:?}");
1179 assert!(msg.contains("{work_dir}"), "actionable token: {msg}");
1180 }
1181
1182 #[test]
1183 fn render_empty_system_is_legal() {
1184 let mut v = vars();
1185 v.system = String::new();
1186 assert_eq!(v.render("[{system}]").expect("render"), "[]");
1187 }
1188
1189 #[test]
1194 fn render_never_substitutes_inside_substituted_values() {
1195 let mut v = vars();
1196 v.prompt = "please mention {model} and {work_dir} literally".to_string();
1197 v.model = None; v.work_dir = None;
1199 let out = v.render("task: {prompt}").expect("render");
1200 assert_eq!(out, "task: please mention {model} and {work_dir} literally");
1201 }
1202
1203 #[test]
1206 fn render_substitutes_placeholder_nested_in_literal_braces() {
1207 let v = vars();
1208 let out = v.render(r#"{"task": "{prompt}", "n": 1}"#).expect("render");
1209 assert_eq!(out, r#"{"task": "do the task", "n": 1}"#);
1210 }
1211
1212 #[test]
1213 fn embed_references_scans_argv_stdin_env_cwd() {
1214 let mut t = EmbedTemplate {
1215 argv: vec!["cat".to_string()],
1216 ..Default::default()
1217 };
1218 assert!(!embed_references(&t, "{system_file}"));
1219 t.stdin = Some("{system_file}".to_string());
1220 assert!(embed_references(&t, "{system_file}"));
1221 t.stdin = None;
1222 t.env.insert("X".to_string(), "{system_file}".to_string());
1223 assert!(embed_references(&t, "{system_file}"));
1224 t.env.clear();
1225 t.cwd = Some("{system_file}".to_string());
1226 assert!(embed_references(&t, "{system_file}"));
1227 }
1228
1229 fn fake_output(stdout: &str, stderr: &str, code: i32) -> std::process::Output {
1230 #[cfg(unix)]
1231 use std::os::unix::process::ExitStatusExt;
1232 #[cfg(windows)]
1233 use std::os::windows::process::ExitStatusExt;
1234 use std::process::ExitStatus;
1235 #[cfg(unix)]
1236 let status = ExitStatus::from_raw(code << 8);
1237 #[cfg(windows)]
1238 let status = ExitStatus::from_raw(code as u32);
1239 std::process::Output {
1240 status,
1241 stdout: stdout.as_bytes().to_vec(),
1242 stderr: stderr.as_bytes().to_vec(),
1243 }
1244 }
1245
1246 #[test]
1250 fn normalize_without_decl_is_historical_json_or_raw() {
1251 let out = fake_output(r#"{"a": 1}"#, "", 0);
1252 let wr = normalize_plain_output(&out, None, None);
1253 assert_eq!(wr.value, serde_json::json!({"a": 1}));
1254 assert!(wr.ok);
1255
1256 let out = fake_output("plain text\n", "warned", 0);
1257 let wr = normalize_plain_output(&out, None, None);
1258 assert_eq!(
1259 wr.value,
1260 serde_json::json!({"raw": "plain text", "stderr": "warned"})
1261 );
1262 assert!(wr.ok);
1263
1264 let out = fake_output("boom", "", 1);
1265 let wr = normalize_plain_output(&out, None, None);
1266 assert!(!wr.ok, "non-zero exit is a failed step");
1267 }
1268
1269 #[test]
1270 fn normalize_result_ptr_extracts_declared_value() {
1271 let decl = SubprocessOutput {
1272 format: Some("json".to_string()),
1273 result_ptr: Some("/result".to_string()),
1274 ok_from: Some("exit_code".to_string()),
1275 stats: None,
1276 };
1277 let out = fake_output(r#"{"result": {"answer": 42}, "noise": true}"#, "", 0);
1278 let wr = normalize_plain_output(&out, Some(&decl), None);
1279 assert_eq!(wr.value, serde_json::json!({"answer": 42}));
1280 assert!(wr.ok);
1281 }
1282
1283 #[test]
1284 fn normalize_declared_json_unparsable_stdout_fails_loud() {
1285 let decl = SubprocessOutput {
1286 format: Some("json".to_string()),
1287 result_ptr: None,
1288 ok_from: None,
1289 stats: None,
1290 };
1291 let out = fake_output("not json at all", "stderr text", 0);
1292 let wr = normalize_plain_output(&out, Some(&decl), None);
1293 assert!(!wr.ok, "declared-JSON unparsable stdout is a failed step");
1294 assert_eq!(wr.value["raw"], "not json at all");
1295 assert_eq!(wr.value["stderr"], "stderr text");
1296 assert!(wr.value["parse_error"].is_string());
1297 }
1298
1299 #[test]
1300 fn normalize_missing_result_ptr_fails_loud_with_actionable_value() {
1301 let decl = SubprocessOutput {
1302 format: Some("json".to_string()),
1303 result_ptr: Some("/missing".to_string()),
1304 ok_from: None,
1305 stats: None,
1306 };
1307 let out = fake_output(r#"{"present": 1}"#, "", 0);
1308 let wr = normalize_plain_output(&out, Some(&decl), None);
1309 assert!(!wr.ok);
1310 assert!(wr.value["error"]
1311 .as_str()
1312 .expect("actionable error message")
1313 .contains("/missing"));
1314 }
1315
1316 #[test]
1317 fn normalize_ok_from_pointer_reads_boolean() {
1318 let decl = SubprocessOutput {
1319 format: Some("json".to_string()),
1320 result_ptr: Some("/result".to_string()),
1321 ok_from: Some("/ok".to_string()),
1322 stats: None,
1323 };
1324 let out = fake_output(r#"{"result": "r", "ok": true}"#, "", 0);
1326 let wr = normalize_plain_output(&out, Some(&decl), None);
1327 assert!(wr.ok);
1328 let out = fake_output(r#"{"result": "r", "ok": false}"#, "", 0);
1330 let wr = normalize_plain_output(&out, Some(&decl), None);
1331 assert!(!wr.ok);
1332 let out = fake_output(r#"{"result": "r", "ok": "yes"}"#, "", 0);
1334 let wr = normalize_plain_output(&out, Some(&decl), None);
1335 assert!(!wr.ok);
1336 }
1337}