Compare commits
21 Commits
feat/70-ro
...
feat/B5-to
| Author | SHA1 | Date | |
|---|---|---|---|
|
2348cc2234
|
|||
|
f2ba12bbc5
|
|||
|
d94c62c143
|
|||
|
cb9e7c7c2e
|
|||
|
2604b9f134
|
|||
|
178e3092d5
|
|||
|
46befde4cd
|
|||
|
cf87e156c5
|
|||
|
79073170ec
|
|||
|
71106afaf1
|
|||
|
d2dcdd6ebb
|
|||
|
222c2a6116
|
|||
|
1115bb0942
|
|||
|
63f578cb15
|
|||
|
76c90fa993
|
|||
|
7984d27553
|
|||
|
43ffffdccb
|
|||
|
5fd7736abd
|
|||
|
03fd4960c3
|
|||
|
5ed6bc3390
|
|||
|
cabec1d08a
|
3
.gitignore
vendored
3
.gitignore
vendored
@@ -1,6 +1,9 @@
|
||||
/target
|
||||
/bench/node_modules
|
||||
/bench/dist
|
||||
/helexa.ai/node_modules
|
||||
/helexa.ai/dist
|
||||
helexa.ai/.env.local
|
||||
*.swp
|
||||
*.swo
|
||||
.idea/
|
||||
|
||||
855
Cargo.lock
generated
855
Cargo.lock
generated
File diff suppressed because it is too large
Load Diff
@@ -8,6 +8,8 @@ members = [
|
||||
"crates/helexa-acp",
|
||||
"crates/helexa-bench",
|
||||
"crates/helexa-router",
|
||||
"crates/helexa-stream",
|
||||
"crates/helexa-upstream",
|
||||
]
|
||||
|
||||
[workspace.package]
|
||||
|
||||
@@ -6,6 +6,7 @@ license.workspace = true
|
||||
|
||||
[dependencies]
|
||||
cortex-core.workspace = true
|
||||
helexa-stream = { path = "../helexa-stream" }
|
||||
async-trait.workspace = true
|
||||
tokio.workspace = true
|
||||
axum.workspace = true
|
||||
|
||||
@@ -1,21 +1,27 @@
|
||||
//! Streaming HTTP reverse proxy to neuron backends.
|
||||
//!
|
||||
//! For streaming requests, SSE chunks are forwarded as they arrive.
|
||||
//! The proxy captures timing information for metrics but does not
|
||||
//! buffer the full response.
|
||||
//! The streaming *mechanism* — forward an SSE body chunk-for-chunk without
|
||||
//! buffering, observing the bytes for metrics — lives in the shared
|
||||
//! [`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 anyhow::Result;
|
||||
use axum::body::Body;
|
||||
use axum::http::{HeaderMap, StatusCode};
|
||||
use axum::http::HeaderMap;
|
||||
use axum::http::StatusCode;
|
||||
use axum::response::{IntoResponse, Response};
|
||||
use futures::Stream;
|
||||
use futures::stream::BoxStream;
|
||||
use helexa_stream::{BodyTail, ChunkObserver, StreamError};
|
||||
use reqwest::Client;
|
||||
use std::pin::Pin;
|
||||
use std::task::{Context, Poll};
|
||||
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.
|
||||
///
|
||||
/// Logging contract: every call emits exactly one structured event at
|
||||
@@ -42,66 +48,41 @@ pub async fn forward_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.
|
||||
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_resp = match req_builder.send().await {
|
||||
Ok(r) => r,
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
let response = helexa_stream::forward_streaming(client, &url, headers, body, observer)
|
||||
.await
|
||||
.map_err(|e| {
|
||||
match &e {
|
||||
StreamError::Upstream(err) => tracing::warn!(
|
||||
node = %route.node_name,
|
||||
url = %url,
|
||||
error = %e,
|
||||
error = %err,
|
||||
"proxy: upstream request failed (network)"
|
||||
);
|
||||
return Err(ProxyError::Upstream(e));
|
||||
),
|
||||
StreamError::ResponseBuild(err) => tracing::warn!(
|
||||
node = %route.node_name,
|
||||
url = %url,
|
||||
error = %err,
|
||||
"proxy: failed to build response"
|
||||
),
|
||||
}
|
||||
};
|
||||
ProxyError::from(e)
|
||||
})?;
|
||||
|
||||
let upstream_status = upstream_resp.status();
|
||||
if !upstream_status.is_success() {
|
||||
if !response.status().is_success() {
|
||||
// Streaming body — can't snippet without breaking the stream
|
||||
// pass-through. Log status + URL; the client still gets the
|
||||
// upstream status, just without the leaked body.
|
||||
tracing::warn!(
|
||||
node = %route.node_name,
|
||||
url = %url,
|
||||
status = upstream_status.as_u16(),
|
||||
status = response.status().as_u16(),
|
||||
"proxy: upstream returned non-2xx"
|
||||
);
|
||||
}
|
||||
|
||||
let status = StatusCode::from_u16(upstream_status.as_u16()).unwrap_or(StatusCode::BAD_GATEWAY);
|
||||
|
||||
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())
|
||||
})
|
||||
Ok(response)
|
||||
}
|
||||
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
@@ -112,6 +93,15 @@ pub enum ProxyError {
|
||||
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 {
|
||||
fn into_response(self) -> Response {
|
||||
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
|
||||
// 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
|
||||
// OpenAI `usage` object (present on the last SSE chunk and on
|
||||
// non-streaming JSON bodies alike) yields engine-truth token counts.
|
||||
// time and keeps a bounded tail of the body text (via the shared
|
||||
// `helexa_stream::BodyTail`), from which the final OpenAI `usage` object
|
||||
// (present on the last SSE chunk and on non-streaming JSON bodies alike)
|
||||
// yields engine-truth token counts.
|
||||
//
|
||||
// Emitted per request, labelled {model, node}:
|
||||
// cortex_time_to_first_token_seconds (histogram) — first body chunk
|
||||
@@ -155,37 +146,15 @@ impl IntoResponse for ProxyError {
|
||||
/// non-streaming bodies.
|
||||
const TAIL_CAP_BYTES: usize = 64 * 1024;
|
||||
|
||||
/// 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`.
|
||||
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 {
|
||||
/// cortex's [`ChunkObserver`]: per-request token metrics plus the
|
||||
/// per-principal reservation settle. Drives cortex policy over the shared
|
||||
/// streaming mechanism.
|
||||
struct CortexMetrics {
|
||||
labels: [(&'static str, String); 2],
|
||||
request_start: Instant,
|
||||
first_chunk: Option<Instant>,
|
||||
last_chunk: Option<Instant>,
|
||||
tail: String,
|
||||
tail: BodyTail,
|
||||
finished: bool,
|
||||
/// Per-principal metering hook (#51). Invoked exactly once in `finish`
|
||||
/// with the observed `(prompt, completion)` so the reservation can be
|
||||
@@ -193,7 +162,7 @@ struct TokenMetrics {
|
||||
usage_sink: Option<crate::metering::UsageSink>,
|
||||
}
|
||||
|
||||
impl TokenMetrics {
|
||||
impl CortexMetrics {
|
||||
fn new(
|
||||
model_id: &str,
|
||||
node_name: &str,
|
||||
@@ -208,26 +177,19 @@ impl TokenMetrics {
|
||||
request_start,
|
||||
first_chunk: None,
|
||||
last_chunk: None,
|
||||
tail: String::new(),
|
||||
tail: BodyTail::new(TAIL_CAP_BYTES),
|
||||
finished: false,
|
||||
usage_sink,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl ChunkObserver for CortexMetrics {
|
||||
fn observe(&mut self, chunk: &[u8]) {
|
||||
let now = Instant::now();
|
||||
self.first_chunk.get_or_insert(now);
|
||||
self.last_chunk = Some(now);
|
||||
self.tail.push_str(&String::from_utf8_lossy(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);
|
||||
}
|
||||
self.tail.push(chunk);
|
||||
}
|
||||
|
||||
/// Emit the metrics exactly once — called on clean stream end and
|
||||
@@ -239,8 +201,8 @@ impl TokenMetrics {
|
||||
}
|
||||
self.finished = true;
|
||||
|
||||
let prompt = last_count_for(&self.tail, "prompt_tokens");
|
||||
let completion = last_count_for(&self.tail, "completion_tokens");
|
||||
let prompt = last_count_for(self.tail.as_str(), "prompt_tokens");
|
||||
let completion = last_count_for(self.tail.as_str(), "completion_tokens");
|
||||
|
||||
// Per-model metrics — only when body chunks actually arrived.
|
||||
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]
|
||||
cortex-core = { workspace = true }
|
||||
helexa-stream = { path = "../helexa-stream" }
|
||||
|
||||
tokio = { workspace = true }
|
||||
axum = { workspace = true }
|
||||
@@ -24,10 +25,17 @@ 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 }
|
||||
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"
|
||||
|
||||
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
|
||||
/// inbound TLS listener.
|
||||
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
|
||||
/// 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 {
|
||||
@@ -58,6 +91,8 @@ impl Default for RouterConfig {
|
||||
Self {
|
||||
router: RouterSettings {
|
||||
listen: "0.0.0.0:8088".into(),
|
||||
poll_interval_secs: default_poll_interval_secs(),
|
||||
region: None,
|
||||
},
|
||||
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 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))
|
||||
}
|
||||
|
||||
/// `GET /health` — liveness plus the configured downstream cortex count.
|
||||
/// Real per-cortex reachability lands with the poller (#72).
|
||||
// ── 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
|
||||
/// 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> {
|
||||
let topo = state.topology.read().await;
|
||||
let reachable = topo.values().filter(|t| t.reachable).count();
|
||||
Json(json!({
|
||||
"status": "ok",
|
||||
"cortexes": {
|
||||
"configured": state.cortexes.len(),
|
||||
"reachable": reachable,
|
||||
}
|
||||
}))
|
||||
}
|
||||
|
||||
/// `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 }))
|
||||
}
|
||||
|
||||
@@ -8,12 +8,15 @@
|
||||
//!
|
||||
//! It holds **zero entitlement logic** — auth/budget stays at cortex
|
||||
//! (epic #47); the router forwards the client bearer unchanged and routes
|
||||
//! on capacity (epic #69). This crate is the binary skeleton (#70):
|
||||
//! a plaintext axum server reusing `cortex-core` types, serving `/health`
|
||||
//! and a stub `/v1/models`.
|
||||
//! 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;
|
||||
|
||||
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.
|
||||
pub async fn run(config: RouterConfig) -> Result<()> {
|
||||
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>()?;
|
||||
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 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
|
||||
/// config. Live multi-operator topology (per-cortex capacity + catalogue)
|
||||
/// is added by the poller (#72), at which point this grows an
|
||||
/// `Arc<RwLock<...>>` topology map alongside the static endpoints.
|
||||
/// This is the router tier of the fractal neuron ← cortex ← router design:
|
||||
/// just as cortex polls each neuron for capacity/catalogue, the router
|
||||
/// polls each cortex's `/health` + `/v1/models`.
|
||||
#[derive(Debug)]
|
||||
pub struct RouterState {
|
||||
/// Downstream cortex endpoints, as configured.
|
||||
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 {
|
||||
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 {
|
||||
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 {
|
||||
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;
|
||||
|
||||
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");
|
||||
}
|
||||
63
crates/helexa-upstream/Cargo.toml
Normal file
63
crates/helexa-upstream/Cargo.toml
Normal file
@@ -0,0 +1,63 @@
|
||||
[package]
|
||||
name = "helexa-upstream"
|
||||
version.workspace = true
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
repository.workspace = true
|
||||
|
||||
[[bin]]
|
||||
name = "helexa-upstream"
|
||||
path = "src/main.rs"
|
||||
|
||||
[lib]
|
||||
name = "helexa_upstream"
|
||||
path = "src/lib.rs"
|
||||
|
||||
[dependencies]
|
||||
tokio = { workspace = true }
|
||||
axum = { workspace = true }
|
||||
tower-http = { workspace = true }
|
||||
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 }
|
||||
chrono = { workspace = true }
|
||||
|
||||
# PostgreSQL — the mesh authority's system of record. Runtime query API
|
||||
# (not the compile-time `query!` macros) so the crate builds in CI without a
|
||||
# live database or a committed `.sqlx` offline cache; correctness is covered
|
||||
# by the gated integration tests. (Macro adoption is a later refinement once
|
||||
# a dev DB + offline cache exist.)
|
||||
sqlx = { version = "0.8", default-features = false, features = [
|
||||
"runtime-tokio",
|
||||
"tls-rustls",
|
||||
"postgres",
|
||||
"macros",
|
||||
"migrate",
|
||||
"chrono",
|
||||
"uuid",
|
||||
] }
|
||||
uuid = { version = "1", features = ["v4", "serde"] }
|
||||
sha2 = "0.10"
|
||||
subtle = "2.6"
|
||||
# Web auth (B4): argon2id password hashing, JWT sessions, CSPRNG secrets,
|
||||
# transactional email.
|
||||
argon2 = "0.5"
|
||||
jsonwebtoken = "9"
|
||||
rand = "0.8"
|
||||
lettre = { version = "0.11", default-features = false, features = [
|
||||
"tokio1-rustls-tls",
|
||||
"smtp-transport",
|
||||
"builder",
|
||||
] }
|
||||
|
||||
# cortex-core for the shared #63 OpenAiError envelope on the authz surface.
|
||||
cortex-core = { workspace = true }
|
||||
|
||||
[dev-dependencies]
|
||||
figment = { workspace = true, features = ["test"] }
|
||||
reqwest = { workspace = true }
|
||||
137
crates/helexa-upstream/migrations/0001_init.sql
Normal file
137
crates/helexa-upstream/migrations/0001_init.sql
Normal file
@@ -0,0 +1,137 @@
|
||||
-- helexa-upstream initial schema (#59): accounts, keys, ledger, top-up
|
||||
-- codes, served-usage. The mesh-level authority's system of record.
|
||||
--
|
||||
-- Token amounts are BIGINT (i64) throughout; the cortex EntitlementProvider
|
||||
-- carries u64 but mesh allocations sit comfortably inside i64 and Postgres
|
||||
-- has no unsigned type.
|
||||
|
||||
CREATE EXTENSION IF NOT EXISTS citext;
|
||||
CREATE EXTENSION IF NOT EXISTS pgcrypto;
|
||||
|
||||
-- ── Users (web auth: email + password) ──────────────────────────────
|
||||
CREATE TABLE users (
|
||||
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
|
||||
email CITEXT NOT NULL UNIQUE,
|
||||
password_hash TEXT NOT NULL, -- argon2id PHC string
|
||||
email_verified BOOLEAN NOT NULL DEFAULT FALSE,
|
||||
-- Browser fingerprint captured at registration (#abuse). Best-effort,
|
||||
-- client-supplied; the primary signal for silent multi-account
|
||||
-- detection. NULL when the client could not produce one.
|
||||
registration_fingerprint TEXT,
|
||||
created_at TIMESTAMPTZ NOT NULL DEFAULT now()
|
||||
);
|
||||
CREATE INDEX users_registration_fingerprint_idx
|
||||
ON users (registration_fingerprint)
|
||||
WHERE registration_fingerprint IS NOT NULL;
|
||||
|
||||
-- Single-use email tokens for verification and password reset. Only the
|
||||
-- sha256 of the emailed secret is stored.
|
||||
CREATE TABLE email_tokens (
|
||||
token_hash BYTEA PRIMARY KEY,
|
||||
user_id UUID NOT NULL REFERENCES users(id) ON DELETE CASCADE,
|
||||
kind TEXT NOT NULL CHECK (kind IN ('verify', 'reset')),
|
||||
expires_at TIMESTAMPTZ NOT NULL,
|
||||
consumed_at TIMESTAMPTZ
|
||||
);
|
||||
CREATE INDEX email_tokens_user_idx ON email_tokens (user_id);
|
||||
|
||||
-- ── Accounts (the billable allocation ledger) ───────────────────────
|
||||
CREATE TABLE accounts (
|
||||
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
|
||||
owner_user_id UUID NOT NULL REFERENCES users(id) ON DELETE CASCADE,
|
||||
allocation_total BIGINT NOT NULL DEFAULT 0,
|
||||
allocation_spent BIGINT NOT NULL DEFAULT 0,
|
||||
allocation_reserved BIGINT NOT NULL DEFAULT 0,
|
||||
-- 'deactivated' is the SILENT abuse flag: keys stop authorizing but no
|
||||
-- surface ever tells the user why (see resolve → 401).
|
||||
status TEXT NOT NULL DEFAULT 'active'
|
||||
CHECK (status IN ('active', 'deactivated')),
|
||||
-- This account shares a registration fingerprint with >= 1 other.
|
||||
fingerprint_flagged BOOLEAN NOT NULL DEFAULT FALSE,
|
||||
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
|
||||
-- The no-overshoot backstop to the atomic reserve UPDATE.
|
||||
CONSTRAINT accounts_no_overshoot
|
||||
CHECK (allocation_spent + allocation_reserved <= allocation_total),
|
||||
CONSTRAINT accounts_nonneg
|
||||
CHECK (allocation_spent >= 0 AND allocation_reserved >= 0)
|
||||
);
|
||||
CREATE INDEX accounts_owner_idx ON accounts (owner_user_id);
|
||||
|
||||
-- ── API keys (Principal.key_id = api_keys.id) ───────────────────────
|
||||
CREATE TABLE api_keys (
|
||||
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
|
||||
account_id UUID NOT NULL REFERENCES accounts(id) ON DELETE CASCADE,
|
||||
key_hash BYTEA NOT NULL, -- sha256(raw key)
|
||||
key_prefix TEXT NOT NULL, -- non-secret display prefix
|
||||
label TEXT NOT NULL DEFAULT '',
|
||||
status TEXT NOT NULL DEFAULT 'active'
|
||||
CHECK (status IN ('active', 'archived')),
|
||||
-- Per-key sub-cap: 'hardcap' = absolute tokens; 'percent' = % of the
|
||||
-- account's allocation_total (resolved to an absolute at reserve time).
|
||||
limit_kind TEXT NOT NULL DEFAULT 'percent'
|
||||
CHECK (limit_kind IN ('percent', 'hardcap')),
|
||||
limit_value BIGINT NOT NULL DEFAULT 100,
|
||||
-- serde of cortex_core::entitlements::CapWindow (Balance | Rolling).
|
||||
cap_window JSONB NOT NULL DEFAULT '{"kind":"balance"}'::jsonb,
|
||||
-- Per-key running ledger (mirrors the account ledger; Balance semantics
|
||||
-- in this migration — rolling-window reset lands with the authz API).
|
||||
key_spent BIGINT NOT NULL DEFAULT 0,
|
||||
key_reserved BIGINT NOT NULL DEFAULT 0,
|
||||
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
|
||||
CONSTRAINT api_keys_key_nonneg
|
||||
CHECK (key_spent >= 0 AND key_reserved >= 0)
|
||||
);
|
||||
-- A raw key resolves only while active; the hash is unique among active keys.
|
||||
CREATE UNIQUE INDEX api_keys_active_hash_idx
|
||||
ON api_keys (key_hash) WHERE status = 'active';
|
||||
CREATE INDEX api_keys_account_idx ON api_keys (account_id);
|
||||
|
||||
-- ── Reservations (reserve → settle/release) ─────────────────────────
|
||||
-- id is BIGSERIAL so it maps to the cortex Reservation.id (u64) verbatim,
|
||||
-- with the Postgres sequence as the sole global authority.
|
||||
CREATE TABLE reservations (
|
||||
id BIGSERIAL PRIMARY KEY,
|
||||
account_id UUID NOT NULL REFERENCES accounts(id) ON DELETE CASCADE,
|
||||
key_id UUID NOT NULL REFERENCES api_keys(id) ON DELETE CASCADE,
|
||||
reserved BIGINT NOT NULL,
|
||||
actual BIGINT,
|
||||
state TEXT NOT NULL DEFAULT 'open'
|
||||
CHECK (state IN ('open', 'settled', 'released')),
|
||||
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
|
||||
settled_at TIMESTAMPTZ
|
||||
);
|
||||
-- The sweeper scans open reservations by age.
|
||||
CREATE INDEX reservations_open_idx
|
||||
ON reservations (created_at) WHERE state = 'open';
|
||||
|
||||
-- ── Top-up codes (hybrid allocation) ────────────────────────────────
|
||||
CREATE TABLE top_up_codes (
|
||||
code_hash BYTEA PRIMARY KEY, -- sha256(raw code)
|
||||
value BIGINT NOT NULL, -- tokens this code grants
|
||||
denomination TEXT, -- human label (e.g. "small")
|
||||
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
|
||||
redeemed_by UUID REFERENCES accounts(id) ON DELETE SET NULL,
|
||||
redeemed_at TIMESTAMPTZ
|
||||
);
|
||||
|
||||
-- ── Served-usage ledger (#58 reconciliation) ────────────────────────
|
||||
-- Absolute per-(operator, account, key, period) served tokens, upserted by
|
||||
-- each cortex; reconciliation rolls these up for operator compensation.
|
||||
CREATE TABLE served_usage (
|
||||
operator_id TEXT NOT NULL,
|
||||
account_id UUID NOT NULL,
|
||||
key_id UUID NOT NULL,
|
||||
period DATE NOT NULL,
|
||||
served_tokens BIGINT NOT NULL DEFAULT 0,
|
||||
reconciled_at TIMESTAMPTZ,
|
||||
PRIMARY KEY (operator_id, account_id, key_id, period)
|
||||
);
|
||||
|
||||
-- ── Web sessions (DB-backed; alt/complement to stateless JWT) ───────
|
||||
CREATE TABLE sessions (
|
||||
token_hash BYTEA PRIMARY KEY, -- sha256(session token)
|
||||
user_id UUID NOT NULL REFERENCES users(id) ON DELETE CASCADE,
|
||||
expires_at TIMESTAMPTZ NOT NULL,
|
||||
created_at TIMESTAMPTZ NOT NULL DEFAULT now()
|
||||
);
|
||||
CREATE INDEX sessions_user_idx ON sessions (user_id);
|
||||
284
crates/helexa-upstream/src/authz.rs
Normal file
284
crates/helexa-upstream/src/authz.rs
Normal file
@@ -0,0 +1,284 @@
|
||||
//! `/authz/v1` — the machine surface cortex's `UpstreamEntitlementProvider`
|
||||
//! (#57) consumes. It mirrors the `cortex_core::entitlements::EntitlementProvider`
|
||||
//! trait 1:1 (resolve / reserve / settle / release / snapshot) over the B1
|
||||
//! ledger.
|
||||
//!
|
||||
//! Contract notes for the cortex client:
|
||||
//! - A **non-2xx** response means the authority could not give an
|
||||
//! authoritative answer (bad caller auth, malformed request, server
|
||||
//! error) → the client should **fail closed**.
|
||||
//! - `reserve` returns **200** whether granted or budget-refused: the body
|
||||
//! carries either `reservation_id` or a `rejected` discriminant. A budget
|
||||
//! refusal is an authoritative answer, not a transport failure.
|
||||
//! - Rejections that are genuinely auth failures use the #63 `OpenAiError`
|
||||
//! envelope so they can be surfaced verbatim.
|
||||
|
||||
use crate::crypto::sha256;
|
||||
use crate::error::envelope_response;
|
||||
use crate::ledger::{self, LedgerError};
|
||||
use crate::state::AppState;
|
||||
use axum::extract::{Request, State};
|
||||
use axum::http::{StatusCode, header};
|
||||
use axum::middleware::Next;
|
||||
use axum::response::{IntoResponse, Response};
|
||||
use axum::routing::post;
|
||||
use axum::{Json, Router};
|
||||
use cortex_core::error_envelope::OpenAiError;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use subtle::ConstantTimeEq;
|
||||
use uuid::Uuid;
|
||||
|
||||
/// The operator a validated client bearer identifies (served-usage
|
||||
/// attribution, #58). Inserted into request extensions by [`client_auth`].
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct OperatorId(pub String);
|
||||
|
||||
/// Build the `/authz/v1` router with the client-auth layer applied.
|
||||
pub fn router(state: &AppState) -> Router<AppState> {
|
||||
Router::new()
|
||||
.route("/authz/v1/resolve", post(resolve))
|
||||
.route("/authz/v1/reserve", post(reserve))
|
||||
.route("/authz/v1/settle", post(settle))
|
||||
.route("/authz/v1/release", post(release))
|
||||
.route("/authz/v1/snapshot", post(snapshot))
|
||||
.layer(axum::middleware::from_fn_with_state(
|
||||
state.clone(),
|
||||
client_auth,
|
||||
))
|
||||
}
|
||||
|
||||
// ── client auth (shared bearer → operator_id) ───────────────────────
|
||||
|
||||
/// Validate the caller's `Authorization: Bearer` against the configured
|
||||
/// client tokens (constant-time) and stamp the `operator_id`. When no tokens
|
||||
/// are configured the surface is open (dev) and a synthetic operator is
|
||||
/// used.
|
||||
async fn client_auth(State(state): State<AppState>, mut req: Request, next: Next) -> Response {
|
||||
let tokens = &state.config.client_auth.tokens;
|
||||
if tokens.is_empty() {
|
||||
req.extensions_mut().insert(OperatorId("dev".into()));
|
||||
return next.run(req).await;
|
||||
}
|
||||
let presented = req
|
||||
.headers()
|
||||
.get(header::AUTHORIZATION)
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.and_then(|v| v.strip_prefix("Bearer "))
|
||||
.map(str::trim)
|
||||
.unwrap_or("");
|
||||
let matched = tokens
|
||||
.iter()
|
||||
.find(|t| t.token.as_bytes().ct_eq(presented.as_bytes()).into());
|
||||
match matched {
|
||||
Some(t) => {
|
||||
req.extensions_mut()
|
||||
.insert(OperatorId(t.operator_id.clone()));
|
||||
next.run(req).await
|
||||
}
|
||||
None => envelope_response(OpenAiError::invalid_api_key(
|
||||
"missing or invalid client credentials",
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
// ── DTOs ────────────────────────────────────────────────────────────
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct ResolveReq {
|
||||
api_key: String,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct PrincipalDto {
|
||||
account_id: String,
|
||||
key_id: String,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct SnapshotDto {
|
||||
hard_cap: Option<i64>,
|
||||
spent: i64,
|
||||
reserved: i64,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct ResolveResp {
|
||||
principal: PrincipalDto,
|
||||
snapshot: SnapshotDto,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct ReserveReq {
|
||||
account_id: String,
|
||||
key_id: String,
|
||||
max_tokens: i64,
|
||||
}
|
||||
|
||||
#[derive(Serialize, Default)]
|
||||
struct ReserveResp {
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
reservation_id: Option<i64>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
rejected: Option<Rejection>,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
#[serde(tag = "kind", rename_all = "snake_case")]
|
||||
enum Rejection {
|
||||
InsufficientQuota {
|
||||
requested: i64,
|
||||
available: i64,
|
||||
},
|
||||
// Part of the frozen wire contract so the cortex client (#57) can map it
|
||||
// without a later breaking change. Not yet constructed: the B1 ledger
|
||||
// implements Balance caps only; rolling-window key sub-caps (which yield
|
||||
// this) land in a follow-up.
|
||||
#[allow(dead_code)]
|
||||
RateLimited {
|
||||
requested: i64,
|
||||
available: i64,
|
||||
retry_after_secs: u64,
|
||||
},
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct SettleReq {
|
||||
reservation_id: i64,
|
||||
actual_tokens: i64,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct ReservationRef {
|
||||
reservation_id: i64,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct SnapshotReq {
|
||||
account_id: String,
|
||||
key_id: String,
|
||||
}
|
||||
|
||||
// ── handlers ────────────────────────────────────────────────────────
|
||||
|
||||
/// `POST /authz/v1/resolve` — bearer key → principal + snapshot, or
|
||||
/// `401 invalid_api_key` (also for a deactivated account: no clue).
|
||||
async fn resolve(State(state): State<AppState>, Json(req): Json<ResolveReq>) -> Response {
|
||||
match ledger::resolve_key(&state.pool, &sha256(&req.api_key)).await {
|
||||
Ok(Some(p)) => Json(ResolveResp {
|
||||
principal: PrincipalDto {
|
||||
account_id: p.account_id.to_string(),
|
||||
key_id: p.key_id.to_string(),
|
||||
},
|
||||
snapshot: SnapshotDto {
|
||||
hard_cap: Some(p.hard_cap),
|
||||
spent: p.key_spent,
|
||||
reserved: p.key_reserved,
|
||||
},
|
||||
})
|
||||
.into_response(),
|
||||
Ok(None) => envelope_response(OpenAiError::invalid_api_key("invalid or unknown API key")),
|
||||
Err(e) => {
|
||||
tracing::error!(error = %e, "resolve query failed");
|
||||
envelope_response(OpenAiError::service_unavailable("authority error", Some(5)))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// `POST /authz/v1/reserve` — 200 with `reservation_id` (granted) or
|
||||
/// `rejected` (budget). Non-2xx only for bad input / server error.
|
||||
async fn reserve(State(state): State<AppState>, Json(req): Json<ReserveReq>) -> Response {
|
||||
let (Ok(account_id), Ok(key_id)) = (
|
||||
Uuid::parse_str(&req.account_id),
|
||||
Uuid::parse_str(&req.key_id),
|
||||
) else {
|
||||
return bad_request("account_id and key_id must be UUIDs");
|
||||
};
|
||||
match ledger::reserve(&state.pool, account_id, key_id, req.max_tokens).await {
|
||||
Ok(reservation_id) => Json(ReserveResp {
|
||||
reservation_id: Some(reservation_id),
|
||||
rejected: None,
|
||||
})
|
||||
.into_response(),
|
||||
Err(LedgerError::InsufficientQuota {
|
||||
requested,
|
||||
available,
|
||||
}) => Json(ReserveResp {
|
||||
reservation_id: None,
|
||||
rejected: Some(Rejection::InsufficientQuota {
|
||||
requested,
|
||||
available,
|
||||
}),
|
||||
})
|
||||
.into_response(),
|
||||
Err(LedgerError::AccountNotFound | LedgerError::KeyNotFound) => {
|
||||
// Resolve succeeded earlier; the principal vanished (archived /
|
||||
// deactivated). Treat as no budget — fail closed at the client.
|
||||
Json(ReserveResp {
|
||||
reservation_id: None,
|
||||
rejected: Some(Rejection::InsufficientQuota {
|
||||
requested: req.max_tokens,
|
||||
available: 0,
|
||||
}),
|
||||
})
|
||||
.into_response()
|
||||
}
|
||||
Err(LedgerError::Db(e)) => {
|
||||
tracing::error!(error = %e, "reserve failed");
|
||||
envelope_response(OpenAiError::service_unavailable("authority error", Some(5)))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// `POST /authz/v1/settle` — idempotent; `204`.
|
||||
async fn settle(State(state): State<AppState>, Json(req): Json<SettleReq>) -> Response {
|
||||
match ledger::settle(&state.pool, req.reservation_id, req.actual_tokens).await {
|
||||
Ok(()) => StatusCode::NO_CONTENT.into_response(),
|
||||
Err(e) => {
|
||||
tracing::error!(error = %e, "settle failed");
|
||||
envelope_response(OpenAiError::service_unavailable("authority error", Some(5)))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// `POST /authz/v1/release` — idempotent; `204`.
|
||||
async fn release(State(state): State<AppState>, Json(req): Json<ReservationRef>) -> Response {
|
||||
match ledger::release(&state.pool, req.reservation_id).await {
|
||||
Ok(()) => StatusCode::NO_CONTENT.into_response(),
|
||||
Err(e) => {
|
||||
tracing::error!(error = %e, "release failed");
|
||||
envelope_response(OpenAiError::service_unavailable("authority error", Some(5)))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// `POST /authz/v1/snapshot` — `{hard_cap, spent, reserved}` or `404`.
|
||||
async fn snapshot(State(state): State<AppState>, Json(req): Json<SnapshotReq>) -> Response {
|
||||
let (Ok(account_id), Ok(key_id)) = (
|
||||
Uuid::parse_str(&req.account_id),
|
||||
Uuid::parse_str(&req.key_id),
|
||||
) else {
|
||||
return bad_request("account_id and key_id must be UUIDs");
|
||||
};
|
||||
match ledger::snapshot(&state.pool, account_id, key_id).await {
|
||||
Ok(Some((hard_cap, spent, reserved))) => Json(SnapshotDto {
|
||||
hard_cap: Some(hard_cap),
|
||||
spent,
|
||||
reserved,
|
||||
})
|
||||
.into_response(),
|
||||
Ok(None) => StatusCode::NOT_FOUND.into_response(),
|
||||
Err(e) => {
|
||||
tracing::error!(error = %e, "snapshot failed");
|
||||
envelope_response(OpenAiError::service_unavailable("authority error", Some(5)))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn bad_request(msg: &str) -> Response {
|
||||
envelope_response(OpenAiError::new(
|
||||
400,
|
||||
"invalid_request_error",
|
||||
"invalid_request",
|
||||
msg,
|
||||
))
|
||||
}
|
||||
261
crates/helexa-upstream/src/config.rs
Normal file
261
crates/helexa-upstream/src/config.rs
Normal file
@@ -0,0 +1,261 @@
|
||||
//! helexa-upstream configuration: loaded from `helexa-upstream.toml` with
|
||||
//! figment, `UPSTREAM_`-prefixed env overrides (mirrors the cortex/router
|
||||
//! convention, e.g. `UPSTREAM_SERVER__LISTEN`, `UPSTREAM_DB__URL`).
|
||||
|
||||
use figment::{
|
||||
Figment,
|
||||
providers::{Env, Format, Toml},
|
||||
};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::path::Path;
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct UpstreamConfig {
|
||||
#[serde(default)]
|
||||
pub server: ServerSettings,
|
||||
pub db: DbSettings,
|
||||
#[serde(default)]
|
||||
pub grant: GrantSettings,
|
||||
#[serde(default)]
|
||||
pub abuse: AbuseSettings,
|
||||
#[serde(default)]
|
||||
pub client_auth: ClientAuthSettings,
|
||||
#[serde(default)]
|
||||
pub authz: AuthzSettings,
|
||||
#[serde(default)]
|
||||
pub auth: AuthSettings,
|
||||
#[serde(default)]
|
||||
pub email: EmailSettings,
|
||||
}
|
||||
|
||||
/// `[auth]` — web-session signing + token lifetimes (B4). Web sessions are
|
||||
/// JWTs, distinct from inference API keys.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct AuthSettings {
|
||||
/// HMAC secret for signing session JWTs. MUST be overridden in prod
|
||||
/// (env `UPSTREAM_AUTH__JWT_SECRET`); the default is dev-only.
|
||||
#[serde(default = "default_jwt_secret")]
|
||||
pub jwt_secret: String,
|
||||
/// Session token lifetime (seconds).
|
||||
#[serde(default = "default_session_ttl")]
|
||||
pub session_ttl_secs: u64,
|
||||
/// Email verification / password-reset token lifetime (seconds).
|
||||
#[serde(default = "default_email_token_ttl")]
|
||||
pub email_token_ttl_secs: u64,
|
||||
/// Public base URL of the frontend, used to build verify/reset links.
|
||||
#[serde(default = "default_app_base_url")]
|
||||
pub app_base_url: String,
|
||||
}
|
||||
|
||||
impl Default for AuthSettings {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
jwt_secret: default_jwt_secret(),
|
||||
session_ttl_secs: default_session_ttl(),
|
||||
email_token_ttl_secs: default_email_token_ttl(),
|
||||
app_base_url: default_app_base_url(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// `[email]` — transactional email transport for verify/reset.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct EmailSettings {
|
||||
/// `"log"` (dev — logs the link) or `"smtp"`.
|
||||
#[serde(default = "default_email_provider")]
|
||||
pub provider: String,
|
||||
/// SMTP relay URL (e.g. "smtp://user:pass@host:587") when provider=smtp.
|
||||
#[serde(default)]
|
||||
pub smtp_url: Option<String>,
|
||||
/// `From:` address.
|
||||
#[serde(default = "default_from_addr")]
|
||||
pub from_addr: String,
|
||||
}
|
||||
|
||||
impl Default for EmailSettings {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
provider: default_email_provider(),
|
||||
smtp_url: None,
|
||||
from_addr: default_from_addr(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// `[client_auth]` — credentials operators' cortexes present to `/authz/v1`.
|
||||
/// Each token maps to an `operator_id` (served-usage attribution, #58). This
|
||||
/// transport credential is distinct from end-user API keys (which ride in
|
||||
/// the `resolve` body). v2 adds mTLS.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
|
||||
pub struct ClientAuthSettings {
|
||||
/// When empty the authz surface is **open** (dev only; logged at warn).
|
||||
#[serde(default)]
|
||||
pub tokens: Vec<ClientToken>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct ClientToken {
|
||||
/// Shared bearer a cortex presents.
|
||||
pub token: String,
|
||||
/// Operator this token identifies.
|
||||
pub operator_id: String,
|
||||
}
|
||||
|
||||
/// `[authz]` — reservation lifecycle knobs.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct AuthzSettings {
|
||||
/// Open reservations older than this are swept (released), self-healing
|
||||
/// a reservation whose settle/release from cortex was lost.
|
||||
#[serde(default = "default_reservation_ttl")]
|
||||
pub reservation_ttl_secs: u64,
|
||||
/// How often the sweeper runs.
|
||||
#[serde(default = "default_sweep_interval")]
|
||||
pub sweep_interval_secs: u64,
|
||||
}
|
||||
|
||||
impl Default for AuthzSettings {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
reservation_ttl_secs: default_reservation_ttl(),
|
||||
sweep_interval_secs: default_sweep_interval(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct ServerSettings {
|
||||
/// Address to listen on (e.g. "0.0.0.0:8090"). Plaintext — edge nginx
|
||||
/// terminates TLS, consistent with the rest of the stack.
|
||||
#[serde(default = "default_listen")]
|
||||
pub listen: String,
|
||||
}
|
||||
|
||||
impl Default for ServerSettings {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
listen: default_listen(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct DbSettings {
|
||||
/// PostgreSQL connection URL (e.g. "postgres://user:pass@host/helexa").
|
||||
pub url: String,
|
||||
/// Max pool connections.
|
||||
#[serde(default = "default_max_connections")]
|
||||
pub max_connections: u32,
|
||||
}
|
||||
|
||||
/// `[grant]` — the flat free token grant every email-verified account
|
||||
/// receives (the floor of the hybrid allocation model; top-up codes extend
|
||||
/// it).
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct GrantSettings {
|
||||
#[serde(default = "default_free_grant")]
|
||||
pub free_token_grant: i64,
|
||||
}
|
||||
|
||||
impl Default for GrantSettings {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
free_token_grant: default_free_grant(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// `[abuse]` — silent multi-account abuse detection. When at least
|
||||
/// `fingerprint_account_threshold` accounts share one registration
|
||||
/// fingerprint, all of them are silently deactivated (no notice to the
|
||||
/// user; deactivation only surfaces as ordinary inference rejections).
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct AbuseSettings {
|
||||
#[serde(default = "default_fingerprint_threshold")]
|
||||
pub fingerprint_account_threshold: i64,
|
||||
}
|
||||
|
||||
impl Default for AbuseSettings {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
fingerprint_account_threshold: default_fingerprint_threshold(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn default_listen() -> String {
|
||||
"0.0.0.0:8090".into()
|
||||
}
|
||||
fn default_max_connections() -> u32 {
|
||||
16
|
||||
}
|
||||
fn default_free_grant() -> i64 {
|
||||
1_000_000
|
||||
}
|
||||
fn default_fingerprint_threshold() -> i64 {
|
||||
5
|
||||
}
|
||||
fn default_reservation_ttl() -> u64 {
|
||||
120
|
||||
}
|
||||
fn default_sweep_interval() -> u64 {
|
||||
60
|
||||
}
|
||||
fn default_jwt_secret() -> String {
|
||||
"dev-insecure-change-me".into()
|
||||
}
|
||||
fn default_session_ttl() -> u64 {
|
||||
7 * 24 * 3600
|
||||
}
|
||||
fn default_email_token_ttl() -> u64 {
|
||||
24 * 3600
|
||||
}
|
||||
fn default_app_base_url() -> String {
|
||||
"http://localhost:5173".into()
|
||||
}
|
||||
fn default_email_provider() -> String {
|
||||
"log".into()
|
||||
}
|
||||
fn default_from_addr() -> String {
|
||||
"helexa <no-reply@helexa.ai>".into()
|
||||
}
|
||||
|
||||
impl UpstreamConfig {
|
||||
/// Load from a TOML file with `UPSTREAM_`-prefixed env overrides
|
||||
/// (`__` nesting separator).
|
||||
pub fn load(path: impl AsRef<Path>) -> Result<Self, Box<figment::Error>> {
|
||||
Figment::new()
|
||||
.merge(Toml::file(path))
|
||||
.merge(Env::prefixed("UPSTREAM_").split("__"))
|
||||
.extract()
|
||||
.map_err(Box::new)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
#[allow(clippy::result_large_err)]
|
||||
fn loads_toml_with_env_override_and_defaults() {
|
||||
figment::Jail::expect_with(|jail| {
|
||||
jail.create_file(
|
||||
"helexa-upstream.toml",
|
||||
r#"
|
||||
[db]
|
||||
url = "postgres://localhost/helexa"
|
||||
"#,
|
||||
)?;
|
||||
jail.set_env("UPSTREAM_SERVER__LISTEN", "127.0.0.1:9099");
|
||||
|
||||
let cfg = UpstreamConfig::load("helexa-upstream.toml").expect("load");
|
||||
assert_eq!(cfg.server.listen, "127.0.0.1:9099");
|
||||
assert_eq!(cfg.db.url, "postgres://localhost/helexa");
|
||||
// Defaults applied when sections omitted.
|
||||
assert_eq!(cfg.grant.free_token_grant, 1_000_000);
|
||||
assert_eq!(cfg.abuse.fingerprint_account_threshold, 5);
|
||||
assert_eq!(cfg.db.max_connections, 16);
|
||||
Ok(())
|
||||
});
|
||||
}
|
||||
}
|
||||
113
crates/helexa-upstream/src/crypto.rs
Normal file
113
crates/helexa-upstream/src/crypto.rs
Normal file
@@ -0,0 +1,113 @@
|
||||
//! Hashing + secret-generation helpers.
|
||||
//!
|
||||
//! - **Passwords** (low-entropy) → argon2id PHC strings.
|
||||
//! - **API keys / top-up codes / email + session tokens** (high-entropy
|
||||
//! secrets minted here) → stored only as their sha256; sha256 is the fast,
|
||||
//! sufficient choice for high-entropy material.
|
||||
|
||||
use argon2::Argon2;
|
||||
use argon2::password_hash::rand_core::OsRng as ArgonOsRng;
|
||||
use argon2::password_hash::{PasswordHash, PasswordHasher, PasswordVerifier, SaltString};
|
||||
use rand::RngCore;
|
||||
use sha2::{Digest, Sha256};
|
||||
|
||||
/// sha256 of `input`, as raw bytes (matches the `BYTEA` columns).
|
||||
pub fn sha256(input: &str) -> Vec<u8> {
|
||||
let mut h = Sha256::new();
|
||||
h.update(input.as_bytes());
|
||||
h.finalize().to_vec()
|
||||
}
|
||||
|
||||
/// Hash a password with argon2id, returning a PHC string for storage.
|
||||
pub fn hash_password(password: &str) -> Result<String, argon2::password_hash::Error> {
|
||||
let salt = SaltString::generate(&mut ArgonOsRng);
|
||||
Ok(Argon2::default()
|
||||
.hash_password(password.as_bytes(), &salt)?
|
||||
.to_string())
|
||||
}
|
||||
|
||||
/// Verify a password against a stored PHC hash. `false` on any mismatch or
|
||||
/// malformed hash (never panics).
|
||||
pub fn verify_password(password: &str, phc: &str) -> bool {
|
||||
match PasswordHash::new(phc) {
|
||||
Ok(parsed) => Argon2::default()
|
||||
.verify_password(password.as_bytes(), &parsed)
|
||||
.is_ok(),
|
||||
Err(_) => false,
|
||||
}
|
||||
}
|
||||
|
||||
/// A fresh URL-safe high-entropy secret (256 bits) for email/session/reset
|
||||
/// tokens. The caller stores only `sha256` of this and emails/returns the
|
||||
/// raw value.
|
||||
pub fn random_token() -> String {
|
||||
let mut bytes = [0u8; 32];
|
||||
rand::rngs::OsRng.fill_bytes(&mut bytes);
|
||||
base62(&bytes)
|
||||
}
|
||||
|
||||
/// Mint a new API key: `(raw, prefix)`. `raw` is shown to the user once;
|
||||
/// only `sha256(raw)` is stored. The prefix is a non-secret display tag.
|
||||
pub fn generate_api_key() -> (String, String) {
|
||||
let mut bytes = [0u8; 32];
|
||||
rand::rngs::OsRng.fill_bytes(&mut bytes);
|
||||
let raw = format!("sk-helexa-{}", base62(&bytes));
|
||||
// Non-secret prefix for the dashboard list (scheme + first few chars).
|
||||
let prefix: String = raw.chars().take(14).collect();
|
||||
(raw, prefix)
|
||||
}
|
||||
|
||||
/// base62 encode (0-9A-Za-z) — URL/clipboard friendly, no padding.
|
||||
fn base62(bytes: &[u8]) -> String {
|
||||
const ALPHABET: &[u8] = b"0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz";
|
||||
// Treat the bytes as a big-endian integer and base62 it. 32 bytes → ~43
|
||||
// chars. Simple repeated-division over a big-uint built from the bytes.
|
||||
let mut digits: Vec<u8> = vec![0];
|
||||
for &byte in bytes {
|
||||
let mut carry = byte as u32;
|
||||
for d in digits.iter_mut() {
|
||||
let v = (*d as u32) * 256 + carry;
|
||||
*d = (v % 62) as u8;
|
||||
carry = v / 62;
|
||||
}
|
||||
while carry > 0 {
|
||||
digits.push((carry % 62) as u8);
|
||||
carry /= 62;
|
||||
}
|
||||
}
|
||||
digits
|
||||
.iter()
|
||||
.rev()
|
||||
.map(|&d| ALPHABET[d as usize] as char)
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn password_round_trips_and_rejects_wrong() {
|
||||
let phc = hash_password("correct horse").unwrap();
|
||||
assert!(verify_password("correct horse", &phc));
|
||||
assert!(!verify_password("wrong", &phc));
|
||||
assert!(!verify_password("correct horse", "not-a-phc-string"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn api_key_has_scheme_prefix_and_unique_body() {
|
||||
let (raw, prefix) = generate_api_key();
|
||||
assert!(raw.starts_with("sk-helexa-"));
|
||||
assert!(prefix.starts_with("sk-helexa-"));
|
||||
let (raw2, _) = generate_api_key();
|
||||
assert_ne!(raw, raw2, "keys are unique");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn random_tokens_are_unique_and_nonempty() {
|
||||
let a = random_token();
|
||||
let b = random_token();
|
||||
assert!(!a.is_empty());
|
||||
assert_ne!(a, b);
|
||||
}
|
||||
}
|
||||
20
crates/helexa-upstream/src/db.rs
Normal file
20
crates/helexa-upstream/src/db.rs
Normal file
@@ -0,0 +1,20 @@
|
||||
//! PostgreSQL pool + embedded migrations.
|
||||
|
||||
use anyhow::{Context, Result};
|
||||
use sqlx::postgres::{PgPool, PgPoolOptions};
|
||||
|
||||
/// Connect to Postgres and run embedded migrations (`./migrations`).
|
||||
pub async fn connect_and_migrate(url: &str, max_connections: u32) -> Result<PgPool> {
|
||||
let pool = PgPoolOptions::new()
|
||||
.max_connections(max_connections)
|
||||
.connect(url)
|
||||
.await
|
||||
.with_context(|| "connecting to PostgreSQL")?;
|
||||
|
||||
sqlx::migrate!("./migrations")
|
||||
.run(&pool)
|
||||
.await
|
||||
.with_context(|| "running migrations")?;
|
||||
|
||||
Ok(pool)
|
||||
}
|
||||
64
crates/helexa-upstream/src/email.rs
Normal file
64
crates/helexa-upstream/src/email.rs
Normal file
@@ -0,0 +1,64 @@
|
||||
//! Transactional email for verification + password-reset links.
|
||||
//!
|
||||
//! Two transports: `Log` (dev — writes the link to the log so flows are
|
||||
//! testable without a relay) and `Smtp` (lettre over rustls). Built from
|
||||
//! `[email]` config.
|
||||
|
||||
use crate::config::EmailSettings;
|
||||
use anyhow::{Context, Result};
|
||||
use lettre::message::Mailbox;
|
||||
use lettre::{AsyncSmtpTransport, AsyncTransport, Message, Tokio1Executor};
|
||||
|
||||
#[derive(Clone)]
|
||||
pub enum EmailSender {
|
||||
/// Dev: log the message instead of sending.
|
||||
Log { from: String },
|
||||
Smtp {
|
||||
from: String,
|
||||
transport: AsyncSmtpTransport<Tokio1Executor>,
|
||||
},
|
||||
}
|
||||
|
||||
impl EmailSender {
|
||||
pub fn from_config(cfg: &EmailSettings) -> Result<Self> {
|
||||
match cfg.provider.as_str() {
|
||||
"smtp" => {
|
||||
let url = cfg
|
||||
.smtp_url
|
||||
.as_deref()
|
||||
.context("[email].smtp_url required when provider = \"smtp\"")?;
|
||||
let transport = AsyncSmtpTransport::<Tokio1Executor>::from_url(url)
|
||||
.context("parsing [email].smtp_url")?
|
||||
.build();
|
||||
Ok(EmailSender::Smtp {
|
||||
from: cfg.from_addr.clone(),
|
||||
transport,
|
||||
})
|
||||
}
|
||||
_ => Ok(EmailSender::Log {
|
||||
from: cfg.from_addr.clone(),
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
/// Send a plaintext email. Errors are returned but the caller treats
|
||||
/// send failures as non-fatal to the request (the user can re-request).
|
||||
pub async fn send(&self, to: &str, subject: &str, body: &str) -> Result<()> {
|
||||
match self {
|
||||
EmailSender::Log { from } => {
|
||||
tracing::info!(%from, %to, %subject, body, "EMAIL (log transport)");
|
||||
Ok(())
|
||||
}
|
||||
EmailSender::Smtp { from, transport } => {
|
||||
let msg = Message::builder()
|
||||
.from(from.parse::<Mailbox>().context("parsing from_addr")?)
|
||||
.to(to.parse::<Mailbox>().context("parsing recipient")?)
|
||||
.subject(subject)
|
||||
.body(body.to_string())
|
||||
.context("building message")?;
|
||||
transport.send(msg).await.context("sending email")?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
21
crates/helexa-upstream/src/error.rs
Normal file
21
crates/helexa-upstream/src/error.rs
Normal file
@@ -0,0 +1,21 @@
|
||||
//! Adapter from the shared, axum-agnostic
|
||||
//! [`cortex_core::error_envelope::OpenAiError`] (#60/#63) to an axum
|
||||
//! response, with `Retry-After`. The `/authz/v1` surface speaks the #63
|
||||
//! envelope so cortex (an OpenAI-compatible proxy) can forward rejections
|
||||
//! verbatim. (The future `/web/v1` surface uses a plain JSON error shape.)
|
||||
|
||||
use axum::http::{HeaderValue, StatusCode, header};
|
||||
use axum::response::{IntoResponse, Json, Response};
|
||||
use cortex_core::error_envelope::OpenAiError;
|
||||
|
||||
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
|
||||
}
|
||||
21
crates/helexa-upstream/src/handlers.rs
Normal file
21
crates/helexa-upstream/src/handlers.rs
Normal file
@@ -0,0 +1,21 @@
|
||||
//! HTTP handlers. B1 ships `/health`; the authz (`/authz/v1`) and web
|
||||
//! (`/web/v1`) surfaces land in later phases.
|
||||
|
||||
use crate::state::AppState;
|
||||
use axum::{Json, Router, extract::State, routing::get};
|
||||
use serde_json::{Value, json};
|
||||
|
||||
pub fn routes() -> Router<AppState> {
|
||||
Router::new()
|
||||
.route("/health", get(health))
|
||||
.route("/", get(health))
|
||||
}
|
||||
|
||||
/// `GET /health` — liveness + a database round-trip (`SELECT 1`).
|
||||
async fn health(State(state): State<AppState>) -> Json<Value> {
|
||||
let db_ok = sqlx::query("SELECT 1").execute(&state.pool).await.is_ok();
|
||||
Json(json!({
|
||||
"status": if db_ok { "ok" } else { "degraded" },
|
||||
"db": if db_ok { "ok" } else { "unreachable" },
|
||||
}))
|
||||
}
|
||||
338
crates/helexa-upstream/src/ledger.rs
Normal file
338
crates/helexa-upstream/src/ledger.rs
Normal file
@@ -0,0 +1,338 @@
|
||||
//! The allocation ledger: reserve → settle/release with the no-overshoot
|
||||
//! guarantee enforced by a row-locked transaction.
|
||||
//!
|
||||
//! Each reserve takes `SELECT … FOR UPDATE` on the account (and key) row, so
|
||||
//! concurrent reserves from many cortexes serialize and `spent + reserved`
|
||||
//! can never exceed the effective cap. The `accounts_no_overshoot` CHECK is
|
||||
//! the DB-level backstop. Settle/release are idempotent (they only act on a
|
||||
//! reservation still in `open`).
|
||||
//!
|
||||
//! Per-key effective cap = `min(resolved key cap, remaining account
|
||||
//! allocation)`. The key cap is resolved from its `limit_kind`:
|
||||
//! `hardcap` → the value verbatim; `percent` → that % of the account's
|
||||
//! `allocation_total`.
|
||||
//!
|
||||
//! Cap-window semantics: this module implements **Balance** (non-resetting)
|
||||
//! caps. Rolling-window key sub-caps (and the `RateLimited` rejection that
|
||||
//! rides them) land with the authz API (B2); today an over-cap is always
|
||||
//! `InsufficientQuota`.
|
||||
|
||||
use sqlx::postgres::PgPool;
|
||||
use uuid::Uuid;
|
||||
|
||||
/// A bearer key resolved to its principal + a budget snapshot.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ResolvedPrincipal {
|
||||
pub account_id: Uuid,
|
||||
pub key_id: Uuid,
|
||||
/// Effective per-key absolute cap (the key sub-cap; the account cap
|
||||
/// still binds at reserve time).
|
||||
pub hard_cap: i64,
|
||||
pub key_spent: i64,
|
||||
pub key_reserved: i64,
|
||||
}
|
||||
|
||||
/// Resolve a key by its `sha256` hash to its principal, or `None` when the
|
||||
/// key is unknown/archived **or its account is deactivated** (the silent
|
||||
/// abuse flag — indistinguishable from an unknown key, by design: no clue).
|
||||
pub async fn resolve_key(
|
||||
pool: &PgPool,
|
||||
key_hash: &[u8],
|
||||
) -> Result<Option<ResolvedPrincipal>, sqlx::Error> {
|
||||
let row = sqlx::query(
|
||||
"SELECT k.id AS key_id, k.account_id, k.limit_kind, k.limit_value, \
|
||||
k.key_spent, k.key_reserved, a.allocation_total \
|
||||
FROM api_keys k JOIN accounts a ON a.id = k.account_id \
|
||||
WHERE k.key_hash = $1 AND k.status = 'active' AND a.status = 'active'",
|
||||
)
|
||||
.bind(key_hash)
|
||||
.fetch_optional(pool)
|
||||
.await?;
|
||||
Ok(row.map(|r| {
|
||||
let total: i64 = sqlx::Row::get(&r, "allocation_total");
|
||||
let limit_kind: String = sqlx::Row::get(&r, "limit_kind");
|
||||
let limit_value: i64 = sqlx::Row::get(&r, "limit_value");
|
||||
ResolvedPrincipal {
|
||||
account_id: sqlx::Row::get(&r, "account_id"),
|
||||
key_id: sqlx::Row::get(&r, "key_id"),
|
||||
hard_cap: resolve_abs_cap(&limit_kind, limit_value, total),
|
||||
key_spent: sqlx::Row::get(&r, "key_spent"),
|
||||
key_reserved: sqlx::Row::get(&r, "key_reserved"),
|
||||
}
|
||||
}))
|
||||
}
|
||||
|
||||
/// Per-key budget snapshot `(hard_cap, spent, reserved)`, or `None` if the
|
||||
/// key/account isn't an active pair.
|
||||
pub async fn snapshot(
|
||||
pool: &PgPool,
|
||||
account_id: Uuid,
|
||||
key_id: Uuid,
|
||||
) -> Result<Option<(i64, i64, i64)>, sqlx::Error> {
|
||||
let row = sqlx::query(
|
||||
"SELECT k.limit_kind, k.limit_value, k.key_spent, k.key_reserved, a.allocation_total \
|
||||
FROM api_keys k JOIN accounts a ON a.id = k.account_id \
|
||||
WHERE k.id = $1 AND k.account_id = $2 AND k.status = 'active' AND a.status = 'active'",
|
||||
)
|
||||
.bind(key_id)
|
||||
.bind(account_id)
|
||||
.fetch_optional(pool)
|
||||
.await?;
|
||||
Ok(row.map(|r| {
|
||||
let total: i64 = sqlx::Row::get(&r, "allocation_total");
|
||||
let limit_kind: String = sqlx::Row::get(&r, "limit_kind");
|
||||
let limit_value: i64 = sqlx::Row::get(&r, "limit_value");
|
||||
let cap = resolve_abs_cap(&limit_kind, limit_value, total);
|
||||
(
|
||||
cap,
|
||||
sqlx::Row::get::<i64, _>(&r, "key_spent"),
|
||||
sqlx::Row::get::<i64, _>(&r, "key_reserved"),
|
||||
)
|
||||
}))
|
||||
}
|
||||
|
||||
/// Release every `open` reservation older than `max_age_secs`, returning
|
||||
/// each one's reserved tokens to its account and key in a single statement.
|
||||
/// The lost-settle self-heal. Returns the number swept.
|
||||
pub async fn sweep_stale(pool: &PgPool, max_age_secs: i64) -> Result<u64, sqlx::Error> {
|
||||
// Data-modifying CTEs: release stale rows, then fold their reserved sums
|
||||
// back into accounts and api_keys. All in one atomic statement.
|
||||
let result = sqlx::query(
|
||||
"WITH stale AS ( \
|
||||
UPDATE reservations SET state = 'released', settled_at = now() \
|
||||
WHERE state = 'open' AND created_at < now() - make_interval(secs => $1) \
|
||||
RETURNING account_id, key_id, reserved \
|
||||
), acct AS ( \
|
||||
UPDATE accounts a SET allocation_reserved = allocation_reserved - s.total \
|
||||
FROM (SELECT account_id, SUM(reserved) AS total FROM stale GROUP BY account_id) s \
|
||||
WHERE a.id = s.account_id \
|
||||
) \
|
||||
UPDATE api_keys k SET key_reserved = key_reserved - s.total \
|
||||
FROM (SELECT key_id, SUM(reserved) AS total FROM stale GROUP BY key_id) s \
|
||||
WHERE k.id = s.key_id",
|
||||
)
|
||||
.bind(max_age_secs as f64)
|
||||
.execute(pool)
|
||||
.await?;
|
||||
Ok(result.rows_affected())
|
||||
}
|
||||
|
||||
/// Resolve a key's per-key cap to an absolute token count.
|
||||
///
|
||||
/// `percent` is `floor(allocation_total * limit_value / 100)`; `hardcap` is
|
||||
/// `limit_value` verbatim. Computed in i128 to avoid overflow, floored at 0.
|
||||
pub fn resolve_abs_cap(limit_kind: &str, limit_value: i64, allocation_total: i64) -> i64 {
|
||||
let cap = match limit_kind {
|
||||
"percent" => (allocation_total as i128 * limit_value as i128) / 100,
|
||||
_ => limit_value as i128, // "hardcap" (and any unknown → treat as absolute)
|
||||
};
|
||||
cap.clamp(0, i64::MAX as i128) as i64
|
||||
}
|
||||
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum LedgerError {
|
||||
#[error("account not found")]
|
||||
AccountNotFound,
|
||||
#[error("api key not found or not active")]
|
||||
KeyNotFound,
|
||||
/// Account balance or a Balance-window key sub-cap is exhausted.
|
||||
#[error("insufficient quota: requested {requested}, available {available}")]
|
||||
InsufficientQuota { requested: i64, available: i64 },
|
||||
#[error(transparent)]
|
||||
Db(#[from] sqlx::Error),
|
||||
}
|
||||
|
||||
/// Reserve `max_tokens` against `account_id`/`key_id`. Returns the
|
||||
/// reservation id (the `BIGSERIAL`, mapped to the cortex `Reservation.id`).
|
||||
pub async fn reserve(
|
||||
pool: &PgPool,
|
||||
account_id: Uuid,
|
||||
key_id: Uuid,
|
||||
max_tokens: i64,
|
||||
) -> Result<i64, LedgerError> {
|
||||
let mut tx = pool.begin().await?;
|
||||
|
||||
// Lock the account row — serializes concurrent reserves on this account.
|
||||
let acct = sqlx::query(
|
||||
"SELECT allocation_total, allocation_spent, allocation_reserved \
|
||||
FROM accounts WHERE id = $1 AND status = 'active' FOR UPDATE",
|
||||
)
|
||||
.bind(account_id)
|
||||
.fetch_optional(&mut *tx)
|
||||
.await?;
|
||||
let Some(acct) = acct else {
|
||||
return Err(LedgerError::AccountNotFound);
|
||||
};
|
||||
let total: i64 = sqlx::Row::get(&acct, "allocation_total");
|
||||
let spent: i64 = sqlx::Row::get(&acct, "allocation_spent");
|
||||
let reserved: i64 = sqlx::Row::get(&acct, "allocation_reserved");
|
||||
let account_avail = total - spent - reserved;
|
||||
|
||||
// Lock the key row and resolve its absolute sub-cap.
|
||||
let key = sqlx::query(
|
||||
"SELECT limit_kind, limit_value, key_spent, key_reserved \
|
||||
FROM api_keys WHERE id = $1 AND account_id = $2 AND status = 'active' FOR UPDATE",
|
||||
)
|
||||
.bind(key_id)
|
||||
.bind(account_id)
|
||||
.fetch_optional(&mut *tx)
|
||||
.await?;
|
||||
let Some(key) = key else {
|
||||
return Err(LedgerError::KeyNotFound);
|
||||
};
|
||||
let limit_kind: String = sqlx::Row::get(&key, "limit_kind");
|
||||
let limit_value: i64 = sqlx::Row::get(&key, "limit_value");
|
||||
let key_spent: i64 = sqlx::Row::get(&key, "key_spent");
|
||||
let key_reserved: i64 = sqlx::Row::get(&key, "key_reserved");
|
||||
let key_cap = resolve_abs_cap(&limit_kind, limit_value, total);
|
||||
let key_avail = key_cap - key_spent - key_reserved;
|
||||
|
||||
let available = account_avail.min(key_avail).max(0);
|
||||
if max_tokens > available {
|
||||
// tx rolls back on drop
|
||||
return Err(LedgerError::InsufficientQuota {
|
||||
requested: max_tokens,
|
||||
available,
|
||||
});
|
||||
}
|
||||
|
||||
let id: i64 = sqlx::Row::get(
|
||||
&sqlx::query(
|
||||
"INSERT INTO reservations (account_id, key_id, reserved, state) \
|
||||
VALUES ($1, $2, $3, 'open') RETURNING id",
|
||||
)
|
||||
.bind(account_id)
|
||||
.bind(key_id)
|
||||
.bind(max_tokens)
|
||||
.fetch_one(&mut *tx)
|
||||
.await?,
|
||||
"id",
|
||||
);
|
||||
sqlx::query("UPDATE accounts SET allocation_reserved = allocation_reserved + $1 WHERE id = $2")
|
||||
.bind(max_tokens)
|
||||
.bind(account_id)
|
||||
.execute(&mut *tx)
|
||||
.await?;
|
||||
sqlx::query("UPDATE api_keys SET key_reserved = key_reserved + $1 WHERE id = $2")
|
||||
.bind(max_tokens)
|
||||
.bind(key_id)
|
||||
.execute(&mut *tx)
|
||||
.await?;
|
||||
|
||||
tx.commit().await?;
|
||||
Ok(id)
|
||||
}
|
||||
|
||||
/// Settle a reservation with the actual tokens used (clamped to
|
||||
/// `[0, reserved]`). Idempotent: a second settle (or settle after release)
|
||||
/// is a no-op.
|
||||
pub async fn settle(
|
||||
pool: &PgPool,
|
||||
reservation_id: i64,
|
||||
actual_tokens: i64,
|
||||
) -> Result<(), LedgerError> {
|
||||
let mut tx = pool.begin().await?;
|
||||
let row = sqlx::query(
|
||||
"UPDATE reservations SET state = 'settled', settled_at = now(), \
|
||||
actual = LEAST(GREATEST($2, 0), reserved) \
|
||||
WHERE id = $1 AND state = 'open' \
|
||||
RETURNING reserved, account_id, key_id, actual",
|
||||
)
|
||||
.bind(reservation_id)
|
||||
.bind(actual_tokens)
|
||||
.fetch_optional(&mut *tx)
|
||||
.await?;
|
||||
let Some(row) = row else {
|
||||
return Ok(()); // already settled/released, or unknown → idempotent no-op
|
||||
};
|
||||
let reserved: i64 = sqlx::Row::get(&row, "reserved");
|
||||
let actual: i64 = sqlx::Row::get(&row, "actual");
|
||||
let account_id: Uuid = sqlx::Row::get(&row, "account_id");
|
||||
let key_id: Uuid = sqlx::Row::get(&row, "key_id");
|
||||
|
||||
sqlx::query(
|
||||
"UPDATE accounts SET allocation_reserved = allocation_reserved - $1, \
|
||||
allocation_spent = allocation_spent + $2 WHERE id = $3",
|
||||
)
|
||||
.bind(reserved)
|
||||
.bind(actual)
|
||||
.bind(account_id)
|
||||
.execute(&mut *tx)
|
||||
.await?;
|
||||
sqlx::query(
|
||||
"UPDATE api_keys SET key_reserved = key_reserved - $1, key_spent = key_spent + $2 WHERE id = $3",
|
||||
)
|
||||
.bind(reserved)
|
||||
.bind(actual)
|
||||
.bind(key_id)
|
||||
.execute(&mut *tx)
|
||||
.await?;
|
||||
|
||||
tx.commit().await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Release a reservation, returning its full reserved amount to the
|
||||
/// allocation. Idempotent.
|
||||
pub async fn release(pool: &PgPool, reservation_id: i64) -> Result<(), LedgerError> {
|
||||
let mut tx = pool.begin().await?;
|
||||
let row = sqlx::query(
|
||||
"UPDATE reservations SET state = 'released', settled_at = now() \
|
||||
WHERE id = $1 AND state = 'open' \
|
||||
RETURNING reserved, account_id, key_id",
|
||||
)
|
||||
.bind(reservation_id)
|
||||
.fetch_optional(&mut *tx)
|
||||
.await?;
|
||||
let Some(row) = row else {
|
||||
return Ok(());
|
||||
};
|
||||
let reserved: i64 = sqlx::Row::get(&row, "reserved");
|
||||
let account_id: Uuid = sqlx::Row::get(&row, "account_id");
|
||||
let key_id: Uuid = sqlx::Row::get(&row, "key_id");
|
||||
|
||||
sqlx::query("UPDATE accounts SET allocation_reserved = allocation_reserved - $1 WHERE id = $2")
|
||||
.bind(reserved)
|
||||
.bind(account_id)
|
||||
.execute(&mut *tx)
|
||||
.await?;
|
||||
sqlx::query("UPDATE api_keys SET key_reserved = key_reserved - $1 WHERE id = $2")
|
||||
.bind(reserved)
|
||||
.bind(key_id)
|
||||
.execute(&mut *tx)
|
||||
.await?;
|
||||
|
||||
tx.commit().await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::resolve_abs_cap;
|
||||
|
||||
#[test]
|
||||
fn hardcap_is_verbatim() {
|
||||
assert_eq!(resolve_abs_cap("hardcap", 50_000, 1_000_000), 50_000);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn percent_is_fraction_of_allocation() {
|
||||
assert_eq!(resolve_abs_cap("percent", 25, 1_000_000), 250_000);
|
||||
assert_eq!(resolve_abs_cap("percent", 100, 1_000_000), 1_000_000);
|
||||
// floor
|
||||
assert_eq!(resolve_abs_cap("percent", 33, 10), 3);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn percent_does_not_overflow_on_large_allocation() {
|
||||
// total * value would overflow i64 if not widened to i128.
|
||||
let cap = resolve_abs_cap("percent", 100, i64::MAX);
|
||||
assert_eq!(cap, i64::MAX);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn negative_or_zero_clamps_to_zero() {
|
||||
assert_eq!(resolve_abs_cap("hardcap", -5, 100), 0);
|
||||
assert_eq!(resolve_abs_cap("percent", 0, 1_000_000), 0);
|
||||
}
|
||||
}
|
||||
89
crates/helexa-upstream/src/lib.rs
Normal file
89
crates/helexa-upstream/src/lib.rs
Normal file
@@ -0,0 +1,89 @@
|
||||
//! helexa-upstream — the mesh-level account/authorization authority (#59).
|
||||
//!
|
||||
//! The clearing house above cortex: it issues accounts and API keys, holds
|
||||
//! the real token-allocation ledger, authorizes inference in real time
|
||||
//! (reserve → settle, fail-closed), and tracks served usage for operator
|
||||
//! reconciliation. cortex's `UpstreamEntitlementProvider` (#57) is a client
|
||||
//! of the `/authz/v1` surface; the helexa.ai frontend is a client of the
|
||||
//! `/web/v1` surface.
|
||||
//!
|
||||
//! Landed so far: B1 — schema + reserve→settle [`ledger`] (no-overshoot) +
|
||||
//! `/health`. B2 — the `/authz/v1` [`authz`] surface (resolve/reserve/
|
||||
//! settle/release/snapshot) with shared-bearer client auth and a
|
||||
//! stale-reservation sweeper.
|
||||
|
||||
pub mod authz;
|
||||
pub mod config;
|
||||
pub mod crypto;
|
||||
pub mod db;
|
||||
pub mod email;
|
||||
pub mod error;
|
||||
pub mod handlers;
|
||||
pub mod ledger;
|
||||
pub mod state;
|
||||
pub mod topup;
|
||||
pub mod web;
|
||||
|
||||
use anyhow::Result;
|
||||
use config::UpstreamConfig;
|
||||
use email::EmailSender;
|
||||
use state::AppState;
|
||||
use std::time::Duration;
|
||||
use tower_http::cors::CorsLayer;
|
||||
use tower_http::trace::TraceLayer;
|
||||
|
||||
/// Build the axum application.
|
||||
pub fn build_app(state: AppState) -> axum::Router {
|
||||
axum::Router::new()
|
||||
.merge(handlers::routes())
|
||||
.merge(authz::router(&state))
|
||||
.merge(web::router(&state))
|
||||
// The /web/v1 surface is called cross-origin by the browser SPA in
|
||||
// dev; same-origin behind nginx in prod. Permissive is fine — these
|
||||
// endpoints authenticate via bearer/JWT, not cookies.
|
||||
.layer(CorsLayer::permissive())
|
||||
.layer(TraceLayer::new_for_http())
|
||||
.with_state(state)
|
||||
}
|
||||
|
||||
/// Start the service: connect Postgres, run migrations, spawn the
|
||||
/// reservation sweeper, bind the listener.
|
||||
pub async fn run(config: UpstreamConfig) -> Result<()> {
|
||||
let pool = db::connect_and_migrate(&config.db.url, config.db.max_connections).await?;
|
||||
let email = EmailSender::from_config(&config.email)?;
|
||||
let listen = config.server.listen.clone();
|
||||
let state = AppState::new(pool, config, email);
|
||||
|
||||
if state.config.client_auth.tokens.is_empty() {
|
||||
tracing::warn!(
|
||||
"no [client_auth] tokens configured — the /authz/v1 surface is OPEN (dev only)"
|
||||
);
|
||||
}
|
||||
|
||||
// Stale-reservation sweeper: releases open reservations whose
|
||||
// settle/release from cortex was lost, self-healing allocation_reserved.
|
||||
spawn_sweeper(&state);
|
||||
|
||||
let addr = listen.parse::<std::net::SocketAddr>()?;
|
||||
tracing::info!("helexa-upstream listening on {addr}");
|
||||
|
||||
let listener = tokio::net::TcpListener::bind(addr).await?;
|
||||
axum::serve(listener, build_app(state)).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn spawn_sweeper(state: &AppState) {
|
||||
let pool = state.pool.clone();
|
||||
let ttl = state.config.authz.reservation_ttl_secs as i64;
|
||||
let interval = Duration::from_secs(state.config.authz.sweep_interval_secs);
|
||||
tokio::spawn(async move {
|
||||
loop {
|
||||
tokio::time::sleep(interval).await;
|
||||
match ledger::sweep_stale(&pool, ttl).await {
|
||||
Ok(n) if n > 0 => tracing::info!(swept = n, "released stale reservations"),
|
||||
Ok(_) => {}
|
||||
Err(e) => tracing::warn!(error = %e, "reservation sweep failed"),
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
81
crates/helexa-upstream/src/main.rs
Normal file
81
crates/helexa-upstream/src/main.rs
Normal file
@@ -0,0 +1,81 @@
|
||||
use anyhow::Result;
|
||||
use clap::{Parser, Subcommand};
|
||||
use helexa_upstream::config::UpstreamConfig;
|
||||
use tracing_subscriber::EnvFilter;
|
||||
|
||||
#[derive(Parser)]
|
||||
#[command(name = "helexa-upstream")]
|
||||
#[command(about = "Mesh-level account & authorization authority for helexa")]
|
||||
#[command(version)]
|
||||
struct Cli {
|
||||
#[command(subcommand)]
|
||||
command: Commands,
|
||||
}
|
||||
|
||||
#[derive(Subcommand)]
|
||||
enum Commands {
|
||||
/// Start the upstream server.
|
||||
Serve {
|
||||
/// Path to the config file.
|
||||
#[arg(short, long, default_value = "helexa-upstream.toml")]
|
||||
config: String,
|
||||
},
|
||||
/// Mint single-use top-up codes and print them (one per line). The raw
|
||||
/// codes are shown only here — only their hash is stored. (The future
|
||||
/// faucet bot calls the same path.)
|
||||
Mint {
|
||||
#[arg(short, long, default_value = "helexa-upstream.toml")]
|
||||
config: String,
|
||||
/// Tokens each code grants.
|
||||
#[arg(long)]
|
||||
value: i64,
|
||||
/// How many codes to mint.
|
||||
#[arg(long, default_value_t = 1)]
|
||||
count: u32,
|
||||
/// Optional human label (e.g. "small", "beta-launch").
|
||||
#[arg(long)]
|
||||
denomination: Option<String>,
|
||||
},
|
||||
}
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() -> Result<()> {
|
||||
tracing_subscriber::fmt()
|
||||
.with_env_filter(
|
||||
EnvFilter::try_from_default_env()
|
||||
.unwrap_or_else(|_| EnvFilter::new("info,helexa_upstream=debug")),
|
||||
)
|
||||
.init();
|
||||
|
||||
let cli = Cli::parse();
|
||||
|
||||
match cli.command {
|
||||
Commands::Serve { config } => {
|
||||
let cfg = UpstreamConfig::load(&config)
|
||||
.map_err(|e| anyhow::anyhow!("failed to load config from '{config}': {e}"))?;
|
||||
tracing::info!(listen = %cfg.server.listen, "starting helexa-upstream");
|
||||
helexa_upstream::run(cfg).await?;
|
||||
}
|
||||
Commands::Mint {
|
||||
config,
|
||||
value,
|
||||
count,
|
||||
denomination,
|
||||
} => {
|
||||
let cfg = UpstreamConfig::load(&config)
|
||||
.map_err(|e| anyhow::anyhow!("failed to load config from '{config}': {e}"))?;
|
||||
let pool =
|
||||
helexa_upstream::db::connect_and_migrate(&cfg.db.url, cfg.db.max_connections)
|
||||
.await?;
|
||||
let codes =
|
||||
helexa_upstream::topup::mint(&pool, value, count, denomination.as_deref()).await?;
|
||||
// Raw codes to stdout (one per line) for the operator to distribute;
|
||||
// logs/diagnostics go to stderr via tracing.
|
||||
for code in codes {
|
||||
println!("{code}");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
23
crates/helexa-upstream/src/state.rs
Normal file
23
crates/helexa-upstream/src/state.rs
Normal file
@@ -0,0 +1,23 @@
|
||||
//! Shared application state.
|
||||
|
||||
use crate::config::UpstreamConfig;
|
||||
use crate::email::EmailSender;
|
||||
use sqlx::postgres::PgPool;
|
||||
use std::sync::Arc;
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct AppState {
|
||||
pub pool: PgPool,
|
||||
pub config: Arc<UpstreamConfig>,
|
||||
pub email: EmailSender,
|
||||
}
|
||||
|
||||
impl AppState {
|
||||
pub fn new(pool: PgPool, config: UpstreamConfig, email: EmailSender) -> Self {
|
||||
Self {
|
||||
pool,
|
||||
config: Arc::new(config),
|
||||
email,
|
||||
}
|
||||
}
|
||||
}
|
||||
82
crates/helexa-upstream/src/topup.rs
Normal file
82
crates/helexa-upstream/src/topup.rs
Normal file
@@ -0,0 +1,82 @@
|
||||
//! Single-use top-up codes (#B5) — the second half of the hybrid allocation
|
||||
//! model. Each code grants `value` tokens to the account that redeems it,
|
||||
//! raising `accounts.allocation_total`. Minting codes is operator/CLI side
|
||||
//! (the future faucet bot calls the same `mint` path); redemption is a
|
||||
//! `/web/v1` action.
|
||||
//!
|
||||
//! Security: only `sha256(code)` is stored. Redemption is **timing-safe and
|
||||
//! single-use** — a conditional `UPDATE … WHERE redeemed_by IS NULL` does
|
||||
//! the claim atomically (concurrent double-redeem → exactly one winner), and
|
||||
//! a not-found code and an already-redeemed code return the **same** generic
|
||||
//! failure with the same code path (no oracle for "valid but spent").
|
||||
|
||||
use crate::crypto::{random_token, sha256};
|
||||
use sqlx::Row;
|
||||
use sqlx::postgres::PgPool;
|
||||
use uuid::Uuid;
|
||||
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum TopUpError {
|
||||
/// Code unknown OR already redeemed — deliberately indistinguishable.
|
||||
#[error("invalid or already-redeemed code")]
|
||||
Invalid,
|
||||
#[error(transparent)]
|
||||
Db(#[from] sqlx::Error),
|
||||
}
|
||||
|
||||
/// Redeem `raw_code` for `account_id`, raising the account's
|
||||
/// `allocation_total` by the code's value. Returns the new total.
|
||||
pub async fn redeem(pool: &PgPool, account_id: Uuid, raw_code: &str) -> Result<i64, TopUpError> {
|
||||
let mut tx = pool.begin().await?;
|
||||
// Atomic single-use claim. `redeemed_by IS NULL` is the guarantee: under
|
||||
// concurrent redemption exactly one UPDATE touches the row.
|
||||
let claimed = sqlx::query(
|
||||
"UPDATE top_up_codes SET redeemed_by = $1, redeemed_at = now() \
|
||||
WHERE code_hash = $2 AND redeemed_by IS NULL RETURNING value",
|
||||
)
|
||||
.bind(account_id)
|
||||
.bind(sha256(raw_code))
|
||||
.fetch_optional(&mut *tx)
|
||||
.await?;
|
||||
let Some(row) = claimed else {
|
||||
// Not found or already redeemed — same path, same error.
|
||||
return Err(TopUpError::Invalid);
|
||||
};
|
||||
let value: i64 = row.get("value");
|
||||
let new_total: i64 = sqlx::query(
|
||||
"UPDATE accounts SET allocation_total = allocation_total + $1 WHERE id = $2 \
|
||||
RETURNING allocation_total",
|
||||
)
|
||||
.bind(value)
|
||||
.bind(account_id)
|
||||
.fetch_one(&mut *tx)
|
||||
.await?
|
||||
.get("allocation_total");
|
||||
tx.commit().await?;
|
||||
Ok(new_total)
|
||||
}
|
||||
|
||||
/// Mint `count` codes each worth `value` tokens, optionally tagged with a
|
||||
/// `denomination` label. Returns the raw codes (shown once — only their
|
||||
/// hash is stored). The CLI prints these; the future faucet bot calls this.
|
||||
pub async fn mint(
|
||||
pool: &PgPool,
|
||||
value: i64,
|
||||
count: u32,
|
||||
denomination: Option<&str>,
|
||||
) -> Result<Vec<String>, sqlx::Error> {
|
||||
let mut codes = Vec::with_capacity(count as usize);
|
||||
for _ in 0..count {
|
||||
let raw = format!("helexa-topup-{}", random_token());
|
||||
sqlx::query(
|
||||
"INSERT INTO top_up_codes (code_hash, value, denomination) VALUES ($1, $2, $3)",
|
||||
)
|
||||
.bind(sha256(&raw))
|
||||
.bind(value)
|
||||
.bind(denomination)
|
||||
.execute(pool)
|
||||
.await?;
|
||||
codes.push(raw);
|
||||
}
|
||||
Ok(codes)
|
||||
}
|
||||
594
crates/helexa-upstream/src/web.rs
Normal file
594
crates/helexa-upstream/src/web.rs
Normal file
@@ -0,0 +1,594 @@
|
||||
//! `/web/v1` — the human-facing account API the helexa.ai frontend (#F4)
|
||||
//! consumes: email+password auth (register / verify / login / reset),
|
||||
//! API-key CRUD with per-key limits, and the account balance. Web sessions
|
||||
//! are JWTs, **distinct** from inference API keys.
|
||||
//!
|
||||
//! Errors use a plain JSON shape `{ "error": { "message", "code" } }` (web
|
||||
//! clients, not OpenAI clients — the #63 envelope is the authz surface).
|
||||
//!
|
||||
//! Silent fingerprint abuse (no clue to the abuser): registration captures
|
||||
//! the browser fingerprint and always succeeds; when ≥ threshold accounts
|
||||
//! share one fingerprint, all are silently `deactivated` (keys then resolve
|
||||
//! as ordinary `401`s at the authz surface — never a "banned" signal).
|
||||
|
||||
use crate::crypto::{generate_api_key, hash_password, random_token, sha256, verify_password};
|
||||
use crate::state::AppState;
|
||||
use axum::extract::{Path, Request, State};
|
||||
use axum::http::{StatusCode, header};
|
||||
use axum::middleware::Next;
|
||||
use axum::response::{IntoResponse, Json, Response};
|
||||
use axum::routing::{get, post};
|
||||
use axum::{Extension, Router};
|
||||
use chrono::{DateTime, Duration, Utc};
|
||||
use jsonwebtoken::{DecodingKey, EncodingKey, Header, Validation, decode, encode};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::json;
|
||||
use sqlx::Row;
|
||||
use uuid::Uuid;
|
||||
|
||||
pub fn router(state: &AppState) -> Router<AppState> {
|
||||
let protected = Router::new()
|
||||
.route("/web/v1/account", get(account))
|
||||
.route("/web/v1/keys", get(list_keys).post(create_key))
|
||||
.route("/web/v1/keys/{id}/archive", post(archive_key))
|
||||
.route(
|
||||
"/web/v1/keys/{id}/limit",
|
||||
axum::routing::patch(update_key_limit),
|
||||
)
|
||||
.route("/web/v1/redeem", post(redeem))
|
||||
.layer(axum::middleware::from_fn_with_state(
|
||||
state.clone(),
|
||||
require_session,
|
||||
));
|
||||
|
||||
Router::new()
|
||||
.route("/web/v1/register", post(register))
|
||||
.route("/web/v1/verify", post(verify))
|
||||
.route("/web/v1/login", post(login))
|
||||
.route("/web/v1/password-reset/request", post(reset_request))
|
||||
.route("/web/v1/password-reset/confirm", post(reset_confirm))
|
||||
.merge(protected)
|
||||
}
|
||||
|
||||
// ── errors ──────────────────────────────────────────────────────────
|
||||
|
||||
enum WebError {
|
||||
BadRequest(&'static str),
|
||||
Unauthorized,
|
||||
Internal,
|
||||
}
|
||||
|
||||
impl IntoResponse for WebError {
|
||||
fn into_response(self) -> Response {
|
||||
let (status, code, message) = match self {
|
||||
WebError::BadRequest(m) => (StatusCode::BAD_REQUEST, "bad_request", m),
|
||||
WebError::Unauthorized => (StatusCode::UNAUTHORIZED, "unauthorized", "unauthorized"),
|
||||
WebError::Internal => (
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
"internal_error",
|
||||
"internal error",
|
||||
),
|
||||
};
|
||||
(
|
||||
status,
|
||||
Json(json!({"error": {"message": message, "code": code}})),
|
||||
)
|
||||
.into_response()
|
||||
}
|
||||
}
|
||||
|
||||
impl From<sqlx::Error> for WebError {
|
||||
fn from(e: sqlx::Error) -> Self {
|
||||
tracing::error!(error = %e, "web db error");
|
||||
WebError::Internal
|
||||
}
|
||||
}
|
||||
|
||||
type WebResult<T> = Result<T, WebError>;
|
||||
|
||||
// ── sessions (JWT) ──────────────────────────────────────────────────
|
||||
|
||||
#[derive(Serialize, Deserialize)]
|
||||
struct Claims {
|
||||
sub: String, // user id
|
||||
exp: usize,
|
||||
}
|
||||
|
||||
fn mint_session(state: &AppState, user_id: Uuid) -> WebResult<String> {
|
||||
let exp = (Utc::now() + Duration::seconds(state.config.auth.session_ttl_secs as i64))
|
||||
.timestamp() as usize;
|
||||
let claims = Claims {
|
||||
sub: user_id.to_string(),
|
||||
exp,
|
||||
};
|
||||
encode(
|
||||
&Header::default(),
|
||||
&claims,
|
||||
&EncodingKey::from_secret(state.config.auth.jwt_secret.as_bytes()),
|
||||
)
|
||||
.map_err(|_| WebError::Internal)
|
||||
}
|
||||
|
||||
/// Authenticated user id, injected by [`require_session`].
|
||||
#[derive(Clone)]
|
||||
struct AuthUser(Uuid);
|
||||
|
||||
async fn require_session(State(state): State<AppState>, mut req: Request, next: Next) -> Response {
|
||||
let token = req
|
||||
.headers()
|
||||
.get(header::AUTHORIZATION)
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.and_then(|v| v.strip_prefix("Bearer "))
|
||||
.map(str::trim);
|
||||
let Some(token) = token else {
|
||||
return WebError::Unauthorized.into_response();
|
||||
};
|
||||
let decoded = decode::<Claims>(
|
||||
token,
|
||||
&DecodingKey::from_secret(state.config.auth.jwt_secret.as_bytes()),
|
||||
&Validation::default(),
|
||||
);
|
||||
match decoded
|
||||
.ok()
|
||||
.and_then(|d| Uuid::parse_str(&d.claims.sub).ok())
|
||||
{
|
||||
Some(uid) => {
|
||||
req.extensions_mut().insert(AuthUser(uid));
|
||||
next.run(req).await
|
||||
}
|
||||
None => WebError::Unauthorized.into_response(),
|
||||
}
|
||||
}
|
||||
|
||||
/// The caller's single account id.
|
||||
async fn account_id_for(state: &AppState, user_id: Uuid) -> WebResult<Uuid> {
|
||||
let row = sqlx::query("SELECT id FROM accounts WHERE owner_user_id = $1")
|
||||
.bind(user_id)
|
||||
.fetch_optional(&state.pool)
|
||||
.await?;
|
||||
row.map(|r| r.get::<Uuid, _>("id"))
|
||||
.ok_or(WebError::Internal)
|
||||
}
|
||||
|
||||
// ── auth lifecycle ──────────────────────────────────────────────────
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct RegisterReq {
|
||||
email: String,
|
||||
password: String,
|
||||
#[serde(default)]
|
||||
fingerprint: Option<String>,
|
||||
}
|
||||
|
||||
/// `POST /web/v1/register` — always returns `202`, regardless of whether the
|
||||
/// email was new, already taken, or fingerprint-flagged (no enumeration, no
|
||||
/// abuse clue).
|
||||
async fn register(State(state): State<AppState>, Json(req): Json<RegisterReq>) -> Response {
|
||||
match register_inner(&state, req).await {
|
||||
Ok(()) | Err(WebError::BadRequest(_)) => {}
|
||||
Err(e) => return e.into_response(),
|
||||
}
|
||||
// Generic 202 whatever happened above (except hard server errors).
|
||||
StatusCode::ACCEPTED.into_response()
|
||||
}
|
||||
|
||||
async fn register_inner(state: &AppState, req: RegisterReq) -> WebResult<()> {
|
||||
if !req.email.contains('@') {
|
||||
return Err(WebError::BadRequest("invalid email"));
|
||||
}
|
||||
if req.password.len() < 8 {
|
||||
return Err(WebError::BadRequest("password too short (min 8)"));
|
||||
}
|
||||
let phc = hash_password(&req.password).map_err(|_| WebError::Internal)?;
|
||||
|
||||
// Insert the user; a duplicate email silently no-ops (no enumeration).
|
||||
let user_id: Option<Uuid> = sqlx::query(
|
||||
"INSERT INTO users (email, password_hash, registration_fingerprint) \
|
||||
VALUES ($1, $2, $3) ON CONFLICT (email) DO NOTHING RETURNING id",
|
||||
)
|
||||
.bind(&req.email)
|
||||
.bind(&phc)
|
||||
.bind(&req.fingerprint)
|
||||
.fetch_optional(&state.pool)
|
||||
.await?
|
||||
.map(|r| r.get("id"));
|
||||
|
||||
let Some(user_id) = user_id else {
|
||||
return Ok(()); // email already registered — say nothing
|
||||
};
|
||||
|
||||
// Account with the flat free grant.
|
||||
sqlx::query("INSERT INTO accounts (owner_user_id, allocation_total) VALUES ($1, $2)")
|
||||
.bind(user_id)
|
||||
.bind(state.config.grant.free_token_grant)
|
||||
.execute(&state.pool)
|
||||
.await?;
|
||||
|
||||
// Silent fingerprint abuse handling.
|
||||
if let Some(fp) = req.fingerprint.as_deref().filter(|f| !f.is_empty()) {
|
||||
apply_fingerprint_policy(state, fp).await?;
|
||||
}
|
||||
|
||||
// Email verification link.
|
||||
let token = random_token();
|
||||
let expires: DateTime<Utc> =
|
||||
Utc::now() + Duration::seconds(state.config.auth.email_token_ttl_secs as i64);
|
||||
sqlx::query(
|
||||
"INSERT INTO email_tokens (token_hash, user_id, kind, expires_at) \
|
||||
VALUES ($1, $2, 'verify', $3)",
|
||||
)
|
||||
.bind(sha256(&token))
|
||||
.bind(user_id)
|
||||
.bind(expires)
|
||||
.execute(&state.pool)
|
||||
.await?;
|
||||
|
||||
let link = format!("{}/verify?token={token}", state.config.auth.app_base_url);
|
||||
let _ = state
|
||||
.email
|
||||
.send(
|
||||
&req.email,
|
||||
"Verify your helexa account",
|
||||
&format!("Welcome to helexa. Verify your email:\n\n{link}\n"),
|
||||
)
|
||||
.await;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Count accounts sharing `fp`; flag them, and silently deactivate all once
|
||||
/// the count reaches the configured threshold. No response difference — the
|
||||
/// abuser gets no signal.
|
||||
async fn apply_fingerprint_policy(state: &AppState, fp: &str) -> WebResult<()> {
|
||||
let count: i64 =
|
||||
sqlx::query_scalar("SELECT count(*) FROM users WHERE registration_fingerprint = $1")
|
||||
.bind(fp)
|
||||
.fetch_one(&state.pool)
|
||||
.await?;
|
||||
if count > 1 {
|
||||
sqlx::query(
|
||||
"UPDATE accounts SET fingerprint_flagged = true \
|
||||
WHERE owner_user_id IN (SELECT id FROM users WHERE registration_fingerprint = $1)",
|
||||
)
|
||||
.bind(fp)
|
||||
.execute(&state.pool)
|
||||
.await?;
|
||||
}
|
||||
if count >= state.config.abuse.fingerprint_account_threshold {
|
||||
let res = sqlx::query(
|
||||
"UPDATE accounts SET status = 'deactivated' \
|
||||
WHERE owner_user_id IN (SELECT id FROM users WHERE registration_fingerprint = $1)",
|
||||
)
|
||||
.bind(fp)
|
||||
.execute(&state.pool)
|
||||
.await?;
|
||||
tracing::warn!(
|
||||
fingerprint = fp,
|
||||
accounts = res.rows_affected(),
|
||||
"silently deactivated fingerprint-abusing accounts"
|
||||
);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct TokenReq {
|
||||
token: String,
|
||||
}
|
||||
|
||||
/// `POST /web/v1/verify` — consume a verification token, mark verified.
|
||||
async fn verify(State(state): State<AppState>, Json(req): Json<TokenReq>) -> WebResult<Response> {
|
||||
let row = sqlx::query(
|
||||
"UPDATE email_tokens SET consumed_at = now() \
|
||||
WHERE token_hash = $1 AND kind = 'verify' AND consumed_at IS NULL AND expires_at > now() \
|
||||
RETURNING user_id",
|
||||
)
|
||||
.bind(sha256(&req.token))
|
||||
.fetch_optional(&state.pool)
|
||||
.await?;
|
||||
let Some(row) = row else {
|
||||
return Err(WebError::BadRequest("invalid or expired token"));
|
||||
};
|
||||
let user_id: Uuid = row.get("user_id");
|
||||
sqlx::query("UPDATE users SET email_verified = true WHERE id = $1")
|
||||
.bind(user_id)
|
||||
.execute(&state.pool)
|
||||
.await?;
|
||||
Ok(StatusCode::OK.into_response())
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct LoginReq {
|
||||
email: String,
|
||||
password: String,
|
||||
}
|
||||
|
||||
/// `POST /web/v1/login` — verify password + email-verified → session JWT.
|
||||
async fn login(State(state): State<AppState>, Json(req): Json<LoginReq>) -> WebResult<Response> {
|
||||
let row = sqlx::query("SELECT id, password_hash, email_verified FROM users WHERE email = $1")
|
||||
.bind(&req.email)
|
||||
.fetch_optional(&state.pool)
|
||||
.await?;
|
||||
// Generic 401 for every failure mode (no enumeration).
|
||||
let Some(row) = row else {
|
||||
return Err(WebError::Unauthorized);
|
||||
};
|
||||
let phc: String = row.get("password_hash");
|
||||
let verified: bool = row.get("email_verified");
|
||||
if !verify_password(&req.password, &phc) || !verified {
|
||||
return Err(WebError::Unauthorized);
|
||||
}
|
||||
let user_id: Uuid = row.get("id");
|
||||
let token = mint_session(&state, user_id)?;
|
||||
Ok(Json(json!({
|
||||
"token": token,
|
||||
"expires_in": state.config.auth.session_ttl_secs,
|
||||
}))
|
||||
.into_response())
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct EmailReq {
|
||||
email: String,
|
||||
}
|
||||
|
||||
/// `POST /web/v1/password-reset/request` — always `202` (no enumeration);
|
||||
/// mints + emails a reset token only if the account exists.
|
||||
async fn reset_request(State(state): State<AppState>, Json(req): Json<EmailReq>) -> Response {
|
||||
// The inner only ever yields `Internal` (DB failure); a missing email is
|
||||
// Ok(()) so there's no enumeration. Surface 500 on a real error, else 202.
|
||||
match reset_request_inner(&state, &req.email).await {
|
||||
Ok(()) => StatusCode::ACCEPTED.into_response(),
|
||||
Err(e) => e.into_response(),
|
||||
}
|
||||
}
|
||||
|
||||
async fn reset_request_inner(state: &AppState, email: &str) -> WebResult<()> {
|
||||
let row = sqlx::query("SELECT id FROM users WHERE email = $1")
|
||||
.bind(email)
|
||||
.fetch_optional(&state.pool)
|
||||
.await?;
|
||||
let Some(row) = row else { return Ok(()) };
|
||||
let user_id: Uuid = row.get("id");
|
||||
let token = random_token();
|
||||
let expires: DateTime<Utc> =
|
||||
Utc::now() + Duration::seconds(state.config.auth.email_token_ttl_secs as i64);
|
||||
sqlx::query(
|
||||
"INSERT INTO email_tokens (token_hash, user_id, kind, expires_at) \
|
||||
VALUES ($1, $2, 'reset', $3)",
|
||||
)
|
||||
.bind(sha256(&token))
|
||||
.bind(user_id)
|
||||
.bind(expires)
|
||||
.execute(&state.pool)
|
||||
.await?;
|
||||
let link = format!("{}/reset?token={token}", state.config.auth.app_base_url);
|
||||
let _ = state
|
||||
.email
|
||||
.send(
|
||||
email,
|
||||
"Reset your helexa password",
|
||||
&format!("Reset your password:\n\n{link}\n"),
|
||||
)
|
||||
.await;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct ResetConfirmReq {
|
||||
token: String,
|
||||
new_password: String,
|
||||
}
|
||||
|
||||
/// `POST /web/v1/password-reset/confirm` — consume reset token, rotate hash.
|
||||
async fn reset_confirm(
|
||||
State(state): State<AppState>,
|
||||
Json(req): Json<ResetConfirmReq>,
|
||||
) -> WebResult<Response> {
|
||||
if req.new_password.len() < 8 {
|
||||
return Err(WebError::BadRequest("password too short (min 8)"));
|
||||
}
|
||||
let row = sqlx::query(
|
||||
"UPDATE email_tokens SET consumed_at = now() \
|
||||
WHERE token_hash = $1 AND kind = 'reset' AND consumed_at IS NULL AND expires_at > now() \
|
||||
RETURNING user_id",
|
||||
)
|
||||
.bind(sha256(&req.token))
|
||||
.fetch_optional(&state.pool)
|
||||
.await?;
|
||||
let Some(row) = row else {
|
||||
return Err(WebError::BadRequest("invalid or expired token"));
|
||||
};
|
||||
let user_id: Uuid = row.get("user_id");
|
||||
let phc = hash_password(&req.new_password).map_err(|_| WebError::Internal)?;
|
||||
sqlx::query("UPDATE users SET password_hash = $1 WHERE id = $2")
|
||||
.bind(phc)
|
||||
.bind(user_id)
|
||||
.execute(&state.pool)
|
||||
.await?;
|
||||
Ok(StatusCode::OK.into_response())
|
||||
}
|
||||
|
||||
// ── account + keys (protected) ──────────────────────────────────────
|
||||
|
||||
async fn account(
|
||||
State(state): State<AppState>,
|
||||
Extension(user): Extension<AuthUser>,
|
||||
) -> WebResult<Response> {
|
||||
let acct = account_id_for(&state, user.0).await?;
|
||||
let row = sqlx::query(
|
||||
"SELECT allocation_total, allocation_spent, allocation_reserved FROM accounts WHERE id = $1",
|
||||
)
|
||||
.bind(acct)
|
||||
.fetch_one(&state.pool)
|
||||
.await?;
|
||||
Ok(Json(json!({
|
||||
"account_id": acct.to_string(),
|
||||
"allocation_total": row.get::<i64, _>("allocation_total"),
|
||||
"allocation_spent": row.get::<i64, _>("allocation_spent"),
|
||||
"allocation_reserved": row.get::<i64, _>("allocation_reserved"),
|
||||
}))
|
||||
.into_response())
|
||||
}
|
||||
|
||||
async fn list_keys(
|
||||
State(state): State<AppState>,
|
||||
Extension(user): Extension<AuthUser>,
|
||||
) -> WebResult<Response> {
|
||||
let acct = account_id_for(&state, user.0).await?;
|
||||
let rows = sqlx::query(
|
||||
"SELECT id, key_prefix, label, status, limit_kind, limit_value, key_spent, key_reserved, \
|
||||
created_at \
|
||||
FROM api_keys WHERE account_id = $1 ORDER BY created_at DESC",
|
||||
)
|
||||
.bind(acct)
|
||||
.fetch_all(&state.pool)
|
||||
.await?;
|
||||
let keys: Vec<_> = rows
|
||||
.iter()
|
||||
.map(|r| {
|
||||
json!({
|
||||
"id": r.get::<Uuid, _>("id").to_string(),
|
||||
"prefix": r.get::<String, _>("key_prefix"),
|
||||
"label": r.get::<String, _>("label"),
|
||||
"status": r.get::<String, _>("status"),
|
||||
"limit_kind": r.get::<String, _>("limit_kind"),
|
||||
"limit_value": r.get::<i64, _>("limit_value"),
|
||||
"spent": r.get::<i64, _>("key_spent"),
|
||||
"reserved": r.get::<i64, _>("key_reserved"),
|
||||
"created_at": r.get::<DateTime<Utc>, _>("created_at").to_rfc3339(),
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
Ok(Json(json!({ "keys": keys })).into_response())
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct CreateKeyReq {
|
||||
#[serde(default)]
|
||||
label: String,
|
||||
/// "percent" | "hardcap" (default percent=100 → full allocation).
|
||||
#[serde(default)]
|
||||
limit_kind: Option<String>,
|
||||
#[serde(default)]
|
||||
limit_value: Option<i64>,
|
||||
}
|
||||
|
||||
async fn create_key(
|
||||
State(state): State<AppState>,
|
||||
Extension(user): Extension<AuthUser>,
|
||||
Json(req): Json<CreateKeyReq>,
|
||||
) -> WebResult<Response> {
|
||||
let acct = account_id_for(&state, user.0).await?;
|
||||
let limit_kind = match req.limit_kind.as_deref() {
|
||||
Some("hardcap") => "hardcap",
|
||||
_ => "percent",
|
||||
};
|
||||
let limit_value = req.limit_value.unwrap_or(100).max(0);
|
||||
let (raw, prefix) = generate_api_key();
|
||||
let id: Uuid = sqlx::query(
|
||||
"INSERT INTO api_keys (account_id, key_hash, key_prefix, label, limit_kind, limit_value) \
|
||||
VALUES ($1, $2, $3, $4, $5, $6) RETURNING id",
|
||||
)
|
||||
.bind(acct)
|
||||
.bind(sha256(&raw))
|
||||
.bind(&prefix)
|
||||
.bind(&req.label)
|
||||
.bind(limit_kind)
|
||||
.bind(limit_value)
|
||||
.fetch_one(&state.pool)
|
||||
.await?
|
||||
.get("id");
|
||||
// The raw key is shown exactly once.
|
||||
Ok((
|
||||
StatusCode::CREATED,
|
||||
Json(json!({
|
||||
"id": id.to_string(),
|
||||
"key": raw,
|
||||
"prefix": prefix,
|
||||
"limit_kind": limit_kind,
|
||||
"limit_value": limit_value,
|
||||
})),
|
||||
)
|
||||
.into_response())
|
||||
}
|
||||
|
||||
async fn archive_key(
|
||||
State(state): State<AppState>,
|
||||
Extension(user): Extension<AuthUser>,
|
||||
Path(id): Path<Uuid>,
|
||||
) -> WebResult<Response> {
|
||||
let acct = account_id_for(&state, user.0).await?;
|
||||
let res = sqlx::query(
|
||||
"UPDATE api_keys SET status = 'archived' WHERE id = $1 AND account_id = $2 AND status = 'active'",
|
||||
)
|
||||
.bind(id)
|
||||
.bind(acct)
|
||||
.execute(&state.pool)
|
||||
.await?;
|
||||
if res.rows_affected() == 0 {
|
||||
return Err(WebError::BadRequest("no such active key"));
|
||||
}
|
||||
Ok(StatusCode::NO_CONTENT.into_response())
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct UpdateLimitReq {
|
||||
limit_kind: String,
|
||||
limit_value: i64,
|
||||
}
|
||||
|
||||
async fn update_key_limit(
|
||||
State(state): State<AppState>,
|
||||
Extension(user): Extension<AuthUser>,
|
||||
Path(id): Path<Uuid>,
|
||||
Json(req): Json<UpdateLimitReq>,
|
||||
) -> WebResult<Response> {
|
||||
if req.limit_kind != "percent" && req.limit_kind != "hardcap" {
|
||||
return Err(WebError::BadRequest(
|
||||
"limit_kind must be percent or hardcap",
|
||||
));
|
||||
}
|
||||
if req.limit_value < 0 {
|
||||
return Err(WebError::BadRequest("limit_value must be >= 0"));
|
||||
}
|
||||
let acct = account_id_for(&state, user.0).await?;
|
||||
let res = sqlx::query(
|
||||
"UPDATE api_keys SET limit_kind = $1, limit_value = $2 WHERE id = $3 AND account_id = $4",
|
||||
)
|
||||
.bind(&req.limit_kind)
|
||||
.bind(req.limit_value)
|
||||
.bind(id)
|
||||
.bind(acct)
|
||||
.execute(&state.pool)
|
||||
.await?;
|
||||
if res.rows_affected() == 0 {
|
||||
return Err(WebError::BadRequest("no such key"));
|
||||
}
|
||||
Ok(StatusCode::NO_CONTENT.into_response())
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct RedeemReq {
|
||||
code: String,
|
||||
}
|
||||
|
||||
/// `POST /web/v1/redeem` — redeem a single-use top-up code, raising the
|
||||
/// account's allocation. Returns the new total. Generic 400 for an invalid
|
||||
/// or already-redeemed code (no oracle).
|
||||
async fn redeem(
|
||||
State(state): State<AppState>,
|
||||
Extension(user): Extension<AuthUser>,
|
||||
Json(req): Json<RedeemReq>,
|
||||
) -> WebResult<Response> {
|
||||
let acct = account_id_for(&state, user.0).await?;
|
||||
match crate::topup::redeem(&state.pool, acct, &req.code).await {
|
||||
Ok(new_total) => Ok(Json(json!({ "allocation_total": new_total })).into_response()),
|
||||
Err(crate::topup::TopUpError::Invalid) => {
|
||||
Err(WebError::BadRequest("invalid or already-redeemed code"))
|
||||
}
|
||||
Err(crate::topup::TopUpError::Db(e)) => {
|
||||
tracing::error!(error = %e, "redeem db error");
|
||||
Err(WebError::Internal)
|
||||
}
|
||||
}
|
||||
}
|
||||
243
crates/helexa-upstream/tests/authz_pg.rs
Normal file
243
crates/helexa-upstream/tests/authz_pg.rs
Normal file
@@ -0,0 +1,243 @@
|
||||
//! Integration tests for the `/authz/v1` surface against a real Postgres,
|
||||
//! driving the built axum app over HTTP. Gated on `UPSTREAM_TEST_DATABASE_URL`
|
||||
//! (skips cleanly when unset, so CI stays green without a DB):
|
||||
//!
|
||||
//! UPSTREAM_TEST_DATABASE_URL=postgres://helexa:helexa@localhost/helexa_test \
|
||||
//! cargo test -p helexa-upstream --test authz_pg
|
||||
|
||||
use helexa_upstream::config::{ClientToken, UpstreamConfig};
|
||||
use helexa_upstream::crypto::sha256;
|
||||
use helexa_upstream::db::connect_and_migrate;
|
||||
use helexa_upstream::state::AppState;
|
||||
use serde_json::{Value, json};
|
||||
use sqlx::Executor;
|
||||
use sqlx::Row;
|
||||
use sqlx::postgres::PgPool;
|
||||
use uuid::Uuid;
|
||||
|
||||
const CLIENT_TOKEN: &str = "test-operator-token";
|
||||
|
||||
async fn spawn_or_skip(test: &str) -> Option<(String, PgPool)> {
|
||||
let Ok(url) = std::env::var("UPSTREAM_TEST_DATABASE_URL") else {
|
||||
eprintln!("skipping {test}: UPSTREAM_TEST_DATABASE_URL not set");
|
||||
return None;
|
||||
};
|
||||
let pool = connect_and_migrate(&url, 16).await.expect("migrate");
|
||||
|
||||
let mut config = UpstreamConfig {
|
||||
server: Default::default(),
|
||||
db: helexa_upstream::config::DbSettings {
|
||||
url,
|
||||
max_connections: 16,
|
||||
},
|
||||
grant: Default::default(),
|
||||
abuse: Default::default(),
|
||||
client_auth: Default::default(),
|
||||
authz: Default::default(),
|
||||
auth: Default::default(),
|
||||
email: Default::default(),
|
||||
};
|
||||
config.client_auth.tokens.push(ClientToken {
|
||||
token: CLIENT_TOKEN.into(),
|
||||
operator_id: "op-test".into(),
|
||||
});
|
||||
|
||||
let email = helexa_upstream::email::EmailSender::from_config(&config.email).unwrap();
|
||||
let state = AppState::new(pool.clone(), config, email);
|
||||
let app = helexa_upstream::build_app(state);
|
||||
let listener = tokio::net::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();
|
||||
});
|
||||
Some((format!("http://{addr}"), pool))
|
||||
}
|
||||
|
||||
/// Seed an account with `total` allocation and an active key with raw value
|
||||
/// `raw` (percent=100). Optionally deactivate the account. Returns
|
||||
/// (account_id, key_id).
|
||||
async fn seed_key(pool: &PgPool, total: i64, raw: &str, deactivated: bool) -> (Uuid, Uuid) {
|
||||
let user_id: Uuid = pool
|
||||
.fetch_one(
|
||||
sqlx::query(
|
||||
"INSERT INTO users (email, password_hash, email_verified) VALUES ($1,'x',true) RETURNING id",
|
||||
)
|
||||
.bind(format!("u-{}@t.local", Uuid::new_v4())),
|
||||
)
|
||||
.await
|
||||
.unwrap()
|
||||
.get("id");
|
||||
let status = if deactivated { "deactivated" } else { "active" };
|
||||
let account_id: Uuid = pool
|
||||
.fetch_one(
|
||||
sqlx::query(
|
||||
"INSERT INTO accounts (owner_user_id, allocation_total, status) VALUES ($1,$2,$3) RETURNING id",
|
||||
)
|
||||
.bind(user_id)
|
||||
.bind(total)
|
||||
.bind(status),
|
||||
)
|
||||
.await
|
||||
.unwrap()
|
||||
.get("id");
|
||||
let key_id: Uuid = pool
|
||||
.fetch_one(
|
||||
sqlx::query(
|
||||
"INSERT INTO api_keys (account_id, key_hash, key_prefix, limit_kind, limit_value) \
|
||||
VALUES ($1,$2,'sk-test','percent',100) RETURNING id",
|
||||
)
|
||||
.bind(account_id)
|
||||
.bind(sha256(raw)),
|
||||
)
|
||||
.await
|
||||
.unwrap()
|
||||
.get("id");
|
||||
(account_id, key_id)
|
||||
}
|
||||
|
||||
fn client() -> reqwest::Client {
|
||||
reqwest::Client::new()
|
||||
}
|
||||
|
||||
async fn post(
|
||||
c: &reqwest::Client,
|
||||
url: String,
|
||||
body: Value,
|
||||
bearer: Option<&str>,
|
||||
) -> reqwest::Response {
|
||||
let mut req = c.post(url).json(&body);
|
||||
if let Some(b) = bearer {
|
||||
req = req.bearer_auth(b);
|
||||
}
|
||||
req.send().await.unwrap()
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn resolve_reserve_settle_round_trip() {
|
||||
let Some((base, pool)) = spawn_or_skip("resolve_reserve_settle_round_trip").await else {
|
||||
return;
|
||||
};
|
||||
let raw = format!("sk-{}", Uuid::new_v4());
|
||||
let (account_id, key_id) = seed_key(&pool, 1000, &raw, false).await;
|
||||
let c = client();
|
||||
|
||||
// resolve
|
||||
let r = post(
|
||||
&c,
|
||||
format!("{base}/authz/v1/resolve"),
|
||||
json!({"api_key": raw}),
|
||||
Some(CLIENT_TOKEN),
|
||||
)
|
||||
.await;
|
||||
assert_eq!(r.status(), 200);
|
||||
let body: Value = r.json().await.unwrap();
|
||||
assert_eq!(body["principal"]["account_id"], account_id.to_string());
|
||||
assert_eq!(body["principal"]["key_id"], key_id.to_string());
|
||||
assert_eq!(body["snapshot"]["hard_cap"], 1000);
|
||||
|
||||
// reserve 400
|
||||
let r = post(
|
||||
&c,
|
||||
format!("{base}/authz/v1/reserve"),
|
||||
json!({"account_id": account_id, "key_id": key_id, "max_tokens": 400}),
|
||||
Some(CLIENT_TOKEN),
|
||||
)
|
||||
.await;
|
||||
assert_eq!(r.status(), 200);
|
||||
let body: Value = r.json().await.unwrap();
|
||||
let rid = body["reservation_id"].as_i64().expect("granted");
|
||||
|
||||
// settle 150
|
||||
let r = post(
|
||||
&c,
|
||||
format!("{base}/authz/v1/settle"),
|
||||
json!({"reservation_id": rid, "actual_tokens": 150}),
|
||||
Some(CLIENT_TOKEN),
|
||||
)
|
||||
.await;
|
||||
assert_eq!(r.status(), 204);
|
||||
|
||||
// snapshot reflects spend
|
||||
let r = post(
|
||||
&c,
|
||||
format!("{base}/authz/v1/snapshot"),
|
||||
json!({"account_id": account_id, "key_id": key_id}),
|
||||
Some(CLIENT_TOKEN),
|
||||
)
|
||||
.await;
|
||||
let body: Value = r.json().await.unwrap();
|
||||
assert_eq!(body["spent"], 150);
|
||||
assert_eq!(body["reserved"], 0);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn over_cap_reserve_is_rejected_not_errored() {
|
||||
let Some((base, pool)) = spawn_or_skip("over_cap_reserve_is_rejected_not_errored").await else {
|
||||
return;
|
||||
};
|
||||
let raw = format!("sk-{}", Uuid::new_v4());
|
||||
let (account_id, key_id) = seed_key(&pool, 100, &raw, false).await;
|
||||
let c = client();
|
||||
let r = post(
|
||||
&c,
|
||||
format!("{base}/authz/v1/reserve"),
|
||||
json!({"account_id": account_id, "key_id": key_id, "max_tokens": 999}),
|
||||
Some(CLIENT_TOKEN),
|
||||
)
|
||||
.await;
|
||||
assert_eq!(r.status(), 200, "budget refusal is an authoritative 200");
|
||||
let body: Value = r.json().await.unwrap();
|
||||
assert!(body["reservation_id"].is_null());
|
||||
assert_eq!(body["rejected"]["kind"], "insufficient_quota");
|
||||
assert_eq!(body["rejected"]["available"], 100);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn deactivated_account_resolves_as_invalid_no_clue() {
|
||||
let Some((base, pool)) = spawn_or_skip("deactivated_account_resolves_as_invalid_no_clue").await
|
||||
else {
|
||||
return;
|
||||
};
|
||||
let raw = format!("sk-{}", Uuid::new_v4());
|
||||
seed_key(&pool, 1000, &raw, true).await; // deactivated
|
||||
let c = client();
|
||||
let r = post(
|
||||
&c,
|
||||
format!("{base}/authz/v1/resolve"),
|
||||
json!({"api_key": raw}),
|
||||
Some(CLIENT_TOKEN),
|
||||
)
|
||||
.await;
|
||||
// Indistinguishable from an unknown key.
|
||||
assert_eq!(r.status(), 401);
|
||||
let body: Value = r.json().await.unwrap();
|
||||
assert_eq!(body["error"]["code"], "invalid_api_key");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn missing_client_auth_is_401_before_db() {
|
||||
let Some((base, pool)) = spawn_or_skip("missing_client_auth_is_401_before_db").await else {
|
||||
return;
|
||||
};
|
||||
let raw = format!("sk-{}", Uuid::new_v4());
|
||||
seed_key(&pool, 1000, &raw, false).await;
|
||||
let c = client();
|
||||
// No bearer → rejected by client_auth.
|
||||
let r = post(
|
||||
&c,
|
||||
format!("{base}/authz/v1/resolve"),
|
||||
json!({"api_key": raw}),
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
assert_eq!(r.status(), 401);
|
||||
// Wrong bearer → also rejected.
|
||||
let r = post(
|
||||
&c,
|
||||
format!("{base}/authz/v1/resolve"),
|
||||
json!({"api_key": raw}),
|
||||
Some("wrong"),
|
||||
)
|
||||
.await;
|
||||
assert_eq!(r.status(), 401);
|
||||
}
|
||||
204
crates/helexa-upstream/tests/ledger_pg.rs
Normal file
204
crates/helexa-upstream/tests/ledger_pg.rs
Normal file
@@ -0,0 +1,204 @@
|
||||
//! Integration tests for the allocation ledger against a real PostgreSQL.
|
||||
//!
|
||||
//! Gated on `UPSTREAM_TEST_DATABASE_URL` — when unset (CI's generic runner,
|
||||
//! local builds without a DB), every test logs a skip and returns, so
|
||||
//! `cargo test --workspace` stays green without Postgres. Point the env var
|
||||
//! at a throwaway database to exercise the no-overshoot guarantee and
|
||||
//! settle/release idempotency:
|
||||
//!
|
||||
//! UPSTREAM_TEST_DATABASE_URL=postgres://helexa:helexa@localhost/helexa_test \
|
||||
//! cargo test -p helexa-upstream --test ledger_pg
|
||||
|
||||
use helexa_upstream::db::connect_and_migrate;
|
||||
use helexa_upstream::ledger::{self, LedgerError};
|
||||
use sqlx::Executor;
|
||||
use sqlx::Row;
|
||||
use sqlx::postgres::PgPool;
|
||||
use uuid::Uuid;
|
||||
|
||||
/// Returns a migrated pool, or `None` (with a skip log) when the env var is
|
||||
/// unset.
|
||||
async fn pool_or_skip(test: &str) -> Option<PgPool> {
|
||||
let Ok(url) = std::env::var("UPSTREAM_TEST_DATABASE_URL") else {
|
||||
eprintln!("skipping {test}: UPSTREAM_TEST_DATABASE_URL not set");
|
||||
return None;
|
||||
};
|
||||
Some(
|
||||
connect_and_migrate(&url, 16)
|
||||
.await
|
||||
.expect("connect + migrate"),
|
||||
)
|
||||
}
|
||||
|
||||
/// Seed a verified user + account (with `total` allocation) + an active key
|
||||
/// (percent=100 so the account cap binds). Returns (account_id, key_id).
|
||||
async fn seed(pool: &PgPool, total: i64) -> (Uuid, Uuid) {
|
||||
let user_id: Uuid = pool
|
||||
.fetch_one(
|
||||
sqlx::query(
|
||||
"INSERT INTO users (email, password_hash, email_verified) \
|
||||
VALUES ($1, 'x', true) RETURNING id",
|
||||
)
|
||||
.bind(format!("u-{}@test.local", Uuid::new_v4())),
|
||||
)
|
||||
.await
|
||||
.unwrap()
|
||||
.get("id");
|
||||
let account_id: Uuid = pool
|
||||
.fetch_one(
|
||||
sqlx::query(
|
||||
"INSERT INTO accounts (owner_user_id, allocation_total) \
|
||||
VALUES ($1, $2) RETURNING id",
|
||||
)
|
||||
.bind(user_id)
|
||||
.bind(total),
|
||||
)
|
||||
.await
|
||||
.unwrap()
|
||||
.get("id");
|
||||
let key_id: Uuid = pool
|
||||
.fetch_one(
|
||||
sqlx::query(
|
||||
"INSERT INTO api_keys (account_id, key_hash, key_prefix, limit_kind, limit_value) \
|
||||
VALUES ($1, $2, 'sk-test', 'percent', 100) RETURNING id",
|
||||
)
|
||||
.bind(account_id)
|
||||
.bind(Uuid::new_v4().as_bytes().to_vec()),
|
||||
)
|
||||
.await
|
||||
.unwrap()
|
||||
.get("id");
|
||||
(account_id, key_id)
|
||||
}
|
||||
|
||||
async fn account_cols(pool: &PgPool, account_id: Uuid) -> (i64, i64) {
|
||||
let row = pool
|
||||
.fetch_one(
|
||||
sqlx::query("SELECT allocation_spent, allocation_reserved FROM accounts WHERE id = $1")
|
||||
.bind(account_id),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
(row.get("allocation_spent"), row.get("allocation_reserved"))
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn concurrent_reserves_never_overshoot() {
|
||||
let Some(pool) = pool_or_skip("concurrent_reserves_never_overshoot").await else {
|
||||
return;
|
||||
};
|
||||
// Allocation admits exactly 5 reservations of 100 (cap 500).
|
||||
let (account_id, key_id) = seed(&pool, 500).await;
|
||||
|
||||
let mut handles = Vec::new();
|
||||
for _ in 0..20 {
|
||||
let pool = pool.clone();
|
||||
handles.push(tokio::spawn(async move {
|
||||
ledger::reserve(&pool, account_id, key_id, 100).await
|
||||
}));
|
||||
}
|
||||
let mut ok = 0;
|
||||
let mut quota = 0;
|
||||
for h in handles {
|
||||
match h.await.unwrap() {
|
||||
Ok(_) => ok += 1,
|
||||
Err(LedgerError::InsufficientQuota { .. }) => quota += 1,
|
||||
Err(e) => panic!("unexpected error: {e}"),
|
||||
}
|
||||
}
|
||||
assert_eq!(ok, 5, "exactly 5 reserves of 100 fit in a 500 allocation");
|
||||
assert_eq!(quota, 15);
|
||||
|
||||
let (spent, reserved) = account_cols(&pool, account_id).await;
|
||||
assert_eq!(spent, 0);
|
||||
assert_eq!(reserved, 500, "reserved exactly the cap, never over");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn settle_is_idempotent_and_reconciles_spend() {
|
||||
let Some(pool) = pool_or_skip("settle_is_idempotent_and_reconciles_spend").await else {
|
||||
return;
|
||||
};
|
||||
let (account_id, key_id) = seed(&pool, 1000).await;
|
||||
let rid = ledger::reserve(&pool, account_id, key_id, 400)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Settle actual=150 (< reserved 400): spent=150, reserved back to 0.
|
||||
ledger::settle(&pool, rid, 150).await.unwrap();
|
||||
let (spent, reserved) = account_cols(&pool, account_id).await;
|
||||
assert_eq!((spent, reserved), (150, 0));
|
||||
|
||||
// Second settle is a no-op.
|
||||
ledger::settle(&pool, rid, 999).await.unwrap();
|
||||
let (spent2, reserved2) = account_cols(&pool, account_id).await;
|
||||
assert_eq!((spent2, reserved2), (150, 0), "settle is idempotent");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn release_returns_reservation_and_is_idempotent() {
|
||||
let Some(pool) = pool_or_skip("release_returns_reservation_and_is_idempotent").await else {
|
||||
return;
|
||||
};
|
||||
let (account_id, key_id) = seed(&pool, 1000).await;
|
||||
let rid = ledger::reserve(&pool, account_id, key_id, 300)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(account_cols(&pool, account_id).await, (0, 300));
|
||||
|
||||
ledger::release(&pool, rid).await.unwrap();
|
||||
assert_eq!(account_cols(&pool, account_id).await, (0, 0));
|
||||
// Idempotent; settle-after-release also a no-op.
|
||||
ledger::release(&pool, rid).await.unwrap();
|
||||
ledger::settle(&pool, rid, 100).await.unwrap();
|
||||
assert_eq!(account_cols(&pool, account_id).await, (0, 0));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn hardcap_key_subcap_binds_below_account() {
|
||||
let Some(pool) = pool_or_skip("hardcap_key_subcap_binds_below_account").await else {
|
||||
return;
|
||||
};
|
||||
// Account has 1000 but the key is hard-capped at 200.
|
||||
let user_id: Uuid = pool
|
||||
.fetch_one(
|
||||
sqlx::query(
|
||||
"INSERT INTO users (email, password_hash, email_verified) \
|
||||
VALUES ($1, 'x', true) RETURNING id",
|
||||
)
|
||||
.bind(format!("u-{}@test.local", Uuid::new_v4())),
|
||||
)
|
||||
.await
|
||||
.unwrap()
|
||||
.get("id");
|
||||
let account_id: Uuid = pool
|
||||
.fetch_one(
|
||||
sqlx::query(
|
||||
"INSERT INTO accounts (owner_user_id, allocation_total) VALUES ($1, 1000) RETURNING id",
|
||||
)
|
||||
.bind(user_id),
|
||||
)
|
||||
.await
|
||||
.unwrap()
|
||||
.get("id");
|
||||
let key_id: Uuid = pool
|
||||
.fetch_one(
|
||||
sqlx::query(
|
||||
"INSERT INTO api_keys (account_id, key_hash, key_prefix, limit_kind, limit_value) \
|
||||
VALUES ($1, $2, 'sk-test', 'hardcap', 200) RETURNING id",
|
||||
)
|
||||
.bind(account_id)
|
||||
.bind(Uuid::new_v4().as_bytes().to_vec()),
|
||||
)
|
||||
.await
|
||||
.unwrap()
|
||||
.get("id");
|
||||
|
||||
ledger::reserve(&pool, account_id, key_id, 200)
|
||||
.await
|
||||
.unwrap();
|
||||
match ledger::reserve(&pool, account_id, key_id, 1).await {
|
||||
Err(LedgerError::InsufficientQuota { available, .. }) => assert_eq!(available, 0),
|
||||
other => panic!("expected InsufficientQuota, got {other:?}"),
|
||||
}
|
||||
}
|
||||
425
crates/helexa-upstream/tests/web_pg.rs
Normal file
425
crates/helexa-upstream/tests/web_pg.rs
Normal file
@@ -0,0 +1,425 @@
|
||||
//! Integration tests for the `/web/v1` account API + the silent fingerprint
|
||||
//! abuse policy, driving the built app over HTTP against a real Postgres.
|
||||
//! Gated on `UPSTREAM_TEST_DATABASE_URL` (skips cleanly when unset).
|
||||
|
||||
use helexa_upstream::config::{ClientToken, UpstreamConfig};
|
||||
use helexa_upstream::crypto::sha256;
|
||||
use helexa_upstream::db::connect_and_migrate;
|
||||
use helexa_upstream::email::EmailSender;
|
||||
use helexa_upstream::state::AppState;
|
||||
use serde_json::{Value, json};
|
||||
use sqlx::Executor;
|
||||
use sqlx::Row;
|
||||
use sqlx::postgres::PgPool;
|
||||
|
||||
const CLIENT_TOKEN: &str = "web-test-operator-token";
|
||||
|
||||
async fn spawn_or_skip(test: &str) -> Option<(String, PgPool)> {
|
||||
let Ok(url) = std::env::var("UPSTREAM_TEST_DATABASE_URL") else {
|
||||
eprintln!("skipping {test}: UPSTREAM_TEST_DATABASE_URL not set");
|
||||
return None;
|
||||
};
|
||||
let pool = connect_and_migrate(&url, 16).await.expect("migrate");
|
||||
let mut config = UpstreamConfig {
|
||||
server: Default::default(),
|
||||
db: helexa_upstream::config::DbSettings {
|
||||
url,
|
||||
max_connections: 16,
|
||||
},
|
||||
grant: Default::default(),
|
||||
abuse: Default::default(),
|
||||
client_auth: Default::default(),
|
||||
authz: Default::default(),
|
||||
auth: Default::default(),
|
||||
email: Default::default(), // Log transport
|
||||
};
|
||||
config.client_auth.tokens.push(ClientToken {
|
||||
token: CLIENT_TOKEN.into(),
|
||||
operator_id: "op-web".into(),
|
||||
});
|
||||
let email = EmailSender::from_config(&config.email).unwrap();
|
||||
let state = AppState::new(pool.clone(), config, email);
|
||||
let app = helexa_upstream::build_app(state);
|
||||
let listener = tokio::net::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();
|
||||
});
|
||||
Some((format!("http://{addr}"), pool))
|
||||
}
|
||||
|
||||
fn unique_email() -> String {
|
||||
format!("u-{}@test.local", uuid::Uuid::new_v4())
|
||||
}
|
||||
|
||||
async fn post(url: String, body: Value, bearer: Option<&str>) -> reqwest::Response {
|
||||
let c = reqwest::Client::new();
|
||||
let mut req = c.post(url).json(&body);
|
||||
if let Some(b) = bearer {
|
||||
req = req.bearer_auth(b);
|
||||
}
|
||||
req.send().await.unwrap()
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn verify_endpoint_consumes_token_once() {
|
||||
let Some((base, pool)) = spawn_or_skip("verify_endpoint_consumes_token_once").await else {
|
||||
return;
|
||||
};
|
||||
let email = unique_email();
|
||||
// Register, then mint a verify token directly (the raw token is only in
|
||||
// the email; here we insert a known one to drive the endpoint).
|
||||
assert_eq!(
|
||||
post(
|
||||
format!("{base}/web/v1/register"),
|
||||
json!({"email": email, "password": "password123"}),
|
||||
None
|
||||
)
|
||||
.await
|
||||
.status(),
|
||||
202
|
||||
);
|
||||
let user_id: uuid::Uuid = pool
|
||||
.fetch_one(sqlx::query("SELECT id FROM users WHERE email = $1").bind(&email))
|
||||
.await
|
||||
.unwrap()
|
||||
.get("id");
|
||||
let raw = "verify-raw-token-xyz";
|
||||
pool.execute(
|
||||
sqlx::query(
|
||||
"INSERT INTO email_tokens (token_hash, user_id, kind, expires_at) \
|
||||
VALUES ($1, $2, 'verify', now() + interval '1 hour')",
|
||||
)
|
||||
.bind(sha256(raw))
|
||||
.bind(user_id),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
post(format!("{base}/web/v1/verify"), json!({"token": raw}), None)
|
||||
.await
|
||||
.status(),
|
||||
200
|
||||
);
|
||||
// Consumed → second attempt fails.
|
||||
assert_eq!(
|
||||
post(format!("{base}/web/v1/verify"), json!({"token": raw}), None)
|
||||
.await
|
||||
.status(),
|
||||
400
|
||||
);
|
||||
|
||||
let verified: bool = pool
|
||||
.fetch_one(sqlx::query("SELECT email_verified FROM users WHERE id = $1").bind(user_id))
|
||||
.await
|
||||
.unwrap()
|
||||
.get("email_verified");
|
||||
assert!(verified);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn account_lifecycle_and_key_resolves_then_archives() {
|
||||
let Some((base, pool)) =
|
||||
spawn_or_skip("account_lifecycle_and_key_resolves_then_archives").await
|
||||
else {
|
||||
return;
|
||||
};
|
||||
let email = unique_email();
|
||||
post(
|
||||
format!("{base}/web/v1/register"),
|
||||
json!({"email": email, "password": "password123"}),
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
// Bypass the email step for the login/key portion.
|
||||
pool.execute(
|
||||
sqlx::query("UPDATE users SET email_verified = true WHERE email = $1").bind(&email),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// login → session JWT
|
||||
let r = post(
|
||||
format!("{base}/web/v1/login"),
|
||||
json!({"email": email, "password": "password123"}),
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
assert_eq!(r.status(), 200);
|
||||
let token = r.json::<Value>().await.unwrap()["token"]
|
||||
.as_str()
|
||||
.unwrap()
|
||||
.to_string();
|
||||
|
||||
// create key (raw shown once)
|
||||
let r = post(
|
||||
format!("{base}/web/v1/keys"),
|
||||
json!({"label": "laptop"}),
|
||||
Some(&token),
|
||||
)
|
||||
.await;
|
||||
assert_eq!(r.status(), 201);
|
||||
let body: Value = r.json().await.unwrap();
|
||||
let raw_key = body["key"].as_str().unwrap().to_string();
|
||||
let key_id = body["id"].as_str().unwrap().to_string();
|
||||
assert!(raw_key.starts_with("sk-helexa-"));
|
||||
|
||||
// account balance reflects the free grant
|
||||
let r = reqwest::Client::new()
|
||||
.get(format!("{base}/web/v1/account"))
|
||||
.bearer_auth(&token)
|
||||
.send()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
r.json::<Value>().await.unwrap()["allocation_total"],
|
||||
1_000_000
|
||||
);
|
||||
|
||||
// list keys shows the prefix, never the raw secret
|
||||
let r = reqwest::Client::new()
|
||||
.get(format!("{base}/web/v1/keys"))
|
||||
.bearer_auth(&token)
|
||||
.send()
|
||||
.await
|
||||
.unwrap();
|
||||
let listed = r.json::<Value>().await.unwrap();
|
||||
let k = &listed["keys"][0];
|
||||
assert_eq!(k["id"], key_id);
|
||||
assert!(k.get("key").is_none(), "raw secret never listed");
|
||||
assert!(k["prefix"].as_str().unwrap().starts_with("sk-helexa-"));
|
||||
|
||||
// the key authorizes at the authz surface
|
||||
let r = post(
|
||||
format!("{base}/authz/v1/resolve"),
|
||||
json!({"api_key": raw_key}),
|
||||
Some(CLIENT_TOKEN),
|
||||
)
|
||||
.await;
|
||||
assert_eq!(r.status(), 200);
|
||||
|
||||
// archive → the key no longer resolves
|
||||
let r = post(
|
||||
format!("{base}/web/v1/keys/{key_id}/archive"),
|
||||
json!({}),
|
||||
Some(&token),
|
||||
)
|
||||
.await;
|
||||
assert_eq!(r.status(), 204);
|
||||
let r = post(
|
||||
format!("{base}/authz/v1/resolve"),
|
||||
json!({"api_key": raw_key}),
|
||||
Some(CLIENT_TOKEN),
|
||||
)
|
||||
.await;
|
||||
assert_eq!(r.status(), 401);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn fingerprint_abuse_silently_deactivates_all_no_clue() {
|
||||
let Some((base, pool)) =
|
||||
spawn_or_skip("fingerprint_abuse_silently_deactivates_all_no_clue").await
|
||||
else {
|
||||
return;
|
||||
};
|
||||
let fp = format!("fp-{}", uuid::Uuid::new_v4());
|
||||
|
||||
// 5 registrations sharing one fingerprint — every one returns a normal 202.
|
||||
let mut emails = Vec::new();
|
||||
for _ in 0..5 {
|
||||
let email = unique_email();
|
||||
let r = post(
|
||||
format!("{base}/web/v1/register"),
|
||||
json!({"email": email, "password": "password123", "fingerprint": fp}),
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
assert_eq!(r.status(), 202, "registration always looks successful");
|
||||
emails.push(email);
|
||||
}
|
||||
|
||||
// Silent effect: all 5 accounts are deactivated + flagged.
|
||||
let (deactivated, flagged): (i64, i64) = {
|
||||
let row = pool
|
||||
.fetch_one(
|
||||
sqlx::query(
|
||||
"SELECT \
|
||||
count(*) FILTER (WHERE a.status = 'deactivated') AS d, \
|
||||
count(*) FILTER (WHERE a.fingerprint_flagged) AS f \
|
||||
FROM accounts a JOIN users u ON u.id = a.owner_user_id \
|
||||
WHERE u.registration_fingerprint = $1",
|
||||
)
|
||||
.bind(&fp),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
(row.get("d"), row.get("f"))
|
||||
};
|
||||
assert_eq!(deactivated, 5, "all sharing accounts silently deactivated");
|
||||
assert_eq!(flagged, 5);
|
||||
|
||||
// No clue at the authz surface: a key on a deactivated account resolves
|
||||
// as an ordinary 401, indistinguishable from an unknown key.
|
||||
let acct: uuid::Uuid = pool
|
||||
.fetch_one(
|
||||
sqlx::query(
|
||||
"SELECT a.id FROM accounts a JOIN users u ON u.id = a.owner_user_id \
|
||||
WHERE u.registration_fingerprint = $1 LIMIT 1",
|
||||
)
|
||||
.bind(&fp),
|
||||
)
|
||||
.await
|
||||
.unwrap()
|
||||
.get("id");
|
||||
let raw = "sk-helexa-deactivated-probe";
|
||||
pool.execute(
|
||||
sqlx::query(
|
||||
"INSERT INTO api_keys (account_id, key_hash, key_prefix) VALUES ($1, $2, 'sk-helexa-')",
|
||||
)
|
||||
.bind(acct)
|
||||
.bind(sha256(raw)),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let r = post(
|
||||
format!("{base}/authz/v1/resolve"),
|
||||
json!({"api_key": raw}),
|
||||
Some(CLIENT_TOKEN),
|
||||
)
|
||||
.await;
|
||||
assert_eq!(
|
||||
r.status(),
|
||||
401,
|
||||
"deactivated account's key looks like any invalid key"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn topup_redeem_raises_allocation_single_use() {
|
||||
let Some((base, pool)) = spawn_or_skip("topup_redeem_raises_allocation_single_use").await
|
||||
else {
|
||||
return;
|
||||
};
|
||||
let email = unique_email();
|
||||
post(
|
||||
format!("{base}/web/v1/register"),
|
||||
json!({"email": email, "password": "password123"}),
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
pool.execute(
|
||||
sqlx::query("UPDATE users SET email_verified = true WHERE email = $1").bind(&email),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let token = post(
|
||||
format!("{base}/web/v1/login"),
|
||||
json!({"email": email, "password": "password123"}),
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.json::<Value>()
|
||||
.await
|
||||
.unwrap()["token"]
|
||||
.as_str()
|
||||
.unwrap()
|
||||
.to_string();
|
||||
|
||||
// Mint a code worth 500_000 (mint path used by the CLI/faucet).
|
||||
let codes = helexa_upstream::topup::mint(&pool, 500_000, 1, Some("test"))
|
||||
.await
|
||||
.unwrap();
|
||||
let code = &codes[0];
|
||||
|
||||
// Redeem → allocation_total rises from the 1_000_000 free grant.
|
||||
let r = post(
|
||||
format!("{base}/web/v1/redeem"),
|
||||
json!({"code": code}),
|
||||
Some(&token),
|
||||
)
|
||||
.await;
|
||||
assert_eq!(r.status(), 200);
|
||||
assert_eq!(
|
||||
r.json::<Value>().await.unwrap()["allocation_total"],
|
||||
1_500_000
|
||||
);
|
||||
|
||||
// Single-use: a second redemption fails generically (no oracle).
|
||||
let r = post(
|
||||
format!("{base}/web/v1/redeem"),
|
||||
json!({"code": code}),
|
||||
Some(&token),
|
||||
)
|
||||
.await;
|
||||
assert_eq!(r.status(), 400);
|
||||
|
||||
// Unknown code: same generic 400.
|
||||
let r = post(
|
||||
format!("{base}/web/v1/redeem"),
|
||||
json!({"code": "helexa-topup-does-not-exist"}),
|
||||
Some(&token),
|
||||
)
|
||||
.await;
|
||||
assert_eq!(r.status(), 400);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn topup_concurrent_double_redeem_one_winner() {
|
||||
let Some((base, pool)) = spawn_or_skip("topup_concurrent_double_redeem_one_winner").await
|
||||
else {
|
||||
return;
|
||||
};
|
||||
// Two verified accounts.
|
||||
let mut tokens = Vec::new();
|
||||
for _ in 0..2 {
|
||||
let email = unique_email();
|
||||
post(
|
||||
format!("{base}/web/v1/register"),
|
||||
json!({"email": email, "password": "password123"}),
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
pool.execute(
|
||||
sqlx::query("UPDATE users SET email_verified = true WHERE email = $1").bind(&email),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let t = post(
|
||||
format!("{base}/web/v1/login"),
|
||||
json!({"email": email, "password": "password123"}),
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.json::<Value>()
|
||||
.await
|
||||
.unwrap()["token"]
|
||||
.as_str()
|
||||
.unwrap()
|
||||
.to_string();
|
||||
tokens.push(t);
|
||||
}
|
||||
let code = helexa_upstream::topup::mint(&pool, 100, 1, None)
|
||||
.await
|
||||
.unwrap()
|
||||
.remove(0);
|
||||
|
||||
// Both accounts race to redeem the same code; exactly one wins.
|
||||
let (a, b) = tokio::join!(
|
||||
post(
|
||||
format!("{base}/web/v1/redeem"),
|
||||
json!({"code": code}),
|
||||
Some(&tokens[0])
|
||||
),
|
||||
post(
|
||||
format!("{base}/web/v1/redeem"),
|
||||
json!({"code": code}),
|
||||
Some(&tokens[1])
|
||||
),
|
||||
);
|
||||
let wins = [a.status(), b.status()]
|
||||
.iter()
|
||||
.filter(|s| s.as_u16() == 200)
|
||||
.count();
|
||||
assert_eq!(wins, 1, "exactly one redemption wins the single-use code");
|
||||
}
|
||||
@@ -11,17 +11,28 @@
|
||||
# the router — the router never owns an inbound TLS listener.
|
||||
listen = "0.0.0.0:8088"
|
||||
|
||||
# How often (seconds) to refresh each cortex's health + /v1/models topology.
|
||||
# poll_interval_secs = 10
|
||||
|
||||
# -- 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"
|
||||
|
||||
52
helexa-upstream.example.toml
Normal file
52
helexa-upstream.example.toml
Normal file
@@ -0,0 +1,52 @@
|
||||
# helexa-upstream.example.toml — mesh-level account/authorization authority
|
||||
#
|
||||
# Copy to helexa-upstream.toml and adjust. Env overrides use the UPSTREAM_
|
||||
# prefix with __ separators, e.g. UPSTREAM_DB__URL=postgres://...
|
||||
|
||||
[server]
|
||||
# Plaintext listener; edge nginx terminates TLS (consistent with the stack).
|
||||
listen = "0.0.0.0:8090"
|
||||
|
||||
[db]
|
||||
# PostgreSQL connection URL. Required.
|
||||
url = "postgres://helexa:helexa@localhost/helexa_upstream"
|
||||
# max_connections = 16
|
||||
|
||||
[grant]
|
||||
# Flat free token grant every email-verified account receives (the floor of
|
||||
# the hybrid allocation; single-use top-up codes extend it).
|
||||
# free_token_grant = 1000000
|
||||
|
||||
[abuse]
|
||||
# When this many accounts share one registration fingerprint, all are
|
||||
# silently deactivated (no notice to the user).
|
||||
# fingerprint_account_threshold = 5
|
||||
|
||||
# -- Client auth: credentials operators' cortexes present to /authz/v1.
|
||||
# Each token maps to an operator_id (served-usage attribution). When no
|
||||
# tokens are configured the authz surface is OPEN (dev only). Distinct from
|
||||
# end-user API keys, which ride inside the resolve request body.
|
||||
# [[client_auth.tokens]]
|
||||
# token = "replace-with-a-strong-shared-secret"
|
||||
# operator_id = "lair-cafe"
|
||||
|
||||
[authz]
|
||||
# Open reservations older than this are swept (released), self-healing a
|
||||
# reservation whose settle/release from a cortex was lost.
|
||||
# reservation_ttl_secs = 120
|
||||
# sweep_interval_secs = 60
|
||||
|
||||
[auth]
|
||||
# HMAC secret for signing web-session JWTs. MUST be overridden in prod via
|
||||
# UPSTREAM_AUTH__JWT_SECRET; the built-in default is dev-only.
|
||||
# jwt_secret = "change-me"
|
||||
# session_ttl_secs = 604800 # 7 days
|
||||
# email_token_ttl_secs = 86400 # 24 hours
|
||||
# Frontend base URL used to build verify/reset links in emails.
|
||||
app_base_url = "https://helexa.ai"
|
||||
|
||||
[email]
|
||||
# "log" (dev: logs the link) or "smtp".
|
||||
provider = "log"
|
||||
# smtp_url = "smtp://user:pass@smtp.example.com:587"
|
||||
from_addr = "helexa <no-reply@helexa.ai>"
|
||||
20
helexa.ai/.env.example
Normal file
20
helexa.ai/.env.example
Normal file
@@ -0,0 +1,20 @@
|
||||
# helexa.ai frontend env. Copy to .env.local for local dev (gitignored).
|
||||
|
||||
# Mesh data-plane (helexa-router, OpenAI-compatible inference). In dev,
|
||||
# vite proxies /v1 and /health here.
|
||||
VITE_ROUTER_BASE_URL=http://localhost:8088
|
||||
|
||||
# Account control-plane (helexa-upstream). In dev, vite proxies /api here
|
||||
# (rewritten to /web/v1).
|
||||
VITE_ACCOUNT_BASE_URL=http://localhost:8090
|
||||
|
||||
# Public-beta banner.
|
||||
VITE_PUBLIC_BETA=true
|
||||
|
||||
# Models for the chat workspace (F3+).
|
||||
# VITE_ANON_MODEL=...
|
||||
# VITE_DEFAULT_MODEL=...
|
||||
|
||||
# Develop the account dashboard (F4) against an in-browser mock before the
|
||||
# upstream account API ships.
|
||||
# VITE_USE_MOCK_ACCOUNT_API=true
|
||||
6
helexa.ai/.gitignore
vendored
Normal file
6
helexa.ai/.gitignore
vendored
Normal file
@@ -0,0 +1,6 @@
|
||||
node_modules
|
||||
dist
|
||||
*.local
|
||||
.env.local
|
||||
.env.*.local
|
||||
*.tsbuildinfo
|
||||
34
helexa.ai/README.md
Normal file
34
helexa.ai/README.md
Normal file
@@ -0,0 +1,34 @@
|
||||
# helexa.ai
|
||||
|
||||
The public-beta frontend for the helexa mesh: a chat-first landing experience
|
||||
(anonymous + authenticated, with all chat history kept client-side in
|
||||
IndexedDB — no server-side history), a `/mission` page on European digital
|
||||
sovereignty, and full account self-service (register, recover, manage API
|
||||
keys, set per-key limits, redeem top-up codes) against `helexa-upstream`.
|
||||
|
||||
Vite + React (SWC) + TypeScript + react-bootstrap + react-router + react-i18next.
|
||||
Lives as a top-level folder in the cortex monorepo; it is **not** a Cargo crate.
|
||||
|
||||
## Develop
|
||||
|
||||
```sh
|
||||
cd helexa.ai
|
||||
npm install
|
||||
cp .env.example .env.local # adjust backend URLs
|
||||
npm run dev # vite dev server, proxies /v1+/health → router, /api → upstream
|
||||
```
|
||||
|
||||
Other scripts: `npm run build` (`tsc -b && vite build` → `dist/`), `npm run
|
||||
preview`, `npm run lint`, `npm run typecheck`.
|
||||
|
||||
In dev, `vite.config.ts` proxies the mesh data-plane (helexa-router) and the
|
||||
account control-plane (helexa-upstream) same-origin. Run a local router
|
||||
(`cargo run -p helexa-router`) for the chat path and a local helexa-upstream
|
||||
for the account path.
|
||||
|
||||
## Status
|
||||
|
||||
F0 scaffold. Theming + i18n (33 languages, usage-ordered selector), the
|
||||
`/mission` page, the chat workspace (Dexie + streaming), and the account
|
||||
dashboard land in subsequent phases — see
|
||||
`~/.claude/plans/we-need-to-plan-modular-graham.md`.
|
||||
23
helexa.ai/eslint.config.js
Normal file
23
helexa.ai/eslint.config.js
Normal file
@@ -0,0 +1,23 @@
|
||||
import js from "@eslint/js";
|
||||
import globals from "globals";
|
||||
import reactHooks from "eslint-plugin-react-hooks";
|
||||
import reactRefresh from "eslint-plugin-react-refresh";
|
||||
import tseslint from "typescript-eslint";
|
||||
import { defineConfig, globalIgnores } from "eslint/config";
|
||||
|
||||
export default defineConfig([
|
||||
globalIgnores(["dist"]),
|
||||
{
|
||||
files: ["**/*.{ts,tsx}"],
|
||||
extends: [
|
||||
js.configs.recommended,
|
||||
tseslint.configs.recommended,
|
||||
reactHooks.configs.flat.recommended,
|
||||
reactRefresh.configs.vite,
|
||||
],
|
||||
languageOptions: {
|
||||
ecmaVersion: 2020,
|
||||
globals: globals.browser,
|
||||
},
|
||||
},
|
||||
]);
|
||||
24
helexa.ai/index.html
Normal file
24
helexa.ai/index.html
Normal file
@@ -0,0 +1,24 @@
|
||||
<!doctype html>
|
||||
<html lang="en">
|
||||
<head>
|
||||
<meta charset="UTF-8" />
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1.0" />
|
||||
<title>helexa.ai</title>
|
||||
<meta name="title" content="helexa.ai" />
|
||||
<meta
|
||||
name="description"
|
||||
content="helexa — near-frontier AI on a sovereign, operator-run mesh. Chat now; bring your own key."
|
||||
/>
|
||||
<meta property="og:type" content="website" />
|
||||
<meta property="og:url" content="https://helexa.ai/" />
|
||||
<meta property="og:title" content="helexa.ai" />
|
||||
<meta
|
||||
property="og:description"
|
||||
content="helexa — near-frontier AI on a sovereign, operator-run mesh."
|
||||
/>
|
||||
</head>
|
||||
<body>
|
||||
<div id="root"></div>
|
||||
<script type="module" src="/src/main.tsx"></script>
|
||||
</body>
|
||||
</html>
|
||||
4115
helexa.ai/package-lock.json
generated
Normal file
4115
helexa.ai/package-lock.json
generated
Normal file
File diff suppressed because it is too large
Load Diff
40
helexa.ai/package.json
Normal file
40
helexa.ai/package.json
Normal file
@@ -0,0 +1,40 @@
|
||||
{
|
||||
"name": "helexa.ai",
|
||||
"private": true,
|
||||
"version": "0.0.0",
|
||||
"type": "module",
|
||||
"scripts": {
|
||||
"dev": "vite",
|
||||
"build": "tsc -b && vite build",
|
||||
"preview": "vite preview",
|
||||
"lint": "eslint .",
|
||||
"typecheck": "tsc -b"
|
||||
},
|
||||
"dependencies": {
|
||||
"@fingerprintjs/fingerprintjs": "^4.6.2",
|
||||
"bootstrap": "^5.3.8",
|
||||
"dexie": "^4.2.0",
|
||||
"dexie-react-hooks": "^4.2.0",
|
||||
"i18next": "^25.7.1",
|
||||
"react": "^19.2.0",
|
||||
"react-bootstrap": "^2.10.10",
|
||||
"react-dom": "^19.2.0",
|
||||
"react-i18next": "^16.4.0",
|
||||
"react-icons": "^5.5.0",
|
||||
"react-router-dom": "^7.10.1"
|
||||
},
|
||||
"devDependencies": {
|
||||
"@eslint/js": "^9.39.1",
|
||||
"@types/node": "^24.10.1",
|
||||
"@types/react": "^19.2.5",
|
||||
"@types/react-dom": "^19.2.3",
|
||||
"@vitejs/plugin-react-swc": "^4.2.0",
|
||||
"eslint": "^9.39.1",
|
||||
"eslint-plugin-react-hooks": "^7.0.1",
|
||||
"eslint-plugin-react-refresh": "^0.4.24",
|
||||
"globals": "^16.5.0",
|
||||
"typescript": "~5.9.3",
|
||||
"typescript-eslint": "^8.46.4",
|
||||
"vite": "^7.2.0"
|
||||
}
|
||||
}
|
||||
12
helexa.ai/src/App.tsx
Normal file
12
helexa.ai/src/App.tsx
Normal file
@@ -0,0 +1,12 @@
|
||||
import { Container } from "react-bootstrap";
|
||||
|
||||
// F0 scaffold shell. Theming, i18n, routing, the chat workspace, mission
|
||||
// page and account dashboard land in the F1+ phases.
|
||||
export default function App() {
|
||||
return (
|
||||
<Container className="py-5">
|
||||
<h1 className="mb-2">helexa.ai</h1>
|
||||
<p className="text-muted">Public beta — coming online.</p>
|
||||
</Container>
|
||||
);
|
||||
}
|
||||
16
helexa.ai/src/index.css
Normal file
16
helexa.ai/src/index.css
Normal file
@@ -0,0 +1,16 @@
|
||||
/* Minimal reset; the full theme (CSS custom properties, light/dark, accent
|
||||
* palette) is ported from the reference site in F1. */
|
||||
:root {
|
||||
color-scheme: light dark;
|
||||
}
|
||||
|
||||
* {
|
||||
box-sizing: border-box;
|
||||
}
|
||||
|
||||
html,
|
||||
body,
|
||||
#root {
|
||||
margin: 0;
|
||||
min-height: 100vh;
|
||||
}
|
||||
11
helexa.ai/src/main.tsx
Normal file
11
helexa.ai/src/main.tsx
Normal file
@@ -0,0 +1,11 @@
|
||||
import { StrictMode } from "react";
|
||||
import { createRoot } from "react-dom/client";
|
||||
import "bootstrap/dist/css/bootstrap.min.css";
|
||||
import "./index.css";
|
||||
import App from "./App.tsx";
|
||||
|
||||
createRoot(document.getElementById("root")!).render(
|
||||
<StrictMode>
|
||||
<App />
|
||||
</StrictMode>,
|
||||
);
|
||||
28
helexa.ai/tsconfig.app.json
Normal file
28
helexa.ai/tsconfig.app.json
Normal file
@@ -0,0 +1,28 @@
|
||||
{
|
||||
"compilerOptions": {
|
||||
"tsBuildInfoFile": "./node_modules/.tmp/tsconfig.app.tsbuildinfo",
|
||||
"target": "ES2022",
|
||||
"useDefineForClassFields": true,
|
||||
"lib": ["ES2022", "DOM", "DOM.Iterable"],
|
||||
"module": "ESNext",
|
||||
"types": ["vite/client"],
|
||||
"skipLibCheck": true,
|
||||
|
||||
/* Bundler mode */
|
||||
"moduleResolution": "bundler",
|
||||
"allowImportingTsExtensions": true,
|
||||
"verbatimModuleSyntax": true,
|
||||
"moduleDetection": "force",
|
||||
"noEmit": true,
|
||||
"jsx": "react-jsx",
|
||||
|
||||
/* Linting */
|
||||
"strict": true,
|
||||
"noUnusedLocals": true,
|
||||
"noUnusedParameters": true,
|
||||
"erasableSyntaxOnly": true,
|
||||
"noFallthroughCasesInSwitch": true,
|
||||
"noUncheckedSideEffectImports": true
|
||||
},
|
||||
"include": ["src"]
|
||||
}
|
||||
7
helexa.ai/tsconfig.json
Normal file
7
helexa.ai/tsconfig.json
Normal file
@@ -0,0 +1,7 @@
|
||||
{
|
||||
"files": [],
|
||||
"references": [
|
||||
{ "path": "./tsconfig.app.json" },
|
||||
{ "path": "./tsconfig.node.json" }
|
||||
]
|
||||
}
|
||||
26
helexa.ai/tsconfig.node.json
Normal file
26
helexa.ai/tsconfig.node.json
Normal file
@@ -0,0 +1,26 @@
|
||||
{
|
||||
"compilerOptions": {
|
||||
"tsBuildInfoFile": "./node_modules/.tmp/tsconfig.node.tsbuildinfo",
|
||||
"target": "ES2023",
|
||||
"lib": ["ES2023"],
|
||||
"module": "ESNext",
|
||||
"types": ["node"],
|
||||
"skipLibCheck": true,
|
||||
|
||||
/* Bundler mode */
|
||||
"moduleResolution": "bundler",
|
||||
"allowImportingTsExtensions": true,
|
||||
"verbatimModuleSyntax": true,
|
||||
"moduleDetection": "force",
|
||||
"noEmit": true,
|
||||
|
||||
/* Linting */
|
||||
"strict": true,
|
||||
"noUnusedLocals": true,
|
||||
"noUnusedParameters": true,
|
||||
"erasableSyntaxOnly": true,
|
||||
"noFallthroughCasesInSwitch": true,
|
||||
"noUncheckedSideEffectImports": true
|
||||
},
|
||||
"include": ["vite.config.ts"]
|
||||
}
|
||||
29
helexa.ai/vite.config.ts
Normal file
29
helexa.ai/vite.config.ts
Normal file
@@ -0,0 +1,29 @@
|
||||
import { defineConfig, loadEnv } from "vite";
|
||||
import react from "@vitejs/plugin-react-swc";
|
||||
|
||||
// During `vite dev`, proxy the mesh data-plane (helexa-router, OpenAI-
|
||||
// compatible) and the account control-plane (helexa-upstream) so the SPA
|
||||
// talks to them same-origin without CORS. Targets are overridable via env.
|
||||
// VITE_ROUTER_BASE_URL — helexa-router (default http://localhost:8088)
|
||||
// VITE_ACCOUNT_BASE_URL — helexa-upstream (default http://localhost:8090)
|
||||
export default defineConfig(({ mode }) => {
|
||||
const env = loadEnv(mode, process.cwd(), "VITE_");
|
||||
const router = env.VITE_ROUTER_BASE_URL || "http://localhost:8088";
|
||||
const account = env.VITE_ACCOUNT_BASE_URL || "http://localhost:8090";
|
||||
return {
|
||||
plugins: [react()],
|
||||
server: {
|
||||
proxy: {
|
||||
"/v1": { target: router, changeOrigin: true },
|
||||
"/health": { target: router, changeOrigin: true },
|
||||
// The frontend calls /api/*; helexa-upstream serves /web/v1/*.
|
||||
"/api": {
|
||||
target: account,
|
||||
changeOrigin: true,
|
||||
rewrite: (p) => p.replace(/^\/api/, "/web/v1"),
|
||||
},
|
||||
},
|
||||
},
|
||||
build: { outDir: "dist" },
|
||||
};
|
||||
});
|
||||
Reference in New Issue
Block a user