vllm_cpp/
callback.rs

1use std::any::Any;
2use std::ffi::CStr;
3use std::os::raw::{c_char, c_void};
4use std::panic::{catch_unwind, AssertUnwindSafe};
5
6use crate::error::Error;
7
8/// Controls whether native streaming continues after a callback.
9#[derive(Clone, Copy, Debug, Eq, PartialEq)]
10pub enum StreamControl {
11    Continue,
12    Stop,
13}
14
15/// One copied streaming delta.
16#[derive(Clone, Debug, Eq, PartialEq)]
17pub struct StreamEvent {
18    pub delta: String,
19    pub finished: bool,
20}
21
22/// How a successful blocking stream ended.
23#[derive(Clone, Copy, Debug, Eq, PartialEq)]
24pub struct StreamOutcome {
25    pub stopped_by_callback: bool,
26}
27
28pub(crate) struct CallbackState<'callback, F> {
29    callback: &'callback mut F,
30    stopped: bool,
31    error: Option<Error>,
32    panic: Option<Box<dyn Any + Send>>,
33}
34
35impl<'callback, F> CallbackState<'callback, F> {
36    pub(crate) fn new(callback: &'callback mut F) -> Self {
37        Self {
38            callback,
39            stopped: false,
40            error: None,
41            panic: None,
42        }
43    }
44
45    pub(crate) const fn stopped(&self) -> bool {
46        self.stopped
47    }
48
49    pub(crate) fn take_error(&mut self) -> Option<Error> {
50        self.error.take()
51    }
52
53    pub(crate) fn take_panic(&mut self) -> Option<Box<dyn Any + Send>> {
54        self.panic.take()
55    }
56
57    fn apply_control(&mut self, control: StreamControl, finished: bool) -> bool {
58        match control {
59            StreamControl::Continue => true,
60            StreamControl::Stop => {
61                if !finished {
62                    self.stopped = true;
63                }
64                false
65            }
66        }
67    }
68}
69
70pub(crate) unsafe extern "C" fn callback_trampoline<F>(
71    delta_text: *const c_char,
72    finished: bool,
73    user_data: *mut c_void,
74) -> bool
75where
76    F: FnMut(StreamEvent) -> StreamControl,
77{
78    // SAFETY: callers pass a stable pointer to CallbackState<F> and the native
79    // blocking function cannot retain it after returning.
80    let state = unsafe { &mut *user_data.cast::<CallbackState<'_, F>>() };
81    if delta_text.is_null() {
82        state.error = Some(Error::InvalidUtf8 {
83            field: "stream delta",
84        });
85        return false;
86    }
87    // SAFETY: vllm.cpp promises a borrowed NUL-terminated string for the callback.
88    let delta = match unsafe { CStr::from_ptr(delta_text) }.to_str() {
89        Ok(delta) => delta.to_owned(),
90        Err(_) => {
91            state.error = Some(Error::InvalidUtf8 {
92                field: "stream delta",
93            });
94            return false;
95        }
96    };
97    let event = StreamEvent { delta, finished };
98    match catch_unwind(AssertUnwindSafe(|| (state.callback)(event))) {
99        Ok(control) => state.apply_control(control, finished),
100        Err(payload) => {
101            state.panic = Some(payload);
102            false
103        }
104    }
105}
106
107#[cfg(test)]
108mod tests {
109    use super::{CallbackState, StreamControl, StreamEvent};
110
111    #[test]
112    fn only_nonterminal_stop_marks_callback_stop() {
113        let mut callback = |_: StreamEvent| StreamControl::Continue;
114        let mut state = CallbackState::new(&mut callback);
115
116        assert!(!state.apply_control(StreamControl::Stop, true));
117        assert!(!state.stopped());
118        assert!(!state.apply_control(StreamControl::Stop, false));
119        assert!(state.stopped());
120    }
121}
vllm_cpp/callback.rs

vllm_cpp/
callback.rs
