1use crate::child_turn;
4use crate::session_runner::AgentSession;
5use crate::subagents::{AgentStatus, ForkTurns};
6use std::sync::atomic::{AtomicU64, Ordering};
7use std::sync::{Arc, Mutex, Weak};
8use std::time::{Duration, Instant};
9use tokio::sync::{Notify, Semaphore, watch};
10use tokio_util::sync::CancellationToken;
11
12pub type JobId = u64;
13
14const MAX_ACTIVE_JOBS: usize = 4;
15const MAX_RETAINED_JOBS: usize = 20;
16pub const MAX_ITERATIONS: u32 = 100;
17
18pub fn format_interval(interval: Duration) -> String {
19 let mut nanos = interval.as_nanos();
20 let units = [
21 (86_400_000_000_000_u128, "d"),
22 (3_600_000_000_000, "h"),
23 (60_000_000_000, "m"),
24 (1_000_000_000, "s"),
25 (1_000_000, "ms"),
26 (1_000, "us"),
27 (1, "ns"),
28 ];
29 let mut output = String::new();
30 for (size, suffix) in units {
31 let count = nanos / size;
32 if count > 0 {
33 output.push_str(&format!("{count}{suffix}"));
34 nanos %= size;
35 }
36 }
37 if output.is_empty() {
38 output.push_str("0s");
39 }
40 output
41}
42
43#[derive(Debug, Clone, Copy, PartialEq, Eq)]
44pub enum JobKind {
45 Loop,
46 Autoresearch,
47}
48
49impl JobKind {
50 pub fn label(self) -> &'static str {
51 match self {
52 Self::Loop => "loop",
53 Self::Autoresearch => "autoresearch",
54 }
55 }
56}
57
58#[derive(Debug, Clone, Copy, PartialEq, Eq)]
59pub enum JobStatus {
60 Queued,
61 Running,
62 Paused,
63 Completed,
64 Failed,
65 Stopped,
66}
67
68impl JobStatus {
69 pub fn label(self) -> &'static str {
70 match self {
71 Self::Queued => "queued",
72 Self::Running => "running",
73 Self::Paused => "paused",
74 Self::Completed => "completed",
75 Self::Failed => "failed",
76 Self::Stopped => "stopped",
77 }
78 }
79
80 pub fn is_finished(self) -> bool {
81 matches!(self, Self::Completed | Self::Failed | Self::Stopped)
82 }
83}
84
85#[derive(Debug, Clone)]
86pub struct JobSnapshot {
87 pub id: JobId,
88 pub kind: JobKind,
89 pub goal: String,
90 pub status: JobStatus,
91 pub iteration: u32,
92 pub limit: Option<u32>,
93 pub interval: Option<Duration>,
94 pub tokens: u64,
95 pub elapsed: Duration,
96 pub session_id: String,
97 pub latest_result: Option<String>,
98 pub error: Option<String>,
99}
100
101#[derive(Debug, Clone)]
102pub struct JobSummary {
103 pub id: JobId,
104 pub kind: JobKind,
105 pub goal: String,
106 pub status: JobStatus,
107 pub iteration: u32,
108 pub limit: Option<u32>,
109 pub interval: Option<Duration>,
110 pub tokens: u64,
111 pub elapsed: Duration,
112}
113
114struct JobState {
115 status: JobStatus,
116 iteration: u32,
117 tokens: u64,
118 latest_result: Option<String>,
119 error: Option<String>,
120}
121
122pub struct JobRecord {
123 id: JobId,
124 kind: JobKind,
125 goal: String,
126 limit: Option<u32>,
127 interval: Option<Duration>,
128 session_id: String,
129 child: Arc<AgentSession>,
130 started: Instant,
131 state: Mutex<JobState>,
132 cancel: CancellationToken,
133 wake: Notify,
134 version: watch::Sender<u64>,
135 parent: Weak<AgentSession>,
136}
137
138impl JobRecord {
139 pub fn snapshot(&self) -> JobSnapshot {
140 let state = self.state.lock().unwrap();
141 JobSnapshot {
142 id: self.id,
143 kind: self.kind,
144 goal: self.goal.clone(),
145 status: state.status,
146 iteration: state.iteration,
147 limit: self.limit,
148 interval: self.interval,
149 tokens: state.tokens,
150 elapsed: self.started.elapsed(),
151 session_id: self.session_id.clone(),
152 latest_result: state.latest_result.clone(),
153 error: state.error.clone(),
154 }
155 }
156
157 fn summary(&self) -> JobSummary {
158 let state = self.state.lock().unwrap();
159 JobSummary {
160 id: self.id,
161 kind: self.kind,
162 goal: self.goal.clone(),
163 status: state.status,
164 iteration: state.iteration,
165 limit: self.limit,
166 interval: self.interval,
167 tokens: state.tokens,
168 elapsed: self.started.elapsed(),
169 }
170 }
171
172 pub fn subscribe(&self) -> watch::Receiver<u64> {
173 self.version.subscribe()
174 }
175
176 pub fn pause(&self) {
177 let mut state = self.state.lock().unwrap();
178 if matches!(state.status, JobStatus::Running | JobStatus::Queued) {
179 state.status = JobStatus::Paused;
180 drop(state);
181 self.signal();
182 }
183 }
184
185 pub fn resume(&self) {
186 let mut state = self.state.lock().unwrap();
187 if state.status == JobStatus::Paused {
188 state.status = JobStatus::Running;
189 drop(state);
190 self.wake.notify_waiters();
191 self.signal();
192 }
193 }
194
195 pub fn is_paused(&self) -> bool {
196 self.state.lock().unwrap().status == JobStatus::Paused
197 }
198
199 pub fn stop(&self) {
200 if self.state.lock().unwrap().status.is_finished() {
201 return;
202 }
203 self.cancel.cancel();
204 self.child.abort();
205 self.wake.notify_waiters();
206 }
207
208 fn signal(&self) {
209 self.version
210 .send_modify(|version| *version = version.wrapping_add(1));
211 if let Some(parent) = self.parent.upgrade() {
212 parent.emit_iterative(self.id, *self.version.borrow());
213 }
214 }
215
216 fn mark_running(&self) {
218 let mut state = self.state.lock().unwrap();
219 if state.status == JobStatus::Queued {
220 state.status = JobStatus::Running;
221 drop(state);
222 self.signal();
223 }
224 }
225
226 fn finish(&self, status: JobStatus) {
228 self.state.lock().unwrap().status = status;
229 self.signal();
230 }
231
232 async fn wait_until_ready(&self) -> bool {
233 loop {
234 let resumed = self.wake.notified();
237 if self.cancel.is_cancelled() {
238 return false;
239 }
240 if !self.is_paused() {
241 return true;
242 }
243 tokio::select! {
244 _ = self.cancel.cancelled() => return false,
245 _ = resumed => {}
246 }
247 }
248 }
249}
250
251pub struct IterativeRuntime {
252 parent: Weak<AgentSession>,
253 jobs: Mutex<Vec<Arc<JobRecord>>>,
254 next_id: AtomicU64,
255 permits: Arc<Semaphore>,
256}
257
258impl IterativeRuntime {
259 pub(crate) fn new(parent: Weak<AgentSession>) -> Arc<Self> {
260 Arc::new(Self {
261 parent,
262 jobs: Mutex::new(Vec::new()),
263 next_id: AtomicU64::new(1),
264 permits: Arc::new(Semaphore::new(MAX_ACTIVE_JOBS)),
265 })
266 }
267
268 pub fn start(
269 self: &Arc<Self>,
270 kind: JobKind,
271 request: JobRequest,
272 ) -> anyhow::Result<Arc<JobRecord>> {
273 let JobRequest {
274 goal,
275 limit,
276 interval,
277 } = request;
278 let parent = self
279 .parent
280 .upgrade()
281 .ok_or_else(|| anyhow::anyhow!("the parent session has closed"))?;
282 let id = self.next_id.fetch_add(1, Ordering::SeqCst);
283 let limit = limit.map(|limit| limit.clamp(1, MAX_ITERATIONS));
284 let task_name = format!("{}_{}", kind.label(), id);
285 let child = parent.create_subagent_session(
286 &task_name,
287 &format!("/root/{task_name}"),
288 ForkTurns::All,
289 None,
290 None,
291 )?;
292 let session_id = child.manager.lock().unwrap().session_id().to_string();
293 let (version, _) = watch::channel(0);
294 let record = Arc::new(JobRecord {
295 id,
296 kind,
297 goal,
298 limit,
299 interval,
300 session_id,
301 child,
302 started: Instant::now(),
303 state: Mutex::new(JobState {
304 status: JobStatus::Queued,
305 iteration: 0,
306 tokens: 0,
307 latest_result: None,
308 error: None,
309 }),
310 cancel: CancellationToken::new(),
311 wake: Notify::new(),
312 version,
313 parent: self.parent.clone(),
314 });
315
316 let mut jobs = self
317 .jobs
318 .lock()
319 .map_err(|_| anyhow::anyhow!("the job list is unavailable"))?;
320 jobs.push(record.clone());
321 while jobs.len() > MAX_RETAINED_JOBS {
322 let Some(position) = jobs
323 .iter()
324 .position(|job| job.state.lock().unwrap().status.is_finished())
325 else {
326 break;
327 };
328 jobs.remove(position);
329 }
330 drop(jobs);
331
332 let runtime = self.clone();
333 let running = record.clone();
334 tokio::spawn(async move { runtime.run(running).await });
335 record.signal();
336 Ok(record)
337 }
338
339 async fn run(self: Arc<Self>, job: Arc<JobRecord>) {
340 let permit = tokio::select! {
341 _ = job.cancel.cancelled() => None,
342 permit = self.permits.clone().acquire_owned() => permit.ok(),
343 };
344 let Some(_permit) = permit else {
345 job.finish(JobStatus::Stopped);
346 return;
347 };
348 if !job.wait_until_ready().await {
349 job.finish(JobStatus::Stopped);
350 return;
351 }
352 job.mark_running();
353
354 let mut iteration = 1_u32;
355 loop {
356 if !job.wait_until_ready().await {
357 job.finish(JobStatus::Stopped);
358 return;
359 }
360 job.state.lock().unwrap().iteration = iteration;
361 job.signal();
362
363 let outcome = child_turn::run_child_turn(
364 &job.parent,
365 &job.child,
366 iteration_prompt(job.kind, &job.goal, iteration, job.limit),
367 Some(job.cancel.clone()),
368 None,
369 )
370 .await;
371 {
372 let mut state = job.state.lock().unwrap();
373 state.tokens = state.tokens.saturating_add(outcome.usage.total_tokens);
374 state.latest_result = outcome.result.clone();
375 state.error = outcome.error.clone();
376 }
377 job.signal();
378
379 if job.cancel.is_cancelled() || outcome.status == AgentStatus::Interrupted {
380 job.finish(JobStatus::Stopped);
381 return;
382 }
383 if outcome.status == AgentStatus::Failed {
384 job.finish(JobStatus::Failed);
385 return;
386 }
387 if outcome
388 .result
389 .as_deref()
390 .is_some_and(contains_completion_marker)
391 {
392 job.finish(JobStatus::Completed);
393 return;
394 }
395 if job.limit == Some(iteration) {
396 job.finish(JobStatus::Completed);
397 return;
398 }
399 if let Some(interval) = job.interval {
400 tokio::select! {
401 _ = job.cancel.cancelled() => {
402 job.finish(JobStatus::Stopped);
403 return;
404 }
405 _ = tokio::time::sleep(interval) => {}
406 }
407 }
408 iteration = iteration.saturating_add(1);
409 }
410 }
411
412 pub fn get(&self, id: JobId) -> Option<Arc<JobRecord>> {
413 self.jobs
414 .lock()
415 .ok()?
416 .iter()
417 .find(|job| job.id == id)
418 .cloned()
419 }
420
421 pub fn summaries(&self) -> Vec<JobSummary> {
422 self.jobs
423 .lock()
424 .map(|jobs| jobs.iter().map(|job| job.summary()).collect())
425 .unwrap_or_default()
426 }
427
428 pub fn has_queued_or_running(&self) -> bool {
429 self.jobs.lock().is_ok_and(|jobs| {
430 jobs.iter().any(|job| {
431 matches!(
432 job.state.lock().unwrap().status,
433 JobStatus::Queued | JobStatus::Running
434 )
435 })
436 })
437 }
438
439 pub fn active_count(&self) -> usize {
440 self.jobs
441 .lock()
442 .map(|jobs| {
443 jobs.iter()
444 .filter(|job| !job.state.lock().unwrap().status.is_finished())
445 .count()
446 })
447 .unwrap_or(0)
448 }
449
450 pub fn latest(&self) -> Option<Arc<JobRecord>> {
451 self.jobs.lock().ok()?.last().cloned()
452 }
453
454 pub(crate) fn stop_all(&self) {
455 if let Ok(jobs) = self.jobs.lock() {
456 for job in jobs.iter() {
457 job.stop();
458 }
459 }
460 }
461}
462
463#[derive(Debug, PartialEq)]
467pub struct JobRequest {
468 goal: String,
469 limit: Option<u32>,
470 interval: Option<Duration>,
471}
472
473impl JobRequest {
474 pub fn parse(kind: JobKind, arguments: &str) -> Result<JobRequest, String> {
475 let arguments = arguments.trim();
476 let (goal, limit) = if let Some((goal, value)) = arguments.rsplit_once("--iterations") {
477 let value = value.trim();
478 if value.is_empty() || value.split_whitespace().count() != 1 {
479 return Err("use --iterations N at the end of the job goal".into());
480 }
481 let limit = value
482 .parse::<u32>()
483 .ok()
484 .filter(|limit| *limit > 0)
485 .ok_or_else(|| "iterations must be a positive number".to_string())?;
486 (goal.trim(), Some(limit))
487 } else {
488 (arguments, None)
489 };
490
491 let mut words = goal.splitn(2, char::is_whitespace);
492 let first = words.next().unwrap_or_default();
493 let parsed_interval = parse_interval(first)?;
494 if kind == JobKind::Autoresearch && parsed_interval.is_some() {
495 return Err("autoresearch does not support an interval".into());
496 }
497 let interval = (kind == JobKind::Loop).then_some(parsed_interval).flatten();
498 if interval.is_some() && limit.is_some() {
499 return Err("an interval and --iterations cannot be used together".into());
500 }
501 let goal = if interval.is_some() {
502 words.next().unwrap_or_default().trim()
503 } else {
504 goal
505 };
506 if goal.is_empty() {
507 return Err("the job goal cannot be empty".into());
508 }
509 Ok(JobRequest {
510 goal: goal.to_string(),
511 limit,
512 interval,
513 })
514 }
515}
516
517fn parse_interval(value: &str) -> Result<Option<Duration>, String> {
518 if !value
519 .as_bytes()
520 .first()
521 .is_some_and(|byte| byte.is_ascii_digit() || *byte == b'.')
522 {
523 return Ok(None);
524 }
525
526 let mut offset = 0;
527 let mut seconds = 0.0;
528 let mut components = 0;
529 while offset < value.len() {
530 let number_start = offset;
531 let mut dot = false;
532 while let Some(byte) = value.as_bytes().get(offset) {
533 if byte.is_ascii_digit() {
534 offset += 1;
535 } else if *byte == b'.' && !dot {
536 dot = true;
537 offset += 1;
538 } else {
539 break;
540 }
541 }
542 let number = value[number_start..offset].parse::<f64>().ok();
543 let Some(number) = number.filter(|number| number.is_finite() && *number >= 0.0) else {
544 return if components == 0 {
545 Ok(None)
546 } else {
547 Err(format!("invalid loop interval: {value}"))
548 };
549 };
550 let rest = &value[offset..];
551 let Some((suffix, scale)) = [
552 ("ns", 0.000_000_001),
553 ("us", 0.000_001),
554 ("µs", 0.000_001),
555 ("ms", 0.001),
556 ("s", 1.0),
557 ("m", 60.0),
558 ("h", 3_600.0),
559 ("d", 86_400.0),
560 ]
561 .into_iter()
562 .find(|(suffix, _)| rest.starts_with(suffix)) else {
563 return if components == 0 {
564 Ok(None)
565 } else {
566 Err(format!("invalid loop interval: {value}"))
567 };
568 };
569 seconds += number * scale;
570 offset += suffix.len();
571 components += 1;
572 }
573
574 let interval = Duration::try_from_secs_f64(seconds)
575 .map_err(|_| format!("loop interval is too large: {value}"))?;
576 if interval.is_zero() {
577 return Err("the loop interval must be greater than zero".into());
578 }
579 Ok(Some(interval))
580}
581
582fn contains_completion_marker(result: &str) -> bool {
583 result.to_ascii_lowercase().contains("[goal-complete]")
584}
585
586fn iteration_prompt(kind: JobKind, goal: &str, iteration: u32, limit: Option<u32>) -> String {
587 let position = limit.map_or_else(
588 || format!("iteration {iteration}"),
589 |limit| format!("iteration {iteration} of {limit}"),
590 );
591 let common = format!(
592 "You are in {position} for this goal:\n\n{goal}\n\nWork directly in the shared repository. Make one useful, verified unit of progress. Do not ask for permission. End the answer with exactly [goal-complete] if the goal is fully achieved and verified. Otherwise end with exactly [continue]."
593 );
594 match kind {
595 JobKind::Loop => common,
596 JobKind::Autoresearch => format!(
597 "Autoresearch job. On the first iteration, establish a repeatable baseline and success metric. On each iteration, test one small idea, run the same verification, keep an improvement, and revert a regression. Protect existing behavior with focused tests. Record the measured result and the decision in the answer.\n\n{common}"
598 ),
599 }
600}
601
602#[cfg(test)]
603mod tests {
604 use super::*;
605
606 #[test]
607 fn iterative_job_arguments_accept_optional_count_or_interval() {
608 assert_eq!(
609 JobRequest::parse(JobKind::Loop, "fix the parser --iterations 7").unwrap(),
610 JobRequest {
611 goal: "fix the parser".into(),
612 limit: Some(7),
613 interval: None,
614 }
615 );
616 assert_eq!(
617 JobRequest::parse(JobKind::Loop, "2d4h measure startup time").unwrap(),
618 JobRequest {
619 goal: "measure startup time".into(),
620 limit: None,
621 interval: Some(Duration::from_secs(2 * 86_400 + 4 * 3_600)),
622 }
623 );
624 assert_eq!(
625 JobRequest::parse(JobKind::Autoresearch, "measure startup time").unwrap(),
626 JobRequest {
627 goal: "measure startup time".into(),
628 limit: None,
629 interval: None,
630 }
631 );
632 assert!(JobRequest::parse(JobKind::Loop, "fix it --iterations 0").is_err());
633 assert!(JobRequest::parse(JobKind::Loop, "--iterations 4").is_err());
634 assert!(JobRequest::parse(JobKind::Loop, "15m fix it --iterations 4").is_err());
635 assert!(JobRequest::parse(JobKind::Autoresearch, "15m fix it").is_err());
636 assert!(JobRequest::parse(JobKind::Loop, "0s fix it").is_err());
637 assert!(JobRequest::parse(JobKind::Loop, "2h30wat fix it").is_err());
638 }
639
640 use kiss_ai::{AssistantEvent, AssistantMessage, ContentBlock, EventStream, StopReason};
641 use std::sync::atomic::AtomicUsize;
642
643 #[test]
644 fn unlimited_prompts_do_not_claim_a_limit() {
645 let prompt = iteration_prompt(JobKind::Loop, "fix tests", 2, None);
646 assert!(prompt.contains("iteration 2 for this goal"));
647 assert!(!prompt.contains("iteration 2 of"));
648 }
649
650 #[test]
651 fn interval_format_uses_compound_units() {
652 assert_eq!(
653 format_interval(Duration::from_secs(2 * 86_400 + 4 * 3_600)),
654 "2d4h"
655 );
656 assert_eq!(format_interval(Duration::from_millis(1_500)), "1s500ms");
657 }
658
659 #[test]
660 fn prompts_define_progress_and_completion() {
661 let loop_prompt = iteration_prompt(JobKind::Loop, "fix tests", 2, Some(10));
662 assert!(loop_prompt.contains("iteration 2 of 10"));
663 assert!(loop_prompt.contains("[goal-complete]"));
664 assert!(loop_prompt.contains("[continue]"));
665
666 let research = iteration_prompt(JobKind::Autoresearch, "make it faster", 1, None);
667 assert!(research.contains("baseline"));
668 assert!(research.contains("success metric"));
669 assert!(research.contains("revert a regression"));
670 }
671
672 #[test]
673 fn completion_marker_is_case_insensitive() {
674 assert!(contains_completion_marker("done\n[GOAL-COMPLETE]"));
675 assert!(!contains_completion_marker("keep going\n[continue]"));
676 }
677
678 fn parent_session() -> Arc<AgentSession> {
679 let registry = kiss_ai::Registry::from_builtin();
680 let model = registry.all().first().unwrap().clone();
681 AgentSession::new(
682 crate::SessionManager::in_memory(std::path::Path::new("/test")),
683 Vec::new(),
684 registry,
685 crate::Settings::default(),
686 "test".into(),
687 model,
688 kiss_ai::ThinkingLevel::Off,
689 None,
690 Arc::new(|_| {}),
691 )
692 }
693
694 #[tokio::test]
695 async fn a_pause_that_arrives_before_the_job_starts_is_kept() {
696 let parent = parent_session();
697 let job = parent
698 .iterative_jobs()
699 .unwrap()
700 .start(
701 JobKind::Loop,
702 JobRequest::parse(JobKind::Loop, "finish it").unwrap(),
703 )
704 .unwrap();
705 job.pause();
706 job.mark_running();
707 assert!(job.is_paused());
708 job.stop();
709 }
710
711 #[tokio::test]
712 async fn a_paused_waiter_resumes_and_a_stopped_job_refuses_work() {
713 let parent = parent_session();
714 let job = parent
715 .iterative_jobs()
716 .unwrap()
717 .start(
718 JobKind::Loop,
719 JobRequest::parse(JobKind::Loop, "finish it").unwrap(),
720 )
721 .unwrap();
722 job.pause();
723 let ready = job.wait_until_ready();
724 tokio::pin!(ready);
725 assert!(futures::poll!(&mut ready).is_pending());
726 job.resume();
727 assert!(
728 tokio::time::timeout(Duration::from_secs(1), ready)
729 .await
730 .unwrap()
731 );
732 job.stop();
733 assert!(!job.wait_until_ready().await);
734 }
735
736 #[tokio::test]
737 async fn one_job_reuses_its_forked_session_and_stops_early() {
738 let parent = parent_session();
739 parent
740 .manager
741 .lock()
742 .unwrap()
743 .append_message(kiss_agent::AgentMessage::user("parent context"))
744 .unwrap();
745 let calls = Arc::new(AtomicUsize::new(0));
746 parent.set_stream_fn(Some(Arc::new({
747 let calls = calls.clone();
748 move |_, _, _| {
749 let call = calls.fetch_add(1, Ordering::SeqCst);
750 let text = if call == 0 {
751 "first change\n[continue]"
752 } else {
753 "verified\n[goal-complete]"
754 };
755 let mut message = AssistantMessage::empty("fake", "fake", "fake");
756 message.content.push(ContentBlock::text(text));
757 message.stop_reason = StopReason::Stop;
758 let (sink, stream) = EventStream::channel();
759 sink.send(AssistantEvent::Start {
760 partial: message.clone(),
761 });
762 sink.done(message);
763 stream
764 }
765 })));
766
767 let runtime = parent.iterative_jobs().unwrap();
768 let job = runtime
769 .start(
770 JobKind::Loop,
771 JobRequest::parse(JobKind::Loop, "finish it").unwrap(),
772 )
773 .unwrap();
774 let mut updates = job.subscribe();
775 tokio::time::timeout(Duration::from_secs(2), async {
776 while !job.snapshot().status.is_finished() {
777 updates.changed().await.unwrap();
778 }
779 })
780 .await
781 .unwrap();
782
783 let snapshot = job.snapshot();
784 assert_eq!(snapshot.status, JobStatus::Completed);
785 assert_eq!(snapshot.iteration, 2);
786 assert_eq!(snapshot.limit, None);
787 assert_eq!(calls.load(Ordering::SeqCst), 2);
788 let messages = job
789 .child
790 .manager
791 .lock()
792 .unwrap()
793 .build_session_context()
794 .messages;
795 assert!(matches!(
796 messages.first(),
797 Some(kiss_agent::AgentMessage::User(user)) if user.content.as_text() == "parent context"
798 ));
799 assert_eq!(
800 messages
801 .iter()
802 .filter(|message| matches!(message, kiss_agent::AgentMessage::Assistant(_)))
803 .count(),
804 2
805 );
806 assert!(!snapshot.session_id.is_empty());
807 }
808}