Compare commits

...

17 Commits

Author SHA1 Message Date
63f578cb15 feat(#75): aggregate /v1/models across operators (federation catalogue)
All checks were successful
CI / Format (push) Successful in 35s
CI / CUDA type-check (push) Successful in 1m38s
CI / Clippy (push) Successful in 2m18s
CI / Test (push) Successful in 4m54s
CI / Build cortex SRPM (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
The router's /v1/models is now the deduped union of every reachable
cortex's catalogue, so an opencode client doing discovery against the
router resolves the whole federation without knowing about operators or
cortexes (resolves #61's "Router/discovery contract").

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

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

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

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

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

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

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

34
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,39 @@ dependencies = [
"tracing-subscriber",
]
[[package]]
name = "helexa-router"
version = "0.1.16"
dependencies = [
"anyhow",
"axum",
"chrono",
"clap",
"cortex-core",
"figment",
"helexa-stream",
"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

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

@@ -32,6 +32,12 @@ pub struct NodeState {
/// 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.
@@ -130,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

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

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
@@ -42,66 +48,41 @@ pub async fn forward_request(
"proxying request"
);
let mut req_builder = client.post(&url).body(body);
let observer = CortexMetrics::new(model_id, &route.node_name, request_start, usage_sink);
// Forward relevant headers.
for (key, value) in headers.iter() {
if key == "host" || key == "content-length" {
continue; // reqwest sets these
}
req_builder = req_builder.header(key, value);
}
let 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 = %err,
"proxy: upstream request failed (network)"
),
StreamError::ResponseBuild(err) => tracing::warn!(
node = %route.node_name,
url = %url,
error = %err,
"proxy: failed to build response"
),
}
ProxyError::from(e)
})?;
let upstream_resp = match req_builder.send().await {
Ok(r) => r,
Err(e) => {
tracing::warn!(
node = %route.node_name,
url = %url,
error = %e,
"proxy: upstream request failed (network)"
);
return Err(ProxyError::Upstream(e));
}
};
let upstream_status = upstream_resp.status();
if !upstream_status.is_success() {
if !response.status().is_success() {
// Streaming body — can't snippet without breaking the stream
// pass-through. Log status + URL; the client still gets the
// upstream status, just without the leaked body.
tracing::warn!(
node = %route.node_name,
url = %url,
status = upstream_status.as_u16(),
status = response.status().as_u16(),
"proxy: upstream returned non-2xx"
);
}
let status = StatusCode::from_u16(upstream_status.as_u16()).unwrap_or(StatusCode::BAD_GATEWAY);
let resp_headers = upstream_resp.headers().clone();
let stream = TokenMetricsStream::new(
Box::pin(upstream_resp.bytes_stream()),
TokenMetrics::new(model_id, &route.node_name, request_start, usage_sink),
);
let body = Body::from_stream(stream);
let mut response = Response::builder().status(status);
for (key, value) in resp_headers.iter() {
response = response.header(key, value);
}
response.body(body).map_err(|e| {
tracing::warn!(
node = %route.node_name,
url = %url,
error = %e,
"proxy: failed to build response"
);
ProxyError::ResponseBuild(e.to_string())
})
Ok(response)
}
#[derive(Debug, thiserror::Error)]
@@ -112,6 +93,15 @@ pub enum ProxyError {
ResponseBuild(String),
}
impl From<StreamError> for ProxyError {
fn from(e: StreamError) -> Self {
match e {
StreamError::Upstream(err) => ProxyError::Upstream(err),
StreamError::ResponseBuild(msg) => ProxyError::ResponseBuild(msg),
}
}
}
impl IntoResponse for ProxyError {
fn into_response(self) -> Response {
let (status, code, message) = match &self {
@@ -139,9 +129,10 @@ impl IntoResponse for ProxyError {
//
// The proxy never buffers or re-serialises the upstream body — chunks
// are forwarded verbatim. For metrics it observes each chunk's arrival
// time and keeps a bounded tail of the body text, from which the final
// OpenAI `usage` object (present on the last SSE chunk and on
// non-streaming JSON bodies alike) yields engine-truth token counts.
// time and keeps a bounded tail of the body text (via the shared
// `helexa_stream::BodyTail`), from which the final OpenAI `usage` object
// (present on the last SSE chunk and on non-streaming JSON bodies alike)
// yields engine-truth token counts.
//
// Emitted per request, labelled {model, node}:
// cortex_time_to_first_token_seconds (histogram) — first body chunk
@@ -155,37 +146,15 @@ impl IntoResponse for ProxyError {
/// non-streaming bodies.
const TAIL_CAP_BYTES: usize = 64 * 1024;
/// Find the value of the LAST `"key": <integer>` occurrence in `tail`.
/// Pure and chunk-boundary-safe (the tail is contiguous appended text).
/// The quoted-needle form means `completion_tokens` never matches
/// `completion_tokens_details`.
pub(crate) fn last_count_for(tail: &str, key: &str) -> Option<u64> {
let needle = format!("\"{key}\"");
let mut result = None;
for (idx, _) in tail.match_indices(&needle) {
let rest = tail[idx + needle.len()..].trim_start();
let Some(rest) = rest.strip_prefix(':') else {
continue;
};
let rest = rest.trim_start();
let digits: &str = &rest[..rest
.char_indices()
.find(|(_, c)| !c.is_ascii_digit())
.map(|(i, _)| i)
.unwrap_or(rest.len())];
if let Ok(v) = digits.parse::<u64>() {
result = Some(v);
}
}
result
}
struct TokenMetrics {
/// cortex's [`ChunkObserver`]: per-request token metrics plus the
/// per-principal reservation settle. Drives cortex policy over the shared
/// streaming mechanism.
struct CortexMetrics {
labels: [(&'static str, String); 2],
request_start: Instant,
first_chunk: Option<Instant>,
last_chunk: Option<Instant>,
tail: String,
tail: BodyTail,
finished: bool,
/// Per-principal metering hook (#51). Invoked exactly once in `finish`
/// with the observed `(prompt, completion)` so the reservation can be
@@ -193,7 +162,7 @@ struct TokenMetrics {
usage_sink: Option<crate::metering::UsageSink>,
}
impl TokenMetrics {
impl CortexMetrics {
fn new(
model_id: &str,
node_name: &str,
@@ -208,26 +177,19 @@ impl TokenMetrics {
request_start,
first_chunk: None,
last_chunk: None,
tail: String::new(),
tail: BodyTail::new(TAIL_CAP_BYTES),
finished: false,
usage_sink,
}
}
}
impl ChunkObserver for CortexMetrics {
fn observe(&mut self, chunk: &[u8]) {
let now = Instant::now();
self.first_chunk.get_or_insert(now);
self.last_chunk = Some(now);
self.tail.push_str(&String::from_utf8_lossy(chunk));
if self.tail.len() > TAIL_CAP_BYTES {
// Keep the newest half; the usage object is always at the
// very end of the body. Split at a char boundary.
let mut cut = self.tail.len() - TAIL_CAP_BYTES / 2;
while !self.tail.is_char_boundary(cut) {
cut += 1;
}
self.tail.drain(..cut);
}
self.tail.push(chunk);
}
/// Emit the metrics exactly once — called on clean stream end and
@@ -239,8 +201,8 @@ impl TokenMetrics {
}
self.finished = true;
let prompt = last_count_for(&self.tail, "prompt_tokens");
let completion = last_count_for(&self.tail, "completion_tokens");
let prompt = last_count_for(self.tail.as_str(), "prompt_tokens");
let completion = last_count_for(self.tail.as_str(), "completion_tokens");
// Per-model metrics — only when body chunks actually arrived.
if let Some(first) = self.first_chunk {
@@ -280,97 +242,3 @@ impl TokenMetrics {
}
}
}
/// Pass-through stream wrapper that feeds [`TokenMetrics`]. Emits on
/// clean end-of-stream; the Drop impl covers client disconnects.
struct TokenMetricsStream {
inner: BoxStream<'static, Result<bytes::Bytes, reqwest::Error>>,
metrics: TokenMetrics,
}
impl TokenMetricsStream {
fn new(
inner: BoxStream<'static, Result<bytes::Bytes, reqwest::Error>>,
metrics: TokenMetrics,
) -> Self {
Self { inner, metrics }
}
}
impl Stream for TokenMetricsStream {
type Item = Result<bytes::Bytes, reqwest::Error>;
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
let this = self.get_mut();
match this.inner.as_mut().poll_next(cx) {
Poll::Ready(Some(Ok(chunk))) => {
this.metrics.observe(&chunk);
Poll::Ready(Some(Ok(chunk)))
}
Poll::Ready(Some(Err(e))) => Poll::Ready(Some(Err(e))),
Poll::Ready(None) => {
this.metrics.finish();
Poll::Ready(None)
}
Poll::Pending => Poll::Pending,
}
}
}
impl Drop for TokenMetricsStream {
fn drop(&mut self) {
self.metrics.finish();
}
}
#[cfg(test)]
mod tests {
use super::last_count_for;
#[test]
fn extracts_counts_from_final_sse_usage_chunk() {
let tail = concat!(
"data: {\"choices\":[{\"delta\":{\"content\":\"hi\"}}]}\n\n",
"data: {\"choices\":[],\"usage\":{\"prompt_tokens\":225,",
"\"completion_tokens\":42,\"total_tokens\":267}}\n\n",
"data: [DONE]\n\n"
);
assert_eq!(last_count_for(tail, "prompt_tokens"), Some(225));
assert_eq!(last_count_for(tail, "completion_tokens"), Some(42));
}
#[test]
fn extracts_counts_from_non_streaming_body() {
let tail = "{\"choices\":[{\"message\":{\"content\":\"hi\"}}],\
\"usage\":{\"prompt_tokens\": 12, \"completion_tokens\": 7}}";
assert_eq!(last_count_for(tail, "prompt_tokens"), Some(12));
assert_eq!(last_count_for(tail, "completion_tokens"), Some(7));
}
#[test]
fn ignores_details_variants_and_takes_last_occurrence() {
// completion_tokens_details must not shadow completion_tokens,
// and the LAST usage object wins (matters when content echoes
// a usage-shaped string earlier in the stream).
let tail = concat!(
"data: {\"usage\":{\"completion_tokens\":1}}\n\n",
"data: {\"usage\":{\"completion_tokens\":99,",
"\"completion_tokens_details\":{\"reasoning_tokens\":3}}}\n\n"
);
assert_eq!(last_count_for(tail, "completion_tokens"), Some(99));
}
#[test]
fn absent_keys_yield_none() {
assert_eq!(
last_count_for("data: [DONE]\n\n", "completion_tokens"),
None
);
assert_eq!(last_count_for("", "prompt_tokens"), None);
// key present but non-numeric value
assert_eq!(
last_count_for("\"completion_tokens\": null", "completion_tokens"),
None
);
}
}

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,
}
@@ -252,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

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

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,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,35 @@
[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 }
helexa-stream = { path = "../helexa-stream" }
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,243 @@
//! Federation catalogue (#75) — the router's aggregate `/v1/models`.
//!
//! Presents the **deduped union** of every reachable cortex's `/v1/models`
//! as the router's own catalogue, so an opencode client doing discovery
//! against the router resolves the whole federation without knowing about
//! operators or cortexes (resolves #61's "Router/discovery contract").
//!
//! Re-tiering: the fractal design is neuron ← cortex ← router. At the
//! router tier the "nodes" are **cortexes**, so the merged entry's
//! `feasible_on` / `locations` are rewritten to **operator names**, not the
//! neuron names a cortex reports. That keeps the federation view honest
//! ("served by these operators") without leaking each operator's internal
//! topology (neuron names, per-device VRAM) to end users.
//!
//! Conflict resolution when operators advertise the same model with
//! different enrichment:
//! - **`limit`** → the *tightest* (smallest `context`), so a client never
//! overflows the most-constrained operator that might serve it (same rule
//! cortex uses across its neurons).
//! - **`cost`** → the *cheapest* (lowest input, then output), the
//! federation "from" price. Richer policy (a range, region/price-aware
//! selection) couples to #68 and is left as a follow-up.
use crate::state::{CortexTopology, entry_feasible};
use cortex_core::harness::{ModelCost, ModelLimit};
use cortex_core::node::{CortexModelEntry, ModelLocation, ModelStatus};
use std::collections::HashMap;
/// Build the federation catalogue: the deduped union of every reachable
/// cortex's serveable models, merged across operators and sorted by id.
pub fn aggregate_models(topology: &HashMap<String, CortexTopology>) -> Vec<CortexModelEntry> {
// Iterate cortexes in name order so `feasible_on` / `locations` and the
// limit/cost tie-breaks are deterministic regardless of map ordering.
let mut cortexes: Vec<(&String, &CortexTopology)> = topology.iter().collect();
cortexes.sort_by(|a, b| a.0.cmp(b.0));
let mut merged: HashMap<String, CortexModelEntry> = HashMap::new();
for (cortex_name, t) in cortexes {
if !t.reachable {
continue;
}
for entry in t.models.values() {
// Only surface models the cortex can actually serve — a
// catalogue-only entry no neuron can host shouldn't appear in
// the federation view.
if !entry_feasible(entry) {
continue;
}
merged
.entry(entry.id.clone())
.and_modify(|acc| merge_into(acc, cortex_name, entry))
.or_insert_with(|| router_entry(cortex_name, entry));
}
}
let mut out: Vec<CortexModelEntry> = merged.into_values().collect();
out.sort_by(|a, b| a.id.cmp(&b.id));
out
}
/// Seed a federation entry from the first cortex that serves the model,
/// re-tiering `feasible_on` / `locations` to the operator name.
fn router_entry(cortex: &str, e: &CortexModelEntry) -> CortexModelEntry {
CortexModelEntry {
id: e.id.clone(),
object: "model".into(),
created: e.created,
owned_by: e.owned_by.clone(),
loaded: e.loaded,
feasible_on: vec![cortex.to_string()],
locations: loaded_location(cortex, e),
capabilities: e.capabilities.clone(),
limit: e.limit.clone(),
cost: e.cost.clone(),
tool_call: e.tool_call,
reasoning: e.reasoning,
}
}
/// Fold another cortex's view of the same model into the merged entry.
fn merge_into(acc: &mut CortexModelEntry, cortex: &str, e: &CortexModelEntry) {
acc.loaded |= e.loaded;
acc.feasible_on.push(cortex.to_string());
acc.locations.extend(loaded_location(cortex, e));
for cap in &e.capabilities {
if !acc.capabilities.contains(cap) {
acc.capabilities.push(cap.clone());
}
}
acc.tool_call |= e.tool_call;
acc.reasoning |= e.reasoning;
acc.limit = tightest_limit(acc.limit.take(), e.limit.clone());
acc.cost = cheapest_cost(acc.cost.take(), e.cost.clone());
}
/// A single cortex-tier location when the model is loaded at that operator;
/// empty when only cold-loadable. Neuron-level VRAM is deliberately dropped.
fn loaded_location(cortex: &str, e: &CortexModelEntry) -> Vec<ModelLocation> {
if e.loaded {
vec![ModelLocation {
node: cortex.to_string(),
status: ModelStatus::Loaded,
vram_estimate_mb: None,
}]
} else {
Vec::new()
}
}
/// Smaller `context` wins — never advertise more headroom than the
/// most-constrained operator can honour.
fn tightest_limit(a: Option<ModelLimit>, b: Option<ModelLimit>) -> Option<ModelLimit> {
match (a, b) {
(None, x) | (x, None) => x,
(Some(a), Some(b)) => Some(if b.context < a.context { b } else { a }),
}
}
/// Cheapest by (input, output) price — the federation "from" price.
fn cheapest_cost(a: Option<ModelCost>, b: Option<ModelCost>) -> Option<ModelCost> {
match (a, b) {
(None, x) | (x, None) => x,
(Some(a), Some(b)) => Some(if (b.input, b.output) < (a.input, a.output) {
b
} else {
a
}),
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::state::CortexTopology;
fn entry(id: &str, loaded: bool, feasible: bool) -> CortexModelEntry {
CortexModelEntry {
id: id.into(),
object: "model".into(),
created: 0,
owned_by: "helexa".into(),
loaded,
feasible_on: if feasible || loaded {
vec!["some-neuron".into()]
} else {
vec![]
},
locations: vec![],
capabilities: vec![],
limit: None,
cost: None,
tool_call: false,
reasoning: false,
}
}
fn cortex(reachable: bool, entries: Vec<CortexModelEntry>) -> CortexTopology {
CortexTopology {
reachable,
consecutive_failures: 0,
last_poll: None,
healthy_nodes: 1,
total_nodes: 1,
models: entries.into_iter().map(|e| (e.id.clone(), e)).collect(),
}
}
#[test]
fn dedupes_and_merges_availability_across_cortexes() {
let mut topo = HashMap::new();
// c-a: model loaded. c-b: same model only cold-loadable.
topo.insert("c-a".into(), cortex(true, vec![entry("m", true, true)]));
topo.insert("c-b".into(), cortex(true, vec![entry("m", false, true)]));
let out = aggregate_models(&topo);
assert_eq!(out.len(), 1, "duplicate model id collapses to one");
let m = &out[0];
assert!(m.loaded, "loaded somewhere → loaded");
// feasible_on re-tiered to operator names, both present, sorted.
assert_eq!(m.feasible_on, vec!["c-a".to_string(), "c-b".to_string()]);
// Only the loaded operator contributes a location, named by operator.
assert_eq!(m.locations.len(), 1);
assert_eq!(m.locations[0].node, "c-a");
assert_eq!(m.locations[0].vram_estimate_mb, None);
}
#[test]
fn unreachable_cortex_is_excluded() {
let mut topo = HashMap::new();
topo.insert("up".into(), cortex(true, vec![entry("m", true, true)]));
topo.insert(
"down".into(),
cortex(false, vec![entry("other", true, true)]),
);
let out = aggregate_models(&topo);
assert_eq!(out.len(), 1);
assert_eq!(out[0].id, "m");
}
#[test]
fn catalogue_only_infeasible_entries_are_hidden() {
let mut topo = HashMap::new();
topo.insert("c".into(), cortex(true, vec![entry("ghost", false, false)]));
assert!(aggregate_models(&topo).is_empty());
}
#[test]
fn preserves_tightest_limit_and_cheapest_cost() {
let mut a = entry("m", true, true);
a.limit = Some(ModelLimit {
context: 32_768,
input: None,
output: 4096,
});
a.cost = Some(ModelCost {
input: 0.50,
output: 1.50,
cache_read: None,
cache_write: None,
});
let mut b = entry("m", true, true);
b.limit = Some(ModelLimit {
context: 16_384, // tighter
input: None,
output: 4096,
});
b.cost = Some(ModelCost {
input: 0.20, // cheaper
output: 0.80,
cache_read: None,
cache_write: None,
});
let mut topo = HashMap::new();
topo.insert("c-a".into(), cortex(true, vec![a]));
topo.insert("c-b".into(), cortex(true, vec![b]));
let out = aggregate_models(&topo);
assert_eq!(out.len(), 1);
assert_eq!(out[0].limit.as_ref().unwrap().context, 16_384);
assert_eq!(out[0].cost.as_ref().unwrap().input, 0.20);
}
}

View File

@@ -0,0 +1,85 @@
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,
/// This router instance's region (e.g. "eu-west"). When set, dispatch
/// (#73) prefers cortexes whose `region` matches, before falling back to
/// any feasible cortex. `None` → no geo affinity.
#[serde(default)]
pub region: Option<String>,
}
fn default_poll_interval_secs() -> u64 {
10
}
/// One downstream cortex the router may proxy to. The router verifies the
/// cortex's outbound TLS cert (#74) and routes on capacity (#73); it holds
/// no entitlement logic of its own and forwards the client bearer verbatim.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct CortexEndpoint {
/// Human-readable label (e.g. "lair-cafe").
pub name: String,
/// Base URL of the cortex gateway (e.g. "https://cortex.example.com").
pub endpoint: String,
/// Optional region tag (e.g. "eu-west") for geo affinity in dispatch
/// (#73). `None` → no region preference applies to this cortex.
#[serde(default)]
pub region: Option<String>,
}
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(),
region: None,
},
cortexes: vec![],
}
}
}

View File

@@ -0,0 +1,215 @@
//! Capacity-aware dispatch (#73) — the router's data path.
//!
//! Given an inbound request's `model`, pick a reachable cortex that can
//! serve it (preferring warm/loaded, region-affine, higher-headroom),
//! forward the client's bearer **unchanged** (auth stays at cortex), and
//! stream the response back verbatim via the shared [`helexa_stream`]
//! module. Cortex's #63-shaped rejections (`429 rate_limit_exceeded`,
//! `400 context_length_exceeded`, …) pass through untouched. Transport
//! failures fail over to the next feasible cortex; a genuine HTTP response —
//! any status — is returned as-is and never retried away.
//!
//! The router holds **no entitlement logic**: it routes on capacity, not
//! budget.
use crate::config::CortexEndpoint;
use crate::error::envelope_response;
use crate::state::RouterState;
use axum::body::Bytes;
use axum::http::HeaderMap;
use axum::response::Response;
use cortex_core::error_envelope::OpenAiError;
use helexa_stream::{ChunkObserver, StreamError};
use std::cmp::Reverse;
use std::collections::HashMap;
/// Retry-After hint (seconds) on the router's own transient rejections.
const RETRY_AFTER_SECS: u64 = 5;
/// Outcome of choosing where to send a request.
#[derive(Debug, PartialEq, Eq)]
pub enum Selection {
/// Feasible reachable cortexes, best-first (failover order).
Candidates(Vec<CortexEndpoint>),
/// Some cortex knows the model but none are reachable right now → 503.
NoReachableCapacity,
/// No configured cortex serves the model at all → 404.
UnknownModel,
}
/// Rank the reachable cortexes that can serve `model`, best-first.
///
/// Ordering (each a tie-break for the next): loaded/warm before cold-loadable
/// · region match before not · more healthy nodes before fewer · name for
/// determinism.
pub async fn select_cortexes(state: &RouterState, model: &str) -> Selection {
let topo = state.topology.read().await;
let by_name: HashMap<&str, &CortexEndpoint> = state
.cortexes
.iter()
.map(|c| (c.name.as_str(), c))
.collect();
let mut ranked: Vec<Ranked> = Vec::new();
let mut known_anywhere = false;
for (name, t) in topo.iter() {
let Some(entry) = t.models.get(model) else {
continue;
};
if !crate::state::entry_feasible(entry) {
continue;
}
// Known even via an unreachable cortex's last-good poll — lets us
// tell "temporarily down" (503) from "nobody serves it" (404).
known_anywhere = true;
if !t.reachable {
continue;
}
let Some(ep) = by_name.get(name.as_str()) else {
continue;
};
let region_match = match (&state.region, &ep.region) {
(Some(r), Some(cr)) => r == cr,
_ => false,
};
ranked.push(Ranked {
loaded: entry.loaded,
region_match,
healthy_nodes: t.healthy_nodes,
endpoint: (*ep).clone(),
});
}
if ranked.is_empty() {
return if known_anywhere {
Selection::NoReachableCapacity
} else {
Selection::UnknownModel
};
}
ranked.sort_by(|a, b| {
// false < true, so negate the "good" booleans to sort good first.
(
!a.loaded,
!a.region_match,
Reverse(a.healthy_nodes),
&a.endpoint.name,
)
.cmp(&(
!b.loaded,
!b.region_match,
Reverse(b.healthy_nodes),
&b.endpoint.name,
))
});
Selection::Candidates(ranked.into_iter().map(|r| r.endpoint).collect())
}
struct Ranked {
loaded: bool,
region_match: bool,
healthy_nodes: u32,
endpoint: CortexEndpoint,
}
/// Proxy an inbound inference request to a capacity-bearing cortex.
///
/// `path` is the inference path to forward to (same on the cortex, e.g.
/// `/v1/chat/completions`). The body is parsed only to extract `model`.
pub async fn dispatch(
state: &RouterState,
path: &str,
headers: HeaderMap,
body: Bytes,
) -> Response {
let Some(model) = extract_model(&body) else {
return envelope_response(OpenAiError::new(
400,
"invalid_request_error",
"missing_model_field",
"missing 'model' field in request body",
));
};
let candidates = match select_cortexes(state, &model).await {
Selection::Candidates(c) => c,
Selection::UnknownModel => {
return envelope_response(
OpenAiError::new(
404,
"invalid_request_error",
"model_not_found",
format!("no operator serves model '{model}'"),
)
.with_param("model"),
);
}
Selection::NoReachableCapacity => {
return envelope_response(OpenAiError::service_unavailable(
format!("model '{model}' is temporarily unavailable on all operators"),
Some(RETRY_AFTER_SECS),
));
}
};
// Try candidates in order, failing over only on transport errors. A
// genuine HTTP response (any status — including cortex's #63 429/400)
// is returned verbatim and never retried away.
for ep in &candidates {
let url = format!("{}{}", ep.endpoint, path);
tracing::info!(cortex = %ep.name, url = %url, model = %model, "dispatching");
match helexa_stream::forward_streaming(
&state.http_client,
&url,
headers.clone(),
body.clone(),
NoopObserver,
)
.await
{
Ok(resp) => return resp,
Err(StreamError::Upstream(e)) => {
tracing::warn!(
cortex = %ep.name,
url = %url,
error = %e,
"cortex unreachable; failing over"
);
continue;
}
Err(StreamError::ResponseBuild(msg)) => {
tracing::error!(cortex = %ep.name, error = %msg, "failed to build proxied response");
return envelope_response(OpenAiError::without_code(
500,
"api_error",
"failed to build proxied response",
));
}
}
}
// Every feasible cortex failed to connect.
tracing::warn!(model = %model, tried = candidates.len(), "all feasible operators unreachable");
envelope_response(OpenAiError::service_unavailable(
format!("all operators able to serve '{model}' are unreachable"),
Some(RETRY_AFTER_SECS),
))
}
/// Pull the `model` field out of a request body without re-serialising it.
fn extract_model(body: &Bytes) -> Option<String> {
let v: serde_json::Value = serde_json::from_slice(body).ok()?;
v.get("model")?.as_str().map(str::to_string)
}
/// The router proxies bytes verbatim and keeps no per-request policy, so it
/// needs no observation hooks. (Token metrics/metering stay at cortex.)
struct NoopObserver;
impl ChunkObserver for NoopObserver {
fn observe(&mut self, _chunk: &[u8]) {}
fn finish(&mut self) {}
}

View File

@@ -0,0 +1,27 @@
//! Router adapter from the shared, axum-agnostic
//! [`cortex_core::error_envelope::OpenAiError`] (#60/#63) to an axum
//! [`Response`], setting `Retry-After` when the envelope carries one.
//!
//! cortex-core owns the envelope shape; this is the only place the router
//! crosses from that data into axum. Mirrors cortex-gateway's adapter so
//! the router's own rejections (no feasible operator, all unreachable) are
//! the same #63-shaped envelopes clients already understand — distinct from
//! cortex's rejections, which the router proxies through verbatim.
use axum::http::{HeaderValue, StatusCode, header};
use axum::response::{IntoResponse, Json, Response};
use cortex_core::error_envelope::OpenAiError;
/// Render an [`OpenAiError`] as an axum response (status + JSON envelope +
/// optional `Retry-After`).
pub fn envelope_response(err: OpenAiError) -> Response {
let status = StatusCode::from_u16(err.status).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
let retry_after = err.retry_after_secs;
let mut response = (status, Json(err.body())).into_response();
if let Some(secs) = retry_after
&& let Ok(value) = HeaderValue::from_str(&secs.to_string())
{
response.headers_mut().insert(header::RETRY_AFTER, value);
}
response
}

View File

@@ -0,0 +1,89 @@
use crate::state::RouterState;
use crate::{catalogue, dispatch};
use axum::body::Bytes;
use axum::http::HeaderMap;
use axum::response::Response;
use axum::{Json, Router, extract::State, routing::get, routing::post};
use serde_json::{Value, json};
use std::sync::Arc;
/// Routes served by the router. Inference paths are capacity-aware-dispatched
/// (#73) to a downstream cortex; `/health` and a stub `/v1/models` are local.
pub fn api_routes() -> Router<Arc<RouterState>> {
Router::new()
.route("/v1/chat/completions", post(chat_completions))
.route("/v1/completions", post(completions))
.route("/v1/responses", post(responses))
.route("/v1/messages", post(messages))
.route("/v1/models", get(list_models))
.route("/health", get(health))
.route("/", get(health))
}
// ── Inference paths — forwarded verbatim to a chosen cortex ──────────
//
// Each handler dispatches to the same path on a capacity-bearing cortex.
// The body is parsed only to read `model`; the bearer and bytes are
// forwarded unchanged, and the SSE response streams back verbatim.
async fn chat_completions(
State(state): State<Arc<RouterState>>,
headers: HeaderMap,
body: Bytes,
) -> Response {
dispatch::dispatch(&state, "/v1/chat/completions", headers, body).await
}
async fn completions(
State(state): State<Arc<RouterState>>,
headers: HeaderMap,
body: Bytes,
) -> Response {
dispatch::dispatch(&state, "/v1/completions", headers, body).await
}
async fn responses(
State(state): State<Arc<RouterState>>,
headers: HeaderMap,
body: Bytes,
) -> Response {
dispatch::dispatch(&state, "/v1/responses", headers, body).await
}
async fn messages(
State(state): State<Arc<RouterState>>,
headers: HeaderMap,
body: Bytes,
) -> Response {
dispatch::dispatch(&state, "/v1/messages", headers, body).await
}
/// `GET /health` — router liveness plus a summary of downstream cortex
/// reachability from the topology poller (#72). `status` reflects the
/// router process itself (always `ok` if it answers); downstream health is
/// 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` — the federation catalogue (#75): the deduped union of
/// every reachable cortex's `/v1/models`, so a client doing discovery
/// against the router resolves the whole federation without knowing about
/// operators or cortexes.
async fn list_models(State(state): State<Arc<RouterState>>) -> Json<Value> {
let topo = state.topology.read().await;
let data: Vec<Value> = catalogue::aggregate_models(&topo)
.iter()
.map(|e| json!(e))
.collect();
Json(json!({ "object": "list", "data": data }))
}

View File

@@ -0,0 +1,60 @@
//! 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 catalogue;
pub mod config;
pub mod dispatch;
pub mod error;
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,145 @@
//! Background poller that refreshes the multi-operator topology (#72).
//!
//! The same pattern as cortex↔neuron, one tier up: periodically poll each
//! configured cortex's `GET /v1/models` (catalogue × topology feasibility +
//! loaded state) and `GET /health` (coarse node-health/load), building the
//! live map the dispatcher (#73) routes on. An unreachable or erroring
//! cortex is debounced over [`POLL_FAILURE_THRESHOLD`] consecutive misses,
//! then flipped unhealthy and excluded from routing; it recovers on the
//! next successful poll.
use crate::state::RouterState;
use chrono::Utc;
use cortex_core::node::CortexModelEntry;
use serde::Deserialize;
use std::time::Duration;
/// Per-cortex HTTP timeout for each poll request.
const POLL_TIMEOUT: Duration = Duration::from_secs(5);
/// Consecutive failed polls before a cortex is marked unreachable. Mirrors
/// cortex's neuron-poll debounce: a single blip (a busy cortex briefly slow
/// to answer) can't yank it — and all its models — out of routing.
pub const POLL_FAILURE_THRESHOLD: u32 = 3;
/// cortex's `/v1/models` envelope — `{ "object": "list", "data": [...] }`.
#[derive(Debug, Deserialize)]
struct ModelsEnvelope {
#[serde(default)]
data: Vec<CortexModelEntry>,
}
/// The subset of cortex's `/health` the router reads.
#[derive(Debug, Deserialize)]
struct CortexHealth {
nodes: CortexHealthNodes,
}
#[derive(Debug, Deserialize)]
struct CortexHealthNodes {
healthy: u32,
total: u32,
}
/// Run forever, polling all cortexes on the configured interval.
pub async fn poll_loop(state: std::sync::Arc<RouterState>) {
loop {
poll_once(&state).await;
tokio::time::sleep(state.poll_interval).await;
}
}
/// Poll every configured cortex once. Public for testing.
pub async fn poll_once(state: &RouterState) {
for cortex in &state.cortexes {
poll_cortex(state, &cortex.name, &cortex.endpoint).await;
}
}
/// Poll one cortex: refresh its model map from `/v1/models`, then its node
/// health from `/health`. A `/v1/models` failure debounces toward
/// unreachable; the `/health` poll is best-effort and never flips
/// reachability on its own (a cortex serving `/v1/models` is routable even
/// if `/health` momentarily isn't).
async fn poll_cortex(state: &RouterState, name: &str, endpoint: &str) {
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| (m.id.clone(), m)).collect();
entry.reachable = true;
entry.consecutive_failures = 0;
entry.last_poll = Some(Utc::now());
tracing::debug!(cortex = name, models = entry.models.len(), "poll ok");
}
Err(reason) => {
entry.consecutive_failures = entry.consecutive_failures.saturating_add(1);
if entry.consecutive_failures >= POLL_FAILURE_THRESHOLD {
entry.reachable = false;
}
tracing::warn!(
cortex = name,
failures = entry.consecutive_failures,
reachable = entry.reachable,
reason,
"cortex poll failed"
);
}
}
drop(topo);
// Best-effort health (node counts). Never flips reachability.
if let Some((healthy, total)) = fetch_health(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,86 @@
use crate::config::{CortexEndpoint, RouterConfig};
use chrono::{DateTime, Utc};
use cortex_core::node::CortexModelEntry;
use std::collections::HashMap;
use std::time::Duration;
use tokio::sync::RwLock;
/// Shared router state: 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 proxying to) cortexes.
pub http_client: reqwest::Client,
/// This router instance's region, for dispatch geo affinity (#73).
pub region: Option<String>,
/// How often the poller refreshes the topology.
pub poll_interval: Duration,
/// Live per-cortex topology, keyed by cortex name. Pre-populated from
/// config (every configured cortex present, `reachable = false`) so the
/// poller and handlers always find an entry; the poller flips
/// reachability and fills the model map.
pub topology: RwLock<HashMap<String, CortexTopology>>,
}
/// Live view of one downstream cortex, refreshed each poll.
#[derive(Debug, Clone, Default)]
pub struct CortexTopology {
/// Whether the cortex is currently routable. Flipped `false` only after
/// [`crate::poller::POLL_FAILURE_THRESHOLD`] consecutive failed polls
/// (debounces transient blips); restored on the next successful poll.
pub reachable: bool,
/// Consecutive failed polls; reset to 0 on success.
pub consecutive_failures: u32,
/// Timestamp of the last successful poll.
pub last_poll: Option<DateTime<Utc>>,
/// Healthy / total neuron counts from the cortex's `/health` (coarse
/// load signal; #73 refines headroom). 0/0 until first health poll.
pub healthy_nodes: u32,
pub total_nodes: u32,
/// The cortex's full `/v1/models` entries, keyed by model id. Stored
/// whole (not distilled to a loaded/feasible bool) so the federation
/// catalogue (#75) can preserve per-model `limit`/`cost`/capabilities.
pub models: HashMap<String, CortexModelEntry>,
}
/// Whether a cortex can serve this model — loaded now, or feasible to
/// cold-load (its catalogue × topology says some neuron can host it).
pub fn entry_feasible(entry: &CortexModelEntry) -> bool {
entry.loaded || !entry.feasible_on.is_empty()
}
impl RouterState {
pub fn from_config(config: &RouterConfig) -> Self {
let topology = config
.cortexes
.iter()
.map(|c| (c.name.clone(), CortexTopology::default()))
.collect();
Self {
cortexes: config.cortexes.clone(),
http_client: reqwest::Client::new(),
region: config.router.region.clone(),
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(entry_feasible))
.map(|(name, _)| name.clone())
.collect()
}
}

View File

@@ -0,0 +1,130 @@
//! End-to-end federation-catalogue test for #75: poll two mock cortexes
//! that overlap on a model, then `GET /v1/models` on the router and verify
//! the deduped union with merged availability and preserved limit/cost.
use axum::Router;
use axum::routing::get;
use helexa_router::config::{CortexEndpoint, RouterConfig};
use helexa_router::poller::poll_once;
use helexa_router::state::RouterState;
use serde_json::{Value, json};
use std::sync::Arc;
use tokio::net::TcpListener;
/// Spawn a mock cortex serving the given `/v1/models` `data` array.
async fn spawn_cortex(models: Value) -> String {
let models = Arc::new(models);
let app = Router::new()
.route(
"/v1/models",
get({
let models = Arc::clone(&models);
move || {
let models = Arc::clone(&models);
async move { axum::Json(json!({ "object": "list", "data": &*models })) }
}
}),
)
.route(
"/health",
get(|| async { axum::Json(json!({"status":"ok","nodes":{"healthy":1,"total":1}})) }),
);
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
tokio::spawn(async move {
axum::serve(listener, app).await.unwrap();
});
format!("http://{addr}")
}
/// Spawn the router (with poller) wired to the given cortex endpoints, and
/// poll once synchronously so the topology is populated before we query.
async fn spawn_router(cortexes: Vec<CortexEndpoint>) -> String {
let cfg = RouterConfig {
cortexes,
..Default::default()
};
let state = Arc::new(RouterState::from_config(&cfg));
poll_once(&state).await; // deterministic: fill topology now
let app = helexa_router::build_app(Arc::clone(&state));
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
tokio::spawn(async move {
axum::serve(listener, app).await.unwrap();
});
format!("http://{addr}")
}
fn model(id: &str, loaded: bool, feasible_on: &[&str], ctx: u64, input_cost: f64) -> Value {
json!({
"id": id,
"object": "model",
"created": 0,
"owned_by": "helexa",
"loaded": loaded,
"feasible_on": feasible_on,
"locations": [],
"limit": { "context": ctx, "output": 4096 },
"cost": { "input": input_cost, "output": input_cost * 3.0 }
})
}
#[tokio::test]
async fn federation_catalogue_dedupes_and_preserves_limit_cost() {
// cortex A: "shared" loaded (ctx 32768, $0.50) + "only-a" loaded.
let a = spawn_cortex(json!([
model("shared", true, &["beast"], 32_768, 0.50),
model("only-a", true, &["beast"], 8_192, 1.00),
]))
.await;
// cortex B: "shared" cold-loadable, tighter ctx (16384), cheaper ($0.20).
let b = spawn_cortex(json!([model("shared", false, &["benjy"], 16_384, 0.20)])).await;
let router = spawn_router(vec![
CortexEndpoint {
name: "op-a".into(),
endpoint: a,
region: None,
},
CortexEndpoint {
name: "op-b".into(),
endpoint: b,
region: None,
},
])
.await;
let body: Value = reqwest::get(format!("{router}/v1/models"))
.await
.unwrap()
.json()
.await
.unwrap();
assert_eq!(body["object"], "list");
let data = body["data"].as_array().unwrap();
// Deduped union: "shared" once + "only-a".
assert_eq!(data.len(), 2);
let shared = data.iter().find(|m| m["id"] == "shared").unwrap();
// Loaded somewhere (op-a) → loaded.
assert_eq!(shared["loaded"], true);
// feasible_on re-tiered to operator names, both present, sorted.
let feasible: Vec<&str> = shared["feasible_on"]
.as_array()
.unwrap()
.iter()
.map(|v| v.as_str().unwrap())
.collect();
assert_eq!(feasible, vec!["op-a", "op-b"]);
// Tightest limit (16384) and cheapest cost ($0.20) win.
assert_eq!(shared["limit"]["context"], 16_384);
assert_eq!(shared["cost"]["input"], 0.20);
// Loaded location named by operator, no neuron VRAM leaked.
let locs = shared["locations"].as_array().unwrap();
assert_eq!(locs.len(), 1);
assert_eq!(locs[0]["node"], "op-a");
assert!(data.iter().any(|m| m["id"] == "only-a"));
}

View File

@@ -0,0 +1,300 @@
//! Capacity-aware dispatch acceptance tests for #73.
//!
//! Covers: a request routes to a cortex serving the model; the client's
//! bearer reaches the cortex; cortex's #63 rejections pass through verbatim
//! and are NOT retried away; transport failure fails over to another
//! feasible cortex; unknown model → 404, no reachable capacity → 503; and
//! the selection ranking (warm/region/headroom).
use axum::body::Bytes;
use axum::extract::State;
use axum::http::{HeaderMap, StatusCode};
use axum::response::{IntoResponse, Response};
use axum::routing::post;
use axum::{Json, Router};
use cortex_core::node::CortexModelEntry;
use helexa_router::config::{CortexEndpoint, RouterConfig};
use helexa_router::dispatch::{Selection, dispatch, select_cortexes};
use helexa_router::state::{CortexTopology, RouterState};
use serde_json::{Value, json};
use std::collections::HashMap;
use tokio::net::TcpListener;
/// A minimal `CortexModelEntry` for MODEL with the given serveability.
fn model_entry(loaded: bool, feasible: bool) -> CortexModelEntry {
CortexModelEntry {
id: MODEL.into(),
object: "model".into(),
created: 0,
owned_by: "helexa".into(),
loaded,
feasible_on: if feasible || loaded {
vec!["n".into()]
} else {
vec![]
},
locations: vec![],
capabilities: vec![],
limit: None,
cost: None,
tool_call: false,
reasoning: false,
}
}
const MODEL: &str = "Qwen/Qwen3-Coder-30B";
// ── Mock cortex backend ──────────────────────────────────────────────
/// Behaviour of a mock cortex, carried in axum State.
#[derive(Clone)]
struct MockCortex {
/// Identifies which cortex answered, echoed in the 200 body.
name: &'static str,
/// When true, return a genuine #63-shaped `429 rate_limit_exceeded`.
rate_limited: bool,
}
async fn mock_handler(State(m): State<MockCortex>, headers: HeaderMap) -> Response {
if m.rate_limited {
return (
StatusCode::TOO_MANY_REQUESTS,
Json(json!({"error":{"type":"rate_limit_error","code":"rate_limit_exceeded","message":"slow down","param":null}})),
)
.into_response();
}
let auth = headers
.get("authorization")
.and_then(|v| v.to_str().ok())
.unwrap_or("")
.to_string();
Json(json!({ "served_by": m.name, "auth_seen": auth })).into_response()
}
async fn spawn_cortex(mock: MockCortex) -> String {
let app = Router::new()
.route("/v1/chat/completions", post(mock_handler))
.with_state(mock);
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
tokio::spawn(async move {
axum::serve(listener, app).await.unwrap();
});
format!("http://{addr}")
}
fn ok_cortex(name: &'static str) -> MockCortex {
MockCortex {
name,
rate_limited: false,
}
}
// ── Helpers to build state with a hand-set topology ──────────────────
fn state_with(cortexes: Vec<CortexEndpoint>, region: Option<String>) -> RouterState {
let cfg = RouterConfig {
cortexes,
..Default::default()
};
let mut state = RouterState::from_config(&cfg);
state.region = region;
state
}
/// Overwrite the topology entry for `name` so tests control reachability and
/// model serveability directly (no live poll).
async fn set_topology(
state: &RouterState,
name: &str,
reachable: bool,
loaded: bool,
feasible: bool,
healthy_nodes: u32,
) {
let mut topo = state.topology.write().await;
let mut models = HashMap::new();
models.insert(MODEL.to_string(), model_entry(loaded, feasible));
topo.insert(
name.to_string(),
CortexTopology {
reachable,
consecutive_failures: 0,
last_poll: None,
healthy_nodes,
total_nodes: healthy_nodes,
models,
},
);
}
fn ep(name: &str, endpoint: &str, region: Option<&str>) -> CortexEndpoint {
CortexEndpoint {
name: name.into(),
endpoint: endpoint.into(),
region: region.map(str::to_string),
}
}
fn chat_body() -> Bytes {
Bytes::from(format!("{{\"model\":\"{MODEL}\",\"stream\":false}}"))
}
async fn body_json(resp: Response) -> (StatusCode, Value) {
let status = resp.status();
let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
.await
.unwrap();
let v = serde_json::from_slice(&bytes).unwrap_or(Value::Null);
(status, v)
}
// ── Tests ────────────────────────────────────────────────────────────
#[tokio::test]
async fn routes_to_serving_cortex_and_forwards_bearer() {
let url = spawn_cortex(ok_cortex("c1")).await;
let state = state_with(vec![ep("c1", &url, None)], None);
set_topology(&state, "c1", true, true, true, 2).await;
let mut headers = HeaderMap::new();
headers.insert("authorization", "Bearer sk-test-123".parse().unwrap());
let resp = dispatch(&state, "/v1/chat/completions", headers, chat_body()).await;
let (status, body) = body_json(resp).await;
assert_eq!(status, StatusCode::OK);
assert_eq!(body["served_by"], "c1");
// Bearer reached the cortex unchanged.
assert_eq!(body["auth_seen"], "Bearer sk-test-123");
}
#[tokio::test]
async fn cortex_429_passes_through_and_is_not_retried() {
// c1 (ranked first: loaded) returns a genuine 429; c2 would return 200.
let c1 = spawn_cortex(MockCortex {
name: "c1",
rate_limited: true,
})
.await;
let c2 = spawn_cortex(ok_cortex("c2")).await;
let state = state_with(vec![ep("c1", &c1, None), ep("c2", &c2, None)], None);
// Both reachable + loaded; c1 has more headroom so it ranks first.
set_topology(&state, "c1", true, true, true, 5).await;
set_topology(&state, "c2", true, true, true, 1).await;
let resp = dispatch(
&state,
"/v1/chat/completions",
HeaderMap::new(),
chat_body(),
)
.await;
let (status, body) = body_json(resp).await;
// The genuine 4xx is returned verbatim — NOT retried to c2.
assert_eq!(status, StatusCode::TOO_MANY_REQUESTS);
assert_eq!(body["error"]["code"], "rate_limit_exceeded");
assert!(body.get("served_by").is_none(), "must not have hit c2");
}
#[tokio::test]
async fn fails_over_to_next_cortex_on_transport_error() {
// c_dead ranks first (more headroom) but its endpoint is a closed port;
// c_live is the fallback. The router must fail over and c_live serves.
let live = spawn_cortex(ok_cortex("c_live")).await;
let state = state_with(
vec![
ep("c_dead", "http://127.0.0.1:1", None),
ep("c_live", &live, None),
],
None,
);
set_topology(&state, "c_dead", true, true, true, 9).await;
set_topology(&state, "c_live", true, true, true, 1).await;
let resp = dispatch(
&state,
"/v1/chat/completions",
HeaderMap::new(),
chat_body(),
)
.await;
let (status, body) = body_json(resp).await;
assert_eq!(status, StatusCode::OK);
assert_eq!(body["served_by"], "c_live");
}
#[tokio::test]
async fn unknown_model_is_404() {
let state = state_with(vec![ep("c1", "http://127.0.0.1:1", None)], None);
// Topology has no entry for MODEL at all.
let resp = dispatch(
&state,
"/v1/chat/completions",
HeaderMap::new(),
chat_body(),
)
.await;
let (status, body) = body_json(resp).await;
assert_eq!(status, StatusCode::NOT_FOUND);
assert_eq!(body["error"]["code"], "model_not_found");
}
#[tokio::test]
async fn known_but_all_unreachable_is_503() {
let state = state_with(vec![ep("c1", "http://127.0.0.1:1", None)], None);
// Cortex knows the model (from a prior good poll) but is now unreachable.
set_topology(&state, "c1", false, true, true, 2).await;
let resp = dispatch(
&state,
"/v1/chat/completions",
HeaderMap::new(),
chat_body(),
)
.await;
let (status, body) = body_json(resp).await;
assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
assert_eq!(body["error"]["code"], "service_unavailable");
}
#[tokio::test]
async fn missing_model_field_is_400() {
let state = state_with(vec![ep("c1", "http://127.0.0.1:1", None)], None);
let resp = dispatch(
&state,
"/v1/chat/completions",
HeaderMap::new(),
Bytes::from_static(b"{\"messages\":[]}"),
)
.await;
let (status, body) = body_json(resp).await;
assert_eq!(status, StatusCode::BAD_REQUEST);
assert_eq!(body["error"]["code"], "missing_model_field");
}
#[tokio::test]
async fn ranking_prefers_loaded_then_region_then_headroom() {
// Router is in eu-west. Candidates:
// warm-eu : loaded, region match, 1 node → best
// warm-us : loaded, no region, 9 nodes
// cold-eu : feasible only, region match → worst (cold)
let state = state_with(
vec![
ep("warm-eu", "http://127.0.0.1:1", Some("eu-west")),
ep("warm-us", "http://127.0.0.1:1", Some("us-east")),
ep("cold-eu", "http://127.0.0.1:1", Some("eu-west")),
],
Some("eu-west".into()),
);
set_topology(&state, "warm-eu", true, true, true, 1).await;
set_topology(&state, "warm-us", true, true, true, 9).await;
set_topology(&state, "cold-eu", true, false, true, 5).await;
let Selection::Candidates(order) = select_cortexes(&state, MODEL).await else {
panic!("expected candidates");
};
let names: Vec<&str> = order.iter().map(|e| e.name.as_str()).collect();
assert_eq!(names, vec!["warm-eu", "warm-us", "cold-eu"]);
}

View File

@@ -0,0 +1,95 @@
//! 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(),
region: None,
},
CortexEndpoint {
name: "b".into(),
endpoint: "https://b.example.com".into(),
region: None,
},
])
.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,171 @@
//! Topology-poller acceptance tests for #72: the router maintains a live
//! map of which cortexes serve which models, marks an unreachable/erroring
//! cortex unhealthy and excludes it from routing, and recovers it once
//! reachable again.
use axum::extract::State;
use axum::http::StatusCode;
use axum::routing::get;
use axum::{Json, Router};
use helexa_router::config::{CortexEndpoint, RouterConfig};
use helexa_router::poller::{POLL_FAILURE_THRESHOLD, poll_once};
use helexa_router::state::{RouterState, entry_feasible};
use serde_json::{Value, json};
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use tokio::net::TcpListener;
/// Shared "is this mock cortex up?" flag, toggled by tests to simulate
/// outage and recovery.
#[derive(Clone)]
struct MockState {
up: Arc<AtomicBool>,
}
async fn mock_models(State(s): State<MockState>) -> Result<Json<Value>, StatusCode> {
if !s.up.load(Ordering::SeqCst) {
return Err(StatusCode::SERVICE_UNAVAILABLE);
}
Ok(Json(json!({
"object": "list",
"data": [
{
"id": "Qwen/Qwen3-Coder-30B",
"object": "model",
"created": 0,
"owned_by": "helexa",
"loaded": true,
"feasible_on": ["beast"],
"locations": [{"node": "beast", "status": "loaded", "vram_estimate_mb": 19000}]
},
{
"id": "Qwen/Qwen3-VL-8B",
"object": "model",
"created": 0,
"owned_by": "helexa",
"loaded": false,
"feasible_on": ["beast"],
"locations": []
}
]
})))
}
async fn mock_health(State(s): State<MockState>) -> Result<Json<Value>, StatusCode> {
if !s.up.load(Ordering::SeqCst) {
return Err(StatusCode::SERVICE_UNAVAILABLE);
}
Ok(Json(json!({
"status": "ok",
"nodes": { "healthy": 2, "total": 3 }
})))
}
/// Spawn a mock cortex; returns (base_url, up_flag).
async fn spawn_mock_cortex() -> (String, Arc<AtomicBool>) {
let up = Arc::new(AtomicBool::new(true));
let state = MockState { up: up.clone() };
let app = Router::new()
.route("/v1/models", get(mock_models))
.route("/health", get(mock_health))
.with_state(state);
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
tokio::spawn(async move {
axum::serve(listener, app).await.unwrap();
});
(format!("http://{addr}"), up)
}
fn state_for(name: &str, endpoint: &str) -> RouterState {
let cfg = RouterConfig {
cortexes: vec![CortexEndpoint {
name: name.into(),
endpoint: endpoint.into(),
region: None,
}],
..Default::default()
};
RouterState::from_config(&cfg)
}
#[tokio::test]
async fn poll_builds_live_topology() {
let (base, _up) = spawn_mock_cortex().await;
let state = state_for("c1", &base);
poll_once(&state).await;
let topo = state.topology.read().await;
let c1 = topo.get("c1").expect("cortex present");
assert!(c1.reachable, "should be reachable after a good poll");
assert_eq!(c1.consecutive_failures, 0);
assert!(c1.last_poll.is_some());
assert_eq!((c1.healthy_nodes, c1.total_nodes), (2, 3));
// Loaded model: loaded + feasible. Catalogue-only model: feasible only
// (not loaded, but feasible_on non-empty).
let coder = c1.models.get("Qwen/Qwen3-Coder-30B").unwrap();
assert!(coder.loaded && entry_feasible(coder));
let vl = c1.models.get("Qwen/Qwen3-VL-8B").unwrap();
assert!(!vl.loaded && entry_feasible(vl));
drop(topo);
// The routing helper sees both serveable models on the reachable cortex.
assert_eq!(
state.cortexes_serving("Qwen/Qwen3-VL-8B").await,
vec!["c1".to_string()]
);
}
#[tokio::test]
async fn unreachable_cortex_excluded_then_recovers() {
let (base, up) = spawn_mock_cortex().await;
let state = state_for("c1", &base);
// Healthy first.
poll_once(&state).await;
assert!(state.topology.read().await["c1"].reachable);
// Take it down. The first failures debounce (stay reachable) until the
// threshold; only then is it excluded.
up.store(false, Ordering::SeqCst);
for i in 1..POLL_FAILURE_THRESHOLD {
poll_once(&state).await;
assert!(
state.topology.read().await["c1"].reachable,
"still reachable after {i} failure(s) (below threshold)"
);
}
poll_once(&state).await; // crosses the threshold
{
let topo = state.topology.read().await;
assert!(!topo["c1"].reachable, "excluded after threshold failures");
assert!(topo["c1"].consecutive_failures >= POLL_FAILURE_THRESHOLD);
}
// Excluded from routing.
assert!(
state
.cortexes_serving("Qwen/Qwen3-Coder-30B")
.await
.is_empty()
);
// Bring it back: the next successful poll restores it.
up.store(true, Ordering::SeqCst);
poll_once(&state).await;
let topo = state.topology.read().await;
assert!(topo["c1"].reachable, "recovered after a good poll");
assert_eq!(topo["c1"].consecutive_failures, 0);
}
#[tokio::test]
async fn unconfigured_endpoint_is_unreachable() {
// Nothing listening on this port → polls fail; below threshold it stays
// at its initial unreachable state, and never panics.
let state = state_for("dead", "http://127.0.0.1:1");
poll_once(&state).await;
let topo = state.topology.read().await;
assert!(!topo["dead"].reachable);
assert_eq!(topo["dead"].consecutive_failures, 1);
}

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

@@ -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};
@@ -205,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(),
),
};
@@ -391,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 {
@@ -503,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")]
@@ -1109,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,
@@ -1282,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(())
}
@@ -1595,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
}
@@ -2181,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)?;
}
@@ -2635,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)?;
}
@@ -2959,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(),
@@ -3209,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!(
@@ -3219,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");
@@ -3469,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() {
@@ -3480,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!(
@@ -3546,8 +3687,11 @@ impl CandleHarness {
}
let tp_for_marker = Arc::clone(&tp);
let handle =
tokio::spawn(chat_completion_tp_inner(tp, request, principal).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)) => {
@@ -3759,7 +3903,13 @@ 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)?;
}
@@ -4280,6 +4430,7 @@ 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();
@@ -4363,7 +4514,13 @@ 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)?;
}
@@ -6680,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

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