Skip to main content

leviath_runtime/
inference_bridge.rs

1//! The async worker side of the ECS inference stage - the sync-ECS ↔ async-I/O
2//! bridge for inference.
3//!
4//! Systems can't `.await`, and a single inference can take up to an hour, so the
5//! inference-dispatch system never runs the network call itself. Instead it
6//! builds an [`InferenceJob`] (an agent's request plus the per-model pool permit
7//! it acquired) and `tokio::spawn`s [`run_inference_job`]. That short-lived task
8//! performs the call with the permit held, reports an [`InferenceOutcome`] on the
9//! results channel, and wakes the tick loop; the inference-collect system drains
10//! outcomes on a later tick and applies them back to the agents.
11//!
12//! One task exists per *in-flight request* (bounded by the per-model
13//! [`InferencePools`](crate::inference_pool::InferencePools) permits), **never**
14//! one per agent - that is what keeps CPU bounded by work, not by agent count.
15
16use std::sync::Arc;
17use std::time::Duration;
18
19use bevy_ecs::entity::Entity;
20use leviath_providers::{InferenceRequest, InferenceResponse, Provider, ProviderError};
21use tokio::sync::Notify;
22use tokio::sync::mpsc::UnboundedSender;
23
24use crate::inference_pool::InferencePermit;
25
26/// How a transient inference failure is retried before the agent is failed.
27///
28/// Transient errors (see [`ProviderError::is_transient`]) are retried with
29/// exponential backoff; a permanent error (auth, invalid request, token limit)
30/// fails immediately. This keeps a passing network blip from marking a stage
31/// `error` and carrying its half-finished work forward.
32#[derive(Debug, Clone, Copy)]
33pub struct RetryPolicy {
34    /// Total attempts including the first (e.g. `4` = one try + three retries).
35    pub max_attempts: u32,
36    /// Base backoff; the retry after attempt `n` waits `base_delay * 2^(n-1)`
37    /// (so 1s, 2s, 4s, … for a 1s base).
38    pub base_delay: Duration,
39    /// Hard ceiling on the total wall-clock time one job (all attempts +
40    /// backoffs) may run before it is aborted and its pool slot freed.
41    ///
42    /// Providers apply this same deadline as their own per-call timeout, so a
43    /// call normally ends there. This outer bound is the backstop: a provider
44    /// timer can be defeated (e.g. a connection that keeps trickling keepalive
45    /// bytes without ever completing resets a read timer forever), which would
46    /// leak the permit until the model's pool fills and new agents never get a
47    /// slot. Wrapping the whole job in a wall-clock timeout guarantees the slot
48    /// is released within a fixed time regardless. Defaults to
49    /// [`leviath_providers::DEFAULT_INFERENCE_TIMEOUT_SECS`]; a stage's
50    /// `request_timeout_secs` overrides it per stage.
51    pub job_timeout: Duration,
52}
53
54impl Default for RetryPolicy {
55    fn default() -> Self {
56        Self {
57            max_attempts: 4,
58            base_delay: Duration::from_secs(1),
59            // The unified default inference deadline. A stage's
60            // `request_timeout_secs` overrides this per stage (see
61            // `pipeline::retry_policy_for`); the providers apply the same value
62            // as their own per-call timeout, so this is the single outer bound
63            // that also frees the pool slot if a provider's timer is defeated.
64            job_timeout: Duration::from_secs(leviath_providers::DEFAULT_INFERENCE_TIMEOUT_SECS),
65        }
66    }
67}
68
69/// A unit of inference work the dispatch system hands to the worker pool.
70pub struct InferenceJob {
71    /// The agent this inference is for.
72    pub entity: Entity,
73    /// The provider to call (already resolved for the agent's model).
74    pub provider: Arc<dyn Provider>,
75    /// The assembled request.
76    pub request: InferenceRequest,
77    /// The per-model pool permit, held for the whole request and released when
78    /// the job finishes.
79    pub permit: InferencePermit,
80    /// When set, count the assembled request's tokens exactly (via the
81    /// provider's `count_tokens`, which uses a remote endpoint where available)
82    /// before calling `infer`, and fail early if it would exceed the model's
83    /// context window. Off by default - the runtime's cheap `len/4` estimates
84    /// drive normal budgeting; this is the opt-in accurate guard.
85    pub exact_token_counting: bool,
86}
87
88/// Flatten a request into the text whose tokens we count for the budget guard:
89/// system blocks, every message's textual content, and each tool's name +
90/// description + JSON schema. This mirrors what the provider sends closely enough
91/// for a context-window check (exact per-message/role overhead is the provider's
92/// to add; the bulk is this text).
93fn flatten_request_text(request: &InferenceRequest) -> String {
94    let mut parts: Vec<String> = Vec::new();
95    for block in &request.system {
96        parts.push(block.text.clone());
97    }
98    for msg in &request.messages {
99        parts.push(msg.content.as_text());
100    }
101    for tool in &request.tools {
102        parts.push(tool.name.clone());
103        parts.push(tool.description.clone());
104        parts.push(tool.parameters.to_string());
105    }
106    parts.join("\n")
107}
108
109/// The completed result of an [`InferenceJob`], applied on a later tick by the
110/// inference-collect system.
111pub struct InferenceOutcome {
112    /// The agent the result belongs to.
113    pub entity: Entity,
114    /// The provider's response, or the error it failed with.
115    pub result: Result<InferenceResponse, ProviderError>,
116    /// Wall-clock time the job took, retries and backoff included. Measured
117    /// here because the ECS only sees the outcome land on a later tick; this
118    /// is the only place the call's real duration exists.
119    pub latency: std::time::Duration,
120}
121
122/// Run one inference job to completion: perform the (possibly hour-long) network
123/// call with the pool permit held, release the slot, report the outcome, and
124/// wake the tick loop.
125///
126/// Meant to be `tokio::spawn`ed by the dispatch system. If the results receiver
127/// has been dropped (the world is shutting down) the send is a harmless no-op.
128pub async fn run_inference_job(
129    job: InferenceJob,
130    results: UnboundedSender<InferenceOutcome>,
131    wake: Arc<Notify>,
132    retry: RetryPolicy,
133    cancel: crate::cancel::CancelToken,
134) {
135    let InferenceJob {
136        entity,
137        provider,
138        request,
139        permit,
140        exact_token_counting,
141    } = job;
142    let started = std::time::Instant::now();
143    // Opt-in accurate pre-flight budget guard: count the assembled request
144    // exactly (remote endpoint where the provider has one, heuristic otherwise)
145    // and refuse a request that would overflow the model's context window,
146    // rather than sending it and letting the provider reject it after the fact.
147    if exact_token_counting {
148        let text = flatten_request_text(&request);
149        let used = provider.count_tokens(&text, &request.model).await;
150        let max = provider.max_context_tokens(&request.model);
151        if used.saturating_add(request.max_tokens) > max {
152            drop(permit);
153            let _ = results.send(InferenceOutcome {
154                entity,
155                result: Err(ProviderError::TokenLimitExceeded { used, max }),
156                latency: started.elapsed(),
157            });
158            wake.notify_one();
159            return;
160        }
161    }
162    // Retry transient failures (connection reset, timeout, 429, 5xx) with
163    // exponential backoff, holding the permit across the backoff; a permanent
164    // error fails immediately. The whole thing is bounded by `job_timeout` so a
165    // never-completing (stalled-stream) call cannot hold the pool slot forever.
166    let attempts = async {
167        let mut attempt = 1u32;
168        loop {
169            match provider.infer(request.clone()).await {
170                Ok(response) => break Ok(response),
171                Err(e) if e.is_transient() && attempt < retry.max_attempts => {
172                    tokio::time::sleep(retry.base_delay * 2u32.pow(attempt - 1)).await;
173                    attempt += 1;
174                }
175                Err(e) => break Err(e),
176            }
177        }
178    };
179    // A cancel drops the whole retry-and-backoff future - aborting the in-flight
180    // HTTP request rather than waiting out the job timeout (up to 15 minutes) -
181    // and reports nothing: the agent is already terminal, so there is no outcome
182    // to apply. Releasing the permit here is the point; a cancelled run used to
183    // hold its model's pool slot for as long as the provider took to answer.
184    let result = tokio::select! {
185        biased;
186        _ = cancel.cancelled() => {
187            drop(permit);
188            return;
189        }
190        outcome = tokio::time::timeout(retry.job_timeout, attempts) => match outcome {
191            Ok(result) => result,
192            Err(_elapsed) => Err(leviath_providers::ProviderError::Other(format!(
193                "inference exceeded the {}s job timeout and was aborted to free the \
194                 pool slot (a stalled or never-completing response)",
195                retry.job_timeout.as_secs()
196            ))),
197        },
198    };
199    drop(permit); // free the pool slot before the collect system runs
200    let _ = results.send(InferenceOutcome {
201        entity,
202        result,
203        latency: started.elapsed(),
204    });
205    wake.notify_one();
206}
207
208#[cfg(test)]
209mod tests {
210    use super::*;
211    use crate::inference_pool::{InferencePoolConfig, InferencePools};
212    use tokio::sync::mpsc;
213
214    fn test_request() -> InferenceRequest {
215        InferenceRequest {
216            system: vec![],
217            messages: vec![],
218            model: "m".to_string(),
219            max_tokens: 100,
220            temperature: 0.0,
221            tools: vec![],
222            extra: serde_json::Value::Null,
223            request_timeout_secs: None,
224        }
225    }
226
227    fn response(text: &str) -> InferenceResponse {
228        InferenceResponse {
229            content: text.to_string(),
230            tool_calls: vec![],
231            tokens_used: leviath_providers::TokenUsage {
232                prompt_tokens: 1,
233                completion_tokens: 1,
234                total_tokens: 2,
235                cached_tokens: 0,
236                cache_write_tokens: 0,
237            },
238            finish_reason: leviath_providers::FinishReason::Complete,
239        }
240    }
241
242    /// A provider that returns a fixed success or error for `infer` (avoids
243    /// cloning `ProviderError`, which isn't `Clone`).
244    enum Fixed {
245        Ok(InferenceResponse),
246        Err(String),
247    }
248
249    #[async_trait::async_trait]
250    impl Provider for Fixed {
251        async fn infer(
252            &self,
253            _req: InferenceRequest,
254        ) -> leviath_providers::Result<InferenceResponse> {
255            match self {
256                Fixed::Ok(r) => Ok(r.clone()),
257                Fixed::Err(m) => Err(ProviderError::Other(m.clone())),
258            }
259        }
260        async fn count_tokens(&self, _text: &str, _model: &str) -> usize {
261            1
262        }
263        fn max_context_tokens(&self, _model: &str) -> usize {
264            100_000
265        }
266        fn name(&self) -> &str {
267            "fixed"
268        }
269        fn capabilities(&self, _model: &str) -> leviath_providers::ModelCapabilities {
270            leviath_providers::ModelCapabilities::default()
271        }
272    }
273
274    fn job(provider: Arc<dyn Provider>) -> InferenceJob {
275        let pools = InferencePools::new(InferencePoolConfig::new());
276        InferenceJob {
277            entity: Entity::from_raw_u32(7)
278                .expect("a small literal index is always a valid entity id"),
279            provider,
280            request: test_request(),
281            permit: pools.try_acquire("m").expect("free pool"),
282            exact_token_counting: false,
283        }
284    }
285
286    /// Cancelling releases the pool slot immediately instead of waiting out the
287    /// job timeout (15 minutes by default), and reports no outcome - the agent is
288    /// already terminal, so there is nothing to apply.
289    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
290    async fn a_cancelled_job_frees_its_pool_slot_without_reporting() {
291        let mut cfg = InferencePoolConfig::new();
292        cfg.set_limit("m", 1);
293        let pools = InferencePools::new(cfg);
294        let permit = pools.try_acquire("m").expect("free pool");
295        assert!(pools.try_acquire("m").is_none(), "pool should be full");
296
297        // A provider that never answers - the stalled-call case a cancel exists
298        // to escape.
299        let provider = Arc::new(Scripted {
300            steps: std::sync::Mutex::new(vec![Step::Hang].into()),
301            calls: std::sync::Mutex::new(0),
302        });
303        let job = InferenceJob {
304            entity: Entity::from_raw_u32(7)
305                .expect("a small literal index is always a valid entity id"),
306            provider,
307            request: test_request(),
308            permit,
309            exact_token_counting: false,
310        };
311        let (tx, mut rx) = mpsc::unbounded_channel();
312        let cancel = crate::cancel::CancelToken::new();
313        let running = tokio::spawn(run_inference_job(
314            job,
315            tx,
316            Arc::new(Notify::new()),
317            // A job timeout far longer than the test: only the cancel can end it.
318            RetryPolicy {
319                max_attempts: 1,
320                base_delay: Duration::ZERO,
321                job_timeout: Duration::from_secs(3600),
322            },
323            cancel.clone(),
324        ));
325        tokio::task::yield_now().await;
326        cancel.cancel();
327
328        tokio::time::timeout(Duration::from_secs(5), running)
329            .await
330            .expect("the cancel ended the job")
331            .unwrap();
332        assert!(
333            pools.try_acquire("m").is_some(),
334            "the pool slot is free for the next agent"
335        );
336        assert!(
337            rx.try_recv().is_err(),
338            "and no outcome is reported for a cancelled run"
339        );
340    }
341
342    #[tokio::test]
343    async fn run_job_aborts_a_hung_call_and_frees_the_pool_slot() {
344        // A model pool of one slot, taken by the (hung) job under test.
345        let mut cfg = InferencePoolConfig::new();
346        cfg.set_limit("m", 1);
347        let pools = InferencePools::new(cfg);
348        let permit = pools.try_acquire("m").expect("free pool");
349        assert!(pools.try_acquire("m").is_none(), "pool should be full");
350
351        let provider = Arc::new(Scripted {
352            steps: std::sync::Mutex::new(vec![Step::Hang].into()),
353            calls: std::sync::Mutex::new(0),
354        });
355        let job = InferenceJob {
356            entity: Entity::from_raw_u32(7)
357                .expect("a small literal index is always a valid entity id"),
358            provider,
359            request: test_request(),
360            permit,
361            exact_token_counting: false,
362        };
363        let (tx, mut rx) = mpsc::unbounded_channel();
364        let policy = RetryPolicy {
365            max_attempts: 1,
366            base_delay: Duration::ZERO,
367            job_timeout: Duration::from_millis(50),
368        };
369        run_inference_job(
370            job,
371            tx,
372            Arc::new(Notify::new()),
373            policy,
374            crate::cancel::CancelToken::new(),
375        )
376        .await;
377
378        // The hung call was aborted with a timeout error…
379        let outcome = rx.try_recv().expect("outcome sent");
380        let err = outcome.result.expect_err("hung call should error");
381        assert!(err.to_string().contains("job timeout"), "got: {err}");
382        // …and its pool slot is free again for the next agent.
383        assert!(
384            pools.try_acquire("m").is_some(),
385            "the slot must be released after the timeout"
386        );
387    }
388
389    #[tokio::test]
390    async fn run_job_reports_ok_and_wakes() {
391        let (tx, mut rx) = mpsc::unbounded_channel();
392        let wake = Arc::new(Notify::new());
393        run_inference_job(
394            job(Arc::new(Fixed::Ok(response("hi")))),
395            tx,
396            wake.clone(),
397            RetryPolicy::default(),
398            crate::cancel::CancelToken::new(),
399        )
400        .await;
401
402        let outcome = rx.try_recv().expect("outcome sent");
403        assert_eq!(
404            outcome.entity,
405            Entity::from_raw_u32(7).expect("a small literal index is always a valid entity id")
406        );
407        assert_eq!(outcome.result.unwrap().content, "hi");
408        // The wake was signalled (a subsequent notified() returns immediately).
409        wake.notified().await;
410    }
411
412    #[tokio::test]
413    async fn run_job_reports_provider_error() {
414        let (tx, mut rx) = mpsc::unbounded_channel();
415        let wake = Arc::new(Notify::new());
416        let err = Arc::new(Fixed::Err("boom".to_string()));
417        run_inference_job(
418            job(err),
419            tx,
420            wake,
421            RetryPolicy::default(),
422            crate::cancel::CancelToken::new(),
423        )
424        .await;
425
426        let outcome = rx.try_recv().expect("outcome sent");
427        assert!(outcome.result.is_err());
428    }
429
430    /// A provider with a fixed `count_tokens` result and context window, used to
431    /// drive the opt-in pre-inference budget guard. `infer` always succeeds.
432    struct Counter {
433        count: usize,
434        max: usize,
435    }
436
437    #[async_trait::async_trait]
438    impl Provider for Counter {
439        async fn infer(
440            &self,
441            _req: InferenceRequest,
442        ) -> leviath_providers::Result<InferenceResponse> {
443            Ok(response("ok"))
444        }
445        async fn count_tokens(&self, _text: &str, _model: &str) -> usize {
446            self.count
447        }
448        fn max_context_tokens(&self, _model: &str) -> usize {
449            self.max
450        }
451        fn name(&self) -> &str {
452            "counter"
453        }
454        fn capabilities(&self, _model: &str) -> leviath_providers::ModelCapabilities {
455            leviath_providers::ModelCapabilities::default()
456        }
457    }
458
459    fn counting_job(provider: Arc<dyn Provider>, exact: bool) -> InferenceJob {
460        let pools = InferencePools::new(InferencePoolConfig::new());
461        InferenceJob {
462            entity: Entity::from_raw_u32(7)
463                .expect("a small literal index is always a valid entity id"),
464            provider,
465            request: test_request(), // max_tokens: 100
466            permit: pools.try_acquire("m").expect("free pool"),
467            exact_token_counting: exact,
468        }
469    }
470
471    #[test]
472    fn flatten_request_text_includes_system_messages_and_tools() {
473        use leviath_providers::{SystemBlock, Tool};
474        let req = InferenceRequest {
475            system: vec![SystemBlock {
476                text: "sys".to_string(),
477                cache_hint: leviath_core::CacheHint::Never,
478            }],
479            messages: vec![leviath_providers::Message {
480                role: "user".to_string(),
481                content: "hello".into(),
482                cache_breakpoint: false,
483            }],
484            model: "m".to_string(),
485            max_tokens: 10,
486            temperature: 0.0,
487            tools: vec![Tool {
488                name: "search".to_string(),
489                description: "find things".to_string(),
490                parameters: serde_json::json!({"type": "object"}),
491            }],
492            extra: serde_json::Value::Null,
493            request_timeout_secs: None,
494        };
495        let text = flatten_request_text(&req);
496        assert!(text.contains("sys"));
497        assert!(text.contains("hello"));
498        assert!(text.contains("search"));
499        assert!(text.contains("find things"));
500        assert!(text.contains("object"));
501    }
502
503    #[tokio::test]
504    async fn guard_rejects_request_over_context_window() {
505        // count(950) + max_tokens(100) = 1050 > context(1000) ⇒ rejected pre-flight.
506        let (tx, mut rx) = mpsc::unbounded_channel();
507        let provider = Arc::new(Counter {
508            count: 950,
509            max: 1000,
510        });
511        run_inference_job(
512            counting_job(provider, true),
513            tx,
514            Arc::new(Notify::new()),
515            RetryPolicy::default(),
516            crate::cancel::CancelToken::new(),
517        )
518        .await;
519        let outcome = rx.try_recv().expect("outcome sent");
520        let err = outcome.result.expect_err("should be rejected");
521        // Assert on the Display string (branch-free) rather than `matches!`,
522        // whose non-matching arm would be an uncovered region.
523        assert_eq!(err.to_string(), "Token limit exceeded: 950 > 1000");
524    }
525
526    #[tokio::test]
527    async fn guard_allows_request_within_context_window() {
528        // count(800) + 100 = 900 ≤ 1000 ⇒ proceeds to infer.
529        let (tx, mut rx) = mpsc::unbounded_channel();
530        let provider = Arc::new(Counter {
531            count: 800,
532            max: 1000,
533        });
534        run_inference_job(
535            counting_job(provider, true),
536            tx,
537            Arc::new(Notify::new()),
538            RetryPolicy::default(),
539            crate::cancel::CancelToken::new(),
540        )
541        .await;
542        let outcome = rx.try_recv().expect("outcome sent");
543        assert_eq!(outcome.result.expect("should succeed").content, "ok");
544    }
545
546    #[tokio::test]
547    async fn counter_provider_metadata_is_exercised() {
548        // Keep the Counter mock's non-`infer` trait methods measured.
549        let p = Counter { count: 5, max: 10 };
550        assert_eq!(p.name(), "counter");
551        assert_eq!(p.max_context_tokens("m"), 10);
552        assert_eq!(p.count_tokens("t", "m").await, 5);
553        assert!(p.capabilities("m").supports_streaming);
554    }
555
556    #[tokio::test]
557    async fn guard_off_skips_the_count_and_proceeds() {
558        // Even wildly over budget, with the flag off the guard never runs.
559        let (tx, mut rx) = mpsc::unbounded_channel();
560        let provider = Arc::new(Counter {
561            count: 1_000_000,
562            max: 1000,
563        });
564        run_inference_job(
565            counting_job(provider, false),
566            tx,
567            Arc::new(Notify::new()),
568            RetryPolicy::default(),
569            crate::cancel::CancelToken::new(),
570        )
571        .await;
572        let outcome = rx.try_recv().expect("outcome sent");
573        assert_eq!(outcome.result.expect("should succeed").content, "ok");
574    }
575
576    #[tokio::test]
577    async fn fixed_provider_metadata_is_exercised() {
578        // Covers the mock's non-`infer` trait methods (the pipeline resolves
579        // these off the provider elsewhere; here we just keep them measured).
580        let p = Fixed::Ok(response("x"));
581        assert_eq!(p.name(), "fixed");
582        assert_eq!(p.count_tokens("t", "m").await, 1);
583        assert_eq!(p.max_context_tokens("m"), 100_000);
584        let _ = p.capabilities("m");
585    }
586
587    #[tokio::test]
588    async fn run_job_survives_dropped_receiver() {
589        let (tx, rx) = mpsc::unbounded_channel();
590        drop(rx); // world shutting down: nobody to receive
591        let wake = Arc::new(Notify::new());
592        // Must not panic even though the send fails.
593        run_inference_job(
594            job(Arc::new(Fixed::Ok(response("x")))),
595            tx,
596            wake,
597            RetryPolicy::default(),
598            crate::cancel::CancelToken::new(),
599        )
600        .await;
601    }
602
603    // ── retry behavior ──
604
605    enum Step {
606        Ok(String),
607        Transient,
608        Permanent,
609        /// Never returns - a stalled/hung call, for the job-timeout test.
610        Hang,
611    }
612
613    /// A provider that plays a scripted sequence of results and counts calls.
614    struct Scripted {
615        steps: std::sync::Mutex<std::collections::VecDeque<Step>>,
616        calls: std::sync::Mutex<u32>,
617    }
618
619    #[async_trait::async_trait]
620    impl Provider for Scripted {
621        async fn infer(
622            &self,
623            _req: InferenceRequest,
624        ) -> leviath_providers::Result<InferenceResponse> {
625            *self.calls.lock().unwrap() += 1;
626            // Pop before matching so the mutex guard is not held across the
627            // `Hang` arm's `.await` (which would make this future non-`Send`).
628            let step = self.steps.lock().unwrap().pop_front();
629            match step {
630                Some(Step::Ok(t)) => Ok(response(&t)),
631                Some(Step::Transient) => Err(ProviderError::RateLimitExceeded),
632                Some(Step::Permanent) => Err(ProviderError::Other("permanent".to_string())),
633                Some(Step::Hang) => std::future::pending().await,
634                None => Err(ProviderError::Other("exhausted".to_string())),
635            }
636        }
637        async fn count_tokens(&self, _t: &str, _m: &str) -> usize {
638            1
639        }
640        fn max_context_tokens(&self, _m: &str) -> usize {
641            100_000
642        }
643        fn name(&self) -> &str {
644            "scripted"
645        }
646        fn capabilities(&self, _m: &str) -> leviath_providers::ModelCapabilities {
647            leviath_providers::ModelCapabilities::default()
648        }
649    }
650
651    fn no_delay(max_attempts: u32) -> RetryPolicy {
652        RetryPolicy {
653            max_attempts,
654            base_delay: Duration::ZERO,
655            job_timeout: Duration::from_secs(30),
656        }
657    }
658
659    #[tokio::test]
660    async fn run_job_retries_transient_then_succeeds() {
661        let provider = Arc::new(Scripted {
662            steps: std::sync::Mutex::new(
663                vec![
664                    Step::Transient,
665                    Step::Transient,
666                    Step::Ok("done".to_string()),
667                ]
668                .into(),
669            ),
670            calls: std::sync::Mutex::new(0),
671        });
672        let (tx, mut rx) = mpsc::unbounded_channel();
673        run_inference_job(
674            job(provider.clone()),
675            tx,
676            Arc::new(Notify::new()),
677            no_delay(4),
678            crate::cancel::CancelToken::new(),
679        )
680        .await;
681        let outcome = rx.try_recv().expect("outcome sent");
682        assert_eq!(outcome.result.unwrap().content, "done");
683        assert_eq!(*provider.calls.lock().unwrap(), 3); // two retries then success
684    }
685
686    #[tokio::test]
687    async fn run_job_gives_up_after_max_attempts() {
688        let provider = Arc::new(Scripted {
689            steps: std::sync::Mutex::new(
690                vec![
691                    Step::Transient,
692                    Step::Transient,
693                    Step::Transient,
694                    Step::Transient,
695                ]
696                .into(),
697            ),
698            calls: std::sync::Mutex::new(0),
699        });
700        let (tx, mut rx) = mpsc::unbounded_channel();
701        run_inference_job(
702            job(provider.clone()),
703            tx,
704            Arc::new(Notify::new()),
705            no_delay(3),
706            crate::cancel::CancelToken::new(),
707        )
708        .await;
709        let outcome = rx.try_recv().expect("outcome sent");
710        assert!(outcome.result.is_err());
711        assert_eq!(*provider.calls.lock().unwrap(), 3); // exhausted the 3 attempts
712    }
713
714    #[tokio::test]
715    async fn run_job_does_not_retry_a_permanent_error() {
716        let provider = Arc::new(Scripted {
717            steps: std::sync::Mutex::new(vec![Step::Permanent, Step::Ok("x".to_string())].into()),
718            calls: std::sync::Mutex::new(0),
719        });
720        let (tx, mut rx) = mpsc::unbounded_channel();
721        run_inference_job(
722            job(provider.clone()),
723            tx,
724            Arc::new(Notify::new()),
725            no_delay(4),
726            crate::cancel::CancelToken::new(),
727        )
728        .await;
729        let outcome = rx.try_recv().expect("outcome sent");
730        assert!(outcome.result.is_err());
731        assert_eq!(*provider.calls.lock().unwrap(), 1); // no retry on a permanent error
732    }
733
734    #[tokio::test]
735    async fn scripted_provider_metadata_is_exercised() {
736        let p = Scripted {
737            steps: std::sync::Mutex::new(std::collections::VecDeque::new()),
738            calls: std::sync::Mutex::new(0),
739        };
740        assert_eq!(p.name(), "scripted");
741        assert_eq!(p.count_tokens("t", "m").await, 1);
742        assert_eq!(p.max_context_tokens("m"), 100_000);
743        let _ = p.capabilities("m");
744    }
745}