1use 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#[derive(Debug, Clone, Copy)]
33pub struct RetryPolicy {
34 pub max_attempts: u32,
36 pub base_delay: Duration,
39 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 job_timeout: Duration::from_secs(leviath_providers::DEFAULT_INFERENCE_TIMEOUT_SECS),
65 }
66 }
67}
68
69pub struct InferenceJob {
71 pub entity: Entity,
73 pub provider: Arc<dyn Provider>,
75 pub request: InferenceRequest,
77 pub permit: InferencePermit,
80 pub exact_token_counting: bool,
86}
87
88fn 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
109pub struct InferenceOutcome {
112 pub entity: Entity,
114 pub result: Result<InferenceResponse, ProviderError>,
116 pub latency: std::time::Duration,
120}
121
122pub 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 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 let attempts = async {
171 let mut attempt = 1u32;
172 loop {
173 match provider.infer(&request).await {
174 Ok(response) => break Ok(response),
175 Err(e) if e.is_transient() && attempt < retry.max_attempts => {
176 tokio::time::sleep(retry.base_delay * 2u32.pow(attempt - 1)).await;
177 attempt += 1;
178 }
179 Err(e) => break Err(e),
180 }
181 }
182 };
183 let result = tokio::select! {
194 biased;
195 _ = cancel.cancelled() => {
196 drop(permit);
197 return;
198 }
199 outcome = tokio::time::timeout(retry.job_timeout, attempts) => match outcome {
200 Ok(result) => result,
201 Err(_elapsed) => Err(leviath_providers::ProviderError::Other(format!(
202 "inference exceeded the {}s job timeout and was aborted to free the \
203 pool slot (a stalled or never-completing response)",
204 retry.job_timeout.as_secs()
205 ))),
206 },
207 };
208 drop(permit); let _ = results.send(InferenceOutcome {
210 entity,
211 result,
212 latency: started.elapsed(),
213 });
214 wake.notify_one();
215}
216
217#[cfg(test)]
218mod tests {
219 use super::*;
220 use crate::inference_pool::{InferencePoolConfig, InferencePools};
221 use tokio::sync::mpsc;
222
223 fn test_request() -> InferenceRequest {
224 InferenceRequest {
225 system: vec![],
226 messages: vec![],
227 model: "m".to_string(),
228 max_tokens: 100,
229 temperature: 0.0,
230 tools: vec![],
231 extra: serde_json::Value::Null,
232 request_timeout_secs: None,
233 }
234 }
235
236 fn response(text: &str) -> InferenceResponse {
237 InferenceResponse {
238 content: text.to_string(),
239 tool_calls: vec![],
240 tokens_used: leviath_providers::TokenUsage {
241 prompt_tokens: 1,
242 completion_tokens: 1,
243 total_tokens: 2,
244 cached_tokens: 0,
245 cache_write_tokens: 0,
246 },
247 finish_reason: leviath_providers::FinishReason::Complete,
248 }
249 }
250
251 enum Fixed {
254 Ok(InferenceResponse),
255 Err(String),
256 }
257
258 #[async_trait::async_trait]
259 impl Provider for Fixed {
260 async fn infer(
261 &self,
262 _req: &InferenceRequest,
263 ) -> leviath_providers::Result<InferenceResponse> {
264 match self {
265 Fixed::Ok(r) => Ok(r.clone()),
266 Fixed::Err(m) => Err(ProviderError::Other(m.clone())),
267 }
268 }
269 async fn count_tokens(&self, _text: &str, _model: &str) -> usize {
270 1
271 }
272 fn max_context_tokens(&self, _model: &str) -> usize {
273 100_000
274 }
275 fn name(&self) -> &str {
276 "fixed"
277 }
278 fn capabilities(&self, _model: &str) -> leviath_providers::ModelCapabilities {
279 leviath_providers::ModelCapabilities::default()
280 }
281 }
282
283 fn job(provider: Arc<dyn Provider>) -> InferenceJob {
284 let pools = InferencePools::new(InferencePoolConfig::new());
285 InferenceJob {
286 entity: Entity::from_raw_u32(7)
287 .expect("a small literal index is always a valid entity id"),
288 provider,
289 request: test_request(),
290 permit: pools.try_acquire("m").expect("free pool"),
291 exact_token_counting: false,
292 }
293 }
294
295 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
299 async fn a_cancelled_job_frees_its_pool_slot_without_reporting() {
300 let mut cfg = InferencePoolConfig::new();
301 cfg.set_limit("m", 1);
302 let pools = InferencePools::new(cfg);
303 let permit = pools.try_acquire("m").expect("free pool");
304 assert!(pools.try_acquire("m").is_none(), "pool should be full");
305
306 let provider = Arc::new(Scripted {
309 steps: std::sync::Mutex::new(vec![Step::Hang].into()),
310 calls: std::sync::Mutex::new(0),
311 });
312 let job = InferenceJob {
313 entity: Entity::from_raw_u32(7)
314 .expect("a small literal index is always a valid entity id"),
315 provider,
316 request: test_request(),
317 permit,
318 exact_token_counting: false,
319 };
320 let (tx, mut rx) = mpsc::unbounded_channel();
321 let cancel = crate::cancel::CancelToken::new();
322 let running = tokio::spawn(run_inference_job(
323 job,
324 tx,
325 Arc::new(Notify::new()),
326 RetryPolicy {
328 max_attempts: 1,
329 base_delay: Duration::ZERO,
330 job_timeout: Duration::from_secs(3600),
331 },
332 cancel.clone(),
333 ));
334 tokio::task::yield_now().await;
335 cancel.cancel();
336
337 tokio::time::timeout(Duration::from_secs(5), running)
338 .await
339 .expect("the cancel ended the job")
340 .unwrap();
341 assert!(
342 pools.try_acquire("m").is_some(),
343 "the pool slot is free for the next agent"
344 );
345 assert!(
346 rx.try_recv().is_err(),
347 "and no outcome is reported for a cancelled run"
348 );
349 }
350
351 #[tokio::test]
352 async fn run_job_aborts_a_hung_call_and_frees_the_pool_slot() {
353 let mut cfg = InferencePoolConfig::new();
355 cfg.set_limit("m", 1);
356 let pools = InferencePools::new(cfg);
357 let permit = pools.try_acquire("m").expect("free pool");
358 assert!(pools.try_acquire("m").is_none(), "pool should be full");
359
360 let provider = Arc::new(Scripted {
361 steps: std::sync::Mutex::new(vec![Step::Hang].into()),
362 calls: std::sync::Mutex::new(0),
363 });
364 let job = InferenceJob {
365 entity: Entity::from_raw_u32(7)
366 .expect("a small literal index is always a valid entity id"),
367 provider,
368 request: test_request(),
369 permit,
370 exact_token_counting: false,
371 };
372 let (tx, mut rx) = mpsc::unbounded_channel();
373 let policy = RetryPolicy {
374 max_attempts: 1,
375 base_delay: Duration::ZERO,
376 job_timeout: Duration::from_millis(50),
377 };
378 run_inference_job(
379 job,
380 tx,
381 Arc::new(Notify::new()),
382 policy,
383 crate::cancel::CancelToken::new(),
384 )
385 .await;
386
387 let outcome = rx.try_recv().expect("outcome sent");
389 let err = outcome.result.expect_err("hung call should error");
390 assert!(err.to_string().contains("job timeout"), "got: {err}");
391 assert!(
393 pools.try_acquire("m").is_some(),
394 "the slot must be released after the timeout"
395 );
396 }
397
398 #[tokio::test]
399 async fn run_job_reports_ok_and_wakes() {
400 let (tx, mut rx) = mpsc::unbounded_channel();
401 let wake = Arc::new(Notify::new());
402 run_inference_job(
403 job(Arc::new(Fixed::Ok(response("hi")))),
404 tx,
405 wake.clone(),
406 RetryPolicy::default(),
407 crate::cancel::CancelToken::new(),
408 )
409 .await;
410
411 let outcome = rx.try_recv().expect("outcome sent");
412 assert_eq!(
413 outcome.entity,
414 Entity::from_raw_u32(7).expect("a small literal index is always a valid entity id")
415 );
416 assert_eq!(outcome.result.unwrap().content, "hi");
417 wake.notified().await;
419 }
420
421 #[tokio::test]
422 async fn run_job_reports_provider_error() {
423 let (tx, mut rx) = mpsc::unbounded_channel();
424 let wake = Arc::new(Notify::new());
425 let err = Arc::new(Fixed::Err("boom".to_string()));
426 run_inference_job(
427 job(err),
428 tx,
429 wake,
430 RetryPolicy::default(),
431 crate::cancel::CancelToken::new(),
432 )
433 .await;
434
435 let outcome = rx.try_recv().expect("outcome sent");
436 assert!(outcome.result.is_err());
437 }
438
439 struct Counter {
442 count: usize,
443 max: usize,
444 }
445
446 #[async_trait::async_trait]
447 impl Provider for Counter {
448 async fn infer(
449 &self,
450 _req: &InferenceRequest,
451 ) -> leviath_providers::Result<InferenceResponse> {
452 Ok(response("ok"))
453 }
454 async fn count_tokens(&self, _text: &str, _model: &str) -> usize {
455 self.count
456 }
457 fn max_context_tokens(&self, _model: &str) -> usize {
458 self.max
459 }
460 fn name(&self) -> &str {
461 "counter"
462 }
463 fn capabilities(&self, _model: &str) -> leviath_providers::ModelCapabilities {
464 leviath_providers::ModelCapabilities::default()
465 }
466 }
467
468 fn counting_job(provider: Arc<dyn Provider>, exact: bool) -> InferenceJob {
469 let pools = InferencePools::new(InferencePoolConfig::new());
470 InferenceJob {
471 entity: Entity::from_raw_u32(7)
472 .expect("a small literal index is always a valid entity id"),
473 provider,
474 request: test_request(), permit: pools.try_acquire("m").expect("free pool"),
476 exact_token_counting: exact,
477 }
478 }
479
480 #[test]
481 fn flatten_request_text_includes_system_messages_and_tools() {
482 use leviath_providers::{SystemBlock, Tool};
483 let req = InferenceRequest {
484 system: vec![SystemBlock {
485 text: "sys".to_string(),
486 cache_hint: leviath_core::CacheHint::Never,
487 }],
488 messages: vec![leviath_providers::Message {
489 role: "user".to_string(),
490 content: "hello".into(),
491 cache_breakpoint: false,
492 }],
493 model: "m".to_string(),
494 max_tokens: 10,
495 temperature: 0.0,
496 tools: vec![Tool {
497 name: "search".to_string(),
498 description: "find things".to_string(),
499 parameters: serde_json::json!({"type": "object"}),
500 }],
501 extra: serde_json::Value::Null,
502 request_timeout_secs: None,
503 };
504 let text = flatten_request_text(&req);
505 assert!(text.contains("sys"));
506 assert!(text.contains("hello"));
507 assert!(text.contains("search"));
508 assert!(text.contains("find things"));
509 assert!(text.contains("object"));
510 }
511
512 #[tokio::test]
513 async fn guard_rejects_request_over_context_window() {
514 let (tx, mut rx) = mpsc::unbounded_channel();
516 let provider = Arc::new(Counter {
517 count: 950,
518 max: 1000,
519 });
520 run_inference_job(
521 counting_job(provider, true),
522 tx,
523 Arc::new(Notify::new()),
524 RetryPolicy::default(),
525 crate::cancel::CancelToken::new(),
526 )
527 .await;
528 let outcome = rx.try_recv().expect("outcome sent");
529 let err = outcome.result.expect_err("should be rejected");
530 assert_eq!(err.to_string(), "Token limit exceeded: 950 > 1000");
533 }
534
535 #[tokio::test]
536 async fn guard_allows_request_within_context_window() {
537 let (tx, mut rx) = mpsc::unbounded_channel();
539 let provider = Arc::new(Counter {
540 count: 800,
541 max: 1000,
542 });
543 run_inference_job(
544 counting_job(provider, true),
545 tx,
546 Arc::new(Notify::new()),
547 RetryPolicy::default(),
548 crate::cancel::CancelToken::new(),
549 )
550 .await;
551 let outcome = rx.try_recv().expect("outcome sent");
552 assert_eq!(outcome.result.expect("should succeed").content, "ok");
553 }
554
555 #[tokio::test]
556 async fn counter_provider_metadata_is_exercised() {
557 let p = Counter { count: 5, max: 10 };
559 assert_eq!(p.name(), "counter");
560 assert_eq!(p.max_context_tokens("m"), 10);
561 assert_eq!(p.count_tokens("t", "m").await, 5);
562 assert!(p.capabilities("m").supports_streaming);
563 }
564
565 #[tokio::test]
566 async fn guard_off_skips_the_count_and_proceeds() {
567 let (tx, mut rx) = mpsc::unbounded_channel();
569 let provider = Arc::new(Counter {
570 count: 1_000_000,
571 max: 1000,
572 });
573 run_inference_job(
574 counting_job(provider, false),
575 tx,
576 Arc::new(Notify::new()),
577 RetryPolicy::default(),
578 crate::cancel::CancelToken::new(),
579 )
580 .await;
581 let outcome = rx.try_recv().expect("outcome sent");
582 assert_eq!(outcome.result.expect("should succeed").content, "ok");
583 }
584
585 #[tokio::test]
586 async fn fixed_provider_metadata_is_exercised() {
587 let p = Fixed::Ok(response("x"));
590 assert_eq!(p.name(), "fixed");
591 assert_eq!(p.count_tokens("t", "m").await, 1);
592 assert_eq!(p.max_context_tokens("m"), 100_000);
593 let _ = p.capabilities("m");
594 }
595
596 #[tokio::test]
597 async fn run_job_survives_dropped_receiver() {
598 let (tx, rx) = mpsc::unbounded_channel();
599 drop(rx); let wake = Arc::new(Notify::new());
601 run_inference_job(
603 job(Arc::new(Fixed::Ok(response("x")))),
604 tx,
605 wake,
606 RetryPolicy::default(),
607 crate::cancel::CancelToken::new(),
608 )
609 .await;
610 }
611
612 enum Step {
615 Ok(String),
616 Transient,
617 Permanent,
618 Hang,
620 }
621
622 struct Scripted {
624 steps: std::sync::Mutex<std::collections::VecDeque<Step>>,
625 calls: std::sync::Mutex<u32>,
626 }
627
628 #[async_trait::async_trait]
629 impl Provider for Scripted {
630 async fn infer(
631 &self,
632 _req: &InferenceRequest,
633 ) -> leviath_providers::Result<InferenceResponse> {
634 *self.calls.lock().unwrap() += 1;
635 let step = self.steps.lock().unwrap().pop_front();
638 match step {
639 Some(Step::Ok(t)) => Ok(response(&t)),
640 Some(Step::Transient) => Err(ProviderError::RateLimitExceeded),
641 Some(Step::Permanent) => Err(ProviderError::Other("permanent".to_string())),
642 Some(Step::Hang) => std::future::pending().await,
643 None => Err(ProviderError::Other("exhausted".to_string())),
644 }
645 }
646 async fn count_tokens(&self, _t: &str, _m: &str) -> usize {
647 1
648 }
649 fn max_context_tokens(&self, _m: &str) -> usize {
650 100_000
651 }
652 fn name(&self) -> &str {
653 "scripted"
654 }
655 fn capabilities(&self, _m: &str) -> leviath_providers::ModelCapabilities {
656 leviath_providers::ModelCapabilities::default()
657 }
658 }
659
660 fn no_delay(max_attempts: u32) -> RetryPolicy {
661 RetryPolicy {
662 max_attempts,
663 base_delay: Duration::ZERO,
664 job_timeout: Duration::from_secs(30),
665 }
666 }
667
668 #[tokio::test]
669 async fn run_job_retries_transient_then_succeeds() {
670 let provider = Arc::new(Scripted {
671 steps: std::sync::Mutex::new(
672 vec![
673 Step::Transient,
674 Step::Transient,
675 Step::Ok("done".to_string()),
676 ]
677 .into(),
678 ),
679 calls: std::sync::Mutex::new(0),
680 });
681 let (tx, mut rx) = mpsc::unbounded_channel();
682 run_inference_job(
683 job(provider.clone()),
684 tx,
685 Arc::new(Notify::new()),
686 no_delay(4),
687 crate::cancel::CancelToken::new(),
688 )
689 .await;
690 let outcome = rx.try_recv().expect("outcome sent");
691 assert_eq!(outcome.result.unwrap().content, "done");
692 assert_eq!(*provider.calls.lock().unwrap(), 3); }
694
695 #[tokio::test]
696 async fn run_job_gives_up_after_max_attempts() {
697 let provider = Arc::new(Scripted {
698 steps: std::sync::Mutex::new(
699 vec![
700 Step::Transient,
701 Step::Transient,
702 Step::Transient,
703 Step::Transient,
704 ]
705 .into(),
706 ),
707 calls: std::sync::Mutex::new(0),
708 });
709 let (tx, mut rx) = mpsc::unbounded_channel();
710 run_inference_job(
711 job(provider.clone()),
712 tx,
713 Arc::new(Notify::new()),
714 no_delay(3),
715 crate::cancel::CancelToken::new(),
716 )
717 .await;
718 let outcome = rx.try_recv().expect("outcome sent");
719 assert!(outcome.result.is_err());
720 assert_eq!(*provider.calls.lock().unwrap(), 3); }
722
723 #[tokio::test]
724 async fn run_job_does_not_retry_a_permanent_error() {
725 let provider = Arc::new(Scripted {
726 steps: std::sync::Mutex::new(vec![Step::Permanent, Step::Ok("x".to_string())].into()),
727 calls: std::sync::Mutex::new(0),
728 });
729 let (tx, mut rx) = mpsc::unbounded_channel();
730 run_inference_job(
731 job(provider.clone()),
732 tx,
733 Arc::new(Notify::new()),
734 no_delay(4),
735 crate::cancel::CancelToken::new(),
736 )
737 .await;
738 let outcome = rx.try_recv().expect("outcome sent");
739 assert!(outcome.result.is_err());
740 assert_eq!(*provider.calls.lock().unwrap(), 1); }
742
743 #[tokio::test]
744 async fn scripted_provider_metadata_is_exercised() {
745 let p = Scripted {
746 steps: std::sync::Mutex::new(std::collections::VecDeque::new()),
747 calls: std::sync::Mutex::new(0),
748 };
749 assert_eq!(p.name(), "scripted");
750 assert_eq!(p.count_tokens("t", "m").await, 1);
751 assert_eq!(p.max_context_tokens("m"), 100_000);
752 let _ = p.capabilities("m");
753 }
754}