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 {
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 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); 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 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 #[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 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 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 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 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 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 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 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(), 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 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_eq!(err.to_string(), "Token limit exceeded: 950 > 1000");
524 }
525
526 #[tokio::test]
527 async fn guard_allows_request_within_context_window() {
528 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 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 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 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); let wake = Arc::new(Notify::new());
592 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 enum Step {
606 Ok(String),
607 Transient,
608 Permanent,
609 Hang,
611 }
612
613 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 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); }
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); }
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); }
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}