inferlab_proxy/
core.rs

1//! Shared HTTP mechanics for the built-in disaggregated-serving proxies.
2//!
3//! Proxy-specific protocol bodies remain in their owning modules.
4
5use crate::error::ProxyError as ProxyLifecycleError;
6use async_stream::try_stream;
7use axum::Json;
8use axum::body::Body;
9use axum::http::{HeaderMap, Response, StatusCode, header};
10use axum::response::IntoResponse;
11use bytes::Bytes;
12use futures_util::{FutureExt, Stream, StreamExt};
13use serde::Serialize;
14use serde_json::Value;
15use std::env;
16use std::fmt;
17use std::future::Future;
18use std::sync::atomic::{AtomicUsize, Ordering};
19use tokio::task::JoinHandle;
20
21/// Identity of a built-in proxy. Each proxy module owns its own [`ProxyMeta`]
22/// so the proxy crate is the authority for the id/version recorded in
23/// `BuiltinProxy` evidence.
24#[derive(Clone, Copy, Debug, Eq, PartialEq)]
25pub struct ProxyMeta {
26    pub id: &'static str,
27    pub version: u32,
28}
29
30/// Build a multi-threaded Tokio runtime and drive `run_async` to completion.
31///
32/// Built-in proxies share this runtime-builder wrapper; the per-proxy
33/// `run` functions call it with their own async entrypoint.
34pub fn run<F, Fut>(run_async: F) -> Result<(), ProxyLifecycleError>
35where
36    F: FnOnce() -> Fut,
37    Fut: Future<Output = Result<(), ProxyLifecycleError>>,
38{
39    let runtime = tokio::runtime::Builder::new_multi_thread()
40        .enable_all()
41        .build()
42        .map_err(|error| ProxyLifecycleError::Lifecycle {
43            message: format!("failed to create proxy tokio runtime: {error}"),
44        })?;
45    runtime.block_on(run_async())
46}
47
48/// Healthcheck response body shared by the proxies.
49#[derive(Serialize)]
50pub struct ProxyHealthcheckResponse {
51    pub ready: bool,
52    pub prefill_instances: usize,
53    pub decode_instances: usize,
54}
55
56/// Forward an upstream response body verbatim, preserving status and
57/// content-type.
58pub async fn forward_response(
59    response: reqwest::Response,
60) -> Result<Response<Body>, ProxyHttpError> {
61    let status = status_code(response.status())?;
62    let content_type = response
63        .headers()
64        .get(reqwest::header::CONTENT_TYPE)
65        .and_then(|value| value.to_str().ok())
66        .map(str::to_owned);
67    let bytes = response
68        .bytes()
69        .await
70        .map_err(|error| ProxyHttpError::upstream("upstream response body read failed", error))?;
71    let mut builder = Response::builder().status(status);
72    if let Some(content_type) = content_type {
73        builder = builder.header(header::CONTENT_TYPE, content_type);
74    }
75    builder.body(Body::from(bytes)).map_err(|error| {
76        ProxyHttpError::internal(format!("failed to build proxy response: {error}"))
77    })
78}
79
80/// Convert an unsuccessful upstream response into a `502 Bad Gateway`
81/// [`ProxyHttpError`] that captures the upstream status and body.
82pub async fn upstream_status_error(context: &str, response: reqwest::Response) -> ProxyHttpError {
83    let status = response.status();
84    let body = match response.text().await {
85        Ok(text) => text,
86        Err(error) => format!("<failed to read upstream error body: {error}>"),
87    };
88    ProxyHttpError::status(
89        StatusCode::BAD_GATEWAY,
90        format!("{context} returned HTTP {status}: {body}"),
91    )
92}
93
94/// Resolve the outbound `Authorization` header from the inbound request or the
95/// `OPENAI_API_KEY` environment variable.
96pub fn outbound_authorization(headers: &HeaderMap) -> Option<String> {
97    headers
98        .get(header::AUTHORIZATION)
99        .and_then(|value| value.to_str().ok())
100        .map(str::to_owned)
101        .or_else(|| {
102            env::var("OPENAI_API_KEY")
103                .ok()
104                .map(|key| format!("Bearer {key}"))
105        })
106}
107
108/// Join a base URL with a path, normalizing a single trailing slash on the
109/// base.
110pub fn join_path(base: &str, path: &str) -> String {
111    format!("{}{}", base.trim_end_matches('/'), path)
112}
113
114/// Convert a `reqwest` status code into an `axum`/`http` status code.
115pub fn status_code(status: reqwest::StatusCode) -> Result<StatusCode, ProxyHttpError> {
116    StatusCode::from_u16(status.as_u16())
117        .map_err(|error| ProxyHttpError::internal(format!("invalid upstream status code: {error}")))
118}
119
120/// Advance a round-robin cursor and return the selected index into a non-empty
121/// target list. Shared by all proxies' prefill/decode selection; each proxy
122/// keeps its own cursor and target list (which stay local).
123pub(crate) fn round_robin_index(cursor: &AtomicUsize, len: usize) -> usize {
124    cursor.fetch_add(1, Ordering::SeqCst) % len
125}
126
127/// Build and send a JSON POST to `url` with an optional `X-Request-Id`, any
128/// `extra_headers`, and an optional `Authorization`, returning the response or a
129/// [`ProxyHttpError`] on transport or non-success status. `context` names the
130/// call in error messages (e.g. "decode request"). Owns the transport for every
131/// built-in proxy POST. Proxy-specific request ids and headers are optional.
132pub(crate) async fn send_json_post(
133    client: reqwest::Client,
134    url: String,
135    body: &Value,
136    request_id: Option<&str>,
137    authorization: Option<&str>,
138    extra_headers: &[(&str, String)],
139    context: &'static str,
140) -> Result<reqwest::Response, ProxyHttpError> {
141    let mut request = client.post(url).json(body);
142    if let Some(request_id) = request_id {
143        request = request.header("X-Request-Id", request_id);
144    }
145    // Extra headers precede `Authorization`: header insertion order reaches
146    // the wire, and Mooncake's prefill always sent its rank header first.
147    for (name, value) in extra_headers {
148        request = request.header(*name, value);
149    }
150    if let Some(authorization) = authorization {
151        request = request.header(reqwest::header::AUTHORIZATION, authorization);
152    }
153    let response = request
154        .send()
155        .await
156        .map_err(|error| ProxyHttpError::upstream(&format!("{context} failed"), error))?;
157    if !response.status().is_success() {
158        return Err(upstream_status_error(context, response).await);
159    }
160    Ok(response)
161}
162
163/// A per-process monotonic request id, `"{pid}-{n}"`, drawn from a proxy-owned
164/// counter. Shared by the vLLM proxies so the id scheme has one home.
165pub(crate) fn next_request_id(counter: &AtomicUsize) -> String {
166    let value = counter.fetch_add(1, Ordering::SeqCst);
167    format!("{}-{value}", std::process::id())
168}
169
170/// Build the outbound HTTP client shared by the proxies, with the pool tuning
171/// (unbounded idle connections per host) both require. Returns the raw
172/// `reqwest` error so each proxy keeps its own construction-failure message.
173pub(crate) fn build_pooled_client() -> reqwest::Result<reqwest::Client> {
174    reqwest::Client::builder()
175        .pool_max_idle_per_host(usize::MAX)
176        .build()
177}
178
179/// Error type shared by both proxies, carrying an HTTP status and a message.
180#[derive(Debug)]
181pub struct ProxyHttpError {
182    status: StatusCode,
183    message: String,
184}
185
186impl ProxyHttpError {
187    pub fn status(status: StatusCode, message: impl Into<String>) -> Self {
188        Self {
189            status,
190            message: message.into(),
191        }
192    }
193
194    pub fn upstream(context: &str, error: reqwest::Error) -> Self {
195        Self::status(StatusCode::BAD_GATEWAY, format!("{context}: {error}"))
196    }
197
198    pub fn internal(message: impl Into<String>) -> Self {
199        Self::status(StatusCode::INTERNAL_SERVER_ERROR, message)
200    }
201}
202
203impl fmt::Display for ProxyHttpError {
204    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
205        write!(formatter, "{}", self.message)
206    }
207}
208
209impl std::error::Error for ProxyHttpError {}
210
211impl IntoResponse for ProxyHttpError {
212    fn into_response(self) -> axum::response::Response {
213        let body = Json(ProxyErrorResponse {
214            error: self.message,
215        });
216        (self.status, body).into_response()
217    }
218}
219
220#[derive(Serialize)]
221pub struct ProxyErrorResponse {
222    pub error: String,
223}
224
225/// Stream a decode response body while a concurrently-running prefill task
226/// completes. Used by proxies whose backend protocol starts both roles
227/// together; the vLLM NIXL proxy instead forwards them sequentially. If the
228/// client drops the response before prefill finishes,
229/// the prefill task is aborted (via [`AbortOnDrop`]); once prefill completes the
230/// abort is disarmed, and a prefill failure surfaces as a stream error.
231pub(crate) fn stream_decode_response(
232    response: reqwest::Response,
233    prefill_task: JoinHandle<Result<(), ProxyHttpError>>,
234) -> Result<Response<Body>, ProxyHttpError> {
235    let status = status_code(response.status())?;
236    let content_type = response
237        .headers()
238        .get(reqwest::header::CONTENT_TYPE)
239        .and_then(|value| value.to_str().ok())
240        .map(str::to_owned);
241    let stream = decode_response_stream(response.bytes_stream(), prefill_task);
242    let mut builder = Response::builder().status(status);
243    if let Some(content_type) = content_type {
244        builder = builder.header(header::CONTENT_TYPE, content_type);
245    }
246    builder.body(Body::from_stream(stream)).map_err(|error| {
247        ProxyHttpError::internal(format!("failed to build proxy response: {error}"))
248    })
249}
250
251/// Stream one successful upstream response without waiting for another role.
252pub(crate) fn stream_response(
253    response: reqwest::Response,
254) -> Result<Response<Body>, ProxyHttpError> {
255    let status = status_code(response.status())?;
256    let content_type = response
257        .headers()
258        .get(reqwest::header::CONTENT_TYPE)
259        .and_then(|value| value.to_str().ok())
260        .map(str::to_owned);
261    let stream = response
262        .bytes_stream()
263        .map(|chunk| chunk.map_err(|error| stream_error(format!("decode stream failed: {error}"))));
264    let mut builder = Response::builder().status(status);
265    if let Some(content_type) = content_type {
266        builder = builder.header(header::CONTENT_TYPE, content_type);
267    }
268    builder.body(Body::from_stream(stream)).map_err(|error| {
269        ProxyHttpError::internal(format!("failed to build proxy response: {error}"))
270    })
271}
272
273/// The decode byte stream, generic over the decode stream and its error type so
274/// it can be exercised without a live `reqwest::Response`. Yields decode bytes
275/// in arrival order; concurrently drives `prefill_task` to completion and aborts
276/// it if the consumer drops the stream before prefill finishes. On decode EOF
277/// before prefill completes, prefill is awaited and its error (if any) surfaces.
278pub(crate) fn decode_response_stream<S, E>(
279    decode_stream: S,
280    prefill_task: JoinHandle<Result<(), ProxyHttpError>>,
281) -> impl Stream<Item = std::result::Result<Bytes, std::io::Error>>
282where
283    S: Stream<Item = std::result::Result<Bytes, E>> + Unpin,
284    E: fmt::Display,
285{
286    let prefill_abort = prefill_task.abort_handle();
287    try_stream! {
288        let mut decode_stream = decode_stream;
289        let mut prefill_task = prefill_task;
290        let mut prefill_abort = AbortOnDrop::new(prefill_abort);
291        let mut prefill_done = false;
292        loop {
293            match next_stream_event(&mut prefill_task, &mut decode_stream, prefill_done).await {
294                StreamEvent::Prefill(prefill) => {
295                    prefill_done = true;
296                    // One-time tie-break: if a decode item was already ready at the
297                    // instant prefill completed, handle that single item before
298                    // surfacing the prefill outcome. `now_or_never()` polls (and so
299                    // consumes) the item, so it must be matched exhaustively: deliver
300                    // a ready chunk, and PROPAGATE a ready decode error (an Ok-only
301                    // match would drop it and truncate the response into a clean 200).
302                    // EOF / not-ready fall through to the
303                    // prefill outcome; decode is never indefinitely preferred.
304                    match decode_stream.next().now_or_never() {
305                        Some(Some(Ok(bytes))) => yield bytes,
306                        Some(Some(Err(error))) => {
307                            Err(stream_error(format!("decode stream failed: {error}")))?;
308                        }
309                        Some(None) | None => {}
310                    }
311                    prefill
312                        .map_err(join_error)?
313                        .map_err(|error| stream_error(error.to_string()))?;
314                    prefill_abort.disarm();
315                }
316                StreamEvent::Decode(Some(Ok(bytes))) => yield bytes,
317                StreamEvent::Decode(Some(Err(error))) => {
318                    Err(stream_error(format!("decode stream failed: {error}")))?;
319                }
320                StreamEvent::Decode(None) => break,
321            }
322        }
323        if !prefill_done {
324            prefill_task
325                .await
326                .map_err(join_error)?
327                .map_err(|error| stream_error(error.to_string()))?;
328            prefill_abort.disarm();
329        }
330    }
331}
332
333enum StreamEvent<E> {
334    Prefill(std::result::Result<Result<(), ProxyHttpError>, tokio::task::JoinError>),
335    Decode(Option<std::result::Result<Bytes, E>>),
336}
337
338async fn next_stream_event<S, E>(
339    prefill_task: &mut JoinHandle<Result<(), ProxyHttpError>>,
340    decode_stream: &mut S,
341    prefill_done: bool,
342) -> StreamEvent<E>
343where
344    S: Stream<Item = std::result::Result<Bytes, E>> + Unpin,
345{
346    // Surface the prefill outcome promptly once the prefill task has COMPLETED, so
347    // a continuously-ready decode stream cannot defer (and thereby suppress) a
348    // prefill failure indefinitely. The caller delivers one already-ready decode
349    // chunk before propagating a prefill error (a one-time tie-break), so a chunk
350    // that was ready at the instant prefill finished is not dropped — but decode
351    // is NOT permanently prioritized.
352    if !prefill_done && prefill_task.is_finished() {
353        return StreamEvent::Prefill(prefill_task.await);
354    }
355    // Prefill is still running: deliver decode bytes as they arrive, and otherwise
356    // await the prefill task's completion (picked up by the `is_finished` check on
357    // the next call). An unbiased race is fine here — there is no completed prefill
358    // outcome to drop, and a ready decode chunk taken by its own branch is yielded,
359    // not lost.
360    tokio::select! {
361        prefill = prefill_task, if !prefill_done => StreamEvent::Prefill(prefill),
362        chunk = decode_stream.next() => StreamEvent::Decode(chunk),
363    }
364}
365
366fn join_error(error: tokio::task::JoinError) -> std::io::Error {
367    stream_error(format!("prefill task failed: {error}"))
368}
369
370fn stream_error(message: String) -> std::io::Error {
371    std::io::Error::other(message)
372}
373
374/// Aborts the held task when dropped unless [`disarm`](AbortOnDrop::disarm)ed.
375struct AbortOnDrop {
376    handle: tokio::task::AbortHandle,
377    armed: bool,
378}
379
380impl AbortOnDrop {
381    fn new(handle: tokio::task::AbortHandle) -> Self {
382        Self {
383            handle,
384            armed: true,
385        }
386    }
387
388    fn disarm(&mut self) {
389        self.armed = false;
390    }
391}
392
393impl Drop for AbortOnDrop {
394    fn drop(&mut self) {
395        if self.armed {
396            self.handle.abort();
397        }
398    }
399}
400
401#[cfg(test)]
402mod tests {
403    use super::*;
404    use anyhow::{Context, Result};
405
406    #[test]
407    fn join_path_normalizes_single_trailing_slash() {
408        assert_eq!(
409            join_path("http://h:1/", "/v1/models"),
410            "http://h:1/v1/models"
411        );
412        assert_eq!(
413            join_path("http://h:1", "/v1/models"),
414            "http://h:1/v1/models"
415        );
416    }
417
418    #[test]
419    fn status_code_maps_reqwest_status() -> Result<()> {
420        let mapped = status_code(reqwest::StatusCode::OK)
421            .map_err(|error| anyhow::anyhow!(error.to_string()))?;
422        assert_eq!(mapped, StatusCode::OK);
423        Ok(())
424    }
425
426    #[test]
427    fn outbound_authorization_prefers_inbound_header() -> Result<()> {
428        let mut headers = HeaderMap::new();
429        headers.insert(header::AUTHORIZATION, "Bearer inbound".parse()?);
430        assert_eq!(
431            outbound_authorization(&headers),
432            Some("Bearer inbound".to_owned())
433        );
434        Ok(())
435    }
436
437    #[test]
438    fn proxy_error_internal_uses_500() {
439        let error = ProxyHttpError::internal("boom");
440        assert_eq!(error.status, StatusCode::INTERNAL_SERVER_ERROR);
441        assert_eq!(error.to_string(), "boom");
442    }
443
444    use std::sync::Arc;
445    use std::sync::atomic::{AtomicBool, Ordering};
446
447    /// Flips a shared flag when dropped — used to observe that an aborted
448    /// prefill task is actually cancelled (its future is dropped).
449    struct SetOnDrop(Arc<AtomicBool>);
450
451    impl Drop for SetOnDrop {
452        fn drop(&mut self) {
453            self.0.store(true, Ordering::SeqCst);
454        }
455    }
456
457    fn proxy_test_runtime() -> Result<tokio::runtime::Runtime> {
458        tokio::runtime::Builder::new_multi_thread()
459            .enable_all()
460            .build()
461            .map_err(|error| anyhow::anyhow!(error.to_string()))
462    }
463
464    #[test]
465    fn streamed_decode_yields_bytes_in_order_when_prefill_succeeds() -> Result<()> {
466        let runtime = proxy_test_runtime()?;
467        let bytes = runtime.block_on(async {
468            let decode = Box::pin(futures_util::stream::iter(vec![
469                std::result::Result::<Bytes, std::io::Error>::Ok(Bytes::from_static(b"hello")),
470                Ok(Bytes::from_static(b" world")),
471            ]));
472            let prefill = tokio::spawn(async { Ok::<(), ProxyHttpError>(()) });
473            let mut stream = Box::pin(decode_response_stream(decode, prefill));
474            let mut out = Vec::new();
475            while let Some(item) = stream.next().await {
476                out.push(item.map_err(|error| anyhow::anyhow!(error.to_string()))?);
477            }
478            anyhow::Ok(out)
479        })?;
480        let joined: Vec<u8> = bytes.into_iter().flatten().collect();
481        assert_eq!(joined, b"hello world");
482        Ok(())
483    }
484
485    #[test]
486    fn streamed_decode_surfaces_prefill_error_after_decode_ends() -> Result<()> {
487        let runtime = proxy_test_runtime()?;
488        let (bytes, error) = runtime.block_on(async {
489            let decode = Box::pin(futures_util::stream::iter(vec![std::result::Result::<
490                Bytes,
491                std::io::Error,
492            >::Ok(
493                Bytes::from_static(b"partial"),
494            )]));
495            let prefill = tokio::spawn(async {
496                Err::<(), ProxyHttpError>(ProxyHttpError::internal("prefill boom"))
497            });
498            let mut stream = Box::pin(decode_response_stream(decode, prefill));
499            let mut bytes = Vec::new();
500            let mut error = None;
501            while let Some(item) = stream.next().await {
502                match item {
503                    Ok(chunk) => bytes.extend_from_slice(&chunk),
504                    Err(stream_error) => {
505                        error = Some(stream_error.to_string());
506                        break;
507                    }
508                }
509            }
510            anyhow::Ok((bytes, error))
511        })?;
512        assert_eq!(bytes, b"partial");
513        let error = error.context("expected a prefill error to surface after decode ended")?;
514        assert!(error.contains("prefill boom"), "got {error}");
515        Ok(())
516    }
517
518    /// A prefill failure must surface even while the decode stream stays
519    /// continuously ready. The one-time tie-break delivers an already-ready chunk
520    /// but must NOT let an always-ready decode stream defer the prefill error
521    /// indefinitely (a permanent decode bias would suppress it).
522    #[test]
523    fn prefill_error_surfaces_even_while_decode_stays_ready() -> Result<()> {
524        let runtime = proxy_test_runtime()?;
525        let error = runtime.block_on(async {
526            // An unbounded, always-synchronously-ready decode stream.
527            let decode = Box::pin(futures_util::stream::repeat_with(|| {
528                std::result::Result::<Bytes, std::io::Error>::Ok(Bytes::from_static(b"x"))
529            }));
530            let prefill = tokio::spawn(async {
531                Err::<(), ProxyHttpError>(ProxyHttpError::internal("prefill boom"))
532            });
533            let mut stream = Box::pin(decode_response_stream(decode, prefill));
534            let mut chunks = 0usize;
535            let mut error = None;
536            while let Some(item) = stream.next().await {
537                match item {
538                    Ok(_) => {
539                        chunks += 1;
540                        // Bound: prove the error is not suppressed forever. The fix
541                        // surfaces it within a handful of chunks; a regression that
542                        // permanently prefers decode would never break out here.
543                        assert!(
544                            chunks < 100_000,
545                            "prefill error was suppressed by a continuously-ready decode stream"
546                        );
547                    }
548                    Err(stream_error) => {
549                        error = Some(stream_error.to_string());
550                        break;
551                    }
552                }
553            }
554            anyhow::Ok(error)
555        })?;
556        let error = error.context("a prefill error must surface even while decode stays ready")?;
557        assert!(error.contains("prefill boom"), "got {error}");
558        Ok(())
559    }
560
561    /// A decode error that is ALREADY ready at the instant prefill completes must be
562    /// propagated by the one-time tie-break, not silently dropped. `now_or_never()`
563    /// polls (and thus consumes) that ready item, so an Ok-only match would discard
564    /// the error; with a successful prefill the stream would then end cleanly,
565    /// turning a decode failure into a truncated 200.
566    #[test]
567    fn decode_error_ready_at_tiebreak_is_not_swallowed() -> Result<()> {
568        let runtime = proxy_test_runtime()?;
569        let error = runtime.block_on(async {
570            // Force prefill to be FINISHED (Ok) so the first stream event is the
571            // prefill outcome and the tie-break is what polls the decode stream.
572            let prefill = tokio::spawn(async { Ok::<(), ProxyHttpError>(()) });
573            while !prefill.is_finished() {
574                tokio::task::yield_now().await;
575            }
576            // A synchronously-ready decode Err waiting at the tie-break instant.
577            let decode = Box::pin(futures_util::stream::iter(vec![std::result::Result::<
578                Bytes,
579                std::io::Error,
580            >::Err(
581                std::io::Error::other("decode boom"),
582            )]));
583            let mut stream = Box::pin(decode_response_stream(decode, prefill));
584            let mut error = None;
585            while let Some(item) = stream.next().await {
586                if let Err(stream_error) = item {
587                    error = Some(stream_error.to_string());
588                    break;
589                }
590            }
591            anyhow::Ok(error)
592        })?;
593        let error =
594            error.context("a decode error ready at the tie-break must surface, not truncate")?;
595        assert!(error.contains("decode boom"), "got {error}");
596        Ok(())
597    }
598
599    #[test]
600    fn dropping_the_stream_before_prefill_finishes_aborts_prefill() -> Result<()> {
601        let runtime = proxy_test_runtime()?;
602        let aborted = Arc::new(AtomicBool::new(false));
603        let flag = aborted.clone();
604        let cancelled = runtime.block_on(async move {
605            let (started_tx, started_rx) = tokio::sync::oneshot::channel::<()>();
606            // Prefill never completes; once aborted, its future is dropped and
607            // SetOnDrop flips the flag. It signals `started` only AFTER the guard
608            // is constructed, so the drop-abort is observed deterministically (no
609            // race on whether the task was polled before the abort fired).
610            let prefill = tokio::spawn(async move {
611                let _guard = SetOnDrop(flag);
612                let _ = started_tx.send(());
613                futures_util::future::pending::<()>().await;
614                Ok::<(), ProxyHttpError>(())
615            });
616            let _ = started_rx.await;
617            // Decode yields one chunk then stays pending, so the loop neither
618            // breaks (decode EOF) nor selects prefill — leaving prefill in flight.
619            let decode = Box::pin(
620                futures_util::stream::once(async {
621                    std::result::Result::<Bytes, std::io::Error>::Ok(Bytes::from_static(b"a"))
622                })
623                .chain(futures_util::stream::pending::<
624                    std::result::Result<Bytes, std::io::Error>,
625                >()),
626            );
627            let mut stream = Box::pin(decode_response_stream(decode, prefill));
628            assert!(matches!(stream.next().await, Some(Ok(_))));
629            drop(stream);
630            for _ in 0..200 {
631                if aborted.load(Ordering::SeqCst) {
632                    return true;
633                }
634                tokio::time::sleep(std::time::Duration::from_millis(10)).await;
635            }
636            false
637        });
638        assert!(
639            cancelled,
640            "prefill task was not aborted when the response stream was dropped"
641        );
642        Ok(())
643    }
644}
inferlab_proxy/core.rs

inferlab_proxy/
core.rs
