inferlab_proxy/
core.rs

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

inferlab_proxy/
core.rs
