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    //
185    // Note this arm sends no outcome and so never reaches the `wake` below: the
186    // tick loop learns the slot is free from the permit's own `Drop` (see
187    // `InferencePools::with_wake`). Without that, this return frees a slot in
188    // silence and every agent queued on this model stays parked (issue #189).
189    let result = tokio::select! {
190        biased;
191        _ = cancel.cancelled() => {
192            drop(permit);
193            return;
194        }
195        outcome = tokio::time::timeout(retry.job_timeout, attempts) => match outcome {
196            Ok(result) => result,
197            Err(_elapsed) => Err(leviath_providers::ProviderError::Other(format!(
198                "inference exceeded the {}s job timeout and was aborted to free the \
199                 pool slot (a stalled or never-completing response)",
200                retry.job_timeout.as_secs()
201            ))),
202        },
203    };
204    drop(permit); // free the pool slot before the collect system runs
205    let _ = results.send(InferenceOutcome {
206        entity,
207        result,
208        latency: started.elapsed(),
209    });
210    wake.notify_one();
211}
212
213#[cfg(test)]
214mod tests {
215    use super::*;
216    use crate::inference_pool::{InferencePoolConfig, InferencePools};
217    use tokio::sync::mpsc;
218
219    fn test_request() -> InferenceRequest {
220        InferenceRequest {
221            system: vec![],
222            messages: vec![],
223            model: "m".to_string(),
224            max_tokens: 100,
225            temperature: 0.0,
226            tools: vec![],
227            extra: serde_json::Value::Null,
228            request_timeout_secs: None,
229        }
230    }
231
232    fn response(text: &str) -> InferenceResponse {
233        InferenceResponse {
234            content: text.to_string(),
235            tool_calls: vec![],
236            tokens_used: leviath_providers::TokenUsage {
237                prompt_tokens: 1,
238                completion_tokens: 1,
239                total_tokens: 2,
240                cached_tokens: 0,
241                cache_write_tokens: 0,
242            },
243            finish_reason: leviath_providers::FinishReason::Complete,
244        }
245    }
246
247    /// A provider that returns a fixed success or error for `infer` (avoids
248    /// cloning `ProviderError`, which isn't `Clone`).
249    enum Fixed {
250        Ok(InferenceResponse),
251        Err(String),
252    }
253
254    #[async_trait::async_trait]
255    impl Provider for Fixed {
256        async fn infer(
257            &self,
258            _req: InferenceRequest,
259        ) -> leviath_providers::Result<InferenceResponse> {
260            match self {
261                Fixed::Ok(r) => Ok(r.clone()),
262                Fixed::Err(m) => Err(ProviderError::Other(m.clone())),
263            }
264        }
265        async fn count_tokens(&self, _text: &str, _model: &str) -> usize {
266            1
267        }
268        fn max_context_tokens(&self, _model: &str) -> usize {
269            100_000
270        }
271        fn name(&self) -> &str {
272            "fixed"
273        }
274        fn capabilities(&self, _model: &str) -> leviath_providers::ModelCapabilities {
275            leviath_providers::ModelCapabilities::default()
276        }
277    }
278
279    fn job(provider: Arc<dyn Provider>) -> InferenceJob {
280        let pools = InferencePools::new(InferencePoolConfig::new());
281        InferenceJob {
282            entity: Entity::from_raw_u32(7)
283                .expect("a small literal index is always a valid entity id"),
284            provider,
285            request: test_request(),
286            permit: pools.try_acquire("m").expect("free pool"),
287            exact_token_counting: false,
288        }
289    }
290
291    /// Cancelling releases the pool slot immediately instead of waiting out the
292    /// job timeout (15 minutes by default), and reports no outcome - the agent is
293    /// already terminal, so there is nothing to apply.
294    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
295    async fn a_cancelled_job_frees_its_pool_slot_without_reporting() {
296        let mut cfg = InferencePoolConfig::new();
297        cfg.set_limit("m", 1);
298        let pools = InferencePools::new(cfg);
299        let permit = pools.try_acquire("m").expect("free pool");
300        assert!(pools.try_acquire("m").is_none(), "pool should be full");
301
302        // A provider that never answers - the stalled-call case a cancel exists
303        // to escape.
304        let provider = Arc::new(Scripted {
305            steps: std::sync::Mutex::new(vec![Step::Hang].into()),
306            calls: std::sync::Mutex::new(0),
307        });
308        let job = InferenceJob {
309            entity: Entity::from_raw_u32(7)
310                .expect("a small literal index is always a valid entity id"),
311            provider,
312            request: test_request(),
313            permit,
314            exact_token_counting: false,
315        };
316        let (tx, mut rx) = mpsc::unbounded_channel();
317        let cancel = crate::cancel::CancelToken::new();
318        let running = tokio::spawn(run_inference_job(
319            job,
320            tx,
321            Arc::new(Notify::new()),
322            // A job timeout far longer than the test: only the cancel can end it.
323            RetryPolicy {
324                max_attempts: 1,
325                base_delay: Duration::ZERO,
326                job_timeout: Duration::from_secs(3600),
327            },
328            cancel.clone(),
329        ));
330        tokio::task::yield_now().await;
331        cancel.cancel();
332
333        tokio::time::timeout(Duration::from_secs(5), running)
334            .await
335            .expect("the cancel ended the job")
336            .unwrap();
337        assert!(
338            pools.try_acquire("m").is_some(),
339            "the pool slot is free for the next agent"
340        );
341        assert!(
342            rx.try_recv().is_err(),
343            "and no outcome is reported for a cancelled run"
344        );
345    }
346
347    #[tokio::test]
348    async fn run_job_aborts_a_hung_call_and_frees_the_pool_slot() {
349        // A model pool of one slot, taken by the (hung) job under test.
350        let mut cfg = InferencePoolConfig::new();
351        cfg.set_limit("m", 1);
352        let pools = InferencePools::new(cfg);
353        let permit = pools.try_acquire("m").expect("free pool");
354        assert!(pools.try_acquire("m").is_none(), "pool should be full");
355
356        let provider = Arc::new(Scripted {
357            steps: std::sync::Mutex::new(vec![Step::Hang].into()),
358            calls: std::sync::Mutex::new(0),
359        });
360        let job = InferenceJob {
361            entity: Entity::from_raw_u32(7)
362                .expect("a small literal index is always a valid entity id"),
363            provider,
364            request: test_request(),
365            permit,
366            exact_token_counting: false,
367        };
368        let (tx, mut rx) = mpsc::unbounded_channel();
369        let policy = RetryPolicy {
370            max_attempts: 1,
371            base_delay: Duration::ZERO,
372            job_timeout: Duration::from_millis(50),
373        };
374        run_inference_job(
375            job,
376            tx,
377            Arc::new(Notify::new()),
378            policy,
379            crate::cancel::CancelToken::new(),
380        )
381        .await;
382
383        // The hung call was aborted with a timeout error…
384        let outcome = rx.try_recv().expect("outcome sent");
385        let err = outcome.result.expect_err("hung call should error");
386        assert!(err.to_string().contains("job timeout"), "got: {err}");
387        // …and its pool slot is free again for the next agent.
388        assert!(
389            pools.try_acquire("m").is_some(),
390            "the slot must be released after the timeout"
391        );
392    }
393
394    #[tokio::test]
395    async fn run_job_reports_ok_and_wakes() {
396        let (tx, mut rx) = mpsc::unbounded_channel();
397        let wake = Arc::new(Notify::new());
398        run_inference_job(
399            job(Arc::new(Fixed::Ok(response("hi")))),
400            tx,
401            wake.clone(),
402            RetryPolicy::default(),
403            crate::cancel::CancelToken::new(),
404        )
405        .await;
406
407        let outcome = rx.try_recv().expect("outcome sent");
408        assert_eq!(
409            outcome.entity,
410            Entity::from_raw_u32(7).expect("a small literal index is always a valid entity id")
411        );
412        assert_eq!(outcome.result.unwrap().content, "hi");
413        // The wake was signalled (a subsequent notified() returns immediately).
414        wake.notified().await;
415    }
416
417    #[tokio::test]
418    async fn run_job_reports_provider_error() {
419        let (tx, mut rx) = mpsc::unbounded_channel();
420        let wake = Arc::new(Notify::new());
421        let err = Arc::new(Fixed::Err("boom".to_string()));
422        run_inference_job(
423            job(err),
424            tx,
425            wake,
426            RetryPolicy::default(),
427            crate::cancel::CancelToken::new(),
428        )
429        .await;
430
431        let outcome = rx.try_recv().expect("outcome sent");
432        assert!(outcome.result.is_err());
433    }
434
435    /// A provider with a fixed `count_tokens` result and context window, used to
436    /// drive the opt-in pre-inference budget guard. `infer` always succeeds.
437    struct Counter {
438        count: usize,
439        max: usize,
440    }
441
442    #[async_trait::async_trait]
443    impl Provider for Counter {
444        async fn infer(
445            &self,
446            _req: InferenceRequest,
447        ) -> leviath_providers::Result<InferenceResponse> {
448            Ok(response("ok"))
449        }
450        async fn count_tokens(&self, _text: &str, _model: &str) -> usize {
451            self.count
452        }
453        fn max_context_tokens(&self, _model: &str) -> usize {
454            self.max
455        }
456        fn name(&self) -> &str {
457            "counter"
458        }
459        fn capabilities(&self, _model: &str) -> leviath_providers::ModelCapabilities {
460            leviath_providers::ModelCapabilities::default()
461        }
462    }
463
464    fn counting_job(provider: Arc<dyn Provider>, exact: bool) -> InferenceJob {
465        let pools = InferencePools::new(InferencePoolConfig::new());
466        InferenceJob {
467            entity: Entity::from_raw_u32(7)
468                .expect("a small literal index is always a valid entity id"),
469            provider,
470            request: test_request(), // max_tokens: 100
471            permit: pools.try_acquire("m").expect("free pool"),
472            exact_token_counting: exact,
473        }
474    }
475
476    #[test]
477    fn flatten_request_text_includes_system_messages_and_tools() {
478        use leviath_providers::{SystemBlock, Tool};
479        let req = InferenceRequest {
480            system: vec![SystemBlock {
481                text: "sys".to_string(),
482                cache_hint: leviath_core::CacheHint::Never,
483            }],
484            messages: vec![leviath_providers::Message {
485                role: "user".to_string(),
486                content: "hello".into(),
487                cache_breakpoint: false,
488            }],
489            model: "m".to_string(),
490            max_tokens: 10,
491            temperature: 0.0,
492            tools: vec![Tool {
493                name: "search".to_string(),
494                description: "find things".to_string(),
495                parameters: serde_json::json!({"type": "object"}),
496            }],
497            extra: serde_json::Value::Null,
498            request_timeout_secs: None,
499        };
500        let text = flatten_request_text(&req);
501        assert!(text.contains("sys"));
502        assert!(text.contains("hello"));
503        assert!(text.contains("search"));
504        assert!(text.contains("find things"));
505        assert!(text.contains("object"));
506    }
507
508    #[tokio::test]
509    async fn guard_rejects_request_over_context_window() {
510        // count(950) + max_tokens(100) = 1050 > context(1000) ⇒ rejected pre-flight.
511        let (tx, mut rx) = mpsc::unbounded_channel();
512        let provider = Arc::new(Counter {
513            count: 950,
514            max: 1000,
515        });
516        run_inference_job(
517            counting_job(provider, true),
518            tx,
519            Arc::new(Notify::new()),
520            RetryPolicy::default(),
521            crate::cancel::CancelToken::new(),
522        )
523        .await;
524        let outcome = rx.try_recv().expect("outcome sent");
525        let err = outcome.result.expect_err("should be rejected");
526        // Assert on the Display string (branch-free) rather than `matches!`,
527        // whose non-matching arm would be an uncovered region.
528        assert_eq!(err.to_string(), "Token limit exceeded: 950 > 1000");
529    }
530
531    #[tokio::test]
532    async fn guard_allows_request_within_context_window() {
533        // count(800) + 100 = 900 ≤ 1000 ⇒ proceeds to infer.
534        let (tx, mut rx) = mpsc::unbounded_channel();
535        let provider = Arc::new(Counter {
536            count: 800,
537            max: 1000,
538        });
539        run_inference_job(
540            counting_job(provider, true),
541            tx,
542            Arc::new(Notify::new()),
543            RetryPolicy::default(),
544            crate::cancel::CancelToken::new(),
545        )
546        .await;
547        let outcome = rx.try_recv().expect("outcome sent");
548        assert_eq!(outcome.result.expect("should succeed").content, "ok");
549    }
550
551    #[tokio::test]
552    async fn counter_provider_metadata_is_exercised() {
553        // Keep the Counter mock's non-`infer` trait methods measured.
554        let p = Counter { count: 5, max: 10 };
555        assert_eq!(p.name(), "counter");
556        assert_eq!(p.max_context_tokens("m"), 10);
557        assert_eq!(p.count_tokens("t", "m").await, 5);
558        assert!(p.capabilities("m").supports_streaming);
559    }
560
561    #[tokio::test]
562    async fn guard_off_skips_the_count_and_proceeds() {
563        // Even wildly over budget, with the flag off the guard never runs.
564        let (tx, mut rx) = mpsc::unbounded_channel();
565        let provider = Arc::new(Counter {
566            count: 1_000_000,
567            max: 1000,
568        });
569        run_inference_job(
570            counting_job(provider, false),
571            tx,
572            Arc::new(Notify::new()),
573            RetryPolicy::default(),
574            crate::cancel::CancelToken::new(),
575        )
576        .await;
577        let outcome = rx.try_recv().expect("outcome sent");
578        assert_eq!(outcome.result.expect("should succeed").content, "ok");
579    }
580
581    #[tokio::test]
582    async fn fixed_provider_metadata_is_exercised() {
583        // Covers the mock's non-`infer` trait methods (the pipeline resolves
584        // these off the provider elsewhere; here we just keep them measured).
585        let p = Fixed::Ok(response("x"));
586        assert_eq!(p.name(), "fixed");
587        assert_eq!(p.count_tokens("t", "m").await, 1);
588        assert_eq!(p.max_context_tokens("m"), 100_000);
589        let _ = p.capabilities("m");
590    }
591
592    #[tokio::test]
593    async fn run_job_survives_dropped_receiver() {
594        let (tx, rx) = mpsc::unbounded_channel();
595        drop(rx); // world shutting down: nobody to receive
596        let wake = Arc::new(Notify::new());
597        // Must not panic even though the send fails.
598        run_inference_job(
599            job(Arc::new(Fixed::Ok(response("x")))),
600            tx,
601            wake,
602            RetryPolicy::default(),
603            crate::cancel::CancelToken::new(),
604        )
605        .await;
606    }
607
608    // ── retry behavior ──
609
610    enum Step {
611        Ok(String),
612        Transient,
613        Permanent,
614        /// Never returns - a stalled/hung call, for the job-timeout test.
615        Hang,
616    }
617
618    /// A provider that plays a scripted sequence of results and counts calls.
619    struct Scripted {
620        steps: std::sync::Mutex<std::collections::VecDeque<Step>>,
621        calls: std::sync::Mutex<u32>,
622    }
623
624    #[async_trait::async_trait]
625    impl Provider for Scripted {
626        async fn infer(
627            &self,
628            _req: InferenceRequest,
629        ) -> leviath_providers::Result<InferenceResponse> {
630            *self.calls.lock().unwrap() += 1;
631            // Pop before matching so the mutex guard is not held across the
632            // `Hang` arm's `.await` (which would make this future non-`Send`).
633            let step = self.steps.lock().unwrap().pop_front();
634            match step {
635                Some(Step::Ok(t)) => Ok(response(&t)),
636                Some(Step::Transient) => Err(ProviderError::RateLimitExceeded),
637                Some(Step::Permanent) => Err(ProviderError::Other("permanent".to_string())),
638                Some(Step::Hang) => std::future::pending().await,
639                None => Err(ProviderError::Other("exhausted".to_string())),
640            }
641        }
642        async fn count_tokens(&self, _t: &str, _m: &str) -> usize {
643            1
644        }
645        fn max_context_tokens(&self, _m: &str) -> usize {
646            100_000
647        }
648        fn name(&self) -> &str {
649            "scripted"
650        }
651        fn capabilities(&self, _m: &str) -> leviath_providers::ModelCapabilities {
652            leviath_providers::ModelCapabilities::default()
653        }
654    }
655
656    fn no_delay(max_attempts: u32) -> RetryPolicy {
657        RetryPolicy {
658            max_attempts,
659            base_delay: Duration::ZERO,
660            job_timeout: Duration::from_secs(30),
661        }
662    }
663
664    #[tokio::test]
665    async fn run_job_retries_transient_then_succeeds() {
666        let provider = Arc::new(Scripted {
667            steps: std::sync::Mutex::new(
668                vec![
669                    Step::Transient,
670                    Step::Transient,
671                    Step::Ok("done".to_string()),
672                ]
673                .into(),
674            ),
675            calls: std::sync::Mutex::new(0),
676        });
677        let (tx, mut rx) = mpsc::unbounded_channel();
678        run_inference_job(
679            job(provider.clone()),
680            tx,
681            Arc::new(Notify::new()),
682            no_delay(4),
683            crate::cancel::CancelToken::new(),
684        )
685        .await;
686        let outcome = rx.try_recv().expect("outcome sent");
687        assert_eq!(outcome.result.unwrap().content, "done");
688        assert_eq!(*provider.calls.lock().unwrap(), 3); // two retries then success
689    }
690
691    #[tokio::test]
692    async fn run_job_gives_up_after_max_attempts() {
693        let provider = Arc::new(Scripted {
694            steps: std::sync::Mutex::new(
695                vec![
696                    Step::Transient,
697                    Step::Transient,
698                    Step::Transient,
699                    Step::Transient,
700                ]
701                .into(),
702            ),
703            calls: std::sync::Mutex::new(0),
704        });
705        let (tx, mut rx) = mpsc::unbounded_channel();
706        run_inference_job(
707            job(provider.clone()),
708            tx,
709            Arc::new(Notify::new()),
710            no_delay(3),
711            crate::cancel::CancelToken::new(),
712        )
713        .await;
714        let outcome = rx.try_recv().expect("outcome sent");
715        assert!(outcome.result.is_err());
716        assert_eq!(*provider.calls.lock().unwrap(), 3); // exhausted the 3 attempts
717    }
718
719    #[tokio::test]
720    async fn run_job_does_not_retry_a_permanent_error() {
721        let provider = Arc::new(Scripted {
722            steps: std::sync::Mutex::new(vec![Step::Permanent, Step::Ok("x".to_string())].into()),
723            calls: std::sync::Mutex::new(0),
724        });
725        let (tx, mut rx) = mpsc::unbounded_channel();
726        run_inference_job(
727            job(provider.clone()),
728            tx,
729            Arc::new(Notify::new()),
730            no_delay(4),
731            crate::cancel::CancelToken::new(),
732        )
733        .await;
734        let outcome = rx.try_recv().expect("outcome sent");
735        assert!(outcome.result.is_err());
736        assert_eq!(*provider.calls.lock().unwrap(), 1); // no retry on a permanent error
737    }
738
739    #[tokio::test]
740    async fn scripted_provider_metadata_is_exercised() {
741        let p = Scripted {
742            steps: std::sync::Mutex::new(std::collections::VecDeque::new()),
743            calls: std::sync::Mutex::new(0),
744        };
745        assert_eq!(p.name(), "scripted");
746        assert_eq!(p.count_tokens("t", "m").await, 1);
747        assert_eq!(p.max_context_tokens("m"), 100_000);
748        let _ = p.capabilities("m");
749    }
750}