use crate::core::{
self, ProxyHealthcheckResponse, ProxyHttpError, ProxyMeta, forward_response, join_path,
outbound_authorization,
};
use crate::error::ProxyError;
use async_stream::stream;
use axum::Json;
use axum::body::Body;
use axum::extract::State;
use axum::http::{HeaderMap, Response, StatusCode, header};
use axum::response::IntoResponse;
use axum::routing::{get, post};
use axum::{Router, serve};
use bytes::Bytes;
use futures_util::{Stream, StreamExt, future::join_all};
use serde::Serialize;
use serde_json::Value;
use std::fmt;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering};
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use tokio::net::TcpListener;
use tokio::task::JoinHandle;
pub const ID: &str = "inferlab-sglang-proxy";
pub const VERSION: u32 = 2;
pub fn meta() -> ProxyMeta {
ProxyMeta {
id: ID,
version: VERSION,
}
}
#[derive(Clone, Debug)]
pub struct Config {
pub host: String,
pub port: u16,
pub prefill: Vec<PrefillTarget>,
pub decode: Vec<String>,
}
#[derive(Clone, Debug)]
pub struct PrefillTarget {
pub url: String,
pub bootstrap_host: String,
pub bootstrap_port: u16,
}
pub fn run(config: Config) -> Result<(), ProxyError> {
core::run(|| run_async(config))
}
pub async fn run_async(config: Config) -> Result<(), ProxyError> {
let host = config.host.clone();
let port = config.port;
let state = ProxyState::new(config)?;
tokio::spawn(await_backends(state.clone()));
let listener = TcpListener::bind((host.as_str(), port))
.await
.map_err(|error| ProxyError::Io {
message: format!("failed to bind SGLang proxy on {host}:{port}: {error}"),
})?;
serve(listener, router(state))
.await
.map_err(|error| ProxyError::Io {
message: format!("SGLang proxy server failed: {error}"),
})
}
fn router(state: ProxyState) -> Router {
Router::new()
.route("/healthcheck", get(healthcheck))
.route("/v1/completions", post(completions))
.route("/v1/chat/completions", post(chat_completions))
.route("/flush_cache", post(flush_cache))
.with_state(state)
}
#[derive(Clone)]
struct ProxyState {
inner: Arc<ProxyStateInner>,
}
struct ProxyStateInner {
client: reqwest::Client,
prefill: Vec<PrefillTarget>,
decode: Vec<String>,
ready: AtomicBool,
prefill_cursor: AtomicUsize,
decode_cursor: AtomicUsize,
room_seed: u64,
room_counter: AtomicU64,
}
impl ProxyState {
fn new(config: Config) -> Result<Self, ProxyError> {
if config.prefill.is_empty() {
return Err(ProxyError::Invalid {
message: "SGLang proxy requires at least one prefill endpoint".to_owned(),
});
}
if config.decode.is_empty() {
return Err(ProxyError::Invalid {
message: "SGLang proxy requires at least one decode endpoint".to_owned(),
});
}
let client = core::build_pooled_client().map_err(|error| ProxyError::Io {
message: format!("failed to create SGLang proxy HTTP client: {error}"),
})?;
Ok(Self {
inner: Arc::new(ProxyStateInner {
client,
prefill: config.prefill,
decode: config.decode,
ready: AtomicBool::new(false),
prefill_cursor: AtomicUsize::new(0),
decode_cursor: AtomicUsize::new(0),
room_seed: room_seed(),
room_counter: AtomicU64::new(0),
}),
})
}
fn client(&self) -> reqwest::Client {
self.inner.client.clone()
}
fn ready(&self) -> bool {
self.inner.ready.load(Ordering::SeqCst)
}
fn set_ready(&self) {
self.inner.ready.store(true, Ordering::SeqCst);
}
fn next_prefill(&self) -> PrefillTarget {
let index = core::round_robin_index(&self.inner.prefill_cursor, self.inner.prefill.len());
self.inner.prefill[index].clone()
}
fn next_decode(&self) -> String {
let index = core::round_robin_index(&self.inner.decode_cursor, self.inner.decode.len());
self.inner.decode[index].clone()
}
fn next_room(&self) -> u64 {
let counter = self.inner.room_counter.fetch_add(1, Ordering::SeqCst);
self.inner.room_seed.wrapping_add(counter) & ((1_u64 << 63) - 1)
}
}
fn room_seed() -> u64 {
let nanos = SystemTime::now()
.duration_since(UNIX_EPOCH)
.map_or(0, |elapsed| elapsed.as_nanos() as u64);
(nanos ^ (u64::from(std::process::id()) << 32)) & ((1_u64 << 63) - 1)
}
async fn await_backends(state: ProxyState) {
let urls = state
.inner
.prefill
.iter()
.map(|target| target.url.clone())
.chain(state.inner.decode.iter().cloned());
let waits = urls.map(|url| await_backend(state.client(), url));
join_all(waits).await;
state.set_ready();
}
async fn await_backend(client: reqwest::Client, url: String) {
loop {
if client
.get(join_path(&url, "/v1/models"))
.send()
.await
.is_ok_and(|response| response.status().is_success())
{
return;
}
tokio::time::sleep(Duration::from_secs(1)).await;
}
}
async fn healthcheck(
State(state): State<ProxyState>,
) -> (StatusCode, Json<ProxyHealthcheckResponse>) {
let ready = state.ready();
let status = if ready {
StatusCode::OK
} else {
StatusCode::SERVICE_UNAVAILABLE
};
(
status,
Json(ProxyHealthcheckResponse {
ready,
prefill_instances: state.inner.prefill.len(),
decode_instances: state.inner.decode.len(),
}),
)
}
async fn completions(
State(state): State<ProxyState>,
headers: HeaderMap,
Json(body): Json<Value>,
) -> Result<Response<Body>, ProxyHttpError> {
completion_route(state, headers, body).await
}
async fn chat_completions(
State(state): State<ProxyState>,
headers: HeaderMap,
Json(body): Json<Value>,
) -> Result<Response<Body>, ProxyHttpError> {
request_route(state, headers, body, "/v1/chat/completions").await
}
async fn completion_route(
state: ProxyState,
headers: HeaderMap,
body: Value,
) -> Result<Response<Body>, ProxyHttpError> {
request_route(state, headers, body, "/v1/completions").await
}
async fn request_route(
state: ProxyState,
headers: HeaderMap,
body: Value,
path: &'static str,
) -> Result<Response<Body>, ProxyHttpError> {
if !state.ready() {
return Err(ProxyHttpError::status(
StatusCode::SERVICE_UNAVAILABLE,
"proxy is not ready",
));
}
let stream = body.get("stream").and_then(Value::as_bool).unwrap_or(false);
let prefill = state.next_prefill();
let decode = state.next_decode();
let request_body = bootstrap_body(&body, &prefill, state.next_room(), path)?;
let authorization = outbound_authorization(&headers);
let client = state.client();
let prefill_url = join_path(&prefill.url, path);
let prefill_body = request_body.clone();
let prefill_authorization = authorization.clone();
let prefill_client = client.clone();
let prefill_task = tokio::spawn(async move {
let response = core::send_json_post(
prefill_client,
prefill_url,
&prefill_body,
None,
prefill_authorization.as_deref(),
&[],
"prefill request",
)
.await?;
drain_response(response, "prefill").await
});
let decode_result = core::send_json_post(
client,
join_path(&decode, path),
&request_body,
None,
authorization.as_deref(),
&[],
"decode request",
)
.await;
let decode_response = match decode_result {
Ok(response) => response,
Err(error) => {
drop(prefill_task);
return Err(error);
}
};
if stream {
if prefill_task.is_finished() {
await_prefill(prefill_task).await?;
core::stream_response(decode_response)
} else if is_text_event_stream(&decode_response) {
stream_sse_decode_response(decode_response, prefill_task)
} else {
stream_detached_decode_response(decode_response, prefill_task)
}
} else {
let (prefill_result, decode_result) = tokio::join!(
await_prefill(prefill_task),
forward_response(decode_response)
);
prefill_result?;
decode_result
}
}
fn is_text_event_stream(response: &reqwest::Response) -> bool {
response
.headers()
.get(reqwest::header::CONTENT_TYPE)
.and_then(|value| value.to_str().ok())
.and_then(|value| value.split(';').next())
.is_some_and(|value| value.trim().eq_ignore_ascii_case("text/event-stream"))
}
fn stream_sse_decode_response(
response: reqwest::Response,
prefill_task: JoinHandle<Result<(), ProxyHttpError>>,
) -> Result<Response<Body>, ProxyHttpError> {
let status = core::status_code(response.status())?;
let content_type = response
.headers()
.get(reqwest::header::CONTENT_TYPE)
.and_then(|value| value.to_str().ok())
.map(str::to_owned);
let stream = sse_decode_response_stream(response.bytes_stream(), prefill_task);
let mut builder = Response::builder().status(status);
if let Some(content_type) = content_type {
builder = builder.header(header::CONTENT_TYPE, content_type);
}
builder.body(Body::from_stream(stream)).map_err(|error| {
ProxyHttpError::internal(format!("failed to build proxy response: {error}"))
})
}
fn stream_detached_decode_response(
response: reqwest::Response,
prefill_task: JoinHandle<Result<(), ProxyHttpError>>,
) -> Result<Response<Body>, ProxyHttpError> {
let status = core::status_code(response.status())?;
let content_type = response
.headers()
.get(reqwest::header::CONTENT_TYPE)
.and_then(|value| value.to_str().ok())
.map(str::to_owned);
let stream = detached_decode_response_stream(response.bytes_stream(), prefill_task);
let mut builder = Response::builder().status(status);
if let Some(content_type) = content_type {
builder = builder.header(header::CONTENT_TYPE, content_type);
}
builder.body(Body::from_stream(stream)).map_err(|error| {
ProxyHttpError::internal(format!("failed to build proxy response: {error}"))
})
}
fn detached_decode_response_stream<S, E>(
decode_stream: S,
prefill_task: JoinHandle<Result<(), ProxyHttpError>>,
) -> impl Stream<Item = Result<Bytes, std::io::Error>>
where
S: Stream<Item = Result<Bytes, E>> + Unpin,
E: fmt::Display,
{
stream! {
let mut decode_stream = decode_stream;
let mut prefill_task = prefill_task;
let mut prefill_done = false;
loop {
if !prefill_done && prefill_task.is_finished() {
prefill_done = true;
if let Err(error) = prefill_stream_outcome((&mut prefill_task).await) {
yield Err(error);
return;
}
}
tokio::select! {
prefill = &mut prefill_task, if !prefill_done => {
prefill_done = true;
if let Err(error) = prefill_stream_outcome(prefill) {
yield Err(error);
return;
}
}
item = decode_stream.next() => match item {
Some(Ok(chunk)) => yield Ok(chunk),
Some(Err(error)) => {
yield Err(std::io::Error::other(format!("decode stream failed: {error}")));
return;
}
None => break,
}
}
}
if !prefill_done
&& let Err(error) = prefill_stream_outcome(prefill_task.await)
{
yield Err(error);
}
}
}
fn sse_decode_response_stream<S, E>(
decode_stream: S,
prefill_task: JoinHandle<Result<(), ProxyHttpError>>,
) -> impl Stream<Item = Result<Bytes, std::io::Error>>
where
S: Stream<Item = Result<Bytes, E>> + Unpin,
E: fmt::Display,
{
stream! {
let mut decode_stream = decode_stream;
let mut prefill_task = prefill_task;
let mut prefill_done = false;
let mut scanner = SseTerminalScanner::default();
loop {
if !prefill_done && prefill_task.is_finished() {
prefill_done = true;
let outcome = prefill_stream_outcome((&mut prefill_task).await);
if let Err(error) = outcome {
yield Err(error);
return;
}
}
tokio::select! {
prefill = &mut prefill_task, if !prefill_done => {
prefill_done = true;
let outcome = prefill_stream_outcome(prefill);
if let Err(error) = outcome {
yield Err(error);
return;
}
}
item = decode_stream.next() => match item {
Some(Ok(chunk)) => {
let scan = scanner.push(&chunk);
if let Some(safe) = scan.safe {
yield Ok(safe);
}
if let Some(terminal) = scan.terminal {
let mut held = terminal.to_vec();
loop {
if !prefill_done && prefill_task.is_finished() {
prefill_done = true;
let outcome = prefill_stream_outcome((&mut prefill_task).await);
if let Err(error) = outcome {
yield Err(error);
return;
}
}
tokio::select! {
prefill = &mut prefill_task, if !prefill_done => {
prefill_done = true;
let outcome = prefill_stream_outcome(prefill);
if let Err(error) = outcome {
yield Err(error);
return;
}
}
item = decode_stream.next() => match item {
Some(Ok(chunk)) => held.extend_from_slice(&chunk),
Some(Err(error)) => {
yield Err(std::io::Error::other(format!(
"decode stream failed: {error}"
)));
return;
}
None => break,
}
}
}
if !prefill_done {
let outcome = prefill_stream_outcome((&mut prefill_task).await);
if let Err(error) = outcome {
yield Err(error);
return;
}
}
yield Ok(Bytes::from(held));
return;
}
}
Some(Err(error)) => {
yield Err(std::io::Error::other(format!("decode stream failed: {error}")));
return;
}
None => {
let scan = scanner.finish();
if let Some(safe) = scan.safe {
yield Ok(safe);
}
if !prefill_done {
let outcome = prefill_stream_outcome((&mut prefill_task).await);
if let Err(error) = outcome {
yield Err(error);
return;
}
}
if let Some(terminal) = scan.terminal {
yield Ok(terminal);
}
return;
}
}
}
}
}
}
fn prefill_stream_outcome(
outcome: Result<Result<(), ProxyHttpError>, tokio::task::JoinError>,
) -> Result<(), std::io::Error> {
outcome
.map_err(|error| std::io::Error::other(format!("prefill task failed: {error}")))?
.map_err(|error| std::io::Error::other(error.to_string()))
}
#[derive(Default)]
struct SseTerminalScanner {
pending: Vec<u8>,
}
impl SseTerminalScanner {
fn push(&mut self, chunk: &[u8]) -> SseScan {
self.pending.extend_from_slice(chunk);
let mut event_start = 0;
while let Some(relative_end) = sse_event_end(&self.pending[event_start..]) {
let event_end = event_start + relative_end;
if is_terminal_sse_event(&self.pending[event_start..event_end]) {
let terminal = self.pending.split_off(event_start);
let safe = std::mem::take(&mut self.pending);
return SseScan::new(safe, terminal);
}
event_start = event_end;
}
if event_start == 0 {
return SseScan::default();
}
let incomplete = self.pending.split_off(event_start);
let safe = std::mem::replace(&mut self.pending, incomplete);
SseScan::safe(safe)
}
fn finish(&mut self) -> SseScan {
let pending = std::mem::take(&mut self.pending);
if pending.is_empty() {
SseScan::default()
} else if is_terminal_sse_event(&pending) {
SseScan::terminal(pending)
} else {
SseScan::safe(pending)
}
}
}
#[derive(Default)]
struct SseScan {
safe: Option<Bytes>,
terminal: Option<Bytes>,
}
impl SseScan {
fn new(safe: Vec<u8>, terminal: Vec<u8>) -> Self {
Self {
safe: (!safe.is_empty()).then(|| Bytes::from(safe)),
terminal: Some(Bytes::from(terminal)),
}
}
fn safe(bytes: Vec<u8>) -> Self {
Self {
safe: Some(Bytes::from(bytes)),
terminal: None,
}
}
fn terminal(bytes: Vec<u8>) -> Self {
Self {
safe: None,
terminal: Some(Bytes::from(bytes)),
}
}
}
fn sse_event_end(bytes: &[u8]) -> Option<usize> {
let mut line_start = 0;
for (index, byte) in bytes.iter().enumerate() {
if *byte != b'\n' {
continue;
}
let line_end = if index > line_start && bytes[index - 1] == b'\r' {
index - 1
} else {
index
};
if line_end == line_start {
return Some(index + 1);
}
line_start = index + 1;
}
None
}
fn is_terminal_sse_event(event: &[u8]) -> bool {
let mut data = Vec::new();
let mut saw_data = false;
for raw_line in event.split(|byte| *byte == b'\n') {
let line = raw_line.strip_suffix(b"\r").unwrap_or(raw_line);
let value = if line == b"data" {
Some(&b""[..])
} else if let Some(value) = line.strip_prefix(b"data:") {
Some(value.strip_prefix(b" ").unwrap_or(value))
} else {
None
};
if let Some(value) = value {
if saw_data {
data.push(b'\n');
}
data.extend_from_slice(value);
saw_data = true;
}
}
saw_data && data == b"[DONE]"
}
fn bootstrap_body(
body: &Value,
prefill: &PrefillTarget,
room: u64,
path: &'static str,
) -> Result<Value, ProxyHttpError> {
let mut body = body.clone();
let object = body.as_object_mut().ok_or_else(|| {
ProxyHttpError::status(
StatusCode::BAD_REQUEST,
"OpenAI completion request body must be a JSON object",
)
})?;
if path == "/v1/completions" && object.get("prompt").is_some_and(Value::is_array) {
return Err(ProxyHttpError::status(
StatusCode::BAD_REQUEST,
"SGLang built-in proxy does not support prompt arrays",
));
}
object.insert(
"bootstrap_host".to_owned(),
Value::String(prefill.bootstrap_host.clone()),
);
object.insert(
"bootstrap_port".to_owned(),
Value::from(prefill.bootstrap_port),
);
object.insert("bootstrap_room".to_owned(), Value::from(room));
Ok(body)
}
async fn drain_response(
response: reqwest::Response,
role: &'static str,
) -> Result<(), ProxyHttpError> {
response.bytes().await.map_err(|error| {
ProxyHttpError::upstream(&format!("{role} response drain failed"), error)
})?;
Ok(())
}
async fn await_prefill(task: JoinHandle<Result<(), ProxyHttpError>>) -> Result<(), ProxyHttpError> {
task.await
.map_err(|error| ProxyHttpError::internal(format!("prefill task failed: {error}")))?
}
#[derive(Serialize)]
struct FlushCacheResponse {
successful: Vec<String>,
failed: Vec<FlushCacheFailure>,
}
#[derive(Serialize)]
struct FlushCacheFailure {
url: String,
error: String,
}
async fn flush_cache(State(state): State<ProxyState>, headers: HeaderMap) -> Response<Body> {
let authorization = outbound_authorization(&headers);
let targets = state
.inner
.prefill
.iter()
.map(|target| target.url.clone())
.chain(state.inner.decode.iter().cloned());
let attempts = targets.map(|url| flush_target(state.client(), url, authorization.clone()));
let mut successful = Vec::new();
let mut failed = Vec::new();
for result in join_all(attempts).await {
match result {
Ok(url) => successful.push(url),
Err(failure) => failed.push(failure),
}
}
let status = if failed.is_empty() {
StatusCode::OK
} else {
StatusCode::PARTIAL_CONTENT
};
(status, Json(FlushCacheResponse { successful, failed })).into_response()
}
async fn flush_target(
client: reqwest::Client,
url: String,
authorization: Option<String>,
) -> Result<String, FlushCacheFailure> {
let endpoint = join_path(&url, "/flush_cache");
let mut request = client.post(endpoint);
if let Some(authorization) = authorization {
request = request.header(reqwest::header::AUTHORIZATION, authorization);
}
let response = request.send().await.map_err(|error| FlushCacheFailure {
url: url.clone(),
error: format!("cache flush request failed: {error}"),
})?;
if response.status().is_success() && response.status() != reqwest::StatusCode::PARTIAL_CONTENT {
Ok(url)
} else {
let status = response.status();
let detail = response
.text()
.await
.unwrap_or_else(|error| format!("failed to read response body: {error}"));
Err(FlushCacheFailure {
url,
error: format!("HTTP {status}: {detail}"),
})
}
}
#[cfg(test)]
mod tests {
use super::*;
use anyhow::{Context, Result, bail};
use axum::body::{Body, to_bytes};
use axum::extract::{Json, State};
use axum::http::{HeaderMap, HeaderValue, Response, StatusCode, header};
use axum::response::IntoResponse;
use axum::routing::{get, post};
use axum::{Router, serve};
use bytes::Bytes;
use futures_util::StreamExt;
use serde_json::{Value, json};
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicU16, AtomicUsize, Ordering};
use tokio::net::TcpListener;
use tokio::sync::{Mutex, Notify};
use tokio::task::JoinHandle;
#[derive(Clone)]
struct BackendState {
completion_requests: Arc<Mutex<Vec<Value>>>,
chat_requests: Arc<Mutex<Vec<Value>>>,
completion_status: Arc<AtomicU16>,
completion_content_type: &'static str,
completion_chunks: Vec<Bytes>,
notify_on_request: Option<Arc<Notify>>,
wait_before_response: Option<Arc<Notify>>,
gate_after_first_chunk: Option<Arc<Notify>>,
body_error: bool,
body_polled: Arc<AtomicBool>,
flush_status: Arc<AtomicU16>,
flush_requests: Arc<AtomicUsize>,
}
impl BackendState {
fn new(body: &'static [u8]) -> Self {
Self {
completion_requests: Arc::new(Mutex::new(Vec::new())),
chat_requests: Arc::new(Mutex::new(Vec::new())),
completion_status: Arc::new(AtomicU16::new(StatusCode::OK.as_u16())),
completion_content_type: "application/json",
completion_chunks: vec![Bytes::from_static(body)],
notify_on_request: None,
wait_before_response: None,
gate_after_first_chunk: None,
body_error: false,
body_polled: Arc::new(AtomicBool::new(false)),
flush_status: Arc::new(AtomicU16::new(StatusCode::OK.as_u16())),
flush_requests: Arc::new(AtomicUsize::new(0)),
}
}
}
async fn mock_completion(
State(state): State<BackendState>,
Json(body): Json<Value>,
) -> Response<Body> {
state.completion_requests.lock().await.push(body);
if let Some(notify) = &state.notify_on_request {
notify.notify_one();
}
if let Some(wait) = &state.wait_before_response {
wait.notified().await;
}
let chunks = state.completion_chunks.clone();
let body_polled = state.body_polled.clone();
let gate = state.gate_after_first_chunk.clone();
let stream = async_stream::stream! {
body_polled.store(true, Ordering::SeqCst);
for (index, chunk) in chunks.into_iter().enumerate() {
yield Ok::<Bytes, std::io::Error>(chunk);
if index == 0
&& let Some(gate) = &gate
{
gate.notified().await;
}
}
if state.body_error {
yield Err(std::io::Error::other("mock response body failed"));
}
};
let status = StatusCode::from_u16(state.completion_status.load(Ordering::SeqCst))
.unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
let mut response = Response::new(Body::from_stream(stream));
*response.status_mut() = status;
response.headers_mut().insert(
header::CONTENT_TYPE,
HeaderValue::from_static(state.completion_content_type),
);
response
}
async fn mock_chat_completion(
State(state): State<BackendState>,
Json(body): Json<Value>,
) -> Response<Body> {
state.chat_requests.lock().await.push(body);
(StatusCode::OK, Json(json!({"route": "chat"}))).into_response()
}
async fn mock_flush(State(state): State<BackendState>) -> Response<Body> {
state.flush_requests.fetch_add(1, Ordering::SeqCst);
let status = StatusCode::from_u16(state.flush_status.load(Ordering::SeqCst))
.unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
(status, "flush").into_response()
}
async fn spawn_backend(state: BackendState) -> Result<(String, JoinHandle<()>)> {
let app = Router::new()
.route(
"/health",
get(|| async { StatusCode::INTERNAL_SERVER_ERROR }),
)
.route("/v1/models", get(|| async { StatusCode::OK }))
.route("/v1/completions", post(mock_completion))
.route("/v1/chat/completions", post(mock_chat_completion))
.route("/flush_cache", post(mock_flush))
.with_state(state);
let listener = TcpListener::bind(("127.0.0.1", 0)).await?;
let address = listener.local_addr()?;
let handle = tokio::spawn(async move {
let _result = serve(listener, app).await;
});
Ok((format!("http://{address}"), handle))
}
fn proxy_state(prefill_url: String, decode_url: String) -> Result<ProxyState> {
ProxyState::new(Config {
host: "127.0.0.1".to_owned(),
port: 8000,
prefill: vec![PrefillTarget {
url: prefill_url,
bootstrap_host: "10.0.0.7".to_owned(),
bootstrap_port: 8998,
}],
decode: vec![decode_url],
})
.map_err(Into::into)
}
#[test]
fn meta_exports_proxy_identity() {
assert_eq!(meta().id, "inferlab-sglang-proxy");
assert_eq!(meta().version, 2);
}
#[tokio::test]
async fn non_streaming_completion_dispatches_both_roles_and_drains_prefill() -> Result<()> {
let decode_seen = Arc::new(Notify::new());
let mut prefill_backend = BackendState::new(br#"{"prefill":true}"#);
prefill_backend.wait_before_response = Some(decode_seen.clone());
let mut decode_backend = BackendState::new(br#"{"decode":true}"#);
decode_backend.notify_on_request = Some(decode_seen);
decode_backend
.completion_status
.store(StatusCode::CREATED.as_u16(), Ordering::SeqCst);
let (prefill_url, prefill_server) = spawn_backend(prefill_backend.clone()).await?;
let (decode_url, decode_server) = spawn_backend(decode_backend.clone()).await?;
let state = proxy_state(prefill_url, decode_url)?;
state.set_ready();
let response = tokio::time::timeout(
std::time::Duration::from_secs(2),
completion_route(
state,
HeaderMap::new(),
json!({"model": "m", "prompt": "hello"}),
),
)
.await
.context("prefill waited for decode instead of both requests being initiated")?
.map_err(|error| anyhow::anyhow!(error.to_string()))?;
assert_eq!(response.status(), StatusCode::CREATED);
assert_eq!(
response.headers().get(header::CONTENT_TYPE),
Some(&HeaderValue::from_static("application/json"))
);
assert_eq!(
to_bytes(response.into_body(), usize::MAX).await?,
Bytes::from_static(br#"{"decode":true}"#)
);
assert!(prefill_backend.body_polled.load(Ordering::SeqCst));
let prefill_requests = prefill_backend.completion_requests.lock().await;
let decode_requests = decode_backend.completion_requests.lock().await;
assert_eq!(prefill_requests.len(), 1);
assert_eq!(*prefill_requests, *decode_requests);
assert_eq!(
prefill_requests[0]
.get("bootstrap_host")
.and_then(Value::as_str),
Some("10.0.0.7")
);
assert_eq!(
prefill_requests[0]
.get("bootstrap_port")
.and_then(Value::as_u64),
Some(8998)
);
assert!(
prefill_requests[0]
.get("bootstrap_room")
.and_then(Value::as_u64)
.is_some()
);
prefill_server.abort();
decode_server.abort();
Ok(())
}
#[tokio::test]
async fn chat_dispatch_preserves_messages_and_unowned_fields_on_both_roles() -> Result<()> {
let prefill_backend = BackendState::new(b"prefill");
let decode_backend = BackendState::new(b"decode");
let (prefill_url, prefill_server) = spawn_backend(prefill_backend.clone()).await?;
let (decode_url, decode_server) = spawn_backend(decode_backend.clone()).await?;
let state = proxy_state(prefill_url, decode_url)?;
state.set_ready();
let request = json!({
"model": "m",
"messages": [{"role": "user", "content": "hello"}],
"temperature": 1.0,
"reasoning_effort": "high",
"chat_template_kwargs": {"enable_thinking": true}
});
let response = request_route(
state,
HeaderMap::new(),
request.clone(),
"/v1/chat/completions",
)
.await
.map_err(|error| anyhow::anyhow!(error.to_string()))?;
assert_eq!(response.status(), StatusCode::OK);
let prefill_requests = prefill_backend.chat_requests.lock().await;
let decode_requests = decode_backend.chat_requests.lock().await;
assert_eq!(prefill_requests.len(), 1);
assert_eq!(*prefill_requests, *decode_requests);
for (key, value) in request.as_object().context("request was not an object")? {
assert_eq!(prefill_requests[0].get(key), Some(value), "changed {key}");
}
assert!(prefill_requests[0]["bootstrap_room"].is_u64());
assert!(prefill_backend.completion_requests.lock().await.is_empty());
assert!(decode_backend.completion_requests.lock().await.is_empty());
prefill_server.abort();
decode_server.abort();
Ok(())
}
#[tokio::test]
async fn prompt_array_is_rejected_before_role_dispatch() -> Result<()> {
let prefill_backend = BackendState::new(b"prefill");
let decode_backend = BackendState::new(b"decode");
let (prefill_url, prefill_server) = spawn_backend(prefill_backend.clone()).await?;
let (decode_url, decode_server) = spawn_backend(decode_backend.clone()).await?;
let state = proxy_state(prefill_url, decode_url)?;
state.set_ready();
let result = completion_route(
state,
HeaderMap::new(),
json!({"model": "m", "prompt": ["one", "two"]}),
)
.await;
let error = match result {
Ok(_) => bail!("prompt arrays must fail"),
Err(error) => error,
};
assert_eq!(error.into_response().status(), StatusCode::BAD_REQUEST);
assert!(prefill_backend.completion_requests.lock().await.is_empty());
assert!(decode_backend.completion_requests.lock().await.is_empty());
prefill_server.abort();
decode_server.abort();
Ok(())
}
#[tokio::test]
async fn streaming_completion_relays_decode_chunks_incrementally() -> Result<()> {
let release_prefill = Arc::new(Notify::new());
let mut prefill_backend = BackendState::new(b"prefill");
prefill_backend.wait_before_response = Some(release_prefill.clone());
let second_chunk = Arc::new(Notify::new());
let mut decode_backend = BackendState::new(b"");
decode_backend.completion_content_type = "text/event-stream";
decode_backend.completion_chunks = vec![
Bytes::from_static(b"data: first\n\n"),
Bytes::from_static(b"data: second\n\n"),
Bytes::from_static(b"data: [DO"),
Bytes::from_static(b"NE]\r"),
Bytes::from_static(b"\n\r"),
Bytes::from_static(b"\n"),
];
decode_backend.gate_after_first_chunk = Some(second_chunk.clone());
let (prefill_url, prefill_server) = spawn_backend(prefill_backend).await?;
let (decode_url, decode_server) = spawn_backend(decode_backend).await?;
let state = proxy_state(prefill_url, decode_url)?;
state.set_ready();
let response = tokio::time::timeout(
std::time::Duration::from_secs(1),
completion_route(
state,
HeaderMap::new(),
json!({"model": "m", "prompt": "hello", "stream": true}),
),
)
.await
.context("a delayed prefill response blocked decode streaming")?
.map_err(|error| anyhow::anyhow!(error.to_string()))?;
assert_eq!(
response.headers().get(header::CONTENT_TYPE),
Some(&HeaderValue::from_static("text/event-stream"))
);
let mut stream = response.into_body().into_data_stream();
let first = tokio::time::timeout(std::time::Duration::from_secs(1), stream.next())
.await?
.context("decode stream ended before its first chunk")??;
assert_eq!(first, Bytes::from_static(b"data: first\n\n"));
release_prefill.notify_one();
second_chunk.notify_one();
let second = stream
.next()
.await
.context("decode stream ended before its second chunk")??;
assert_eq!(second, Bytes::from_static(b"data: second\n\n"));
let terminal = stream
.next()
.await
.context("decode stream ended before its terminal event")??;
assert_eq!(terminal, Bytes::from_static(b"data: [DONE]\r\n\r\n"));
prefill_server.abort();
decode_server.abort();
Ok(())
}
#[tokio::test]
async fn late_prefill_failure_prevents_terminal_sse_event() -> Result<()> {
let release_prefill_failure = Arc::new(Notify::new());
let mut prefill_backend = BackendState::new(b"prefill");
prefill_backend.body_error = true;
prefill_backend.gate_after_first_chunk = Some(release_prefill_failure.clone());
let release_terminal = Arc::new(Notify::new());
let mut decode_backend = BackendState::new(b"");
decode_backend.completion_content_type = "text/event-stream; charset=utf-8";
decode_backend.completion_chunks = vec![
Bytes::from_static(b"data: first\r\n\r\n"),
Bytes::from_static(b"data: [DO"),
Bytes::from_static(b"NE]\r"),
Bytes::from_static(b"\n\r"),
Bytes::from_static(b"\n"),
];
decode_backend.gate_after_first_chunk = Some(release_terminal.clone());
let (prefill_url, prefill_server) = spawn_backend(prefill_backend).await?;
let (decode_url, decode_server) = spawn_backend(decode_backend).await?;
let state = proxy_state(prefill_url, decode_url)?;
state.set_ready();
let response = completion_route(
state,
HeaderMap::new(),
json!({"model": "m", "prompt": "hello", "stream": true}),
)
.await
.map_err(|error| anyhow::anyhow!(error.to_string()))?;
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(
response.headers().get(header::CONTENT_TYPE),
Some(&HeaderValue::from_static(
"text/event-stream; charset=utf-8"
))
);
let mut stream = response.into_body().into_data_stream();
let first = stream
.next()
.await
.context("decode stream ended before its first event")??;
assert_eq!(first, Bytes::from_static(b"data: first\r\n\r\n"));
release_terminal.notify_one();
assert!(
tokio::time::timeout(Duration::from_millis(50), stream.next())
.await
.is_err(),
"terminal SSE bytes were forwarded before prefill completed"
);
release_prefill_failure.notify_one();
let result = stream
.next()
.await
.context("stream completed after a late prefill failure")?;
let error = match result {
Ok(bytes) => bail!(
"late prefill failure forwarded terminal bytes: {:?}",
String::from_utf8_lossy(&bytes)
),
Err(error) => error,
};
assert!(error.to_string().contains("prefill response drain failed"));
prefill_server.abort();
decode_server.abort();
Ok(())
}
#[tokio::test]
async fn prefill_failure_before_headers_returns_non_success() -> Result<()> {
let prefill_backend = BackendState::new(b"prefill failed");
prefill_backend
.completion_status
.store(StatusCode::INTERNAL_SERVER_ERROR.as_u16(), Ordering::SeqCst);
let release_decode = Arc::new(Notify::new());
let mut decode_backend = BackendState::new(b"decode");
decode_backend.wait_before_response = Some(release_decode.clone());
let (prefill_url, prefill_server) = spawn_backend(prefill_backend.clone()).await?;
let (decode_url, decode_server) = spawn_backend(decode_backend.clone()).await?;
let state = proxy_state(prefill_url, decode_url)?;
state.set_ready();
let request = tokio::spawn(completion_route(
state,
HeaderMap::new(),
json!({"model": "m", "prompt": "hello", "stream": true}),
));
wait_until(&prefill_backend.body_polled).await?;
tokio::task::yield_now().await;
release_decode.notify_one();
let error = match request.await? {
Ok(_) => bail!("prefill failure must fail before public headers"),
Err(error) => error,
};
assert_eq!(error.into_response().status(), StatusCode::BAD_GATEWAY);
assert_eq!(decode_backend.completion_requests.lock().await.len(), 1);
prefill_server.abort();
decode_server.abort();
Ok(())
}
#[tokio::test]
async fn prefill_body_failure_before_headers_returns_non_success() -> Result<()> {
let mut prefill_backend = BackendState::new(b"prefill");
prefill_backend.body_error = true;
let release_decode = Arc::new(Notify::new());
let mut decode_backend = BackendState::new(b"decode");
decode_backend.wait_before_response = Some(release_decode.clone());
let (prefill_url, prefill_server) = spawn_backend(prefill_backend.clone()).await?;
let (decode_url, decode_server) = spawn_backend(decode_backend).await?;
let state = proxy_state(prefill_url, decode_url)?;
state.set_ready();
let request = tokio::spawn(completion_route(
state,
HeaderMap::new(),
json!({"model": "m", "prompt": "hello", "stream": true}),
));
wait_until(&prefill_backend.body_polled).await?;
tokio::task::yield_now().await;
release_decode.notify_one();
let error = match request.await? {
Ok(_) => bail!("prefill body failure must fail before public headers"),
Err(error) => error,
};
assert_eq!(error.into_response().status(), StatusCode::BAD_GATEWAY);
prefill_server.abort();
decode_server.abort();
Ok(())
}
#[tokio::test]
async fn decode_failure_does_not_wait_for_slow_prefill() -> Result<()> {
let prefill_seen = Arc::new(Notify::new());
let release_prefill = Arc::new(Notify::new());
let mut prefill_backend = BackendState::new(b"prefill");
prefill_backend.notify_on_request = Some(prefill_seen.clone());
prefill_backend.wait_before_response = Some(release_prefill.clone());
let mut decode_backend = BackendState::new(b"decode failed");
decode_backend.wait_before_response = Some(prefill_seen);
decode_backend
.completion_status
.store(StatusCode::INTERNAL_SERVER_ERROR.as_u16(), Ordering::SeqCst);
let (prefill_url, prefill_server) = spawn_backend(prefill_backend.clone()).await?;
let (decode_url, decode_server) = spawn_backend(decode_backend).await?;
let state = proxy_state(prefill_url, decode_url)?;
state.set_ready();
let result = tokio::time::timeout(
std::time::Duration::from_secs(1),
completion_route(
state,
HeaderMap::new(),
json!({"model": "m", "prompt": "hello", "stream": true}),
),
)
.await
.context("decode failure waited for a slow prefill response")?;
let error = match result {
Ok(_) => bail!("decode failure must return a non-success response"),
Err(error) => error,
};
assert_eq!(error.into_response().status(), StatusCode::BAD_GATEWAY);
assert_eq!(prefill_backend.completion_requests.lock().await.len(), 1);
release_prefill.notify_one();
prefill_server.abort();
decode_server.abort();
Ok(())
}
#[tokio::test]
async fn dropping_public_stream_detaches_prefill_drain() -> Result<()> {
let release_prefill = Arc::new(Notify::new());
let prefill_drained = Arc::new(AtomicBool::new(false));
let drained = prefill_drained.clone();
let release = release_prefill.clone();
let prefill = tokio::spawn(async move {
release.notified().await;
drained.store(true, Ordering::SeqCst);
Ok::<(), ProxyHttpError>(())
});
let decode = Box::pin(
futures_util::stream::once(async {
Ok::<Bytes, std::io::Error>(Bytes::from_static(b"data: first\n\n"))
})
.chain(futures_util::stream::pending()),
);
let mut stream = Box::pin(sse_decode_response_stream(decode, prefill));
assert!(matches!(stream.next().await, Some(Ok(_))));
drop(stream);
release_prefill.notify_one();
wait_until(&prefill_drained).await?;
Ok(())
}
#[tokio::test]
async fn decode_failure_after_streaming_starts_fails_the_public_stream() -> Result<()> {
let prefill_backend = BackendState::new(b"prefill");
let release_error = Arc::new(Notify::new());
let mut decode_backend = BackendState::new(b"");
decode_backend.completion_content_type = "text/event-stream";
decode_backend.completion_chunks = vec![
Bytes::from_static(b"data: first\n\n"),
Bytes::from_static(b"data: [DONE]\n\n"),
];
decode_backend.body_error = true;
decode_backend.gate_after_first_chunk = Some(release_error.clone());
let (prefill_url, prefill_server) = spawn_backend(prefill_backend).await?;
let (decode_url, decode_server) = spawn_backend(decode_backend).await?;
let state = proxy_state(prefill_url, decode_url)?;
state.set_ready();
let response = completion_route(
state,
HeaderMap::new(),
json!({"model": "m", "prompt": "hello", "stream": true}),
)
.await
.map_err(|error| anyhow::anyhow!(error.to_string()))?;
let mut stream = response.into_body().into_data_stream();
assert!(matches!(stream.next().await, Some(Ok(_))));
release_error.notify_one();
let result = stream
.next()
.await
.context("decode stream ended successfully after an upstream body failure")?;
let error = match result {
Ok(_) => bail!("decode body failure must fail the public stream"),
Err(error) => error,
};
assert!(error.to_string().contains("decode stream failed"));
prefill_server.abort();
decode_server.abort();
Ok(())
}
#[tokio::test]
async fn terminal_sse_waits_for_clean_decode_eof() -> Result<()> {
let decode = Box::pin(futures_util::stream::iter(vec![
Ok::<Bytes, std::io::Error>(Bytes::from_static(b"data: first\n\n")),
Ok(Bytes::from_static(b"data: [DONE]\n\n")),
Err(std::io::Error::other("decode failed after terminal event")),
]));
let prefill = tokio::spawn(async { Ok::<(), ProxyHttpError>(()) });
let mut stream = Box::pin(sse_decode_response_stream(decode, prefill));
let first = stream
.next()
.await
.context("stream ended before its first event")??;
assert_eq!(first, Bytes::from_static(b"data: first\n\n"));
let result = stream
.next()
.await
.context("stream completed after a late decode failure")?;
let error = match result {
Ok(bytes) => bail!(
"decode failure forwarded terminal bytes: {:?}",
String::from_utf8_lossy(&bytes)
),
Err(error) => error,
};
assert!(error.to_string().contains("decode stream failed"));
Ok(())
}
async fn wait_until(flag: &AtomicBool) -> Result<()> {
tokio::time::timeout(std::time::Duration::from_secs(1), async {
while !flag.load(Ordering::SeqCst) {
tokio::task::yield_now().await;
}
})
.await
.context("expected asynchronous condition was not observed")?;
Ok(())
}
#[tokio::test]
async fn flush_cache_attempts_all_targets_and_reports_partial_failure() -> Result<()> {
let prefill_backend = BackendState::new(b"prefill");
let decode_backend = BackendState::new(b"decode");
let (prefill_url, prefill_server) = spawn_backend(prefill_backend.clone()).await?;
let (decode_url, decode_server) = spawn_backend(decode_backend.clone()).await?;
let state = proxy_state(prefill_url, decode_url)?;
let all_succeeded = flush_cache(State(state.clone()), HeaderMap::new()).await;
assert_eq!(all_succeeded.status(), StatusCode::OK);
decode_backend
.flush_status
.store(StatusCode::PARTIAL_CONTENT.as_u16(), Ordering::SeqCst);
let partial = flush_cache(State(state), HeaderMap::new()).await;
assert_eq!(partial.status(), StatusCode::PARTIAL_CONTENT);
let body: Value =
serde_json::from_slice(&to_bytes(partial.into_body(), usize::MAX).await?)?;
assert_eq!(body["successful"].as_array().map(Vec::len), Some(1));
assert_eq!(body["failed"].as_array().map(Vec::len), Some(1));
assert_eq!(prefill_backend.flush_requests.load(Ordering::SeqCst), 2);
assert_eq!(decode_backend.flush_requests.load(Ordering::SeqCst), 2);
prefill_server.abort();
decode_server.abort();
Ok(())
}
#[tokio::test]
async fn healthcheck_is_unsuccessful_until_all_backends_are_ready() -> Result<()> {
let prefill_backend = BackendState::new(b"prefill");
let decode_backend = BackendState::new(b"decode");
let (prefill_url, prefill_server) = spawn_backend(prefill_backend).await?;
let (decode_url, decode_server) = spawn_backend(decode_backend).await?;
let state = proxy_state(prefill_url, decode_url)?;
let (status, Json(body)) = healthcheck(State(state.clone())).await;
assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
assert!(!body.ready);
tokio::time::timeout(
Duration::from_secs(1),
tokio::spawn(await_backends(state.clone())),
)
.await
.context("backend readiness did not use the responsive model endpoint")??;
let (status, Json(body)) = healthcheck(State(state)).await;
assert_eq!(status, StatusCode::OK);
assert!(body.ready);
prefill_server.abort();
decode_server.abort();
Ok(())
}
}