Skip to main content

litellm_rs/core/integrations/
callback_runtime.rs

1//! Non-blocking ordered runtime for integration lifecycle callbacks.
2
3use std::sync::Arc;
4
5use tokio::sync::{mpsc, oneshot};
6use tokio::task::JoinHandle;
7use tracing::{error, warn};
8
9use super::IntegrationManager;
10use crate::core::traits::integration::{
11    EmbeddingEndEvent, EmbeddingErrorEvent, EmbeddingStartEvent, GuardrailDecisionEvent,
12    IntegrationError, IntegrationResult, LlmEndEvent, LlmErrorEvent, LlmStartEvent, LlmStreamEvent,
13};
14
15pub(crate) trait CallbackMetrics: Send + Sync {
16    #[cfg(any(feature = "gateway", test))]
17    fn begin_llm_lifecycle(&self, event: &LlmStartEvent);
18    #[cfg(any(feature = "gateway", test))]
19    fn finish_llm_lifecycle(&self, event: &LlmEndEvent);
20    #[cfg(feature = "gateway")]
21    fn fail_llm_lifecycle(&self, event: &LlmErrorEvent);
22    #[cfg(any(feature = "gateway", test))]
23    fn cancel_llm_lifecycle(&self, event: &LlmStartEvent);
24    #[cfg(feature = "gateway")]
25    fn record_embedding_start(&self, event: &EmbeddingStartEvent);
26    #[cfg(feature = "gateway")]
27    fn record_embedding_end(&self, event: &EmbeddingEndEvent);
28    fn render(&self) -> String;
29}
30
31pub(crate) type CallbackMetricsRecorder = Arc<dyn CallbackMetrics>;
32
33#[derive(Debug)]
34enum CallbackEvent {
35    Start(LlmStartEvent),
36    End(LlmEndEvent),
37    Error(LlmErrorEvent),
38    Stream(LlmStreamEvent),
39    EmbeddingStart(EmbeddingStartEvent),
40    EmbeddingEnd(EmbeddingEndEvent),
41    EmbeddingError(EmbeddingErrorEvent),
42    GuardrailDecision(GuardrailDecisionEvent),
43}
44
45/// Error returned when an event cannot enter the non-blocking callback queue.
46#[derive(Debug, thiserror::Error, PartialEq, Eq)]
47pub enum CallbackDispatchError {
48    /// The bounded queue is at capacity.
49    #[error("callback queue is full")]
50    QueueFull,
51    /// The callback worker is no longer accepting events.
52    #[error("callback queue is closed")]
53    QueueClosed,
54}
55
56/// Cloneable request-path handle for non-blocking callback delivery.
57#[derive(Clone, Default)]
58pub struct CallbackDispatcher {
59    sender: Option<mpsc::Sender<CallbackEvent>>,
60    manager: Option<Arc<IntegrationManager>>,
61    callback_metrics: Option<CallbackMetricsRecorder>,
62}
63
64/// Capacity reserved for the terminal event of an admitted lifecycle.
65pub struct CallbackTerminalPermit {
66    permit: Option<mpsc::OwnedPermit<CallbackEvent>>,
67}
68
69impl CallbackTerminalPermit {
70    /// Enqueue an LLM success using its reserved capacity.
71    pub fn emit_end(self, event: LlmEndEvent) {
72        if let Some(permit) = self.permit {
73            permit.send(CallbackEvent::End(event));
74        }
75    }
76
77    /// Enqueue an LLM error using its reserved capacity.
78    pub fn emit_error(self, event: LlmErrorEvent) {
79        if let Some(permit) = self.permit {
80            permit.send(CallbackEvent::Error(event));
81        }
82    }
83
84    /// Enqueue an embedding success using its reserved capacity.
85    pub fn emit_embedding_end(self, event: EmbeddingEndEvent) {
86        if let Some(permit) = self.permit {
87            permit.send(CallbackEvent::EmbeddingEnd(event));
88        }
89    }
90
91    /// Enqueue an embedding error using its reserved capacity.
92    pub fn emit_embedding_error(self, event: EmbeddingErrorEvent) {
93        if let Some(permit) = self.permit {
94            permit.send(CallbackEvent::EmbeddingError(event));
95        }
96    }
97}
98
99#[cfg(any(feature = "gateway", test))]
100enum CallbackMetricsKind {
101    Llm(Box<LlmStartEvent>),
102    #[cfg(feature = "gateway")]
103    Embedding,
104}
105
106/// Request lifecycle guard for synchronous in-process callback metrics.
107///
108/// This is deliberately separate from [`CallbackTerminalPermit`]: exporter
109/// queue admission keeps its public failure semantics, while metrics never
110/// enter that queue. Dropping an unfinished LLM guard releases the active
111/// request gauge during cancellation, task abortion, or panic unwinding.
112#[cfg(any(feature = "gateway", test))]
113pub(crate) struct CallbackMetricsPermit {
114    recorder: CallbackMetricsRecorder,
115    kind: CallbackMetricsKind,
116    completed: bool,
117}
118
119#[cfg(any(feature = "gateway", test))]
120impl CallbackMetricsPermit {
121    #[cfg(feature = "gateway")]
122    pub(crate) fn update_llm_target(&mut self, model: &str, provider: &str) {
123        if let CallbackMetricsKind::Llm(event) = &mut self.kind {
124            event.model = model.to_string();
125            event.provider = Some(provider.to_string());
126        }
127    }
128
129    pub(crate) fn emit_end(mut self, event: &LlmEndEvent) {
130        self.recorder.finish_llm_lifecycle(event);
131        self.completed = true;
132    }
133
134    #[cfg(feature = "gateway")]
135    pub(crate) fn emit_error(mut self, event: &LlmErrorEvent) {
136        if matches!(self.kind, CallbackMetricsKind::Llm(_)) {
137            self.recorder.fail_llm_lifecycle(event);
138        }
139        self.completed = true;
140    }
141
142    #[cfg(feature = "gateway")]
143    pub(crate) fn emit_embedding_end(mut self, event: &EmbeddingEndEvent) {
144        self.recorder.record_embedding_end(event);
145        self.completed = true;
146    }
147
148    #[cfg(feature = "gateway")]
149    pub(crate) fn emit_embedding_error(mut self) {
150        self.completed = true;
151    }
152}
153
154#[cfg(any(feature = "gateway", test))]
155impl Drop for CallbackMetricsPermit {
156    fn drop(&mut self) {
157        if !self.completed
158            && let CallbackMetricsKind::Llm(event) = &self.kind
159        {
160            self.recorder.cancel_llm_lifecycle(event);
161        }
162    }
163}
164
165impl CallbackDispatcher {
166    /// Create a dispatcher that performs no external callback work.
167    pub fn disabled() -> Self {
168        Self::default()
169    }
170
171    /// Whether this dispatcher has exporter delivery or in-process metrics enabled.
172    pub fn is_enabled(&self) -> bool {
173        self.sender.is_some() || self.callback_metrics.is_some()
174    }
175
176    /// List the callback integrations registered with this dispatcher.
177    pub async fn registered_integrations(&self) -> Vec<&'static str> {
178        let mut integrations = match &self.manager {
179            Some(manager) => manager.list_integrations().await,
180            None => Vec::new(),
181        };
182        if self.callback_metrics.is_some() {
183            integrations.push("prometheus");
184        }
185        integrations
186    }
187
188    /// Render metrics from the configured Prometheus callback backend.
189    pub fn render_prometheus_metrics(&self) -> Option<String> {
190        self.callback_metrics
191            .as_ref()
192            .map(|metrics| metrics.render())
193    }
194
195    /// Atomically admit an LLM start and reserve capacity for its terminal event.
196    pub fn begin_llm(
197        &self,
198        event: LlmStartEvent,
199    ) -> Result<CallbackTerminalPermit, CallbackDispatchError> {
200        self.begin(CallbackEvent::Start(event))
201    }
202
203    /// Atomically admit an embedding start and reserve its terminal capacity.
204    pub fn begin_embedding(
205        &self,
206        event: EmbeddingStartEvent,
207    ) -> Result<CallbackTerminalPermit, CallbackDispatchError> {
208        self.begin(CallbackEvent::EmbeddingStart(event))
209    }
210
211    /// Enqueue an LLM start event without waiting for exporter I/O.
212    pub fn emit_start(&self, event: LlmStartEvent) -> Result<(), CallbackDispatchError> {
213        self.try_send(CallbackEvent::Start(event))
214    }
215
216    /// Enqueue an LLM success event without waiting for exporter I/O.
217    pub fn emit_end(&self, event: LlmEndEvent) -> Result<(), CallbackDispatchError> {
218        self.try_send(CallbackEvent::End(event))
219    }
220
221    /// Enqueue an LLM error event without waiting for exporter I/O.
222    pub fn emit_error(&self, event: LlmErrorEvent) -> Result<(), CallbackDispatchError> {
223        self.try_send(CallbackEvent::Error(event))
224    }
225
226    /// Enqueue an LLM stream event without waiting for exporter I/O.
227    pub fn emit_stream(&self, event: LlmStreamEvent) -> Result<(), CallbackDispatchError> {
228        self.try_send(CallbackEvent::Stream(event))
229    }
230
231    /// Enqueue a sanitized guardrail decision. Queue pressure is observable only.
232    pub fn emit_guardrail_decision(&self, event: GuardrailDecisionEvent) {
233        if let Err(error) = self.try_send(CallbackEvent::GuardrailDecision(event)) {
234            warn!(
235                error = %error,
236                "Guardrail decision callback dispatch failed"
237            );
238        }
239    }
240
241    #[cfg(any(feature = "gateway", test))]
242    pub(crate) fn begin_llm_metrics(&self, event: &LlmStartEvent) -> Option<CallbackMetricsPermit> {
243        let recorder = self.callback_metrics.clone()?;
244        recorder.begin_llm_lifecycle(event);
245        Some(CallbackMetricsPermit {
246            recorder,
247            kind: CallbackMetricsKind::Llm(Box::new(event.clone())),
248            completed: false,
249        })
250    }
251
252    #[cfg(feature = "gateway")]
253    pub(crate) fn begin_embedding_metrics(
254        &self,
255        event: &EmbeddingStartEvent,
256    ) -> Option<CallbackMetricsPermit> {
257        let recorder = self.callback_metrics.clone()?;
258        recorder.record_embedding_start(event);
259        Some(CallbackMetricsPermit {
260            recorder,
261            kind: CallbackMetricsKind::Embedding,
262            completed: false,
263        })
264    }
265
266    fn try_send(&self, event: CallbackEvent) -> Result<(), CallbackDispatchError> {
267        let Some(sender) = &self.sender else {
268            return Ok(());
269        };
270        sender.try_send(event).map_err(|error| match error {
271            mpsc::error::TrySendError::Full(_) => CallbackDispatchError::QueueFull,
272            mpsc::error::TrySendError::Closed(_) => CallbackDispatchError::QueueClosed,
273        })
274    }
275
276    fn begin(
277        &self,
278        start_event: CallbackEvent,
279    ) -> Result<CallbackTerminalPermit, CallbackDispatchError> {
280        let Some(sender) = &self.sender else {
281            return Ok(CallbackTerminalPermit { permit: None });
282        };
283        let start_permit = sender
284            .clone()
285            .try_reserve_owned()
286            .map_err(map_reserve_error)?;
287        let terminal_permit = sender
288            .clone()
289            .try_reserve_owned()
290            .map_err(map_reserve_error)?;
291        start_permit.send(start_event);
292        Ok(CallbackTerminalPermit {
293            permit: Some(terminal_permit),
294        })
295    }
296}
297
298fn map_reserve_error<T>(error: mpsc::error::TrySendError<T>) -> CallbackDispatchError {
299    match error {
300        mpsc::error::TrySendError::Full(_) => CallbackDispatchError::QueueFull,
301        mpsc::error::TrySendError::Closed(_) => CallbackDispatchError::QueueClosed,
302    }
303}
304
305/// Owns the callback queue worker and its graceful-shutdown signal.
306pub struct CallbackRuntime {
307    dispatcher: CallbackDispatcher,
308    shutdown_tx: Option<oneshot::Sender<()>>,
309    worker: Option<JoinHandle<IntegrationResult<()>>>,
310}
311
312impl CallbackRuntime {
313    /// Start a callback worker around an existing integration manager.
314    pub fn new(manager: Arc<IntegrationManager>, queue_capacity: usize) -> IntegrationResult<Self> {
315        if queue_capacity < 2 {
316            return Err(IntegrationError::config(
317                "callback queue capacity must be at least 2",
318            ));
319        }
320
321        let (sender, receiver) = mpsc::channel(queue_capacity);
322        let (shutdown_tx, shutdown_rx) = oneshot::channel();
323        let worker_manager = Arc::clone(&manager);
324        let worker = tokio::spawn(async move {
325            run_callback_worker(worker_manager, receiver, shutdown_rx).await
326        });
327
328        Ok(Self {
329            dispatcher: CallbackDispatcher {
330                sender: Some(sender),
331                manager: Some(manager),
332                callback_metrics: None,
333            },
334            shutdown_tx: Some(shutdown_tx),
335            worker: Some(worker),
336        })
337    }
338
339    /// Construct a disabled runtime.
340    pub fn disabled() -> Self {
341        Self {
342            dispatcher: CallbackDispatcher::disabled(),
343            shutdown_tx: None,
344            worker: None,
345        }
346    }
347
348    /// Obtain the request-path dispatcher.
349    pub fn dispatcher(&self) -> CallbackDispatcher {
350        self.dispatcher.clone()
351    }
352
353    #[cfg(any(feature = "gateway", test))]
354    pub(crate) fn with_callback_metrics(
355        mut self,
356        metrics: Option<CallbackMetricsRecorder>,
357    ) -> Self {
358        self.dispatcher.callback_metrics = metrics;
359        self
360    }
361
362    /// Drain pending events, flush integrations, and stop the worker.
363    pub async fn shutdown(mut self) -> IntegrationResult<()> {
364        if let Some(shutdown_tx) = self.shutdown_tx.take() {
365            let _ = shutdown_tx.send(());
366        }
367        let Some(worker) = self.worker.take() else {
368            return Ok(());
369        };
370        worker
371            .await
372            .map_err(|error| IntegrationError::other(format!("callback worker failed: {error}")))?
373    }
374}
375
376impl Drop for CallbackRuntime {
377    fn drop(&mut self) {
378        if let Some(shutdown_tx) = self.shutdown_tx.take() {
379            let _ = shutdown_tx.send(());
380        }
381    }
382}
383
384async fn run_callback_worker(
385    manager: Arc<IntegrationManager>,
386    mut receiver: mpsc::Receiver<CallbackEvent>,
387    mut shutdown_rx: oneshot::Receiver<()>,
388) -> IntegrationResult<()> {
389    loop {
390        tokio::select! {
391            biased;
392            _ = &mut shutdown_rx => {
393                receiver.close();
394                while let Some(event) = receiver.recv().await {
395                    dispatch_event(manager.as_ref(), event).await;
396                }
397                break;
398            }
399            event = receiver.recv() => {
400                let Some(event) = event else {
401                    break;
402                };
403                dispatch_event(manager.as_ref(), event).await;
404            }
405        }
406    }
407
408    if let Err(error) = manager.flush().await {
409        warn!("Callback integration flush failed: {}", error);
410    }
411    manager.shutdown().await
412}
413
414async fn dispatch_event(manager: &IntegrationManager, event: CallbackEvent) {
415    let result = match event {
416        CallbackEvent::Start(event) => manager.on_llm_start(&event).await,
417        CallbackEvent::End(event) => manager.on_llm_end(&event).await,
418        CallbackEvent::Error(event) => manager.on_llm_error(&event).await,
419        CallbackEvent::Stream(event) => manager.on_llm_stream(&event).await,
420        CallbackEvent::EmbeddingStart(event) => manager.on_embedding_start(&event).await,
421        CallbackEvent::EmbeddingEnd(event) => manager.on_embedding_end(&event).await,
422        CallbackEvent::EmbeddingError(event) => manager.on_embedding_error(&event).await,
423        CallbackEvent::GuardrailDecision(event) => manager.on_guardrail_decision(&event).await,
424    };
425    if let Err(error) = result {
426        error!("Callback event dispatch failed: {}", error);
427    }
428}
429
430#[cfg(test)]
431mod tests {
432    use std::sync::atomic::{AtomicBool, Ordering};
433
434    use async_trait::async_trait;
435    use parking_lot::Mutex;
436    use tokio::sync::Notify;
437
438    use super::*;
439    use crate::core::integrations::{IntegrationManagerConfig, PrometheusIntegration};
440    use crate::core::traits::integration::{Integration, LlmErrorEvent};
441
442    struct RecordingIntegration {
443        events: Arc<Mutex<Vec<&'static str>>>,
444        flushed: Arc<AtomicBool>,
445    }
446
447    struct BlockingIntegration {
448        entered: Arc<Notify>,
449        release: Arc<Notify>,
450        events: Arc<Mutex<Vec<&'static str>>>,
451    }
452
453    struct FailingIntegration;
454
455    #[async_trait]
456    impl Integration for FailingIntegration {
457        fn name(&self) -> &'static str {
458            "failing"
459        }
460
461        fn is_enabled(&self) -> bool {
462            true
463        }
464
465        async fn on_llm_start(&self, _event: &LlmStartEvent) -> IntegrationResult<()> {
466            Err(IntegrationError::other("test callback failure"))
467        }
468
469        async fn on_llm_end(&self, _event: &LlmEndEvent) -> IntegrationResult<()> {
470            Err(IntegrationError::other("test callback failure"))
471        }
472
473        async fn on_llm_error(&self, _event: &LlmErrorEvent) -> IntegrationResult<()> {
474            Err(IntegrationError::other("test callback failure"))
475        }
476
477        async fn flush(&self) -> IntegrationResult<()> {
478            Ok(())
479        }
480
481        async fn shutdown(&self) -> IntegrationResult<()> {
482            Ok(())
483        }
484    }
485
486    #[async_trait]
487    impl Integration for BlockingIntegration {
488        fn name(&self) -> &'static str {
489            "blocking"
490        }
491
492        fn is_enabled(&self) -> bool {
493            true
494        }
495
496        async fn on_llm_start(&self, _event: &LlmStartEvent) -> IntegrationResult<()> {
497            self.events.lock().push("start");
498            self.entered.notify_one();
499            self.release.notified().await;
500            Ok(())
501        }
502
503        async fn on_llm_end(&self, _event: &LlmEndEvent) -> IntegrationResult<()> {
504            self.events.lock().push("end");
505            Ok(())
506        }
507
508        async fn on_llm_error(&self, _event: &LlmErrorEvent) -> IntegrationResult<()> {
509            Ok(())
510        }
511
512        async fn flush(&self) -> IntegrationResult<()> {
513            Ok(())
514        }
515
516        async fn shutdown(&self) -> IntegrationResult<()> {
517            Ok(())
518        }
519    }
520
521    #[async_trait]
522    impl Integration for RecordingIntegration {
523        fn name(&self) -> &'static str {
524            "recording"
525        }
526
527        fn is_enabled(&self) -> bool {
528            true
529        }
530
531        async fn on_llm_start(&self, _event: &LlmStartEvent) -> IntegrationResult<()> {
532            self.events.lock().push("start");
533            Ok(())
534        }
535
536        async fn on_llm_end(&self, _event: &LlmEndEvent) -> IntegrationResult<()> {
537            self.events.lock().push("end");
538            Ok(())
539        }
540
541        async fn on_llm_error(&self, _event: &LlmErrorEvent) -> IntegrationResult<()> {
542            self.events.lock().push("error");
543            Ok(())
544        }
545
546        async fn flush(&self) -> IntegrationResult<()> {
547            self.flushed.store(true, Ordering::SeqCst);
548            Ok(())
549        }
550
551        async fn shutdown(&self) -> IntegrationResult<()> {
552            Ok(())
553        }
554    }
555
556    #[tokio::test]
557    async fn runtime_preserves_order_and_flushes_on_shutdown() {
558        let events = Arc::new(Mutex::new(Vec::new()));
559        let flushed = Arc::new(AtomicBool::new(false));
560        let manager = Arc::new(IntegrationManager::with_defaults());
561        manager
562            .register(Arc::new(RecordingIntegration {
563                events: Arc::clone(&events),
564                flushed: Arc::clone(&flushed),
565            }))
566            .await;
567        let runtime = match CallbackRuntime::new(manager, 8) {
568            Ok(runtime) => runtime,
569            Err(error) => panic!("callback runtime should start: {error}"),
570        };
571        let dispatcher = runtime.dispatcher();
572
573        assert!(
574            dispatcher
575                .emit_start(LlmStartEvent::new("req-1", "model"))
576                .is_ok()
577        );
578        assert!(
579            dispatcher
580                .emit_end(LlmEndEvent::new("req-1", "model"))
581                .is_ok()
582        );
583        assert!(runtime.shutdown().await.is_ok());
584        assert_eq!(*events.lock(), vec!["start", "end"]);
585        assert!(flushed.load(Ordering::SeqCst));
586    }
587
588    #[tokio::test]
589    async fn runtime_rejects_capacity_below_lifecycle_pair() {
590        let manager = Arc::new(IntegrationManager::with_defaults());
591        for capacity in [0, 1] {
592            let error = match CallbackRuntime::new(Arc::clone(&manager), capacity) {
593                Ok(_) => panic!("callback runtime capacity {capacity} must fail"),
594                Err(error) => error,
595            };
596            assert!(error.to_string().contains("at least 2"));
597        }
598    }
599
600    #[tokio::test]
601    async fn dispatcher_reports_full_and_closed_queue() {
602        let entered = Arc::new(Notify::new());
603        let release = Arc::new(Notify::new());
604        let manager = Arc::new(IntegrationManager::new(
605            IntegrationManagerConfig::default().parallel(false),
606        ));
607        manager
608            .register(Arc::new(BlockingIntegration {
609                entered: Arc::clone(&entered),
610                release: Arc::clone(&release),
611                events: Arc::new(Mutex::new(Vec::new())),
612            }))
613            .await;
614        let runtime = match CallbackRuntime::new(manager, 2) {
615            Ok(runtime) => runtime,
616            Err(error) => panic!("callback runtime should start: {error}"),
617        };
618        let dispatcher = runtime.dispatcher();
619        let entered_wait = entered.notified();
620
621        assert!(
622            dispatcher
623                .emit_start(LlmStartEvent::new("req-1", "model"))
624                .is_ok()
625        );
626        entered_wait.await;
627        assert!(
628            dispatcher
629                .emit_end(LlmEndEvent::new("req-1", "model"))
630                .is_ok()
631        );
632        assert!(
633            dispatcher
634                .emit_end(LlmEndEvent::new("req-1", "model"))
635                .is_ok()
636        );
637        assert_eq!(
638            dispatcher.emit_error(LlmErrorEvent::new("req-1", "model", "failed")),
639            Err(CallbackDispatchError::QueueFull)
640        );
641
642        release.notify_one();
643        assert!(runtime.shutdown().await.is_ok());
644        assert_eq!(
645            dispatcher.emit_start(LlmStartEvent::new("req-2", "model")),
646            Err(CallbackDispatchError::QueueClosed)
647        );
648    }
649
650    #[tokio::test]
651    async fn lifecycle_admission_reserves_terminal_capacity_as_a_pair() {
652        let entered = Arc::new(Notify::new());
653        let release = Arc::new(Notify::new());
654        let events = Arc::new(Mutex::new(Vec::new()));
655        let manager = Arc::new(IntegrationManager::new(
656            IntegrationManagerConfig::default().parallel(false),
657        ));
658        manager
659            .register(Arc::new(BlockingIntegration {
660                entered: Arc::clone(&entered),
661                release: Arc::clone(&release),
662                events: Arc::clone(&events),
663            }))
664            .await;
665        let runtime = CallbackRuntime::new(manager, 2).unwrap();
666        let dispatcher = runtime.dispatcher();
667        let entered_wait = entered.notified();
668
669        let terminal = dispatcher
670            .begin_llm(LlmStartEvent::new("req-1", "model"))
671            .unwrap();
672        entered_wait.await;
673        assert!(matches!(
674            dispatcher.begin_llm(LlmStartEvent::new("req-2", "model")),
675            Err(CallbackDispatchError::QueueFull)
676        ));
677        terminal.emit_end(LlmEndEvent::new("req-1", "model"));
678        release.notify_one();
679        runtime.shutdown().await.unwrap();
680
681        assert_eq!(*events.lock(), vec!["start", "end"]);
682    }
683
684    #[tokio::test]
685    async fn metrics_lifecycle_is_independent_without_hiding_queue_errors() {
686        let entered = Arc::new(Notify::new());
687        let release = Arc::new(Notify::new());
688        let manager = Arc::new(IntegrationManager::new(
689            IntegrationManagerConfig::default().parallel(false),
690        ));
691        manager
692            .register(Arc::new(BlockingIntegration {
693                entered: Arc::clone(&entered),
694                release: Arc::clone(&release),
695                events: Arc::new(Mutex::new(Vec::new())),
696            }))
697            .await;
698        let prometheus = Arc::new(PrometheusIntegration::with_defaults());
699        let runtime = CallbackRuntime::new(manager, 2)
700            .unwrap()
701            .with_callback_metrics(Some(prometheus));
702        let dispatcher = runtime.dispatcher();
703        let entered_wait = entered.notified();
704
705        let first = dispatcher
706            .begin_llm(LlmStartEvent::new("req-1", "model"))
707            .unwrap();
708        entered_wait.await;
709
710        let second_start = LlmStartEvent::new("req-2", "model");
711        let second_metrics = dispatcher
712            .begin_llm_metrics(&second_start)
713            .expect("configured metrics should admit independently");
714        assert!(matches!(
715            dispatcher.begin_llm(second_start),
716            Err(CallbackDispatchError::QueueFull)
717        ));
718        second_metrics.emit_end(&LlmEndEvent::new("req-2", "model").latency(25));
719        assert_eq!(
720            dispatcher.emit_end(LlmEndEvent::new("queued", "model")),
721            Ok(())
722        );
723        for _ in 0..2 {
724            assert_eq!(
725                dispatcher.emit_end(LlmEndEvent::new("retry", "model")),
726                Err(CallbackDispatchError::QueueFull)
727            );
728        }
729
730        let cancelled = dispatcher
731            .begin_llm_metrics(&LlmStartEvent::new("req-3", "model"))
732            .expect("configured metrics should admit independently");
733        drop(cancelled);
734
735        let rendered = dispatcher.render_prometheus_metrics().unwrap();
736        assert!(rendered.contains("litellm_requests_total{model=\"model\"} 2"));
737        assert!(rendered.contains("litellm_requests_success_total{model=\"model\"} 1"));
738        assert!(rendered.contains("litellm_active_requests 0"));
739
740        first.emit_end(LlmEndEvent::new("req-1", "model"));
741        release.notify_one();
742        runtime.shutdown().await.unwrap();
743    }
744
745    #[tokio::test]
746    async fn failing_backend_does_not_block_healthy_backend() {
747        let events = Arc::new(Mutex::new(Vec::new()));
748        let manager = Arc::new(IntegrationManager::new(
749            IntegrationManagerConfig::default()
750                .parallel(false)
751                .fail_fast(false),
752        ));
753        manager.register(Arc::new(FailingIntegration)).await;
754        manager
755            .register(Arc::new(RecordingIntegration {
756                events: Arc::clone(&events),
757                flushed: Arc::new(AtomicBool::new(false)),
758            }))
759            .await;
760        let runtime = match CallbackRuntime::new(manager, 8) {
761            Ok(runtime) => runtime,
762            Err(error) => panic!("callback runtime should start: {error}"),
763        };
764        let dispatcher = runtime.dispatcher();
765
766        assert!(
767            dispatcher
768                .emit_start(LlmStartEvent::new("req-1", "model"))
769                .is_ok()
770        );
771        assert!(
772            dispatcher
773                .emit_end(LlmEndEvent::new("req-1", "model"))
774                .is_ok()
775        );
776        assert!(runtime.shutdown().await.is_ok());
777        assert_eq!(*events.lock(), vec!["start", "end"]);
778    }
779}