kiss_coding/
iterative.rs

1//! Bounded loop and autoresearch jobs in persistent child sessions.
2
3use 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 set_status(&self, status: JobStatus) {
217        self.state.lock().unwrap().status = status;
218        self.signal();
219    }
220
221    async fn wait_until_ready(&self) -> bool {
222        loop {
223            if self.cancel.is_cancelled() {
224                return false;
225            }
226            if !self.is_paused() {
227                return true;
228            }
229            tokio::select! {
230                _ = self.cancel.cancelled() => return false,
231                _ = self.wake.notified() => {}
232            }
233        }
234    }
235}
236
237pub struct IterativeRuntime {
238    parent: Weak<AgentSession>,
239    jobs: Mutex<Vec<Arc<JobRecord>>>,
240    next_id: AtomicU64,
241    permits: Arc<Semaphore>,
242}
243
244impl IterativeRuntime {
245    pub(crate) fn new(parent: Weak<AgentSession>) -> Arc<Self> {
246        Arc::new(Self {
247            parent,
248            jobs: Mutex::new(Vec::new()),
249            next_id: AtomicU64::new(1),
250            permits: Arc::new(Semaphore::new(MAX_ACTIVE_JOBS)),
251        })
252    }
253
254    pub fn start(
255        self: &Arc<Self>,
256        kind: JobKind,
257        goal: impl Into<String>,
258        limit: Option<u32>,
259        interval: Option<Duration>,
260    ) -> anyhow::Result<Arc<JobRecord>> {
261        let goal = goal.into();
262        if goal.trim().is_empty() {
263            anyhow::bail!("the job goal cannot be empty");
264        }
265        let parent = self
266            .parent
267            .upgrade()
268            .ok_or_else(|| anyhow::anyhow!("the parent session has closed"))?;
269        let id = self.next_id.fetch_add(1, Ordering::SeqCst);
270        if limit.is_some() && interval.is_some() {
271            anyhow::bail!("an interval and --iterations cannot be used together");
272        }
273        if kind == JobKind::Autoresearch && interval.is_some() {
274            anyhow::bail!("autoresearch does not support an interval");
275        }
276        if interval.is_some_and(|interval| interval.is_zero()) {
277            anyhow::bail!("the interval must be greater than zero");
278        }
279        let limit = limit.map(|limit| limit.clamp(1, MAX_ITERATIONS));
280        let task_name = format!("{}_{}", kind.label(), id);
281        let child = parent.create_subagent_session(
282            &task_name,
283            &format!("/root/{task_name}"),
284            ForkTurns::All,
285            None,
286            None,
287        )?;
288        let session_id = child.manager.lock().unwrap().session_id().to_string();
289        let (version, _) = watch::channel(0);
290        let record = Arc::new(JobRecord {
291            id,
292            kind,
293            goal,
294            limit,
295            interval,
296            session_id,
297            child,
298            started: Instant::now(),
299            state: Mutex::new(JobState {
300                status: JobStatus::Queued,
301                iteration: 0,
302                tokens: 0,
303                latest_result: None,
304                error: None,
305            }),
306            cancel: CancellationToken::new(),
307            wake: Notify::new(),
308            version,
309            parent: self.parent.clone(),
310        });
311
312        let mut jobs = self
313            .jobs
314            .lock()
315            .map_err(|_| anyhow::anyhow!("the job list is unavailable"))?;
316        jobs.push(record.clone());
317        while jobs.len() > MAX_RETAINED_JOBS {
318            let Some(position) = jobs
319                .iter()
320                .position(|job| job.state.lock().unwrap().status.is_finished())
321            else {
322                break;
323            };
324            jobs.remove(position);
325        }
326        drop(jobs);
327
328        let runtime = self.clone();
329        let running = record.clone();
330        tokio::spawn(async move { runtime.run(running).await });
331        record.signal();
332        Ok(record)
333    }
334
335    async fn run(self: Arc<Self>, job: Arc<JobRecord>) {
336        let permit = tokio::select! {
337            _ = job.cancel.cancelled() => None,
338            permit = self.permits.clone().acquire_owned() => permit.ok(),
339        };
340        let Some(_permit) = permit else {
341            job.set_status(JobStatus::Stopped);
342            return;
343        };
344        if !job.wait_until_ready().await {
345            job.set_status(JobStatus::Stopped);
346            return;
347        }
348        job.set_status(JobStatus::Running);
349
350        let mut iteration = 1_u32;
351        loop {
352            if !job.wait_until_ready().await {
353                job.set_status(JobStatus::Stopped);
354                return;
355            }
356            job.state.lock().unwrap().iteration = iteration;
357            job.signal();
358
359            let outcome = child_turn::run_child_turn(
360                &job.parent,
361                &job.child,
362                iteration_prompt(job.kind, &job.goal, iteration, job.limit),
363                Some(job.cancel.clone()),
364                None,
365            )
366            .await;
367            {
368                let mut state = job.state.lock().unwrap();
369                state.tokens = state.tokens.saturating_add(outcome.usage.total_tokens);
370                state.latest_result = outcome.result.clone();
371                state.error = outcome.error.clone();
372            }
373            job.signal();
374
375            if job.cancel.is_cancelled() || outcome.status == AgentStatus::Interrupted {
376                job.set_status(JobStatus::Stopped);
377                return;
378            }
379            if outcome.status == AgentStatus::Failed {
380                job.set_status(JobStatus::Failed);
381                return;
382            }
383            if outcome
384                .result
385                .as_deref()
386                .is_some_and(contains_completion_marker)
387            {
388                job.set_status(JobStatus::Completed);
389                return;
390            }
391            if job.limit == Some(iteration) {
392                job.set_status(JobStatus::Completed);
393                return;
394            }
395            if let Some(interval) = job.interval {
396                tokio::select! {
397                    _ = job.cancel.cancelled() => {
398                        job.set_status(JobStatus::Stopped);
399                        return;
400                    }
401                    _ = tokio::time::sleep(interval) => {}
402                }
403            }
404            iteration = iteration.saturating_add(1);
405        }
406    }
407
408    pub fn get(&self, id: JobId) -> Option<Arc<JobRecord>> {
409        self.jobs
410            .lock()
411            .ok()?
412            .iter()
413            .find(|job| job.id == id)
414            .cloned()
415    }
416
417    pub fn summaries(&self) -> Vec<JobSummary> {
418        self.jobs
419            .lock()
420            .map(|jobs| jobs.iter().map(|job| job.summary()).collect())
421            .unwrap_or_default()
422    }
423
424    pub fn active_count(&self) -> usize {
425        self.jobs
426            .lock()
427            .map(|jobs| {
428                jobs.iter()
429                    .filter(|job| !job.state.lock().unwrap().status.is_finished())
430                    .count()
431            })
432            .unwrap_or(0)
433    }
434
435    pub fn latest(&self) -> Option<Arc<JobRecord>> {
436        self.jobs.lock().ok()?.last().cloned()
437    }
438
439    pub(crate) fn stop_all(&self) {
440        if let Ok(jobs) = self.jobs.lock() {
441            for job in jobs.iter() {
442                job.stop();
443            }
444        }
445    }
446}
447
448fn contains_completion_marker(result: &str) -> bool {
449    result.to_ascii_lowercase().contains("[goal-complete]")
450}
451
452fn iteration_prompt(kind: JobKind, goal: &str, iteration: u32, limit: Option<u32>) -> String {
453    let position = limit.map_or_else(
454        || format!("iteration {iteration}"),
455        |limit| format!("iteration {iteration} of {limit}"),
456    );
457    let common = format!(
458        "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]."
459    );
460    match kind {
461        JobKind::Loop => common,
462        JobKind::Autoresearch => format!(
463            "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}"
464        ),
465    }
466}
467
468#[cfg(test)]
469mod tests {
470    use super::*;
471    use kiss_ai::{AssistantEvent, AssistantMessage, ContentBlock, EventStream, StopReason};
472    use std::sync::atomic::AtomicUsize;
473
474    #[test]
475    fn unlimited_prompts_do_not_claim_a_limit() {
476        let prompt = iteration_prompt(JobKind::Loop, "fix tests", 2, None);
477        assert!(prompt.contains("iteration 2 for this goal"));
478        assert!(!prompt.contains("iteration 2 of"));
479    }
480
481    #[test]
482    fn interval_format_uses_compound_units() {
483        assert_eq!(
484            format_interval(Duration::from_secs(2 * 86_400 + 4 * 3_600)),
485            "2d4h"
486        );
487        assert_eq!(format_interval(Duration::from_millis(1_500)), "1s500ms");
488    }
489
490    #[test]
491    fn prompts_define_progress_and_completion() {
492        let loop_prompt = iteration_prompt(JobKind::Loop, "fix tests", 2, Some(10));
493        assert!(loop_prompt.contains("iteration 2 of 10"));
494        assert!(loop_prompt.contains("[goal-complete]"));
495        assert!(loop_prompt.contains("[continue]"));
496
497        let research = iteration_prompt(JobKind::Autoresearch, "make it faster", 1, None);
498        assert!(research.contains("baseline"));
499        assert!(research.contains("success metric"));
500        assert!(research.contains("revert a regression"));
501    }
502
503    #[test]
504    fn completion_marker_is_case_insensitive() {
505        assert!(contains_completion_marker("done\n[GOAL-COMPLETE]"));
506        assert!(!contains_completion_marker("keep going\n[continue]"));
507    }
508
509    #[tokio::test]
510    async fn one_job_reuses_its_forked_session_and_stops_early() {
511        let registry = kiss_ai::Registry::from_builtin();
512        let model = registry.all().first().unwrap().clone();
513        let parent = AgentSession::new(
514            crate::SessionManager::in_memory(std::path::Path::new("/test")),
515            Vec::new(),
516            registry,
517            crate::Settings::default(),
518            "test".into(),
519            model,
520            kiss_ai::ThinkingLevel::Off,
521            None,
522            Arc::new(|_| {}),
523        );
524        parent
525            .manager
526            .lock()
527            .unwrap()
528            .append_message(kiss_agent::AgentMessage::user("parent context"))
529            .unwrap();
530        let calls = Arc::new(AtomicUsize::new(0));
531        parent.set_stream_fn(Some(Arc::new({
532            let calls = calls.clone();
533            move |_, _, _| {
534                let call = calls.fetch_add(1, Ordering::SeqCst);
535                let text = if call == 0 {
536                    "first change\n[continue]"
537                } else {
538                    "verified\n[goal-complete]"
539                };
540                let mut message = AssistantMessage::empty("fake", "fake", "fake");
541                message.content.push(ContentBlock::text(text));
542                message.stop_reason = StopReason::Stop;
543                let (sink, stream) = EventStream::channel();
544                sink.send(AssistantEvent::Start {
545                    partial: message.clone(),
546                });
547                sink.done(message);
548                stream
549            }
550        })));
551
552        let runtime = parent.iterative_jobs().unwrap();
553        let job = runtime
554            .start(JobKind::Loop, "finish it", None, None)
555            .unwrap();
556        let mut updates = job.subscribe();
557        tokio::time::timeout(Duration::from_secs(2), async {
558            while !job.snapshot().status.is_finished() {
559                updates.changed().await.unwrap();
560            }
561        })
562        .await
563        .unwrap();
564
565        let snapshot = job.snapshot();
566        assert_eq!(snapshot.status, JobStatus::Completed);
567        assert_eq!(snapshot.iteration, 2);
568        assert_eq!(snapshot.limit, None);
569        assert_eq!(calls.load(Ordering::SeqCst), 2);
570        let messages = job
571            .child
572            .manager
573            .lock()
574            .unwrap()
575            .build_session_context()
576            .messages;
577        assert!(matches!(
578            messages.first(),
579            Some(kiss_agent::AgentMessage::User(user)) if user.content.as_text() == "parent context"
580        ));
581        assert_eq!(
582            messages
583                .iter()
584                .filter(|message| matches!(message, kiss_agent::AgentMessage::Assistant(_)))
585                .count(),
586            2
587        );
588        assert!(!snapshot.session_id.is_empty());
589    }
590}
kiss_coding/iterative.rs

kiss_coding/
iterative.rs
