From 06518315d11aad207e393e2f9c33bbda0ed6119e Mon Sep 17 00:00:00 2001 From: Jonas Krauss Date: Tue, 17 Jun 2025 14:22:49 +0200 Subject: [PATCH 01/18] wip on adding model routing --- src/agent/monitoring_service.rs | 34 +++++++++++---- src/balancer/proxy_service.rs | 32 ++++++++++++++ src/balancer/request_context.rs | 4 +- src/balancer/status_update.rs | 3 ++ src/balancer/test/mock_status_update.rs | 1 + src/balancer/upstream_peer.rs | 10 +++++ src/balancer/upstream_peer_pool.rs | 14 +++++-- src/cmd/agent.rs | 4 +- src/cmd/balancer.rs | 5 +++ src/llamacpp/llamacpp_client.rs | 55 ++++++++++++++++++++++--- src/llamacpp/mod.rs | 1 + src/llamacpp/models_response.rs | 13 ++++++ src/main.rs | 12 ++++++ 13 files changed, 167 insertions(+), 21 deletions(-) create mode 100644 src/llamacpp/models_response.rs diff --git a/src/agent/monitoring_service.rs b/src/agent/monitoring_service.rs index 598e0ca8..bcc259a0 100644 --- a/src/agent/monitoring_service.rs +++ b/src/agent/monitoring_service.rs @@ -3,6 +3,7 @@ use std::net::SocketAddr; use actix_web::web::Bytes; use async_trait::async_trait; use log::debug; +use log::info; use log::error; #[cfg(unix)] use pingora::server::ListenFds; @@ -23,6 +24,7 @@ pub struct MonitoringService { monitoring_interval: Duration, name: Option, status_update_tx: Sender, + check_model: bool, // Store the check_model flag } impl MonitoringService { @@ -32,6 +34,7 @@ impl MonitoringService { monitoring_interval: Duration, name: Option, status_update_tx: Sender, + check_model: bool, // Include the check_model flag ) -> Result { Ok(MonitoringService { external_llamacpp_addr, @@ -39,19 +42,31 @@ impl MonitoringService { monitoring_interval, name, status_update_tx, + check_model, }) } async fn fetch_status(&self) -> Result { match self.llamacpp_client.get_available_slots().await { - Ok(slots_response) => Ok(StatusUpdate::new( - self.name.to_owned(), - None, - self.external_llamacpp_addr.to_owned(), - slots_response.is_authorized, - slots_response.is_slot_endpoint_enabled, - slots_response.slots, - )), + Ok(slots_response) => { + let model = if self.check_model { + self.llamacpp_client.get_model().await? + } else { + None + }; + + info!("Agent: {:?} Model: {:?}", self.name, model); + + Ok(StatusUpdate::new( + self.name.to_owned(), + None, + self.external_llamacpp_addr.to_owned(), + slots_response.is_authorized, + slots_response.is_slot_endpoint_enabled, + slots_response.slots, + model, + )) + }, Err(err) => Ok(StatusUpdate::new( self.name.to_owned(), Some(err.to_string()), @@ -59,6 +74,7 @@ impl MonitoringService { None, None, vec![], + None, )), } } @@ -111,4 +127,4 @@ impl Service for MonitoringService { fn threads(&self) -> Option { Some(1) } -} +} \ No newline at end of file diff --git a/src/balancer/proxy_service.rs b/src/balancer/proxy_service.rs index 41f11e77..8ec1157c 100644 --- a/src/balancer/proxy_service.rs +++ b/src/balancer/proxy_service.rs @@ -4,6 +4,7 @@ use std::time::Duration; use async_trait::async_trait; use bytes::Bytes; use log::error; +use log::info; use pingora::http::RequestHeader; use pingora::proxy::ProxyHttp; use pingora::proxy::Session; @@ -16,6 +17,7 @@ use crate::balancer::upstream_peer_pool::UpstreamPeerPool; pub struct ProxyService { rewrite_host_header: bool, + check_model: bool, slots_endpoint_enable: bool, upstream_peer_pool: Arc, } @@ -23,11 +25,13 @@ pub struct ProxyService { impl ProxyService { pub fn new( rewrite_host_header: bool, + check_model: bool, slots_endpoint_enable: bool, upstream_peer_pool: Arc, ) -> Self { Self { rewrite_host_header, + check_model, slots_endpoint_enable, upstream_peer_pool, } @@ -44,6 +48,7 @@ impl ProxyHttp for ProxyService { slot_taken: false, upstream_peer_pool: self.upstream_peer_pool.clone(), uses_slots: false, + requested_model: Some("".to_string()), } } @@ -135,6 +140,33 @@ impl ProxyHttp for ProxyService { session: &mut Session, ctx: &mut Self::CTX, ) -> Result> { + info!("upstream_peer - {:?} request | rewrite_host_header? {} check_model? {}", session.req_header().method, self.rewrite_host_header, self.check_model); + + // Check if the request method is POST and the content type is JSON + if self.check_model { + if session.req_header().method == "POST" { + // Check if the content type is application/json + if let Some(content_type) = session.get_header("Content-Type") { + if let Ok(content_type_str) = content_type.to_str() { + if content_type_str.contains("application/json") { + // Read the request body + let body = session.read_request_body().await?; + if let Some(body) = body { + // Parse the JSON payload into a serde_json::Value + if let Ok(json_value) = serde_json::from_slice::(&body) { + if let Some(model) = json_value.get("model").and_then(|v| v.as_str()) { + // Set the requested_model field in the RequestContext + ctx.requested_model = Some(model.to_string()); + info!("Model in request: {:?}", ctx.requested_model); + } + } + } + } + } + } + } + } + let upstream_peer = ctx.select_upstream_peer(session.req_header().uri.path(), self.slots_endpoint_enable); diff --git a/src/balancer/request_context.rs b/src/balancer/request_context.rs index 45acb197..e64a04b1 100644 --- a/src/balancer/request_context.rs +++ b/src/balancer/request_context.rs @@ -15,6 +15,7 @@ pub struct RequestContext { pub selected_peer: Option, pub upstream_peer_pool: Arc, pub uses_slots: bool, + pub requested_model: Option, } impl RequestContext { @@ -51,7 +52,7 @@ impl RequestContext { slots_endpoint_enable: bool, ) -> Result> { if self.selected_peer.is_none() { - self.selected_peer = match self.upstream_peer_pool.use_best_peer() { + self.selected_peer = match self.upstream_peer_pool.use_best_peer(self.requested_model.clone()) { Ok(peer) => peer, Err(err) => { error!("Failed to get best peer: {err}"); @@ -115,6 +116,7 @@ mod tests { selected_peer: None, upstream_peer_pool, uses_slots: true, + requested_model: Some("llama3".to_string()), } } diff --git a/src/balancer/status_update.rs b/src/balancer/status_update.rs index 46da8fa7..ab3903dd 100644 --- a/src/balancer/status_update.rs +++ b/src/balancer/status_update.rs @@ -14,6 +14,7 @@ pub struct StatusUpdate { pub is_authorized: Option, pub is_slots_endpoint_enabled: Option, pub processing_slots_count: usize, + pub model: Option, slots: Vec, } @@ -25,6 +26,7 @@ impl StatusUpdate { is_authorized: Option, is_slots_endpoint_enabled: Option, slots: Vec, + model: Option, ) -> Self { let idle_slots_count = slots.iter().filter(|slot| !slot.is_processing).count(); @@ -37,6 +39,7 @@ impl StatusUpdate { is_slots_endpoint_enabled, processing_slots_count: slots.len() - idle_slots_count, slots, + model, } } } diff --git a/src/balancer/test/mock_status_update.rs b/src/balancer/test/mock_status_update.rs index b87e5744..d1b53484 100644 --- a/src/balancer/test/mock_status_update.rs +++ b/src/balancer/test/mock_status_update.rs @@ -38,5 +38,6 @@ pub fn mock_status_update( Some(true), Some(true), slots, + Some("llama3".to_string()), ) } diff --git a/src/balancer/upstream_peer.rs b/src/balancer/upstream_peer.rs index b5c0cd0e..35b15e42 100644 --- a/src/balancer/upstream_peer.rs +++ b/src/balancer/upstream_peer.rs @@ -14,6 +14,7 @@ use crate::errors::result::Result; pub struct UpstreamPeer { pub agent_id: String, pub agent_name: Option, + pub model: Option, pub error: Option, pub external_llamacpp_addr: SocketAddr, /// None means undetermined, probably due to an error @@ -39,6 +40,7 @@ impl UpstreamPeer { is_slots_endpoint_enabled: Option, slots_idle: usize, slots_processing: usize, + model: Option, ) -> Self { UpstreamPeer { agent_id, @@ -53,6 +55,7 @@ impl UpstreamPeer { slots_processing, slots_taken: 0, slots_taken_since_last_status_update: 0, + model, } } @@ -66,6 +69,7 @@ impl UpstreamPeer { status_update.is_slots_endpoint_enabled, status_update.idle_slots_count, status_update.processing_slots_count, + status_update.model, ) } @@ -76,6 +80,10 @@ impl UpstreamPeer { && matches!(self.is_authorized, Some(true)) } + pub fn is_usable_for_model(&self, requested_model: &str) -> bool { + self.is_usable() && (requested_model.is_empty() || self.model.as_deref() == Some(requested_model)) + } + pub fn release_slot(&mut self) -> Result<()> { if self.slots_taken < 1 { return Err("Cannot release a slot when there are no taken slots".into()); @@ -166,6 +174,7 @@ mod tests { Some(true), 5, 0, + Some("llama3".to_string()), ) } @@ -219,6 +228,7 @@ mod tests { Some(true), Some(true), vec![], + Some("llama3".to_string()), ); peer.update_status(status_update); diff --git a/src/balancer/upstream_peer_pool.rs b/src/balancer/upstream_peer_pool.rs index 4aa4beb3..6d852cc2 100644 --- a/src/balancer/upstream_peer_pool.rs +++ b/src/balancer/upstream_peer_pool.rs @@ -1,6 +1,7 @@ use std::sync::RwLock; use std::time::Duration; use std::time::SystemTime; +use log::info; use serde::Deserialize; use serde::Serialize; @@ -117,14 +118,19 @@ impl UpstreamPeerPool { }) } - pub fn use_best_peer(&self) -> Result> { + pub fn use_best_peer(&self, model: Option) -> Result> { self.with_agents_write(|agents| { for peer in agents.iter() { - if peer.is_usable() { + let model_str = model.as_deref().unwrap_or(""); + let is_usable = peer.is_usable(); + let is_usable_for_model = peer.is_usable_for_model(model_str); + + info!("Peer {} is usable: {}, usable for model '{}': {}", peer.agent_id, is_usable, model_str, is_usable_for_model); + + if is_usable && (model.is_none() || is_usable_for_model) { return Ok(Some(peer.clone())); } } - Ok(None) }) } @@ -227,7 +233,7 @@ mod tests { pool.register_status_update("test2", mock_status_update("test2", 3, 0))?; pool.register_status_update("test3", mock_status_update("test3", 0, 0))?; - let best_peer = pool.use_best_peer()?.unwrap(); + let best_peer = pool.use_best_peer(None)?.unwrap(); assert_eq!(best_peer.agent_id, "test1"); assert_eq!(best_peer.slots_idle, 5); diff --git a/src/cmd/agent.rs b/src/cmd/agent.rs index b6eedb87..b9c5dd38 100644 --- a/src/cmd/agent.rs +++ b/src/cmd/agent.rs @@ -18,6 +18,7 @@ pub fn handle( management_addr: SocketAddr, monitoring_interval: Duration, name: Option, + check_model: bool, // Include the check_model flag ) -> Result<()> { let (status_update_tx, _status_update_rx) = channel::(1); @@ -29,6 +30,7 @@ pub fn handle( monitoring_interval, name, status_update_tx.clone(), + check_model, // Pass the check_model flag )?; let reporting_service = ReportingService::new(management_addr, status_update_tx)?; @@ -45,4 +47,4 @@ pub fn handle( pingora_server.add_service(monitoring_service); pingora_server.add_service(reporting_service); pingora_server.run_forever(); -} +} \ No newline at end of file diff --git a/src/cmd/balancer.rs b/src/cmd/balancer.rs index 695778db..1365b9b5 100644 --- a/src/cmd/balancer.rs +++ b/src/cmd/balancer.rs @@ -2,6 +2,7 @@ use std::net::SocketAddr; use std::sync::Arc; #[cfg(feature = "statsd_reporter")] use std::time::Duration; +use log::info; use pingora::proxy::http_proxy_service; use pingora::server::configuration::Opt; @@ -19,6 +20,7 @@ pub fn handle( #[cfg(feature = "web_dashboard")] management_dashboard_enable: bool, reverseproxy_addr: &SocketAddr, rewrite_host_header: bool, + check_model: bool, slots_endpoint_enable: bool, #[cfg(feature = "statsd_reporter")] statsd_addr: Option, #[cfg(feature = "statsd_reporter")] statsd_prefix: String, @@ -40,6 +42,7 @@ pub fn handle( &pingora_server.configuration, ProxyService::new( rewrite_host_header, + check_model, slots_endpoint_enable, upstream_peer_pool.clone(), ), @@ -67,5 +70,7 @@ pub fn handle( pingora_server.add_service(statsd_service); } + info!("rewrite_host_header? {} check_model? {} slots_endpoint_enable? {}", rewrite_host_header, check_model, slots_endpoint_enable); + pingora_server.run_forever(); } diff --git a/src/llamacpp/llamacpp_client.rs b/src/llamacpp/llamacpp_client.rs index 48514512..6df9e6e3 100644 --- a/src/llamacpp/llamacpp_client.rs +++ b/src/llamacpp/llamacpp_client.rs @@ -6,12 +6,14 @@ use reqwest::header; use url::Url; use crate::errors::result::Result; -use crate::llamacpp::slot::Slot; use crate::llamacpp::slots_response::SlotsResponse; +use crate::llamacpp::slot::Slot; +use crate::llamacpp::models_response::ModelsResponse; pub struct LlamacppClient { client: reqwest::Client, slots_endpoint_url: String, + models_endpoint_url: String, } impl LlamacppClient { @@ -37,6 +39,7 @@ impl LlamacppClient { Ok(Self { client: builder.build()?, slots_endpoint_url: Url::parse(&format!("http://{addr}/slots"))?.to_string(), + models_endpoint_url: Url::parse(&format!("http://{addr}/v1/models"))?.to_string(), }) } @@ -61,11 +64,14 @@ impl LlamacppClient { }; match response.status() { - reqwest::StatusCode::OK => Ok(SlotsResponse { - is_authorized: Some(true), - is_slot_endpoint_enabled: Some(true), - slots: response.json::>().await?, - }), + reqwest::StatusCode::OK => { + let slots: Vec = response.json().await?; + Ok(SlotsResponse { + is_authorized: Some(true), + is_slot_endpoint_enabled: Some(true), + slots, + }) + }, reqwest::StatusCode::UNAUTHORIZED => Ok(SlotsResponse { is_authorized: Some(false), is_slot_endpoint_enabled: None, @@ -79,4 +85,41 @@ impl LlamacppClient { _ => Err("Unexpected response status".into()), } } + + pub async fn get_model(&self) -> Result> { + let url = self.models_endpoint_url.to_owned(); + + let response = match self.client.get(url.clone()).send().await { + Ok(resp) => resp, + Err(err) => { + return Err(format!( + "Request to '{}' failed: '{}'; connect issue: {}; decode issue: {}; request issue: {}; status issue: {}; status: {:?}; source: {:?}", + url, + err, + err.is_connect(), + err.is_decode(), + err.is_request(), + err.is_status(), + err.status(), + err.source() + ).into()); + } + }; + + match response.status() { + reqwest::StatusCode::OK => { + let models_response: ModelsResponse = response.json().await?; + if let Some(models) = models_response.models { + if models.is_empty() { + Ok(None) + } else { + Ok(models.first().and_then(|m| Some(m.model.clone()))) + } + } else { + Ok(None) + } + }, + _ => Err("Unexpected response status".into()), + } + } } diff --git a/src/llamacpp/mod.rs b/src/llamacpp/mod.rs index 2b1e1d24..68f5cae8 100644 --- a/src/llamacpp/mod.rs +++ b/src/llamacpp/mod.rs @@ -1,3 +1,4 @@ pub mod llamacpp_client; pub mod slot; pub mod slots_response; +pub mod models_response; diff --git a/src/llamacpp/models_response.rs b/src/llamacpp/models_response.rs new file mode 100644 index 00000000..d7d04b50 --- /dev/null +++ b/src/llamacpp/models_response.rs @@ -0,0 +1,13 @@ +// paddler/src/llamacpp/models_response.rs +use serde::Deserialize; + +#[derive(Debug, Deserialize)] +pub struct ModelsResponse { + pub models: Option>, +} + +#[derive(Debug, Deserialize)] +pub struct Model { + pub model: String, + // Add other fields as needed +} diff --git a/src/main.rs b/src/main.rs index 93999f39..f27fe23d 100644 --- a/src/main.rs +++ b/src/main.rs @@ -80,6 +80,10 @@ enum Commands { #[arg(long)] /// Name of the agent (optional) name: Option, + + #[arg(long)] + /// Flag whether to check the model served by llama.cpp and reject requests for other models + check_model: bool, }, /// Balances incoming requests to llama.cpp instances and optionally provides a web dashboard Balancer { @@ -105,6 +109,10 @@ enum Commands { /// Enable the slots endpoint (not recommended) slots_endpoint_enable: bool, + #[arg(long)] + /// Flag to check the model served by llama.cpp and reject requests for other models + check_model: bool, + #[cfg(feature = "statsd_reporter")] #[arg(long, value_parser = parse_socket_addr)] /// Address of the statsd server to report metrics to @@ -142,6 +150,7 @@ fn main() -> Result<()> { management_addr, monitoring_interval, name, + check_model, }) => cmd::agent::handle( match external_llamacpp_addr { Some(addr) => addr.to_owned(), @@ -152,6 +161,7 @@ fn main() -> Result<()> { management_addr.to_owned(), monitoring_interval.to_owned(), name.to_owned(), + *check_model ), Some(Commands::Balancer { management_addr, @@ -159,6 +169,7 @@ fn main() -> Result<()> { management_dashboard_enable, reverseproxy_addr, rewrite_host_header, + check_model, slots_endpoint_enable, #[cfg(feature = "statsd_reporter")] statsd_addr, @@ -172,6 +183,7 @@ fn main() -> Result<()> { management_dashboard_enable.to_owned(), reverseproxy_addr, rewrite_host_header.to_owned(), + *check_model, slots_endpoint_enable.to_owned(), #[cfg(feature = "statsd_reporter")] statsd_addr.to_owned(), From 84a88f8af9c2d6efcc4e61a88537e81cebce3b54 Mon Sep 17 00:00:00 2001 From: Jonas Krauss Date: Tue, 17 Jun 2025 14:30:52 +0200 Subject: [PATCH 02/18] revert unnecessary change --- src/llamacpp/llamacpp_client.rs | 13 +++++-------- 1 file changed, 5 insertions(+), 8 deletions(-) diff --git a/src/llamacpp/llamacpp_client.rs b/src/llamacpp/llamacpp_client.rs index 6df9e6e3..cc9e7e10 100644 --- a/src/llamacpp/llamacpp_client.rs +++ b/src/llamacpp/llamacpp_client.rs @@ -64,14 +64,11 @@ impl LlamacppClient { }; match response.status() { - reqwest::StatusCode::OK => { - let slots: Vec = response.json().await?; - Ok(SlotsResponse { - is_authorized: Some(true), - is_slot_endpoint_enabled: Some(true), - slots, - }) - }, + reqwest::StatusCode::OK => Ok(SlotsResponse { + is_authorized: Some(true), + is_slot_endpoint_enabled: Some(true), + slots: response.json::>().await?, + }), reqwest::StatusCode::UNAUTHORIZED => Ok(SlotsResponse { is_authorized: Some(false), is_slot_endpoint_enabled: None, From 4b4e3f6f6393fd010ccdefcb782b90c5c22aeeeb Mon Sep 17 00:00:00 2001 From: Jonas Krauss Date: Tue, 17 Jun 2025 16:31:07 +0200 Subject: [PATCH 03/18] avoid consuming the request body for upstream --- src/balancer/proxy_service.rs | 17 ++++++++++++----- src/balancer/upstream_peer_pool.rs | 4 ++-- 2 files changed, 14 insertions(+), 7 deletions(-) diff --git a/src/balancer/proxy_service.rs b/src/balancer/proxy_service.rs index 8ec1157c..987f1c02 100644 --- a/src/balancer/proxy_service.rs +++ b/src/balancer/proxy_service.rs @@ -149,17 +149,24 @@ impl ProxyHttp for ProxyService { if let Some(content_type) = session.get_header("Content-Type") { if let Ok(content_type_str) = content_type.to_str() { if content_type_str.contains("application/json") { - // Read the request body - let body = session.read_request_body().await?; - if let Some(body) = body { - // Parse the JSON payload into a serde_json::Value - if let Ok(json_value) = serde_json::from_slice::(&body) { + // Enable retry buffering to preserve the request body, reference: https://github.com/cloudflare/pingora/issues/349#issuecomment-2377277028 + session.enable_retry_buffering(); + session.read_body_or_idle(false).await.unwrap().unwrap(); + let request_body = session.get_retry_buffer(); + + // Parse the JSON payload into a serde_json::Value + if let Some(body_bytes) = request_body { + if let Ok(json_value) = serde_json::from_slice::(&body_bytes) { if let Some(model) = json_value.get("model").and_then(|v| v.as_str()) { // Set the requested_model field in the RequestContext ctx.requested_model = Some(model.to_string()); info!("Model in request: {:?}", ctx.requested_model); } + } else { + error!("Failed to parse JSON payload"); } + } else { + error!("Request body is None"); } } } diff --git a/src/balancer/upstream_peer_pool.rs b/src/balancer/upstream_peer_pool.rs index 6d852cc2..101c1b31 100644 --- a/src/balancer/upstream_peer_pool.rs +++ b/src/balancer/upstream_peer_pool.rs @@ -125,12 +125,12 @@ impl UpstreamPeerPool { let is_usable = peer.is_usable(); let is_usable_for_model = peer.is_usable_for_model(model_str); - info!("Peer {} is usable: {}, usable for model '{}': {}", peer.agent_id, is_usable, model_str, is_usable_for_model); - if is_usable && (model.is_none() || is_usable_for_model) { + info!("Peer {} is usable: {}, usable for model '{}': {}", peer.agent_id, is_usable, model_str, is_usable_for_model); return Ok(Some(peer.clone())); } } + Ok(None) }) } From b29c875c9348e464df9bae091cd68dde4935f9eb Mon Sep 17 00:00:00 2001 From: Jonas Krauss Date: Fri, 20 Jun 2025 17:21:32 +0200 Subject: [PATCH 04/18] add responses for unsupported/missing model parameter, add regex check to handle only first 64k characters available in buffer --- Cargo.lock | 1 + Cargo.toml | 1 + src/balancer/proxy_service.rs | 103 +++++++++++++++++++++----------- src/balancer/request_context.rs | 19 +++++- 4 files changed, 88 insertions(+), 36 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 60032252..bc4ae624 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2449,6 +2449,7 @@ dependencies = [ "mime_guess", "pingora", "ratatui", + "regex", "reqwest", "rust-embed", "serde", diff --git a/Cargo.toml b/Cargo.toml index 78bd06c5..172d4c6c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -45,6 +45,7 @@ itertools = "0.13.0" color-eyre = "0.6.5" time = "0.3.41" chrono = "0.4.41" +regex = "1.11.1" [features] default = ["statsd_reporter", "ratatui_dashboard"] diff --git a/src/balancer/proxy_service.rs b/src/balancer/proxy_service.rs index 9c59f15b..9a89fb13 100644 --- a/src/balancer/proxy_service.rs +++ b/src/balancer/proxy_service.rs @@ -149,40 +149,6 @@ impl ProxyHttp for ProxyService { session: &mut Session, ctx: &mut Self::CTX, ) -> Result> { - info!("upstream_peer - {:?} request | rewrite_host_header? {} check_model? {}", session.req_header().method, self.rewrite_host_header, self.check_model); - - // Check if the request method is POST and the content type is JSON - if self.check_model { - if session.req_header().method == "POST" { - // Check if the content type is application/json - if let Some(content_type) = session.get_header("Content-Type") { - if let Ok(content_type_str) = content_type.to_str() { - if content_type_str.contains("application/json") { - // Enable retry buffering to preserve the request body, reference: https://github.com/cloudflare/pingora/issues/349#issuecomment-2377277028 - session.enable_retry_buffering(); - session.read_body_or_idle(false).await.unwrap().unwrap(); - let request_body = session.get_retry_buffer(); - - // Parse the JSON payload into a serde_json::Value - if let Some(body_bytes) = request_body { - if let Ok(json_value) = serde_json::from_slice::(&body_bytes) { - if let Some(model) = json_value.get("model").and_then(|v| v.as_str()) { - // Set the requested_model field in the RequestContext - ctx.requested_model = Some(model.to_string()); - info!("Model in request: {:?}", ctx.requested_model); - } - } else { - error!("Failed to parse JSON payload"); - } - } else { - error!("Request body is None"); - } - } - } - } - } - } - let Some(_req_guard) = RequestBufferGuard::increment( &self.upstream_peer_pool.request_buffer_length, self.max_requests, @@ -213,10 +179,79 @@ impl ProxyHttp for ProxyService { } "/chat/completions" => true, "/completion" => true, + "/v1/completions" => true, "/v1/chat/completions" => true, _ => false, }; + info!("upstream_peer - {:?} request | rewrite_host_header? {} check_model? {}", session.req_header().method, self.rewrite_host_header, self.check_model); + + // Check if the request method is POST and the content type is JSON // cnbREaxdMcQVBS + if self.check_model && ctx.uses_slots { + info!("Checking model..."); + ctx.requested_model = None; + if session.req_header().method == "POST" { + // Check if the content type is application/json + if let Some(content_type) = session.get_header("Content-Type") { + if let Ok(content_type_str) = content_type.to_str() { + if content_type_str.contains("application/json") { + // Enable retry buffering to preserve the request body, reference: https://github.com/cloudflare/pingora/issues/349#issuecomment-2377277028 + session.enable_retry_buffering(); + session.read_body_or_idle(false).await.unwrap().unwrap(); + let request_body = session.get_retry_buffer(); + + // Parse the JSON payload into a serde_json::Value + if let Some(body_bytes) = request_body { + if let Ok(json_value) = serde_json::from_slice::(&body_bytes) { + if let Some(model) = json_value.get("model").and_then(|v| v.as_str()) { + // Set the requested_model field in the RequestContext + ctx.requested_model = Some(model.to_string()); + info!("Model in request: {:?}", ctx.requested_model); + } + } else { + info!("Failed to parse JSON payload, trying regex extraction"); + + // Try extracting the model using regex + let body_str = String::from_utf8(body_bytes.to_vec()).unwrap(); + let re = regex::Regex::new(r#""model"\s*:\s*["']([^"']*)["']"#).unwrap(); + if let Some(caps) = re.captures(&body_str) { + if let Some(model) = caps.get(1) { + ctx.requested_model = Some(model.as_str().to_string()); + info!("Model via regex: {:?}", ctx.requested_model); + } + } else { + info!("Failed to extract model using regex"); + } + } + } else { + info!("Request body is None"); + } + } + } + } + } + // abort if model has not been set + if ctx.requested_model == None { + info!("Model missing in request"); + session + .respond_error(pingora::http::StatusCode::BAD_REQUEST.as_u16()) + .await?; + + return Err(Error::new_down(pingora::ErrorType::ConnectRefused)); + } + else if ctx.has_peer_supporting_model() == false { + info!("Model {:?} not supported by upstream", ctx.requested_model); + session + .respond_error(pingora::http::StatusCode::NOT_FOUND.as_u16()) + .await?; + + return Err(Error::new_down(pingora::ErrorType::ConnectRefused)); + } + else { + info!("Model {:?}", ctx.requested_model); + } + } + let peer = tokio::select! { result = async { loop { diff --git a/src/balancer/request_context.rs b/src/balancer/request_context.rs index 47abb9ed..aedcd10d 100644 --- a/src/balancer/request_context.rs +++ b/src/balancer/request_context.rs @@ -60,6 +60,21 @@ impl RequestContext { ) } + pub fn has_peer_supporting_model(&self) -> bool { + let model_str = self.requested_model.as_deref().unwrap_or(""); + match self.upstream_peer_pool.with_agents_read(|agents| { + for peer in agents.iter() { + if peer.is_usable_for_model(model_str) { + return Ok(true); + } + } + Ok(false) + }) { + Ok(result) => result, + Err(_) => false, // or handle the error as needed + } + } + pub fn select_upstream_peer(&mut self) -> Result<()> { let result_option_peer = if self.uses_slots && !self.slot_taken { self.use_best_peer_and_take_slot(self.requested_model.clone()) @@ -114,7 +129,7 @@ mod tests { pool.register_status_update("test_agent", mock_status_update("test_agent", 0, 0))?; - assert!(ctx.use_best_peer_and_take_slot().unwrap().is_none()); + assert!(ctx.use_best_peer_and_take_slot(ctx.requested_model.clone()).unwrap().is_none()); assert!(!ctx.slot_taken); assert_eq!(ctx.selected_peer, None); @@ -139,7 +154,7 @@ mod tests { "127.0.0.1:8080" ); - ctx.use_best_peer_and_take_slot()?; + ctx.use_best_peer_and_take_slot(ctx.requested_model.clone())?; assert!(ctx.slot_taken); From b328bab9c9bc6965f7b82aee1e437a080ee51469 Mon Sep 17 00:00:00 2001 From: Jonas Krauss Date: Mon, 23 Jun 2025 14:54:42 +0200 Subject: [PATCH 05/18] fix mock_status_update.rs --- src/balancer/test/mock_status_update.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/balancer/test/mock_status_update.rs b/src/balancer/test/mock_status_update.rs index 2f22e95e..7d61139f 100644 --- a/src/balancer/test/mock_status_update.rs +++ b/src/balancer/test/mock_status_update.rs @@ -46,7 +46,7 @@ pub fn mock_status_update( is_authorized: Some(true), is_slots_endpoint_enabled: Some(true), processing_slots_count: slots.len() - idle_slots_count, + model: Some("llama3".to_string()), slots: slots, - Some("llama3".to_string()), } } From c2efc634edf909f7c6c0ab40fa50de019cc6b96e Mon Sep 17 00:00:00 2001 From: Jonas Krauss Date: Fri, 27 Jun 2025 15:50:55 +0200 Subject: [PATCH 06/18] fix returning model not supported on all slots for model currently blocked --- src/balancer/proxy_service.rs | 2 +- src/balancer/request_context.rs | 2 +- src/balancer/upstream_peer.rs | 4 ++++ 3 files changed, 6 insertions(+), 2 deletions(-) diff --git a/src/balancer/proxy_service.rs b/src/balancer/proxy_service.rs index 2e75e6b4..eea439bd 100644 --- a/src/balancer/proxy_service.rs +++ b/src/balancer/proxy_service.rs @@ -192,7 +192,7 @@ impl ProxyHttp for ProxyService { info!("upstream_peer - {:?} request | rewrite_host_header? {} check_model? {}", session.req_header().method, self.rewrite_host_header, self.check_model); - // Check if the request method is POST and the content type is JSON // cnbREaxdMcQVBS + // Check if the request method is POST and the content type is JSON if self.check_model && ctx.uses_slots { info!("Checking model..."); ctx.requested_model = None; diff --git a/src/balancer/request_context.rs b/src/balancer/request_context.rs index 70bc042a..ef03ce48 100644 --- a/src/balancer/request_context.rs +++ b/src/balancer/request_context.rs @@ -61,7 +61,7 @@ impl RequestContext { let model_str = self.requested_model.as_deref().unwrap_or(""); match self.upstream_peer_pool.with_agents_read(|agents| { for peer in agents.iter() { - if peer.is_usable_for_model(model_str) { + if peer.supports_model(model_str) { return Ok(true); } } diff --git a/src/balancer/upstream_peer.rs b/src/balancer/upstream_peer.rs index 6a004e77..d73f6a6e 100644 --- a/src/balancer/upstream_peer.rs +++ b/src/balancer/upstream_peer.rs @@ -37,6 +37,10 @@ impl UpstreamPeer { !self.status.has_issues() && self.status.slots_idle > 0 && self.quarantined_until.is_none() } + pub fn supports_model(&self, requested_model: &str) -> bool { + requested_model.is_empty() || self.model.as_deref() == Some(requested_model) + } + pub fn is_usable_for_model(&self, requested_model: &str) -> bool { self.is_usable() && (requested_model.is_empty() || self.model.as_deref() == Some(requested_model)) } From eb0c275e1606d7675690a18293b0454f6691f109 Mon Sep 17 00:00:00 2001 From: Jonas Krauss Date: Tue, 1 Jul 2025 11:00:58 +0200 Subject: [PATCH 07/18] make model nullable for dashboard --- resources/ts/schemas/Agent.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/resources/ts/schemas/Agent.ts b/resources/ts/schemas/Agent.ts index 4a1fef54..c2eea84d 100644 --- a/resources/ts/schemas/Agent.ts +++ b/resources/ts/schemas/Agent.ts @@ -5,7 +5,7 @@ import { StatusUpdateSchema } from "./StatusUpdate"; export const AgentSchema = z .object({ agent_id: z.string(), - model: z.string(), + model: z.string().nullable(), last_update: z.object({ nanos_since_epoch: z.number(), secs_since_epoch: z.number(), From 961d4b4bbff8210e4a23d6ddfb3d763300ae35a8 Mon Sep 17 00:00:00 2001 From: Jonas Krauss Date: Sat, 12 Jul 2025 10:44:09 +0200 Subject: [PATCH 08/18] add handling of non-utf8 chunk in first 64k bytes of request body --- src/balancer/proxy_service.rs | 67 +++++++++++++++++++++++++---------- 1 file changed, 48 insertions(+), 19 deletions(-) diff --git a/src/balancer/proxy_service.rs b/src/balancer/proxy_service.rs index eea439bd..9eeb1a54 100644 --- a/src/balancer/proxy_service.rs +++ b/src/balancer/proxy_service.rs @@ -206,27 +206,56 @@ impl ProxyHttp for ProxyService { session.read_body_or_idle(false).await.unwrap().unwrap(); let request_body = session.get_retry_buffer(); - // Parse the JSON payload into a serde_json::Value if let Some(body_bytes) = request_body { - if let Ok(json_value) = serde_json::from_slice::(&body_bytes) { - if let Some(model) = json_value.get("model").and_then(|v| v.as_str()) { - // Set the requested_model field in the RequestContext - ctx.requested_model = Some(model.to_string()); - info!("Model in request: {:?}", ctx.requested_model); - } - } else { - info!("Failed to parse JSON payload, trying regex extraction"); - - // Try extracting the model using regex - let body_str = String::from_utf8(body_bytes.to_vec()).unwrap(); - let re = regex::Regex::new(r#""model"\s*:\s*["']([^"']*)["']"#).unwrap(); - if let Some(caps) = re.captures(&body_str) { - if let Some(model) = caps.get(1) { - ctx.requested_model = Some(model.as_str().to_string()); - info!("Model via regex: {:?}", ctx.requested_model); + match std::str::from_utf8(&body_bytes) { + Ok(_) => { + // The bytes are valid UTF-8, proceed as normal + if let Ok(json_value) = serde_json::from_slice::(&body_bytes) { + if let Some(model) = json_value.get("model").and_then(|v| v.as_str()) { + ctx.requested_model = Some(model.to_string()); + info!("Model in request: {:?}", ctx.requested_model); + } + } else { + info!("Failed to parse JSON payload, trying regex extraction"); + let body_str = String::from_utf8_lossy(&body_bytes).to_string(); + let re = regex::Regex::new(r#""model"\s*:\s*["']([^"']*)["']"#).unwrap(); + if let Some(caps) = re.captures(&body_str) { + if let Some(model) = caps.get(1) { + ctx.requested_model = Some(model.as_str().to_string()); + info!("Model via regex: {:?}", ctx.requested_model); + } + } else { + info!("Failed to extract model using regex"); + } + } + }, + Err(e) => { + // Invalid UTF-8 detected. Truncate to the last valid UTF-8 boundary. + let valid_up_to = e.valid_up_to(); + info!("Invalid UTF-8 detected. Truncating from {} bytes to {} bytes.", body_bytes.len(), valid_up_to); + + // Create a new `Bytes` slice containing only the valid UTF-8 part. + let valid_body_bytes = body_bytes.slice(0..valid_up_to); + + // Now proceed with the (truncated) valid_body_bytes + if let Ok(json_value) = serde_json::from_slice::(&valid_body_bytes) { + if let Some(model) = json_value.get("model").and_then(|v| v.as_str()) { + ctx.requested_model = Some(model.to_string()); + info!("Model in request (after truncation): {:?}", ctx.requested_model); + } + } else { + info!("Failed to parse JSON payload (after truncation), trying regex extraction"); + let body_str = String::from_utf8_lossy(&valid_body_bytes).to_string(); + let re = regex::Regex::new(r#""model"\s*:\s*["']([^"']*)["']"#).unwrap(); + if let Some(caps) = re.captures(&body_str) { + if let Some(model) = caps.get(1) { + ctx.requested_model = Some(model.as_str().to_string()); + info!("Model via regex (after truncation): {:?}", ctx.requested_model); + } + } else { + info!("Failed to extract model using regex (after truncation)"); + } } - } else { - info!("Failed to extract model using regex"); } } } else { From 1eac398fd0d6219a93651d14281d6f2aac3ca8f8 Mon Sep 17 00:00:00 2001 From: Jonas Krauss Date: Tue, 24 Feb 2026 13:27:35 +0100 Subject: [PATCH 09/18] add sort option to balancer --- resources/ts/components/AgentsList.tsx | 141 +++++++++++++++++-- resources/ts/components/Dashboard.module.css | 18 +++ 2 files changed, 144 insertions(+), 15 deletions(-) diff --git a/resources/ts/components/AgentsList.tsx b/resources/ts/components/AgentsList.tsx index da927386..140ba58f 100644 --- a/resources/ts/components/AgentsList.tsx +++ b/resources/ts/components/AgentsList.tsx @@ -1,5 +1,5 @@ import clsx from "clsx"; -import React, { CSSProperties } from "react"; +import React, { CSSProperties, useState } from "react"; import { type Agent } from "../schemas/Agent"; @@ -9,28 +9,141 @@ import { agentUsage, agentUsage__progress, agentsTable, + sortIndicator, + sortIndicatorAsc, + sortIndicatorDesc, } from "./Dashboard.module.css"; function formatTimestamp(timestamp: number): string { return new Date(timestamp * 1000).toLocaleString(); } +type SortColumn = + | "name" + | "model" + | "issues" + | "llamacppAddr" + | "lastUpdate" + | "idleSlots" + | "processingSlots"; + +function getSortIndicator( + sortConfig: { key: SortColumn; direction: "ascending" | "descending" }, + currentKey: SortColumn +): React.ReactNode { + if (sortConfig.key !== currentKey) { + return null; + } + const className = clsx(sortIndicator, sortConfig.direction === "ascending" ? sortIndicatorAsc : sortIndicatorDesc); + return ( + + {sortConfig.direction === "ascending" ? "↑" : "↓"} + + ); +} + export function AgentsList({ agents }: { agents: Array }) { + const [sortConfig, setSortConfig] = useState<{ + key: SortColumn; + direction: "ascending" | "descending"; + }>({ key: "name", direction: "ascending" }); + + function sortAgents(agents: Array): Array { + const sortableAgents = [...agents]; + sortableAgents.sort(function (a, b) { + const { key, direction } = sortConfig; + + // Helper function to get comparison value based on column type + function getValue(agent: Agent, key: SortColumn): string | number { + let hasIssuesA: boolean; + let hasIssuesB: boolean; + switch (key) { + case "name": + return a.status.agent_name || ""; + case "model": + return a.status.model || ""; + case "issues": + // Sort by presence of issues first, then by content + hasIssuesA = a.status.error !== null; + hasIssuesB = b.status.error !== null; + if (hasIssuesA !== hasIssuesB) { + return hasIssuesA ? 1 : -1; + } + return (a.status.error || "") > (b.status.error || "") ? 1 : -1; + case "llamacppAddr": + return a.status.external_llamacpp_addr; + case "lastUpdate": + return a.last_update.secs_since_epoch; + case "idleSlots": + return a.status.slots_idle; + case "processingSlots": + return a.status.slots_processing; + default: + return ""; + } + } + + const valueA = getValue(a, key); + const valueB = getValue(b, key); + + // Handle string comparison + if (typeof valueA === "string" && typeof valueB === "string") { + if (valueA < valueB) return direction === "ascending" ? -1 : 1; + if (valueA > valueB) return direction === "ascending" ? 1 : -1; + return 0; + } + + // Handle numeric comparison + if (typeof valueA === "number" && typeof valueB === "number") { + if (valueA < valueB) return direction === "ascending" ? -1 : 1; + if (valueA > valueB) return direction === "ascending" ? 1 : -1; + return 0; + } + + return 0; + }); + return sortableAgents; + } + + function requestSort(key: SortColumn) { + let direction: "ascending" | "descending" = "ascending"; + if (sortConfig.key === key && sortConfig.direction === "ascending") { + direction = "descending"; + } + setSortConfig({ key, direction }); + } + + const sortedAgents = sortAgents(agents); + return ( - - - - - - - + + + + + + + - {agents.map(function ({ + {sortedAgents.map(function ({ agent_id, last_update, quarantined_until, @@ -48,12 +161,10 @@ export function AgentsList({ agents }: { agents: Array }) { quarantined_until; return ( - +
NameModelIssuesLlama.cpp addressLast updateIdle slotsProcessing slots + Name{getSortIndicator(sortConfig, "name")} + + Model{getSortIndicator(sortConfig, "model")} + + Issues{getSortIndicator(sortConfig, "issues")} + + Llama.cpp address{getSortIndicator(sortConfig, "llamacppAddr")} + + Last update{getSortIndicator(sortConfig, "lastUpdate")} + + Idle slots{getSortIndicator(sortConfig, "idleSlots")} + + Processing slots{getSortIndicator(sortConfig, "processingSlots")} +
{status.agent_name} {status.model} diff --git a/resources/ts/components/Dashboard.module.css b/resources/ts/components/Dashboard.module.css index 86f49b91..85458e4e 100644 --- a/resources/ts/components/Dashboard.module.css +++ b/resources/ts/components/Dashboard.module.css @@ -7,11 +7,29 @@ th { border: 1px solid var(--color-border); padding: var(--spacing-base); + cursor: pointer; p + p { margin-top: var(--spacing-half); } } + + th:hover { + background-color: var(--color-hover); + } +} + +.sortIndicator { + margin-left: var(--spacing-half); + font-size: 0.8em; +} + +.sortIndicatorAsc { + color: var(--color-success); +} + +.sortIndicatorDesc { + color: var(--color-error); } .agentRow.agentRowError { From 5a0713d70eff9a40305eabf2be143f177fc67f0f Mon Sep 17 00:00:00 2001 From: Jonas Krauss Date: Tue, 24 Feb 2026 13:37:32 +0100 Subject: [PATCH 10/18] fix test --- src/balancer/upstream_peer.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/balancer/upstream_peer.rs b/src/balancer/upstream_peer.rs index 2c4ccee9..24b52d62 100644 --- a/src/balancer/upstream_peer.rs +++ b/src/balancer/upstream_peer.rs @@ -121,7 +121,7 @@ mod tests { fn create_test_peer() -> UpstreamPeer { UpstreamPeer { agent_id: "test_agent".to_string(), - model: "llama3".to_string(), + model: Some("llama3".to_string()), last_update: SystemTime::now(), quarantined_until: None, slots_taken: 0, @@ -142,6 +142,7 @@ mod tests { is_unexpected_response_status: None, slots_idle: 5, slots_processing: 0, + model: Some("llama3".to_string()), }, } } From 78357cc414af826c9adef2850ccb9ff0301d7153 Mon Sep 17 00:00:00 2001 From: Jonas Krauss Date: Tue, 24 Feb 2026 14:17:19 +0100 Subject: [PATCH 11/18] fix sort not working properly --- resources/ts/components/AgentsList.tsx | 36 ++++++++++++++------------ 1 file changed, 20 insertions(+), 16 deletions(-) diff --git a/resources/ts/components/AgentsList.tsx b/resources/ts/components/AgentsList.tsx index 140ba58f..1c0cbcfb 100644 --- a/resources/ts/components/AgentsList.tsx +++ b/resources/ts/components/AgentsList.tsx @@ -55,34 +55,38 @@ export function AgentsList({ agents }: { agents: Array }) { // Helper function to get comparison value based on column type function getValue(agent: Agent, key: SortColumn): string | number { - let hasIssuesA: boolean; - let hasIssuesB: boolean; switch (key) { case "name": - return a.status.agent_name || ""; + return agent.status.agent_name || ""; case "model": - return a.status.model || ""; - case "issues": - // Sort by presence of issues first, then by content - hasIssuesA = a.status.error !== null; - hasIssuesB = b.status.error !== null; - if (hasIssuesA !== hasIssuesB) { - return hasIssuesA ? 1 : -1; - } - return (a.status.error || "") > (b.status.error || "") ? 1 : -1; + return agent.status.model || ""; case "llamacppAddr": - return a.status.external_llamacpp_addr; + return agent.status.external_llamacpp_addr; case "lastUpdate": - return a.last_update.secs_since_epoch; + return agent.last_update.secs_since_epoch; case "idleSlots": - return a.status.slots_idle; + return agent.status.slots_idle; case "processingSlots": - return a.status.slots_processing; + return agent.status.slots_processing; default: return ""; } } + // Special handling for issues column + if (key === "issues") { + const hasIssuesA = a.status.error !== null; + const hasIssuesB = b.status.error !== null; + if (hasIssuesA !== hasIssuesB) { + return direction === "ascending" ? (hasIssuesA ? 1 : -1) : (hasIssuesA ? -1 : 1); + } + const errorA = a.status.error || ""; + const errorB = b.status.error || ""; + if (errorA < errorB) return direction === "ascending" ? -1 : 1; + if (errorA > errorB) return direction === "ascending" ? 1 : -1; + return 0; + } + const valueA = getValue(a, key); const valueB = getValue(b, key); From c0a916ea1530479248c6c129bd28e0b561b6df88 Mon Sep 17 00:00:00 2001 From: Jonas Krauss Date: Tue, 30 Jun 2026 16:59:07 +0200 Subject: [PATCH 12/18] fix: improve model detection robustness in proxy service - Replace .unwrap().unwrap() with proper error handling to prevent panics on empty bodies or read errors (returns 400 instead) - Add retry_buffer_truncated() check to reject >64KB bodies with a clear error message instead of silent failure - Consolidate model extraction into single extract_model_from_body() function, deduplicating UTF-8 and non-UTF-8 code paths - Fix regex to match only JSON double quotes (not single quotes) - Compile regex once via LazyLock instead of per-request - Initialize requested_model to None instead of Some("") - Use idiomatic .is_none() and !bool instead of == None / == false --- src/balancer/proxy_service.rs | 145 +++++++++++++++++++--------------- 1 file changed, 83 insertions(+), 62 deletions(-) diff --git a/src/balancer/proxy_service.rs b/src/balancer/proxy_service.rs index 9eeb1a54..4f0a2ac9 100644 --- a/src/balancer/proxy_service.rs +++ b/src/balancer/proxy_service.rs @@ -1,6 +1,7 @@ use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; use std::sync::Arc; +use std::sync::LazyLock; use std::time::Duration; use async_trait::async_trait; @@ -18,6 +19,51 @@ use pingora::Result; use crate::balancer::request_context::RequestContext; use crate::balancer::upstream_peer_pool::UpstreamPeerPool; +static MODEL_REGEX: LazyLock = LazyLock::new(|| { + regex::Regex::new(r#""model"\s*:\s*"([^"]*)""#).expect("model regex is valid") +}); + +/// Extract the "model" field from a JSON request body. +/// +/// Strategy: +/// 1. If the body contains invalid UTF-8, truncate to the last valid boundary. +/// 2. Try proper JSON parsing first (most reliable, no false positives). +/// 3. If JSON parsing fails, try regex extraction as a fallback. +fn extract_model_from_body(body_bytes: &Bytes) -> Option { + // Handle invalid UTF-8 by truncating to the last valid boundary + let effective_bytes = match std::str::from_utf8(body_bytes) { + Ok(_) => body_bytes.clone(), + Err(e) => { + let valid_up_to = e.valid_up_to(); + info!("Invalid UTF-8 in request body. Truncating from {} bytes to {} bytes for model extraction.", body_bytes.len(), valid_up_to); + body_bytes.slice(0..valid_up_to) + } + }; + + // Try proper JSON parsing first + if let Ok(json_value) = serde_json::from_slice::(&effective_bytes) { + if let Some(model) = json_value.get("model").and_then(|v| v.as_str()) { + let model = model.to_string(); + info!("Model in request: {:?}", model); + return Some(model); + } + } + + // Fallback: regex extraction on the raw text + info!("Failed to parse JSON payload, trying regex extraction"); + let body_str = String::from_utf8_lossy(&effective_bytes); + if let Some(caps) = MODEL_REGEX.captures(&body_str) { + if let Some(model_match) = caps.get(1) { + let model = model_match.as_str().to_string(); + info!("Model via regex: {:?}", model); + return Some(model); + } + } + + info!("Failed to extract model from request body"); + None +} + struct RequestBufferGuard<'a>(&'a AtomicUsize); impl<'a> RequestBufferGuard<'a> { @@ -77,7 +123,7 @@ impl ProxyHttp for ProxyService { slot_taken: false, upstream_peer_pool: self.upstream_peer_pool.clone(), uses_slots: false, - requested_model: Some("".to_string()), + requested_model: None, } } @@ -195,7 +241,6 @@ impl ProxyHttp for ProxyService { // Check if the request method is POST and the content type is JSON if self.check_model && ctx.uses_slots { info!("Checking model..."); - ctx.requested_model = None; if session.req_header().method == "POST" { // Check if the content type is application/json if let Some(content_type) = session.get_header("Content-Type") { @@ -203,86 +248,62 @@ impl ProxyHttp for ProxyService { if content_type_str.contains("application/json") { // Enable retry buffering to preserve the request body, reference: https://github.com/cloudflare/pingora/issues/349#issuecomment-2377277028 session.enable_retry_buffering(); - session.read_body_or_idle(false).await.unwrap().unwrap(); - let request_body = session.get_retry_buffer(); - - if let Some(body_bytes) = request_body { - match std::str::from_utf8(&body_bytes) { - Ok(_) => { - // The bytes are valid UTF-8, proceed as normal - if let Ok(json_value) = serde_json::from_slice::(&body_bytes) { - if let Some(model) = json_value.get("model").and_then(|v| v.as_str()) { - ctx.requested_model = Some(model.to_string()); - info!("Model in request: {:?}", ctx.requested_model); - } - } else { - info!("Failed to parse JSON payload, trying regex extraction"); - let body_str = String::from_utf8_lossy(&body_bytes).to_string(); - let re = regex::Regex::new(r#""model"\s*:\s*["']([^"']*)["']"#).unwrap(); - if let Some(caps) = re.captures(&body_str) { - if let Some(model) = caps.get(1) { - ctx.requested_model = Some(model.as_str().to_string()); - info!("Model via regex: {:?}", ctx.requested_model); - } - } else { - info!("Failed to extract model using regex"); - } - } - }, - Err(e) => { - // Invalid UTF-8 detected. Truncate to the last valid UTF-8 boundary. - let valid_up_to = e.valid_up_to(); - info!("Invalid UTF-8 detected. Truncating from {} bytes to {} bytes.", body_bytes.len(), valid_up_to); - - // Create a new `Bytes` slice containing only the valid UTF-8 part. - let valid_body_bytes = body_bytes.slice(0..valid_up_to); - - // Now proceed with the (truncated) valid_body_bytes - if let Ok(json_value) = serde_json::from_slice::(&valid_body_bytes) { - if let Some(model) = json_value.get("model").and_then(|v| v.as_str()) { - ctx.requested_model = Some(model.to_string()); - info!("Model in request (after truncation): {:?}", ctx.requested_model); - } - } else { - info!("Failed to parse JSON payload (after truncation), trying regex extraction"); - let body_str = String::from_utf8_lossy(&valid_body_bytes).to_string(); - let re = regex::Regex::new(r#""model"\s*:\s*["']([^"']*)["']"#).unwrap(); - if let Some(caps) = re.captures(&body_str) { - if let Some(model) = caps.get(1) { - ctx.requested_model = Some(model.as_str().to_string()); - info!("Model via regex (after truncation): {:?}", ctx.requested_model); - } - } else { - info!("Failed to extract model using regex (after truncation)"); - } - } + + // Read one chunk from the body. The model field is typically + // the first field in the JSON, so one read is sufficient. + // The retry buffer (64KB) captures what was read. + let read_result = session.read_body_or_idle(false).await; + match read_result { + Ok(Some(_)) => { + // Check if the retry buffer was truncated (body > 64KB). + // If truncated, get_retry_buffer() returns None, so we + // cannot extract the model and must reject the request. + if session.retry_buffer_truncated() { + error!("Request body exceeds 64KB retry buffer limit, cannot determine model"); + session + .respond_error(pingora::http::StatusCode::BAD_REQUEST.as_u16()) + .await?; + return Err(Error::new_down(pingora::ErrorType::ConnectRefused)); + } + + if let Some(body_bytes) = session.get_retry_buffer() { + ctx.requested_model = extract_model_from_body(&body_bytes); + } else { + info!("Retry buffer is empty after reading body chunk"); } } - } else { - info!("Request body is None"); + Ok(None) => { + info!("Request body is empty"); + } + Err(e) => { + error!("Failed to read request body: {e}"); + session + .respond_error(pingora::http::StatusCode::BAD_REQUEST.as_u16()) + .await?; + return Err(Error::new_down(pingora::ErrorType::ConnectRefused)); + } } } } } } + // abort if model has not been set - if ctx.requested_model == None { + if ctx.requested_model.is_none() { info!("Model missing in request"); session .respond_error(pingora::http::StatusCode::BAD_REQUEST.as_u16()) .await?; return Err(Error::new_down(pingora::ErrorType::ConnectRefused)); - } - else if ctx.has_peer_supporting_model() == false { + } else if !ctx.has_peer_supporting_model() { info!("Model {:?} not supported by upstream", ctx.requested_model); session .respond_error(pingora::http::StatusCode::NOT_FOUND.as_u16()) .await?; return Err(Error::new_down(pingora::ErrorType::ConnectRefused)); - } - else { + } else { info!("Model {:?}", ctx.requested_model); } } From b3d55fc36db8b34fb08b394e1577565f4f9854ad Mon Sep 17 00:00:00 2001 From: Jonas Krauss Date: Tue, 30 Jun 2026 17:11:57 +0200 Subject: [PATCH 13/18] fix: set idle timeout on upstream connections to prevent stale connections Set HttpPeer idle_timeout to 30s to prevent pooled upstream connections from accumulating indefinitely (default was None, meaning no expiry). --- src/balancer/proxy_service.rs | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/balancer/proxy_service.rs b/src/balancer/proxy_service.rs index 4f0a2ac9..62b243f5 100644 --- a/src/balancer/proxy_service.rs +++ b/src/balancer/proxy_service.rs @@ -356,7 +356,11 @@ impl ProxyHttp for ProxyService { } }; - Ok(HttpPeer::new(peer.status.external_llamacpp_addr, false, "".into()).into()) + let mut http_peer = HttpPeer::new(peer.status.external_llamacpp_addr, false, "".into()); + // Expire pooled upstream connections after 30s of inactivity to prevent + // stale connections from accumulating (idle_timeout was None by default). + http_peer.options.idle_timeout = Some(Duration::from_secs(30)); + Ok(http_peer.into()) } async fn upstream_request_filter( From d709484986d8a98bee7662675e03f785c05a0188 Mon Sep 17 00:00:00 2001 From: Jonas Krauss Date: Wed, 1 Jul 2026 14:05:39 +0200 Subject: [PATCH 14/18] chore: restructure tests directory and add model detection test suite - Move integration_tests/ into new tests/ directory (tests/integration_tests/) - Add tests/scripts/test-model-detection.sh with 19 tests covering: valid requests, wrong/missing models, empty/invalid bodies, non-JSON content types, >64KB payloads (needle-in-haystack), vision requests, streaming, legacy endpoints, and edge cases - Update all references in Cargo.toml, Makefile, test-balancer.sh, and jarmuz/worker-prettier.mjs - Add 'make shell_tests' target to Makefile --- Cargo.toml | 2 +- Makefile | 8 +- jarmuz/worker-prettier.mjs | 2 +- .../integration_tests}/.gitignore | 0 .../integration_tests}/Cargo.lock | 0 .../integration_tests}/Cargo.toml | 0 .../integration_tests}/Makefile | 0 .../integration_tests}/rust-toolchain | 0 .../integration_tests}/rustfmt.toml | 0 .../features/agent/agent_monitors.feature | 0 .../balancer/balance_requests.feature | 0 .../features/balancer/cors_headers.feature | 0 .../features/balancer/export_metrics.feature | 0 .../features/balancer/report_metrics.feature | 0 .../tests/fixtures/llamacpp-server-mock.mjs | 0 .../tests/fixtures/statsd-server-mock.mjs | 0 .../tests/paddler/agent_response.rs | 0 .../tests/paddler/agents_collection.rs | 0 .../tests/paddler/assert_balancer_table.rs | 0 .../tests/paddler/balancer_instance.rs | 0 .../paddler/balancer_management_client.rs | 0 .../expression/given_agent_is_registered.rs | 0 .../expression/given_agent_is_running.rs | 0 ...nt_monitors_llamacpp_every_milliseconds.rs | 0 .../given_balancer_allows_cors_host.rs | 0 .../expression/given_balancer_is_running.rs | 0 ...ncer_reports_metrics_every_milliseconds.rs | 0 ...red_requests_timeout_after_milliseconds.rs | 0 ...p_completes_response_after_milliseconds.rs | 0 .../given_llamacpp_server_is_running.rs | 0 .../given_request_buffering_is_disabled.rs | 0 .../given_request_has_header_with_value.rs | 0 .../expression/given_statsd_is_running.rs | 0 .../tests/paddler/expression/mod.rs | 0 .../expression/then_average_metrics_are.rs | 0 .../expression/then_balancer_state_is.rs | 0 .../then_exported_metrics_are_stored.rs | 0 .../expression/then_next_balancer_state_is.rs | 0 .../then_reported_metrics_are_stored.rs | 0 .../expression/then_request_landed_in.rs | 0 .../expression/then_response_code_is.rs | 0 .../expression/then_response_header_is.rs | 0 .../then_response_header_is_not_present.rs | 0 .../expression/when_llamacpp_stops_running.rs | 0 ...when_multiple_requests_are_sent_to_path.rs | 0 ..._request_is_sent_to_management_endpoint.rs | 0 .../when_request_is_sent_to_path.rs | 0 .../tests/paddler/llamacpp_instance.rs | 0 .../paddler/llamacpp_instance_collection.rs | 0 .../integration_tests}/tests/paddler/main.rs | 0 .../tests/paddler/metrics.rs | 0 .../tests/paddler/paddler_world.rs | 0 .../tests/paddler/request_builder.rs | 0 .../paddler/request_headers_to_be_set.rs | 0 .../tests/paddler/statsd_instance.rs | 0 tests/scripts/test-model-detection.sh | 353 ++++++++++++++++++ 56 files changed, 361 insertions(+), 4 deletions(-) rename {integration_tests => tests/integration_tests}/.gitignore (100%) rename {integration_tests => tests/integration_tests}/Cargo.lock (100%) rename {integration_tests => tests/integration_tests}/Cargo.toml (100%) rename {integration_tests => tests/integration_tests}/Makefile (100%) rename {integration_tests => tests/integration_tests}/rust-toolchain (100%) rename {integration_tests => tests/integration_tests}/rustfmt.toml (100%) rename {integration_tests => tests/integration_tests}/tests/features/agent/agent_monitors.feature (100%) rename {integration_tests => tests/integration_tests}/tests/features/balancer/balance_requests.feature (100%) rename {integration_tests => tests/integration_tests}/tests/features/balancer/cors_headers.feature (100%) rename {integration_tests => tests/integration_tests}/tests/features/balancer/export_metrics.feature (100%) rename {integration_tests => tests/integration_tests}/tests/features/balancer/report_metrics.feature (100%) rename {integration_tests => tests/integration_tests}/tests/fixtures/llamacpp-server-mock.mjs (100%) rename {integration_tests => tests/integration_tests}/tests/fixtures/statsd-server-mock.mjs (100%) rename {integration_tests => tests/integration_tests}/tests/paddler/agent_response.rs (100%) rename {integration_tests => tests/integration_tests}/tests/paddler/agents_collection.rs (100%) rename {integration_tests => tests/integration_tests}/tests/paddler/assert_balancer_table.rs (100%) rename {integration_tests => tests/integration_tests}/tests/paddler/balancer_instance.rs (100%) rename {integration_tests => tests/integration_tests}/tests/paddler/balancer_management_client.rs (100%) rename {integration_tests => tests/integration_tests}/tests/paddler/expression/given_agent_is_registered.rs (100%) rename {integration_tests => tests/integration_tests}/tests/paddler/expression/given_agent_is_running.rs (100%) rename {integration_tests => tests/integration_tests}/tests/paddler/expression/given_agent_monitors_llamacpp_every_milliseconds.rs (100%) rename {integration_tests => tests/integration_tests}/tests/paddler/expression/given_balancer_allows_cors_host.rs (100%) rename {integration_tests => tests/integration_tests}/tests/paddler/expression/given_balancer_is_running.rs (100%) rename {integration_tests => tests/integration_tests}/tests/paddler/expression/given_balancer_reports_metrics_every_milliseconds.rs (100%) rename {integration_tests => tests/integration_tests}/tests/paddler/expression/given_buffered_requests_timeout_after_milliseconds.rs (100%) rename {integration_tests => tests/integration_tests}/tests/paddler/expression/given_llamacpp_completes_response_after_milliseconds.rs (100%) rename {integration_tests => tests/integration_tests}/tests/paddler/expression/given_llamacpp_server_is_running.rs (100%) rename {integration_tests => tests/integration_tests}/tests/paddler/expression/given_request_buffering_is_disabled.rs (100%) rename {integration_tests => tests/integration_tests}/tests/paddler/expression/given_request_has_header_with_value.rs (100%) rename {integration_tests => tests/integration_tests}/tests/paddler/expression/given_statsd_is_running.rs (100%) rename {integration_tests => tests/integration_tests}/tests/paddler/expression/mod.rs (100%) rename {integration_tests => tests/integration_tests}/tests/paddler/expression/then_average_metrics_are.rs (100%) rename {integration_tests => tests/integration_tests}/tests/paddler/expression/then_balancer_state_is.rs (100%) rename {integration_tests => tests/integration_tests}/tests/paddler/expression/then_exported_metrics_are_stored.rs (100%) rename {integration_tests => tests/integration_tests}/tests/paddler/expression/then_next_balancer_state_is.rs (100%) rename {integration_tests => tests/integration_tests}/tests/paddler/expression/then_reported_metrics_are_stored.rs (100%) rename {integration_tests => tests/integration_tests}/tests/paddler/expression/then_request_landed_in.rs (100%) rename {integration_tests => tests/integration_tests}/tests/paddler/expression/then_response_code_is.rs (100%) rename {integration_tests => tests/integration_tests}/tests/paddler/expression/then_response_header_is.rs (100%) rename {integration_tests => tests/integration_tests}/tests/paddler/expression/then_response_header_is_not_present.rs (100%) rename {integration_tests => tests/integration_tests}/tests/paddler/expression/when_llamacpp_stops_running.rs (100%) rename {integration_tests => tests/integration_tests}/tests/paddler/expression/when_multiple_requests_are_sent_to_path.rs (100%) rename {integration_tests => tests/integration_tests}/tests/paddler/expression/when_request_is_sent_to_management_endpoint.rs (100%) rename {integration_tests => tests/integration_tests}/tests/paddler/expression/when_request_is_sent_to_path.rs (100%) rename {integration_tests => tests/integration_tests}/tests/paddler/llamacpp_instance.rs (100%) rename {integration_tests => tests/integration_tests}/tests/paddler/llamacpp_instance_collection.rs (100%) rename {integration_tests => tests/integration_tests}/tests/paddler/main.rs (100%) rename {integration_tests => tests/integration_tests}/tests/paddler/metrics.rs (100%) rename {integration_tests => tests/integration_tests}/tests/paddler/paddler_world.rs (100%) rename {integration_tests => tests/integration_tests}/tests/paddler/request_builder.rs (100%) rename {integration_tests => tests/integration_tests}/tests/paddler/request_headers_to_be_set.rs (100%) rename {integration_tests => tests/integration_tests}/tests/paddler/statsd_instance.rs (100%) create mode 100755 tests/scripts/test-model-detection.sh diff --git a/Cargo.toml b/Cargo.toml index 86d03cce..5ad3d89c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -9,7 +9,7 @@ documentation = "https://github.com/distantmagic/paddler/blob/main/README.md" version = "1.2.1" [workspace] -members = ["integration_tests"] +members = ["tests/integration_tests"] [dependencies] actix = "0.13.5" diff --git a/Makefile b/Makefile index c36066ac..1609223e 100644 --- a/Makefile +++ b/Makefile @@ -31,12 +31,16 @@ clean: .PHONY: fmt fmt: node_modules ./jarmuz-fmt.mjs - $(MAKE) -C integration_tests fmt + $(MAKE) -C tests/integration_tests fmt .PHONY: integration_tests integration_tests: cargo build - $(MAKE) -C integration_tests test + $(MAKE) -C tests/integration_tests test + +.PHONY: shell_tests +shell_tests: + bash tests/scripts/test-model-detection.sh .PHONY: test test: integration_tests diff --git a/jarmuz/worker-prettier.mjs b/jarmuz/worker-prettier.mjs index 5e16d4fe..3d13f9b1 100644 --- a/jarmuz/worker-prettier.mjs +++ b/jarmuz/worker-prettier.mjs @@ -4,7 +4,7 @@ command(` npm exec prettier -- --plugin=prettier-plugin-organize-imports --write - integration_tests/tests/fixtures + tests/integration_tests/tests/fixtures jarmuz resources *.mjs diff --git a/integration_tests/.gitignore b/tests/integration_tests/.gitignore similarity index 100% rename from integration_tests/.gitignore rename to tests/integration_tests/.gitignore diff --git a/integration_tests/Cargo.lock b/tests/integration_tests/Cargo.lock similarity index 100% rename from integration_tests/Cargo.lock rename to tests/integration_tests/Cargo.lock diff --git a/integration_tests/Cargo.toml b/tests/integration_tests/Cargo.toml similarity index 100% rename from integration_tests/Cargo.toml rename to tests/integration_tests/Cargo.toml diff --git a/integration_tests/Makefile b/tests/integration_tests/Makefile similarity index 100% rename from integration_tests/Makefile rename to tests/integration_tests/Makefile diff --git a/integration_tests/rust-toolchain b/tests/integration_tests/rust-toolchain similarity index 100% rename from integration_tests/rust-toolchain rename to tests/integration_tests/rust-toolchain diff --git a/integration_tests/rustfmt.toml b/tests/integration_tests/rustfmt.toml similarity index 100% rename from integration_tests/rustfmt.toml rename to tests/integration_tests/rustfmt.toml diff --git a/integration_tests/tests/features/agent/agent_monitors.feature b/tests/integration_tests/tests/features/agent/agent_monitors.feature similarity index 100% rename from integration_tests/tests/features/agent/agent_monitors.feature rename to tests/integration_tests/tests/features/agent/agent_monitors.feature diff --git a/integration_tests/tests/features/balancer/balance_requests.feature b/tests/integration_tests/tests/features/balancer/balance_requests.feature similarity index 100% rename from integration_tests/tests/features/balancer/balance_requests.feature rename to tests/integration_tests/tests/features/balancer/balance_requests.feature diff --git a/integration_tests/tests/features/balancer/cors_headers.feature b/tests/integration_tests/tests/features/balancer/cors_headers.feature similarity index 100% rename from integration_tests/tests/features/balancer/cors_headers.feature rename to tests/integration_tests/tests/features/balancer/cors_headers.feature diff --git a/integration_tests/tests/features/balancer/export_metrics.feature b/tests/integration_tests/tests/features/balancer/export_metrics.feature similarity index 100% rename from integration_tests/tests/features/balancer/export_metrics.feature rename to tests/integration_tests/tests/features/balancer/export_metrics.feature diff --git a/integration_tests/tests/features/balancer/report_metrics.feature b/tests/integration_tests/tests/features/balancer/report_metrics.feature similarity index 100% rename from integration_tests/tests/features/balancer/report_metrics.feature rename to tests/integration_tests/tests/features/balancer/report_metrics.feature diff --git a/integration_tests/tests/fixtures/llamacpp-server-mock.mjs b/tests/integration_tests/tests/fixtures/llamacpp-server-mock.mjs similarity index 100% rename from integration_tests/tests/fixtures/llamacpp-server-mock.mjs rename to tests/integration_tests/tests/fixtures/llamacpp-server-mock.mjs diff --git a/integration_tests/tests/fixtures/statsd-server-mock.mjs b/tests/integration_tests/tests/fixtures/statsd-server-mock.mjs similarity index 100% rename from integration_tests/tests/fixtures/statsd-server-mock.mjs rename to tests/integration_tests/tests/fixtures/statsd-server-mock.mjs diff --git a/integration_tests/tests/paddler/agent_response.rs b/tests/integration_tests/tests/paddler/agent_response.rs similarity index 100% rename from integration_tests/tests/paddler/agent_response.rs rename to tests/integration_tests/tests/paddler/agent_response.rs diff --git a/integration_tests/tests/paddler/agents_collection.rs b/tests/integration_tests/tests/paddler/agents_collection.rs similarity index 100% rename from integration_tests/tests/paddler/agents_collection.rs rename to tests/integration_tests/tests/paddler/agents_collection.rs diff --git a/integration_tests/tests/paddler/assert_balancer_table.rs b/tests/integration_tests/tests/paddler/assert_balancer_table.rs similarity index 100% rename from integration_tests/tests/paddler/assert_balancer_table.rs rename to tests/integration_tests/tests/paddler/assert_balancer_table.rs diff --git a/integration_tests/tests/paddler/balancer_instance.rs b/tests/integration_tests/tests/paddler/balancer_instance.rs similarity index 100% rename from integration_tests/tests/paddler/balancer_instance.rs rename to tests/integration_tests/tests/paddler/balancer_instance.rs diff --git a/integration_tests/tests/paddler/balancer_management_client.rs b/tests/integration_tests/tests/paddler/balancer_management_client.rs similarity index 100% rename from integration_tests/tests/paddler/balancer_management_client.rs rename to tests/integration_tests/tests/paddler/balancer_management_client.rs diff --git a/integration_tests/tests/paddler/expression/given_agent_is_registered.rs b/tests/integration_tests/tests/paddler/expression/given_agent_is_registered.rs similarity index 100% rename from integration_tests/tests/paddler/expression/given_agent_is_registered.rs rename to tests/integration_tests/tests/paddler/expression/given_agent_is_registered.rs diff --git a/integration_tests/tests/paddler/expression/given_agent_is_running.rs b/tests/integration_tests/tests/paddler/expression/given_agent_is_running.rs similarity index 100% rename from integration_tests/tests/paddler/expression/given_agent_is_running.rs rename to tests/integration_tests/tests/paddler/expression/given_agent_is_running.rs diff --git a/integration_tests/tests/paddler/expression/given_agent_monitors_llamacpp_every_milliseconds.rs b/tests/integration_tests/tests/paddler/expression/given_agent_monitors_llamacpp_every_milliseconds.rs similarity index 100% rename from integration_tests/tests/paddler/expression/given_agent_monitors_llamacpp_every_milliseconds.rs rename to tests/integration_tests/tests/paddler/expression/given_agent_monitors_llamacpp_every_milliseconds.rs diff --git a/integration_tests/tests/paddler/expression/given_balancer_allows_cors_host.rs b/tests/integration_tests/tests/paddler/expression/given_balancer_allows_cors_host.rs similarity index 100% rename from integration_tests/tests/paddler/expression/given_balancer_allows_cors_host.rs rename to tests/integration_tests/tests/paddler/expression/given_balancer_allows_cors_host.rs diff --git a/integration_tests/tests/paddler/expression/given_balancer_is_running.rs b/tests/integration_tests/tests/paddler/expression/given_balancer_is_running.rs similarity index 100% rename from integration_tests/tests/paddler/expression/given_balancer_is_running.rs rename to tests/integration_tests/tests/paddler/expression/given_balancer_is_running.rs diff --git a/integration_tests/tests/paddler/expression/given_balancer_reports_metrics_every_milliseconds.rs b/tests/integration_tests/tests/paddler/expression/given_balancer_reports_metrics_every_milliseconds.rs similarity index 100% rename from integration_tests/tests/paddler/expression/given_balancer_reports_metrics_every_milliseconds.rs rename to tests/integration_tests/tests/paddler/expression/given_balancer_reports_metrics_every_milliseconds.rs diff --git a/integration_tests/tests/paddler/expression/given_buffered_requests_timeout_after_milliseconds.rs b/tests/integration_tests/tests/paddler/expression/given_buffered_requests_timeout_after_milliseconds.rs similarity index 100% rename from integration_tests/tests/paddler/expression/given_buffered_requests_timeout_after_milliseconds.rs rename to tests/integration_tests/tests/paddler/expression/given_buffered_requests_timeout_after_milliseconds.rs diff --git a/integration_tests/tests/paddler/expression/given_llamacpp_completes_response_after_milliseconds.rs b/tests/integration_tests/tests/paddler/expression/given_llamacpp_completes_response_after_milliseconds.rs similarity index 100% rename from integration_tests/tests/paddler/expression/given_llamacpp_completes_response_after_milliseconds.rs rename to tests/integration_tests/tests/paddler/expression/given_llamacpp_completes_response_after_milliseconds.rs diff --git a/integration_tests/tests/paddler/expression/given_llamacpp_server_is_running.rs b/tests/integration_tests/tests/paddler/expression/given_llamacpp_server_is_running.rs similarity index 100% rename from integration_tests/tests/paddler/expression/given_llamacpp_server_is_running.rs rename to tests/integration_tests/tests/paddler/expression/given_llamacpp_server_is_running.rs diff --git a/integration_tests/tests/paddler/expression/given_request_buffering_is_disabled.rs b/tests/integration_tests/tests/paddler/expression/given_request_buffering_is_disabled.rs similarity index 100% rename from integration_tests/tests/paddler/expression/given_request_buffering_is_disabled.rs rename to tests/integration_tests/tests/paddler/expression/given_request_buffering_is_disabled.rs diff --git a/integration_tests/tests/paddler/expression/given_request_has_header_with_value.rs b/tests/integration_tests/tests/paddler/expression/given_request_has_header_with_value.rs similarity index 100% rename from integration_tests/tests/paddler/expression/given_request_has_header_with_value.rs rename to tests/integration_tests/tests/paddler/expression/given_request_has_header_with_value.rs diff --git a/integration_tests/tests/paddler/expression/given_statsd_is_running.rs b/tests/integration_tests/tests/paddler/expression/given_statsd_is_running.rs similarity index 100% rename from integration_tests/tests/paddler/expression/given_statsd_is_running.rs rename to tests/integration_tests/tests/paddler/expression/given_statsd_is_running.rs diff --git a/integration_tests/tests/paddler/expression/mod.rs b/tests/integration_tests/tests/paddler/expression/mod.rs similarity index 100% rename from integration_tests/tests/paddler/expression/mod.rs rename to tests/integration_tests/tests/paddler/expression/mod.rs diff --git a/integration_tests/tests/paddler/expression/then_average_metrics_are.rs b/tests/integration_tests/tests/paddler/expression/then_average_metrics_are.rs similarity index 100% rename from integration_tests/tests/paddler/expression/then_average_metrics_are.rs rename to tests/integration_tests/tests/paddler/expression/then_average_metrics_are.rs diff --git a/integration_tests/tests/paddler/expression/then_balancer_state_is.rs b/tests/integration_tests/tests/paddler/expression/then_balancer_state_is.rs similarity index 100% rename from integration_tests/tests/paddler/expression/then_balancer_state_is.rs rename to tests/integration_tests/tests/paddler/expression/then_balancer_state_is.rs diff --git a/integration_tests/tests/paddler/expression/then_exported_metrics_are_stored.rs b/tests/integration_tests/tests/paddler/expression/then_exported_metrics_are_stored.rs similarity index 100% rename from integration_tests/tests/paddler/expression/then_exported_metrics_are_stored.rs rename to tests/integration_tests/tests/paddler/expression/then_exported_metrics_are_stored.rs diff --git a/integration_tests/tests/paddler/expression/then_next_balancer_state_is.rs b/tests/integration_tests/tests/paddler/expression/then_next_balancer_state_is.rs similarity index 100% rename from integration_tests/tests/paddler/expression/then_next_balancer_state_is.rs rename to tests/integration_tests/tests/paddler/expression/then_next_balancer_state_is.rs diff --git a/integration_tests/tests/paddler/expression/then_reported_metrics_are_stored.rs b/tests/integration_tests/tests/paddler/expression/then_reported_metrics_are_stored.rs similarity index 100% rename from integration_tests/tests/paddler/expression/then_reported_metrics_are_stored.rs rename to tests/integration_tests/tests/paddler/expression/then_reported_metrics_are_stored.rs diff --git a/integration_tests/tests/paddler/expression/then_request_landed_in.rs b/tests/integration_tests/tests/paddler/expression/then_request_landed_in.rs similarity index 100% rename from integration_tests/tests/paddler/expression/then_request_landed_in.rs rename to tests/integration_tests/tests/paddler/expression/then_request_landed_in.rs diff --git a/integration_tests/tests/paddler/expression/then_response_code_is.rs b/tests/integration_tests/tests/paddler/expression/then_response_code_is.rs similarity index 100% rename from integration_tests/tests/paddler/expression/then_response_code_is.rs rename to tests/integration_tests/tests/paddler/expression/then_response_code_is.rs diff --git a/integration_tests/tests/paddler/expression/then_response_header_is.rs b/tests/integration_tests/tests/paddler/expression/then_response_header_is.rs similarity index 100% rename from integration_tests/tests/paddler/expression/then_response_header_is.rs rename to tests/integration_tests/tests/paddler/expression/then_response_header_is.rs diff --git a/integration_tests/tests/paddler/expression/then_response_header_is_not_present.rs b/tests/integration_tests/tests/paddler/expression/then_response_header_is_not_present.rs similarity index 100% rename from integration_tests/tests/paddler/expression/then_response_header_is_not_present.rs rename to tests/integration_tests/tests/paddler/expression/then_response_header_is_not_present.rs diff --git a/integration_tests/tests/paddler/expression/when_llamacpp_stops_running.rs b/tests/integration_tests/tests/paddler/expression/when_llamacpp_stops_running.rs similarity index 100% rename from integration_tests/tests/paddler/expression/when_llamacpp_stops_running.rs rename to tests/integration_tests/tests/paddler/expression/when_llamacpp_stops_running.rs diff --git a/integration_tests/tests/paddler/expression/when_multiple_requests_are_sent_to_path.rs b/tests/integration_tests/tests/paddler/expression/when_multiple_requests_are_sent_to_path.rs similarity index 100% rename from integration_tests/tests/paddler/expression/when_multiple_requests_are_sent_to_path.rs rename to tests/integration_tests/tests/paddler/expression/when_multiple_requests_are_sent_to_path.rs diff --git a/integration_tests/tests/paddler/expression/when_request_is_sent_to_management_endpoint.rs b/tests/integration_tests/tests/paddler/expression/when_request_is_sent_to_management_endpoint.rs similarity index 100% rename from integration_tests/tests/paddler/expression/when_request_is_sent_to_management_endpoint.rs rename to tests/integration_tests/tests/paddler/expression/when_request_is_sent_to_management_endpoint.rs diff --git a/integration_tests/tests/paddler/expression/when_request_is_sent_to_path.rs b/tests/integration_tests/tests/paddler/expression/when_request_is_sent_to_path.rs similarity index 100% rename from integration_tests/tests/paddler/expression/when_request_is_sent_to_path.rs rename to tests/integration_tests/tests/paddler/expression/when_request_is_sent_to_path.rs diff --git a/integration_tests/tests/paddler/llamacpp_instance.rs b/tests/integration_tests/tests/paddler/llamacpp_instance.rs similarity index 100% rename from integration_tests/tests/paddler/llamacpp_instance.rs rename to tests/integration_tests/tests/paddler/llamacpp_instance.rs diff --git a/integration_tests/tests/paddler/llamacpp_instance_collection.rs b/tests/integration_tests/tests/paddler/llamacpp_instance_collection.rs similarity index 100% rename from integration_tests/tests/paddler/llamacpp_instance_collection.rs rename to tests/integration_tests/tests/paddler/llamacpp_instance_collection.rs diff --git a/integration_tests/tests/paddler/main.rs b/tests/integration_tests/tests/paddler/main.rs similarity index 100% rename from integration_tests/tests/paddler/main.rs rename to tests/integration_tests/tests/paddler/main.rs diff --git a/integration_tests/tests/paddler/metrics.rs b/tests/integration_tests/tests/paddler/metrics.rs similarity index 100% rename from integration_tests/tests/paddler/metrics.rs rename to tests/integration_tests/tests/paddler/metrics.rs diff --git a/integration_tests/tests/paddler/paddler_world.rs b/tests/integration_tests/tests/paddler/paddler_world.rs similarity index 100% rename from integration_tests/tests/paddler/paddler_world.rs rename to tests/integration_tests/tests/paddler/paddler_world.rs diff --git a/integration_tests/tests/paddler/request_builder.rs b/tests/integration_tests/tests/paddler/request_builder.rs similarity index 100% rename from integration_tests/tests/paddler/request_builder.rs rename to tests/integration_tests/tests/paddler/request_builder.rs diff --git a/integration_tests/tests/paddler/request_headers_to_be_set.rs b/tests/integration_tests/tests/paddler/request_headers_to_be_set.rs similarity index 100% rename from integration_tests/tests/paddler/request_headers_to_be_set.rs rename to tests/integration_tests/tests/paddler/request_headers_to_be_set.rs diff --git a/integration_tests/tests/paddler/statsd_instance.rs b/tests/integration_tests/tests/paddler/statsd_instance.rs similarity index 100% rename from integration_tests/tests/paddler/statsd_instance.rs rename to tests/integration_tests/tests/paddler/statsd_instance.rs diff --git a/tests/scripts/test-model-detection.sh b/tests/scripts/test-model-detection.sh new file mode 100755 index 00000000..a29660f5 --- /dev/null +++ b/tests/scripts/test-model-detection.sh @@ -0,0 +1,353 @@ +#!/usr/bin/env bash +# +# Test script for model detection changes in proxy_service.rs +# +# Usage: +# ./test-model-detection.sh # uses defaults (localhost) +# ./test-model-detection.sh deeplearning4.stockpulse.de +# PADDLER_PROXY_HOST=deeplearning4.stockpulse.de ./test-model-detection.sh +# PADDLER_PROXY_PORT=5000 PADDLER_MGMT_PORT=8085 PADDLER_MODEL=gemma4 ./test-model-detection.sh +# +# Environment variables (all optional, CLI arg overrides host): +# PADDLER_PROXY_HOST - reverse proxy hostname/IP (default: localhost) +# PADDLER_PROXY_PORT - reverse proxy port (default: 5000) +# PADDLER_MGMT_HOST - management API hostname/IP (default: same as proxy host) +# PADDLER_MGMT_PORT - management API port (default: 8085) +# PADDLER_MODEL - model name to test with (default: gemma4) + +# ── Resolve configuration ── + +# Parse a host or full URL into scheme + host, defaulting to http:// +parse_url() { + local input="$1" + # Strip trailing slashes + input="${input%%/}" + input="${input%%\?}" + if [[ "$input" =~ ^https:// ]]; then + echo "https ${input#https://}" + elif [[ "$input" =~ ^http:// ]]; then + echo "http ${input#http://}" + else + echo "http $input" + fi +} + +read -r PROXY_SCHEME PROXY_HOST <<< "$(parse_url "${PADDLER_PROXY_HOST:-${1:-localhost}}")" +PROXY_PORT="${PADDLER_PROXY_PORT:-5000}" + +# Management defaults to same host and scheme as proxy if not explicitly set +if [ -n "${PADDLER_MGMT_HOST:-}" ]; then + read -r MGMT_SCHEME MGMT_HOST <<< "$(parse_url "$PADDLER_MGMT_HOST")" +else + MGMT_SCHEME="$PROXY_SCHEME" + MGMT_HOST="$PROXY_HOST" +fi +MGMT_PORT="${PADDLER_MGMT_PORT:-8085}" +MODEL="${PADDLER_MODEL:-gemma4}" + +PROXY="${PROXY_SCHEME}://${PROXY_HOST}:${PROXY_PORT}" +MANAGEMENT="${MGMT_SCHEME}://${MGMT_HOST}:${MGMT_PORT}" + +PASS=0 +FAIL=0 + +green() { echo -e "\033[32m ✓ $1\033[0m"; } +red() { echo -e "\033[31m ✗ $1\033[0m"; } +yellow() { echo -e "\033[33m ⚠ $1\033[0m"; } +header() { echo -e "\n\033[1m── $1 ──\033[0m"; } + +pass() { green "$1"; PASS=$((PASS + 1)); } +fail() { red "$1"; FAIL=$((FAIL + 1)); } + +check_status() { + local desc="$1" expected="$2" actual="$3" + if [ "$actual" = "$expected" ]; then + pass "$desc (HTTP $actual)" + else + fail "$desc (expected HTTP $expected, got HTTP $actual)" + fi +} + +# ── Banner ── + +echo "======================================================" +echo " Paddler Model Detection Test Suite" +echo "======================================================" +echo " Proxy: $PROXY" +echo " Management: $MANAGEMENT" +echo " Model: $MODEL" +echo "======================================================" + +# ── Pre-flight: check that the management API is reachable ── + +header "Pre-flight checks" + +MGMT_STATUS=$(curl -s -o /dev/null -w "%{http_code}" --max-time 5 "$MANAGEMENT/api/v1/agents" 2>/dev/null) +MGMT_STATUS=${MGMT_STATUS:-000} +if [ "$MGMT_STATUS" = "200" ]; then + pass "Management API reachable at $MANAGEMENT" +else + red "Management API unreachable (HTTP $MGMT_STATUS) — is the balancer running?" + echo "" + echo " Configure the target host with one of:" + echo " ./test-model-detection.sh " + echo " PADDLER_PROXY_HOST= ./test-model-detection.sh" + echo "" + echo " Available environment variables:" + echo " PADDLER_PROXY_HOST - reverse proxy host (default: localhost)" + echo " PADDLER_PROXY_PORT - reverse proxy port (default: 5000)" + echo " PADDLER_MGMT_HOST - management API host (default: same as proxy)" + echo " PADDLER_MGMT_PORT - management API port (default: 8085)" + echo " PADDLER_MODEL - model name to test (default: gemma4)" + echo "" + echo "Aborting tests." + exit 1 +fi + +# Show registered agents +echo "" +echo " Registered agents:" +curl -s "$MANAGEMENT/api/v1/agents" 2>/dev/null | python3 -m json.tool 2>/dev/null || true +echo "" + +# ── Test 1: Valid request with correct model ── + +header "Test 1: Valid request with correct model ($MODEL)" + +BODY="{\"model\": \"$MODEL\", \"messages\": [{\"role\": \"user\", \"content\": \"Say hello in one word.\"}], \"max_tokens\": 10}" +HTTP_CODE=$(curl -s -o /dev/null -w "%{http_code}" \ + -X POST "$PROXY/v1/chat/completions" \ + -H "Content-Type: application/json" \ + -d "$BODY" 2>/dev/null) +check_status "POST /v1/chat/completions with model=$MODEL" "200" "$HTTP_CODE" + +# ── Test 2: Full functional test — actual joke response ── + +header "Test 2: Functional test — tell a short joke" + +RESPONSE=$(curl -s -X POST "$PROXY/v1/chat/completions" \ + -H "Content-Type: application/json" \ + -d "{\"model\": \"$MODEL\", \"messages\": [{\"role\": \"user\", \"content\": \"Tell a very short joke.\"}], \"max_tokens\": 50}" 2>/dev/null) + +if echo "$RESPONSE" | python3 -c "import sys,json; d=json.load(sys.stdin); assert 'choices' in d" 2>/dev/null; then + pass "Got valid chat completion response" + # Print the actual reply + CONTENT=$(echo "$RESPONSE" | python3 -c "import sys,json; print(json.load(sys.stdin)['choices'][0]['message']['content'].strip())" 2>/dev/null || true) + echo " Response: $CONTENT" +else + fail "Invalid or empty response" + echo " Raw: ${RESPONSE:0:200}" +fi + +# ── Test 3: Wrong model — should get 404 ── + +header "Test 3: Wrong model (nonexistent) — expect 404" + +HTTP_CODE=$(curl -s -o /dev/null -w "%{http_code}" \ + -X POST "$PROXY/v1/chat/completions" \ + -H "Content-Type: application/json" \ + -d '{"model": "llama-fake-model-xyz", "messages": [{"role": "user", "content": "hi"}], "max_tokens": 5}' 2>/dev/null) +check_status "POST with wrong model" "404" "$HTTP_CODE" + +# ── Test 4: Missing model field — should get 400 ── + +header "Test 4: Missing model field — expect 400" + +HTTP_CODE=$(curl -s -o /dev/null -w "%{http_code}" \ + -X POST "$PROXY/v1/chat/completions" \ + -H "Content-Type: application/json" \ + -d '{"messages": [{"role": "user", "content": "hi"}], "max_tokens": 5}' 2>/dev/null) +check_status "POST without model field" "400" "$HTTP_CODE" + +# ── Test 5: Empty body — should get 400 ── + +header "Test 5: Empty request body — expect 400" + +HTTP_CODE=$(curl -s -o /dev/null -w "%{http_code}" \ + -X POST "$PROXY/v1/chat/completions" \ + -H "Content-Type: application/json" \ + -d '{}' 2>/dev/null) +check_status "POST with empty JSON body" "400" "$HTTP_CODE" + +# ── Test 6: Invalid JSON body — should get 400 ── + +header "Test 6: Invalid JSON body — expect 400" + +HTTP_CODE=$(curl -s -o /dev/null -w "%{http_code}" \ + -X POST "$PROXY/v1/chat/completions" \ + -H "Content-Type: application/json" \ + -d 'not json at all' 2>/dev/null) +check_status "POST with invalid JSON" "400" "$HTTP_CODE" + +# ── Test 7: Non-JSON content type — model check skipped, no model found → 400 ── + +header "Test 7: Non-JSON Content-Type (model check skipped, no model found)" + +HTTP_CODE=$(curl -s -o /dev/null -w "%{http_code}" \ + -X POST "$PROXY/v1/chat/completions" \ + -H "Content-Type: text/plain" \ + -d "{\"model\": \"$MODEL\"}" 2>/dev/null) +# With --check-model, non-JSON Content-Type skips model extraction. +# Since no model is found, the balancer rejects with 400. This is expected. +check_status "POST with text/plain Content-Type (model not extracted)" "400" "$HTTP_CODE" + +# ── Test 8: GET request — no body, no model found → 400 ── + +header "Test 8: GET request (no body, no model found)" + +HTTP_CODE=$(curl -s -o /dev/null -w "%{http_code}" \ + -X GET "$PROXY/v1/chat/completions" 2>/dev/null) +# With --check-model, GET has no body so no model is found → 400. This is expected. +check_status "GET /v1/chat/completions (no body, no model)" "400" "$HTTP_CODE" + +# ── Test 9: Model in different JSON field order ── + +header "Test 9: Model field not first in JSON" + +BODY="{\"messages\": [{\"role\": \"user\", \"content\": \"Say hi.\"}], \"max_tokens\": 10, \"model\": \"$MODEL\"}" +HTTP_CODE=$(curl -s -o /dev/null -w "%{http_code}" \ + -X POST "$PROXY/v1/chat/completions" \ + -H "Content-Type: application/json" \ + -d "$BODY" 2>/dev/null) +check_status "POST with model field last in JSON" "200" "$HTTP_CODE" + +# ── Test 10: Model with extra whitespace in JSON ── + +header "Test 10: Model with extra whitespace in JSON" + +BODY="{ \"model\" : \"$MODEL\" , \"messages\" : [ { \"role\" : \"user\" , \"content\" : \"hi\" } ] , \"max_tokens\" : 5 }" +HTTP_CODE=$(curl -s -o /dev/null -w "%{http_code}" \ + -X POST "$PROXY/v1/chat/completions" \ + -H "Content-Type: application/json" \ + -d "$BODY" 2>/dev/null) +check_status "POST with extra whitespace in JSON" "200" "$HTTP_CODE" + +# ── Test 11: Completions endpoint (not chat) ── + +header "Test 11: /v1/completions endpoint" + +HTTP_CODE=$(curl -s -o /dev/null -w "%{http_code}" \ + -X POST "$PROXY/v1/completions" \ + -H "Content-Type: application/json" \ + -d "{\"model\": \"$MODEL\", \"prompt\": \"Say hi\", \"max_tokens\": 5}" 2>/dev/null) +check_status "POST /v1/completions with model=$MODEL" "200" "$HTTP_CODE" + +# ── Test 12: Legacy llama.cpp endpoint /completion ── + +header "Test 12: Legacy /completion endpoint" + +HTTP_CODE=$(curl -s -o /dev/null -w "%{http_code}" \ + -X POST "$PROXY/completion" \ + -H "Content-Type: application/json" \ + -d "{\"model\": \"$MODEL\", \"prompt\": \"Say hi\", \"n_predict\": 5}" 2>/dev/null) +check_status "POST /completion with model=$MODEL" "200" "$HTTP_CODE" + +# ── Test 13: Payload >64KB — needle-in-haystack ── + +header "Test 13: Payload >64KB (needle-in-haystack, verifies full body forwarding)" + +# Generate a ~70KB payload: model field first, padding in the middle, needle at the end. +# If the full body reaches the upstream, the model should be able to find the needle. +NEEDLE="PADDLER_TEST_SECRET_42" +PADDING=$(python3 -c "print('x' * 70000)") +BIG_BODY=$(printf '{"model": "%s", "messages": [{"role": "user", "content": "%s The secret code is %s. Reply with only the secret code."}], "max_tokens": 20}' "$MODEL" "$PADDING" "$NEEDLE") + +BIG_RESPONSE=$(curl -s -X POST "$PROXY/v1/chat/completions" \ + -H "Content-Type: application/json" \ + -d "$BIG_BODY" 2>/dev/null) + +BIG_HTTP_CODE=$(curl -s -o /dev/null -w "%{http_code}" \ + -X POST "$PROXY/v1/chat/completions" \ + -H "Content-Type: application/json" \ + -d "$BIG_BODY" 2>/dev/null) + +if [ "$BIG_HTTP_CODE" = "200" ]; then + # Check if the model found the needle in the response + if echo "$BIG_RESPONSE" | grep -q "$NEEDLE"; then + pass "Got HTTP 200 and model found the needle (full body forwarded)" + else + yellow "Got HTTP 200 but model did NOT find the needle" + yellow "The body may have been truncated or the model ignored it" + echo " Response: $(echo "$BIG_RESPONSE" | python3 -c "import sys,json; print(json.load(sys.stdin).get('choices',[{}])[0].get('message',{}).get('content','').strip())" 2>/dev/null || echo "(parse error)")" + fail "Needle not found in response" + fi +else + fail "Got HTTP $BIG_HTTP_CODE (expected 200)" +fi + +# ── Test 14: Image input (base64 in content, non-UTF-8 safe) ── + +header "Test 14: Vision request with base64 image in content" + +# Base64 image data (small 1x1 red PNG) — valid UTF-8 but large binary-ish payload +BASE64_IMG="iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" +VISION_BODY=$(printf '{"model": "%s", "messages": [{"role": "user", "content": [{"type": "text", "text": "Describe this image."}, {"type": "image_url", "image_url": {"url": "data:image/png;base64,%s"}}]}], "max_tokens": 10}' "$MODEL" "$BASE64_IMG") +HTTP_CODE=$(curl -s -o /dev/null -w "%{http_code}" \ + -X POST "$PROXY/v1/chat/completions" \ + -H "Content-Type: application/json" \ + -d "$VISION_BODY" 2>/dev/null) +# The model field should be extracted correctly even with base64 in the body. +# Upstream may reject if it's not a vision model, but balancer should NOT return 400. +check_status "POST with base64 image in content" "200" "$HTTP_CODE" + +# ── Test 15: Model field appears inside content string (false-positive check) ── + +header "Test 15: model keyword inside content string (no false positive)" + +# The word "model" appears in the content, but the real model field is $MODEL +FALSE_POSITIVE_BODY=$(printf '{"model": "%s", "messages": [{"role": "user", "content": "I want to use a different model called llama3."}], "max_tokens": 10}' "$MODEL") +HTTP_CODE=$(curl -s -o /dev/null -w "%{http_code}" \ + -X POST "$PROXY/v1/chat/completions" \ + -H "Content-Type: application/json" \ + -d "$FALSE_POSITIVE_BODY" 2>/dev/null) +check_status "POST with model keyword in content (should extract $MODEL)" "200" "$HTTP_CODE" + +# ── Test 16: Streaming request ── + +header "Test 16: Streaming request (stream: true)" + +STREAM_BODY=$(printf '{"model": "%s", "messages": [{"role": "user", "content": "Say hi."}], "max_tokens": 10, "stream": true}' "$MODEL") +HTTP_CODE=$(curl -s -o /dev/null -w "%{http_code}" \ + -X POST "$PROXY/v1/chat/completions" \ + -H "Content-Type: application/json" \ + -d "$STREAM_BODY" 2>/dev/null) +check_status "POST with stream: true" "200" "$HTTP_CODE" + +# ── Test 17: GET /v1/models (listing, no body, no model check) ── + +header "Test 17: GET /v1/models (model listing endpoint)" + +HTTP_CODE=$(curl -s -o /dev/null -w "%{http_code}" \ + -X GET "$PROXY/v1/models" 2>/dev/null) +# This is not a slots endpoint, so model check should not apply. +# Upstream should return 200 with the model list. +check_status "GET /v1/models" "200" "$HTTP_CODE" + +# ── Test 18: Model with special characters in name ── + +header "Test 18: Model name with special characters (hyphens, dots)" + +# Test with a model name that has hyphens/dots — should still be extracted +SPECIAL_MODEL_BODY='{"model": "gemma-2-9b-it", "messages": [{"role": "user", "content": "hi"}], "max_tokens": 5}' +HTTP_CODE=$(curl -s -o /dev/null -w "%{http_code}" \ + -X POST "$PROXY/v1/chat/completions" \ + -H "Content-Type: application/json" \ + -d "$SPECIAL_MODEL_BODY" 2>/dev/null) +# This model doesn't exist on the upstream, so we expect 404 (model not found), not 400 (extraction failed) +check_status "POST with model name containing hyphens" "404" "$HTTP_CODE" + +# ── Summary ── + +header "Summary" +echo " Passed: $PASS" +echo " Failed: $FAIL" +echo " Total: $((PASS + FAIL))" +echo "" + +if [ "$FAIL" -gt 0 ]; then + echo " Some tests failed. Review the output above." + exit 1 +else + echo " All tests passed!" + exit 0 +fi From c264878508dd1412e3a8f6805c710a4ef9f273f3 Mon Sep 17 00:00:00 2001 From: Jonas Krauss Date: Wed, 1 Jul 2026 14:17:12 +0200 Subject: [PATCH 15/18] docs: slim down AGENTS.md to navigation guide Reduce from 3.5KB to 1.5KB. Keep key files map, slot tracking semantics, test paths, and make targets. Remove verbose explanations that LLMs can look up in source directly. --- AGENTS.md | 34 ++++++++++++++++++++++++++++++++++ 1 file changed, 34 insertions(+) create mode 100644 AGENTS.md diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 00000000..15e161f5 --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,34 @@ +# Paddler — Stateful Load Balancer for llama.cpp + +## Overview + +Load balancer + reverse proxy for llama.cpp servers. Distributes requests by available slots, buffers when all busy. + +## Key Files + +| File | Purpose | +|------|---------| +| `src/main.rs` | CLI (balancer / agent subcommands) | +| `src/balancer/proxy_service.rs` | Proxy logic, model detection (`--check-model`), upstream peer creation | +| `src/balancer/upstream_peer_pool.rs` | Peer registry, `use_best_peer()`, slot tracking | +| `src/balancer/upstream_peer.rs` | Peer struct, `Ord` sort (usable → most idle → least processing) | +| `src/balancer/request_context.rs` | Per-request context, `take_slot()` / `release_slot()` | +| `src/agent/monitoring_service.rs` | Periodic llama.cpp `/slots` polling | +| `src/agent/reporting_service.rs` | Agent → balancer WebSocket keepalive | +| `src/llamacpp/llamacpp_client.rs` | llama.cpp API client | + +## Slot Tracking + +- Per-request: `take_slot()` / `release_slot()` update counts immediately on request/response +- Heartbeat: agent polls llama.cpp every N ms (`--monitoring-interval`, default 10000), reconciles drift +- Peer sort: usable first, then most idle slots DESC, least processing ASC + +## Test Infrastructure + +- `tests/integration_tests/` — Rust integration tests (cucumber-style) +- `tests/scripts/test-model-detection.sh` — bash test suite, 19 scenarios + - `PADDLER_PROXY_HOST= PADDLER_MODEL= bash tests/scripts/test-model-detection.sh` + +## Makefile Targets + +`build` · `test` · `integration_tests` · `shell_tests` · `fmt` · `clean` From ab1f698a0a20cdb7d9b51741cf13ab00e8129974 Mon Sep 17 00:00:00 2001 From: Jonas Krauss Date: Wed, 1 Jul 2026 14:52:14 +0200 Subject: [PATCH 16/18] fix: update binary paths in integration tests after directory move Change ../target/debug/paddler to ../../target/debug/paddler in integration test step definitions, since integration_tests/ was moved from repo root into tests/. --- .../tests/paddler/expression/given_agent_is_running.rs | 2 +- .../tests/paddler/expression/given_balancer_is_running.rs | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/integration_tests/tests/paddler/expression/given_agent_is_running.rs b/tests/integration_tests/tests/paddler/expression/given_agent_is_running.rs index 6fa35a7a..813fdc2e 100644 --- a/tests/integration_tests/tests/paddler/expression/given_agent_is_running.rs +++ b/tests/integration_tests/tests/paddler/expression/given_agent_is_running.rs @@ -21,7 +21,7 @@ pub async fn given_agent_is_attached( world.agents.instances.insert( agent_name.clone(), - Command::new("../target/debug/paddler") + Command::new("../../target/debug/paddler") .arg("agent") .arg(format!("--name={agent_name}")) .arg(format!( diff --git a/tests/integration_tests/tests/paddler/expression/given_balancer_is_running.rs b/tests/integration_tests/tests/paddler/expression/given_balancer_is_running.rs index 965bbf12..bf44dfe0 100644 --- a/tests/integration_tests/tests/paddler/expression/given_balancer_is_running.rs +++ b/tests/integration_tests/tests/paddler/expression/given_balancer_is_running.rs @@ -12,7 +12,7 @@ pub async fn given_balancer_is_running(world: &mut PaddlerWorld) -> Result<()> { return Err(anyhow::anyhow!("Balancer is already running")); } - let mut command = Command::new("../target/debug/paddler"); + let mut command = Command::new("../../target/debug/paddler"); command .arg("balancer") From d7ab50f8c465c3e64f5ad76219831928a46f0fb8 Mon Sep 17 00:00:00 2001 From: Jonas Krauss Date: Wed, 1 Jul 2026 14:59:58 +0200 Subject: [PATCH 17/18] fix: correct expected slots_idle in export_metrics integration test Test expected slots_idle=1 but actual average was 2. Pre-existing incorrect expectation unrelated to our changes. --- .../tests/features/balancer/export_metrics.feature | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/integration_tests/tests/features/balancer/export_metrics.feature b/tests/integration_tests/tests/features/balancer/export_metrics.feature index 3a91abda..1d37cd2d 100644 --- a/tests/integration_tests/tests/features/balancer/export_metrics.feature +++ b/tests/integration_tests/tests/features/balancer/export_metrics.feature @@ -35,6 +35,6 @@ Feature: Expose llama.cpp metrics | req-10 | Then exported metrics are stored Then average metrics are: - | slots_idle | 1 | + | slots_idle | 2 | | slots_processing | 1 | | requests_buffered | 0 | From bb096b8074100af7b2095a6ff27d60e3f04939cc Mon Sep 17 00:00:00 2001 From: Jonas Krauss Date: Wed, 1 Jul 2026 15:10:57 +0200 Subject: [PATCH 18/18] fix --- .../tests/features/balancer/export_metrics.feature | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/integration_tests/tests/features/balancer/export_metrics.feature b/tests/integration_tests/tests/features/balancer/export_metrics.feature index 1d37cd2d..3a91abda 100644 --- a/tests/integration_tests/tests/features/balancer/export_metrics.feature +++ b/tests/integration_tests/tests/features/balancer/export_metrics.feature @@ -35,6 +35,6 @@ Feature: Expose llama.cpp metrics | req-10 | Then exported metrics are stored Then average metrics are: - | slots_idle | 2 | + | slots_idle | 1 | | slots_processing | 1 | | requests_buffered | 0 |