Compare commits

...

8 Commits

Author SHA1 Message Date
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
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
19 changed files with 1141 additions and 224 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)

14
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,19 @@ dependencies = [
"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,7 @@ members = [
"crates/neuron",
"crates/helexa-acp",
"crates/helexa-bench",
"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,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

@@ -1346,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(())
}
@@ -2250,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)?;
}
@@ -2704,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)?;
}
@@ -3633,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)) => {
@@ -3846,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)?;
}
@@ -4367,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();
@@ -4450,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)?;
}
@@ -6767,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

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