1use std::sync::Arc;
16use std::sync::atomic::{AtomicBool, Ordering};
17use std::time::Duration;
18
19use tokio::sync::Semaphore;
20use tracing::Instrument;
21
22use crate::ast::{Command, Value};
23use crate::dispatch::CommandDispatcher;
24use crate::duration::parse_duration;
25use crate::interpreter::ExecResult;
26use crate::tools::{ExecContext, ToolRegistry};
27
28use super::pipeline::PipelineRunner;
29
30#[derive(Debug, Clone)]
32pub struct ScatterOptions {
33 pub var_name: String,
35 pub limit: usize,
37 pub timeout: Option<Duration>,
41}
42
43#[derive(Debug, Clone)]
45pub struct GatherOptions {
46 pub progress: bool,
48 pub first: usize,
50 pub format: String,
52}
53
54impl Default for ScatterOptions {
55 fn default() -> Self {
56 Self {
57 var_name: "ITEM".to_string(),
58 limit: 8,
59 timeout: None,
60 }
61 }
62}
63
64impl Default for GatherOptions {
65 fn default() -> Self {
66 Self {
67 progress: false,
68 first: 0,
69 format: "lines".to_string(),
70 }
71 }
72}
73
74#[derive(Debug, Clone)]
76pub struct ScatterResult {
77 pub item: String,
79 pub result: ExecResult,
81 pub timed_out: bool,
83}
84
85pub struct ScatterGatherRunner {
92 tools: Arc<ToolRegistry>,
93 sequential_dispatcher: Arc<dyn CommandDispatcher>,
96}
97
98impl ScatterGatherRunner {
99 pub fn new(
104 tools: Arc<ToolRegistry>,
105 dispatcher: Arc<dyn CommandDispatcher>,
106 ) -> Self {
107 Self { tools, sequential_dispatcher: dispatcher }
108 }
109
110 #[tracing::instrument(level = "info", skip(self, pre_scatter, scatter_opts, parallel, gather_opts, post_gather, ctx), fields(item_count = tracing::field::Empty, parallelism = scatter_opts.limit))]
119 pub async fn run(
120 &self,
121 pre_scatter: &[Command],
122 scatter_opts: ScatterOptions,
123 parallel: &[Command],
124 gather_opts: GatherOptions,
125 post_gather: &[Command],
126 ctx: &mut ExecContext,
127 ) -> ExecResult {
128 let runner = PipelineRunner::new(self.tools.clone());
129
130 let (text, data) = if pre_scatter.is_empty() {
133 let data = ctx.take_stdin_data();
137 let text = match ctx.read_stdin_to_text().await {
138 Ok(s) => s.unwrap_or_default(),
139 Err(e) => return ExecResult::failure(2, format!("scatter: {e}")),
140 };
141 (text, data)
142 } else {
143 let result = runner.run_sequential(pre_scatter, ctx, &*self.sequential_dispatcher).await;
144 if !result.ok() {
145 return result;
146 }
147 (result.text_out().into_owned(), result.data)
148 };
149
150 let items = match extract_items(data.as_ref(), &text) {
152 Ok(items) => items,
153 Err(msg) => return ExecResult::failure(1, msg),
154 };
155 if items.is_empty() {
156 return ExecResult::success("");
157 }
158
159 tracing::Span::current().record("item_count", items.len());
160
161 let results = self
163 .run_parallel(&items, &scatter_opts, parallel, ctx)
164 .await;
165
166 let GatherOutput {
168 text: gathered,
169 dropped_failures,
170 } = gather_results(&results, &gather_opts);
171
172 if !dropped_failures.is_empty() {
178 let err = format!(
179 "gather: {} task(s) failed and were omitted from line output: {} (use --json to capture per-task status)",
180 dropped_failures.len(),
181 dropped_failures.join(", ")
182 );
183 return ExecResult::from_output(1, gathered, err);
184 }
185
186 if post_gather.is_empty() {
188 ExecResult::success(gathered)
189 } else {
190 ctx.set_stdin(gathered);
191 runner.run_sequential(post_gather, ctx, &*self.sequential_dispatcher).await
192 }
193 }
194
195 #[tracing::instrument(level = "debug", skip(self, items, opts, commands, base_ctx), fields(worker_count = items.len()))]
204 async fn run_parallel(
205 &self,
206 items: &[String],
207 opts: &ScatterOptions,
208 commands: &[Command],
209 base_ctx: &ExecContext,
210 ) -> Vec<ScatterResult> {
211 let semaphore = Arc::new(Semaphore::new(opts.limit));
212 let tools = self.tools.clone();
213 let var_name = opts.var_name.clone();
214
215 let mut handles = Vec::with_capacity(items.len());
217
218 for item in items.iter().cloned() {
219 let permit = semaphore.clone().acquire_owned().await;
220 let tools = tools.clone();
221 let worker_dispatcher = self.sequential_dispatcher.fork_attached().await;
226 let commands = commands.to_vec();
227 let var_name = var_name.clone();
228 let base_scope = base_ctx.scope.clone();
229 let backend = base_ctx.backend.clone();
230 let cwd = base_ctx.cwd.clone();
231 let parent_token = base_ctx.cancel.clone();
232 let worker_token = parent_token.child_token();
233
234 let timed_out_flag = Arc::new(AtomicBool::new(false));
240 let timer_handle: Option<tokio::task::JoinHandle<()>> = opts.timeout.map(|d| {
241 let cancel = worker_token.clone();
242 let flag = timed_out_flag.clone();
243 tokio::spawn(async move {
244 tokio::time::sleep(d).await;
245 flag.store(true, Ordering::SeqCst);
246 cancel.cancel();
247 })
248 });
249 let timed_out_check = timed_out_flag.clone();
250
251 let item_label = if item.len() > 64 {
252 format!("{}...", &item[..64])
253 } else {
254 item.clone()
255 };
256 let worker_span = tracing::debug_span!("scatter_worker", item = %item_label);
257 let handle = tokio::spawn(crate::telemetry::bind_current_context(async move {
261 let _permit = permit; let mut scope = base_scope;
265 scope.set(&var_name, Value::String(item.clone()));
266
267 let mut ctx = ExecContext::with_backend_and_scope(backend, scope);
268 ctx.set_cwd(cwd);
269 ctx.cancel = worker_token;
270
271 let runner = PipelineRunner::new(tools);
274 let result = runner.run_sequential(&commands, &mut ctx, &*worker_dispatcher).await;
275
276 if let Some(h) = timer_handle {
279 h.abort();
280 }
281
282 let timed_out = timed_out_check.load(Ordering::SeqCst);
283 ScatterResult { item, result, timed_out }
284 }.instrument(worker_span)));
285
286 handles.push(handle);
287 }
288
289 let mut results = Vec::with_capacity(handles.len());
291 for handle in handles {
292 match handle.await {
293 Ok(result) => results.push(result),
294 Err(e) => {
295 results.push(ScatterResult {
296 item: String::new(),
297 result: ExecResult::failure(1, format!("Task panicked: {}", e)),
298 timed_out: false,
299 });
300 }
301 }
302 }
303
304 results
305 }
306}
307
308pub fn extract_items(data: Option<&Value>, text: &str) -> Result<Vec<String>, String> {
317 match data {
319 Some(Value::Json(serde_json::Value::Array(arr))) => {
321 return Ok(arr.iter().map(|v| match v {
322 serde_json::Value::String(s) => s.clone(),
323 other => other.to_string(),
324 }).collect());
325 }
326 Some(Value::String(s)) => return Ok(vec![s.clone()]),
328 Some(Value::Int(i)) => return Ok(vec![i.to_string()]),
330 Some(Value::Float(f)) => return Ok(vec![f.to_string()]),
331 Some(Value::Bool(b)) => return Ok(vec![b.to_string()]),
332 Some(Value::Null) => return Ok(vec!["null".to_string()]),
333 Some(Value::Json(json)) => return Ok(vec![json.to_string()]),
338 Some(Value::Bytes(b)) => return Ok(vec![format!("[binary: {} bytes]", b.len())]),
340 None => {}
342 }
343
344 let trimmed = text.trim_end_matches(['\n', '\r']);
346 if trimmed.is_empty() {
347 return Ok(vec![]);
348 }
349 Ok(trimmed
350 .split('\n')
351 .map(|line| line.trim_end_matches('\r').to_string())
352 .collect())
353}
354
355struct GatherOutput {
358 text: String,
359 dropped_failures: Vec<String>,
363}
364
365fn gather_results(results: &[ScatterResult], opts: &GatherOptions) -> GatherOutput {
373 let results_to_use = if opts.first > 0 && opts.first < results.len() {
374 &results[..opts.first]
375 } else {
376 results
377 };
378
379 if opts.format == "json" {
380 let json_results: Vec<serde_json::Value> = results_to_use
382 .iter()
383 .map(|r| {
384 serde_json::json!({
385 "item": r.item,
386 "ok": r.result.ok(),
387 "code": r.result.code,
388 "out": r.result.text_out().trim(),
389 "err": r.result.err.trim(),
390 "timed_out": r.timed_out,
391 })
392 })
393 .collect();
394
395 GatherOutput {
396 text: serde_json::to_string_pretty(&json_results).unwrap_or_default(),
397 dropped_failures: Vec::new(),
398 }
399 } else {
400 let text = results_to_use
405 .iter()
406 .filter(|r| r.result.ok())
407 .map(|r| r.result.text_out())
408 .map(|t| t.trim().to_string())
409 .collect::<Vec<_>>()
410 .join("\n");
411 let dropped_failures = results_to_use
412 .iter()
413 .filter(|r| !r.result.ok())
414 .map(|r| r.item.clone())
415 .collect();
416 GatherOutput {
417 text,
418 dropped_failures,
419 }
420 }
421}
422
423pub fn parse_scatter_options(args: &crate::tools::ToolArgs) -> ScatterOptions {
425 let mut opts = ScatterOptions::default();
426
427 if let Some(Value::String(name)) = args.named.get("as") {
428 opts.var_name = name.clone();
429 }
430
431 if let Some(Value::Int(n)) = args.named.get("limit") {
432 let requested = *n;
433 let clamped = requested.clamp(1, SCATTER_LIMIT_MAX as i64);
434 if requested > SCATTER_LIMIT_MAX as i64 {
435 tracing::warn!(
436 target: "kaish::scatter",
437 requested = requested,
438 ceiling = SCATTER_LIMIT_MAX,
439 "scatter limit clamped to ceiling"
440 );
441 }
442 opts.limit = clamped as usize;
443 }
444
445 if let Some(Value::String(s)) = args.named.get("timeout") {
449 match parse_duration(s) {
450 Some(d) => opts.timeout = Some(d),
451 None => tracing::warn!(
452 target: "kaish::scatter",
453 value = %s,
454 "scatter --timeout: invalid duration (try: 30, 5s, 500ms, 2m, 1h)"
455 ),
456 }
457 } else if let Some(Value::Int(n)) = args.named.get("timeout") {
458 if *n >= 0 {
459 opts.timeout = Some(Duration::from_secs(*n as u64));
460 }
461 }
462
463 opts
464}
465
466pub const SCATTER_LIMIT_MAX: usize = 10_000;
470
471pub fn parse_gather_options(args: &crate::tools::ToolArgs) -> GatherOptions {
473 let mut opts = GatherOptions::default();
474
475 if args.has_flag("progress") {
476 opts.progress = true;
477 }
478
479 if let Some(Value::Int(n)) = args.named.get("first") {
480 opts.first = (*n).max(0) as usize;
481 }
482
483 if let Some(Value::String(fmt)) = args.named.get("format") {
484 opts.format = fmt.clone();
485 }
486
487 opts
488}
489
490#[cfg(test)]
491mod tests {
492 use super::*;
493
494 #[test]
495 fn test_extract_items_structured_json_array() {
496 let data = Value::Json(serde_json::json!(["a", "b", "c"]));
497 let items = extract_items(Some(&data), "").unwrap();
498 assert_eq!(items, vec!["a", "b", "c"]);
499 }
500
501 #[test]
502 fn test_extract_items_structured_mixed_types() {
503 let data = Value::Json(serde_json::json!([1, "two", true]));
504 let items = extract_items(Some(&data), "").unwrap();
505 assert_eq!(items, vec!["1", "two", "true"]);
506 }
507
508 #[test]
509 fn test_extract_items_structured_string() {
510 let data = Value::String("single".into());
511 let items = extract_items(Some(&data), "").unwrap();
512 assert_eq!(items, vec!["single"]);
513 }
514
515 #[test]
516 fn test_extract_items_single_line_text() {
517 let items = extract_items(None, "hello").unwrap();
518 assert_eq!(items, vec!["hello"]);
519 }
520
521 #[test]
522 fn test_extract_items_empty() {
523 let items = extract_items(None, "").unwrap();
524 assert!(items.is_empty());
525 }
526
527 #[test]
528 fn test_extract_items_multiline_fans_out_per_line() {
529 let items = extract_items(None, "one\ntwo\nthree").unwrap();
532 assert_eq!(items, vec!["one", "two", "three"]);
533 }
534
535 #[test]
536 fn test_extract_items_trailing_newline_no_phantom_item() {
537 let items = extract_items(None, "one\ntwo\n").unwrap();
539 assert_eq!(items, vec!["one", "two"]);
540 }
541
542 #[test]
543 fn test_extract_items_crlf_per_line() {
544 let items = extract_items(None, "one\r\ntwo\r\n").unwrap();
546 assert_eq!(items, vec!["one", "two"]);
547 }
548
549 #[test]
550 fn test_extract_items_interior_blank_line_preserved() {
551 let items = extract_items(None, "a\n\nb").unwrap();
553 assert_eq!(items, vec!["a", "", "b"]);
554 }
555
556 #[test]
557 fn test_extract_items_whitespace_within_line_not_split() {
558 let items = extract_items(None, "a b\nc d").unwrap();
560 assert_eq!(items, vec!["a b", "c d"]);
561 }
562
563 #[test]
564 fn test_extract_items_only_newlines_is_empty() {
565 let items = extract_items(None, "\n\n").unwrap();
566 assert!(items.is_empty());
567 }
568
569 #[test]
570 fn test_extract_items_structured_overrides_text() {
571 let data = Value::Json(serde_json::json!(["x", "y"]));
573 let items = extract_items(Some(&data), "ignored\ntext").unwrap();
574 assert_eq!(items, vec!["x", "y"]);
575 }
576
577 #[test]
578 fn test_gather_results_lines() {
579 let results = vec![
580 ScatterResult {
581 item: "a".to_string(),
582 result: ExecResult::success("result_a"),
583 timed_out: false,
584 },
585 ScatterResult {
586 item: "b".to_string(),
587 result: ExecResult::success("result_b"),
588 timed_out: false,
589 },
590 ];
591
592 let opts = GatherOptions::default();
593 let output = gather_results(&results, &opts);
594 assert_eq!(output.text, "result_a\nresult_b");
595 assert!(output.dropped_failures.is_empty());
596 }
597
598 #[test]
599 fn test_gather_results_lines_reports_dropped_failures() {
600 let results = vec![
603 ScatterResult {
604 item: "a".to_string(),
605 result: ExecResult::success("result_a"),
606 timed_out: false,
607 },
608 ScatterResult {
609 item: "b".to_string(),
610 result: ExecResult::failure(1, "boom"),
611 timed_out: false,
612 },
613 ];
614
615 let opts = GatherOptions::default();
616 let output = gather_results(&results, &opts);
617 assert_eq!(output.text, "result_a");
619 assert_eq!(output.dropped_failures, vec!["b".to_string()]);
620 }
621
622 #[test]
623 fn test_gather_results_json_keeps_failures_as_rows() {
624 let results = vec![ScatterResult {
626 item: "b".to_string(),
627 result: ExecResult::failure(2, "boom"),
628 timed_out: false,
629 }];
630 let opts = GatherOptions {
631 format: "json".to_string(),
632 ..Default::default()
633 };
634 let output = gather_results(&results, &opts);
635 assert!(output.dropped_failures.is_empty());
636 assert!(output.text.contains("\"ok\": false"));
637 assert!(output.text.contains("\"code\": 2"));
638 }
639
640 #[test]
641 fn test_gather_results_json() {
642 let results = vec![ScatterResult {
643 item: "test".to_string(),
644 result: ExecResult::success("output"),
645 timed_out: false,
646 }];
647
648 let opts = GatherOptions {
649 format: "json".to_string(),
650 ..Default::default()
651 };
652 let output = gather_results(&results, &opts);
653 assert!(output.text.contains("\"item\": \"test\""));
654 assert!(output.text.contains("\"ok\": true"));
655 }
656
657 #[test]
658 fn test_gather_results_first_n() {
659 let results = vec![
660 ScatterResult {
661 item: "a".to_string(),
662 result: ExecResult::success("1"),
663 timed_out: false,
664 },
665 ScatterResult {
666 item: "b".to_string(),
667 result: ExecResult::success("2"),
668 timed_out: false,
669 },
670 ScatterResult {
671 item: "c".to_string(),
672 result: ExecResult::success("3"),
673 timed_out: false,
674 },
675 ];
676
677 let opts = GatherOptions {
678 first: 2,
679 ..Default::default()
680 };
681 let output = gather_results(&results, &opts);
682 assert_eq!(output.text, "1\n2");
683 }
684
685 #[test]
686 fn test_parse_scatter_options() {
687 use crate::tools::ToolArgs;
688
689 let mut args = ToolArgs::new();
690 args.named.insert("as".to_string(), Value::String("URL".to_string()));
691 args.named.insert("limit".to_string(), Value::Int(4));
692
693 let opts = parse_scatter_options(&args);
694 assert_eq!(opts.var_name, "URL");
695 assert_eq!(opts.limit, 4);
696 }
697
698 #[test]
699 fn test_parse_gather_options() {
700 use crate::tools::ToolArgs;
701
702 let mut args = ToolArgs::new();
703 args.named.insert("first".to_string(), Value::Int(5));
704 args.named.insert("format".to_string(), Value::String("json".to_string()));
705
706 let opts = parse_gather_options(&args);
707 assert_eq!(opts.first, 5);
708 assert_eq!(opts.format, "json");
709 }
710
711 #[test]
712 fn scatter_limit_clamps_to_ceiling() {
713 use crate::tools::ToolArgs;
714
715 let mut args = ToolArgs::new();
716 args.named.insert("limit".to_string(), Value::Int(999_999));
717 let opts = parse_scatter_options(&args);
718 assert_eq!(opts.limit, SCATTER_LIMIT_MAX);
719 }
720
721 #[test]
722 fn scatter_limit_raises_zero_to_one() {
723 use crate::tools::ToolArgs;
724
725 let mut args = ToolArgs::new();
726 args.named.insert("limit".to_string(), Value::Int(0));
727 let opts = parse_scatter_options(&args);
728 assert_eq!(opts.limit, 1);
729 }
730
731 #[test]
732 fn scatter_limit_raises_negative_to_one() {
733 use crate::tools::ToolArgs;
734
735 let mut args = ToolArgs::new();
736 args.named.insert("limit".to_string(), Value::Int(-42));
737 let opts = parse_scatter_options(&args);
738 assert_eq!(opts.limit, 1);
739 }
740
741 #[test]
742 fn scatter_limit_preserves_valid_values() {
743 use crate::tools::ToolArgs;
744
745 let mut args = ToolArgs::new();
746 args.named.insert("limit".to_string(), Value::Int(500));
747 let opts = parse_scatter_options(&args);
748 assert_eq!(opts.limit, 500);
749 }
750}