Compare commits

..

24 Commits

Author SHA1 Message Date
5fd7736abd feat(#72): router↔cortex topology poller (multi-operator capacity map)
All checks were successful
CI / Format (push) Successful in 39s
CI / CUDA type-check (push) Successful in 1m39s
CI / Clippy (push) Successful in 2m14s
CI / Test (push) Successful in 4m43s
CI / Build cortex SRPM (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
Builds the live topology the dispatcher (#73) will route on — the same
pattern as cortex↔neuron, one tier up.

- `poller.rs`: background loop polls each configured cortex's
  `GET /v1/models` (deserialised straight into the shared
  `cortex_core::node::CortexModelEntry`) and `GET /health`, on a
  configurable `poll_interval_secs` (default 10).
- `state.rs`: `RouterState` gains an `http_client`, `poll_interval`, and a
  `RwLock<HashMap<cortex_name, CortexTopology>>` pre-populated from config
  so the poller/handlers always find an entry. Per cortex: `reachable`,
  `consecutive_failures`, `last_poll`, healthy/total node counts, and a
  per-model `{loaded, feasible}` map (feasible = loaded OR cortex reports
  `feasible_on`, i.e. cold-loadable). `cortexes_serving(model)` returns the
  reachable cortexes that can serve a model — groundwork for #73.
- Debounce: a cortex flips unreachable only after
  `POLL_FAILURE_THRESHOLD` (3) consecutive failed polls, and recovers on
  the next good poll — mirrors cortex's neuron-poll debounce so a blip
  can't yank a whole operator out of routing. `/health` poll is
  best-effort and never flips reachability on its own.
- `lib.rs` spawns the poll loop in `run()`. `/health` now surfaces
  `cortexes.reachable`; `status` stays router-liveness (always `ok`).

Tests (`topology.rs`): live-map build (loaded vs catalogue-only feasible,
node counts, routing helper); unreachable→excluded→recovers across the
debounce threshold; dead endpoint never panics. Skeleton tests unchanged.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01F6o3ddqmYNh9kzdwq6eowh
2026-06-21 19:01:29 +03:00
03fd4960c3 Merge fix/71-shared-streaming-proxy: shared helexa-stream SSE proxy (#71)
All checks were successful
build-prerelease / Resolve version stamps + change detection (push) Successful in 35s
build-prerelease / Lint (fmt + clippy) (push) Successful in 2m24s
build-prerelease / Build neuron-blackwell (push) Successful in 1m40s
build-prerelease / Build neuron-ampere (push) Successful in 2m17s
build-prerelease / Build helexa-bench binary (push) Successful in 2m11s
build-prerelease / Build neuron-ada (push) Successful in 2m18s
build-prerelease / Build cortex binary (push) Successful in 2m25s
build-prerelease / Test (push) Successful in 4m53s
build-prerelease / Package helexa-bench RPM (push) Successful in 1m16s
build-prerelease / Package cortex RPM (push) Successful in 1m15s
build-prerelease / Package helexa-neuron-ampere RPM (push) Successful in 1m36s
build-prerelease / Package helexa-neuron-ada RPM (push) Successful in 1m39s
build-prerelease / Package helexa-neuron-blackwell RPM (push) Successful in 1m39s
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Successful in 54s
# Conflicts:
#	Cargo.lock
#	Cargo.toml
2026-06-21 18:15:12 +03:00
5ed6bc3390 Merge feat/70-router-skeleton: helexa-router binary skeleton (#70)
Some checks failed
build-prerelease / Test (push) Blocked by required conditions
build-prerelease / Resolve version stamps + change detection (push) Successful in 56s
build-prerelease / Build neuron-blackwell (push) Successful in 1m31s
build-prerelease / Build neuron-ada (push) Successful in 2m4s
build-prerelease / Build neuron-ampere (push) Successful in 2m11s
build-prerelease / Lint (fmt + clippy) (push) Successful in 2m29s
build-prerelease / Build cortex binary (push) Successful in 2m42s
build-prerelease / Build helexa-bench binary (push) Successful in 2m56s
build-prerelease / Package cortex RPM (push) Successful in 1m18s
build-prerelease / Package helexa-neuron-ampere RPM (push) Successful in 1m42s
build-prerelease / Package helexa-bench RPM (push) Successful in 1m47s
build-prerelease / Package helexa-neuron-ada RPM (push) Successful in 2m24s
build-prerelease / Package helexa-neuron-blackwell RPM (push) Successful in 2m25s
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Has been cancelled
2026-06-21 18:06:07 +03:00
cabec1d08a fix(#71): extract SSE streaming passthrough into shared helexa-stream
All checks were successful
CI / Format (push) Successful in 40s
CI / CUDA type-check (push) Successful in 1m38s
CI / Clippy (push) Successful in 2m15s
CI / Test (push) Successful in 6m28s
CI / Build cortex SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
The true-streaming SSE passthrough (Body::from_stream, no full-response
buffering, with chunk-observation hooks) was cortex-only. helexa-router
(#69) needs the same mechanism to proxy a chat-completions/messages
stream verbatim to a selected cortex. Extract it once.

New `crates/helexa-stream` owns the *mechanism* (kept HTTP-free
cortex-core untouched — it would have forced axum/reqwest/futures onto
every cortex-core consumer):

- `forward_streaming(client, url, headers, body, observer)` — POST and
  stream the response back chunk-for-chunk; status-agnostic, so a
  non-2xx (e.g. cortex 429) is passed through with status+headers
  intact (the #69 backpressure-passthrough requirement).
- `ChunkObserver` trait + `ObservedStream` wrapper — feeds each chunk to
  the observer, calls `finish` exactly once on clean end or on drop
  (client disconnect).
- `BodyTail` (bounded tail accumulator) + `last_count_for` (trailing
  OpenAI `usage` extraction) — the reusable pieces an observer uses.

cortex keeps its *policy*: `proxy.rs` now supplies a `CortexMetrics`
observer (per-request token metrics + per-principal reservation settle),
its logging contract, and the error envelope, driving the shared
mechanism. `proxy::last_count_for` is re-exported so `handlers`/
`anthropic_sse` call sites are unchanged. No behaviour change — the
existing cortex `streaming.rs` tests pass as-is.

helexa-stream tests prove chunk-for-chunk incremental delivery, observer
finish-once, usage extraction, and non-2xx passthrough.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01F6o3ddqmYNh9kzdwq6eowh
2026-06-21 18:05:35 +03:00
881fc85a4c feat(#70): helexa-router binary skeleton — plaintext axum server
All checks were successful
CI / Format (push) Successful in 38s
CI / CUDA type-check (push) Successful in 1m39s
CI / Clippy (push) Successful in 2m19s
CI / Test (push) Successful in 4m40s
CI / Build cortex SRPM (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
Foundation for epic #69 (public multi-operator ingress proxy). New
`crates/helexa-router` workspace binary: a plaintext axum server that
reuses cortex-core types and serves the two stub endpoints the rest of
#69 builds on.

- `[router] listen` + `[[cortexes]]` config via figment + `HELEXA_ROUTER_`
  env overrides, matching the cortex/neuron convention.
- `GET /health` reports the configured downstream cortex count.
- `GET /v1/models` returns an empty OpenAI list (real cross-operator
  aggregation is #75).
- No inbound TLS listener (edge nginx terminates client TLS per #69's
  posture); no auth layer — the router forwards the client bearer to
  cortex and holds zero entitlement logic (#47 stays additive).
- 3 tests: both endpoints over a real ephemeral-port server, plus
  TOML+env config load.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01F6o3ddqmYNh9kzdwq6eowh
2026-06-21 17:52:28 +03:00
b2ed20b55a docs(CLAUDE.md): document the branch → CI → merge-on-green workflow
All checks were successful
build-prerelease / Resolve version stamps + change detection (push) Successful in 31s
build-prerelease / Lint (fmt + clippy) (push) Has been skipped
build-prerelease / Test (push) Has been skipped
build-prerelease / Build cortex binary (push) Has been skipped
build-prerelease / Package cortex RPM (push) Has been skipped
build-prerelease / Build helexa-bench binary (push) Has been skipped
build-prerelease / Package helexa-bench RPM (push) Has been skipped
build-prerelease / Build neuron-blackwell (push) Has been skipped
build-prerelease / Build neuron-ampere (push) Has been skipped
build-prerelease / Build neuron-ada (push) Has been skipped
build-prerelease / Package helexa-neuron-ada RPM (push) Has been skipped
build-prerelease / Package helexa-neuron-ampere RPM (push) Has been skipped
build-prerelease / Package helexa-neuron-blackwell RPM (push) Has been skipped
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Has been skipped
Capture the development loop in-repo so it's available to any session, not
just via personal agent memory: feature branch per change; local CI triad is
CPU-only so the branch CI's CUDA type-check is the real gate for neuron/TP
changes; push on local-green and background-watch; merge when the four
validation jobs are green (not the SRPM/COPR deploy jobs); docs-only changes
can go straight to main. Notes the core.sshCommand key-pinning gotcha.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01M5aNfNzS2fSZ5wnMeSQ9Wg
2026-06-21 15:05:46 +03:00
bee27e9b9c Merge fix/68-cost-schema-wire-contract: pin the /v1/models cost wire contract
All checks were successful
build-prerelease / Resolve version stamps + change detection (push) Successful in 31s
build-prerelease / Lint (fmt + clippy) (push) Successful in 2m20s
build-prerelease / Build neuron-blackwell (push) Successful in 1m42s
build-prerelease / Build helexa-bench binary (push) Successful in 2m2s
build-prerelease / Build neuron-ampere (push) Successful in 2m17s
build-prerelease / Build neuron-ada (push) Successful in 2m19s
build-prerelease / Build cortex binary (push) Successful in 2m31s
build-prerelease / Test (push) Successful in 4m49s
build-prerelease / Package helexa-bench RPM (push) Successful in 1m22s
build-prerelease / Package cortex RPM (push) Successful in 1m20s
build-prerelease / Package helexa-neuron-ampere RPM (push) Successful in 1m44s
build-prerelease / Package helexa-neuron-ada RPM (push) Successful in 1m45s
build-prerelease / Package helexa-neuron-blackwell RPM (push) Successful in 1m47s
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Successful in 54s
Closes #68. Documents ModelCost as the source-of-truth pricing field
(USD per 1M tokens, JSON numbers — models.dev/opencode shape), defines the
absent-vs-0.0 distinction (not-priced vs intentionally-free), adds a wire
test locking it, and documents cost.* in models.example.toml. The cost code
path already existed; this pins the contract. Branch CI green.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01M5aNfNzS2fSZ5wnMeSQ9Wg
2026-06-20 12:08:23 +03:00
87d9c291ce fix(#68): pin the /v1/models cost wire contract — units + absent-vs-zero
All checks were successful
CI / Format (push) Successful in 44s
CI / CUDA type-check (push) Successful in 1m37s
CI / Clippy (push) Successful in 2m15s
CI / Test (push) Successful in 4m51s
CI / Build cortex SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
The cost code path already exists (cortex list_models populates
cost: profile.cost from the catalogue; aliases inherit it), so opencode's
$0.00 is a config gap (no cost in the live models.toml), not missing
plumbing. What was missing is the *contract*: units pinned against a wire
test, and a defined meaning for "free".

- Document ModelCost as the load-bearing source of truth: USD per 1,000,000
  tokens as JSON numbers (models.dev/opencode shape) — NOT per-token, NOT
  decimal strings (OpenRouter's pricing shape, which helexa deliberately
  does not emit). Define the absent-vs-zero distinction: cost omitted = "not
  priced / unknown"; cost present with 0.0 = "intentionally free". Note the
  advertised rate must equal what metering (#51) / reconciliation (#58/#59)
  bill against — today both read this catalogue value.
- New wire test (model_cost.rs): a priced model with cache tiers flows
  through as per-million numbers; an explicit-0.0 free model keeps its cost
  block with cache tiers omitted; an unpriced model omits `cost` entirely.
- models.example.toml: document cost.* in the field reference and show all
  three cases (priced-free explicit 0.0 vs the unpriced Qwen3-8B with no
  cost block).

Decisions recorded on #68: source of truth = operator models.toml for now
(marketplace clearing house #59 later, same value); no OpenRouter-style
`pricing` (opencode/models.dev alignment is sufficient); end-to-end
non-zero $ spent needs operators to populate cost in the live catalogue.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01M5aNfNzS2fSZ5wnMeSQ9Wg
2026-06-20 12:02:03 +03:00
d4742467e0 Merge fix/65-text-prefill-vram-backstop: request-time length-aware VRAM backstop for text prefill
All checks were successful
build-prerelease / Resolve version stamps + change detection (push) Successful in 32s
build-prerelease / Build cortex binary (push) Has been skipped
build-prerelease / Build helexa-bench binary (push) Has been skipped
build-prerelease / Package cortex RPM (push) Has been skipped
build-prerelease / Package helexa-bench RPM (push) Has been skipped
build-prerelease / Build neuron-blackwell (push) Successful in 1m41s
build-prerelease / Build neuron-ampere (push) Successful in 2m16s
build-prerelease / Build neuron-ada (push) Successful in 2m17s
build-prerelease / Lint (fmt + clippy) (push) Successful in 2m38s
build-prerelease / Package helexa-neuron-ada RPM (push) Successful in 1m39s
build-prerelease / Test (push) Successful in 4m59s
build-prerelease / Package helexa-neuron-ampere RPM (push) Successful in 1m36s
build-prerelease / Package helexa-neuron-blackwell RPM (push) Successful in 1m41s
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Successful in 51s
Closes #65. Gives the text prefill path a request-time, length-aware
VRAM guard (reusing #67's ContextProfile KV cost against current free
VRAM), closing the poll-vs-request snapshot staleness gap and the
vision/text asymmetry. Branch CI green (fmt, clippy, test, CUDA type-check).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01M5aNfNzS2fSZ5wnMeSQ9Wg
2026-06-20 11:53:33 +03:00
e7f7e376fc fix(#65): request-time length-aware VRAM backstop for text prefill
All checks were successful
CI / Format (push) Successful in 33s
CI / CUDA type-check (push) Successful in 1m41s
CI / Clippy (push) Successful in 2m24s
CI / Test (push) Successful in 4m41s
CI / Build cortex SRPM (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
Close the poll-vs-request snapshot gap #67 left open. The text prefill
guard (validate_request) only checked the static min_free_vram floor;
the derived input cap (effective_prompt_cap) is computed at /models poll
time from the tightest card's free VRAM *then*. If free VRAM drops
between that poll and the request — a co-resident model loads, a
concurrent prefill grows its KV — a prompt at-or-below the now-stale cap
clears the floor yet no longer fits, OOMing mid-prefill and poisoning the
device context (the 2026-05-26 beast incident #47 exists to eliminate).

validate_request now re-runs #67's length×KV-vs-VRAM physics against
request-time free VRAM, reusing the model's ContextProfile
(kv_bytes_per_token_per_card, full-attention-layer-only, TP-sharded)
rather than re-deriving the cost. Footprint = KV(prompt + output_reserve)
+ activation_headroom + static floor, all per card and commensurable with
the tightest-card free VRAM on both single-GPU and TP loads. Degenerate
zero-KV / no-profile models ride the existing floor check, mirroring
derive_limit's VRAM-ceiling fallback; CPU loads (vram_free_mb == 0) skip
all VRAM checks unchanged.

This closes the vision/text asymmetry: the text path now has the
live-VRAM guard validate_vision_prefill already gave the vision path.

5 unit tests incl. the acceptance staleness test: a cap derived against
ample free VRAM, applied at request time against tightened VRAM, rejects
a prompt sized at the stale cap with a clean InsufficientVram (503)
instead of an OOM. Threaded context_limit_cfg into chat_completion_tp_inner
(spawned, no &self) and used &self.context_limit_cfg at the three method
call sites.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01M5aNfNzS2fSZ5wnMeSQ9Wg
2026-06-20 11:45:53 +03:00
3b9a6e37f6 Merge fix/cortex-poll-debounce-retryable: poll debounce + retryable 503 for feasible-but-unhealthy node
All checks were successful
build-prerelease / Resolve version stamps + change detection (push) Successful in 58s
build-prerelease / Build neuron-blackwell (push) Successful in 1m31s
build-prerelease / Lint (fmt + clippy) (push) Successful in 2m31s
build-prerelease / Build cortex binary (push) Successful in 3m0s
build-prerelease / Test (push) Successful in 5m3s
build-prerelease / Package cortex RPM (push) Successful in 1m20s
build-prerelease / Build neuron-ada (push) Successful in 2m2s
build-prerelease / Build neuron-ampere (push) Successful in 2m12s
build-prerelease / Build helexa-bench binary (push) Successful in 2m9s
build-prerelease / Package helexa-neuron-ada RPM (push) Successful in 1m39s
build-prerelease / Package helexa-bench RPM (push) Successful in 1m20s
build-prerelease / Package helexa-neuron-ampere RPM (push) Successful in 1m41s
build-prerelease / Package helexa-neuron-blackwell RPM (push) Successful in 1m43s
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Successful in 53s
2026-06-18 12:46:30 +03:00
526b662c5e fix(cortex): poll-failure debounce + retryable 503 for feasible-but-unhealthy node
All checks were successful
CI / Format (push) Successful in 44s
CI / CUDA type-check (push) Successful in 1m31s
CI / Clippy (push) Successful in 2m13s
CI / Test (push) Successful in 5m9s
CI / Build cortex SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
Defense-in-depth for the agent0 NoFeasibleNeuron storm (root cause fixed in
neuron). Two cortex resilience gaps this incident exposed:

1. Brittle health flip: the poller marked a node unhealthy on a SINGLE missed
   /models poll, instantly yanking the node and all its models from routing.
   A busy neuron briefly slow to answer shouldn't be declared dead. Now
   debounced: NodeState.consecutive_poll_failures must reach
   POLL_FAILURE_THRESHOLD (3) before the node flips unhealthy (~20s at the 10s
   poll interval); any successful poll resets it. A never-healthy node stays
   unhealthy (the counter only protects an already-healthy node from blips).

2. Transient surfaced as permanent: when a catalogued model's only feasible
   neuron is momentarily unhealthy, the router returned 404 NoFeasibleNeuron —
   which litellm/clients treat as non-retryable, so agent0 hard-failed.
   pick_feasible_neuron now distinguishes "a feasible node exists but is
   unhealthy right now" → new RouteError::FeasibleNodeUnhealthy (503 +
   Retry-After: 3, retryable) from "no node could ever satisfy the topology" →
   404 NoFeasibleNeuron (permanent). Mirrors the beast case exactly: healthy
   1-GPU nodes + an unhealthy 2-GPU node → retry, don't fail.

Tests: poller test updated to assert debounce (1 miss keeps healthy, 3 flip);
new feasibility_routing tests cover transient-503 vs permanent-404. Local
fmt/clippy/test green.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-06-18 12:39:18 +03:00
db7e373b90 fix(neuron): decouple GET /models from the inference worker (control-plane starvation)
Some checks failed
CI / CUDA type-check (push) Successful in 1m40s
CI / Format (push) Successful in 36s
CI / Clippy (push) Successful in 2m20s
CI / Test (push) Successful in 4m28s
CI / Build cortex SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
build-prerelease / Test (push) Blocked by required conditions
build-prerelease / Package helexa-neuron-ada RPM (push) Blocked by required conditions
build-prerelease / Package helexa-neuron-ampere RPM (push) Blocked by required conditions
build-prerelease / Resolve version stamps + change detection (push) Successful in 36s
build-prerelease / Build cortex binary (push) Has been skipped
build-prerelease / Package cortex RPM (push) Has been skipped
build-prerelease / Build helexa-bench binary (push) Has been skipped
build-prerelease / Package helexa-bench RPM (push) Has been skipped
build-prerelease / Build neuron-blackwell (push) Successful in 1m43s
build-prerelease / Build neuron-ampere (push) Successful in 2m20s
build-prerelease / Build neuron-ada (push) Successful in 2m21s
build-prerelease / Lint (fmt + clippy) (push) Successful in 2m34s
build-prerelease / Package helexa-neuron-blackwell RPM (push) Has been cancelled
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Has been cancelled
Root cause of the agent0 `NoFeasibleNeuron` 404 storm: `GET /models` →
`LoadedHandle::derived_limit` (#67) queried free VRAM *synchronously through
the per-device worker thread* on every poll. During inference that worker is
saturated serially processing forward jobs, so the VRAM query queued behind
them and `/models` blocked for seconds. cortex's poller timed out on `/models`,
marked the (sole-feasible) node unhealthy, and the model fell out of routing →
404. Confirmed live: under load, `/version` and `/health` stayed ~4ms while
`/models` hit the 5s timeout.

Fix — the HTTP control plane never touches the inference worker:
- LoadedModel / TpLoadedModel gain `last_free_mb: AtomicU64`, a cached free-VRAM
  reading.
- `derived_limit` is now sync and reads `last_free_mb` instead of awaiting a
  worker query — so `/models` is a pure cache read regardless of inference load.
- The cache is refreshed off the request path: seeded at load (worker idle),
  then by a background `vram_cache_refresh_loop` every 5s. Single-GPU caches the
  device's free VRAM; TP caches the tightest free across ranks — the exact
  values `derived_limit` used before, just no longer on the request path. A
  transient `0` (worker gone/poisoned) never clobbers a good cached value.
- The request-path live VRAM check in `validate_request` is unchanged, so the
  real prefill OOM guard still uses fresh readings.

226 neuron unit tests pass; non-CUDA build + fmt + clippy green. CUDA/TP paths
validated by branch CI; live acceptance = `/models` stays responsive under
concurrent inference (re-run of the repro).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-06-18 12:32:15 +03:00
5c1623a817 fix(#49): allow-anonymous mode must ignore unrecognized keys, not 401
All checks were successful
CI / Format (push) Successful in 38s
CI / CUDA type-check (push) Successful in 1m41s
CI / Clippy (push) Successful in 2m20s
CI / Test (push) Successful in 4m37s
CI / Build cortex SRPM (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
build-prerelease / Resolve version stamps + change detection (push) Successful in 32s
build-prerelease / Lint (fmt + clippy) (push) Successful in 2m21s
build-prerelease / Build neuron-blackwell (push) Has been skipped
build-prerelease / Build neuron-ampere (push) Has been skipped
build-prerelease / Build neuron-ada (push) Has been skipped
build-prerelease / Package helexa-neuron-ada RPM (push) Has been skipped
build-prerelease / Package helexa-neuron-ampere RPM (push) Has been skipped
build-prerelease / Package helexa-neuron-blackwell RPM (push) Has been skipped
build-prerelease / Build helexa-bench binary (push) Has been skipped
build-prerelease / Package helexa-bench RPM (push) Has been skipped
build-prerelease / Build cortex binary (push) Successful in 2m35s
build-prerelease / Test (push) Successful in 4m39s
build-prerelease / Package cortex RPM (push) Successful in 1m25s
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Successful in 52s
Regression from #49: the auth middleware rejected ANY present-but-
unresolvable bearer token with 401 invalid_api_key, even when
require_auth=false. But OpenAI-compatible clients (opencode, Open WebUI,
Agent Zero, litellm) send a placeholder bearer by default — so enabling
the build broke every existing client even though the operator never
opted into auth. Pre-#49 the bearer was never inspected at all.

Fix: in allow-anonymous mode (require_auth=false, the default) an
unrecognized key is now ignored and the request is served anonymously,
restoring pre-#49 behaviour. A bad key only 401s when require_auth=true.
A valid key is still resolved + metered in both modes.

Test renamed/split: unrecognized_key_is_ignored_when_auth_not_required
(now 200, served anonymously) + invalid_key_is_401_when_auth_required.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-06-17 21:40:34 +03:00
3b60dd7a31 Merge #56 (phase 3): fail-fast prompt pre-validation + advisory hints
All checks were successful
build-prerelease / Resolve version stamps + change detection (push) Successful in 36s
build-prerelease / Lint (fmt + clippy) (push) Successful in 2m44s
build-prerelease / Build neuron-blackwell (push) Successful in 1m42s
build-prerelease / Build helexa-bench binary (push) Successful in 2m2s
build-prerelease / Build neuron-ada (push) Successful in 2m19s
build-prerelease / Build neuron-ampere (push) Successful in 2m20s
build-prerelease / Build cortex binary (push) Successful in 2m28s
build-prerelease / Test (push) Successful in 4m45s
build-prerelease / Package helexa-bench RPM (push) Successful in 1m24s
build-prerelease / Package cortex RPM (push) Successful in 1m21s
build-prerelease / Package helexa-neuron-ada RPM (push) Successful in 1m42s
build-prerelease / Package helexa-neuron-ampere RPM (push) Successful in 1m41s
build-prerelease / Package helexa-neuron-blackwell RPM (push) Successful in 1m44s
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Successful in 55s
2026-06-17 20:57:55 +03:00
4feaaf1cfb Merge #55 (phase 2d): cortex load-aware routing across replicas
Some checks failed
build-prerelease / Test (push) Blocked by required conditions
build-prerelease / Package cortex RPM (push) Blocked by required conditions
build-prerelease / Package helexa-neuron-ampere RPM (push) Blocked by required conditions
build-prerelease / Package helexa-neuron-blackwell RPM (push) Blocked by required conditions
build-prerelease / Resolve version stamps + change detection (push) Successful in 34s
build-prerelease / Build neuron-blackwell (push) Successful in 1m39s
build-prerelease / Build neuron-ampere (push) Successful in 2m20s
build-prerelease / Build neuron-ada (push) Successful in 2m21s
build-prerelease / Build cortex binary (push) Successful in 2m52s
build-prerelease / Lint (fmt + clippy) (push) Successful in 3m1s
build-prerelease / Package helexa-neuron-ada RPM (push) Successful in 2m9s
build-prerelease / Build helexa-bench binary (push) Has been cancelled
build-prerelease / Package helexa-bench RPM (push) Has been cancelled
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Has been cancelled
2026-06-17 20:51:26 +03:00
057bc71e80 feat(#47 #56 phase 3): fail-fast prompt pre-validation + advisory hints
All checks were successful
CI / Format (push) Successful in 29s
CI / CUDA type-check (push) Successful in 1m37s
CI / Clippy (push) Successful in 2m35s
CI / Test (push) Successful in 5m4s
CI / Build cortex SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
Stage 3 (DX): A0 burned an hour then failed deep in litellm with
prompt_too_long (35544 > 32768). cortex knows each model's real context
window (#62/#67) and can pre-empt that at the edge.

- Pre-validate the prompt against the model's advertised limit.context
  before dispatch (in proxy_with_metrics, covering chat/completions/
  responses). Over → 400 context_length_exceeded in the #60 envelope — the
  same shape neuron emits on overflow, just earlier and without burning a
  cold-load/queue slot. cortex has no tokenizer, so estimate_prompt_tokens
  under-counts (~4 chars/token over message text); neuron stays the exact
  wall and we only catch gross overages. Skipped when no limit is known.
- Advisory X-Helexa-Advice header: fingerprints User-Agent
  (litellm / Agent-Zero / Zed) and attaches client-specific guidance.
  Strictly advisory — header only, never in the error envelope, behaviour
  never depends on it; unknown clients get nothing.

3 integration tests: over-long prompt → 400 context_length_exceeded with
the advice header, refused before neuron is hit; within-context passes
through; unknown client gets a clean 400 with no advice header. cortex-side
(no CUDA); local fmt/clippy/test green.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-06-17 20:50:38 +03:00
dd31c3cd49 feat(#47 #55 phase 2d): cortex load-aware routing across replicas
All checks were successful
CI / Format (push) Successful in 39s
CI / CUDA type-check (push) Successful in 1m50s
CI / Clippy (push) Successful in 2m24s
CI / Test (push) Successful in 4m51s
CI / Build cortex SRPM (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
Stage 2 completes: when a model is loaded on more than one healthy neuron,
the router picks the least-busy replica instead of always taking the first,
and neuron backpressure propagates to the client intact.

- NodeState.model_load: per-model admission load (in_flight + queue_depth),
  stashed by the poller from neuron's /health (#53/#2b).
- router::resolve collects all loaded replicas and picks the one with the
  lowest in_flight+queue_depth (ties break by node name for determinism),
  replacing the previous first-match-wins.
- Backpressure passthrough: the existing streaming proxy already forwards
  the upstream status + all headers verbatim, so a neuron 503/429 +
  Retry-After + #60 envelope reaches the client unmodified — now covered by
  a regression test so a future change can't silently unwrap it.

Tests (tests/load_routing.rs): routes to the idle replica and follows the
lighter load when it flips; ties break by name; a saturated neuron's 503 +
Retry-After + envelope propagates through the gateway intact. All
cortex-side (no CUDA); local fmt/clippy/test green.

Retry-route-to-another-replica-on-backpressure (the issue's stretch goal)
is deferred — least-busy spread + honest passthrough is the substantive win.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-06-17 20:45:50 +03:00
c83f1eb98c feat(#47 #54 phase 2c): neuron per-principal in-flight cap (fair-share)
Some checks failed
CI / Format (push) Successful in 37s
CI / CUDA type-check (push) Successful in 1m37s
CI / Clippy (push) Successful in 2m13s
CI / Test (push) Successful in 4m50s
CI / Build cortex SRPM (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
build-prerelease / Test (push) Blocked by required conditions
build-prerelease / Build neuron-ampere (push) Blocked by required conditions
build-prerelease / Build neuron-ada (push) Blocked by required conditions
build-prerelease / Resolve version stamps + change detection (push) Successful in 37s
build-prerelease / Build neuron-blackwell (push) Successful in 1m28s
build-prerelease / Lint (fmt + clippy) (push) Successful in 3m0s
build-prerelease / Build cortex binary (push) Has been skipped
build-prerelease / Build helexa-bench binary (push) Has been skipped
build-prerelease / Package cortex RPM (push) Has been skipped
build-prerelease / Package helexa-bench RPM (push) Has been skipped
build-prerelease / Package helexa-neuron-ada RPM (push) Has been cancelled
build-prerelease / Package helexa-neuron-ampere RPM (push) Has been cancelled
build-prerelease / Package helexa-neuron-blackwell RPM (push) Has been cancelled
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Has been cancelled
Budget caps total spend over time (#52); this caps instantaneous
starvation so one principal's burst can't monopolize a model while others
wait.

- AdmissionController gains per-principal accounting (moved from a lone
  atomic to a Mutex<AdmissionState> holding the overall pending count + a
  per-principal map). enter(principal) now also fast-rejects when a
  principal already has max_per_principal requests in flight/queued →
  AdmissionRejection::PrincipalCap. Anonymous (None) requests are exempt.
- Config [harness.candle.admission].max_per_principal (default 2 = one
  running + one queued; 0 disables). A bursting principal's overflow is
  refused while a different principal still gets a queue slot.
- The principal (account/key) is reconstructed on the neuron side from the
  x-helexa-account-id/key-id headers cortex stamps (#49) — trusted over
  WireGuard, never from the request body — and threaded explicitly through
  all inference entry points (chat_completion, *_stream(_with),
  responses_stream, and the TP variants) to the admission gate.
- InferenceError::PerPrincipalLimit → 429 rate_limit_exceeded + Retry-After
  (distinct from load-shedding's 503 Overloaded); opencode/AI SDK self-pace.

Tests: fair-share unit test (A floods → A's 2nd is PrincipalCap, B still
queues + is served) + the existing admission tests adapted to enter(None).
Non-CUDA build green locally; TP entry points (cuda-gated) validated by CI.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-06-17 20:40:25 +03:00
a60c9f1075 feat(#47 #53 phase 2b): expose per-model admission load in GET /health
All checks were successful
CI / Format (push) Successful in 30s
CI / CUDA type-check (push) Successful in 1m30s
CI / Clippy (push) Successful in 2m18s
CI / Test (push) Successful in 4m17s
CI / Build cortex SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
build-prerelease / Resolve version stamps + change detection (push) Successful in 33s
build-prerelease / Build neuron-blackwell (push) Successful in 1m42s
build-prerelease / Build neuron-ampere (push) Successful in 2m18s
build-prerelease / Build neuron-ada (push) Successful in 2m19s
build-prerelease / Build helexa-bench binary (push) Successful in 2m18s
build-prerelease / Lint (fmt + clippy) (push) Successful in 2m27s
build-prerelease / Build cortex binary (push) Successful in 2m45s
build-prerelease / Package helexa-neuron-ada RPM (push) Successful in 2m2s
build-prerelease / Test (push) Successful in 4m50s
build-prerelease / Package helexa-bench RPM (push) Successful in 1m18s
build-prerelease / Package cortex RPM (push) Successful in 1m22s
build-prerelease / Package helexa-neuron-ampere RPM (push) Successful in 1m37s
build-prerelease / Package helexa-neuron-blackwell RPM (push) Successful in 1m43s
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Successful in 56s
Completes #53: the bounded scheduler's lock-free counters are now visible
to the fleet, which is what cortex's load-aware router (#55) consumes to
spread traffic across replicas and propagate honest backpressure.

- cortex-core::discovery: HealthResponse gains `models: Vec<ModelLoad>`
  (#[serde(default)] — back-compatible; older gateways/neurons interop).
  ModelLoad { id, in_flight, queue_depth }.
- LoadedHandle::load() → (in_flight, queue_depth), lock-free for both
  single-GPU and TP; CandleHarness::load_snapshot() enumerates resident
  models; the /health handler overlays it from the candle harness.

Tests: /health always exposes a models array (api integration test); a
pre-#53 payload without `models` still deserializes, and ModelLoad
round-trips (cortex-core serde tests). Local fmt/clippy/test green.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-06-17 20:13:07 +03:00
b2bd86bfa5 feat(#47 #53 phase 2a): neuron admission control — bounded queue + backpressure
All checks were successful
CI / Format (push) Successful in 41s
CI / CUDA type-check (push) Successful in 1m40s
CI / Clippy (push) Successful in 2m18s
CI / Test (push) Successful in 4m53s
CI / Build cortex SRPM (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
build-prerelease / Resolve version stamps + change detection (push) Successful in 32s
build-prerelease / Build cortex binary (push) Has been skipped
build-prerelease / Build helexa-bench binary (push) Has been skipped
build-prerelease / Package cortex RPM (push) Has been skipped
build-prerelease / Package helexa-bench RPM (push) Has been skipped
build-prerelease / Build neuron-blackwell (push) Successful in 1m43s
build-prerelease / Build neuron-ampere (push) Successful in 2m18s
build-prerelease / Build neuron-ada (push) Successful in 2m19s
build-prerelease / Lint (fmt + clippy) (push) Successful in 2m29s
build-prerelease / Package helexa-neuron-ada RPM (push) Successful in 1m46s
build-prerelease / Test (push) Successful in 4m48s
build-prerelease / Package helexa-neuron-ampere RPM (push) Successful in 1m49s
build-prerelease / Package helexa-neuron-blackwell RPM (push) Successful in 1m53s
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Successful in 1m7s
Replaces the per-model unbounded, untimed FIFO of inference-lock waiters
(a busy model made new requests hang ~300s until the client gave up with
an opaque error) with an explicit bounded scheduler.

- harness::admission::AdmissionController: batch-1 scheduler — max_in_flight
  running (1) + a bounded queue (max_queue_depth) with a max_wait. enter()
  fast-rejects when the queue is full (QueueFull) or the wait elapses
  (Timeout); the returned AdmissionPermit is held for the request and frees
  both slots on drop. Pure async (no CUDA), lock-free in_flight/queue_depth
  counters for future /health reporting. Configurable via
  [harness.candle.admission] (max_in_flight=1, max_queue_depth=8,
  max_wait_secs=30).
- Gated at all four inference entry points before the inference_lock/pool
  lock: single-GPU non-streaming + streaming, TP non-streaming + streaming.
  The streaming paths acquire the permit before opening the SSE (so a
  rejection is a clean error, not a half-open stream) and move it into the
  inference task.
- InferenceError::Overloaded { retry_after_secs } → 503 rate_limit_exceeded
  + Retry-After via the #60/#63 envelope: a fast, retryable "busy" signal
  opencode/AI SDK back off on, not a stall.

Scope: this branch is the admission *core* (the hang→backpressure fix).
Exposing in_flight/queue_depth in GET /health (consumed by cortex
load-aware routing #55) is the next focused branch under #53.

4 unit tests (admit/report load, queue-full reject, wait-timeout reject)
+ Overloaded envelope mapping test. Non-CUDA build green locally; the
CUDA + TP sites are validated by branch CI.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-06-17 20:03:07 +03:00
cdf87284af feat(#47 phase 1d): budget enforcement — hard caps, reserve→settle, 429
All checks were successful
CI / Format (push) Successful in 1s
CI / CUDA type-check (push) Successful in 1m40s
CI / Clippy (push) Successful in 2m40s
CI / Test (push) Successful in 6m23s
CI / Build cortex SRPM (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
build-prerelease / Resolve version stamps + change detection (push) Successful in 34s
build-prerelease / Lint (fmt + clippy) (push) Successful in 2m19s
build-prerelease / Test (push) Successful in 4m28s
build-prerelease / Build neuron-blackwell (push) Has been skipped
build-prerelease / Build neuron-ampere (push) Has been skipped
build-prerelease / Build neuron-ada (push) Has been skipped
build-prerelease / Package helexa-neuron-ada RPM (push) Has been skipped
build-prerelease / Package helexa-neuron-ampere RPM (push) Has been skipped
build-prerelease / Package helexa-neuron-blackwell RPM (push) Has been skipped
build-prerelease / Build helexa-bench binary (push) Has been skipped
build-prerelease / Package helexa-bench RPM (push) Has been skipped
build-prerelease / Build cortex binary (push) Successful in 2m27s
build-prerelease / Package cortex RPM (push) Successful in 1m23s
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Successful in 50s
Stage 1 complete: the A0 seatbelt (#52). Flips the metering-only reserve(0)
from #51 to the request's real upper-bound cost and refuses over-cap
requests *before* neuron is hit.

- metering::reservation_estimate: prompt estimate (~4 chars/token over the
  body — cortex has no tokenizer, so a conservative over-estimate; neuron
  stays the exact context wall) + max output. Max output comes from
  max_completion_tokens / legacy max_tokens, else the model's advertised
  limit.output (#62), else FALLBACK_MAX_OUTPUT. Over-reserving is safe —
  settle reconciles to actual.
- metering::reserve_or_reject: reserve the estimate; on BudgetError map to
  the #63 envelope and the caller refuses before dispatch — rolling window →
  429 rate_limit_exceeded + Retry-After (until reset); hard balance → 429
  insufficient_quota (no Retry-After). Never 402.
- Wired into both the OpenAI proxy path (proxy_with_metrics) and the
  Anthropic path (estimate from the translated body). advertised_output_limit
  reads the loaded model's limit.output from fleet state.
- Reservation prevents overshoot under concurrency: a successful reserve
  gates on spent+reserved+estimate ≤ cap, and settle records actual ≤
  reserved, so spend can never exceed the hard cap.

4 integration tests with a hit-counting mock neuron: balance over-cap →
429 insufficient_quota (no Retry-After, not dispatched); rolling over-cap →
429 rate_limit_exceeded + Retry-After (not dispatched); within-cap served;
**A0 repro** — a capped key's 20-request fan-out drains the cap, then is
refused, neuron only saw the served ones, and spend never exceeds the cap.
Plus 5 metering unit tests. Local fmt/clippy/test all green.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-06-17 19:35:04 +03:00
4f16b8c541 feat(#47 phase 1c): per-request token metering + spend ledger
All checks were successful
CI / Format (push) Successful in 40s
CI / CUDA type-check (push) Successful in 1m41s
CI / Clippy (push) Successful in 2m15s
CI / Test (push) Successful in 4m28s
CI / Build cortex SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
build-prerelease / Resolve version stamps + change detection (push) Successful in 32s
build-prerelease / Build neuron-blackwell (push) Has been skipped
build-prerelease / Build neuron-ampere (push) Has been skipped
build-prerelease / Build neuron-ada (push) Has been skipped
build-prerelease / Package helexa-neuron-ada RPM (push) Has been skipped
build-prerelease / Package helexa-neuron-ampere RPM (push) Has been skipped
build-prerelease / Package helexa-neuron-blackwell RPM (push) Has been skipped
build-prerelease / Lint (fmt + clippy) (push) Successful in 2m30s
build-prerelease / Build cortex binary (push) Successful in 2m49s
build-prerelease / Package cortex RPM (push) Successful in 1m24s
build-prerelease / Test (push) Successful in 5m59s
build-prerelease / Build helexa-bench binary (push) Has been skipped
build-prerelease / Package helexa-bench RPM (push) Has been skipped
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Successful in 49s
Stage 1 accounting (#51): capture real per-request usage and feed it to
the spend ledger + per-principal metrics. Establishes the reserve→settle
lifecycle that budget enforcement (#52) will tighten.

- cortex-gateway::metering: ReservationGuard makes reservation leaks
  impossible — settle() records actual spend + releases the remainder;
  dropping an un-settled guard releases the whole reservation, so any
  early return / error / dropped stream resolves it. UsageSink is the
  completion hook; principal_from_headers reconstructs the principal from
  the middleware-stamped headers (uniform across all proxy paths, no
  handler-signature churn); record_spend emits per-principal counters.
- proxy::TokenMetrics gains an optional usage_sink, invoked exactly once
  in finish() with the observed (prompt, completion) — restructured so it
  always runs (even when no body/usage arrived → settle 0 → release),
  while preserving the existing per-model metric emissions unchanged.
- All proxy paths metered: chat/completions/responses via
  proxy_with_metrics (reserve 0 → forward_request → settle in finish);
  Anthropic non-streaming settles from the buffered body; Anthropic
  streaming (anthropic_sse) now scans the upstream frames for the usage
  object (#48) — it captured none before — and settles at pump end.
- This phase reserves 0 tokens (metering only, no enforcement); #52 flips
  the reserved amount to prompt+max_output and surfaces BudgetError. The
  settle/release plumbing is identical, so that change is localized.
- New Prometheus counters: cortex_spend_tokens_total (+ prompt/completion
  splits), labelled by account/key.

2 integration tests: cumulative per-key spend after N requests with
reservations settled to zero outstanding; anonymous requests record no
spend. Local fmt/clippy/test all green.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-06-17 19:29:51 +03:00
486d7e9a8f feat(#47 phase 1b): API-key auth + principal resolution
All checks were successful
CI / Format (push) Successful in 36s
CI / CUDA type-check (push) Successful in 1m51s
CI / Clippy (push) Successful in 2m40s
CI / Test (push) Successful in 5m50s
CI / Build cortex SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
build-prerelease / Resolve version stamps + change detection (push) Successful in 31s
build-prerelease / Build neuron-blackwell (push) Successful in 1m41s
build-prerelease / Build neuron-ada (push) Successful in 2m15s
build-prerelease / Build neuron-ampere (push) Successful in 2m18s
build-prerelease / Build helexa-bench binary (push) Successful in 2m20s
build-prerelease / Build cortex binary (push) Successful in 2m22s
build-prerelease / Lint (fmt + clippy) (push) Successful in 3m10s
build-prerelease / Test (push) Successful in 5m19s
build-prerelease / Package helexa-bench RPM (push) Successful in 1m18s
build-prerelease / Package cortex RPM (push) Successful in 1m20s
build-prerelease / Package helexa-neuron-ampere RPM (push) Successful in 1m40s
build-prerelease / Package helexa-neuron-ada RPM (push) Successful in 1m44s
build-prerelease / Package helexa-neuron-blackwell RPM (push) Successful in 1m45s
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Successful in 57s
Stage 1 identity (#49): cortex now knows who a request is for. Identity
rides standard bearer auth only (Authorization: Bearer <key>) — no custom
required headers or body fields — which is what keeps every tier
OpenAI-compatible by construction.

- cortex-gateway::auth: `require_principal` axum middleware
  (from_fn_with_state), wired in build_app outer-to-inner as
  trace → CORS → auth → handlers (CORS outer so preflight short-circuits).
  It resolves the bearer key via the EntitlementProvider, inserts the
  typed Principal into request extensions (for metering #51 / enforcement
  #52), and stamps internal x-helexa-account-id / x-helexa-key-id headers
  so the principal reaches neuron, which trusts cortex over WireGuard (#54).
- Anti-spoofing: client-supplied principal headers are stripped before the
  authoritative value is stamped — a client can never assert a principal
  it didn't authenticate as.
- Rejection contract (#63): missing key under require_auth, or any present
  but unresolvable key, → 401 invalid_api_key in the #60 envelope. /health
  and / stay public. require_auth=false (default) allows anonymous through
  but still 401s a present-but-invalid key.
- Header-name constants (HEADER_ACCOUNT_ID/KEY_ID) live in cortex-core so
  neuron (#54) shares them. The chat/completions/responses paths forward
  the stamped headers automatically via proxy::forward_request; the
  Anthropic streaming + non-streaming paths forward them explicitly via
  auth::forward_principal_headers (they build their own upstream requests).

5 integration tests: missing-key 401, invalid-key 401 (even when auth not
required, not dispatched), valid key reaches neuron with principal headers
+ spoofed header stripped, anonymous allowed when not required, /health
public. Local fmt/clippy/test all green.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-06-17 19:07:10 +03:00
48 changed files with 4449 additions and 283 deletions

View File

@@ -185,6 +185,32 @@ Run these locally before pushing. `cargo fmt --all` fixes formatting
automatically. Clippy warnings must be resolved, not suppressed with
`#[allow(...)]` unless there is a clear rationale.
## Development workflow
Work each change on its own branch; `main` stays releasable.
1. Implement on a feature branch (`fix/<issue>-…`, `feat/<issue>-…`).
2. Run the CI triad locally (`cargo fmt --check --all`,
`cargo clippy --workspace -- -D warnings`, `cargo test --workspace`).
Local builds are **CPU-only** — the `#[cfg(feature = "cuda")]` neuron/TP
paths do NOT compile locally. The branch CI's **CUDA type-check** job is
the only thing that validates them, so for any neuron change the push to
Gitea is the real gate, not a rubber stamp.
3. Push the branch on local-green (no need to ask first), and background-watch
its CI run via the gitea-mcp `actions_run_read` tools. Start the next piece
of work meanwhile.
4. Merge to `main` when the four **validation** jobs are green — Format,
Clippy, Test, CUDA type-check. The SRPM / COPR / version-bump jobs are the
deploy pipeline (they run on `main`), not validation — don't wait on them.
5. Merging/pushing to `main` triggers the auto-deploy pipeline.
Docs-only changes (no `#[cfg(feature = "cuda")]` impact) can go straight to
`main` — there's nothing for the CUDA type-check to prove.
SSH note: the gitea remote host offers multiple agent keys and cuts the
connection before reaching the right one. This repo pins the working key via
`git config core.sshCommand "ssh -i ~/.ssh/id_grenade -o IdentitiesOnly=yes"`.
## Environment
- Targets Fedora 43 (systemd, SELinux enforcing)

33
Cargo.lock generated
View File

@@ -800,6 +800,7 @@ dependencies = [
"cortex-core",
"eventsource-stream",
"futures",
"helexa-stream",
"metrics",
"metrics-exporter-prometheus",
"reqwest",
@@ -1922,6 +1923,38 @@ dependencies = [
"tracing-subscriber",
]
[[package]]
name = "helexa-router"
version = "0.1.16"
dependencies = [
"anyhow",
"axum",
"chrono",
"clap",
"cortex-core",
"figment",
"reqwest",
"serde",
"serde_json",
"tokio",
"tower-http",
"tracing",
"tracing-subscriber",
]
[[package]]
name = "helexa-stream"
version = "0.1.16"
dependencies = [
"async-stream",
"axum",
"futures",
"reqwest",
"thiserror 2.0.18",
"tokio",
"tokio-stream",
]
[[package]]
name = "hermit-abi"
version = "0.5.2"

View File

@@ -7,6 +7,8 @@ members = [
"crates/neuron",
"crates/helexa-acp",
"crates/helexa-bench",
"crates/helexa-router",
"crates/helexa-stream",
]
[workspace.package]

View File

@@ -68,6 +68,57 @@ pub struct HealthResponse {
pub devices: Vec<DeviceHealth>,
#[serde(default)]
pub activation: ActivationStatus,
/// Per-model admission load (#53): how many requests are running vs.
/// queued on each loaded model right now. Cortex's load-aware router
/// (#55) reads this to spread traffic across replicas and to propagate
/// honest backpressure. `#[serde(default)]` keeps older gateways/neurons
/// interoperable (absent → empty → treated as no load info).
#[serde(default)]
pub models: Vec<ModelLoad>,
}
/// Live admission load for one loaded model (#53).
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModelLoad {
pub id: String,
/// Requests currently running (batch-1 → 0 or 1).
pub in_flight: usize,
/// Requests waiting in the bounded admission queue.
pub queue_depth: usize,
}
#[cfg(test)]
mod health_load_tests {
use super::*;
#[test]
fn health_response_without_models_field_still_deserializes() {
// A pre-#53 neuron's /health payload omits `models`; the gateway
// must still parse it (serde default → empty).
let json = r#"{"uptime_secs":42,"devices":[]}"#;
let resp: HealthResponse = serde_json::from_str(json).expect("back-compat parse");
assert_eq!(resp.uptime_secs, 42);
assert!(resp.models.is_empty());
}
#[test]
fn health_response_round_trips_model_load() {
let resp = HealthResponse {
uptime_secs: 1,
devices: vec![],
activation: ActivationStatus::default(),
models: vec![ModelLoad {
id: "Qwen/Qwen3.6-27B".into(),
in_flight: 1,
queue_depth: 3,
}],
};
let s = serde_json::to_string(&resp).unwrap();
let back: HealthResponse = serde_json::from_str(&s).unwrap();
assert_eq!(back.models.len(), 1);
assert_eq!(back.models[0].in_flight, 1);
assert_eq!(back.models[0].queue_depth, 3);
}
}
/// High-level activation state of the neuron daemon. The HTTP listener

View File

@@ -23,6 +23,14 @@
use async_trait::async_trait;
use serde::{Deserialize, Serialize};
/// Internal header carrying the resolved account id from cortex to neuron.
/// neuron trusts these over the WireGuard link (#54); cortex **strips** any
/// client-supplied copy before stamping the authoritative value, so a client
/// can never assert a principal directly.
pub const HEADER_ACCOUNT_ID: &str = "x-helexa-account-id";
/// Internal header carrying the resolved key id from cortex to neuron.
pub const HEADER_KEY_ID: &str = "x-helexa-key-id";
/// Who a request is for. Resolved once at the edge from the bearer key and
/// carried through the request context. `account_id` is the billable owner
/// (spendable at any operator, by decision); `key_id` identifies the

View File

@@ -54,10 +54,26 @@ pub struct ModelLimit {
pub output: usize,
}
/// Operator-set pricing in USD per 1M tokens.
/// Operator-set pricing, **USD per 1,000,000 tokens, as JSON numbers**
/// (`float`) — the models.dev/opencode `cost` convention, which is what
/// helexa's primary client reads. NOT per-token, NOT decimal strings (that
/// is OpenRouter's `pricing` shape, which helexa deliberately does not emit
/// — see #68). A client must not rescale by 10⁶.
///
/// Self-hosted deployments typically leave both at `0.0`. Cache fields are
/// optional — set when the backend supports a prefix-cache discount tier.
/// `cost` is sourced from the operator's `models.toml` catalogue profile and
/// surfaced verbatim on `/v1/models`. The *absent* vs *zero* distinction is
/// intentional and load-bearing (#68):
/// - **`cost` absent** (the whole object omitted) — the model is **not
/// priced**: the operator has not declared a rate. Clients should treat
/// spend as unknown, not free.
/// - **`cost` present with `input`/`output` = `0.0`** — the model is
/// **intentionally free** (self-hosted, no charge). opencode renders `$0`.
///
/// Cache fields are optional — set them only when the backend supports a
/// prefix-cache discount tier (relevant once cache-token reporting, #64,
/// lands). The advertised rate here must equal the rate metering (#51) and
/// reconciliation (#58/#59) bill against; today both read this catalogue
/// value.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModelCost {
/// USD per 1M input (prompt) tokens.
@@ -98,7 +114,8 @@ pub struct ModelInfo {
/// `None` when neither the catalogue nor the loaded model can provide it.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub limit: Option<ModelLimit>,
/// Operator-set pricing in USD per 1M tokens (0.0 = free/self-hosted).
/// Operator-set pricing — see [`ModelCost`] for units and the
/// absent (not priced) vs `0.0` (intentionally free) distinction.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub cost: Option<ModelCost>,
/// `true` when the model's tokenizer contains recognised tool-call

View File

@@ -1,4 +1,4 @@
use crate::discovery::{ActivationStatus, DiscoveryResponse};
use crate::discovery::{ActivationStatus, DiscoveryResponse, ModelLoad};
use crate::harness::{ModelCost, ModelLimit};
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
@@ -27,6 +27,17 @@ pub struct NodeState {
/// to synthesize `Loading` locations so clients see a catalogued
/// model that's mid-prewarm as "loading", not "missing".
pub activation: Option<ActivationStatus>,
/// Last-seen per-model admission load from this neuron's `/health`
/// (#53), keyed by model id. The router (#55) reads it to pick the
/// least-busy replica when a model is loaded on more than one neuron.
/// Empty until the first /health poll reports load.
pub model_load: HashMap<String, ModelLoad>,
/// Consecutive failed `/models` polls. The poller marks a node
/// unhealthy only once this crosses a threshold, so a single transient
/// miss (e.g. a neuron momentarily slow to answer while busy) doesn't
/// yank the node — and all its models — out of routing. Reset to 0 on
/// any successful poll.
pub consecutive_poll_failures: u32,
}
/// A model registered on a node, with its runtime status.
@@ -125,7 +136,9 @@ pub struct CortexModelEntry {
/// at load time. `None` when neither source provides it.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub limit: Option<ModelLimit>,
/// Operator-set pricing in USD per 1M tokens (0.0 = free/self-hosted).
/// Operator-set pricing from the catalogue profile — see
/// [`cortex_core::harness::ModelCost`] for units (USD per 1M tokens) and
/// the absent (not priced) vs `0.0` (intentionally free) distinction.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub cost: Option<ModelCost>,
/// `true` when any neuron reports this model supports tool calls.

View File

@@ -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

View File

@@ -32,6 +32,8 @@ pub async fn stream_translated(
openai_body: axum::body::Bytes,
model_id: &str,
node_name: &str,
inbound_headers: &axum::http::HeaderMap,
usage_sink: Option<crate::metering::UsageSink>,
) -> Response {
let url = format!("{endpoint}/v1/chat/completions");
tracing::info!(
@@ -42,13 +44,14 @@ pub async fn stream_translated(
"proxying streaming request (anthropic SSE translation)"
);
let upstream = match client
let request = crate::auth::forward_principal_headers(
client
.post(&url)
.header("content-type", "application/json")
.body(openai_body)
.send()
.await
{
.body(openai_body),
inbound_headers,
);
let upstream = match request.send().await {
Ok(r) => r,
Err(e) => {
tracing::warn!(
@@ -94,6 +97,10 @@ pub async fn stream_translated(
let mut saw_tool_call = false;
let mut last_finish: Option<String> = None;
let mut frames = 0u64;
// Engine-truth usage for metering (#51), scanned from the upstream
// frames (neuron emits a final `usage` object on the stream, #48).
let mut usage_prompt = 0u64;
let mut usage_completion = 0u64;
'outer: while let Some(block) = upstream.next().await {
let block = match block {
@@ -121,6 +128,15 @@ pub async fn stream_translated(
continue;
}
tracing::trace!(node = %node, frame = %data, "anthropic stream: upstream frame");
// Capture usage for metering before translation — the
// usage object rides on a late frame (often after the
// last content delta).
if let Some(p) = crate::proxy::last_count_for(data, "prompt_tokens") {
usage_prompt = p;
}
if let Some(c) = crate::proxy::last_count_for(data, "completion_tokens") {
usage_completion = c;
}
let Ok(chunk) = serde_json::from_str::<ChatCompletionChunk>(data) else {
tracing::debug!(node = %node, "anthropic stream: unparsable upstream frame skipped");
continue;
@@ -162,6 +178,14 @@ pub async fn stream_translated(
terminated = done,
"anthropic stream complete"
);
// Settle metering with the observed usage (#51). Runs on every exit
// path of the pump — clean end, early break, or upstream error — so
// the reservation is always resolved. `(0, 0)` when no usage frame
// was seen, which releases without recording spend.
if let Some(sink) = usage_sink {
sink(usage_prompt, usage_completion);
}
});
Response::builder()

View File

@@ -0,0 +1,133 @@
//! API-key authentication + principal resolution (#49).
//!
//! Identity rides standard bearer auth only — `Authorization: Bearer <key>`
//! — which is what keeps every tier OpenAI-compatible by construction (no
//! custom required headers or body fields, per #47). The middleware resolves
//! the key to a [`Principal`] via the [`EntitlementProvider`], carries it in
//! the request extensions for cortex-side metering/enforcement (#51/#52), and
//! stamps it as internal headers on the request so it reaches neuron, which
//! trusts cortex's assertion over WireGuard (#54).
//!
//! Anti-spoofing: any client-supplied principal header is **stripped** before
//! the authoritative value is stamped, so a client can never assert a
//! principal it didn't authenticate as.
//!
//! Rejection contract (#63): missing key under `require_auth`, or any present
//! but unresolvable key, yields `401 invalid_api_key` in the #60 envelope.
use crate::error::envelope_response;
use crate::state::CortexState;
use axum::extract::{Request, State};
use axum::http::header::AUTHORIZATION;
use axum::http::{HeaderMap, HeaderValue};
use axum::middleware::Next;
use axum::response::Response;
use cortex_core::entitlements::{HEADER_ACCOUNT_ID, HEADER_KEY_ID};
use cortex_core::error_envelope::OpenAiError;
use std::sync::Arc;
/// Endpoints that never require auth: liveness/readiness probes. Everything
/// else flows through resolution.
fn is_public(path: &str) -> bool {
path == "/health" || path == "/"
}
/// Extract the bearer token from an `Authorization` header value, if present
/// and well-formed. Scheme match is case-insensitive per RFC 7235.
fn parse_bearer(headers: &HeaderMap) -> Option<String> {
let raw = headers.get(AUTHORIZATION)?.to_str().ok()?;
let (scheme, token) = raw.split_once(' ')?;
if scheme.eq_ignore_ascii_case("bearer") {
let token = token.trim();
(!token.is_empty()).then(|| token.to_string())
} else {
None
}
}
/// Axum middleware: resolve the bearer key, attach the principal, stamp the
/// internal headers. Wired in `build_app` via `from_fn_with_state`.
pub async fn require_principal(
State(fleet): State<Arc<CortexState>>,
mut req: Request,
next: Next,
) -> Response {
if is_public(req.uri().path()) {
return next.run(req).await;
}
// Anti-spoof: drop any client-supplied principal headers up front.
{
let headers = req.headers_mut();
headers.remove(HEADER_ACCOUNT_ID);
headers.remove(HEADER_KEY_ID);
}
match parse_bearer(req.headers()) {
Some(key) => match fleet.entitlements.resolve(&key).await {
Ok(principal) => {
// Stamp the authoritative principal for neuron. Account/key
// ids come from operator config, so they're valid header
// values; guard anyway and skip a malformed one rather than
// panic.
if let (Ok(account), Ok(key_id)) = (
HeaderValue::from_str(&principal.account_id),
HeaderValue::from_str(&principal.key_id),
) {
let headers = req.headers_mut();
headers.insert(HEADER_ACCOUNT_ID, account);
headers.insert(HEADER_KEY_ID, key_id);
}
// Carry the typed principal for cortex-side metering (#51)
// and budget enforcement (#52).
req.extensions_mut().insert(principal);
next.run(req).await
}
// An unrecognized key only hard-fails when auth is *required*.
// In allow-anonymous mode (the default) we must IGNORE it and
// serve the request unauthenticated — otherwise the placeholder
// keys that OpenAI-compatible clients send by default (opencode,
// Open WebUI, Agent Zero, litellm) would all break, even though
// the operator never opted into auth. Pre-#49 the bearer was
// never inspected at all; this preserves that for require_auth=false.
Err(_) => {
if fleet.require_auth {
unauthorized("invalid API key")
} else {
tracing::debug!(
"ignoring unrecognized bearer token (require_auth=false): serving anonymously"
);
next.run(req).await
}
}
},
None => {
if fleet.require_auth {
unauthorized("missing API key; supply 'Authorization: Bearer <key>'")
} else {
next.run(req).await
}
}
}
}
/// `401 invalid_api_key` in the standard envelope (#63).
fn unauthorized(message: &str) -> Response {
envelope_response(OpenAiError::invalid_api_key(message))
}
/// Copy the cortex-stamped principal headers from an inbound [`HeaderMap`]
/// onto an outbound reqwest builder. Used by the Anthropic proxy paths,
/// which construct their own upstream requests instead of going through
/// [`crate::proxy::forward_request`] (which forwards all headers verbatim).
pub fn forward_principal_headers(
mut builder: reqwest::RequestBuilder,
headers: &HeaderMap,
) -> reqwest::RequestBuilder {
for name in [HEADER_ACCOUNT_ID, HEADER_KEY_ID] {
if let Some(value) = headers.get(name) {
builder = builder.header(name, value);
}
}
builder
}

View File

@@ -190,7 +190,7 @@ async fn completions(
/// `POST /v1/messages` — accept Anthropic format, translate, proxy, translate back.
async fn anthropic_messages(
State(fleet): State<Arc<CortexState>>,
_headers: HeaderMap,
headers: HeaderMap,
body: Bytes,
) -> Response {
// Parse as Anthropic request.
@@ -306,6 +306,29 @@ async fn anthropic_messages(
}
let start = Instant::now();
// Per-request metering + budget enforcement (#51/#52), same lifecycle as
// the OpenAI paths. Estimate from the translated OpenAI body (what neuron
// sees). Refuse over-cap before dispatch via the #63 envelope; otherwise
// build the sink consumed by whichever branch runs below.
let usage_sink = match crate::metering::principal_from_headers(&headers) {
Some(principal) => {
let advertised =
advertised_output_limit(&fleet, &route.node_name, &route.resolved_model_id).await;
let max_tokens = crate::metering::reservation_estimate(&openai_body, advertised);
match crate::metering::reserve_or_reject(
Arc::clone(&fleet.entitlements),
&principal,
max_tokens,
)
.await
{
Ok(guard) => Some(crate::metering::usage_sink(principal, guard)),
Err(env) => return crate::error::envelope_response(env),
}
}
None => None,
};
if is_streaming {
// Anthropic SSE translation (#24): upstream speaks OpenAI SSE;
// re-frame it event-by-event into Anthropic's message_start /
@@ -316,6 +339,8 @@ async fn anthropic_messages(
openai_body,
&model_id,
&route.node_name,
&headers,
usage_sink,
)
.await;
metrics::histogram!("cortex_request_duration_seconds", &labels)
@@ -335,11 +360,14 @@ async fn anthropic_messages(
cold_start = route.cold_start,
"proxying request"
);
let upstream_resp = fleet
let upstream_resp = crate::auth::forward_principal_headers(
fleet
.http_client
.post(&target_url)
.body(openai_body)
.header("content-type", "application/json")
.header("content-type", "application/json"),
&headers,
)
.send()
.await;
@@ -437,6 +465,15 @@ async fn anthropic_messages(
metrics::histogram!("cortex_request_duration_seconds", &labels)
.record(start.elapsed().as_secs_f64());
// Settle metering with the upstream usage (#51). Scanned from the
// raw body — same engine-truth source as the streaming path — so we
// don't depend on the typed usage struct's optionality.
if let Some(sink) = usage_sink {
let tail = String::from_utf8_lossy(&body_bytes);
let prompt = proxy::last_count_for(&tail, "prompt_tokens").unwrap_or(0);
let completion = proxy::last_count_for(&tail, "completion_tokens").unwrap_or(0);
sink(prompt, completion);
}
// Did the model actually produce a structured tool call, or just
// text? This is the single most useful signal for "is tool
// calling working end-to-end" — a `false` here alongside a
@@ -724,6 +761,19 @@ async fn proxy_with_metrics(
body: Bytes,
model_id: &str,
) -> Response {
// Fail-fast prompt pre-validation (#56): refuse a prompt that already
// exceeds the model's advertised context window *before* dispatching to
// neuron — the same `400 context_length_exceeded` neuron would emit on
// overflow, just earlier and without burning a cold-load/queue slot.
// cortex has no tokenizer, so the estimate under-counts and neuron stays
// the exact wall; we only catch gross overages (the A0 failure mode).
if let Some(context) = advertised_context(fleet, &route.node_name, model_id).await {
let est = estimate_prompt_tokens(&body);
if est > context {
return context_length_exceeded_response(context, est, &headers);
}
}
let labels = [
("model", model_id.to_string()),
("node", route.node_name.clone()),
@@ -734,9 +784,42 @@ async fn proxy_with_metrics(
metrics::counter!("cortex_cold_starts_total", &labels).increment(1);
}
// Per-request metering + budget enforcement (#51/#52): reconstruct the
// principal from the middleware-stamped headers, reserve the request's
// upper-bound cost (prompt estimate + max output), and build the
// completion sink that settles actual spend when the response finishes.
// A reservation over the hard cap is refused *before* dispatch with the
// #63 envelope. Anonymous requests skip all of this. Must happen before
// `headers`/`body` are moved into the proxy.
let usage_sink = match crate::metering::principal_from_headers(&headers) {
Some(principal) => {
let advertised = advertised_output_limit(fleet, &route.node_name, model_id).await;
let max_tokens = crate::metering::reservation_estimate(&body, advertised);
match crate::metering::reserve_or_reject(
Arc::clone(&fleet.entitlements),
&principal,
max_tokens,
)
.await
{
Ok(guard) => Some(crate::metering::usage_sink(principal, guard)),
Err(env) => return crate::error::envelope_response(env),
}
}
None => None,
};
let start = Instant::now();
let result =
proxy::forward_request(&fleet.http_client, route, path, headers, body, model_id).await;
let result = proxy::forward_request(
&fleet.http_client,
route,
path,
headers,
body,
model_id,
usage_sink,
)
.await;
let duration = start.elapsed();
match result {
@@ -755,6 +838,117 @@ async fn proxy_with_metrics(
}
}
/// The model's advertised `limit.output` (#62) on a given node, used as the
/// default output budget for budget reservations (#52) when the request
/// omits `max_(completion_)tokens`. `None` when the node/model/limit is
/// unknown — callers fall back to [`crate::metering::FALLBACK_MAX_OUTPUT`].
async fn advertised_output_limit(
fleet: &CortexState,
node_name: &str,
model_id: &str,
) -> Option<u64> {
let nodes = fleet.nodes.read().await;
nodes
.get(node_name)?
.models
.get(model_id)?
.limit
.as_ref()
.map(|l| l.output as u64)
}
/// The model's advertised hard context window (`limit.context`, #62/#67) on a
/// node, used for fail-fast prompt pre-validation (#56). `None` when no limit
/// is known — pre-validation is then skipped and neuron remains the wall.
async fn advertised_context(fleet: &CortexState, node_name: &str, model_id: &str) -> Option<u64> {
let nodes = fleet.nodes.read().await;
nodes
.get(node_name)?
.models
.get(model_id)?
.limit
.as_ref()
.map(|l| l.context as u64)
}
/// Conservative prompt-token estimate (~4 chars/token over message text).
/// cortex has no tokenizer; under-counting is the safe direction — we only
/// pre-reject gross overages (#56), and neuron enforces the exact wall.
fn estimate_prompt_tokens(body: &[u8]) -> u64 {
let Ok(v) = serde_json::from_slice::<Value>(body) else {
return (body.len() as u64 / 4).max(1);
};
let mut chars = 0usize;
if let Some(messages) = v.get("messages").and_then(Value::as_array) {
for m in messages {
match m.get("content") {
Some(Value::String(s)) => chars += s.len(),
Some(Value::Array(parts)) => {
for p in parts {
if let Some(t) = p.get("text").and_then(Value::as_str) {
chars += t.len();
}
}
}
_ => {}
}
chars += 8; // rough per-message role/formatting overhead
}
} else if let Some(prompt) = v.get("prompt").and_then(Value::as_str) {
chars += prompt.len(); // legacy /v1/completions
} else {
return (body.len() as u64 / 4).max(1);
}
(chars as u64 / 4).max(1)
}
/// Client-specific, advisory guidance for an over-long prompt (#56),
/// fingerprinted from `User-Agent`. Strictly advisory: it rides the
/// `X-Helexa-Advice` header only, never the error envelope, and behaviour
/// never depends on it. Unknown clients get nothing.
fn client_advice(headers: &HeaderMap) -> Option<&'static str> {
let ua = headers
.get(axum::http::header::USER_AGENT)?
.to_str()
.ok()?
.to_ascii_lowercase();
if ua.contains("litellm") {
Some(
"litellm forwards the full context; lower the configured context window or enable client-side compaction",
)
} else if ua.contains("agent-zero") || ua.contains("agent zero") {
Some("reduce the conversation/context size or summarize earlier turns before resending")
} else if ua.contains("zed") {
Some("reduce the assistant context window in Zed's settings")
} else {
None
}
}
/// `400 context_length_exceeded` for an over-long prompt caught at the edge
/// (#56), in the #60 envelope — the same shape neuron emits on overflow, so
/// clients (opencode auto-compacts) handle it identically. Attaches the
/// advisory `X-Helexa-Advice` header for fingerprinted clients.
fn context_length_exceeded_response(
context: u64,
prompt_est: u64,
headers: &HeaderMap,
) -> Response {
let env = OpenAiError::context_length_exceeded(format!(
"This model's maximum context length is {context} tokens. Your request is \
estimated at ~{prompt_est} tokens. Please reduce the length of the messages."
))
.with_extra("max", json!(context))
.with_extra("estimated_prompt_tokens", json!(prompt_est));
let mut response = crate::error::envelope_response(env);
if let Some(advice) = client_advice(headers)
&& let Ok(value) = axum::http::HeaderValue::from_str(advice)
{
response.headers_mut().insert("x-helexa-advice", value);
}
response
}
/// Update `last_accessed` timestamp for a model on a node (drives LRU eviction).
async fn touch_model(fleet: &CortexState, node_name: &str, model_id: &str) {
let mut nodes = fleet.nodes.write().await;

View File

@@ -1,8 +1,10 @@
pub mod anthropic_sse;
pub mod auth;
pub mod entitlements_local;
pub mod error;
pub mod evictor;
pub mod handlers;
pub mod metering;
pub mod metrics;
pub mod poller;
pub mod proxy;
@@ -11,15 +13,26 @@ pub mod state;
use anyhow::Result;
use axum::Router;
use axum::middleware::from_fn_with_state;
use cortex_core::config::GatewayConfig;
use std::sync::Arc;
use tower_http::cors::CorsLayer;
use tower_http::trace::TraceLayer;
/// Build the Axum application router with all routes wired up.
///
/// Layer order (outermost first): trace → CORS → auth → handlers. CORS is
/// outer to auth so preflight `OPTIONS` short-circuits before resolution;
/// auth (`require_principal`) resolves the bearer key, attaches the
/// principal, and stamps the internal principal headers before any handler
/// runs.
pub fn build_app(fleet: Arc<state::CortexState>) -> Router {
Router::new()
.merge(handlers::api_routes())
.layer(from_fn_with_state(
Arc::clone(&fleet),
auth::require_principal,
))
.layer(CorsLayer::permissive())
.layer(TraceLayer::new_for_http())
.with_state(fleet)

View File

@@ -0,0 +1,219 @@
//! Per-request token metering (#51).
//!
//! Captures the real `(prompt, completion)` usage of every request and feeds
//! it to two places: the [`EntitlementProvider`] spend ledger (via
//! reserve→settle) and per-principal Prometheus counters. The principal is
//! reconstructed from the internal headers the auth middleware stamped (#49),
//! so this works uniformly across every proxy path without threading the
//! typed principal through each handler.
//!
//! The reserve→settle lifecycle is established here but, in this phase,
//! reserves **zero** tokens — metering only, no enforcement. Budget
//! enforcement (#52) flips the reserved amount to the real
//! `prompt + max_output` and handles the [`BudgetError`] rejection; the
//! settle/release plumbing is identical, so that change is localized.
//!
//! [`ReservationGuard`] makes leaks impossible: settling records actual
//! spend and releases the unused remainder; dropping a guard that was never
//! settled releases the whole reservation. So an early return, error path,
//! or dropped stream can't strand a reservation.
use axum::http::HeaderMap;
use cortex_core::entitlements::{
BudgetError, EntitlementProvider, HEADER_ACCOUNT_ID, HEADER_KEY_ID, Principal,
};
use cortex_core::error_envelope::OpenAiError;
use std::sync::Arc;
/// Fallback output-token budget when neither the request nor the model's
/// advertised limit gives one. Bounds the reservation so a capped key is
/// still gated even on under-specified requests (#52).
pub const FALLBACK_MAX_OUTPUT: u64 = 4096;
/// Invoked exactly once at request completion with best-effort
/// `(prompt_tokens, completion_tokens)`. When no usage could be observed
/// (e.g. a pre-dispatch failure or a dropped stream) it is dropped unused —
/// which releases the held reservation via [`ReservationGuard`]'s `Drop`.
pub type UsageSink = Box<dyn FnOnce(u64, u64) + Send>;
/// Reconstruct the principal from the cortex-stamped internal headers. The
/// auth middleware strips any client copy and stamps the authoritative value,
/// so these headers are trustworthy within cortex. `None` for anonymous
/// (unauthenticated) requests.
pub fn principal_from_headers(headers: &HeaderMap) -> Option<Principal> {
let account_id = headers.get(HEADER_ACCOUNT_ID)?.to_str().ok()?.to_string();
let key_id = headers.get(HEADER_KEY_ID)?.to_str().ok()?.to_string();
Some(Principal { account_id, key_id })
}
/// Emit per-principal spend counters (#51). Labelled by account/key only —
/// both are operator-bounded, so cardinality is controlled.
pub fn record_spend(principal: &Principal, prompt: u64, completion: u64) {
let labels = [
("account", principal.account_id.clone()),
("key", principal.key_id.clone()),
];
metrics::counter!("cortex_spend_tokens_total", &labels).increment(prompt + completion);
metrics::counter!("cortex_spend_prompt_tokens_total", &labels).increment(prompt);
metrics::counter!("cortex_spend_completion_tokens_total", &labels).increment(completion);
}
/// Holds a budget reservation for the life of a request. [`settle`] records
/// actual spend and releases the remainder; an un-settled guard releases the
/// whole reservation when dropped. Anonymous requests carry an empty guard,
/// where every operation is a no-op.
///
/// [`settle`]: ReservationGuard::settle
pub struct ReservationGuard {
provider: Arc<dyn EntitlementProvider>,
reservation: Option<cortex_core::entitlements::Reservation>,
}
impl ReservationGuard {
/// An empty guard for an anonymous request — no reservation to resolve.
pub fn anonymous(provider: Arc<dyn EntitlementProvider>) -> Self {
Self {
provider,
reservation: None,
}
}
/// Wrap an already-acquired reservation.
fn held(
provider: Arc<dyn EntitlementProvider>,
reservation: cortex_core::entitlements::Reservation,
) -> Self {
Self {
provider,
reservation: Some(reservation),
}
}
/// Settle with the tokens actually consumed, disarming the drop-release.
/// Spawns the (fast, in-process for the local provider) settle so the
/// caller — which may be a sync stream-completion callback — needn't
/// await.
pub fn settle(mut self, actual_tokens: u64) {
if let Some(reservation) = self.reservation.take() {
let provider = Arc::clone(&self.provider);
tokio::spawn(async move {
provider.settle(reservation, actual_tokens).await;
});
}
}
}
impl Drop for ReservationGuard {
fn drop(&mut self) {
if let Some(reservation) = self.reservation.take() {
let provider = Arc::clone(&self.provider);
tokio::spawn(async move {
provider.release(reservation).await;
});
}
}
}
/// Build the completion sink for an authenticated request: record spend and
/// settle the reservation with the observed total. Dropping it unused (no
/// usage observed) releases the reservation via the guard.
pub fn usage_sink(principal: Principal, guard: ReservationGuard) -> UsageSink {
Box::new(move |prompt, completion| {
record_spend(&principal, prompt, completion);
guard.settle(prompt + completion);
})
}
/// Reserve the request's upper-bound token cost for the principal, refusing
/// *before* dispatch if it would exceed the hard cap (#52). On success
/// returns a guard the caller settles with actual usage; on refusal returns
/// the #63 envelope (`rate_limit_exceeded` + `Retry-After` for a resetting
/// window, `insufficient_quota` for a hard balance — never `402`).
pub async fn reserve_or_reject(
provider: Arc<dyn EntitlementProvider>,
principal: &Principal,
max_tokens: u64,
) -> Result<ReservationGuard, OpenAiError> {
match provider.reserve(principal, max_tokens).await {
Ok(reservation) => Ok(ReservationGuard::held(provider, reservation)),
Err(err) => Err(budget_error_to_envelope(err)),
}
}
/// Map a [`BudgetError`] to the #63 envelope. The provider chose the window
/// semantics; this only translates them to HTTP.
fn budget_error_to_envelope(err: BudgetError) -> OpenAiError {
match err {
BudgetError::RateLimited {
retry_after_secs, ..
} => OpenAiError::rate_limit_exceeded(err.to_string(), retry_after_secs),
BudgetError::InsufficientQuota { .. } => OpenAiError::insufficient_quota(err.to_string()),
}
}
/// Upper-bound tokens to reserve for a request (#52): an over-estimate of
/// the prompt plus the maximum output. `advertised_output` is the model's
/// `limit.output` (#62), used when the request omits `max_(completion_)tokens`.
/// Over-reserving is safe — settle corrects spend to the actual usage.
pub fn reservation_estimate(body: &[u8], advertised_output: Option<u64>) -> u64 {
let max_output = requested_max_output(body)
.or(advertised_output)
.unwrap_or(FALLBACK_MAX_OUTPUT);
estimate_prompt_tokens(body).saturating_add(max_output)
}
/// The client's requested output cap, from `max_completion_tokens` (or the
/// legacy `max_tokens`). `None` when unspecified.
fn requested_max_output(body: &[u8]) -> Option<u64> {
let v: serde_json::Value = serde_json::from_slice(body).ok()?;
v.get("max_completion_tokens")
.or_else(|| v.get("max_tokens"))
.and_then(serde_json::Value::as_u64)
}
/// Rough prompt-token estimate at ~4 chars/token over the whole body. cortex
/// has no tokenizer; JSON overhead makes this a conservative over-estimate,
/// and neuron remains the exact context wall (#56/#60). Settle reconciles to
/// the real usage afterward.
fn estimate_prompt_tokens(body: &[u8]) -> u64 {
(body.len() as u64 / 4).max(1)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn requested_max_output_prefers_max_completion_tokens() {
let body = br#"{"model":"m","max_completion_tokens":256,"max_tokens":99}"#;
assert_eq!(requested_max_output(body), Some(256));
}
#[test]
fn requested_max_output_falls_back_to_legacy_max_tokens() {
let body = br#"{"model":"m","max_tokens":128}"#;
assert_eq!(requested_max_output(body), Some(128));
}
#[test]
fn estimate_uses_requested_output_when_present() {
// Requested output dominates; prompt estimate is small for a tiny body.
let body = br#"{"model":"m","max_tokens":1000}"#;
let est = reservation_estimate(body, Some(8192));
assert!(est >= 1000 && est < 1100, "est was {est}");
}
#[test]
fn estimate_uses_advertised_output_when_request_omits_it() {
let body = br#"{"model":"m","messages":[]}"#;
let est = reservation_estimate(body, Some(8192));
assert!(est >= 8192, "est was {est}");
}
#[test]
fn estimate_falls_back_when_nothing_advertised() {
let body = br#"{"model":"m"}"#;
let est = reservation_estimate(body, None);
assert!(est >= FALLBACK_MAX_OUTPUT, "est was {est}");
}
}

View File

@@ -63,4 +63,16 @@ fn describe_metrics() {
"cortex_cold_starts_total",
"Total number of cold-start model loads"
);
metrics::describe_counter!(
"cortex_spend_tokens_total",
"Total metered tokens (prompt + completion) per principal, labelled by account/key (#51)"
);
metrics::describe_counter!(
"cortex_spend_prompt_tokens_total",
"Metered prompt tokens per principal, labelled by account/key (#51)"
);
metrics::describe_counter!(
"cortex_spend_completion_tokens_total",
"Metered completion tokens per principal, labelled by account/key (#51)"
);
}

View File

@@ -5,12 +5,29 @@ use crate::state::CortexState;
use chrono::Utc;
use cortex_core::discovery::{DiscoveryResponse, HealthResponse};
use cortex_core::harness::ModelInfo;
use cortex_core::node::{ModelEntry, ModelStatus};
use cortex_core::node::{ModelEntry, ModelStatus, NodeState};
use std::sync::Arc;
use std::time::Duration;
const POLL_INTERVAL: Duration = Duration::from_secs(10);
/// Consecutive failed `/models` polls before a node is marked unhealthy.
/// Debounces transient misses (a busy neuron briefly slow to answer) so a
/// single blip can't yank a node — and its models — out of routing. At the
/// 10s poll interval this tolerates ~20s of flapping before evicting.
const POLL_FAILURE_THRESHOLD: u32 = 3;
/// Record a failed poll for `node`, marking it unhealthy only once failures
/// reach [`POLL_FAILURE_THRESHOLD`]. Below the threshold the node keeps its
/// last-known health, riding over transient misses. A successful poll resets
/// the counter (see the success arm in `poll_once`).
fn record_poll_failure(node: &mut NodeState) {
node.consecutive_poll_failures = node.consecutive_poll_failures.saturating_add(1);
if node.consecutive_poll_failures >= POLL_FAILURE_THRESHOLD {
node.healthy = false;
}
}
/// Runs forever, polling all neurons on a fixed interval.
pub async fn poll_loop(fleet: Arc<CortexState>) {
loop {
@@ -138,13 +155,14 @@ async fn poll_neuron(fleet: &CortexState, name: &str, endpoint: &str) {
// Remove models no longer reported by the neuron.
node.models.retain(|id, _| seen.contains(id));
node.consecutive_poll_failures = 0;
node.healthy = true;
node.last_poll = Some(Utc::now());
tracing::debug!(node = name, models = models.len(), "poll ok");
}
Err(e) => {
tracing::warn!(node = name, error = %e, "failed to parse /models response");
node.healthy = false;
record_poll_failure(node);
}
}
}
@@ -154,11 +172,11 @@ async fn poll_neuron(fleet: &CortexState, name: &str, endpoint: &str) {
status = %resp.status(),
"neuron returned non-success status"
);
node.healthy = false;
record_poll_failure(node);
}
Err(e) => {
tracing::warn!(node = name, error = %e, "failed to reach neuron");
node.healthy = false;
record_poll_failure(node);
}
}
@@ -200,6 +218,9 @@ async fn poll_health(fleet: &CortexState, name: &str, endpoint: &str) {
let mut nodes = fleet.nodes.write().await;
if let Some(node) = nodes.get_mut(name) {
node.activation = Some(h.activation);
// Per-model admission load (#53) → keyed by id for the
// load-aware router (#55).
node.model_load = h.models.into_iter().map(|m| (m.id.clone(), m)).collect();
}
}
Err(e) => {

View File

@@ -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
@@ -31,6 +37,7 @@ pub async fn forward_request(
headers: HeaderMap,
body: bytes::Bytes,
model_id: &str,
usage_sink: Option<crate::metering::UsageSink>,
) -> Result<Response, ProxyError> {
let request_start = Instant::now();
let url = format!("{}{}", route.endpoint, path);
@@ -41,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),
);
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)]
@@ -111,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 {
@@ -138,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
@@ -154,42 +146,29 @@ 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
/// settled and spend recorded. `None` for anonymous requests.
usage_sink: Option<crate::metering::UsageSink>,
}
impl TokenMetrics {
fn new(model_id: &str, node_name: &str, request_start: Instant) -> Self {
impl CortexMetrics {
fn new(
model_id: &str,
node_name: &str,
request_start: Instant,
usage_sink: Option<crate::metering::UsageSink>,
) -> Self {
Self {
labels: [
("model", model_id.to_string()),
@@ -198,25 +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
@@ -227,28 +200,28 @@ impl TokenMetrics {
return;
}
self.finished = true;
let Some(first) = self.first_chunk else {
return; // no body ever arrived — nothing to record
};
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 {
let ttft = first.duration_since(self.request_start).as_secs_f64();
metrics::histogram!("cortex_time_to_first_token_seconds", &self.labels).record(ttft);
if let Some(prompt) = last_count_for(&self.tail, "prompt_tokens") {
if let Some(prompt) = prompt {
metrics::counter!("cortex_prompt_tokens_total", &self.labels).increment(prompt);
}
let Some(completion) = last_count_for(&self.tail, "completion_tokens") else {
return;
};
if completion == 0 {
return;
}
metrics::counter!("cortex_completion_tokens_total", &self.labels).increment(completion);
if let Some(completion) = completion.filter(|c| *c > 0) {
metrics::counter!("cortex_completion_tokens_total", &self.labels)
.increment(completion);
let last = self.last_chunk.unwrap_or(first);
let decode_window = last.duration_since(first).as_secs_f64();
// Streaming: rate over the decode window (first→last chunk).
// Non-streaming bodies arrive as ~one chunk (window ≈ 0), where
// the only honest denominator is the full request duration.
// Non-streaming bodies arrive as ~one chunk (window ≈ 0),
// where the only honest denominator is the full request
// duration.
let secs = if decode_window >= 0.1 {
decode_window
} else {
@@ -261,96 +234,11 @@ 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,
// Per-principal metering + reservation settle (#51). Always runs so
// the reservation is resolved even when no usage/body was observed
// (sink with (0, 0) → settle 0 → release).
if let Some(sink) = self.usage_sink.take() {
sink(prompt.unwrap_or(0), completion.unwrap_or(0));
}
}
}
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
);
}
}

View File

@@ -50,6 +50,10 @@ pub enum RouteError {
"model '{model_id}' is in the catalogue but no healthy neuron's topology satisfies its constraints"
)]
NoFeasibleNeuron { model_id: String },
#[error(
"model '{model_id}' is feasible on a neuron that is currently unhealthy — retry shortly"
)]
FeasibleNodeUnhealthy { model_id: String },
#[error("cold-load of '{model_id}' on '{node}' failed: {message}")]
ColdLoadFailed {
model_id: String,
@@ -68,7 +72,9 @@ impl RouteError {
/// safe to retry the same request); everything else is 404.
pub fn http_status(&self) -> u16 {
match self {
RouteError::NoHealthyNodes | RouteError::ModelRecovering { .. } => 503,
RouteError::NoHealthyNodes
| RouteError::ModelRecovering { .. }
| RouteError::FeasibleNodeUnhealthy { .. } => 503,
_ => 404,
}
}
@@ -81,7 +87,8 @@ impl RouteError {
| RouteError::EndpointResolveFailed(_, _)
| RouteError::NoFeasibleNeuron { .. }
| RouteError::ColdLoadFailed { .. }
| RouteError::ModelRecovering { .. } => "api_error",
| RouteError::ModelRecovering { .. }
| RouteError::FeasibleNodeUnhealthy { .. } => "api_error",
}
}
@@ -94,6 +101,7 @@ impl RouteError {
RouteError::NoFeasibleNeuron { .. } => "service_unavailable",
RouteError::ColdLoadFailed { .. } => "service_unavailable",
RouteError::ModelRecovering { .. } => "service_unavailable",
RouteError::FeasibleNodeUnhealthy { .. } => "service_unavailable",
}
}
@@ -105,6 +113,7 @@ impl RouteError {
pub fn retry_after_secs(&self) -> Option<u64> {
match self {
RouteError::ModelRecovering { .. } => Some(2),
RouteError::FeasibleNodeUnhealthy { .. } => Some(3),
RouteError::NoHealthyNodes => Some(5),
_ => None,
}
@@ -132,7 +141,9 @@ pub async fn resolve(
// Snapshot loaded / unloaded / recovering state from the poller cache.
let (loaded_route, unloaded_route, recovering_node, any_healthy) = {
let nodes = fleet.nodes.read().await;
let mut loaded_route = None;
// All healthy nodes with the model loaded, each with its current
// admission load (#53) so we can pick the least-busy replica (#55).
let mut loaded_candidates: Vec<(String, String, usize)> = Vec::new();
let mut unloaded_route = None;
let mut recovering_node = None;
let mut any_healthy = false;
@@ -144,8 +155,15 @@ pub async fn resolve(
if let Some(entry) = node.models.get(model_id) {
match entry.status {
ModelStatus::Loaded | ModelStatus::Reloading => {
loaded_route = Some((node.name.clone(), node.endpoint.clone(), false));
break;
// Least-busy score: in-flight + queued from the
// neuron's last /health (#53). Unknown load (no poll
// yet) scores 0 so the replica stays eligible.
let score = node
.model_load
.get(model_id)
.map(|l| l.in_flight + l.queue_depth)
.unwrap_or(0);
loaded_candidates.push((node.name.clone(), node.endpoint.clone(), score));
}
ModelStatus::Unloaded => {
if unloaded_route.is_none() {
@@ -175,6 +193,12 @@ pub async fn resolve(
}
}
}
// Pick the least-busy loaded replica; ties break by node name for
// deterministic routing. `false` = not a cold start.
let loaded_route = loaded_candidates
.into_iter()
.min_by(|a, b| a.2.cmp(&b.2).then_with(|| a.0.cmp(&b.0)))
.map(|(name, endpoint, _score)| (name, endpoint, false));
(loaded_route, unloaded_route, recovering_node, any_healthy)
};
@@ -237,11 +261,32 @@ async fn pick_feasible_neuron(
b.2.cmp(&a.2) // pinned first (true > false)
.then(a.0.cmp(&b.0))
});
let pick = candidates.into_iter().next();
pick.map(|(n, e, _)| (n, e))
.ok_or_else(|| RouteError::NoFeasibleNeuron {
if let Some((n, e, _)) = candidates.into_iter().next() {
return Ok((n, e));
}
// No *healthy* feasible neuron. Distinguish a transient outage from a
// permanent misconfiguration: if some neuron is topologically feasible
// but currently unhealthy (e.g. it briefly missed polls while busy),
// this is retryable — return 503 + Retry-After so the client backs off
// and retries instead of treating a 404 as a hard failure. Only when no
// neuron could *ever* satisfy the topology is it a permanent 404.
let feasible_but_unhealthy = nodes.values().any(|node| {
!node.healthy
&& node
.discovery
.as_ref()
.is_some_and(|disc| profile.is_feasible_on(&node.name, &disc.devices))
});
if feasible_but_unhealthy {
Err(RouteError::FeasibleNodeUnhealthy {
model_id: profile.id.clone(),
})
} else {
Err(RouteError::NoFeasibleNeuron {
model_id: profile.id.clone(),
})
}
}
/// Issue `POST {endpoint}/models/load` for this profile on this neuron,

View File

@@ -37,6 +37,8 @@ impl CortexState {
last_poll: None,
discovery: None,
activation: None,
model_load: HashMap::new(),
consecutive_poll_failures: 0,
},
);
}

View File

@@ -0,0 +1,272 @@
//! Integration tests for API-key auth + principal resolution (#49).
//!
//! Verifies the #63 rejection contract (401 invalid_api_key via the #60
//! envelope) and that an authenticated request reaches neuron carrying the
//! internal principal headers — while a client-supplied principal header is
//! stripped (anti-spoofing).
use axum::Json;
use axum::extract::Path;
use axum::http::HeaderMap;
use axum::routing::{get, post};
use cortex_core::config::{
ApiKeyConfig, EntitlementsConfig, EvictionSettings, EvictionStrategy, GatewayConfig,
GatewaySettings, NeuronEndpoint,
};
use cortex_core::entitlements::{CapWindow, HEADER_ACCOUNT_ID, HEADER_KEY_ID};
use cortex_core::node::{ModelEntry, ModelStatus};
use cortex_gateway::state::CortexState;
use serde_json::{Value, json};
use std::sync::{Arc, Mutex};
use tokio::net::TcpListener;
/// What the mock neuron observed on the inbound `/v1/chat/completions`
/// request: the principal headers cortex stamped (or didn't).
#[derive(Default)]
struct Seen {
account_id: Option<String>,
key_id: Option<String>,
}
/// Spawn a mock neuron that records the principal headers it receives and
/// returns a trivial chat completion. Returns (base_url, observed).
async fn spawn_capturing_neuron() -> (String, Arc<Mutex<Seen>>) {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let base_url = format!("http://{addr}");
let inference_url = base_url.clone();
let seen: Arc<Mutex<Seen>> = Arc::new(Mutex::new(Seen::default()));
let sink = Arc::clone(&seen);
let app = axum::Router::new()
.route(
"/models/{model_id}/endpoint",
get(move |Path(_): Path<String>| {
let url = inference_url.clone();
async move { Json(json!({ "url": url })) }
}),
)
.route(
"/v1/chat/completions",
post(move |headers: HeaderMap, Json(body): Json<Value>| {
let sink = Arc::clone(&sink);
async move {
{
let mut s = sink.lock().unwrap();
s.account_id = headers
.get(HEADER_ACCOUNT_ID)
.and_then(|v| v.to_str().ok())
.map(str::to_string);
s.key_id = headers
.get(HEADER_KEY_ID)
.and_then(|v| v.to_str().ok())
.map(str::to_string);
}
let model = body.get("model").and_then(Value::as_str).unwrap_or("m");
Json(json!({
"id": "chatcmpl-auth-001",
"object": "chat.completion",
"created": 1700000000_u64,
"model": model,
"choices": [{
"index": 0,
"message": {"role": "assistant", "content": "ok"},
"finish_reason": "stop"
}],
"usage": {"prompt_tokens": 3, "completion_tokens": 1, "total_tokens": 4}
}))
}
}),
)
.with_state(());
tokio::spawn(async move {
axum::serve(listener, app).await.unwrap();
});
(base_url, seen)
}
/// Spawn a gateway with the given entitlements config, a single neuron, and
/// `test-model` seeded as loaded (build_app spawns no poller).
async fn spawn_gateway(neuron_url: &str, entitlements: EntitlementsConfig) -> String {
let config = GatewayConfig {
gateway: GatewaySettings {
listen: "127.0.0.1:0".into(),
metrics_listen: "127.0.0.1:0".into(),
},
eviction: EvictionSettings {
strategy: EvictionStrategy::Lru,
defrag_after_cycles: 0,
},
neurons: vec![NeuronEndpoint {
name: "mock-node".into(),
endpoint: neuron_url.to_string(),
}],
models_config: "/dev/null".into(),
entitlements,
};
let fleet = Arc::new(CortexState::from_config(&config));
{
let mut nodes = fleet.nodes.write().await;
let node = nodes.get_mut("mock-node").unwrap();
node.healthy = true;
node.models.insert(
"test-model".into(),
ModelEntry {
id: "test-model".into(),
status: ModelStatus::Loaded,
last_accessed: None,
vram_estimate_mb: Some(8000),
capabilities: Vec::new(),
tool_call: false,
reasoning: false,
limit: None,
},
);
}
let app = cortex_gateway::build_app(Arc::clone(&fleet));
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 one_key_config(require_auth: bool) -> EntitlementsConfig {
EntitlementsConfig {
require_auth,
keys: vec![ApiKeyConfig {
key: "sk-good".into(),
account_id: "acct-1".into(),
key_id: Some("key-1".into()),
hard_cap: None,
window: CapWindow::Balance,
}],
}
}
fn chat_body() -> Value {
json!({
"model": "test-model",
"messages": [{"role": "user", "content": "hi"}]
})
}
#[tokio::test]
async fn missing_key_when_required_is_401_invalid_api_key() {
let (neuron, _seen) = spawn_capturing_neuron().await;
let gateway = spawn_gateway(&neuron, one_key_config(true)).await;
let resp = reqwest::Client::new()
.post(format!("{gateway}/v1/chat/completions"))
.json(&chat_body())
.send()
.await
.unwrap();
assert_eq!(resp.status(), reqwest::StatusCode::UNAUTHORIZED);
let body: Value = resp.json().await.unwrap();
assert_eq!(body["error"]["code"], "invalid_api_key");
assert_eq!(body["error"]["type"], "invalid_request_error");
}
#[tokio::test]
async fn unrecognized_key_is_ignored_when_auth_not_required() {
let (neuron, seen) = spawn_capturing_neuron().await;
// allow-anonymous mode: a placeholder/unknown bearer (as opencode,
// Open WebUI, Agent Zero, litellm all send by default) must NOT be
// rejected — it's ignored and the request is served anonymously.
let gateway = spawn_gateway(&neuron, one_key_config(false)).await;
let resp = reqwest::Client::new()
.post(format!("{gateway}/v1/chat/completions"))
.bearer_auth("sk-dummy-placeholder")
.json(&chat_body())
.send()
.await
.unwrap();
assert_eq!(resp.status(), reqwest::StatusCode::OK);
let _ = resp.bytes().await.unwrap();
// Served, but anonymous — no principal stamped from the bogus key.
assert!(seen.lock().unwrap().account_id.is_none());
}
#[tokio::test]
async fn invalid_key_is_401_when_auth_required() {
let (neuron, seen) = spawn_capturing_neuron().await;
// With auth required, a present-but-wrong credential is rejected.
let gateway = spawn_gateway(&neuron, one_key_config(true)).await;
let resp = reqwest::Client::new()
.post(format!("{gateway}/v1/chat/completions"))
.bearer_auth("sk-wrong")
.json(&chat_body())
.send()
.await
.unwrap();
assert_eq!(resp.status(), reqwest::StatusCode::UNAUTHORIZED);
let body: Value = resp.json().await.unwrap();
assert_eq!(body["error"]["code"], "invalid_api_key");
// Rejected before dispatch — neuron never saw the request.
assert!(seen.lock().unwrap().account_id.is_none());
}
#[tokio::test]
async fn valid_key_reaches_neuron_with_principal_headers() {
let (neuron, seen) = spawn_capturing_neuron().await;
let gateway = spawn_gateway(&neuron, one_key_config(true)).await;
let resp = reqwest::Client::new()
.post(format!("{gateway}/v1/chat/completions"))
.bearer_auth("sk-good")
// A spoofed principal header must be stripped, not forwarded.
.header(HEADER_ACCOUNT_ID, "attacker")
.json(&chat_body())
.send()
.await
.unwrap();
assert_eq!(resp.status(), reqwest::StatusCode::OK);
let s = seen.lock().unwrap();
assert_eq!(s.account_id.as_deref(), Some("acct-1"));
assert_eq!(s.key_id.as_deref(), Some("key-1"));
}
#[tokio::test]
async fn anonymous_allowed_when_auth_not_required() {
let (neuron, seen) = spawn_capturing_neuron().await;
let gateway = spawn_gateway(&neuron, EntitlementsConfig::default()).await;
let resp = reqwest::Client::new()
.post(format!("{gateway}/v1/chat/completions"))
.json(&chat_body())
.send()
.await
.unwrap();
assert_eq!(resp.status(), reqwest::StatusCode::OK);
// No principal resolved → no principal headers stamped.
let s = seen.lock().unwrap();
assert!(s.account_id.is_none());
assert!(s.key_id.is_none());
}
#[tokio::test]
async fn health_is_public_even_when_auth_required() {
let (neuron, _seen) = spawn_capturing_neuron().await;
let gateway = spawn_gateway(&neuron, one_key_config(true)).await;
let resp = reqwest::Client::new()
.get(format!("{gateway}/health"))
.send()
.await
.unwrap();
assert_eq!(resp.status(), reqwest::StatusCode::OK);
}

View File

@@ -0,0 +1,253 @@
//! Integration tests for budget enforcement (#52) — the A0 seatbelt.
//!
//! A reservation over the key's hard cap is refused *before* neuron is hit,
//! with the #63 code matching the cap-window semantics (rate_limit_exceeded
//! + Retry-After for a resetting window, insufficient_quota for a hard
//! balance). Spend never exceeds the cap. No 402, ever.
use axum::Json;
use axum::extract::Path;
use axum::routing::{get, post};
use cortex_core::config::{
ApiKeyConfig, EntitlementsConfig, EvictionSettings, EvictionStrategy, GatewayConfig,
GatewaySettings, NeuronEndpoint,
};
use cortex_core::entitlements::{CapWindow, Principal};
use cortex_core::node::{ModelEntry, ModelStatus};
use cortex_gateway::state::CortexState;
use serde_json::{Value, json};
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use tokio::net::TcpListener;
/// Mock neuron with a hit counter on the inference path, so a test can prove
/// a request was (or wasn't) dispatched.
async fn spawn_counting_neuron() -> (String, Arc<AtomicU64>) {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let base_url = format!("http://{addr}");
let inference_url = base_url.clone();
let hits = Arc::new(AtomicU64::new(0));
let sink = Arc::clone(&hits);
let app = axum::Router::new()
.route(
"/models/{model_id}/endpoint",
get(move |Path(_): Path<String>| {
let url = inference_url.clone();
async move { Json(json!({ "url": url })) }
}),
)
.route(
"/v1/chat/completions",
post(move |Json(body): Json<Value>| {
let sink = Arc::clone(&sink);
async move {
sink.fetch_add(1, Ordering::SeqCst);
let model = body.get("model").and_then(Value::as_str).unwrap_or("m");
Json(json!({
"id": "chatcmpl-budget",
"object": "chat.completion",
"created": 1700000000_u64,
"model": model,
"choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}],
"usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}
}))
}
}),
);
tokio::spawn(async move {
axum::serve(listener, app).await.unwrap();
});
(base_url, hits)
}
async fn spawn_gateway(neuron_url: &str, key: ApiKeyConfig) -> (Arc<CortexState>, String) {
let config = GatewayConfig {
gateway: GatewaySettings {
listen: "127.0.0.1:0".into(),
metrics_listen: "127.0.0.1:0".into(),
},
eviction: EvictionSettings {
strategy: EvictionStrategy::Lru,
defrag_after_cycles: 0,
},
neurons: vec![NeuronEndpoint {
name: "mock-node".into(),
endpoint: neuron_url.to_string(),
}],
models_config: "/dev/null".into(),
entitlements: EntitlementsConfig {
require_auth: true,
keys: vec![key],
},
};
let fleet = Arc::new(CortexState::from_config(&config));
{
let mut nodes = fleet.nodes.write().await;
let node = nodes.get_mut("mock-node").unwrap();
node.healthy = true;
node.models.insert(
"test-model".into(),
ModelEntry {
id: "test-model".into(),
status: ModelStatus::Loaded,
last_accessed: None,
vram_estimate_mb: Some(8000),
capabilities: Vec::new(),
tool_call: false,
reasoning: false,
limit: None,
},
);
}
let app = cortex_gateway::build_app(Arc::clone(&fleet));
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();
});
(fleet, format!("http://{addr}"))
}
fn key(window: CapWindow, hard_cap: u64) -> ApiKeyConfig {
ApiKeyConfig {
key: "sk-cap".into(),
account_id: "acct-cap".into(),
key_id: Some("key-cap".into()),
hard_cap: Some(hard_cap),
window,
}
}
fn chat(max_tokens: u64) -> Value {
json!({
"model": "test-model",
"max_tokens": max_tokens,
"messages": [{"role": "user", "content": "hi"}]
})
}
#[tokio::test]
async fn balance_over_cap_is_429_insufficient_quota_before_dispatch() {
let (neuron, hits) = spawn_counting_neuron().await;
// Cap far below a single request's reservation (max_tokens 1000).
let (_fleet, gateway) = spawn_gateway(&neuron, key(CapWindow::Balance, 10)).await;
let resp = reqwest::Client::new()
.post(format!("{gateway}/v1/chat/completions"))
.bearer_auth("sk-cap")
.json(&chat(1000))
.send()
.await
.unwrap();
assert_eq!(resp.status(), reqwest::StatusCode::TOO_MANY_REQUESTS);
// Hard balance → no Retry-After.
assert!(resp.headers().get(reqwest::header::RETRY_AFTER).is_none());
let body: Value = resp.json().await.unwrap();
assert_eq!(body["error"]["code"], "insufficient_quota");
// Refused before dispatch — neuron never saw it.
assert_eq!(hits.load(Ordering::SeqCst), 0);
}
#[tokio::test]
async fn rolling_over_cap_is_429_rate_limited_with_retry_after() {
let (neuron, hits) = spawn_counting_neuron().await;
let (_fleet, gateway) =
spawn_gateway(&neuron, key(CapWindow::Rolling { seconds: 3600 }, 10)).await;
let resp = reqwest::Client::new()
.post(format!("{gateway}/v1/chat/completions"))
.bearer_auth("sk-cap")
.json(&chat(1000))
.send()
.await
.unwrap();
assert_eq!(resp.status(), reqwest::StatusCode::TOO_MANY_REQUESTS);
let retry = resp
.headers()
.get(reqwest::header::RETRY_AFTER)
.expect("rolling-window rejection must carry Retry-After");
assert!(retry.to_str().unwrap().parse::<u64>().unwrap() >= 1);
let body: Value = resp.json().await.unwrap();
assert_eq!(body["error"]["code"], "rate_limit_exceeded");
assert_eq!(hits.load(Ordering::SeqCst), 0);
}
#[tokio::test]
async fn within_cap_is_served() {
let (neuron, hits) = spawn_counting_neuron().await;
let (_fleet, gateway) = spawn_gateway(&neuron, key(CapWindow::Balance, 1_000_000)).await;
let resp = reqwest::Client::new()
.post(format!("{gateway}/v1/chat/completions"))
.bearer_auth("sk-cap")
.json(&chat(50))
.send()
.await
.unwrap();
assert_eq!(resp.status(), reqwest::StatusCode::OK);
let _ = resp.bytes().await.unwrap();
assert_eq!(hits.load(Ordering::SeqCst), 1);
}
#[tokio::test]
async fn a0_seatbelt_caps_a_runaway_fan_out() {
// An Agent-Zero-style key with a modest cap: a burst of requests drains
// it, then further requests are refused — the account stops draining and
// spend never exceeds the cap.
let (neuron, hits) = spawn_counting_neuron().await;
let (fleet, gateway) = spawn_gateway(&neuron, key(CapWindow::Balance, 100)).await;
let client = reqwest::Client::new();
let mut ok = 0;
let mut refused = 0;
for _ in 0..20 {
let resp = client
.post(format!("{gateway}/v1/chat/completions"))
.bearer_auth("sk-cap")
.json(&chat(20))
.send()
.await
.unwrap();
match resp.status() {
reqwest::StatusCode::OK => {
ok += 1;
let _ = resp.bytes().await.unwrap();
}
reqwest::StatusCode::TOO_MANY_REQUESTS => {
refused += 1;
let body: Value = resp.json().await.unwrap();
assert_eq!(body["error"]["code"], "insufficient_quota");
}
other => panic!("unexpected status {other}"),
}
}
assert!(ok >= 1, "some requests should be served");
assert!(refused >= 1, "the cap must eventually refuse the fan-out");
assert_eq!(
hits.load(Ordering::SeqCst),
ok,
"refused requests never dispatched"
);
// Spend never exceeded the hard cap (reservation prevents overshoot).
// Poll briefly for in-flight settles to land.
let principal = Principal {
account_id: "acct-cap".into(),
key_id: "key-cap".into(),
};
for _ in 0..50 {
let snap = fleet.entitlements.snapshot(&principal).await.unwrap();
if snap.reserved == 0 {
break;
}
tokio::time::sleep(std::time::Duration::from_millis(20)).await;
}
let snap = fleet.entitlements.snapshot(&principal).await.unwrap();
assert!(snap.spent <= 100, "spent {} exceeded cap", snap.spent);
}

View File

@@ -0,0 +1,124 @@
//! Router: a catalogued model whose only topologically-feasible neuron is
//! currently unhealthy is a *transient* condition (retryable 503), not a
//! permanent 404. This is the exact shape of the beast incident: benjy/
//! quadbrat (1 GPU, healthy) can't host the 27B, and beast (2 GPU) — the
//! sole feasible node — briefly drops out → clients must back off and retry,
//! not hard-fail.
use cortex_core::config::{
EvictionSettings, EvictionStrategy, GatewayConfig, GatewaySettings, NeuronEndpoint,
};
use cortex_core::discovery::{DeviceInfo, DiscoveryResponse};
use cortex_gateway::router::{self, RouteError};
use cortex_gateway::state::CortexState;
use std::sync::Arc;
fn devices(n: usize) -> Vec<DeviceInfo> {
(0..n)
.map(|i| DeviceInfo {
index: i as u32,
name: "RTX 5090".into(),
vram_total_mb: 32_768,
compute_capability: "9.0".into(),
})
.collect()
}
fn discovery(host: &str, n_devices: usize) -> DiscoveryResponse {
DiscoveryResponse {
hostname: host.into(),
os: "Linux".into(),
kernel: "7.0".into(),
cuda_version: Some("13.0".into()),
driver_version: Some("999".into()),
devices: devices(n_devices),
harnesses: vec!["candle".into()],
cuda_unavailable_reason: None,
max_prompt_tokens: 49_152,
}
}
/// Catalogue with one model needing 2 devices. Returns a temp path.
fn write_catalogue() -> std::path::PathBuf {
let toml = r#"
[[models]]
id = "big-model"
harness = "candle"
min_devices = 2
"#;
let path = std::env::temp_dir().join("cortex_test_feasibility_models.toml");
std::fs::write(&path, toml).unwrap();
path
}
async fn fleet_with(big_healthy: bool, big_devices: usize) -> Arc<CortexState> {
let cat = write_catalogue();
let config = GatewayConfig {
gateway: GatewaySettings {
listen: "127.0.0.1:0".into(),
metrics_listen: "127.0.0.1:0".into(),
},
eviction: EvictionSettings {
strategy: EvictionStrategy::Lru,
defrag_after_cycles: 0,
},
neurons: vec![
NeuronEndpoint {
name: "small".into(),
endpoint: "http://127.0.0.1:1".into(),
},
NeuronEndpoint {
name: "big".into(),
endpoint: "http://127.0.0.1:2".into(),
},
],
models_config: cat.to_string_lossy().into_owned(),
entitlements: Default::default(),
};
let fleet = Arc::new(CortexState::from_config(&config));
{
let mut nodes = fleet.nodes.write().await;
// "small" is healthy but only has 1 GPU → not feasible for the model.
let small = nodes.get_mut("small").unwrap();
small.healthy = true;
small.discovery = Some(discovery("small", 1));
// "big" has enough GPUs but its health is the variable under test.
let big = nodes.get_mut("big").unwrap();
big.healthy = big_healthy;
big.discovery = Some(discovery("big", big_devices));
}
fleet
}
#[tokio::test]
async fn feasible_node_unhealthy_is_transient_503() {
// big (2 GPU, the only feasible node) is unhealthy; small (1 GPU) is
// healthy but can't host the model → retryable, not a permanent 404.
let fleet = fleet_with(false, 2).await;
let err = router::resolve(&fleet, "big-model")
.await
.expect_err("model can't be served right now");
assert!(
matches!(err, RouteError::FeasibleNodeUnhealthy { .. }),
"expected FeasibleNodeUnhealthy, got {err:?}"
);
assert_eq!(err.http_status(), 503);
assert_eq!(err.retry_after_secs(), Some(3));
assert_eq!(err.code(), "service_unavailable");
}
#[tokio::test]
async fn no_node_can_ever_satisfy_is_permanent_404() {
// big is healthy but only has 1 GPU now (e.g. topology genuinely can't
// satisfy min_devices=2 anywhere) → permanent, non-retryable 404.
let fleet = fleet_with(true, 1).await;
let err = router::resolve(&fleet, "big-model")
.await
.expect_err("no feasible topology");
assert!(
matches!(err, RouteError::NoFeasibleNeuron { .. }),
"expected NoFeasibleNeuron, got {err:?}"
);
assert_eq!(err.http_status(), 404);
assert_eq!(err.retry_after_secs(), None);
}

View File

@@ -0,0 +1,189 @@
//! Load-aware routing across replicas (#55).
//!
//! When a model is loaded on more than one healthy neuron, the router picks
//! the least-busy replica using the per-model admission load each neuron
//! reports on `GET /health` (#53), rather than always taking the first.
mod common;
use axum::Json;
use axum::extract::Path;
use axum::http::{StatusCode, header};
use axum::response::IntoResponse;
use axum::routing::{get, post};
use cortex_core::config::{
EvictionSettings, EvictionStrategy, GatewayConfig, GatewaySettings, NeuronEndpoint,
};
use cortex_core::discovery::ModelLoad;
use cortex_core::node::{ModelEntry, ModelStatus};
use cortex_gateway::state::CortexState;
use serde_json::{Value, json};
use std::sync::Arc;
use tokio::net::TcpListener;
/// Seed a node as healthy with `test-model` loaded and a given admission load.
async fn seed_loaded(fleet: &CortexState, node: &str, in_flight: usize, queue_depth: usize) {
let mut nodes = fleet.nodes.write().await;
let n = nodes.get_mut(node).expect("node exists");
n.healthy = true;
n.models.insert(
"test-model".into(),
ModelEntry {
id: "test-model".into(),
status: ModelStatus::Loaded,
last_accessed: None,
vram_estimate_mb: Some(8000),
capabilities: Vec::new(),
tool_call: false,
reasoning: false,
limit: None,
},
);
n.model_load.insert(
"test-model".into(),
ModelLoad {
id: "test-model".into(),
in_flight,
queue_depth,
},
);
}
/// Build a gateway state over two mock neurons (no poller; we seed state).
async fn two_neuron_fleet(endpoint_a: &str, endpoint_b: &str) -> Arc<CortexState> {
let config = GatewayConfig {
gateway: GatewaySettings {
listen: "127.0.0.1:0".into(),
metrics_listen: "127.0.0.1:0".into(),
},
eviction: EvictionSettings {
strategy: EvictionStrategy::Lru,
defrag_after_cycles: 0,
},
neurons: vec![
NeuronEndpoint {
name: "node-a".into(),
endpoint: endpoint_a.to_string(),
},
NeuronEndpoint {
name: "node-b".into(),
endpoint: endpoint_b.to_string(),
},
],
models_config: "/dev/null".into(),
entitlements: Default::default(),
};
Arc::new(CortexState::from_config(&config))
}
#[tokio::test]
async fn routes_to_least_busy_replica() {
let neuron_a = common::spawn_mock_neuron().await;
let neuron_b = common::spawn_mock_neuron().await;
let fleet = two_neuron_fleet(&neuron_a, &neuron_b).await;
// A is busy (1 running + 3 queued), B is idle.
seed_loaded(&fleet, "node-a", 1, 3).await;
seed_loaded(&fleet, "node-b", 0, 0).await;
let route = cortex_gateway::router::resolve(&fleet, "test-model")
.await
.expect("model is loaded on both nodes");
assert_eq!(route.node_name, "node-b", "should pick the idle replica");
// Flip the load: now B is the busy one.
seed_loaded(&fleet, "node-a", 0, 0).await;
seed_loaded(&fleet, "node-b", 1, 5).await;
let route = cortex_gateway::router::resolve(&fleet, "test-model")
.await
.expect("still loaded");
assert_eq!(route.node_name, "node-a", "should follow the lighter load");
}
/// Mock neuron whose inference endpoint always returns a #63 backpressure
/// envelope (503 + Retry-After) — simulating a saturated neuron.
async fn spawn_busy_neuron() -> String {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let base_url = format!("http://{addr}");
let inference_url = base_url.clone();
let app = axum::Router::new()
.route(
"/models/{model_id}/endpoint",
get(move |Path(_): Path<String>| {
let url = inference_url.clone();
async move { Json(json!({ "url": url })) }
}),
)
.route(
"/v1/chat/completions",
post(|| async {
let body = json!({"error": {
"message": "model is busy (admission queue full); retry shortly",
"type": "rate_limit_error",
"code": "rate_limit_exceeded",
"param": null
}});
(
StatusCode::SERVICE_UNAVAILABLE,
[(header::RETRY_AFTER, "6")],
Json(body),
)
.into_response()
}),
);
tokio::spawn(async move {
axum::serve(listener, app).await.unwrap();
});
base_url
}
#[tokio::test]
async fn neuron_backpressure_is_propagated_intact() {
// A saturated neuron's 503 + Retry-After + envelope must reach the client
// verbatim — not unwrapped, remapped, or stripped (#55 / #63).
let neuron = spawn_busy_neuron().await;
let fleet = two_neuron_fleet(&neuron, &neuron).await;
seed_loaded(&fleet, "node-a", 1, 8).await;
let app = cortex_gateway::build_app(Arc::clone(&fleet));
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();
});
let resp = reqwest::Client::new()
.post(format!("http://{addr}/v1/chat/completions"))
.json(&json!({"model": "test-model", "messages": [{"role": "user", "content": "hi"}]}))
.send()
.await
.unwrap();
assert_eq!(resp.status(), reqwest::StatusCode::SERVICE_UNAVAILABLE);
assert_eq!(
resp.headers()
.get(reqwest::header::RETRY_AFTER)
.and_then(|v| v.to_str().ok()),
Some("6"),
"Retry-After must survive the proxy"
);
let body: Value = resp.json().await.unwrap();
assert_eq!(body["error"]["code"], "rate_limit_exceeded");
}
#[tokio::test]
async fn ties_break_deterministically_by_name() {
let neuron_a = common::spawn_mock_neuron().await;
let neuron_b = common::spawn_mock_neuron().await;
let fleet = two_neuron_fleet(&neuron_a, &neuron_b).await;
// Equal load on both → stable pick (lowest node name).
seed_loaded(&fleet, "node-a", 0, 0).await;
seed_loaded(&fleet, "node-b", 0, 0).await;
let route = cortex_gateway::router::resolve(&fleet, "test-model")
.await
.expect("loaded");
assert_eq!(route.node_name, "node-a", "ties break by name");
}

View File

@@ -0,0 +1,207 @@
//! Integration tests for per-request token metering (#51).
//!
//! Drives authenticated requests through the gateway to a mock neuron that
//! reports a fixed `usage` object, then asserts the EntitlementProvider's
//! spend ledger reflects cumulative per-key spend and that reservations
//! settle to actual (no outstanding reserved tokens once requests complete).
mod common;
use cortex_core::config::{
ApiKeyConfig, EntitlementsConfig, EvictionSettings, EvictionStrategy, GatewayConfig,
GatewaySettings, NeuronEndpoint,
};
use cortex_core::entitlements::{CapWindow, Principal};
use cortex_core::node::{ModelEntry, ModelStatus};
use cortex_gateway::state::CortexState;
use serde_json::json;
use std::sync::Arc;
use std::time::Duration;
use tokio::net::TcpListener;
const ACCOUNT: &str = "acct-meter";
const KEY_ID: &str = "key-meter";
const BEARER: &str = "sk-meter";
/// The mock neuron (common::spawn_mock_neuron) reports this fixed usage on
/// every chat completion.
const PROMPT_PER_REQ: u64 = 10;
const COMPLETION_PER_REQ: u64 = 5;
async fn spawn_metered_gateway(neuron_url: &str) -> (Arc<CortexState>, String) {
let config = GatewayConfig {
gateway: GatewaySettings {
listen: "127.0.0.1:0".into(),
metrics_listen: "127.0.0.1:0".into(),
},
eviction: EvictionSettings {
strategy: EvictionStrategy::Lru,
defrag_after_cycles: 0,
},
neurons: vec![NeuronEndpoint {
name: "mock-node".into(),
endpoint: neuron_url.to_string(),
}],
models_config: "/dev/null".into(),
entitlements: EntitlementsConfig {
require_auth: true,
keys: vec![ApiKeyConfig {
key: BEARER.into(),
account_id: ACCOUNT.into(),
key_id: Some(KEY_ID.into()),
hard_cap: Some(1_000_000),
window: CapWindow::Balance,
}],
},
};
let fleet = Arc::new(CortexState::from_config(&config));
{
let mut nodes = fleet.nodes.write().await;
let node = nodes.get_mut("mock-node").unwrap();
node.healthy = true;
node.models.insert(
"test-model".into(),
ModelEntry {
id: "test-model".into(),
status: ModelStatus::Loaded,
last_accessed: None,
vram_estimate_mb: Some(8000),
capabilities: Vec::new(),
tool_call: false,
reasoning: false,
limit: None,
},
);
}
let app = cortex_gateway::build_app(Arc::clone(&fleet));
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();
});
(fleet, format!("http://{addr}"))
}
fn principal() -> Principal {
Principal {
account_id: ACCOUNT.into(),
key_id: KEY_ID.into(),
}
}
/// Poll the provider ledger until settled spend reaches `expected` (settle
/// runs in a spawned task after the response stream finishes) or time out.
async fn await_spent(fleet: &CortexState, expected: u64) -> u64 {
let principal = principal();
for _ in 0..100 {
let snap = fleet.entitlements.snapshot(&principal).await.unwrap();
if snap.spent >= expected {
return snap.spent;
}
tokio::time::sleep(Duration::from_millis(20)).await;
}
fleet.entitlements.snapshot(&principal).await.unwrap().spent
}
#[tokio::test]
async fn cumulative_spend_is_metered_per_key() {
let neuron = common::spawn_mock_neuron().await;
let (fleet, gateway) = spawn_metered_gateway(&neuron).await;
let client = reqwest::Client::new();
const N: u64 = 3;
for _ in 0..N {
let resp = client
.post(format!("{gateway}/v1/chat/completions"))
.bearer_auth(BEARER)
.json(&json!({"model": "test-model", "messages": [{"role": "user", "content": "hi"}]}))
.send()
.await
.unwrap();
assert_eq!(resp.status(), reqwest::StatusCode::OK);
// Drain the body so the response stream finishes and metering settles.
let _ = resp.bytes().await.unwrap();
}
let expected = N * (PROMPT_PER_REQ + COMPLETION_PER_REQ);
let spent = await_spent(&fleet, expected).await;
assert_eq!(
spent, expected,
"ledger must reflect cumulative per-key spend"
);
// Reservations settled to actual — nothing left outstanding.
let snap = fleet.entitlements.snapshot(&principal()).await.unwrap();
assert_eq!(snap.reserved, 0, "all reservations must settle/release");
assert_eq!(snap.hard_cap, Some(1_000_000));
}
#[tokio::test]
async fn anonymous_request_records_no_spend() {
// require_auth=false so the unauthenticated request is served, but with
// no principal it must not touch any ledger.
let neuron = common::spawn_mock_neuron().await;
let config = GatewayConfig {
gateway: GatewaySettings {
listen: "127.0.0.1:0".into(),
metrics_listen: "127.0.0.1:0".into(),
},
eviction: EvictionSettings {
strategy: EvictionStrategy::Lru,
defrag_after_cycles: 0,
},
neurons: vec![NeuronEndpoint {
name: "mock-node".into(),
endpoint: neuron.clone(),
}],
models_config: "/dev/null".into(),
entitlements: EntitlementsConfig::default(),
};
let fleet = Arc::new(CortexState::from_config(&config));
{
let mut nodes = fleet.nodes.write().await;
let node = nodes.get_mut("mock-node").unwrap();
node.healthy = true;
node.models.insert(
"test-model".into(),
ModelEntry {
id: "test-model".into(),
status: ModelStatus::Loaded,
last_accessed: None,
vram_estimate_mb: Some(8000),
capabilities: Vec::new(),
tool_call: false,
reasoning: false,
limit: None,
},
);
}
let app = cortex_gateway::build_app(Arc::clone(&fleet));
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();
});
let resp = reqwest::Client::new()
.post(format!("http://{addr}/v1/chat/completions"))
.json(&json!({"model": "test-model", "messages": [{"role": "user", "content": "hi"}]}))
.send()
.await
.unwrap();
assert_eq!(resp.status(), reqwest::StatusCode::OK);
let _ = resp.bytes().await.unwrap();
// An unconfigured principal has a zeroed snapshot — nothing was metered.
let snap = fleet
.entitlements
.snapshot(&Principal {
account_id: "nobody".into(),
key_id: "nobody".into(),
})
.await
.unwrap();
assert_eq!(snap.spent, 0);
}

View File

@@ -0,0 +1,131 @@
//! Issue #68: the `cost` wire contract on `GET /v1/models`.
//!
//! `cost` is operator-set pricing sourced from the `models.toml` catalogue
//! profile (the source of truth today; the marketplace clearing house #59
//! later — both must read the same value metering/#51 bills against). The
//! shape is the models.dev/opencode convention: **USD per 1,000,000 tokens,
//! as JSON numbers**, with optional `cache_read`/`cache_write` tiers. This
//! test pins:
//! - the units/shape (per-million floats, not per-token, not strings);
//! - that cache fields flow through when present and are omitted otherwise;
//! - the load-bearing **absent vs `0.0`** distinction (#68): a model with
//! no catalogue `cost` omits the key entirely (price unknown), distinct
//! from an explicit `0.0` (intentionally free).
//!
//! Catalogue-only models surface via Pass 1 of `list_models` even with no
//! feasible neuron, so this is hermetic — no nodes or poller needed.
use cortex_core::config::{
EvictionSettings, EvictionStrategy, GatewayConfig, GatewaySettings, NeuronEndpoint,
};
use cortex_gateway::state::CortexState;
use std::sync::Arc;
use tokio::net::TcpListener;
#[tokio::test]
async fn v1_models_cost_units_shape_and_absent_vs_zero() {
// Three catalogue models exercise the whole contract: a priced model
// with cache tiers, an intentionally-free model (explicit 0.0), and an
// unpriced model (no `cost` block at all).
let models_toml = r#"
[[models]]
id = "priced-model"
harness = "candle"
cost.input = 0.5
cost.output = 1.5
cost.cache_read = 0.05
cost.cache_write = 0.6
[[models]]
id = "free-model"
harness = "candle"
cost.input = 0.0
cost.output = 0.0
[[models]]
id = "unpriced-model"
harness = "candle"
"#;
let cat_path = std::env::temp_dir().join("cortex_test_issue68_models.toml");
std::fs::write(&cat_path, models_toml).unwrap();
let config = GatewayConfig {
gateway: GatewaySettings {
listen: "127.0.0.1:0".into(),
metrics_listen: "127.0.0.1:0".into(),
},
eviction: EvictionSettings {
strategy: EvictionStrategy::Lru,
defrag_after_cycles: 0,
},
// Never contacted: build_app does not spawn the poller, so the
// catalogue alone drives /v1/models.
neurons: vec![NeuronEndpoint {
name: "mock-node".into(),
endpoint: "http://127.0.0.1:1".into(),
}],
models_config: cat_path.to_string_lossy().into_owned(),
entitlements: Default::default(),
};
let fleet = Arc::new(CortexState::from_config(&config));
let app = cortex_gateway::build_app(Arc::clone(&fleet));
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();
});
let body: serde_json::Value = reqwest::Client::new()
.get(format!("http://{addr}/v1/models"))
.send()
.await
.unwrap()
.json()
.await
.unwrap();
let data = body["data"].as_array().expect("data is an array");
let entry = |id: &str| {
data.iter()
.find(|m| m["id"] == id)
.unwrap_or_else(|| panic!("{id} present in /v1/models"))
.clone()
};
// Priced model: exact values flow through as JSON numbers (USD per 1M
// tokens). If anything rescaled by 10⁶ or stringified, these fail.
let priced = entry("priced-model");
assert_eq!(priced["cost"]["input"], 0.5);
assert_eq!(priced["cost"]["output"], 1.5);
assert_eq!(priced["cost"]["cache_read"], 0.05);
assert_eq!(priced["cost"]["cache_write"], 0.6);
assert!(
priced["cost"]["input"].is_number(),
"cost.input must be a JSON number, not a string"
);
// Intentionally free: cost present, rates explicitly 0.0. Unset cache
// tiers are omitted (skip_serializing_if), not emitted as null/0.
let free = entry("free-model");
assert_eq!(free["cost"]["input"], 0.0);
assert_eq!(free["cost"]["output"], 0.0);
assert!(
free["cost"].get("cache_read").is_none(),
"absent cache tiers must be omitted, not null"
);
assert!(free["cost"].get("cache_write").is_none());
// Unpriced: the whole `cost` object is omitted — "price unknown",
// distinct from the free model's explicit 0.0. This is the #68
// distinction opencode needs to avoid showing $0 for a model whose
// price simply hasn't been declared.
let unpriced = entry("unpriced-model");
assert!(
unpriced.get("cost").is_none(),
"a model with no catalogue cost must omit `cost` entirely, got {:?}",
unpriced.get("cost")
);
let _ = std::fs::remove_file(&cat_path);
}

View File

@@ -228,10 +228,26 @@ async fn test_poller_marks_unreachable_node_unhealthy() {
nodes.get_mut("dead-node").unwrap().healthy = true;
}
// Debounce (#53 follow-up): a single missed poll must NOT evict a
// previously-healthy node — a busy neuron briefly slow to answer
// shouldn't yank its models out of routing.
cortex_gateway::poller::poll_once(&fleet).await;
assert!(
fleet.nodes.read().await.get("dead-node").unwrap().healthy,
"one failed poll should not mark a healthy node unhealthy"
);
let nodes = fleet.nodes.read().await;
assert!(!nodes.get("dead-node").unwrap().healthy);
// It flips unhealthy only after POLL_FAILURE_THRESHOLD (3) consecutive
// failures.
cortex_gateway::poller::poll_once(&fleet).await;
cortex_gateway::poller::poll_once(&fleet).await;
assert!(
!fleet.nodes.read().await.get("dead-node").unwrap().healthy,
"three consecutive failed polls should mark the node unhealthy"
);
// A subsequent successful poll would reset the counter and restore
// health; covered implicitly by the discovery tests above.
}
#[tokio::test]

View File

@@ -0,0 +1,174 @@
//! Fail-fast prompt pre-validation + advisory client hints (#56).
//!
//! cortex refuses a prompt that already exceeds the model's advertised
//! context window before dispatching to neuron — the same #60
//! `context_length_exceeded` envelope neuron would emit, just earlier — and
//! attaches an advisory `X-Helexa-Advice` header for fingerprinted clients.
use axum::Json;
use axum::extract::Path;
use axum::routing::{get, post};
use cortex_core::config::{
EvictionSettings, EvictionStrategy, GatewayConfig, GatewaySettings, NeuronEndpoint,
};
use cortex_core::harness::ModelLimit;
use cortex_core::node::{ModelEntry, ModelStatus};
use cortex_gateway::state::CortexState;
use serde_json::{Value, json};
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use tokio::net::TcpListener;
/// Mock neuron with a hit counter, so a test can prove a request was (or
/// wasn't) dispatched past the gateway's pre-validation.
async fn spawn_counting_neuron() -> (String, Arc<AtomicU64>) {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let base_url = format!("http://{addr}");
let inference_url = base_url.clone();
let hits = Arc::new(AtomicU64::new(0));
let sink = Arc::clone(&hits);
let app = axum::Router::new()
.route(
"/models/{model_id}/endpoint",
get(move |Path(_): Path<String>| {
let url = inference_url.clone();
async move { Json(json!({ "url": url })) }
}),
)
.route(
"/v1/chat/completions",
post(move || {
let sink = Arc::clone(&sink);
async move {
sink.fetch_add(1, Ordering::SeqCst);
Json(json!({
"id": "c", "object": "chat.completion", "created": 1_700_000_000_u64,
"model": "test-model",
"choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}],
"usage": {"prompt_tokens": 3, "completion_tokens": 1, "total_tokens": 4}
}))
}
}),
);
tokio::spawn(async move {
axum::serve(listener, app).await.unwrap();
});
(base_url, hits)
}
/// Gateway over one neuron with `test-model` loaded and a tiny advertised
/// context window (so a modest prompt overflows it).
async fn spawn_gateway(neuron: &str, context: usize) -> String {
let config = GatewayConfig {
gateway: GatewaySettings {
listen: "127.0.0.1:0".into(),
metrics_listen: "127.0.0.1:0".into(),
},
eviction: EvictionSettings {
strategy: EvictionStrategy::Lru,
defrag_after_cycles: 0,
},
neurons: vec![NeuronEndpoint {
name: "mock-node".into(),
endpoint: neuron.to_string(),
}],
models_config: "/dev/null".into(),
entitlements: Default::default(),
};
let fleet = Arc::new(CortexState::from_config(&config));
{
let mut nodes = fleet.nodes.write().await;
let n = nodes.get_mut("mock-node").unwrap();
n.healthy = true;
n.models.insert(
"test-model".into(),
ModelEntry {
id: "test-model".into(),
status: ModelStatus::Loaded,
last_accessed: None,
vram_estimate_mb: Some(8000),
capabilities: Vec::new(),
tool_call: false,
reasoning: false,
limit: Some(ModelLimit {
context,
input: None,
output: 16,
}),
},
);
}
let app = cortex_gateway::build_app(Arc::clone(&fleet));
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}")
}
#[tokio::test]
async fn over_long_prompt_is_rejected_before_dispatch() {
let (neuron, hits) = spawn_counting_neuron().await;
let gateway = spawn_gateway(&neuron, 50).await; // tiny 50-token window
// ~1200 chars → ~300 est tokens, well over 50.
let big = "word ".repeat(240);
let resp = reqwest::Client::new()
.post(format!("{gateway}/v1/chat/completions"))
.header("user-agent", "litellm/1.0")
.json(&json!({"model": "test-model", "messages": [{"role": "user", "content": big}]}))
.send()
.await
.unwrap();
assert_eq!(resp.status(), reqwest::StatusCode::BAD_REQUEST);
// Advisory hint for the fingerprinted client (header only, never body).
assert!(
resp.headers().get("x-helexa-advice").is_some(),
"litellm should get advice"
);
let body: Value = resp.json().await.unwrap();
assert_eq!(body["error"]["code"], "context_length_exceeded");
assert_eq!(body["error"]["max"], 50);
// Refused at the edge — neuron never saw it.
assert_eq!(hits.load(Ordering::SeqCst), 0);
}
#[tokio::test]
async fn within_context_passes_through() {
let (neuron, hits) = spawn_counting_neuron().await;
let gateway = spawn_gateway(&neuron, 4096).await;
let resp = reqwest::Client::new()
.post(format!("{gateway}/v1/chat/completions"))
.json(&json!({"model": "test-model", "messages": [{"role": "user", "content": "hi"}]}))
.send()
.await
.unwrap();
assert_eq!(resp.status(), reqwest::StatusCode::OK);
let _ = resp.bytes().await.unwrap();
assert_eq!(hits.load(Ordering::SeqCst), 1, "served by neuron");
}
#[tokio::test]
async fn unknown_client_gets_no_advice_header() {
let (neuron, _hits) = spawn_counting_neuron().await;
let gateway = spawn_gateway(&neuron, 50).await;
let big = "word ".repeat(240);
let resp = reqwest::Client::new()
.post(format!("{gateway}/v1/chat/completions"))
// no/unknown User-Agent → no advice, but still a clean 400
.json(&json!({"model": "test-model", "messages": [{"role": "user", "content": big}]}))
.send()
.await
.unwrap();
assert_eq!(resp.status(), reqwest::StatusCode::BAD_REQUEST);
assert!(resp.headers().get("x-helexa-advice").is_none());
let body: Value = resp.json().await.unwrap();
assert_eq!(body["error"]["code"], "context_length_exceeded");
}

View File

@@ -0,0 +1,34 @@
[package]
name = "helexa-router"
version.workspace = true
edition.workspace = true
license.workspace = true
repository.workspace = true
[[bin]]
name = "helexa-router"
path = "src/main.rs"
[lib]
name = "helexa_router"
path = "src/lib.rs"
[dependencies]
cortex-core = { workspace = true }
tokio = { workspace = true }
axum = { workspace = true }
tower-http = { workspace = true }
reqwest = { workspace = true }
serde = { workspace = true }
serde_json = { workspace = true }
figment = { workspace = true }
anyhow = { 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"] }

View File

@@ -0,0 +1,75 @@
use figment::{
Figment,
providers::{Env, Format, Toml},
};
use serde::{Deserialize, Serialize};
use std::path::Path;
/// Top-level `helexa-router` configuration.
///
/// Loaded from TOML with `HELEXA_ROUTER_`-prefixed env overrides (using
/// `__` as the nesting separator), matching the cortex/neuron convention.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RouterConfig {
pub router: RouterSettings,
/// Downstream cortex endpoints the router can dispatch to. The skeleton
/// (#70) only loads these; capacity/catalogue polling (#72) and
/// capacity-aware dispatch (#73) consume them later.
#[serde(default)]
pub cortexes: Vec<CortexEndpoint>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RouterSettings {
/// Address to listen on for the inbound API (e.g. "0.0.0.0:8088").
///
/// Plaintext only — operator/edge nginx terminates client TLS in front
/// 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,
}
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)]
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,
}
impl RouterConfig {
/// Load configuration from a TOML file, with environment variable
/// overrides prefixed with `HELEXA_ROUTER_` and `__` as the separator
/// (e.g. `HELEXA_ROUTER_ROUTER__LISTEN=0.0.0.0:8088`).
pub fn load(path: impl AsRef<Path>) -> Result<Self, Box<figment::Error>> {
Figment::new()
.merge(Toml::file(path))
.merge(Env::prefixed("HELEXA_ROUTER_").split("__"))
.extract()
.map_err(Box::new)
}
}
impl Default for RouterConfig {
fn default() -> Self {
Self {
router: RouterSettings {
listen: "0.0.0.0:8088".into(),
poll_interval_secs: default_poll_interval_secs(),
},
cortexes: vec![],
}
}
}

View File

@@ -0,0 +1,43 @@
use crate::state::RouterState;
use axum::{Json, Router, extract::State, routing::get};
use cortex_core::openai::ModelsResponse;
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`.
pub fn api_routes() -> Router<Arc<RouterState>> {
Router::new()
.route("/v1/models", get(list_models))
.route("/health", get(health))
.route("/", get(health))
}
/// `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![],
})
}

View File

@@ -0,0 +1,57 @@
//! helexa-router — public multi-operator ingress proxy (router.helexa.ai).
//!
//! The router is the data-plane *ingress* tier: a geo-distributed,
//! capacity-aware, OpenAI/Anthropic-compatible reverse proxy in front of
//! many operator-run cortexes ("cortex-of-cortexes"). End users configure
//! one `baseURL` and the router forwards their request to a cortex with
//! capacity, proxying #63-shaped rejections back verbatim.
//!
//! 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). A background [`poller`] keeps a live
//! per-cortex topology (#72) that the dispatcher (#73) will route on.
pub mod config;
pub mod handlers;
pub mod poller;
pub mod state;
use anyhow::Result;
use config::RouterConfig;
use std::sync::Arc;
use tower_http::cors::CorsLayer;
use tower_http::trace::TraceLayer;
/// Build the axum application: handlers + CORS + tracing. No auth layer —
/// the router asserts no identity of its own and forwards the client bearer
/// to the downstream cortex, which authenticates it (#69).
pub fn build_app(state: Arc<state::RouterState>) -> axum::Router {
axum::Router::new()
.merge(handlers::api_routes())
.layer(CorsLayer::permissive())
.layer(TraceLayer::new_for_http())
.with_state(state)
}
/// Start the router: build state from config and bind the plaintext HTTP
/// 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));
// 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}");
let listener = tokio::net::TcpListener::bind(listen_addr).await?;
axum::serve(listener, app).await?;
Ok(())
}

View File

@@ -0,0 +1,52 @@
use anyhow::Result;
use clap::{Parser, Subcommand};
use helexa_router::config::RouterConfig;
use tracing_subscriber::EnvFilter;
#[derive(Parser)]
#[command(name = "helexa-router")]
#[command(about = "Public multi-operator ingress proxy for helexa")]
#[command(version)]
struct Cli {
#[command(subcommand)]
command: Commands,
}
#[derive(Subcommand)]
enum Commands {
/// Start the router server.
Serve {
/// Path to the router config file.
#[arg(short, long, default_value = "helexa-router.toml")]
config: 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_router=debug")),
)
.init();
let cli = Cli::parse();
match cli.command {
Commands::Serve { config } => {
let cfg = RouterConfig::load(&config)
.map_err(|e| anyhow::anyhow!("failed to load config from '{config}': {e}"))?;
tracing::info!(
cortexes = cfg.cortexes.len(),
listen = %cfg.router.listen,
"starting helexa-router"
);
helexa_router::run(cfg).await?;
}
}
Ok(())
}

View File

@@ -0,0 +1,157 @@
//! 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::{RouterModelStatus, 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) {
let models = fetch_models(state, 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| {
let feasible = m.loaded || !m.feasible_on.is_empty();
(
m.id,
RouterModelStatus {
loaded: m.loaded,
feasible,
},
)
})
.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(state, 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(
state: &RouterState,
endpoint: &str,
) -> Result<Vec<CortexModelEntry>, &'static str> {
let url = format!("{endpoint}/v1/models");
let resp = state
.http_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(state: &RouterState, endpoint: &str) -> Option<(u32, u32)> {
let url = format!("{endpoint}/health");
let resp = state
.http_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))
}

View File

@@ -0,0 +1,85 @@
use crate::config::{CortexEndpoint, RouterConfig};
use chrono::{DateTime, Utc};
use std::collections::HashMap;
use std::time::Duration;
use tokio::sync::RwLock;
/// Shared router state: the configured cortex list plus the live topology
/// map the poller (#72) maintains and the dispatcher (#73) will route on.
///
/// 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>,
/// Shared client for polling (and, later, proxying to) cortexes.
pub http_client: reqwest::Client,
/// 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,
/// Per-model serveability, keyed by model id, from `/v1/models`.
pub models: HashMap<String, RouterModelStatus>,
}
/// What a cortex can do with one model, distilled from its `/v1/models`
/// entry to the two facts the router routes on.
#[derive(Debug, Clone)]
pub struct RouterModelStatus {
/// The model is loaded on at least one of the cortex's neurons.
pub loaded: bool,
/// The cortex can serve it — loaded now, or feasible to cold-load
/// (catalogue × topology says some neuron can host it).
pub feasible: bool,
}
impl RouterState {
pub fn from_config(config: &RouterConfig) -> Self {
let topology = config
.cortexes
.iter()
.map(|c| (c.name.clone(), CortexTopology::default()))
.collect();
Self {
cortexes: config.cortexes.clone(),
http_client: reqwest::Client::new(),
poll_interval: Duration::from_secs(config.router.poll_interval_secs),
topology: RwLock::new(topology),
}
}
/// 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(|m| m.feasible))
.map(|(name, _)| name.clone())
.collect()
}
}

View File

@@ -0,0 +1,93 @@
//! Skeleton acceptance tests for #70: the router builds, serves `/health`
//! and `/v1/models` on a plaintext port, and loads its cortex-endpoint list
//! from TOML with env overrides.
use helexa_router::config::{CortexEndpoint, RouterConfig};
use helexa_router::state::RouterState;
use std::sync::Arc;
use tokio::net::TcpListener;
/// Bind the router app on an ephemeral port and return its base URL.
async fn spawn_router(cortexes: Vec<CortexEndpoint>) -> String {
let cfg = RouterConfig {
cortexes,
..Default::default()
};
let state = Arc::new(RouterState::from_config(&cfg));
let app = helexa_router::build_app(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}")
}
#[tokio::test]
async fn health_reports_configured_cortex_count() {
let base = spawn_router(vec![
CortexEndpoint {
name: "a".into(),
endpoint: "https://a.example.com".into(),
},
CortexEndpoint {
name: "b".into(),
endpoint: "https://b.example.com".into(),
},
])
.await;
let body: serde_json::Value = reqwest::get(format!("{base}/health"))
.await
.unwrap()
.json()
.await
.unwrap();
assert_eq!(body["status"], "ok");
assert_eq!(body["cortexes"]["configured"], 2);
}
#[tokio::test]
async fn models_returns_empty_openai_list() {
let base = spawn_router(vec![]).await;
let resp = reqwest::get(format!("{base}/v1/models")).await.unwrap();
assert!(resp.status().is_success());
let body: serde_json::Value = resp.json().await.unwrap();
assert_eq!(body["object"], "list");
assert_eq!(body["data"].as_array().unwrap().len(), 0);
}
#[test]
#[allow(clippy::result_large_err)]
fn config_loads_from_toml_with_env_override() {
figment::Jail::expect_with(|jail| {
jail.create_file(
"helexa-router.toml",
r#"
[router]
listen = "127.0.0.1:8088"
[[cortexes]]
name = "lair-cafe"
endpoint = "https://cortex.lair.cafe"
"#,
)?;
// Env override wins over the TOML value.
jail.set_env("HELEXA_ROUTER_ROUTER__LISTEN", "0.0.0.0:9099");
let cfg = RouterConfig::load("helexa-router.toml").expect("load config");
assert_eq!(cfg.router.listen, "0.0.0.0:9099");
assert_eq!(cfg.cortexes.len(), 1);
assert_eq!(cfg.cortexes[0].name, "lair-cafe");
assert_eq!(cfg.cortexes[0].endpoint, "https://cortex.lair.cafe");
Ok(())
});
}

View File

@@ -0,0 +1,169 @@
//! 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;
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(),
}],
..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.
let coder = c1.models.get("Qwen/Qwen3-Coder-30B").unwrap();
assert!(coder.loaded && coder.feasible);
let vl = c1.models.get("Qwen/Qwen3-VL-8B").unwrap();
assert!(!vl.loaded && vl.feasible);
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);
}

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

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

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

View File

@@ -13,6 +13,7 @@ use axum::response::sse::{Event, KeepAlive, Sse};
use axum::response::{IntoResponse, Json};
use axum::routing::{get, post};
use cortex_core::discovery::{DiscoveryResponse, HealthResponse};
use cortex_core::entitlements::{HEADER_ACCOUNT_ID, HEADER_KEY_ID};
use cortex_core::harness::ModelSpec;
use cortex_core::openai::{ChatCompletionRequest, MessageContent};
use cortex_core::responses::{ResponsesRequest, ResponsesUsage};
@@ -71,6 +72,12 @@ async fn health_handler(State(state): State<Arc<NeuronState>>) -> Json<HealthRes
// know about activation lifecycle.
let mut snapshot = state.health_cache.snapshot().await;
snapshot.activation = state.activation.snapshot().await;
// Per-model admission load (#53) — read live from the candle harness so
// cortex's load-aware router (#55) can spread traffic and propagate
// backpressure. Absent when no candle harness is present.
if let Some(candle) = &state.candle {
snapshot.models = candle.load_snapshot().await;
}
Json(snapshot)
}
@@ -228,6 +235,17 @@ fn default_enable_thinking(req: &mut ChatCompletionRequest, include_thinking: bo
}
}
/// The request's principal for fair-share admission (#54), reconstructed
/// from the internal headers cortex stamps (#49). cortex strips any
/// client-supplied copy and asserts the authoritative value, so over the
/// trusted WireGuard link these are safe to key fair-share on. `None` for an
/// unauthenticated/direct request — exempt from the per-principal cap.
fn principal_key(headers: &axum::http::HeaderMap) -> Option<String> {
let account = headers.get(HEADER_ACCOUNT_ID)?.to_str().ok()?;
let key = headers.get(HEADER_KEY_ID)?.to_str().ok()?;
Some(format!("{account}/{key}"))
}
/// OpenAI-compatible chat completions. Dispatches to streaming SSE when
/// `stream: true` is set on the request; otherwise returns a single
/// `ChatCompletionResponse`.
@@ -271,8 +289,14 @@ async fn chat_completions(
// true`) keep reasoning on.
default_enable_thinking(&mut req, include_thinking);
// Fair-share admission principal (#54), from cortex's stamped headers.
let principal = principal_key(&headers);
if req.stream.unwrap_or(false) {
match candle.chat_completion_stream_with(req, chat_config).await {
match candle
.chat_completion_stream_with(req, chat_config, principal)
.await
{
Ok(rx) => {
// Each chunk → one SSE `data: {json}` line. After the
// channel closes, append the OpenAI [DONE] terminator.
@@ -289,7 +313,7 @@ async fn chat_completions(
Err(e) => inference_error_response(e),
}
} else {
match candle.chat_completion(req).await {
match candle.chat_completion(req, principal).await {
Ok(resp) => Json(resp).into_response(),
Err(e) => inference_error_response(e),
}
@@ -302,6 +326,7 @@ async fn chat_completions(
/// event stream into the Responses event family.
async fn responses(
State(state): State<Arc<NeuronState>>,
headers: axum::http::HeaderMap,
Json(req): Json<ResponsesRequest>,
) -> impl IntoResponse {
let Some(candle) = state.candle.as_ref().map(Arc::clone) else {
@@ -336,9 +361,12 @@ async fn responses(
};
chat_req.stream = Some(stream_requested);
// Fair-share admission principal (#54), from cortex's stamped headers.
let principal = principal_key(&headers);
if stream_requested {
match candle
.responses_stream(chat_req, response_id, message_item_id)
.responses_stream(chat_req, response_id, message_item_id, principal)
.await
{
Ok(rx) => {
@@ -362,7 +390,7 @@ async fn responses(
// and translate the result. We don't currently re-tokenise
// to compute usage; the harness returns it via the chat
// response and we pass it through.
match candle.chat_completion(chat_req).await {
match candle.chat_completion(chat_req, principal).await {
Ok(chat_resp) => {
// Extract the assistant text (chat completions
// always emits one choice on the candle path).
@@ -486,6 +514,24 @@ fn inference_error_response(err: InferenceError) -> axum::response::Response {
"template_render_failed",
format!("chat template could not render this request: {detail}"),
),
// Admission control refused on load (#53): a fast, retryable "busy"
// signal. 503 (service busy) + Retry-After; opencode/AI SDK back off.
InferenceError::Overloaded { retry_after_secs } => OpenAiError::new(
503,
"rate_limit_error",
"rate_limit_exceeded",
"model is busy (admission queue full); retry shortly",
)
.with_retry_after(retry_after_secs),
// Per-principal fair-share cap (#54): 429 rate_limit_exceeded +
// Retry-After — the caller is sending too many concurrent requests.
InferenceError::PerPrincipalLimit { retry_after_secs } => OpenAiError::new(
429,
"rate_limit_error",
"rate_limit_exceeded",
"too many concurrent requests for this key; retry shortly",
)
.with_retry_after(retry_after_secs),
InferenceError::Other(e) => OpenAiError::without_code(500, "api_error", format!("{e:#}")),
};
envelope_response(env)
@@ -660,6 +706,26 @@ mod error_envelope_tests {
assert_eq!(error["required_mb"], 8_192);
}
#[tokio::test]
async fn overloaded_is_503_rate_limited_with_retry_after() {
// Admission rejection (#53) → fast, retryable backpressure.
let resp = inference_error_response(InferenceError::Overloaded {
retry_after_secs: 7,
});
assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
let retry = resp
.headers()
.get(axum::http::header::RETRY_AFTER)
.expect("admission rejection must advertise Retry-After");
assert_eq!(retry.to_str().unwrap(), "7");
let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
.await
.unwrap();
let body: Value = serde_json::from_slice(&bytes).unwrap();
assert_eq!(body["error"]["code"], "rate_limit_exceeded");
}
#[tokio::test]
async fn insufficient_vram_carries_retry_after() {
// Transient 503 — VRAM frees as in-flight requests finish, so the

View File

@@ -85,6 +85,68 @@ pub struct CandleHarnessConfig {
/// `/models`, and enforces it. These knobs tune that derivation.
#[serde(default)]
pub context_limit: ContextLimitConfig,
/// Admission control (#53): bounds the per-model wait queue so a busy
/// model returns a fast, retryable `429`/`503` instead of stalling new
/// requests until their client times out.
#[serde(default)]
pub admission: AdmissionConfig,
}
/// `[harness.candle.admission]` settings (#53).
///
/// Inference is batch-1, so `max_in_flight` is 1 in practice; the queue
/// (`max_queue_depth`) absorbs short bursts, and `max_wait_secs` caps how
/// long a queued request waits before it's refused with backpressure.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AdmissionConfig {
/// Concurrent running requests per model. Batch-1 inference → 1.
#[serde(default = "default_admission_max_in_flight")]
pub max_in_flight: usize,
/// Queued (waiting) requests allowed beyond the in-flight one. The
/// `(max_in_flight + max_queue_depth + 1)`-th request is refused
/// immediately with `429`/`503` + `Retry-After`.
#[serde(default = "default_admission_max_queue_depth")]
pub max_queue_depth: usize,
/// Maximum seconds a queued request waits for the in-flight slot before
/// it is refused (turns the old ~300s client-side hang into a fast,
/// honest signal).
#[serde(default = "default_admission_max_wait_secs")]
pub max_wait_secs: u64,
/// Per-principal fair-share cap (#54): max in-flight + queued requests
/// for any single principal (resolved from the `x-helexa-*` headers
/// cortex stamps), so one client can't monopolize the queue while others
/// wait. Over-cap → `429 rate_limit_exceeded` + `Retry-After`. `0`
/// disables the cap; anonymous requests are always exempt.
#[serde(default = "default_admission_max_per_principal")]
pub max_per_principal: usize,
}
impl Default for AdmissionConfig {
fn default() -> Self {
Self {
max_in_flight: default_admission_max_in_flight(),
max_queue_depth: default_admission_max_queue_depth(),
max_wait_secs: default_admission_max_wait_secs(),
max_per_principal: default_admission_max_per_principal(),
}
}
}
fn default_admission_max_in_flight() -> usize {
1
}
fn default_admission_max_queue_depth() -> usize {
8
}
fn default_admission_max_wait_secs() -> u64 {
30
}
fn default_admission_max_per_principal() -> usize {
2
}
/// `[harness.candle.prefix_cache]` settings.

View File

@@ -0,0 +1,298 @@
//! Per-model admission control (#53).
//!
//! Inference against a loaded model is batch-1: one request runs at a time,
//! serialized by the model's `inference_lock` (single-GPU) / `pool` mutex
//! (TP). Before this, the wait for that lock was an **unbounded FIFO of
//! mutex waiters with no timeout** — a busy model made every new request
//! hang until its client gave up (~300s) with an opaque error.
//!
//! [`AdmissionController`] replaces that implicit unbounded wait with an
//! explicit bounded scheduler: at most `max_in_flight` running (1, batch-1)
//! plus a bounded queue of `max_queue_depth` waiters, each waiting at most
//! `max_wait`. When the queue is full or the wait elapses, the request is
//! rejected *immediately* — an honest, fast, retryable "busy" signal
//! (`429`/`503` + `Retry-After` per #63) instead of a silent stall.
//!
//! The controller is pure async (no CUDA), so the inference paths just call
//! [`AdmissionController::enter`] before taking the inference lock and hold
//! the returned [`AdmissionPermit`] for the request's lifetime. Its counters
//! ([`in_flight`](AdmissionController::in_flight) /
//! [`queue_depth`](AdmissionController::queue_depth)) are lock-free, so
//! `/health` can read live load without contending with inference.
use crate::config::AdmissionConfig;
use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use std::time::Duration;
use tokio::sync::{OwnedSemaphorePermit, Semaphore};
/// Why admission was refused. All map to the #63 backpressure envelope
/// (`rate_limit_exceeded` + `Retry-After`); they differ in cause (and HTTP
/// status — load → `503`, per-principal → `429`).
#[derive(Debug, Clone, Copy)]
pub enum AdmissionRejection {
/// The bounded wait queue was already full (server-side load).
QueueFull { retry_after_secs: u64 },
/// A queue slot was taken but the in-flight slot didn't free within
/// `max_wait` (server-side load).
Timeout { retry_after_secs: u64 },
/// This principal already has `max_per_principal` requests in flight or
/// queued (#54 fair-share) — one principal can't monopolize the model.
PrincipalCap { retry_after_secs: u64 },
}
impl AdmissionRejection {
pub fn retry_after_secs(&self) -> u64 {
match self {
AdmissionRejection::QueueFull { retry_after_secs }
| AdmissionRejection::Timeout { retry_after_secs }
| AdmissionRejection::PrincipalCap { retry_after_secs } => *retry_after_secs,
}
}
}
/// Admission accounting, mutated under a brief lock (never held across an
/// await). `pending` is queued + in-flight overall; `per_principal` is the
/// same count keyed by principal for fair-share (#54).
#[derive(Default, Debug)]
struct AdmissionState {
pending: usize,
per_principal: HashMap<String, usize>,
}
/// Bounded batch-1 scheduler for one loaded model, with per-principal
/// fair-share.
pub struct AdmissionController {
/// In-flight slots — `max_in_flight` permits (1 for batch-1).
slots: Arc<Semaphore>,
/// Queued + in-flight accounting (overall + per principal).
state: Arc<Mutex<AdmissionState>>,
/// `max_in_flight + max_queue_depth` — the overall rejection threshold.
max_pending: usize,
/// Max in-flight + queued for any single principal (#54). `0` disables.
max_per_principal: usize,
max_in_flight: usize,
max_wait: Duration,
}
impl AdmissionController {
pub fn new(cfg: &AdmissionConfig) -> Self {
// A controller with zero in-flight slots would deadlock; clamp.
let max_in_flight = cfg.max_in_flight.max(1);
Self {
slots: Arc::new(Semaphore::new(max_in_flight)),
state: Arc::new(Mutex::new(AdmissionState::default())),
max_pending: max_in_flight + cfg.max_queue_depth,
max_per_principal: cfg.max_per_principal,
max_in_flight,
max_wait: Duration::from_secs(cfg.max_wait_secs),
}
}
/// Admit a request for `principal` (`None` = anonymous, exempt from the
/// per-principal cap). Reserves a queue slot — fast-rejecting if the
/// overall queue is full or the principal is over its fair-share cap —
/// then waits up to `max_wait` for an in-flight slot. The returned permit
/// must be held for the request's lifetime; dropping it frees the slots.
pub async fn enter(
&self,
principal: Option<&str>,
) -> Result<AdmissionPermit, AdmissionRejection> {
// Decision + reservation under one brief lock so concurrent callers
// can't both slip past the thresholds. No await is held here.
{
let mut st = self.state.lock().expect("admission state poisoned");
if st.pending >= self.max_pending {
return Err(AdmissionRejection::QueueFull {
retry_after_secs: self.retry_hint(st.pending),
});
}
if let Some(p) = principal
&& self.max_per_principal > 0
&& st.per_principal.get(p).copied().unwrap_or(0) >= self.max_per_principal
{
return Err(AdmissionRejection::PrincipalCap {
retry_after_secs: self.retry_hint(st.pending),
});
}
st.pending += 1;
if let Some(p) = principal {
*st.per_principal.entry(p.to_string()).or_insert(0) += 1;
}
}
match tokio::time::timeout(self.max_wait, Arc::clone(&self.slots).acquire_owned()).await {
Ok(Ok(permit)) => Ok(AdmissionPermit {
_permit: permit,
state: Arc::clone(&self.state),
principal: principal.map(str::to_string),
}),
// Semaphore is never closed; treat a closed/elapsed wait the same.
Ok(Err(_)) | Err(_) => {
self.release(principal);
Err(AdmissionRejection::Timeout {
retry_after_secs: self.retry_hint(self.max_pending),
})
}
}
}
/// Roll back a reserved-but-not-admitted slot (wait timed out).
fn release(&self, principal: Option<&str>) {
let mut st = self.state.lock().expect("admission state poisoned");
st.pending = st.pending.saturating_sub(1);
decrement_principal(&mut st.per_principal, principal);
}
/// Requests currently running (holding an in-flight slot).
pub fn in_flight(&self) -> usize {
self.max_in_flight
.saturating_sub(self.slots.available_permits())
}
/// Requests waiting for an in-flight slot.
pub fn queue_depth(&self) -> usize {
let pending = self.state.lock().expect("admission state poisoned").pending;
pending.saturating_sub(self.in_flight())
}
/// Rough `Retry-After`: scale with how backed-up the model is, clamped to
/// a sane band. Without per-request timing this is a heuristic, but it
/// gives well-behaved clients (opencode/AI SDK) a sensible backoff.
fn retry_hint(&self, pending: usize) -> u64 {
let queued = pending.saturating_sub(self.max_in_flight) as u64;
((queued + 1) * 2).clamp(1, 120)
}
}
/// Decrement (and prune at zero) a principal's outstanding count.
fn decrement_principal(map: &mut HashMap<String, usize>, principal: Option<&str>) {
if let Some(p) = principal
&& let Some(count) = map.get_mut(p)
{
*count -= 1;
if *count == 0 {
map.remove(p);
}
}
}
/// Held for a request's lifetime; frees the in-flight + queue slot (and the
/// principal's fair-share slot) on drop.
#[derive(Debug)]
pub struct AdmissionPermit {
_permit: OwnedSemaphorePermit,
state: Arc<Mutex<AdmissionState>>,
principal: Option<String>,
}
impl Drop for AdmissionPermit {
fn drop(&mut self) {
let mut st = self.state.lock().expect("admission state poisoned");
st.pending = st.pending.saturating_sub(1);
decrement_principal(&mut st.per_principal, self.principal.as_deref());
}
}
#[cfg(test)]
mod tests {
use super::*;
/// Config with the per-principal cap disabled (0) — most tests exercise
/// the overall queue with anonymous (`None`) callers.
fn cfg(max_in_flight: usize, max_queue_depth: usize, max_wait_secs: u64) -> AdmissionConfig {
AdmissionConfig {
max_in_flight,
max_queue_depth,
max_wait_secs,
max_per_principal: 0,
}
}
#[tokio::test]
async fn admits_up_to_in_flight_and_reports_load() {
let ctrl = AdmissionController::new(&cfg(1, 4, 30));
assert_eq!(ctrl.in_flight(), 0);
let p = ctrl.enter(None).await.expect("first admits");
assert_eq!(ctrl.in_flight(), 1);
assert_eq!(ctrl.queue_depth(), 0);
drop(p);
assert_eq!(ctrl.in_flight(), 0);
}
#[tokio::test]
async fn rejects_when_queue_full() {
// 1 in-flight + 1 queue slot = capacity 2; the 3rd is refused fast.
let ctrl = Arc::new(AdmissionController::new(&cfg(1, 1, 30)));
let _running = ctrl.enter(None).await.expect("admit running");
// Fill the single queue slot with a waiter that parks on the semaphore.
let ctrl2 = Arc::clone(&ctrl);
let waiter = tokio::spawn(async move { ctrl2.enter(None).await.map(|p| drop(p)) });
// Give the waiter a moment to occupy the queue slot.
tokio::time::sleep(Duration::from_millis(50)).await;
assert_eq!(ctrl.queue_depth(), 1);
// Queue full → immediate QueueFull with a Retry-After hint.
match ctrl.enter(None).await {
Err(AdmissionRejection::QueueFull { retry_after_secs }) => {
assert!(retry_after_secs >= 1)
}
other => panic!("expected QueueFull, got {other:?}"),
}
// Release the runner so the parked waiter can proceed and finish.
drop(_running);
waiter.await.unwrap().unwrap();
}
#[tokio::test]
async fn rejects_on_wait_timeout() {
// Zero queue depth + a runner holding the only slot → a second
// request can't even queue, so it's QueueFull, not Timeout. Use a
// queue of 1 and a tiny max_wait to exercise the timeout path.
let ctrl = Arc::new(AdmissionController::new(&cfg(1, 1, 0)));
let _running = ctrl.enter(None).await.expect("admit running");
// max_wait 0 → the queued request times out almost immediately.
match ctrl.enter(None).await {
Err(AdmissionRejection::Timeout { .. }) => {}
other => panic!("expected Timeout, got {other:?}"),
}
// The timed-out request released its queue slot.
assert_eq!(ctrl.queue_depth(), 0);
}
#[tokio::test]
async fn per_principal_cap_protects_other_principals() {
// Generous overall queue, but each principal capped at 1 in-flight+
// queued. Principal A holds the running slot; A's second request is
// refused (PrincipalCap) rather than occupying the queue, so B's
// single request still gets a queue slot and proceeds.
let cfg = AdmissionConfig {
max_in_flight: 1,
max_queue_depth: 8,
max_wait_secs: 30,
max_per_principal: 1,
};
let ctrl = Arc::new(AdmissionController::new(&cfg));
let _a1 = ctrl.enter(Some("acct-a/key-a")).await.expect("A admits");
// A is over its fair-share cap → fast PrincipalCap, no queue slot taken.
match ctrl.enter(Some("acct-a/key-a")).await {
Err(AdmissionRejection::PrincipalCap { retry_after_secs }) => {
assert!(retry_after_secs >= 1)
}
other => panic!("expected PrincipalCap, got {other:?}"),
}
// B (a different principal) is admitted to the queue and proceeds
// once A releases — it was never stuck behind A's backlog.
let ctrl2 = Arc::clone(&ctrl);
let b = tokio::spawn(async move { ctrl2.enter(Some("acct-b/key-b")).await.map(drop) });
tokio::time::sleep(Duration::from_millis(50)).await;
assert_eq!(ctrl.queue_depth(), 1, "B is queued, not rejected");
drop(_a1);
b.await.unwrap().expect("B is served after A releases");
}
}

View File

@@ -33,7 +33,7 @@ use crate::wire::{
use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering};
#[cfg(feature = "cuda")]
use std::time::Duration;
use std::time::{SystemTime, UNIX_EPOCH};
@@ -81,6 +81,9 @@ pub struct CandleHarness {
/// Context-limit derivation settings (#67), read in `list_models`
/// to compute each model's advertised `limit{context,input,output}`.
context_limit_cfg: crate::config::ContextLimitConfig,
/// Admission-control settings (#53), used to build each loaded model's
/// [`super::admission::AdmissionController`] at load time.
admission_cfg: crate::config::AdmissionConfig,
}
/// Devices/capabilities snapshot of a model entering auto-recovery
@@ -146,6 +149,16 @@ impl LoadedHandle {
}
}
/// Current admission load (#53): `(in_flight, queue_depth)`. Lock-free,
/// so `/health` can read it without contending with inference.
pub fn load(&self) -> (usize, usize) {
match self {
LoadedHandle::Single(m) => (m.admission.in_flight(), m.admission.queue_depth()),
#[cfg(feature = "cuda")]
LoadedHandle::Tp(m) => (m.admission.in_flight(), m.admission.queue_depth()),
}
}
/// Modalities the loaded model supports. Stage B7 (single-GPU) +
/// TP-vision (#12) — both single-GPU and TP loads advertise
/// `"vision"` when a replicated vision tower materialised.
@@ -192,23 +205,50 @@ impl LoadedHandle {
/// `NEURON_MAX_PROMPT_TOKENS`, when explicitly set, is applied as a
/// clamp-only upper bound on the derived `context` — a backstop, not
/// the authority. Unset → no clamp; the derivation stands alone.
pub async fn derived_limit(
/// Refresh the cached free-VRAM reading used by [`Self::derived_limit`]
/// (#53). Queries the device worker — so it MUST run off the request
/// path (background refresher / load-time seed), never from a control
/// endpoint, since the query queues behind inference on the worker.
/// Single-GPU caches the device's free VRAM; TP caches the tightest
/// free across ranks (the same value `derived_limit` used pre-cache).
pub async fn refresh_free_mb(&self) {
let free = match self {
LoadedHandle::Single(m) => m.query_vram().await.0,
#[cfg(feature = "cuda")]
LoadedHandle::Tp(m) => m.query_vram_tightest_free_mb().await,
};
// Don't clobber a good cached value with a transient `0`
// (worker gone/poisoned sentinel).
if free > 0 {
match self {
LoadedHandle::Single(m) => m.last_free_mb.store(free, Ordering::Release),
#[cfg(feature = "cuda")]
LoadedHandle::Tp(m) => m.last_free_mb.store(free, Ordering::Release),
}
}
}
pub fn derived_limit(
&self,
cfg: &crate::config::ContextLimitConfig,
) -> Option<cortex_core::harness::ModelLimit> {
if !cfg.enabled {
return None;
}
// Read the *cached* free VRAM — never query the device worker here.
// This runs on `GET /models`; a live query would queue behind
// inference on the worker thread and stall the control plane (#53).
// The cache is refreshed off the request path (load + background task).
let (profile, free_mb, rate) = match self {
LoadedHandle::Single(m) => (
m.context_profile?,
m.query_vram().await.0,
m.last_free_mb.load(Ordering::Acquire),
m.prefill_rate.get(),
),
#[cfg(feature = "cuda")]
LoadedHandle::Tp(m) => (
m.context_profile?,
m.query_vram_tightest_free_mb().await,
m.last_free_mb.load(Ordering::Acquire),
m.prefill_rate.get(),
),
};
@@ -305,6 +345,10 @@ pub struct LoadedModel {
/// for the TP path (which already had this invariant by accident
/// because the pool lock covered the same window).
pub inference_lock: tokio::sync::Mutex<()>,
/// Bounded admission scheduler (#53). Gated *before* `inference_lock`
/// so a busy model refuses overflow fast instead of growing an
/// unbounded, untimed queue of lock waiters.
pub admission: super::admission::AdmissionController,
/// Open/close token IDs for the reasoning marker this model
/// emits, populated once at load time by probing the tokenizer's
/// added-tokens table. `None` for non-reasoning models or
@@ -374,6 +418,13 @@ pub struct LoadedModel {
/// request-path enforcement reads this — `0` means "not derived yet"
/// → fall back to the static `NEURON_MAX_PROMPT_TOKENS`.
pub derived_input_cap: AtomicUsize,
/// Cached free VRAM (MiB) for the control plane (#53). `derived_limit`
/// (served by `GET /models`) reads this instead of querying the device
/// worker, which during inference is saturated processing forward jobs —
/// a live query would queue behind them and stall `/models`, tripping
/// cortex's health poller into marking the node unhealthy. Refreshed off
/// the request path: seeded at load, then by a background task.
pub last_free_mb: AtomicU64,
}
impl LoadedModel {
@@ -422,6 +473,10 @@ pub struct TpLoadedModel {
/// serialises subprocess RPC traffic on the pool's
/// `Vec<Worker>` channels.
pub pool: tokio::sync::Mutex<super::tp::WorkerPool>,
/// Bounded admission scheduler (#53), mirroring the single-GPU path.
/// Gated before the pool lock so an overloaded TP model returns fast
/// backpressure instead of an unbounded, untimed wait.
pub admission: super::admission::AdmissionController,
/// Handle into the leader device worker's TP slab. The boxed
/// `TpLeaderModel` (with its embedded `Arc<Comm>` clones and
/// per-rank CUDA tensors) lives on the worker thread; we hold an
@@ -482,6 +537,10 @@ pub struct TpLoadedModel {
/// Mint for pool-wide snapshot ids. Plain counter; uniqueness only
/// needs to hold per model lifetime (snapshots die with the model).
pub next_snapshot_id: std::sync::atomic::AtomicU64,
/// Cached tightest free VRAM (MiB) for the control plane (#53) — see
/// [`LoadedModel::last_free_mb`]. Read by `derived_limit` so `GET /models`
/// never fans a VRAM query out to the (inference-saturated) TP workers.
pub last_free_mb: AtomicU64,
}
#[cfg(feature = "cuda")]
@@ -1088,6 +1147,32 @@ fn debug_poison_armed(model_id: &str) -> bool {
armed && !FIRED.swap(true, Ordering::Relaxed)
}
/// Background control-plane VRAM cache refresher (#53). Every few seconds,
/// refreshes each loaded model's `last_free_mb` so `derived_limit` (served
/// by `GET /models`) reads a cached value and never queries the device
/// worker on the request path — a live query would queue behind inference
/// forward jobs on the worker thread, stalling `/models` for seconds and
/// tripping cortex's health poller into evicting the node from routing.
/// Holds a `Weak` so a shutting-down harness lets the task exit. The query
/// itself may queue behind inference, but that only delays this background
/// refresh — no request-path caller is ever blocked.
async fn vram_cache_refresh_loop(weak: std::sync::Weak<CandleHarness>) {
const REFRESH_INTERVAL: std::time::Duration = std::time::Duration::from_secs(5);
loop {
tokio::time::sleep(REFRESH_INTERVAL).await;
let Some(this) = weak.upgrade() else {
return; // harness dropped — exit
};
// Snapshot handles, then release the read lock before awaiting the
// (possibly slow) worker queries so we never hold it across an await.
let handles: Vec<LoadedHandle> = this.models.read().await.values().cloned().collect();
drop(this);
for handle in handles {
handle.refresh_free_mb().await;
}
}
}
/// Background auto-recovery task (#17). Drains poisoned model ids and
/// rebuilds each via [`CandleHarness::recover_one`]. Holds a `Weak` so a
/// shutting-down harness lets the task exit; processes one id at a time,
@@ -1261,25 +1346,67 @@ fn validate_vision_prefill(prompt_len: usize, vram_free_mb: u64) -> Result<(), I
/// the caller as `max`), or if free VRAM is below the floor. Enforcing
/// the *derived* cap means a VRAM-tight host rejects a prompt that
/// wouldn't fit, instead of accepting it and OOMing mid-prefill.
///
/// The third VRAM check — the length-aware backstop (#65) — closes the
/// poll-vs-request snapshot gap #67 leaves open. `max` is
/// `effective_prompt_cap()`, the input budget derived at **/models poll
/// time** from the tightest card's free VRAM *then*. If free VRAM has
/// since dropped (a co-resident model loaded, a concurrent prefill grew
/// its KV), a prompt at-or-below that now-stale cap still clears the
/// static floor yet no longer fits — and OOMs mid-prefill, poisoning the
/// device context (the 2026-05-26 beast incident the #47 work exists to
/// eliminate). So we re-run the same length×KV-vs-VRAM physics #67 uses
/// for the cap, but against **request-time** free VRAM, reusing the
/// model's [`ContextProfile`] rather than re-deriving the KV cost. This
/// gives the text path the live-VRAM guard the vision path already has
/// (`validate_vision_prefill`). `profile`/`kv_bytes_per_token_per_card`
/// are per-card and `vram_free_mb` is the tightest card's free VRAM, so
/// the two are commensurable on both single-GPU and TP loads.
fn validate_request(
prompt_len: usize,
vram_free_mb: u64,
max: usize,
profile: Option<&super::context_limit::ContextProfile>,
cfg: &crate::config::ContextLimitConfig,
) -> Result<(), InferenceError> {
if prompt_len > max {
return Err(InferenceError::PromptTooLong { prompt_len, max });
}
// VRAM check is skipped on CPU loads (vram_free_mb == 0 sentinel)
// VRAM checks are skipped on CPU loads (vram_free_mb == 0 sentinel)
// because the (0, 0) reply from `query_vram` is also what a missing
// worker returns. The CPU path has no per-GPU memory limit anyway —
// host RAM is bounded by the OOM killer, not this check.
if vram_free_mb == 0 {
return Ok(());
}
let min = min_free_vram_mb();
if vram_free_mb != 0 && vram_free_mb < min {
if vram_free_mb < min {
return Err(InferenceError::InsufficientVram {
free_mb: vram_free_mb,
required_mb: min,
});
}
// Length-aware backstop (#65): KV the whole sequence (prompt +
// generation reserve) will occupy, plus the prefill activation
// headroom, plus the static floor as an additive cushion — all per
// card. A degenerate zero-KV profile (no full-attention layers) or a
// model with no captured profile skips this and rides the floor
// check above, mirroring `derive_limit`'s VRAM-ceiling fallback.
if let Some(profile) = profile
&& profile.kv_bytes_per_token_per_card > 0
{
let tokens = (prompt_len as u64).saturating_add(cfg.output_reserve_tokens as u64);
let kv_mb = profile.kv_bytes_per_token_per_card.saturating_mul(tokens) / (1024 * 1024);
let required_mb = kv_mb
.saturating_add(cfg.activation_headroom_mb)
.saturating_add(min);
if required_mb > vram_free_mb {
return Err(InferenceError::InsufficientVram {
free_mb: vram_free_mb,
required_mb,
});
}
}
Ok(())
}
@@ -1565,6 +1692,7 @@ impl CandleHarness {
recovery_tx,
prefix_cache_cfg: config.prefix_cache.clone(),
context_limit_cfg: config.context_limit.clone(),
admission_cfg: config.admission.clone(),
});
// Background auto-recovery task (#17). Holds a `Weak` so it can't
// keep the harness alive. Spawned only when a tokio runtime is
@@ -1573,6 +1701,11 @@ impl CandleHarness {
if tokio::runtime::Handle::try_current().is_ok() {
let weak = Arc::downgrade(&this);
tokio::spawn(recovery_loop(weak, recovery_rx));
// Control-plane VRAM cache refresher (#53): keeps each loaded
// model's `last_free_mb` current off the request path, so
// `derived_limit` / `GET /models` never query the device worker
// (which is saturated during inference) and never stall.
tokio::spawn(vram_cache_refresh_loop(Arc::downgrade(&this)));
}
this
}
@@ -2006,6 +2139,7 @@ impl CandleHarness {
pub async fn chat_completion(
&self,
request: ChatCompletionRequest,
principal: Option<String>,
) -> Result<ChatCompletionResponse, InferenceError> {
let handle = {
let models = self.models.read().await;
@@ -2030,7 +2164,7 @@ impl CandleHarness {
LoadedHandle::Single(m) => m,
#[cfg(feature = "cuda")]
LoadedHandle::Tp(m) => {
return self.chat_completion_tp(m, request).await;
return self.chat_completion_tp(m, request, principal).await;
}
};
@@ -2059,6 +2193,15 @@ impl CandleHarness {
return Err(self.trigger_recovery(&model_id).await);
}
// Admission control (#53): refuse fast if the bounded queue is full
// or the wait elapses, rather than joining an unbounded lock-wait.
// The permit is held for the whole request (released on drop).
let _admit = loaded
.admission
.enter(principal.as_deref())
.await
.map_err(InferenceError::from)?;
// Serialise concurrent requests against this model. Holds for
// the duration of clear_kv_cache → prefill → decode so two
// requests' chunked-prefill sequences can't interleave on the
@@ -2149,7 +2292,13 @@ impl CandleHarness {
"chat_completion: starting"
);
validate_request(prompt_len, vram_free_mb, loaded.effective_prompt_cap())?;
validate_request(
prompt_len,
vram_free_mb,
loaded.effective_prompt_cap(),
loaded.context_profile.as_ref(),
&self.context_limit_cfg,
)?;
if vision_route.is_some() {
validate_vision_prefill(prompt_len, vram_free_mb)?;
}
@@ -2378,8 +2527,13 @@ impl CandleHarness {
pub async fn chat_completion_stream(
&self,
request: ChatCompletionRequest,
principal: Option<String>,
) -> Result<mpsc::Receiver<ChatCompletionChunk>, InferenceError> {
self.chat_completion_stream_with(request, wire_chat::ChatProjectionConfig::default())
self.chat_completion_stream_with(
request,
wire_chat::ChatProjectionConfig::default(),
principal,
)
.await
}
@@ -2391,8 +2545,9 @@ impl CandleHarness {
&self,
request: ChatCompletionRequest,
mut config: wire_chat::ChatProjectionConfig,
principal: Option<String>,
) -> Result<mpsc::Receiver<ChatCompletionChunk>, InferenceError> {
let stream = self.inference_stream(request).await?;
let stream = self.inference_stream(request, principal).await?;
// Fill in the model's reasoning markers if the caller
// didn't pre-populate them — they're a property of the
// loaded model (which the HTTP handler doesn't reach into
@@ -2419,9 +2574,10 @@ impl CandleHarness {
request: ChatCompletionRequest,
response_id: String,
message_item_id: String,
principal: Option<String>,
) -> Result<mpsc::Receiver<crate::wire::openai_responses::ResponseStreamFrame>, InferenceError>
{
let stream = self.inference_stream(request).await?;
let stream = self.inference_stream(request, principal).await?;
let meta = crate::wire::openai_responses::ResponseMeta {
response_id,
created_at: stream.created,
@@ -2442,6 +2598,7 @@ impl CandleHarness {
async fn inference_stream(
&self,
request: ChatCompletionRequest,
principal: Option<String>,
) -> Result<InferenceStream, InferenceError> {
let handle = {
let models = self.models.read().await;
@@ -2466,7 +2623,7 @@ impl CandleHarness {
LoadedHandle::Single(m) => m,
#[cfg(feature = "cuda")]
LoadedHandle::Tp(m) => {
return self.inference_tp_stream(m, request).await;
return self.inference_tp_stream(m, request, principal).await;
}
};
@@ -2595,7 +2752,13 @@ impl CandleHarness {
);
}
validate_request(prompt_len, vram_free_mb, loaded.effective_prompt_cap())?;
validate_request(
prompt_len,
vram_free_mb,
loaded.effective_prompt_cap(),
loaded.context_profile.as_ref(),
&self.context_limit_cfg,
)?;
if vision_route.is_some() {
validate_vision_prefill(prompt_len, vram_free_mb)?;
}
@@ -2610,6 +2773,15 @@ impl CandleHarness {
// role chunk was already sent above, so the client sees
// immediate "stream open" feedback even when this request
// queues behind another for the lock.
// Admission control (#53): refuse before opening the stream if the
// model's bounded queue is full / the wait elapses. The permit moves
// into the inference task and is held until it completes.
let admit = loaded
.admission
.enter(principal.as_deref())
.await
.map_err(InferenceError::from)?;
let tool_schemas = build_tool_schemas(&request);
if let (Some(worker), Some(handle)) = (loaded.worker.clone(), loaded.arch_handle) {
#[cfg(feature = "cuda")]
@@ -2620,6 +2792,7 @@ impl CandleHarness {
let tool_schemas_inner = tool_schemas.clone();
tokio::spawn(
async move {
let _admit = admit;
let _inference_guard = loaded_for_task.inference_lock.lock().await;
match stream_inference_via_worker(
worker,
@@ -2680,6 +2853,7 @@ impl CandleHarness {
let tool_call_tokens_inner = loaded.tool_call_tokens.clone();
let tool_schemas_inner = tool_schemas.clone();
tokio::task::spawn_blocking(move || {
let _admit = admit;
let _g = span_for_task.enter();
// `blocking_lock` is safe here: spawn_blocking runs on
// a dedicated thread, not on the async runtime, so
@@ -2779,6 +2953,24 @@ pub struct InferenceStream {
/// Auto-recovery (#17) — rebuild a poisoned model's device context
/// automatically instead of leaving it bricked until a human reloads.
impl CandleHarness {
/// Per-model admission load for `GET /health` (#53): in-flight + queued
/// counts for every resident model. Lock-free per-model reads, so this
/// only briefly holds the registry read lock to enumerate handles.
pub async fn load_snapshot(&self) -> Vec<cortex_core::discovery::ModelLoad> {
let models = self.models.read().await;
models
.values()
.map(|handle| {
let (in_flight, queue_depth) = handle.load();
cortex_core::discovery::ModelLoad {
id: handle.model_id().to_string(),
in_flight,
queue_depth,
}
})
.collect()
}
/// True while `model_id` is being auto-recovered (its slot is briefly
/// absent from the registry during the reload).
pub async fn is_recovering(&self, model_id: &str) -> bool {
@@ -2890,7 +3082,7 @@ impl Harness for CandleHarness {
// physics + live free VRAM + measured prefill rate. `None`
// for arches without a context profile. `cost` stays
// operator-set in the catalogue, filled by the gateway.
let limit = h.derived_limit(&self.context_limit_cfg).await;
let limit = h.derived_limit(&self.context_limit_cfg);
out.push(ModelInfo {
id: h.model_id().into(),
harness: "candle".into(),
@@ -3128,6 +3320,7 @@ impl Harness for CandleHarness {
worker,
arch_handle,
inference_lock: tokio::sync::Mutex::new(()),
admission: super::admission::AdmissionController::new(&self.admission_cfg),
reasoning_tokens,
tool_call_tokens,
chat_template,
@@ -3139,6 +3332,7 @@ impl Harness for CandleHarness {
context_profile,
prefill_rate: super::context_limit::PrefillRateEma::new(),
derived_input_cap: AtomicUsize::new(0),
last_free_mb: AtomicU64::new(0),
});
if loaded.prefix_cache.is_some() {
tracing::info!(
@@ -3149,6 +3343,14 @@ impl Harness for CandleHarness {
);
}
// Seed the control-plane VRAM cache (#53) while the worker is idle
// (load just finished), so `/models` has a value before the
// background refresher's first tick and never queries the worker.
let (free_mb, _) = loaded.query_vram().await;
if free_mb > 0 {
loaded.last_free_mb.store(free_mb, Ordering::Release);
}
let mut models = self.models.write().await;
models.insert(spec.model_id.clone(), LoadedHandle::Single(loaded));
tracing::info!(model = %spec.model_id, "model loaded");
@@ -3372,6 +3574,7 @@ impl CandleHarness {
tokenizer,
devices: devices.clone(),
pool: TMutex::new(pool),
admission: super::admission::AdmissionController::new(&self.admission_cfg),
leader_handle,
leader_device: leader_device.clone(),
poisoned: AtomicBool::new(false),
@@ -3398,6 +3601,7 @@ impl CandleHarness {
),
prefill_rate: super::context_limit::PrefillRateEma::new(),
derived_input_cap: AtomicUsize::new(0),
last_free_mb: AtomicU64::new(0),
next_snapshot_id: std::sync::atomic::AtomicU64::new(1),
});
if tp_loaded.prefix_cache.is_some() {
@@ -3409,6 +3613,14 @@ impl CandleHarness {
);
}
// Seed the control-plane VRAM cache (#53) — tightest free across
// ranks, while the workers are idle post-load — so `/models` never
// fans a query out to the inference-busy TP workers.
let free_mb = tp_loaded.query_vram_tightest_free_mb().await;
if free_mb > 0 {
tp_loaded.last_free_mb.store(free_mb, Ordering::Release);
}
let mut models = self.models.write().await;
models.insert(spec.model_id.clone(), LoadedHandle::Tp(tp_loaded));
tracing::info!(
@@ -3438,6 +3650,7 @@ impl CandleHarness {
&self,
tp: Arc<TpLoadedModel>,
request: ChatCompletionRequest,
principal: Option<String>,
) -> Result<ChatCompletionResponse, InferenceError> {
// Tag every line of this request with a short req_id so a
// grep over journalctl reconstructs one request even when
@@ -3474,7 +3687,11 @@ impl CandleHarness {
}
let tp_for_marker = Arc::clone(&tp);
let handle = tokio::spawn(chat_completion_tp_inner(tp, request).instrument(span.clone()));
let context_limit_cfg = self.context_limit_cfg.clone();
let handle = tokio::spawn(
chat_completion_tp_inner(tp, request, principal, context_limit_cfg)
.instrument(span.clone()),
);
match handle.await {
Ok(Ok(resp)) => Ok(resp),
Ok(Err(e)) => {
@@ -3545,6 +3762,7 @@ impl CandleHarness {
&self,
tp: Arc<TpLoadedModel>,
request: ChatCompletionRequest,
principal: Option<String>,
) -> Result<InferenceStream, InferenceError> {
if tp.poisoned.load(Ordering::Acquire) {
return Err(self.trigger_recovery(&request.model).await);
@@ -3685,15 +3903,30 @@ impl CandleHarness {
"TP chat_completion (stream): starting"
);
validate_request(prompt_len, vram_free_mb, tp.effective_prompt_cap())?;
validate_request(
prompt_len,
vram_free_mb,
tp.effective_prompt_cap(),
tp.context_profile.as_ref(),
&self.context_limit_cfg,
)?;
if vision_route.is_some() {
validate_vision_prefill(prompt_len, vram_free_mb)?;
}
// Admission control (#53): refuse before opening the stream; the
// permit moves into the orchestration task and is held for its life.
let admit = tp
.admission
.enter(principal.as_deref())
.await
.map_err(InferenceError::from)?;
let tool_schemas = build_tool_schemas(&request);
let tp_for_task = Arc::clone(&tp);
tokio::spawn(
async move {
let _admit = admit;
let mut failure: Option<String> = None;
let mut pool = acquire_pool_lock(&tp_for_task.pool, &model_id).await;
let leader_handle = tp_for_task.leader_handle;
@@ -4196,6 +4429,8 @@ impl CandleHarness {
async fn chat_completion_tp_inner(
tp: Arc<TpLoadedModel>,
request: ChatCompletionRequest,
principal: Option<String>,
context_limit_cfg: crate::config::ContextLimitConfig,
) -> Result<ChatCompletionResponse, InferenceError> {
let req_start = std::time::Instant::now();
let model_id = request.model.clone();
@@ -4279,11 +4514,25 @@ async fn chat_completion_tp_inner(
"TP chat_completion: starting"
);
validate_request(prompt_len, vram_free_mb, tp.effective_prompt_cap())?;
validate_request(
prompt_len,
vram_free_mb,
tp.effective_prompt_cap(),
tp.context_profile.as_ref(),
&context_limit_cfg,
)?;
if vision_route.is_some() {
validate_vision_prefill(prompt_len, vram_free_mb)?;
}
// Admission control (#53): bounded queue + fast reject before joining
// the pool-lock wait. Held for the whole request (released on drop).
let _admit = tp
.admission
.enter(principal.as_deref())
.await
.map_err(InferenceError::from)?;
// Acquire the pool lock for the duration of the request. After
// Phase 3 the leader's TpLeaderModel lives in the device worker
// thread, so the pool lock now serialises only subprocess RPC
@@ -4826,10 +5075,35 @@ pub enum InferenceError {
/// failure mode that hid several client-compat bugs. Maps to 422.
#[error("chat template could not render this request: {detail}")]
TemplateRenderFailed { detail: String },
/// Admission control (#53) refused on load: the model's bounded queue is
/// full or the wait elapsed. Maps to `503 rate_limit_exceeded` +
/// `Retry-After` — a fast, retryable "busy" signal, not a stall.
#[error("model is busy; retry after {retry_after_secs}s")]
Overloaded { retry_after_secs: u64 },
/// Per-principal fair-share cap (#54) exceeded: this principal already
/// has its max requests in flight/queued. Maps to `429
/// rate_limit_exceeded` + `Retry-After`; a well-behaved client self-paces.
#[error("per-principal in-flight limit reached; retry after {retry_after_secs}s")]
PerPrincipalLimit { retry_after_secs: u64 },
#[error(transparent)]
Other(#[from] anyhow::Error),
}
impl From<super::admission::AdmissionRejection> for InferenceError {
fn from(rejection: super::admission::AdmissionRejection) -> Self {
use super::admission::AdmissionRejection;
match rejection {
AdmissionRejection::QueueFull { retry_after_secs }
| AdmissionRejection::Timeout { retry_after_secs } => {
InferenceError::Overloaded { retry_after_secs }
}
AdmissionRejection::PrincipalCap { retry_after_secs } => {
InferenceError::PerPrincipalLimit { retry_after_secs }
}
}
}
}
/// Build the model's prompt from a [`ChatCompletionRequest`].
///
/// Prefers the model's own `chat_template` when one was loaded
@@ -6563,6 +6837,110 @@ mod tests {
assert!(validate_vision_prefill(12_960, 12_445).is_ok());
}
// ── #65: request-time length-aware VRAM backstop (text prefill) ──
/// A beast-like profile: 16 full-attn layers, 4 kv heads, head_dim
/// 256, f16, TP=2 → 32 KiB/token/card (same numbers as the
/// `context_limit` unit tests). At defaults this makes the
/// length-aware footprint `(prompt_len + 8192)/32 + 2048 + 1500` MiB
/// per card.
fn backstop_profile() -> super::super::context_limit::ContextProfile {
super::super::context_limit::ContextProfile {
max_position_embeddings: 262_144,
kv_bytes_per_token_per_card: super::super::context_limit::kv_bytes_per_token(
16, 4, 256, 2, 2,
),
world_size: 2,
}
}
/// A prompt under the cap with ample free VRAM passes; the same
/// prompt over the cap is `PromptTooLong` before any VRAM math.
#[test]
fn validate_request_cap_and_fit() {
let cfg = crate::config::ContextLimitConfig::default();
let profile = backstop_profile();
// Under cap, 40 GB free → fits.
assert!(validate_request(8_000, 40_000, 100_000, Some(&profile), &cfg).is_ok());
// Over the cap → PromptTooLong, independent of VRAM.
assert!(matches!(
validate_request(100_001, 40_000, 100_000, Some(&profile), &cfg),
Err(InferenceError::PromptTooLong { .. })
));
}
/// The CPU sentinel (`vram_free_mb == 0`) skips every VRAM check,
/// including the new length-aware one — host RAM is the OOM killer's
/// problem, not this guard's.
#[test]
fn validate_request_cpu_sentinel_skips_vram() {
let cfg = crate::config::ContextLimitConfig::default();
let profile = backstop_profile();
assert!(validate_request(1_000_000, 0, 2_000_000, Some(&profile), &cfg).is_ok());
}
/// The static floor remains a backstop: free VRAM below
/// `min_free_vram_mb()` is rejected before the length-aware estimate
/// even runs (so `required_mb` is the floor, not the KV footprint).
#[test]
fn validate_request_static_floor_still_binds() {
let cfg = crate::config::ContextLimitConfig::default();
let profile = backstop_profile();
assert!(matches!(
validate_request(10, 800, 100_000, Some(&profile), &cfg),
Err(InferenceError::InsufficientVram {
free_mb: 800,
required_mb: 1500
})
));
}
/// A model with no captured profile (non-qwen3_5 arch) has no
/// length-aware physics to apply, so it rides only the static floor —
/// a fitting prompt with VRAM above the floor passes.
#[test]
fn validate_request_no_profile_rides_floor() {
let cfg = crate::config::ContextLimitConfig::default();
assert!(validate_request(500_000, 5_000, 1_000_000, None, &cfg).is_ok());
}
/// The acceptance test (#65): a cap derived against *ample* free VRAM
/// is later applied at request time against *tightened* free VRAM. A
/// prompt sized exactly at the now-stale `effective_prompt_cap()`
/// clears the cap and the static floor, yet no longer fits — the
/// length-aware backstop catches it with a clean `InsufficientVram`
/// instead of an OOM-poisoned context. Same prompt with the original
/// ample VRAM still passes, proving the guard only bites on staleness.
#[test]
fn validate_request_catches_poll_vs_request_staleness() {
let cfg = crate::config::ContextLimitConfig::default();
let profile = backstop_profile();
// Cap derived at /models poll time with 40 GB free on the tightest
// card — throughput binds, giving input = 87040 (the issue's
// worked beast figure).
let limit = super::super::context_limit::derive_limit(&profile, 40_000, 800.0, None, &cfg);
let cap = limit.input.expect("input budget derived");
assert_eq!(cap, 87_040);
// With that same ample VRAM, a prompt at the cap still fits.
assert!(validate_request(cap, 40_000, cap, Some(&profile), &cfg).is_ok());
// Now free VRAM has dropped to 5 GB between the poll and the
// request (a co-resident model loaded). The prompt is still ≤ cap
// and clears the 1500 MiB floor, but its footprint —
// (87040 + 8192)/32 + 2048 + 1500 = 6524 MiB — exceeds 5000 MiB.
let err = validate_request(cap, 5_000, cap, Some(&profile), &cfg)
.expect_err("stale cap must not let an over-VRAM prompt through");
assert!(matches!(
err,
InferenceError::InsufficientVram {
free_mb: 5_000,
required_mb: 6_524
}
));
}
// ── Tool-call body parsing ───────────────────────────────────────
fn weather_schemas() -> ToolSchemas {

View File

@@ -100,9 +100,9 @@ pub const KV_CACHE_DTYPE_BYTES: usize = 2;
/// state, not a growing cache). Sharded across the TP world: per-rank
/// KV-head count is `n_kv_heads / world_size`.
///
/// `2 ×` accounts for K and V. Shared by the limit derivation here and
/// the per-rank load-time logging in the TP paths (and, in future, by
/// #65's length-aware pre-flight guard).
/// `2 ×` accounts for K and V. Shared by the limit derivation here, the
/// per-rank load-time logging in the TP paths, and #65's request-time
/// length-aware pre-flight guard (`candle::validate_request`).
pub fn kv_bytes_per_token(
n_full_attn_layers: usize,
n_kv_heads: usize,

View File

@@ -1,5 +1,6 @@
//! Harness registry — maps harness names to trait implementations.
pub mod admission;
pub mod arch;
pub mod candle;
pub mod chat_template;

View File

@@ -30,6 +30,9 @@ impl HealthCache {
// direct read from the cache stays a well-typed
// HealthResponse on the wire.
activation: Default::default(),
// Per-model admission load is overlaid by the api handler
// from the candle harness (#53); the cache doesn't own it.
models: Vec::new(),
}),
has_gpus: RwLock::new(false),
}

View File

@@ -114,6 +114,12 @@ async fn test_health_endpoint() {
let body: serde_json::Value = resp.json().await.unwrap();
assert_eq!(body["uptime_secs"], 0);
// Per-model admission load (#53) is always present, even with no models
// loaded (empty array) — cortex's load-aware router (#55) relies on it.
assert!(
body["models"].is_array(),
"/health must expose a models load array"
);
}
#[tokio::test]

View File

@@ -0,0 +1,31 @@
# helexa-router.example.toml — example configuration
#
# Copy to helexa-router.toml and adjust for your environment.
#
# Environment variable overrides use the HELEXA_ROUTER_ prefix with __
# separators:
# HELEXA_ROUTER_ROUTER__LISTEN=0.0.0.0:8088
[router]
# Plaintext listener. Operator/edge nginx terminates client TLS in front of
# 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.
#
# The skeleton only loads this list; capacity/catalogue polling and
# capacity-aware dispatch arrive in later issues.
# [[cortexes]]
# name = "lair-cafe"
# endpoint = "https://cortex.lair.cafe"
# [[cortexes]]
# name = "example-operator"
# endpoint = "https://cortex.example.com"

View File

@@ -26,6 +26,18 @@
# the load to neuron as `scheme:id` so the daemon
# fetches from the right registry. Omit to let
# neuron substitute its own `default_source`.
# cost.* - optional operator-set pricing, surfaced verbatim on
# GET /v1/models for clients (opencode) to display
# spend. USD per 1,000,000 tokens, as numbers:
# cost.input prompt tokens
# cost.output completion tokens
# cost.cache_read cache-hit tokens (optional tier)
# cost.cache_write cache-write tokens (optional tier)
# Absent vs zero is intentional (#68): OMIT the whole
# cost block to mean "price not declared / unknown";
# set cost.input/output = 0.0 to mean "intentionally
# free" (self-hosted). The advertised rate must match
# what metering bills against.
# Tensor-parallel target — needs a neuron with at least 2 large GPUs.
# The example pins to a specific neuron name; adjust or remove the
@@ -41,13 +53,16 @@ pinned_on = ["your-multi-gpu-neuron"]
limit.context = 32768
limit.input = 28672
limit.output = 4096
# Pricing in USD per 1M tokens — 0.0 for self-hosted.
# Pricing in USD per 1M tokens. Explicit 0.0 = intentionally free
# (self-hosted) — distinct from omitting `cost`, which means "not priced".
cost.input = 0.0
cost.output = 0.0
# Static capability hints (unioned with runtime-detected flags).
capabilities = ["text", "reasoning"]
# Mid-size dense model — fits on any single GPU with ≥16 GB VRAM.
# No `cost` block here: this model is "not priced" — /v1/models omits the
# `cost` key for it, so opencode shows spend as unknown rather than $0.
[[models]]
id = "Qwen/Qwen3-8B"
harness = "candle"