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

luft_core/scheduler/
mod.rs
