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.init_run_with_cancel(run_id, events, CancellationToken::new());
103    }
104
105    /// Initialise per-run state with a caller-owned cancellation token.
106    /// The runtime and scheduler then observe the same cancellation signal.
107    pub fn init_run_with_cancel(
108        &self,
109        run_id: RunId,
110        events: EventSender,
111        run_cancel: CancellationToken,
112    ) {
113        self.runs.insert(
114            run_id,
115            RunState {
116                quota_used: Arc::new(AtomicU32::new(0)),
117                run_cancel,
118                events,
119                agent_cancels: DashMap::new(),
120            },
121        );
122    }
123
124    /// Schedule and run a single agent task: quota check → permit → retry loop →
125    /// events. Cancellation flows via `RunContext::cancel`.
126    ///
127    /// The `agent` span carries `run_id`/`agent_id`/`phase_id`/`model` so every
128    /// log emitted on this task's async path inherits them (see
129    /// `docs/design/program-logging.md`).
130    #[tracing::instrument(
131        name = "agent",
132        skip_all,
133        fields(
134            run_id = %run_id,
135            agent_id = %task.agent_id,
136            phase_id = task.phase_id,
137            model = task.model.as_deref().unwrap_or("default"),
138        )
139    )]
140    pub async fn run_agent(
141        &self,
142        run_id: RunId,
143        mut task: AgentTask,
144        backend_id: Option<&str>,
145    ) -> Result<AgentResult, SchedulerError> {
146        let backend = match backend_id {
147            Some(id) => self.registry.get(id)?,
148            None => {
149                let from_current = crate::contract::current_backend()
150                    .and_then(|cb| self.registry.get(&cb.id).ok());
151                match from_current {
152                    Some(b) => b,
153                    None => self.registry.default_backend()?,
154                }
155            }
156        };
157
158        // Snapshot per-run handles without holding the DashMap guard across await.
159        let (quota_used, run_cancel, events) = {
160            let rs = self
161                .runs
162                .get(&run_id)
163                .ok_or(SchedulerError::RunNotFound(run_id))?;
164            (
165                rs.quota_used.clone(),
166                rs.run_cancel.clone(),
167                rs.events.clone(),
168            )
169        };
170
171        // Quota.
172        let used = quota_used.fetch_add(1, Ordering::Relaxed) + 1;
173        if used > self.config.quota_per_run {
174            tracing::warn!(
175                used,
176                limit = self.config.quota_per_run,
177                "run quota exceeded"
178            );
179            let _ = events.send(AgentEvent::AgentDone {
180                run_id,
181                agent_id: task.agent_id,
182                status: AgentStatus::Error,
183                tokens: TokenUsage::default(),
184                elapsed_ms: 0,
185                name: task.name.clone(),
186                agent_seq: task.agent_seq,
187                output: serde_json::Value::Null,
188                findings: Vec::new(),
189                prompt: task.prompt.clone(),
190                retry_count: 0,
191                ts: Utc::now(),
192            });
193            return Err(SchedulerError::QuotaExceeded {
194                limit: self.config.quota_per_run,
195                used,
196            });
197        }
198
199        // Per-agent cancel token: a child of the run token, so it fires when the
200        // run is cancelled OR this agent is cancelled individually.
201        let agent_token = run_cancel.child_token();
202        if let Some(rs) = self.runs.get(&run_id) {
203            rs.agent_cancels.insert(task.agent_id, agent_token.clone());
204        }
205
206        // Acquire a permit (cancellable while waiting).
207        let permit = tokio::select! {
208            p = self.semaphore.clone().acquire_owned() => p.expect("semaphore never closed"),
209            _ = agent_token.cancelled() => {
210                let _ = events.send(AgentEvent::AgentDone {
211                    run_id,
212                    agent_id: task.agent_id,
213                    status: AgentStatus::Cancelled,
214                    tokens: TokenUsage::default(),
215                    elapsed_ms: 0,
216                    name: task.name.clone(),
217                    agent_seq: task.agent_seq,
218                    output: serde_json::Value::Null,
219                    findings: Vec::new(),
220                    prompt: task.prompt.clone(),
221                    retry_count: 0,
222                    ts: Utc::now(),
223                });
224                self.cleanup_agent(run_id, task.agent_id);
225                return Err(cancel_kind(&run_cancel));
226            }
227        };
228
229        let _ = events.send(AgentEvent::AgentStarted {
230            run_id,
231            phase_id: task.phase_id,
232            agent_id: task.agent_id,
233            prompt_preview: preview(&task.prompt),
234            model: task.model.clone(),
235            description: task.description.clone(),
236            role: task.role.clone(),
237            name: task.name.clone(),
238            agent_seq: task.agent_seq,
239            ts: Utc::now(),
240        });
241
242        let start = Instant::now();
243        let mut attempt = 0u32;
244        let original_prompt = task.prompt.clone();
245        let mut schema_retry_count = 0u32;
246        let outcome: Result<AgentResult, SchedulerError> = loop {
247            let ctx = RunContext {
248                run_id,
249                cancel: agent_token.clone(),
250                events: events.clone(),
251            };
252            let run_fut = backend.run(task.clone(), ctx);
253            let res = match task.timeout {
254                Some(t) => match tokio::time::timeout(t, run_fut).await {
255                    Ok(r) => r,
256                    Err(_) => Err(BackendError::Timeout),
257                },
258                None => run_fut.await,
259            };
260
261            match res {
262                Ok(result) => {
263                    if let Some(ref schema) = task.output_schema {
264                        let fallback_text = match &result.output {
265                            serde_json::Value::String(s) => Some(s.clone()),
266                            obj if obj.get("_agent_fallback_text").is_some() => obj
267                                .get("text")
268                                .and_then(|v| v.as_str())
269                                .map(|s| s.to_string()),
270                            _ => None,
271                        };
272                        let validation_err = if fallback_text.is_some() {
273                            Some(
274                                "agent returned text instead of calling workflow_validate_schema tool"
275                                    .to_string(),
276                            )
277                        } else {
278                            validate_output(&result.output, schema)
279                                .err()
280                                .map(|e| e.to_string())
281                        };
282
283                        if let Some(error) = validation_err {
284                            schema_retry_count += 1;
285                            if schema_retry_count > self.config.retry.schema_retry_max {
286                                tracing::error!(
287                                    error = %error,
288                                    attempts = schema_retry_count,
289                                    "agent output failed schema validation, retries exhausted"
290                                );
291                                break Err(SchedulerError::SchemaValidation(error));
292                            }
293                            let _ = events.send(AgentEvent::SchemaRetry {
294                                run_id,
295                                agent_id: task.agent_id,
296                                attempt: schema_retry_count,
297                                max: self.config.retry.schema_retry_max,
298                            });
299                            tracing::warn!(
300                                error = %error,
301                                attempt = schema_retry_count,
302                                "schema validation failed, retrying with feedback"
303                            );
304                            let schema_json =
305                                serde_json::to_string_pretty(schema).unwrap_or_default();
306                            let last_output = fallback_text
307                                .clone()
308                                .unwrap_or_else(|| {
309                                    serde_json::to_string_pretty(&result.output)
310                                        .unwrap_or_default()
311                                });
312                            task.prompt = if fallback_text.is_some() {
313                                format!(
314                                    "{original_prompt}\n\n\
315                                     ---\n\
316                                     You returned your result as plain text instead of calling the `workflow_validate_schema` tool.\n\
317                                     You MUST call the `workflow_validate_schema` tool to submit your result.\n\
318                                     Do NOT return the result as a text message.\n\
319                                     \n\
320                                     Your text output was:\n\
321                                     ```\n{last_output}\n```\n\
322                                     \n\
323                                     Required JSON Schema:\n\
324                                     ```json\n{schema}\n```",
325                                    original_prompt = original_prompt,
326                                    last_output = last_output,
327                                    schema = schema_json,
328                                )
329                            } else {
330                                format!(
331                                    "{original_prompt}\n\n\
332                                     ---\n\
333                                     Your previous response did not match the required schema.\n\
334                                     Error: {error}\n\
335                                     \n\
336                                     Your output was:\n\
337                                     ```json\n{last_output}\n```\n\
338                                     \n\
339                                     Required JSON Schema:\n\
340                                     ```json\n{schema}\n```\n\
341                                     \n\
342                                     Call the `workflow_validate_schema` tool with a JSON object that\n\
343                                     matches this schema exactly. Include ALL required fields.",
344                                    original_prompt = original_prompt,
345                                    error = error,
346                                    last_output = last_output,
347                                    schema = schema_json,
348                                )
349                            };
350                            continue;
351                        }
352                    }
353                    break Ok(result);
354                }
355                Err(e) => {
356                    if agent_token.is_cancelled() || matches!(e, BackendError::Cancelled) {
357                        tracing::debug!("agent cancelled");
358                        break Err(cancel_kind(&run_cancel));
359                    }
360                    if !e.is_retryable() {
361                        tracing::error!(error = %e, "non-retryable backend error");
362                        break Err(SchedulerError::NonRetryable(e));
363                    }
364                    attempt += 1;
365                    if attempt > self.config.retry.max_attempts {
366                        tracing::error!(attempts = attempt, error = %e, "agent exhausted retries");
367                        break Err(SchedulerError::Exhausted {
368                            attempts: attempt,
369                            source: e,
370                        });
371                    }
372                    let backoff = self.config.retry.backoff(attempt);
373                    tracing::warn!(
374                        attempt, backoff_ms = backoff.as_millis() as u64, error = %e,
375                        "retryable backend error; retrying"
376                    );
377                    tokio::select! {
378                        _ = tokio::time::sleep(backoff) => {}
379                        _ = agent_token.cancelled() => break Err(cancel_kind(&run_cancel)),
380                    }
381                }
382            }
383        };
384
385        let elapsed_ms = start.elapsed().as_millis() as u64;
386        let (status, tokens) = match &outcome {
387            Ok(r) => (r.status.clone(), r.tokens_used),
388            Err(SchedulerError::AgentCancelled) | Err(SchedulerError::RunCancelled) => {
389                (AgentStatus::Cancelled, TokenUsage::default())
390            }
391            Err(_) => (AgentStatus::Error, TokenUsage::default()),
392        };
393        let _ = events.send(AgentEvent::AgentDone {
394            run_id,
395            agent_id: task.agent_id,
396            status: status.clone(),
397            tokens,
398            elapsed_ms,
399            name: task.name.clone(),
400            agent_seq: task.agent_seq,
401            output: match &outcome {
402                Ok(r) => r.output.clone(),
403                Err(_) => serde_json::Value::Null,
404            },
405            findings: match &outcome {
406                Ok(r) => r.findings.clone(),
407                Err(_) => Vec::new(),
408            },
409            prompt: task.prompt.clone(),
410            retry_count: attempt,
411            ts: Utc::now(),
412        });
413        tracing::info!(?status, elapsed_ms, "agent finished");
414
415        // Invoke journal callback if configured (M1 transparent persistence).
416        if let Some(ref cb) = self.journal_callback {
417            let output = match &outcome {
418                Ok(r) => r.output.clone(),
419                Err(_) => serde_json::Value::Null,
420            };
421            let agent_status = status.clone();
422            let tokens_used = tokens;
423            let agent_id = task.agent_id;
424            let phase_id = task.phase_id;
425            cb.on_agent_done(agent_id, phase_id, agent_status, output, tokens_used)
426                .await;
427        }
428
429        drop(permit);
430        self.cleanup_agent(run_id, task.agent_id);
431        outcome
432    }
433
434    /// Run a batch of tasks concurrently (the `parallel()` primitive). Bounded
435    /// by the same global semaphore; does not short-circuit on failure — results
436    /// preserve input order.
437    pub async fn run_parallel(
438        &self,
439        run_id: RunId,
440        tasks: Vec<(AgentTask, Option<String>)>,
441    ) -> Vec<Result<AgentResult, SchedulerError>> {
442        let futs = tasks.into_iter().map(|(task, backend)| async move {
443            self.run_agent(run_id, task, backend.as_deref()).await
444        });
445        futures::future::join_all(futs).await
446    }
447
448    /// Cancel one agent (fires its token; the backend observes `ctx.cancel`).
449    pub fn cancel_agent(&self, run_id: RunId, agent_id: AgentId) {
450        if let Some(rs) = self.runs.get(&run_id) {
451            if let Some(tok) = rs.agent_cancels.get(&agent_id) {
452                tok.cancel();
453            }
454        }
455    }
456
457    /// Cancel the whole run (all child agent tokens fire).
458    pub fn cancel_run(&self, run_id: RunId) {
459        if let Some(rs) = self.runs.get(&run_id) {
460            rs.run_cancel.cancel();
461        }
462    }
463
464    /// Quota consumed by a run, if initialised.
465    pub fn quota_used(&self, run_id: RunId) -> Option<u32> {
466        self.runs
467            .get(&run_id)
468            .map(|rs| rs.quota_used.load(Ordering::Relaxed))
469    }
470
471    fn cleanup_agent(&self, run_id: RunId, agent_id: AgentId) {
472        if let Some(rs) = self.runs.get(&run_id) {
473            rs.agent_cancels.remove(&agent_id);
474        }
475    }
476}
477
478fn cancel_kind(run_cancel: &CancellationToken) -> SchedulerError {
479    if run_cancel.is_cancelled() {
480        SchedulerError::RunCancelled
481    } else {
482        SchedulerError::AgentCancelled
483    }
484}
485
486fn preview(s: &str) -> String {
487    s.chars().take(60).collect()
488}
489
490#[cfg(test)]
491mod tests {
492    use super::*;
493    use crate::mock_backend::{FailKind, MockBackend, MockBehavior};
494    use std::path::PathBuf;
495    use std::sync::atomic::AtomicUsize;
496    use std::time::Duration;
497    use uuid::Uuid;
498
499    fn fast_config(max_concurrency: usize, quota: u32) -> SchedulerConfig {
500        SchedulerConfig {
501            max_concurrency,
502            quota_per_run: quota,
503            retry: RetryPolicy {
504                max_attempts: 2,
505                initial_backoff: Duration::from_millis(1),
506                backoff_multiplier: 2.0,
507                max_backoff: Duration::from_millis(5),
508                schema_retry_max: 1,
509            },
510        }
511    }
512
513    fn mk_task(prompt: &str) -> AgentTask {
514        AgentTask {
515            agent_id: Uuid::now_v7(),
516            phase_id: 0,
517            prompt: prompt.to_string(),
518            model: None,
519            allowlist: None,
520            workdir: PathBuf::from("."),
521            mcp_endpoint: None,
522            timeout: None,
523            output_schema: None,
524            workdir_override: None,
525            description: None,
526            role: None,
527            name: None,
528            agent_seq: 0,
529            thread_id: None,
530        }
531    }
532
533    fn mk_task_with_schema(prompt: &str) -> AgentTask {
534        let mut task = mk_task(prompt);
535        task.output_schema = Some(serde_json::json!({
536            "type": "object",
537            "properties": {
538                "answer": { "type": "string" }
539            },
540            "required": ["answer"]
541        }));
542        task
543    }
544
545    fn fallback_output(text: &str) -> serde_json::Value {
546        serde_json::json!({
547            "_agent_fallback_text": true,
548            "text": text,
549        })
550    }
551
552    fn ok_result(id: AgentId) -> AgentResult {
553        AgentResult {
554            agent_id: id,
555            status: AgentStatus::Ok,
556            output: serde_json::Value::Null,
557            findings: vec![],
558            tokens_used: TokenUsage::default(),
559            artifacts: vec![],
560            logs: LogRef::default(),
561            thread_id: None,
562        }
563    }
564
565    fn sched_with(backend: Arc<dyn AgentBackend>, cfg: SchedulerConfig) -> Arc<Scheduler> {
566        Scheduler::new(cfg, BackendRegistry::new().with(backend), None)
567    }
568
569    // A backend that records peak concurrency.
570    struct ProbeBackend {
571        cur: Arc<AtomicUsize>,
572        peak: Arc<AtomicUsize>,
573        delay: Duration,
574    }
575
576    #[async_trait::async_trait]
577    impl AgentBackend for ProbeBackend {
578        fn id(&self) -> &'static str {
579            "probe"
580        }
581        fn capabilities(&self) -> AgentCapabilities {
582            AgentCapabilities::default()
583        }
584        fn as_any(&self) -> &dyn std::any::Any {
585            self
586        }
587        async fn run(
588            &self,
589            task: AgentTask,
590            _ctx: RunContext,
591        ) -> Result<AgentResult, BackendError> {
592            let c = self.cur.fetch_add(1, Ordering::SeqCst) + 1;
593            self.peak.fetch_max(c, Ordering::SeqCst);
594            tokio::time::sleep(self.delay).await;
595            self.cur.fetch_sub(1, Ordering::SeqCst);
596            Ok(ok_result(task.agent_id))
597        }
598    }
599
600    #[tokio::test]
601    async fn test_default_config_concurrency() {
602        let c = SchedulerConfig::default().max_concurrency;
603        assert_eq!(c, 1);
604    }
605
606    /// A backend whose `id()` is parameterised and which records that id in the
607    /// agent result output, so per-task routing can be observed.
608    struct IdBackend {
609        id: &'static str,
610    }
611
612    #[async_trait::async_trait]
613    impl AgentBackend for IdBackend {
614        fn id(&self) -> &'static str {
615            self.id
616        }
617        fn capabilities(&self) -> AgentCapabilities {
618            AgentCapabilities::default()
619        }
620        fn as_any(&self) -> &dyn std::any::Any {
621            self
622        }
623        async fn run(&self, task: AgentTask, _ctx: RunContext) -> Result<AgentResult, BackendError> {
624            Ok(AgentResult {
625                agent_id: task.agent_id,
626                status: AgentStatus::Ok,
627                output: serde_json::Value::String(self.id.to_string()),
628                findings: vec![],
629                tokens_used: TokenUsage::default(),
630                artifacts: vec![],
631                logs: LogRef::default(),
632                thread_id: None,
633            })
634        }
635    }
636
637    #[tokio::test]
638    #[serial_test::serial]
639    async fn per_task_backend_routes_to_named_backend() {
640        // No captured current_backend -> the omitted-backend case must fall
641        // back to the registry default (alpha, first-registered).
642        crate::contract::clear_current_backend();
643        let a = Arc::new(IdBackend { id: "alpha" }) as Arc<dyn AgentBackend>;
644        let b = Arc::new(IdBackend { id: "beta" }) as Arc<dyn AgentBackend>;
645        // alpha is the default (first-registered); beta is selectable.
646        let sched = Arc::new(Scheduler::new(
647            fast_config(4, 1000),
648            BackendRegistry::new().with(a).with(b),
649            None,
650        ));
651        let run_id = Uuid::now_v7();
652        let _rx = sched.init_run(run_id, 256);
653
654        // Explicit backend id -> routes to that backend.
655        let r = sched
656            .run_agent(run_id, mk_task("t1"), Some("beta"))
657            .await
658            .unwrap();
659        assert_eq!(r.output, serde_json::Value::String("beta".to_string()));
660
661        // Omitted backend, no current_backend captured -> registry default (alpha).
662        let r = sched.run_agent(run_id, mk_task("t2"), None).await.unwrap();
663        assert_eq!(r.output, serde_json::Value::String("alpha".to_string()));
664        crate::contract::clear_current_backend();
665    }
666
667    #[tokio::test]
668    #[serial_test::serial]
669    async fn per_task_backend_follows_current_backend() {
670        // When the ACP handshake has captured a current backend, an omitted
671        // `backend` routes to it (via its registry id), even if it is not the
672        // registry default.
673        crate::contract::clear_current_backend();
674        let a = Arc::new(IdBackend { id: "alpha" }) as Arc<dyn AgentBackend>;
675        let b = Arc::new(IdBackend { id: "beta" }) as Arc<dyn AgentBackend>;
676        let sched = Arc::new(Scheduler::new(
677            fast_config(4, 1000),
678            BackendRegistry::new().with(a).with(b),
679            None,
680        ));
681        let run_id = Uuid::now_v7();
682        let _rx = sched.init_run(run_id, 256);
683
684        // Simulate the ACP handshake capturing "beta" as the current backend.
685        crate::contract::set_current_backend(crate::contract::CurrentBackend {
686            id: "beta".to_string(),
687            name: "beta".to_string(),
688            version: "0".to_string(),
689            title: None,
690            client: crate::contract::ClientIdentity {
691                name: "luft".to_string(),
692                version: "test".to_string(),
693                title: None,
694            },
695        });
696
697        // Omitted backend -> routes to current_backend ("beta"), not the
698        // registry default ("alpha").
699        let r = sched.run_agent(run_id, mk_task("t"), None).await.unwrap();
700        assert_eq!(r.output, serde_json::Value::String("beta".to_string()));
701        crate::contract::clear_current_backend();
702    }
703
704    #[tokio::test]
705    #[serial_test::serial]
706    async fn per_task_backend_falls_back_when_current_unregistered() {
707        // current_backend captured but its id is not registered -> fall back
708        // to the registry default instead of erroring.
709        crate::contract::clear_current_backend();
710        let a = Arc::new(IdBackend { id: "alpha" }) as Arc<dyn AgentBackend>;
711        let sched = Arc::new(Scheduler::new(
712            fast_config(4, 1000),
713            BackendRegistry::new().with(a),
714            None,
715        ));
716        let run_id = Uuid::now_v7();
717        let _rx = sched.init_run(run_id, 256);
718
719        crate::contract::set_current_backend(crate::contract::CurrentBackend {
720            id: "gone".to_string(),
721            name: "gone".to_string(),
722            version: "0".to_string(),
723            title: None,
724            client: crate::contract::ClientIdentity {
725                name: "luft".to_string(),
726                version: "test".to_string(),
727                title: None,
728            },
729        });
730        let r = sched.run_agent(run_id, mk_task("t"), None).await.unwrap();
731        assert_eq!(r.output, serde_json::Value::String("alpha".to_string()));
732        crate::contract::clear_current_backend();
733    }
734
735    #[tokio::test]
736    async fn per_task_backend_unknown_id_errors() {
737        let a = Arc::new(IdBackend { id: "alpha" }) as Arc<dyn AgentBackend>;
738        let sched = Arc::new(Scheduler::new(
739            fast_config(4, 1000),
740            BackendRegistry::new().with(a),
741            None,
742        ));
743        let run_id = Uuid::now_v7();
744        let _rx = sched.init_run(run_id, 256);
745        assert!(sched
746            .run_agent(run_id, mk_task("t"), Some("nope"))
747            .await
748            .is_err());
749    }
750
751    #[tokio::test]
752    async fn test_concurrency_limit() {
753        let cur = Arc::new(AtomicUsize::new(0));
754        let peak = Arc::new(AtomicUsize::new(0));
755        let backend = Arc::new(ProbeBackend {
756            cur: cur.clone(),
757            peak: peak.clone(),
758            delay: Duration::from_millis(40),
759        });
760        let sched = sched_with(backend, fast_config(2, 1000));
761        let run_id = Uuid::now_v7();
762        let _rx = sched.init_run(run_id, 256);
763
764        let tasks: Vec<_> = (0..6).map(|i| (mk_task(&format!("t{i}")), None)).collect();
765        let results = sched.run_parallel(run_id, tasks).await;
766
767        assert!(results.iter().all(|r| r.is_ok()));
768        assert!(
769            peak.load(Ordering::SeqCst) <= 2,
770            "peak {}",
771            peak.load(Ordering::SeqCst)
772        );
773    }
774
775    #[tokio::test]
776    async fn test_quota_exceeded() {
777        let backend = Arc::new(MockBackend::new(
778            "mock",
779            vec![MockBehavior::Success {
780                output: serde_json::Value::Null,
781                tokens: TokenUsage::default(),
782                delay: Duration::from_millis(5),
783            }],
784        ));
785        let sched = sched_with(backend, fast_config(8, 3));
786        let run_id = Uuid::now_v7();
787        let _rx = sched.init_run(run_id, 256);
788
789        let tasks: Vec<_> = (0..4).map(|i| (mk_task(&format!("t{i}")), None)).collect();
790        let results = sched.run_parallel(run_id, tasks).await;
791
792        let ok = results.iter().filter(|r| r.is_ok()).count();
793        let quota_err = results
794            .iter()
795            .filter(|r| matches!(r, Err(SchedulerError::QuotaExceeded { .. })))
796            .count();
797        assert_eq!(ok, 3);
798        assert_eq!(quota_err, 1);
799    }
800
801    #[tokio::test]
802    async fn test_retry_on_retryable_error() {
803        let backend = Arc::new(MockBackend::new(
804            "mock",
805            vec![
806                MockBehavior::fail(FailKind::Spawn),
807                MockBehavior::fail(FailKind::Spawn),
808                MockBehavior::Success {
809                    output: serde_json::Value::Null,
810                    tokens: TokenUsage::default(),
811                    delay: Duration::ZERO,
812                },
813            ],
814        ));
815        let probe = backend.clone();
816        let sched = sched_with(backend, fast_config(4, 1000));
817        let run_id = Uuid::now_v7();
818        let _rx = sched.init_run(run_id, 64);
819
820        let r = sched.run_agent(run_id, mk_task("x"), None).await;
821        assert!(r.is_ok(), "{r:?}");
822        assert_eq!(probe.call_count(), 3);
823    }
824
825    #[tokio::test]
826    async fn test_no_retry_on_non_retryable() {
827        let backend = Arc::new(MockBackend::new(
828            "mock",
829            vec![MockBehavior::fail(FailKind::Protocol)],
830        ));
831        let probe = backend.clone();
832        let sched = sched_with(backend, fast_config(4, 1000));
833        let run_id = Uuid::now_v7();
834        let _rx = sched.init_run(run_id, 64);
835
836        let r = sched.run_agent(run_id, mk_task("x"), None).await;
837        assert!(matches!(r, Err(SchedulerError::NonRetryable(_))), "{r:?}");
838        assert_eq!(probe.call_count(), 1);
839    }
840
841    #[tokio::test]
842    async fn test_retry_exhausted() {
843        let backend = Arc::new(MockBackend::new(
844            "mock",
845            vec![MockBehavior::fail(FailKind::Spawn)],
846        ));
847        let probe = backend.clone();
848        let sched = sched_with(backend, fast_config(4, 1000));
849        let run_id = Uuid::now_v7();
850        let _rx = sched.init_run(run_id, 64);
851
852        let r = sched.run_agent(run_id, mk_task("x"), None).await;
853        assert!(
854            matches!(r, Err(SchedulerError::Exhausted { attempts: 3, .. })),
855            "{r:?}"
856        );
857        assert_eq!(probe.call_count(), 3);
858    }
859
860    #[tokio::test]
861    async fn test_schema_fallback_then_succeeds() {
862        let backend = Arc::new(MockBackend::new(
863            "mock",
864            vec![
865                MockBehavior::Success {
866                    output: fallback_output("i forgot the tool"),
867                    tokens: TokenUsage::default(),
868                    delay: Duration::ZERO,
869                },
870                MockBehavior::Success {
871                    output: serde_json::json!({"answer": "ok"}),
872                    tokens: TokenUsage::default(),
873                    delay: Duration::ZERO,
874                },
875            ],
876        ));
877        let probe = backend.clone();
878        let sched = sched_with(backend, fast_config(4, 1000));
879        let run_id = Uuid::now_v7();
880        let mut rx = sched.init_run(run_id, 64);
881
882        let task = mk_task_with_schema("respond");
883        let r = sched.run_agent(run_id, task, None).await;
884        assert!(r.is_ok(), "{r:?}");
885        assert_eq!(probe.call_count(), 2);
886
887        let mut prompt_with_feedback = None;
888        while let Ok(event) = rx.try_recv() {
889            if let AgentEvent::AgentDone { prompt, .. } = event {
890                prompt_with_feedback = Some(prompt);
891            }
892        }
893        let prompt = prompt_with_feedback.expect("AgentDone event with prompt");
894        assert!(prompt.contains("workflow_validate_schema"));
895        assert!(prompt.contains("Required JSON Schema"));
896    }
897
898    #[tokio::test]
899    async fn test_schema_mismatch_then_succeeds() {
900        let backend = Arc::new(MockBackend::new(
901            "mock",
902            vec![
903                MockBehavior::Success {
904                    output: serde_json::json!({"wrong": "field"}),
905                    tokens: TokenUsage::default(),
906                    delay: Duration::ZERO,
907                },
908                MockBehavior::Success {
909                    output: serde_json::json!({"answer": "ok"}),
910                    tokens: TokenUsage::default(),
911                    delay: Duration::ZERO,
912                },
913            ],
914        ));
915        let probe = backend.clone();
916        let sched = sched_with(backend, fast_config(4, 1000));
917        let run_id = Uuid::now_v7();
918        let _rx = sched.init_run(run_id, 64);
919
920        let task = mk_task_with_schema("respond");
921        let r = sched.run_agent(run_id, task, None).await;
922        assert!(r.is_ok(), "{r:?}");
923        assert_eq!(probe.call_count(), 2);
924    }
925
926    #[tokio::test]
927    async fn test_schema_fallback_exhausted() {
928        let backend = Arc::new(MockBackend::new(
929            "mock",
930            vec![MockBehavior::Success {
931                output: fallback_output("still no tool"),
932                tokens: TokenUsage::default(),
933                delay: Duration::ZERO,
934            }],
935        ));
936        let probe = backend.clone();
937        let sched = sched_with(backend, fast_config(4, 1000));
938        let run_id = Uuid::now_v7();
939        let _rx = sched.init_run(run_id, 64);
940
941        let task = mk_task_with_schema("respond");
942        let r = sched.run_agent(run_id, task, None).await;
943        assert!(
944            matches!(r, Err(SchedulerError::SchemaValidation(_))),
945            "{r:?}"
946        );
947        assert_eq!(probe.call_count(), 2);
948    }
949
950    #[tokio::test]
951    async fn test_cancel_run() {
952        let backend = Arc::new(MockBackend::new("mock", vec![MockBehavior::Hang]));
953        let sched = sched_with(backend, fast_config(8, 1000));
954        let run_id = Uuid::now_v7();
955        let _rx = sched.init_run(run_id, 64);
956
957        let s2 = sched.clone();
958        let handle = tokio::spawn(async move {
959            let tasks: Vec<_> = (0..3).map(|i| (mk_task(&format!("h{i}")), None)).collect();
960            s2.run_parallel(run_id, tasks).await
961        });
962        tokio::time::sleep(Duration::from_millis(20)).await;
963        sched.cancel_run(run_id);
964
965        let results = handle.await.unwrap();
966        assert_eq!(results.len(), 3);
967        assert!(results
968            .iter()
969            .all(|r| matches!(r, Err(SchedulerError::RunCancelled))));
970    }
971
972    #[tokio::test]
973    async fn test_cancel_agent() {
974        let backend = Arc::new(MockBackend::new("mock", vec![MockBehavior::Hang]));
975        let sched = sched_with(backend, fast_config(8, 1000));
976        let run_id = Uuid::now_v7();
977        let _rx = sched.init_run(run_id, 64);
978
979        let task = mk_task("hang");
980        let agent_id = task.agent_id;
981        let s2 = sched.clone();
982        let handle = tokio::spawn(async move { s2.run_agent(run_id, task, None).await });
983        tokio::time::sleep(Duration::from_millis(20)).await;
984        sched.cancel_agent(run_id, agent_id);
985
986        let r = handle.await.unwrap();
987        assert!(matches!(r, Err(SchedulerError::AgentCancelled)), "{r:?}");
988    }
989
990    #[tokio::test]
991    async fn test_parallel_partial_failure() {
992        let backend = Arc::new(MockBackend::new(
993            "mock",
994            vec![
995                MockBehavior::Success {
996                    output: serde_json::Value::Null,
997                    tokens: TokenUsage::default(),
998                    delay: Duration::ZERO,
999                },
1000                MockBehavior::fail(FailKind::Protocol),
1001                MockBehavior::Success {
1002                    output: serde_json::Value::Null,
1003                    tokens: TokenUsage::default(),
1004                    delay: Duration::ZERO,
1005                },
1006            ],
1007        ));
1008        let sched = sched_with(backend, fast_config(1, 1000)); // serialize for deterministic behavior order
1009        let run_id = Uuid::now_v7();
1010        let _rx = sched.init_run(run_id, 64);
1011
1012        let tasks: Vec<_> = (0..3).map(|i| (mk_task(&format!("p{i}")), None)).collect();
1013        let results = sched.run_parallel(run_id, tasks).await;
1014
1015        assert_eq!(results.len(), 3);
1016        assert_eq!(results.iter().filter(|r| r.is_ok()).count(), 2);
1017        assert_eq!(results.iter().filter(|r| r.is_err()).count(), 1);
1018    }
1019
1020    #[tokio::test]
1021    async fn test_event_sequence() {
1022        let backend = Arc::new(MockBackend::new(
1023            "mock",
1024            vec![MockBehavior::Success {
1025                output: serde_json::Value::Null,
1026                tokens: TokenUsage {
1027                    input: 10,
1028                    output: 5,
1029                    ..Default::default()
1030                },
1031                delay: Duration::ZERO,
1032            }],
1033        ));
1034        let sched = sched_with(backend, fast_config(4, 1000));
1035        let run_id = Uuid::now_v7();
1036        let mut rx = sched.init_run(run_id, 64);
1037
1038        let r = sched.run_agent(run_id, mk_task("x"), None).await;
1039        assert!(r.is_ok());
1040
1041        let e1 = rx.recv().await.unwrap();
1042        assert!(matches!(e1, AgentEvent::AgentStarted { .. }), "{e1:?}");
1043        let e2 = rx.recv().await.unwrap();
1044        match e2 {
1045            AgentEvent::AgentDone { status, tokens, .. } => {
1046                assert_eq!(status, AgentStatus::Ok);
1047                assert_eq!(tokens.input, 10);
1048            }
1049            other => panic!("expected AgentDone, got {other:?}"),
1050        }
1051    }
1052}
luft_core/scheduler/mod.rs

luft_core/scheduler/
mod.rs
