luft_core/scheduler/
mod.rs

1//! Concurrent scheduler (M1, §2): concurrency limiting, per-run quota, retry,
2//! cancellation, and event reporting.
3//!
4//! Design note (§9.2 C1): the public `run_agent` returns
5//! `Result<AgentResult, SchedulerError>` rather than the design doc's
6//! `(AgentResult, TaskHandle)` tuple — per-agent cancellation is keyed by
7//! `agent_id` via [`Scheduler::cancel_agent`], so a handle is unnecessary.
8
9mod config;
10mod error;
11mod registry;
12
13pub use config::{RetryPolicy, SchedulerConfig};
14pub use error::SchedulerError;
15pub use registry::BackendRegistry;
16
17use crate::contract::*;
18use chrono::Utc;
19use dashmap::DashMap;
20use std::sync::atomic::{AtomicU32, Ordering};
21use std::sync::Arc;
22use std::time::Instant;
23use tokio::sync::{broadcast, Semaphore};
24use tokio_util::sync::CancellationToken;
25
26/// Callback invoked by the scheduler when an agent completes.
27/// Implemented by JournalStore to enable transparent persistence.
28/// Defined here (not in journal.rs) to avoid circular dependency:
29///   scheduler → journal → scheduler.
30#[async_trait::async_trait]
31pub trait JournalCallback: Send + Sync {
32    /// Called when an agent completes (success or non-retryable failure).
33    async fn on_agent_done(
34        &self,
35        agent_id: AgentId,
36        phase_id: PhaseId,
37        status: AgentStatus,
38        output: serde_json::Value,
39        tokens: TokenUsage,
40    );
41}
42
43/// Per-run state held inside the scheduler.
44struct RunState {
45    quota_used: Arc<AtomicU32>,
46    run_cancel: CancellationToken,
47    events: EventSender,
48    /// Per-agent cancel tokens (children of `run_cancel`), keyed by agent id.
49    agent_cancels: DashMap<AgentId, CancellationToken>,
50}
51
52/// Concurrency-controlled agent scheduler. Held as `Arc<Scheduler>` and shared
53/// across orchestration coroutines.
54pub struct Scheduler {
55    config: SchedulerConfig,
56    semaphore: Arc<Semaphore>,
57    registry: BackendRegistry,
58    runs: DashMap<RunId, RunState>,
59    /// Optional journal callback invoked after each agent completes.
60    /// Used by JournalStore for transparent checkpoint persistence.
61    journal_callback: Option<Arc<dyn JournalCallback>>,
62}
63
64impl Scheduler {
65    pub fn new(
66        config: SchedulerConfig,
67        registry: BackendRegistry,
68        journal_callback: Option<Arc<dyn JournalCallback>>,
69    ) -> Arc<Self> {
70        let semaphore = Arc::new(Semaphore::new(config.max_concurrency));
71        Arc::new(Self {
72            config,
73            semaphore,
74            registry,
75            runs: DashMap::new(),
76            journal_callback,
77        })
78    }
79
80    pub fn config(&self) -> &SchedulerConfig {
81        &self.config
82    }
83
84    /// Initialise per-run state. Must be called before any `run_agent`.
85    /// Returns the broadcast receiver; further consumers use `resubscribe()`.
86    pub fn init_run(
87        &self,
88        run_id: RunId,
89        event_capacity: usize,
90    ) -> broadcast::Receiver<AgentEvent> {
91        let (tx, rx) = broadcast::channel(event_capacity);
92        self.init_run_with(run_id, tx);
93        rx
94    }
95
96    /// Initialise per-run state using an externally-owned event sender.
97    ///
98    /// This lets the orchestration layer share a single event bus between the
99    /// scheduler (`AgentStarted`/`AgentDone`, plus the [`RunContext`] handed to
100    /// backends) and the runtime SDK (`phase`/`log`/`pipeline`/`RunDone`).
101    pub fn init_run_with(&self, run_id: RunId, events: EventSender) {
102        self.runs.insert(
103            run_id,
104            RunState {
105                quota_used: Arc::new(AtomicU32::new(0)),
106                run_cancel: CancellationToken::new(),
107                events,
108                agent_cancels: DashMap::new(),
109            },
110        );
111    }
112
113    /// Schedule and run a single agent task: quota check → permit → retry loop →
114    /// events. Cancellation flows via `RunContext::cancel`.
115    ///
116    /// The `agent` span carries `run_id`/`agent_id`/`phase_id`/`model` so every
117    /// log emitted on this task's async path inherits them (see
118    /// `docs/design/program-logging.md`).
119    #[tracing::instrument(
120        name = "agent",
121        skip_all,
122        fields(
123            run_id = %run_id,
124            agent_id = %task.agent_id,
125            phase_id = task.phase_id,
126            model = task.model.as_deref().unwrap_or("default"),
127        )
128    )]
129    pub async fn run_agent(
130        &self,
131        run_id: RunId,
132        mut task: AgentTask,
133        backend_id: Option<&str>,
134    ) -> Result<AgentResult, SchedulerError> {
135        let backend = match backend_id {
136            Some(id) => self.registry.get(id)?,
137            None => self.registry.default_backend()?,
138        };
139
140        // Snapshot per-run handles without holding the DashMap guard across await.
141        let (quota_used, run_cancel, events) = {
142            let rs = self
143                .runs
144                .get(&run_id)
145                .ok_or(SchedulerError::RunNotFound(run_id))?;
146            (
147                rs.quota_used.clone(),
148                rs.run_cancel.clone(),
149                rs.events.clone(),
150            )
151        };
152
153        // Quota.
154        let used = quota_used.fetch_add(1, Ordering::Relaxed) + 1;
155        if used > self.config.quota_per_run {
156            tracing::warn!(
157                used,
158                limit = self.config.quota_per_run,
159                "run quota exceeded"
160            );
161            let _ = events.send(AgentEvent::AgentDone {
162                run_id,
163                agent_id: task.agent_id,
164                status: AgentStatus::Error,
165                tokens: TokenUsage::default(),
166                elapsed_ms: 0,
167                name: task.name.clone(),
168                agent_seq: task.agent_seq,
169                output: serde_json::Value::Null,
170                findings: Vec::new(),
171                prompt: task.prompt.clone(),
172                retry_count: 0,
173                ts: Utc::now(),
174            });
175            return Err(SchedulerError::QuotaExceeded {
176                limit: self.config.quota_per_run,
177                used,
178            });
179        }
180
181        // Per-agent cancel token: a child of the run token, so it fires when the
182        // run is cancelled OR this agent is cancelled individually.
183        let agent_token = run_cancel.child_token();
184        if let Some(rs) = self.runs.get(&run_id) {
185            rs.agent_cancels.insert(task.agent_id, agent_token.clone());
186        }
187
188        // Acquire a permit (cancellable while waiting).
189        let permit = tokio::select! {
190            p = self.semaphore.clone().acquire_owned() => p.expect("semaphore never closed"),
191            _ = agent_token.cancelled() => {
192                let _ = events.send(AgentEvent::AgentDone {
193                    run_id,
194                    agent_id: task.agent_id,
195                    status: AgentStatus::Cancelled,
196                    tokens: TokenUsage::default(),
197                    elapsed_ms: 0,
198                    name: task.name.clone(),
199                    agent_seq: task.agent_seq,
200                    output: serde_json::Value::Null,
201                    findings: Vec::new(),
202                    prompt: task.prompt.clone(),
203                    retry_count: 0,
204                    ts: Utc::now(),
205                });
206                self.cleanup_agent(run_id, task.agent_id);
207                return Err(cancel_kind(&run_cancel));
208            }
209        };
210
211        let _ = events.send(AgentEvent::AgentStarted {
212            run_id,
213            phase_id: task.phase_id,
214            agent_id: task.agent_id,
215            prompt_preview: preview(&task.prompt),
216            model: task.model.clone(),
217            description: task.description.clone(),
218            role: task.role.clone(),
219            name: task.name.clone(),
220            agent_seq: task.agent_seq,
221            ts: Utc::now(),
222        });
223
224        let start = Instant::now();
225        let mut attempt = 0u32;
226        let original_prompt = task.prompt.clone();
227        let mut schema_retry_count = 0u32;
228        let outcome: Result<AgentResult, SchedulerError> = loop {
229            let ctx = RunContext {
230                run_id,
231                cancel: agent_token.clone(),
232                events: events.clone(),
233            };
234            let run_fut = backend.run(task.clone(), ctx);
235            let res = match task.timeout {
236                Some(t) => match tokio::time::timeout(t, run_fut).await {
237                    Ok(r) => r,
238                    Err(_) => Err(BackendError::Timeout),
239                },
240                None => run_fut.await,
241            };
242
243            match res {
244                Ok(result) => {
245                    if let Some(ref schema) = task.output_schema {
246                        let fallback = result.output.get("_agent_fallback_text").is_some();
247                        let validation_err = if fallback {
248                            Some(
249                                "agent returned text instead of calling structured_output tool"
250                                    .to_string(),
251                            )
252                        } else {
253                            validate_output(&result.output, schema)
254                                .err()
255                                .map(|e| e.to_string())
256                        };
257
258                        if let Some(error) = validation_err {
259                            schema_retry_count += 1;
260                            if schema_retry_count > self.config.retry.schema_retry_max {
261                                tracing::error!(
262                                    error = %error,
263                                    attempts = schema_retry_count,
264                                    "agent output failed schema validation, retries exhausted"
265                                );
266                                break Err(SchedulerError::SchemaValidation(error));
267                            }
268                            let _ = events.send(AgentEvent::SchemaRetry {
269                                run_id,
270                                agent_id: task.agent_id,
271                                attempt: schema_retry_count,
272                                max: self.config.retry.schema_retry_max,
273                            });
274                            tracing::warn!(
275                                error = %error,
276                                attempt = schema_retry_count,
277                                "schema validation failed, retrying with feedback"
278                            );
279                            let schema_json =
280                                serde_json::to_string_pretty(schema).unwrap_or_default();
281                            let last_output = if fallback {
282                                result
283                                    .output
284                                    .get("text")
285                                    .and_then(|v| v.as_str())
286                                    .unwrap_or("")
287                                    .to_string()
288                            } else {
289                                serde_json::to_string_pretty(&result.output).unwrap_or_default()
290                            };
291                            task.prompt = if fallback {
292                                format!(
293                                    "{original_prompt}\n\n\
294                                     ---\n\
295                                     You returned your result as plain text instead of calling the `structured_output` tool.\n\
296                                     You MUST call the `structured_output` tool to submit your result.\n\
297                                     Do NOT return the result as a text message.\n\
298                                     \n\
299                                     Your text output was:\n\
300                                     ```\n{last_output}\n```\n\
301                                     \n\
302                                     Required JSON Schema:\n\
303                                     ```json\n{schema}\n```",
304                                    original_prompt = original_prompt,
305                                    last_output = last_output,
306                                    schema = schema_json,
307                                )
308                            } else {
309                                format!(
310                                    "{original_prompt}\n\n\
311                                     ---\n\
312                                     Your previous response did not match the required schema.\n\
313                                     Error: {error}\n\
314                                     \n\
315                                     Your output was:\n\
316                                     ```json\n{last_output}\n```\n\
317                                     \n\
318                                     Required JSON Schema:\n\
319                                     ```json\n{schema}\n```\n\
320                                     \n\
321                                     Call the `structured_output` tool with a JSON object that\n\
322                                     matches this schema exactly. Include ALL required fields.",
323                                    original_prompt = original_prompt,
324                                    error = error,
325                                    last_output = last_output,
326                                    schema = schema_json,
327                                )
328                            };
329                            continue;
330                        }
331                    }
332                    break Ok(result);
333                }
334                Err(e) => {
335                    if agent_token.is_cancelled() || matches!(e, BackendError::Cancelled) {
336                        tracing::debug!("agent cancelled");
337                        break Err(cancel_kind(&run_cancel));
338                    }
339                    if !e.is_retryable() {
340                        tracing::error!(error = %e, "non-retryable backend error");
341                        break Err(SchedulerError::NonRetryable(e));
342                    }
343                    attempt += 1;
344                    if attempt > self.config.retry.max_attempts {
345                        tracing::error!(attempts = attempt, error = %e, "agent exhausted retries");
346                        break Err(SchedulerError::Exhausted {
347                            attempts: attempt,
348                            source: e,
349                        });
350                    }
351                    let backoff = self.config.retry.backoff(attempt);
352                    tracing::warn!(
353                        attempt, backoff_ms = backoff.as_millis() as u64, error = %e,
354                        "retryable backend error; retrying"
355                    );
356                    tokio::select! {
357                        _ = tokio::time::sleep(backoff) => {}
358                        _ = agent_token.cancelled() => break Err(cancel_kind(&run_cancel)),
359                    }
360                }
361            }
362        };
363
364        let elapsed_ms = start.elapsed().as_millis() as u64;
365        let (status, tokens) = match &outcome {
366            Ok(r) => (r.status.clone(), r.tokens_used),
367            Err(SchedulerError::AgentCancelled) | Err(SchedulerError::RunCancelled) => {
368                (AgentStatus::Cancelled, TokenUsage::default())
369            }
370            Err(_) => (AgentStatus::Error, TokenUsage::default()),
371        };
372        let _ = events.send(AgentEvent::AgentDone {
373            run_id,
374            agent_id: task.agent_id,
375            status: status.clone(),
376            tokens,
377            elapsed_ms,
378            name: task.name.clone(),
379            agent_seq: task.agent_seq,
380            output: match &outcome {
381                Ok(r) => r.output.clone(),
382                Err(_) => serde_json::Value::Null,
383            },
384            findings: match &outcome {
385                Ok(r) => r.findings.clone(),
386                Err(_) => Vec::new(),
387            },
388            prompt: task.prompt.clone(),
389            retry_count: attempt,
390            ts: Utc::now(),
391        });
392        tracing::info!(?status, elapsed_ms, "agent finished");
393
394        // Invoke journal callback if configured (M1 transparent persistence).
395        if let Some(ref cb) = self.journal_callback {
396            let output = match &outcome {
397                Ok(r) => r.output.clone(),
398                Err(_) => serde_json::Value::Null,
399            };
400            let agent_status = status.clone();
401            let tokens_used = tokens;
402            let agent_id = task.agent_id;
403            let phase_id = task.phase_id;
404            cb.on_agent_done(agent_id, phase_id, agent_status, output, tokens_used)
405                .await;
406        }
407
408        drop(permit);
409        self.cleanup_agent(run_id, task.agent_id);
410        outcome
411    }
412
413    /// Run a batch of tasks concurrently (the `parallel()` primitive). Bounded
414    /// by the same global semaphore; does not short-circuit on failure — results
415    /// preserve input order.
416    pub async fn run_parallel(
417        &self,
418        run_id: RunId,
419        tasks: Vec<(AgentTask, Option<String>)>,
420    ) -> Vec<Result<AgentResult, SchedulerError>> {
421        let futs = tasks.into_iter().map(|(task, backend)| async move {
422            self.run_agent(run_id, task, backend.as_deref()).await
423        });
424        futures::future::join_all(futs).await
425    }
426
427    /// Cancel one agent (fires its token; the backend observes `ctx.cancel`).
428    pub fn cancel_agent(&self, run_id: RunId, agent_id: AgentId) {
429        if let Some(rs) = self.runs.get(&run_id) {
430            if let Some(tok) = rs.agent_cancels.get(&agent_id) {
431                tok.cancel();
432            }
433        }
434    }
435
436    /// Cancel the whole run (all child agent tokens fire).
437    pub fn cancel_run(&self, run_id: RunId) {
438        if let Some(rs) = self.runs.get(&run_id) {
439            rs.run_cancel.cancel();
440        }
441    }
442
443    /// Current global active concurrency.
444    pub fn active_concurrency(&self) -> usize {
445        self.config.max_concurrency - self.semaphore.available_permits()
446    }
447
448    /// Quota consumed by a run, if initialised.
449    pub fn quota_used(&self, run_id: RunId) -> Option<u32> {
450        self.runs
451            .get(&run_id)
452            .map(|rs| rs.quota_used.load(Ordering::Relaxed))
453    }
454
455    fn cleanup_agent(&self, run_id: RunId, agent_id: AgentId) {
456        if let Some(rs) = self.runs.get(&run_id) {
457            rs.agent_cancels.remove(&agent_id);
458        }
459    }
460}
461
462fn cancel_kind(run_cancel: &CancellationToken) -> SchedulerError {
463    if run_cancel.is_cancelled() {
464        SchedulerError::RunCancelled
465    } else {
466        SchedulerError::AgentCancelled
467    }
468}
469
470fn preview(s: &str) -> String {
471    s.chars().take(60).collect()
472}
473
474#[cfg(test)]
475mod tests {
476    use super::*;
477    use crate::mock_backend::{FailKind, MockBackend, MockBehavior};
478    use std::path::PathBuf;
479    use std::sync::atomic::AtomicUsize;
480    use std::time::Duration;
481    use uuid::Uuid;
482
483    fn fast_config(max_concurrency: usize, quota: u32) -> SchedulerConfig {
484        SchedulerConfig {
485            max_concurrency,
486            quota_per_run: quota,
487            retry: RetryPolicy {
488                max_attempts: 2,
489                initial_backoff: Duration::from_millis(1),
490                backoff_multiplier: 2.0,
491                max_backoff: Duration::from_millis(5),
492                schema_retry_max: 1,
493            },
494        }
495    }
496
497    fn mk_task(prompt: &str) -> AgentTask {
498        AgentTask {
499            agent_id: Uuid::now_v7(),
500            phase_id: 0,
501            prompt: prompt.to_string(),
502            model: None,
503            allowlist: None,
504            workdir: PathBuf::from("."),
505            mcp_endpoint: None,
506            timeout: None,
507            output_schema: None,
508            workdir_override: None,
509            description: None,
510            role: None,
511            name: None,
512            agent_seq: 0,
513            thread_id: None,
514        }
515    }
516
517    fn mk_task_with_schema(prompt: &str) -> AgentTask {
518        let mut task = mk_task(prompt);
519        task.output_schema = Some(serde_json::json!({
520            "type": "object",
521            "properties": {
522                "answer": { "type": "string" }
523            },
524            "required": ["answer"]
525        }));
526        task
527    }
528
529    fn fallback_output(text: &str) -> serde_json::Value {
530        serde_json::json!({
531            "_agent_fallback_text": true,
532            "text": text,
533        })
534    }
535
536    fn ok_result(id: AgentId) -> AgentResult {
537        AgentResult {
538            agent_id: id,
539            status: AgentStatus::Ok,
540            output: serde_json::Value::Null,
541            findings: vec![],
542            tokens_used: TokenUsage::default(),
543            artifacts: vec![],
544            logs: LogRef::default(),
545            thread_id: None,
546        }
547    }
548
549    fn sched_with(backend: Arc<dyn AgentBackend>, cfg: SchedulerConfig) -> Arc<Scheduler> {
550        Scheduler::new(cfg, BackendRegistry::new().with(backend), None)
551    }
552
553    // A backend that records peak concurrency.
554    struct ProbeBackend {
555        cur: Arc<AtomicUsize>,
556        peak: Arc<AtomicUsize>,
557        delay: Duration,
558    }
559
560    #[async_trait::async_trait]
561    impl AgentBackend for ProbeBackend {
562        fn id(&self) -> &'static str {
563            "probe"
564        }
565        fn capabilities(&self) -> AgentCapabilities {
566            AgentCapabilities::default()
567        }
568        fn as_any(&self) -> &dyn std::any::Any {
569            self
570        }
571        async fn run(
572            &self,
573            task: AgentTask,
574            _ctx: RunContext,
575        ) -> Result<AgentResult, BackendError> {
576            let c = self.cur.fetch_add(1, Ordering::SeqCst) + 1;
577            self.peak.fetch_max(c, Ordering::SeqCst);
578            tokio::time::sleep(self.delay).await;
579            self.cur.fetch_sub(1, Ordering::SeqCst);
580            Ok(ok_result(task.agent_id))
581        }
582    }
583
584    #[tokio::test]
585    async fn test_default_config_concurrency() {
586        let c = SchedulerConfig::default().max_concurrency;
587        assert!((4..=16).contains(&c), "got {c}");
588    }
589
590    #[tokio::test]
591    async fn test_concurrency_limit() {
592        let cur = Arc::new(AtomicUsize::new(0));
593        let peak = Arc::new(AtomicUsize::new(0));
594        let backend = Arc::new(ProbeBackend {
595            cur: cur.clone(),
596            peak: peak.clone(),
597            delay: Duration::from_millis(40),
598        });
599        let sched = sched_with(backend, fast_config(2, 1000));
600        let run_id = Uuid::now_v7();
601        let _rx = sched.init_run(run_id, 256);
602
603        let tasks: Vec<_> = (0..6).map(|i| (mk_task(&format!("t{i}")), None)).collect();
604        let results = sched.run_parallel(run_id, tasks).await;
605
606        assert!(results.iter().all(|r| r.is_ok()));
607        assert!(
608            peak.load(Ordering::SeqCst) <= 2,
609            "peak {}",
610            peak.load(Ordering::SeqCst)
611        );
612    }
613
614    #[tokio::test]
615    async fn test_quota_exceeded() {
616        let backend = Arc::new(MockBackend::new(
617            "mock",
618            vec![MockBehavior::Success {
619                output: serde_json::Value::Null,
620                tokens: TokenUsage::default(),
621                delay: Duration::from_millis(5),
622            }],
623        ));
624        let sched = sched_with(backend, fast_config(8, 3));
625        let run_id = Uuid::now_v7();
626        let _rx = sched.init_run(run_id, 256);
627
628        let tasks: Vec<_> = (0..4).map(|i| (mk_task(&format!("t{i}")), None)).collect();
629        let results = sched.run_parallel(run_id, tasks).await;
630
631        let ok = results.iter().filter(|r| r.is_ok()).count();
632        let quota_err = results
633            .iter()
634            .filter(|r| matches!(r, Err(SchedulerError::QuotaExceeded { .. })))
635            .count();
636        assert_eq!(ok, 3);
637        assert_eq!(quota_err, 1);
638    }
639
640    #[tokio::test]
641    async fn test_retry_on_retryable_error() {
642        let backend = Arc::new(MockBackend::new(
643            "mock",
644            vec![
645                MockBehavior::fail(FailKind::Spawn),
646                MockBehavior::fail(FailKind::Spawn),
647                MockBehavior::Success {
648                    output: serde_json::Value::Null,
649                    tokens: TokenUsage::default(),
650                    delay: Duration::ZERO,
651                },
652            ],
653        ));
654        let probe = backend.clone();
655        let sched = sched_with(backend, fast_config(4, 1000));
656        let run_id = Uuid::now_v7();
657        let _rx = sched.init_run(run_id, 64);
658
659        let r = sched.run_agent(run_id, mk_task("x"), None).await;
660        assert!(r.is_ok(), "{r:?}");
661        assert_eq!(probe.call_count(), 3);
662    }
663
664    #[tokio::test]
665    async fn test_no_retry_on_non_retryable() {
666        let backend = Arc::new(MockBackend::new(
667            "mock",
668            vec![MockBehavior::fail(FailKind::Protocol)],
669        ));
670        let probe = backend.clone();
671        let sched = sched_with(backend, fast_config(4, 1000));
672        let run_id = Uuid::now_v7();
673        let _rx = sched.init_run(run_id, 64);
674
675        let r = sched.run_agent(run_id, mk_task("x"), None).await;
676        assert!(matches!(r, Err(SchedulerError::NonRetryable(_))), "{r:?}");
677        assert_eq!(probe.call_count(), 1);
678    }
679
680    #[tokio::test]
681    async fn test_retry_exhausted() {
682        let backend = Arc::new(MockBackend::new(
683            "mock",
684            vec![MockBehavior::fail(FailKind::Spawn)],
685        ));
686        let probe = backend.clone();
687        let sched = sched_with(backend, fast_config(4, 1000));
688        let run_id = Uuid::now_v7();
689        let _rx = sched.init_run(run_id, 64);
690
691        let r = sched.run_agent(run_id, mk_task("x"), None).await;
692        assert!(
693            matches!(r, Err(SchedulerError::Exhausted { attempts: 3, .. })),
694            "{r:?}"
695        );
696        assert_eq!(probe.call_count(), 3);
697    }
698
699    #[tokio::test]
700    async fn test_schema_fallback_then_succeeds() {
701        let backend = Arc::new(MockBackend::new(
702            "mock",
703            vec![
704                MockBehavior::Success {
705                    output: fallback_output("i forgot the tool"),
706                    tokens: TokenUsage::default(),
707                    delay: Duration::ZERO,
708                },
709                MockBehavior::Success {
710                    output: serde_json::json!({"answer": "ok"}),
711                    tokens: TokenUsage::default(),
712                    delay: Duration::ZERO,
713                },
714            ],
715        ));
716        let probe = backend.clone();
717        let sched = sched_with(backend, fast_config(4, 1000));
718        let run_id = Uuid::now_v7();
719        let mut rx = sched.init_run(run_id, 64);
720
721        let task = mk_task_with_schema("respond");
722        let r = sched.run_agent(run_id, task, None).await;
723        assert!(r.is_ok(), "{r:?}");
724        assert_eq!(probe.call_count(), 2);
725
726        let mut prompt_with_feedback = None;
727        while let Ok(event) = rx.try_recv() {
728            if let AgentEvent::AgentDone { prompt, .. } = event {
729                prompt_with_feedback = Some(prompt);
730            }
731        }
732        let prompt = prompt_with_feedback.expect("AgentDone event with prompt");
733        assert!(prompt.contains("structured_output"));
734        assert!(prompt.contains("Required JSON Schema"));
735    }
736
737    #[tokio::test]
738    async fn test_schema_mismatch_then_succeeds() {
739        let backend = Arc::new(MockBackend::new(
740            "mock",
741            vec![
742                MockBehavior::Success {
743                    output: serde_json::json!({"wrong": "field"}),
744                    tokens: TokenUsage::default(),
745                    delay: Duration::ZERO,
746                },
747                MockBehavior::Success {
748                    output: serde_json::json!({"answer": "ok"}),
749                    tokens: TokenUsage::default(),
750                    delay: Duration::ZERO,
751                },
752            ],
753        ));
754        let probe = backend.clone();
755        let sched = sched_with(backend, fast_config(4, 1000));
756        let run_id = Uuid::now_v7();
757        let _rx = sched.init_run(run_id, 64);
758
759        let task = mk_task_with_schema("respond");
760        let r = sched.run_agent(run_id, task, None).await;
761        assert!(r.is_ok(), "{r:?}");
762        assert_eq!(probe.call_count(), 2);
763    }
764
765    #[tokio::test]
766    async fn test_schema_fallback_exhausted() {
767        let backend = Arc::new(MockBackend::new(
768            "mock",
769            vec![MockBehavior::Success {
770                output: fallback_output("still no tool"),
771                tokens: TokenUsage::default(),
772                delay: Duration::ZERO,
773            }],
774        ));
775        let probe = backend.clone();
776        let sched = sched_with(backend, fast_config(4, 1000));
777        let run_id = Uuid::now_v7();
778        let _rx = sched.init_run(run_id, 64);
779
780        let task = mk_task_with_schema("respond");
781        let r = sched.run_agent(run_id, task, None).await;
782        assert!(
783            matches!(r, Err(SchedulerError::SchemaValidation(_))),
784            "{r:?}"
785        );
786        assert_eq!(probe.call_count(), 2);
787    }
788
789    #[tokio::test]
790    async fn test_cancel_run() {
791        let backend = Arc::new(MockBackend::new("mock", vec![MockBehavior::Hang]));
792        let sched = sched_with(backend, fast_config(8, 1000));
793        let run_id = Uuid::now_v7();
794        let _rx = sched.init_run(run_id, 64);
795
796        let s2 = sched.clone();
797        let handle = tokio::spawn(async move {
798            let tasks: Vec<_> = (0..3).map(|i| (mk_task(&format!("h{i}")), None)).collect();
799            s2.run_parallel(run_id, tasks).await
800        });
801        tokio::time::sleep(Duration::from_millis(20)).await;
802        sched.cancel_run(run_id);
803
804        let results = handle.await.unwrap();
805        assert_eq!(results.len(), 3);
806        assert!(results
807            .iter()
808            .all(|r| matches!(r, Err(SchedulerError::RunCancelled))));
809    }
810
811    #[tokio::test]
812    async fn test_cancel_agent() {
813        let backend = Arc::new(MockBackend::new("mock", vec![MockBehavior::Hang]));
814        let sched = sched_with(backend, fast_config(8, 1000));
815        let run_id = Uuid::now_v7();
816        let _rx = sched.init_run(run_id, 64);
817
818        let task = mk_task("hang");
819        let agent_id = task.agent_id;
820        let s2 = sched.clone();
821        let handle = tokio::spawn(async move { s2.run_agent(run_id, task, None).await });
822        tokio::time::sleep(Duration::from_millis(20)).await;
823        sched.cancel_agent(run_id, agent_id);
824
825        let r = handle.await.unwrap();
826        assert!(matches!(r, Err(SchedulerError::AgentCancelled)), "{r:?}");
827    }
828
829    #[tokio::test]
830    async fn test_parallel_partial_failure() {
831        let backend = Arc::new(MockBackend::new(
832            "mock",
833            vec![
834                MockBehavior::Success {
835                    output: serde_json::Value::Null,
836                    tokens: TokenUsage::default(),
837                    delay: Duration::ZERO,
838                },
839                MockBehavior::fail(FailKind::Protocol),
840                MockBehavior::Success {
841                    output: serde_json::Value::Null,
842                    tokens: TokenUsage::default(),
843                    delay: Duration::ZERO,
844                },
845            ],
846        ));
847        let sched = sched_with(backend, fast_config(1, 1000)); // serialize for deterministic behavior order
848        let run_id = Uuid::now_v7();
849        let _rx = sched.init_run(run_id, 64);
850
851        let tasks: Vec<_> = (0..3).map(|i| (mk_task(&format!("p{i}")), None)).collect();
852        let results = sched.run_parallel(run_id, tasks).await;
853
854        assert_eq!(results.len(), 3);
855        assert_eq!(results.iter().filter(|r| r.is_ok()).count(), 2);
856        assert_eq!(results.iter().filter(|r| r.is_err()).count(), 1);
857    }
858
859    #[tokio::test]
860    async fn test_event_sequence() {
861        let backend = Arc::new(MockBackend::new(
862            "mock",
863            vec![MockBehavior::Success {
864                output: serde_json::Value::Null,
865                tokens: TokenUsage {
866                    input: 10,
867                    output: 5,
868                    ..Default::default()
869                },
870                delay: Duration::ZERO,
871            }],
872        ));
873        let sched = sched_with(backend, fast_config(4, 1000));
874        let run_id = Uuid::now_v7();
875        let mut rx = sched.init_run(run_id, 64);
876
877        let r = sched.run_agent(run_id, mk_task("x"), None).await;
878        assert!(r.is_ok());
879
880        let e1 = rx.recv().await.unwrap();
881        assert!(matches!(e1, AgentEvent::AgentStarted { .. }), "{e1:?}");
882        let e2 = rx.recv().await.unwrap();
883        match e2 {
884            AgentEvent::AgentDone { status, tokens, .. } => {
885                assert_eq!(status, AgentStatus::Ok);
886                assert_eq!(tokens.input, 10);
887            }
888            other => panic!("expected AgentDone, got {other:?}"),
889        }
890    }
891}
luft_core/scheduler/mod.rs

luft_core/scheduler/
mod.rs
