Compare commits

...

5 Commits

Author SHA1 Message Date
1115bb0942 feat(#74): verify downstream cortex TLS certs (outbound pinning)
All checks were successful
CI / Format (push) Successful in 41s
CI / CUDA type-check (push) Successful in 1m34s
CI / Clippy (push) Successful in 2m23s
CI / Test (push) Successful in 5m19s
CI / Build cortex SRPM (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
The router is a TLS client to cortexes; the router->cortex hop crosses
the helexa->operator boundary carrying the client's bearer. This pins
that hop to an enrolled cert.

Trust mechanism (the open question): per-cortex enrolled trust anchor.
Each [[cortexes]] entry gets an optional `tls_ca` — a PEM CA (or
self-signed cert) the cortex's TLS cert must chain to. When set, the
router builds a client that trusts ONLY that anchor (platform roots
disabled), so the cortex must present the expected cert and a rogue
endpoint with any other (even publicly-valid) cert is rejected at the
handshake. Enrolment = the operator hands helexa the cortex's cert,
referenced by path in router config. This is the natural model for
self-hosted operators behind their own nginx/private CA, and reuses the
reqwest public API (no custom rustls verifier, no new TLS backend).

- `RouterState` now holds a per-cortex `reqwest::Client` map
  (`client_for`), replacing the single shared client; poller and dispatch
  use the per-cortex client. `build_client(tls_ca)` is the builder.
- Fail closed: a `tls_ca` that can't load omits the cortex from the
  client map — it's never polled or routed to, rather than silently
  degrading to unpinned TLS. The poller treats a missing client (and a
  rejected handshake) as a failed poll, so #72's existing reachability
  debounce excludes it.

Tests (`tls.rs`, 4): a live tokio-rustls HTTPS server proves a client
enrolled with the server's cert is accepted (200) while clients pinned to
a different cert — or using default roots — are rejected; the poller
marks a wrong-cert cortex unreachable while a correctly-enrolled one is
reachable; a missing pin file disables the cortex (fail closed); garbage
PEM is rejected at build. Existing suites updated for the per-cortex
client + new config field.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01F6o3ddqmYNh9kzdwq6eowh
2026-06-21 21:23:20 +03:00
63f578cb15 feat(#75): aggregate /v1/models across operators (federation catalogue)
All checks were successful
CI / Format (push) Successful in 35s
CI / CUDA type-check (push) Successful in 1m38s
CI / Clippy (push) Successful in 2m18s
CI / Test (push) Successful in 4m54s
CI / Build cortex SRPM (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
The router's /v1/models is now the deduped union of every reachable
cortex's 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").

To preserve per-model limit/cost, the topology poller now retains each
cortex's full `cortex_core::node::CortexModelEntry` (was distilled to a
{loaded, feasible} bool). `entry_feasible()` replaces the dropped field;
dispatch (#73) and `cortexes_serving` use it — no routing behaviour
change.

`catalogue.rs::aggregate_models`:
- Dedupe by model id; a model served by >=1 reachable cortex appears once.
- Merge availability: `loaded` OR across operators; only feasible
  (loaded-or-cold-loadable) entries surface — a catalogue-only model no
  neuron can host is hidden.
- Re-tier to operator names: `feasible_on` becomes the cortexes that can
  serve it and `locations` the operators it's loaded on (node = cortex
  name), so the federation view doesn't leak each operator's neuron names
  or per-device VRAM.
- Conflict resolution: `limit` → tightest (smallest context, so a client
  never overflows the most-constrained operator); `cost` → cheapest
  (the federation "from" price). Richer range/region policy couples to
  #68, noted as follow-up.

Tests: 4 unit (dedupe+merge, unreachable excluded, infeasible hidden,
tightest-limit+cheapest-cost) + 1 end-to-end (two mock cortexes
overlapping on a model → GET /v1/models over HTTP asserts the merged
union). dispatch/topology suites updated for the entry-storage change.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01F6o3ddqmYNh9kzdwq6eowh
2026-06-21 21:08:16 +03:00
76c90fa993 Merge feat/73-capacity-aware-dispatch: capacity-aware dispatch + region affinity + failover (#73)
All checks were successful
build-prerelease / Resolve version stamps + change detection (push) Successful in 32s
build-prerelease / Build neuron-blackwell (push) Successful in 1m38s
build-prerelease / Build helexa-bench binary (push) Successful in 2m13s
build-prerelease / Build neuron-ampere (push) Successful in 2m19s
build-prerelease / Build neuron-ada (push) Successful in 2m20s
build-prerelease / Build cortex binary (push) Successful in 2m24s
build-prerelease / Lint (fmt + clippy) (push) Successful in 3m6s
build-prerelease / Test (push) Successful in 5m41s
build-prerelease / Package helexa-bench RPM (push) Successful in 1m21s
build-prerelease / Package cortex RPM (push) Successful in 1m19s
build-prerelease / Package helexa-neuron-ampere RPM (push) Successful in 1m33s
build-prerelease / Package helexa-neuron-ada RPM (push) Successful in 1m34s
build-prerelease / Package helexa-neuron-blackwell RPM (push) Successful in 1m37s
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Successful in 53s
2026-06-21 19:48:13 +03:00
7984d27553 feat(#73): capacity-aware dispatch with region affinity + failover
All checks were successful
CI / Format (push) Successful in 40s
CI / CUDA type-check (push) Successful in 1m37s
CI / Clippy (push) Successful in 2m16s
CI / Test (push) Successful in 4m53s
CI / Build cortex SRPM (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
The router's data path. Wires the topology poller (#72) and the shared
streaming proxy (#71) into real request routing.

- `dispatch.rs`: `select_cortexes(model)` ranks reachable cortexes that
  can serve the model, best-first — loaded/warm before cold-loadable,
  region match before not, more healthy nodes before fewer, name for
  determinism. `dispatch()` extracts `model`, picks candidates, and
  forwards via `helexa_stream::forward_streaming` (bearer + bytes
  verbatim, SSE streamed back). Cortex's #63 rejections (429/400/…) pass
  through untouched; transport failures fail over to the next candidate;
  a genuine HTTP response — any status — is returned as-is, never retried
  away.
- Router-originated rejections use the #63 envelope: 404 model_not_found
  (no operator serves it), 503 service_unavailable + Retry-After (known
  but all unreachable / all candidates failed to connect), 400
  missing_model_field. `error.rs` is the router's envelope→axum adapter
  (mirrors cortex-gateway's).
- `handlers.rs`: `/v1/chat/completions`, `/v1/completions`,
  `/v1/responses`, `/v1/messages` dispatch to the same path on a chosen
  cortex. The router holds zero entitlement logic — routes on capacity,
  not budget.
- Config: optional `region` on the router and per-cortex for geo affinity.

Tests (`dispatch.rs`): routes to a serving cortex + forwards the bearer;
cortex 429 passes through and is NOT retried; transport failure fails
over to a live cortex; unknown→404, known-but-unreachable→503,
missing-model→400; ranking order (warm/region/headroom). 7 new, existing
skeleton/topology suites unchanged.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01F6o3ddqmYNh9kzdwq6eowh
2026-06-21 19:40:07 +03:00
43ffffdccb Merge feat/72-router-topology-poller: router↔cortex capacity & catalogue poller (#72)
All checks were successful
build-prerelease / Resolve version stamps + change detection (push) Successful in 34s
build-prerelease / Lint (fmt + clippy) (push) Successful in 2m18s
build-prerelease / Build neuron-blackwell (push) Successful in 1m43s
build-prerelease / Build helexa-bench binary (push) Successful in 2m6s
build-prerelease / Build neuron-ada (push) Successful in 2m17s
build-prerelease / Build neuron-ampere (push) Successful in 2m18s
build-prerelease / Build cortex binary (push) Successful in 2m18s
build-prerelease / Test (push) Successful in 4m50s
build-prerelease / Package helexa-bench RPM (push) Successful in 1m19s
build-prerelease / Package cortex RPM (push) Successful in 1m25s
build-prerelease / Package helexa-neuron-ada RPM (push) Successful in 1m39s
build-prerelease / Package helexa-neuron-ampere RPM (push) Successful in 1m41s
build-prerelease / Package helexa-neuron-blackwell RPM (push) Successful in 1m40s
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Successful in 55s
2026-06-21 19:09:13 +03:00
16 changed files with 1383 additions and 64 deletions

37
Cargo.lock generated
View File

@@ -1933,10 +1933,15 @@ dependencies = [
"clap",
"cortex-core",
"figment",
"helexa-stream",
"rcgen",
"reqwest",
"rustls",
"serde",
"serde_json",
"thiserror 2.0.18",
"tokio",
"tokio-rustls",
"tower-http",
"tracing",
"tracing-subscriber",
@@ -2995,6 +3000,16 @@ dependencies = [
"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]]
name = "percent-encoding"
version = "2.3.2"
@@ -3361,6 +3376,19 @@ dependencies = [
"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]]
name = "reborrow"
version = "0.5.5"
@@ -5248,6 +5276,15 @@ version = "1.0.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "cfe53a6657fd280eaa890a3bc59152892ffa3e30101319d168b781ed6529b049"
[[package]]
name = "yasna"
version = "0.5.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e17bb3549cc1321ae1296b9cdc2698e2b6cb1992adfa19a8c72e5b7a738f44cd"
dependencies = [
"time",
]
[[package]]
name = "yoke"
version = "0.7.5"

View File

@@ -15,6 +15,7 @@ path = "src/lib.rs"
[dependencies]
cortex-core = { workspace = true }
helexa-stream = { path = "../helexa-stream" }
tokio = { workspace = true }
axum = { workspace = true }
@@ -24,6 +25,7 @@ serde = { workspace = true }
serde_json = { workspace = true }
figment = { workspace = true }
anyhow = { workspace = true }
thiserror = { workspace = true }
clap = { workspace = true }
tracing = { workspace = true }
tracing-subscriber = { workspace = true }
@@ -32,3 +34,8 @@ chrono = { workspace = true }
[dev-dependencies]
# Jail (isolated cwd + env) for config tests.
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"

View 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);
}
}

View File

@@ -32,6 +32,11 @@ pub struct RouterSettings {
/// 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 {
@@ -41,12 +46,31 @@ fn default_poll_interval_secs() -> u64 {
/// 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
/// 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 {
/// Human-readable label (e.g. "lair-cafe").
pub name: String,
/// Base URL of the cortex gateway (e.g. "https://cortex.example.com").
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 {
@@ -68,6 +92,7 @@ impl Default for RouterConfig {
router: RouterSettings {
listen: "0.0.0.0:8088".into(),
poll_interval_secs: default_poll_interval_secs(),
region: None,
},
cortexes: vec![],
}

View 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) {}
}

View 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
}

View File

@@ -1,20 +1,63 @@
use crate::state::RouterState;
use axum::{Json, Router, extract::State, routing::get};
use cortex_core::openai::ModelsResponse;
use crate::{catalogue, dispatch};
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 std::sync::Arc;
/// Routes served by the router skeleton. The inference paths
/// (`/v1/chat/completions`, `/v1/messages`, …) arrive with capacity-aware
/// dispatch (#73); for now the router only answers `/health` and a stub
/// `/v1/models`.
/// Routes served by the router. Inference paths are capacity-aware-dispatched
/// (#73) to a downstream cortex; `/health` and a stub `/v1/models` are local.
pub fn api_routes() -> Router<Arc<RouterState>> {
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("/health", get(health))
.route("/", get(health))
}
// ── Inference paths — forwarded verbatim to a chosen cortex ──────────
//
// 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
@@ -32,12 +75,15 @@ async fn health(State(state): State<Arc<RouterState>>) -> Json<Value> {
}))
}
/// `GET /v1/models` — empty catalogue stub. The real cross-operator union
/// (catalogue × topology feasibility, aggregated from each cortex) is the
/// federation-catalogue issue (#75).
async fn list_models() -> Json<ModelsResponse> {
Json(ModelsResponse {
object: "list".into(),
data: vec![],
})
/// `GET /v1/models` — the federation catalogue (#75): the deduped union of
/// every reachable cortex's `/v1/models`, so a client doing discovery
/// against the router resolves the whole federation without knowing about
/// operators or cortexes.
async fn list_models(State(state): State<Arc<RouterState>>) -> Json<Value> {
let topo = state.topology.read().await;
let data: Vec<Value> = catalogue::aggregate_models(&topo)
.iter()
.map(|e| json!(e))
.collect();
Json(json!({ "object": "list", "data": data }))
}

View File

@@ -11,7 +11,10 @@
//! on capacity (epic #69). A background [`poller`] keeps a live
//! per-cortex topology (#72) that the dispatcher (#73) will route on.
pub mod catalogue;
pub mod config;
pub mod dispatch;
pub mod error;
pub mod handlers;
pub mod poller;
pub mod state;

View File

@@ -8,7 +8,7 @@
//! then flipped unhealthy and excluded from routing; it recovers on the
//! next successful poll.
use crate::state::{RouterModelStatus, RouterState};
use crate::state::RouterState;
use chrono::Utc;
use cortex_core::node::CortexModelEntry;
use serde::Deserialize;
@@ -62,7 +62,19 @@ pub async fn poll_once(state: &RouterState) {
/// 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) {
let models = fetch_models(state, endpoint).await;
// 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 {
@@ -71,19 +83,7 @@ async fn poll_cortex(state: &RouterState, name: &str, endpoint: &str) {
match models {
Ok(models) => {
entry.models = models
.into_iter()
.map(|m| {
let feasible = m.loaded || !m.feasible_on.is_empty();
(
m.id,
RouterModelStatus {
loaded: m.loaded,
feasible,
},
)
})
.collect();
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());
@@ -106,7 +106,7 @@ async fn poll_cortex(state: &RouterState, name: &str, endpoint: &str) {
drop(topo);
// Best-effort health (node counts). Never flips reachability.
if let Some((healthy, total)) = fetch_health(state, endpoint).await {
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;
@@ -117,12 +117,11 @@ async fn poll_cortex(state: &RouterState, name: &str, endpoint: &str) {
/// GET `/v1/models`, returning the parsed entries or a short failure reason.
async fn fetch_models(
state: &RouterState,
client: &reqwest::Client,
endpoint: &str,
) -> Result<Vec<CortexModelEntry>, &'static str> {
let url = format!("{endpoint}/v1/models");
let resp = state
.http_client
let resp = client
.get(&url)
.timeout(POLL_TIMEOUT)
.send()
@@ -140,15 +139,9 @@ async fn fetch_models(
/// GET `/health`, returning `(healthy, total)` node counts. `None` on any
/// failure — the caller leaves the previous counts in place.
async fn fetch_health(state: &RouterState, endpoint: &str) -> Option<(u32, u32)> {
async fn fetch_health(client: &reqwest::Client, endpoint: &str) -> Option<(u32, u32)> {
let url = format!("{endpoint}/health");
let resp = state
.http_client
.get(&url)
.timeout(POLL_TIMEOUT)
.send()
.await
.ok()?;
let resp = client.get(&url).timeout(POLL_TIMEOUT).send().await.ok()?;
if !resp.status().is_success() {
return None;
}

View File

@@ -1,5 +1,6 @@
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;
@@ -14,8 +15,15 @@ use tokio::sync::RwLock;
pub struct RouterState {
/// Downstream cortex endpoints, as configured.
pub cortexes: Vec<CortexEndpoint>,
/// Shared client for polling (and, later, proxying to) cortexes.
pub http_client: reqwest::Client,
/// 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
@@ -40,19 +48,16 @@ pub struct CortexTopology {
/// load signal; #73 refines headroom). 0/0 until first health poll.
pub healthy_nodes: u32,
pub total_nodes: u32,
/// Per-model serveability, keyed by model id, from `/v1/models`.
pub models: HashMap<String, RouterModelStatus>,
/// 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>,
}
/// What a cortex can do with one model, distilled from its `/v1/models`
/// entry to the two facts the router routes on.
#[derive(Debug, Clone)]
pub struct RouterModelStatus {
/// The model is loaded on at least one of the cortex's neurons.
pub loaded: bool,
/// The cortex can serve it — loaded now, or feasible to cold-load
/// (catalogue × topology says some neuron can host it).
pub feasible: bool,
/// 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 {
@@ -63,14 +68,42 @@ impl RouterState {
.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 {
cortexes: config.cortexes.clone(),
http_client: reqwest::Client::new(),
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.
@@ -78,8 +111,34 @@ impl RouterState {
let topo = self.topology.read().await;
topo.iter()
.filter(|(_, t)| t.reachable)
.filter(|(_, t)| t.models.get(model_id).is_some_and(|m| m.feasible))
.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),
}

View 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"));
}

View 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"]);
}

View File

@@ -31,10 +31,14 @@ async fn health_reports_configured_cortex_count() {
CortexEndpoint {
name: "a".into(),
endpoint: "https://a.example.com".into(),
region: None,
tls_ca: None,
},
CortexEndpoint {
name: "b".into(),
endpoint: "https://b.example.com".into(),
region: None,
tls_ca: None,
},
])
.await;

View 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());
}

View File

@@ -9,7 +9,7 @@ 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;
use helexa_router::state::{RouterState, entry_feasible};
use serde_json::{Value, json};
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
@@ -82,6 +82,8 @@ fn state_for(name: &str, endpoint: &str) -> RouterState {
cortexes: vec![CortexEndpoint {
name: name.into(),
endpoint: endpoint.into(),
region: None,
tls_ca: None,
}],
..Default::default()
};
@@ -102,11 +104,12 @@ async fn poll_builds_live_topology() {
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.
// 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 && coder.feasible);
assert!(coder.loaded && entry_feasible(coder));
let vl = c1.models.get("Qwen/Qwen3-VL-8B").unwrap();
assert!(!vl.loaded && vl.feasible);
assert!(!vl.loaded && entry_feasible(vl));
drop(topo);
// The routing helper sees both serveable models on the reachable cortex.

View File

@@ -17,14 +17,22 @@ listen = "0.0.0.0:8088"
# -- Downstream cortexes -------------------------------------------------
# Each [[cortexes]] entry is an operator-run cortex the router may dispatch
# 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
# capacity-aware dispatch arrive in later issues.
# Outbound TLS pinning (optional): set `tls_ca` to a PEM trust anchor that
# 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]]
# name = "lair-cafe"
# endpoint = "https://cortex.lair.cafe"
# region = "eu-west"
# tls_ca = "/etc/helexa-router/pins/lair-cafe.pem"
# [[cortexes]]
# name = "example-operator"