Compare commits
9 Commits
feat/70-ro
...
feat/74-ou
| Author | SHA1 | Date | |
|---|---|---|---|
|
1115bb0942
|
|||
|
63f578cb15
|
|||
|
76c90fa993
|
|||
|
7984d27553
|
|||
|
43ffffdccb
|
|||
|
5fd7736abd
|
|||
|
03fd4960c3
|
|||
|
5ed6bc3390
|
|||
|
cabec1d08a
|
52
Cargo.lock
generated
52
Cargo.lock
generated
@@ -800,6 +800,7 @@ dependencies = [
|
|||||||
"cortex-core",
|
"cortex-core",
|
||||||
"eventsource-stream",
|
"eventsource-stream",
|
||||||
"futures",
|
"futures",
|
||||||
|
"helexa-stream",
|
||||||
"metrics",
|
"metrics",
|
||||||
"metrics-exporter-prometheus",
|
"metrics-exporter-prometheus",
|
||||||
"reqwest",
|
"reqwest",
|
||||||
@@ -1928,18 +1929,37 @@ version = "0.1.16"
|
|||||||
dependencies = [
|
dependencies = [
|
||||||
"anyhow",
|
"anyhow",
|
||||||
"axum",
|
"axum",
|
||||||
|
"chrono",
|
||||||
"clap",
|
"clap",
|
||||||
"cortex-core",
|
"cortex-core",
|
||||||
"figment",
|
"figment",
|
||||||
|
"helexa-stream",
|
||||||
|
"rcgen",
|
||||||
"reqwest",
|
"reqwest",
|
||||||
|
"rustls",
|
||||||
"serde",
|
"serde",
|
||||||
"serde_json",
|
"serde_json",
|
||||||
|
"thiserror 2.0.18",
|
||||||
"tokio",
|
"tokio",
|
||||||
|
"tokio-rustls",
|
||||||
"tower-http",
|
"tower-http",
|
||||||
"tracing",
|
"tracing",
|
||||||
"tracing-subscriber",
|
"tracing-subscriber",
|
||||||
]
|
]
|
||||||
|
|
||||||
|
[[package]]
|
||||||
|
name = "helexa-stream"
|
||||||
|
version = "0.1.16"
|
||||||
|
dependencies = [
|
||||||
|
"async-stream",
|
||||||
|
"axum",
|
||||||
|
"futures",
|
||||||
|
"reqwest",
|
||||||
|
"thiserror 2.0.18",
|
||||||
|
"tokio",
|
||||||
|
"tokio-stream",
|
||||||
|
]
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "hermit-abi"
|
name = "hermit-abi"
|
||||||
version = "0.5.2"
|
version = "0.5.2"
|
||||||
@@ -2980,6 +3000,16 @@ dependencies = [
|
|||||||
"syn",
|
"syn",
|
||||||
]
|
]
|
||||||
|
|
||||||
|
[[package]]
|
||||||
|
name = "pem"
|
||||||
|
version = "3.0.6"
|
||||||
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
|
checksum = "1d30c53c26bc5b31a98cd02d20f25a7c8567146caf63ed593a9d87b2775291be"
|
||||||
|
dependencies = [
|
||||||
|
"base64 0.22.1",
|
||||||
|
"serde_core",
|
||||||
|
]
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "percent-encoding"
|
name = "percent-encoding"
|
||||||
version = "2.3.2"
|
version = "2.3.2"
|
||||||
@@ -3346,6 +3376,19 @@ dependencies = [
|
|||||||
"crossbeam-utils",
|
"crossbeam-utils",
|
||||||
]
|
]
|
||||||
|
|
||||||
|
[[package]]
|
||||||
|
name = "rcgen"
|
||||||
|
version = "0.13.2"
|
||||||
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
|
checksum = "75e669e5202259b5314d1ea5397316ad400819437857b90861765f24c4cf80a2"
|
||||||
|
dependencies = [
|
||||||
|
"pem",
|
||||||
|
"ring",
|
||||||
|
"rustls-pki-types",
|
||||||
|
"time",
|
||||||
|
"yasna",
|
||||||
|
]
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "reborrow"
|
name = "reborrow"
|
||||||
version = "0.5.5"
|
version = "0.5.5"
|
||||||
@@ -5233,6 +5276,15 @@ version = "1.0.1"
|
|||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
checksum = "cfe53a6657fd280eaa890a3bc59152892ffa3e30101319d168b781ed6529b049"
|
checksum = "cfe53a6657fd280eaa890a3bc59152892ffa3e30101319d168b781ed6529b049"
|
||||||
|
|
||||||
|
[[package]]
|
||||||
|
name = "yasna"
|
||||||
|
version = "0.5.2"
|
||||||
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
|
checksum = "e17bb3549cc1321ae1296b9cdc2698e2b6cb1992adfa19a8c72e5b7a738f44cd"
|
||||||
|
dependencies = [
|
||||||
|
"time",
|
||||||
|
]
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "yoke"
|
name = "yoke"
|
||||||
version = "0.7.5"
|
version = "0.7.5"
|
||||||
|
|||||||
@@ -8,6 +8,7 @@ members = [
|
|||||||
"crates/helexa-acp",
|
"crates/helexa-acp",
|
||||||
"crates/helexa-bench",
|
"crates/helexa-bench",
|
||||||
"crates/helexa-router",
|
"crates/helexa-router",
|
||||||
|
"crates/helexa-stream",
|
||||||
]
|
]
|
||||||
|
|
||||||
[workspace.package]
|
[workspace.package]
|
||||||
|
|||||||
@@ -6,6 +6,7 @@ license.workspace = true
|
|||||||
|
|
||||||
[dependencies]
|
[dependencies]
|
||||||
cortex-core.workspace = true
|
cortex-core.workspace = true
|
||||||
|
helexa-stream = { path = "../helexa-stream" }
|
||||||
async-trait.workspace = true
|
async-trait.workspace = true
|
||||||
tokio.workspace = true
|
tokio.workspace = true
|
||||||
axum.workspace = true
|
axum.workspace = true
|
||||||
|
|||||||
@@ -1,21 +1,27 @@
|
|||||||
//! Streaming HTTP reverse proxy to neuron backends.
|
//! Streaming HTTP reverse proxy to neuron backends.
|
||||||
//!
|
//!
|
||||||
//! For streaming requests, SSE chunks are forwarded as they arrive.
|
//! The streaming *mechanism* — forward an SSE body chunk-for-chunk without
|
||||||
//! The proxy captures timing information for metrics but does not
|
//! buffering, observing the bytes for metrics — lives in the shared
|
||||||
//! buffer the full response.
|
//! [`helexa_stream`] crate (#71), so cortex and helexa-router use one
|
||||||
|
//! implementation. This module supplies cortex's *policy*: the
|
||||||
|
//! [`CortexMetrics`] observer (per-request token metrics + per-principal
|
||||||
|
//! reservation settle), cortex's logging contract, and the cortex error
|
||||||
|
//! envelope. The usage-extraction helper is re-exported from the shared
|
||||||
|
//! crate so existing call sites keep working.
|
||||||
|
|
||||||
use crate::router::RouteDecision;
|
use crate::router::RouteDecision;
|
||||||
use anyhow::Result;
|
use axum::http::HeaderMap;
|
||||||
use axum::body::Body;
|
use axum::http::StatusCode;
|
||||||
use axum::http::{HeaderMap, StatusCode};
|
|
||||||
use axum::response::{IntoResponse, Response};
|
use axum::response::{IntoResponse, Response};
|
||||||
use futures::Stream;
|
use helexa_stream::{BodyTail, ChunkObserver, StreamError};
|
||||||
use futures::stream::BoxStream;
|
|
||||||
use reqwest::Client;
|
use reqwest::Client;
|
||||||
use std::pin::Pin;
|
|
||||||
use std::task::{Context, Poll};
|
|
||||||
use std::time::Instant;
|
use std::time::Instant;
|
||||||
|
|
||||||
|
/// Re-export the shared usage-extraction helper. Several cortex modules
|
||||||
|
/// (`handlers`, `anthropic_sse`) pull token counts out of a buffered body
|
||||||
|
/// tail via this function; it lives in `helexa-stream` now.
|
||||||
|
pub use helexa_stream::last_count_for;
|
||||||
|
|
||||||
/// Proxy a request body to the resolved backend node and stream the response.
|
/// Proxy a request body to the resolved backend node and stream the response.
|
||||||
///
|
///
|
||||||
/// Logging contract: every call emits exactly one structured event at
|
/// Logging contract: every call emits exactly one structured event at
|
||||||
@@ -42,66 +48,41 @@ pub async fn forward_request(
|
|||||||
"proxying request"
|
"proxying request"
|
||||||
);
|
);
|
||||||
|
|
||||||
let mut req_builder = client.post(&url).body(body);
|
let observer = CortexMetrics::new(model_id, &route.node_name, request_start, usage_sink);
|
||||||
|
|
||||||
// Forward relevant headers.
|
let response = helexa_stream::forward_streaming(client, &url, headers, body, observer)
|
||||||
for (key, value) in headers.iter() {
|
.await
|
||||||
if key == "host" || key == "content-length" {
|
.map_err(|e| {
|
||||||
continue; // reqwest sets these
|
match &e {
|
||||||
}
|
StreamError::Upstream(err) => tracing::warn!(
|
||||||
req_builder = req_builder.header(key, value);
|
node = %route.node_name,
|
||||||
}
|
url = %url,
|
||||||
|
error = %err,
|
||||||
|
"proxy: upstream request failed (network)"
|
||||||
|
),
|
||||||
|
StreamError::ResponseBuild(err) => tracing::warn!(
|
||||||
|
node = %route.node_name,
|
||||||
|
url = %url,
|
||||||
|
error = %err,
|
||||||
|
"proxy: failed to build response"
|
||||||
|
),
|
||||||
|
}
|
||||||
|
ProxyError::from(e)
|
||||||
|
})?;
|
||||||
|
|
||||||
let upstream_resp = match req_builder.send().await {
|
if !response.status().is_success() {
|
||||||
Ok(r) => r,
|
|
||||||
Err(e) => {
|
|
||||||
tracing::warn!(
|
|
||||||
node = %route.node_name,
|
|
||||||
url = %url,
|
|
||||||
error = %e,
|
|
||||||
"proxy: upstream request failed (network)"
|
|
||||||
);
|
|
||||||
return Err(ProxyError::Upstream(e));
|
|
||||||
}
|
|
||||||
};
|
|
||||||
|
|
||||||
let upstream_status = upstream_resp.status();
|
|
||||||
if !upstream_status.is_success() {
|
|
||||||
// Streaming body — can't snippet without breaking the stream
|
// Streaming body — can't snippet without breaking the stream
|
||||||
// pass-through. Log status + URL; the client still gets the
|
// pass-through. Log status + URL; the client still gets the
|
||||||
// upstream status, just without the leaked body.
|
// upstream status, just without the leaked body.
|
||||||
tracing::warn!(
|
tracing::warn!(
|
||||||
node = %route.node_name,
|
node = %route.node_name,
|
||||||
url = %url,
|
url = %url,
|
||||||
status = upstream_status.as_u16(),
|
status = response.status().as_u16(),
|
||||||
"proxy: upstream returned non-2xx"
|
"proxy: upstream returned non-2xx"
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
let status = StatusCode::from_u16(upstream_status.as_u16()).unwrap_or(StatusCode::BAD_GATEWAY);
|
Ok(response)
|
||||||
|
|
||||||
let resp_headers = upstream_resp.headers().clone();
|
|
||||||
let stream = TokenMetricsStream::new(
|
|
||||||
Box::pin(upstream_resp.bytes_stream()),
|
|
||||||
TokenMetrics::new(model_id, &route.node_name, request_start, usage_sink),
|
|
||||||
);
|
|
||||||
|
|
||||||
let body = Body::from_stream(stream);
|
|
||||||
|
|
||||||
let mut response = Response::builder().status(status);
|
|
||||||
for (key, value) in resp_headers.iter() {
|
|
||||||
response = response.header(key, value);
|
|
||||||
}
|
|
||||||
|
|
||||||
response.body(body).map_err(|e| {
|
|
||||||
tracing::warn!(
|
|
||||||
node = %route.node_name,
|
|
||||||
url = %url,
|
|
||||||
error = %e,
|
|
||||||
"proxy: failed to build response"
|
|
||||||
);
|
|
||||||
ProxyError::ResponseBuild(e.to_string())
|
|
||||||
})
|
|
||||||
}
|
}
|
||||||
|
|
||||||
#[derive(Debug, thiserror::Error)]
|
#[derive(Debug, thiserror::Error)]
|
||||||
@@ -112,6 +93,15 @@ pub enum ProxyError {
|
|||||||
ResponseBuild(String),
|
ResponseBuild(String),
|
||||||
}
|
}
|
||||||
|
|
||||||
|
impl From<StreamError> for ProxyError {
|
||||||
|
fn from(e: StreamError) -> Self {
|
||||||
|
match e {
|
||||||
|
StreamError::Upstream(err) => ProxyError::Upstream(err),
|
||||||
|
StreamError::ResponseBuild(msg) => ProxyError::ResponseBuild(msg),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
impl IntoResponse for ProxyError {
|
impl IntoResponse for ProxyError {
|
||||||
fn into_response(self) -> Response {
|
fn into_response(self) -> Response {
|
||||||
let (status, code, message) = match &self {
|
let (status, code, message) = match &self {
|
||||||
@@ -139,9 +129,10 @@ impl IntoResponse for ProxyError {
|
|||||||
//
|
//
|
||||||
// The proxy never buffers or re-serialises the upstream body — chunks
|
// The proxy never buffers or re-serialises the upstream body — chunks
|
||||||
// are forwarded verbatim. For metrics it observes each chunk's arrival
|
// are forwarded verbatim. For metrics it observes each chunk's arrival
|
||||||
// time and keeps a bounded tail of the body text, from which the final
|
// time and keeps a bounded tail of the body text (via the shared
|
||||||
// OpenAI `usage` object (present on the last SSE chunk and on
|
// `helexa_stream::BodyTail`), from which the final OpenAI `usage` object
|
||||||
// non-streaming JSON bodies alike) yields engine-truth token counts.
|
// (present on the last SSE chunk and on non-streaming JSON bodies alike)
|
||||||
|
// yields engine-truth token counts.
|
||||||
//
|
//
|
||||||
// Emitted per request, labelled {model, node}:
|
// Emitted per request, labelled {model, node}:
|
||||||
// cortex_time_to_first_token_seconds (histogram) — first body chunk
|
// cortex_time_to_first_token_seconds (histogram) — first body chunk
|
||||||
@@ -155,37 +146,15 @@ impl IntoResponse for ProxyError {
|
|||||||
/// non-streaming bodies.
|
/// non-streaming bodies.
|
||||||
const TAIL_CAP_BYTES: usize = 64 * 1024;
|
const TAIL_CAP_BYTES: usize = 64 * 1024;
|
||||||
|
|
||||||
/// Find the value of the LAST `"key": <integer>` occurrence in `tail`.
|
/// cortex's [`ChunkObserver`]: per-request token metrics plus the
|
||||||
/// Pure and chunk-boundary-safe (the tail is contiguous appended text).
|
/// per-principal reservation settle. Drives cortex policy over the shared
|
||||||
/// The quoted-needle form means `completion_tokens` never matches
|
/// streaming mechanism.
|
||||||
/// `completion_tokens_details`.
|
struct CortexMetrics {
|
||||||
pub(crate) fn last_count_for(tail: &str, key: &str) -> Option<u64> {
|
|
||||||
let needle = format!("\"{key}\"");
|
|
||||||
let mut result = None;
|
|
||||||
for (idx, _) in tail.match_indices(&needle) {
|
|
||||||
let rest = tail[idx + needle.len()..].trim_start();
|
|
||||||
let Some(rest) = rest.strip_prefix(':') else {
|
|
||||||
continue;
|
|
||||||
};
|
|
||||||
let rest = rest.trim_start();
|
|
||||||
let digits: &str = &rest[..rest
|
|
||||||
.char_indices()
|
|
||||||
.find(|(_, c)| !c.is_ascii_digit())
|
|
||||||
.map(|(i, _)| i)
|
|
||||||
.unwrap_or(rest.len())];
|
|
||||||
if let Ok(v) = digits.parse::<u64>() {
|
|
||||||
result = Some(v);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
result
|
|
||||||
}
|
|
||||||
|
|
||||||
struct TokenMetrics {
|
|
||||||
labels: [(&'static str, String); 2],
|
labels: [(&'static str, String); 2],
|
||||||
request_start: Instant,
|
request_start: Instant,
|
||||||
first_chunk: Option<Instant>,
|
first_chunk: Option<Instant>,
|
||||||
last_chunk: Option<Instant>,
|
last_chunk: Option<Instant>,
|
||||||
tail: String,
|
tail: BodyTail,
|
||||||
finished: bool,
|
finished: bool,
|
||||||
/// Per-principal metering hook (#51). Invoked exactly once in `finish`
|
/// Per-principal metering hook (#51). Invoked exactly once in `finish`
|
||||||
/// with the observed `(prompt, completion)` so the reservation can be
|
/// with the observed `(prompt, completion)` so the reservation can be
|
||||||
@@ -193,7 +162,7 @@ struct TokenMetrics {
|
|||||||
usage_sink: Option<crate::metering::UsageSink>,
|
usage_sink: Option<crate::metering::UsageSink>,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl TokenMetrics {
|
impl CortexMetrics {
|
||||||
fn new(
|
fn new(
|
||||||
model_id: &str,
|
model_id: &str,
|
||||||
node_name: &str,
|
node_name: &str,
|
||||||
@@ -208,26 +177,19 @@ impl TokenMetrics {
|
|||||||
request_start,
|
request_start,
|
||||||
first_chunk: None,
|
first_chunk: None,
|
||||||
last_chunk: None,
|
last_chunk: None,
|
||||||
tail: String::new(),
|
tail: BodyTail::new(TAIL_CAP_BYTES),
|
||||||
finished: false,
|
finished: false,
|
||||||
usage_sink,
|
usage_sink,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl ChunkObserver for CortexMetrics {
|
||||||
fn observe(&mut self, chunk: &[u8]) {
|
fn observe(&mut self, chunk: &[u8]) {
|
||||||
let now = Instant::now();
|
let now = Instant::now();
|
||||||
self.first_chunk.get_or_insert(now);
|
self.first_chunk.get_or_insert(now);
|
||||||
self.last_chunk = Some(now);
|
self.last_chunk = Some(now);
|
||||||
self.tail.push_str(&String::from_utf8_lossy(chunk));
|
self.tail.push(chunk);
|
||||||
if self.tail.len() > TAIL_CAP_BYTES {
|
|
||||||
// Keep the newest half; the usage object is always at the
|
|
||||||
// very end of the body. Split at a char boundary.
|
|
||||||
let mut cut = self.tail.len() - TAIL_CAP_BYTES / 2;
|
|
||||||
while !self.tail.is_char_boundary(cut) {
|
|
||||||
cut += 1;
|
|
||||||
}
|
|
||||||
self.tail.drain(..cut);
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Emit the metrics exactly once — called on clean stream end and
|
/// Emit the metrics exactly once — called on clean stream end and
|
||||||
@@ -239,8 +201,8 @@ impl TokenMetrics {
|
|||||||
}
|
}
|
||||||
self.finished = true;
|
self.finished = true;
|
||||||
|
|
||||||
let prompt = last_count_for(&self.tail, "prompt_tokens");
|
let prompt = last_count_for(self.tail.as_str(), "prompt_tokens");
|
||||||
let completion = last_count_for(&self.tail, "completion_tokens");
|
let completion = last_count_for(self.tail.as_str(), "completion_tokens");
|
||||||
|
|
||||||
// Per-model metrics — only when body chunks actually arrived.
|
// Per-model metrics — only when body chunks actually arrived.
|
||||||
if let Some(first) = self.first_chunk {
|
if let Some(first) = self.first_chunk {
|
||||||
@@ -280,97 +242,3 @@ impl TokenMetrics {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Pass-through stream wrapper that feeds [`TokenMetrics`]. Emits on
|
|
||||||
/// clean end-of-stream; the Drop impl covers client disconnects.
|
|
||||||
struct TokenMetricsStream {
|
|
||||||
inner: BoxStream<'static, Result<bytes::Bytes, reqwest::Error>>,
|
|
||||||
metrics: TokenMetrics,
|
|
||||||
}
|
|
||||||
|
|
||||||
impl TokenMetricsStream {
|
|
||||||
fn new(
|
|
||||||
inner: BoxStream<'static, Result<bytes::Bytes, reqwest::Error>>,
|
|
||||||
metrics: TokenMetrics,
|
|
||||||
) -> Self {
|
|
||||||
Self { inner, metrics }
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
impl Stream for TokenMetricsStream {
|
|
||||||
type Item = Result<bytes::Bytes, reqwest::Error>;
|
|
||||||
|
|
||||||
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
|
|
||||||
let this = self.get_mut();
|
|
||||||
match this.inner.as_mut().poll_next(cx) {
|
|
||||||
Poll::Ready(Some(Ok(chunk))) => {
|
|
||||||
this.metrics.observe(&chunk);
|
|
||||||
Poll::Ready(Some(Ok(chunk)))
|
|
||||||
}
|
|
||||||
Poll::Ready(Some(Err(e))) => Poll::Ready(Some(Err(e))),
|
|
||||||
Poll::Ready(None) => {
|
|
||||||
this.metrics.finish();
|
|
||||||
Poll::Ready(None)
|
|
||||||
}
|
|
||||||
Poll::Pending => Poll::Pending,
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
impl Drop for TokenMetricsStream {
|
|
||||||
fn drop(&mut self) {
|
|
||||||
self.metrics.finish();
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
#[cfg(test)]
|
|
||||||
mod tests {
|
|
||||||
use super::last_count_for;
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn extracts_counts_from_final_sse_usage_chunk() {
|
|
||||||
let tail = concat!(
|
|
||||||
"data: {\"choices\":[{\"delta\":{\"content\":\"hi\"}}]}\n\n",
|
|
||||||
"data: {\"choices\":[],\"usage\":{\"prompt_tokens\":225,",
|
|
||||||
"\"completion_tokens\":42,\"total_tokens\":267}}\n\n",
|
|
||||||
"data: [DONE]\n\n"
|
|
||||||
);
|
|
||||||
assert_eq!(last_count_for(tail, "prompt_tokens"), Some(225));
|
|
||||||
assert_eq!(last_count_for(tail, "completion_tokens"), Some(42));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn extracts_counts_from_non_streaming_body() {
|
|
||||||
let tail = "{\"choices\":[{\"message\":{\"content\":\"hi\"}}],\
|
|
||||||
\"usage\":{\"prompt_tokens\": 12, \"completion_tokens\": 7}}";
|
|
||||||
assert_eq!(last_count_for(tail, "prompt_tokens"), Some(12));
|
|
||||||
assert_eq!(last_count_for(tail, "completion_tokens"), Some(7));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn ignores_details_variants_and_takes_last_occurrence() {
|
|
||||||
// completion_tokens_details must not shadow completion_tokens,
|
|
||||||
// and the LAST usage object wins (matters when content echoes
|
|
||||||
// a usage-shaped string earlier in the stream).
|
|
||||||
let tail = concat!(
|
|
||||||
"data: {\"usage\":{\"completion_tokens\":1}}\n\n",
|
|
||||||
"data: {\"usage\":{\"completion_tokens\":99,",
|
|
||||||
"\"completion_tokens_details\":{\"reasoning_tokens\":3}}}\n\n"
|
|
||||||
);
|
|
||||||
assert_eq!(last_count_for(tail, "completion_tokens"), Some(99));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn absent_keys_yield_none() {
|
|
||||||
assert_eq!(
|
|
||||||
last_count_for("data: [DONE]\n\n", "completion_tokens"),
|
|
||||||
None
|
|
||||||
);
|
|
||||||
assert_eq!(last_count_for("", "prompt_tokens"), None);
|
|
||||||
// key present but non-numeric value
|
|
||||||
assert_eq!(
|
|
||||||
last_count_for("\"completion_tokens\": null", "completion_tokens"),
|
|
||||||
None
|
|
||||||
);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|||||||
@@ -15,6 +15,7 @@ path = "src/lib.rs"
|
|||||||
|
|
||||||
[dependencies]
|
[dependencies]
|
||||||
cortex-core = { workspace = true }
|
cortex-core = { workspace = true }
|
||||||
|
helexa-stream = { path = "../helexa-stream" }
|
||||||
|
|
||||||
tokio = { workspace = true }
|
tokio = { workspace = true }
|
||||||
axum = { workspace = true }
|
axum = { workspace = true }
|
||||||
@@ -24,10 +25,17 @@ serde = { workspace = true }
|
|||||||
serde_json = { workspace = true }
|
serde_json = { workspace = true }
|
||||||
figment = { workspace = true }
|
figment = { workspace = true }
|
||||||
anyhow = { workspace = true }
|
anyhow = { workspace = true }
|
||||||
|
thiserror = { workspace = true }
|
||||||
clap = { workspace = true }
|
clap = { workspace = true }
|
||||||
tracing = { workspace = true }
|
tracing = { workspace = true }
|
||||||
tracing-subscriber = { workspace = true }
|
tracing-subscriber = { workspace = true }
|
||||||
|
chrono = { workspace = true }
|
||||||
|
|
||||||
[dev-dependencies]
|
[dev-dependencies]
|
||||||
# Jail (isolated cwd + env) for config tests.
|
# Jail (isolated cwd + env) for config tests.
|
||||||
figment = { workspace = true, features = ["test"] }
|
figment = { workspace = true, features = ["test"] }
|
||||||
|
# Self-signed cert generation + a minimal HTTPS server for the outbound
|
||||||
|
# TLS-pinning tests (#74).
|
||||||
|
rcgen = "0.13"
|
||||||
|
rustls = "0.23"
|
||||||
|
tokio-rustls = "0.26"
|
||||||
|
|||||||
243
crates/helexa-router/src/catalogue.rs
Normal file
243
crates/helexa-router/src/catalogue.rs
Normal file
@@ -0,0 +1,243 @@
|
|||||||
|
//! Federation catalogue (#75) — the router's aggregate `/v1/models`.
|
||||||
|
//!
|
||||||
|
//! Presents the **deduped union** of every reachable cortex's `/v1/models`
|
||||||
|
//! as the router's own catalogue, so an opencode client doing discovery
|
||||||
|
//! against the router resolves the whole federation without knowing about
|
||||||
|
//! operators or cortexes (resolves #61's "Router/discovery contract").
|
||||||
|
//!
|
||||||
|
//! Re-tiering: the fractal design is neuron ← cortex ← router. At the
|
||||||
|
//! router tier the "nodes" are **cortexes**, so the merged entry's
|
||||||
|
//! `feasible_on` / `locations` are rewritten to **operator names**, not the
|
||||||
|
//! neuron names a cortex reports. That keeps the federation view honest
|
||||||
|
//! ("served by these operators") without leaking each operator's internal
|
||||||
|
//! topology (neuron names, per-device VRAM) to end users.
|
||||||
|
//!
|
||||||
|
//! Conflict resolution when operators advertise the same model with
|
||||||
|
//! different enrichment:
|
||||||
|
//! - **`limit`** → the *tightest* (smallest `context`), so a client never
|
||||||
|
//! overflows the most-constrained operator that might serve it (same rule
|
||||||
|
//! cortex uses across its neurons).
|
||||||
|
//! - **`cost`** → the *cheapest* (lowest input, then output), the
|
||||||
|
//! federation "from" price. Richer policy (a range, region/price-aware
|
||||||
|
//! selection) couples to #68 and is left as a follow-up.
|
||||||
|
|
||||||
|
use crate::state::{CortexTopology, entry_feasible};
|
||||||
|
use cortex_core::harness::{ModelCost, ModelLimit};
|
||||||
|
use cortex_core::node::{CortexModelEntry, ModelLocation, ModelStatus};
|
||||||
|
use std::collections::HashMap;
|
||||||
|
|
||||||
|
/// Build the federation catalogue: the deduped union of every reachable
|
||||||
|
/// cortex's serveable models, merged across operators and sorted by id.
|
||||||
|
pub fn aggregate_models(topology: &HashMap<String, CortexTopology>) -> Vec<CortexModelEntry> {
|
||||||
|
// Iterate cortexes in name order so `feasible_on` / `locations` and the
|
||||||
|
// limit/cost tie-breaks are deterministic regardless of map ordering.
|
||||||
|
let mut cortexes: Vec<(&String, &CortexTopology)> = topology.iter().collect();
|
||||||
|
cortexes.sort_by(|a, b| a.0.cmp(b.0));
|
||||||
|
|
||||||
|
let mut merged: HashMap<String, CortexModelEntry> = HashMap::new();
|
||||||
|
for (cortex_name, t) in cortexes {
|
||||||
|
if !t.reachable {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
for entry in t.models.values() {
|
||||||
|
// Only surface models the cortex can actually serve — a
|
||||||
|
// catalogue-only entry no neuron can host shouldn't appear in
|
||||||
|
// the federation view.
|
||||||
|
if !entry_feasible(entry) {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
merged
|
||||||
|
.entry(entry.id.clone())
|
||||||
|
.and_modify(|acc| merge_into(acc, cortex_name, entry))
|
||||||
|
.or_insert_with(|| router_entry(cortex_name, entry));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
let mut out: Vec<CortexModelEntry> = merged.into_values().collect();
|
||||||
|
out.sort_by(|a, b| a.id.cmp(&b.id));
|
||||||
|
out
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Seed a federation entry from the first cortex that serves the model,
|
||||||
|
/// re-tiering `feasible_on` / `locations` to the operator name.
|
||||||
|
fn router_entry(cortex: &str, e: &CortexModelEntry) -> CortexModelEntry {
|
||||||
|
CortexModelEntry {
|
||||||
|
id: e.id.clone(),
|
||||||
|
object: "model".into(),
|
||||||
|
created: e.created,
|
||||||
|
owned_by: e.owned_by.clone(),
|
||||||
|
loaded: e.loaded,
|
||||||
|
feasible_on: vec![cortex.to_string()],
|
||||||
|
locations: loaded_location(cortex, e),
|
||||||
|
capabilities: e.capabilities.clone(),
|
||||||
|
limit: e.limit.clone(),
|
||||||
|
cost: e.cost.clone(),
|
||||||
|
tool_call: e.tool_call,
|
||||||
|
reasoning: e.reasoning,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Fold another cortex's view of the same model into the merged entry.
|
||||||
|
fn merge_into(acc: &mut CortexModelEntry, cortex: &str, e: &CortexModelEntry) {
|
||||||
|
acc.loaded |= e.loaded;
|
||||||
|
acc.feasible_on.push(cortex.to_string());
|
||||||
|
acc.locations.extend(loaded_location(cortex, e));
|
||||||
|
for cap in &e.capabilities {
|
||||||
|
if !acc.capabilities.contains(cap) {
|
||||||
|
acc.capabilities.push(cap.clone());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
acc.tool_call |= e.tool_call;
|
||||||
|
acc.reasoning |= e.reasoning;
|
||||||
|
acc.limit = tightest_limit(acc.limit.take(), e.limit.clone());
|
||||||
|
acc.cost = cheapest_cost(acc.cost.take(), e.cost.clone());
|
||||||
|
}
|
||||||
|
|
||||||
|
/// A single cortex-tier location when the model is loaded at that operator;
|
||||||
|
/// empty when only cold-loadable. Neuron-level VRAM is deliberately dropped.
|
||||||
|
fn loaded_location(cortex: &str, e: &CortexModelEntry) -> Vec<ModelLocation> {
|
||||||
|
if e.loaded {
|
||||||
|
vec![ModelLocation {
|
||||||
|
node: cortex.to_string(),
|
||||||
|
status: ModelStatus::Loaded,
|
||||||
|
vram_estimate_mb: None,
|
||||||
|
}]
|
||||||
|
} else {
|
||||||
|
Vec::new()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Smaller `context` wins — never advertise more headroom than the
|
||||||
|
/// most-constrained operator can honour.
|
||||||
|
fn tightest_limit(a: Option<ModelLimit>, b: Option<ModelLimit>) -> Option<ModelLimit> {
|
||||||
|
match (a, b) {
|
||||||
|
(None, x) | (x, None) => x,
|
||||||
|
(Some(a), Some(b)) => Some(if b.context < a.context { b } else { a }),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Cheapest by (input, output) price — the federation "from" price.
|
||||||
|
fn cheapest_cost(a: Option<ModelCost>, b: Option<ModelCost>) -> Option<ModelCost> {
|
||||||
|
match (a, b) {
|
||||||
|
(None, x) | (x, None) => x,
|
||||||
|
(Some(a), Some(b)) => Some(if (b.input, b.output) < (a.input, a.output) {
|
||||||
|
b
|
||||||
|
} else {
|
||||||
|
a
|
||||||
|
}),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
use crate::state::CortexTopology;
|
||||||
|
|
||||||
|
fn entry(id: &str, loaded: bool, feasible: bool) -> CortexModelEntry {
|
||||||
|
CortexModelEntry {
|
||||||
|
id: id.into(),
|
||||||
|
object: "model".into(),
|
||||||
|
created: 0,
|
||||||
|
owned_by: "helexa".into(),
|
||||||
|
loaded,
|
||||||
|
feasible_on: if feasible || loaded {
|
||||||
|
vec!["some-neuron".into()]
|
||||||
|
} else {
|
||||||
|
vec![]
|
||||||
|
},
|
||||||
|
locations: vec![],
|
||||||
|
capabilities: vec![],
|
||||||
|
limit: None,
|
||||||
|
cost: None,
|
||||||
|
tool_call: false,
|
||||||
|
reasoning: false,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn cortex(reachable: bool, entries: Vec<CortexModelEntry>) -> CortexTopology {
|
||||||
|
CortexTopology {
|
||||||
|
reachable,
|
||||||
|
consecutive_failures: 0,
|
||||||
|
last_poll: None,
|
||||||
|
healthy_nodes: 1,
|
||||||
|
total_nodes: 1,
|
||||||
|
models: entries.into_iter().map(|e| (e.id.clone(), e)).collect(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn dedupes_and_merges_availability_across_cortexes() {
|
||||||
|
let mut topo = HashMap::new();
|
||||||
|
// c-a: model loaded. c-b: same model only cold-loadable.
|
||||||
|
topo.insert("c-a".into(), cortex(true, vec![entry("m", true, true)]));
|
||||||
|
topo.insert("c-b".into(), cortex(true, vec![entry("m", false, true)]));
|
||||||
|
|
||||||
|
let out = aggregate_models(&topo);
|
||||||
|
assert_eq!(out.len(), 1, "duplicate model id collapses to one");
|
||||||
|
let m = &out[0];
|
||||||
|
assert!(m.loaded, "loaded somewhere → loaded");
|
||||||
|
// feasible_on re-tiered to operator names, both present, sorted.
|
||||||
|
assert_eq!(m.feasible_on, vec!["c-a".to_string(), "c-b".to_string()]);
|
||||||
|
// Only the loaded operator contributes a location, named by operator.
|
||||||
|
assert_eq!(m.locations.len(), 1);
|
||||||
|
assert_eq!(m.locations[0].node, "c-a");
|
||||||
|
assert_eq!(m.locations[0].vram_estimate_mb, None);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn unreachable_cortex_is_excluded() {
|
||||||
|
let mut topo = HashMap::new();
|
||||||
|
topo.insert("up".into(), cortex(true, vec![entry("m", true, true)]));
|
||||||
|
topo.insert(
|
||||||
|
"down".into(),
|
||||||
|
cortex(false, vec![entry("other", true, true)]),
|
||||||
|
);
|
||||||
|
let out = aggregate_models(&topo);
|
||||||
|
assert_eq!(out.len(), 1);
|
||||||
|
assert_eq!(out[0].id, "m");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn catalogue_only_infeasible_entries_are_hidden() {
|
||||||
|
let mut topo = HashMap::new();
|
||||||
|
topo.insert("c".into(), cortex(true, vec![entry("ghost", false, false)]));
|
||||||
|
assert!(aggregate_models(&topo).is_empty());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn preserves_tightest_limit_and_cheapest_cost() {
|
||||||
|
let mut a = entry("m", true, true);
|
||||||
|
a.limit = Some(ModelLimit {
|
||||||
|
context: 32_768,
|
||||||
|
input: None,
|
||||||
|
output: 4096,
|
||||||
|
});
|
||||||
|
a.cost = Some(ModelCost {
|
||||||
|
input: 0.50,
|
||||||
|
output: 1.50,
|
||||||
|
cache_read: None,
|
||||||
|
cache_write: None,
|
||||||
|
});
|
||||||
|
let mut b = entry("m", true, true);
|
||||||
|
b.limit = Some(ModelLimit {
|
||||||
|
context: 16_384, // tighter
|
||||||
|
input: None,
|
||||||
|
output: 4096,
|
||||||
|
});
|
||||||
|
b.cost = Some(ModelCost {
|
||||||
|
input: 0.20, // cheaper
|
||||||
|
output: 0.80,
|
||||||
|
cache_read: None,
|
||||||
|
cache_write: None,
|
||||||
|
});
|
||||||
|
|
||||||
|
let mut topo = HashMap::new();
|
||||||
|
topo.insert("c-a".into(), cortex(true, vec![a]));
|
||||||
|
topo.insert("c-b".into(), cortex(true, vec![b]));
|
||||||
|
|
||||||
|
let out = aggregate_models(&topo);
|
||||||
|
assert_eq!(out.len(), 1);
|
||||||
|
assert_eq!(out[0].limit.as_ref().unwrap().context, 16_384);
|
||||||
|
assert_eq!(out[0].cost.as_ref().unwrap().input, 0.20);
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -27,17 +27,50 @@ pub struct RouterSettings {
|
|||||||
/// of the router (see #69's TLS posture). The router never owns an
|
/// of the router (see #69's TLS posture). The router never owns an
|
||||||
/// inbound TLS listener.
|
/// inbound TLS listener.
|
||||||
pub listen: String,
|
pub listen: String,
|
||||||
|
/// How often (seconds) the background poller refreshes each cortex's
|
||||||
|
/// health + `/v1/models` topology (#72). Defaults to 10s, matching the
|
||||||
|
/// cortex↔neuron poll cadence one tier down.
|
||||||
|
#[serde(default = "default_poll_interval_secs")]
|
||||||
|
pub poll_interval_secs: u64,
|
||||||
|
/// This router instance's region (e.g. "eu-west"). When set, dispatch
|
||||||
|
/// (#73) prefers cortexes whose `region` matches, before falling back to
|
||||||
|
/// any feasible cortex. `None` → no geo affinity.
|
||||||
|
#[serde(default)]
|
||||||
|
pub region: Option<String>,
|
||||||
|
}
|
||||||
|
|
||||||
|
fn default_poll_interval_secs() -> u64 {
|
||||||
|
10
|
||||||
}
|
}
|
||||||
|
|
||||||
/// One downstream cortex the router may proxy to. The router verifies the
|
/// One downstream cortex the router may proxy to. The router verifies the
|
||||||
/// cortex's outbound TLS cert (#74) and routes on capacity (#73); it holds
|
/// cortex's outbound TLS cert (#74) and routes on capacity (#73); it holds
|
||||||
/// no entitlement logic of its own and forwards the client bearer verbatim.
|
/// no entitlement logic of its own and forwards the client bearer verbatim.
|
||||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||||
pub struct CortexEndpoint {
|
pub struct CortexEndpoint {
|
||||||
/// Human-readable label (e.g. "lair-cafe").
|
/// Human-readable label (e.g. "lair-cafe").
|
||||||
pub name: String,
|
pub name: String,
|
||||||
/// Base URL of the cortex gateway (e.g. "https://cortex.example.com").
|
/// Base URL of the cortex gateway (e.g. "https://cortex.example.com").
|
||||||
pub endpoint: String,
|
pub endpoint: String,
|
||||||
|
/// Optional region tag (e.g. "eu-west") for geo affinity in dispatch
|
||||||
|
/// (#73). `None` → no region preference applies to this cortex.
|
||||||
|
#[serde(default)]
|
||||||
|
pub region: Option<String>,
|
||||||
|
/// Path to a PEM trust anchor that **enrols** this cortex (#74): the
|
||||||
|
/// expected CA (or self-signed cert) the cortex's TLS cert must chain
|
||||||
|
/// to. When set on an `https://` endpoint, the router builds a client
|
||||||
|
/// that trusts **only** this anchor (platform roots disabled), so the
|
||||||
|
/// outbound router→cortex hop — which carries the client's bearer —
|
||||||
|
/// reaches a cert the router was told to expect, and a rogue endpoint
|
||||||
|
/// presenting any other (even publicly-valid) cert is rejected at the
|
||||||
|
/// TLS handshake. A rejected handshake surfaces as a connection error,
|
||||||
|
/// which the poller (#72) already treats as unreachable → excluded.
|
||||||
|
///
|
||||||
|
/// `None` → standard platform-root validation (use for cortexes behind
|
||||||
|
/// a publicly-trusted cert, or plaintext `http://` on a private network
|
||||||
|
/// where the WireGuard mesh is the trust boundary).
|
||||||
|
#[serde(default)]
|
||||||
|
pub tls_ca: Option<String>,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl RouterConfig {
|
impl RouterConfig {
|
||||||
@@ -58,6 +91,8 @@ impl Default for RouterConfig {
|
|||||||
Self {
|
Self {
|
||||||
router: RouterSettings {
|
router: RouterSettings {
|
||||||
listen: "0.0.0.0:8088".into(),
|
listen: "0.0.0.0:8088".into(),
|
||||||
|
poll_interval_secs: default_poll_interval_secs(),
|
||||||
|
region: None,
|
||||||
},
|
},
|
||||||
cortexes: vec![],
|
cortexes: vec![],
|
||||||
}
|
}
|
||||||
|
|||||||
221
crates/helexa-router/src/dispatch.rs
Normal file
221
crates/helexa-router/src/dispatch.rs
Normal file
@@ -0,0 +1,221 @@
|
|||||||
|
//! Capacity-aware dispatch (#73) — the router's data path.
|
||||||
|
//!
|
||||||
|
//! Given an inbound request's `model`, pick a reachable cortex that can
|
||||||
|
//! serve it (preferring warm/loaded, region-affine, higher-headroom),
|
||||||
|
//! forward the client's bearer **unchanged** (auth stays at cortex), and
|
||||||
|
//! stream the response back verbatim via the shared [`helexa_stream`]
|
||||||
|
//! module. Cortex's #63-shaped rejections (`429 rate_limit_exceeded`,
|
||||||
|
//! `400 context_length_exceeded`, …) pass through untouched. Transport
|
||||||
|
//! failures fail over to the next feasible cortex; a genuine HTTP response —
|
||||||
|
//! any status — is returned as-is and never retried away.
|
||||||
|
//!
|
||||||
|
//! The router holds **no entitlement logic**: it routes on capacity, not
|
||||||
|
//! budget.
|
||||||
|
|
||||||
|
use crate::config::CortexEndpoint;
|
||||||
|
use crate::error::envelope_response;
|
||||||
|
use crate::state::RouterState;
|
||||||
|
use axum::body::Bytes;
|
||||||
|
use axum::http::HeaderMap;
|
||||||
|
use axum::response::Response;
|
||||||
|
use cortex_core::error_envelope::OpenAiError;
|
||||||
|
use helexa_stream::{ChunkObserver, StreamError};
|
||||||
|
use std::cmp::Reverse;
|
||||||
|
use std::collections::HashMap;
|
||||||
|
|
||||||
|
/// Retry-After hint (seconds) on the router's own transient rejections.
|
||||||
|
const RETRY_AFTER_SECS: u64 = 5;
|
||||||
|
|
||||||
|
/// Outcome of choosing where to send a request.
|
||||||
|
#[derive(Debug, PartialEq, Eq)]
|
||||||
|
pub enum Selection {
|
||||||
|
/// Feasible reachable cortexes, best-first (failover order).
|
||||||
|
Candidates(Vec<CortexEndpoint>),
|
||||||
|
/// Some cortex knows the model but none are reachable right now → 503.
|
||||||
|
NoReachableCapacity,
|
||||||
|
/// No configured cortex serves the model at all → 404.
|
||||||
|
UnknownModel,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Rank the reachable cortexes that can serve `model`, best-first.
|
||||||
|
///
|
||||||
|
/// Ordering (each a tie-break for the next): loaded/warm before cold-loadable
|
||||||
|
/// · region match before not · more healthy nodes before fewer · name for
|
||||||
|
/// determinism.
|
||||||
|
pub async fn select_cortexes(state: &RouterState, model: &str) -> Selection {
|
||||||
|
let topo = state.topology.read().await;
|
||||||
|
let by_name: HashMap<&str, &CortexEndpoint> = state
|
||||||
|
.cortexes
|
||||||
|
.iter()
|
||||||
|
.map(|c| (c.name.as_str(), c))
|
||||||
|
.collect();
|
||||||
|
|
||||||
|
let mut ranked: Vec<Ranked> = Vec::new();
|
||||||
|
let mut known_anywhere = false;
|
||||||
|
|
||||||
|
for (name, t) in topo.iter() {
|
||||||
|
let Some(entry) = t.models.get(model) else {
|
||||||
|
continue;
|
||||||
|
};
|
||||||
|
if !crate::state::entry_feasible(entry) {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
// Known even via an unreachable cortex's last-good poll — lets us
|
||||||
|
// tell "temporarily down" (503) from "nobody serves it" (404).
|
||||||
|
known_anywhere = true;
|
||||||
|
if !t.reachable {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
let Some(ep) = by_name.get(name.as_str()) else {
|
||||||
|
continue;
|
||||||
|
};
|
||||||
|
let region_match = match (&state.region, &ep.region) {
|
||||||
|
(Some(r), Some(cr)) => r == cr,
|
||||||
|
_ => false,
|
||||||
|
};
|
||||||
|
ranked.push(Ranked {
|
||||||
|
loaded: entry.loaded,
|
||||||
|
region_match,
|
||||||
|
healthy_nodes: t.healthy_nodes,
|
||||||
|
endpoint: (*ep).clone(),
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
if ranked.is_empty() {
|
||||||
|
return if known_anywhere {
|
||||||
|
Selection::NoReachableCapacity
|
||||||
|
} else {
|
||||||
|
Selection::UnknownModel
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
ranked.sort_by(|a, b| {
|
||||||
|
// false < true, so negate the "good" booleans to sort good first.
|
||||||
|
(
|
||||||
|
!a.loaded,
|
||||||
|
!a.region_match,
|
||||||
|
Reverse(a.healthy_nodes),
|
||||||
|
&a.endpoint.name,
|
||||||
|
)
|
||||||
|
.cmp(&(
|
||||||
|
!b.loaded,
|
||||||
|
!b.region_match,
|
||||||
|
Reverse(b.healthy_nodes),
|
||||||
|
&b.endpoint.name,
|
||||||
|
))
|
||||||
|
});
|
||||||
|
|
||||||
|
Selection::Candidates(ranked.into_iter().map(|r| r.endpoint).collect())
|
||||||
|
}
|
||||||
|
|
||||||
|
struct Ranked {
|
||||||
|
loaded: bool,
|
||||||
|
region_match: bool,
|
||||||
|
healthy_nodes: u32,
|
||||||
|
endpoint: CortexEndpoint,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Proxy an inbound inference request to a capacity-bearing cortex.
|
||||||
|
///
|
||||||
|
/// `path` is the inference path to forward to (same on the cortex, e.g.
|
||||||
|
/// `/v1/chat/completions`). The body is parsed only to extract `model`.
|
||||||
|
pub async fn dispatch(
|
||||||
|
state: &RouterState,
|
||||||
|
path: &str,
|
||||||
|
headers: HeaderMap,
|
||||||
|
body: Bytes,
|
||||||
|
) -> Response {
|
||||||
|
let Some(model) = extract_model(&body) else {
|
||||||
|
return envelope_response(OpenAiError::new(
|
||||||
|
400,
|
||||||
|
"invalid_request_error",
|
||||||
|
"missing_model_field",
|
||||||
|
"missing 'model' field in request body",
|
||||||
|
));
|
||||||
|
};
|
||||||
|
|
||||||
|
let candidates = match select_cortexes(state, &model).await {
|
||||||
|
Selection::Candidates(c) => c,
|
||||||
|
Selection::UnknownModel => {
|
||||||
|
return envelope_response(
|
||||||
|
OpenAiError::new(
|
||||||
|
404,
|
||||||
|
"invalid_request_error",
|
||||||
|
"model_not_found",
|
||||||
|
format!("no operator serves model '{model}'"),
|
||||||
|
)
|
||||||
|
.with_param("model"),
|
||||||
|
);
|
||||||
|
}
|
||||||
|
Selection::NoReachableCapacity => {
|
||||||
|
return envelope_response(OpenAiError::service_unavailable(
|
||||||
|
format!("model '{model}' is temporarily unavailable on all operators"),
|
||||||
|
Some(RETRY_AFTER_SECS),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
// Try candidates in order, failing over only on transport errors. A
|
||||||
|
// genuine HTTP response (any status — including cortex's #63 429/400)
|
||||||
|
// is returned verbatim and never retried away.
|
||||||
|
for ep in &candidates {
|
||||||
|
// A candidate whose pinned TLS client failed to build (#74) is
|
||||||
|
// disabled — skip it and fail over, same as an unreachable cortex.
|
||||||
|
let Some(client) = state.client_for(&ep.name) else {
|
||||||
|
tracing::warn!(cortex = %ep.name, "no TLS client (disabled); skipping candidate");
|
||||||
|
continue;
|
||||||
|
};
|
||||||
|
let url = format!("{}{}", ep.endpoint, path);
|
||||||
|
tracing::info!(cortex = %ep.name, url = %url, model = %model, "dispatching");
|
||||||
|
match helexa_stream::forward_streaming(
|
||||||
|
client,
|
||||||
|
&url,
|
||||||
|
headers.clone(),
|
||||||
|
body.clone(),
|
||||||
|
NoopObserver,
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
{
|
||||||
|
Ok(resp) => return resp,
|
||||||
|
Err(StreamError::Upstream(e)) => {
|
||||||
|
tracing::warn!(
|
||||||
|
cortex = %ep.name,
|
||||||
|
url = %url,
|
||||||
|
error = %e,
|
||||||
|
"cortex unreachable; failing over"
|
||||||
|
);
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
Err(StreamError::ResponseBuild(msg)) => {
|
||||||
|
tracing::error!(cortex = %ep.name, error = %msg, "failed to build proxied response");
|
||||||
|
return envelope_response(OpenAiError::without_code(
|
||||||
|
500,
|
||||||
|
"api_error",
|
||||||
|
"failed to build proxied response",
|
||||||
|
));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Every feasible cortex failed to connect.
|
||||||
|
tracing::warn!(model = %model, tried = candidates.len(), "all feasible operators unreachable");
|
||||||
|
envelope_response(OpenAiError::service_unavailable(
|
||||||
|
format!("all operators able to serve '{model}' are unreachable"),
|
||||||
|
Some(RETRY_AFTER_SECS),
|
||||||
|
))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Pull the `model` field out of a request body without re-serialising it.
|
||||||
|
fn extract_model(body: &Bytes) -> Option<String> {
|
||||||
|
let v: serde_json::Value = serde_json::from_slice(body).ok()?;
|
||||||
|
v.get("model")?.as_str().map(str::to_string)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// The router proxies bytes verbatim and keeps no per-request policy, so it
|
||||||
|
/// needs no observation hooks. (Token metrics/metering stay at cortex.)
|
||||||
|
struct NoopObserver;
|
||||||
|
|
||||||
|
impl ChunkObserver for NoopObserver {
|
||||||
|
fn observe(&mut self, _chunk: &[u8]) {}
|
||||||
|
fn finish(&mut self) {}
|
||||||
|
}
|
||||||
27
crates/helexa-router/src/error.rs
Normal file
27
crates/helexa-router/src/error.rs
Normal file
@@ -0,0 +1,27 @@
|
|||||||
|
//! Router adapter from the shared, axum-agnostic
|
||||||
|
//! [`cortex_core::error_envelope::OpenAiError`] (#60/#63) to an axum
|
||||||
|
//! [`Response`], setting `Retry-After` when the envelope carries one.
|
||||||
|
//!
|
||||||
|
//! cortex-core owns the envelope shape; this is the only place the router
|
||||||
|
//! crosses from that data into axum. Mirrors cortex-gateway's adapter so
|
||||||
|
//! the router's own rejections (no feasible operator, all unreachable) are
|
||||||
|
//! the same #63-shaped envelopes clients already understand — distinct from
|
||||||
|
//! cortex's rejections, which the router proxies through verbatim.
|
||||||
|
|
||||||
|
use axum::http::{HeaderValue, StatusCode, header};
|
||||||
|
use axum::response::{IntoResponse, Json, Response};
|
||||||
|
use cortex_core::error_envelope::OpenAiError;
|
||||||
|
|
||||||
|
/// Render an [`OpenAiError`] as an axum response (status + JSON envelope +
|
||||||
|
/// optional `Retry-After`).
|
||||||
|
pub fn envelope_response(err: OpenAiError) -> Response {
|
||||||
|
let status = StatusCode::from_u16(err.status).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
|
||||||
|
let retry_after = err.retry_after_secs;
|
||||||
|
let mut response = (status, Json(err.body())).into_response();
|
||||||
|
if let Some(secs) = retry_after
|
||||||
|
&& let Ok(value) = HeaderValue::from_str(&secs.to_string())
|
||||||
|
{
|
||||||
|
response.headers_mut().insert(header::RETRY_AFTER, value);
|
||||||
|
}
|
||||||
|
response
|
||||||
|
}
|
||||||
@@ -1,37 +1,89 @@
|
|||||||
use crate::state::RouterState;
|
use crate::state::RouterState;
|
||||||
use axum::{Json, Router, extract::State, routing::get};
|
use crate::{catalogue, dispatch};
|
||||||
use cortex_core::openai::ModelsResponse;
|
use axum::body::Bytes;
|
||||||
|
use axum::http::HeaderMap;
|
||||||
|
use axum::response::Response;
|
||||||
|
use axum::{Json, Router, extract::State, routing::get, routing::post};
|
||||||
use serde_json::{Value, json};
|
use serde_json::{Value, json};
|
||||||
use std::sync::Arc;
|
use std::sync::Arc;
|
||||||
|
|
||||||
/// Routes served by the router skeleton. The inference paths
|
/// Routes served by the router. Inference paths are capacity-aware-dispatched
|
||||||
/// (`/v1/chat/completions`, `/v1/messages`, …) arrive with capacity-aware
|
/// (#73) to a downstream cortex; `/health` and a stub `/v1/models` are local.
|
||||||
/// dispatch (#73); for now the router only answers `/health` and a stub
|
|
||||||
/// `/v1/models`.
|
|
||||||
pub fn api_routes() -> Router<Arc<RouterState>> {
|
pub fn api_routes() -> Router<Arc<RouterState>> {
|
||||||
Router::new()
|
Router::new()
|
||||||
|
.route("/v1/chat/completions", post(chat_completions))
|
||||||
|
.route("/v1/completions", post(completions))
|
||||||
|
.route("/v1/responses", post(responses))
|
||||||
|
.route("/v1/messages", post(messages))
|
||||||
.route("/v1/models", get(list_models))
|
.route("/v1/models", get(list_models))
|
||||||
.route("/health", get(health))
|
.route("/health", get(health))
|
||||||
.route("/", get(health))
|
.route("/", get(health))
|
||||||
}
|
}
|
||||||
|
|
||||||
/// `GET /health` — liveness plus the configured downstream cortex count.
|
// ── Inference paths — forwarded verbatim to a chosen cortex ──────────
|
||||||
/// Real per-cortex reachability lands with the poller (#72).
|
//
|
||||||
|
// Each handler dispatches to the same path on a capacity-bearing cortex.
|
||||||
|
// The body is parsed only to read `model`; the bearer and bytes are
|
||||||
|
// forwarded unchanged, and the SSE response streams back verbatim.
|
||||||
|
|
||||||
|
async fn chat_completions(
|
||||||
|
State(state): State<Arc<RouterState>>,
|
||||||
|
headers: HeaderMap,
|
||||||
|
body: Bytes,
|
||||||
|
) -> Response {
|
||||||
|
dispatch::dispatch(&state, "/v1/chat/completions", headers, body).await
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn completions(
|
||||||
|
State(state): State<Arc<RouterState>>,
|
||||||
|
headers: HeaderMap,
|
||||||
|
body: Bytes,
|
||||||
|
) -> Response {
|
||||||
|
dispatch::dispatch(&state, "/v1/completions", headers, body).await
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn responses(
|
||||||
|
State(state): State<Arc<RouterState>>,
|
||||||
|
headers: HeaderMap,
|
||||||
|
body: Bytes,
|
||||||
|
) -> Response {
|
||||||
|
dispatch::dispatch(&state, "/v1/responses", headers, body).await
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn messages(
|
||||||
|
State(state): State<Arc<RouterState>>,
|
||||||
|
headers: HeaderMap,
|
||||||
|
body: Bytes,
|
||||||
|
) -> Response {
|
||||||
|
dispatch::dispatch(&state, "/v1/messages", headers, body).await
|
||||||
|
}
|
||||||
|
|
||||||
|
/// `GET /health` — router liveness plus a summary of downstream cortex
|
||||||
|
/// reachability from the topology poller (#72). `status` reflects the
|
||||||
|
/// router process itself (always `ok` if it answers); downstream health is
|
||||||
|
/// the informational `cortexes` block, so a fully-degraded fleet doesn't
|
||||||
|
/// make the router look dead to its own liveness probe.
|
||||||
async fn health(State(state): State<Arc<RouterState>>) -> Json<Value> {
|
async fn health(State(state): State<Arc<RouterState>>) -> Json<Value> {
|
||||||
|
let topo = state.topology.read().await;
|
||||||
|
let reachable = topo.values().filter(|t| t.reachable).count();
|
||||||
Json(json!({
|
Json(json!({
|
||||||
"status": "ok",
|
"status": "ok",
|
||||||
"cortexes": {
|
"cortexes": {
|
||||||
"configured": state.cortexes.len(),
|
"configured": state.cortexes.len(),
|
||||||
|
"reachable": reachable,
|
||||||
}
|
}
|
||||||
}))
|
}))
|
||||||
}
|
}
|
||||||
|
|
||||||
/// `GET /v1/models` — empty catalogue stub. The real cross-operator union
|
/// `GET /v1/models` — the federation catalogue (#75): the deduped union of
|
||||||
/// (catalogue × topology feasibility, aggregated from each cortex) is the
|
/// every reachable cortex's `/v1/models`, so a client doing discovery
|
||||||
/// federation-catalogue issue (#75).
|
/// against the router resolves the whole federation without knowing about
|
||||||
async fn list_models() -> Json<ModelsResponse> {
|
/// operators or cortexes.
|
||||||
Json(ModelsResponse {
|
async fn list_models(State(state): State<Arc<RouterState>>) -> Json<Value> {
|
||||||
object: "list".into(),
|
let topo = state.topology.read().await;
|
||||||
data: vec![],
|
let data: Vec<Value> = catalogue::aggregate_models(&topo)
|
||||||
})
|
.iter()
|
||||||
|
.map(|e| json!(e))
|
||||||
|
.collect();
|
||||||
|
Json(json!({ "object": "list", "data": data }))
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -8,12 +8,15 @@
|
|||||||
//!
|
//!
|
||||||
//! It holds **zero entitlement logic** — auth/budget stays at cortex
|
//! It holds **zero entitlement logic** — auth/budget stays at cortex
|
||||||
//! (epic #47); the router forwards the client bearer unchanged and routes
|
//! (epic #47); the router forwards the client bearer unchanged and routes
|
||||||
//! on capacity (epic #69). This crate is the binary skeleton (#70):
|
//! on capacity (epic #69). A background [`poller`] keeps a live
|
||||||
//! a plaintext axum server reusing `cortex-core` types, serving `/health`
|
//! per-cortex topology (#72) that the dispatcher (#73) will route on.
|
||||||
//! and a stub `/v1/models`.
|
|
||||||
|
|
||||||
|
pub mod catalogue;
|
||||||
pub mod config;
|
pub mod config;
|
||||||
|
pub mod dispatch;
|
||||||
|
pub mod error;
|
||||||
pub mod handlers;
|
pub mod handlers;
|
||||||
|
pub mod poller;
|
||||||
pub mod state;
|
pub mod state;
|
||||||
|
|
||||||
use anyhow::Result;
|
use anyhow::Result;
|
||||||
@@ -37,7 +40,15 @@ pub fn build_app(state: Arc<state::RouterState>) -> axum::Router {
|
|||||||
/// listener. TLS is terminated by edge nginx ahead of this process.
|
/// listener. TLS is terminated by edge nginx ahead of this process.
|
||||||
pub async fn run(config: RouterConfig) -> Result<()> {
|
pub async fn run(config: RouterConfig) -> Result<()> {
|
||||||
let state = Arc::new(state::RouterState::from_config(&config));
|
let state = Arc::new(state::RouterState::from_config(&config));
|
||||||
let app = build_app(state);
|
|
||||||
|
// Background topology poller (#72): refresh each cortex's health +
|
||||||
|
// catalogue so routing decisions see live capacity.
|
||||||
|
let poller_state = Arc::clone(&state);
|
||||||
|
tokio::spawn(async move {
|
||||||
|
poller::poll_loop(poller_state).await;
|
||||||
|
});
|
||||||
|
|
||||||
|
let app = build_app(Arc::clone(&state));
|
||||||
|
|
||||||
let listen_addr = config.router.listen.parse::<std::net::SocketAddr>()?;
|
let listen_addr = config.router.listen.parse::<std::net::SocketAddr>()?;
|
||||||
tracing::info!("helexa-router listening on {listen_addr}");
|
tracing::info!("helexa-router listening on {listen_addr}");
|
||||||
|
|||||||
150
crates/helexa-router/src/poller.rs
Normal file
150
crates/helexa-router/src/poller.rs
Normal file
@@ -0,0 +1,150 @@
|
|||||||
|
//! Background poller that refreshes the multi-operator topology (#72).
|
||||||
|
//!
|
||||||
|
//! The same pattern as cortex↔neuron, one tier up: periodically poll each
|
||||||
|
//! configured cortex's `GET /v1/models` (catalogue × topology feasibility +
|
||||||
|
//! loaded state) and `GET /health` (coarse node-health/load), building the
|
||||||
|
//! live map the dispatcher (#73) routes on. An unreachable or erroring
|
||||||
|
//! cortex is debounced over [`POLL_FAILURE_THRESHOLD`] consecutive misses,
|
||||||
|
//! then flipped unhealthy and excluded from routing; it recovers on the
|
||||||
|
//! next successful poll.
|
||||||
|
|
||||||
|
use crate::state::RouterState;
|
||||||
|
use chrono::Utc;
|
||||||
|
use cortex_core::node::CortexModelEntry;
|
||||||
|
use serde::Deserialize;
|
||||||
|
use std::time::Duration;
|
||||||
|
|
||||||
|
/// Per-cortex HTTP timeout for each poll request.
|
||||||
|
const POLL_TIMEOUT: Duration = Duration::from_secs(5);
|
||||||
|
|
||||||
|
/// Consecutive failed polls before a cortex is marked unreachable. Mirrors
|
||||||
|
/// cortex's neuron-poll debounce: a single blip (a busy cortex briefly slow
|
||||||
|
/// to answer) can't yank it — and all its models — out of routing.
|
||||||
|
pub const POLL_FAILURE_THRESHOLD: u32 = 3;
|
||||||
|
|
||||||
|
/// cortex's `/v1/models` envelope — `{ "object": "list", "data": [...] }`.
|
||||||
|
#[derive(Debug, Deserialize)]
|
||||||
|
struct ModelsEnvelope {
|
||||||
|
#[serde(default)]
|
||||||
|
data: Vec<CortexModelEntry>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// The subset of cortex's `/health` the router reads.
|
||||||
|
#[derive(Debug, Deserialize)]
|
||||||
|
struct CortexHealth {
|
||||||
|
nodes: CortexHealthNodes,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Deserialize)]
|
||||||
|
struct CortexHealthNodes {
|
||||||
|
healthy: u32,
|
||||||
|
total: u32,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Run forever, polling all cortexes on the configured interval.
|
||||||
|
pub async fn poll_loop(state: std::sync::Arc<RouterState>) {
|
||||||
|
loop {
|
||||||
|
poll_once(&state).await;
|
||||||
|
tokio::time::sleep(state.poll_interval).await;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Poll every configured cortex once. Public for testing.
|
||||||
|
pub async fn poll_once(state: &RouterState) {
|
||||||
|
for cortex in &state.cortexes {
|
||||||
|
poll_cortex(state, &cortex.name, &cortex.endpoint).await;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Poll one cortex: refresh its model map from `/v1/models`, then its node
|
||||||
|
/// health from `/health`. A `/v1/models` failure debounces toward
|
||||||
|
/// unreachable; the `/health` poll is best-effort and never flips
|
||||||
|
/// reachability on its own (a cortex serving `/v1/models` is routable even
|
||||||
|
/// if `/health` momentarily isn't).
|
||||||
|
async fn poll_cortex(state: &RouterState, name: &str, endpoint: &str) {
|
||||||
|
// A cortex whose pinned TLS client failed to build (#74) is disabled:
|
||||||
|
// there is no client to poll with, so it stays unreachable.
|
||||||
|
let Some(client) = state.client_for(name) else {
|
||||||
|
let mut topo = state.topology.write().await;
|
||||||
|
if let Some(entry) = topo.get_mut(name) {
|
||||||
|
entry.consecutive_failures = entry.consecutive_failures.saturating_add(1);
|
||||||
|
entry.reachable = false;
|
||||||
|
}
|
||||||
|
tracing::warn!(cortex = name, "no TLS client (disabled); skipping poll");
|
||||||
|
return;
|
||||||
|
};
|
||||||
|
|
||||||
|
let models = fetch_models(client, endpoint).await;
|
||||||
|
|
||||||
|
let mut topo = state.topology.write().await;
|
||||||
|
let Some(entry) = topo.get_mut(name) else {
|
||||||
|
return; // not a configured cortex (shouldn't happen)
|
||||||
|
};
|
||||||
|
|
||||||
|
match models {
|
||||||
|
Ok(models) => {
|
||||||
|
entry.models = models.into_iter().map(|m| (m.id.clone(), m)).collect();
|
||||||
|
entry.reachable = true;
|
||||||
|
entry.consecutive_failures = 0;
|
||||||
|
entry.last_poll = Some(Utc::now());
|
||||||
|
tracing::debug!(cortex = name, models = entry.models.len(), "poll ok");
|
||||||
|
}
|
||||||
|
Err(reason) => {
|
||||||
|
entry.consecutive_failures = entry.consecutive_failures.saturating_add(1);
|
||||||
|
if entry.consecutive_failures >= POLL_FAILURE_THRESHOLD {
|
||||||
|
entry.reachable = false;
|
||||||
|
}
|
||||||
|
tracing::warn!(
|
||||||
|
cortex = name,
|
||||||
|
failures = entry.consecutive_failures,
|
||||||
|
reachable = entry.reachable,
|
||||||
|
reason,
|
||||||
|
"cortex poll failed"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
drop(topo);
|
||||||
|
|
||||||
|
// Best-effort health (node counts). Never flips reachability.
|
||||||
|
if let Some((healthy, total)) = fetch_health(client, endpoint).await {
|
||||||
|
let mut topo = state.topology.write().await;
|
||||||
|
if let Some(entry) = topo.get_mut(name) {
|
||||||
|
entry.healthy_nodes = healthy;
|
||||||
|
entry.total_nodes = total;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// GET `/v1/models`, returning the parsed entries or a short failure reason.
|
||||||
|
async fn fetch_models(
|
||||||
|
client: &reqwest::Client,
|
||||||
|
endpoint: &str,
|
||||||
|
) -> Result<Vec<CortexModelEntry>, &'static str> {
|
||||||
|
let url = format!("{endpoint}/v1/models");
|
||||||
|
let resp = client
|
||||||
|
.get(&url)
|
||||||
|
.timeout(POLL_TIMEOUT)
|
||||||
|
.send()
|
||||||
|
.await
|
||||||
|
.map_err(|_| "unreachable")?;
|
||||||
|
if !resp.status().is_success() {
|
||||||
|
return Err("non-success status");
|
||||||
|
}
|
||||||
|
let envelope = resp
|
||||||
|
.json::<ModelsEnvelope>()
|
||||||
|
.await
|
||||||
|
.map_err(|_| "bad json")?;
|
||||||
|
Ok(envelope.data)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// GET `/health`, returning `(healthy, total)` node counts. `None` on any
|
||||||
|
/// failure — the caller leaves the previous counts in place.
|
||||||
|
async fn fetch_health(client: &reqwest::Client, endpoint: &str) -> Option<(u32, u32)> {
|
||||||
|
let url = format!("{endpoint}/health");
|
||||||
|
let resp = client.get(&url).timeout(POLL_TIMEOUT).send().await.ok()?;
|
||||||
|
if !resp.status().is_success() {
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
let health = resp.json::<CortexHealth>().await.ok()?;
|
||||||
|
Some((health.nodes.healthy, health.nodes.total))
|
||||||
|
}
|
||||||
@@ -1,21 +1,144 @@
|
|||||||
use crate::config::{CortexEndpoint, RouterConfig};
|
use crate::config::{CortexEndpoint, RouterConfig};
|
||||||
|
use chrono::{DateTime, Utc};
|
||||||
|
use cortex_core::node::CortexModelEntry;
|
||||||
|
use std::collections::HashMap;
|
||||||
|
use std::time::Duration;
|
||||||
|
use tokio::sync::RwLock;
|
||||||
|
|
||||||
/// Shared router state.
|
/// Shared router state: the configured cortex list plus the live topology
|
||||||
|
/// map the poller (#72) maintains and the dispatcher (#73) will route on.
|
||||||
///
|
///
|
||||||
/// The skeleton (#70) holds only the static downstream cortex list from
|
/// This is the router tier of the fractal neuron ← cortex ← router design:
|
||||||
/// config. Live multi-operator topology (per-cortex capacity + catalogue)
|
/// just as cortex polls each neuron for capacity/catalogue, the router
|
||||||
/// is added by the poller (#72), at which point this grows an
|
/// polls each cortex's `/health` + `/v1/models`.
|
||||||
/// `Arc<RwLock<...>>` topology map alongside the static endpoints.
|
|
||||||
#[derive(Debug)]
|
#[derive(Debug)]
|
||||||
pub struct RouterState {
|
pub struct RouterState {
|
||||||
/// Downstream cortex endpoints, as configured.
|
/// Downstream cortex endpoints, as configured.
|
||||||
pub cortexes: Vec<CortexEndpoint>,
|
pub cortexes: Vec<CortexEndpoint>,
|
||||||
|
/// Per-cortex HTTP client, keyed by cortex name (#74). A cortex enrolled
|
||||||
|
/// with a `tls_ca` gets a client that trusts only that anchor; others
|
||||||
|
/// get a default client. A cortex whose `tls_ca` failed to load is
|
||||||
|
/// **absent** here — `client_for` returns `None` and it is never
|
||||||
|
/// polled or routed to (fail closed: a misconfigured pin must not
|
||||||
|
/// silently fall back to unpinned TLS).
|
||||||
|
clients: HashMap<String, reqwest::Client>,
|
||||||
|
/// This router instance's region, for dispatch geo affinity (#73).
|
||||||
|
pub region: Option<String>,
|
||||||
|
/// How often the poller refreshes the topology.
|
||||||
|
pub poll_interval: Duration,
|
||||||
|
/// Live per-cortex topology, keyed by cortex name. Pre-populated from
|
||||||
|
/// config (every configured cortex present, `reachable = false`) so the
|
||||||
|
/// poller and handlers always find an entry; the poller flips
|
||||||
|
/// reachability and fills the model map.
|
||||||
|
pub topology: RwLock<HashMap<String, CortexTopology>>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Live view of one downstream cortex, refreshed each poll.
|
||||||
|
#[derive(Debug, Clone, Default)]
|
||||||
|
pub struct CortexTopology {
|
||||||
|
/// Whether the cortex is currently routable. Flipped `false` only after
|
||||||
|
/// [`crate::poller::POLL_FAILURE_THRESHOLD`] consecutive failed polls
|
||||||
|
/// (debounces transient blips); restored on the next successful poll.
|
||||||
|
pub reachable: bool,
|
||||||
|
/// Consecutive failed polls; reset to 0 on success.
|
||||||
|
pub consecutive_failures: u32,
|
||||||
|
/// Timestamp of the last successful poll.
|
||||||
|
pub last_poll: Option<DateTime<Utc>>,
|
||||||
|
/// Healthy / total neuron counts from the cortex's `/health` (coarse
|
||||||
|
/// load signal; #73 refines headroom). 0/0 until first health poll.
|
||||||
|
pub healthy_nodes: u32,
|
||||||
|
pub total_nodes: u32,
|
||||||
|
/// The cortex's full `/v1/models` entries, keyed by model id. Stored
|
||||||
|
/// whole (not distilled to a loaded/feasible bool) so the federation
|
||||||
|
/// catalogue (#75) can preserve per-model `limit`/`cost`/capabilities.
|
||||||
|
pub models: HashMap<String, CortexModelEntry>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Whether a cortex can serve this model — loaded now, or feasible to
|
||||||
|
/// cold-load (its catalogue × topology says some neuron can host it).
|
||||||
|
pub fn entry_feasible(entry: &CortexModelEntry) -> bool {
|
||||||
|
entry.loaded || !entry.feasible_on.is_empty()
|
||||||
}
|
}
|
||||||
|
|
||||||
impl RouterState {
|
impl RouterState {
|
||||||
pub fn from_config(config: &RouterConfig) -> Self {
|
pub fn from_config(config: &RouterConfig) -> Self {
|
||||||
|
let topology = config
|
||||||
|
.cortexes
|
||||||
|
.iter()
|
||||||
|
.map(|c| (c.name.clone(), CortexTopology::default()))
|
||||||
|
.collect();
|
||||||
|
|
||||||
|
// One client per cortex. A `tls_ca` that fails to load omits the
|
||||||
|
// cortex from the map (fail closed) rather than degrading to an
|
||||||
|
// unpinned client.
|
||||||
|
let mut clients = HashMap::new();
|
||||||
|
for c in &config.cortexes {
|
||||||
|
match build_client(c.tls_ca.as_deref()) {
|
||||||
|
Ok(client) => {
|
||||||
|
clients.insert(c.name.clone(), client);
|
||||||
|
}
|
||||||
|
Err(e) => {
|
||||||
|
tracing::error!(
|
||||||
|
cortex = %c.name,
|
||||||
|
tls_ca = c.tls_ca.as_deref().unwrap_or(""),
|
||||||
|
error = %e,
|
||||||
|
"failed to build pinned TLS client; cortex disabled (fail closed)"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
Self {
|
Self {
|
||||||
cortexes: config.cortexes.clone(),
|
cortexes: config.cortexes.clone(),
|
||||||
|
clients,
|
||||||
|
region: config.router.region.clone(),
|
||||||
|
poll_interval: Duration::from_secs(config.router.poll_interval_secs),
|
||||||
|
topology: RwLock::new(topology),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// The HTTP client to use for `name`, or `None` if the cortex is
|
||||||
|
/// disabled (its `tls_ca` failed to load). Callers must treat `None` as
|
||||||
|
/// "not routable / not pollable".
|
||||||
|
pub fn client_for(&self, name: &str) -> Option<&reqwest::Client> {
|
||||||
|
self.clients.get(name)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Names of reachable cortexes that can serve `model_id` (loaded or
|
||||||
|
/// feasible to cold-load). Groundwork for capacity-aware dispatch (#73);
|
||||||
|
/// unreachable cortexes are excluded by construction.
|
||||||
|
pub async fn cortexes_serving(&self, model_id: &str) -> Vec<String> {
|
||||||
|
let topo = self.topology.read().await;
|
||||||
|
topo.iter()
|
||||||
|
.filter(|(_, t)| t.reachable)
|
||||||
|
.filter(|(_, t)| t.models.get(model_id).is_some_and(entry_feasible))
|
||||||
|
.map(|(name, _)| name.clone())
|
||||||
|
.collect()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Build a cortex HTTP client. With `tls_ca` set, the client trusts **only**
|
||||||
|
/// that PEM anchor (platform roots disabled) — pinning the router→cortex hop
|
||||||
|
/// to an enrolled cert (#74). Without it, standard platform-root validation.
|
||||||
|
pub fn build_client(tls_ca: Option<&str>) -> Result<reqwest::Client, BuildClientError> {
|
||||||
|
let mut builder = reqwest::Client::builder();
|
||||||
|
if let Some(path) = tls_ca {
|
||||||
|
let pem = std::fs::read(path).map_err(|e| BuildClientError::Read(path.to_string(), e))?;
|
||||||
|
let cert = reqwest::Certificate::from_pem(&pem).map_err(BuildClientError::Parse)?;
|
||||||
|
builder = builder
|
||||||
|
.tls_built_in_root_certs(false)
|
||||||
|
.add_root_certificate(cert);
|
||||||
|
}
|
||||||
|
builder.build().map_err(BuildClientError::Build)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Why a cortex's pinned client could not be built (→ cortex disabled).
|
||||||
|
#[derive(Debug, thiserror::Error)]
|
||||||
|
pub enum BuildClientError {
|
||||||
|
#[error("reading TLS anchor '{0}'")]
|
||||||
|
Read(String, #[source] std::io::Error),
|
||||||
|
#[error("parsing TLS anchor PEM")]
|
||||||
|
Parse(#[source] reqwest::Error),
|
||||||
|
#[error("building HTTP client")]
|
||||||
|
Build(#[source] reqwest::Error),
|
||||||
}
|
}
|
||||||
|
|||||||
132
crates/helexa-router/tests/catalogue.rs
Normal file
132
crates/helexa-router/tests/catalogue.rs
Normal file
@@ -0,0 +1,132 @@
|
|||||||
|
//! End-to-end federation-catalogue test for #75: poll two mock cortexes
|
||||||
|
//! that overlap on a model, then `GET /v1/models` on the router and verify
|
||||||
|
//! the deduped union with merged availability and preserved limit/cost.
|
||||||
|
|
||||||
|
use axum::Router;
|
||||||
|
use axum::routing::get;
|
||||||
|
use helexa_router::config::{CortexEndpoint, RouterConfig};
|
||||||
|
use helexa_router::poller::poll_once;
|
||||||
|
use helexa_router::state::RouterState;
|
||||||
|
use serde_json::{Value, json};
|
||||||
|
use std::sync::Arc;
|
||||||
|
use tokio::net::TcpListener;
|
||||||
|
|
||||||
|
/// Spawn a mock cortex serving the given `/v1/models` `data` array.
|
||||||
|
async fn spawn_cortex(models: Value) -> String {
|
||||||
|
let models = Arc::new(models);
|
||||||
|
let app = Router::new()
|
||||||
|
.route(
|
||||||
|
"/v1/models",
|
||||||
|
get({
|
||||||
|
let models = Arc::clone(&models);
|
||||||
|
move || {
|
||||||
|
let models = Arc::clone(&models);
|
||||||
|
async move { axum::Json(json!({ "object": "list", "data": &*models })) }
|
||||||
|
}
|
||||||
|
}),
|
||||||
|
)
|
||||||
|
.route(
|
||||||
|
"/health",
|
||||||
|
get(|| async { axum::Json(json!({"status":"ok","nodes":{"healthy":1,"total":1}})) }),
|
||||||
|
);
|
||||||
|
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||||
|
let addr = listener.local_addr().unwrap();
|
||||||
|
tokio::spawn(async move {
|
||||||
|
axum::serve(listener, app).await.unwrap();
|
||||||
|
});
|
||||||
|
format!("http://{addr}")
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Spawn the router (with poller) wired to the given cortex endpoints, and
|
||||||
|
/// poll once synchronously so the topology is populated before we query.
|
||||||
|
async fn spawn_router(cortexes: Vec<CortexEndpoint>) -> String {
|
||||||
|
let cfg = RouterConfig {
|
||||||
|
cortexes,
|
||||||
|
..Default::default()
|
||||||
|
};
|
||||||
|
let state = Arc::new(RouterState::from_config(&cfg));
|
||||||
|
poll_once(&state).await; // deterministic: fill topology now
|
||||||
|
|
||||||
|
let app = helexa_router::build_app(Arc::clone(&state));
|
||||||
|
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||||
|
let addr = listener.local_addr().unwrap();
|
||||||
|
tokio::spawn(async move {
|
||||||
|
axum::serve(listener, app).await.unwrap();
|
||||||
|
});
|
||||||
|
format!("http://{addr}")
|
||||||
|
}
|
||||||
|
|
||||||
|
fn model(id: &str, loaded: bool, feasible_on: &[&str], ctx: u64, input_cost: f64) -> Value {
|
||||||
|
json!({
|
||||||
|
"id": id,
|
||||||
|
"object": "model",
|
||||||
|
"created": 0,
|
||||||
|
"owned_by": "helexa",
|
||||||
|
"loaded": loaded,
|
||||||
|
"feasible_on": feasible_on,
|
||||||
|
"locations": [],
|
||||||
|
"limit": { "context": ctx, "output": 4096 },
|
||||||
|
"cost": { "input": input_cost, "output": input_cost * 3.0 }
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn federation_catalogue_dedupes_and_preserves_limit_cost() {
|
||||||
|
// cortex A: "shared" loaded (ctx 32768, $0.50) + "only-a" loaded.
|
||||||
|
let a = spawn_cortex(json!([
|
||||||
|
model("shared", true, &["beast"], 32_768, 0.50),
|
||||||
|
model("only-a", true, &["beast"], 8_192, 1.00),
|
||||||
|
]))
|
||||||
|
.await;
|
||||||
|
// cortex B: "shared" cold-loadable, tighter ctx (16384), cheaper ($0.20).
|
||||||
|
let b = spawn_cortex(json!([model("shared", false, &["benjy"], 16_384, 0.20)])).await;
|
||||||
|
|
||||||
|
let router = spawn_router(vec![
|
||||||
|
CortexEndpoint {
|
||||||
|
name: "op-a".into(),
|
||||||
|
endpoint: a,
|
||||||
|
region: None,
|
||||||
|
tls_ca: None,
|
||||||
|
},
|
||||||
|
CortexEndpoint {
|
||||||
|
name: "op-b".into(),
|
||||||
|
endpoint: b,
|
||||||
|
region: None,
|
||||||
|
tls_ca: None,
|
||||||
|
},
|
||||||
|
])
|
||||||
|
.await;
|
||||||
|
|
||||||
|
let body: Value = reqwest::get(format!("{router}/v1/models"))
|
||||||
|
.await
|
||||||
|
.unwrap()
|
||||||
|
.json()
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
assert_eq!(body["object"], "list");
|
||||||
|
let data = body["data"].as_array().unwrap();
|
||||||
|
// Deduped union: "shared" once + "only-a".
|
||||||
|
assert_eq!(data.len(), 2);
|
||||||
|
|
||||||
|
let shared = data.iter().find(|m| m["id"] == "shared").unwrap();
|
||||||
|
// Loaded somewhere (op-a) → loaded.
|
||||||
|
assert_eq!(shared["loaded"], true);
|
||||||
|
// feasible_on re-tiered to operator names, both present, sorted.
|
||||||
|
let feasible: Vec<&str> = shared["feasible_on"]
|
||||||
|
.as_array()
|
||||||
|
.unwrap()
|
||||||
|
.iter()
|
||||||
|
.map(|v| v.as_str().unwrap())
|
||||||
|
.collect();
|
||||||
|
assert_eq!(feasible, vec!["op-a", "op-b"]);
|
||||||
|
// Tightest limit (16384) and cheapest cost ($0.20) win.
|
||||||
|
assert_eq!(shared["limit"]["context"], 16_384);
|
||||||
|
assert_eq!(shared["cost"]["input"], 0.20);
|
||||||
|
// Loaded location named by operator, no neuron VRAM leaked.
|
||||||
|
let locs = shared["locations"].as_array().unwrap();
|
||||||
|
assert_eq!(locs.len(), 1);
|
||||||
|
assert_eq!(locs[0]["node"], "op-a");
|
||||||
|
|
||||||
|
assert!(data.iter().any(|m| m["id"] == "only-a"));
|
||||||
|
}
|
||||||
301
crates/helexa-router/tests/dispatch.rs
Normal file
301
crates/helexa-router/tests/dispatch.rs
Normal file
@@ -0,0 +1,301 @@
|
|||||||
|
//! Capacity-aware dispatch acceptance tests for #73.
|
||||||
|
//!
|
||||||
|
//! Covers: a request routes to a cortex serving the model; the client's
|
||||||
|
//! bearer reaches the cortex; cortex's #63 rejections pass through verbatim
|
||||||
|
//! and are NOT retried away; transport failure fails over to another
|
||||||
|
//! feasible cortex; unknown model → 404, no reachable capacity → 503; and
|
||||||
|
//! the selection ranking (warm/region/headroom).
|
||||||
|
|
||||||
|
use axum::body::Bytes;
|
||||||
|
use axum::extract::State;
|
||||||
|
use axum::http::{HeaderMap, StatusCode};
|
||||||
|
use axum::response::{IntoResponse, Response};
|
||||||
|
use axum::routing::post;
|
||||||
|
use axum::{Json, Router};
|
||||||
|
use cortex_core::node::CortexModelEntry;
|
||||||
|
use helexa_router::config::{CortexEndpoint, RouterConfig};
|
||||||
|
use helexa_router::dispatch::{Selection, dispatch, select_cortexes};
|
||||||
|
use helexa_router::state::{CortexTopology, RouterState};
|
||||||
|
use serde_json::{Value, json};
|
||||||
|
use std::collections::HashMap;
|
||||||
|
use tokio::net::TcpListener;
|
||||||
|
|
||||||
|
/// A minimal `CortexModelEntry` for MODEL with the given serveability.
|
||||||
|
fn model_entry(loaded: bool, feasible: bool) -> CortexModelEntry {
|
||||||
|
CortexModelEntry {
|
||||||
|
id: MODEL.into(),
|
||||||
|
object: "model".into(),
|
||||||
|
created: 0,
|
||||||
|
owned_by: "helexa".into(),
|
||||||
|
loaded,
|
||||||
|
feasible_on: if feasible || loaded {
|
||||||
|
vec!["n".into()]
|
||||||
|
} else {
|
||||||
|
vec![]
|
||||||
|
},
|
||||||
|
locations: vec![],
|
||||||
|
capabilities: vec![],
|
||||||
|
limit: None,
|
||||||
|
cost: None,
|
||||||
|
tool_call: false,
|
||||||
|
reasoning: false,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
const MODEL: &str = "Qwen/Qwen3-Coder-30B";
|
||||||
|
|
||||||
|
// ── Mock cortex backend ──────────────────────────────────────────────
|
||||||
|
|
||||||
|
/// Behaviour of a mock cortex, carried in axum State.
|
||||||
|
#[derive(Clone)]
|
||||||
|
struct MockCortex {
|
||||||
|
/// Identifies which cortex answered, echoed in the 200 body.
|
||||||
|
name: &'static str,
|
||||||
|
/// When true, return a genuine #63-shaped `429 rate_limit_exceeded`.
|
||||||
|
rate_limited: bool,
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn mock_handler(State(m): State<MockCortex>, headers: HeaderMap) -> Response {
|
||||||
|
if m.rate_limited {
|
||||||
|
return (
|
||||||
|
StatusCode::TOO_MANY_REQUESTS,
|
||||||
|
Json(json!({"error":{"type":"rate_limit_error","code":"rate_limit_exceeded","message":"slow down","param":null}})),
|
||||||
|
)
|
||||||
|
.into_response();
|
||||||
|
}
|
||||||
|
let auth = headers
|
||||||
|
.get("authorization")
|
||||||
|
.and_then(|v| v.to_str().ok())
|
||||||
|
.unwrap_or("")
|
||||||
|
.to_string();
|
||||||
|
Json(json!({ "served_by": m.name, "auth_seen": auth })).into_response()
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn spawn_cortex(mock: MockCortex) -> String {
|
||||||
|
let app = Router::new()
|
||||||
|
.route("/v1/chat/completions", post(mock_handler))
|
||||||
|
.with_state(mock);
|
||||||
|
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||||
|
let addr = listener.local_addr().unwrap();
|
||||||
|
tokio::spawn(async move {
|
||||||
|
axum::serve(listener, app).await.unwrap();
|
||||||
|
});
|
||||||
|
format!("http://{addr}")
|
||||||
|
}
|
||||||
|
|
||||||
|
fn ok_cortex(name: &'static str) -> MockCortex {
|
||||||
|
MockCortex {
|
||||||
|
name,
|
||||||
|
rate_limited: false,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── Helpers to build state with a hand-set topology ──────────────────
|
||||||
|
|
||||||
|
fn state_with(cortexes: Vec<CortexEndpoint>, region: Option<String>) -> RouterState {
|
||||||
|
let cfg = RouterConfig {
|
||||||
|
cortexes,
|
||||||
|
..Default::default()
|
||||||
|
};
|
||||||
|
let mut state = RouterState::from_config(&cfg);
|
||||||
|
state.region = region;
|
||||||
|
state
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Overwrite the topology entry for `name` so tests control reachability and
|
||||||
|
/// model serveability directly (no live poll).
|
||||||
|
async fn set_topology(
|
||||||
|
state: &RouterState,
|
||||||
|
name: &str,
|
||||||
|
reachable: bool,
|
||||||
|
loaded: bool,
|
||||||
|
feasible: bool,
|
||||||
|
healthy_nodes: u32,
|
||||||
|
) {
|
||||||
|
let mut topo = state.topology.write().await;
|
||||||
|
let mut models = HashMap::new();
|
||||||
|
models.insert(MODEL.to_string(), model_entry(loaded, feasible));
|
||||||
|
topo.insert(
|
||||||
|
name.to_string(),
|
||||||
|
CortexTopology {
|
||||||
|
reachable,
|
||||||
|
consecutive_failures: 0,
|
||||||
|
last_poll: None,
|
||||||
|
healthy_nodes,
|
||||||
|
total_nodes: healthy_nodes,
|
||||||
|
models,
|
||||||
|
},
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
fn ep(name: &str, endpoint: &str, region: Option<&str>) -> CortexEndpoint {
|
||||||
|
CortexEndpoint {
|
||||||
|
name: name.into(),
|
||||||
|
endpoint: endpoint.into(),
|
||||||
|
region: region.map(str::to_string),
|
||||||
|
tls_ca: None,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn chat_body() -> Bytes {
|
||||||
|
Bytes::from(format!("{{\"model\":\"{MODEL}\",\"stream\":false}}"))
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn body_json(resp: Response) -> (StatusCode, Value) {
|
||||||
|
let status = resp.status();
|
||||||
|
let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
let v = serde_json::from_slice(&bytes).unwrap_or(Value::Null);
|
||||||
|
(status, v)
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── Tests ────────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn routes_to_serving_cortex_and_forwards_bearer() {
|
||||||
|
let url = spawn_cortex(ok_cortex("c1")).await;
|
||||||
|
let state = state_with(vec![ep("c1", &url, None)], None);
|
||||||
|
set_topology(&state, "c1", true, true, true, 2).await;
|
||||||
|
|
||||||
|
let mut headers = HeaderMap::new();
|
||||||
|
headers.insert("authorization", "Bearer sk-test-123".parse().unwrap());
|
||||||
|
|
||||||
|
let resp = dispatch(&state, "/v1/chat/completions", headers, chat_body()).await;
|
||||||
|
let (status, body) = body_json(resp).await;
|
||||||
|
|
||||||
|
assert_eq!(status, StatusCode::OK);
|
||||||
|
assert_eq!(body["served_by"], "c1");
|
||||||
|
// Bearer reached the cortex unchanged.
|
||||||
|
assert_eq!(body["auth_seen"], "Bearer sk-test-123");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn cortex_429_passes_through_and_is_not_retried() {
|
||||||
|
// c1 (ranked first: loaded) returns a genuine 429; c2 would return 200.
|
||||||
|
let c1 = spawn_cortex(MockCortex {
|
||||||
|
name: "c1",
|
||||||
|
rate_limited: true,
|
||||||
|
})
|
||||||
|
.await;
|
||||||
|
let c2 = spawn_cortex(ok_cortex("c2")).await;
|
||||||
|
let state = state_with(vec![ep("c1", &c1, None), ep("c2", &c2, None)], None);
|
||||||
|
// Both reachable + loaded; c1 has more headroom so it ranks first.
|
||||||
|
set_topology(&state, "c1", true, true, true, 5).await;
|
||||||
|
set_topology(&state, "c2", true, true, true, 1).await;
|
||||||
|
|
||||||
|
let resp = dispatch(
|
||||||
|
&state,
|
||||||
|
"/v1/chat/completions",
|
||||||
|
HeaderMap::new(),
|
||||||
|
chat_body(),
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
let (status, body) = body_json(resp).await;
|
||||||
|
|
||||||
|
// The genuine 4xx is returned verbatim — NOT retried to c2.
|
||||||
|
assert_eq!(status, StatusCode::TOO_MANY_REQUESTS);
|
||||||
|
assert_eq!(body["error"]["code"], "rate_limit_exceeded");
|
||||||
|
assert!(body.get("served_by").is_none(), "must not have hit c2");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn fails_over_to_next_cortex_on_transport_error() {
|
||||||
|
// c_dead ranks first (more headroom) but its endpoint is a closed port;
|
||||||
|
// c_live is the fallback. The router must fail over and c_live serves.
|
||||||
|
let live = spawn_cortex(ok_cortex("c_live")).await;
|
||||||
|
let state = state_with(
|
||||||
|
vec![
|
||||||
|
ep("c_dead", "http://127.0.0.1:1", None),
|
||||||
|
ep("c_live", &live, None),
|
||||||
|
],
|
||||||
|
None,
|
||||||
|
);
|
||||||
|
set_topology(&state, "c_dead", true, true, true, 9).await;
|
||||||
|
set_topology(&state, "c_live", true, true, true, 1).await;
|
||||||
|
|
||||||
|
let resp = dispatch(
|
||||||
|
&state,
|
||||||
|
"/v1/chat/completions",
|
||||||
|
HeaderMap::new(),
|
||||||
|
chat_body(),
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
let (status, body) = body_json(resp).await;
|
||||||
|
|
||||||
|
assert_eq!(status, StatusCode::OK);
|
||||||
|
assert_eq!(body["served_by"], "c_live");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn unknown_model_is_404() {
|
||||||
|
let state = state_with(vec![ep("c1", "http://127.0.0.1:1", None)], None);
|
||||||
|
// Topology has no entry for MODEL at all.
|
||||||
|
let resp = dispatch(
|
||||||
|
&state,
|
||||||
|
"/v1/chat/completions",
|
||||||
|
HeaderMap::new(),
|
||||||
|
chat_body(),
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
let (status, body) = body_json(resp).await;
|
||||||
|
assert_eq!(status, StatusCode::NOT_FOUND);
|
||||||
|
assert_eq!(body["error"]["code"], "model_not_found");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn known_but_all_unreachable_is_503() {
|
||||||
|
let state = state_with(vec![ep("c1", "http://127.0.0.1:1", None)], None);
|
||||||
|
// Cortex knows the model (from a prior good poll) but is now unreachable.
|
||||||
|
set_topology(&state, "c1", false, true, true, 2).await;
|
||||||
|
let resp = dispatch(
|
||||||
|
&state,
|
||||||
|
"/v1/chat/completions",
|
||||||
|
HeaderMap::new(),
|
||||||
|
chat_body(),
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
let (status, body) = body_json(resp).await;
|
||||||
|
assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
|
||||||
|
assert_eq!(body["error"]["code"], "service_unavailable");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn missing_model_field_is_400() {
|
||||||
|
let state = state_with(vec![ep("c1", "http://127.0.0.1:1", None)], None);
|
||||||
|
let resp = dispatch(
|
||||||
|
&state,
|
||||||
|
"/v1/chat/completions",
|
||||||
|
HeaderMap::new(),
|
||||||
|
Bytes::from_static(b"{\"messages\":[]}"),
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
let (status, body) = body_json(resp).await;
|
||||||
|
assert_eq!(status, StatusCode::BAD_REQUEST);
|
||||||
|
assert_eq!(body["error"]["code"], "missing_model_field");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn ranking_prefers_loaded_then_region_then_headroom() {
|
||||||
|
// Router is in eu-west. Candidates:
|
||||||
|
// warm-eu : loaded, region match, 1 node → best
|
||||||
|
// warm-us : loaded, no region, 9 nodes
|
||||||
|
// cold-eu : feasible only, region match → worst (cold)
|
||||||
|
let state = state_with(
|
||||||
|
vec![
|
||||||
|
ep("warm-eu", "http://127.0.0.1:1", Some("eu-west")),
|
||||||
|
ep("warm-us", "http://127.0.0.1:1", Some("us-east")),
|
||||||
|
ep("cold-eu", "http://127.0.0.1:1", Some("eu-west")),
|
||||||
|
],
|
||||||
|
Some("eu-west".into()),
|
||||||
|
);
|
||||||
|
set_topology(&state, "warm-eu", true, true, true, 1).await;
|
||||||
|
set_topology(&state, "warm-us", true, true, true, 9).await;
|
||||||
|
set_topology(&state, "cold-eu", true, false, true, 5).await;
|
||||||
|
|
||||||
|
let Selection::Candidates(order) = select_cortexes(&state, MODEL).await else {
|
||||||
|
panic!("expected candidates");
|
||||||
|
};
|
||||||
|
let names: Vec<&str> = order.iter().map(|e| e.name.as_str()).collect();
|
||||||
|
assert_eq!(names, vec!["warm-eu", "warm-us", "cold-eu"]);
|
||||||
|
}
|
||||||
@@ -31,10 +31,14 @@ async fn health_reports_configured_cortex_count() {
|
|||||||
CortexEndpoint {
|
CortexEndpoint {
|
||||||
name: "a".into(),
|
name: "a".into(),
|
||||||
endpoint: "https://a.example.com".into(),
|
endpoint: "https://a.example.com".into(),
|
||||||
|
region: None,
|
||||||
|
tls_ca: None,
|
||||||
},
|
},
|
||||||
CortexEndpoint {
|
CortexEndpoint {
|
||||||
name: "b".into(),
|
name: "b".into(),
|
||||||
endpoint: "https://b.example.com".into(),
|
endpoint: "https://b.example.com".into(),
|
||||||
|
region: None,
|
||||||
|
tls_ca: None,
|
||||||
},
|
},
|
||||||
])
|
])
|
||||||
.await;
|
.await;
|
||||||
|
|||||||
210
crates/helexa-router/tests/tls.rs
Normal file
210
crates/helexa-router/tests/tls.rs
Normal file
@@ -0,0 +1,210 @@
|
|||||||
|
//! Outbound TLS pinning tests for #74.
|
||||||
|
//!
|
||||||
|
//! Proves the router, as a TLS client to cortexes, reaches a cortex
|
||||||
|
//! presenting its **enrolled** cert and rejects one presenting an
|
||||||
|
//! unexpected (or untrusted) cert — and that a rejected handshake flows
|
||||||
|
//! through the existing reachability path (#72) to exclude the cortex.
|
||||||
|
//!
|
||||||
|
//! A minimal `tokio-rustls` HTTPS server presents a self-signed cert; the
|
||||||
|
//! router's `reqwest` client (native-tls) validates against the PEM anchor
|
||||||
|
//! enrolled in config. Server (rustls) and client (native-tls) interoperate
|
||||||
|
//! at the protocol level — what matters is the trust decision.
|
||||||
|
|
||||||
|
use helexa_router::config::{CortexEndpoint, RouterConfig};
|
||||||
|
use helexa_router::poller::poll_once;
|
||||||
|
use helexa_router::state::{RouterState, build_client};
|
||||||
|
use std::io::Write;
|
||||||
|
use std::sync::Arc;
|
||||||
|
use tokio::io::{AsyncReadExt, AsyncWriteExt};
|
||||||
|
use tokio::net::TcpListener;
|
||||||
|
use tokio_rustls::TlsAcceptor;
|
||||||
|
|
||||||
|
/// A self-signed cert: PEM (for the reqwest pin file) + DER cert/key (for
|
||||||
|
/// the rustls server).
|
||||||
|
struct TestCert {
|
||||||
|
cert_pem: String,
|
||||||
|
cert_der: rustls::pki_types::CertificateDer<'static>,
|
||||||
|
key_der: Vec<u8>,
|
||||||
|
}
|
||||||
|
|
||||||
|
fn make_cert() -> TestCert {
|
||||||
|
let key = rcgen::generate_simple_self_signed(vec!["127.0.0.1".to_string()]).unwrap();
|
||||||
|
TestCert {
|
||||||
|
cert_pem: key.cert.pem(),
|
||||||
|
cert_der: key.cert.der().clone(),
|
||||||
|
key_der: key.key_pair.serialize_der(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Write a cert PEM to a unique temp file (named by `tag`) and return the
|
||||||
|
/// path. `tag` is caller-unique (we use the bound port), so no randomness.
|
||||||
|
fn write_pem(tag: &str, pem: &str) -> String {
|
||||||
|
let path = std::env::temp_dir().join(format!("helexa-router-tls-{tag}.pem"));
|
||||||
|
let mut f = std::fs::File::create(&path).unwrap();
|
||||||
|
f.write_all(pem.as_bytes()).unwrap();
|
||||||
|
path.to_string_lossy().into_owned()
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Spawn a minimal HTTPS server presenting `cert`, answering every request
|
||||||
|
/// with a canned `/v1/models`-shaped 200. Returns its `https://` base URL.
|
||||||
|
async fn spawn_https(cert: &TestCert) -> String {
|
||||||
|
let _ = rustls::crypto::aws_lc_rs::default_provider().install_default();
|
||||||
|
|
||||||
|
let key = rustls::pki_types::PrivateKeyDer::Pkcs8(rustls::pki_types::PrivatePkcs8KeyDer::from(
|
||||||
|
cert.key_der.clone(),
|
||||||
|
));
|
||||||
|
let config = rustls::ServerConfig::builder()
|
||||||
|
.with_no_client_auth()
|
||||||
|
.with_single_cert(vec![cert.cert_der.clone()], key)
|
||||||
|
.unwrap();
|
||||||
|
let acceptor = TlsAcceptor::from(Arc::new(config));
|
||||||
|
|
||||||
|
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||||
|
let addr = listener.local_addr().unwrap();
|
||||||
|
tokio::spawn(async move {
|
||||||
|
loop {
|
||||||
|
let Ok((stream, _)) = listener.accept().await else {
|
||||||
|
continue;
|
||||||
|
};
|
||||||
|
let acceptor = acceptor.clone();
|
||||||
|
tokio::spawn(async move {
|
||||||
|
if let Ok(mut tls) = acceptor.accept(stream).await {
|
||||||
|
let mut buf = [0u8; 2048];
|
||||||
|
let _ = tls.read(&mut buf).await; // consume request line/headers
|
||||||
|
let body = "{\"object\":\"list\",\"data\":[]}";
|
||||||
|
let resp = format!(
|
||||||
|
"HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}",
|
||||||
|
body.len(),
|
||||||
|
body
|
||||||
|
);
|
||||||
|
let _ = tls.write_all(resp.as_bytes()).await;
|
||||||
|
let _ = tls.shutdown().await;
|
||||||
|
}
|
||||||
|
});
|
||||||
|
}
|
||||||
|
});
|
||||||
|
format!("https://{addr}")
|
||||||
|
}
|
||||||
|
|
||||||
|
fn tag_for(url: &str) -> String {
|
||||||
|
url.rsplit(':').next().unwrap_or("0").to_string()
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn pinned_client_accepts_enrolled_cert_and_rejects_others() {
|
||||||
|
let server_cert = make_cert();
|
||||||
|
let other_cert = make_cert();
|
||||||
|
let url = spawn_https(&server_cert).await;
|
||||||
|
let tag = tag_for(&url);
|
||||||
|
|
||||||
|
let good_pin = write_pem(&format!("{tag}-good"), &server_cert.cert_pem);
|
||||||
|
let bad_pin = write_pem(&format!("{tag}-bad"), &other_cert.cert_pem);
|
||||||
|
|
||||||
|
// Enrolled with the server's own cert → handshake trusted → 200.
|
||||||
|
let good = build_client(Some(&good_pin)).unwrap();
|
||||||
|
let resp = good.get(format!("{url}/v1/models")).send().await;
|
||||||
|
assert!(resp.is_ok(), "enrolled cert must be accepted: {resp:?}");
|
||||||
|
assert_eq!(resp.unwrap().status(), 200);
|
||||||
|
|
||||||
|
// Enrolled with a different cert → server's cert is unexpected → reject.
|
||||||
|
let bad = build_client(Some(&bad_pin)).unwrap();
|
||||||
|
assert!(
|
||||||
|
bad.get(format!("{url}/v1/models")).send().await.is_err(),
|
||||||
|
"unexpected cert must be rejected"
|
||||||
|
);
|
||||||
|
|
||||||
|
// No enrollment (default platform roots) → self-signed cert untrusted.
|
||||||
|
let default = build_client(None).unwrap();
|
||||||
|
assert!(
|
||||||
|
default
|
||||||
|
.get(format!("{url}/v1/models"))
|
||||||
|
.send()
|
||||||
|
.await
|
||||||
|
.is_err(),
|
||||||
|
"un-enrolled self-signed cert must be rejected by default roots"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn poller_excludes_cortex_with_unexpected_cert() {
|
||||||
|
let server_cert = make_cert();
|
||||||
|
let other_cert = make_cert();
|
||||||
|
let url = spawn_https(&server_cert).await;
|
||||||
|
let tag = tag_for(&url);
|
||||||
|
|
||||||
|
let good_pin = write_pem(&format!("{tag}-pgood"), &server_cert.cert_pem);
|
||||||
|
let bad_pin = write_pem(&format!("{tag}-pbad"), &other_cert.cert_pem);
|
||||||
|
|
||||||
|
// Cortex A enrolled correctly → reachable. Cortex B enrolled with the
|
||||||
|
// wrong cert → TLS handshake fails → excluded.
|
||||||
|
let cfg = RouterConfig {
|
||||||
|
cortexes: vec![
|
||||||
|
CortexEndpoint {
|
||||||
|
name: "good".into(),
|
||||||
|
endpoint: url.clone(),
|
||||||
|
region: None,
|
||||||
|
tls_ca: Some(good_pin),
|
||||||
|
},
|
||||||
|
CortexEndpoint {
|
||||||
|
name: "bad".into(),
|
||||||
|
endpoint: url.clone(),
|
||||||
|
region: None,
|
||||||
|
tls_ca: Some(bad_pin),
|
||||||
|
},
|
||||||
|
],
|
||||||
|
..Default::default()
|
||||||
|
};
|
||||||
|
let state = RouterState::from_config(&cfg);
|
||||||
|
poll_once(&state).await;
|
||||||
|
|
||||||
|
let topo = state.topology.read().await;
|
||||||
|
assert!(
|
||||||
|
topo["good"].reachable,
|
||||||
|
"correctly-enrolled cortex reachable"
|
||||||
|
);
|
||||||
|
assert!(
|
||||||
|
!topo["bad"].reachable,
|
||||||
|
"cortex presenting an unexpected cert is excluded"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn misconfigured_pin_disables_cortex_fail_closed() {
|
||||||
|
// A `tls_ca` pointing at a nonexistent file must NOT fall back to an
|
||||||
|
// unpinned client — the cortex is disabled entirely.
|
||||||
|
let cfg = RouterConfig {
|
||||||
|
cortexes: vec![
|
||||||
|
CortexEndpoint {
|
||||||
|
name: "broken".into(),
|
||||||
|
endpoint: "https://127.0.0.1:1".into(),
|
||||||
|
region: None,
|
||||||
|
tls_ca: Some("/no/such/anchor.pem".into()),
|
||||||
|
},
|
||||||
|
CortexEndpoint {
|
||||||
|
name: "plain".into(),
|
||||||
|
endpoint: "http://127.0.0.1:1".into(),
|
||||||
|
region: None,
|
||||||
|
tls_ca: None,
|
||||||
|
},
|
||||||
|
],
|
||||||
|
..Default::default()
|
||||||
|
};
|
||||||
|
let state = RouterState::from_config(&cfg);
|
||||||
|
assert!(
|
||||||
|
state.client_for("broken").is_none(),
|
||||||
|
"a cortex with an unloadable pin is disabled (fail closed)"
|
||||||
|
);
|
||||||
|
assert!(
|
||||||
|
state.client_for("plain").is_some(),
|
||||||
|
"an un-pinned cortex still gets a client"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn build_client_rejects_garbage_pem() {
|
||||||
|
let path = write_pem(
|
||||||
|
"garbage",
|
||||||
|
"-----BEGIN CERTIFICATE-----\nnope\n-----END CERTIFICATE-----",
|
||||||
|
);
|
||||||
|
assert!(build_client(Some(&path)).is_err());
|
||||||
|
}
|
||||||
172
crates/helexa-router/tests/topology.rs
Normal file
172
crates/helexa-router/tests/topology.rs
Normal file
@@ -0,0 +1,172 @@
|
|||||||
|
//! Topology-poller acceptance tests for #72: the router maintains a live
|
||||||
|
//! map of which cortexes serve which models, marks an unreachable/erroring
|
||||||
|
//! cortex unhealthy and excludes it from routing, and recovers it once
|
||||||
|
//! reachable again.
|
||||||
|
|
||||||
|
use axum::extract::State;
|
||||||
|
use axum::http::StatusCode;
|
||||||
|
use axum::routing::get;
|
||||||
|
use axum::{Json, Router};
|
||||||
|
use helexa_router::config::{CortexEndpoint, RouterConfig};
|
||||||
|
use helexa_router::poller::{POLL_FAILURE_THRESHOLD, poll_once};
|
||||||
|
use helexa_router::state::{RouterState, entry_feasible};
|
||||||
|
use serde_json::{Value, json};
|
||||||
|
use std::sync::Arc;
|
||||||
|
use std::sync::atomic::{AtomicBool, Ordering};
|
||||||
|
use tokio::net::TcpListener;
|
||||||
|
|
||||||
|
/// Shared "is this mock cortex up?" flag, toggled by tests to simulate
|
||||||
|
/// outage and recovery.
|
||||||
|
#[derive(Clone)]
|
||||||
|
struct MockState {
|
||||||
|
up: Arc<AtomicBool>,
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn mock_models(State(s): State<MockState>) -> Result<Json<Value>, StatusCode> {
|
||||||
|
if !s.up.load(Ordering::SeqCst) {
|
||||||
|
return Err(StatusCode::SERVICE_UNAVAILABLE);
|
||||||
|
}
|
||||||
|
Ok(Json(json!({
|
||||||
|
"object": "list",
|
||||||
|
"data": [
|
||||||
|
{
|
||||||
|
"id": "Qwen/Qwen3-Coder-30B",
|
||||||
|
"object": "model",
|
||||||
|
"created": 0,
|
||||||
|
"owned_by": "helexa",
|
||||||
|
"loaded": true,
|
||||||
|
"feasible_on": ["beast"],
|
||||||
|
"locations": [{"node": "beast", "status": "loaded", "vram_estimate_mb": 19000}]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"id": "Qwen/Qwen3-VL-8B",
|
||||||
|
"object": "model",
|
||||||
|
"created": 0,
|
||||||
|
"owned_by": "helexa",
|
||||||
|
"loaded": false,
|
||||||
|
"feasible_on": ["beast"],
|
||||||
|
"locations": []
|
||||||
|
}
|
||||||
|
]
|
||||||
|
})))
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn mock_health(State(s): State<MockState>) -> Result<Json<Value>, StatusCode> {
|
||||||
|
if !s.up.load(Ordering::SeqCst) {
|
||||||
|
return Err(StatusCode::SERVICE_UNAVAILABLE);
|
||||||
|
}
|
||||||
|
Ok(Json(json!({
|
||||||
|
"status": "ok",
|
||||||
|
"nodes": { "healthy": 2, "total": 3 }
|
||||||
|
})))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Spawn a mock cortex; returns (base_url, up_flag).
|
||||||
|
async fn spawn_mock_cortex() -> (String, Arc<AtomicBool>) {
|
||||||
|
let up = Arc::new(AtomicBool::new(true));
|
||||||
|
let state = MockState { up: up.clone() };
|
||||||
|
let app = Router::new()
|
||||||
|
.route("/v1/models", get(mock_models))
|
||||||
|
.route("/health", get(mock_health))
|
||||||
|
.with_state(state);
|
||||||
|
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||||
|
let addr = listener.local_addr().unwrap();
|
||||||
|
tokio::spawn(async move {
|
||||||
|
axum::serve(listener, app).await.unwrap();
|
||||||
|
});
|
||||||
|
(format!("http://{addr}"), up)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn state_for(name: &str, endpoint: &str) -> RouterState {
|
||||||
|
let cfg = RouterConfig {
|
||||||
|
cortexes: vec![CortexEndpoint {
|
||||||
|
name: name.into(),
|
||||||
|
endpoint: endpoint.into(),
|
||||||
|
region: None,
|
||||||
|
tls_ca: None,
|
||||||
|
}],
|
||||||
|
..Default::default()
|
||||||
|
};
|
||||||
|
RouterState::from_config(&cfg)
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn poll_builds_live_topology() {
|
||||||
|
let (base, _up) = spawn_mock_cortex().await;
|
||||||
|
let state = state_for("c1", &base);
|
||||||
|
|
||||||
|
poll_once(&state).await;
|
||||||
|
|
||||||
|
let topo = state.topology.read().await;
|
||||||
|
let c1 = topo.get("c1").expect("cortex present");
|
||||||
|
assert!(c1.reachable, "should be reachable after a good poll");
|
||||||
|
assert_eq!(c1.consecutive_failures, 0);
|
||||||
|
assert!(c1.last_poll.is_some());
|
||||||
|
assert_eq!((c1.healthy_nodes, c1.total_nodes), (2, 3));
|
||||||
|
|
||||||
|
// Loaded model: loaded + feasible. Catalogue-only model: feasible only
|
||||||
|
// (not loaded, but feasible_on non-empty).
|
||||||
|
let coder = c1.models.get("Qwen/Qwen3-Coder-30B").unwrap();
|
||||||
|
assert!(coder.loaded && entry_feasible(coder));
|
||||||
|
let vl = c1.models.get("Qwen/Qwen3-VL-8B").unwrap();
|
||||||
|
assert!(!vl.loaded && entry_feasible(vl));
|
||||||
|
drop(topo);
|
||||||
|
|
||||||
|
// The routing helper sees both serveable models on the reachable cortex.
|
||||||
|
assert_eq!(
|
||||||
|
state.cortexes_serving("Qwen/Qwen3-VL-8B").await,
|
||||||
|
vec!["c1".to_string()]
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn unreachable_cortex_excluded_then_recovers() {
|
||||||
|
let (base, up) = spawn_mock_cortex().await;
|
||||||
|
let state = state_for("c1", &base);
|
||||||
|
|
||||||
|
// Healthy first.
|
||||||
|
poll_once(&state).await;
|
||||||
|
assert!(state.topology.read().await["c1"].reachable);
|
||||||
|
|
||||||
|
// Take it down. The first failures debounce (stay reachable) until the
|
||||||
|
// threshold; only then is it excluded.
|
||||||
|
up.store(false, Ordering::SeqCst);
|
||||||
|
for i in 1..POLL_FAILURE_THRESHOLD {
|
||||||
|
poll_once(&state).await;
|
||||||
|
assert!(
|
||||||
|
state.topology.read().await["c1"].reachable,
|
||||||
|
"still reachable after {i} failure(s) (below threshold)"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
poll_once(&state).await; // crosses the threshold
|
||||||
|
{
|
||||||
|
let topo = state.topology.read().await;
|
||||||
|
assert!(!topo["c1"].reachable, "excluded after threshold failures");
|
||||||
|
assert!(topo["c1"].consecutive_failures >= POLL_FAILURE_THRESHOLD);
|
||||||
|
}
|
||||||
|
// Excluded from routing.
|
||||||
|
assert!(
|
||||||
|
state
|
||||||
|
.cortexes_serving("Qwen/Qwen3-Coder-30B")
|
||||||
|
.await
|
||||||
|
.is_empty()
|
||||||
|
);
|
||||||
|
|
||||||
|
// Bring it back: the next successful poll restores it.
|
||||||
|
up.store(true, Ordering::SeqCst);
|
||||||
|
poll_once(&state).await;
|
||||||
|
let topo = state.topology.read().await;
|
||||||
|
assert!(topo["c1"].reachable, "recovered after a good poll");
|
||||||
|
assert_eq!(topo["c1"].consecutive_failures, 0);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn unconfigured_endpoint_is_unreachable() {
|
||||||
|
// Nothing listening on this port → polls fail; below threshold it stays
|
||||||
|
// at its initial unreachable state, and never panics.
|
||||||
|
let state = state_for("dead", "http://127.0.0.1:1");
|
||||||
|
poll_once(&state).await;
|
||||||
|
let topo = state.topology.read().await;
|
||||||
|
assert!(!topo["dead"].reachable);
|
||||||
|
assert_eq!(topo["dead"].consecutive_failures, 1);
|
||||||
|
}
|
||||||
21
crates/helexa-stream/Cargo.toml
Normal file
21
crates/helexa-stream/Cargo.toml
Normal file
@@ -0,0 +1,21 @@
|
|||||||
|
[package]
|
||||||
|
name = "helexa-stream"
|
||||||
|
version.workspace = true
|
||||||
|
edition.workspace = true
|
||||||
|
license.workspace = true
|
||||||
|
repository.workspace = true
|
||||||
|
|
||||||
|
[lib]
|
||||||
|
name = "helexa_stream"
|
||||||
|
path = "src/lib.rs"
|
||||||
|
|
||||||
|
[dependencies]
|
||||||
|
axum = { workspace = true }
|
||||||
|
reqwest = { workspace = true }
|
||||||
|
futures = { workspace = true }
|
||||||
|
thiserror = { workspace = true }
|
||||||
|
|
||||||
|
[dev-dependencies]
|
||||||
|
tokio = { workspace = true }
|
||||||
|
tokio-stream = { workspace = true }
|
||||||
|
async-stream = "0.3"
|
||||||
290
crates/helexa-stream/src/lib.rs
Normal file
290
crates/helexa-stream/src/lib.rs
Normal file
@@ -0,0 +1,290 @@
|
|||||||
|
//! Shared streaming reverse-proxy mechanism (#71).
|
||||||
|
//!
|
||||||
|
//! cortex and helexa-router both need to proxy an OpenAI/Anthropic SSE
|
||||||
|
//! response from a downstream backend **verbatim** — chunks forwarded as
|
||||||
|
//! they arrive, never buffering the full body — while observing the bytes
|
||||||
|
//! for metrics/metering. This crate owns that mechanism so there is one
|
||||||
|
//! implementation, not one per tier.
|
||||||
|
//!
|
||||||
|
//! The split is mechanism vs policy:
|
||||||
|
//!
|
||||||
|
//! - **Mechanism (here):** [`forward_streaming`] POSTs to a backend and
|
||||||
|
//! streams the response body back through an [`ObservedStream`], which
|
||||||
|
//! feeds every chunk to a caller-supplied [`ChunkObserver`] and calls
|
||||||
|
//! [`ChunkObserver::finish`] exactly once on clean end-of-stream or on
|
||||||
|
//! drop (client disconnect mid-stream). [`BodyTail`] and
|
||||||
|
//! [`last_count_for`] are the reusable pieces an observer uses to pull
|
||||||
|
//! the trailing OpenAI `usage` object out of the streamed bytes.
|
||||||
|
//! - **Policy (caller):** what to *do* with the observed bytes — which
|
||||||
|
//! metric names to emit, which labels, whether to settle a per-principal
|
||||||
|
//! reservation — lives in the consumer's `ChunkObserver` impl, not here.
|
||||||
|
//!
|
||||||
|
//! The proxy is status-agnostic: a non-2xx upstream response (e.g. a
|
||||||
|
//! cortex `429 rate_limit_exceeded`) is streamed back with its status and
|
||||||
|
//! headers intact, so honest backpressure reaches the client unchanged.
|
||||||
|
//! Only a network failure or a malformed response build is an error.
|
||||||
|
|
||||||
|
use axum::body::{Body, Bytes};
|
||||||
|
use axum::http::{HeaderMap, StatusCode};
|
||||||
|
use axum::response::Response;
|
||||||
|
use futures::Stream;
|
||||||
|
use futures::stream::BoxStream;
|
||||||
|
use reqwest::Client;
|
||||||
|
use std::pin::Pin;
|
||||||
|
use std::task::{Context, Poll};
|
||||||
|
|
||||||
|
/// Observes the bytes of a streamed proxy response without altering them.
|
||||||
|
///
|
||||||
|
/// `observe` is called for each forwarded chunk; `finish` is called
|
||||||
|
/// exactly once — on clean end-of-stream or on drop — and implementations
|
||||||
|
/// must be idempotent (the [`ObservedStream`] guards against a double call,
|
||||||
|
/// but a `finish` that runs side effects should still self-guard).
|
||||||
|
pub trait ChunkObserver: Send + Unpin + 'static {
|
||||||
|
/// A body chunk has been forwarded downstream. The slice is the exact
|
||||||
|
/// bytes the client receives.
|
||||||
|
fn observe(&mut self, chunk: &[u8]);
|
||||||
|
|
||||||
|
/// The stream has ended (cleanly or via client disconnect). Called once.
|
||||||
|
fn finish(&mut self);
|
||||||
|
}
|
||||||
|
|
||||||
|
/// A bounded accumulator for the tail of a streamed body.
|
||||||
|
///
|
||||||
|
/// The OpenAI `usage` object rides on the final SSE chunk (and sits at the
|
||||||
|
/// end of a non-streaming JSON body), so retaining a generous tail is
|
||||||
|
/// enough to recover token counts via [`last_count_for`]; the cap bounds
|
||||||
|
/// memory on huge bodies. Appends are char-boundary-safe.
|
||||||
|
#[derive(Debug)]
|
||||||
|
pub struct BodyTail {
|
||||||
|
tail: String,
|
||||||
|
cap: usize,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl BodyTail {
|
||||||
|
/// Create a tail retaining at most `cap` bytes.
|
||||||
|
pub fn new(cap: usize) -> Self {
|
||||||
|
Self {
|
||||||
|
tail: String::new(),
|
||||||
|
cap,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Append a chunk, trimming from the front past the cap. When trimming,
|
||||||
|
/// the newest half is kept (the usage object is always at the very end).
|
||||||
|
pub fn push(&mut self, chunk: &[u8]) {
|
||||||
|
self.tail.push_str(&String::from_utf8_lossy(chunk));
|
||||||
|
if self.tail.len() > self.cap {
|
||||||
|
let mut cut = self.tail.len() - self.cap / 2;
|
||||||
|
while !self.tail.is_char_boundary(cut) {
|
||||||
|
cut += 1;
|
||||||
|
}
|
||||||
|
self.tail.drain(..cut);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// The retained tail text.
|
||||||
|
pub fn as_str(&self) -> &str {
|
||||||
|
&self.tail
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Find the value of the LAST `"key": <integer>` occurrence in `tail`.
|
||||||
|
///
|
||||||
|
/// Pure and chunk-boundary-safe (the tail is contiguous appended text).
|
||||||
|
/// The quoted-needle form means `completion_tokens` never matches
|
||||||
|
/// `completion_tokens_details`, and taking the last occurrence means the
|
||||||
|
/// final `usage` object wins even if content earlier in the stream echoed
|
||||||
|
/// a usage-shaped string.
|
||||||
|
pub fn last_count_for(tail: &str, key: &str) -> Option<u64> {
|
||||||
|
let needle = format!("\"{key}\"");
|
||||||
|
let mut result = None;
|
||||||
|
for (idx, _) in tail.match_indices(&needle) {
|
||||||
|
let rest = tail[idx + needle.len()..].trim_start();
|
||||||
|
let Some(rest) = rest.strip_prefix(':') else {
|
||||||
|
continue;
|
||||||
|
};
|
||||||
|
let rest = rest.trim_start();
|
||||||
|
let digits: &str = &rest[..rest
|
||||||
|
.char_indices()
|
||||||
|
.find(|(_, c)| !c.is_ascii_digit())
|
||||||
|
.map(|(i, _)| i)
|
||||||
|
.unwrap_or(rest.len())];
|
||||||
|
if let Ok(v) = digits.parse::<u64>() {
|
||||||
|
result = Some(v);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
result
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Error from [`forward_streaming`]. Distinguishes a network/transport
|
||||||
|
/// failure reaching the backend from a failure assembling the downstream
|
||||||
|
/// response. A non-2xx upstream *status* is not an error — it is streamed
|
||||||
|
/// through verbatim.
|
||||||
|
#[derive(Debug, thiserror::Error)]
|
||||||
|
pub enum StreamError {
|
||||||
|
#[error("upstream request failed")]
|
||||||
|
Upstream(reqwest::Error),
|
||||||
|
#[error("failed to build response")]
|
||||||
|
ResponseBuild(String),
|
||||||
|
}
|
||||||
|
|
||||||
|
/// POST `body` to `url` and stream the response back verbatim through
|
||||||
|
/// `observer`.
|
||||||
|
///
|
||||||
|
/// Request headers are forwarded except `host` / `content-length` (reqwest
|
||||||
|
/// sets these). The returned [`Response`] carries the upstream status and
|
||||||
|
/// headers unchanged — including non-2xx — with a body that streams the
|
||||||
|
/// upstream bytes chunk-for-chunk, feeding each chunk to `observer`.
|
||||||
|
pub async fn forward_streaming<O: ChunkObserver>(
|
||||||
|
client: &Client,
|
||||||
|
url: &str,
|
||||||
|
headers: HeaderMap,
|
||||||
|
body: Bytes,
|
||||||
|
observer: O,
|
||||||
|
) -> Result<Response, StreamError> {
|
||||||
|
let mut req_builder = client.post(url).body(body);
|
||||||
|
for (key, value) in headers.iter() {
|
||||||
|
if key == "host" || key == "content-length" {
|
||||||
|
continue; // reqwest sets these
|
||||||
|
}
|
||||||
|
req_builder = req_builder.header(key, value);
|
||||||
|
}
|
||||||
|
|
||||||
|
let upstream = req_builder.send().await.map_err(StreamError::Upstream)?;
|
||||||
|
|
||||||
|
let status =
|
||||||
|
StatusCode::from_u16(upstream.status().as_u16()).unwrap_or(StatusCode::BAD_GATEWAY);
|
||||||
|
let resp_headers = upstream.headers().clone();
|
||||||
|
|
||||||
|
let stream = ObservedStream::new(Box::pin(upstream.bytes_stream()), observer);
|
||||||
|
let body = Body::from_stream(stream);
|
||||||
|
|
||||||
|
let mut response = Response::builder().status(status);
|
||||||
|
for (key, value) in resp_headers.iter() {
|
||||||
|
response = response.header(key, value);
|
||||||
|
}
|
||||||
|
response
|
||||||
|
.body(body)
|
||||||
|
.map_err(|e| StreamError::ResponseBuild(e.to_string()))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Pass-through stream wrapper that feeds a [`ChunkObserver`]. Forwards
|
||||||
|
/// each chunk verbatim, calls `observe` per chunk, and `finish` once on
|
||||||
|
/// clean end-of-stream; the `Drop` impl covers client disconnects.
|
||||||
|
pub struct ObservedStream<O: ChunkObserver> {
|
||||||
|
inner: BoxStream<'static, Result<Bytes, reqwest::Error>>,
|
||||||
|
observer: O,
|
||||||
|
finished: bool,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl<O: ChunkObserver> ObservedStream<O> {
|
||||||
|
/// Wrap a byte stream with an observer.
|
||||||
|
pub fn new(inner: BoxStream<'static, Result<Bytes, reqwest::Error>>, observer: O) -> Self {
|
||||||
|
Self {
|
||||||
|
inner,
|
||||||
|
observer,
|
||||||
|
finished: false,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn finish(&mut self) {
|
||||||
|
if self.finished {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
self.finished = true;
|
||||||
|
self.observer.finish();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl<O: ChunkObserver> Stream for ObservedStream<O> {
|
||||||
|
type Item = Result<Bytes, reqwest::Error>;
|
||||||
|
|
||||||
|
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
|
||||||
|
let this = self.get_mut();
|
||||||
|
match this.inner.as_mut().poll_next(cx) {
|
||||||
|
Poll::Ready(Some(Ok(chunk))) => {
|
||||||
|
this.observer.observe(&chunk);
|
||||||
|
Poll::Ready(Some(Ok(chunk)))
|
||||||
|
}
|
||||||
|
Poll::Ready(Some(Err(e))) => Poll::Ready(Some(Err(e))),
|
||||||
|
Poll::Ready(None) => {
|
||||||
|
this.finish();
|
||||||
|
Poll::Ready(None)
|
||||||
|
}
|
||||||
|
Poll::Pending => Poll::Pending,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl<O: ChunkObserver> Drop for ObservedStream<O> {
|
||||||
|
fn drop(&mut self) {
|
||||||
|
self.finish();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn extracts_counts_from_final_sse_usage_chunk() {
|
||||||
|
let tail = concat!(
|
||||||
|
"data: {\"choices\":[{\"delta\":{\"content\":\"hi\"}}]}\n\n",
|
||||||
|
"data: {\"choices\":[],\"usage\":{\"prompt_tokens\":225,",
|
||||||
|
"\"completion_tokens\":42,\"total_tokens\":267}}\n\n",
|
||||||
|
"data: [DONE]\n\n"
|
||||||
|
);
|
||||||
|
assert_eq!(last_count_for(tail, "prompt_tokens"), Some(225));
|
||||||
|
assert_eq!(last_count_for(tail, "completion_tokens"), Some(42));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn extracts_counts_from_non_streaming_body() {
|
||||||
|
let tail = "{\"choices\":[{\"message\":{\"content\":\"hi\"}}],\
|
||||||
|
\"usage\":{\"prompt_tokens\": 12, \"completion_tokens\": 7}}";
|
||||||
|
assert_eq!(last_count_for(tail, "prompt_tokens"), Some(12));
|
||||||
|
assert_eq!(last_count_for(tail, "completion_tokens"), Some(7));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn ignores_details_variants_and_takes_last_occurrence() {
|
||||||
|
// completion_tokens_details must not shadow completion_tokens,
|
||||||
|
// and the LAST usage object wins (matters when content echoes
|
||||||
|
// a usage-shaped string earlier in the stream).
|
||||||
|
let tail = concat!(
|
||||||
|
"data: {\"usage\":{\"completion_tokens\":1}}\n\n",
|
||||||
|
"data: {\"usage\":{\"completion_tokens\":99,",
|
||||||
|
"\"completion_tokens_details\":{\"reasoning_tokens\":3}}}\n\n"
|
||||||
|
);
|
||||||
|
assert_eq!(last_count_for(tail, "completion_tokens"), Some(99));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn absent_keys_yield_none() {
|
||||||
|
assert_eq!(
|
||||||
|
last_count_for("data: [DONE]\n\n", "completion_tokens"),
|
||||||
|
None
|
||||||
|
);
|
||||||
|
assert_eq!(last_count_for("", "prompt_tokens"), None);
|
||||||
|
// key present but non-numeric value
|
||||||
|
assert_eq!(
|
||||||
|
last_count_for("\"completion_tokens\": null", "completion_tokens"),
|
||||||
|
None
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn body_tail_retains_usage_after_cap_trim() {
|
||||||
|
// Cap small enough that the filler forces several front-trims, but
|
||||||
|
// (as in production, where cap ≫ the usage object) large enough that
|
||||||
|
// the trailing usage object survives the newest-half retention.
|
||||||
|
let mut tail = BodyTail::new(512);
|
||||||
|
for _ in 0..100 {
|
||||||
|
tail.push(b"data: {\"choices\":[{\"delta\":{\"content\":\"x\"}}]}\n\n");
|
||||||
|
}
|
||||||
|
assert!(tail.as_str().len() <= 512, "cap must bound the tail");
|
||||||
|
tail.push(b"data: {\"usage\":{\"prompt_tokens\":5,\"completion_tokens\":9}}\n\n");
|
||||||
|
assert_eq!(last_count_for(tail.as_str(), "prompt_tokens"), Some(5));
|
||||||
|
assert_eq!(last_count_for(tail.as_str(), "completion_tokens"), Some(9));
|
||||||
|
}
|
||||||
|
}
|
||||||
162
crates/helexa-stream/tests/streaming.rs
Normal file
162
crates/helexa-stream/tests/streaming.rs
Normal file
@@ -0,0 +1,162 @@
|
|||||||
|
//! Integration tests for the shared streaming proxy (#71): proves a backend
|
||||||
|
//! SSE response is forwarded chunk-for-chunk (no buffering), the observer
|
||||||
|
//! sees every byte and finishes once, and non-2xx is streamed through with
|
||||||
|
//! its status intact — the behaviours both cortex and helexa-router rely on.
|
||||||
|
|
||||||
|
use axum::Router;
|
||||||
|
use axum::body::Body;
|
||||||
|
use axum::http::{HeaderMap, StatusCode};
|
||||||
|
use axum::response::Response;
|
||||||
|
use axum::routing::post;
|
||||||
|
use helexa_stream::{BodyTail, ChunkObserver, forward_streaming, last_count_for};
|
||||||
|
use std::sync::{Arc, Mutex};
|
||||||
|
use std::time::{Duration, Instant};
|
||||||
|
use tokio::net::TcpListener;
|
||||||
|
|
||||||
|
/// Observer that records what it saw, for assertions.
|
||||||
|
#[derive(Clone, Default)]
|
||||||
|
struct RecordingObserver {
|
||||||
|
inner: Arc<Mutex<Recorded>>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Default)]
|
||||||
|
struct Recorded {
|
||||||
|
chunks: usize,
|
||||||
|
finished: usize,
|
||||||
|
tail: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl ChunkObserver for RecordingObserver {
|
||||||
|
fn observe(&mut self, chunk: &[u8]) {
|
||||||
|
let mut r = self.inner.lock().unwrap();
|
||||||
|
r.chunks += 1;
|
||||||
|
r.tail.push_str(&String::from_utf8_lossy(chunk));
|
||||||
|
}
|
||||||
|
fn finish(&mut self) {
|
||||||
|
self.inner.lock().unwrap().finished += 1;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Mock backend that streams 5 SSE chunks with 30ms gaps, then a usage
|
||||||
|
/// chunk and `[DONE]`.
|
||||||
|
async fn sse_handler() -> Response {
|
||||||
|
let chunks: Vec<&'static str> = vec![
|
||||||
|
"data: {\"choices\":[{\"delta\":{\"content\":\"a\"}}]}\n\n",
|
||||||
|
"data: {\"choices\":[{\"delta\":{\"content\":\"b\"}}]}\n\n",
|
||||||
|
"data: {\"choices\":[{\"delta\":{\"content\":\"c\"}}]}\n\n",
|
||||||
|
"data: {\"choices\":[{\"delta\":{\"content\":\"d\"}}]}\n\n",
|
||||||
|
"data: {\"choices\":[{\"delta\":{\"content\":\"e\"}}]}\n\n",
|
||||||
|
"data: {\"choices\":[],\"usage\":{\"prompt_tokens\":11,\"completion_tokens\":5}}\n\n",
|
||||||
|
"data: [DONE]\n\n",
|
||||||
|
];
|
||||||
|
let stream = async_stream::stream! {
|
||||||
|
for c in chunks {
|
||||||
|
tokio::time::sleep(Duration::from_millis(30)).await;
|
||||||
|
yield Ok::<_, std::io::Error>(axum::body::Bytes::from_static(c.as_bytes()));
|
||||||
|
}
|
||||||
|
};
|
||||||
|
Response::new(Body::from_stream(stream))
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn rate_limited_handler() -> Response {
|
||||||
|
Response::builder()
|
||||||
|
.status(StatusCode::TOO_MANY_REQUESTS)
|
||||||
|
.body(Body::from("{\"error\":{\"type\":\"rate_limit_exceeded\"}}"))
|
||||||
|
.unwrap()
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn spawn_backend(router: Router) -> String {
|
||||||
|
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||||
|
let addr = listener.local_addr().unwrap();
|
||||||
|
tokio::spawn(async move {
|
||||||
|
axum::serve(listener, router).await.unwrap();
|
||||||
|
});
|
||||||
|
format!("http://{addr}")
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn streams_chunks_incrementally_and_observes_usage() {
|
||||||
|
let base = spawn_backend(Router::new().route("/v1/chat/completions", post(sse_handler))).await;
|
||||||
|
let observer = RecordingObserver::default();
|
||||||
|
let probe = observer.clone();
|
||||||
|
|
||||||
|
let client = reqwest::Client::new();
|
||||||
|
let resp = forward_streaming(
|
||||||
|
&client,
|
||||||
|
&format!("{base}/v1/chat/completions"),
|
||||||
|
HeaderMap::new(),
|
||||||
|
axum::body::Bytes::from_static(b"{\"model\":\"x\",\"stream\":true}"),
|
||||||
|
observer,
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
.expect("forward ok");
|
||||||
|
|
||||||
|
assert_eq!(resp.status(), StatusCode::OK);
|
||||||
|
|
||||||
|
// Read the proxied body as a stream, timestamping arrivals.
|
||||||
|
let mut body = resp.into_body().into_data_stream();
|
||||||
|
let mut arrivals: Vec<Instant> = Vec::new();
|
||||||
|
let mut collected = String::new();
|
||||||
|
use futures::StreamExt;
|
||||||
|
while let Some(item) = body.next().await {
|
||||||
|
let bytes = item.unwrap();
|
||||||
|
arrivals.push(Instant::now());
|
||||||
|
collected.push_str(&String::from_utf8_lossy(&bytes));
|
||||||
|
}
|
||||||
|
|
||||||
|
// Incremental delivery: first and last chunk are meaningfully apart
|
||||||
|
// (5×30ms gaps), proving no full-response buffering.
|
||||||
|
let spread = *arrivals.last().unwrap() - arrivals[0];
|
||||||
|
assert!(
|
||||||
|
spread >= Duration::from_millis(100),
|
||||||
|
"expected incremental delivery, spread was {spread:?}"
|
||||||
|
);
|
||||||
|
|
||||||
|
// The client received the terminator and the usage object verbatim.
|
||||||
|
assert!(collected.contains("data: [DONE]"));
|
||||||
|
|
||||||
|
// The observer saw the bytes and finished exactly once.
|
||||||
|
let r = probe.inner.lock().unwrap();
|
||||||
|
assert!(r.chunks >= 5, "observer saw {} chunks", r.chunks);
|
||||||
|
assert_eq!(r.finished, 1, "finish must run exactly once");
|
||||||
|
assert_eq!(last_count_for(&r.tail, "prompt_tokens"), Some(11));
|
||||||
|
assert_eq!(last_count_for(&r.tail, "completion_tokens"), Some(5));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn non_2xx_is_streamed_through_verbatim() {
|
||||||
|
let base =
|
||||||
|
spawn_backend(Router::new().route("/v1/chat/completions", post(rate_limited_handler)))
|
||||||
|
.await;
|
||||||
|
let observer = RecordingObserver::default();
|
||||||
|
let probe = observer.clone();
|
||||||
|
|
||||||
|
let client = reqwest::Client::new();
|
||||||
|
let resp = forward_streaming(
|
||||||
|
&client,
|
||||||
|
&format!("{base}/v1/chat/completions"),
|
||||||
|
HeaderMap::new(),
|
||||||
|
axum::body::Bytes::new(),
|
||||||
|
observer,
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
.expect("forward ok");
|
||||||
|
|
||||||
|
// Backpressure status reaches the client unchanged.
|
||||||
|
assert_eq!(resp.status(), StatusCode::TOO_MANY_REQUESTS);
|
||||||
|
let body = axum::body::to_bytes(resp.into_body(), usize::MAX)
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
assert!(String::from_utf8_lossy(&body).contains("rate_limit_exceeded"));
|
||||||
|
|
||||||
|
// finish still runs once even with a tiny non-streaming body.
|
||||||
|
assert_eq!(probe.inner.lock().unwrap().finished, 1);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn body_tail_smoke() {
|
||||||
|
let mut tail = BodyTail::new(128);
|
||||||
|
tail.push(b"hello ");
|
||||||
|
tail.push(b"world");
|
||||||
|
assert_eq!(tail.as_str(), "hello world");
|
||||||
|
}
|
||||||
@@ -11,17 +11,28 @@
|
|||||||
# the router — the router never owns an inbound TLS listener.
|
# the router — the router never owns an inbound TLS listener.
|
||||||
listen = "0.0.0.0:8088"
|
listen = "0.0.0.0:8088"
|
||||||
|
|
||||||
|
# How often (seconds) to refresh each cortex's health + /v1/models topology.
|
||||||
|
# poll_interval_secs = 10
|
||||||
|
|
||||||
# -- Downstream cortexes -------------------------------------------------
|
# -- Downstream cortexes -------------------------------------------------
|
||||||
# Each [[cortexes]] entry is an operator-run cortex the router may dispatch
|
# Each [[cortexes]] entry is an operator-run cortex the router may dispatch
|
||||||
# to. The router forwards the client's bearer verbatim (auth stays at
|
# to. The router forwards the client's bearer verbatim (auth stays at
|
||||||
# cortex) and routes on capacity. Outbound TLS to each cortex is verified.
|
# cortex) and routes on capacity (preferring matching `region`).
|
||||||
#
|
#
|
||||||
# The skeleton only loads this list; capacity/catalogue polling and
|
# Outbound TLS pinning (optional): set `tls_ca` to a PEM trust anchor that
|
||||||
# capacity-aware dispatch arrive in later issues.
|
# enrols this cortex — the CA (or self-signed cert) its TLS cert must chain
|
||||||
|
# to. The router then trusts ONLY that anchor for this cortex (platform
|
||||||
|
# roots disabled), so the router->cortex hop (which carries the client's
|
||||||
|
# bearer) reaches the cert you expect and a rogue endpoint presenting any
|
||||||
|
# other cert is rejected at the handshake. A cortex whose `tls_ca` fails to
|
||||||
|
# load is disabled (fail closed). Omit `tls_ca` for a publicly-trusted cert
|
||||||
|
# or plaintext http:// on a private (e.g. WireGuard) network.
|
||||||
|
|
||||||
# [[cortexes]]
|
# [[cortexes]]
|
||||||
# name = "lair-cafe"
|
# name = "lair-cafe"
|
||||||
# endpoint = "https://cortex.lair.cafe"
|
# endpoint = "https://cortex.lair.cafe"
|
||||||
|
# region = "eu-west"
|
||||||
|
# tls_ca = "/etc/helexa-router/pins/lair-cafe.pem"
|
||||||
|
|
||||||
# [[cortexes]]
|
# [[cortexes]]
|
||||||
# name = "example-operator"
|
# name = "example-operator"
|
||||||
|
|||||||
Reference in New Issue
Block a user